strided-basic 0.6.0

Shared typed strided CPU primitives and copy/reduction execution.
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
//! Reduce operations on dynamic-rank strided views.

#[cfg(feature = "parallel")]
use crate::kernel::same_contiguous_layout;
use crate::kernel::{
    build_plan_fused, for_each_inner_block_preordered, sequential_contiguous_layout, total_len,
};
use crate::maybe_sync::{MaybeSendSync, MaybeSync};
use crate::simd;
use crate::view::{col_major_strides, StridedArray, StridedView, StridedViewMut};
use crate::{Result, StridedError};
use std::ops::Range;
use strided_view::ElementOp;

#[cfg(feature = "parallel")]
use crate::fuse::compute_costs;
use crate::inline::InlineVec;
#[cfg(feature = "parallel")]
use crate::threading::{
    for_each_inner_block_with_offsets, mapreduce_threaded_fold, SendPtr, MINTHREADLENGTH,
};

/// Number of independent accumulators of [`fold_contiguous`].
const FOLD_LANES: usize = 16;

/// Minimum outputs along a unit-stride kept axis for the column sweep of
/// [`reduce_axis`].
const AXIS_SWEEP_MIN_LEN: usize = 16;

/// Byte budget of one output block of the column sweep of [`reduce_axis`].
const AXIS_SWEEP_BLOCK_BYTES: usize = 64 * 1024;

/// Folds a contiguous run with [`FOLD_LANES`] independent accumulators.
///
/// Lane `i` starts from element `i` and folds elements `i + FOLD_LANES`,
/// `i + 2 * FOLD_LANES`, ... in order; the remainder is folded into the first
/// lanes. `init` is combined exactly once, before lane 0, and the lanes are
/// then combined left to right. Independent lanes let the compiler keep
/// several vector accumulators in flight instead of one serial dependency
/// chain, which is what makes a closure-based sum or max approach the speed
/// of a hand written lane loop.
#[inline(always)]
fn fold_contiguous<T, Op, M, R, U>(src: &[T], init: U, map_fn: &M, reduce_fn: &R) -> U
where
    T: Copy,
    Op: ElementOp<T>,
    M: Fn(T) -> U,
    R: Fn(U, U) -> U,
    U: Clone,
{
    if src.len() < 2 * FOLD_LANES {
        let mut acc = init;
        for &value in src {
            acc = reduce_fn(acc, map_fn(Op::apply(value)));
        }
        return acc;
    }
    let (head, rest) = src.split_at(FOLD_LANES);
    let mut lanes: [U; FOLD_LANES] = core::array::from_fn(|lane| map_fn(Op::apply(head[lane])));
    let mut chunks = rest.chunks_exact(FOLD_LANES);
    for chunk in chunks.by_ref() {
        for (lane, &value) in lanes.iter_mut().zip(chunk) {
            *lane = reduce_fn(lane.clone(), map_fn(Op::apply(value)));
        }
    }
    for (lane, &value) in lanes.iter_mut().zip(chunks.remainder()) {
        *lane = reduce_fn(lane.clone(), map_fn(Op::apply(value)));
    }
    lanes.into_iter().fold(init, reduce_fn)
}

/// Folds `len` elements starting at `ptr` with step `stride`.
///
/// # Safety
/// `ptr.offset(i * stride)` must be a readable element for every `i < len`.
#[inline(always)]
unsafe fn fold_run<T, Op, M, R, U>(
    ptr: *const T,
    stride: isize,
    len: usize,
    init: U,
    map_fn: &M,
    reduce_fn: &R,
) -> U
where
    T: Copy,
    Op: ElementOp<T>,
    M: Fn(T) -> U,
    R: Fn(U, U) -> U,
    U: Clone,
{
    if stride == 1 {
        // SAFETY: the caller proves `len` readable unit-stride elements.
        let run = unsafe { std::slice::from_raw_parts(ptr, len) };
        return fold_contiguous::<T, Op, M, R, U>(run, init, map_fn, reduce_fn);
    }
    let mut acc = init;
    let mut cursor = ptr;
    for index in 0..len {
        // SAFETY: the caller proves every element of the run is readable.
        acc = reduce_fn(acc, map_fn(Op::apply(unsafe { *cursor })));
        if index + 1 < len {
            cursor = cursor.wrapping_offset(stride);
        }
    }
    acc
}

