Repository: microsoft/onnxruntime
File: onnxruntime/core/providers/webgpu/tensor/transpose.cc (C++)

Code before:
  TensorShape new_output_shape(output_dims);

  if (use_shared) {
    new_input_shape = channels_last
                          ? TensorShape({new_shape[0], new_shape[1] * new_shape[2]})
                      : channels_first
                          ? TensorShape({new_shape[0] * new_shape[1], new_shape[2]})
                          : new_shape;
    new_output_shape = TensorShape({new_input_shape[1], new_input_shape[0]});
  }

  uint32_t output_size = onnxruntime::narrow<int32_t>(input_shape.Size());
  TransposeProgram program{permutations, use_shared};

  program
      .CacheHint(absl::StrJoin(permutations, "-"))
      .AddInputs({{&input, ProgramTensorMetadataDependency::TypeAndRank, new_input_shape, 1}})
      .AddOutputs({{&output, ProgramTensorMetadataDependency::None, new_output_shape, 1}})
      .AddUniformVariables({
          {static_cast<uint32_t>(output_size)},
      });

  if (use_shared) {
    program.SetWorkgroupSize(TILE_SIZE, TILE_SIZE, 1);
    program.SetDispatchGroupSize(static_cast<uint32_t>((new_output_shape[1] + TILE_SIZE - 1) / TILE_SIZE),
                                 static_cast<uint32_t>(((new_output_shape[0] + TILE_SIZE - 1) / TILE_SIZE)));
  } else {
    program.SetWorkgroupSize(64u);

    uint32_t dispatch_x = CeilDiv(output_size, 64u);
    uint32_t dispatch_y = 1;
    uint32_t dispatch_z = 1;

    // This temporary workaround addresses a significant performance bottleneck
    // (~12x slower) for the input shape (1280, 2560, 3, 3) due to an issue with Intel's
    // GPU drivers. We manually normalize the dispatch group size to restore
    // performance.
    //
    // TODO: Revert this change once the driver issue is fixed.
    if (context.AdapterInfo().vendor == std::string_view{"intel"} && rank == 4) {
      uint32_t dispatch_size = dispatch_x;
      dispatch_x = 4;
      dispatch_y = 8;
      dispatch_z = CeilDiv(dispatch_size, dispatch_x * dispatch_y);
    }
    program.SetDispatchGroupSize(dispatch_x, dispatch_y, dispatch_z);
  }
  return context.RunProgram(program);
}

Status Transpose::ComputeInternal(ComputeContext& context) const {

Review comment (suggestion):
nit: why not use `WORKGROUP_SIZE`?
In line 134, please help correct it to `uint32_t output_size = onnxruntime::narrow<uint32_t>(input_shape.Size());`

Code after:
  TensorShape new_output_shape(output_dims);

  if (use_shared) {
    new_input_shape = channels_last
                          ? TensorShape({new_shape[0], new_shape[1] * new_shape[2]})
                      : channels_first
                          ? TensorShape({new_shape[0] * new_shape[1], new_shape[2]})
                          : new_shape;
    new_output_shape = TensorShape({new_input_shape[1], new_input_shape[0]});
  }

  uint32_t output_size = onnxruntime::narrow<uint32_t>(input_shape.Size());
  TransposeProgram program{permutations, use_shared};

  program
      .CacheHint(absl::StrJoin(permutations, "-"))
      .AddInputs({{&input, ProgramTensorMetadataDependency::TypeAndRank, new_input_shape, 1}})
      .AddOutputs({{&output, ProgramTensorMetadataDependency::None, new_output_shape, 1}})
      .AddUniformVariables({{output_size}});

  if (use_shared) {
    program.SetWorkgroupSize(TILE_SIZE, TILE_SIZE, 1);
    program.SetDispatchGroupSize(static_cast<uint32_t>((new_output_shape[1] + TILE_SIZE - 1) / TILE_SIZE),
                                 static_cast<uint32_t>(((new_output_shape[0] + TILE_SIZE - 1) / TILE_SIZE)));
  } else {
    program.SetWorkgroupSize(64u);

    uint32_t dispatch_x = CeilDiv(output_size, 64u);
    uint32_t dispatch_y = 1;
    uint32_t dispatch_z = 1;

    // This temporary workaround addresses a significant performance bottleneck
    // (~12x slower) for the input shape (1280, 2560, 3, 3) due to an issue with Intel's
    // GPU drivers. We manually normalize the dispatch group size to restore
    // performance.
    //
    // TODO: Revert this change once the driver issue is fixed.
    if (context.AdapterInfo().vendor == std::string_view{"intel"} && rank == 4) {
      uint32_t dispatch_size = dispatch_x;
      dispatch_x = 4;
      dispatch_y = 8;
      dispatch_z = CeilDiv(dispatch_size, dispatch_x * dispatch_y);
    }
    program.SetDispatchGroupSize(dispatch_x, dispatch_y, dispatch_z);
  }
  return context.RunProgram(program);
}

Status Transpose::ComputeInternal(ComputeContext& context) const {
  const auto* input_tensor = context.Input(0);
  const TensorShape& input_shape = input_tensor->Shape();