feat: add __dlpack__, from_dlpack support
#2649
Conversation
agoose77
changed the title
from_dlpack operation to backends
Codecov Report
Additional details and impacted files
|
| nplike: NumpyLike | ||
| if device_type == DLPackDevice.CPU: | ||
| nplike = Numpy.instance() | ||
| elif device_type == DLPackDevice.CUDA: | ||
| nplike = Cupy.instance() | ||
| elif device_type == DLPackDevice.CUDA_PINNED: | ||
| # TODO: this should support GPU | ||
| # nplike = (Numpy if prefer_cpu else Cupy).instance() | ||
| nplike = Numpy.instance() | ||
| elif device_type == DLPackDevice.ROCM: | ||
| nplike = Cupy.instance() | ||
| elif device_type == DLPackDevice.ROCM_PINNED: | ||
| nplike = (Numpy if prefer_cpu else Cupy).instance() | ||
| elif device_type == DLPackDevice.CUDA_MANAGED: | ||
| nplike = (Numpy if prefer_cpu else Cupy).instance() | ||
| else: | ||
| raise AssertionError |
There was a problem hiding this comment.
We could generalise this to include JAX. However, I think this is a lot simpler to reason about for now.
agoose77
changed the title
from_dlpack operation to backends__dlpack__, from_dlpack support
…ce__` And expose CUDA
agoose77
force-pushed
the
agoose77/feat-from-dlpack
branch
from
884bed5
to
cc4b00f
Compare
There was a problem hiding this comment.
This converts from DLPack arrays (ak.from_dlpack) and converts single-buffer Awkward Arrays to DLPack.
Just to be clear, that "single-buffer Awkward Arrays" is a proper subset of "Awkward Arrays that won't fail in ak.to_numpy":
>>> array = ak.Array([
... [{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}],
... [{"x": 3, "y": 3.3}, {"x": 4, "y": 4.4}],
... ])
>>> ak.to_numpy(array) # also np.asarray(array)
array([[(1, 1.1), (2, 2.2)],
[(3, 3.3), (4, 4.4)]], dtype=[('x', '<i8'), ('y', '<f8')])
>>> array = ak.Array([1.1, 2.2, None, 4.4, None, None, 7.7])
>>> ak.to_numpy(array) # NOT np.asarray(array)
masked_array(data=[1.1, 2.2, --, 4.4, --, --, 7.7],
mask=[False, False, True, False, True, True, False],
fill_value=1e+20)In both of the above, array contains multiple buffers that might be on different GPUs in the "cuda" backend, so it's good that they'll fail to be converted to DLPack.
It looks good to me; go ahead and merge it if you don't have any last-minute changes.
| if not backend.nplike.supports_structured_dtypes: | ||
| raise TypeError( | ||
| f"backend {backend.name!r} does not support structured " | ||
| f"dtypes required for converting {type(self).__name__} " | ||
| f"into a backend array" | ||
| ) |
There was a problem hiding this comment.
You noticed this: okay, so maybe I didn't need to tell you.
There was a problem hiding this comment.
This is a cleanup more than a feature; NumPy does support both DLPack and structured dtypes. However, NumPy doesn't support structured dtypes with DLPack, and will report that the conversion fails, e.g.
>>> import awkward as ak
>>> array = ak.Array([{'x': 1}]).to_numpy()
>>> ptr = array.__dlpack__()
Traceback (most recent call last):
File "...", line 5, in <module>
ptr = array.__dlpack__()
^^^^^^^^^^^^^^^^^^
BufferError: DLPack only supports signed/unsigned integers, float and complex dtypes.
[('x', '<i8')]
This PR fixes #2648.
Array.__dlpack__(stream=None)andArray.__dlpack_device__()ak.from_dlpackto_nplikeArray.__array__withArray.__array_interface__andArray.__cuda_array_interface__(as these are higher precedence, and symmetric to one-another. It's arbitrary, we could keep__array__, but I think__array_interface__should be fairly well supported).