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

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

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

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

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

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

7
cccl_upstream/docs/python/api_reference.rst Normal file
View File

@@ -0,0 +1,7 @@
API Reference
=============
.. toctree::
:maxdepth: 1
compute_api

751
cccl_upstream/docs/python/compute/developer_overview.rst Normal file
View File

@@ -0,0 +1,751 @@
``cuda.compute`` Developer Overview
===================================
This document provides an overview of the internal structure of
``cuda.compute``. At a high level, ``cuda.compute`` exposes CUDA C++
parallel algorithms through a Python API. Internally, it combines
Python-side operator compilation, CUDA C++ source generation, and
runtime just-in-time (JIT) compilation and linking.
We start with a simplified prototype. As we encounter the limitations
of that simplified model, we introduce the additional mechanisms
needed by the full implementation, referring to the relevant source
code where useful.
We begin with a minimal example that invokes a CUDA C++ kernel from
Python. In this simplified prototype, the kernel takes a single
integer argument and prints it.
.. code-block:: c++
#include <cstdio>
__global__ void kernel(int value) {
std::printf("thread %d: %d\n", threadIdx.x, value);
}
extern "C" void launcher(int value) {
kernel<<<1, 4>>>(value);
cudaDeviceSynchronize();
}
We can compile this code using ``nvcc``:
.. code-block:: bash
nvcc -Xcompiler=-fPIC -x cu kernel.cu -shared -o libkernel.so
The resulting shared library exports the host function ``launcher``,
which we can call from Python using ``ctypes``:
.. code-block:: python
import ctypes
bindings = ctypes.CDLL('libkernel.so')
bindings.launcher.argtypes = [ctypes.c_int]
bindings.launcher(42)
Running that Python code produces:
.. code-block:: bash
thread 0: 42
thread 1: 42
thread 2: 42
thread 3: 42
The example above works because all of the behavior is fixed ahead of
time in the CUDA C++ source. The kernel and the operation it performs
are both known in advance.
A library primitive such as reduction is different. Its behavior
depends not only on the input type, but also on the operator being
applied. A practical Python API therefore cannot be limited to a
single built-in case such as summing ``float`` values. It needs to
support many data types and user-provided operators.
That means the CUDA C++ side must be able to invoke device code that
originates in Python. Reduction is a useful motivating example, but to
keep the mechanics simple we will start with a much smaller building
block: compiling a simple Python function and making it callable from
CUDA C++. The same technique later applies to user-provided reduction
operators.
We can compile such a Python function to PTX using
`Numba-CUDA <https://nvidia.github.io/numba-cuda/>`_ as follows:
.. code-block:: python
import numba.cuda
def op(value):
return 2 * value
ptx, _ = numba.cuda.compile(op, sig=numba.int32(numba.int32))
That'd give us the following PTX code:
.. code-block:: bash
.visible .func (.param .b32 func_retval0) op(.param .b32 op_param_0)
{
.reg .b32 %r<3>;
ld.param.u32 %r1, [op_param_0];
shl.b32 %r2, %r1, 1;
st.param.b32 [func_retval0+0], %r2;
ret;
}
At this point, the Python function has been compiled to device code,
but the CUDA C++ side still needs a way to refer to it.
Conceptually, we would like to treat the operator as an externally
defined device function and call it from CUDA C++:
.. code-block:: c++
#include <cstdio>
extern "C" __device__ int op(int a); // defined in Python
extern "C" __global__ void kernel(int value) {
std::printf("thread %d: %d\n", threadIdx.x, op(value));
}
extern "C" void launcher(int value) {
kernel<<<1, 4>>>(value);
cudaDeviceSynchronize();
}
This raises the next question: how do we combine device code produced
from Python with CUDA C++ code that calls it?
The difficulty is not just that the operator's implementation comes
from Python. The CUDA C++ side must also declare and call that
operator with the correct signature.
In the code above, the operator has the fixed signature ``int
op(int)``. A real API cannot assume that. The user might supply an
operator on ``float``, ``complex``, or some user-defined type, and the
generated CUDA C++ code has to match that interface exactly. In other
words, the declaration of ``op`` and the CUDA C++ source that calls it
depend on the user's types and operator signature.
That means the CUDA C++ side must be generated and compiled at
runtime. Using ``nvcc`` for that would make the API depend on an
external compiler toolchain being available on every user machine.
Instead, we use NVRTC, which is designed for runtime compilation of
CUDA C++.
Our Python code is now:
.. code-block:: python
import ctypes
import numba.cuda
def op(value):
return 2 * value
ptx, _ = numba.cuda.compile(op, sig=numba.int32(numba.int32))
bindings = ctypes.CDLL('./build/libkernel.so')
bindings.launcher.argtypes = [ctypes.c_int, ctypes.c_char_p, ctypes.c_int]
bindings.launcher(42, ptx.encode('utf-8'), len(ptx))
Correspondingly, the C++ launcher now accepts the operator PTX as an
additional argument. Inside the launcher, the CUDA C++ kernel is now
assembled as a source string and compiled with NVRTC:
.. code-block:: c++
extern "C" void launcher(int value,
const char* op_ptx, int op_ptx_size)
{
cudaSetDevice(0);
// Kernel is now a string!
std::string kernel_source = R"XXX(
extern "C" __device__ int op(int a);
extern "C" __global__ void kernel(int value) {
printf("thread %d prints value %d\n", threadIdx.x, op(value));
}
)XXX";
Once that source string has been assembled, we compile it to PTX with
NVRTC:
.. code-block:: c++
nvrtcProgram prog;
const char *name = "test_kernel";
nvrtcCreateProgram(&prog, kernel_source.c_str(), name, 0, nullptr, nullptr);
cudaDeviceProp deviceProp;
cudaGetDeviceProperties(&deviceProp, 0);
const int cc_major = deviceProp.major;
const int cc_minor = deviceProp.minor;
const std::string arch = std::string("-arch=sm_") + std::to_string(cc_major) + std::to_string(cc_minor);
const char* args[] = { arch.c_str(), "-rdc=true" };
const int num_args = sizeof(args) / sizeof(args[0]);
// Compile the CUDA C++ kernel to PTX
std::size_t ptx_size{};
nvrtcResult compile_result = nvrtcCompileProgram(prog, num_args, args);
nvrtcGetPTXSize(prog, &ptx_size);
std::unique_ptr<char[]> ptx{new char[ptx_size]};
nvrtcGetPTX(prog, ptx.get());
nvrtcDestroyProgram(&prog);
At this point, we have two PTX inputs: PTX for the generated CUDA C++
kernel and PTX for the Python-defined operator. We can combine them
using nvJitLink:
.. code-block:: c++
const char* link_options[] = { arch.c_str() };
// Link PTX comping from kernel and PTX coming from Python operator
nvJitLinkHandle handle;
nvJitLinkCreate(&handle, 1, link_options);
nvJitLinkAddData(handle, NVJITLINK_INPUT_PTX, ptx.get(), ptx_size, name);
nvJitLinkAddData(handle, NVJITLINK_INPUT_PTX, op_ptx, op_ptx_size, name);
nvJitLinkComplete(handle);
// Get resulting cubin
std::size_t cubin_size{};
nvJitLinkGetLinkedCubinSize(handle, &cubin_size);
std::unique_ptr<char[]> cubin{new char[cubin_size]};
nvJitLinkGetLinkedCubin(handle, cubin.get());
nvJitLinkDestroy(&handle);
The result of linking is a cubin containing the generated kernel and
the Python-defined operator. We can load that cubin as a CUDA
library, retrieve the kernel from it, and launch it:
.. code-block:: c++
// Load cubin
CUlibrary library;
cuLibraryLoadData(&library, cubin.get(), nullptr, nullptr, 0, nullptr, nullptr, 0);
// Get kernel pointer out of the library
CUkernel kernel;
cuLibraryGetKernel(&kernel, library, "kernel");
// Launch the kernel
void *kernel_args[] = { &value };
cuLaunchKernel((CUfunction)kernel, 1, 1, 1, 4, 1, 1, 0, 0, kernel_args, nullptr);
Now the output of the Python program would be:
.. code-block:: bash
thread 0 prints value 84
thread 1 prints value 84
thread 2 prints value 84
thread 3 prints value 84
This works, but it is still not optimal from a performance
perspective. If the operator were compiled as part of the same CUDA
C++ translation unit as the kernel, the compiler could inline it
directly. In the PTX-linked version above, however, the generated
cubin still contains a call to ``op`` instead of the operator body
itself.
To address this, we switch to a different intermediate representation.
Instead of PTX, we use `LTO-IR
<https://developer.nvidia.com/blog/cuda-12-0-compiler-support-for-runtime-lto-using-nvjitlink-library/>`_.
LTO-IR preserves enough information for link-time optimization, which
allows the operator to be inlined into the generated kernel.
On the Python side, switching from PTX to LTO-IR requires only a small
change:
.. code-block:: python
ltoir, _ = numba.cuda.compile(op, sig=numba.int32(numba.int32), output="ltoir")
On the C++ side, we make the same switch from PTX to LTO-IR:
.. code-block:: c++
const char* args[] = { arch.c_str(), "-rdc=true", "-dlto" };
const int num_args = sizeof(args) / sizeof(args[0]);
nvrtcResult compile_result = nvrtcCompileProgram(prog, num_args, args);
std::size_t ltoir_size{};
nvrtcGetLTOIRSize(prog, &ltoir_size);
std::unique_ptr<char[]> ltoir{new char[ltoir_size]};
nvrtcGetLTOIR(prog, ltoir.get());
nvrtcDestroyProgram(&prog);
const char* link_options[] = { "-lto", arch.c_str() };
nvJitLinkHandle handle;
nvJitLinkCreate(&handle, 2, link_options);
nvJitLinkAddData(handle, NVJITLINK_INPUT_LTOIR, ltoir.get(), ltoir_size, name);
nvJitLinkAddData(handle, NVJITLINK_INPUT_LTOIR, op_ltoir, op_ltoir_size, name);
If you inspect the generated cubin now, you will no longer see a call
to ``op``. Instead, the operator has been inlined into the kernel,
which improves performance. That is the key benefit of switching from
PTX to LTO-IR.
At this point, we have a working prototype that can pass
Python-defined operators into CUDA C++ kernels without sacrificing
performance. The next problem is user-defined data types. So far, the
examples have used built-in scalar types, but a practical API also
needs to support types whose layout is only known on the Python side.
Fortunately, the kernel source is already being assembled as a string at
runtime. That means we can also generate the type information needed by
the CUDA C++ side.
As a concrete example, suppose we want to pass a ``numba.complex128``
value into the kernel. The C++ side does not see the original Python
type definition, but that is not an issue. It only needs a storage
type with matching size and alignment, and can type-erase everything
else.
.. code-block:: c++
extern "C" void launcher(void *value_ptr, int type_size, int type_alignment,
const char* op_ltoir, int op_ltoir_size)
{
std::string storage_t = "struct __align__(" + std::to_string(type_alignment) + ")"
+ "storage_t { char data[" + std::to_string(type_size) + "]; };";
std::string kernel_source = storage_t + R"XXX(
extern "C" __device__ int op(char *state);
extern "C" __global__ void kernel(storage_t value) {
printf("thread %d prints value %d\n", threadIdx.x, op(value.data));
}
)XXX";
// ...
void *kernel_args[] = { value_ptr };
cuLaunchKernel((CUfunction)kernel, 1, 1, 1, 4, 1, 1, 0, 0, kernel_args, nullptr);
In this version, the operator takes a type-erased pointer. On the
Python side, we therefore pass a pointer to the ``numba.complex128``
value, together with the size and alignment needed to construct a
matching storage type on the C++ side:
.. code-block:: python
import ctypes
import numba
import numba.cuda
import numpy as np
def op(value):
return numba.int32(value[0].real + value[0].imag)
value_type = numba.complex128
context = numba.cuda.descriptor.cuda_target.target_context
size = context.get_value_type(value_type).get_abi_size(context.target_data)
alignment = context.get_value_type(value_type).get_abi_alignment(context.target_data)
ltoir, _ = numba.cuda.compile(op, sig=numba.int32(numba.types.CPointer(value_type)), output='ltoir')
value = np.array([1 + 2j], dtype=np.complex128)
type_erased_value_ptr = value.ctypes.data_as(ctypes.c_void_p)
bindings = ctypes.CDLL('./build/libkernel.so')
bindings.launcher.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_int, ctypes.c_char_p, ctypes.c_int]
bindings.launcher(type_erased_value_ptr, size, alignment, ltoir, len(ltoir))
In this example, we obtain the size and alignment of
``numba.complex128`` from Numba's type system. The remaining detail is
how to pass the value to ``cuLaunchKernel``. Kernel arguments are
described to ``cuLaunchKernel`` as pointers to host memory from which
the launch parameters are copied. In Python, that host-memory pointer
can be obtained in a few ways, for example with ``ctypes.byref`` or by
placing the value in a ``numpy.array`` and retrieving the array's
address with ``value.ctypes.data_as(ctypes.c_void_p)``.
One more ingredient is needed to get closer to the full
``cuda.compute`` implementation. The kernels in the CUDA C++ Core
Compute Libraries are templates, so our generated kernel must be a
template as well.
.. code-block:: c++
std::string kernel_source = storage_t + R"XXX(
extern "C" __device__ int op(char *state);
template <class T>
__global__ void kernel(T value) {
printf("thread %d prints value %d\n", threadIdx.x, op(value.data));
}
)XXX";
Defining the kernel as a template is still not enough. We also need to
instantiate that template for the generated storage type. NVRTC
provides the necessary API for that:
.. code-block:: c++
nvrtcProgram prog;
const char *name = "test_kernel";
nvrtcCreateProgram(&prog, kernel_source.c_str(), name, 0, nullptr, nullptr);
// Get the name of the instantiated kernel
std::string kernel_name = "kernel<storage_t>";
// Instantiate kernel template
nvrtcAddNameExpression(prog, kernel_name.c_str());
// ...
// Get lowered name of the kernel
const char* kernel_lowered_name; // _Z6kernelI9storage_tEvT_
nvrtcGetLoweredName(prog, kernel_name.c_str(), &kernel_lowered_name);
// ...
// Use it to get kernel pointer
cuLibraryGetKernel(&kernel, library, kernel_lowered_name);
With these pieces in place, we can connect the simplified prototype back
to ``cuda.compute``.
At a high level, the ``cuda.compute`` API follows the same overall
structure, but packages it into three stages. Using parallel reduction
as an example:
#. In the first stage, ``cuda.compute.make_reduce_into(...)`` constructs
a reusable reduction object:
``reducer = cuda.compute.make_reduce_into(d_in=d_in, d_out=d_out, op=op, h_init=h_init)``
Here ``op`` is a Python function that must be made available to the
CUDA kernel. As in the simplified prototype above, this stage
compiles ``op`` to LTO-IR, generates the corresponding CUDA C++
source, instantiates the necessary kernels, and compiles them with
NVRTC. The resulting build state is stored inside the returned
reduction object. At this stage, the concrete runtime values of the
provided arrays do not matter yet; later calls may use different
pointers or sizes, as long as the interface remains compatible.
#. In the second stage, that reduction object is used to query the
amount of temporary storage required by the algorithm:
``temp_storage_size = reducer(temp_storage=None, d_in=d_input, d_out=d_output, num_items=num_items, op=op, h_init=h_init)``
This returns the size of the temporary storage buffer, which must be
allocated in device-accessible memory. No kernels are launched at
this stage.
#. In the third stage, the algorithm is executed using the allocated
temporary storage:
``reducer(temp_storage=temp_storage, d_in=d_input, d_out=d_output, num_items=num_items, op=op, h_init=h_init)``
At this point, the kernels stored in the reduction object are
launched and the reduction is performed.
Build results and device state
------------------------------
An algorithm is built in one of two ways. A **default build**
(``compute_capability=None``, the common path) targets the current CUDA
device — it queries that device's compute capability, then compiles and loads
for it in one step. An **explicit ahead-of-time (AOT) build** (a
``compute_capability=`` argument naming one or more compute capabilities) names
its targets directly, compiles for each without loading, and needs no GPU, so it
can run on a build machine with no device (see
:ref:`cuda.compute.ahead_of_time_compilation`).
Either way, building produces a native build result — a Cython build-result
object wrapping the corresponding C runtime struct — that carries two kinds of
state with different device affinity:
* The **compiled payload** (the compiled device code and its launch policy)
depends only on the target compute capability. It is device-independent: the
same payload is valid on any device of that compute capability.
* **Loaded state** is created when the build result is loaded for execution —
the registered ``CUlibrary`` and the kernel handles resolved from it. It
belongs to the device (and context) it was loaded on. CUB launch paths
resolve a ``CUkernel`` to the current-context ``CUfunction`` and may get or
set kernel attributes on it, and CUDA kernel-attribute behavior is
device-specific.
Consequently, a *loaded* build result cannot be shared across two devices even
when they have the same compute capability: its handles are device-specific.
The compiled payload can be reused, but each device needs its own loaded result
— built directly, or reconstructed from the shared payload. This is a property
of CUDA and the build struct itself, independent of caching or free-threaded
Python. The next section describes how ``cuda.compute`` caches build results to
reuse the payload while giving each device its own loaded state.
Caching and free-threaded Python
--------------------------------
The user-facing cache behavior is described in :ref:`cuda.compute.caching`. This
section describes the implementation contracts that keep that behavior correct
for free-threaded Python and multi-GPU use.
Two cache layers
++++++++++++++++
Internally, ``cuda.compute`` separates two kinds of cached state:
* **Wrapper objects** are the Python objects returned by ``make_*`` APIs, such as
``make_reduce_into``. They own per-call descriptor state and are cached per
Python thread by ``cache_with_registered_key_functions`` in
``cuda/compute/_caching.py``. Keeping wrapper caches thread-local avoids
sharing mutable wrapper state across concurrent calls from free-threaded
Python.
* **Per-cc build results** (``_PerCCBuildResults``) hold one *canonical* Cython
build result per target compute capability — the single authoritative result
for that cc, carrying the device-independent compiled payload described in
`Build results and device state`_. They are cached by ``cache_build_results``
and may be shared by wrapper objects in different Python threads — and, for
default builds, across same-cc devices (except on the v2 HostJIT backend
today; see `Device keying`_). Each device's loaded result is tracked
separately within the entry, so sharing an entry never shares device-specific
state.
The normal cache-hit path is intentionally cheap. A wrapper-cache hit is
thread-local and does not consult the process-wide build-result cache. When a
wrapper is constructed, a completed build-result hit requires one process-wide
dictionary lookup and does not take an explicit cache lock. The two build
kinds then diverge because they differ in whether the target device is known
when the wrapper is constructed. A default build already knows its device — the
wrapper cache queried it to build and keys the wrapper to it — so the wrapper
resolves that device's loaded result once at construction and stores a direct
reference; executing it then needs no current-device query and no shared-cache
lookup. An explicit AOT or deserialized wrapper has no such binding — an AOT
build targets compute capabilities with no GPU queried, and a deserialized
wrapper is reconstructed without a device binding — so its device is known only
at call time. Each call resolves the per-device loaded result from a dictionary
inside the per-cc build results, where a completed lookup also takes no explicit
lock.
Design requirements
+++++++++++++++++++
The free-threading design is constrained by the following requirements:
* Importing ``cuda.compute`` in a free-threaded CPython interpreter must not
re-enable the GIL.
* Free-threading support should not add global locking or shared-state
contention to the normal single-threaded execution path. Wrapper cache hits
should be thread-local, and normal algorithm execution should not take a
global cache lock.
* Mutable wrapper state must not be shared across threads.
* Expensive native build results should still be shared across threads when they
are safe to share.
* Same-key concurrent cold builds should build once; waiters should receive the
same result or observe the same exception.
The current free-threading support boundary is the ``minimal-cu12`` and
``minimal-cu13`` extras. These extras omit Numba and Numba CUDA. Consequently,
free-threaded support currently covers built-in ``OpKind`` operations and
externally compiled ``RawOp`` operations, but not Python-callable operators.
The full ``cu12`` and ``cu13`` extras remain
outside the support claim until the Numba CUDA dependency is replaced by a
free-threading-compatible implementation.
CI runs ``test_free_threading_stress.py`` directly from the minimal test job.
The v1 backend is covered across the supported CUDA 12 and 13 lanes, and a
separate CTK 13.X minimal job runs the same suite against the v2 HostJIT
backend. Pytest runs each suite in one process while the stress tests create
and synchronize their own worker threads.
Build and validation requirements
+++++++++++++++++++++++++++++++++
The Cython extension that backs ``cuda.compute`` must opt in to free-threaded
execution:
.. code-block:: cython
# cython: freethreading_compatible=True
Without this marker, importing the extension in a free-threaded CPython process
can cause CPython to re-enable the GIL. The generated extension should advertise
``Py_MOD_GIL_NOT_USED`` and importing ``cuda.compute`` should leave
``sys._is_gil_enabled()`` false.
The free-threaded wheel must also keep its free-threaded ABI tag after repair and
merge steps. For CPython 3.14, the expected wheel tag contains
``cp314-cp314t`` rather than the regular ``cp314-cp314`` tag. The acceptance
criteria for a free-threaded build are:
* the wheel has the expected ``cp314-cp314t`` ABI tag;
* importing ``cuda.compute`` does not re-enable the GIL;
* the free-threading stress suite passes without forcing ``PYTHON_GIL=0`` or
``-X gil=0``.
Device keying
+++++++++++++
User-facing multi-GPU behavior and requirements are described in
:ref:`cuda.compute.multi_gpu`; this section covers the keying mechanism.
For the default build path, the wrapper cache includes
the current CUDA runtime device ordinal and compute capability in its key:
wrapper objects hold device-bound state, so each device (and thread) receives
its own wrapper. The shared build-result cache is keyed by compute capability
alone — the compiled payload depends only on the cc — so one shared entry
serves every same-cc device ordinal.
Per-device loaded state lives inside the shared entry. The device that built
the entry loads the canonical result in place; each additional same-cc device
loads its own clone of the compiled payload through serialization (serialize,
deserialize without loading, then load on the new device) instead of running
a full native compilation, and the clone re-validates the payload against the
current device. When the backend cannot serialize build results (the v2 HostJIT
backend today), sharing is not possible, so default builds are keyed per device
ordinal instead and each device builds its own entry.
Explicit AOT builds cannot include a device ordinal in their compilation key —
they build with no GPU queried — so their canonical results are shared
process-wide by specialization and target compute capabilities. Unlike a default
build, an AOT build compiles without loading, so no device owns the canonical
result until first execution: the first device to run claims and loads it, and
other same-cc devices load their own clone, exactly as above.
The first implementation intentionally keys shared build results by CUDA runtime
device ordinal rather than by CUDA context handle. User-managed CUDA driver
contexts are not a target use case for ``cuda.compute``. CUDA runtime,
``cuda.core``, CuPy, and PyTorch-style applications are expected to use the
primary-context model, and language frontends generally prefer that model.
Concurrent build coordination
+++++++++++++++++++++++++++++
When several threads miss the same cache key at once, only one should run the
expensive build and the rest should wait for its result. A shared helper
provides this coordination, a pattern called *single-flight*. The cache
dictionary stores either a completed value or a temporary ``_InFlightBuild``
entry. On a miss, each caller creates a candidate in-flight
entry, and ``dict.setdefault`` elects one caller to run the builder. Other
callers receive the winning entry and wait on its ``threading.Event``. If the
operation succeeds, the in-flight entry is replaced by the completed result and
all waiting threads receive that same object. If it fails, the exception is
propagated to the waiting threads and the failed entry is removed so that a
later call can retry. Completed-result hits do not allocate an in-flight entry
or take an explicit cache lock.
The same helper coordinates two kinds of misses: a compilation miss in the
process-wide build cache, where ``cache_build_results`` runs the native build
once per specialization, and a per-device load miss inside a per-cc build
result, where ``resolve`` loads (or clones and loads) the result once per
device.
When adding a new algorithm, the factory that returns the reusable wrapper object
should use ``cache_with_registered_key_functions``. The wrapper constructor
should pass the expensive native build operation to ``cache_build_results``,
which returns two values: the shared build results, and the loaded result bound
to the constructing device (``None`` for an explicit AOT build, which has no
constructing device). Store both; ``__call__`` passes them to
``resolve_build_result`` (see any algorithm class for the pattern).
Do not perform an expensive native build before entering
``cache_build_results``; otherwise same-key cold factory calls can duplicate the
build and bypass single-flight coordination.
The specialization key must include every argument that can affect generated
code, type layout, policy selection, or native build state. It should not include
runtime-only values such as array pointers, array contents, item counts, streams,
or temporary-storage pointers unless those values change the compiled interface.
User-object and descriptor contracts
++++++++++++++++++++++++++++++++++++
Wrapper objects returned by ``make_*`` APIs are not safe for concurrent calls
from multiple threads. If two threads need the same algorithm specialization,
each thread should call the
factory and receive its own wrapper object, or the caller must externally
serialize access to a shared wrapper. The wrapper updates its Cython
``Iterator``, ``Op``, ``Value``, and algorithm-specific descriptors before each
native call, so concurrent calls through the same wrapper could overwrite the
descriptor state another thread is about to use.
The same contract applies to wrappers reconstructed by ``deserialize()`` —
they are the same classes with the same mutable descriptors. Unlike the
factories, ``deserialize()`` does not hand each calling thread its own object
through the per-thread wrapper cache: every call constructs a fresh, uncached
wrapper. The natural deserialize-once-and-share pattern therefore reintroduces
exactly the descriptor races the per-thread factory cache prevents. Threads
that need a deserialized algorithm concurrently should each deserialize the
blob themselves; that performs no recompilation, at the cost of an independent
native load per object.
Read-only iterator and operator objects may be shared across threads. The
iterator base class uses a per-iterator lock for first-time lazy construction of
advance, input-dereference, and output-dereference ``Op`` objects; cached access
after that remains lock-free. This lock does not make arbitrary mutation safe:
concurrent mutation of iterator state, operator state, captured state, or child
iterators remains unsupported unless the caller synchronizes externally.
Mutable execution state belongs to one thread at a time unless the caller
provides synchronization. This includes output arrays, temporary-storage buffers,
streams, ``DoubleBuffer`` instances, and other objects whose state changes as
part of a launch.
Backend-specific notes
++++++++++++++++++++++
The v1 NVRTC/nvJitLink backend and the v2 HostJIT backend have different
free-threading risk surfaces and must be audited independently. v1 stresses
NVRTC, nvJitLink, CUDA library loading, and CUB host dispatch. v2 adds HostJIT
compiler state, LLVM/Clang initialization, persistent PCH paths, generated
source/cubin artifacts, and dynamic loader lifetime.
Transform has one additional v1 native-cache rule. Each transform build result
owns a native cache of launch configurations (``async_config`` /
``prefetch_config``) in ``c/parallel/src/transform.cu``. Because one build
result is shared by every thread using the same specialization, and the Cython
bindings release the GIL around the native call, each configuration is filled
exactly once through ``std::call_once``; later calls on any thread only pay the
``once_flag`` fast-path check. This holds on every interpreter build — regular
GIL builds also execute the native call concurrently once the GIL is released,
so the cache must be thread-safe unconditionally.
The v2 backend addresses the same transform concern differently, and only on
Windows. HostJIT compiles generated code with ``-fno-threadsafe-statics``
because the Windows CRT guard support that thread-safe function-local statics
require is unavailable. Generated CUB code still initializes function-local
statics lazily — transform's launch configuration among them — so a
per-build-result ``first_call_gate``
(``c/parallel.v2/src/util/first_call_gate.h``) serializes the first successful
call into each generated function; after it completes, an atomic fast-path check
lets later concurrent calls proceed without locking. Empty calls bypass the gate
because they return before CUB initializes the static. This covers transform and
binary search; other platforms keep thread-safe statics and need no gate.
Clearing caches
+++++++++++++++
``clear_all_caches()`` is process-local. It clears all known per-thread wrapper
caches through a weak registry of live thread cache containers, and it clears the
shared build-result cache. Separate Python processes build and cache
independently.
Calling ``clear_all_caches()`` concurrently with active factory calls or
algorithm execution is not supported unless the caller synchronizes externally.
Source map
----------
For readers who want to connect this overview back to the source tree:
* The Python-facing API, operator compilation, and the logic for
constructing and invoking reusable algorithm objects live under
``python/cuda_cccl/cuda/compute/``.
* The lower-level C/C++ runtime compilation and kernel-building
machinery lives under ``c/parallel/`` (and ``c/parallel.v2/`` for the v2
HostJIT backend).
* User-facing examples for ``cuda.compute`` live under
``python/cuda_cccl/tests/compute/examples/``.

555
cccl_upstream/docs/python/compute/index.rst Normal file
View File

@@ -0,0 +1,555 @@
.. _cccl-python-compute:
``cuda.compute``: Parallel Computing Primitives
===============================================
.. toctree::
:hidden:
:maxdepth: 2
Overview <self>
developer_overview
The ``cuda.compute`` library provides composable primitives for building custom
parallel algorithms on the GPU—without writing CUDA kernels directly.
Algorithms
----------
Algorithms are the core of ``cuda.compute``. They operate on arrays or
:ref:`iterators <cuda.compute.iterators>` and can be composed to build specialized
GPU operations—reductions, scans, sorts, transforms, and more.
Typical usage of an algorithm looks like this:
.. code-block:: python
cuda.compute.reduce_into(
d_in=..., # input array or iterator
d_out=..., # output array or iterator
op=..., # binary operator (built-in or user-defined)
num_items=..., # number of input elements
h_init=..., # initial value for the reduction
)
API conventions
+++++++++++++++
* **Keyword-only parameters** — All algorithm parameters are **keyword-only**.
They must always be passed by name, not by position:
.. code-block:: python
# correct
cuda.compute.reduce_into(d_in=d_input, d_out=d_output, num_items=n, op=OpKind.PLUS, h_init=h_init)
# incorrect — positional arguments are not accepted
cuda.compute.reduce_into(d_input, d_output, n, OpKind.PLUS, h_init) # TypeError
* **Naming** — The ``d_`` prefix denotes *device* memory (e.g., CuPy arrays, PyTorch tensors);
``h_`` denotes *host* memory (NumPy arrays). Some scalar values must be passed as
host arrays.
* **Output semantics** — Algorithms write results into a user-provided array or iterator
rather than returning them. This keeps memory ownership explicit and lifetimes under
your control.
* **Operators** — Many algorithms accept an ``op`` parameter. This can be a built-in
:class:`OpKind <cuda.compute.op.OpKind>` value or a
:ref:`user-defined operator <cuda.compute.user_defined_operators>`.
When possible, prefer built-in operators (e.g., ``OpKind.PLUS``) over the equivalent
user-defined operation (e.g., ``lambda a, b: a + b``) for better performance.
* **Iterators** — Inputs and outputs can be :ref:`iterators <cuda.compute.iterators>`
instead of arrays, enabling lazy evaluation and operation fusion.
Full Example
++++++++++++
The following example uses :func:`reduce_into <cuda.compute.algorithms.reduce_into>`
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
:caption: Sum reduction example.
Object-based API (expert mode)
++++++++++++++++++++++++++++++
Many algorithms allocate temporary device memory for intermediate results. For finer
control over allocation—or to reuse buffers across calls—use the object-based API.
For example, :func:`make_reduce_into <cuda.compute.algorithms.make_reduce_into>`
returns a reusable reduction object that lets you manage memory explicitly.
.. code-block:: python
:caption: Controlling temporary memory.
# create a reducer object:
reducer = cuda.compute.make_reduce_into(d_in=d_in, d_out=d_out, op=op, h_init=h_init)
# get the temporary storage size by passing None for the temp_storage argument:
temp_storage_bytes = reducer(temp_storage=None, d_in=d_in, d_out=d_out, num_items=num_items, op=op, h_init=h_init)
# allocate the temporary storage as any array-like object
# (e.g., CuPy array, Torch tensor):
temp_storage = cp.empty(temp_storage_bytes, dtype=np.uint8)
# perform the reduction:
reducer(temp_storage=temp_storage, d_in=d_in, d_out=d_out, num_items=num_items, op=op, h_init=h_init)
The object-based API splits the algorithm invocation into three phases,
1. Constructing an algorithm object
2. Determining the amount of temporary memory needed by the computation
3. Performing the computation
It is important that the type of arguments passed during construction (step 1) match
those passed during invocation (step 2 and 3). Otherwise you may see unexpected errors
or silent bugs.
- Data types of arrays/iterators must match. If you pass an array of `int32` data type
as the `d_in=` argument during construction of a reducer object, you must pass
an array of dtype `int32` when invoking it. The array can be of a different size.
- Bytecode instructions of functions must match. If you pass a function/lambda for
the operator during construction, you must pass a function with the same bytecode
instructions during invocation. This means you _can_ pass a different function
referencing different global/closures, but the operations within the functions
must be the same.
.. _cuda.compute.user_defined_operators:
User-Defined Operators
----------------------
A powerful feature is the ability to use algorithms with user-defined operators.
For example, to compute the sum of only the even values in a sequence,
we can use :func:`reduce_into <cuda.compute.algorithms.reduce_into>` with a custom binary operation:
.. literalinclude:: ../../../python/cuda_cccl/tests/compute/examples/reduction/sum_custom_reduction.py
:language: python
:start-after: # example-begin
:caption: Reduction with a custom binary operation.
Features and Restrictions
+++++++++++++++++++++++++
User-defined operations are just-in-time (JIT) compiled into device code using
`Numba CUDA <https://nvidia.github.io/numba-cuda/>`_, so they inherit many
of the same features and restrictions as Numba CUDA functions:
* `Python features <https://nvidia.github.io/numba-cuda/user/cudapysupported.html>`_
and `atomic operations <https://nvidia.github.io/numba-cuda/user/intrinsics.html>`_
supported by Numba CUDA are also supported within user-defined operators.
* Nested functions must be decorated with ``@numba.cuda.jit``.
* Variables captured in closures or globals follow
`Numba CUDA semantics <https://nvidia.github.io/numba-cuda/user/globals.html>`_:
scalars and host arrays are captured by value (as constants),
while device arrays are captured by reference.
.. _cuda.compute.iterators:
Iterators
---------
Iterators represent sequences whose elements are computed **on the fly**. They can
be used in place of arrays in most algorithms, enabling lazy evaluation, operation
fusion, and custom data access patterns.
A :func:`CountingIterator <cuda.compute.iterators.CountingIterator>`, for example,
represents an integer sequence starting from a given value:
.. code-block:: python
it = CountingIterator(np.int32(1)) # represents [1, 2, 3, 4, ...]
To compute the sum of the first 100 integers, we can pass a
:func:`CountingIterator <cuda.compute.iterators.CountingIterator>` directly to
:func:`reduce_into <cuda.compute.algorithms.reduce_into>`. No memory is allocated
to store the input sequence—the values are generated as needed.
.. literalinclude:: ../../../python/cuda_cccl/tests/compute/examples/iterator/counting_iterator_basic.py
:language: python
:start-after: # example-begin
:caption: Counting iterator example.
Iterators can also be used to *fuse* operations. In the example below, a
:func:`TransformIterator <cuda.compute.iterators.TransformIterator>` lazily applies
the square operation to each element of the input sequence. The resulting iterator
is then passed to :func:`reduce_into <cuda.compute.algorithms.reduce_into>` to compute
the sum of squares.
Because the square is evaluated on demand during the reduction, there is no need
to create or store an intermediate array of squared values. The transform and the
reduction are fused into a single pass over the data.
.. literalinclude:: ../../../python/cuda_cccl/tests/compute/examples/iterator/transform_iterator_basic.py
:language: python
:start-after: # example-begin
:caption: Transform iterator example.
Some iterators can also be used as the output of an algorithm. In the example below,
a :func:`TransformOutputIterator <cuda.compute.iterators.TransformOutputIterator>`
applies the square-root operation to the result of a reduction before writing
it into the underlying array.
.. literalinclude:: ../../../python/cuda_cccl/tests/compute/examples/iterator/transform_output_iterator.py
:language: python
:start-after: # example-begin
:caption: Transform output iterator example.
As another example, :func:`ZipIterator <cuda.compute.iterators.ZipIterator>` combines multiple
arrays or iterators into a single logical sequence. In the example below, we combine
a counting iterator and an array, creating an iterator that yields ``(index, value)``
pairs. This combined iterator is then used as the input to
:func:`reduce_into <cuda.compute.algorithms.reduce_into>` to compute the index of
the maximum value in the array.
.. literalinclude:: ../../../python/cuda_cccl/tests/compute/examples/iterator/zip_iterator_counting.py
:language: python
:start-after: # example-begin
:caption: Argmax using a zip iterator.
These examples illustrate a few of the patterns enabled by iterators. See the
:ref:`API reference <cuda_compute-module>` for the full set of available iterators.
.. _cuda.compute.custom_types:
Struct Types
------------
The :func:`gpu_struct <cuda.compute.struct.gpu_struct>` decorator defines
GPU-compatible struct types. These are useful when you have data laid out
as an "array of structures", similar to `NumPy structured arrays <https://numpy.org/doc/stable/user/basics.rec.html>`_.
.. literalinclude:: ../../../python/cuda_cccl/tests/compute/examples/struct/struct_reduction.py
:language: python
:start-after: # example-begin
:caption: Custom struct type in a reduction.
Array of Structures vs Structure of Arrays
++++++++++++++++++++++++++++++++++++++++++
When working with structured data, there are two common memory layouts:
* **Array of Structures (AoS)** — each element is a complete struct, stored
contiguously. For example, an array of ``Point`` structs where each point's
``x`` and ``y`` are adjacent in memory.
* **Structure of Arrays (SoA)** — each field is stored in its own array.
For example, separate ``x_coords`` and ``y_coords`` arrays.
``cuda.compute`` supports both layouts:
* **``gpu_struct``** — defines a true AoS type with named fields
* **``ZipIterator``** — combines separate arrays into tuples on the fly, letting
you work with SoA data as if it were AoS
.. _cuda.compute.caching:
Caching
-------
Algorithms in ``cuda.compute`` are compiled to GPU code at runtime. To
avoid recompiling on every call, build results are cached in memory.
When you invoke an algorithm with the same configuration—same dtypes,
iterator kinds, operator, compute capability, and current device—the
cached build is reused. On systems with multiple GPUs, GPUs with the same
compute capability may share one compiled build, while each GPU keeps its
own loaded state. Compiled build results may be reused by
multiple threads in the same process on any interpreter build;
free-threaded Python additionally runs such threads in parallel.
What determines the cache key
+++++++++++++++++++++++++++++
Each algorithm computes a cache key from:
* **Array dtypes** — the data types of input and output arrays
* **Iterator kinds** — for iterator inputs/outputs, a descriptor of the iterator type
* **Operator identity** — for user-defined functions, the function's bytecode,
constants, and closure contents (see below)
* **Compute capability** — the GPU architecture of the current device
* **Current device** — determines the algorithm object and its loaded state;
the compiled code itself is shared across devices with the same compute
capability (see Multi-GPU behavior below). Not used when an explicit
``compute_capability=`` is given, which keys on the requested compute
capabilities instead (see :ref:`cuda.compute.ahead_of_time_compilation`)
* **Algorithm-specific parameters** — such as initial value dtype or determinism mode
Note that array *contents* or *pointers* are not part of the cache key—only
the array's dtype. This means you can reuse a cached algorithm across different
arrays of the same type.
.. _cuda.compute.multi_gpu:
Multi-GPU behavior
++++++++++++++++++
Loaded builds are device-specific. With the default current-device build path,
``cuda.compute`` compiles once per compute capability and reuses the compiled
payload on other GPUs with the same compute capability; each GPU still receives
its own loaded native state. An explicit ahead-of-time build behaves the same
way across same-compute-capability GPUs. Set the intended current CUDA device
before invoking an algorithm, and pass arrays — and a stream — that belong to
that device; currently, the stream selects a queue on the current device and
does not select the device itself.
.. _cuda.compute.free_threading:
Free-threaded Python
++++++++++++++++++++
.. important::
Free-threaded Python support is currently validated with the
``minimal-cu12`` and ``minimal-cu13`` extras, which do not install Numba or
Numba CUDA:
.. code-block:: bash
pip install cuda-cccl[minimal-cu13] # or minimal-cu12
The full ``cu12`` and ``cu13`` extras and Python-callable operators that
require Numba CUDA are not currently
supported in free-threaded Python. Use built-in
:class:`OpKind <cuda.compute.op.OpKind>` operations or externally compiled
:class:`RawOp <cuda.compute.op.RawOp>` operations with the minimal
installation.
Independent calls from multiple Python threads reuse compiled build results
within the same process on any interpreter build. A free-threaded interpreter
additionally runs those calls in parallel instead of interleaving them under
the GIL.
The cache is local to the current Python process. Separate Python processes build
and cache independently, even if they use the same GPU and algorithm
configuration.
This does not make user-provided memory or CUDA work automatically safe to share.
Users are still responsible for avoiding data races, such as two threads writing
to the same output array at the same time. Read-only iterator and operator
objects may be shared across threads, but concurrent mutation of those objects,
captured state, or underlying arrays requires external synchronization. For
concurrent use, prefer the direct
algorithm APIs, such as
:func:`reduce_into <cuda.compute.algorithms.reduce_into>`, or create a separate
reusable algorithm object in each thread (for example, the object returned by
:func:`make_reduce_into <cuda.compute.algorithms.make_reduce_into>`). If multiple
threads share one of these objects, serialize access to that object.
The examples below additionally use CuPy for device arrays. CuPy is not part
of the ``minimal`` extras, so install it separately (``pip install
cupy-cuda13x`` or ``cupy-cuda12x``; free-threaded Linux wheels are available
starting with CuPy 14.1).
.. literalinclude:: ../../../python/cuda_cccl/tests/compute/examples/free_threading/direct_api.py
:language: python
:start-after: # example-begin
:caption: Concurrent reductions through the direct API from multiple threads.
.. literalinclude:: ../../../python/cuda_cccl/tests/compute/examples/free_threading/object_api.py
:language: python
:start-after: # example-begin
:caption: Per-thread algorithm objects created with a ``make_*`` factory.
How user-defined functions are cached
+++++++++++++++++++++++++++++++++++++
User-defined operators and predicates are hashed based on their bytecode, constants,
and closure contents. Two functions with identical bytecode and closures produce
the same cache key, even if defined at different source locations.
Closure contents are recursively hashed:
* **Scalars and host arrays** — hashed by value
* **Device arrays** — hashed by pointer, shape, and dtype (not contents)
* **Nested functions** — hashed by their own bytecode and closures
Because device arrays captured in closures are hashed by pointer, changing the
array's contents does not invalidate the cache—only reassigning the variable to
a different array does.
Memory considerations
+++++++++++++++++++++
The cache persists for the lifetime of the process and grows with the number of
unique algorithm configurations. In long-running applications or exploratory
notebooks, this can accumulate significant memory.
To clear all caches and free memory:
.. code-block:: python
import cuda.compute
cuda.compute.clear_all_caches()
This forces recompilation on the next algorithm invocation—useful for benchmarking
compilation time or reclaiming memory.
In multi-threaded programs, make sure no other thread is building or running an
algorithm while you call it: ``clear_all_caches()`` does not synchronize with
concurrent use, and a build that is already in progress may finish afterwards
and place its result back into the cache.
.. _cuda.compute.serialization:
Serialization
-------------
:ref:`Caching <cuda.compute.caching>` reuses build results *within* a single
process. Serialization goes one step further: it lets you persist a built
algorithm to a blob of bytes and reconstruct it later—in another process, or on
another machine—**without recompiling**.
Use :func:`serialize <cuda.compute.algorithms.serialize>` on any object returned by a
``make_*`` factory to obtain a ``bytes`` blob, and
:func:`deserialize <cuda.compute.algorithms.deserialize>` to reconstruct it. The blob stores
the *compiled* build result, so :func:`deserialize <cuda.compute.algorithms.deserialize>`
performs no JIT compilation—it neither invokes Numba nor recompiles device code.
In practice you would write the blob to a file and load it in a later run or on
another machine.
.. literalinclude:: ../../../python/cuda_cccl/tests/compute/examples/serialization/serialize_roundtrip.py
:language: python
:start-after: # example-begin
:caption: Serializing a built algorithm and reconstructing it without recompiling.
The same argument-matching rules described above for the object-based API apply
to a deserialized algorithm: the dtypes, iterator kinds, and operator you pass
when invoking it must match those used when it was originally built.
Blobs are versioned and self-describing, but they are not a long-term storage
format: compatibility across ``cuda-cccl`` versions is not guaranteed, and
loading a blob produced by a different version may be rejected with a clear
error. Persist the inputs needed to rebuild (or re-serialize) rather than
relying on old blobs surviving an upgrade.
The same threading rules also apply: like any other reusable algorithm object,
a deserialized algorithm must be used by one thread at a time unless access is
externally serialized (see :ref:`Free-threaded Python <cuda.compute.free_threading>`).
For concurrent use, call :func:`deserialize <cuda.compute.algorithms.deserialize>`
in each thread — reconstruction performs no recompilation, so per-thread
deserialization from one shared blob is cheap. (Currently each deserialized
object loads its native build state independently; a future release may share
that state behind the scenes.)
.. _cuda.compute.ahead_of_time_compilation:
Ahead-of-Time Compilation
-------------------------
By default, an algorithm is compiled for the compute capability of the *current*
device. To build for other GPUs—or to build on a machine with no GPU at all—pass
``compute_capability=`` to any ``make_*`` factory. Combined with
:ref:`serialization <cuda.compute.serialization>`, this lets you compile once on
a build machine and ship a ready-to-run artifact to your deployment targets.
Building for specific architectures
+++++++++++++++++++++++++++++++++++
Pass a single compute capability, or a list of them, to build one artifact that
runs on any of the listed architectures:
.. code-block:: python
reducer = cuda.compute.make_reduce_into(
d_in=d_in, d_out=d_out, op=OpKind.PLUS, h_init=h_init,
compute_capability=[80, 90], # build for sm_80 and sm_90
)
A compute capability may be given as an integer (``90``, ``75``), a
``(major, minor)`` pair (``(9, 0)``), or a string (``"9.0"``). When the argument
is omitted, the current device's architecture is used.
When a multi-architecture artifact is invoked, the build result matching the
running GPU is selected and loaded on the first call. Invoking it on a GPU whose
architecture was not built raises an error.
The compiled artifact may be reused by multiple devices with the same compute
capability. Loading remains device-specific — each device loads its own copy of
the compiled payload — but no device recompiles.
Building without a GPU
++++++++++++++++++++++
A ``make_*`` factory normally inspects its input arrays to determine dtypes,
which requires real device allocations. To build on a machine that has no
GPU—for example, in CI—pass :class:`ProxyArray <cuda.compute.ProxyArray>` and
:class:`ProxyValue <cuda.compute.ProxyValue>` placeholders in place of real
arrays and scalars. A proxy describes *only* the dtype (and, for arrays, shape
and contiguity); it holds no GPU memory.
.. literalinclude:: ../../../python/cuda_cccl/tests/compute/examples/serialization/ahead_of_time_compilation.py
:language: python
:start-after: # example-begin
:caption: Compiling for two architectures with no GPU present.
When building without a GPU you must pass ``compute_capability=`` explicitly:
with no device to query, there is no architecture to default to. A proxy is a
build-time placeholder only—supply the real device arrays and scalars when you
invoke the (possibly deserialized) algorithm. Passing a proxy to a compiled
algorithm's ``__call__`` raises ``RuntimeError``.
.. _cuda.compute.externally_compiled_operators:
Externally Compiled Operators
-----------------------------
:class:`RawOp <cuda.compute.op.RawOp>` can be used to directly pass compiled device code
(LTO-IR) implementing custom operators.
This is useful for users who wish to use a different compilation pipeline than the default
used by ``cuda.compute`` (JIT compilation of Python callables using Numba CUDA).
The example below shows how to compile a C++ device function
to LTO-IR using `cuda.core <https://nvidia.github.io/cuda-python/cuda-core/latest/>`_,
:func:`reduce_into <cuda.compute.algorithms.reduce_into>`:
.. literalinclude:: ../../../python/cuda_cccl/tests/compute/examples/raw_op/cpp_stateless.py
:language: python
:start-after: # example-begin
.. important::
**Required calling convention**: Compiled functions must use untyped pointers for
all parameters, with manual type casting inside the function. In C++, this means
all arguments (and the return value) must be passed as ``void*`` pointers:
.. code-block:: cpp
extern "C" __device__ void my_binary_op(void* a, void* b, void* result) {
*static_cast<int*>(result) = *static_cast<int*>(a) + *static_cast<int*>(b);
}
You must ensure that:
* All parameters are untyped pointers with manual casting in the function body
* Type casts match the actual data types passed at runtime
* For stateful operators, state is the first parameter (also an untyped pointer)
* State bytes have the correct layout and alignment
Type mismatches can cause crashes, memory corruption, or silent incorrect results.
If you wish to use ``cuda.compute`` solely with externally compiled operators
(i.e., without native JIT support), you can install a
minimal version of the `cuda-cccl` package that ships without Numba/Numba CUDA dependencies:
.. code-block:: bash
pip install cuda-cccl[minimal-cu13] # or minimal-cu12 (pip-installed cuda-toolkit)
pip install cuda-cccl[minimal-sysctk13] # or minimal-sysctk12 (system CUDA toolkit)
Examples
--------
For complete runnable examples and additional usage patterns, see the
`examples directory <https://github.com/NVIDIA/CCCL/tree/main/python/cuda_cccl/tests/compute/examples>`_.
API Reference
-------------
- :ref:`cuda_compute-module`

