Add array-api copy semantics to dlpack MakePjrtBuffer

Author: pearu

jax/_src/dlpack.py

@@ -22,6 +22,7 @@
 from jax._src import xla_bridge
 from jax._src.api import device_put
 from jax._src.lax.lax import _array_copy
+from jax._src.lib import version as jaxlib_version
 from jax._src.lib import xla_client
 from jax._src.sharding import Sharding
 from jax._src.typing import Array, DLDeviceType, DTypeLike
@@ -168,16 +169,18 @@ def to_dlpack(x: Array, stream: int | Any | None = None,
       f"version ({max_version}) was requested."
     )
 
-def _place_array(_arr, device, dlpack_device, copy):
+def _check_device(device, dlpack_device, copy):
   if device and dlpack_device != device:
     if copy is not None and not copy:
       raise ValueError(
         f"Specified {device=} which requires a copy since the source device "
         f"is {repr(dlpack_device)}, however copy=False. Set copy=True or "
         "copy=None to perform the requested operation."
       )
-    else:
-      return device_put(_arr, device)
+
+def _place_array(_arr, device, dlpack_device, copy):
+  if device and dlpack_device != device:
+    return device_put(_arr, device)
   if copy:
     return jnp.array(_arr, copy=True)
   return _arr
@@ -215,6 +218,7 @@ def _legacy_from_dlpack(dlpack, device: xla_client.Device | None = None,
   _arr = jnp.asarray(xla_client._xla.dlpack_managed_tensor_to_buffer(
       dlpack, cpu_backend, gpu_backend))
   dlpack_device, = _arr.devices()
+  _check_device(device, dlpack_device, copy)
   return _place_array(_arr, device, dlpack_device, copy)
 
 def _from_dlpack(external_array, device: xla_client.Device | None = None,
@@ -230,6 +234,8 @@ def _from_dlpack(external_array, device: xla_client.Device | None = None,
 
   backend = xla_bridge.get_backend(dl_device_platform)
   dlpack_device = backend.device_from_local_hardware_id(device_id)
+  _check_device(device, dlpack_device, copy)
+
   try:
     stream = dlpack_device.get_stream_for_external_ready_events()
   except xla_client.XlaRuntimeError as err:
@@ -239,8 +245,29 @@ def _from_dlpack(external_array, device: xla_client.Device | None = None,
       raise
   dlpack = external_array.__dlpack__(stream=stream)
 
-  _arr = jnp.asarray(xla_client._xla.dlpack_managed_tensor_to_buffer(
-      dlpack, dlpack_device, stream))
+  try:
+    if jaxlib_version < (0, 6, 3):
+      _buf = xla_client._xla.dlpack_managed_tensor_to_buffer(
+        dlpack, dlpack_device, stream)
+    else:
+      _buf = xla_client._xla.dlpack_managed_tensor_to_buffer(
+        dlpack, dlpack_device, stream, copy)
+  except xla_client.XlaRuntimeError as e:
+    se = str(e)
+    i = se.index("is not aligned to")
+    if i > 0:
+      raise ValueError(
+        "Specified input which requires a copy since the source data "
+        f"buffer {se[i:]} However copy=False. Set copy=True or "
+        "copy=None to perform the requested operation."
+      )
+    else:
+      raise
+  _arr = jnp.asarray(_buf)
+  if copy and jaxlib_version >= (0, 6, 3):
+    # dlpack_managed_tensor_to_buffer implements array-api copy
+    # semantics, so resetting copy to avoid an unnecessary copy
+    copy = None
   return _place_array(_arr, device, dlpack_device, copy)
 
 def from_dlpack(external_array,

jaxlib/_jax/__init__.pyi

@@ -738,7 +738,7 @@ def buffer_to_dlpack_managed_tensor(
 ) -> Any: ...
 @overload
 def dlpack_managed_tensor_to_buffer(
-    tensor: Any, device: Device, stream: int | None
+    tensor: Any, device: Device, stream: int | None, copy: bool | None
 ) -> ArrayImpl: ...
 @overload
 def dlpack_managed_tensor_to_buffer(  # Legacy overload

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",

tests/array_interoperability_test.py

@@ -205,14 +205,21 @@ def testTensorFlowToJaxInt64(self):
   @jtu.sample_product(
     shape=all_shapes,
     dtype=numpy_dtypes,
-    copy=[False, True],
+    copy=[False, True, None] if jaxlib_version >= (0, 6, 3) else [False, True],
   )
   def testNumpyToJax(self, shape, dtype, copy):
     rng = jtu.rand_default(self.rng())
     x_np = rng(shape, dtype)
     device = jax.devices()[0]
+    if jaxlib_version < (0, 6, 3):
+      # dlpack_managed_tensor_to_buffer of older jaxlib version always
+      # copies input buffers that are unaligned
+      buf_has_64bit_alignment = True
+    else:
+      buf_has_64bit_alignment = (x_np.__array_interface__["data"][0] & 63 == 0)
     _from_dlpack = lambda: jnp.from_dlpack(x_np, device=device, copy=copy)
-    if jax.default_backend() == 'gpu' and not copy:
+    if copy is not None and not copy and (jax.default_backend() == 'gpu'
+                                          or not buf_has_64bit_alignment):
       self.assertRaisesRegex(
         ValueError,
         r"Specified .* which requires a copy",