Tracking: CUDA support in cppjit #176

Opening this issue to track my ongoing work on cppjit's GPU programming support with CUDA. An experimental API surface incorporating the above functionality can be expected by the end of the year (2026) except those that depend on clang-repl improvements available in LLVM24

With CUDA mode the cppjit interpreter into a host/device compiler pair. Allowing kernels definitions to compile to PTX for the GPU and launch directly with a Pythonic API. They share a session with host C++ and with modules compiled elsewhere, and interoperate with CuPy, torch and cuda.core arrays, streams and graphs. This builds on top of the current basic functionality of JITting device code with CppInterOp and allows cppjit users to build a GPU programming DSL by composing device kernels and operations.

An example of this work used in an application context:

#include <cuda_runtime.h>

constexpr int K_BINS_V = 128;

__device__ __forceinline__ int wrap1v(int i, int n) { return i + 1 == n ? 0 : i + 1; }

// The eight births of voxel (i, j, k) and their tokens into the block's shared histogram.
__device__ __forceinline__ void ecc_count_voxel(const int* __restrict__ q, int* sh, int i, int j, int k,
                                                int i1, int j1, int k1, int n, long long nn)
{
    const int v000 = q[i * nn + j * n + k],   v100 = q[i1 * nn + j * n + k];
    const int v010 = q[i * nn + j1 * n + k],  v110 = q[i1 * nn + j1 * n + k];
    const int v001 = q[i * nn + j * n + k1],  v101 = q[i1 * nn + j * n + k1];
// more operations..
    atomicAdd(&sh[fyz], 1);
    atomicSub(&sh[c], 1);
}

__global__ void ecc_hist_rows_kernel(const int* __restrict__ q, int* __restrict__ hist, int n)
{
    __shared__ int sh[K_BINS_V];
    for (int b = threadIdx.x; b < K_BINS_V; b += blockDim.x) sh[b] = 0;
    __syncthreads();
    const int j = blockIdx.y, i = blockIdx.z;
    const int i1 = wrap1v(i, n), j1 = wrap1v(j, n);
    const long long nn = static_cast<long long>(n) * n;
    for (int k = blockIdx.x * blockDim.x + threadIdx.x; k < n; k += gridDim.x * blockDim.x)
        ecc_count_voxel(q, sh, i, j, k, i1, j1, wrap1v(k, n), n, nn);
    __syncthreads();
    for (int b = threadIdx.x; b < K_BINS_V; b += blockDim.x)
        if (sh[b] != 0) atomicAdd(&hist[b], sh[b]);
}

and the Python surface processing a point cloud dataset:

def ecc_hist_cppjit(q):
    '''Signed histogram of births (int32, length K) on the device: one launch on the current stream.'''
    n = q.shape[0]
    hist = cp.zeros(K, dtype=cp.int32)
    stream = cp.cuda.get_current_stream()
    grid = (n ** 3 + BLOCK - 1) // BLOCK
    cppjit.gbl.ecc_hist_kernel[grid, BLOCK, 0, stream](q, hist, n)
    return hist


def run_cppjit(n):
    hist = ecc_hist_cppjit(cw.field(n))
    cp.cuda.get_current_stream().synchronize()
    return cp.asnumpy(hist).astype(np.int64)

The work is experimental and covers Linux x86_64 only. It spans three layers: clang-repl's incremental device pipeline in LLVM, the CppInterOp API, and the Python surface here. Checked cppjit items are implemented on a development branch and are not yet usptreamed.

LLVM (clang-repl)

• Flush the device bootstrap module before the first PTU (PR 226975)
• Fix the TargetMachine leak in device PTX generation (PR 228208)
• Fix PTX emission for inputs with no device code (PR 226977)
• Link builtin bitcode into every PTU, which enables device math (PR 229788)
• Change the optimization level at runtime with setOptLevel and the repl pragma (PR 214793)
• Print the device module IR under -emit-llvm (PR 235803)
• Run the CUDA interpreter tests on CI with a stub runtime (PR 236206)
• Pass the target features to the device TargetMachine (PR 235754)
• Register the device binary before static initializers, so one input can define and launch a kernel (PR 234643, in review)
• Take the setOptLevel defaults from the driver and cc1
• Apply the optimization level and the pragma to the device compilation
• Keep value printing out of the device compilation (Done, to be upstreamed)
• Find libcudart in lib64 toolkit layouts (Done, to be upstreamed)
• Embed cubins in the incremental fatbin through nvJitLink, removing the driver JIT at first launch (Done, needs improvement)
• Expose per-PTU device artifacts for incremental device linking (WIP, need to figure out how to test)

All of the pending items have atomic patches that are currently being upstreamed.

CppInterOp

• Fix the CUDA test fixture and re-enable the unit tests on the self-hosted GPU runner (PR 1161)
• Cpp::HasAttr to query the CUDA attributes of a declaration (PR 1163, in review)
• Propagate include paths to the device compiler instance (PR 1164, in review)
• Getters for the device architecture and the per-PTU fatbins (WIP)
• Runtime opt level control through the C API (WIP)
• Toolkit and runtime detection, if the libcudart lookup lands as a library API (Done)

cppjit:

• CUDA mode through CPPJIT_ENABLE_CUDA, CPPJIT_CUDA_PATH and CPPJIT_OFFLOAD_ARCH, with a clean boot failure on a wrong toolkit path
• Kernel launches with the subscript launch configuration, k[grid, block](...), including runtime configs and namespaced kernels
• User include paths reach the device parse
• Load and launch nvrtc modules through cppjit.cuda
• Accept __cuda_stream__ objects as the launch stream
• Record launches into CUDA graphs, including cuda.core's GraphBuilder
• Accept foreign buffers through the array interface and DLPack, as views and as kernel arguments (WIP)
• Validate foreign buffers, and bind read-only buffers only to const parameters (WIP)

All of the above items are being reworked for upstreaming and will have PR's opened soon once the CppInterOp surface has landed.

Blockers:

• Current self-hosted runner has a pre-volta GPU making it tricky to test certain functions from CUB
• Not a main blocker but a decision on whether we launch a single interpreter with DeviceCI or a host/device pair in case kernels are designed to be used with a boundary. The cost paid is that parsing large host-only headers like Boost would also go through the DeviceCI which may not be optimal.

Note: This issue has an assignee and the first phase of the work is almost complete modulo upstreaming and code review.