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

PiperOrigin-RevId: 806070441

Author: ezhulenev

third_party/xla/xla/pjrt/c/pjrt_c_api_helpers.cc

@@ -467,14 +467,15 @@ xla::PjRtFuture<> ConvertCEventToCppFuture(PJRT_Event* c_event,
   event_onready_args.extension_start = nullptr;
   event_onready_args.event = c_event;
 
-  PjRtFuture<>::Promise promise = PjRtFuture<>::CreatePromise();
+  auto [promise, future] = PjRtFuture<>::MakePromise();
   event_onready_args.user_arg = new std::function<void(PJRT_Error*)>(
-      [promise, c_event, c_api](PJRT_Error* error) mutable {
+      [promise = std::make_shared<PjRtFuture<>::Promise>(std::move(promise)),
+       c_event, c_api](PJRT_Error* error) mutable {
         if (error != nullptr) {
-          promise.Set(::pjrt::PjrtErrorToStatus(error, c_api));
+          promise->Set(::pjrt::PjrtErrorToStatus(error, c_api));
           ::pjrt::MakeErrorDeleter(c_api)(error);
         } else {
-          promise.Set();
+          promise->Set();
         }
         ::pjrt::MakeEventDeleter(c_api)(c_event);
       });
@@ -489,7 +490,7 @@ xla::PjRtFuture<> ConvertCEventToCppFuture(PJRT_Event* c_event,
   if (error != nullptr) {
     return PjRtFuture<>(::pjrt::PjrtErrorToStatus(error, c_api));
   }
-  return PjRtFuture<>(std::move(promise));
+  return std::move(future);
 }
 
 static absl::StatusOr<PJRT_NamedValue> ConvertToPjRtNamedValue(

third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc

@@ -692,15 +692,14 @@ PJRT_Error* PJRT_AsyncHostToDeviceTransferManager_TransferData(
       "PJRT_AsyncHostToDeviceTransferManager_TransferData_Args",
       PJRT_AsyncHostToDeviceTransferManager_TransferData_Args_STRUCT_SIZE,
       args->struct_size));
-  xla::PjRtFuture<>::Promise promise = xla::PjRtFuture<>::CreatePromise();
+  auto [promise, future] = xla::PjRtFuture<>::MakePromise();
   absl::AnyInvocable<void() &&> on_done_with_d2h_transfer =
-      [promise]() mutable { promise.Set(); };
+      [promise = std::move(promise)]() mutable { promise.Set(); };
   PJRT_RETURN_IF_ERROR(
       args->transfer_manager->transfer_manager->TransferRawDataToSubBuffer(
           args->buffer_index, args->data, args->offset, args->transfer_size,
           args->is_last_transfer, std::move(on_done_with_d2h_transfer)));
-  args->done_with_h2d_transfer =
-      new PJRT_Event{xla::PjRtFuture<>(std::move(promise))};
+  args->done_with_h2d_transfer = new PJRT_Event{std::move(future)};
   return nullptr;
 }
 

third_party/xla/xla/pjrt/se_raw_buffer.cc

@@ -37,20 +37,19 @@ limitations under the License.
 namespace xla {
 
 PjRtFuture<> PjRtStreamExecutorDeviceEvent::GetReadyFuture() {
-  PjRtFuture<>::Promise promise = PjRtFuture<>::CreatePromise();
-  event_.AndThen([promise, event = event_]() mutable {
+  auto [promise, future] = PjRtFuture<>::MakePromise();
+  event_.AndThen([promise = std::move(promise), event = event_]() mutable {
     if (auto* error = event.GetErrorIfPresent()) {
       promise.Set(*error);
     } else {
       promise.Set();
     }
   });
 
-  return PjRtFuture<>(
-      promise,
+  return PjRtFutureHelpers::WithProfiling(
+      std::move(future),
       /*on_block_start=*/
-      [ready_event = FormRef(promise.async_value()),
-       callee_method = callee_method_, callee_type = callee_type_]() {
+      [callee_method = callee_method_, callee_type = callee_type_]() {
         tsl::profiler::TraceMeProducer traceme(
             [&] { return absl::StrCat(callee_type, "::", callee_method); });
         return PjRtFutureHelpers::ProfilingKeys({traceme.GetContextId()});