Skip to content
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

fix: generate array view #3115

Merged
merged 15 commits into from
May 21, 2024
17 changes: 13 additions & 4 deletions src/awkward/_connect/numba/arrayview_cuda.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
40 changes: 40 additions & 0 deletions tests-cuda/test_3115_array_typed_cuda_jit.py
Original file line number Diff line number Diff line change
@@ -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)
Loading