typing ak.Array for numba.cuda.jit signature

[essoca]

Version of Awkward Array

2.6.2

Description and code to reproduce

Hey guys, I followed a hint from the discussion in #696 to type ak.Array for numba signatures. So I tried something like

import awkward as ak
import numba as nb
from numba import types

cpu_arr_type = ak.Array([[[0, 1], [2, 3]], [[4, 5]]], backend='cpu').numba_type

@nb.njit(types.void(cpu_arr_type))
def cpu_kernel(arr):
   do_something_with_arr

and this works like a charm.

However, I'm interested in the same case but with a cuda kernel. So I tried what appeared more natural to do:

gpu_arr_type = ak.Array([[[0, 1], [2, 3]], [[4, 5]]], backend='cuda').numba_type

@nb.cuda.jit(types.void(gpu_arr_type), extensions=[ak.numba.cuda])
def cuda_kernel(arr):
   do_something_with_arr

This time, I get the error:

self = <awkward._connect.numba.arrayview_cuda.ArrayViewArgHandler object at 0x784afbc13fa0>
ty = ak.ArrayView(ak.ListArrayType(array(int64, 1d, C), ak.ListArrayType(array(int64, 1d, C), ak.NumpyArrayType(array(int64, 1d, C), {}), {}), {}), None, ())
val = <Array [[[4, 1], [2, -1]], [...], [[4, 0]]] type='3 * var * var * int64'>
stream = 0, retr = []

    def prepare_args(self, ty, val, stream, retr):
        if isinstance(val, ak.Array):
            if isinstance(val.layout.backend, CupyBackend):
                # Use uint64 for pos, start, stop, the array pointers values, and the pylookup value
                tys = numba.types.UniTuple(numba.types.uint64, 5)
    
>               start = val._numbaview.start
E               AttributeError: 'NoneType' object has no attribute 'start'

.../site-packages/awkward/_connect/numba/arrayview_cuda.py:21: AttributeError

How should this latter case be correctly treated? Note that, without typing, the thing works as expected:

@nb.cuda.jit(extensions=[ak.numba.cuda])
def cuda_kernel_no_typing(arr):
   do_something_with_arr

However, I'm interested in ak.Arrays with the 3D layout of integers (as above) and would like to take advantage of numba's eager compilation. I'm passing the arr for testing as

backend = 'cpu'  # or 'cuda'
arr = ak.to_backend(
    ak.Array([
        [[4, 1], [2, -1]],
        [[0, -1], [1, 1], [3, -1]],
        [[4, 0]]
    ]),
    backend
)

Any help is appreciated!

[jpivarski]

Maybe the signature is wrong?

If I declare a function without specifying the signature,

>>> @nb.cuda.jit(extensions=[ak.numba.cuda])
... def cuda_kernel_no_typing(arr):
...     return None
... 

then run it,

>>> cuda_kernel_no_typing[1024, 1024](arr)

Now it's made a signature because it JIT-compiled before running:

>>> cuda_kernel_no_typing.signatures[0]
(ak.ArrayView(ak.ListArrayType(array(int64, 1d, C), ak.ListArrayType(array(int64, 1d, C), ak.NumpyArrayType(array(int64, 1d, C), {}), {}), {}), None, ()),)

Using exactly that signature to precompile a function:

>>> sig = cuda_kernel_no_typing.signatures[0]
>>> @nb.cuda.jit(sig, extensions=[ak.numba.cuda])
... def cuda_kernel(arr):
...     return None
... 

It seems to work (no errors):

>>> cuda_kernel[1024, 1024](arr)

[essoca]

Many thanks @jpivarski for your quick response! Your same procedure (grabbing the sig after JIT-compilation and using it to type) doesn't work for me. It throws the error above. I'm using the numba version 0.59.1

It would be great to generate the signature without having to run the 'untyped` kernel first. That's actually the idea behind typing: we know in advance what is coming into the kernel.

Now, one can see here some inconsistency:

gpu_arr_type = ak.Array([[[0, 1], [2, 3]], [[4, 5]]], backend='cuda').numba_type

>>> gpu_arr_type
ak.ArrayView(ak.ListArrayType(array(int64, 1d, C), ak.ListArrayType(array(int64, 1d, C), ak.NumpyArrayType(array(int64, 1d, C), {}), {}), {}), None, ())

which doesn't look the same as

>>> cuda_kernel_no_typing.signatures[0]
(ak.ArrayView(ak.ListArrayType(array(int64, 1d, C), ak.ListArrayType(array(int64, 1d, C), ak.NumpyArrayType(array(int64, 1d, C), {}), {}), {}), None, ()),)

The former is the unpackedversion of the latter one.

[jpivarski]

@ianna, do you see what this might be?

[ianna]

@essoca - Could you, please, try to change it as follows to pin the ArrayView:

