From 1b826c9324344390c43309b195cad85065b733e2 Mon Sep 17 00:00:00 2001 From: Ianna Osborne Date: Wed, 18 Sep 2024 20:36:45 +0200 Subject: [PATCH 01/18] test: add slicing test for CPU and GPU in test_3140_cuda_slicing.py --- tests-cuda/test_3140_cuda_slicing.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests-cuda/test_3140_cuda_slicing.py b/tests-cuda/test_3140_cuda_slicing.py index 59e2cfcb67..147f7d7a8d 100644 --- a/tests-cuda/test_3140_cuda_slicing.py +++ b/tests-cuda/test_3140_cuda_slicing.py @@ -677,3 +677,21 @@ def test_0127_tomask_operation(): [None], [6.6, None, None, 9.9], ] + +arrg_cpu = ak.Array([[1, 2, 3], [0], [4, 5]]) +arrg_gpu = ak.Array([[1, 2, 3], [0], [4, 5]], backend="cuda") + +def test_simple_slice_cpu(): + print("slice on CPU") + out = arrg_cpu[:, 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(): + print("slice on GPU") + out = arrg_gpu[:, 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}") From ba9b4d10dcc163965d82d6e851d2c361a253437a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 18 Sep 2024 18:37:48 +0000 Subject: [PATCH 02/18] style: pre-commit fixes --- tests-cuda/test_3140_cuda_slicing.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/tests-cuda/test_3140_cuda_slicing.py b/tests-cuda/test_3140_cuda_slicing.py index 147f7d7a8d..134ac357e8 100644 --- a/tests-cuda/test_3140_cuda_slicing.py +++ b/tests-cuda/test_3140_cuda_slicing.py @@ -678,15 +678,21 @@ def test_0127_tomask_operation(): [6.6, None, None, 9.9], ] + arrg_cpu = ak.Array([[1, 2, 3], [0], [4, 5]]) arrg_gpu = ak.Array([[1, 2, 3], [0], [4, 5]], backend="cuda") + def test_simple_slice_cpu(): print("slice on CPU") out = arrg_cpu[:, 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}") + 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(): @@ -694,4 +700,8 @@ def test_simple_slice_gpu(): out = arrg_gpu[:, 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}") + 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}", + ) From ef2609e5e791707403f44b03a3ad1d738ad4f374 Mon Sep 17 00:00:00 2001 From: Ianna Osborne Date: Thu, 19 Sep 2024 15:20:34 +0200 Subject: [PATCH 03/18] 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)' --- src/awkward/contents/listarray.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/awkward/contents/listarray.py b/src/awkward/contents/listarray.py index a05eeaea55..fbaebb8aed 100644 --- a/src/awkward/contents/listarray.py +++ b/src/awkward/contents/listarray.py @@ -728,7 +728,7 @@ def _getitem_next( self._starts.data, self._stops.data, lenstarts, - head, + int(head), # NOTE: when 'at' is an integer it is always 'int64' for both CPU and GPU kernels ), slicer=head, ) From ed55752882f944a56f70c7c8c2f838de688fe27e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 19 Sep 2024 13:21:34 +0000 Subject: [PATCH 04/18] style: pre-commit fixes --- src/awkward/contents/listarray.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/awkward/contents/listarray.py b/src/awkward/contents/listarray.py index fbaebb8aed..47b70b7188 100644 --- a/src/awkward/contents/listarray.py +++ b/src/awkward/contents/listarray.py @@ -728,7 +728,9 @@ def _getitem_next( self._starts.data, self._stops.data, lenstarts, - int(head), # NOTE: when 'at' is an integer it is always 'int64' for both CPU and GPU kernels + int( + head + ), # NOTE: when 'at' is an integer it is always 'int64' for both CPU and GPU kernels ), slicer=head, ) From bd7b125691c10fd52e12cc11843c8bb0598e0b51 Mon Sep 17 00:00:00 2001 From: Ianna Osborne Date: Thu, 19 Sep 2024 15:45:46 +0200 Subject: [PATCH 05/18] use ak._slicing.normalize_integer_like(head) --- src/awkward/contents/listarray.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/awkward/contents/listarray.py b/src/awkward/contents/listarray.py index 47b70b7188..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 @@ -728,9 +729,7 @@ def _getitem_next( self._starts.data, self._stops.data, lenstarts, - int( - head - ), # NOTE: when 'at' is an integer it is always 'int64' for both CPU and GPU kernels + head, ), slicer=head, ) From 298ecf9cbb93048ef833093098fb01df496db682 Mon Sep 17 00:00:00 2001 From: Ianna Osborne Date: Thu, 19 Sep 2024 16:32:41 +0200 Subject: [PATCH 06/18] convert head ndarray to scalar --- src/awkward/contents/listarray.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/awkward/contents/listarray.py b/src/awkward/contents/listarray.py index 8f0a6f5e4b..e6ff1d9cf1 100644 --- a/src/awkward/contents/listarray.py +++ b/src/awkward/contents/listarray.py @@ -729,7 +729,7 @@ def _getitem_next( self._starts.data, self._stops.data, lenstarts, - head, + head.item(), ), slicer=head, ) From 2627a5158aa420f6ba14c1d3d8ed5cc5677b2392 Mon Sep 17 00:00:00 2001 From: Ianna Osborne Date: Thu, 19 Sep 2024 16:52:35 +0200 Subject: [PATCH 07/18] add item attribute to TypeTracerArray --- src/awkward/_nplikes/typetracer.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/awkward/_nplikes/typetracer.py b/src/awkward/_nplikes/typetracer.py index 64b99abf01..14db7f49a0 100644 --- a/src/awkward/_nplikes/typetracer.py +++ b/src/awkward/_nplikes/typetracer.py @@ -562,6 +562,9 @@ def __setitem__( if isinstance(value, TypeTracerArray) and value.ndim > existing_value.ndim: raise ValueError("cannot assign shape larger than destination") + def item(self): + return self + def copy(self): self.touch_data() return self From 6333a6c5e5b5053c135bcbd1a3b8339ec6a3d821 Mon Sep 17 00:00:00 2001 From: Ianna Osborne Date: Thu, 19 Sep 2024 16:56:00 +0200 Subject: [PATCH 08/18] cleanup tests --- tests-cuda/test_3140_cuda_slicing.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tests-cuda/test_3140_cuda_slicing.py b/tests-cuda/test_3140_cuda_slicing.py index 134ac357e8..34bc28f84b 100644 --- a/tests-cuda/test_3140_cuda_slicing.py +++ b/tests-cuda/test_3140_cuda_slicing.py @@ -679,12 +679,8 @@ def test_0127_tomask_operation(): ] -arrg_cpu = ak.Array([[1, 2, 3], [0], [4, 5]]) -arrg_gpu = ak.Array([[1, 2, 3], [0], [4, 5]], backend="cuda") - - def test_simple_slice_cpu(): - print("slice on CPU") + arrg_cpu = ak.Array([[1, 2, 3], [0], [4, 5]]) out = arrg_cpu[:, 0] expected = [1, 0, 4] result = out.tolist() @@ -696,7 +692,7 @@ def test_simple_slice_cpu(): def test_simple_slice_gpu(): - print("slice on GPU") + arrg_gpu = ak.Array([[1, 2, 3], [0], [4, 5]], backend="cuda") out = arrg_gpu[:, 0] expected = [1, 0, 4] result = out.tolist() From da434dc75eb695c21d8c771dc0f0ce82a1cd6d8f Mon Sep 17 00:00:00 2001 From: Ianna Osborne Date: Thu, 19 Sep 2024 19:03:03 +0200 Subject: [PATCH 09/18] add more tests --- tests-cuda/test_3140_cuda_slicing.py | 72 ++++++++++++++++++++++++++-- 1 file changed, 68 insertions(+), 4 deletions(-) diff --git a/tests-cuda/test_3140_cuda_slicing.py b/tests-cuda/test_3140_cuda_slicing.py index 34bc28f84b..828bfe552b 100644 --- a/tests-cuda/test_3140_cuda_slicing.py +++ b/tests-cuda/test_3140_cuda_slicing.py @@ -680,8 +680,8 @@ def test_0127_tomask_operation(): def test_simple_slice_cpu(): - arrg_cpu = ak.Array([[1, 2, 3], [0], [4, 5]]) - out = arrg_cpu[:, 0] + 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( @@ -692,8 +692,8 @@ def test_simple_slice_cpu(): def test_simple_slice_gpu(): - arrg_gpu = ak.Array([[1, 2, 3], [0], [4, 5]], backend="cuda") - out = arrg_gpu[:, 0] + 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( @@ -701,3 +701,67 @@ def test_simple_slice_gpu(): 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}") From 5ed207262904f802655f81259b14b288fb50fd3d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 19 Sep 2024 17:03:27 +0000 Subject: [PATCH 10/18] style: pre-commit fixes --- tests-cuda/test_3140_cuda_slicing.py | 68 ++++++++++++++++++++-------- 1 file changed, 50 insertions(+), 18 deletions(-) diff --git a/tests-cuda/test_3140_cuda_slicing.py b/tests-cuda/test_3140_cuda_slicing.py index 828bfe552b..da1cd4bef7 100644 --- a/tests-cuda/test_3140_cuda_slicing.py +++ b/tests-cuda/test_3140_cuda_slicing.py @@ -705,63 +705,95 @@ def test_simple_slice_gpu(): def test_simple_slice_cpu1(): arr = ak.Array([[1, 2, 3], [0], [4, 5]]) - out = arr[:,1:] - expected = [[2, 3], [], [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}") + 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]] + 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}") + 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] + 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}") + 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] + 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}") + 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] + 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}") + 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] + 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}") + 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] + 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}") + 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] + 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}") + 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}", + ) From 2fdbb2a0bf07c36658396f45fb21fc472fa7cb72 Mon Sep 17 00:00:00 2001 From: Ianna Osborne Date: Fri, 20 Sep 2024 16:36:17 +0200 Subject: [PATCH 11/18] use 'ListArray-at' role --- dev/generate-kernel-signatures.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dev/generate-kernel-signatures.py b/dev/generate-kernel-signatures.py index 74bc3c4fe1..c18843bf62 100644 --- a/dev/generate-kernel-signatures.py +++ b/dev/generate-kernel-signatures.py @@ -429,7 +429,7 @@ 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( From 095a2ed8940c938ad8a0c56b516829702caa8bf1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 20 Sep 2024 14:36:55 +0000 Subject: [PATCH 12/18] style: pre-commit fixes --- dev/generate-kernel-signatures.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/dev/generate-kernel-signatures.py b/dev/generate-kernel-signatures.py index c18843bf62..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"] or 'ListArray-at' == x.get('role', None)) 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( From 408a24b2387c256b54f1dbeed8dd5c4341f1ca19 Mon Sep 17 00:00:00 2001 From: Ianna Osborne Date: Fri, 20 Sep 2024 16:37:17 +0200 Subject: [PATCH 13/18] use pointer awkward_ListArray_getitem_next_at.cu --- .../cuda/cuda_kernels/awkward_ListArray_getitem_next_at.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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; } From 2a12275d1767e085b74bb64813c1ac09c289582d Mon Sep 17 00:00:00 2001 From: Ianna Osborne Date: Fri, 20 Sep 2024 16:38:20 +0200 Subject: [PATCH 14/18] use role listarray.py --- src/awkward/contents/listarray.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/awkward/contents/listarray.py b/src/awkward/contents/listarray.py index e6ff1d9cf1..8f0a6f5e4b 100644 --- a/src/awkward/contents/listarray.py +++ b/src/awkward/contents/listarray.py @@ -729,7 +729,7 @@ def _getitem_next( self._starts.data, self._stops.data, lenstarts, - head.item(), + head, ), slicer=head, ) From 13d3a49d41d950b1efd360a820fc837f7d5c90c2 Mon Sep 17 00:00:00 2001 From: Ianna Osborne Date: Fri, 20 Sep 2024 16:39:25 +0200 Subject: [PATCH 15/18] revert changes --- src/awkward/_nplikes/typetracer.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/awkward/_nplikes/typetracer.py b/src/awkward/_nplikes/typetracer.py index 14db7f49a0..64b99abf01 100644 --- a/src/awkward/_nplikes/typetracer.py +++ b/src/awkward/_nplikes/typetracer.py @@ -562,9 +562,6 @@ def __setitem__( if isinstance(value, TypeTracerArray) and value.ndim > existing_value.ndim: raise ValueError("cannot assign shape larger than destination") - def item(self): - return self - def copy(self): self.touch_data() return self From 11f3f4e35b9f3b6a4e1c24c010460a71e7460f94 Mon Sep 17 00:00:00 2001 From: Ianna Osborne Date: Fri, 20 Sep 2024 17:13:40 +0200 Subject: [PATCH 16/18] same for awkward_RegularArray_getitem_next_at.cu --- .../cuda/cuda_kernels/awkward_RegularArray_getitem_next_at.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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; } From a7f8f3b4a1d08d8cc51ed07f64f26c9ae7d7740d Mon Sep 17 00:00:00 2001 From: Ianna Osborne Date: Fri, 20 Sep 2024 17:15:03 +0200 Subject: [PATCH 17/18] remove special case for cuda --- src/awkward/contents/regulararray.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) 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 From a8845ee91e93abf117e0f42860225be3312e0c1b Mon Sep 17 00:00:00 2001 From: Ianna Osborne Date: Mon, 23 Sep 2024 15:51:43 +0200 Subject: [PATCH 18/18] correct role for jagged size --- kernel-specification.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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: |