From 13ac5ec870a67f006c0f30ac8ab4757a9123a03f Mon Sep 17 00:00:00 2001
From: lkdvos <ldevos98@gmail.com>
Date: Tue, 23 Jun 2026 21:04:44 -0400
Subject: [PATCH] Specialize non-reducing kernel: drop reduction-only iszero
 hoist

The dim-1 inner core in `_mapreduce_kernel_expr` carried an
`iszero(stride_1_1)` branch that hoists `A1[I1]` out of the `@simd`
loop. That hoist only matters for reductions, where the destination
stride along the inner dim is zero. For the non-reducing path
(`op === nothing`: map!/permute/copy!/fill!/...) every destination
element is written exactly once, so the branch is dead and the loop
body was needlessly generated twice.

Emit the plain `@simd` loop (one body, no runtime branch) for
`op === nothing`, keeping the hoist for reductions. Runtime-neutral
(no permute regression) and trims a bit of non-reducing compile time.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 src/mapreduce.jl | 79 +++++++++++++++++++++++++++++++-----------------
 1 file changed, 51 insertions(+), 28 deletions(-)

diff --git a/src/mapreduce.jl b/src/mapreduce.jl
index d7c4659..5de8e6c 100644
--- a/src/mapreduce.jl
+++ b/src/mapreduce.jl
@@ -351,40 +351,63 @@ function _mapreduce_kernel_expr(f, op, initop, N::Int, M::Int)
     firststrides = Expr(:tuple, (stridevars[1, j] for j in 1:M)...)
     unitstridecond = :(all(==(1), $firststrides))
 
+    i = 1
     if op == Nothing
+        # Non-reducing (map!/permute/...): each destination element is written once,
+        # so skip the reduction-only `iszero` hoist below and emit just the
+        # unit-stride and generic `@simd` loops.
         ex = Expr(:(=), lhsex, fcallex)
-        exa = Expr(:(=), :a, fcallex)
+        if N >= 1
+            ex = quote
+                if $unitstridecond
+                    @simd for $(innerloopvars[i]) in Base.OneTo($(blockdimvars[i]))
+                        $ex
+                        $unitstep1ex
+                        $unitstep2ex
+                    end
+                    $(returnstride1ex[i])
+                    $(returnstride2ex[i])
+                else
+                    @simd for $(innerloopvars[i]) in Base.OneTo($(blockdimvars[i]))
+                        $ex
+                        $(stepstride1ex[i])
+                        $(stepstride2ex[i])
+                    end
+                    $(returnstride1ex[i])
+                    $(returnstride2ex[i])
+                end
+            end
+        end
     else
         ex = Expr(:(=), lhsex, Expr(:call, :op, lhsex, fcallex))
         exa = Expr(:(=), :a, Expr(:call, :op, :a, fcallex))
-    end
-    i = 1
-    if N >= 1
-        ex = quote
-            if iszero($(stridevars[1, 1])) # explicitly hoist A1[I1] out of loop
-                a = $lhsex
-                @simd for $(innerloopvars[i]) in Base.OneTo($(blockdimvars[i]))
-                    $exa
-                    $(stepstride2ex[i])
-                end
-                $lhsex = a
-                $(returnstride2ex[i])
-            elseif $unitstridecond
-                @simd for $(innerloopvars[i]) in Base.OneTo($(blockdimvars[i]))
-                    $ex
-                    $unitstep1ex
-                    $unitstep2ex
-                end
-                $(returnstride1ex[i])
-                $(returnstride2ex[i])
-            else
-                @simd for $(innerloopvars[i]) in Base.OneTo($(blockdimvars[i]))
-                    $ex
-                    $(stepstride1ex[i])
-                    $(stepstride2ex[i])
+        if N >= 1
+            ex = quote
+                if iszero($(stridevars[1, 1])) # explicitly hoist A1[I1] out of loop
+                    a = $lhsex
+                    @simd for $(innerloopvars[i]) in Base.OneTo($(blockdimvars[i]))
+                        $exa
+                        $(stepstride2ex[i])
+                    end
+                    $lhsex = a
+                    $(returnstride2ex[i])
+                elseif $unitstridecond
+                    @simd for $(innerloopvars[i]) in Base.OneTo($(blockdimvars[i]))
+                        $ex
+                        $unitstep1ex
+                        $unitstep2ex
+                    end
+                    $(returnstride1ex[i])
+                    $(returnstride2ex[i])
+                else
+                    @simd for $(innerloopvars[i]) in Base.OneTo($(blockdimvars[i]))
+                        $ex
+                        $(stepstride1ex[i])
+                        $(stepstride2ex[i])
+                    end
+                    $(returnstride1ex[i])
+                    $(returnstride2ex[i])
                 end
-                $(returnstride1ex[i])
-                $(returnstride2ex[i])
             end
         end
     end