/// Full reduction with map function: `reduce(init, op, map.(src))`.
///
/// # Evaluation order
///
/// `reduce_fn` must be associative and commutative, and `init` must be an
/// identity of `reduce_fn` whenever the reduction runs on several threads.
/// A contiguous source is folded with sixteen independent lanes (see the
/// crate source for the exact lane assignment), non-contiguous sources are
/// traversed in a cache-friendly loop order, and parallel execution folds
/// chunks independently and combines them in chunk order. Floating point
/// results can therefore differ from a strict left to right fold in the last
/// bits; for a fixed shape, layout and thread count the result is repeatable.
///
/// Because `map_fn` and `reduce_fn` are opaque closures this entry cannot
/// select the dtype-specific SIMD kernels of [`crate::ErasedReducePlan`];
/// callers that reduce with a known operation (sum, product, max, min) on a
/// supported dtype can use that plan for the fastest path.
pub fn reduce<T: Copy + MaybeSendSync, Op: ElementOp<T>, M, R, U>(
    src: &StridedView<T, Op>,
    map_fn: M,
    reduce_fn: R,
    init: U,
) -> Result<U>
where
    M: Fn(T) -> U + MaybeSync,
    R: Fn(U, U) -> U + MaybeSync,
    U: Clone + MaybeSendSync,
{
    reduce_impl(src, map_fn, reduce_fn, init)
}

fn reduce_impl<T: Copy + MaybeSendSync, Op: ElementOp<T>, M, R, U>(
    src: &StridedView<T, Op>,
    map_fn: M,
    reduce_fn: R,
    init: U,
) -> Result<U>
where
    M: Fn(T) -> U + MaybeSync,
    R: Fn(U, U) -> U + MaybeSync,
    U: Clone + MaybeSendSync,
{
    let src_ptr = src.ptr();
    let src_dims = src.dims();
    let src_strides = src.strides();

    let contiguous = sequential_contiguous_layout(src_dims, &[src_strides])?;
    if contiguous.is_some() {
        let len = total_len(src_dims)?;
        let src = unsafe { std::slice::from_raw_parts(src_ptr, len) };
        return Ok(simd::dispatch_if_large(len, || {
            fold_contiguous::<T, Op, M, R, U>(src, init, &map_fn, &reduce_fn)
        }));
    }

    // Parallel contiguous fast path: split into scheduler chunks with slice-based iteration.
    // This enables LLVM auto-vectorization on each chunk, unlike the general threaded path
    // which uses scalar pointer-offset loops.
    #[cfg(feature = "parallel")]
    {
        let total = total_len(src_dims)?;
        let nthreads = crate::execution_policy::rayon_threads();
        if total > MINTHREADLENGTH
            && nthreads > 1
            && same_contiguous_layout(src_dims, &[src_strides]).is_some()
        {
            let src_slice = unsafe { std::slice::from_raw_parts(src_ptr, total) };
            let result = crate::threading::parallel_map_reduce(
                0..total,
                nthreads,
                &|range| {
                    simd::dispatch_if_large(range.len(), || {
                        fold_contiguous::<T, Op, M, R, U>(
                            &src_slice[range],
                            init.clone(),
                            &map_fn,
                            &reduce_fn,
                        )
                    })
                },
                &|a, b| reduce_fn(a, b),
            );
            return Ok(result);
        }
    }

    let strides_list: [&[isize]; 1] = [src_strides];

    let (fused_dims, ordered_strides, plan) =
        build_plan_fused(src_dims, &strides_list, None, std::mem::size_of::<T>());

    #[cfg(feature = "parallel")]
    {
        let total = total_len(&fused_dims)?;
        let nthreads = crate::execution_policy::rayon_threads();
        if total > MINTHREADLENGTH && nthreads > 1 {
            let src_send = SendPtr(src_ptr as *mut T);
            let costs = compute_costs(&ordered_strides);
            let initial_offsets: crate::inline::OperandVec<isize> =
                ordered_strides.iter().map(|_| 0isize).collect();
            let map_fn_ref = &map_fn;
            let reduce_fn_ref = &reduce_fn;
            let init_ref = &init;
            // The reduction tree combines each split's partial result, so no
            // worker-count-sized output-slot vector is needed (#306 scope).
            let result = mapreduce_threaded_fold(
                &fused_dims,
                &plan.block,
                &ordered_strides,
                &initial_offsets,
                &costs,
                nthreads,
                &|dims, blocks, strides_list, offsets| {
                    let mut acc = init_ref.clone();
                    for_each_inner_block_with_offsets(
                        dims,
                        blocks,
                        strides_list,
                        offsets,
                        |offsets, len, strides| {
                            let mut ptr = unsafe { src_send.as_const().offset(offsets[0]) };
                            let stride = strides[0];
                            for _ in 0..len {
                                let val = Op::apply(unsafe { *ptr });
                                let mapped = map_fn_ref(val);
                                acc = reduce_fn_ref(acc.clone(), mapped);
                                unsafe {
                                    ptr = ptr.offset(stride);
                                }
                            }
                            Ok(())
                        },
                    )?;
                    Ok(acc)
                },
                &|a, b| reduce_fn_ref(a, b),
            )?;
            return Ok(result);
        }
    }

    let mut acc = init;
    let initial_offsets: crate::inline::OperandVec<isize> =
        ordered_strides.iter().map(|_| 0isize).collect();
    for_each_inner_block_preordered(
        &fused_dims,
        &plan.block,
        &ordered_strides,
        &initial_offsets,
        |offsets, len, strides| {
            let mut ptr = unsafe { src_ptr.offset(offsets[0]) };
            let stride = strides[0];
            for _ in 0..len {
                let val = Op::apply(unsafe { *ptr });
                let mapped = map_fn(val);
                acc = reduce_fn(acc.clone(), mapped);
                unsafe {
                    ptr = ptr.offset(stride);
                }
            }
            Ok(())
        },
    )?;

    Ok(acc)
}

/// Reduce along a single axis, returning a new `StridedArray`.
///
/// This allocates a column-major output and delegates to
/// [`reduce_axis_into`].
///
/// Every output folds `init` once and then its reduced elements. When the
/// reduced axis has unit stride those elements are folded with the lane
/// scheme of [`reduce`], so `reduce_fn` must be associative and commutative.
/// When the leading kept axis has unit stride, a block of adjacent outputs
/// is accumulated one contiguous source column at a time, which keeps the
/// left to right order per output. With the `parallel` feature, outputs are
/// split across threads once the number of source elements read exceeds the
/// threading threshold; each output is computed by one thread, so the result
/// does not depend on the thread count.
pub fn reduce_axis<T: Copy + MaybeSendSync, Op: ElementOp<T>, M, R, U>(
    src: &StridedView<T, Op>,
    axis: usize,
    map_fn: M,
    reduce_fn: R,
    init: U,
) -> Result<StridedArray<U>>
where
    M: Fn(T) -> U + MaybeSync,
    R: Fn(U, U) -> U + MaybeSync,
    U: Clone + MaybeSendSync,
{
    let rank = src.ndim();
    if axis >= rank {
        return Err(StridedError::InvalidAxis { axis, rank });
    }
    // Reject an element count beyond `usize` before allocating the output.
    total_len(src.dims())?;
    let out_dims: InlineVec<usize> = src
        .dims()
        .iter()
        .enumerate()
        .filter(|(i, _)| *i != axis)
        .map(|(_, &d)| d)
        .collect();
    // A full reduction keeps the historical rank-1 shape-[1] output; the
    // `_into` form accepts any single-element destination for that case.
    let (out_dims, out_strides) = if out_dims.is_empty() {
        let mut dims = InlineVec::new();
        dims.push(1usize);
        let strides = col_major_strides(&dims);
        (dims, strides)
    } else {
        let strides = col_major_strides(&out_dims);
        (out_dims, strides)
    };
    let total_out = total_len(&out_dims)?;
    let mut out =
        StridedArray::<U>::from_parts(vec![init.clone(); total_out], &out_dims, &out_strides, 0)?;
    {
        let mut dest = out.view_mut();
        // The fresh column-major destination needs no validation.
        reduce_axis_into_impl(&mut dest, src, axis, map_fn, reduce_fn, init)?;
    }
    Ok(out)
}

