Ahead-of-Time (AOT) Compilation#

This guide demonstrates how to use CuTe DSL’s Ahead-of-Time (AOT) compilation features to export compiled kernels for use in production environments.

Overview#

CuTe DSL Ahead-of-Time (hereinafter referred to as AOT) compilation allows you to:

  • Compile once, enable cross-compilation: Write kernels in Python and cross-compile them for multiple GPU architectures, and cross-compile the host object for other CPU architectures (see Host Cross-Compilation for AArch64).

  • Remove JIT overhead: Eliminate compilation delays in production by pre-compiling kernels.

  • Flexible integration: Easily integrate compiled kernels into both Python and C/C++ codebases using flexible deployment options.

We provide 2 levels of AOT ABI:

  1. Low-Level CuTe ABI: This ABI is expressed using CuTe DSL types and tensors, mirroring the original Python function.

  2. High-Level Apache TVM FFI ABI: For interop with various frameworks (e.g., PyTorch, JAX), and offer high-level stable ABI access.

This guide will focus on the CuTe ABI AOT. For the Apache TVM FFI AOT, please refer to the section “Exporting Compiled Module” in TVM FFI Compilation.

CuTe ABI AOT Workflow#

Export Interface#

The export_to_c interface is provided by the JitCompiledFunction class. It accepts the following parameters:

  • file_path: The path to the directory where the header and object files will be saved.

  • file_name: The base name for the header and object files. The same file name will always overwrite existing files.

  • function_prefix: The prefix of the function symbol in the generated object file. This should be a unique identifier to avoid symbol conflicts. Users should ensure the function prefix is unique for each exported function. Defaults to the file_name.

It generates the following files:

  • {file_path}/{file_name}.h: A C header file containing API function declarations. This header specifies the runtime function signatures in C, mirroring the original Python function interfaces.

  • {file_path}/{file_name}.o: A standard object file containing the compiled kernel code. You can link this object file into either a static or shared library. It includes the host entry function, fatbin data, and helper functions such as cuda_init and cuda_load_to_device. Additionally, it embeds metadata for runtime loading and version verification.

Example:

import cutlass.cute as cute
import cutlass.cute.cuda as cuda

@cute.kernel
def print_tensor_kernel(a: cute.Tensor):
    cute.printf("a: {}", a)

@cute.jit
def print_tensor(a: cute.Tensor, stream: cuda.CUstream):
    print_tensor_kernel(a).launch(grid=(1, 1, 1), block=(1, 1, 1), stream=stream)

compiled_func = cute.compile(print_tensor)
# Export compiled functions to object files and headers
compiled_func.export_to_c(file_path="./artifacts", file_name="print_tensor_example", function_prefix="print_tensor")

Loading in Python#

Load pre-compiled object files or shared libraries into Python for execution.

import cutlass.cute as cute
import torch
from cutlass.cute import from_dlpack
import cutlass.cute.cuda as cuda

# Load module from object file
module = cute.runtime.load_module("./artifacts/print_tensor_example.o")
# or
module = cute.runtime.load_module("./artifacts/libprint_tensor_example.so")

# Prepare data
a = torch.arange(160, dtype=torch.float32, device="cuda").reshape(16, 10)
a_cute = from_dlpack(a).mark_layout_dynamic()
stream = cuda.CUstream(0)

# Call the function (no JIT compilation needed!)
module.print_tensor(a_cute, stream=stream)

# This will fail because 'non_existing_api' was not exported:
# module.non_existing_api()

C++ Integration with Static Linking#

Integrate compiled kernels directly into your C++ executable during the build process. The generated header file supplies the necessary API for loading the module and invoking the function.

Example:

#include "print_tensor_example.h"
#include <cuda_runtime.h>

