[ExecuTorch][WebGPU] Add optimized addmm op#20855
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20855
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (2 Unrelated Failures)As of commit 30a4b34 with merge base aceeb40 ( FLAKY - The following job failed but was likely due to flakiness present on trunk:
BROKEN TRUNK - The following job failed but was present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
meta-cla
Bot
added
the
CLA Signed
2 participants
Stack from ghstack (oldest at bottom):
fp32
aten.addmm.defaultnow runs on the WebGPU backend as a shared-memory-tiled GEMM, so the fused bias-plus-matmul that HuggingFaceLinearlowers to executes on-device. In the Florence-2 BART and DaViT graphs the HFnn.Linearlayers lower toaten.addmm(notaten.linear), making addmm the dominant dense GEMM in those stacks.Problem — The WebGPU backend had no
aten.addmm.defaultkernel, so any graph whose linears lowered through addmm delegated at export but failed at runtime load. As with plain matmul, a naive per-output kernel would re-stream both operands from global memory with no reuse.Solution — Before:
aten.addmm.defaulthad no runtime kernel (load-time failure). After: the handler bindsself(bias),mat1 [M,K],mat2 [K,N], and the output, and dispatches a tiled GEMM computingout = beta*self + alpha*(mat1 @ mat2). Each workgroup cooperatively stages 32x32 slabs ofmat1andmat2into workgroup shared memory,workgroupBarrier-syncs, and accumulates acrossceil(K/32)k-tiles so the coalesced tile load is amortized over the whole output tile. Thebeta*self + alpha*accepilogue broadcastsselffrom either[N](the common HF bias case) or a full[M,N].Implementation — Shares the tiled-GEMM approach with
linear_fp32:TILE = 32,RPT = 4,@workgroup_size(8, 8, 1), twovar<workgroup>a_sub/b_subarrays, each thread accumulating a 4x4 register sub-tile overceil(K/32)tiles. It re-derives the B-side read format2's actual[K,N]layout:read_breadst_mat2[krow*N + col](N contiguous), unlikelinear's transposed[N,K]weight read. No vec4 variant is provided: withmat2 [K,N]the N dimension is contiguous, so a vec4-over-K view would require a strided gather on the mat2 side that erodes the benefit, and the cooperative shared-memory tile load already closes the coalescing gap. Epilogue selectst_self[r*N + c]whenself_2d, else the broadcastt_self[c], and writesbeta*self_val + alpha*acc[ir][ic].beta/alphaare read with the sharedutils::scalar_or(Double/Int/Null-permissive, default1.0). Dispatch is a 2D gridceil(N/32) x ceil(M/32), throwing if either dimension exceeds 65535. Uses the sharedutils::make_compute_pipelineandutils::make_uniform(32-byteAddmmParams {M,N,K,self_2d,beta,alpha,_pad[2]}). Validates non-null buffers, 2Dmat1/mat2/out,mat2shape== [K,N], andselfbroadcastable from[N]or[M,N].Constraints — fp32 only (output byte-size check);
mat1/mat2/outmust be 2D withmat2exactly[K,N];selfmust be[N]or[M,N]; scalar-only tiled path (no vec4); 2D dispatch capped at 65535 workgroups per dimension; fixed@workgroup_size(8, 8, 1).Co-authored-with: Claude Code.
Differential Revision: D110836673