/// Reduce along a single axis into caller-owned storage.
///
/// `dest` must have the source's shape with `axis` removed, or any
/// single-element destination for a full reduction, and an injective layout.
/// Every destination element is written: an empty reduced axis fills each
/// output with `init`.
///
/// # Errors
///
/// Returns [`StridedError::InvalidAxis`] for `axis >= src.ndim()`,
/// [`StridedError::RankMismatch`] / [`StridedError::ShapeMismatch`] for a
/// destination shape that is not the reduced source shape,
/// [`StridedError::NonInjectiveOutputLayout`] for an overlapping destination,
/// and [`StridedError::OffsetOverflow`] for unrepresentable sizes.
///
/// # Examples
///
/// ```
/// use strided_basic::{reduce_axis_into, StridedArray};
/// # let src = StridedArray::<f64>::from_fn_col_major(&[2, 3], |i| (i[0] + 2 * i[1]) as f64);
/// let mut out = StridedArray::<f64>::col_major(&[3]);
/// reduce_axis_into(&mut out.view_mut(), &src.view(), 0, |x| x, |a, b| a + b, 0.0).unwrap();
/// assert_eq!(out.get(&[0]), 1.0);
/// ```
///
/// Validation runs here; [`reduce_axis_into_impl`] is the shared kernel. The
/// array-returning wrapper owns a fresh column-major destination and calls the
/// kernel directly, so the common path does not re-validate.
pub fn reduce_axis_into<T: Copy + MaybeSendSync, Op: ElementOp<T>, M, R, U>(
    dest: &mut StridedViewMut<U>,
    src: &StridedView<T, Op>,
    axis: usize,
    map_fn: M,
    reduce_fn: R,
    init: U,
) -> Result<()>
where
    M: Fn(T) -> U + MaybeSync,
    R: Fn(U, U) -> U + MaybeSync,
    U: Clone + MaybeSendSync,
{
    let rank = src.ndim();
    if axis >= rank {
        return Err(StridedError::InvalidAxis { axis, rank });
    }
    let src_dims = src.dims();
    // Reject an element count beyond usize (a huge stride-0 broadcast) up
    // front instead of looping over it.
    total_len(src_dims)?;
    let kept_len = rank - 1;
    let dest_len = total_len(dest.dims())?;
    if kept_len == 0 {
        if dest_len != 1 {
            return Err(StridedError::ShapeMismatch(vec![dest_len], vec![1]));
        }
    } else {
        if dest.ndim() != kept_len {
            return Err(StridedError::RankMismatch(dest.ndim(), kept_len));
        }
        let dest_dims = dest.dims();
        let mut d = 0usize;
        for (i, &src_dim) in src_dims.iter().enumerate() {
            if i == axis {
                continue;
            }
            if src_dim != dest_dims[d] {
                let expected: Vec<usize> = src_dims
                    .iter()
                    .enumerate()
                    .filter(|(i, _)| *i != axis)
                    .map(|(_, &d)| d)
                    .collect();
                return Err(StridedError::ShapeMismatch(dest_dims.to_vec(), expected));
            }
            d += 1;
        }
        if !crate::layout_check::is_injective_layout(dest.dims(), dest.strides()) {
            return Err(StridedError::NonInjectiveOutputLayout);
        }
    }
    reduce_axis_into_impl(dest, src, axis, map_fn, reduce_fn, init)
}

