Skip to content

Faster fallback matrix multiplication - #590

Closed
christiangnrd wants to merge 9 commits into
JuliaGPU:mainfrom
christiangnrd:fastmatmul
Closed

Faster fallback matrix multiplication#590
christiangnrd wants to merge 9 commits into
JuliaGPU:mainfrom
christiangnrd:fastmatmul

Conversation

@christiangnrd

@christiangnrd christiangnrd commented Apr 13, 2025 •
edited
Loading

Copy link
Copy Markdown
Member

Adapted from the Metal scalar gemm kernel and modified for the macro-based KernelAbstractions interface based on the KA "performant_matmul.jl" example but with some fixes.

@maleadt

This comment was marked as resolved.

Sign in to view
@christiangnrd

This comment was marked as outdated.

Sign in to view
@christiangnrd christiangnrd mentioned this pull request Apr 18, 2025
Closed
@christiangnrd

This comment was marked as outdated.

Sign in to view
@christiangnrd christiangnrd reopened this Apr 19, 2025
@christiangnrd
christiangnrd force-pushed the fastmatmul branch 2 times, most recently from 3892d1b to 1499c12 Compare April 25, 2025 15:03
@christiangnrd christiangnrd mentioned this pull request May 15, 2025
Merged
@christiangnrd

christiangnrd commented Aug 25, 2026 •
edited
Loading

Copy link
Copy Markdown
Member Author

OpenCL POCL before and after:

eltype size main PR speedup
Float16 512³ 23 45 2.0×
Float16 1024³ 22 47 2.1×
Float16 2048³ 13 45 3.5×
Float16 4096³ 2 45 22.5×
Int32 512³ 24 60 2.5×
Int32 1024³ 13 60 4.6×
Int32 2048³ 4 59 14.8×
Int32 4096³ 2 37 18.5×
Float32 512³ 21 44 2.1×
Float32 1024³ 13 46 3.5×
Float32 2048³ 4 44 11.0×
Float32 4096³ 2 31 15.5×

OpenCL RUSTICL before and after:

eltype size main PR speedup
Float16 512³ 159 357 2.0×
Float16 1024³ 214 476 2.1×
Float16 2048³ 230 508 3.5×
Float16 4096³ 212 490 22.5×
Int32 512³ 151 329 2.5×
Int32 1024³ 210 455 4.6×
Int32 2048³ 216 472 14.8×
Int32 4096³ 194 432 18.5×
Float32 512³ 151 330 2.1×
Float32 1024³ 208 457 3.5×
Float32 2048³ 214 472 11.0×
Float32 4096³ 194 331 15.5×

@christiangnrd
christiangnrd marked this pull request as ready for review August 25, 2026 02:18
@christiangnrd
christiangnrd requested a review from maleadt August 25, 2026 02:19
@christiangnrd

christiangnrd commented Aug 25, 2026

Copy link
Copy Markdown
Member Author

I updated this to include the overflow and other various improvements recently applied to the fallback matmul, and reimplemented on top of KernelAbstractions instead of KernelInterface so we can merge sooner rather than later.

@maleadt

maleadt commented Aug 26, 2026

Copy link
Copy Markdown
Member

I had to make a few changes, such as using unsafe_indices, since the algorithm itself does the bounds checking, and I was getting wrong results until I added that.

Oof, that's bad, and unexpected.

I had Claude look into the unsafe_indices question. It's not a bug in KA or Metal — it's KA's index masking doing more than one might expect.

Without unsafe_indices, @kernel splits the body at every @synchronize and wraps each chunk of user code in if __validindex(ctx), leaving only the barriers unconditional. So when ndrange = size(C) isn't a multiple of the tile, out-of-range lanes still hit the barriers (good), but they never execute their tile[li, lj] = ... loads. In-range lanes then read stale threadgroup memory for those slots — 3 3 2 / 3 3 2 / 2 2 2 is one real product plus one stale zero per tile, and the 16.6 on Metal is the same slot containing garbage (Metal doesn't zero local memory; CUDA happened to).

Confirmed on Metal with TILE=2, 3×3: masked kernel + ndrange=(3,3) is wrong; masked + padded ndrange=(4,4) is fine; unsafe_indices=true is fine either way.

So for a cooperative tiled kernel that does its own bounds checks, unsafe_indices=true is the right declaration — all lanes have to participate in the loads. Padding the ndrange alone would also work, but I'd keep the flag (with a comment) so it doesn't silently regress. The I/J-undefined-after-@synchronize issue is the same split machinery: each chunk is its own scope, and only the CPU path still splits when unsafe_indices=true, which is why it showed up on some backends only.

