[XLA:GPU] Rename uint32_count to uint8_count in GPUDriver::Asynchrono…

…usMemsetUint8
openxla · Sep 28, 2024 · 17575b2 · 17575b2
1 parent 0f30f33
commit 17575b2
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 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
@@ -499,7 +499,7 @@ class GpuDriver {
   static absl::Status AsynchronousMemsetUint8(Context* context,
                                               GpuDevicePtr location,
                                               uint8_t value,
-                                              size_t uint32_count,
+                                              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
@@ -956,11 +956,11 @@ absl::Status GpuDriver::SynchronousMemsetUint32(Context* context,
 absl::Status GpuDriver::AsynchronousMemsetUint8(Context* context,
                                                 hipDeviceptr_t location,
                                                 uint8 value,
-                                                size_t uint32_count,
+                                                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();
 }