webgpu: using subgroup matrix to optimize Conv operator by xhcao · Pull Request #32252 · microsoft/onnxruntime

xhcao · 2026-08-25T01:27:52Z

Description

Adds a subgroup-matrix Conv path for the WebGPU EP that computes the convolution as
an implicit GEMM. Covers F16, channels-last, group == 1 convs on devices reporting
an 8x16x16 subgroup-matrix config, and falls back to the existing Conv paths
otherwise.

Motivation and Context

Conv dominates runtime in vision models, but the existing WebGPU Conv paths do not
use the subgroup-matrix units that the MatMul path already exploits on Intel
Xe2/Xe3. This brings the same acceleration to Conv.

azure-pipelines · 2026-08-25T01:28:00Z

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

Jie Chen (jchen10) · 2026-09-17T01:39:49Z

/run bots

Jiajia Qin (qjia7) · 2026-09-17T02:06:33Z

The /run bots request was received from Jie Chen (@jchen10), but no CI run was started. No pending action_required workflow runs matched the current PR head. This pull request currently has merge conflicts; please resolve them before requesting CI again.

xhcao · 2026-09-21T02:28:06Z

Jie Chen (@jchen10) I had also fixed out-of-bound issue, you could start to review the PR, thanks.

azure-pipelines · 2026-10-16T02:53:25Z

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

Copilot

🟡 Changes recommended

Segmented-weight loads and FP16 partial-sum overflow can produce incorrect convolution results.

2 open findings

What changed in this PR

Adds subgroup-matrix acceleration for eligible FP16, channels-last convolutions in the WebGPU execution provider, retaining existing fallback paths.

Changes:

Adds an implicit-GEMM shader with bias and fused activations.
Integrates capability checks, weight caching, and Intel-specific tiling.
Adds convolution correctness tests across shapes and activations.

File	Description
onnxruntime/test/providers/webgpu/subgroup_matrix_conv_test.cc	Adds convolution correctness coverage.
onnxruntime/core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.h	Declares Conv selector and shared scratch budget.
onnxruntime/core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.cc	Adds Conv-specific tiling and safety clamps.
onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h	Declares optimized Conv interfaces.
onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc	Implements eligibility, weight caching, and dispatch.
onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv_8x16x16.wgsl.template	Implements subgroup-matrix convolution.
onnxruntime/core/providers/webgpu/nn/conv.h	Adds optimized implementation state.
onnxruntime/core/providers/webgpu/nn/conv.cc	Integrates optimized dispatch and prepacking decisions.

🧠 Review effort: Balanced

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

+    return false;
+  }
+
+  const std::optional<SubgroupMatrixConfig> config = GetSupportedConfig(context);


+    let a_scratch_base = sg_index * kTileM * kSgMatK;  // this subgroup's A staging slot
+
+    // Accumulators: kSgMatCountM M blocks x kSgMatCountN N blocks (sg_mat_c<mi>_<ni>).
+    var sg_mat_c0_0: subgroup_matrix_result<f16, sg_mat_n, sg_mat_m>;


xhcao · 2026-10-19T08:31:28Z

There were concerns regarding the Tuned Tree on the special devices, so I had removed the Conv Tuned Tree, now the PR only has Heuristic Tiling function. If there is a better way, we could add the Conv Tuned Tree in another PR.
Below table is performance comparation among Exact decision tree, Opitmized decision tree and Heuristic tiling + win + vulkan, Heuristic tiling + win + D3D.