/// Shared kernel of [`reduce_axis_into`] and [`reduce_axis`].
///
/// Assumes the destination has the reduced source shape (or one element for a
/// full reduction) and an injective layout; the public entry validates that,
/// while the array-returning wrapper owns a freshly allocated column-major
/// destination and can skip the re-validation.
fn reduce_axis_into_impl<T: Copy + MaybeSendSync, Op: ElementOp<T>, M, R, U>(
    dest: &mut StridedViewMut<U>,
    src: &StridedView<T, Op>,
    axis: usize,
    map_fn: M,
    reduce_fn: R,
    init: U,
) -> Result<()>
where
    M: Fn(T) -> U + MaybeSync,
    R: Fn(U, U) -> U + MaybeSync,
    U: Clone + MaybeSendSync,
{
    let rank = src.ndim();
    let src_dims = src.dims();
    let src_strides = src.strides();
    let src_ptr = src.ptr();
    let dest_len = total_len(dest.dims())?;
    let axis_len = src_dims[axis];
    let axis_stride = src_strides[axis];

    if axis_len == 0 {
        // No source elements to fold: every output is `init`. Fill the
        // destination directly rather than traversing the source, because an
        // empty source view carries a dangling base pointer and `offset` on it
        // is UB even when the fold reads zero elements.
        fill_outputs_with(dest, &init);
        return Ok(());
    }

    if rank == 1 {
        // Full reduction into a single destination element.
        // SAFETY: the view proves `axis_len` readable elements along `axis`.
        let acc = unsafe {
            fold_run::<T, Op, M, R, U>(src_ptr, axis_stride, axis_len, init, &map_fn, &reduce_fn)
        };
        unsafe { *dest.as_mut_ptr() = acc };
        return Ok(());
    }
    if dest_len == 0 {
        return Ok(());
    }

    // Build the kept-axis metadata in one pass. A column-major destination is
    // the common case (the array-returning wrapper always allocates one); its
    // offset equals the linear output index, so its stride list is only
    // materialized for a strided destination.
    let dest_stride_view = dest.strides();
    let mut kept: InlineVec<(usize, isize)> = InlineVec::new();
    let mut dest_linear = true;
    let mut prefix = 1usize;
    for (i, (&dim, &src_stride)) in src_dims.iter().zip(src_strides).enumerate() {
        if i == axis {
            continue;
        }
        let d = kept.len();
        kept.push((dim, src_stride));
        if dest_linear && dest_stride_view[d] != prefix as isize {
            dest_linear = false;
        }
        prefix = prefix.saturating_mul(dim);
    }
    let dest_strides: InlineVec<isize> = if dest_linear {
        InlineVec::new()
    } else {
        dest_stride_view.iter().copied().collect()
    };
    let out_ptr = dest.as_mut_ptr();
    let parts = AxisParts {
        kept: &kept,
        dest_strides: &dest_strides,
        axis_len,
        axis_stride,
        map_fn: &map_fn,
        reduce_fn: &reduce_fn,
        init: &init,
    };

    #[cfg(feature = "parallel")]
    {
        let work = dest_len.saturating_mul(axis_len);
        let nthreads = crate::threading::parallel_threads_for_len(work).min(dest_len);
        if nthreads > 1 {
            let src_send = SendPtr(src_ptr as *mut T);
            let out_send = SendPtr(out_ptr);
            let parts = &parts;
            if dest_linear {
                crate::threading::parallel_for_each(0..dest_len, nthreads, &|range| {
                    // SAFETY: worker ranges are disjoint, the destination is
                    // injective, and the view proves every source offset
                    // reachable from its kept coordinates.
                    unsafe {
                        reduce_axis_range_linear::<T, Op, M, R, U>(
                            src_send.as_const(),
                            out_send.as_ptr(),
                            range,
                            parts,
                        )
                    }
                });
            } else {
                crate::threading::parallel_for_each(0..dest_len, nthreads, &|range| {
                    // SAFETY: as for the linear path.
                    unsafe {
                        reduce_axis_range_strided::<T, Op, M, R, U>(
                            src_send.as_const(),
                            out_send.as_ptr(),
                            range,
                            parts,
                        )
                    }
                });
            }
            return Ok(());
        }
    }

    // SAFETY: as for the parallel ranges, with one range covering all outputs.
    if dest_linear {
        unsafe {
            reduce_axis_range_linear::<T, Op, M, R, U>(src_ptr, out_ptr, 0..dest_len, &parts)
        };
    } else {
        unsafe {
            reduce_axis_range_strided::<T, Op, M, R, U>(src_ptr, out_ptr, 0..dest_len, &parts)
        };
    }
    Ok(())
}