void run_print_tensor() {
    // Prepare tensor, the tensor declaration is in the header file
    print_tensor_Tensor_a_t tensor_a;
    tensor_a.data = nullptr; // GPU memory is set to nullptr.
    // Set dynamic shapes and strides
    tensor_a.dynamic_shapes[0] = 32;
    tensor_a.dynamic_shapes[1] = 16;
    tensor_a.dynamic_strides[0] = 16;

    // Create stream
    cudaStream_t stream;
    cudaStreamCreate(&stream);

    // Load module before calling the kernel
    print_tensor_Kernel_Module_t module;
    print_tensor_Kernel_Module_Load(&module);

    // Call the kernel; the kernel wrapper function is defined in the header file
    cute_dsl_print_tensor_wrapper(&module, &tensor_a, stream);

    // Cleanup
    print_tensor_Kernel_Module_Unload(&module);
    cudaStreamDestroy(stream);
}

The print_tensor_example.h header file is generated by the export_to_c interface. It includes:

  • The print_tensor_Kernel_Module_t type: Represents the kernel module.

  • The print_tensor_Tensor_a_t type: A tensor-specific type that defines the ABI for a particular CuTe tensor.

  • The cute_dsl_print_tensor_wrapper function: The user-facing entry point to invoke the kernel.

The compilation of the C++ executable requires the libcuda_dialect_runtime.so or libcuda_dialect_runtime_static.a library which is involved in <wheel_install_path>/lib, along with the CUDA driver and runtime libraries, to function properly.

C++ Integration with Dynamic Loading#

Dynamically load pre-compiled object files or shared libraries at runtime. By including the CuteDSLRuntime.h header, you can load the module, look up exported functions, and invoke them.

#include "CuteDSLRuntime.h"
#include <cuda_runtime.h>
#include <cstdio>

void run_print_tensor() {
    // Load module from shared library
    CuteDSLRT_Module_t *module = nullptr;
    CuteDSLRT_Error_t err = CuteDSLRT_Module_Load(
        &module,
        "./artifacts/libprint_tensor_example.so"
    );
    // or
    CuteDSLRT_Error_t err = CuteDSLRT_Module_Load(
        &module,
        "./artifacts/print_tensor_example.o"
    );
    check_error(err);

    // Lookup function
    CuteDSLRT_Function_t *func = nullptr;
    err = CuteDSLRT_Module_Get_Function(&func, module, "print_tensor");
    check_error(err);

    // Prepare arguments, matching the argument type defined in the header file
    typedef struct {
        void *data;
        int32_t dynamic_shapes[2];
        int64_t dynamic_strides[1];
    } print_tensor_Tensor_a_t;

    print_tensor_Tensor_a_t tensor_a;
    tensor_a.data = nullptr;
    tensor_a.dynamic_shapes[0] = 32;
    tensor_a.dynamic_shapes[1] = 16;
    tensor_a.dynamic_strides[0] = 16;

    // Create stream
    cudaStream_t stream;
    cudaStreamCreate(&stream);

    // Call the function; the runtime function accepts packed arguments, refer to the wrapper in the header file
    // The trailing packed argument receives the CUDA error code of the kernel launch
    int32_t ret = 0;
    void* args[] = {&tensor_a, &stream, &ret};
    err = CuteDSLRT_Function_Run(func, args, 3);
    if (ret != cudaSuccess) {
        fprintf(stderr, "kernel launch failed: %s\n",
                cudaGetErrorName(static_cast<cudaError_t>(ret)));
    }
    check_error(err);
    cudaStreamSynchronize(stream);

    // Cleanup
    CuteDSLRT_Module_Destroy(module);
    cudaStreamDestroy(stream);
}

The CuteDSLRuntime.h header file can be found in <wheel_install_path>/include. It includes:

  • The CuteDSLRT_Error_t type: Indicates the status of the runtime API call itself, not the CUDA error code of the kernel launch. See Return Values and Error Handling.

  • The CuteDSLRT_Module_Load function: Loads the module.

  • The CuteDSLRT_Module_Get_Function function: Gets a function from the loaded module. The runtime API will load the CUDA module for kernel execution.

  • The CuteDSLRT_Function_Run function: Runs the function.

  • The CuteDSLRT_Module_Destroy function: Destroys the module.

