Use the explicit gemm for wide convs on neural accelerators - #4214
Open
erwinzhang7 wants to merge 1 commit into
Open
Use the explicit gemm for wide convs on neural accelerators#4214erwinzhang7 wants to merge 1 commit into
erwinzhang7 wants to merge 1 commit into
Conversation
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.
conv.cpppicks between the implicit gemm and unfolding into an explicit gemm. That choice wastuned when both paths ran on
BlockMMA.steel_matmulnow takes the NAX path on hardware withneural accelerators, so one side of the tradeoff got a lot faster and the boundary hasn't moved.
Measured on an M5 Max (
applegpu_g17s), before and after, over eight shape families thatcurrently reach the implicit gemm, so nothing winograd takes. n=8 per cell:
The 128 and 256 rows are a control: they are below the threshold, so both builds run the same
code for them and the true ratio is 1.00x. They measure 0.97x to 1.10x, which puts the noise
floor around 10% and leaves every real result well clear of it.
The axis is the output channel count, not the filter size or the spatial extent. Unfolding costs
an extra pass over the input, roughly
2*M*Kelements, while the gemm that follows does2*M*N*K, so the work per materialized element scales withN. HoldingC_outfixed and movingKfrom 1152 to 9216 leaves the ratio flat (2.62x, 2.47x, 2.58x, 2.58x); holdingKfixed andmoving
C_outwalks it across the whole range in the table above. The same behaviour shows upfor 3x3, 5x5 and 7x7 filters, stride 1 and 2, dilation 1 and 2, aligned and unaligned channels,
and
Mfrom 1024 to 65536, all within about 0.1x.Gating on
is_nax_available()because without the accelerators the explicit path isworse. On an M4 Pro (
applegpu_g16s) it loses at every size, reaching 1.01x only atC_out = 4096:Restricted to float16 and bfloat16 because for float32 the gain is precision, not speed. NAX runs
float32 at TF32 mantissa, and conv doesn't do that today. With
MLX_ENABLE_TF32=0so both pathsrun true float32, the explicit path loses everywhere: 0.52x, 0.67x, 0.78x, 0.88x for the same
four sizes.
Verified
M5 Max, macOS 26.6, and M4 Pro, both against a CPU reference: 144 cases and 36 cases over filter
1/3/5/7, stride 1/2, padding 0/1/2, dilation 1/2, aligned and unaligned channels, two batch
shapes.
float16 and bfloat16 sit at dtype rounding. float32 is two orders tighter, which is the check
that it stays off the TF32 path.
Separately, 504 cases per dtype comparing the two dispatch paths directly against each other:
worst relative error 7.79e-04 for float16, 6.06e-03 for bfloat16, no shape mismatches.
On an M4 Pro the branch is unreachable, so it runs exactly what it does today.
Notes
Winograd, depthwise, grouped and the small-channel paths are untouched; the branch sits after
those and only catches what would otherwise go to the implicit gemm.
The explicit path already bounds its own memory.
max_unfold_rowstiles the unfold againstmaxBufferLengthand reuses one buffer, so this doesn't introduce an unbounded allocation.C_out >= 512rather than 256 because at 256 the win is thin onceMgets large (1.03x to 1.19xat
M = 16384and above), which isn't worth the extra pass. Happy to move it if you'd ratherhave the 1.2x.