Open lgray opened 1 month ago
ah - I was missing a awkward.numba.register_and_check()
call! that was it. Seems to work fine now.
Ah - no that was for a flattened array! I still get the error above even after doing the register_and_check()
.
So if I wrap in a flatten/unflatten I'm able to get this to work, which is a bit clunky. It seems like some aspect of awkward array are lost when dealing with the "cuda"
target of numba.vectorize
?
Here's the code that functions:
import awkward as ak
import cupy as cp
import numba as nb
ak.numba.register_and_check()
@nb.vectorize(
[
nb.float32(nb.float32),
nb.float64(nb.float64),
]
)
def _square(x):
return x * x
@nb.vectorize(
[
nb.float32(nb.float32),
nb.float64(nb.float64),
],
target="cuda",
)
def _square_cuda(x):
return x * x
def square_cuda_wrapped(x):
counts = x.layout.offsets.data[1:] - x.layout.offsets.data[:-1]
return ak.unflatten(cp.array(_square_cuda(ak.flatten(x))), counts)
counts = cp.random.poisson(lam=3, size=5000000)
flat_values = cp.random.normal(size=int(counts.sum()))
values = ak.unflatten(flat_values, counts)
values2_cpu = _square(ak.to_backend(values, "cpu"))
print(values2_cpu)
values2 = square_cuda_wrapped(values)
print(values2)
@ianna Have you been able to gain any understanding as to what is going on here? Mostly just curious, really. It's holding up progress with getting coffea going on GPUs.
@lgray and @jpivarski - it looks like there could be a simple solution to this issue. However, it needs to be implemented in Numba. I'm checking with the developers and keep you posted.
(I'm not too surprised that the fix needs to go into Numba itself, since Awkward doesn't get much control over how functions that call an Awkward Array get dispatched. If the calling function doesn't call __array_ufunc__
or some equivalent, we have no entry into the code that eventually fails.)
Version of Awkward Array
2.6.6
Description and code to reproduce
The following code fails:
resulting in the output: