Add array-api copy semantics to dlpack MakePjrtBuffer

Author: pearu

jax/_src/dlpack.py

@@ -240,7 +240,12 @@ def _from_dlpack(external_array, device: xla_client.Device | None = None,
   dlpack = external_array.__dlpack__(stream=stream)
 
   _arr = jnp.asarray(xla_client._xla.dlpack_managed_tensor_to_buffer(
-      dlpack, dlpack_device, stream))
+      dlpack, dlpack_device, stream, copy))
+  if copy:
+    # dlpack_managed_tensor_to_buffer already handles the array-api
+    # copy semantics, so resetting copy to False to avoid a second
+    # copy
+    copy = False
   return _place_array(_arr, device, dlpack_device, copy)
 
 def from_dlpack(external_array,

jaxlib/dlpack.cc

@@ -211,39 +211,49 @@ absl::StatusOr<std::vector<int64_t>> GetByteStrides(const DLTensor& dl_tensor) {
 absl::StatusOr<std::unique_ptr<PjRtBuffer>> MakePjrtBuffer(
     PjRtDevice& device, ::DLManagedTensor* dlmt, const Shape& shape,
     PrimitiveType element_type, absl::Span<int64_t const> dimensions,
+    std::optional<bool> copy = std::nullopt,
     std::optional<std::intptr_t> stream = std::nullopt) {
   std::function<void()> on_delete_callback;
   if (dlmt->deleter) {
     on_delete_callback = [dlmt]() { dlmt->deleter(dlmt); };
   }
 
-  // First try to create a view.
   void* data =
       static_cast<char*>(dlmt->dl_tensor.data) + dlmt->dl_tensor.byte_offset;
-  auto result = device.client()->CreateViewOfDeviceBuffer(
-      data, shape, *device.default_memory_space(), on_delete_callback, stream);
-
-  // If that fails with invalid argument, it's possibly because of the incorrect
-  // alignment. If we're on CPU, we can create a copy of buffer.
-  if (result.status().code() == absl::StatusCode::kInvalidArgument &&
-      dlmt->dl_tensor.device.device_type == kDLCPU) {
-    LOG(WARNING) << "DLPack buffer is not aligned (data at: " << data
-                 << "). Creating a copy.";
-
-    // Convert tensor strides (expressed in number of elements) to byte strides.
-    std::optional<std::vector<int64_t>> byte_strides;
-    if (dlmt->dl_tensor.strides) {
-      TF_ASSIGN_OR_RETURN(byte_strides, GetByteStrides(dlmt->dl_tensor));
-    }
 
-    TF_ASSIGN_OR_RETURN(auto* memory_space, device.default_memory_space());
+  // On CPU, creating a view may fail because of unaligned data buffer
+  // in which case we'll fallback to copy. On non-CPU, array-api copy
+  // semantics is handled in dlpack._place_array function.
+  bool fallback_to_copy = !copy.has_value() && dlmt->dl_tensor.device.device_type == kDLCPU;
+
+  if (!copy.value_or(false)) {
+    // require a view (copy == False) or try create a view (copy == None)
+    auto result = device.client()->CreateViewOfDeviceBuffer(
+        data, shape, *device.default_memory_space(), on_delete_callback, stream);
+    if ((result.status().code() != absl::StatusCode::kInvalidArgument)
+        || (!fallback_to_copy))
+    {
+      // succesful view or return error when copy == False or copy ==
+      // None with fallback_to_copy == False
+      return result;
+    }
+  }
+  // require copy (copy == True) or fallback to copy (copy == None)
 
-    // Create a copy.
-    result = device.client()->BufferFromHostBuffer(
-        data, element_type, dimensions, byte_strides,
-        PjRtClient::HostBufferSemantics::kMutableZeroCopy, on_delete_callback,
-        memory_space, /*device_layout=*/nullptr);
+  // Convert tensor strides (expressed in number of elements) to byte strides.
+  std::optional<std::vector<int64_t>> byte_strides;
+  if (dlmt->dl_tensor.strides) {
+    TF_ASSIGN_OR_RETURN(byte_strides, GetByteStrides(dlmt->dl_tensor));
   }
+
+  TF_ASSIGN_OR_RETURN(auto* memory_space, device.default_memory_space());
+
+  // Create a copy.
+  auto result = device.client()->BufferFromHostBuffer(
+      data, element_type, dimensions, byte_strides,
+      PjRtClient::HostBufferSemantics::kMutableZeroCopy, on_delete_callback,
+      memory_space, /*device_layout=*/nullptr);
+
   return result;
 }
 
@@ -424,7 +434,7 @@ absl::StatusOr<nb::object> DLPackManagedTensorToBuffer(
 
 absl::StatusOr<nb::object> DLPackManagedTensorToBuffer(
     const nb::capsule& tensor, ifrt::Device* ifrt_device,
-    nb_class_ptr<PyClient> client, std::optional<std::intptr_t> stream) {
+    nb_class_ptr<PyClient> client, std::optional<std::intptr_t> stream, std::optional<bool> copy) {
   ifrt::PjRtDevice* device =
       llvm::dyn_cast_or_null<ifrt::PjRtDevice>(ifrt_device);
   if (device == nullptr) {
@@ -469,7 +479,7 @@ absl::StatusOr<nb::object> DLPackManagedTensorToBuffer(
 
   TF_ASSIGN_OR_RETURN(auto pjrt_buffer,
                       MakePjrtBuffer(*device->pjrt_device(), dlmt, shape,
-                                     element_type, dimensions, stream));
+                                     element_type, dimensions, copy, stream));
 
   // We have taken ownership of the array inside the capsule; make sure the
   // capsule it cannot be used again.

jaxlib/dlpack.h

@@ -46,7 +46,8 @@ absl::StatusOr<nanobind::object> DLPackManagedTensorToBuffer(
 
 absl::StatusOr<nanobind::object> DLPackManagedTensorToBuffer(
     const nanobind::capsule& tensor, ifrt::Device* device,
-    nb_class_ptr<PyClient> client, std::optional<std::intptr_t> stream);
+    nb_class_ptr<PyClient> client, std::optional<std::intptr_t> stream,
+    std::optional<bool> copy);
 
 // Converts a PrimitiveType to the nanobind specific implementation of
 // DLDataType.

jaxlib/xla.cc

@@ -575,11 +575,11 @@ NB_MODULE(_jax, m) {
   m.def(
       "dlpack_managed_tensor_to_buffer",
       [](const nb::capsule& tensor, nb_class_ptr<PyDevice> device,
-         std::optional<std::intptr_t> stream) {
+         std::optional<std::intptr_t> stream, std::optional<bool> copy) {
         return xla::ValueOrThrow(DLPackManagedTensorToBuffer(
-            tensor, device->device(), device->client(), stream));
+            tensor, device->device(), device->client(), stream, copy));
       },
-      nb::arg("dlpack"), nb::arg("device"), nb::arg("stream").none());
+      nb::arg("dlpack"), nb::arg("device"), nb::arg("stream").none(), nb::arg("copy").none() = nb::none());
   // Legacy overload
   m.def(
       "dlpack_managed_tensor_to_buffer",