Remove the need for StreamExecutorInterface::DeallocateEvent.

PiperOrigin-RevId: 635910366
tensorflow · May 21, 2024 · 883175b · 883175b
1 parent bb21389
commit 883175b
Show file tree

Hide file tree

Showing 17 changed files with 18 additions and 49 deletions.
diff --git a/tensorflow/c/experimental/stream_executor/stream_executor.cc b/tensorflow/c/experimental/stream_executor/stream_executor.cc
@@ -407,10 +407,6 @@ class CStreamExecutor : public StreamExecutor {
     return stream_executor_->host_callback(&device_, stream_handle,
                                            &HostCallbackTrampoline, ctx);
   }
-  absl::Status DeallocateEvent(Event* event) override {
-    static_cast<CEvent*>(event->implementation())->Destroy();
-    return absl::OkStatus();
-  }
   absl::Status RecordEvent(Stream* stream, Event* event) override {
     SP_Stream stream_handle =
         static_cast<CStream*>(stream->implementation())->Handle();

diff --git a/third_party/xla/xla/backends/interpreter/executor.h b/third_party/xla/xla/backends/interpreter/executor.h
@@ -114,10 +114,6 @@ class XlaInterpreterExecutor : public StreamExecutor {
   bool HostCallback(Stream *stream,
                     absl::AnyInvocable<absl::Status() &&> callback) override;
 
-  absl::Status DeallocateEvent(Event *event) override {
-    return absl::OkStatus();
-  }
-
   absl::Status RecordEvent(Stream *stream, Event *event) override {
     return absl::Status{absl::StatusCode::kUnimplemented, "RecordEvent"};
   }

diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc b/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc
@@ -758,10 +758,6 @@ bool GpuExecutor::HostCallback(Stream* stream,
   delete callback;
 }
 
-absl::Status GpuExecutor::DeallocateEvent(Event* event) {
-  return AsGpuEvent(event)->Destroy();
-}
-
 absl::Status GpuExecutor::RecordEvent(Stream* stream, Event* event) {
   return AsGpuEvent(event)->Record(AsGpuStream(stream));
 }

diff --git a/third_party/xla/xla/stream_executor/event.cc b/third_party/xla/xla/stream_executor/event.cc
@@ -30,15 +30,7 @@ Event::Event(StreamExecutorInterface* stream_exec,
              std::unique_ptr<EventInterface> implementation)
     : stream_exec_(stream_exec), implementation_(std::move(implementation)) {}
 
-Event::~Event() {
-  // Deal with nullptr implementation_, as this event may have been std::moved.
-  if (stream_exec_ && implementation_) {
-    auto status = stream_exec_->DeallocateEvent(this);
-    if (!status.ok()) {
-      LOG(ERROR) << status.message();
-    }
-  }
-}
+Event::~Event() = default;
 
 Event::Event(Event&&) = default;
 Event& Event::operator=(Event&&) = default;

diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_event.cc b/third_party/xla/xla/stream_executor/gpu/gpu_event.cc
@@ -27,7 +27,7 @@ namespace gpu {
 GpuEvent::GpuEvent(GpuExecutor* parent)
     : parent_(parent), gpu_event_(nullptr) {}
 
-GpuEvent::~GpuEvent() {}
+GpuEvent::~GpuEvent() { Destroy().IgnoreError(); }
 
 absl::Status GpuEvent::Init() {
   return GpuDriver::InitEvent(parent_->gpu_context(), &gpu_event_,

diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_executor.h b/third_party/xla/xla/stream_executor/gpu/gpu_executor.h
@@ -233,8 +233,6 @@ class GpuExecutor : public StreamExecutor {
 
   bool CreateStreamDependency(Stream* dependent, Stream* other) override;
 
-  absl::Status DeallocateEvent(Event* event) override;
-
   absl::Status RecordEvent(Stream* stream, Event* event) override;
 
   absl::Status WaitForEvent(Stream* stream, Event* event) override;

diff --git a/third_party/xla/xla/stream_executor/host/host_executor.cc b/third_party/xla/xla/stream_executor/host/host_executor.cc
@@ -244,10 +244,6 @@ static HostEvent* AsHostEvent(Event* event) {
   return static_cast<HostEvent*>(event->implementation());
 }
 
-absl::Status HostExecutor::DeallocateEvent(Event* /*event*/) {
-  return absl::OkStatus();
-}
-
 absl::Status HostExecutor::RecordEvent(Stream* stream, Event* event) {
   std::shared_ptr<absl::Notification> notification =
       AsHostEvent(event)->notification();

diff --git a/third_party/xla/xla/stream_executor/host/host_executor.h b/third_party/xla/xla/stream_executor/host/host_executor.h
@@ -119,7 +119,6 @@ class HostExecutor : public StreamExecutor {
   bool HostCallback(Stream* stream,
                     absl::AnyInvocable<absl::Status() &&> callback) override;
 
-  absl::Status DeallocateEvent(Event* event) override;
   absl::Status RecordEvent(Stream* stream, Event* event) override;
   absl::Status WaitForEvent(Stream* stream, Event* event) override;
   Event::Status PollForEventStatus(Event* event) override;

diff --git a/third_party/xla/xla/stream_executor/mock_stream_executor.h b/third_party/xla/xla/stream_executor/mock_stream_executor.h
@@ -133,7 +133,6 @@ class MockStreamExecutor : public StreamExecutorInterface {
   MOCK_METHOD(bool, HostCallback,
               (Stream * stream, absl::AnyInvocable<absl::Status() &&> callback),
               (override));
-  MOCK_METHOD(absl::Status, DeallocateEvent, (Event * event), (override));
   MOCK_METHOD(absl::Status, RecordEvent, (Stream * stream, Event* event),
               (override));
   MOCK_METHOD(absl::Status, WaitForEvent, (Stream * stream, Event* event),

diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_executor.cc b/third_party/xla/xla/stream_executor/rocm/rocm_executor.cc
@@ -660,10 +660,6 @@ bool GpuExecutor::HostCallback(Stream* stream,
   delete callback;
 }
 
-absl::Status GpuExecutor::DeallocateEvent(Event* event) {
-  return AsGpuEvent(event)->Destroy();
-}
-
 absl::Status GpuExecutor::RecordEvent(Stream* stream, Event* event) {
   return AsGpuEvent(event)->Record(AsGpuStream(stream));
 }

diff --git a/third_party/xla/xla/stream_executor/stream_executor_interface.h b/third_party/xla/xla/stream_executor/stream_executor_interface.h
@@ -263,9 +263,6 @@ class StreamExecutorInterface {
   virtual bool HostCallback(Stream* stream,
                             absl::AnyInvocable<absl::Status() &&> callback) = 0;
 
-  // Performs platform-specific deallocation and cleanup of an event.
-  virtual absl::Status DeallocateEvent(Event* event) = 0;
-
   // Inserts the specified event at the end of the specified stream.
   virtual absl::Status RecordEvent(Stream* stream, Event* event) = 0;
 

diff --git a/third_party/xla/xla/stream_executor/tpu/BUILD b/third_party/xla/xla/stream_executor/tpu/BUILD
@@ -262,6 +262,7 @@ cc_library(
         ":tpu_executor_c_api_hdrs",
         ":tpu_topology_external",
         "//xla/stream_executor",
+        "//xla/stream_executor:event_interface",
         "//xla/stream_executor:stream_executor_interface",
         "//xla/stream_executor:stream_interface",
         "@com_google_absl//absl/container:flat_hash_map",
@@ -480,6 +481,7 @@ cc_library(
         ":c_api_decl",
         ":tpu_topology_external",
         "//xla/stream_executor",
+        "//xla/stream_executor:event_interface",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/synchronization",

diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_event.h b/third_party/xla/xla/stream_executor/tpu/tpu_event.h
@@ -17,19 +17,25 @@ limitations under the License.
 
 #include "xla/stream_executor/event_interface.h"
 #include "xla/stream_executor/tpu/c_api_decl.h"
-#include "xla/stream_executor/tpu/tpu_api.h"
 #include "xla/stream_executor/tpu/tpu_executor_api.h"
+#include "xla/stream_executor/tpu/tpu_platform_interface.h"
 
 namespace stream_executor {
 namespace tpu {
 
 class TpuEvent : public EventInterface {
  public:
-  explicit TpuEvent(SE_Event* event) : event_(event) {}
-  ~TpuEvent() override { ExecutorApiFn()->TpuEvent_FreeFn(event_); }
+  explicit TpuEvent(SE_Event* event,
+                    tensorflow::tpu::TpuPlatformInterface* platform)
+      : event_(event), platform_(platform) {}
+  ~TpuEvent() override {
+    platform_->EraseEvent(this);
+    ExecutorApiFn()->TpuEvent_FreeFn(event_);
+  }
 
  private:
   SE_Event* event_;
+  tensorflow::tpu::TpuPlatformInterface* platform_;
 };
 
 }  // namespace tpu

diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_executor.cc b/third_party/xla/xla/stream_executor/tpu/tpu_executor.cc
@@ -92,11 +92,6 @@ bool TpuExecutor::CreateStreamDependency(Stream* dependent, Stream* other) {
       get_stream(other->implementation()));
 }
 
-absl::Status TpuExecutor::DeallocateEvent(Event* event) {
-  tpu_platform().EraseEvent(event->implementation());
-  return absl::OkStatus();
-}
-
 stream_executor::Event::Status TpuExecutor::PollForEventStatus(
     stream_executor::Event* event) {
   auto se_event = tpu_platform().LookupEvent(event->implementation());
@@ -137,7 +132,7 @@ absl::StatusOr<std::unique_ptr<Stream>> TpuExecutor::CreateStream(
 
 absl::StatusOr<std::unique_ptr<Event>> TpuExecutor::CreateEvent() {
   SE_Event* se_event = ExecutorApiFn()->TpuEvent_NewFn(executor_);
-  auto tpu_event = std::make_unique<TpuEvent>(se_event);
+  auto tpu_event = std::make_unique<TpuEvent>(se_event, platform_);
   tpu_platform().InsertEvent(tpu_event.get(), se_event);
 
   StatusHelper status;

diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_executor.h b/third_party/xla/xla/stream_executor/tpu/tpu_executor.h
@@ -79,8 +79,6 @@ class TpuExecutor : public tensorflow::tpu::TpuExecutorInterface {
 
   void Deallocate(DeviceMemoryBase* memory) override;
 
-  absl::Status DeallocateEvent(Event* event) override;
-
   bool DeviceMemoryUsage(int64_t* free, int64_t* total) const override;
 
   void DequeueOutfeed(int32_t outfeed_queue_index, absl::Span<uint8_t> bytes,

diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_platform.h b/third_party/xla/xla/stream_executor/tpu/tpu_platform.h
@@ -108,7 +108,7 @@ class TpuPlatform : public ::tensorflow::tpu::TpuPlatformInterface {
     mutex().Unlock();
     return stream;
   }
-  void EraseEvent(stream_executor::EventInterface* key);
+  void EraseEvent(stream_executor::EventInterface* key) override;
 
   SE_Platform* se_platform() const { return platform_; }
 

diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_platform_interface.h b/third_party/xla/xla/stream_executor/tpu/tpu_platform_interface.h
@@ -1,3 +1,4 @@
+#include "xla/stream_executor/event_interface.h"
 /* Copyright 2020 The OpenXLA Authors.
 
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -58,6 +59,8 @@ class TpuPlatformInterface : public stream_executor::Platform {
 
   virtual TpuRuntimeVersion version() const = 0;
 
+  virtual void EraseEvent(stream_executor::EventInterface* key) {};
+
   TpuTopologyExternal topology() {
     return TpuTopologyExternal(GetTopologyPtr());
   }