diff --git a/dev/generate-kernel-signatures.py b/dev/generate-kernel-signatures.py index 98c992b45f..57c0dabd02 100644 --- a/dev/generate-kernel-signatures.py +++ b/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", diff --git a/dev/generate-tests.py b/dev/generate-tests.py index 0872f8113b..6971c4a1a2 100644 --- a/dev/generate-tests.py +++ b/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: diff --git a/kernel-specification.yml b/kernel-specification.yml index 14abc7166a..24f5944417 100644 --- a/kernel-specification.yml +++ b/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 diff --git a/kernel-test-data.json b/kernel-test-data.json index 876b57b89b..7069528879 100644 --- a/kernel-test-data.json +++ b/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": "" } ] }, diff --git a/src/awkward/_backends/cupy.py b/src/awkward/_backends/cupy.py index a4b6aae327..e57773ac00 100644 --- a/src/awkward/_backends/cupy.py +++ b/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, diff --git a/src/awkward/_connect/cuda/_compute.py b/src/awkward/_connect/cuda/_compute.py index f7565d67af..aafe526cea 100644 --- a/src/awkward/_connect/cuda/_compute.py +++ b/src/awkward/_connect/cuda/_compute.py @@ -2,6 +2,7 @@ import functools +from types import SimpleNamespace from cuda.compute import ( CountingIterator, @@ -4159,37 +4160,90 @@ def is_null(x): ) -# KERNEL IS NOT USED: checking for out of range errors takes too much time if I take it outside the closure -# composes two index arrays into one by resolving outerindex through innerindex: -# toindex[i] = innerindex[outerindex[i]], preserving -1 (missing) entries from outerindex def awkward_IndexedArray_simplify( toindex, outerindex, outerlength, innerindex, innerlength ): if outerlength == 0: - return + return None + if hasattr(outerindex, "__cuda_array_interface__"): + from awkward._connect.cuda._indexedarray import simplify_status + + row, attempt = ( + simplify_status(toindex, outerindex, outerlength, innerindex, innerlength) + .get() + .tolist() + ) + if row < outerlength: + return SimpleNamespace( + str=b"index out of range", + id=row, + attempt=attempt, + filename=b"cuda.compute (awkward_IndexedArray_simplify)", + ) + return None + outer = outerindex[:outerlength] + inner = innerindex[:innerlength] + parameters = cp.asarray([innerlength], dtype=cp.int64) + current_stream = cp.cuda.get_current_stream() + # Adapt older CuPy streams to the CUDA Stream Protocol as well. + stream = SimpleNamespace(__cuda_stream__=lambda: (0, current_stream.ptr)) + + def error_position(i): + index = np.int64(outer[i]) + return i if index >= 0 and index >= parameters[0] else np.int64(outer.size) + + indices = CountingIterator(np.int64(0)) + # Plain array input avoids stateful TransformIterator launch failures. + positions = cp.empty(outerlength, dtype=cp.int64) + unary_transform( + d_in=indices, + d_out=positions, + op=error_position, + num_items=outerlength, + stream=stream, + ) + first = cp.empty(1, dtype=cp.int64) + reduce_into( + d_in=positions, + d_out=first, + op=OpKind.MINIMUM, + num_items=outerlength, + h_init=np.asarray(outerlength, dtype=np.int64), + stream=stream, + ) - def simplify_op(j): - if j < 0: - return -1 - if j >= innerlength: - # raise inside a JIT closure is silently swallowed - raise IndexError( - "index out of range in compiled CUDA code (awkward_IndexedArray_simplify)" + status = cp.empty(2, dtype=cp.int64) + + def write(i): + # Match the CPU's sequential output prefix, without ever reading an + # invalid inner index. Leave the error row and remaining capacity alone. + if i < first[0]: + index = np.int64(outer[i]) + toindex[i] = -1 if index < 0 else np.int64(inner[index]) + if i == 0: + row = first[0] + status[0] = row + status[1] = ( + outer[row] if row < outer.size else np.int64(9223372036854775807) ) - return innerindex[j] + return np.int8(0) - # to prevent out-of-bounds error - out_buf = ( - toindex[:outerlength] - if len(toindex) >= outerlength - else cp.empty(outerlength, dtype=toindex.dtype) - ) unary_transform( - d_in=outerindex[:outerlength], - d_out=out_buf, - op=simplify_op, + d_in=CountingIterator(np.int64(0)), + d_out=DiscardIterator(), + op=write, num_items=outerlength, + stream=stream, ) + row, attempt = status.get().tolist() + if row < outerlength: + return SimpleNamespace( + str=b"index out of range", + id=row, + attempt=attempt, + filename=b"cuda.compute (awkward_IndexedArray_simplify)", + ) + return None # filters an indexed array by collecting all valid (non-negative) indices into a carry array diff --git a/src/awkward/_connect/cuda/_indexedarray.py b/src/awkward/_connect/cuda/_indexedarray.py new file mode 100644 index 0000000000..74e5b7e19e --- /dev/null +++ b/src/awkward/_connect/cuda/_indexedarray.py @@ -0,0 +1,204 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE + +"""Prepared IndexedArray simplification with canonical first-error writes.""" + +from functools import lru_cache +from threading import RLock + +from cuda.compute import ( + CountingIterator, + DiscardIterator, + make_reduce_into, + make_unary_transform, +) +from cuda.compute.op import make_op_adapter + +from awkward._nplikes.cupy import Cupy +from awkward._nplikes.numpy import Numpy + +cp = Cupy.instance()._module +np = Numpy.instance()._module +_lock = RLock() + + +class _Stream: + def __init__(self): + self.stream = cp.cuda.get_current_stream() + + def __cuda_stream__(self): + return (0, self.stream.ptr) + + +class _ArrayState: + # CCCL rewrites captured CUDA-array-interface objects into device arrays. + # These proxies carry launch metadata only; indexing in compiled operators + # acts on the replacement arrays, not on this Python class. + def __init__(self, dtype): + self.dtype = np.dtype(dtype) + self.clear() + + def bind(self, array): + interface = array.__cuda_array_interface__ + self.__cuda_array_interface__ = { + "shape": tuple(interface["shape"]), + "strides": interface.get("strides"), + "typestr": interface["typestr"], + "data": (interface["data"][0], False), + "version": 3, + } + + def __len__(self): + return self.__cuda_array_interface__["shape"][0] + + def data_ptr(self): + return self.__cuda_array_interface__["data"][0] + + def clear(self): + self.__cuda_array_interface__ = { + "shape": (0,), + "strides": None, + "typestr": self.dtype.str, + "data": (0, False), + "version": 3, + } + + +class _Simplify: + def __init__(self, output_dtype, outer_dtype, inner_dtype, length, innerlength): + self.output = output = _ArrayState(output_dtype) + self.outer = outer = _ArrayState(outer_dtype) + self.inner = inner = _ArrayState(inner_dtype) + self.status = status = _ArrayState(np.int64) + + def small(_): + # CCCL substitutes device arrays for metadata proxies. + # pylint: disable=unsupported-assignment-operation + status[0] = length + status[1] = np.int64(9223372036854775807) + for row in range(length): + index = np.int64(outer[row]) + if index >= 0 and index >= innerlength: + status[0] = row + status[1] = index + break + output[row] = -1 if index < 0 else np.int64(inner[index]) + return np.int8(0) + + def combine(left, right): + # Map valid row numbers to the length sentinel, then take MIN. + # The map is idempotent, making this associative/commutative on + # row numbers plus the sentinel. No stateful iterator is needed. + a = np.int64(length) + if left < length: + index = np.int64(outer[left]) + if index >= 0 and index >= innerlength: + a = left + b = np.int64(length) + if right < length: + index = np.int64(outer[right]) + if index >= 0 and index >= innerlength: + b = right + return a if a < b else b + + def write(row): + # pylint: disable=unsupported-assignment-operation + first = status[0] + if row < first: + index = np.int64(outer[row]) + output[row] = -1 if index < 0 else np.int64(inner[index]) + if row == 0: + status[1] = ( + np.int64(outer[first]) + if first < length + else np.int64(9223372036854775807) + ) + return np.int8(0) + + self.small, self.combine, self.write = small, combine, write + self.length = length + self.iterator = CountingIterator(np.int64(0)) + self.discard = DiscardIterator() + self.transforms = {} + self.reducer = None + self.reduce_storage_bytes = None + self.reduce_operation = None + self.initial = np.asarray(length, dtype=np.int64) + + def transform(self, name, operation, output, length, stream): + if name not in self.transforms: + adapter = make_op_adapter(operation) + prepared = make_unary_transform( + d_in=self.iterator, d_out=output, op=adapter + ) + self.transforms[name] = (adapter, prepared) + adapter, prepared = self.transforms[name] + prepared( + d_in=self.iterator, + d_out=output, + op=adapter, + num_items=length, + stream=stream, + ) + + def __call__(self, output, outer, inner, stream): + # Status and reduction workspace belong to this invocation; cached + # operators retain only descriptors. The small path remains fused. + status = cp.empty(2, dtype=cp.int64) + states = (self.output, self.outer, self.inner, self.status) + for state, array in zip(states, (output, outer, inner, status), strict=True): + state.bind(array) + try: + if self.length <= 64: + self.transform("small", self.small, self.discard, 1, stream) + else: + first = status[:1] + if self.reducer is None: + self.reduce_operation = make_op_adapter(self.combine) + self.reducer = make_reduce_into( + d_in=self.iterator, + d_out=first, + op=self.reduce_operation, + h_init=self.initial, + ) + arguments = { + "d_in": self.iterator, + "d_out": first, + "num_items": self.length, + "op": self.reduce_operation, + "h_init": self.initial, + "stream": stream, + } + # The cache key fixes device/context, dtypes and item count. + # Only sizing is reused; workspace itself remains per-call. + if self.reduce_storage_bytes is None: + self.reduce_storage_bytes = self.reducer( + temp_storage=None, **arguments + ) + storage = cp.empty(self.reduce_storage_bytes, dtype=cp.uint8) + self.reducer(temp_storage=storage, **arguments) + self.transform("write", self.write, self.discard, self.length, stream) + finally: + for state in states: + state.clear() + return status + + +@lru_cache(maxsize=32) +def _setup( + device, context, output_dtype, outer_dtype, inner_dtype, length, innerlength +): + return _Simplify(output_dtype, outer_dtype, inner_dtype, length, innerlength) + + +def simplify_status(output, outer, length, inner, innerlength): + with _lock: + setup = _setup( + cp.cuda.runtime.getDevice(), + cp.cuda.driver.ctxGetCurrent(), + output.dtype.str, + outer.dtype.str, + inner.dtype.str, + length, + innerlength, + ) + return setup(output, outer, inner, _Stream()) diff --git a/src/awkward/_connect/cuda/cuda_kernels/awkward_IndexedArray_simplify.cu b/src/awkward/_connect/cuda/cuda_kernels/awkward_IndexedArray_simplify.cu deleted file mode 100644 index ca7ec9b107..0000000000 --- a/src/awkward/_connect/cuda/cuda_kernels/awkward_IndexedArray_simplify.cu +++ /dev/null @@ -1,33 +0,0 @@ -// BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE - -enum class INDEXEDARRAY_SIMPLIFY_ERRORS { - IND_OUT_OF_RANGE // message: "index out of range" -}; - -template -__global__ void -awkward_IndexedArray_simplify( - T* toindex, - const C* outerindex, - int64_t outerlength, - const U* innerindex, - int64_t innerlength, - 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 < outerlength) { - C j = outerindex[thread_id]; - if (j < 0) { - toindex[thread_id] = -1; - - } else { - if (j >= innerlength) { - RAISE_ERROR(INDEXEDARRAY_SIMPLIFY_ERRORS::IND_OUT_OF_RANGE) - } else { - toindex[thread_id] = innerindex[j]; - } - } - } - } -} diff --git a/src/awkward/_errors.py b/src/awkward/_errors.py index c4b8c2428f..e4f76b4a62 100644 --- a/src/awkward/_errors.py +++ b/src/awkward/_errors.py @@ -38,6 +38,33 @@ def frozen_primary(cls): context._frozen = True return context + @classmethod + def cuda_primary(cls): + """Capture bounded, host-only diagnostic metadata without retaining buffers.""" + context = cls.primary() + if context is None: + return None + snapshot = context.__dict__.get("_cuda_snapshot") + if snapshot is None: + if isinstance(context, OperationErrorContext): + snapshot = OperationErrorContext(context.name, (), {}) + snapshot.args = [_cuda_argument(value) for value in context._raw_args] + snapshot.kwargs = { + key: _cuda_argument(value) + for key, value in context._raw_kwargs.items() + if isinstance(key, str) + } + elif isinstance(context, SlicingErrorContext): + snapshot = SlicingErrorContext(None, None) + snapshot.array = _cuda_argument(context._raw_array) + snapshot.where = _cuda_argument(context._raw_where) + else: + # Unknown context subclasses may retain arbitrary arguments. + # Preserve the context name without calling custom formatters. + snapshot = OperationErrorContext(type(context).__name__, (), {}) + context._cuda_snapshot = snapshot + return snapshot + def freeze(self) -> None: """Format the note now and drop the references to the raw arguments.""" @@ -71,7 +98,11 @@ def _add_note(exception: E, note: str) -> E: exception.__notes__ = [note] return exception - note = self.note + # CUDA dispatch already captured a host-only note. Reuse it for an + # immediate error too: formatting the original context would read device + # values while unwinding the public call. + context = self.__dict__.get("_cuda_snapshot", self) + note = context.note if issubclass(cls, (NotImplementedError, AssertionError)): note = "\n\nSee if this has been reported at https://github.com/scikit-hep/awkward/issues" return _add_note(exception, note) @@ -325,3 +356,31 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T: return func(*args, **kwargs) return wrapper + + +def _cuda_argument(value, *, depth=0): + """Never inspect array values: formatting a CUDA array can synchronize.""" + from awkward import contents, highlevel, index, record + + if isinstance(value, (highlevel.Array, highlevel.Record)): + return f"<{type(value).__name__} type='{value.type}'>" + if isinstance(value, contents.Content): + return f"<{type(value).__name__} length={len(value)}>" + if isinstance(value, record.Record): + return "" + if isinstance(value, index.Index): + return f"<{type(value).__name__} length={len(value)} dtype={value.dtype}>" + if isinstance(value, slice): + return SlicingErrorContext.format_slice(value) + if value is None or type(value) in (bool, int, float, str, bytes): + if isinstance(value, (str, bytes)): + value = value[:72] + return repr(value) + if type(value) in (tuple, list) and depth < 3: + parts = [_cuda_argument(x, depth=depth + 1) for x in value[:8]] + if len(value) > 8: + parts.append("...") + if isinstance(value, tuple): + return "(" + ", ".join(parts) + ("," if len(value) == 1 else "") + ")" + return "[" + ", ".join(parts) + "]" + return f"<{type(value).__name__}>" diff --git a/src/awkward/_kernels.py b/src/awkward/_kernels.py index 0b18580358..409115b989 100644 --- a/src/awkward/_kernels.py +++ b/src/awkward/_kernels.py @@ -231,7 +231,7 @@ def __call__(self, *args) -> None: ak_cuda.cuda_streamptr_to_contexts[cupy_stream_ptr][1].append( ak_cuda.Invocation( name=self.key[0], - error_context=ak._errors.ErrorContext.frozen_primary(), + error_context=ak._errors.ErrorContext.cuda_primary(), ) ) @@ -251,7 +251,7 @@ def __init__(self, impl: Callable[..., Any], key: KernelKeyType): super().__init__(impl, key) self._cupy = Cupy.instance() - def __call__(self, *args) -> None: + def __call__(self, *args) -> KernelError | None: import awkward._connect.cuda as ak_cuda args = maybe_materialize(*args) @@ -275,7 +275,7 @@ def __call__(self, *args) -> None: ak_cuda.cuda_streamptr_to_contexts[cupy_stream_ptr][1].append( ak_cuda.Invocation( name=self.key[0], - error_context=ak._errors.ErrorContext.frozen_primary(), + error_context=ak._errors.ErrorContext.cuda_primary(), ) ) diff --git a/tests-cuda/test_4481_cuda_error_context.py b/tests-cuda/test_4481_cuda_error_context.py new file mode 100644 index 0000000000..537c8b69d9 --- /dev/null +++ b/tests-cuda/test_4481_cuda_error_context.py @@ -0,0 +1,32 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE + +import gc +import weakref + +import pytest + +import awkward as ak + +cp = pytest.importorskip("cupy") + + +def test_successful_slicing_does_not_format_or_retain_arrays(monkeypatch): + import awkward._connect.cuda as cuda + + array = ak.to_backend( + ak.to_regular([[[1, 2], [3, 4]], [[5, 6], [7, 8]]], axis=-1), "cuda" + ) + indices = ak.to_backend([[0], [1]], "cuda") + refs = [weakref.ref(array), weakref.ref(indices)] + + def forbidden(*args, **kwargs): + raise AssertionError("successful CUDA slicing must not format values") + + with monkeypatch.context() as patch: + patch.setattr(ak.Array, "_repr", forbidden) + result = array[:, indices] + del array, indices + gc.collect() + assert all(ref() is None for ref in refs) + cuda.synchronize_cuda(cp.cuda.get_current_stream()) + assert ak.to_list(result) == [[[1], [4]], [[5], [8]]] diff --git a/tests-cuda/test_4484_cuda_error_context.py b/tests-cuda/test_4484_cuda_error_context.py new file mode 100644 index 0000000000..2de811e226 --- /dev/null +++ b/tests-cuda/test_4484_cuda_error_context.py @@ -0,0 +1,37 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE + +import numpy as np +import pytest + +import awkward as ak +from awkward._errors import ErrorContext + +cp = pytest.importorskip("cupy") + + +@pytest.mark.parametrize("option", [False, True]) +@pytest.mark.parametrize("length", [32, 4096]) +def test_invalid_field_projection_uses_metadata_note(monkeypatch, option, length): + inner_type = ak.contents.IndexedOptionArray if option else ak.contents.IndexedArray + inner = inner_type( + ak.index.Index64(np.array([1, 0], dtype=np.int64)), + ak.contents.NumpyArray(np.array([10, 20])), + ) + indices = np.zeros(length, dtype=np.int64) + indices[length // 2] = 2 + array = ak.to_backend( + ak.Array( + ak.contents.IndexedArray( + ak.index.Index64(indices), ak.contents.RecordArray([inner], ["x"]) + ) + ), + "cuda", + ) + + def forbidden(*args, **kwargs): + raise AssertionError("Invalid CUDA projection must not format device values") + + monkeypatch.setattr(ErrorContext, "format_argument", forbidden) + with pytest.raises(ValueError, match="index out of range") as caught: + array["x"] + assert any("type='" in note for note in caught.value.__notes__) diff --git a/tests-cuda/test_4484_indexedarray_simplify_compute.py b/tests-cuda/test_4484_indexedarray_simplify_compute.py new file mode 100644 index 0000000000..8114022956 --- /dev/null +++ b/tests-cuda/test_4484_indexedarray_simplify_compute.py @@ -0,0 +1,153 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE + +import importlib.util +import sys +from pathlib import Path + +import numpy as np +import pytest + +import awkward as ak +from awkward._backends.cupy import CupyBackend +from awkward._backends.numpy import NumpyBackend + +cp = pytest.importorskip("cupy") +pytest.importorskip("cuda.compute") + +spec = importlib.util.spec_from_file_location( + "indexedarray_simplify_cases", + Path(__file__).resolve().parents[1] + / "tests/test_4484_indexedarray_simplify_compute.py", +) +contract = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = contract +spec.loader.exec_module(contract) + + +@pytest.mark.parametrize("name,specialization,case", contract.CASES) +def test_cuda_cpu_parity(name, specialization, case): + with cp.cuda.Stream(non_blocking=True): + _cuda_cpu_parity(name, specialization, case) + + +def _cuda_cpu_parity(name, specialization, case): + args, outputs, key = contract.arguments(specialization, case) + gpu_args = [cp.asarray(x) if isinstance(x, np.ndarray) else x for x in args] + expected = NumpyBackend.instance()[(name, *key)](*args) + kernel = CupyBackend.instance()[(name, *key)] + assert isinstance(kernel, ak._kernels.CudaComputeKernel) + actual = kernel(*gpu_args) + contract.assert_error(actual, expected) + for i in outputs: + np.testing.assert_array_equal(cp.asnumpy(gpu_args[i]), args[i]) + + +@pytest.mark.parametrize("outer_dtype", [np.int32, np.uint32, np.int64]) +@pytest.mark.parametrize("inner_dtype", [np.int32, np.uint32, np.int64]) +@pytest.mark.parametrize("option", [False, True]) +def test_public_simplification(outer_dtype, inner_dtype, option): + inner_class = ( + ak.contents.IndexedOptionArray + if option and inner_dtype != np.uint32 + else ak.contents.IndexedArray + ) + inner_values = ( + [2, -1, 0] if inner_class is ak.contents.IndexedOptionArray else [2, 1, 0] + ) + inner = inner_class( + ak.index.Index(np.asarray(inner_values, dtype=inner_dtype)), + ak.contents.NumpyArray(np.asarray([10, 20, 30])), + ) + outer_values = [1, -1, 0] if option and outer_dtype != np.uint32 else [1, 0, 2] + outer_class = ak.contents.IndexedOptionArray if option else ak.contents.IndexedArray + expected = outer_class.simplified( + ak.index.Index(np.asarray(outer_values, dtype=outer_dtype)), inner + ).to_list() + with cp.cuda.Stream(non_blocking=True) as stream: + gpu_inner = ak.to_backend(inner, "cuda", highlevel=False) + outer = ak.index.Index(cp.asarray(outer_values, dtype=outer_dtype)) + assert outer_class.simplified(outer, gpu_inner).to_list() == expected + with pytest.raises(ValueError, match="index out of range"): + outer_class.simplified( + ak.index.Index(cp.asarray([1, 3], dtype=outer_dtype)), gpu_inner + ) + ak._connect.cuda.synchronize_cuda(stream) + assert outer_class.simplified(outer, gpu_inner).to_list() == expected + + +@pytest.mark.parametrize( + "length,error_at", [(33, None), (257, 0), (4097, 256), (4097, 4096)] +) +def test_first_error_across_blocks(length, error_at): + args = contract.large_case(length, error_at) + key = ("awkward_IndexedArray_simplify", np.int64, np.int64, np.int64) + with cp.cuda.Stream(non_blocking=True): + gpu_args = [cp.asarray(x) if isinstance(x, np.ndarray) else x for x in args] + expected = NumpyBackend.instance()[key](*args) + actual = CupyBackend.instance()[key](*gpu_args) + contract.assert_error(actual, expected) + np.testing.assert_array_equal(cp.asnumpy(gpu_args[0]), args[0]) + + +@pytest.mark.parametrize("outer_dtype", [np.int32, np.uint32, np.int64]) +@pytest.mark.parametrize("inner_dtype", [np.int32, np.uint32, np.int64]) +def test_prepared_rebinding_and_chunk_boundary(outer_dtype, inner_dtype): + from awkward._connect.cuda._indexedarray import simplify_status + + pending = [] + key = ("awkward_IndexedArray_simplify", np.int64, outer_dtype, inner_dtype) + for length in (32, 64, 65, 129): + for invalid in (False, True): + outer = (np.arange(length) % 3).astype(outer_dtype) + inner = np.asarray([2, 1, 0], dtype=inner_dtype) + if invalid: + outer[min(length - 1, 64)] = 3 + outer[-1] = 99 + expected = np.full(length + 3, 123, dtype=np.int64) + error = NumpyBackend.instance()[key](expected, outer, length, inner, 3) + stream = cp.cuda.Stream(non_blocking=True) + with stream: + gpu_outer, gpu_inner = cp.asarray(outer), cp.asarray(inner) + output = cp.full(length + 3, 123, dtype=cp.int64) + status = simplify_status(output, gpu_outer, length, gpu_inner, 3) + pending.append( + (stream, gpu_outer, gpu_inner, output, status, expected, error, length) + ) + for stream, _outer, _inner, output, status, expected, error, length in pending: + stream.synchronize() + np.testing.assert_array_equal(cp.asnumpy(output), expected) + row, attempt = cp.asnumpy(status) + if error.str: + assert (row, attempt) == (error.id, error.attempt) + else: + assert row == length + + +@pytest.mark.parametrize("outer_dtype", [np.int32, np.uint32, np.int64]) +@pytest.mark.parametrize("inner_dtype", [np.int32, np.uint32, np.int64]) +@pytest.mark.parametrize("length", [64, 65, 129, 4096, 4097]) +@pytest.mark.parametrize("nondefault", [False, True]) +def test_fused_dtype_stream_and_rebinding(outer_dtype, inner_dtype, length, nondefault): + stream = cp.cuda.Stream(non_blocking=True) if nondefault else cp.cuda.Stream.null + key = ("awkward_IndexedArray_simplify", np.int64, outer_dtype, inner_dtype) + with stream: + kernel = CupyBackend.instance()[key] + for error_row in (None, 0, min(64, length - 1), length - 1): + for repeat in range(2): + outer = (np.arange(length) % 3).astype(outer_dtype) + if outer_dtype is not np.uint32: + outer[::5] = -1 + inner = np.asarray([repeat + 2, repeat + 1, repeat], dtype=inner_dtype) + if inner_dtype is not np.uint32: + inner[0] = -7 - repeat + if error_row is not None: + outer[error_row] = 3 + if error_row < length - 1: + outer[-1] = 99 + expected = np.full(length + 3, 123, dtype=np.int64) + error = NumpyBackend.instance()[key](expected, outer, length, inner, 3) + output = cp.full(length + 3, 123, dtype=cp.int64) + actual = kernel(output, cp.asarray(outer), length, cp.asarray(inner), 3) + contract.assert_error(actual, error) + np.testing.assert_array_equal(cp.asnumpy(output), expected) + ak._connect.cuda.synchronize_cuda(stream) diff --git a/tests/test_4481_cuda_error_context.py b/tests/test_4481_cuda_error_context.py new file mode 100644 index 0000000000..df87380373 --- /dev/null +++ b/tests/test_4481_cuda_error_context.py @@ -0,0 +1,57 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE + +import gc +import weakref + +import pytest + +import awkward as ak +from awkward._errors import ErrorContext, OperationErrorContext, SlicingErrorContext + + +@pytest.mark.parametrize("kind", ["operation", "slice"]) +def test_cuda_capture_no_value_formatting_or_buffer_retention(monkeypatch, kind): + array = ak.Array([[1, 2], []]) + index = ak.Array([[0], []]) + references = [weakref.ref(array), weakref.ref(index)] + + def forbidden(*args, **kwargs): + raise AssertionError("CUDA diagnostics must not format array values") + + monkeypatch.setattr(ak.Array, "_repr", forbidden) + monkeypatch.setattr(ak.Array, "__repr__", forbidden) + monkeypatch.setattr(ErrorContext, "format_argument", forbidden) + if kind == "operation": + context = OperationErrorContext( + "ak.example", (array,), {"indices": index, "axis": 1} + ) + else: + context = SlicingErrorContext(array, (slice(None), index)) + with context: + snapshot = ErrorContext.cuda_primary() + assert ErrorContext.cuda_primary() is snapshot + assert "var * int64" in snapshot.note + assert "[1, 2]" not in snapshot.note + del context, array, index + gc.collect() + assert all(ref() is None for ref in references) + error = snapshot.decorate_exception(ValueError, ValueError("bad index")) + assert "type='2 * var * int64'" in error.__notes__[0] + if kind == "operation": + assert "ak.example(" in error.__notes__[0] + assert "axis = 1" in error.__notes__[0] + else: + assert "(:, " + assert "Index64" in _cuda_argument(ak.index.Index64(np.array([0]))) + assert _cuda_argument((1,)) == "(1,)" + assert _cuda_argument(b"x" * 100) == repr(b"x" * 72) + assert _cuda_argument("x" * 100) == repr("x" * 72) + assert _cuda_argument(object()) == "" + assert _cuda_argument([[[[1]]]]) == "[[[]]]" + assert _cuda_argument(None) == "None" + + +def test_unknown_cuda_context_uses_name_without_custom_formatting(): + class Custom(ErrorContext): + @property + def note(self): + raise AssertionError("Custom context must not be formatted") + + with Custom(): + snapshot = ErrorContext.cuda_primary() + assert "Custom(" in snapshot.note + assert ErrorContext.cuda_primary() is snapshot diff --git a/tests/test_4484_cuda_error_context.py b/tests/test_4484_cuda_error_context.py new file mode 100644 index 0000000000..031881c7c5 --- /dev/null +++ b/tests/test_4484_cuda_error_context.py @@ -0,0 +1,46 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE + +import pytest + +import awkward as ak +from awkward._errors import ErrorContext, OperationErrorContext, SlicingErrorContext + + +@pytest.mark.parametrize("kind", ["operation", "slice"]) +def test_immediate_cuda_error_reuses_snapshot(monkeypatch, kind): + array = ak.Array([[1, 2], []]) + context = ( + OperationErrorContext("ak.example", (array,), {"axis": 1}) + if kind == "operation" + else SlicingErrorContext(array, "x") + ) + + def forbidden(*args, **kwargs): + raise AssertionError("Immediate CUDA errors must not format device values") + + monkeypatch.setattr(ErrorContext, "format_argument", forbidden) + monkeypatch.setattr(SlicingErrorContext, "format_slice", forbidden) + error = ValueError("index out of range in compiled CUDA code") + with pytest.raises(ValueError) as caught: + with context: + snapshot = ErrorContext.cuda_primary() + raise error + assert caught.value is error + assert caught.value.__notes__ == [snapshot.note] + assert "2 * var * int64" in caught.value.__notes__[0] + assert ErrorContext.primary() is None + + +def test_cpu_error_keeps_value_formatting(): + with pytest.raises(ValueError) as caught: + with SlicingErrorContext(ak.Array([123, 456]), "x"): + raise ValueError("bad field") + assert "[123, 456]" in caught.value.__notes__[0] + + +def test_cuda_snapshot_preserves_assertion_guidance(): + with pytest.raises(AssertionError) as caught: + with OperationErrorContext("ak.example", (), {}): + ErrorContext.cuda_primary() + raise AssertionError("internal error") + assert "github.com/scikit-hep/awkward/issues" in caught.value.__notes__[0] diff --git a/tests/test_4484_indexedarray_simplify_compute.py b/tests/test_4484_indexedarray_simplify_compute.py new file mode 100644 index 0000000000..0a8fdf5a4a --- /dev/null +++ b/tests/test_4484_indexedarray_simplify_compute.py @@ -0,0 +1,602 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE + +"""CPU contract checks and host execution of the actual CCCL operator bodies. + +Host substitutes do not test GPU compilation or execution; tests-cuda runs the +same case inventory through real CUDA dispatch. +""" + +import ast +import copy +import json +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest + +import awkward as ak +from awkward._backends.numpy import NumpyBackend + +yaml = pytest.importorskip("yaml") +ROOT = Path(__file__).resolve().parents[1] +# Match the installed package covered by --cov=awkward in CI. +CUDA_SOURCE = Path(ak.__file__).resolve().parent / "_connect/cuda" +FUNCTIONS = {"awkward_IndexedArray_simplify": "awkward_IndexedArray_simplify"} + +DTYPES = {"int32_t": np.int32, "uint32_t": np.uint32, "int64_t": np.int64} + + +def cases(): + specifications = yaml.safe_load((ROOT / "kernel-specification.yml").read_text())[ + "kernels" + ] + data = { + x["name"]: x["tests"] + for x in json.loads((ROOT / "kernel-test-data.json").read_text())["unit-tests"] + } + # Width-specific extremes are filtered per specialization below. + data["awkward_IndexedArray_simplify"].extend( + [ + { + "inputs": { + "outerindex": [-1, -9, -9223372036854775808], + "innerindex": [], + "outerlength": 3, + "innerlength": 0, + }, + "outputs": {"toindex": [-1, -1, -1, 123, 123, 123]}, + "error": False, + "message": "", + }, + { + "inputs": { + "outerindex": [4294967295, 0], + "innerindex": [1], + "outerlength": 2, + "innerlength": 1, + }, + "outputs": {"toindex": [123, 123, 123, 123, 123]}, + "error": True, + "message": "index out of range", + }, + { + "inputs": { + "outerindex": [9223372036854775807, 0], + "innerindex": [1], + "outerlength": 2, + "innerlength": 1, + }, + "outputs": {"toindex": [123, 123, 123, 123, 123]}, + "error": True, + "message": "index out of range", + }, + { + "inputs": { + "outerindex": [0, 1, 99], + "innerindex": [4294967295, 2147483648], + "outerlength": 2, + "innerlength": 2, + }, + "outputs": {"toindex": [4294967295, 2147483648, 123, 123, 123]}, + "error": False, + "message": "", + }, + { + "inputs": { + "outerindex": [0, 1], + "innerindex": [-9223372036854775808, 9223372036854775807], + "outerlength": 2, + "innerlength": 2, + }, + "outputs": { + "toindex": [ + -9223372036854775808, + 9223372036854775807, + 123, + 123, + 123, + ] + }, + "error": False, + "message": "", + }, + ] + ) + output = [] + for spec in specifications: + if spec["name"] not in FUNCTIONS: + continue + for specialization in spec["specializations"]: + for i, case in enumerate(data[spec["name"]]): + valid = True + for argument in specialization["args"]: + if "List[" in argument["type"]: + dtype = next( + v for k, v in DTYPES.items() if k in argument["type"] + ) + values = case["inputs"].get( + argument["name"], case["outputs"].get(argument["name"], []) + ) + limits = np.iinfo(dtype) + if any(x < limits.min or x > limits.max for x in values): + valid = False + if valid: + output.append( + pytest.param( + spec["name"], + specialization, + case, + id=f"{specialization['name']}-{i}", + ) + ) + return output + + +CASES = cases() + + +def arguments(specialization, case): + values = [] + outputs = [] + key = [] + for argument in specialization["args"]: + name = argument["name"] + if "List[" in argument["type"]: + dtype = next(v for k, v in DTYPES.items() if k in argument["type"]) + key.append(dtype) + if argument["dir"] == "out": + sentinel = np.iinfo(dtype).max - 2 + value = np.full(len(case["outputs"][name]) + 8, sentinel, dtype=dtype) + outputs.append(len(values)) + else: + value = np.array(case["inputs"][name], dtype=dtype) + values.append(value) + else: + values.append(case["inputs"][name]) + return values, outputs, key + + +class _DeviceArray(np.ndarray): + reads = 0 + + def get(self): + type(self).reads += 1 + return np.asarray(self).copy() + + +class _Counting: + def __init__(self, start): + self.start = start + + +class _Discard: + pass + + +def host_implementation(): + def array_factory(function): + def allocate(*args, **kwargs): + return function(*args, **kwargs).view(_DeviceArray) + + return allocate + + cp = SimpleNamespace( + **{name: array_factory(getattr(np, name)) for name in ("empty", "asarray")}, + int64=np.int64, + cuda=SimpleNamespace(get_current_stream=lambda: SimpleNamespace(ptr=123)), + ) + + def transform(*, d_in, d_out, op, num_items, stream): + assert stream.__cuda_stream__() == (0, 123) + for i in range(num_items): + value = op(np.int64(d_in.start + i)) + if not isinstance(d_out, _Discard): + d_out[i] = value + + def reduce(*, d_in, d_out, num_items, h_init, **kwargs): + assert isinstance(d_in, np.ndarray) + assert d_in.dtype == np.dtype(np.int64) + assert kwargs["stream"].__cuda_stream__() == (0, 123) + d_out[0] = min([int(h_init), *(int(d_in[i]) for i in range(num_items))]) + + tree = ast.parse((CUDA_SOURCE / "_compute.py").read_text()) + body = [ + node + for node in tree.body + if isinstance(node, ast.FunctionDef) + and node.name == "awkward_IndexedArray_simplify" + ] + namespace = { + "__name__": "awkward._connect.cuda._compute", + "np": np, + "cp": cp, + "SimpleNamespace": SimpleNamespace, + "CountingIterator": _Counting, + "DiscardIterator": _Discard, + "OpKind": SimpleNamespace(MINIMUM=1), + "unary_transform": transform, + "reduce_into": reduce, + } + exec( + compile( + ast.Module(body=body, type_ignores=[]), + str(CUDA_SOURCE / "_compute.py"), + "exec", + ), + namespace, + ) + return namespace + + +@pytest.fixture(scope="module") +def implementation(): + return host_implementation() + + +def assert_error(actual, expected): + assert (actual is None or actual.str is None) == ( + expected is None or expected.str is None + ) + if expected is not None and expected.str is not None: + assert actual.str == expected.str + assert actual.id == expected.id + assert actual.attempt == expected.attempt + + +@pytest.mark.parametrize("name,specialization,case", CASES) +def test_cpu_parity(name, specialization, case, implementation): + args, outputs, key = arguments(specialization, case) + expected_args = copy.deepcopy(args) + expected = NumpyBackend.instance()[(name, *key)](*expected_args) + assert bool(expected.str) == bool(case["error"]) + _DeviceArray.reads = 0 + actual = implementation[FUNCTIONS[name]](*args) + assert_error(actual, expected) + for i in outputs: + np.testing.assert_array_equal(args[i], expected_args[i]) + assert _DeviceArray.reads <= 1 + + +def test_raw_dispatch_removed(): + tree = ast.parse((ROOT / "dev/generate-kernel-signatures.py").read_text()) + node = next( + n + for n in tree.body + if isinstance(n, ast.Assign) + and any( + isinstance(t, ast.Name) and t.id == "cuda_kernels_impl" for t in n.targets + ) + ) + assert FUNCTIONS.keys().isdisjoint(ast.literal_eval(node.value)) + + +@pytest.fixture(scope="module") +def compute_backend(implementation): + class ComputeOnHost(NumpyBackend): + def _new_kernel(self, key): + if key[0] in FUNCTIONS: + return implementation[FUNCTIONS[key[0]]] + return super()._new_kernel(key) + + backend = object.__new__(ComputeOnHost) + backend.__init__() + return backend + + +@pytest.mark.parametrize("outer_dtype", [np.int32, np.uint32, np.int64]) +@pytest.mark.parametrize("inner_dtype", [np.int32, np.uint32, np.int64]) +@pytest.mark.parametrize("option", [False, True]) +def test_public_simplification(compute_backend, outer_dtype, inner_dtype, option): + inner_class = ( + ak.contents.IndexedOptionArray + if option and inner_dtype != np.uint32 + else ak.contents.IndexedArray + ) + inner_values = ( + [2, -1, 0] if inner_class is ak.contents.IndexedOptionArray else [2, 1, 0] + ) + inner = inner_class( + ak.index.Index(np.asarray(inner_values, dtype=inner_dtype)), + ak.contents.NumpyArray(np.asarray([10, 20, 30])), + ) + outer_values = [1, -1, 0] if option and outer_dtype != np.uint32 else [1, 0, 2] + outer = ak.index.Index(np.asarray(outer_values, dtype=outer_dtype)) + outer_class = ak.contents.IndexedOptionArray if option else ak.contents.IndexedArray + expected = outer_class.simplified(outer, inner).to_list() + inner._backend = compute_backend + assert outer_class.simplified(outer, inner).to_list() == expected + with pytest.raises(ValueError, match="index out of range"): + outer_class.simplified( + ak.index.Index(np.asarray([1, 3], dtype=outer_dtype)), inner + ) + + +def large_case(length, error_at): + rng = np.random.default_rng(2026) + outer = rng.integers(-1, 64, size=length, dtype=np.int64) + if error_at is not None: + outer[error_at] = 64 + if error_at + 1 < length: + outer[error_at + 1] = 99 + inner = rng.integers(-7, 100, size=64, dtype=np.int64) + return [np.full(length + 7, 123, dtype=np.int64), outer, length, inner, inner.size] + + +@pytest.mark.parametrize( + "length,error_at", [(33, None), (257, 0), (4097, 256), (4097, 4096)] +) +def test_first_error_across_blocks(implementation, length, error_at): + args = large_case(length, error_at) + expected_args = copy.deepcopy(args) + expected = NumpyBackend.instance()[ + ("awkward_IndexedArray_simplify", np.int64, np.int64, np.int64) + ](*expected_args) + actual = implementation["awkward_IndexedArray_simplify"](*args) + assert_error(actual, expected) + np.testing.assert_array_equal(args[0], expected_args[0]) + + +@pytest.fixture(scope="module") +def prepared_implementation(): + import weakref + from functools import lru_cache + from threading import RLock + + registry = weakref.WeakValueDictionary() + events = [] + + class Buffer(_DeviceArray): + @property + def __cuda_array_interface__(self): + registry[self.ctypes.data] = self + return { + "shape": self.shape, + "strides": None, + "typestr": self.dtype.str, + "data": (self.ctypes.data, False), + "version": 3, + } + + def prepare(**kwargs): + events.append(("prepare",)) + operation = kwargs["op"] + + def launch(**kwargs): + import time + + time.sleep(0.0001) + events.append(("transform", kwargs["num_items"])) + for i in range(kwargs["num_items"]): + value = operation(i) + if not isinstance(kwargs["d_out"], _Discard): + kwargs["d_out"][i] = value + + return launch + + def prepare_reduce(**kwargs): + events.append(("prepare_reduce",)) + + def launch(*, temp_storage, **kwargs): + if temp_storage is None: + events.append(("workspace_query",)) + return 64 + events.append(("reduce", kwargs["num_items"])) + result = int(kwargs["h_init"]) + assert isinstance(kwargs["d_in"], _Counting) + assert kwargs["d_in"].start == 0 + for row in range(kwargs["num_items"]): + result = kwargs["op"](result, np.int64(row)) + kwargs["d_out"][0] = result + return 64 + + return launch + + def empty(length, dtype): + events.append(("allocate", length, np.dtype(dtype).str)) + return np.empty(length, dtype=dtype).view(Buffer) + + namespace = { + "np": np, + "cp": SimpleNamespace( + empty=empty, + int64=np.int64, + uint8=np.uint8, + cuda=SimpleNamespace( + get_current_stream=lambda: SimpleNamespace(ptr=123), + runtime=SimpleNamespace(getDevice=lambda: 0), + driver=SimpleNamespace(ctxGetCurrent=lambda: 1), + ), + ), + "lru_cache": lru_cache, + "_lock": RLock(), + "CountingIterator": _Counting, + "DiscardIterator": _Discard, + "OpKind": SimpleNamespace(MINIMUM=1), + "make_op_adapter": lambda op: op, + "make_unary_transform": prepare, + "make_reduce_into": prepare_reduce, + } + namespace["__name__"] = "awkward._connect.cuda._indexedarray" + path = CUDA_SOURCE / "_indexedarray.py" + tree = ast.parse(path.read_text()) + body = [ + node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.ClassDef)) + ] + exec(compile(ast.Module(body=body, type_ignores=[]), str(path), "exec"), namespace) + state = namespace["_ArrayState"] + state.__getitem__ = lambda self, i: registry[self.data_ptr()][i] + state.__setitem__ = lambda self, i, value: registry[self.data_ptr()].__setitem__( + i, value + ) + return namespace, Buffer, events + + +@pytest.mark.parametrize("name,specialization,case", CASES) +def test_prepared_cpu_parity(name, specialization, case, prepared_implementation): + namespace, buffer, _ = prepared_implementation + args, outputs, key = arguments(specialization, case) + expected_args = copy.deepcopy(args) + expected = NumpyBackend.instance()[(name, *key)](*expected_args) + device = [x.copy().view(buffer) if isinstance(x, np.ndarray) else x for x in args] + if args[2] == 0: + return + row, attempt = namespace["simplify_status"](*device).get().tolist() + assert (row < args[2]) == bool(expected.str) + if expected.str: + assert (row, attempt) == (expected.id, expected.attempt) + for i in outputs: + np.testing.assert_array_equal(device[i], expected_args[i]) + + +@pytest.mark.parametrize("length", [1, 32, 64, 65, 129, 4097]) +@pytest.mark.parametrize("error_row", [None, 0, 63, 64, 128]) +def test_prepared_partial_output_and_scratch( + length, error_row, prepared_implementation +): + namespace, buffer, events = prepared_implementation + outer = (np.arange(length, dtype=np.int64) % 3).view(buffer) + inner = np.asarray([-7, 10, 20], dtype=np.int64).view(buffer) + if error_row is not None and error_row < length: + outer[error_row] = 3 + outer[-1] = 99 + output = np.full(length + 3, 123, dtype=np.int64).view(buffer) + expected = output.copy() + error = NumpyBackend.instance()[ + "awkward_IndexedArray_simplify", np.int64, np.int64, np.int64 + ](expected, outer, length, inner, len(inner)) + events.clear() + row, attempt = namespace["simplify_status"]( + output, outer, length, inner, len(inner) + ).get() + np.testing.assert_array_equal(output, expected) + if error.str: + assert (row, attempt) == (error.id, error.attempt) + else: + assert row == length + transforms = [event for event in events if event[0] == "transform"] + if length <= 64: + assert transforms == [("transform", 1)] + assert not any(event[0] == "reduce" for event in events) + else: + assert transforms == [("transform", length)] + assert ("reduce", length) in events + assert ("allocate", 2, np.dtype(np.int64).str) in events + + +def test_prepared_reuse_lifetimes_and_concurrent_calls(prepared_implementation): + import gc + import weakref + from concurrent.futures import ThreadPoolExecutor + + namespace, buffer, events = prepared_implementation + + def run(value): + outer = (np.arange(129, dtype=np.int64) % 3).view(buffer) + inner = np.asarray([value, -7, 20], dtype=np.int64).view(buffer) + output = np.full(132, 123, dtype=np.int64).view(buffer) + if value % 2: + outer[64] = 3 + outer[100] = 99 + expected = output.copy() + error = NumpyBackend.instance()[ + "awkward_IndexedArray_simplify", np.int64, np.int64, np.int64 + ](expected, outer, 129, inner, 3) + references = [weakref.ref(x) for x in (outer, inner, output)] + row, attempt = namespace["simplify_status"](output, outer, 129, inner, 3).get() + np.testing.assert_array_equal(output, expected) + if error.str: + assert (row, attempt) == (error.id, error.attempt) + else: + assert row == 129 + return references + + run(0) + events.clear() + with ThreadPoolExecutor(max_workers=4) as executor: + references = list(executor.map(run, range(12))) + gc.collect() + assert all(reference() is None for group in references for reference in group) + assert not any(event[0].startswith("prepare") for event in events) + + +@pytest.mark.parametrize("name,specialization,case", CASES) +def test_prepared_production_dispatch( + name, specialization, case, prepared_implementation, implementation, monkeypatch +): + import sys + from types import ModuleType + + namespace, buffer, _ = prepared_implementation + module = ModuleType("awkward._connect.cuda._indexedarray") + module.simplify_status = namespace["simplify_status"] + monkeypatch.setitem(sys.modules, module.__name__, module) + args, outputs, key = arguments(specialization, case) + expected = copy.deepcopy(args) + error = NumpyBackend.instance()[(name, *key)](*expected) + device = [x.view(buffer) if isinstance(x, np.ndarray) else x for x in args] + actual = implementation[name](*device) + assert_error(actual, error) + for i in outputs: + np.testing.assert_array_equal(device[i], expected[i]) + + +def test_prepared_descriptor_lengths(prepared_implementation): + namespace, buffer, _ = prepared_implementation + assert namespace["_Stream"]().__cuda_stream__() == (0, 123) + state = namespace["_ArrayState"](np.int64) + assert len(state) == 0 + array = np.arange(7, dtype=np.int64).view(buffer) + state.bind(array) + assert len(state) == 7 + assert state.data_ptr() == array.ctypes.data + state.clear() + assert len(state) == 0 + assert state.data_ptr() == 0 + + +@pytest.mark.parametrize("length", [32, 129, 4096]) +def test_compact_scratch_and_workspace_sizing_reuse(length, prepared_implementation): + namespace, buffer, events = prepared_implementation + namespace["_setup"].cache_clear() + outer = (np.arange(length, dtype=np.int64) % 3).view(buffer) + inner = np.array([-7, 10, 20], dtype=np.int64).view(buffer) + output = np.full(length + 3, 123, dtype=np.int64).view(buffer) + for iteration in range(2): + events.clear() + status = namespace["simplify_status"](output, outer, length, inner, len(inner)) + assert status.shape == (2,) + assert status.get().tolist() == [length, np.iinfo(np.int64).max] + allocations = [event for event in events if event[0] == "allocate"] + assert allocations == ( + [("allocate", 2, np.dtype(np.int64).str)] + if length <= 64 + else [ + ("allocate", 2, np.dtype(np.int64).str), + ("allocate", 64, np.dtype(np.uint8).str), + ] + ) + queries = [event for event in events if event[0] == "workspace_query"] + assert len(queries) == (1 if length > 64 and iteration == 0 else 0) + np.testing.assert_array_equal(output[:length], inner[outer]) + np.testing.assert_array_equal(output[length:], [123, 123, 123]) + + +@pytest.mark.parametrize("dtype", [np.int32, np.uint32, np.int64]) +def test_actual_fused_operator_is_associative(dtype, prepared_implementation): + import itertools + + namespace, buffer, _ = prepared_implementation + values = np.array([0, 1, 3, 0, 99, 2], dtype=dtype).view(buffer) + if dtype is not np.uint32: + values[1] = -1 + setup = namespace["_Simplify"](np.int64, dtype, np.int64, len(values), 3) + setup.outer.bind(values) + try: + domain = range(len(values) + 1) + op = setup.combine + for a, b, c in itertools.product(domain, repeat=3): + assert op(op(a, b), c) == op(a, op(b, c)) + assert op(a, b) == op(b, a) + finally: + setup.outer.clear()