feat: add more cuda kernels

scikit-hep · Jan 15, 2024 · effcef9 · effcef9
1 parent 327b1bf
commit effcef9
Show file tree

Hide file tree

Showing 6 changed files with 81 additions and 2 deletions.
diff --git a/awkward-cpp/src/cpu-kernels/awkward_ListOffsetArray_drop_none_indexes.cpp b/awkward-cpp/src/cpu-kernels/awkward_ListOffsetArray_drop_none_indexes.cpp
@@ -41,7 +41,6 @@ ERROR awkward_ListOffsetArray_drop_none_indexes_64(
     length_offsets,
     length_indexes);
 }
-
 ERROR awkward_ListOffsetArray_drop_none_indexes_32(
   int32_t* tooffsets,
   const int32_t* noneindexes,

diff --git a/dev/generate-kernel-signatures.py b/dev/generate-kernel-signatures.py
@@ -45,7 +45,9 @@
     "awkward_missing_repeat",
     "awkward_RegularArray_getitem_jagged_expand",
     "awkward_ListArray_getitem_jagged_expand",
+    "awkward_ListArray_getitem_next_array_advanced",
     "awkward_ListArray_getitem_next_array",
+    "awkward_ListArray_getitem_next_at",
     "awkward_NumpyArray_reduce_adjust_starts_64",
     "awkward_NumpyArray_reduce_adjust_starts_shifts_64",
     "awkward_RegularArray_getitem_next_at",

diff --git a/dev/generate-tests.py b/dev/generate-tests.py
@@ -494,7 +494,9 @@ def gencpukerneltests(specdict):
     "awkward_missing_repeat",
     "awkward_RegularArray_getitem_jagged_expand",
     "awkward_ListArray_getitem_jagged_expand",
+    "awkward_ListArray_getitem_next_array_advanced",
     "awkward_ListArray_getitem_next_array",
+    "awkward_ListArray_getitem_next_at",
     "awkward_NumpyArray_reduce_adjust_starts_64",
     "awkward_NumpyArray_reduce_adjust_starts_shifts_64",
     "awkward_RegularArray_getitem_next_at",

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
@@ -0,0 +1,45 @@
+// BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+enum class LISTARRAY_GETITEM_NEXT_ARRAY_ADVANCED_ERRORS {
+  STOP_LT_START,     // message: "stops[i] < starts[i]"
+  STOP_GET_LEN,      // message: "stops[i] > len(content)"
+  IND_OUT_OF_RANGE,  // message: "index out of range"
+};
+
+template <typename T, typename C, typename U, typename V, typename W, typename X>
+__global__ void
+awkward_ListArray_getitem_next_array_advanced(T* tocarry,
+                                              C* toadvanced,
+                                              const U* fromstarts,
+                                              const V* fromstops,
+                                              const W* fromarray,
+                                              const X* fromadvanced,
+                                              int64_t lenstarts,
+                                              int64_t lenarray,
+                                              int64_t lencontent,
+                                              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 < lenstarts) {
+      if (fromstops[thread_id] < fromstarts[thread_id]) {
+        RAISE_ERROR(LISTARRAY_GETITEM_NEXT_ARRAY_ADVANCED_ERRORS::STOP_LT_START)
+      }
+      if ((fromstarts[thread_id] != fromstops[thread_id]) &&
+          (fromstops[thread_id] > lencontent)) {
+        RAISE_ERROR(LISTARRAY_GETITEM_NEXT_ARRAY_ADVANCED_ERRORS::STOP_GET_LEN)
+      }
+      int64_t length = fromstops[thread_id] - fromstarts[thread_id];
+      int64_t regular_at = fromarray[fromadvanced[thread_id]];
+      if (regular_at < 0) {
+        regular_at += length;
+      }
+      if (!(0 <= regular_at && regular_at < length)) {
+        RAISE_ERROR(LISTARRAY_GETITEM_NEXT_ARRAY_ADVANCED_ERRORS::IND_OUT_OF_RANGE)
+      }
+      tocarry[thread_id] = fromstarts[thread_id] + regular_at;
+      toadvanced[thread_id] = thread_id;
+    }
+  }
+}
diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_next_at.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_next_at.cu
@@ -0,0 +1,31 @@
+// BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+enum class LISTARRAY_GETITEM_NEXT_AT_ERRORS {
+  IND_OUT_OF_RANGE,  // message: "index out of range"
+};
+
+template <typename T, typename C, typename U>
+__global__ void
+awkward_ListArray_getitem_next_at(T* tocarry,
+                                  const C* fromstarts,
+                                  const U* fromstops,
+                                  int64_t lenstarts,
+                                  int64_t at,
+                                  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 < lenstarts) {
+      int64_t length = fromstops[thread_id] - fromstarts[thread_id];
+      int64_t regular_at = at;
+      if (regular_at < 0) {
+        regular_at += length;
+      }
+      if (!(0 <= regular_at && regular_at < length)) {
+        RAISE_ERROR(LISTARRAY_GETITEM_NEXT_AT_ERRORS::IND_OUT_OF_RANGE)
+      }
+      tocarry[thread_id] = fromstarts[thread_id] + regular_at;
+    }
+  }
+}
diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_validity.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_validity.cu
@@ -14,8 +14,8 @@ awkward_ListArray_validity(const C* starts,
                            int64_t lencontent,
                            uint64_t invocation_index,
                            uint64_t* err_code) {
-  int64_t thread_id = blockIdx.x * blockDim.x + threadIdx.x;
   if (err_code[0] == NO_ERROR) {
+    int64_t thread_id = blockIdx.x * blockDim.x + threadIdx.x;
     if (thread_id < length) {
       C start = starts[thread_id];
       T stop = stops[thread_id];