Should cupy ufuncs work on awkward arrays that are correctly on-device?

[lgray]

Version of Awkward Array

2.6.7

Description and code to reproduce

Awkward knows how to dispatch through numpy ufuncs, and will dispatch to cupy ufuncs when encountering the "cuda" backend.

However, when using cupy ufuncs on awkward arrays with cuda backing it doesn't work out of the box!

Here is a reproducer:

import awkward as ak
import cupy as cp
import numpy as np

counts = cp.random.poisson(lam=3, size=5000000)
flat_values = cp.random.normal(size=int(counts.sum()))

values = ak.unflatten(flat_values, counts)

np_vals = np.abs(values)
print(np_vals, ak.backend(np_vals))

cp_vals = cp.abs(values)
print(cp_vals, ak.backend(cp_vals))

it outputs:

[[0.045683051144471104, 0.27795220821182787, ..., 0.8520697217330707], ..., []] cuda

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[16], line 13
     10 np_vals = np.abs(values)
     11 print(np_vals, ak.backend(np_vals))
---> 13 cp_vals = cp.abs(values)
     14 print(cp_vals, ak.backend(cp_vals))

File cupy[/_core/_kernel.pyx:1283](https://analytics-hub.fnal.gov/_core/_kernel.pyx#line=1282), in cupy._core._kernel.ufunc.__call__()

File cupy[/_core/_kernel.pyx:159](https://analytics-hub.fnal.gov/_core/_kernel.pyx#line=158), in cupy._core._kernel._preprocess_args()

File cupy[/_core/_kernel.pyx:145](https://analytics-hub.fnal.gov/_core/_kernel.pyx#line=144), in cupy._core._kernel._preprocess_arg()

TypeError: Unsupported type <class 'awkward.highlevel.Array'>

It would be really nice if I could use cupy ufuncs on awkward arrays just as easily I can use numpy ufuncs.
I am not sure if this is a bug or a feature request.

[jpivarski]

I don't know if CuPy has an equivalent of NEP 13 (or NEP 18) to dispatch from cp.* functions to whatever. CuPy uses NumPy's dispatch from np.* to cp.*, but you want cp.* to something that we write.

In a web-search, I didn't find anything.

[lgray]

The reason I am asking about this is because it's fairly straightforward to define custom ufunc kernels in cupy to accelerate things like invariant mass and deltaphi (fusing loops, etc). It would then be nice to have the ability to apply custom cupy functions directly to awkward arrays.

[jpivarski]

It is the case that if you call np.whatever on a CUDA-backed Awkward Array, it will run cp.whatever.

[lgray]

Yes, I am aware. But that doesn't work with custom ufuncs as far as I know?

[jpivarski]

Custom ufuncs from @numba.vectorize work (I had to get a PR into Numba to make that happen). As I understood from our conversation (you, me, and @ianna), something needs to be done to make that work for the CUDA target.

[lgray]

Correct, I was trying to find a work-around for that and ran into this problem instead. :-)

[kmaehashi]

Hi, CuPy maintainer here! I'd like to understand what exactly is the want here. Specifically, isn't it possible to give CuPy ndarrays to NumPy to enjoy NEP 13? (e.g. np.abs(cupy_ndarray) dispatches to CuPy).

[lgray]

@kmaehashi Thanks for your interest! I would like to use cupy ufuncs with awkward arrays, which are a generalization of numpy and it supports cpu and gpu backends.

Right now it can use NEP-13 to dispatch from numpy to cupy when the array is on the GPU, however you cannot use cp.<some ufunc>(an_awkward_array) because the cupy ufunc itself does not have a dispatch mechanism!

The example that caused this to crop up actually arose when trying to use numba-cuda, and you can see a bit more of what's going on in this PR: NVIDIA/numba-cuda#37.

To put it succinctly: I would like to be able to dispatch cupy ufuncs to awkward arrays, which right now only works if there is a numpy equivalent of the ufunc. However, that is not generally the case and quickly causes pain!