Nit

axinging · Sep 15, 2024 · 2ddd64b · 2ddd64b
1 parent 0021013
commit 2ddd64b
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/onnxruntime/core/providers/webgpu/tensor/transpose.cc b/onnxruntime/core/providers/webgpu/tensor/transpose.cc
@@ -148,14 +148,25 @@ Status Transpose::ComputeInternal(ComputeContext& context) const {
 
   uint32_t output_size = gsl::narrow_cast<int32_t>(input_tensor->Shape().Size());
   TransposeProgram program{*p_perm, use_shared};
+  const auto tileSize = 16; //todo
+  program.SetWorkgroupSize(tileSize,tileSize,1);
+
+  /*
+  useShared
+  .SetDispatchGroupSize(use_shared ? (new_output_shape[1] + tileSize-1 / tileSize, (new_output_shape[0] +tileSize -1 / tileSize): (output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE))
+  */
+  // const auto dispatch_group_size = use_shared ? (new_output_shape[1] + tileSize-1 / tileSize, (new_output_shape[0] +tileSize -1 / tileSize): (output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE);
   program
       .CacheHint(absl::StrJoin(*p_perm, "-"))
       .AddInputs({{input_tensor, ProgramTensorMetadataDependency::TypeAndRank, new_input_shape, 1}})
       .AddOutputs({{output_tensor, ProgramTensorMetadataDependency::TypeAndRank, new_output_shape, 1}})
-      .SetDispatchGroupSize((output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE)
+      .SetDispatchGroupSize(static_cast<uint32_t>((new_output_shape[1] + tileSize-1) / tileSize), static_cast<uint32_t>(((new_output_shape[0] +tileSize -1) / tileSize)))
       .AddUniformVariables({
           {static_cast<uint32_t>(output_size)},
       });
+
+  use_shared? program.SetDispatchGroupSize(static_cast<uint32_t>((new_output_shape[1] + tileSize-1) / tileSize), static_cast<uint32_t>(((new_output_shape[0] +tileSize -1) / tileSize)))
+  : program.SetDispatchGroupSize((output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE);
   return context.RunProgram(program);
 }