Add support for ops that capture device arrays from enclosing scope

python/cuda_cccl/cuda/compute/__init__.py

@@ -89,3 +89,7 @@
     "unique_by_key",
     "ZipIterator",
 ]
+
+# import extensions to numba CUDA enabling global capture of
+# device arrays.
+from . import _device_array_capture  # noqa: F403, F401

python/cuda_cccl/cuda/compute/_caching.py

@@ -77,7 +77,11 @@ def __eq__(self, other):
         return self._identity == other._identity
 
     def __hash__(self):
-        return hash(self._identity)
+        try:
+            return hash(self._identity)
+        except TypeError:
+            # if we can't hash the identity, use its id
+            return id(self._identity)
 
     def __repr__(self):
         return str(self._func)

python/cuda_cccl/cuda/compute/_device_array_capture.py

@@ -0,0 +1,197 @@
+"""
+Extends Numba CUDA to enable global capture of objects implementing
+__cuda_array_interface__.
+
+HOW IT WORKS:
+1. Patches typeof_impl to check for DeviceArrayLike objects
+2. Defines a DeviceArrayType (subclass of Array) to represent DeviceArrayLike objects
+3. Registers data model and lower_constant for the type
+"""
+
+import numpy as np
+from numba.cuda import types
+from numba.cuda.core.imputils import builtin_registry
+from numba.cuda.datamodel.models import StructModel
+from numba.cuda.datamodel.registry import register_default
+from numba.cuda.np import numpy_support
+from numba.cuda.typing import typeof as typeof_module
+
+from .typing import DeviceArrayLike
+
+# =============================================================================
+# Step 1: Define a custom Array subclass for DeviceArrayLike objects
+# =============================================================================
+
+
+class DeviceArrayType(types.Array):
+    """
+    Type for any DeviceArrayLike object.
+
+    By subclassing Array and overriding copy() to ignore readonly,
+    we ensure these arrays remain mutable when captured from globals.
+    """
+
+    def __init__(self, dtype, ndim, layout, readonly=False, aligned=True):
+        type_name = "device_array"
+        if readonly:
+            type_name = "readonly " + type_name
+        name = f"{type_name}({dtype}, {ndim}d, {layout})"
+        super().__init__(
+            dtype, ndim, layout, readonly=readonly, name=name, aligned=aligned
+        )
+
+    def copy(self, dtype=None, ndim=None, layout=None, readonly=None):
+        # IGNORE readonly parameter - CAI objects are device memory references
+        # that should remain mutable
+        return DeviceArrayType(
+            dtype=dtype if dtype is not None else self.dtype,
+            ndim=ndim if ndim is not None else self.ndim,
+            layout=layout if layout is not None else self.layout,
+            readonly=not self.mutable,  # Preserve original mutability
+            aligned=self.aligned,
+        )
+
+
+# =============================================================================
+# Step 2: Define _typeof_device_array to type DeviceArrayLike objects
+# =============================================================================
+
+
+def _typeof_device_array(val, c):
+    """Type of a DevieArrayLike object."""
+    interface = val.__cuda_array_interface__
+
+    dtype = numpy_support.from_dtype(np.dtype(interface["typestr"]))
+    shape = interface["shape"]
+    ndim = len(shape)
+    strides = interface.get("strides")
+
+    # Determine layout
+    if ndim == 0:
+        layout = "C"
+    elif strides is None:
+        layout = "C"
+    else:
+        itemsize = np.dtype(interface["typestr"]).itemsize
+        # Check C-contiguous
+        c_strides = []
+        stride = itemsize
+        for i in range(ndim - 1, -1, -1):
+            c_strides.insert(0, stride)
+            stride *= shape[i]
+
+        if tuple(strides) == tuple(c_strides):
+            layout = "C"
+        else:
+            # Check F-contiguous
+            f_strides = []
+            stride = itemsize
+            for i in range(ndim):
+                f_strides.append(stride)
+                stride *= shape[i]
+            if tuple(strides) == tuple(f_strides):
+                layout = "F"
+            else:
+                layout = "A"
+
+    readonly = interface["data"][1]
+    return DeviceArrayType(dtype, ndim, layout, readonly=readonly)
+
+
+# =============================================================================
+# Step 3: Patch typeof_impl to handle CAI objects
+# =============================================================================
+
+# Get the original generic fallback
+_original_typeof_impl = typeof_module.typeof_impl
+_original_generic = _original_typeof_impl.dispatch(object)
+
+
+def _patched_generic(val, c):
+    """
+    Patched generic typeof handler that checks for DeviceArrayLike objects.
+    """
+    # Check for CAI FIRST
+    if isinstance(val, DeviceArrayLike):
+        return _typeof_device_array(val, c)
+
+    # Fall back to original behavior
+    return _original_generic(val, c)
+
+
+# Register our patched handler for the generic object case
+_original_typeof_impl.register(object)(_patched_generic)
+
+
+# =============================================================================
+# Step 4: Register data model for our type
+# =============================================================================
+
+
+@register_default(DeviceArrayType)
+class DeviceArrayModel(StructModel):
+    """Data model for DeviceArrayType - same as regular Array."""
+
+    def __init__(self, dmm, fe_type):
+        ndim = fe_type.ndim
+        members = [
+            ("meminfo", types.MemInfoPointer(fe_type.dtype)),
+            ("parent", types.pyobject),
+            ("nitems", types.intp),
+            ("itemsize", types.intp),
+            ("data", types.CPointer(fe_type.dtype)),
+            ("shape", types.UniTuple(types.intp, ndim)),
+            ("strides", types.UniTuple(types.intp, ndim)),
+        ]
+        super().__init__(dmm, fe_type, members)
+
+
+# =============================================================================
+# Step 5: Register lower_constant for our type
+# =============================================================================
+
+
+@builtin_registry.lower_constant(DeviceArrayType)
+def lower_device_array(context, builder, ty, pyval):
+    """
+    Lower DeviceArrayLike objects by embedding the device pointer as a constant.
+    """
+    interface = pyval.__cuda_array_interface__
+
+    shape = interface["shape"]
+    strides = interface.get("strides")
+    data_ptr = interface["data"][0]
+    typestr = interface["typestr"]
+
+    # Calculate strides if not provided (C-contiguous)
+    if strides is None:
+        itemsize = np.dtype(typestr).itemsize
+        ndim = len(shape)
+        strides = []
+        stride = itemsize
+        for i in range(ndim - 1, -1, -1):
+            strides.insert(0, stride)
+            stride *= shape[i]
+        strides = tuple(strides)
+
+    # Embed device pointer as constant
+    llvoidptr = context.get_value_type(types.voidptr)
+    data = context.get_constant(types.uintp, data_ptr).inttoptr(llvoidptr)
+
+    # Build array structure
+    ary = context.make_array(ty)(context, builder)
+    kshape = [context.get_constant(types.intp, s) for s in shape]
+    kstrides = [context.get_constant(types.intp, s) for s in strides]
+    itemsize = np.dtype(typestr).itemsize
+
+    context.populate_array(
+        ary,
+        data=builder.bitcast(data, ary.data.type),
+        shape=kshape,
+        strides=kstrides,
+        itemsize=context.get_constant(types.intp, itemsize),
+        parent=None,
+        meminfo=None,
+    )
+
+    return ary._getvalue()

