Skip to main content

strided_basic/
kernel.rs

1//! Kernel iteration engine ported from Julia's Strided.jl/src/mapreduce.jl
2//!
3//! This module implements the core iteration engine that follows Julia's
4//! `_mapreduce_kernel!` pattern for cache-optimized strided array operations.
5
6use crate::fuse::{compress_dims, fuse_dims};
7use crate::inline::{InlineVec, OperandVec};
8use crate::{block, order, Result, StridedError};
9
10/// Maximum total elements for the small tensor fast path.
11/// Tensors at or below this size skip `compute_order` and `compute_block_sizes`
12/// since they fit in L1 cache and blocking provides no benefit.
13pub const SMALL_TENSOR_THRESHOLD: usize = 1024;
14
15/// Blocking metadata produced by the execution planners; not a layout-validity proof.
16#[derive(Debug)]
17pub struct KernelPlan {
18    #[cfg_attr(not(test), allow(dead_code))]
19    pub(crate) order: InlineVec<usize>, // outer -> inner
20    pub block: InlineVec<usize>,
21}
22
23/// Build an execution plan for strided iteration (used only in tests).
24///
25/// This follows Julia's `_mapreduce_fuse!` -> `_mapreduce_order!` -> `_mapreduce_block!` pipeline:
26/// 1. Fuse contiguous dimensions
27/// 2. Compute optimal iteration order
28/// 3. Compute block sizes for cache efficiency
29#[cfg(test)]
30pub(crate) fn build_plan(
31    dims: &[usize],
32    strides_list: &[&[isize]],
33    dest_index: Option<usize>,
34    elem_size: usize,
35) -> KernelPlan {
36    let order = order::compute_order(dims, strides_list, dest_index);
37    let block = block::compute_block_sizes(dims, &order, strides_list, elem_size);
38    KernelPlan { order, block }
39}
40
41/// Build an execution plan with dimension fusion.
42///
43/// Pipeline: order -> reorder -> fuse -> block.
44///
45/// Ordering first ensures that dimensions are sorted by stride importance
46/// (smallest stride innermost). Fusing *after* ordering catches contiguous
47/// dimensions regardless of the original memory layout (column-major,
48/// row-major, or any permutation).
49///
50/// Returns `(fused_dims, ordered_strides, KernelPlan)` where dimensions
51/// and strides are already in iteration order (plan.order is identity).
52pub(crate) fn build_plan_fused(
53    dims: &[usize],
54    strides_list: &[&[isize]],
55    dest_index: Option<usize>,
56    elem_size: usize,
57) -> (InlineVec<usize>, OperandVec<InlineVec<isize>>, KernelPlan) {
58    // 1. Compute optimal iteration order on original dims
59    let order = order::compute_order(dims, strides_list, dest_index);
60
61    // 2. Reorder dims and strides
62    let ordered_dims: InlineVec<usize> = order.iter().map(|&d| dims[d]).collect();
63    let ordered_strides: OperandVec<InlineVec<isize>> = strides_list
64        .iter()
65        .map(|strides| order.iter().map(|&d| strides[d]).collect())
66        .collect();
67    let ordered_strides_refs: OperandVec<&[isize]> =
68        ordered_strides.iter().map(|s| s.as_slice()).collect();
69
70    // 3. Fuse contiguous dimensions in ordered space
71    let fused_dims = fuse_dims(&ordered_dims, &ordered_strides_refs);
72
73    // 4. Compress: remove size-1 dimensions to reduce loop depth
74    let (compressed_dims, compressed_strides) = compress_dims(&fused_dims, &ordered_strides_refs);
75    let compressed_strides_refs: OperandVec<&[isize]> =
76        compressed_strides.iter().map(|s| s.as_slice()).collect();
77
78    // 5. Compute blocks with identity ordering (already ordered)
79    let identity: InlineVec<usize> = (0..compressed_dims.len()).collect();
80    let block = block::compute_block_sizes(
81        &compressed_dims,
82        &identity,
83        &compressed_strides_refs,
84        elem_size,
85    );
86
87    (
88        compressed_dims,
89        compressed_strides,
90        KernelPlan {
91            order: identity,
92            block,
93        },
94    )
95}
96
97/// Simplified plan for small tensors that fit in L1 cache.
98///
99/// Skips `compute_order` and `compute_block_sizes` since blocking is
100/// unnecessary for small data. Only fuses and compresses dimensions
101/// to reduce loop depth.
102///
103/// Returns `(fused_dims, fused_strides, KernelPlan)` where the plan's
104/// order is identity and block sizes equal the full dimension sizes.
105pub(crate) fn build_plan_fused_small(
106    dims: &[usize],
107    strides_list: &[&[isize]],
108) -> (InlineVec<usize>, OperandVec<InlineVec<isize>>, KernelPlan) {
109    // Single fuse + compress pass (no ordering needed)
110    let fused = fuse_dims(dims, strides_list);
111    let (fused_dims, fused_strides) = compress_dims(&fused, strides_list);
112
113    // Block = full dimension (no tiling)
114    let block = fused_dims.clone();
115    let identity: InlineVec<usize> = (0..fused_dims.len()).collect();
116
117    (
118        fused_dims,
119        fused_strides,
120        KernelPlan {
121            order: identity,
122            block,
123        },
124    )
125}
126
127// ============================================================================
128// Block-based iteration with inner stride callback
129// ============================================================================
130
131/// Iterate over blocks, calling f with (offsets, block_len, inner_strides).
132///
133/// The callback receives the current byte offsets for each array, the number
134/// of elements in the innermost block, and the innermost strides for each array.
135/// This allows the caller to implement vectorized inner loops.
136#[cfg(test)]
137#[inline]
138pub(crate) fn for_each_inner_block<F>(
139    dims: &[usize],
140    plan: &KernelPlan,
141    strides_list: &[&[isize]],
142    mut f: F,
143) -> Result<()>
144where
145    F: FnMut(&[isize], usize, &[isize]) -> Result<()>,
146{
147    let rank = dims.len();
148    if rank == 0 {
149        let offsets: OperandVec<isize> = strides_list.iter().map(|_| 0isize).collect();
150        return f(&offsets, 1, &[]);
151    }
152
153    // Reorder dimensions and strides according to plan
154    let ordered_dims: InlineVec<usize> = plan.order.iter().map(|&d| dims[d]).collect();
155    let ordered_blocks: InlineVec<usize> = plan.block.clone();
156
157    // Build stride vectors for each array, ordered by plan
158    let mut ordered_strides: OperandVec<InlineVec<isize>> = OperandVec::new();
159    for strides in strides_list {
160        ordered_strides.push(plan.order.iter().map(|&d| strides[d]).collect());
161    }
162    let ordered_strides_refs: OperandVec<&[isize]> =
163        ordered_strides.iter().map(|s| s.as_slice()).collect();
164
165    // Initial offsets (all zero)
166    let mut offsets: OperandVec<isize> = strides_list.iter().map(|_| 0isize).collect();
167
168    // Call the specialized kernel based on rank
169    match rank {
170        1 => kernel_1d_inner(
171            &ordered_dims,
172            &ordered_blocks,
173            &ordered_strides_refs,
174            &mut offsets,
175            &mut f,
176        ),
177        2 => kernel_2d_inner(
178            &ordered_dims,
179            &ordered_blocks,
180            &ordered_strides_refs,
181            &mut offsets,
182            &mut f,
183        ),
184        3 => kernel_3d_inner(
185            &ordered_dims,
186            &ordered_blocks,
187            &ordered_strides_refs,
188            &mut offsets,
189            &mut f,
190        ),
191        4 => kernel_4d_inner(
192            &ordered_dims,
193            &ordered_blocks,
194            &ordered_strides_refs,
195            &mut offsets,
196            &mut f,
197        ),
198        5 => kernel_5d_inner(
199            &ordered_dims,
200            &ordered_blocks,
201            &ordered_strides_refs,
202            &mut offsets,
203            &mut f,
204        ),
205        6 => kernel_6d_inner(
206            &ordered_dims,
207            &ordered_blocks,
208            &ordered_strides_refs,
209            &mut offsets,
210            &mut f,
211        ),
212        7 => kernel_7d_inner(
213            &ordered_dims,
214            &ordered_blocks,
215            &ordered_strides_refs,
216            &mut offsets,
217            &mut f,
218        ),
219        8 => kernel_8d_inner(
220            &ordered_dims,
221            &ordered_blocks,
222            &ordered_strides_refs,
223            &mut offsets,
224            &mut f,
225        ),
226        _ => kernel_nd_inner_iterative(
227            &ordered_dims,
228            &ordered_blocks,
229            &ordered_strides_refs,
230            &mut offsets,
231            &mut f,
232        ),
233    }
234}
235
236// ============================================================================
237// Macro-generated rank-specialized kernels (inner-block callback)
238// ============================================================================
239
240// Every loop below advances an offset only when another position along its
241// level follows, and each loop rewinds its own level on completion. Offsets
242// therefore only ever hold reachable positions of the validated layouts (plus
243// the caller's initial offsets), so no intermediate value can overflow even
244// when a layout ends within one stride of `isize::MAX` (issue #243
245// follow-up). The rewind distance is the difference of two reachable offsets,
246// so it is representable as well.
247
248/// Element-level nested loops for ranks >= 2.
249///
250/// Iterates over the element block, calling `f` at each innermost position.
251/// Levels are listed from outermost to innermost (excluding level 0 which is
252/// the callback level). For a rank-N kernel, element levels are N-1, N-2, ..., 1.
253/// Each level leaves the offsets where it found them.
254macro_rules! elem_loops {
255    // Base case: single level (innermost element level).
256    // Iterates blens[$lv] times, calling f and advancing by stride[$lv]
257    // between calls.
258    ($offsets:ident, $strides:ident, $f:ident, $blens:ident, $is:ident; $lv:literal) => {
259        for _i in 0..$blens[$lv] {
260            $f($offsets, $blens[0], &$is)?;
261            if _i + 1 < $blens[$lv] {
262                for (o, s) in $offsets.iter_mut().zip($strides.iter()) {
263                    *o += s[$lv];
264                }
265            }
266        }
267        let _back = $blens[$lv].saturating_sub(1) as isize;
268        for (o, s) in $offsets.iter_mut().zip($strides.iter()) {
269            *o -= _back * s[$lv];
270        }
271    };
272    // Recursive case: outermost level wraps inner levels, which rewind
273    // themselves; this level advances between inner passes.
274    ($offsets:ident, $strides:ident, $f:ident, $blens:ident, $is:ident;
275     $lv:literal, $next:literal $(, $rest:literal)*) => {
276        for _i in 0..$blens[$lv] {
277            elem_loops!($offsets, $strides, $f, $blens, $is; $next $(, $rest)*);
278            if _i + 1 < $blens[$lv] {
279                for (o, s) in $offsets.iter_mut().zip($strides.iter()) {
280                    *o += s[$lv];
281                }
282            }
283        }
284        let _back = $blens[$lv].saturating_sub(1) as isize;
285        for (o, s) in $offsets.iter_mut().zip($strides.iter()) {
286            *o -= _back * s[$lv];
287        }
288    };
289}
290
291/// Block-level nested while loops for ranks >= 2.
292///
293/// Each block level iterates over tiles of the corresponding dimension.
294/// At level 0 (innermost), element loops are executed. Each block level
295/// advances to the next tile only when one follows and rewinds itself on
296/// completion; inner levels rewind themselves.
297macro_rules! block_loop {
298    // Base case: single block level (the innermost).
299    ($dims:ident, $blocks:ident, $strides:ident, $offsets:ident, $f:ident,
300     $blens:ident, $is:ident; elem=[$($el:literal),+]; $lv0:literal; top=$top:literal) => {{
301        let mut _j = 0usize;
302        let mut _advanced = 0usize;
303        while _j < $dims[$lv0] {
304            $blens[$lv0] = $blocks[$lv0].max(1).min($dims[$lv0]).min($dims[$lv0] - _j);
305            elem_loops!($offsets, $strides, $f, $blens, $is; $($el),+);
306            _j += $blens[$lv0];
307            if _j < $dims[$lv0] {
308                for (o, s) in $offsets.iter_mut().zip($strides.iter()) {
309                    *o += ($blens[$lv0] as isize) * s[$lv0];
310                }
311                _advanced = _j;
312            }
313        }
314        for (o, s) in $offsets.iter_mut().zip($strides.iter()) {
315            *o -= (_advanced as isize) * s[$lv0];
316        }
317    }};
318    // Recursive case: outer block level wraps inner block levels.
319    ($dims:ident, $blocks:ident, $strides:ident, $offsets:ident, $f:ident,
320     $blens:ident, $is:ident; elem=[$($el:literal),+];
321     $lv:literal, $next:literal $(, $rest:literal)*; top=$top:literal) => {{
322        let mut _j = 0usize;
323        let mut _advanced = 0usize;
324        while _j < $dims[$lv] {
325            $blens[$lv] = $blocks[$lv].max(1).min($dims[$lv]).min($dims[$lv] - _j);
326            block_loop!($dims, $blocks, $strides, $offsets, $f, $blens, $is;
327                elem=[$($el),+]; $next $(, $rest)*; top=$top);
328            _j += $blens[$lv];
329            if _j < $dims[$lv] {
330                for (o, s) in $offsets.iter_mut().zip($strides.iter()) {
331                    *o += ($blens[$lv] as isize) * s[$lv];
332                }
333                _advanced = _j;
334            }
335        }
336        for (o, s) in $offsets.iter_mut().zip($strides.iter()) {
337            *o -= (_advanced as isize) * s[$lv];
338        }
339    }};
340}
341
342/// Generate a rank-specialized kernel function.
343///
344/// For rank 1 there are no element loops; only a single block loop on dim 0.
345/// For rank >= 2, block loops nest from outermost to innermost (dim 0), and
346/// element loops nest from dim N-1 down to dim 1 inside the innermost block.
347/// The offsets are back at their starting values on return.
348macro_rules! make_kernel {
349    // Rank 1: single block loop, no element nesting.
350    ($name:ident, rank=1) => {
351        #[inline]
352        fn $name<F>(
353            dims: &[usize],
354            blocks: &[usize],
355            strides: &[&[isize]],
356            offsets: &mut [isize],
357            f: &mut F,
358        ) -> Result<()>
359        where
360            F: FnMut(&[isize], usize, &[isize]) -> Result<()>,
361        {
362            let d0 = dims[0];
363            let b0 = blocks[0].max(1).min(d0);
364            let inner_strides: OperandVec<isize> = strides.iter().map(|s| s[0]).collect();
365
366            let mut j0 = 0usize;
367            let mut advanced = 0usize;
368            while j0 < d0 {
369                let blen0 = b0.min(d0 - j0);
370                f(offsets, blen0, &inner_strides)?;
371                j0 += blen0;
372                if j0 < d0 {
373                    for (o, s) in offsets.iter_mut().zip(strides.iter()) {
374                        *o += (blen0 as isize) * s[0];
375                    }
376                    advanced = j0;
377                }
378            }
379            for (o, s) in offsets.iter_mut().zip(strides.iter()) {
380                *o -= (advanced as isize) * s[0];
381            }
382            Ok(())
383        }
384    };
385    // Rank >= 2: nested block loops + element loops.
386    //   block=[outermost, ..., 0]  - block loop levels from outermost to innermost
387    //   elem=[N-1, ..., 1]         - element loop levels from outermost to innermost
388    //   top=N-1                    - topmost element level (= rank - 1)
389    ($name:ident, rank=$rank:literal,
390     block=[$($blk:literal),+], elem=[$($el:literal),+], top=$top:literal) => {
391        #[inline]
392        fn $name<F>(
393            dims: &[usize],
394            blocks: &[usize],
395            strides: &[&[isize]],
396            offsets: &mut [isize],
397            f: &mut F,
398        ) -> Result<()>
399        where
400            F: FnMut(&[isize], usize, &[isize]) -> Result<()>,
401        {
402            let inner_strides: OperandVec<isize> = strides.iter().map(|s| s[0]).collect();
403            let mut blens = [0usize; $rank];
404            block_loop!(dims, blocks, strides, offsets, f, blens, inner_strides;
405                elem=[$($el),+]; $($blk),+; top=$top);
406            Ok(())
407        }
408    };
409}
410
411make_kernel!(kernel_1d_inner, rank = 1);
412make_kernel!(
413    kernel_2d_inner,
414    rank = 2,
415    block = [1, 0],
416    elem = [1],
417    top = 1
418);
419make_kernel!(
420    kernel_3d_inner,
421    rank = 3,
422    block = [2, 1, 0],
423    elem = [2, 1],
424    top = 2
425);
426make_kernel!(
427    kernel_4d_inner,
428    rank = 4,
429    block = [3, 2, 1, 0],
430    elem = [3, 2, 1],
431    top = 3
432);
433make_kernel!(
434    kernel_5d_inner,
435    rank = 5,
436    block = [4, 3, 2, 1, 0],
437    elem = [4, 3, 2, 1],
438    top = 4
439);
440make_kernel!(
441    kernel_6d_inner,
442    rank = 6,
443    block = [5, 4, 3, 2, 1, 0],
444    elem = [5, 4, 3, 2, 1],
445    top = 5
446);
447make_kernel!(
448    kernel_7d_inner,
449    rank = 7,
450    block = [6, 5, 4, 3, 2, 1, 0],
451    elem = [6, 5, 4, 3, 2, 1],
452    top = 6
453);
454make_kernel!(
455    kernel_8d_inner,
456    rank = 8,
457    block = [7, 6, 5, 4, 3, 2, 1, 0],
458    elem = [7, 6, 5, 4, 3, 2, 1],
459    top = 7
460);
461
462/// N-dimensional kernel with inner block callback (iterative form).
463///
464/// Fallback for rank >= 9. Uses a carry-style increment over outer levels
465/// instead of compile-time unrolled nested loops.
466#[inline]
467fn kernel_nd_inner_iterative<F>(
468    dims: &[usize],
469    blocks: &[usize],
470    strides: &[&[isize]],
471    offsets: &mut [isize],
472    f: &mut F,
473) -> Result<()>
474where
475    F: FnMut(&[isize], usize, &[isize]) -> Result<()>,
476{
477    let rank = dims.len();
478    debug_assert!(rank >= 9);
479
480    let d0 = dims[0];
481    let b0 = blocks[0].max(1).min(d0);
482    let inner_strides: OperandVec<isize> = strides.iter().map(|s| s[0]).collect();
483
484    // Current position for each outer level (1..rank-1). Level 0 uses block loop.
485    let mut idx: Vec<usize> = vec![0usize; rank];
486
487    // As in the rank-specialized kernels, offsets advance only when another
488    // position follows and each level rewinds from its last position, so they
489    // never leave the reachable span.
490    if dims.contains(&0) {
491        return Ok(());
492    }
493    loop {
494        // Level 0: callback over contiguous block fragments.
495        let mut j0 = 0usize;
496        let mut advanced = 0usize;
497        while j0 < d0 {
498            let blen0 = b0.min(d0 - j0);
499            f(offsets, blen0, &inner_strides)?;
500            j0 += blen0;
501            if j0 < d0 {
502                for (offset, s) in offsets.iter_mut().zip(strides.iter()) {
503                    *offset += (blen0 as isize) * s[0];
504                }
505                advanced = j0;
506            }
507        }
508        for (offset, s) in offsets.iter_mut().zip(strides.iter()) {
509            *offset -= (advanced as isize) * s[0];
510        }
511
512        // Carry-style increment for outer levels.
513        let mut level = 1usize;
514        loop {
515            if idx[level] + 1 < dims[level] {
516                idx[level] += 1;
517                for (offset, s) in offsets.iter_mut().zip(strides.iter()) {
518                    *offset += s[level];
519                }
520                break;
521            }
522
523            let last = (dims[level] - 1) as isize;
524            idx[level] = 0;
525            for (offset, s) in offsets.iter_mut().zip(strides.iter()) {
526                *offset -= last * s[level];
527            }
528            level += 1;
529            if level == rank {
530                return Ok(());
531            }
532        }
533    }
534}
535
536// ============================================================================
537// Pre-ordered iteration (for threaded leaf functions)
538// ============================================================================
539
540/// Iterate over blocks with pre-ordered dimensions and initial offsets.
541///
542/// Unlike `for_each_inner_block`, this function assumes that `dims`, `blocks`,
543/// and `strides` are **already in iteration order** (i.e., identity ordering).
544/// It also accepts `initial_offsets` which are added to the starting offsets
545/// before iteration begins.
546///
547/// This avoids the redundant re-ordering and per-callback `Vec` allocation
548/// that `for_each_inner_block_with_offsets` previously incurred.
549///
550/// The block-walking machinery is compiled once (`preordered_dyn`): a generic
551/// caller only instantiates its own callback, reached through one indirect call
552/// per inner block (a run of elements, never per element). Keeping the walk
553/// generic would duplicate it for every operation, dtype and element-op
554/// combination of every downstream crate.
555#[inline]
556pub(crate) fn for_each_inner_block_preordered<S: AsRef<[isize]>, F>(
557    dims: &[usize],
558    blocks: &[usize],
559    strides: &[S],
560    initial_offsets: &[isize],
561    mut f: F,
562) -> Result<()>
563where
564    F: FnMut(&[isize], usize, &[isize]) -> Result<()>,
565{
566    let refs: OperandVec<&[isize]> = strides.iter().map(|s| s.as_ref()).collect();
567    preordered_dyn(dims, blocks, &refs, initial_offsets, &mut f)
568}
569
570type BlockFn<'a> = &'a mut dyn FnMut(&[isize], usize, &[isize]) -> Result<()>;
571
572#[inline(never)]
573fn preordered_dyn(
574    dims: &[usize],
575    blocks: &[usize],
576    strides: &[&[isize]],
577    initial_offsets: &[isize],
578    mut f: BlockFn<'_>,
579) -> Result<()> {
580    let rank = dims.len();
581    if rank == 0 {
582        return f(initial_offsets, 1, &[]);
583    }
584
585    // Start from initial_offsets (kernel functions reset to starting values at end)
586    let mut offsets: OperandVec<isize> = initial_offsets.iter().copied().collect();
587
588    match rank {
589        1 => kernel_1d_inner(dims, blocks, strides, &mut offsets, &mut f),
590        2 => kernel_2d_inner(dims, blocks, strides, &mut offsets, &mut f),
591        3 => kernel_3d_inner(dims, blocks, strides, &mut offsets, &mut f),
592        4 => kernel_4d_inner(dims, blocks, strides, &mut offsets, &mut f),
593        5 => kernel_5d_inner(dims, blocks, strides, &mut offsets, &mut f),
594        6 => kernel_6d_inner(dims, blocks, strides, &mut offsets, &mut f),
595        7 => kernel_7d_inner(dims, blocks, strides, &mut offsets, &mut f),
596        8 => kernel_8d_inner(dims, blocks, strides, &mut offsets, &mut f),
597        _ => kernel_nd_inner_iterative(dims, blocks, strides, &mut offsets, &mut f),
598    }
599}
600
601// ============================================================================
602// Utility functions
603// ============================================================================
604
605/// Check exact shape agreement before prevalidated replay.
606///
607/// # Examples
608///
609/// ```
610/// use strided_basic::execution::ensure_same_shape;
611/// ensure_same_shape(&[2, 3], &[2, 3]).unwrap();
612/// assert!(ensure_same_shape(&[2, 3], &[3, 2]).is_err());
613/// ```
614///
615/// # Errors
616/// Returns a rank or shape mismatch.
617pub fn ensure_same_shape(a: &[usize], b: &[usize]) -> Result<()> {
618    if a.len() != b.len() {
619        return Err(crate::StridedError::RankMismatch(a.len(), b.len()));
620    }
621    if a != b {
622        return Err(crate::StridedError::ShapeMismatch(a.to_vec(), b.to_vec()));
623    }
624    Ok(())
625}
626
627#[derive(Copy, Clone, Debug, Eq, PartialEq)]
628pub(crate) enum ContiguousLayout {
629    /// C-like layout: last axis varies fastest.
630    RowMajor,
631    /// Julia/Fortran-like layout: first axis varies fastest.
632    ColMajor,
633}
634
635/// Returns the contiguous memory layout kind for the given (dims, strides).
636///
637/// Notes:
638/// - Ignores axes with `dim <= 1` since they do not affect addressability.
639/// - Does not treat negative-stride views as contiguous for fast-path purposes.
640pub(crate) fn contiguous_layout(dims: &[usize], strides: &[isize]) -> Option<ContiguousLayout> {
641    if dims.len() != strides.len() {
642        return None;
643    }
644    if dims.is_empty() {
645        return Some(ContiguousLayout::RowMajor);
646    }
647
648    // Row-major: check from last to first.
649    let mut expected = 1isize;
650    let mut row_ok = true;
651    for (&dim, &stride) in dims.iter().rev().zip(strides.iter().rev()) {
652        if dim <= 1 {
653            continue;
654        }
655        if stride != expected {
656            row_ok = false;
657            break;
658        }
659        expected = expected.saturating_mul(dim as isize);
660    }
661    if row_ok {
662        return Some(ContiguousLayout::RowMajor);
663    }
664
665    // Col-major: check from first to last.
666    let mut expected = 1isize;
667    for (&dim, &stride) in dims.iter().zip(strides.iter()) {
668        if dim <= 1 {
669            continue;
670        }
671        if stride != expected {
672            return None;
673        }
674        expected = expected.saturating_mul(dim as isize);
675    }
676    Some(ContiguousLayout::ColMajor)
677}
678
679/// Checked element count of `dims`.
680///
681/// A shape containing a zero extent has zero elements even when the product
682/// of its other extents would overflow; such views pass
683/// `strided_view::validate_bounds`, so this must not form that product. A
684/// nonempty shape whose element count exceeds `usize::MAX` (for example a
685/// huge stride-0 broadcast) returns [`StridedError::OffsetOverflow`].
686pub(crate) fn total_len(dims: &[usize]) -> Result<usize> {
687    if dims.contains(&0) {
688        return Ok(0);
689    }
690    dims.iter()
691        .try_fold(1usize, |acc, &dim| acc.checked_mul(dim))
692        .ok_or(StridedError::OffsetOverflow)
693}
694
695/// Whether the sequential contiguous fast path should be used.
696///
697/// When the `parallel` feature is enabled and the total element count exceeds
698/// the threading threshold, we normally must *not* take the contiguous fast path
699/// so that the parallel kernel path can be reached. If the active Rayon pool has
700/// only one worker, however, there is no threaded path to reach and the
701/// sequential slice loop is the right kernel.
702#[inline]
703pub(crate) fn use_sequential_fast_path(total: usize) -> bool {
704    #[cfg(feature = "parallel")]
705    {
706        total <= crate::threading::MINTHREADLENGTH || crate::execution_policy::rayon_threads() <= 1
707    }
708    #[cfg(not(feature = "parallel"))]
709    {
710        let _ = total;
711        true
712    }
713}
714
715/// Returns the common contiguous layout if **all** provided stride arrays
716/// share the same contiguous layout for the given `dims`.
717///
718/// Returns `None` if `strides_list` is empty, any array is not contiguous,
719/// or any two arrays have different contiguous layouts.
720#[inline]
721pub(crate) fn same_contiguous_layout(
722    dims: &[usize],
723    strides_list: &[&[isize]],
724) -> Option<ContiguousLayout> {
725    let first = contiguous_layout(dims, strides_list.first()?)?;
726    for strides in &strides_list[1..] {
727        if contiguous_layout(dims, strides)? != first {
728            return None;
729        }
730    }
731    Some(first)
732}
733
734/// Returns the common contiguous layout only when the sequential fast path
735/// should be used (total elements <= threading threshold).
736///
737/// Returns [`StridedError::OffsetOverflow`] when the element count of `dims`
738/// does not fit in `usize`.
739#[inline]
740pub(crate) fn sequential_contiguous_layout(
741    dims: &[usize],
742    strides_list: &[&[isize]],
743) -> Result<Option<ContiguousLayout>> {
744    if !use_sequential_fast_path(total_len(dims)?) {
745        return Ok(None);
746    }
747    Ok(same_contiguous_layout(dims, strides_list))
748}
749
750#[cfg(test)]
751#[path = "kernel/tests/tests.rs"]
752mod tests;