Skip to content

Commit

Permalink
Just use memoryview for NumPy as well
Browse files Browse the repository at this point in the history
NumPy is the main user of `__array_interface__`. However NumPy also
supports the Python Buffer Protocol, which is what is typically used for
this purpose. This can easily be used with `memoryview`s as we have done
here. So simply drop the `__array_interface__` check and use
`memoryview` for NumPy as well.
  • Loading branch information
jakirkham committed Aug 11, 2020
1 parent ecf870b commit 84eedca
Showing 1 changed file with 11 additions and 15 deletions.
26 changes: 11 additions & 15 deletions ucp/_libs/utils.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -15,9 +15,7 @@ cpdef uintptr_t get_buffer_data(buffer, bint check_writable=False) except *:
is read only and check_writable=True is set.
"""

cdef dict iface = getattr(buffer, "__array_interface__", None)
if iface is None:
iface = getattr(buffer, "__cuda_array_interface__", None)
cdef dict iface = getattr(buffer, "__cuda_array_interface__", None)

cdef uintptr_t data_ptr
cdef bint data_readonly
Expand Down Expand Up @@ -45,18 +43,16 @@ cpdef Py_ssize_t get_buffer_nbytes(buffer, check_min_size, bint cuda_support) ex
if `check_min_size` is greater than the size of the buffer
"""

cdef dict iface = getattr(buffer, "__array_interface__", None)
if iface is None:
iface = getattr(buffer, "__cuda_array_interface__", None)
if not cuda_support and iface is not None:
raise ValueError(
"UCX is not configured with CUDA support, please add "
"`cuda_copy` and/or `cuda_ipc` to the UCX_TLS environment"
"variable and that the ucx-proc=*=gpu package is "
"installed. See "
"https://ucx-py.readthedocs.io/en/latest/install.html for "
"more information."
)
cdef dict iface = getattr(buffer, "__cuda_array_interface__", None)
if not cuda_support and iface is not None:
raise ValueError(
"UCX is not configured with CUDA support, please add "
"`cuda_copy` and/or `cuda_ipc` to the UCX_TLS environment"
"variable and that the ucx-proc=*=gpu package is "
"installed. See "
"https://ucx-py.readthedocs.io/en/latest/install.html for "
"more information."
)

cdef tuple shape, strides
cdef Py_ssize_t i, s, itemsize, ndim, nbytes, min_size
Expand Down

0 comments on commit 84eedca

Please sign in to comment.