Remove GpuStream::platform_specific_stream method in favor of the more widely-used ::gpu_stream method.

PiperOrigin-RevId: 658127980

Author: klucke

xla/stream_executor/cuda/cuda_executor.cc

@@ -665,13 +665,13 @@ void GpuExecutor::DeallocateStream(Stream* stream) {
       dnn_->NotifyStreamDestroyed(stream);
     }
   }
-  GpuStream* cuda_stream = AsGpuStream(stream);
+  GpuStream* gpu_stream = AsGpuStream(stream);
   absl::MutexLock l(&alive_gpu_streams_mu_);
-  alive_gpu_streams_.erase(cuda_stream->platform_specific_stream());
-  if (!cuda_stream->IsIdle()) {
+  alive_gpu_streams_.erase(gpu_stream->gpu_stream());
+  if (!gpu_stream->IsIdle()) {
     LOG(ERROR) << "Deallocating stream with pending work";
   }
-  cuda_stream->Destroy();
+  gpu_stream->Destroy();
 }
 
 absl::Status GpuExecutor::BlockHostUntilDone(Stream* stream) {
@@ -805,20 +805,20 @@ absl::StatusOr<std::unique_ptr<Event>> GpuExecutor::CreateEvent() {
 
 absl::StatusOr<std::unique_ptr<Stream>> GpuExecutor::CreateStream(
     std::optional<std::variant<StreamPriority, int>> priority) {
-  auto gpu_stream = std::make_unique<GpuStream>(this);
+  auto stream = std::make_unique<GpuStream>(this);
   if (priority.has_value()) {
     if (std::holds_alternative<StreamPriority>(*priority)) {
-      gpu_stream->SetPriority(std::get<StreamPriority>(*priority));
+      stream->SetPriority(std::get<StreamPriority>(*priority));
     } else {
-      gpu_stream->SetPriority(std::get<int>(*priority));
+      stream->SetPriority(std::get<int>(*priority));
     }
   }
   absl::MutexLock l(&alive_gpu_streams_mu_);
-  bool init_worked = gpu_stream->Init();
+  bool init_worked = stream->Init();
   if (init_worked) {
-    auto platform_specific_stream = gpu_stream->platform_specific_stream();
-    alive_gpu_streams_[platform_specific_stream] = gpu_stream.get();
-    return std::move(gpu_stream);
+    auto gpu_stream = stream->gpu_stream();
+    alive_gpu_streams_[gpu_stream] = stream.get();
+    return std::move(stream);
   } else {
     return absl::InvalidArgumentError("Failed to initialize gpu stream");
   }

xla/stream_executor/gpu/gpu_stream.h

@@ -58,12 +58,6 @@ class GpuStream : public StreamCommon {
     parent()->DeallocateStream(this);
   }
 
-  // Returns a pointer to a platform specific stream associated with this object
-  // if it exists, or nullptr otherwise. This is available via Stream public API
-  // as Stream::PlatformSpecificHandle, and should not be accessed directly
-  // outside of a StreamExecutor package.
-  void* platform_specific_stream() const { return gpu_stream_; }
-
   // Explicitly initialize the CUDA resources associated with this stream.
   bool Init();
 

xla/stream_executor/rocm/rocm_executor.cc

@@ -499,7 +499,7 @@ void GpuExecutor::DeallocateStream(Stream* stream) {
   }
   GpuStream* rocm_stream = AsGpuStream(stream);
   absl::MutexLock l(&alive_gpu_streams_mu_);
-  alive_gpu_streams_.erase(rocm_stream->platform_specific_stream());
+  alive_gpu_streams_.erase(rocm_stream->gpu_stream());
   if (!rocm_stream->IsIdle()) {
     LOG(ERROR) << "Deallocating stream with pending work";
   }
@@ -642,20 +642,20 @@ absl::StatusOr<std::unique_ptr<Event>> GpuExecutor::CreateEvent() {
 
 absl::StatusOr<std::unique_ptr<Stream>> GpuExecutor::CreateStream(
     std::optional<std::variant<StreamPriority, int>> priority) {
-  auto gpu_stream = std::make_unique<GpuStream>(this);
+  auto stream = std::make_unique<GpuStream>(this);
   if (priority.has_value()) {
     if (std::holds_alternative<StreamPriority>(*priority)) {
-      gpu_stream->SetPriority(std::get<StreamPriority>(*priority));
+      stream->SetPriority(std::get<StreamPriority>(*priority));
     } else {
-      gpu_stream->SetPriority(std::get<int>(*priority));
+      stream->SetPriority(std::get<int>(*priority));
     }
   }
   absl::MutexLock l(&alive_gpu_streams_mu_);
-  bool init_worked = gpu_stream->Init();
+  bool init_worked = stream->Init();
   if (init_worked) {
-    auto platform_specific_stream = gpu_stream->platform_specific_stream();
-    alive_gpu_streams_[platform_specific_stream] = gpu_stream.get();
-    return std::move(gpu_stream);
+    auto gpu_stream = stream->gpu_stream();
+    alive_gpu_streams_[gpu_stream] = stream.get();
+    return std::move(stream);
   } else {
     return absl::InvalidArgumentError("Failed to initialize GPU stream");
   }