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}",
+    )