feat: migrate IndexedArray simplify to cuda.compute

dev/generate-kernel-signatures.py

@@ -57,7 +57,6 @@ def strip_stamp(text):
     # "awkward_IndexedArray_reduce_next_fix_offsets_64",
     # "awkward_RegularArray_getitem_next_array_advanced",
     # "awkward_ByteMaskedArray_toIndexedOptionArray",
-    "awkward_IndexedArray_simplify",
     # "awkward_UnionArray_validity",
     # "awkward_IndexedArray_validity",
     # "awkward_ByteMaskedArray_overlay_mask",

dev/generate-tests.py

@@ -17,12 +17,14 @@
 
 CURRENT_DIR = os.path.dirname(os.path.realpath(__file__))
 
-# Kernels dispatched to `cuda.compute` rather than compiled CUDA: these raise
-# errors eagerly (inside the call) instead of after `ak_cu.synchronize_cuda()`.
+# Kernels dispatched to `cuda.compute` rather than compiled CUDA: their errors
+# are raised eagerly (or checked by the generated KernelError adapter), rather
+# than after `ak_cu.synchronize_cuda()`.
 # IMPORTANT: whenever an error-raising kernel moves to `cuda.compute` (or back),
 # this set must be updated to match, or the generated CUDA unit tests will
 # wrap `pytest.raises` around the wrong statement.
 CUDA_COMPUTE_KERNELS = {
+    "awkward_IndexedArray_simplify",
     "awkward_ListOffsetArray_reduce_nonlocal_outstartsstops_64",
     "awkward_ListArray_compact_offsets",
     "awkward_ListArray_broadcast_tooffsets",
@@ -37,6 +39,25 @@
     "awkward_UnionArray_validity",
 }
 
+
+# These compute kernels return the CPU KernelError protocol. High-level callers
+# consume that record; standalone generated tests turn it into an exception.
+CUDA_RETURN_ERROR_KERNELS = {
+    "awkward_IndexedArray_simplify",
+}
+
+CUDA_CHECK_ERROR_HELPER = """def _checked_kernel(kernel):
+    def checked(*args):
+        error = kernel(*args)
+        if error is not None and error.str is not None:
+            message = error.str.decode(errors="surrogateescape")
+            raise ValueError(f"{message} in compiled CUDA code ({kernel.key[0]})")
+        return error
+    return checked
+
+
+"""
+
 # The C loader may be missing (pyyaml built without libyaml) and is avoided on
 # free-threaded builds.
 if getattr(sys, "_is_gil_enabled", lambda: True)():
@@ -1061,6 +1082,7 @@ def gencudakerneltests(specdict):
                     "import awkward._connect.cuda as ak_cu\n"
                     "from awkward._backends.cupy import CupyBackend\n\n"
                     "cupy_backend = CupyBackend.instance()\n\n"
+                    + CUDA_CHECK_ERROR_HELPER
                 )
                 num = 1
                 if spec.tests == []:
@@ -1101,6 +1123,8 @@ def gencudakerneltests(specdict):
                         + "', {}]\n".format(", ".join(dtypes))
                     )
                     f.write(" " * 4 + cuda_string)
+                    if spec.templatized_kernel_name in CUDA_RETURN_ERROR_KERNELS:
+                        f.write("    funcC = _checked_kernel(funcC)\n")
                     args = ""
                     count = 0
                     for arg in spec.args:
@@ -1212,7 +1236,7 @@ def gencudaunittests_new(specdict):
                 "import awkward as ak\n"
                 "import awkward._connect.cuda as ak_cu\n"
                 "from awkward._backends.cupy import CupyBackend\n\n"
-                "cupy_backend = CupyBackend.instance()\n\n"
+                "cupy_backend = CupyBackend.instance()\n\n" + CUDA_CHECK_ERROR_HELPER
             )
 
             # CASES list
@@ -1231,9 +1255,13 @@ def gencudaunittests_new(specdict):
             f.write("@pytest.fixture(scope='session')\n")
             f.write(f"def funcC_{spec.name}():\n")
             f.write(
-                f"    return cupy_backend['{spec.templatized_kernel_name}',"
-                f" {', '.join(dtypes)}]\n\n"
+                f"    kernel = cupy_backend['{spec.templatized_kernel_name}',"
+                f" {', '.join(dtypes)}]\n"
             )
+            if spec.templatized_kernel_name in CUDA_RETURN_ERROR_KERNELS:
+                f.write("    return _checked_kernel(kernel)\n\n")
+            else:
+                f.write("    return kernel\n\n")
 
             if not status:
                 f.write(
@@ -1366,6 +1394,14 @@ def gencudaunittests_new(specdict):
 
             # --- Kernel call ---
             args_str = ", ".join(arg.name for arg in spec.args)
+            if spec.templatized_kernel_name in CUDA_COMPUTE_KERNELS:
+                f.write(
+                    "\n    if case['error']:\n"
+                    "        error_message = re.escape(case['message'])\n"
+                    "        with pytest.raises(ValueError, match=error_message):\n"
+                    f"            funcC_{spec.name}({args_str})\n"
+                    "        return\n"
+                )
             f.write(f"\n    funcC_{spec.name}({args_str})\n")
 
             # --- Error vs success path ---
@@ -1499,6 +1535,7 @@ def gencudaunittests(specdict):
                     "import awkward._connect.cuda as ak_cu\n"
                     "from awkward._backends.cupy import CupyBackend\n\n"
                     "cupy_backend = CupyBackend.instance()\n\n"
+                    + CUDA_CHECK_ERROR_HELPER
                 )
                 unit_test_values = unit_test_map[spec.templatized_kernel_name]
                 tests = unit_test_values["tests"]
