Repository: scikit-hep/awkward

Issue:
typing ak.Array for numba.cuda.jit signature
### Version of Awkward Array

2.6.2

### Description and code to reproduce

Hey guys, I followed a hint from the discussion in [#696](https://github.com/scikit-hep/awkward/discussions/696#discussion-2571850) to type `ak.Array` for numba signatures. So I tried something like

```python
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:
```python
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:
```python
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:
```python
@nb.cuda.jit(extensions=[ak.numba.cuda])
def cuda_kernel_no_typing(arr):
   do_something_with_arr
```
However, I'm interested in `ak.Array`s 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
```python
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!



Discussion:
Maybe the signature is wrong?

If I declare a function without specifying the signature,

```python
>>> @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:

```python
>>> 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):

```python
>>> cuda_kernel[1024, 1024](arr)
```
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:

```python
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 
```python
>>> 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 `unpacked`version of the latter one.
@ianna, do you see what this might be?
@essoca - Could you, please, try to change it as follows to pin the `ArrayView`:
```python
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:
```python
@nb.cuda.jit(types.void(gpu_arr_type), extensions=[ak.numba.cuda])
def cuda_kernel(arr):
    return None

cuda_kernel[1024, 1024](array)
```
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:
```python
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
```

@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:
```python
array2.numba_type
array2._numbaview
array2._numbaview is None
```
then this should work:
```python
cuda_kernel[1024, 1024](array2)
```
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.
> 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.
@ianna: it **partially** works. But executing `array2._numbaview` before passing `array2` to the kernel has no effect at all.

![Screenshot from 2024-05-11 10-42-51](https://github.com/scikit-hep/awkward/assets/28766807/bb0c9bde-09b1-48a0-9a5d-d548b103f14c)

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`?


Resolving patch:
diff --git a/src/awkward/_connect/numba/arrayview_cuda.py b/src/awkward/_connect/numba/arrayview_cuda.py
--- a/src/awkward/_connect/numba/arrayview_cuda.py
+++ b/src/awkward/_connect/numba/arrayview_cuda.py
@@ -15,13 +15,22 @@ class ArrayViewArgHandler:
     def prepare_args(self, ty, val, stream, retr):
         if isinstance(val, ak.Array):
             if isinstance(val.layout.backend, CupyBackend):
+                if ty is not val.numba_type:
+                    raise NumbaTypeError(
+                        f"the array type: {val.numba_type} does not match "
+                        f"the kernel signature type: {ty}"
+                    )
+
                 # 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
-                stop = val._numbaview.stop
-                pos = val._numbaview.pos
-                arrayptrs = val._numbaview.lookup.arrayptrs.data.ptr
+                view = val._numbaview
+                assert view is not None
+
+                start = view.start
+                stop = view.stop
+                pos = view.pos
+                arrayptrs = view.lookup.arrayptrs.data.ptr
                 pylookup = 0
 
                 return tys, (pos, start, stop, arrayptrs, pylookup)


Test patch:
diff --git a/tests-cuda/test_3115_array_typed_cuda_jit.py b/tests-cuda/test_3115_array_typed_cuda_jit.py
new file mode 100644
--- /dev/null
+++ b/tests-cuda/test_3115_array_typed_cuda_jit.py
@@ -0,0 +1,40 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
+
+from __future__ import annotations
+
+import pytest
+
+import awkward as ak
+
+nb = pytest.importorskip("numba")
+nb_cuda = pytest.importorskip("numba.cuda")
+
+from numba import types  # noqa: E402
+
+nb.config.CUDA_LOW_OCCUPANCY_WARNINGS = False
+nb.config.CUDA_WARN_ON_IMPLICIT_COPY = False
+
+
+try:
+    ak.numba.register_and_check()
+except ImportError:
+    pytest.skip(reason="too old Numba version", allow_module_level=True)
+
+
+def test_array_typed():
+    # create an ak.Array with a cuda backend:
+    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):
+        return None
+
+    array = ak.Array([[[0, 1], [2, 3]], [[4, 5]]], backend="cuda")
+    cuda_kernel[1024, 1024](array)
+
+    other_array = ak.Array(
+        [[1.1, 2.2, 3.3], [None, 4.4], None, [None, 5.5]], backend="cuda"
+    )
+
+    with pytest.raises(nb.core.errors.NumbaTypeError):
+        cuda_kernel[1024, 1024](other_array)