[WebGPU] Intel SubgroupMatrix MatMul: batched-B support

[Jie Chen (jchen10)]

• Create the Intel subgroup-matrix impl lazily in ComputeInternal (via std::call_once) instead of PrePackInternal, which only fires for constant initializers and so skipped the subgroup-matrix path for runtime/dynamic B (batched matmul).
• Fall back to the generic MatMul path when N is odd: Intel's f16 subgroup-matrix B load requires an even row stride (4-byte-aligned K-rows); an odd N corrupts odd output columns.
• Add batched f16 MatMul tests (matmul_large_test.cc).

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Jie Chen (jchen10)]

Jiajia Qin (@qjia7) PTAL

[Copilot]

Pull request overview

This PR updates the WebGPU MatMul implementation to better utilize Intel’s subgroup-matrix path for batched (true-BMM) workloads by (1) initializing the optimized implementation lazily at first compute, (2) extending the subgroup-matrix kernel to dispatch batch slices on z, and (3) tuning the Intel tiling selector to account for batch-driven occupancy while falling back for unsupported odd N cases.

Changes:

• Lazily initialize the subgroup-matrix optimized MatMul path in ComputeInternal using std::call_once so it can apply even when B is not a constant initializer.
• Add batched-B support to the subgroup-matrix MatMul path by dispatching slices on z and applying per-slice A/B/output offsets in WGSL.
• Extend Intel tiling selection to incorporate batch as an occupancy multiplier and clamp split-K accordingly.

Reviewed changes

Copilot reviewed 6 out of 7 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
onnxruntime/test/providers/webgpu/matmul_large_test.cc	Adds disabled large tests covering true batched-B MatMul, unaligned shapes, and 4D batching.
onnxruntime/core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.cc	Adds batch-aware heuristic tiling and split-K clamping for batched dispatch.
onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.h	Updates the tiling-selector callback signature to include batch.
onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc	Implements true batched-B handling (z-dispatch), odd-N fallback, and passes batch into the tiling selector.
onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul_8x16x16.wgsl.template	Adds batch-id recovery from flattened workgroup_idx and per-slice offsets for A/B/output addressing.
onnxruntime/core/providers/webgpu/math/matmul.h	Adds std::once_flag and removes PrePackInternal override to support lazy optimized-path creation.
onnxruntime/core/providers/webgpu/math/matmul.cc	Performs lazy subgroup-matrix impl creation in ComputeInternal via std::call_once; removes PrePackInternal implementation.

[Jie Chen (jchen10)]

Jiajia Qin (@qjia7) Thanks for reviewing.

Hariharan Seshadri (@hariharans29) Could you please take a look, thanks!

[Hariharan Seshadri (hariharans29)]

Review: PR #29749 — [WebGPU] Intel SubgroupMatrix MatMul: batched-B support

Verdict — approve. Small, well-scoped fix that (1) unblocks the Intel subgroup-matrix MatMul path when B is a runtime tensor (batched matmul), (2) extends the kernel and tiling policy to handle batched B correctly via z-dispatch, (3) adds a real correctness guard for Intel's f16 subgroup-load 4-byte-alignment requirement, and (4) is defended by a targeted test suite that specifically pins the identical-batch-dims requirement so a future "product-only" refactor cannot silently mis-pair slices. Prior review by Jiajia Qin (@qjia7) was thorough and each point was addressed cleanly.

