dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

feat(cccl): integrate missing CCCL directories — python/, ci/, .agent/, docs/, test/

Browse Source 
Sparse-checkout from NVIDIA/cccl main branch to complete cccl_upstream:

Added:
- python/cuda_cccl/ (226 files) — Python bindings for device-level algorithms
  Critical for muh toolchain: cuda.compute.reduce_into, scan, radix_sort, etc.
  Includes 204 .py files with full test coverage for all 27 algorithms
- ci/ (163 files) — Build/test infrastructure
  build_cub.sh, test_cub.sh, build_and_test_targets.sh, matrix.yaml
  Directly maps to our [INFRA-CI] and [INFRA-BUILD] items
- .agent/skills/ (7 files) — NVIDIA's own agent skills for CCCL
  cccl-style/SKILL.md, cccl-test/SKILL.md, sass-diff/SKILL.md
- docs/ (491 files) — Official CCCL documentation
  CI references, CMake guides, Python compute docs, libcudacxx PTX docs
- test/ (12 files) — Top-level integration tests (cuda_smoke, stdpar)
- Root configs: .clang-format, .clang-tidy, CONTRIBUTING.md, pyproject.toml
- CLAUDE.md symlink → AGENTS.md (NVIDIA's standard)

cccl_upstream now mirrors full NVIDIA/cccl structure:
  Before: 42M (cub + thrust + libcudacxx + cudax + c + examples + benchmarks)
  After:  53M (+python +ci +docs +.agent +test +configs)

This completes the CCCL base needed for:
- [muh-bench] items: ci/util/build_and_test_targets.sh for targeted builds
- [CCCL-verify] items: python/cuda_cccl/tests/ as reference implementations
- [CCCL-test] items: ci/test_cub.sh, ci/test_thrust.sh
- Agent workflow: .agent/skills/ for consistent style and test patterns
This commit is contained in:
muh-bot
2026-08-07 02:34:33 +00:00
parent 3f97dca7ad
commit 2a7ca101d7
908 changed files with 121615 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
cccl_upstream/python/cuda_cccl/tests/compute/examples/binary_search/__init__.py Normal file
View File

@@ -0,0 +1,3 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

28
cccl_upstream/python/cuda_cccl/tests/compute/examples/binary_search/lower_bound_basic.py Normal file
View File

@@ -0,0 +1,28 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
import cupy as cp
import numpy as np
import cuda.compute
h_data = np.array([1, 3, 3, 5, 7, 9], dtype=np.int32)
h_values = np.array([0, 3, 4, 10], dtype=np.int32)
d_data = cp.asarray(h_data)
d_values = cp.asarray(h_values)
d_out = cp.empty(len(h_values), dtype=np.uintp)
cuda.compute.lower_bound(
d_data=d_data,
num_items=len(d_data),
d_values=d_values,
num_values=len(d_values),
d_out=d_out,
)
expected = np.searchsorted(h_data, h_values, side="left").astype(np.uintp)
got = cp.asnumpy(d_out)
assert np.array_equal(got, expected)

30
cccl_upstream/python/cuda_cccl/tests/compute/examples/binary_search/lower_bound_object.py Normal file
View File

@@ -0,0 +1,30 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
import cupy as cp
import numpy as np
import cuda.compute
h_data = np.array([1, 3, 3, 5, 7, 9], dtype=np.int32)
h_values = np.array([0, 3, 4, 10], dtype=np.int32)
d_data = cp.asarray(h_data)
d_values = cp.asarray(h_values)
d_out = cp.empty(len(h_values), dtype=np.uintp)
searcher = cuda.compute.make_lower_bound(d_data=d_data, d_values=d_values, d_out=d_out)
searcher(
d_data=d_data,
num_items=len(d_data),
d_values=d_values,
num_values=len(d_values),
d_out=d_out,
comp=None,
)
expected = np.searchsorted(h_data, h_values, side="left").astype(np.uintp)
got = cp.asnumpy(d_out)
assert np.array_equal(got, expected)

28
cccl_upstream/python/cuda_cccl/tests/compute/examples/binary_search/upper_bound_basic.py Normal file
View File

@@ -0,0 +1,28 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
import cupy as cp
import numpy as np
import cuda.compute
h_data = np.array([1, 3, 3, 5, 7, 9], dtype=np.int32)
h_values = np.array([0, 3, 4, 10], dtype=np.int32)
d_data = cp.asarray(h_data)
d_values = cp.asarray(h_values)
d_out = cp.empty(len(h_values), dtype=np.uintp)
cuda.compute.upper_bound(
d_data=d_data,
num_items=len(d_data),
d_values=d_values,
num_values=len(d_values),
d_out=d_out,
)
expected = np.searchsorted(h_data, h_values, side="right").astype(np.uintp)
got = cp.asnumpy(d_out)
assert np.array_equal(got, expected)

30
cccl_upstream/python/cuda_cccl/tests/compute/examples/binary_search/upper_bound_object.py Normal file
View File

@@ -0,0 +1,30 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
import cupy as cp
import numpy as np
import cuda.compute
h_data = np.array([1, 3, 3, 5, 7, 9], dtype=np.int32)
h_values = np.array([0, 3, 4, 10], dtype=np.int32)
d_data = cp.asarray(h_data)
d_values = cp.asarray(h_values)
d_out = cp.empty(len(h_values), dtype=np.uintp)
searcher = cuda.compute.make_upper_bound(d_data=d_data, d_values=d_values, d_out=d_out)
searcher(
d_data=d_data,
num_items=len(d_data),
d_values=d_values,
num_values=len(d_values),
d_out=d_out,
comp=None,
)
expected = np.searchsorted(h_data, h_values, side="right").astype(np.uintp)
got = cp.asnumpy(d_out)
assert np.array_equal(got, expected)

3
cccl_upstream/python/cuda_cccl/tests/compute/examples/free_threading/__init__.py Normal file
View File

@@ -0,0 +1,3 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

51
cccl_upstream/python/cuda_cccl/tests/compute/examples/free_threading/direct_api.py Normal file
View File

@@ -0,0 +1,51 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Run independent direct API calls from multiple Python threads.
The direct algorithm APIs (``cuda.compute.reduce_into`` and friends) are safe
to call concurrently from any thread: each calling thread transparently
receives its own per-thread algorithm object, while the expensive compiled
build result is shared across threads.
"""
from concurrent.futures import ThreadPoolExecutor
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import OpKind
def reduce_values(h_input):
dtype = np.int32
h_init = np.array([0], dtype=dtype)
d_input = cp.asarray(h_input, dtype=dtype)
d_output = cp.empty(1, dtype=dtype)
cuda.compute.reduce_into(
d_in=d_input,
d_out=d_output,
num_items=len(h_input),
op=OpKind.PLUS,
h_init=h_init,
)
return int(d_output.get()[0])
inputs = [
np.array([1, 2, 3, 4], dtype=np.int32),
np.array([5, 6, 7, 8], dtype=np.int32),
]
with ThreadPoolExecutor(max_workers=len(inputs)) as executor:
results = list(executor.map(reduce_values, inputs))
expected = [int(np.sum(h_input)) for h_input in inputs]
assert results == expected
print(f"Free-threaded direct API results: {results}")

69
cccl_upstream/python/cuda_cccl/tests/compute/examples/free_threading/object_api.py Normal file
View File

@@ -0,0 +1,69 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Run independent object-based API calls from multiple Python threads.
This demonstrates the supported object-API pattern for concurrency: each
thread calls the ``make_*`` factory itself and uses the returned object only
on that thread. Algorithm objects must be used by one thread at a time, so a
per-thread factory call (cheap: objects are cached per thread and the compiled
build result is shared process-wide) is the way to use them concurrently.
"""
from concurrent.futures import ThreadPoolExecutor
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import OpKind
def reduce_values(h_input):
dtype = np.int32
h_init = np.array([0], dtype=dtype)
d_input = cp.asarray(h_input, dtype=dtype)
d_output = cp.empty(1, dtype=dtype)
reducer = cuda.compute.make_reduce_into(
d_in=d_input,
d_out=d_output,
op=OpKind.PLUS,
h_init=h_init,
)
temp_storage_size = reducer(
temp_storage=None,
d_in=d_input,
d_out=d_output,
num_items=len(h_input),
op=OpKind.PLUS,
h_init=h_init,
)
d_temp_storage = cp.empty(temp_storage_size, dtype=np.uint8)
reducer(
temp_storage=d_temp_storage,
d_in=d_input,
d_out=d_output,
num_items=len(h_input),
op=OpKind.PLUS,
h_init=h_init,
)
return int(d_output.get()[0])
inputs = [
np.array([1, 2, 3, 4], dtype=np.int32),
np.array([5, 6, 7, 8], dtype=np.int32),
]
with ThreadPoolExecutor(max_workers=len(inputs)) as executor:
results = list(executor.map(reduce_values, inputs))
expected = [int(np.sum(h_input)) for h_input in inputs]
assert results == expected
print(f"Free-threaded object API results: {results}")

5
cccl_upstream/python/cuda_cccl/tests/compute/examples/histogram/__init__.py Normal file
View File

@@ -0,0 +1,5 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Parallel histogram algorithms examples package."""

44
cccl_upstream/python/cuda_cccl/tests/compute/examples/histogram/histogram_even_basic.py Normal file
View File

@@ -0,0 +1,44 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use histogram_even to bin a sequence of samples.
"""
import cupy as cp
import numpy as np
import cuda.compute
# Prepare the input and output arrays.
num_samples = 10
h_samples = np.array(
[2.2, 6.1, 7.1, 2.9, 3.5, 0.3, 2.9, 2.1, 6.1, 999.5], dtype="float32"
)
d_samples = cp.asarray(h_samples)
num_levels = 7
d_histogram = cp.empty(num_levels - 1, dtype="int32")
lower_level = np.float32(0)
upper_level = np.float32(12)
# Perform the histogram even.
cuda.compute.histogram_even(
d_samples=d_samples,
d_histogram=d_histogram,
num_output_levels=num_levels,
lower_level=lower_level,
upper_level=upper_level,
num_samples=num_samples,
)
# Verify the result.
h_actual_histogram = cp.asnumpy(d_histogram)
h_expected_histogram, _ = np.histogram(
h_samples, bins=num_levels - 1, range=(lower_level, upper_level)
)
h_expected_histogram = h_expected_histogram.astype("int32")
np.testing.assert_array_equal(h_actual_histogram, h_expected_histogram)
print(f"Histogram even basic result: {h_actual_histogram}")

71
cccl_upstream/python/cuda_cccl/tests/compute/examples/histogram/histogram_object.py Normal file
View File

@@ -0,0 +1,71 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use histogram object API to bin a sequence of samples.
"""
import cupy as cp
import numpy as np
import cuda.compute
# Prepare the input and output arrays.
h_samples = np.array(
[1.5, 2.3, 4.7, 6.2, 7.8, 3.1, 5.5, 8.9, 2.7, 6.4], dtype="float32"
)
d_samples = cp.asarray(h_samples)
num_levels = 6
# note that the object API requires passing numpy arrays
# rather than scalars:
h_num_output_levels = np.array([num_levels], dtype=np.int32)
h_lower_level = np.array([0.0], dtype=np.float32)
h_upper_level = np.array([10.0], dtype=np.float32)
d_histogram = cp.zeros(num_levels - 1, dtype="int32")
# Create the histogram object.
histogrammer = cuda.compute.make_histogram_even(
d_samples=d_samples,
d_histogram=d_histogram,
h_num_output_levels=h_num_output_levels,
h_lower_level=h_lower_level,
h_upper_level=h_upper_level,
num_samples=len(h_samples),
)
# Get the temporary storage size.
temp_storage_size = histogrammer(
temp_storage=None,
d_samples=d_samples,
d_histogram=d_histogram,
h_num_output_levels=h_num_output_levels,
h_lower_level=h_lower_level,
h_upper_level=h_upper_level,
num_samples=len(h_samples),
)
# Allocate the temporary storage.
d_temp_storage = cp.empty(temp_storage_size, dtype=np.uint8)
# Perform the histogram.
histogrammer(
temp_storage=d_temp_storage,
d_samples=d_samples,
d_histogram=d_histogram,
h_num_output_levels=h_num_output_levels,
h_lower_level=h_lower_level,
h_upper_level=h_upper_level,
num_samples=len(h_samples),
)
# Verify the result.
h_result = cp.asnumpy(d_histogram)
expected_histogram = np.array([1, 3, 2, 3, 1], dtype="int32")
np.testing.assert_array_equal(h_result, expected_histogram)
print("Histogram object example completed successfully")

5
cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/__init__.py Normal file
View File

@@ -0,0 +1,5 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Iterator examples for reduction operations."""

40
cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/cache_modified_iterator_basic.py Normal file
View File

@@ -0,0 +1,40 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use cache_modified_iterator.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
CacheModifiedInputIterator,
OpKind,
)
# Prepare the input array.
h_input = np.array([1, 2, 3, 4, 5], dtype=np.int32)
d_input = cp.asarray(h_input)
# Create the cache modified iterator.
cache_it = CacheModifiedInputIterator(d_input, "stream")
# Prepare the initial value for the reduction.
h_init = np.array([0], dtype=np.int32)
# Prepare the output array.
d_output = cp.empty(1, dtype=np.int32)
# Perform the reduction.
cuda.compute.reduce_into(
d_in=cache_it, d_out=d_output, num_items=len(d_input), op=OpKind.PLUS, h_init=h_init
)
# Verify the result.
expected_output = sum(h_input) # 1 + 2 + 3 + 4 + 5 = 15
assert (d_output == expected_output).all()
print(f"Cache modified iterator result: {d_output[0]} (expected: {expected_output})")

40
cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/constant_iterator_basic.py Normal file
View File

