Repository: microsoft/onnxruntime
File: onnxruntime/core/providers/webgpu/math/gemm_utils.cc (C++)

Code before:
  } else {
    shader.AdditionalImplementation() << "    if(row < i32(" << rows << ") && col < i32(" << components_per_row << ")) {\n";
  }

  const std::string input_indices = input_name + "_indices";
  shader.AdditionalImplementation() << "        var " << input_indices << ": " << input_name << "_indices_t" << ";\n";

  if (batch_dims) {
    shader.AdditionalImplementation() << ConvertOutputBatchIndicesToInputBatchIndices(input_name, input, input.Rank() - 2, batch_dims ? batch_dims->Rank() : 0, " batch_indices ") << "\n";
  }

  shader.AdditionalImplementation() << input.IndicesSet(input_indices, input.Rank() - 2, "u32(row)") << "\n"
                                    << input.IndicesSet(input_indices, input.Rank() - 1, "u32(colIn)") << "\n"
                                    << "        value = " << input.GetByIndices(input_indices) << ";\n"
                                    << "    }\n"
                                    << "    return value;\n"
                                    << "}\n\n";
}

void MatMulReadFnSource(ShaderHelper& shader,
                        const ShaderVariableHelper& a,
                        const ShaderVariableHelper& b,
                        const ShaderIndicesHelper* batch_dims,
                        bool transA,
                        bool transB,
                        bool is_vec4) {
  MatMulReadFnSource(shader, "mm_readA", a, "a", batch_dims, "uniforms.dim_a_outer", "uniforms.dim_inner", transA, is_vec4);
  MatMulReadFnSource(shader, "mm_readB", b, "b", batch_dims, "uniforms.dim_inner", "uniforms.dim_b_outer", transB, is_vec4);
}

void MatMulWriteFnSource(ShaderHelper& shader,
                         const ShaderVariableHelper& output,
                         const ShaderVariableHelper* bias,
                         bool is_gemm,
                         int c_components,
                         int output_components,
                         bool c_is_scalar,
                         std::string activation_snippet,
                         bool is_channels_last,
                         bool use_split_k,
                         ProgramVariableDataType output_variable_type) {
  shader.AdditionalImplementation()
      << "fn mm_write(batch: i32, row: i32, colIn: i32, valueIn: output_value_t) { \n";

  shader.AdditionalImplementation() << "  let col = colIn * " << output_components << ";\n";

  shader.AdditionalImplementation() << "if(row < i32(uniforms.dim_a_outer) && col < i32(uniforms.dim_b_outer)) { \n"
                                    << "    var value = valueIn; \n";

  if (use_split_k) {
    // Set output when MatMul is performed with Split-K.

Review comment (performance):
nit:
```
int components = is_vec4 ? 4 : 1;
MatMulReadFnSource(shader, "mm_readA", a, "a", batch_dims, "uniforms.dim_a_outer", "uniforms.dim_inner", transA, components );
MatMulReadFnSource(shader, "mm_readB", b, "b", batch_dims, "uniforms.dim_inner", "uniforms.dim_b_outer", transB, components );

Code after:
    shader.AdditionalImplementation() << "    if(row < i32(" << rows << ") && col < i32(" << components_per_row << ")) {\n";
  }

  const std::string input_indices = input_name + "_indices";
  shader.AdditionalImplementation() << "        var " << input_indices << ": " << input_name << "_indices_t" << ";\n";

  if (batch_dims) {
    shader.AdditionalImplementation() << ConvertOutputBatchIndicesToInputBatchIndices(input_name, input, input.Rank() - 2, batch_dims ? batch_dims->Rank() : 0, " batch_indices ") << "\n";
  }

  shader.AdditionalImplementation() << input.IndicesSet(input_indices, input.Rank() - 2, "u32(row)") << "\n"
                                    << input.IndicesSet(input_indices, input.Rank() - 1, "u32(colIn)") << "\n"
                                    << "        value = " << input.GetByIndices(input_indices) << ";\n"
                                    << "    }\n"
                                    << "    return value;\n"
                                    << "}\n\n";
}

void MatMulReadFnSource(ShaderHelper& shader,
                        const ShaderVariableHelper& a,
                        const ShaderVariableHelper& b,
                        const ShaderIndicesHelper* batch_dims,
                        bool transA,
                        bool transB) {
  MatMulReadFnSource(shader, "mm_readA", a, "a", batch_dims, "uniforms.dim_a_outer", "uniforms.dim_inner", transA);
  MatMulReadFnSource(shader, "mm_readB", b, "b", batch_dims, "uniforms.dim_inner", "uniforms.dim_b_outer", transB);
}

void MatMulWriteFnSource(ShaderHelper& shader,
                         const ShaderVariableHelper& output,
                         const ShaderVariableHelper* bias,
                         bool is_gemm,
                         int c_components,
                         bool c_is_scalar,
                         std::string activation_snippet,
                         bool is_channels_last,
                         bool use_split_k,
                         ProgramVariableDataType output_variable_type) {
  const int output_components = output.NumComponents();
  shader.AdditionalImplementation()
      << "fn mm_write(batch: i32, row: i32, colIn: i32, valueIn: output_value_t) {\n";

  shader.AdditionalImplementation() << "  let col = colIn * " << output_components << ";\n";

  shader.AdditionalImplementation() << "  if(row < i32(uniforms.dim_a_outer) && col < i32(uniforms.dim_b_outer)) {\n"
                                    << "    var value = valueIn;\n";

  if (use_split_k) {
    // Set output when MatMul is performed with Split-K.
    // When Split-K is used in MatMul, the bias will be handled in `MatMulFillBiasOrZeroBeforeSplitKProgram`
    // instead of here, so `bias` and `is_channels_last` is not used for Split-K. Note that we