python/cuda_cccl/tests/compute/test_select.py

@@ -420,3 +420,126 @@ def condition(pair):
     expected_count = np.sum(h_sums < 70)
 
     assert num_selected == expected_count
+
+
+def test_select_stateful_threshold():
+    """Test stateful select that uses state for threshold"""
+    num_items = 1000
+    h_in = random_array(num_items, np.int32, max_value=100)
+
+    # Create device state containing threshold value
+    threshold_value = 50
+    threshold_state = cp.array([threshold_value], dtype=np.int32)
+
+    # Define condition that references state as closure
+    def threshold_select(x):
+        return x > threshold_state[0]
+
+    d_in = cp.asarray(h_in)
+    d_out = cp.empty_like(d_in)
+    d_num_selected = cp.empty(2, dtype=np.uint64)
+
+    cuda.compute.select(
+        d_in,
+        d_out,
+        d_num_selected,
+        threshold_select,
+        num_items,
+    )
+
+    # Check selected output
+    num_selected = int(d_num_selected[0].get())
+    got = d_out.get()[:num_selected]
+
+    # Verify all output values are > threshold
+    assert np.all(got > threshold_value)
+
+    # Verify we got the expected number of items
+    expected_selected = h_in[h_in > threshold_value]
+    expected_count = len(expected_selected)
+
+    assert num_selected == expected_count
+
+    # Verify exact results
+    assert np.array_equal(got, expected_selected)
+
+
+def test_select_stateful_atomic():
+    """Test stateful select with atomic operations to count rejected items"""
+    from numba import cuda as numba_cuda
+
+    num_items = 1000
+    h_in = random_array(num_items, np.int32, max_value=100)
+
+    # Create device state for counting rejected items
+    reject_counter = cp.zeros(1, dtype=np.int32)
+
+    # Define condition that references state as closure
+    def count_rejects(x):
+        if x > 50:
+            return True
+        else:
+            numba_cuda.atomic.add(reject_counter, 0, 1)
+            return False
+
+    d_in = cp.asarray(h_in)
+    d_out = cp.empty_like(d_in)
+    d_num_selected = cp.empty(2, dtype=np.uint64)
+
+    cuda.compute.select(
+        d_in,
+        d_out,
+        d_num_selected,
+        count_rejects,
+        num_items,
+    )
+
+    # Check selected output
+    num_selected = int(d_num_selected[0].get())
+    got = d_out.get()[:num_selected]
+
+    # Verify all output values are > 50
+    assert np.all(got > 50)
+
+    # Verify we got the expected number of items
+    expected_selected = h_in[h_in > 50]
+    expected_count = len(expected_selected)
+
+    assert num_selected == expected_count
+
+    # Verify exact results
+    assert np.array_equal(got, expected_selected)
+
+    # Verify state contains count of rejected items
+    rejected_count = int(reject_counter[0].get())
+    expected_rejected = len(h_in[h_in <= 50])
+    assert rejected_count == expected_rejected, (
+        f"Expected {expected_rejected} rejections, got {rejected_count}"
+    )
+
+
+def test_select_stateful_counting_rejects():
+    """Test select that counts rejected items using state"""
+    from numba import cuda as numba_cuda
+
+    d_in = cp.arange(100, dtype=np.int32)
+    d_out = cp.empty_like(d_in)
+    d_num_selected = cp.empty(1, dtype=np.uint64)
+
+    reject_count = cp.zeros(1, dtype=np.int32)
+
+    # Define condition that references state as closure
+    def count_rejects(x):
+        if x >= 50:
+            return True
+        else:
+            numba_cuda.atomic.add(reject_count, 0, 1)
+            return False
+
+    cuda.compute.select(d_in, d_out, d_num_selected, count_rejects, len(d_in))
+
+    num_selected = int(d_num_selected.get()[0])
+    num_rejected = int(reject_count.get()[0])
+
+    assert num_selected == 50  # Values 50-99
+    assert num_rejected == 50  # Values 0-49

python/cuda_cccl/tests/compute/test_transform.py

@@ -380,3 +380,27 @@ def add_vectors(v1: Vec2D, v2: Vec2D) -> Vec2D:
 
     np.testing.assert_equal(result["x"], h_in1["x"] + h_in2["x"])
     np.testing.assert_equal(result["y"], h_in1["y"] + h_in2["y"])
+
+
+def test_unary_transform_stateful_counting():
+    """Test unary_transform with state that counts even numbers."""
+    from numba import cuda as numba_cuda
+
+    d_in = cp.arange(100, dtype=np.int32)
+    d_out = cp.empty_like(d_in)
+
+    even_count = cp.zeros(1, dtype=np.int32)
+
+    # Define op that references state as closure
+    def count_evens(x):
+        if x % 2 == 0:
+            numba_cuda.atomic.add(even_count, 0, 1)
+        return x * 2
+
+    cuda.compute.unary_transform(d_in, d_out, count_evens, len(d_in))
+
+    expected_output = cp.arange(100, dtype=np.int32) * 2
+    np.testing.assert_array_equal(d_out.get(), expected_output.get())
+
+    num_evens = int(even_count.get()[0])
+    assert num_evens == 50  # 0, 2, 4, ..., 98