models	Exact decision tree (%)	Opitmized decision tree(%)	Heuristic tiling(%)+ win + vulkan	Heuristic tiling(%)+ win + D3D
depth-anything-base-fp16-pow-fp32	-6	-5	-2	-1.5
detr-resnet-50-fp16	-34	-28	-25	-18
modnet-fp16	-38	-40	-33	-28
resnet50-v1-f16-demo	-50	-50	-51	-35
sd-turbo-unet-fp16-demo-layernorm	-30	-11	-6	-12
sd-turbo-vae-decoder-fp16-demo	-18	-10	-3	-4
sdunet-v1.5-demo-layernorm	-19	-8	-24	-12
yolov8m-fp16	-41	-41	-29	-30
yolov8n-fp16	-39	-32	-28	-26
yolov8n-pose-fp16	-30	-28	-24	-32
yolov8x-pose-fp16	-30	-30	-20	-20

Jie Chen (jchen10) · 2026-10-20T01:13:16Z

+    let channel_i_vec = uniforms.channel_i / kVecSize;
+    let oh = m / uniforms.output_w;
+    let ow = m % uniforms.output_w;
+    let cin_vec = k_vec % channel_i_vec;


Use cin or channel_i naming consistently.

Do you mean the same name as uniforms.channel_i? But channel_i_vec is related with kVecSize in line 99.

How do you like this way:

let cin_vec = uniforms.cin / kVecSize; let oh = m / uniforms.output_w; let ow = m % uniforms.output_w; let k_cin = k_vec % cin_vec; let k_hw = k_vec / cin_vec;

Jie Chen (jchen10) · 2026-10-20T01:27:37Z

+    if (ih < 0 || ih >= i32(uniforms.src_h) || iw < 0 || iw >= i32(uniforms.src_w)) {
+        return input_src_value_t();
+    }
+    let src_vec_idx = src_vec_batch_offset +
+                      (u32(ih) * uniforms.src_w + u32(iw)) * channel_i_vec + cin_vec;


Similarly, use input_w/h and output_w/h, or src_w/h and dst_w/h. Don't mix them, like src_w, output_w.

Jie Chen (jchen10) · 2026-10-20T02:03:49Z

+            let k_vec_base = kb * kSgMatK / kVecSize;
+            let k_vec_per_row = kSgMatK / kVecSize;


let k_vec_per_sgmat_row = kSgMatK / kVecSize; let k_vec_base = kb * k_vec_per_sgmat_row;

Jie Chen (jchen10) · 2026-10-20T06:17:55Z

+        return;
+    }
+    let slice_idx = workgroup_idx % tiles_per_slice;    // tile index within the slice
+    let n_tile = slice_idx % uniforms.num_n_tile;       // which kTileN-wide column tile


tiles_per_slice is actually tiles_per_batch, right? So, slice_idx should be tile_idx. It seems there is no need to introduce a distracting concept of slice.

Jie Chen (jchen10) · 2026-10-20T06:33:28Z

+    let src_vec_batch_offset = batch_id * uniforms.src_h * uniforms.src_w * (uniforms.channel_i / kVecSize);
+    let out_batch_offset = batch_id * uniforms.M * uniforms.N;


base, start, and offset are used for naming a bit too casually. we'd better cleanup a little bit. Personally I am in favor in_batch_base_vec over src_vec_batch_offset, and out_batch_base over out_batch_offset.

Jie Chen (jchen10) · 2026-10-20T11:32:11Z

+    if (parent.PrepackedKernel() != nullptr) {
+      return Status::OK();
+    }


Use ORT_ENFORCE instead?

Cannot use ORT_ENFORCE here, I update the code comments

Jie Chen (jchen10) · 2026-10-20T11:45:43Z

+    // Safety net: the Conv selector already enforces the Conv scratch budget
+    // (output tiles + im2col A staging), so this clamp normally does nothing; it
+    // guards against a stray oversized split-K from the selector. Only split_k is
+    // clamped, not the tile area: ConvHeuristicTiling always picks the smallest
+    // tile_m (8), so the largest candidate tile needs 8 * (64 + 16) = 640 elements
+    // at split_k 1, or 8 * 8 * 80 = 5120 at the largest split_k (8), well inside
+    // the budget.
+    uint32_t split_k = tiling->split_k;
+    while (split_k > 1 &&
+           split_k * tile_m * (tile_n + config.K) > intel::kMaxScratchElems) {
+      split_k /= 2;
+    }


