feat: add boolean cases for CPU kernels (#3059)

* feat: add awkward_ListArray_getitem_jagged_apply.cu and awkward_ListArray_getitem_jagged_shrink * fix: use cupy.ones instead of multiple kernels * feat: add awkward_NumpyArray_rearrange_shifted CUDA kernel * feat: add new CPU kernel for awkward_unique_bool_64 * fix: test outputs * feat: add bool case for awkward_unique * fix: tolength assignment * fix: test cases * feat: add bool case for some kernels * fix: test case errors * fix: remove templates for bool * fix: segfault in ctypes double pointer * fix: assume sorted array for awkward_unique tests * fix: clean up PR * fix: remove awkward_NumpyArray_fill kernel * fix: remove awkward_unique and awkward_unique_copy kernels * fix: lint error * fix: unexpected keyword argument equal_nan error
scikit-hep · Apr 20, 2024 · 8a701a5 · 8a701a5
1 parent 28d8917
commit 8a701a5
Show file tree

Hide file tree

Showing 24 changed files with 793 additions and 2,197 deletions.
diff --git a/awkward-cpp/src/cpu-kernels/awkward_NumpyArray_fill.cpp b/awkward-cpp/src/cpu-kernels/awkward_NumpyArray_fill.cpp
diff --git a/awkward-cpp/src/cpu-kernels/awkward_NumpyArray_rearrange_shifted.cpp b/awkward-cpp/src/cpu-kernels/awkward_NumpyArray_rearrange_shifted.cpp
@@ -8,25 +8,25 @@ template <typename FROM, typename TO>
 ERROR
 awkward_NumpyArray_rearrange_shifted(
   TO* toptr,
-  const FROM* shifts,
+  const FROM* fromshifts,
   int64_t length,
-  const FROM* offsets,
+  const FROM* fromoffsets,
   int64_t offsetslength,
-  const FROM* parents,
+  const FROM* fromparents,
   int64_t /* parentslength */,  // FIXME: these arguments are not needed
-  const FROM* starts,
+  const FROM* fromstarts,
   int64_t /* startslength */) {
   int64_t k = 0;
   for (int64_t i = 0; i < offsetslength - 1; i++) {
-    for (int64_t j = 0; j < offsets[i + 1] - offsets[i]; j++) {
-      toptr[k] = toptr[k] + offsets[i];
+    for (int64_t j = 0; j < fromoffsets[i + 1] - fromoffsets[i]; j++) {
+      toptr[k] = toptr[k] + fromoffsets[i];
       k++;
     }
   }
   for (int64_t i = 0;  i < length;  i++) {
-    int64_t parent = parents[i];
-    int64_t start = starts[parent];
-    toptr[i] = toptr[i] + shifts[toptr[i]] - start;
+    int64_t parent = fromparents[i];
+    int64_t start = fromstarts[parent];
+    toptr[i] = toptr[i] + fromshifts[toptr[i]] - start;
   }
 
   return success();

diff --git a/awkward-cpp/src/cpu-kernels/awkward_NumpyArray_subrange_equal.cpp b/awkward-cpp/src/cpu-kernels/awkward_NumpyArray_subrange_equal.cpp
@@ -36,20 +36,6 @@ ERROR awkward_NumpyArray_subrange_equal(
 
   return success();
 }
-
-ERROR awkward_NumpyArray_subrange_equal_bool(
-  bool* tmpptr,
-  const int64_t* fromstarts,
-  const int64_t* fromstops,
-  int64_t length,
-  bool* toequal) {
-    return awkward_NumpyArray_subrange_equal<bool>(
-      tmpptr,
-      fromstarts,
-      fromstops,
-      length,
-      toequal);
-}
 ERROR awkward_NumpyArray_subrange_equal_int8(
   int8_t* tmpptr,
   const int64_t* fromstarts,

diff --git a/awkward-cpp/src/cpu-kernels/awkward_NumpyArray_subrange_equal_bool.cpp b/awkward-cpp/src/cpu-kernels/awkward_NumpyArray_subrange_equal_bool.cpp
@@ -0,0 +1,37 @@
+// BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
+
+#define FILENAME(line) FILENAME_FOR_EXCEPTIONS_C("src/cpu-kernels/awkward_NumpyArray_subrange_equal_bool.cpp", line)
+
+#include "awkward/kernels.h"
+
+ERROR awkward_NumpyArray_subrange_equal_bool(
+  bool* tmpptr,
+  const int64_t* fromstarts,
+  const int64_t* fromstops,
+  int64_t length,
+  bool* toequal) {
+
+  bool differ = true;
+  int64_t leftlen;
+  int64_t rightlen;
+
+  for (int64_t i = 0;  i < length - 1;  i++) {
+    leftlen = fromstops[i] - fromstarts[i];
+    for (int64_t ii = i + 1; ii < length - 1;  ii++) {
+      rightlen = fromstops[ii] - fromstarts[ii];
+      if (leftlen == rightlen) {
+        differ = false;
+        for (int64_t j = 0; j < leftlen; j++) {
+          if ((tmpptr[fromstarts[i] + j] != 0) != (tmpptr[fromstarts[ii] + j] != 0)) {
+            differ = true;
+            break;
+          }
+        }
+      }
+    }
+  }
+
+  *toequal = !differ;
+
+  return success();
+}
diff --git a/awkward-cpp/src/cpu-kernels/awkward_unique.cpp b/awkward-cpp/src/cpu-kernels/awkward_unique.cpp
diff --git a/awkward-cpp/src/cpu-kernels/awkward_unique_copy.cpp b/awkward-cpp/src/cpu-kernels/awkward_unique_copy.cpp
diff --git a/awkward-cpp/src/cpu-kernels/awkward_unique_ranges.cpp b/awkward-cpp/src/cpu-kernels/awkward_unique_ranges.cpp
@@ -26,20 +26,6 @@ ERROR awkward_unique_ranges(
   return success();
 }
 
-ERROR awkward_unique_ranges_bool(
-  bool* toptr,
-  int64_t length,
-  const int64_t* fromoffsets,
-  int64_t offsetslength,
-  int64_t* tooffsets) {
-    return awkward_unique_ranges<bool>(
-      toptr,
-      length,
-      fromoffsets,
-      offsetslength,
-      tooffsets);
-}
-
 ERROR awkward_unique_ranges_int8(
   int8_t* toptr,
   int64_t length,