The compilation of the C++ executable requires the libcute_dsl_runtime.so library which is involved in <wheel_install_path>/lib, along with the CUDA driver and runtime libraries, to function properly.

Return Values and Error Handling#

The wrapper function in the generated header returns an int32_t:

static inline int32_t cute_dsl_print_tensor_wrapper(
    print_tensor_Kernel_Module_t *module,
    print_tensor_Tensor_a_t *a,
    cudaStream_t stream);

The returned value is a CUDA runtime cudaError_t code:

  • 0 (cudaSuccess) means every kernel launch in the exported function was submitted successfully.

  • Any other value is the code returned by cudaLaunchKernelExC for the first launch that failed. The exported function returns at that point, so kernels launched later in the same @cute.jit function do not run.

Because the value is an ordinary cudaError_t, the CUDA runtime helpers cudaGetErrorName and cudaGetErrorString translate it into a human-readable message. The generated header also defines a CUTE_DSL_CUDA_ERROR_CHECK macro, defined in terms of those two helpers, that reports the code this way:

#include "print_tensor_example.h"

int32_t ret = cute_dsl_print_tensor_wrapper(&module, &tensor_a, stream);
CUTE_DSL_CUDA_ERROR_CHECK(ret);

// Or inspect the code directly
if (ret != cudaSuccess) {
    cudaError_t err = static_cast<cudaError_t>(ret);
    fprintf(stderr, "kernel launch failed: %s: %s\n",
            cudaGetErrorName(err), cudaGetErrorString(err));
}

Note that:

  • The return value only covers launch submission. Kernel execution is asynchronous, so faults such as cudaErrorIllegalAddress do not appear in it. Check cudaStreamSynchronize or cudaDeviceSynchronize separately for those.

  • Dynamic loading reports two independent statuses. CuteDSLRT_Error_t describes the runtime API call itself, such as module loading, symbol lookup and invocation, and is decoded with CuteDSLRT_GetErrorName and CuteDSLRT_GetErrorString. It reports every kernel failure as the single value CuteDSLRT_Error_CudaError and does not carry the underlying code; that code is written to the trailing packed argument instead, as shown in the dynamic loading example above.

  • Loading in Python raises instead of returning. A non-zero code is raised as a DSLCudaRuntimeError carrying the cudaError_t name.

  • The Apache TVM FFI ABI uses a different contract. Its exported functions return 0 on success and -1 on failure, and the message is retrieved through the TVM FFI error object rather than from the return value. See TVM FFI Compilation.

Host Cross-Compilation for AArch64#

By default, the object file produced by export_to_c targets the same machine that runs cute.compile (the build host, typically x86_64). CuTe DSL can instead emit the host portion of the object for a different CPU architecture, so you can build a kernel on an x86_64 machine and deploy it to an AArch64 Linux machine. Only the host code is affected; the GPU device code is still selected by --gpu-arch as described above.

Build Host vs Target#

  • Build host: the machine where you run Python, cute.compile, and export_to_c.

  • Target: the machine where the resulting shared library or executable will run.

When the host target is left empty (the default), the .o is compiled for the build host. When a target is requested, the .o is an ELF object for that target’s triple, and the accompanying .h uses only fixed-width and opaque types, so the C ABI is identical regardless of any word-size difference between the build host and the target.

Selecting the Host Target#

The host target is chosen at cute.compile time through the --host-target option. Two input formats are accepted:

  • A preset tag. linux-aarch64 maps to the triple aarch64-unknown-linux-gnu.

  • A long form for explicit CPU and feature tuning: llvm -mtriple=<triple> [-mcpu=<cpu>] [-mattr=<features>].