/// Writes `value` into every logical destination element.
///
/// Used when the reduced axis is empty: the fold would read zero source
/// elements, so the destination is exactly `init` everywhere and no source
/// pointer may be formed. The odometer keeps the fill allocation-free within
/// [`strided_view::INLINE_RANK`] and spills above it.
fn fill_outputs_with<U: Clone>(dest: &mut StridedViewMut<U>, value: &U) {
    let dims = dest.dims();
    if dims.iter().any(|&dim| dim == 0) {
        return;
    }
    let strides = dest.strides();
    let rank = dims.len();
    let total: usize = dims.iter().product();
    let base = dest.as_mut_ptr();
    if rank == 0 {
        // SAFETY: a rank-0 destination is one writable element.
        unsafe { *base = value.clone() };
        return;
    }
    let mut coords_inline = [0usize; strided_view::INLINE_RANK];
    let mut coords_heap: Vec<usize> = Vec::new();
    let coords: &mut [usize] = if rank <= strided_view::INLINE_RANK {
        &mut coords_inline[..rank]
    } else {
        coords_heap.resize(rank, 0);
        &mut coords_heap[..]
    };
    let mut offset = 0isize;
    for _ in 0..total {
        // SAFETY: the destination is an injective view validated by the
        // caller, and the cursor visits each logical element exactly once, so
        // every written offset is in bounds.
        unsafe { *base.offset(offset) = value.clone() };
        for axis in 0..rank {
            coords[axis] += 1;
            offset += strides[axis];
            if coords[axis] < dims[axis] {
                break;
            }
            coords[axis] = 0;
            offset -= dims[axis] as isize * strides[axis];
        }
    }
}

/// Shared inputs of [`reduce_axis_range_linear`] and
/// [`reduce_axis_range_strided`].
struct AxisParts<'a, M, R, U> {
    /// `(extent, source stride)` of every kept axis, in output order.
    kept: &'a [(usize, isize)],
    /// Destination stride of every kept axis, in the same order.
    dest_strides: &'a [isize],
    axis_len: usize,
    axis_stride: isize,
    map_fn: &'a M,
    reduce_fn: &'a R,
    init: &'a U,
}

