Remove StreamExecutorInterface::WaitForEventOnExternalStream and replace with a proper virtual method on Event.

third_party/xla/xla/stream_executor/cuda/cuda_executor.cc

@@ -772,18 +772,6 @@ absl::Status GpuExecutor::WaitForEvent(Stream* stream, Event* event) {
   }
 }
 
-absl::Status GpuExecutor::WaitForEventOnExternalStream(std::intptr_t stream,
-                                                       Event* event) {
-  if (GpuDriver::WaitStreamOnEvent(context_,
-                                   absl::bit_cast<GpuStreamHandle>(stream),
-                                   AsGpuEvent(event)->gpu_event())) {
-    return absl::OkStatus();
-  } else {
-    return absl::InternalError(
-        "error waiting for CUDA event on external stream");
-  }
-}
-
 Event::Status GpuExecutor::PollForEventStatus(Event* event) {
   return AsGpuEvent(event)->PollForStatus();
 }

third_party/xla/xla/stream_executor/event.cc

@@ -32,8 +32,4 @@ Event::Status Event::PollForStatus() {
   return stream_exec_->PollForEventStatus(this);
 }
 
-absl::Status Event::WaitForEventOnExternalStream(std::intptr_t stream) {
-  return stream_exec_->WaitForEventOnExternalStream(stream, this);
-}
-
 }  // namespace stream_executor

third_party/xla/xla/stream_executor/event.h

@@ -41,7 +41,7 @@ class Event {
     kComplete,
   };
 
-  Event(StreamExecutorInterface* stream_exec);
+  explicit Event(StreamExecutorInterface* stream_exec);
 
   // Releases any resources held by the Event object.
   virtual ~Event() = default;
@@ -51,7 +51,9 @@ class Event {
 
   // Blocks `stream` on this event. `stream` is a raw platform-specific
   // stream (e.g. GpuStreamHandle).
-  absl::Status WaitForEventOnExternalStream(std::intptr_t stream);
+  virtual absl::Status WaitForEventOnExternalStream(std::intptr_t stream) {
+    return absl::UnimplementedError("Not supported for this Event.");
+  }
 
   Event(Event&&) = default;
   Event& operator=(Event&&) = default;

third_party/xla/xla/stream_executor/gpu/BUILD

@@ -197,6 +197,7 @@ gpu_only_cc_library(
         ":gpu_stream",
         ":gpu_types_header",
         "//xla/stream_executor:stream_executor_headers",
+        "@com_google_absl//absl/base",
         "@com_google_absl//absl/status",
     ],
 )

third_party/xla/xla/stream_executor/gpu/gpu_event.cc

@@ -15,7 +15,11 @@ limitations under the License.
 
 #include "xla/stream_executor/gpu/gpu_event.h"
 
+#include <cstdint>
+
+#include "absl/base/casts.h"
 #include "absl/status/status.h"
+#include "xla/stream_executor/event.h"
 #include "xla/stream_executor/gpu/gpu_driver.h"
 #include "xla/stream_executor/gpu/gpu_executor.h"
 #include "xla/stream_executor/gpu/gpu_stream.h"
@@ -45,5 +49,15 @@ absl::Status GpuEvent::Record(GpuStream* stream) {
 
 GpuEventHandle GpuEvent::gpu_event() { return gpu_event_; }
 
+absl::Status GpuEvent::WaitForEventOnExternalStream(std::intptr_t stream) {
+  if (GpuDriver::WaitStreamOnEvent(parent_->gpu_context(),
+                                   absl::bit_cast<GpuStreamHandle>(stream),
+                                   gpu_event_)) {
+    return absl::OkStatus();
+  } else {
+    return absl::InternalError("Error waiting for event on external stream");
+  }
+}
+
 }  // namespace gpu
 }  // namespace stream_executor

third_party/xla/xla/stream_executor/gpu/gpu_event.h

@@ -16,6 +16,8 @@ limitations under the License.
 #ifndef XLA_STREAM_EXECUTOR_GPU_GPU_EVENT_H_
 #define XLA_STREAM_EXECUTOR_GPU_GPU_EVENT_H_
 
+#include <cstdint>
+
 #include "absl/status/status.h"
 #include "xla/stream_executor/event.h"
 #include "xla/stream_executor/gpu/gpu_stream.h"
@@ -47,6 +49,8 @@ class GpuEvent : public Event {
   // The underlying CUDA event element.
   GpuEventHandle gpu_event();
 
+  absl::Status WaitForEventOnExternalStream(std::intptr_t stream) override;
+
  private:
   // The Executor used to which this object and GpuEventHandle are bound.
   GpuExecutor* parent_;

third_party/xla/xla/stream_executor/gpu/gpu_executor.h

@@ -237,9 +237,6 @@ class GpuExecutor : public StreamExecutor {
 
   absl::Status WaitForEvent(Stream* stream, Event* event) override;
 
-  absl::Status WaitForEventOnExternalStream(std::intptr_t stream,
-                                            Event* event) override;
-
   Event::Status PollForEventStatus(Event* event) override;
 
   absl::Status BlockHostUntilDone(Stream* stream) override;

third_party/xla/xla/stream_executor/mock_stream_executor.h

@@ -136,8 +136,6 @@ class MockStreamExecutor : public StreamExecutorInterface {
               (override));
   MOCK_METHOD(absl::Status, WaitForEvent, (Stream * stream, Event* event),
               (override));
-  MOCK_METHOD(absl::Status, WaitForEventOnExternalStream,
-              (std::intptr_t stream, Event* event), (override));
   MOCK_METHOD(Event::Status, PollForEventStatus, (Event * event), (override));
   MOCK_METHOD(void, DeallocateStream, (Stream * stream), (override));
   MOCK_METHOD(bool, CreateStreamDependency, (Stream * dependent, Stream* other),

third_party/xla/xla/stream_executor/rocm/rocm_executor.cc

@@ -676,18 +676,6 @@ absl::Status GpuExecutor::WaitForEvent(Stream* stream, Event* event) {
   }
 }
 
-absl::Status GpuExecutor::WaitForEventOnExternalStream(std::intptr_t stream,
-                                                       Event* event) {
-  if (GpuDriver::WaitStreamOnEvent(context_,
-                                   absl::bit_cast<GpuStreamHandle>(stream),
-                                   AsGpuEvent(event)->gpu_event())) {
-    return absl::OkStatus();
-  } else {
-    return absl::InternalError(
-        "error waiting for ROCM event on external stream");
-  }
-}
-
 Event::Status GpuExecutor::PollForEventStatus(Event* event) {
   return AsGpuEvent(event)->PollForStatus();
 }

third_party/xla/xla/stream_executor/stream_executor_interface.h

@@ -269,14 +269,6 @@ class StreamExecutorInterface {
   // Waits for the specified event at the end of the specified stream.
   virtual absl::Status WaitForEvent(Stream* stream, Event* event) = 0;
 
-  // Waits for the specified event at the end of the raw platform-specific
-  // stream.
-  virtual absl::Status WaitForEventOnExternalStream(std::intptr_t stream,
-                                                    Event* event) {
-    return absl::UnimplementedError(
-        "WaitForEventOnExternalStream not supported on this executor.");
-  }
-
   // Requests the current status of the event from the underlying platform.
   virtual Event::Status PollForEventStatus(Event* event) = 0;