Collapse GpuStream class into StreamCommon.

PiperOrigin-RevId: 687352019
openxla · Oct 18, 2024 · 873f95f · 873f95f
1 parent 8761b7e
commit 873f95f
Show file tree

Hide file tree

Showing 17 changed files with 62 additions and 124 deletions.
diff --git a/xla/service/gpu/runtime/nccl_api.cc b/xla/service/gpu/runtime/nccl_api.cc
@@ -37,6 +37,7 @@ limitations under the License.
 #include "xla/stream_executor/device_memory_allocator.h"
 #include "xla/stream_executor/gpu/gpu_stream.h"
 #include "xla/stream_executor/stream.h"
+#include "xla/stream_executor/stream_executor.h"
 #include "xla/tsl/concurrency/ref_count.h"
 #include "xla/xla_data.pb.h"
 #include "tsl/platform/errors.h"

diff --git a/xla/stream_executor/cuda/BUILD b/xla/stream_executor/cuda/BUILD
@@ -1215,8 +1215,6 @@ cc_library(
         "//xla/stream_executor:platform",
         "//xla/stream_executor:stream",
         "//xla/stream_executor:stream_common",
-        "//xla/stream_executor/gpu:gpu_driver_header",
-        "//xla/stream_executor/gpu:gpu_stream",
         "@com_google_absl//absl/base",
         "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/log",

diff --git a/xla/stream_executor/cuda/cuda_stream.cc b/xla/stream_executor/cuda/cuda_stream.cc
@@ -40,7 +40,6 @@ limitations under the License.
 #include "xla/stream_executor/cuda/cuda_status.h"
 #include "xla/stream_executor/device_memory.h"
 #include "xla/stream_executor/event.h"
-#include "xla/stream_executor/gpu/gpu_driver.h"
 #include "xla/stream_executor/kernel.h"
 #include "xla/stream_executor/launch_dim.h"
 #include "xla/stream_executor/platform.h"

diff --git a/xla/stream_executor/cuda/cuda_stream.h b/xla/stream_executor/cuda/cuda_stream.h
@@ -31,16 +31,16 @@ limitations under the License.
 #include "xla/stream_executor/device_memory.h"
 #include "xla/stream_executor/event.h"
 #include "xla/stream_executor/event_based_timer.h"
-#include "xla/stream_executor/gpu/gpu_stream.h"
 #include "xla/stream_executor/kernel.h"
 #include "xla/stream_executor/launch_dim.h"
 #include "xla/stream_executor/platform.h"
 #include "xla/stream_executor/stream.h"
+#include "xla/stream_executor/stream_common.h"
 
 namespace stream_executor {
 namespace gpu {
 
-class CudaStream : public GpuStream {
+class CudaStream : public StreamCommon {
  public:
   absl::Status WaitFor(Stream* other) override;
   absl::Status RecordEvent(Event* event) override;
@@ -82,7 +82,7 @@ class CudaStream : public GpuStream {
   CudaStream(StreamExecutor* executor, CudaEvent completed_event,
              std::optional<std::variant<StreamPriority, int>> priority,
              CUstream stream_handle)
-      : GpuStream(executor, priority),
+      : StreamCommon(executor, priority),
         executor_(executor),
         completed_event_(std::move(completed_event)),
         stream_handle_(stream_handle) {}

diff --git a/xla/stream_executor/cuda/cuda_timer_test.cc b/xla/stream_executor/cuda/cuda_timer_test.cc
@@ -26,7 +26,6 @@ limitations under the License.
 #include "xla/stream_executor/cuda/cuda_executor.h"
 #include "xla/stream_executor/cuda/cuda_platform_id.h"
 #include "xla/stream_executor/device_memory.h"
-#include "xla/stream_executor/gpu/gpu_stream.h"
 #include "xla/stream_executor/gpu/gpu_test_kernels.h"
 #include "xla/stream_executor/kernel.h"
 #include "xla/stream_executor/kernel_spec.h"
@@ -73,7 +72,6 @@ class CudaTimerTest : public ::testing::TestWithParam<CudaTimer::TimerType> {
 
   StreamExecutor* executor_;
   std::unique_ptr<Stream> stream_;
-  GpuStream* gpu_stream_;
 
  private:
   void SetUp() override {
@@ -82,13 +80,12 @@ class CudaTimerTest : public ::testing::TestWithParam<CudaTimer::TimerType> {
                                 stream_executor::cuda::kCudaPlatformId));
     TF_ASSERT_OK_AND_ASSIGN(executor_, platform->ExecutorForDevice(0));
     TF_ASSERT_OK_AND_ASSIGN(stream_, executor_->CreateStream(std::nullopt));
-    gpu_stream_ = AsGpuStream(stream_.get());
   }
 };
 
 TEST_P(CudaTimerTest, Create) {
   TF_ASSERT_OK_AND_ASSIGN(
-      CudaTimer timer, CudaTimer::Create(executor_, gpu_stream_, GetParam()));
+      CudaTimer timer, CudaTimer::Create(executor_, stream_.get(), GetParam()));
 
   // We don't really care what kernel we launch here as long as it takes a
   // non-zero amount of time.

diff --git a/xla/stream_executor/gpu/BUILD b/xla/stream_executor/gpu/BUILD
@@ -356,16 +356,9 @@ gpu_only_cc_library(
     hdrs = ["gpu_stream.h"],
     deps = [
         ":gpu_types_header",
-        "//xla/stream_executor:kernel",
-        "//xla/stream_executor:launch_dim",
-        "//xla/stream_executor:platform",
         "//xla/stream_executor:stream",
-        "//xla/stream_executor:stream_common",
-        "//xla/stream_executor:stream_executor_h",
         "@com_google_absl//absl/base",
-        "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
-        "@com_google_absl//absl/status",
     ],
 )
 

diff --git a/xla/stream_executor/gpu/gpu_stream.cc b/xla/stream_executor/gpu/gpu_stream.cc
@@ -15,43 +15,18 @@ limitations under the License.
 
 #include "xla/stream_executor/gpu/gpu_stream.h"
 
-#include <optional>
-
 #include "absl/base/casts.h"
 #include "absl/log/check.h"
-#include "absl/log/log.h"
-#include "absl/status/status.h"
 #include "xla/stream_executor/gpu/gpu_types.h"
-#include "xla/stream_executor/kernel.h"
-#include "xla/stream_executor/launch_dim.h"
 #include "xla/stream_executor/stream.h"
 
 namespace stream_executor {
 namespace gpu {
 
-absl::Status GpuStream::Launch(const ThreadDim& thread_dims,
-                               const BlockDim& block_dims, const Kernel& kernel,
-                               const KernelArgs& args) {
-  return Launch(thread_dims, block_dims, std::nullopt, kernel, args);
-}
-
-absl::Status GpuStream::Launch(const ThreadDim& thread_dims,
-                               const BlockDim& block_dims,
-                               const ClusterDim& cluster_dims,
-                               const Kernel& kernel, const KernelArgs& args) {
-  return Launch(thread_dims, block_dims, std::make_optional(cluster_dims),
-                kernel, args);
-}
-
-GpuStream* AsGpuStream(Stream* stream) {
-  DCHECK(stream != nullptr);
-  return static_cast<GpuStream*>(stream);
-}
-
 GpuStreamHandle AsGpuStreamValue(Stream* stream) {
   DCHECK(stream != nullptr);
   return absl::bit_cast<GpuStreamHandle>(
-      AsGpuStream(stream)->platform_specific_handle().stream);
+      stream->platform_specific_handle().stream);
 }
 
 }  // namespace gpu

diff --git a/xla/stream_executor/gpu/gpu_stream.h b/xla/stream_executor/gpu/gpu_stream.h
@@ -19,59 +19,12 @@ limitations under the License.
 #ifndef XLA_STREAM_EXECUTOR_GPU_GPU_STREAM_H_
 #define XLA_STREAM_EXECUTOR_GPU_GPU_STREAM_H_
 
-#include <optional>
-#include <variant>
-
-#include "absl/log/check.h"
 #include "xla/stream_executor/gpu/gpu_types.h"
-#include "xla/stream_executor/kernel.h"
-#include "xla/stream_executor/launch_dim.h"
-#include "xla/stream_executor/platform.h"
 #include "xla/stream_executor/stream.h"
-#include "xla/stream_executor/stream_common.h"
-#include "xla/stream_executor/stream_executor.h"
 
 namespace stream_executor {
 namespace gpu {
 
-// Wraps a GpuStreamHandle in order to satisfy the platform-independent
-// StreamInterface.
-//
-// Thread-safe post-initialization.
-class GpuStream : public StreamCommon {
- public:
-  GpuStream(StreamExecutor* parent,
-            std::optional<std::variant<StreamPriority, int>> priority)
-      : StreamCommon(parent) {
-    if (priority.has_value()) {
-      stream_priority_ = priority.value();
-    }
-  }
-
-  std::variant<StreamPriority, int> priority() const override {
-    return stream_priority_;
-  }
-
-  absl::Status Launch(const ThreadDim& thread_dims, const BlockDim& block_dims,
-                      const Kernel& k, const KernelArgs& args) override;
-  absl::Status Launch(const ThreadDim& thread_dims, const BlockDim& block_dims,
-                      const ClusterDim& cluster_dims, const Kernel& k,
-                      const KernelArgs& args) override;
-
- private:
-  // Helper method to launch a kernel with optional cluster dimensions.
-  virtual absl::Status Launch(const ThreadDim& thread_dims,
-                              const BlockDim& block_dims,
-                              const std::optional<ClusterDim>& cluster_dims,
-                              const Kernel& kernel, const KernelArgs& args) = 0;
-
-  std::variant<StreamPriority, int> stream_priority_;
-};
-
-// Helper functions to simplify extremely common flows.
-// Converts a Stream to the underlying GpuStream implementation.
-GpuStream* AsGpuStream(Stream* stream);
-
 // Extracts a GpuStreamHandle from a GpuStream-backed Stream object.
 GpuStreamHandle AsGpuStreamValue(Stream* stream);
 }  // namespace gpu

diff --git a/xla/stream_executor/host/host_stream.cc b/xla/stream_executor/host/host_stream.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include <cfenv>  // NOLINT
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <queue>
 #include <utility>
 
@@ -197,7 +198,13 @@ absl::Status HostStream::BlockUntilDone() {
 
 absl::Status HostStream::Launch(const ThreadDim& thread_dims,
                                 const BlockDim& block_dims,
+                                const std::optional<ClusterDim>& cluster_dims,
                                 const Kernel& kernel, const KernelArgs& args) {
+  if (cluster_dims.has_value()) {
+    if (cluster_dims->x != 1 || cluster_dims->y != 1 || cluster_dims->z != 1) {
+      return absl::UnimplementedError("Not implemented for Host");
+    }
+  }
   const HostKernel* host_kernel = AsHostKernel(&kernel);
 
   const KernelArgsDeviceMemoryArray* device_mem =

diff --git a/xla/stream_executor/host/host_stream.h b/xla/stream_executor/host/host_stream.h
@@ -73,6 +73,7 @@ class HostStream : public StreamCommon {
   absl::Status DoHostCallbackWithStatus(
       absl::AnyInvocable<absl::Status() &&> callback) override;
   absl::Status Launch(const ThreadDim& thread_dims, const BlockDim& block_dims,
+                      const std::optional<ClusterDim>& cluster_dims,
                       const Kernel& kernel, const KernelArgs& args) override;
 
  private:

diff --git a/xla/stream_executor/mock_stream.h b/xla/stream_executor/mock_stream.h
@@ -18,13 +18,13 @@ limitations under the License.
 
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <string>
 #include <variant>
 
 #include "absl/functional/any_invocable.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "absl/strings/string_view.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/device_memory.h"
 #include "xla/stream_executor/event.h"
@@ -77,11 +77,7 @@ class MockStream : public Stream {
               (const, override));
   MOCK_METHOD(absl::Status, Launch,
               (const ThreadDim &thread_dims, const BlockDim &block_dims,
-               const Kernel &k, const KernelArgs &args),
-              (override));
-  MOCK_METHOD(absl::Status, Launch,
-              (const ThreadDim &thread_dims, const BlockDim &block_dims,
-               const ClusterDim &cluster_dims, const Kernel &k,
+               const std::optional<ClusterDim> &cluster_dims, const Kernel &k,
                const KernelArgs &args),
               (override));
   MOCK_METHOD(const std::string &, GetName, (), (const, override));

diff --git a/xla/stream_executor/rocm/BUILD b/xla/stream_executor/rocm/BUILD
@@ -1002,7 +1002,7 @@ cc_library(
         "//xla/stream_executor:launch_dim",
         "//xla/stream_executor:platform",
         "//xla/stream_executor:stream",
-        "//xla/stream_executor/gpu:gpu_stream",
+        "//xla/stream_executor:stream_common",
         "@com_google_absl//absl/base",
         "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/log",
@@ -1097,7 +1097,6 @@ xla_test(
         "//xla/stream_executor:platform_manager",
         "//xla/stream_executor:stream",
         "//xla/stream_executor:typed_kernel_factory",
-        "//xla/stream_executor/gpu:gpu_stream",
         "//xla/stream_executor/gpu:gpu_test_kernels_rocm",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/time",

diff --git a/xla/stream_executor/rocm/rocm_stream.h b/xla/stream_executor/rocm/rocm_stream.h
@@ -17,29 +17,27 @@ limitations under the License.
 #define XLA_STREAM_EXECUTOR_ROCM_ROCM_STREAM_H_
 
 #include <cstdint>
-#include <memory>
 #include <optional>
 #include <utility>
 #include <variant>
 
 #include "absl/functional/any_invocable.h"
 #include "absl/status/status.h"
-#include "absl/status/statusor.h"
 #include "rocm/include/hip/hip_runtime.h"
 #include "xla/stream_executor/device_memory.h"
 #include "xla/stream_executor/event.h"
 #include "xla/stream_executor/event_based_timer.h"
-#include "xla/stream_executor/gpu/gpu_stream.h"
 #include "xla/stream_executor/kernel.h"
 #include "xla/stream_executor/launch_dim.h"
 #include "xla/stream_executor/platform.h"
 #include "xla/stream_executor/rocm/rocm_event.h"
 #include "xla/stream_executor/stream.h"
+#include "xla/stream_executor/stream_common.h"
 
 namespace stream_executor {
 namespace gpu {
 
-class RocmStream : public GpuStream {
+class RocmStream : public StreamCommon {
  public:
   absl::Status WaitFor(Stream* other) override;
   absl::Status RecordEvent(Event* event) override;
@@ -79,7 +77,7 @@ class RocmStream : public GpuStream {
   RocmStream(StreamExecutor* executor, RocmEvent completed_event,
              std::optional<std::variant<StreamPriority, int>> priority,
              hipStream_t stream_handle)
-      : GpuStream(executor, priority),
+      : StreamCommon(executor, priority),
         executor_(executor),
         completed_event_(std::move(completed_event)),
         stream_handle_(stream_handle) {}

diff --git a/xla/stream_executor/rocm/rocm_timer_test.cc b/xla/stream_executor/rocm/rocm_timer_test.cc
@@ -24,7 +24,6 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/time/time.h"
 #include "xla/stream_executor/device_memory.h"
-#include "xla/stream_executor/gpu/gpu_stream.h"
 #include "xla/stream_executor/gpu/gpu_test_kernels.h"
 #include "xla/stream_executor/kernel.h"
 #include "xla/stream_executor/kernel_spec.h"
@@ -73,7 +72,6 @@ class RocmTimerTest : public ::testing::Test {
 
   RocmExecutor* executor_;
   std::unique_ptr<Stream> stream_;
-  GpuStream* gpu_stream_;
 
  private:
   void SetUp() override {
@@ -84,13 +82,12 @@ class RocmTimerTest : public ::testing::Test {
                             platform->ExecutorForDevice(0));
     executor_ = reinterpret_cast<RocmExecutor*>(executor);
     TF_ASSERT_OK_AND_ASSIGN(stream_, executor_->CreateStream(std::nullopt));
-    gpu_stream_ = AsGpuStream(stream_.get());
   }
 };
 
 TEST_F(RocmTimerTest, Create) {
   TF_ASSERT_OK_AND_ASSIGN(RocmTimer timer,
-                          RocmTimer::Create(executor_, gpu_stream_));
+                          RocmTimer::Create(executor_, stream_.get()));
 
   // We don't really care what kernel we launch here as long as it takes a
   // non-zero amount of time.