/// Column-major destination cursor: the destination offset is the linear output
/// index, so no destination stride bookkeeping is needed.
///
/// Kept out of line so the common path stays the same single-cursor loop as the
/// pre-split kernel.
///
/// # Safety
/// `out` must be a column-major destination of `parts.kept` extents, `range` a
/// subrange of its logical output indices that no other thread writes, and
/// every source offset reachable from a kept coordinate plus
/// `k * parts.axis_stride` for `k < parts.axis_len` must be a readable element
/// of `src`.
#[inline(never)]
unsafe fn reduce_axis_range_linear<T, Op, M, R, U>(
    src: *const T,
    out: *mut U,
    range: Range<usize>,
    parts: &AxisParts<'_, M, R, U>,
) where
    T: Copy,
    Op: ElementOp<T>,
    M: Fn(T) -> U,
    R: Fn(U, U) -> U,
    U: Clone,
{
    let kept = parts.kept;
    let (lead_extent, lead_stride) = kept[0];
    let sweep = lead_stride == 1 && lead_extent >= AXIS_SWEEP_MIN_LEN && parts.axis_len > 1;
    let block = (AXIS_SWEEP_BLOCK_BYTES / std::mem::size_of::<U>().max(1)).max(AXIS_SWEEP_MIN_LEN);

    // Decode the first output once, then advance incrementally. A fixed stack
    // array with a spill covers every inline rank without the per-access enum
    // match a `TinyVec` cursor would add to the hot loop, and allocates only
    // above `INLINE_RANK` (#306 P1).
    let coords_len = kept.len();
    let mut coords_inline = [0usize; strided_view::INLINE_RANK];
    let mut coords_heap: Vec<usize> = Vec::new();
    let coords: &mut [usize] = if coords_len <= strided_view::INLINE_RANK {
        &mut coords_inline[..coords_len]
    } else {
        coords_heap.resize(coords_len, 0);
        &mut coords_heap[..]
    };
    let mut rest = range.start;
    let mut src_off = 0isize;
    for (coord, &(extent, stride)) in coords.iter_mut().zip(kept) {
        *coord = rest % extent;
        rest /= extent;
        src_off += *coord as isize * stride;
    }

    let mut output = range.start;
    while output < range.end {
        let len = if sweep {
            (lead_extent - coords[0]).min(range.end - output).min(block)
        } else {
            1
        };
        if sweep {
            // SAFETY: the `len` outputs are inside the caller's range, their
            // unit-stride destination elements are contiguous, and their
            // unit-stride source columns are readable.
            unsafe { sweep_block::<T, Op, M, R, U>(src, src_off, out.add(output), len, parts) };
        } else {
            // SAFETY: one reachable output and its reduced run.
            unsafe {
                let acc = fold_run::<T, Op, M, R, U>(
                    src.offset(src_off),
                    parts.axis_stride,
                    parts.axis_len,
                    parts.init.clone(),
                    parts.map_fn,
                    parts.reduce_fn,
                );
                *out.add(output) = acc;
            }
        }
        output += len;
        if output >= range.end {
            break;
        }
        // Advance the kept coordinates by `len` along the leading axis; a
        // block never crosses the end of the leading axis.
        coords[0] += len;
        src_off += len as isize * lead_stride;
        let mut axis = 0;
        while coords[axis] == kept[axis].0 {
            src_off -= kept[axis].0 as isize * kept[axis].1;
            coords[axis] = 0;
            axis += 1;
            coords[axis] += 1;
            src_off += kept[axis].1;
        }
    }
}

/// General strided destination cursor, tracking a destination offset next to
/// the source offset.
///
/// # Safety
/// `out` must be a destination of `parts.kept` extents with injective layout,
/// `range` a subrange of its logical output indices that no other thread
/// writes, and every source offset reachable from a kept coordinate plus
/// `k * parts.axis_stride` for `k < parts.axis_len` must be a readable element
/// of `src`.
#[inline(never)]
unsafe fn reduce_axis_range_strided<T, Op, M, R, U>(
    src: *const T,
    out: *mut U,
    range: Range<usize>,
    parts: &AxisParts<'_, M, R, U>,
) where
    T: Copy,
    Op: ElementOp<T>,
    M: Fn(T) -> U,
    R: Fn(U, U) -> U,
    U: Clone,
{
    let kept = parts.kept;
    let dest_strides = parts.dest_strides;
    let (lead_extent, lead_stride) = kept[0];
    let lead_dest_stride = dest_strides[0];
    let sweep = lead_stride == 1
        && lead_dest_stride == 1
        && lead_extent >= AXIS_SWEEP_MIN_LEN
        && parts.axis_len > 1;
    let block = (AXIS_SWEEP_BLOCK_BYTES / std::mem::size_of::<U>().max(1)).max(AXIS_SWEEP_MIN_LEN);

    // Same fixed-array-plus-spill cursor as `reduce_axis_range_linear`: the
    // `TinyVec` enum match is not paid per output access in the hot loop and
    // only ranks above `INLINE_RANK` allocate (#306 P1).
    let coords_len = kept.len();
    let mut coords_inline = [0usize; strided_view::INLINE_RANK];
    let mut coords_heap: Vec<usize> = Vec::new();
    let coords: &mut [usize] = if coords_len <= strided_view::INLINE_RANK {
        &mut coords_inline[..coords_len]
    } else {
        coords_heap.resize(coords_len, 0);
        &mut coords_heap[..]
    };
    let mut rest = range.start;
    let mut src_off = 0isize;
    let mut dest_off = 0isize;
    for (coord, (&(extent, stride), &dest_stride)) in
        coords.iter_mut().zip(kept.iter().zip(dest_strides))
    {
        *coord = rest % extent;
        rest /= extent;
        src_off += *coord as isize * stride;
        dest_off += *coord as isize * dest_stride;
    }

    let mut output = range.start;
    while output < range.end {
        let len = if sweep {
            (lead_extent - coords[0]).min(range.end - output).min(block)
        } else {
            1
        };
        if sweep {
            // SAFETY: as for the linear cursor; `lead_dest_stride == 1` makes
            // the `len` destinations contiguous.
            unsafe {
                sweep_block::<T, Op, M, R, U>(src, src_off, out.offset(dest_off), len, parts)
            };
        } else {
            // SAFETY: one reachable output and its reduced run.
            unsafe {
                let acc = fold_run::<T, Op, M, R, U>(
                    src.offset(src_off),
                    parts.axis_stride,
                    parts.axis_len,
                    parts.init.clone(),
                    parts.map_fn,
                    parts.reduce_fn,
                );
                *out.offset(dest_off) = acc;
            }
        }
        output += len;
        if output >= range.end {
            break;
        }
        coords[0] += len;
        src_off += len as isize * lead_stride;
        dest_off += len as isize * lead_dest_stride;
        let mut axis = 0;
        while coords[axis] == kept[axis].0 {
            src_off -= kept[axis].0 as isize * kept[axis].1;
            dest_off -= kept[axis].0 as isize * dest_strides[axis];
            coords[axis] = 0;
            axis += 1;
            coords[axis] += 1;
            src_off += kept[axis].1;
            dest_off += dest_strides[axis];
        }
    }
}

