From 6e8a6f897387af7a01f98d6984f7ab34e93b98ec Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Thu, 11 Dec 2025 14:23:34 +0000 Subject: [PATCH 1/9] Move algorithm cache to a central registry --- python/cuda_cccl/cuda/compute/_caching.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/cuda_cccl/cuda/compute/_caching.py b/python/cuda_cccl/cuda/compute/_caching.py index ad675c05341..d7956effef6 100644 --- a/python/cuda_cccl/cuda/compute/_caching.py +++ b/python/cuda_cccl/cuda/compute/_caching.py @@ -13,6 +13,9 @@ # Central registry of all algorithm caches _cache_registry: dict[str, object] = {} +# Central registry of all algorithm caches +_cache_registry: dict[str, object] = {} + def cache_with_key(key): """ From b9d3f9cc969f102c448c1e52c6b56ee4d417b656 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Thu, 11 Dec 2025 21:18:14 +0000 Subject: [PATCH 2/9] Add bench_select.py --- .../benchmarks/compute/bench_select.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/python/cuda_cccl/benchmarks/compute/bench_select.py b/python/cuda_cccl/benchmarks/compute/bench_select.py index 7f3f38609f6..12c7de015fd 100644 --- a/python/cuda_cccl/benchmarks/compute/bench_select.py +++ b/python/cuda_cccl/benchmarks/compute/bench_select.py @@ -30,7 +30,8 @@ def select_iterator(size, d_in, out, num_selected, build_only): def less_than_50(x): return x < 50 - selector = cuda.compute.make_select(d_in_iter, out, num_selected, less_than_50) + selector = cuda.compute.make_select( + d_in_iter, out, num_selected, less_than_50) if not build_only: temp_bytes = selector(None, d_in_iter, out, num_selected, size) temp_storage = cp.empty(temp_bytes, dtype=np.uint8) @@ -51,7 +52,8 @@ def select_struct(inp, out, num_selected, build_only): def in_first_quadrant(p: Point) -> np.uint8: return (p.x > 50) and (p.y > 50) - selector = cuda.compute.make_select(inp, out, num_selected, in_first_quadrant) + selector = cuda.compute.make_select( + inp, out, num_selected, in_first_quadrant) if not build_only: temp_bytes = selector(None, inp, out, num_selected, size) temp_storage = cp.empty(temp_bytes, dtype=np.uint8) @@ -66,7 +68,8 @@ def select_stateful(inp, out, num_selected, threshold_state, build_only): def threshold_select(x): return x > threshold_state[0] - selector = cuda.compute.make_select(inp, out, num_selected, threshold_select) + selector = cuda.compute.make_select( + inp, out, num_selected, threshold_select) if not build_only: temp_bytes = selector(None, inp, out, num_selected, size) temp_storage = cp.empty(temp_bytes, dtype=np.uint8) @@ -84,7 +87,8 @@ def bench_select_pointer(bench_fixture, request, size): def run(): select_pointer( - inp, out, num_selected, build_only=(bench_fixture == "compile_benchmark") + inp, out, num_selected, build_only=( + bench_fixture == "compile_benchmark") ) fixture = request.getfixturevalue(bench_fixture) @@ -114,13 +118,15 @@ def run(): @pytest.mark.parametrize("bench_fixture", ["compile_benchmark", "benchmark"]) def bench_select_struct(bench_fixture, request, size): actual_size = 100 if bench_fixture == "compile_benchmark" else size - inp = cp.random.randint(0, 100, (actual_size, 2), dtype=np.int32).view(Point.dtype) + inp = cp.random.randint(0, 100, (actual_size, 2), + dtype=np.int32).view(Point.dtype) out = cp.empty_like(inp) num_selected = cp.empty(2, dtype=np.uint64) def run(): select_struct( - inp, out, num_selected, build_only=(bench_fixture == "compile_benchmark") + inp, out, num_selected, build_only=( + bench_fixture == "compile_benchmark") ) fixture = request.getfixturevalue(bench_fixture) From 036816578449cebfbd79406acf4bf515a5a17699 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Wed, 17 Dec 2025 16:00:18 +0000 Subject: [PATCH 3/9] Add tests for stateful select and transform --- python/cuda_cccl/tests/compute/test_select.py | 123 ++++++++++++++++++ .../cuda_cccl/tests/compute/test_transform.py | 24 ++++ 2 files changed, 147 insertions(+) diff --git a/python/cuda_cccl/tests/compute/test_select.py b/python/cuda_cccl/tests/compute/test_select.py index b6b8f311f53..13b3128262e 100644 --- a/python/cuda_cccl/tests/compute/test_select.py +++ b/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 diff --git a/python/cuda_cccl/tests/compute/test_transform.py b/python/cuda_cccl/tests/compute/test_transform.py index dc3aceabcb3..a4712e46249 100644 --- a/python/cuda_cccl/tests/compute/test_transform.py +++ b/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 From 43f0ff8b0cb756fdef1d685a71f3df072db76c15 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Fri, 12 Dec 2025 20:45:59 +0000 Subject: [PATCH 4/9] For the purposes of caching, hash DeviceArrayLike objects by pointer, shape, and dtype --- python/cuda_cccl/cuda/compute/_caching.py | 39 +++++++++++++++++++---- 1 file changed, 32 insertions(+), 7 deletions(-) diff --git a/python/cuda_cccl/cuda/compute/_caching.py b/python/cuda_cccl/cuda/compute/_caching.py index d7956effef6..2c08090d65f 100644 --- a/python/cuda_cccl/cuda/compute/_caching.py +++ b/python/cuda_cccl/cuda/compute/_caching.py @@ -58,6 +58,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. @@ -86,8 +113,6 @@ class CachableFunction: """ def __init__(self, func): - import numba.cuda.dispatcher - self._func = func closure = func.__closure__ if func.__closure__ is not None else [] @@ -95,16 +120,16 @@ def __init__(self, func): # 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): From 547704bf73579c6e698ef09b33d06f0c4fe2bce5 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Wed, 17 Dec 2025 16:12:27 +0000 Subject: [PATCH 5/9] Update select benchmark --- .../benchmarks/compute/bench_select.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/python/cuda_cccl/benchmarks/compute/bench_select.py b/python/cuda_cccl/benchmarks/compute/bench_select.py index 12c7de015fd..e8f62042fe2 100644 --- a/python/cuda_cccl/benchmarks/compute/bench_select.py +++ b/python/cuda_cccl/benchmarks/compute/bench_select.py @@ -131,3 +131,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) From d67eac4a2397e777038c60ee0c7d8d1291fec52e Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Thu, 18 Dec 2025 10:08:16 +0000 Subject: [PATCH 6/9] Bump numba-cuda dependency to 0.23.0 --- python/cuda_cccl/pyproject.toml | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/python/cuda_cccl/pyproject.toml b/python/cuda_cccl/pyproject.toml index e414dfda371..a401fbe4bcd 100644 --- a/python/cuda_cccl/pyproject.toml +++ b/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. From 5e29ec9f6c3be057fdced34c09295374a1c1e465 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Thu, 18 Dec 2025 10:13:46 +0000 Subject: [PATCH 7/9] Add select example --- .../select/select_with_side_effect.py | 48 +++++++++++++++++++ python/cuda_cccl/tests/compute/test_select.py | 4 +- 2 files changed, 50 insertions(+), 2 deletions(-) create mode 100644 python/cuda_cccl/tests/compute/examples/select/select_with_side_effect.py diff --git a/python/cuda_cccl/tests/compute/examples/select/select_with_side_effect.py b/python/cuda_cccl/tests/compute/examples/select/select_with_side_effect.py new file mode 100644 index 00000000000..a39f28d3bb4 --- /dev/null +++ b/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 diff --git a/python/cuda_cccl/tests/compute/test_select.py b/python/cuda_cccl/tests/compute/test_select.py index 13b3128262e..e0c2e8f2c14 100644 --- a/python/cuda_cccl/tests/compute/test_select.py +++ b/python/cuda_cccl/tests/compute/test_select.py @@ -518,8 +518,8 @@ def count_rejects(x): ) -def test_select_stateful_counting_rejects(): - """Test select that counts rejected items using state""" +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) From 7f26f980e656200508ce7ce6f5338da44c9684e4 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Thu, 18 Dec 2025 10:19:37 +0000 Subject: [PATCH 8/9] Lint --- .../benchmarks/compute/bench_select.py | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/python/cuda_cccl/benchmarks/compute/bench_select.py b/python/cuda_cccl/benchmarks/compute/bench_select.py index e8f62042fe2..8d284c63a74 100644 --- a/python/cuda_cccl/benchmarks/compute/bench_select.py +++ b/python/cuda_cccl/benchmarks/compute/bench_select.py @@ -30,8 +30,7 @@ def select_iterator(size, d_in, out, num_selected, build_only): def less_than_50(x): return x < 50 - selector = cuda.compute.make_select( - d_in_iter, out, num_selected, less_than_50) + selector = cuda.compute.make_select(d_in_iter, out, num_selected, less_than_50) if not build_only: temp_bytes = selector(None, d_in_iter, out, num_selected, size) temp_storage = cp.empty(temp_bytes, dtype=np.uint8) @@ -52,8 +51,7 @@ def select_struct(inp, out, num_selected, build_only): def in_first_quadrant(p: Point) -> np.uint8: return (p.x > 50) and (p.y > 50) - selector = cuda.compute.make_select( - inp, out, num_selected, in_first_quadrant) + selector = cuda.compute.make_select(inp, out, num_selected, in_first_quadrant) if not build_only: temp_bytes = selector(None, inp, out, num_selected, size) temp_storage = cp.empty(temp_bytes, dtype=np.uint8) @@ -68,8 +66,7 @@ def select_stateful(inp, out, num_selected, threshold_state, build_only): def threshold_select(x): return x > threshold_state[0] - selector = cuda.compute.make_select( - inp, out, num_selected, threshold_select) + selector = cuda.compute.make_select(inp, out, num_selected, threshold_select) if not build_only: temp_bytes = selector(None, inp, out, num_selected, size) temp_storage = cp.empty(temp_bytes, dtype=np.uint8) @@ -87,8 +84,7 @@ def bench_select_pointer(bench_fixture, request, size): def run(): select_pointer( - inp, out, num_selected, build_only=( - bench_fixture == "compile_benchmark") + inp, out, num_selected, build_only=(bench_fixture == "compile_benchmark") ) fixture = request.getfixturevalue(bench_fixture) @@ -118,15 +114,13 @@ def run(): @pytest.mark.parametrize("bench_fixture", ["compile_benchmark", "benchmark"]) def bench_select_struct(bench_fixture, request, size): actual_size = 100 if bench_fixture == "compile_benchmark" else size - inp = cp.random.randint(0, 100, (actual_size, 2), - dtype=np.int32).view(Point.dtype) + inp = cp.random.randint(0, 100, (actual_size, 2), dtype=np.int32).view(Point.dtype) out = cp.empty_like(inp) num_selected = cp.empty(2, dtype=np.uint64) def run(): select_struct( - inp, out, num_selected, build_only=( - bench_fixture == "compile_benchmark") + inp, out, num_selected, build_only=(bench_fixture == "compile_benchmark") ) fixture = request.getfixturevalue(bench_fixture) From 3f13f4274c5d0fb2a823fc4c662e257288192d86 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Fri, 19 Dec 2025 10:01:20 +0000 Subject: [PATCH 9/9] Remove duplicate cache registry --- python/cuda_cccl/cuda/compute/_caching.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/python/cuda_cccl/cuda/compute/_caching.py b/python/cuda_cccl/cuda/compute/_caching.py index 2c08090d65f..7c6bcd0ec02 100644 --- a/python/cuda_cccl/cuda/compute/_caching.py +++ b/python/cuda_cccl/cuda/compute/_caching.py @@ -10,8 +10,6 @@ except ImportError: from cuda.core.experimental import Device -# Central registry of all algorithm caches -_cache_registry: dict[str, object] = {} # Central registry of all algorithm caches _cache_registry: dict[str, object] = {}