feat: add more kernels with cumulative sum

ManasviGoyal · ManasviGoyal · commit 432a11d6d71c · 2024-01-30T15:24:32.000+01:00
diff --git a/dev/generate-kernel-signatures.py b/dev/generate-kernel-signatures.py
@@ -50,9 +50,12 @@
     "awkward_missing_repeat",
     "awkward_RegularArray_getitem_jagged_expand",
     "awkward_ListArray_getitem_jagged_expand",
+    "awkward_ListArray_getitem_jagged_carrylen",
     "awkward_ListArray_getitem_next_array_advanced",
     "awkward_ListArray_getitem_next_array",
     "awkward_ListArray_getitem_next_at",
+    "awkward_ListArray_getitem_next_range_counts",
+    "awkward_ListArray_rpad_and_clip_length_axis1",
     "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
@@ -690,9 +690,12 @@ def gencpuunittests(specdict):
     "awkward_missing_repeat",
     "awkward_RegularArray_getitem_jagged_expand",
     "awkward_ListArray_getitem_jagged_expand",
+    "awkward_ListArray_getitem_jagged_carrylen",
     "awkward_ListArray_getitem_next_array_advanced",
     "awkward_ListArray_getitem_next_array",
     "awkward_ListArray_getitem_next_at",
+    "awkward_ListArray_getitem_next_range_counts",
+    "awkward_ListArray_rpad_and_clip_length_axis1",
     "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/__init__.py b/src/awkward/_connect/cuda/__init__.py
@@ -84,12 +84,15 @@ def fetch_template_specializations(kernel_dict):
         "awkward_IndexedArray_flatten_nextcarry",
         "awkward_IndexedArray_getitem_nextcarry",
         "awkward_IndexedArray_getitem_nextcarry_outindex",
+        "awkward_ListArray_getitem_next_range_counts",
         "awkward_IndexedArray_index_of_nulls",
         "awkward_IndexedArray_reduce_next_64",
         "awkward_IndexedArray_reduce_next_nonlocal_nextshifts_64",
         "awkward_IndexedArray_reduce_next_nonlocal_nextshifts_fromshifts_64",
         "awkward_IndexedOptionArray_rpad_and_clip_mask_axis1",
         "awkward_ListArray_compact_offsets",
+        "awkward_ListArray_getitem_jagged_carrylen",
+        "awkward_ListArray_rpad_and_clip_length_axis1",
         "awkward_MaskedArray_getitem_next_jagged_project",
         "awkward_UnionArray_project",
         "awkward_reduce_count_64",
diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_ByteMaskedArray_numnull.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_ByteMaskedArray_numnull.cu
@@ -24,7 +24,6 @@ awkward_ByteMaskedArray_numnull_a(T* numnull,
     int64_t thread_id = blockIdx.x * blockDim.x + threadIdx.x;
 
     if (thread_id < length) {
-      *numnull = 0;
       if ((mask[thread_id] != 0) != validwhen) {
         scan_in_array[thread_id] = 1;
       }
@@ -45,6 +44,6 @@ awkward_ByteMaskedArray_numnull_b(T* numnull,
                                   uint64_t invocation_index,
                                   uint64_t* err_code) {
   if (err_code[0] == NO_ERROR) {
-    *numnull = scan_in_array[length - 1];
+    *numnull = (T)scan_in_array[length - 1];
   }
 }
diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_IndexedArray_numnull.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_IndexedArray_numnull.cu
@@ -14,11 +14,11 @@
 template <typename T, typename C>
 __global__ void
 awkward_IndexedArray_numnull_a(T* numnull,
-                                  const C* fromindex,
-                                  int64_t lenindex,
-                                  int64_t* scan_in_array,
-                                  uint64_t invocation_index,
-                                  uint64_t* err_code) {
+                               const C* fromindex,
+                               int64_t lenindex,
+                               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;
 
@@ -36,12 +36,12 @@ awkward_IndexedArray_numnull_a(T* numnull,
 template <typename T, typename C>
 __global__ void
 awkward_IndexedArray_numnull_b(T* numnull,
-                                  const C* fromindex,
-                                  int64_t lenindex,
-                                  int64_t* scan_in_array,
-                                  uint64_t invocation_index,
-                                  uint64_t* err_code) {
+                               const C* fromindex,
+                               int64_t lenindex,
+                               int64_t* scan_in_array,
+                               uint64_t invocation_index,
+                               uint64_t* err_code) {
   if (err_code[0] == NO_ERROR) {
-    *numnull = scan_in_array[lenindex - 1];
+    *numnull = (T)scan_in_array[lenindex - 1];
   }
 }
diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_IndexedArray_numnull_parents.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_IndexedArray_numnull_parents.cu
@@ -54,7 +54,7 @@ awkward_IndexedArray_numnull_parents_b(T* numnull,
         numnull[thread_id] = 0;
       }
     }
-    *tolength = scan_in_array[lenindex - 1];
+    *tolength = (T)scan_in_array[lenindex - 1];
   }
 }
 
diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_carrylen.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_jagged_carrylen.cu
@@ -0,0 +1,46 @@
+// BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+// BEGIN PYTHON
+// def f(grid, block, args):
+//     (carrylen, slicestarts, slicestops, sliceouterlen, invocation_index, err_code) = args
+//     scan_in_array = cupy.empty(sliceouterlen, dtype=cupy.int64)
+//     cuda_kernel_templates.get_function(fetch_specialization(["awkward_ListArray_getitem_jagged_carrylen_a", carrylen.dtype, slicestarts.dtype, slicestops.dtype]))(grid, block, (carrylen, slicestarts, slicestops, sliceouterlen, scan_in_array, invocation_index, err_code))
+//     scan_in_array = inclusive_scan(grid, block, (scan_in_array, invocation_index, err_code))
+//     cuda_kernel_templates.get_function(fetch_specialization(["awkward_ListArray_getitem_jagged_carrylen_b", carrylen.dtype, slicestarts.dtype, slicestops.dtype]))(grid, block, (carrylen, slicestarts, slicestops, sliceouterlen, scan_in_array, invocation_index, err_code))
+// out["awkward_ListArray_getitem_jagged_carrylen_a", {dtype_specializations}] = None
+// out["awkward_ListArray_getitem_jagged_carrylen_b", {dtype_specializations}] = None
+// END PYTHON
+
+template <typename T, typename C, typename U>
+__global__ void
+awkward_ListArray_getitem_jagged_carrylen_a(T* carrylen,
+                                            const C* slicestarts,
+                                            const U* slicestops,
+                                            int64_t sliceouterlen,
+                                            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 < sliceouterlen) {
+        scan_in_array[thread_id] = (T)(slicestops[thread_id] - slicestarts[thread_id]);
+    }
+  }
+}
+
+template <typename T, typename C, typename U>
+__global__ void
+awkward_ListArray_getitem_jagged_carrylen_b(T* carrylen,
+                                            const C* slicestarts,
+                                            const U* slicestops,
+                                            int64_t sliceouterlen,
+                                            int64_t* scan_in_array,
+                                            uint64_t invocation_index,
+                                            uint64_t* err_code) {
+  if (err_code[0] == NO_ERROR) {
+      *carrylen = scan_in_array[sliceouterlen - 1];
+  }
+}
+
+// fails for sliceouterlen = 1
diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_next_range_counts.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_getitem_next_range_counts.cu
@@ -0,0 +1,46 @@
+// BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+// BEGIN PYTHON
+// def f(grid, block, args):
+//     (total, fromoffsets, lenstarts, invocation_total, err_code) = args
+//     scan_in_array = cupy.empty(lenstarts, dtype=cupy.int64)
+//     cuda_kernel_templates.get_function(fetch_specialization(["awkward_ListArray_getitem_next_range_counts_a", total.dtype, fromoffsets.dtype]))(grid, block, (total, fromoffsets, lenstarts, scan_in_array, invocation_total, err_code))
+//     scan_in_array = inclusive_scan(grid, block, (scan_in_array, invocation_total, err_code))
+//     cuda_kernel_templates.get_function(fetch_specialization(["awkward_ListArray_getitem_next_range_counts_b", total.dtype, fromoffsets.dtype]))(grid, block, (total, fromoffsets, lenstarts, scan_in_array, invocation_total, err_code))
+// out["awkward_ListArray_getitem_next_range_counts_a", {dtype_specializations}] = None
+// out["awkward_ListArray_getitem_next_range_counts_b", {dtype_specializations}] = None
+// END PYTHON
+
+template <typename T, typename C>
+__global__ void
+awkward_ListArray_getitem_next_range_counts_a(T* total,
+                                            const C* fromoffsets,
+                                            int64_t lenstarts,
+                                            int64_t* scan_in_array,
+                                            uint64_t invocation_total,
+                                            uint64_t* err_code) {
+  if (err_code[0] == NO_ERROR) {
+    int64_t thread_id = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (thread_id < lenstarts) {
+      scan_in_array[thread_id] = (T)(fromoffsets[thread_id + 1] - fromoffsets[thread_id]);
+    }
+  }
+}
+
+template <typename T, typename C>
+__global__ void
+awkward_ListArray_getitem_next_range_counts_b(T* total,
+                                            const C* fromoffsets,
+                                            int64_t lenstarts,
+                                            int64_t* scan_in_array,
+                                            uint64_t invocation_total,
+                                            uint64_t* err_code) {
+  if (err_code[0] == NO_ERROR) {
+    int64_t thread_id = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (thread_id < lenstarts) {
+      *total = scan_in_array[lenstarts - 1];
+    }
+  }
+}
diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_rpad_and_clip_length_axis1.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_ListArray_rpad_and_clip_length_axis1.cu
@@ -0,0 +1,49 @@
+// BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+// BEGIN PYTHON
+// def f(grid, block, args):
+//     (tomin, fromstarts, fromstops, target, lenstarts, invocation_index, err_code) = args
+//     scan_in_array = cupy.empty(lenstarts, dtype=cupy.int64)
+//     cuda_kernel_templates.get_function(fetch_specialization(["awkward_ListArray_rpad_and_clip_length_axis1_a", tomin.dtype, fromstarts.dtype, fromstops.dtype]))(grid, block, (tomin, fromstarts, fromstops, target, lenstarts, scan_in_array, invocation_index, err_code))
+//     scan_in_array = inclusive_scan(grid, block, (scan_in_array, invocation_index, err_code))
+//     cuda_kernel_templates.get_function(fetch_specialization(["awkward_ListArray_rpad_and_clip_length_axis1_b", tomin.dtype, fromstarts.dtype, fromstops.dtype]))(grid, block, (tomin, fromstarts, fromstops, target, lenstarts, scan_in_array, invocation_index, err_code))
+// out["awkward_ListArray_rpad_and_clip_length_axis1_a", {dtype_specializations}] = None
+// out["awkward_ListArray_rpad_and_clip_length_axis1_b", {dtype_specializations}] = None
+// END PYTHON
+
+template <typename T, typename C, typename U>
+__global__ void
+awkward_ListArray_rpad_and_clip_length_axis1_a(T* tomin,
+                                               const C* fromstarts,
+                                               const U* fromstops,
+                                               int64_t target,
+                                               int64_t lenstarts,
+                                               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 < lenstarts) {
+      int64_t rangeval = fromstops[thread_id] - fromstarts[thread_id];
+      scan_in_array[thread_id] = (target > rangeval) ? target : rangeval;
+    }
+  }
+}
+
+template <typename T, typename C, typename U>
+__global__ void
+awkward_ListArray_rpad_and_clip_length_axis1_b(T* tomin,
+                                               const C* fromstarts,
+                                               const U* fromstops,
+                                               int64_t target,
+                                               int64_t lenstarts,
+                                               int64_t* scan_in_array,
+                                               uint64_t invocation_index,
+                                               uint64_t* err_code) {
+  if (err_code[0] == NO_ERROR) {
+      *tomin = scan_in_array[lenstarts - 1];
+  }
+}
+
+// fails for lenstarts = 1

