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

Code before:
      });

  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(WORKGROUP_SIZE);

    uint32_t dispatch_x = ceil_div(output_size, WORKGROUP_SIZE);
    uint32_t dispatch_y = 1;
    uint32_t dispatch_z = 1;

    // This temporary workaround addresses a significant performance bottleneck
    // (~12x slower) for the shape (3, 3, 2560, 1280) 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"}) {
      uint32_t dispatch_size = dispatch_x;
      dispatch_x = 4;
      dispatch_y = 8;
      dispatch_z = ceil_div(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();
  int32_t rank = static_cast<int32_t>(input_shape.NumDimensions());

  TensorShapeVector output_dims(rank);
  InlinedVector<size_t> default_perm(rank);
  const InlinedVector<size_t>* p_perm = nullptr;
  ORT_RETURN_IF_ERROR(ComputeOutputShape(*input_tensor, output_dims, default_perm, p_perm));
  TensorShape output_shape(output_dims);
  auto* output_tensor = context.Output(0, output_shape);

  int64_t output_size = output_shape.Size();
  if (output_size == 0) {
    return Status::OK();
  }

  return DoTranspose(context, *p_perm, *input_tensor, *output_tensor);
}

}  // namespace webgpu

Review comment (suggestion):
When running `detr-resnet-50` model, I found there are redundant/invalid work groups. For example, the input shape is [512, 512, 3, 3], the original dispatch size is [36864, 1, 1]. After these codes, the dispatch size is [131072, 1, 1],  nearly ~3x invalid work groups. The PR changed the dispatch size to [4, 8, 1152], was still 36864 work groups. And the worst situation only has 31 invalid work groups.
And The Transpose with [512, 512, 3, 3] shape cost from 3.86 ms to 1.29 ms.

Code after:
      });

  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(WORKGROUP_SIZE);

    uint32_t dispatch_x = ceil_div(output_size, WORKGROUP_SIZE);
    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 = ceil_div(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();
  int32_t rank = static_cast<int32_t>(input_shape.NumDimensions());

  TensorShapeVector output_dims(rank);
  InlinedVector<size_t> default_perm(rank);
  const InlinedVector<size_t>* p_perm = nullptr;
  ORT_RETURN_IF_ERROR(ComputeOutputShape(*input_tensor, output_dims, default_perm, p_perm));
  TensorShape output_shape(output_dims);
  auto* output_tensor = context.Output(0, output_shape);

  int64_t output_size = output_shape.Size();
  if (output_size == 0) {
    return Status::OK();
  }

  return DoTranspose(context, *p_perm, *input_tensor, *output_tensor);
}

}  // namespace webgpu