Inline the strided-view mul! routing and return C explicitly - #764
Merged
Conversation
The routing methods added in JuliaGPU#760 sit in front of every GPU mul!, but the invoke back to LinearAlgebra's original implementation goes through an edge specialized on the abstract invoke signature, which severs constant propagation of alpha/beta. For plain 3-arg mul! the (true, false) constants previously folded MulAddMul into a concrete type; now the inlined generic body union-splits the MulAddMul constructor, the phi-merge widens it back to the abstract UnionAll, and the generic_matmatmul! call becomes a dynamic dispatch on every 3-arg GPU matmul. This also broke Enzyme's rule interception, which relies on that call devirtualizing (see EnzymeAD/Enzyme.jl#3468). Marking the wrappers @inline restores the chain: the routing check inlines into the caller, the invoke sees the constant scalars, invoke-level constant propagation fires, and MulAddMul folds concrete again, so generic_matmatmul! devirtualizes exactly as before JuliaGPU#760. It also removes an extra non-inlined call from the hot path. Returning C explicitly instead of the invoke's result makes the return value independent of inference precision on the invoke edge. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WgVTTES3ytMLYhZGyvQvf8
kshyatt
force-pushed
the
wsm/mul-routing-inline
branch
from
August 20, 2026 06:54
512d6e6 to
7510126
Compare
kshyatt
approved these changes
Aug 20, 2026
rsenne
approved these changes
Aug 20, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
AI-generated fix for GPUArrays on Enzyme which was slightly broken due to the recent mul! change [cc @kshyatt @rsenne ]
Follow-up to #760.
Problem
The routing methods from #760 sit in front of every GPU
mul!, and in the common no-view case forward viainvokeback to LinearAlgebra's original implementation. Thatinvokeedge is specialized on the abstract invoke signature, which severs constant propagation ofα/β.Concretely, for plain 3-arg
mul!the(true, false)constants previously const-foldedMulAddMul(α, β)into a concrete type, sogeneric_matmatmul!devirtualized. With the wrapper in between, the inlined generic body union-splits theMulAddMulconstructor into its 4 concrete types, the φ-merge widens them back to the abstractMulAddMul{ais1, bis0}UnionAll (union limit is 3), and thegeneric_matmatmul!call becomes a dynamic dispatch on every 3-arg GPU matmul (Julia ≤ 1.11; on 1.12+ the equivalent precision loss happens abovegeneric_matmatmul_wrapper!):Besides the dispatch overhead, this broke Enzyme's custom-rule interception, which relies on that call devirtualizing so the rule applies with correct activity info: every Enzyme.jl CI run fails since v11.5.11 (see EnzymeAD/Enzyme.jl#3468 for the downstream analysis).
Fix
Mark the four routing wrappers
@inlineandreturn Cexplicitly:@inline: the routing check inlines into the caller, theinvokethen sees the constant scalars at its own call site, invoke-level constant propagation fires, andMulAddMulfolds concrete again —generic_matmatmul!devirtualizes exactly as before Fix multiplication with strided GPU array views #760. It also removes an extra non-inlined call from the hot path of every GPUmul!.return C: makes the wrapper's return value independent of inference precision on the invoke edge, rather than relying on the abstract-signatureinvokeinferring a concrete result.Validation
With JLArrays on Julia 1.10 and 1.11 (both sides of the version gate):
mul!onviews of JLArrays still routes to the generic kernels and produces correct results (3- and 5-arg);generic_matmatmul!call in the 3-arg chain is static again (verified via Enzyme, whose static rule interception depends on it: its full GPUArrays-linalg test suite passes 28/28 against this branch with unmodified Enzyme, where v11.5.11 fails).🤖 Generated with Claude Code
https://claude.ai/code/session_01WgVTTES3ytMLYhZGyvQvf8