@@ -0,0 +1,40 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use constant_iterator.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
ConstantIterator,
OpKind,
)
# Prepare the input and output arrays.
constant_value = 42
num_items = 5
# Create the constant iterator.
constant_it = ConstantIterator(np.int32(constant_value))
# Prepare the initial value for the reduction.
h_init = np.array([0], dtype=np.int32)
# Prepare the output array.
d_output = cp.empty(1, dtype=np.int32)
# Perform the reduction.
cuda.compute.reduce_into(
d_in=constant_it, d_out=d_output, num_items=num_items, op=OpKind.PLUS, h_init=h_init
)
# Verify the result.
expected_output = constant_value * num_items
assert (d_output == expected_output).all()
print(f"Constant iterator result: {d_output[0]} (expected: {expected_output})")

44
cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/counting_iterator_basic.py Normal file
View File

@@ -0,0 +1,44 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use counting_iterator.
"""
import functools
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
CountingIterator,
OpKind,
)
# Prepare the input and output arrays.
first_item = 1
num_items = 100
# Create the counting iterator.
first_it = CountingIterator(np.int32(first_item))
# Prepare the initial value for the reduction.
h_init = np.array([0], dtype=np.int32)
# Prepare the output array.
d_output = cp.empty(1, dtype=np.int32)
# Perform the reduction.
cuda.compute.reduce_into(
d_in=first_it, d_out=d_output, num_items=num_items, op=OpKind.PLUS, h_init=h_init
)
# Verify the result.
expected_output = functools.reduce(
lambda a, b: a + b, range(first_item, first_item + num_items)
)
assert (d_output == expected_output).all()
print(f"Counting iterator result: {d_output[0]} (expected: {expected_output})")

47
cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/discard_iterator_basic.py Normal file
View File

@@ -0,0 +1,47 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use DiscardIterator.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
DiscardIterator,
OpKind,
)
# Prepare the input and output arrays.
h_in_keys = np.array([1, 1, 2, 3, 3, 7, 8, 8], dtype="int32")
d_in_keys = cp.asarray(h_in_keys)
d_out_keys = cp.empty_like(d_in_keys)
d_out_num_selected = cp.empty(1, np.int32)
# Prepare the discard iterator for values.
d_in_values = DiscardIterator()
d_out_values = DiscardIterator()
# Perform the unique by key operation.
cuda.compute.unique_by_key(
d_in_keys=d_in_keys,
d_in_items=d_in_values,
d_out_keys=d_out_keys,
d_out_items=d_out_values,
d_out_num_selected=d_out_num_selected,
op=OpKind.EQUAL_TO,
num_items=d_in_keys.size,
)
# Verify the result.
num_selected = cp.asnumpy(d_out_num_selected)[0]
h_out_keys = cp.asnumpy(d_out_keys)[:num_selected]
expected_keys = np.array([1, 2, 3, 7, 8])
assert np.array_equal(h_out_keys, expected_keys)
print(f"Discard iterator result - keys: {h_out_keys}, count: {num_selected}")

37
cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/permutation_iterator_basic.py Normal file
View File

@@ -0,0 +1,37 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Demonstrate reduction with permutation iterator as input.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
OpKind,
PermutationIterator,
)
# Create a permutation iterator which selects values at the given indices:
d_values = cp.asarray([10, 20, 30, 40, 50], dtype=np.int32)
d_indices = cp.asarray([2, 0, 4, 1], dtype=np.int32) # permutation indices
perm_it = PermutationIterator(d_values, d_indices)
# Prepare the initial value and output for the reduction.
h_init = np.array([0], dtype=np.int32)
d_output = cp.empty(1, dtype=np.int32)
# Perform the reduction on the permuted values.
num_items = len(d_indices)
cuda.compute.reduce_into(
d_in=perm_it, d_out=d_output, num_items=num_items, op=OpKind.PLUS, h_init=h_init
)
# Verify the result:
expected_output = d_values[d_indices].sum()
assert d_output[0] == expected_output
print(f"Permutation iterator result: {d_output[0]} (expected: {expected_output})")

54
cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/permutation_iterator_composed.py Normal file
View File

@@ -0,0 +1,54 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Demonstrate composed permutation iterator with transform iterator.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
CountingIterator,
OpKind,
PermutationIterator,
TransformIterator,
)
def square_op(x):
return x * x
# Create a CountingIterator that generates: 0, 1, 2, 3, 4, 5, 6, 7, 8, 9
counting_it = CountingIterator(np.int32(0))
# Wrap it in a TransformIterator to square the values: 0, 1, 4, 9, 16, 25, 36, 49, 64, 81
transform_it = TransformIterator(counting_it, square_op)
# Create indices to permute the squared values
d_indices = cp.asarray([3, 1, 5, 2], dtype=np.int32)
# Create permutation iterator that accesses the squared counting iterator
# This will access: squares[3]=9, squares[1]=1, squares[5]=25, squares[2]=4
perm_it = PermutationIterator(transform_it, d_indices)
# Prepare the initial value and output for the reduction
h_init = np.array([0], dtype=np.int32)
d_output = cp.empty(1, dtype=np.int32)
# Perform the reduction on the composed iterator
num_items = len(d_indices)
cuda.compute.reduce_into(
d_in=perm_it, d_out=d_output, num_items=num_items, op=OpKind.PLUS, h_init=h_init
)
# Verify the result: 9 + 1 + 25 + 4 = 39
expected_output = 9 + 1 + 25 + 4
assert d_output[0] == expected_output
print(
f"Composed permutation iterator result: {d_output[0]} (expected: {expected_output})"
)

49
cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/permutation_iterator_output.py Normal file
View File

@@ -0,0 +1,49 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Demonstrate transform with permutation iterator as output (scatter operation).
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
CountingIterator,
PermutationIterator,
)
def square_op(x):
return x * x
# Prepare the output values array and indices.
d_values = cp.zeros(10, dtype=np.int32) # Output array
d_indices = cp.asarray([9, 3, 7, 1, 5], dtype=np.int32) # Scatter indices
# Create input iterator that generates: 0, 1, 2, 3, 4
input_it = CountingIterator(np.int32(0))
# Create permutation iterator for output (scatter).
# This will write to: values[9], values[3], values[7], values[1], values[5]
perm_it = PermutationIterator(d_values, d_indices)
# Perform the transform, scattering squared values to permuted locations.
num_items = len(d_indices)
cuda.compute.unary_transform(
d_in=input_it, d_out=perm_it, op=square_op, num_items=num_items
)
# Verify the result: values[9]=0, values[3]=1, values[7]=4, values[1]=9, values[5]=16
# Other positions should remain 0
expected = np.zeros(10, dtype=np.int32)
for i, idx in enumerate(d_indices.get()):
expected[idx] = i * i
assert np.array_equal(d_values.get(), expected)
print(f"Permutation output iterator result: {d_values.get()}")
print(f"Expected: {expected}")

46
cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/reverse_input_iterator.py Normal file
View File

@@ -0,0 +1,46 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use reverse_input_iterator.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
OpKind,
ReverseIterator,
)
# Prepare the input and output arrays.
h_input = np.array([1, 2, 3, 4, 5], dtype=np.int32)
d_input = cp.asarray(h_input)
# Create the reverse input iterator.
reverse_it = ReverseIterator(d_input)
d_output = cp.empty(len(d_input), dtype=np.int32)
# Prepare the initial value for the reduction.
h_init = np.array(0, dtype=np.int32)
# Perform the reduction.
cuda.compute.inclusive_scan(
d_in=reverse_it,
d_out=d_output,
op=OpKind.PLUS,
init_value=h_init,
num_items=len(d_input),
)
# Verify the result.
expected_output = np.array([5, 9, 12, 14, 15], dtype=np.int32)
result = d_output.get()
np.testing.assert_array_equal(result, expected_output)
print(f"Original input: {h_input}")
print(f"Reverse scan result: {result}")
print(f"Expected result: {expected_output}")

46
cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/reverse_output_iterator.py Normal file
View File

@@ -0,0 +1,46 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use reverse_output_iterator.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
OpKind,
ReverseIterator,
)
# Prepare the input and output arrays.
h_input = np.array([1, 2, 3, 4, 5], dtype=np.int32)
d_input = cp.asarray(h_input)
# Prepare the output array.
d_output = cp.empty(len(d_input), dtype=np.int32)
h_init = np.array(0, dtype=np.int32)
# Create the reverse output iterator.
reverse_out_it = ReverseIterator(d_output)
# Perform the reduction.
cuda.compute.inclusive_scan(
d_in=d_input,
d_out=reverse_out_it,
op=OpKind.PLUS,
init_value=h_init,
num_items=len(d_input),
)
# Verify the result.
expected_output = np.array([15, 10, 6, 3, 1], dtype=np.int32)
result = d_output.get()
np.testing.assert_array_equal(result, expected_output)
print(f"Original input: {h_input}")
print(f"Reverse output result: {result}")
print(f"Expected result: {expected_output}")

37
cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/shuffle_iterator_basic.py Normal file
View File

@@ -0,0 +1,37 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Using ShuffleIterator to obtain a random permutation of an array
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import PermutationIterator, ShuffleIterator
# Input data and output array:
d_input = cp.asarray([10, 20, 30, 40, 50, 60, 70, 80, 100], dtype=np.int32)
d_output = cp.empty_like(d_input)
num_items = len(d_input)
# Create a shuffle iterator that produces a random permutation of [0, num_items)
shuffle_it = ShuffleIterator(num_items, seed=42)
# Use PermutationIterator to permute the data according to the random indices
perm_it = PermutationIterator(d_input, shuffle_it)
# Use unary_transform to write values into d_output
cuda.compute.unary_transform(
d_in=perm_it, d_out=d_output, op=lambda x: x, num_items=num_items
)
# Verify it is a valid permutation of the input data:
cp.testing.assert_array_equal(cp.sort(d_output), d_input)
# Print the values
print(f"Input data: {d_input}")
print(f"Shuffled data: {d_output}")

36
cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/transform_iterator_basic.py Normal file
View File

@@ -0,0 +1,36 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Using ``reduce_into`` with a ``TransformIterator`` to compute the
sum of squares of a sequence of numbers.
"""
import cupy as cp
import numpy as np
from cuda.compute import (
OpKind,
TransformIterator,
reduce_into,
)
# Prepare the input and output arrays.
d_input = cp.arange(10, dtype=np.int32)
d_output = cp.empty(1, dtype=np.int32)
h_init = np.array([0], dtype=np.int32) # Initial value for the reduction
# Create a TransformIterator to (lazily) apply the square
it_input = TransformIterator(d_input, lambda a: a**2)
# Use `reduce_into` to compute the sum of the squares of the input.
reduce_into(
d_in=it_input, d_out=d_output, num_items=len(d_input), op=OpKind.PLUS, h_init=h_init
)
# Verify the result.
expected_output = cp.sum(d_input**2).get()
assert d_output[0] == expected_output
print(f"Transform iterator result: {d_output[0]} (expected: {expected_output})")

51
cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/transform_iterator_lambda.py Normal file
View File

@@ -0,0 +1,51 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Demonstrate TransformIterator with a lambda function.
This example shows how to use a lambda function with TransformIterator
to apply a transformation on-the-fly during reduction, without needing
to define a named function.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
CountingIterator,
OpKind,
TransformIterator,
)
# Prepare the parameters.
first_item = 1
num_items = 10
# Create a TransformIterator that squares each value from a CountingIterator
# using a lambda function.
transform_it = TransformIterator(
CountingIterator(np.int32(first_item)), lambda x: x * x
)
h_init = np.array([0], dtype=np.int32)
d_output = cp.empty(1, dtype=np.int32)
# Perform the reduction: sum of squares from 1 to 10.
cuda.compute.reduce_into(
d_in=transform_it,
d_out=d_output,
num_items=num_items,
op=OpKind.PLUS,
h_init=h_init,
)
# Verify the result: 1^2 + 2^2 + ... + 10^2 = 385
expected_output = sum(x * x for x in range(first_item, first_item + num_items))
assert d_output[0] == expected_output
print(
f"Sum of squares with lambda TransformIterator: {d_output[0]} (expected: {expected_output})"
)

44
cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/transform_output_iterator.py Normal file
View File

@@ -0,0 +1,44 @@
# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
TransformOutputIterator example demonstrating reduction with transform output iterator.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
OpKind,
TransformOutputIterator,
)
# Create input and output arrays
d_input = cp.array([1, 2, 3, 4, 5.0], dtype=np.float32)
d_output = cp.empty(shape=1, dtype=np.float32)
# Define the transform operation to be applied
# to the result of the sum reduction.
# TransformOutputIterator requires type annotations:
def sqrt(x: np.float32) -> np.float32:
return x**0.5
# Create transform output iterator
d_out_it = TransformOutputIterator(d_output, sqrt)
# Apply a sum reduction into the transform output iterator
cuda.compute.reduce_into(
d_in=d_input,
d_out=d_out_it,
num_items=len(d_input),
op=OpKind.PLUS,
h_init=np.asarray([0], dtype=np.float32),
)
assert cp.allclose(d_output, cp.sqrt(cp.sum(d_input)), atol=1e-6)

58
cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/zip_iterator_counting.py Normal file
View File

@@ -0,0 +1,58 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use zip_iterator with counting iterator to
find the index with maximum value in an array.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
CountingIterator,
ZipIterator,
)
def max_by_value(p1, p2):
"""Reduction operation that returns the pair with the larger value."""
return p1 if p1[1] > p2[1] else p2
# Create the counting iterator.
counting_it = CountingIterator(np.int32(0))
# Prepare the input array.
arr = cp.asarray([0, 1, 2, 4, 7, 3, 5, 6], dtype=np.int32)
# Create the zip iterator.
zip_it = ZipIterator(counting_it, arr)
num_items = 8
# Note: initial value passed as a numpy struct
dtype = np.dtype([("index", np.int32), ("value", np.int32)], align=True)
h_init = np.asarray([(-1, -1)], dtype=dtype)
d_output = cp.empty(1, dtype=dtype)
# Perform the reduction.
cuda.compute.reduce_into(
d_in=zip_it, d_out=d_output, num_items=num_items, op=max_by_value, h_init=h_init
)
result = d_output.get()[0]
expected_index = 4
expected_value = 7
assert result["index"] == expected_index
assert result["value"] == expected_value
print(
f"Zip iterator with counting result: index={result['index']} "
f"(expected: {expected_index}), value={result['value']} (expected: {expected_value})"
)

