[WebGPU] Use shared memory tiles in subgroup matrix MatMul on NVIDIA

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

@@ -3,6 +3,7 @@
 
 #include "core/providers/webgpu/math/subgroup_matrix_tiling.h"
 
+#include <algorithm>
 #include <cstdint>
 #include <limits>
 #include <memory>
@@ -26,14 +27,18 @@
 class SubgroupMatrixMatMulProgram final : public Program<SubgroupMatrixMatMulProgram> {
  public:
   SubgroupMatrixMatMulProgram(const Activation& activation, bool has_bias, SubgroupMatrixConfig config,
-                              uint32_t sg_mat_count_m, uint32_t sg_mat_count_n, uint32_t split_k)
+                              uint32_t sg_mat_count_m, uint32_t sg_mat_count_n, uint32_t split_k,
+                              uint32_t sg_grid_m, uint32_t sg_grid_n, uint32_t tile_k)
       : Program{"SubgroupMatrixMatMul"},
         activation_(activation),
         has_bias_(has_bias),
         config_(config),
         sg_mat_count_m_(sg_mat_count_m),
         sg_mat_count_n_(sg_mat_count_n),
-        split_k_(split_k) {}
+        split_k_(split_k),
+        sg_grid_m_(sg_grid_m),
+        sg_grid_n_(sg_grid_n),
+        tile_k_(tile_k) {}
   Status GenerateShaderCode(ShaderHelper& sh) const override;
   WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"M", ProgramUniformVariableDataType::Uint32},
                                           {"N", ProgramUniformVariableDataType::Uint32},
@@ -49,6 +54,9 @@
   const uint32_t sg_mat_count_m_;
   const uint32_t sg_mat_count_n_;
   const uint32_t split_k_;
+  const uint32_t sg_grid_m_;
+  const uint32_t sg_grid_n_;
+  const uint32_t tile_k_;
 };
 
 namespace {
@@ -71,8 +79,28 @@
                                           {"N_b", ProgramUniformVariableDataType::Uint32});
 };
 
+// A staged tiling reads A and B as vec4 so N must be a multiple of 4 and the template
+// unrolls at most 2x2 subgroup matrices per subgroup.
+bool IsStagedTilingSupported(const wgpu::Limits& limits, const SubgroupMatrixConfig& config,
+                             const SubgroupMatrixTiling& tiling, uint32_t N) {
+  const uint32_t cell_m = config.M * tiling.sg_grid_m;
+  const uint32_t cell_n = config.N * tiling.sg_grid_n;
+  if (N % 4 != 0 || tiling.split_k != 1 || cell_m == 0 || cell_n == 0 || tiling.tile_k == 0 ||
+      tiling.tile_k % config.K != 0 || tiling.tile_m % cell_m != 0 || tiling.tile_n % cell_n != 0 ||
+      tiling.tile_m / cell_m > 2 || tiling.tile_n / cell_n > 2) {
+    return false;
+  }
+  const uint64_t threads = uint64_t{config.subgroupSize} * tiling.sg_grid_m * tiling.sg_grid_n;
+  // Matches the scratch size in the template where each staged row is padded by 8 elements.
+  const uint64_t staged_elems = uint64_t{tiling.tile_m} * (tiling.tile_k + 8) + uint64_t{tiling.tile_k} * (tiling.tile_n + 8);
+  const uint64_t scratch_bytes = sizeof(uint16_t) * std::max(staged_elems, uint64_t{tiling.tile_m} * tiling.tile_n);
+  return threads <= limits.maxComputeInvocationsPerWorkgroup && threads <= limits.maxComputeWorkgroupSizeX &&
+         scratch_bytes <= limits.maxComputeWorkgroupStorageSize;
+}
+
 // Subgroup-matrix MatMul implementation. Loads both A and B directly from global
-// memory and runs the subgroup-matrix kernel during Compute. The class is
+// memory (or stages them in shared memory when the tiling sets tile_k) and runs the
+// subgroup-matrix kernel during Compute. The class is
 // intended to support all subgroup-matrix configs. For now 8x16x16 and 16x16x16
 // are implemented. The per-problem output tiling is supplied by a vendor-specific
 // selector kept internal to this impl.
@@ -135,6 +163,9 @@
     }
 
     const auto& config = config_;
+    if (tiling->tile_k > 0) {
+      return config.K != 0 && K % config.K == 0 && IsStagedTilingSupported(context.DeviceLimits(), config, *tiling, N);
+    }
     const uint32_t n_align = config.Is(16, 16, 16) ? 8 : 2;
     const bool needs_padded_b = N % n_align != 0;
     if (needs_padded_b && (!b_is_constant || N > std::numeric_limits<uint32_t>::max() - n_align)) {
@@ -194,10 +225,15 @@
                       "MatMul algorithm subgroup_matrix has no tiling for these inputs.");
 
     const auto& config = config_;
