[cuda.compute] Fix device memory leaks caused by caching

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 <cuda.compute.op.OpKind>` value or a

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

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,15 +838,15 @@ 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.
         """
         cdef Pointer cached = self._cached_ptr_obj
         if cached is None:
             cached = Pointer(0)
             self._cached_ptr_obj = cached
-        cached.rebind(ptr, owner)
+        cached.rebind(ptr)
         self.state = cached
 
     @property

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
@@ -19,7 +20,13 @@
 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_data_pointer,
+    get_dtype,
+    get_shape,
+    is_contiguous,
+    is_device_array,
+)
 from .struct import _Struct
 
 try:
@@ -499,6 +506,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.
@@ -573,8 +606,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
@@ -604,19 +638,65 @@ 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 _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
+    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.
+
+    ``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 _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.
         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, 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``/
@@ -625,12 +705,19 @@ 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, in_progress) 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, in_progress),
+                    _make_hashable(v, keep_alive, in_progress),
+                )
+                for k, v in value.items()
+            )
         )
     else:
+        keep_alive.append(value)
         return id(value)
 
 
@@ -662,7 +749,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,25 +761,31 @@ def clear_all_caches():
 
 class CachableFunction:
     """
-    A type that wraps a function and provides custom comparison
-    (__eq__) and hash (__hash__) implementations.
-
-    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.
+    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.
+
+    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
     # has to do with caching of Python callables that will be
     # JIT compiled.
-    def __init__(self, func):
-        self._func = 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))
+            contents.append(_make_hashable(cell.cell_contents, keep_alive, in_progress))
         self._identity = (
             func.__name__,
             func.__code__.co_code,
@@ -706,10 +799,13 @@ 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, in_progress
+                )
                 for name in func.__code__.co_names
             ),
         )
+        self._keep_alive = tuple(keep_alive)
 
     def __eq__(self, other):
         return self._identity == other._identity
@@ -718,7 +814,12 @@ def __hash__(self):
         return hash(self._identity)
 
     def __repr__(self):
-        return str(self._func)
+        # 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"<CachableFunction {self._name} {digest}>"
 
 
 # Register keyers for built-in types

python/cuda_cccl/cuda/compute/_cccl_interop.py

@@ -97,18 +97,22 @@ 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.")
     dtype = get_dtype(array)
 
     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,
@@ -228,15 +232,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()