Subgroup matrix

All major platform APIs have now released a similar extensions for cooperative matrix:

Metal introduced simdgroup_matrix in MSL 3.1
HLSL has support in SM6.8 (currently experimental release)
SPIR-V/Vulkan released SPV_KHR_cooperative_matrix

This feature is very useful for accelerating ML operations by performing matrix calculations across multiple threads.

Would a wgpu extension providing an interface for this make it feasible?

Update on Google's thinking:

This is important, with clear use cases
Platform support is uneven so far. HLSL's SM 6.8 is still experimental.
We think the portability issues for subgroups also affect coopartive matrix multiplies. So we'd like to resolve subgroup portability first.

Google has an internal investigation that we'd like to publish here. It describes the intersection across underlying APIs, and a proposal. We need some time to post it.

Update: We've started prototyping. The design is still a bit fluid because some things are more work than we had hoped, and we're not sure some bits of the design are worth the cost.

We hope to have initial benchmark results to help motivate the feature, before restarting the discussion in public.

@dneto0

@dneto0 any updates in the meantime? I'm very interested in this :)

@dneto0 any updates in the meantime? I'm very interested in this :)

Yes, sorry about the radio silence.

We've made progress with a bunch of caveats:

We have a prototype guarded in Chrome by "enable-unsafe-webgpu", and requiringenable chromium_experimental_subgroup_matrix; in the WGSL shaders.
Validation code is incomplete. We have not broadly tested correctness or stability.
Our prototype exposes the corresponding features on Metal (macOS) and Vulkan (Android, ChromeOS)
We've just given it to internal teams to experiment with and give us feedback. Our team hasn't measured the performance gains.
Microsoft's Edge team has landed experimental support in ONNX, with reported results on Mac. [webgpu] Implement SubGroupMatrix based MatMulNBits for Metal microsoft/onnxruntime#23729
We are monitoring the development the corresponding feature in HLSL and D3D

At the WebGPU face-to-face meeting last week, Alan presented slides that included this link to a detailed design proposal:

https://dawn.googlesource.com/dawn/+/refs/heads/main/docs/dawn/features/subgroup_matrix.md

@if

Minutes from WGSL committee meeting, Wed Sep 30

