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

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

cccl_upstream/python/cuda_cccl/tests/compute/conftest.py

@@ -0,0 +1,169 @@
+import builtins
+from collections.abc import Generator
+
+import numpy as np
+import pytest
+
+from cuda.core import Device, Stream
+
+try:
+    from cuda.compute._build_info import USING_V2
+except ImportError:
+    USING_V2 = False
+
+check_ldl_stl_in_sass = False
+
+
+# Define a pytest fixture that returns random arrays with different dtypes
+@pytest.fixture(
+    params=[
+        np.int8,
+        np.int16,
+        np.int32,
+        np.int64,
+        np.uint8,
+        np.uint16,
+        np.uint32,
+        np.uint64,
+        np.float16,
+        np.float32,
+        np.float64,
+        np.complex64,
+        np.complex128,
+    ]
+)
+def input_array(request):
+    dtype = np.dtype(request.param)
+    sample_size = 1000
+
+    # Generate random values based on the dtype
+    if np.issubdtype(dtype, np.integer):
+        is_unsigned = dtype.kind == "u"
+        # For integer types, use np.random.randint for random integers
+        if is_unsigned:
+            low_inclusive, high_exclusive = 0, 8
+        else:
+            low_inclusive, high_exclusive = -5, 6
+        array = np.random.randint(
+            low=low_inclusive, high=high_exclusive, size=sample_size, dtype=dtype
+        )
+    elif np.issubdtype(dtype, np.floating):
+        # For floating-point types, use np.random.random and cast to the required dtype
+        array = np.random.random(sample_size).astype(dtype)
+    elif np.issubdtype(dtype, np.complexfloating):
+        # For complex types, generate random real and imaginary parts
+        packed = np.random.random(2 * sample_size)
+        real_part = packed[:sample_size]
+        imag_part = packed[sample_size:]
+        array = (real_part + 1j * imag_part).astype(dtype)
+
+    return array
+
+
+# Define a pytest fixture that returns random floating-point arrays only
+@pytest.fixture(
+    params=[
+        np.float32,
+        np.float64,
+    ]
+)
+def floating_array(request):
+    dtype = np.dtype(request.param)
+    sample_size = 1000
+
+    # Generate random floating-point values
+    array = np.random.random(sample_size).astype(dtype)
+    return array
+
+
+@pytest.fixture(scope="function")
+def cuda_stream() -> Generator[Stream, None, None]:
+    device = Device()
+    device.set_current()
+    stream = device.create_stream()
+    try:
+        yield stream
+    finally:
+        stream.close()
+
+
+@pytest.fixture(scope="function", autouse=True)
+def verify_sass(request):
+    if request.node.get_closest_marker("no_verify_sass"):
+        return
+
+    if not check_ldl_stl_in_sass:
+        return
+
+    # Pull monkeypatch dynamically rather than as a fixture parameter so this
+    # autouse fixture does not add monkeypatch to every test's static fixture
+    # closure. pytest-run-parallel treats monkeypatch as thread-unsafe based on
+    # that closure, so a parameter here would serialize the entire free-threaded
+    # parallel sweep -- even though this fixture only patches on the opt-in
+    # SASS-check path (check_ldl_stl_in_sass, off by default and in CI).
+    monkeypatch = request.getfixturevalue("monkeypatch")
+
+    import cuda.compute._cccl_interop
+
+    monkeypatch.setattr(
+        cuda.compute._cccl_interop,
+        "_check_sass",
+        True,
+    )
+
+
+@pytest.fixture
+def raise_on_numba_import(monkeypatch):
+    """This fixture will raise if a test attempts to import numba"""
+    real_import = builtins.__import__
+
+    def guarded_import(name, *args, **kwargs):
+        if name == "numba" or name.startswith("numba."):
+            raise ModuleNotFoundError(
+                "This test is marked 'no_numba' but attempted to import it"
+            )
+        return real_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(builtins, "__import__", guarded_import)
+
+
+def pytest_collection_modifyitems(config, items):
+    """Runs after pytest collects the tests. Makes a test marked no_numba fail
+    if it imports numba, and skips a test marked serialization when running on
+    the v2 (HostJIT) backend."""
+    serialization_skip = pytest.mark.skip(
+        reason="serialization not supported on v2 (HostJIT) backend"
+    )
+
+    # Tests marked no_numba must not import numba. We enforce that by attaching
+    # the raise_on_numba_import fixture defined above to each one; it raises if
+    # numba is imported.
+    #
+    # We skip attaching it during a real pytest-run-parallel sweep of more than
+    # one thread: the fixture uses monkeypatch, which pytest-run-parallel
+    # serializes as thread-unsafe, so attaching it to every no_numba test would
+    # make the whole sweep run serially and defeat its purpose. A single-threaded
+    # run has no sweep to protect, so we attach it and keep the check.
+    #
+    # config.getoption gives the --parallel-threads value as an int for the
+    # default but a str when passed on the command line; normalize to str and
+    # count the run as parallel only for an explicit number > 1:
+    #
+    #   pytest ...                          -> 1       single-threaded, attach
+    #   pytest --parallel-threads=1 ...     -> "1"     single-threaded, attach
+    #   pytest --parallel-threads=8 ...     -> "8"     parallel, skip (CI sweep)
+    #   pytest --parallel-threads=auto ...  -> "auto"  single-threaded, attach
+    #
+    # "auto" counts as single-threaded on purpose: it can resolve to one CPU, so
+    # keeping the check is safer than dropping it on a non-parallel run. isdigit
+    # also stops int() from raising on the non-numeric "auto".
+    parallel_threads = str(config.getoption("parallel_threads", 1))
+    running_parallel = parallel_threads.isdigit() and int(parallel_threads) > 1
+    for item in items:
+        # no_numba: add raise_on_numba_import unless we skip it for the sweep
+        if item.get_closest_marker("no_numba") and not running_parallel:
+            if "raise_on_numba_import" not in item.fixturenames:
+                item.fixturenames.append("raise_on_numba_import")
+        # serialization is unsupported on v2 (HostJIT); skip those tests there
+        if USING_V2 and item.get_closest_marker("serialization"):
+            item.add_marker(serialization_skip)

cccl_upstream/python/cuda_cccl/tests/compute/examples/binary_search/__init__.py

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

cccl_upstream/python/cuda_cccl/tests/compute/examples/binary_search/lower_bound_basic.py

@@ -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)

cccl_upstream/python/cuda_cccl/tests/compute/examples/binary_search/lower_bound_object.py

@@ -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)

cccl_upstream/python/cuda_cccl/tests/compute/examples/binary_search/upper_bound_basic.py

@@ -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)

cccl_upstream/python/cuda_cccl/tests/compute/examples/binary_search/upper_bound_object.py

@@ -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)

cccl_upstream/python/cuda_cccl/tests/compute/examples/free_threading/__init__.py

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

cccl_upstream/python/cuda_cccl/tests/compute/examples/free_threading/direct_api.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/free_threading/object_api.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/histogram/__init__.py

@@ -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."""

cccl_upstream/python/cuda_cccl/tests/compute/examples/histogram/histogram_even_basic.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/histogram/histogram_object.py

@@ -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")

cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/__init__.py

@@ -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."""

cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/cache_modified_iterator_basic.py

@@ -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})")

cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/constant_iterator_basic.py

@@ -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})")

cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/counting_iterator_basic.py

@@ -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})")

cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/discard_iterator_basic.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/permutation_iterator_basic.py

@@ -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})")

cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/permutation_iterator_composed.py

@@ -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})"
+)

cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/permutation_iterator_output.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/reverse_input_iterator.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/reverse_output_iterator.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/shuffle_iterator_basic.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/transform_iterator_basic.py

@@ -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})")

cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/transform_iterator_lambda.py

@@ -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})"
+)

cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/transform_output_iterator.py

@@ -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)

cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/zip_iterator_counting.py

@@ -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})"
+)

cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/zip_iterator_elementwise.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/iterator/zip_iterator_reduction.py

@@ -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})"
+)

cccl_upstream/python/cuda_cccl/tests/compute/examples/partition/__init__.py

@@ -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."""

cccl_upstream/python/cuda_cccl/tests/compute/examples/partition/three_way_partition_basic.py

@@ -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")

cccl_upstream/python/cuda_cccl/tests/compute/examples/partition/three_way_partition_object.py

@@ -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")

cccl_upstream/python/cuda_cccl/tests/compute/examples/raw_op/__init__.py

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

cccl_upstream/python/cuda_cccl/tests/compute/examples/raw_op/cpp_stateful.py

@@ -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!")

cccl_upstream/python/cuda_cccl/tests/compute/examples/raw_op/cpp_stateless.py

@@ -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!")

cccl_upstream/python/cuda_cccl/tests/compute/examples/raw_op/llvm_stateless.py

@@ -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!")

cccl_upstream/python/cuda_cccl/tests/compute/examples/reduction/__init__.py

@@ -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."""

cccl_upstream/python/cuda_cccl/tests/compute/examples/reduction/min_reduction.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/reduction/minmax_reduction.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/reduction/reduce_object.py

@@ -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")

cccl_upstream/python/cuda_cccl/tests/compute/examples/reduction/sum_custom_reduction.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/reduction/sum_reduction.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/reduction/sum_reduction_lambda.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/scan/__init__.py

@@ -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."""

cccl_upstream/python/cuda_cccl/tests/compute/examples/scan/ema_example.py

@@ -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")

cccl_upstream/python/cuda_cccl/tests/compute/examples/scan/exclusive_scan_max.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/scan/exclusive_scan_object.py

@@ -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")

cccl_upstream/python/cuda_cccl/tests/compute/examples/scan/inclusive_scan_custom.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/scan/inclusive_scan_object.py

@@ -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")

cccl_upstream/python/cuda_cccl/tests/compute/examples/scan/logcdf_example.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/scan/running_average.py

@@ -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

cccl_upstream/python/cuda_cccl/tests/compute/examples/scan/segmented_sum.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/segmented/__init__.py

@@ -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."""

cccl_upstream/python/cuda_cccl/tests/compute/examples/segmented/segmented_reduce_basic.py

@@ -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()}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/segmented/segmented_reduce_object.py

@@ -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")

cccl_upstream/python/cuda_cccl/tests/compute/examples/select/__init__.py

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

cccl_upstream/python/cuda_cccl/tests/compute/examples/select/select_basic.py

@@ -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()

cccl_upstream/python/cuda_cccl/tests/compute/examples/select/select_object.py

@@ -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()

cccl_upstream/python/cuda_cccl/tests/compute/examples/select/select_with_iterator.py

@@ -0,0 +1,48 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+#
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+# example-begin
+import cupy as cp
+
+from 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()

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

@@ -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

cccl_upstream/python/cuda_cccl/tests/compute/examples/serialization/__init__.py

@@ -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."""

cccl_upstream/python/cuda_cccl/tests/compute/examples/serialization/ahead_of_time_compilation.py

@@ -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")

cccl_upstream/python/cuda_cccl/tests/compute/examples/serialization/serialize_roundtrip.py

@@ -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])

cccl_upstream/python/cuda_cccl/tests/compute/examples/sort/__init__.py

@@ -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."""

cccl_upstream/python/cuda_cccl/tests/compute/examples/sort/merge_sort_basic.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/sort/merge_sort_object.py

@@ -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")

cccl_upstream/python/cuda_cccl/tests/compute/examples/sort/radix_sort_basic.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/sort/radix_sort_buffer.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/sort/radix_sort_object.py

@@ -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")

cccl_upstream/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_basic.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_buffer.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/sort/segmented_sort_object.py

@@ -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")

cccl_upstream/python/cuda_cccl/tests/compute/examples/struct/__init__.py

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

cccl_upstream/python/cuda_cccl/tests/compute/examples/struct/nested_struct_reduction.py

@@ -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})")

cccl_upstream/python/cuda_cccl/tests/compute/examples/struct/nested_struct_tuple_construction.py

@@ -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})")

cccl_upstream/python/cuda_cccl/tests/compute/examples/struct/nested_struct_zip_iterator.py

@@ -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})")

cccl_upstream/python/cuda_cccl/tests/compute/examples/struct/struct_reduction.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/struct/struct_transform.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/transform/binary_transform_basic.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/transform/binary_transform_object.py

@@ -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")

cccl_upstream/python/cuda_cccl/tests/compute/examples/transform/unary_transform_basic.py

@@ -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}")

cccl_upstream/python/cuda_cccl/tests/compute/examples/transform/unary_transform_object.py

@@ -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")

cccl_upstream/python/cuda_cccl/tests/compute/examples/unique/__init__.py

