feat: add some of the kernels using cuda.compute

dev/generate-kernel-signatures.py

@@ -16,11 +16,11 @@
     "awkward_ListArray_min_range",
     "awkward_ListArray_validity",
     "awkward_BitMaskedArray_to_ByteMaskedArray",
-    "awkward_ListArray_broadcast_tooffsets",
-    "awkward_ListArray_compact_offsets",
+    # "awkward_ListArray_broadcast_tooffsets",
+    # "awkward_ListArray_compact_offsets",
     "awkward_ListOffsetArray_flatten_offsets",
-    "awkward_IndexedArray_overlay_mask",
-    "awkward_ByteMaskedArray_numnull",
+    # "awkward_IndexedArray_overlay_mask",
+    # "awkward_ByteMaskedArray_numnull",
     "awkward_IndexedArray_numnull",
     "awkward_IndexedArray_numnull_parents",
     "awkward_IndexedArray_numnull_unique_64",
@@ -51,9 +51,9 @@
     "awkward_RegularArray_reduce_local_nextparents_64",
     "awkward_RegularArray_reduce_nonlocal_preparenext_64",
     # "awkward_missing_repeat",
-    "awkward_RegularArray_getitem_jagged_expand",
-    "awkward_ListArray_combinations_length",
-    "awkward_ListArray_combinations",
+    # "awkward_RegularArray_getitem_jagged_expand",
+    # "awkward_ListArray_combinations_length",
+    # "awkward_ListArray_combinations",
     "awkward_RegularArray_combinations_64",
     "awkward_ListArray_getitem_jagged_apply",
     "awkward_ListArray_getitem_jagged_carrylen",
@@ -72,14 +72,14 @@
     "awkward_UnionArray_regular_index",
     "awkward_ListOffsetArray_reduce_nonlocal_nextstarts_64",
     "awkward_ListArray_getitem_next_range_spreadadvanced",
-    "awkward_ListArray_localindex",
+    # "awkward_ListArray_localindex",
     "awkward_NumpyArray_pad_zero_to_length",
     "awkward_NumpyArray_reduce_adjust_starts_64",
     "awkward_NumpyArray_rearrange_shifted",
     "awkward_NumpyArray_reduce_adjust_starts_shifts_64",
     "awkward_RegularArray_getitem_next_at",
     "awkward_BitMaskedArray_to_IndexedOptionArray",
-    "awkward_ByteMaskedArray_getitem_nextcarry",
+    # "awkward_ByteMaskedArray_getitem_nextcarry",
     "awkward_ByteMaskedArray_getitem_nextcarry_outindex",
     "awkward_ByteMaskedArray_reduce_next_64",
     "awkward_ByteMaskedArray_reduce_next_nonlocal_nextshifts_64",
@@ -97,8 +97,8 @@
     "awkward_IndexedArray_local_preparenext_64",
     "awkward_IndexedArray_ranges_next_64",
     "awkward_IndexedArray_ranges_carry_next_64",
-    "awkward_IndexedArray_reduce_next_64",
-    "awkward_IndexedArray_reduce_next_nonlocal_nextshifts_64",
+    # "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_ListOffsetArray_local_preparenext_64",
@@ -118,7 +118,7 @@
     "awkward_UnionArray_nestedfill_tags_index",
     "awkward_UnionArray_regular_index_getsize",
     "awkward_UnionArray_simplify",
-    "awkward_UnionArray_simplify_one",
+    # "awkward_UnionArray_simplify_one",
     "awkward_RecordArray_reduce_nonlocal_outoffsets_64",
     # "awkward_reduce_count_64",
     # "awkward_reduce_max",

dev/generate-tests.py

@@ -1567,19 +1567,40 @@ def gencudaunittests(specdict):
                                 count += 1
                             else:
                                 args += ", " + arg.name
-                        f.write(" " * 4 + "funcC(" + args + ")\n")
+                        # Determine if this is a cuda.compute kernel (raises errors eagerly)
+                        # or compiled CUDA kernel (raises errors after `ak_cu.synchronize_cuda()`)
+                        CUDA_COPUTE_KERNELS = {
+                            "awkward_ListArray_compact_offsets",
+                            "awkward_ListArray_broadcast_tooffsets",
+                        }
+
+                        raises_error_eagerly = (
+                            spec.templatized_kernel_name in CUDA_COPUTE_KERNELS
+                        )
+
                         if test["error"]:
