webgpu: fix OOB subgroup-matrix loads for trailing partial tiles - #32364
Merged
Hariharan Seshadri (hariharans29) merged 2 commits intoSep 8, 2026
Merged
Conversation
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
The affected Gemm and MatMul paths lack enabled numerical regression coverage.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Prevents WebGPU subgroup-matrix Gemm/MatMul kernels from loading beyond trailing partial tiles.
Changes:
- Shifts partial M/N tiles in bounds and skips overlapping writes.
- Falls back when matrices are smaller than selected tiles.
- Regenerates WGSL template goldens.
File summaries
| File | Description |
|---|---|
tools/python/wgsl_template/test/in_tree_golden/static-cpp/string_table.h |
Updates generated shader strings. |
tools/python/wgsl_template/test/in_tree_golden/static-cpp/index_impl.h |
Refreshes generated hashes. |
tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/pad.h |
Reindexes generated strings. |
tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/oihw_to_ohwi.h |
Reindexes generated strings. |
tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/im2col_matmul.h |
Reindexes generated strings. |
tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_matmul_pad_b.h |
Reindexes generated strings. |
tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/math/subgroup_matrix_matmul_8x16x16.h |
Captures generated MatMul changes. |
tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/index_impl.h |
Refreshes literal-generator hashes. |
onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc |
Adds small-matrix fallback. |
onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul_8x16x16.wgsl.template |
Makes trailing MatMul tiles safe. |
onnxruntime/core/providers/webgpu/math/subgroup_matrix_gemm.cc |
Adds small-matrix fallback. |
onnxruntime/core/providers/webgpu/math/subgroup_matrix_gemm_8x16x16.wgsl.template |
Makes trailing Gemm tiles safe. |
Review details
- Files reviewed: 7/15 changed files
- Comments generated: 2
- Review effort level: Balanced
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Contributor
Author
…tiles The 8x16x16 gemm/matmul kernels let a trailing partial M/N tile's subgroupMatrixLoad read past A/B. Tint's robustness transform resets such an out-of-bounds load to offset 0 with the minimum stride instead of zero-filling, so it silently returns a different valid tile and corrupts otherwise in-range data. Shift the trailing tile back so it always ends exactly at M/N, and skip the rows/columns the shift re-covers at write-out. The host now falls back to the generic path when the tile doesn't fit M/N at all, since the shift only works when it does. Regenerated the wgsl_template golden snapshots (in_tree_golden) to match the updated templates.
Jie Chen (jchen10)
force-pushed
the
sgmm_robust
branch
from
September 7, 2026 06:21
118c0b4 to
dab4e36
Compare
Compute N_b before dispatching the pad program so the tile-fit check can bail out early without wasting a pad dispatch.
Jiajia Qin (qjia7)
approved these changes
Sep 7, 2026
Hariharan Seshadri (hariharans29)
merged commit Sep 8, 2026
2a2833a
into
microsoft:main
97 of 100 checks passed
This was referenced Sep 10, 2026
Open
This was referenced Sep 14, 2026
Open
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.
The 8x16x16 gemm/matmul kernels let a trailing partial M/N tile's subgroupMatrixLoad read past A/B. Tint's robustness transform resets such an out-of-bounds load to offset 0 with the minimum stride instead of zero-filling, so it silently returns a different valid tile and corrupts otherwise in-range data.
Shift the trailing tile back so it always ends exactly at M/N, and skip the rows/columns the shift re-covers at write-out. The host now falls back to the generic path when the tile doesn't fit M/N at all, since the shift only works when it does.