feat(cccl): integrate missing CCCL directories — python/, ci/, .agent/, docs/, test/
Sparse-checkout from NVIDIA/cccl main branch to complete cccl_upstream: Added: - python/cuda_cccl/ (226 files) — Python bindings for device-level algorithms Critical for muh toolchain: cuda.compute.reduce_into, scan, radix_sort, etc. Includes 204 .py files with full test coverage for all 27 algorithms - ci/ (163 files) — Build/test infrastructure build_cub.sh, test_cub.sh, build_and_test_targets.sh, matrix.yaml Directly maps to our [INFRA-CI] and [INFRA-BUILD] items - .agent/skills/ (7 files) — NVIDIA's own agent skills for CCCL cccl-style/SKILL.md, cccl-test/SKILL.md, sass-diff/SKILL.md - docs/ (491 files) — Official CCCL documentation CI references, CMake guides, Python compute docs, libcudacxx PTX docs - test/ (12 files) — Top-level integration tests (cuda_smoke, stdpar) - Root configs: .clang-format, .clang-tidy, CONTRIBUTING.md, pyproject.toml - CLAUDE.md symlink → AGENTS.md (NVIDIA's standard) cccl_upstream now mirrors full NVIDIA/cccl structure: Before: 42M (cub + thrust + libcudacxx + cudax + c + examples + benchmarks) After: 53M (+python +ci +docs +.agent +test +configs) This completes the CCCL base needed for: - [muh-bench] items: ci/util/build_and_test_targets.sh for targeted builds - [CCCL-verify] items: python/cuda_cccl/tests/ as reference implementations - [CCCL-test] items: ci/test_cub.sh, ci/test_thrust.sh - Agent workflow: .agent/skills/ for consistent style and test patterns
This commit is contained in:
muh-bot
751
cccl_upstream/docs/python/compute/developer_overview.rst
Normal file
751
cccl_upstream/docs/python/compute/developer_overview.rst
Normal 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, <oir_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/``.
|
||||
Reference in New Issue
Block a user