Add CUDA Array Interface consumer support

[pearu]

This PR adds CUDA Array Interface (versions 2 and 3) consumer support to JAX.

In addition, the PR enables constructing JAX arrays from objects that implement dlpack provider support.

Fixes #1100

Requires openxla/xla#8237

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[jakevdp]

Thanks for the contribution – it's not clear to me that asarray and array should transparently handle dlpack and other data interchange formats. For example, the Python Array API standard explicitly decided not to do this (see data-apis/array-api#301 and linked issues). This suggests that we should stick with explicit data interchange functions like jnp.from_dlpack. What do you think?

[pearu]

This suggests that we should stick with explicit data interchange functions like jnp.from_dlpack. What do you think?

@jakevdp Yes, I agree. I'll update the PR accordingly.

[jakevdp]

Similar to from_dlpack, I wonder if an explicit from_cuda_interface function would better achieve the goal here?

[pearu]

Similar to from_dlpack, I wonder if an explicit from_cuda_interface function would better achieve the goal here?

While array-api does not mention CUDA Array Interface nor NumPy Array Interface, these are kind of legacy buffer protocols and objects implementing the Buffer Protocol are allowed inputs to asarray according to the array-api.

Here's a summary how the objects implementing CUDA Array Interface protocol are resolved elsewhere:

• cupy resolves __cuda_array_interface__ in array (asarray typically calls array). Btw, it has internal method _array_from_cuda_array_interface for that.
• PyTorch resolves __cuda_array_interface__ in tensor (analogue of array) and in as_tensor (analogue of asarray) and in asarray.
• mpi4py resolves __cuda_array_interface__ in frombuffer and getbuffer.

Considering the above, my first choice would be to resolve CAI objects in the array function (as in this PR). However, the explicit from_cai or similar also makes sense as this is about importing a view of an existing buffer similar to from_dlpack. What do you think?

[jakevdp]

That sounds good, we can do it in asarray. Though it introduces some complexities, e.g. should lists of cuda-compatible objects be treated as lists of arrays? That's possible currently with our handling of buffer protocol objects.

[pearu]

Though it introduces some complexities, e.g. should lists of cuda-compatible objects be treated as lists of arrays? That's possible currently with our handling of buffer protocol objects.

Currently, jax, numpy, and cupy implement this using the following principle: if input to asarray is a native array object then asarray([input]) is an array with ndim larger than input.ndim. (PyTorch does not support list of tensors as an input to asarray).

Your concern corresponds to the case where input is not a native array object, say, it is anything that implements one of the data exchange protocols. Currently, all above mentioned libraries will raise an exception on asarray([input]):

• numpy raises TypeError with a FutureWarning:

  The input object of type 'W' is an array-like implementing one of the corresponding protocols (`__array__`, `__array_interface__` or `__array_struct__`); but not a sequence (or 0-D).\
  In the future, this object will be coerced as if it was first converted using `np.array(obj) ...
  

• cupy raises ValueError: Unsupported dtype object
• jax raises TypeError: Cannot interpret '<__main__.W object at 0x7fc54a55b090>' as a data type

Considering the numpy FutureWarning, the answer might be affirmative, but notice that https://data-apis.org/array-api/latest/API_specification/generated/array_api.asarray.html#array_api.asarray does not support this nor does it support the current behavior where input list items are native array objects.

[jakevdp]

Thanks - before we proceed with the review, can you sign the CLA?

[rgommers]

Similar to from_dlpack, I wonder if an explicit from_cuda_interface function would better achieve the goal here?

I'm not 100% sure about this one. DLPack should provide a superset of functionality of __cuda_array_interface__ (and for CUDA specifically, the stream support in DLPack was improved based on the lessons of __cuda_array_interface__, so I'd consider the latter legacy and nudge any remaining libraries that don't have DLPack support that's on par with their __cuda_array_interface__ implementation to fix that.

There's something to be said for this either way of course, given that not all libraries today have support that's fully equivalent between the two protocols. But there isn't too much traffic on gh-1100 in 3+ years. And we have too many different protocols in Python that are overlapping - with DLPack being the most viable one long-term, as the only one that has multi-device support. __cuda_array_interface__ doesn't even seem to support ROCm, so unless I'm missing something, it may not be needed long-term? Keeping its support inside another function (or multiple) makes it easier to evolve / phase out.

[jakevdp]

Circling back here: after chatting with @hawkinsp a bit, I think this is probably the right approach. We'll have to wait on openxla/xla#8237 before we can do anything here, but once that's in let's plan to proceed with this review.

[jakevdp]

OK, looks good! Last thing: could you please squash the changes into a single commit? Thanks!

[pearu]

OK, looks good! Last thing: could you please squash the changes into a single commit? Thanks!

Done.

[pearu]

I guess before merging this PR, openxla/xla#8237 needs to land soon.

[pearu]

@jakevdp I think we can move forward with this PR as openxla/xla#8237 has landed

[jakevdp]

Looks good - we also should add a CHANGELOG entry.

[jakevdp]

Looks good!

[jakevdp]

Sorry for the delay in mergind: we're running into some test failures when testing this against PJRT runtimes; I'm trying to debug internally.

[jakevdp]

We need a skip condition here for PJRT runtimes, similar to others in this file. Something like this should work:

Suggested change

	def testCaiToJax(self, shape, dtype):
	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)

[pearu]

Applied.

IIUC, PJRT is about running code on Intel GPUs but this test is decorated with jtu.run_on_devices("cuda") that would imply that the test ought to be skipped anyway when using a non-CUDA device. Is there an explanation why this test is still executed for PJRT runtimes?

Another question is if CAI support can ever be used within PJRT because CUDA is Nvidia device-specific?