50
cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/zip_iterator_elementwise.py Normal file
View File

@@ -0,0 +1,50 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use zip_iterator to perform elementwise sum of two arrays.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
ZipIterator,
)
# Prepare the input arrays.
d_input1 = cp.array([1, 2, 3, 4, 5], dtype=np.int32)
d_input2 = cp.array([10, 20, 30, 40, 50], dtype=np.int32)
# Create the zip iterator.
zip_it = ZipIterator(d_input1, d_input2)
# Prepare the output array.
num_items = len(d_input1)
d_output = cp.empty(num_items, dtype=np.int32)
def sum_paired_values(pair):
"""Extract values from the zip iterator pair and sum them."""
return pair[0] + pair[1]
# Perform the unary transform.
cuda.compute.unary_transform(
d_in=zip_it, d_out=d_output, op=sum_paired_values, num_items=num_items
)
# Calculate the expected results.
expected = d_input1.get() + d_input2.get()
result = d_output.get()
# Verify the result.
np.testing.assert_allclose(result, expected)
print(f"Input array 1: {d_input1.get()}")
print(f"Input array 2: {d_input2.get()}")
print(f"Elementwise sum result: {result}")
print(f"Expected result: {expected}")

61
cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/zip_iterator_reduction.py Normal file
View File

@@ -0,0 +1,61 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use zip_iterator to simultaneously perform a reduction
operation on two arrays.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
ZipIterator,
gpu_struct,
)
@gpu_struct
class Pair:
first: np.int32
second: np.float32
def sum_pairs(p1, p2):
"""Reduction operation that adds corresponding elements of pairs."""
return Pair(p1[0] + p2[0], p1[1] + p2[1])
# Prepare the input arrays.
d_input1 = cp.array([1, 2, 3, 4, 5], dtype=np.int32)
d_input2 = cp.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float32)
# Create the zip iterator.
zip_it = ZipIterator(d_input1, d_input2)
# Prepare the initial value for the reduction.
h_init = Pair(0, 0.0)
# Prepare the output array.
d_output = cp.empty(1, dtype=Pair.dtype)
# Perform the reduction.
cuda.compute.reduce_into(
d_in=zip_it, d_out=d_output, num_items=len(d_input1), op=sum_pairs, h_init=h_init
)
# Calculate the expected results.
expected_first = sum(d_input1.get())
expected_second = sum(d_input2.get())
result = d_output.get()[0]
assert result["first"] == expected_first
assert result["second"] == expected_second
print(
f"Zip iterator result: first={result['first']} (expected: {expected_first}), "
f"second={result['second']} (expected: {expected_second})"
)

5
cccl_upstream/python/cuda_cccl/tests/compute/examples/partition/__init__.py Normal file
View File

@@ -0,0 +1,5 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Parallel partitioning algorithms examples package."""

65
cccl_upstream/python/cuda_cccl/tests/compute/examples/partition/three_way_partition_basic.py Normal file
View File

@@ -0,0 +1,65 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use three_way_partition to partition a sequence of integers into three parts.
"""
import cupy as cp
import numpy as np
import cuda.compute
# Prepare the input and output arrays.
dtype = np.int32
h_input = np.array([0, 2, 9, 1, 5, 6, 7, -3, 17, 10], dtype=dtype)
d_input = cp.asarray(h_input)
d_first_part = cp.empty_like(d_input)
d_second_part = cp.empty_like(d_input)
d_unselected = cp.empty_like(d_input)
d_num_selected = cp.empty(2, dtype=np.int64)
def less_than_op(x):
return x < 8 and x >= 0
def greater_than_equal_op(x):
return x >= 8
# Perform the three_way_partition.
cuda.compute.three_way_partition(
d_in=d_input,
d_first_part_out=d_first_part,
d_second_part_out=d_second_part,
d_unselected_out=d_unselected,
d_num_selected_out=d_num_selected,
select_first_part_op=less_than_op,
select_second_part_op=greater_than_equal_op,
num_items=len(h_input),
)
# Verify the result.
expected_first_part = np.array([0, 2, 1, 5, 6, 7], dtype=dtype)
expected_second_part = np.array([9, 17, 10], dtype=dtype)
expected_unselected = np.array([-3], dtype=dtype)
expected_num_selected = np.array([6, 3], dtype=np.int64)
actual_num_selected = d_num_selected.get()
num_selected_first_part = int(actual_num_selected[0])
num_selected_second_part = int(actual_num_selected[1])
actual_first_part = d_first_part.get()[:num_selected_first_part]
actual_second_part = d_second_part.get()[:num_selected_second_part]
actual_unselected = d_unselected.get()[
: d_input.size - num_selected_first_part - num_selected_second_part
]
np.testing.assert_array_equal(actual_first_part, expected_first_part)
np.testing.assert_array_equal(actual_second_part, expected_second_part)
np.testing.assert_array_equal(actual_unselected, expected_unselected)
np.testing.assert_array_equal(actual_num_selected, expected_num_selected)
print("Three way partition basic example completed successfully")

91
cccl_upstream/python/cuda_cccl/tests/compute/examples/partition/three_way_partition_object.py Normal file
View File

@@ -0,0 +1,91 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use three_way_partition with the object API.
"""
import cupy as cp
import numpy as np
import cuda.compute
# Prepare the input and output arrays.
dtype = np.int32
h_input = np.array([0, 2, 9, 1, 5, 6, 7, -3, 17, 10], dtype=dtype)
d_input = cp.asarray(h_input)
d_first_part = cp.empty_like(d_input)
d_second_part = cp.empty_like(d_input)
d_unselected = cp.empty_like(d_input)
d_num_selected = cp.empty(2, dtype=np.int64)
def less_than_op(x):
return x < 8 and x >= 0
def greater_than_equal_op(x):
return x >= 8
# Create the three_way_partition object.
partitioner = cuda.compute.make_three_way_partition(
d_in=d_input,
d_first_part_out=d_first_part,
d_second_part_out=d_second_part,
d_unselected_out=d_unselected,
d_num_selected_out=d_num_selected,
select_first_part_op=less_than_op,
select_second_part_op=greater_than_equal_op,
)
# Get the temporary storage size.
temp_storage_size = partitioner(
temp_storage=None,
d_in=d_input,
d_first_part_out=d_first_part,
d_second_part_out=d_second_part,
d_unselected_out=d_unselected,
d_num_selected_out=d_num_selected,
select_first_part_op=less_than_op,
select_second_part_op=greater_than_equal_op,
num_items=len(h_input),
)
d_temp_storage = cp.empty(temp_storage_size, dtype=np.uint8)
# Perform the three_way_partition.
partitioner(
temp_storage=d_temp_storage,
d_in=d_input,
d_first_part_out=d_first_part,
d_second_part_out=d_second_part,
d_unselected_out=d_unselected,
d_num_selected_out=d_num_selected,
select_first_part_op=less_than_op,
select_second_part_op=greater_than_equal_op,
num_items=len(h_input),
)
# Verify the result.
expected_first_part = np.array([0, 2, 1, 5, 6, 7], dtype=dtype)
expected_second_part = np.array([9, 17, 10], dtype=dtype)
expected_unselected = np.array([-3], dtype=dtype)
expected_num_selected = np.array([6, 3], dtype=np.int64)
actual_num_selected = d_num_selected.get()
num_selected_first_part = int(actual_num_selected[0])
num_selected_second_part = int(actual_num_selected[1])
actual_first_part = d_first_part.get()[:num_selected_first_part]
actual_second_part = d_second_part.get()[:num_selected_second_part]
actual_unselected = d_unselected.get()[
: d_input.size - num_selected_first_part - num_selected_second_part
]
np.testing.assert_array_equal(actual_first_part, expected_first_part)
np.testing.assert_array_equal(actual_second_part, expected_second_part)
np.testing.assert_array_equal(actual_unselected, expected_unselected)
np.testing.assert_array_equal(actual_num_selected, expected_num_selected)
print("Three way partition object example completed successfully")

3
cccl_upstream/python/cuda_cccl/tests/compute/examples/raw_op/__init__.py Normal file
View File

@@ -0,0 +1,3 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

125
cccl_upstream/python/cuda_cccl/tests/compute/examples/raw_op/cpp_stateful.py Normal file
View File

@@ -0,0 +1,125 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Create a stateful custom operator using RawOp.
This example demonstrates how to create a stateful operator that maintains
runtime state (in this case, a counter on the device). The operator selects
even numbers and atomically increments a counter for each selected item.
"""
import struct
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute.op import RawOp
from cuda.core import Device, Program, ProgramOptions
def get_arch():
"""Get the SM architecture string for the current device."""
device = Device()
device.set_current()
cc_major, cc_minor = device.compute_capability
return f"sm_{cc_major}{cc_minor}"
def compile_cpp_to_ltoir(source: str, arch: str) -> bytes:
"""Compile C++ source to LTOIR using cuda.core."""
opts = ProgramOptions(
arch=arch,
relocatable_device_code=True,
link_time_optimization=True,
)
prog = Program(source, "c++", options=opts)
return prog.compile("ltoir").code
# Create a device counter initialized to 0
d_counter = cp.zeros(1, dtype=np.int32)
# Pack the counter pointer as state bytes
counter_ptr = d_counter.__cuda_array_interface__["data"][0]
state_bytes = struct.pack("P", counter_ptr)
state_alignment = np.dtype(np.intp).alignment
# Define a C++ stateful select operator
# The operator selects even numbers and counts them using atomic operations
cpp_source = """
extern "C" __device__ void select_even_with_count(void* state, void* input, void* result) {
// Extract counter pointer from state
int* counter = *reinterpret_cast<int**>(state);
// Get input value
int value = *static_cast<int*>(input);
// Check if even
bool is_even = (value % 2 == 0);
// If selected, atomically increment the counter
if (is_even) {
atomicAdd(counter, 1);
}
// Store result as bool (uint8)
*static_cast<unsigned char*>(result) = is_even ? 1 : 0;
}
"""
# Compile C++ to LTOIR
arch = get_arch()
ltoir_bytes = compile_cpp_to_ltoir(cpp_source, arch)
# Create a stateful RawOp with the state bytes
select_op = RawOp(
ltoir=ltoir_bytes,
name="select_even_with_count",
state=state_bytes,
state_alignment=state_alignment,
)
# Prepare test data: numbers 0 to 19
num_items = 20
h_input = np.arange(num_items, dtype=np.int32)
d_input = cp.array(h_input)
# Allocate output arrays
d_output = cp.empty(num_items, dtype=np.int32)
d_num_selected = cp.empty(1, dtype=np.int32)
# Run select with the stateful operator
cuda.compute.select(
d_in=d_input,
d_out=d_output,
d_num_selected_out=d_num_selected,
cond=select_op,
num_items=num_items,
)
# Get results
num_selected = d_num_selected.get()[0]
counter_value = d_counter.get()[0]
# Verify: should have selected 10 even numbers (0, 2, 4, ..., 18)
expected_count = 10
assert num_selected == expected_count, (
f"Expected {expected_count} selected, got {num_selected}"
)
assert counter_value == expected_count, (
f"Expected counter={expected_count}, got {counter_value}"
)
# Verify the selected values are correct
selected_values = d_output.get()[:num_selected]
expected_selected = np.arange(0, 20, 2, dtype=np.int32)
assert np.array_equal(selected_values, expected_selected), "Selected values don't match"
print(f"Selected {num_selected} even numbers")
print(f"Counter value: {counter_value}")
print(f"Selected values: {selected_values}")
print("RawOp stateful example completed successfully!")

71
cccl_upstream/python/cuda_cccl/tests/compute/examples/raw_op/cpp_stateless.py Normal file
View File

@@ -0,0 +1,71 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Create a custom C++ operator from LTOIR bytecode using RawOp.
This example demonstrates how to compile C++ device code to LTOIR and use it
as a custom operator.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute.op import RawOp
from cuda.core import Device, Program, ProgramOptions
def get_arch():
"""Get the SM architecture string for the current device."""
device = Device()
device.set_current()
cc_major, cc_minor = device.compute_capability
return f"sm_{cc_major}{cc_minor}"
def compile_cpp_to_ltoir(source: str, arch: str) -> bytes:
"""Compile C++ source to LTOIR using cuda.core."""
opts = ProgramOptions(
arch=arch,
relocatable_device_code=True,
link_time_optimization=True,
)
prog = Program(source, "c++", options=opts)
return prog.compile("ltoir").code
# Define a C++ custom multiply operator
cpp_source = """
extern "C" __device__ void multiply_op(void* a, void* b, void* result) {
*static_cast<int*>(result) = *static_cast<int*>(a) * *static_cast<int*>(b);
}
"""
# Compile C++ to LTOIR
arch = get_arch()
ltoir_bytes = compile_cpp_to_ltoir(cpp_source, arch)
# Create a RawOp from the LTOIR bytecode
multiply_op = RawOp(ltoir=ltoir_bytes, name="multiply_op")
# Prepare test data
h_input = np.array([1, 2, 3, 4, 5], dtype=np.int32)
d_input = cp.array(h_input)
d_output = cp.empty(1, dtype=np.int32)
h_init = np.array(1, dtype=np.int32)
# Use the custom operator with reduce_into
cuda.compute.reduce_into(
d_in=d_input, d_out=d_output, num_items=len(d_input), op=multiply_op, h_init=h_init
)
# Verify the result
result = d_output.get()[0]
expected = np.prod(h_input) # 1 * 2 * 3 * 4 * 5 = 120
assert result == expected, f"Expected {expected}, got {result}"
print(f"Custom multiply reduction result: {result}")
print("RawOp stateless example completed successfully!")

85
cccl_upstream/python/cuda_cccl/tests/compute/examples/raw_op/llvm_stateless.py Normal file
View File

