Add Intel WebGPU subgroup-matrix MatMul implementation

onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc

@@ -5,6 +5,7 @@
 
 #include "contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h"
 #include "contrib_ops/webgpu/quantization/matmul_nbits_common.h"
+#include "core/providers/webgpu/vendor/intel/intel_device_info.h"
 
 namespace onnxruntime {
 namespace contrib {
@@ -299,13 +300,8 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te
   // For large M on Intel Xe, cap dispatch_y so each workgroup processes multiple
   // M-tiles sequentially, reducing scheduling overhead.
   if (M > 2048 && context.AdapterInfo().vendor == std::string_view{"intel"}) {
-    // Each XeCore has 4 XVE x 8 SIMD-32 hardware threads = 32 subgroups.
-    uint32_t hw_subgroups = 0;
-    if (context.AdapterInfo().architecture == std::string_view{"xe-3lpg"}) {
-      hw_subgroups = 384;  // 12 XeCore x 32
-    } else if (context.AdapterInfo().architecture == std::string_view{"xe-2lpg"}) {
-      hw_subgroups = 256;  // 8 XeCore x 32
-    }
+    const uint32_t hw_subgroups =
+        ::onnxruntime::webgpu::intel::HwSubgroups(std::string_view{context.AdapterInfo().architecture});
     if (hw_subgroups > 0) {
       constexpr uint32_t kOccupancyFactor = 16;  // empirically tuned on Xe2/Xe3 devices
       uint32_t target_wgs = hw_subgroups * kOccupancyFactor / (work_group_size / 32);

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

@@ -13,6 +13,9 @@
 #include "core/providers/webgpu/data_transfer.h"
 #include "core/providers/webgpu/vendor/intel/math/matmul.h"
 #include "core/providers/webgpu/webgpu_utils.h"
+#if !defined(__wasm__)
+#include "core/providers/webgpu/vendor/intel/math/subgroup_matrix_matmul.h"
+#endif
 
 namespace onnxruntime {
 namespace webgpu {
@@ -120,6 +123,17 @@ Status MatMul::ComputeInternal(ComputeContext& context) const {
   }
   bool has_bias = context.InputCount() > 2;
 
+#if !defined(__wasm__)
+  // Try a vendor-optimized implementation (e.g. Intel subgroup-matrix) first.
+  if (impl_) {
+    bool handled = false;
+    ORT_RETURN_IF_ERROR(impl_->Compute(context, handled));
+    if (handled) {
+      return Status::OK();
+    }
+  }
+#endif
+
   if (helper.N() < 8 && helper.K() < 8) {  // call MatMulNaiveProgram
 
     const uint32_t m = narrow<uint32_t>(helper.M());  // left matrix first dimension
@@ -174,6 +188,26 @@ Status MatMul::ComputeInternal(ComputeContext& context) const {
   return ComputeMatMul(&context, Activation(), inputs, output_tensor);
 }
 
+Status MatMul::PrePackInternal(ComputeContextBase& context,
+                               const Tensor& tensor,
+                               int input_idx,
+                               AllocatorPtr alloc,
+                               /*out*/ bool& is_packed) {
+  is_packed = false;
+#if !defined(__wasm__)
+  // Create the vendor-optimized implementation once based on device capabilities.
+  if (!impl_) {
+    impl_ = intel::CreateSubgroupMatrixMatMulImpl(context, *this);
+  }
+#else
+  ORT_UNUSED_PARAMETER(context);
+#endif
+  ORT_UNUSED_PARAMETER(tensor);
+  ORT_UNUSED_PARAMETER(input_idx);
+  ORT_UNUSED_PARAMETER(alloc);
+  return Status::OK();
+}
+
 Status ComputeMatMul(ComputeContext* context,
                      const Activation& activation, std::vector<const Tensor*>& inputs, Tensor* output_tensor, bool is_channels_last,
                      const TensorShape& input_a_reshape,

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

@@ -3,6 +3,8 @@
 
 #pragma once
 
+#include <memory>
+
 #include "core/providers/webgpu/webgpu_kernel.h"
 #include "core/providers/webgpu/program.h"
 #include "core/providers/cpu/math/matmul_helper.h"
@@ -29,12 +31,41 @@ MatMulFillBiasOrZeroBeforeSplitKProgram CreateMatMulFillBiasOrZeroBeforeSplitKPr
 
 class MatMul final : public WebGpuKernel {
  public:
+  // Abstract base class for alternative optimized MatMul implementations.
+  // Implementations can provide optimized computation paths by deriving from this class.
+  class MatMulOptImpl {
+   public:
+    explicit MatMulOptImpl(const MatMul& parent) : parent_(parent) {}
+    virtual ~MatMulOptImpl() = default;
+
+    // Called during Compute phase to execute implementation-specific computation.
+    // @param context       The WebGPU compute context.
+    // @param handled       Output parameter. Set to true if this implementation handled the computation.
+    // @return Status::OK() on success, or an error status on failure.
+    virtual Status Compute(ComputeContext& context,
+                           /*out*/ bool& handled) = 0;
+
+   protected:
+    const MatMul& parent_;
+  };
+
   MatMul(const OpKernelInfo& info) : WebGpuKernel{info} {}
 
   Status ComputeInternal(ComputeContext& context) const override;
+
+  Status PrePackInternal(ComputeContextBase& context,
+                         const Tensor& tensor,
+                         int input_idx,
+                         AllocatorPtr alloc,
+                         /*out*/ bool& is_packed) override;
+
   constexpr static uint32_t MATMUL_PACKED_WORKGROUP_SIZE_X = 8;
   constexpr static uint32_t MATMUL_PACKED_WORKGROUP_SIZE_Y = 8;
   constexpr static uint32_t MATMUL_PACKED_WORKGROUP_SIZE_Z = 1;
+
+ private:
+  // Alternative optimized implementation (created during PrePack if applicable)
+  mutable std::unique_ptr<MatMulOptImpl> impl_;
 };
 
 class MatMulNaiveProgram final : public Program<MatMulNaiveProgram> {

onnxruntime/core/providers/webgpu/vendor/intel/intel_device_info.cc

@@ -0,0 +1,45 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#if !defined(__wasm__)
+
+#include "core/providers/webgpu/vendor/intel/intel_device_info.h"
+
+#include <algorithm>
+
+namespace onnxruntime {
+namespace webgpu {
+namespace intel {
+
+uint32_t HwSubgroups(std::string_view arch) {
+  if (arch == std::string_view{"xe-3lpg"}) {
+    return 384;  // 12 Xe cores x 32
+  }
+  if (arch == std::string_view{"xe-2lpg"}) {
+    return 256;  // 8 Xe cores x 32
+  }
+  return 0;  // unknown architecture; caller decides the fallback
+}
+
+bool IsSubgroupMatrixConfigSupported(const ComputeContextBase& context, uint32_t m, uint32_t n, uint32_t k) {
+  if (!context.HasFeature(wgpu::FeatureName::ChromiumExperimentalSubgroupMatrix)) {
+    return false;
+  }
+  const wgpu::AdapterInfo& adapter_info = context.AdapterInfo();
+  if (adapter_info.subgroupMinSize != 16 || adapter_info.subgroupMaxSize != 32) {
+    return false;
+  }
+  const wgpu::AdapterPropertiesSubgroupMatrixConfigs& configs = context.SubgroupMatrixConfigs();
+  return std::any_of(configs.configs, configs.configs + configs.configCount,
+                     [m, n, k](const auto& c) {
+                       return c.componentType == wgpu::SubgroupMatrixComponentType::F16 &&
+                              c.resultComponentType == wgpu::SubgroupMatrixComponentType::F16 &&
+                              c.M == m && c.N == n && c.K == k;
+                     });
+}
+
+}  // namespace intel
+}  // namespace webgpu
+}  // namespace onnxruntime
+
+#endif  // !defined(__wasm__)

onnxruntime/core/providers/webgpu/vendor/intel/intel_device_info.h

@@ -0,0 +1,31 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#if !defined(__wasm__)
+
+#include <cstdint>
+#include <string_view>
+
+#include "core/providers/webgpu/compute_context.h"
+
+namespace onnxruntime {
+namespace webgpu {
+namespace intel {
+
+// Approximate number of subgroups an Intel Xe GPU keeps resident at once. Each Xe
+// core runs 4 XVE x 8 SIMD32 hardware threads = 32 subgroups, so the count is
+// (number of Xe cores) x 32. Returns 0 for an unrecognized architecture so each
+// caller can apply its own fallback policy.
+uint32_t HwSubgroups(std::string_view arch);
+
+// Returns true if the device exposes an m x n x k F16/F16 subgroup matrix config
+// within the Intel Xe subgroup size range (min 16, max 32).
+bool IsSubgroupMatrixConfigSupported(const ComputeContextBase& context, uint32_t m, uint32_t n, uint32_t k);
+
+}  // namespace intel
+}  // namespace webgpu
+}  // namespace onnxruntime
+
+#endif  // !defined(__wasm__)