[ExecuTorch][WebGPU] Add fp32 matmul ops (mm, bmm, linear) to the WebGPU backend#20917
[ExecuTorch][WebGPU] Add fp32 matmul ops (mm, bmm, linear) to the WebGPU backend#20917JCNTH wants to merge 3 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20917
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New Failure, 1 Unrelated Failure, 6 Unclassified FailuresAs of commit 7ce3554 with merge base 21554e5 ( NEW FAILURE - The following job has failed:
UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:
FLAKY - The following job failed but was likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
SS-JIA
left a comment
There was a problem hiding this comment.
Review automatically exported from Phabricator review in Meta.
Stack from ghstack (oldest at bottom):
Add fp32 GEMM handlers (
aten.mm,aten.bmm,aten.linear) — the dense matmuls the on-device training tail needs (tiled + vec4 WGSL).Key changes:
runtime/ops/{mm,bmm,linear}/— tiled + vec4 fp32 GEMM WGSL kernels + handlersCMakeLists.txtWEBGPU_SRCS— wire the three sourcesReuses the shared Vulkan partitioner (
aten.mm/bmm/linearalready have VulkanOpFeatures); this adds the WebGPU kernels only.Co-authored-with: Claude Code.
@exported-using-ghexport
Differential Revision: D111755135
Differential Revision: D111755135