Adapt Stablehlo to recent upstream Tosa changes.

Changes needed for llvm/llvm-project#130337,
llvm/llvm-project#129720,
and llvm/llvm-project#130340

FUTURE_COPYBARA_INTEGRATE_REVIEW=#23793 from Tixxx:tixxx/cp_sync 3a8a643
PiperOrigin-RevId: 738783395

Author: akuegel

third_party/stablehlo/temporary.patch

@@ -190,6 +190,11 @@ absl::Status CollectivePermuteStartThunk::Initialize(
                             params.executor->CreateEvent());
         receiver_barrier_events_.emplace(current_id, std::move(receiver_event));
       }
+      if (sender_barrier_events_.find(current_id) ==
+          sender_barrier_events_.end()) {
+        TF_ASSIGN_OR_RETURN(auto sender_event, params.executor->CreateEvent());
+        sender_barrier_events_.emplace(current_id, std::move(sender_event));
+      }
     }
     TF_ASSIGN_OR_RETURN(
         std::vector<DeviceBufferPair> device_buffers,
@@ -273,7 +278,7 @@ absl::Status CollectivePermuteStartThunk::RunCollective(
                                 config().replica_groups, config().group_mode));
 
     auto rendezvous_name = absl::StrFormat(
-        "rendezvous of collective-permute; run_id=%d; op id:%d; "
+        "rendezvous before calling collective-permute; run_id=%d; op id:%d; "
         "num_local_participants:%d",
         params.collective_params->run_id.ToInt(), config_.config.op_id,
         num_local_participants);
@@ -293,9 +298,48 @@ absl::Status CollectivePermuteStartThunk::RunCollective(
     }
   }
 
