From 464697a01ec19e4878b665bb3a862e27f726773e Mon Sep 17 00:00:00 2001 From: Manasvi Goyal <55101825+ManasviGoyal@users.noreply.github.com> Date: Thu, 2 May 2024 17:11:01 +0200 Subject: [PATCH 1/3] feat: add ak.flatten and ak.concatenate kernels (#3086) * feat: add awkward_UnionArray_nestedfill_tags_index * fix: backend error * test: add cuda integration tests * test: add more tests * fix: test names * feat: add ak.flatten CUDA kernels * fix: send 2D array to CUDA * test: fixes * fix: backend issue in indexoptionarray.py * fix: pointer to pointer for Cupy * fix: awkward_UnionArray_flatten_combine * fix: remove changes * refactor: remove dead code * fix: cases * fix: error * fix: pointer to pointer cases * fix: kernel names and merge reduce sum int bool kernels * feat: add spec for awkward_sort * fix: spec for awkward_sort * fix: spec order * fix: tests with bool output * fix: Cupy case for double pointer * fix: unmerge awkward_reduce_sum_int_bool kernels * refactor: remove FIXME for unused arguments * 'awkward_regularize_rangeslice' is used by multiple kernels --------- Co-authored-by: Jim Pivarski --- ...wkward_IndexedArray_numnull_unique_64.cpp} | 0 ..._ListArray_getitem_next_array_advanced.cpp | 7 - ...wkward_ListOffsetArray_flatten_offsets.cpp | 21 +- .../awkward_NumpyArray_rearrange_shifted.cpp | 10 +- ...gularArray_getitem_next_array_advanced.cpp | 3 - dev/generate-kernel-signatures.py | 7 +- dev/generate-tests.py | 146 ++- kernel-specification.yml | 44 +- kernel-test-data.json | 428 ++++++--- src/awkward/_connect/cuda/__init__.py | 3 + ...d_ListArray_getitem_next_array_advanced.cu | 1 - ...awkward_ListOffsetArray_flatten_offsets.cu | 1 - .../awkward_NumpyArray_rearrange_shifted.cu | 10 +- ...egularArray_getitem_next_array_advanced.cu | 1 - ...gularArray_reduce_local_nextparents_64.cu} | 6 +- ...arArray_reduce_nonlocal_preparenext_64.cu} | 6 +- .../awkward_UnionArray_flatten_combine.cu | 69 ++ .../awkward_UnionArray_flatten_length.cu | 51 + ...wkward_UnionArray_nestedfill_tags_index.cu | 65 ++ src/awkward/contents/indexedoptionarray.py | 4 +- src/awkward/contents/listarray.py | 1 - src/awkward/contents/listoffsetarray.py | 2 - src/awkward/contents/numpyarray.py | 2 - src/awkward/contents/regulararray.py | 5 +- src/awkward/contents/unionarray.py | 112 ++- src/awkward/index.py | 5 +- src/awkward/operations/ak_angle.py | 2 - src/awkward/operations/ak_argcartesian.py | 2 - src/awkward/operations/ak_concatenate.py | 5 +- src/awkward/operations/ak_fill_none.py | 2 - src/awkward/operations/ak_imag.py | 2 - src/awkward/operations/ak_linear_fit.py | 2 - .../operations/ak_merge_option_of_records.py | 2 - .../operations/ak_merge_union_of_records.py | 2 - src/awkward/operations/ak_real.py | 2 - src/awkward/operations/ak_run_lengths.py | 2 - src/awkward/operations/ak_where.py | 2 - src/awkward/operations/ak_with_field.py | 2 - tests-cuda/test_3086_cuda_concatenate.py | 877 ++++++++++++++++++ tests-cuda/test_3086_cuda_flatten.py | 476 ++++++++++ 40 files changed, 2087 insertions(+), 303 deletions(-) rename awkward-cpp/src/cpu-kernels/{awkward_IndexedArray_numnull_unique.cpp => awkward_IndexedArray_numnull_unique_64.cpp} (100%) rename src/awkward/_connect/cuda/cuda_kernels/{awkward_RegularArray_reduce_local_nextparents.cu => awkward_RegularArray_reduce_local_nextparents_64.cu} (75%) rename src/awkward/_connect/cuda/cuda_kernels/{awkward_RegularArray_reduce_nonlocal_preparenext.cu => awkward_RegularArray_reduce_nonlocal_preparenext_64.cu} (75%) create mode 100644 src/awkward/_connect/cuda/cuda_kernels/awkward_UnionArray_flatten_combine.cu create mode 100644 src/awkward/_connect/cuda/cuda_kernels/awkward_UnionArray_flatten_length.cu create mode 100644 src/awkward/_connect/cuda/cuda_kernels/awkward_UnionArray_nestedfill_tags_index.cu create mode 100644 tests-cuda/test_3086_cuda_concatenate.py create mode 100644 tests-cuda/test_3086_cuda_flatten.py diff --git a/awkward-cpp/src/cpu-kernels/awkward_IndexedArray_numnull_unique.cpp b/awkward-cpp/src/cpu-kernels/awkward_IndexedArray_numnull_unique_64.cpp similarity index 100% rename from awkward-cpp/src/cpu-kernels/awkward_IndexedArray_numnull_unique.cpp rename to awkward-cpp/src/cpu-kernels/awkward_IndexedArray_numnull_unique_64.cpp diff --git a/awkward-cpp/src/cpu-kernels/awkward_ListArray_getitem_next_array_advanced.cpp b/awkward-cpp/src/cpu-kernels/awkward_ListArray_getitem_next_array_advanced.cpp index 316fb4d9dd..36348e1c08 100644 --- a/awkward-cpp/src/cpu-kernels/awkward_ListArray_getitem_next_array_advanced.cpp +++ b/awkward-cpp/src/cpu-kernels/awkward_ListArray_getitem_next_array_advanced.cpp @@ -13,7 +13,6 @@ ERROR awkward_ListArray_getitem_next_array_advanced( const T* fromarray, const T* fromadvanced, int64_t lenstarts, - int64_t /* lenarray */, // FIXME: this argument is not needed int64_t lencontent) { for (int64_t i = 0; i < lenstarts; i++) { if (fromstops[i] < fromstarts[i]) { @@ -44,7 +43,6 @@ ERROR awkward_ListArray32_getitem_next_array_advanced_64( const int64_t* fromarray, const int64_t* fromadvanced, int64_t lenstarts, - int64_t lenarray, int64_t lencontent) { return awkward_ListArray_getitem_next_array_advanced( tocarry, @@ -54,7 +52,6 @@ ERROR awkward_ListArray32_getitem_next_array_advanced_64( fromarray, fromadvanced, lenstarts, - lenarray, lencontent); } ERROR awkward_ListArrayU32_getitem_next_array_advanced_64( @@ -65,7 +62,6 @@ ERROR awkward_ListArrayU32_getitem_next_array_advanced_64( const int64_t* fromarray, const int64_t* fromadvanced, int64_t lenstarts, - int64_t lenarray, int64_t lencontent) { return awkward_ListArray_getitem_next_array_advanced( tocarry, @@ -75,7 +71,6 @@ ERROR awkward_ListArrayU32_getitem_next_array_advanced_64( fromarray, fromadvanced, lenstarts, - lenarray, lencontent); } ERROR awkward_ListArray64_getitem_next_array_advanced_64( @@ -86,7 +81,6 @@ ERROR awkward_ListArray64_getitem_next_array_advanced_64( const int64_t* fromarray, const int64_t* fromadvanced, int64_t lenstarts, - int64_t lenarray, int64_t lencontent) { return awkward_ListArray_getitem_next_array_advanced( tocarry, @@ -96,6 +90,5 @@ ERROR awkward_ListArray64_getitem_next_array_advanced_64( fromarray, fromadvanced, lenstarts, - lenarray, lencontent); } diff --git a/awkward-cpp/src/cpu-kernels/awkward_ListOffsetArray_flatten_offsets.cpp b/awkward-cpp/src/cpu-kernels/awkward_ListOffsetArray_flatten_offsets.cpp index 683dab7682..0e9a4facab 100644 --- a/awkward-cpp/src/cpu-kernels/awkward_ListOffsetArray_flatten_offsets.cpp +++ b/awkward-cpp/src/cpu-kernels/awkward_ListOffsetArray_flatten_offsets.cpp @@ -9,8 +9,7 @@ ERROR awkward_ListOffsetArray_flatten_offsets( T* tooffsets, const C* outeroffsets, int64_t outeroffsetslen, - const T* inneroffsets, - int64_t /* inneroffsetslen */) { // FIXME: this argument is not needed + const T* inneroffsets) { for (int64_t i = 0; i < outeroffsetslen; i++) { tooffsets[i] = inneroffsets[outeroffsets[i]]; @@ -21,38 +20,32 @@ ERROR awkward_ListOffsetArray32_flatten_offsets_64( int64_t* tooffsets, const int32_t* outeroffsets, int64_t outeroffsetslen, - const int64_t* inneroffsets, - int64_t inneroffsetslen) { + const int64_t* inneroffsets) { return awkward_ListOffsetArray_flatten_offsets( tooffsets, outeroffsets, outeroffsetslen, - inneroffsets, - inneroffsetslen); + inneroffsets); } ERROR awkward_ListOffsetArrayU32_flatten_offsets_64( int64_t* tooffsets, const uint32_t* outeroffsets, int64_t outeroffsetslen, - const int64_t* inneroffsets, - int64_t inneroffsetslen) { + const int64_t* inneroffsets) { return awkward_ListOffsetArray_flatten_offsets( tooffsets, outeroffsets, outeroffsetslen, - inneroffsets, - inneroffsetslen); + inneroffsets); } ERROR awkward_ListOffsetArray64_flatten_offsets_64( int64_t* tooffsets, const int64_t* outeroffsets, int64_t outeroffsetslen, - const int64_t* inneroffsets, - int64_t inneroffsetslen) { + const int64_t* inneroffsets) { return awkward_ListOffsetArray_flatten_offsets( tooffsets, outeroffsets, outeroffsetslen, - inneroffsets, - inneroffsetslen); + inneroffsets); } diff --git a/awkward-cpp/src/cpu-kernels/awkward_NumpyArray_rearrange_shifted.cpp b/awkward-cpp/src/cpu-kernels/awkward_NumpyArray_rearrange_shifted.cpp index c616baed9a..01229cc852 100644 --- a/awkward-cpp/src/cpu-kernels/awkward_NumpyArray_rearrange_shifted.cpp +++ b/awkward-cpp/src/cpu-kernels/awkward_NumpyArray_rearrange_shifted.cpp @@ -13,9 +13,7 @@ awkward_NumpyArray_rearrange_shifted( const FROM* fromoffsets, int64_t offsetslength, const FROM* fromparents, - int64_t /* parentslength */, // FIXME: these arguments are not needed - const FROM* fromstarts, - int64_t /* startslength */) { + const FROM* fromstarts) { int64_t k = 0; for (int64_t i = 0; i < offsetslength - 1; i++) { for (int64_t j = 0; j < fromoffsets[i + 1] - fromoffsets[i]; j++) { @@ -39,9 +37,7 @@ awkward_NumpyArray_rearrange_shifted_toint64_fromint64( const int64_t* fromoffsets, int64_t offsetslength, const int64_t* fromparents, - int64_t parentslength, - const int64_t* fromstarts, - int64_t startslength) { + const int64_t* fromstarts) { return awkward_NumpyArray_rearrange_shifted( - toptr, fromshifts, length, fromoffsets, offsetslength, fromparents, parentslength, fromstarts, startslength); + toptr, fromshifts, length, fromoffsets, offsetslength, fromparents, fromstarts); } diff --git a/awkward-cpp/src/cpu-kernels/awkward_RegularArray_getitem_next_array_advanced.cpp b/awkward-cpp/src/cpu-kernels/awkward_RegularArray_getitem_next_array_advanced.cpp index 3e1f17ead8..b3966c212e 100644 --- a/awkward-cpp/src/cpu-kernels/awkward_RegularArray_getitem_next_array_advanced.cpp +++ b/awkward-cpp/src/cpu-kernels/awkward_RegularArray_getitem_next_array_advanced.cpp @@ -11,7 +11,6 @@ ERROR awkward_RegularArray_getitem_next_array_advanced( const T* fromadvanced, const T* fromarray, int64_t length, - int64_t /* lenarray */, // FIXME: this argument is not needed int64_t size) { for (int64_t i = 0; i < length; i++) { tocarry[i] = i*size + fromarray[fromadvanced[i]]; @@ -25,7 +24,6 @@ ERROR awkward_RegularArray_getitem_next_array_advanced_64( const int64_t* fromadvanced, const int64_t* fromarray, int64_t length, - int64_t lenarray, int64_t size) { return awkward_RegularArray_getitem_next_array_advanced( tocarry, @@ -33,6 +31,5 @@ ERROR awkward_RegularArray_getitem_next_array_advanced_64( fromadvanced, fromarray, length, - lenarray, size); } diff --git a/dev/generate-kernel-signatures.py b/dev/generate-kernel-signatures.py index 539b0ef39a..b64946626c 100644 --- a/dev/generate-kernel-signatures.py +++ b/dev/generate-kernel-signatures.py @@ -48,8 +48,8 @@ "awkward_RegularArray_getitem_next_range_spreadadvanced", "awkward_RegularArray_getitem_next_array", "awkward_RegularArray_getitem_next_array_regularize", - "awkward_RegularArray_reduce_local_nextparents", - "awkward_RegularArray_reduce_nonlocal_preparenext", + "awkward_RegularArray_reduce_local_nextparents_64", + "awkward_RegularArray_reduce_nonlocal_preparenext_64", "awkward_missing_repeat", "awkward_RegularArray_getitem_jagged_expand", "awkward_ListArray_combinations_length", @@ -101,6 +101,9 @@ "awkward_ListOffsetArray_drop_none_indexes", "awkward_ListOffsetArray_reduce_local_nextparents_64", "awkward_ListOffsetArray_reduce_nonlocal_maxcount_offsetscopy_64", + "awkward_UnionArray_flatten_length", + "awkward_UnionArray_flatten_combine", + "awkward_UnionArray_nestedfill_tags_index", "awkward_UnionArray_regular_index_getsize", "awkward_UnionArray_simplify", "awkward_UnionArray_simplify_one", diff --git a/dev/generate-tests.py b/dev/generate-tests.py index 01c221e24d..99420b7939 100644 --- a/dev/generate-tests.py +++ b/dev/generate-tests.py @@ -239,12 +239,17 @@ def getdtypes(args): return dtypes +def gettype(arg, args): + typename = remove_const( + next(argument for argument in args if argument.name == arg).typename + ) + return typename + + def checkuint(test_args, args): flag = True for arg, val in test_args: - typename = remove_const( - next(argument for argument in args if argument.name == arg).typename - ) + typename = gettype(arg, args) if "List[uint" in typename and (any(n < 0 for n in val)): flag = False return flag @@ -254,9 +259,7 @@ def checkintrange(test_args, error, args): flag = True if not error: for arg, val in test_args: - typename = remove_const( - next(argument for argument in args if argument.name == arg).typename - ) + typename = gettype(arg, args) if "int" in typename or "uint" in typename: dtype = gettypename(typename) min_val, max_val = np.iinfo(dtype).min, np.iinfo(dtype).max @@ -305,6 +308,35 @@ def genpykernels(): from numpy import uint8 kMaxInt64 = 9223372036854775806 kSliceNone = kMaxInt64 + 1 + +def awkward_regularize_rangeslice( + start, stop, posstep, hasstart, hasstop, length, +): + if posstep: + if not hasstart: start = 0 + elif start < 0: start += length + if start < 0: start = 0 + if start > length: start = length + + if not hasstop: stop = length + elif stop < 0: stop += length + if stop < 0: stop = 0 + if stop > length: stop = length + if stop < start: stop = start + + else: + if not hasstart: start = length - 1 + elif start < 0: start += length + if start < -1: start = -1 + if start > length - 1: start = length - 1 + + if not hasstop: stop = -1 + elif stop < 0: stop += length + if stop < -1: stop = -1 + if stop > length - 1: stop = length - 1 + if stop > start: stop = start + return start, stop + """ tests_spec = os.path.join(CURRENT_DIR, "..", "awkward-cpp", "tests-spec") @@ -503,11 +535,7 @@ def gencpukerneltests(specdict): num += 1 for arg, val in test["inargs"].items(): f.write(" " * 4 + arg + " = " + str(val) + "\n") - typename = remove_const( - next( - argument for argument in spec.args if argument.name == arg - ).typename - ) + typename = gettype(arg, spec.args) if "List" in typename: count = typename.count("List") typename = gettypename(typename) @@ -615,13 +643,7 @@ def gencpuunittests(specdict): num += 1 f.write(funcName) for arg, val in test["outputs"].items(): - typename = remove_const( - next( - argument - for argument in spec.args - if argument.name == arg - ).typename - ) + typename = gettype(arg, spec.args) f.write( " " * 4 + arg @@ -649,14 +671,7 @@ def gencpuunittests(specdict): + f"{arg} = {arg}_ptr\n" ) for arg, val in test["inputs"].items(): - typename = remove_const( - next( - argument - for argument in spec.args - if argument.name == arg - ).typename - ) - + typename = gettype(arg, spec.args) f.write(" " * 4 + arg + " = " + str(val) + "\n") if "List" in typename: count = typename.count("List") @@ -690,6 +705,26 @@ def gencpuunittests(specdict): if not test["error"]: f.write(" " * 4 + "ret_pass = funcC(" + args + ")\n") for arg, val in test["outputs"].items(): + typename = gettype(arg, spec.args) + if "bool" in typename: + output = [] + if "List[List" in typename: + for row in val: + for data in row: + if data >= 1: + data = 1 + output.append(data) + val = output + elif "List" in typename: + for data in val: + if data >= 1: + data = 1 + output.append(data) + val = output + else: + if val >= 1: + val = 1 + f.write( " " * 4 + "pytest_" + arg + " = " + str(val) + "\n" ) @@ -746,8 +781,8 @@ def gencpuunittests(specdict): "awkward_RegularArray_getitem_next_range_spreadadvanced", "awkward_RegularArray_getitem_next_array", "awkward_RegularArray_getitem_next_array_regularize", - "awkward_RegularArray_reduce_local_nextparents", - "awkward_RegularArray_reduce_nonlocal_preparenext", + "awkward_RegularArray_reduce_local_nextparents_64", + "awkward_RegularArray_reduce_nonlocal_preparenext_64", "awkward_missing_repeat", "awkward_RegularArray_getitem_jagged_expand", "awkward_ListArray_combinations_length", @@ -799,6 +834,9 @@ def gencpuunittests(specdict): "awkward_ListOffsetArray_drop_none_indexes", "awkward_ListOffsetArray_reduce_local_nextparents_64", "awkward_ListOffsetArray_reduce_nonlocal_maxcount_offsetscopy_64", + "awkward_UnionArray_flatten_length", + "awkward_UnionArray_flatten_combine", + "awkward_UnionArray_nestedfill_tags_index", "awkward_UnionArray_regular_index_getsize", "awkward_UnionArray_simplify", "awkward_UnionArray_simplify_one", @@ -882,13 +920,7 @@ def gencudakerneltests(specdict): num += 1 dtypes = [] for arg, val in test["inargs"].items(): - typename = remove_const( - next( - argument - for argument in spec.args - if argument.name == arg - ).typename - ) + typename = gettype(arg, spec.args) if "List" not in typename: f.write(" " * 4 + arg + " = " + str(val) + "\n") if "List" in typename: @@ -1021,13 +1053,7 @@ def gencudaunittests(specdict): ) f.write(funcName) for arg, val in test["outputs"].items(): - typename = remove_const( - next( - argument - for argument in spec.args - if argument.name == arg - ).typename - ) + typename = gettype(arg, spec.args) if "List" not in typename: f.write(" " * 4 + arg + " = " + str(val) + "\n") if "List" in typename: @@ -1046,15 +1072,13 @@ def gencudaunittests(specdict): f.write( " " * 4 + f"{arg} = cupy.array({val}, dtype=cupy.{typename})\n" + + " " * 4 + + f"{arg}_array = [cupy.array(row, dtype=cupy.{typename}) for row in {arg}]\n" + + " " * 4 + + f"{arg} = cupy.array([row.data.ptr for row in {arg}_array])\n" ) for arg, val in test["inputs"].items(): - typename = remove_const( - next( - argument - for argument in spec.args - if argument.name == arg - ).typename - ) + typename = gettype(arg, spec.args) if "List" not in typename: f.write(" " * 4 + arg + " = " + str(val) + "\n") if "List" in typename: @@ -1073,6 +1097,10 @@ def gencudaunittests(specdict): f.write( " " * 4 + f"{arg} = cupy.array({val}, dtype=cupy.{typename})\n" + + " " * 4 + + f"{arg}_array = [cupy.array(row, dtype=cupy.{typename}) for row in {arg}]\n" + + " " * 4 + + f"{arg} = cupy.array([row.data.ptr for row in {arg}_array])\n" ) cuda_string = ( "funcC = cupy_backend['" @@ -1110,6 +1138,26 @@ def gencudaunittests(specdict): """ ) for arg, val in test["outputs"].items(): + typename = gettype(arg, spec.args) + if "bool" in typename: + output = [] + if "List[List" in typename: + for row in val: + for data in row: + if data >= 1: + data = 1 + output.append(data) + val = output + elif "List" in typename: + for data in val: + if data >= 1: + data = 1 + output.append(data) + val = output + else: + if val >= 1: + val = 1 + f.write( " " * 4 + "pytest_" + arg + " = " + str(val) + "\n" ) diff --git a/kernel-specification.yml b/kernel-specification.yml index 8d0e94b330..ad3c5b5265 100644 --- a/kernel-specification.yml +++ b/kernel-specification.yml @@ -558,6 +558,7 @@ kernels: def awkward_IndexedArray_local_preparenext( tocarry, starts, parents, parentslength, nextparents, nextlen ): + j = 0 for i in range(parentslength): parent = parents[i]; if j < nextlen and parent == nextparents[j]: @@ -1638,7 +1639,6 @@ kernels: - {name: fromarray, type: "Const[List[int64_t]]", dir: in, role: ListOffsetArray-offsets} - {name: fromadvanced, type: "Const[List[int64_t]]", dir: in, role: ListOffsetArray2-offsets} - {name: lenstarts, type: "int64_t", dir: in, role: default} - - {name: lenarray, type: "int64_t", dir: in, role: default} - {name: lencontent, type: "int64_t", dir: in, role: ListArray-length} - name: awkward_ListArray64_getitem_next_array_advanced_64 args: @@ -1649,7 +1649,6 @@ kernels: - {name: fromarray, type: "Const[List[int64_t]]", dir: in, role: ListOffsetArray-offsets} - {name: fromadvanced, type: "Const[List[int64_t]]", dir: in, role: ListOffsetArray2-offsets} - {name: lenstarts, type: "int64_t", dir: in, role: default} - - {name: lenarray, type: "int64_t", dir: in, role: default} - {name: lencontent, type: "int64_t", dir: in, role: ListArray-length} - name: awkward_ListArrayU32_getitem_next_array_advanced_64 args: @@ -1660,7 +1659,6 @@ kernels: - {name: fromarray, type: "Const[List[int64_t]]", dir: in, role: ListOffsetArray-offsets} - {name: fromadvanced, type: "Const[List[int64_t]]", dir: in, role: ListOffsetArray2-offsets} - {name: lenstarts, type: "int64_t", dir: in, role: default} - - {name: lenarray, type: "int64_t", dir: in, role: default} - {name: lencontent, type: "int64_t", dir: in, role: ListArray-length} description: null definition: | @@ -1672,7 +1670,6 @@ kernels: fromarray, fromadvanced, lenstarts, - lenarray, lencontent, ): for i in range(lenstarts): @@ -1770,7 +1767,7 @@ kernels: length = fromstops[i] - fromstarts[i] regular_start = start regular_stop = stop - awkward_regularize_rangeslice( + regular_start, regular_stop = awkward_regularize_rangeslice( regular_start, regular_stop, step > 0, @@ -1789,7 +1786,7 @@ kernels: length = fromstops[i] - fromstarts[i] regular_start = start regular_stop = stop - awkward_regularize_rangeslice( + regular_start, regular_stop = awkward_regularize_rangeslice( regular_start, regular_stop, step > 0, @@ -1844,7 +1841,7 @@ kernels: length = fromstops[i] - fromstarts[i] regular_start = start regular_stop = stop - awkward_regularize_rangeslice( + regular_start, regular_stop = awkward_regularize_rangeslice( regular_start, regular_stop, step > 0, @@ -2131,25 +2128,22 @@ kernels: - {name: outeroffsets, type: "Const[List[int32_t]]", dir: in, role: ListOffsetArray-offsets} - {name: outeroffsetslen, type: "int64_t", dir: in, role: default} - {name: inneroffsets, type: "Const[List[int64_t]]", dir: in, role: ListOffsetArray2-offsets} - - {name: inneroffsetslen, type: "int64_t", dir: in, role: default} - name: awkward_ListOffsetArray64_flatten_offsets_64 args: - {name: tooffsets, type: "List[int64_t]", dir: out} - {name: outeroffsets, type: "Const[List[int64_t]]", dir: in, role: ListOffsetArray-offsets} - {name: outeroffsetslen, type: "int64_t", dir: in, role: default} - {name: inneroffsets, type: "Const[List[int64_t]]", dir: in, role: ListOffsetArray2-offsets} - - {name: inneroffsetslen, type: "int64_t", dir: in, role: default} - name: awkward_ListOffsetArrayU32_flatten_offsets_64 args: - {name: tooffsets, type: "List[int64_t]", dir: out} - {name: outeroffsets, type: "Const[List[uint32_t]]", dir: in, role: ListOffsetArray-offsets} - {name: outeroffsetslen, type: "int64_t", dir: in, role: default} - {name: inneroffsets, type: "Const[List[int64_t]]", dir: in, role: ListOffsetArray2-offsets} - - {name: inneroffsetslen, type: "int64_t", dir: in, role: default} description: null definition: | def awkward_ListOffsetArray_flatten_offsets( - tooffsets, outeroffsets, outeroffsetslen, inneroffsets, inneroffsetslen + tooffsets, outeroffsets, outeroffsetslen, inneroffsets ): for i in range(outeroffsetslen): tooffsets[i] = inneroffsets[outeroffsets[i]] @@ -2591,12 +2585,10 @@ kernels: - {name: fromoffsets, type: "Const[List[int64_t]]", dir: in, role: IndexedArray-index} - {name: offsetslength, type: "int64_t", dir: in, role: default} - {name: fromparents, type: "Const[List[int64_t]]", dir: in, role: IndexedArray-index} - - {name: parentslength, type: "int64_t", dir: in, role: default} - {name: fromstarts, type: "Const[List[int64_t]]", dir: in, role: IndexedArray-index} - - {name: startslength, type: "int64_t", dir: in, role: default} description: null definition: | - def awkward_NumpyArray_rearrange_shifted(toptr, fromshifts, length, fromoffsets, offsetslength, fromparents, parentslength, fromstarts, startslength): + def awkward_NumpyArray_rearrange_shifted(toptr, fromshifts, length, fromoffsets, offsetslength, fromparents, fromstarts): k = 0 for i in range(offsetslength - 1): for j in range(fromoffsets[i + 1] - fromoffsets[i]): @@ -3022,12 +3014,11 @@ kernels: - {name: fromadvanced, type: "Const[List[int64_t]]", dir: in, role: ListOffsetArray-offsets} - {name: fromarray, type: "Const[List[int64_t]]", dir: in, role: ListOffsetArray2-offsets} - {name: length, type: "int64_t", dir: in, role: default} - - {name: lenarray, type: "int64_t", dir: in, role: default} - {name: size, type: "int64_t", dir: in, role: RegularArray-size} description: null definition: | def awkward_RegularArray_getitem_next_array_advanced( - tocarry, toadvanced, fromadvanced, fromarray, length, lenarray, size + tocarry, toadvanced, fromadvanced, fromarray, length, size ): for i in range(length): tocarry[i] = (i * size) + fromarray[fromadvanced[i]] @@ -3128,7 +3119,7 @@ kernels: toindex[(i * size) + j] = j automatic-tests: true - - name: awkward_RegularArray_reduce_local_nextparents + - name: awkward_RegularArray_reduce_local_nextparents_64 specializations: - name: awkward_RegularArray_reduce_local_nextparents_64 args: @@ -3137,7 +3128,7 @@ kernels: - {name: length, type: "int64_t", dir: in, role: reducer-length} description: null definition: | - def awkward_RegularArray_reduce_local_nextparents(nextparents, size, length): + def awkward_RegularArray_reduce_local_nextparents_64(nextparents, size, length): k = 0 for i in range(length): for _ in range(size): @@ -3146,7 +3137,7 @@ kernels: k += 1 automatic-tests: false - - name: awkward_RegularArray_reduce_nonlocal_preparenext + - name: awkward_RegularArray_reduce_nonlocal_preparenext_64 specializations: - name: awkward_RegularArray_reduce_nonlocal_preparenext_64 args: @@ -3157,7 +3148,7 @@ kernels: - {name: length, type: "int64_t", dir: in} description: null definition: | - def awkward_RegularArray_reduce_nonlocal_preparenext(nextcarry, nextparents, parents, size, length): + def awkward_RegularArray_reduce_nonlocal_preparenext_64(nextcarry, nextparents, parents, size, length): k = 0 for j in range(size): for i in range(length): @@ -5126,7 +5117,18 @@ kernels: - {name: stable, type: "bool", dir: in, role: ListArray-replacement} description: null definition: | - Insert Python definition here + def awkward_sort( + toptr, fromptr, length, offsets, offsetslength, parentslength, ascending, stable + ): + result = [] + for i in range(offsetslength - 1): + sub_list = sorted(fromptr[offsets[i]:offsets[i+1]], reverse=not ascending) + if stable: + result += sub_list + else: + result += sorted(sub_list, reverse=not ascending) + for i in range(parentslength): + toptr[i] = result[i] automatic-tests: false - name: awkward_unique_offsets diff --git a/kernel-test-data.json b/kernel-test-data.json index a49090d6a3..b3e27a323a 100644 --- a/kernel-test-data.json +++ b/kernel-test-data.json @@ -5085,7 +5085,7 @@ ] }, { - "name": "awkward_RegularArray_reduce_local_nextparents", + "name": "awkward_RegularArray_reduce_local_nextparents_64", "status": true, "tests": [ { @@ -5146,7 +5146,7 @@ ] }, { - "name": "awkward_RegularArray_reduce_nonlocal_preparenext", + "name": "awkward_RegularArray_reduce_nonlocal_preparenext_64", "status": true, "tests": [ { @@ -7680,7 +7680,6 @@ "inputs": { "fromadvanced": [], "fromarray": [], - "lenarray": 0, "length": 0, "size": 0 }, @@ -7695,7 +7694,6 @@ "inputs": { "fromadvanced": [0, 1, 2, 3], "fromarray": [0, 1, 0, 1], - "lenarray": 4, "length": 4, "size": 2 }, @@ -7710,7 +7708,6 @@ "inputs": { "fromadvanced": [0, 1, 2, 3], "fromarray": [0, 0, 0, 0], - "lenarray": 4, "length": 4, "size": 3 }, @@ -7725,7 +7722,6 @@ "inputs": { "fromadvanced": [0, 1, 2, 3], "fromarray": [0, 1, 2, 1], - "lenarray": 4, "length": 4, "size": 3 }, @@ -7740,7 +7736,6 @@ "inputs": { "fromadvanced": [0, 1, 2, 3], "fromarray": [0, 0, 0, 0], - "lenarray": 4, "length": 4, "size": 5 }, @@ -7755,7 +7750,6 @@ "inputs": { "fromadvanced": [0, 1, 2, 3], "fromarray": [0, 1, 4, 1], - "lenarray": 4, "length": 4, "size": 5 }, @@ -7770,7 +7764,6 @@ "inputs": { "fromadvanced": [0, 1, 2, 3], "fromarray": [1, 0, 1, 0], - "lenarray": 4, "length": 4, "size": 2 }, @@ -7785,7 +7778,6 @@ "inputs": { "fromadvanced": [0, 1, 2, 3], "fromarray": [1, 0, 0, 1], - "lenarray": 4, "length": 4, "size": 3 }, @@ -7800,7 +7792,6 @@ "inputs": { "fromadvanced": [0, 1, 2, 3], "fromarray": [1, 0, 0, 1], - "lenarray": 4, "length": 4, "size": 5 }, @@ -7815,7 +7806,6 @@ "inputs": { "fromadvanced": [0, 1, 2, 3], "fromarray": [1, 3, 0, 4], - "lenarray": 4, "length": 4, "size": 5 }, @@ -7830,7 +7820,6 @@ "inputs": { "fromadvanced": [0, 1, 2, 3], "fromarray": [2, 0, 0, 1], - "lenarray": 4, "length": 4, "size": 5 }, @@ -7845,7 +7834,6 @@ "inputs": { "fromadvanced": [0, 1, 2, 3], "fromarray": [2, 2, 2, 2], - "lenarray": 4, "length": 4, "size": 3 }, @@ -7860,7 +7848,6 @@ "inputs": { "fromadvanced": [0, 1, 2, 3], "fromarray": [2, 2, 2, 2], - "lenarray": 4, "length": 4, "size": 5 }, @@ -7875,7 +7862,6 @@ "inputs": { "fromadvanced": [0, 1, 2, 3], "fromarray": [3, 3, 3, 3], - "lenarray": 4, "length": 4, "size": 5 }, @@ -7890,7 +7876,6 @@ "inputs": { "fromadvanced": [0, 1, 2, 3], "fromarray": [4, 4, 4, 4], - "lenarray": 4, "length": 4, "size": 5 }, @@ -10241,7 +10226,6 @@ "message": "", "inputs": { "inneroffsets": [0], - "inneroffsetslen": 1, "outeroffsets": [], "outeroffsetslen": 0 }, @@ -10254,7 +10238,6 @@ "message": "", "inputs": { "inneroffsets": [0, 1, 2, 3], - "inneroffsetslen": 4, "outeroffsets": [0, 0, 0, 1, 3], "outeroffsetslen": 5 }, @@ -10267,7 +10250,6 @@ "message": "", "inputs": { "inneroffsets": [0, 1, 2, 3, 4, 5, 6], - "inneroffsetslen": 7, "outeroffsets": [0, 0, 1, 3, 6], "outeroffsetslen": 5 }, @@ -10280,7 +10262,6 @@ "message": "", "inputs": { "inneroffsets": [0, 1, 1, 5], - "inneroffsetslen": 4, "outeroffsets": [0, 0, 1, 3], "outeroffsetslen": 4 }, @@ -10293,7 +10274,6 @@ "message": "", "inputs": { "inneroffsets": [0, 1, 1, 6, 6], - "inneroffsetslen": 5, "outeroffsets": [0, 0, 1, 4], "outeroffsetslen": 4 }, @@ -10306,7 +10286,6 @@ "message": "", "inputs": { "inneroffsets": [0, 4, 8, 12, 14, 16], - "inneroffsetslen": 6, "outeroffsets": [0, 3, 3, 5], "outeroffsetslen": 4 }, @@ -10319,7 +10298,6 @@ "message": "", "inputs": { "inneroffsets": [0, 1, 2, 5, 5, 7, 7, 11], - "inneroffsetslen": 8, "outeroffsets": [0, 1, 2, 2, 5, 7], "outeroffsetslen": 6 }, @@ -10332,7 +10310,6 @@ "message": "", "inputs": { "inneroffsets": [0, 1, 2, 3, 4, 5, 6], - "inneroffsetslen": 7, "outeroffsets": [0, 1, 3, 6], "outeroffsetslen": 4 }, @@ -10345,7 +10322,6 @@ "message": "", "inputs": { "inneroffsets": [0, 5, 10, 15, 20, 25, 30], - "inneroffsetslen": 7, "outeroffsets": [0, 3, 6], "outeroffsetslen": 3 }, @@ -10358,7 +10334,6 @@ "message": "", "inputs": { "inneroffsets": [0, 2, 6], - "inneroffsetslen": 3, "outeroffsets": [0, 1, 1, 1, 2], "outeroffsetslen": 5 }, @@ -10371,7 +10346,6 @@ "message": "", "inputs": { "inneroffsets": [0, 0, 2, 6], - "inneroffsetslen": 4, "outeroffsets": [0, 2, 2, 2, 3], "outeroffsetslen": 5 }, @@ -10384,7 +10358,6 @@ "message": "", "inputs": { "inneroffsets": [0, 0, 0, 2, 6], - "inneroffsetslen": 5, "outeroffsets": [0, 3, 3, 3, 4], "outeroffsetslen": 5 }, @@ -10397,7 +10370,6 @@ "message": "", "inputs": { "inneroffsets": [0, 0, 0, 0, 2, 7, 7], - "inneroffsetslen": 7, "outeroffsets": [0, 4, 4, 4, 6], "outeroffsetslen": 5 }, @@ -10410,7 +10382,6 @@ "message": "", "inputs": { "inneroffsets": [0, 3, 5, 6, 6], - "inneroffsetslen": 5, "outeroffsets": [0, 1, 2, 3, 4], "outeroffsetslen": 5 }, @@ -10423,7 +10394,6 @@ "message": "", "inputs": { "inneroffsets": [0, 7, 14, 21, 28, 35, 42, 49, 56, 63, 70, 77, 84, 91, 98, 105, 112, 119, 126, 133, 140, 147, 154, 161, 168, 175, 182, 189, 196, 203, 210], - "inneroffsetslen": 31, "outeroffsets": [0, 5, 10, 15, 20, 25, 30], "outeroffsetslen": 7 }, @@ -10436,7 +10406,6 @@ "message": "", "inputs": { "inneroffsets": [0, 2, 4, 6, 8, 10, 12, 13, 14, 15, 16], - "inneroffsetslen": 11, "outeroffsets": [0, 2, 4, 6, 8, 10], "outeroffsetslen": 6 }, @@ -10449,7 +10418,6 @@ "message": "", "inputs": { "inneroffsets": [0, 3, 5, 6, 6, 10], - "inneroffsetslen": 6, "outeroffsets": [0, 2, 2, 3, 5], "outeroffsetslen": 5 }, @@ -10462,7 +10430,6 @@ "message": "", "inputs": { "inneroffsets": [0, 3, 3, 5, 6, 6, 10], - "inneroffsetslen": 7, "outeroffsets": [0, 3, 3, 4, 6], "outeroffsetslen": 5 }, @@ -10475,7 +10442,6 @@ "message": "", "inputs": { "inneroffsets": [0, 3, 3, 3, 5, 6, 6, 10], - "inneroffsetslen": 8, "outeroffsets": [0, 4, 4, 5, 7], "outeroffsetslen": 5 }, @@ -10488,7 +10454,6 @@ "message": "", "inputs": { "inneroffsets": [0, 3, 3, 5, 5, 8], - "inneroffsetslen": 6, "outeroffsets": [0, 3, 3, 5], "outeroffsetslen": 4 }, @@ -10501,7 +10466,6 @@ "message": "", "inputs": { "inneroffsets": [0, 3, 6, 9, 12, 14, 16], - "inneroffsetslen": 7, "outeroffsets": [0, 2, 4, 5, 6, 6, 6], "outeroffsetslen": 7 }, @@ -10514,7 +10478,6 @@ "message": "", "inputs": { "inneroffsets": [0, 2, 4, 6, 8, 10], - "inneroffsetslen": 6, "outeroffsets": [0, 3, 3, 5], "outeroffsetslen": 4 }, @@ -10527,7 +10490,6 @@ "message": "", "inputs": { "inneroffsets": [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20], - "inneroffsetslen": 11, "outeroffsets": [0, 3, 3, 5, 8, 8, 10], "outeroffsetslen": 7 }, @@ -10540,7 +10502,6 @@ "message": "", "inputs": { "inneroffsets": [0, 4, 4, 4, 4, 6, 7, 7, 12, 12], - "inneroffsetslen": 10, "outeroffsets": [0, 5, 5, 6, 9], "outeroffsetslen": 5 }, @@ -10553,7 +10514,6 @@ "message": "", "inputs": { "inneroffsets": [0, 3, 6, 9, 11, 13, 14], - "inneroffsetslen": 7, "outeroffsets": [0, 3, 5, 6], "outeroffsetslen": 4 }, @@ -10566,7 +10526,6 @@ "message": "", "inputs": { "inneroffsets": [0, 3, 5, 6, 6], - "inneroffsetslen": 5, "outeroffsets": [1, 2, 3, 4], "outeroffsetslen": 4 }, @@ -10579,7 +10538,6 @@ "message": "", "inputs": { "inneroffsets": [0, 3, 3, 5, 6, 6, 10], - "inneroffsetslen": 7, "outeroffsets": [3, 3, 4, 6], "outeroffsetslen": 4 }, @@ -10592,7 +10550,6 @@ "message": "", "inputs": { "inneroffsets": [0, 3, 3, 3, 5, 6, 6, 10], - "inneroffsetslen": 8, "outeroffsets": [4, 4, 5, 7], "outeroffsetslen": 4 }, @@ -13815,6 +13772,57 @@ } ] }, + { + "name": "awkward_ListArray_getitem_next_range_carrylength", + "status": true, + "tests": [ + { + "error": false, + "message": "", + "inputs": { + "fromstarts": [], + "fromstops": [], + "lenstarts": 0, + "start": 0, + "stop": 0, + "step": 0 + }, + "outputs": { + "carrylength": [0] + } + }, + { + "error": false, + "message": "", + "inputs": { + "fromstarts": [0, 2, 2, 3, 5], + "fromstops": [2, 2, 3, 5, 7], + "lenstarts": 5, + "start": 0, + "stop": 3, + "step": 1 + }, + "outputs": { + "carrylength": [7] + } + }, + { + "error": false, + "message": "", + "inputs": { + "fromstarts": [0, 2, 2, 3, 5], + "fromstops": [2, 2, 3, 5, 7], + "lenstarts": 5, + "start": 0, + "stop": 2, + "step": 0 + }, + "outputs": { + "carrylength": [0] + } + } + ] + }, { "name": "awkward_ListArray_getitem_next_array", "status": true, @@ -15909,7 +15917,7 @@ }, { "name": "awkward_UnionArray_flatten_length", - "status": false, + "status": true, "tests": [ { "error": false, @@ -15937,7 +15945,6 @@ "total_length": [3] } }, - { "error": false, "message": "", @@ -15950,12 +15957,25 @@ "outputs": { "total_length": [7] } + }, + { + "error": false, + "message": "", + "inputs": { + "fromtags": [1, 1, 1, 1], + "fromindex": [0, 1, 2, 3], + "length": 4, + "offsetsraws": [[0, 1, 3, 5, 7], [1, 3, 5, 7, 9]] + }, + "outputs": { + "total_length": [8] + } } ] }, { "name": "awkward_UnionArray_flatten_combine", - "status": false, + "status": true, "tests": [ { "error": false, @@ -16398,6 +16418,68 @@ } ] }, + { + "name": "awkward_IndexedArray_local_preparenext", + "status": true, + "tests": [ + { + "error": false, + "message": "", + "inputs": { + "nextlen": 4, + "nextparents": [0, 0, 0, 0], + "parentslength": 5, + "parents": [0, 0, 0, 0, 0], + "starts": [0] + }, + "outputs": { + "tocarry": [0, 1, 2, 3, -1] + } + }, + { + "error": false, + "message": "", + "inputs": { + "nextlen": 7, + "nextparents": [0, 0, 0, 0, 1, 1, 1], + "parentslength": 11, + "parents": [0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1], + "starts": [0, 2, 5] + }, + "outputs": { + "tocarry": [0, 1, 2, 3, -1, -1, 4, 5, 6, -1, -1] + } + }, + { + "error": false, + "message": "", + "inputs": { + "nextlen": 0, + "nextparents": [], + "parentslength": 0, + "parents": [], + "starts": [] + }, + "outputs": { + "tocarry": [] + } + }, + { + "error": false, + "message": "", + "inputs": { + "nextlen": 9, + "nextparents": [0, 0, 0, 2, 2, 3, 4, 4, 4], + "parentslength": 17, + "parents": [0, 0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4], + "starts": [0, 5, 8, 11, 14] + }, + "outputs": { + "tocarry": [0, 1, 2, -1, -1, -1, -1, -1, 3, 4, -1, 5, -1, -1, 6, 7, 8] + } + } + ] + }, { "name": "awkward_IndexedArray_index_of_nulls", "status": true, @@ -16866,7 +16948,6 @@ "fromarray": [0, 0], "fromstarts": [0, 0], "fromstops": [-1, 1], - "lenarray": 2, "lencontent": 5, "lenstarts": 2 }, @@ -16883,7 +16964,6 @@ "fromarray": [0], "fromstarts": [0], "fromstops": [5], - "lenarray": 1, "lencontent": 4, "lenstarts": 1 }, @@ -16900,7 +16980,6 @@ "fromarray": [0, 1, -4, 1], "fromstarts": [0, 0, 0, 0], "fromstops": [3, 3, 3, 3], - "lenarray": 4, "lencontent": 6, "lenstarts": 4 }, @@ -16917,7 +16996,6 @@ "fromarray": [0, 0], "fromstarts": [0, 0], "fromstops": [1, 1], - "lenarray": 2, "lencontent": 5, "lenstarts": 2 }, @@ -16934,7 +17012,6 @@ "fromarray": [0], "fromstarts": [0], "fromstops": [2], - "lenarray": 1, "lencontent": 4, "lenstarts": 1 }, @@ -16951,7 +17028,6 @@ "fromarray": [], "fromstarts": [], "fromstops": [], - "lenarray": 0, "lencontent": 0, "lenstarts": 0 }, @@ -16968,7 +17044,6 @@ "fromarray": [0, 1, -1, 1], "fromstarts": [0, 0, 0, 0], "fromstops": [3, 3, 3, 3], - "lenarray": 4, "lencontent": 6, "lenstarts": 4 }, @@ -16985,7 +17060,6 @@ "fromarray": [0, 0, 0, 0], "fromstarts": [0, 3, 3, 3], "fromstops": [3, 6, 6, 6], - "lenarray": 4, "lencontent": 6, "lenstarts": 4 }, @@ -17002,7 +17076,6 @@ "fromarray": [1, 0], "fromstarts": [0, 0], "fromstops": [3, 3], - "lenarray": 2, "lencontent": 5, "lenstarts": 2 }, @@ -17019,7 +17092,6 @@ "fromarray": [1, 0, 0, 1], "fromstarts": [0, 0, 0, 0], "fromstops": [3, 3, 3, 3], - "lenarray": 4, "lencontent": 6, "lenstarts": 4 }, @@ -17036,7 +17108,6 @@ "fromarray": [0, 0], "fromstarts": [1, 0], "fromstops": [3, 2], - "lenarray": 2, "lencontent": 3, "lenstarts": 2 }, @@ -17053,7 +17124,6 @@ "fromarray": [0, 1, -1, 1], "fromstarts": [10, 10, 10, 10], "fromstops": [15, 15, 15, 15], - "lenarray": 4, "lencontent": 30, "lenstarts": 4 }, @@ -17070,7 +17140,6 @@ "fromarray": [0], "fromstarts": [1], "fromstops": [3], - "lenarray": 1, "lencontent": 3, "lenstarts": 1 }, @@ -17087,7 +17156,6 @@ "fromarray": [1, 0, 0, 1], "fromstarts": [10, 10, 10, 10], "fromstops": [15, 15, 15, 15], - "lenarray": 4, "lencontent": 30, "lenstarts": 4 }, @@ -17104,7 +17172,6 @@ "fromarray": [1, -2, 0, -1], "fromstarts": [10, 10, 10, 10], "fromstops": [15, 15, 15, 15], - "lenarray": 4, "lencontent": 30, "lenstarts": 4 }, @@ -17121,7 +17188,6 @@ "fromarray": [1, -2, 0, -1], "fromstarts": [10, 0, 0, 5], "fromstops": [15, 5, 5, 10], - "lenarray": 4, "lencontent": 30, "lenstarts": 4 }, @@ -17138,7 +17204,6 @@ "fromarray": [1, 0, 0, 1], "fromstarts": [0, 15, 15, 15], "fromstops": [5, 20, 20, 20], - "lenarray": 4, "lencontent": 30, "lenstarts": 4 }, @@ -17155,7 +17220,6 @@ "fromarray": [1, 2], "fromstarts": [0, 0], "fromstops": [3, 3], - "lenarray": 2, "lencontent": 5, "lenstarts": 2 }, @@ -17172,7 +17236,6 @@ "fromarray": [2, 0, 0, 1], "fromstarts": [10, 10, 10, 10], "fromstops": [15, 15, 15, 15], - "lenarray": 4, "lencontent": 30, "lenstarts": 4 }, @@ -17189,7 +17252,6 @@ "fromarray": [2, 2, 2, 2], "fromstarts": [10, 0, 0, 5], "fromstops": [15, 5, 5, 10], - "lenarray": 4, "lencontent": 30, "lenstarts": 4 }, @@ -17206,7 +17268,6 @@ "fromarray": [-2, -2, -2, -2], "fromstarts": [10, 10, 10, 10], "fromstops": [15, 15, 15, 15], - "lenarray": 4, "lencontent": 30, "lenstarts": 4 }, @@ -17223,7 +17284,6 @@ "fromarray": [-2, -2, -2, -2], "fromstarts": [10, 0, 0, 5], "fromstops": [15, 5, 5, 10], - "lenarray": 4, "lencontent": 30, "lenstarts": 4 }, @@ -17240,7 +17300,6 @@ "fromarray": [1, 0, 0, 1], "fromstarts": [0, 3, 3, 3], "fromstops": [3, 6, 6, 6], - "lenarray": 4, "lencontent": 6, "lenstarts": 4 }, @@ -17257,7 +17316,6 @@ "fromarray": [1, 0, 0, 1], "fromstarts": [0, 5, 10, 5], "fromstops": [5, 10, 15, 10], - "lenarray": 4, "lencontent": 30, "lenstarts": 4 }, @@ -17274,7 +17332,6 @@ "fromarray": [0, 1, -1, 1], "fromstarts": [15, 0, 0, 15], "fromstops": [20, 5, 5, 20], - "lenarray": 4, "lencontent": 30, "lenstarts": 4 }, @@ -17291,7 +17348,6 @@ "fromarray": [0, 1, -1, 1], "fromstarts": [15, 15, 15, 15], "fromstops": [20, 20, 20, 20], - "lenarray": 4, "lencontent": 30, "lenstarts": 4 }, @@ -17308,7 +17364,6 @@ "fromarray": [1, 0, 0, 1], "fromstarts": [15, 15, 15, 15], "fromstops": [20, 20, 20, 20], - "lenarray": 4, "lencontent": 30, "lenstarts": 4 }, @@ -17325,7 +17380,6 @@ "fromarray": [-1, -1, -1, -1], "fromstarts": [15, 0, 0, 15], "fromstops": [20, 5, 5, 20], - "lenarray": 4, "lencontent": 30, "lenstarts": 4 }, @@ -17342,7 +17396,6 @@ "fromarray": [-1, 0], "fromstarts": [0, 0], "fromstops": [3, 3], - "lenarray": 2, "lencontent": 5, "lenstarts": 2 }, @@ -17359,7 +17412,6 @@ "fromarray": [-1, -1, -1, -1], "fromstarts": [0, 0, 0, 0], "fromstops": [3, 3, 3, 3], - "lenarray": 4, "lencontent": 6, "lenstarts": 4 }, @@ -17376,7 +17428,6 @@ "fromarray": [1, 2], "fromstarts": [1, 1], "fromstops": [4, 4], - "lenarray": 2, "lencontent": 4, "lenstarts": 2 }, @@ -17393,7 +17444,6 @@ "fromarray": [-1, -1, -1, -1], "fromstarts": [0, 3, 3, 3], "fromstops": [3, 6, 6, 6], - "lenarray": 4, "lencontent": 6, "lenstarts": 4 }, @@ -17410,7 +17460,6 @@ "fromarray": [0, 0, 0, 0], "fromstarts": [3, 0, 0, 3], "fromstops": [6, 3, 3, 6], - "lenarray": 4, "lencontent": 6, "lenstarts": 4 }, @@ -17427,7 +17476,6 @@ "fromarray": [1], "fromstarts": [2], "fromstops": [5], - "lenarray": 1, "lencontent": 5, "lenstarts": 1 }, @@ -17444,7 +17492,6 @@ "fromarray": [-1, 0], "fromstarts": [1, 1], "fromstops": [4, 4], - "lenarray": 2, "lencontent": 4, "lenstarts": 2 }, @@ -17461,7 +17508,6 @@ "fromarray": [0, 1, -1, 1], "fromstarts": [3, 0, 0, 3], "fromstops": [6, 3, 3, 6], - "lenarray": 4, "lencontent": 6, "lenstarts": 4 }, @@ -17478,7 +17524,6 @@ "fromarray": [0, 0, 0, 0], "fromstarts": [3, 3, 3, 3], "fromstops": [6, 6, 6, 6], - "lenarray": 4, "lencontent": 6, "lenstarts": 4 }, @@ -17495,7 +17540,6 @@ "fromarray": [0, 1, -1, 1], "fromstarts": [3, 3, 3, 3], "fromstops": [6, 6, 6, 6], - "lenarray": 4, "lencontent": 6, "lenstarts": 4 }, @@ -17512,7 +17556,6 @@ "fromarray": [1, 1], "fromstarts": [3, 0], "fromstops": [5, 3], - "lenarray": 2, "lencontent": 5, "lenstarts": 2 }, @@ -17529,7 +17572,6 @@ "fromarray": [1, 1, 0, 0], "fromstarts": [3, 0, 0, 6], "fromstops": [5, 3, 3, 9], - "lenarray": 4, "lencontent": 9, "lenstarts": 4 }, @@ -17546,7 +17588,6 @@ "fromarray": [-1, -1, -1, -1], "fromstarts": [0, 15, 15, 15], "fromstops": [5, 20, 20, 20], - "lenarray": 4, "lencontent": 30, "lenstarts": 4 }, @@ -17563,7 +17604,6 @@ "fromarray": [1, -1, 0, 0], "fromstarts": [3, 0, 0, 3], "fromstops": [5, 3, 3, 5], - "lenarray": 4, "lencontent": 5, "lenstarts": 4 }, @@ -17580,7 +17620,6 @@ "fromarray": [1, 0, 0, 1], "fromstarts": [3, 3, 3, 3], "fromstops": [6, 6, 6, 6], - "lenarray": 4, "lencontent": 6, "lenstarts": 4 }, @@ -17597,7 +17636,6 @@ "fromarray": [-1, -1, -1, -1], "fromstarts": [0, 5, 10, 5], "fromstops": [5, 10, 15, 10], - "lenarray": 4, "lencontent": 30, "lenstarts": 4 }, @@ -17614,7 +17652,6 @@ "fromarray": [2, 0, 1, 1, 2, 0], "fromstarts": [3, 0, 3, 3, 3, 0], "fromstops": [6, 3, 6, 6, 6, 3], - "lenarray": 6, "lencontent": 6, "lenstarts": 6 }, @@ -17631,7 +17668,6 @@ "fromarray": [0, 1, -1, 1], "fromstarts": [5, 0, 0, 5], "fromstops": [10, 5, 5, 10], - "lenarray": 4, "lencontent": 30, "lenstarts": 4 }, @@ -17648,7 +17684,6 @@ "fromarray": [-1, -1, -1, -1], "fromstarts": [3, 0, 0, 3], "fromstops": [6, 3, 3, 6], - "lenarray": 4, "lencontent": 6, "lenstarts": 4 }, @@ -17665,7 +17700,6 @@ "fromarray": [1, 1, 0, 0], "fromstarts": [6, 0, 0, 6], "fromstops": [9, 3, 3, 9], - "lenarray": 4, "lencontent": 9, "lenstarts": 4 }, @@ -17682,7 +17716,6 @@ "fromarray": [-1, -1, -1, -1], "fromstarts": [5, 0, 0, 5], "fromstops": [10, 5, 5, 10], - "lenarray": 4, "lencontent": 30, "lenstarts": 4 }, @@ -20642,9 +20675,7 @@ "fromoffsets": [], "offsetslength": 0, "fromparents": [], - "parentslength": 0, - "fromstarts": [], - "startslength": 0 + "fromstarts": [] }, "outputs": { "toptr": [] @@ -20660,9 +20691,7 @@ "fromoffsets": [0, 1, 3, 3, 5, 7, 9], "offsetslength": 7, "fromparents": [0, 1, 3, 6], - "parentslength": 4, - "fromstarts": [0, 1, 2, 3, 4, 5, 6], - "startslength": 7 + "fromstarts": [0, 1, 2, 3, 4, 5, 6] }, "outputs": { "toptr": [0, 3, 3, 6, 7, 10, 11, 14, 15] @@ -20678,9 +20707,7 @@ "fromoffsets": [0, 2, 5, 8], "offsetslength": 4, "fromparents": [0, 1, 3, 6], - "parentslength": 4, - "fromstarts": [0, 1, 2, 3, 4, 5, 6], - "startslength": 7 + "fromstarts": [0, 1, 2, 3, 4, 5, 6] }, "outputs": { "toptr": [0, -1, 1, -2, 2, 5, 5, 5] @@ -21760,15 +21787,30 @@ "error": false, "message": "", "inputs": { - "tmpstarts": [0, 2, 3, 5, 7], - "fromcounts": [2, 2, 3, 4, 5], + "tmpstarts": [0, 5, 5, 6, 8], + "fromcounts": [5, 0, 2, 3, 1], "length": 5, "tag": 1 }, "outputs": { - "totags": [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], - "toindex": [0, 1, 2, 4, 5, 7, 8, 11, 12, 13, 14, 15], - "tmpstarts": [2, 4, 6, 9, 12] + "totags": [1, 1, 1, 1, 1, 1, 1, 1, 1], + "toindex": [0, 1, 2, 3, 4, 5, 7, 8, 10], + "tmpstarts": [5, 5, 7, 9, 9] + } + }, + { + "error": false, + "message": "", + "inputs": { + "tmpstarts": [0, 2, 4, 5, 7], + "fromcounts": [2, 3, 2, 2, 0], + "length": 5, + "tag": 1 + }, + "outputs": { + "totags": [1, 1, 1, 1, 1, 1, 1], + "toindex": [0, 1, 2, 3, 5, 7, 8], + "tmpstarts": [2, 5, 6, 7, 7] } } ] @@ -26142,6 +26184,45 @@ "outputs": { "toequal": [0] } + }, + { + "error": false, + "message": "", + "inputs": { + "tmpptr": [0, 2, 2, 3, 5], + "fromstarts": [0, 2, 3, 3], + "fromstops": [2, 3, 3, 5], + "length": 4 + }, + "outputs": { + "toequal": [0] + } + }, + { + "error": false, + "message": "", + "inputs": { + "tmpptr": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + "fromstarts": [0, 2, 4, 6, 8, 10], + "fromstops": [2, 4, 6, 8, 10, 12], + "length": 6 + }, + "outputs": { + "toequal": [1] + } + }, + { + "error": false, + "message": "", + "inputs": { + "tmpptr": [1, 1, 1, 1, 1, 1], + "fromstarts": [2, 2, 2, 2, 2, 2], + "fromstops": [4, 4, 4, 4, 4, 4], + "length": 6 + }, + "outputs": { + "toequal": [1] + } } ] }, @@ -26174,6 +26255,45 @@ "outputs": { "toequal": [1] } + }, + { + "error": false, + "message": "", + "inputs": { + "tmpptr": [0, 2, 2, 3, 5], + "fromstarts": [0, 2, 3, 3], + "fromstops": [2, 3, 3, 5], + "length": 4 + }, + "outputs": { + "toequal": [0] + } + }, + { + "error": false, + "message": "", + "inputs": { + "tmpptr": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + "fromstarts": [0, 2, 4, 6, 8, 10], + "fromstops": [2, 4, 6, 8, 10, 12], + "length": 6 + }, + "outputs": { + "toequal": [1] + } + }, + { + "error": false, + "message": "", + "inputs": { + "tmpptr": [1, 1, 1, 1, 1, 1], + "fromstarts": [2, 2, 2, 2, 2, 2], + "fromstops": [4, 4, 4, 4, 4, 4], + "length": 6 + }, + "outputs": { + "toequal": [1] + } } ] }, @@ -26358,6 +26478,92 @@ } } ] + }, + { + "name": "awkward_sort", + "status": false, + "tests": [ + { + "error": false, + "message": "", + "inputs": { + "fromptr": [], + "offsets": [], + "offsetslength": 0, + "parentslength": 0, + "length": 1, + "ascending": true, + "stable": true + }, + "outputs": { + "toptr": [] + } + }, + { + "error": false, + "message": "", + "inputs": { + "fromptr": [8, 6, 7, 5, 3, 0, 9], + "offsets": [0, 3, 3, 7], + "offsetslength": 4, + "parentslength": 7, + "length": 7, + "ascending": true, + "stable": true + }, + "outputs": { + "toptr": [6, 7, 8, 0, 3, 5, 9] + } + }, + { + "error": false, + "message": "", + "inputs": { + "fromptr": [8, 6, 7, 5, 3, 0, 9], + "offsets": [0, 3, 3, 7], + "offsetslength": 4, + "parentslength": 7, + "length": 7, + "ascending": false, + "stable": true + }, + "outputs": { + "toptr": [8, 7, 6, 9, 5, 3, 0] + } + }, + { + "error": false, + "message": "", + "inputs": { + "fromptr": [8, 6, 7, 5, 3, 0, 9], + "offsets": [0, 3, 3, 7], + "offsetslength": 4, + "parentslength": 7, + "length": 7, + "ascending": true, + "stable": false + }, + "outputs": { + "toptr": [6, 7, 8, 0, 3, 5, 9] + } + }, + { + "error": false, + "message": "", + "inputs": { + "fromptr": [8, 6, 7, 5, 3, 0, 9], + "offsets": [0, 3, 3, 7], + "offsetslength": 4, + "parentslength": 7, + "length": 7, + "ascending": false, + "stable": false + }, + "outputs": { + "toptr": [8, 7, 6, 9, 5, 3, 0] + } + } + ] } ] } diff --git a/src/awkward/_connect/cuda/__init__.py b/src/awkward/_connect/cuda/__init__.py index 499bafa485..354fdcd217 100644 --- a/src/awkward/_connect/cuda/__init__.py +++ b/src/awkward/_connect/cuda/__init__.py @@ -107,7 +107,10 @@ def fetch_template_specializations(kernel_dict): "awkward_ListOffsetArray_reduce_nonlocal_maxcount_offsetscopy_64", "awkward_ListOffsetArray_rpad_length_axis1", "awkward_MaskedArray_getitem_next_jagged_project", + "awkward_UnionArray_nestedfill_tags_index", "awkward_NumpyArray_rearrange_shifted", + "awkward_UnionArray_flatten_length", + "awkward_UnionArray_flatten_combine", "awkward_UnionArray_project", "awkward_reduce_count_64", "awkward_reduce_sum", diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_next_array_advanced.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_next_array_advanced.cu index 349a63c0f4..75e34ee928 100644 --- a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_next_array_advanced.cu +++ b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_next_array_advanced.cu @@ -16,7 +16,6 @@ awkward_ListArray_getitem_next_array_advanced( const W* fromarray, const X* fromadvanced, int64_t lenstarts, - int64_t lenarray, int64_t lencontent, uint64_t invocation_index, uint64_t* err_code) { diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListOffsetArray_flatten_offsets.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListOffsetArray_flatten_offsets.cu index 5bf7bab47e..bfc285486c 100644 --- a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListOffsetArray_flatten_offsets.cu +++ b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListOffsetArray_flatten_offsets.cu @@ -7,7 +7,6 @@ awkward_ListOffsetArray_flatten_offsets( const C* outeroffsets, int64_t outeroffsetslen, const U* inneroffsets, - int64_t inneroffsetslen, uint64_t invocation_index, uint64_t* err_code) { if (err_code[0] == NO_ERROR) { diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_NumpyArray_rearrange_shifted.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_NumpyArray_rearrange_shifted.cu index a7df7a5e55..28eb039552 100644 --- a/src/awkward/_connect/cuda/cuda_kernels/awkward_NumpyArray_rearrange_shifted.cu +++ b/src/awkward/_connect/cuda/cuda_kernels/awkward_NumpyArray_rearrange_shifted.cu @@ -2,9 +2,9 @@ // BEGIN PYTHON // def f(grid, block, args): -// (toptr, fromshifts, length, fromoffsets, offsetslength, fromparents, parentslength, fromstarts, startslength, invocation_index, err_code) = args -// cuda_kernel_templates.get_function(fetch_specialization(["awkward_NumpyArray_rearrange_shifted_a", toptr.dtype, fromshifts.dtype, fromoffsets.dtype, fromparents.dtype, fromstarts.dtype]))(grid, block, (toptr, fromshifts, length, fromoffsets, offsetslength, fromparents, parentslength, fromstarts, startslength, invocation_index, err_code)) -// cuda_kernel_templates.get_function(fetch_specialization(["awkward_NumpyArray_rearrange_shifted_b", toptr.dtype, fromshifts.dtype, fromoffsets.dtype, fromparents.dtype, fromstarts.dtype]))(grid, block, (toptr, fromshifts, length, fromoffsets, offsetslength, fromparents, parentslength, fromstarts, startslength, invocation_index, err_code)) +// (toptr, fromshifts, length, fromoffsets, offsetslength, fromparents, fromstarts, invocation_index, err_code) = args +// cuda_kernel_templates.get_function(fetch_specialization(["awkward_NumpyArray_rearrange_shifted_a", toptr.dtype, fromshifts.dtype, fromoffsets.dtype, fromparents.dtype, fromstarts.dtype]))(grid, block, (toptr, fromshifts, length, fromoffsets, offsetslength, fromparents, fromstarts, invocation_index, err_code)) +// cuda_kernel_templates.get_function(fetch_specialization(["awkward_NumpyArray_rearrange_shifted_b", toptr.dtype, fromshifts.dtype, fromoffsets.dtype, fromparents.dtype, fromstarts.dtype]))(grid, block, (toptr, fromshifts, length, fromoffsets, offsetslength, fromparents, fromstarts, invocation_index, err_code)) // out["awkward_NumpyArray_rearrange_shifted_a", {dtype_specializations}] = None // out["awkward_NumpyArray_rearrange_shifted_b", {dtype_specializations}] = None // END PYTHON @@ -19,9 +19,7 @@ awkward_NumpyArray_rearrange_shifted_a( U* fromoffsets, int64_t offsetslength, V* fromparents, - int64_t /* parentslength */, W* fromstarts, - int64_t /* startslength */, uint64_t invocation_index, uint64_t* err_code) { if (err_code[0] == NO_ERROR) { @@ -45,9 +43,7 @@ awkward_NumpyArray_rearrange_shifted_b( U* fromoffsets, int64_t offsetslength, V* fromparents, - int64_t /* parentslength */, W* fromstarts, - int64_t /* startslength */, uint64_t invocation_index, uint64_t* err_code) { if (err_code[0] == NO_ERROR) { diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_RegularArray_getitem_next_array_advanced.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_RegularArray_getitem_next_array_advanced.cu index 79382736d2..56340bbee3 100644 --- a/src/awkward/_connect/cuda/cuda_kernels/awkward_RegularArray_getitem_next_array_advanced.cu +++ b/src/awkward/_connect/cuda/cuda_kernels/awkward_RegularArray_getitem_next_array_advanced.cu @@ -8,7 +8,6 @@ awkward_RegularArray_getitem_next_array_advanced( const U* fromadvanced, const V* fromarray, int64_t length, - int64_t lenarray, int64_t size, uint64_t invocation_index, uint64_t* err_code) { diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_RegularArray_reduce_local_nextparents.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_RegularArray_reduce_local_nextparents_64.cu similarity index 75% rename from src/awkward/_connect/cuda/cuda_kernels/awkward_RegularArray_reduce_local_nextparents.cu rename to src/awkward/_connect/cuda/cuda_kernels/awkward_RegularArray_reduce_local_nextparents_64.cu index 1d66dc53ef..79da5441f5 100644 --- a/src/awkward/_connect/cuda/cuda_kernels/awkward_RegularArray_reduce_local_nextparents.cu +++ b/src/awkward/_connect/cuda/cuda_kernels/awkward_RegularArray_reduce_local_nextparents_64.cu @@ -5,13 +5,13 @@ // (nextparents, size, length, invocation_index, err_code) = args // scan_in_array = cupy.ones(length * size, dtype=cupy.int64) // scan_in_array = cupy.cumsum(scan_in_array) -// cuda_kernel_templates.get_function(fetch_specialization(['awkward_RegularArray_reduce_local_nextparents', nextparents.dtype]))(grid, block, (nextparents, size, length, scan_in_array, invocation_index, err_code)) -// out["awkward_RegularArray_reduce_local_nextparents", {dtype_specializations}] = None +// cuda_kernel_templates.get_function(fetch_specialization(['awkward_RegularArray_reduce_local_nextparents_64', nextparents.dtype]))(grid, block, (nextparents, size, length, scan_in_array, invocation_index, err_code)) +// out["awkward_RegularArray_reduce_local_nextparents_64", {dtype_specializations}] = None // END PYTHON template __global__ void -awkward_RegularArray_reduce_local_nextparents( +awkward_RegularArray_reduce_local_nextparents_64( T* nextparents, int64_t size, int64_t length, diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_RegularArray_reduce_nonlocal_preparenext.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_RegularArray_reduce_nonlocal_preparenext_64.cu similarity index 75% rename from src/awkward/_connect/cuda/cuda_kernels/awkward_RegularArray_reduce_nonlocal_preparenext.cu rename to src/awkward/_connect/cuda/cuda_kernels/awkward_RegularArray_reduce_nonlocal_preparenext_64.cu index 20f9d50e8d..c6d6b425ce 100644 --- a/src/awkward/_connect/cuda/cuda_kernels/awkward_RegularArray_reduce_nonlocal_preparenext.cu +++ b/src/awkward/_connect/cuda/cuda_kernels/awkward_RegularArray_reduce_nonlocal_preparenext_64.cu @@ -5,13 +5,13 @@ // (nextcarry, nextparents, parents, size, length, invocation_index, err_code) = args // scan_in_array = cupy.ones(length * size, dtype=cupy.int64) // scan_in_array = cupy.cumsum(scan_in_array) -// cuda_kernel_templates.get_function(fetch_specialization(['awkward_RegularArray_reduce_nonlocal_preparenext', nextcarry.dtype, nextparents.dtype, parents.dtype]))(grid, block, (nextcarry, nextparents, parents, size, length, scan_in_array, invocation_index, err_code)) -// out["awkward_RegularArray_reduce_nonlocal_preparenext", {dtype_specializations}] = None +// cuda_kernel_templates.get_function(fetch_specialization(['awkward_RegularArray_reduce_nonlocal_preparenext_64', nextcarry.dtype, nextparents.dtype, parents.dtype]))(grid, block, (nextcarry, nextparents, parents, size, length, scan_in_array, invocation_index, err_code)) +// out["awkward_RegularArray_reduce_nonlocal_preparenext_64", {dtype_specializations}] = None // END PYTHON template __global__ void -awkward_RegularArray_reduce_nonlocal_preparenext( +awkward_RegularArray_reduce_nonlocal_preparenext_64( T* nextcarry, C* nextparents, const U* parents, diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_UnionArray_flatten_combine.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_UnionArray_flatten_combine.cu new file mode 100644 index 0000000000..18dab536c4 --- /dev/null +++ b/src/awkward/_connect/cuda/cuda_kernels/awkward_UnionArray_flatten_combine.cu @@ -0,0 +1,69 @@ +// BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE + +// BEGIN PYTHON +// def f(grid, block, args): +// (totags, toindex, tooffsets, fromtags, fromindex, length, offsetsraws, invocation_index, err_code) = args +// scan_in_array_tooffsets = cupy.zeros(length + 1, dtype=cupy.int64) +// cuda_kernel_templates.get_function(fetch_specialization(['awkward_UnionArray_flatten_combine_a', totags.dtype, toindex.dtype, tooffsets.dtype, fromtags.dtype, fromindex.dtype, offsetsraws[0].dtype]))(grid, block, (totags, toindex, tooffsets, fromtags, fromindex, length, offsetsraws, scan_in_array_tooffsets, invocation_index, err_code)) +// scan_in_array_tooffsets = cupy.cumsum(scan_in_array_tooffsets) +// cuda_kernel_templates.get_function(fetch_specialization(['awkward_UnionArray_flatten_combine_b', totags.dtype, toindex.dtype, tooffsets.dtype, fromtags.dtype, fromindex.dtype, offsetsraws[0].dtype]))(grid, block, (totags, toindex, tooffsets, fromtags, fromindex, length, offsetsraws, scan_in_array_tooffsets, invocation_index, err_code)) +// out["awkward_UnionArray_flatten_combine_a", {dtype_specializations}] = None +// out["awkward_UnionArray_flatten_combine_b", {dtype_specializations}] = None +// END PYTHON + +template +__global__ void +awkward_UnionArray_flatten_combine_a( + T* totags, + C* toindex, + U* tooffsets, + const V* fromtags, + const W* fromindex, + int64_t length, + X** offsetsraws, + int64_t* scan_in_array_tooffsets, + 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) { + V tag = fromtags[thread_id]; + W idx = fromindex[thread_id]; + X start = offsetsraws[tag][idx]; + X stop = offsetsraws[tag][idx + 1]; + scan_in_array_tooffsets[thread_id + 1] = stop - start; + } + } +} + +template +__global__ void +awkward_UnionArray_flatten_combine_b( + T* totags, + C* toindex, + U* tooffsets, + const V* fromtags, + const W* fromindex, + int64_t length, + X** offsetsraws, + int64_t* scan_in_array_tooffsets, + 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) { + V tag = fromtags[thread_id]; + W idx = fromindex[thread_id]; + X start = offsetsraws[tag][idx]; + X stop = offsetsraws[tag][idx + 1]; + int64_t k = scan_in_array_tooffsets[thread_id]; + for (int64_t j = start; j < stop; j++) { + totags[k] = tag; + toindex[k] = j; + k++; + } + tooffsets[thread_id] = scan_in_array_tooffsets[thread_id]; + } + tooffsets[length] = scan_in_array_tooffsets[length]; + } +} diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_UnionArray_flatten_length.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_UnionArray_flatten_length.cu new file mode 100644 index 0000000000..4ffb7d2e3b --- /dev/null +++ b/src/awkward/_connect/cuda/cuda_kernels/awkward_UnionArray_flatten_length.cu @@ -0,0 +1,51 @@ +// BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE + +// BEGIN PYTHON +// def f(grid, block, args): +// (total_length, fromtags, fromindex, length, offsetsraws, invocation_index, err_code) = args +// scan_in_array = cupy.zeros(length, dtype=cupy.int64) +// cuda_kernel_templates.get_function(fetch_specialization(['awkward_UnionArray_flatten_length_a', total_length.dtype, fromtags.dtype, fromindex.dtype, offsetsraws[0].dtype]))(grid, block, (total_length, fromtags, fromindex, length, offsetsraws, scan_in_array, invocation_index, err_code)) +// scan_in_array = cupy.cumsum(scan_in_array) +// cuda_kernel_templates.get_function(fetch_specialization(['awkward_UnionArray_flatten_length_b', total_length.dtype, fromtags.dtype, fromindex.dtype, offsetsraws[0].dtype]))(grid, block, (total_length, fromtags, fromindex, length, offsetsraws, scan_in_array, invocation_index, err_code)) +// out["awkward_UnionArray_flatten_length_a", {dtype_specializations}] = None +// out["awkward_UnionArray_flatten_length_b", {dtype_specializations}] = None +// END PYTHON + +template +__global__ void +awkward_UnionArray_flatten_length_a( + T* total_length, + const C* fromtags, + const U* fromindex, + int64_t length, + V** offsetsraws, + 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 < length) { + C tag = fromtags[thread_id]; + U idx = fromindex[thread_id]; + V start = offsetsraws[tag][idx]; + V stop = offsetsraws[tag][idx + 1]; + scan_in_array[thread_id] = stop - start; + } + } +} + +template +__global__ void +awkward_UnionArray_flatten_length_b( + T* total_length, + const C* fromtags, + const U* fromindex, + int64_t length, + V** offsetsraws, + int64_t* scan_in_array, + uint64_t invocation_index, + uint64_t* err_code) { + if (err_code[0] == NO_ERROR) { + *total_length = length > 0 ? scan_in_array[length - 1] : 0; + } +} diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_UnionArray_nestedfill_tags_index.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_UnionArray_nestedfill_tags_index.cu new file mode 100644 index 0000000000..d27ecba5e0 --- /dev/null +++ b/src/awkward/_connect/cuda/cuda_kernels/awkward_UnionArray_nestedfill_tags_index.cu @@ -0,0 +1,65 @@ +// BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE + +// BEGIN PYTHON +// def f(grid, block, args): +// (totags, toindex, tmpstarts, tag, fromcounts, length, invocation_index, err_code) = args +// if length > 0: +// scan_in_array = cupy.zeros(int(tmpstarts[length -1] + fromcounts[length - 1]), dtype=cupy.int64) +// else: +// scan_in_array = cupy.zeros(length, dtype=cupy.int64) +// cuda_kernel_templates.get_function(fetch_specialization(["awkward_UnionArray_nestedfill_tags_index_a", totags.dtype, toindex.dtype, tmpstarts.dtype, fromcounts.dtype]))(grid, block, (totags, toindex, tmpstarts, tag, fromcounts, length, scan_in_array, invocation_index, err_code)) +// scan_in_array = cupy.cumsum(scan_in_array) +// cuda_kernel_templates.get_function(fetch_specialization(["awkward_UnionArray_nestedfill_tags_index_b", totags.dtype, toindex.dtype, tmpstarts.dtype, fromcounts.dtype]))(grid, block, (totags, toindex, tmpstarts, tag, fromcounts, length, scan_in_array, invocation_index, err_code)) +// out["awkward_UnionArray_nestedfill_tags_index_a", {dtype_specializations}] = None +// out["awkward_UnionArray_nestedfill_tags_index_b", {dtype_specializations}] = None +// END PYTHON + +template +__global__ void +awkward_UnionArray_nestedfill_tags_index_a( + T* totags, + C* toindex, + U* tmpstarts, + T tag, + const V* fromcounts, + int64_t length, + 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 < length) { + U start = tmpstarts[thread_id]; + V stop = start + fromcounts[thread_id]; + for (int64_t j = start; j < stop; j++) { + scan_in_array[j] += 1; + } + } + } +} + +template +__global__ void +awkward_UnionArray_nestedfill_tags_index_b( + T* totags, + C* toindex, + U* tmpstarts, + T tag, + const V* fromcounts, + int64_t length, + 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 < length) { + U start = tmpstarts[thread_id]; + V stop = start + fromcounts[thread_id]; + for (int64_t j = start; j < stop; j++) { + totags[j] = tag; + toindex[j] = (C)(scan_in_array[j] - 1); + } + tmpstarts[thread_id] = stop; + } + } +} diff --git a/src/awkward/contents/indexedoptionarray.py b/src/awkward/contents/indexedoptionarray.py index 55e3a3a326..408f4b3028 100644 --- a/src/awkward/contents/indexedoptionarray.py +++ b/src/awkward/contents/indexedoptionarray.py @@ -1739,8 +1739,8 @@ def _to_list(self, behavior, json_conversions): index = self._index.raw(numpy) not_missing = index >= 0 - - nextcontent = self._content._carry(ak.index.Index(index[not_missing]), False) + content = ak.to_backend(self._content, "cpu", highlevel=False) + nextcontent = content._carry(ak.index.Index(index[not_missing]), False) out = nextcontent._to_list(behavior, json_conversions) for i, isvalid in enumerate(not_missing): diff --git a/src/awkward/contents/listarray.py b/src/awkward/contents/listarray.py index f51ce95180..722b9044dd 100644 --- a/src/awkward/contents/listarray.py +++ b/src/awkward/contents/listarray.py @@ -985,7 +985,6 @@ def _getitem_next( regular_flathead.data, advanced.data, lenstarts, - regular_flathead.length, self._content.length, ), slicer=head, diff --git a/src/awkward/contents/listoffsetarray.py b/src/awkward/contents/listoffsetarray.py index cc704e78a8..4aa149b69d 100644 --- a/src/awkward/contents/listoffsetarray.py +++ b/src/awkward/contents/listoffsetarray.py @@ -680,7 +680,6 @@ def _getitem_next( regular_flathead.data, advanced.data, lenstarts, - regular_flathead.length, self._content.length, ), slicer=head, @@ -769,7 +768,6 @@ def _offsets_and_flattened(self, axis: int, depth: int) -> tuple[Index, Content] self._offsets.data, self._offsets.length, inneroffsets.data, - inneroffsets.length, ) ) return ( diff --git a/src/awkward/contents/numpyarray.py b/src/awkward/contents/numpyarray.py index 2428e617a8..3cd0bdf113 100644 --- a/src/awkward/contents/numpyarray.py +++ b/src/awkward/contents/numpyarray.py @@ -982,9 +982,7 @@ def _argsort_next( offsets.data, offsets_length, parents.data, - parents_length, starts.data, - starts.length, ) ) out = NumpyArray(nextcarry.data, parameters=None, backend=self._backend) diff --git a/src/awkward/contents/regulararray.py b/src/awkward/contents/regulararray.py index 91e63b3133..2d144a0b7a 100644 --- a/src/awkward/contents/regulararray.py +++ b/src/awkward/contents/regulararray.py @@ -647,7 +647,6 @@ def _getitem_next( advanced.data, regular_flathead.data, self._length, - regular_flathead.length, self._size, ), slicer=head, @@ -996,7 +995,7 @@ def _reduce_next( ) self._backend.maybe_kernel_error( self._backend[ - "awkward_RegularArray_reduce_nonlocal_preparenext", + "awkward_RegularArray_reduce_nonlocal_preparenext_64", nextcarry.dtype.type, nextparents.dtype.type, parents.dtype.type, @@ -1067,7 +1066,7 @@ def _reduce_next( assert nextparents.nplike is index_nplike self._backend.maybe_kernel_error( self._backend[ - "awkward_RegularArray_reduce_local_nextparents", + "awkward_RegularArray_reduce_local_nextparents_64", nextparents.dtype.type, ]( nextparents.data, diff --git a/src/awkward/contents/unionarray.py b/src/awkward/contents/unionarray.py index c29f3de94b..bb3b6ca929 100644 --- a/src/awkward/contents/unionarray.py +++ b/src/awkward/contents/unionarray.py @@ -12,6 +12,7 @@ from awkward._layout import maybe_posaxis from awkward._meta.unionmeta import UnionMeta from awkward._nplikes.array_like import ArrayLike +from awkward._nplikes.cupy import Cupy from awkward._nplikes.numpy import Numpy from awkward._nplikes.numpy_like import IndexType, NumpyMetadata from awkward._nplikes.placeholder import PlaceholderArray @@ -897,23 +898,40 @@ def _offsets_and_flattened(self, axis: int, depth: int) -> tuple[Index, Content] and self._tags.nplike is self._backend.index_nplike and self._index.nplike is self._backend.index_nplike ) - self._backend.maybe_kernel_error( - self._backend[ - "awkward_UnionArray_flatten_length", - total_length.dtype.type, - self._tags.dtype.type, - self._index.dtype.type, - np.int64, - ]( - total_length.data, - self._tags.data, - self._index.data, - self._tags.length, - offsetsraws.ctypes.data_as( - ctypes.POINTER(ctypes.POINTER(ctypes.c_int64)) - ), + if self._backend.nplike == Numpy.instance(): + self._backend.maybe_kernel_error( + self._backend[ + "awkward_UnionArray_flatten_length", + total_length.dtype.type, + self._tags.dtype.type, + self._index.dtype.type, + np.int64, + ]( + total_length.data, + self._tags.data, + self._index.data, + self._tags.length, + offsetsraws.ctypes.data_as( + ctypes.POINTER(ctypes.POINTER(ctypes.c_int64)) + ), + ) + ) + elif self._backend.nplike == Cupy.instance(): + self._backend.maybe_kernel_error( + self._backend[ + "awkward_UnionArray_flatten_length", + total_length.dtype.type, + self._tags.dtype.type, + self._index.dtype.type, + np.int64, + ]( + total_length.data, + self._tags.data, + self._index.data, + self._tags.length, + offsetsraws, + ) ) - ) totags = ak.index.Index8.empty( total_length[0], nplike=self._backend.index_nplike @@ -932,28 +950,48 @@ def _offsets_and_flattened(self, axis: int, depth: int) -> tuple[Index, Content] and self._tags.nplike is self._backend.index_nplike and self._index.nplike is self._backend.index_nplike ) - self._backend.maybe_kernel_error( - self._backend[ - "awkward_UnionArray_flatten_combine", - totags.dtype.type, - toindex.dtype.type, - tooffsets.dtype.type, - self._tags.dtype.type, - self._index.dtype.type, - np.int64, - ]( - totags.data, - toindex.data, - tooffsets.data, - self._tags.data, - self._index.data, - self._tags.length, - offsetsraws.ctypes.data_as( - ctypes.POINTER(ctypes.POINTER(ctypes.c_int64)) - ), + if self._backend.nplike == Numpy.instance(): + self._backend.maybe_kernel_error( + self._backend[ + "awkward_UnionArray_flatten_combine", + totags.dtype.type, + toindex.dtype.type, + tooffsets.dtype.type, + self._tags.dtype.type, + self._index.dtype.type, + np.int64, + ]( + totags.data, + toindex.data, + tooffsets.data, + self._tags.data, + self._index.data, + self._tags.length, + offsetsraws.ctypes.data_as( + ctypes.POINTER(ctypes.POINTER(ctypes.c_int64)) + ), + ) + ) + elif self._backend.nplike == Cupy.instance(): + self._backend.maybe_kernel_error( + self._backend[ + "awkward_UnionArray_flatten_combine", + totags.dtype.type, + toindex.dtype.type, + tooffsets.dtype.type, + self._tags.dtype.type, + self._index.dtype.type, + np.int64, + ]( + totags.data, + toindex.data, + tooffsets.data, + self._tags.data, + self._index.data, + self._tags.length, + offsetsraws, + ) ) - ) - return ( tooffsets, UnionArray( diff --git a/src/awkward/index.py b/src/awkward/index.py index 8f13e9d204..3e6f879cfe 100644 --- a/src/awkward/index.py +++ b/src/awkward/index.py @@ -136,7 +136,10 @@ def metadata(self) -> dict: @property def ptr(self): - return self._data.ctypes.data + if self._nplike == Numpy.instance(): + return self._data.ctypes.data + elif self._nplike == Cupy.instance(): + return self._data.data.ptr @property def length(self) -> ShapeItem: diff --git a/src/awkward/operations/ak_angle.py b/src/awkward/operations/ak_angle.py index 2bf34fa39c..49128bd944 100644 --- a/src/awkward/operations/ak_angle.py +++ b/src/awkward/operations/ak_angle.py @@ -3,7 +3,6 @@ from __future__ import annotations import awkward as ak -from awkward._backends.numpy import NumpyBackend from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext from awkward._nplikes.numpy_like import NumpyMetadata @@ -11,7 +10,6 @@ __all__ = ("angle",) np = NumpyMetadata.instance() -cpu = NumpyBackend.instance() @ak._connect.numpy.implements("angle") diff --git a/src/awkward/operations/ak_argcartesian.py b/src/awkward/operations/ak_argcartesian.py index 44facda474..12deed5749 100644 --- a/src/awkward/operations/ak_argcartesian.py +++ b/src/awkward/operations/ak_argcartesian.py @@ -5,7 +5,6 @@ from collections.abc import Mapping import awkward as ak -from awkward._backends.numpy import NumpyBackend from awkward._dispatch import high_level_function from awkward._nplikes.numpy_like import NumpyMetadata from awkward._regularize import regularize_axis @@ -13,7 +12,6 @@ __all__ = ("argcartesian",) np = NumpyMetadata.instance() -cpu = NumpyBackend.instance() @high_level_function() diff --git a/src/awkward/operations/ak_concatenate.py b/src/awkward/operations/ak_concatenate.py index 42df7bc4f4..f35f1baf22 100644 --- a/src/awkward/operations/ak_concatenate.py +++ b/src/awkward/operations/ak_concatenate.py @@ -6,7 +6,6 @@ import awkward as ak from awkward._backends.dispatch import backend_of_obj -from awkward._backends.numpy import NumpyBackend from awkward._dispatch import high_level_function from awkward._do import mergeable from awkward._layout import HighLevelContext, ensure_same_backend, maybe_posaxis @@ -22,7 +21,6 @@ __all__ = ("concatenate",) np = NumpyMetadata.instance() -cpu = NumpyBackend.instance() @ak._connect.numpy.implements("concatenate") @@ -194,7 +192,8 @@ def action(inputs, depth, backend, **kwargs): nextinputs = [] for x in inputs: if x.is_option and x.content.is_list: - nextinputs.append(fill_none(x, [], axis=0, highlevel=False)) + empty = ak.to_backend([], backend) + nextinputs.append(fill_none(x, empty, axis=0, highlevel=False)) else: nextinputs.append(x) inputs = nextinputs diff --git a/src/awkward/operations/ak_fill_none.py b/src/awkward/operations/ak_fill_none.py index 739bed11a7..89834689cd 100644 --- a/src/awkward/operations/ak_fill_none.py +++ b/src/awkward/operations/ak_fill_none.py @@ -3,7 +3,6 @@ from __future__ import annotations import awkward as ak -from awkward._backends.numpy import NumpyBackend from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext, ensure_same_backend, maybe_posaxis from awkward._nplikes.numpy_like import NumpyMetadata @@ -13,7 +12,6 @@ __all__ = ("fill_none",) np = NumpyMetadata.instance() -cpu = NumpyBackend.instance() @high_level_function() diff --git a/src/awkward/operations/ak_imag.py b/src/awkward/operations/ak_imag.py index 0bc4836e54..b90d2dbfa4 100644 --- a/src/awkward/operations/ak_imag.py +++ b/src/awkward/operations/ak_imag.py @@ -3,7 +3,6 @@ from __future__ import annotations import awkward as ak -from awkward._backends.numpy import NumpyBackend from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext from awkward._nplikes.numpy_like import NumpyMetadata @@ -11,7 +10,6 @@ __all__ = ("imag",) np = NumpyMetadata.instance() -cpu = NumpyBackend.instance() @ak._connect.numpy.implements("imag") diff --git a/src/awkward/operations/ak_linear_fit.py b/src/awkward/operations/ak_linear_fit.py index dc5234a7bf..971fea64fe 100644 --- a/src/awkward/operations/ak_linear_fit.py +++ b/src/awkward/operations/ak_linear_fit.py @@ -3,7 +3,6 @@ from __future__ import annotations import awkward as ak -from awkward._backends.numpy import NumpyBackend from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext, ensure_same_backend from awkward._nplikes import ufuncs @@ -12,7 +11,6 @@ __all__ = ("linear_fit",) -cpu = NumpyBackend.instance() np = NumpyMetadata.instance() diff --git a/src/awkward/operations/ak_merge_option_of_records.py b/src/awkward/operations/ak_merge_option_of_records.py index 80b6dfdfb2..c3e1095ba4 100644 --- a/src/awkward/operations/ak_merge_option_of_records.py +++ b/src/awkward/operations/ak_merge_option_of_records.py @@ -3,7 +3,6 @@ from __future__ import annotations import awkward as ak -from awkward._backends.numpy import NumpyBackend from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext, maybe_posaxis from awkward._nplikes.numpy_like import NumpyMetadata @@ -13,7 +12,6 @@ __all__ = ("merge_option_of_records",) np = NumpyMetadata.instance() -cpu = NumpyBackend.instance() @high_level_function() diff --git a/src/awkward/operations/ak_merge_union_of_records.py b/src/awkward/operations/ak_merge_union_of_records.py index 1418c747b6..d523c0b5f8 100644 --- a/src/awkward/operations/ak_merge_union_of_records.py +++ b/src/awkward/operations/ak_merge_union_of_records.py @@ -3,7 +3,6 @@ from __future__ import annotations import awkward as ak -from awkward._backends.numpy import NumpyBackend from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext, maybe_posaxis from awkward._nplikes.numpy_like import ArrayLike, NumpyMetadata @@ -13,7 +12,6 @@ __all__ = ("merge_union_of_records",) np = NumpyMetadata.instance() -cpu = NumpyBackend.instance() @high_level_function() diff --git a/src/awkward/operations/ak_real.py b/src/awkward/operations/ak_real.py index 2c49c98c42..6d52971dab 100644 --- a/src/awkward/operations/ak_real.py +++ b/src/awkward/operations/ak_real.py @@ -3,7 +3,6 @@ from __future__ import annotations import awkward as ak -from awkward._backends.numpy import NumpyBackend from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext from awkward._nplikes.numpy_like import NumpyMetadata @@ -11,7 +10,6 @@ __all__ = ("real",) np = NumpyMetadata.instance() -cpu = NumpyBackend.instance() @ak._connect.numpy.implements("real") diff --git a/src/awkward/operations/ak_run_lengths.py b/src/awkward/operations/ak_run_lengths.py index ba5702ea54..673627fdbc 100644 --- a/src/awkward/operations/ak_run_lengths.py +++ b/src/awkward/operations/ak_run_lengths.py @@ -3,7 +3,6 @@ from __future__ import annotations import awkward as ak -from awkward._backends.numpy import NumpyBackend from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext from awkward._nplikes.numpy_like import NumpyMetadata @@ -12,7 +11,6 @@ __all__ = ("run_lengths",) np = NumpyMetadata.instance() -cpu = NumpyBackend.instance() @high_level_function() diff --git a/src/awkward/operations/ak_where.py b/src/awkward/operations/ak_where.py index 5cd80cd3f2..1d2312a47f 100644 --- a/src/awkward/operations/ak_where.py +++ b/src/awkward/operations/ak_where.py @@ -3,7 +3,6 @@ from __future__ import annotations import awkward as ak -from awkward._backends.numpy import NumpyBackend from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext, ensure_same_backend from awkward._nplikes.numpy_like import NumpyMetadata @@ -11,7 +10,6 @@ __all__ = ("where",) np = NumpyMetadata.instance() -cpu = NumpyBackend.instance() @ak._connect.numpy.implements("where") diff --git a/src/awkward/operations/ak_with_field.py b/src/awkward/operations/ak_with_field.py index 9998356be6..671a061978 100644 --- a/src/awkward/operations/ak_with_field.py +++ b/src/awkward/operations/ak_with_field.py @@ -5,7 +5,6 @@ import copy import awkward as ak -from awkward._backends.numpy import NumpyBackend from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext, ensure_same_backend from awkward._nplikes.numpy_like import NumpyMetadata @@ -14,7 +13,6 @@ __all__ = ("with_field",) np = NumpyMetadata.instance() -cpu = NumpyBackend.instance() @high_level_function() diff --git a/tests-cuda/test_3086_cuda_concatenate.py b/tests-cuda/test_3086_cuda_concatenate.py new file mode 100644 index 0000000000..0a69de935a --- /dev/null +++ b/tests-cuda/test_3086_cuda_concatenate.py @@ -0,0 +1,877 @@ +from __future__ import annotations + +import cupy as cp +import numpy as np +import pytest + +import awkward as ak +from awkward.types import ArrayType, ListType, NumpyType, OptionType, RegularType + +to_list = ak.operations.to_list + + +def test_0184_concatenate_number(): + a1 = ak.highlevel.Array([[1, 2, 3], [], [4, 5]]).layout + a2 = ak.highlevel.Array([[[1.1], [2.2, 3.3]], [[]], [[4.4], [5.5]]]).layout + a3 = ak.highlevel.Array([[123], [223], [323]]).layout + + cuda_a1 = ak.to_backend(a1, "cuda") + cuda_a2 = ak.to_backend(a2, "cuda") + cuda_a3 = ak.to_backend(a3, "cuda") + + assert to_list(ak.operations.concatenate([cuda_a1, 999], axis=1)) == [ + [1, 2, 3, 999], + [999], + [4, 5, 999], + ] + + assert to_list(ak.operations.concatenate([cuda_a2, 999], axis=2)) == [ + [[1.1, 999.0], [2.2, 3.3, 999.0]], + [[999.0]], + [[4.4, 999.0], [5.5, 999.0]], + ] + + assert ( + str(ak.operations.type(ak.operations.concatenate([cuda_a1, cuda_a3], axis=1))) + == "3 * var * int64" + ) + + assert to_list(ak.operations.concatenate([cuda_a1, cuda_a3], axis=1)) == [ + [1, 2, 3, 123], + [223], + [4, 5, 323], + ] + + +def test_0184_negative_axis_concatenate(): + one = ak.highlevel.Array( + [[[0.0, 1.1, 2.2], []], [[3.3, 4.4]], [[5.5], [6.6, 7.7, 8.8, 9.9]]] + ).layout + two = ak.highlevel.Array( + [[[10, 20], [30]], [[40]], [[50, 60, 70], [80, 90]]] + ).layout + arrays = [one, two] + + cuda_arrays = ak.to_backend(arrays, "cuda") + + assert ak.operations.concatenate(cuda_arrays, axis=-1).to_list() == [ + [[0.0, 1.1, 2.2, 10, 20], [30]], + [[3.3, 4.4, 40]], + [[5.5, 50, 60, 70], [6.6, 7.7, 8.8, 9.9, 80, 90]], + ] + + assert ak.operations.concatenate(cuda_arrays, axis=-2).to_list() == [ + [[0.0, 1.1, 2.2], [], [10, 20], [30]], + [[3.3, 4.4], [40]], + [[5.5], [6.6, 7.7, 8.8, 9.9], [50, 60, 70], [80, 90]], + ] + + assert ak.operations.concatenate(cuda_arrays, axis=-3).to_list() == [ + [[0.0, 1.1, 2.2], []], + [[3.3, 4.4]], + [[5.5], [6.6, 7.7, 8.8, 9.9]], + [[10, 20], [30]], + [[40]], + [[50, 60, 70], [80, 90]], + ] + + +def test_0184_broadcast_and_apply_levels_concatenate(): + arrays = [ + ak.highlevel.Array( + [[[0.0, 1.1, 2.2], []], [[3.3, 4.4]], [[5.5], [6.6, 7.7, 8.8, 9.9]]] + ).layout, + ak.highlevel.Array([[[10, 20], [30]], [[40]], [[50, 60, 70], [80, 90]]]).layout, + ] + + cuda_arrays = ak.to_backend(arrays, "cuda") + + # nothing is required to have the same length + assert ak.operations.concatenate(cuda_arrays, axis=0).to_list() == [ + [[0.0, 1.1, 2.2], []], + [[3.3, 4.4]], + [[5.5], [6.6, 7.7, 8.8, 9.9]], + [[10, 20], [30]], + [[40]], + [[50, 60, 70], [80, 90]], + ] + # the outermost arrays are required to have the same length, but nothing deeper than that + assert ak.operations.concatenate(cuda_arrays, axis=1).to_list() == [ + [[0.0, 1.1, 2.2], [], [10, 20], [30]], + [[3.3, 4.4], [40]], + [[5.5], [6.6, 7.7, 8.8, 9.9], [50, 60, 70], [80, 90]], + ] + # the outermost arrays and the first level are required to have the same length, but nothing deeper + assert ak.operations.concatenate(cuda_arrays, axis=2).to_list() == [ + [[0.0, 1.1, 2.2, 10, 20], [30]], + [[3.3, 4.4, 40]], + [[5.5, 50, 60, 70], [6.6, 7.7, 8.8, 9.9, 80, 90]], + ] + + +def test_0184_list_array_concatenate(): + one = ak.highlevel.Array([[1, 2, 3], [], [4, 5]]).layout + two = ak.highlevel.Array([[1.1, 2.2], [3.3, 4.4], [5.5]]).layout + + one = ak.contents.ListArray(one.starts, one.stops, one.content) + two = ak.contents.ListArray(two.starts, two.stops, two.content) + + cuda_one = ak.to_backend(one, "cuda") + cuda_two = ak.to_backend(two, "cuda") + + assert to_list(ak.operations.concatenate([cuda_one, cuda_two], 0)) == [ + [1, 2, 3], + [], + [4, 5], + [1.1, 2.2], + [3.3, 4.4], + [5.5], + ] + assert to_list(ak.operations.concatenate([cuda_one, cuda_two], 1)) == [ + [1, 2, 3, 1.1, 2.2], + [3.3, 4.4], + [4, 5, 5.5], + ] + + +def test_0184_indexed_array_concatenate(): + one = ak.highlevel.Array([[1, 2, 3], [None, 4], None, [None, 5]]).layout + two = ak.highlevel.Array([6, 7, 8]).layout + three = ak.highlevel.Array([[6.6], [7.7, 8.8]]).layout + four = ak.highlevel.Array([[6.6], [7.7, 8.8], None, [9.9]]).layout + + cuda_one = ak.to_backend(one, "cuda") + cuda_two = ak.to_backend(two, "cuda") + cuda_three = ak.to_backend(three, "cuda") + cuda_four = ak.to_backend(four, "cuda") + + assert to_list(ak.operations.concatenate([cuda_one, cuda_two], 0)) == [ + [1, 2, 3], + [None, 4], + None, + [None, 5], + 6, + 7, + 8, + ] + + with pytest.raises(ValueError): + to_list(ak.operations.concatenate([cuda_one, cuda_three], 1)) + + assert to_list(ak.operations.concatenate([cuda_one, cuda_four], 1)) == [ + [1, 2, 3, 6.6], + [None, 4, 7.7, 8.8], + [], + [None, 5, 9.9], + ] + + +def test_0184_listoffsetarray_concatenate(): + content_one = ak.contents.NumpyArray(np.array([1, 2, 3, 4, 5, 6, 7, 8, 9])) + offsets_one = ak.index.Index64(np.array([0, 3, 3, 5, 9])) + one = ak.contents.ListOffsetArray(offsets_one, content_one) + + cuda_one = ak.to_backend(one, "cuda") + + assert to_list(cuda_one) == [[1, 2, 3], [], [4, 5], [6, 7, 8, 9]] + + content_two = ak.contents.NumpyArray(np.array([100, 200, 300, 400, 500])) + offsets_two = ak.index.Index64(np.array([0, 2, 4, 4, 5])) + two = ak.contents.ListOffsetArray(offsets_two, content_two) + + cuda_two = ak.to_backend(two, "cuda") + + assert to_list(cuda_two) == [[100, 200], [300, 400], [], [500]] + assert to_list(ak.operations.concatenate([cuda_one, cuda_two], 0)) == [ + [1, 2, 3], + [], + [4, 5], + [6, 7, 8, 9], + [100, 200], + [300, 400], + [], + [500], + ] + assert to_list(ak.operations.concatenate([cuda_one, cuda_two], 1)) == [ + [1, 2, 3, 100, 200], + [300, 400], + [4, 5], + [6, 7, 8, 9, 500], + ] + + +def test_0184_even_more(): + dim1 = ak.highlevel.Array([1.1, 2.2, 3.3, 4.4, 5.5]).layout + dim1a = ak.highlevel.Array([[1.1], [2.2], [3.3], [4.4], [5.5]]).layout + dim1b = ak.highlevel.Array(np.array([[1.1], [2.2], [3.3], [4.4], [5.5]])).layout + dim2 = ak.highlevel.Array([[0, 1, 2], [], [3, 4], [5], [6, 7, 8, 9]]).layout + dim3 = ak.highlevel.Array( + [[[0, 1, 2], []], [[3, 4]], [], [[5], [6, 7, 8, 9]], []] + ).layout + + cuda_dim1 = ak.to_backend(dim1, "cuda") + cuda_dim1a = ak.to_backend(dim1a, "cuda") + cuda_dim1b = ak.to_backend(dim1b, "cuda") + cuda_dim2 = ak.to_backend(dim2, "cuda") + cuda_dim3 = ak.to_backend(dim3, "cuda") + + num = cp.array([999]) + + assert ak.operations.concatenate([cuda_dim1, num]).to_list() == [ + 1.1, + 2.2, + 3.3, + 4.4, + 5.5, + 999, + ] + assert ak.operations.concatenate([num, cuda_dim1]).to_list() == [ + 999, + 1.1, + 2.2, + 3.3, + 4.4, + 5.5, + ] + assert ak.operations.concatenate([cuda_dim1, cuda_dim2]).to_list() == [ + 1.1, + 2.2, + 3.3, + 4.4, + 5.5, + [0, 1, 2], + [], + [3, 4], + [5], + [6, 7, 8, 9], + ] + assert ak.operations.concatenate([cuda_dim2, cuda_dim1]).to_list() == [ + [0, 1, 2], + [], + [3, 4], + [5], + [6, 7, 8, 9], + 1.1, + 2.2, + 3.3, + 4.4, + 5.5, + ] + + with pytest.raises(ValueError): + ak.operations.concatenate([cuda_dim1, 999], axis=1) + + assert ak.operations.concatenate([cuda_dim2, 999], axis=1).to_list() == [ + [0, 1, 2, 999], + [999], + [3, 4, 999], + [5, 999], + [6, 7, 8, 9, 999], + ] + + assert ak.operations.concatenate([999, cuda_dim2], axis=1).to_list() == [ + [999, 0, 1, 2], + [999], + [999, 3, 4], + [999, 5], + [999, 6, 7, 8, 9], + ] + + with pytest.raises(ValueError): + ak.operations.concatenate([cuda_dim1, cuda_dim2], axis=1) + with pytest.raises(ValueError): + ak.operations.concatenate([cuda_dim2, cuda_dim1], axis=1) + assert ak.operations.concatenate([cuda_dim1a, cuda_dim2], axis=1).to_list() == [ + [1.1, 0, 1, 2], + [2.2], + [3.3, 3, 4], + [4.4, 5], + [5.5, 6, 7, 8, 9], + ] + assert ak.operations.concatenate([cuda_dim2, cuda_dim1a], axis=1).to_list() == [ + [0, 1, 2, 1.1], + [2.2], + [3, 4, 3.3], + [5, 4.4], + [6, 7, 8, 9, 5.5], + ] + assert ak.operations.concatenate([cuda_dim1b, cuda_dim2], axis=1).to_list() == [ + [1.1, 0, 1, 2], + [2.2], + [3.3, 3, 4], + [4.4, 5], + [5.5, 6, 7, 8, 9], + ] + assert ak.operations.concatenate([cuda_dim2, cuda_dim1b], axis=1).to_list() == [ + [0, 1, 2, 1.1], + [2.2], + [3, 4, 3.3], + [5, 4.4], + [6, 7, 8, 9, 5.5], + ] + num = cp.array([123]) + + assert ak.operations.concatenate([num, cuda_dim1, cuda_dim2]).to_list() == [ + 123, + 1.1, + 2.2, + 3.3, + 4.4, + 5.5, + [0, 1, 2], + [], + [3, 4], + [5], + [6, 7, 8, 9], + ] + assert ak.operations.concatenate([num, cuda_dim2, cuda_dim1]).to_list() == [ + 123, + [0, 1, 2], + [], + [3, 4], + [5], + [6, 7, 8, 9], + 1.1, + 2.2, + 3.3, + 4.4, + 5.5, + ] + assert ak.operations.concatenate([cuda_dim1, num, cuda_dim2]).to_list() == [ + 1.1, + 2.2, + 3.3, + 4.4, + 5.5, + 123, + [0, 1, 2], + [], + [3, 4], + [5], + [6, 7, 8, 9], + ] + assert ak.operations.concatenate([cuda_dim1, cuda_dim2, num]).to_list() == [ + 1.1, + 2.2, + 3.3, + 4.4, + 5.5, + [0, 1, 2], + [], + [3, 4], + [5], + [6, 7, 8, 9], + 123, + ] + assert ak.operations.concatenate([cuda_dim2, num, cuda_dim1]).to_list() == [ + [0, 1, 2], + [], + [3, 4], + [5], + [6, 7, 8, 9], + 123, + 1.1, + 2.2, + 3.3, + 4.4, + 5.5, + ] + assert ak.operations.concatenate([cuda_dim2, cuda_dim1, num]).to_list() == [ + [0, 1, 2], + [], + [3, 4], + [5], + [6, 7, 8, 9], + 1.1, + 2.2, + 3.3, + 4.4, + 5.5, + 123, + ] + + with pytest.raises(ValueError): + ak.operations.concatenate([123, cuda_dim1, cuda_dim2], axis=1) + with pytest.raises(ValueError): + ak.operations.concatenate([123, cuda_dim2, cuda_dim1], axis=1) + with pytest.raises(ValueError): + ak.operations.concatenate([cuda_dim1, 123, cuda_dim2], axis=1) + with pytest.raises(ValueError): + ak.operations.concatenate([cuda_dim1, cuda_dim2, 123], axis=1) + with pytest.raises(ValueError): + ak.operations.concatenate([cuda_dim2, 123, cuda_dim1], axis=1) + with pytest.raises(ValueError): + ak.operations.concatenate([cuda_dim2, cuda_dim1, 123], axis=1) + + assert ak.operations.concatenate( + [123, cuda_dim1a, cuda_dim2], axis=1 + ).to_list() == [ + [123, 1.1, 0, 1, 2], + [123, 2.2], + [123, 3.3, 3, 4], + [123, 4.4, 5], + [123, 5.5, 6, 7, 8, 9], + ] + assert ak.operations.concatenate( + [123, cuda_dim2, cuda_dim1a], axis=1 + ).to_list() == [ + [123, 0, 1, 2, 1.1], + [123, 2.2], + [123, 3, 4, 3.3], + [123, 5, 4.4], + [123, 6, 7, 8, 9, 5.5], + ] + assert ak.operations.concatenate( + [cuda_dim1a, 123, cuda_dim2], axis=1 + ).to_list() == [ + [1.1, 123, 0, 1, 2], + [2.2, 123], + [3.3, 123, 3, 4], + [4.4, 123, 5], + [5.5, 123, 6, 7, 8, 9], + ] + assert ak.operations.concatenate( + [cuda_dim1a, cuda_dim2, 123], axis=1 + ).to_list() == [ + [1.1, 0, 1, 2, 123], + [2.2, 123], + [3.3, 3, 4, 123], + [4.4, 5, 123], + [5.5, 6, 7, 8, 9, 123], + ] + assert ak.operations.concatenate( + [cuda_dim2, 123, cuda_dim1a], axis=1 + ).to_list() == [ + [0, 1, 2, 123, 1.1], + [123, 2.2], + [3, 4, 123, 3.3], + [5, 123, 4.4], + [6, 7, 8, 9, 123, 5.5], + ] + assert ak.operations.concatenate( + [cuda_dim2, cuda_dim1a, 123], axis=1 + ).to_list() == [ + [0, 1, 2, 1.1, 123], + [2.2, 123], + [3, 4, 3.3, 123], + [5, 4.4, 123], + [6, 7, 8, 9, 5.5, 123], + ] + + assert ak.operations.concatenate( + [123, cuda_dim1b, cuda_dim2], axis=1 + ).to_list() == [ + [123, 1.1, 0, 1, 2], + [123, 2.2], + [123, 3.3, 3, 4], + [123, 4.4, 5], + [123, 5.5, 6, 7, 8, 9], + ] + assert ak.operations.concatenate( + [123, cuda_dim2, cuda_dim1b], axis=1 + ).to_list() == [ + [123, 0, 1, 2, 1.1], + [123, 2.2], + [123, 3, 4, 3.3], + [123, 5, 4.4], + [123, 6, 7, 8, 9, 5.5], + ] + assert ak.operations.concatenate( + [cuda_dim1b, 123, cuda_dim2], axis=1 + ).to_list() == [ + [1.1, 123, 0, 1, 2], + [2.2, 123], + [3.3, 123, 3, 4], + [4.4, 123, 5], + [5.5, 123, 6, 7, 8, 9], + ] + assert ak.operations.concatenate( + [cuda_dim1b, cuda_dim2, 123], axis=1 + ).to_list() == [ + [1.1, 0, 1, 2, 123], + [2.2, 123], + [3.3, 3, 4, 123], + [4.4, 5, 123], + [5.5, 6, 7, 8, 9, 123], + ] + assert ak.operations.concatenate( + [cuda_dim2, 123, cuda_dim1b], axis=1 + ).to_list() == [ + [0, 1, 2, 123, 1.1], + [123, 2.2], + [3, 4, 123, 3.3], + [5, 123, 4.4], + [6, 7, 8, 9, 123, 5.5], + ] + assert ak.operations.concatenate( + [cuda_dim2, cuda_dim1b, 123], axis=1 + ).to_list() == [ + [0, 1, 2, 1.1, 123], + [2.2, 123], + [3, 4, 3.3, 123], + [5, 4.4, 123], + [6, 7, 8, 9, 5.5, 123], + ] + + # FIX ME + # assert ak.operations.concatenate([cuda_dim3, 123]).to_list() == [ + # [[0, 1, 2], []], + # [[3, 4]], + # [], + # [[5], [6, 7, 8, 9]], + # [], + # 123, + # ] + + assert ak.operations.concatenate([cuda_dim3, num]).to_list() == [ + [[0, 1, 2], []], + [[3, 4]], + [], + [[5], [6, 7, 8, 9]], + [], + 123, + ] + assert ak.operations.concatenate([num, cuda_dim3]).to_list() == [ + 123, + [[0, 1, 2], []], + [[3, 4]], + [], + [[5], [6, 7, 8, 9]], + [], + ] + + assert ak.operations.concatenate([cuda_dim3, 123], axis=1).to_list() == [ + [[0, 1, 2], [], 123], + [[3, 4], 123], + [123], + [[5], [6, 7, 8, 9], 123], + [123], + ] + assert ak.operations.concatenate([123, cuda_dim3], axis=1).to_list() == [ + [123, [0, 1, 2], []], + [123, [3, 4]], + [123], + [123, [5], [6, 7, 8, 9]], + [123], + ] + + with pytest.raises(ValueError): + ak.operations.concatenate([cuda_dim3, cuda_dim1], axis=1) + with pytest.raises(ValueError): + ak.operations.concatenate([cuda_dim1, cuda_dim3], axis=1) + + assert ak.operations.concatenate([cuda_dim3, cuda_dim2], axis=1).to_list() == [ + [[0, 1, 2], [], 0, 1, 2], + [[3, 4]], + [3, 4], + [[5], [6, 7, 8, 9], 5], + [6, 7, 8, 9], + ] + assert ak.operations.concatenate([cuda_dim2, cuda_dim3], axis=1).to_list() == [ + [0, 1, 2, [0, 1, 2], []], + [[3, 4]], + [3, 4], + [5, [5], [6, 7, 8, 9]], + [6, 7, 8, 9], + ] + + assert ak.operations.concatenate([cuda_dim3, 123], axis=2).to_list() == [ + [[0, 1, 2, 123], [123]], + [[3, 4, 123]], + [], + [[5, 123], [6, 7, 8, 9, 123]], + [], + ] + assert ak.operations.concatenate([123, cuda_dim3], axis=2).to_list() == [ + [[123, 0, 1, 2], [123]], + [[123, 3, 4]], + [], + [[123, 5], [123, 6, 7, 8, 9]], + [], + ] + + assert ak.operations.concatenate([cuda_dim3, cuda_dim3], axis=2).to_list() == [ + [[0, 1, 2, 0, 1, 2], []], + [[3, 4, 3, 4]], + [], + [[5, 5], [6, 7, 8, 9, 6, 7, 8, 9]], + [], + ] + + rec1 = ak.highlevel.Array( + [ + {"x": [1, 2], "y": [1.1]}, + {"x": [], "y": [2.2, 3.3]}, + {"x": [3], "y": []}, + {"x": [5, 6, 7], "y": []}, + {"x": [8, 9], "y": [4.4, 5.5]}, + ] + ).layout + rec2 = ak.highlevel.Array( + [ + {"x": [100], "y": [10, 20]}, + {"x": [200], "y": []}, + {"x": [300, 400], "y": [30]}, + {"x": [], "y": [40, 50]}, + {"x": [400, 500], "y": [60]}, + ] + ).layout + + assert ak.operations.concatenate([rec1, rec2]).to_list() == [ + {"x": [1, 2], "y": [1.1]}, + {"x": [], "y": [2.2, 3.3]}, + {"x": [3], "y": []}, + {"x": [5, 6, 7], "y": []}, + {"x": [8, 9], "y": [4.4, 5.5]}, + {"x": [100], "y": [10, 20]}, + {"x": [200], "y": []}, + {"x": [300, 400], "y": [30]}, + {"x": [], "y": [40, 50]}, + {"x": [400, 500], "y": [60]}, + ] + + assert ak.operations.concatenate([rec1, rec2], axis=1).to_list() == [ + {"x": [1, 2, 100], "y": [1.1, 10, 20]}, + {"x": [200], "y": [2.2, 3.3]}, + {"x": [3, 300, 400], "y": [30]}, + {"x": [5, 6, 7], "y": [40, 50]}, + {"x": [8, 9, 400, 500], "y": [4.4, 5.5, 60]}, + ] + + +def test_1586_numpy_regular(): + a1 = ak.Array(np.array([[0.0, 1.1], [2.2, 3.3]])) + a2 = ak.from_json("[[4.4, 5.5], [6.6, 7.7], [8.8, 9.9]]") + cuda_a1 = ak.to_backend(a1, "cuda") + cuda_a2 = ak.to_backend(a2, "cuda") + + assert isinstance(cuda_a1.layout, ak.contents.NumpyArray) + + cuda_a2 = ak.to_regular(cuda_a2, axis=1) + + cuda_c = ak.concatenate([cuda_a1, cuda_a2]) + + assert cuda_c.to_list() == [ + [0.0, 1.1], + [2.2, 3.3], + [4.4, 5.5], + [6.6, 7.7], + [8.8, 9.9], + ] + assert cuda_c.type == ArrayType(RegularType(NumpyType("float64"), 2), 5) + + +def test_1586_regular_option(): + a1 = ak.from_json("[[0.0, 1.1], [2.2, 3.3]]") + a2 = ak.from_json("[[4.4, 5.5], [6.6, 7.7], null, [8.8, 9.9]]") + cuda_a1 = ak.to_backend(a1, "cuda") + cuda_a2 = ak.to_backend(a2, "cuda") + + cuda_a1 = ak.to_regular(cuda_a1, axis=1) + cuda_a2 = ak.to_regular(cuda_a2, axis=1) + cuda_c = ak.concatenate([cuda_a1, cuda_a2]) + + assert cuda_c.to_list() == [ + [0.0, 1.1], + [2.2, 3.3], + [4.4, 5.5], + [6.6, 7.7], + None, + [8.8, 9.9], + ] + assert cuda_c.type == ArrayType(OptionType(RegularType(NumpyType("float64"), 2)), 6) + + +def test_1586_regular_regular(): + a1 = ak.from_json("[[0.0, 1.1], [2.2, 3.3]]") + a2 = ak.from_json("[[4.4, 5.5], [6.6, 7.7], [8.8, 9.9]]") + cuda_a1 = ak.to_backend(a1, "cuda") + cuda_a2 = ak.to_backend(a2, "cuda") + + cuda_a1 = ak.to_regular(cuda_a1, axis=1) + cuda_a2 = ak.to_regular(cuda_a2, axis=1) + + cuda_c = ak.concatenate([cuda_a1, cuda_a2]) + + assert cuda_c.to_list() == [ + [0.0, 1.1], + [2.2, 3.3], + [4.4, 5.5], + [6.6, 7.7], + [8.8, 9.9], + ] + assert cuda_c.type == ArrayType(RegularType(NumpyType("float64"), 2), 5) + + +def test_1586_option_option(): + a1 = ak.from_json("[[0.0, 1.1], null, [2.2, 3.3]]") + a2 = ak.from_json("[[4.4, 5.5], [6.6, 7.7], null, [8.8, 9.9]]") + cuda_a1 = ak.to_backend(a1, "cuda") + cuda_a2 = ak.to_backend(a2, "cuda") + + cuda_a1 = ak.to_regular(cuda_a1, axis=1) + cuda_a2 = ak.to_regular(cuda_a2, axis=1) + cuda_c = ak.concatenate([cuda_a1, cuda_a2]) + + assert cuda_c.to_list() == [ + [0.0, 1.1], + None, + [2.2, 3.3], + [4.4, 5.5], + [6.6, 7.7], + None, + [8.8, 9.9], + ] + assert cuda_c.type == ArrayType(OptionType(RegularType(NumpyType("float64"), 2)), 7) + + +def test_1586_option_option_axis1(): + a1 = ak.from_json("[[0.0, 1.1], null, [2.2, 3.3]]") + a2 = ak.from_json("[[4.4, 5.5, 6.6], null, [7.7, 8.8, 9.9]]") + cuda_a1 = ak.to_backend(a1, "cuda") + cuda_a2 = ak.to_backend(a2, "cuda") + + cuda_a1 = ak.to_regular(cuda_a1, axis=1) + cuda_a2 = ak.to_regular(cuda_a2, axis=1) + cuda_c = ak.concatenate([cuda_a1, cuda_a2], axis=1) + assert cuda_c.to_list() == [ + [0.0, 1.1, 4.4, 5.5, 6.6], + [], + [2.2, 3.3, 7.7, 8.8, 9.9], + ] + assert cuda_c.type == ArrayType(ListType(NumpyType("float64")), 3) + + +def test_1586_regular_option_axis1(): + a1 = ak.from_json("[[0.0, 1.1], [7, 8], [2.2, 3.3]]") + a2 = ak.from_json("[[4.4, 5.5, 6.6], null, [7.7, 8.8, 9.9]]") + cuda_a1 = ak.to_backend(a1, "cuda") + cuda_a2 = ak.to_backend(a2, "cuda") + + cuda_a1 = ak.to_regular(cuda_a1, axis=1) + cuda_a2 = ak.to_regular(cuda_a2, axis=1) + cuda_c = ak.concatenate([cuda_a1, cuda_a2], axis=1) + assert cuda_c.to_list() == [ + [0.0, 1.1, 4.4, 5.5, 6.6], + [7, 8], + [2.2, 3.3, 7.7, 8.8, 9.9], + ] + assert cuda_c.type == ArrayType(ListType(NumpyType("float64")), 3) + + +def test_1586_option_regular(): + a1 = ak.from_json("[[0.0, 1.1], null, [2.2, 3.3]]") + a2 = ak.from_json("[[4.4, 5.5], [6.6, 7.7], [8.8, 9.9]]") + cuda_a1 = ak.to_backend(a1, "cuda") + cuda_a2 = ak.to_backend(a2, "cuda") + + cuda_a1 = ak.to_regular(cuda_a1, axis=1) + cuda_a2 = ak.to_regular(cuda_a2, axis=1) + cuda_c = ak.concatenate([cuda_a1, cuda_a2]) + assert cuda_c.to_list() == [ + [0.0, 1.1], + None, + [2.2, 3.3], + [4.4, 5.5], + [6.6, 7.7], + [8.8, 9.9], + ] + assert cuda_c.type == ArrayType(OptionType(RegularType(NumpyType("float64"), 2)), 6) + + +def test_1586_option_regular_axis1(): + a1 = ak.from_json("[[0.0, 1.1], null, [2.2, 3.3]]") + a2 = ak.from_json("[[4.4, 5.5, 6.6], [7, 8, 9], [7.7, 8.8, 9.9]]") + cuda_a1 = ak.to_backend(a1, "cuda") + cuda_a2 = ak.to_backend(a2, "cuda") + + cuda_a1 = ak.to_regular(cuda_a1, axis=1) + cuda_a2 = ak.to_regular(cuda_a2, axis=1) + cuda_c = ak.concatenate([cuda_a1, cuda_a2], axis=1) + assert cuda_c.to_list() == [ + [0.0, 1.1, 4.4, 5.5, 6.6], + [7, 8, 9], + [2.2, 3.3, 7.7, 8.8, 9.9], + ] + assert cuda_c.type == ArrayType(ListType(NumpyType("float64")), 3) + + +def test_1586_regular_numpy(): + a1 = ak.from_json("[[0.0, 1.1], [2.2, 3.3]]") + a2 = ak.Array(np.array([[4.4, 5.5], [6.6, 7.7], [8.8, 9.9]])) + cuda_a1 = ak.to_backend(a1, "cuda") + cuda_a2 = ak.to_backend(a2, "cuda") + + cuda_a1 = ak.to_regular(cuda_a1, axis=1) + assert isinstance(cuda_a2.layout, ak.contents.NumpyArray) + cuda_c = ak.concatenate([cuda_a1, cuda_a2]) + assert cuda_c.to_list() == [ + [0.0, 1.1], + [2.2, 3.3], + [4.4, 5.5], + [6.6, 7.7], + [8.8, 9.9], + ] + assert cuda_c.type == ArrayType(RegularType(NumpyType("float64"), 2), 5) + + +def test_2663_broadcast_tuples_record_record(): + record_1 = ak.Array( + [ + [ + {"0": [1, 5, 1], "1": [2, 5, 1]}, + {"0": [3, 5, 1], "1": [4, 5, 1]}, + ] + ] + ) + record_2 = ak.Array( + [ + [ + {"0": [1, 5, 1], "1": [9, 10, 11]}, + {"0": [6, 7, 8], "1": [4, 5, 1]}, + ] + ] + ) + cuda_record_1 = ak.to_backend(record_1, "cuda") + cuda_record_2 = ak.to_backend(record_2, "cuda") + cuda_result = ak.concatenate([cuda_record_1, cuda_record_2], axis=-1) + + assert ak.almost_equal( + to_list(cuda_result), + [ + [ + {"0": [1, 5, 1, 1, 5, 1], "1": [2, 5, 1, 9, 10, 11]}, + {"0": [3, 5, 1, 6, 7, 8], "1": [4, 5, 1, 4, 5, 1]}, + ] + ], + ) + + +def test_2663_broadcast_tuples_tuple_tuple(): + tuple_1 = ak.Array( + [ + [ + ([1, 5, 1], [2, 5, 1]), + ([3, 5, 1], [4, 5, 1]), + ] + ] + ) + tuple_2 = ak.Array( + [ + [ + ([1, 5, 1], [9, 10, 11]), + ([6, 7, 8], [4, 5, 1]), + ] + ] + ) + + result = ak.concatenate([tuple_1, tuple_2], axis=-1) + assert ak.almost_equal( + result, + [ + [ + ([1, 5, 1, 1, 5, 1], [2, 5, 1, 9, 10, 11]), + ([3, 5, 1, 6, 7, 8], [4, 5, 1, 4, 5, 1]), + ] + ], + ) diff --git a/tests-cuda/test_3086_cuda_flatten.py b/tests-cuda/test_3086_cuda_flatten.py new file mode 100644 index 0000000000..d0542693ce --- /dev/null +++ b/tests-cuda/test_3086_cuda_flatten.py @@ -0,0 +1,476 @@ +from __future__ import annotations + +import numpy as np + +import awkward as ak + +to_list = ak.operations.to_list + + +def test_0724(): + a = ak.contents.NumpyArray(np.empty(0)) + idx = ak.index.Index64([]) + a = ak.contents.IndexedOptionArray(idx, a) + idx = ak.index.Index64([0]) + a = ak.contents.ListOffsetArray(idx, a) + idx = ak.index.Index64([175990832]) + a = ak.contents.ListOffsetArray(idx, a) + + assert ak.operations.flatten(a, axis=2).to_list() == [] + assert str(ak.operations.flatten(a, axis=2).type) == "0 * var * ?float64" + + +def test_ByteMaskedArray_flatten(): + content = ak.operations.from_iter( + [ + [[0.0, 1.1, 2.2], [], [3.3, 4.4]], + [], + [[5.5]], + [[6.6, 7.7, 8.8, 9.9]], + [[], [10.0, 11.1, 12.2]], + ], + highlevel=False, + ) + mask = ak.index.Index8(np.array([0, 0, 1, 1, 0], dtype=np.int8)) + array = ak.contents.ByteMaskedArray(mask, content, valid_when=False) + + assert to_list(array) == [ + [[0.0, 1.1, 2.2], [], [3.3, 4.4]], + [], + None, + None, + [[], [10.0, 11.1, 12.2]], + ] + assert ak.operations.to_list(ak.operations.flatten(array, axis=1)) == [ + [0.0, 1.1, 2.2], + [], + [3.3, 4.4], + [], + [10.0, 11.1, 12.2], + ] + assert ak.operations.to_list(ak.operations.flatten(array, axis=-2)) == [ + [0.0, 1.1, 2.2], + [], + [3.3, 4.4], + [], + [10.0, 11.1, 12.2], + ] + assert ak.operations.to_list(ak.operations.flatten(array, axis=2)) == [ + [0.0, 1.1, 2.2, 3.3, 4.4], + [], + None, + None, + [10.0, 11.1, 12.2], + ] + assert ak.operations.to_list(ak.operations.flatten(array, axis=-1)) == [ + [0.0, 1.1, 2.2, 3.3, 4.4], + [], + None, + None, + [10.0, 11.1, 12.2], + ] + + +def test_flatten_ListOffsetArray(): + array = ak.highlevel.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]]).layout + + assert ak.operations.to_list(ak.operations.flatten(array)) == [ + 1.1, + 2.2, + 3.3, + 4.4, + 5.5, + ] + assert ak.operations.to_list(ak.operations.flatten(array[1:])) == [4.4, 5.5] + + array = ak.highlevel.Array( + [[[0.0, 1.1, 2.2], [], [3.3, 4.4]], [], [[5.5]], [[], [6.6, 7.7, 8.8, 9.9]]] + ).layout + + assert ak.operations.to_list(ak.operations.flatten(array)) == [ + [0.0, 1.1, 2.2], + [], + [3.3, 4.4], + [5.5], + [], + [6.6, 7.7, 8.8, 9.9], + ] + assert ak.operations.to_list(ak.operations.flatten(array[1:])) == [ + [5.5], + [], + [6.6, 7.7, 8.8, 9.9], + ] + assert ak.operations.to_list(ak.operations.flatten(array[:, 1:])) == [ + [], + [3.3, 4.4], + [6.6, 7.7, 8.8, 9.9], + ] + assert ak.operations.to_list(ak.operations.flatten(array, axis=2)) == [ + [0.0, 1.1, 2.2, 3.3, 4.4], + [], + [5.5], + [6.6, 7.7, 8.8, 9.9], + ] + assert ak.operations.to_list(ak.operations.flatten(array[1:], axis=2)) == [ + [], + [5.5], + [6.6, 7.7, 8.8, 9.9], + ] + assert ak.operations.to_list(ak.operations.flatten(array[:, 1:], axis=2)) == [ + [3.3, 4.4], + [], + [], + [6.6, 7.7, 8.8, 9.9], + ] + + array = ak.highlevel.Array( + np.arange(2 * 3 * 5 * 7).reshape(2, 3, 5, 7).tolist() + ).layout + + assert ( + ak.operations.to_list(ak.operations.flatten(array, axis=1)) + == np.arange(2 * 3 * 5 * 7).reshape(2 * 3, 5, 7).tolist() + ) + assert ( + ak.operations.to_list(ak.operations.flatten(array, axis=2)) + == np.arange(2 * 3 * 5 * 7).reshape(2, 3 * 5, 7).tolist() + ) + assert ( + ak.operations.to_list(ak.operations.flatten(array, axis=3)) + == np.arange(2 * 3 * 5 * 7).reshape(2, 3, 5 * 7).tolist() + ) + + array = ak.highlevel.Array( + ak.operations.from_iter( + np.arange(2 * 3 * 5 * 7).reshape(2, 3, 5, 7).tolist(), highlevel=False + ) + ).layout + + assert ( + ak.operations.to_list(ak.operations.flatten(array, axis=1)) + == np.arange(2 * 3 * 5 * 7).reshape(2 * 3, 5, 7).tolist() + ) + assert ( + ak.operations.to_list(ak.operations.flatten(array, axis=2)) + == np.arange(2 * 3 * 5 * 7).reshape(2, 3 * 5, 7).tolist() + ) + assert ( + ak.operations.to_list(ak.operations.flatten(array, axis=3)) + == np.arange(2 * 3 * 5 * 7).reshape(2, 3, 5 * 7).tolist() + ) + + array = ak.highlevel.Array(np.arange(2 * 3 * 5 * 7).reshape(2, 3, 5, 7)).layout + + assert ( + ak.operations.to_list(ak.operations.flatten(array, axis=1)) + == np.arange(2 * 3 * 5 * 7).reshape(2 * 3, 5, 7).tolist() + ) + assert ( + ak.operations.to_list(ak.operations.flatten(array, axis=2)) + == np.arange(2 * 3 * 5 * 7).reshape(2, 3 * 5, 7).tolist() + ) + assert ( + ak.operations.to_list(ak.operations.flatten(array, axis=3)) + == np.arange(2 * 3 * 5 * 7).reshape(2, 3, 5 * 7).tolist() + ) + + +def test_flatten_IndexedArray(): + array = ak.highlevel.Array( + [[1.1, 2.2, None, 3.3], None, [], None, [4.4, 5.5], None] + ).layout + + assert ak.operations.to_list(ak.operations.flatten(array)) == [ + 1.1, + 2.2, + None, + 3.3, + 4.4, + 5.5, + ] + assert ak.operations.to_list(ak.operations.flatten(array[1:])) == [4.4, 5.5] + + array = ak.highlevel.Array( + [ + [[0.0, 1.1, 2.2], None, None, [3.3, 4.4]], + [], + [[5.5]], + [[], [6.6, 7.7, 8.8, 9.9]], + ] + ).layout + + assert ak.operations.to_list(ak.operations.flatten(array, axis=2)) == [ + [0.0, 1.1, 2.2, 3.3, 4.4], + [], + [5.5], + [6.6, 7.7, 8.8, 9.9], + ] + assert ak.operations.to_list(ak.operations.flatten(array[1:], axis=2)) == [ + [], + [5.5], + [6.6, 7.7, 8.8, 9.9], + ] + assert ak.operations.to_list(ak.operations.flatten(array[:, 1:], axis=2)) == [ + [3.3, 4.4], + [], + [], + [6.6, 7.7, 8.8, 9.9], + ] + + array = ak.highlevel.Array( + [ + [[0.0, 1.1, 2.2], [3.3, 4.4]], + [], + [[5.5]], + None, + None, + [[], [6.6, 7.7, 8.8, 9.9]], + ] + ).layout + + assert ak.operations.to_list(ak.operations.flatten(array, axis=2)) == [ + [0.0, 1.1, 2.2, 3.3, 4.4], + [], + [5.5], + None, + None, + [6.6, 7.7, 8.8, 9.9], + ] + assert ak.operations.to_list(ak.operations.flatten(array[1:], axis=2)) == [ + [], + [5.5], + None, + None, + [6.6, 7.7, 8.8, 9.9], + ] + assert ak.operations.to_list(ak.operations.flatten(array[:, 1:], axis=2)) == [ + [3.3, 4.4], + [], + [], + None, + None, + [6.6, 7.7, 8.8, 9.9], + ] + + array = ak.highlevel.Array( + [ + [[0.0, 1.1, None, 2.2], None, [], None, [3.3, 4.4]], + None, + [], + [[5.5]], + None, + [[], [6.6, None, 7.7, 8.8, 9.9], None], + ] + ).layout + + assert ak.operations.to_list(ak.operations.flatten(array)) == [ + [0.0, 1.1, None, 2.2], + None, + [], + None, + [3.3, 4.4], + [5.5], + [], + [6.6, None, 7.7, 8.8, 9.9], + None, + ] + assert ak.operations.to_list(ak.operations.flatten(array, axis=2)) == [ + [0.0, 1.1, None, 2.2, 3.3, 4.4], + None, + [], + [5.5], + None, + [6.6, None, 7.7, 8.8, 9.9], + ] + assert ak.operations.to_list(ak.operations.flatten(array[1:], axis=2)) == [ + None, + [], + [5.5], + None, + [6.6, None, 7.7, 8.8, 9.9], + ] + assert ak.operations.to_list(ak.operations.flatten(array[:, 1:], axis=2)) == [ + [3.3, 4.4], + None, + [], + [], + None, + [6.6, None, 7.7, 8.8, 9.9], + ] + + content = 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([2, 1, 0, 3, 3, 4], dtype=np.int64)) + array = ak.contents.IndexedArray(index, content) + + assert to_list(array) == [ + [3.3, 4.4], + [], + [0.0, 1.1, 2.2], + [5.5], + [5.5], + [6.6, 7.7, 8.8, 9.9], + ] + assert ak.operations.to_list(ak.operations.flatten(array)) == [ + 3.3, + 4.4, + 0.0, + 1.1, + 2.2, + 5.5, + 5.5, + 6.6, + 7.7, + 8.8, + 9.9, + ] + + content = 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([2, 2, 1, 0, 3], dtype=np.int64)) + array = ak.contents.IndexedArray(index, content) + + assert to_list(array) == [ + [[5.5]], + [[5.5]], + [], + [[0.0, 1.1, 2.2], [], [3.3, 4.4]], + [[], [6.6, 7.7, 8.8, 9.9]], + ] + assert ak.operations.to_list(ak.operations.flatten(array, axis=2)) == [ + [5.5], + [5.5], + [], + [0.0, 1.1, 2.2, 3.3, 4.4], + [6.6, 7.7, 8.8, 9.9], + ] + + +def test_flatten_RecordArray(): + array = ak.highlevel.Array( + [ + {"x": [], "y": [[3, 3, 3]]}, + {"x": [[1]], "y": [[2, 2]]}, + {"x": [[2], [2]], "y": [[1]]}, + {"x": [[3], [3], [3]], "y": [[]]}, + ] + ).layout + + assert ak.operations.to_list(ak.operations.flatten(array, axis=2)) == [ + {"x": [], "y": [3, 3, 3]}, + {"x": [1], "y": [2, 2]}, + {"x": [2, 2], "y": [1]}, + {"x": [3, 3, 3], "y": []}, + ] + assert ak.operations.to_list(ak.operations.flatten(array[1:], axis=2)) == [ + {"x": [1], "y": [2, 2]}, + {"x": [2, 2], "y": [1]}, + {"x": [3, 3, 3], "y": []}, + ] + assert ak.operations.to_list(ak.operations.flatten(array[:, 1:], axis=2)) == [ + {"x": [], "y": []}, + {"x": [], "y": []}, + {"x": [2], "y": []}, + {"x": [3, 3], "y": []}, + ] + + +def test_flatten_UnionArray(): + content1 = ak.operations.from_iter( + [[1.1], [2.2, 2.2], [3.3, 3.3, 3.3]], highlevel=False + ) + content2 = ak.operations.from_iter( + [[[3, 3, 3], [3, 3, 3], [3, 3, 3]], [[2, 2], [2, 2]], [[1]]], highlevel=False + ) + content3 = ak.operations.from_iter( + [ + [["3", "3", "3"], ["3", "3", "3"], ["3", "3", "3"]], + [["2", "2"], ["2", "2"]], + [["1"]], + ], + highlevel=False, + ) + tags = ak.index.Index8(np.array([0, 1, 0, 1, 0, 1], dtype=np.int8)) + index = ak.index.Index64(np.array([0, 0, 1, 1, 2, 2], dtype=np.int64)) + array = ak.contents.UnionArray(tags, index, [content1, content2]) + array = ak.to_backend(array, "cuda") + + assert to_list(array) == [ + [1.1], + [[3, 3, 3], [3, 3, 3], [3, 3, 3]], + [2.2, 2.2], + [[2, 2], [2, 2]], + [3.3, 3.3, 3.3], + [[1]], + ] + assert to_list(array[1:]) == [ + [[3, 3, 3], [3, 3, 3], [3, 3, 3]], + [2.2, 2.2], + [[2, 2], [2, 2]], + [3.3, 3.3, 3.3], + [[1]], + ] + assert ak.operations.to_list(ak.operations.flatten(array)) == [ + 1.1, + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + 2.2, + 2.2, + [2, 2], + [2, 2], + 3.3, + 3.3, + 3.3, + [1], + ] + assert ak.operations.to_list(ak.operations.flatten(array[1:])) == [ + [3, 3, 3], + [3, 3, 3], + [3, 3, 3], + 2.2, + 2.2, + [2, 2], + [2, 2], + 3.3, + 3.3, + 3.3, + [1], + ] + + array = ak.contents.UnionArray(tags, index, [content2, content3]) + + assert to_list(array) == [ + [[3, 3, 3], [3, 3, 3], [3, 3, 3]], + [["3", "3", "3"], ["3", "3", "3"], ["3", "3", "3"]], + [[2, 2], [2, 2]], + [["2", "2"], ["2", "2"]], + [[1]], + [["1"]], + ] + assert ak.operations.to_list(ak.operations.flatten(array, axis=2)) == [ + [3, 3, 3, 3, 3, 3, 3, 3, 3], + ["3", "3", "3", "3", "3", "3", "3", "3", "3"], + [2, 2, 2, 2], + ["2", "2", "2", "2"], + [1], + ["1"], + ] + assert ak.operations.to_list(ak.operations.flatten(array[1:], axis=2)) == [ + ["3", "3", "3", "3", "3", "3", "3", "3", "3"], + [2, 2, 2, 2], + ["2", "2", "2", "2"], + [1], + ["1"], + ] + assert ak.operations.to_list(ak.operations.flatten(array[:, 1:], axis=2)) == [ + [3, 3, 3, 3, 3, 3], + ["3", "3", "3", "3", "3", "3"], + [2, 2], + ["2", "2"], + [], + [], + ] From 7f1d261da4e7f98769079c78548d05eccebbb56f Mon Sep 17 00:00:00 2001 From: Ianna Osborne Date: Thu, 2 May 2024 19:07:46 +0200 Subject: [PATCH 2/3] fix: intermittent segfault in PyPy (#3089) * debug: add debugging flag * test: debug * fix: allocate memory for 2d arrays in tests * fix: remove legacy parameter from api * debug: exclude some tests that rely on 2d array parameters in cpu kernels * debug: exclude cpu kernel unt tests with 2d arrays as parameters * fix: typo * test: remove coverage for pypy * fix: restore tests * fix: restore cpu kernel unit tests * fix: remove debugger * fix: restore tests * fix: add more tests --------- Co-authored-by: Jim Pivarski --- .DS_Store | Bin 8196 -> 8196 bytes .github/workflows/test.yml | 9 +- .../awkward_UnionArray_flatten_length.cpp | 6 +- .../src/cpu-kernels/awkward_unique_ranges.cpp | 21 ----- .../awkward_unique_ranges_bool.cpp | 1 - awkward-cpp/src/cpu-kernels/kernel-utils.cpp | 2 +- dev/generate-tests.py | 80 +++++++++++++++--- kernel-specification.yml | 15 +--- kernel-test-data.json | 24 ++---- src/awkward/contents/numpyarray.py | 2 - 10 files changed, 87 insertions(+), 73 deletions(-) diff --git a/.DS_Store b/.DS_Store index 40072e236494d6318ae89fda2260a2129e9dd017..f0c6683eef7b709a38f347afd2bb5a398c6cbc52 100644 GIT binary patch delta 97 zcmZp1XmOa}&nUMsU^hRb++-es{K>8Y;v8uy#mPze`8k`*1tb{xnHf?TQW?tNg583X dSU6!k#?AaM<873^c#oFz7wabkn-W_F2hESrr*6&Rr+ FUI6wT5?ufQ diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 47a025756d..a86ff5940a 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -139,16 +139,21 @@ jobs: - name: Test CPU kernels with explicitly defined values run: python -m pytest -vv -rs awkward-cpp/tests-cpu-kernels-explicit - - name: Test non-kernels + - name: Test non-kernels (Python) run: >- python -m pytest -vv -rs tests --cov=awkward --cov-report=term --cov-report=xml + if: startsWith(matrix.python-version, '3.') + + - name: Test non-kernels (PyPy) + run: >- + python -m pytest -vv -rs tests + if: startsWith(matrix.python-version, 'pypy') - name: Upload Codecov results if: (matrix.python-version == '3.9') && (matrix.runs-on == 'ubuntu-latest') uses: codecov/codecov-action@v4 - Linux-ROOT: runs-on: ubuntu-20.04 diff --git a/awkward-cpp/src/cpu-kernels/awkward_UnionArray_flatten_length.cpp b/awkward-cpp/src/cpu-kernels/awkward_UnionArray_flatten_length.cpp index 1afe377de1..ae49efde44 100644 --- a/awkward-cpp/src/cpu-kernels/awkward_UnionArray_flatten_length.cpp +++ b/awkward-cpp/src/cpu-kernels/awkward_UnionArray_flatten_length.cpp @@ -13,11 +13,11 @@ ERROR awkward_UnionArray_flatten_length( T** offsetsraws) { *total_length = 0; for (int64_t i = 0; i < length; i++) { - FROMTAGS tag = fromtags[i]; - FROMINDEX idx = fromindex[i]; + int64_t tag = (int64_t)fromtags[i]; + int64_t idx = (int64_t)fromindex[i]; T start = offsetsraws[tag][idx]; T stop = offsetsraws[tag][idx + 1]; - *total_length = *total_length + (stop - start); + *total_length = *total_length + (int64_t)(stop - start); } return success(); } diff --git a/awkward-cpp/src/cpu-kernels/awkward_unique_ranges.cpp b/awkward-cpp/src/cpu-kernels/awkward_unique_ranges.cpp index 6b8fa08024..a9edf83553 100644 --- a/awkward-cpp/src/cpu-kernels/awkward_unique_ranges.cpp +++ b/awkward-cpp/src/cpu-kernels/awkward_unique_ranges.cpp @@ -7,7 +7,6 @@ template ERROR awkward_unique_ranges( T* toptr, - int64_t /* length */, // FIXME: this argument is not needed const int64_t* fromoffsets, int64_t offsetslength, int64_t* tooffsets) { @@ -28,13 +27,11 @@ ERROR awkward_unique_ranges( ERROR awkward_unique_ranges_int8( int8_t* toptr, - int64_t length, const int64_t* fromoffsets, int64_t offsetslength, int64_t* tooffsets) { return awkward_unique_ranges( toptr, - length, fromoffsets, offsetslength, tooffsets); @@ -42,13 +39,11 @@ ERROR awkward_unique_ranges_int8( ERROR awkward_unique_ranges_uint8( uint8_t* toptr, - int64_t length, const int64_t* fromoffsets, int64_t offsetslength, int64_t* tooffsets) { return awkward_unique_ranges( toptr, - length, fromoffsets, offsetslength, tooffsets); @@ -56,13 +51,11 @@ ERROR awkward_unique_ranges_uint8( ERROR awkward_unique_ranges_int16( int16_t* toptr, - int64_t length, const int64_t* fromoffsets, int64_t offsetslength, int64_t* tooffsets) { return awkward_unique_ranges( toptr, - length, fromoffsets, offsetslength, tooffsets); @@ -70,13 +63,11 @@ ERROR awkward_unique_ranges_int16( ERROR awkward_unique_ranges_uint16( uint16_t* toptr, - int64_t length, const int64_t* fromoffsets, int64_t offsetslength, int64_t* tooffsets) { return awkward_unique_ranges( toptr, - length, fromoffsets, offsetslength, tooffsets); @@ -84,13 +75,11 @@ ERROR awkward_unique_ranges_uint16( ERROR awkward_unique_ranges_int32( int32_t* toptr, - int64_t length, const int64_t* fromoffsets, int64_t offsetslength, int64_t* tooffsets) { return awkward_unique_ranges( toptr, - length, fromoffsets, offsetslength, tooffsets); @@ -98,13 +87,11 @@ ERROR awkward_unique_ranges_int32( ERROR awkward_unique_ranges_uint32( uint32_t* toptr, - int64_t length, const int64_t* fromoffsets, int64_t offsetslength, int64_t* tooffsets) { return awkward_unique_ranges( toptr, - length, fromoffsets, offsetslength, tooffsets); @@ -112,13 +99,11 @@ ERROR awkward_unique_ranges_uint32( ERROR awkward_unique_ranges_int64( int64_t* toptr, - int64_t length, const int64_t* fromoffsets, int64_t offsetslength, int64_t* tooffsets) { return awkward_unique_ranges( toptr, - length, fromoffsets, offsetslength, tooffsets); @@ -126,13 +111,11 @@ ERROR awkward_unique_ranges_int64( ERROR awkward_unique_ranges_uint64( uint64_t* toptr, - int64_t length, const int64_t* fromoffsets, int64_t offsetslength, int64_t* tooffsets) { return awkward_unique_ranges( toptr, - length, fromoffsets, offsetslength, tooffsets); @@ -140,13 +123,11 @@ ERROR awkward_unique_ranges_uint64( ERROR awkward_unique_ranges_float32( float* toptr, - int64_t length, const int64_t* fromoffsets, int64_t offsetslength, int64_t* tooffsets) { return awkward_unique_ranges( toptr, - length, fromoffsets, offsetslength, tooffsets); @@ -154,13 +135,11 @@ ERROR awkward_unique_ranges_float32( ERROR awkward_unique_ranges_float64( double* toptr, - int64_t length, const int64_t* fromoffsets, int64_t offsetslength, int64_t* tooffsets) { return awkward_unique_ranges( toptr, - length, fromoffsets, offsetslength, tooffsets); diff --git a/awkward-cpp/src/cpu-kernels/awkward_unique_ranges_bool.cpp b/awkward-cpp/src/cpu-kernels/awkward_unique_ranges_bool.cpp index 4636338041..6ddc9d8268 100644 --- a/awkward-cpp/src/cpu-kernels/awkward_unique_ranges_bool.cpp +++ b/awkward-cpp/src/cpu-kernels/awkward_unique_ranges_bool.cpp @@ -6,7 +6,6 @@ ERROR awkward_unique_ranges_bool( bool* toptr, - int64_t /* length */, // FIXME: this argument is not needed const int64_t* fromoffsets, int64_t offsetslength, int64_t* tooffsets) { diff --git a/awkward-cpp/src/cpu-kernels/kernel-utils.cpp b/awkward-cpp/src/cpu-kernels/kernel-utils.cpp index 22ca679629..d4b4d07000 100644 --- a/awkward-cpp/src/cpu-kernels/kernel-utils.cpp +++ b/awkward-cpp/src/cpu-kernels/kernel-utils.cpp @@ -60,7 +60,7 @@ void awkward_ListArray_combinations_step( } if (j + 1 == n) { for (int64_t k = 0; k < n; k++) { - tocarry[k][toindex[k]] = fromindex[k]; + tocarry[k][toindex[k]] = (T)fromindex[k]; toindex[k]++; } } diff --git a/dev/generate-tests.py b/dev/generate-tests.py index 99420b7939..37dc859b9a 100644 --- a/dev/generate-tests.py +++ b/dev/generate-tests.py @@ -518,7 +518,10 @@ def gencpukerneltests(specdict): ) f.write( - "import ctypes\nimport pytest\n\nfrom awkward_cpp.cpu_kernels import lib\n\n" + "import ctypes\n" + "import numpy as np\n" + "import pytest\n\n" + "from awkward_cpp.cpu_kernels import lib\n\n" ) num = 1 if spec.tests == []: @@ -547,13 +550,29 @@ def gencpukerneltests(specdict): elif count == 2: f.write( " " * 4 - + f"{arg}_ptr = (ctypes.POINTER(ctypes.c_{typename}) * len({arg}))()\n" + + f"{typename}Ptr = ctypes.POINTER(ctypes.c_{typename})\n" + " " * 4 - + f"for i in range(len({arg})):\n" + + f"{typename}PtrPtr = ctypes.POINTER({typename}Ptr)\n" + + " " * 4 + + f"dim0 = len({arg})\n" + + " " * 4 + + f"dim1 = len({arg}[0])\n" + + " " * 4 + + f"{arg}_np_arr_2d = np.empty([dim0, dim1], dtype=np.{typename})\n" + + " " * 4 + + "for i in range(dim0):\n" + " " * 8 - + f"{arg}_ptr[i] = (ctypes.c_{typename} * len({arg}[i]))(*{arg}[i])\n" + + "for j in range(dim1):\n" + + " " * 12 + + f"{arg}_np_arr_2d[i][j] = {arg}[i][j]\n" + + " " * 4 + + f"{arg}_ct_arr = np.ctypeslib.as_ctypes({arg}_np_arr_2d)\n" + + " " * 4 + + f"{typename}PtrArr = {typename}Ptr * {arg}_ct_arr._length_\n" + " " * 4 - + f"{arg} = {arg}_ptr\n" + + f"{arg}_ct_ptr = ctypes.cast({typename}PtrArr(*(ctypes.cast(row, {typename}Ptr) for row in {arg}_ct_arr)), {typename}PtrPtr)\n" + + " " * 4 + + f"{arg} = {arg}_ct_ptr\n" ) f.write(" " * 4 + "funcC = getattr(lib, '" + spec.name + "')\n") args = "" @@ -627,6 +646,7 @@ def gencpuunittests(specdict): f.write( "import ctypes\n" + "import numpy as np\n" "import pytest\n\n" "from awkward_cpp.cpu_kernels import lib\n\n" ) @@ -662,13 +682,29 @@ def gencpuunittests(specdict): elif count == 2: f.write( " " * 4 - + f"{arg}_ptr = (ctypes.POINTER(ctypes.c_{typename}) * len({arg}))()\n" + + f"{typename}Ptr = ctypes.POINTER(ctypes.c_{typename})\n" + + " " * 4 + + f"{typename}PtrPtr = ctypes.POINTER({typename}Ptr)\n" + + " " * 4 + + f"dim0 = len({arg})\n" + + " " * 4 + + f"dim1 = len({arg}[0])\n" + + " " * 4 + + f"{arg}_np_arr_2d = np.empty([dim0, dim1], dtype=np.{typename})\n" + " " * 4 - + f"for i in range(len({arg})):\n" + + "for i in range(dim0):\n" + " " * 8 - + f"{arg}_ptr[i] = (ctypes.c_{typename} * len({arg}[i]))(*{arg}[i])\n" + + "for j in range(dim1):\n" + + " " * 12 + + f"{arg}_np_arr_2d[i][j] = {arg}[i][j]\n" + " " * 4 - + f"{arg} = {arg}_ptr\n" + + f"{arg}_ct_arr = np.ctypeslib.as_ctypes({arg}_np_arr_2d)\n" + + " " * 4 + + f"{typename}PtrArr = {typename}Ptr * {arg}_ct_arr._length_\n" + + " " * 4 + + f"{arg}_ct_ptr = ctypes.cast({typename}PtrArr(*(ctypes.cast(row, {typename}Ptr) for row in {arg}_ct_arr)), {typename}PtrPtr)\n" + + " " * 4 + + f"{arg} = {arg}_ct_ptr\n" ) for arg, val in test["inputs"].items(): typename = gettype(arg, spec.args) @@ -684,13 +720,29 @@ def gencpuunittests(specdict): elif count == 2: f.write( " " * 4 - + f"{arg}_ptr = (ctypes.POINTER(ctypes.c_{typename}) * len({arg}))()\n" + + f"{typename}Ptr = ctypes.POINTER(ctypes.c_{typename})\n" + + " " * 4 + + f"{typename}PtrPtr = ctypes.POINTER({typename}Ptr)\n" + + " " * 4 + + f"dim0 = len({arg})\n" + " " * 4 - + f"for i in range(len({arg})):\n" + + f"dim1 = len({arg}[0])\n" + + " " * 4 + + f"{arg}_np_arr_2d = np.empty([dim0, dim1], dtype=np.{typename})\n" + + " " * 4 + + "for i in range(dim0):\n" + " " * 8 - + f"{arg}_ptr[i] = (ctypes.c_{typename} * len({arg}[i]))(*{arg}[i])\n" + + "for j in range(dim1):\n" + + " " * 12 + + f"{arg}_np_arr_2d[i][j] = {arg}[i][j]\n" + + " " * 4 + + f"{arg}_ct_arr = np.ctypeslib.as_ctypes({arg}_np_arr_2d)\n" + + " " * 4 + + f"{typename}PtrArr = {typename}Ptr * {arg}_ct_arr._length_\n" + + " " * 4 + + f"{arg}_ct_ptr = ctypes.cast({typename}PtrArr(*(ctypes.cast(row, {typename}Ptr) for row in {arg}_ct_arr)), {typename}PtrPtr)\n" + " " * 4 - + f"{arg} = {arg}_ptr\n" + + f"{arg} = {arg}_ct_ptr\n" ) f.write(" " * 4 + "funcC = getattr(lib, '" + spec.name + "')\n") @@ -1183,7 +1235,7 @@ def genunittests(): os.path.join(CURRENT_DIR, "..", "awkward-cpp", "tests-spec-explicit", func), "w", ) as file: - file.write("import pytest\nimport numpy\nimport kernels\n\n") + file.write("import pytest\n" "import numpy\n" "import kernels\n\n") for test in function["tests"]: num += 1 funcName = "def test_" + function["name"] + "_" + str(num) + "():\n" diff --git a/kernel-specification.yml b/kernel-specification.yml index ad3c5b5265..c58bc4d287 100644 --- a/kernel-specification.yml +++ b/kernel-specification.yml @@ -5182,14 +5182,13 @@ kernels: - name: awkward_unique_ranges_bool args: - {name: toptr, type: "List[bool]", dir: out} - - {name: length, type: "int64_t", dir: in, role: default} - {name: fromoffsets, type: "Const[List[int64_t]]", dir: in, role: default} - {name: offsetslength, type: "int64_t", dir: in, role: default} - {name: tooffsets, type: "List[int64_t]", dir: out, role: default} description: null definition: | def awkward_unique_ranges_bool( - toptr, length, fromoffsets, offsetslength, tooffsets + toptr, fromoffsets, offsetslength, tooffsets ): m = 0 for i in range(offsetslength - 1): @@ -5208,77 +5207,67 @@ kernels: - name: awkward_unique_ranges_int8 args: - {name: toptr, type: "List[int8_t]", dir: out} - - {name: length, type: "int64_t", dir: in, role: default} - {name: fromoffsets, type: "Const[List[int64_t]]", dir: in, role: default} - {name: offsetslength, type: "int64_t", dir: in, role: default} - {name: tooffsets, type: "List[int64_t]", dir: out, role: default} - name: awkward_unique_ranges_int16 args: - {name: toptr, type: "List[int16_t]", dir: out} - - {name: length, type: "int64_t", dir: in, role: default} - {name: fromoffsets, type: "Const[List[int64_t]]", dir: in, role: default} - {name: offsetslength, type: "int64_t", dir: in, role: default} - {name: tooffsets, type: "List[int64_t]", dir: out, role: default} - name: awkward_unique_ranges_int32 args: - {name: toptr, type: "List[int32_t]", dir: out} - - {name: length, type: "int64_t", dir: in, role: default} - {name: fromoffsets, type: "Const[List[int64_t]]", dir: in, role: default} - {name: offsetslength, type: "int64_t", dir: in, role: default} - {name: tooffsets, type: "List[int64_t]", dir: out, role: default} - name: awkward_unique_ranges_int64 args: - {name: toptr, type: "List[int64_t]", dir: out} - - {name: length, type: "int64_t", dir: in, role: default} - {name: fromoffsets, type: "Const[List[int64_t]]", dir: in, role: default} - {name: offsetslength, type: "int64_t", dir: in, role: default} - {name: tooffsets, type: "List[int64_t]", dir: out, role: default} - name: awkward_unique_ranges_uint8 args: - {name: toptr, type: "List[uint8_t]", dir: out} - - {name: length, type: "int64_t", dir: in, role: default} - {name: fromoffsets, type: "Const[List[int64_t]]", dir: in, role: default} - {name: offsetslength, type: "int64_t", dir: in, role: default} - {name: tooffsets, type: "List[int64_t]", dir: out, role: default} - name: awkward_unique_ranges_uint16 args: - {name: toptr, type: "List[uint16_t]", dir: out} - - {name: length, type: "int64_t", dir: in, role: default} - {name: fromoffsets, type: "Const[List[int64_t]]", dir: in, role: default} - {name: offsetslength, type: "int64_t", dir: in, role: default} - {name: tooffsets, type: "List[int64_t]", dir: out, role: default} - name: awkward_unique_ranges_uint32 args: - {name: toptr, type: "List[uint32_t]", dir: out} - - {name: length, type: "int64_t", dir: in, role: default} - {name: fromoffsets, type: "Const[List[int64_t]]", dir: in, role: default} - {name: offsetslength, type: "int64_t", dir: in, role: default} - {name: tooffsets, type: "List[int64_t]", dir: out, role: default} - name: awkward_unique_ranges_uint64 args: - {name: toptr, type: "List[uint64_t]", dir: out} - - {name: length, type: "int64_t", dir: in, role: default} - {name: fromoffsets, type: "Const[List[int64_t]]", dir: in, role: default} - {name: offsetslength, type: "int64_t", dir: in, role: default} - {name: tooffsets, type: "List[int64_t]", dir: out, role: default} - name: awkward_unique_ranges_float32 args: - {name: toptr, type: "List[float]", dir: out} - - {name: length, type: "int64_t", dir: in, role: default} - {name: fromoffsets, type: "Const[List[int64_t]]", dir: in, role: default} - {name: offsetslength, type: "int64_t", dir: in, role: default} - {name: tooffsets, type: "List[int64_t]", dir: out, role: default} - name: awkward_unique_ranges_float64 args: - {name: toptr, type: "List[double]", dir: out} - - {name: length, type: "int64_t", dir: in, role: default} - {name: fromoffsets, type: "Const[List[int64_t]]", dir: in, role: default} - {name: offsetslength, type: "int64_t", dir: in, role: default} - {name: tooffsets, type: "List[int64_t]", dir: out, role: default} description: null definition: | def awkward_unique_ranges( - toptr, length, fromoffsets, offsetslength, tooffsets + toptr, fromoffsets, offsetslength, tooffsets ): m = 0 for i in range(offsetslength - 1): diff --git a/kernel-test-data.json b/kernel-test-data.json index b3e27a323a..fde02211fa 100644 --- a/kernel-test-data.json +++ b/kernel-test-data.json @@ -26365,8 +26365,7 @@ "inputs": { "toptr": [], "fromoffsets": [0], - "offsetslength": 1, - "length": 0 + "offsetslength": 1 }, "outputs": { "toptr": [], @@ -26379,8 +26378,7 @@ "inputs": { "toptr": [1, 2], "fromoffsets": [0, 2], - "offsetslength": 2, - "length": 2 + "offsetslength": 2 }, "outputs": { "toptr": [1, 2], @@ -26393,8 +26391,7 @@ "inputs": { "toptr": [0, 1, 2, 3, 4, 5], "fromoffsets": [0, 3, 5, 6], - "offsetslength": 4, - "length": 6 + "offsetslength": 4 }, "outputs": { "toptr": [0, 1, 2, 3, 4, 5], @@ -26407,8 +26404,7 @@ "inputs": { "toptr": [3, 2, 1, 0], "fromoffsets": [0, 0, 3, 3], - "offsetslength": 4, - "length": 4 + "offsetslength": 4 }, "outputs": { "toptr": [3, 3, 1, 0], @@ -26427,8 +26423,7 @@ "inputs": { "toptr": [], "fromoffsets": [0], - "offsetslength": 1, - "length": 1 + "offsetslength": 1 }, "outputs": { "toptr": [], @@ -26441,8 +26436,7 @@ "inputs": { "toptr": [1, 1], "fromoffsets": [0, 2], - "offsetslength": 2, - "length": 1 + "offsetslength": 2 }, "outputs": { "toptr": [1, 1], @@ -26455,8 +26449,7 @@ "inputs": { "toptr": [0, 1, 1, 1, 1, 1], "fromoffsets": [0, 3, 5, 6], - "offsetslength": 4, - "length": 6 + "offsetslength": 4 }, "outputs": { "toptr": [0, 1, 1, 1, 1, 1], @@ -26469,8 +26462,7 @@ "inputs": { "toptr": [1, 1, 1, 0], "fromoffsets": [0, 0, 3, 3], - "offsetslength": 4, - "length": 4 + "offsetslength": 4 }, "outputs": { "toptr": [1, 1, 0, 0], diff --git a/src/awkward/contents/numpyarray.py b/src/awkward/contents/numpyarray.py index 3cd0bdf113..b76584e11e 100644 --- a/src/awkward/contents/numpyarray.py +++ b/src/awkward/contents/numpyarray.py @@ -831,7 +831,6 @@ def _unique(self, negaxis, starts, parents, outlength): nextoffsets.dtype.type, ]( out, - out.shape[0], offsets.data, offsets.length, nextoffsets.data, @@ -846,7 +845,6 @@ def _unique(self, negaxis, starts, parents, outlength): nextoffsets.dtype.type, ]( out, - out.shape[0], offsets.data, offsets.length, nextoffsets.data, From 082e1b7e117c8912cd1b4beb2b4f5919d696f3cf Mon Sep 17 00:00:00 2001 From: Andrea Zonca Date: Fri, 3 May 2024 12:20:45 -0700 Subject: [PATCH 3/3] fix: is_valid checks in header-only library (#3091) * feat: set an explicit index in Indexed and IndexedOption Useful in particular for categorical builders * test: new tests checking presence of a bug in is_valid * style: pre-commit fixes * test: renamed test file with pull request number * test: use const instead of define * test: use ! instead of not operator (issue with Windows) * test: switch bug_fixed to 1, tests will fail until bug fixed * fix: reimplement is_valid to check max of index_ against len of content * test: print out error message * fix: is_valid in IndexedOption now checks index_ is less than len of content * test: IndexOption should have signed index to accommodate -1 * test: print out error from is_valid * fix: reimplemented is_valid for Union, now checks index is below length of content for each tag * feat: implement max_index tracking in Indexed * fix: set max_index_ also in extend_index * feat: append_valid() now calls append_valid(i) * fix: set max_index_ in extend_valid * feat: implement is_valid with max_index_ * feat: check index of Indexed is >0, if not set max_index_ to max int * feat: custom error message for negative index --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../layout-builder/awkward/LayoutBuilder.h | 93 ++++++++++---- header-only/tests/CMakeLists.txt | 4 + .../test_3091-layout-builder-is-valid.cpp | 116 ++++++++++++++++++ 3 files changed, 191 insertions(+), 22 deletions(-) create mode 100644 header-only/tests/test_3091-layout-builder-is-valid.cpp diff --git a/header-only/layout-builder/awkward/LayoutBuilder.h b/header-only/layout-builder/awkward/LayoutBuilder.h index 03ab0ba0d6..89df74022a 100644 --- a/header-only/layout-builder/awkward/LayoutBuilder.h +++ b/header-only/layout-builder/awkward/LayoutBuilder.h @@ -1121,7 +1121,8 @@ namespace awkward { /// buffer, using `AWKWARD_LAYOUTBUILDER_DEFAULT_OPTIONS` for initializing the buffer. Indexed() : index_( - awkward::GrowableBuffer(AWKWARD_LAYOUTBUILDER_DEFAULT_OPTIONS)) { + awkward::GrowableBuffer(AWKWARD_LAYOUTBUILDER_DEFAULT_OPTIONS)), + max_index_(0) { size_t id = 0; set_id(id); } @@ -1132,7 +1133,8 @@ namespace awkward { /// /// @param options Initial size configuration of a buffer. Indexed(const awkward::BuilderOptions& options) - : index_(awkward::GrowableBuffer(options)) { + : index_(awkward::GrowableBuffer(options)), + max_index_(0) { size_t id = 0; set_id(id); } @@ -1151,6 +1153,19 @@ namespace awkward { return content_; } + /// @brief Inserts an explicit value in the `index` buffer and + /// returns the reference to the builder content. + BUILDER& + append_index(size_t i) noexcept { + index_.append(i); + if (i > max_index_) { + max_index_ = i; + } else if (i < 0) { + max_index_ = UINTMAX_MAX; + } + return content_; + } + /// @brief Inserts `size` number indices in the `index` buffer /// and returns the reference to the builder content. /// @@ -1159,6 +1174,9 @@ namespace awkward { extend_index(size_t size) noexcept { size_t start = content_.length(); size_t stop = start + size; + if (stop - 1 > max_index_) { + max_index_ = stop - 1; + } for (size_t i = start; i < stop; i++) { index_.append(i); } @@ -1189,6 +1207,7 @@ namespace awkward { /// of the builder. void clear() noexcept { + max_index_ = 0; index_.clear(); content_.clear(); } @@ -1211,17 +1230,21 @@ namespace awkward { /// @brief Checks for validity and consistency. bool is_valid(std::string& error) const noexcept { - if (content_.length() != index_.length()) { + + if (max_index_ == UINTMAX_MAX) { std::stringstream out; - out << "Indexed node" << id_ << " has content length " - << content_.length() << " but index has length " << index_.length() - << "\n"; + out << "Indexed node" << id_ << " has a negative index\n"; + error.append(out.str()); + return false; + } else if (max_index_ >= content_.length()) { + std::stringstream out; + out << "Indexed node" << id_ << " has index " << max_index_ + << " but content has length " + << content_.length() << "\n"; error.append(out.str()); - return false; - } else { - return content_.is_valid(error); } + return content_.is_valid(error); } /// @brief Copies and concatenates all the accumulated data in each of the @@ -1290,6 +1313,9 @@ namespace awkward { /// @brief Unique form ID. size_t id_; + + /// @brief Keep track of maximum index value. + size_t max_index_; }; /// @class IndexedOption @@ -1310,7 +1336,8 @@ namespace awkward { IndexedOption() : index_( awkward::GrowableBuffer(AWKWARD_LAYOUTBUILDER_DEFAULT_OPTIONS)), - last_valid_(-1) { + last_valid_(-1), + max_index_(0) { size_t id = 0; set_id(id); } @@ -1322,7 +1349,8 @@ namespace awkward { /// @param options Initial size configuration of a buffer. IndexedOption(const awkward::BuilderOptions& options) : index_(awkward::GrowableBuffer(options)), - last_valid_(-1) { + last_valid_(-1), + max_index_(0) { size_t id = 0; set_id(id); } @@ -1338,7 +1366,18 @@ namespace awkward { BUILDER& append_valid() noexcept { last_valid_ = content_.length(); - index_.append(last_valid_); + return append_valid(last_valid_); + } + + /// @brief Inserts an explicit value in the `index` buffer and + /// returns the reference to the builder content. + BUILDER& + append_valid(size_t i) noexcept { + last_valid_ = content_.length(); + index_.append(i); + if (i > max_index_) { + max_index_ = i; + } return content_; } @@ -1351,6 +1390,9 @@ namespace awkward { size_t start = content_.length(); size_t stop = start + size; last_valid_ = stop - 1; + if (last_valid_ > max_index_) { + max_index_ = last_valid_; + } for (size_t i = start; i < stop; i++) { index_.append(i); } @@ -1411,17 +1453,16 @@ namespace awkward { /// @brief Checks for validity and consistency. bool is_valid(std::string& error) const noexcept { - if (content_.length() != last_valid_ + 1) { + if (max_index_ >= content_.length()) { std::stringstream out; - out << "IndexedOption node" << id_ << " has content length " - << content_.length() << " but last valid index is " << last_valid_ - << "\n"; + out << "IndexedOption node" << id_ << " has index " << max_index_ + << " but content has length " + << content_.length() << "\n"; error.append(out.str()); return false; - } else { - return content_.is_valid(error); } + return content_.is_valid(error); } /// @brief Retrieves the names and sizes (in bytes) of the buffers used @@ -1502,6 +1543,9 @@ namespace awkward { /// @brief Last valid index. size_t last_valid_; + + /// @brief Keep track of maximum index value. + size_t max_index_; }; /// @class Unmasked @@ -2298,11 +2342,16 @@ namespace awkward { auto index_sequence((std::index_sequence_for())); std::vector lengths = content_lengths(index_sequence); - for (size_t tag = 0; tag < contents_count_; tag++) { - if (lengths[tag] != last_valid_index_[tag] + 1) { + std::unique_ptr index_ptr(new INDEX[index_.length()]); + index_.concatenate(index_ptr.get()); + std::unique_ptr tags_ptr(new TAGS[tags_.length()]); + tags_.concatenate(tags_ptr.get()); + for (size_t i = 0; i < index_.length(); i++) { + if (index_ptr.get()[i] < 0 || index_ptr.get()[i] >= lengths[tags_ptr.get()[i]]) { std::stringstream out; - out << "Union node" << id_ << " has content length " << lengths[tag] - << " but index length " << last_valid_index_[tag] << "\n"; + out << "Union node" << id_ << " has index " << index_ptr.get()[i] + << " at position " << i << " but content has length " + << lengths[tags_ptr.get()[i]] << "\n"; error.append(out.str()); return false; diff --git a/header-only/tests/CMakeLists.txt b/header-only/tests/CMakeLists.txt index 3d16722ce7..8fe000acee 100644 --- a/header-only/tests/CMakeLists.txt +++ b/header-only/tests/CMakeLists.txt @@ -14,9 +14,13 @@ endmacro(addtest_nolibs) addtest_nolibs(test_1494-layout-builder test_1494-layout-builder.cpp) addtest_nolibs(test_1542-growable-buffer test_1542-growable-buffer.cpp) addtest_nolibs(test_1560-builder-options test_1560-builder-options.cpp) +addtest_nolibs(test_3091-layout-builder-is-valid + test_3091-layout-builder-is-valid.cpp) target_link_libraries(test_1494-layout-builder PRIVATE awkward::layout-builder) target_link_libraries(test_1542-growable-buffer PRIVATE awkward::growable-buffer) target_link_libraries(test_1560-builder-options PRIVATE awkward::builder-options) +target_link_libraries(test_3091-layout-builder-is-valid + PRIVATE awkward::layout-builder) diff --git a/header-only/tests/test_3091-layout-builder-is-valid.cpp b/header-only/tests/test_3091-layout-builder-is-valid.cpp new file mode 100644 index 0000000000..c18e757539 --- /dev/null +++ b/header-only/tests/test_3091-layout-builder-is-valid.cpp @@ -0,0 +1,116 @@ +// BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE + +#include "awkward/LayoutBuilder.h" + +// if BUG_FIXED is false, the tests confirm the presence of the bugs in is_valid +const bool BUG_FIXED=1; + +template +using IndexedBuilder = awkward::LayoutBuilder::Indexed; + +template +using IndexedOptionBuilder = awkward::LayoutBuilder::IndexedOption; + +template +using StringBuilder = awkward::LayoutBuilder::String; + +void +test_Indexed_categorical() { + IndexedBuilder> builder; + assert(builder.length() == 0); + builder.set_parameters("\"__array__\": \"categorical\""); + + auto& subbuilder = builder.append_index(0); + builder.append_index(1); + + subbuilder.append("zero"); + subbuilder.append("one"); + + std::string error; + assert(builder.is_valid(error) == true); + + // index and content could have different lengths + builder.append_index(1); + assert(builder.is_valid(error) == BUG_FIXED); +} + +void +test_Indexed_categorical_invalid_index() { + IndexedBuilder> builder; + assert(builder.length() == 0); + builder.set_parameters("\"__array__\": \"categorical\""); + + auto& subbuilder = builder.append_index(0); + builder.append_index(1); + + subbuilder.append("zero"); + subbuilder.append("one"); + + std::string error; + assert(builder.is_valid(error) == true); + + // index should be less than the length of content + subbuilder.append("two"); + builder.append_index(9); + bool assertion = builder.is_valid(error) == !BUG_FIXED; + std::cout << error << std::endl; + assert(assertion); +} + +void +test_IndexedOption_categorical() { + IndexedOptionBuilder> builder; + assert(builder.length() == 0); + builder.set_parameters("\"__array__\": \"categorical\""); + + builder.append_invalid(); + auto& subbuilder = builder.append_valid(1); + subbuilder.append("zero"); + builder.append_valid(1); + subbuilder.append("one"); + + std::string error; + bool assertion = builder.is_valid(error); + std::cout << error << std::endl; + assert(assertion); + + // index and content could have different lengths + builder.append_valid(1); + builder.append_valid(1); + builder.append_valid(1); + assert(builder.is_valid(error) == BUG_FIXED); +} + +void +test_IndexedOption_categorical_invalid_index() { + IndexedOptionBuilder> builder; + assert(builder.length() == 0); + builder.set_parameters("\"__array__\": \"categorical\""); + + builder.append_invalid(); + auto& subbuilder = builder.append_valid(1); + subbuilder.append("zero"); + builder.append_valid(1); + subbuilder.append("one"); + + std::string error; + bool assertion = builder.is_valid(error); + std::cout << error << std::endl; + assert(assertion); + + // index should be less than the length of content + builder.append_valid(9); + subbuilder.append("two"); + assertion = builder.is_valid(error) == !BUG_FIXED; + std::cout << error << std::endl; + assert(assertion); +} + +int main(int /* argc */, char ** /* argv */) { + test_Indexed_categorical(); + test_Indexed_categorical_invalid_index(); + test_IndexedOption_categorical(); + test_IndexedOption_categorical_invalid_index(); + + return 0; +}