@@ -0,0 +1,85 @@
# Copyright (c) 2026 NVIDIA CORPORATION.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#
# ruff: noqa: E402 — the v1-skip block below intentionally precedes the
# example's imports so the imports stay grouped at the start of the example
# body (after `# example-begin`).
# Skip cleanly on v1 — the LLVM-bitcode code-path (DeviceCode(kind="llvm_ir"))
# requires cccl.c.parallel built with CCCL_PYTHON_USE_V2=ON. On a v1 wheel the
# binding silently treats it as LTO-IR and nvJitLink rejects the bitcode bytes
# as malformed LTO-IR.
import sys
try:
from cuda.compute._build_info import USING_V2
except ImportError:
USING_V2 = False
if not USING_V2:
print("llvm_stateless requires cccl.c.parallel v2; skipping.")
sys.exit(0)
# example-begin
"""
Create a custom operator from LLVM bitcode using RawOp.
This example demonstrates how to supply pre-compiled LLVM bitcode to RawOp,
which is the preferred path for cccl.parallel v2 because the bitcode is
linked into the CUB module at the LLVM IR level and the optimizer inlines
the operator through kernel inner loops.
"""
import cupy as cp
import llvmlite.binding as llvm
import numpy as np
import cuda.compute
from cuda.compute._device_code import DeviceCode
from cuda.compute.op import RawOp
# Hand-written LLVM IR for an extern "C" multiply operator with the
# void(void*, void*, void*) ABI RawOp expects. llvmlite parses the text and
# serializes it to LLVM bitcode (the binary form, starting with magic "BC")
# that the v2 backend's LLVM linker accepts.
llvm_ir = """
target triple = "nvptx64-nvidia-cuda"
define void @multiply_op(ptr %a, ptr %b, ptr %result) {
entry:
%x = load i32, ptr %a, align 4
%y = load i32, ptr %b, align 4
%r = mul i32 %x, %y
store i32 %r, ptr %result, align 4
ret void
}
"""
mod = llvm.parse_assembly(llvm_ir)
mod.verify()
bitcode = bytes(mod.as_bitcode())
# Wrap the bitcode in DeviceCode so RawOp knows to treat it as LLVM bitcode
# rather than the default LTO-IR. (Raw `bytes` is accepted too and treated
# as LTO-IR — the legacy form.)
multiply_op = RawOp(
ltoir=DeviceCode(op_bytes=bitcode, kind="llvm_ir"),
name="multiply_op",
)
h_input = np.array([1, 2, 3, 4, 5], dtype=np.int32)
d_input = cp.array(h_input)
d_output = cp.empty(1, dtype=np.int32)
h_init = np.array(1, dtype=np.int32)
cuda.compute.reduce_into(
d_in=d_input, d_out=d_output, num_items=len(d_input), op=multiply_op, h_init=h_init
)
result = d_output.get()[0]
expected = np.prod(h_input) # 1 * 2 * 3 * 4 * 5 = 120
assert result == expected, f"Expected {expected}, got {result}"
print(f"Custom multiply reduction result: {result}")
print("RawOp LLVM-IR example completed successfully!")

5
cccl_upstream/python/cuda_cccl/tests/compute/examples/reduction/__init__.py Normal file
View File

@@ -0,0 +1,5 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Parallel reduction algorithms examples package."""

37
cccl_upstream/python/cuda_cccl/tests/compute/examples/reduction/min_reduction.py Normal file
View File

@@ -0,0 +1,37 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Computing the minimum value of a sequence using `reduce_into`.
"""
import cupy as cp
import numpy as np
import cuda.compute
def min_op(a, b):
# the binary operation for the reduction
return a if a < b else b
# Prepare the input and output arrays.
dtype = np.int32
h_init = np.array([42], dtype=dtype)
d_input = cp.array([8, 6, 7, 5, 3, 0, 9], dtype=dtype)
d_output = cp.empty(1, dtype=dtype)
# Perform the reduction.
cuda.compute.reduce_into(
d_in=d_input, d_out=d_output, num_items=len(d_input), op=min_op, h_init=h_init
)
# Verify the result.
expected_output = 0
result = d_output.get()[0]
assert result == expected_output
print(f"Min reduction result: {result}")

61
cccl_upstream/python/cuda_cccl/tests/compute/examples/reduction/minmax_reduction.py Normal file
View File

@@ -0,0 +1,61 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Simultaneously computing the minimum and maximum values of a sequence using `reduce_into`
with a custom data type.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
TransformIterator,
gpu_struct,
)
# Define a custom data type for the accumulator.
@gpu_struct
class MinMax:
min_val: np.float64
max_val: np.float64
# Define the binary operation for the reduction.
def minmax_op(v1: MinMax, v2: MinMax):
c_min = min(v1.min_val, v2.min_val)
c_max = max(v1.max_val, v2.max_val)
return MinMax(c_min, c_max)
# Define a transform operation to convert a value `x` to MinMax(abs(x), abs(x)).
def transform_op(v):
av = abs(v)
return MinMax(av, av)
# Prepare the input and output data.
nelems = 4096
d_in = cp.random.randn(nelems)
tr_it = TransformIterator(d_in, transform_op)
d_out = cp.empty(tuple(), dtype=MinMax.dtype)
h_init = MinMax(np.inf, -np.inf)
# Perform the reduction.
cuda.compute.reduce_into(
d_in=tr_it, d_out=d_out, num_items=nelems, op=minmax_op, h_init=h_init
)
# Verify the result.
actual = d_out.get()
h = np.abs(d_in.get())
expected = np.asarray([(h.min(), h.max())], dtype=MinMax.dtype)
assert actual == expected
print(f"MinMax reduction result: {actual}")

58
cccl_upstream/python/cuda_cccl/tests/compute/examples/reduction/reduce_object.py Normal file
View File

@@ -0,0 +1,58 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Reduction example using the object API.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
OpKind,
)
# Prepare the input and output arrays.
dtype = np.int32
init_value = 5
h_init = np.array([init_value], dtype=dtype)
h_input = np.array([1, 2, 3, 4], dtype=dtype)
d_input = cp.asarray(h_input)
d_output = cp.empty(1, dtype=dtype)
# Create a reducer object.
reducer = cuda.compute.make_reduce_into(
d_in=d_input, d_out=d_output, op=OpKind.PLUS, h_init=h_init
)
# Get the temporary storage size.
temp_storage_size = reducer(
temp_storage=None,
d_in=d_input,
d_out=d_output,
num_items=len(h_input),
op=OpKind.PLUS,
h_init=h_init,
)
# Allocate temporary storage using any user-defined allocator.
# The result must be an object exposing `__cuda_array_interface__`.
d_temp_storage = cp.empty(temp_storage_size, dtype=np.uint8)
# Perform the reduction.
reducer(
temp_storage=d_temp_storage,
d_in=d_input,
d_out=d_output,
num_items=len(h_input),
op=OpKind.PLUS,
h_init=h_init,
)
expected_result = np.sum(h_input) + init_value
actual_result = d_output.get()[0]
assert actual_result == expected_result
print("Reduce object example completed successfully")

37
cccl_upstream/python/cuda_cccl/tests/compute/examples/reduction/sum_custom_reduction.py Normal file
View File

@@ -0,0 +1,37 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Sum only even values in an array using reduction with custom operation.
"""
import cupy as cp
import numpy as np
import cuda.compute
# Prepare the input and output arrays.
dtype = np.int32
h_init = np.array([0], dtype=dtype)
d_input = cp.array([1, 2, 3, 4, 5], dtype=dtype)
d_output = cp.empty(1, dtype=dtype)
# Define the binary operation for the reduction.
def add_op(a, b):
return (a if a % 2 == 0 else 0) + (b if b % 2 == 0 else 0)
# Perform the reduction.
cuda.compute.reduce_into(
d_in=d_input, d_out=d_output, num_items=len(d_input), op=add_op, h_init=h_init
)
# Verify the result.
expected_output = 6
assert (d_output == expected_output).all()
result = d_output[0]
print(f"Custom sum reduction result: {result}")

31
cccl_upstream/python/cuda_cccl/tests/compute/examples/reduction/sum_reduction.py Normal file
View File

@@ -0,0 +1,31 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Sum all values in an array using reduction with PLUS operation.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import OpKind
# Prepare the input and output arrays.
dtype = np.int32
h_init = np.array([0], dtype=dtype)
d_input = cp.array([1, 2, 3, 4, 5], dtype=dtype)
d_output = cp.empty(1, dtype=dtype)
# Perform the reduction.
cuda.compute.reduce_into(
d_in=d_input, d_out=d_output, num_items=len(d_input), op=OpKind.PLUS, h_init=h_init
)
# Verify the result.
expected_output = 15
assert (d_output == expected_output).all()
result = d_output[0]
print(f"Sum reduction result: {result}")

38
cccl_upstream/python/cuda_cccl/tests/compute/examples/reduction/sum_reduction_lambda.py Normal file
View File

@@ -0,0 +1,38 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Sum all values in an array using reduction with a lambda function.
This example demonstrates that lambda functions can be used directly
as reduction operators, providing a concise alternative to defining
named functions.
"""
import cupy as cp
import numpy as np
import cuda.compute
# Prepare the input and output arrays.
dtype = np.int32
h_init = np.array([0], dtype=dtype)
d_input = cp.array([1, 2, 3, 4, 5], dtype=dtype)
d_output = cp.empty(1, dtype=dtype)
# Perform the reduction using a lambda function.
cuda.compute.reduce_into(
d_in=d_input,
d_out=d_output,
num_items=len(d_input),
op=lambda a, b: a + b,
h_init=h_init,
)
# Verify the result.
expected_output = 15
assert (d_output == expected_output).all()
result = d_output[0]
print(f"Sum reduction with lambda result: {result}")

5
cccl_upstream/python/cuda_cccl/tests/compute/examples/scan/__init__.py Normal file
View File

@@ -0,0 +1,5 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Parallel scan algorithms examples package."""

101
cccl_upstream/python/cuda_cccl/tests/compute/examples/scan/ema_example.py Normal file
View File

@@ -0,0 +1,101 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
CountingIterator,
TransformIterator,
ZipIterator,
gpu_struct,
)
# Given input vector u and smoothing parameter 0<alpha<1
# the exponential moving average sequence ma is defined
# by ma[t] = (1 - alpha) * u[t] + alpha * ma[t-1] with
# initial condition ma[0] = u[0].
#
# This recurrence equation admits a closed form solution:
#
# ma[t] = alpha ** (t + 1) * u[0] + (1 - alpha) * (alpha ** t) *
# sum(u[s] * alpha ** (-s) for s in range(t))
#
# Closed form solution could be computed using inclusive_scan,
# except naive implementation suffers from underflow/overflow
# problem for long sequences.
#
# This implementation solves this problem by representing number
# using (fp_value, int_exponent) to extend representable range.
u = 3.0 + cp.sin(cp.linspace(-3.0, 3.0, num=1024, dtype=cp.double))
u += cp.random.normal(0.0, 0.1, size=u.size)
alpha = 0.05
assert 0.0 < alpha < 1.0
@gpu_struct
class ValueScale:
value: cp.float64
scale: cp.int64
def add_op(v1: ValueScale, v2: ValueScale) -> ValueScale:
if v1.scale > v2.scale:
s = v2.scale
v = v2.value + v1.value * (alpha ** (v1.scale - v2.scale))
else:
s = v1.scale
v = v1.value + v2.value * (alpha ** (v2.scale - v1.scale))
return ValueScale(v, s)
def negative_op(i: cp.int64) -> cp.int64:
return -i
seq_it = CountingIterator(cp.int64(0))
negative_exponents_it = TransformIterator(seq_it, negative_op)
d_inp = ZipIterator(u, negative_exponents_it)
d_cumsum = cp.empty(u.shape, dtype=ValueScale.dtype)
h_init = ValueScale(0.0, 0)
cuda.compute.inclusive_scan(
d_in=d_inp, d_out=d_cumsum, op=add_op, init_value=h_init, num_items=u.size
)
it_seq = CountingIterator(cp.int64(0))
d_ema = cp.empty_like(u)
def combine_op(v: ValueScale, t: cp.int64) -> cp.float64:
return (1 - alpha) * v.value * alpha ** (t + v.scale)
cuda.compute.binary_transform(
d_in1=d_cumsum, d_in2=it_seq, d_out=d_ema, op=combine_op, num_items=u.size
)
d_ema += (alpha ** cp.arange(1, u.size + 1)) * u[0]
def ema_ref(u: np.ndarray, alpha: float):
"Sequential reference implementation of EMA"
a = np.empty_like(u)
a[0] = u[0]
for t in range(1, u.size):
a[t] = (1 - alpha) * u[t] + alpha * a[t - 1]
return a
h_u = u.get()
h_ema = ema_ref(h_u, alpha)
assert np.allclose(h_ema, d_ema.get())
print("Exponential moving average example completed successfully")

37
cccl_upstream/python/cuda_cccl/tests/compute/examples/scan/exclusive_scan_max.py Normal file
View File

@@ -0,0 +1,37 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Exclusive scan using custom maximum operation.
"""
import cupy as cp
import numpy as np
import cuda.compute
# Define the binary operation for the scan.
def max_op(a, b):
return max(a, b)
# Prepare the input and output arrays.
h_init = np.array([1], dtype="int32")
d_input = cp.array([-5, 0, 2, -3, 2, 4, 0, -1, 2, 8], dtype="int32")
d_output = cp.empty_like(d_input, dtype="int32")
# Perform the exclusive scan.
cuda.compute.exclusive_scan(
d_in=d_input, d_out=d_output, op=max_op, init_value=h_init, num_items=d_input.size
)
# Verify the result.
expected = np.asarray([1, 1, 1, 2, 2, 2, 4, 4, 4, 4])
result = d_output.get()
np.testing.assert_equal(result, expected)
print(f"Exclusive scan max result: {result}")

53
cccl_upstream/python/cuda_cccl/tests/compute/examples/scan/exclusive_scan_object.py Normal file
View File