+    // Staged tilings read B through shared memory so they never need the padded copy.
+    const bool staged = tiling->tile_k > 0;
+    ORT_RETURN_IF_NOT(staged ? IsStagedTilingSupported(context.DeviceLimits(), config, *tiling, N)
+                             : tiling->sg_grid_m == 1 && tiling->sg_grid_n == 1,
+                      "MatMul algorithm subgroup_matrix got a tiling this device cannot run.");
     // Vulkan cooperative matrix loads need 16 byte aligned rows so the
     // 16x16x16 config pads B to a multiple of 8 columns.
     const uint32_t n_align = config.Is(16, 16, 16) ? 8 : 2;
-    const bool needs_padded_b = N % n_align != 0;
+    const bool needs_padded_b = !staged && N % n_align != 0;
     // Require whole subgroup-matrix K blocks. An unaligned B must be constant
     // because its padded copy is cached by this implementation.
     ORT_RETURN_IF_NOT(config.K != 0 && K % config.K == 0 &&
@@ -215,7 +251,7 @@
 
     // The kernel keeps its operand loads in bounds by shifting a trailing partial
     // tile back, which is only possible when the tile fits within M and N.
-    ORT_ENFORCE(M >= tiling->tile_m && N_b >= tiling->tile_n);
+    ORT_ENFORCE(staged || (M >= tiling->tile_m && N_b >= tiling->tile_n));
 
     // The optimized path will run: now materialize the aligned B for unaligned N.
     const Tensor* b_used = b;
@@ -229,17 +265,21 @@
     const uint32_t tile_m = tiling->tile_m;
     const uint32_t tile_n = tiling->tile_n;
     const uint32_t split_k = tiling->split_k;
-    const uint32_t sg_mat_count_m = tile_m / config.M;
-    const uint32_t sg_mat_count_n = tile_n / config.N;
-    ORT_ENFORCE(tile_m % config.M == 0 && tile_n % config.N == 0,
+    const uint32_t sg_grid_m = tiling->sg_grid_m;
+    const uint32_t sg_grid_n = tiling->sg_grid_n;
+    const uint32_t tile_k = tiling->tile_k;
+    const uint32_t sg_mat_count_m = tile_m / (config.M * sg_grid_m);
+    const uint32_t sg_mat_count_n = tile_n / (config.N * sg_grid_n);
+    ORT_ENFORCE(tile_m % (config.M * sg_grid_m) == 0 && tile_n % (config.N * sg_grid_n) == 0,
                 "Tiling must be a multiple of the subgroup-matrix shape: ",
                 tile_m, "x", tile_n, " vs ", config.M, "x", config.N);
     const uint32_t dispatch_x = (N + tile_n - 1) / tile_n;
     const uint32_t dispatch_y = (M + tile_m - 1) / tile_m;
 
     SubgroupMatrixMatMulProgram program{activation, has_bias, config,
-                                        sg_mat_count_m, sg_mat_count_n, split_k};
-    program.SetWorkgroupSize(config.subgroupSize * split_k);
+                                        sg_mat_count_m, sg_mat_count_n, split_k,
+                                        sg_grid_m, sg_grid_n, tile_k};
+    program.SetWorkgroupSize(config.subgroupSize * split_k * sg_grid_m * sg_grid_n);
     if (context.HasFeature(wgpu::FeatureName::SubgroupSizeControl)) {
       program.SetSubgroupSize(config.subgroupSize);
     }
@@ -248,9 +288,9 @@
                       static_cast<uint32_t>(config.componentType),
                       static_cast<uint32_t>(config.resultComponentType),
                       config.M, config.N, config.K, config.subgroupSize,
-                      sg_mat_count_m, sg_mat_count_n, split_k)
-        .AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, 1},
-                    {b_used, ProgramTensorMetadataDependency::TypeAndRank, 1}})
+                      sg_mat_count_m, sg_mat_count_n, split_k, sg_grid_m, sg_grid_n, tile_k)
+        .AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, staged ? 4 : 1},
+                    {b_used, ProgramTensorMetadataDependency::TypeAndRank, staged ? 4 : 1}})
         .AddOutput({output, ProgramTensorMetadataDependency::Rank, 1})
         .AddUniformVariables({{M}, {N}, {K}, {dispatch_x}, {N_b}});
     // Activation uniforms must remain last because definitions and values are matched by index.