As the comment says, this logic is unnecessary. Simply remove it or add an ORT_ENFORCE.

Jie Chen (jchen10) · 2026-10-20T12:19:27Z

+// Combined workgroup-memory (SLM) budget in f16 elements for the two scratch
+// buffers the conv kernel uses: the split-K output tiles
+// (split_k * tile_m * tile_n) plus the per-subgroup im2col A staging
+// (split_k * tile_m * sg_mat_k). split_k is clamped to fit this. Shares
+// intel::kMaxScratchElems rather than hardcoding a second copy of the same
+// budget the selector already sizes tilings against.


The comment seems irrelevant to the code lines nearby.

Jie Chen (jchen10) · 2026-10-20T12:25:11Z

+// Largest per-tile subgroup-matrix counts the WGSL template unrolls: it generates the
+// accumulators sg_mat_c<mi>_<ni> for mi < 8 and ni < 4. A tiling above these would be
+// accepted by the selector but silently mis-computed -- the kernel would accumulate
+// only the blocks it generated while the write-out still covers the whole tile. The
+// tile candidates the selector picks from stay inside this, so this is enforced as a
+// safety net rather than something expected to fire.
+constexpr uint32_t kMaxSgMatCountM = 8;
+constexpr uint32_t kMaxSgMatCountN = 4;


The selector should ensure this, and we don't need to bother to check it again here?

If a future change to the tiling selector ever violated this, the WGSL template would silently under-accumulate (it only unrolls ni < 4 accumulators) instead of crashing — a wrong answer with no visible error.

The guard itself somewhat makes sense, but should we check it here? The values are specific to the 8x16x16 WGSL template. Perhaps we'd better put it somewhere template specific. Not mandatory in this PR. We can do it later.

xhcao

Jie, please take a look again, thanks

xhcao · 2026-10-21T02:45:55Z

+    if (parent.PrepackedKernel() != nullptr) {
+      return Status::OK();
+    }


Cannot use ORT_ENFORCE here, I update the code comments

xhcao · 2026-10-21T02:46:35Z

+// Largest per-tile subgroup-matrix counts the WGSL template unrolls: it generates the
+// accumulators sg_mat_c<mi>_<ni> for mi < 8 and ni < 4. A tiling above these would be
+// accepted by the selector but silently mis-computed -- the kernel would accumulate
+// only the blocks it generated while the write-out still covers the whole tile. The
+// tile candidates the selector picks from stay inside this, so this is enforced as a
+// safety net rather than something expected to fire.
+constexpr uint32_t kMaxSgMatCountM = 8;
+constexpr uint32_t kMaxSgMatCountN = 4;


If a future change to the tiling selector ever violated this, the WGSL template would silently under-accumulate (it only unrolls ni < 4 accumulators) instead of crashing — a wrong answer with no visible error.

xhcao · 2026-10-21T02:59:35Z

+    // Safety net: the Conv selector already enforces the Conv scratch budget
+    // (output tiles + im2col A staging), so this clamp normally does nothing; it
+    // guards against a stray oversized split-K from the selector. Only split_k is
+    // clamped, not the tile area: ConvHeuristicTiling always picks the smallest
+    // tile_m (8), so the largest candidate tile needs 8 * (64 + 16) = 640 elements
+    // at split_k 1, or 8 * 8 * 80 = 5120 at the largest split_k (8), well inside
+    // the budget.
+    uint32_t split_k = tiling->split_k;
+    while (split_k > 1 &&
+           split_k * tile_m * (tile_n + config.K) > intel::kMaxScratchElems) {
+      split_k /= 2;
+    }


xhcao · 2026-10-21T06:05:20Z

