strided-kernel
Cache-optimized compute kernels over strided-view tensors.
Scope
- Unary/Binary/N-ary map kernels (
map_into,zip_map*_into) - Reductions (
reduce,reduce_axis) - Utility ops (
copy_into,add,dot,sum,symmetrize_into) - Optional parallel execution via
parallelfeature (Rayon)
Quick Example
use ;
let src = from_fn_row_major;
let mut dst = row_major;
map_into.unwrap;
assert_eq!;
Map and Reduce Operations
use ;
let a = from_fn_row_major;
let b = from_fn_row_major;
let mut out = row_major;
// Unary map: dest[i] = f(src[i])
map_into.unwrap;
// Binary zip map: dest[i] = f(a[i], b[i])
zip_map2_into.unwrap;
// Full reduction
let total = reduce.unwrap;
Built-in Erased Reduction Order
The built-in erased Sum, Product, and SumSquares reductions have a
narrower, explicit determinism contract than the generic closure-based
reduce API:
- Compact full reductions traverse consecutive physical storage from the logical origin. Other full reductions use the stable fused/block traversal. Axis reductions enumerate reduced coordinates in declared axis order, with the first reduced axis varying fastest.
- Serial compact
f32/f64full reductions use four SIMD accumulators when thesimdfeature is enabled. Other dtype/feature combinations use a fixed eight-lane scalar accumulator. Accumulator lanes are merged in stable order. Generic closure reductions retain their existing association. SumSquaresacceptsf32andf64. Each value is multiplied by itself and rounded in that dtype before the result enters theSumaccumulator. The multiply and add are not contracted into FMA, so overflow and underflow are classified before accumulation. The kernel does not materialize a squared tensor.i32andi64use wrapping sum/product. Floating and complex reductions use same-dtype arithmetic without widening, compensation, conjugation, or fast-math. Reassociation can change roundoff, signed zero, NaN details, and the point at which an intermediate overflows or underflows. Consequently, floating product may differ in finite/zero/infinite/NaN classification from a strict left fold.ExecContext::serial()andExecContext::max_threads(1)use the same serial algorithm. A fixed nonzero thread budget has deterministic partition and merge order for a fixed executable, target, input, and layout; worker completion order does not affect the result. Different budgets may produce different floating or complex roundoff.- Empty sum, product, and sum-of-squares return zero, one, and zero respectively. No tolerance or correctness gate is relaxed for the optimized association.
High-Level Operations
use ;
let a = from_fn_row_major;
let mut out = row_major;
// Copy
copy_into.unwrap;
// Element-wise add: dest[i] += src[i]
add.unwrap;
// Dot product
let d = dot.unwrap;
// Symmetrize: dest = (src + src^T) / 2
symmetrize_into.unwrap;
Cache Optimization
The library automatically optimizes iteration order for cache efficiency:
- Dimension Fusion: Contiguous dimensions are fused to reduce loop overhead
- Dimension Reordering: Dimensions are sorted by stride magnitude for optimal memory access
- Tiled Iteration: Operations are blocked to fit in L1 cache (32KB)
- Contiguous Fast Paths: Contiguous arrays bypass blocking for direct iteration
Parallel Feature
[]
= { = "0.4", = ["parallel"] }
The default ExecutionPolicy::AmbientRayon behavior uses the current installed
Rayon pool (or the global pool). Explicit runtimes can bound strided-owned
fanout without creating another pool:
use NonZeroUsize;
use ;
let source = from_fn_col_major;
let mut destination = col_major;
let max_threads = new.unwrap;
with_execution_policy;
assert_eq!;
ExecutionPolicy controls fanout, not CPU placement or pool construction.
Nested strided operations inside a bounded worker partition run sequentially.
The opaque parallel permutation-copy path is used under an explicit Rayon
policy only when the installed pool size fits the requested budget; otherwise
that copy falls back to its serial implementation because the external copy
cannot be capped per call.
The policy is worker-local; callbacks that depend on isolation from unrelated
work must not invoke their own Rayon scheduling or yielding APIs.
Benchmarks
Run all benchmarks (single-threaded + multi-threaded, Rust + Julia):
Or individually:
# Single-threaded Rust
# Single-threaded Julia
JULIA_NUM_THREADS=1
# Multi-threaded Rust (N threads)
RAYON_NUM_THREADS=N
# Multi-threaded comparison script
# Scaling benchmarks (sum + permute, 1/2/4 threads)
# Rank-25 tensor permutation (quantum circuit simulation workload)
RAYON_NUM_THREADS=1
# Rank-25 Julia comparison
JULIA_NUM_THREADS=1
# Mul kernel comparison against PyTorch CPU at 1T and 4T.
# Requires uv; the script uses uv run --with torch --with numpy so PyTorch
# is installed only in the benchmark environment.
# The PyTorch runner uses torch.mul(..., out=...) to avoid allocator/autograd overhead.
# Noncompact batched outer product includes both compact output and a
# torchlike_output case whose non-contiguous output strides match torch.einsum.
STRIDED_KERNEL_MUL_BENCH_PROFILE=full
Published benchmark programs and current measured results live in
strided-rs-benchmark-suite.
Algorithm Comparison: Julia Strided.jl vs Rust strided-rs
Both implementations share the same core algorithm ported from Strided.jl:
- Dimension fusion — merge contiguous dimensions to reduce loop depth
- Importance-weighted ordering — bit-pack stride orders with output array weighted 2× to determine optimal iteration order
- L1 cache blocking — iteratively halve block sizes until the working set fits in 32 KB
- Reversed loop nesting — innermost loop operates on the highest-importance dimension (smallest stride) for optimal cache access
The key architectural differences are:
| Feature | Julia | Rust |
|---|---|---|
| Kernel generation | @generated unrolls loops per (rank, num_arrays) at compile time |
Handwritten 1D/2D/3D/4D specializations + generic N-D fallback |
| Inner-loop SIMD | Explicit @simd pragma on innermost loop |
Stride-specialized inner loops: slice-based when stride=1, raw pointer otherwise; relies on LLVM auto-vectorization |
| Threading | Recursive dimension-splitting via Threads.@spawn |
Recursive dimension-splitting via rayon::join; order-before-fuse pipeline enables layout-agnostic parallelization |
Note: Strided.jl threading bug for non-column-major views. Julia's pipeline fuses before ordering (
fuse → order → block), so_mapreduce_fuse!only detects column-major contiguity. Permuted views (e.g.PermutedDimsArray(A, (2,1))) with row-major strides are never fused, causing_mapreduce_threaded!to fall through to the single-threaded kernel. strided-rs fixes this by simply reordering the pipeline toorder → fuse → block: ordering first puts smallest-stride dimensions adjacent, and a single fusion pass then catches contiguity regardless of memory layout. See docs/strided_jl_threading_bug.md for a minimal reproduction and root cause analysis.