From 82b2ae211c1c160b1518c99caa00ab76bbf090fb Mon Sep 17 00:00:00 2001 From: Pearu Peterson Date: Sun, 7 Jan 2024 19:15:32 +0200 Subject: [PATCH] Add CUDA Array Interface consumer support --- CHANGELOG.md | 5 ++ jax/_src/numpy/lax_numpy.py | 23 ++++++-- tests/array_interoperability_test.py | 81 ++++++++++++++++++++++++++++ 3 files changed, 106 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 224e7a67cec1..83be5f478ec9 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py index be9b77034016..e29b31c46b9b 100644 --- a/jax/_src/numpy/lax_numpy.py +++ b/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) diff --git a/tests/array_interoperability_test.py b/tests/array_interoperability_test.py index 6a9ace853b26..865fca228daa 100644 --- a/tests/array_interoperability_test.py +++ b/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):