feat: complete CPU kernel migration from parents to offsets (#3988)

awkward-cpp/CMakeLists.txt

@@ -84,6 +84,18 @@ set_target_properties(
              VISIBILITY_INLINES_HIDDEN ON
              CXX_EXTENSIONS NO)
 
+# Optional OpenMP threading for the reducers / sorts. The kernels gate every `#pragma omp parallel
+# for` on `_OPENMP`, so the same source compiles cleanly without OpenMP and runs serially in that
+# case. We don't fail the build if OpenMP isn't available — users who lack it (e.g. macOS Apple
+# Clang without a separate `libomp`) just get the serial path.
+find_package(OpenMP COMPONENTS CXX)
+if(OpenMP_CXX_FOUND)
+  target_link_libraries(awkward-cpu-kernels PRIVATE OpenMP::OpenMP_CXX)
+  message(STATUS "OpenMP found: cpu-kernels will use threading where it pays")
+else()
+  message(STATUS "OpenMP not found: cpu-kernels will run serially")
+endif()
+
 # Second tier: libawkward
 add_library(awkward SHARED ${LIBAWKWARD_SOURCES})
 if(${CMAKE_CXX_COMPILER_ID} MATCHES "Clang")

awkward-cpp/include/awkward/common.h

@@ -9,6 +9,7 @@
 #ifdef _MSC_VER
   #define EXPORT_SYMBOL __declspec(dllexport)
   #define ERROR Error
+  #define __restrict__ __restrict
   using ssize_t = std::ptrdiff_t;
 #else
   #define EXPORT_SYMBOL __attribute__((visibility("default")))

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

@@ -6,26 +6,27 @@
 
 ERROR awkward_ByteMaskedArray_reduce_next_64(
   int64_t* nextcarry,
-  int64_t* nextparents,
+  int64_t* nextoffsets,   // length outlength + 1
   int64_t* outindex,
   const int8_t* mask,
-  const int64_t* parents,
-  int64_t length,
+  const int64_t* offsets,
+  int64_t outlength,
   bool validwhen) {
   int64_t k = 0;
-  for (int64_t i = 0; i < length; i++) {
-    bool is_valid = ((mask[i] != 0) == validwhen);
-
-    if (is_valid) {
+  nextoffsets[0] = 0;
+  for (int64_t bin = 0; bin < outlength; bin++) {
+    for (int64_t i = offsets[bin]; i < offsets[bin + 1]; i++) {
+      bool is_valid = ((mask[i] != 0) == validwhen);
+      if (is_valid) {
         nextcarry[k] = i;
-        nextparents[k] = parents[i];
-        outindex[i] = k;
+        outindex[i]  = k;
         k++;
-    }
-    else {
+      }
+      else {
         outindex[i] = -1;
+      }
     }
+    nextoffsets[bin + 1] = k;
   }
-
   return success();
 }

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

@@ -0,0 +1,41 @@
+// 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_Index_parents_to_offsets.cpp", line)
+
+#include "awkward/kernels.h"
+
+template <typename T>
+ERROR awkward_Index_parents_to_offsets(
+  T* tooffsets,
+  const T* fromparents,
+  int64_t lenparents,
+  int64_t outlength) {
+  // offsets are initialized to 0, so
+  // we only need to count the number of
+  // parents for each offset.
+  for (int64_t i = 0; i < lenparents; i++) {
+    int64_t p = fromparents[i];
+    if (p < outlength) {
+      tooffsets[p + 1]++;
+    }
+  }
+  // offsets length is outlength + 1, so
+  // we need to do a cumulative sum
+  // to get the correct offsets.
+  for (int64_t i = 1; i <= outlength; i++) {
+    tooffsets[i] += tooffsets[i - 1];
+  }
+
+  return success();
+}
+ERROR awkward_Index_parents_to_offsets_64(
+  int64_t* tooffsets,
+  const int64_t* fromparents,
+  int64_t lenparents,
+  int64_t outlength) {
+  return awkward_Index_parents_to_offsets<int64_t>(
+    tooffsets,
+    fromparents,
+    lenparents,
+    outlength);
+}

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

@@ -17,42 +17,12 @@ ERROR awkward_IndexedArray_fill(
   }
   return success();
 }
-ERROR awkward_IndexedArray_fill_to64_from32(
-  int64_t* toindex,
-  int64_t toindexoffset,
-  const int32_t* fromindex,
-  int64_t length,
-  int64_t base) {
-  return awkward_IndexedArray_fill<int32_t, int64_t>(
-    toindex,
-    toindexoffset,
-    fromindex,
-    length,
-    base);
-}
-ERROR awkward_IndexedArray_fill_to64_fromU32(
-  int64_t* toindex,
-  int64_t toindexoffset,
-  const uint32_t* fromindex,
-  int64_t length,
-  int64_t base) {
-  return awkward_IndexedArray_fill<uint32_t, int64_t>(
-    toindex,
-    toindexoffset,
-    fromindex,
-    length,
-    base);
-}
-ERROR awkward_IndexedArray_fill_to64_from64(
-  int64_t* toindex,
-  int64_t toindexoffset,
-  const int64_t* fromindex,
-  int64_t length,
-  int64_t base) {
-  return awkward_IndexedArray_fill<int64_t, int64_t>(
-    toindex,
-    toindexoffset,
-    fromindex,
-    length,
-    base);
-}
+
+#define WRAPPER(SUFFIX, FROM, TO) \
+  ERROR awkward_IndexedArray_fill_to64_from##SUFFIX(TO* toindex, int64_t toindexoffset, const FROM* fromindex, int64_t length, int64_t base) { \
+    return awkward_IndexedArray_fill<FROM, TO>(toindex, toindexoffset, fromindex, length, base); \
+  }
+
+WRAPPER(32, int32_t, int64_t)
+WRAPPER(U32, uint32_t, int64_t)
+WRAPPER(64, int64_t, int64_t)

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

@@ -23,36 +23,12 @@ ERROR awkward_IndexedArray_flatten_nextcarry(
   }
   return success();
 }
-ERROR awkward_IndexedArray32_flatten_nextcarry_64(
-  int64_t* tocarry,
-  const int32_t* fromindex,
-  int64_t lenindex,
-  int64_t lencontent) {
-  return awkward_IndexedArray_flatten_nextcarry<int32_t, int64_t>(
-    tocarry,
-    fromindex,
-    lenindex,
-    lencontent);
-}
-ERROR awkward_IndexedArrayU32_flatten_nextcarry_64(
-  int64_t* tocarry,
-  const uint32_t* fromindex,
-  int64_t lenindex,
-  int64_t lencontent) {
-  return awkward_IndexedArray_flatten_nextcarry<uint32_t, int64_t>(
-    tocarry,
-    fromindex,
-    lenindex,
-    lencontent);
-}
-ERROR awkward_IndexedArray64_flatten_nextcarry_64(
-  int64_t* tocarry,
-  const int64_t* fromindex,
-  int64_t lenindex,
-  int64_t lencontent) {
-  return awkward_IndexedArray_flatten_nextcarry<int64_t, int64_t>(
-    tocarry,
-    fromindex,
-    lenindex,
-    lencontent);
-}
+
+#define WRAPPER(SUFFIX, C, T) \
+  ERROR awkward_IndexedArray##SUFFIX(T* tocarry, const C* fromindex, int64_t lenindex, int64_t lencontent) { \
+    return awkward_IndexedArray_flatten_nextcarry<C, T>(tocarry, fromindex, lenindex, lencontent); \
+  }
+
+WRAPPER(32_flatten_nextcarry_64, int32_t, int64_t)
+WRAPPER(U32_flatten_nextcarry_64, uint32_t, int64_t)
+WRAPPER(64_flatten_nextcarry_64, int64_t, int64_t)

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

@@ -31,42 +31,12 @@ ERROR awkward_IndexedArray_flatten_none2empty(
   }
   return success();
 }
-ERROR awkward_IndexedArray32_flatten_none2empty_64(
-  int64_t* outoffsets,
-  const int32_t* outindex,
-  int64_t outindexlength,
-  const int64_t* offsets,
-  int64_t offsetslength) {
-  return awkward_IndexedArray_flatten_none2empty<int64_t, int32_t>(
-    outoffsets,
-    outindex,
-    outindexlength,
-    offsets,
-    offsetslength);
-}
-ERROR awkward_IndexedArrayU32_flatten_none2empty_64(
-  int64_t* outoffsets,
-  const uint32_t* outindex,
-  int64_t outindexlength,
-  const int64_t* offsets,
-  int64_t offsetslength) {
-  return awkward_IndexedArray_flatten_none2empty<int64_t, uint32_t>(
-    outoffsets,
-    outindex,
-    outindexlength,
-    offsets,
-    offsetslength);
-}
-ERROR awkward_IndexedArray64_flatten_none2empty_64(
-  int64_t* outoffsets,
-  const int64_t* outindex,
-  int64_t outindexlength,
-  const int64_t* offsets,
-  int64_t offsetslength) {
-  return awkward_IndexedArray_flatten_none2empty<int64_t, int64_t>(
-    outoffsets,
-    outindex,
-    outindexlength,
-    offsets,
-    offsetslength);
-}
+
+#define WRAPPER(SUFFIX, T, C) \
+  ERROR awkward_IndexedArray##SUFFIX(T* outoffsets, const C* outindex, int64_t outindexlength, const T* offsets, int64_t offsetslength) { \
+    return awkward_IndexedArray_flatten_none2empty<T, C>(outoffsets, outindex, outindexlength, offsets, offsetslength); \
+  }
+
+WRAPPER(32_flatten_none2empty_64, int64_t, int32_t)
+WRAPPER(U32_flatten_none2empty_64, int64_t, uint32_t)
+WRAPPER(64_flatten_none2empty_64, int64_t, int64_t)

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

@@ -23,36 +23,12 @@ ERROR awkward_IndexedArray_getitem_nextcarry(
   }
   return success();
 }
-ERROR awkward_IndexedArray32_getitem_nextcarry_64(
-  int64_t* tocarry,
-  const int32_t* fromindex,
-  int64_t lenindex,
-  int64_t lencontent) {
-  return awkward_IndexedArray_getitem_nextcarry<int32_t, int64_t>(
-    tocarry,
-    fromindex,
-    lenindex,
-    lencontent);
-}
-ERROR awkward_IndexedArrayU32_getitem_nextcarry_64(
-  int64_t* tocarry,
-  const uint32_t* fromindex,
-  int64_t lenindex,
-  int64_t lencontent) {
-  return awkward_IndexedArray_getitem_nextcarry<uint32_t, int64_t>(
-    tocarry,
-    fromindex,
-    lenindex,
-    lencontent);
-}
-ERROR awkward_IndexedArray64_getitem_nextcarry_64(
-  int64_t* tocarry,
-  const int64_t* fromindex,
-  int64_t lenindex,
-  int64_t lencontent) {
-  return awkward_IndexedArray_getitem_nextcarry<int64_t, int64_t>(
-    tocarry,
-    fromindex,
-    lenindex,
-    lencontent);
-}
+
+#define WRAPPER(SUFFIX, C, T) \
+  ERROR awkward_IndexedArray##SUFFIX(T* tocarry, const C* fromindex, int64_t lenindex, int64_t lencontent) { \
+    return awkward_IndexedArray_getitem_nextcarry<C, T>(tocarry, fromindex, lenindex, lencontent); \
+  }
+
+WRAPPER(32_getitem_nextcarry_64, int32_t, int64_t)
+WRAPPER(U32_getitem_nextcarry_64, uint32_t, int64_t)
+WRAPPER(64_getitem_nextcarry_64, int64_t, int64_t)

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

@@ -28,42 +28,12 @@ ERROR awkward_IndexedArray_getitem_nextcarry_outindex(
   }
   return success();
 }
-ERROR awkward_IndexedArray32_getitem_nextcarry_outindex_64(
-  int64_t* tocarry,
-  int32_t* toindex,
-  const int32_t* fromindex,
-  int64_t lenindex,
-  int64_t lencontent) {
-  return awkward_IndexedArray_getitem_nextcarry_outindex<int32_t, int64_t>(
-    tocarry,
-    toindex,
-    fromindex,
-    lenindex,
-    lencontent);
-}
-ERROR awkward_IndexedArrayU32_getitem_nextcarry_outindex_64(
-  int64_t* tocarry,
-  uint32_t* toindex,
-  const uint32_t* fromindex,
-  int64_t lenindex,
-  int64_t lencontent) {
-  return awkward_IndexedArray_getitem_nextcarry_outindex<uint32_t, int64_t>(
-    tocarry,
-    toindex,
-    fromindex,
-    lenindex,
-    lencontent);
-}
-ERROR awkward_IndexedArray64_getitem_nextcarry_outindex_64(
-  int64_t* tocarry,
-  int64_t* toindex,
-  const int64_t* fromindex,
-  int64_t lenindex,
-  int64_t lencontent) {
-  return awkward_IndexedArray_getitem_nextcarry_outindex<int64_t, int64_t>(
-    tocarry,
-    toindex,
-    fromindex,
-    lenindex,
-    lencontent);
-}
+
+#define WRAPPER(SUFFIX, C, T) \
+  ERROR awkward_IndexedArray##SUFFIX(T* tocarry, C* toindex, const C* fromindex, int64_t lenindex, int64_t lencontent) { \
+    return awkward_IndexedArray_getitem_nextcarry_outindex<C, T>(tocarry, toindex, fromindex, lenindex, lencontent); \
+  }
+
+WRAPPER(32_getitem_nextcarry_outindex_64, int32_t, int64_t)
+WRAPPER(U32_getitem_nextcarry_outindex_64, uint32_t, int64_t)
+WRAPPER(64_getitem_nextcarry_outindex_64, int64_t, int64_t)