@@ -0,0 +1,53 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Exclusive scan example demonstrating the object API.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
OpKind,
)
# Prepare the input and output arrays.
dtype = np.int32
h_init = np.array([0], dtype=dtype)
h_input = np.array([1, 2, 3, 4], dtype=dtype)
d_input = cp.asarray(h_input)
d_output = cp.empty(len(h_input), dtype=dtype)
# Create the scanner object and allocate temporary storage.
scanner = cuda.compute.make_exclusive_scan(
d_in=d_input, d_out=d_output, op=OpKind.PLUS, init_value=h_init
)
temp_storage_size = scanner(
temp_storage=None,
d_in=d_input,
d_out=d_output,
op=OpKind.PLUS,
num_items=len(h_input),
init_value=h_init,
)
d_temp_storage = cp.empty(temp_storage_size, dtype=np.uint8)
# Perform the exclusive scan.
scanner(
temp_storage=d_temp_storage,
d_in=d_input,
d_out=d_output,
op=OpKind.PLUS,
num_items=len(h_input),
init_value=h_init,
)
# Verify the result.
expected_result = np.array([0, 1, 3, 6], dtype=dtype)
actual_result = d_output.get()
np.testing.assert_array_equal(actual_result, expected_result)
print("Exclusive scan object example completed successfully")

36
cccl_upstream/python/cuda_cccl/tests/compute/examples/scan/inclusive_scan_custom.py Normal file
View File

@@ -0,0 +1,36 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Inclusive scan with custom operation (prefix sum of even values).
"""
import cupy as cp
import numpy as np
import cuda.compute
# Prepare the input and output arrays.
h_init = np.array([0], dtype="int32")
d_input = cp.array([1, 2, 3, 4, 5], dtype="int32")
d_output = cp.empty_like(d_input, dtype="int32")
# Define the binary operation for the scan.
def add_op(a, b):
return (a if a % 2 == 0 else 0) + (b if b % 2 == 0 else 0)
# Perform the inclusive scan.
cuda.compute.inclusive_scan(
d_in=d_input, d_out=d_output, op=add_op, init_value=h_init, num_items=d_input.size
)
# Verify the result.
expected = np.asarray([0, 2, 2, 6, 6])
assert np.array_equal(d_output.get(), expected)
result = d_output.get()
print(f"Inclusive scan custom result: {result}")

53
cccl_upstream/python/cuda_cccl/tests/compute/examples/scan/inclusive_scan_object.py Normal file
View File

@@ -0,0 +1,53 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Inclusive scan example demonstrating the object API.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
OpKind,
)
# Prepare the input and output arrays.
dtype = np.int32
h_init = np.array([0], dtype=dtype)
h_input = np.array([1, 2, 3, 4], dtype=dtype)
d_input = cp.asarray(h_input)
d_output = cp.empty(len(h_input), dtype=dtype)
# Create the scanner object and allocate temporary storage.
scanner = cuda.compute.make_inclusive_scan(
d_in=d_input, d_out=d_output, op=OpKind.PLUS, init_value=h_init
)
temp_storage_size = scanner(
temp_storage=None,
d_in=d_input,
d_out=d_output,
op=OpKind.PLUS,
num_items=len(h_input),
init_value=h_init,
)
d_temp_storage = cp.empty(temp_storage_size, dtype=np.uint8)
# Perform the inclusive scan.
scanner(
temp_storage=d_temp_storage,
d_in=d_input,
d_out=d_output,
op=OpKind.PLUS,
num_items=len(h_input),
init_value=h_init,
)
# Verify the result.
expected_result = np.array([1, 3, 6, 10], dtype=dtype)
actual_result = d_output.get()
np.testing.assert_array_equal(actual_result, expected_result)
print("Inclusive scan object example completed successfully")

79
cccl_upstream/python/cuda_cccl/tests/compute/examples/scan/logcdf_example.py Normal file
View File

@@ -0,0 +1,79 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Given a vector of log-probabilities, compute a vector of logarithms of cumulative density function.
"""
import math
import cupy as cp
import cupyx.scipy.special as cp_special
import numpy as np
import cuda.compute
# Prepare the input data and compute log-probabilities.
# Use log-add-exp binary operation to sidestep flush-to-zero
# and numerical exceptions computing logarithms of that.
n = 500
p = 0.31
m = cp.arange(n + 1, dtype=cp.float64)
nm = n - m
lognorm = (
cp_special.loggamma(1 + n)
- cp_special.loggamma(1 + m)
- cp_special.loggamma(1 + nm)
)
logpdf = lognorm + m * cp.log(p) + nm * cp.log1p(-p)
assert n + 1 == logpdf.size
# Define the binary operations for the scans.
def logaddexp(logp1: cp.float64, logp2: cp.float64):
m_max = max(logp1, logp2)
m_min = min(logp1, logp2)
return m_max + math.log(1.0 + math.exp(m_min - m_max))
def maximum(v1: cp.float64, v2: cp.float64):
return max(v1, v2)
# Prepare the output arrays and initial value.
logcdf = cp.empty_like(logpdf)
h_init = np.array(-np.inf, dtype=np.float64)
logcdf2 = cp.empty_like(logpdf)
# Perform the first inclusive scan (log-add-exp).
cuda.compute.inclusive_scan(
d_in=logpdf, d_out=logcdf, op=logaddexp, init_value=h_init, num_items=logpdf.size
)
# Perform the second inclusive scan (maximum).
cuda.compute.inclusive_scan(
d_in=logcdf, d_out=logcdf2, op=maximum, init_value=h_init, num_items=logpdf.size
)
# Verify the results and compute quantiles.
assert cp.all(logcdf2[:-1] <= logcdf2[1:])
assert float(cp.max(logcdf2)) <= 0.0
q25, q75 = cp.searchsorted(logcdf2, cp.asarray(np.log([0.25, 0.75])))
try:
from scipy.stats.distributions import binom
q25_ref, q75_ref = binom(n, p).isf([0.75, 0.25])
assert q25 == q25_ref
assert q75 == q75_ref
except ImportError:
print("scipy not found, skipping assertions")
print(f"Log CDF example completed. q25: {q25}, q75: {q75}")

50
cccl_upstream/python/cuda_cccl/tests/compute/examples/scan/running_average.py Normal file
View File

@@ -0,0 +1,50 @@
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
ConstantIterator,
TransformOutputIterator,
ZipIterator,
gpu_struct,
)
# example-begin
"""
Inclusive scan using zip iterator and output transform iterator to compute running average.
"""
@gpu_struct
class SumAndCount: # data type to store the running sum and the count
sum: np.float32
count: np.int32
# binary operation for the scan computes the running sum and running count
def add_op(x1: SumAndCount, x2: SumAndCount) -> SumAndCount:
return SumAndCount(x1.sum + x2.sum, x1.count + x2.count)
# output transform operation divides the sum by the count to get the running average
def write_op(x: SumAndCount) -> np.float32:
return x.sum / x.count
# construct a zip iterator to pair the input with the sequence [1, 1, ..., 1]
d_input = cp.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float32)
it_input = ZipIterator(d_input, ConstantIterator(np.int32(1)))
# output transform iterator divides the sum by the count to get the running average
d_output = cp.empty_like(d_input)
it_output = TransformOutputIterator(d_output, write_op)
h_init = SumAndCount(0.0, 0)
cuda.compute.inclusive_scan(
d_in=it_input, d_out=it_output, op=add_op, init_value=h_init, num_items=len(d_input)
)
expected = np.array([1.0, 1.5, 2.0, 2.5, 3.0], dtype=np.float32)
np.testing.assert_allclose(d_output.get(), expected)
# example-end

66
cccl_upstream/python/cuda_cccl/tests/compute/examples/scan/segmented_sum.py Normal file
View File

@@ -0,0 +1,66 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Implement segmented scan using zip iterator and ordinary scan.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
ZipIterator,
gpu_struct,
)
# Prepare the input data and head flags.
# Segmented inclusive sum on array of values and head-flags
# array demarkating locations of start of segments can be implemented
# using ordinary inclusive scan using Schwarz operator acting
# of value-flag pairs. `ZipIterator` can be used to efficiently
# load data from pair of input arrays, instead of copying them
# to array of structs.
#
# For example, for data = [1, 1, 1, 1, 1, 1, 1, 1] with
# 3 segments encoded by head_flags = [0, 0, 1, 0, 0, 1, 1, 0]
# corresponding to segmented data [[1, 1], [1, 1, 1], [1], [1, 1]],
# the expected prefix-sum values are [1, 2, 1, 2, 3, 1, 1, 2]
data = cp.asarray([1, 1, 1, 1, 1, 1, 1, 1], dtype=cp.int64)
hflg = cp.asarray([0, 0, 1, 0, 0, 1, 1, 0], dtype=cp.int32)
# Define the custom data type and binary operation.
@gpu_struct
class ValueFlag:
value: cp.int64
flag: cp.int32
def schwartz_sum(op1: ValueFlag, op2: ValueFlag) -> ValueFlag:
f1: cp.int32 = 1 if op1.flag else 0
f2: cp.int32 = 1 if op2.flag else 0
f: cp.int32 = f1 | f2
v: cp.int64 = op2.value if f2 else op1.value + op2.value
return ValueFlag(v, f)
# Prepare the output array and initial value.
zip_it = ZipIterator(data, hflg)
d_output = cp.empty(data.shape, dtype=ValueFlag.dtype)
h_init = ValueFlag(0, 0)
# Perform the segmented scan.
cuda.compute.inclusive_scan(
d_in=zip_it, d_out=d_output, op=schwartz_sum, init_value=h_init, num_items=data.size
)
# Verify the result.
expected_prefix = np.asarray([1, 2, 1, 2, 3, 1, 1, 2], dtype=np.int64)
result = d_output.get()
assert np.array_equal(result["value"], expected_prefix)
print(f"Segmented sum result: {result}")

5
cccl_upstream/python/cuda_cccl/tests/compute/examples/segmented/__init__.py Normal file
View File

@@ -0,0 +1,5 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Parallel segmented algorithms examples package."""

58
cccl_upstream/python/cuda_cccl/tests/compute/examples/segmented/segmented_reduce_basic.py Normal file
View File

@@ -0,0 +1,58 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use segmented_reduce to find the minimum in each segment.
"""
import cupy as cp
import numpy as np
import cuda.compute
def min_op(a, b):
return a if a < b else b
dtype = np.dtype(np.int32)
max_val = np.iinfo(dtype).max
h_init = np.asarray(max_val, dtype=dtype)
# Prepare the offsets.
offsets = cp.array([0, 7, 11, 16], dtype=np.int64)
first_segment = (8, 6, 7, 5, 3, 0, 9)
second_segment = (-4, 3, 0, 1)
third_segment = (3, 1, 11, 25, 8)
# Prepare the input array.
d_input = cp.array(
[*first_segment, *second_segment, *third_segment],
dtype=dtype,
)
# Prepare the start and end offsets.
start_o = offsets[:-1]
end_o = offsets[1:]
# Prepare the output array.
n_segments = start_o.size
d_output = cp.empty(n_segments, dtype=dtype)
# Perform the segmented reduce.
cuda.compute.segmented_reduce(
d_in=d_input,
d_out=d_output,
num_segments=n_segments,
start_offsets_in=start_o,
end_offsets_in=end_o,
op=min_op,
h_init=h_init,
)
# Verify the result.
expected_output = cp.asarray([0, -4, 1], dtype=d_output.dtype)
assert (d_output == expected_output).all()
print(f"Segmented reduce basic result: {d_output.get()}")

69
cccl_upstream/python/cuda_cccl/tests/compute/examples/segmented/segmented_reduce_object.py Normal file
View File

@@ -0,0 +1,69 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Segmented reduction using the object API.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
OpKind,
)
# Prepare the input and output arrays.
dtype = np.int32
h_init = np.array([0], dtype=dtype)
h_input = np.array([1, 2, 3, 4, 5, 6], dtype=dtype)
d_input = cp.asarray(h_input)
d_output = cp.empty(2, dtype=dtype)
start_offsets = cp.array([0, 3], dtype=np.int64)
end_offsets = cp.array([3, 6], dtype=np.int64)
# Create the segmented reduce object.
reducer = cuda.compute.make_segmented_reduce(
d_in=d_input,
d_out=d_output,
start_offsets_in=start_offsets,
end_offsets_in=end_offsets,
op=OpKind.PLUS,
h_init=h_init,
)
# Get the temporary storage size.
temp_storage_size = reducer(
temp_storage=None,
d_in=d_input,
d_out=d_output,
num_segments=2,
start_offsets_in=start_offsets,
end_offsets_in=end_offsets,
op=OpKind.PLUS,
h_init=h_init,
)
# Allocate the temporary storage.
d_temp_storage = cp.empty(temp_storage_size, dtype=np.uint8)
# Perform the segmented reduce.
reducer(
temp_storage=d_temp_storage,
d_in=d_input,
d_out=d_output,
num_segments=2,
start_offsets_in=start_offsets,
end_offsets_in=end_offsets,
op=OpKind.PLUS,
h_init=h_init,
)
# Verify the result.
expected_result = np.array([6, 15], dtype=dtype)
actual_result = d_output.get()
np.testing.assert_array_equal(actual_result, expected_result)
print("Segmented reduce object example completed successfully")

4
cccl_upstream/python/cuda_cccl/tests/compute/examples/select/__init__.py Normal file
View File

@@ -0,0 +1,4 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

39
cccl_upstream/python/cuda_cccl/tests/compute/examples/select/select_basic.py Normal file
View File

@@ -0,0 +1,39 @@
# 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 cuda.compute.algorithms import select
# Create input data
d_in = cp.array([1, 2, 3, 4, 5, 6, 7, 8], dtype=cp.int32)
d_out = cp.empty_like(d_in)
d_num_selected = cp.zeros(2, dtype=cp.uint64)
# Define select condition (keep even numbers)
def is_even(x):
return x % 2 == 0
# Execute select
select(
d_in=d_in,
d_out=d_out,
d_num_selected_out=d_num_selected,
cond=is_even,
num_items=len(d_in),
)
# Get results
num_selected = int(d_num_selected[0])
result = d_out[:num_selected].get()
print(f"Selected {num_selected} items: {result}")
# Output: Selected 4 items: [2 4 6 8]
# example-end
assert num_selected == 4
assert (result == [2, 4, 6, 8]).all()

78
cccl_upstream/python/cuda_cccl/tests/compute/examples/select/select_object.py Normal file
View File

