[IFRT] Add "PjRt-compatible" interfaces for PjRt escape hatches

This change helps migrate various PjRt implementations that do not directly
use/wrap around a low-level PjRt runtime. Implementing the interface will
preserve the compatibility between the JAX Python binding and such
PjRt-compatible IFRT implementations.

PiperOrigin-RevId: 494186705

Author: hyeontaek

xla/python/ifrt/client.h

@@ -99,7 +99,7 @@ class Client : public llvm::RTTIExtends<Client, llvm::RTTIRoot> {
 
   // TODO(hyeontaek): Potentially remove this method to encourage supporting
   // only ahead-of-time compilation.
-  virtual Compiler* GetDefaultCompiler() const = 0;
+  virtual Compiler* GetDefaultCompiler() = 0;
 
   static char ID;  // NOLINT
 };

xla/python/pjrt_ifrt/pjrt_array.cc

@@ -91,8 +91,8 @@ StatusOr<DType> ToDType(xla::PrimitiveType primitive_type) {
 StatusOr<std::unique_ptr<Array>> PjRtArray::Create(
     Client* client, DType dtype, Shape shape,
     std::shared_ptr<const Sharding> sharding, PjRtBuffers pjrt_buffers) {
-  if (!llvm::isa_and_nonnull<PjRtClient>(client)) {
-    return InvalidArgument("PjRtClient expected");
+  if (!llvm::isa_and_nonnull<PjRtCompatibleClient>(client)) {
+    return InvalidArgument("PjRtCompatibleClient expected");
   }
   if (pjrt_buffers.empty()) {
     return InvalidArgument("pjrt_buffers must be non-empty");
@@ -101,22 +101,22 @@ StatusOr<std::unique_ptr<Array>> PjRtArray::Create(
     return InvalidArgument("device and buffer counts mismatch: %d vs. %d",
                            sharding->devices().size(), pjrt_buffers.size());
   }
-  return std::unique_ptr<Array>(
-      new PjRtArray(static_cast<PjRtClient*>(client), dtype, std::move(shape),
-                    std::move(sharding), std::move(pjrt_buffers)));
+  return std::unique_ptr<Array>(new PjRtArray(
+      static_cast<PjRtCompatibleClient*>(client), dtype, std::move(shape),
+      std::move(sharding), std::move(pjrt_buffers)));
 }
 
 StatusOr<std::unique_ptr<Array>> PjRtArray::Create(
     Client* client, std::shared_ptr<PjRtBuffer> pjrt_buffer) {
-  if (!llvm::isa_and_nonnull<PjRtClient>(client)) {
-    return InvalidArgument("PjRtClient expected");
+  if (!llvm::isa_and_nonnull<PjRtCompatibleClient>(client)) {
+    return InvalidArgument("PjRtCompatibleClient expected");
   }
   TF_ASSIGN_OR_RETURN(auto dtype,
                       ToDType(pjrt_buffer->on_device_shape().element_type()));
   Shape shape(pjrt_buffer->on_device_shape().dimensions());
   auto sharding = SingleDeviceSharding::Create(pjrt_buffer->device());
   return std::unique_ptr<Array>(new PjRtArray(
-      static_cast<PjRtClient*>(client), dtype, std::move(shape),
+      static_cast<PjRtCompatibleClient*>(client), dtype, std::move(shape),
       std::move(sharding), PjRtBuffers({std::move(pjrt_buffer)})));
 }
 
@@ -149,7 +149,7 @@ StatusOr<std::unique_ptr<Array>> PjRtArray::Create(Client* client, Shape shape,
       std::move(pjrt_buffers));
 }
 
-PjRtArray::PjRtArray(PjRtClient* client, DType dtype, Shape shape,
+PjRtArray::PjRtArray(PjRtCompatibleClient* client, DType dtype, Shape shape,
                      std::shared_ptr<const Sharding> sharding,
                      PjRtBuffers pjrt_buffers)
     : client_(client),

xla/python/pjrt_ifrt/pjrt_array.h

@@ -36,8 +36,19 @@ StatusOr<xla::PrimitiveType> ToPrimitiveType(DType dtype);
 // Converts `xla::PrimitiveType` into IFRT `DType`.
 StatusOr<DType> ToDType(xla::PrimitiveType primitive_type);
 
+// PjRt-compatible `Array` interface that wraps a list of `xla::PjRtBuffer`s.
+class PjRtCompatibleArray
+    : public llvm::RTTIExtends<PjRtCompatibleArray, Array> {
+ public:
+  // APIs that allow direct access to `PjRtBuffer` for PjRt-only operations.
+  virtual absl::Span<const std::shared_ptr<PjRtBuffer>> pjrt_buffers() = 0;
+  virtual StatusOr<absl::Span<std::shared_ptr<PjRtBuffer>>>
+  mutable_pjrt_buffers() = 0;
+};
+
 // `Array` implementation that wraps a list of `xla::PjRtBuffer`s.
-class PjRtArray final : public llvm::RTTIExtends<PjRtArray, Array> {
+class PjRtArray final
+    : public llvm::RTTIExtends<PjRtArray, PjRtCompatibleArray> {
  public:
   static constexpr int kPjRtBufferInlineSize = 1;
   using PjRtBuffers =
@@ -60,26 +71,25 @@ class PjRtArray final : public llvm::RTTIExtends<PjRtArray, Array> {
   static StatusOr<std::unique_ptr<Array>> Create(Client* client, Shape shape,
                                                  PjRtBuffers pjrt_buffers);
 
-  absl::Span<const std::shared_ptr<PjRtBuffer>> pjrt_buffers() const {
+  // PjRtCompatibleArray implementation.
+
+  absl::Span<const std::shared_ptr<PjRtBuffer>> pjrt_buffers() override {
     DCHECK(this);
     return pjrt_buffers_;
   }
-  absl::Span<std::shared_ptr<PjRtBuffer>> pjrt_buffers() {
+  StatusOr<absl::Span<std::shared_ptr<PjRtBuffer>>> mutable_pjrt_buffers()
+      override {
     DCHECK(this);
     return absl::MakeSpan(pjrt_buffers_);
   }
-  PjRtBuffer* pjrt_buffer(int device_id) const {
-    DCHECK(this);
-    return pjrt_buffers_[device_id].get();
-  }
 
   // Array implementation.
 
   ~PjRtArray() override = default;
 
   Client* client() const override {
     DCHECK(this);
-    return const_cast<PjRtClient*>(client_);
+    return client_;
   }
 
   DType dtype() const override {
@@ -121,10 +131,10 @@ class PjRtArray final : public llvm::RTTIExtends<PjRtArray, Array> {
   static char ID;  // NOLINT
 
  private:
-  PjRtArray(PjRtClient* client, DType dtype, Shape shape,
+  PjRtArray(PjRtCompatibleClient* client, DType dtype, Shape shape,
             std::shared_ptr<const Sharding> sharding, PjRtBuffers pjrt_buffers);
 
-  PjRtClient* client_;
+  PjRtCompatibleClient* client_;
   DType dtype_;
   Shape shape_;
   std::shared_ptr<const Sharding> sharding_;

xla/python/pjrt_ifrt/pjrt_client.cc

@@ -85,7 +85,7 @@ PjRtClient::AssembleArrayFromSingleDeviceArrays(
       return InvalidArgument("Only PjRtArray is supported: arrays[%d]=%s", i,
                              arrays[i]->DebugString());
     }
-    const auto* array = static_cast<const PjRtArray*>(arrays[i]);
+    auto* array = static_cast<PjRtArray*>(arrays[i]);
     if (array->dtype() != dtype) {
       return InvalidArgument(
           "Every input must have the same dtype: %s (shard 0) vs. %s (shard "

xla/python/pjrt_ifrt/pjrt_client.h

@@ -31,16 +31,29 @@ limitations under the License.
 namespace xla {
 namespace ifrt {
 
+// PjRt-compatible `Client` interface.
+class PjRtCompatibleClient
+    : public llvm::RTTIExtends<PjRtCompatibleClient, Client> {
+ public:
+  // APIs that allow direct access to `xla::PjRtClient` for PjRt-only
+  // operations.
+  virtual xla::PjRtClient* pjrt_client() = 0;
+  virtual std::shared_ptr<xla::PjRtClient> shared_ptr_pjrt_client() = 0;
+};
+
 // `Client` implementation that wraps `xla::PjRtClient`.
-class PjRtClient final : public llvm::RTTIExtends<PjRtClient, Client> {
+class PjRtClient final
+    : public llvm::RTTIExtends<PjRtClient, PjRtCompatibleClient> {
  public:
   static std::unique_ptr<ifrt::Client> Create(
       std::shared_ptr<xla::PjRtClient> pjrt_client);
   static std::unique_ptr<ifrt::Client> Create(
       std::unique_ptr<xla::PjRtClient> pjrt_client);
 
-  xla::PjRtClient* pjrt_client() const { return pjrt_client_.get(); }
-  std::shared_ptr<xla::PjRtClient> shared_ptr_pjrt_client() const {
+  // PjRtCompatibleClient implementation.
+
+  xla::PjRtClient* pjrt_client() override { return pjrt_client_.get(); }
+  std::shared_ptr<xla::PjRtClient> shared_ptr_pjrt_client() override {
     return pjrt_client_;
   }
 
@@ -114,9 +127,9 @@ class PjRtClient final : public llvm::RTTIExtends<PjRtClient, Client> {
     return pjrt_client_->CreateHostToDeviceChannelHandle();
   }
 
-  Compiler* GetDefaultCompiler() const override {
+  Compiler* GetDefaultCompiler() override {
     DCHECK(this);
-    return const_cast<PjRtCompiler*>(&default_compiler_);
+    return &default_compiler_;
   }
 
   static char ID;  // NOLINT

xla/python/pjrt_ifrt/pjrt_executable.h

@@ -31,17 +31,40 @@ limitations under the License.
 namespace xla {
 namespace ifrt {
 
+// PjRt-compatible `Executable` interface.
+class PjRtCompatibleExecutable
+    : public llvm::RTTIExtends<PjRtCompatibleExecutable, Executable> {
+ public:
+  // APIs that allow direct access to `xla::PjRtExecutable` for PjRt-only
+  // operations.
+  virtual xla::PjRtExecutable* pjrt_executable() = 0;
+};
+
+// PjRt-compatible `LoadedExecutable` interface.
+class PjRtCompatibleLoadedExecutable
+    : public llvm::RTTIExtends<PjRtCompatibleLoadedExecutable,
+                               LoadedExecutable> {
+ public:
+  // APIs that allow direct access to `xla::PjRtLoadedExecutable` for PjRt-only
+  // operations.
+  virtual xla::PjRtLoadedExecutable* pjrt_loaded_executable() = 0;
+  virtual std::shared_ptr<xla::PjRtLoadedExecutable>
+  shared_ptr_pjrt_loaded_executable() = 0;
+};
+
 // `Executable` implementation that wraps a `xla::PjRtExecutable`.
 class PjRtExecutable final
-    : public llvm::RTTIExtends<PjRtExecutable, Executable> {
+    : public llvm::RTTIExtends<PjRtExecutable, PjRtCompatibleExecutable> {
  public:
   // Creates PjRtExecutable from xla::PjRtExecutable.
   static StatusOr<std::unique_ptr<Executable>> Create(
       std::unique_ptr<xla::PjRtExecutable> pjrt_executable);
   static StatusOr<std::unique_ptr<Executable>> Create(
       std::shared_ptr<xla::PjRtExecutable> pjrt_executable);
 
-  xla::PjRtExecutable* pjrt_executable() const {
+  // PjRtCompatibleExecutable implementation.
+
+  xla::PjRtExecutable* pjrt_executable() override {
     DCHECK(this);
     return pjrt_executable_.get();
   }
@@ -101,7 +124,8 @@ class PjRtExecutable final
 
 // `LoadedExecutable` implementation that wraps a `xla::PjRtLoadedExecutable`.
 class PjRtLoadedExecutable final
-    : public llvm::RTTIExtends<PjRtLoadedExecutable, LoadedExecutable> {
+    : public llvm::RTTIExtends<PjRtLoadedExecutable,
+                               PjRtCompatibleLoadedExecutable> {
  public:
   using LoadedExecutable::ExecuteOptions;
   using LoadedExecutable::ExecuteResult;
@@ -125,12 +149,14 @@ class PjRtLoadedExecutable final
       PjRtClient* client, const XlaComputation& computation,
       CompileOptions options);
 
-  xla::PjRtLoadedExecutable* pjrt_loaded_executable() const {
+  // PjRtCompatibleLoadedExecutable implementation.
+
+  xla::PjRtLoadedExecutable* pjrt_loaded_executable() override {
     DCHECK(this);
     return pjrt_loaded_executable_.get();
   }
   std::shared_ptr<xla::PjRtLoadedExecutable> shared_ptr_pjrt_loaded_executable()
-      const {
+      override {
     DCHECK(this);
     return pjrt_loaded_executable_;
   }
@@ -181,7 +207,7 @@ class PjRtLoadedExecutable final
 
   Client* client() const override {
     DCHECK(this);
-    return const_cast<PjRtClient*>(client_);
+    return client_;
   }
   StatusOr<ExecuteResult> Execute(absl::Span<Array* const> args,
                                   const ExecuteOptions& options,
@@ -190,21 +216,21 @@ class PjRtLoadedExecutable final
   Future<Status> Delete() override;
   bool IsDeleted() const override {
     DCHECK(this);
-    return pjrt_loaded_executable()->IsDeleted();
+    return pjrt_loaded_executable_->IsDeleted();
   }
 
   const DeviceAssignment& device_assignment() const override {
     DCHECK(this);
-    return pjrt_loaded_executable()->device_assignment();
+    return pjrt_loaded_executable_->device_assignment();
   }
   absl::Span<const LoadedExecutable::LogicalDeviceIds>
   addressable_device_logical_ids() const override {
     DCHECK(this);
-    return pjrt_loaded_executable()->addressable_device_logical_ids();
+    return pjrt_loaded_executable_->addressable_device_logical_ids();
   }
   absl::Span<Device* const> addressable_devices() const override {
     DCHECK(this);
-    return pjrt_loaded_executable()->addressable_devices();
+    return pjrt_loaded_executable_->addressable_devices();
   }
 
   static char ID;  // NOLINT

xla/python/py_array.h

@@ -181,7 +181,8 @@ class PyArray : public pybind11::object {
     if (ifrt_array_ptr == nullptr) {
       return {};
     }
-    auto* arr = llvm::dyn_cast_or_null<ifrt::PjRtArray>(ifrt_array_ptr);
+    auto* arr =
+        llvm::dyn_cast_or_null<ifrt::PjRtCompatibleArray>(ifrt_array_ptr);
     if (arr == nullptr) {
       throw XlaRuntimeError(
           "This operation is implemented for a PjRt-compatible backend only.");

xla/python/py_buffer.h

@@ -87,7 +87,8 @@ class PyBuffer {
   // Short-term escape hatch to get PjRtBuffer from PyBuffer.
   // TODO(hyeontaek): Migrate all users of this method to be agnostic of PjRt.
   PjRtBuffer* pjrt_buffer() const {
-    auto* arr = llvm::dyn_cast_or_null<ifrt::PjRtArray>(ifrt_array_.get());
+    auto* arr =
+        llvm::dyn_cast_or_null<ifrt::PjRtCompatibleArray>(ifrt_array_.get());
     if (arr == nullptr) {
       throw XlaRuntimeError(
           "This operation is implemented for a PjRt-compatible backend only.");
@@ -98,7 +99,8 @@ class PyBuffer {
   // Short-term escape hatch to get PjRtBuffer from PyBuffer.
   // TODO(hyeontaek): Migrate all users of this method to be agnostic of PjRt.
   std::shared_ptr<PjRtBuffer> shared_ptr_pjrt_buffer() const {
-    auto* arr = llvm::dyn_cast_or_null<ifrt::PjRtArray>(ifrt_array_.get());
+    auto* arr =
+        llvm::dyn_cast_or_null<ifrt::PjRtCompatibleArray>(ifrt_array_.get());
     if (arr == nullptr) {
       throw XlaRuntimeError(
           "This operation is implemented for a PjRt-compatible backend only.");
@@ -406,7 +408,8 @@ class PyShardedBuffer {
 #ifdef JAX_ENABLE_IFRT
   PyBuffer::object GetPyBuffer(int device_id) const {
     // TODO(hyeontaek): Remove this method. This method will not scale well.
-    auto* arr = llvm::dyn_cast_or_null<ifrt::PjRtArray>(ifrt_array_.get());
+    auto* arr =
+        llvm::dyn_cast_or_null<ifrt::PjRtCompatibleArray>(ifrt_array_.get());
     if (arr == nullptr) {
       throw XlaRuntimeError(
           "This operation is implemented for a PjRt-compatible backend only.");
@@ -446,7 +449,8 @@ class PyShardedBuffer {
   // Short-term escape hatch to get PjRtBuffer from PyShardedBuffer.
   // TODO(hyeontaek): Migrate all users of this method to be agnostic of PjRt.
   PjRtBuffer* pjrt_buffer(int device_id) const {
-    auto* arr = llvm::dyn_cast_or_null<ifrt::PjRtArray>(ifrt_array_.get());
+    auto* arr =
+        llvm::dyn_cast_or_null<ifrt::PjRtCompatibleArray>(ifrt_array_.get());
     if (arr == nullptr) {
       throw XlaRuntimeError(
           "This operation is implemented for a PjRt-compatible backend only.");

xla/python/py_client.cc

@@ -206,15 +206,15 @@ Status PyClient::Defragment() {
       if (array->ifrt_array == nullptr) {
         continue;
       }
-      auto* arr =
-          llvm::dyn_cast_or_null<ifrt::PjRtArray>(array->ifrt_array.get());
+      auto* arr = llvm::dyn_cast_or_null<ifrt::PjRtCompatibleArray>(
+          array->ifrt_array.get());
       if (arr == nullptr) {
         throw XlaRuntimeError(
             "This operation is implemented for a PjRt-compatible backend "
             "only.");
       }
-      absl::Span<std::shared_ptr<PjRtBuffer>> pjrt_buffers =
-          arr->pjrt_buffers();
+      TF_ASSIGN_OR_RETURN(absl::Span<std::shared_ptr<PjRtBuffer>> pjrt_buffers,
+                          arr->mutable_pjrt_buffers());
 #else
       absl::Span<std::shared_ptr<PjRtBuffer>> pjrt_buffers =
           absl::MakeSpan(array->pjrt_buffers);
@@ -669,8 +669,8 @@ StatusOr<py::bytes> PyClient::HeapProfile() {
     if (array->ifrt_array == nullptr) {
       continue;
     }
-    auto* arr =
-        llvm::dyn_cast_or_null<ifrt::PjRtArray>(array->ifrt_array.get());
+    auto* arr = llvm::dyn_cast_or_null<ifrt::PjRtCompatibleArray>(
+        array->ifrt_array.get());
     // TODO(hyeontaek): Support non-PjRt Arrays.
     if (arr == nullptr) {
       throw XlaRuntimeError(

xla/python/py_client.h

@@ -112,7 +112,7 @@ class PyClient : public std::enable_shared_from_this<PyClient> {
   // TODO(hyeontaek): Migrate all users of this method to be agnostic of PjRt.
   xla::PjRtClient* pjrt_client() const {
     auto* pjrt_client =
-        llvm::dyn_cast_or_null<ifrt::PjRtClient>(ifrt_client_.get());
+        llvm::dyn_cast_or_null<ifrt::PjRtCompatibleClient>(ifrt_client_.get());
     if (pjrt_client == nullptr) {
       throw XlaRuntimeError(
           "This operation is implemented for a PjRt-compatible backend only.");
@@ -121,7 +121,7 @@ class PyClient : public std::enable_shared_from_this<PyClient> {
   }
   std::shared_ptr<PjRtClient> shared_ptr_pjrt_client() {
     auto* pjrt_client =
-        llvm::dyn_cast_or_null<ifrt::PjRtClient>(ifrt_client_.get());
+        llvm::dyn_cast_or_null<ifrt::PjRtCompatibleClient>(ifrt_client_.get());
     if (pjrt_client == nullptr) {
       throw XlaRuntimeError(
           "This operation is implemented for a PjRt-compatible backend only.");