PR #17737: [XLA:GPU] Rename uint32_count to uint8_count in `GPUDr…

…iver::AsynchronousMemsetUint8` Imported from GitHub PR #17737 In `GpuDriver::AsynchronousMemsetUint8`, the count should be in bytes, instead of in uint32. Copybara import of the project: -- 17575b2 by PragmaTwice <[email protected]>: [XLA:GPU] Rename uint32_count to uint8_count in GPUDriver::AsynchronousMemsetUint8 Merging this change closes #17737 COPYBARA_INTEGRATE_REVIEW=#17737 from PragmaTwice:fix-typo 17575b2 PiperOrigin-RevId: 680656988
openxla · Sep 30, 2024 · b24e4a4 · b24e4a4
1 parent 21306df
commit b24e4a4
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 7 deletions.
diff --git a/xla/stream_executor/cuda/cuda_driver.cc b/xla/stream_executor/cuda/cuda_driver.cc
@@ -1079,10 +1079,10 @@ absl::Status GpuDriver::SynchronousMemsetUint32(Context* context,
 absl::Status GpuDriver::AsynchronousMemsetUint8(Context* context,
                                                 CUdeviceptr location,
                                                 uint8_t value,
-                                                size_t uint32_count,
+                                                size_t uint8_count,
                                                 CUstream stream) {
   ScopedActivateContext activation(context);
-  return cuda::ToStatus(cuMemsetD8Async(location, value, uint32_count, stream),
+  return cuda::ToStatus(cuMemsetD8Async(location, value, uint8_count, stream),
                         "Failed to enqueue async memset operation");
 }
 

diff --git a/xla/stream_executor/gpu/gpu_driver.h b/xla/stream_executor/gpu/gpu_driver.h
@@ -498,8 +498,7 @@ class GpuDriver {
   // http://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MEM.html#group__CUDA__MEM_1gaef08a7ccd61112f94e82f2b30d43627
   static absl::Status AsynchronousMemsetUint8(Context* context,
                                               GpuDevicePtr location,
-                                              uint8_t value,
-                                              size_t uint32_count,
+                                              uint8_t value, size_t uint8_count,
                                               GpuStreamHandle stream);
 
   // Performs an asynchronous memset of the device memory segment via

diff --git a/xla/stream_executor/rocm/rocm_driver.cc b/xla/stream_executor/rocm/rocm_driver.cc
@@ -955,12 +955,11 @@ absl::Status GpuDriver::SynchronousMemsetUint32(Context* context,
 
 absl::Status GpuDriver::AsynchronousMemsetUint8(Context* context,
                                                 hipDeviceptr_t location,
-                                                uint8 value,
-                                                size_t uint32_count,
+                                                uint8 value, size_t uint8_count,
                                                 GpuStreamHandle stream) {
   ScopedActivateContext activation{context};
   RETURN_IF_ROCM_ERROR(
-      wrap::hipMemsetAsync(location, value, uint32_count, stream),
+      wrap::hipMemsetAsync(location, value, uint8_count, stream),
       "Failed to enqueue async memset operation");
   return absl::OkStatus();
 }