How to use Awkward Arrays on GPUs

[soleti]

First of all, this is an amazing project! I am trying to use awkward arrays with CUDA kernels implemented with Numba. It would be good to have more information in general on how to use awkward arrays on GPU (on this page https://awkward-array.org/how-to-math-gpu.html) and, if possible, an example with Numba cuda-jitted function.

Thanks!

[jpivarski]

Just to let you know, GPU support won't be ready for some months, and then only on Linux for Nvidia GPUs. You might be able to

pip install awkward[cuda]

and do some simple things like

>>> one = ak.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]], kernels="cuda")
>>> two = ak.Array([100, 200, 300], kernels="cuda")
>>> one, two
(<Array:cuda [[1.1, 2.2, 3.3], ... [4.4, 5.5]] type='3 * var * float64'>,
 <Array:cuda [100, 200, 300] type='3 * int64'>)
>>> three = one + two
>>> three
<Array:cuda [[101, 102, 103], ... [304, 306]] type='3 * var * float64'>
>>> ak.flatten(three)
<Array:cuda [101, 102, 103, 304, 306] type='5 * float64'>
>>> ak.to_cupy(ak.flatten(three))
array([101.1, 102.2, 103.3, 304.4, 305.5])
>>> type(ak.to_cupy(ak.flatten(three)))
<class 'cupy.core.core.ndarray'>

but a lot hasn't been implemented and it's a minefield for segmentation faults. (Throughout the codebase, there's a // DERIVE comment indicating all the places we need to determine whether an array is in main memory or on the GPU, instead of just assuming it's in main memory. At least we'll be able to do that systematically, but it hasn't been done yet.)

I just tried ak.unflatten and put in a minor tweak to make it turn CuPy arrays into jagged arrays (#754), but quickly ran into some segmentation faults when trying to do some math on it. So, just to let you know: it's not ready yet.

[aghoshpub]

It's probably a lot of work but indeed would be a pretty useful to be able to accelerate awkward array computations on GPUs !