[ROCm] Fix grid dimension issue introduced in d8909c9

openxla · Nov 11, 2024 · dd86f21 · dd86f21
1 parent b0aae98
commit dd86f21
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/xla/stream_executor/rocm/rocm_command_buffer.cc b/xla/stream_executor/rocm/rocm_command_buffer.cc
@@ -273,7 +273,7 @@ absl::StatusOr<GraphNodeHandle> RocmCommandBuffer::CreateKernelNode(
       static_cast<const RocmKernel&>(kernel).gpu_function();
   params.func = function;
   params.gridDim.x = blocks.x;
-  params.gridDim.z = blocks.y;
+  params.gridDim.y = blocks.y;
   params.gridDim.z = blocks.z;
   params.blockDim.x = threads.x;
   params.blockDim.y = threads.y;
@@ -319,7 +319,7 @@ absl::Status RocmCommandBuffer::UpdateKernelNode(
       static_cast<const RocmKernel&>(kernel).gpu_function();
   params.func = function;
   params.gridDim.x = blocks.x;
-  params.gridDim.z = blocks.y;
+  params.gridDim.y = blocks.y;
   params.gridDim.z = blocks.z;
   params.blockDim.x = threads.x;
   params.blockDim.y = threads.y;