/// Accumulates `len` adjacent outputs, one contiguous source column per
/// reduced index, keeping the left to right order per output.
///
/// # Safety
/// `out` must be `len` writable initialized outputs and `src + src_off +
/// k * parts.axis_stride` must start `len` readable elements for every
/// `k < parts.axis_len`.
#[inline(always)]
unsafe fn sweep_block<T, Op, M, R, U>(
    src: *const T,
    src_off: isize,
    out: *mut U,
    len: usize,
    parts: &AxisParts<'_, M, R, U>,
) where
    T: Copy,
    Op: ElementOp<T>,
    M: Fn(T) -> U,
    R: Fn(U, U) -> U,
    U: Clone,
{
    let map_fn = parts.map_fn;
    let reduce_fn = parts.reduce_fn;
    // SAFETY: forwarded caller contract.
    let out = unsafe { std::slice::from_raw_parts_mut(out, len) };
    let column = |k: usize| {
        // SAFETY: forwarded caller contract for reduced index `k`.
        unsafe {
            std::slice::from_raw_parts(src.offset(src_off + k as isize * parts.axis_stride), len)
        }
    };
    let first = column(0);
    for (slot, &value) in out.iter_mut().zip(first) {
        *slot = reduce_fn(parts.init.clone(), map_fn(Op::apply(value)));
    }
    let mut k = 1;
    while k + 4 <= parts.axis_len {
        let (c0, c1, c2, c3) = (column(k), column(k + 1), column(k + 2), column(k + 3));
        for (index, slot) in out.iter_mut().enumerate() {
            let mut acc = reduce_fn(slot.clone(), map_fn(Op::apply(c0[index])));
            acc = reduce_fn(acc, map_fn(Op::apply(c1[index])));
            acc = reduce_fn(acc, map_fn(Op::apply(c2[index])));
            acc = reduce_fn(acc, map_fn(Op::apply(c3[index])));
            *slot = acc;
        }
        k += 4;
    }
    while k < parts.axis_len {
        let values = column(k);
        for (slot, &value) in out.iter_mut().zip(values) {
            *slot = reduce_fn(slot.clone(), map_fn(Op::apply(value)));
        }
        k += 1;
    }
}