import cutlass.cute as cute

# Preset tag
compiled = cute.compile(
    my_function, *args,
    options="--gpu-arch sm_100a --host-target linux-aarch64")

# Long form with explicit CPU and features
compiled = cute.compile(
    my_function, *args,
    options="--gpu-arch sm_100a "
            "--host-target 'llvm -mtriple=aarch64-unknown-linux-gnu "
            "-mcpu=neoverse-n1 -mattr=+lse'")

Invalid input, such as an unknown preset tag or a malformed long form, is rejected immediately when cute.compile is called, rather than later during export.

Exporting the Cross-Compiled Object#

Once the function is compiled with a host target, export it exactly as in the native workflow; no additional arguments are required.

compiled.export_to_c(file_path="./artifacts", file_name="kernel",
                     function_prefix="kernel")

The resulting kernel.o is an AArch64 ELF object, and kernel.h is portable across the build host and the target.

Cross-Linking the Shared Library#

Link the object with your AArch64 cross toolchain (for example aarch64-linux-gnu-g++) and a sysroot that provides the CUDA headers and libraries for the target:

aarch64-linux-gnu-g++ -shared -o kernel.so kernel.o \
    $(python -m cutlass.cute.export.aot_config --ldflags --target aarch64-unknown-linux-gnu) \
    $(python -m cutlass.cute.export.aot_config --libs    --target aarch64-unknown-linux-gnu)

Linking succeeds against the stub on the build host. Deploy kernel.so to the AArch64 target, where it binds against the real runtime library.

Limitations#

  • AArch64 only. Only aarch64-unknown-linux-gnu and aarch64-unknown-nto-qnx8.0.0 (see below) are currently supported. Other architectures fail with a “target not registered” error during code generation.

  • Not compatible with TVM FFI. Combining --enable-tvm-ffi with --host-target raises an error; drop --enable-tvm-ffi and use the plain AOT export path described here.

  • CUDA runtime version must match. The exported object depends on the CUDA runtime; the target’s CUDA runtime/toolkit version must match the one CuTe DSL was built against (see Object File Compatibility Issues).

  • Linking is your responsibility. You must supply your own cross toolchain and a target sysroot with the CUDA headers and libraries; the stub only resolves the CuTe DSL runtime symbols.

Host Cross-Compilation for QNX 8.0#

The same AOT export targets QNX 8.0 on AArch64. Only the static-linking integration is supported: the exported object is linked into your final QNX shared library or executable at build time. Dynamic loading (CuteDSLRT_Module_*, cute.runtime.load_module) is not available on QNX, because the module loader is built on LLVM ORC JIT, which is not cross-built for that platform. Those entry points still exist and return CuteDSLRT_Error_UnsupportedOnPlatform.

Select the target with the qnx8-aarch64 preset:

compiled = cute.compile(
    my_function, *args,
    options="--gpu-arch sm_110a --host-target qnx8-aarch64")

compiled.export_to_c(file_path="./artifacts", file_name="kernel",
                     function_prefix="kernel")

The preset maps to the triple aarch64-unknown-nto-qnx8.0.0. LLVM has no QNX target, so that triple resolves to an unknown OS and generic AArch64 ELF codegen: the emitted object is identical to the one linux-aarch64 produces. The distinct triple exists so that the runtime-library lookup below can tell the two targets apart. Because the object is plain AArch64 ELF and QNX uses the same AAPCS64 ABI, the QNX linker consumes it directly.

Cross-link with the QNX 8.0 toolchain:

q++ -Vgcc_ntoaarch64le -shared -o kernel.so kernel.o \
    $(python -m cutlass.cute.export.aot_config --ldflags --target aarch64-unknown-nto-qnx8.0.0) \
    $(python -m cutlass.cute.export.aot_config --libs    --target aarch64-unknown-nto-qnx8.0.0)

Obtaining the Target Runtime#

The wheel by default ships only the link-time stub for QNX, under lib/stubs/aarch64-unknown-nto-qnx8.0.0/. It resolves the CuTe DSL runtime symbols on the build host and is never executed. This mirrors how the CUDA toolkit ships lib/stubs/libcuda.so.

The real libcute_dsl_runtime.so for QNX is obtained separately via the extras nvidia-cutlass-dsl[qnx] or nvidia-cutlass-dsl[qnx-cu13] on x86_64 Linux, and is installed at lib/aarch64-unknown-nto-qnx8.0.0/. It is only needed on the target: deploy it into the QNX root filesystem as part of your image build, where the dynamic loader binds against it at run time.

Version Compatibility#

Because the runtime and the wheel are obtained separately, they must be matched explicitly:

  • The QNX runtime artifact must come from the same |DSL| build as the wheel that produced the object. The object embeds a version that is checked when a module is loaded, and the two are only guaranteed consistent when they are built together.

  • The QNX CUDA toolkit on the target must match the CUDA version CuTe DSL was built against, exactly as in the AArch64 Linux case (see Object File Compatibility Issues).

QNX Limitations#

  • Static linking only. Dynamic loading is unavailable; see above.

  • AArch64 only. QNX on other architectures is not supported.

  • Restricted CUDA surface. A safety or otherwise restricted QNX CUDA build may not provide every CUDA entry point the CuTe DSL runtime references, which surfaces as an unresolved symbol when linking the runtime for QNX.

Supported Argument Types#

CuTe DSL supports the following argument types:

  • cute.Tensor

  • cute.Shape / cute.Coord / cute.Tile / cute.IntTuple / cute.Stride

  • cuda.CUstream

  • cutlass.Int8 / cutlass.Int16 / cutlass.Int32 / cutlass.Int64 / cutlass.Boolean

  • cutlass.Uint8 / cutlass.Uint16 / cutlass.Uint32 / cutlass.Uint64

  • cutlass.Float32 / cutlass.TFloat32 / cutlass.Float64 / cutlass.Float16

Note that:

  1. cute.Tensor is a dynamic tensor type that only contains dynamic shapes and strides in its ABI representation. As a result, different compilations may produce different tensor ABIs. This is why declarations for each tensor type are included in the generated header file.

  2. strides in cute.Tensor are determined by the use_32bit_strides compile argument. When use_32bit_strides is set to True, the strides are 32-bit; when set to False, they are 64-bit.

  3. Currently, custom types are not supported for AOT compilation.

Object File Compatibility Issues#

The object file generated by CuTe DSL depends on the CUDA runtime library. Therefore, ensure that the version of the CUDA runtime/toolkit library matches the version used by CuTe DSL. Otherwise, ABI compatibility with the CUDA runtime cannot be guaranteed.

When using C++ static linking integration, compatibility is assured because the header and object files are generated together and guaranteed to match.

For C++ dynamic loading integration and Python loading, the binary file is loaded at runtime. To ensure compatibility, version information is embedded in the metadata of the generated binary file. At runtime, this version information is checked, and if it does not match the expected version, the binary file will be rejected.

Relation to Apache TVM FFI AOT#

Apache TVM FFI AOT offers a comparable capability, enabling TVM functions to be compiled into binary files that can be loaded and executed at runtime. For more information, see the section “Exporting Compiled Module” in TVM FFI Compilation.

The primary distinction is that, when TVM FFI is enabled, CuTe DSL generates a dedicated wrapper function on top of the underlying CuTe ABI. This wrapper adheres to the calling conventions defined by TVM FFI. In contrast, the CuTe ABI entry function is specified directly in the generated header file, which affects how arguments must be provided.

For instance, with the TVM FFI wrapper function, users are able to pass in arguments such as torch.Tensor directly. However, when calling the CuTe ABI entry function, arguments should be provided as cute.Tensor types.