Simplify how device events are handled. Because definition_stream() is always

provided and only used if the stream is defined, we can avoid tracking
usage_stream separately.

PiperOrigin-RevId: 814765407

Author: pschuh

xla/pjrt/pjrt_stream_executor_client.cc

@@ -1420,7 +1420,9 @@ PjRtStreamExecutorBuffer::Release(bool wait_for_operations_to_complete) {
       for (const auto& stream_and_event : events) {
         VLOG(4)
             << "Checking whether need to wait for stream_and_event: stream: "
-            << stream_and_event.stream
+            << (stream_and_event.event->IsDefined()
+                    ? stream_and_event.event->definition_stream()
+                    : nullptr)
             << "; event: " << &*stream_and_event.event
             << "; reference_held: " << stream_and_event.reference_held
             << "; is_predetermined_error: "
@@ -1494,7 +1496,7 @@ void PjRtStreamExecutorBuffer::ConvertUsageHold(TrackedDeviceBuffer* buffer,
                                                 bool reference_held) {
   absl::MutexLock lock(&mu_);
   CHECK(device_buffer() == buffer || device_buffer() == nullptr);
-  buffer->AddUsageEvent(usage_stream, std::move(event), reference_held);
+  buffer->AddUsageEvent(std::move(event), reference_held);
   DecrementUsage();
 }
 

xla/pjrt/tracked_device_buffer.cc

@@ -32,6 +32,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
+#include "xla/pjrt/device_event.h"
 #include "xla/pjrt/event_pool.h"
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/pjrt/pjrt_common.h"
@@ -257,31 +258,31 @@ void TrackedDeviceBuffer::ConfirmDonation() {
   ReleaseDeviceMemory();
 }
 
-void TrackedDeviceBuffer::AddUsageEvent(se::Stream* usage_stream,
-                                        BufferSequencingEventRef event,
+void TrackedDeviceBuffer::AddUsageEvent(BufferSequencingEventRef event,
                                         bool reference_held) {
   CHECK(in_use_);
 
   // If the event is 0, it means that the event is not recorded yet and the task
   // related to this event is deferred, so just add it.
   if (!event->IsDefined()) {
-    usage_events_.push_back({usage_stream, event, reference_held});
+    usage_events_.push_back({event, reference_held});
     return;
   }
+  auto* usage_stream = event->definition_stream();
 
   for (auto& existing : usage_events_) {
     // If the existing event is 0, it means that the event is not recorded yet
     // and the task related to this event is deferred, so don't replace it.
     if (!existing.event->IsDefined()) continue;
-    if (existing.stream == usage_stream) {
+    if (existing.event->definition_stream() == usage_stream) {
       if (*existing.event < *event) {
         existing.event = event;
         existing.reference_held = reference_held;
       }
       return;
     }
   }
-  usage_events_.push_back({usage_stream, event, reference_held});
+  usage_events_.push_back({event, reference_held});
 }
 
 TrackedDeviceBuffer::StreamAndEventContainer
@@ -312,6 +313,16 @@ tsl::RCReference<CommonPjRtRawBuffer> TrackedDeviceBuffer::GetRawBuffer(
       device_memory_);
 }
 
+void TrackedDeviceBuffer::AddUsageEvent(
+    tsl::RCReference<PjRtDeviceEvent> event) {
+  if (event) {
+    AddUsageEvent(
+        tensorflow::down_cast<PjRtStreamExecutorDeviceEvent*>(event.get())
+            ->event(),
+        true);
+  }
+}
+
 void GetDeviceBufferEvents(
     const TrackedDeviceBuffer& buffer, bool get_usage_events,
     absl::flat_hash_set<BufferSequencingEvent*>* events) {

xla/pjrt/tracked_device_buffer.h

@@ -163,6 +163,8 @@ class BufferSequencingEvent : tsl::AsyncPayload::KeepOnError {
     se::Stream* definition_stream;
   };
 
+  se::Stream* definition_stream() const { return event_->definition_stream; }
+
  private:
   uint64_t sequence_number() const;
 
@@ -218,8 +220,6 @@ class TrackedDeviceBuffer : public AbstractTrackedDeviceBuffer {
  public:
   // Helper object to keep track of usage of the buffer on streams.
   struct StreamAndEvent {
-    // A stream the buffer has been used on.
-    se::Stream* stream;
     // An event that is later than the most recent usage of the buffer on
     // stream.
     BufferSequencingEventRef event;
@@ -282,8 +282,7 @@ class TrackedDeviceBuffer : public AbstractTrackedDeviceBuffer {
   //                   reference to *this to stay live until after the host
   //                   is sure that the usage (transfer or execution) has
   //                   completed.
-  void AddUsageEvent(se::Stream* usage_stream, BufferSequencingEventRef event,
-                     bool reference_held);
+  void AddUsageEvent(BufferSequencingEventRef event, bool reference_held);
 
   using StreamAndEventContainer = absl::InlinedVector<StreamAndEvent, 3>;
   // Returns the set of streams that the buffer was used on, and for each stream
@@ -303,9 +302,7 @@ class TrackedDeviceBuffer : public AbstractTrackedDeviceBuffer {
   tsl::RCReference<CommonPjRtRawBuffer> GetRawBuffer(
       PjRtMemorySpace* memory_space) override;
 
-  void AddUsageEvent(tsl::RCReference<PjRtDeviceEvent> event) override {
-    LOG(FATAL) << "Implement";
-  }
+  void AddUsageEvent(tsl::RCReference<PjRtDeviceEvent> event) override;
 
   void Delete(PjRtMemorySpace* memory_space) override {
     LOG(FATAL) << "Implement";