-  return ::xla::gpu::RunCollectivePermute(
+  auto status = ::xla::gpu::RunCollectivePermute(
       collectives, source_target, device_buffers, stream, comm_handle.comm,
       device_string, current_id, use_memcpy, recv_ptr_map_);
+
+  if (use_memcpy) {
+    std::optional<int64_t> source_id = source_target.source;
+    std::optional<int64_t> target_id = source_target.target;
+    // After the memcpy p2p is dispatched, the receiver needs to
+    // wait for the sender's event before proceeding to ensure
+    // data has been copied.
+    if (target_id) {
+      absl::MutexLock lock(&barrier_mutex_);
+      auto sender_event = sender_barrier_events_.find(current_id);
+      TF_RETURN_IF_ERROR(stream.RecordEvent(sender_event->second.get()));
+    }
+    TF_ASSIGN_OR_RETURN(
+        size_t num_local_participants,
+        GetNumLocalParticipants(*params.collective_params,
+                                config().replica_groups, config().group_mode));
+
+    auto rendezvous_name = absl::StrFormat(
+        "rendezvous after calling collective-permute; run_id=%d; op id:%d; "
+        "num_local_participants:%d",
+        params.collective_params->run_id.ToInt(), config_.config.op_id,
+        num_local_participants);
+    auto rendezvous_key = CallRendezvousKey{params.collective_params->run_id};
+
+    // Perform a rendezvous to make sure all senders have their events
+    // recorded.
+    Rendezvous(rendezvous_name, rendezvous_key, num_local_participants,
+               /*warn_stuck_timeout=*/absl::Seconds(20),
+               /*terminate_timeout=*/absl::Seconds(40));
+
+    // For receiving side, wait for the recorded event from the sending side.
+    if (source_id) {
+      absl::MutexLock lock(&barrier_mutex_);
+      auto sender_event = sender_barrier_events_.find(*source_id);
+      TF_RETURN_IF_ERROR(stream.WaitFor(sender_event->second.get()));
+    }
+  }
+
+  return status;
 }
 
 absl::Status RunCollectivePermute(

xla/backends/gpu/runtime/collective_permute_thunk.cc

@@ -122,6 +122,9 @@ class CollectivePermuteStartThunk : public CollectiveThunk {
   absl::Mutex barrier_mutex_;
   absl::flat_hash_map<int64_t, std::unique_ptr<se::Event>>
       receiver_barrier_events_;
+  absl::flat_hash_map<int64_t, std::unique_ptr<se::Event>>
+      sender_barrier_events_;
+
   bool p2p_memcpy_enabled_ = false;
   int64_t device_count_;
 };

xla/backends/gpu/runtime/collective_permute_thunk.h

@@ -237,6 +237,7 @@ cc_library(
         ":global_flags",
         ":rpc_helper",
         "//xla:status_macros",
+        "//xla/pjrt:pjrt_layout",
         "//xla/python/ifrt",
         "//xla/python/ifrt:client_impl_util",
         "//xla/python/ifrt:sharding_serdes",
@@ -258,6 +259,7 @@ cc_library(
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:cord",
         "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
         "@llvm-project//llvm:Support",
         "@tsl//tsl/profiler/lib:traceme",
@@ -400,6 +402,7 @@ cc_library(
         "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/cleanup",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/container:node_hash_set",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",

xla/python/ifrt_proxy/client/BUILD

@@ -36,6 +36,7 @@
 #include "absl/strings/str_join.h"
 #include "absl/strings/string_view.h"
 #include "absl/strings/substitute.h"
+#include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
 #include "llvm/Support/Casting.h"
 #include "xla/python/ifrt/array.h"
@@ -393,6 +394,13 @@ Future<> Array::GetReadyFuture() const {
   tsl::profiler::TraceMe traceme_ifrt_entrypoint(
       "IfrtProxyEntrypointArrayGetReadyFuture");
 
+  {
+    absl::MutexLock lock(&mu_);
+    if (deleted_ == DeletionState::kDeleted) {
+      return Future<>(absl::InvalidArgumentError("Already deleted array."));
+    }
+  }
+
   auto req = std::make_unique<CheckValueReadyRequest>();
   req->add_value_handles(handle_.handle);
 
@@ -405,6 +413,10 @@ Future<> Array::GetReadyFuture() const {
 }
 
 Future<> Array::Delete() {
+  {
+    absl::MutexLock lock(&mu_);
+    deleted_ = DeletionState::kDeleted;
+  }
   if (rpc_helper_->version().protocol_version() >= 5) {
     rpc_helper_->Batch(RpcHelper::kDeleteArray, handle_);
     return Future<>(absl::OkStatus());
@@ -429,6 +441,15 @@ Future<> Array::Delete() {
 bool Array::IsDeleted() const {
   tsl::profiler::TraceMe traceme_ifrt_entrypoint(
       "IfrtProxyEntrypointIsDeleted");
+  {
+    absl::MutexLock lock(&mu_);
+    if (deleted_ == DeletionState::kDeleted) {
+      return true;
+    }
+    if (deleted_ == DeletionState::kAlive) {
+      return false;
+    }
+  }
   if (GetGlobalClientFlags()->array_is_deleted_hack) {
     return false;
   }
@@ -438,6 +459,12 @@ bool Array::IsDeleted() const {
   absl::StatusOr<std::shared_ptr<IsArrayDeletedResponse>> response =
       rpc_helper_->IsArrayDeleted(std::move(req)).Await();
   if (response.ok()) {
+    absl::MutexLock lock(&mu_);
+    if ((*response)->deleted()) {
+      deleted_ = DeletionState::kDeleted;
+    } else {
+      deleted_ = DeletionState::kAlive;
+    }
     return (*response)->deleted();
   } else {
     LOG(ERROR) << "Internal error from proxy server during Array::IsDeleted(): "
@@ -486,7 +513,9 @@ Array::AssembleArrayFromSingleDeviceArrays(
           "not a xla::ifrt::proxy::Array.",
           rcref.get()));
     }
-    req->add_single_device_array_handles(array->handle_.handle);
+    TF_ASSIGN_OR_RETURN(ArrayHandle handle,
+                        array->GetHandle(array_copy_semantics));
+    req->add_single_device_array_handles(handle.handle);
   }
 
   ArrayHandle result_handle;
@@ -569,8 +598,8 @@ Array::RemapArrays(xla::ifrt::Client* client,
                           *arrays[i]->sharding().devices(),
                           arrays[i]->sharding().memory_kind()));
     }
-
-    req->add_array_handles(array->handle_.handle);
+    TF_ASSIGN_OR_RETURN(ArrayHandle handle, array->GetHandle(semantics));
+    req->add_array_handles(handle.handle);
   }
 
   std::vector<ArrayHandle> result_handles;
@@ -617,7 +646,8 @@ Array::DisassembleIntoSingleDeviceArrays(
         "version < 8");
   }
   auto req = std::make_unique<DisassembleIntoSingleDeviceArraysRequest>();
-  req->set_array_handle(handle_.handle);
+  TF_ASSIGN_OR_RETURN(ArrayHandle handle, GetHandle(array_copy_semantics));
+  req->set_array_handle(handle.handle);
   req->set_copy_semantics(ToArrayCopySemanticsProto(array_copy_semantics));
   req->set_single_device_shard_semantics(
       ToSingleDeviceShardSemanticsProto(single_device_shard_semantics));
@@ -665,7 +695,8 @@ absl::StatusOr<tsl::RCReference<xla::ifrt::Array>> Array::FullyReplicatedShard(
   tsl::profiler::TraceMe traceme_ifrt_entrypoint(
       "IfrtProxyEntrypointFullyReplicatedShard");
   auto req = std::make_unique<FullyReplicatedShardRequest>();
-  req->set_array_handle(handle_.handle);
+  TF_ASSIGN_OR_RETURN(ArrayHandle handle, GetHandle(semantics));
+  req->set_array_handle(handle.handle);
   req->set_copy_semantics(ToArrayCopySemanticsProto(semantics));
 
   ArrayHandle result_handle;
@@ -706,7 +737,11 @@ Future<> Array::CopyToStringHostBuffer(
         "String arrays are not supported in ifrt-proxy version < 9"));
   }
   auto req = std::make_unique<CopyToHostBufferRequest>();
-  req->set_array_handle(handle_.handle);
+  absl::StatusOr<ArrayHandle> handle = GetHandle(semantics);
+  if (!handle.ok()) {
+    return Future<>(handle.status());
+  }
+  req->set_array_handle(handle->handle);
   if (byte_strides.has_value()) {
     return Future<>(absl::InvalidArgumentError(
         "Byte strides are not supported for string arrays."));
@@ -768,7 +803,11 @@ Future<> Array::CopyToHostBuffer(
   }
 
   auto req = std::make_unique<CopyToHostBufferRequest>();
-  req->set_array_handle(handle_.handle);
+  absl::StatusOr<ArrayHandle> handle = GetHandle(semantics);
+  if (!handle.ok()) {
+    return Future<>(handle.status());
+  }
+  req->set_array_handle(handle->handle);
   if (byte_strides.has_value()) {
     *req->mutable_byte_strides() = ToByteStridesProto(*byte_strides);
   }
@@ -829,8 +868,23 @@ Future<> Array::CopyToHostBuffer(
 xla::ifrt::Client* Array::client() const { return client_; }
 
 std::string Array::DebugString() const {
-  return absl::Substitute("proxy::Array, this=$0, handle=$1", this,
-                          handle_.handle);
+  std::string is_deleted;
+  {
+    absl::MutexLock l(&mu_);
+    switch (deleted_) {
+      case DeletionState::kUnknown:
+        is_deleted = "unknown";
+        break;
+      case DeletionState::kDeleted:
+        is_deleted = "true";
+        break;
+      case DeletionState::kAlive:
+        is_deleted = "false";
+        break;
+    }
+  }
+  return absl::Substitute("proxy::Array, this=$0, handle=$1, deleted=$2", this,
+                          handle_.handle, is_deleted);
 }
 
 }  // namespace proxy

xla/python/ifrt_proxy/client/array.cc

@@ -26,10 +26,13 @@
 #include <vector>
 
 #include "absl/base/attributes.h"
+#include "absl/base/thread_annotations.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
+#include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
 #include "llvm/Support/ExtensibleRTTI.h"
+#include "xla/pjrt/pjrt_layout.h"
 #include "xla/python/ifrt/array.h"
 #include "xla/python/ifrt/client.h"
 #include "xla/python/ifrt/dtype.h"
@@ -106,7 +109,31 @@ class Array final : public llvm::RTTIExtends<Array, xla::ifrt::Array> {
 
   ~Array() override { Destruct(rpc_helper_.get(), handle_); }
 
-  ArrayHandle handle() const { return handle_; }
+  absl::StatusOr<ArrayHandle> GetHandle(ArrayCopySemantics semantics) {
+    absl::MutexLock l(&mu_);
+    if (deleted_ == DeletionState::kDeleted) {
+      return absl::InvalidArgumentError("Array already deleted.");
+    }
+    if (semantics == ArrayCopySemantics::kDonateInput) {
+      deleted_ = DeletionState::kDeleted;
+    }
+    return handle_;
+  }
+
+  // Fetches the ArrayHandle when the ArrayCopySemantics (i.e., whether the
+  // array is meant to be donated or copied) is not known.
+  //
+  // Calling this function may cause `IsDelete()` calls to result in a
+  // synchronous RPC to the proxy-server. To avoid such performance overhead,
+  // prefer using `GetHandle(semantics)` whenever the semantics are known.
+  absl::StatusOr<ArrayHandle> GetHandleUnknownIfBeingDonated() {
+    absl::MutexLock l(&mu_);
+    if (deleted_ == DeletionState::kDeleted) {
+      return absl::InvalidArgumentError("Array already deleted.");
+    }
+    deleted_ = DeletionState::kUnknown;
+    return handle_;
+  }
 
   xla::ifrt::Client* client() const override;
   Future<> GetReadyFuture() const override;
@@ -158,7 +185,17 @@ class Array final : public llvm::RTTIExtends<Array, xla::ifrt::Array> {
   const DType dtype_;
   const Shape shape_;
   const std::shared_ptr<const Sharding> sharding_;
-  const ArrayHandle handle_;
+
+  const ArrayHandle handle_
+      ABSL_DEPRECATED("Use GetHandle() function instead.");
+
+  mutable absl::Mutex mu_;
+  enum class DeletionState {
+    kUnknown,  // Need to ask the proxy-server whether the array is deleted.
+    kDeleted,  // IsDeleted() will return true.
+    kAlive     // IsDeleted() will return false.
+  };
+  mutable DeletionState deleted_ ABSL_GUARDED_BY(mu_) = DeletionState::kAlive;
 };
 
 }  // namespace proxy

xla/python/ifrt_proxy/client/array.h

@@ -277,7 +277,9 @@ Client::CopyArrays(absl::Span<tsl::RCReference<xla::ifrt::Array>> arrays,
   for (const auto& array : arrays) {
     if (auto* proxy_array =
             llvm::dyn_cast<xla::ifrt::proxy::Array>(array.get())) {
-      req->add_array_handles(proxy_array->handle().handle);
+      TF_ASSIGN_OR_RETURN(ArrayHandle handle,
+                          proxy_array->GetHandle(semantics));
+      req->add_array_handles(handle.handle);
     } else {
       return absl::InvalidArgumentError(
           "CopyArrays only supports arrays created via IFRT Proxy client");
@@ -351,7 +353,13 @@ xla::ifrt::Future<> Client::GetReadyFuture(
     // type, but this may be extended later to other types such as Tuples.
     if (auto proxy_array =
             llvm::dyn_cast<xla::ifrt::proxy::Array>(value.get())) {
-      req->add_value_handles(proxy_array->handle().handle);
+      absl::StatusOr<ArrayHandle> handle =
+          proxy_array->GetHandle(ArrayCopySemantics::kAlwaysCopy);
+      if (!handle.ok()) {
+        futures.push_back(Future<>(handle.status()));
+      } else {
+        req->add_value_handles(handle->handle);
+      }
     } else {
       futures.push_back(value->GetReadyFuture());
     }