@@ -322,14 +362,18 @@
 
 Status GenerateSubgroupMatrixShaderCode(ShaderHelper& shader, const SubgroupMatrixConfig& config,
                                         uint32_t sg_mat_count_m, uint32_t sg_mat_count_n,
-                                        uint32_t split_k) {
+                                        uint32_t split_k, uint32_t sg_grid_m, uint32_t sg_grid_n,
+                                        uint32_t tile_k) {
   return WGSL_TEMPLATE_APPLY(shader, "math/subgroup_matrix_matmul.wgsl.template",
+                             WGSL_TEMPLATE_PARAMETER(sg_grid_m, sg_grid_m),
+                             WGSL_TEMPLATE_PARAMETER(sg_grid_n, sg_grid_n),
                              WGSL_TEMPLATE_PARAMETER(sg_mat_count_m, sg_mat_count_m),
                              WGSL_TEMPLATE_PARAMETER(sg_mat_count_n, sg_mat_count_n),
                              WGSL_TEMPLATE_PARAMETER(sg_mat_k, config.K),
                              WGSL_TEMPLATE_PARAMETER(sg_mat_m, config.M),
                              WGSL_TEMPLATE_PARAMETER(sg_mat_n, config.N),
-                             WGSL_TEMPLATE_PARAMETER(split_k, split_k));
+                             WGSL_TEMPLATE_PARAMETER(split_k, split_k),
+                             WGSL_TEMPLATE_PARAMETER(tile_k, tile_k));
 }
 
 // Default tiling used on any vendor without a specialized policy: a 64x64 output
@@ -345,6 +389,30 @@
   };
 }
 
+// A 4x4 grid of subgroups stages K slices 64 deep and each subgroup accumulates 2x2 subgroup
+// matrices when the problem has at least 128 rows and 256 tiles of 64x64 and 1x1 otherwise.
+// Staging starts at 64 rows and below that only for outputs at least 8192 wide because the
+// default split K tiling fills the GPU better on small problems.
+SubgroupMatrixTilingSelector MakeStagedTilingSelector(const SubgroupMatrixConfig& config) {
+  return [config, fallback = MakeDefaultTilingSelector()](
+             const ComputeContext& context, uint32_t M, uint32_t N,
+             uint32_t K, uint32_t batch) -> std::optional<SubgroupMatrixTiling> {
+    if ((M >= 64 || N >= 8192) && K % config.K == 0) {
+      const auto& limits = context.DeviceLimits();
+      const bool large = M >= 128 && uint64_t{(M + 63) / 64} * ((N + 63) / 64) >= 256;
+      const SubgroupMatrixTiling big{config.M * 8, config.N * 8, 1, 4, 4, 64};
+      const SubgroupMatrixTiling small{config.M * 4, config.N * 4, 1, 4, 4, 64};
+      if (large && IsStagedTilingSupported(limits, config, big, N)) {
+        return big;
+      }
+      if (IsStagedTilingSupported(limits, config, small, N)) {
+        return small;
+      }
+    }
+    return fallback(context, M, N, K, batch);
+  };
+}
+
 }  // namespace
 
 Status SubgroupMatrixMatMulProgram::GenerateShaderCode(ShaderHelper& shader) const {
@@ -369,7 +437,8 @@
       << "}\n";
 
   if (config_.Is(8, 16, 16) || config_.Is(16, 16, 16)) {
-    return GenerateSubgroupMatrixShaderCode(shader, config_, sg_mat_count_m_, sg_mat_count_n_, split_k_);
+    return GenerateSubgroupMatrixShaderCode(shader, config_, sg_mat_count_m_, sg_mat_count_n_, split_k_,
+                                            sg_grid_m_, sg_grid_n_, tile_k_);
   }
   return Status(onnxruntime::common::ONNXRUNTIME, onnxruntime::common::NOT_IMPLEMENTED,
                 "Unsupported subgroup matrix config dimensions.");
@@ -384,11 +453,13 @@
   if (!config) {
     return nullptr;
   }
-  // Intel GPUs use a tuned/heuristic tiling policy; every other vendor falls back
+  // Intel and NVIDIA GPUs use a tuned/heuristic tiling policy; every other vendor falls back
   // to a default tiling.
-  const bool is_intel = context.AdapterInfo().vendor == std::string_view{"intel"};
+  const auto vendor = context.AdapterInfo().vendor;
   SubgroupMatrixTilingSelector tiling_selector =
-      is_intel ? intel::CreateSubgroupMatrixTilingSelector(context) : MakeDefaultTilingSelector();
+      vendor == std::string_view{"intel"}    ? intel::CreateSubgroupMatrixTilingSelector(context)
+      : vendor == std::string_view{"nvidia"} ? MakeStagedTilingSelector(*config)
+                                             : MakeDefaultTilingSelector();
   if (!tiling_selector) {
     return nullptr;
   }