From 1fa9461af2eaa498c9c2d73f516e45565b0ad7f2 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 29 Oct 2026 09:53:04 -0400 Subject: [PATCH 1/3] 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 --- 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 Date: Thu, 29 Oct 2026 09:53:21 -0400 Subject: [PATCH 2/3] 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 --- 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} From 65820a424785b8e0fee08da9189da3979e39be2d Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 29 Oct 2026 09:54:11 -0400 Subject: [PATCH 3/3] 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 --- 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}