[lkdvos]

This PR reduces the fixed cost of the mapreduce kernel, which dominates for small arrays such as the blocks of a block-sparse tensor.

• Block loops as while loops. for J in 1:blocks[i]:dims[i] constructs a step range for every block loop, which computes its last element with an integer division. A while loop only needs comparisons and additions.
• Unit-stride check hoisted. For non-reducing kernels, the check whether the innermost dimension has unit stride in all arrays is loop invariant, but was evaluated every time the innermost loop started. The kernel now generates both variants of the loop nest and selects one up front.

Neither change affects first-call compile time noticeably.

Benchmarks

C = α * permutedims(A, (1, 4, 2, 5, 3, 6)) through _mapreducedim!, as TensorOperations calls it, with β = 0, single-threaded, time in ns:

shape	elements	main Float64	PR Float64	main ComplexF64	PR ComplexF64
(2, 2, 1, 1, 1, 1)	4	214	182	234	208
(2, 2, 2, 2, 1, 1)	16	248	208	281	233
(2, 2, 2, 2, 2, 2)	64	382	301	436	344
(4, 4, 2, 2, 2, 2)	256	479	409	554	486
(8, 8, 2, 2, 2, 2)	1024	636	552	1014	963
(16, 16, 4, 1, 4, 1)	4096	1261	1173	2740	2697
(32, 32, 8, 1, 8, 1)	65536	18434	20071	76501	77292

The largest case is within the noise of the machine. The reducing kernel that is used for β ≠ 0 only gains from the first change, a few percent.

First calls in a fresh session (a 5-d permutation, a broadcast, a reduction and a map!): 2.27 s on this branch against 2.35 s on main.

Large arrays

To check for regressions on large arrays, this branch and main were loaded side by side in one process and sampled alternately, which keeps load on the machine from biasing the comparison. Median time of this PR relative to main, single-threaded, Intel Xeon Gold 6244 (AVX-512):

case	Float64	ComplexF64
permutation 64⁴, (3, 1, 4, 2), β = 0	0.98–0.99	0.78–0.82
permutation 16⁵, (5, 3, 1, 4, 2), β = 0	1.09–1.13	1.03–1.12
transpose 4096², β = 0	1.04	0.70–0.76
permutation 32⁴, (4, 3, 2, 1), β = 0	0.74–0.86	0.64–0.70
permutations above, β ≠ 0	0.99–1.01	1.00–1.04
map! over 3 contiguous arrays, 2²² elements	1.00	1.00
map! over 3 stride-2 arrays, 1024 × 4096	1.00	0.99
sum over dims (1, 3) and 2, 256 × 256 × 64	1.00	1.00

The while loops are neutral here; the differences come from the hoisted unit-stride check. With the check inside the loop nest, LLVM versions the strided inner loop for a destination with unit stride. Without it, the strided nest has no such version:

• Float64: the 16⁵ case loses the loop that gathers from the source and stores contiguously, and scatters into the destination instead.
• ComplexF64: the strided inner loop is no longer vectorized at all. That is faster on this machine, because vectorized gathers of complex numbers are slow.

Hoisting a third, destination-unit-stride variant fixes the Float64 16⁵ case (1.00), but also brings back the slow ComplexF64 loops (0.91–1.15 relative to main), so this PR does not do that. The balance may differ on AVX2 machines, which have no scatter instructions.

With 8 threads the measurements are noisier; the minimum times are within ±5% of main except for the same ComplexF64 improvements.

The tests pass, single- and multi-threaded (4 threads), including the JLArray and Aqua tests.

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 90.00000% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/mapreduce.jl	90.00%	2 Missing ⚠️

Files with missing lines	Coverage Δ
src/mapreduce.jl	93.75% <90.00%> (-0.22%)	⬇️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.