Specialize non-reducing kernel: drop reduction-only iszero hoist#69
Conversation
Codecov Report❌ Patch coverage is
... and 2 files with indirect coverage changes 🚀 New features to boost your workflow:
|
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>
lkdvos
force-pushed
the
ld-nonreduce-kernel
branch
from
26b0013 to
13ac5ec
Compare
| outerstepstrideex = Vector{Expr}(undef, N) | ||
| outerreturnstrideex = Vector{Expr}(undef, N) | ||
| for i in 1:N | ||
| stepex = Expr(:block) | ||
| returnex = Expr(:block) | ||
| for j in 1:M | ||
| push!(stepex.args, :($(Ivars[j]) += $(blockdimvars[i]) * $(stridevars[i, j]))) | ||
| push!(returnex.args, :($(Ivars[j]) -= dims[$i] * $(stridevars[i, j]))) | ||
| end | ||
| outerstepstrideex[i] = stepex | ||
| outerreturnstrideex[i] = returnex | ||
| end |
There was a problem hiding this comment.
Do we still use these variables? I am a bit rusty on the logic of this generated function.
There was a problem hiding this comment.
I think so, somewhere in the outer block loop ( although I'm honestly rusty on the logic every time I don't look at it for a while, I have to just generate some of these expressions to really see what is going on... )
There was a problem hiding this comment.
Ah yes, the github view conveniently ended before the end of the function, where the outer part blocks started appearing.
There was a problem hiding this comment.
Maybe Claude can be asked if this generating function body can be simplified, while still producing the same generated code. The current approach is quite cumbersome.
There was a problem hiding this comment.
I tried this multiple times already, with little success...
Jutho
left a comment
Jutho
left a comment
There was a problem hiding this comment.
Actual changes look good to me; but I added a question on some other part of the generated function.
2 participants
The generated inner loop of
_mapreduce_kernel!always includes aniszero(stride_1_1)branch that hoistsA1[I1]out of the@simdloop. This hoist only matters for reductions, where the destination has stride 0 along the inner dimension. On the non-reducing path (op === nothing:map!,permutedims!,copy!,conj!, ...) every destination element is written exactly once, so the branch is dead code — yet it was still emitted, and can't be eliminated because the condition is not compile time constant, therefore duplicating the inner loop body and leaving a runtime branch in the hot loop.This PR keys the inner-loop generation on
op: the non-reducing case now emits only the unit-stride and generic@simdloops, while the reducing case is unchanged. This is a pure generation-time specialization, so should not affect actual runtimes.Benchmarks
Compile time of 12 fresh kernel instantiations (map! + permute over various eltypes and ranks), median of 5 alternating trials against
main, with the untouched reducing kernels as drift control:Net effect: roughly 10% less compile time for non-reducing kernel instantiations, runtime-neutral.
🤖 Generated with assistance of Claude Code