webgpu: route Conv2D to the subgroup-matrix MatMul kernel

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

@@ -8,25 +8,19 @@
 #include <cstdint>
 #include <memory>
 #include <optional>
-#include <string_view>
 #include <utility>
 
 #include "core/common/narrow.h"
 #include "core/providers/webgpu/compute_context.h"
+#include "core/providers/webgpu/subgroup_matrix_common.h"
 #include "core/providers/webgpu/math/subgroup_matrix_config.h"
 #include "core/providers/webgpu/shader_helper.h"
-#include "core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.h"
 
 namespace onnxruntime {
 namespace webgpu {
 
 namespace {
 
-// Lanes per subgroup assumed by the subgroup-matrix kernel. The workgroup runs
-// split_k subgroups, so its size is kSubgroupMatrixSubgroupSize * split_k.
-// TODO: use subgroup-size-control to enforce the subgroup size is 32.
-constexpr uint32_t kSubgroupMatrixSubgroupSize = 32;
-
 // Subgroup-matrix Gemm implementation. Loads A and B directly from global memory
 // (transposed operands via column-major loads) and runs the cooperative
 // subgroup-matrix kernel during Compute. Y = alpha * op(A) @ op(B) + beta * C.
@@ -173,15 +167,6 @@ Status GenerateShaderCode8x16x16(ShaderHelper& shader, const ShaderVariableHelpe
                              WGSL_TEMPLATE_VARIABLE(output, output));
 }
 
-// Default tiling used on any vendor without a specialized policy: a fixed 32x32
-// output tile with no split-K.
-SubgroupMatrixTilingSelector MakeDefaultTilingSelector() {
-  return [](const ComputeContext&, uint32_t /*M*/, uint32_t /*N*/,
-            uint32_t /*K*/, uint32_t /*batch*/) -> std::optional<SubgroupMatrixTiling> {
-    return SubgroupMatrixTiling{32, 32, 1};
-  };
-}
-
 }  // namespace
 
 Status SubgroupMatrixGemmProgram::GenerateShaderCode(ShaderHelper& shader) const {
@@ -203,19 +188,9 @@ Status SubgroupMatrixGemmProgram::GenerateShaderCode(ShaderHelper& shader) const
 
 std::unique_ptr<Gemm::GemmOptImpl> CreateSubgroupMatrixGemmImpl(
     const Gemm& parent, const ComputeContextBase& context) {
-  // Only run on devices that report the fixed 8x16x16 F16 subgroup-matrix config
-  // this kernel is implemented for.
   int32_t config_index = 0;
-  if (!IsSubgroupMatrixConfigSupported(context, /*is_fp16=*/true, config_index) ||
-      !supported_subgroup_matrix_configs[config_index].Is(8, 16, 16)) {
-    return nullptr;
-  }
-  // Intel GPUs use a tuned/heuristic tiling policy; every other vendor falls back
-  // to a fixed default tiling.
-  const bool is_intel = context.AdapterInfo().vendor == std::string_view{"intel"};
-  SubgroupMatrixTilingSelector tiling_selector =
-      is_intel ? intel::CreateSubgroupMatrixTilingSelector(context) : MakeDefaultTilingSelector();
-  if (!tiling_selector) {
+  SubgroupMatrixTilingSelector tiling_selector;
+  if (!TrySelectSubgroupMatrixConfig(context, config_index, tiling_selector)) {
     return nullptr;
   }
   return std::make_unique<SubgroupMatrixGemmImpl>(parent, config_index, std::move(tiling_selector));

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

@@ -7,33 +7,17 @@
 
 #include <functional>
 #include <memory>
+#include <mutex>
 #include <optional>
 
 #include "core/providers/webgpu/math/matmul.h"
+#include "core/providers/webgpu/subgroup_matrix_common.h"
 #include "core/providers/webgpu/program.h"
 #include "core/providers/webgpu/shader_helper.h"
 
 namespace onnxruntime {
 namespace webgpu {
 
-// Per-workgroup output tiling for one MatMul problem: the tile shape and split-K
-// factor chosen by a vendor-specific policy. The subgroup-matrix shape itself is
-// separate from this selection.
-struct SubgroupMatrixTiling {
-  uint32_t tile_m;   // output rows per workgroup
-  uint32_t tile_n;   // output cols per workgroup
-  uint32_t split_k;  // subgroups cooperating along K (1 = no split)
-};
-
-// Vendor-supplied callback that selects the output tiling for a given problem.
-// batch is the number of z-dispatched slices (1 for a shared 2D weight), used by
-// the policy to scale occupancy. Returning nullopt declines the problem, so
-// MatMul falls back to another compute path. An empty selector likewise yields
-// no implementation.
-using SubgroupMatrixTilingSelector =
-    std::function<std::optional<SubgroupMatrixTiling>(const ComputeContext& context,
-                                                      uint32_t M, uint32_t N, uint32_t K, uint32_t batch)>;
-
 // Creates a MatMulOptImpl that runs the subgroup-matrix kernel on devices whose
 // vendor policy supports it. The per-problem output tiling comes from a
 // vendor-specific selector chosen internally from the device context. Returns
@@ -74,6 +58,58 @@ class SubgroupMatrixMatMulProgram final : public Program<SubgroupMatrixMatMulPro
   const uint32_t split_k_;
 };
 
+// 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 even row stride so the
+// subgroup-matrix f16 load's 4-byte row-start alignment holds for odd N. Shared by
+// the subgroup-matrix MatMul and Conv 1x1 paths.
+class SubgroupMatrixMatMulPadBProgram final : public Program<SubgroupMatrixMatMulPadBProgram> {
+ public:
+  SubgroupMatrixMatMulPadBProgram() : Program{"SubgroupMatrixMatMulPadB"} {}
+  Status GenerateShaderCode(ShaderHelper& sh) const override;
+  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"output_size", ProgramUniformVariableDataType::Uint32},
+                                          {"N", ProgramUniformVariableDataType::Uint32},
+                                          {"N_b", ProgramUniformVariableDataType::Uint32});
+};
+
+// Caches an even-strided (N_b = N + 1) copy of a constant weight B with odd N so
+// the column-pad pass runs once and is reused across inference steps. Shared by the
+// subgroup-matrix MatMul and Conv 1x1 paths, which need the same odd-N alignment
+// fix. Only valid when B is a constant initializer - a runtime B changes per run
+// and must not be cached.
+class SubgroupMatrixPadBCache {
+ public:
+  // Builds the padded copy of `b` (whose last dim must equal N) on first use and
+  // returns it via `b_used` with its row stride via `n_b`.
+  Status EnsurePaddedB(ComputeContext& context, const Tensor& b, uint32_t N,
+                       /*out*/ const Tensor*& b_used, /*out*/ uint32_t& n_b) const;
+
+ private:
+  mutable std::once_flag pad_once_;
+  mutable std::unique_ptr<Tensor> padded_b_;
+  mutable uint32_t padded_b_stride_ = 0;
+};
+
+// Runs the subgroup-matrix kernel for Y = A @ B (+ optional bias) given resolved
+// operands and their logical shapes. a_shape/b_shape may differ from the tensors'
+// own shapes when the caller reshaped them (e.g. a 1x1 Conv folding N,H,W into M).
+// Handles the shared 2D-weight and batched-B cases, odd-N even-stride padding (via
+// pad_cache when b_is_constant), vendor tiling and dispatch. When `output` is
+// non-null the caller's pre-allocated tensor is used as a flat buffer (its element
+// count must equal batch*M*N); when null the result is allocated via
+// context.Output(0, ...) shaped like a_shape with its trailing dim set to N. Sets
+// handled=true on success; leaves handled=false (touching nothing) on an
+// unsupported device/problem so the caller can fall back. Shared by the MatMul
+// operator and the Conv 1x1 path.
+Status DispatchSubgroupMatrixMatMul(ComputeContext& context,
+                                    int32_t config_index,
+                                    const SubgroupMatrixTilingSelector& tiling_selector,
+                                    const SubgroupMatrixPadBCache& pad_cache,
+                                    const Tensor* a, const Tensor* b, const Tensor* bias,
+                                    Tensor* output,
+                                    const TensorShape& a_shape, const TensorShape& b_shape,
+                                    bool b_is_constant,
+                                    /*out*/ bool& handled);
+
 }  // namespace webgpu
 }  // namespace onnxruntime
 

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

@@ -10,6 +10,9 @@
 #include "core/providers/webgpu/nn/grouped_conv.h"
 #include "core/providers/webgpu/webgpu_utils.h"
 #include "core/providers/webgpu/math/matmul.h"
+#if !defined(__wasm__)
+#include "core/providers/webgpu/nn/subgroup_matrix_conv.h"
+#endif
 
 namespace onnxruntime {
 namespace webgpu {
@@ -29,6 +32,23 @@ Status TransposeKernel(ComputeContext& context, const Tensor* kernel, const Tens
 
 template <bool is_channels_last, bool is_fused>
 Status Conv<is_channels_last, is_fused>::ComputeInternal(ComputeContext& context) const {
+#if !defined(__wasm__)
+  // Subgroup-matrix 1x1 Conv has top priority: create the impl once (device query),
+  // then let it decide inside Compute whether the kernel/attrs qualify. It declines
+  // (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 (impl_ != nullptr) {
+    bool handled = false;
+    ORT_RETURN_IF_ERROR(impl_->Compute(context, handled));
+    if (handled) {
+      return Status::OK();
+    }
+  }
+#endif
+
   bool has_bias = context.InputCount() > 2;
   const auto* input = context.Input<Tensor>(0);
   const Tensor* kernel = nullptr;
@@ -204,8 +224,11 @@ Status Conv<is_channels_last, is_fused>::ComputeInternal(ComputeContext& context
     return context.RunProgram(program);
   }
 
-  const auto same_size = is_channels_last && input_height == kernel_height && input_width == kernel_width && pads[0] == 0 && pads[1] == 0;
-  if (same_size || (kernel_height == 1 && kernel_width == 1 && pads[0] == 0 && pads[1] == 0 && strides[0] == 1 && strides[1] == 1)) {
+  // Both reshapes need all pads (not just the leading ones) and the inferred output
+  // geometry to qualify; see the predicates in conv.h.
+  const auto same_size = is_channels_last && IsConvSameSizeMatMul(input_height, input_width, kernel_height,
+                                                                  kernel_width, output_height, output_width, pads);
+  if (same_size || IsConv1x1MatMul(kernel_height, kernel_width, pads, strides)) {
     Tensor transposed_kernel;
     TensorShape input_reshape;
     TensorShape kernel_reshape;
@@ -367,9 +390,10 @@ Status Conv<is_channels_last, is_fused>::PrePackInternal(ComputeContextBase& con
   const int64_t kernel_height = dims[2];
   const int64_t kernel_width = dims[3];
 
-  const bool is_1x1_conv =
-      (kernel_height == 1 && kernel_width == 1 && pads[0] == 0 && pads[1] == 0 &&
-       strides[0] == 1 && strides[1] == 1);
+  // Must stay in sync with the gate in ComputeInternal: if this predicate accepted a
+  // Conv that ComputeInternal sends down the general path instead, that path would read
+  // a kernel we deliberately left untransposed here.
+  const bool is_1x1_conv = IsConv1x1MatMul(kernel_height, kernel_width, pads, strides);
 
   if constexpr (!is_channels_last) {
     if (is_1x1_conv) {

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

@@ -3,6 +3,11 @@
 
 #pragma once
 
+#include <algorithm>
+#include <cstdint>
+#include <memory>
+#include <mutex>
+
 #include "core/common/optional.h"
 #include "core/providers/webgpu/webgpu_kernel.h"
 #include "core/providers/cpu/nn/conv_attributes.h"
@@ -13,13 +18,77 @@
 namespace onnxruntime {
 namespace webgpu {
 
+// Eligibility for the MatMul Conv fast paths, which reshape a Conv into a plain 2D
+// matmul (see Conv::ComputeInternal and the subgroup-matrix Conv impl). A reshape is
+// only equivalent to the Conv when nothing is padded: a leading pad shifts the window,
+// and a trailing pad adds output positions the matmul never computes. Comparing the
+// inferred output geometry on top of that also covers dilations, which change the
+// effective kernel size without appearing anywhere in the reshape.
+// The containers hold the spatial attributes in ONNX order, i.e. pads is
+// [begin..., end...] and strides is one entry per spatial dim; both must already be
+// defaulted/normalized to the 2D layout the callers use.
+
+// True when every pad is zero, leading and trailing, across all spatial dims.
+template <typename PadContainer>
+bool HasNoConvPadding(const PadContainer& pads) {
+  return std::all_of(pads.begin(), pads.end(), [](auto p) { return p == 0; });
+}
+
+// 1x1 reshape: each N,H,W position becomes an independent matmul row, so the output
+// grid has to match the input grid one-to-one in row-major order. With a 1x1 kernel, no
+// padding and unit strides that holds for any dilation (the effective kernel stays
+// 1x1), so this predicate needs only the attributes -- PrePackInternal, which cannot
+// see the input shape, uses the same one to stay in sync with ComputeInternal.
+template <typename PadContainer, typename StrideContainer>
+bool IsConv1x1MatMul(int64_t kernel_height, int64_t kernel_width,
+                     const PadContainer& pads, const StrideContainer& strides) {
+  return kernel_height == 1 && kernel_width == 1 && HasNoConvPadding(pads) &&
+         strides[0] == 1 && strides[1] == 1;
+}
+
+// same-size reshape: the whole window folds into a single matmul row per batch element,
+// so the kernel must cover the entire input and leave exactly one output position.
+// Only valid for channels-last (the caller checks that).
+template <typename PadContainer>
+bool IsConvSameSizeMatMul(int64_t input_height, int64_t input_width,
+                          int64_t kernel_height, int64_t kernel_width,
+                          int64_t output_height, int64_t output_width,
+                          const PadContainer& pads) {
+  return input_height == kernel_height && input_width == kernel_width &&
+         output_height == 1 && output_width == 1 && HasNoConvPadding(pads);
+}
+
 template <bool is_channels_last, bool is_fused>
 class Conv : public WebGpuKernel {
  public:
+#if !defined(__wasm__)
+  // Abstract base class for alternative optimized Conv implementations (currently the
+  // subgroup-matrix 1x1 / same-size path). Implementations hold the parent Conv and read
+  // conv_attrs/activation/... from it, mirroring MatMul::MatMulOptImpl.
+  class ConvOptImpl {
+   public:
+    explicit ConvOptImpl(const Conv& parent) : parent_(parent) {}
+    virtual ~ConvOptImpl() = default;
+
+    // Attempts the optimized path, reading the Conv operands from `context` and the Conv
+    // attributes from the parent. Sets handled=true when it ran; leaves handled=false
+    // (allocating nothing) so the caller falls back to the normal Conv path.
+    virtual Status Compute(ComputeContext& context, /*out*/ bool& handled) = 0;
+
+   protected:
+    const Conv& parent_;
+  };
+#endif
+
   Conv(const OpKernelInfo& info) : WebGpuKernel(info), conv_attrs_(info) {
     if (is_fused) {
       ORT_ENFORCE(GetFusedActivationAttr(info, activation_).IsOK());
     }
+    // Whether the weight input (index 1) is a constant initializer. The 1x1 Conv
+    // subgroup-matrix path uses this to decide it can pad an odd-N weight once and
+    // cache the result; a non-constant weight must not be cached.
+    const Tensor* w = nullptr;
+    w_is_constant_ = info.TryGetConstantInput(1, &w);
   }
   Status ComputeInternal(ComputeContext& context) const override;
 
@@ -29,10 +98,27 @@ class Conv : public WebGpuKernel {
                          AllocatorPtr alloc,
                          /*out*/ bool& is_packed) override;
 
+  // State an optimized implementation (ConvOptImpl) reads from its parent Conv.
+  const ConvAttributes& ConvAttrs() const { return conv_attrs_; }
+  const Activation& ConvActivation() const { return activation_; }
+  // True when input 1 (the weight) is a constant initializer. See w_is_constant_.
+  bool IsWeightConstant() const { return w_is_constant_; }
+  // The prepacked (OIHW->HWIO transposed) weight, or null when it is not prepacked.
+  const Tensor* PrepackedKernel() const { return transposed_kernel_.get(); }
+
  protected:
   ConvAttributes conv_attrs_;
   Activation activation_;
   std::unique_ptr<Tensor> transposed_kernel_;  // should only have value when `is_initializer` AND `is_4D` AND `is_NHWC`
+  bool w_is_constant_ = false;                 // whether the weight input (index 1) is a constant initializer
+
+  // Optional subgroup-matrix implementation for the 1x1 Conv path, lazily created
+  // on the first Compute call (once device capabilities can be queried). Null after
+  // initialization means the device has no subgroup-matrix support.
+#if !defined(__wasm__)
+  mutable std::unique_ptr<ConvOptImpl> impl_;
+  mutable std::once_flag impl_init_flag_;
+#endif
 };
 
 Status TransposeKernel(ComputeContext& context, const Tensor* kernel, const TensorShape& kernel_shape, Tensor* transposed_kernel, const InlinedVector<size_t>& perm);