JB: Alan provided a detailed presentation at the F2F and there is a link to a detailed design doc. (subgroup_matrix.md file in dawn features). That includes a section at teh end called proposal which talks about WGSL stuff and defines an extension (chromium_experimental_subgroup_matrix). This seems like the intended audience is internal google which talks about chromium. Thing I want to talk about is the next steps. David posted comment in March on status update. Looks like a lot of progress. How would we update for the present.
AB: Thought was to do a chromium proposal and clean up to do a WGSL proposal with what we expect the HLSL changes to be or do we make subgroup_id a language feature or an enable. HLSL will get it in a future version but we could emulate now if we wanted. Not sure if it's a good idea. Every platform has the ability to have that builtin already. Would look like the chromium one minus the chromium_experimental naming. Can bikeshed the names. If folks are happy we can do an enable, or an enable and a separate language extension if we want to have the builtin.
JB: The builtin is valuable for a lot of reasons
AB: Agreed
JB: Seems like it would make sense to break subgroup_id out to a proposal which would be an extension because it isn't available everywhere
AB: Depends how you would do it. You could emulate it with an atomic; might provide a different value from what a real subgroup_id feature would give. If you update your system and get a new D3D that supported it natively then the resulting value might change. We haven’t come to consensus internally on this. If you’re willing to allow that wobblineness then you can make it a language feature and pretend it’s available on all hardware.
JB: That's polyfilled on systems which don't support natively
AB: Which is just Windows.
CF: So this would be a workgroup atomic?
AB: Yea it would have to be.
RC: If we did emulate on D3D windows your answer didn't give me a lot of hope the emulation would give the same result as the real implementation. That's my hesitation.
AB: That's my hesitation as well because you're using the atomic. Can't guarantee the order the id is generated in. No the most satisfying but was floated as an idea by the HLSL idea at the f2f.
CF: If this is emulatable by an atomic, couldn’t a user emulate it as well if they really needed to, and were willing to take that compromise (of using a workgroup atomic).
AB: that's effectively how the cts tests work. We create a value in place of this value. THe difference is if we emulate you get the real value if available and emulated on one platform
JB: Making the user do the emulation means it's ore controlled for them .They can choose what implementation they get and decide to choose to transition. Means implementations only provide it when they offer the real platform feature. This feels like even if we offer subgroup_id as a device conditional feature it is still valuable. It's still going to be useful and that people can wrap in something and provide portability.
DS: Would that require some kind of @if or if constexpr. You’d need two separate shaders, or “don’t compile this part” conditional on the feature.
JP: in terms of user emulation, they are already doing it and they're doing it incorrectly. Not doing atomics, assuming linear layout and it mostly works until it doesn't. If we want to push emulation to them then requires education on devices that don't do this as they expected. So advantage of us doing it is giving something we know will work.
CF: Generally quite against emulation particularly in cases where it's hardware information and it just doesn't exist. Generally of the opinion we should expose our constraints directly and not try to paper over too much. Sympathetic to users doing this wrong. Emulation is a cost that every user using this would have to bear and wouldn't know they were bearing.
RC: Asked Jesse about emulation and said yes it would give you an id but he said the values you get out may not relate to the other values you get. The native wave index would probably correlate to sb_group_id. Emulated doesn't necessarily correlate. We'll give an id but not not relate. Depends on shader author assumptions on how it relates to there things.
JB: Right so you could emulate it but clarifies you will get different results
RC: Just that it may not be correlated when a native would.
PM: In emulation we need to use local storage which may add complexity as we don't use local storage for anything else that I'm aware of.
JB: Might effect limits as we're consuming space
PM: at least one u32
AB: The Naga implementation limits to 1d workgroups. The case where it does work out to emulate well. If you have a linear you probably only assign one way. Gets complicated with 2 and 3 as hardware does everything. If you stay in the linear box things work out. Once you get into upper dimensions then hardware could do things differently depending on how you access memory. IDs could be set to help with memory collessing.
JB: Lots of interesting discussion even if we break subgroup_id out. So, next steps, split subgroup_id to own proposal which seems limited and we can maybe just take as a PR.
AB: The way we do experimental in chrome for metal uses a feature which really didn't extend in metal 4. One is going to have broader implementation then the other. Not sure if MW has feedback from metal. Seems like going forward metal uses different objects so how forward looking do we want to be for this feature. We were planning on pushing out with simdgroup matrix. newer stuff is the performance primitives and tensor objects.
MW: Can look into that. Anything with metal4 won't work on intel macs.
AB: So, 2 proposals. 1 for subgroup_matrix and one for subgroup_id and we can talk about the enable vs language feature for subgroup_id and subgroup_matrix assumes subgroup_id
DN: We'll write up PRs for proposals to discuss and then prospectively land them after agree on direction.
JB: SGTM

My personal summary of today's committee meeting:

Detailed design doc looks good
Resolved: split out PR to add subgroup_id. Much discussion on pros and cons of polyfilling this with atomics on systems that can't provide it natively.
Resolved: split out rest into PR to add proposal

Proposal for subgroup matrix feature

Proposal for subgroup matrix feature (#5335)

Navigation Menu

Repository navigation

Subgroup matrix #4195

1 remaining item

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Subgroup matrix #4195

Description

Activity

PWhiddy commented on Jan 5, 2024

dneto0 commented on Mar 1, 2024

dneto0 commented on Jun 13, 2024

dneto0 commented on Sep 10, 2024

janhohenheim commented on Mar 3, 2025

dneto0 commented on Mar 28, 2025

jimblandy commented on Sep 22, 2025

jimblandy commented on Sep 30, 2025

1 remaining item

jimblandy commented on Sep 30, 2025

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions