[webgpu] Apply template to MatMulNBitsWideTile (#25353)

### Description
This commit applies WGSL template to `MatMulNBitsWideTile` to improve
code readability and enables more flexible data handling.

As part of this change, support for 4-bit and 8-bit shaders has been
consolidated, and a common `CEIL_DIV` utility has been introduced. The
previous `ShaderUsage::UseUniform` and
`ShaderUsage::UseIndicesTypeAlias` flags are no longer necessary and
have been removed.

### Motivation and Context
See above

Author: daijh

cmake/onnxruntime_providers_webgpu.cmake

@@ -172,12 +172,24 @@
     file(MAKE_DIRECTORY ${WGSL_GENERATED_DIR})
 
     # Find all WGSL template input files
-    file(GLOB_RECURSE WGSL_TEMPLATE_FILES
-      "${ONNXRUNTIME_ROOT}/core/providers/webgpu/*.wgsl.template"
-      "${ONNXRUNTIME_ROOT}/contrib_ops/webgpu/*.wgsl.template")
+    set(WGSL_SEARCH_PATHS "${ONNXRUNTIME_ROOT}/core/providers/webgpu/*.wgsl.template")
+    if(NOT onnxruntime_DISABLE_CONTRIB_OPS)
+        list(APPEND WGSL_SEARCH_PATHS "${ONNXRUNTIME_ROOT}/contrib_ops/webgpu/*.wgsl.template")
+    endif()
+    file(GLOB_RECURSE WGSL_TEMPLATE_FILES ${WGSL_SEARCH_PATHS})
 
     # Set wgsl-gen command line options as a list
-    set(WGSL_GEN_OPTIONS "-i" "${ONNXRUNTIME_ROOT}/core/providers/webgpu/" "-i" "${ONNXRUNTIME_ROOT}/contrib_ops/webgpu/" "--output" "${WGSL_GENERATED_DIR}" "-I" "wgsl_template_gen/" "--preserve-code-ref" "--verbose")
+    set(WGSL_GEN_OPTIONS
+        "--output" "${WGSL_GENERATED_DIR}"
+        "-I" "wgsl_template_gen/"
+        "--preserve-code-ref"
+        "--verbose"
+        "-i" "${ONNXRUNTIME_ROOT}/core/providers/webgpu"
+    )
+    if(NOT onnxruntime_DISABLE_CONTRIB_OPS)
+        list(APPEND WGSL_GEN_OPTIONS "-i" "${ONNXRUNTIME_ROOT}/contrib_ops/webgpu")
+    endif()
+
     if (onnxruntime_WGSL_TEMPLATE STREQUAL "static")
       if (CMAKE_BUILD_TYPE STREQUAL "Debug")
         list(APPEND WGSL_GEN_OPTIONS "--generator" "static-cpp-literal")

onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc

@@ -22,6 +22,11 @@ namespace {
 
 constexpr unsigned int kMinMForTileOptimization = 4;
 
+template <typename T>
+inline T ceil_div(T numerator, T denominator) {
+  return (numerator + denominator - 1) / denominator;
+}
+
 }  // namespace
 
 ONNX_OPERATOR_KERNEL_EX(
@@ -37,165 +42,24 @@ ONNX_OPERATOR_KERNEL_EX(
     MatMulNBits);
 
 Status MatMulNBitsWideTileProgram::GenerateShaderCode(ShaderHelper& shader) const {
-  const auto& a = shader.AddInput("input_a", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | ShaderUsage::UseValueTypeAlias);
-  const auto& b = shader.AddInput("input_b", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | ShaderUsage::UseValueTypeAlias);
-  shader.AddInput("scales", ShaderUsage::UseUniform);
+  shader.AddInput("input_a", ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
+  shader.AddInput("input_b", ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
+  shader.AddInput("scales", ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
   if (has_zero_points_) {
-    shader.AddInput("zero_points", ShaderUsage::UseUniform);
+    shader.AddInput("zero_points", ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
   }
-  const auto& y = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias | ShaderUsage::UseIndicesTypeAlias);
-
-  // Bock size 32, `a` component size 4, 8 `a` components per block.
-  constexpr uint32_t kAComponentsForBlock32 = 8;
+  shader.AddOutput("output", ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
 
   const uint32_t workgroup_size = WorkgroupSizeX() * WorkgroupSizeY();
   ORT_ENFORCE(tile_m_ == workgroup_size / 8, "tile_m must be workgroup_size / 8.");
   ORT_ENFORCE(tile_n_ == workgroup_size, "tile_n must be workgroup_size.");
+  ORT_ENFORCE(nbits_ == 4 || nbits_ == 8, "Only 4/8 bits are supported for webgpu matmulnbits.");
 
-  // memory read/write helpers
-  shader.AdditionalImplementation() << "fn mm_read_a(batch : u32, row : u32, col : u32) -> input_a_value_t {\n"
-                                    << "  if (batch < uniforms.input_a_shape[0] && row < uniforms.input_a_shape[1] && col < uniforms.input_a_shape[2]) {\n"
-                                    << "    return " << a.GetByIndices("input_a_indices_t(batch, row, col)") << ";\n"
-                                    << "  }\n"
-                                    << "  return input_a_value_t(0);\n"
-                                    << "}\n";
-  if (nbits_ == 4) {
-    shader.AdditionalImplementation() << "\n"
-                                      << "fn mm_read_b(row : u32, col : u32) -> input_b_value_t {\n"
-                                      << "  if (row < uniforms.input_b_shape[0] && col < uniforms.input_b_shape[1]) {\n"
-                                      << "    return " << b.GetByIndices("input_b_indices_t(row, col, 0)") << ";\n"
-                                      << "  }\n"
-                                      << "  return input_b_value_t(0);\n"
-                                      << "}\n";
-
-    shader.AdditionalImplementation() << R"(
-fn dequantize_packed8xU4(packed_value : u32, zero_point : output_element_t, scale : output_element_t) -> mat2x4<output_element_t> {
-  let lower_values: vec4<u32> = unpack4xU8(packed_value & 0x0F0F0F0Fu);
-  let upper_values: vec4<u32> = unpack4xU8((packed_value >> 4u) & 0x0F0F0F0Fu);
-
-  let zero_matrix: mat2x4<output_element_t> = mat2x4<output_element_t>(
-      zero_point, zero_point, zero_point, zero_point,
-      zero_point, zero_point, zero_point, zero_point
-  );
-
-  var dequantized_values: mat2x4<output_element_t> = mat2x4<output_element_t>(
-      output_element_t(lower_values[0]), output_element_t(upper_values[0]),
-      output_element_t(lower_values[1]), output_element_t(upper_values[1]),
-      output_element_t(lower_values[2]), output_element_t(upper_values[2]),
-      output_element_t(lower_values[3]), output_element_t(upper_values[3])
-  );
-
-  dequantized_values = (dequantized_values - zero_matrix) * scale;
-  return dequantized_values;
-}
-)";
-  }
-
-  shader.AdditionalImplementation() << "\n"
-                                    << "fn mm_read_scale(row : u32, col : u32) -> output_element_t {\n"
-                                    << "  if (row < uniforms.input_b_shape[0] && col < uniforms.input_b_shape[1]) {\n"
-                                    << "    return scales[row * uniforms.input_b_shape[1] + col];\n"
-                                    << "  }\n"
-                                    << "  return output_element_t(0);\n"
-                                    << "}\n"
-                                    << GenerateZeroPointReadingCode(nbits_, has_zero_points_);
-
-  shader.AdditionalImplementation() << "\n"
-                                    << "fn mm_write_y(batch : u32, row : u32, col : u32, value : output_value_t) {\n"
-                                    << "  if (row < uniforms.output_shape[1] && col < uniforms.output_shape[2]) {\n"
-                                    << "    " << y.SetByIndices("output_indices_t(batch, row, col)", "value") << "\n"
-                                    << "  }\n"
-                                    << "}\n";
-
-  // declare const variables
-  shader.AdditionalImplementation() << "\n"
-                                    << "// A block32 containing 8 components of `a`." << "\n"
-                                    << "const kAComponentsForBlock32 = " << kAComponentsForBlock32 << "u;\n"
-                                    << "const kTileM = " << tile_m_ << "u;\n"
-                                    << "const kTileN = " << tile_n_ << "u;\n";
-
-  // declare workgroup memory
-  shader.AdditionalImplementation() << "\n"
-                                    << "var<workgroup> a_data_tile: array<array<input_a_value_t, kAComponentsForBlock32>, kTileM>;\n"
-                                    << "\n";
-
-  // main
-  shader.MainFunctionBody() << R"MAIN_FN(
-  let batch = workgroup_idx / (uniforms.num_M_tile * uniforms.num_N_tile);
-  let row = ((workgroup_idx / uniforms.num_N_tile) % uniforms.num_M_tile) * kTileM;
-  let col = (workgroup_idx % uniforms.num_N_tile) * kTileN;
-
-  let a_elements_per_col = uniforms.input_a_shape[2];
-  let a_blocks_per_col = (a_elements_per_col + kAComponentsForBlock32 - 1) / kAComponentsForBlock32;
-
-  // Utilizing an f32 accumulator mitigated precision loss with minimal
-  // performance impact compared to an f16 accumulator.
-  var results : array<f32, kTileM>;
-  for (var a_block_idx = 0u; a_block_idx < a_blocks_per_col; a_block_idx++) {
-    // Load `a` elements into workgroup memory, TileM x kAComponentsForBlock32 (block32)
-    let a_row_idx = local_idx / kAComponentsForBlock32;
-    let a_col_idx = local_idx % kAComponentsForBlock32;
-    a_data_tile[a_row_idx][a_col_idx] = mm_read_a(batch, row + a_row_idx, a_block_idx * kAComponentsForBlock32 + a_col_idx);
-    workgroupBarrier();
-
-    let b_row = col + local_idx;
-    let b_col = a_block_idx;
-
-    let scale = mm_read_scale(b_row, b_col);
-    let zero_point = mm_read_zero(b_row, b_col, uniforms.input_b_shape[0], uniforms.zero_blocks_per_col);
-)MAIN_FN";
-
-  if (nbits_ == 4) {
-    shader.MainFunctionBody() << R"MAIN_FN(
-    let b_data = mm_read_b(b_row, b_col);
-    // `b` component size is 4.
-    for (var b_idx = 0u; b_idx < 4u; b_idx++) {
-      let b_dequantized = dequantize_packed8xU4(b_data[b_idx], zero_point, scale);
-      for (var m_idx = 0u; m_idx < kTileM; m_idx++) {
-        let a_data0 = a_data_tile[m_idx][b_idx * 2u];
-        let a_data1 = a_data_tile[m_idx][b_idx * 2u + 1u];
-
-        results[m_idx] += f32(dot(a_data0, b_dequantized[0])) + f32(dot(a_data1, b_dequantized[1]));
-      }
-    }
-)MAIN_FN";
-  } else {
-    shader.MainFunctionBody() << "    var b_data0 = vec4<u32>(0);\n"
-                                 "    var b_data1 = vec4<u32>(0);\n"
-                                 "    if (b_row < uniforms.input_b_shape[0] && b_col < uniforms.input_b_shape[1]) {\n"
-                              << "      b_data0 = " << b.GetByIndices("input_b_indices_t(b_row, b_col, 0)") << ";\n"
-                              << "      b_data1 = " << b.GetByIndices("input_b_indices_t(b_row, b_col, 1)") << ";\n"
-                                                                                                               "    }"
-                              << R"MAIN_FN(
-    for (var b_idx = 0u; b_idx < 4u; b_idx++) {
-      let b_dequantized0 = (vec4<output_element_t>(unpack4xU8(b_data0[b_idx])) - vec4<output_element_t>(zero_point)) * scale;
-      let b_dequantized1 = (vec4<output_element_t>(unpack4xU8(b_data1[b_idx])) - vec4<output_element_t>(zero_point)) * scale;
-      for (var m_idx = 0u; m_idx < kTileM; m_idx++) {
-        let a_data0 = a_data_tile[m_idx][b_idx];
-        let a_data1 = a_data_tile[m_idx][b_idx + 4u];
-
-        results[m_idx] += f32(dot(a_data0, b_dequantized0)) + f32(dot(a_data1, b_dequantized1));
-      }
-    }
-)MAIN_FN";
-  }
-
-  shader.MainFunctionBody() << R"MAIN_FN(
-
-    workgroupBarrier();
-  }
-
-  if (batch >= uniforms.input_a_shape[0]) {
-    return;
-  }
-
-  // Write the results.
-  for (var m_idx = 0u; m_idx < kTileM; m_idx++) {
-    mm_write_y(batch, row + m_idx, col + local_idx, output_value_t(results[m_idx]));
-  }
-)MAIN_FN";
-
-  return Status::OK();
+  return WGSL_TEMPLATE_APPLY(shader, "quantization/matmul_nbits_wide_tile.wgsl.template",
+                             WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_),
+                             WGSL_TEMPLATE_PARAMETER(nbits, nbits_),
+                             WGSL_TEMPLATE_PARAMETER(tile_m, tile_m_),
+                             WGSL_TEMPLATE_PARAMETER(tile_n, tile_n_));
 }
 
 // Apply similar idea with DP4AMatMulNBitsSmallMProgram algorithm.
@@ -408,38 +272,55 @@ Status MatMulNBits::ComputeInternal(onnxruntime::webgpu::ComputeContext& context
 
   // WideTileProgram
   // This program is optimized for Block32 prefill using Tile16x128.
-  const bool use_wide_tile_program = block_size == 32 && components_a == 4 && components_b == 4 && M >= kMinMForTileOptimization;
+  const bool use_wide_tile_program = block_size == 32 &&
+                                     components_a == 4 &&
+                                     components_b == 4 &&
+                                     M >= kMinMForTileOptimization;
   if (use_wide_tile_program) {
     // Enforce output components to 1.
     components = 1;
 
     constexpr uint32_t workgroup_size = 128;
     constexpr uint32_t tile_m = workgroup_size / 8;
     constexpr uint32_t tile_n = workgroup_size;
-    uint32_t num_N_tile = (N + tile_n - 1) / tile_n;
-    uint32_t num_M_tile = (M + tile_m - 1) / tile_m;
+    const uint32_t num_N_tile = ceil_div(N, tile_n);
+    const uint32_t num_M_tile = ceil_div(M, tile_m);
 
     MatMulNBitsWideTileProgram program{has_zero_points, tile_m, tile_n, nbits};
     program.SetWorkgroupSize(workgroup_size);
-    program.SetDispatchGroupSize((N + tile_n - 1) / tile_n,
-                                 (M + tile_m - 1) / tile_m,
-                                 batch_count);
-    program.CacheHint("Tile" + std::to_string(tile_m) + "x" + std::to_string(tile_n) + "_Block32");
-
-    TensorShape reshaped_a_shape{batch_count, M, K / components_a};
-    TensorShape reshaped_b_shape{N, n_blocks_per_col, blob_size_in_words / components_b};
-    TensorShape reshaped_y_shape{batch_count, M, N / components};
-
-    program
-        .AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, reshaped_a_shape, onnxruntime::narrow<int>(components_a)},
-                    {b, ProgramTensorMetadataDependency::TypeAndRank, reshaped_b_shape, onnxruntime::narrow<int>(components_b * 4)},
-                    {scales, ProgramTensorMetadataDependency::None}})
-        .AddOutput({y, ProgramTensorMetadataDependency::TypeAndRank, reshaped_y_shape, onnxruntime::narrow<int>(components)})
-        .AddUniformVariables({{block_size}, {zero_blocks_per_col}, {num_N_tile}, {num_M_tile}})
-        .CacheHint(nbits, has_zero_points);
+    program.SetDispatchGroupSize(num_N_tile, num_M_tile, batch_count);
+
+    constexpr uint32_t kU32Components = 4;
+    const uint32_t components_b_with_u32 = components_b * kU32Components;
+    const uint32_t K_of_b = n_blocks_per_col * blob_size / components_b_with_u32;
+    const uint32_t K_of_a = K / components_a;
+
+    program.AddInput({a,
+                      ProgramTensorMetadataDependency::TypeAndRank,
+                      onnxruntime::narrow<int>(components_a)});
+    program.AddInput({b,
+                      ProgramTensorMetadataDependency::TypeAndRank,
+                      onnxruntime::narrow<int>(components_b_with_u32)});
+    program.AddInput({scales, ProgramTensorMetadataDependency::TypeAndRank});
     if (has_zero_points) {
-      program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4});
+      program.AddInput({zero_points,
+                        ProgramTensorMetadataDependency::TypeAndRank,
+                        {ceil_div(zero_points->Shape().Size(), static_cast<int64_t>(4))},
+                        4});
     }
