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

Code before:
    Tensor* output,
    bool is_gemm,
    float beta,
    uint32_t output_components,
    bool bias_is_scalar,
    const TensorShape& output_shape) {
  const bool has_bias = bias != nullptr;

  // Currently we only support GEMM and channels last format for MatMul with Split-K.
  MatMulFillBiasOrZeroBeforeSplitKProgram program(is_gemm, has_bias, output_components, bias_is_scalar);

  const uint32_t dim_a_outer = narrow<uint32_t>(output_shape[output_shape.NumDimensions() - 2]);
  const uint32_t dim_b_outer = narrow<uint32_t>(output_shape[output_shape.NumDimensions() - 1]);

  // Fill one value per invocation. Now we use default workgroup size (64) for this program.
  const uint32_t total_outputs = dim_a_outer * dim_b_outer;
  const uint32_t dispatch_x = (total_outputs + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE;

  TensorShape output_shape_temp;
  if (is_gemm) {
    // GEMM doesn't have `batch` in its output shape.
    output_shape_temp = TensorShape({dim_a_outer, dim_b_outer});
  } else {
    const uint32_t batch_size = 1;
    output_shape_temp = TensorShape({batch_size, dim_a_outer, dim_b_outer});
  }

  // To reuse `MatMulWriteFnSource()` we need to set `dim_b_outer` in components when `output_shape`
  // is in `vec4`, while use `output_shape` directly as the output shape.
  const uint32_t dim_b_outer_components = narrow<uint32_t>(dim_b_outer * output_components);
  program.CacheHint(is_gemm, has_bias, output_components, bias_is_scalar)
      .AddOutput({output, ProgramTensorMetadataDependency::TypeAndRank, output_shape_temp, static_cast<int32_t>(output_components)})
      .AddUniformVariables({{dim_a_outer}, {dim_b_outer_components}, {beta}})
      .SetDispatchGroupSize(dispatch_x);

  if (has_bias) {
    // We always use `c_components` as `output_components` in GEMM, and 4 in MatMul.
    const TensorShape reduced_bias_shape = ReduceShapeByComponents(bias->Shape(), output_components);
    program.AddInput({bias, ProgramTensorMetadataDependency::TypeAndRank, reduced_bias_shape, static_cast<int32_t>(output_components)});
  }

  return program;
}

}  // namespace webgpu
}  // namespace onnxruntime


Review comment (performance):
Why not directly use `output_shape` since splitK has ensured that batch_size must be 1? And pass the expected 2d shape for is_gemm=true?

Code after:
    Tensor* output,
    bool is_gemm,
    float beta,
    uint32_t output_components,
    bool bias_is_scalar,
    const TensorShape& output_shape) {
  const bool has_bias = bias != nullptr;

  // Currently we only support GEMM and channels last format for MatMul with Split-K.
  MatMulFillBiasOrZeroBeforeSplitKProgram program(is_gemm, has_bias, output_components, bias_is_scalar);

  const uint32_t dim_a_outer = narrow<uint32_t>(output_shape[output_shape.NumDimensions() - 2]);
  const uint32_t dim_b_outer = narrow<uint32_t>(output_shape[output_shape.NumDimensions() - 1]);

  // Fill one value per invocation. Now we use default workgroup size (64) for this program.
  const uint32_t total_outputs = dim_a_outer * dim_b_outer;
  const uint32_t dispatch_x = (total_outputs + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE;

  // To reuse `MatMulWriteFnSource()` we need to set `dim_b_outer` in components when `output_shape`
  // is in `vec4`, while use `output_shape` directly as the output shape.
  const uint32_t dim_b_outer_components = narrow<uint32_t>(dim_b_outer * output_components);
  program.CacheHint(is_gemm, has_bias, output_components, bias_is_scalar)
      .AddOutput({output, ProgramTensorMetadataDependency::TypeAndRank, output_shape, static_cast<int32_t>(output_components)})
      .AddUniformVariables({{dim_a_outer}, {dim_b_outer_components}, {beta}})
      .SetDispatchGroupSize(dispatch_x);

  if (has_bias) {
    // We always use `c_components` as `output_components` in GEMM, and 4 in MatMul.
    const TensorShape reduced_bias_shape = ReduceShapeByComponents(bias->Shape(), output_components);
    program.AddInput({bias, ProgramTensorMetadataDependency::TypeAndRank, reduced_bias_shape, static_cast<int32_t>(output_components)});
  }

  return program;
}

}  // namespace webgpu
}  // namespace onnxruntime