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

Code before:
  const uint32_t dim_b_outer = narrow<uint32_t>(b_shape[b_shape.NumDimensions() - 1]);  // right matrix first dimension

  const bool is_vec4 = dim_inner % 4 == 0 && dim_b_outer % 4 == 0;

  InlinedVector<int64_t> elements_per_thread = dim_a_outer <= 8
                                                   ? InlinedVector<int64_t>({4, 1, 1})
                                                   : InlinedVector<int64_t>({4, 4, 1});

  const uint32_t dispatch_x = narrow<uint32_t>((dim_b_outer + MatMul::MATMUL_PACKED_WORKGROUP_SIZE_X * elements_per_thread[0] - 1) /
                                               (MatMul::MATMUL_PACKED_WORKGROUP_SIZE_X * elements_per_thread[0]));
  const uint32_t dispatch_y = narrow<uint32_t>((dim_a_outer + MatMul::MATMUL_PACKED_WORKGROUP_SIZE_Y * elements_per_thread[1] - 1) /
                                               (MatMul::MATMUL_PACKED_WORKGROUP_SIZE_Y * elements_per_thread[1]));
  uint32_t dispatch_z = narrow<uint32_t>((static_cast<uint32_t>(batch_size) + MatMul::MATMUL_PACKED_WORKGROUP_SIZE_Z * elements_per_thread[2] - 1) /
                                         (MatMul::MATMUL_PACKED_WORKGROUP_SIZE_Z * elements_per_thread[2]));

  const int components = is_vec4 ? 4 : 1;
  const TensorShape a_shape_temp = CreateMatMulIntermediateShape(outer_dims_a, dim_a_outer, dim_inner, components);
  const TensorShape b_shape_temp = CreateMatMulIntermediateShape(outer_dims_b, dim_inner, dim_b_outer, components);
  const TensorShape output_shape_temp = TensorShape({batch_size, dim_a_outer, dim_b_outer / components});

  ProgramOutput output(output_tensor, ProgramTensorMetadataDependency::Rank, output_shape_temp, components);
  const Tensor* bias = has_bias ? inputs[2] : nullptr;
  bool use_bias_in_matmul = has_bias;
  uint32_t split_dim_inner = 1;

  const SplitKConfig& split_k_config = context->GetSplitKConfig();
  const bool need_split_k = split_k_config.UseSplitK(context, is_vec4, activation.activation_kind_, batch_size, /*is_gemm*/ false, is_channels_last, dim_a_outer, dim_b_outer, dim_inner);
  if (need_split_k) {
    ORT_ENFORCE(batch_size == 1, "Split-K MatMul only supports batch_size == 1.");
    ORT_ENFORCE(is_vec4, "Split-K MatMul only supports bias in vec4 format.");
    ORT_ENFORCE(is_channels_last, "Split-K MatMul only supports channels-last format.");

    // Initialize `output_tensor` with 0 or bias before MatMulProgram with Split-K enabled.
    const auto fill_bias_program = CreateMatMulFillBiasOrZeroBeforeSplitKProgram(bias, output_tensor, /*is_gemm*/ false, /*beta*/ 1.0f, /*bias_components*/ 4, /*bias_is_scalar*/ false, output_shape_temp);
    ORT_RETURN_IF_ERROR(context->RunProgram(fill_bias_program));

    // `bias` has been handled in the execution of `fill_bias_program` so we don't need to set
    // `bias` again in `MatMulProgram`.
    use_bias_in_matmul = false;

    // With Split-K, `dim_inner` will be split into multiple parts and `dispatch_z` will be the
    // number of splits along `dim_inner`.
    split_dim_inner = split_k_config.GetSplitDimInner();
    dispatch_z = (dim_inner + split_dim_inner - 1) / split_dim_inner;

    // The output should be declared in atomic types in `MatMulProgram` for the use of atomic
    // built-in functions.
    output.is_atomic = true;
  }

  MatMulProgram matmul_program{activation, use_bias_in_matmul, is_vec4, elements_per_thread, is_channels_last, split_dim_inner};

