msl: fix cooperative matrix loading semantics#9451
Open
kvark wants to merge 1 commit intogfx-rs:trunkfrom
Open
msl: fix cooperative matrix loading semantics#9451kvark wants to merge 1 commit intogfx-rs:trunkfrom
kvark wants to merge 1 commit intogfx-rs:trunkfrom
Conversation
The naga MSL backend forwarded the WGSL `coopLoad(T)` / `coopStore(T)` `row_major` flag straight through to Metal's `simdgroup_load` / `simdgroup_store` `transpose` argument. Metal's `transpose` means "memory is transposed relative to the simdgroup_matrix canonical layout" — and on Apple silicon that canonical layout is row-major, so `transpose=false` is the correct load for row-major memory. Naga's IR uses `row_major=true` to mean "memory is row-major"; the emission must therefore negate the flag when targeting Metal, matching the SPIRV backend which already correctly maps `row_major=true → RowMajorKHR`. Symptom: any WGSL using `coopLoadT` / `coopStoreT` on row-major memory produced numerically-wrong results on Apple silicon. The existing wgpu cooperative-matrix example masked this because it used `coopLoad` on row-major memory (the matching error cancelled the MSL one on Metal), and because its A-matrix init happened to be symmetric so the corresponding bug on Vulkan was also invisible. Neither regression path will be silent after this commit: - `naga/src/back/msl/writer.rs` emits the correct `transpose` flag for `coopLoad` / `coopStore` and `coopLoadT` / `coopStoreT`. The snapshot in `naga/tests/out/msl/wgsl-cooperative-matrix.metal` is regenerated accordingly (2-line diff). - `examples/features/.../shader.wgsl` now uses `coopLoadT` / `coopStoreT` so that the WGSL matches the "Row-major tiles" comment in the surrounding host code. This also serves as a working reference for the naga convention. - `examples/features/.../mod.rs` reinitialises matrix A with a row/column-weighted formula so neither A nor B is symmetric in (i, j); the test can no longer pass by coincidence regardless of which load variant a future regression picks. - `docs/api-specs/cooperative_matrix.md` documents both variants (`coopLoad` = column-major, `coopLoadT` = row-major), updates the worked example to use the T variants with its row-major buffers, and adds "variant does not match storage" as an undefined-behaviour case.
Member
Author
|
@jimblandy this is quite important to land ASAP. I was having issues going between Linux and MacOS development, and now realized the implementation had been inconsistent. Need to pick this up to continue cross-platform development. |
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
2 participants
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.
Connections
Following up on support from #8251.
Description
Looks like this was broken and it slipped through because the example was using a symmetric matrix.
MSL loaded cooperative matrix data transposed compared to SPV.
The internal spec document also needs to be fixed - specifying that it's ColMajor by default is consistent with the rest of matrix logic in WGSL.
Testing
Example, updated test
Squash or Rebase?
Squash is fine.
Checklist
wgpumay be affected behaviorally.CHANGELOG.mdentries for the user-facing effects of this change are present.