-
Notifications
You must be signed in to change notification settings - Fork 97
WIP: Allow CuPy #212
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
WIP: Allow CuPy #212
Conversation
As CuPy arrays (currently) do not have a `tobytes` method, workaround this in `ensure_bytes` by coercing CuPy arrays back to NumPy arrays before calling `tobytes`. This way it can just use NumPy's `tobytes` method. Should be unneeded in a future version of CuPy.
As the decoding algorithm requires something that supports the buffer protocol and assumes it is working on host memory, this won't work on a CuPy array. So convert the CuPy array to a NumPy array first so that it can proceed through the usual decoding path.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Have written a few things below to save some thoughts here and hopefully guide the discussion should others be interested.
if isinstance(buf, np.ndarray): | ||
# already a numpy array | ||
if isinstance(buf, (np.ndarray, cp.ndarray)): | ||
# already an array | ||
arr = buf |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is easy enough to do, but it would be nice if we didn't need to learn about CuPy here.
We will probably need to find another solution for Numcodecs in the near term, but wanted to mention some related things. Namely there is some discussion about generally checking array types in issue ( numpy/numpy#14891 ).
Also mentioned by Matt in the aforementioned issue, Dask has is_arraylike
, which presents a way of solving this problem. I'm not sure it totally helps us here as we really need to check that there is a data buffer we can grab and use, but it is at least a good starting point.
Another thought I had is we could registered ndarray
types and their conversion mechanisms with Numcodecs. This could allow users to add their own array types (like CuPy), but not complicate things too much here.
Other thoughts on this would generally be welcome. 🙂
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
as we really need to check that there is a data buffer we can grab and use, but it is at least a good starting point.
Perhaps one could check for the presence of a __array__
or __cuda_array_interface__
protocol?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
That's a good point. Thanks! Will give that a look.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yep, if goal here is to ultimately enable codecs with GPU implementations then using __cuda_array_interface__
sounds good. E.g., when using numba cuda functions you get back an instance of some other class (DeviceArray or something like that IIRC) but that exports __cuda_array_interface__
too.
# because they don't have a `tobytes` method (yet) | ||
# xref: https://github.com/cupy/cupy/pull/2617 | ||
if isinstance(arr, cp.ndarray): | ||
arr = arr.get() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
As noted, this is a workaround. The real solution here is PR ( cupy/cupy#2617 ), which should land in CuPy 7.0.0. IOW this can probably be ignored for the sake of this discussion.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Ooh, +1 to cupy.ndarray..tobytes
# Force CuPy arrays to NumPy arrays | ||
# as they support the buffer protocol | ||
if isinstance(s, cp.ndarray): | ||
s = s.get() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
One way to solve this might be to call ensure_bytes
here. That would guarantee we have something we can perform text decoding on. It would also generalize for our use case here.
On the flip side that does incur a copy. Though this may be pretty inexpensive in the grand scheme of things.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think in Dask we do a check like if hasattr(s, "get"): s = s.get()
. This is a little non-specific, but it works well in practice and doesn't require the cupy import.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Just for context, this is what the ensure_bytes
change would look like ( #213 ).
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It's also possible we can get away with dropping this part entirely. This function is mainly used in conjunction with reading things (like JSON). So ensuring it can work with CuPy is probably unnecessary.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes not sure you need to worry about cupy arrays here as this is mainly for JSON reading IIRC.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Some thoughts:
- I'm shocked at how small this is
- I agree that it would be good to avoid explicitly importing cupy
- My guess is that we can get around this with mild creativiity
if isinstance(buf, np.ndarray): | ||
# already a numpy array | ||
if isinstance(buf, (np.ndarray, cp.ndarray)): | ||
# already an array | ||
arr = buf |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
as we really need to check that there is a data buffer we can grab and use, but it is at least a good starting point.
Perhaps one could check for the presence of a __array__
or __cuda_array_interface__
protocol?
# Force CuPy arrays to NumPy arrays | ||
# as they support the buffer protocol | ||
if isinstance(s, cp.ndarray): | ||
s = s.get() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think in Dask we do a check like if hasattr(s, "get"): s = s.get()
. This is a little non-specific, but it works well in practice and doesn't require the cupy import.
# because they don't have a `tobytes` method (yet) | ||
# xref: https://github.com/cupy/cupy/pull/2617 | ||
if isinstance(arr, cp.ndarray): | ||
arr = arr.get() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Ooh, +1 to cupy.ndarray..tobytes
Yeah the short nature of this change is more a consequence of work Alistair and I did last year (interestingly also in November 🙂). ( #128 ) Thanks for the tips. I'm inclined to agree some small tweaks probably make this easier to integrate. |
This adds a small change to handle CuPy
ndarrays
in thecompat
functions. Admittedly we may come up with a better way to address the underlying concern here, but I wanted to have this here for tracking. Also it would be useful to have somewhere to discuss this.TODO:
tox -e py38
passes locallytox -e py27
passes locallytox -e docs
passes locally