Review comment (suggestion):
I would prefer to move the check of `if (context->KernelContext().GetUseDeterministicCompute())` here instead of inside `UseSplitK()`. Technically the function should not depend on the compute context

Code after:
  const uint32_t dim_b_outer = narrow<uint32_t>(b_shape[b_shape.NumDimensions() - 1]);  // right matrix first dimension

  const bool is_vec4 = dim_inner % 4 == 0 && dim_b_outer % 4 == 0;

  InlinedVector<int64_t> elements_per_thread = dim_a_outer <= 8
                                                   ? InlinedVector<int64_t>({4, 1, 1})
                                                   : InlinedVector<int64_t>({4, 4, 1});

  const uint32_t dispatch_x = narrow<uint32_t>((dim_b_outer + MatMul::MATMUL_PACKED_WORKGROUP_SIZE_X * elements_per_thread[0] - 1) /
                                               (MatMul::MATMUL_PACKED_WORKGROUP_SIZE_X * elements_per_thread[0]));
  const uint32_t dispatch_y = narrow<uint32_t>((dim_a_outer + MatMul::MATMUL_PACKED_WORKGROUP_SIZE_Y * elements_per_thread[1] - 1) /
                                               (MatMul::MATMUL_PACKED_WORKGROUP_SIZE_Y * elements_per_thread[1]));
  uint32_t dispatch_z = narrow<uint32_t>((static_cast<uint32_t>(batch_size) + MatMul::MATMUL_PACKED_WORKGROUP_SIZE_Z * elements_per_thread[2] - 1) /
                                         (MatMul::MATMUL_PACKED_WORKGROUP_SIZE_Z * elements_per_thread[2]));

  const int components = is_vec4 ? 4 : 1;
  const TensorShape a_shape_temp = CreateMatMulIntermediateShape(outer_dims_a, dim_a_outer, dim_inner, components);
  const TensorShape b_shape_temp = CreateMatMulIntermediateShape(outer_dims_b, dim_inner, dim_b_outer, components);
  const TensorShape output_shape_temp = TensorShape({batch_size, dim_a_outer, dim_b_outer / components});

  ProgramOutput output(output_tensor, ProgramTensorMetadataDependency::Rank, output_shape_temp, components);
  const Tensor* bias = has_bias ? inputs[2] : nullptr;
  bool use_bias_in_matmul = has_bias;
  uint32_t split_dim_inner = 1;

  // Current Split-K implementation relies on atomic operations, which are not deterministic.
  if (!context->KernelContext().GetUseDeterministicCompute()) {
    const SplitKConfig& split_k_config = context->GetSplitKConfig();
    const bool need_split_k = split_k_config.UseSplitK(is_vec4, activation.activation_kind_, batch_size, /*is_gemm*/ false, is_channels_last, dim_a_outer, dim_b_outer, dim_inner);
    if (need_split_k) {
      ORT_ENFORCE(batch_size == 1, "Split-K MatMul only supports batch_size == 1.");
      ORT_ENFORCE(is_vec4, "Split-K MatMul only supports bias in vec4 format.");
      ORT_ENFORCE(is_channels_last, "Split-K MatMul only supports channels-last format.");

      // Initialize `output_tensor` with 0 or bias before MatMulProgram with Split-K enabled.
      const auto fill_bias_program = CreateMatMulFillBiasOrZeroBeforeSplitKProgram(bias, output_tensor, /*is_gemm*/ false, /*beta*/ 1.0f, /*bias_components*/ 4, /*bias_is_scalar*/ false, output_shape_temp);
      ORT_RETURN_IF_ERROR(context->RunProgram(fill_bias_program));

      // `bias` has been handled in the execution of `fill_bias_program` so we don't need to set
      // `bias` again in `MatMulProgram`.
      use_bias_in_matmul = false;

      // With Split-K, `dim_inner` will be split into multiple parts and `dispatch_z` will be the
      // number of splits along `dim_inner`.
      split_dim_inner = split_k_config.GetSplitDimInner();
      dispatch_z = (dim_inner + split_dim_inner - 1) / split_dim_inner;

      // The output should be declared in atomic types in `MatMulProgram` for the use of atomic
      // built-in functions.
      output.is_atomic = true;
    }