@@ -0,0 +1,78 @@
# 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 cuda.compute.algorithms import make_select
# Create input data
d_in = cp.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], dtype=cp.int32)
d_out = cp.empty_like(d_in)
d_num_selected = cp.zeros(2, dtype=cp.uint64)
# Define select condition (keep values > 5)
def greater_than_5(x):
return x > 5
# Create select object (can be reused)
selector = make_select(
d_in=d_in,
d_out=d_out,
d_num_selected_out=d_num_selected,
cond=greater_than_5,
)
# Get required temp storage
temp_storage_bytes = selector(
temp_storage=None,
d_in=d_in,
d_out=d_out,
d_num_selected_out=d_num_selected,
cond=greater_than_5,
num_items=len(d_in),
)
d_temp_storage = cp.empty(temp_storage_bytes, dtype=cp.uint8)
# Execute select
selector(
temp_storage=d_temp_storage,
d_in=d_in,
d_out=d_out,
d_num_selected_out=d_num_selected,
cond=greater_than_5,
num_items=len(d_in),
)
# Get results
num_selected = int(d_num_selected[0])
result = d_out[:num_selected].get()
print(f"Selected {num_selected} items: {result}")
# Output: Selected 5 items: [ 6 7 8 9 10]
# Reuse the same select object with different input
d_in2 = cp.array([10, 20, 3, 15, 2, 8, 30], dtype=cp.int32)
d_out2 = cp.empty_like(d_in2)
d_num_selected2 = cp.zeros(2, dtype=cp.uint64)
selector(
temp_storage=d_temp_storage,
d_in=d_in2,
d_out=d_out2,
d_num_selected_out=d_num_selected2,
cond=greater_than_5,
num_items=len(d_in2),
)
num_selected2 = int(d_num_selected2[0])
result2 = d_out2[:num_selected2].get()
print(f"Second select: {num_selected2} items: {result2}")
# Output: Second select: 5 items: [10 20 15 8 30]
# example-end
assert num_selected == 5
assert (result == [6, 7, 8, 9, 10]).all()

48
cccl_upstream/python/cuda_cccl/tests/compute/examples/select/select_with_iterator.py Normal file
View File

@@ -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 cuda.compute.algorithms import select
from cuda.compute.iterators import TransformIterator
# Create input data
d_in = cp.array([1, 2, 3, 4, 5, 6, 7, 8], dtype=cp.int32)
d_out = cp.empty_like(d_in)
d_num_selected = cp.zeros(2, dtype=cp.uint64)
# Create iterator that squares each value
def square(x):
return x * x
squared_iter = TransformIterator(d_in, square)
# Select squared values that are greater than 20
def greater_than_20(x):
return x > 20
select(
d_in=squared_iter,
d_out=d_out,
d_num_selected_out=d_num_selected,
cond=greater_than_20,
num_items=len(d_in),
)
# Get results
num_selected = int(d_num_selected[0])
result = d_out[:num_selected].get()
print(f"Selected {num_selected} items: {result}")
# Output: Selected 4 items: [25 36 49 64]
# (5^2=25, 6^2=36, 7^2=49, 8^2=64, all > 20)
# example-end
assert num_selected == 4
assert (result == [25, 36, 49, 64]).all()

54
cccl_upstream/python/cuda_cccl/tests/compute/examples/select/select_with_side_effect.py Normal file
View File

@@ -0,0 +1,54 @@
# 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_in,
d_out=d_out,
d_num_selected_out=d_num_selected,
cond=count_rejects,
num_items=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

5
cccl_upstream/python/cuda_cccl/tests/compute/examples/serialization/__init__.py Normal file
View File

@@ -0,0 +1,5 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Serialization and ahead-of-time compilation examples package."""

50
cccl_upstream/python/cuda_cccl/tests/compute/examples/serialization/ahead_of_time_compilation.py Normal file
View File

@@ -0,0 +1,50 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#
# ruff: noqa: E402 — the v2-skip block below intentionally precedes the
# example's imports so the imports stay grouped at the start of the example
# body (after `# example-begin`).
# Ahead-of-time compilation is only supported on the default (v1) backend; the
# HostJIT (v2) backend raises NotImplementedError. Skip cleanly there so the
# example runner treats it as a pass.
import sys
try:
from cuda.compute._build_info import USING_V2
except ImportError:
USING_V2 = False
if USING_V2:
print("ahead-of-time build is unsupported on the HostJIT (v2) backend; skipping.")
sys.exit(0)
# example-begin
"""
Compile a reduction ahead of time for multiple GPU architectures without a GPU
present, using dtype-only placeholders, and serialize the result for deployment.
"""
import numpy as np
from cuda.compute import OpKind, ProxyArray, ProxyValue, make_reduce_into, serialize
# ProxyArray / ProxyValue describe only the dtype of each argument; they hold no
# GPU memory, so no device is required to build. compute_capability must be given
# explicitly, since there is no device whose architecture we could default to.
reducer = make_reduce_into(
d_in=ProxyArray(np.int32),
d_out=ProxyArray(np.int32),
op=OpKind.PLUS,
h_init=ProxyValue(np.int32),
compute_capability=[80, 90], # build for sm_80 and sm_90
)
# Serialize the multi-architecture build for shipping to deployment targets. On
# a target GPU, cuda.compute.deserialize(blob) reconstructs the object, and the
# build result matching the running architecture is loaded on the first call.
blob = serialize(reducer)
assert len(blob) > 0
print(f"Compiled ahead of time for sm_80 and sm_90; serialized {len(blob)} bytes")

74
cccl_upstream/python/cuda_cccl/tests/compute/examples/serialization/serialize_roundtrip.py Normal file
View File

@@ -0,0 +1,74 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#
# ruff: noqa: E402 — the v2-skip block below intentionally precedes the
# example's imports so the imports stay grouped at the start of the example
# body (after `# example-begin`).
# Serialization is only supported on the default (v1) backend; the HostJIT (v2)
# backend raises NotImplementedError. Skip cleanly there so the example runner
# treats it as a pass.
import sys
try:
from cuda.compute._build_info import USING_V2
except ImportError:
USING_V2 = False
if USING_V2:
print("serialization is unsupported on the HostJIT (v2) backend; skipping.")
sys.exit(0)
# example-begin
"""
Serialize a built reduction to bytes, then reconstruct and run it without
recompiling.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import OpKind
# Build a reduction object for the current device, as usual.
dtype = np.int32
h_init = np.array([0], dtype=dtype)
d_input = cp.array([1, 2, 3, 4, 5], dtype=dtype)
d_output = cp.empty(1, dtype=dtype)
reducer = cuda.compute.make_reduce_into(
d_in=d_input, d_out=d_output, op=OpKind.PLUS, h_init=h_init
)
# Serialize the compiled build result to a blob of bytes. In practice you would
# write this to a file and load it in a later run or on another machine; here we
# keep it in memory to stay self-contained.
blob = cuda.compute.serialize(reducer)
# Reconstruct the reduction from the blob. This performs no JIT compilation.
restored = cuda.compute.deserialize(blob)
# Invoke the restored object exactly as if it had just been built.
temp_storage_size = restored(
temp_storage=None,
d_in=d_input,
d_out=d_output,
num_items=len(d_input),
op=OpKind.PLUS,
h_init=h_init,
)
d_temp_storage = cp.empty(temp_storage_size, dtype=np.uint8)
restored(
temp_storage=d_temp_storage,
d_in=d_input,
d_out=d_output,
num_items=len(d_input),
op=OpKind.PLUS,
h_init=h_init,
)
# The reconstructed object produces the same result as the original.
assert d_output.get()[0] == 15
print("Serialize/deserialize round-trip result:", d_output.get()[0])

5
cccl_upstream/python/cuda_cccl/tests/compute/examples/sort/__init__.py Normal file
View File

@@ -0,0 +1,5 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Parallel sorting algorithms examples package."""

47
cccl_upstream/python/cuda_cccl/tests/compute/examples/sort/merge_sort_basic.py Normal file
View File

@@ -0,0 +1,47 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Demonstrate basic merge sort with keys and values.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
OpKind,
)
# Prepare the input and output arrays.
h_in_keys = np.array([-5, 0, 2, -3, 2, 4, 0, -1, 2, 8], dtype="int32")
h_in_values = np.array(
[-3.2, 2.2, 1.9, 4.0, -3.9, 2.7, 0, 8.3 - 1, 2.9, 5.4], dtype="float32"
)
d_in_keys = cp.asarray(h_in_keys)
d_in_values = cp.asarray(h_in_values)
# Perform the merge sort.
cuda.compute.merge_sort(
d_in_keys=d_in_keys,
d_in_values=d_in_values,
d_out_keys=d_in_keys,
d_out_values=d_in_values,
op=OpKind.LESS,
num_items=d_in_keys.size,
)
# Verify the result.
h_out_keys = cp.asnumpy(d_in_keys)
h_out_values = cp.asnumpy(d_in_values)
argsort = np.argsort(h_in_keys, stable=True)
expected_keys = np.array(h_in_keys)[argsort]
expected_values = np.array(h_in_values)[argsort]
assert np.array_equal(h_out_keys, expected_keys)
assert np.array_equal(h_out_values, expected_values)
print(f"Merge sort basic result - keys: {h_out_keys}, values: {h_out_values}")

68
cccl_upstream/python/cuda_cccl/tests/compute/examples/sort/merge_sort_object.py Normal file
View File

@@ -0,0 +1,68 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Merge sort example demonstrating the object API.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
OpKind,
)
# Prepare the input and output arrays.
dtype = np.int32
h_input_keys = np.array([4, 2, 3, 1], dtype=dtype)
h_input_values = np.array([40, 20, 30, 10], dtype=dtype)
d_input_keys = cp.asarray(h_input_keys)
d_input_values = cp.asarray(h_input_values)
d_output_keys = cp.empty_like(d_input_keys)
d_output_values = cp.empty_like(d_input_values)
# Create the merge sort object.
sorter = cuda.compute.make_merge_sort(
d_in_keys=d_input_keys,
d_in_values=d_input_values,
d_out_keys=d_output_keys,
d_out_values=d_output_values,
op=OpKind.LESS,
)
# Get the temporary storage size.
temp_storage_size = sorter(
temp_storage=None,
d_in_keys=d_input_keys,
d_in_values=d_input_values,
d_out_keys=d_output_keys,
d_out_values=d_output_values,
op=OpKind.LESS,
num_items=len(h_input_keys),
)
# Allocate the temporary storage.
d_temp_storage = cp.empty(temp_storage_size, dtype=np.uint8)
# Perform the merge sort.
sorter(
temp_storage=d_temp_storage,
d_in_keys=d_input_keys,
d_in_values=d_input_values,
d_out_keys=d_output_keys,
d_out_values=d_output_values,
op=OpKind.LESS,
num_items=len(h_input_keys),
)
# Verify the result.
expected_keys = np.array([1, 2, 3, 4], dtype=dtype)
expected_values = np.array([10, 20, 30, 40], dtype=dtype)
actual_keys = d_output_keys.get()
actual_values = d_output_values.get()
np.testing.assert_array_equal(actual_keys, expected_keys)
np.testing.assert_array_equal(actual_values, expected_values)
print("Merge sort object example completed successfully")

51
cccl_upstream/python/cuda_cccl/tests/compute/examples/sort/radix_sort_basic.py Normal file
View File

@@ -0,0 +1,51 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use radix_sort to sort keys and values.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
SortOrder,
)
# Prepare the input and output arrays.
h_in_keys = np.array([-5, 0, 2, -3, 2, 4, 0, -1, 2, 8], dtype="int32")
h_in_values = np.array(
[-3.2, 2.2, 1.9, 4.0, -3.9, 2.7, 0, 8.3 - 1, 2.9, 5.4], dtype="float32"
)
d_in_keys = cp.asarray(h_in_keys)
d_in_values = cp.asarray(h_in_values)
# Prepare the output arrays.
d_out_keys = cp.empty_like(d_in_keys)
d_out_values = cp.empty_like(d_in_values)
# Perform the radix sort.
cuda.compute.radix_sort(
d_in_keys=d_in_keys,
d_out_keys=d_out_keys,
d_in_values=d_in_values,
d_out_values=d_out_values,
order=SortOrder.ASCENDING,
num_items=d_in_keys.size,
)
# Verify the result.
h_out_keys = cp.asnumpy(d_out_keys)
h_out_values = cp.asnumpy(d_out_values)
argsort = np.argsort(h_in_keys, stable=True)
expected_keys = np.array(h_in_keys)[argsort]
expected_values = np.array(h_in_values)[argsort]
assert np.array_equal(h_out_keys, expected_keys)
assert np.array_equal(h_out_values, expected_values)
print(f"Radix sort basic result - keys: {h_out_keys}, values: {h_out_values}")

55
cccl_upstream/python/cuda_cccl/tests/compute/examples/sort/radix_sort_buffer.py Normal file
View File

@@ -0,0 +1,55 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use radix_sort with DoubleBuffer for reduced temporary storage.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
DoubleBuffer,
SortOrder,
)
# Prepare the input and output arrays.
h_in_keys = np.array([-5, 0, 2, -3, 2, 4, 0, -1, 2, 8], dtype="int32")
h_in_values = np.array(
[-3.2, 2.2, 1.9, 4.0, -3.9, 2.7, 0, 8.3 - 1, 2.9, 5.4], dtype="float32"
)
d_in_keys = cp.asarray(h_in_keys)
d_in_values = cp.asarray(h_in_values)
d_out_keys = cp.empty_like(d_in_keys)
d_out_values = cp.empty_like(d_in_values)
# Create the double buffer.
keys_double_buffer = DoubleBuffer(d_in_keys, d_out_keys)
values_double_buffer = DoubleBuffer(d_in_values, d_out_values)
# Perform the radix sort.
cuda.compute.radix_sort(
d_in_keys=keys_double_buffer,
d_out_keys=None,
d_in_values=values_double_buffer,
d_out_values=None,
order=SortOrder.ASCENDING,
num_items=d_in_keys.size,
)
# Verify the result.
h_out_keys = cp.asnumpy(keys_double_buffer.current())
h_out_values = cp.asnumpy(values_double_buffer.current())
argsort = np.argsort(h_in_keys, stable=True)
h_expected_keys = np.array(h_in_keys)[argsort]
h_expected_values = np.array(h_in_values)[argsort]
assert np.array_equal(h_out_keys, h_expected_keys)
assert np.array_equal(h_out_values, h_expected_values)
print(f"Radix sort buffer result - keys: {h_out_keys}, values: {h_out_values}")

64
cccl_upstream/python/cuda_cccl/tests/compute/examples/sort/radix_sort_object.py Normal file
View File

@@ -0,0 +1,64 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use radix_sort with the object API.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
SortOrder,
)
# Prepare the input and output arrays.
dtype = np.int32
h_input_keys = np.array([4, 2, 3, 1], dtype=dtype)
h_input_values = np.array([40, 20, 30, 10], dtype=dtype)
d_input_keys = cp.asarray(h_input_keys)
d_input_values = cp.asarray(h_input_values)
d_output_keys = cp.empty_like(d_input_keys)
d_output_values = cp.empty_like(d_input_values)
# Create the radix sort object.
sorter = cuda.compute.make_radix_sort(
d_in_keys=d_input_keys,
d_out_keys=d_output_keys,
d_in_values=d_input_values,
d_out_values=d_output_values,
order=SortOrder.ASCENDING,
)
# Get the temporary storage size.
temp_storage_size = sorter(
temp_storage=None,
d_in_keys=d_input_keys,
d_out_keys=d_output_keys,
d_in_values=d_input_values,
d_out_values=d_output_values,
num_items=len(h_input_keys),
)
d_temp_storage = cp.empty(temp_storage_size, dtype=np.uint8)
# Perform the radix sort.
sorter(
temp_storage=d_temp_storage,
d_in_keys=d_input_keys,
d_out_keys=d_output_keys,
d_in_values=d_input_values,
d_out_values=d_output_values,
num_items=len(h_input_keys),
)
# Verify the result.
expected_keys = np.array([1, 2, 3, 4], dtype=dtype)
expected_values = np.array([10, 20, 30, 40], dtype=dtype)
actual_keys = d_output_keys.get()
actual_values = d_output_values.get()
np.testing.assert_array_equal(actual_keys, expected_keys)
np.testing.assert_array_equal(actual_values, expected_values)
print("Radix sort object example completed successfully")

55
cccl_upstream/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_basic.py Normal file
View File

@@ -0,0 +1,55 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use segmented_sort to sort keys and values within segments.
"""
import cupy as cp
import numpy as np
import cuda.compute
# Prepare input keys and values, and segment offsets.
h_in_keys = np.array([9, 1, 5, 4, 2, 8, 7, 3, 6], dtype="int32")
h_in_vals = np.array([90, 10, 50, 40, 20, 80, 70, 30, 60], dtype="int32")
# 3 segments: [0,3), [3,5), [5,9)
start_offsets = np.array([0, 3, 5], dtype=np.int64)
end_offsets = np.array([3, 5, 9], dtype=np.int64)
d_in_keys = cp.asarray(h_in_keys)
d_in_vals = cp.asarray(h_in_vals)
d_out_keys = cp.empty_like(d_in_keys)
d_out_vals = cp.empty_like(d_in_vals)
# Perform the segmented sort (ascending within each segment).
cuda.compute.segmented_sort(
d_in_keys=d_in_keys,
d_out_keys=d_out_keys,
d_in_values=d_in_vals,
d_out_values=d_out_vals,
num_items=d_in_keys.size,
num_segments=start_offsets.size,
start_offsets_in=cp.asarray(start_offsets),
end_offsets_in=cp.asarray(end_offsets),
order=cuda.compute.SortOrder.ASCENDING,
)
# Verify the result.
h_out_keys = cp.asnumpy(d_out_keys)
h_out_vals = cp.asnumpy(d_out_vals)
expected_pairs = []
for s, e in zip(start_offsets, end_offsets):
seg_pairs = sorted(zip(h_in_keys[s:e], h_in_vals[s:e]), key=lambda kv: kv[0])
expected_pairs.extend(seg_pairs)
expected_keys = np.array([k for k, _ in expected_pairs], dtype=h_in_keys.dtype)
expected_vals = np.array([v for _, v in expected_pairs], dtype=h_in_vals.dtype)
assert np.array_equal(h_out_keys, expected_keys)
assert np.array_equal(h_out_vals, expected_vals)
print(f"Segmented sort basic result - keys: {h_out_keys}, values: {h_out_vals}")

