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

[CCCL] 瘦身 + 补全: 移除 cudax/python/libcudacxx-tests 冗余文件, 新增 c2h 测试助手 + cmake 构建系统 + 8 个 CUDA thrust examples

Browse Source 
变更摘要:
- 删除: cudax/ (783 files, 7.2M) — 实验性组件，竞赛不需要
- 删除: python/ (226 files, 2.0M) — Python 绑定，竞赛不需要
- 删除: libcudacxx/{test,benchmarks,codegen,cmake,share} (4432 files, 31M)
  保留: libcudacxx/include/ (1463 headers, cuda::std 编译依赖)
- 新增: c2h/ (27 files) — CUB Catch2 测试辅助头文件，编译 243 个测试必需
- 新增: cmake/ (29 files) — CCCL 原生 CMake 构建系统
- 新增: thrust/examples/cuda/ (7 files) + cpp_integration/ (1 file)
  async_reduce, custom_temporary_allocation, explicit_cuda_stream,
  global_device_vector, range_view, unwrap_pointer, wrap_pointer, device

结果: cccl_upstream 从 74M→35M (瘦身 53%), 核心内容 100% 保留:
  27/27 tuning headers, 78 benchmarks, 243 tests,
  60 thrust examples, 18 CUB examples, 全部编译头文件
This commit is contained in:
muh-bot
2026-08-03 12:39:26 +00:00
parent a2a5dd8f00
commit 24ef6a91b5
5439 changed files with 0 additions and 719516 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,27 +0,0 @@
"""
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
View File

@@ -1,24 +0,0 @@
# 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
View File

@@ -1,7 +0,0 @@
# 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
View File

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

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

@@ -1,51 +0,0 @@
# 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
View File

@@ -1,9 +0,0 @@
# 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
View File

@@ -1,24 +0,0 @@
# 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
View File

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

@@ -1,162 +0,0 @@
# 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
View File

@@ -1,87 +0,0 @@
# 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
View File

@@ -1,658 +0,0 @@
# 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
View File

@@ -1,17 +0,0 @@
# 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
View File

@@ -1,15 +0,0 @@
# 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
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
View File

@@ -1,14 +0,0 @@
# 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
View File

@@ -1,17 +0,0 @@
# 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
View File

@@ -1,40 +0,0 @@
# 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
View File

@@ -1,38 +0,0 @@
# 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
View File

File diff suppressed because it is too large Load Diff

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

@@ -1,57 +0,0 @@
# 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
View File

@@ -1,740 +0,0 @@
# 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
View File

@@ -1,471 +0,0 @@
# 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
View File

@@ -1,172 +0,0 @@
# 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
View File

@@ -1,45 +0,0 @@
# 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
View File

File diff suppressed because it is too large Load Diff

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

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

@@ -1,370 +0,0 @@
# 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
View File

@@ -1,132 +0,0 @@
# 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
View File

@@ -1,41 +0,0 @@
# 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
View File

@@ -1,217 +0,0 @@
# 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
View File

@@ -1,60 +0,0 @@
# 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
View File

@@ -1,285 +0,0 @@
# 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
View File

@@ -1,58 +0,0 @@
# 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
View File

@@ -1,22 +0,0 @@
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
View File

@@ -1,168 +0,0 @@
# 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
View File

@@ -1,94 +0,0 @@
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
View File

@@ -1,76 +0,0 @@
# 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
View File

@@ -1,345 +0,0 @@
# 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
View File

@@ -1,330 +0,0 @@
# 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
View File

@@ -1,277 +0,0 @@
# 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
View File

@@ -1,436 +0,0 @@
# 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
View File

@@ -1,293 +0,0 @@
# 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
View File

@@ -1,239 +0,0 @@
# 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
View File

@@ -1,25 +0,0 @@
# 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
View File

@@ -1,261 +0,0 @@
# 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
View File

@@ -1,289 +0,0 @@
# 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
View File

@@ -1,298 +0,0 @@
# 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
View File

@@ -1,62 +0,0 @@
# 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
View File

@@ -1,289 +0,0 @@
# 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
View File

@@ -1,371 +0,0 @@
# 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
View File

@@ -1,253 +0,0 @@
# 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
View File

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

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

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

@@ -1,31 +0,0 @@
# 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
View File

@@ -1,293 +0,0 @@
# 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
View File

@@ -1,156 +0,0 @@
# 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
View File

@@ -1,21 +0,0 @@
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
View File

@@ -1,103 +0,0 @@
# 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
View File

@@ -1,108 +0,0 @@
# 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
View File

@@ -1,103 +0,0 @@
# 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
View File

@@ -1,220 +0,0 @@
# 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
View File

@@ -1,191 +0,0 @@
# 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
View File

@@ -1,140 +0,0 @@
# 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
View File

@@ -1,148 +0,0 @@
# 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,
)

292
cccl_upstream/python/cuda_cccl/cuda/compute/iterators/_transform.py
View File

@@ -1,292 +0,0 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""TransformIterator 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 ..op import make_op_adapter
from ..types import TypeDescriptor, signature_from_annotations
from ._base import IteratorBase
from ._common import CUDA_PREAMBLE, ensure_iterator
class TransformIterator(IteratorBase):
"""
An iterator that applies a unary function to elements as they are read from an underlying iterator.
Similar to `thrust::transform_iterator <https://nvidia.github.io/cccl/thrust/api/classthrust_1_1transform__iterator.html>`_.
For input iteration (default): reads from underlying, applies transform, returns result.
For output iteration: applies transform to input values, writes to underlying.
Example:
The code snippet below demonstrates the usage of a ``TransformIterator`` composed with a ``CountingIterator``
to transform the input before performing a reduction:
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/iterator/transform_iterator_basic.py
:language: python
:start-after: # example-begin
Args:
underlying: The underlying iterator or device array
transform_op: The unary operation to apply
output_value_type: TypeDescriptor for the output type (optional, will be inferred if not provided)
is_input: True for input iterator (default), False for output iterator
"""
__slots__ = [
"_underlying",
"_transform_op",
"_value_type",
"_is_input",
"_compiled_op",
]
def __init__(
self,
underlying,
transform_op,
value_type: TypeDescriptor | None = None,
is_input: bool = True,
):
"""
Create a transform iterator.
Args:
underlying: The underlying iterator or device array to transform
transform_op: The unary transform operation (callable or OpKind)
value_type: TypeDescriptor for the transformed value type.
For input iterators: inferred if None.
For output iterators: must be provided or have annotations.
is_input: True for input iterator, False for output iterator
"""
underlying = ensure_iterator(underlying)
self._underlying = underlying
self._transform_op = make_op_adapter(transform_op)
self._is_input = is_input
# Lazily compiled transform Op, keyed on the build's target compute
# capability: its LTO-IR is arch-specific, so reusing one iterator
# instance across builds targeting different arches must not reuse the
# first arch's op (nvJitLink rejects a newer-arch input linked into an
# older-arch result). Mirrors the per-cc op caches in IteratorBase.
self._compiled_op: dict = {}
# Determine value type
if value_type is None:
if is_input:
# value_type is the return type
_, value_type = signature_from_annotations(transform_op)
if value_type is None:
value_type = self._transform_op.get_return_type(
(underlying.value_type,)
)
else:
# value_type is the input type
input_types, _ = signature_from_annotations(transform_op)
if len(input_types) != 1:
raise ValueError(
"TransformOutputIterator transform function must take exactly one argument with type annotation"
)
value_type = input_types[0]
assert value_type is not None
self._value_type = value_type
super().__init__(
state_bytes=bytes(self._underlying.state),
state_alignment=self._underlying.state_alignment,
value_type=value_type,
)
def _get_compiled_op(self):
"""Get the compiled Op for the current target cc, compiling lazily if needed."""
from .._target_cc import get_target_cc
key = get_target_cc()
if key not in self._compiled_op:
if self._is_input:
input_type = self._underlying.value_type
output_type = self._value_type
else:
input_type = self._value_type
output_type = self._underlying.value_type
self._compiled_op[key] = self._transform_op.compile(
(input_type,),
output_type,
)
return self._compiled_op[key]
def _make_advance_op(self) -> Op:
"""Provide Op for advance that delegates to underlying iterator."""
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) {{
{child_op.name}(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 dereference that reads from underlying then transforms."""
if not self._is_input:
return None
child_op = self._underlying.get_input_deref_op()
if child_op is None:
raise ValueError("Underlying iterator must support input dereference")
compiled_op = self._get_compiled_op()
symbol = self._make_input_deref_symbol()
temp_decl = make_variable_declaration(self._underlying.value_type, "temp")
source = dedent(f"""
{CUDA_PREAMBLE}
extern "C" __device__ void {child_op.name}(void* state, void* result);
extern "C" __device__ void {compiled_op.name}(void* input, void* output);
extern "C" __device__ void {symbol}(void* state, void* result) {{
{temp_decl}
{child_op.name}(state, &temp);
{compiled_op.name}(&temp, result);
}}
""").strip()
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[
compiled_op.code,
*compiled_op.extra_code,
child_op.code,
*child_op.extra_code,
],
)
def _make_output_deref_op(self) -> Op | None:
"""Provide Op for output dereference that transforms then writes to underlying."""
if self._is_input:
return None
child_op = self._underlying.get_output_deref_op()
if child_op is None:
raise ValueError("Underlying iterator must support output dereference")
compiled_op = self._get_compiled_op()
symbol = self._make_output_deref_symbol()
temp_decl = make_variable_declaration(self._underlying.value_type, "temp")
source = dedent(f"""
{CUDA_PREAMBLE}
extern "C" __device__ void {child_op.name}(void* state, void* value);
extern "C" __device__ void {compiled_op.name}(void* input, void* output);
extern "C" __device__ void {symbol}(void* state, void* value) {{
{temp_decl}
{compiled_op.name}(value, &temp);
{child_op.name}(state, &temp);
}}
""").strip()
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[
compiled_op.code,
*compiled_op.extra_code,
child_op.code,
*child_op.extra_code,
],
)
def advance(self, offset: int) -> "TransformIterator":
"""Return a new iterator advanced by offset elements."""
if not hasattr(self._underlying, "__add__"):
raise AttributeError("Underlying iterator does not support advance")
return TransformIterator(
self._underlying + offset, # type: ignore[operator, arg-type]
self._transform_op,
self._value_type,
is_input=self._is_input,
)
def __add__(self, offset: int) -> "TransformIterator":
return self.advance(offset)
def __radd__(self, offset: int) -> "TransformIterator":
return self.advance(offset)
@property
def children(self):
return (self._underlying,)
@property
def kind(self):
"""Return a hashable kind for caching purposes."""
# Convert _value_type to tuple if it's a list (for output iterators)
value_type = (
tuple(self._value_type)
if isinstance(self._value_type, list)
else self._value_type
)
return (
"TransformIterator",
self._is_input,
self._transform_op,
self._underlying.kind,
value_type,
)
class TransformOutputIterator(TransformIterator):
"""
An iterator that applies a unary function to values before writing them to an underlying iterator.
Similar to `thrust::transform_output_iterator <https://nvidia.github.io/cccl/thrust/api/classthrust_1_1transform__output__iterator.html>`_.
This is a convenience subclass of TransformIterator configured for output mode.
Example:
The code snippet below demonstrates the usage of a ``TransformOutputIterator`` to transform the output
of a reduction before writing to an output array:
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/iterator/transform_output_iterator.py
:language: python
:start-after: # example-begin
Args:
underlying: The underlying iterator or device array
transform_op: The operation to be applied to values before they are written
output_value_type: TypeDescriptor for the input value type (optional, will be extracted from annotations if not provided)
"""
def __init__(self, underlying, transform_op, output_value_type=None):
super().__init__(underlying, transform_op, output_value_type, is_input=False)

