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

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

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

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

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

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

42
cccl_upstream/python/cuda_cccl/.gitignore vendored Normal file
View File

@@ -0,0 +1,42 @@
# Python
__pycache__/
*.py[cod]
*$py.class
*.so
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
*.egg-info/
.installed.cfg
*.egg
build.log
.python-version
# CMake
CMakeFiles/
CMakeCache.txt
cmake_install.cmake
Makefile
*.cmake
!CMakeLists.txt
# IDE
.idea/
.vscode/
*.swp
*.swo
# Generated by CMake at build time
cuda/compute/_build_info.py
cuda/compute/cu12/
cuda/compute/cu13/

187
cccl_upstream/python/cuda_cccl/CMakeLists.txt Normal file
View File

@@ -0,0 +1,187 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
cmake_minimum_required(VERSION 3.30)
# Must be set before project() initializes the CUDA language; otherwise CMake
# < 3.23 defaults to sm_52, which is below CCCL's minimum supported arch.
include(../../cmake/CCCLCheckCudaArchitectures.cmake)
set(
CMAKE_CUDA_ARCHITECTURES
"${minimum_cccl_arch}"
CACHE STRING
"CUDA architectures for CCCL"
)
project(cuda_cccl DESCRIPTION "Python package cuda_cccl" LANGUAGES CUDA CXX C)
find_package(CUDAToolkit REQUIRED)
set(CUDA_VERSION_MAJOR ${CUDAToolkit_VERSION_MAJOR})
set(CUDA_VERSION_DIR "cu${CUDA_VERSION_MAJOR}")
message(
STATUS
"Building for CUDA ${CUDA_VERSION_MAJOR}, output directory: ${CUDA_VERSION_DIR}"
)
# Build cuda_cccl against either cccl.c.parallel (v1, NVRTC) by default or
# cccl.c.parallel.v2 (HostJIT) when CCCL_PYTHON_USE_V2=ON. v2 is opt-in until
# it replaces v1 across the matrix.
set(_cccl_root ../..)
set(CCCL_TOPLEVEL_PROJECT ON) # Enable the developer builds
option(
CCCL_PYTHON_USE_V2
"Build cuda_cccl against cccl.c.parallel.v2 (HostJIT)."
OFF
)
if (CCCL_PYTHON_USE_V2)
set(CCCL_ENABLE_C_PARALLEL_V2 ON)
set(CCCL_C_PARALLEL_V2_LIBRARY_OUTPUT_DIRECTORY ${SKBUILD_PROJECT_NAME})
set(_cccl_c_parallel_target cccl.c.parallel.v2)
set(_using_v2_py "True")
else()
set(CCCL_ENABLE_C_PARALLEL ON)
set(CCCL_C_PARALLEL_LIBRARY_OUTPUT_DIRECTORY ${SKBUILD_PROJECT_NAME})
set(_cccl_c_parallel_target cccl.c.parallel)
set(_using_v2_py "False")
endif()
# Surface the v1/v2 choice to Python (tests use it to skip v2-only failures,
# and __init__.py uses it to wire up wheel-bundled hostjit header paths).
# Generated into the build dir and installed via CMake — writing into the
# source tree would miss scikit-build-core's package-file snapshot.
set(_build_info_py "${CMAKE_CURRENT_BINARY_DIR}/_build_info.py")
file(
WRITE "${_build_info_py}"
"# Auto-generated by CMakeLists.txt; do not edit.\nUSING_V2 = ${_using_v2_py}\n"
)
install(FILES "${_build_info_py}" DESTINATION cuda/compute)
# Just install the rest:
set(libcudacxx_ENABLE_INSTALL_RULES ON)
set(CUB_ENABLE_INSTALL_RULES ON)
set(Thrust_ENABLE_INSTALL_RULES ON)
# Install to our output location:
include(GNUInstallDirs)
set(old_libdir "${CMAKE_INSTALL_LIBDIR}") # push
set(old_includedir "${CMAKE_INSTALL_INCLUDEDIR}") # push
set(CMAKE_INSTALL_LIBDIR "cuda/cccl/headers/lib")
set(CMAKE_INSTALL_INCLUDEDIR "cuda/cccl/headers/include")
add_subdirectory(${_cccl_root} _parent_cccl)
set(CMAKE_INSTALL_LIBDIR "${old_libdir}") # pop
set(CMAKE_INSTALL_INCLUDEDIR "${old_includedir}") # pop
# Install version-specific binaries
set(_cccl_c_parallel_install_targets ${_cccl_c_parallel_target})
if (CCCL_PYTHON_USE_V2)
list(APPEND _cccl_c_parallel_install_targets libnvcc)
endif()
install(
TARGETS ${_cccl_c_parallel_install_targets}
DESTINATION cuda/compute/${CUDA_VERSION_DIR}/cccl
)
# Build and install Cython extension
find_package(Python3 COMPONENTS Interpreter Development.Module REQUIRED)
set(CYTHON_version_command "${Python3_EXECUTABLE}" -m cython --version)
execute_process(
COMMAND ${CYTHON_version_command}
OUTPUT_VARIABLE CYTHON_version_output
ERROR_VARIABLE CYTHON_version_output
OUTPUT_STRIP_TRAILING_WHITESPACE
ERROR_STRIP_TRAILING_WHITESPACE
COMMAND_ERROR_IS_FATAL ANY
)
if ("${CYTHON_version_output}" MATCHES "^[Cc]ython version ([^,]+)")
set(CYTHON_VERSION "${CMAKE_MATCH_1}")
else()
message(
FATAL_ERROR
"Failed to parse Cython version from:\n${CYTHON_version_output}"
)
endif()
# -3 generates source for Python 3
# -M generates depfile
# -t cythonizes if PYX is newer than preexisting output
# -w sets working directory
set(
CYTHON_FLAGS
-3
-M
-t
-w
"${cuda_cccl_SOURCE_DIR}"
)
message(STATUS "Using Cython ${CYTHON_VERSION}")
set(pyx_source_file "${cuda_cccl_SOURCE_DIR}/cuda/compute/_bindings_impl.pyx")
set(_generated_extension_src "${cuda_cccl_BINARY_DIR}/_bindings_impl.c")
set(_depfile "${cuda_cccl_BINARY_DIR}/_bindings_impl.c.dep")
# Backend-conditional Cython .pxi files. Where v1 and v2 expose different
# struct layouts or call signatures, the .pyx `include`s a generated .pxi
# whose source is chosen here. The helpers inside present a uniform interface
# so the rest of _bindings_impl.pyx stays backend-agnostic.
if (CCCL_PYTHON_USE_V2)
set(_backend_suffix "v2")
else()
set(_backend_suffix "v1")
endif()
foreach (
_pxi_stem
segmented_reduce_backend
binary_search_backend
op_code_type
serialization
)
configure_file(
"${CMAKE_CURRENT_SOURCE_DIR}/cuda/compute/_bindings_${_pxi_stem}_${_backend_suffix}.pxi"
"${CMAKE_CURRENT_BINARY_DIR}/_bindings_${_pxi_stem}.pxi"
COPYONLY
)
endforeach()
# Custom Cython compilation command. `-I ${BINARY_DIR}` lets the .pyx's
# `include "_bindings_..._backend.pxi"` resolve to the file we configured
# above.
add_custom_command(
OUTPUT "${_generated_extension_src}"
COMMAND
"${Python3_EXECUTABLE}" -m cython
# gersemi: off
${CYTHON_FLAGS}
-I "${CMAKE_CURRENT_BINARY_DIR}"
"${pyx_source_file}"
--output-file "${_generated_extension_src}"
# gersemi: on
DEPENDS "${pyx_source_file}"
DEPFILE "${_depfile}"
COMMENT "Cythonizing ${pyx_source_file} for CUDA ${CUDA_VERSION_MAJOR}"
)
add_custom_target(
cythonize_bindings_impl
ALL
DEPENDS "${_generated_extension_src}"
)
python3_add_library(
_bindings_impl
MODULE
WITH_SOABI
"${_generated_extension_src}"
)
add_dependencies(_bindings_impl cythonize_bindings_impl)
target_link_libraries(
_bindings_impl
PRIVATE #
${_cccl_c_parallel_target}
CUDA::cuda_driver
)
set_target_properties(_bindings_impl PROPERTIES INSTALL_RPATH "$ORIGIN/cccl")
install(TARGETS _bindings_impl DESTINATION cuda/compute/${CUDA_VERSION_DIR})

1
cccl_upstream/python/cuda_cccl/LICENSE Normal file
View File

@@ -0,0 +1 @@
../../LICENSE

49
cccl_upstream/python/cuda_cccl/README.md Normal file
View File

@@ -0,0 +1,49 @@
# CUDA CCCL Python Package
[`cuda.cccl`](https://nvidia.github.io/cccl/unstable/python)
provides a Pythonic interface to the
[CUDA Core Compute Libraries](https://nvidia.github.io/cccl/unstable/cpp.html#cccl-cpp-libraries).
It provides the following modules:
- **`cuda.compute`** - Device-level parallel algorithms (reduce, scan, sort, etc.) and iterators
- **`cuda.cccl.headers`** - Programmatic access to CCCL headers
## Installation
Install from PyPI:
```bash
pip install cuda-cccl[cu13] # For CUDA 13.x (pip-installed cuda-toolkit)
pip install cuda-cccl[cu12] # For CUDA 12.x (pip-installed cuda-toolkit)
```
If you already have a CUDA toolkit on your system and do not want pip to
install it, use the `sysctk` variants:
```bash
pip install cuda-cccl[sysctk13] # For CUDA 13.x (system CUDA toolkit)
pip install cuda-cccl[sysctk12] # For CUDA 12.x (system CUDA toolkit)
```
For a minimal install without Numba (useful when supplying pre-compiled operators):
```bash
pip install cuda-cccl[minimal-cu13] # pip-installed cuda-toolkit
pip install cuda-cccl[minimal-sysctk13] # system CUDA toolkit
```
Install from conda-forge:
```bash
conda install -c conda-forge cccl-python
```
**Requirements:** Python 3.10+, CUDA Toolkit 12.x or 13.x, NVIDIA GPU with Compute Capability 7.5+
## Documentation
For complete documentation, examples, and API reference, visit:
- **Full Documentation**: [nvidia.github.io/cccl/unstable/python](https://nvidia.github.io/cccl/unstable/python)
- **Repository**: [github.com/NVIDIA/cccl](https://github.com/NVIDIA/cccl)
- **Examples**: [github.com/NVIDIA/cccl/tree/main/python/cuda_cccl/tests/compute/examples](https://github.com/NVIDIA/cccl/tree/main/python/cuda_cccl/tests/compute/examples)

19
cccl_upstream/python/cuda_cccl/benchmarks/compute/.gitignore vendored Normal file
View File

@@ -0,0 +1,19 @@
# Build artifacts
build/
bin/
*.o
*.so
*.a
compile_commands.json
# Python cache
__pycache__/
*.pyc
*.pyo
# Benchmark results
results*/
# Pixi
.pixi/
pixi.lock

12
cccl_upstream/python/cuda_cccl/benchmarks/compute/AGENTS.md Normal file
View File

@@ -0,0 +1,12 @@
# CCCL Python cuda.compute benchmarks
This directory contains the code for the Python cuda.compute benchmarks.
They are migrated from the original C++ benchmarks and they should match the C++ implementations
as closely as possible.
The original C++ benchmarks are available in this repository in: ../../../../cub/benchmarks/bench/
We follow the same directory structure and naming conventions converting to Python were appropriate.
The code for cuda.compute is in this repository under: `../../../../python/cuda_cccl/cuda/compute/`. Look into this directory when searching for existing APIs in Python.
The benchmarks use nvbench to run the benchmarks and report the results.

147
cccl_upstream/python/cuda_cccl/benchmarks/compute/README.md Normal file
View File

@@ -0,0 +1,147 @@
# cuda.compute Benchmarks
Compare Python `cuda.compute` performance against C++ CUB implementations.
## Setup
This project uses [pixi](https://pixi.sh) to manage environments and dependencies.
Two environments are available:
- **`wheel`** - Uses the released `cuda-cccl` package
- **`source`** - Builds `cuda-cccl` from the local repository
### Build C++ Benchmarks
Build CUB benchmarks using the CI script (one-time, ~13 minutes):
```bash
cd /path/to/cccl
./ci/build_cub.sh -arch 89 # Use your GPU arch (89=RTX 4090, 80=A100, 90=H100)
```
Binaries are built to: `build/cub/bin/`
## Run Benchmarks
### Using pixi tasks
```bash
# Run Python benchmarks (released cuda-cccl)
pixi run -e wheel bench
# Run Python benchmarks (local source build)
pixi run -e source bench
# Run Python benchmarks with reduced parameter set
pixi run -e wheel bench-quick
# Run just one benchmark
pixi run -e wheel bench -b transform/fill
# Run C++ benchmarks
pixi run -e wheel bench-cpp
# Run both Python and C++ benchmarks
pixi run -e wheel bench-all
```
### Using run_benchmarks.py directly
```bash
# Run both C++ and Python (default)
pixi run -e wheel python run_benchmarks.py -b transform/fill -d 0
# Run only C++
pixi run -e wheel python run_benchmarks.py -b transform/fill --cpp
# Run only Python
pixi run -e wheel python run_benchmarks.py -b transform/fill --py
# Show help
pixi run -e wheel python run_benchmarks.py --help
```
To run the benchmarks using the "quick" configuration:
```bash
pixi run -e wheel python run_benchmarks.py --quick
```
## Compare Results
```bash
pixi run -e wheel python analysis/python_vs_cpp_summary.py -b transform/fill
```
## Web Report
A simple page used to visualize a set of results.
- Requires `results/` to be populated with benchmark results.
First generate a manifest:
```bash
pixi run -e wheel python analysis/generate_web_report_manifest.py \
--results-dir results \
--output results/manifest.json
```
Build the web report single file app:
```bash
cd analysis/web-report
npm install
npm run build
```
This will output a single file app to `analysis/web-report/dist/` copy it to the `results/` directory and:
```bash
cd results/
python3 -m http.server
```
Now its possible to share the results directory as a zip/tar file.
## Manual Usage
### List benchmark configurations
```bash
# Python
pixi run -e wheel python transform/fill.py --list
# C++
/path/to/cccl/build/cub/bin/cub.bench.transform.fill.base --list
```
### Run with custom options
```bash
# Python - specific type and size
pixi run -e wheel python transform/fill.py --axis "T=I32" --axis "Elements[pow2]=20" --devices 0
# C++ - save JSON
/path/to/cccl/build/cub/bin/cub.bench.transform.fill.base \
--json results/transform/fill_cpp.json \
--devices 0
```
### Compare manually
```bash
pixi run -e wheel python analysis/python_vs_cpp_summary.py \
results/transform/fill_py.json \
results/transform/fill_cpp.json \
--device 0
```
## AI commands
These are using the .opencode folder but can be moved to other Agents.
### /migration-status
Generates a report of the migration status for each benchmark in CUB.

132
cccl_upstream/python/cuda_cccl/benchmarks/compute/histogram/even.py Normal file
View File

@@ -0,0 +1,132 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for histogram_even using cuda.compute.
C++ equivalent: cub/benchmarks/bench/histogram/even.cu
Notes:
- The C++ benchmark uses Entropy axis with nvbench_helper bit entropy generation
- Migration: Python matches the bitwise-AND entropy approach and skips some I8/I16 large-bin cases due to CUDA errors.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import cupy as cp
import numpy as np
from utils import FUNDAMENTAL_TYPES as TYPE_MAP
from utils import as_cupy_stream, generate_data_with_entropy
import cuda.bench as bench
from cuda.compute import make_histogram_even
def get_upper_level(dtype, num_bins, num_elements):
"""
Compute upper level for histogram bins.
Mirrors C++ get_upper_level() from histogram_common.cuh
"""
if np.issubdtype(dtype, np.integer):
# For integer types, upper_level = min(num_bins, max_value_for_type)
max_val = np.iinfo(dtype).max
return dtype(min(num_bins, max_val))
else:
# For floating point types, upper_level = num_elements
return dtype(num_elements)
def bench_histogram_even(state: bench.State):
type_str = state.get_string("SampleT{ct}")
dtype = TYPE_MAP[type_str]
num_elements = int(state.get_int64("Elements{io}"))
num_bins = int(state.get_int64("Bins"))
entropy_str = state.get_string("Entropy")
# Skip invalid configurations (like C++ does)
# For integer types, skip if num_bins > max value representable by SampleT
if np.issubdtype(dtype, np.integer):
max_val = np.iinfo(dtype).max
if num_bins > max_val:
state.skip("Number of bins exceeds what SampleT can represent")
return
num_levels = num_bins + 1
lower_level = dtype(0)
upper_level = get_upper_level(dtype, num_bins, num_elements)
alloc_stream = as_cupy_stream(state.get_stream())
d_samples = generate_data_with_entropy(
num_elements,
dtype,
entropy_str,
alloc_stream,
min_val=lower_level,
max_val=upper_level,
)
# Output histogram (counter type is int32 in C++)
with alloc_stream:
d_histogram = cp.zeros(num_bins, dtype=np.int32)
alloc_stream.synchronize()
h_num_output_levels = np.array([num_levels], dtype=np.int32)
h_lower_level = np.array([lower_level], dtype=dtype)
h_upper_level = np.array([upper_level], dtype=dtype)
histogrammer = 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_elements,
)
temp_storage_bytes = 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=num_elements,
)
with alloc_stream:
temp_storage = cp.empty(temp_storage_bytes, dtype=np.uint8)
state.add_element_count(num_elements)
state.add_global_memory_reads(num_elements * d_samples.dtype.itemsize)
state.add_global_memory_writes(num_bins * d_histogram.dtype.itemsize)
def launcher(launch: bench.Launch):
histogrammer(
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_elements,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
if __name__ == "__main__":
b = bench.register(bench_histogram_even)
b.set_name("base")
b.add_string_axis("SampleT{ct}", list(TYPE_MAP.keys()))
b.add_int64_power_of_two_axis("Elements{io}", range(16, 29, 4))
b.add_int64_axis("Bins", [32, 128, 2048, 2097152])
b.add_string_axis("Entropy", ["0.201", "1.000"])
bench.run_all_benchmarks(sys.argv)

1562
cccl_upstream/python/cuda_cccl/benchmarks/compute/host/host_benchmark_cases.py Normal file
View File

File diff suppressed because it is too large Load Diff

95
cccl_upstream/python/cuda_cccl/benchmarks/compute/host/test_host_pytest_benchmark.py Normal file
View File

@@ -0,0 +1,95 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
import pytest
from host_benchmark_cases import (
CALL_CASES,
CASES,
HostBenchmarkCase,
patch_wrapper_to_skip_native_compute,
synchronize,
)
import cuda.compute as cc
pytest.importorskip("pytest_benchmark")
BUILD_TIME_ROUNDS = 10
ONESHOT_ROUNDS = 20
ONESHOT_ITERATIONS = 100
TWOSHOT_ROUNDS = 20
TWOSHOT_ITERATIONS = 1000
def _case_params(cases: list[HostBenchmarkCase]) -> list[pytest.ParameterSet]:
params = []
for case in cases:
marks = []
if case.skip_reason is not None:
marks.append(pytest.mark.skip(reason=case.skip_reason))
params.append(pytest.param(case, id=case.name, marks=marks))
return params
@pytest.mark.benchmark(group="cuda.compute.host.build_time")
@pytest.mark.parametrize("case", _case_params(CASES))
def test_build_time(benchmark, case: HostBenchmarkCase):
state = case.setup()
synchronize()
def setup() -> None:
cc.clear_all_caches()
def build():
return case.make_wrapper(state)
benchmark.pedantic(
build,
setup=setup,
rounds=BUILD_TIME_ROUNDS,
iterations=1,
warmup_rounds=0,
)
@pytest.mark.benchmark(group="cuda.compute.host.oneshot_cached")
@pytest.mark.parametrize("case", _case_params(CALL_CASES))
def test_oneshot_cached_host_overhead(benchmark, case: HostBenchmarkCase):
cc.clear_all_caches()
state = case.setup()
wrapper = case.make_wrapper(state)
patch_wrapper_to_skip_native_compute(wrapper, case.noop_return_kind)
synchronize()
def call() -> None:
case.oneshot(state)
benchmark.pedantic(
call,
rounds=ONESHOT_ROUNDS,
iterations=ONESHOT_ITERATIONS,
warmup_rounds=0,
)
@pytest.mark.benchmark(group="cuda.compute.host.twoshot_call")
@pytest.mark.parametrize("case", _case_params(CALL_CASES))
def test_twoshot_call_host_overhead(benchmark, case: HostBenchmarkCase):
cc.clear_all_caches()
state = case.setup()
wrapper = case.make_wrapper(state)
patch_wrapper_to_skip_native_compute(wrapper, case.noop_return_kind)
synchronize()
def call() -> None:
case.twoshot(state, wrapper)
benchmark.pedantic(
call,
rounds=TWOSHOT_ROUNDS,
iterations=TWOSHOT_ITERATIONS,
warmup_rounds=0,
)

97
cccl_upstream/python/cuda_cccl/benchmarks/compute/merge_sort/keys.py Normal file
View File

@@ -0,0 +1,97 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for merge_sort keys using cuda.compute.
C++ equivalent: cub/benchmarks/bench/merge_sort/keys.cu
Notes:
- The C++ benchmark uses Entropy axis to control data distribution
- Uses less_t comparison operator (ascending sort)
- Keys only (no values) - see pairs.cu for key-value sorting
- Migration: Python fixes offsets and approximates entropy generation.
- OffsetT axis is omitted because the Python API does not expose offset type.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import cupy as cp
import numpy as np
from utils import SIGNED_TYPES, as_cupy_stream, generate_data_with_entropy
import cuda.bench as bench
from cuda.compute import OpKind, make_merge_sort
def bench_merge_sort_keys(state: bench.State):
type_str = state.get_string("T{ct}")
dtype = SIGNED_TYPES[type_str]
num_elements = int(state.get_int64("Elements{io}"))
entropy_str = state.get_string("Entropy")
alloc_stream = as_cupy_stream(state.get_stream())
d_in_keys = generate_data_with_entropy(
num_elements, dtype, entropy_str, alloc_stream
)
# Output array for sorted keys (merge_sort requires separate output)
with alloc_stream:
d_out_keys = cp.empty(num_elements, dtype=dtype)
alloc_stream.synchronize()
sorter = 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_elements,
)
with alloc_stream:
temp_storage = cp.empty(temp_storage_bytes, dtype=np.uint8)
state.add_element_count(num_elements)
state.add_global_memory_reads(num_elements * d_in_keys.dtype.itemsize, "Size")
state.add_global_memory_writes(num_elements * d_out_keys.dtype.itemsize)
def launcher(launch: bench.Launch):
sorter(
temp_storage=temp_storage,
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_elements,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
if __name__ == "__main__":
b = bench.register(bench_merge_sort_keys)
b.set_name("base")
b.add_string_axis("T{ct}", list(SIGNED_TYPES.keys()))
b.add_int64_power_of_two_axis("Elements{io}", range(16, 29, 4))
b.add_string_axis("Entropy", ["1.000", "0.201"])
# Note: OffsetT axis from C++ is not exposed in Python API
bench.run_all_benchmarks(sys.argv)

110
cccl_upstream/python/cuda_cccl/benchmarks/compute/merge_sort/pairs.py Normal file
View File

@@ -0,0 +1,110 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for merge_sort pairs using cuda.compute.
C++ equivalent: cub/benchmarks/bench/merge_sort/pairs.cu
Notes:
- Uses Entropy axis to control key distribution
- Keys and values are sorted together
- Migration: Python omits int128 values and OffsetT axis.
- OffsetT axis is omitted because the Python API does not expose offset type.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import cupy as cp
import numpy as np
from utils import (
INTEGRAL_TYPES,
SIGNED_TYPES,
as_cupy_stream,
generate_data_with_entropy,
)
import cuda.bench as bench
from cuda.compute import OpKind, make_merge_sort
KEY_TYPE_MAP = SIGNED_TYPES
VALUE_TYPE_MAP = INTEGRAL_TYPES
def bench_merge_sort_pairs(state: bench.State):
key_type_str = state.get_string("KeyT{ct}")
value_type_str = state.get_string("ValueT{ct}")
key_dtype = KEY_TYPE_MAP[key_type_str]
value_dtype = VALUE_TYPE_MAP[value_type_str]
num_elements = int(state.get_int64("Elements{io}"))
entropy_str = state.get_string("Entropy")
alloc_stream = as_cupy_stream(state.get_stream())
d_in_keys = generate_data_with_entropy(
num_elements, key_dtype, entropy_str, alloc_stream
)
with alloc_stream:
d_in_values = generate_data_with_entropy(
num_elements, value_dtype, "1.000", alloc_stream
)
d_out_keys = cp.empty(num_elements, dtype=key_dtype)
d_out_values = cp.empty(num_elements, dtype=value_dtype)
alloc_stream.synchronize()
sorter = 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,
)
temp_storage_bytes = 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,
op=OpKind.LESS,
num_items=num_elements,
)
with alloc_stream:
temp_storage = cp.empty(temp_storage_bytes, dtype=np.uint8)
state.add_element_count(num_elements)
state.add_global_memory_reads(num_elements * d_in_keys.dtype.itemsize)
state.add_global_memory_reads(num_elements * d_in_values.dtype.itemsize)
state.add_global_memory_writes(num_elements * d_out_keys.dtype.itemsize)
state.add_global_memory_writes(num_elements * d_out_values.dtype.itemsize)
def launcher(launch: bench.Launch):
sorter(
temp_storage=temp_storage,
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,
num_items=num_elements,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
if __name__ == "__main__":
b = bench.register(bench_merge_sort_pairs)
b.set_name("base")
b.add_string_axis("KeyT{ct}", list(KEY_TYPE_MAP.keys()))
b.add_string_axis("ValueT{ct}", list(VALUE_TYPE_MAP.keys()))
b.add_int64_power_of_two_axis("Elements{io}", range(16, 29, 4))
b.add_string_axis("Entropy", ["1.000", "0.201"])
bench.run_all_benchmarks(sys.argv)

138
cccl_upstream/python/cuda_cccl/benchmarks/compute/partition/three_way.py Normal file
View File

@@ -0,0 +1,138 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for three_way_partition using cuda.compute.
C++ equivalent: cub/benchmarks/bench/partition/three_way.cu
Notes:
- The C++ benchmark uses Entropy axis to control data distribution
- Uses less_then_t<T> predicate operators to divide data into three partitions:
- First partition: items < left_border (max/3)
- Second partition: items < right_border (max*2/3)
- Third partition (unselected): items >= right_border
- T axis covers fundamental types (C++ fundamental_types minus int128)
- Migration: Python uses FUNDAMENTAL_TYPES; omits OffsetT axis.
- OffsetT axis is omitted because the Python API does not expose offset type.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import cupy as cp
import numpy as np
from utils import FUNDAMENTAL_TYPES, as_cupy_stream, generate_data_with_entropy
import cuda.bench as bench
from cuda.compute import make_three_way_partition
def bench_three_way_partition(state: bench.State):
type_str = state.get_string("T{ct}")
dtype = FUNDAMENTAL_TYPES[type_str]
num_elements = int(state.get_int64("Elements{io}"))
entropy_str = state.get_string("Entropy")
alloc_stream = as_cupy_stream(state.get_stream())
if np.issubdtype(dtype, np.integer):
info = np.iinfo(dtype)
min_val = 0
max_val = info.max
else:
info = np.finfo(dtype)
min_val = 0.0
max_val = info.max
left_border = max_val // 3 if np.issubdtype(dtype, np.integer) else max_val / 3
right_border = left_border * 2
d_in = generate_data_with_entropy(
num_elements,
dtype,
entropy_str,
alloc_stream,
min_val=min_val,
max_val=max_val,
)
with alloc_stream:
d_first_part_out = cp.empty(num_elements, dtype=dtype)
d_second_part_out = cp.empty(num_elements, dtype=dtype)
d_unselected_out = cp.empty(num_elements, dtype=dtype)
# d_num_selected_out stores [num_first_part, num_second_part]
d_num_selected_out = cp.empty(2, dtype=np.int32)
alloc_stream.synchronize()
# Convert borders to the correct type for closure capture
left_thresh = dtype(left_border)
right_thresh = dtype(right_border)
def select_first_part(x):
return x < left_thresh
def select_second_part(x):
return x < right_thresh
partitioner = make_three_way_partition(
d_in=d_in,
d_first_part_out=d_first_part_out,
d_second_part_out=d_second_part_out,
d_unselected_out=d_unselected_out,
d_num_selected_out=d_num_selected_out,
select_first_part_op=select_first_part,
select_second_part_op=select_second_part,
)
temp_storage_bytes = partitioner(
temp_storage=None,
d_in=d_in,
d_first_part_out=d_first_part_out,
d_second_part_out=d_second_part_out,
d_unselected_out=d_unselected_out,
d_num_selected_out=d_num_selected_out,
select_first_part_op=select_first_part,
select_second_part_op=select_second_part,
num_items=num_elements,
)
with alloc_stream:
temp_storage = cp.empty(temp_storage_bytes, dtype=np.uint8)
state.add_element_count(num_elements)
state.add_global_memory_reads(num_elements * d_in.dtype.itemsize)
state.add_global_memory_writes(num_elements * d_in.dtype.itemsize)
# C++ reports add_global_memory_writes<offset_t>(1) — 1 element of offset type.
state.add_global_memory_writes(d_num_selected_out.dtype.itemsize)
def launcher(launch: bench.Launch):
partitioner(
temp_storage=temp_storage,
d_in=d_in,
d_first_part_out=d_first_part_out,
d_second_part_out=d_second_part_out,
d_unselected_out=d_unselected_out,
d_num_selected_out=d_num_selected_out,
select_first_part_op=select_first_part,
select_second_part_op=select_second_part,
num_items=num_elements,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
if __name__ == "__main__":
b = bench.register(bench_three_way_partition)
b.set_name("base")
b.add_string_axis("T{ct}", list(FUNDAMENTAL_TYPES.keys()))
b.add_int64_power_of_two_axis("Elements{io}", range(16, 29, 4))
b.add_string_axis("Entropy", ["1.000", "0.544", "0.000"])
# Note: OffsetT axis from C++ is not exposed in Python API
bench.run_all_benchmarks(sys.argv)

89
cccl_upstream/python/cuda_cccl/benchmarks/compute/pixi.toml Normal file
View File

@@ -0,0 +1,89 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0
[workspace]
channels = ["conda-forge"]
platforms = ["linux-64"]
channel-priority = "disabled"
# CUDA 13.1 system requirement
[feature.cu13.system-requirements]
cuda = "13"
[feature.cu13.dependencies]
cuda-version = "13.1.*"
# Python benchmark dependencies
[feature.bench.dependencies]
python = "3.13.*"
numpy = "*"
cupy = "*"
pytest-benchmark = "*"
pyyaml = "*"
pre-commit = "*"
[feature.bench.pypi-dependencies]
cuda-bench = ">=0.2.0"
# C++ benchmark build dependencies
[feature.cpp-bench.dependencies]
cmake = "*"
ninja = "*"
cxx-compiler = "*"
fmt = "*"
cuda-cudart-dev = "*"
[feature.cpp-bench.target.linux-64.dependencies]
cuda-crt-dev_linux-64 = "*"
cuda-driver-dev_linux-64 = "*"
[feature.cpp-bench.target.linux-64.activation.env]
CUDA_HOME = "$CONDA_PREFIX/targets/x86_64-linux"
# Important: cuda-cccl installation variants
# Released version from PyPI
[feature.cccl-wheel.pypi-dependencies]
cuda-cccl = { version = ">=0.1.0", extras = ["cu13"] }
# Local source build (editable install from repo)
# Needs nvcc, compilers, and CUDA dev libraries for scikit-build-core to build cuda-cccl
[feature.cccl-source.dependencies]
cuda-nvcc = "*"
cuda-nvrtc-dev = "*"
libnvjitlink-dev = "*"
cuda-cudart-dev = "*"
cuda-driver-dev = "*"
c-compiler = "*"
cxx-compiler = "*"
[feature.cccl-source.pypi-dependencies]
cuda-cccl = { path = "../..", editable = true, extras = ["cu13"] }
# Environments
[environments]
wheel = { features = ["cu13", "bench", "cpp-bench", "cccl-wheel"] }
source = { features = ["cu13", "bench", "cpp-bench", "cccl-source"] }
# Tasks
[tasks.bench]
cmd = ["python", "run_benchmarks.py", "--py"]
description = "Run Python cuda.compute benchmarks"
[tasks.bench-quick]
cmd = ["python", "run_benchmarks.py", "--py", "--quick"]
description = "Run Python benchmarks with reduced parameter set"
[tasks.bench-cpp]
cmd = ["python", "run_benchmarks.py", "--cpp"]
description = "Run C++ CUB benchmarks"
[tasks.bench-all]
cmd = ["python", "run_benchmarks.py"]
description = "Run both Python and C++ benchmarks"
[tasks.pre-commit]
cmd = ["pre-commit", "run", "--all-files"]
cwd = "../../../.."
description = "Run pre-commit checks on the entire repo"

130
cccl_upstream/python/cuda_cccl/benchmarks/compute/quick_configs.yaml Normal file
View File

@@ -0,0 +1,130 @@
# Quick mode configurations for fast benchmark testing
# Used by run_benchmarks.py with --quick flag
#
# Each benchmark lists axis name -> value pairs.
# Use C++ axis names (with {ct}/{io} suffixes where applicable).
# The script strips these suffixes for Python benchmarks.
#
# For power-of-two axes (Elements{io}, MaxSegSize, MaxSegmentSize, SegmentSize, Segments{io}),
# values are exponents (e.g., 16 means 2^16 = 65536). SegmentSize is converted to
# actual values for Python benchmarks.
#
# NOTE: Only specify axes that exist in BOTH C++ and Python benchmarks.
# Some C++ benchmarks have extra axes (like OffsetT{ct}) not in Python.
transform/fill:
"T{ct}": "I32"
"Elements{io}": 16
transform/babelstream:
"T{ct}": "F32"
"Elements{io}": 16
transform/heavy:
"Heaviness{ct}": "64"
"Elements{io}": 16
transform/fib:
"Elements{io}": 16
transform/grayscale:
"T{ct}": "F32"
"Elements{io}": 16
transform/complex_cmp:
"Elements{io}": 16
transform_reduce/sum:
"T{ct}": "I32"
"Elements{io}": 16
reduce/sum:
"T{ct}": "I32"
"Elements{io}": 16
reduce/min:
"T{ct}": "I32"
"Elements{io}": 16
reduce/custom:
"T{ct}": "I32"
"Elements{io}": 16
reduce/nondeterministic:
"T{ct}": "I32"
"Elements{io}": 16
scan/exclusive/sum:
"T{ct}": "I32"
"Elements{io}": 16
scan/exclusive/custom:
"T{ct}": "I32"
"Elements{io}": 16
histogram/even:
"SampleT{ct}": "I32"
"Elements{io}": 16
"Bins": 128
"Entropy": "1.000"
select/if:
"T{ct}": "I32"
"Elements{io}": 16
"Entropy": "1.000"
select/unique_by_key:
"KeyT{ct}": "I32"
"ValueT{ct}": "I32"
"Elements{io}": 16
"MaxSegSize": 4
radix_sort/keys:
"T{ct}": "I32"
"Elements{io}": 16
"Entropy": "1.000"
radix_sort/pairs:
"KeyT{ct}": "I32"
"ValueT{ct}": "I32"
"Elements{io}": 16
"Entropy": "1.000"
merge_sort/keys:
"T{ct}": "I32"
"Elements{io}": 16
"Entropy": "1.000"
merge_sort/pairs:
"KeyT{ct}": "I32"
"ValueT{ct}": "I32"
"Elements{io}": 16
"Entropy": "1.000"
segmented_sort/keys:
benchmarks:
power:
"T{ct}": "I32"
"Elements{io}": 22
"Segments{io}": 12
"Entropy": "1.000"
small:
"T{ct}": "I32"
"Elements{io}": 22
"MaxSegmentSize": 1
large:
"T{ct}": "I32"
"Elements{io}": 22
"MaxSegmentSize": 10
segmented_reduce/variable_sum:
benchmarks:
variable_default:
"T{ct}": "I32"
"Elements{io}": 16
"MaxSegmentSize": 4
partition/three_way:
"T{ct}": "I32"
"Elements{io}": 16
"Entropy": "1.000"

97
cccl_upstream/python/cuda_cccl/benchmarks/compute/radix_sort/keys.py Normal file
View File

@@ -0,0 +1,97 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for radix_sort keys using cuda.compute.
C++ equivalent: cub/benchmarks/bench/radix_sort/keys.cu
Notes:
- The C++ benchmark uses Entropy axis to control data distribution
- Sort order is always ascending (C++ benchmark hardcodes this)
- Keys only (no values) - see radix_sort/pairs.cu for key-value sorting
- begin_bit=0, end_bit=sizeof(T)*8 (full key comparison)
- Migration: Python fixes offsets, excludes int128, and approximates entropy generation.
- OffsetT axis is omitted because the Python API does not expose offset type.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import cupy as cp
import numpy as np
from utils import FUNDAMENTAL_TYPES as TYPE_MAP
from utils import as_cupy_stream, generate_data_with_entropy
import cuda.bench as bench
from cuda.compute import SortOrder, make_radix_sort
def bench_radix_sort_keys(state: bench.State):
type_str = state.get_string("T{ct}")
dtype = TYPE_MAP[type_str]
num_elements = int(state.get_int64("Elements{io}"))
entropy_str = state.get_string("Entropy")
alloc_stream = as_cupy_stream(state.get_stream())
d_in_keys = generate_data_with_entropy(
num_elements, dtype, entropy_str, alloc_stream
)
# Output array for sorted keys
with alloc_stream:
d_out_keys = cp.empty(num_elements, dtype=dtype)
alloc_stream.synchronize()
sorter = make_radix_sort(
d_in_keys=d_in_keys,
d_out_keys=d_out_keys,
d_in_values=None,
d_out_values=None,
order=SortOrder.ASCENDING,
)
temp_storage_bytes = sorter(
temp_storage=None,
d_in_keys=d_in_keys,
d_out_keys=d_out_keys,
d_in_values=None,
d_out_values=None,
num_items=num_elements,
)
with alloc_stream:
temp_storage = cp.empty(temp_storage_bytes, dtype=np.uint8)
state.add_element_count(num_elements)
state.add_global_memory_reads(num_elements * d_in_keys.dtype.itemsize, "Size")
state.add_global_memory_writes(num_elements * d_out_keys.dtype.itemsize)
def launcher(launch: bench.Launch):
sorter(
temp_storage=temp_storage,
d_in_keys=d_in_keys,
d_out_keys=d_out_keys,
d_in_values=None,
d_out_values=None,
num_items=num_elements,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
if __name__ == "__main__":
b = bench.register(bench_radix_sort_keys)
b.set_name("base")
b.add_string_axis("T{ct}", list(TYPE_MAP.keys()))
b.add_int64_power_of_two_axis("Elements{io}", range(16, 29, 4))
b.add_string_axis("Entropy", ["1.000", "0.544", "0.201"])
# Note: OffsetT axis from C++ is not exposed in Python API
bench.run_all_benchmarks(sys.argv)

109
cccl_upstream/python/cuda_cccl/benchmarks/compute/radix_sort/pairs.py Normal file
View File

@@ -0,0 +1,109 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for radix_sort pairs (key-value) using cuda.compute.
C++ equivalent: cub/benchmarks/bench/radix_sort/pairs.cu
Notes:
- The C++ benchmark uses Entropy axis to control key data distribution
- Sort order is always ascending (C++ benchmark hardcodes this)
- Keys and values are sorted together (values rearranged by key order)
- begin_bit=0, end_bit=sizeof(KeyT)*8 (full key comparison)
- C++ uses integral_types for keys and int8/16/32/64(+int128) for values
- Migration: Python matches C++ integral_types for both keys and values; omits int128 and OffsetT axis.
- OffsetT axis is omitted because the Python API does not expose offset type.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import cupy as cp
import numpy as np
from utils import INTEGRAL_TYPES, as_cupy_stream, generate_data_with_entropy
import cuda.bench as bench
from cuda.compute import SortOrder, make_radix_sort
KEY_TYPE_MAP = INTEGRAL_TYPES
VALUE_TYPE_MAP = INTEGRAL_TYPES
def bench_radix_sort_pairs(state: bench.State):
key_type_str = state.get_string("KeyT{ct}")
value_type_str = state.get_string("ValueT{ct}")
key_dtype = KEY_TYPE_MAP[key_type_str]
value_dtype = VALUE_TYPE_MAP[value_type_str]
num_elements = int(state.get_int64("Elements{io}"))
entropy_str = state.get_string("Entropy")
alloc_stream = as_cupy_stream(state.get_stream())
d_in_keys = generate_data_with_entropy(
num_elements, key_dtype, entropy_str, alloc_stream
)
d_in_values = generate_data_with_entropy(
num_elements, value_dtype, "1.000", alloc_stream
)
with alloc_stream:
d_out_keys = cp.empty(num_elements, dtype=key_dtype)
d_out_values = cp.empty(num_elements, dtype=value_dtype)
alloc_stream.synchronize()
sorter = 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,
)
temp_storage_bytes = 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_elements,
)
with alloc_stream:
temp_storage = cp.empty(temp_storage_bytes, dtype=np.uint8)
state.add_element_count(num_elements)
state.add_global_memory_reads(num_elements * d_in_keys.dtype.itemsize)
state.add_global_memory_reads(num_elements * d_in_values.dtype.itemsize)
state.add_global_memory_writes(num_elements * d_out_keys.dtype.itemsize)
state.add_global_memory_writes(num_elements * d_out_values.dtype.itemsize)
def launcher(launch: bench.Launch):
sorter(
temp_storage=temp_storage,
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_elements,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
if __name__ == "__main__":
b = bench.register(bench_radix_sort_pairs)
b.set_name("base")
b.add_string_axis("KeyT{ct}", list(KEY_TYPE_MAP.keys()))
b.add_string_axis("ValueT{ct}", list(VALUE_TYPE_MAP.keys()))
b.add_int64_power_of_two_axis("Elements{io}", range(16, 29, 4))
b.add_string_axis("Entropy", ["1.000", "0.201"])
# Note: OffsetT axis from C++ is not exposed in Python API
bench.run_all_benchmarks(sys.argv)

82
cccl_upstream/python/cuda_cccl/benchmarks/compute/reduce/custom.py Normal file
View File

@@ -0,0 +1,82 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for reduce custom operation using cuda.compute.reduce_into.
C++ equivalent: cub/benchmarks/bench/reduce/custom.cu
Notes:
- Uses a custom max operator (not OpKind) to exercise generic path
- int128 and complex32 are not supported by cupy
- Migration: Python limits to basic numeric types; C++ includes int128/complex.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import cupy as cp
import numpy as np
from utils import SIGNED_TYPES as TYPE_MAP
from utils import as_cupy_stream, generate_data_with_entropy
import cuda.bench as bench
from cuda.compute import make_reduce_into
def max_op(a, b):
return a if a > b else b
def bench_reduce_custom(state: bench.State):
type_str = state.get_string("T{ct}")
dtype = TYPE_MAP[type_str]
num_items = int(state.get_int64("Elements{io}"))
alloc_stream = as_cupy_stream(state.get_stream())
with alloc_stream:
d_in = generate_data_with_entropy(num_items, dtype, "1.000", alloc_stream)
d_out = cp.empty(1, dtype=dtype)
h_init = np.zeros(1, dtype=dtype)
reducer = make_reduce_into(d_in=d_in, d_out=d_out, op=max_op, h_init=h_init)
temp_storage_bytes = reducer(
temp_storage=None,
d_in=d_in,
d_out=d_out,
num_items=num_items,
op=max_op,
h_init=h_init,
)
with alloc_stream:
temp_storage = cp.empty(temp_storage_bytes, dtype=np.uint8)
state.add_element_count(num_items)
state.add_global_memory_reads(num_items * d_in.dtype.itemsize, "Size")
state.add_global_memory_writes(d_out.dtype.itemsize)
def launcher(launch: bench.Launch):
reducer(
temp_storage=temp_storage,
d_in=d_in,
d_out=d_out,
num_items=num_items,
op=max_op,
h_init=h_init,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
if __name__ == "__main__":
b = bench.register(bench_reduce_custom)
b.set_name("base")
b.add_string_axis("T{ct}", list(TYPE_MAP.keys()))
b.add_int64_power_of_two_axis("Elements{io}", range(16, 29, 4))
bench.run_all_benchmarks(sys.argv)

85
cccl_upstream/python/cuda_cccl/benchmarks/compute/reduce/min.py Normal file
View File

@@ -0,0 +1,85 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for reduce min operation using cuda.compute.reduce_into.
C++ equivalent: cub/benchmarks/bench/reduce/min.cu
Notes:
- Uses OpKind.MINIMUM for minimum reduction
- C++ uses cuda::minimum<> which CUB recognizes for optimized code paths (DPX on Hopper+)
- int128 and complex32 are not supported by cupy
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import cupy as cp
import numpy as np
from utils import FUNDAMENTAL_TYPES as TYPE_MAP
from utils import as_cupy_stream, generate_data_with_entropy
import cuda.bench as bench
from cuda.compute import OpKind, make_reduce_into
def bench_reduce_min(state: bench.State):
type_str = state.get_string("T{ct}")
dtype = TYPE_MAP[type_str]
num_items = int(state.get_int64("Elements{io}"))
alloc_stream = as_cupy_stream(state.get_stream())
with alloc_stream:
d_in = generate_data_with_entropy(num_items, dtype, "1.000", alloc_stream)
d_out = cp.empty(1, dtype=dtype)
# Initial value for min reduction (max value of type)
if np.issubdtype(dtype, np.integer):
init_val = np.iinfo(dtype).max
else:
init_val = np.finfo(dtype).max
h_init = np.array([init_val], dtype=dtype)
reducer = make_reduce_into(d_in=d_in, d_out=d_out, op=OpKind.MINIMUM, h_init=h_init)
temp_storage_bytes = reducer(
temp_storage=None,
d_in=d_in,
d_out=d_out,
num_items=num_items,
op=OpKind.MINIMUM,
h_init=h_init,
)
with alloc_stream:
temp_storage = cp.empty(temp_storage_bytes, dtype=np.uint8)
state.add_element_count(num_items)
state.add_global_memory_reads(num_items * d_in.dtype.itemsize, "Size")
state.add_global_memory_writes(d_out.dtype.itemsize)
def launcher(launch: bench.Launch):
reducer(
temp_storage=temp_storage,
d_in=d_in,
d_out=d_out,
num_items=num_items,
op=OpKind.MINIMUM,
h_init=h_init,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
if __name__ == "__main__":
b = bench.register(bench_reduce_min)
b.set_name("base")
b.add_string_axis("T{ct}", list(TYPE_MAP.keys()))
b.add_int64_power_of_two_axis("Elements{io}", range(16, 29, 4))
bench.run_all_benchmarks(sys.argv)

87
cccl_upstream/python/cuda_cccl/benchmarks/compute/reduce/nondeterministic.py Normal file
View File

@@ -0,0 +1,87 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for nondeterministic reduce sum using cuda.compute.reduce_into.
C++ equivalent: cub/benchmarks/bench/reduce/nondeterministic.cu
Notes:
- Uses Determinism.NOT_GUARANTEED
- C++ tests int32, int64, float, double
- Migration: Python fixes offsets; C++ exposes an OffsetT axis.
- OffsetT axis is omitted because the Python API does not expose offset type.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import cupy as cp
import numpy as np
from utils import ALL_TYPES as _ALL_TYPES
from utils import as_cupy_stream, generate_data_with_entropy
import cuda.bench as bench
from cuda.compute import Determinism, OpKind, make_reduce_into
TYPE_MAP = {k: _ALL_TYPES[k] for k in ("I32", "I64", "F32", "F64")}
def bench_reduce_nondeterministic(state: bench.State):
type_str = state.get_string("T{ct}")
dtype = TYPE_MAP[type_str]
num_items = int(state.get_int64("Elements{io}"))
alloc_stream = as_cupy_stream(state.get_stream())
with alloc_stream:
d_in = generate_data_with_entropy(num_items, dtype, "1.000", alloc_stream)
d_out = cp.empty(1, dtype=dtype)
h_init = np.zeros(1, dtype=dtype)
reducer = make_reduce_into(
d_in=d_in,
d_out=d_out,
op=OpKind.PLUS,
h_init=h_init,
determinism=Determinism.NOT_GUARANTEED,
)
temp_storage_bytes = reducer(
temp_storage=None,
d_in=d_in,
d_out=d_out,
num_items=num_items,
op=OpKind.PLUS,
h_init=h_init,
)
with alloc_stream:
temp_storage = cp.empty(temp_storage_bytes, dtype=np.uint8)
state.add_element_count(num_items)
state.add_global_memory_reads(num_items * d_in.dtype.itemsize, "Size")
state.add_global_memory_writes(1 * d_out.dtype.itemsize)
def launcher(launch: bench.Launch):
reducer(
temp_storage=temp_storage,
d_in=d_in,
d_out=d_out,
num_items=num_items,
op=OpKind.PLUS,
h_init=h_init,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
if __name__ == "__main__":
b = bench.register(bench_reduce_nondeterministic)
b.set_name("base")
b.add_string_axis("T{ct}", list(TYPE_MAP.keys()))
b.add_int64_power_of_two_axis("Elements{io}", range(16, 29, 4))
bench.run_all_benchmarks(sys.argv)

80
cccl_upstream/python/cuda_cccl/benchmarks/compute/reduce/sum.py Normal file
View File

@@ -0,0 +1,80 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for reduce sum operation using cuda.compute.reduce_into.
C++ equivalent: cub/benchmarks/bench/reduce/sum.cu
Notes:
- int128 and complex32 are not supported by cupy
- Migration: Python excludes int128/complex; C++ supports more types/tuning.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import cupy as cp
import numpy as np
from utils import SIGNED_TYPES as TYPE_MAP
from utils import as_cupy_stream, generate_data_with_entropy
import cuda.bench as bench
from cuda.compute import OpKind, make_reduce_into
def bench_reduce_sum(state: bench.State):
type_str = state.get_string("T{ct}")
dtype = TYPE_MAP[type_str]
num_items = int(state.get_int64("Elements{io}"))
alloc_stream = as_cupy_stream(state.get_stream())
with alloc_stream:
d_in = generate_data_with_entropy(num_items, dtype, "1.000", alloc_stream)
d_out = cp.empty(1, dtype=dtype)
# Initial value for reduction
h_init = np.zeros(1, dtype=dtype)
reducer = make_reduce_into(d_in=d_in, d_out=d_out, op=OpKind.PLUS, h_init=h_init)
temp_storage_bytes = reducer(
temp_storage=None,
d_in=d_in,
d_out=d_out,
num_items=num_items,
op=OpKind.PLUS,
h_init=h_init,
)
with alloc_stream:
temp_storage = cp.empty(temp_storage_bytes, dtype=np.uint8)
state.add_element_count(num_items)
state.add_global_memory_reads(num_items * d_in.dtype.itemsize, "Size")
state.add_global_memory_writes(d_out.dtype.itemsize)
def launcher(launch: bench.Launch):
reducer(
temp_storage=temp_storage,
d_in=d_in,
d_out=d_out,
num_items=num_items,
op=OpKind.PLUS,
h_init=h_init,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
if __name__ == "__main__":
b = bench.register(bench_reduce_sum)
b.set_name("base")
b.add_string_axis("T{ct}", list(TYPE_MAP.keys()))
b.add_int64_power_of_two_axis("Elements{io}", range(16, 29, 4))
bench.run_all_benchmarks(sys.argv)

528
cccl_upstream/python/cuda_cccl/benchmarks/compute/run_benchmarks.py Executable file
View File

@@ -0,0 +1,528 @@
#!/usr/bin/env python3
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Run Python cuda.compute and C++ CUB benchmarks.
Prerequisites: C++ benchmarks must be built first via:
cd /path/to/cccl && ./ci/build_cub.sh -arch 89
Usage:
python run_benchmarks.py [options]
Options:
-d, --device ID GPU device ID [default: 0]
-b, --benchmark NAME Run specific benchmark only [default: all]
--py Only run Python benchmarks
--cpp Only run C++ benchmarks
--quick, -q Run with reduced parameter set for fast testing
--profile Run each config once (nvbench profile mode), no sampling
-h, --help Show this help message
Benchmark names follow CUB structure:
e.g. transform/fill, transform/babelstream
Examples:
python run_benchmarks.py # Run all benchmarks
python run_benchmarks.py -b transform/fill # Only fill benchmark
python run_benchmarks.py -b reduce/sum -d 0 # Reduce sum on device 0
python run_benchmarks.py -b scan/exclusive/sum --py # Only Python
python run_benchmarks.py --quick # Quick mode (reduced params)
python run_benchmarks.py -b merge_sort/keys -q # Quick mode, single benchmark
"""
import argparse
import os
import shlex
import subprocess
import sys
from pathlib import Path
# ============================================================================
# Configuration
# ============================================================================
SCRIPT_DIR = Path(__file__).parent
CCCL_ROOT = SCRIPT_DIR.parents[3]
RESULTS_DIR = SCRIPT_DIR / "results"
CUB_BENCH_DIR = CCCL_ROOT / "build" / "cub" / "bin"
QUICK_CONFIG_FILE = SCRIPT_DIR / "quick_configs.yaml"
# Supported benchmarks (Python implementations available)
SUPPORTED_BENCHMARKS = [
"transform/fill",
"transform/babelstream",
"transform/heavy",
"transform/fib",
"transform/grayscale",
"transform/complex_cmp",
"transform_reduce/sum",
"reduce/sum",
"reduce/min",
"reduce/custom",
"reduce/nondeterministic",
"scan/exclusive/sum",
"scan/exclusive/custom",
"histogram/even",
"select/if",
"select/unique_by_key",
"radix_sort/keys",
"radix_sort/pairs",
"merge_sort/keys",
"merge_sort/pairs",
"segmented_sort/keys",
"segmented_reduce/variable_sum",
"partition/three_way",
]
# Axes that use power-of-two values (need [pow2] suffix for nvbench)
# These are the base names (without {ct}/{io} suffixes)
POW2_AXES_CPP = {
"Elements",
"GuaranteedMaxSegSize",
"MaxSegSize",
"MaxSegmentSize",
"SegmentSize",
"Segments",
}
POW2_AXES_PY = {
"Elements",
"GuaranteedMaxSegSize",
"MaxSegSize",
"MaxSegmentSize",
"Segments",
}
# Axis name mappings from C++ to Python.
# Keep type axes in their C++ form (`{ct}`) to match benchmark axis names.
CPP_TO_PY_AXIS_MAP = {
"T{ct}": "T{ct}",
"KeyT{ct}": "KeyT{ct}",
"ValueT{ct}": "ValueT{ct}",
"SampleT{ct}": "SampleT{ct}",
"Heaviness{ct}": "Heaviness{ct}",
}
def strip_axis_suffix(axis_name: str) -> str:
"""Strip {ct} or {io} suffix from axis name for Python benchmarks.
e.g., "T{ct}" -> "T{ct}", "Elements{io}" -> "Elements{io}".
"""
if axis_name in CPP_TO_PY_AXIS_MAP:
return CPP_TO_PY_AXIS_MAP[axis_name]
# Generic suffix stripping for any other axes
if axis_name.endswith("{ct}"):
return axis_name.rsplit("{", 1)[0]
if axis_name.endswith("{io}"):
return axis_name
return axis_name
def get_base_axis_name(axis_name: str) -> str:
"""Get base axis name (without suffix) for POW2 check.
e.g., "Elements{io}" -> "Elements"
"""
if axis_name.endswith("{ct}") or axis_name.endswith("{io}"):
return axis_name.rsplit("{", 1)[0]
return axis_name
# ============================================================================
# Helper Functions
# ============================================================================
def print_banner(msg: str) -> None:
"""Print a banner message."""
print()
print("=" * 72)
print(msg)
print("=" * 72)
print()
def print_section(msg: str) -> None:
"""Print a section header."""
print()
print("-" * 72)
print(msg)
print("-" * 72)
def load_quick_configs() -> dict:
"""Load quick mode configurations from YAML file."""
import yaml # only needed for --quick
with open(QUICK_CONFIG_FILE) as f:
return yaml.safe_load(f)
def build_axis_args_from_config(axis_config: dict, for_python: bool) -> list:
"""Build --axis arguments for nvbench CLI from an axis config dict.
Args:
axis_config: Dict of axis_name -> value
for_python: If True, strip C++ suffixes for Python benchmarks
"""
args = []
for axis_name, value in axis_config.items():
# For Python, strip C++ suffixes from axis names
if for_python:
axis_name = strip_axis_suffix(axis_name)
# Check if this is a power-of-two axis (using base name)
base_name = get_base_axis_name(axis_name)
if for_python and base_name == "SegmentSize":
actual_value = 2 ** int(value)
args.extend(["--axis", f"{axis_name}={actual_value}"])
continue
pow2_axes = POW2_AXES_PY if for_python else POW2_AXES_CPP
if base_name in pow2_axes:
args.extend(["--axis", f"{axis_name}[pow2]={value}"])
else:
args.extend(["--axis", f"{axis_name}={value}"])
return args
def get_quick_config_entry(benchmark: str, quick_configs: dict) -> dict:
"""Get quick config entry for a benchmark, raising if missing."""
if benchmark not in quick_configs:
raise ValueError(
f"Benchmark '{benchmark}' not found in quick_configs.yaml.\n"
f"Cannot run in --quick mode. Add configuration for this benchmark."
)
return quick_configs[benchmark]
def get_cpp_binary(benchmark: str) -> str:
"""Get C++ binary name from benchmark path.
e.g., "transform/fill" -> "cub.bench.transform.fill.base"
"""
return f"cub.bench.{benchmark.replace('/', '.')}.base"
def get_py_script(benchmark: str) -> Path:
"""Get Python script path from benchmark path.
e.g., "transform/fill" -> "transform/fill.py"
"""
return SCRIPT_DIR / f"{benchmark}.py"
def get_result_path(benchmark: str, suffix: str) -> Path:
"""Get results file path.
e.g., "transform/fill", "cpp" -> "results/transform/fill_cpp.json"
"""
bench_path = Path(benchmark)
return RESULTS_DIR / bench_path.parent / f"{bench_path.name}_{suffix}.json"
def get_log_path(benchmark: str, suffix: str) -> Path:
"""Get log file path under results/logs.
e.g., "transform/fill", "cpp" -> "results/logs/transform/fill_cpp.log"
"""
bench_path = Path(benchmark)
return RESULTS_DIR / "logs" / bench_path.parent / f"{bench_path.name}_{suffix}.log"
def run_and_log(cmd: list, log_path: Path, env: dict | None = None) -> dict:
"""Run command and write stdout/stderr to log file."""
log_path.parent.mkdir(parents=True, exist_ok=True)
with open(log_path, "w", encoding="utf-8") as log_file:
log_file.write(f"Command: {shlex.join(cmd)}\n\n")
log_file.flush()
try:
result = subprocess.run(
cmd, check=False, stdout=log_file, stderr=log_file, env=env
)
except Exception as exc: # noqa: BLE001
log_file.write("\nERROR: Runner failed to execute command.\n")
log_file.write(f"{exc}\n")
return {"status": "error", "returncode": None, "error": str(exc)}
status = "ok" if result.returncode == 0 else "failed"
return {"status": status, "returncode": result.returncode}
# ============================================================================
# Benchmark Runner
# ============================================================================
def run_benchmark(
benchmark: str,
device: str,
run_py: bool,
run_cpp: bool,
quick_mode: bool,
quick_configs: dict,
profile: bool,
) -> dict:
"""Run a single benchmark.
Returns dict with paths to generated result files.
"""
cpp_binary = get_cpp_binary(benchmark)
py_script = get_py_script(benchmark)
cpp_result = get_result_path(benchmark, "cpp")
py_result = get_result_path(benchmark, "py")
cpp_log = get_log_path(benchmark, "cpp")
py_log = get_log_path(benchmark, "py")
# Ensure results subdirectory exists
cpp_result.parent.mkdir(parents=True, exist_ok=True)
# Build axis arguments for quick mode
cpp_axis_args = []
py_axis_args = []
if quick_mode:
config_entry = get_quick_config_entry(benchmark, quick_configs)
if "benchmarks" in config_entry:
for bench_name, axis_config in config_entry["benchmarks"].items():
cpp_axis_args.extend(["--benchmark", bench_name])
cpp_axis_args.extend(
build_axis_args_from_config(axis_config, for_python=False)
)
py_axis_args.extend(["--benchmark", bench_name])
py_axis_args.extend(
build_axis_args_from_config(axis_config, for_python=True)
)
else:
cpp_axis_args = build_axis_args_from_config(config_entry, for_python=False)
py_axis_args = build_axis_args_from_config(config_entry, for_python=True)
results = {}
# Run C++ benchmark
if run_cpp:
print(f"Running C++ benchmark: {cpp_binary}")
cpp_bin = CUB_BENCH_DIR / cpp_binary
if not cpp_bin.exists():
print(f"ERROR: C++ binary not found: {cpp_bin}")
print()
print("Please build C++ benchmarks first:")
print(f" cd {CCCL_ROOT}")
print(" ./ci/build_cub.sh -arch <your_gpu_arch> # e.g., 89 for RTX 4090")
print()
print("Available benchmarks in build directory:")
if CUB_BENCH_DIR.exists():
binaries = list(CUB_BENCH_DIR.glob("*.base"))[:10]
for b in binaries:
print(f" {b.name}")
if len(list(CUB_BENCH_DIR.glob("*.base"))) > 10:
print(" ...")
else:
print(f" (directory not found: {CUB_BENCH_DIR})")
sys.exit(1)
cmd = [str(cpp_bin), "--json", str(cpp_result), "--devices", device]
cmd.extend(cpp_axis_args)
if profile:
cmd.append("--profile")
# Ensure the CUB build lib dir (containing libnvbench.so) is on the
# dynamic linker search path for the child process.
cpp_lib_dir = str(CUB_BENCH_DIR.parent / "lib")
cpp_env = os.environ.copy()
existing_ld = cpp_env.get("LD_LIBRARY_PATH", "")
cpp_env["LD_LIBRARY_PATH"] = (
f"{cpp_lib_dir}:{existing_ld}" if existing_ld else cpp_lib_dir
)
cpp_status = run_and_log(cmd, cpp_log, env=cpp_env)
print(f" Results: {cpp_result}")
print(f" Log: {cpp_log}")
if cpp_status["status"] != "ok":
print(f" WARNING: C++ benchmark failed (exit {cpp_status['returncode']}).")
results["cpp"] = cpp_result
results["cpp_status"] = cpp_status
# Run Python benchmark
if run_py:
print(f"Running Python benchmark: {py_script.relative_to(SCRIPT_DIR)}")
if not py_script.exists():
print(f"ERROR: Python script not found: {py_script}")
sys.exit(1)
cmd = [
sys.executable,
str(py_script),
"--json",
str(py_result),
"--devices",
device,
]
cmd.extend(py_axis_args)
if profile:
cmd.append("--profile")
py_status = run_and_log(cmd, py_log)
print(f" Results: {py_result}")
print(f" Log: {py_log}")
if py_status["status"] != "ok":
print(
f" WARNING: Python benchmark failed (exit {py_status['returncode']})."
)
results["py"] = py_result
results["py_status"] = py_status
return results
# ============================================================================
# Main
# ============================================================================
def main():
parser = argparse.ArgumentParser(
description="Run Python cuda.compute and C++ CUB benchmarks",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog=f"""
Supported benchmarks:
{chr(10).join(f" {b}" for b in SUPPORTED_BENCHMARKS)}
""",
)
parser.add_argument(
"-d", "--device", default="0", help="GPU device ID [default: 0]"
)
parser.add_argument(
"-b", "--benchmark", help="Run specific benchmark only [default: all]"
)
parser.add_argument("--py", action="store_true", help="Only run Python benchmarks")
parser.add_argument("--cpp", action="store_true", help="Only run C++ benchmarks")
parser.add_argument(
"-q",
"--quick",
action="store_true",
help="Run with reduced parameter set for fast testing",
)
parser.add_argument(
"--profile",
action="store_true",
help="Run each benchmark configuration once (nvbench profile mode) with "
"no sampling -- for smoke-testing that benchmarks execute without error, "
"not for measuring performance.",
)
args = parser.parse_args()
# Determine what to run
run_py = True
run_cpp = True
if args.py and not args.cpp:
run_cpp = False
elif args.cpp and not args.py:
run_py = False
# Load quick configs if needed
quick_configs = {}
if args.quick:
quick_configs = load_quick_configs()
# Validate and determine benchmarks to run
if args.benchmark:
if args.benchmark not in SUPPORTED_BENCHMARKS:
print(f"ERROR: Benchmark '{args.benchmark}' not supported.")
print()
print("Available benchmarks:")
for b in SUPPORTED_BENCHMARKS:
print(f" {b}")
sys.exit(1)
benchmarks_to_run = [args.benchmark]
else:
benchmarks_to_run = SUPPORTED_BENCHMARKS
# Print configuration
print_banner("CCCL Benchmark Runner")
print("Configuration:")
print(f" CCCL Root: {CCCL_ROOT}")
print(f" C++ Binaries: {CUB_BENCH_DIR}")
print(f" Results Dir: {RESULTS_DIR}")
print(f" Device: {args.device}")
print(f" Benchmarks: {' '.join(benchmarks_to_run)}")
print(f" Run C++: {run_cpp}")
print(f" Run Python: {run_py}")
print(f" Quick Mode: {args.quick}")
print(f" Profile Mode: {args.profile}")
print()
# Check C++ binaries directory exists (if running C++)
if run_cpp and not CUB_BENCH_DIR.exists():
print(f"ERROR: C++ benchmark directory not found: {CUB_BENCH_DIR}")
print()
print("Please build C++ benchmarks first:")
print(f" cd {CCCL_ROOT}")
print(" ./ci/build_cub.sh -arch <your_gpu_arch> # e.g., 89 for RTX 4090")
sys.exit(1)
# Run benchmarks
all_results = {}
for bench in benchmarks_to_run:
print_section(f"Benchmark: {bench}")
results = run_benchmark(
bench, args.device, run_py, run_cpp, args.quick, quick_configs, args.profile
)
all_results[bench] = results
print()
# Print summary
print_banner("Summary")
print(f"Results directory: {RESULTS_DIR}")
print()
print("Generated files:")
for bench in benchmarks_to_run:
cpp_result = get_result_path(bench, "cpp")
py_result = get_result_path(bench, "py")
cpp_log = get_log_path(bench, "cpp")
py_log = get_log_path(bench, "py")
print(f" {bench}:")
if cpp_result.exists():
print(f" C++ results: {cpp_result}")
if cpp_log.exists():
print(f" C++ log: {cpp_log}")
if py_result.exists():
print(f" Python results: {py_result}")
if py_log.exists():
print(f" Python log: {py_log}")
status = all_results.get(bench, {})
cpp_status = status.get("cpp_status")
py_status = status.get("py_status")
if cpp_status and cpp_status.get("status") != "ok":
print(
" C++ status: "
f"{cpp_status.get('status')} (exit {cpp_status.get('returncode')})"
)
if py_status and py_status.get("status") != "ok":
print(
" Python status: "
f"{py_status.get('status')} (exit {py_status.get('returncode')})"
)
print()
# Exit non-zero if any benchmark failed so CI (e.g. the --profile smoke gate)
# catches benchmark rot instead of silently passing on printed warnings.
failed = []
for bench in benchmarks_to_run:
results = all_results.get(bench, {})
for key in ("cpp_status", "py_status"):
if results.get(key, {}).get("status") not in (None, "ok"):
failed.append(bench)
break
if failed:
print(f"ERROR: {len(failed)} benchmark(s) failed: {failed}")
sys.exit(1)
if __name__ == "__main__":
main()

83
cccl_upstream/python/cuda_cccl/benchmarks/compute/scan/exclusive/custom.py Normal file
View File

@@ -0,0 +1,83 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for exclusive scan custom operation using cuda.compute.exclusive_scan.
C++ equivalent: cub/benchmarks/bench/scan/exclusive/custom.cu
Notes:
- Uses a custom max operator (not OpKind)
- int128 and complex32 are not supported by cupy
- Migration: Python fixes offsets; C++ exposes an OffsetT axis.
- OffsetT axis is omitted because the Python API does not expose offset type.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
import cupy as cp
import numpy as np
from utils import SIGNED_TYPES as TYPE_MAP
from utils import as_cupy_stream, generate_data_with_entropy
import cuda.bench as bench
from cuda.compute import make_exclusive_scan
def max_op(a, b):
return a if a > b else b
def bench_scan_exclusive_custom(state: bench.State):
type_str = state.get_string("T{ct}")
dtype = TYPE_MAP[type_str]
num_items = int(state.get_int64("Elements{io}"))
alloc_stream = as_cupy_stream(state.get_stream())
with alloc_stream:
d_in = generate_data_with_entropy(num_items, dtype, "1.000", alloc_stream)
d_out = cp.empty(num_items, dtype=dtype)
h_init = np.zeros(1, dtype=dtype)
scanner = make_exclusive_scan(d_in=d_in, d_out=d_out, op=max_op, init_value=h_init)
temp_storage_bytes = scanner(
temp_storage=None,
d_in=d_in,
d_out=d_out,
op=max_op,
init_value=h_init,
num_items=num_items,
)
with alloc_stream:
temp_storage = cp.empty(temp_storage_bytes, dtype=np.uint8)
state.add_element_count(num_items)
state.add_global_memory_reads(num_items * d_in.dtype.itemsize, "Size")
state.add_global_memory_writes(num_items * d_out.dtype.itemsize)
def launcher(launch: bench.Launch):
scanner(
temp_storage=temp_storage,
d_in=d_in,
d_out=d_out,
op=max_op,
init_value=h_init,
num_items=num_items,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
if __name__ == "__main__":
b = bench.register(bench_scan_exclusive_custom)
b.set_name("base")
b.add_string_axis("T{ct}", list(TYPE_MAP.keys()))
b.add_int64_power_of_two_axis("Elements{io}", range(16, 33, 4))
bench.run_all_benchmarks(sys.argv)

84
cccl_upstream/python/cuda_cccl/benchmarks/compute/scan/exclusive/sum.py Normal file
View File

@@ -0,0 +1,84 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for exclusive scan sum operation using cuda.compute.exclusive_scan.
C++ equivalent: cub/benchmarks/bench/scan/exclusive/sum.cu
Notes:
- int128 and complex32 are not supported by cupy
- Migration: Python fixes offsets; C++ exposes an OffsetT axis.
- OffsetT axis is omitted because the Python API does not expose offset type.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
import cupy as cp
import numpy as np
from utils import SIGNED_TYPES as TYPE_MAP
from utils import as_cupy_stream, generate_data_with_entropy
import cuda.bench as bench
from cuda.compute import OpKind, make_exclusive_scan
def bench_scan_exclusive_sum(state: bench.State):
type_str = state.get_string("T{ct}")
dtype = TYPE_MAP[type_str]
num_items = int(state.get_int64("Elements{io}"))
alloc_stream = as_cupy_stream(state.get_stream())
with alloc_stream:
d_in = generate_data_with_entropy(num_items, dtype, "1.000", alloc_stream)
# Output is same size as input for scan
d_out = cp.empty(num_items, dtype=dtype)
# Initial value for scan (identity for addition)
h_init = np.zeros(1, dtype=dtype)
scanner = make_exclusive_scan(
d_in=d_in, d_out=d_out, op=OpKind.PLUS, init_value=h_init
)
temp_storage_bytes = scanner(
temp_storage=None,
d_in=d_in,
d_out=d_out,
op=OpKind.PLUS,
init_value=h_init,
num_items=num_items,
)
with alloc_stream:
temp_storage = cp.empty(temp_storage_bytes, dtype=np.uint8)
state.add_element_count(num_items)
state.add_global_memory_reads(num_items * d_in.dtype.itemsize, "Size")
state.add_global_memory_writes(num_items * d_out.dtype.itemsize)
def launcher(launch: bench.Launch):
scanner(
temp_storage=temp_storage,
d_in=d_in,
d_out=d_out,
op=OpKind.PLUS,
init_value=h_init,
num_items=num_items,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
if __name__ == "__main__":
b = bench.register(bench_scan_exclusive_sum)
b.set_name("base")
b.add_string_axis("T{ct}", list(TYPE_MAP.keys()))
b.add_int64_power_of_two_axis("Elements{io}", range(16, 33, 4))
bench.run_all_benchmarks(sys.argv)

165
cccl_upstream/python/cuda_cccl/benchmarks/compute/segmented_reduce/variable_sum.py Normal file
View File

@@ -0,0 +1,165 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for segmented_reduce (sum) with variable-size segments.
C++ equivalent: cub/benchmarks/bench/segmented_reduce/variable_sum.cu (uses variable_base.cuh)
Notes:
- Implements four sub-benchmarks: variable_default, variable_small_dynamic,
variable_medium_dynamic, variable_large_dynamic.
- OffsetT axis is omitted because the Python API does not expose offset type.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import cupy as cp
import numpy as np
from utils import (
ALL_TYPES,
as_cupy_stream,
generate_data_with_entropy,
generate_uniform_segment_offsets,
)
import cuda.bench as bench
from cuda.compute import OpKind, make_segmented_reduce
TYPE_MAP = {k: ALL_TYPES[k] for k in ("I32", "I64", "F32", "F64")}
def run_segmented_reduce(
state: bench.State,
d_in,
d_out,
h_init,
start_offsets,
end_offsets,
num_segments,
guaranteed_max_seg_size,
):
reducer = make_segmented_reduce(
d_in=d_in,
d_out=d_out,
start_offsets_in=start_offsets,
end_offsets_in=end_offsets,
op=OpKind.PLUS,
h_init=h_init,
)
temp_storage_bytes = reducer(
temp_storage=None,
d_in=d_in,
d_out=d_out,
num_segments=num_segments,
start_offsets_in=start_offsets,
end_offsets_in=end_offsets,
op=OpKind.PLUS,
h_init=h_init,
max_segment_size=guaranteed_max_seg_size,
)
alloc_stream = as_cupy_stream(state.get_stream())
with alloc_stream:
temp_storage = cp.empty(temp_storage_bytes, dtype=np.uint8)
def launcher(launch: bench.Launch):
reducer(
temp_storage=temp_storage,
d_in=d_in,
d_out=d_out,
num_segments=num_segments,
start_offsets_in=start_offsets,
end_offsets_in=end_offsets,
op=OpKind.PLUS,
h_init=h_init,
max_segment_size=guaranteed_max_seg_size,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False, sync=True)
def bench_variable_segmented_reduce(state: bench.State):
type_str = state.get_string("T{ct}")
dtype = TYPE_MAP[type_str]
num_elements = int(state.get_int64("Elements{io}"))
max_segment_size = int(state.get_int64("MaxSegmentSize"))
guaranteed_max_seg_size = int(state.get_int64("GuaranteedMaxSegSize"))
# Skip cases where hint would be incorrect (max > guaranteed)
if guaranteed_max_seg_size != 0 and max_segment_size > guaranteed_max_seg_size:
state.skip("max_segment_size > guaranteed_max_seg_size")
return
min_segment_size = 1
offsets = generate_uniform_segment_offsets(
num_elements, min_segment_size, max_segment_size
)
alloc_stream = as_cupy_stream(state.get_stream())
h_init = np.zeros(1, dtype=dtype)
d_in = generate_data_with_entropy(num_elements, dtype, "1.000", alloc_stream)
with alloc_stream:
start_offsets = cp.asarray(offsets[:-1], dtype=np.int64)
end_offsets = cp.asarray(offsets[1:], dtype=np.int64)
d_out = cp.empty(int(start_offsets.size), dtype=dtype)
alloc_stream.synchronize()
num_segments = int(start_offsets.size)
state.add_element_count(num_elements)
state.add_global_memory_reads(num_elements * d_in.dtype.itemsize)
state.add_global_memory_writes(num_segments * d_out.dtype.itemsize)
state.add_global_memory_reads((num_segments + 1) * start_offsets.dtype.itemsize)
run_segmented_reduce(
state,
d_in,
d_out,
h_init,
start_offsets,
end_offsets,
num_segments,
guaranteed_max_seg_size,
)
if __name__ == "__main__":
# Default: no size hint — uses generic large-reduce kernel regardless of segment size
b_default = bench.register(bench_variable_segmented_reduce)
b_default.set_name("variable_default")
b_default.add_string_axis("T{ct}", list(TYPE_MAP.keys()))
b_default.add_int64_power_of_two_axis("Elements{io}", range(16, 28, 4))
b_default.add_int64_power_of_two_axis("MaxSegmentSize", range(1, 17, 1))
b_default.add_int64_axis("GuaranteedMaxSegSize", [0])
# Small segments (1–16 items): hint enables warp-level reduction
b_small = bench.register(bench_variable_segmented_reduce)
b_small.set_name("variable_small_dynamic")
b_small.add_string_axis("T{ct}", list(TYPE_MAP.keys()))
b_small.add_int64_power_of_two_axis("Elements{io}", range(16, 28, 4))
b_small.add_int64_power_of_two_axis("MaxSegmentSize", range(1, 5, 1))
b_small.add_int64_power_of_two_axis("GuaranteedMaxSegSize", range(1, 5, 1))
# Medium segments (32–256 items): hint enables warp-level reduction
b_medium = bench.register(bench_variable_segmented_reduce)
b_medium.set_name("variable_medium_dynamic")
b_medium.add_string_axis("T{ct}", list(TYPE_MAP.keys()))
b_medium.add_int64_power_of_two_axis("Elements{io}", range(16, 28, 4))
b_medium.add_int64_power_of_two_axis("MaxSegmentSize", range(5, 9, 1))
b_medium.add_int64_power_of_two_axis("GuaranteedMaxSegSize", range(5, 9, 1))
# Large segments (512+ items): hint enables block-level reduction
b_large = bench.register(bench_variable_segmented_reduce)
b_large.set_name("variable_large_dynamic")
b_large.add_string_axis("T{ct}", list(TYPE_MAP.keys()))
b_large.add_int64_power_of_two_axis("Elements{io}", range(16, 28, 4))
b_large.add_int64_power_of_two_axis("MaxSegmentSize", range(9, 17, 1))
b_large.add_int64_power_of_two_axis("GuaranteedMaxSegSize", range(9, 17, 1))
bench.run_all_benchmarks(sys.argv)

167
cccl_upstream/python/cuda_cccl/benchmarks/compute/segmented_sort/keys.py Normal file
View File

@@ -0,0 +1,167 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for segmented_sort keys using cuda.compute.
C++ equivalent: cub/benchmarks/bench/segmented_sort/keys.cu
Notes:
- Implements three sub-benchmarks: power, small, large
- Power uses power-law segment sizes with Entropy axis
- Small/large use uniform segment sizes with MaxSegmentSize axis
- Migration: uniform offsets use min_segment_size ~ max/2.
- OffsetT axis is omitted because the Python API does not expose offset type.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import cupy as cp
import numpy as np
from utils import (
FUNDAMENTAL_TYPES as TYPE_MAP,
)
from utils import (
as_cupy_stream,
generate_data_with_entropy,
generate_power_law_offsets,
generate_uniform_segment_offsets,
)
import cuda.bench as bench
from cuda.compute import SortOrder, make_segmented_sort
def run_segmented_sort(
state: bench.State,
d_in_keys,
d_out_keys,
start_offsets,
end_offsets,
num_items,
num_segments,
):
sorter = make_segmented_sort(
d_in_keys=d_in_keys,
d_out_keys=d_out_keys,
d_in_values=None,
d_out_values=None,
start_offsets_in=start_offsets,
end_offsets_in=end_offsets,
order=SortOrder.ASCENDING,
)
temp_storage_bytes = sorter(
temp_storage=None,
d_in_keys=d_in_keys,
d_out_keys=d_out_keys,
d_in_values=None,
d_out_values=None,
num_items=num_items,
num_segments=num_segments,
start_offsets_in=start_offsets,
end_offsets_in=end_offsets,
)
alloc_stream = as_cupy_stream(state.get_stream())
with alloc_stream:
temp_storage = cp.empty(temp_storage_bytes, dtype=np.uint8)
def launcher(launch: bench.Launch):
sorter(
temp_storage=temp_storage,
d_in_keys=d_in_keys,
d_out_keys=d_out_keys,
d_in_values=None,
d_out_values=None,
num_items=num_items,
num_segments=num_segments,
start_offsets_in=start_offsets,
end_offsets_in=end_offsets,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False, sync=True)
def bench_segmented_sort(state: bench.State, use_power_law: bool):
type_str = state.get_string("T{ct}")
dtype = TYPE_MAP[type_str]
num_elements = int(state.get_int64("Elements{io}"))
alloc_stream = as_cupy_stream(state.get_stream())
if use_power_law:
num_segments = int(state.get_int64("Segments{io}"))
entropy_str = state.get_string("Entropy")
else:
max_segment_size = int(state.get_int64("MaxSegmentSize"))
min_segment_size = max(1, max_segment_size // 2)
entropy_str = "1.000"
if use_power_law:
offsets = generate_power_law_offsets(num_elements, num_segments)
else:
offsets = generate_uniform_segment_offsets(
num_elements, min_segment_size, max_segment_size
)
d_in_keys = generate_data_with_entropy(
num_elements, dtype, entropy_str, alloc_stream
)
with alloc_stream:
d_out_keys = cp.empty(num_elements, dtype=dtype)
start_offsets = cp.asarray(offsets[:-1], dtype=np.int64)
end_offsets = cp.asarray(offsets[1:], dtype=np.int64)
alloc_stream.synchronize()
num_segments = int(start_offsets.size)
state.add_element_count(num_elements)
state.add_global_memory_reads(num_elements * d_in_keys.dtype.itemsize)
state.add_global_memory_writes(num_elements * d_out_keys.dtype.itemsize)
state.add_global_memory_reads((num_segments + 1) * start_offsets.dtype.itemsize)
run_segmented_sort(
state,
d_in_keys,
d_out_keys,
start_offsets,
end_offsets,
num_elements,
num_segments,
)
def bench_segmented_sort_power(state: bench.State):
bench_segmented_sort(state, use_power_law=True)
def bench_segmented_sort_uniform(state: bench.State):
bench_segmented_sort(state, use_power_law=False)
if __name__ == "__main__":
b_power = bench.register(bench_segmented_sort_power)
b_power.set_name("power")
b_power.add_string_axis("T{ct}", list(TYPE_MAP.keys()))
b_power.add_int64_power_of_two_axis("Elements{io}", range(22, 31, 4))
b_power.add_int64_power_of_two_axis("Segments{io}", range(12, 21, 4))
b_power.add_string_axis("Entropy", ["1.000", "0.201"])
b_small = bench.register(bench_segmented_sort_uniform)
b_small.set_name("small")
b_small.add_string_axis("T{ct}", list(TYPE_MAP.keys()))
b_small.add_int64_power_of_two_axis("Elements{io}", range(22, 31, 4))
b_small.add_int64_power_of_two_axis("MaxSegmentSize", range(1, 9, 1))
b_large = bench.register(bench_segmented_sort_uniform)
b_large.set_name("large")
b_large.add_string_axis("T{ct}", list(TYPE_MAP.keys()))
b_large.add_int64_power_of_two_axis("Elements{io}", range(22, 31, 4))
b_large.add_int64_power_of_two_axis("MaxSegmentSize", range(10, 19, 2))
bench.run_all_benchmarks(sys.argv)

119
cccl_upstream/python/cuda_cccl/benchmarks/compute/select/if.py Normal file
View File

@@ -0,0 +1,119 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for select_if using cuda.compute.
C++ equivalent: cub/benchmarks/bench/select/if.cu
Notes:
- The C++ benchmark uses a `less_then_t<T>` predicate with threshold based on entropy
- Entropy controls what fraction of elements are selected:
- 1.000 → selects ~100% (threshold = max value)
- 0.544 → selects ~54.4% (threshold at 54.4% of range)
- 0.000 → selects ~0% (threshold = min value)
- InPlace axis controls whether output can alias input (not exposed in Python API)
- Migration: Python cannot expose InPlace axis; output is sized to num_elements but metrics use actual selected count.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import cupy as cp
import numpy as np
from utils import (
ENTROPY_TO_PROB,
as_cupy_stream,
generate_data_with_entropy,
lerp_min_max,
)
from utils import (
FUNDAMENTAL_TYPES as TYPE_MAP,
)
import cuda.bench as bench
from cuda.compute import make_select
# Entropy values from C++ benchmark
# These control the selection threshold and thus how many elements are selected
def bench_select_if(state: bench.State):
type_str = state.get_string("T{ct}")
dtype = TYPE_MAP[type_str]
num_elements = int(state.get_int64("Elements{io}"))
entropy_str = state.get_string("Entropy")
probability = ENTROPY_TO_PROB[entropy_str]
threshold = lerp_min_max(dtype, probability)
alloc_stream = as_cupy_stream(state.get_stream())
# Match C++ benchmark: input data generation is independent of Entropy.
# Entropy only controls the selection threshold.
d_in = generate_data_with_entropy(num_elements, dtype, "1.000", alloc_stream)
with alloc_stream:
selected_elements = int(cp.count_nonzero(d_in < threshold).get())
d_out = cp.empty(selected_elements, dtype=dtype)
d_num_selected = cp.zeros(1, dtype=np.int64)
alloc_stream.synchronize()
# Create predicate: select elements less than threshold
# For numba device functions, we need to use the value directly in closure
thresh_val = threshold
def less_than_threshold(x):
return x < thresh_val
selector = make_select(
d_in=d_in,
d_out=d_out,
d_num_selected_out=d_num_selected,
cond=less_than_threshold,
)
temp_storage_bytes = selector(
temp_storage=None,
d_in=d_in,
d_out=d_out,
d_num_selected_out=d_num_selected,
cond=less_than_threshold,
num_items=num_elements,
)
with alloc_stream:
temp_storage = cp.empty(temp_storage_bytes, dtype=np.uint8)
state.add_element_count(num_elements)
state.add_global_memory_reads(num_elements * d_in.dtype.itemsize)
state.add_global_memory_writes(selected_elements * d_out.dtype.itemsize)
state.add_global_memory_writes(1 * d_num_selected.dtype.itemsize)
def launcher(launch: bench.Launch):
selector(
temp_storage=temp_storage,
d_in=d_in,
d_out=d_out,
d_num_selected_out=d_num_selected,
cond=less_than_threshold,
num_items=num_elements,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
if __name__ == "__main__":
b = bench.register(bench_select_if)
b.set_name("base")
b.add_string_axis("T{ct}", list(TYPE_MAP.keys()))
b.add_int64_power_of_two_axis("Elements{io}", range(16, 29, 4))
b.add_string_axis("Entropy", ["1.000", "0.544", "0.000"])
# Note: InPlace axis is not exposed in Python API, so we skip it
bench.run_all_benchmarks(sys.argv)

137
cccl_upstream/python/cuda_cccl/benchmarks/compute/select/unique_by_key.py Normal file
View File

@@ -0,0 +1,137 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for unique_by_key using cuda.compute.
C++ equivalent: cub/benchmarks/bench/select/unique_by_key.cu
Notes:
- The C++ benchmark uses MaxSegSize axis to control segment sizes
- Uses equal_to comparison operator for key equality
- Generates key segments with sizes between 1 and MaxSegSize
- Both keys and values are processed
- Migration: Python fixes offsets and generates key segments on GPU to mirror C++.
- OffsetT axis is omitted because the Python API does not expose offset type.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import cupy as cp
import numpy as np
from utils import INTEGRAL_TYPES, SIGNED_TYPES, as_cupy_stream, generate_key_segments
import cuda.bench as bench
from cuda.compute import OpKind, make_unique_by_key
KEY_TYPE_MAP = INTEGRAL_TYPES
VALUE_TYPE_MAP = {**SIGNED_TYPES, "C32": np.complex64}
def bench_unique_by_key(state: bench.State):
key_type_str = state.get_string("KeyT{ct}")
value_type_str = state.get_string("ValueT{ct}")
key_dtype = KEY_TYPE_MAP[key_type_str]
value_dtype = VALUE_TYPE_MAP[value_type_str]
num_elements = int(state.get_int64("Elements{io}"))
max_seg_size = int(state.get_int64("MaxSegSize"))
if num_elements > np.iinfo(np.int32).max:
state.skip("Skipping: num_elements exceeds int32 limits")
return
alloc_stream = as_cupy_stream(state.get_stream())
d_in_keys = generate_key_segments(
num_elements,
key_dtype,
min_segment_size=1,
max_segment_size=max_seg_size,
stream=alloc_stream,
)
with alloc_stream:
d_in_values = cp.zeros(num_elements, dtype=value_dtype)
d_out_keys = cp.empty(num_elements, dtype=key_dtype)
d_out_values = cp.empty(num_elements, dtype=value_dtype)
d_num_selected = cp.empty(1, dtype=np.int32)
alloc_stream.synchronize()
uniquer = make_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_num_selected,
op=OpKind.EQUAL_TO,
)
temp_storage_bytes = uniquer(
temp_storage=None,
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_num_selected,
op=OpKind.EQUAL_TO,
num_items=num_elements,
)
with alloc_stream:
temp_storage = cp.empty(temp_storage_bytes, dtype=np.uint8)
# Run once before timing to materialize the number of selected runs,
# matching the C++ metric accounting flow.
uniquer(
temp_storage=temp_storage,
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_num_selected,
op=OpKind.EQUAL_TO,
num_items=num_elements,
stream=alloc_stream,
)
alloc_stream.synchronize()
num_runs = int(d_num_selected.get()[0])
state.add_element_count(num_elements)
state.add_global_memory_reads(int(num_elements * d_in_keys.dtype.itemsize))
state.add_global_memory_reads(int(num_elements * d_in_values.dtype.itemsize))
state.add_global_memory_writes(int(num_runs * d_out_keys.dtype.itemsize))
state.add_global_memory_writes(int(num_runs * d_out_values.dtype.itemsize))
state.add_global_memory_writes(int(d_num_selected.dtype.itemsize))
def launcher(launch: bench.Launch):
uniquer(
temp_storage=temp_storage,
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_num_selected,
op=OpKind.EQUAL_TO,
num_items=num_elements,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
if __name__ == "__main__":
b = bench.register(bench_unique_by_key)
b.set_name("base")
b.add_string_axis("KeyT{ct}", list(KEY_TYPE_MAP.keys()))
b.add_string_axis("ValueT{ct}", list(VALUE_TYPE_MAP.keys()))
b.add_int64_power_of_two_axis("Elements{io}", range(16, 29, 4))
b.add_int64_power_of_two_axis("MaxSegSize", [1, 4, 8])
# Note: OffsetT axis from C++ is not exposed in Python API
bench.run_all_benchmarks(sys.argv)

237
cccl_upstream/python/cuda_cccl/benchmarks/compute/transform/babelstream.py Normal file
View File

@@ -0,0 +1,237 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for BabelStream operations using cuda.compute transforms.
C++ equivalent: cub/benchmarks/bench/transform/babelstream.cu
Notes:
- Migration: Python omits OffsetT axis and int128 types.
- OffsetT axis is omitted because the Python API does not expose offset type.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import cupy as cp
from utils import ALL_TYPES as _ALL_TYPES
from utils import as_cupy_stream
import cuda.bench as bench
import cuda.compute
from cuda.compute import ZipIterator
TYPE_MAP = {k: _ALL_TYPES[k] for k in ("I8", "I16", "F32", "F64")}
START_A = 11
START_B = 2
START_C = 1
START_SCALAR = -2
assert START_A == START_A + START_B + START_SCALAR * START_C
def _reset_pools():
cp.get_default_memory_pool().free_all_blocks()
cp.get_default_pinned_memory_pool().free_all_blocks()
def bench_mul(state: bench.State):
"""
Benchmark: b[i] = c[i] * scalar
Unary transform with scalar multiplication.
"""
type_str = state.get_string("T{ct}")
dtype = TYPE_MAP[type_str]
num_items = int(state.get_int64("Elements{io}"))
_reset_pools()
alloc_stream = as_cupy_stream(state.get_stream())
try:
with alloc_stream:
c = cp.full(num_items, START_C, dtype=dtype)
b = cp.full(num_items, START_B, dtype=dtype)
except (MemoryError, cp.cuda.memory.OutOfMemoryError):
state.skip("Skipping: out of memory.")
return
scalar = dtype(START_SCALAR)
def mul_op(ci):
return ci * scalar
transform = cuda.compute.make_unary_transform(d_in=c, d_out=b, op=mul_op)
state.add_element_count(num_items)
state.add_global_memory_reads(num_items * c.dtype.itemsize)
state.add_global_memory_writes(num_items * b.dtype.itemsize)
def launcher(launch: bench.Launch):
transform(
d_in=c, d_out=b, op=mul_op, num_items=num_items, stream=launch.get_stream()
)
state.exec(launcher, batched=False)
def bench_add(state: bench.State):
"""
Benchmark: c[i] = a[i] + b[i]
Binary transform with addition.
"""
type_str = state.get_string("T{ct}")
dtype = TYPE_MAP[type_str]
num_items = int(state.get_int64("Elements{io}"))
_reset_pools()
alloc_stream = as_cupy_stream(state.get_stream())
try:
with alloc_stream:
a = cp.full(num_items, START_A, dtype=dtype)
b = cp.full(num_items, START_B, dtype=dtype)
c = cp.full(num_items, START_C, dtype=dtype)
except (MemoryError, cp.cuda.memory.OutOfMemoryError):
state.skip("Skipping: out of memory.")
return
def add_op(ai, bi):
return ai + bi
transform = cuda.compute.make_binary_transform(d_in1=a, d_in2=b, d_out=c, op=add_op)
state.add_element_count(num_items)
state.add_global_memory_reads(2 * num_items * a.dtype.itemsize)
state.add_global_memory_writes(num_items * c.dtype.itemsize)
def launcher(launch: bench.Launch):
transform(
d_in1=a,
d_in2=b,
d_out=c,
op=add_op,
num_items=num_items,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
def bench_triad(state: bench.State):
"""
Benchmark: a[i] = b[i] + scalar * c[i]
Binary transform with fused multiply-add.
"""
type_str = state.get_string("T{ct}")
dtype = TYPE_MAP[type_str]
num_items = int(state.get_int64("Elements{io}"))
_reset_pools()
alloc_stream = as_cupy_stream(state.get_stream())
try:
with alloc_stream:
a = cp.full(num_items, START_A, dtype=dtype)
b = cp.full(num_items, START_B, dtype=dtype)
c = cp.full(num_items, START_C, dtype=dtype)
except (MemoryError, cp.cuda.memory.OutOfMemoryError):
state.skip("Skipping: out of memory.")
return
scalar = dtype(START_SCALAR)
def triad_op(bi, ci):
return bi + scalar * ci
transform = cuda.compute.make_binary_transform(
d_in1=b, d_in2=c, d_out=a, op=triad_op
)
state.add_element_count(num_items)
state.add_global_memory_reads(2 * num_items * a.dtype.itemsize)
state.add_global_memory_writes(num_items * a.dtype.itemsize)
def launcher(launch: bench.Launch):
transform(
d_in1=b,
d_in2=c,
d_out=a,
op=triad_op,
num_items=num_items,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
def bench_nstream(state: bench.State):
"""
Benchmark: a[i] = a[i] + b[i] + scalar * c[i]
Ternary transform using ZipIterator to combine (a, b, c) as input.
"""
type_str = state.get_string("T{ct}")
dtype = TYPE_MAP[type_str]
num_items = int(state.get_int64("Elements{io}"))
_reset_pools()
alloc_stream = as_cupy_stream(state.get_stream())
try:
with alloc_stream:
a = cp.full(num_items, START_A, dtype=dtype)
b = cp.full(num_items, START_B, dtype=dtype)
c = cp.full(num_items, START_C, dtype=dtype)
except (MemoryError, cp.cuda.memory.OutOfMemoryError):
state.skip("Skipping: out of memory.")
return
scalar = dtype(START_SCALAR)
# Use ZipIterator to combine 3 inputs into one for unary transform
zip_in = ZipIterator(a, b, c)
def nstream_op(abc):
return abc[0] + abc[1] + scalar * abc[2]
transform = cuda.compute.make_unary_transform(d_in=zip_in, d_out=a, op=nstream_op)
state.add_element_count(num_items)
state.add_global_memory_reads(3 * num_items * a.dtype.itemsize)
state.add_global_memory_writes(num_items * a.dtype.itemsize)
def launcher(launch: bench.Launch):
# Update ZipIterator state for each iteration
zip_in_iter = ZipIterator(a, b, c)
transform(
d_in=zip_in_iter,
d_out=a,
op=nstream_op,
num_items=num_items,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
# Registry of all BabelStream benchmarks
BENCHMARKS = {
"mul": bench_mul,
"add": bench_add,
"triad": bench_triad,
"nstream": bench_nstream,
}
if __name__ == "__main__":
for name, bench_fn in BENCHMARKS.items():
b = bench.register(bench_fn)
b.set_name(name)
b.add_string_axis("T{ct}", list(TYPE_MAP.keys()))
b.add_int64_power_of_two_axis("Elements{io}", range(16, 33, 4))
bench.run_all_benchmarks(sys.argv)

105
cccl_upstream/python/cuda_cccl/benchmarks/compute/transform/complex_cmp.py Normal file
View File

@@ -0,0 +1,105 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for complex comparison using cuda.compute.
C++ equivalent: cub/benchmarks/bench/transform/complex_cmp.cu
Notes:
- Uses two overlapping input ranges (in[0:n-1], in[1:n])
- Output is boolean array of size n-1
- Benchmark name is "compare_complex" to match C++
- Migration: Python uses explicit lexicographic compare for complex64.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import math
import cupy as cp
import numpy as np
from utils import as_cupy_stream, generate_data_with_entropy
import cuda.bench as bench
from cuda.compute import make_binary_transform
_COMPLEX_EPS = np.finfo(np.float32).eps
_COMPLEX_THRESHOLD = _COMPLEX_EPS * 2.0
def less_complex(a, b):
mag0 = math.sqrt(a.real * a.real + a.imag * a.imag)
mag1 = math.sqrt(b.real * b.real + b.imag * b.imag)
if math.isnan(mag0) or math.isnan(mag1):
return False
if math.isinf(mag0) or math.isinf(mag1):
scaler = 0.5
mag0 = math.sqrt(
(a.real * scaler) * (a.real * scaler)
+ (a.imag * scaler) * (a.imag * scaler)
)
mag1 = math.sqrt(
(b.real * scaler) * (b.real * scaler)
+ (b.imag * scaler) * (b.imag * scaler)
)
if abs(mag0 - mag1) < _COMPLEX_THRESHOLD:
phase0 = math.atan2(a.imag, a.real)
phase1 = math.atan2(b.imag, b.real)
return phase0 < phase1
return mag0 < mag1
def bench_compare_complex(state: bench.State):
num_elements = int(state.get_int64("Elements{io}"))
alloc_stream = as_cupy_stream(state.get_stream())
try:
with alloc_stream:
real = generate_data_with_entropy(
num_elements, np.float32, "1.000", alloc_stream
)
imag = generate_data_with_entropy(
num_elements, np.float32, "1.000", alloc_stream
)
d_in = (real + 1j * imag).astype(np.complex64)
d_out = cp.empty(num_elements - 1, dtype=np.bool_)
except (MemoryError, cp.cuda.memory.OutOfMemoryError):
state.skip("Skipping: out of memory.")
return
num_items = num_elements - 1
transformer = make_binary_transform(
d_in1=d_in[:-1], d_in2=d_in[1:], d_out=d_out, op=less_complex
)
state.add_element_count(num_elements)
state.add_global_memory_reads(num_elements * d_in.dtype.itemsize)
state.add_global_memory_writes(num_elements * d_out.dtype.itemsize)
def launcher(launch: bench.Launch):
transformer(
d_in1=d_in[:-1],
d_in2=d_in[1:],
d_out=d_out,
op=less_complex,
num_items=num_items,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
if __name__ == "__main__":
b = bench.register(bench_compare_complex)
b.set_name("compare_complex")
b.add_int64_power_of_two_axis("Elements{io}", range(16, 33, 4))
bench.run_all_benchmarks(sys.argv)

94
cccl_upstream/python/cuda_cccl/benchmarks/compute/transform/fib.py Normal file
View File

@@ -0,0 +1,94 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for transform fibonacci using cuda.compute.
C++ equivalent: cub/benchmarks/bench/transform/fib.cu
Notes:
- Input values are int64 in [0, 42]
- Output values are uint32
- Benchmark name is "fibonacci" to match C++
- Migration: Python fixes offsets to int64.
- OffsetT axis is omitted because the Python API does not expose offset type.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import cupy as cp
import numpy as np
from utils import as_cupy_stream, generate_data_with_entropy
import cuda.bench as bench
from cuda.compute import make_unary_transform
def fib_op(n):
t1 = 0
t2 = 1
if n < 1:
return t1
if n == 1:
return t1
if n == 2:
return t2
i = 3
while i <= n:
next_val = t1 + t2
t1 = t2
t2 = next_val
i += 1
return t2
def bench_transform_fib(state: bench.State):
# Axes
num_elements = int(state.get_int64("Elements{io}"))
alloc_stream = as_cupy_stream(state.get_stream())
try:
with alloc_stream:
d_in = generate_data_with_entropy(
num_elements,
np.int64,
"1.000",
alloc_stream,
min_val=np.int64(0),
max_val=np.int64(42),
)
d_out = cp.empty(num_elements, dtype=np.uint32)
except (MemoryError, cp.cuda.memory.OutOfMemoryError):
state.skip("Skipping: out of memory.")
return
transformer = make_unary_transform(d_in=d_in, d_out=d_out, op=fib_op)
state.add_element_count(num_elements)
state.add_global_memory_reads(num_elements * d_in.dtype.itemsize)
state.add_global_memory_writes(num_elements * d_out.dtype.itemsize)
def launcher(launch: bench.Launch):
transformer(
d_in=d_in,
d_out=d_out,
op=fib_op,
num_items=num_elements,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
if __name__ == "__main__":
b = bench.register(bench_transform_fib)
b.set_name("fibonacci")
b.add_int64_power_of_two_axis("Elements{io}", range(16, 33, 4))
bench.run_all_benchmarks(sys.argv)

72
cccl_upstream/python/cuda_cccl/benchmarks/compute/transform/fill.py Normal file
View File

@@ -0,0 +1,72 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for fill operation using cuda.compute.ConstantIterator.
C++ equivalent: cub/benchmarks/bench/transform/fill.cu
Notes:
- Migration: Python matches C++ integral_types (I8-I64); no tune parameters.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import cupy as cp
from utils import INTEGRAL_TYPES as TYPE_MAP
from utils import as_cupy_stream
import cuda.bench as bench
import cuda.compute
from cuda.compute import ConstantIterator, OpKind
def bench_fill(state: bench.State):
type_str = state.get_string("T{ct}")
dtype = TYPE_MAP[type_str]
num_items = int(state.get_int64("Elements{io}"))
# Setup data
alloc_stream = as_cupy_stream(state.get_stream())
try:
with alloc_stream:
d_out = cp.empty(num_items, dtype=dtype)
except (MemoryError, cp.cuda.memory.OutOfMemoryError):
state.skip("Skipping: out of memory.")
return
# Python equivalent of C++ return_constant<T>{42}
constant_it = ConstantIterator(dtype(42))
transform = cuda.compute.make_unary_transform(
d_in=constant_it, d_out=d_out, op=OpKind.IDENTITY
)
state.add_element_count(num_items)
state.add_global_memory_reads(0)
state.add_global_memory_writes(num_items * d_out.dtype.itemsize)
def launcher(launch: bench.Launch):
transform(
d_in=constant_it,
d_out=d_out,
op=OpKind.IDENTITY,
num_items=num_items,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
if __name__ == "__main__":
b = bench.register(bench_fill)
b.set_name("fill")
b.add_string_axis("T{ct}", list(TYPE_MAP.keys()))
b.add_int64_power_of_two_axis("Elements{io}", range(16, 33, 4))
bench.run_all_benchmarks(sys.argv)

89
cccl_upstream/python/cuda_cccl/benchmarks/compute/transform/grayscale.py Normal file
View File

@@ -0,0 +1,89 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for transform grayscale using cuda.compute.
C++ equivalent: cub/benchmarks/bench/transform/grayscale.cu
Notes:
- Input is an RGB struct with three channels
- Output is grayscale value of the same type
- Benchmark name is "grayscale" to match C++
- Migration: Python uses AoS (`gpu_struct`) to mirror C++ `rgb_t<T>`.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import cupy as cp
from utils import FLOAT_TYPES as TYPE_MAP
from utils import as_cupy_stream, generate_data_with_entropy
import cuda.bench as bench
from cuda.compute import gpu_struct, make_unary_transform
def bench_transform_grayscale(state: bench.State):
type_str = state.get_string("T{ct}")
dtype = TYPE_MAP[type_str]
num_elements = int(state.get_int64("Elements{io}"))
# Grayscale weights
w_r = dtype(0.2989)
w_g = dtype(0.587)
w_b = dtype(0.114)
RGB = gpu_struct({"r": dtype, "g": dtype, "b": dtype})
def to_grayscale(pixel: RGB):
return w_r * pixel.r + w_g * pixel.g + w_b * pixel.b
alloc_stream = as_cupy_stream(state.get_stream())
try:
with alloc_stream:
r_data = generate_data_with_entropy(
num_elements, dtype, "1.000", alloc_stream
)
g_data = generate_data_with_entropy(
num_elements, dtype, "1.000", alloc_stream
)
b_data = generate_data_with_entropy(
num_elements, dtype, "1.000", alloc_stream
)
d_pixels = cp.empty(num_elements, dtype=RGB.dtype)
d_pixels["r"] = r_data
d_pixels["g"] = g_data
d_pixels["b"] = b_data
d_out = cp.empty(num_elements, dtype=dtype)
except (MemoryError, cp.cuda.memory.OutOfMemoryError):
state.skip("Skipping: out of memory.")
return
transformer = make_unary_transform(d_in=d_pixels, d_out=d_out, op=to_grayscale)
state.add_element_count(num_elements)
state.add_global_memory_reads(num_elements * d_pixels.dtype.itemsize)
state.add_global_memory_writes(num_elements * d_out.dtype.itemsize)
def launcher(launch: bench.Launch):
transformer(
d_in=d_pixels,
d_out=d_out,
op=to_grayscale,
num_items=num_elements,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
if __name__ == "__main__":
b = bench.register(bench_transform_grayscale)
b.set_name("grayscale")
b.add_string_axis("T{ct}", list(TYPE_MAP.keys()))
b.add_int64_power_of_two_axis("Elements{io}", range(16, 33, 4))
bench.run_all_benchmarks(sys.argv)

138
cccl_upstream/python/cuda_cccl/benchmarks/compute/transform/heavy.py Normal file
View File

@@ -0,0 +1,138 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Python benchmark for heavy transform using cuda.compute.
C++ equivalent: cub/benchmarks/bench/transform/heavy.cu
Notes:
- Migration: Python uses Numba local arrays to emulate register pressure.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import cupy as cp
import numba
import numpy as np
from numba import cuda as lang
from utils import as_cupy_stream, generate_data_with_entropy
import cuda.bench as bench
import cuda.compute
def _heavy_op_32(data):
reg = lang.local.array(shape=32, dtype=numba.uint32)
reg[0] = data
for i in range(1, 32):
x = reg[i - 1]
reg[i] = x * x + 1
for i in range(32):
x = reg[i]
reg[i] = (x * x) % 19
for i in range(32):
reg[i] = reg[32 - i - 1] * reg[i]
out = data - data # uint32(0)
for i in range(32):
out += reg[i]
return out
def _heavy_op_64(data):
reg = lang.local.array(shape=64, dtype=numba.uint32)
reg[0] = data
for i in range(1, 64):
x = reg[i - 1]
reg[i] = x * x + 1
for i in range(64):
x = reg[i]
reg[i] = (x * x) % 19
for i in range(64):
reg[i] = reg[64 - i - 1] * reg[i]
out = data - data
for i in range(64):
out += reg[i]
return out
def _heavy_op_128(data):
reg = lang.local.array(shape=128, dtype=numba.uint32)
reg[0] = data
for i in range(1, 128):
x = reg[i - 1]
reg[i] = x * x + 1
for i in range(128):
x = reg[i]
reg[i] = (x * x) % 19
for i in range(128):
reg[i] = reg[128 - i - 1] * reg[i]
out = data - data
for i in range(128):
out += reg[i]
return out
def _heavy_op_256(data):
reg = lang.local.array(shape=256, dtype=numba.uint32)
reg[0] = data
for i in range(1, 256):
x = reg[i - 1]
reg[i] = x * x + 1
for i in range(256):
x = reg[i]
reg[i] = (x * x) % 19
for i in range(256):
reg[i] = reg[256 - i - 1] * reg[i]
out = data - data
for i in range(256):
out += reg[i]
return out
_HEAVY_OPS = {
32: _heavy_op_32,
64: _heavy_op_64,
128: _heavy_op_128,
256: _heavy_op_256,
}
def bench_heavy(state: bench.State):
# Axes
n_regs = int(state.get_string("Heaviness{ct}"))
size = int(state.get_int64("Elements{io}"))
alloc_stream = as_cupy_stream(state.get_stream())
try:
with alloc_stream:
d_in = generate_data_with_entropy(size, np.uint32, "1.000", alloc_stream)
d_out = cp.empty(size, dtype=np.uint32)
except (MemoryError, cp.cuda.memory.OutOfMemoryError):
state.skip("Skipping: out of memory.")
return
op = _HEAVY_OPS[n_regs]
transform = cuda.compute.make_unary_transform(d_in=d_in, d_out=d_out, op=op)
state.add_element_count(size)
state.add_global_memory_reads(size * d_in.dtype.itemsize)
state.add_global_memory_writes(size * d_out.dtype.itemsize)
def launcher(launch: bench.Launch):
transform(
d_in=d_in, d_out=d_out, op=op, num_items=size, stream=launch.get_stream()
)
state.exec(launcher, batched=False)
if __name__ == "__main__":
b = bench.register(bench_heavy)
b.set_name("heavy")
b.add_string_axis("Heaviness{ct}", [str(v) for v in (32, 64, 128, 256)])
b.add_int64_power_of_two_axis("Elements{io}", range(16, 33, 4))
bench.run_all_benchmarks(sys.argv)

86
cccl_upstream/python/cuda_cccl/benchmarks/compute/transform_reduce/sum.py Normal file
View File

@@ -0,0 +1,86 @@
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Python benchmark for transform reduce sum using cuda.compute.
C++ equivalent: cub/benchmarks/bench/transform_reduce/sum.cu
Notes:
- Uses TransformIterator with a square operation
- OffsetT axis from C++ is fixed to Python default (int64)
- Migration: Python fixes offsets and omits int128/complex types.
- OffsetT axis is omitted because the Python API does not expose offset type.
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import cupy as cp
import numpy as np
from utils import SIGNED_TYPES as TYPE_MAP
from utils import as_cupy_stream, generate_data_with_entropy
import cuda.bench as bench
from cuda.compute import OpKind, TransformIterator, make_reduce_into
def square_op(x):
return x * x
def bench_transform_reduce_sum(state: bench.State):
type_str = state.get_string("T{ct}")
dtype = TYPE_MAP[type_str]
num_items = int(state.get_int64("Elements{io}"))
alloc_stream = as_cupy_stream(state.get_stream())
with alloc_stream:
d_in = generate_data_with_entropy(num_items, dtype, "1.000", alloc_stream)
d_out = cp.empty(1, dtype=dtype)
transform_it = TransformIterator(d_in, square_op)
h_init = np.zeros(1, dtype=dtype)
reducer = make_reduce_into(
d_in=transform_it, d_out=d_out, op=OpKind.PLUS, h_init=h_init
)
temp_storage_bytes = reducer(
temp_storage=None,
d_in=transform_it,
d_out=d_out,
num_items=num_items,
op=OpKind.PLUS,
h_init=h_init,
)
with alloc_stream:
temp_storage = cp.empty(temp_storage_bytes, dtype=np.uint8)
state.add_element_count(num_items)
state.add_global_memory_reads(num_items * d_in.dtype.itemsize, "Size")
state.add_global_memory_writes(1 * d_out.dtype.itemsize)
def launcher(launch: bench.Launch):
reducer(
temp_storage=temp_storage,
d_in=transform_it,
d_out=d_out,
num_items=num_items,
op=OpKind.PLUS,
h_init=h_init,
stream=launch.get_stream(),
)
state.exec(launcher, batched=False)
if __name__ == "__main__":
b = bench.register(bench_transform_reduce_sum)
b.set_name("base")
b.add_string_axis("T{ct}", list(TYPE_MAP.keys()))
b.add_int64_power_of_two_axis("Elements{io}", range(16, 29, 4))
bench.run_all_benchmarks(sys.argv)

252
cccl_upstream/python/cuda_cccl/benchmarks/compute/utils.py Normal file
View File

@@ -0,0 +1,252 @@
import cupy as cp
import numpy as np
import cuda.bench as bench
ALL_TYPES = {
"I8": np.int8,
"I16": np.int16,
"I32": np.int32,
"I64": np.int64,
"U8": np.uint8,
"U16": np.uint16,
"U32": np.uint32,
"U64": np.uint64,
"F32": np.float32,
"F64": np.float64,
}
SIGNED_TYPES = {k: ALL_TYPES[k] for k in ("I8", "I16", "I32", "I64", "F32", "F64")}
FLOAT_TYPES = {k: ALL_TYPES[k] for k in ("F32", "F64")}
# Matches C++ integral_types = {int8_t, int16_t, int32_t, int64_t}
INTEGRAL_TYPES = {k: ALL_TYPES[k] for k in ("I8", "I16", "I32", "I64")}
# Matches C++ fundamental_types = {int8..int64, [int128,] float, double}
# int128 is excluded because it is not supported by numpy/cupy.
FUNDAMENTAL_TYPES = {k: ALL_TYPES[k] for k in ("I8", "I16", "I32", "I64", "F32", "F64")}
ENTROPY_TO_STEPS = {
"1.000": 0,
"0.811": 1,
"0.544": 2,
"0.337": 3,
"0.201": 4,
"0.000": 0,
}
ENTROPY_TO_PROB = {
"1.000": 1.0,
"0.811": 0.811,
"0.544": 0.544,
"0.337": 0.337,
"0.201": 0.201,
"0.000": 0.0,
}
def as_cupy_stream(cs: bench.CudaStream) -> cp.cuda.Stream:
"""Convert nvbench CudaStream to CuPy Stream."""
return cp.cuda.ExternalStream(cs.addressof())
def lerp_min_max(dtype, probability):
"""Interpolate between min/max for dtype like nvbench_helper.cuh."""
if probability == 1.0:
if np.issubdtype(dtype, np.integer):
return np.iinfo(dtype).max
return np.finfo(dtype).max
if np.issubdtype(dtype, np.integer):
min_val = float(np.iinfo(dtype).min)
max_val = float(np.iinfo(dtype).max)
else:
min_val = float(np.finfo(dtype).min)
max_val = float(np.finfo(dtype).max)
return dtype(min_val + probability * (max_val - min_val))
def _bitwise_and(a, b, dtype):
if np.issubdtype(dtype, np.floating):
view_dtype = cp.uint32 if dtype == np.float32 else cp.uint64
return (a.view(view_dtype) & b.view(view_dtype)).view(dtype)
return a & b
def _uniform_random(num_elements, dtype, min_val, max_val):
rand = cp.random.random(num_elements)
if np.issubdtype(dtype, np.floating):
return ((float(max_val) - float(min_val)) * rand + float(min_val)).astype(dtype)
min_f = float(min_val)
max_f = float(max_val)
return cp.floor((max_f - min_f + 1) * rand + min_f).astype(dtype)
def generate_data_with_entropy(
num_elements, dtype, entropy_str, stream, min_val=None, max_val=None
):
"""Generate data with nvbench_helper-style bit entropy."""
if min_val is None or max_val is None:
if np.issubdtype(dtype, np.integer):
info = np.iinfo(dtype)
default_min = info.min
default_max = info.max
else:
info = np.finfo(dtype)
default_min = info.tiny
default_max = info.max
min_val = default_min if min_val is None else min_val
max_val = default_max if max_val is None else max_val
steps = ENTROPY_TO_STEPS[entropy_str]
with stream:
if entropy_str == "0.000":
scalar = _uniform_random(1, dtype, min_val, max_val)[0]
data = cp.full(num_elements, scalar, dtype=dtype)
else:
data = _uniform_random(num_elements, dtype, min_val, max_val)
for _ in range(steps):
tmp = _uniform_random(num_elements, dtype, min_val, max_val)
data = _bitwise_and(data, tmp, dtype)
return data
def generate_uniform_segment_offsets(num_elements, min_segment_size, max_segment_size):
num_elements = int(num_elements)
if min_segment_size <= 0:
raise ValueError("min_segment_size must be positive")
if max_segment_size < min_segment_size:
raise ValueError("max_segment_size must be >= min_segment_size")
num_segments_est = int(np.ceil(num_elements / min_segment_size))
if num_segments_est > np.iinfo(np.int32).max:
raise MemoryError("Too many segments for int32 offsets")
sizes = cp.random.randint(
min_segment_size,
max_segment_size + 1,
size=num_segments_est,
dtype=cp.int64,
)
cumsum = cp.cumsum(sizes)
cutoff = int(
cp.searchsorted(
cumsum, cp.asarray(num_elements, dtype=cp.int64), side="left"
).item()
)
sizes = sizes[: cutoff + 1]
prev = 0 if cutoff == 0 else int(cumsum[cutoff - 1].item())
sizes[cutoff] = num_elements - prev
offsets = cp.empty(cutoff + 2, dtype=cp.int64)
offsets[0] = 0
offsets[1:] = cp.cumsum(sizes)
offsets[-1] = num_elements
return offsets
def generate_power_law_offsets(num_elements, num_segments):
if num_segments <= 0:
return cp.asarray([0, num_elements], dtype=cp.int64)
# Mirror nvbench_helper power-law generation:
# draw log-normal samples, normalize to total elements,
# floor to integer segment sizes, then distribute remainder
# across the first `diff` segments.
samples = cp.random.lognormal(3.0, 1.2, size=num_segments)
if int(cp.count_nonzero(samples).item()) == 0:
samples = cp.ones(num_segments, dtype=cp.float64)
sample_sum = float(samples.sum().item())
sizes = cp.floor(samples * num_elements / sample_sum).astype(cp.int64)
diff = int(num_elements - sizes.sum().item())
if diff > 0:
sizes[:diff] += 1
offsets = cp.empty(num_segments + 1, dtype=cp.int64)
offsets[0] = 0
offsets[1:] = cp.cumsum(sizes)
return offsets
def generate_fixed_segment_offsets(num_elements, segment_size, stream):
num_segments = max(1, num_elements // segment_size)
actual_elements = num_segments * segment_size
with stream:
start_offsets = cp.arange(0, actual_elements, segment_size, dtype=np.int64)
end_offsets = cp.arange(
segment_size, actual_elements + 1, segment_size, dtype=np.int64
)
return start_offsets, end_offsets, num_segments, actual_elements
def generate_key_segments(
num_elements, key_dtype, min_segment_size, max_segment_size, stream
):
"""Generate GPU key segments (runs of equal keys) matching C++ generate.uniform.key_segments.
All computation stays on GPU via CuPy. We avoid ``cp.repeat`` (which
doesn't accept a device array for *repeats*) by building a segment-id
array through ``cp.searchsorted`` on cumulative offsets instead.
"""
num_elements = int(num_elements)
if min_segment_size <= 0:
raise ValueError("min_segment_size must be positive")
if max_segment_size < min_segment_size:
raise ValueError("max_segment_size must be >= min_segment_size")
num_segments_est = int(np.ceil(num_elements / min_segment_size))
if num_segments_est > np.iinfo(np.int32).max:
raise MemoryError("Too many segments for int32 offsets")
with stream:
sizes = cp.random.randint(
min_segment_size,
max_segment_size + 1,
size=num_segments_est,
dtype=cp.int64,
)
cumsum = cp.cumsum(sizes)
# Find how many full segments fit within num_elements
cutoff = int(
cp.searchsorted(
cumsum, cp.asarray(num_elements, dtype=cp.int64), side="left"
).item()
)
sizes = sizes[: cutoff + 1]
prev = 0 if cutoff == 0 else int(cumsum[cutoff - 1].item())
sizes[cutoff] = num_elements - prev
# Build cumulative offsets for the final segments
offsets = cp.empty(cutoff + 2, dtype=cp.int64)
offsets[0] = 0
offsets[1:] = cp.cumsum(sizes)
# Instead of cp.repeat (which doesn't support device repeats),
# use searchsorted to map each element index to its segment id.
indices = cp.arange(num_elements, dtype=cp.int64)
# searchsorted(offsets[1:], indices, side="right") gives the segment id
segment_ids = cp.searchsorted(offsets[1:], indices, side="right")
# Map segment ids to key values, wrapping within dtype range
if np.issubdtype(key_dtype, np.integer):
info = np.iinfo(key_dtype)
if np.dtype(key_dtype).itemsize < 8:
range_size = int(info.max) - int(info.min) + 1
keys = ((segment_ids % range_size) + int(info.min)).astype(
key_dtype, copy=False
)
else:
# For int64, avoid Python overflow: just cast directly
keys = segment_ids.astype(key_dtype, copy=False)
else:
keys = segment_ids.astype(key_dtype, copy=False)
return keys

27
cccl_upstream/python/cuda_cccl/cuda/cccl/__init__.py Normal file
View File

@@ -0,0 +1,27 @@
"""
CUDA Core Library (CCCL) Python Package
"""
import importlib.metadata
try:
__version__ = importlib.metadata.version("cuda-cccl")
except Exception:
__version__ = "0.0.0"
from .headers.include_paths import get_include_paths
# cuda.bindings is required, but instead of being listed as a required dependency,
# it is installed via an extra (e.g., [cu12] or [cu13]).
#
# One of the first things we should do is check that it is available, and raise
# a helpful error message if it is not.
try:
import cuda.bindings as _cuda_bindings # type: ignore
except ImportError:
raise ImportError(
"cuda.bindings is not installed. Please install the appropriate extra cuda-cccl[cu12] or cuda-cccl[cu13]."
) from None
del _cuda_bindings
__all__ = ["get_include_paths", "__version__"]

24
cccl_upstream/python/cuda_cccl/cuda/cccl/_cuda_version_utils.py Normal file
View File

@@ -0,0 +1,24 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
CUDA version detection utilities shared across the cccl package.
"""
from typing import Optional
import cuda.bindings
def detect_cuda_version() -> Optional[int]:
cuda_version = cuda.bindings.__version__
return int(cuda_version.split(".")[0])
def get_recommended_extra(cuda_version: Optional[int]) -> str:
"""Get the recommended pip extra for the detected CUDA version."""
if cuda_version == 13:
return "cu13"
else:
return "cu12"

7
cccl_upstream/python/cuda_cccl/cuda/cccl/headers/__init__.py Normal file
View File

@@ -0,0 +1,7 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from .include_paths import get_include_paths
__all__ = ["__version__", "get_include_paths"]

1
cccl_upstream/python/cuda_cccl/cuda/cccl/headers/include/__init__.py Normal file
View File

@@ -0,0 +1 @@
# Intentionally empty

51
cccl_upstream/python/cuda_cccl/cuda/cccl/headers/include_paths.py Normal file
View File

@@ -0,0 +1,51 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
import sys
from dataclasses import dataclass
from functools import lru_cache
from importlib.resources import as_file, files
from pathlib import Path
from typing import Optional
# type: ignore[import-not-found]
from cuda.pathfinder import find_nvidia_header_directory
@dataclass
class IncludePaths:
cuda: Optional[Path]
libcudacxx: Optional[Path]
cub: Optional[Path]
thrust: Optional[Path]
def as_tuple(self):
# Note: higher-level ... lower-level order:
return (self.thrust, self.cub, self.libcudacxx, self.cuda)
@lru_cache()
def get_include_paths(probe_file: str = "cub/version.cuh") -> IncludePaths:
cuda_incl = find_nvidia_header_directory("cudart")
if cuda_incl is None:
raise RuntimeError("Unable to locate CUDA include directory.")
with as_file(files("cuda.cccl.headers.include")) as f:
cccl_incl = Path(f)
probe_file_path = Path(probe_file)
if not (cccl_incl / probe_file_path).exists():
for sp in sys.path:
cccl_incl = Path(sp).resolve() / "cuda" / "cccl" / "headers" / "include"
if (cccl_incl / probe_file_path).exists():
break
else:
raise RuntimeError("Unable to locate CCCL include directory.")
return IncludePaths(
cuda=cuda_incl,
libcudacxx=cccl_incl,
cub=cccl_incl,
thrust=cccl_incl,
)

9
cccl_upstream/python/cuda_cccl/cuda/cccl/parallel/__init__.py Normal file
View File

@@ -0,0 +1,9 @@
# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License -Identifier: Apache-2.0 WITH LLVM-exception
from . import experimental
__all__ = [
"experimental",
]

24
cccl_upstream/python/cuda_cccl/cuda/cccl/parallel/experimental/__init__.py Normal file
View File

@@ -0,0 +1,24 @@
# Copyright (c) 2025, NVIDIA CORPORATION.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# alias for backwards compatibility
from warnings import warn
from cuda.compute import * # noqa: F403
warn(
"The module cuda.cccl.parallel.experimental is deprecated. Use cuda.compute instead.",
FutureWarning,
)

0
cccl_upstream/python/cuda_cccl/cuda/cccl/py.typed Normal file
View File

162
cccl_upstream/python/cuda_cccl/cuda/compute/__init__.py Normal file
View File

@@ -0,0 +1,162 @@
# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
# When built against the v2 (HostJIT) backend, the JIT loads Clang's CUDA
# headers and our cuda_minimal stubs from paths that don't exist on the
# user's machine. The wheel bundles both under cuda/cccl/headers/{clang,…};
# point hostjit at them via the env vars its detectDefaultConfig() reads.
# Only sets vars that aren't already configured by the user, and skips
# silently if the bundled directories are absent (e.g. v1 builds).
def _configure_hostjit_paths() -> None:
import os
from pathlib import Path
try:
from ._build_info import USING_V2 # type: ignore[import-not-found]
except ImportError:
return
if not USING_V2:
return
# Probe for actual file presence, not just directory existence: editable
# (`pip install -e`) installs leave behind empty placeholder dirs in the
# source tree (with just `__pycache__`), so `is_dir()` succeeds but the
# bundled headers are absent. In that case, leave the env vars unset and
# let the C library use its build-time CLANG_HEADERS_DIR / HOSTJIT_INCLUDE_DIR
# macros (pointing at the LLVM source tree under the CMake build dir).
headers_dir = Path(__file__).resolve().parent.parent / "cccl" / "headers"
clang_dir = headers_dir / "clang"
if (
clang_dir / "__clang_cuda_math_forward_declares.h"
).is_file() and not os.environ.get("HOSTJIT_CLANG_PATH"):
os.environ["HOSTJIT_CLANG_PATH"] = str(clang_dir)
if (
headers_dir / "hostjit" / "cuda_minimal" / "__clang_cuda_runtime_wrapper.h"
).is_file() and not os.environ.get("HOSTJIT_INCLUDE_PATH"):
os.environ["HOSTJIT_INCLUDE_PATH"] = str(headers_dir)
_configure_hostjit_paths()
from ._bindings import _BINDINGS_AVAILABLE # type: ignore[attr-defined] # noqa: E402
if not _BINDINGS_AVAILABLE:
__all__ = ["_BINDINGS_AVAILABLE"]
def __getattr__(name):
raise AttributeError(
f"Cannot access 'cuda.compute.{name}' because CUDA bindings are not available."
"This typically means you're running on a CPU-only machine without CUDA drivers installed."
)
else:
from ._caching import clear_all_caches
from ._proxy import ProxyArray, ProxyValue
from .algorithms import (
DoubleBuffer,
SortOrder,
binary_transform,
deserialize,
exclusive_scan,
histogram_even,
inclusive_scan,
lower_bound,
make_binary_transform,
make_exclusive_scan,
make_histogram_even,
make_inclusive_scan,
make_lower_bound,
make_merge_sort,
make_radix_sort,
make_reduce_into,
make_segmented_reduce,
make_segmented_sort,
make_select,
make_three_way_partition,
make_unary_transform,
make_unique_by_key,
make_upper_bound,
merge_sort,
radix_sort,
reduce_into,
segmented_reduce,
segmented_sort,
select,
serialize,
three_way_partition,
unary_transform,
unique_by_key,
upper_bound,
)
from .determinism import Determinism
from .iterators import (
CacheModifiedInputIterator,
ConstantIterator,
CountingIterator,
DiscardIterator,
PermutationIterator,
ReverseIterator,
ShuffleIterator,
TransformIterator,
TransformOutputIterator,
ZipIterator,
)
from .op import OpKind
from .struct import gpu_struct
__all__ = [
"_BINDINGS_AVAILABLE",
"serialize",
"deserialize",
"ProxyArray",
"ProxyValue",
"binary_transform",
"clear_all_caches",
"CacheModifiedInputIterator",
"ConstantIterator",
"CountingIterator",
"DiscardIterator",
"DoubleBuffer",
"exclusive_scan",
"gpu_struct",
"histogram_even",
"inclusive_scan",
"lower_bound",
"make_binary_transform",
"make_exclusive_scan",
"make_select",
"make_histogram_even",
"make_inclusive_scan",
"make_lower_bound",
"make_merge_sort",
"make_radix_sort",
"make_reduce_into",
"make_segmented_reduce",
"make_segmented_sort",
"make_three_way_partition",
"make_unary_transform",
"make_unique_by_key",
"make_upper_bound",
"merge_sort",
"OpKind",
"Determinism",
"PermutationIterator",
"radix_sort",
"reduce_into",
"ReverseIterator",
"ShuffleIterator",
"segmented_reduce",
"segmented_sort",
"select",
"SortOrder",
"TransformIterator",
"TransformOutputIterator",
"three_way_partition",
"unary_transform",
"unique_by_key",
"upper_bound",
"ZipIterator",
]

87
cccl_upstream/python/cuda_cccl/cuda/compute/_bindings.py Normal file
View File

@@ -0,0 +1,87 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# _bindings.py is a shim module that imports symbols from a
# _bindings_impl extension module. The shim serves the following purposes:
#
# 1. Import a CUDA-specific extension. The cuda.cccl wheel ships with multiple
# extensions, one for each CUDA version. At runtime, this shim chooses the
# appropriate extension based on the detected CUDA version, and imports all
# symbols from it.
#
# 2. Preload `nvrtc` and `nvJitLink` before importing the extension.
# These shared libraries are indirect dependencies, pulled in via the direct
# dependency `cccl.c.parallel`. To ensure reliable symbol resolution at
# runtime, we explicitly load them first using `cuda.pathfinder`.
# Without this step, importing the Cython extension directly may fail or behave
# inconsistently depending on environment setup and dynamic linker behavior.
# This indirection ensures the right loading order, regardless of how
# `_bindings` is first imported across the codebase.
#
# 3. On Windows, add the directory containing cccl.c.parallel's dependent DLL
# (e.g. cuda/cccl/parallel/experimental/cu13/_bindings_impl.cp312-win_amd64.pyd)
# to the current process's DLL search path using `os.add_dll_directory`.
from __future__ import annotations
import importlib
import os
from cuda.cccl._cuda_version_utils import detect_cuda_version, get_recommended_extra
from cuda.pathfinder import ( # type: ignore[import-not-found]
load_nvidia_dynamic_lib,
)
def _load_cuda_libraries():
# Load appropriate libraries for the detected CUDA version
for libname in ("nvrtc", "nvJitLink"):
load_nvidia_dynamic_lib(libname)
_load_cuda_libraries()
# Import the appropriate bindings implementation depending on what
# CUDA version is available:
cuda_version = detect_cuda_version()
if cuda_version not in [12, 13]:
raise RuntimeError(
f"Unsupported CUDA version: {cuda_version}. Only CUDA 12 and 13 are supported."
)
# `extra_name` is one of "cu12", "cu13", etc.
extra_name = get_recommended_extra(cuda_version)
module_suffix = f".{extra_name}._bindings_impl"
module_fullname = __package__ + module_suffix
# On Windows, ensure the dependent DLLs next to the extension are discoverable.
# The extension lives at .../experimental/<extra_name>/_bindings_impl.*.pyd
# and its dependent DLLs are under .../experimental/<extra_name>/cccl/.
if os.name == "nt":
spec = importlib.util.find_spec(module_fullname)
if spec and spec.origin:
dll_dir = os.path.join(os.path.dirname(spec.origin), "cccl")
if os.path.isdir(dll_dir):
# Assign the DLL directory handle to a global such that it stays
# alive for the lifetime of this module (and thus, keeps the DLL
# directory in the search path).
try:
_cccl_dll_dir_handle = os.add_dll_directory(dll_dir) # noqa: F841
except Exception:
pass
_BINDINGS_AVAILABLE = False
try:
bindings_module = importlib.import_module(module_suffix, __package__)
# Import all symbols from the module
globals().update(bindings_module.__dict__)
_BINDINGS_AVAILABLE = True
except ImportError as e:
import warnings
warnings.warn(
f"CUDA CCCL bindings for CUDA {cuda_version} not available: {e}",
RuntimeWarning,
)

658
cccl_upstream/python/cuda_cccl/cuda/compute/_bindings.pyi Normal file
View File

@@ -0,0 +1,658 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
import ctypes
from enum import IntEnum
from typing import Any, Optional
from typing_extensions import Buffer
class OpKind(IntEnum):
_value_: int
STATELESS = ...
STATEFUL = ...
PLUS = ...
MINUS = ...
MULTIPLIES = ...
DIVIDES = ...
MODULUS = ...
EQUAL_TO = ...
NOT_EQUAL_TO = ...
GREATER = ...
LESS = ...
GREATER_EQUAL = ...
LESS_EQUAL = ...
LOGICAL_AND = ...
LOGICAL_OR = ...
LOGICAL_NOT = ...
BIT_AND = ...
BIT_OR = ...
BIT_XOR = ...
BIT_NOT = ...
IDENTITY = ...
NEGATE = ...
MINIMUM = ...
MAXIMUM = ...
class TypeEnum(IntEnum):
_value_: int
INT8 = ...
INT16 = ...
INT32 = ...
INT64 = ...
UINT8 = ...
UINT16 = ...
UINT32 = ...
UINT64 = ...
FLOAT16 = ...
FLOAT32 = ...
FLOAT64 = ...
STORAGE = ...
BOOLEAN = ...
class IteratorKind(IntEnum):
_value_: int
POINTER = ...
ITERATOR = ...
class SortOrder(IntEnum):
_value_: int
ASCENDING = ...
DESCENDING = ...
class InitKind(IntEnum):
_value_: int
NO_INIT = ...
FUTURE_VALUE_INIT = ...
VALUE_INIT = ...
class Determinism(IntEnum):
_value_: int
NOT_GUARANTEED = ...
RUN_TO_RUN = ...
GPU_TO_GPU = ...
class BinarySearchMode(IntEnum):
_value_: int
LOWER_BOUND = ...
UPPER_BOUND = ...
class Op:
def __init__(
self,
name: Optional[str] = ...,
operator_type: OpKind = ...,
ltoir=None,
state=None,
state_alignment: int = 1,
extra_ltoirs=None,
): ...
@property
def state(self) -> bytes: ...
@state.setter
def state(self, new_value: bytes) -> None: ...
@property
def name(self) -> str: ...
@property
def ltoir(self) -> bytes: ...
@property
def state_alignment(self) -> int: ...
@property
def operator_type(self) -> OpKind: ...
@property
def code(self) -> Any: ...
@property
def extra_code(self) -> list: ...
class TypeInfo:
def __init__(self, size: int, alignment: int, type_enum: TypeEnum): ...
@property
def size(self) -> int: ...
@property
def alignment(self) -> int: ...
@property
def typenum(self) -> int: ...
def as_bytes(self) -> bytes: ...
class Value:
def __init__(self, type: TypeInfo, state: memoryview): ...
@property
def type(self) -> TypeInfo: ...
@property
def state(self) -> memoryview: ...
@state.setter
def state(self, new_value: memoryview) -> None: ...
def as_bytes(self) -> bytes: ...
class Pointer:
def __init__(self, arg): ...
def make_pointer_object(ptr: int | ctypes.c_void_p, owner: Any) -> Pointer: ...
class IteratorState(Buffer):
def __init__(self, arg): ...
@property
def size(self) -> int: ...
class Iterator:
def __init__(
self,
alignment: int,
iterator_type: IteratorKind,
advance_fn: Op,
dereference_fn: Op,
value_type: TypeInfo,
state=None,
host_advance_fn=None,
):
pass
@property
def advance_op(self) -> Op: ...
@property
def dereference_op(self) -> Op: ...
@property
def dereference_or_assign_op(self) -> Op: ...
@property
def state(self): ...
@state.setter
def state(self, value) -> None: ...
@property
def type(self) -> IteratorKind: ...
@property
def alignment(self) -> int: ...
@property
def value_type(self) -> TypeInfo: ...
def as_bytes(self) -> bytes: ...
def is_kind_pointer(self) -> bool: ...
def is_kind_iterator(self) -> bool: ...
@property
def host_advance_fn(self): ...
@host_advance_fn.setter
def host_advance_fn(self, value) -> None: ...
class CommonData:
def __init__(
self,
cc_major: int,
cc_minor: int,
cub_path: str,
thrust_path: str,
libcudacxx_path: str,
ctk_path: str,
): ...
@property
def compute_capability(self) -> tuple[int, int]: ...
@property
def cub_path(self) -> str: ...
@property
def thrust_path(self) -> str: ...
@property
def libcudacxx_path(self) -> str: ...
@property
def ctk_path(self) -> str: ...
# ------------
# DeviceReduce
# ------------
class DeviceReduceBuildResult:
def __init__(
self,
d_in: Iterator,
d_out: Iterator,
binary_op: Op,
h_init: Value,
determinism: Determinism,
info: CommonData,
): ...
def compute(
self,
temp_storage_ptr: int | None,
temp_storage_nbytes: int,
d_in: Iterator,
d_out: Iterator,
num_items: int,
binary_op: Op,
h_init: Value,
stream,
) -> int: ...
def compute_nondeterministic(
self,
temp_storage_ptr: int | None,
temp_storage_nbytes: int,
d_in: Iterator,
d_out: Iterator,
num_items: int,
binary_op: Op,
h_init: Value,
stream,
) -> int: ...
def serialize(self) -> bytes: ...
@staticmethod
def deserialize(
blob: bytes, load: bool = ..., check_cc: bool = ...
) -> DeviceReduceBuildResult: ...
@staticmethod
def compile(*args) -> DeviceReduceBuildResult: ...
def load(self) -> None: ...
@property
def determinism(self) -> int: ...
# ----------
# DeviceScan
# ----------
class DeviceScanBuildResult:
def __init__(
self,
d_in: Iterator,
d_out: Iterator,
binary_op: Op,
init_type: TypeInfo,
force_inclusive: bool,
init_kind: InitKind,
info: CommonData,
): ...
def compute_inclusive(
self,
temp_storage_ptr: int | None,
temp_storage_nbytes: int,
d_in: Iterator,
d_out: Iterator,
num_items: int,
binary_op: Op,
h_init: Value,
stream,
) -> int: ...
def compute_exclusive(
self,
temp_storage_ptr: int | None,
temp_storage_nbytes: int,
d_in: Iterator,
d_out: Iterator,
num_items: int,
binary_op: Op,
h_init: Value,
stream,
) -> int: ...
def compute_inclusive_future_value(
self,
temp_storage_ptr: int | None,
temp_storage_nbytes: int,
d_in: Iterator,
d_out: Iterator,
num_items: int,
binary_op: Op,
h_init: Iterator,
stream,
) -> int: ...
def compute_exclusive_future_value(
self,
temp_storage_ptr: int | None,
temp_storage_nbytes: int,
d_in: Iterator,
d_out: Iterator,
num_items: int,
binary_op: Op,
h_init: Iterator,
stream,
) -> int: ...
def compute_inclusive_no_init(
self,
temp_storage_ptr: int | None,
temp_storage_nbytes: int,
d_in: Iterator,
d_out: Iterator,
num_items: int,
binary_op: Op,
h_init: None,
stream,
) -> int: ...
def serialize(self) -> bytes: ...
@staticmethod
def deserialize(
blob: bytes, load: bool = ..., check_cc: bool = ...
) -> DeviceScanBuildResult: ...
@staticmethod
def compile(*args) -> DeviceScanBuildResult: ...
def load(self) -> None: ...
# ---------------------
# DeviceSegmentedReduce
# ---------------------
class DeviceSegmentedReduceBuildResult:
def __init__(
self,
d_in: Iterator,
d_out: Iterator,
start_offsets: Iterator,
end_offsets: Iterator,
binary_op: Op,
h_init: Value,
info: CommonData,
): ...
def compute(
self,
temp_storage_ptr: int | None,
temp_storage_nbytes: int,
d_in: Iterator,
d_out: Iterator,
num_items: int,
start_offsets: Iterator,
end_offsets: Iterator,
binary_op: Op,
h_init: Value,
max_segment_size: int | None = None,
stream=None,
) -> int: ...
def serialize(self) -> bytes: ...
@staticmethod
def deserialize(
blob: bytes, load: bool = ..., check_cc: bool = ...
) -> DeviceSegmentedReduceBuildResult: ...
@staticmethod
def compile(*args) -> DeviceSegmentedReduceBuildResult: ...
def load(self) -> None: ...
# ---------------
# DeviceMergeSort
# ---------------
class DeviceMergeSortBuildResult:
def __init__(
self,
d_in_keys: Iterator,
d_in_items: Iterator,
d_out_keys: Iterator,
d_out_items: Iterator,
binary_op: Op,
info: CommonData,
): ...
def compute(
self,
temp_storage_ptr: int | None,
temp_storage_nbytes: int,
d_in_keys: Iterator,
d_in_items: Iterator,
d_out_keys: Iterator,
d_out_items: Iterator,
num_items: int,
binary_op: Op,
stream,
) -> int: ...
def serialize(self) -> bytes: ...
@staticmethod
def deserialize(
blob: bytes, load: bool = ..., check_cc: bool = ...
) -> DeviceMergeSortBuildResult: ...
@staticmethod
def compile(*args) -> DeviceMergeSortBuildResult: ...
def load(self) -> None: ...
# -----------------
# DeviceUniqueByKey
# -----------------
class DeviceUniqueByKeyBuildResult:
def __init__(
self,
d_keys_in: Iterator,
d_values_in: Iterator,
d_keys_out: Iterator,
d_values_out: Iterator,
d_num_selected_out: Iterator,
binary_op: Op,
info: CommonData,
): ...
def compute(
self,
temp_storage_ptr: int | None,
temp_storage_nbytes: int,
d_keys_in: Iterator,
d_values_in: Iterator,
d_keys_out: Iterator,
d_values_out: Iterator,
d_num_selected_out: Iterator,
binary_op: Op,
num_items: int,
stream,
) -> int: ...
def serialize(self) -> bytes: ...
@staticmethod
def deserialize(
blob: bytes, load: bool = ..., check_cc: bool = ...
) -> DeviceUniqueByKeyBuildResult: ...
@staticmethod
def compile(*args) -> DeviceUniqueByKeyBuildResult: ...
def load(self) -> None: ...
# -----------------
# DeviceRadixSort
# -----------------
class DeviceRadixSortBuildResult:
def __init__(self): ...
def compute(
self,
temp_storage_ptr: int | None,
temp_storage_nbytes: int,
d_keys_in: Iterator,
d_keys_out: Iterator,
d_values_in: Iterator,
d_values_out: Iterator,
decomposer_op: Op,
num_items: int,
begin_bit: int,
end_bit: int,
is_overwrite_okay: bool,
selector: int,
stream,
) -> tuple[int, int]: ...
def serialize(self) -> bytes: ...
@staticmethod
def deserialize(
blob: bytes, load: bool = ..., check_cc: bool = ...
) -> DeviceRadixSortBuildResult: ...
@staticmethod
def compile(*args) -> DeviceRadixSortBuildResult: ...
def load(self) -> None: ...
# --------------------
# DeviceUnaryTransform
# --------------------
class DeviceUnaryTransform:
def __init__(
self,
d_in: Iterator,
d_out: Iterator,
op: Op,
info: CommonData,
): ...
def compute(
self,
d_in: Iterator,
d_out: Iterator,
num_items: int,
stream,
) -> None: ...
def serialize(self) -> bytes: ...
@staticmethod
def deserialize(
blob: bytes, load: bool = ..., check_cc: bool = ...
) -> DeviceUnaryTransform: ...
@staticmethod
def compile(*args) -> DeviceUnaryTransform: ...
def load(self) -> None: ...
# ---------------------
# DeviceBinaryTransform
# ---------------------
class DeviceBinaryTransform:
def __init__(
self,
d_in1: Iterator,
d_in2: Iterator,
d_out: Iterator,
op: Op,
info: CommonData,
): ...
def compute(
self,
d_in1: Iterator,
d_in2: Iterator,
d_out: Iterator,
num_items: int,
stream,
) -> None: ...
def serialize(self) -> bytes: ...
@staticmethod
def deserialize(
blob: bytes, load: bool = ..., check_cc: bool = ...
) -> DeviceBinaryTransform: ...
@staticmethod
def compile(*args) -> DeviceBinaryTransform: ...
def load(self) -> None: ...
# ---------------
# DeviceHistogram
# ---------------
class DeviceHistogramBuildResult:
def __init__(
self,
num_channels: int,
num_active_channels: int,
d_samples: Iterator,
num_levels: int,
d_histogram: Iterator,
level_type: TypeInfo,
num_rows: int,
row_stride_samples: int,
is_evenly_segmented: bool,
info: CommonData,
): ...
def compute_even(
self,
d_samples: Iterator,
d_histogram: Iterator,
h_num_output_levels: Value,
h_lower_level: Value,
h_upper_level: Value,
num_row_pixels: int,
num_rows: int,
row_stride_samples: int,
stream,
) -> None: ...
def serialize(self) -> bytes: ...
@staticmethod
def deserialize(
blob: bytes, load: bool = ..., check_cc: bool = ...
) -> DeviceHistogramBuildResult: ...
@staticmethod
def compile(*args) -> DeviceHistogramBuildResult: ...
def load(self) -> None: ...
# -------------------
# DeviceBinarySearch
# -------------------
class DeviceBinarySearchBuildResult:
def __init__(
self,
mode: BinarySearchMode,
d_data: Iterator,
d_values: Iterator,
d_out: Iterator,
comparison_op: Op,
info: CommonData,
): ...
def compute(
self,
d_data: Iterator,
num_items: int,
d_values: Iterator,
num_values: int,
d_out: Iterator,
comparison_op: Op,
stream,
) -> None: ...
def serialize(self) -> bytes: ...
@staticmethod
def deserialize(
blob: bytes, load: bool = ..., check_cc: bool = ...
) -> DeviceBinarySearchBuildResult: ...
@staticmethod
def compile(*args) -> DeviceBinarySearchBuildResult: ...
def load(self) -> None: ...
# -----------------
# DeviceSegmentedSort
# -----------------
class DeviceSegmentedSortBuildResult:
def __init__(self): ...
def compute(
self,
temp_storage_ptr: int | None,
temp_storage_nbytes: int,
d_in_keys: Iterator,
d_out_keys: Iterator,
d_in_values: Iterator,
d_out_values: Iterator,
num_items: int,
num_segments: int,
d_begin_offsets: Iterator,
d_end_offsets: Iterator,
is_overwrite_okay: bool,
selector: int,
stream,
) -> tuple[int, int]: ...
def serialize(self) -> bytes: ...
@staticmethod
def deserialize(
blob: bytes, load: bool = ..., check_cc: bool = ...
) -> DeviceSegmentedSortBuildResult: ...
@staticmethod
def compile(*args) -> DeviceSegmentedSortBuildResult: ...
def load(self) -> None: ...
# ---------------------
# DeviceThreeWayPartition
# ---------------------
class DeviceThreeWayPartitionBuildResult:
def __init__(
self,
d_in: Iterator,
d_first_part_out: Iterator,
d_second_part_out: Iterator,
d_unselected_out: Iterator,
d_num_selected_out: Iterator,
select_first_part_op: Op,
select_second_part_op: Op,
info: CommonData,
): ...
def compute(
self,
d_in: Iterator,
d_first_part_out: Iterator,
d_second_part_out: Iterator,
d_unselected_out: Iterator,
d_num_selected_out: Iterator,
num_items: int,
stream,
) -> int: ...
def serialize(self) -> bytes: ...
@staticmethod
def deserialize(
blob: bytes, load: bool = ..., check_cc: bool = ...
) -> DeviceThreeWayPartitionBuildResult: ...
@staticmethod
def compile(*args) -> DeviceThreeWayPartitionBuildResult: ...
def load(self) -> None: ...

17
cccl_upstream/python/cuda_cccl/cuda/compute/_bindings_binary_search_backend_v1.pxi Normal file
View File

@@ -0,0 +1,17 @@
# v1 (cccl.c.parallel, NVRTC) — binary_search build_result_t struct +
# uniform cubin-bytes helper. v1 nests a transform build_result and carries
# op-state metadata; v2 (sibling file) flattens to top-level cubin fields.
cdef extern from "cccl/c/binary_search.h":
cdef struct cccl_device_binary_search_build_result_t 'cccl_device_binary_search_build_result_t':
cccl_device_transform_build_result_t transform
size_t op_state_size
size_t op_state_alignment
cdef inline bytes _binary_search_cubin_bytes(
cccl_device_binary_search_build_result_t* b,
):
return PyBytes_FromStringAndSize(
<const char*>b.transform.payload, b.transform.payload_size
)

15
cccl_upstream/python/cuda_cccl/cuda/compute/_bindings_binary_search_backend_v2.pxi Normal file
View File

@@ -0,0 +1,15 @@
# v2 (cccl.c.parallel.v2, HostJIT) — binary_search build_result_t struct +
# uniform cubin-bytes helper. v2 uses payload/payload_size matching v1.
cdef extern from "cccl/c/binary_search.h":
cdef struct cccl_device_binary_search_build_result_t 'cccl_device_binary_search_build_result_t':
void* payload
size_t payload_size
cdef inline bytes _binary_search_cubin_bytes(
cccl_device_binary_search_build_result_t* b,
):
return PyBytes_FromStringAndSize(
<const char*>b.payload, b.payload_size
)

2992
cccl_upstream/python/cuda_cccl/cuda/compute/_bindings_impl.pyx Normal file
View File

File diff suppressed because it is too large Load Diff

14
cccl_upstream/python/cuda_cccl/cuda/compute/_bindings_op_code_type_v1.pxi Normal file
View File

@@ -0,0 +1,14 @@
# v1 (cccl.c.parallel, NVRTC) — cccl_op_code_type enum + string-to-enum helper.
# Selected at CMake configure time and configure_file'd to the build dir as
# `_bindings_op_code_type.pxi`. v1's types.h does not define CCCL_OP_LLVM_IR.
cdef extern from "cccl/c/types.h":
cdef enum cccl_op_code_type:
CCCL_OP_LTOIR
CCCL_OP_CPP_SOURCE
cdef inline cccl_op_code_type _parse_code_type(str s) noexcept:
if s == "cpp_source":
return CCCL_OP_CPP_SOURCE
return CCCL_OP_LTOIR

17
cccl_upstream/python/cuda_cccl/cuda/compute/_bindings_op_code_type_v2.pxi Normal file
View File

@@ -0,0 +1,17 @@
# v2 (cccl.c.parallel.v2, HostJIT) — cccl_op_code_type enum + string-to-enum
# helper. Selected at CMake configure time and configure_file'd to the build
# dir as `_bindings_op_code_type.pxi`. v2's types.h adds CCCL_OP_LLVM_IR.
cdef extern from "cccl/c/types.h":
cdef enum cccl_op_code_type:
CCCL_OP_LTOIR
CCCL_OP_CPP_SOURCE
CCCL_OP_LLVM_IR
cdef inline cccl_op_code_type _parse_code_type(str s) noexcept:
if s == "llvm_ir":
return CCCL_OP_LLVM_IR
if s == "cpp_source":
return CCCL_OP_CPP_SOURCE
return CCCL_OP_LTOIR

40
cccl_upstream/python/cuda_cccl/cuda/compute/_bindings_segmented_reduce_backend_v1.pxi Normal file
View File

@@ -0,0 +1,40 @@
# v1 (cccl.c.parallel, NVRTC) — segmented_reduce extern + uniform call helper.
# Selected at CMake configure time and configure_file'd to the build dir as
# `_bindings_segmented_reduce_backend.pxi`. v1's signature takes
# `size_t max_segment_size` between `init` and `stream`.
cdef extern from "cccl/c/segmented_reduce.h":
cdef CUresult cccl_device_segmented_reduce(
cccl_device_segmented_reduce_build_result_t,
void *,
size_t *,
cccl_iterator_t,
cccl_iterator_t,
uint64_t,
cccl_iterator_t,
cccl_iterator_t,
cccl_op_t,
cccl_value_t,
size_t,
CUstream
) nogil
cdef inline CUresult _call_segmented_reduce(
cccl_device_segmented_reduce_build_result_t bld,
void* storage_ptr,
size_t* storage_sz,
cccl_iterator_t d_in,
cccl_iterator_t d_out,
uint64_t num_items,
cccl_iterator_t start_offsets,
cccl_iterator_t end_offsets,
cccl_op_t op_data,
cccl_value_t init,
size_t max_segment_size,
CUstream stream,
) nogil:
return cccl_device_segmented_reduce(
bld, storage_ptr, storage_sz, d_in, d_out, num_items,
start_offsets, end_offsets, op_data, init, max_segment_size, stream
)

38
cccl_upstream/python/cuda_cccl/cuda/compute/_bindings_segmented_reduce_backend_v2.pxi Normal file
View File

@@ -0,0 +1,38 @@
# v2 (cccl.c.parallel.v2, HostJIT) — segmented_reduce extern + uniform call
# helper. v2 dropped `size_t max_segment_size`; the helper accepts it for
# signature-compatibility with v1 and silently ignores it.
cdef extern from "cccl/c/segmented_reduce.h":
cdef CUresult cccl_device_segmented_reduce(
cccl_device_segmented_reduce_build_result_t,
void *,
size_t *,
cccl_iterator_t,
cccl_iterator_t,
uint64_t,
cccl_iterator_t,
cccl_iterator_t,
cccl_op_t,
cccl_value_t,
CUstream
) nogil
cdef inline CUresult _call_segmented_reduce(
cccl_device_segmented_reduce_build_result_t bld,
void* storage_ptr,
size_t* storage_sz,
cccl_iterator_t d_in,
cccl_iterator_t d_out,
uint64_t num_items,
cccl_iterator_t start_offsets,
cccl_iterator_t end_offsets,
cccl_op_t op_data,
cccl_value_t init,
size_t max_segment_size,
CUstream stream,
) nogil:
return cccl_device_segmented_reduce(
bld, storage_ptr, storage_sz, d_in, d_out, num_items,
start_offsets, end_offsets, op_data, init, stream
)

1363
cccl_upstream/python/cuda_cccl/cuda/compute/_bindings_serialization_v1.pxi Normal file
View File

File diff suppressed because it is too large Load Diff

57
cccl_upstream/python/cuda_cccl/cuda/compute/_bindings_serialization_v2.pxi Normal file
View File

@@ -0,0 +1,57 @@
# v2 (HostJIT) backend — serialize/deserialize/compile/load not supported.
# Included at the end of _bindings_impl.pyx; provides stub functions so that
# the class methods exist but raise a clear error if called.
_NOT_SUPPORTED_MSG = (
"serialize/deserialize (and ahead-of-time compile/load) is not supported "
"with the HostJIT (v2) backend."
)
def _reduce_serialize(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _reduce_deserialize(blob, load=True, check_cc=True): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _reduce_compile(*args): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _reduce_load(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _scan_serialize(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _scan_deserialize(blob, load=True, check_cc=True): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _scan_compile(*args): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _scan_load(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _segmented_reduce_serialize(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _segmented_reduce_deserialize(blob, load=True, check_cc=True): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _segmented_reduce_compile(*args): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _segmented_reduce_load(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _merge_sort_serialize(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _merge_sort_deserialize(blob, load=True, check_cc=True): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _merge_sort_compile(*args): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _merge_sort_load(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _unique_by_key_serialize(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _unique_by_key_deserialize(blob, load=True, check_cc=True): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _unique_by_key_compile(*args): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _unique_by_key_load(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _radix_sort_serialize(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _radix_sort_deserialize(blob, load=True, check_cc=True): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _radix_sort_compile(*args): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _radix_sort_load(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _unary_transform_serialize(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _unary_transform_deserialize(blob, load=True, check_cc=True): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _unary_transform_compile(*args): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _unary_transform_load(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _binary_transform_serialize(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _binary_transform_deserialize(blob, load=True, check_cc=True): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _binary_transform_compile(*args): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _binary_transform_load(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _histogram_serialize(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _histogram_deserialize(blob, load=True, check_cc=True): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _histogram_compile(*args): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _histogram_load(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _binary_search_serialize(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _binary_search_deserialize(blob, load=True, check_cc=True): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _binary_search_compile(*args): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _binary_search_load(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _three_way_partition_serialize(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _three_way_partition_deserialize(blob, load=True, check_cc=True): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _three_way_partition_compile(*args): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _three_way_partition_load(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _segmented_sort_serialize(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _segmented_sort_deserialize(blob, load=True, check_cc=True): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _segmented_sort_compile(*args): raise NotImplementedError(_NOT_SUPPORTED_MSG)
def _segmented_sort_load(self): raise NotImplementedError(_NOT_SUPPORTED_MSG)

740
cccl_upstream/python/cuda_cccl/cuda/compute/_caching.py Normal file
View File

@@ -0,0 +1,740 @@
# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
import functools
import threading
import types
import weakref
from typing import Any, Callable, Hashable, NamedTuple, TypeVar
import numpy as np
from cuda.core import Device
try:
from cuda.core._utils.cuda_utils import CUDAError
except ImportError:
from cuda.core.experimental._utils.cuda_utils import CUDAError
from ._utils.protocols import get_dtype, get_shape, is_device_array
from .struct import _Struct
try:
from ._build_info import USING_V2 # type: ignore[import-not-found]
except ImportError:
USING_V2 = False
# Whether the backend can serialize build results (the v2 HostJIT backend
# cannot, today). Serialization is what lets same-cc devices share one default
# build — a non-owner device loads by cloning the shared entry through
# serialize -> deserialize -> load — so backends without it key default builds
# per device ordinal and build independently per device instead.
# TODO: delete this flag (and every branch on it) once v2 supports
# build-result serialization.
_BACKEND_SERIALIZES_BUILD_RESULTS = not USING_V2
# Registry thet maps type -> key function for extracting cache key
# from a value of that type.
_KEY_FUNCTIONS: dict[type, Callable[[Any], Hashable]] = {}
def _type_fqn(v):
# fully-qualified type name to distinguish np.ndarray from cp.ndarray from GpuStruct
return f"{type(v).__module__}.{type(v).__name__}"
def _key_for(value: Any) -> Hashable:
"""
Extract a cache key from a value using the registered KEY_FUNCTIONS.
This function checks the type of the value and delegates to the
appropriate registered keyer. Falls back to using the value
directly if no keyer is registered.
Args:
value: The value to extract a cache key from
Returns:
A hashable cache key
"""
# Handle sequences (lists, tuples) by recursively converting to tuple
if isinstance(value, (list, tuple)):
return tuple(_key_for(item) for item in value)
# Check for exact type match first
value_type = type(value)
if value_type in _KEY_FUNCTIONS:
return _KEY_FUNCTIONS[value_type](value)
# DeviceArrayLike is not a runtime-checkable protocol, so
# we cannot isinstance() with it.
if is_device_array(value):
return (_type_fqn(value), get_dtype(value))
# Check for instance match (handles inheritance)
for registered_type, keyer in _KEY_FUNCTIONS.items():
if isinstance(value, registered_type):
return keyer(value)
# Fallback: use value directly (assumes it's hashable)
return value
# The specialization part of every cache key: what _make_cache_key_from_args
# extracts from the user-facing arguments (dtypes, op identity, iterator
# kinds, ...).
_SpecializationKey = tuple[Hashable, ...]
def _make_cache_key_from_args(*args, **kwargs) -> _SpecializationKey:
"""
Create a cache key from function arguments.
Args:
*args: Positional arguments
**kwargs: Keyword arguments
Returns:
A tuple containing the extracted cache keys
"""
positional_keys = tuple(_key_for(arg) for arg in args)
# Sort kwargs by key name for consistent ordering
if kwargs:
sorted_kwargs = sorted(kwargs.items())
kwarg_keys = tuple((k, _key_for(v)) for k, v in sorted_kwargs)
return positional_keys + (kwarg_keys,)
return positional_keys
# Process-wide registry of all algorithm caches.
_process_wide_cache_registry: dict[str, object] = {}
class _ThreadLocalCaches:
"""
Container for wrapper caches owned by a single Python thread.
Each thread gets its own instance via ``threading.local()``. We use
``__weakref__`` to enable the process-wide registry of caches to hold weak
references to the thread's caches. That way, if a thread exits, its caches
will be garbage collected and removed from the registry even if the
process-wide registry still references them.
"""
__slots__ = ("wrapper_caches", "__weakref__")
def __init__(self) -> None:
# Outer key: decorated algorithm factory name, e.g.,
# make_reduce_into.__qualname__. Inner key: _WrapperCacheKey. No
# thread id is needed in either key because each thread holds a
# separate thread local object of this class.
self.wrapper_caches: dict[str, dict[_WrapperCacheKey, Any]] = {}
class _InFlightBuild:
"""
Coordination state for one shared build-result currently being built.
The first thread for a cache key runs the builder. Other threads wait on
``event`` and receive either the completed build result or the builder's
exception.
"""
def __init__(self) -> None:
self.event = threading.Event()
self.result: Any = None
self.exception: BaseException | None = None
# A compute capability packed into one int as major * 10 + minor, e.g.
# (9, 0) -> 90 and (12, 0) -> 120; cc_to_key / key_to_cc in _cccl_interop
# convert to and from the (major, minor) pair form. Purely documentary:
# mypy treats it as int.
_PackedCCKey = int
class _DeviceBuildTarget(NamedTuple):
"""
Target identity of a per-device wrapper or build entry.
Two roles. In the wrapper cache it keys every default-build wrapper:
wrappers hold device-bound state (their construction-time binding), so
each device gets its own. In the build-results cache it is used only when
the backend cannot serialize build results (the v2 HostJIT backend
today): sharing one entry across same-cc devices requires cloning it
through serialization, so such backends key default builds per device
ordinal instead.
TODO: once v2 supports build-result serialization, delete the build-cache
role (the _BACKEND_SERIALIZES_BUILD_RESULTS branch in
cache_build_results); the wrapper-cache role remains.
NamedTuples compare as plain tuples, so all target kinds must keep
structurally disjoint layouts (arity or element types) to never compare
equal to each other.
"""
device_id: int
cc: tuple[int, int]
class _DefaultBuildTarget(NamedTuple):
"""
Target identity of a default build shared across same-cc devices.
Holds the packed cc alone: the compiled payload depends only on the cc,
and _PerCCBuildResults.resolve() gives each device its own loaded state.
See _DeviceBuildTarget for the cross-kind equality constraint.
"""
cc_key: _PackedCCKey
class _AOTBuildTarget(NamedTuple):
"""
Target identity of an explicit AOT build, with no device attached.
Holds the normalized, sorted, packed compute-capability keys. See
_DeviceBuildTarget for the cross-kind equality constraint.
"""
cc_keys: tuple[_PackedCCKey, ...]
# Inner key of a thread's per-factory wrapper cache; see _ThreadLocalCaches.
# The target is None for explicit AOT builds: the decorator performs no device
# query or cc normalization there, and the raw compute_capability kwarg is
# already part of the specialization. Differently spelled ccs (80 vs (8, 0))
# therefore yield distinct wrappers, which still share one compiled build
# because cache_build_results normalizes its own key.
_WrapperCacheKey = tuple[_DeviceBuildTarget | None, _SpecializationKey]
# Composite key of the process-wide build-results cache:
# (build-result type, build target (device or AOT), specialization). The
# build-result type is the algorithm's Cython class from _bindings (e.g.
# DeviceReduceBuildResult), namespacing entries per algorithm.
# TODO: drop _DeviceBuildTarget from this union once v2 supports build-result
# serialization; it then only keys the wrapper cache.
_BuildResultsCacheKey = tuple[
type,
_DefaultBuildTarget | _DeviceBuildTarget | _AOTBuildTarget,
_SpecializationKey,
]
_thread_local = threading.local()
# Process wide registry of per-thread caches. It enables a thread to call
# clear_all_caches() to clear all caches across all threads.
_process_wide_thread_cache_registry: weakref.WeakSet[_ThreadLocalCaches] = (
weakref.WeakSet()
)
_process_wide_thread_cache_registry_lock = threading.Lock()
# _InFlightBuild entries are temporary: replaced by the completed build
# results or removed on builder failure.
_process_wide_build_results_cache: dict[
_BuildResultsCacheKey, _PerCCBuildResults | _InFlightBuild
] = {}
_CACHE_MISS = object()
_KeyT = TypeVar("_KeyT", bound=Hashable)
def _cache_single_flight(
cache: dict[_KeyT, Any], cache_key: _KeyT, builder: Callable[[], Any]
) -> Any:
"""Return a cached value, coalescing concurrent builds for the same key.
``cache`` may be any single-flight dict — currently the process-wide
build-results cache and the per-device loaded results inside each
_PerCCBuildResults. Its entries are one of two things: a completed value,
which is terminal, or a temporary _InFlightBuild while the one elected
caller runs ``builder``; other callers wait on its event and receive the
same result or exception. A failed builder's entry is removed so a later
call retries.
"""
cache_entry = cache.get(cache_key, _CACHE_MISS)
if cache_entry is _CACHE_MISS:
in_flight = _InFlightBuild()
# setdefault elects one builder without an explicit lock on cache hits.
cache_entry = cache.setdefault(cache_key, in_flight)
if cache_entry is in_flight:
try:
result = builder()
in_flight.result = result
cache[cache_key] = result
except BaseException as exc:
in_flight.exception = exc
cache.pop(cache_key, None)
raise
finally:
in_flight.event.set()
return result
if isinstance(cache_entry, _InFlightBuild):
cache_entry.event.wait()
if cache_entry.exception is not None:
raise cache_entry.exception
return cache_entry.result
return cache_entry
class _PerCCBuildResults(dict[_PackedCCKey, Any]):
"""One algorithm specialization's compiled build results, keyed by target cc.
Instances may be shared process-wide across threads: the factory build
cache hands all same-specialization wrappers one instance. The compiled
payload depends only on the cc, but a loaded build result holds
device-specific native state, so devices never share one. The first
device to execute claims and loads the canonical build result in place;
each additional device lazily loads its own clone of the compiled payload
(serialize -> deserialize -> load). This class tracks the loaded result —
canonical or clone — assigned to each device.
"""
def __init__(
self,
build_results: dict[_PackedCCKey, Any],
*,
loaded_device_id: int | None = None,
) -> None:
super().__init__(build_results)
# The single device that claimed each cc's canonical result and loads
# it in place; every other device clones instead. The atomic claim
# (resolve()'s setdefault) is what prevents double-loading the
# canonical object.
self._owner_devices: dict[_PackedCCKey, int] = {}
# The loaded result each (cc key, device ordinal) pair executes: the
# canonical result for its owner device; an independent clone — or,
# when cloning fails, an independently built result — for every other
# device. Also holds temporary _InFlightBuild entries while a first
# load is in flight.
self._loaded_results: dict[tuple[_PackedCCKey, int], Any] = {}
# Loading the canonical result mutates its native handle fields, while
# cloning serializes its payload. Serialize those source operations, but
# keep completed-result lookups lock-free.
self._source_locks = {cc: threading.Lock() for cc in self}
if loaded_device_id is not None:
# The caller already built and loaded the single entry on
# loaded_device_id: pre-record the post-conditions resolve()'s
# owner-load path would otherwise produce on first use.
if len(self) != 1:
raise ValueError("A device-bound _PerCCBuildResults must be singular")
for cc, build_result in self.items():
self._owner_devices[cc] = loaded_device_id
self._loaded_results[(cc, loaded_device_id)] = build_result
def resolve(self, cc: _PackedCCKey, device_id: int) -> Any:
"""Return the build result loaded for ``device_id`` without recompiling."""
# Completed loads are terminal — never removed or replaced — so the
# warm path is one lock-free lookup with no dict mutation and no
# closure allocation. Misses and in-flight loads (which only exist
# around the first load per device) fall through to the single-flight
# machinery below.
loaded = self._loaded_results.get((cc, device_id))
if loaded is not None and not isinstance(loaded, _InFlightBuild):
return loaded
source = self[cc]
owner_device = self._owner_devices.setdefault(cc, device_id)
def load_for_device():
if owner_device == device_id:
with self._source_locks[cc]:
source.load()
return source
with self._source_locks[cc]:
blob = source.serialize()
result = type(source).deserialize(blob, load=False, check_cc=True)
result.load()
return result
return _cache_single_flight(
self._loaded_results, (cc, device_id), load_for_device
)
def serialize_build_result(self, cc: _PackedCCKey) -> bytes:
"""Serialize a canonical result without racing its first device load."""
with self._source_locks[cc]:
return self[cc].serialize()
def _get_current_device_info() -> tuple[int, tuple[int, int]]:
device = Device()
cc_major, cc_minor = device.compute_capability
return device.device_id, (cc_major, cc_minor)
def _get_thread_caches() -> _ThreadLocalCaches:
caches = getattr(_thread_local, "caches", None)
if caches is None:
caches = _ThreadLocalCaches()
_thread_local.caches = caches
with _process_wide_thread_cache_registry_lock:
_process_wide_thread_cache_registry.add(caches)
return caches
def _clear_wrapper_caches(cache_name: str | None = None) -> None:
with _process_wide_thread_cache_registry_lock:
thread_caches = list(_process_wide_thread_cache_registry)
for caches in thread_caches:
if cache_name is None:
caches.wrapper_caches.clear()
else:
caches.wrapper_caches.pop(cache_name, None)
def cache_build_results(
build_result_type: type,
*key_args,
compute_capability,
builder: Callable[[], Any],
) -> Any:
"""
Cache the shared Cython build results for one specialization.
Current-device builds are keyed by compute capability alone and shared
across same-cc device ordinals: the compiled payload only depends on the
cc, and _PerCCBuildResults.resolve() gives each device its own loaded
state. When the backend cannot serialize build results (the v2 HostJIT
backend today), a non-owner device has no way to load the shared entry —
loading it clones through serialization — so default builds are keyed per
device ordinal instead and each device builds its own entry. Explicit AOT
builds have no current device and are keyed by their normalized target
compute capabilities. The key intentionally excludes the current Python
thread so wrappers can share compiled results.
Args:
build_result_type: Cython build-result type. This separates entries
that may otherwise have identical specialization keys.
*key_args: Positional values used to form the specialization part of
the cache key.
compute_capability: Explicit AOT target or ``None`` for the current
device.
builder: Callable that creates the _PerCCBuildResults on a cache miss.
Exactly one thread runs this callable for a given key while other
threads wait for the result.
Returns:
``(build_results, bound_result)``: the cached or newly built
_PerCCBuildResults and, for current-device builds, the loaded result
bound to the constructing device — resolved once here so ``__call__``
needs no device query. ``bound_result`` is ``None`` for explicit AOT
builds, which resolve per call.
"""
from ._cccl_interop import cc_to_key, normalize_compute_capabilities
if compute_capability is None:
# The factory decorator already queried the device on the wrapper-cache
# miss path and hands the result through thread-local state; fall back
# to a fresh query for direct construction (e.g. deserialization).
device_info = getattr(_thread_local, "factory_device_info", None)
if device_info is None:
device_info = _get_current_device_info()
device_id, cc = device_info
packed_cc = cc_to_key(cc)
# TODO: reduce to _DefaultBuildTarget(packed_cc) once v2 supports
# build-result serialization.
target_key = (
_DefaultBuildTarget(packed_cc)
if _BACKEND_SERIALIZES_BUILD_RESULTS
else _DeviceBuildTarget(device_id, cc)
)
user_cache_key = _make_cache_key_from_args(*key_args)
cache_key = (build_result_type, target_key, user_cache_key)
build_results = _cache_single_flight(
_process_wide_build_results_cache, cache_key, builder
)
return build_results, _bind_default_build(
build_results, packed_cc, device_id, builder
)
target_ccs = normalize_compute_capabilities(compute_capability)
assert target_ccs is not None
aot_target_key = _AOTBuildTarget(tuple(cc_to_key(cc) for cc in target_ccs))
user_cache_key = _make_cache_key_from_args(*key_args)
aot_cache_key = (build_result_type, aot_target_key, user_cache_key)
return (
_cache_single_flight(_process_wide_build_results_cache, aot_cache_key, builder),
None,
)
def _bind_default_build(
build_results, packed_cc: _PackedCCKey, device_id: int, builder
):
"""Resolve the loaded result the constructing device executes.
Default wrappers are bound to the device that was current at factory-call
time (their wrapper-cache key includes it), so the binding is resolved
once here and reused by every ``__call__`` with no device query. For the
device that built the shared entry this is a warm lookup; another same-cc
device loads its own clone of the compiled payload here (serialize ->
deserialize -> load, milliseconds where a build costs a second).
Cloning is only an optimization and its failures have no stable exception
type, so any resolve failure falls back to a full build for this device —
recorded in the shared per-device slot so same-device threads share it —
and genuine errors surface from the build path itself.
"""
try:
return build_results.resolve(packed_cc, device_id)
except Exception:
def build_privately():
(result,) = builder().values()
return result
return _cache_single_flight(
build_results._loaded_results, (packed_cc, device_id), build_privately
)
class _CacheWithRegisteredKeyFunctions:
"""
Decorator to cache the result of the decorated function.
The cache key is automatically computed from the decorated function's
arguments using the registered key functions.
"""
def __call__(self, func: Callable) -> Callable:
"""
Decorator to cache the result of the decorated function.
Args:
func: The function whose result is to be cached.
Notes
-----
Default builds append the current CUDA device and compute capability to
the cache key. Explicit AOT builds include their normalized target
compute capabilities without querying a device.
"""
cache_name = func.__qualname__
@functools.wraps(func)
def inner(*args, **kwargs):
user_cache_key = _make_cache_key_from_args(*args, **kwargs)
# When the caller targets explicit compute capabilities, that value
# is already part of user_cache_key (it arrives as a kwarg) and we
# must NOT query a device — the whole point is to build without a
# GPU. Otherwise, salt the key with the current device's cc so a
# build cached on one device isn't reused on another.
if kwargs.get("compute_capability") is None:
# Only device-availability failures should be reinterpreted as
# "pass compute_capability": no driver / no device raises
# CUDAError, and querying device 0 on a machine with zero
# devices raises ValueError. Anything else (a real bug) must
# propagate untouched. The original error is chained and echoed
# so a genuine driver/permission failure isn't hidden behind a
# misleading "no device" message.
try:
device_id, cc = _get_current_device_info()
except (CUDAError, ValueError) as e:
raise RuntimeError(
"make_<algo> was called without compute_capability and the "
f"current CUDA device could not be queried ({e}). Pass "
"compute_capability=<cc or list of ccs> to compile without "
"a GPU (e.g. with ProxyArray / ProxyValue)."
) from e
target = _DeviceBuildTarget(device_id, cc)
target_cc_arg = cc
else:
target = None
target_cc_arg = kwargs.get("compute_capability")
# No thread id in the key: the containing cache is threading.local,
# so each thread only ever sees its own entries.
cache_key = (target, user_cache_key)
thread_caches = _get_thread_caches()
cache = thread_caches.wrapper_caches.setdefault(cache_name, {})
if cache_key not in cache:
# Shared device code (operators, iterators) is compiled to LTO-IR
# once and linked into every per-arch build result, so it must target
# the lowest requested cc (nvJitLink requires final SM >= each
# linked input's arch). Set that target around the build.
from ._target_cc import target_cc
# Hand the device info queried above to cache_build_results
# (reached through the wrapper's __init__) so the miss path
# does not construct a second cuda.core Device. Saved/restored
# so nested factory calls fall back to their own query.
previous_device_info = getattr(
_thread_local, "factory_device_info", None
)
_thread_local.factory_device_info = target
try:
with target_cc(target_cc_arg):
result = func(*args, **kwargs)
finally:
_thread_local.factory_device_info = previous_device_info
cache[cache_key] = result
return cache[cache_key]
inner.cache_clear = lambda: _clear_wrapper_caches(cache_name) # type: ignore[attr-defined]
# Register the cache in the central registry
_process_wide_cache_registry[func.__qualname__] = inner
return inner
def register(self, type_: type, key_function: Callable[[Any], Hashable]) -> None:
"""
Register a key function for a specific type.
A key function extracts a hashable cache key from a value.
Args:
type_: The type to register
key_function: A callable that takes an instance of type_ and
returns a hashable cache key
"""
_KEY_FUNCTIONS[type_] = key_function
cache_with_registered_key_functions = _CacheWithRegisteredKeyFunctions()
def _make_hashable(value):
# duck-type check for numba.cuda.CUDADispatcher:
if hasattr(value, "py_func") and callable(value.py_func):
return CachableFunction(value.py_func)
elif is_device_array(value):
# Ops with device arrays in globals/closures will be handled
# by stateful op machinery, which enables updating the state
# (pointers). Thus, we only cache on the dtype and shape of
# the referenced array, but not its pointer.
return (get_dtype(value), get_shape(value))
elif isinstance(value, (np.number, np.bool_)):
return ("numpy.scalar", value.dtype.str, value.tobytes())
elif isinstance(value, (bool, int, float)):
# Python scalars are immutable values; key them by type and value so
# equal-valued scalars share a cache entry. Without this they fall
# through to ``id(value)`` below, and a fresh (non-interned) ``int``/
# ``float`` with the same value misses the build cache on every call.
# ``_type_fqn`` keeps ``True`` distinct from ``1``/``1.0`` (and avoids
# collisions between like-named scalar subclasses from other modules).
return ("python.scalar", _type_fqn(value), value)
elif isinstance(value, (list, tuple)):
return tuple(_make_hashable(v) for v in value)
elif isinstance(value, dict):
return tuple(
sorted((_make_hashable(k), _make_hashable(v)) for k, v in value.items())
)
else:
return id(value)
def clear_all_caches():
"""
Clear all algorithm caches.
This function clears cached algorithm wrappers and completed build results
in the current process, forcing recompilation on the next invocation.
Useful for benchmarking compilation time.
This function is not synchronized with active factory calls or algorithm
execution. Callers that use it in a multi-threaded program must externally
synchronize with all threads that may create or use cuda.compute algorithm
objects. If a build is already in progress, that build may complete after
this function returns and repopulate the completed build-result cache.
Example
-------
>>> import cuda.compute
>>> cuda.compute.clear_all_caches()
"""
_clear_wrapper_caches()
_process_wide_build_results_cache.clear()
# Auxiliary caches registered process-wide (e.g. _jit._infer_return_type)
# must be cleared too, so builds after a clear really are cold. Factory
# entries' cache_clear is idempotent with _clear_wrapper_caches above.
for cached_func in _process_wide_cache_registry.values():
cached_func.cache_clear()
class CachableFunction:
"""
A type that wraps a function and provides custom comparison
(__eq__) and hash (__hash__) implementations.
The purpose of this class is to enable caching and comparison of
functions based on their bytecode, constants, and closures, while
ignoring other attributes such as their names or docstrings.
"""
# TODO: eventually, move this class to _jit.py as it only
# has to do with caching of Python callables that will be
# JIT compiled.
def __init__(self, func):
self._func = func
closure = func.__closure__ if func.__closure__ is not None else []
contents = []
# Make closure contents hashable
for cell in closure:
contents.append(_make_hashable(cell.cell_contents))
self._identity = (
func.__name__,
func.__code__.co_code,
func.__code__.co_consts,
tuple(contents),
tuple(
# if `name` is found in __globals__, try and hash
# the referenced object. If `name` is not found in
# __globals__, (e.g., `name` is part of a dotted
# name like `np.argmax`), for caching purposes we
# use the hash of the name itself. Assumes numba
# known how to interpret the dotted name at JIT
# time.
_make_hashable(func.__globals__.get(name, name))
for name in func.__code__.co_names
),
)
def __eq__(self, other):
return self._identity == other._identity
def __hash__(self):
return hash(self._identity)
def __repr__(self):
return str(self._func)
# Register keyers for built-in types
cache_with_registered_key_functions.register(
np.ndarray, lambda arr: ("numpy.ndarray", arr.dtype)
)
cache_with_registered_key_functions.register(
types.FunctionType, lambda fn: CachableFunction(fn)
)
cache_with_registered_key_functions.register(_Struct, lambda v: (_type_fqn(v), v.dtype))
def _register_proxy_types():
# Registered lazily to avoid importing _proxy (and numpy-dtype construction)
# at module import time; the keys are dtype-only so equal-dtype proxies share
# a cache entry.
from ._proxy import ProxyArray, ProxyValue
cache_with_registered_key_functions.register(
ProxyArray, lambda v: ("ProxyArray", v.dtype)
)
cache_with_registered_key_functions.register(
ProxyValue, lambda v: ("ProxyValue", v.dtype)
)
_register_proxy_types()

471
cccl_upstream/python/cuda_cccl/cuda/compute/_cccl_interop.py Normal file
View File

@@ -0,0 +1,471 @@
# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
import enum
import functools
import os
import subprocess
import tempfile
import warnings
from typing import Callable, List
try:
from cuda.core import Device as CudaDevice
except ImportError:
from cuda.core.experimental import Device as CudaDevice
import numpy as np
# TODO: adding a type-ignore here because `cuda` being a
# namespace package confuses mypy when `cuda.<something_else>`
# is installed, but not `cuda.cccl`. For namespace packages,
# it appears we need to actually install the sub-package
# in order for mypy to find its py.typed file. However, CI
# does type checking of `cuda.cccl` without actually installing
# it.
#
# We need to find a better solution for this.
from cuda.cccl import get_include_paths # type: ignore
from . import types
from ._bindings import (
CommonData,
Iterator,
IteratorKind,
IteratorState,
Op,
OpKind,
Pointer,
TypeEnum,
TypeInfo,
Value,
make_pointer_object,
)
from ._caching import _PerCCBuildResults
from ._utils.protocols import get_data_pointer, get_dtype, is_contiguous
from .iterators._base import IteratorBase
from .typing import DeviceArrayLike, GpuStruct
# Mapping from numpy dtype to TypeEnum for creating TypeInfo
_NUMPY_DTYPE_TO_ENUM = {
np.dtype("int8"): TypeEnum.INT8,
np.dtype("int16"): TypeEnum.INT16,
np.dtype("int32"): TypeEnum.INT32,
np.dtype("int64"): TypeEnum.INT64,
np.dtype("uint8"): TypeEnum.UINT8,
np.dtype("uint16"): TypeEnum.UINT16,
np.dtype("uint32"): TypeEnum.UINT32,
np.dtype("uint64"): TypeEnum.UINT64,
np.dtype("float16"): TypeEnum.FLOAT16,
np.dtype("float32"): TypeEnum.FLOAT32,
np.dtype("float64"): TypeEnum.FLOAT64,
np.dtype("bool"): TypeEnum.BOOLEAN,
}
@functools.lru_cache(maxsize=256)
def _type_info_from_dtype(dtype: np.dtype) -> TypeInfo:
"""
Create a TypeInfo from a numpy dtype.
Handles both primitive types and structured dtypes.
"""
dtype = np.dtype(dtype)
# Handle structured dtypes
if dtype.type == np.void and dtype.fields is not None:
return TypeInfo(dtype.itemsize, dtype.alignment, TypeEnum.STORAGE)
if dtype.kind == "c":
return TypeInfo(dtype.itemsize, dtype.alignment, TypeEnum.STORAGE)
# Fallback for any other type
type_enum = _NUMPY_DTYPE_TO_ENUM.get(dtype, TypeEnum.STORAGE)
return TypeInfo(dtype.itemsize, dtype.alignment, type_enum)
def _is_well_known_op(op: OpKind) -> bool:
return isinstance(op, OpKind) and op not in (OpKind.STATELESS, OpKind.STATEFUL)
def _device_array_to_cccl_iter(array: DeviceArrayLike) -> Iterator:
from ._proxy import ProxyArray
if not is_contiguous(array):
raise ValueError("Non-contiguous arrays are not supported.")
dtype = get_dtype(array)
info = _type_info_from_dtype(dtype)
state_info = _type_info_from_dtype(np.intp)
# A ProxyArray has no GPU allocation: leave the pointer NULL for build-time
# (ahead-of-time) compilation. The real pointer is bound at __call__ via
# set_cccl_iterator_state().
state = None if isinstance(array, ProxyArray) else get_data_pointer(array)
return Iterator(
state_info.alignment,
IteratorKind.POINTER,
Op(),
Op(),
info,
# Note: this is slightly slower, but supports all ndarray-like objects
# as long as they support CAI
# TODO: switch to use gpumemoryview once it's ready
state=state,
)
def _none_to_cccl_iter() -> Iterator:
# Any type could be used here, we just need to pass NULL.
info = _type_info_from_dtype(np.uint8)
return Iterator(info.alignment, IteratorKind.POINTER, Op(), Op(), info, state=None)
class _IteratorIO(enum.Enum):
INPUT = 0
OUTPUT = 1
def _to_cccl_iter(
it: DeviceArrayLike | IteratorBase | None, io_kind: _IteratorIO
) -> Iterator:
if it is None:
return _none_to_cccl_iter()
if isinstance(it, IteratorBase):
return it.to_cccl_iter(io_kind == _IteratorIO.OUTPUT)
return _device_array_to_cccl_iter(it)
def to_cccl_input_iter(array_or_iterator) -> Iterator:
return _to_cccl_iter(array_or_iterator, _IteratorIO.INPUT)
def to_cccl_output_iter(array_or_iterator) -> Iterator:
return _to_cccl_iter(array_or_iterator, _IteratorIO.OUTPUT)
def to_cccl_value_state(array_or_struct: np.ndarray | GpuStruct) -> memoryview:
from ._proxy import _PROXY_VALUE_DATA_ERROR, ProxyValue
if isinstance(array_or_struct, ProxyValue):
# Reached only if a proxy leaks into an execute call — proxies describe
# types for build, they carry no data to run with.
raise RuntimeError(_PROXY_VALUE_DATA_ERROR)
if isinstance(array_or_struct, np.ndarray):
assert array_or_struct.flags.contiguous
data = array_or_struct.data.cast("B")
return data
else:
# it's a GpuStruct, use the array underlying it
return to_cccl_value_state(array_or_struct._data)
def to_cccl_value(array_or_struct: np.ndarray | GpuStruct) -> Value:
from ._proxy import ProxyValue
if isinstance(array_or_struct, ProxyValue):
# Build-time placeholder: describe the type with a correctly sized zero
# buffer. The real value bytes are bound at __call__ via
# set_cccl_value_state().
info = _type_info_from_dtype(array_or_struct.dtype)
zero_bytes = memoryview(bytearray(array_or_struct.dtype.itemsize))
return Value(info, zero_bytes)
if isinstance(array_or_struct, np.ndarray):
info = _type_info_from_dtype(array_or_struct.dtype)
return Value(info, array_or_struct.data.cast("B"))
else:
# it's a GpuStruct, use the array underlying it
return to_cccl_value(array_or_struct._data)
def set_cccl_value_state(cccl_value: Value, array_or_struct: np.ndarray | GpuStruct):
"""
Set the state of a CCCL Value object from a numpy array or GpuStruct.
Args:
cccl_value: The CCCL Value binding object
array_or_struct: The numpy array or GpuStruct to get the state from
"""
cccl_value.state = to_cccl_value_state(array_or_struct)
def get_value_type(
d_in: DeviceArrayLike | IteratorBase | GpuStruct | np.ndarray,
):
from ._proxy import ProxyValue
from .struct import _Struct
if isinstance(d_in, IteratorBase):
return d_in.value_type
if isinstance(d_in, ProxyValue):
return types.from_numpy_dtype(d_in.dtype)
if isinstance(d_in, _Struct):
return type(d_in)._type_descriptor # type: ignore[union-attr]
dtype = get_dtype(d_in)
if dtype.type == np.void:
return types.from_numpy_dtype(dtype)
return types.from_numpy_dtype(dtype)
def set_cccl_iterator_state(cccl_it: Iterator, input_it):
if cccl_it.is_kind_pointer():
ptr = get_data_pointer(input_it)
ptr_obj = make_pointer_object(ptr, input_it)
cccl_it.state = ptr_obj
else:
state_ = input_it.state
if isinstance(state_, (IteratorState, Pointer)):
cccl_it.state = state_
else:
cccl_it.state = make_pointer_object(state_, input_it)
@functools.lru_cache()
def get_includes() -> List[str]:
def as_option(p):
if p is None:
return ""
return f"-I{p}"
paths = get_include_paths().as_tuple()
opts = [as_option(path) for path in paths]
return opts
def _check_compile_result(cubin: bytes):
# check compiled code for LDL/STL instructions
temp_cubin_file = tempfile.NamedTemporaryFile(delete=False)
try:
temp_cubin_file.write(cubin)
out = subprocess.run(
["nvdisasm", "-gi", temp_cubin_file.name], capture_output=True
)
if out.returncode != 0:
raise RuntimeError("nvdisasm failed")
sass = out.stdout.decode("utf-8")
except FileNotFoundError:
sass = "nvdiasm not found, skipping SASS validation"
warnings.warn(sass)
assert "LDL" not in sass, "LDL instruction found in SASS"
assert "STL" not in sass, "STL instruction found in SASS"
return temp_cubin_file.name
# this global variable controls whether the compile result is checked
# for LDL/STL instructions. Should be set to `True` for testing only.
_check_sass: bool = False
def _common_data_for_cc(cc):
"""Build a ``CommonData`` for a given compute capability.
``cc`` is a ``(major, minor)`` pair. When ``None``, the current device's
compute capability is queried (requires a live GPU).
"""
if cc is None:
cc_major, cc_minor = CudaDevice().compute_capability
else:
cc_major, cc_minor = cc
cub_path, thrust_path, libcudacxx_path, cuda_include_path = get_includes()
return CommonData(
cc_major, cc_minor, cub_path, thrust_path, libcudacxx_path, cuda_include_path
)
def call_build(build_impl_fn: Callable, *args, cc=None, **kwargs):
"""Build (compile + load) via ``build_impl_fn``, supplying compute capability and paths.
``cc`` is an optional ``(major, minor)`` pair; when ``None`` the current
device's compute capability is used (the default, load-bearing behavior).
Returns the loaded build result.
"""
global _check_sass
common_data = _common_data_for_cc(cc)
result = build_impl_fn(
*args,
common_data,
**kwargs,
)
if _check_sass:
cubin = result._get_cubin()
temp_cubin_file_name = _check_compile_result(cubin)
os.unlink(temp_cubin_file_name)
return result
def call_compile(build_impl_cls: Callable, *args, cc, **kwargs):
"""Compile only (no load) for an explicit compute capability ``cc``.
``build_impl_cls`` is a ``Device<Algo>BuildResult`` type exposing a
``compile(...)`` staticmethod. Unlike :func:`call_build`, this never touches
the CUDA driver — it can run on a machine with no GPU. The returned build
result is *not* loaded; call ``.load()`` (once, on a matching device) before
executing. ``cc`` is a ``(major, minor)`` pair and is required.
"""
common_data = _common_data_for_cc(cc)
# build_impl_cls is a Device<Algo>BuildResult class exposing a compile()
# staticmethod; it's typed Callable here, so silence the attr check.
return build_impl_cls.compile(*args, common_data, **kwargs) # type: ignore[attr-defined]
def build_for_ccs(build_impl_cls: Callable, *args, compute_capability=None, **kwargs):
"""Build the ``{cc_key: build_result}`` map for an algorithm.
With ``compute_capability=None`` (the default), this performs a fused
build+load for the current device and returns a single-entry map whose
result is already loaded. Otherwise it compiles (without loading) for each
requested compute capability and returns ``{cc_key: build_result}``, with
each result loaded lazily on first use by ``resolve_build_result``.
"""
ccs = normalize_compute_capabilities(compute_capability)
if ccs is None:
# Fused build+load for the current device. Query its cc once (clear error
# if no device) and pass it through, so call_build doesn't re-query.
device_id, cc_key = current_device_info()
build_result = call_build(build_impl_cls, *args, cc=key_to_cc(cc_key), **kwargs)
# The fused build already loaded the kernels; mark it so the lazy
# load() in resolve_build_result() is a no-op (a second C load would leak /
# re-register the library).
build_result._loaded = True
return _PerCCBuildResults({cc_key: build_result}, loaded_device_id=device_id)
return _PerCCBuildResults(
{
cc_to_key(cc): call_compile(build_impl_cls, *args, cc=cc, **kwargs)
for cc in ccs
}
)
def cc_to_key(cc) -> int:
"""Normalize a compute capability to the integer key ``major * 10 + minor``.
Accepts an int (``90``, ``75``), a ``(major, minor)`` pair, or a string
like ``"90"`` / ``"9.0"``.
"""
if isinstance(cc, (tuple, list)):
major, minor = cc
return int(major) * 10 + int(minor)
if isinstance(cc, str):
cc = cc.replace(".", "")
return int(cc)
return int(cc)
def key_to_cc(key: int):
"""Inverse of :func:`cc_to_key`: integer key -> ``(major, minor)`` pair."""
return (key // 10, key % 10)
def normalize_compute_capabilities(compute_capability):
"""Normalize the ``compute_capability=`` argument of ``make_<algo>``.
Returns a sorted list of unique ``(major, minor)`` pairs, or ``None`` to
mean "use the current device" (the default build path). Accepts a single
cc (int / pair / str) or a list thereof.
"""
if compute_capability is None:
return None
if isinstance(compute_capability, (int, str)):
ccs = [compute_capability]
elif (
isinstance(compute_capability, tuple)
and len(compute_capability) == 2
and all(isinstance(x, int) for x in compute_capability)
):
# a single (major, minor) pair
ccs = [compute_capability]
else:
ccs = list(compute_capability)
keys = sorted({cc_to_key(cc) for cc in ccs})
if not keys:
raise ValueError("compute_capability list is empty")
return [key_to_cc(k) for k in keys]
def current_device_info() -> tuple[int, int]:
"""Return the current device ordinal and packed compute-capability key.
Raises a clear, actionable error if no CUDA device is available: building
without a GPU has no device to infer the target arch from, so the caller
must pass an explicit ``compute_capability=``.
"""
try:
device = CudaDevice()
cc = device.compute_capability
except Exception as e:
raise RuntimeError(
"No compute_capability was given and no CUDA device is available to target."
) from e
return device.device_id, cc_to_key(tuple(cc))
def current_device_cc_key() -> int:
"""The current device's compute capability as a ``major * 10 + minor`` key."""
return current_device_info()[1]
def current_device_id() -> int:
"""The current CUDA device ordinal, without a compute-capability query.
The compute-capability query roughly doubles the cost of
``current_device_info()``, and callers that only key per-device state
(see ``resolve_build_result``) run on every algorithm invocation.
"""
try:
return CudaDevice().device_id
except Exception as e:
raise RuntimeError("No CUDA device is available to execute on.") from e
def resolve_build_result(build_results: dict, bound_result=None):
"""Load the build result for the current device.
``bound_result`` is a default-build wrapper's construction-time binding
(see cache_build_results): already the loaded result for the wrapper's
device, returned without any device query. Deserialized wrappers have no
binding and resolve per call.
"""
if bound_result is not None:
return bound_result
# Wrappers always hold a _PerCCBuildResults (build_for_ccs and
# deserialization both produce one); the per-device ownership/clone
# protocol in resolve() relies on it, so fail loudly on anything else
# rather than fall back to an unprotected load.
assert isinstance(build_results, _PerCCBuildResults)
if len(build_results) == 1:
# A singular _PerCCBuildResults is used as-is whatever the current device's
# compute capability is (single-target blobs were already cc-checked at
# deserialization), so only the device ordinal is needed to key the
# per-device loaded state. This path runs on every call of AOT and
# deserialized wrappers; skip the costlier compute-capability query.
(build_result_cc,) = build_results
device_id = current_device_id()
else:
device_id, device_cc_key = current_device_info()
build_result_cc = device_cc_key
if build_result_cc not in build_results:
available = ", ".join(
f"{maj}.{minor}"
for maj, minor in (key_to_cc(k) for k in sorted(build_results))
)
major, minor = key_to_cc(build_result_cc)
raise RuntimeError(
f"This algorithm was compiled for compute capabilities [{available}], "
f"but the current device has compute capability {major}.{minor}. "
f"Rebuild with compute_capability including {major}{minor}."
)
return build_results.resolve(build_result_cc, device_id)

172
cccl_upstream/python/cuda_cccl/cuda/compute/_cpp_compile.py Normal file
View File

@@ -0,0 +1,172 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
C++ code generation and compilation infrastructure.
"""
from __future__ import annotations
import functools
from cuda.cccl import get_include_paths
from cuda.core import Device, Program, ProgramOptions
from ._bindings import TypeEnum
from ._device_code import DeviceCode
try:
from ._build_info import USING_V2 # type: ignore[import-not-found]
except ImportError:
USING_V2 = False
def _get_arch_string() -> str:
"""Target arch string for iterator LTO-IR compilation.
Honors the build's target compute capability (set for multi-arch / no-GPU
builds) so iterator device code is compiled for the lowest target arch and
links into every build result; falls back to the current device otherwise.
"""
from ._target_cc import get_target_cc
cc = get_target_cc()
if cc is None:
cc = Device().compute_capability
cc_major, cc_minor = cc
return f"sm_{cc_major}{cc_minor}"
@functools.lru_cache(maxsize=1)
def _get_include_paths() -> list[str]:
"""Get include paths for CCCL headers."""
paths = get_include_paths().as_tuple()
return [p for p in paths if p is not None]
def compile_cpp_to_ltoir(
source: str,
arch: str | None = None,
) -> bytes:
"""
Compile C++ source code to LTOIR.
Args:
source: C++ source code string
arch: Target architecture (e.g., "sm_80"). If None, uses current device.
Returns:
LTOIR bytes
Example:
source = '''
extern "C" __device__ void my_add(void* a, void* b, void* result) {
*static_cast<int*>(result) = *static_cast<int*>(a) + *static_cast<int*>(b);
}
'''
ltoir = compile_cpp_to_ltoir(source)
"""
# Resolve the concrete arch before the cache lookup so the key reflects the
# compute capability compiled for. If arch stays None (the usual iterator/op
# call, resolved from target_cc), every target collapses to one key and
# LTO-IR built for one arch can be reused for another, which nvJitLink
# rejects.
if arch is None:
arch = _get_arch_string()
return _compile_cpp_to_ltoir_cached(source, arch)
@functools.lru_cache(maxsize=256)
def _compile_cpp_to_ltoir_cached(source: str, arch: str) -> bytes:
# Get include paths
include_paths = _get_include_paths()
# Configure compilation options for LTO
opts = ProgramOptions(
arch=arch,
relocatable_device_code=True,
link_time_optimization=True,
std="c++20",
define_macro="__NV_NO_VECTOR_DEPRECATION_DIAG",
include_path=include_paths,
)
# Compile to LTOIR
program = Program(source, "c++", options=opts)
result = program.compile("ltoir")
return result.code
# Expose the cached-callable surface (cache_info/cache_clear) on the public
# entry point, backed by the arch-aware inner cache.
compile_cpp_to_ltoir.cache_clear = _compile_cpp_to_ltoir_cached.cache_clear # type: ignore[attr-defined]
compile_cpp_to_ltoir.cache_info = _compile_cpp_to_ltoir_cached.cache_info # type: ignore[attr-defined]
def compile_cpp_op_code(source: str, arch: str | None = None) -> DeviceCode:
"""Compile C++ wrapper source to whatever form the active backend prefers.
Returns a :class:`DeviceCode` wrapping the bytes and the matching format tag.
Cached so identical iterator structures produce identical code bytes —
callers can inspect ``cache_info()`` to verify symbol determinism.
"""
# v2 keeps the C++ source verbatim (arch-independent); v1 resolves the
# concrete arch before caching (see compile_cpp_to_ltoir).
if USING_V2:
return _compile_cpp_op_code_cached(source, None)
if arch is None:
arch = _get_arch_string()
return _compile_cpp_op_code_cached(source, arch)
@functools.lru_cache(maxsize=256)
def _compile_cpp_op_code_cached(source: str, arch: str | None) -> DeviceCode:
if USING_V2:
return DeviceCode(op_bytes=source.encode("utf-8"), kind="cpp_source")
return DeviceCode(op_bytes=compile_cpp_to_ltoir(source, arch=arch), kind="ltoir")
compile_cpp_op_code.cache_clear = _compile_cpp_op_code_cached.cache_clear # type: ignore[attr-defined]
compile_cpp_op_code.cache_info = _compile_cpp_op_code_cached.cache_info # type: ignore[attr-defined]
def cpp_type_from_descriptor(type_desc) -> str | None:
"""
Get the C++ type name from a TypeDescriptor.
Important: for efficiency, this function returns None
for non-primitive types. Callers must take care
to handle that case appropriately.
"""
# Map TypeEnum to C++ types
type_map = {
TypeEnum.INT8: "int8_t",
TypeEnum.INT16: "int16_t",
TypeEnum.INT32: "int32_t",
TypeEnum.INT64: "int64_t",
TypeEnum.UINT8: "uint8_t",
TypeEnum.UINT16: "uint16_t",
TypeEnum.UINT32: "uint32_t",
TypeEnum.UINT64: "uint64_t",
TypeEnum.FLOAT16: "__half",
TypeEnum.FLOAT32: "float",
TypeEnum.FLOAT64: "double",
TypeEnum.BOOLEAN: "bool",
TypeEnum.STORAGE: None,
}
return type_map[type_desc.info.typenum]
def make_variable_declaration(type_desc, name: str) -> str:
"""
Generate a C++ variable declaration, like "int32_t temp;"
or "alignas(8) char temp[16];"
"""
cpp_type = cpp_type_from_descriptor(type_desc)
if cpp_type is not None:
return f"{cpp_type} {name};"
# STORAGE type - use aligned char array
return f"alignas({type_desc.alignment}) char {name}[{type_desc.size}];"

45
cccl_upstream/python/cuda_cccl/cuda/compute/_device_code.py Normal file
View File

@@ -0,0 +1,45 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Public type for passing compiled device-side operator code into ``Op`` /
``RawOp``. Wraps the bytes together with their format tag so the two cannot
get out of sync as they flow through the binding layer.
Lives in its own module to keep ``op.py`` and the Cython ``_bindings_impl``
free of import-cycle headaches: the Cython side duck-types on
``(op_bytes, kind)`` attributes and never imports the class directly.
"""
from __future__ import annotations
from dataclasses import dataclass
# Tag values mirror the C-side ``cccl_op_code_type`` enum.
_VALID_KINDS = ("ltoir", "llvm_ir", "cpp_source")
@dataclass(frozen=True)
class DeviceCode:
"""A compiled-or-source device-code blob ready to hand to ``Op``.
Args:
op_bytes: the raw blob (LTO-IR, LLVM bitcode, or C++ source bytes).
kind: one of ``"ltoir"`` (default), ``"llvm_ir"``, ``"cpp_source"``;
tells the backend how to interpret ``op_bytes``.
For most uses you don't construct ``DeviceCode`` directly — the internal
JIT-compile helpers return one, and the iterator/algorithm machinery
forwards them. Construct explicitly when feeding a ``RawOp`` from outside
the default pipeline.
"""
op_bytes: bytes
kind: str = "ltoir"
def __post_init__(self):
if self.kind not in _VALID_KINDS:
raise ValueError(
f"DeviceCode.kind must be one of {_VALID_KINDS!r}; got {self.kind!r}"
)

1053
cccl_upstream/python/cuda_cccl/cuda/compute/_jit.py Normal file
View File

File diff suppressed because it is too large Load Diff

0
cccl_upstream/python/cuda_cccl/cuda/compute/_nvtx.py Normal file
View File

370
cccl_upstream/python/cuda_cccl/cuda/compute/_odr_helpers.py Normal file
View File

@@ -0,0 +1,370 @@
# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
ODR (One Definition Rule) Helpers for CCCL Python Interop.
This module provides utilities to create wrapper functions for
device functions that are defined in Python and JIT compiled by Numba.
On the C++ side, these functions are declared as `extern "C"` functions with
void* parameters - the arguments types can not be known at C++ compile time.
Thus, the helpers in this module generate wrapper device functions that accept
void* arguments (matching C++ declarations), cast them to the correct
typed arguments, load/store values as needed, and call the original
function with properly typed arguments.
Example flow:
User provides: def add(x: int32, y: int32) -> int32
Wrapper signature: void(void*, void*, void*) # x_ptr, y_ptr, result_ptr
C++ sees: extern "C" void wrapped_add(void*, void*, void*);
"""
from __future__ import annotations
import enum
import itertools
import textwrap
import threading
from typing import TYPE_CHECKING
from numba import types
from numba.core.extending import intrinsic
from ._utils import sanitize_identifier
if TYPE_CHECKING:
from numba.core.typing import Signature
# Global counter to generate unique symbol names even when the same function
# is used multiple times (e.g., as both selectors in `three_way_partition`).
_wrapper_name_counter = itertools.count()
_wrapper_name_lock = threading.Lock()
__all__ = [
"create_op_void_ptr_wrapper",
"create_stateful_op_void_ptr_wrapper",
"create_advance_void_ptr_wrapper",
"create_input_dereference_void_ptr_wrapper",
"create_output_dereference_void_ptr_wrapper",
]
class _ArgMode(enum.Enum):
"""How a void* argument should be handled in wrapper codegen."""
LOAD = "load" # Cast to typed pointer, load value
PTR = "ptr" # Cast to typed pointer, pass pointer directly
STORE = "store" # Cast to typed pointer, store return value here
# Unpack packed data pointers into array structs
STATE = "state"
class _ArgSpec:
"""Specification for a wrapper argument."""
__slots__ = ("numba_type", "mode")
def __init__(self, numba_type, mode: _ArgMode):
self.numba_type = numba_type
self.mode = mode
def _build_numba_array_struct(context, builder, array_type, data_ptr, info):
"""Build a numba array struct from a data pointer and array info.
Args:
context: Numba codegen context
builder: LLVM IR builder
array_type: Numba Array type for the array
data_ptr: LLVM value for the data pointer
info: Dict with 'shape', 'itemsize', 'strides' for the array
Returns:
LLVM value representing the array struct
"""
import llvmlite.ir as ir
from numba.cuda.np.arrayobj import make_array, populate_array
out_ary = make_array(array_type)(context, builder)
populate_array(
out_ary,
data=data_ptr,
shape=[ir.Constant(ir.IntType(64), info["shape"])],
strides=[ir.Constant(ir.IntType(64), info["strides"])],
itemsize=info["itemsize"],
meminfo=None,
)
return out_ary._getvalue()
def _unpack_state_arrays(context, builder, packed_ptr, type_info_pairs):
"""Unpack packed data pointers into numba array structs.
Args:
context: Numba codegen context
builder: LLVM IR builder
packed_ptr: void* pointing to an array of data pointers
type_info_pairs: List of (array_type, info) tuples
Returns:
List of LLVM values representing the unpacked array structs
"""
import llvmlite.ir as ir
# Cast void* to pointer-to-pointer (array of pointers)
ptr_type = ir.IntType(64).as_pointer()
base_ptr = builder.bitcast(packed_ptr, ptr_type.as_pointer())
result = []
for j, (array_type, info) in enumerate(type_info_pairs):
# Load j-th pointer from the array and cast to correct type
elem_ptr = builder.gep(base_ptr, [ir.Constant(ir.IntType(32), j)])
dtype_llvm = context.get_value_type(array_type.dtype)
typed_ptr_ptr = builder.bitcast(elem_ptr, dtype_llvm.as_pointer().as_pointer())
data_ptr = builder.load(typed_ptr_ptr)
# Build array struct from pointer
array_val = _build_numba_array_struct(
context, builder, array_type, data_ptr, info
)
result.append(array_val)
return result
def _codegen_void_ptr_wrapper(
context, builder, args, arg_specs, func_device, inner_sig
):
"""Generate LLVM IR for a void* wrapper function.
This is the codegen implementation shared by all void* wrappers.
It processes each argument according to its _ArgSpec mode, calls
the inner function, and stores the result if needed.
Args:
context: Numba codegen context
builder: LLVM IR builder
args: LLVM values for the void* arguments
arg_specs: List of _ArgSpec describing each argument
func_device: The device function to call
inner_sig: Numba signature for the inner function
Returns:
LLVM dummy value (for void return)
"""
input_vals = []
state_array_vals = []
ret_ptr = None
for i, (arg, spec) in enumerate(zip(args, arg_specs)):
match spec.mode:
case _ArgMode.LOAD:
# Cast void* to typed pointer and load value
llvm_type = context.get_value_type(spec.numba_type)
typed_ptr = builder.bitcast(arg, llvm_type.as_pointer())
val = builder.load(typed_ptr)
input_vals.append(val)
case _ArgMode.PTR:
# Cast void* to typed pointer, pass pointer directly
llvm_type = context.get_value_type(spec.numba_type.dtype)
typed_ptr = builder.bitcast(arg, llvm_type.as_pointer())
input_vals.append(typed_ptr)
case _ArgMode.STORE:
# Cast void* to typed pointer for storing result
llvm_type = context.get_value_type(spec.numba_type)
ret_ptr = builder.bitcast(arg, llvm_type.as_pointer())
case _ArgMode.STATE:
# Cast void* to a packed array of pointers and unpack them
array_vals = _unpack_state_arrays(
context, builder, arg, spec.numba_type
)
state_array_vals.extend(array_vals)
case _:
raise ValueError(f"Invalid arg mode: {spec.mode}")
# Prepend state arrays at the beginning (inner_sig expects state args first)
input_vals = state_array_vals + input_vals
# Call the inner function
cres = context.compile_subroutine(builder, func_device, inner_sig, caching=False)
result = context.call_internal(builder, cres.fndesc, inner_sig, input_vals)
# Store result if needed
if ret_ptr is not None:
builder.store(result, ret_ptr)
return context.get_dummy_value()
def _create_void_ptr_wrapper(
func, name: str, arg_specs: list[_ArgSpec], inner_sig: "Signature"
):
"""
Given a function and a list of _ArgSpec, create a wrapper function
that takes all void* arguments, bitcasts them to the
appropriate typed pointers, and calls the inner function with
the typed arguments. Each void* argument is handled according
to its _ArgSpec.
Args:
func: The function to wrap (will be compiled as device function)
name: Base name for the wrapper function
arg_specs: List of _ArgSpec describing each void* argument
inner_sig: Numba signature for the inner function call
Returns:
Tuple of (wrapper_func, wrapper_sig)
"""
from numba.cuda import jit as cuda_jit
# Wrap function as device function
func_device = cuda_jit(device=True)(func)
# Generate argument names and signature
arg_names = [f"arg_{i}" for i in range(len(arg_specs))]
arg_str = ", ".join(arg_names)
void_sig = types.void(*(types.voidptr for _ in arg_specs))
# Create unique wrapper name using global counter
sanitized_name = sanitize_identifier(name)
if not sanitized_name.isidentifier():
raise ValueError(
f"Function name '{name}' cannot be sanitized into a valid identifier"
)
for arg_name in arg_names:
if not arg_name.isidentifier():
raise ValueError(
f"Invalid argument name '{arg_name}' - must be a valid identifier"
)
with _wrapper_name_lock:
unique_suffix = next(_wrapper_name_counter)
wrapper_name = f"wrapped_{sanitized_name}_{unique_suffix}"
# We need exec() here because Numba's @intrinsic decorator requires:
# 1. A function with a specific signature visible at parse time
# 2. The number of arguments must match the wrapper signature
# The actual codegen logic is in _codegen_void_ptr_wrapper - this just
# creates the minimal intrinsic shell that delegates to it.
wrapper_src = textwrap.dedent(f"""
@intrinsic
def impl(typingctx, {arg_str}):
def codegen(context, builder, impl_sig, args):
return codegen_helper(context, builder, args, arg_specs, func_device, inner_sig)
return void_sig, codegen
def {wrapper_name}({arg_str}):
return impl({arg_str})
""")
local_dict = {
"intrinsic": intrinsic,
"void_sig": void_sig,
"arg_specs": arg_specs,
"func_device": func_device,
"inner_sig": inner_sig,
"codegen_helper": _codegen_void_ptr_wrapper,
}
exec(wrapper_src, {}, local_dict)
wrapper_func = local_dict[wrapper_name]
wrapper_func.__globals__.update(local_dict)
return wrapper_func, void_sig
def create_op_void_ptr_wrapper(op, sig: "Signature"):
"""Creates a wrapper function for user-defined operators like unary or binary operators.
The wrapper takes N+1 arguments where N is the number of input arguments to `op`, the last
argument is a pointer to the result.
"""
arg_specs = [_ArgSpec(t, _ArgMode.LOAD) for t in sig.args]
arg_specs.append(_ArgSpec(sig.return_type, _ArgMode.STORE))
return _create_void_ptr_wrapper(op, op.__name__, arg_specs, sig)
def create_stateful_op_void_ptr_wrapper(
op, sig: "Signature", state_array_types, state_info
):
"""Creates a wrapper function for a stateful operator with void* arguments.
The wrapper takes N+2 void* arguments:
- states_ptr: pointer to packed array of data pointers for state arrays
- N input args: one for each regular input argument
- result: pointer where result is stored
Args:
op: The user's callable operator
sig: The signature of the operator (state_array1, state_array2, ..., regular_arg1, regular_arg2, ...) -> return_type
state_array_types: List/tuple of numba Array types for the state parameters
state_info: List/tuple of dicts with 'shape', 'itemsize', 'strides' for each state array
Returns:
Tuple of (wrapper_func, wrapper_sig)
"""
num_states = len(state_array_types)
# Build arg_specs: states_ptr + regular inputs + result
# The packed state arrays spec goes first, then regular LOAD args, then STORE for result
# numba_type is a list of (array_type, info) tuples
type_info_pairs = list(zip(state_array_types, state_info))
arg_specs = [_ArgSpec(type_info_pairs, _ArgMode.STATE)]
for i in range(num_states, len(sig.args)):
arg_specs.append(_ArgSpec(sig.args[i], _ArgMode.LOAD))
arg_specs.append(_ArgSpec(sig.return_type, _ArgMode.STORE))
return _create_void_ptr_wrapper(op, op.__name__, arg_specs, sig)
def create_advance_void_ptr_wrapper(advance_fn, state_ptr_type):
"""Creates a wrapper function for iterator advance method.
The wrapper takes 2 void* arguments:
- state pointer
- offset pointer (points to uint64 value)
"""
arg_specs = [
_ArgSpec(state_ptr_type, _ArgMode.PTR),
_ArgSpec(types.uint64, _ArgMode.LOAD), # uint64 is the offset type
]
inner_sig = types.void(state_ptr_type, types.uint64)
return _create_void_ptr_wrapper(
advance_fn, advance_fn.__name__, arg_specs, inner_sig
)
def create_input_dereference_void_ptr_wrapper(deref_fn, state_ptr_type, value_type):
"""Creates a wrapper function for input iterator dereference method.
The wrapper takes 2 void* arguments:
- state pointer
- result pointer (function writes result here)
"""
arg_specs = [
_ArgSpec(state_ptr_type, _ArgMode.PTR),
_ArgSpec(types.CPointer(value_type), _ArgMode.PTR),
]
inner_sig = types.void(state_ptr_type, types.CPointer(value_type))
return _create_void_ptr_wrapper(deref_fn, deref_fn.__name__, arg_specs, inner_sig)
def create_output_dereference_void_ptr_wrapper(deref_fn, state_ptr_type, value_type):
"""Creates a wrapper function for output iterator dereference method.
The wrapper takes 2 void* arguments:
- state pointer
- value pointer (value to write)
"""
arg_specs = [
_ArgSpec(state_ptr_type, _ArgMode.PTR),
_ArgSpec(value_type, _ArgMode.LOAD),
]
inner_sig = types.void(state_ptr_type, value_type)
return _create_void_ptr_wrapper(deref_fn, deref_fn.__name__, arg_specs, inner_sig)

132
cccl_upstream/python/cuda_cccl/cuda/compute/_proxy.py Normal file
View File

@@ -0,0 +1,132 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Device-less placeholders for ahead-of-time (no-GPU) compilation.
``ProxyArray`` / ``ProxyValue`` describe *only* the dtype (and, for arrays,
shape/contiguity) of an argument, with no backing GPU allocation. Pass them to
``make_<algo>(...)`` together with ``compute_capability=`` to compile an
algorithm on a machine that has no GPU (or no live data), then ``serialize()``
the result. The real device arrays / scalars are supplied later at ``__call__``.
Accessing a proxy's data pointer raises ``RuntimeError`` — a proxy can be used
to *build* an algorithm but never to *run* one.
"""
from __future__ import annotations
import numpy as np
_PROXY_DATA_ERROR = (
"ProxyArray has no GPU data — it is a build-time placeholder only. "
"Pass a real device array when calling the compiled algorithm."
)
_PROXY_VALUE_DATA_ERROR = (
"ProxyValue has no data — it is a build-time placeholder only. "
"Pass a real scalar or numpy array when calling the compiled algorithm."
)
class _ProxyCAI(dict):
"""CAI dict whose 'data' key raises on access."""
def __missing__(self, key):
if key == "data":
raise RuntimeError(_PROXY_DATA_ERROR)
raise KeyError(key)
def get(self, key, default=None):
# dict.get() bypasses __missing__, so guard it too: a consumer that
# defensively does cai.get("data") must still hit the loud failure
# rather than silently receiving a null pointer.
if key == "data":
raise RuntimeError(_PROXY_DATA_ERROR)
return super().get(key, default)
class ProxyArray:
"""Dtype-only placeholder for a device array.
Use in place of a real device array when calling ``make_<algo>()`` to
trigger ahead-of-time compilation without allocating GPU memory — for
example, on a build machine that has no GPU or no live data.
Satisfies the ``DeviceArrayLike`` protocol:
* ``is_device_array(proxy)`` -> ``True``
* ``get_dtype(proxy)`` -> the dtype supplied at construction
* ``get_data_pointer(proxy)``-> raises ``RuntimeError``
* ``is_contiguous(proxy)`` -> ``True``
Accessing the data pointer raises ``RuntimeError``; passing a
``ProxyArray`` to a compiled algorithm's ``__call__`` is not supported.
Example::
from cuda.compute import ProxyArray, make_reduce_into, OpKind
import numpy as np
reducer = make_reduce_into(
d_in=ProxyArray(np.float32),
d_out=ProxyArray(np.float32),
op=OpKind.PLUS,
h_init=np.zeros(1, dtype=np.float32),
compute_capability=[80, 90],
)
reducer.serialize()
"""
__slots__ = ("_dtype",)
def __init__(self, dtype):
self._dtype = np.dtype(dtype)
@property
def dtype(self) -> np.dtype:
return self._dtype
@property
def __cuda_array_interface__(self) -> dict:
return _ProxyCAI(
{
"shape": (1,),
"typestr": self._dtype.str,
"version": 3,
"strides": None, # C-contiguous
# "data" is intentionally absent — accessing it raises RuntimeError
}
)
def __repr__(self) -> str:
return f"ProxyArray(dtype={self._dtype})"
class ProxyValue:
"""Dtype-only placeholder for a scalar / initial-value argument.
Use in place of a real numpy scalar or array when calling ``make_<algo>()``
to trigger ahead-of-time compilation without real data — for example, for
the ``h_init`` argument of :func:`~cuda.compute.make_reduce_into`.
Accessing the data of a ``ProxyValue`` raises ``RuntimeError``; passing
one to a compiled algorithm's ``__call__`` is not supported.
"""
__slots__ = ("_dtype",)
def __init__(self, dtype):
self._dtype = np.dtype(dtype)
@property
def dtype(self) -> np.dtype:
return self._dtype
def __repr__(self) -> str:
return f"ProxyValue(dtype={self._dtype})"
def is_proxy(obj) -> bool:
"""True if *obj* is a build-time placeholder (ProxyArray/ProxyValue)."""
return isinstance(obj, (ProxyArray, ProxyValue))

41
cccl_upstream/python/cuda_cccl/cuda/compute/_serialization/__init__.py Normal file
View File

@@ -0,0 +1,41 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Serialization for cuda.compute algorithms."""
from __future__ import annotations
from .dispatch import deserialize as deserialize
from .dispatch import serialize as serialize
from .serializable import BOOL as BOOL
from .serializable import BUILD_RESULT as BUILD_RESULT
from .serializable import BUILD_RESULTS as BUILD_RESULTS
from .serializable import CONDITIONAL as CONDITIONAL
from .serializable import ENUM as ENUM
from .serializable import ITER as ITER
from .serializable import NESTED as NESTED
from .serializable import OP as OP
from .serializable import U8 as U8
from .serializable import U32 as U32
from .serializable import U64 as U64
from .serializable import VALUE as VALUE
from .serializable import Serializable as Serializable
__all__ = [
"serialize",
"deserialize",
"Serializable",
"ITER",
"OP",
"VALUE",
"U8",
"U32",
"U64",
"BOOL",
"ENUM",
"CONDITIONAL",
"BUILD_RESULT",
"BUILD_RESULTS",
"NESTED",
]

217
cccl_upstream/python/cuda_cccl/cuda/compute/_serialization/codec.py Normal file
View File

@@ -0,0 +1,217 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Byte serialization of CCCL descriptors (``cccl_op_t``, ``cccl_iterator_t``, etc.)."""
from __future__ import annotations
import struct
import numpy as np
from cuda.cccl import __version__ as _PKG_VERSION
from .._bindings import Iterator, IteratorKind, Op, OpKind, TypeEnum, TypeInfo, Value
from .._device_code import DeviceCode
# An opaque 8-byte marker identifying a cuda.compute serialization blob.
_MAGIC = b"CCCLPYS1"
class Writer:
"""Append-only little-endian byte buffer."""
def __init__(self) -> None:
self.buf = bytearray()
def u8(self, v: int) -> None:
self.buf += struct.pack("<B", v)
def u32(self, v: int) -> None:
self.buf += struct.pack("<I", v)
def u64(self, v: int) -> None:
self.buf += struct.pack("<Q", v)
def blob(self, b: bytes) -> None:
self.u64(len(b))
self.buf += b
def text(self, s: str) -> None:
self.blob(s.encode("utf-8"))
def getvalue(self) -> bytes:
return bytes(self.buf)
class Reader:
"""Bounds-checked little-endian reader over a bytes blob."""
def __init__(self, data: bytes) -> None:
self._data = memoryview(data)
self.pos = 0
def _take(self, n: int) -> memoryview:
end = self.pos + n
if end > len(self._data):
raise ValueError("serialization descriptor blob truncated")
out = self._data[self.pos : end]
self.pos = end
return out
def u8(self) -> int:
return struct.unpack("<B", self._take(1))[0]
def u32(self) -> int:
return struct.unpack("<I", self._take(4))[0]
def u64(self) -> int:
return struct.unpack("<Q", self._take(8))[0]
def blob(self) -> bytes:
return bytes(self._take(self.u64()))
def text(self) -> str:
return self.blob().decode("utf-8")
def remaining(self) -> bytes:
"""Bytes after the descriptor region: the C build_result blob."""
return bytes(self._data[self.pos :])
# --- framing -----------------------------------------------------------------
def _check_header(r: Reader) -> None:
"""Validate the magic and package-version stamp at the start of a blob."""
if bytes(r._take(len(_MAGIC))) != _MAGIC:
raise ValueError(
"serialization blob: bad magic (not a cuda.compute serialization blob)"
)
version = r.text()
if version != _PKG_VERSION:
raise ValueError(
"serialization blob: cuda-cccl version mismatch "
f"(blob={version!r}, current={_PKG_VERSION!r}); "
"re-serialize with this version of cuda-cccl"
)
def begin(algo_tag: str) -> Writer:
"""Start a descriptor sidecar with the magic/version/algo header.
``algo_tag`` is the algorithm class's ``__qualname__`` (e.g. ``"_Reduce"``).
"""
w = Writer()
w.buf += _MAGIC
w.text(_PKG_VERSION)
w.text(algo_tag)
return w
def open(blob: bytes, expected_algo: str) -> Reader:
"""Validate the header and return a reader positioned at the first field."""
r = Reader(blob)
_check_header(r)
algo = r.text()
if algo != expected_algo:
raise ValueError(
f"serialization blob: wrong algorithm (blob tag={algo!r}, expected={expected_algo!r})"
)
return r
def peek_algo(blob: bytes) -> str:
"""Return the algorithm tag (class ``__qualname__``) from a blob header
without consuming the blob.
Validates magic + version. Used by the generic ``deserialize`` dispatcher to
pick the right algorithm reconstructor.
"""
r = Reader(blob)
_check_header(r)
return r.text()
# --- descriptor (de)serialization --------------------------------------------
def write_type_info(w: Writer, ti: TypeInfo) -> None:
w.u64(ti.size)
w.u64(ti.alignment)
w.u32(int(ti.typenum))
def read_type_info(r: Reader) -> TypeInfo:
size = r.u64()
alignment = r.u64()
type_enum = r.u32()
return TypeInfo(size, alignment, TypeEnum(type_enum))
def write_op(w: Writer, op: Op) -> None:
# Serialize the operator's device code in full so reconstruction needs no
# JIT; only per-call op state is omitted.
w.u32(int(op.operator_type))
w.text(op.name)
w.blob(op.ltoir)
w.text(op.code.kind)
w.u32(op.state_alignment)
# State size is structural: it fixes op_data.size at construction. The state
# bytes themselves are bound per-call.
w.u64(len(op.state))
extras = op.extra_code
w.u32(len(extras))
for dc in extras:
w.blob(dc.op_bytes)
w.text(dc.kind)
def read_op(r: Reader) -> Op:
operator_type = OpKind(r.u32())
name = r.text()
code = r.blob()
code_kind = r.text()
state_alignment = r.u32()
state_size = r.u64()
n_extra = r.u32()
extras = [DeviceCode(op_bytes=r.blob(), kind=r.text()) for _ in range(n_extra)]
return Op(
name=name,
operator_type=operator_type,
ltoir=DeviceCode(op_bytes=code, kind=code_kind),
state=bytes(state_size), # zero placeholder; real bytes bound per-call
state_alignment=state_alignment,
extra_ltoirs=extras,
)
def write_iterator(w: Writer, it: Iterator) -> None:
w.u8(1 if it.is_kind_pointer() else 0)
w.u32(it.alignment)
write_type_info(w, it.value_type)
write_op(w, it.advance_op)
write_op(w, it.dereference_or_assign_op)
def read_iterator(r: Reader) -> Iterator:
kind = IteratorKind.POINTER if r.u8() else IteratorKind.ITERATOR
alignment = r.u32()
value_type = read_type_info(r)
advance = read_op(r)
deref = read_op(r)
# state is bound per-call (set_cccl_iterator_state); start with none.
return Iterator(alignment, kind, advance, deref, value_type, state=None)
def write_value(w: Writer, val: Value) -> None:
# Only the type is static; the value bytes are bound per-call.
write_type_info(w, val.type)
def read_value(r: Reader) -> Value:
value_type = read_type_info(r)
# Placeholder state sized to the value type; __call__ rebinds the real bytes.
placeholder = np.zeros(value_type.size, dtype=np.uint8)
return Value(value_type, placeholder)

60
cccl_upstream/python/cuda_cccl/cuda/compute/_serialization/dispatch.py Normal file
View File

@@ -0,0 +1,60 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Public, free-standing serialize/deserialize entry points."""
from __future__ import annotations
from typing import Any
from . import codec
from .serializable import Serializable
def serialize(algorithm: Any) -> bytes:
"""Serialize a built algorithm into a blob of bytes.
Args:
algorithm: An object returned by a ``make_*`` factory (e.g.
:func:`make_reduce_into`, :func:`make_exclusive_scan`).
Returns:
A versioned, self-describing byte blob. Reconstruct it with
:func:`deserialize`; no objects required at load time.
"""
if not callable(getattr(type(algorithm), "serialize", None)):
raise TypeError(
f"{type(algorithm).__name__} is not a serializable algorithm "
"(expected an object from a make_* factory)."
)
return algorithm.serialize()
def deserialize(blob: bytes):
"""Reconstruct a built algorithm from a blob produced by :func:`serialize`.
Warning:
The returned object is **not safe to use from multiple threads
concurrently**. Do not deserialize once and share the object across
threads: every call writes its arguments (array pointers, sizes,
operator and initial-value state) into the object before launching, so
overlapping calls can launch kernels with another thread's arguments —
silently wrong results or CUDA errors, with no exception raised at the
point of misuse. Unlike the ``make_*`` factories, which hand each
calling thread its own cached object, ``deserialize`` returns a fresh
uncached object with no per-thread protection. For concurrent use,
call :func:`deserialize` in each thread — reconstruction performs no
recompilation, so per-thread deserialization from one shared blob is
cheap. One thread at a time (for example, handing the object between
threads with proper ordering) is fine.
Raises:
ValueError: if the blob is malformed or its algorithm tag is unknown.
"""
tag = codec.peek_algo(blob)
try:
cls = Serializable._registry[tag]
except KeyError:
raise ValueError(f"serialization blob: unknown algorithm tag {tag!r}") from None
return cls.deserialize(blob)

285
cccl_upstream/python/cuda_cccl/cuda/compute/_serialization/serializable.py Normal file
View File

@@ -0,0 +1,285 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Schema-driven serialize/deserialize base for cuda.compute algorithms.
A built-algorithm class declares a ``__serialization_schema__`` listing its
serialized members as ``(attr_name, kind)`` pairs, including its ``build_result``
as a ``BUILD_RESULT(<type>)`` member. ``Serializable`` provides generic
``serialize``/``deserialize`` that walk the schema, so subclasses need no
hand-written codec and both directions share one field order. Subclasses
auto-register by their ``__qualname__`` for the free-function
``deserialize`` dispatcher.
"""
from __future__ import annotations
from typing import Any, Callable, TypeVar
from . import codec
class _Kind:
"""A serialized member kind: writes/reads one value.
``obj`` is the instance being (de)serialized; most kinds ignore it, but
``CONDITIONAL`` uses it to read a selector member deserialized earlier.
"""
__slots__ = ()
def write(self, w: codec.Writer, value: Any, obj: Any) -> None:
raise NotImplementedError
def read(self, r: codec.Reader, obj: Any) -> Any:
raise NotImplementedError
class _Descriptor(_Kind):
"""Iterator / Op / Value descriptor, delegating to the codec codec."""
__slots__ = ("_write", "_read")
def __init__(self, writer: Callable, reader: Callable) -> None:
self._write = writer
self._read = reader
def write(self, w: codec.Writer, value: Any, obj: Any) -> None:
self._write(w, value)
def read(self, r: codec.Reader, obj: Any) -> Any:
return self._read(r)
ITER = _Descriptor(codec.write_iterator, codec.read_iterator)
OP = _Descriptor(codec.write_op, codec.read_op)
VALUE = _Descriptor(codec.write_value, codec.read_value)
class _Scalar(_Kind):
"""A fixed-width little-endian unsigned integer (u8/u32/u64)."""
__slots__ = ("width",)
def __init__(self, width: int) -> None:
self.width = width
def write(self, w: codec.Writer, value: Any, obj: Any) -> None:
{1: w.u8, 4: w.u32, 8: w.u64}[self.width](int(value))
def read(self, r: codec.Reader, obj: Any) -> int:
return {1: r.u8, 4: r.u32, 8: r.u64}[self.width]()
U8, U32, U64 = _Scalar(1), _Scalar(4), _Scalar(8)
class _Bool(_Kind):
"""A boolean, stored as a u8 (0/1)."""
__slots__ = ()
def write(self, w: codec.Writer, value: Any, obj: Any) -> None:
w.u8(1 if value else 0)
def read(self, r: codec.Reader, obj: Any) -> bool:
return bool(r.u8())
BOOL = _Bool()
class _Enum(_Kind):
"""An IntEnum member, stored as a u8 and reconstructed as the enum type."""
__slots__ = ("enum_cls",)
def __init__(self, enum_cls: Any) -> None:
self.enum_cls = enum_cls
def write(self, w: codec.Writer, value: Any, obj: Any) -> None:
w.u8(int(value))
def read(self, r: codec.Reader, obj: Any) -> Any:
return self.enum_cls(r.u8())
def ENUM(enum_cls: Any) -> _Enum:
"""Schema kind for a u8-backed IntEnum member."""
return _Enum(enum_cls)
class _SubObject(_Kind):
"""A sub-object with its own ``serialize()``/``deserialize()``, carried as a
length-prefixed blob. Used for the C ``build_result`` and for a nested
``Serializable`` (e.g. select wrapping three_way_partition)."""
__slots__ = ("cls",)
def __init__(self, cls: Any) -> None:
self.cls = cls
def write(self, w: codec.Writer, value: Any, obj: Any) -> None:
w.blob(value.serialize())
def read(self, r: codec.Reader, obj: Any) -> Any:
return self.cls.deserialize(r.blob())
def BUILD_RESULT(cls: type) -> _SubObject:
"""Schema kind for an algorithm's ``Device<Algo>BuildResult`` member."""
return _SubObject(cls)
class _BuildResults(_Kind):
"""A ``{cc: Device<Algo>BuildResult}`` mapping — one compiled build result per
target compute capability.
Wire form: ``u32`` count, then for each entry a ``u32`` cc key
(``cc_major * 10 + cc_minor``) followed by the length-prefixed build_result
blob. Entries are written in sorted-key order so the encoding is
deterministic. On read, each build_result is deserialized *without* loading
(``load=False``); the matching build result is loaded lazily on first call, so a
multi-arch artifact stays portable across GPUs and needs no live device to
deserialize.
"""
__slots__ = ("cls",)
def __init__(self, cls: Any) -> None:
self.cls = cls
def write(self, w: codec.Writer, value: Any, obj: Any) -> None:
from .._caching import _PerCCBuildResults
# Wrappers always hold a _PerCCBuildResults (build_for_ccs and
# read() below both produce one). serialize_build_result takes the
# per-cc source lock, so serialization cannot observe a source whose
# first device load is still in progress; a plain dict here would
# silently bypass that lock.
assert isinstance(value, _PerCCBuildResults)
ccs = sorted(value)
w.u32(len(ccs))
for cc in ccs:
w.u32(int(cc))
w.blob(value.serialize_build_result(cc))
def read(self, r: codec.Reader, obj: Any) -> Any:
count = r.u32()
# A single-target blob must match this device, so validate its cc-major
# eagerly (clear error at deserialize). A multi-arch blob legitimately
# carries build results for other archs, so defer the cc check — resolve_build_result
# picks the matching one at call time. Kernel load stays lazy either way.
check_cc = count == 1
entries = [(r.u32(), r.blob()) for _ in range(count)]
result: dict[int, Any] = {}
for cc, blob in entries:
if cc in result:
raise ValueError(
f"duplicate compute-capability key {cc} in build_results blob"
)
result[cc] = self.cls.deserialize(blob, load=False, check_cc=check_cc)
from .._caching import _PerCCBuildResults
return _PerCCBuildResults(result)
def BUILD_RESULTS(cls: type) -> _BuildResults:
"""Schema kind for a ``{cc: Device<Algo>BuildResult}`` build result mapping."""
return _BuildResults(cls)
def NESTED(cls: type) -> _SubObject:
"""Schema kind for a nested ``Serializable`` member (its blob is embedded)."""
return _SubObject(cls)
class _Conditional(_Kind):
"""A member whose kind depends on an earlier member's value.
``selector`` names a member deserialized *before* this one; ``branches``
maps each possible selector value to the kind to use (or ``None`` for an
absent member that (de)serializes to ``None``).
"""
__slots__ = ("selector", "branches")
def __init__(self, selector: str, branches: dict) -> None:
self.selector = selector
self.branches = branches
def _kind(self, obj: Any) -> "_Kind | None":
return self.branches[getattr(obj, self.selector)]
def write(self, w: codec.Writer, value: Any, obj: Any) -> None:
kind = self._kind(obj)
if kind is not None:
kind.write(w, value, obj)
def read(self, r: codec.Reader, obj: Any) -> Any:
kind = self._kind(obj)
return None if kind is None else kind.read(r, obj)
def CONDITIONAL(selector: str, branches: dict) -> _Conditional:
"""Schema kind for a member whose kind is chosen by ``selector``'s value."""
return _Conditional(selector, branches)
_S = TypeVar("_S", bound="Serializable")
class Serializable:
"""Mixin providing schema-driven serialize/deserialize + registration."""
__slots__ = ()
# __qualname__ -> subclass, populated as algorithm modules are imported.
_registry: dict[str, type[Serializable]] = {}
# Subclasses declare their serialized members here.
__serialization_schema__: tuple = ()
# Construction-time binding of a default-build wrapper's loaded result
# (see cache_build_results). Annotation only: storage comes from each
# subclass's __slots__; __init__ or deserialize() below assigns it.
_bound_build_result: Any
def __init_subclass__(cls, **kwargs: Any) -> None:
super().__init_subclass__(**kwargs)
Serializable._registry[cls.__qualname__] = cls
def _after_deserialize(self) -> None:
"""Hook to bind derived, non-serialized state after schema members are read.
Called once at the end of ``deserialize``. Subclasses that keep a cached
attribute derived from serialized members (and set it in ``__init__``)
override this to rebind it; the default is a no-op.
"""
def serialize(self) -> bytes:
"""Serialize this built algorithm to a self-contained serialization blob."""
w = codec.begin(type(self).__qualname__)
for attr, kind in self.__serialization_schema__:
kind.write(w, getattr(self, attr), self)
return w.getvalue()
@classmethod
def deserialize(cls: type[_S], blob: bytes) -> _S:
"""Reconstruct a built algorithm from a blob; no objects required.
Members are read in schema order and set on the instance as they are
read, so a ``CONDITIONAL`` member can consult a selector read earlier.
"""
r = codec.open(blob, cls.__qualname__)
obj = cls.__new__(cls)
# deserialize() bypasses __init__, which is where default-build
# wrappers bind their loaded result (see cache_build_results); an
# unbound wrapper resolves per call instead.
obj._bound_build_result = None
for attr, kind in cls.__serialization_schema__:
setattr(obj, attr, kind.read(r, obj))
obj._after_deserialize()
return obj

58
cccl_upstream/python/cuda_cccl/cuda/compute/_target_cc.py Normal file
View File

@@ -0,0 +1,58 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Target compute-capability context for device-code (LTO-IR) compilation.
Operators (numba) and iterators (NVRTC C++) are compiled to LTO-IR once and
linked into every per-arch build result of an algorithm. nvJitLink requires the final
target SM to be **at least as new** as every linked LTO/PTX input's arch. So for
a multi-arch build, this shared device code must be compiled for the **lowest**
target arch; otherwise linking a (say) sm_89 operator into an sm_80 cubin fails
with ``nvJitLink error``.
This module holds a context-local "target cc" that the leaf compilers
(``_jit`` for operators, ``_cpp_compile`` for iterators) consult. It is set
around a build by the caching decorator that wraps every ``make_<algo>``.
``None`` means "use the current device" — the default single-target behavior,
unchanged.
"""
from __future__ import annotations
import contextlib
import contextvars
# (major, minor) tuple, or None to mean "current device default".
_target_cc: contextvars.ContextVar = contextvars.ContextVar(
"cccl_target_cc", default=None
)
def get_target_cc():
"""The current build's target cc as ``(major, minor)``, or ``None``.
``None`` means device code should target the current device (the default).
"""
return _target_cc.get()
@contextlib.contextmanager
def target_cc(compute_capability):
"""Set the shared-device-code target cc for the duration of a build.
``compute_capability`` is the ``make_<algo>`` argument (``None`` / int /
``(major, minor)`` / list). For a multi-arch build the shared operator /
iterator LTO-IR is compiled for the **lowest** requested arch so it links
into every build result. ``None`` leaves the current-device default in place.
"""
from ._cccl_interop import normalize_compute_capabilities
ccs = normalize_compute_capabilities(compute_capability)
# normalized list is sorted ascending, so ccs[0] is the minimum target.
cc = ccs[0] if ccs else None
token = _target_cc.set(cc)
try:
yield
finally:
_target_cc.reset(token)

22
cccl_upstream/python/cuda_cccl/cuda/compute/_utils/__init__.py Normal file
View File

@@ -0,0 +1,22 @@
from __future__ import annotations
import re
__all__ = ["sanitize_identifier"]
def sanitize_identifier(name: str) -> str:
"""Sanitize a name to be a valid Python/LLVM identifier.
This replaces any character that isn't alphanumeric or underscore with
an underscore. This is needed because:
- Lambda functions have __name__ = "<lambda>" which contains angle brackets
- Python identifiers and LLVM/NVVM global names don't allow special characters
Args:
name: The name to sanitize (e.g., function __name__)
Returns:
A sanitized name safe for use as a Python identifier or LLVM symbol
"""
return re.sub(r"[^a-zA-Z0-9_]", "_", name)

168
cccl_upstream/python/cuda_cccl/cuda/compute/_utils/protocols.py Normal file
View File

@@ -0,0 +1,168 @@
# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
from typing import TYPE_CHECKING, List, Optional, Tuple
import numpy as np
"""
Utilities for extracting information from protocols such as `__cuda_array_interface__` and `__cuda_stream__`.
"""
if TYPE_CHECKING:
from ..typing import DeviceArrayLike, GpuStruct
def is_device_array(obj: object) -> bool:
"""Check if an object implements the `__cuda_array_interface__` protocol."""
return hasattr(obj, "__cuda_array_interface__")
def get_data_pointer(arr: DeviceArrayLike) -> int:
# TODO: these are fast paths for CuPy and PyTorch until
# we have a more general solution.
# Fast path for PyTorch (arr.data_ptr())
try:
return arr.data_ptr() # type: ignore
except AttributeError:
pass
# Fast path for CuPy (arr.data.ptr)
try:
return arr.data.ptr # type: ignore
except AttributeError:
pass
# Fall back to __cuda_array_interface__
return arr.__cuda_array_interface__["data"][0]
def get_dtype(arr: DeviceArrayLike | GpuStruct | np.ndarray) -> np.dtype:
# Try the fast path via .dtype attribute (works for np.ndarray, GpuStruct, and most device arrays)
try:
return np.dtype(arr.dtype) # type: ignore
except (AttributeError, TypeError):
pass
# Fall back to __cuda_array_interface__ for DeviceArrayLike
cai = arr.__cuda_array_interface__ # type: ignore
typestr = cai["typestr"]
if typestr.startswith("|V"):
# it's a structured dtype, use the descr field:
return np.dtype(cai["descr"])
else:
# a simple dtype, use the typestr field:
return np.dtype(typestr)
def get_shape(arr: DeviceArrayLike) -> Tuple[int]:
try:
# TODO: this is a fast path for CuPy until
# we have a more general solution.
return arr.shape # type: ignore
except AttributeError:
return arr.__cuda_array_interface__["shape"]
def get_size(arr: DeviceArrayLike) -> int:
"""Get the total number of elements in an array."""
# Try fast path via .size attribute
try:
return int(arr.size) # type: ignore
except AttributeError:
pass
# Fall back to computing from shape
shape = get_shape(arr)
import math
return math.prod(shape)
def is_contiguous(arr: DeviceArrayLike) -> bool:
cai = arr.__cuda_array_interface__
strides = cai["strides"]
if strides is None:
return True
shape = cai["shape"]
if any(dim == 0 for dim in shape):
# array has no elements
return True
if all(dim == 1 for dim in shape):
# there is a single element:
return True
itemsize = get_dtype(arr).itemsize
if strides[-1] == itemsize:
# assume C-contiguity
expected_stride = itemsize
for dim, stride in zip(reversed(shape), reversed(strides)):
if stride != expected_stride:
return False
expected_stride *= dim
return True
elif strides[0] == itemsize:
# assume F-contiguity
expected_stride = itemsize
for dim, stride in zip(shape, strides):
if stride != expected_stride:
return False
expected_stride *= dim
return True
else:
# not contiguous
return False
def compute_c_contiguous_strides_in_bytes(
shape: Tuple[int], itemsize: int
) -> Tuple[int, ...]:
"""Return C-contiguous strides in bytes for a given shape and itemsize (compatible with NumPy .strides)."""
strides: List[int] = []
acc = itemsize
for dim in reversed(shape):
strides.insert(0, acc)
acc *= dim
return tuple(strides)
def validate_and_get_stream(stream) -> Optional[int]:
# null stream is allowed
if stream is None:
return None
try:
stream_property = stream.__cuda_stream__()
except AttributeError as e:
raise TypeError(
f"stream argument {stream} does not implement the '__cuda_stream__' protocol"
) from e
try:
version, handle, *_ = stream_property
except (TypeError, ValueError) as e:
raise TypeError(
f"could not obtain __cuda_stream__ protocol version and handle from {stream_property}"
) from e
if version == 0:
if not isinstance(handle, int):
raise TypeError(f"invalid stream handle {handle}")
return handle
raise TypeError(f"unsupported __cuda_stream__ version {version}")

94
cccl_upstream/python/cuda_cccl/cuda/compute/_utils/temp_storage_buffer.py Normal file
View File

@@ -0,0 +1,94 @@
from __future__ import annotations
import functools
import weakref
from types import SimpleNamespace
from typing import Optional
from cuda.bindings import driver, runtime
try:
from cuda.core import Device
from cuda.core._utils.cuda_utils import handle_return
except ImportError:
from cuda.core.experimental import Device
from cuda.core.experimental._utils.cuda_utils import handle_return
from ..typing import StreamLike
@functools.cache
def _set_default_mempool_threshold(device_id: int):
"""
Set the release threshold for the default memory pool on this device,
if we haven't already done so. This prevents the driver from attempting
to shrink the pool after every sync, which can be slow.
"""
default_pool = handle_return(driver.cuDeviceGetDefaultMemPool(device_id))
threshold = handle_return(
driver.cuMemPoolGetAttribute(
default_pool, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD
)
)
if int(threshold) == 0:
handle_return(
driver.cuMemPoolSetAttribute(
default_pool,
driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
driver.cuuint64_t(0xFFFFFFFFFFFFFFFF),
)
)
def _finalize_buffer(ptr: int, stream_handle: Optional[int] = None):
"""Cleanup function for weakref finalizer."""
if ptr != 0:
try:
handle_return(runtime.cudaFreeAsync(ptr, stream_handle))
except Exception as e:
# Don't raise in finalizer, just print warning
print(f"Warning: Failed to free CUDA memory: {e}")
class TempStorageBuffer:
"""
Simple wrapper type around the memory allocation used for temporary storage,
exposing __cuda_array_interface__ and some other attributes for fast access.
This implementation uses cuda.bindings.runtime.cudaMallocAsync and
cudaFreeAsync for allocation and deallocation.
"""
def __init__(self, size: int, stream: Optional[StreamLike] = None):
# Get the current device
dev = Device()
stream_handle = stream.__cuda_stream__()[1] if stream is not None else None
# Set the release threshold for the default memory pool on this device
_set_default_mempool_threshold(dev.device_id)
# Allocate memory using cudaMallocAsync
device_ptr_int = handle_return(runtime.cudaMallocAsync(size, stream_handle))
self._ptr = int(device_ptr_int)
self._stream_handle = stream_handle
self._size = size
# attributes for fast path access in protocols.py
self.nbytes = size
self.data = SimpleNamespace(ptr=self._ptr)
# Set up weakref finalizer for cleanup
self._finalizer = weakref.finalize(
self, _finalize_buffer, self._ptr, self._stream_handle
)
@property
def __cuda_array_interface__(self):
return {
"data": (self._ptr, False),
"shape": (self._size,),
"strides": (1,),
"typestr": "|u1",
"version": 3,
}

76
cccl_upstream/python/cuda_cccl/cuda/compute/algorithms/__init__.py Normal file
View File

@@ -0,0 +1,76 @@
# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
from .._serialization import deserialize as deserialize
from .._serialization import serialize as serialize
from ._binary_search import lower_bound as lower_bound
from ._binary_search import make_lower_bound as make_lower_bound
from ._binary_search import make_upper_bound as make_upper_bound
from ._binary_search import upper_bound as upper_bound
from ._histogram import histogram_even as histogram_even
from ._histogram import make_histogram_even as make_histogram_even
from ._reduce import make_reduce_into as make_reduce_into
from ._reduce import reduce_into as reduce_into
from ._scan import exclusive_scan as exclusive_scan
from ._scan import inclusive_scan as inclusive_scan
from ._scan import make_exclusive_scan as make_exclusive_scan
from ._scan import make_inclusive_scan as make_inclusive_scan
from ._segmented_reduce import make_segmented_reduce as make_segmented_reduce
from ._segmented_reduce import segmented_reduce
from ._select import make_select as make_select
from ._select import select as select
from ._sort import DoubleBuffer, SortOrder
from ._sort import make_merge_sort as make_merge_sort
from ._sort import make_radix_sort as make_radix_sort
from ._sort import make_segmented_sort as make_segmented_sort
from ._sort import merge_sort as merge_sort
from ._sort import radix_sort as radix_sort
from ._sort import segmented_sort as segmented_sort
from ._three_way_partition import make_three_way_partition as make_three_way_partition
from ._three_way_partition import three_way_partition as three_way_partition
from ._transform import binary_transform, unary_transform
from ._transform import make_binary_transform as make_binary_transform
from ._transform import make_unary_transform as make_unary_transform
from ._unique_by_key import make_unique_by_key as make_unique_by_key
from ._unique_by_key import unique_by_key as unique_by_key
__all__ = [
"serialize",
"deserialize",
"reduce_into",
"make_reduce_into",
"lower_bound",
"make_lower_bound",
"upper_bound",
"make_upper_bound",
"inclusive_scan",
"make_inclusive_scan",
"exclusive_scan",
"make_exclusive_scan",
"unary_transform",
"make_unary_transform",
"binary_transform",
"make_binary_transform",
"histogram_even",
"make_histogram_even",
"merge_sort",
"make_merge_sort",
"radix_sort",
"make_radix_sort",
"segmented_reduce",
"make_segmented_reduce",
"unique_by_key",
"make_unique_by_key",
"segmented_sort",
"make_segmented_sort",
"three_way_partition",
"make_three_way_partition",
"select",
"make_select",
"DoubleBuffer",
"SortOrder",
]

345
cccl_upstream/python/cuda_cccl/cuda/compute/algorithms/_binary_search.py Normal file
View File

@@ -0,0 +1,345 @@
# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
from typing import ClassVar
import numpy as np
from .. import _bindings, types
from .. import _cccl_interop as cccl
from .._caching import cache_build_results, cache_with_registered_key_functions
from .._cccl_interop import set_cccl_iterator_state
from .._serialization import (
BUILD_RESULTS,
ITER,
OP,
Serializable,
)
from .._utils import protocols
from ..op import OpAdapter, OpKind, make_op_adapter
from ..typing import DeviceArrayLike, IteratorT, Operator
def _data_pointer_or_none(array) -> int | None:
# A ProxyArray is a build-time placeholder with no GPU allocation and thus no
# data pointer; return None for it (these pointers are only cache-key
# discriminators) so binary_search can be built without a GPU.
from .._proxy import is_proxy
return None if is_proxy(array) else protocols.get_data_pointer(array)
class _BinarySearch:
# Shared implementation for the lower/upper bound searchers.
_MODE: ClassVar[_bindings.BinarySearchMode]
__slots__ = [
"_bound_build_result",
"build_results",
"loaded_build_result",
"d_data_cccl",
"d_values_cccl",
"d_out_cccl",
"op_cccl",
"data_ptr",
"out_ptr",
]
__serialization_schema__ = (
("d_data_cccl", ITER),
("d_values_cccl", ITER),
("d_out_cccl", ITER),
("op_cccl", OP),
("build_results", BUILD_RESULTS(_bindings.DeviceBinarySearchBuildResult)),
)
def __init__(
self,
d_data: DeviceArrayLike,
d_values: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike,
comp: OpAdapter,
compute_capability=None,
):
if not protocols.is_device_array(d_data):
raise ValueError("d_data must be a device array for index outputs.")
if not protocols.is_device_array(d_out):
raise ValueError("d_out must be a device array for index outputs.")
out_dtype = protocols.get_dtype(d_out)
if out_dtype.kind != "u":
raise TypeError("d_out must use an unsigned integer dtype for indices.")
if out_dtype.itemsize != np.dtype(np.uintp).itemsize:
raise ValueError(
"d_out must use a pointer-sized unsigned integer dtype (np.uintp)."
)
self.data_ptr = _data_pointer_or_none(d_data)
self.out_ptr = _data_pointer_or_none(d_out)
self.d_data_cccl = cccl.to_cccl_input_iter(d_data)
self.d_values_cccl = cccl.to_cccl_input_iter(d_values)
data_value_type = cccl.get_value_type(d_data)
self.d_out_cccl = cccl.to_cccl_output_iter(d_out)
self.op_cccl = comp.compile((data_value_type, data_value_type), types.uint8)
self.build_results, self._bound_build_result = cache_build_results(
_bindings.DeviceBinarySearchBuildResult,
d_data,
d_values,
d_out,
comp,
self._MODE,
compute_capability=compute_capability,
builder=lambda: cccl.build_for_ccs(
_bindings.DeviceBinarySearchBuildResult,
self._MODE,
self.d_data_cccl,
self.d_values_cccl,
self.d_out_cccl,
self.op_cccl,
compute_capability=compute_capability,
),
)
def __call__(
self,
*,
d_data,
num_items: int,
d_values,
num_values: int,
d_out,
comp: Operator | None,
stream=None,
):
# Select (and lazily load) the build result for the current device.
self.loaded_build_result = cccl.resolve_build_result(
self.build_results, self._bound_build_result
)
set_cccl_iterator_state(self.d_data_cccl, d_data)
set_cccl_iterator_state(self.d_values_cccl, d_values)
set_cccl_iterator_state(self.d_out_cccl, d_out)
# Update op state for stateful ops
comp_adapter = make_op_adapter(OpKind.LESS if comp is None else comp)
self.op_cccl.state = comp_adapter.get_state()
stream_handle = protocols.validate_and_get_stream(stream)
self.loaded_build_result.compute(
self.d_data_cccl,
num_items,
self.d_values_cccl,
num_values,
self.d_out_cccl,
self.op_cccl,
stream_handle,
)
class _LowerBound(_BinarySearch, Serializable):
__slots__ = ()
_MODE = _bindings.BinarySearchMode.LOWER_BOUND
class _UpperBound(_BinarySearch, Serializable):
__slots__ = ()
_MODE = _bindings.BinarySearchMode.UPPER_BOUND
@cache_with_registered_key_functions
def _make_binary_search(
d_data: DeviceArrayLike,
d_values: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike,
comp: OpAdapter,
mode: _bindings.BinarySearchMode,
data_ptr: int,
out_ptr: int,
compute_capability=None,
):
"""Cached factory for the binary_search searchers."""
cls = _LowerBound if mode == _bindings.BinarySearchMode.LOWER_BOUND else _UpperBound
return cls(d_data, d_values, d_out, comp, compute_capability=compute_capability)
def make_lower_bound(
*,
d_data: DeviceArrayLike,
d_values: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike,
comp: Operator | None = None,
compute_capability=None,
):
"""
Create a lower_bound object that can be called to find insertion positions.
Example:
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/binary_search/lower_bound_object.py
:language: python
:start-after: # example-begin
Args:
d_data: Device array containing the sorted input range.
d_values: Device array or iterator containing the search values.
d_out: Device array to store the index results.
comp: Optional comparison operator (default: ``OpKind.LESS``).
compute_capability: Compute capability, or list of capabilities, to
build for ahead of time. Accepts a packed int (e.g. ``90``), a
``(major, minor)`` pair, a string (e.g. ``"9.0"``), or a list
thereof. When ``None`` (the default), the current device's
architecture is used.
Returns:
A callable object that performs lower_bound.
See Also:
:func:`lower_bound`
"""
comp_adapter = make_op_adapter(OpKind.LESS if comp is None else comp)
return _make_binary_search(
d_data,
d_values,
d_out,
comp_adapter,
_bindings.BinarySearchMode.LOWER_BOUND,
_data_pointer_or_none(d_data),
_data_pointer_or_none(d_out),
compute_capability=compute_capability,
)
def make_upper_bound(
*,
d_data: DeviceArrayLike,
d_values: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike,
comp: Operator | None = None,
compute_capability=None,
):
"""
Create an upper_bound object that can be called to find insertion positions.
Example:
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/binary_search/upper_bound_object.py
:language: python
:start-after: # example-begin
Args:
d_data: Device array containing the sorted input range.
d_values: Device array or iterator containing the search values.
d_out: Device array to store the index results.
comp: Optional comparison operator (default: ``OpKind.LESS``).
compute_capability: Compute capability, or list of capabilities, to
build for ahead of time. Accepts a packed int (e.g. ``90``), a
``(major, minor)`` pair, a string (e.g. ``"9.0"``), or a list
thereof. When ``None`` (the default), the current device's
architecture is used.
Returns:
A callable object that performs upper_bound.
See Also:
:func:`upper_bound`
"""
comp_adapter = make_op_adapter(OpKind.LESS if comp is None else comp)
return _make_binary_search(
d_data,
d_values,
d_out,
comp_adapter,
_bindings.BinarySearchMode.UPPER_BOUND,
_data_pointer_or_none(d_data),
_data_pointer_or_none(d_out),
compute_capability=compute_capability,
)
def lower_bound(
*,
d_data: DeviceArrayLike,
num_items: int,
d_values: DeviceArrayLike | IteratorT,
num_values: int,
d_out: DeviceArrayLike,
comp: Operator | None = None,
stream=None,
):
"""
Find the *first* position that each value in ``d_values`` would be inserted into
``d_data`` to maintain sorted order.
Example:
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/binary_search/lower_bound_basic.py
:language: python
:start-after: # example-begin
Args:
d_data: Device array containing the sorted input range.
num_items: Number of items in ``d_data``.
d_values: Device array or iterator containing the search values.
num_values: Number of items in ``d_values``.
d_out: Device array to store the index results.
comp: Optional comparison operator (default: ``OpKind.LESS``).
stream: CUDA stream for the operation (optional).
"""
searcher = make_lower_bound(
d_data=d_data, d_values=d_values, d_out=d_out, comp=comp
)
searcher(
d_data=d_data,
num_items=num_items,
d_values=d_values,
num_values=num_values,
d_out=d_out,
comp=comp,
stream=stream,
)
def upper_bound(
*,
d_data: DeviceArrayLike,
num_items: int,
d_values: DeviceArrayLike | IteratorT,
num_values: int,
d_out: DeviceArrayLike,
comp: Operator | None = None,
stream=None,
):
"""
Find the *last* position that each value in ``d_values`` would be inserted into
``d_data`` to maintain sorted order.
Example:
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/binary_search/upper_bound_basic.py
:language: python
:start-after: # example-begin
Args:
d_data: Device array containing the sorted input range.
num_items: Number of items in ``d_data``.
d_values: Device array or iterator containing the search values.
num_values: Number of items in ``d_values``.
d_out: Device array to store the index results.
comp: Optional comparison operator (default: ``OpKind.LESS``).
stream: CUDA stream for the operation (optional).
"""
searcher = make_upper_bound(
d_data=d_data, d_values=d_values, d_out=d_out, comp=comp
)
searcher(
d_data=d_data,
num_items=num_items,
d_values=d_values,
num_values=num_values,
d_out=d_out,
comp=comp,
stream=stream,
)

330
cccl_upstream/python/cuda_cccl/cuda/compute/algorithms/_histogram.py Normal file
View File

@@ -0,0 +1,330 @@
# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
import math
from typing import Union
import numpy as np
from .. import _bindings
from .. import _cccl_interop as cccl
from .._caching import cache_build_results, cache_with_registered_key_functions
from .._cccl_interop import set_cccl_iterator_state, to_cccl_value_state
from .._serialization import BUILD_RESULTS, ITER, U64, VALUE, Serializable
from .._utils.protocols import get_data_pointer, validate_and_get_stream
from .._utils.temp_storage_buffer import TempStorageBuffer
from ..typing import DeviceArrayLike, IteratorT
class _Histogram(Serializable):
__slots__ = [
"_bound_build_result",
"num_rows",
"d_samples_cccl",
"d_histogram_cccl",
"h_num_output_levels_cccl",
"h_lower_level_cccl",
"h_upper_level_cccl",
"build_results",
"loaded_build_result",
]
__serialization_schema__ = (
("num_rows", U64),
("d_samples_cccl", ITER),
("d_histogram_cccl", ITER),
("h_num_output_levels_cccl", VALUE),
("h_lower_level_cccl", VALUE),
("h_upper_level_cccl", VALUE),
("build_results", BUILD_RESULTS(_bindings.DeviceHistogramBuildResult)),
)
def __init__(
self,
d_samples: DeviceArrayLike | IteratorT,
d_histogram: DeviceArrayLike,
h_num_output_levels: np.ndarray,
h_lower_level: np.ndarray,
h_upper_level: np.ndarray,
num_samples: int,
compute_capability=None,
):
num_channels = 1
num_active_channels = 1
is_evenly_segmented = True
self.num_rows = 1
num_levels = h_num_output_levels[0]
row_stride_samples = num_samples
self.d_samples_cccl = cccl.to_cccl_input_iter(d_samples)
self.d_histogram_cccl = cccl.to_cccl_output_iter(d_histogram)
self.h_num_output_levels_cccl = cccl.to_cccl_value(h_num_output_levels)
self.h_lower_level_cccl = cccl.to_cccl_value(h_lower_level)
self.h_upper_level_cccl = cccl.to_cccl_value(h_upper_level)
self.build_results, self._bound_build_result = cache_build_results(
_bindings.DeviceHistogramBuildResult,
d_samples,
d_histogram,
int(num_levels),
h_lower_level.dtype,
num_samples,
is_evenly_segmented,
compute_capability=compute_capability,
builder=lambda: cccl.build_for_ccs(
_bindings.DeviceHistogramBuildResult,
num_channels,
num_active_channels,
self.d_samples_cccl,
num_levels,
self.d_histogram_cccl,
self.h_lower_level_cccl.type,
self.num_rows,
row_stride_samples,
is_evenly_segmented,
compute_capability=compute_capability,
),
)
def __call__(
self,
*,
temp_storage,
d_samples: DeviceArrayLike | IteratorT,
d_histogram: DeviceArrayLike,
h_num_output_levels: np.ndarray,
h_lower_level: np.ndarray,
h_upper_level: np.ndarray,
num_samples: int,
stream=None,
):
# Select (and lazily load) the build result for the current device.
self.loaded_build_result = cccl.resolve_build_result(
self.build_results, self._bound_build_result
)
set_cccl_iterator_state(self.d_samples_cccl, d_samples)
set_cccl_iterator_state(self.d_histogram_cccl, d_histogram)
self.h_num_output_levels_cccl.state = to_cccl_value_state(h_num_output_levels)
self.h_lower_level_cccl.state = to_cccl_value_state(h_lower_level)
self.h_upper_level_cccl.state = to_cccl_value_state(h_upper_level)
stream_handle = validate_and_get_stream(stream)
if temp_storage is None:
temp_storage_bytes = 0
d_temp_storage = 0
else:
temp_storage_bytes = temp_storage.nbytes
# Note: this is slightly slower, but supports all ndarray-like objects as long as they support CAI
# TODO: switch to use gpumemoryview once it's ready
d_temp_storage = get_data_pointer(temp_storage)
temp_storage_bytes = self.loaded_build_result.compute_even(
d_temp_storage,
temp_storage_bytes,
self.d_samples_cccl,
self.d_histogram_cccl,
self.h_num_output_levels_cccl,
self.h_lower_level_cccl,
self.h_upper_level_cccl,
num_samples,
self.num_rows,
num_samples,
stream_handle,
)
return temp_storage_bytes
@cache_with_registered_key_functions
def _make_histogram_even_impl(
d_samples: DeviceArrayLike | IteratorT,
d_histogram: DeviceArrayLike,
num_output_levels_val: int,
level_dtype,
uses_64bit_offset: bool,
uses_privatized_smem: bool,
compute_capability=None,
):
"""Internal cached implementation of make_histogram_even.
The uses_64bit_offset and uses_privatized_smem parameters ensure
kernels compiled for different offset and bin count regimes aren't reused.
"""
# Reconstruct the numpy arrays expected by _Histogram
h_num_output_levels = np.array([num_output_levels_val], dtype=np.int32)
# Bounds are runtime values. These placeholders only provide storage for
# cccl_value_t wrappers; build receives only the level type.
h_lower_level = np.zeros(1, dtype=level_dtype)
h_upper_level = np.ones(1, dtype=level_dtype)
# v1 only needs num_samples to select the generated offset type, so use a
# representative value for the requested offset-width regime.
if uses_64bit_offset:
sample_size = cccl.get_value_type(d_samples).size
int_max = np.iinfo(np.int32).max
# Smallest representative sample count that still selects long long
# offsets in v1's build-time offset type check.
build_num_samples = math.ceil(int_max / sample_size)
else:
build_num_samples = 1
return _Histogram(
d_samples,
d_histogram,
h_num_output_levels,
h_lower_level,
h_upper_level,
build_num_samples,
compute_capability=compute_capability,
)
def make_histogram_even(
*,
d_samples: DeviceArrayLike | IteratorT,
d_histogram: DeviceArrayLike,
h_num_output_levels: np.ndarray,
h_lower_level: np.ndarray,
h_upper_level: np.ndarray,
num_samples: int,
compute_capability=None,
):
"""Implements a device-wide histogram that places ``d_samples`` into evenly-spaced bins.
Example:
Below, ``make_histogram_even`` is used to create a histogram object that can be reused.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/histogram/histogram_object.py
:language: python
:start-after: # example-begin
Args:
d_samples: Device array or iterator containing the input samples to be histogrammed
d_histogram: Device array to store the histogram
h_num_output_levels: Host array containing the number of output levels
h_lower_level: Host array containing the lower level
h_upper_level: Host array containing the upper level
num_samples: Number of samples to be histogrammed
compute_capability: Compute capability, or list of capabilities, to
build for ahead of time. Accepts a packed int (e.g. ``90``), a
``(major, minor)`` pair, a string (e.g. ``"9.0"``), or a list
thereof. When ``None`` (the default), the current device's
architecture is used.
Returns:
A callable object that can be used to perform the histogram
"""
# Extract compile-relevant cache inputs from arrays.
num_output_levels_val = int(h_num_output_levels[0])
if h_lower_level.dtype != h_upper_level.dtype:
raise TypeError(
"h_lower_level and h_upper_level must have the same dtype; "
f"got {h_lower_level.dtype} and {h_upper_level.dtype}"
)
level_dtype = h_lower_level.dtype
# Mirrors v1 c/parallel/src/histogram.cu offset_cpp selection:
# (num_rows * row_stride_samples * sample_size) < INT_MAX selects int,
# otherwise long long. cuda.compute currently builds one-row histograms,
# so row_stride_samples is num_samples.
sample_size = cccl.get_value_type(d_samples).size
int_max = np.iinfo(np.int32).max
uses_64bit_offset = num_samples * sample_size >= int_max
# Mirrors CUB's even-histogram dispatch:
# detail::histogram::max_privatized_smem_bins is 256, and
# dispatch_histogram.cuh uses PRIVATIZED_SMEM_BINS=256 for <=256 bins
# and 0 for >256 bins.
num_bins = num_output_levels_val - 1
uses_privatized_smem = num_bins <= 256
# TODO: Once v2 is the default, remove uses_64bit_offset,
# num_output_levels_val, and uses_privatized_smem from this cache key;
# v2 passes row sizing and num_output_levels at runtime.
return _make_histogram_even_impl(
d_samples,
d_histogram,
num_output_levels_val,
level_dtype,
uses_64bit_offset,
uses_privatized_smem,
compute_capability=compute_capability,
)
def histogram_even(
*,
d_samples: DeviceArrayLike | IteratorT,
d_histogram: DeviceArrayLike,
num_output_levels: int,
lower_level: Union[np.floating, np.integer],
upper_level: Union[np.floating, np.integer],
num_samples: int,
stream=None,
):
"""
Performs device-wide histogram computation with evenly-spaced bins.
This function automatically handles temporary storage allocation and execution.
Example:
Below, ``histogram_even`` is used to compute a histogram with evenly-spaced bins.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/histogram/histogram_even_basic.py
:language: python
:start-after: # example-begin
:caption: Basic histogram example.
Args:
d_samples: Device array or iterator containing the input sequence of data samples
d_histogram: Device array to store the computed histogram
num_output_levels: Number of histogram bin levels (num_bins = num_output_levels - 1)
lower_level: Lower sample value bound (inclusive)
upper_level: Upper sample value bound (exclusive)
num_samples: Number of input samples
stream: CUDA stream for the operation (optional)
"""
# Histogram can accept multiple channels, with one value per channel for
# each of these parameters. The API only supports one channel for now but we
# pass arrays to make_histogram_even to support multiple channels in the
# future.
h_num_output_levels = np.array([num_output_levels], dtype=np.int32)
h_lower_level = np.array([lower_level], dtype=type(lower_level))
h_upper_level = np.array([upper_level], dtype=type(upper_level))
histogram = 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_storage_bytes = 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,
stream=stream,
)
temp_storage = TempStorageBuffer(temp_storage_bytes, stream)
histogram(
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,
stream=stream,
)

277
cccl_upstream/python/cuda_cccl/cuda/compute/algorithms/_reduce.py Normal file
View File

@@ -0,0 +1,277 @@
# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
from typing import Callable
import numpy as np
from .. import _bindings
from .. import _cccl_interop as cccl
from .._caching import cache_build_results, cache_with_registered_key_functions
from .._cccl_interop import (
get_value_type,
set_cccl_iterator_state,
to_cccl_value_state,
)
from .._serialization import BUILD_RESULTS, ITER, OP, VALUE, Serializable
from .._utils.protocols import get_data_pointer, get_dtype, validate_and_get_stream
from .._utils.temp_storage_buffer import TempStorageBuffer
from ..determinism import Determinism
from ..op import OpAdapter, make_op_adapter
from ..typing import (
DeviceArrayLike,
GpuStruct,
IteratorBase,
IteratorT,
Operator,
_Struct,
)
class _Reduce(Serializable):
__slots__ = [
"_bound_build_result",
"d_in_cccl",
"d_out_cccl",
"h_init_cccl",
"op_cccl",
"build_results",
"loaded_build_result",
"device_reduce_fn",
]
__serialization_schema__ = (
("d_in_cccl", ITER),
("d_out_cccl", ITER),
("op_cccl", OP),
("h_init_cccl", VALUE),
("build_results", BUILD_RESULTS(_bindings.DeviceReduceBuildResult)),
)
# TODO: constructor shouldn't require concrete `d_in`, `d_out`:
def __init__(
self,
d_in: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike | IteratorT,
op: OpAdapter,
h_init: np.ndarray | GpuStruct,
determinism: Determinism,
compute_capability=None,
):
self.d_in_cccl = cccl.to_cccl_input_iter(d_in)
self.d_out_cccl = cccl.to_cccl_output_iter(d_out)
self.h_init_cccl = cccl.to_cccl_value(h_init)
# Compile the op with value types
value_type = get_value_type(h_init)
self.op_cccl = op.compile((value_type, value_type), value_type)
# loaded_build_result / device_reduce_fn are bound lazily on the first
# __call__ (see _bind_device_reduce_fn).
self.build_results, self._bound_build_result = cache_build_results(
_bindings.DeviceReduceBuildResult,
d_in,
d_out,
op,
h_init,
determinism,
compute_capability=compute_capability,
builder=lambda: cccl.build_for_ccs(
_bindings.DeviceReduceBuildResult,
self.d_in_cccl,
self.d_out_cccl,
self.op_cccl,
self.h_init_cccl,
determinism,
compute_capability=compute_capability,
),
)
def _bind_device_reduce_fn(self) -> None:
# Derived from the loaded build result (not serialized); bound at __call__
# once resolve_build_result picks + loads the current device's build result.
if (
Determinism(self.loaded_build_result.determinism)
is Determinism.NOT_GUARANTEED
):
self.device_reduce_fn = self.loaded_build_result.compute_nondeterministic
else:
self.device_reduce_fn = self.loaded_build_result.compute
def __call__(
self,
*,
temp_storage,
d_in,
d_out,
num_items: int,
op: Callable | OpAdapter,
h_init: np.ndarray | GpuStruct,
stream=None,
):
# Select (and lazily load) the current device's build result, then bind the
# derived compute fn from it.
self.loaded_build_result = cccl.resolve_build_result(
self.build_results, self._bound_build_result
)
self._bind_device_reduce_fn()
set_cccl_iterator_state(self.d_in_cccl, d_in)
set_cccl_iterator_state(self.d_out_cccl, d_out)
# Update op state for stateful ops
op_adapter = make_op_adapter(op)
self.op_cccl.state = op_adapter.get_state()
self.h_init_cccl.state = to_cccl_value_state(h_init)
stream_handle = validate_and_get_stream(stream)
if temp_storage is None:
temp_storage_bytes = 0
d_temp_storage = 0
else:
temp_storage_bytes = temp_storage.nbytes
d_temp_storage = get_data_pointer(temp_storage)
temp_storage_bytes = self.device_reduce_fn(
d_temp_storage,
temp_storage_bytes,
self.d_in_cccl,
self.d_out_cccl,
num_items,
self.op_cccl,
self.h_init_cccl,
stream_handle,
)
return temp_storage_bytes
@cache_with_registered_key_functions
def make_reduce_into(
*,
d_in: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike | IteratorT,
op: Operator,
h_init: np.ndarray | GpuStruct,
**kwargs,
):
"""Computes a device-wide reduction using the specified binary ``op`` and initial value ``init``.
Example:
Below, ``make_reduce_into`` is used to create a reduction object that can be reused.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/reduction/reduce_object.py
:language: python
:start-after: # example-begin
Args:
d_in: Device array or iterator containing the input sequence of data items
d_out: Device array (of size 1) or iterator that will store the result of the reduction
op: Binary operator to apply.
The signature is ``(T, T) -> T``, where ``T`` is
the data type of the initial value ``h_init``.
init: Numpy array storing initial value of the reduction
compute_capability: Compute capability, or list of capabilities, to
build for ahead of time. Accepts a packed int (e.g. ``90``), a
``(major, minor)`` pair, a string (e.g. ``"9.0"``), or a list
thereof. When ``None`` (the default), the current device's
architecture is used.
Returns:
A callable object that can be used to perform the reduction
"""
try:
accum_dtype = get_dtype(h_init)
except (AttributeError, TypeError) as e:
raise TypeError(
"Could not determine accumulator dtype from h_init; "
"expected numpy array or object with .dtype"
) from e
# Validate d_in and d_out if they are device arrays (iterators may not expose
# dtype reliably here). Additionally, only require equality of dtypes for
# struct objects; mixed scalar dtypes (e.g. int8 input with int64 output)
# is acceptable
if isinstance(h_init, _Struct):
for arr, name in ((d_in, "input"), (d_out, "output")):
if isinstance(arr, IteratorBase):
continue
dtype = get_dtype(arr)
if dtype != accum_dtype:
raise TypeError(
f"reduce_into dtype mismatch: {name} dtype {dtype} != "
f"accumulator dtype {accum_dtype}. "
f"Ensure {name} elements and h_init have identical dtype to "
"avoid truncation or misinterpretation."
)
op_adapter = make_op_adapter(op)
return _Reduce(
d_in,
d_out,
op_adapter,
h_init,
kwargs.get("determinism", Determinism.RUN_TO_RUN),
compute_capability=kwargs.get("compute_capability"),
)
def reduce_into(
*,
d_in: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike | IteratorT,
num_items: int,
op: Operator,
h_init: np.ndarray | GpuStruct,
stream=None,
**kwargs,
):
"""
Performs device-wide reduction.
This function automatically handles temporary storage allocation and execution.
Example:
Below, ``reduce_into`` is used to compute the sum of a sequence of integers.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/reduction/sum_reduction.py
:language: python
:start-after: # example-begin
Args:
d_in: Device array or iterator containing the input sequence of data items
d_out: Device array or iterator to store the result of the reduction
num_items: Number of items to reduce
op: Binary operator to apply.
The signature is ``(T, T) -> T``, where ``T`` is
the data type of the initial value ``h_init``.
h_init: Initial value for the reduction
stream: CUDA stream for the operation (optional)
"""
reducer = make_reduce_into(d_in=d_in, d_out=d_out, op=op, h_init=h_init, **kwargs)
tmp_storage_bytes = reducer(
temp_storage=None,
d_in=d_in,
d_out=d_out,
num_items=num_items,
op=op,
h_init=h_init,
stream=stream,
)
tmp_storage = TempStorageBuffer(tmp_storage_bytes, stream)
reducer(
temp_storage=tmp_storage,
d_in=d_in,
d_out=d_out,
num_items=num_items,
op=op,
h_init=h_init,
stream=stream,
)

436
cccl_upstream/python/cuda_cccl/cuda/compute/algorithms/_scan.py Normal file
View File

@@ -0,0 +1,436 @@
# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
from typing import Callable, cast
import numpy as np
from .. import _bindings
from .. import _cccl_interop as cccl
from .._caching import cache_build_results, cache_with_registered_key_functions
from .._cccl_interop import (
get_value_type,
set_cccl_iterator_state,
to_cccl_value_state,
)
from .._serialization import (
BOOL,
BUILD_RESULTS,
CONDITIONAL,
ENUM,
ITER,
OP,
VALUE,
Serializable,
)
from .._utils.protocols import (
get_data_pointer,
is_device_array,
validate_and_get_stream,
)
from .._utils.temp_storage_buffer import TempStorageBuffer
from ..op import OpAdapter, make_op_adapter
from ..typing import DeviceArrayLike, GpuStruct, IteratorT, Operator
def get_init_kind(
init_value: np.ndarray | DeviceArrayLike | GpuStruct | None,
) -> _bindings.InitKind:
match init_value:
case None:
return _bindings.InitKind.NO_INIT
case _ if is_device_array(init_value):
return _bindings.InitKind.FUTURE_VALUE_INIT
case _:
return _bindings.InitKind.VALUE_INIT
class _Scan(Serializable):
__slots__ = [
"_bound_build_result",
"build_results",
"loaded_build_result",
"d_in_cccl",
"d_out_cccl",
"init_value_cccl",
"op_cccl",
"init_kind",
"force_inclusive",
"device_scan_fn",
]
__serialization_schema__ = (
("init_kind", ENUM(_bindings.InitKind)),
("force_inclusive", BOOL),
("d_in_cccl", ITER),
("d_out_cccl", ITER),
("op_cccl", OP),
(
"init_value_cccl",
CONDITIONAL(
"init_kind",
{
_bindings.InitKind.NO_INIT: None,
_bindings.InitKind.FUTURE_VALUE_INIT: ITER,
_bindings.InitKind.VALUE_INIT: VALUE,
},
),
),
("build_results", BUILD_RESULTS(_bindings.DeviceScanBuildResult)),
)
# TODO: constructor shouldn't require concrete `d_in`, `d_out`:
def __init__(
self,
d_in: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike | IteratorT,
op: OpAdapter,
init_value: np.ndarray | DeviceArrayLike | GpuStruct | None,
force_inclusive: bool,
compute_capability=None,
):
self.d_in_cccl = cccl.to_cccl_input_iter(d_in)
self.d_out_cccl = cccl.to_cccl_output_iter(d_out)
self.init_kind = get_init_kind(init_value)
self.init_value_cccl: _bindings.Iterator | _bindings.Value | None
match self.init_kind:
case _bindings.InitKind.NO_INIT:
self.init_value_cccl = None
value_type = get_value_type(d_in)
init_value_type_info = self.d_in_cccl.value_type
case _bindings.InitKind.FUTURE_VALUE_INIT:
self.init_value_cccl = cccl.to_cccl_input_iter(
cast(DeviceArrayLike, init_value)
)
value_type = get_value_type(cast(DeviceArrayLike, init_value))
init_value_type_info = self.init_value_cccl.value_type
case _bindings.InitKind.VALUE_INIT:
init_value_typed = cast(np.ndarray | GpuStruct, init_value)
self.init_value_cccl = cccl.to_cccl_value(init_value_typed)
value_type = get_value_type(init_value_typed)
init_value_type_info = self.init_value_cccl.type
self.force_inclusive = force_inclusive
# Compile the op with value types
self.op_cccl = op.compile((value_type, value_type), value_type)
# loaded_build_result / device_scan_fn are bound lazily on the first
# __call__ (see _bind_device_scan_fn).
self.build_results, self._bound_build_result = cache_build_results(
_bindings.DeviceScanBuildResult,
d_in,
d_out,
op,
init_value,
force_inclusive,
self.init_kind,
compute_capability=compute_capability,
builder=lambda: cccl.build_for_ccs(
_bindings.DeviceScanBuildResult,
self.d_in_cccl,
self.d_out_cccl,
self.op_cccl,
init_value_type_info,
force_inclusive,
self.init_kind,
compute_capability=compute_capability,
),
)
def _bind_device_scan_fn(self) -> None:
# Derived from force_inclusive + init_kind, from the loaded build result (not
# serialized as a function); bound at __call__ once resolve_build_result picks
# + loads the current device's build result.
match (self.force_inclusive, self.init_kind):
case (True, _bindings.InitKind.FUTURE_VALUE_INIT):
self.device_scan_fn = (
self.loaded_build_result.compute_inclusive_future_value
)
case (True, _bindings.InitKind.VALUE_INIT):
self.device_scan_fn = self.loaded_build_result.compute_inclusive
case (True, _bindings.InitKind.NO_INIT):
self.device_scan_fn = self.loaded_build_result.compute_inclusive_no_init
case (False, _bindings.InitKind.FUTURE_VALUE_INIT):
self.device_scan_fn = (
self.loaded_build_result.compute_exclusive_future_value
)
case (False, _bindings.InitKind.VALUE_INIT):
self.device_scan_fn = self.loaded_build_result.compute_exclusive
case (False, _bindings.InitKind.NO_INIT):
raise ValueError("Exclusive scan with No init value is not supported")
def __call__(
self,
*,
temp_storage,
d_in,
d_out,
op: Callable | OpAdapter,
init_value: np.ndarray | DeviceArrayLike | GpuStruct | None,
num_items: int,
stream=None,
):
# Select (and lazily load) the current device's build result, then bind the
# derived compute fn from it.
self.loaded_build_result = cccl.resolve_build_result(
self.build_results, self._bound_build_result
)
self._bind_device_scan_fn()
set_cccl_iterator_state(self.d_in_cccl, d_in)
set_cccl_iterator_state(self.d_out_cccl, d_out)
# Update op state for stateful ops
op_adapter = make_op_adapter(op)
self.op_cccl.state = op_adapter.get_state()
match self.init_kind:
case _bindings.InitKind.FUTURE_VALUE_INIT:
# We know that the init_value_cccl is an Iterator, so this cast
# tells MyPy what the actual type is. cast() is a no-op at runtime,
# which makes it better than isinstance() since this is a hot path
# and we have to minimize the work we do prior to calling the
# kernel.
self.init_value_cccl = cast(_bindings.Iterator, self.init_value_cccl)
set_cccl_iterator_state(self.init_value_cccl, init_value)
case _bindings.InitKind.VALUE_INIT:
self.init_value_cccl = cast(_bindings.Value, self.init_value_cccl)
self.init_value_cccl.state = to_cccl_value_state(
cast(np.ndarray | GpuStruct, init_value)
)
stream_handle = validate_and_get_stream(stream)
if temp_storage is None:
temp_storage_bytes = 0
d_temp_storage = 0
else:
temp_storage_bytes = temp_storage.nbytes
d_temp_storage = get_data_pointer(temp_storage)
temp_storage_bytes = self.device_scan_fn(
d_temp_storage,
temp_storage_bytes,
self.d_in_cccl,
self.d_out_cccl,
num_items,
self.op_cccl,
self.init_value_cccl,
stream_handle,
)
return temp_storage_bytes
# TODO Figure out `sum` without operator and initial value
# TODO Accept stream
@cache_with_registered_key_functions
def make_exclusive_scan(
*,
d_in: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike | IteratorT,
op: Operator,
init_value: np.ndarray | DeviceArrayLike | GpuStruct,
compute_capability=None,
):
"""Computes a device-wide scan using the specified binary ``op`` and initial value ``init``.
Example:
Below, ``make_exclusive_scan`` is used to create an exclusive scan object that can be reused.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/scan/exclusive_scan_object.py
:language: python
:start-after: # example-begin
Args:
d_in: Device array or iterator containing the input sequence of data items
d_out: Device array that will store the result of the scan
op: Binary scan operator.
The signature is ``(T, T) -> T``, where ``T`` is the data type of
the initial value ``init_value``.
init_value: Numpy array, device array, or GPU struct storing initial value of the scan
compute_capability: Compute capability, or list of capabilities, to
build for ahead of time. Accepts a packed int (e.g. ``90``), a
``(major, minor)`` pair, a string (e.g. ``"9.0"``), or a list
thereof. When ``None`` (the default), the current device's
architecture is used.
Returns:
A callable object that can be used to perform the scan
"""
op_adapter = make_op_adapter(op)
return _Scan(
d_in,
d_out,
op_adapter,
init_value,
False,
compute_capability=compute_capability,
)
def exclusive_scan(
*,
d_in: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike | IteratorT,
op: Operator,
init_value: np.ndarray | DeviceArrayLike | GpuStruct,
num_items: int,
stream=None,
):
"""
Performs device-wide exclusive scan.
This function automatically handles temporary storage allocation and execution.
Example:
Below, ``exclusive_scan`` is used to compute an exclusive scan with max operation.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/scan/exclusive_scan_max.py
:language: python
:start-after: # example-begin
Args:
d_in: Device array or iterator containing the input sequence of data items
d_out: Device array or iterator to store the result of the scan
op: Binary scan operator.
The signature is ``(T, T) -> T``, where ``T`` is the data type of
the initial value ``init_value``.
init_value: Initial value for the scan
num_items: Number of items to scan
stream: CUDA stream for the operation (optional)
"""
scanner = make_exclusive_scan(d_in=d_in, d_out=d_out, op=op, init_value=init_value)
tmp_storage_bytes = scanner(
temp_storage=None,
d_in=d_in,
d_out=d_out,
op=op,
init_value=init_value,
num_items=num_items,
stream=stream,
)
tmp_storage = TempStorageBuffer(tmp_storage_bytes, stream)
scanner(
temp_storage=tmp_storage,
d_in=d_in,
d_out=d_out,
op=op,
init_value=init_value,
num_items=num_items,
stream=stream,
)
# TODO Figure out `sum` without operator and initial value
# TODO Accept stream
@cache_with_registered_key_functions
def make_inclusive_scan(
*,
d_in: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike | IteratorT,
op: Operator,
init_value: np.ndarray | DeviceArrayLike | GpuStruct | None = None,
compute_capability=None,
):
"""Computes a device-wide scan using the specified binary ``op`` and initial value ``init``.
Example:
Below, ``make_inclusive_scan`` is used to create an inclusive scan object that can be reused.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/scan/inclusive_scan_object.py
:language: python
:start-after: # example-begin
Args:
d_in: Device array or iterator containing the input sequence of data items
d_out: Device array that will store the result of the scan
op: Binary scan operator.
The signature is ``(T, T) -> T``, where ``T`` is the data type of
the initial value ``init_value``.
init_value: Numpy array, device array, or GPU struct storing initial value of the scan, or None for no initial value
compute_capability: Compute capability, or list of capabilities, to
build for ahead of time. Accepts a packed int (e.g. ``90``), a
``(major, minor)`` pair, a string (e.g. ``"9.0"``), or a list
thereof. When ``None`` (the default), the current device's
architecture is used.
Returns:
A callable object that can be used to perform the scan
"""
op_adapter = make_op_adapter(op)
return _Scan(
d_in,
d_out,
op_adapter,
init_value,
True,
compute_capability=compute_capability,
)
def inclusive_scan(
*,
d_in: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike | IteratorT,
op: Operator,
init_value: np.ndarray | DeviceArrayLike | GpuStruct | None = None,
num_items: int,
stream=None,
):
"""
Performs device-wide inclusive scan.
This function automatically handles temporary storage allocation and execution.
Example:
Below, ``inclusive_scan`` is used to compute an inclusive scan (prefix sum).
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/scan/inclusive_scan_custom.py
:language: python
:start-after: # example-begin
Args:
d_in: Device array or iterator containing the input sequence of data items
d_out: Device array or iterator to store the result of the scan
op: Binary scan operator.
The signature is ``(T, T) -> T``, where ``T`` is the data type of
the initial value ``init_value``.
init_value: Initial value for the scan
num_items: Number of items to scan
stream: CUDA stream for the operation (optional)
"""
scanner = make_inclusive_scan(d_in=d_in, d_out=d_out, op=op, init_value=init_value)
tmp_storage_bytes = scanner(
temp_storage=None,
d_in=d_in,
d_out=d_out,
op=op,
init_value=init_value,
num_items=num_items,
stream=stream,
)
tmp_storage = TempStorageBuffer(tmp_storage_bytes, stream)
scanner(
temp_storage=tmp_storage,
d_in=d_in,
d_out=d_out,
op=op,
init_value=init_value,
num_items=num_items,
stream=stream,
)

293
cccl_upstream/python/cuda_cccl/cuda/compute/algorithms/_segmented_reduce.py Normal file
View File

@@ -0,0 +1,293 @@
# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
from typing import Callable
import numpy as np
from .. import _bindings
from .. import _cccl_interop as cccl
from .._caching import cache_build_results, cache_with_registered_key_functions
from .._cccl_interop import (
get_value_type,
set_cccl_iterator_state,
to_cccl_value_state,
)
from .._serialization import BUILD_RESULTS, ITER, OP, VALUE, Serializable
from .._utils.protocols import (
get_data_pointer,
validate_and_get_stream,
)
from .._utils.temp_storage_buffer import TempStorageBuffer
from ..op import OpAdapter, make_op_adapter
from ..typing import DeviceArrayLike, GpuStruct, IteratorT, Operator
class _SegmentedReduce(Serializable):
__slots__ = [
"_bound_build_result",
"build_results",
"loaded_build_result",
"d_in_cccl",
"d_out_cccl",
"start_offsets_in_cccl",
"end_offsets_in_cccl",
"h_init_cccl",
"op_cccl",
]
__serialization_schema__ = (
("d_in_cccl", ITER),
("d_out_cccl", ITER),
("start_offsets_in_cccl", ITER),
("end_offsets_in_cccl", ITER),
("h_init_cccl", VALUE),
("op_cccl", OP),
("build_results", BUILD_RESULTS(_bindings.DeviceSegmentedReduceBuildResult)),
)
def __init__(
self,
d_in: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike | IteratorT,
start_offsets_in: DeviceArrayLike | IteratorT,
end_offsets_in: DeviceArrayLike | IteratorT,
op: OpAdapter,
h_init: np.ndarray | GpuStruct,
compute_capability=None,
):
self.d_in_cccl = cccl.to_cccl_input_iter(d_in)
self.d_out_cccl = cccl.to_cccl_output_iter(d_out)
self.start_offsets_in_cccl = cccl.to_cccl_input_iter(start_offsets_in)
self.end_offsets_in_cccl = cccl.to_cccl_input_iter(end_offsets_in)
self.h_init_cccl = cccl.to_cccl_value(h_init)
# Compile the op with value types
value_type = get_value_type(h_init)
self.op_cccl = op.compile((value_type, value_type), value_type)
self.build_results, self._bound_build_result = cache_build_results(
_bindings.DeviceSegmentedReduceBuildResult,
d_in,
d_out,
start_offsets_in,
end_offsets_in,
op,
h_init,
compute_capability=compute_capability,
builder=lambda: cccl.build_for_ccs(
_bindings.DeviceSegmentedReduceBuildResult,
self.d_in_cccl,
self.d_out_cccl,
self.start_offsets_in_cccl,
self.end_offsets_in_cccl,
self.op_cccl,
self.h_init_cccl,
compute_capability=compute_capability,
),
)
def __call__(
self,
*,
temp_storage,
d_in,
d_out,
num_segments: int,
start_offsets_in,
end_offsets_in,
op: Callable | OpAdapter,
h_init,
max_segment_size: int | None = None,
stream=None,
):
# Select (and lazily load) the build result for the current device.
self.loaded_build_result = cccl.resolve_build_result(
self.build_results, self._bound_build_result
)
if num_segments > np.iinfo(np.int32).max:
raise RuntimeError(
"Segmented sort does not currently support more than 2^31-1 segments."
)
if max_segment_size is None:
max_segment_size = 0 # CCCL.c treats 0 as "not specified"
if max_segment_size > 0:
try:
from .._build_info import USING_V2 # type: ignore[import-not-found]
except ImportError:
USING_V2 = False
if USING_V2:
import warnings
warnings.warn(
"max_segment_size is not used by the v2 backend and will be ignored",
stacklevel=4,
)
set_cccl_iterator_state(self.d_in_cccl, d_in)
set_cccl_iterator_state(self.d_out_cccl, d_out)
set_cccl_iterator_state(self.start_offsets_in_cccl, start_offsets_in)
set_cccl_iterator_state(self.end_offsets_in_cccl, end_offsets_in)
op_adapter = make_op_adapter(op)
self.op_cccl.state = op_adapter.get_state()
self.h_init_cccl.state = to_cccl_value_state(h_init)
stream_handle = validate_and_get_stream(stream)
if temp_storage is None:
temp_storage_bytes = 0
d_temp_storage = 0
else:
temp_storage_bytes = temp_storage.nbytes
d_temp_storage = get_data_pointer(temp_storage)
temp_storage_bytes = self.loaded_build_result.compute(
d_temp_storage,
temp_storage_bytes,
self.d_in_cccl,
self.d_out_cccl,
num_segments,
self.start_offsets_in_cccl,
self.end_offsets_in_cccl,
self.op_cccl,
self.h_init_cccl,
max_segment_size,
stream_handle,
)
return temp_storage_bytes
@cache_with_registered_key_functions
def make_segmented_reduce(
*,
d_in: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike | IteratorT,
start_offsets_in: DeviceArrayLike | IteratorT,
end_offsets_in: DeviceArrayLike | IteratorT,
op: Operator,
h_init: np.ndarray | GpuStruct,
compute_capability=None,
):
"""Computes a device-wide segmented reduction using the specified binary ``op`` and initial value ``init``.
Example:
Below, ``make_segmented_reduce`` is used to create a segmented reduction object that can be reused.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/segmented/segmented_reduce_object.py
:language: python
:start-after: # example-begin
Args:
d_in: Device array or iterator containing the input sequence of data items
d_out: Device array that will store the result of the reduction
start_offsets_in: Device array or iterator containing offsets to start of segments
end_offsets_in: Device array or iterator containing offsets to end of segments
op: Binary operator to apply.
The signature is ``(T, T) -> T``, where ``T`` is
the data type of the initial value ``h_init``.
init: Numpy array storing initial value of the reduction
compute_capability: Compute capability, or list of capabilities, to
build for ahead of time. Accepts a packed int (e.g. ``90``), a
``(major, minor)`` pair, a string (e.g. ``"9.0"``), or a list
thereof. When ``None`` (the default), the current device's
architecture is used.
Returns:
A callable object that can be used to perform the reduction
"""
op_adapter = make_op_adapter(op)
return _SegmentedReduce(
d_in,
d_out,
start_offsets_in,
end_offsets_in,
op_adapter,
h_init,
compute_capability=compute_capability,
)
def segmented_reduce(
*,
d_in: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike | IteratorT,
num_segments: int,
start_offsets_in: DeviceArrayLike | IteratorT,
end_offsets_in: DeviceArrayLike | IteratorT,
op: Operator,
h_init: np.ndarray | GpuStruct,
max_segment_size: int | None = None,
stream=None,
):
"""
Performs device-wide segmented reduction.
This function automatically handles temporary storage allocation and execution.
Example:
Below, ``segmented_reduce`` is used to compute the minimum value of segments in a sequence of integers.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/segmented/segmented_reduce_basic.py
:language: python
:start-after: # example-begin
Args:
d_in: Device array or iterator containing the input sequence of data items
d_out: Device array to store the result of the reduction for each segment
num_segments: Number of segments to reduce
start_offsets_in: Device array or iterator containing the sequence of beginning offsets
end_offsets_in: Device array or iterator containing the sequence of ending offsets
op: Binary operator to apply.
The signature is ``(T, T) -> T``, where ``T`` is
the data type of the initial value ``h_init``.
h_init: Initial value for the reduction
max_segment_size: The number of elements in the largest segment (optional)
If provided, this information is used to dispatch to the
optimal kernel for best performance.
stream: CUDA stream for the operation (optional)
"""
reducer = make_segmented_reduce(
d_in=d_in,
d_out=d_out,
start_offsets_in=start_offsets_in,
end_offsets_in=end_offsets_in,
op=op,
h_init=h_init,
)
tmp_storage_bytes = reducer(
temp_storage=None,
d_in=d_in,
d_out=d_out,
num_segments=num_segments,
start_offsets_in=start_offsets_in,
end_offsets_in=end_offsets_in,
op=op,
h_init=h_init,
max_segment_size=max_segment_size,
stream=stream,
)
tmp_storage = TempStorageBuffer(tmp_storage_bytes, stream)
reducer(
temp_storage=tmp_storage,
d_in=d_in,
d_out=d_out,
num_segments=num_segments,
start_offsets_in=start_offsets_in,
end_offsets_in=end_offsets_in,
op=op,
h_init=h_init,
max_segment_size=max_segment_size,
stream=stream,
)

239
cccl_upstream/python/cuda_cccl/cuda/compute/algorithms/_select.py Normal file
View File

@@ -0,0 +1,239 @@
# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
from functools import cache
from .._caching import cache_with_registered_key_functions
from .._cpp_compile import compile_cpp_op_code
from .._serialization import NESTED, Serializable
from .._utils.temp_storage_buffer import TempStorageBuffer
from ..iterators import DiscardIterator
from ..op import OpAdapter, RawOp, make_op_adapter
from ..typing import DeviceArrayLike, IteratorT, Operator
from ._three_way_partition import _ThreeWayPartition, make_three_way_partition
@cache
def _always_false_op(_target_cc):
# ``_target_cc`` (the build's get_target_cc()) is part of the cache key so the
# predicate's LTO-IR is recompiled per target arch: this RawOp is linked into
# the three-way-partition build, and nvJitLink rejects a newer-arch input in
# an older-arch result. Without the key, the first build's arch would leak
# into every later build (module-global cache). compile_cpp_op_code() reads
# the same target internally; the arg only distinguishes cache entries.
source = """
extern "C" __device__ void always_false(void*, void* result) {{
*static_cast<bool*>(result) = false;
}}
"""
code = compile_cpp_op_code(source)
return RawOp(ltoir=code, name="always_false")
def _get_always_false_op():
"""The always-false predicate compiled for the current build's target cc."""
from .._target_cc import get_target_cc
return _always_false_op(get_target_cc())
class _Select(Serializable):
__slots__ = ["_bound_build_result", "partitioner", "always_false_op", "_discards"]
__serialization_schema__ = (("partitioner", NESTED(_ThreeWayPartition)),)
def __init__(
self,
d_in: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike | IteratorT,
d_num_selected_out: DeviceArrayLike,
cond: OpAdapter,
compute_capability=None,
):
self.always_false_op = _get_always_false_op()
d_second, d_unselected = self._discard_iterators(d_out)
self.partitioner = make_three_way_partition(
d_in=d_in,
d_first_part_out=d_out,
d_second_part_out=d_second,
d_unselected_out=d_unselected,
d_num_selected_out=d_num_selected_out,
select_first_part_op=cond,
select_second_part_op=self.always_false_op,
compute_capability=compute_capability,
)
def _discard_iterators(self, d_out):
# The second/unselected outputs are discarded; their iterators depend
# only on d_out's type, so build the pair once and cache it. Bound
# lazily (on first construction or first call) so a deserialized
# _Select, which has no construction d_out, builds them on first use.
try:
return self._discards
except AttributeError:
self._discards = (DiscardIterator(d_out), DiscardIterator(d_out))
return self._discards
def _after_deserialize(self) -> None:
# always_false_op (the always-false second predicate) is not serialized.
# Its compiled LTO-IR is already baked into the (serialized) three-way
# partition build result, and __call__ reads only this op's runtime state
# (which is empty — the predicate is stateless). So reconstruct an
# empty-state stand-in WITHOUT compiling: deserialize() must neither
# recompile nor require a GPU, and calling _get_always_false_op() here
# would do both (cold cache -> compile_cpp_op_code -> Device() fallback).
self.always_false_op = RawOp(ltoir=b"", name="always_false")
def __call__(
self,
*,
temp_storage,
d_in,
d_out,
d_num_selected_out,
cond,
num_items: int,
stream=None,
):
d_second, d_unselected = self._discard_iterators(d_out)
return self.partitioner(
temp_storage=temp_storage,
d_in=d_in,
d_first_part_out=d_out,
d_second_part_out=d_second,
d_unselected_out=d_unselected,
d_num_selected_out=d_num_selected_out,
select_first_part_op=make_op_adapter(cond),
select_second_part_op=self.always_false_op,
num_items=num_items,
stream=stream,
)
@cache_with_registered_key_functions
def make_select(
*,
d_in: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike | IteratorT,
d_num_selected_out: DeviceArrayLike,
cond: Operator,
compute_capability=None,
):
"""
Create a select object that can be called to select elements matching a condition.
This is the object-oriented API that allows explicit control over temporary
storage allocation. For simpler usage, consider using :func:`select`.
Example:
Below, ``make_select`` is used to create a select object that can be reused.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/select/select_object.py
:language: python
:start-after: # example-begin
Args:
d_in: Device array or iterator containing the input sequence of data items.
d_out: Device array or iterator to store the selected output items.
d_num_selected_out: Device array to store the number of items that passed the selection.
The count is stored in ``d_num_selected_out[0]``.
cond: Selection condition (predicate).
The signature is ``(T) -> uint8``, where ``T`` is the input data type.
Returns 1 (selected) or 0 (not selected).
compute_capability: Compute capability, or list of capabilities, to
build for ahead of time. Accepts a packed int (e.g. ``90``), a
``(major, minor)`` pair, a string (e.g. ``"9.0"``), or a list
thereof. When ``None`` (the default), the current device's
architecture is used.
Returns:
A callable object that performs the selection operation.
"""
cond_adapter = make_op_adapter(cond)
# Note: _Select internally calls make_three_way_partition which will
# normalize the cond. But we've already normalized it, so the Op
# will be passed through make_op unchanged.
return _Select(
d_in,
d_out,
d_num_selected_out,
cond_adapter,
compute_capability=compute_capability,
)
def select(
*,
d_in: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike | IteratorT,
d_num_selected_out: DeviceArrayLike,
cond: Operator,
num_items: int,
stream=None,
):
"""
Performs device-wide selection of elements based on a condition.
Given an input sequence, this function selects all elements for which the condition
function ``cond`` returns true (non-zero) and writes them to the output in a
compacted form. The number of selected elements is written to ``d_num_selected_out[0]``.
This function automatically handles temporary storage allocation and execution.
The ``cond`` function can reference device arrays as globals or closures - they will
be automatically captured as state arrays, enabling stateful operations like counting.
Example:
Below, ``select`` is used to select even numbers from an input array:
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/select/select_basic.py
:language: python
:start-after: # example-begin
You can also use iterators for more complex selection patterns:
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/select/select_with_iterator.py
:language: python
:start-after: # example-begin
Args:
d_in: Device array or iterator containing the input sequence of data items.
d_out: Device array or iterator to store the selected output items.
d_num_selected_out: Device array to store the number of items that passed the selection.
The count is stored in ``d_num_selected_out[0]``.
cond: Selection condition (predicate).
The signature is ``(T) -> uint8``, where ``T`` is the input data type.
Returns 1 (selected) or 0 (not selected).
Can reference device arrays as globals/closures - they will be automatically captured.
num_items: Number of items in the input sequence.
stream: CUDA stream to use for the operation (optional).
"""
# Create adapter to support stateful ops
cond_adapter = make_op_adapter(cond)
selector = make_select(
d_in=d_in, d_out=d_out, d_num_selected_out=d_num_selected_out, cond=cond_adapter
)
tmp_storage_bytes = selector(
temp_storage=None,
d_in=d_in,
d_out=d_out,
d_num_selected_out=d_num_selected_out,
cond=cond_adapter,
num_items=num_items,
stream=stream,
)
tmp_storage = TempStorageBuffer(tmp_storage_bytes, stream)
selector(
temp_storage=tmp_storage,
d_in=d_in,
d_out=d_out,
d_num_selected_out=d_num_selected_out,
cond=cond_adapter,
num_items=num_items,
stream=stream,
)

25
cccl_upstream/python/cuda_cccl/cuda/compute/algorithms/_sort/__init__.py Normal file
View File

@@ -0,0 +1,25 @@
# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
from ._merge_sort import make_merge_sort as make_merge_sort
from ._merge_sort import merge_sort as merge_sort
from ._radix_sort import make_radix_sort as make_radix_sort
from ._radix_sort import radix_sort as radix_sort
from ._segmented_sort import make_segmented_sort as make_segmented_sort
from ._segmented_sort import segmented_sort as segmented_sort
from ._sort_common import DoubleBuffer, SortOrder
__all__ = [
"make_merge_sort",
"merge_sort",
"make_radix_sort",
"radix_sort",
"make_segmented_sort",
"segmented_sort",
"DoubleBuffer",
"SortOrder",
]

261
cccl_upstream/python/cuda_cccl/cuda/compute/algorithms/_sort/_merge_sort.py Normal file
View File

@@ -0,0 +1,261 @@
# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
from ... import _bindings, types
from ... import _cccl_interop as cccl
from ..._caching import cache_build_results, cache_with_registered_key_functions
from ..._cccl_interop import set_cccl_iterator_state
from ..._serialization import BUILD_RESULTS, ITER, OP, Serializable
from ..._utils.protocols import (
get_data_pointer,
validate_and_get_stream,
)
from ..._utils.temp_storage_buffer import TempStorageBuffer
from ...op import OpAdapter, make_op_adapter
from ...typing import DeviceArrayLike, IteratorT, Operator
class _MergeSort(Serializable):
__slots__ = [
"_bound_build_result",
"d_in_keys_cccl",
"d_in_values_cccl",
"d_out_keys_cccl",
"d_out_values_cccl",
"op_adapter",
"op_cccl",
"build_results",
"loaded_build_result",
]
__serialization_schema__ = (
("d_in_keys_cccl", ITER),
("d_in_values_cccl", ITER),
("d_out_keys_cccl", ITER),
("d_out_values_cccl", ITER),
("op_cccl", OP),
("build_results", BUILD_RESULTS(_bindings.DeviceMergeSortBuildResult)),
)
def __init__(
self,
d_in_keys: DeviceArrayLike | IteratorT,
d_in_values: DeviceArrayLike | IteratorT | None,
d_out_keys: DeviceArrayLike,
d_out_values: DeviceArrayLike | None,
op: OpAdapter,
compute_capability=None,
):
present_in_values = d_in_values is not None
present_out_values = d_out_values is not None
assert present_in_values == present_out_values
self.d_in_keys_cccl = cccl.to_cccl_input_iter(d_in_keys)
self.d_in_values_cccl = cccl.to_cccl_input_iter(d_in_values)
self.d_out_keys_cccl = cccl.to_cccl_output_iter(d_out_keys)
self.d_out_values_cccl = cccl.to_cccl_output_iter(d_out_values)
self.op_adapter = op
# Compile the op - merge_sort expects int8 return (comparison)
value_type = cccl.get_value_type(d_in_keys)
self.op_cccl = op.compile((value_type, value_type), types.int8)
self.build_results, self._bound_build_result = cache_build_results(
_bindings.DeviceMergeSortBuildResult,
d_in_keys,
d_in_values,
d_out_keys,
d_out_values,
op,
compute_capability=compute_capability,
builder=lambda: cccl.build_for_ccs(
_bindings.DeviceMergeSortBuildResult,
self.d_in_keys_cccl,
self.d_in_values_cccl,
self.d_out_keys_cccl,
self.d_out_values_cccl,
self.op_cccl,
compute_capability=compute_capability,
),
)
def __call__(
self,
*,
temp_storage,
d_in_keys: DeviceArrayLike | IteratorT,
d_in_values: DeviceArrayLike | IteratorT | None,
d_out_keys: DeviceArrayLike,
d_out_values: DeviceArrayLike | None,
num_items: int,
op: Operator,
stream=None,
):
# Select (and lazily load) the build result for the current device.
self.loaded_build_result = cccl.resolve_build_result(
self.build_results, self._bound_build_result
)
present_in_values = d_in_values is not None
present_out_values = d_out_values is not None
assert present_in_values == present_out_values
set_cccl_iterator_state(self.d_in_keys_cccl, d_in_keys)
if present_in_values:
set_cccl_iterator_state(self.d_in_values_cccl, d_in_values)
set_cccl_iterator_state(self.d_out_keys_cccl, d_out_keys)
if present_out_values:
set_cccl_iterator_state(self.d_out_values_cccl, d_out_values)
op_adapter = make_op_adapter(op)
self.op_cccl.state = op_adapter.get_state()
stream_handle = validate_and_get_stream(stream)
if temp_storage is None:
temp_storage_bytes = 0
d_temp_storage = 0
else:
temp_storage_bytes = temp_storage.nbytes
# Note: this is slightly slower, but supports all ndarray-like objects as long as they support CAI
# TODO: switch to use gpumemoryview once it's ready
d_temp_storage = get_data_pointer(temp_storage)
temp_storage_bytes = self.loaded_build_result.compute(
d_temp_storage,
temp_storage_bytes,
self.d_in_keys_cccl,
self.d_in_values_cccl,
self.d_out_keys_cccl,
self.d_out_values_cccl,
num_items,
self.op_cccl,
stream_handle,
)
return temp_storage_bytes
@cache_with_registered_key_functions
def make_merge_sort(
*,
d_in_keys: DeviceArrayLike | IteratorT,
d_in_values: DeviceArrayLike | IteratorT | None = None,
d_out_keys: DeviceArrayLike,
d_out_values: DeviceArrayLike | None = None,
op: Operator,
compute_capability=None,
):
"""Implements a device-wide merge sort using ``d_in_keys`` and the comparison operator ``op``.
Example:
Below, ``make_merge_sort`` is used to create a merge sort object that can be reused.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/sort/merge_sort_object.py
:language: python
:start-after: # example-begin
Args:
d_in_keys: Device array or iterator containing the input keys to be sorted
d_in_values: Optional device array or iterator that contains each key's corresponding value
d_out_keys: Device array to store the sorted keys
d_out_values: Device array to store the sorted values
op: The comparison operator for sorting. The signature is ``(T, T) -> int8``, where ``T`` is the input data type. See notes below.
compute_capability: Compute capability, or list of capabilities, to
build for ahead of time. Accepts a packed int (e.g. ``90``), a
``(major, minor)`` pair, a string (e.g. ``"9.0"``), or a list
thereof. When ``None`` (the default), the current device's
architecture is used.
Returns:
A callable object that can be used to perform the merge sort
.. important::
The provided comparison operator must follow `strict weak ordering <https://en.cppreference.com/w/cpp/concepts/strict_weak_order.html>`_
semantics. For example, the comparator ``lambda lhs, rhs: lhs < rhs`` follows strict weak ordering, but the comparator
``lambda lhs, rhs: rhs >= lhs`` does not, because it is reflexive: ``r(x, x) == True``. Providing a comparator that does not
follow the required semantics can lead to incorrect results, silent memory corruption, or crashes.
"""
op_adapter = make_op_adapter(op)
return _MergeSort(
d_in_keys,
d_in_values,
d_out_keys,
d_out_values,
op_adapter,
compute_capability=compute_capability,
)
def merge_sort(
*,
d_in_keys: DeviceArrayLike | IteratorT,
d_in_values: DeviceArrayLike | IteratorT | None = None,
d_out_keys: DeviceArrayLike,
d_out_values: DeviceArrayLike | None = None,
num_items: int,
op: Operator,
stream=None,
):
"""
Performs device-wide merge sort.
This function automatically handles temporary storage allocation and execution.
Example:
Below, ``merge_sort`` is used to sort a sequence of keys inplace. It also rearranges the values according to the keys' order.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/sort/merge_sort_basic.py
:language: python
:start-after: # example-begin
Args:
d_in_keys: Device array or iterator containing the input sequence of keys
d_in_values: Device array or iterator containing the input sequence of values (optional)
d_out_keys: Device array to store the sorted keys
d_out_values: Device array to store the sorted values (optional)
num_items: Number of items to sort
op: The comparison operator for sorting. The signature is ``(T, T) -> int8``, where ``T`` is the input data type. See notes below.
stream: CUDA stream for the operation (optional)
.. important::
The provided comparison operator must follow `strict weak ordering <https://en.cppreference.com/w/cpp/concepts/strict_weak_order.html>`_
semantics. For example, the comparator ``lambda lhs, rhs: lhs < rhs`` follows strict weak ordering, but the comparator
``lambda lhs, rhs: rhs >= lhs`` does not, because it is reflexive: ``r(x, x) == True``. Providing a comparator that does not
follow the required semantics can lead to incorrect results, silent memory corruption, or crashes.
"""
sorter = 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=op,
)
tmp_storage_bytes = 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,
stream=stream,
)
tmp_storage = TempStorageBuffer(tmp_storage_bytes, stream)
sorter(
temp_storage=tmp_storage,
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,
stream=stream,
)

289
cccl_upstream/python/cuda_cccl/cuda/compute/algorithms/_sort/_radix_sort.py Normal file
View File

@@ -0,0 +1,289 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
from ... import _bindings
from ... import _cccl_interop as cccl
from ..._caching import cache_build_results, cache_with_registered_key_functions
from ..._cccl_interop import set_cccl_iterator_state
from ..._serialization import BUILD_RESULTS, ITER, OP, Serializable
from ..._utils.protocols import (
get_data_pointer,
get_dtype,
validate_and_get_stream,
)
from ..._utils.temp_storage_buffer import TempStorageBuffer
from ...typing import DeviceArrayLike
from ._sort_common import DoubleBuffer, SortOrder, _get_arrays
class _RadixSort(Serializable):
__slots__ = [
"_bound_build_result",
"d_in_keys_cccl",
"d_out_keys_cccl",
"d_in_values_cccl",
"d_out_values_cccl",
"decomposer_op",
"build_results",
"loaded_build_result",
]
__serialization_schema__ = (
("d_in_keys_cccl", ITER),
("d_out_keys_cccl", ITER),
("d_in_values_cccl", ITER),
("d_out_values_cccl", ITER),
("decomposer_op", OP),
("build_results", BUILD_RESULTS(_bindings.DeviceRadixSortBuildResult)),
)
def __init__(
self,
d_in_keys: DeviceArrayLike | DoubleBuffer,
d_out_keys: DeviceArrayLike | None,
d_in_values: DeviceArrayLike | DoubleBuffer | None,
d_out_values: DeviceArrayLike | None,
order: SortOrder,
compute_capability=None,
):
d_in_keys_array, d_out_keys_array, d_in_values_array, d_out_values_array = (
_get_arrays(d_in_keys, d_out_keys, d_in_values, d_out_values)
)
self.d_in_keys_cccl = cccl.to_cccl_input_iter(d_in_keys_array)
self.d_out_keys_cccl = cccl.to_cccl_output_iter(d_out_keys_array)
self.d_in_values_cccl = cccl.to_cccl_input_iter(d_in_values_array)
self.d_out_values_cccl = cccl.to_cccl_output_iter(d_out_values_array)
# TODO: decomposer op is not supported for now
self.decomposer_op = cccl.Op(
name="",
operator_type=cccl.OpKind.STATELESS,
ltoir=b"",
state_alignment=1,
state=b"", # explicit empty bytes so the serialize path is byte-safe
)
decomposer_return_type = "".encode("utf-8")
build_order = (
_bindings.SortOrder.ASCENDING
if order is SortOrder.ASCENDING
else _bindings.SortOrder.DESCENDING
)
self.build_results, self._bound_build_result = cache_build_results(
_bindings.DeviceRadixSortBuildResult,
d_in_keys,
d_out_keys,
d_in_values,
d_out_values,
order,
compute_capability=compute_capability,
builder=lambda: cccl.build_for_ccs(
_bindings.DeviceRadixSortBuildResult,
build_order,
self.d_in_keys_cccl,
self.d_in_values_cccl,
self.decomposer_op,
decomposer_return_type,
compute_capability=compute_capability,
),
)
def __call__(
self,
*,
temp_storage,
d_in_keys: DeviceArrayLike | DoubleBuffer,
d_out_keys: DeviceArrayLike | None,
d_in_values: DeviceArrayLike | DoubleBuffer | None,
d_out_values: DeviceArrayLike | None,
num_items: int,
begin_bit: int | None = None,
end_bit: int | None = None,
stream=None,
):
# Select (and lazily load) the build result for the current device.
self.loaded_build_result = cccl.resolve_build_result(
self.build_results, self._bound_build_result
)
d_in_keys_array, d_out_keys_array, d_in_values_array, d_out_values_array = (
_get_arrays(d_in_keys, d_out_keys, d_in_values, d_out_values)
)
set_cccl_iterator_state(self.d_in_keys_cccl, d_in_keys_array)
if d_in_values_array is not None:
set_cccl_iterator_state(self.d_in_values_cccl, d_in_values_array)
set_cccl_iterator_state(self.d_out_keys_cccl, d_out_keys_array)
if d_out_values_array is not None:
set_cccl_iterator_state(self.d_out_values_cccl, d_out_values_array)
is_overwrite_okay = isinstance(d_in_keys, DoubleBuffer)
stream_handle = validate_and_get_stream(stream)
if temp_storage is None:
temp_storage_bytes = 0
d_temp_storage = 0
else:
temp_storage_bytes = temp_storage.nbytes
# Note: this is slightly slower, but supports all ndarray-like objects as long as they support CAI
# TODO: switch to use gpumemoryview once it's ready
d_temp_storage = get_data_pointer(temp_storage)
if begin_bit is None:
begin_bit = 0
if end_bit is None:
key_type = get_dtype(d_in_keys_array)
end_bit = key_type.itemsize * 8
selector = -1
temp_storage_bytes, selector = self.loaded_build_result.compute(
d_temp_storage,
temp_storage_bytes,
self.d_in_keys_cccl,
self.d_out_keys_cccl,
self.d_in_values_cccl,
self.d_out_values_cccl,
self.decomposer_op,
num_items,
begin_bit,
end_bit,
is_overwrite_okay,
selector,
stream_handle,
)
if is_overwrite_okay and temp_storage is not None:
assert selector in (0, 1)
assert isinstance(d_in_keys, DoubleBuffer)
d_in_keys.selector = selector
if d_in_values is not None:
assert isinstance(d_in_values, DoubleBuffer)
d_in_values.selector = selector
return temp_storage_bytes
@cache_with_registered_key_functions
def make_radix_sort(
*,
d_in_keys: DeviceArrayLike | DoubleBuffer,
d_out_keys: DeviceArrayLike | None,
d_in_values: DeviceArrayLike | DoubleBuffer | None,
d_out_values: DeviceArrayLike | None,
order: SortOrder,
compute_capability=None,
):
"""Implements a device-wide radix sort using ``d_in_keys`` in the requested order.
Example:
Below, ``make_radix_sort`` is used to create a radix sort object that can be reused.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/sort/radix_sort_object.py
:language: python
:start-after: # example-begin
Args:
d_in_keys: Device array or DoubleBuffer containing the input keys to be sorted
d_out_keys: Device array to store the sorted keys
d_in_values: Optional Device array or DoubleBuffer containing the input keys to be sorted
d_out_values: Device array to store the sorted values
op: Callable representing the comparison operator
compute_capability: Compute capability, or list of capabilities, to
build for ahead of time. Accepts a packed int (e.g. ``90``), a
``(major, minor)`` pair, a string (e.g. ``"9.0"``), or a list
thereof. When ``None`` (the default), the current device's
architecture is used.
Returns:
A callable object that can be used to perform the radix sort
"""
return _RadixSort(
d_in_keys,
d_out_keys,
d_in_values,
d_out_values,
order,
compute_capability=compute_capability,
)
def radix_sort(
*,
d_in_keys: DeviceArrayLike | DoubleBuffer,
d_out_keys: DeviceArrayLike | None,
d_in_values: DeviceArrayLike | DoubleBuffer | None = None,
d_out_values: DeviceArrayLike | None = None,
num_items: int,
order: SortOrder,
begin_bit: int | None = None,
end_bit: int | None = None,
stream=None,
):
"""
Performs device-wide radix sort.
This function automatically handles temporary storage allocation and execution.
Example:
Below, ``radix_sort`` is used to sort a sequence of keys. It also rearranges the values according to the keys' order.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/sort/radix_sort_basic.py
:language: python
:start-after: # example-begin
In the following example, ``radix_sort`` is used to sort a sequence of keys with a ``DoubleBuffer`` for reduced temporary storage.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/sort/radix_sort_buffer.py
:language: python
:start-after: # example-begin
Args:
d_in_keys: Device array or DoubleBuffer containing the input sequence of keys
d_out_keys: Device array to store the sorted keys (optional)
d_in_values: Device array or DoubleBuffer containing the input sequence of values (optional)
d_out_values: Device array to store the sorted values (optional)
num_items: Number of items to sort
order: Sort order (ascending or descending)
begin_bit: Beginning bit position for comparison (optional)
end_bit: Ending bit position for comparison (optional)
stream: CUDA stream for the operation (optional)
"""
sorter = 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=order,
)
tmp_storage_bytes = 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,
begin_bit=begin_bit,
end_bit=end_bit,
stream=stream,
)
tmp_storage = TempStorageBuffer(tmp_storage_bytes, stream)
sorter(
temp_storage=tmp_storage,
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,
begin_bit=begin_bit,
end_bit=end_bit,
stream=stream,
)

298
cccl_upstream/python/cuda_cccl/cuda/compute/algorithms/_sort/_segmented_sort.py Normal file
View File

@@ -0,0 +1,298 @@
# Copyright (c) 2024-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 ... import _bindings
from ... import _cccl_interop as cccl
from ..._caching import cache_build_results, cache_with_registered_key_functions
from ..._cccl_interop import set_cccl_iterator_state
from ..._serialization import BUILD_RESULTS, ITER, Serializable
from ..._utils.protocols import (
get_data_pointer,
validate_and_get_stream,
)
from ..._utils.temp_storage_buffer import TempStorageBuffer
from ...typing import DeviceArrayLike
from ._sort_common import DoubleBuffer, SortOrder, _get_arrays
class _SegmentedSort(Serializable):
__slots__ = [
"_bound_build_result",
"build_results",
"loaded_build_result",
"d_in_keys_cccl",
"d_out_keys_cccl",
"d_in_values_cccl",
"d_out_values_cccl",
"start_offsets_in_cccl",
"end_offsets_in_cccl",
]
__serialization_schema__ = (
("d_in_keys_cccl", ITER),
("d_out_keys_cccl", ITER),
("d_in_values_cccl", ITER),
("d_out_values_cccl", ITER),
("start_offsets_in_cccl", ITER),
("end_offsets_in_cccl", ITER),
("build_results", BUILD_RESULTS(_bindings.DeviceSegmentedSortBuildResult)),
)
def __init__(
self,
d_in_keys: DeviceArrayLike | DoubleBuffer,
d_out_keys: DeviceArrayLike | None,
d_in_values: DeviceArrayLike | DoubleBuffer | None,
d_out_values: DeviceArrayLike | None,
start_offsets_in: DeviceArrayLike,
end_offsets_in: DeviceArrayLike,
order: SortOrder,
compute_capability=None,
):
d_in_keys_array, d_out_keys_array, d_in_values_array, d_out_values_array = (
_get_arrays(d_in_keys, d_out_keys, d_in_values, d_out_values)
)
self.d_in_keys_cccl = cccl.to_cccl_input_iter(d_in_keys_array)
self.d_out_keys_cccl = cccl.to_cccl_output_iter(d_out_keys_array)
self.d_in_values_cccl = cccl.to_cccl_input_iter(d_in_values_array)
self.d_out_values_cccl = cccl.to_cccl_output_iter(d_out_values_array)
self.start_offsets_in_cccl = cccl.to_cccl_input_iter(start_offsets_in)
self.end_offsets_in_cccl = cccl.to_cccl_input_iter(end_offsets_in)
build_order = (
_bindings.SortOrder.ASCENDING
if order is SortOrder.ASCENDING
else _bindings.SortOrder.DESCENDING
)
self.build_results, self._bound_build_result = cache_build_results(
_bindings.DeviceSegmentedSortBuildResult,
d_in_keys,
d_out_keys,
d_in_values,
d_out_values,
start_offsets_in,
end_offsets_in,
order,
compute_capability=compute_capability,
builder=lambda: cccl.build_for_ccs(
_bindings.DeviceSegmentedSortBuildResult,
build_order,
self.d_in_keys_cccl,
self.d_in_values_cccl,
self.start_offsets_in_cccl,
self.end_offsets_in_cccl,
compute_capability=compute_capability,
),
)
def __call__(
self,
*,
temp_storage,
d_in_keys,
d_out_keys,
d_in_values,
d_out_values,
num_items,
num_segments,
start_offsets_in,
end_offsets_in,
stream=None,
):
# Select (and lazily load) the build result for the current device.
self.loaded_build_result = cccl.resolve_build_result(
self.build_results, self._bound_build_result
)
if num_segments > np.iinfo(np.int32).max:
raise RuntimeError(
"Segmented sort does not currently support more than 2^31-1 segments."
)
d_in_keys_array, d_out_keys_array, d_in_values_array, d_out_values_array = (
_get_arrays(d_in_keys, d_out_keys, d_in_values, d_out_values)
)
set_cccl_iterator_state(self.d_in_keys_cccl, d_in_keys_array)
set_cccl_iterator_state(self.d_out_keys_cccl, d_out_keys_array)
if d_in_values_array is not None:
set_cccl_iterator_state(self.d_in_values_cccl, d_in_values_array)
if d_out_values_array is not None:
set_cccl_iterator_state(self.d_out_values_cccl, d_out_values_array)
set_cccl_iterator_state(self.start_offsets_in_cccl, start_offsets_in)
set_cccl_iterator_state(self.end_offsets_in_cccl, end_offsets_in)
stream_handle = validate_and_get_stream(stream)
if temp_storage is None:
temp_storage_bytes = 0
d_temp_storage = 0
else:
temp_storage_bytes = temp_storage.nbytes
d_temp_storage = get_data_pointer(temp_storage)
# Detect overwrite mode and selector, similar to radix sort
is_overwrite_okay = isinstance(d_in_keys, DoubleBuffer)
selector = -1
temp_storage_bytes, selector = self.loaded_build_result.compute(
d_temp_storage,
temp_storage_bytes,
self.d_in_keys_cccl,
self.d_out_keys_cccl,
self.d_in_values_cccl,
self.d_out_values_cccl,
num_items,
num_segments,
self.start_offsets_in_cccl,
self.end_offsets_in_cccl,
is_overwrite_okay,
selector,
stream_handle,
)
if is_overwrite_okay and temp_storage is not None:
assert selector in (0, 1)
assert isinstance(d_in_keys, DoubleBuffer)
d_in_keys.selector = selector
if d_in_values is not None:
assert isinstance(d_in_values, DoubleBuffer)
d_in_values.selector = selector
return temp_storage_bytes
@cache_with_registered_key_functions
def make_segmented_sort(
*,
d_in_keys: DeviceArrayLike | DoubleBuffer,
d_out_keys: DeviceArrayLike | None = None,
d_in_values: DeviceArrayLike | DoubleBuffer | None = None,
d_out_values: DeviceArrayLike | None = None,
start_offsets_in: DeviceArrayLike,
end_offsets_in: DeviceArrayLike,
order: SortOrder,
compute_capability=None,
):
"""
Performs a device-wide segmented sort using the specified keys and values.
Example:
Below, ``make_segmented_sort`` is used to create a segmented sort object that can be reused.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/sort/segmented_sort_object.py
:language: python
:start-after: # example-begin
Args:
d_in_keys: Device array or DoubleBuffer containing the input keys to be sorted
d_out_keys: Device array to store the sorted keys
d_in_values: Optional Device array or DoubleBuffer containing the input values to be sorted
d_out_values: Device array to store the sorted values
start_offsets_in: Device array or iterator containing the sequence of beginning offsets
end_offsets_in: Device array or iterator containing the sequence of ending offsets
order: SortOrder specifying the order of the sort
compute_capability: Compute capability, or list of capabilities, to
build for ahead of time. Accepts a packed int (e.g. ``90``), a
``(major, minor)`` pair, a string (e.g. ``"9.0"``), or a list
thereof. When ``None`` (the default), the current device's
architecture is used.
Returns:
A callable object that can be used to perform the segmented sort
"""
return _SegmentedSort(
d_in_keys,
d_out_keys,
d_in_values,
d_out_values,
start_offsets_in,
end_offsets_in,
order,
compute_capability=compute_capability,
)
def segmented_sort(
*,
d_in_keys: DeviceArrayLike | DoubleBuffer,
d_out_keys: DeviceArrayLike | None = None,
d_in_values: DeviceArrayLike | DoubleBuffer | None = None,
d_out_values: DeviceArrayLike | None = None,
num_items: int,
num_segments: int,
start_offsets_in: DeviceArrayLike,
end_offsets_in: DeviceArrayLike,
order: SortOrder,
stream=None,
):
"""
Performs device-wide segmented sort.
This function automatically handles temporary storage allocation and execution.
Example:
Below, ``segmented_sort`` is used to perform a segmented sort. It also rearranges the values according to the keys' order.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/sort/segmented_sort_basic.py
:language: python
:start-after: # example-begin
In the following example, ``segmented_sort`` is used to perform a segmented sort with a ``DoubleBuffer`` for reduced temporary storage.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/sort/segmented_sort_buffer.py
:language: python
:start-after: # example-begin
Args:
d_in_keys: Device array or DoubleBuffer containing the input keys to be sorted
d_out_keys: Device array to store the sorted keys (optional)
d_in_values: Device array or DoubleBuffer containing the input values to be sorted (optional)
d_out_values: Device array to store the sorted values (optional)
num_items: Total number of items to sort
num_segments: Number of segments to sort
start_offsets_in: Device array or iterator containing the sequence of beginning offsets
end_offsets_in: Device array or iterator containing the sequence of ending offsets
order: Sort order (ascending or descending)
stream: CUDA stream for the operation (optional)
"""
sorter = make_segmented_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,
start_offsets_in=start_offsets_in,
end_offsets_in=end_offsets_in,
order=order,
)
tmp_storage_bytes = 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,
num_segments=num_segments,
start_offsets_in=start_offsets_in,
end_offsets_in=end_offsets_in,
stream=stream,
)
tmp_storage = TempStorageBuffer(tmp_storage_bytes, stream)
sorter(
temp_storage=tmp_storage,
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,
num_segments=num_segments,
start_offsets_in=start_offsets_in,
end_offsets_in=end_offsets_in,
stream=stream,
)

62
cccl_upstream/python/cuda_cccl/cuda/compute/algorithms/_sort/_sort_common.py Normal file
View File

@@ -0,0 +1,62 @@
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
from enum import Enum
from typing import Tuple
from ..._caching import cache_with_registered_key_functions
from ..._utils.protocols import get_dtype
from ...typing import DeviceArrayLike
class SortOrder(Enum):
ASCENDING = 0
DESCENDING = 1
class DoubleBuffer:
def __init__(self, d_current: DeviceArrayLike, d_alternate: DeviceArrayLike):
self.d_buffers = [d_current, d_alternate]
self.selector = 0
def current(self):
return self.d_buffers[self.selector]
def alternate(self):
return self.d_buffers[1 - self.selector]
def _get_arrays(
d_in_keys: DeviceArrayLike | DoubleBuffer,
d_out_keys: DeviceArrayLike | None,
d_in_values: DeviceArrayLike | DoubleBuffer | None,
d_out_values: DeviceArrayLike | None,
) -> Tuple[DeviceArrayLike, DeviceArrayLike, DeviceArrayLike, DeviceArrayLike]:
if isinstance(d_in_keys, DoubleBuffer):
d_in_keys_array = d_in_keys.current()
d_out_keys_array = d_in_keys.alternate()
if d_in_values is not None:
assert isinstance(d_in_values, DoubleBuffer)
d_in_values_array = d_in_values.current()
d_out_values_array = d_in_values.alternate()
else:
d_in_values_array = None
d_out_values_array = None
else:
d_in_keys_array = d_in_keys
d_in_values_array = d_in_values
d_out_keys_array = d_out_keys
d_out_values_array = d_out_values
return d_in_keys_array, d_out_keys_array, d_in_values_array, d_out_values_array
# DoubleBuffer: extract dtype from current buffer
cache_with_registered_key_functions.register(
DoubleBuffer, lambda buf: get_dtype(buf.current())
)

289
cccl_upstream/python/cuda_cccl/cuda/compute/algorithms/_three_way_partition.py Normal file
View File

@@ -0,0 +1,289 @@
# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
from typing import Callable
from .. import _bindings, types
from .. import _cccl_interop as cccl
from .._caching import cache_build_results, cache_with_registered_key_functions
from .._cccl_interop import set_cccl_iterator_state
from .._serialization import BUILD_RESULTS, ITER, OP, Serializable
from .._utils import protocols
from .._utils.temp_storage_buffer import TempStorageBuffer
from ..op import OpAdapter, make_op_adapter
from ..typing import DeviceArrayLike, IteratorT, Operator
class _ThreeWayPartition(Serializable):
__slots__ = [
"_bound_build_result",
"build_results",
"loaded_build_result",
"d_in_cccl",
"d_first_part_out_cccl",
"d_second_part_out_cccl",
"d_unselected_out_cccl",
"d_num_selected_out_cccl",
"select_first_part_op_cccl",
"select_second_part_op_cccl",
]
__serialization_schema__ = (
("d_in_cccl", ITER),
("d_first_part_out_cccl", ITER),
("d_second_part_out_cccl", ITER),
("d_unselected_out_cccl", ITER),
("d_num_selected_out_cccl", ITER),
("select_first_part_op_cccl", OP),
("select_second_part_op_cccl", OP),
("build_results", BUILD_RESULTS(_bindings.DeviceThreeWayPartitionBuildResult)),
)
def __init__(
self,
d_in: DeviceArrayLike | IteratorT,
d_first_part_out: DeviceArrayLike | IteratorT,
d_second_part_out: DeviceArrayLike | IteratorT,
d_unselected_out: DeviceArrayLike | IteratorT,
d_num_selected_out: DeviceArrayLike | IteratorT,
select_first_part_op: OpAdapter,
select_second_part_op: OpAdapter,
compute_capability=None,
):
self.d_in_cccl = cccl.to_cccl_input_iter(d_in)
self.d_first_part_out_cccl = cccl.to_cccl_output_iter(d_first_part_out)
self.d_second_part_out_cccl = cccl.to_cccl_output_iter(d_second_part_out)
self.d_unselected_out_cccl = cccl.to_cccl_output_iter(d_unselected_out)
self.d_num_selected_out_cccl = cccl.to_cccl_output_iter(d_num_selected_out)
# Compile ops - partition predicates return uint8 (boolean)
value_type = cccl.get_value_type(d_in)
self.select_first_part_op_cccl = select_first_part_op.compile(
(value_type,), types.uint8
)
self.select_second_part_op_cccl = select_second_part_op.compile(
(value_type,), types.uint8
)
self.build_results, self._bound_build_result = cache_build_results(
_bindings.DeviceThreeWayPartitionBuildResult,
d_in,
d_first_part_out,
d_second_part_out,
d_unselected_out,
d_num_selected_out,
select_first_part_op,
select_second_part_op,
compute_capability=compute_capability,
builder=lambda: cccl.build_for_ccs(
_bindings.DeviceThreeWayPartitionBuildResult,
self.d_in_cccl,
self.d_first_part_out_cccl,
self.d_second_part_out_cccl,
self.d_unselected_out_cccl,
self.d_num_selected_out_cccl,
self.select_first_part_op_cccl,
self.select_second_part_op_cccl,
compute_capability=compute_capability,
),
)
def __call__(
self,
*,
temp_storage,
d_in,
d_first_part_out,
d_second_part_out,
d_unselected_out,
d_num_selected_out,
select_first_part_op: Callable | OpAdapter,
select_second_part_op: Callable | OpAdapter,
num_items: int,
stream=None,
):
# Select (and lazily load) the build result for the current device.
self.loaded_build_result = cccl.resolve_build_result(
self.build_results, self._bound_build_result
)
set_cccl_iterator_state(self.d_in_cccl, d_in)
set_cccl_iterator_state(self.d_first_part_out_cccl, d_first_part_out)
set_cccl_iterator_state(self.d_second_part_out_cccl, d_second_part_out)
set_cccl_iterator_state(self.d_unselected_out_cccl, d_unselected_out)
set_cccl_iterator_state(self.d_num_selected_out_cccl, d_num_selected_out)
first_op_adapter = make_op_adapter(select_first_part_op)
second_op_adapter = make_op_adapter(select_second_part_op)
self.select_first_part_op_cccl.state = first_op_adapter.get_state()
self.select_second_part_op_cccl.state = second_op_adapter.get_state()
stream_handle = protocols.validate_and_get_stream(stream)
if temp_storage is None:
temp_storage_bytes = 0
d_temp_storage = 0
else:
temp_storage_bytes = temp_storage.nbytes
d_temp_storage = protocols.get_data_pointer(temp_storage)
temp_storage_bytes = self.loaded_build_result.compute(
d_temp_storage,
temp_storage_bytes,
self.d_in_cccl,
self.d_first_part_out_cccl,
self.d_second_part_out_cccl,
self.d_unselected_out_cccl,
self.d_num_selected_out_cccl,
self.select_first_part_op_cccl,
self.select_second_part_op_cccl,
num_items,
stream_handle,
)
return temp_storage_bytes
@cache_with_registered_key_functions
def make_three_way_partition(
*,
d_in: DeviceArrayLike | IteratorT,
d_first_part_out: DeviceArrayLike | IteratorT,
d_second_part_out: DeviceArrayLike | IteratorT,
d_unselected_out: DeviceArrayLike | IteratorT,
d_num_selected_out: DeviceArrayLike | IteratorT,
select_first_part_op: Operator,
select_second_part_op: Operator,
compute_capability=None,
):
"""
Computes a device-wide three-way partition using the specified unary ``select_first_part_op`` and ``select_second_part_op`` operators.
Example:
Below, ``make_three_way_partition`` is used to create a three-way partition object that can be reused.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/partition/three_way_partition_object.py
:language: python
:start-after: # example-begin
Args:
d_in: Device array or iterator containing the input sequence of data items
d_first_part_out: Device array or iterator to store the first part of the output
d_second_part_out: Device array or iterator to store the second part of the output
d_unselected_out: Device array or iterator to store the unselected items
d_num_selected_out: Device array to store the number of items selected. The total number of items selected by ``select_first_part_op`` and ``select_second_part_op`` is stored in ``d_num_selected_out[0]`` and ``d_num_selected_out[1]``, respectively.
select_first_part_op: Unary operator to select the first part.
The signature is ``(T) -> uint8``, where ``T`` is the input data type.
Returns 1 (selected) or 0 (not selected).
Can reference device arrays as globals/closures - they will be automatically captured.
select_second_part_op: Unary operator to select the second part.
The signature is ``(T) -> uint8``, where ``T`` is the input data type.
Returns 1 (selected) or 0 (not selected).
Can reference device arrays as globals/closures - they will be automatically captured.
compute_capability: Compute capability, or list of capabilities, to
build for ahead of time. Accepts a packed int (e.g. ``90``), a
``(major, minor)`` pair, a string (e.g. ``"9.0"``), or a list
thereof. When ``None`` (the default), the current device's
architecture is used.
Returns:
A callable object that can be used to perform the three-way partition
"""
first_op_adapter = make_op_adapter(select_first_part_op)
second_op_adapter = make_op_adapter(select_second_part_op)
return _ThreeWayPartition(
d_in,
d_first_part_out,
d_second_part_out,
d_unselected_out,
d_num_selected_out,
first_op_adapter,
second_op_adapter,
compute_capability=compute_capability,
)
def three_way_partition(
*,
d_in: DeviceArrayLike | IteratorT,
d_first_part_out: DeviceArrayLike | IteratorT,
d_second_part_out: DeviceArrayLike | IteratorT,
d_unselected_out: DeviceArrayLike | IteratorT,
d_num_selected_out: DeviceArrayLike | IteratorT,
select_first_part_op: Operator,
select_second_part_op: Operator,
num_items: int,
stream=None,
):
"""
Performs device-wide three-way partition. Given an input sequence of data items, it partitions the items into three parts:
- The first part is selected by the ``select_first_part_op`` operator.
- The second part is selected by the ``select_second_part_op`` operator.
- The unselected items are not selected by either operator.
This function automatically handles temporary storage allocation and execution.
Example:
Below, ``three_way_partition`` is used to partition a sequence of integers into three parts.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/partition/three_way_partition_basic.py
:language: python
:start-after: # example-begin
Args:
d_in: Device array or iterator containing the input sequence of data items
d_first_part_out: Device array or iterator to store the first part of the output
d_second_part_out: Device array or iterator to store the second part of the output
d_unselected_out: Device array or iterator to store the unselected items
d_num_selected_out: Device array to store the number of items selected. The total number of items selected by ``select_first_part_op`` and ``select_second_part_op`` is stored in ``d_num_selected_out[0]`` and ``d_num_selected_out[1]``, respectively.
select_first_part_op: Unary operator to select the first part.
The signature is ``(T) -> uint8``, where ``T`` is the input data type.
Returns 1 (selected) or 0 (not selected).
select_second_part_op: Unary operator to select the second part.
The signature is ``(T) -> uint8``, where ``T`` is the input data type.
Returns 1 (selected) or 0 (not selected).
num_items: Number of items to partition
stream: CUDA stream for the operation (optional)
"""
# Create adapters to support stateful ops
first_op_adapter = make_op_adapter(select_first_part_op)
second_op_adapter = make_op_adapter(select_second_part_op)
partitioner = make_three_way_partition(
d_in=d_in,
d_first_part_out=d_first_part_out,
d_second_part_out=d_second_part_out,
d_unselected_out=d_unselected_out,
d_num_selected_out=d_num_selected_out,
select_first_part_op=first_op_adapter,
select_second_part_op=second_op_adapter,
)
tmp_storage_bytes = partitioner(
temp_storage=None,
d_in=d_in,
d_first_part_out=d_first_part_out,
d_second_part_out=d_second_part_out,
d_unselected_out=d_unselected_out,
d_num_selected_out=d_num_selected_out,
select_first_part_op=first_op_adapter,
select_second_part_op=second_op_adapter,
num_items=num_items,
stream=stream,
)
tmp_storage = TempStorageBuffer(tmp_storage_bytes, stream)
partitioner(
temp_storage=tmp_storage,
d_in=d_in,
d_first_part_out=d_first_part_out,
d_second_part_out=d_second_part_out,
d_unselected_out=d_unselected_out,
d_num_selected_out=d_num_selected_out,
select_first_part_op=first_op_adapter,
select_second_part_op=second_op_adapter,
num_items=num_items,
stream=stream,
)

371
cccl_upstream/python/cuda_cccl/cuda/compute/algorithms/_transform.py Normal file
View File

@@ -0,0 +1,371 @@
# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
from typing import Callable
from .. import _bindings
from .. import _cccl_interop as cccl
from .._caching import cache_build_results, cache_with_registered_key_functions
from .._cccl_interop import set_cccl_iterator_state
from .._serialization import BUILD_RESULTS, ITER, OP, Serializable
from .._utils import protocols
from ..op import OpAdapter, make_op_adapter
from ..typing import DeviceArrayLike, IteratorT, Operator
class _UnaryTransform(Serializable):
__slots__ = [
"_bound_build_result",
"d_in_cccl",
"d_out_cccl",
"op_cccl",
"build_results",
"loaded_build_result",
]
__serialization_schema__ = (
("d_in_cccl", ITER),
("d_out_cccl", ITER),
("op_cccl", OP),
("build_results", BUILD_RESULTS(_bindings.DeviceUnaryTransform)),
)
def __init__(
self,
d_in: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike | IteratorT,
op: OpAdapter,
compute_capability=None,
):
self.d_in_cccl = cccl.to_cccl_input_iter(d_in)
self.d_out_cccl = cccl.to_cccl_output_iter(d_out)
# Compile the op with input/output types
in_type = cccl.get_value_type(d_in)
out_type = cccl.get_value_type(d_out)
self.op_cccl = op.compile((in_type,), out_type)
self.build_results, self._bound_build_result = cache_build_results(
_bindings.DeviceUnaryTransform,
d_in,
d_out,
op,
compute_capability=compute_capability,
builder=lambda: cccl.build_for_ccs(
_bindings.DeviceUnaryTransform,
self.d_in_cccl,
self.d_out_cccl,
self.op_cccl,
compute_capability=compute_capability,
),
)
def __call__(
self,
*,
d_in,
d_out,
op: Callable | OpAdapter,
num_items: int,
stream=None,
):
# Select (and lazily load) the build result for the current device.
self.loaded_build_result = cccl.resolve_build_result(
self.build_results, self._bound_build_result
)
op_adapter = make_op_adapter(op)
set_cccl_iterator_state(self.d_in_cccl, d_in)
set_cccl_iterator_state(self.d_out_cccl, d_out)
self.op_cccl.state = op_adapter.get_state()
stream_handle = protocols.validate_and_get_stream(stream)
self.loaded_build_result.compute(
self.d_in_cccl,
self.d_out_cccl,
num_items,
self.op_cccl,
stream_handle,
)
return None
class _BinaryTransform(Serializable):
__slots__ = [
"_bound_build_result",
"d_in1_cccl",
"d_in2_cccl",
"d_out_cccl",
"op_cccl",
"build_results",
"loaded_build_result",
]
__serialization_schema__ = (
("d_in1_cccl", ITER),
("d_in2_cccl", ITER),
("d_out_cccl", ITER),
("op_cccl", OP),
("build_results", BUILD_RESULTS(_bindings.DeviceBinaryTransform)),
)
def __init__(
self,
d_in1: DeviceArrayLike | IteratorT,
d_in2: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike | IteratorT,
op: OpAdapter,
compute_capability=None,
):
self.d_in1_cccl = cccl.to_cccl_input_iter(d_in1)
self.d_in2_cccl = cccl.to_cccl_input_iter(d_in2)
self.d_out_cccl = cccl.to_cccl_output_iter(d_out)
# Compile the op with input/output types
in1_type = cccl.get_value_type(d_in1)
in2_type = cccl.get_value_type(d_in2)
out_type = cccl.get_value_type(d_out)
self.op_cccl = op.compile((in1_type, in2_type), out_type)
self.build_results, self._bound_build_result = cache_build_results(
_bindings.DeviceBinaryTransform,
d_in1,
d_in2,
d_out,
op,
compute_capability=compute_capability,
builder=lambda: cccl.build_for_ccs(
_bindings.DeviceBinaryTransform,
self.d_in1_cccl,
self.d_in2_cccl,
self.d_out_cccl,
self.op_cccl,
compute_capability=compute_capability,
),
)
def __call__(
self,
*,
d_in1,
d_in2,
d_out,
op: Callable | OpAdapter,
num_items: int,
stream=None,
):
# Select (and lazily load) the build result for the current device.
self.loaded_build_result = cccl.resolve_build_result(
self.build_results, self._bound_build_result
)
set_cccl_iterator_state(self.d_in1_cccl, d_in1)
set_cccl_iterator_state(self.d_in2_cccl, d_in2)
set_cccl_iterator_state(self.d_out_cccl, d_out)
op_adapter = make_op_adapter(op)
self.op_cccl.state = op_adapter.get_state()
stream_handle = protocols.validate_and_get_stream(stream)
self.loaded_build_result.compute(
self.d_in1_cccl,
self.d_in2_cccl,
self.d_out_cccl,
num_items,
self.op_cccl,
stream_handle,
)
return None
@cache_with_registered_key_functions
def make_unary_transform(
*,
d_in: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike | IteratorT,
op: Operator,
compute_capability=None,
):
"""
Create a unary transform object that can be called to apply a transformation
to each element of the input according to the unary operation ``op``.
This is the object-oriented API that allows explicit control over temporary
storage allocation. For simpler usage, consider using :func:`unary_transform`.
Example:
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/transform/unary_transform_object.py
:language: python
:start-after: # example-begin
Args:
d_in: Device array or iterator containing the input sequence of data items.
d_out: Device array or iterator to store the result of the transformation.
op: Unary operation to apply to each element.
The signature is ``(T) -> U``, where ``T`` is
the input data type and ``U`` is the output data type.
compute_capability: Compute capability, or list of capabilities, to
build for ahead of time. Accepts a packed int (e.g. ``90``), a
``(major, minor)`` pair, a string (e.g. ``"9.0"``), or a list
thereof. When ``None`` (the default), the current device's
architecture is used.
Returns:
A callable object that performs the transformation.
"""
op_adapter = make_op_adapter(op)
return _UnaryTransform(
d_in, d_out, op_adapter, compute_capability=compute_capability
)
@cache_with_registered_key_functions
def make_binary_transform(
*,
d_in1: DeviceArrayLike | IteratorT,
d_in2: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike | IteratorT,
op: Operator,
compute_capability=None,
):
"""
Create a binary transform object that can be called to apply a transformation
to the given pair of input sequences according to the binary operation ``op``.
This is the object-oriented API that allows explicit control over temporary
storage allocation. For simpler usage, consider using :func:`binary_transform`.
Example:
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/transform/binary_transform_object.py
:language: python
:start-after: # example-begin
Args:
d_in1: Device array or iterator containing the first input sequence of data items.
d_in2: Device array or iterator containing the second input sequence of data items.
d_out: Device array or iterator to store the result of the transformation.
op: Binary operation.
The signature is ``(T1, T2) -> U``, where ``T1`` and ``T2`` are the input data types and
``U`` is the output data type.
compute_capability: Compute capability, or list of capabilities, to
build for ahead of time. Accepts a packed int (e.g. ``90``), a
``(major, minor)`` pair, a string (e.g. ``"9.0"``), or a list
thereof. When ``None`` (the default), the current device's
architecture is used.
Returns:
A callable object that performs the transformation.
"""
op_adapter = make_op_adapter(op)
return _BinaryTransform(
d_in1, d_in2, d_out, op_adapter, compute_capability=compute_capability
)
def unary_transform(
*,
d_in: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike | IteratorT,
op: Operator,
num_items: int,
stream=None,
):
"""
Performs device-wide unary transform.
This function automatically handles temporary storage allocation and execution.
The ``op`` function can reference device arrays as globals or closures - they will
be automatically captured as state arrays, enabling stateful operations like counting.
Example:
Below, ``unary_transform`` is used to apply a transformation to each element of the input.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/transform/unary_transform_basic.py
:language: python
:start-after: # example-begin
When working with custom struct types, you need to provide type annotations
to help with type inference. See the binary transform struct example for reference:
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/struct/struct_transform.py
:language: python
:start-after: # example-begin
Args:
d_in: Device array or iterator containing the input sequence of data items.
d_out: Device array or iterator to store the result of the transformation.
op: Unary operation to apply to each element.
The signature is ``(T) -> U``, where ``T`` is
the input data type and ``U`` is the output data type.
Can reference device arrays as globals/closures - they will be automatically captured.
num_items: Number of items to transform.
stream: CUDA stream to use for the operation.
"""
op_adapter = make_op_adapter(op)
transformer = make_unary_transform(d_in=d_in, d_out=d_out, op=op_adapter)
transformer(
d_in=d_in, d_out=d_out, op=op_adapter, num_items=num_items, stream=stream
)
def binary_transform(
*,
d_in1: DeviceArrayLike | IteratorT,
d_in2: DeviceArrayLike | IteratorT,
d_out: DeviceArrayLike | IteratorT,
op: Operator,
num_items: int,
stream=None,
):
"""
Performs device-wide binary transform.
This function automatically handles temporary storage allocation and execution.
Example:
Below, ``binary_transform`` is used to apply a transformation to pairs of elements from two input sequences.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/transform/binary_transform_basic.py
:language: python
:start-after: # example-begin
When working with custom struct types, you need to provide type annotations
to help with type inference. See the following example:
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/struct/struct_transform.py
:language: python
:start-after: # example-begin
Args:
d_in1: Device array or iterator containing the first input sequence of data items.
d_in2: Device array or iterator containing the second input sequence of data items.
d_out: Device array or iterator to store the result of the transformation.
op: Binary operation.
The signature is ``(T1, T2) -> U``, where ``T1`` and ``T2`` are the input data types and
``U`` is the output data type.
Can reference device arrays as globals/closures - they will be automatically captured.
num_items: Number of items to transform.
stream: CUDA stream to use for the operation.
"""
op_adapter = make_op_adapter(op)
transformer = make_binary_transform(
d_in1=d_in1, d_in2=d_in2, d_out=d_out, op=op_adapter
)
transformer(
d_in1=d_in1,
d_in2=d_in2,
d_out=d_out,
op=op_adapter,
num_items=num_items,
stream=stream,
)

253
cccl_upstream/python/cuda_cccl/cuda/compute/algorithms/_unique_by_key.py Normal file
View File

@@ -0,0 +1,253 @@
# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
from .. import _bindings, types
from .. import _cccl_interop as cccl
from .._caching import cache_build_results, cache_with_registered_key_functions
from .._cccl_interop import set_cccl_iterator_state
from .._serialization import BUILD_RESULTS, ITER, OP, Serializable
from .._utils.protocols import (
get_data_pointer,
validate_and_get_stream,
)
from .._utils.temp_storage_buffer import TempStorageBuffer
from ..op import OpAdapter, make_op_adapter
from ..typing import DeviceArrayLike, IteratorT, Operator
class _UniqueByKey(Serializable):
__slots__ = [
"_bound_build_result",
"build_results",
"loaded_build_result",
"d_in_keys_cccl",
"d_in_items_cccl",
"d_out_keys_cccl",
"d_out_items_cccl",
"d_out_num_selected_cccl",
"op_cccl",
]
__serialization_schema__ = (
("d_in_keys_cccl", ITER),
("d_in_items_cccl", ITER),
("d_out_keys_cccl", ITER),
("d_out_items_cccl", ITER),
("d_out_num_selected_cccl", ITER),
("op_cccl", OP),
("build_results", BUILD_RESULTS(_bindings.DeviceUniqueByKeyBuildResult)),
)
def __init__(
self,
d_in_keys: DeviceArrayLike | IteratorT,
d_in_items: DeviceArrayLike | IteratorT,
d_out_keys: DeviceArrayLike | IteratorT,
d_out_items: DeviceArrayLike | IteratorT,
d_out_num_selected: DeviceArrayLike,
op: OpAdapter,
compute_capability=None,
):
self.d_in_keys_cccl = cccl.to_cccl_input_iter(d_in_keys)
self.d_in_items_cccl = cccl.to_cccl_input_iter(d_in_items)
self.d_out_keys_cccl = cccl.to_cccl_output_iter(d_out_keys)
self.d_out_items_cccl = cccl.to_cccl_output_iter(d_out_items)
self.d_out_num_selected_cccl = cccl.to_cccl_output_iter(d_out_num_selected)
# Compile the op - unique_by_key expects bool return (comparison)
value_type = cccl.get_value_type(d_in_keys)
self.op_cccl = op.compile((value_type, value_type), types.uint8)
self.build_results, self._bound_build_result = cache_build_results(
_bindings.DeviceUniqueByKeyBuildResult,
d_in_keys,
d_in_items,
d_out_keys,
d_out_items,
d_out_num_selected,
op,
compute_capability=compute_capability,
builder=lambda: cccl.build_for_ccs(
_bindings.DeviceUniqueByKeyBuildResult,
self.d_in_keys_cccl,
self.d_in_items_cccl,
self.d_out_keys_cccl,
self.d_out_items_cccl,
self.d_out_num_selected_cccl,
self.op_cccl,
compute_capability=compute_capability,
),
)
def __call__(
self,
*,
temp_storage,
d_in_keys: DeviceArrayLike | IteratorT,
d_in_items: DeviceArrayLike | IteratorT,
d_out_keys: DeviceArrayLike | IteratorT,
d_out_items: DeviceArrayLike | IteratorT,
d_out_num_selected: DeviceArrayLike,
op: Operator,
num_items: int,
stream=None,
):
# Select (and lazily load) the build result for the current device.
self.loaded_build_result = cccl.resolve_build_result(
self.build_results, self._bound_build_result
)
set_cccl_iterator_state(self.d_in_keys_cccl, d_in_keys)
set_cccl_iterator_state(self.d_in_items_cccl, d_in_items)
set_cccl_iterator_state(self.d_out_keys_cccl, d_out_keys)
set_cccl_iterator_state(self.d_out_items_cccl, d_out_items)
set_cccl_iterator_state(self.d_out_num_selected_cccl, d_out_num_selected)
# Update op state for stateful ops
op_adapter = make_op_adapter(op)
self.op_cccl.state = op_adapter.get_state()
stream_handle = validate_and_get_stream(stream)
if temp_storage is None:
temp_storage_bytes = 0
d_temp_storage = 0
else:
temp_storage_bytes = temp_storage.nbytes
# Note: this is slightly slower, but supports all ndarray-like objects as long as they support CAI
# TODO: switch to use gpumemoryview once it's ready
d_temp_storage = get_data_pointer(temp_storage)
temp_storage_bytes = self.loaded_build_result.compute(
d_temp_storage,
temp_storage_bytes,
self.d_in_keys_cccl,
self.d_in_items_cccl,
self.d_out_keys_cccl,
self.d_out_items_cccl,
self.d_out_num_selected_cccl,
self.op_cccl,
num_items,
stream_handle,
)
return temp_storage_bytes
@cache_with_registered_key_functions
def make_unique_by_key(
*,
d_in_keys: DeviceArrayLike | IteratorT,
d_in_items: DeviceArrayLike | IteratorT,
d_out_keys: DeviceArrayLike | IteratorT,
d_out_items: DeviceArrayLike | IteratorT,
d_out_num_selected: DeviceArrayLike,
op: Operator,
compute_capability=None,
):
"""Implements a device-wide unique by key operation using ``d_in_keys`` and the comparison operator ``op``. Only the first key and its value from each run is selected and the total number of items selected is also reported.
Example:
Below, ``make_unique_by_key`` is used to create a unique by key object that can be reused.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/unique/unique_by_key_object.py
:language: python
:start-after: # example-begin
Args:
d_in_keys: Device array or iterator containing the input sequence of keys
d_in_items: Device array or iterator that contains each key's corresponding item
d_out_keys: Device array or iterator to store the outputted keys
d_out_items: Device array or iterator to store each outputted key's item
d_out_num_selected: Device array to store how many items were selected
op: Callable or OpKind representing the equality operator
compute_capability: Compute capability, or list of capabilities, to
build for ahead of time. Accepts a packed int (e.g. ``90``), a
``(major, minor)`` pair, a string (e.g. ``"9.0"``), or a list
thereof. When ``None`` (the default), the current device's
architecture is used.
Returns:
A callable object that can be used to perform unique by key
"""
op_adapter = make_op_adapter(op)
return _UniqueByKey(
d_in_keys,
d_in_items,
d_out_keys,
d_out_items,
d_out_num_selected,
op_adapter,
compute_capability=compute_capability,
)
def unique_by_key(
*,
d_in_keys: DeviceArrayLike | IteratorT,
d_in_items: DeviceArrayLike | IteratorT,
d_out_keys: DeviceArrayLike | IteratorT,
d_out_items: DeviceArrayLike | IteratorT,
d_out_num_selected: DeviceArrayLike,
op: Operator,
num_items: int,
stream=None,
):
"""
Performs device-wide unique by key operation using the single-phase API.
This function automatically handles temporary storage allocation and execution.
Example:
Below, ``unique_by_key`` is used to populate the arrays of output keys and items with the first key and its corresponding item from each sequence of equal keys. It also outputs the number of items selected.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/unique/unique_by_key_basic.py
:language: python
:start-after: # example-begin
Args:
d_in_keys: Device array or iterator containing the input sequence of keys
d_in_items: Device array or iterator that contains each key's corresponding item
d_out_keys: Device array or iterator to store the outputted keys
d_out_items: Device array or iterator to store each outputted key's item
d_out_num_selected: Device array to store how many items were selected
op: Callable or OpKind representing the equality operator
num_items: Number of items to process
stream: CUDA stream for the operation (optional)
"""
uniquer = make_unique_by_key(
d_in_keys=d_in_keys,
d_in_items=d_in_items,
d_out_keys=d_out_keys,
d_out_items=d_out_items,
d_out_num_selected=d_out_num_selected,
op=op,
)
tmp_storage_bytes = uniquer(
temp_storage=None,
d_in_keys=d_in_keys,
d_in_items=d_in_items,
d_out_keys=d_out_keys,
d_out_items=d_out_items,
d_out_num_selected=d_out_num_selected,
op=op,
num_items=num_items,
stream=stream,
)
tmp_storage = TempStorageBuffer(tmp_storage_bytes, stream)
uniquer(
temp_storage=tmp_storage,
d_in_keys=d_in_keys,
d_in_items=d_in_items,
d_out_keys=d_out_keys,
d_out_items=d_out_items,
d_out_num_selected=d_out_num_selected,
op=op,
num_items=num_items,
stream=stream,
)

0
cccl_upstream/python/cuda_cccl/cuda/compute/cccl/.gitkeep Normal file
View File

5
cccl_upstream/python/cuda_cccl/cuda/compute/determinism.py Normal file
View File

@@ -0,0 +1,5 @@
from __future__ import annotations
from ._bindings import Determinism
__all__ = ["Determinism"]

31
cccl_upstream/python/cuda_cccl/cuda/compute/iterators/__init__.py Normal file
View File

@@ -0,0 +1,31 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
from ._base import IteratorBase
from ._cache_modified import CacheModifiedInputIterator
from ._constant import ConstantIterator
from ._counting import CountingIterator
from ._discard import DiscardIterator
from ._permutation import PermutationIterator
from ._reverse import ReverseIterator
from ._shuffle import ShuffleIterator
from ._transform import TransformIterator, TransformOutputIterator
from ._zip import ZipIterator
__all__ = [
"CacheModifiedInputIterator",
"ConstantIterator",
"CountingIterator",
"DiscardIterator",
"IteratorBase",
"PermutationIterator",
"ReverseIterator",
"ShuffleIterator",
"TransformIterator",
"TransformOutputIterator",
"ZipIterator",
]

293
cccl_upstream/python/cuda_cccl/cuda/compute/iterators/_base.py Normal file
View File

@@ -0,0 +1,293 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Base classes for iterators.
"""
from __future__ import annotations
import hashlib
import threading
from typing import Hashable
from .._bindings import Iterator, IteratorKind, IteratorState, Op
from .._caching import cache_with_registered_key_functions
from ..types import TypeDescriptor
class IteratorBase:
"""
Iterators represent streams of data computed on the fly.
See :py:mod:`cuda.compute.iterators` for available iterators.
"""
# Subclassing
# -----------
#
# Subclasses must implement the following methods that return
# Op objects.
#
# - _make_advance_op() -> Op
# - _make_input_deref_op() -> Op | None
# - _make_output_deref_op() -> Op | None
#
# Iterators composed of other iterators must also implement:
#
# - children property to return tuple of child iterators for dependency tracking
#
# Examples of such "compound" iterators include TransformIterator,
# PermutationIterator, ReverseIterator and ZipIterator.
#
# The base class provides public cached accessors:
#
# - get_advance_op() -> Op (cached)
# - get_input_deref_op() -> Op | None (cached)
# - get_output_deref_op() -> Op | None (cached)
__slots__ = [
"_state_bytes",
"_state_alignment",
"_value_type",
"_advance_op",
"_input_deref_op",
"_output_deref_op",
"_uid_cached",
"_op_lock",
]
def __init__(
self,
state_bytes: bytes,
state_alignment: int,
value_type: TypeDescriptor,
):
"""
Args:
state_bytes: bytes object representing iterator's state
state_alignment: Alignment of the state
value_type: Type of dereferenced values
"""
self._state_bytes = state_bytes
self._state_alignment = state_alignment
self._value_type = value_type
# Per-cc caches: the compiled Op is arch-specific (its LTO-IR is
# built for the current build's target compute capability), so the memo
# must be keyed on that cc. Reusing one iterator instance across builds
# targeting different arches otherwise leaks the first arch's LTO-IR into
# the others, which nvJitLink rejects. Keyed on get_target_cc() (None ==
# current device); see get_advance_op() below.
self._advance_op: dict[Hashable, Op] = {}
self._input_deref_op: dict[Hashable, Op | None] = {}
self._output_deref_op: dict[Hashable, Op | None] = {}
self._uid_cached: str | None = None
# Free-threaded Python can let multiple threads share a read-only
# iterator object and race during the first lazy Op construction.
# The lock only protects that cache miss path; cached access stays
# lock-free and iterator mutation remains the caller's responsibility.
self._op_lock = threading.Lock()
@property
def state(self) -> IteratorState:
"""Return the iterator state for CCCL interop."""
return IteratorState(self._state_bytes)
@property
def state_alignment(self) -> int:
"""Return the alignment of the iterator state."""
return self._state_alignment
@property
def value_type(self) -> TypeDescriptor:
"""Return the TypeDescriptor for dereferenced values."""
return self._value_type
@property
def children(self) -> tuple["IteratorBase", ...]:
"""Return child iterators for automatic dependency tracking. Override in subclasses."""
return ()
def _get_uid(self) -> str:
"""Return a deterministic unique identifier for this iterator type."""
if self._uid_cached is None:
self._uid_cached = _deterministic_suffix(self.kind)
return self._uid_cached
def _make_advance_symbol(self) -> str:
"""Generate symbol name for advance operation."""
return f"{self.__class__.__name__}_advance_{self._get_uid()}"
def _make_input_deref_symbol(self) -> str:
"""Generate symbol name for input dereference operation."""
return f"{self.__class__.__name__}_input_deref_{self._get_uid()}"
def _make_output_deref_symbol(self) -> str:
"""Generate symbol name for output dereference operation."""
return f"{self.__class__.__name__}_output_deref_{self._get_uid()}"
def get_advance_op(self) -> Op:
"""Get the cached Op for the advance operation."""
from .._target_cc import get_target_cc
key = get_target_cc()
if key not in self._advance_op:
with self._op_lock:
if key not in self._advance_op:
self._advance_op[key] = self._make_advance_op()
return self._advance_op[key]
def get_input_deref_op(self) -> Op | None:
"""Get the cached Op for input dereference operation, or None if not supported."""
from .._target_cc import get_target_cc
key = get_target_cc()
if key not in self._input_deref_op:
with self._op_lock:
if key not in self._input_deref_op:
self._input_deref_op[key] = self._make_input_deref_op()
return self._input_deref_op[key]
def get_output_deref_op(self) -> Op | None:
"""Get the cached Op for output dereference operation, or None if not supported."""
from .._target_cc import get_target_cc
key = get_target_cc()
if key not in self._output_deref_op:
with self._op_lock:
if key not in self._output_deref_op:
self._output_deref_op[key] = self._make_output_deref_op()
return self._output_deref_op[key]
@property
def is_input_iterator(self) -> bool:
"""Return True if this iterator supports input dereference."""
return self.get_input_deref_op() is not None
@property
def is_output_iterator(self) -> bool:
"""Return True if this iterator supports output dereference."""
return self.get_output_deref_op() is not None
def to_cccl_iter(self, is_output: bool = False) -> Iterator:
"""
Convert this iterator to a CCCL Iterator for algorithm interop.
Args:
is_output: If True, use output_dereference; otherwise use input_dereference
Returns:
CCCL Iterator object
"""
# Get advance op
advance_op = self.get_advance_op()
# Get dereference op based on direction
if is_output:
deref_op = self.get_output_deref_op()
if deref_op is None:
raise ValueError("This iterator does not support output operations")
else:
deref_op = self.get_input_deref_op()
if deref_op is None:
raise ValueError("This iterator does not support input operations")
# Create the CCCL Iterator
return Iterator(
self._state_alignment,
IteratorKind.ITERATOR,
advance_op,
deref_op,
self._value_type.info,
state=self.state,
)
@property
def kind(self) -> Hashable:
"""Return a hashable kind for caching purposes.
Note: state_bytes is intentionally excluded - iterators with the same
type structure but different runtime state should share cached reducers.
"""
return (type(self).__name__, self._value_type)
# Abstract methods for subclasses
def _make_advance_op(self) -> Op:
"""
Create Op object for advance operation.
Returns:
Op object with compiled LTOIR
"""
raise NotImplementedError
def _make_input_deref_op(self) -> Op | None:
"""
Create Op object for input dereference operation.
Returns:
Op object with compiled LTOIR, or None if not supported
"""
raise NotImplementedError
def _make_output_deref_op(self) -> Op | None:
"""
Create Op object for output dereference operation.
Returns:
Op object with compiled LTOIR, or None if not supported
"""
raise NotImplementedError
def _deterministic_suffix(kind: Hashable) -> str:
kind_str = str(kind)
return hashlib.sha256(kind_str.encode()).hexdigest()[:16]
def compose_iterator_states(
iterators: list[IteratorBase],
) -> tuple[bytes, int, list[int]]:
"""
Concatenate multiple iterator states with proper alignment.
This is used by composite iterators (like ZipIterator and PermutationIterator)
that need to store multiple child iterator states in their own state.
Args:
iterators: List of child iterators whose states should be composed
Returns:
Tuple of:
- combined_state_bytes: Concatenated state bytes with padding
- combined_alignment: Maximum alignment requirement
- offsets: List of byte offsets for each iterator's state
"""
if not iterators:
return (b"", 1, [])
states = [bytes(memoryview(it.state)) for it in iterators]
alignments = [it.state_alignment for it in iterators]
offsets = []
current_offset = 0
combined = b""
for state, align in zip(states, alignments):
# Add padding to meet alignment requirement
padding = (align - (current_offset % align)) % align
combined += b"\x00" * padding
current_offset += padding
offsets.append(current_offset)
combined += state
current_offset += len(state)
max_alignment = max(alignments)
return (combined, max_alignment, offsets)
cache_with_registered_key_functions.register(IteratorBase, lambda it: it.kind)
__all__ = ["IteratorBase"]

156
cccl_upstream/python/cuda_cccl/cuda/compute/iterators/_cache_modified.py Normal file
View File

@@ -0,0 +1,156 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""CacheModifiedInputIterator implementation."""
from __future__ import annotations
import struct
from textwrap import dedent
from typing import Literal
from .._bindings import Op, OpKind
from .._cpp_compile import compile_cpp_op_code, cpp_type_from_descriptor
from .._utils.protocols import get_data_pointer, get_dtype
from ..types import from_numpy_dtype
from ._base import IteratorBase
from ._common import CUDA_PREAMBLE
# Map modifier names to PTX cache operators and C++ intrinsics
_CACHE_MODIFIERS = {
"stream": ("cs", "__ldcs"), # Cache streaming (evict first)
"global": ("cg", "__ldcg"), # Cache at L2 only
"volatile": ("cv", "__ldcv"), # Don't cache, always fetch
}
class CacheModifiedInputIterator(IteratorBase):
"""
Iterator that wraps a device pointer with cache-modified loads.
This iterator uses PTX cache modifiers to control how data is loaded:
- "stream": Uses streaming loads (ld.global.cs) - hints that data will not be reused
- "global": Uses global cache loads (ld.global.cg) - caches only at L2
- "volatile": Uses volatile loads (ld.global.cv) - always fetches from memory
Supports element types of size 1, 2, 4, 8, or 16 bytes.
"""
__slots__ = [
"_modifier",
"_array",
"_ptr",
]
def __init__(
self,
array,
modifier: Literal["stream", "global", "volatile"] = "stream",
):
"""
Create a cache-modified input iterator.
Args:
array: Device array to wrap (must support __cuda_array_interface__)
modifier: Cache modifier - "stream", "global", or "volatile"
"""
if modifier not in _CACHE_MODIFIERS:
raise ValueError(
f"Unknown modifier: {modifier}. Must be one of {list(_CACHE_MODIFIERS.keys())}"
)
self._modifier = modifier
self._array = array # Keep reference to prevent GC
ptr = get_data_pointer(array)
dtype = get_dtype(array)
self._ptr = ptr
value_type = from_numpy_dtype(dtype)
# Cache-modified loads only supported for power-of-two sizes up to 16 bytes
# These correspond to PTX instructions: ld.global.{modifier}.b{8,16,32,64,128}
if value_type.size not in (1, 2, 4, 8, 16):
raise ValueError(
f"CacheModifiedInputIterator only supports types of size 1, 2, 4, 8, or 16 bytes. "
f"Got type with size {value_type.size} bytes. "
f"This matches PTX cache-modified load instruction limitations."
)
# State is just the pointer (8 bytes on 64-bit)
state_bytes = struct.pack("Q", ptr)
super().__init__(
state_bytes=state_bytes,
state_alignment=8, # Pointer alignment
value_type=value_type,
)
def _make_advance_op(self) -> Op:
symbol = self._make_advance_symbol()
cpp_type = cpp_type_from_descriptor(self._value_type)
source = dedent(f"""
{CUDA_PREAMBLE}
extern "C" __device__ void {symbol}(void* state, void* offset) {{
auto* s = static_cast<{cpp_type}**>(state);
auto dist = *static_cast<uint64_t*>(offset);
*s += dist;
}}
""").strip()
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[],
)
def _make_input_deref_op(self) -> Op | None:
symbol = self._make_input_deref_symbol()
cpp_type = cpp_type_from_descriptor(self._value_type)
_, intrinsic = _CACHE_MODIFIERS[self._modifier]
# Use cache-modified intrinsic for all supported sizes (1, 2, 4, 8, 16 bytes)
# These correspond to PTX instructions: ld.global.{modifier}.b{8,16,32,64,128}
# Note: __ldcs, __ldcg, __ldcv intrinsics work for all these sizes
source = dedent(f"""
{CUDA_PREAMBLE}
extern "C" __device__ void {symbol}(void* state, void* result) {{
auto* ptr = *static_cast<{cpp_type}**>(state);
*static_cast<{cpp_type}*>(result) = {intrinsic}(ptr);
}}
""").strip()
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[],
)
def _make_output_deref_op(self) -> Op | None:
# Cache-modified iterator is input-only
return None
def __add__(self, offset: int) -> "CacheModifiedInputIterator":
"""Advance the iterator by offset elements."""
out = CacheModifiedInputIterator(self._array, self._modifier)
offset_ptr = self._ptr + offset * get_dtype(out._array).itemsize
out._ptr = offset_ptr
out._state_bytes = struct.pack("Q", offset_ptr)
out._uid_cached = None
return out
@property
def kind(self):
"""Return a hashable kind for caching purposes."""
return (
"CacheModifiedInputIterator",
self._modifier,
self._value_type,
)

21
cccl_upstream/python/cuda_cccl/cuda/compute/iterators/_common.py Normal file
View File

@@ -0,0 +1,21 @@
from __future__ import annotations
from .._utils.protocols import is_device_array
from ._base import IteratorBase
CUDA_PREAMBLE = """#include <cuda/std/cstdint>
#include <cuda_fp16.h>
#include <cuda/std/cstring>
using namespace cuda::std;
"""
def ensure_iterator(obj):
"""Wrap array in PointerIterator if needed."""
from ._pointer import PointerIterator
if isinstance(obj, IteratorBase):
return obj
if is_device_array(obj):
return PointerIterator(obj)
raise TypeError("Expected an iterator or a device array")

103
cccl_upstream/python/cuda_cccl/cuda/compute/iterators/_constant.py Normal file
View File

@@ -0,0 +1,103 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""ConstantIterator implementation."""
from __future__ import annotations
from textwrap import dedent
import numpy as np
from .._bindings import Op, OpKind
from .._cpp_compile import compile_cpp_op_code, cpp_type_from_descriptor
from ..types import from_numpy_dtype
from ._base import IteratorBase
from ._common import CUDA_PREAMBLE
class ConstantIterator(IteratorBase):
"""
Iterator representing a sequence of constant values.
Similar to `thrust::constant_iterator <https://nvidia.github.io/cccl/thrust/api/classthrust_1_1constant__iterator.html>`_.
Every dereference returns the same constant value.
Example:
The code snippet below demonstrates the usage of a ``ConstantIterator``
representing a sequence of constant values:
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/iterator/constant_iterator_basic.py
:language: python
:start-after: # example-begin
Args:
value: The value of every item in the sequence
"""
def __init__(self, value: np.number):
"""
Create a constant iterator with the given value.
Args:
value: The constant value (must be a numpy scalar)
"""
if not isinstance(value, np.generic):
value = np.array(value).flatten()[0]
self._constant_value = value
value_type = from_numpy_dtype(value.dtype)
state_bytes = value.tobytes()
super().__init__(
state_bytes=state_bytes,
state_alignment=value_type.alignment,
value_type=value_type,
)
def _make_advance_op(self) -> Op:
symbol = self._make_advance_symbol()
source = dedent(f"""
{CUDA_PREAMBLE}
extern "C" __device__ void {symbol}(void*, void*) {{
}}
""").strip()
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[],
)
def _make_input_deref_op(self) -> Op | None:
symbol = self._make_input_deref_symbol()
cpp_type = cpp_type_from_descriptor(self._value_type)
source = dedent(f"""
{CUDA_PREAMBLE}
extern "C" __device__ void {symbol}(void* state, void* result) {{
*static_cast<{cpp_type}*>(result) = *static_cast<{cpp_type}*>(state);
}}
""").strip()
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[],
)
def _make_output_deref_op(self) -> Op | None:
return None
def __add__(self, offset: int) -> "ConstantIterator":
"""Return a new ConstantIterator (value doesn't change with position)."""
return ConstantIterator(self._constant_value)

108
cccl_upstream/python/cuda_cccl/cuda/compute/iterators/_counting.py Normal file
View File

@@ -0,0 +1,108 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""CountingIterator implementation."""
from __future__ import annotations
from textwrap import dedent
import numpy as np
from .._bindings import Op, OpKind
from .._cpp_compile import compile_cpp_op_code, cpp_type_from_descriptor
from ..types import from_numpy_dtype
from ._base import IteratorBase
from ._common import CUDA_PREAMBLE
class CountingIterator(IteratorBase):
"""
Iterator representing a sequence of incrementing values.
Similar to `thrust::counting_iterator <https://nvidia.github.io/cccl/thrust/api/classthrust_1_1counting__iterator.html>`_.
The iterator starts at `start` and increments by 1 for each advance.
Example:
The code snippet below demonstrates the usage of a ``CountingIterator``
representing the sequence ``[10, 11, 12]``:
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/iterator/counting_iterator_basic.py
:language: python
:start-after: # example-begin
Args:
start: The initial value of the sequence
"""
def __init__(self, start: np.number):
"""
Create a counting iterator starting at `start`.
Args:
start: The initial value (must be a numpy scalar)
"""
if not isinstance(start, np.generic):
start = np.array(start).flatten()[0]
self._start_value = start
value_type = from_numpy_dtype(start.dtype)
state_bytes = start.tobytes()
super().__init__(
state_bytes=state_bytes,
state_alignment=value_type.alignment,
value_type=value_type,
)
def _make_advance_op(self) -> Op:
symbol = self._make_advance_symbol()
cpp_type = cpp_type_from_descriptor(self._value_type)
source = dedent(f"""
{CUDA_PREAMBLE}
extern "C" __device__ void {symbol}(void* state, void* offset) {{
auto* s = static_cast<{cpp_type}*>(state);
auto dist = *static_cast<uint64_t*>(offset);
*s += static_cast<{cpp_type}>(dist);
}}
""").strip()
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[],
)
def _make_input_deref_op(self) -> Op | None:
symbol = self._make_input_deref_symbol()
cpp_type = cpp_type_from_descriptor(self._value_type)
source = dedent(f"""
{CUDA_PREAMBLE}
extern "C" __device__ void {symbol}(void* state, void* result) {{
*static_cast<{cpp_type}*>(result) = *static_cast<{cpp_type}*>(state);
}}
""").strip()
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[],
)
def _make_output_deref_op(self) -> Op | None:
return None
def __add__(self, offset: int) -> "CountingIterator":
"""Return a new CountingIterator advanced by offset elements."""
new_start = self._start_value + offset
return CountingIterator(new_start)

103
cccl_upstream/python/cuda_cccl/cuda/compute/iterators/_discard.py Normal file
View File

@@ -0,0 +1,103 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""DiscardIterator implementation."""
from __future__ import annotations
from .._bindings import Op, OpKind
from .._cpp_compile import compile_cpp_op_code
from .._utils.protocols import get_dtype
from .._utils.temp_storage_buffer import TempStorageBuffer
from ..types import TypeDescriptor, from_numpy_dtype
from ._base import IteratorBase
from ._common import CUDA_PREAMBLE
class DiscardIterator(IteratorBase):
"""
Iterator that discards all reads and writes.
"""
def __init__(self, reference_iterator=None):
"""
Create a discard iterator.
Args:
reference_iterator: Optional iterator or device array used to infer
value_type/state_type. Defaults to a temporary byte buffer.
"""
if reference_iterator is None:
reference_iterator = TempStorageBuffer(1)
self._reference_iterator = reference_iterator
if hasattr(reference_iterator, "__cuda_array_interface__"):
value_type = from_numpy_dtype(get_dtype(reference_iterator))
state_bytes = bytes(value_type.dtype.itemsize)
elif isinstance(reference_iterator, IteratorBase):
value_type = reference_iterator.value_type
if isinstance(value_type, TypeDescriptor):
state_bytes = bytes(value_type.dtype.itemsize)
else:
state_bytes = bytes(value_type.info.size)
else:
raise TypeError("reference_iterator must be a device array or iterator")
super().__init__(
state_bytes=state_bytes,
state_alignment=value_type.alignment,
value_type=value_type,
)
def _make_advance_op(self) -> Op:
symbol = self._make_advance_symbol()
source = f"""{CUDA_PREAMBLE}
extern "C" __device__ void {symbol}(void*, void*) {{
}}
"""
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[],
)
def _make_input_deref_op(self) -> Op | None:
symbol = self._make_input_deref_symbol()
source = f"""{CUDA_PREAMBLE}
extern "C" __device__ void {symbol}(void*, void*) {{
}}
"""
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[],
)
def _make_output_deref_op(self) -> Op | None:
symbol = self._make_output_deref_symbol()
source = f"""{CUDA_PREAMBLE}
extern "C" __device__ void {symbol}(void*, void*) {{
}}
"""
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[],
)
def __add__(self, offset: int) -> "DiscardIterator":
return DiscardIterator(self._reference_iterator)

220
cccl_upstream/python/cuda_cccl/cuda/compute/iterators/_permutation.py Normal file
View File

@@ -0,0 +1,220 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""PermutationIterator implementation."""
from __future__ import annotations
from textwrap import dedent
from .._bindings import Op, OpKind
from .._cpp_compile import compile_cpp_op_code, make_variable_declaration
from ._base import IteratorBase, compose_iterator_states
from ._common import CUDA_PREAMBLE, ensure_iterator
class PermutationIterator(IteratorBase):
"""
Iterator that accesses values through an index mapping.
At position i, yields values[indices[i]].
Similar to `thrust::permutation_iterator <https://nvidia.github.io/cccl/thrust/api/classthrust_1_1permutation__iterator.html>`_.
Example:
The code snippet below demonstrates accessing values through an index mapping.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/iterator/permutation_iterator_basic.py
:language: python
:start-after: # example-begin
"""
__slots__ = [
"_values",
"_indices",
"_values_offset",
"_indices_offset",
]
def __init__(
self,
values,
indices,
):
"""
Create a permutation iterator.
Args:
values: Iterator or array providing the values to be permuted
indices: Iterator or array providing the indices for permutation
"""
# Wrap arrays in PointerIterator
self._values = ensure_iterator(values)
self._indices = ensure_iterator(indices)
# Compose states from both iterators
state_bytes, state_alignment, offsets = compose_iterator_states(
[self._values, self._indices]
)
self._values_offset = offsets[0]
self._indices_offset = offsets[1]
super().__init__(
state_bytes=state_bytes,
state_alignment=state_alignment,
value_type=self._values.value_type,
)
def _make_advance_op(self) -> Op:
"""Provide Op for advance that only advances indices iterator."""
child_op = self._indices.get_advance_op()
symbol = self._make_advance_symbol()
source = dedent(f"""
{CUDA_PREAMBLE}
extern "C" __device__ void {child_op.name}(void* state, void* offset);
extern "C" __device__ void {symbol}(void* state, void* offset) {{
char* indices_state = static_cast<char*>(state) + {self._indices_offset};
{child_op.name}(indices_state, offset);
}}
""").strip()
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[child_op.code, *child_op.extra_code],
)
def _make_input_deref_op(self) -> Op | None:
"""Provide Op for input deref that reads index then accesses values."""
indices_deref_op = self._indices.get_input_deref_op()
if indices_deref_op is None:
raise ValueError("Indices iterator must support input dereference")
values_deref_op = self._values.get_input_deref_op()
if values_deref_op is None:
return None
# Also need values advance for random access
values_advance_op = self._values.get_advance_op()
symbol = self._make_input_deref_symbol()
idx_decl = make_variable_declaration(self._indices.value_type, "idx")
values_state_size = len(bytes(memoryview(self._values.state)))
source = dedent(f"""
{CUDA_PREAMBLE}
extern "C" __device__ void {indices_deref_op.name}(void* state, void* result);
extern "C" __device__ void {values_advance_op.name}(void* state, void* offset);
extern "C" __device__ void {values_deref_op.name}(void* state, void* result);
extern "C" __device__ void {symbol}(void* state, void* result) {{
char* values_state = static_cast<char*>(state) + {self._values_offset};
char* indices_state = static_cast<char*>(state) + {self._indices_offset};
{idx_decl}
{indices_deref_op.name}(indices_state, &idx);
alignas({self._values.state_alignment}) char temp_values[{values_state_size}];
memcpy(temp_values, values_state, {values_state_size});
uint64_t offset = static_cast<uint64_t>(idx);
{values_advance_op.name}(temp_values, &offset);
{values_deref_op.name}(temp_values, result);
}}
""").strip()
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[
values_advance_op.code,
*values_advance_op.extra_code,
indices_deref_op.code,
*indices_deref_op.extra_code,
values_deref_op.code,
*values_deref_op.extra_code,
],
)
def _make_output_deref_op(self) -> Op | None:
"""Provide Op for output deref that reads index then writes to values."""
indices_deref_op = self._indices.get_input_deref_op()
if indices_deref_op is None:
raise ValueError("Indices iterator must support input dereference")
values_deref_op = self._values.get_output_deref_op()
if values_deref_op is None:
return None
# Also need values advance for random access
values_advance_op = self._values.get_advance_op()
symbol = self._make_output_deref_symbol()
idx_decl = make_variable_declaration(self._indices.value_type, "idx")
values_state_size = len(bytes(memoryview(self._values.state)))
source = dedent(f"""
{CUDA_PREAMBLE}
extern "C" __device__ void {indices_deref_op.name}(void* state, void* result);
extern "C" __device__ void {values_advance_op.name}(void* state, void* offset);
extern "C" __device__ void {values_deref_op.name}(void* state, void* value);
extern "C" __device__ void {symbol}(void* state, void* value) {{
char* values_state = static_cast<char*>(state) + {self._values_offset};
char* indices_state = static_cast<char*>(state) + {self._indices_offset};
{idx_decl}
{indices_deref_op.name}(indices_state, &idx);
alignas({self._values.state_alignment}) char temp_values[{values_state_size}];
memcpy(temp_values, values_state, {values_state_size});
uint64_t offset = static_cast<uint64_t>(idx);
{values_advance_op.name}(temp_values, &offset);
{values_deref_op.name}(temp_values, value);
}}
""").strip()
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[
values_advance_op.code,
*values_advance_op.extra_code,
indices_deref_op.code,
*indices_deref_op.extra_code,
values_deref_op.code,
*values_deref_op.extra_code,
],
)
@property
def children(self):
return (self._values, self._indices)
def __add__(self, offset: int) -> "PermutationIterator":
"""Advance the indices iterator by offset, keeping values at base."""
return PermutationIterator(
self._values, # values stays at base for random access
self._indices + offset, # only indices advances # type: ignore[operator]
)
@property
def kind(self):
"""Return a hashable kind for caching purposes."""
return ("PermutationIterator", self._values.kind, self._indices.kind)

191
cccl_upstream/python/cuda_cccl/cuda/compute/iterators/_pointer.py Normal file
View File

@@ -0,0 +1,191 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""PointerIterator implementation - simple bidirectional iterator for device arrays."""
from __future__ import annotations
import ctypes
import sys
from textwrap import dedent
from .._bindings import Op, OpKind
from .._cpp_compile import compile_cpp_op_code, cpp_type_from_descriptor
from .._utils.protocols import get_data_pointer, get_dtype
from ..types import from_numpy_dtype
from ._base import IteratorBase
from ._common import CUDA_PREAMBLE
class PointerIterator(IteratorBase):
"""
Simple iterator wrapping a device array pointer.
Supports both input (reading) and output (writing) operations.
Handles both scalar types (using typed C++ code) and struct types
(using byte-level memcpy).
"""
def __init__(self, array):
"""
Create a pointer iterator from a device array.
Args:
array: Device array with __cuda_array_interface__
"""
# Get pointer and dtype from array
ptr = get_data_pointer(array)
dtype = get_dtype(array)
value_type = from_numpy_dtype(dtype)
# State is just the pointer
state_bytes = ctypes.c_void_p(ptr)
state_bytes_buffer = (ctypes.c_char * 8)()
ctypes.memmove(state_bytes_buffer, ctypes.byref(state_bytes), 8)
state_bytes = bytes(state_bytes_buffer)
self._cpp_type = cpp_type_from_descriptor(value_type) # None for struct types
self._element_size = value_type.info.size
self._array = array # Keep reference to prevent GC
super().__init__(
state_bytes=state_bytes,
state_alignment=8, # pointer alignment
value_type=value_type,
)
@property
def array(self):
return self._array
def _make_advance_op(self) -> Op:
symbol = self._make_advance_symbol()
if self._cpp_type:
# Scalar type - use typed pointer arithmetic
source = dedent(f"""
{CUDA_PREAMBLE}
extern "C" __device__ void {symbol}(void* state, void* offset) {{
auto* ptr_state = static_cast<{self._cpp_type}**>(state);
auto dist = *static_cast<int64_t*>(offset);
*ptr_state += dist;
}}
""").strip()
else:
# Struct type - use byte-level pointer arithmetic
source = dedent(f"""
{CUDA_PREAMBLE}
extern "C" __device__ void {symbol}(void* state, void* offset) {{
auto* ptr_state = static_cast<char**>(state);
auto dist = *static_cast<int64_t*>(offset);
*ptr_state += dist * {self._element_size};
}}
""").strip()
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[],
)
def _make_input_deref_op(self) -> Op | None:
symbol = self._make_input_deref_symbol()
if self._cpp_type:
# Scalar type - use typed dereference
source = dedent(f"""
{CUDA_PREAMBLE}
extern "C" __device__ void {symbol}(void* state, void* result) {{
auto* ptr_state = static_cast<{self._cpp_type}**>(state);
*static_cast<{self._cpp_type}*>(result) = **ptr_state;
}}
""").strip()
else:
# Struct type - use memcpy
source = dedent(f"""
{CUDA_PREAMBLE}
extern "C" __device__ void {symbol}(void* state, void* result) {{
auto* ptr_state = static_cast<char**>(state);
memcpy(result, *ptr_state, {self._element_size});
}}
""").strip()
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[],
)
def _make_output_deref_op(self) -> Op | None:
symbol = self._make_output_deref_symbol()
if self._cpp_type:
# Scalar type - use typed dereference
source = dedent(f"""
{CUDA_PREAMBLE}
extern "C" __device__ void {symbol}(void* state, void* value) {{
auto* ptr_state = static_cast<{self._cpp_type}**>(state);
**ptr_state = *static_cast<{self._cpp_type}*>(value);
}}
""").strip()
else:
# Struct type - use memcpy
source = dedent(f"""
{CUDA_PREAMBLE}
extern "C" __device__ void {symbol}(void* state, void* value) {{
auto* ptr_state = static_cast<char**>(state);
memcpy(*ptr_state, value, {self._element_size});
}}
""").strip()
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[],
)
def __add__(self, offset: int):
dtype = get_dtype(self._array)
offset_ptr = self._current_pointer() + offset * dtype.itemsize
return self._clone_with_pointer(offset_ptr)
def _current_pointer(self) -> int:
return int.from_bytes(self._state_bytes, sys.byteorder, signed=False)
def _clone_with_pointer(self, pointer_value: int):
"""Clone this iterator with a different pointer value."""
clone = PointerIterator(self._array)
state_bytes_buffer = (ctypes.c_char * 8)()
ptr_obj = ctypes.c_void_p(pointer_value)
ctypes.memmove(state_bytes_buffer, ctypes.byref(ptr_obj), 8)
clone._state_bytes = bytes(state_bytes_buffer)
clone._uid_cached = None
return clone
@property
def kind(self):
"""
Return a hashable kind for caching purposes.
Include _cpp_type and _element_size since they affect generated code.
Different code paths are taken for scalar vs struct types.
"""
return (
type(self).__name__,
self._value_type,
self._cpp_type,
self._element_size,
)

140
cccl_upstream/python/cuda_cccl/cuda/compute/iterators/_reverse.py Normal file
View File

@@ -0,0 +1,140 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""ReverseIterator implementation."""
from __future__ import annotations
from textwrap import dedent
from .._bindings import Op, OpKind
from .._cpp_compile import compile_cpp_op_code
from .._utils.protocols import get_size, is_device_array
from ._base import IteratorBase
from ._common import CUDA_PREAMBLE, ensure_iterator
class ReverseIterator(IteratorBase):
"""
Iterator that reverses the direction of an underlying iterator.
Advance with positive offset moves backward in the underlying iterator.
"""
__slots__ = [
"_underlying",
]
def __init__(self, underlying):
"""
Create a reverse iterator.
Args:
underlying: The underlying iterator or array to reverse
"""
if is_device_array(underlying):
# TODO: this is probably incorrect behaviour. In C++, initializing
# with a pointer to the end of the array is left to be done explicitly
# by the user.
self._underlying = ensure_iterator(underlying) + (get_size(underlying) - 1)
else:
self._underlying = ensure_iterator(underlying)
super().__init__(
state_bytes=bytes(self._underlying.state),
state_alignment=self._underlying.state_alignment,
value_type=self._underlying.value_type,
)
def _make_advance_op(self) -> Op:
"""Provide Op for advance that negates offset direction."""
child_op = self._underlying.get_advance_op()
symbol = self._make_advance_symbol()
source = dedent(f"""
{CUDA_PREAMBLE}
extern "C" __device__ void {child_op.name}(void* state, void* offset);
extern "C" __device__ void {symbol}(void* state, void* offset) {{
int64_t neg_offset = -static_cast<int64_t>(*static_cast<uint64_t*>(offset));
{child_op.name}(state, &neg_offset);
}}
""").strip()
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[child_op.code, *child_op.extra_code],
)
def _make_input_deref_op(self) -> Op | None:
"""Provide Op for input dereference that delegates to underlying."""
child_op = self._underlying.get_input_deref_op()
if child_op is None:
return None
symbol = self._make_input_deref_symbol()
source = dedent(f"""
{CUDA_PREAMBLE}
extern "C" __device__ void {child_op.name}(void* state, void* result);
extern "C" __device__ void {symbol}(void* state, void* result) {{
{child_op.name}(state, result);
}}
""").strip()
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[child_op.code, *child_op.extra_code],
)
def _make_output_deref_op(self) -> Op | None:
"""Provide Op for output dereference that delegates to underlying."""
child_op = self._underlying.get_output_deref_op()
if child_op is None:
return None
symbol = self._make_output_deref_symbol()
source = dedent(f"""
{CUDA_PREAMBLE}
extern "C" __device__ void {child_op.name}(void* state, void* value);
extern "C" __device__ void {symbol}(void* state, void* value) {{
{child_op.name}(state, value);
}}
""").strip()
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[child_op.code, *child_op.extra_code],
)
@property
def children(self):
return (self._underlying,)
@property
def kind(self):
"""Return a hashable kind for caching purposes."""
return ("ReverseIterator", self._underlying.kind)
def __add__(self, offset: int):
return ReverseIterator(self._underlying + offset)

148
cccl_upstream/python/cuda_cccl/cuda/compute/iterators/_shuffle.py Normal file
View File

@@ -0,0 +1,148 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""ShuffleIterator implementation."""
from __future__ import annotations
import struct
from textwrap import dedent
import numpy as np
from .._bindings import Op, OpKind
from .._cpp_compile import compile_cpp_op_code
from ..types import from_numpy_dtype
from ._base import IteratorBase
from ._common import CUDA_PREAMBLE
_SHUFFLE_STATE_STRUCT = """\
struct ShuffleState {
int64_t current_index;
uint64_t num_items;
uint64_t seed;
};"""
class ShuffleIterator(IteratorBase):
"""
Iterator that produces a deterministic random permutation of indices.
At position ``i``, yields ``bijection(i)`` where the bijection is a random
permutation of ``[0, num_items)`` parameterized by ``seed``.
Example:
The code snippet below demonstrates the usage of a ``ShuffleIterator``
to randomly permute indices:
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/iterator/shuffle_iterator_basic.py
:language: python
:start-after: # example-begin
Args:
num_items: Number of elements in the domain to permute. Must be > 0.
seed: Seed for the random permutation. Different seeds produce
different (deterministic) permutations. Defaults to 0.
"""
__slots__ = ["_num_items", "_seed", "_current_index"]
def __init__(self, num_items: int, seed: int = 0, *, _current_index: int = 0):
if num_items <= 0:
raise ValueError("num_items must be > 0")
self._num_items = int(num_items)
self._seed = int(seed)
self._current_index = int(_current_index)
# State layout matches C++ ShuffleState:
# int64_t current_index (offset 0, size 8)
# uint64_t num_items (offset 8, size 8)
# uint64_t seed (offset 16, size 8)
state_bytes = struct.pack(
"<qQQ", self._current_index, self._num_items, self._seed
)
super().__init__(
state_bytes=state_bytes,
state_alignment=8,
value_type=from_numpy_dtype(np.dtype("int64")),
)
def _make_advance_op(self) -> Op:
symbol = self._make_advance_symbol()
source = dedent(f"""
{CUDA_PREAMBLE}
{_SHUFFLE_STATE_STRUCT}
extern "C" __device__ void {symbol}(void* state, void* offset) {{
auto s = static_cast<ShuffleState*>(state);
auto dist = *static_cast<int64_t*>(offset);
s->current_index += dist;
}}
""")
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[],
)
def _make_input_deref_op(self) -> Op | None:
symbol = self._make_input_deref_symbol()
# Note: a potential optimization is to avoid constructing
# `cuda::random_bijection` objects upon every dereference,
# instead constructing it once and using it as the state
# object. The tradeoff is that it would require a C++
# extension providing a constructor for
# `cuda::random_bijection` objects, since we would now be
# doing it on the host. See discussion in #7721.
source = dedent(f"""
#include <cuda/__random/random_bijection.h>
#include <cuda/__random/pcg_engine.h>
{CUDA_PREAMBLE}
{_SHUFFLE_STATE_STRUCT}
// __noinline__ is required to prevent the compiler from merging
// this function's register usage into the calling kernel during LTO
// inlining. feistel_bijection constructs 24 round keys with
// UNROLL_FULL, which exhausts the kernel's register budget and
// causes spilling to local memory (LDL/STL instructions).
// Keeping it non-inlined gives it an isolated register frame.
__device__ __noinline__ int64_t __shuffle_apply(uint64_t num_items, uint64_t seed, uint64_t idx) {{
cuda::pcg64 rng(seed);
cuda::random_bijection<uint64_t> bijection(num_items, rng);
return static_cast<int64_t>(bijection(idx));
}}
extern "C" __device__ void {symbol}(void* state, void* result) {{
const auto* s = static_cast<const ShuffleState*>(state);
*static_cast<int64_t*>(result) = __shuffle_apply(
s->num_items, s->seed, static_cast<uint64_t>(s->current_index));
}}
""")
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[],
)
def _make_output_deref_op(self) -> Op | None:
return None
def __add__(self, offset: int) -> "ShuffleIterator":
return ShuffleIterator(
self._num_items,
self._seed,
_current_index=self._current_index + offset,
)

Some files were not shown because too many files have changed in this diff Show More