Add CUDA Array Interface consumer support

CHANGELOG.md

@@ -12,6 +12,11 @@ Remember to align the itemized text with the first line of an item within a list
 
 ## jax 0.4.24 (Feb 6, 2024)
 
+* New Features
+  * Added [CUDA Array
+    Interface](https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html)
+    import support.
+
 * Changes
 
   * JAX lowering to StableHLO does not depend on physical devices anymore.

jax/_src/numpy/lax_numpy.py

@@ -53,13 +53,14 @@
 from jax._src.custom_derivatives import custom_jvp
 from jax._src import dispatch
 from jax._src import dtypes
+from jax._src import xla_bridge
 from jax._src.api_util import _ensure_index_tuple
 from jax._src.array import ArrayImpl
 from jax._src.core import ShapedArray, ConcreteArray
 from jax._src.lax.lax import (_array_copy, _sort_lt_comparator,
                               _sort_le_comparator, PrecisionLike)
 from jax._src.lax import lax as lax_internal
-from jax._src.lib import xla_client as xc
+from jax._src.lib import xla_client as xc, xla_extension_version
 from jax._src.numpy import reductions
 from jax._src.numpy import ufuncs
 from jax._src.numpy import util
@@ -2121,9 +2122,25 @@ def array(object: Any, dtype: DTypeLike | None = None, copy: bool = True,
   # to be used for type inference below.
   if isinstance(object, (bool, int, float, complex)):
     _ = dtypes.coerce_to_array(object, dtype)
+  elif not isinstance(object, Array):
+    # Check if object supports any of the data exchange protocols
+    # (except dlpack, see data-apis/array-api#301). If it does,
+    # consume the object as jax array and continue (but not return) so
+    # that other array() arguments get processed against the input
+    # object.
+    #
+    # Notice that data exchange protocols define dtype in the
+    # corresponding data structures and it may not be available as
+    # object.dtype. So, we'll resolve the protocols here before
+    # evaluating object.dtype.
+    if hasattr(object, '__jax_array__'):
+      object = object.__jax_array__()
+    elif hasattr(object, '__cuda_array_interface__'):
+      if xla_extension_version >= 237:
+        cai = object.__cuda_array_interface__
+        backend = xla_bridge.get_backend("cuda")
+        object = xc._xla.cuda_array_interface_to_buffer(cai, backend)
 
-  if hasattr(object, '__jax_array__'):
-    object = object.__jax_array__()
   object = tree_map(lambda leaf: leaf.__jax_array__()
                     if hasattr(leaf, "__jax_array__") else leaf, object)
   leaves = tree_leaves(object, is_leaf=lambda x: x is None)

tests/array_interoperability_test.py

@@ -274,6 +274,87 @@ def testJaxToCuPy(self, shape, dtype):
                      z.__cuda_array_interface__["data"][0])
     self.assertAllClose(x, cupy.asnumpy(z))
 
+  @unittest.skipIf(xla_extension_version < 237, "Requires newer jaxlib")
+  @jtu.sample_product(
+    shape=all_shapes,
+    dtype=jtu.dtypes.supported(cuda_array_interface_dtypes),
+  )
+  @unittest.skipIf(not cupy, "Test requires CuPy")
+  @jtu.run_on_devices("cuda")
+  def testCuPyToJax(self, shape, dtype):
+    rng = jtu.rand_default(self.rng())
+    x = rng(shape, dtype)
+    y = cupy.asarray(x)
+    z = jnp.array(y, copy=False)  # this conversion uses dlpack protocol
+    self.assertEqual(z.dtype, dtype)
+    self.assertEqual(y.__cuda_array_interface__["data"][0],
+                     z.__cuda_array_interface__["data"][0])
+    self.assertAllClose(np.asarray(z), cupy.asnumpy(y))
+
+  @unittest.skipIf(xla_extension_version < 237, "Requires newer jaxlib")
+  @jtu.sample_product(
+    shape=all_shapes,
+    dtype=jtu.dtypes.supported(cuda_array_interface_dtypes),
+  )
+  @jtu.run_on_devices("cuda")
+  def testCaiToJax(self, shape, dtype):
+    if xb.using_pjrt_c_api():
+      self.skipTest("CUDA Array Interface support is incomplete in the PJRT C API")  # TODO(jakevdp)
+    rng = jtu.rand_default(self.rng())
+    x = rng(shape, dtype)
+
+    # using device with highest device_id for testing the correctness
+    # of detecting the device id from a pointer value
+    device = jax.devices('cuda')[-1]
+    with jax.default_device(device):
+      y = jnp.array(x, dtype=dtype)
+    self.assertEqual(y.dtype, dtype)
+
+    # Using a jax array CAI provider support to construct an object
+    # that implements the CUDA Array Interface, versions 2 and 3.
+    cai = y.__cuda_array_interface__
+    stream = tuple(y.devices())[0].get_stream_for_external_ready_events()
+
+    class CAIWithoutStridesV2:
+      __cuda_array_interface__ = cai.copy()
+      __cuda_array_interface__["version"] = 2
+      # CAI version 2 may not define strides and does not define stream
+      __cuda_array_interface__.pop("strides", None)
+      __cuda_array_interface__.pop("stream", None)
+
+    class CAIWithoutStrides:
+      __cuda_array_interface__ = cai.copy()
+      __cuda_array_interface__["version"] = 3
+      __cuda_array_interface__["strides"] = None
+      __cuda_array_interface__["stream"] = None  # default stream
+
+    class CAIWithStrides:
+      __cuda_array_interface__ = cai.copy()
+      __cuda_array_interface__["version"] = 3
+      strides = (dtype.dtype.itemsize,) if shape else ()
+      for s in reversed(shape[1:]):
+        strides = (strides[0] * s, *strides)
+      __cuda_array_interface__['strides'] = strides
+      __cuda_array_interface__["stream"] = stream
+
+    for CAIObject in [CAIWithoutStridesV2, CAIWithoutStrides,
+                      CAIWithStrides]:
+      z = jnp.array(CAIObject(), copy=False)
+      self.assertEqual(y.__cuda_array_interface__["data"][0],
+                       z.__cuda_array_interface__["data"][0])
+      self.assertAllClose(x, z)
+      if 0 in shape:
+        # the device id detection from a zero pointer value is not
+        # possible
+        pass
+      else:
+        self.assertEqual(y.devices(), z.devices())
+
+      z = jnp.array(CAIObject(), copy=True)
+      if 0 not in shape:
+        self.assertNotEqual(y.__cuda_array_interface__["data"][0],
+                            z.__cuda_array_interface__["data"][0])
+      self.assertAllClose(x, z)
 
 class Bfloat16Test(jtu.JaxTestCase):