Skip to content

Inline the strided-view mul! routing and return C explicitly - #764

Merged
kshyatt merged 1 commit into
JuliaGPU:mainfrom
wsmoses:wsm/mul-routing-inline
Aug 21, 2026
Merged

Inline the strided-view mul! routing and return C explicitly#764
kshyatt merged 1 commit into
JuliaGPU:mainfrom
wsmoses:wsm/mul-routing-inline

Conversation

@wsmoses

@wsmoses wsmoses commented Aug 19, 2026

Copy link
Copy Markdown
Member

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 via invoke back to LinearAlgebra's original implementation. That invoke edge is specialized on the abstract invoke signature, which severs constant propagation of α/β.

Concretely, for plain 3-arg mul! the (true, false) constants previously const-folded MulAddMul(α, β) into a concrete type, so generic_matmatmul! devirtualized. With the wrapper in between, the inlined generic body union-splits the MulAddMul constructor into its 4 concrete types, the φ-merge widens them back to the abstract MulAddMul{ais1, bis0} UnionAll (union limit is 3), and the generic_matmatmul! call becomes a dynamic dispatch on every 3-arg GPU matmul (Julia ≤ 1.11; on 1.12+ the equivalent precision loss happens above generic_matmatmul_wrapper!):

23%41 = φ (#18 => %31, #19 => %33, #21 => %37, #22 => %39)::LinearAlgebra.MulAddMul{ais1, bis0, Bool, Bool} where {ais1, bis0}%42 = LinearAlgebra.generic_matmatmul!(C, 'N', 'N', A, B, %41)::JLArray{Float64, 2}   # dynamic

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 @inline and return C explicitly:

  • @inline: the routing check inlines into the caller, the invoke then sees the constant scalars at its own call site, invoke-level constant propagation fires, and MulAddMul folds 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 GPU mul!.
  • return C: makes the wrapper's return value independent of inference precision on the invoke edge, rather than relying on the abstract-signature invoke inferring a concrete result.

Validation

With JLArrays on Julia 1.10 and 1.11 (both sides of the version gate):

  • strided-view routing unchanged: mul! on views of JLArrays still routes to the generic kernels and produces correct results (3- and 5-arg);
  • the 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

@wsmoses
wsmoses requested a review from maleadt August 19, 2026 09:14
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
kshyatt force-pushed the wsm/mul-routing-inline branch from 512d6e6 to 7510126 Compare August 20, 2026 06:54
@kshyatt
kshyatt merged commit 2fb8460 into JuliaGPU:main Aug 21, 2026
18 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants