From 3decf2d554be2e85130fe71c388a18daefa0d2fb Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 22 Oct 2026 16:47:41 -0500 Subject: [PATCH 1/8] Fix device memory leaks --- python/cuda_cccl/cuda/compute/_bindings.pyi | 6 +- .../cuda_cccl/cuda/compute/_bindings_impl.pyx | 34 +++--- python/cuda_cccl/cuda/compute/_caching.py | 15 ++- .../cuda_cccl/cuda/compute/_cccl_interop.py | 13 +- python/cuda_cccl/cuda/compute/_jit.py | 47 +++++--- .../cuda_cccl/tests/compute/test_caching.py | 114 ++++++++++++++++++ .../compute/test_free_threading_stress.py | 5 +- .../tests/compute/test_jit_backend.py | 14 ++- .../compute/test_multi_cc_serialization.py | 5 +- 9 files changed, 192 insertions(+), 61 deletions(-) create mode 100644 python/cuda_cccl/tests/compute/test_caching.py diff --git a/python/cuda_cccl/cuda/compute/_bindings.pyi b/python/cuda_cccl/cuda/compute/_bindings.pyi index 3b407ba44cd7..8545cd71d301 100644 --- a/python/cuda_cccl/cuda/compute/_bindings.pyi +++ b/python/cuda_cccl/cuda/compute/_bindings.pyi @@ -129,9 +129,9 @@ class Value: class Pointer: def __init__(self, arg): ... - def rebind(self, ptr: int | ctypes.c_void_p, owner: Any) -> None: ... + def rebind(self, ptr: int | ctypes.c_void_p) -> None: ... -def make_pointer_object(ptr: int | ctypes.c_void_p, owner: Any) -> Pointer: ... +def make_pointer_object(ptr: int | ctypes.c_void_p) -> Pointer: ... class IteratorState(Buffer): def __init__(self, arg): ... @@ -170,7 +170,7 @@ class Iterator: def as_bytes(self) -> bytes: ... @property def is_ptr_kind(self) -> bool: ... - def bind_pointer_state(self, ptr: int | ctypes.c_void_p, owner: Any) -> None: ... + def bind_pointer_state(self, ptr: int | ctypes.c_void_p) -> None: ... @property def host_advance_fn(self): ... @host_advance_fn.setter diff --git a/python/cuda_cccl/cuda/compute/_bindings_impl.pyx b/python/cuda_cccl/cuda/compute/_bindings_impl.pyx index 67f593ebb028..8ac25aec70a5 100644 --- a/python/cuda_cccl/cuda/compute/_bindings_impl.pyx +++ b/python/cuda_cccl/cuda/compute/_bindings_impl.pyx @@ -609,33 +609,31 @@ cdef class Pointer(StateBase): ) self.set_state(ptr, ref) - def rebind(self, ptr, owner): - """Update the ptr and ref in place""" + def rebind(self, ptr): + """Update the ptr in place from a raw address. + + Clears ``ref``: a raw address has no Python object to keep alive, and + recording one here (e.g. the device array the address came from) would + keep it alive for as long as this Pointer does, which for the Pointer + cached on an Iterator is for as long as the algorithm's wrapper cache + entry lives. + """ if isinstance(ptr, int): self.ptr = int_as_ptr(ptr) elif isinstance(ptr, ctypes.c_void_p): self.ptr = int_as_ptr(ptr.value) else: raise TypeError( - "First argument must be an integer, or ctypes.c_void_p, " + "Argument must be an integer, or ctypes.c_void_p, " f"got {type(ptr)}" ) - self.ref = owner + self.ref = None -def make_pointer_object(ptr, owner): +def make_pointer_object(ptr): + """Make a Pointer from a raw address.""" cdef Pointer res = Pointer(0) - - if isinstance(ptr, int): - res.ptr = int_as_ptr(ptr) - elif isinstance(ptr, ctypes.c_void_p): - res.ptr = int_as_ptr(ptr.value) - else: - raise TypeError( - "First argument must be an integer, or ctypes.c_void_p, " - f"got {type(ptr)}" - ) - res.ref = owner + res.rebind(ptr) return res @@ -840,7 +838,7 @@ cdef class Iterator: self.is_ptr_kind = (it_kind == cccl_iterator_kind_t.POINTER) self._cached_ptr_obj = None - def bind_pointer_state(self, ptr, owner): + def bind_pointer_state(self, ptr): """Set state from a raw pointer, reusing a cached Pointer instead of allocating one each call. Only valid when is_ptr_kind is True. """ @@ -848,7 +846,7 @@ cdef class Iterator: if cached is None: cached = Pointer(0) self._cached_ptr_obj = cached - cached.rebind(ptr, owner) + cached.rebind(ptr) self.state = cached @property diff --git a/python/cuda_cccl/cuda/compute/_caching.py b/python/cuda_cccl/cuda/compute/_caching.py index 25b869c1fb22..07ef7fa15f0f 100644 --- a/python/cuda_cccl/cuda/compute/_caching.py +++ b/python/cuda_cccl/cuda/compute/_caching.py @@ -662,7 +662,7 @@ def clear_all_caches(): _clear_wrapper_caches() _process_wide_build_results_cache.clear() # Auxiliary caches registered process-wide (_jit._infer_return_type, the - # compiled-operator memo _jit._compile_op_impl, and the NVRTC memos in + # compiled-operator memo _jit.compile_op, and the NVRTC memos in # _cpp_compile) must be cleared too, so builds after a clear really are # cold. Factory entries' cache_clear is idempotent with # _clear_wrapper_caches above. Snapshot the registry first: a lazy import @@ -674,19 +674,18 @@ def clear_all_caches(): class CachableFunction: """ - A type that wraps a function and provides custom comparison - (__eq__) and hash (__hash__) implementations. + A cache key for a function, with comparison (__eq__) and hash (__hash__) + based on its bytecode, constants, and closures, while ignoring other + attributes such as its docstring. - The purpose of this class is to enable caching and comparison of - functions based on their bytecode, constants, and closures, while - ignoring other attributes such as their names or docstrings. + It does not hold a reference to the function it was built from. """ # TODO: eventually, move this class to _jit.py as it only # has to do with caching of Python callables that will be # JIT compiled. def __init__(self, func): - self._func = func + self._name = func.__name__ closure = func.__closure__ if func.__closure__ is not None else [] contents = [] @@ -718,7 +717,7 @@ def __hash__(self): return hash(self._identity) def __repr__(self): - return str(self._func) + return f"" # Register keyers for built-in types diff --git a/python/cuda_cccl/cuda/compute/_cccl_interop.py b/python/cuda_cccl/cuda/compute/_cccl_interop.py index 9011948552b2..2346290842e5 100644 --- a/python/cuda_cccl/cuda/compute/_cccl_interop.py +++ b/python/cuda_cccl/cuda/compute/_cccl_interop.py @@ -228,15 +228,22 @@ def get_value_type( def set_cccl_iterator_state(cccl_it: Iterator, input_it): + """Bind the per-call state of ``input_it`` to the cached ``cccl_it``. + + Only the pointer value (or state bytes) is bound; ``input_it`` itself is + not kept alive by ``cccl_it``. The caller holds ``input_it`` for the + duration of the call, which is all the launch needs, whereas ``cccl_it`` + belongs to an algorithm object that lives in the process-wide wrapper + cache. + """ if cccl_it.is_ptr_kind: - ptr = get_data_pointer(input_it) - cccl_it.bind_pointer_state(ptr, input_it) + cccl_it.bind_pointer_state(get_data_pointer(input_it)) else: state_ = input_it.state if isinstance(state_, (IteratorState, Pointer)): cccl_it.state = state_ else: - cccl_it.state = make_pointer_object(state_, input_it) + cccl_it.state = make_pointer_object(state_) @functools.lru_cache() diff --git a/python/cuda_cccl/cuda/compute/_jit.py b/python/cuda_cccl/cuda/compute/_jit.py index 1c2205110de0..988d18149edc 100644 --- a/python/cuda_cccl/cuda/compute/_jit.py +++ b/python/cuda_cccl/cuda/compute/_jit.py @@ -24,6 +24,7 @@ from ._bindings import Op, OpKind, TypeEnum from ._caching import ( CachableFunction, + _cache_single_flight, _make_cache_key_from_args, _process_wide_cache_registry, cache_with_registered_key_functions, @@ -870,25 +871,20 @@ def _compile_wrapper_to_device_code(wrapped_op, wrapper_sig, cc): # ----------------------------------------------------------------------------- -@functools.lru_cache(maxsize=256) -def _compile_op_impl(cachable_op, input_types_tuple: tuple, output_type, cc=None): - """Cached implementation of op compilation. +def _compile_op_impl(op, input_types_tuple: tuple, output_type, cc=None): + """Compile an operator to device code. Uncached; see compile_op. Args: - cachable_op: CachableFunction wrapper around the operator + op: The Python operator input_types_tuple: Tuple of input TypeDescriptors output_type: Output TypeDescriptor cc: Target compute capability ``(major, minor)`` for the device code, or None to target the current device (which requires a GPU to be - present). Part of the cache key so the same operator compiled for - different arches does not collide. + present). """ from ._bindings import Op, OpKind from ._odr_helpers import create_op_void_ptr_wrapper - # Extract the actual function from CachableFunction - op = cachable_op._func - # Ensure any gpu_struct classes referenced in the op are registered _ensure_function_structs_registered(op) @@ -916,11 +912,14 @@ def _compile_op_impl(cachable_op, input_types_tuple: tuple, output_type, cc=None ) -# Keep clear_all_caches() covering compiled operators: this is the memo a build -# must miss after a clear for the JIT to actually run again. (The struct -# registration caches above stay unregistered: clearing them would only make the -# next build re-register the same types, overhead with no benefit.) -_process_wide_cache_registry["_jit._compile_op_impl"] = _compile_op_impl +# Memo of compiled operators, keyed by (operator cache key, input types, +# output type, target cc). The operator's key is a CachableFunction, which holds +# no reference to the operator, so the memo never keeps an operator's closure, +# and with it any device array a stateful operator captured, alive. The +# operator itself only travels into _compile_op_impl on a miss, from the +# compile_op caller that holds it. Concurrent misses for one key are +# coalesced by _cache_single_flight. +_compile_op_cache: dict = {} def compile_op(op, input_types, output_type=None): @@ -930,14 +929,24 @@ def compile_op(op, input_types, output_type=None): types produce the same compiled result (same symbol names and LTOIR), which allows proper deduplication during linking. """ - from ._caching import CachableFunction - - cachable_op = CachableFunction(op) - return _compile_op_impl( - cachable_op, tuple(input_types), output_type, get_target_cc() + input_types_tuple = tuple(input_types) + cc = get_target_cc() + key = (CachableFunction(op), input_types_tuple, output_type, cc) + return _cache_single_flight( + _compile_op_cache, + key, + lambda: _compile_op_impl(op, input_types_tuple, output_type, cc), ) +compile_op.cache_clear = _compile_op_cache.clear # type: ignore[attr-defined] +# Keep clear_all_caches() covering compiled operators: this is the memo a build +# must miss after a clear for the JIT to actually run again. (The struct +# registration caches above stay unregistered: clearing them would only make the +# next build re-register the same types, overhead with no benefit.) +_process_wide_cache_registry["_jit.compile_op"] = compile_op + + class _StatelessOp(OpAdapter): """Adapter for stateless callables.""" diff --git a/python/cuda_cccl/tests/compute/test_caching.py b/python/cuda_cccl/tests/compute/test_caching.py new file mode 100644 index 000000000000..c1d61dbe28b4 --- /dev/null +++ b/python/cuda_cccl/tests/compute/test_caching.py @@ -0,0 +1,114 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +"""Behavioral tests for the cached algorithm wrappers.""" + +import gc +import weakref + +import numpy as np +import pytest +from _utils.device_array import DeviceArray +from cuda.core import Device + +import cuda.compute +from cuda.compute import OpKind, TransformIterator + +NUM_ITEMS = 1024 + + +def _double(x): + return x * 2 + + +def _run_reduce_into(): + d_in = DeviceArray.from_numpy(np.arange(NUM_ITEMS, dtype=np.int32)) + d_out = DeviceArray.empty(1, np.int32) + h_init = np.zeros(1, dtype=np.int32) + cuda.compute.reduce_into( + d_in=d_in, d_out=d_out, num_items=NUM_ITEMS, op=OpKind.PLUS, h_init=h_init + ) + return [d_in, d_out] + + +def _run_reduce_into_transform_iterator(): + d_in = DeviceArray.from_numpy(np.arange(NUM_ITEMS, dtype=np.int32)) + d_out = DeviceArray.empty(1, np.int32) + h_init = np.zeros(1, dtype=np.int32) + cuda.compute.reduce_into( + d_in=TransformIterator(d_in, _double), + d_out=d_out, + num_items=NUM_ITEMS, + op=OpKind.PLUS, + h_init=h_init, + ) + return [d_in, d_out] + + +def _run_unary_transform_with_captured_array(): + # ``scale`` is captured by the operator's closure, which makes this a + # stateful operator. The operator is keyed in the caches through + # CachableFunction, which must not keep the function, and with it the + # captured array, alive. + scale = DeviceArray.from_numpy(np.array([3], dtype=np.int32)) + + def scale_op(x): + return x * scale[0] + + d_in = DeviceArray.from_numpy(np.arange(NUM_ITEMS, dtype=np.int32)) + d_out = DeviceArray.empty(NUM_ITEMS, np.int32) + cuda.compute.unary_transform( + d_in=d_in, d_out=d_out, op=scale_op, num_items=NUM_ITEMS + ) + return [d_in, d_out, scale] + + +def _run_unique_by_key(): + d_in_keys = DeviceArray.from_numpy(np.arange(NUM_ITEMS, dtype=np.int32) // 4) + d_in_items = DeviceArray.from_numpy(np.arange(NUM_ITEMS, dtype=np.float32)) + d_out_keys = DeviceArray.empty(NUM_ITEMS, np.int32) + d_out_items = DeviceArray.empty(NUM_ITEMS, np.float32) + d_out_num_selected = DeviceArray.empty(1, np.int32) + cuda.compute.unique_by_key( + d_in_keys=d_in_keys, + d_in_items=d_in_items, + d_out_keys=d_out_keys, + d_out_items=d_out_items, + d_out_num_selected=d_out_num_selected, + op=OpKind.EQUAL_TO, + num_items=NUM_ITEMS, + ) + return [d_in_keys, d_in_items, d_out_keys, d_out_items, d_out_num_selected] + + +@pytest.mark.parametrize( + "run_algorithm", + [ + _run_reduce_into, + _run_reduce_into_transform_iterator, + _run_unary_transform_with_captured_array, + _run_unique_by_key, + ], + ids=[ + "reduce_into", + "reduce_into_transform_iterator", + "unary_transform_with_captured_array", + "unique_by_key", + ], +) +def test_cached_wrapper_does_not_retain_caller_arrays(run_algorithm): + # The single-phase API caches the algorithm object it builds, keyed by + # argument types, for the rest of the process. That object must not keep + # the arrays of any particular call alive: a user iterating over many type + # combinations would otherwise accumulate device memory that cannot be + # released (NVIDIA/cccl#12187). + arrays = run_algorithm() + Device().sync() + array_refs = [weakref.ref(array) for array in arrays] + + del arrays + gc.collect() + + still_alive = [ref() for ref in array_refs if ref() is not None] + assert still_alive == [] diff --git a/python/cuda_cccl/tests/compute/test_free_threading_stress.py b/python/cuda_cccl/tests/compute/test_free_threading_stress.py index 5a597f6a9bf5..b32082d6ecef 100644 --- a/python/cuda_cccl/tests/compute/test_free_threading_stress.py +++ b/python/cuda_cccl/tests/compute/test_free_threading_stress.py @@ -1880,8 +1880,9 @@ def test_distinct_op_cold_build_storm(jit_compute_module): _cache_single_flight elects a separate builder per thread and the backend compiles run in parallel rather than coalescing. This is the test that exercises whether numba-cuda-mlir's compile pipeline -- MLIR contexts, - lowering registries, the NVVM path -- tolerates concurrent entry, and - whether _compile_op_impl's non-coalescing lru_cache is safe under it. + lowering registries, the NVVM path -- tolerates concurrent entry: the + compile_op memo coalesces only same-key misses, so distinct keys compile + concurrently. """ cc = jit_compute_module diff --git a/python/cuda_cccl/tests/compute/test_jit_backend.py b/python/cuda_cccl/tests/compute/test_jit_backend.py index e5e7b4b055b3..eac65f6fce96 100644 --- a/python/cuda_cccl/tests/compute/test_jit_backend.py +++ b/python/cuda_cccl/tests/compute/test_jit_backend.py @@ -293,10 +293,9 @@ def touch(a): # Only now import cuda.compute and compile an operator through it. from cuda.compute import _mlir from cuda.compute._jit import _compile_op_impl - from cuda.compute._caching import CachableFunction int32 = _mlir.types.int32 - _compile_op_impl(CachableFunction(lambda x: x + 1), (int32,), int32, (8, 9)) + _compile_op_impl(lambda x: x + 1, (int32,), int32, (8, 9)) """ ) result = subprocess.run( @@ -360,7 +359,7 @@ def test_clear_all_caches_drops_compiled_device_code(): import cuda.compute from cuda.compute import CountingIterator, OpKind, TransformIterator from cuda.compute._cpp_compile import compile_cpp_op_code, compile_cpp_to_ltoir - from cuda.compute._jit import _compile_op_impl + from cuda.compute._jit import _compile_op_cache from cuda.compute.algorithms._select import _always_false_op def add_one(x): @@ -393,14 +392,18 @@ def run_select(): ) assert int(d_num.copy_to_host()[0]) == 4 - memos = (_compile_op_impl, compile_cpp_op_code, _always_false_op) + # The compiled-operator memo is a plain dict (every miss adds an entry); + # the other two are lru_caches. + memos = (compile_cpp_op_code, _always_false_op) first = make_iter() run_reduce(first) run_select() + assert len(_compile_op_cache) > 0 assert all(m.cache_info().currsize > 0 for m in memos) cuda.compute.clear_all_caches() + assert len(_compile_op_cache) == 0 assert all(m.cache_info().currsize == 0 for m in memos) assert compile_cpp_to_ltoir.cache_info().currsize == 0 @@ -408,10 +411,11 @@ def run_select(): # is cold exactly when every memo misses again. run_reduce(make_iter()) run_select() + assert len(_compile_op_cache) > 0 assert all(m.cache_info().misses > 0 for m in memos) # Reusing the live iterator relinks the ops memoized on it: no recompile. cuda.compute.clear_all_caches() run_reduce(first) - assert _compile_op_impl.cache_info().misses == 0 + assert len(_compile_op_cache) == 0 assert compile_cpp_op_code.cache_info().misses == 0 diff --git a/python/cuda_cccl/tests/compute/test_multi_cc_serialization.py b/python/cuda_cccl/tests/compute/test_multi_cc_serialization.py index 2d84c707621c..253b35c20bec 100644 --- a/python/cuda_cccl/tests/compute/test_multi_cc_serialization.py +++ b/python/cuda_cccl/tests/compute/test_multi_cc_serialization.py @@ -1111,7 +1111,6 @@ def test_operator_device_code_targets_the_requested_cc(monkeypatch): follow the build machine instead. """ from cuda.compute import _jit - from cuda.compute._caching import CachableFunction if not _jit.USING_V2: pytest.skip("v1 compiles the operator through a different entry point") @@ -1131,8 +1130,8 @@ def add_one(x): return x + 1 int32 = from_numpy_dtype(np.dtype(np.int32)) - _jit._compile_op_impl.cache_clear() - _jit._compile_op_impl(CachableFunction(add_one), (int32,), int32, (7, 5)) + _jit.compile_op.cache_clear() + _jit._compile_op_impl(add_one, (int32,), int32, (7, 5)) assert recorded["chip"] == "sm_75" From cec1cece87bd485140f4b1202f77ac101389e320 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 22 Oct 2026 19:02:00 -0500 Subject: [PATCH 2/8] Further fixes --- python/cuda_cccl/cuda/compute/_caching.py | 8 +++- .../cuda/compute/iterators/_transform.py | 7 +++- .../cuda_cccl/tests/compute/test_caching.py | 37 ++++++++++++++++++- .../cuda_cccl/tests/compute/test_iterators.py | 14 +++++++ 4 files changed, 63 insertions(+), 3 deletions(-) diff --git a/python/cuda_cccl/cuda/compute/_caching.py b/python/cuda_cccl/cuda/compute/_caching.py index 07ef7fa15f0f..bfdcf753358c 100644 --- a/python/cuda_cccl/cuda/compute/_caching.py +++ b/python/cuda_cccl/cuda/compute/_caching.py @@ -6,6 +6,7 @@ from __future__ import annotations import functools +import hashlib import threading import types import weakref @@ -717,7 +718,12 @@ def __hash__(self): return hash(self._identity) def __repr__(self): - return f"" + # Iterators derive device symbol names from str() of their cache key + # (see IteratorBase._get_uid), so the repr must distinguish functions + # that compare unequal: two same-named operators with different bodies + # must not share a symbol. Hash the identity rather than printing it. + digest = hashlib.sha256(repr(self._identity).encode()).hexdigest()[:16] + return f"" # Register keyers for built-in types diff --git a/python/cuda_cccl/cuda/compute/iterators/_transform.py b/python/cuda_cccl/cuda/compute/iterators/_transform.py index 9dcf5eec3d9a..995a4cd71663 100644 --- a/python/cuda_cccl/cuda/compute/iterators/_transform.py +++ b/python/cuda_cccl/cuda/compute/iterators/_transform.py @@ -9,6 +9,7 @@ from textwrap import dedent from .._bindings import Op, OpKind +from .._caching import _key_for from .._cpp_compile import compile_cpp_op_code, make_variable_declaration from ..op import make_op_adapter from ..types import TypeDescriptor, signature_from_annotations @@ -302,10 +303,14 @@ def kind(self): if isinstance(self._value_type, list) else self._value_type ) + # Key the op adapter the way algorithm arguments are keyed, rather than + # placing the adapter itself in the key: a JIT adapter holds the Python + # function (and, for a stateful op, its captured device arrays), and + # this kind ends up in process-wide cache keys. return ( "TransformIterator", self._is_input, - self._transform_op, + _key_for(self._transform_op), self._underlying.kind, value_type, ) diff --git a/python/cuda_cccl/tests/compute/test_caching.py b/python/cuda_cccl/tests/compute/test_caching.py index c1d61dbe28b4..1fefff683ad4 100644 --- a/python/cuda_cccl/tests/compute/test_caching.py +++ b/python/cuda_cccl/tests/compute/test_caching.py @@ -13,11 +13,19 @@ from cuda.core import Device import cuda.compute -from cuda.compute import OpKind, TransformIterator +from cuda.compute import OpKind, TransformIterator, types NUM_ITEMS = 1024 +def _require_jit_backend(): + """Skip when the JIT backend is absent, as on the minimal extras.""" + pytest.importorskip( + "numba_cuda_mlir", + reason="Python-callable operators require numba-cuda-mlir", + ) + + def _double(x): return x * 2 @@ -33,6 +41,7 @@ def _run_reduce_into(): def _run_reduce_into_transform_iterator(): + _require_jit_backend() d_in = DeviceArray.from_numpy(np.arange(NUM_ITEMS, dtype=np.int32)) d_out = DeviceArray.empty(1, np.int32) h_init = np.zeros(1, dtype=np.int32) @@ -46,7 +55,31 @@ def _run_reduce_into_transform_iterator(): return [d_in, d_out] +def _run_reduce_into_transform_iterator_with_captured_array(): + _require_jit_backend() + # A stateful transform op: the TransformIterator's cache key must not + # embed the op adapter, which holds the function and its captured array. + # Stateful ops cannot infer their return type, so it is given explicitly. + scale = DeviceArray.from_numpy(np.array([3], dtype=np.int32)) + + def scale_op(x): + return x * scale[0] + + d_in = DeviceArray.from_numpy(np.arange(NUM_ITEMS, dtype=np.int32)) + d_out = DeviceArray.empty(1, np.int32) + h_init = np.zeros(1, dtype=np.int32) + cuda.compute.reduce_into( + d_in=TransformIterator(d_in, scale_op, value_type=types.int32), + d_out=d_out, + num_items=NUM_ITEMS, + op=OpKind.PLUS, + h_init=h_init, + ) + return [d_in, d_out, scale] + + def _run_unary_transform_with_captured_array(): + _require_jit_backend() # ``scale`` is captured by the operator's closure, which makes this a # stateful operator. The operator is keyed in the caches through # CachableFunction, which must not keep the function, and with it the @@ -87,12 +120,14 @@ def _run_unique_by_key(): [ _run_reduce_into, _run_reduce_into_transform_iterator, + _run_reduce_into_transform_iterator_with_captured_array, _run_unary_transform_with_captured_array, _run_unique_by_key, ], ids=[ "reduce_into", "reduce_into_transform_iterator", + "reduce_into_transform_iterator_with_captured_array", "unary_transform_with_captured_array", "unique_by_key", ], diff --git a/python/cuda_cccl/tests/compute/test_iterators.py b/python/cuda_cccl/tests/compute/test_iterators.py index 85f85a05812e..eec96d735f37 100644 --- a/python/cuda_cccl/tests/compute/test_iterators.py +++ b/python/cuda_cccl/tests/compute/test_iterators.py @@ -92,6 +92,20 @@ def op3(x): assert it4.kind != it7.kind +def test_transform_iterator_uid_distinguishes_same_named_ops(): + # Device symbol names are derived from str(kind). Two ops with the same + # __name__ but different bodies must not share a symbol, or linking the + # two iterators into one algorithm would collide; equal ops must share + # one so identical iterators deduplicate. + ary = DeviceArray.from_numpy(np.asarray([0, 1, 2], dtype="int32")) + it1 = TransformIterator(ary, lambda x: x + 1) + it2 = TransformIterator(ary, lambda x: x + 1) + it3 = TransformIterator(ary, lambda x: x * 2) + + assert it1._get_uid() == it2._get_uid() + assert it1._get_uid() != it3._get_uid() + + def test_reverse_input_iterator_equality(): ary1 = DeviceArray.from_numpy(np.asarray([0, 1, 2], dtype="int32")) ary2 = DeviceArray.from_numpy(np.asarray([3, 4, 5], dtype="int32")) From 65a33698f5a63dae71c14e01123cb8b97733e747 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 22 Oct 2026 19:05:36 -0500 Subject: [PATCH 3/8] Update docs --- docs/python/compute/index.rst | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/python/compute/index.rst b/docs/python/compute/index.rst index 6507f779faee..f07a494ccee7 100644 --- a/docs/python/compute/index.rst +++ b/docs/python/compute/index.rst @@ -52,7 +52,9 @@ API conventions * **Output semantics** — Algorithms write results into a user-provided array or iterator rather than returning them. This keeps memory ownership explicit and lifetimes under - your control. + your control. As with any asynchronous CUDA launch, ``cuda.compute`` does not keep + the arrays you pass alive: keep your own references until the work submitted on the + stream has completed. * **Operators** — Many algorithms accept an ``op`` parameter. This can be a built-in :class:`OpKind ` value or a From 9be4474935ae07bf487d2b6fed8a43aef07bce24 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Fri, 23 Oct 2026 16:43:02 -0500 Subject: [PATCH 4/8] Use proxy arrays instead --- python/cuda_cccl/cuda/compute/_caching.py | 31 +++++++++++++++++-- .../cuda_cccl/cuda/compute/_cccl_interop.py | 14 ++++++--- python/cuda_cccl/cuda/compute/_proxy.py | 8 +++++ .../compute/test_multi_cc_serialization.py | 26 ++++++++++++++++ 4 files changed, 72 insertions(+), 7 deletions(-) diff --git a/python/cuda_cccl/cuda/compute/_caching.py b/python/cuda_cccl/cuda/compute/_caching.py index bfdcf753358c..8c5c670c7293 100644 --- a/python/cuda_cccl/cuda/compute/_caching.py +++ b/python/cuda_cccl/cuda/compute/_caching.py @@ -20,7 +20,7 @@ except ImportError: from cuda.core.experimental._utils.cuda_utils import CUDAError -from ._utils.protocols import get_dtype, get_shape, is_device_array +from ._utils.protocols import get_dtype, get_shape, is_contiguous, is_device_array from .struct import _Struct try: @@ -500,6 +500,32 @@ def build_privately(): ) +def _proxy_device_arrays(args: tuple, kwargs: dict) -> tuple[tuple, dict]: + """Replace device-array arguments with dtype-only ``ProxyArray`` placeholders. + + A factory builds an algorithm from its arguments' *types*; the arrays' data + is bound at ``__call__``. Building from proxies therefore produces the same + algorithm, but the wrapper, which lives in the process-wide wrapper cache, + is never constructed from (and so can never retain) a caller's allocation, + and the single-phase and ahead-of-time (``ProxyArray``) build paths are + identical. Non-contiguous arrays are left as is so the build keeps raising + the same error it does today. Objects that merely contain arrays (iterators, + ``DoubleBuffer``) are not device arrays and pass through unchanged. + """ + from ._proxy import ProxyArray + + def proxy(value): + if ( + is_device_array(value) + and not isinstance(value, ProxyArray) + and is_contiguous(value) + ): + return ProxyArray(get_dtype(value)) + return value + + return tuple(proxy(a) for a in args), {k: proxy(v) for k, v in kwargs.items()} + + class _CacheWithRegisteredKeyFunctions: """ Decorator to cache the result of the decorated function. @@ -574,8 +600,9 @@ def inner(*args, **kwargs): ) _thread_local.factory_device_info = target try: + build_args, build_kwargs = _proxy_device_arrays(args, kwargs) with target_cc(target_cc_arg): - result = func(*args, **kwargs) + result = func(*build_args, **build_kwargs) finally: _thread_local.factory_device_info = previous_device_info cache[cache_key] = result diff --git a/python/cuda_cccl/cuda/compute/_cccl_interop.py b/python/cuda_cccl/cuda/compute/_cccl_interop.py index 2346290842e5..751cfd1484a1 100644 --- a/python/cuda_cccl/cuda/compute/_cccl_interop.py +++ b/python/cuda_cccl/cuda/compute/_cccl_interop.py @@ -97,7 +97,7 @@ def _is_well_known_op(op: OpKind) -> bool: def _device_array_to_cccl_iter(array: DeviceArrayLike) -> Iterator: - from ._proxy import ProxyArray + from ._proxy import PROXY_PLACEHOLDER_ADDRESS, ProxyArray if not is_contiguous(array): raise ValueError("Non-contiguous arrays are not supported.") @@ -105,10 +105,14 @@ def _device_array_to_cccl_iter(array: DeviceArrayLike) -> Iterator: info = _type_info_from_dtype(dtype) state_info = _type_info_from_dtype(np.intp) - # A ProxyArray has no GPU allocation: leave the pointer NULL for build-time - # (ahead-of-time) compilation. The real pointer is bound at __call__ via - # set_cccl_iterator_state(). - state = None if isinstance(array, ProxyArray) else get_data_pointer(array) + # A ProxyArray has no GPU allocation: bind a placeholder address for the + # build (see PROXY_PLACEHOLDER_ADDRESS for why it is not NULL). The real + # pointer is bound at __call__ via set_cccl_iterator_state(). + state = ( + PROXY_PLACEHOLDER_ADDRESS + if isinstance(array, ProxyArray) + else get_data_pointer(array) + ) return Iterator( state_info.alignment, IteratorKind.POINTER, diff --git a/python/cuda_cccl/cuda/compute/_proxy.py b/python/cuda_cccl/cuda/compute/_proxy.py index 4f498dde7cc8..ad4898272906 100644 --- a/python/cuda_cccl/cuda/compute/_proxy.py +++ b/python/cuda_cccl/cuda/compute/_proxy.py @@ -23,6 +23,14 @@ "Pass a real device array when calling the compiled algorithm." ) +# Build-time placeholder address for a ProxyArray's CCCL iterator. The C library +# reads a pointer iterator's NULL state as "this optional array was not passed" +# (keys-only sorts), so a proxy must present a non-NULL address to be compiled +# as a present array. The value is never dereferenced: a build does not launch, +# and a proxy cannot be executed (get_data_pointer raises), so the real address +# is always bound by __call__ before any launch. +PROXY_PLACEHOLDER_ADDRESS = 0x1000 + _PROXY_VALUE_DATA_ERROR = ( "ProxyValue has no data — it is a build-time placeholder only. " "Pass a real scalar or numpy array when calling the compiled algorithm." diff --git a/python/cuda_cccl/tests/compute/test_multi_cc_serialization.py b/python/cuda_cccl/tests/compute/test_multi_cc_serialization.py index 253b35c20bec..1605c2d09dba 100644 --- a/python/cuda_cccl/tests/compute/test_multi_cc_serialization.py +++ b/python/cuda_cccl/tests/compute/test_multi_cc_serialization.py @@ -31,8 +31,10 @@ import cuda.compute._cccl_interop as cccl_interop from cuda.compute import ( + OpKind, ProxyArray, deserialize, + make_merge_sort, make_select, make_unary_transform, serialize, @@ -48,6 +50,7 @@ key_to_cc, normalize_compute_capabilities, ) +from cuda.compute._proxy import PROXY_PLACEHOLDER_ADDRESS from cuda.compute._target_cc import target_cc from cuda.compute.iterators import CountingIterator, TransformIterator from cuda.compute.types import from_numpy_dtype @@ -462,6 +465,29 @@ def _compilable_ccs(want=2): # ---------------------------------------------------------------------------- +@requires_serialization +@pytest.mark.parametrize("with_values", [False, True], ids=["keys_only", "pairs"]) +def test_proxy_compile_sorts_with_optional_values(with_values): + # The C library tells a keys-only sort from a key-value sort by whether the + # values iterator's pointer state is NULL. A ProxyArray therefore binds a + # non-NULL placeholder address, so a proxy values array compiles as present + # and a proxy keys array is never mistaken for an absent one. + ccs = _compilable_ccs(1) + if not ccs: + pytest.skip("toolchain compiles for no target arch") + values = ProxyArray(np.float32) if with_values else None + sorter = make_merge_sort( + d_in_keys=ProxyArray(np.int32), + d_in_values=values, + d_out_keys=ProxyArray(np.int32), + d_out_values=values, + op=OpKind.LESS, + compute_capability=ccs[0], + ) + assert sorter.d_in_keys_cccl.state == PROXY_PLACEHOLDER_ADDRESS + assert (sorter.d_in_values_cccl.state == PROXY_PLACEHOLDER_ADDRESS) is with_values + + @requires_serialization def test_proxy_single_cc_compile_is_not_loaded(): ccs = _compilable_ccs(1) From 508da50bd2e12a5219c9267e35814743c3c2ed75 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Fri, 30 Oct 2026 17:27:44 -0500 Subject: [PATCH 5/8] Address review --- python/cuda_cccl/cuda/compute/_caching.py | 41 +++++++++--- .../cuda_cccl/tests/compute/test_caching.py | 62 +++++++++++++++++++ 2 files changed, 95 insertions(+), 8 deletions(-) diff --git a/python/cuda_cccl/cuda/compute/_caching.py b/python/cuda_cccl/cuda/compute/_caching.py index 8c5c670c7293..22e319effcdd 100644 --- a/python/cuda_cccl/cuda/compute/_caching.py +++ b/python/cuda_cccl/cuda/compute/_caching.py @@ -632,19 +632,35 @@ def register(self, type_: type, key_function: Callable[[Any], Hashable]) -> None cache_with_registered_key_functions = _CacheWithRegisteredKeyFunctions() -def _make_hashable(value): - """Hash a global/closure variable for caching purposes.""" +def _make_hashable(value, keep_alive: list): + """Hash a global/closure variable for caching purposes. + + Values that cannot be keyed by content fall back to ``id()``. Such a value + is also appended to ``keep_alive`` so the caller can keep it alive. This + prevents cache key collisions: once the object is freed, Python may give + its address to an unrelated object, and two different operators would + then produce equal keys. + """ # duck-type check for a JIT backend's dispatcher: if hasattr(value, "py_func") and callable(value.py_func): return CachableFunction(value.py_func) + elif isinstance(value, types.FunctionType): + # A captured plain Python function (e.g. a helper called from the + # operator) is keyed by its own content, recursively. + return CachableFunction(value) elif is_device_array(value): # The stateful op machinery threads the array's pointer and # shape through the *runtime* state buffer. Thus, the key # includes only dtype and rank, not pointer or shape. return (get_dtype(value), len(get_shape(value))) + elif isinstance(value, np.ndarray): + # A captured host array is frozen into the compiled operator as a + # constant, so arrays with different contents are different operators. + digest = hashlib.sha256(np.ascontiguousarray(value).tobytes()).digest() + return ("numpy.ndarray", value.dtype.str, value.shape, digest) elif isinstance(value, (np.number, np.bool_)): return ("numpy.scalar", value.dtype.str, value.tobytes()) - elif isinstance(value, (bool, int, float)): + elif isinstance(value, (bool, int, float, str, bytes)): # Python scalars are immutable values; key them by type and value so # equal-valued scalars share a cache entry. Without this they fall # through to ``id(value)`` below, and a fresh (non-interned) ``int``/ @@ -653,12 +669,16 @@ def _make_hashable(value): # collisions between like-named scalar subclasses from other modules). return ("python.scalar", _type_fqn(value), value) elif isinstance(value, (list, tuple)): - return tuple(_make_hashable(v) for v in value) + return tuple(_make_hashable(v, keep_alive) for v in value) elif isinstance(value, dict): return tuple( - sorted((_make_hashable(k), _make_hashable(v)) for k, v in value.items()) + sorted( + (_make_hashable(k, keep_alive), _make_hashable(v, keep_alive)) + for k, v in value.items() + ) ) else: + keep_alive.append(value) return id(value) @@ -706,7 +726,10 @@ class CachableFunction: based on its bytecode, constants, and closures, while ignoring other attributes such as its docstring. - It does not hold a reference to the function it was built from. + It does not hold a reference to the function it was built from. It does + hold a reference to any captured object that could only be keyed by its + ``id()`` (see :func:`_make_hashable`), so that address stays valid for + the lifetime of the key. """ # TODO: eventually, move this class to _jit.py as it only @@ -716,10 +739,11 @@ def __init__(self, func): self._name = func.__name__ closure = func.__closure__ if func.__closure__ is not None else [] + keep_alive: list = [] contents = [] # Make closure contents hashable for cell in closure: - contents.append(_make_hashable(cell.cell_contents)) + contents.append(_make_hashable(cell.cell_contents, keep_alive)) self._identity = ( func.__name__, func.__code__.co_code, @@ -733,10 +757,11 @@ def __init__(self, func): # use the hash of the name itself. Assumes numba # known how to interpret the dotted name at JIT # time. - _make_hashable(func.__globals__.get(name, name)) + _make_hashable(func.__globals__.get(name, name), keep_alive) for name in func.__code__.co_names ), ) + self._keep_alive = tuple(keep_alive) def __eq__(self, other): return self._identity == other._identity diff --git a/python/cuda_cccl/tests/compute/test_caching.py b/python/cuda_cccl/tests/compute/test_caching.py index 1fefff683ad4..c3579c8f2721 100644 --- a/python/cuda_cccl/tests/compute/test_caching.py +++ b/python/cuda_cccl/tests/compute/test_caching.py @@ -147,3 +147,65 @@ def test_cached_wrapper_does_not_retain_caller_arrays(run_algorithm): still_alive = [ref() for ref in array_refs if ref() is not None] assert still_alive == [] + + +# Operators keyed for the caches must not collide once the user has dropped +# them. The key does not keep the operator alive, so anything it identifies by +# address would otherwise be free to be reused by a later, different operator. + + +def _make_helper_op(bias): + def helper(x): + return x + bias + + def op(x): + return helper(x) + + return op + + +def _make_host_table_op(bias): + table = np.array([bias], dtype=np.int32) + + def op(x): + return x + table[0] + + return op + + +class _Opaque: + def __init__(self, bias): + self.bias = bias + + +def _make_opaque_op(bias): + cfg = _Opaque(bias) + + def op(x): + return x + cfg.bias + + return op + + +@pytest.mark.parametrize( + "factory", [_make_helper_op, _make_host_table_op, _make_opaque_op] +) +def test_operator_cache_identity_does_not_collide_after_operator_is_dropped(factory): + from cuda.compute._caching import CachableFunction + + seen = {} + for bias in range(2000): + op = factory(bias) + key = CachableFunction(op) + assert key not in seen, f"key for bias {bias} collided with {seen[key]}" + seen[key] = bias + del op + gc.collect() + + +@pytest.mark.parametrize("factory", [_make_helper_op, _make_host_table_op]) +def test_operator_cache_identity_is_content_based(factory): + from cuda.compute._caching import CachableFunction + + assert CachableFunction(factory(1)) == CachableFunction(factory(1)) + assert CachableFunction(factory(1)) != CachableFunction(factory(2)) From 7a571544dffca2d1e347a472f39a1dba00f04f04 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 2 Nov 2026 10:50:41 -0600 Subject: [PATCH 6/8] Address failures and reviews --- python/cuda_cccl/cuda/compute/_caching.py | 2 +- .../tests/compute/test_func_caching.py | 22 ++++++++++++------- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/python/cuda_cccl/cuda/compute/_caching.py b/python/cuda_cccl/cuda/compute/_caching.py index 22e319effcdd..df23ffffd963 100644 --- a/python/cuda_cccl/cuda/compute/_caching.py +++ b/python/cuda_cccl/cuda/compute/_caching.py @@ -657,7 +657,7 @@ def _make_hashable(value, keep_alive: list): # A captured host array is frozen into the compiled operator as a # constant, so arrays with different contents are different operators. digest = hashlib.sha256(np.ascontiguousarray(value).tobytes()).digest() - return ("numpy.ndarray", value.dtype.str, value.shape, digest) + return ("numpy.ndarray", value.dtype, value.shape, digest) elif isinstance(value, (np.number, np.bool_)): return ("numpy.scalar", value.dtype.str, value.tobytes()) elif isinstance(value, (bool, int, float, str, bytes)): diff --git a/python/cuda_cccl/tests/compute/test_func_caching.py b/python/cuda_cccl/tests/compute/test_func_caching.py index d57e17cb6448..d6cbb7c8f646 100644 --- a/python/cuda_cccl/tests/compute/test_func_caching.py +++ b/python/cuda_cccl/tests/compute/test_func_caching.py @@ -7,6 +7,12 @@ global_x = 1 +def _key(value): + # ``_make_hashable`` collects objects it keys by ``id()`` into a keep-alive + # list; the tests below compare keys only, so a throwaway list is fine. + return _make_hashable(value, []) + + def test_func_caching_basic(): def func(x): return x @@ -90,16 +96,16 @@ def test_make_hashable_python_scalars_keyed_by_value(): a = int(str(10**6)) b = int(str(10**6)) assert a is not b - assert _make_hashable(a) == _make_hashable(b) + assert _key(a) == _key(b) x = float(str(3.5)) y = float(str(3.5)) - assert _make_hashable(x) == _make_hashable(y) + assert _key(x) == _key(y) # Distinct values, types, and bool-vs-int must not collide. - assert _make_hashable(a) != _make_hashable(int(str(10**6 + 1))) - assert _make_hashable(1) != _make_hashable(1.0) - assert _make_hashable(True) != _make_hashable(1) + assert _key(a) != _key(int(str(10**6 + 1))) + assert _key(1) != _key(1.0) + assert _key(True) != _key(1) def test_func_caching_with_python_scalar_closure(): @@ -221,13 +227,13 @@ def test_make_hashable_device_array_keys_on_dtype_and_rank_only(): a = ProxyArray(np.int64) b = ProxyArray(np.int64) assert a is not b - assert _make_hashable(a) == _make_hashable(b) + assert _key(a) == _key(b) c = ProxyArray(np.int32) - assert _make_hashable(a) != _make_hashable(c) + assert _key(a) != _key(c) d = _RankedProxyArray(np.int64, (1, 1)) - assert _make_hashable(a) != _make_hashable(d) + assert _key(a) != _key(d) def test_func_caching_with_device_array_closure_ignores_shape(): From dc60e839b4a2f8ffa3f5033aab15c0e865ad7d70 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Tue, 10 Nov 2026 09:58:24 -0600 Subject: [PATCH 7/8] Address review --- python/cuda_cccl/cuda/compute/_caching.py | 40 +++++++++++++++---- .../tests/compute/test_func_caching.py | 31 ++++++++++++++ 2 files changed, 63 insertions(+), 8 deletions(-) diff --git a/python/cuda_cccl/cuda/compute/_caching.py b/python/cuda_cccl/cuda/compute/_caching.py index df23ffffd963..2dc7114f7604 100644 --- a/python/cuda_cccl/cuda/compute/_caching.py +++ b/python/cuda_cccl/cuda/compute/_caching.py @@ -632,7 +632,20 @@ def register(self, type_: type, key_function: Callable[[Any], Hashable]) -> None cache_with_registered_key_functions = _CacheWithRegisteredKeyFunctions() -def _make_hashable(value, keep_alive: list): +def _function_key(func, in_progress: tuple): + """Key a captured function by content, unless its key is already being + built further up the walk (the operator refers to itself, directly or via + another function). Such a back-reference is keyed by the position of that + function in ``in_progress``; identical walks produce identical positions, + and descending again would never terminate. + """ + for depth, active in enumerate(in_progress): + if active is func: + return ("recursive", depth) + return CachableFunction(func, _in_progress=in_progress) + + +def _make_hashable(value, keep_alive: list, in_progress: tuple = ()): """Hash a global/closure variable for caching purposes. Values that cannot be keyed by content fall back to ``id()``. Such a value @@ -640,14 +653,17 @@ def _make_hashable(value, keep_alive: list): prevents cache key collisions: once the object is freed, Python may give its address to an unrelated object, and two different operators would then produce equal keys. + + ``in_progress`` holds the functions whose keys are being built up the + walk; see :func:`_function_key`. """ # duck-type check for a JIT backend's dispatcher: if hasattr(value, "py_func") and callable(value.py_func): - return CachableFunction(value.py_func) + return _function_key(value.py_func, in_progress) elif isinstance(value, types.FunctionType): # A captured plain Python function (e.g. a helper called from the # operator) is keyed by its own content, recursively. - return CachableFunction(value) + return _function_key(value, in_progress) elif is_device_array(value): # The stateful op machinery threads the array's pointer and # shape through the *runtime* state buffer. Thus, the key @@ -669,11 +685,14 @@ def _make_hashable(value, keep_alive: list): # collisions between like-named scalar subclasses from other modules). return ("python.scalar", _type_fqn(value), value) elif isinstance(value, (list, tuple)): - return tuple(_make_hashable(v, keep_alive) for v in value) + return tuple(_make_hashable(v, keep_alive, in_progress) for v in value) elif isinstance(value, dict): return tuple( sorted( - (_make_hashable(k, keep_alive), _make_hashable(v, keep_alive)) + ( + _make_hashable(k, keep_alive, in_progress), + _make_hashable(v, keep_alive, in_progress), + ) for k, v in value.items() ) ) @@ -735,15 +754,18 @@ class CachableFunction: # TODO: eventually, move this class to _jit.py as it only # has to do with caching of Python callables that will be # JIT compiled. - def __init__(self, func): + def __init__(self, func, _in_progress: tuple = ()): self._name = func.__name__ + # Functions whose keys are being built, outermost first, so a + # reference back to one of them does not recurse forever. + in_progress = (*_in_progress, func) closure = func.__closure__ if func.__closure__ is not None else [] keep_alive: list = [] contents = [] # Make closure contents hashable for cell in closure: - contents.append(_make_hashable(cell.cell_contents, keep_alive)) + contents.append(_make_hashable(cell.cell_contents, keep_alive, in_progress)) self._identity = ( func.__name__, func.__code__.co_code, @@ -757,7 +779,9 @@ def __init__(self, func): # use the hash of the name itself. Assumes numba # known how to interpret the dotted name at JIT # time. - _make_hashable(func.__globals__.get(name, name), keep_alive) + _make_hashable( + func.__globals__.get(name, name), keep_alive, in_progress + ) for name in func.__code__.co_names ), ) diff --git a/python/cuda_cccl/tests/compute/test_func_caching.py b/python/cuda_cccl/tests/compute/test_func_caching.py index d6cbb7c8f646..bc932e245fd2 100644 --- a/python/cuda_cccl/tests/compute/test_func_caching.py +++ b/python/cuda_cccl/tests/compute/test_func_caching.py @@ -7,6 +7,24 @@ global_x = 1 +def add(x, y): + # Module-level operator whose body names itself: ``add`` is both this + # function's global name and the attribute name in ``np.add``. + return np.add(x, y) + + +def fact(n): + return 1 if n <= 1 else n * fact(n - 1) + + +def _ping(n): + return _pong(n - 1) if n else 0 + + +def _pong(n): + return _ping(n - 1) if n else 1 + + def _key(value): # ``_make_hashable`` collects objects it keys by ``id()`` into a keep-alive # list; the tests below compare keys only, so a throwaway list is fine. @@ -275,3 +293,16 @@ def func(x): func2 = make_func2() assert CachableFunction(func1) != CachableFunction(func2) + + +def test_self_referencing_operator_key_terminates(): + # Regression test: keying a function by content must not recurse into + # itself when its own name appears among the globals it references. + assert CachableFunction(add) == CachableFunction(add) + assert CachableFunction(fact) == CachableFunction(fact) + assert CachableFunction(add) != CachableFunction(fact) + + +def test_mutually_recursive_operator_keys_terminate_and_differ(): + assert CachableFunction(_ping) == CachableFunction(_ping) + assert CachableFunction(_ping) != CachableFunction(_pong) From 88dea82d9022187ed9fc5918c06449616b7e6019 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Tue, 10 Nov 2026 12:56:24 -0600 Subject: [PATCH 8/8] Address review --- python/cuda_cccl/cuda/compute/_caching.py | 22 ++++++++++- .../cuda_cccl/tests/compute/test_caching.py | 38 +++++++++++++++++++ 2 files changed, 59 insertions(+), 1 deletion(-) diff --git a/python/cuda_cccl/cuda/compute/_caching.py b/python/cuda_cccl/cuda/compute/_caching.py index 2dc7114f7604..31c726350ae0 100644 --- a/python/cuda_cccl/cuda/compute/_caching.py +++ b/python/cuda_cccl/cuda/compute/_caching.py @@ -20,7 +20,13 @@ except ImportError: from cuda.core.experimental._utils.cuda_utils import CUDAError -from ._utils.protocols import get_dtype, get_shape, is_contiguous, is_device_array +from ._utils.protocols import ( + get_data_pointer, + get_dtype, + get_shape, + is_contiguous, + is_device_array, +) from .struct import _Struct try: @@ -665,6 +671,20 @@ def _make_hashable(value, keep_alive: list, in_progress: tuple = ()): # operator) is keyed by its own content, recursively. return _function_key(value, in_progress) elif is_device_array(value): + if len(in_progress) > 1: + # A device array captured by a *nested* function (a helper the + # operator calls) is not threaded through the runtime state; the + # compiler bakes its pointer, shape and strides into the device + # code as constants. Key it by exactly those, so two operators + # whose helpers capture different arrays never share a kernel. + cai = value.__cuda_array_interface__ + return ( + "device_array.baked", + get_dtype(value), + tuple(cai["shape"]), + cai.get("strides"), + get_data_pointer(value), + ) # The stateful op machinery threads the array's pointer and # shape through the *runtime* state buffer. Thus, the key # includes only dtype and rank, not pointer or shape. diff --git a/python/cuda_cccl/tests/compute/test_caching.py b/python/cuda_cccl/tests/compute/test_caching.py index c3579c8f2721..0a9c34b2eeb9 100644 --- a/python/cuda_cccl/tests/compute/test_caching.py +++ b/python/cuda_cccl/tests/compute/test_caching.py @@ -209,3 +209,41 @@ def test_operator_cache_identity_is_content_based(factory): assert CachableFunction(factory(1)) == CachableFunction(factory(1)) assert CachableFunction(factory(1)) != CachableFunction(factory(2)) + + +def _make_op_with_registered_helper(table): + # A helper made callable from device code captures ``table``. The outer + # operator captures no array itself, so it is stateless and the compiler + # bakes the helper's array pointer into the kernel as a constant. + from numba_cuda_mlir.extending import register_jitable + + from cuda.compute import _mlir + + @register_jitable(typing_registry=_mlir.typing_registry) + def helper(x): + return x + table[0] + + def op(x): + return helper(x) + + return op + + +def test_operators_with_helper_captured_arrays_do_not_share_a_kernel(): + _require_jit_backend() + from cuda.compute._caching import CachableFunction + + t1 = DeviceArray.from_numpy(np.array([100], dtype=np.int32)) + t2 = DeviceArray.from_numpy(np.array([200], dtype=np.int32)) + op1 = _make_op_with_registered_helper(t1) + op2 = _make_op_with_registered_helper(t2) + assert CachableFunction(op1) != CachableFunction(op2) + + d_in = DeviceArray.from_numpy(np.arange(4, dtype=np.int32)) + d_out = DeviceArray.empty(4, np.int32) + cuda.compute.unary_transform(d_in=d_in, d_out=d_out, op=op1, num_items=4) + Device().sync() + np.testing.assert_array_equal(d_out.copy_to_host(), np.arange(4) + 100) + cuda.compute.unary_transform(d_in=d_in, d_out=d_out, op=op2, num_items=4) + Device().sync() + np.testing.assert_array_equal(d_out.copy_to_host(), np.arange(4) + 200)