fix: ListArray slicing on GPU (#3248)

* test: add slicing test for CPU and GPU in test_3140_cuda_slicing.py

* style: pre-commit fixes

* cast 'at' to int

head in this case can be an array and it can be regularized to a proper backend, then the GPU kernel needs to be updated to handle a 'cp.array(0)'

* style: pre-commit fixes

* use ak._slicing.normalize_integer_like(head)

* convert head ndarray to scalar

* add item attribute to TypeTracerArray

* cleanup tests

* add more tests

* style: pre-commit fixes

* use 'ListArray-at' role

* style: pre-commit fixes

* use pointer awkward_ListArray_getitem_next_at.cu

* use role listarray.py

* revert changes

* same for awkward_RegularArray_getitem_next_at.cu

* remove special case for cuda

* correct role for jagged size

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

dev/generate-kernel-signatures.py

@@ -429,7 +429,10 @@ def by_signature(cuda_kernel_templates):
  special = [repr(spec["name"])]
  [type_to_pytype(x["type"], special) for x in childfunc["args"]]
  dirlist = [repr(x["dir"]) for x in childfunc["args"]]
- ispointerlist = [repr("List" in x["type"]) for x in childfunc["args"]]
+ ispointerlist = [
+ repr("List" in x["type"] or "ListArray-at" == x.get("role", None))
+ for x in childfunc["args"]
+ ]
  if spec["name"] in cuda_kernels_impl:
  with open(
  os.path.join(

kernel-specification.yml

@@ -1466,7 +1466,7 @@ kernels:
  - {name: tocarry, type: "List[int64_t]", dir: out}
  - {name: fromstarts, type: "Const[List[int32_t]]", dir: in, role: ListArray-starts}
  - {name: fromstops, type: "Const[List[int32_t]]", dir: in, role: ListArray-stops}
- - {name: jaggedsize, type: "int64_t", dir: in, role: ListArray-at}
+ - {name: jaggedsize, type: "int64_t", dir: in, role: ListArray-length}
  - {name: length, type: "int64_t", dir: in, role: default}
  - name: awkward_ListArray64_getitem_jagged_expand_64
  args:
@@ -1476,7 +1476,7 @@ kernels:
  - {name: tocarry, type: "List[int64_t]", dir: out}
  - {name: fromstarts, type: "Const[List[int64_t]]", dir: in, role: ListArray-starts}
  - {name: fromstops, type: "Const[List[int64_t]]", dir: in, role: ListArray-stops}
- - {name: jaggedsize, type: "int64_t", dir: in, role: ListArray-at}
+ - {name: jaggedsize, type: "int64_t", dir: in, role: ListArray-length}
  - {name: length, type: "int64_t", dir: in, role: default}
  - name: awkward_ListArrayU32_getitem_jagged_expand_64
  args:
@@ -1486,7 +1486,7 @@ kernels:
  - {name: tocarry, type: "List[int64_t]", dir: out}
  - {name: fromstarts, type: "Const[List[uint32_t]]", dir: in, role: ListArray-starts}
  - {name: fromstops, type: "Const[List[uint32_t]]", dir: in, role: ListArray-stops}
- - {name: jaggedsize, type: "int64_t", dir: in, role: ListArray-at}
+ - {name: jaggedsize, type: "int64_t", dir: in, role: ListArray-length}
  - {name: length, type: "int64_t", dir: in, role: default}
  description: null
  definition: |

src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_next_at.cu

@@ -11,15 +11,15 @@ awkward_ListArray_getitem_next_at(
  const C* fromstarts,
  const U* fromstops,
  int64_t lenstarts,
- int64_t at,
+ int64_t* at,
  uint64_t invocation_index,
  uint64_t* err_code) {
  if (err_code[0] == NO_ERROR) {
  int64_t thread_id = blockIdx.x * blockDim.x + threadIdx.x;
 
  if (thread_id < lenstarts) {
  int64_t length = fromstops[thread_id] - fromstarts[thread_id];
- int64_t regular_at = at;
+ int64_t regular_at = at[0];
  if (regular_at < 0) {
  regular_at += length;
  }

src/awkward/_connect/cuda/cuda_kernels/awkward_RegularArray_getitem_next_at.cu

@@ -8,14 +8,14 @@ template <typename T>
 __global__ void
 awkward_RegularArray_getitem_next_at(
  T* tocarry,
- int64_t at,
+ int64_t* at,
  int64_t length,
  int64_t size,
  uint64_t invocation_index,
  uint64_t* err_code) {
  if (err_code[0] == NO_ERROR) {
  int64_t thread_id = blockIdx.x * blockDim.x + threadIdx.x;
- int64_t regular_at = at;
+ int64_t regular_at = at[0];
  if (regular_at < 0) {
  regular_at += size;
  }

src/awkward/contents/listarray.py

@@ -712,6 +712,7 @@ def _getitem_next(
  nexthead, nexttail = ak._slicing.head_tail(tail)
  lenstarts = self._starts.length
  nextcarry = ak.index.Index64.empty(lenstarts, self._backend.index_nplike)
+ head = ak._slicing.normalize_integer_like(head)
  assert (
  nextcarry.nplike is self._backend.index_nplike
  and self._starts.nplike is self._backend.index_nplike

src/awkward/contents/regulararray.py

@@ -471,8 +471,7 @@ def _getitem_next(
  nexthead, nexttail = ak._slicing.head_tail(tail)
  nextcarry = ak.index.Index64.empty(self._length, index_nplike)
  assert nextcarry.nplike is index_nplike
- if ak.backend(head) == "cuda":
- head = int(ak.to_backend(head, backend=self._backend)[0])
+ head = ak._slicing.normalize_integer_like(head)
  self._maybe_index_error(
  self._backend[
  "awkward_RegularArray_getitem_next_at", nextcarry.dtype.type

tests-cuda/test_3140_cuda_slicing.py

@@ -677,3 +677,123 @@ def test_0127_tomask_operation():
  [None],
  [6.6, None, None, 9.9],
  ]
+
+
+def test_simple_slice_cpu():
+ arr = ak.Array([[1, 2, 3], [0], [4, 5]])
+ out = arr[:, 0]
+ expected = [1, 0, 4]
+ result = out.tolist()
+ cp.testing.assert_array_list_equal(
+ result,
+ expected,
+ err_msg=f"Slice of [[1, 2, 3], [0], [4, 5]] should be {expected}, but got {result}",
+ )
+
+
+def test_simple_slice_gpu():
+ arr = ak.Array([[1, 2, 3], [0], [4, 5]], backend="cuda")
+ out = arr[:, 0]
+ expected = [1, 0, 4]
+ result = out.tolist()
+ cp.testing.assert_array_list_equal(
+ result,
+ expected,
+ err_msg=f"Slice of [[1, 2, 3], [0], [4, 5]] should be {expected}, but got {result}",
+ )
+
+
+def test_simple_slice_cpu1():
+ arr = ak.Array([[1, 2, 3], [0], [4, 5]])
+ out = arr[:, 1:]
+ expected = [[2, 3], [], [5]]
+ result = out.tolist()
+ cp.testing.assert_array_list_equal(
+ result,
+ expected,
+ err_msg=f"Slice of [[1, 2, 3], [0], [4, 5]] should be {expected}, but got {result}",
+ )
+
+
+def test_simple_slice_gpu1():
+ arr = ak.Array([[1, 2, 3], [0], [4, 5]], backend="cuda")
+ out = arr[:, 1:]
+ expected = [[2, 3], [], [5]]
+ result = out.tolist()
+ cp.testing.assert_array_list_equal(
+ result,
+ expected,
+ err_msg=f"Slice of [[1, 2, 3], [0], [4, 5]] should be {expected}, but got {result}",
+ )
+
+
+def test_simple_slice_cpu2():
+ arr = ak.Array([[1, 2, 3], [0], [4, 5]])
+ out = arr[:, :1]
+ expected = [[1], [0], [4]]
+ result = out.tolist()
+ cp.testing.assert_array_list_equal(
+ result,
+ expected,
+ err_msg=f"Slice of [[1, 2, 3], [0], [4, 5]] should be {expected}, but got {result}",
+ )
+
+
+def test_simple_slice_gpu2():
+ arr = ak.Array([[1, 2, 3], [0], [4, 5]], backend="cuda")
+ out = arr[:, :1]
+ expected = [[1], [0], [4]]
+ result = out.tolist()
+ cp.testing.assert_array_list_equal(
+ result,
+ expected,
+ err_msg=f"Slice of [[1, 2, 3], [0], [4, 5]] should be {expected}, but got {result}",
+ )
+
+
+def test_simple_slice_cpu3():
+ arr = ak.Array([[1, 2, 3], [0], [4, 5]])
+ out = arr[:, 1::2]
+ expected = [[2], [], [5]]
+ result = out.tolist()
+ cp.testing.assert_array_list_equal(
+ result,
+ expected,
+ err_msg=f"Slice of [[1, 2, 3], [0], [4, 5]] should be {expected}, but got {result}",
+ )
+
+
+def test_simple_slice_gpu3():
+ arr = ak.Array([[1, 2, 3], [0], [4, 5]], backend="cuda")
+ out = arr[:, 1::2]
+ expected = [[2], [], [5]]
+ result = out.tolist()
+ cp.testing.assert_array_list_equal(
+ result,
+ expected,
+ err_msg=f"Slice of [[1, 2, 3], [0], [4, 5]] should be {expected}, but got {result}",
+ )
+
+
+def test_simple_slice_cpu4():
+ arr = ak.Array([[1, 2, 3], [0], [4, 5]])
+ out = arr[:, ::-1]
+ expected = [[3, 2, 1], [0], [5, 4]]
+ result = out.tolist()
+ cp.testing.assert_array_list_equal(
+ result,
+ expected,
+ err_msg=f"Slice of [[1, 2, 3], [0], [4, 5]] should be {expected}, but got {result}",
+ )
+
+
+def test_simple_slice_gpu4():
+ arr = ak.Array([[1, 2, 3], [0], [4, 5]], backend="cuda")
+ out = arr[:, ::-1]
+ expected = [[3, 2, 1], [0], [5, 4]]
+ result = out.tolist()
+ cp.testing.assert_array_list_equal(
+ result,
+ expected,
+ err_msg=f"Slice of [[1, 2, 3], [0], [4, 5]] should be {expected}, but got {result}",
+ )