feat: complex reducers (#3149)

* feat: sum complex

* feat: add sum_complex reducer

* fix: remove min_complex

* fix: test data

* fix: spec tests

* test: sum complex integration tests

* feat: add sum_bool, prod_bool, countnonzero complex

* test: add complex tests for ak.all and ak.any

* fix: remove extra kernel from list

* feat: add min and max complex

* feat: added awkward_IndexedArray_local_preparenext_64

* test: add tests

* feat: add awkward_ListOffsetArray_reduce_nonlocal_outstartsstops_64 kernel

* feat: added kernels and tests

* test: makrk missing kernel tests as skip instead of xfail

* fix: minor fixes

* test: add few tests

* fix: minor fix

* Update dev/generate-tests.py

revert merge change

---------

Co-authored-by: Ianna Osborne <[email protected]>

awkward-cpp/src/cpu-kernels/awkward_IndexedArray_local_preparenext_64.cpp

@@ -6,7 +6,7 @@
 
 ERROR awkward_IndexedArray_local_preparenext_64(
     int64_t* tocarry,
-    const int64_t* /* starts */,   // FIXME: this argument is not needed
+    const int64_t* /* starts */,   // used in CUDA kernels
     const int64_t* parents,
     const int64_t parentslength,
     const int64_t* nextparents,

dev/generate-kernel-signatures.py

@@ -94,6 +94,7 @@
     "awkward_IndexedArray_getitem_nextcarry",
     "awkward_IndexedArray_getitem_nextcarry_outindex",
     "awkward_IndexedArray_index_of_nulls",
+    "awkward_IndexedArray_local_preparenext_64",
     "awkward_IndexedArray_ranges_next_64",
     "awkward_IndexedArray_ranges_carry_next_64",
     "awkward_IndexedArray_reduce_next_64",
@@ -110,25 +111,32 @@
     "awkward_ListOffsetArray_drop_none_indexes",
     "awkward_ListOffsetArray_reduce_local_nextparents_64",
     "awkward_ListOffsetArray_reduce_nonlocal_maxcount_offsetscopy_64",
+    "awkward_ListOffsetArray_reduce_nonlocal_outstartsstops_64",
     "awkward_ListOffsetArray_reduce_local_outoffsets_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",
-    "awkward_reduce_argmax",
-    "awkward_reduce_argmin",
+    "awkward_RecordArray_reduce_nonlocal_outoffsets_64",
     "awkward_reduce_count_64",
     "awkward_reduce_max",
+    "awkward_reduce_max_complex",
     "awkward_reduce_min",
+    "awkward_reduce_min_complex",
     "awkward_reduce_sum",
+    "awkward_reduce_sum_bool",
+    "awkward_reduce_sum_bool_complex",
+    "awkward_reduce_sum_complex",
     "awkward_reduce_sum_int32_bool_64",
     "awkward_reduce_sum_int64_bool_64",
-    "awkward_reduce_sum_bool",
     "awkward_reduce_prod",
     "awkward_reduce_prod_bool",
+    "awkward_reduce_prod_bool_complex",
+    "awkward_reduce_prod_complex",
     "awkward_reduce_countnonzero",
+    "awkward_reduce_countnonzero_complex",
     "awkward_sorting_ranges",
     "awkward_sorting_ranges_length",
 ]

dev/generate-tests.py

@@ -444,6 +444,7 @@ def genspectests(specdict):
 
 """
             )
+
             f.write("import pytest\nimport numpy as np\nimport kernels\n\n")
             num = 1
             if spec.tests == []:
@@ -931,6 +932,7 @@ def gencpuunittests(specdict):
     "awkward_IndexedArray_getitem_nextcarry",
     "awkward_IndexedArray_getitem_nextcarry_outindex",
     "awkward_IndexedArray_index_of_nulls",
+    "awkward_IndexedArray_local_preparenext_64",
     "awkward_IndexedArray_ranges_next_64",
     "awkward_IndexedArray_ranges_carry_next_64",
     "awkward_IndexedArray_reduce_next_64",
@@ -947,25 +949,32 @@ def gencpuunittests(specdict):
     "awkward_ListOffsetArray_drop_none_indexes",
     "awkward_ListOffsetArray_reduce_local_nextparents_64",
     "awkward_ListOffsetArray_reduce_nonlocal_maxcount_offsetscopy_64",
+    "awkward_ListOffsetArray_reduce_nonlocal_outstartsstops_64",
     "awkward_ListOffsetArray_reduce_local_outoffsets_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",
-    "awkward_reduce_argmax",
-    "awkward_reduce_argmin",
+    "awkward_RecordArray_reduce_nonlocal_outoffsets_64",
     "awkward_reduce_count_64",
     "awkward_reduce_max",
+    "awkward_reduce_max_complex",
     "awkward_reduce_min",
+    "awkward_reduce_min_complex",
     "awkward_reduce_sum",
+    "awkward_reduce_sum_bool",
+    "awkward_reduce_sum_bool_complex",
+    "awkward_reduce_sum_complex",
     "awkward_reduce_sum_int32_bool_64",
     "awkward_reduce_sum_int64_bool_64",
-    "awkward_reduce_sum_bool",
     "awkward_reduce_prod",
     "awkward_reduce_prod_bool",
+    "awkward_reduce_prod_bool_complex",
+    "awkward_reduce_prod_complex",
     "awkward_reduce_countnonzero",
+    "awkward_reduce_countnonzero_complex",
     "awkward_sorting_ranges",
     "awkward_sorting_ranges_length",
 ]

kernel-specification.yml

@@ -2915,7 +2915,7 @@ kernels:
 
           # Close open sublist
           if lenparents > 0:
-              outoffsets[k_sublist + 1] = j
+              outoffsets[k_sublist + 1] = j + 1
               outcarry[parents[i]] = k_sublist
               k_sublist += 1
 
@@ -3829,7 +3829,22 @@ kernels:
           - {name: stable, type: "bool", dir: in, role: ListArray-replacement}
     description: null
     definition: |
-      Insert Python definition here
+      def awkward_argsort(
+        toptr, fromptr, length, offsets, offsetslength, ascending, stable
+      ):
+          result = []
+          indices = list(range(length))
+          for i in range(offsetslength - 1):
+              sub_indices = indices[offsets[i]:offsets[i+1]]
+              sorted_indices = sorted(
+                  sub_indices,
+                  key=lambda x: fromptr[x],
+                  reverse=not ascending
+              )
+              local_sorted_indices = [idx - offsets[i] for idx in sorted_indices]
+              result.extend(local_sorted_indices)
+          for i in range(length):
+              toptr[i] = result[i]
     automatic-tests: false
 
   - name: awkward_index_rpad_and_clip_axis0
@@ -4018,6 +4033,15 @@ kernels:
           - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength}
     description: null
     definition: |
+      def awkward_reduce_argmax_complex(toptr, fromptr, parents, lenparents, outlength):
+          for k in range(outlength):
+              toptr[k] = -1
+          for i in range(lenparents):
+              parent = parents[i]
+              if (toptr[parent] == -1  or (fromptr[i * 2] > fromptr[toptr[parent * 2]] or
+                 (fromptr[i * 2] == fromptr[toptr[parent * 2]]  and
+                  fromptr[i * 2 + 1] > fromptr[toptr[parent * 2 + 1]]))):
+                  toptr[parent] = i
     automatic-tests: false
 
   - name: awkward_reduce_argmin
@@ -4121,6 +4145,15 @@ kernels:
           - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength}
     description: null
     definition: |
+      def awkward_reduce_argmin_complex(toptr, fromptr, parents, lenparents, outlength):
+          for k in range(outlength):
+              toptr[k] = -1
+          for i in range(lenparents):
+              parent = parents[i]
+              if (toptr[parent] == -1  or (fromptr[i * 2] < fromptr[toptr[parent * 2]] or
+                 (fromptr[i * 2] == fromptr[toptr[parent * 2]]  and
+                  fromptr[i * 2 + 1] < fromptr[toptr[parent * 2 + 1]]))):
+                  toptr[parent] = i
     automatic-tests: false
 
   - name: awkward_reduce_count_64
@@ -4246,6 +4279,11 @@ kernels:
           - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength}
     description: null
     definition: |
+      def awkward_reduce_countnonzero_complex(toptr, fromptr, parents, lenparents, outlength):
+          for i in range(outlength):
+              toptr[i] = 0
+          for i in range(lenparents):
+              toptr[parents[i]] += fromptr[i * 2] != 0 or fromptr[i * 2 + 1] != 0
     automatic-tests: false
 
   - name: awkward_reduce_max
@@ -4360,6 +4398,17 @@ kernels:
           - {name: identity, type: "double", dir: in, role: reducer-identity}
     description: null
     definition: |
+      def awkward_reduce_max_complex(toptr, fromptr, parents, lenparents, outlength, identity):
+          for i in range(outlength):
+              toptr[i * 2] = identity
+              toptr[i * 2 + 1] = 0
+          for i in range(lenparents):
+              x = fromptr[i * 2]
+              y = fromptr[i * 2 + 1]
+              parent = parents[i]
+              if x > toptr[parent * 2] or (x == toptr[parent * 2]  and  y > toptr[parent * 2 + 1]):
+                  toptr[parent * 2] = x
+                  toptr[parent * 2 + 1] = y
     automatic-tests: false
 
   - name: awkward_reduce_min
@@ -4474,6 +4523,17 @@ kernels:
           - {name: identity, type: "double", dir: in, role: reducer-identity}
     description: null
     definition: |
+      def awkward_reduce_min_complex(toptr, fromptr, parents, lenparents, outlength, identity):
+          for i in range(outlength):
+              toptr[i * 2] = identity
+              toptr[i * 2 + 1] = 0
+          for i in range(lenparents):
+              x = fromptr[i * 2]
+              y = fromptr[i * 2 + 1]
+              parent = parents[i]
+              if x < toptr[parent * 2] or (x == toptr[parent * 2]  and  y < toptr[parent * 2 + 1]):
+                  toptr[parent * 2] = x
+                  toptr[parent * 2 + 1] = y
     automatic-tests: false
 
   - name: awkward_reduce_prod
@@ -4617,6 +4677,16 @@ kernels:
           - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength}
     description: null
     definition: |
+      def awkward_reduce_prod_complex(toptr, fromptr, parents, lenparents, outlength):
+          for i in range(outlength):
+              toptr[i * 2] = 1
+              toptr[i * 2 + 1] = 0
+          for i in range(lenparents):
+              parent = parents[i]
+              real = toptr[parent * 2] * fromptr[i * 2] - toptr[parent * 2 + 1] * fromptr[i * 2 + 1]
+              imag = toptr[parent * 2] * fromptr[i * 2 + 1] + toptr[parent * 2 + 1] * fromptr[i * 2]
+              toptr[parents[i] * 2] = real
+              toptr[parents[i] * 2 + 1] = imag
     automatic-tests: false
 
   - name: awkward_reduce_prod_bool
@@ -4725,6 +4795,11 @@ kernels:
           - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength}
     description: null
     definition: |
+      def awkward_reduce_prod_bool_complex(toptr, fromptr, parents, lenparents, outlength):
+          for i in range(outlength):
+              toptr[i] = 1
+          for i in range(lenparents):
+              toptr[parents[i]] = fromptr[i * 2] != 0 or fromptr[i * 2 + 1] != 0
     automatic-tests: false
 
   - name: awkward_reduce_sum
@@ -4868,6 +4943,13 @@ kernels:
           - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength}
     description: null
     definition: |
+      def awkward_reduce_sum_complex(toptr, fromptr, parents, lenparents, outlength):
+          for i in range(outlength):
+              toptr[i * 2] = 0
+              toptr[i * 2 + 1] = 0
+          for i in range(lenparents):
+              toptr[parents[i] * 2] += fromptr[i * 2]
+              toptr[parents[i] * 2 + 1] += fromptr[i * 2 + 1]
     automatic-tests: false
 
   - name: awkward_reduce_sum_bool
@@ -4976,6 +5058,11 @@ kernels:
           - {name: outlength, type: "int64_t", dir: in, role: reducer-outlength}
     description: null
     definition: |
+      def awkward_reduce_sum_bool_complex(toptr, fromptr, parents, lenparents, outlength):
+          for i in range(outlength):
+              toptr[i] = 0
+          for i in range(lenparents):
+              toptr[parents[i]] |= fromptr[i * 2] != 0 or fromptr[i * 2 + 1] != 0
     automatic-tests: false
 
   - name: awkward_reduce_sum_int32_bool_64