105
cccl_upstream/docs/python/compute_api.rst Normal file
View File

@@ -0,0 +1,105 @@
.. _cuda_compute-module:
``cuda.compute`` API Reference
==============================
.. warning::
``cuda.compute`` is in public beta.
The API is subject to change without notice.
Algorithms
----------
.. automodule:: cuda.compute.algorithms
:members:
:undoc-members:
:imported-members:
Iterators
---------
.. automodule:: cuda.compute.iterators
:members:
:undoc-members:
:imported-members:
Operators
---------
.. py:currentmodule:: cuda.compute.op
.. Unfortunately, we need to manually document the OpKind enum here because
.. the `._bindings` module, where OpKind is defined, is mocked out when building
.. docs. The mock out is needed to avoid the need for CUDA to be installed
.. at docs build time.
.. py:class:: OpKind
Enumeration of operator kinds for CUDA parallel algorithms.
This enum defines the types of operations that can be performed
in parallel algorithms, including arithmetic, logical, and bitwise operations.
.. py:attribute:: PLUS
.. py:attribute:: MINUS
.. py:attribute:: MULTIPLIES
.. py:attribute:: DIVIDES
.. py:attribute:: MODULUS
.. py:attribute:: EQUAL_TO
.. py:attribute:: NOT_EQUAL_TO
.. py:attribute:: GREATER
.. py:attribute:: LESS
.. py:attribute:: GREATER_EQUAL
.. py:attribute:: LESS_EQUAL
.. py:attribute:: LOGICAL_AND
.. py:attribute:: LOGICAL_OR
.. py:attribute:: LOGICAL_NOT
.. py:attribute:: BIT_AND
.. py:attribute:: BIT_OR
.. py:attribute:: BIT_XOR
.. py:attribute:: BIT_NOT
.. py:attribute:: IDENTITY
.. py:attribute:: NEGATE
.. py:attribute:: MINIMUM
.. py:attribute:: MAXIMUM
.. autoclass:: cuda.compute.op.RawOp
:members:
:undoc-members:
Ahead-of-Time Compilation
-------------------------
The :func:`serialize <cuda.compute.algorithms.serialize>` and
:func:`deserialize <cuda.compute.algorithms.deserialize>` functions (listed under
`Algorithms`_ above) persist and restore built algorithms. To build ahead of time
for architectures other than the current device's—or with no GPU present—pass the
following dtype-only placeholders to a ``make_*`` factory in place of real arrays
and scalars:
.. py:class:: cuda.compute.ProxyArray(dtype)
A dtype-only placeholder for a device array. Use in place of a real device
array when calling a ``make_*`` factory to compile an algorithm without
allocating GPU memory. Satisfies the ``DeviceArrayLike`` protocol; accessing
its data pointer raises ``RuntimeError``. See
:ref:`cuda.compute.ahead_of_time_compilation`.
.. py:class:: cuda.compute.ProxyValue(dtype)
A dtype-only placeholder for a scalar or initial-value argument (such as
``h_init``). Use in place of a real numpy scalar or array to compile an
algorithm without real data; accessing its data raises ``RuntimeError``. See
:ref:`cuda.compute.ahead_of_time_compilation`.
Utilities
---------
.. automodule:: cuda.compute.struct
:members:
Typing
------
.. automodule:: cuda.compute.typing
:members:

