From 1fa9461af2eaa498c9c2d73f516e45565b0ad7f2 Mon Sep 17 00:00:00 2001
From: lkdvos <ldevos98@gmail.com>
Date: Thu, 29 Oct 2026 09:53:04 -0400
Subject: [PATCH 1/2] Use while loops for the block loops in the kernel

Constructing the step range `1:blocks[i]:dims[i]` computes its last element with an integer
division, which is paid for every block loop in the nest, also when the arrays fit in a
single block. A while loop only needs comparisons and additions.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/mapreduce.jl | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/mapreduce.jl b/src/mapreduce.jl
index 824e5c5..72eb34f 100644
--- a/src/mapreduce.jl
+++ b/src/mapreduce.jl
@@ -497,10 +497,12 @@ function _mapreduce_kernel_expr(f, op, initop, N::Int, M::Int)
     if initop === Nothing
         for outer i in 1:N
             ex = quote
-                for $(blockloopvars[i]) in 1:blocks[$i]:dims[$i]
+                $(blockloopvars[i]) = 1
+                while $(blockloopvars[i]) <= dims[$i]
                     $(blockdimvars[i]) = min(blocks[$i], dims[$i] - $(blockloopvars[i]) + 1)
                     $ex
                     $(outerstepstrideex[i])
+                    $(blockloopvars[i]) += blocks[$i]
                 end
                 $(outerreturnstrideex[i])
             end
@@ -516,11 +518,13 @@ function _mapreduce_kernel_expr(f, op, initop, N::Int, M::Int)
         for outer i in 1:N
             ex = quote
                 $(initvars[i]) = $(initvars[i + 1])
-                for $(blockloopvars[i]) in 1:blocks[$i]:dims[$i]
+                $(blockloopvars[i]) = 1
+                while $(blockloopvars[i]) <= dims[$i]
                     $(blockdimvars[i]) = min(blocks[$i], dims[$i] - $(blockloopvars[i]) + 1)
                     $ex
                     $(initvars[i]) &= $(stridevars[i, 1]) > 0
                     $(outerstepstrideex[i])
+                    $(blockloopvars[i]) += blocks[$i]
                 end
                 $(outerreturnstrideex[i])
             end

From 38b51de8cba8dbfae85bd1e723830b3a1d4f317b Mon Sep 17 00:00:00 2001
From: lkdvos <ldevos98@gmail.com>
Date: Thu, 29 Oct 2026 09:53:21 -0400
Subject: [PATCH 2/2] Hoist the unit-stride check out of the loop nest for
 non-reducing kernels

Whether the innermost dimension has unit stride in all arrays is loop invariant, but it was
checked every time the innermost loop started. For arrays with short innermost dimensions
this is a large part of the work. Non-reducing kernels now generate the loop nest for both
cases and select one up front.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/mapreduce.jl | 40 ++++++++++++++++++++--------------------
 1 file changed, 20 insertions(+), 20 deletions(-)

diff --git a/src/mapreduce.jl b/src/mapreduce.jl
index 72eb34f..b94e544 100644
--- a/src/mapreduce.jl
+++ b/src/mapreduce.jl
@@ -297,7 +297,7 @@ function _mapreduce_threaded!(
     return nothing
 end
 
-function _mapreduce_kernel_expr(f, op, initop, N::Int, M::Int)
+function _mapreduce_kernel_expr(f, op, initop, N::Int, M::Int; unit::Bool = false)
     @assert N > 0 "Cannot generate 0-dim kernel"
 
     # many variables
@@ -388,28 +388,19 @@ function _mapreduce_kernel_expr(f, op, initop, N::Int, M::Int)
     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.
+        # so skip the reduction-only `iszero` hoist below. Whether the innermost loop
+        # has unit stride is loop invariant, and is selected by the caller via `unit`.
         ex = Expr(:(=), lhsex, fcallex)
         if N >= 1
+            step1ex, step2ex = unit ? (unitstep1ex, unitstep2ex) : (stepstride1ex[i], stepstride2ex[i])
             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])
+                @simd for $(innerloopvars[i]) in Base.OneTo($(blockdimvars[i]))
+                    $ex
+                    $step1ex
+                    $step2ex
                 end
+                $(returnstride1ex[i])
+                $(returnstride2ex[i])
             end
         end
     else
@@ -549,7 +540,16 @@ end
         strides::NTuple{M, NTuple{N, Int}},
         offsets::NTuple{M, Int}
     ) where {N, M}
-    return _mapreduce_kernel_expr(f, op, initop, N, M)
+    op === Nothing || return _mapreduce_kernel_expr(f, op, initop, N, M)
+    # select the unit-stride variant of the innermost loop once, instead of in every iteration
+    unitcond = Expr(:call, :all, :(==(1)), Expr(:tuple, (:(strides[$j][1]) for j in 1:M)...))
+    return quote
+        if $unitcond
+            $(_mapreduce_kernel_expr(f, op, initop, N, M; unit = true))
+        else
+            $(_mapreduce_kernel_expr(f, op, initop, N, M; unit = false))
+        end
+    end
 end
 
 function indexorder(strides::NTuple{N, Int}) where {N}