Move CommonPjRtRawBuffer::CopyRawHostToDeviceAndReturnEvent and

CopyRawDeviceToHostAndReturnEvent to PjRtRawBuffer and introduce PjRtRawBufferInterface
which hides the common c-api implementations.

PiperOrigin-RevId: 931432961

Author: pschuh

third_party/xla/xla/pjrt/BUILD

@@ -1536,6 +1536,7 @@ cc_library(
         "//xla:literal",
         "//xla:shape_util",
         "//xla/pjrt/c:pjrt_c_api_raw_buffer_extension_hdrs",
+        "//xla/pjrt/c:pjrt_c_api_status_utils",
         "//xla/tsl/concurrency:async_value",
         "//xla/tsl/concurrency:ref_count",
         "//xla/tsl/platform:status_macros",

third_party/xla/xla/pjrt/c/BUILD

@@ -314,7 +314,10 @@ cc_library(
     name = "pjrt_c_api_raw_buffer_extension_hdrs",
     hdrs = ["pjrt_c_api_raw_buffer_extension.h"],
     visibility = ["//visibility:public"],
-    deps = [":pjrt_c_api_hdrs"],
+    deps = [
+        ":pjrt_c_api_device_event_hdrs",
+        ":pjrt_c_api_hdrs",
+    ],
 )
 
 cc_library(

third_party/xla/xla/pjrt/c/pjrt_c_api_raw_buffer_extension.h

@@ -20,6 +20,7 @@ limitations under the License.
 #include <stdint.h>
 
 #include "xla/pjrt/c/pjrt_c_api.h"
+#include "xla/pjrt/c/pjrt_c_api_device_event.h"
 
 #ifdef __cplusplus
 extern "C" {
@@ -34,18 +35,43 @@ typedef struct PJRT_RawBuffer_FunctionTable {
   void (*inc_ref)(PJRT_RawBuffer* raw_buffer);
   void (*dec_ref)(PJRT_RawBuffer* raw_buffer);
   // Gets the number of bytes of the buffer storage on the device
-  size_t (*get_on_device_size_in_bytes)(PJRT_RawBuffer* raw_buffer);
+  size_t (*get_on_device_size_in_bytes)(const PJRT_RawBuffer* raw_buffer);
   // Gets the memory space that this buffer is attached to.
-  PJRT_Memory* (*get_memory_space)(PJRT_RawBuffer* raw_buffer);
+  PJRT_Memory* (*get_memory_space)(const PJRT_RawBuffer* raw_buffer);
   // If visible to the host, returns the base pointer for direct access.
-  void* (*get_host_pointer)(PJRT_RawBuffer* raw_buffer);
+  void* (*get_host_pointer)(const PJRT_RawBuffer* raw_buffer);
+  // Transfers the buffer to a sub-range of the on-device representation.
+  // offset+transfer_size must be less than get_on_device_size_in_bytes. The
+  // returned event transitions to ready on error, or after the transfer has
+  // completed.
+  //
+  // Note that the underlying driver may have requirements
+  // on the alignment of `src` and `offset` as well. Look at implementations of
+  // this method for specific alignment requirements.
+  PJRT_Error* (*copy_raw_host_to_device_and_return_event)(
+      PJRT_RawBuffer* raw_buffer, const void* src, int64_t offset,
+      int64_t transfer_size, PJRT_DeviceEvent* event);
+  // Transfers a sub-range of the on-device representation of the buffer.
+  // offset+transfer_size must be less than get_on_device_size_in_bytes. The
+  // returned event transitions to ready on error, or after the transfer has
+  // completed.
+  //
+  // Note that the underlying driver may have requirements
+  // on the alignment of `dst` and `offset` as well. Look at implementations of
+  // this method for specific alignment requirements.
+  PJRT_Error* (*copy_raw_device_to_host_and_return_event)(
+      PJRT_RawBuffer* raw_buffer, void* dst, int64_t offset,
+      int64_t transfer_size, PJRT_DeviceEvent* event);
+  // Return opaque device memory pointer to the underlying memory.
+  void* (*opaque_device_memory_data_pointer)(const PJRT_RawBuffer* raw_buffer);
 } PJRT_RawBuffer_FunctionTable;
 
 struct PJRT_RawBuffer {
   const PJRT_RawBuffer_FunctionTable* vtable;
 };
 
-PJRT_DEFINE_STRUCT_TRAITS(PJRT_RawBuffer_FunctionTable, get_host_pointer);
+PJRT_DEFINE_STRUCT_TRAITS(PJRT_RawBuffer_FunctionTable,
+                          opaque_device_memory_data_pointer);
 PJRT_DEFINE_STRUCT_TRAITS(PJRT_RawBuffer, vtable);
 
 struct PJRT_RawBuffer_CreateRawAliasOfBuffer_Args {

third_party/xla/xla/pjrt/c/pjrt_c_api_raw_buffer_external.cc

@@ -177,6 +177,26 @@ Future<> PjRtCApiRawBuffer::CopyRawDeviceToHost(void* dst, int64_t offset,
       c_api_, c_extension_, c_buffer_, dst, offset, transfer_size);
 }
 
+absl::StatusOr<PjRtDeviceEventRef>
+PjRtCApiRawBuffer::CopyRawHostToDeviceAndReturnEvent(const void* src,
+                                                     int64_t offset,
+                                                     int64_t transfer_size) {
+  return static_cast<PjRtRawBufferInterface*>(c_buffer_)
+      ->CopyRawHostToDeviceAndReturnEvent(src, offset, transfer_size);
+}
+
+absl::StatusOr<PjRtDeviceEventRef>
+PjRtCApiRawBuffer::CopyRawDeviceToHostAndReturnEvent(void* dst, int64_t offset,
+                                                     int64_t transfer_size) {
+  return static_cast<PjRtRawBufferInterface*>(c_buffer_)
+      ->CopyRawDeviceToHostAndReturnEvent(dst, offset, transfer_size);
+}
+
+void* PjRtCApiRawBuffer::OpaqueDeviceMemoryDataPointer() const {
+  return static_cast<PjRtRawBufferInterface*>(c_buffer_)
+      ->OpaqueDeviceMemoryDataPointer();
+}
+
 static std::optional<absl::StatusOr<tsl::RCReference<PjRtRawBuffer>>>
 PjRtCApiBuffer_CreateRawAliasOfBuffer_Factory(PjRtBuffer* buffer) {
   if (auto* c_api_buffer = dynamic_cast<xla::PjRtCApiBuffer*>(buffer)) {

third_party/xla/xla/pjrt/c/pjrt_c_api_raw_buffer_external.h

@@ -75,6 +75,11 @@ class PjRtCApiRawBuffer : public PjRtRawBuffer {
                                int64_t transfer_size) override;
   Future<> CopyRawDeviceToHost(void* dst, int64_t offset,
                                int64_t transfer_size) override;
+  absl::StatusOr<PjRtDeviceEventRef> CopyRawHostToDeviceAndReturnEvent(
+      const void* src, int64_t offset, int64_t transfer_size) override;
+  absl::StatusOr<PjRtDeviceEventRef> CopyRawDeviceToHostAndReturnEvent(
+      void* dst, int64_t offset, int64_t transfer_size) override;
+  void* OpaqueDeviceMemoryDataPointer() const override;
 
  private:
   PJRT_RawBuffer* c_buffer_;

third_party/xla/xla/pjrt/c/pjrt_c_api_raw_buffer_internal.cc

@@ -86,10 +86,10 @@ PJRT_Error* PJRT_RawBuffer_CopyRawHostToDevice(
   PJRT_RETURN_IF_ERROR(ActualStructSizeIsGreaterOrEqual(
       "PJRT_RawBuffer_CopyRawHostToDevice_Args",
       PJRT_RawBuffer_CopyRawHostToDevice_Args_STRUCT_SIZE, args->struct_size));
-  auto* cpp_buffer = static_cast<xla::PjRtRawBuffer*>(args->buffer);
-  auto result = cpp_buffer->CopyRawHostToDevice(args->src, args->offset,
-                                                args->transfer_size);
-  args->event = new PJRT_Event{std::move(result)};
+  auto* interface = static_cast<xla::PjRtRawBufferInterface*>(args->buffer);
+  auto future = interface->CopyRawHostToDevice(args->src, args->offset,
+                                               args->transfer_size);
+  args->event = new PJRT_Event{std::move(future)};
   return nullptr;
 }
 
@@ -98,10 +98,10 @@ PJRT_Error* PJRT_RawBuffer_CopyRawDeviceToHost(
   PJRT_RETURN_IF_ERROR(ActualStructSizeIsGreaterOrEqual(
       "PJRT_RawBuffer_CopyRawDeviceToHost_Args",
       PJRT_RawBuffer_CopyRawDeviceToHost_Args_STRUCT_SIZE, args->struct_size));
-  auto* cpp_buffer = static_cast<xla::PjRtRawBuffer*>(args->buffer);
-  auto result = cpp_buffer->CopyRawDeviceToHost(args->dst, args->offset,
-                                                args->transfer_size);
-  args->event = new PJRT_Event{std::move(result)};
+  auto* interface = static_cast<xla::PjRtRawBufferInterface*>(args->buffer);
+  auto future = interface->CopyRawDeviceToHost(args->dst, args->offset,
+                                               args->transfer_size);
+  args->event = new PJRT_Event{std::move(future)};
   return nullptr;
 }
 

third_party/xla/xla/pjrt/raw_buffer.cc

@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "xla/pjrt/raw_buffer.h"
 
+#include <cstddef>
 #include <utility>
 #include <vector>
 
@@ -27,6 +28,7 @@ limitations under the License.
 #include "xla/future.h"
 #include "xla/pjrt/async_work_runner.h"
 #include "xla/pjrt/c/pjrt_c_api_raw_buffer_extension.h"
+#include "xla/pjrt/c/pjrt_c_api_status_utils.h"
 #include "xla/pjrt/device_event.h"
 #include "xla/pjrt/device_event_utils.h"
 #include "xla/pjrt/pjrt_client.h"
@@ -73,6 +75,19 @@ class WrappedPjRtStagingBuffer : public PjRtStagingBuffer {
  private:
   tsl::AsyncValueRef<PjRtStagingBuffer> real_buffer_;
 };
+
+Future<> ConvertEventToFuture(PjRtDeviceEventRef event) {
+  auto [promise, future] = tsl::MakePromise<void>();
+  event.AndThen([promise = std::move(promise), event]() mutable {
+    auto error = event.GetErrorIfPresent();
+    if (error.has_value() && !error->ok()) {
+      promise.Set(*error);
+    } else {
+      promise.Set();
+    }
+  });
+  return std::move(future);
+}
 }  // namespace
 
 std::vector<RegisterRawBufferFactory::FactoryFuncT>& GetFactoryFuncs() {
@@ -224,21 +239,110 @@ const PJRT_RawBuffer_FunctionTable PjRtRawBuffer::kRawBufferVtable = {
       static_cast<PjRtRawBuffer*>(raw_buffer)->DropRef();
     },
     /*get_on_device_size_in_bytes=*/
-    +[](PJRT_RawBuffer* raw_buffer) -> size_t {
-      return static_cast<PjRtRawBuffer*>(raw_buffer)->GetOnDeviceSizeInBytes();
+    +[](const PJRT_RawBuffer* raw_buffer) -> size_t {
+      return static_cast<const PjRtRawBuffer*>(raw_buffer)
+          ->GetOnDeviceSizeInBytes();
     },
     /*get_memory_space=*/
-    +[](PJRT_RawBuffer* raw_buffer) -> PJRT_Memory* {
-      return static_cast<PjRtRawBuffer*>(raw_buffer)
+    +[](const PJRT_RawBuffer* raw_buffer) -> PJRT_Memory* {
+      return static_cast<const PjRtRawBuffer*>(raw_buffer)
           ->memory_space()
           ->ToCApiPtr();
     },
     /*get_host_pointer=*/
-    +[](PJRT_RawBuffer* raw_buffer) -> void* {
-      return static_cast<PjRtRawBuffer*>(raw_buffer)->GetHostPointer();
+    +[](const PJRT_RawBuffer* raw_buffer) -> void* {
+      return static_cast<const PjRtRawBuffer*>(raw_buffer)->GetHostPointer();
+    },
+    /*copy_raw_host_to_device_and_return_event=*/
+    +[](PJRT_RawBuffer* raw_buffer, const void* src, int64_t offset,
+        int64_t transfer_size, PJRT_DeviceEvent* event) -> PJRT_Error* {
+      auto result =
+          static_cast<PjRtRawBuffer*>(raw_buffer)
+              ->CopyRawHostToDeviceAndReturnEvent(src, offset, transfer_size);
+      if (!result.ok()) {
+        return pjrt::StatusToPjRtError(result.status());
+      }
+      *event = std::move(*result).release().ToC();
+      return nullptr;
+    },
+    /*copy_raw_device_to_host_and_return_event=*/
+    +[](PJRT_RawBuffer* raw_buffer, void* dst, int64_t offset,
+        int64_t transfer_size, PJRT_DeviceEvent* event) -> PJRT_Error* {
+      auto result =
+          static_cast<PjRtRawBuffer*>(raw_buffer)
+              ->CopyRawDeviceToHostAndReturnEvent(dst, offset, transfer_size);
+      if (!result.ok()) {
+        return pjrt::StatusToPjRtError(result.status());
+      }
+      *event = std::move(*result).release().ToC();
+      return nullptr;
+    },
+    /*opaque_device_memory_data_pointer=*/
+    +[](const PJRT_RawBuffer* raw_buffer) -> void* {
+      return static_cast<const PjRtRawBuffer*>(raw_buffer)
+          ->OpaqueDeviceMemoryDataPointer();
     },
 };
 
 PjRtRawBuffer::PjRtRawBuffer() { vtable = &kRawBufferVtable; }
 
+absl::StatusOr<PjRtDeviceEventRef>
+PjRtRawBufferInterface::CopyRawHostToDeviceAndReturnEvent(
+    const void* src, int64_t offset, int64_t transfer_size) {
+  PJRT_DeviceEvent device_event;
+  PJRT_Error* error = vtable->copy_raw_host_to_device_and_return_event(
+      this, src, offset, transfer_size, &device_event);
+  if (error != nullptr) {
+    return pjrt::PjrtErrorToStatus(error);
+  }
+  return PjRtDeviceEventRef::TakeRef(PjRtDeviceEventPtr(device_event));
+}
+
+absl::StatusOr<PjRtDeviceEventRef>
+PjRtRawBufferInterface::CopyRawDeviceToHostAndReturnEvent(
+    void* dst, int64_t offset, int64_t transfer_size) {
+  PJRT_DeviceEvent device_event;
+  PJRT_Error* error = vtable->copy_raw_device_to_host_and_return_event(
+      this, dst, offset, transfer_size, &device_event);
+  if (error != nullptr) {
+    return pjrt::PjrtErrorToStatus(error);
+  }
+  return PjRtDeviceEventRef::TakeRef(PjRtDeviceEventPtr(device_event));
+}
+
+void* PjRtRawBufferInterface::OpaqueDeviceMemoryDataPointer() const {
+  return vtable->opaque_device_memory_data_pointer(this);
+}
+
+void PjRtRawBufferInterface::AddRef() { vtable->inc_ref(this); }
+
+void PjRtRawBufferInterface::DropRef() { vtable->dec_ref(this); }
+
+PjRtMemorySpace* PjRtRawBufferInterface::memory_space() const {
+  return PjRtMemorySpace::FromC(vtable->get_memory_space(this));
+}
+
+void* PjRtRawBufferInterface::GetHostPointer() const {
+  return vtable->get_host_pointer(this);
+}
+
+size_t PjRtRawBufferInterface::GetOnDeviceSizeInBytes() const {
+  return vtable->get_on_device_size_in_bytes(this);
+}
+
+Future<> PjRtRawBufferInterface::CopyRawHostToDevice(const void* src,
+                                                     int64_t offset,
+                                                     int64_t transfer_size) {
+  ASSIGN_OR_RETURN(auto event, CopyRawHostToDeviceAndReturnEvent(
+                                   src, offset, transfer_size));
+  return ConvertEventToFuture(std::move(event));
+}
+
+Future<> PjRtRawBufferInterface::CopyRawDeviceToHost(void* dst, int64_t offset,
+                                                     int64_t transfer_size) {
+  ASSIGN_OR_RETURN(auto event, CopyRawDeviceToHostAndReturnEvent(
+                                   dst, offset, transfer_size));
+  return ConvertEventToFuture(std::move(event));
+}
+
 }  // namespace xla

third_party/xla/xla/pjrt/raw_buffer.h

@@ -44,6 +44,30 @@ namespace xla {
 class PjRtMemorySpace;
 class PjRtBuffer;
 
+class PjRtRawBufferInterface : public PJRT_RawBuffer {
+ public:
+  using RemoteSendCallback =
+      std::function<void(absl::Status status, bool sends_were_enqueued)>;
+  void AddRef();
+  void DropRef();
+
+  PjRtMemorySpace* memory_space() const;
+  void* GetHostPointer() const;
+  size_t GetOnDeviceSizeInBytes() const;
+
+  Future<> CopyRawHostToDevice(const void* src, int64_t offset,
+                               int64_t transfer_size);
+  Future<> CopyRawDeviceToHost(void* dst, int64_t offset,
+                               int64_t transfer_size);
+
+  absl::StatusOr<PjRtDeviceEventRef> CopyRawHostToDeviceAndReturnEvent(
+      const void* src, int64_t offset, int64_t transfer_size);
+  absl::StatusOr<PjRtDeviceEventRef> CopyRawDeviceToHostAndReturnEvent(
+      void* dst, int64_t offset, int64_t transfer_size);
+
+  void* OpaqueDeviceMemoryDataPointer() const;
+};
+
 // Experimental. Don't use unless you know what you're doing.
 // A raw buffer is an unsafe API for directly transferring into device
 // memory while existing processes are consuming or mutating the same buffer.
@@ -87,23 +111,12 @@ class PjRtRawBuffer : public PJRT_RawBuffer,
   virtual Future<> CopyRawDeviceToHost(void* dst, int64_t offset,
                                        int64_t transfer_size) = 0;
 
- private:
-  static const PJRT_RawBuffer_FunctionTable kRawBufferVtable;
-};
-
-class CommonPjRtRawBuffer;
-using PjRtRawBufferRef = tsl::RCReference<CommonPjRtRawBuffer>;
-
-// Adds methods common to all implementations of PjRtRawBuffer based on device
-// events.
-class CommonPjRtRawBuffer : public PjRtRawBuffer {
- public:
   // Return opaque device memory pointer to the underlying memory.
   virtual void* OpaqueDeviceMemoryDataPointer() const = 0;
 
   // Transfers the buffer to a sub-range of the on-device representation.
   // offset+transfer_size must be less than GetOnDeviceSizeInBytes. The
-  // returned future transitions to ready on error, or after the transfer has
+  // returned event transitions to ready on error, or after the transfer has
   // completed.
   //
   // Note that the underlying driver may have requirements
@@ -114,7 +127,7 @@ class CommonPjRtRawBuffer : public PjRtRawBuffer {
 
   // Transfers a sub-range of the on-device representation of the buffer.
   // offset+transfer_size must be less than GetOnDeviceSizeInBytes. The
-  // returned future transitions to ready on error, or after the transfer has
+  // returned event transitions to ready on error, or after the transfer has
   // completed.
   //
   // Note that the underlying driver may have requirements
@@ -123,6 +136,15 @@ class CommonPjRtRawBuffer : public PjRtRawBuffer {
   virtual absl::StatusOr<PjRtDeviceEventRef> CopyRawDeviceToHostAndReturnEvent(
       void* dst, int64_t offset, int64_t transfer_size) = 0;
 
+ private:
+  static const PJRT_RawBuffer_FunctionTable kRawBufferVtable;
+};
+
+class CommonPjRtRawBuffer;
+using PjRtRawBufferRef = tsl::RCReference<CommonPjRtRawBuffer>;
+
+class CommonPjRtRawBuffer : public PjRtRawBuffer {
+ public:
   // Copies the buffer to a remote device.
   // The serialized_descriptor contains metadata about the buffer on the remote
   // device. The on_done callback is called when the transfer is complete or