61
cccl_upstream/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_buffer.py Normal file
View File

@@ -0,0 +1,61 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use segmented_sort with DoubleBuffer for reduced temporary storage.
"""
import cupy as cp
import numpy as np
import cuda.compute
# Prepare input keys and values, and segment offsets.
h_in_keys = np.array([9, 1, 5, 4, 2, 8, 7, 3, 6], dtype="int32")
h_in_vals = np.array([90, 10, 50, 40, 20, 80, 70, 30, 60], dtype="int32")
# 3 segments: [0,3), [3,5), [5,9)
start_offsets = np.array([0, 3, 5], dtype=np.int64)
end_offsets = np.array([3, 5, 9], dtype=np.int64)
d_in_keys = cp.asarray(h_in_keys)
d_in_vals = cp.asarray(h_in_vals)
d_tmp_keys = cp.empty_like(d_in_keys)
d_tmp_vals = cp.empty_like(d_in_vals)
# Create double buffers for keys and values.
keys_db = cuda.compute.DoubleBuffer(d_in_keys, d_tmp_keys)
vals_db = cuda.compute.DoubleBuffer(d_in_vals, d_tmp_vals)
# Perform the segmented sort (descending within each segment).
cuda.compute.segmented_sort(
d_in_keys=keys_db,
d_out_keys=None,
d_in_values=vals_db,
d_out_values=None,
num_items=d_in_keys.size,
num_segments=start_offsets.size,
start_offsets_in=cp.asarray(start_offsets),
end_offsets_in=cp.asarray(end_offsets),
order=cuda.compute.SortOrder.DESCENDING,
)
# Verify the result.
h_out_keys = cp.asnumpy(keys_db.current())
h_out_vals = cp.asnumpy(vals_db.current())
expected_pairs = []
for s, e in zip(start_offsets, end_offsets):
seg_pairs = sorted(
zip(h_in_keys[s:e], h_in_vals[s:e]), key=lambda kv: kv[0], reverse=True
)
expected_pairs.extend(seg_pairs)
expected_keys = np.array([k for k, _ in expected_pairs], dtype=h_in_keys.dtype)
expected_vals = np.array([v for _, v in expected_pairs], dtype=h_in_vals.dtype)
assert np.array_equal(h_out_keys, expected_keys)
assert np.array_equal(h_out_vals, expected_vals)
print(f"Segmented sort buffer result - keys: {h_out_keys}, values: {h_out_vals}")

78
cccl_upstream/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_object.py Normal file
View File

@@ -0,0 +1,78 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use segmented_sort with the object API.
"""
import cupy as cp
import numpy as np
import cuda.compute
# Prepare the input and segment offsets.
dtype = np.int32
h_input_keys = np.array([9, 1, 5, 4, 2, 8, 7, 3, 6], dtype=dtype)
h_input_vals = np.array([90, 10, 50, 40, 20, 80, 70, 30, 60], dtype=dtype)
start_offsets = np.array([0, 3, 5], dtype=np.int64)
end_offsets = np.array([3, 5, 9], dtype=np.int64)
d_input_keys = cp.asarray(h_input_keys)
d_input_vals = cp.asarray(h_input_vals)
d_output_keys = cp.empty_like(d_input_keys)
d_output_vals = cp.empty_like(d_input_vals)
# Create the segmented sort object.
sorter = cuda.compute.make_segmented_sort(
d_in_keys=d_input_keys,
d_out_keys=d_output_keys,
d_in_values=d_input_vals,
d_out_values=d_output_vals,
start_offsets_in=cp.asarray(start_offsets),
end_offsets_in=cp.asarray(end_offsets),
order=cuda.compute.SortOrder.ASCENDING,
)
# Get the temporary storage size.
temp_storage_size = sorter(
temp_storage=None,
d_in_keys=d_input_keys,
d_out_keys=d_output_keys,
d_in_values=d_input_vals,
d_out_values=d_output_vals,
num_items=h_input_keys.size,
num_segments=start_offsets.size,
start_offsets_in=cp.asarray(start_offsets),
end_offsets_in=cp.asarray(end_offsets),
)
d_temp_storage = cp.empty(temp_storage_size, dtype=np.uint8)
# Perform the segmented sort.
sorter(
temp_storage=d_temp_storage,
d_in_keys=d_input_keys,
d_out_keys=d_output_keys,
d_in_values=d_input_vals,
d_out_values=d_output_vals,
num_items=h_input_keys.size,
num_segments=start_offsets.size,
start_offsets_in=cp.asarray(start_offsets),
end_offsets_in=cp.asarray(end_offsets),
)
# Verify the result.
expected_pairs = []
for s, e in zip(start_offsets, end_offsets):
seg_pairs = sorted(zip(h_input_keys[s:e], h_input_vals[s:e]), key=lambda kv: kv[0])
expected_pairs.extend(seg_pairs)
expected_keys = np.array([k for k, _ in expected_pairs], dtype=dtype)
expected_values = np.array([v for _, v in expected_pairs], dtype=dtype)
actual_keys = d_output_keys.get()
actual_values = d_output_vals.get()
np.testing.assert_array_equal(actual_keys, expected_keys)
np.testing.assert_array_equal(actual_values, expected_values)
print("Segmented sort object example completed successfully")

3
cccl_upstream/python/cuda_cccl/tests/compute/examples/struct/__init__.py Normal file
View File

@@ -0,0 +1,3 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

77
cccl_upstream/python/cuda_cccl/tests/compute/examples/struct/nested_struct_reduction.py Normal file
View File

@@ -0,0 +1,77 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example demonstrating reductions with nested gpu_struct types.
This example shows how to define nested structs and use them in reduction
operations. The reduction combines values from both the outer and inner
struct fields.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import gpu_struct
# Define an inner struct to hold coordinate data
@gpu_struct
class Point:
x: np.int32
y: np.int32
# Define an outer struct that contains the inner struct
@gpu_struct
class Particle:
id: np.int64
position: Point
def sum_particles(p1, p2):
"""Reduction operation that sums all fields of two particles."""
return Particle(
p1.id + p2.id,
Point(p1.position.x + p2.position.x, p1.position.y + p2.position.y),
)
# Prepare the input data
num_items = 10
h_data = np.zeros(num_items, dtype=Particle.dtype)
for i in range(num_items):
h_data[i]["id"] = i * 10
h_data[i]["position"]["x"] = i
h_data[i]["position"]["y"] = i * 2
# Copy to device
d_input = cp.empty(num_items, dtype=Particle.dtype)
d_input.set(h_data)
# Prepare output and initial value
d_output = cp.empty(1, dtype=Particle.dtype)
h_init = Particle(0, Point(0, 0))
# Perform the reduction
cuda.compute.reduce_into(
d_in=d_input, d_out=d_output, num_items=num_items, op=sum_particles, h_init=h_init
)
# Verify the result
result = d_output.get()[0]
expected_id = sum(i * 10 for i in range(num_items))
expected_x = sum(range(num_items))
expected_y = sum(i * 2 for i in range(num_items))
assert result["id"] == expected_id
assert result["position"]["x"] == expected_x
assert result["position"]["y"] == expected_y
print("Nested struct reduction result:")
print(f" id: {result['id']} (expected: {expected_id})")
print(f" position.x: {result['position']['x']} (expected: {expected_x})")
print(f" position.y: {result['position']['y']} (expected: {expected_y})")

82
cccl_upstream/python/cuda_cccl/tests/compute/examples/struct/nested_struct_tuple_construction.py Normal file
View File

@@ -0,0 +1,82 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing tuple syntax for constructing nested gpu_struct types.
When working with nested structs in device functions, you can use tuple syntax
as a convenient shorthand for constructing the nested struct values. This can
make code more concise while maintaining the same functionality.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import gpu_struct
# Define nested structs
@gpu_struct
class Stats:
count: np.int32
sum: np.float32
@gpu_struct
class DataPoint:
value: np.int64
stats: Stats
def sum_with_tuples(d1, d2):
"""
Reduction operation using tuple syntax for nested struct construction.
Instead of writing: Stats(d1.stats.count + d2.stats.count, ...)
We can use tuple syntax: (d1.stats.count + d2.stats.count, ...)
"""
return DataPoint(
d1.value + d2.value,
# Tuple syntax for constructing the nested Stats struct
(d1.stats.count + d2.stats.count, d1.stats.sum + d2.stats.sum),
)
# Prepare the input data
num_items = 10
h_data = np.zeros(num_items, dtype=DataPoint.dtype)
for i in range(num_items):
h_data[i]["value"] = i * 10
h_data[i]["stats"]["count"] = 1
h_data[i]["stats"]["sum"] = float(i)
# Copy to device
d_input = cp.empty(num_items, dtype=DataPoint.dtype)
d_input.set(h_data)
# Prepare output and initial value
d_output = cp.empty(1, dtype=DataPoint.dtype)
h_init = DataPoint(0, Stats(0, 0.0))
# Perform the reduction
cuda.compute.reduce_into(
d_in=d_input, d_out=d_output, num_items=num_items, op=sum_with_tuples, h_init=h_init
)
# Verify the result
result = d_output.get()[0]
expected_value = sum(i * 10 for i in range(num_items))
expected_count = num_items
expected_sum = sum(float(i) for i in range(num_items))
assert result["value"] == expected_value
assert result["stats"]["count"] == expected_count
assert np.isclose(result["stats"]["sum"], expected_sum)
print("Nested struct with tuple construction result:")
print(f" value: {result['value']} (expected: {expected_value})")
print(f" stats.count: {result['stats']['count']} (expected: {expected_count})")
print(f" stats.sum: {result['stats']['sum']:.2f} (expected: {expected_sum:.2f})")

