[xla] Migrate to PjRtFuture<>::MakePromise() API

PiperOrigin-RevId: 806057470

Author: ezhulenev

third_party/xla/xla/pjrt/BUILD

@@ -109,6 +109,7 @@ cc_library(
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/synchronization",
         "@local_tsl//tsl/profiler/lib:traceme",
     ],

third_party/xla/xla/pjrt/abstract_tracked_device_buffer.h

@@ -18,15 +18,21 @@ limitations under the License.
 
 #include <array>
 #include <memory>
+#include <vector>
 
 #include "absl/base/thread_annotations.h"
+#include "absl/functional/any_invocable.h"
 #include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
 #include "absl/synchronization/mutex.h"
+#include "xla/pjrt/device_event.h"
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/pjrt/pjrt_future.h"
 #include "xla/pjrt/raw_buffer.h"
+#include "xla/tsl/concurrency/async_value.h"
+#include "xla/tsl/concurrency/ref_count.h"
 
 namespace xla {
 
@@ -224,9 +230,8 @@ class CommonPjRtBuffer : public PjRtBuffer {
 
   absl::Status AcquireScopedRawBuffer(
       absl::AnyInvocable<absl::StatusOr<tsl::RCReference<PjRtDeviceEvent>>(
-                             tsl::RCReference<CommonPjRtRawBuffer> raw_buffer,
-                             std::vector<tsl::RCReference<tsl::AsyncValue>>
-                                 definition_events) &&>
+          tsl::RCReference<CommonPjRtRawBuffer> raw_buffer,
+          std::vector<tsl::RCReference<tsl::AsyncValue>> definition_events) &&>
           scoped_acquire,
       const char* caller_name = "AcquireScopedRawBuffer");
 
@@ -294,7 +299,7 @@ class CommonPjRtBuffer : public PjRtBuffer {
   }
 
   mutable absl::Mutex mu_;
-  PjRtFuture<>::Promise definition_promise_ ABSL_GUARDED_BY(mu_);
+  PjRtFuture<> definition_future_ ABSL_GUARDED_BY(mu_);
   PjRtMemorySpace* const memory_space_;
 
  private:

third_party/xla/xla/pjrt/common_pjrt_client.cc

@@ -1194,12 +1194,12 @@ PjRtFuture<> CommonPjRtBufferImpl::GetReadyFuture() {
     return PjRtFuture<>(InvalidArgument(
         "GetReadyFuture() called on deleted or donated buffer"));
   }
-  if (!definition_promise_) {
-    definition_promise_ =
-        device_buffer()->GetReadyFuturePromise(memory_space());
+  if (!definition_future_) {
+    auto promise = device_buffer()->GetReadyFuturePromise(memory_space());
+    definition_future_ = client()->CreateFutureFromUserPromise(
+        memory_space(), "CommonPjRtBuffer", "Await", std::move(promise));
   }
-  return client()->CreateFutureFromUserPromise(
-      memory_space(), "CommonPjRtBuffer", "Await", definition_promise_);
+  return definition_future_;
 }
 
 }  // namespace xla

third_party/xla/xla/pjrt/pjrt_future.h

@@ -217,6 +217,8 @@ class PjRtFutureBase : public PjRtFutureMoveControl<is_move_only> {
     }
   }
 
+  explicit operator bool() const { return static_cast<bool>(promise_); }
+
  protected:
   static constexpr bool IsMoveOnly() { return is_move_only; }
 

third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc

@@ -2194,18 +2194,20 @@ void PjRtStreamExecutorBuffer::CopyToRemoteDevice(
 
 PjRtFuture<> PjRtStreamExecutorBuffer::GetReadyFuture() {
   absl::InlinedVector<BufferSequencingEventRef, 2> definition_events;
-  PjRtFuture<>::Promise definition_promise;
+  PjRtFuture<>::MoveOnlyPromise definition_promise;
+  PjRtFuture<> definition_future;
   {
     absl::MutexLock lock(&mu_);
     if (device_buffer() == nullptr) {
       return PjRtFuture<>(InvalidArgument(
           "GetReadyFuture() called on deleted or donated buffer"));
     }
-    if (!definition_promise_) {
+    if (!definition_future_) {
       definition_events = device_buffer()->definition_events();
-      definition_promise_ = PjRtFuture<>::CreatePromise();
+      std::tie(definition_promise, definition_future_) =
+          PjRtFuture<>::MakePromise();
     }
-    definition_promise = definition_promise_;
+    definition_future = definition_future_;
   }
 
   if (!definition_events.empty()) {
@@ -2214,12 +2216,13 @@ PjRtFuture<> PjRtStreamExecutorBuffer::GetReadyFuture() {
     auto async_wait_for_events =
         [definition_events = std::move(definition_events),
          local_device_state = std::move(local_device_state),
-         definition_promise]() mutable {
+         definition_promise = std::make_shared<PjRtFuture<>::Promise>(
+             std::move(definition_promise))]() mutable {
           std::unique_ptr<se::Stream> stream;
           absl::Status defined_status =
               definition_events[0]->GetDefinedStatus();
           if (!defined_status.ok()) {
-            definition_promise.Set(defined_status);
+            definition_promise->Set(defined_status);
             return;
           }
           for (auto& event : definition_events) {
@@ -2242,28 +2245,29 @@ PjRtFuture<> PjRtStreamExecutorBuffer::GetReadyFuture() {
                  event_with_status = definition_events[0]]() mutable {
                   local_device_state->ReturnStreamToPool(
                       std::unique_ptr<se::Stream>(stream_ptr));
-                  definition_promise.Set(event_with_status->GetDefinedStatus());
+                  definition_promise->Set(
+                      event_with_status->GetDefinedStatus());
                 });
             if (!status.ok()) {
-              definition_promise.Set(status);
+              definition_promise->Set(status);
               return;
             }
           } else {
             // All events are already complete; set the `definition_promise`
             // with the status of the buffer's first definition event which may
             // have error status to propagate.
-            definition_promise.Set(definition_events[0]->GetDefinedStatus());
+            definition_promise->Set(definition_events[0]->GetDefinedStatus());
           }
         };
     first_definition_event->ExecuteOrAddToFutureTasks(
         absl::StrFormat("async_wait_for_events_%p", &async_wait_for_events),
         std::move(async_wait_for_events));
   }
 
-  return PjRtFuture<>(
-      std::move(definition_promise),
+  return PjRtFutureHelpers::WithProfiling(
+      std::move(definition_future),
       /*on_block_start=*/
-      []() {
+      [] {
         tsl::profiler::TraceMeProducer traceme(
             "PjRtStreamExecutorBuffer::Await");
         VLOG(3) << "PjRtStreamExecutorBuffer::Await";