+// Combined workgroup-memory (SLM) budget in f16 elements for the two scratch
+// buffers the conv kernel uses: the split-K output tiles
+// (split_k * tile_m * tile_n) plus the per-subgroup im2col A staging
+// (split_k * tile_m * sg_mat_k). split_k is clamped to fit this. Shares
+// intel::kMaxScratchElems rather than hardcoding a second copy of the same
+// budget the selector already sizes tilings against.


xhcao · 2026-10-21T06:08:49Z

+    let channel_i_vec = uniforms.channel_i / kVecSize;
+    let oh = m / uniforms.output_w;
+    let ow = m % uniforms.output_w;
+    let cin_vec = k_vec % channel_i_vec;


Do you mean the same name as uniforms.channel_i? But channel_i_vec is related with kVecSize in line 99.

xhcao · 2026-10-21T06:19:51Z

+    if (ih < 0 || ih >= i32(uniforms.src_h) || iw < 0 || iw >= i32(uniforms.src_w)) {
+        return input_src_value_t();
+    }
+    let src_vec_idx = src_vec_batch_offset +
+                      (u32(ih) * uniforms.src_w + u32(iw)) * channel_i_vec + cin_vec;


xhcao · 2026-10-21T06:34:18Z

+        return;
+    }
+    let slice_idx = workgroup_idx % tiles_per_slice;    // tile index within the slice
+    let n_tile = slice_idx % uniforms.num_n_tile;       // which kTileN-wide column tile


xhcao · 2026-10-21T06:37:34Z

+    let src_vec_batch_offset = batch_id * uniforms.src_h * uniforms.src_w * (uniforms.channel_i / kVecSize);
+    let out_batch_offset = batch_id * uniforms.M * uniforms.N;


xhcao · 2026-10-21T06:42:20Z

+            let k_vec_base = kb * kSgMatK / kVecSize;
+            let k_vec_per_row = kSgMatK / kVecSize;


Jie Chen (jchen10) · 2026-10-21T11:35:46Z

+constexpr uint32_t kMatMulScratchNPad = 0;
+constexpr uint32_t kConvScratchNPad = kSubgroupMatrixK;


ScratchNPad doesn't look like a straightforward and informative name. May ask AI to suggest a better one.

Jie Chen (jchen10) · 2026-10-21T11:39:21Z