+    program.AddOutput({y,
+                       ProgramTensorMetadataDependency::TypeAndRank,
+                       onnxruntime::narrow<int>(components)});
+    program.AddUniformVariables({{batch_count},
+                                 {M},
+                                 {N},
+                                 {K_of_a},
+                                 {K_of_b},
+                                 {n_blocks_per_col},
+                                 {zero_blocks_per_col},
+                                 {num_N_tile},
+                                 {num_M_tile}});
+    program.CacheHint(nbits, has_zero_points);
 
     return context.RunProgram(program);
   }

onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h

@@ -15,10 +15,15 @@ using namespace onnxruntime::webgpu;
 class MatMulNBitsWideTileProgram final : public Program<MatMulNBitsWideTileProgram> {
  public:
   MatMulNBitsWideTileProgram(bool has_zero_points, uint32_t tile_m, uint32_t tile_n, uint32_t nbits)
-      : Program{"MatMulNBitsWideTileProgram"}, has_zero_points_{has_zero_points}, tile_m_(tile_m), tile_n_(tile_n), nbits_(nbits) {}
+      : Program{"MatMulNBitsWideTile"}, has_zero_points_{has_zero_points}, tile_m_(tile_m), tile_n_(tile_n), nbits_(nbits) {}
 
   Status GenerateShaderCode(ShaderHelper& sh) const override;
-  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"block_size", ProgramUniformVariableDataType::Uint32},
+  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"Batch", ProgramUniformVariableDataType::Uint32},
+                                          {"M", ProgramUniformVariableDataType::Uint32},
+                                          {"N", ProgramUniformVariableDataType::Uint32},
+                                          {"K_of_a", ProgramUniformVariableDataType::Uint32},
+                                          {"K_of_b", ProgramUniformVariableDataType::Uint32},
+                                          {"n_blocks_per_col", ProgramUniformVariableDataType::Uint32},
                                           {"zero_blocks_per_col", ProgramUniformVariableDataType::Uint32},
                                           {"num_N_tile", ProgramUniformVariableDataType::Uint32},
                                           {"num_M_tile", ProgramUniformVariableDataType::Uint32});