-                            f.write(
-                                f"""
-    error_message = re.escape("{test["message"]} in compiled CUDA code ({spec.templatized_kernel_name})")
-"""
-                            )
-                            f.write(
-                                """    with pytest.raises(ValueError, match=rf"{error_message}"):
-        ak_cu.synchronize_cuda()
-"""
-                            )
+                            error_message_line = f'    error_message = re.escape("{test["message"]} in compiled CUDA code ({spec.templatized_kernel_name})")\n'
+                            if raises_error_eagerly:
+                                # call a kernel directly inside `pytest.raises()`
+                                f.write(
+                                    "\n"
+                                    + error_message_line
+                                    + '    with pytest.raises(ValueError, match=rf"{error_message}"):\n'
+                                    + " " * 8
+                                    + "funcC("
+                                    + args
+                                    + ")\n"
+                                )
+                            else:
+                                f.write(" " * 4 + "funcC(" + args + ")\n")
+                                f.write(
+                                    "\n"
+                                    + error_message_line
+                                    + '    with pytest.raises(ValueError, match=rf"{error_message}"):\n'
+                                    "        ak_cu.synchronize_cuda()\n"
+                                )
                         else:
+                            f.write(" " * 4 + "funcC(" + args + ")\n")
                             f.write(
                                 """
     try:

src/awkward/_backends/cupy.py

@@ -85,6 +85,18 @@ def _supports_cuda_compute(self, kernel_name: str) -> bool:
             "awkward_reduce_argmin_complex",
             "awkward_reduce_count_64",
             "awkward_reduce_countnonzero",
+            "awkward_IndexedArray_overlay_mask",
+            "awkward_IndexedArray_reduce_next_64",
+            "awkward_IndexedArray_reduce_next_nonlocal_nextshifts_64",
+            "awkward_ByteMaskedArray_getitem_nextcarry",
+            "awkward_ByteMaskedArray_numnull",
+            "awkward_RegularArray_getitem_jagged_expand",
+            "awkward_UnionArray_simplify_one",
+            "awkward_ListArray_broadcast_tooffsets",
+            "awkward_ListArray_localindex",
+            "awkward_ListArray_compact_offsets",
+            "awkward_ListArray_combinations_length",
+            "awkward_ListArray_combinations",
             "awkward_reduce_countnonzero_complex",
             # indexing / structure
             "awkward_missing_repeat",
@@ -128,6 +140,18 @@ def _get_cuda_compute_impl(self, kernel_name: str):
             "awkward_missing_repeat": cuda_compute.awkward_missing_repeat,
             "awkward_index_rpad_and_clip_axis0": cuda_compute.awkward_index_rpad_and_clip_axis0,
             "awkward_index_rpad_and_clip_axis1": cuda_compute.awkward_index_rpad_and_clip_axis1,
+            "awkward_IndexedArray_overlay_mask": cuda_compute.awkward_IndexedArray_overlay_mask,
+            "awkward_IndexedArray_reduce_next_64": cuda_compute.awkward_IndexedArray_reduce_next_64,
+            "awkward_IndexedArray_reduce_next_nonlocal_nextshifts_64": cuda_compute.awkward_IndexedArray_reduce_next_nonlocal_nextshifts_64,
+            "awkward_ByteMaskedArray_getitem_nextcarry": cuda_compute.awkward_ByteMaskedArray_getitem_nextcarry,
+            "awkward_ByteMaskedArray_numnull": cuda_compute.awkward_ByteMaskedArray_numnull,
+            "awkward_RegularArray_getitem_jagged_expand": cuda_compute.awkward_RegularArray_getitem_jagged_expand,
+            "awkward_UnionArray_simplify_one": cuda_compute.awkward_UnionArray_simplify_one,
+            "awkward_ListArray_broadcast_tooffsets": cuda_compute.awkward_ListArray_broadcast_tooffsets,
+            "awkward_ListArray_localindex": cuda_compute.awkward_ListArray_localindex,
+            "awkward_ListArray_compact_offsets": cuda_compute.awkward_ListArray_compact_offsets,
+            "awkward_ListArray_combinations_length": cuda_compute.awkward_ListArray_combinations_length,
+            "awkward_ListArray_combinations": cuda_compute.awkward_ListArray_combinations,
         }.get(kernel_name)
 
     def prepare_reducer(self, reducer: ak._reducers.Reducer) -> ak._reducers.Reducer: