webgpu: route Conv2D to the subgroup-matrix MatMul kernel by xhcao · Pull Request #32004 · microsoft/onnxruntime

xhcao · 2026-08-12T03:32:11Z

Description

Reuse the existing subgroup-matrix MatMul kernel for Conv2D on WebGPU.

Motivation and Context

While MatMul/Gemm already had a much faster subgroup-matrix (cooperative matrix) kernel with vendor-tuned tiling, so routing some convolution computation to the same kernel closes an easy performance gap without touching the general Conv path.

azure-pipelines · 2026-08-12T03:32:20Z

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

Extract the subgroup-matrix MatMul into a reusable SubgroupMatrixMatMulContext (device capability check + vendor tiling policy) dispatched via the shared ComputeMatMul(). The 1x1 Conv path now caches this context and runs the subgroup-matrix kernel, threading w_is_constant/b_is_constant so odd-N constant weights are still padded once and cached.

Jie Chen (jchen10) · 2026-08-12T11:04:32Z

+      std::call_once(impl_init_flag_, [&]() {
+        impl_ = CreateSubgroupMatrixConv1x1Impl(context);
+      });
+      if (impl_ != nullptr && activation_.activation_kind_ == ActivationKind::None) {


Let impl_ decide whether to reject activation_.activation_kind_ == ActivationKind::None

Jie Chen (jchen10) · 2026-08-12T11:12:12Z

+#if !defined(__wasm__)
+      // Lazily create the subgroup-matrix 1x1-Conv impl once so the device
+      // capability query is not repeated on every inference step. It computes a
+      // plain (non-activation-fused) matmul, so only attempt it when no activation
+      // is fused; on the channels-last path B is the (possibly constant) weight, so
+      // an odd-N weight can be padded/cached, whereas channels-first B is the
+      // activation input and is never constant.
+      std::call_once(impl_init_flag_, [&]() {
+        impl_ = CreateSubgroupMatrixConv1x1Impl(context);
+      });
+      if (impl_ != nullptr && activation_.activation_kind_ == ActivationKind::None) {
+        bool handled = false;
+        ORT_RETURN_IF_ERROR(impl_->Compute(context, matmul_inputs, output,
+                                           matmul_input_reshapes[0], matmul_input_reshapes[1],
+                                           is_channels_last && w_is_constant_, handled));
+        if (handled) {
+          return Status::OK();
+        }
+      }
+#endif


Move this at the beginning of Compute(). Check the 1x1 condition by impl itself.

Jie Chen (jchen10) · 2026-08-12T11:56:50Z

+    // When B is a shared 2D weight (its batch folds to 1), fold A's leading dims
+    // into M so the problem runs as a single 2D-weight matmul (e.g. the 1x1 Conv
+    // [batch,H*W,C] @ [1,C,N] -> [batch*H*W, C] @ [C, N]).
+    const int64_t batch_a = a_shape.SizeToDimension(a_shape.NumDimensions() - 2);
+    const int64_t batch_b = b_shape.SizeToDimension(b_shape.NumDimensions() - 2);
+    if (batch_a != 1 && batch_b == 1) {
+      const int64_t k = a_shape[a_shape.NumDimensions() - 1];
+      const int64_t n = b_shape[b_shape.NumDimensions() - 1];
+      const int64_t batch_and_m = a_shape.SizeToDimension(a_shape.NumDimensions() - 1);
+      a_shape = TensorShape({batch_and_m, k});
+      b_shape = TensorShape({k, n});
+    }


It seems we are duplicating most Compute() code lines with SubgroupMatrixMatMul except this block here and input a, b reshape, right? Maybe we can refactor a bit.

Jie Chen (jchen10) · 2026-08-17T01:42:25Z

+  });
+  if (impl_ != nullptr) {
+    bool handled = false;
+    ORT_RETURN_IF_ERROR(impl_->Compute(context, conv_attrs_, activation_, is_channels_last,


Define an abstract class ConvOptImpl as MatMulOptImpl does, so that we needn't pass conv_attrs_, activation_, ... to impl here. Instead impl can access them via parent.

Jie Chen (jchen10) · 2026-08-17T02:51:48Z

+    if (!is_channels_last || prepacked_kernel != nullptr || !w_is_constant ||
+        activation.activation_kind_ != ActivationKind::None || conv_attrs.group != 1) {
+      return Status::OK();
+    }


prepacked_kernel != nullptr is hard to understand, as it's logically complicated which implies quite a few conditions. As for conv(matmul) here, the major conditions are:

// 3. MatMul optimization (same_size || is_1x1_conv): // - is_channels_last: transposes // - !is_channels_last: does NOT transpose (uses kernel directly)

As you only serve channels-last, should we suppose or just assert prepacked_kernel != nullptr?

Jie Chen (jchen10) · 2026-08-17T02:53:14Z

+    const auto kernel_width = kernel_shape[3];
+
+    const bool same_size = input_height == kernel_height && input_width == kernel_width && pads[0] == 0 && pads[1] == 0;
+    const bool is_conv1x1_matmul = same_size || (kernel_height == 1 && kernel_width == 1 && pads[0] == 0 && pads[1] == 0 && strides[0] == 1 && strides[1] == 1);


As we also support same size, conv1x1 is not accurate to be used as the file/class name.

Jie Chen (jchen10) · 2026-08-17T11:15:27Z

+  virtual Status Compute(ComputeContext& context, /*out*/ bool& handled) = 0;
+
+ protected:
+  const ConvOptImplParent& parent_;


No need to define a new class ConvOptImplParent. Simply use Conv class instead.

Jie Chen (jchen10) · 2026-08-17T11:21:41Z

is_conv_matmul instead of is_conv1x1_matmul

azure-pipelines · 2026-08-18T03:22:53Z

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

Jie Chen (jchen10)

LGTM

Copilot

Pull request overview

Routes eligible WebGPU Conv2D operations through the shared subgroup-matrix MatMul kernel.

Changes:

Adds subgroup-matrix Conv implementation and shared configuration utilities.
Refactors MatMul/Gemm to reuse common tiling and padding logic.
Adds disabled large Conv test cases.

Reviewed changes

Copilot reviewed 10 out of 12 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
`onnxruntime/test/providers/webgpu/conv_large_test.cc`	Adds large Conv correctness cases.
`onnxruntime/core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.h`	Uses shared subgroup definitions.
`onnxruntime/core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.cc`	Updates shared-header dependency.
`onnxruntime/core/providers/webgpu/subgroup_matrix_common.h`	Declares common configuration APIs.
`onnxruntime/core/providers/webgpu/subgroup_matrix_common.cc`	Implements device and tiling selection.
`onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.h`	Declares optimized Conv factory.
`onnxruntime/core/providers/webgpu/nn/subgroup_matrix_conv.cc`	Implements Conv-to-MatMul routing.
`onnxruntime/core/providers/webgpu/nn/conv.h`	Adds optimized implementation state.
`onnxruntime/core/providers/webgpu/nn/conv.cc`	Invokes the subgroup path first.
`onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.h`	Exposes shared dispatch and padding.
`onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc`	Refactors reusable MatMul execution.
`onnxruntime/core/providers/webgpu/math/subgroup_matrix_gemm.cc`	Reuses common device selection.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

xhcao · 2026-08-20T08:41:38Z

+  return [](const ComputeContext&, uint32_t /*M*/, uint32_t /*N*/,
+            uint32_t /*K*/, uint32_t /*batch*/) -> std::optional<SubgroupMatrixTiling> {
+    return SubgroupMatrixTiling{32, 32, 1};
+  };


xhcao · 2026-08-20T08:41:52Z

+    const bool same_size = input_height == kernel_height && input_width == kernel_width && pads[0] == 0 && pads[1] == 0;
+    const bool is_conv_matmul = same_size || (kernel_height == 1 && kernel_width == 1 && pads[0] == 0 && pads[1] == 0 && strides[0] == 1 && strides[1] == 1);


xhcao · 2026-08-20T08:41:08Z

+  RunConvTest<MLFloat16, version>(p);
+}
+
+TEST(Conv_Large, DISABLED_Conv1x1) {


xhcao · 2026-08-20T08:41:24Z

+// the MatMul, Gemm and Conv 1x1 factories. On a supported device sets config_index /
+// tiling_selector and returns true; returns false (leaving the outputs untouched) when
+// the device or vendor is unsupported, so the caller yields no implementation.


Jiajia Qin (qjia7) · 2026-08-24T09:13:19Z

+  // (handled=false, allocating nothing) otherwise and we fall through to the normal
+  // path order below.
+  std::call_once(impl_init_flag_, [&]() {
+    impl_ = CreateSubgroupMatrixConvImpl(*this, context);


If subgroupMatrix is only used for 1x1 Conv, I would like to integrate SubgroupMatrix support into [a] ComputeMatMul instead of here. That's an unified entrance both for Conv and MatMul ops.

Jiajia Qin (@qjia7), I will apply SubgroupMatrix to general Conv, #32252, will let these two PRs share some code. Furthermore, it would be inconvenient in many ways if integrate these code into ComputeMatMul, take the cached transposed tensor for an example.

Jiajia Qin (qjia7) · 2026-08-24T09:17:26Z

          .AddUniformVariables({{output_size}, {static_cast<uint32_t>(matmul_output_shape[1])}, {static_cast<uint32_t>(matmul_output_shape[2])}, {static_cast<uint32_t>(K)}});
      return context.RunProgram(program);
    } else {
      return ComputeMatMul(&context, activation_, matmul_inputs, output, is_channels_last, matmul_input_reshapes[0], matmul_input_reshapes[1]);


[a] Please also integrate MatMulNaiveProgram into ComputeMatMul so that all point-wise conv just go to ComputeMatMul path. And in ComputeMatMul path, do all necessary optimizations. Conv itself doesn't need to take care what kind of algorithms it uses underlying.

xhcao · 2026-09-11T08:05:03Z

#32304 has routed conv2d 1x1 to subgroup-matrix MatMul kernel. Close it.

xhcao force-pushed the subgroup_matrix_conv1x1 branch from e462a32 to 7f5aba8 Compare August 12, 2026 07:56

Jie Chen (jchen10) reviewed Aug 12, 2026

View reviewed changes

xhcao force-pushed the subgroup_matrix_conv1x1 branch from 8380a1b to 438b814 Compare August 14, 2026 01:45

Address Jie's comments

56ca415

xhcao force-pushed the subgroup_matrix_conv1x1 branch from 438b814 to 56ca415 Compare August 14, 2026 01:47

Restrict the scope of use

abbc93f

xhcao force-pushed the subgroup_matrix_conv1x1 branch from 49166ec to abbc93f Compare August 14, 2026 08:26

Jie Chen (jchen10) reviewed Aug 17, 2026

View reviewed changes

Merge branch 'main' into subgroup_matrix_conv1x1

969cd8e

Jie Chen (jchen10) reviewed Aug 17, 2026

View reviewed changes

rename class/files' names

9ce96b3

Jie Chen (jchen10) reviewed Aug 17, 2026

View reviewed changes

xhcao changed the title ~~webgpu: reuse subgroup-matrix MatMul for 1x1 Conv~~ webgpu: route 1x1 Conv to the subgroup-matrix MatMul kernel Aug 18, 2026

xhcao changed the title ~~webgpu: route 1x1 Conv to the subgroup-matrix MatMul kernel~~ webgpu: route Conv2D to the subgroup-matrix MatMul kernel Aug 18, 2026

Remove ConvOptImplParent class

cc6f8a0

xhcao marked this pull request as ready for review August 18, 2026 03:22

Jie Chen (jchen10) approved these changes Aug 18, 2026

View reviewed changes

Jiajia Qin (qjia7) requested a balanced review from Copilot August 19, 2026 08:35

Copilot started reviewing on behalf of Jiajia Qin (qjia7) August 19, 2026 08:36

View session

Copilot AI reviewed Aug 19, 2026

View reviewed changes

Address copilot's comments

075068c

Jiajia Qin (qjia7) reviewed Aug 24, 2026

View reviewed changes

xhcao closed this Sep 11, 2026

		const bool same_size = input_height == kernel_height && input_width == kernel_width && pads[0] == 0 && pads[1] == 0;
		const bool is_conv_matmul = same_size \|\| (kernel_height == 1 && kernel_width == 1 && pads[0] == 0 && pads[1] == 0 && strides[0] == 1 && strides[1] == 1);

Conversation

xhcao commented Aug 12, 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 12, 2026

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!

azure-pipelines Bot commented Aug 18, 2026

Uh oh!

Jie Chen (jchen10) left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

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!

xhcao commented Sep 11, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

xhcao commented Aug 12, 2026 •

edited

Loading