webgpu: using subgroup matrix to optimize Conv operator

onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc

@@ -15,10 +15,10 @@
 #include "core/providers/webgpu/compute_context.h"
 #include "core/providers/webgpu/math/matmul.h"
 #include "core/providers/webgpu/math/subgroup_matrix_config.h"
+#include "core/providers/webgpu/math/subgroup_matrix_pad_b.h"
 #include "core/providers/webgpu/program.h"
 #include "core/providers/webgpu/shader_helper.h"
 #include "core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.h"
-#include "core/providers/webgpu/webgpu_utils.h"
 namespace onnxruntime {
 namespace webgpu {
 
@@ -53,24 +53,6 @@ class SubgroupMatrixMatMulProgram final : public Program<SubgroupMatrixMatMulPro
 
 namespace {
 
-// Copies a row-major f16 weight B [K, N] into a column-padded [K, N_b] buffer
-// (N_b >= N), zero-filling columns [N, N_b). Gives B an aligned row stride so the
-// subgroup-matrix f16 load's row-start alignment holds for any N.
-class SubgroupMatrixMatMulPadBProgram final : public Program<SubgroupMatrixMatMulPadBProgram> {
- public:
-  SubgroupMatrixMatMulPadBProgram() : Program{"SubgroupMatrixMatMulPadB"} {}
-  Status GenerateShaderCode(ShaderHelper& shader) const override {
-    const auto& input_b = shader.AddInput("input_b", ShaderUsage::UseValueTypeAlias);
-    const auto& output = shader.AddOutput("output", ShaderUsage::UseValueTypeAlias);
-    return WGSL_TEMPLATE_APPLY(shader, "math/subgroup_matrix_matmul_pad_b.wgsl.template",
-                               WGSL_TEMPLATE_VARIABLE(input_b, input_b),
-                               WGSL_TEMPLATE_VARIABLE(output, output));
-  }
-  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"output_size", ProgramUniformVariableDataType::Uint32},
-                                          {"N", ProgramUniformVariableDataType::Uint32},
-                                          {"N_b", ProgramUniformVariableDataType::Uint32});
-};
-
 // Subgroup-matrix MatMul implementation. Loads both A and B directly from global
 // memory and runs the subgroup-matrix kernel during Compute. The class is
 // intended to support all subgroup-matrix configs. For now 8x16x16 and 16x16x16
@@ -264,45 +246,18 @@ class SubgroupMatrixMatMulImplInternal final : public SubgroupMatrixMatMulImpl {
   }
 
  private:
-  // Lazily builds an aligned copy of a constant weight B [..., K, N] with unaligned N
-  // by widening its last dim to an aligned N_b (zero-filling the extra columns) and
-  // caches it, so the per-run pad cost is paid once. Works for a 2D weight [K, N]
-  // and a batched weight [batch, K, N] alike: the pad pass treats B as a flat
-  // [rows, N] -> [rows, N_b] copy over rows = numel / N (= K, or batch*K), which is
-  // exactly the even-stride layout the kernel indexes via N_b. Runs the GPU pad pass
-  // on first use under call_once; the cached tensor is held for the kernel's
-  // lifetime. Only valid when B is a constant initializer (checked by the caller) -
-  // a runtime B changes per run and must not be cached.
+  // Lazily builds an aligned copy of a constant weight B with unaligned N, caching
+  // it so the per-run pad cost is paid once. Only valid when B is a constant
+  // initializer (checked by the caller) - a runtime B changes per run and must not
+  // be cached.
   Status EnsurePaddedB(ComputeContext& context,
                        const Tensor& b,
                        const TensorShape& b_shape,
                        uint32_t N,
                        uint32_t n_b) const {
-    TensorShapeVector padded_dims{b_shape.GetDims().begin(), b_shape.GetDims().end()};
-    padded_dims.back() = static_cast<int64_t>(n_b);
-    const TensorShape padded_shape{padded_dims};
-    const int64_t output_size_i64 = padded_shape.Size();
-    ORT_RETURN_IF_NOT(output_size_i64 <= static_cast<int64_t>(std::numeric_limits<uint32_t>::max()),
-                      "Cannot pad B because the padded tensor has ", output_size_i64,
-                      " elements, exceeding uint32_t shader indexing range.");
-    const uint32_t output_size = narrow<uint32_t>(output_size_i64);
-
     std::call_once(pad_once_, [&]() {
-      auto padded = std::make_unique<Tensor>(context.CreateGPUTensor(b.DataType(), padded_shape));
-      Status s = Status::OK();
-      // A zero-element padded tensor (e.g. a zero-batch or empty constant B) needs no
-      // pad pass - dispatching 0 workgroups is pointless and some drivers reject it.
-      // Just cache the empty tensor; the main kernel dispatches nothing for it.
-      if (output_size != 0) {
-        SubgroupMatrixMatMulPadBProgram program;
-        program.SetWorkgroupSize(WORKGROUP_SIZE)
-            .SetDispatchGroupSize(CeilDiv<uint32_t>(output_size, WORKGROUP_SIZE))
-            .AddInput({&b, ProgramTensorMetadataDependency::TypeAndRank, b_shape, 1})
-            .AddOutput({padded.get(), ProgramTensorMetadataDependency::TypeAndRank, padded->Shape(), 1})
-            .AddUniformVariables({{output_size}, {N}, {n_b}});
-        s = context.RunProgram(program);
-      }
-      if (s.IsOK()) {
+      std::unique_ptr<Tensor> padded;
+      if (PadSubgroupMatrixWeightColumns(context, b, b_shape, N, n_b, padded).IsOK()) {
         padded_b_ = std::move(padded);
       }
     });

onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul_pad_b.wgsl.template

@@ -4,7 +4,8 @@
 // Copies a row-major f16 weight B [K, N] into a column-padded [K, N_b] buffer
 // (N_b >= N), zero-filling columns [N, N_b). One thread per padded-output element.
 // Gives B an aligned row stride so the subgroup-matrix f16 load's row-start
-// alignment holds for any N. See EnsurePaddedB in subgroup_matrix_matmul.cc.
+// alignment holds for any N. Shared by MatMul and Conv; see
+// PadSubgroupMatrixWeightColumns in subgroup_matrix_pad_b.cc.
 
 #use guardAgainstOutOfBoundsWorkgroupSizes
 #use .getByOffset .setByOffset

onnxruntime/core/providers/webgpu/math/subgroup_matrix_pad_b.cc

@@ -0,0 +1,62 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/webgpu/math/subgroup_matrix_pad_b.h"
+
+#include <limits>
+
+#include "core/common/narrow.h"
+#include "core/providers/webgpu/program.h"
+#include "core/providers/webgpu/shader_helper.h"
+#include "core/providers/webgpu/webgpu_utils.h"
+
+namespace onnxruntime {
+namespace webgpu {
+
+namespace {
+
+class SubgroupMatrixPadBProgram final : public Program<SubgroupMatrixPadBProgram> {
+ public:
+  SubgroupMatrixPadBProgram() : Program{"SubgroupMatrixPadB"} {}
+  Status GenerateShaderCode(ShaderHelper& shader) const override {
+    const auto& input_b = shader.AddInput("input_b", ShaderUsage::UseValueTypeAlias);
+    const auto& output = shader.AddOutput("output", ShaderUsage::UseValueTypeAlias);
+    return WGSL_TEMPLATE_APPLY(shader, "math/subgroup_matrix_matmul_pad_b.wgsl.template",
+                               WGSL_TEMPLATE_VARIABLE(input_b, input_b),
+                               WGSL_TEMPLATE_VARIABLE(output, output));
+  }
+  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"output_size", ProgramUniformVariableDataType::Uint32},
+                                          {"N", ProgramUniformVariableDataType::Uint32},
+                                          {"N_b", ProgramUniformVariableDataType::Uint32});
+};
+
+}  // namespace
+
+Status PadSubgroupMatrixWeightColumns(ComputeContext& context, const Tensor& b, const TensorShape& b_shape,
+                                      uint32_t N, uint32_t n_b, std::unique_ptr<Tensor>& padded) {
+  TensorShapeVector padded_dims{b_shape.GetDims().begin(), b_shape.GetDims().end()};
+  padded_dims.back() = static_cast<int64_t>(n_b);
+  const TensorShape padded_shape{padded_dims};
+  const int64_t output_size_i64 = padded_shape.Size();
+  ORT_RETURN_IF_NOT(output_size_i64 <= static_cast<int64_t>(std::numeric_limits<uint32_t>::max()),
+                    "Cannot pad the subgroup-matrix weight because the padded tensor has ", output_size_i64,
+                    " elements, exceeding uint32_t shader indexing range.");
+  const uint32_t output_size = narrow<uint32_t>(output_size_i64);
+
+  padded = std::make_unique<Tensor>(context.CreateGPUTensor(b.DataType(), padded_shape));
+  // A zero-element padded tensor needs no pad pass -- dispatching 0 workgroups is
+  // pointless and some drivers reject it.
+  if (output_size == 0) {
+    return Status::OK();
+  }
+  SubgroupMatrixPadBProgram program;
+  program.SetWorkgroupSize(WORKGROUP_SIZE)
+      .SetDispatchGroupSize(CeilDiv<uint32_t>(output_size, WORKGROUP_SIZE))
+      .AddInput({&b, ProgramTensorMetadataDependency::TypeAndRank, b_shape, 1})
+      .AddOutput({padded.get(), ProgramTensorMetadataDependency::TypeAndRank, padded->Shape(), 1})
+      .AddUniformVariables({{output_size}, {N}, {n_b}});
+  return context.RunProgram(program);
+}
+
+}  // namespace webgpu
+}  // namespace onnxruntime

onnxruntime/core/providers/webgpu/math/subgroup_matrix_pad_b.h

@@ -0,0 +1,27 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <cstdint>
+#include <memory>
+
+#include "core/framework/tensor.h"
+#include "core/framework/tensor_shape.h"
+#include "core/providers/webgpu/compute_context.h"
+
+namespace onnxruntime {
+namespace webgpu {
+
+// Pads a row-major f16 weight B [..., K, N] along its last dim to N_b (>= N),
+// zero-filling the new columns, giving B an aligned row stride/offset for a
+// subgroup-matrix row-major load. Treats B as a flat [rows, N] -> [rows, N_b] copy
+// with rows = numel / N, so a 2D [K, N] weight and a batched [batch, K, N] one (or,
+// for Conv, an HWIO [kh, kw, Cin, Cout] one) all work. Always runs the GPU pad
+// pass; callers needing to pay this once for a constant weight should cache the
+// result themselves (e.g. via std::call_once), as subgroup_matrix_matmul.cc does.
+Status PadSubgroupMatrixWeightColumns(ComputeContext& context, const Tensor& b, const TensorShape& b_shape,
+                                      uint32_t N, uint32_t n_b, std::unique_ptr<Tensor>& padded);
+
+}  // namespace webgpu
+}  // namespace onnxruntime

onnxruntime/core/providers/webgpu/nn/conv.cc

@@ -1,9 +1,11 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
+
 #include "core/providers/webgpu/nn/conv.h"
 #include "core/providers/webgpu/nn/conv2d_mm.h"
 #include "core/providers/webgpu/nn/conv3d_naive.h"
 #include "core/providers/webgpu/nn/im2col_matmul.h"
+#include "core/providers/webgpu/nn/subgroup_matrix_conv.h"
 #include "core/providers/webgpu/shader_helper.h"
 #include "core/providers/webgpu/webgpu_supported_types.h"
 #include "core/providers/webgpu/tensor/transpose.h"
@@ -26,31 +28,15 @@ Status TransposeKernel(ComputeContext& context, const Tensor* kernel, const Tens
   return Transpose::DoTranspose(context, perm, reshaped_kernel, *transposed_kernel);
 }
 
-template <bool is_channels_last, bool is_fused>
-Status Conv<is_channels_last, is_fused>::ComputeInternal(ComputeContext& context) const {
-  bool has_bias = context.InputCount() > 2;
-  const auto* input = context.Input<Tensor>(0);
-  const Tensor* kernel = prepacked_kernel_ ? prepacked_kernel_.get() : context.Input<Tensor>(1);
-  const auto* bias = has_bias ? context.Input<Tensor>(2) : nullptr;
-  TensorShape input_shape = input->Shape();
-  ORT_ENFORCE(kernel != nullptr, "Conv kernel tensor is required.");
-  // Prepacked kernels are stored permuted; recover the logical OIHW shape.
-  TensorShape kernel_shape = kernel->Shape();
-  switch (kernel_layout_) {
-    case KernelLayout::OIHW:
-      break;
-    case KernelLayout::HWIO:
-      kernel_shape = TensorShape(TensorShapeVector{kernel_shape[3], kernel_shape[2], kernel_shape[0], kernel_shape[1]});
-      break;
-    case KernelLayout::OHWI:
-      kernel_shape = TensorShape(TensorShapeVector{kernel_shape[0], kernel_shape[3], kernel_shape[1], kernel_shape[2]});
-      break;
-  }
-  ConvAttributes::ConvPadVector local_pads(conv_attrs_.pads.begin(), conv_attrs_.pads.end());
-  TensorShapeVector local_dilations(conv_attrs_.dilations.begin(), conv_attrs_.dilations.end());
-  TensorShapeVector local_strides(conv_attrs_.strides.begin(), conv_attrs_.strides.end());
+Status ComputeConvGeometry(const ConvAttributes& conv_attrs, const TensorShape& input_shape,
+                           const TensorShape& kernel_shape, bool is_channels_last,
+                           std::vector<uint32_t>& pads, std::vector<uint32_t>& strides,
+                           std::vector<uint32_t>& dilations, TensorShape& output_shape) {
+  ConvAttributes::ConvPadVector local_pads(conv_attrs.pads.begin(), conv_attrs.pads.end());
+  TensorShapeVector local_dilations(conv_attrs.dilations.begin(), conv_attrs.dilations.end());
+  TensorShapeVector local_strides(conv_attrs.strides.begin(), conv_attrs.strides.end());
   TensorShapeVector kernel_spacial_shape_vector;
-  ORT_RETURN_IF_ERROR(conv_attrs_.ComputeKernelShape(kernel_shape, kernel_spacial_shape_vector, false));
+  ORT_RETURN_IF_ERROR(conv_attrs.ComputeKernelShape(kernel_shape, kernel_spacial_shape_vector, false));
   if (local_pads.empty()) {
     local_pads.resize(kernel_spacial_shape_vector.size() * 2, 0);
   }
@@ -63,23 +49,80 @@ Status Conv<is_channels_last, is_fused>::ComputeInternal(ComputeContext& context
   TensorShapeVector input_shape_vector = input_shape.AsShapeVector();
   auto batch = input_shape[0];
   TensorShapeVector output_shape_vector = {batch};
-  TensorShape input_spacial_shape = is_channels_last ? TensorShape(TensorShapeVector(std::next(input_shape_vector.begin()), std::prev(input_shape_vector.end()))) : input_shape.Slice(2);
-  ORT_RETURN_IF_ERROR(conv_attrs_.InferPadsAndOutputShape(input_spacial_shape, kernel_spacial_shape_vector, local_strides, local_dilations, local_pads, output_shape_vector));
+  TensorShape input_spacial_shape = is_channels_last
+                                        ? TensorShape(TensorShapeVector(std::next(input_shape_vector.begin()), std::prev(input_shape_vector.end())))
+                                        : input_shape.Slice(2);
+  ORT_RETURN_IF_ERROR(conv_attrs.InferPadsAndOutputShape(input_spacial_shape, kernel_spacial_shape_vector,
+                                                         local_strides, local_dilations, local_pads,
+                                                         output_shape_vector));
   auto output_channels = kernel_shape[0];
   if (is_channels_last) {
     output_shape_vector.push_back(output_channels);
   } else {
     output_shape_vector.insert(output_shape_vector.begin() + 1, output_channels);
   }
-  auto output_shape = TensorShape(output_shape_vector);
-  auto* output = context.Output(0, output_shape);
-  std::vector<uint32_t> strides;
-  std::vector<uint32_t> pads;
-  std::vector<uint32_t> dilations;
+  output_shape = TensorShape(output_shape_vector);
+
   auto transform_dim = [](int64_t dim) { return static_cast<int32_t>(dim); };
+  pads.clear();
+  strides.clear();
+  dilations.clear();
   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);
+  return Status::OK();
+}
+
+template <bool is_channels_last, bool is_fused>
+typename Conv<is_channels_last, is_fused>::ConvOptImpl*
+Conv<is_channels_last, is_fused>::ConvOptImplCache::GetOrCreate(
+    const Conv<is_channels_last, is_fused>& parent, const ComputeContextBase& context) {
+  std::call_once(init_flag_, [&]() {
+    impl_ = CreateSubgroupMatrixConvImpl(parent, context);
+  });
+  return impl_.get();
+}
+
+template <bool is_channels_last, bool is_fused>
+Status Conv<is_channels_last, is_fused>::ComputeInternal(ComputeContext& context) const {
+  // Preferred fast path on devices with a vendor subgroup-matrix tiling policy;
+  // declines (handled=false) when the kernel/attrs don't qualify.
+  if (ConvOptImpl* opt_impl = opt_impl_cache_.GetOrCreate(*this, context); opt_impl != nullptr) {
+    bool handled = false;
+    ORT_RETURN_IF_ERROR(opt_impl->Compute(context, handled));
+    if (handled) {
+      return Status::OK();
+    }
+  }
+
+  bool has_bias = context.InputCount() > 2;
+  const auto* input = context.Input<Tensor>(0);
+  const Tensor* kernel = prepacked_kernel_ ? prepacked_kernel_.get() : context.Input<Tensor>(1);
+  const auto* bias = has_bias ? context.Input<Tensor>(2) : nullptr;
+  TensorShape input_shape = input->Shape();
+  ORT_ENFORCE(kernel != nullptr, "Conv kernel tensor is required.");
+  // Prepacked kernels are stored permuted; recover the logical OIHW shape.
+  TensorShape kernel_shape = kernel->Shape();
+  switch (kernel_layout_) {
+    case KernelLayout::OIHW:
+      break;
+    case KernelLayout::HWIO:
+      kernel_shape = TensorShape(TensorShapeVector{kernel_shape[3], kernel_shape[2], kernel_shape[0], kernel_shape[1]});
+      break;
+    case KernelLayout::OHWI:
+      kernel_shape = TensorShape(TensorShapeVector{kernel_shape[0], kernel_shape[3], kernel_shape[1], kernel_shape[2]});
+      break;
+  }
+  std::vector<uint32_t> pads;
+  std::vector<uint32_t> strides;
+  std::vector<uint32_t> dilations;
+  TensorShape output_shape;
+  ORT_RETURN_IF_ERROR(ComputeConvGeometry(conv_attrs_, input_shape, kernel_shape, is_channels_last,
+                                          pads, strides, dilations, output_shape));
+  auto* output = context.Output(0, output_shape);
+  auto batch = input_shape[0];
+  auto output_channels = kernel_shape[0];
+  TensorShapeVector output_shape_vector = output_shape.AsShapeVector();
   auto rank = input_shape.NumDimensions();
   const InlinedVector<size_t> perm = {2, 3, 1, 0};
   if (rank > 5) {
@@ -125,6 +168,7 @@ Status Conv<is_channels_last, is_fused>::ComputeInternal(ComputeContext& context
   } else if (rank == 3) {
     // Conv1D
     TensorShapeVector kernel_shape_vector = kernel_shape.AsShapeVector();
+    TensorShapeVector input_shape_vector = input_shape.AsShapeVector();
     input_shape_vector.insert(input_shape_vector.begin() + (is_channels_last ? 1 : 2), 1, 1);
     output_shape_vector.insert(output_shape_vector.begin() + (is_channels_last ? 1 : 2), 1, 1);
     kernel_shape_vector.insert(kernel_shape_vector.begin() + 2, 1);
@@ -323,6 +367,7 @@ Status Conv<is_channels_last, is_fused>::PrePackInternal(ComputeContextBase& con
   }
 
   // Im2ColMatMul path: transpose OIHW -> OHWI once here instead of on every inference.
+  // Subgroup-matrix Conv (tried first in ComputeInternal) shares this same prepack.
   //
   // Placed before the auto_pad check below on purpose:
   //   - Safe: CanApplyIm2ColMatMulProgram() only looks at the adapter, dtype, layout,
@@ -332,10 +377,10 @@ Status Conv<is_channels_last, is_fused>::PrePackInternal(ComputeContextBase& con
   //   - Necessary: otherwise every auto_pad != NOTSET model would bail out below and
   //     keep paying for the transpose on every inference.
   //
-  // This call and the one in ComputeInternal must stay in agreement: only the im2col
-  // path can read the OHWI layout, so a decision made here that ComputeInternal later
-  // reverses would corrupt the weights. If CanApplyIm2ColMatMulProgram() ever gains a
-  // condition that is not known at prepack time (pads, strides, input shape), this
+  // This call and the one in ComputeInternal must stay in agreement: only im2col and
+  // subgroup-matrix Conv can read OHWI, so a decision made here that ComputeInternal
+  // later reverses would corrupt the weights. If CanApplyIm2ColMatMulProgram() ever
+  // gains a condition not known at prepack time (pads, strides, input shape), this
   // shortcut must go away. ComputeInternal ORT_ENFORCEs the invariant.
   if (CanApplyIm2ColMatMulProgram(context, is_channels_last, activation_,
                                   kernel_shape, onnxruntime::narrow<uint32_t>(conv_attrs_.group),
@@ -367,6 +412,8 @@ Status Conv<is_channels_last, is_fused>::PrePackInternal(ComputeContextBase& con
 
   // Analyze execution paths in ComputeInternal to determine if kernel transpose is needed:
   //
+  // 0. Subgroup-matrix Conv path: no dedicated branch -- reads whatever layout falls
+  //    out below (OIHW, OHWI from Im2ColMatMul, or HWIO from the general transpose).
   // 1. Im2ColMatMul path: handled above (prepacked as OHWI)
   // 2. Grouped conv (group > 1): handled above (skip if !is_channels_last)
   // 3. MatMul optimization (same_size || is_1x1_conv):