array = ak.Array([[[0, 1], [2, 3]], [[4, 5]]], backend='cuda')
gpu_arr_type = array.numba_type

@nb.cuda.jit(types.void(gpu_arr_type), extensions=[ak.numba.cuda])
def cuda_kernel(arr):
   do_something_with_arr

It seems, it works for me:

@nb.cuda.jit(types.void(gpu_arr_type), extensions=[ak.numba.cuda])
def cuda_kernel(arr):
    return None

cuda_kernel[1024, 1024](array)

[essoca]

Thanks for trying @ianna. You're almost there. Up to that point, it also works for me. Now test it with a different array with the same nested structure:

array2 = ak.Array([[[4, 1], [2, -1]], [[0, -1], [1, 1], [3, -1]], [[4, 0]]], backend='cuda')

>>> cuda_kernel[1024, 1024](array2)

../../miniforge3/envs/qb/lib/python3.10/site-packages/numba/cuda/dispatcher.py:539: in __call__
    return self.dispatcher.call(args, self.griddim, self.blockdim,
../../miniforge3/envs/qb/lib/python3.10/site-packages/numba/cuda/dispatcher.py:683: in call
    kernel.launch(args, griddim, blockdim, stream, sharedmem)
../../miniforge3/envs/qb/lib/python3.10/site-packages/numba/cuda/dispatcher.py:317: in launch
    self._prepare_args(t, v, stream, retr, kernelargs)
../../miniforge3/envs/qb/lib/python3.10/site-packages/numba/cuda/dispatcher.py:380: in _prepare_args
    ty, val = extension.prepare_args(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <awkward._connect.numba.arrayview_cuda.ArrayViewArgHandler object at 0x7406bcc4c160>
ty = ak.ArrayView(ak.ListArrayType(array(int64, 1d, C), ak.ListArrayType(array(int64, 1d, C), ak.NumpyArrayType(array(int64, 1d, C), {}), {}), {}), None, ())
val = <Array [[[4, 1], [2, -1]], [...], [[4, 0]]] type='3 * var * var * int64'>
stream = 0, retr = []

    def prepare_args(self, ty, val, stream, retr):
        if isinstance(val, ak.Array):
            if isinstance(val.layout.backend, CupyBackend):
                # Use uint64 for pos, start, stop, the array pointers values, and the pylookup value
                tys = numba.types.UniTuple(numba.types.uint64, 5)
    
>               start = val._numbaview.start
E               AttributeError: 'NoneType' object has no attribute 'start'

../../miniforge3/envs/qb/lib/python3.10/site-packages/awkward/_connect/numba/arrayview_cuda.py:21: AttributeError

[ianna]

@essoca - thanks, I can reproduce the error. It is definitely a bug!

I think it is related to our way of a generated C++ ArrayView typing that is unique for each Python Awkward array. The view holds a pointer to the array data. The view is generated on demand. It looks like when the type is given we miss to identify that it is has not been done yet.

Please, try to force it to be generated and check that:

array2.numba_type
array2._numbaview
array2._numbaview is None

then this should work:

cuda_kernel[1024, 1024](array2)

[essoca]

Thanks @ianna. Could you be a bit more explicit about what should be done for that to work? When array is initially created, it has a numba_type, its _numbaview field is initially None, but then it gets asynchronously self-assigned. I have tried to set the _numbaview to None (in order to extract the sig and JIT-compile) but it gets self-assigned again. So the issue persists.

[ianna]

Thanks @ianna. Could you be a bit more explicit about what should be done for that to work? When array is initially created, it has a numba_type, its _numbaview field is initially None, but then it gets asynchronously self-assigned. I have tried to set the _numbaview to None (in order to extract the sig and JIT-compile) but it gets self-assigned again. So the issue persists.

Sorry, I was not clear. You should not try setting the _numbaview to None. On the contrary, I think, it might work if you try to force it to be generated (e.g. to be not None) by executing array2._numbaview before passing array2 to the kernel.

[essoca]

@ianna: it partially works. But executing array2._numbaview before passing array2 to the kernel has no effect at all.

After putting a breakpoint at line 518, the only way for execution to reach 519 successfully is by inspecting array2 (i.e. hovering over it). Actually, if you remove 517 and stop at the kernel line, you will notice at first inspection that the _numbaview field is None. The second inspection somehow triggers having a _numbaview ... and the execution reaches the print() statement successfully. So what is the callable which is indeed triggering the generation of a _numbaview?

[ianna]

@essoca - I think, the PR #3115 fixes the issue. What would be your expectations if the signature type does not correspond to the array type passed to the kernel?

[essoca]

@ianna - Thanks for the PR! Things now work as expected. If the signature type does not correspond to the array type passed to the kernel, the Numba dispatcher will raise the TypeError: No matching definition for argument type(s) .... This is a problem from the user (and not the libraries) though :)