Skip to content

Commit 26b0013

Browse files
lkdvosclaude
andcommitted
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>
1 parent 5baa949 commit 26b0013

1 file changed

Lines changed: 30 additions & 17 deletions

File tree

src/mapreduce.jl

Lines changed: 30 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -340,25 +340,14 @@ function _mapreduce_kernel_expr(f, op, initop, N::Int, M::Int)
340340
outerreturnstrideex[i] = returnex
341341
end
342342

343+
i = 1
343344
if op == Nothing
345+
# Non-reducing (map!/permute/...): each destination element is written once,
346+
# so skip the reduction-only `iszero` hoist below and emit just the direct
347+
# `@simd` loop.
344348
ex = Expr(:(=), lhsex, fcallex)
345-
exa = Expr(:(=), :a, fcallex)
346-
else
347-
ex = Expr(:(=), lhsex, Expr(:call, :op, lhsex, fcallex))
348-
exa = Expr(:(=), :a, Expr(:call, :op, :a, fcallex))
349-
end
350-
i = 1
351-
if N >= 1
352-
ex = quote
353-
if iszero($(stridevars[1, 1])) # explicitly hoist A1[I1] out of loop
354-
a = $lhsex
355-
@simd for $(innerloopvars[i]) in Base.OneTo($(blockdimvars[i]))
356-
$exa
357-
$(stepstride2ex[i])
358-
end
359-
$lhsex = a
360-
$(returnstride2ex[i])
361-
else
349+
if N >= 1
350+
ex = quote
362351
@simd for $(innerloopvars[i]) in Base.OneTo($(blockdimvars[i]))
363352
$ex
364353
$(stepstride1ex[i])
@@ -368,6 +357,30 @@ function _mapreduce_kernel_expr(f, op, initop, N::Int, M::Int)
368357
$(returnstride2ex[i])
369358
end
370359
end
360+
else
361+
ex = Expr(:(=), lhsex, Expr(:call, :op, lhsex, fcallex))
362+
exa = Expr(:(=), :a, Expr(:call, :op, :a, fcallex))
363+
if N >= 1
364+
ex = quote
365+
if iszero($(stridevars[1, 1])) # explicitly hoist A1[I1] out of loop
366+
a = $lhsex
367+
@simd for $(innerloopvars[i]) in Base.OneTo($(blockdimvars[i]))
368+
$exa
369+
$(stepstride2ex[i])
370+
end
371+
$lhsex = a
372+
$(returnstride2ex[i])
373+
else
374+
@simd for $(innerloopvars[i]) in Base.OneTo($(blockdimvars[i]))
375+
$ex
376+
$(stepstride1ex[i])
377+
$(stepstride2ex[i])
378+
end
379+
$(returnstride1ex[i])
380+
$(returnstride2ex[i])
381+
end
382+
end
383+
end
371384
end
372385
for outer i in 2:N
373386
ex = quote

0 commit comments

Comments
 (0)