@@ -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."""

cccl_upstream/python/cuda_cccl/tests/compute/examples/unique/unique_by_key_basic.py

@@ -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}"
+)

cccl_upstream/python/cuda_cccl/tests/compute/examples/unique/unique_by_key_object.py

@@ -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")

cccl_upstream/python/cuda_cccl/tests/compute/test_binary_search.py

@@ -0,0 +1,287 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+import numpy as np
+import pytest
+from _utils.device_array import DeviceArray
+
+import cuda.compute
+from cuda.compute import (
+    OpKind,
+    deserialize,
+    make_lower_bound,
+    make_upper_bound,
+    serialize,
+)
+
+DTYPE_LIST = [
+    np.int32,
+    np.int64,
+    np.uint32,
+    np.uint64,
+    np.float32,
+    np.float64,
+]
+
+
+def random_sorted_array(size, dtype, max_value=1000):
+    rng = np.random.default_rng()
+    if np.isdtype(dtype, "integral"):
+        data = rng.integers(max_value, size=size, dtype=dtype)
+    else:
+        if dtype == np.float16:  # pragma: no cover - float16 not used here
+            data = rng.random(size=size, dtype=np.float32).astype(dtype)
+        else:
+            data = rng.random(size=size, dtype=dtype)
+    data.sort()
+    return data
+
+
+@pytest.fixture(scope="function", autouse=True)
+def disable_sass_check(monkeypatch):
+    monkeypatch.setattr(
+        cuda.compute._cccl_interop,
+        "_check_sass",
+        False,
+    )
+
+
+@pytest.mark.parametrize(
+    "search, side",
+    [
+        (cuda.compute.lower_bound, "left"),
+        (cuda.compute.upper_bound, "right"),
+    ],
+)
+def test_binary_search_explicit_opkind_less(search, side):
+    h_data = np.array([1, 3, 3, 7, 9], dtype=np.int32)
+    h_values = np.array([0, 3, 4, 10], dtype=np.int32)
+
+    d_data = DeviceArray.from_numpy(h_data)
+    d_values = DeviceArray.from_numpy(h_values)
+    d_out = DeviceArray.empty(len(h_values), np.uintp)
+
+    search(
+        d_data=d_data,
+        num_items=len(h_data),
+        d_values=d_values,
+        num_values=len(h_values),
+        d_out=d_out,
+        comp=OpKind.LESS,
+    )
+
+    expected = np.searchsorted(h_data, h_values, side=side).astype(np.uintp)
+    np.testing.assert_array_equal(d_out.copy_to_host(), expected)
+
+
+@pytest.mark.parametrize(
+    "search, side",
+    [
+        (cuda.compute.lower_bound, "left"),
+        (cuda.compute.upper_bound, "right"),
+    ],
+)
+def test_binary_search_custom_comparator(search, side):
+    h_data = np.array([9, 7, 3, 3, 1], dtype=np.int32)
+    h_values = np.array([10, 4, 3, 0], dtype=np.int32)
+
+    def greater(lhs, rhs):
+        return lhs > rhs
+
+    d_data = DeviceArray.from_numpy(h_data)
+    d_values = DeviceArray.from_numpy(h_values)
+    d_out = DeviceArray.empty(len(h_values), np.uintp)
+
+    search(
+        d_data=d_data,
+        num_items=len(h_data),
+        d_values=d_values,
+        num_values=len(h_values),
+        d_out=d_out,
+        comp=greater,
+    )
+
+    expected = np.searchsorted(-h_data, -h_values, side=side).astype(np.uintp)
+    np.testing.assert_array_equal(d_out.copy_to_host(), expected)
+
+
+@pytest.mark.parametrize("dtype", DTYPE_LIST)
+@pytest.mark.parametrize(
+    "num_items,num_values", [(0, 0), (0, 128), (128, 0), (512, 128)]
+)
+def test_lower_bound_basic(dtype, num_items, num_values):
+    h_data = random_sorted_array(num_items, dtype)
+    h_values = random_sorted_array(num_values, dtype)
+
+    d_data = DeviceArray.from_numpy(h_data)
+    d_values = DeviceArray.from_numpy(h_values)
+    d_out = DeviceArray.empty(num_values, np.uintp)
+
+    cuda.compute.lower_bound(
+        d_data=d_data,
+        num_items=num_items,
+        d_values=d_values,
+        num_values=num_values,
+        d_out=d_out,
+    )
+
+    expected = np.searchsorted(h_data, h_values, side="left").astype(np.uintp)
+    got = d_out.copy_to_host()
+    assert np.array_equal(got, expected)
+
+
+@pytest.mark.parametrize("dtype", DTYPE_LIST)
+@pytest.mark.parametrize(
+    "num_items,num_values", [(0, 0), (0, 128), (128, 0), (512, 128)]
+)
+def test_upper_bound_basic(dtype, num_items, num_values):
+    h_data = random_sorted_array(num_items, dtype)
+    h_values = random_sorted_array(num_values, dtype)
+
+    d_data = DeviceArray.from_numpy(h_data)
+    d_values = DeviceArray.from_numpy(h_values)
+    d_out = DeviceArray.empty(num_values, np.uintp)
+
+    cuda.compute.upper_bound(
+        d_data=d_data,
+        num_items=num_items,
+        d_values=d_values,
+        num_values=num_values,
+        d_out=d_out,
+    )
+
+    expected = np.searchsorted(h_data, h_values, side="right").astype(np.uintp)
+    got = d_out.copy_to_host()
+    assert np.array_equal(got, expected)
+
+
+@pytest.mark.parametrize("dtype", DTYPE_LIST)
+def test_binary_search_with_duplicates(dtype):
+    rng = np.random.default_rng()
+    h_data = (
+        rng.integers(10, size=1024, dtype=dtype)
+        if np.isdtype(dtype, "integral")
+        else rng.random(1024, dtype=dtype)
+    )
+    h_data.sort()
+    h_values = (
+        rng.integers(10, size=128, dtype=dtype)
+        if np.isdtype(dtype, "integral")
+        else rng.random(128, dtype=dtype)
+    )
+
+    d_data = DeviceArray.from_numpy(h_data)
+    d_values = DeviceArray.from_numpy(h_values)
+    d_out = DeviceArray.empty(len(h_values), np.uintp)
+
+    cuda.compute.lower_bound(
+        d_data=d_data,
+        num_items=len(h_data),
+        d_values=d_values,
+        num_values=len(h_values),
+        d_out=d_out,
+    )
+    expected = np.searchsorted(h_data, h_values, side="left").astype(np.uintp)
+    got = d_out.copy_to_host()
+    assert np.array_equal(got, expected)
+
+    cuda.compute.upper_bound(
+        d_data=d_data,
+        num_items=len(h_data),
+        d_values=d_values,
+        num_values=len(h_values),
+        d_out=d_out,
+    )
+    expected = np.searchsorted(h_data, h_values, side="right").astype(np.uintp)
+    got = d_out.copy_to_host()
+    assert np.array_equal(got, expected)
+
+
+def test_binary_search_requires_unsigned_output():
+    """Output must be unsigned integer dtype for indices."""
+    h_data = np.array([1, 2, 3, 4], dtype=np.int32)
+    h_values = np.array([2, 3], dtype=np.int32)
+    d_data = DeviceArray.from_numpy(h_data)
+    d_values = DeviceArray.from_numpy(h_values)
+    d_out = DeviceArray.empty(len(h_values), np.int32)  # signed, should fail
+
+    with pytest.raises(TypeError, match="unsigned integer"):
+        cuda.compute.lower_bound(
+            d_data=d_data,
+            num_items=len(h_data),
+            d_values=d_values,
+            num_values=len(h_values),
+            d_out=d_out,
+        )
+
+
+def test_binary_search_requires_pointer_sized_output():
+    """Output must be pointer-sized (np.uintp) to hold any valid index."""
+    h_data = np.array([1, 2, 3, 4], dtype=np.int32)
+    h_values = np.array([2, 3], dtype=np.int32)
+    d_data = DeviceArray.from_numpy(h_data)
+    d_values = DeviceArray.from_numpy(h_values)
+    d_out = DeviceArray.empty(
+        len(h_values), np.uint32
+    )  # unsigned but not pointer-sized (on 64-bit)
+
+    with pytest.raises(ValueError, match="pointer-sized"):
+        cuda.compute.lower_bound(
+            d_data=d_data,
+            num_items=len(h_data),
+            d_values=d_values,
+            num_values=len(h_values),
+            d_out=d_out,
+        )
+
+
+@pytest.mark.serialization
+def test_serialize_deserialize_lower_bound_round_trip():
+    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 = DeviceArray.from_numpy(h_data)
+    d_values = DeviceArray.from_numpy(h_values)
+    d_out = DeviceArray.empty(len(h_values), np.uintp)
+
+    builder = make_lower_bound(d_data=d_data, d_values=d_values, d_out=d_out)
+    blob = serialize(builder)
+    assert len(blob) > 0
+
+    loaded = deserialize(blob)
+    loaded(
+        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)
+    np.testing.assert_array_equal(d_out.copy_to_host(), expected)
+
+
+@pytest.mark.serialization
+def test_serialize_deserialize_upper_bound_round_trip():
+    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 = DeviceArray.from_numpy(h_data)
+    d_values = DeviceArray.from_numpy(h_values)
+    d_out = DeviceArray.empty(len(h_values), np.uintp)
+
+    builder = make_upper_bound(d_data=d_data, d_values=d_values, d_out=d_out)
+    blob = serialize(builder)
+    assert len(blob) > 0
+
+    loaded = deserialize(blob)
+    loaded(
+        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)
+    np.testing.assert_array_equal(d_out.copy_to_host(), expected)

cccl_upstream/python/cuda_cccl/tests/compute/test_bindings.py

@@ -0,0 +1,165 @@
+import ctypes
+
+import pytest
+
+import cuda.compute._bindings as bindings
+
+
+@pytest.fixture(
+    params=[
+        "INT8",
+        "INT16",
+        "INT32",
+        "INT64",
+        "UINT8",
+        "UINT16",
+        "UINT32",
+        "UINT64",
+        "FLOAT32",
+        "FLOAT64",
+        "STORAGE",
+    ]
+)
+def cccl_type_enum(request):
+    return getattr(bindings.TypeEnum, request.param)
+
+
+@pytest.fixture(params=["STATEFUL", "STATELESS"])
+def cccl_op_kind(request):
+    return getattr(bindings.OpKind, request.param)
+
+
+@pytest.fixture(params=["POINTER", "ITERATOR"])
+def cccl_iterator_kind(request):
+    return getattr(bindings.IteratorKind, request.param)
+
+
+def test_TypeEnum_positive(cccl_type_enum):
+    assert isinstance(cccl_type_enum, bindings.TypeEnum)
+    assert isinstance(cccl_type_enum.value, int)
+
+
+def test_TypeEnum_negative(cccl_iterator_kind):
+    assert not isinstance(cccl_iterator_kind, bindings.TypeEnum)
+
+
+def test_OpKind_positive(cccl_op_kind):
+    assert isinstance(cccl_op_kind, bindings.OpKind)
+    assert isinstance(cccl_op_kind.value, int)
+
+
+def test_OpKind_negative(cccl_iterator_kind):
+    assert not isinstance(cccl_iterator_kind, bindings.OpKind)
+
+
+def test_IteratorKind_positive(cccl_iterator_kind):
+    assert isinstance(cccl_iterator_kind, bindings.IteratorKind)
+    assert isinstance(cccl_iterator_kind.value, int)
+
+
+def test_IteratorKind_negative(cccl_type_enum):
+    assert not isinstance(cccl_type_enum, bindings.IteratorKind)
+
+
+def test_Op_default():
+    res = bindings.Op()
+    assert isinstance(res, bindings.Op)
+
+
+def test_Op_state_setter():
+    res = bindings.Op()
+    bytes = b"\x00" * 20
+    res.state = bytes
+
+
+def test_Op_state_getter():
+    res = bindings.Op()
+    assert isinstance(res.state, bytes)
+
+
+def test_Op_params_stateless():
+    fake_ltoir = b"\x00" * 127
+    res = bindings.Op(
+        name="fn", operator_type=bindings.OpKind.STATELESS, ltoir=fake_ltoir
+    )
+
+    assert isinstance(res, bindings.Op)
+
+
+def test_Op_params_stateful():
+    fake_ltoir = b"\x42" * 127
+    fake_state = b"\x01" * 16
+    res = bindings.Op(
+        name="fn",
+        operator_type=bindings.OpKind.STATEFUL,
+        ltoir=fake_ltoir,
+        state=fake_state,
+        state_alignment=16,
+    )
+
+    assert isinstance(res, bindings.Op)
+    assert res.state == fake_state
+    assert res.ltoir == fake_ltoir
+
+
+def test_TypeInfo_ctor(cccl_type_enum):
+    ti = bindings.TypeInfo(4, 1, cccl_type_enum)
+    assert isinstance(ti, bindings.TypeInfo)
+
+
+def test_TypeInfo_validate(cccl_type_enum):
+    # size must positive
+    with pytest.raises(ValueError):
+        bindings.TypeInfo(0, 1, bindings.TypeEnum.FLOAT32)
+
+    # alignment must be positive, power of two
+    with pytest.raises(ValueError):
+        bindings.TypeInfo(8, 3, bindings.TypeEnum.FLOAT32)
+
+
+def test_Value_ctor():
+    ti = bindings.TypeInfo(64, 64, bindings.TypeEnum.UINT64)
+    state = bytearray(ctypes.c_uint64(2**63 + 17))
+    v = bindings.Value(ti, state)
+    assert isinstance(v, bindings.Value)
+
+
+def test_Iterator_ctor1():
+    fake_ptr = 42
+    type_info = bindings.TypeInfo(32, 32, bindings.TypeEnum.INT32)
+    cccl_it = bindings.Iterator(
+        1,  # state alignment
+        bindings.IteratorKind.POINTER,
+        bindings.Op(),
+        bindings.Op(),
+        type_info,
+        bindings.Pointer(fake_ptr),
+    )
+    assert isinstance(cccl_it, bindings.Iterator)
+
+
+def test_Iterator_ctor2():
+    fake_ptr = 42
+    type_info = bindings.TypeInfo(32, 32, bindings.TypeEnum.INT32)
+    cccl_it = bindings.Iterator(
+        1,  # state alignment
+        bindings.IteratorKind.ITERATOR,
+        bindings.Op(),
+        bindings.Op(),
+        type_info,
+        bindings.IteratorState(
+            ctypes.c_void_p(fake_ptr),
+        ),
+    )
+    assert isinstance(cccl_it, bindings.Iterator)
+
+
+def test_CommonData():
+    cub_path = "/example/path/to/cub/includes"
+    thrust_path = "/example/path/to/thrust/includes"
+    libcudacxx_path = "/example/path/to/libcudacxx/includes"
+    gtk_path = "/usr/local/cuda/include"
+    common_data = bindings.CommonData(
+        8, 6, cub_path, thrust_path, libcudacxx_path, gtk_path
+    )
+    assert isinstance(common_data, bindings.CommonData)

cccl_upstream/python/cuda_cccl/tests/compute/test_deferred_annotations.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+#
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from __future__ import annotations
+
+import numpy as np
+from _utils.device_array import DeviceArray
+
+from cuda.compute import OpKind, TransformIterator, gpu_struct, reduce_into
+
+
+def test_deferred_annotations():
+    # test that we can use @gpu_struct with deferred annotations
+    # GH: #6421
+
+    @gpu_struct
+    class MyStruct:
+        x: np.int32
+        y: np.int32
+
+
+def test_transform_iterator_future_annotations():
+    def add_one(x: "np.int32") -> "np.int32":
+        return x + np.int32(1)
+
+    h_in = np.arange(8, dtype=np.int32)
+    d_in = DeviceArray.from_numpy(h_in)
+    d_out = DeviceArray.empty(1, np.int32)
+    h_init = np.array([0], dtype=np.int32)
+
+    transform_it = TransformIterator(d_in, add_one)
+    reduce_into(
+        d_in=transform_it,
+        d_out=d_out,
+        num_items=h_in.size,
+        op=OpKind.PLUS,
+        h_init=h_init,
+    )
+
+    expected = int(np.sum(h_in + 1))
+    assert int(d_out.copy_to_host()[0]) == expected

cccl_upstream/python/cuda_cccl/tests/compute/test_func_caching.py

@@ -0,0 +1,209 @@
+import numpy as np
+
+from cuda.compute._caching import CachableFunction, _make_hashable
+
+global_x = 1
+
+
+def test_func_caching_basic():
+    def func(x):
+        return x
+
+    f1 = CachableFunction(func)
+
+    def func(x):
+        return x
+
+    f2 = CachableFunction(func)
+
+    assert f1 == f2
+
+
+def test_func_caching_different_names():
+    def func(x):
+        return x
+
+    f1 = CachableFunction(func)
+
+    def func2(x):
+        return x
+
+    f2 = CachableFunction(func2)
+
+    assert f1 != f2
+
+
+def test_func_caching_different_code():
+    def func(x):
+        return x
+
+    f1 = CachableFunction(func)
+
+    def func(x):
+        return x + 1
+
+    f2 = CachableFunction(func)
+    assert f1 != f2
+
+
+def test_func_caching_with_closure():
+    def factory(x):
+        def func(y):
+            return x + y
+
+        return func
+
+    f1 = CachableFunction(factory(1))
+    f2 = CachableFunction(factory(1))
+    assert f1 == f2
+
+    f3 = CachableFunction(factory(2))
+    assert f1 != f3
+
+
+def test_func_caching_with_numpy_numeric_scalar_closure():
+    def factory(indexlength, regularsize):
+        index_dtype = np.int64
+        idx_len = index_dtype(indexlength)
+        reg_size = index_dtype(regularsize)
+
+        def func(counter):
+            return counter % idx_len + reg_size
+
+        return func
+
+    f1 = CachableFunction(factory(100_000, 16))
+    f2 = CachableFunction(factory(100_000, 16))
+    assert f1 == f2
+
+    f3 = CachableFunction(factory(100_000, 32))
+    assert f1 != f3
+
+
+def test_make_hashable_python_scalars_keyed_by_value():
+    # Regression test for gh-9626: plain Python int/float/bool scalars used to
+    # fall through to ``id(value)``, so two equal-valued but distinct (non-
+    # interned) objects produced different cache keys and missed the cache.
+    # int(str(...)) forces fresh, non-interned objects.
+    a = int(str(10**6))
+    b = int(str(10**6))
+    assert a is not b
+    assert _make_hashable(a) == _make_hashable(b)
+
+    x = float(str(3.5))
+    y = float(str(3.5))
+    assert _make_hashable(x) == _make_hashable(y)
+
+    # Distinct values, types, and bool-vs-int must not collide.
+    assert _make_hashable(a) != _make_hashable(int(str(10**6 + 1)))
+    assert _make_hashable(1) != _make_hashable(1.0)
+    assert _make_hashable(True) != _make_hashable(1)
+
+
+def test_func_caching_with_python_scalar_closure():
+    # gh-9626: closures capturing equal-valued Python scalars must compare
+    # equal so the algorithm build cache hits instead of rebuilding every call.
+    def factory(indexlength, regularsize):
+        # int(str(...)) forces fresh, non-interned int objects.
+        idx_len = int(str(indexlength))
+        reg_size = int(str(regularsize))
+
+        def func(counter):
+            return counter % idx_len + reg_size
+
+        return func
+
+    f1 = CachableFunction(factory(100_000, 16))
+    f2 = CachableFunction(factory(100_000, 16))
+    assert f1 == f2
+
+    f3 = CachableFunction(factory(100_000, 32))
+    assert f1 != f3
+
+
+def test_func_caching_with_global_variable():
+    global global_x
+
+    def func(y):
+        return global_x + y
+
+    f1 = CachableFunction(func)
+    f2 = CachableFunction(func)
+    assert f1 == f2
+
+    global_x = 2
+    f3 = CachableFunction(func)
+    assert f1 != f3
+
+
+def test_func_caching_wrapped_cuda_jit_function():
+    import numba.cuda
+
+    def make_func():
+        @numba.cuda.jit
+        def inner(x):
+            return x
+
+        def func(x):
+            return inner(x) + 1
+
+        return func
+
+    def make_func2():
+        @numba.cuda.jit
+        def inner(x):
+            return 2 * x
+
+        def func(x):
+            return inner(x) + 1
+
+        return func
+
+    func1 = make_func()
+    func2 = make_func()
+    func3 = make_func2()
+
+    assert CachableFunction(func1) == CachableFunction(func2)
+    assert CachableFunction(func1) != CachableFunction(func3)
+
+
+def test_func_caching_with_global_np_ufunc():
+    def make_func():
+        def func(x):
+            return np.argmin(x) + 1
+
+        return func
+
+    def make_func2():
+        def func(x):
+            return np.argmax(x) + 1
+
+        return func
+
+    func1 = make_func()
+    func2 = make_func2()
+
+    assert CachableFunction(func1) != CachableFunction(func2)
+
+
+def test_func_caching_with_aliased_np_ufunc():
+    def make_func1():
+        amin = np.argmin
+
+        def func(x):
+            return amin(x) + 1
+
+        return func
+
+    def make_func2():
+        amax = np.argmax
+
+        def func(x):
+            return amax(x) + 1
+
+        return func
+
+    func1 = make_func1()
+    func2 = make_func2()
+
+    assert CachableFunction(func1) != CachableFunction(func2)

cccl_upstream/python/cuda_cccl/tests/compute/test_histogram.py

@@ -0,0 +1,614 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+import math
+
+import numpy as np
+import pytest
+from _utils.device_array import DeviceArray
+
+import cuda.compute
+from cuda.compute import (
+    ConstantIterator,
+    CountingIterator,
+    deserialize,
+    make_histogram_even,
+    serialize,
+)
+from cuda.compute._utils.temp_storage_buffer import TempStorageBuffer
+
+DTYPE_LIST = [
+    np.uint8,
+    np.uint16,
+    np.uint32,
+    np.uint64,
+    np.int8,
+    np.int16,
+    np.int32,
+    np.int64,
+    np.float16,
+    np.float32,
+    np.float64,
+]
+
+
+def get_mark(dt, log_size):
+    if log_size + np.log2(np.dtype(dt).itemsize) < 21:
+        return tuple()
+    return pytest.mark.large
+
+
+def type_to_problem_sizes(dtype):
+    if dtype in [np.uint8, np.int8]:
+        return [8, 10, 12, 14]
+    elif dtype in [np.float16, np.uint16, np.int16]:
+        return [10, 12, 14, 16]
+    elif dtype in [np.uint32, np.int32, np.float32]:
+        return [12, 14, 16, 18]
+    elif dtype in [np.uint64, np.int64, np.float64]:
+        return [12, 14, 16, 18]
+    else:
+        raise ValueError("Unsupported dtype")
+
+
+dtype_size_pairs = [
+    pytest.param(dt, 2**log_size, marks=get_mark(dt, log_size))
+    for dt in DTYPE_LIST
+    for log_size in type_to_problem_sizes(dt)
+]
+
+
+def random_int_array(size, dtype):
+    if np.issubdtype(dtype, np.integer):
+        if dtype in [np.uint8, np.int8]:
+            max_val = 126
+        else:
+            max_val = 1024
+        return np.random.randint(0, max_val, size=size).astype(dtype)
+    else:
+        # For floating point, generate values in similar range
+        return (np.random.random(size) * 1024).astype(dtype)
+
+
+def compute_reference_histogram(h_samples, num_levels, lower_level, upper_level):
+    # Filter samples within range [lower_level, upper_level)
+    valid_mask = (h_samples >= lower_level) & (h_samples < upper_level)
+    valid_samples = h_samples[valid_mask]
+
+    if len(valid_samples) == 0:
+        return np.zeros(num_levels - 1, dtype=np.int32)
+
+    # Compute bin indices for valid samples using float arithmetic
+    bin_indices = (
+        (valid_samples.astype(np.float64) - lower_level)
+        * (num_levels - 1)
+        / (upper_level - lower_level)
+    ).astype(int)
+
+    # Ensure indices are within valid range [0, num_levels-2]
+    bin_indices = np.clip(bin_indices, 0, num_levels - 2)
+
+    # Use bincount to get histogram
+    histogram = np.bincount(bin_indices, minlength=num_levels - 1)
+
+    return histogram.astype(np.int32)
+
+
+@pytest.mark.no_verify_sass(reason="LDL/STL instructions emitted for this test.")
+@pytest.mark.parametrize("dtype,num_samples", dtype_size_pairs)
+def test_device_histogram_basic_use(dtype, num_samples):
+    if dtype in [np.uint8, np.int8]:
+        max_level = 126.0
+        max_level_count = 127
+    else:
+        max_level = 1024.0
+        max_level_count = 1025
+
+    num_levels = max_level_count
+    lower_level = dtype(0.0)
+    upper_level = dtype(max_level)
+
+    h_samples = random_int_array(num_samples, dtype)
+    d_samples = DeviceArray.from_numpy(h_samples)
+
+    d_histogram = DeviceArray.from_numpy(np.zeros(num_levels - 1, dtype=np.int32))
+
+    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,
+    )
+
+    h_expected = compute_reference_histogram(
+        h_samples, num_levels, lower_level, upper_level
+    )
+    h_result = d_histogram.copy_to_host()
+
+    np.testing.assert_array_equal(h_result, h_expected)
+
+
+@pytest.mark.no_verify_sass(reason="LDL/STL instructions emitted for this test.")
+def test_device_histogram_sample_iterator():
+    max_level_count = 1025
+    num_levels = max_level_count
+    num_bins = num_levels - 1
+
+    samples_per_bin = 10
+    adjusted_total_samples = num_bins * samples_per_bin
+
+    counting_it = CountingIterator(np.int32(0))
+
+    d_histogram = DeviceArray.from_numpy(np.zeros(num_levels - 1, dtype=np.int32))
+
+    # Set up levels so that values 0 to adjusted_total_samples-1 are evenly distributed
+    lower_level = np.int32(0.0)
+    upper_level = np.int32(adjusted_total_samples)
+
+    cuda.compute.histogram_even(
+        d_samples=counting_it,
+        d_histogram=d_histogram,
+        num_output_levels=num_levels,
+        lower_level=lower_level,
+        upper_level=upper_level,
+        num_samples=adjusted_total_samples,
+    )
+
+    # Each bin should have exactly samples_per_bin elements
+    h_expected = np.full(num_bins, samples_per_bin, dtype=np.int32)
+    h_result = d_histogram.copy_to_host()
+
+    np.testing.assert_array_equal(h_result, h_expected)
+
+
+def test_device_histogram_single_sample():
+    h_samples = np.array([5.0], dtype=np.float32)
+    d_samples = DeviceArray.from_numpy(h_samples)
+
+    num_levels = 5
+    lower_level = np.float32(0.0)
+    upper_level = np.float32(10.0)
+
+    d_histogram = DeviceArray.from_numpy(np.zeros(num_levels - 1, dtype=np.int32))
+
+    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=1,
+    )
+
+    # Sample 5.0 should go into bin 2 (bins: [0,2.5), [2.5,5), [5,7.5), [7.5,10))
+    h_expected = np.array([0, 0, 1, 0], dtype=np.int32)
+    h_result = d_histogram.copy_to_host()
+
+    np.testing.assert_array_equal(h_result, h_expected)
+
+
+def test_device_histogram_out_of_range():
+    h_samples = np.array([-1.0, 0.5, 5.5, 10.5, 15.0], dtype=np.float32)
+    d_samples = DeviceArray.from_numpy(h_samples)
+
+    num_levels = 3  # 2 bins: [0,5), [5,10)
+    lower_level = np.float32(0.0)
+    upper_level = np.float32(10.0)
+
+    d_histogram = DeviceArray.from_numpy(np.zeros(num_levels - 1, dtype=np.int32))
+
+    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=len(h_samples),
+    )
+
+    # Only 0.5 (bin 0) and 5.5 (bin 1) should be counted
+    # -1.0, 10.5, and 15.0 are out of range
+    h_expected = np.array([1, 1], dtype=np.int32)
+    h_result = d_histogram.copy_to_host()
+
+    np.testing.assert_array_equal(h_result, h_expected)
+
+
+def test_device_histogram_with_stream(cuda_stream):
+    h_samples = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], dtype=np.float32)
+    d_samples = DeviceArray.from_numpy(h_samples, stream=cuda_stream)
+
+    num_levels = 5  # 4 bins: [0,2), [2,4), [4,6), [6,8)
+    lower_level = np.float32(0.0)
+    upper_level = np.float32(8.0)
+
+    d_histogram = DeviceArray.from_numpy(
+        np.zeros(num_levels - 1, dtype=np.int32), stream=cuda_stream
+    )
+
+    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=len(h_samples),
+        stream=cuda_stream,
+    )
+
+    h_result = d_histogram.copy_to_host(stream=cuda_stream)
+
+    # Expected: bin 0: [1.0, 2.0), bin 1: [2.0, 4.0), bin 2: [4.0, 6.0), bin 3: [6.0, 8.0)
+    # Values: 1.0->bin0, 2.0->bin1, 3.0->bin1, 4.0->bin2, 5.0->bin2, 6.0->bin3, 7.0->bin3, 8.0->out_of_range
+    h_expected = np.array([1, 2, 2, 2], dtype=np.int32)
+
+    np.testing.assert_array_equal(h_result, h_expected)
+
+
+@pytest.mark.no_verify_sass(reason="LDL/STL instructions emitted for this test.")
+def test_device_histogram_with_constant_iterator():
+    constant_it = ConstantIterator(np.float32(3.0))
+
+    num_samples = 10
+    num_levels = 5  # 4 bins: [0,2), [2,4), [4,6), [6,8)
+    lower_level = np.float32(0.0)
+    upper_level = np.float32(8.0)
+
+    d_histogram = DeviceArray.from_numpy(np.zeros(num_levels - 1, dtype=np.int32))
+
+    cuda.compute.histogram_even(
+        d_samples=constant_it,
+        d_histogram=d_histogram,
+        num_output_levels=num_levels,
+        lower_level=lower_level,
+        upper_level=upper_level,
+        num_samples=num_samples,
+    )
+
+    h_result = d_histogram.copy_to_host()
+
+    # Expected: All 10 samples have value 3.0, which falls in bin 1 [2,4)
+    h_expected = np.array([0, 10, 0, 0], dtype=np.int32)
+
+    np.testing.assert_array_equal(h_result, h_expected)
+
+
+def test_histogram_even():
+    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 = DeviceArray.from_numpy(h_samples)
+    num_levels = 7
+    d_histogram = DeviceArray.empty(num_levels - 1, np.int32)
+    lower_level = np.float32(0)
+    upper_level = np.float32(12)
+
+    # Run histogram with automatic temp storage allocation
+    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,
+    )
+
+    # Check the result is correct
+    h_actual_histogram = d_histogram.copy_to_host()
+    # Calculate expected histogram using numpy
+    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)
+
+
+def test_histogram_cache_bug_crosses_256_bin_threshold():
+    # GH:#7622
+    # Regression test for a bug where the histogram build artifact for
+    # num_bins <= 256 would be reused for larger bin counts, resulting
+    # in invalid shared memory accesses, because a different shared
+    # memory strategy is used for num_bins > 256.
+    num_samples = 128
+    h_num_output_levels = np.array([0], dtype=np.int32)
+    h_lower_level = np.array([0], dtype=np.int32)
+    h_upper_level = np.array([0], dtype=np.int32)
+
+    # First: 128 bins (uses shared memory, privatized_smem_bins=256)
+    num_bins_1 = 128
+    h_num_output_levels[0] = num_bins_1 + 1
+    h_lower_level[0] = 0
+    h_upper_level[0] = num_bins_1
+
+    h_samples = np.random.randint(0, num_bins_1, size=num_samples, dtype=np.int32)
+    d_samples = DeviceArray.from_numpy(h_samples)
+    d_histogram = DeviceArray.from_numpy(np.zeros(num_bins_1, dtype=np.int32))
+
+    hist = 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=num_samples,
+    )
+    temp_bytes = hist(
+        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=num_samples,
+    )
+    temp_storage = DeviceArray.empty(temp_bytes, np.uint8)
+    hist(
+        temp_storage=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=num_samples,
+    )
+    assert int(d_histogram.copy_to_host().sum()) == num_samples
+
+    num_bins_2 = 2048
+    h_num_output_levels[0] = num_bins_2 + 1
+    h_lower_level[0] = 0
+    h_upper_level[0] = num_bins_2
+
+    h_samples = np.random.randint(0, num_bins_2, size=num_samples, dtype=np.int32)
+    d_samples = DeviceArray.from_numpy(h_samples)
+    d_histogram = DeviceArray.from_numpy(np.zeros(num_bins_2, dtype=np.int32))
+
+    hist2 = 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=num_samples,
+    )
+
+    temp_bytes2 = hist2(
+        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=num_samples,
+    )
+    temp_storage2 = DeviceArray.empty(temp_bytes2, np.uint8)
+    hist2(
+        temp_storage=temp_storage2,
+        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=num_samples,
+    )
+    assert int(d_histogram.copy_to_host().sum()) == num_samples
+
+
+def test_histogram_cache_reuses_artifact_when_bounds_change():
+    cuda.compute.clear_all_caches()
+
+    num_samples = 8
+    num_levels = 5
+    d_samples = DeviceArray.from_numpy(np.arange(num_samples, dtype=np.float32))
+    d_histogram = DeviceArray.empty(num_levels - 1, np.int32)
+    h_num_output_levels = np.array([num_levels], dtype=np.int32)
+
+    h_lower_level_1 = np.array([0], dtype=np.float32)
+    h_upper_level_1 = np.array([8], dtype=np.float32)
+    h_lower_level_2 = np.array([10], dtype=np.float32)
+    h_upper_level_2 = np.array([18], dtype=np.float32)
+
+    hist1 = 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_1,
+        h_upper_level=h_upper_level_1,
+        num_samples=num_samples,
+    )
+    hist2 = 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_2,
+        h_upper_level=h_upper_level_2,
+        num_samples=num_samples,
+    )
+    assert hist1 is hist2
+
+    d_samples = DeviceArray.from_numpy(np.arange(10, 18, dtype=np.float32))
+    d_histogram.copy_from_host(np.zeros(num_levels - 1, dtype=np.int32))
+    temp_bytes = hist2(
+        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_2,
+        h_upper_level=h_upper_level_2,
+        num_samples=num_samples,
+    )
+    temp_storage = DeviceArray.empty(temp_bytes, np.uint8)
+    hist2(
+        temp_storage=temp_storage,
+        d_samples=d_samples,
+        d_histogram=d_histogram,
+        h_num_output_levels=h_num_output_levels,
+        h_lower_level=h_lower_level_2,
+        h_upper_level=h_upper_level_2,
+        num_samples=num_samples,
+    )
+    np.testing.assert_array_equal(
+        d_histogram.copy_to_host(), np.array([2, 2, 2, 2], dtype=np.int32)
+    )
+
+
+def test_histogram_cache_reuses_artifact_for_same_offset_width():
+    cuda.compute.clear_all_caches()
+
+    num_levels = 5
+    d_histogram = DeviceArray.empty(num_levels - 1, np.int32)
+    h_num_output_levels = np.array([num_levels], dtype=np.int32)
+    h_lower_level = np.array([0], dtype=np.float32)
+    h_upper_level = np.array([12], dtype=np.float32)
+
+    hist1 = cuda.compute.make_histogram_even(
+        d_samples=DeviceArray.from_numpy(np.arange(8, dtype=np.float32)),
+        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=8,
+    )
+    hist2 = cuda.compute.make_histogram_even(
+        d_samples=DeviceArray.from_numpy(np.arange(12, dtype=np.float32)),
+        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=12,
+    )
+    assert hist1 is hist2
+
+    # Build-only call that crosses v1's offset-width threshold without
+    # allocating or executing that many samples.
+    large_num_samples = math.ceil(
+        np.iinfo(np.int32).max / np.dtype(np.float32).itemsize
+    )
+    hist3 = cuda.compute.make_histogram_even(
+        d_samples=DeviceArray.from_numpy(np.arange(12, dtype=np.float32)),
+        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=large_num_samples,
+    )
+    assert hist3 is not hist1
+
+    d_samples = DeviceArray.from_numpy(np.arange(12, dtype=np.float32))
+    d_histogram.copy_from_host(np.zeros(num_levels - 1, dtype=np.int32))
+    temp_bytes = hist2(
+        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=12,
+    )
+    temp_storage = DeviceArray.empty(temp_bytes, np.uint8)
+    hist2(
+        temp_storage=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=12,
+    )
+    np.testing.assert_array_equal(
+        d_histogram.copy_to_host(), np.array([3, 3, 3, 3], dtype=np.int32)
+    )
+
+
+def test_make_histogram_even_rejects_mismatched_bound_dtypes():
+    num_samples = 8
+    d_samples = DeviceArray.from_numpy(np.arange(num_samples, dtype=np.int32))
+    d_histogram = DeviceArray.empty(4, np.int32)
+
+    with pytest.raises(TypeError, match="must have the same dtype"):
+        cuda.compute.make_histogram_even(
+            d_samples=d_samples,
+            d_histogram=d_histogram,
+            h_num_output_levels=np.array([5], dtype=np.int32),
+            h_lower_level=np.array([0], dtype=np.int32),
+            h_upper_level=np.array([8], dtype=np.float32),
+            num_samples=num_samples,
+        )
+
+
+def _run(
+    histogram,
+    *,
+    d_samples,
+    d_histogram,
+    h_num_output_levels,
+    h_lower_level,
+    h_upper_level,
+    num_samples,
+):
+    bytes_needed = histogram(
+        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=num_samples,
+    )
+    tmp = TempStorageBuffer(bytes_needed, None)
+    histogram(
+        temp_storage=tmp,
+        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=num_samples,
+    )
+
+
+@pytest.mark.serialization
+def test_serialize_deserialize_histogram_even_round_trip():
+    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 = DeviceArray.from_numpy(h_samples)
+    num_levels = 7
+    d_histogram = DeviceArray.empty(num_levels - 1, "int32")
+    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([12.0], dtype=np.float32)
+
+    builder = 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=num_samples,
+    )
+    blob = serialize(builder)
+    assert len(blob) > 0
+
+    loaded = deserialize(blob)
+    _run(
+        loaded,
+        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=num_samples,
+    )
+
+    expected, _ = np.histogram(
+        h_samples,
+        bins=num_levels - 1,
+        range=(float(h_lower_level[0]), float(h_upper_level[0])),
+    )
+    np.testing.assert_array_equal(d_histogram.copy_to_host(), expected.astype(np.int32))

cccl_upstream/python/cuda_cccl/tests/compute/test_iterators.py

@@ -0,0 +1,223 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+#
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+import numpy as np
+import pytest
+from _utils.device_array import DeviceArray
+
+import cuda.compute
+from cuda.compute import OpKind
+from cuda.compute._utils.protocols import (
+    compute_c_contiguous_strides_in_bytes,
+)
+from cuda.compute.iterators import (
+    CacheModifiedInputIterator,
+    ConstantIterator,
+    CountingIterator,
+    ReverseIterator,
+    TransformIterator,
+)
+
+
+def test_constant_iterator_equality():
+    it1 = ConstantIterator(np.int32(0))
+    it2 = ConstantIterator(np.int32(0))
+    it3 = ConstantIterator(np.int32(1))
+    it4 = ConstantIterator(np.int64(0))
+
+    assert it1.kind == it2.kind == it3.kind
+    assert it1.kind != it4.kind
+
+
+def test_counting_iterator_equality():
+    it1 = CountingIterator(np.int32(0))
+    it2 = CountingIterator(np.int32(0))
+    it3 = CountingIterator(np.int32(1))
+    it4 = CountingIterator(np.int64(0))
+
+    assert it1.kind == it2.kind == it3.kind
+    assert it1.kind != it4.kind
+
+
+def test_cache_modified_input_iterator_equality():
+    ary1 = DeviceArray.from_numpy(np.asarray([0, 1, 2], dtype="int32"))
+    ary2 = DeviceArray.from_numpy(np.asarray([3, 4, 5], dtype="int32"))
+    ary3 = DeviceArray.from_numpy(np.asarray([0, 1, 2], dtype="int64"))
+
+    it1 = CacheModifiedInputIterator(ary1, "stream")
+    it2 = CacheModifiedInputIterator(ary1, "stream")
+    it3 = CacheModifiedInputIterator(ary2, "stream")
+    it4 = CacheModifiedInputIterator(ary3, "stream")
+
+    assert it1.kind == it2.kind == it3.kind
+    assert it1.kind != it4.kind
+
+
+def test_equality_transform_iterator():
+    def op1(x):
+        return x
+
+    def op2(x):
+        return 2 * x
+
+    def op3(x):
+        return x
+
+    it = CountingIterator(np.int32(0))
+    it = CountingIterator(np.int32(1))
+    it1 = TransformIterator(it, op1)
+    it2 = TransformIterator(it, op1)
+    it3 = TransformIterator(it, op3)
+
+    assert it1.kind == it2.kind
+    # op3 has a different name than op1, so should have a different kind
+    assert it1.kind != it3.kind
+
+    ary1 = DeviceArray.from_numpy(np.asarray([0, 1, 2]))
+    ary2 = DeviceArray.from_numpy(np.asarray([3, 4, 5]))
+    it4 = TransformIterator(ary1, op1)
+    it5 = TransformIterator(ary1, op1)
+    it6 = TransformIterator(ary1, op2)
+    it7 = TransformIterator(ary1, op3)
+    it8 = TransformIterator(ary2, op1)
+
+    assert it4.kind == it5.kind == it8.kind
+    # op2 has different bytecode, so should have a different kind
+    assert it4.kind != it6.kind
+    # op3 has a different name than op1, so should have a different kind
+    assert it4.kind != it7.kind
+
+
+def test_reverse_input_iterator_equality():
+    ary1 = DeviceArray.from_numpy(np.asarray([0, 1, 2], dtype="int32"))
+    ary2 = DeviceArray.from_numpy(np.asarray([3, 4, 5], dtype="int32"))
+    ary3 = DeviceArray.from_numpy(np.asarray([0, 1, 2], dtype="int64"))
+
+    it1 = ReverseIterator(ary1)
+    it2 = ReverseIterator(ary1)
+    it3 = ReverseIterator(ary2)
+    it4 = ReverseIterator(ary3)
+
+    assert it1.kind == it2.kind == it3.kind
+    assert it1.kind != it4.kind
+
+
+def test_reverse_output_iterator_equality():
+    ary1 = DeviceArray.from_numpy(np.asarray([0, 1, 2], dtype="int32"))
+    ary2 = DeviceArray.from_numpy(np.asarray([3, 4, 5], dtype="int32"))
+    ary3 = DeviceArray.from_numpy(np.asarray([0, 1, 2], dtype="int64"))
+
+    it1 = ReverseIterator(ary1)
+    it2 = ReverseIterator(ary1)
+    it3 = ReverseIterator(ary2)
+    it4 = ReverseIterator(ary3)
+
+    assert it1.kind == it2.kind == it3.kind
+    assert it1.kind != it4.kind
+
+
+@pytest.mark.parametrize(
+    "shape, itemsize, expected",
+    [
+        # Basic 1D
+        ((5,), 4, (4,)),
+        ((10,), 1, (1,)),
+        # Basic 2D
+        ((2, 3), 4, (12, 4)),
+        ((3, 2), 8, (16, 8)),
+        # Basic 3D
+        ((2, 3, 4), 1, (12, 4, 1)),
+        ((2, 3, 4), 2, (24, 8, 2)),
+        # Scalars (0D array)
+        ((), 4, ()),
+        # Shape with a zero-length dimension
+        ((0, 3), 4, (12, 4)),
+        ((3, 0), 4, (0, 4)),
+    ],
+)
+def test_compute_c_contiguous_strides_in_bytes(shape, itemsize, expected):
+    result = compute_c_contiguous_strides_in_bytes(shape, itemsize)
+    assert result == expected
+
+
+@pytest.mark.parametrize(
+    "shape, dtype",
+    [
+        ((2, 3), np.int32),
+        ((4, 5, 6), np.float64),
+        ((10,), np.uint8),
+        ((1,), np.float16),
+    ],
+)
+def test_matches_numpy_strides_for_c_contiguous_arrays(shape, dtype):
+    arr = np.zeros(shape, dtype=dtype, order="C")
+    expected = arr.strides
+    result = compute_c_contiguous_strides_in_bytes(shape, dtype().itemsize)
+    assert result == expected
+
+
+def test_transform_iterator_with_lambda():
+    """Test TransformIterator with a lambda function."""
+    first_item = 10
+    num_items = 100
+
+    # Use a lambda function directly with TransformIterator
+    transform_it = TransformIterator(
+        CountingIterator(np.int32(first_item)), lambda x: x * 2
+    )
+    h_init = np.array([0], dtype=np.int32)
+    d_output = DeviceArray.empty(1, np.int32)
+
+    # Perform reduction on the transformed iterator
+    cuda.compute.reduce_into(
+        d_in=transform_it,
+        d_out=d_output,
+        num_items=num_items,
+        op=OpKind.PLUS,
+        h_init=h_init,
+    )
+
+    # Expected: sum of (10*2, 11*2, ..., 109*2) = 2 * sum(10..109)
+    expected = 2 * sum(range(first_item, first_item + num_items))
+    assert d_output.copy_to_host()[0] == expected
+
+
+def test_transform_iterator_with_zip_iterator():
+    """Test TransformIterator wrapping ZipIterator (struct types)."""
+    from cuda.compute.iterators import ZipIterator
+
+    # Create a ZipIterator with two int32 arrays
+    h_a = np.arange(10, dtype=np.int32)
+    h_b = np.arange(100, 110, dtype=np.int32)
+    d_a = DeviceArray.from_numpy(h_a)
+    d_b = DeviceArray.from_numpy(h_b)
+
+    zip_it = ZipIterator(d_a, d_b)
+
+    # Create a transform that sums the two fields
+    # Input is a struct with two int32 fields, output is a single int32
+    def sum_fields(pair):
+        return pair[0] + pair[1]
+
+    # Create TransformIterator wrapping ZipIterator
+    # This tests that cpp_type_from_descriptor handles struct types correctly
+    transform_it = TransformIterator(zip_it, sum_fields)
+
+    # Use it in a reduction
+    h_init = np.array([0], dtype=np.int32)
+    d_output = DeviceArray.empty(1, np.int32)
+
+    cuda.compute.reduce_into(
+        d_in=transform_it,
+        d_out=d_output,
+        num_items=len(h_a),
+        op=OpKind.PLUS,
+        h_init=h_init,
+    )
+
+    result = d_output.copy_to_host()[0]
+    expected = (h_a + h_b).sum()
+
+    assert result == expected, f"Expected {expected}, got {result}"

cccl_upstream/python/cuda_cccl/tests/compute/test_merge_sort.py

@@ -0,0 +1,496 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from typing import List
+
+import numpy as np
+import pytest
+from _utils.device_array import DeviceArray
+
+import cuda.compute
+from cuda.compute import (
+    CacheModifiedInputIterator,
+    OpKind,
+    deserialize,
+    gpu_struct,
+    make_merge_sort,
+    serialize,
+)
+from cuda.compute._utils.temp_storage_buffer import TempStorageBuffer
+
+DTYPE_LIST = [
+    np.uint8,
+    np.uint16,
+    np.uint32,
+    np.uint64,
+    np.int8,
+    np.int16,
+    np.int32,
+    np.int64,
+    np.float16,
+    np.float32,
+    np.float64,
+]
+
+
+def random_array(size, dtype, max_value=None) -> np.typing.NDArray:
+    rng = np.random.default_rng()
+    if np.isdtype(dtype, "integral"):
+        if max_value is None:
+            max_value = np.iinfo(dtype).max
+        return rng.integers(max_value, size=size, dtype=dtype)
+    elif np.isdtype(dtype, "real floating"):
+        if dtype == np.float16:  # Cannot generate float16 directly
+            return rng.random(size=size, dtype=np.float32).astype(dtype)
+        else:
+            return rng.random(size=size, dtype=dtype)
+    else:
+        raise ValueError(f"Unsupported dtype {dtype}")
+
+
+def type_to_problem_sizes(dtype) -> List[int]:
+    if dtype in DTYPE_LIST:
+        return [2, 4, 6, 8, 10, 16, 20]
+    else:
+        raise ValueError("Unsupported dtype")
+
+
+def merge_sort_device(
+    d_in_keys, d_in_items, d_out_keys, d_out_items, op, num_items, stream=None
+):
+    cuda.compute.merge_sort(
+        d_in_keys=d_in_keys,
+        d_in_values=d_in_items,
+        d_out_keys=d_out_keys,
+        d_out_values=d_out_items,
+        num_items=num_items,
+        op=op,
+        stream=stream,
+    )
+
+
+def compare_op(lhs, rhs):
+    return np.uint8(lhs < rhs)
+
+
+merge_sort_params = [
+    (dt, 2**log_size, OpKind.LESS if dt == np.float16 else compare_op)
+    for dt in DTYPE_LIST
+    for log_size in type_to_problem_sizes(dt)
+]
+
+
+@pytest.mark.parametrize("dtype,num_items,op", merge_sort_params)
+def test_merge_sort_keys(dtype, num_items, op):
+    h_in_keys = random_array(num_items, dtype)
+
+    d_in_keys = DeviceArray.from_numpy(h_in_keys)
+
+    merge_sort_device(d_in_keys, None, d_in_keys, None, op, num_items)
+
+    h_out_keys = d_in_keys.copy_to_host()
+    h_in_keys.sort()
+
+    np.testing.assert_array_equal(h_out_keys, h_in_keys)
+
+
+@pytest.mark.parametrize("dtype,num_items,op", merge_sort_params)
+def test_merge_sort_pairs(dtype, num_items, op, monkeypatch):
+    if dtype == np.float16:
+        import cuda.compute._cccl_interop
+
+        monkeypatch.setattr(cuda.compute._cccl_interop, "_check_sass", False)
+
+    h_in_keys = random_array(num_items, dtype)
+    h_in_items = random_array(num_items, np.float32)
+
+    d_in_keys = DeviceArray.from_numpy(h_in_keys)
+    d_in_items = DeviceArray.from_numpy(h_in_items)
+
+    merge_sort_device(d_in_keys, d_in_items, d_in_keys, d_in_items, op, num_items)
+
+    h_out_keys = d_in_keys.copy_to_host()
+    h_out_items = d_in_items.copy_to_host()
+
+    argsort = np.argsort(h_in_keys, stable=True)
+    h_in_keys = np.array(h_in_keys)[argsort]
+    h_in_items = np.array(h_in_items)[argsort]
+
+    np.testing.assert_array_equal(h_out_keys, h_in_keys)
+    np.testing.assert_array_equal(h_out_items, h_in_items)
+
+
+@pytest.mark.parametrize("dtype,num_items,op", merge_sort_params)
+def test_merge_sort_keys_copy(dtype, num_items, op):
+    h_in_keys = random_array(num_items, dtype)
+    h_out_keys = np.empty(num_items, dtype=dtype)
+
+    d_in_keys = DeviceArray.from_numpy(h_in_keys)
+    d_out_keys = DeviceArray.empty(h_out_keys.shape, h_out_keys.dtype)
+
+    merge_sort_device(d_in_keys, None, d_out_keys, None, op, num_items)
+
+    h_out_keys = d_out_keys.copy_to_host()
+    h_in_keys.sort()
+
+    np.testing.assert_array_equal(h_out_keys, h_in_keys)
+
+
+@pytest.mark.parametrize("dtype,num_items,op", merge_sort_params)
+def test_merge_sort_pairs_copy(dtype, num_items, op, monkeypatch):
+    if dtype == np.float16:
+        import cuda.compute._cccl_interop
+
+        monkeypatch.setattr(cuda.compute._cccl_interop, "_check_sass", False)
+
+    h_in_keys = random_array(num_items, dtype)
+    h_in_items = random_array(num_items, np.float32)
+    h_out_keys = np.empty(num_items, dtype=dtype)
+    h_out_items = np.empty(num_items, dtype=np.float32)
+
+    d_in_keys = DeviceArray.from_numpy(h_in_keys)
+    d_in_items = DeviceArray.from_numpy(h_in_items)
+    d_out_keys = DeviceArray.empty(h_out_keys.shape, h_out_keys.dtype)
+    d_out_items = DeviceArray.empty(h_out_items.shape, h_out_items.dtype)
+
+    merge_sort_device(d_in_keys, d_in_items, d_out_keys, d_out_items, op, num_items)
+
+    h_out_keys = d_out_keys.copy_to_host()
+    h_out_items = d_out_items.copy_to_host()
+
+    argsort = np.argsort(h_in_keys, stable=True)
+    h_in_keys = np.array(h_in_keys)[argsort]
+    h_in_items = np.array(h_in_items)[argsort]
+
+    np.testing.assert_array_equal(h_out_keys, h_in_keys)
+    np.testing.assert_array_equal(h_out_items, h_in_items)
+
+
+def test_merge_sort_pairs_struct_type():
+    @gpu_struct
+    class key_pair:
+        a: np.int16
+        b: np.uint64
+
+    @gpu_struct
+    class item_pair:
+        a: np.int32
+        b: np.float32
+
+    def struct_compare_op(lhs, rhs):
+        return np.uint8(lhs.b < rhs.b) if lhs.a == rhs.a else np.uint8(lhs.a < rhs.a)
+
+    num_items = 1000
+
+    a_keys = np.random.randint(0, 100, num_items, dtype=np.int16)
+    b_keys = np.random.randint(0, 100, num_items, dtype=np.uint64)
+
+    a_items = np.random.randint(0, 100, num_items, dtype=np.int32)
+    b_items = np.random.rand(num_items).astype(np.float32)
+
+    h_in_keys = np.empty(num_items, dtype=key_pair.dtype)
+    h_in_items = np.empty(num_items, dtype=item_pair.dtype)
+
+    h_in_keys["a"] = a_keys
+    h_in_keys["b"] = b_keys
+
+    h_in_items["a"] = a_items
+    h_in_items["b"] = b_items
+
+    d_in_keys = DeviceArray.from_numpy(h_in_keys)
+    d_in_items = DeviceArray.from_numpy(h_in_items)
+
+    merge_sort_device(
+        d_in_keys, d_in_items, d_in_keys, d_in_items, struct_compare_op, num_items
+    )
+
+    h_out_keys = d_in_keys.copy_to_host()
+    h_out_items = d_in_items.copy_to_host()
+
+    argsort = np.argsort(h_in_keys, stable=True)
+    h_in_keys = np.array(h_in_keys)[argsort]
+    h_in_items = np.array(h_in_items)[argsort]
+
+    np.testing.assert_array_equal(h_out_keys, h_in_keys)
+    np.testing.assert_array_equal(h_out_items, h_in_items)
+
+
+def test_merge_sort_keys_complex():
+    def compare_complex(lhs, rhs):
+        return np.uint8(lhs.real < rhs.real)
+
+    num_items = 10000
+    max_value = 20  # To ensure that the stability property is being tested
+    real = random_array(num_items, np.int64, max_value)
+    imaginary = random_array(num_items, np.int64, max_value)
+
+    h_in_keys = real + 1j * imaginary
+    d_in_keys = DeviceArray.from_numpy(h_in_keys)
+
+    merge_sort_device(d_in_keys, None, d_in_keys, None, compare_complex, num_items)
+
+    h_out_keys = d_in_keys.copy_to_host()
+    h_in_keys = h_in_keys[np.argsort(h_in_keys.real, stable=True)]
+
+    np.testing.assert_array_equal(h_out_keys, h_in_keys)
+
+
+@pytest.mark.parametrize("dtype,num_items,op", merge_sort_params)
+def test_merge_sort_keys_copy_iterator_input(dtype, num_items, op):
+    h_in_keys = random_array(num_items, dtype)
+    h_out_keys = np.empty(num_items, dtype=dtype)
+
+    d_in_keys = DeviceArray.from_numpy(h_in_keys)
+    d_out_keys = DeviceArray.empty(h_out_keys.shape, h_out_keys.dtype)
+
+    i_input = CacheModifiedInputIterator(d_in_keys, modifier="stream")
+
+    merge_sort_device(i_input, None, d_out_keys, None, op, num_items)
+
+    h_in_keys.sort()
+    h_out_keys = d_out_keys.copy_to_host()
+
+    np.testing.assert_array_equal(h_out_keys, h_in_keys)
+
+
+@pytest.mark.parametrize("dtype,num_items,op", merge_sort_params)
+def test_merge_sort_pairs_copy_iterator_input(dtype, num_items, op, monkeypatch):
+    if dtype == np.float16:
+        import cuda.compute._cccl_interop
+
+        monkeypatch.setattr(cuda.compute._cccl_interop, "_check_sass", False)
+
+    h_in_keys = random_array(num_items, dtype)
+    h_in_items = random_array(num_items, np.float32)
+    h_out_keys = np.empty(num_items, dtype=dtype)
+    h_out_items = np.empty(num_items, dtype=np.float32)
+
+    d_in_keys = DeviceArray.from_numpy(h_in_keys)
+    d_in_items = DeviceArray.from_numpy(h_in_items)
+    d_out_keys = DeviceArray.empty(h_out_keys.shape, h_out_keys.dtype)
+    d_out_items = DeviceArray.empty(h_out_items.shape, h_out_items.dtype)
+
+    i_input_keys = CacheModifiedInputIterator(d_in_keys, modifier="stream")
+    i_input_items = CacheModifiedInputIterator(d_in_items, modifier="stream")
+
+    merge_sort_device(
+        i_input_keys, i_input_items, d_out_keys, d_out_items, op, num_items
+    )
+
+    h_out_keys = d_out_keys.copy_to_host()
+    h_out_items = d_out_items.copy_to_host()
+
+    argsort = np.argsort(h_in_keys, stable=True)
+    h_in_keys = np.array(h_in_keys)[argsort]
+    h_in_items = np.array(h_in_items)[argsort]
+
+    np.testing.assert_array_equal(h_out_keys, h_in_keys)
+    np.testing.assert_array_equal(h_out_items, h_in_items)
+
+
+def test_merge_sort_with_stream(cuda_stream):
+    num_items = 10000
+
+    h_in_keys = random_array(num_items, np.int32)
+    d_in_keys = DeviceArray.from_numpy(h_in_keys, stream=cuda_stream)
+    d_out_keys = DeviceArray.empty(h_in_keys.shape, h_in_keys.dtype, stream=cuda_stream)
+
+    merge_sort_device(
+        d_in_keys, None, d_out_keys, None, compare_op, num_items, stream=cuda_stream
+    )
+
+    got = d_out_keys.copy_to_host(stream=cuda_stream)
+    h_in_keys.sort()
+
+    np.testing.assert_array_equal(got, h_in_keys)
+
+
+def test_merge_sort_well_known_less():
+    dtype = np.int32
+
+    h_in_keys = np.array([5, 2, 8, 1, 9, 3], dtype=dtype)
+    d_in_keys = DeviceArray.from_numpy(h_in_keys)
+    d_out_keys = DeviceArray.empty(h_in_keys.shape, h_in_keys.dtype)
+
+    cuda.compute.merge_sort(
+        d_in_keys=d_in_keys,
+        d_in_values=None,
+        d_out_keys=d_out_keys,
+        d_out_values=None,
+        num_items=len(h_in_keys),
+        op=OpKind.LESS,
+    )
+
+    expected = np.array([1, 2, 3, 5, 8, 9])
+    np.testing.assert_equal(d_out_keys.copy_to_host(), expected)
+
+
+def test_merge_sort_well_known_greater():
+    dtype = np.int32
+
+    h_in_keys = np.array([5, 2, 8, 1, 9, 3], dtype=dtype)
+    d_in_keys = DeviceArray.from_numpy(h_in_keys)
+    d_out_keys = DeviceArray.empty(h_in_keys.shape, h_in_keys.dtype)
+
+    cuda.compute.merge_sort(
+        d_in_keys=d_in_keys,
+        d_in_values=None,
+        d_out_keys=d_out_keys,
+        d_out_values=None,
+        num_items=len(h_in_keys),
+        op=OpKind.GREATER,
+    )
+
+    expected = np.array([9, 8, 5, 3, 2, 1])
+    np.testing.assert_equal(d_out_keys.copy_to_host(), expected)
+
+
+def test_merge_sort_large_temp_storage_not_negative():
+    """Regression test for https://github.com/NVIDIA/cccl/issues/7911.
+
+    temp_storage_bytes was returned as a signed 32-bit int, overflowing
+    to a negative value for large inputs requiring >2GB temp storage.
+    """
+    num_items = 2**28
+    dtype = np.int64
+    d_in_keys = DeviceArray.empty(num_items, dtype)
+    d_out_keys = DeviceArray.empty(num_items, dtype)
+
+    sorter = cuda.compute.make_merge_sort(
+        d_in_keys=d_in_keys,
+        d_in_values=None,
+        d_out_keys=d_out_keys,
+        d_out_values=None,
+        op=OpKind.LESS,
+    )
+
+    temp_storage_bytes = sorter(
+        temp_storage=None,
+        d_in_keys=d_in_keys,
+        d_in_values=None,
+        d_out_keys=d_out_keys,
+        d_out_values=None,
+        op=OpKind.LESS,
+        num_items=num_items,
+    )
+
+    assert temp_storage_bytes > 0
+
+
+def test_merge_sort_with_values_well_known():
+    dtype = np.int32
+
+    h_in_keys = np.array([3, 1, 4, 2], dtype=dtype)
+    h_in_values = np.array([30, 10, 40, 20], dtype=dtype)
+    d_in_keys = DeviceArray.from_numpy(h_in_keys)
+    d_in_values = DeviceArray.from_numpy(h_in_values)
+    d_out_keys = DeviceArray.empty(h_in_keys.shape, h_in_keys.dtype)
+    d_out_values = DeviceArray.empty(h_in_values.shape, h_in_values.dtype)
+
+    cuda.compute.merge_sort(
+        d_in_keys=d_in_keys,
+        d_in_values=d_in_values,
+        d_out_keys=d_out_keys,
+        d_out_values=d_out_values,
+        num_items=len(h_in_keys),
+        op=OpKind.LESS,
+    )
+
+    expected_keys = np.array([1, 2, 3, 4])
+    expected_values = np.array([10, 20, 30, 40])
+    np.testing.assert_equal(d_out_keys.copy_to_host(), expected_keys)
+    np.testing.assert_equal(d_out_values.copy_to_host(), expected_values)
+
+
+def _run(sorter, *, d_in_keys, d_in_values, d_out_keys, d_out_values, num_items, op):
+    bytes_needed = sorter(
+        temp_storage=None,
+        d_in_keys=d_in_keys,
+        d_in_values=d_in_values,
+        d_out_keys=d_out_keys,
+        d_out_values=d_out_values,
+        num_items=num_items,
+        op=op,
+    )
+    tmp = TempStorageBuffer(bytes_needed, None)
+    sorter(
+        temp_storage=tmp,
+        d_in_keys=d_in_keys,
+        d_in_values=d_in_values,
+        d_out_keys=d_out_keys,
+        d_out_values=d_out_values,
+        num_items=num_items,
+        op=op,
+    )
+
+
+@pytest.mark.serialization
+def test_serialize_deserialize_merge_sort_keys_values():
+    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 = DeviceArray.from_numpy(h_in_keys)
+    d_in_values = DeviceArray.from_numpy(h_in_values)
+    d_out_keys = DeviceArray.empty(h_in_keys.shape, h_in_keys.dtype)
+    d_out_values = DeviceArray.empty(h_in_values.shape, h_in_values.dtype)
+
+    builder = make_merge_sort(
+        d_in_keys=d_in_keys,
+        d_in_values=d_in_values,
+        d_out_keys=d_out_keys,
+        d_out_values=d_out_values,
+        op=OpKind.LESS,
+    )
+    blob = serialize(builder)
+    assert len(blob) > 0
+
+    loaded = deserialize(blob)
+    _run(
+        loaded,
+        d_in_keys=d_in_keys,
+        d_in_values=d_in_values,
+        d_out_keys=d_out_keys,
+        d_out_values=d_out_values,
+        num_items=h_in_keys.size,
+        op=OpKind.LESS,
+    )
+
+    # kind="stable" works on all supported NumPy versions; the stable= keyword
+    # was only added in NumPy 2.0 and cuda-cccl pins no numpy floor.
+    argsort = np.argsort(h_in_keys, kind="stable")
+    np.testing.assert_array_equal(d_out_keys.copy_to_host(), h_in_keys[argsort])
+    np.testing.assert_array_equal(d_out_values.copy_to_host(), h_in_values[argsort])
+
+
+@pytest.mark.serialization
+def test_serialize_deserialize_merge_sort_keys_only():
+    # Keys-only: d_in_values / d_out_values are None, which become "none"
+    # iterators — the plain ITER schema members round-trip them fine.
+    h_in_keys = np.array([5, 2, 8, 1, 9, 3, 7, 0, 6, 4], dtype="int32")
+    d_in_keys = DeviceArray.from_numpy(h_in_keys)
+    d_out_keys = DeviceArray.empty(h_in_keys.shape, h_in_keys.dtype)
+
+    builder = make_merge_sort(
+        d_in_keys=d_in_keys,
+        d_in_values=None,
+        d_out_keys=d_out_keys,
+        d_out_values=None,
+        op=OpKind.LESS,
+    )
+    blob = serialize(builder)
+    assert len(blob) > 0
+
+    loaded = deserialize(blob)
+    _run(
+        loaded,
+        d_in_keys=d_in_keys,
+        d_in_values=None,
+        d_out_keys=d_out_keys,
+        d_out_values=None,
+        num_items=h_in_keys.size,
+        op=OpKind.LESS,
+    )
+
+    np.testing.assert_array_equal(d_out_keys.copy_to_host(), np.sort(h_in_keys))

cccl_upstream/python/cuda_cccl/tests/compute/test_nested_struct.py

@@ -0,0 +1,499 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+import numpy as np
+from _utils.device_array import DeviceArray
+
+import cuda.compute
+from cuda.compute import ZipIterator, gpu_struct
+
+
+def test_reduce_nested_struct_direct():
+    Inner = gpu_struct({"a": np.int32, "b": np.float32})
+    Outer = gpu_struct({"x": np.int64, "inner": Inner})
+
+    def sum_nested(s1, s2):
+        return Outer(
+            s1.x + s2.x, Inner(s1.inner.a + s2.inner.a, s1.inner.b + s2.inner.b)
+        )
+
+    num_items = 10
+
+    h_data = np.zeros(num_items, dtype=Outer.dtype)
+    for i in range(num_items):
+        h_data[i]["x"] = i
+        h_data[i]["inner"]["a"] = i * 2
+        h_data[i]["inner"]["b"] = float(i * 3)
+
+    d_input = DeviceArray.from_numpy(h_data)
+    d_output = DeviceArray.empty(1, Outer.dtype)
+
+    h_init = Outer(0, Inner(0, 0.0))
+
+    cuda.compute.reduce_into(
+        d_in=d_input, d_out=d_output, num_items=num_items, op=sum_nested, h_init=h_init
+    )
+
+    result = d_output.copy_to_host()[0]
+
+    expected_x = sum(range(num_items))
+    expected_a = sum(i * 2 for i in range(num_items))
+    expected_b = sum(float(i * 3) for i in range(num_items))
+
+    assert result["x"] == expected_x
+    assert result["inner"]["a"] == expected_a
+    assert np.isclose(result["inner"]["b"], expected_b)
+
+
+def test_nested_struct_inline():
+    """Test creating nested structs using inline dictionary syntax."""
+    # Create a struct with an inline nested struct definition
+    Outer = gpu_struct({"x": np.int64, "inner": {"a": np.int32, "b": np.float32}})
+
+    # Get the nested struct type from the outer struct for construction
+    Inner = type(Outer(0, (0, 0.0)).inner)
+
+    def sum_nested(s1, s2):
+        return Outer(
+            s1.x + s2.x, Inner(s1.inner.a + s2.inner.a, s1.inner.b + s2.inner.b)
+        )
+
+    num_items = 10
+
+    h_data = np.zeros(num_items, dtype=Outer.dtype)
+    for i in range(num_items):
+        h_data[i]["x"] = i
+        h_data[i]["inner"]["a"] = i * 2
+        h_data[i]["inner"]["b"] = float(i * 3)
+
+    d_input = DeviceArray.from_numpy(h_data)
+    d_output = DeviceArray.empty(1, Outer.dtype)
+
+    h_init = Outer(0, Inner(0, 0.0))
+
+    cuda.compute.reduce_into(
+        d_in=d_input, d_out=d_output, num_items=num_items, op=sum_nested, h_init=h_init
+    )
+
+    result = d_output.copy_to_host()[0]
+
+    expected_x = sum(range(num_items))
+    expected_a = sum(i * 2 for i in range(num_items))
+    expected_b = sum(float(i * 3) for i in range(num_items))
+
+    assert result["x"] == expected_x
+    assert result["inner"]["a"] == expected_a
+    assert np.isclose(result["inner"]["b"], expected_b)
+
+
+def test_nested_struct_in_zip_iterator():
+    Point = gpu_struct({"x": np.int32, "y": np.int32})
+    Color = gpu_struct({"r": np.uint8, "g": np.uint8, "b": np.uint8})
+    Pixel = gpu_struct({"position": Point, "color": Color})
+
+    def sum_pixels(p1, p2):
+        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,
+            ),
+        )
+
+    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 = DeviceArray.from_numpy(h_points)
+    d_colors = DeviceArray.from_numpy(h_colors)
+
+    zip_it = ZipIterator(d_points, d_colors)
+
+    d_output = DeviceArray.empty(1, Pixel.dtype)
+    h_init = Pixel(Point(0, 0), Color(0, 0, 0))
+
+    cuda.compute.reduce_into(
+        d_in=zip_it, d_out=d_output, num_items=num_items, op=sum_pixels, h_init=h_init
+    )
+
+    result = d_output.copy_to_host()[0]
+
+    expected_x = sum(i for i in 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
+
+
+def test_dict_init_nested_struct():
+    """Test initializing a nested struct with a dictionary."""
+    Inner = gpu_struct({"a": np.int32, "b": np.float32})
+    Outer = gpu_struct({"x": np.int64, "inner": Inner})
+
+    # Initialize with nested dictionary
+    obj = Outer({"x": 42, "inner": {"a": 10, "b": 3.14}})
+
+    assert obj.x == 42
+    assert obj.inner.a == 10
+    assert np.isclose(obj.inner.b, 3.14)
+
+
+def test_dict_init_per_field():
+    """Test initializing a struct with a dictionary for a nested field."""
+    Inner = gpu_struct({"a": np.int32, "b": np.float32})
+    Outer = gpu_struct({"x": np.int64, "inner": Inner})
+
+    # Mix positional value with dictionary for nested field
+    obj = Outer(42, {"a": 10, "b": 3.14})
+
+    assert obj.x == 42
+    assert obj.inner.a == 10
+    assert np.isclose(obj.inner.b, 3.14)
+
+
+def test_dict_init_deeply_nested():
+    """Test initializing deeply nested structs (3+ levels) with dictionaries."""
+    Level1 = gpu_struct({"value": np.int32})
+    Level2 = gpu_struct({"data": np.float32, "nested": Level1})
+    Level3 = gpu_struct({"id": np.int64, "middle": Level2})
+
+    # Initialize with deeply nested dictionary
+    obj = Level3({"id": 100, "middle": {"data": 2.5, "nested": {"value": 42}}})
+
+    assert obj.id == 100
+    assert np.isclose(obj.middle.data, 2.5)
+    assert obj.middle.nested.value == 42
+
+
+def test_dict_init_mixed():
+    """Test mixed initialization with some dicts and some direct values."""
+    Inner1 = gpu_struct({"a": np.int32, "b": np.int32})
+    Inner2 = gpu_struct({"c": np.float32, "d": np.float32})
+    Outer = gpu_struct({"x": np.int64, "inner1": Inner1, "inner2": Inner2})
+
+    # Mix different initialization styles
+    inner1_obj = Inner1(1, 2)  # Direct instantiation
+    # Mix direct object and dict
+    obj = Outer(100, inner1_obj, {"c": 3.0, "d": 4.0})
+
+    assert obj.x == 100
+    assert obj.inner1.a == 1
+    assert obj.inner1.b == 2
+    assert np.isclose(obj.inner2.c, 3.0)
+    assert np.isclose(obj.inner2.d, 4.0)
+
+
+def test_dict_init_with_reduction():
+    """Test that dict-initialized structs work correctly in reductions."""
+    Inner = gpu_struct({"a": np.int32, "b": np.float32})
+    Outer = gpu_struct({"x": np.int64, "inner": Inner})
+
+    def sum_nested(s1, s2):
+        return Outer(
+            s1.x + s2.x, Inner(s1.inner.a + s2.inner.a, s1.inner.b + s2.inner.b)
+        )
+
+    num_items = 10
+    h_data = np.zeros(num_items, dtype=Outer.dtype)
+    for i in range(num_items):
+        h_data[i]["x"] = i
+        h_data[i]["inner"]["a"] = i * 2
+        h_data[i]["inner"]["b"] = float(i * 3)
+
+    d_input = DeviceArray.from_numpy(h_data)
+    d_output = DeviceArray.empty(1, Outer.dtype)
+
+    # Use dictionary initialization for the init value
+    h_init = Outer({"x": 0, "inner": {"a": 0, "b": 0.0}})
+
+    cuda.compute.reduce_into(
+        d_in=d_input, d_out=d_output, num_items=num_items, op=sum_nested, h_init=h_init
+    )
+
+    result = d_output.copy_to_host()[0]
+
+    expected_x = sum(range(num_items))
+    expected_a = sum(i * 2 for i in range(num_items))
+    expected_b = sum(float(i * 3) for i in range(num_items))
+
+    assert result["x"] == expected_x
+    assert result["inner"]["a"] == expected_a
+    assert np.isclose(result["inner"]["b"], expected_b)
+
+
+def test_nested_struct_tuple_construction():
+    """Test constructing nested structs using tuple syntax in device functions."""
+    Inner = gpu_struct({"a": np.int32, "b": np.float32})
+    Outer = gpu_struct({"x": np.int64, "inner": Inner})
+
+    def sum_nested_with_tuples(s1, s2):
+        # Use tuple syntax instead of Inner(...)
+        return Outer(s1.x + s2.x, (s1.inner.a + s2.inner.a, s1.inner.b + s2.inner.b))
+
+    num_items = 10
+
+    h_data = np.zeros(num_items, dtype=Outer.dtype)
+    for i in range(num_items):
+        h_data[i]["x"] = i
+        h_data[i]["inner"]["a"] = i * 2
+        h_data[i]["inner"]["b"] = float(i * 3)
+
+    d_input = DeviceArray.from_numpy(h_data)
+    d_output = DeviceArray.empty(1, Outer.dtype)
+
+    h_init = Outer(0, Inner(0, 0.0))
+
+    cuda.compute.reduce_into(
+        d_in=d_input,
+        d_out=d_output,
+        num_items=num_items,
+        op=sum_nested_with_tuples,
+        h_init=h_init,
+    )
+
+    result = d_output.copy_to_host()[0]
+
+    expected_x = sum(range(num_items))
+    expected_a = sum(i * 2 for i in range(num_items))
+    expected_b = sum(float(i * 3) for i in range(num_items))
+
+    assert result["x"] == expected_x
+    assert result["inner"]["a"] == expected_a
+    assert np.isclose(result["inner"]["b"], expected_b)
+
+
+def test_deeply_nested_tuple_construction():
+    """Test constructing deeply nested structs (3 levels) using tuple syntax."""
+    Level1 = gpu_struct({"value": np.int32})
+    Level2 = gpu_struct({"data": np.float32, "nested": Level1})
+    Level3 = gpu_struct({"id": np.int64, "middle": Level2})
+
+    def sum_deeply_nested(v1, v2):
+        # Use nested tuple syntax: (float, (int,))
+        return Level3(
+            v1.id + v2.id,
+            (
+                v1.middle.data + v2.middle.data,
+                (v1.middle.nested.value + v2.middle.nested.value,),
+            ),
+        )
+
+    num_items = 10
+
+    h_data = np.zeros(num_items, dtype=Level3.dtype)
+    for i in range(num_items):
+        h_data[i]["id"] = i * 10
+        h_data[i]["middle"]["data"] = float(i * 2.5)
+        h_data[i]["middle"]["nested"]["value"] = i * 3
+
+    d_input = DeviceArray.from_numpy(h_data)
+    d_output = DeviceArray.empty(1, Level3.dtype)
+
+    h_init = Level3(0, Level2(0.0, Level1(0)))
+
+    cuda.compute.reduce_into(
+        d_in=d_input,
+        d_out=d_output,
+        num_items=num_items,
+        op=sum_deeply_nested,
+        h_init=h_init,
+    )
+
+    result = d_output.copy_to_host()[0]
+
+    expected_id = sum(i * 10 for i in range(num_items))
+    expected_data = sum(float(i * 2.5) for i in range(num_items))
+    expected_value = sum(i * 3 for i in range(num_items))
+
+    assert result["id"] == expected_id
+    assert np.isclose(result["middle"]["data"], expected_data)
+    assert result["middle"]["nested"]["value"] == expected_value
+
+
+def test_mixed_tuple_and_direct_construction():
+    """Test mixing tuple construction with direct struct construction."""
+    Inner1 = gpu_struct({"a": np.int32, "b": np.int32})
+    Inner2 = gpu_struct({"c": np.float32, "d": np.float32})
+    Outer = gpu_struct({"x": np.int64, "inner1": Inner1, "inner2": Inner2})
+
+    def sum_mixed(s1, s2):
+        # Mix direct struct construction with tuple construction
+        return Outer(
+            s1.x + s2.x,
+            Inner1(s1.inner1.a + s2.inner1.a, s1.inner1.b + s2.inner1.b),
+            (s1.inner2.c + s2.inner2.c, s1.inner2.d + s2.inner2.d),
+        )
+
+    num_items = 10
+
+    h_data = np.zeros(num_items, dtype=Outer.dtype)
+    for i in range(num_items):
+        h_data[i]["x"] = i
+        h_data[i]["inner1"]["a"] = i * 2
+        h_data[i]["inner1"]["b"] = i * 3
+        h_data[i]["inner2"]["c"] = float(i * 4)
+        h_data[i]["inner2"]["d"] = float(i * 5)
+
+    d_input = DeviceArray.from_numpy(h_data)
+    d_output = DeviceArray.empty(1, Outer.dtype)
+
+    h_init = Outer(0, Inner1(0, 0), Inner2(0.0, 0.0))
+
+    cuda.compute.reduce_into(
+        d_in=d_input, d_out=d_output, num_items=num_items, op=sum_mixed, h_init=h_init
+    )
+
+    result = d_output.copy_to_host()[0]
+
+    expected_x = sum(range(num_items))
+    expected_a = sum(i * 2 for i in range(num_items))
+    expected_b = sum(i * 3 for i in range(num_items))
+    expected_c = sum(float(i * 4) for i in range(num_items))
+    expected_d = sum(float(i * 5) for i in range(num_items))
+
+    assert result["x"] == expected_x
+    assert result["inner1"]["a"] == expected_a
+    assert result["inner1"]["b"] == expected_b
+    assert np.isclose(result["inner2"]["c"], expected_c)
+    assert np.isclose(result["inner2"]["d"], expected_d)
+
+
+def test_tuple_construction_in_zip_iterator():
+    """Test tuple construction with ZipIterator combining nested structs."""
+    Point = gpu_struct({"x": np.int32, "y": np.int32})
+    Color = gpu_struct({"r": np.uint8, "g": np.uint8, "b": np.uint8})
+    Pixel = gpu_struct({"position": Point, "color": Color})
+
+    def sum_pixels_with_tuples(p1, p2):
+        # Use tuple syntax for both nested structs
+        return Pixel(
+            (p1.position.x + p2.position.x, p1.position.y + p2.position.y),
+            (
+                p1.color.r + p2.color.r,
+                p1.color.g + p2.color.g,
+                p1.color.b + p2.color.b,
+            ),
+        )
+
+    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 = DeviceArray.from_numpy(h_points)
+    d_colors = DeviceArray.from_numpy(h_colors)
+
+    zip_it = ZipIterator(d_points, d_colors)
+
+    d_output = DeviceArray.empty(1, Pixel.dtype)
+    h_init = Pixel(Point(0, 0), Color(0, 0, 0))
+
+    cuda.compute.reduce_into(
+        d_in=zip_it,
+        d_out=d_output,
+        num_items=num_items,
+        op=sum_pixels_with_tuples,
+        h_init=h_init,
+    )
+
+    result = d_output.copy_to_host()[0]
+
+    expected_x = sum(i for i in 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
+
+
+def test_all_tuple_construction():
+    """Test constructing a struct where all fields use tuple syntax."""
+    Inner1 = gpu_struct({"a": np.int32})
+    Inner2 = gpu_struct({"b": np.float32})
+    Outer = gpu_struct({"field1": Inner1, "field2": Inner2})
+
+    def sum_all_tuples(s1, s2):
+        # All fields use tuple syntax
+        return Outer((s1.field1.a + s2.field1.a,), (s1.field2.b + s2.field2.b,))
+
+    num_items = 5
+
+    h_data = np.zeros(num_items, dtype=Outer.dtype)
+    for i in range(num_items):
+        h_data[i]["field1"]["a"] = i
+        h_data[i]["field2"]["b"] = float(i * 2)
+
+    d_input = DeviceArray.from_numpy(h_data)
+    d_output = DeviceArray.empty(1, Outer.dtype)
+
+    h_init = Outer(Inner1(0), Inner2(0.0))
+
+    cuda.compute.reduce_into(
+        d_in=d_input,
+        d_out=d_output,
+        num_items=num_items,
+        op=sum_all_tuples,
+        h_init=h_init,
+    )
+
+    result = d_output.copy_to_host()[0]
+
+    expected_a = sum(range(num_items))
+    expected_b = sum(float(i * 2) for i in range(num_items))
+
+    assert result["field1"]["a"] == expected_a
+    assert np.isclose(result["field2"]["b"], expected_b)
+
+
+def test_struct_field_order_matters():
+    """Test that struct types with same fields in different order are not equal.
+
+    This is a regression test for a bug where StructTypeDescriptor.__hash__()
+    sorted fields before hashing, causing structs with the same fields but
+    different order to hash to the same value and be considered equal.
+    """
+
+    # Create two structs with identical field types but different order
+    @gpu_struct
+    class SumAndCount:
+        sum: np.float32
+        count: np.int32
+
+    @gpu_struct
+    class CountAndSum:
+        count: np.int32
+        sum: np.float32
+
+    # These should NOT be equal (field order matters for struct layout)
+    assert SumAndCount._type_descriptor != CountAndSum._type_descriptor
+
+    # They should have different hashes
+    assert hash(SumAndCount._type_descriptor) != hash(CountAndSum._type_descriptor)
+
+    # Verify they can be used as distinct dict keys
+    cache = {
+        SumAndCount._type_descriptor: "sum_first",
+        CountAndSum._type_descriptor: "count_first",
+    }
+    assert len(cache) == 2
+    assert cache[SumAndCount._type_descriptor] == "sum_first"
+    assert cache[CountAndSum._type_descriptor] == "count_first"

cccl_upstream/python/cuda_cccl/tests/compute/test_no_numba.py

@@ -0,0 +1,606 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+import numpy as np
+import pytest
+from _utils.device_array import DeviceArray
+
+import cuda.compute
+from cuda.compute import (
+    CacheModifiedInputIterator,
+    ConstantIterator,
+    CountingIterator,
+    DiscardIterator,
+    OpKind,
+    PermutationIterator,
+    ReverseIterator,
+    ShuffleIterator,
+    SortOrder,
+    TransformIterator,
+    TransformOutputIterator,
+    ZipIterator,
+)
+from cuda.compute._cpp_compile import compile_cpp_op_code
+from cuda.compute.op import RawOp
+from cuda.compute.types import int16 as cccl_int16
+from cuda.compute.types import int32 as cccl_int32
+
+# These tests define the minimal-extra integration contract. They intentionally
+# use small fixed inputs and avoid the Python-callable operator path.
+pytestmark = pytest.mark.no_numba
+
+
+def _raw_op(source: str, name: str) -> RawOp:
+    return RawOp(ltoir=compile_cpp_op_code(source), name=name)
+
+
+def _raw_even_i32_op() -> RawOp:
+    source = """
+extern "C" __device__ void no_numba_even_i32(void* x, void* result) {
+    int value = *static_cast<int*>(x);
+    *static_cast<bool*>(result) = (value % 2) == 0;
+}
+"""
+    return _raw_op(source, "no_numba_even_i32")
+
+
+def _raw_less_than_i32_op(name: str, threshold: int) -> RawOp:
+    source = f"""
+extern "C" __device__ void {name}(void* x, void* result) {{
+    int value = *static_cast<int*>(x);
+    *static_cast<unsigned char*>(result) = value < {threshold} ? 1 : 0;
+}}
+"""
+    return _raw_op(source, name)
+
+
+def _raw_plus_i64_op() -> RawOp:
+    source = """
+extern "C" __device__ void no_numba_plus_i64(
+    void* lhs,
+    void* rhs,
+    void* result
+) {
+    *static_cast<long long*>(result) =
+        *static_cast<long long*>(lhs) + *static_cast<long long*>(rhs);
+}
+"""
+    return _raw_op(source, "no_numba_plus_i64")
+
+
+def _raw_square_i32_op() -> RawOp:
+    source = """
+extern "C" __device__ void no_numba_square_i32(void* x, void* result) {
+    int value = *static_cast<int*>(x);
+    *static_cast<int*>(result) = value * value;
+}
+"""
+    return _raw_op(source, "no_numba_square_i32")
+
+
+def _raw_zip_sum_i32_op() -> RawOp:
+    source = """
+struct Zip2I32 {
+    int field_0;
+    int field_1;
+};
+
+extern "C" __device__ void no_numba_zip_sum_i32(void* x, void* result) {
+    auto values = static_cast<Zip2I32*>(x);
+    *static_cast<int*>(result) = values->field_0 + values->field_1;
+}
+"""
+    return _raw_op(source, "no_numba_zip_sum_i32")
+
+
+def _raw_negate_i16_op() -> RawOp:
+    source = """
+extern "C" __device__ void no_numba_negate_i16(void* x, void* result) {
+    *static_cast<short*>(result) = -*static_cast<short*>(x);
+}
+"""
+    return _raw_op(source, "no_numba_negate_i16")
+
+
+def test_import_numba_raises(raise_on_numba_import):
+    # Request the guard explicitly rather than relying on the no_numba
+    # auto-injection (skipped under the parallel sweep, see conftest
+    # pytest_collection_modifyitems): this test exercises the guard directly, and
+    # the monkeypatch it depends on keeps it single-threaded under the sweep.
+    with pytest.raises(
+        ImportError, match="This test is marked 'no_numba' but attempted to import it"
+    ):
+        import numba.cuda  # noqa: F401
+
+
+def test_reduce_well_known_plus():
+    h_input = np.arange(1, 14, dtype=np.int32)
+    d_input = DeviceArray.from_numpy(h_input)
+    d_output = DeviceArray.empty(1, np.int32)
+    h_init = np.array([5], dtype=np.int32)
+
+    cuda.compute.reduce_into(
+        d_in=d_input,
+        d_out=d_output,
+        num_items=h_input.size,
+        op=OpKind.PLUS,
+        h_init=h_init,
+    )
+
+    assert d_output.copy_to_host()[0] == np.sum(h_input, initial=h_init[0])
+
+
+def test_exclusive_scan_well_known_plus():
+    h_input = np.asarray([2, 4, 6, 8, 10, 12], dtype=np.uint16)
+    d_input = DeviceArray.from_numpy(h_input)
+    d_output = DeviceArray.empty(h_input.shape, h_input.dtype)
+    h_init = np.array([1], dtype=np.uint16)
+
+    cuda.compute.exclusive_scan(
+        d_in=d_input,
+        d_out=d_output,
+        op=OpKind.PLUS,
+        init_value=h_init,
+        num_items=h_input.size,
+    )
+
+    expected = np.asarray([1, 3, 7, 13, 21, 31], dtype=np.uint16)
+    np.testing.assert_array_equal(d_output.copy_to_host(), expected)
+
+
+def test_binary_transform_well_known_plus():
+    h_lhs = np.asarray([1.5, 2.5, 3.5, 4.5], dtype=np.float32)
+    h_rhs = np.asarray([10.0, 20.0, 30.0, 40.0], dtype=np.float32)
+    d_lhs = DeviceArray.from_numpy(h_lhs)
+    d_rhs = DeviceArray.from_numpy(h_rhs)
+    d_output = DeviceArray.empty(h_lhs.shape, h_lhs.dtype)
+
+    cuda.compute.binary_transform(
+        d_in1=d_lhs,
+        d_in2=d_rhs,
+        d_out=d_output,
+        op=OpKind.PLUS,
+        num_items=h_lhs.size,
+    )
+
+    np.testing.assert_allclose(d_output.copy_to_host(), h_lhs + h_rhs)
+
+
+def test_unary_transform_well_known_negate():
+    h_input = np.asarray([-4, -2, 0, 2, 4], dtype=np.int8)
+    d_input = DeviceArray.from_numpy(h_input)
+    d_output = DeviceArray.empty(h_input.shape, h_input.dtype)
+
+    cuda.compute.unary_transform(
+        d_in=d_input,
+        d_out=d_output,
+        op=OpKind.NEGATE,
+        num_items=h_input.size,
+    )
+
+    np.testing.assert_array_equal(
+        d_output.copy_to_host(), np.asarray([4, 2, 0, -2, -4])
+    )
+
+
+@pytest.mark.parametrize(
+    "search, side",
+    [
+        (cuda.compute.lower_bound, "left"),
+        (cuda.compute.upper_bound, "right"),
+    ],
+)
+def test_binary_search_explicit_opkind_less(search, side):
+    h_data = np.asarray([1, 3, 3, 7, 9, 11], dtype=np.int64)
+    h_values = np.asarray([0, 3, 4, 10, 12], dtype=np.int64)
+    d_out = DeviceArray.empty(h_values.shape, np.uintp)
+
+    search(
+        d_data=DeviceArray.from_numpy(h_data),
+        num_items=h_data.size,
+        d_values=DeviceArray.from_numpy(h_values),
+        num_values=h_values.size,
+        d_out=d_out,
+        comp=OpKind.LESS,
+    )
+
+    expected = np.searchsorted(h_data, h_values, side=side).astype(np.uintp)
+    np.testing.assert_array_equal(d_out.copy_to_host(), expected)
+
+
+@pytest.mark.no_verify_sass
+def test_segmented_reduce_well_known_plus():
+    h_input = np.asarray([1, 2, 3, 4, 5, 6, 7, 8], dtype=np.uint32)
+    h_starts = np.asarray([0, 3, 5], dtype=np.int32)
+    h_ends = np.asarray([3, 5, 8], dtype=np.int32)
+    d_input = DeviceArray.from_numpy(h_input)
+    d_starts = DeviceArray.from_numpy(h_starts)
+    d_ends = DeviceArray.from_numpy(h_ends)
+    d_output = DeviceArray.empty(3, np.uint32)
+    h_init = np.array([0], dtype=np.uint32)
+
+    cuda.compute.segmented_reduce(
+        d_in=d_input,
+        d_out=d_output,
+        num_segments=3,
+        start_offsets_in=d_starts,
+        end_offsets_in=d_ends,
+        op=OpKind.PLUS,
+        h_init=h_init,
+    )
+
+    np.testing.assert_array_equal(d_output.copy_to_host(), np.asarray([6, 9, 21]))
+
+
+def test_merge_sort_well_known_less():
+    h_input = np.asarray([3.5, -1.0, 2.25, 2.0, 7.0], dtype=np.float64)
+    d_input = DeviceArray.from_numpy(h_input)
+    d_output = DeviceArray.empty(h_input.shape, h_input.dtype)
+
+    cuda.compute.merge_sort(
+        d_in_keys=d_input,
+        d_in_values=None,
+        d_out_keys=d_output,
+        d_out_values=None,
+        num_items=h_input.size,
+        op=OpKind.LESS,
+    )
+
+    np.testing.assert_array_equal(d_output.copy_to_host(), np.sort(h_input))
+
+
+def test_radix_sort_key_value_pairs():
+    h_keys = np.asarray([4, -2, 7, 1, -2, 0], dtype=np.int16)
+    h_values = np.asarray([40, 20, 70, 10, 21, 0], dtype=np.uint8)
+    d_out_keys = DeviceArray.empty(h_keys.shape, h_keys.dtype)
+    d_out_values = DeviceArray.empty(h_values.shape, h_values.dtype)
+
+    cuda.compute.radix_sort(
+        d_in_keys=DeviceArray.from_numpy(h_keys),
+        d_out_keys=d_out_keys,
+        d_in_values=DeviceArray.from_numpy(h_values),
+        d_out_values=d_out_values,
+        num_items=h_keys.size,
+        order=SortOrder.ASCENDING,
+    )
+
+    order = np.argsort(h_keys, stable=True)
+    np.testing.assert_array_equal(d_out_keys.copy_to_host(), h_keys[order])
+    np.testing.assert_array_equal(d_out_values.copy_to_host(), h_values[order])
+
+
+def test_segmented_sort_keys():
+    h_keys = np.asarray([3, 1, 2, 9, 7, 8, 6, 5], dtype=np.uint64)
+    h_offsets = np.asarray([0, 3, 6, 8], dtype=np.int64)
+    d_output = DeviceArray.empty(h_keys.shape, h_keys.dtype)
+
+    cuda.compute.segmented_sort(
+        d_in_keys=DeviceArray.from_numpy(h_keys),
+        d_out_keys=d_output,
+        d_in_values=None,
+        d_out_values=None,
+        num_items=h_keys.size,
+        num_segments=h_offsets.size - 1,
+        start_offsets_in=DeviceArray.from_numpy(h_offsets[:-1]),
+        end_offsets_in=DeviceArray.from_numpy(h_offsets[1:]),
+        order=SortOrder.ASCENDING,
+    )
+
+    expected = np.asarray([1, 2, 3, 7, 8, 9, 5, 6], dtype=np.uint64)
+    np.testing.assert_array_equal(d_output.copy_to_host(), expected)
+
+
+@pytest.mark.no_verify_sass
+def test_unique_by_key_well_known_equal_to():
+    h_keys = np.asarray([1, 1, 2, 2, 2, 3, 4, 4], dtype=np.int16)
+    h_values = np.asarray([10, 11, 20, 21, 22, 30, 40, 41], dtype=np.int8)
+    d_keys = DeviceArray.from_numpy(h_keys)
+    d_values = DeviceArray.from_numpy(h_values)
+    d_out_keys = DeviceArray.empty(h_keys.shape, h_keys.dtype)
+    d_out_values = DeviceArray.empty(h_values.shape, h_values.dtype)
+    d_num_selected = DeviceArray.empty(1, np.int64)
+
+    cuda.compute.unique_by_key(
+        d_in_keys=d_keys,
+        d_in_items=d_values,
+        d_out_keys=d_out_keys,
+        d_out_items=d_out_values,
+        d_out_num_selected=d_num_selected,
+        op=OpKind.EQUAL_TO,
+        num_items=h_keys.size,
+    )
+
+    num_selected = int(d_num_selected.copy_to_host()[0])
+    np.testing.assert_array_equal(
+        d_out_keys.copy_to_host()[:num_selected], [1, 2, 3, 4]
+    )
+    np.testing.assert_array_equal(
+        d_out_values.copy_to_host()[:num_selected], [10, 20, 30, 40]
+    )
+
+
+def test_histogram_even_small_range():
+    h_samples = np.asarray([0.5, 1.5, 2.5, 2.75, 3.0, 3.5], dtype=np.float32)
+    d_histogram = DeviceArray.empty(4, np.int32)
+
+    cuda.compute.histogram_even(
+        d_samples=DeviceArray.from_numpy(h_samples),
+        d_histogram=d_histogram,
+        num_output_levels=5,
+        lower_level=np.float32(0.0),
+        upper_level=np.float32(4.0),
+        num_samples=h_samples.size,
+    )
+
+    expected, _ = np.histogram(h_samples, bins=4, range=(0.0, 4.0))
+    np.testing.assert_array_equal(d_histogram.copy_to_host(), expected.astype(np.int32))
+
+
+def test_select_raw_op():
+    h_input = np.arange(12, dtype=np.int32)
+    d_output = DeviceArray.empty(h_input.shape, h_input.dtype)
+    d_num_selected = DeviceArray.empty(1, np.uint64)
+
+    cuda.compute.select(
+        d_in=DeviceArray.from_numpy(h_input),
+        d_out=d_output,
+        d_num_selected_out=d_num_selected,
+        cond=_raw_even_i32_op(),
+        num_items=h_input.size,
+    )
+
+    num_selected = int(d_num_selected.copy_to_host()[0])
+    np.testing.assert_array_equal(d_output.copy_to_host()[:num_selected], h_input[::2])
+
+
+def test_three_way_partition_raw_op():
+    h_input = np.arange(12, dtype=np.int32)
+    d_first = DeviceArray.empty(h_input.shape, h_input.dtype)
+    d_second = DeviceArray.empty(h_input.shape, h_input.dtype)
+    d_unselected = DeviceArray.empty(h_input.shape, h_input.dtype)
+    d_num_selected = DeviceArray.empty(2, np.uint64)
+
+    cuda.compute.three_way_partition(
+        d_in=DeviceArray.from_numpy(h_input),
+        d_first_part_out=d_first,
+        d_second_part_out=d_second,
+        d_unselected_out=d_unselected,
+        d_num_selected_out=d_num_selected,
+        select_first_part_op=_raw_less_than_i32_op("no_numba_less_than_4_i32", 4),
+        select_second_part_op=_raw_less_than_i32_op("no_numba_less_than_8_i32", 8),
+        num_items=h_input.size,
+    )
+
+    selected = d_num_selected.copy_to_host()
+    first_count = int(selected[0])
+    second_count = int(selected[1])
+    unselected_count = h_input.size - first_count - second_count
+
+    np.testing.assert_array_equal(d_first.copy_to_host()[:first_count], h_input[:4])
+    np.testing.assert_array_equal(d_second.copy_to_host()[:second_count], h_input[4:8])
+    np.testing.assert_array_equal(
+        d_unselected.copy_to_host()[:unselected_count], h_input[8:]
+    )
+
+
+def test_raw_op_reduce():
+    h_input = np.asarray([10, 20, 30, 40], dtype=np.int64)
+    d_output = DeviceArray.empty(1, np.int64)
+
+    cuda.compute.reduce_into(
+        d_in=DeviceArray.from_numpy(h_input),
+        d_out=d_output,
+        num_items=h_input.size,
+        op=_raw_plus_i64_op(),
+        h_init=np.array([5], dtype=np.int64),
+    )
+
+    assert d_output.copy_to_host()[0] == 105
+
+
+def test_stream_argument(cuda_stream):
+    h_lhs = np.asarray([2, 4, 6, 8, 10], dtype=np.int32)
+    h_rhs = np.asarray([1, 3, 5, 7, 9], dtype=np.int32)
+    d_lhs = DeviceArray.from_numpy(h_lhs, stream=cuda_stream)
+    d_rhs = DeviceArray.from_numpy(h_rhs, stream=cuda_stream)
+    d_output = DeviceArray.empty(h_lhs.shape, h_lhs.dtype, stream=cuda_stream)
+
+    cuda.compute.binary_transform(
+        d_in1=d_lhs,
+        d_in2=d_rhs,
+        d_out=d_output,
+        op=OpKind.PLUS,
+        num_items=h_lhs.size,
+        stream=cuda_stream,
+    )
+
+    np.testing.assert_array_equal(
+        d_output.copy_to_host(stream=cuda_stream),
+        np.asarray([3, 7, 11, 15, 19]),
+    )
+
+
+def test_counting_iterator_reduce():
+    d_output = DeviceArray.empty(1, np.int32)
+
+    cuda.compute.reduce_into(
+        d_in=CountingIterator(np.int32(3)),
+        d_out=d_output,
+        num_items=8,
+        op=OpKind.PLUS,
+        h_init=np.array([0], dtype=np.int32),
+    )
+
+    assert d_output.copy_to_host()[0] == 52
+
+
+def test_constant_iterator_reduce():
+    d_output = DeviceArray.empty(1, np.float32)
+
+    cuda.compute.reduce_into(
+        d_in=ConstantIterator(np.float32(1.5)),
+        d_out=d_output,
+        num_items=8,
+        op=OpKind.PLUS,
+        h_init=np.array([0], dtype=np.float32),
+    )
+
+    np.testing.assert_allclose(d_output.copy_to_host()[0], np.float32(12.0))
+
+
+def test_cache_modified_input_iterator_reduce():
+    h_input = np.asarray([2, 4, 6, 8, 10], dtype=np.uint16)
+    d_input = DeviceArray.from_numpy(h_input)
+    d_output = DeviceArray.empty(1, np.uint16)
+    iterator = CacheModifiedInputIterator(d_input, modifier="stream")
+
+    cuda.compute.reduce_into(
+        d_in=iterator,
+        d_out=d_output,
+        num_items=h_input.size,
+        op=OpKind.PLUS,
+        h_init=np.array([0], dtype=np.uint16),
+    )
+
+    assert d_output.copy_to_host()[0] == 30
+
+
+def test_reverse_input_iterator_scan():
+    h_input = np.asarray([1, 2, 3, 4, 5], dtype=np.int32)
+    d_input = DeviceArray.from_numpy(h_input)
+    d_output = DeviceArray.empty(h_input.shape, h_input.dtype)
+
+    cuda.compute.inclusive_scan(
+        d_in=ReverseIterator(d_input),
+        d_out=d_output,
+        op=OpKind.PLUS,
+        init_value=np.array([0], dtype=np.int32),
+        num_items=h_input.size,
+    )
+
+    np.testing.assert_array_equal(
+        d_output.copy_to_host(), np.asarray([5, 9, 12, 14, 15])
+    )
+
+
+def test_reverse_output_iterator_scan():
+    h_input = np.asarray([1, 2, 3, 4, 5], dtype=np.int32)
+    d_input = DeviceArray.from_numpy(h_input)
+    d_output = DeviceArray.empty(h_input.shape, h_input.dtype)
+
+    cuda.compute.inclusive_scan(
+        d_in=d_input,
+        d_out=ReverseIterator(d_output),
+        op=OpKind.PLUS,
+        init_value=np.array([0], dtype=np.int32),
+        num_items=h_input.size,
+    )
+
+    np.testing.assert_array_equal(
+        d_output.copy_to_host(), np.asarray([15, 10, 6, 3, 1])
+    )
+
+
+def test_permutation_iterator_reduce():
+    h_values = np.asarray([10, 20, 30, 40, 50, 60], dtype=np.int64)
+    h_indices = np.asarray([4, 2, 5, 1], dtype=np.int32)
+    d_values = DeviceArray.from_numpy(h_values)
+    d_indices = DeviceArray.from_numpy(h_indices)
+    d_output = DeviceArray.empty(1, np.int64)
+
+    cuda.compute.reduce_into(
+        d_in=PermutationIterator(d_values, d_indices),
+        d_out=d_output,
+        num_items=h_indices.size,
+        op=OpKind.PLUS,
+        h_init=np.array([0], dtype=np.int64),
+    )
+
+    assert d_output.copy_to_host()[0] == 160
+
+
+def test_transform_iterator_reduce():
+    d_output = DeviceArray.empty(1, np.int32)
+    iterator = TransformIterator(
+        CountingIterator(np.int32(1)), _raw_square_i32_op(), value_type=cccl_int32
+    )
+
+    cuda.compute.reduce_into(
+        d_in=iterator,
+        d_out=d_output,
+        num_items=6,
+        op=OpKind.PLUS,
+        h_init=np.array([0], dtype=np.int32),
+    )
+
+    assert d_output.copy_to_host()[0] == 91
+
+
+def test_transform_output_iterator_reduce():
+    h_input = np.asarray([1, 2, 3, 4], dtype=np.int16)
+    d_input = DeviceArray.from_numpy(h_input)
+    d_output = DeviceArray.empty(1, np.int16)
+    output_iterator = TransformOutputIterator(
+        d_output, _raw_negate_i16_op(), output_value_type=cccl_int16
+    )
+
+    cuda.compute.reduce_into(
+        d_in=d_input,
+        d_out=output_iterator,
+        num_items=h_input.size,
+        op=OpKind.PLUS,
+        h_init=np.array([0], dtype=np.int16),
+    )
+
+    assert d_output.copy_to_host()[0] == -10
+
+
+def test_zip_iterator_transform():
+    h_lhs = np.asarray([1, 2, 3, 4, 5], dtype=np.int32)
+    h_rhs = np.asarray([10, 20, 30, 40, 50], dtype=np.int32)
+    d_lhs = DeviceArray.from_numpy(h_lhs)
+    d_rhs = DeviceArray.from_numpy(h_rhs)
+    d_output = DeviceArray.empty(h_lhs.shape, h_lhs.dtype)
+
+    cuda.compute.unary_transform(
+        d_in=ZipIterator(d_lhs, d_rhs),
+        d_out=d_output,
+        op=_raw_zip_sum_i32_op(),
+        num_items=h_lhs.size,
+    )
+
+    np.testing.assert_array_equal(d_output.copy_to_host(), h_lhs + h_rhs)
+
+
+def test_shuffle_iterator_transform():
+    num_items = 17
+    d_output = DeviceArray.empty(num_items, np.int64)
+
+    cuda.compute.unary_transform(
+        d_in=ShuffleIterator(num_items, seed=123),
+        d_out=d_output,
+        op=OpKind.IDENTITY,
+        num_items=num_items,
+    )
+
+    result = d_output.copy_to_host()
+    assert sorted(result.tolist()) == list(range(num_items))
+
+
+def test_discard_iterator_transform():
+    h_input = np.asarray([1, 2, 3, 4, 5], dtype=np.int32)
+    h_reference = np.full_like(h_input, -1)
+    d_input = DeviceArray.from_numpy(h_input)
+    d_reference = DeviceArray.from_numpy(h_reference)
+
+    cuda.compute.unary_transform(
+        d_in=d_input,
+        d_out=DiscardIterator(d_reference),
+        op=OpKind.IDENTITY,
+        num_items=h_input.size,
+    )
+
+    np.testing.assert_array_equal(
+        d_reference.copy_to_host(), np.full(5, -1, dtype=np.int32)
+    )

cccl_upstream/python/cuda_cccl/tests/compute/test_numba_path_thread_safety.py

@@ -0,0 +1,387 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+"""Thread-safety of the shared build/cache machinery for the numba-dependent
+op paths, on regular (GIL) interpreters.
+
+cuda.compute releases the GIL around native builds and
+launches, so concurrent calls overlap in the C layer even under the GIL.
+So the shared caching/native machinery must be thread-safe on every interpreter
+build, and these tests validate that for the numba op paths.
+
+They cover the paths the free-threading stress suite
+(test_free_threading_stress.py) cannot: that suite runs on the minimal extra,
+which omits numba, so Python-callable ops, stateful closure ops, gpu_struct
+types, and return-type inference are untested there. This file exercises them
+from multiple threads on a regular GIL interpreter. The C-layer build/launch
+machinery itself is op-source-agnostic and already covered by the stress suite;
+the unique surface here is the numba frontend (CachableFunction hashing,
+inference, closure/struct handling).
+
+Only the native build/launch phase genuinely overlaps (the GIL is released
+around it); numba's own compilation serializes on its compiler lock and the
+Python-side key hashing serializes under the GIL. So the goal is not to prove
+numba runs in parallel, but that cuda.compute's caching, key hashing, and
+descriptor machinery stay correct when these paths are entered concurrently.
+"""
+
+import concurrent.futures
+import threading
+
+import numpy as np
+import pytest
+from _utils.device_array import DeviceArray
+
+import cuda.compute
+from cuda.compute import (
+    CountingIterator,
+    OpKind,
+    TransformIterator,
+    gpu_struct,
+    make_reduce_into,
+    make_unary_transform,
+)
+from cuda.core import Device
+
+pytestmark = pytest.mark.no_verify_sass(
+    reason="Concurrency tests intentionally run concurrent workers."
+)
+
+# Every input in this file is tiny (32-128 elements), so raising the thread count
+# adds no meaningful GPU-memory pressure -- it just exercises the concurrent
+# build/cache machinery harder.
+THREADS = 8
+ITERATIONS = 3
+
+# The distinct-op build storm is the canonical stress for concurrent NATIVE
+# builds. On the v2 (HostJIT) backend clang is embedded in-process and is not
+# thread-safe, so a regression in the compile-serialization lock corrupts the
+# heap and crashes here -- but only at high concurrency (at 4 threads it fires
+# only under external load). Use many threads so this stays a reliable gate.
+# One barrier-synced round of that many simultaneous builds already corrupts the
+# heap pre-fix (it crashed on the first round every time), and extra rounds are
+# expensive on v2 (builds serialize), so a single iteration is enough here.
+BUILD_STORM_THREADS = 24
+BUILD_STORM_ITERATIONS = 1
+
+
+def _run_threaded(workers):
+    # The default timeout turns a worker that dies before reaching the barrier
+    # into a BrokenBarrierError in its peers instead of hanging the CI job.
+    barrier = threading.Barrier(len(workers), timeout=60)
+    with concurrent.futures.ThreadPoolExecutor(max_workers=len(workers)) as executor:
+        futures = [executor.submit(worker, barrier) for worker in workers]
+        errors = []
+        for future in futures:
+            try:
+                future.result()
+            except BaseException as exc:  # noqa: BLE001
+                errors.append(exc)
+        if errors:
+            # Surface the root-cause worker failure, not the barrier breakage
+            # it caused in the other workers.
+            raise next(
+                (
+                    error
+                    for error in errors
+                    if not isinstance(error, threading.BrokenBarrierError)
+                ),
+                errors[0],
+            )
+
+
+def _reduce_with_temp(reducer, **kwargs):
+    temp_storage_bytes = reducer(temp_storage=None, **kwargs)
+    temp_storage = DeviceArray.empty(temp_storage_bytes, np.uint8)
+    return reducer(temp_storage=temp_storage, **kwargs)
+
+
+def _single_build_result(algorithm):
+    assert len(algorithm.build_results) == 1
+    return next(iter(algorithm.build_results.values()))
+
+
+def _make_clamped_max_op(k):
+    def clamped_max(a, b):
+        m = a if a > b else b
+        return m if m > k else k
+
+    return clamped_max
+
+
+def test_concurrent_distinct_python_ops_build_storm():
+    """Distinct Python callables force a separate numba+native build per worker.
+
+    Distinct closure constants give each worker its own cache key, so
+    _cache_single_flight elects a separate builder per thread. The native
+    builds overlap (the GIL is released around them) while the CachableFunction
+    hashing and numba compilation serialize (under the GIL and numba's own
+    compiler lock); the target is that this concurrent entry keeps each
+    worker's op and build uncontaminated. The op computes max(a, b, k) with a
+    per-worker k that dominates every input, so the expected result is exactly
+    k however CUB shapes the reduction tree — and a wrong k directly exposes
+    any cross-thread op/build contamination.
+    """
+    num_items = 64
+    for iteration in range(BUILD_STORM_ITERATIONS):
+        cuda.compute.clear_all_caches()
+        returned_reducers = [None] * BUILD_STORM_THREADS
+
+        def make_thread(worker_id):
+            k = 10_000 + worker_id * 7 + iteration
+            op = _make_clamped_max_op(k)
+            h_in = np.arange(num_items, dtype=np.int64) + worker_id
+            h_init = np.array([0], dtype=np.int64)
+            d_in = DeviceArray.from_numpy(h_in)
+            d_out = DeviceArray.empty(1, np.int64)
+
+            def thread(barrier):
+                # cuda.core device state is per-thread; initialize explicitly
+                # rather than relying on DeviceArray-construction side effects.
+                Device().set_current()
+                barrier.wait()
+                reducer = make_reduce_into(d_in=d_in, d_out=d_out, op=op, h_init=h_init)
+                returned_reducers[worker_id] = reducer
+                _reduce_with_temp(
+                    reducer,
+                    d_in=d_in,
+                    d_out=d_out,
+                    op=op,
+                    h_init=h_init,
+                    num_items=num_items,
+                )
+                Device().sync()
+                assert int(d_out.copy_to_host()[0]) == k
+
+            return thread
+
+        _run_threaded(
+            [make_thread(worker_id) for worker_id in range(BUILD_STORM_THREADS)]
+        )
+
+        build_ids = {id(_single_build_result(r)) for r in returned_reducers}
+        assert len(build_ids) == BUILD_STORM_THREADS
+
+
+def test_concurrent_shared_python_op_coalesces():
+    """One shared Python callable from all threads coalesces to a single build.
+
+    All workers hash the same function concurrently (CachableFunction walks
+    bytecode, constants, and closures on every factory call) and race the
+    same cache key; exactly one build must result, with one wrapper per thread.
+    """
+
+    def add(a, b):
+        return a + b
+
+    num_items = 64
+    for iteration in range(ITERATIONS):
+        cuda.compute.clear_all_caches()
+        returned_reducers = [None] * THREADS
+
+        def make_thread(worker_id):
+            h_in = np.arange(num_items, dtype=np.int32) + worker_id + iteration
+            h_init = np.array([worker_id], dtype=np.int32)
+            d_in = DeviceArray.from_numpy(h_in)
+            d_out = DeviceArray.empty(1, np.int32)
+
+            def thread(barrier):
+                # cuda.core device state is per-thread; initialize explicitly
+                # rather than relying on DeviceArray-construction side effects.
+                Device().set_current()
+                barrier.wait()
+                reducer = make_reduce_into(
+                    d_in=d_in, d_out=d_out, op=add, h_init=h_init
+                )
+                returned_reducers[worker_id] = reducer
+                _reduce_with_temp(
+                    reducer,
+                    d_in=d_in,
+                    d_out=d_out,
+                    op=add,
+                    h_init=h_init,
+                    num_items=num_items,
+                )
+                Device().sync()
+                assert int(d_out.copy_to_host()[0]) == int(h_in.sum()) + worker_id
+
+            return thread
+
+        _run_threaded([make_thread(worker_id) for worker_id in range(THREADS)])
+
+        # Per-thread wrappers, one shared native build.
+        assert len({id(r) for r in returned_reducers}) == THREADS
+        assert len({id(_single_build_result(r)) for r in returned_reducers}) == 1
+
+
+def _make_scale_op(k):
+    def scale(x):
+        return x * k
+
+    return scale
+
+
+def test_concurrent_transform_iterator_return_type_inference():
+    """Unannotated TransformIterator ops infer return types concurrently.
+
+    TransformIterator without value_type routes through _infer_return_type
+    (numba typing). Ops are distinct per worker and per iteration so the
+    inference cache stays cold and each thread runs the typing path itself.
+    """
+    num_items = 32
+    for iteration in range(ITERATIONS):
+        cuda.compute.clear_all_caches()
+
+        def make_thread(worker_id):
+            k = np.int32(worker_id + 1 + iteration * THREADS)
+            op = _make_scale_op(k)
+            h_init = np.array([0], dtype=np.int32)
+            d_out = DeviceArray.empty(1, np.int32)
+
+            def thread(barrier):
+                # cuda.core device state is per-thread; initialize explicitly
+                # rather than relying on DeviceArray-construction side effects.
+                Device().set_current()
+                barrier.wait()
+                # Inference widens int32 * int32-closure to an int64 value type;
+                # the int64-iterator/int32-accumulator mix is a supported,
+                # CI-exercised configuration and all expected values fit int32.
+                d_in = TransformIterator(CountingIterator(np.int32(0)), op)
+                reducer = make_reduce_into(
+                    d_in=d_in, d_out=d_out, op=OpKind.PLUS, h_init=h_init
+                )
+                _reduce_with_temp(
+                    reducer,
+                    d_in=d_in,
+                    d_out=d_out,
+                    op=OpKind.PLUS,
+                    h_init=h_init,
+                    num_items=num_items,
+                )
+                Device().sync()
+                expected = sum(i * int(k) for i in range(num_items))
+                assert int(d_out.copy_to_host()[0]) == expected
+
+            return thread
+
+        _run_threaded([make_thread(worker_id) for worker_id in range(THREADS)])
+
+
+def test_concurrent_stateful_closure_ops_isolate_state():
+    """Same op code with different captured device arrays across threads.
+
+    The captured arrays do not change the cache key (stateful op machinery
+    updates pointers per call), so all threads share one native build while
+    each thread's wrapper must carry its own captured-state pointer. If state
+    isolation broke, one thread's output would use another thread's offset.
+    """
+
+    def make_adder(arr):
+        def add_offset(x):
+            return x + arr[0]
+
+        return add_offset
+
+    num_items = 64
+    for iteration in range(ITERATIONS):
+        cuda.compute.clear_all_caches()
+        returned_transformers = [None] * THREADS
+
+        def make_thread(worker_id):
+            offset = worker_id * 10 + iteration
+            d_offset = DeviceArray.from_numpy(np.array([offset], dtype=np.int32))
+            op = make_adder(d_offset)
+            h_in = np.arange(num_items, dtype=np.int32)
+            d_in = DeviceArray.from_numpy(h_in)
+            d_out = DeviceArray.empty(h_in.shape, h_in.dtype)
+
+            def thread(barrier):
+                # cuda.core device state is per-thread; initialize explicitly
+                # rather than relying on DeviceArray-construction side effects.
+                Device().set_current()
+                barrier.wait()
+                transformer = make_unary_transform(d_in=d_in, d_out=d_out, op=op)
+                returned_transformers[worker_id] = transformer
+                transformer(d_in=d_in, d_out=d_out, op=op, num_items=num_items)
+                Device().sync()
+                np.testing.assert_array_equal(d_out.copy_to_host(), h_in + offset)
+
+            return thread
+
+        _run_threaded([make_thread(worker_id) for worker_id in range(THREADS)])
+
+        assert len({id(t) for t in returned_transformers}) == THREADS
+        assert len({id(_single_build_result(t)) for t in returned_transformers}) == 1
+
+
+def test_concurrent_gpu_struct_reduce():
+    """gpu_struct types and struct-typed Python ops used from multiple threads."""
+
+    @gpu_struct
+    class MinMax:
+        min_val: np.int32
+        max_val: np.int32
+
+    def minmax_op(a, b):
+        c_min = min(a.min_val, b.min_val)
+        c_max = max(a.max_val, b.max_val)
+        return MinMax(c_min, c_max)
+
+    num_items = 128
+    info = np.iinfo(np.int32)
+
+    # Warm the numba struct registration once on the main thread (mirroring the
+    # stress suite's warm pass). First-time registration goes through
+    # non-coalescing lru_caches in _jit.py that mutate numba's global
+    # registries, so racing it cold from four threads is a separate library
+    # hardening question, not this test's target.
+    warm_in = DeviceArray.from_numpy(
+        np.zeros((1, 2), dtype=np.int32).view(MinMax.dtype)
+    )
+    warm_out = DeviceArray.empty(1, MinMax.dtype)
+    cuda.compute.reduce_into(
+        d_in=warm_in,
+        d_out=warm_out,
+        op=minmax_op,
+        h_init=MinMax(info.max, info.min),
+        num_items=1,
+    )
+    Device().sync()
+
+    for iteration in range(ITERATIONS):
+        cuda.compute.clear_all_caches()
+
+        def make_thread(worker_id):
+            base = worker_id * 1000 + iteration
+            h_pairs = np.stack(
+                [
+                    np.arange(num_items, dtype=np.int32) + base,
+                    np.arange(num_items, dtype=np.int32) * 2 + base,
+                ],
+                axis=1,
+            )
+            d_in = DeviceArray.from_numpy(h_pairs.view(MinMax.dtype))
+            d_out = DeviceArray.empty(1, MinMax.dtype)
+            h_init = MinMax(info.max, info.min)
+
+            def thread(barrier):
+                # cuda.core device state is per-thread; initialize explicitly
+                # rather than relying on DeviceArray-construction side effects.
+                Device().set_current()
+                barrier.wait()
+                cuda.compute.reduce_into(
+                    d_in=d_in,
+                    d_out=d_out,
+                    op=minmax_op,
+                    h_init=h_init,
+                    num_items=num_items,
+                )
+                Device().sync()
+                result = d_out.copy_to_host()
+                assert int(result["min_val"][0]) == int(h_pairs[:, 0].min())
+                assert int(result["max_val"][0]) == int(h_pairs[:, 1].max())
+
+            return thread
+
+        _run_threaded([make_thread(worker_id) for worker_id in range(THREADS)])

cccl_upstream/python/cuda_cccl/tests/compute/test_permutation_iterator.py

@@ -0,0 +1,263 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+import numpy as np
+from _utils.device_array import DeviceArray
+
+import cuda.compute
+from cuda.compute.iterators import (
+    CountingIterator,
+    PermutationIterator,
+    ZipIterator,
+)
+
+
+def test_permutation_iterator_equality():
+    d_values1 = DeviceArray.from_numpy(np.asarray([10, 20, 30, 40, 50], dtype="int32"))
+    d_values2 = DeviceArray.from_numpy(np.asarray([100, 200, 300], dtype="int32"))
+    d_values3 = DeviceArray.from_numpy(np.asarray([10, 20, 30, 40, 50], dtype="int64"))
+
+    d_indices1 = DeviceArray.from_numpy(np.asarray([0, 2, 1], dtype="int32"))
+    d_indices2 = DeviceArray.from_numpy(np.asarray([1, 0, 2], dtype="int32"))
+    d_indices3 = DeviceArray.from_numpy(np.asarray([0, 2, 1], dtype="int64"))
+
+    # Same value and index types should have same kind
+    it1 = PermutationIterator(d_values1, d_indices1)
+    it2 = PermutationIterator(d_values1, d_indices2)
+    it3 = PermutationIterator(d_values2, d_indices1)
+
+    assert it1.kind == it2.kind == it3.kind
+
+    # Different value type should have different kind
+    it4 = PermutationIterator(d_values3, d_indices1)
+    assert it1.kind != it4.kind
+
+    # Different index type should have different kind
+    it5 = PermutationIterator(d_values1, d_indices3)
+    assert it1.kind != it5.kind
+
+
+def test_permutation_iterator_with_array_values():
+    h_values = np.asarray([10, 20, 30, 40, 50], dtype="int32")
+    h_indices = np.asarray([2, 0, 4, 1], dtype="int32")
+    d_values = DeviceArray.from_numpy(h_values)
+    d_indices = DeviceArray.from_numpy(h_indices)
+    perm_it = PermutationIterator(d_values, d_indices)
+
+    h_init = np.array([0], dtype="int32")
+    d_output = DeviceArray.empty(1, np.int32)
+    cuda.compute.reduce_into(
+        d_in=perm_it,
+        d_out=d_output,
+        num_items=len(h_indices),
+        op=cuda.compute.OpKind.PLUS,
+        h_init=h_init,
+    )
+    assert d_output.copy_to_host()[0] == h_values[h_indices].sum()
+
+
+def test_permutation_iterator_with_iterator_values():
+    values_it = CountingIterator(np.int32(10))
+    h_indices = np.asarray([2, 0, 4, 1], dtype="int32")
+    d_indices = DeviceArray.from_numpy(h_indices)
+    perm_it = PermutationIterator(values_it, d_indices)
+
+    h_init = np.array([0], dtype="int32")
+    d_output = DeviceArray.empty(1, np.int32)
+
+    cuda.compute.reduce_into(
+        d_in=perm_it,
+        d_out=d_output,
+        num_items=len(h_indices),
+        op=cuda.compute.OpKind.PLUS,
+        h_init=h_init,
+    )
+
+    expected = np.arange(10, 20)[h_indices].sum()
+    assert d_output.copy_to_host()[0] == expected
+
+
+def test_permutation_iterator_of_zip_iterator():
+    @cuda.compute.gpu_struct
+    class Pair:
+        value_0: np.int32
+        value_1: np.int32
+
+    h_values1 = np.asarray([10, 20, 30, 40, 50], dtype="int32")
+    h_values2 = np.asarray([1, 2, 3, 4, 5], dtype="int32")
+    d_values1 = DeviceArray.from_numpy(h_values1)
+    d_values2 = DeviceArray.from_numpy(h_values2)
+    zip_it = ZipIterator(d_values1, d_values2)
+    h_indices = np.asarray([2, 0, 4], dtype="int32")
+    d_indices = DeviceArray.from_numpy(h_indices)
+    perm_it = PermutationIterator(zip_it, d_indices)
+
+    def sum_both_fields(a, b):
+        return Pair(a.value_0 + b.value_0, a.value_1 + b.value_1)
+
+    h_init = Pair(0, 0)
+    d_output = DeviceArray.empty(1, Pair.dtype)
+
+    cuda.compute.reduce_into(
+        d_in=perm_it,
+        d_out=d_output,
+        num_items=len(h_indices),
+        op=sum_both_fields,
+        h_init=h_init,
+    )
+
+    result = d_output.copy_to_host()[0]
+    assert result["value_0"] == h_values1[h_indices].sum()
+    assert result["value_1"] == h_values2[h_indices].sum()
+
+
+def test_zip_iterator_of_permutation_iterators():
+    @cuda.compute.gpu_struct
+    class Pair:
+        value_0: np.int32
+        value_1: np.int32
+
+    h_values1 = np.asarray([10, 20, 30, 40, 50], dtype="int32")
+    h_values2 = np.asarray([100, 200, 300, 400, 500], dtype="int32")
+    h_indices1 = np.asarray([4, 1, 3, 0], dtype="int32")
+    h_indices2 = np.asarray([2, 4, 0, 1], dtype="int32")
+    d_values1 = DeviceArray.from_numpy(h_values1)
+    d_values2 = DeviceArray.from_numpy(h_values2)
+    d_indices1 = DeviceArray.from_numpy(h_indices1)
+    d_indices2 = DeviceArray.from_numpy(h_indices2)
+    perm_it1 = PermutationIterator(d_values1, d_indices1)
+    perm_it2 = PermutationIterator(d_values2, d_indices2)
+
+    zip_it = ZipIterator(perm_it1, perm_it2)
+
+    def sum_both_fields(a, b):
+        return Pair(a.value_0 + b.value_0, a.value_1 + b.value_1)
+
+    h_init = Pair(0, 0)
+    d_output = DeviceArray.empty(1, Pair.dtype)
+
+    num_items = len(h_indices1)
+    cuda.compute.reduce_into(
+        d_in=zip_it,
+        d_out=d_output,
+        num_items=num_items,
+        op=sum_both_fields,
+        h_init=h_init,
+    )
+
+    result = d_output.copy_to_host()[0]
+    assert result["value_0"] == h_values1[h_indices1].sum()
+    assert result["value_1"] == h_values2[h_indices2].sum()
+
+
+def test_unary_transform_of_permutation_iterator():
+    h_values = np.asarray([10, 20, 30, 40, 50], dtype="int32")
+    h_indices = np.asarray([2, 0, 4, 1], dtype="int32")
+    d_values = DeviceArray.from_numpy(h_values)
+    d_indices = DeviceArray.from_numpy(h_indices)
+    perm_it = PermutationIterator(d_values, d_indices)
+
+    def op(a):
+        return a + 1
+
+    d_out = DeviceArray.empty(len(h_indices), h_values.dtype)
+    cuda.compute.unary_transform(
+        d_in=perm_it, d_out=d_out, op=op, num_items=len(h_indices)
+    )
+
+    expected = h_values[h_indices] + 1
+    np.testing.assert_array_equal(d_out.copy_to_host(), expected)
+
+
+def test_caching_permutation_iterator():
+    """Test that iterator compilation is cached across instances with the same structure."""
+    from cuda.compute._cpp_compile import compile_cpp_op_code
+
+    # Test 1: Same structure → same kind
+    it1 = PermutationIterator(
+        DeviceArray.from_numpy(np.arange(10, dtype=np.int32)),
+        DeviceArray.from_numpy(np.arange(10, dtype=np.int32)),
+    )
+    it2 = PermutationIterator(
+        DeviceArray.from_numpy(np.arange(20, dtype=np.int32)),
+        DeviceArray.from_numpy(np.arange(5, dtype=np.int32)),
+    )
+    assert it1.kind == it2.kind, "Same structure should have same kind"
+
+    # Test 2: Different index type → different kind
+    it3 = PermutationIterator(
+        DeviceArray.from_numpy(np.arange(10, dtype=np.int32)),
+        DeviceArray.from_numpy(np.arange(10, dtype=np.int64)),
+    )
+    assert it1.kind != it3.kind, "Different index type should have different kind"
+
+    # Test 3: Different value type → different kind
+    it4 = PermutationIterator(
+        DeviceArray.from_numpy(np.arange(10, dtype=np.int64)),
+        DeviceArray.from_numpy(np.arange(10, dtype=np.int32)),
+    )
+    assert it1.kind != it4.kind, "Different value type should have different kind"
+
+    # Test 4: Verify compilation caching with cache statistics
+    compile_cpp_op_code.cache_clear()
+
+    # Create multiple instances with same structure
+    iterators = []
+    for i in range(3):
+        it = PermutationIterator(
+            DeviceArray.from_numpy(np.arange(i * 10, (i + 1) * 10, dtype=np.float32)),
+            DeviceArray.from_numpy(np.arange(5, dtype=np.int32)),
+        )
+        # Trigger compilation by accessing Op objects
+        it.get_advance_op()
+        it.get_input_deref_op()
+        iterators.append(it)
+
+    cache_info = compile_cpp_op_code.cache_info()
+    assert cache_info.hits >= 2, (
+        f"Expected cache hits for same structure, got {cache_info.hits} hits, "
+        f"{cache_info.misses} misses"
+    )
+
+
+def test_permutation_iterator_advance():
+    """Test PermutationIterator.__add__ only advances indices, not values."""
+    # Create values array [10, 20, 30, 40, 50, 60, 70]
+    h_values = np.asarray([10, 20, 30, 40, 50, 60, 70], dtype="int32")
+    d_values = DeviceArray.from_numpy(h_values)
+
+    # Create indices array [2, 0, 4, 1, 3, 5]
+    # indices[0] = 2 -> values[2] = 30
+    # indices[1] = 0 -> values[0] = 10
+    # indices[2] = 4 -> values[4] = 50
+    # indices[3] = 1 -> values[1] = 20
+    # indices[4] = 3 -> values[3] = 40
+    # indices[5] = 5 -> values[5] = 60
+    h_indices = np.asarray([2, 0, 4, 1, 3, 5], dtype="int32")
+    d_indices = DeviceArray.from_numpy(h_indices)
+
+    perm_it = PermutationIterator(d_values, d_indices)
+
+    # Advance by 2 positions (should skip first 2 indices)
+    offset = 2
+    advanced_perm_it = perm_it + offset
+
+    # Reduce from the advanced position
+    # Should process indices[2:] = [4, 1, 3, 5]
+    # Which accesses values[4, 1, 3, 5] = [50, 20, 40, 60]
+    h_init = np.array([0], dtype="int32")
+    d_output = DeviceArray.empty(1, np.int32)
+
+    remaining_items = len(h_indices) - offset
+    cuda.compute.reduce_into(
+        d_in=advanced_perm_it,
+        d_out=d_output,
+        num_items=remaining_items,
+        op=cuda.compute.OpKind.PLUS,
+        h_init=h_init,
+    )
+
+    # Expected: values[indices[2:]] = values[[4, 1, 3, 5]] = [50, 20, 40, 60]
+    expected = h_values[h_indices[offset:]].sum()
+    result = d_output.copy_to_host()[0]
+    assert result == expected, f"Expected {expected}, got {result}"

cccl_upstream/python/cuda_cccl/tests/compute/test_radix_sort.py

@@ -0,0 +1,629 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+import itertools
+from typing import Tuple
+
+import numpy as np
+import pytest
+from _utils.device_array import DeviceArray, get_compute_capability
+
+import cuda.compute
+from cuda.compute import (
+    DoubleBuffer,
+    SortOrder,
+    deserialize,
+    make_radix_sort,
+    serialize,
+)
+from cuda.compute._utils.temp_storage_buffer import TempStorageBuffer
+
+
+def get_mark(dt, log_size):
+    if log_size < 20:
+        return tuple()
+    return pytest.mark.large
+
+
+DTYPE_LIST = [
+    np.uint8,
+    np.uint16,
+    np.uint32,
+    np.uint64,
+    np.int8,
+    np.int16,
+    np.int32,
+    np.int64,
+    np.float16,
+    np.float32,
+    np.float64,
+]
+
+PROBLEM_SIZES = [2, 10, 20]
+
+DTYPE_SIZE = [
+    pytest.param(dt, 2**log_size, marks=get_mark(dt, log_size))
+    for dt in DTYPE_LIST
+    for log_size in PROBLEM_SIZES
+]
+
+
+def random_array(size, dtype, max_value=None) -> np.typing.NDArray:
+    rng = np.random.default_rng()
+    if np.isdtype(dtype, "integral"):
+        if max_value is None:
+            max_value = np.iinfo(dtype).max
+        return rng.integers(max_value, size=size, dtype=dtype)
+    elif np.isdtype(dtype, "real floating"):
+        return np.random.uniform(low=-10.0, high=10.0, size=size).astype(dtype)
+    else:
+        raise ValueError(f"Unsupported dtype {dtype}")
+
+
+def radix_sort_device(
+    d_in_keys,
+    d_out_keys,
+    d_in_values,
+    d_out_values,
+    order,
+    num_items,
+    begin_bit=None,
+    end_bit=None,
+    stream=None,
+):
+    # Use the new single-phase API with automatic temp storage allocation
+    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,
+        num_items=num_items,
+        order=order,
+        begin_bit=begin_bit,
+        end_bit=end_bit,
+        stream=stream,
+    )
+
+
+def get_floating_point_keys(array):
+    """
+    This function computes the keys for floating point types.
+    From the cub docs, this is the required behavior:
+
+    For positive floating point values, the sign bit is inverted.
+    For negative floating point values, the full key is inverted.
+    """
+    if array.dtype == np.float32:
+        uint_type = np.uint32
+        sign_mask = np.uint32(0x80000000)
+    elif array.dtype == np.float64:
+        uint_type = np.uint64
+        sign_mask = np.uint64(0x8000000000000000)
+
+    # Get the binary representation as unsigned integers
+    binary = array.copy().view(uint_type)
+
+    # Create masks for positive and negative numbers
+    is_positive = array >= 0
+    is_negative = ~is_positive
+
+    # For positive numbers: flip the sign bit (leftmost bit)
+    binary[is_positive] ^= sign_mask
+
+    # For negative numbers: invert all bits
+    binary[is_negative] = ~binary[is_negative]
+
+    return binary
+
+
+def host_sort(h_in_keys, h_in_values, order, begin_bit=None, end_bit=None) -> Tuple:
+    if begin_bit is not None and end_bit is not None:
+        num_bits = end_bit - begin_bit
+        mask = np.array(((1 << (num_bits)) - 1) << begin_bit, dtype=np.uint64)
+
+        if np.issubdtype(h_in_keys.dtype, np.floating):
+            h_in_keys_copy = get_floating_point_keys(h_in_keys)
+        else:
+            h_in_keys_copy = h_in_keys
+
+        h_in_keys_copy = h_in_keys_copy.astype(np.uint64)
+        h_in_keys_copy = (h_in_keys_copy & mask) >> begin_bit
+    else:
+        h_in_keys_copy = h_in_keys
+
+    if order is SortOrder.DESCENDING:
+        # We do this for stability. We need to cast to a signed integer to properly negate the keys.
+        signed_dtype = np.dtype(h_in_keys_copy.dtype.name.replace("uint", "int"))
+        argsort = np.argsort(-h_in_keys_copy.astype(signed_dtype), stable=True)
+    else:
+        argsort = np.argsort(h_in_keys_copy, stable=True)
+
+    h_in_keys = h_in_keys[argsort]
+    if h_in_values is not None:
+        h_in_values = h_in_values[argsort]
+
+    return h_in_keys, h_in_values
+
+
+@pytest.mark.parametrize(
+    "dtype, num_items",
+    DTYPE_SIZE,
+)
+def test_radix_sort_keys(dtype, num_items, monkeypatch):
+    cc_major, _ = get_compute_capability()
+    # Skip sass verification for CC 9.0+ due to a bug in NVRTC.
+    # TODO: add NVRTC version check, ref nvbug 5243118
+    if cc_major >= 9:
+        import cuda.compute._cccl_interop
+
+        monkeypatch.setattr(
+            cuda.compute._cccl_interop,
+            "_check_sass",
+            False,
+        )
+
+    order = SortOrder.ASCENDING
+    h_in_keys = random_array(num_items, dtype, max_value=20)
+    h_out_keys = np.empty(num_items, dtype=dtype)
+
+    d_in_keys = DeviceArray.from_numpy(h_in_keys)
+    d_out_keys = DeviceArray.empty(h_out_keys.shape, h_out_keys.dtype)
+
+    radix_sort_device(d_in_keys, d_out_keys, None, None, order, num_items)
+
+    h_out_keys = d_out_keys.copy_to_host()
+
+    h_in_keys, _ = host_sort(h_in_keys, None, order)
+
+    np.testing.assert_array_equal(h_out_keys, h_in_keys)
+
+
+@pytest.mark.parametrize(
+    "dtype, num_items",
+    DTYPE_SIZE,
+)
+def test_radix_sort_pairs(dtype, num_items, monkeypatch):
+    import cuda.compute._cccl_interop
+
+    monkeypatch.setattr(
+        cuda.compute._cccl_interop,
+        "_check_sass",
+        False,
+    )
+
+    order = SortOrder.DESCENDING
+    h_in_keys = random_array(num_items, dtype, max_value=20)
+    h_in_values = random_array(num_items, np.float32)
+    h_out_keys = np.empty(num_items, dtype=dtype)
+    h_out_values = np.empty(num_items, dtype=np.float32)
+
+    d_in_keys = DeviceArray.from_numpy(h_in_keys)
+    d_in_values = DeviceArray.from_numpy(h_in_values)
+    d_out_keys = DeviceArray.empty(h_out_keys.shape, h_out_keys.dtype)
+    d_out_values = DeviceArray.empty(h_out_values.shape, h_out_values.dtype)
+
+    radix_sort_device(
+        d_in_keys, d_out_keys, d_in_values, d_out_values, order, num_items
+    )
+
+    h_out_keys = d_out_keys.copy_to_host()
+    h_out_values = d_out_values.copy_to_host()
+
+    h_in_keys, h_in_values = host_sort(h_in_keys, h_in_values, order)
+
+    np.testing.assert_array_equal(h_out_keys, h_in_keys)
+    np.testing.assert_array_equal(h_out_values, h_in_values)
+
+
+@pytest.mark.parametrize(
+    "dtype, num_items",
+    DTYPE_SIZE,
+)
+def test_radix_sort_keys_double_buffer(dtype, num_items, monkeypatch):
+    cc_major, _ = get_compute_capability()
+    # Skip sass verification for CC 9.0+ due to a bug in NVRTC.
+    # TODO: add NVRTC version check, ref nvbug 5243118
+    if cc_major >= 9:
+        import cuda.compute._cccl_interop
+
+        monkeypatch.setattr(
+            cuda.compute._cccl_interop,
+            "_check_sass",
+            False,
+        )
+
+    order = SortOrder.DESCENDING
+    h_in_keys = random_array(num_items, dtype, max_value=20)
+    h_out_keys = np.empty(num_items, dtype=dtype)
+
+    d_in_keys = DeviceArray.from_numpy(h_in_keys)
+    d_out_keys = DeviceArray.empty(h_out_keys.shape, h_out_keys.dtype)
+
+    keys_double_buffer = DoubleBuffer(d_in_keys, d_out_keys)
+
+    radix_sort_device(keys_double_buffer, None, None, None, order, num_items)
+
+    h_out_keys = keys_double_buffer.current().copy_to_host()
+
+    h_in_keys, _ = host_sort(h_in_keys, None, order)
+
+    np.testing.assert_array_equal(h_out_keys, h_in_keys)
+
+
+@pytest.mark.parametrize(
+    "dtype, num_items",
+    DTYPE_SIZE,
+)
+def test_radix_sort_pairs_double_buffer(dtype, num_items, monkeypatch):
+    cc_major, _ = get_compute_capability()
+    # NOTE: int16 failures seen only with NVRTC 13.1:
+    if cc_major >= 9 or np.isdtype(dtype, (np.int16, np.uint32)):
+        import cuda.compute._cccl_interop
+
+        monkeypatch.setattr(
+            cuda.compute._cccl_interop,
+            "_check_sass",
+            False,
+        )
+
+    order = SortOrder.ASCENDING
+    h_in_keys = random_array(num_items, dtype, max_value=20)
+    h_in_values = random_array(num_items, np.float32)
+    h_out_keys = np.empty(num_items, dtype=dtype)
+    h_out_values = np.empty(num_items, dtype=np.float32)
+
+    d_in_keys = DeviceArray.from_numpy(h_in_keys)
+    d_in_values = DeviceArray.from_numpy(h_in_values)
+    d_out_keys = DeviceArray.empty(h_out_keys.shape, h_out_keys.dtype)
+    d_out_values = DeviceArray.empty(h_out_values.shape, h_out_values.dtype)
+
+    keys_double_buffer = DoubleBuffer(d_in_keys, d_out_keys)
+    values_double_buffer = DoubleBuffer(d_in_values, d_out_values)
+
+    radix_sort_device(
+        keys_double_buffer, None, values_double_buffer, None, order, num_items
+    )
+
+    h_out_keys = keys_double_buffer.current().copy_to_host()
+    h_out_values = values_double_buffer.current().copy_to_host()
+
+    h_in_keys, h_in_values = host_sort(h_in_keys, h_in_values, order)
+
+    np.testing.assert_array_equal(h_out_keys, h_in_keys)
+    np.testing.assert_array_equal(h_out_values, h_in_values)
+
+
+# These tests take longer to execute so we reduce the number of test cases
+DTYPE_SIZE_BIT_WINDOW = [
+    pytest.param(dt, 2**log_size, marks=get_mark(dt, log_size))
+    for dt in [np.uint8, np.int16, np.uint32, np.int64, np.float64]
+    for log_size in [2, 24]
+]
+
+
+@pytest.mark.parametrize(
+    "dtype, num_items",
+    DTYPE_SIZE_BIT_WINDOW,
+)
+def test_radix_sort_pairs_bit_window(dtype, num_items, monkeypatch):
+    cc_major, _ = get_compute_capability()
+    # NOTE: int16 failures seen only with NVRTC 13.1:
+    if cc_major >= 9 or np.isdtype(dtype, (np.int16, np.uint32)):
+        import cuda.compute._cccl_interop
+
+        monkeypatch.setattr(
+            cuda.compute._cccl_interop,
+            "_check_sass",
+            False,
+        )
+
+    order = SortOrder.ASCENDING
+    num_bits = dtype().itemsize
+    begin_bits = [0, num_bits // 3, 3 * num_bits // 4, num_bits]
+    end_bits = [0, num_bits // 3, 3 * num_bits // 4, num_bits]
+
+    for begin_bit, end_bit in itertools.product(begin_bits, end_bits):
+        if end_bit < begin_bit:
+            continue
+
+        h_in_keys = random_array(num_items, dtype)
+        h_in_values = random_array(num_items, np.float32)
+        h_out_keys = np.empty(num_items, dtype=dtype)
+        h_out_values = np.empty(num_items, dtype=np.float32)
+
+        d_in_keys = DeviceArray.from_numpy(h_in_keys)
+        d_in_values = DeviceArray.from_numpy(h_in_values)
+        d_out_keys = DeviceArray.empty(h_out_keys.shape, h_out_keys.dtype)
+        d_out_values = DeviceArray.empty(h_out_values.shape, h_out_values.dtype)
+
+        radix_sort_device(
+            d_in_keys,
+            d_out_keys,
+            d_in_values,
+            d_out_values,
+            order,
+            num_items,
+            begin_bit,
+            end_bit,
+        )
+
+        h_out_keys = d_out_keys.copy_to_host()
+        h_out_values = d_out_values.copy_to_host()
+
+        h_in_keys, h_in_values = host_sort(
+            h_in_keys, h_in_values, order, begin_bit, end_bit
+        )
+
+        np.testing.assert_array_equal(h_out_keys, h_in_keys)
+        np.testing.assert_array_equal(h_out_values, h_in_values)
+
+
+@pytest.mark.parametrize(
+    "dtype, num_items",
+    DTYPE_SIZE_BIT_WINDOW,
+)
+def test_radix_sort_pairs_double_buffer_bit_window(dtype, num_items, monkeypatch):
+    if np.isdtype(dtype, (np.uint8, np.int16, np.uint32)):
+        import cuda.compute._cccl_interop
+
+        monkeypatch.setattr(
+            cuda.compute._cccl_interop,
+            "_check_sass",
+            False,
+        )
+
+    order = SortOrder.DESCENDING
+    num_bits = dtype().itemsize
+    begin_bits = [0, num_bits // 3, 3 * num_bits // 4, num_bits]
+    end_bits = [0, num_bits // 3, 3 * num_bits // 4, num_bits]
+
+    for begin_bit, end_bit in itertools.product(begin_bits, end_bits):
+        if end_bit < begin_bit:
+            continue
+
+        h_in_keys = random_array(num_items, dtype)
+        h_in_values = random_array(num_items, np.float32)
+        h_out_keys = np.empty(num_items, dtype=dtype)
+        h_out_values = np.empty(num_items, dtype=np.float32)
+
+        d_in_keys = DeviceArray.from_numpy(h_in_keys)
+        d_in_values = DeviceArray.from_numpy(h_in_values)
+        d_out_keys = DeviceArray.empty(h_out_keys.shape, h_out_keys.dtype)
+        d_out_values = DeviceArray.empty(h_out_values.shape, h_out_values.dtype)
+
+        keys_double_buffer = DoubleBuffer(d_in_keys, d_out_keys)
+        values_double_buffer = DoubleBuffer(d_in_values, d_out_values)
+
+        radix_sort_device(
+            keys_double_buffer,
+            None,
+            values_double_buffer,
+            None,
+            order,
+            num_items,
+            begin_bit,
+            end_bit,
+        )
+
+        h_out_keys = keys_double_buffer.current().copy_to_host()
+        h_out_values = values_double_buffer.current().copy_to_host()
+
+        h_in_keys, h_in_values = host_sort(
+            h_in_keys, h_in_values, order, begin_bit, end_bit
+        )
+
+        np.testing.assert_array_equal(h_out_keys, h_in_keys)
+        np.testing.assert_array_equal(h_out_values, h_in_values)
+
+
+@pytest.mark.large
+@pytest.mark.parametrize("dtype", [np.int32, np.float32])
+def test_radix_sort_large_num_items(dtype, monkeypatch):
+    """Regression test for https://github.com/NVIDIA/cccl/issues/7938.
+
+    Radix sort produces incorrect output for large inputs that require
+    multiple "portions" internally (roughly >= 2**28 elements, depending
+    on the tuning policy chosen for the current GPU).
+    """
+    import cuda.compute._cccl_interop
+
+    monkeypatch.setattr(
+        cuda.compute._cccl_interop,
+        "_check_sass",
+        False,
+    )
+    num_items = 2**28
+
+    h_in_keys = np.arange(num_items - 1, -1, -1, dtype=dtype)
+
+    d_in_keys = DeviceArray.from_numpy(h_in_keys)
+    d_out_keys = DeviceArray.empty(num_items, dtype)
+
+    cuda.compute.radix_sort(
+        d_in_keys=d_in_keys,
+        d_out_keys=d_out_keys,
+        d_in_values=None,
+        d_out_values=None,
+        num_items=num_items,
+        order=SortOrder.ASCENDING,
+    )
+
+    h_out_keys = d_out_keys.copy_to_host()
+    h_expected, _ = host_sort(h_in_keys, None, SortOrder.ASCENDING)
+
+    np.testing.assert_array_equal(h_out_keys, h_expected)
+
+
+def test_radix_sort_with_stream(cuda_stream):
+    num_items = 10000
+
+    h_in_keys = random_array(num_items, np.int32)
+    d_in_keys = DeviceArray.from_numpy(h_in_keys, stream=cuda_stream)
+    d_out_keys = DeviceArray.empty(h_in_keys.shape, h_in_keys.dtype, stream=cuda_stream)
+
+    radix_sort_device(
+        d_in_keys,
+        d_out_keys,
+        None,
+        None,
+        SortOrder.ASCENDING,
+        num_items,
+        stream=cuda_stream,
+    )
+
+    got = d_out_keys.copy_to_host(stream=cuda_stream)
+    h_in_keys.sort()
+
+    np.testing.assert_array_equal(got, h_in_keys)
+
+
+def test_radix_sort(monkeypatch):
+    cc_major, _ = get_compute_capability()
+    # Skip sass verification for CC 9.0+ due to a bug in NVRTC.
+    # TODO: add NVRTC version check, ref nvbug 5243118
+    if cc_major >= 9:
+        import cuda.compute._cccl_interop as cccl_interop
+
+        monkeypatch.setattr(
+            cccl_interop,
+            "_check_sass",
+            False,
+        )
+
+    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 = DeviceArray.from_numpy(h_in_keys)
+    d_in_values = DeviceArray.from_numpy(h_in_values)
+
+    d_out_keys = DeviceArray.empty(h_in_keys.shape, h_in_keys.dtype)
+    d_out_values = DeviceArray.empty(h_in_values.shape, h_in_values.dtype)
+
+    # Call single-phase API directly with num_items parameter
+    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,
+        num_items=h_in_keys.size,
+        order=SortOrder.ASCENDING,
+    )
+
+    # Check the result is correct
+    h_out_keys = d_out_keys.copy_to_host()
+    h_out_items = d_out_values.copy_to_host()
+
+    argsort = np.argsort(h_in_keys, stable=True)
+    h_in_keys = np.array(h_in_keys)[argsort]
+    h_in_values = np.array(h_in_values)[argsort]
+
+    np.testing.assert_array_equal(h_out_keys, h_in_keys)
+    np.testing.assert_array_equal(h_out_items, h_in_values)
+
+
+def test_radix_sort_double_buffer(monkeypatch):
+    cc_major, _ = get_compute_capability()
+    # Skip sass verification for CC 9.0+ due to a bug in NVRTC.
+    # TODO: add NVRTC version check, ref nvbug 5243118
+    if cc_major >= 9:
+        import cuda.compute._cccl_interop as cccl_interop
+
+        monkeypatch.setattr(
+            cccl_interop,
+            "_check_sass",
+            False,
+        )
+
+    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 = DeviceArray.from_numpy(h_in_keys)
+    d_in_values = DeviceArray.from_numpy(h_in_values)
+
+    d_out_keys = DeviceArray.empty(h_in_keys.shape, h_in_keys.dtype)
+    d_out_values = DeviceArray.empty(h_in_values.shape, h_in_values.dtype)
+
+    keys_double_buffer = DoubleBuffer(d_in_keys, d_out_keys)
+    values_double_buffer = DoubleBuffer(d_in_values, d_out_values)
+
+    # Call single-phase API directly with num_items parameter
+    cuda.compute.radix_sort(
+        d_in_keys=keys_double_buffer,
+        d_out_keys=None,
+        d_in_values=values_double_buffer,
+        d_out_values=None,
+        num_items=h_in_keys.size,
+        order=SortOrder.ASCENDING,
+    )
+
+    # Check the result is correct
+    h_out_keys = keys_double_buffer.current().copy_to_host()
+    h_out_values = values_double_buffer.current().copy_to_host()
+
+    argsort = np.argsort(h_in_keys, stable=True)
+    h_in_keys = np.array(h_in_keys)[argsort]
+    h_in_values = np.array(h_in_values)[argsort]
+
+    np.testing.assert_array_equal(h_out_keys, h_in_keys)
+    np.testing.assert_array_equal(h_out_values, h_in_values)
+
+
+def _run(sorter, *, d_in_keys, d_out_keys, d_in_values, d_out_values, num_items):
+    bytes_needed = sorter(
+        temp_storage=None,
+        d_in_keys=d_in_keys,
+        d_out_keys=d_out_keys,
+        d_in_values=d_in_values,
+        d_out_values=d_out_values,
+        num_items=num_items,
+    )
+    tmp = TempStorageBuffer(bytes_needed, None)
+    sorter(
+        temp_storage=tmp,
+        d_in_keys=d_in_keys,
+        d_out_keys=d_out_keys,
+        d_in_values=d_in_values,
+        d_out_values=d_out_values,
+        num_items=num_items,
+    )
+
+
+@pytest.mark.serialization
+def test_serialize_deserialize_radix_sort_keys_values():
+    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 = DeviceArray.from_numpy(h_in_keys)
+    d_in_values = DeviceArray.from_numpy(h_in_values)
+    d_out_keys = DeviceArray.empty(h_in_keys.shape, h_in_keys.dtype)
+    d_out_values = DeviceArray.empty(h_in_values.shape, h_in_values.dtype)
+
+    builder = make_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,
+    )
+    blob = serialize(builder)
+    assert len(blob) > 0
+
+    loaded = deserialize(blob)
+    _run(
+        loaded,
+        d_in_keys=d_in_keys,
+        d_out_keys=d_out_keys,
+        d_in_values=d_in_values,
+        d_out_values=d_out_values,
+        num_items=h_in_keys.size,
+    )
+
+    argsort = np.argsort(h_in_keys, stable=True)
+    np.testing.assert_array_equal(d_out_keys.copy_to_host(), h_in_keys[argsort])
+    np.testing.assert_array_equal(d_out_values.copy_to_host(), h_in_values[argsort])

cccl_upstream/python/cuda_cccl/tests/compute/test_raw_op.py

@@ -0,0 +1,525 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+#
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+import re
+import struct
+
+import numpy as np
+import pytest
+from _utils.device_array import DeviceArray, get_compute_capability
+
+import cuda.compute
+from cuda.compute import types
+from cuda.compute._cpp_compile import _get_include_paths
+from cuda.compute.op import RawOp
+from cuda.core import Program, ProgramOptions
+
+# Mark all tests in this module as no_numba
+pytestmark = pytest.mark.no_numba
+
+
+def get_arch():
+    """Get the SM architecture string for the current device."""
+    cc_major, cc_minor = get_compute_capability()
+    return f"sm_{cc_major}{cc_minor}"
+
+
+def compile_to_ltoir(
+    source: str, arch: str, include_paths: list | None = None
+) -> bytes:
+    """Compile C++ source to LTOIR using cuda.core."""
+    opts = ProgramOptions(
+        arch=arch,
+        relocatable_device_code=True,
+        link_time_optimization=True,
+        include_path=include_paths,
+    )
+    prog = Program(source, "c++", options=opts)
+    return prog.compile("ltoir").code
+
+
+def extract_function_name(source: str) -> str:
+    """Extract function name from C++ source."""
+    match = re.search(r'extern\s+"C"\s+__device__\s+\w+\s+(\w+)\s*\(', source)
+    if match:
+        return match.group(1)
+    raise ValueError("Could not extract function name from C++ source.")
+
+
+def make_cpp_op(source: str, name: str = None, include_paths=None) -> RawOp:
+    """
+    Compile C++ source to LTO-IR and create a stateless RawOp. For v2 users
+    the LLVM-bitcode path is preferred (see ``examples/raw_op/llvm_stateless.py``);
+    this helper exercises the LTO-IR escape hatch.
+
+    Args:
+        source: C++ source code containing the operator function
+        name: Optional function name. If not provided, will be extracted from source
+
+    Returns:
+        RawOp instance ready to use with CCCL algorithms
+    """
+    if name is None:
+        name = extract_function_name(source)
+
+    arch = get_arch()
+    ltoir = compile_to_ltoir(source, arch, include_paths=include_paths)
+
+    return RawOp(ltoir=ltoir, name=name)
+
+
+def make_cpp_stateful_op(
+    source: str, state: bytes, name: str = None, state_alignment: int = 8
+) -> RawOp:
+    """
+    Compile C++ source to LTO-IR and create a stateful RawOp. See the note on
+    :func:`make_cpp_op` regarding the v2-preferred LLVM-bitcode path.
+
+    Args:
+        source: C++ source code containing the operator function
+        state: State data as bytes
+        name: Optional function name. If not provided, will be extracted from source
+        state_alignment: Memory alignment for state (default: 8)
+
+    Returns:
+        RawOp instance ready to use with CCCL algorithms
+    """
+    if name is None:
+        name = extract_function_name(source)
+
+    arch = get_arch()
+    ltoir = compile_to_ltoir(source, arch)
+
+    return RawOp(
+        ltoir=ltoir,
+        name=name,
+        state=state,
+        state_alignment=state_alignment,
+    )
+
+
+def test_cpp_op_basic_add():
+    """Test a basic C++ addition operator with reduce_into."""
+    cpp_source = """
+    extern "C" __device__ void add_op(void* a, void* b, void* result) {
+        *static_cast<int*>(result) = *static_cast<int*>(a) + *static_cast<int*>(b);
+    }
+    """
+
+    op = make_cpp_op(cpp_source, "add_op")
+
+    # Create test data
+    num_items = 100
+    h_input = np.arange(num_items, dtype=np.int32)
+    d_input = DeviceArray.from_numpy(h_input)
+    d_output = DeviceArray.empty(1, np.int32)
+
+    # Use the custom op with reduce_into
+    h_init = np.array(0, dtype=np.int32)
+    cuda.compute.reduce_into(
+        d_in=d_input, d_out=d_output, num_items=num_items, op=op, h_init=h_init
+    )
+
+    # Verify result
+    result = d_output.copy_to_host()[0]
+    expected = np.sum(h_input)
+    assert result == expected, f"Expected {expected}, got {result}"
+
+
+def test_cpp_op_max():
+    """Test a C++ max operator with reduce_into."""
+    cpp_source = """
+    #include <cuda/std/algorithm>
+
+    extern "C" __device__ void max_op(void* a, void* b, void* result) {
+        float va = *static_cast<float*>(a);
+        float vb = *static_cast<float*>(b);
+        *static_cast<float*>(result) = cuda::std::max(va, vb);
+    }
+    """
+    # _get_include_paths() returns list of CCCL include paths;
+    # needed for libcudacxx headers like cuda/std/algorithm
+    op = make_cpp_op(cpp_source, include_paths=_get_include_paths())
+
+    # Create test data
+    num_items = 100
+    h_input = np.random.randn(num_items).astype(np.float32)
+    d_input = DeviceArray.from_numpy(h_input)
+    d_output = DeviceArray.empty(1, np.float32)
+
+    # Use the custom op with reduce_into
+    h_init = np.array(-np.inf, dtype=np.float32)
+    cuda.compute.reduce_into(
+        d_in=d_input, d_out=d_output, num_items=num_items, op=op, h_init=h_init
+    )
+
+    # Verify result
+    result = d_output.copy_to_host()[0]
+    expected = np.max(h_input)
+    assert np.isclose(result, expected), f"Expected {expected}, got {result}"
+
+
+def test_cpp_op_multiply():
+    """Test a C++ multiply operator."""
+    cpp_source = """
+    extern "C" __device__ void multiply(void* a, void* b, void* result) {
+        *static_cast<int*>(result) = *static_cast<int*>(a) * *static_cast<int*>(b);
+    }
+    """
+
+    op = make_cpp_op(cpp_source, "multiply")
+
+    # Create test data - use small numbers to avoid overflow
+    num_items = 5
+    h_input = np.array([1, 2, 3, 4, 5], dtype=np.int32)
+    d_input = DeviceArray.from_numpy(h_input)
+    d_output = DeviceArray.empty(1, np.int32)
+
+    # Use the custom op with reduce_into
+    h_init = np.array(1, dtype=np.int32)
+    cuda.compute.reduce_into(
+        d_in=d_input, d_out=d_output, num_items=num_items, op=op, h_init=h_init
+    )
+
+    # Verify result
+    result = d_output.copy_to_host()[0]
+    expected = np.prod(h_input)
+    assert result == expected, f"Expected {expected}, got {result}"
+
+
+def test_cpp_op_complex_logic():
+    """Test a C++ operator with more complex logic - bitwise OR (associative)."""
+    cpp_source = """
+    extern "C" __device__ void bitwise_or(void* a, void* b, void* result) {
+        // Bitwise OR is associative: (a | b) | c = a | (b | c)
+        int va = *static_cast<int*>(a);
+        int vb = *static_cast<int*>(b);
+        *static_cast<int*>(result) = va | vb;
+    }
+    """
+
+    op = make_cpp_op(cpp_source, "bitwise_or")
+
+    # Create test data with specific bit patterns
+    num_items = 5
+    h_input = np.array([1, 2, 4, 8, 16], dtype=np.int32)  # Powers of 2
+    d_input = DeviceArray.from_numpy(h_input)
+    d_output = DeviceArray.empty(1, np.int32)
+
+    # Use the custom op with reduce_into
+    h_init = np.array(0, dtype=np.int32)
+    cuda.compute.reduce_into(
+        d_in=d_input, d_out=d_output, num_items=num_items, op=op, h_init=h_init
+    )
+
+    # Expected: 1 | 2 | 4 | 8 | 16 = 31 (all bits set)
+    result = d_output.copy_to_host()[0]
+    expected = 31
+    assert result == expected, f"Expected {expected}, got {result}"
+
+
+def test_cpp_op_different_types():
+    """Test C++ operator with different numeric types."""
+    cpp_source = """
+    extern "C" __device__ void add_doubles(void* a, void* b, void* result) {
+        *static_cast<double*>(result) = *static_cast<double*>(a) + *static_cast<double*>(b);
+    }
+    """
+
+    op = make_cpp_op(cpp_source, "add_doubles")
+
+    # Create test data
+    num_items = 50
+    h_input = np.random.randn(num_items).astype(np.float64)
+    d_input = DeviceArray.from_numpy(h_input)
+    d_output = DeviceArray.empty(1, np.float64)
+
+    # Use the custom op with reduce_into
+    h_init = np.array(0.0, dtype=np.float64)
+    cuda.compute.reduce_into(
+        d_in=d_input, d_out=d_output, num_items=num_items, op=op, h_init=h_init
+    )
+
+    # Verify result
+    result = d_output.copy_to_host()[0]
+    expected = np.sum(h_input)
+    assert np.isclose(result, expected), f"Expected {expected}, got {result}"
+
+
+def test_cpp_op_name_extraction():
+    """Test that function name is correctly extracted from C++ source."""
+    cpp_source = """
+    extern "C" __device__ void my_function_name(void* a, void* b, void* result) {
+        *static_cast<int*>(result) = *static_cast<int*>(a) + *static_cast<int*>(b);
+    }
+    """
+
+    # Don't provide name - it should be extracted
+    op = make_cpp_op(cpp_source)
+
+    # Create test data
+    num_items = 10
+    h_input = np.arange(num_items, dtype=np.int32)
+    d_input = DeviceArray.from_numpy(h_input)
+    d_output = DeviceArray.empty(1, np.int32)
+
+    # Use the custom op with reduce_into
+    h_init = np.array(0, dtype=np.int32)
+    cuda.compute.reduce_into(
+        d_in=d_input, d_out=d_output, num_items=num_items, op=op, h_init=h_init
+    )
+
+    # Verify result
+    result = d_output.copy_to_host()[0]
+    expected = np.sum(h_input)
+    assert result == expected, f"Expected {expected}, got {result}"
+
+
+def test_cpp_op_min():
+    """Test a C++ min operator."""
+    cpp_source = """
+    extern "C" __device__ void min_op(void* a, void* b, void* result) {
+        int va = *static_cast<int*>(a);
+        int vb = *static_cast<int*>(b);
+        *static_cast<int*>(result) = va < vb ? va : vb;
+    }
+    """
+
+    op = make_cpp_op(cpp_source, "min_op")
+
+    # Create test data
+    num_items = 100
+    h_input = np.random.randint(-1000, 1000, num_items, dtype=np.int32)
+    d_input = DeviceArray.from_numpy(h_input)
+    d_output = DeviceArray.empty(1, np.int32)
+
+    # Use the custom op with reduce_into
+    h_init = np.array(np.iinfo(np.int32).max, dtype=np.int32)
+    cuda.compute.reduce_into(
+        d_in=d_input, d_out=d_output, num_items=num_items, op=op, h_init=h_init
+    )
+
+    # Verify result
+    result = d_output.copy_to_host()[0]
+    expected = np.min(h_input)
+    assert result == expected, f"Expected {expected}, got {result}"
+
+
+def test_cpp_op_with_struct():
+    """Test a C++ operator that works with struct types."""
+    from cuda.compute import gpu_struct
+
+    # Define a simple 2D point struct
+    Point = gpu_struct({"x": np.int32, "y": np.int32})
+
+    # C++ operator to add two points (field by field)
+    cpp_source = """
+    struct Point {
+        int x;
+        int y;
+    };
+
+    extern "C" __device__ void add_points(void* a, void* b, void* result) {
+        Point* pa = static_cast<Point*>(a);
+        Point* pb = static_cast<Point*>(b);
+        Point* pr = static_cast<Point*>(result);
+        pr->x = pa->x + pb->x;
+        pr->y = pa->y + pb->y;
+    }
+    """
+
+    op = make_cpp_op(cpp_source, "add_points")
+
+    # Create test data
+    num_items = 10
+    h_data = np.zeros(num_items, dtype=Point.dtype)
+    for i in range(num_items):
+        h_data[i]["x"] = i
+        h_data[i]["y"] = i * 2
+
+    d_input = DeviceArray.from_numpy(h_data)
+    d_output = DeviceArray.empty(1, Point.dtype)
+
+    # Initial point (0, 0)
+    h_init = Point(0, 0)
+
+    # Use the custom op with reduce_into
+    cuda.compute.reduce_into(
+        d_in=d_input, d_out=d_output, num_items=num_items, op=op, h_init=h_init
+    )
+
+    # Verify result
+    result = d_output.copy_to_host()[0]
+    expected_x = sum(range(num_items))  # 0+1+2+...+9 = 45
+    expected_y = sum(i * 2 for i in range(num_items))  # 0+2+4+...+18 = 90
+
+    assert result["x"] == expected_x, f"Expected x={expected_x}, got {result['x']}"
+    assert result["y"] == expected_y, f"Expected y={expected_y}, got {result['y']}"
+
+
+def test_cpp_op_with_transform_iterator():
+    """Test that RawOp works with TransformIterator."""
+    from cuda.compute import OpKind, TransformIterator
+
+    # C++ unary operator that doubles a value
+    cpp_source = """
+    extern "C" __device__ void double_op(void* input, void* result) {
+        *static_cast<int*>(result) = *static_cast<int*>(input) * 2;
+    }
+    """
+
+    op = make_cpp_op(cpp_source, "double_op")
+
+    # Create input data
+    num_items = 10
+    h_input = np.arange(num_items, dtype=np.int32)
+    d_input = DeviceArray.from_numpy(h_input)
+
+    # Create transform iterator with RawOp
+    transform_iter = TransformIterator(d_input, op, value_type=types.int32)
+
+    # Use the transform iterator with reduce
+    d_output = DeviceArray.empty(1, np.int32)
+    h_init = np.array(0, dtype=np.int32)
+
+    # Sum the doubled values using built-in PLUS operator
+    cuda.compute.reduce_into(
+        d_in=transform_iter,
+        d_out=d_output,
+        num_items=num_items,
+        op=OpKind.PLUS,
+        h_init=h_init,
+    )
+
+    # Verify result: sum of (0*2, 1*2, 2*2, ..., 9*2) = 2 * sum(0..9) = 2 * 45 = 90
+    result = d_output.copy_to_host()[0]
+    expected = 2 * np.sum(h_input)
+    assert result == expected, f"Expected {expected}, got {result}"
+
+
+def test_cpp_stateful_op_reduce_with_constant():
+    """Test stateful RawOp with a simple stateful reduce."""
+    # State: a single int32 constant value (10) on device
+    d_constant = DeviceArray.from_numpy(np.array([10], dtype=np.int32))
+    constant_ptr = d_constant.__cuda_array_interface__["data"][0]
+    state_data = struct.pack("P", constant_ptr)
+    state_alignment = np.dtype(np.intp).alignment
+
+    # C++ operator that adds inputs plus reads constant from state
+    cpp_source = """
+    extern "C" __device__ void add_with_state_constant(void* state, void* a, void* b, void* result) {
+        // Extract constant pointer from state
+        int* constant_ptr = *reinterpret_cast<int**>(state);
+        int constant = *constant_ptr;
+
+        int va = *static_cast<int*>(a);
+        int vb = *static_cast<int*>(b);
+        *static_cast<int*>(result) = va + vb + constant;
+    }
+    """
+
+    op = make_cpp_stateful_op(
+        cpp_source, state_data, "add_with_state_constant", state_alignment
+    )
+
+    # Create test data
+    num_items = 5
+    h_input = np.array([1, 2, 3, 4, 5], dtype=np.int32)
+    d_input = DeviceArray.from_numpy(h_input)
+    d_output = DeviceArray.empty(1, np.int32)
+
+    # Use the stateful op with reduce_into
+    h_init = np.array(0, dtype=np.int32)
+    cuda.compute.reduce_into(
+        d_in=d_input, d_out=d_output, num_items=num_items, op=op, h_init=h_init
+    )
+
+    # Get result
+    result = d_output.copy_to_host()[0]
+    # Each reduction adds 10, so we expect input sum + some multiple of 10
+    # The exact value depends on tree structure, but should be > sum(inputs)
+    sum_inputs = np.sum(h_input)
+    assert result > sum_inputs, f"Expected result > {sum_inputs}, got {result}"
+
+
+def test_cpp_stateful_op_select_with_counter():
+    """Test stateful RawOp with select_if that atomically updates a counter."""
+    # Create a device counter initialized to 0
+    d_counter = DeviceArray.from_numpy(np.zeros(1, dtype=np.int32))
+
+    # State: pointer to the counter
+    counter_ptr = d_counter.__cuda_array_interface__["data"][0]
+    state_data = struct.pack("P", counter_ptr)  # Pack pointer as bytes
+
+    # Use proper pointer alignment for the platform
+    state_alignment = np.dtype(np.intp).alignment
+
+    # C++ select operator that counts selected items
+    # Selects even numbers and atomically increments counter for each selection
+    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;
+    }
+    """
+
+    op = make_cpp_stateful_op(
+        cpp_source,
+        state_data,
+        "select_even_with_count",
+        state_alignment,
+    )
+
+    # Create test data: 0 to 19
+    num_items = 20
+    h_input = np.arange(num_items, dtype=np.int32)
+    d_input = DeviceArray.from_numpy(h_input)
+
+    # Allocate output arrays
+    d_output = DeviceArray.empty(num_items, np.int32)
+    d_num_selected = DeviceArray.empty(1, np.int32)
+
+    # Run select
+    cuda.compute.select(
+        d_in=d_input,
+        d_out=d_output,
+        d_num_selected_out=d_num_selected,
+        cond=op,
+        num_items=num_items,
+    )
+
+    # Get results
+    num_selected = d_num_selected.copy_to_host()[0]
+    counter_value = d_counter.copy_to_host()[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.copy_to_host()[: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"
+    )