Allow libraries that implement __array_ufunc__ to override CUDAUFuncDispatcher

Hi,

I wonder if similar approach implemented in this PR:
Allow libraries that implement array_ufunc to override DUFunc.__c… numba/numba#8995 would work for CUDAUFuncDispatcher?

The first work on awkward arrays, the latter does not. We want to be able to do the same for both:

>>> import numba as nb
>>> @nb.vectorize()
... def _square(x):
...     return x * x
... 
>>> _square        
<numba._DUFunc '_square'>

we need to add similar functionality to a CUDAUFuncDispatcher that handles ufuncs on CUDA to allow the following:

>>> @nb.vectorize(
...     target="cuda",
... )
... def _square_cuda(x):
...     return x * x
... 
>>> _square_cuda
<numba.cuda.vectorizers.CUDAUFuncDispatcher object at 0x7259f82fa790>

Please, see the issue: `"cuda"` target in numba.vectorize not working correctly? · Issue #3179 · scikit-hep/awkward · GitHub

Thanks,

Ianna

Thanks for bringing this up - I’ve created a new issue in the numba-cuda repo for further discussion / notes: [FEA] Allow libraries that implement __array_ufunc__ to override CUDAUFuncDispatcher · Issue #36 · NVIDIA/numba-cuda · GitHub

1 Like

Thanks @gmarkall ! I’ve closed mine as duplicate :slight_smile:

1 Like

Your writeup was much better than mine, so I’ve reopened yours and closed mine - many thanks for writing it up!