98
cccl_upstream/python/cuda_cccl/tests/compute/examples/struct/nested_struct_zip_iterator.py Normal file
View File

@@ -0,0 +1,98 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing ZipIterator with nested gpu_struct types.
This example demonstrates combining separate arrays of nested structs using
ZipIterator, then performing a reduction that operates on the combined data.
This is useful when you have related data stored in separate arrays that need
to be processed together.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import ZipIterator, gpu_struct
# Define nested structs for geometric and color data
@gpu_struct
class Point:
x: np.int32
y: np.int32
@gpu_struct
class Color:
r: np.uint8
g: np.uint8
b: np.uint8
@gpu_struct
class Pixel:
position: Point
color: Color
def sum_pixels(p1, p2):
"""Reduction operation that sums all fields of two pixels."""
return Pixel(
Point(p1.position.x + p2.position.x, p1.position.y + p2.position.y),
Color(
p1.color.r + p2.color.r, p1.color.g + p2.color.g, p1.color.b + p2.color.b
),
)
# Prepare separate arrays for points and colors
num_items = 100
h_points = np.array([(i, i * 2) for i in range(num_items)], dtype=Point.dtype)
h_colors = np.array(
[(i % 256, (i * 2) % 256, (i * 3) % 256) for i in range(num_items)],
dtype=Color.dtype,
)
d_points = cp.empty(num_items, dtype=Point.dtype)
d_points.set(h_points)
d_colors = cp.empty(num_items, dtype=Color.dtype)
d_colors.set(h_colors)
# Create a zip iterator to combine the points and colors
zip_it = ZipIterator(d_points, d_colors)
# Prepare output and initial value
d_output = cp.empty(1, dtype=Pixel.dtype)
h_init = Pixel(Point(0, 0), Color(0, 0, 0))
# Perform the reduction on the zipped data
cuda.compute.reduce_into(
d_in=zip_it, d_out=d_output, num_items=num_items, op=sum_pixels, h_init=h_init
)
# Verify the result
result = d_output.get()[0]
expected_x = sum(range(num_items))
expected_y = sum(i * 2 for i in range(num_items))
expected_r = sum(i % 256 for i in range(num_items)) % 256
expected_g = sum((i * 2) % 256 for i in range(num_items)) % 256
expected_b = sum((i * 3) % 256 for i in range(num_items)) % 256
assert result["position"]["x"] == expected_x
assert result["position"]["y"] == expected_y
assert result["color"]["r"] == expected_r
assert result["color"]["g"] == expected_g
assert result["color"]["b"] == expected_b
print("Nested struct with ZipIterator result:")
print(f" position.x: {result['position']['x']} (expected: {expected_x})")
print(f" position.y: {result['position']['y']} (expected: {expected_y})")
print(f" color.r: {result['color']['r']} (expected: {expected_r})")
print(f" color.g: {result['color']['g']} (expected: {expected_g})")
print(f" color.b: {result['color']['b']} (expected: {expected_b})")

50
cccl_upstream/python/cuda_cccl/tests/compute/examples/struct/struct_reduction.py Normal file
View File

@@ -0,0 +1,50 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Finding the maximum green value in a sequence of pixels using `reduce_into`
with a custom data type.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import gpu_struct
# Define a custom data type to store the pixel values.
@gpu_struct
class Pixel:
r: np.int32
g: np.int32
b: np.int32
# Define a reduction operation that returns the pixel with the maximum green value.
def max_g_value(x, y):
return x if x.g > y.g else y
# Prepare the input and output arrays.
d_rgb = cp.random.randint(0, 256, (10, 3), dtype=np.int32).view(Pixel.dtype)
d_out = cp.empty(1, Pixel.dtype)
# Prepare the initial value for the reduction.
h_init = Pixel(0, 0, 0)
# Perform the reduction.
cuda.compute.reduce_into(
d_in=d_rgb, d_out=d_out, num_items=d_rgb.size, op=max_g_value, h_init=h_init
)
# Calculate the expected result.
h_rgb = d_rgb.get()
expected = h_rgb[h_rgb.view("int32")[:, 1].argmax()]
# Verify the result.
assert expected["g"] == d_out.get()["g"]
result = d_out.get()
print(f"Pixel reduction result: {result}")

60
cccl_upstream/python/cuda_cccl/tests/compute/examples/struct/struct_transform.py Normal file
View File

@@ -0,0 +1,60 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example demonstrating binary_transform with custom struct types.
When working with struct inputs in transform operations, you need to provide
type annotations to help Numba infer the correct types. Unlike reduce_into
which can infer types from h_init, transform operations require explicit
annotations when using struct inputs.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import gpu_struct
@gpu_struct
class Point2D:
x: np.float32
y: np.float32
def add_points(p1: Point2D, p2: Point2D) -> Point2D:
return Point2D(p1.x + p2.x, p1.y + p2.y)
num_items = 1000
h_in1 = np.empty(num_items, dtype=Point2D.dtype)
h_in1["x"] = np.random.rand(num_items).astype(np.float32)
h_in1["y"] = np.random.rand(num_items).astype(np.float32)
h_in2 = np.empty(num_items, dtype=Point2D.dtype)
h_in2["x"] = np.random.rand(num_items).astype(np.float32)
h_in2["y"] = np.random.rand(num_items).astype(np.float32)
d_in1 = cp.empty_like(h_in1)
d_in1.set(h_in1)
d_in2 = cp.empty_like(h_in2)
d_in2.set(h_in2)
d_out = cp.empty_like(d_in1)
cuda.compute.binary_transform(
d_in1=d_in1, d_in2=d_in2, d_out=d_out, op=add_points, num_items=num_items
)
result = d_out.get()
np.testing.assert_allclose(result["x"], h_in1["x"] + h_in2["x"], rtol=1e-5)
np.testing.assert_allclose(result["y"], h_in1["y"] + h_in2["y"], rtol=1e-5)
print("Binary transform with structs completed successfully")
print(f"First result point: x={result[0]['x']:.4f}, y={result[0]['y']:.4f}")

35
cccl_upstream/python/cuda_cccl/tests/compute/examples/transform/binary_transform_basic.py Normal file
View File

@@ -0,0 +1,35 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use binary_transform to perform elementwise addition.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
OpKind,
)
# Prepare the input and output arrays.
input1_data = np.array([1, 2, 3, 4], dtype=np.int32)
input2_data = np.array([10, 20, 30, 40], dtype=np.int32)
d_in1 = cp.asarray(input1_data)
d_in2 = cp.asarray(input2_data)
d_out = cp.empty_like(d_in1)
# Perform the binary transform.
cuda.compute.binary_transform(
d_in1=d_in1, d_in2=d_in2, d_out=d_out, op=OpKind.PLUS, num_items=len(d_in1)
)
# Verify the result.
result = d_out.get()
expected = input1_data + input2_data
np.testing.assert_array_equal(result, expected)
print(f"Binary transform result: {result}")

44
cccl_upstream/python/cuda_cccl/tests/compute/examples/transform/binary_transform_object.py Normal file
View File

@@ -0,0 +1,44 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Binary transform examples demonstrating the transform object API.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
OpKind,
)
# Prepare the input and output arrays.
dtype = np.int32
h_input1 = np.array([1, 2, 3, 4], dtype=dtype)
h_input2 = np.array([10, 20, 30, 40], dtype=dtype)
d_input1 = cp.asarray(h_input1)
d_input2 = cp.asarray(h_input2)
d_output = cp.empty_like(d_input1)
# Create the binary transform object.
transformer = cuda.compute.make_binary_transform(
d_in1=d_input1, d_in2=d_input2, d_out=d_output, op=OpKind.PLUS
)
# Perform the binary transform.
transformer(
d_in1=d_input1,
d_in2=d_input2,
d_out=d_output,
op=OpKind.PLUS,
num_items=len(h_input1),
)
# Verify the result.
expected_result = np.array([11, 22, 33, 44], dtype=dtype)
actual_result = d_output.get()
np.testing.assert_array_equal(actual_result, expected_result)
print("Binary transform object example completed successfully")

34
cccl_upstream/python/cuda_cccl/tests/compute/examples/transform/unary_transform_basic.py Normal file
View File

@@ -0,0 +1,34 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use unary_transform to apply a unary operation to each element.
"""
import cupy as cp
import numpy as np
import cuda.compute
# Prepare the input and output arrays.
input_data = np.array([1, 2, 3, 4, 5], dtype=np.int32)
d_in = cp.asarray(input_data)
d_out = cp.empty_like(d_in)
# Define the unary operation.
def op(a):
return a + 1
# Perform the unary transform.
cuda.compute.unary_transform(d_in=d_in, d_out=d_out, op=op, num_items=len(d_in))
# Verify the result.
result = d_out.get()
expected = input_data + 1
np.testing.assert_array_equal(result, expected)
print(f"Unary transform result: {result}")

39
cccl_upstream/python/cuda_cccl/tests/compute/examples/transform/unary_transform_object.py Normal file
View File

@@ -0,0 +1,39 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Unary transform examples demonstrating the object API and well-known operations.
"""
import cupy as cp
import numpy as np
import cuda.compute
# Prepare the input and output arrays.
dtype = np.int32
h_input = np.array([1, 2, 3, 4], dtype=dtype)
d_input = cp.asarray(h_input)
d_output = cp.empty_like(d_input)
# Define the unary operation.
def add_one_op(a):
return a + 1
# Create the unary transform object.
transformer = cuda.compute.make_unary_transform(
d_in=d_input, d_out=d_output, op=add_one_op
)
# Perform the unary transform.
transformer(d_in=d_input, d_out=d_output, op=add_one_op, num_items=len(h_input))
# Verify the result.
expected_result = np.array([2, 3, 4, 5], dtype=dtype)
actual_result = d_output.get()
np.testing.assert_array_equal(actual_result, expected_result)
print("Unary transform object example completed successfully")

5
cccl_upstream/python/cuda_cccl/tests/compute/examples/unique/__init__.py Normal file
View File

@@ -0,0 +1,5 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Parallel unique algorithms examples package."""

52
cccl_upstream/python/cuda_cccl/tests/compute/examples/unique/unique_by_key_basic.py Normal file
View File

@@ -0,0 +1,52 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use unique_by_key to remove all
but the first value for each group of consecutive keys.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
OpKind,
)
# Prepare the input and output arrays.
h_in_keys = np.array([0, 2, 2, 9, 5, 5, 5, 8], dtype="int32")
h_in_values = np.array([1, 2, 3, 4, 5, 6, 7, 8], dtype="float32")
d_in_keys = cp.asarray(h_in_keys)
d_in_values = cp.asarray(h_in_values)
d_out_keys = cp.empty_like(d_in_keys)
d_out_values = cp.empty_like(d_in_values)
d_out_num_selected = cp.empty(1, np.int32)
# Perform the unique by key operation.
cuda.compute.unique_by_key(
d_in_keys=d_in_keys,
d_in_items=d_in_values,
d_out_keys=d_out_keys,
d_out_items=d_out_values,
d_out_num_selected=d_out_num_selected,
op=OpKind.EQUAL_TO,
num_items=d_in_keys.size,
)
# Verify the result.
num_selected = cp.asnumpy(d_out_num_selected)[0]
h_out_keys = cp.asnumpy(d_out_keys)[:num_selected]
h_out_values = cp.asnumpy(d_out_values)[:num_selected]
expected_keys = np.array([0, 2, 9, 5, 8])
expected_values = np.array([1, 2, 4, 5, 8])
assert np.array_equal(h_out_keys, expected_keys)
assert np.array_equal(h_out_values, expected_values)
print(
f"Unique by key basic result - keys: {h_out_keys}, values: {h_out_values}, count: {num_selected}"
)

73
cccl_upstream/python/cuda_cccl/tests/compute/examples/unique/unique_by_key_object.py Normal file
View File

@@ -0,0 +1,73 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# example-begin
"""
Example showing how to use unique_by_key with the object API.
"""
import cupy as cp
import numpy as np
import cuda.compute
from cuda.compute import (
OpKind,
)
# Unique by key example demonstrating the object API
dtype = np.int32
h_input_keys = np.array([1, 1, 2, 3, 3], dtype=dtype)
h_input_values = np.array([10, 20, 30, 40, 50], dtype=dtype)
d_input_keys = cp.asarray(h_input_keys)
d_input_values = cp.asarray(h_input_values)
d_output_keys = cp.empty_like(d_input_keys)
d_output_values = cp.empty_like(d_input_values)
d_num_selected = cp.empty(1, dtype=np.int32)
# Create the unique by key object.
uniquer = cuda.compute.make_unique_by_key(
d_in_keys=d_input_keys,
d_in_items=d_input_values,
d_out_keys=d_output_keys,
d_out_items=d_output_values,
d_out_num_selected=d_num_selected,
op=OpKind.EQUAL_TO,
)
# Get the temporary storage size.
temp_storage_size = uniquer(
temp_storage=None,
d_in_keys=d_input_keys,
d_in_items=d_input_values,
d_out_keys=d_output_keys,
d_out_items=d_output_values,
d_out_num_selected=d_num_selected,
op=OpKind.EQUAL_TO,
num_items=len(h_input_keys),
)
# Allocate the temporary storage.
d_temp_storage = cp.empty(temp_storage_size, dtype=np.uint8)
# Perform the unique by key operation.
uniquer(
temp_storage=d_temp_storage,
d_in_keys=d_input_keys,
d_in_items=d_input_values,
d_out_keys=d_output_keys,
d_out_items=d_output_values,
d_out_num_selected=d_num_selected,
op=OpKind.EQUAL_TO,
num_items=len(h_input_keys),
)
# Verify the result.
num_selected = d_num_selected.get()[0]
expected_keys = np.array([1, 2, 3], dtype=dtype)
expected_values = np.array([10, 30, 40], dtype=dtype)
actual_keys = d_output_keys.get()[:num_selected]
actual_values = d_output_values.get()[:num_selected]
np.testing.assert_array_equal(actual_keys, expected_keys)
np.testing.assert_array_equal(actual_values, expected_values)
print("Unique by key object example completed successfully")