@@ -1577,6 +1614,8 @@ def gencudaunittests(specdict):
                             + "', {}]\n".format(", ".join(dtypes))
                         )
                         f.write(" " * 4 + cuda_string)
+                        if spec.templatized_kernel_name in CUDA_RETURN_ERROR_KERNELS:
+                            f.write("    funcC = _checked_kernel(funcC)\n")
                         args = ""
                         count = 0
                         for arg in spec.args:

kernel-specification.yml

@@ -953,7 +953,12 @@ kernels:
           - {name: outerlength, type: "int64_t", dir: in, role: default}
           - {name: innerindex, type: "Const[List[uint32_t]]", dir: in, role: IndexedArray2-index}
           - {name: innerlength, type: "int64_t", dir: in, role: IndexedArray-length}
-    description: null
+    description: |
+      Compose outer and inner indices into int64 output. Every negative outer
+      index maps to -1, while negative inner indices are copied unchanged.
+      On failure, report the first outer row and its original index, preserving
+      the already-written prefix and leaving the failing row and tail untouched.
+      Buffers must cover declared lengths; input and output do not overlap.
     definition: |
       def awkward_IndexedArray_simplify(
           toindex, outerindex, outerlength, innerindex, innerlength

kernel-test-data.json

@@ -12289,6 +12289,137 @@
           "outputs": {
             "toindex": [6,4,2,0]
           }
+        },
+        {
+          "inputs": {
+            "outerindex": [
+              1,
+              -9,
+              0
+            ],
+            "innerindex": [
+              2,
+              -7,
+              1
+            ],
+            "outerlength": 3,
+            "innerlength": 3
+          },
+          "outputs": {
+            "toindex": [
+              -7,
+              -1,
+              2,
+              123,
+              123,
+              123
+            ]
+          },
+          "error": false,
+          "message": ""
+        },
+        {
+          "inputs": {
+            "outerindex": [
+              0,
+              2,
+              1
+            ],
+            "innerindex": [
+              9,
+              7
+            ],
+            "outerlength": 3,
+            "innerlength": 2
+          },
+          "outputs": {
+            "toindex": [
+              9,
+              123,
+              123,
+              123,
+              123,
+              123
+            ]
+          },
+          "error": true,
+          "message": "index out of range"
+        },
+        {
+          "inputs": {
+            "outerindex": [
+              -7,
+              1,
+              9,
+              0
+            ],
+            "innerindex": [
+              5,
+              8
+            ],
+            "outerlength": 4,
+            "innerlength": 2
+          },
+          "outputs": {
+            "toindex": [
+              -1,
+              8,
+              123,
+              123,
+              123,
+              123,
+              123
+            ]
+          },
+          "error": true,
+          "message": "index out of range"
+        },
+        {
+          "inputs": {
+            "outerindex": [
+              2,
+              0,
+              3
+            ],
+            "innerindex": [
+              4,
+              5,
+              6
+            ],
+            "outerlength": 3,
+            "innerlength": 2
+          },
+          "outputs": {
+            "toindex": [
+              123,
+              123,
+              123,
+              123,
+              123,
+              123
+            ]
+          },
+          "error": true,
+          "message": "index out of range"
+        },
+        {
+          "inputs": {
+            "outerindex": [
+              99
+            ],
+            "innerindex": [],
+            "outerlength": 0,
+            "innerlength": 0
+          },
+          "outputs": {
+            "toindex": [
+              123,
+              123,
+              123
+            ]
+          },
+          "error": false,
+          "message": ""
         }
       ]
     },

src/awkward/_backends/cupy.py

@@ -62,6 +62,7 @@ def _new_kernel(
     # ---------------------------------------------------------
     def _supports_cuda_compute(self, kernel_name: str) -> bool:
         return kernel_name in (
+            "awkward_IndexedArray_simplify",
             "awkward_ListOffsetArray_reduce_nonlocal_outstartsstops_64",
             # core reducers
             "awkward_reduce_sum",
@@ -187,6 +188,7 @@ def _get_cuda_compute_impl(self, kernel_name: str):
         from awkward._connect.cuda import _compute as cuda_compute
 
         return {
+            "awkward_IndexedArray_simplify": cuda_compute.awkward_IndexedArray_simplify,
             "awkward_ListOffsetArray_reduce_nonlocal_outstartsstops_64": cuda_compute.awkward_ListOffsetArray_reduce_nonlocal_outstartsstops_64,
             "awkward_sort": cuda_compute.segmented_sort,
             "awkward_argsort": cuda_compute.segmented_argsort,