From 65820a424785b8e0fee08da9189da3979e39be2d Mon Sep 17 00:00:00 2001
From: lkdvos <ldevos98@gmail.com>
Date: Thu, 29 Oct 2026 09:54:11 -0400
Subject: [PATCH] Forward to a lower-dimensional kernel for trailing singleton
 dimensions

Fused dimensions are kept as trailing dimensions of size 1, which still cost a loop each in
the kernel. The kernel now forwards to the kernel for its non-singleton dimensions. The views
are rewrapped with the lower dimensionality, so that this kernel is the same one that arrays
of that dimensionality compile, rather than one extra kernel per dimensionality and type.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 Project.toml     |  2 +-
 src/mapreduce.jl | 41 +++++++++++++++++++++++++++++++++++++++--
 2 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/Project.toml b/Project.toml
index 5cb347a..df7199d 100644
--- a/Project.toml
+++ b/Project.toml
@@ -32,7 +32,7 @@ LinearAlgebra = "1.6"
 Metal = "1.9"
 PrecompileTools = "1"
 Random = "1.6"
-StridedViews = "0.5.1"
+StridedViews = "0.5.5"
 Test = "1.6"
 TupleTools = "1.6"
 julia = "1.10"
diff --git a/src/mapreduce.jl b/src/mapreduce.jl
index b94e544..fb7fb93 100644
--- a/src/mapreduce.jl
+++ b/src/mapreduce.jl
@@ -532,6 +532,42 @@ function _mapreduce_kernel_expr(f, op, initop, N::Int, M::Int; unit::Bool = fals
     return ex
 end
 
+# Trailing singleton dimensions, e.g. those left by fusing dimensions, still cost a loop each.
+# Forward to the kernel of lower dimensionality instead, with the views rewrapped such that
+# this kernel is shared with arrays that have that dimensionality to begin with.
+function _mapreduce_trim_expr(N::Int, ex)
+    N > 1 || return ex
+    for k in 1:(N - 1)
+        ex = quote
+            if nactive == $k
+                dk = _firstk(dims, Val($k))
+                sk = _firstk(strides, Val($k))
+                return _mapreduce_kernel!(
+                    f, op, initop, dk, _firstk(blocks, Val($k)), _rewrap(arrays, dk, sk), sk, offsets
+                )
+            end
+            $ex
+        end
+    end
+    return quote
+        nactive = _lastnonsingleton(dims)
+        $ex
+    end
+end
+_firstk(t::Tuple, ::Val{k}) where {k} = ntuple(i -> t[i], Val(k))
+_firstk(ts::Tuple{Vararg{Tuple}}, v::Val) = map(t -> _firstk(t, v), ts)
+function _lastnonsingleton(dims::Dims)
+    n = length(dims)
+    while n > 1 && dims[n] == 1
+        n -= 1
+    end
+    return n
+end
+# the kernel only indexes the parents, so the strides of the views need no normalization
+_rewrap(arrays, dims, strides) = map(arrays, strides) do a, s
+    return StridedViews.unsafe_stridedview(eltype(a), parent(a), dims, s, offset(a), a.op)
+end
+
 @generated function _mapreduce_kernel!(
         (f), (op),
         (initop), dims::NTuple{N, Int},
@@ -540,16 +576,17 @@ end
         strides::NTuple{M, NTuple{N, Int}},
         offsets::NTuple{M, Int}
     ) where {N, M}
-    op === Nothing || return _mapreduce_kernel_expr(f, op, initop, N, M)
+    op === Nothing || return _mapreduce_trim_expr(N, _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
+    ex = 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
+    return _mapreduce_trim_expr(N, ex)
 end
 
 function indexorder(strides::NTuple{N, Int}) where {N}