diff --git a/python/cuda_cccl/benchmarks/compute/bench_select.py b/python/cuda_cccl/benchmarks/compute/bench_select.py index 7f3f38609f6..8d284c63a74 100644 --- a/python/cuda_cccl/benchmarks/compute/bench_select.py +++ b/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) diff --git a/python/cuda_cccl/cuda/compute/_caching.py b/python/cuda_cccl/cuda/compute/_caching.py index ad675c05341..7c6bcd0ec02 100644 --- a/python/cuda_cccl/cuda/compute/_caching.py +++ b/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,8 +111,6 @@ class CachableFunction: """ def __init__(self, func): - import numba.cuda.dispatcher - self._func = func closure = func.__closure__ if func.__closure__ is not None else [] @@ -92,16 +118,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): 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. 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 b6b8f311f53..e0c2e8f2c14 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_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 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