From e946646c6315ff9bde0c8a2255ee6f7fc6a31539 Mon Sep 17 00:00:00 2001 From: Ianna Osborne Date: Thu, 26 Sep 2024 18:01:35 +0200 Subject: [PATCH] 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 | 5 +- kernel-specification.yml | 6 +- .../awkward_ListArray_getitem_next_at.cu | 4 +- .../awkward_RegularArray_getitem_next_at.cu | 4 +- src/awkward/contents/listarray.py | 1 + src/awkward/contents/regulararray.py | 3 +- tests-cuda/test_3140_cuda_slicing.py | 120 ++++++++++++++++++ 7 files changed, 133 insertions(+), 10 deletions(-) diff --git a/dev/generate-kernel-signatures.py b/dev/generate-kernel-signatures.py index 74bc3c4fe1..7038d95f7f 100644 --- a/dev/generate-kernel-signatures.py +++ b/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( diff --git a/kernel-specification.yml b/kernel-specification.yml index 2838b8db5c..5c901f00e2 100644 --- a/kernel-specification.yml +++ b/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: | diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_next_at.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_next_at.cu index 421f0d15c1..7fc13f2ae8 100644 --- a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_next_at.cu +++ b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_next_at.cu @@ -11,7 +11,7 @@ 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) { @@ -19,7 +19,7 @@ awkward_ListArray_getitem_next_at( 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; } diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_RegularArray_getitem_next_at.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_RegularArray_getitem_next_at.cu index 8f1282974d..1b8bd53b38 100644 --- a/src/awkward/_connect/cuda/cuda_kernels/awkward_RegularArray_getitem_next_at.cu +++ b/src/awkward/_connect/cuda/cuda_kernels/awkward_RegularArray_getitem_next_at.cu @@ -8,14 +8,14 @@ template __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; } diff --git a/src/awkward/contents/listarray.py b/src/awkward/contents/listarray.py index a05eeaea55..8f0a6f5e4b 100644 --- a/src/awkward/contents/listarray.py +++ b/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 diff --git a/src/awkward/contents/regulararray.py b/src/awkward/contents/regulararray.py index a5a16fcdff..318a21bca5 100644 --- a/src/awkward/contents/regulararray.py +++ b/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 diff --git a/tests-cuda/test_3140_cuda_slicing.py b/tests-cuda/test_3140_cuda_slicing.py index 59e2cfcb67..da1cd4bef7 100644 --- a/tests-cuda/test_3140_cuda_slicing.py +++ b/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}", + )