Original file line number	Diff line number	Diff line change
`@@ -24,7 +24,6 @@ awkward_ByteMaskedArray_numnull_a(T* numnull,`
`24`	`24`	`int64_t thread_id = blockIdx.x * blockDim.x + threadIdx.x;`
`25`	`25`
`26`	`26`	`if (thread_id < length) {`
`27`		`- *numnull = 0;`
`28`	`27`	`if ((mask[thread_id] != 0) != validwhen) {`
`29`	`28`	`scan_in_array[thread_id] = 1;`
`30`	`29`	`}`
`@@ -45,6 +44,6 @@ awkward_ByteMaskedArray_numnull_b(T* numnull,`
`45`	`44`	`uint64_t invocation_index,`
`46`	`45`	`uint64_t* err_code) {`
`47`	`46`	`if (err_code[0] == NO_ERROR) {`
`48`		`- *numnull = scan_in_array[length - 1];`
	`47`	`+ *numnull = (T)scan_in_array[length - 1];`
`49`	`48`	`}`
`50`	`49`	`}`
Original file line number	Diff line number	Diff line change
`@@ -54,7 +54,7 @@ awkward_IndexedArray_numnull_parents_b(T* numnull,`
`54`	`54`	`numnull[thread_id] = 0;`
`55`	`55`	`}`
`56`	`56`	`}`
`57`		`- *tolength = scan_in_array[lenindex - 1];`
	`57`	`+ *tolength = (T)scan_in_array[lenindex - 1];`
`58`	`58`	`}`
`59`	`59`	`}`
`60`	`60`