mapreduce: add exact-arithmetic regression test for strided mixed reductions

Follow-up to the #7 resolution, addressing the adversarial-review caveats: the
prior comment overstated coalescing for a small leading dim, and the empirical
sweep left nothing committed and used Float32 (which can pass or fail for the
wrong reason on large sums).

Tighten the comment to the precise mechanism (the miscompile needs adjacent
lanes a *full* stride apart on every lane; reducing dim 1 keeps a contiguous
innermost axis, so lanes read stride-1 within each length-n1 run) and add an
Int32 regression test covering dims=(1,3)/(1,4)/(1,2,3) plus the coalesced
dims=2/3 cases. Int32 is exact and associative, so it is immune to the
Float32 accumulation-order rounding that made a stress sweep of >2^24-length
sums look corrupted when it was not; verified all cases match the CPU exactly,
including small leading dims and large strided sizes.

Refines PR_REVIEW.md finding #7.

Author: michel2323

src/mapreduce.jl

@@ -199,13 +199,18 @@ function GPUArrays.mapreducedim!(f::F, op::OP, R::oneWrappedArray{T},
     # reductions get less parallelism but stay correct; the common many-slice case is also fast.
     #
     # `size(Rreduce, 1) == 1` (i.e. dim 1 kept) is the correct predicate, not just a special
-    # case: in `partial_mapreduce_device` consecutive work-item lanes step through consecutive
-    # `Rreduce` entries, whose *first* axis is dim 1. When dim 1 is among the reduced dims the
-    # innermost reduced axis is contiguous, so those lane reads coalesce and the result is
-    # correct even if a strided dim (e.g. dim 3) is *also* reduced — the miscompile only bites
-    # when the innermost reduced axis is itself strided (dim 1 kept). Empirically verified: a
-    # sweep of `dims=(1,3)`-style mixed reductions with small leading dims (n1 ∈ 2..7) matched
-    # the CPU exactly, so no extra materialization is needed for those (see PR_REVIEW.md #7).
+    # case. In `partial_mapreduce_device` adjacent work-item lanes step through adjacent
+    # `Rreduce` entries, whose *fastest* axis is array dim 1. The miscompile bites only when
+    # adjacent lanes land a *full stride* apart on every lane — i.e. when the innermost reduced
+    # axis is itself strided, which is exactly the dim-1-kept case routed to the coalesced
+    # kernel above. When dim 1 is among the reduced dims, adjacent lanes instead read
+    # consecutive memory (stride 1) within each length-n1 run, with only an occasional jump at
+    # a run boundary — not the every-lane full-stride pattern — so the reads stay well-formed
+    # even when a strided dim (e.g. dim 3) is *also* reduced. Verified with exact Int32
+    # reductions (immune to the accumulation-order rounding that makes >2^24-length Float32
+    # sums differ between GPU and CPU): `dims=(1,3)`, `(1,4)` and full `(1,2,3)` reductions at
+    # large strided sizes match the CPU exactly, so these need no extra materialization (see
+    # PR_REVIEW.md #7 and the "strided mixed reductions" regression test).
     if oneL0.LTS[] && size(Rreduce, 1) == 1
         # cap the group size at what this kernel actually supports on this device (queried
         # from the driver), rather than a hardcoded 256 that can exceed the kernel's max

test/array.jl

@@ -104,3 +104,24 @@ end
   resize!(b, 1)
   @test length(b) == 1
 end
+
+@testset "strided mixed reductions" begin
+  # The Aurora LTS IGC miscompiles a reduction kernel's global reads when the *innermost*
+  # reduced axis is strided (dim 1 kept, e.g. `dims=2`); mapreducedim! routes those to a
+  # coalesced kernel. Reductions that also reduce dim 1 (e.g. `dims=(1,3)`) keep a contiguous
+  # innermost axis and stay correct on the workgroup-per-slice kernel — including with a small
+  # leading dim, where the contiguous run is short. Use Int32 (exact, associative) so the
+  # check is immune to Float32 accumulation-order rounding. Regression for PR_REVIEW.md #7.
+  for (sz, dts) in (
+      ((2, 512, 64),    ((1, 3), (2,), (3,), (2, 3), (1, 2, 3), 1)),
+      ((3, 256, 48),    ((1, 3), (2,), (1, 2, 3))),
+      ((7, 300, 20),    ((1, 3), (2,), (3,))),
+      ((2, 128, 8, 32), ((1, 3), (1, 4), (2, 4), (1, 2, 4))),
+    )
+    A = rand(Int32(1):Int32(4), sz...)
+    dA = oneArray(A)
+    for dt in dts
+      @test Array(sum(dA; dims = dt)) == sum(A; dims = dt)
+    end
+  end
+end