1
cccl_upstream/python/cuda_cccl/cuda/compute/iterators/_utils.py
View File

@@ -1 +0,0 @@
from __future__ import annotations

216
cccl_upstream/python/cuda_cccl/cuda/compute/iterators/_zip.py
View File

@@ -1,216 +0,0 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""ZipIterator implementation."""
from __future__ import annotations
from textwrap import dedent
from .._bindings import Op, OpKind
from .._cpp_compile import compile_cpp_op_code
from ..types import struct
from ._base import IteratorBase, compose_iterator_states
from ._common import CUDA_PREAMBLE, ensure_iterator
class ZipIterator(IteratorBase):
"""
Iterator that zips multiple iterators together.
At each position, yields a tuple of values from all underlying iterators.
Similar to `thrust::zip_iterator <https://nvidia.github.io/cccl/thrust/api/classthrust_1_1zip__iterator.html>`_.
Example:
The code snippet below demonstrates how to zip together an array and a
:class:`CountingIterator <cuda.compute.iterators.CountingIterator>` to
find the index of the maximum value of the array.
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/iterator/zip_iterator_counting.py
:language: python
:start-after: # example-begin
"""
__slots__ = [
"_iterators",
"_field_names",
"_value_offsets",
"_state_offsets",
"_advance_result",
"_input_deref_result",
"_output_deref_result",
]
def __init__(self, *args):
"""
Create a zip iterator.
Args:
*args: Iterators or arrays to zip together. Can be:
- Multiple iterators/arrays: ZipIterator(it1, it2, it3)
- A single sequence of iterators: ZipIterator([it1, it2, it3])
"""
# Handle both ZipIterator(it1, it2) and ZipIterator([it1, it2])
if len(args) == 1 and isinstance(args[0], (list, tuple)):
iterators = args[0]
else:
iterators = args
if len(iterators) < 1:
raise ValueError("ZipIterator requires at least one iterator")
# Wrap arrays in PointerIterator
iterators = [ensure_iterator(it) for it in iterators]
self._iterators = list(iterators)
# Compose states from all iterators
self._state_bytes, self._state_alignment, self._state_offsets = (
compose_iterator_states(self._iterators)
)
# Build combined value type (struct layout)
self._field_names = [f"field_{i}" for i in range(len(self._iterators))]
fields = {
name: it.value_type for name, it in zip(self._field_names, self._iterators)
}
self._value_type = struct(fields, name=f"Zip{len(iterators)}")
self._value_offsets = [
self._value_type.dtype.fields[name][1] for name in self._field_names
]
super().__init__(
state_bytes=self._state_bytes,
state_alignment=self._state_alignment,
value_type=self._value_type,
)
def _make_advance_op(self) -> Op:
"""Provide Op for advance that calls all child iterator advances."""
child_ops = [it.get_advance_op() for it in self._iterators]
symbol = self._make_advance_symbol()
externs = "\n".join(
f'extern "C" __device__ void {op.name}(void* state, void* offset);'
for op in child_ops
)
calls = "\n ".join(
f"{op.name}(static_cast<char*>(state) + {offset}, offset);"
for op, offset in zip(child_ops, self._state_offsets)
)
source = dedent(f"""
{CUDA_PREAMBLE}
{externs}
extern "C" __device__ void {symbol}(void* state, void* offset) {{
{calls}
}}
""").strip()
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[c for op in child_ops for c in [op.code, *op.extra_code]],
)
def _make_input_deref_op(self) -> Op | None:
"""Provide Op for input deref that calls all child iterator input derefs."""
child_ops = [it.get_input_deref_op() for it in self._iterators]
if not all(op is not None for op in child_ops):
return None
symbol = self._make_input_deref_symbol()
externs = "\n".join(
f'extern "C" __device__ void {op.name}(void* state, void* result);'
for op in child_ops
)
calls = "\n ".join(
f"{op.name}(static_cast<char*>(state) + {state_off}, "
f"static_cast<char*>(result) + {val_off});"
for op, state_off, val_off in zip(
child_ops, self._state_offsets, self._value_offsets
)
)
source = dedent(f"""
{CUDA_PREAMBLE}
{externs}
extern "C" __device__ void {symbol}(void* state, void* result) {{
{calls}
}}
""").strip()
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[c for op in child_ops for c in [op.code, *op.extra_code]],
)
def _make_output_deref_op(self) -> Op | None:
"""Provide Op for output deref that calls all child iterator output derefs."""
child_ops = [it.get_output_deref_op() for it in self._iterators]
if not all(op is not None for op in child_ops):
return None
symbol = self._make_output_deref_symbol()
externs = "\n".join(
f'extern "C" __device__ void {op.name}(void* state, void* value);'
for op in child_ops
)
calls = "\n ".join(
f"{op.name}(static_cast<char*>(state) + {state_off}, "
f"static_cast<char*>(value) + {val_off});"
for op, state_off, val_off in zip(
child_ops, self._state_offsets, self._value_offsets
)
)
source = dedent(f"""
{CUDA_PREAMBLE}
{externs}
extern "C" __device__ void {symbol}(void* state, void* value) {{
{calls}
}}
""").strip()
code = compile_cpp_op_code(source)
return Op(
operator_type=OpKind.STATELESS,
name=symbol,
ltoir=code,
extra_ltoirs=[c for op in child_ops for c in [op.code, *op.extra_code]],
)
@property
def children(self):
return tuple(self._iterators)
def __add__(self, offset: int) -> "ZipIterator":
"""Advance all child iterators by offset."""
advanced_iterators = [it + offset for it in self._iterators] # type: ignore[operator]
return ZipIterator(*advanced_iterators)
@property
def kind(self):
"""Return a hashable kind for caching purposes."""
return ("ZipIterator", tuple(it.kind for it in self._iterators))

253
cccl_upstream/python/cuda_cccl/cuda/compute/op.py
View File

@@ -1,253 +0,0 @@
# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
from ._bindings import Op, OpKind
from ._caching import CachableFunction, cache_with_registered_key_functions
from ._device_code import DeviceCode
def _is_well_known_op(op: OpKind) -> bool:
return isinstance(op, OpKind) and op not in (OpKind.STATELESS, OpKind.STATEFUL)
class _OpAdapter:
"""
Provides a unified interface for operators, whether they are:
- Well-known operations (OpKind.PLUS, OpKind.MAXIMUM, etc.)
- Stateless user-provided callables
- Stateful user-provided callables
"""
def compile(self, input_types, output_type=None) -> Op:
"""
Compile this operator to an Op for CCCL interop.
Args:
input_types: Tuple of TypeDescriptors for input arguments
output_type: Optional TypeDescriptor for return value (inferred if None)
Returns:
Compiled Op object for C++ interop
"""
raise NotImplementedError("Subclasses must implement this method")
@property
def is_stateful(self) -> bool:
"""Return True if this op has runtime state."""
return False
def get_state(self) -> bytes:
"""
Return the op's state bytes.
"""
return b""
def get_return_type(self, input_types):
"""Get the return type for this op given input types."""
raise NotImplementedError(
f"get_return_type not implemented for {self.__class__.__name__}"
)
class _WellKnownOp(_OpAdapter):
"""Internal wrapper for well-known OpKind values."""
__slots__ = ["_kind"]
def __init__(self, kind: OpKind):
if not _is_well_known_op(kind):
raise ValueError(
f"OpKind.{kind.name} is not a well-known operation. "
"Use OpKind.PLUS, OpKind.MAXIMUM, etc."
)
self._kind = kind
def compile(self, input_types, output_type=None) -> Op:
return Op(
operator_type=self._kind,
name="",
ltoir=b"",
state_alignment=1,
state=b"",
)
@property
def kind(self) -> OpKind:
"""The underlying OpKind."""
return self._kind
def __eq__(self, other):
if not isinstance(other, _WellKnownOp):
return False
return self._kind == other._kind
def __hash__(self):
return hash(self._kind)
class RawOp(_OpAdapter):
"""
``RawOp`` lets you supply pre-compiled device code (LTO-IR) implementing a
custom operator, bypassing the default Numba-based JIT pipeline.
Example:
Supplying C++ device code compiled to LTO-IR via NVRTC:
.. literalinclude:: ../../python/cuda_cccl/tests/compute/examples/raw_op/cpp_stateless.py
:language: python
:start-after: # example-begin
Args:
name: The ABI name of the operator.
ltoir: Raw ``bytes`` of pre-compiled LTO-IR implementing the operator
(for example, produced by ``nvcc -dlto`` or NVRTC).
state: Optional bytes representing the operator's state.
state_alignment: Alignment requirement for the state bytes (default: 1).
extra_ltoirs: Optional list of additional LTO-IR ``bytes`` to link.
Notes:
- The provided code must define a function with the specified name and the correct signature.
- The function must use untyped pointers for all parameters and return type. The function body
is responsible for correctly interpreting the pointer arguments based on the expected input and output types.
For stateless operators, the signature is
void func(void* arg1, void* arg2, ..., void* result)`
For stateful operators, the first parameter must be a pointer to the state:
void func(void* state, void* arg1, void* arg2, ...)
"""
__slots__ = [
"_ltoir",
"_name",
"_state",
"_state_alignment",
"_extra_ltoirs",
]
def __init__(
self,
*,
ltoir: bytes | DeviceCode,
name: str,
state: bytes = b"",
state_alignment: int = 1,
extra_ltoirs: list[bytes | DeviceCode] | None = None,
):
self._ltoir = ltoir
self._name = name
self._state = state
self._state_alignment = state_alignment
self._extra_ltoirs = extra_ltoirs or []
def compile(self, input_types, output_type=None) -> Op:
# Determine if stateful based on whether state is provided
op_kind = OpKind.STATEFUL if self._state else OpKind.STATELESS
return Op(
operator_type=op_kind,
name=self._name,
ltoir=self._ltoir,
state=self._state,
state_alignment=self._state_alignment,
extra_ltoirs=self._extra_ltoirs,
)
def get_state(self) -> bytes:
"""Return the op's state bytes."""
return self._state
@property
def _identity(self):
return (
self._ltoir,
self._name,
self._state,
self._state_alignment,
tuple(self._extra_ltoirs),
)
def __eq__(self, other):
if not isinstance(other, RawOp):
return False
return self._identity == other._identity
def __hash__(self):
return hash(self._identity)
# Public aliases
OpAdapter = _OpAdapter
def _jit_op_adapter_factory():
# helper that tries to import `_jit.py`. If it fails,
# returns a function that raises an appropriate error when called.
try:
from ._jit import to_jit_op_adapter
return to_jit_op_adapter
except ModuleNotFoundError as e:
if "numba" in str(e):
def _missing_jit_adapter(op):
raise ImportError(
"numba-cuda is required to JIT compile Python callables"
)
return _missing_jit_adapter
raise
to_jit_op_adapter = _jit_op_adapter_factory()
def make_op_adapter(op) -> OpAdapter:
"""
Create an Op from a callable or well-known OpKind.
Args:
op: Callable or OpKind
Returns:
A value with appropriate subtype of _BaseOp
"""
# Already an _OpAdapter instance:
if isinstance(op, _OpAdapter):
return op
# Well-known operation
if isinstance(op, OpKind):
return _WellKnownOp(op)
# It's a Python callable
return to_jit_op_adapter(op)
cache_with_registered_key_functions.register(
_WellKnownOp, lambda op: (op._kind.name, op._kind.value)
)
cache_with_registered_key_functions.register(
OpKind, lambda kind: (kind.name, kind.value)
)
cache_with_registered_key_functions.register(
type(lambda: None), lambda func: CachableFunction(func)
)
cache_with_registered_key_functions.register(RawOp, lambda op: op._identity)
__all__ = [
"OpAdapter",
"OpKind",
"make_op_adapter",
"RawOp",
]

194
cccl_upstream/python/cuda_cccl/cuda/compute/struct.py
View File

@@ -1,194 +0,0 @@
# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
import functools
from types import new_class
from typing import Any, ClassVar, TypeGuard, Union, cast, get_type_hints
import numpy as np
from . import types
"""
This module provides `gpu_struct`, a factory for producing struct types.
"""
def gpu_struct(
field_dict: Union[dict, np.dtype, type],
name: str = "AnonymousStruct",
):
"""
A factory for creating struct types.
Args:
field_dict
A dictionary, numpy dtype, or annotated class providing the
mapping of field names to data types.
name
The name of the struct type that will be returned.
Returns:
A struct class helpful for writing operations on struct values.
"""
# Handle numpy dtype input
if isinstance(field_dict, np.dtype):
if field_dict.type != np.void or field_dict.fields is None:
field_dict = {}
else:
field_dict = {
name: field_info[0] for name, field_info in field_dict.fields.items()
}
# Handle annotated class (decorator usage)
if isinstance(field_dict, type) and hasattr(field_dict, "__annotations__"):
name = field_dict.__name__
field_dict = get_type_hints(field_dict)
# At this point, field_dict must be a dict
assert isinstance(field_dict, dict)
# Validate field names are valid Python identifiers
for key in field_dict:
if not isinstance(key, str) or not key.isidentifier():
raise ValueError(
f"gpu_struct field name {key!r} is not a valid Python identifier"
)
# Normalize fields for storage on the struct class
field_spec = {}
for key, val in field_dict.items():
if _is_struct_type(val):
field_spec[key] = val
elif isinstance(val, dict):
# Nested struct definition - recursively create inner struct
field_spec[key] = gpu_struct(val, name=key)
else:
field_spec[key] = val
# Create a simple Python class for user-facing struct values
struct_class = cast(type[_Struct], new_class(name, bases=(_Struct,)))
struct_class._field_spec = field_spec
struct_class._type_descriptor = _get_struct_type_descriptor(struct_class) # type: ignore[arg-type]
struct_class.dtype = _get_struct_record_dtype(struct_class) # type: ignore[arg-type]
return struct_class
class _Struct:
"""Internal base class for all gpu_structs."""
_field_spec: ClassVar[dict[str, Any]]
_type_descriptor: ClassVar[types.StructTypeDescriptor]
dtype: ClassVar[np.dtype]
_fields: dict[str, Any]
@classmethod
def _fields_from_args(cls, *args, **kwargs):
field_spec = cls._field_spec
if args and isinstance(args[0], dict):
fields = args[0]
elif args:
assert len(args) == len(field_spec), (
f"Expected {len(field_spec)} arguments, got {len(args)}"
)
fields = dict(zip(field_spec.keys(), args))
else:
fields = kwargs
assert fields.keys() == field_spec.keys()
return {
name: _coerce_value(field_spec[name], fields[name]) for name in field_spec
}
def __init__(self, *args, **kwargs):
"""Supporting construction from positional, keyword, and dict arguments."""
self._fields = self._fields_from_args(*args, **kwargs)
for name, value in self._fields.items():
setattr(self, name, value)
# NumPy array representation:
self._data = np.asarray(_as_numpy_record_value(self))
self.__array_interface__ = self._data.__array_interface__
def _as_numpy_record_value(val) -> np.void:
"""Convert a gpu_struct *value* to a numpy record."""
def _fields_to_tuples(fields_dict: dict[str, Any]) -> tuple[Any, ...]:
return tuple(
_fields_to_tuples(v._fields) if isinstance(v, _Struct) else v
for v in fields_dict.values()
)
return np.void(
_fields_to_tuples(val._fields),
dtype=_get_struct_record_dtype(type(val)), # type: ignore[arg-type]
)
@functools.cache
def _get_struct_record_dtype(struct_class: type) -> np.dtype:
return _get_struct_type_descriptor(struct_class).dtype
def _coerce_value(field_type, value: Any) -> Any:
if isinstance(value, _Struct):
return value
if isinstance(field_type, np.dtype):
return field_type.type(value)
if isinstance(field_type, type) and issubclass(field_type, np.generic):
return field_type(value) # type: ignore[call-arg]
if isinstance(value, tuple):
return field_type(*value)
if isinstance(value, dict):
return field_type(**value)
# field_type is a class (e.g., another gpu_struct)
raise TypeError(f"Cannot coerce {type(value).__name__} into {field_type.__name__}")
def _is_struct_type(typ: Any) -> TypeGuard[type[_Struct]]:
"""Check if a type is a GPU struct class."""
return isinstance(typ, type) and issubclass(typ, _Struct)
@functools.cache
def _get_struct_type_descriptor(
struct_class: type,
) -> types.StructTypeDescriptor:
type_descriptors = _field_spec_to_type_descriptors(
struct_class._field_spec # type: ignore[attr-defined]
)
return types.struct(type_descriptors, name=struct_class.__name__)
def _field_spec_to_type_descriptors(
field_spec: dict[str, Any],
) -> dict[str, types.TypeDescriptor]:
type_descriptors = {}
for key, val in field_spec.items():
if isinstance(val, types.TypeDescriptor):
type_descriptors[key] = val
elif _is_struct_type(val):
type_descriptors[key] = val._type_descriptor
elif isinstance(val, np.dtype):
type_descriptors[key] = types.from_numpy_dtype(val)
else:
type_descriptors[key] = types.from_numpy_dtype(np.dtype(val))
return type_descriptors

292
cccl_upstream/python/cuda_cccl/cuda/compute/types.py
View File

@@ -1,292 +0,0 @@
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from __future__ import annotations
import inspect
from typing import get_type_hints
import numpy as np
from ._bindings import TypeEnum, TypeInfo
_ENUM_TO_DTYPE: dict[TypeEnum, np.dtype] = {
TypeEnum.INT8: np.dtype("int8"),
TypeEnum.INT16: np.dtype("int16"),
TypeEnum.INT32: np.dtype("int32"),
TypeEnum.INT64: np.dtype("int64"),
TypeEnum.UINT8: np.dtype("uint8"),
TypeEnum.UINT16: np.dtype("uint16"),
TypeEnum.UINT32: np.dtype("uint32"),
TypeEnum.UINT64: np.dtype("uint64"),
TypeEnum.FLOAT16: np.dtype("float16"),
TypeEnum.FLOAT32: np.dtype("float32"),
TypeEnum.FLOAT64: np.dtype("float64"),
TypeEnum.BOOLEAN: np.dtype("bool"),
}
class TypeDescriptor:
def __init__(self, size: int, alignment: int, type_enum: TypeEnum):
self._type_info = TypeInfo(size, alignment, type_enum)
self._dtype = _ENUM_TO_DTYPE.get(type_enum, np.dtype(f"V{size}"))
@property
def info(self) -> TypeInfo:
"""Return the TypeInfo for this type."""
return self._type_info
@property
def size(self):
return self._type_info.size
@property
def alignment(self):
return self._type_info.alignment
@property
def dtype(self) -> np.dtype:
"""Return the numpy dtype for this type."""
return self._dtype
def pointer(self) -> "PointerTypeDescriptor":
"""Create a pointer type to this type."""
return PointerTypeDescriptor(self)
def __repr__(self) -> str:
return f"TypeDescriptor({self._dtype})"
def __eq__(self, other):
if not isinstance(other, TypeDescriptor):
return False
return self._dtype == other._dtype
def __hash__(self):
return hash(self._dtype)
class StructTypeDescriptor(TypeDescriptor):
def __init__(
self,
fields: dict[str, "TypeDescriptor"],
name: str = "AnonStruct",
):
dtype = _build_struct_dtype(fields)
if not dtype.isalignedstruct:
raise ValueError(f"dtype {dtype} must be aligned")
# Structs use STORAGE type enum
super().__init__(dtype.itemsize, dtype.alignment, TypeEnum.STORAGE)
self._dtype = dtype
self._fields = fields
self._name = name
def __repr__(self) -> str:
return f"StructTypeDescriptor({self._name}, dtype={self._dtype})"
@property
def name(self) -> str:
return self._name
@property
def fields(self) -> dict[str, "TypeDescriptor"]:
return self._fields
def layout_key(self) -> tuple[tuple[str, "TypeDescriptor"], ...]:
"""Return a stable, hashable key for this struct layout."""
return tuple(self._fields.items())
def __eq__(self, other):
if not isinstance(other, StructTypeDescriptor):
return False
# Compare by fields (TypeDescriptors) to correctly distinguish structs
# with different pointer pointee types, which have identical numpy dtypes
# (both uint64) but different semantics.
# Must compare as ordered items because field order matters for struct layout.
return list(self._fields.items()) == list(other._fields.items())
def __hash__(self):
return hash(tuple(self._fields.items()))
class PointerTypeDescriptor(TypeDescriptor):
def __init__(self, pointee: TypeDescriptor):
# Pointer is 8 bytes, 8-byte aligned, represented as uint64
super().__init__(8, 8, TypeEnum.UINT64)
self._pointee = pointee
@property
def pointee(self) -> TypeDescriptor:
"""Return the type this pointer points to."""
return self._pointee
def __repr__(self) -> str:
return f"PointerTypeDescriptor({self._pointee})"
def __eq__(self, other):
if not isinstance(other, PointerTypeDescriptor):
return False
return self._pointee == other._pointee
def __hash__(self):
return hash(("PointerTypeDescriptor", self._pointee))
def _build_struct_dtype(
fields: dict[str, TypeDescriptor],
) -> np.dtype:
dtype_list = []
for field_name, field_type in fields.items():
dtype_list.append((field_name, field_type.dtype))
return np.dtype(dtype_list, align=True)
def struct(
fields: dict[str, TypeDescriptor],
name: str = "AnonStruct",
) -> StructTypeDescriptor:
"""Create a type descriptor for a struct."""
return StructTypeDescriptor(fields, name=name)
def pointer(pointee: TypeDescriptor) -> PointerTypeDescriptor:
"""
Create a pointer to the given type.
"""
return PointerTypeDescriptor(pointee)
def from_numpy_dtype(dtype: np.dtype | type) -> TypeDescriptor:
"""
Convert a numpy dtype (or numpy type) to a TypeDescriptor.
Handles POD types and structured dtypes (recursively for nested structs).
"""
dtype = np.dtype(dtype)
# Check if it's a known POD type
td = _DTYPE_TO_TD.get(dtype)
if td is not None:
return td
# Handle structured dtypes (structs)
if dtype.names is not None:
fields: dict[str, TypeDescriptor] = {}
assert dtype.fields is not None
for name in dtype.names:
field_info = dtype.fields[name]
field_dtype = field_info[0]
fields[name] = from_numpy_dtype(field_dtype)
return struct(fields) # type: ignore[arg-type]
# Some other NumPy type (e.g., complex64) for which we don't
# have a specific TypeDescriptor. Use STORAGE and preserve the original dtype.
td = TypeDescriptor(dtype.itemsize, dtype.alignment, TypeEnum.STORAGE)
td._dtype = dtype
return td
# Signed integer types
int8 = TypeDescriptor(1, 1, TypeEnum.INT8)
int16 = TypeDescriptor(2, 2, TypeEnum.INT16)
int32 = TypeDescriptor(4, 4, TypeEnum.INT32)
int64 = TypeDescriptor(8, 8, TypeEnum.INT64)
# Unsigned integer types
uint8 = TypeDescriptor(1, 1, TypeEnum.UINT8)
uint16 = TypeDescriptor(2, 2, TypeEnum.UINT16)
uint32 = TypeDescriptor(4, 4, TypeEnum.UINT32)
uint64 = TypeDescriptor(8, 8, TypeEnum.UINT64)
# Floating point types
float16 = TypeDescriptor(2, 2, TypeEnum.FLOAT16)
float32 = TypeDescriptor(4, 4, TypeEnum.FLOAT32)
float64 = TypeDescriptor(8, 8, TypeEnum.FLOAT64)
# Boolean
boolean = TypeDescriptor(1, 1, TypeEnum.BOOLEAN)
# Mapping from numpy dtype to TypeDescriptor for POD types
_DTYPE_TO_TD: dict[np.dtype, TypeDescriptor] = {
np.dtype("int8"): int8,
np.dtype("int16"): int16,
np.dtype("int32"): int32,
np.dtype("int64"): int64,
np.dtype("uint8"): uint8,
np.dtype("uint16"): uint16,
np.dtype("uint32"): uint32,
np.dtype("uint64"): uint64,
np.dtype("float16"): float16,
np.dtype("float32"): float32,
np.dtype("float64"): float64,
np.dtype("bool"): boolean,
}
def to_ctypes_type(td: TypeDescriptor):
"""Convert a TypeDescriptor to a ctypes type."""
return np.ctypeslib.as_ctypes_type(td.dtype)
def _annotation_to_type_descriptor(annotation):
"""
Convert a type annotation to a TypeDescriptor.
Handles:
- TypeDescriptor: returns as-is
- gpu_struct classes: returns their _type_descriptor
- numpy dtypes/types: converts via from_numpy_dtype
"""
from .struct import _is_struct_type
if isinstance(annotation, TypeDescriptor):
return annotation
if _is_struct_type(annotation):
return annotation._type_descriptor # type: ignore[union-attr]
# numpy dtype or type
return from_numpy_dtype(np.dtype(annotation))
def signature_from_annotations(py_func):
try:
annotations = get_type_hints(py_func)
except Exception:
annotations = py_func.__annotations__
spec = inspect.getfullargspec(py_func)
arg_names = list(spec.args)
input_tds = []
# Try to get input types from annotations
for name in arg_names:
if name in annotations:
input_tds.append(_annotation_to_type_descriptor(annotations[name]))
break
if "return" in annotations:
output_td = _annotation_to_type_descriptor(annotations["return"])
else:
output_td = None
return input_tds, output_td
__all__ = [
"int8",
"int16",
"int32",
"int64",
"uint8",
"uint16",
"uint32",
"uint64",
"float16",
"float32",
"float64",
"boolean",
"struct",
"pointer",
"from_numpy_dtype",
"to_ctypes_type",
]

64
cccl_upstream/python/cuda_cccl/cuda/compute/typing.py
View File

@@ -1,64 +0,0 @@
# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from typing import Callable, Protocol, TypeVar
from .iterators import IteratorBase
from .op import OpAdapter, OpKind, RawOp
from .struct import _Struct
class DeviceArrayLike(Protocol):
"""Protocol for array-like objects that expose device memory via CUDA Array Interface.
Any object implementing the ``__cuda_array_interface__`` attribute can be used
where a :class:`DeviceArrayLike` is expected. This includes CuPy arrays, Numba
device arrays, PyTorch CUDA tensors, and other GPU array types.
See `CUDA Array Interface specification <https://nvidia.github.io/numba-cuda/user/cuda_array_interface.html>`_
for details.
"""
__cuda_array_interface__: dict
class StreamLike(Protocol):
"""Protocol for CUDA stream objects.
Any object implementing the ``__cuda_stream__()`` method can be used where a
:class:`StreamLike` is expected. This includes stream objects from CuPy,
Numba CUDA, PyTorch, and other CUDA libraries.
"""
def __cuda_stream__(self) -> tuple[int, int]: ...
GpuStruct = TypeVar("GpuStruct", bound=_Struct)
"""
Instance of types created with :class:`cuda.compute.struct.gpu_struct`.
"""
IteratorT = TypeVar("IteratorT", bound=IteratorBase)
"""Type variable for iterator objects.
Represents any subclass of :class:`IteratorBase <cuda.compute.iterators.IteratorBase>`.
See :py:mod:`cuda.compute.iterators` for all available iterators.
"""
Operator = Callable | OpKind | RawOp | OpAdapter
"""Type alias for operator objects passed to algorithm functions.
Algorithms accept the following objects as operators:
* Python functions or lambdas implementing the operator. This function will be JIT
compiled into device code using `numba.cuda <https://nvidia.github.io/numba-cuda/>`_.
* :class:`OpKind <cuda.compute.op.OpKind>` enumerators which are pre-defined constants
for common operations.
* :class:`RawOp <cuda.compute.op.RawOp` objects containing pre-compiled device code.
"""
__all__ = ["DeviceArrayLike", "GpuStruct", "IteratorT", "Operator"]