What the change is doing (and why it's the right shape)

• Lazy init in ComputeInternal instead of PrePackInternal. matmul.cc — PrePack fires only for constant initializers, so any MatMul with a non-constant B (batched bmm, dynamic weights) was permanently missing the subgroup-matrix path. Moving creation to ComputeInternal behind std::call_once fixes that without paying an init cost per Compute. mutable std::once_flag impl_init_flag_; on the shared kernel plus the once-guarded creation is the right primitive here — the kernel object is shared across concurrent Compute calls on the same session.
• Batched-B shape enumeration in subgroup_matrix_matmul.cc is disciplined about what it accepts:
  • Shared 2D weight B [K, N]: fold leading A dims into M, batch = 1. This is the pre-existing case; no regression.
  • Batched B [..., K, N]: same rank as A, identical leading dims (no broadcasting), each (A_slice_i, B_slice_i) pair becomes one z-slice. The kernel pairs slice-i-with-slice-i, so the identical-dims requirement is a real correctness precondition, not a "we'd rather not deal with broadcasting."
  • Anything else (rank mismatch, non-identical leading dims, odd N) falls back to the generic MatMul path.
• Odd-N fallback with a concrete reason. The comment at subgroup_matrix_matmul.cc explains: Intel's f16 subgroupMatrixLoad for the right operand reads columns in 32-bit (2×f16) pairs, so each K-row must start 4-byte aligned. An odd N shifts every other K-row by 2 bytes and corrupts odd output columns. That's a specific, testable claim, and the test DISABLED_BatchedB_Unaligned covers N = 1023, 33 shapes to exercise the fallback while N = 130 (even, but odd-M) exercises the kept-on-the-fast-path variant. Nice separation.
• WGSL kernel batching in subgroup_matrix_matmul_8x16x16.wgsl.template is straightforward:

  let num_m_tile = (uniforms.M + kTileM - 1u) / kTileM;
  let tiles_per_slice = uniforms.num_n_tile * num_m_tile;
  let batch_id = workgroup_idx / tiles_per_slice;
  let slice_idx = workgroup_idx % tiles_per_slice;

  The comment explicitly notes that workgroup_idx is the true linear index — invariant under x/y/z normalization — and that batch slices are contiguous blocks of tiles_per_slice tiles. Per-slice offsets are batch_id * M * K, batch_id * K * N, batch_id * M * N; for the shared-weight case batch = 1 collapses b_batch_offset to zero, so the 2D path stays byte-for-byte the same load pattern it was. Bit-identity for the pre-existing 2D case is preserved by construction.
• Tiling selector's batch parameter is a pure occupancy multiplier in intel/subgroup_matrix_tiling_selector.cc:
  • HeuristicTiling multiplies tile_count(M) * tile_count(N) by batch when checking against hw — larger batch fills the machine at bigger tiles.
  • ClampSplitKForBatch retires split-K factors from the pretuned table (which was tuned at batch = 1) when batch * tile_count_M * tile_count_N * split_k > 2 * hw. Retiring by halving is fine because the pretuned split_k values are powers-of-two.
  • EffectiveHwSubgroups extraction is a small clean-up that removes the duplicated "if HwSubgroups returns 0, use 256" fallback.

Correctness observations

1. std::call_once captures context by reference in the lambda but does not persist it. matmul.cc:

   std::call_once(impl_init_flag_, [&]() {
     impl_ = CreateSubgroupMatrixMatMulImpl(*this, context);
   });

   The lambda captures context by reference, but context is a stack local of the current ComputeInternal call. As long as CreateSubgroupMatrixMatMulImpl does not stash the reference into the returned MatMulOptImpl (it currently queries adapter info at each Compute invocation via the tiling selector callback that receives the per-call context), this is safe — and it matches the pre-existing PrePackInternal contract. Worth a quick manual re-check that no path inside CreateSubgroupMatrixMatMulImpl stores anything from context beyond the lambda's lifetime. My read of the diff says no, but it's the kind of thing that regresses silently.

2. The ORT_ENFORCE(narrow<uint32_t>(b_shape[0]) == K, ...) in the shared-2D branch is a behavior change — the pre-PR code silently return Status::OK() (fallback) when B's K didn't match A's K. Now it hard-fails. Per Jiajia Qin (@qjia7)'s discussion and the author's response, this case is essentially never reachable on a valid MatMul (ONNX validates it), so ORT_ENFORCE is appropriate. Just calling it out as a behavior delta for anyone auditing the diff: an invalid shape now throws where it used to silently fall through to the generic path (which would also have thrown, just further downstream).

3. The for (size_t i = 0; i + 2 < rank; ++i) batch-dim comparison correctly avoids size_t underflow when rank < 2 — good defensive style.

4. The batch = narrow<uint32_t>(a_shape.SizeToDimension(rank - 2)) fits into uint32_t for any reasonable ONNX shape. narrow<> throws on overflow, which is the right failure mode.

5. The WGSL contract "batch slices are contiguous blocks of tiles_per_slice tiles" under workgroup normalization depends on the ORT WebGPU EP normalizing the (x, y, z) dispatch such that z-index maps to the slowest-varying stride of the flattened workgroup_idx. This is a pre-existing invariant used elsewhere in the EP but worth pinning down with a comment in the host code too (currently the invariant is spelled out only in the WGSL template comment). Not blocking.

6. The output shape derivation for batched B — because the subgroup path enters only when a_shape and b_shape have identical rank and identical leading dims, "copy A's shape and swap the last dim to N" is exactly correct. The DISABLED_BatchedB_BroadcastEqualProduct test (A=[2,1,M,K] × B=[1,2,K,N], A=[1,4,M,K] × B=[4,1,K,N]) is the key regression guard: it verifies that a future refactor cannot naively use "batch product" instead of "batch-dim equality" without breaking correctness. This is exactly the right test to have.

Non-blocking observations

1. std::call_once performance implication for kernels that will never have a subgroup path (e.g., non-Intel, non-AVX-capable devices). Every ComputeInternal invocation on such a device still pays for one atomic-check on impl_init_flag_ before hitting the if (impl_) guard. Uncontended call_once is a single relaxed atomic load on all major implementations, so this is well below noise, but worth mentioning if MatMul becomes a critical hot path on non-subgroup devices. An alternative is a hand-rolled std::atomic<bool> sentinel, but call_once is fine.

2. Tests are all DISABLED_. Consistent with the existing convention in matmul_large_test.cc — these are opt-in slow tests requiring a real WebGPU device with the appropriate subgroup extensions. Please confirm they've actually been run locally on Intel hardware (the PR description says "Add batched f16 MatMul tests" but doesn't state whether they were executed). Given the guard-specific nature of DISABLED_BatchedB_BroadcastEqualProduct and DISABLED_BatchedB_Unaligned, a "yes, I ran these and they pass" note in a follow-up comment would tie a bow on this.