Note that KA's performant_matmul.jl example has the same latent problem (ndrange = size(C), no unsafe_indices); it only works for tile-aligned sizes.

@maleadt @claude
Respect β in generic_matmatmul! when the inner dimension is empty
8e48910 
The early return zeroed C whenever A or B was empty, but for K == 0 the
result is β*C. The tiled kernel already handles that (no tiles to process),
so only skip the launch when C itself is empty.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@christiangnrd christiangnrd changed the title Faster (still slow) fallback matrix multiplication Faster fallback matrix multiplication Aug 26, 2026
@maleadt maleadt self-assigned this Aug 26, 2026
@christiangnrd

christiangnrd commented Aug 26, 2026

Copy link
Copy Markdown
Member Author

Note that KA's performant_matmul.jl example has the same latent problem (ndrange = size(C), no unsafe_indices); it only works for tile-aligned sizes.

The 0.9 example hasn't been fixed, but the 0.10 example got fixed (and then I switched it over to KernelInterface so it was no longer an issue)

I've updated the PR description to reflect the current state

maleadt and others added 5 commits August 26, 2026 16:19
@maleadt @claude
Merge opA/opB into a single opelem
018ea2f 
Both computed the same thing, element [i, j] of op(X); only the call sites
differ in which index is the contraction one.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@maleadt @claude
Drop the BANK parameter from the matmul kernel
a67d216 
A fallback kernel does not need a tunable here; hardcode the one-row padding
that avoids shared-memory bank conflicts.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@maleadt @claude
Rename MAX_TILE_DIM to TILE_DIM
cfd2452 
It is the tile dimension, not an upper bound.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@maleadt @claude
Keep the strided-view intercepts on the 5-arg generic_matmatmul!
2bc3f5f 
It already unwraps the operands via wrapper_char/_unwrap, instead of
indexing through the wrapper with 'N' flags.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@maleadt @claude
Exercise the matmul fallback with non-tile-aligned shapes
d58ca26 
The gemm tests only used 4×4 inputs, which never exercise partial tiles or
an inner dimension spanning several tiles, i.e. exactly the failure mode of
a tiled kernel. Use 33×17×5 (and n=33 for Symmetric/Hermitian, which was a
tile-aligned 128) so every dimension has a partial tile; the test load does
not increase.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@maleadt

maleadt commented Aug 26, 2026 •
edited
Loading

Copy link
Copy Markdown
Member

Pushed a handful of clean-ups/minimizations by Fable.

@christiangnrd

christiangnrd commented Aug 26, 2026

Copy link
Copy Markdown
Member Author

Nice! Should we stack this on top of #772 to make sure those tests pass before merging?

Also, do you mind applying the 018ea2f change in Metal? That code was ripped from there

@maleadt

maleadt commented Aug 26, 2026

Copy link
Copy Markdown
Member

Nice! Should we stack this on top of #772 to make sure those tests pass before merging?

Yeah that's reasonable.

Also, do you mind applying the 018ea2f change in Metal? That code was ripped from there

JuliaGPU/Metal.jl#935

@christiangnrd

christiangnrd commented Aug 26, 2026

Copy link
Copy Markdown
Member Author

Julia 1.11 Metal failure seems unrelated?

@christiangnrd
christiangnrd changed the base branch from main to stmmtest August 26, 2026 14:53
@christiangnrd
christiangnrd changed the base branch from stmmtest to main August 26, 2026 14:53
@christiangnrd

christiangnrd commented Aug 26, 2026 •
edited
Loading

Copy link
Copy Markdown
Member Author

This PR on from my fork lets just wait for #772 to be merged and rebase.

@maleadt maleadt mentioned this pull request Aug 26, 2026
Closed
@maleadt

maleadt commented Aug 26, 2026

Copy link
Copy Markdown
Member

CI failure is JuliaGPU/Metal.jl#937

@maleadt maleadt assigned christiangnrd and unassigned maleadt Aug 26, 2026
@maleadt

maleadt commented Aug 26, 2026

Copy link
Copy Markdown
Member

This PR on from my fork lets just wait for #772 to be merged and rebase.

OK, I'll let you apply the stack if you don't mind.

@christiangnrd christiangnrd mentioned this pull request Aug 26, 2026
Merged
@christiangnrd

christiangnrd commented Aug 26, 2026 •
edited
Loading

Copy link
Copy Markdown
Member Author

OK, I'll let you apply the stack if you don't mind.

#773

@maleadt maleadt closed this Aug 27, 2026
@christiangnrd
christiangnrd deleted the fastmatmul branch August 27, 2026 10:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@maleadt maleadt maleadt approved these changes

Assignees

@christiangnrd christiangnrd

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@christiangnrd @maleadt