Support operations with side-effects (state) in cuda.compute

python/cuda_cccl/benchmarks/compute/bench_select.py

@@ -125,3 +125,24 @@ def run():
 
     fixture = request.getfixturevalue(bench_fixture)
     fixture(run)
+
+
+@pytest.mark.parametrize("bench_fixture", ["compile_benchmark", "benchmark"])
+def bench_select_stateful(bench_fixture, request, size):
+    actual_size = 100 if bench_fixture == "compile_benchmark" else size
+    inp = cp.random.randint(0, 100, actual_size, dtype=np.int32)
+    out = cp.empty_like(inp)
+    num_selected = cp.empty(2, dtype=np.uint64)
+    threshold_state = cp.array([50], dtype=np.int32)
+
+    def run():
+        select_stateful(
+            inp,
+            out,
+            num_selected,
+            threshold_state,
+            build_only=(bench_fixture == "compile_benchmark"),
+        )
+
+    fixture = request.getfixturevalue(bench_fixture)
+    fixture(run)

python/cuda_cccl/cuda/compute/_caching.py

@@ -10,6 +10,7 @@
 except ImportError:
     from cuda.core.experimental import Device
 
+
 # Central registry of all algorithm caches
 _cache_registry: dict[str, object] = {}
 
@@ -55,6 +56,33 @@ def cache_clear():
     return deco
 
 
+def _hash_device_array_like(value):
+    # hash based on pointer, shape, and dtype
+    ptr = value.__cuda_array_interface__["data"][0]
+    shape = value.__cuda_array_interface__["shape"]
+    dtype = value.__cuda_array_interface__["typestr"]
+    return hash((ptr, shape, dtype))
+
+
+def _make_hashable(value):
+    import numba.cuda.dispatcher
+
+    from .typing import DeviceArrayLike
+
+    if isinstance(value, numba.cuda.dispatcher.CUDADispatcher):
+        return CachableFunction(value.py_func)
+    elif isinstance(value, DeviceArrayLike):
+        return _hash_device_array_like(value)
+    elif isinstance(value, (list, tuple)):
+        return tuple(_make_hashable(v) for v in value)
+    elif isinstance(value, dict):
+        return tuple(
+            sorted((_make_hashable(k), _make_hashable(v)) for k, v in value.items())
+        )
+    else:
+        return id(value)
+
+
 def clear_all_caches():
     """
     Clear all algorithm caches.
@@ -83,25 +111,23 @@ class CachableFunction:
     """
 
     def __init__(self, func):
-        import numba.cuda.dispatcher
-
         self._func = func
 
         closure = func.__closure__ if func.__closure__ is not None else []
         contents = []
         # if any of the contents is a numba.cuda.dispatcher.CUDADispatcher
         # use the function for caching purposes:
         for cell in closure:
-            if isinstance(cell.cell_contents, numba.cuda.dispatcher.CUDADispatcher):
-                contents.append(CachableFunction(cell.cell_contents.py_func))
-            else:
-                contents.append(cell.cell_contents)
+            contents.append(_make_hashable(cell.cell_contents))
         self._identity = (
             func.__name__,
             func.__code__.co_code,
             func.__code__.co_consts,
             tuple(contents),
-            tuple(func.__globals__.get(name, None) for name in func.__code__.co_names),
+            tuple(
+                _make_hashable(func.__globals__.get(name, None))
+                for name in func.__code__.co_names
+            ),
         )
 
     def __eq__(self, other):

python/cuda_cccl/pyproject.toml

@@ -31,7 +31,6 @@ dependencies = [
   "numpy",
   "cuda-pathfinder>=1.2.3",
   "cuda-core",
-  "numba-cuda>=0.20.0,!=0.21.2",
   "typing_extensions",
 ]
 
@@ -42,12 +41,12 @@ readme = { file = "README.md", content-type = "text/markdown" }
 cu12 = [
   "cuda-bindings>=12.9.1,<13.0.0",
   "cuda-toolkit[nvrtc,nvjitlink,cudart,nvcc]==12.*",
-  "numba-cuda[cu12]>=0.20.0,!=0.21.2",
+  "numba-cuda[cu12]>=0.23.0",
 ]
 cu13 = [
   "cuda-bindings>=13.0.0,<14.0.0",
   "cuda-toolkit[nvrtc,nvjitlink,cudart,nvcc,nvvm]==13.*",
-  "numba-cuda[cu13]>=0.20.0,!=0.21.2",
+  "numba-cuda[cu13]>=0.23.0",
 ]
 test-cu12 = [
   # an undocumented way to inherit the dependencies of the cu12 extra.

python/cuda_cccl/tests/compute/examples/select/select_with_side_effect.py

@@ -0,0 +1,48 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+#
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+# example-begin
+import cupy as cp
+from numba import cuda as numba_cuda
+
+from cuda.compute.algorithms import select
+
+# Create input data: values 0 to 99
+d_in = cp.arange(100, dtype=cp.int32)
+d_out = cp.empty_like(d_in)
+d_num_selected = cp.empty(1, dtype=cp.uint64)
+
+# Counter for rejected items (side effect state)
+reject_count = cp.zeros(1, dtype=cp.int32)
+
+
+# Define condition that counts rejected items as a side effect
+def count_rejects(x):
+    if x % 2 == 0:
+        return True
+    else:
+        numba_cuda.atomic.add(reject_count, 0, 1)
+        return False
+
+
+# Execute select - selects even numbers, counts rejections
+select(d_in, d_out, d_num_selected, count_rejects, len(d_in))
+
+# Get results
+num_selected = int(d_num_selected.get()[0])
+num_rejected = int(reject_count.get()[0])
+result = d_out[:num_selected].get()
+
+print(f"Selected {num_selected} items (values % 2 == 0)")
+print(f"Rejected {num_rejected} items (values % 2 != 0)")
+print(f"First 5 selected: {result[:5]}")
+# Output:
+# Selected 50 items (even numbers)
+# Rejected 50 items (odd numbers)
+# First 5 selected: [0 2 4 6 8]
+# example-end
+
+assert num_selected == 50  # Even numbers
+assert num_rejected == 50  # Odd numbers

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_with_side_effect_counting_rejects():
+    """Select with side effect that counts rejected items"""
+    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