3. The batched-B path relies on the caller (WebGpuKernel infra or MatMul::ComputeInternal) having already allocated the output tensor with a shape that matches A's shape with the last dim replaced by N. This is a pre-existing convention in the EP and is spelled out in the diff comment, but a small assertion inside the batched branch — e.g., ORT_ENFORCE(output->Shape().SizeToDimension(rank - 2) == batch) — would harden the kernel against a future output-shape derivation change. Minor.

4. Naming. EffectiveHwSubgroups reads well; ClampSplitKForBatch is descriptive. tiles_per_slice in the WGSL is a nice name. No complaints.

5. The comment "batch is the number of z-dispatched slices (1 for a shared 2D weight)" in subgroup_matrix_matmul.h is exactly the right level of docstring for the selector callback. Good.

CI

Head ee2794a at 1/1 latest check green; the substantive-content commit 312ead7 was 83/86. Please confirm that the 3 non-green in 312ead7 are known flakes (not related to this PR) — a quick look at what those 3 were would close the loop. No obvious concern from the diff.

Bottom line

Approve. The design (lazy call_once init, identical-batch-dims requirement instead of product-based, N-alignment fallback, batch as pure occupancy multiplier) is right, the WGSL changes are minimal and preserve the 2D fast path exactly, and the test suite is specifically shaped to catch the correctness landmines this PR could have introduced. The two things I'd like from the author in follow-up:

• A one-line confirmation that CreateSubgroupMatrixMatMulImpl does not stash the ComputeContext& reference beyond the lambda's lifetime (i.e., the once-init lambda's context capture is truly transient).
• A one-line note that the new DISABLED_BatchedB* tests were run locally on an Intel GPU and pass.

Neither is blocking.

[Hariharan Seshadri (hariharans29)]

Approving as Jiajia Qin (@qjia7) is OOF and her comments have been resolved.

[Jie Chen (jchen10)]

Thanks, Hariharan Seshadri (@hariharans29)! Just rebased and squashed the commits.

[Jie Chen (jchen10)]

The failing checks look like a CI infra issue.
C:\Program Files\Microsoft Visual Studio\2022\Enterprise\MSBuild\Microsoft\VC\v170\Microsoft.CppBuild.targets(511,5): error MSB8040: Spectre-mitigated libraries are required for this project. Install them from the Visual Studio installer (Individual components tab) for any toolsets and architectures being used. Learn more: https://aka.ms/Ofhn4c [D:\a\_work\onnxruntime\onnxruntime\RelWithDebInfo\CMakeFiles\CMakeScratch\TryCompile-o1i4cl\cmTC_44c59.vcxproj] Done Building Project "D:\a\_work\onnxruntime\onnxruntime\RelWithDebInfo\CMakeFiles\CMakeScratch\TryCompile-o1i4cl\cmTC_44c59.vcxproj" (default targets) -- FAILED.

[Hariharan Seshadri (hariharans29)]

The failing checks look like a CI infra issue. C:\Program Files\Microsoft Visual Studio\2022\Enterprise\MSBuild\Microsoft\VC\v170\Microsoft.CppBuild.targets(511,5): error MSB8040: Spectre-mitigated libraries are required for this project. Install them from the Visual Studio installer (Individual components tab) for any toolsets and architectures being used. Learn more: https://aka.ms/Ofhn4c [D:\a\_work\onnxruntime\onnxruntime\RelWithDebInfo\CMakeFiles\CMakeScratch\TryCompile-o1i4cl\cmTC_44c59.vcxproj] Done Building Project "D:\a\_work\onnxruntime\onnxruntime\RelWithDebInfo\CMakeFiles\CMakeScratch\TryCompile-o1i4cl\cmTC_44c59.vcxproj" (default targets) -- FAILED.

Yeah we have some CI issues. I ll follow up on this PR and merge once they are resolved.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.