34
cccl_upstream/docs/python/index.rst Normal file
View File

@@ -0,0 +1,34 @@
CCCL Python Libraries
======================
Overview
--------
The CUDA Core Compute Libraries (CCCL) for Python are a collection of modules
with the shared goal of providing **high-quality, high-performance, and easy-to-use**
abstractions for CUDA Python developers.
* :doc:`cuda.compute <compute/index>` — Composable device-level primitives for building
custom parallel algorithms, without writing CUDA kernels directly.
These libraries expose the generic, highly-optimized algorithms from the
`CCCL C++ libraries <https://nvidia.github.io/cccl/cpp.html>`_,
which have been tuned to provide optimal performance across GPU architectures.
Who is this for?
----------------
- **Library authors** building parallel algorithms that need portable performance
across GPU architectures—without dropping to CUDA C++.
- **Application developers** using PyTorch, CuPy, or other GPU-accelerated frameworks
who need custom algorithms beyond what those libraries provide.
.. toctree::
:maxdepth: 2
:caption: CCCL Python Libraries
setup
compute/index
resources
api_reference

38
cccl_upstream/docs/python/resources.rst Normal file
View File

@@ -0,0 +1,38 @@
Resources
=========
Examples
--------
For recipes and patterns, see our examples:
* ``cuda.compute`` `examples <https://github.com/NVIDIA/cccl/tree/main/python/cuda_cccl/tests/compute/examples>`_
CUB and Thrust Documentation
----------------------------
The CCCL Python libraries are built on top of the CUB and Thrust libraries.
See the `CUB documentation <https://nvlabs.github.io/cub/>`_ and `Thrust documentation <https://thrust.github.io/>`_
for more information regarding the underlying libraries.
Asking for Help
---------------
If you have a question, run into an issue, or have a feature request,
please raise an issue or start a discussion on our `GitHub repository <https://github.com/NVIDIA/cccl/issues>`_.
Contributing
------------
We welcome contributions! Please see the
`contributing guide <https://github.com/NVIDIA/cccl/blob/main/CONTRIBUTING.md>`_
for instructions on how to set up a development environment and submit a pull request.
Once you have a development environment set up, see :doc:`setup` for instructions
on how to install `cuda.cccl` in development mode.
License
-------
The CCCL Python libraries are licensed under the `Apache License 2.0 <https://www.apache.org/licenses/LICENSE-2.0>`_.

106
cccl_upstream/docs/python/setup.rst Normal file
View File

@@ -0,0 +1,106 @@
.. _cccl-python-setup:
Setup and Installation
======================
This guide walks you through installing and setting up the CUDA Python Core Libraries (CCCL).
Prerequisites
-------------
Before installing cuda-cccl, ensure you have:
* **Python 3.10 or later**
* **CUDA Toolkit 12.x or 13.x**
* **Compatible NVIDIA GPU** with Compute Capability 7.5 or higher
* **Operating Systems:** Linux (tested on Ubuntu 20.04+) or Windows 10/11 (with WSL2 support)
Installation
------------
Install from PyPI
~~~~~~~~~~~~~~~~~
The easiest way to install ``cuda-cccl`` is using pip:
.. code-block:: bash
pip install cuda-cccl[cu13] # or cuda-cccl[cu12]
This will install ``cuda-cccl`` along with all required dependencies, including
the ``cuda-toolkit`` pip packages for the chosen CUDA major version.
If you already have a CUDA toolkit installed on your system (e.g., via the
NVIDIA runfile, package manager, or Conda) and do not want pip to install it,
use the ``sysctk`` variants instead:
.. code-block:: bash
pip install cuda-cccl[sysctk13] # or cuda-cccl[sysctk12]
These install the same dependencies except ``cuda-toolkit``; it is your
responsibility to ensure a compatible CUDA toolkit is on ``PATH`` and
``LD_LIBRARY_PATH``.
For a minimal install without Numba (useful when you supply your own
:ref:`pre-compiled operators <cuda.compute.externally_compiled_operators>`), use:
.. code-block:: bash
pip install cuda-cccl[minimal-cu13] # pip-installed CUDA toolkit
pip install cuda-cccl[minimal-sysctk13] # system CUDA toolkit
Free-threaded Python support is currently validated with the ``minimal-cu12``
and ``minimal-cu13`` extras. The full ``cu12`` and ``cu13`` extras depend on
Numba CUDA and are not currently supported in free-threaded Python.
Install from conda-forge
~~~~~~~~~~~~~~~~~~~~~~~~~
Alternatively, you can install ``cuda-cccl`` using conda:
.. code-block:: bash
conda install -c conda-forge cccl-python
This will install the CCCL Python libraries and their dependencies from the conda-forge channel.
Install from Source
~~~~~~~~~~~~~~~~~~~
For development or to access the latest features:
.. code-block:: bash
git clone https://github.com/NVIDIA/cccl.git
cd cccl/python/cuda_cccl
pip install -e .[test-cu13] # or .[test-cu12], .[test-sysctk13], .[test-sysctk12]
The test extras do not install CuPy. To also run the CuPy-based
``cuda.compute`` examples, install CuPy separately, for example
``pip install cupy-cuda13x``.
Development Setup
~~~~~~~~~~~~~~~~~~
For contributing to cuda-cccl or advanced development:
.. code-block:: bash
# Clone the repository
git clone https://github.com/NVIDIA/cccl.git
cd cccl/python/cuda_cccl
# Install in development mode with test dependencies
pip install -e .[test-cu13] # or .[test-cu12], .[test-sysctk13], .[test-sysctk12]
# Run tests to verify everything works
pytest tests/
Next Steps
----------
Now that you have ``cuda-cccl`` installed, check out:
* :doc:`compute/index` - Parallel computing primitives for operations on arrays or data ranges