-        // This tile fills the machine on its own. Among all such tiles, keep the
-        // largest-area one (best operand reuse); this also supersedes any
-        // under-filling pick made earlier.
        if (!filled || area > best_area) {
          filled = true;
          best_area = area;
          tile_m = tm;
          tile_n = tn;
        }
      } else if (!filled && (tiles > best_tiles || (tiles == best_tiles && area > best_area))) {
-        // No tile has filled the machine yet: chase occupancy by maximizing the
-        // number of independent tiles, breaking ties toward larger area.


Keep the original comments as the corresponding code logic is not easy to read.

Jie Chen (jchen10) · 2026-10-21T11:40:38Z

-  // Add split-K only while the independent tiles under-fill the machine; cap the
-  // total at 2x hardware and keep enough K work per subgroup and scratch budget.


Keep the comment as well

Jie Chen (jchen10) · 2026-10-21T11:57:28Z

+// Halves split_k until it fits K's block count and the scratch budget. No-op today
+// (ConvHeuristicTiling already guarantees both); a safety net for future callers.
+void ClampSplitKToProblem(SubgroupMatrixTiling& t, uint32_t K, uint32_t scratch_n_pad) {
+  const uint32_t k_blocks = K / kSubgroupMatrixK;
+  while (t.split_k > 1 && (t.split_k > k_blocks || ScratchElems(t, scratch_n_pad) > kMaxScratchElems)) {
+    t.split_k /= 2;
+  }
+}


Remove the unused logic to be less distracting.

Jie Chen (jchen10) · 2026-10-21T12:16:07Z

+    if (N < tiling->tile_n) {
+      return Status::OK();
+    }


Use ORT_ENFORCE?

No, decline means that subgroup matrix path cannot handle the current Conv node, let the others path to handle it. It does not need to throw an error here.

Jie Chen (jchen10) · 2026-10-21T12:40:42Z

+// Largest per-tile subgroup-matrix counts the WGSL template unrolls: it generates the
+// accumulators sg_mat_c<mi>_<ni> for mi < 8 and ni < 4. A tiling above these would be
+// accepted by the selector but silently mis-computed -- the kernel would accumulate
+// only the blocks it generated while the write-out still covers the whole tile. The
+// tile candidates the selector picks from stay inside this, so this is enforced as a
+// safety net rather than something expected to fire.
+constexpr uint32_t kMaxSgMatCountM = 8;
+constexpr uint32_t kMaxSgMatCountN = 4;


The guard itself somewhat makes sense, but should we check it here? The values are specific to the 8x16x16 WGSL template. Perhaps we'd better put it somewhere template specific. Not mandatory in this PR. We can do it later.

Jie Chen (jchen10) · 2026-10-21T12:58:34Z

+    let channel_i_vec = uniforms.channel_i / kVecSize;
+    let oh = m / uniforms.output_w;
+    let ow = m % uniforms.output_w;
+    let cin_vec = k_vec % channel_i_vec;


How do you like this way:

let cin_vec = uniforms.cin / kVecSize; let oh = m / uniforms.output_w; let ow = m % uniforms.output_w; let k_cin = k_vec % cin_vec; let k_hw = k_vec / cin_vec;

xhcao

Jie, PTAL.

xhcao · 2026-10-22T02:26:36Z

+    if (N < tiling->tile_n) {
+      return Status::OK();
+    }


No, decline means that subgroup matrix path cannot handle the current Conv node, let the others path to handle it. It does not need to throw an error here.

xhcao · 2026-10-22T02:41:21Z

+    let channel_i_vec = uniforms.channel_i / kVecSize;
+    let oh = m / uniforms.output_w;
+    let ow = m % uniforms.output_w;
+    let cin_vec = k_vec % channel_i_vec;


xhcao · 2026-10-22T02:54:14Z

+constexpr uint32_t kMatMulScratchNPad = 0;
+constexpr uint32_t kConvScratchNPad = kSubgroupMatrixK;


xhcao · 2026-10-22T02:56:30Z

-        // This tile fills the machine on its own. Among all such tiles, keep the
-        // largest-area one (best operand reuse); this also supersedes any
-        // under-filling pick made earlier.
        if (!filled || area > best_area) {
          filled = true;
          best_area = area;
          tile_m = tm;
          tile_n = tn;
        }
      } else if (!filled && (tiles > best_tiles || (tiles == best_tiles && area > best_area))) {
-        // No tile has filled the machine yet: chase occupancy by maximizing the
-        // number of independent tiles, breaking ties toward larger area.


xhcao · 2026-10-22T02:57:50Z

-  // Add split-K only while the independent tiles under-fill the machine; cap the
-  // total at 2x hardware and keep enough K work per subgroup and scratch budget.


xhcao · 2026-10-22T03:04:47Z

+// Halves split_k until it fits K's block count and the scratch budget. No-op today
+// (ConvHeuristicTiling already guarantees both); a safety net for future callers.
+void ClampSplitKToProblem(SubgroupMatrixTiling& t, uint32_t K, uint32_t scratch_n_pad) {
+  const uint32_t k_blocks = K / kSubgroupMatrixK;
+  while (t.split_k > 1 && (t.split_k > k_blocks || ScratchElems(t, scratch_n_pad) > kMaxScratchElems)) {
+    t.split_k /= 2;
+  }
+}


Jie Chen (jchen10) · 2026-10-22T05:29:33Z

LGTM, thanks!

Jiajia Qin (@qjia7) PTAL

Jie Chen (jchen10) · 2026-10-26T01:45:56Z

/run bots

Jiajia Qin (qjia7) · 2026-10-26T01:52:15Z

The request was received. Approve workflow runs. Thanks.

xhcao · 2026-10-26T02:18:23Z

/run bots

Jiajia Qin (qjia7) · 2026-10-26T02:22:16Z

The request was received. Approve workflow runs. Thanks.

xhcao · 2026-10-26T07:53:22Z

/run failed bots

Jiajia Qin (qjia7) · 2026-10-26T08:01:21Z

The request was received. Failed workflow jobs were queued for rerun. Thanks.

Jiajia Qin (qjia7) · 2026-10-26T07:23:33Z

+// trailing partial N tile back to N - tile_n, which needs tile_n <= N, and N
+// itself need not be a multiple of kSubgroupMatrixN. Callers must have already
+// rejected N < kSubgroupMatrixN.
+void ClampTileNToProblem(SubgroupMatrixTiling& t, uint32_t N) {


How about renaming to CapTileNToAlignedN?

aligned_N = (N / subgroup_N) * subgroup_N; tile_n = min(tile_n, aligned_N);

Jiajia Qin (qjia7) · 2026-10-26T08:05:55Z

+    auto transform_dim = [](int64_t dim) { return static_cast<int32_t>(dim); };
+    std::transform(local_pads.begin(), local_pads.end(), std::back_inserter(pads), transform_dim);
+    std::transform(local_strides.begin(), local_strides.end(), std::back_inserter(strides), transform_dim);
+    std::transform(local_dilations.begin(), local_dilations.end(), std::back_inserter(dilations), transform_dim);


Could we extract lines 116–147 into a shared Conv geometry helper and reuse it in conv.cc? Both paths duplicate attribute defaulting and padding/output-shape inference. Keeping that logic shared would prevent differences in auto_pad handling. Please leave output allocation outside the helper so the optimized path can still decline safely.

Or Could we move the subgroup matrix Conv attempt after the common geometry calculation in conv.cc and pass in the resolved geometry? That would avoid duplicating shape/padding inference. We should keep the attempt before output allocation and path-specific transformations so fallback can reuse the same geometry safely.

Done. The first option was adopted.

Jiajia Qin (qjia7) · 2026-10-26T08:25:53Z

+    // this layout) -- blocked on PrePackInternal committing from the weight alone
+    // while Compute can still decline at runtime onto the im2col path.
+    Tensor run_local_ohwi_weight;
+    const Tensor* ohwi_weight = nullptr;


Could we reuse the existing OHWI prepacking for constant weights instead of introducing a separate lazy transpose cache? Both paths consume the same layout. We would need to let subgroup Conv accept the packed weights and ensure runtime fallback remains valid after the original initializer is released. Non-constant weights would still need a per-run transpose.

Done. Now SubgroupMatrixConv supported OHWI and HWIO layout, and removed the related code in original initializer.

Jiajia Qin (qjia7) · 2026-10-26T08:44:56Z

+    return false;
+  }
+
+  LOGS_DEFAULT(VERBOSE) << "SubgroupMatrixConv ACCEPTED: kernel_shape[OIHW]=" << kernel_shape.ToString()


Could we rename “ACCEPTED” to “initial eligibility checks passed”? Compute() can still decline afterward. Please put the actual dispatch message immediately before RunProgram() so it identifies the selected implementation reliably.

xhcao

Jiajia Qin (@qjia7) Please take a look, again. Thanks.

xhcao · 2026-10-27T05:10:46Z

+// trailing partial N tile back to N - tile_n, which needs tile_n <= N, and N
+// itself need not be a multiple of kSubgroupMatrixN. Callers must have already
+// rejected N < kSubgroupMatrixN.
+void ClampTileNToProblem(SubgroupMatrixTiling& t, uint32_t N) {


xhcao · 2026-10-27T05:13:24Z

+    auto transform_dim = [](int64_t dim) { return static_cast<int32_t>(dim); };
+    std::transform(local_pads.begin(), local_pads.end(), std::back_inserter(pads), transform_dim);
+    std::transform(local_strides.begin(), local_strides.end(), std::back_inserter(strides), transform_dim);
+    std::transform(local_dilations.begin(), local_dilations.end(), std::back_inserter(dilations), transform_dim);


Done. The first option was adopted.

xhcao · 2026-10-27T05:19:18Z

+    // this layout) -- blocked on PrePackInternal committing from the weight alone
+    // while Compute can still decline at runtime onto the im2col path.
+    Tensor run_local_ohwi_weight;
+    const Tensor* ohwi_weight = nullptr;


Done. Now SubgroupMatrixConv supported OHWI and HWIO layout, and removed the related code in original initializer.

xhcao · 2026-10-27T05:19:43Z

+    return false;
+  }
+
+  LOGS_DEFAULT(VERBOSE) << "SubgroupMatrixConv ACCEPTED: kernel_shape[OIHW]=" << kernel_shape.ToString()


xhcao · 2026-10-27T05:23:57Z

/run bots

Jiajia Qin (qjia7) · 2026-10-27T05:32:17Z

The request was received. Approve workflow runs. Thanks.

Jiajia Qin (qjia7) · 2026-10-27T07:57:57Z

+            let b_base = kb * kSgMatK * uniforms.N + global_base_n;
+            var sg_mat_b0: subgroup_matrix_right<f16, sg_mat_n, sg_mat_k> =
+                subgroupMatrixLoad<subgroup_matrix_right<f16, sg_mat_n, sg_mat_k>, row_major>(
+                    &input_b, b_base + 0 * kSgMatN, uniforms.N);


C2 [P1, new]: Guard odd-N HWIO weights before row-major subgroup-matrix loads

Location: onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv_8x16x16.wgsl.template:233-238; layout selection at subgroup_matrix_conv.cc:164-173.

The new shared-prepack route accepts HWIO weights and loads them as a row-major FP16 matrix with stride uniforms.N. On a supported Intel D3D12 adapter, take a constant-weight FusedConv with HardSwish, auto_pad=NOTSET, X NCHW [1,16,10,10], and W OIHW [17,16,3,3]. HardSwish is excluded from Im2ColMatMul, so PrePackInternal packs HWIO, while subgroup Conv still accepts K=144 and N=17. The new branch passes a 34-byte row stride; its shifted second N tile also starts at byte offset 2.

The DirectX matrix-load contract requires Wave-scope device-memory load strides and starting addresses to be at least 4-byte aligned. This is not just an old MatMul comment: the pinned Dawn v20261013.193955 storage-access lowering converts the FP16 offset/stride to bytes without padding, and its printer uses MatrixScope::Wave. The sibling MatMul explicitly pads odd FP16 N for this requirement.

Consequently, a valid model now reaches an unsupported/undefined matrix load, so numerical correctness is not guaranteed. At the previous reviewed head the same eligible Conv skipped prepacking and transposed to OHWI, whose K=144 column stride and tile offsets are aligned. This regression is introduced by the new HWIO path, not by sharing OHWI itself.

Please decline odd-N HWIO weights before output allocation, preserving the existing HWIO fallback, or provide an aligned/staged weight representation and matching tile addressing. Keep odd-N OHWI support. Add a path-verified HardSwish/odd-Cout constant-weight regression; the new HWIO test uses only Cout=32. The alignment violation is established by code and contract, not a claimed GPU reproduction.

xhcao force-pushed the subgroup_matrix_conv_2 branch from 8260374 to c6f937c Compare August 25, 2026 02:08

xhcao mentioned this pull request Aug 25, 2026

webgpu: route Conv2D to the subgroup-matrix MatMul kernel #32004

Closed

webgpu: using subgroup matrix to optimize Conv operator

3702f0f

xhcao force-pushed the subgroup_matrix_conv_2 branch from 5d0db8c to 3702f0f Compare September 20, 2026 04:53

xhcao added 5 commits October 8, 2026 15:35

Merge branch 'main' into pr32252

4fc9dbb

Add a ConvHeuristicTiling function and rebase code

76a1a26

modify the ConvHeuristicTiling

e16e3d6

Merge remote-tracking branch 'upstream/main' into subgroup_matrix_conv_2

7ca4886

Optimize some code

57ed167

xhcao marked this pull request as ready for review October 16, 2026 02:53

Copilot AI balanced review requested due to automatic review settings October 16, 2026 02:53

Copilot started reviewing on behalf of xhcao October 16, 2026 02:53

View session

Copilot AI reviewed Oct 16, 2026

View reviewed changes

xhcao added 3 commits October 19, 2026 14:52

Reject weights exceeding maxStorageBufferBindingSize

9c98929

Merge remote-tracking branch 'upstream/main' into subgroup_matrix_conv_2

14a4e7a

Remove conv tuned tree

2b3fabb

Remove kConvWideNSplitKBlocks

fd5af66

Jie Chen (jchen10) reviewed Oct 20, 2026

View reviewed changes

Address comments and optimize ConvHeuristicTiling

0e3f41a

xhcao commented Oct 21, 2026

View reviewed changes

Jie Chen (jchen10) approved these changes Oct 21, 2026

View reviewed changes

xhcao added 2 commits October 22, 2026 13:07

Address Jie's comments

1110cd1

Merge remote-tracking branch 'upstream/main' into subgroup_matrix_conv_2

7354fdc

xhcao commented Oct 22, 2026

View reviewed changes

github-advanced-security Bot found potential problems Oct 26, 2026

View reviewed changes

Comment thread onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc Fixed

xhcao added 2 commits October 26, 2026 10:14

Fix bots' errors

9ac6a2f

Merge remote-tracking branch 'upstream/main' into subgroup_matrix_conv_2

74e62fe

Jiajia Qin (qjia7) reviewed Oct 26, 2026

View reviewed changes

xhcao added 2 commits October 27, 2026 13:07

Address Jiajia's comments

92f5aaf

Merge remote-tracking branch 'upstream/main' into subgroup_matrix_conv_2

a46dc56

xhcao commented Oct 27, 2026

View reviewed changes

Jiajia Qin (qjia7) reviewed Oct 27, 2026

View reviewed changes

		let k_vec_base = kb * kSgMatK / kVecSize;
		let k_vec_per_row = kSgMatK / kVecSize;

		let src_vec_batch_offset = batch_id * uniforms.src_h * uniforms.src_w * (uniforms.channel_i / kVecSize);
		let out_batch_offset = batch_id * uniforms.M * uniforms.N;

		constexpr uint32_t kMatMulScratchNPad = 0;
		constexpr uint32_t kConvScratchNPad = kSubgroupMatrixK;

		// Add split-K only while the independent tiles under-fill the machine; cap the
		// total at 2x hardware and keep enough K work per subgroup and scratch budget.

Conversation

xhcao commented Aug 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Motivation and Context

Uh oh!

azure-pipelines Bot commented Aug 25, 2026

Uh oh!

Jie Chen (jchen10) commented Sep 17, 2026

Uh oh!

Jiajia Qin (qjia7) commented Sep 17, 2026

Uh oh!

xhcao commented Sep 21, 2026

Uh oh!

azure-pipelines Bot commented Oct 16, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

xhcao commented Oct 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Jie Chen (jchen10) Oct 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

xhcao left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

xhcao commented Aug 25, 2026 •

edited

Loading

xhcao commented Oct 19, 2026 •

edited

Loading

Jie Chen (jchen10) Oct 21, 2026 •

edited

Loading