From cbf5fc7b0f44f011c2ad67a6822ce18ea7b52372 Mon Sep 17 00:00:00 2001 From: ManasviGoyal Date: Tue, 4 Jun 2024 11:34:22 +0200 Subject: [PATCH 1/7] feat add awkward_ListArray_getitem_jagged_apply kernel --- dev/generate-kernel-signatures.py | 1 + dev/generate-tests.py | 1 + src/awkward/_connect/cuda/__init__.py | 1 + .../awkward_ListArray_getitem_jagged_apply.cu | 123 ++++++++++++++++++ 4 files changed, 126 insertions(+) create mode 100644 src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_apply.cu diff --git a/dev/generate-kernel-signatures.py b/dev/generate-kernel-signatures.py index 7b1c55c8e9..6a61074505 100644 --- a/dev/generate-kernel-signatures.py +++ b/dev/generate-kernel-signatures.py @@ -53,6 +53,7 @@ "awkward_missing_repeat", "awkward_RegularArray_getitem_jagged_expand", "awkward_ListArray_combinations_length", + "awkward_ListArray_getitem_jagged_apply", "awkward_ListArray_getitem_jagged_carrylen", "awkward_ListArray_getitem_jagged_descend", "awkward_ListArray_getitem_jagged_expand", diff --git a/dev/generate-tests.py b/dev/generate-tests.py index 5e544ff36d..6672460c99 100644 --- a/dev/generate-tests.py +++ b/dev/generate-tests.py @@ -838,6 +838,7 @@ def gencpuunittests(specdict): "awkward_missing_repeat", "awkward_RegularArray_getitem_jagged_expand", "awkward_ListArray_combinations_length", + "awkward_ListArray_getitem_jagged_apply", "awkward_ListArray_getitem_jagged_carrylen", "awkward_ListArray_getitem_jagged_descend", "awkward_ListArray_getitem_jagged_expand", diff --git a/src/awkward/_connect/cuda/__init__.py b/src/awkward/_connect/cuda/__init__.py index d48aab5ccd..a9a590fecc 100644 --- a/src/awkward/_connect/cuda/__init__.py +++ b/src/awkward/_connect/cuda/__init__.py @@ -96,6 +96,7 @@ def fetch_template_specializations(kernel_dict): "awkward_ListArray_broadcast_tooffsets", "awkward_ListArray_combinations_length", "awkward_ListArray_compact_offsets", + "awkward_ListArray_getitem_jagged_apply", "awkward_ListArray_getitem_jagged_carrylen", "awkward_ListArray_getitem_jagged_descend", "awkward_ListArray_getitem_jagged_numvalid", diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_apply.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_apply.cu new file mode 100644 index 0000000000..83cccee4ec --- /dev/null +++ b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_apply.cu @@ -0,0 +1,123 @@ +// BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE + +// BEGIN PYTHON +// def f(grid, block, args): +// (tooffsets, tocarry, slicestarts, slicestops, sliceouterlen, sliceindex, sliceinnerlen, fromstarts, fromstops, contentlen, invocation_index, err_code) = args +// scan_in_array = cupy.zeros(sliceouterlen + 1, dtype=cupy.int64) +// cuda_kernel_templates.get_function(fetch_specialization(["awkward_ListArray_getitem_jagged_apply_a", tooffsets.dtype, tocarry.dtype, slicestarts.dtype, slicestops.dtype, sliceindex.dtype, fromstarts.dtype, fromstops.dtype]))(grid, block, (tooffsets, tocarry, slicestarts, slicestops, sliceouterlen, sliceindex, sliceinnerlen, fromstarts, fromstops, contentlen, scan_in_array, invocation_index, err_code)) +// scan_in_array = cupy.cumsum(scan_in_array) +// print(scan_in_array) +// cuda_kernel_templates.get_function(fetch_specialization(["awkward_ListArray_getitem_jagged_apply_b", tooffsets.dtype, tocarry.dtype, slicestarts.dtype, slicestops.dtype, sliceindex.dtype, fromstarts.dtype, fromstops.dtype]))(grid, block, (tooffsets, tocarry, slicestarts, slicestops, sliceouterlen, sliceindex, sliceinnerlen, fromstarts, fromstops, contentlen, scan_in_array, invocation_index, err_code)) +// out["awkward_ListArray_getitem_jagged_apply_a", {dtype_specializations}] = None +// out["awkward_ListArray_getitem_jagged_apply_b", {dtype_specializations}] = None +// END PYTHON + +enum class LISTARRAY_GETITEM_JAGGED_APPLY_ERRORS { + JAG_STOP_LT_START, // message: "jagged slice's stops[i] < starts[i]" + OFF_GET_CON, // message: "jagged slice's offsets extend beyond its content" + STOP_LT_START, // message: "stops[i] < starts[i]" + STOP_GET_LEN, // message: "stops[i] > len(content)" + IND_OUT_OF_RANGE, // message: "index out of range" +}; + +template +__global__ void +awkward_ListArray_getitem_jagged_apply_a( + T* tooffsets, + C* tocarry, + const U* slicestarts, + const V* slicestops, + int64_t sliceouterlen, + const W* sliceindex, + int64_t sliceinnerlen, + const X* fromstarts, + const Y* fromstops, + int64_t contentlen, + int64_t* scan_in_array, + uint64_t invocation_index, + uint64_t* err_code) { + if (err_code[0] == NO_ERROR) { + int64_t thread_id = blockIdx.x * blockDim.x + threadIdx.x; + scan_in_array[0] = 0; + + if (thread_id < sliceouterlen) { + U slicestart = slicestarts[thread_id]; + V slicestop = slicestops[thread_id]; + + if (slicestart != slicestop) { + if (slicestop < slicestart) { + RAISE_ERROR(LISTARRAY_GETITEM_JAGGED_APPLY_ERRORS::JAG_STOP_LT_START) + } + if (slicestop > sliceinnerlen) { + RAISE_ERROR(LISTARRAY_GETITEM_JAGGED_APPLY_ERRORS::OFF_GET_CON) + } + int64_t start = (int64_t)fromstarts[thread_id]; + int64_t stop = (int64_t)fromstops[thread_id]; + if (stop < start) { + RAISE_ERROR(LISTARRAY_GETITEM_JAGGED_APPLY_ERRORS::STOP_LT_START) + } + if (start != stop && stop > contentlen) { + RAISE_ERROR(LISTARRAY_GETITEM_JAGGED_APPLY_ERRORS::STOP_GET_LEN) + } + int64_t count = stop - start; + scan_in_array[thread_id + 1] = slicestop - slicestart; + } + } + } +} + +template +__global__ void +awkward_ListArray_getitem_jagged_apply_b( + T* tooffsets, + C* tocarry, + const U* slicestarts, + const V* slicestops, + int64_t sliceouterlen, + const W* sliceindex, + int64_t sliceinnerlen, + const X* fromstarts, + const Y* fromstops, + int64_t contentlen, + int64_t* scan_in_array, + 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 < sliceouterlen) { + U slicestart = slicestarts[thread_id]; + V slicestop = slicestops[thread_id]; + tooffsets[thread_id] = (T)(scan_in_array[thread_id]); + if (slicestart != slicestop) { + if (slicestop < slicestart) { + RAISE_ERROR(LISTARRAY_GETITEM_JAGGED_APPLY_ERRORS::JAG_STOP_LT_START) + } + if (slicestop > sliceinnerlen) { + RAISE_ERROR(LISTARRAY_GETITEM_JAGGED_APPLY_ERRORS::OFF_GET_CON) + } + int64_t start = (int64_t)fromstarts[thread_id]; + int64_t stop = (int64_t)fromstops[thread_id]; + if (stop < start) { + RAISE_ERROR(LISTARRAY_GETITEM_JAGGED_APPLY_ERRORS::STOP_LT_START) + } + if (start != stop && stop > contentlen) { + RAISE_ERROR(LISTARRAY_GETITEM_JAGGED_APPLY_ERRORS::STOP_GET_LEN) + } + int64_t count = stop - start; + for (int64_t j = slicestart; j < slicestop; j++) { + int64_t index = (int64_t) sliceindex[j]; + if (index < -count || index > count) { + RAISE_ERROR(LISTARRAY_GETITEM_JAGGED_APPLY_ERRORS::IND_OUT_OF_RANGE) + } + if (index < 0) { + index += count; + } + tocarry[scan_in_array[thread_id] + j - slicestart] = start + index; + printf("%d, ", scan_in_array[thread_id] + j - slicestart); + } + } + } + tooffsets[sliceouterlen] = scan_in_array[sliceouterlen]; + } +} From 0e76fba54269e27bb002375c7cc89d2561629f38 Mon Sep 17 00:00:00 2001 From: ManasviGoyal Date: Tue, 4 Jun 2024 11:46:08 +0200 Subject: [PATCH 2/7] fix: remove print statements --- .../cuda/cuda_kernels/awkward_ListArray_getitem_jagged_apply.cu | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_apply.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_apply.cu index 83cccee4ec..88a7398352 100644 --- a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_apply.cu +++ b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_apply.cu @@ -6,7 +6,6 @@ // scan_in_array = cupy.zeros(sliceouterlen + 1, dtype=cupy.int64) // cuda_kernel_templates.get_function(fetch_specialization(["awkward_ListArray_getitem_jagged_apply_a", tooffsets.dtype, tocarry.dtype, slicestarts.dtype, slicestops.dtype, sliceindex.dtype, fromstarts.dtype, fromstops.dtype]))(grid, block, (tooffsets, tocarry, slicestarts, slicestops, sliceouterlen, sliceindex, sliceinnerlen, fromstarts, fromstops, contentlen, scan_in_array, invocation_index, err_code)) // scan_in_array = cupy.cumsum(scan_in_array) -// print(scan_in_array) // cuda_kernel_templates.get_function(fetch_specialization(["awkward_ListArray_getitem_jagged_apply_b", tooffsets.dtype, tocarry.dtype, slicestarts.dtype, slicestops.dtype, sliceindex.dtype, fromstarts.dtype, fromstops.dtype]))(grid, block, (tooffsets, tocarry, slicestarts, slicestops, sliceouterlen, sliceindex, sliceinnerlen, fromstarts, fromstops, contentlen, scan_in_array, invocation_index, err_code)) // out["awkward_ListArray_getitem_jagged_apply_a", {dtype_specializations}] = None // out["awkward_ListArray_getitem_jagged_apply_b", {dtype_specializations}] = None @@ -114,7 +113,6 @@ awkward_ListArray_getitem_jagged_apply_b( index += count; } tocarry[scan_in_array[thread_id] + j - slicestart] = start + index; - printf("%d, ", scan_in_array[thread_id] + j - slicestart); } } } From c429c5266aaeb3e792b50bd3433eb3235f1d922a Mon Sep 17 00:00:00 2001 From: ManasviGoyal Date: Tue, 4 Jun 2024 14:04:04 +0200 Subject: [PATCH 3/7] feat: add awkward_ListArray_getitem_jagged_shrink kernel --- dev/generate-kernel-signatures.py | 1 + dev/generate-tests.py | 1 + src/awkward/_connect/cuda/__init__.py | 1 + .../awkward_ListArray_getitem_jagged_apply.cu | 2 +- ...awkward_ListArray_getitem_jagged_shrink.cu | 105 ++++++++++++++++++ ...istArray_getitem_next_range_carrylength.cu | 2 - 6 files changed, 109 insertions(+), 3 deletions(-) create mode 100644 src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_shrink.cu diff --git a/dev/generate-kernel-signatures.py b/dev/generate-kernel-signatures.py index 6a61074505..67d1e69a9f 100644 --- a/dev/generate-kernel-signatures.py +++ b/dev/generate-kernel-signatures.py @@ -58,6 +58,7 @@ "awkward_ListArray_getitem_jagged_descend", "awkward_ListArray_getitem_jagged_expand", "awkward_ListArray_getitem_jagged_numvalid", + "awkward_ListArray_getitem_jagged_shrink", "awkward_ListArray_getitem_next_array_advanced", "awkward_ListArray_getitem_next_array", "awkward_ListArray_getitem_next_at", diff --git a/dev/generate-tests.py b/dev/generate-tests.py index 6672460c99..e434c2093e 100644 --- a/dev/generate-tests.py +++ b/dev/generate-tests.py @@ -843,6 +843,7 @@ def gencpuunittests(specdict): "awkward_ListArray_getitem_jagged_descend", "awkward_ListArray_getitem_jagged_expand", "awkward_ListArray_getitem_jagged_numvalid", + "awkward_ListArray_getitem_jagged_shrink", "awkward_ListArray_getitem_next_array_advanced", "awkward_ListArray_getitem_next_array", "awkward_ListArray_getitem_next_at", diff --git a/src/awkward/_connect/cuda/__init__.py b/src/awkward/_connect/cuda/__init__.py index a9a590fecc..098d1c4525 100644 --- a/src/awkward/_connect/cuda/__init__.py +++ b/src/awkward/_connect/cuda/__init__.py @@ -100,6 +100,7 @@ def fetch_template_specializations(kernel_dict): "awkward_ListArray_getitem_jagged_carrylen", "awkward_ListArray_getitem_jagged_descend", "awkward_ListArray_getitem_jagged_numvalid", + "awkward_ListArray_getitem_jagged_shrink", "awkward_ListArray_getitem_next_range", "awkward_ListArray_getitem_next_range_carrylength", "awkward_ListArray_min_range", diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_apply.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_apply.cu index 88a7398352..b3017113b4 100644 --- a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_apply.cu +++ b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_apply.cu @@ -104,7 +104,7 @@ awkward_ListArray_getitem_jagged_apply_b( RAISE_ERROR(LISTARRAY_GETITEM_JAGGED_APPLY_ERRORS::STOP_GET_LEN) } int64_t count = stop - start; - for (int64_t j = slicestart; j < slicestop; j++) { + for (int64_t j = slicestart + threadIdx.y; j < slicestop; j += blockDim.y) { int64_t index = (int64_t) sliceindex[j]; if (index < -count || index > count) { RAISE_ERROR(LISTARRAY_GETITEM_JAGGED_APPLY_ERRORS::IND_OUT_OF_RANGE) diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_shrink.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_shrink.cu new file mode 100644 index 0000000000..477af8e65c --- /dev/null +++ b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_shrink.cu @@ -0,0 +1,105 @@ +// BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE + +// BEGIN PYTHON +// def f(grid, block, args): +// (tocarry, tosmalloffsets, tolargeoffsets, slicestarts, slicestops, length, missing, invocation_index, err_code) = args +// if length > 0 and length < int(slicestops[length - 1]): +// len_array = int(slicestops[length - 1]) +// else: +// len_array = length +// scan_in_array_k = cupy.zeros(len_array, dtype=cupy.int64) +// scan_in_array_tosmalloffsets = cupy.zeros(length + 1, dtype=cupy.int64) +// scan_in_array_tolargeoffsets = cupy.zeros(length + 1, dtype=cupy.int64) +// cuda_kernel_templates.get_function(fetch_specialization(["awkward_ListArray_getitem_jagged_shrink_a", tocarry.dtype, tosmalloffsets.dtype, tolargeoffsets.dtype, slicestarts.dtype, slicestops.dtype, missing.dtype]))(grid, block, (tocarry, tosmalloffsets, tolargeoffsets, slicestarts, slicestops, length, missing, scan_in_array_k, scan_in_array_tosmalloffsets, scan_in_array_tolargeoffsets, invocation_index, err_code)) +// scan_in_array_k = cupy.cumsum(scan_in_array_k) +// scan_in_array_tosmalloffsets = cupy.cumsum(scan_in_array_tosmalloffsets) +// scan_in_array_tolargeoffsets = cupy.cumsum(scan_in_array_tolargeoffsets) +// cuda_kernel_templates.get_function(fetch_specialization(["awkward_ListArray_getitem_jagged_shrink_b", tocarry.dtype, tosmalloffsets.dtype, tolargeoffsets.dtype, slicestarts.dtype, slicestops.dtype, missing.dtype]))(grid, block, (tocarry, tosmalloffsets, tolargeoffsets, slicestarts, slicestops, length, missing, scan_in_array_k, scan_in_array_tosmalloffsets, scan_in_array_tolargeoffsets, invocation_index, err_code)) +// out["awkward_ListArray_getitem_jagged_shrink_a", {dtype_specializations}] = None +// out["awkward_ListArray_getitem_jagged_shrink_b", {dtype_specializations}] = None +// END PYTHON + +template +__global__ void +awkward_ListArray_getitem_jagged_shrink_a( + T* tocarry, + C* tosmalloffsets, + U* tolargeoffsets, + const V* slicestarts, + const W* slicestops, + int64_t length, + const X* missing, + int64_t* scan_in_array_k, + int64_t* scan_in_array_tosmalloffsets, + int64_t* scan_in_array_tolargeoffsets, + 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 < length) { + if (thread_id == 0) { + scan_in_array_tosmalloffsets[0] = slicestarts[0]; + scan_in_array_tolargeoffsets[0] = slicestarts[0]; + } + V slicestart = slicestarts[thread_id]; + W slicestop = slicestops[thread_id]; + if (slicestart != slicestop) { + C smallcount = 0; + for (int64_t j = slicestart + threadIdx.y; j < slicestop; j += blockDim.y) { + if (missing[j] >= 0) { + smallcount++; + } + } + scan_in_array_k[thread_id + 1] = smallcount; + scan_in_array_tosmalloffsets[thread_id + 1] = smallcount; + } + scan_in_array_tolargeoffsets[thread_id + 1] = slicestop - slicestart; + } + } +} + +template +__global__ void +awkward_ListArray_getitem_jagged_shrink_b( + T* tocarry, + C* tosmalloffsets, + U* tolargeoffsets, + const V* slicestarts, + const W* slicestops, + int64_t length, + const X* missing, + int64_t* scan_in_array_k, + int64_t* scan_in_array_tosmalloffsets, + int64_t* scan_in_array_tolargeoffsets, + 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 (length == 0) { + tosmalloffsets[0] = 0; + tolargeoffsets[0] = 0; + } + else { + tosmalloffsets[0] = slicestarts[0]; + tolargeoffsets[0] = slicestarts[0]; + } + if (thread_id < length) { + V slicestart = slicestarts[thread_id]; + W slicestop = slicestops[thread_id]; + int64_t k = scan_in_array_k[thread_id] - scan_in_array_k[0]; + if (slicestart != slicestop) { + for (int64_t j = slicestart + threadIdx.y; j < slicestop; j += blockDim.y) { + if (missing[j] >= 0) { + tocarry[k] = j; + k++; + } + } + tosmalloffsets[thread_id + 1] = scan_in_array_tosmalloffsets[thread_id + 1]; + } + else { + tosmalloffsets[thread_id + 1] = scan_in_array_tosmalloffsets[thread_id]; + } + tolargeoffsets[thread_id + 1] = scan_in_array_tolargeoffsets[thread_id + 1]; + } + } +} diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_next_range_carrylength.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_next_range_carrylength.cu index dde19b6e8b..a69aecb732 100644 --- a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_next_range_carrylength.cu +++ b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_next_range_carrylength.cu @@ -26,8 +26,6 @@ awkward_ListArray_getitem_next_range_carrylength_a( uint64_t* err_code) { if (err_code[0] == NO_ERROR) { int64_t thread_id = blockIdx.x * blockDim.x + threadIdx.x; - const int64_t kMaxInt64 = 9223372036854775806; // 2**63 - 2: see below - const int64_t kSliceNone = kMaxInt64 + 1; // for Slice::none() if (thread_id < lenstarts) { int64_t length = fromstops[thread_id] - fromstarts[thread_id]; int64_t regular_start = start; From 3d11b6eeafc1dfdf41e87fa90c2f5449c6c93229 Mon Sep 17 00:00:00 2001 From: ManasviGoyal Date: Tue, 4 Jun 2024 16:23:03 +0200 Subject: [PATCH 4/7] test: cuda integration tests --- ...est_3140_cuda_jagged_and_masked_getitem.py | 1255 +++++++++++++++++ 1 file changed, 1255 insertions(+) create mode 100644 tests-cuda/test_3140_cuda_jagged_and_masked_getitem.py diff --git a/tests-cuda/test_3140_cuda_jagged_and_masked_getitem.py b/tests-cuda/test_3140_cuda_jagged_and_masked_getitem.py new file mode 100644 index 0000000000..884c6a6498 --- /dev/null +++ b/tests-cuda/test_3140_cuda_jagged_and_masked_getitem.py @@ -0,0 +1,1255 @@ +from __future__ import annotations + +import cupy as cp +import numpy as np +import pytest + +import awkward as ak + +to_list = ak.operations.to_list + + +def test_0111_jagged_and_masked_getitem_bitmaskedarray2b(): + array = ak.operations.from_iter( + [[0.0, 1.1, 2.2], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]], highlevel=False + ) + + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + + mask = ak.index.IndexU8(cp.array([66], dtype=cp.uint8)) + maskedarray = ak.contents.BitMaskedArray( + mask, cuda_array, valid_when=False, length=4, lsb_order=True + ) # lsb_order is irrelevant in this example + cuda_array = ak.highlevel.Array([[1, 1], None, [0], [0, -1]], backend="cuda").layout + + assert to_list(maskedarray) == [ + [0.0, 1.1, 2.2], + None, + [5.5], + [6.6, 7.7, 8.8, 9.9], + ] + assert to_list(maskedarray[cuda_array]) == [ + [1.1, 1.1], + None, + [5.5], + [6.6, 9.9], + ] + assert maskedarray.to_typetracer()[cuda_array].form == maskedarray[cuda_array].form + + +def test_0111_jagged_and_masked_getitem_bytemaskedarray2b(): + array = ak.operations.from_iter( + [[0.0, 1.1, 2.2], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]], highlevel=False + ) + + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + + mask = ak.index.Index8(cp.array([0, 1, 0, 0], dtype=cp.int8)) + maskedarray = ak.contents.ByteMaskedArray(mask, cuda_array, valid_when=False) + + cuda_array = ak.highlevel.Array([[1, 1], None, [0], [0, -1]], backend="cuda").layout + + assert to_list(maskedarray) == [ + [0.0, 1.1, 2.2], + None, + [5.5], + [6.6, 7.7, 8.8, 9.9], + ] + assert to_list(maskedarray[cuda_array]) == [ + [1.1, 1.1], + None, + [5.5], + [6.6, 9.9], + ] + assert maskedarray.to_typetracer()[cuda_array].form == maskedarray[cuda_array].form + + +def test_0111_jagged_and_masked_getitem_emptyarray(): + content = ak.contents.EmptyArray() + offsets = ak.index.Index64(np.array([0, 0, 0, 0, 0], dtype=np.int64)) + listoffsetarray = ak.contents.ListOffsetArray(offsets, content) + + cuda_listoffsetarray = ak.to_backend(listoffsetarray, "cuda", highlevel=False) + + array1 = ak.highlevel.Array([[], [], [], []], check_valid=True).layout + cuda_array1 = ak.to_backend(array1, "cuda") + + array2 = ak.highlevel.Array([[], [None], [], []], check_valid=True).layout + cuda_array2 = ak.to_backend(array2, "cuda") + + array3 = ak.highlevel.Array([[], [], None, []], check_valid=True).layout + cuda_array3 = ak.to_backend(array3, "cuda") + + array4 = ak.highlevel.Array([[], [None], None, []], check_valid=True).layout + cuda_array4 = ak.to_backend(array4, "cuda") + + array5 = ak.highlevel.Array([[], [0], [], []], check_valid=True).layout + cuda_array5 = ak.to_backend(array5, "cuda") + + assert to_list(cuda_listoffsetarray) == [[], [], [], []] + + assert to_list(cuda_listoffsetarray[cuda_array1]) == [[], [], [], []] + assert ( + cuda_listoffsetarray.to_typetracer()[cuda_array1].form + == cuda_listoffsetarray[cuda_array1].form + ) + + assert to_list(cuda_listoffsetarray[cuda_array2]) == [[], [None], [], []] + assert ( + cuda_listoffsetarray.to_typetracer()[cuda_array2].form + == cuda_listoffsetarray[cuda_array2].form + ) + assert to_list(cuda_listoffsetarray[cuda_array3]) == [[], [], None, []] + assert ( + cuda_listoffsetarray.to_typetracer()[cuda_array3].form + == cuda_listoffsetarray[cuda_array3].form + ) + assert to_list(cuda_listoffsetarray[cuda_array4]) == [[], [None], None, []] + assert ( + cuda_listoffsetarray.to_typetracer()[cuda_array4].form + == cuda_listoffsetarray[cuda_array4].form + ) + + with pytest.raises(IndexError): + cuda_listoffsetarray[cuda_array5] + + +def test_0111_jagged_and_masked_getitem_indexedarray(): + array = ak.operations.from_iter( + [[0.0, 1.1, 2.2], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]], highlevel=False + ) + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + + index = ak.index.Index64(cp.array([3, 2, 1, 0], dtype=cp.int64)) + cuda_indexedarray = ak.contents.IndexedArray(index, cuda_array) + + assert to_list(cuda_indexedarray) == [ + [6.6, 7.7, 8.8, 9.9], + [5.5], + [3.3, 4.4], + [0.0, 1.1, 2.2], + ] + + array1 = ak.highlevel.Array([[0, -1], [0], [], [1, 1]], check_valid=True).layout + cuda_array1 = ak.to_backend(array1, "cuda") + + assert to_list(cuda_indexedarray[cuda_array1]) == [ + [6.6, 9.9], + [5.5], + [], + [1.1, 1.1], + ] + assert ( + cuda_indexedarray.to_typetracer()[cuda_array1].form + == cuda_indexedarray[cuda_array1].form + ) + + array1 = ak.highlevel.Array( + [[0, -1], [0], [None], [1, None, 1]], check_valid=True + ).layout + cuda_array1 = ak.to_backend(array1, "cuda") + + assert to_list(cuda_indexedarray[cuda_array1]) == [ + [6.6, 9.9], + [5.5], + [None], + [1.1, None, 1.1], + ] + assert ( + cuda_indexedarray.to_typetracer()[cuda_array1].form + == cuda_indexedarray[cuda_array1].form + ) + + array1 = ak.highlevel.Array([[0, -1], [0], None, [1, 1]], check_valid=True).layout + cuda_array1 = ak.to_backend(array1, "cuda") + + assert to_list(cuda_indexedarray[cuda_array1]) == [ + [6.6, 9.9], + [5.5], + None, + [1.1, 1.1], + ] + assert ( + cuda_indexedarray.to_typetracer()[cuda_array1].form + == cuda_indexedarray[cuda_array1].form + ) + + array1 = ak.highlevel.Array([[0, -1], [0], None, [None]], check_valid=True).layout + cuda_array1 = ak.to_backend(array1, "cuda") + + assert to_list(cuda_indexedarray[cuda_array1]) == [[6.6, 9.9], [5.5], None, [None]] + assert ( + cuda_indexedarray.to_typetracer()[cuda_array1].form + == cuda_indexedarray[cuda_array1].form + ) + + index = ak.index.Index64(cp.array([3, 2, 1, 0], dtype=cp.int64)) + cuda_indexedarray = ak.contents.IndexedOptionArray(index, cuda_array) + + assert to_list(cuda_indexedarray) == [ + [6.6, 7.7, 8.8, 9.9], + [5.5], + [3.3, 4.4], + [0.0, 1.1, 2.2], + ] + + array1 = ak.highlevel.Array([[0, -1], [0], [], [1, 1]], check_valid=True) + cuda_array1 = ak.to_backend(array1, "cuda") + + assert to_list(cuda_indexedarray[cuda_array1]) == [ + [6.6, 9.9], + [5.5], + [], + [1.1, 1.1], + ] + assert ( + cuda_indexedarray.to_typetracer()[cuda_array1].form + == cuda_indexedarray[cuda_array1].form + ) + + array1 = ak.highlevel.Array( + [[0, -1], [0], [None], [1, None, 1]], check_valid=True + ).layout + cuda_array1 = ak.to_backend(array1, "cuda") + + assert to_list(cuda_indexedarray[cuda_array1]) == [ + [6.6, 9.9], + [5.5], + [None], + [1.1, None, 1.1], + ] + assert ( + cuda_indexedarray.to_typetracer()[cuda_array1].form + == cuda_indexedarray[cuda_array1].form + ) + + array1 = ak.highlevel.Array([[0, -1], [0], None, []], check_valid=True).layout + cuda_array1 = ak.to_backend(array1, "cuda") + + assert to_list(cuda_indexedarray[cuda_array1]) == [[6.6, 9.9], [5.5], None, []] + assert ( + cuda_indexedarray.to_typetracer()[array1].form + == cuda_indexedarray[cuda_array1].form + ) + + array1 = ak.highlevel.Array( + [[0, -1], [0], None, [1, None, 1]], check_valid=True + ).layout + cuda_array1 = ak.to_backend(array1, "cuda") + + assert to_list(cuda_indexedarray[cuda_array1]) == [ + [6.6, 9.9], + [5.5], + None, + [1.1, None, 1.1], + ] + assert ( + cuda_indexedarray.to_typetracer()[cuda_array1].form + == cuda_indexedarray[cuda_array1].form + ) + + +def test_0111_jagged_and_masked_getitem_indexedarray2(): + array = ak.operations.from_iter( + [[0.0, 1.1, 2.2], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]], highlevel=False + ) + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + + index = ak.index.Index64(cp.array([3, 2, -1, 0], dtype=cp.int64)) + cuda_indexedarray = ak.contents.IndexedOptionArray(index, cuda_array) + cuda_array = ak.highlevel.Array([[0, -1], [0], None, [1, 1]], backend="cuda").layout + + assert to_list(cuda_indexedarray) == [ + [6.6, 7.7, 8.8, 9.9], + [5.5], + None, + [0.0, 1.1, 2.2], + ] + assert to_list(cuda_indexedarray[cuda_array]) == [ + [6.6, 9.9], + [5.5], + None, + [1.1, 1.1], + ] + assert ( + cuda_indexedarray.to_typetracer()[cuda_array].form + == cuda_indexedarray[cuda_array].form + ) + + +def test_0111_jagged_and_masked_getitem_indexedarray2b(): + array = ak.operations.from_iter( + [[0.0, 1.1, 2.2], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]], highlevel=False + ) + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + + index = ak.index.Index64(cp.array([0, -1, 2, 3], dtype=cp.int64)) + cuda_indexedarray = ak.contents.IndexedOptionArray(index, cuda_array) + cuda_array = ak.highlevel.Array([[1, 1], None, [0], [0, -1]], backend="cuda").layout + + assert to_list(cuda_indexedarray) == [ + [0.0, 1.1, 2.2], + None, + [5.5], + [6.6, 7.7, 8.8, 9.9], + ] + assert to_list(cuda_indexedarray[cuda_array]) == [ + [1.1, 1.1], + None, + [5.5], + [6.6, 9.9], + ] + assert ( + cuda_indexedarray.to_typetracer()[cuda_array].form + == cuda_indexedarray[cuda_array].form + ) + + +def test_0111_jagged_and_masked_getitem_indexedarray3(): + array = ak.highlevel.Array([0.0, 1.1, 2.2, None, 4.4, None, None, 7.7]).layout + cuda_array = ak.to_backend(array, "cuda") + + assert to_list(cuda_array[ak.highlevel.Array([4, 3, 2], backend="cuda")]) == [ + 4.4, + None, + 2.2, + ] + assert to_list( + cuda_array[ak.highlevel.Array([4, 3, 2, None, 1], backend="cuda")] + ) == [ + 4.4, + None, + 2.2, + None, + 1.1, + ] + + array = ak.highlevel.Array([[0.0, 1.1, None, 2.2], [3.3, None, 4.4], [5.5]]).layout + array2 = ak.highlevel.Array([[3, 2, 2, 1], [1, 2], []]).layout + + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + cuda_array2 = ak.to_backend(array2, "cuda") + + assert to_list(cuda_array[cuda_array2]) == [ + [2.2, None, None, 1.1], + [None, 4.4], + [], + ] + assert cuda_array.to_typetracer()[cuda_array2].form == cuda_array[cuda_array2].form + + cuda_array = ak.highlevel.Array( + [[0.0, 1.1, 2.2], [3.3, 4.4], None, [5.5]], backend="cuda" + ).layout + cuda_array2 = ak.highlevel.Array([3, 2, 1], backend="cuda").layout + cuda_array3 = ak.highlevel.Array([3, 2, 1, None, 0], backend="cuda").layout + cuda_array4 = ak.highlevel.Array( + [[2, 1, 1, 0], [1], None, [0]], backend="cuda" + ).layout + cuda_array5 = ak.highlevel.Array( + [[2, 1, 1, 0], None, [1], [0]], backend="cuda" + ).layout + cuda_array6 = ak.highlevel.Array( + [[2, 1, 1, 0], None, [1], [0], None], backend="cuda" + ).layout + + assert to_list(cuda_array[cuda_array2]) == [[5.5], None, [3.3, 4.4]] + assert cuda_array.to_typetracer()[cuda_array2].form == cuda_array[cuda_array2].form + assert to_list(cuda_array[cuda_array3]) == [ + [5.5], + None, + [3.3, 4.4], + None, + [0.0, 1.1, 2.2], + ] + assert cuda_array.to_typetracer()[cuda_array3].form == cuda_array[cuda_array3].form + + assert (to_list(cuda_array[cuda_array4])) == [ + [2.2, 1.1, 1.1, 0.0], + [4.4], + None, + [5.5], + ] + assert cuda_array.to_typetracer()[cuda_array4].form == cuda_array[cuda_array4].form + + assert to_list(cuda_array[cuda_array5]) == [ + [2.2, 1.1, 1.1, 0], + None, + None, + [5.5], + ] + assert cuda_array.to_typetracer()[cuda_array5].form == cuda_array[cuda_array5].form + with pytest.raises(IndexError): + cuda_array[cuda_array6] + + +def test_0111_jagged_and_masked_getitem_jagged(): + array = ak.highlevel.Array( + [[1.1, 2.2, 3.3], [], [4.4, 5.5], [6.6], [7.7, 8.8, 9.9]], check_valid=True + ).layout + array2 = ak.highlevel.Array( + [[0, -1], [], [-1, 0], [-1], [1, 1, -2, 0]], check_valid=True + ).layout + + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + cuda_array2 = ak.to_backend(array2, "cuda") + + assert to_list(cuda_array[cuda_array2]) == [ + [1.1, 3.3], + [], + [5.5, 4.4], + [6.6], + [8.8, 8.8, 8.8, 7.7], + ] + assert cuda_array.to_typetracer()[cuda_array2].form == cuda_array[cuda_array2].form + + +def test_0111_jagged_and_masked_getitem_double_jagged(): + array = ak.highlevel.Array( + [[[0, 1, 2, 3], [4, 5]], [[6, 7, 8], [9, 10, 11, 12, 13]]], check_valid=True + ).layout + array2 = ak.highlevel.Array( + [[[2, 1, 0], [-1]], [[-1, -2, -3], [2, 1, 1, 3]]], check_valid=True + ).layout + + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + cuda_array2 = ak.to_backend(array2, "cuda") + + assert to_list(cuda_array[cuda_array2]) == [ + [[2, 1, 0], [5]], + [[8, 7, 6], [11, 10, 10, 12]], + ] + assert cuda_array.to_typetracer()[cuda_array2].form == cuda_array[cuda_array2].form + + content = ak.operations.from_iter( + [[0, 1, 2, 3], [4, 5], [6, 7, 8], [9, 10, 11, 12, 13]], highlevel=False + ) + regulararray = ak.contents.RegularArray(content, 2, zeros_length=0) + cuda_regulararray = ak.to_backend(regulararray, "cuda", highlevel=False) + + array1 = ak.highlevel.Array([[2, 1, 0], [-1]], check_valid=True).layout + cuda_array1 = ak.to_backend(array1, "cuda") + + assert to_list(cuda_regulararray[:, cuda_array1]) == [ + [[2, 1, 0], [5]], + [[8, 7, 6], [13]], + ] + assert ( + cuda_regulararray.to_typetracer()[:, cuda_array1].form + == cuda_regulararray[:, cuda_array1].form + ) + assert to_list(cuda_regulararray[1:, cuda_array1]) == [[[8, 7, 6], [13]]] + assert ( + cuda_regulararray.to_typetracer()[1:, cuda_array1].form + == cuda_regulararray[1:, cuda_array1].form + ) + + offsets = ak.index.Index64(np.array([0, 2, 4], dtype=np.int64)) + listoffsetarray = ak.contents.ListOffsetArray(offsets, content) + cuda_listoffsetarray = ak.to_backend(listoffsetarray, "cuda", highlevel=False) + + assert to_list(cuda_listoffsetarray[:, cuda_array1]) == [ + [[2, 1, 0], [5]], + [[8, 7, 6], [13]], + ] + assert ( + cuda_listoffsetarray.to_typetracer()[:, cuda_array1].form + == cuda_listoffsetarray[:, cuda_array1].form + ) + assert to_list(cuda_listoffsetarray[1:, cuda_array1]) == [[[8, 7, 6], [13]]] + assert ( + cuda_listoffsetarray.to_typetracer()[1:, cuda_array1].form + == cuda_listoffsetarray[1:, cuda_array1].form + ) + + +def test_0111_jagged_and_masked_getitem_masked_jagged(): + array = ak.highlevel.Array( + [[1.1, 2.2, 3.3], [], [4.4, 5.5], [6.6], [7.7, 8.8, 9.9]], check_valid=True + ).layout + array1 = ak.highlevel.Array( + [[-1, -2], None, [], None, [-2, 0]], check_valid=True + ).layout + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + cuda_array1 = ak.to_backend(array1, "cuda") + + assert to_list(cuda_array[cuda_array1]) == [[3.3, 2.2], None, [], None, [8.8, 7.7]] + assert cuda_array.to_typetracer()[cuda_array1].form == cuda_array[cuda_array1].form + + +def test_0111_jagged_and_masked_getitem_jagged_masked(): + array = ak.highlevel.Array( + [[1.1, 2.2, 3.3], [], [4.4, 5.5], [6.6], [7.7, 8.8, 9.9]], check_valid=True + ).layout + array1 = ak.highlevel.Array( + [[-1, None], [], [None, 0], [None], [1]], check_valid=True + ).layout + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + cuda_array1 = ak.to_backend(array1, "cuda") + + assert to_list(cuda_array[cuda_array1]) == [ + [3.3, None], + [], + [None, 4.4], + [None], + [8.8], + ] + assert cuda_array.to_typetracer()[cuda_array1].form == cuda_array[cuda_array1].form + + +def test_0111_jagged_and_masked_getitem_array_boolean_to_int(): + a = ak.operations.from_iter( + [[True, True, True], [], [True, True], [True], [True, True, True, True]], + highlevel=False, + ) + cuda_a = ak.to_backend(a, "cuda", highlevel=False) + b = ak._slicing._normalise_item_bool_to_int(cuda_a, backend=cuda_a.backend) + assert to_list(b) == [[0, 1, 2], [], [0, 1], [0], [0, 1, 2, 3]] + + a = ak.operations.from_iter( + [ + [True, True, False], + [], + [True, False], + [False], + [True, True, True, False], + ], + highlevel=False, + ) + cuda_a = ak.to_backend(a, "cuda", highlevel=False) + b = ak._slicing._normalise_item_bool_to_int(cuda_a, backend=cuda_a.backend) + assert to_list(b) == [[0, 1], [], [0], [], [0, 1, 2]] + + a = ak.operations.from_iter( + [ + [False, True, True], + [], + [False, True], + [False], + [False, True, True, True], + ], + highlevel=False, + ) + cuda_a = ak.to_backend(a, "cuda", highlevel=False) + b = ak._slicing._normalise_item_bool_to_int(cuda_a, backend=cuda_a.backend) + assert to_list(b) == [[1, 2], [], [1], [], [1, 2, 3]] + + # a = ak.operations.from_iter( + # [[True, True, None], [], [True, None], [None], [True, True, True, None]], + # highlevel=False, + # ) + # cuda_a = ak.to_backend(a, "cuda", highlevel=False) + # # b = ak._slicing._normalise_item_bool_to_int(cuda_a, backend=cuda_a.backend) + # # error in _slicing line 553 - FIXME + # assert to_list(b) == [[0, 1, None], [], [0, None], [None], [0, 1, 2, None]] + # assert ( + # b.content.index.data[b.content.index.data >= 0].tolist() + # == np.arange(6).tolist() # kernels expect nonnegative entries to be arange + # ) + + # a = ak.operations.from_iter( + # [[None, True, True], [], [None, True], [None], [None, True, True, True]], + # highlevel=False, + # ) + # cuda_a = ak.to_backend(a, "cuda", highlevel=False) + # b = ak._slicing._normalise_item_bool_to_int(cuda_a, backend=cuda_a.backend) + # assert to_list(b) == [[None, 1, 2], [], [None, 1], [None], [None, 1, 2, 3]] + # assert ( + # b.content.index.data[b.content.index.data >= 0].tolist() + # == np.arange(6).tolist() # kernels expect nonnegative entries to be arange + # ) + + # a = ak.operations.from_iter( + # [[False, True, None], [], [False, None], [None], [False, True, True, None]], + # highlevel=False, + # ) + # cuda_a = ak.to_backend(a, "cuda", highlevel=False) + # b = ak._slicing._normalise_item_bool_to_int(cuda_a, backend=cuda_a.backend) + # assert to_list(b) == [[1, None], [], [None], [None], [1, 2, None]] + # assert ( + # b.content.index.data[b.content.index.data >= 0].tolist() + # == np.arange(3).tolist() # kernels expect nonnegative entries to be arange + # ) + + # a = ak.operations.from_iter( + # [[None, True, False], [], [None, False], [None], [None, True, True, False]], + # highlevel=False, + # ) + # cuda_a = ak.to_backend(a, "cuda", highlevel=False) + # b = ak._slicing._normalise_item_bool_to_int(cuda_a, backend=cuda_a.backend) + # assert to_list(b) == [[None, 1], [], [None], [None], [None, 1, 2]] + # assert ( + # b.content.index.data[b.content.index.data >= 0].tolist() + # == np.arange(3).tolist() # kernels expect nonnegative entries to be arange + # ) + + +def test_0111_jagged_and_masked_getitem_array_slice(): + array = ak.highlevel.Array( + [0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9], check_valid=True + ).layout + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + + assert to_list(cuda_array[[5, 2, 2, 3, 9, 0, 1]]) == [ + 5.5, + 2.2, + 2.2, + 3.3, + 9.9, + 0.0, + 1.1, + ] + assert ( + cuda_array.to_typetracer()[[5, 2, 2, 3, 9, 0, 1]].form + == cuda_array[[5, 2, 2, 3, 9, 0, 1]].form + ) + assert to_list(cuda_array[cp.array([5, 2, 2, 3, 9, 0, 1])]) == [ + 5.5, + 2.2, + 2.2, + 3.3, + 9.9, + 0.0, + 1.1, + ] + assert ( + cuda_array.to_typetracer()[cp.array([5, 2, 2, 3, 9, 0, 1])].form + == cuda_array[cp.array([5, 2, 2, 3, 9, 0, 1])].form + ) + + array2 = ak.contents.NumpyArray(np.array([5, 2, 2, 3, 9, 0, 1], dtype=np.int32)) + cuda_array2 = ak.to_backend(array2, "cuda") + + assert to_list(cuda_array[cuda_array2]) == [5.5, 2.2, 2.2, 3.3, 9.9, 0.0, 1.1] + assert cuda_array.to_typetracer()[cuda_array2].form == cuda_array[cuda_array2].form + assert to_list( + cuda_array[ + ak.highlevel.Array( + cp.array([5, 2, 2, 3, 9, 0, 1], dtype=cp.int32), check_valid=True + ) + ] + ) == [5.5, 2.2, 2.2, 3.3, 9.9, 0.0, 1.1] + assert ( + cuda_array.to_typetracer()[ + ak.highlevel.Array( + cp.array([5, 2, 2, 3, 9, 0, 1], dtype=cp.int32), check_valid=True + ) + ].form + == cuda_array[ + ak.highlevel.Array( + cp.array([5, 2, 2, 3, 9, 0, 1], dtype=cp.int32), check_valid=True + ) + ].form + ) + assert to_list( + cuda_array[ + ak.highlevel.Array(cp.array([5, 2, 2, 3, 9, 0, 1]), check_valid=True) + ] + ) == [ + 5.5, + 2.2, + 2.2, + 3.3, + 9.9, + 0.0, + 1.1, + ] + assert ( + cuda_array.to_typetracer()[ + cp.array(ak.highlevel.Array([5, 2, 2, 3, 9, 0, 1])) + ].form + == cuda_array[cp.array(ak.highlevel.Array([5, 2, 2, 3, 9, 0, 1]))].form + ) + + array3 = ak.contents.NumpyArray( + np.array([False, False, False, False, False, True, False, True, False, True]) + ) + cuda_array3 = ak.to_backend(array3, "cuda") + + assert to_list(cuda_array[cuda_array3]) == [5.5, 7.7, 9.9] + assert cuda_array.to_typetracer()[cuda_array3].form == cuda_array[cuda_array3].form + + content = ak.contents.NumpyArray(np.array([1, 0, 9, 3, 2, 2, 5], dtype=np.int64)) + index = ak.index.Index64(np.array([6, 5, 4, 3, 2, 1, 0], dtype=np.int64)) + indexedarray = ak.contents.IndexedArray(index, content) + cuda_indexedarray = ak.to_backend(indexedarray, "cuda") + + assert to_list(cuda_array[cuda_indexedarray]) == [5.5, 2.2, 2.2, 3.3, 9.9, 0.0, 1.1] + assert ( + cuda_array.to_typetracer()[cuda_indexedarray].form + == cuda_array[cuda_indexedarray].form + ) + assert to_list( + cuda_array[cp.array(ak.highlevel.Array(indexedarray, check_valid=True))] + ) == [ + 5.5, + 2.2, + 2.2, + 3.3, + 9.9, + 0.0, + 1.1, + ] + assert ( + cuda_array.to_typetracer()[ak.highlevel.Array(cuda_indexedarray)].form + == cuda_array[ak.highlevel.Array(cuda_indexedarray)].form + ) + + emptyarray = ak.contents.EmptyArray() + cuda_emptyarray = ak.to_backend(emptyarray, "cuda") + + assert to_list(cuda_array[cuda_emptyarray]) == [] + assert ( + cuda_array.to_typetracer()[cuda_emptyarray].form + == cuda_array[cuda_emptyarray].form + ) + + array = ak.highlevel.Array( + np.array([[0.0, 1.1, 2.2, 3.3, 4.4], [5.5, 6.6, 7.7, 8.8, 9.9]]), + check_valid=True, + ).layout + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + + numpyarray1 = ak.contents.NumpyArray(np.array([[0, 1], [1, 0]])) + numpyarray2 = ak.contents.NumpyArray(np.array([[2, 4], [3, 3]])) + + cuda_numpyarray1 = ak.to_backend(numpyarray1, "cuda") + cuda_numpyarray2 = ak.to_backend(numpyarray2, "cuda") + + assert to_list( + cuda_array[ + cuda_numpyarray1, + cuda_numpyarray2, + ] + ) == [[2.2, 9.9], [8.8, 3.3]] + assert ( + cuda_array.to_typetracer()[ + cuda_numpyarray1, + cuda_numpyarray2, + ].form + == cuda_array[ + cuda_numpyarray1, + cuda_numpyarray2, + ].form + ) + assert to_list(cuda_array[cuda_numpyarray1]) == [ + [[0.0, 1.1, 2.2, 3.3, 4.4], [5.5, 6.6, 7.7, 8.8, 9.9]], + [[5.5, 6.6, 7.7, 8.8, 9.9], [0.0, 1.1, 2.2, 3.3, 4.4]], + ] + assert ( + cuda_array.to_typetracer()[cuda_numpyarray1].form + == cuda_array[cuda_numpyarray1].form + ) + + +def test_0111_jagged_and_masked_getitem_array_slice_2(): + array = ak.highlevel.Array( + [0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9], check_valid=True + ).layout + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + + content0 = ak.contents.NumpyArray(np.array([5, 2, 2])) + content1 = ak.contents.NumpyArray(np.array([3, 9, 0, 1])) + tags = ak.index.Index8(np.array([0, 0, 0, 1, 1, 1, 1], dtype=np.int8)) + index2 = ak.index.Index64(np.array([0, 1, 2, 0, 1, 2, 3], dtype=np.int64)) + unionarray = ak.contents.UnionArray.simplified(tags, index2, [content0, content1]) + cuda_unionarray = ak.to_backend(unionarray, "cuda") + + assert to_list(cuda_array[cuda_unionarray]) == [5.5, 2.2, 2.2, 3.3, 9.9, 0.0, 1.1] + assert ( + cuda_array.to_typetracer()[cuda_unionarray].form + == cuda_array[cuda_unionarray].form + ) + + +def test_0111_jagged_and_masked_getitem_array_slice_with_union(): + array = ak.highlevel.Array( + [0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9], check_valid=True + ).layout + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + + content0 = ak.contents.NumpyArray(np.array([5, 2, 2])) + content1 = ak.contents.NumpyArray(np.array([3, 9, 0, 1])) + tags = ak.index.Index8(np.array([0, 0, 0, 1, 1, 1, 1], dtype=np.int8)) + index2 = ak.index.Index64(np.array([0, 1, 2, 0, 1, 2, 3], dtype=np.int64)) + unionarray = ak.contents.UnionArray.simplified(tags, index2, [content0, content1]) + cuda_unionarray = ak.to_backend(unionarray, "cuda") + + assert to_list( + cuda_array[ak.highlevel.Array(cuda_unionarray, check_valid=True)] + ) == [ + 5.5, + 2.2, + 2.2, + 3.3, + 9.9, + 0.0, + 1.1, + ] + assert ( + cuda_array.to_typetracer()[ak.highlevel.Array(cuda_unionarray)].form + == cuda_array[ak.highlevel.Array(cuda_unionarray)].form + ) + + +def test_0111_jagged_and_masked_getitem_jagged_mask(): + array = ak.highlevel.Array( + [[1.1, 2.2, 3.3], [], [4.4, 5.5], [6.6], [7.7, 8.8, 9.9]], check_valid=True + ).layout + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + + assert to_list( + cuda_array[[[True, True, True], [], [True, True], [True], [True, True, True]]] + ) == [[1.1, 2.2, 3.3], [], [4.4, 5.5], [6.6], [7.7, 8.8, 9.9]] + assert ( + cuda_array.to_typetracer()[ + [[True, True, True], [], [True, True], [True], [True, True, True]] + ].form + == cuda_array[ + [[True, True, True], [], [True, True], [True], [True, True, True]] + ].form + ) + assert to_list( + cuda_array[[[False, True, True], [], [True, True], [True], [True, True, True]]] + ) == [[2.2, 3.3], [], [4.4, 5.5], [6.6], [7.7, 8.8, 9.9]] + assert ( + cuda_array.to_typetracer()[ + [[False, True, True], [], [True, True], [True], [True, True, True]] + ].form + == cuda_array[ + [[False, True, True], [], [True, True], [True], [True, True, True]] + ].form + ) + assert to_list( + cuda_array[[[True, False, True], [], [True, True], [True], [True, True, True]]] + ) == [[1.1, 3.3], [], [4.4, 5.5], [6.6], [7.7, 8.8, 9.9]] + assert ( + cuda_array.to_typetracer()[ + [[True, False, True], [], [True, True], [True], [True, True, True]] + ].form + == cuda_array[ + [[True, False, True], [], [True, True], [True], [True, True, True]] + ].form + ) + assert to_list( + cuda_array[[[True, True, True], [], [False, True], [True], [True, True, True]]] + ) == [[1.1, 2.2, 3.3], [], [5.5], [6.6], [7.7, 8.8, 9.9]] + assert ( + cuda_array.to_typetracer()[ + [[True, True, True], [], [False, True], [True], [True, True, True]] + ].form + == cuda_array[ + [[True, True, True], [], [False, True], [True], [True, True, True]] + ].form + ) + assert to_list( + cuda_array[[[True, True, True], [], [False, False], [True], [True, True, True]]] + ) == [[1.1, 2.2, 3.3], [], [], [6.6], [7.7, 8.8, 9.9]] + assert ( + cuda_array.to_typetracer()[ + [[True, True, True], [], [False, False], [True], [True, True, True]] + ].form + == cuda_array[ + [[True, True, True], [], [False, False], [True], [True, True, True]] + ].form + ) + + +def test_0111_jagged_and_masked_getitem_jagged_missing_mask(): + array = ak.highlevel.Array( + [[1.1, 2.2, 3.3], [], [4.4, 5.5]], check_valid=True + ).layout + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + + assert to_list(cuda_array[[[True, True, True], [], [True, True]]]) == [ + [1.1, 2.2, 3.3], + [], + [4.4, 5.5], + ] + assert ( + cuda_array.to_typetracer()[[[True, True, True], [], [True, True]]].form + == cuda_array[[[True, True, True], [], [True, True]]].form + ) + assert to_list(cuda_array[[[True, False, True], [], [True, True]]]) == [ + [1.1, 3.3], + [], + [4.4, 5.5], + ] + assert ( + cuda_array.to_typetracer()[[[True, False, True], [], [True, True]]].form + == cuda_array[[[True, False, True], [], [True, True]]].form + ) + assert to_list(cuda_array[[[True, True, False], [], [False, None]]]) == [ + [1.1, 2.2], + [], + [None], + ] + assert ( + cuda_array.to_typetracer()[[[True, True, False], [], [False, None]]].form + == cuda_array[[[True, True, False], [], [False, None]]].form + ) + assert to_list(cuda_array[[[True, True, False], [], [True, None]]]) == [ + [1.1, 2.2], + [], + [4.4, None], + ] + assert ( + cuda_array.to_typetracer()[[[True, True, False], [], [True, None]]].form + == cuda_array[[[True, True, False], [], [True, None]]].form + ) + + assert to_list(cuda_array[[[True, None, True], [], [True, True]]]) == [ + [1.1, None, 3.3], + [], + [4.4, 5.5], + ] + assert ( + cuda_array.to_typetracer()[[[True, None, True], [], [True, True]]].form + == cuda_array[[[True, None, True], [], [True, True]]].form + ) + assert to_list(cuda_array[[[True, None, False], [], [True, True]]]) == [ + [1.1, None], + [], + [4.4, 5.5], + ] + assert ( + cuda_array.to_typetracer()[[[True, None, False], [], [True, True]]].form + == cuda_array[[[True, None, False], [], [True, True]]].form + ) + + assert to_list(cuda_array[[[False, None, False], [], [True, True]]]) == [ + [None], + [], + [4.4, 5.5], + ] + assert ( + cuda_array.to_typetracer()[[[False, None, False], [], [True, True]]].form + == cuda_array[[[False, None, False], [], [True, True]]].form + ) + assert to_list(cuda_array[[[True, True, False], [], [False, True]]]) == [ + [1.1, 2.2], + [], + [5.5], + ] + assert ( + cuda_array.to_typetracer()[[[True, True, False], [], [False, True]]].form + == cuda_array[[[True, True, False], [], [False, True]]].form + ) + assert to_list(cuda_array[[[True, True, None], [], [False, True]]]) == [ + [1.1, 2.2, None], + [], + [5.5], + ] + assert ( + cuda_array.to_typetracer()[[[True, True, None], [], [False, True]]].form + == cuda_array[[[True, True, None], [], [False, True]]].form + ) + assert to_list(cuda_array[[[True, True, False], [None], [False, True]]]) == [ + [1.1, 2.2], + [None], + [5.5], + ] + assert ( + cuda_array.to_typetracer()[[[True, True, False], [None], [False, True]]].form + == cuda_array[[[True, True, False], [None], [False, True]]].form + ) + + +def test_0111_jagged_and_masked_getitem_masked_of_jagged_of_whatever(): + content = ak.contents.NumpyArray(np.arange(2 * 3 * 5)) + regulararray1 = ak.contents.RegularArray(content, 5, zeros_length=0) + regulararray2 = ak.contents.RegularArray(regulararray1, 3, zeros_length=0) + cuda_regulararray2 = ak.to_backend(regulararray2, "cuda", highlevel=False) + + array1 = ak.highlevel.Array( + [[[2], None, [-1, 2, 0]], [[-3], None, [-5, -3, 4]]], check_valid=True + ).layout + array2 = ak.highlevel.Array( + [[[2], None, [-1, None, 0]], [[-3], None, [-5, None, 4]]], + check_valid=True, + ).layout + + cuda_array1 = ak.to_backend(array1, backend="cuda") + cuda_array2 = ak.to_backend(array2, backend="cuda") + + assert to_list(cuda_regulararray2[cuda_array1]) == [ + [[2], None, [14, 12, 10]], + [[17], None, [25, 27, 29]], + ] + assert ( + cuda_regulararray2.to_typetracer()[cuda_array1].form + == cuda_regulararray2[cuda_array1].form + ) + + assert to_list(cuda_regulararray2[cuda_array2]) == [ + [[2], None, [14, None, 10]], + [[17], None, [25, None, 29]], + ] + assert ( + cuda_regulararray2.to_typetracer()[cuda_array2].form + == cuda_regulararray2[cuda_array2].form + ) + + +def test_0111_jagged_and_masked_getitem_missing(): + array = ak.highlevel.Array( + [0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9], check_valid=True + ).layout + array2 = ak.highlevel.Array([3, 6, None, None, -2, 6], check_valid=True).layout + + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + cuda_array2 = ak.to_backend(array2, "cuda") + + assert to_list(cuda_array[cuda_array2]) == [ + 3.3, + 6.6, + None, + None, + 8.8, + 6.6, + ] + assert cuda_array.to_typetracer()[cuda_array2].form == cuda_array[cuda_array2].form + + content = ak.contents.NumpyArray( + np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9, 10.0, 11.1, 999]) + ) + regulararray = ak.contents.RegularArray(content, 4, zeros_length=0) + cuda_regulararray = ak.to_backend(regulararray, "cuda", highlevel=False) + + assert to_list(cuda_regulararray) == [ + [0.0, 1.1, 2.2, 3.3], + [4.4, 5.5, 6.6, 7.7], + [8.8, 9.9, 10.0, 11.1], + ] + array3 = ak.highlevel.Array([2, 1, 1, None, -1], check_valid=True).layout + cuda_array3 = ak.to_backend(array3, "cuda") + + assert to_list(cuda_regulararray[cuda_array3]) == [ + [8.8, 9.9, 10.0, 11.1], + [4.4, 5.5, 6.6, 7.7], + [4.4, 5.5, 6.6, 7.7], + None, + [8.8, 9.9, 10.0, 11.1], + ] + assert ( + cuda_regulararray.to_typetracer()[cuda_array3].form + == cuda_regulararray[cuda_array3].form + ) + assert to_list(cuda_regulararray[:, cuda_array3]) == [ + [2.2, 1.1, 1.1, None, 3.3], + [6.6, 5.5, 5.5, None, 7.7], + [10.0, 9.9, 9.9, None, 11.1], + ] + assert ( + cuda_regulararray.to_typetracer()[:, cuda_array3].form + == cuda_regulararray[:, cuda_array3].form + ) + assert to_list(cuda_regulararray[1:, cuda_array3]) == [ + [6.6, 5.5, 5.5, None, 7.7], + [10.0, 9.9, 9.9, None, 11.1], + ] + assert ( + cuda_regulararray.to_typetracer()[1:, cuda_array3].form + == cuda_regulararray[1:, cuda_array3].form + ) + + maskedarray = np.ma.MaskedArray( + [2, 1, 1, 999, -1], [False, False, False, True, False] + ) + cuda_maskedarray = ak.to_backend(maskedarray, backend="cuda") + + assert to_list(cuda_regulararray[cuda_maskedarray]) == [ + [8.8, 9.9, 10.0, 11.1], + [4.4, 5.5, 6.6, 7.7], + [4.4, 5.5, 6.6, 7.7], + None, + [8.8, 9.9, 10.0, 11.1], + ] + assert ( + cuda_regulararray.to_typetracer()[cuda_maskedarray].form + == cuda_regulararray[cuda_maskedarray].form + ) + assert to_list( + cuda_regulararray[ + :, + cuda_maskedarray, + ] + ) == [ + [2.2, 1.1, 1.1, None, 3.3], + [6.6, 5.5, 5.5, None, 7.7], + [10.0, 9.9, 9.9, None, 11.1], + ] + assert ( + cuda_regulararray.to_typetracer()[ + :, + cuda_maskedarray, + ].form + == cuda_regulararray[ + :, + cuda_maskedarray, + ].form + ) + + assert to_list( + cuda_regulararray[ + 1:, + cuda_maskedarray, + ] + ) == [[6.6, 5.5, 5.5, None, 7.7], [10.0, 9.9, 9.9, None, 11.1]] + assert ( + cuda_regulararray.to_typetracer()[ + 1:, + cuda_maskedarray, + ].form + == cuda_regulararray[ + 1:, + cuda_maskedarray, + ].form + ) + + content = ak.contents.NumpyArray( + np.array([[0.0, 1.1, 2.2, 3.3], [4.4, 5.5, 6.6, 7.7], [8.8, 9.9, 10.0, 11.1]]) + ) + cuda_content = ak.to_backend(content, "cuda", highlevel=False) + + assert to_list(cuda_content[cuda_array3]) == [ + [8.8, 9.9, 10.0, 11.1], + [4.4, 5.5, 6.6, 7.7], + [4.4, 5.5, 6.6, 7.7], + None, + [8.8, 9.9, 10.0, 11.1], + ] + assert ( + cuda_content.to_typetracer()[cuda_array3].form == cuda_content[cuda_array3].form + ) + assert to_list(cuda_content[:, cuda_array3]) == [ + [2.2, 1.1, 1.1, None, 3.3], + [6.6, 5.5, 5.5, None, 7.7], + [10.0, 9.9, 9.9, None, 11.1], + ] + assert ( + cuda_content.to_typetracer()[:, cuda_array3].form + == cuda_content[:, cuda_array3].form + ) + assert to_list(cuda_content[1:, cuda_array3]) == [ + [6.6, 5.5, 5.5, None, 7.7], + [10.0, 9.9, 9.9, None, 11.1], + ] + assert ( + cuda_content.to_typetracer()[1:, cuda_array3].form + == cuda_content[1:, cuda_array3].form + ) + + assert to_list(cuda_content[cuda_maskedarray]) == [ + [8.8, 9.9, 10.0, 11.1], + [4.4, 5.5, 6.6, 7.7], + [4.4, 5.5, 6.6, 7.7], + None, + [8.8, 9.9, 10.0, 11.1], + ] + assert ( + cuda_content.to_typetracer()[cuda_maskedarray].form + == cuda_content[cuda_maskedarray].form + ) + assert to_list( + cuda_content[ + :, + cuda_maskedarray, + ] + ) == [ + [2.2, 1.1, 1.1, None, 3.3], + [6.6, 5.5, 5.5, None, 7.7], + [10.0, 9.9, 9.9, None, 11.1], + ] + assert ( + cuda_content.to_typetracer()[ + :, + cuda_maskedarray, + ].form + == cuda_content[ + :, + cuda_maskedarray, + ].form + ) + assert to_list( + cuda_content[ + 1:, + cuda_maskedarray, + ] + ) == [[6.6, 5.5, 5.5, None, 7.7], [10.0, 9.9, 9.9, None, 11.1]] + assert ( + cuda_content.to_typetracer()[ + 1:, + cuda_maskedarray, + ].form + == cuda_content[ + 1:, + cuda_maskedarray, + ].form + ) + + content = ak.contents.NumpyArray( + np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9, 10.0, 11.1, 999]) + ) + offsets = ak.index.Index64(np.array([0, 4, 8, 12], dtype=np.int64)) + listoffsetarray = ak.contents.ListOffsetArray(offsets, content) + cuda_listoffsetarray = ak.to_backend(listoffsetarray, "cuda", highlevel=False) + + assert to_list(cuda_listoffsetarray) == [ + [0.0, 1.1, 2.2, 3.3], + [4.4, 5.5, 6.6, 7.7], + [8.8, 9.9, 10.0, 11.1], + ] + assert to_list(cuda_listoffsetarray[:, cuda_array3]) == [ + [2.2, 1.1, 1.1, None, 3.3], + [6.6, 5.5, 5.5, None, 7.7], + [10.0, 9.9, 9.9, None, 11.1], + ] + assert ( + cuda_listoffsetarray.to_typetracer()[:, cuda_array3].form + == cuda_listoffsetarray[:, cuda_array3].form + ) + assert to_list(cuda_listoffsetarray[1:, cuda_array3]) == [ + [6.6, 5.5, 5.5, None, 7.7], + [10.0, 9.9, 9.9, None, 11.1], + ] + assert ( + cuda_listoffsetarray.to_typetracer()[1:, cuda_array3].form + == cuda_listoffsetarray[1:, cuda_array3].form + ) + + assert to_list( + cuda_listoffsetarray[ + :, + cuda_maskedarray, + ] + ) == [ + [2.2, 1.1, 1.1, None, 3.3], + [6.6, 5.5, 5.5, None, 7.7], + [10.0, 9.9, 9.9, None, 11.1], + ] + assert ( + cuda_listoffsetarray.to_typetracer()[ + :, + cuda_maskedarray, + ].form + == cuda_listoffsetarray[ + :, + cuda_maskedarray, + ].form + ) + assert to_list( + cuda_listoffsetarray[ + 1:, + cuda_maskedarray, + ] + ) == [[6.6, 5.5, 5.5, None, 7.7], [10.0, 9.9, 9.9, None, 11.1]] + assert ( + cuda_listoffsetarray.to_typetracer()[ + 1:, + cuda_maskedarray, + ].form + == cuda_listoffsetarray[ + 1:, + cuda_maskedarray, + ].form + ) From 0e8d0bfe5f93d057ec8dc4b6f45bd9f4f616c131 Mon Sep 17 00:00:00 2001 From: ManasviGoyal Date: Wed, 5 Jun 2024 09:43:45 +0200 Subject: [PATCH 5/7] test: more slicing integration tests --- ...est_3140_cuda_jagged_and_masked_getitem.py | 276 +++++++++++ tests-cuda/test_3140_cuda_slicing.py | 464 ++++++++++++++++++ 2 files changed, 740 insertions(+) create mode 100644 tests-cuda/test_3140_cuda_slicing.py diff --git a/tests-cuda/test_3140_cuda_jagged_and_masked_getitem.py b/tests-cuda/test_3140_cuda_jagged_and_masked_getitem.py index 884c6a6498..fff5417c03 100644 --- a/tests-cuda/test_3140_cuda_jagged_and_masked_getitem.py +++ b/tests-cuda/test_3140_cuda_jagged_and_masked_getitem.py @@ -1253,3 +1253,279 @@ def test_0111_jagged_and_masked_getitem_missing(): cuda_maskedarray, ].form ) + + +def test_0111_jagged_and_masked_getitem_new_slices(): + content = ak.contents.NumpyArray(np.array([1, 0, 9, 3, 2, 2, 5], dtype=np.int64)) + index = ak.index.Index64(np.array([6, 5, -1, 3, 2, -1, 0], dtype=np.int64)) + indexedarray = ak.contents.IndexedOptionArray(index, content) + cuda_indexedarray = ak.to_backend(indexedarray, backend="cuda") + + assert to_list(cuda_indexedarray) == [5, 2, None, 3, 9, None, 1] + + offsets = ak.index.Index64(np.array([0, 4, 4, 7], dtype=np.int64)) + listoffsetarray = ak.contents.ListOffsetArray(offsets, content) + cuda_listoffsetarray = ak.to_backend(listoffsetarray, backend="cuda") + + assert to_list(cuda_listoffsetarray) == [[1, 0, 9, 3], [], [2, 2, 5]] + + offsets = ak.index.Index64(np.array([1, 4, 4, 6], dtype=np.int64)) + listoffsetarray = ak.contents.ListOffsetArray(offsets, content) + cuda_listoffsetarray = ak.to_backend(listoffsetarray, backend="cuda") + + assert to_list(cuda_listoffsetarray) == [[0, 9, 3], [], [2, 2]] + + starts = ak.index.Index64(np.array([1, 99, 5], dtype=np.int64)) + stops = ak.index.Index64(np.array([4, 99, 7], dtype=np.int64)) + listarray = ak.contents.ListArray(starts, stops, content) + cuda_listarray = ak.to_backend(listarray, backend="cuda") + + assert to_list(cuda_listarray) == [[0, 9, 3], [], [2, 5]] + + +def test_0111_jagged_and_masked_getitem_record(): + array = ak.highlevel.Array( + [ + {"x": [0, 1, 2], "y": [0.0, 1.1, 2.2, 3.3]}, + {"x": [3, 4, 5, 6], "y": [4.4, 5.5]}, + {"x": [7, 8], "y": [6.6, 7.7, 8.8, 9.9]}, + ], + check_valid=True, + ).layout + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + + array2 = ak.highlevel.Array([[-1, 1], [0, 0, 1], [-1, -2]], check_valid=True).layout + array3 = ak.highlevel.Array( + [[-1, 1], [0, 0, None, 1], [-1, -2]], check_valid=True + ).layout + array4 = ak.highlevel.Array([[-1, 1], None, [-1, -2]], check_valid=True).layout + cuda_array2 = ak.to_backend(array2, "cuda") + cuda_array3 = ak.to_backend(array3, "cuda") + cuda_array4 = ak.to_backend(array4, "cuda") + + assert to_list(cuda_array[cuda_array2]) == [ + {"x": [2, 1], "y": [3.3, 1.1]}, + {"x": [3, 3, 4], "y": [4.4, 4.4, 5.5]}, + {"x": [8, 7], "y": [9.9, 8.8]}, + ] + assert cuda_array.to_typetracer()[cuda_array2].form == cuda_array[cuda_array2].form + assert to_list(cuda_array[cuda_array3]) == [ + {"x": [2, 1], "y": [3.3, 1.1]}, + {"x": [3, 3, None, 4], "y": [4.4, 4.4, None, 5.5]}, + {"x": [8, 7], "y": [9.9, 8.8]}, + ] + assert cuda_array.to_typetracer()[cuda_array3].form == cuda_array[cuda_array3].form + assert to_list(cuda_array[cuda_array4]) == [ + {"x": [2, 1], "y": [3.3, 1.1]}, + None, + {"x": [8, 7], "y": [9.9, 8.8]}, + ] + assert cuda_array.to_typetracer()[cuda_array4].form == cuda_array[cuda_array4].form + + +def test_0111_jagged_and_masked_getitem_records_missing(): + array = ak.highlevel.Array( + [ + {"x": 0, "y": 0.0}, + {"x": 1, "y": 1.1}, + {"x": 2, "y": 2.2}, + {"x": 3, "y": 3.3}, + {"x": 4, "y": 4.4}, + {"x": 5, "y": 5.5}, + {"x": 6, "y": 6.6}, + {"x": 7, "y": 7.7}, + {"x": 8, "y": 8.8}, + {"x": 9, "y": 9.9}, + ], + check_valid=True, + ).layout + array2 = ak.highlevel.Array([3, 1, None, 1, 7], check_valid=True).layout + + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + cuda_array2 = ak.to_backend(array2, "cuda") + + assert to_list(cuda_array[cuda_array2]) == [ + {"x": 3, "y": 3.3}, + {"x": 1, "y": 1.1}, + None, + {"x": 1, "y": 1.1}, + {"x": 7, "y": 7.7}, + ] + assert cuda_array.to_typetracer()[cuda_array2].form == cuda_array[cuda_array2].form + + array = ak.highlevel.Array( + [ + [ + {"x": 0, "y": 0.0}, + {"x": 1, "y": 1.1}, + {"x": 2, "y": 2.2}, + {"x": 3, "y": 3.3}, + ], + [ + {"x": 4, "y": 4.4}, + {"x": 5, "y": 5.5}, + {"x": 6, "y": 6.6}, + {"x": 7, "y": 7.7}, + {"x": 8, "y": 8.8}, + {"x": 9, "y": 9.9}, + ], + ], + check_valid=True, + ).layout + array2 = ak.highlevel.Array([1, None, 2, -1], check_valid=True).layout + + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + cuda_array2 = ak.to_backend(array2, "cuda") + + assert to_list(cuda_array[:, cuda_array2]) == [ + [{"x": 1, "y": 1.1}, None, {"x": 2, "y": 2.2}, {"x": 3, "y": 3.3}], + [{"x": 5, "y": 5.5}, None, {"x": 6, "y": 6.6}, {"x": 9, "y": 9.9}], + ] + assert ( + cuda_array.to_typetracer()[:, cuda_array2].form + == cuda_array[:, cuda_array2].form + ) + + array = ak.highlevel.Array( + [ + {"x": [0, 1, 2, 3], "y": [0.0, 1.1, 2.2, 3.3]}, + {"x": [4, 5, 6, 7], "y": [4.4, 5.5, 6.6, 7.7]}, + {"x": [8, 9, 10, 11], "y": [8.8, 9.9, 10.0, 11.1]}, + ], + check_valid=True, + ).layout + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + + assert to_list(cuda_array[:, cuda_array2]) == [ + {"x": [1, None, 2, 3], "y": [1.1, None, 2.2, 3.3]}, + {"x": [5, None, 6, 7], "y": [5.5, None, 6.6, 7.7]}, + {"x": [9, None, 10, 11], "y": [9.9, None, 10.0, 11.1]}, + ] + assert ( + cuda_array.to_typetracer()[:, cuda_array2].form + == cuda_array[:, cuda_array2].form + ) + assert to_list(cuda_array[1:, cuda_array2]) == [ + {"x": [5, None, 6, 7], "y": [5.5, None, 6.6, 7.7]}, + {"x": [9, None, 10, 11], "y": [9.9, None, 10.0, 11.1]}, + ] + assert ( + cuda_array.to_typetracer()[1:, cuda_array2].form + == cuda_array[1:, cuda_array2].form + ) + + +def test_0111_jagged_and_masked_getitem_regular_regular(): + content = ak.contents.NumpyArray(np.arange(2 * 3 * 5)) + regulararray1 = ak.contents.RegularArray(content, 5, zeros_length=0) + regulararray2 = ak.contents.RegularArray(regulararray1, 3, zeros_length=0) + + cuda_regulararray2 = ak.to_backend(regulararray2, "cuda", highlevel=False) + + array1 = ak.highlevel.Array( + [[[2], [1, -2], [-1, 2, 0]], [[-3], [-4, 3], [-5, -3, 4]]], + check_valid=True, + ).layout + array2 = ak.highlevel.Array( + [[[2], [1, -2], [-1, None, 0]], [[-3], [-4, 3], [-5, None, 4]]], + check_valid=True, + ).layout + + cuda_array1 = ak.to_backend(array1, "cuda") + cuda_array2 = ak.to_backend(array2, "cuda") + + assert to_list(cuda_regulararray2[cuda_array1]) == [ + [[2], [6, 8], [14, 12, 10]], + [[17], [21, 23], [25, 27, 29]], + ] + assert ( + cuda_regulararray2.to_typetracer()[cuda_array1].form + == cuda_regulararray2[cuda_array1].form + ) + + assert to_list(cuda_regulararray2[cuda_array2]) == [ + [[2], [6, 8], [14, None, 10]], + [[17], [21, 23], [25, None, 29]], + ] + assert ( + cuda_regulararray2.to_typetracer()[cuda_array2].form + == cuda_regulararray2[cuda_array2].form + ) + + +def test_0111_jagged_and_masked_getitem_sequential(): + array = ak.highlevel.Array( + np.arange(2 * 3 * 5).reshape(2, 3, 5).tolist(), check_valid=True + ).layout + array2 = ak.highlevel.Array([[2, 1, 0], [2, 1, 0]], check_valid=True).layout + + cuda_array = ak.to_backend(array, "cuda", highlevel=False) + cuda_array2 = ak.to_backend(array2, "cuda") + + assert to_list(cuda_array[cuda_array2]) == [ + [[10, 11, 12, 13, 14], [5, 6, 7, 8, 9], [0, 1, 2, 3, 4]], + [[25, 26, 27, 28, 29], [20, 21, 22, 23, 24], [15, 16, 17, 18, 19]], + ] + assert cuda_array.to_typetracer()[cuda_array2].form == cuda_array[cuda_array2].form + assert to_list(cuda_array[cuda_array2, :2]) == [ + [[10, 11], [5, 6], [0, 1]], + [[25, 26], [20, 21], [15, 16]], + ] + assert ( + cuda_array.to_typetracer()[cuda_array2, :2].form + == cuda_array[cuda_array2, :2].form + ) + + +def test_0111_jagged_and_masked_getitem_union(): + one = ak.operations.from_iter( + [["1.1", "2.2", "3.3"], [], ["4.4", "5.5"]], highlevel=False + ) + two = ak.operations.from_iter( + [[6.6], [7.7, 8.8], [], [9.9, 10.0, 11.1, 12.2]], highlevel=False + ) + tags = ak.index.Index8(np.array([0, 0, 0, 1, 1, 1, 1], dtype=np.int8)) + index = ak.index.Index64(np.array([0, 1, 2, 0, 1, 2, 3], dtype=np.int64)) + unionarray = ak.contents.UnionArray(tags, index, [one, two]) + + cuda_unionarray = ak.to_backend(unionarray, "cuda") + + assert to_list(cuda_unionarray) == [ + ["1.1", "2.2", "3.3"], + [], + ["4.4", "5.5"], + [6.6], + [7.7, 8.8], + [], + [9.9, 10.0, 11.1, 12.2], + ] + + +def test_0111_jagged_and_masked_getitem_union_2(): + one = ak.operations.from_iter([[1.1, 2.2, 3.3], [], [4.4, 5.5]], highlevel=False) + two = ak.operations.from_iter( + [[6.6], [7.7, 8.8], [], [9.9, 10.0, 11.1, 12.2]], highlevel=False + ) + tags = ak.index.Index8(np.array([0, 0, 0, 1, 1, 1, 1], dtype=np.int8)) + index = ak.index.Index64(np.array([0, 1, 2, 0, 1, 2, 3], dtype=np.int64)) + unionarray = ak.contents.UnionArray.simplified(tags, index, [one, two]) + array = ak.highlevel.Array( + [[0, -1], [], [1, 1], [], [-1], [], [1, -2, -1]], check_valid=True + ).layout + + cuda_unionarray = ak.to_backend(unionarray, "cuda", highlevel=False) + cuda_array = ak.to_backend(array, "cuda") + + assert to_list(cuda_unionarray[cuda_array]) == [ + [1.1, 3.3], + [], + [5.5, 5.5], + [], + [8.8], + [], + [10.0, 11.1, 12.2], + ] + assert ( + cuda_unionarray.to_typetracer()[cuda_array].form + == cuda_unionarray[cuda_array].form + ) diff --git a/tests-cuda/test_3140_cuda_slicing.py b/tests-cuda/test_3140_cuda_slicing.py new file mode 100644 index 0000000000..3a5c2233bf --- /dev/null +++ b/tests-cuda/test_3140_cuda_slicing.py @@ -0,0 +1,464 @@ +from __future__ import annotations + +import numpy as np + +import awkward as ak + +to_list = ak.operations.to_list + + +def test_0315_integerindex_null_more(): + f = ak.highlevel.Array([[0, None, 2], None, [3, 4], []], backend="cuda").layout + g1 = ak.highlevel.Array([[1, 2, None], None, [], [None]], backend="cuda").layout + g2 = ak.highlevel.Array([[], None, None, []], backend="cuda").layout + g3 = ak.highlevel.Array([[], [], [], []], backend="cuda").layout + + assert to_list(f[g1]) == [[None, 2, None], None, [], [None]] + assert to_list(f[g2]) == [[], None, None, []] + assert to_list(f[g3]) == [[], None, [], []] + assert f.to_typetracer()[g1].form == f[g1].form + assert f.to_typetracer()[g2].form == f[g2].form + assert f.to_typetracer()[g3].form == f[g3].form + + a = ak.highlevel.Array([[0, 1, 2, None], None], backend="cuda").layout + b = ak.highlevel.Array([[2, 1, None, 3], None], backend="cuda").layout + + assert to_list(a[b]) == [[2, 1, None, None], None] + assert a.to_typetracer()[b].form == a[b].form + + b = ak.highlevel.Array([[2, 1, None, 3], []], backend="cuda").layout + + assert to_list(a[b]) == [[2, 1, None, None], None] + assert a.to_typetracer()[b].form == a[b].form + + b = ak.highlevel.Array([[2, 1, None, 3], [0, 1]], backend="cuda").layout + assert to_list(a[b]) == [[2, 1, None, None], None] + assert a.to_typetracer()[b].form == a[b].form + + +def test_0315_integerindex_null_more_2(): + a = ak.highlevel.Array( + [[[0, 1, 2, None], None], [[3, 4], [5]], None, [[6]]], backend="cuda" + ).layout + b = ak.highlevel.Array( + [[[2, 1, None, 3], [0, 1]], [[0], None], None, [None]], backend="cuda" + ).layout + c = ak.highlevel.Array( + [ + [[False, True, None, False], [False, True]], + [[True, False], None], + None, + [None], + ], + backend="cuda", + ).layout + + assert to_list(a[b]) == [ + [[2, 1, None, None], None], + [[3], None], + None, + [None], + ] + assert to_list(a[c]) == [[[1, None], None], [[4], None], None, [None]] + assert a.to_typetracer()[c].form == a[c].form + + +def test_1405_slicing_untested_cases_list_option_list(): + """Check that non-offset list(option(list indexes correctly""" + content = ak.contents.ListOffsetArray( + ak.index.Index64(np.array([0, 2], dtype=np.int64)), + ak.contents.ListOffsetArray( + ak.index.Index64(np.array([2, 2, 3], dtype=np.int64)), + ak.contents.NumpyArray(np.array([2, 2, 2], dtype=np.int64)), + ), + ) + + index = ak.contents.ListOffsetArray( + ak.index.Index64(np.array([0, 2], dtype=np.int64)), + ak.contents.IndexedOptionArray( + ak.index.Index64(np.array([0, 1], dtype=np.int64)), + ak.contents.ListOffsetArray( + ak.index.Index64(np.array([0, 0, 1], dtype=np.int64)), + ak.contents.NumpyArray(np.array([0], dtype=np.int64)), + ), + ), + ) + + cuda_content = ak.to_backend(content, "cuda") + cuda_index = ak.to_backend(index, "cuda") + + assert cuda_content[cuda_index].to_list() == [[[], [2]]] + + +def test_1405_slicing_untested_cases_list_option_list_offset(): + """Check that offset list(option(list indexes correctly""" + content = ak.contents.ListOffsetArray( + ak.index.Index64(np.array([1, 3], dtype=np.int64)), + ak.contents.ListOffsetArray( + ak.index.Index64(np.array([0, 2, 2, 3], dtype=np.int64)), + ak.contents.NumpyArray(np.array([2, 2, 2], dtype=np.int64)), + ), + ) + + index = ak.contents.ListOffsetArray( + ak.index.Index64(np.array([0, 2], dtype=np.int64)), + ak.contents.IndexedOptionArray( + ak.index.Index64(np.array([0, 1], dtype=np.int64)), + ak.contents.ListOffsetArray( + ak.index.Index64(np.array([0, 0, 1], dtype=np.int64)), + ak.contents.NumpyArray(np.array([0], dtype=np.int64)), + ), + ), + ) + + cuda_content = ak.to_backend(content, "cuda") + cuda_index = ak.to_backend(index, "cuda") + + assert cuda_content[cuda_index].to_list() == [[[], [2]]] + + +def test_1502_getitem_jagged_issue1406(): + array = ak.Array( + ak.contents.ListOffsetArray( + ak.index.Index64(np.array([1, 3], dtype=np.int64)), + ak.contents.ListOffsetArray( + ak.index.Index64(np.array([0, 2, 2, 3], dtype=np.int64)), + ak.contents.NumpyArray(np.array([0, 1, 2], dtype=np.int64)), + ), + ), + check_valid=True, + ) + + index = ak.Array( + ak.contents.ListOffsetArray( + ak.index.Index64(np.array([0, 2], dtype=np.int64)), + ak.contents.IndexedOptionArray( + ak.index.Index64(np.array([0, 1], dtype=np.int64)), + ak.contents.ListOffsetArray( + ak.index.Index64(np.array([0, 0, 1], dtype=np.int64)), + ak.contents.NumpyArray(np.array([0], dtype=np.int64)), + ), + ), + ), + check_valid=True, + ) + + cuda_array = ak.to_backend(array, "cuda") + cuda_index = ak.to_backend(index, "cuda") + + assert to_list(cuda_array[cuda_index]) == [[[], [2]]] + + +def test_1502_getitem_jagged_issue1406_success_start_offset0(): + array = ak.Array( + ak.contents.ListOffsetArray( + ak.index.Index64(np.array([0, 2], dtype=np.int64)), + ak.contents.ListOffsetArray( + ak.index.Index64(np.array([2, 2, 3], dtype=np.int64)), + ak.contents.NumpyArray(np.array([0, 1, 2], dtype=np.int64)), + ), + ), + check_valid=True, + ) + + index = ak.Array( + ak.contents.ListOffsetArray( + ak.index.Index64(np.array([0, 2], dtype=np.int64)), + ak.contents.IndexedOptionArray( + ak.index.Index64(np.array([0, 1], dtype=np.int64)), + ak.contents.ListOffsetArray( + ak.index.Index64(np.array([0, 0, 1], dtype=np.int64)), + ak.contents.NumpyArray(np.array([0], dtype=np.int64)), + ), + ), + ), + check_valid=True, + ) + + cuda_array = ak.to_backend(array, "cuda") + cuda_index = ak.to_backend(index, "cuda") + + assert to_list(cuda_array[cuda_index]) == [[[], [2]]] + + +def test_1502_getitem_jagged_issue1406_success_remove_option_type(): + array = ak.Array( + ak.contents.ListOffsetArray( + ak.index.Index64(np.array([1, 3], dtype=np.int64)), + ak.contents.ListOffsetArray( + ak.index.Index64(np.array([0, 2, 2, 3], dtype=np.int64)), + ak.contents.NumpyArray(np.array([0, 1, 2], dtype=np.int64)), + ), + ), + check_valid=True, + ) + + index = ak.Array( + ak.contents.ListOffsetArray( + ak.index.Index64(np.array([0, 2], dtype=np.int64)), + ak.contents.ListOffsetArray( + ak.index.Index64(np.array([0, 0, 1], dtype=np.int64)), + ak.contents.NumpyArray(np.array([0], dtype=np.int64)), + ), + ), + check_valid=True, + ) + + cuda_array = ak.to_backend(array, "cuda") + cuda_index = ak.to_backend(index, "cuda") + + assert to_list(cuda_array[cuda_index]) == [[[], [2]]] + + +def test_1502_getitem_jagged_issue1406_success_nonempty_list(): + array = ak.Array( + ak.contents.ListOffsetArray( + ak.index.Index64(np.array([1, 3], dtype=np.int64)), + ak.contents.ListOffsetArray( + ak.index.Index64(np.array([0, 1, 2, 3], dtype=np.int64)), + ak.contents.NumpyArray(np.array([0, 1, 2], dtype=np.int64)), + ), + ), + check_valid=True, + ) + + index = ak.Array( + ak.contents.ListOffsetArray( + ak.index.Index64(np.array([0, 2], dtype=np.int64)), + ak.contents.IndexedOptionArray( + ak.index.Index64(np.array([0, 1], dtype=np.int64)), + ak.contents.ListOffsetArray( + ak.index.Index64(np.array([0, 1, 2], dtype=np.int64)), + ak.contents.NumpyArray(np.array([0, 0], dtype=np.int64)), + ), + ), + ), + check_valid=True, + ) + + cuda_array = ak.to_backend(array, "cuda") + cuda_index = ak.to_backend(index, "cuda") + + assert to_list(cuda_array[cuda_index]) == [[[1], [2]]] + + +def test_1904_drop_none_ListArray_and_axis_None(): + index = ak.index.Index64(np.asarray([0, -1, 1, -1, 4, -1, 5])) + content = ak.contents.recordarray.RecordArray( + [ + ak.contents.numpyarray.NumpyArray( + np.array([6.6, 1.1, 2.2, 3.3, 4.4, 5.5, 7.7]) + ) + ], + ["nest"], + ) + indexoptionarray = ak.contents.IndexedOptionArray(index, content) + a = ak.contents.listarray.ListArray( + ak.index.Index(np.array([4, 100, 1], np.int64)), + ak.index.Index(np.array([7, 100, 3, 200], np.int64)), + indexoptionarray, + ) + + cuda_a = ak.to_backend(a, "cuda") + + assert ( + to_list(ak.drop_none(cuda_a)) + == to_list(cuda_a[~ak.is_none(cuda_a, axis=1)]) + == [ + [{"nest": 4.4}, {"nest": 5.5}], + [], + [{"nest": 1.1}], + ] + ) + + +def test_1904_drop_none_ListOffsetArray_IndexedOptionArray_NumpyArray_outoforder(): + index = ak.index.Index64(np.asarray([0, -1, 1, 5, 4, 2, 5])) + content = ak.contents.numpyarray.NumpyArray( + np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6]) + ) + indexoptionarray = ak.contents.IndexedOptionArray(index, content) + offsets = ak.index.Index64(np.asarray([0, 4, 5, 6])) + listoffset = ak.contents.ListOffsetArray(offsets, indexoptionarray) + + cuda_listoffset = ak.to_backend(listoffset, "cuda") + + assert to_list(cuda_listoffset) == [[0.0, None, 1.1, 5.5], [4.4], [2.2]] + assert ( + to_list(ak.drop_none(cuda_listoffset, axis=1)) + == to_list(cuda_listoffset[~ak.is_none(cuda_listoffset, axis=1)]) + == [[0.0, 1.1, 5.5], [4.4], [2.2]] + ) + assert to_list(ak.drop_none(cuda_listoffset)) == [[0.0, 1.1, 5.5], [4.4], [2.2]] + + +def test_1904_drop_none_from_iter(): + a = ak.Array([[1], [2, None]], backend="cuda") + assert to_list(ak.drop_none(a)) == [[1], [2]] + + a = ak.Array([[2, None]], backend="cuda") + assert to_list(ak.drop_none(a)) == [[2]] + + a = ak.Array([[[None]]], backend="cuda") + assert to_list(ak.drop_none(a)) == [[[]]] + + a = ak.Array([1, 2, None], backend="cuda") + assert to_list(ak.drop_none(a, axis=0)) + + a = ak.Array([[[1, None]], [[3, 4]], [[5, 6]], [[7.8]]], backend="cuda") + assert ( + to_list(ak.drop_none(a, axis=2)) + == to_list(a[~ak.is_none(a, axis=2)]) + == [[[1.0]], [[3.0, 4.0]], [[5.0, 6.0]], [[7.8]]] + ) + + a = ak.Array([[[0]], [[None]], [[1], None], [[2, None]]], backend="cuda") + assert ( + to_list(ak.drop_none(a, axis=1)) + == to_list(a[~ak.is_none(a, axis=1)]) + == [[[0]], [[None]], [[1]], [[2, None]]] + ) + + a = ak.Array( + [[[0]], [None, 34], [[1], None, 31], [[2, [[None]]]], [[[None]]]], + backend="cuda", + ) + assert ( + to_list(ak.drop_none(a, axis=0)) + == to_list(a[~ak.is_none(a, axis=0)]) + == [[[0]], [None, 34], [[1], None, 31], [[2, [[None]]]], [[[None]]]] + ) + + a = ak.Array([[[1, None]], [[3, None]], [[5, 6]], [[7.8]]], backend="cuda") + assert ( + to_list(ak.drop_none(a, axis=2)) + == to_list(a[~ak.is_none(a, axis=2)]) + == [[[1.0]], [[3.0]], [[5.0, 6.0]], [[7.8]]] + ) + + a = ak.Array([[[1, None]], [[None, 4]], [[5, 6]], [[7.8]]], backend="cuda") + assert ( + to_list(ak.drop_none(a, axis=2)) + == to_list(a[~ak.is_none(a, axis=2)]) + == [[[1.0]], [[4.0]], [[5.0, 6.0]], [[7.8]]] + ) + + a = ak.Array([[[1, None]], [[None, None]], [[5, 6]], [[7.8]]], backend="cuda") + assert ( + to_list(ak.drop_none(a, axis=2)) + == to_list(a[~ak.is_none(a, axis=2)]) + == [[[1.0]], [[]], [[5.0, 6.0]], [[7.8]]] + ) + + a = ak.Array([[[1, None]], [[None, None]], [[None, 6]], [[7.8]]], backend="cuda") + assert ( + to_list(ak.drop_none(a, axis=2)) + == to_list(a[~ak.is_none(a, axis=2)]) + == [[[1.0]], [[]], [[6.0]], [[7.8]]] + ) + + a = ak.Array( + [[{"x": [1], "y": [[2]]}], [{"x": [None], "y": [[None]]}], None], backend="cuda" + ) + assert to_list(a) == [ + [{"x": [1], "y": [[2]]}], + [{"x": [None], "y": [[None]]}], + None, + ] + assert to_list(ak.drop_none(a)) == [ + [{"x": [1], "y": [[2]]}], + [{"x": [], "y": [[]]}], + ] + assert ( + to_list(ak.drop_none(a, axis=0)) + == to_list(a[~ak.is_none(a, axis=0)]) + == [[{"x": [1], "y": [[2]]}], [{"x": [None], "y": [[None]]}]] + ) + assert to_list(ak.drop_none(a, axis=1)) == [ + [{"x": [1], "y": [[2]]}], + [{"x": [], "y": [[None]]}], + None, + ] + + +def test_1904_drop_none_List_ByteMaskedArray_NumpyArray(): + a = ak.contents.listarray.ListArray( + ak.index.Index(np.array([1, 3], np.int64)), + ak.index.Index(np.array([3, 4], np.int64)), + ak.contents.bytemaskedarray.ByteMaskedArray( + ak.index.Index(np.array([1, 0, 1, 0, 1], np.int8)), + ak.contents.numpyarray.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6])), + valid_when=True, + ), + ) + + cuda_a = ak.to_backend(a, "cuda") + + assert to_list(cuda_a) == [[None, 3.3], [None]] + assert to_list(ak.drop_none(cuda_a)) == [[3.3], []] + assert to_list(ak.drop_none(cuda_a, axis=1)) == to_list( + cuda_a[~ak.is_none(cuda_a, axis=1)] + ) + + +def test_1904_drop_none_ListOffsetArray_ByteMaskedArray_NumpyArray(): + a = ak.contents.listoffsetarray.ListOffsetArray( + ak.index.Index(np.array([1, 4, 4, 6, 7], np.int64)), + ak.contents.bytemaskedarray.ByteMaskedArray( + ak.index.Index(np.array([1, 0, 1, 0, 1], np.int8)), + ak.contents.numpyarray.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6])), + valid_when=True, + ), + ) + + cuda_a = ak.to_backend(a, "cuda") + + assert to_list(cuda_a) == [[None, 3.3, None], [], [5.5], []] + assert to_list(ak.drop_none(cuda_a, axis=1)) == [[3.3], [], [5.5], []] + + +def test_1904_drop_none_RegularArray_RecordArray_NumpyArray(): + index = ak.index.Index64(np.asarray([0, -1, 1, 2, 3, 4, -1, 6, 7, 8, -1, 10])) + content = ak.contents.numpyarray.NumpyArray( + np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9, 10.10]) + ) + indexoptionarray = ak.contents.IndexedOptionArray(index, content) + offsets = ak.index.Index64(np.array([0, 3, 3, 5, 6, 10, 10])) + listoffsetarray = ak.contents.listoffsetarray.ListOffsetArray( + offsets, indexoptionarray + ) + regulararray = ak.contents.regulararray.RegularArray(listoffsetarray, 2) + + cuda_regulararray = ak.to_backend(regulararray, "cuda") + + assert to_list(cuda_regulararray) == [ + [[0.0, None, 1.1], []], + [[2.2, 3.3], [4.4]], + [[None, 6.6, 7.7, 8.8], []], + ] + assert to_list(ak.drop_none(cuda_regulararray, axis=2)) == to_list( + cuda_regulararray[~ak.is_none(cuda_regulararray, axis=2)] + ) + + +def test_1904_drop_none_RecordArray(): + a = ak.Array( + [ + [{"x": [1], "y": [[2]]}], + None, + [None], + [{"x": None, "y": None}], + [{"x": [None], "y": [None]}], + [{"x": [11], "y": [[None]]}], + ], + backend="cuda", + ) + assert to_list(ak.drop_none(a, axis=1)) == to_list(a[~ak.is_none(a, axis=1)]) + assert to_list(ak.drop_none(a, axis=2)) == [ + [{"x": [1], "y": [[2]]}], + None, + [None], + [{"x": None, "y": None}], + [{"x": [], "y": []}], + [{"x": [11], "y": [[None]]}], + ] From bd969bc3b491b2a7a25193721be87165525f1cd4 Mon Sep 17 00:00:00 2001 From: ManasviGoyal Date: Wed, 5 Jun 2024 10:59:46 +0200 Subject: [PATCH 6/7] fix: ndarray error for cupy array shape --- src/awkward/_slicing.py | 2 +- tests-cuda/test_3140_cuda_slicing.py | 238 +++++++++++++++++++++++++-- 2 files changed, 223 insertions(+), 17 deletions(-) diff --git a/src/awkward/_slicing.py b/src/awkward/_slicing.py index dc9b0c0567..bd72e99b6a 100644 --- a/src/awkward/_slicing.py +++ b/src/awkward/_slicing.py @@ -550,7 +550,7 @@ def _normalise_item_bool_to_int(item: Content, backend: Backend) -> Content: # outindex fits into the lists; non-missing are sequential outindex = ak.index.Index64( - item_backend.index_nplike.full(nextoffsets.data[-1], -1, dtype=np.int64) + item_backend.index_nplike.full(nextoffsets[-1], -1, dtype=np.int64) ) outindex.data[~isnegative[expanded]] = item_backend.index_nplike.arange( nextcontent.shape[0], dtype=np.int64 diff --git a/tests-cuda/test_3140_cuda_slicing.py b/tests-cuda/test_3140_cuda_slicing.py index 3a5c2233bf..047fc7977c 100644 --- a/tests-cuda/test_3140_cuda_slicing.py +++ b/tests-cuda/test_3140_cuda_slicing.py @@ -401,22 +401,6 @@ def test_1904_drop_none_List_ByteMaskedArray_NumpyArray(): ) -def test_1904_drop_none_ListOffsetArray_ByteMaskedArray_NumpyArray(): - a = ak.contents.listoffsetarray.ListOffsetArray( - ak.index.Index(np.array([1, 4, 4, 6, 7], np.int64)), - ak.contents.bytemaskedarray.ByteMaskedArray( - ak.index.Index(np.array([1, 0, 1, 0, 1], np.int8)), - ak.contents.numpyarray.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6])), - valid_when=True, - ), - ) - - cuda_a = ak.to_backend(a, "cuda") - - assert to_list(cuda_a) == [[None, 3.3, None], [], [5.5], []] - assert to_list(ak.drop_none(cuda_a, axis=1)) == [[3.3], [], [5.5], []] - - def test_1904_drop_none_RegularArray_RecordArray_NumpyArray(): index = ak.index.Index64(np.asarray([0, -1, 1, 2, 3, 4, -1, 6, 7, 8, -1, 10])) content = ak.contents.numpyarray.NumpyArray( @@ -462,3 +446,225 @@ def test_1904_drop_none_RecordArray(): [{"x": [], "y": []}], [{"x": [11], "y": [[None]]}], ] + + +def test_2246_slice_not_packed(): + index = ak.Array( + ak.contents.ListOffsetArray( + ak.index.Index64([0, 3, 5]), + ak.contents.NumpyArray( + np.array([True, False, False, True, True, False, False], dtype=np.bool_) + ), + ) + ) + array = ak.Array([[0, 1, 2], [3, 4]]) + + cuda_index = ak.to_backend(index, "cuda") + cuda_array = ak.to_backend(array, "cuda") + + result = cuda_array[cuda_index] + assert result.tolist() == [[0], [3, 4]] + + +def test_0127_tomask_operation_ByteMaskedArray_jaggedslice0(): + array = ak.operations.from_iter( + [[0.0, 1.1, 2.2], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]], highlevel=False + ) + index = ak.index.Index64(np.array([0, 1, 2, 3], dtype=np.int64)) + indexedarray = ak.contents.IndexedOptionArray(index, array) + + cuda_indexedarray = ak.to_backend(indexedarray, "cuda") + + assert to_list(cuda_indexedarray) == [ + [0.0, 1.1, 2.2], + [3.3, 4.4], + [5.5], + [6.6, 7.7, 8.8, 9.9], + ] + assert to_list( + cuda_indexedarray[ + ak.highlevel.Array([[0, -1], [0], [], [1, 1]], backend="cuda") + ] + ) == [ + [0.0, 2.2], + [3.3], + [], + [7.7, 7.7], + ] + + mask = ak.index.Index8(np.array([0, 0, 0, 0], dtype=np.int8)) + maskedarray = ak.contents.ByteMaskedArray(mask, array, valid_when=False) + + cuda_maskedarray = ak.to_backend(maskedarray, "cuda") + + assert to_list(cuda_maskedarray) == [ + [0.0, 1.1, 2.2], + [3.3, 4.4], + [5.5], + [6.6, 7.7, 8.8, 9.9], + ] + assert to_list( + cuda_maskedarray[ak.highlevel.Array([[0, -1], [0], [], [1, 1]], backend="cuda")] + ) == [ + [0.0, 2.2], + [3.3], + [], + [7.7, 7.7], + ] + + +def test_0127_tomask_operation_ByteMaskedArray_jaggedslice1(): + model = ak.highlevel.Array( + [ + [0.0, 1.1, None, 2.2], + [], + [3.3, None, 4.4], + [5.5], + [6.6, 7.7, None, 8.8, 9.9], + ], + backend="cuda", + ) + assert to_list( + model[ + ak.highlevel.Array( + [[3, 2, 1, 1, 0], [], [1], [0, 0], [1, 2]], backend="cuda" + ) + ] + ) == [ + [2.2, None, 1.1, 1.1, 0.0], + [], + [None], + [5.5, 5.5], + [7.7, None], + ] + + content = ak.contents.NumpyArray( + np.array([0.0, 1.1, 999, 2.2, 3.3, 123, 4.4, 5.5, 6.6, 7.7, 321, 8.8, 9.9]) + ) + mask = ak.index.Index8( + np.array([0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0], dtype=np.int8) + ) + maskedarray = ak.contents.ByteMaskedArray(mask, content, valid_when=False) + offsets = ak.index.Index64(np.array([0, 4, 4, 7, 8, 13], dtype=np.int64)) + listarray = ak.highlevel.Array( + ak.contents.ListOffsetArray(offsets, maskedarray), backend="cuda" + ) + + assert to_list(listarray) == to_list(model) + assert to_list( + listarray[ + ak.highlevel.Array( + [[3, 2, 1, 1, 0], [], [1], [0, 0], [1, 2]], backend="cuda" + ) + ] + ) == [[2.2, None, 1.1, 1.1, 0.0], [], [None], [5.5, 5.5], [7.7, None]] + + +def test_0127_tomask_operation_ByteMaskedArray_jaggedslice2(): + model = ak.highlevel.Array( + [ + [[0.0, 1.1, None, 2.2], [], [3.3, None, 4.4]], + [], + [[5.5]], + [[6.6, 7.7, None, 8.8, 9.9]], + ], + backend="cuda", + ) + assert to_list( + model[ + ak.highlevel.Array( + [[[3, 2, 1, 1, 0], [], [1]], [], [[0, 0]], [[1, 2]]], backend="cuda" + ) + ] + ) == [[[2.2, None, 1.1, 1.1, 0.0], [], [None]], [], [[5.5, 5.5]], [[7.7, None]]] + + content = ak.contents.NumpyArray( + np.array([0.0, 1.1, 999, 2.2, 3.3, 123, 4.4, 5.5, 6.6, 7.7, 321, 8.8, 9.9]) + ) + mask = ak.index.Index8( + np.array([0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0], dtype=np.int8) + ) + maskedarray = ak.contents.ByteMaskedArray(mask, content, valid_when=False) + offsets = ak.index.Index64(np.array([0, 4, 4, 7, 8, 13], dtype=np.int64)) + sublistarray = ak.contents.ListOffsetArray(offsets, maskedarray) + offsets2 = ak.index.Index64(np.array([0, 3, 3, 4, 5], dtype=np.int64)) + listarray = ak.highlevel.Array( + ak.contents.ListOffsetArray(offsets2, sublistarray), backend="cuda" + ) + assert to_list(listarray) == to_list(model) + assert to_list( + listarray[ + ak.highlevel.Array( + [[[3, 2, 1, 1, 0], [], [1]], [], [[0, 0]], [[1, 2]]], backend="cuda" + ) + ] + ) == [[[2.2, None, 1.1, 1.1, 0.0], [], [None]], [], [[5.5, 5.5]], [[7.7, None]]] + + +def test_0127_tomask_operation_ByteMaskedArray_jaggedslice3(): + model = ak.highlevel.Array( + [ + [[[0.0, 1.1, None, 2.2], [], [3.3, None, 4.4]], []], + [[[5.5]], [[6.6, 7.7, None, 8.8, 9.9]]], + ], + backend="cuda", + ) + assert to_list( + model[ + ak.highlevel.Array( + [[[[3, 2, 1, 1, 0], [], [1]], []], [[[0, 0]], [[1, 2]]]], backend="cuda" + ) + ] + ) == [[[[2.2, None, 1.1, 1.1, 0.0], [], [None]], []], [[[5.5, 5.5]], [[7.7, None]]]] + + content = ak.contents.NumpyArray( + np.array([0.0, 1.1, 999, 2.2, 3.3, 123, 4.4, 5.5, 6.6, 7.7, 321, 8.8, 9.9]) + ) + mask = ak.index.Index8( + np.array([0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0], dtype=np.int8) + ) + maskedarray = ak.contents.ByteMaskedArray(mask, content, valid_when=False) + offsets = ak.index.Index64(np.array([0, 4, 4, 7, 8, 13], dtype=np.int64)) + subsublistarray = ak.contents.ListOffsetArray(offsets, maskedarray) + offsets2 = ak.index.Index64(np.array([0, 3, 3, 4, 5], dtype=np.int64)) + sublistarray = ak.contents.ListOffsetArray(offsets2, subsublistarray) + offsets3 = ak.index.Index64(np.array([0, 2, 4], dtype=np.int64)) + listarray = ak.highlevel.Array( + ak.contents.ListOffsetArray(offsets3, sublistarray), backend="cuda" + ) + assert to_list(listarray) == to_list(model) + assert to_list( + listarray[ + ak.highlevel.Array( + [[[[3, 2, 1, 1, 0], [], [1]], []], [[[0, 0]], [[1, 2]]]], backend="cuda" + ) + ] + ) == [[[[2.2, None, 1.1, 1.1, 0.0], [], [None]], []], [[[5.5, 5.5]], [[7.7, None]]]] + + +def test_0127_tomask_operation(): + array = ak.highlevel.Array( + [[0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]], backend="cuda" + ) + mask1 = ak.highlevel.Array([True, True, False, False, True], backend="cuda") + assert to_list(array[mask1]) == [[0.0, 1.1, 2.2], [], [6.6, 7.7, 8.8, 9.9]] + assert to_list(ak.operations.mask(array, mask1)) == [ + [0.0, 1.1, 2.2], + [], + None, + None, + [6.6, 7.7, 8.8, 9.9], + ] + + mask2 = ak.highlevel.Array( + [[False, True, False], [], [True, True], [False], [True, False, False, True]], + backend="cuda", + ) + assert to_list(array[mask2]) == [[1.1], [], [3.3, 4.4], [], [6.6, 9.9]] + assert to_list(ak.operations.mask(array, mask2)) == [ + [None, 1.1, None], + [], + [3.3, 4.4], + [None], + [6.6, None, None, 9.9], + ] From 10d8c80334a707a62ed462a94fe2d2ba9018e867 Mon Sep 17 00:00:00 2001 From: Manasvi Goyal <55101825+ManasviGoyal@users.noreply.github.com> Date: Wed, 5 Jun 2024 13:41:34 +0200 Subject: [PATCH 7/7] fix: remove unused variable --- .../cuda/cuda_kernels/awkward_ListArray_getitem_jagged_apply.cu | 1 - 1 file changed, 1 deletion(-) diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_apply.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_apply.cu index b3017113b4..ce5e8c940c 100644 --- a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_apply.cu +++ b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_apply.cu @@ -58,7 +58,6 @@ awkward_ListArray_getitem_jagged_apply_a( if (start != stop && stop > contentlen) { RAISE_ERROR(LISTARRAY_GETITEM_JAGGED_APPLY_ERRORS::STOP_GET_LEN) } - int64_t count = stop - start; scan_in_array[thread_id + 1] = slicestop - slicestart; } }