-
Notifications
You must be signed in to change notification settings - Fork 87
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
typing ak.Array for numba.cuda.jit signature #3105
Comments
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,
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:
It seems to work (no errors): >>> cuda_kernel[1024, 1024](arr) |
Many thanks @jpivarski for your quick response! Your same procedure (grabbing the 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 |
@ianna, do you see what this might be? |
@essoca - Could you, please, try to change it as follows to pin the 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) |
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 |
@essoca - thanks, I can reproduce the error. It is definitely a bug! I think it is related to our way of a generated C++ 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) |
Thanks @ianna. Could you be a bit more explicit about what should be done for that to work? When |
Sorry, I was not clear. You should not try setting the |
@ianna: it partially works. But executing After putting a breakpoint at line 518, the only way for execution to reach 519 successfully is by inspecting |
@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 |
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 likeand 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:
This time, I get the error:
How should this latter case be correctly treated? Note that, without typing, the thing works as expected:
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 thearr
for testing asAny help is appreciated!
The text was updated successfully, but these errors were encountered: