Just-In-Time (JIT) Compilation in cuEST#

Note

JIT compilation is available only in the cuEST CUDA 13 package (built against CUDA Toolkit 13.0.2 or later). In the cuEST CUDA 12 package, JIT initialization is skipped entirely at handle creation, so affected routines fall back to their precompiled kernels instead of a JIT-specialized one; explicitly requesting JIT mode with the CUDA 12 package returns CUEST_STATUS_INVALID_ARGUMENT. Install the cuEST CUDA 13 package (with a matching CUDA 13.x toolkit and driver) to enable JIT specialization.

JIT compilation also requires the CUDA runtime-compilation libraries NVRTC and nvJitLink to be available at run time. Both ship with the CUDA Toolkit but are not bundled with cuEST, so a CUDA Toolkit installation providing libnvrtc and libnvJitLink must be present and discoverable (for example, on LD_LIBRARY_PATH) on every machine that runs cuEST. See the Installation Guide for the full prerequisites.

Warning

Behavior change: the first compute call into a JIT-compiled routine now compiles the kernels it needs at run time, so that call is markedly slower than later ones and requires the CUDA Toolkit’s NVRTC and nvJitLink libraries to be present at run time. Compiled kernels are cached to disk (by default under ~/.cuest_cache/), so the cost is paid once per machine and reused by every later run. No source or API change is required to benefit; the cache directory and the number of compile threads are configurable on the handle (see Controlling JIT Compilation).

Overview#

cuEST uses JIT compilation for a number of its compute routines, currently: cuestPotentialCompute, cuestPotentialDerivativeCompute, cuestDFCoulombCompute, cuestDFSymmetricDerivativeCompute, cuestPCMPotentialCompute, cuestPCMDerivativeCompute, and cuestPCMRadiiDerivativeCompute; the set is expected to grow across releases. The GPU kernels for these routines are specialized to the characteristics of a given calculation and compiled at run time rather than shipped fully pre-built. The space of possible specializations is too large to pre-compile and ship in its entirety, so cuEST compiles only the specific kernels a given calculation needs, on demand.

JIT compilation is not applied uniformly to every function. To find out whether a particular routine uses JIT compilation, consult that function’s entry in the C API reference, where JIT-compiled routines are identified as such.

Compilation happens during the compute phase: the first compute call that needs a given specialized kernel triggers its compilation, and the result is cached (see Disk Caching) so that later calls reuse it. Creating handles and plans, and querying workspaces, do not compile kernels. The compilation runs on the host and does not use the CUDA stream held by the cuEST handle – it neither enqueues work on nor synchronizes that stream.

JIT compilation changes only how these kernels are built, not the numerical results they produce: in principle a JIT-compiled kernel yields the same output as the equivalent statically compiled kernel.

Note

The JIT compilation described on this page is algorithm-level JIT: cuEST compiles specialized kernels tailored to a given calculation from source or LTO IR at run time, and caches the resulting kernels in its own disk cache (see Disk Caching).

This is distinct from the CUDA driver’s PTX JIT compilation, which translates PTX into machine code at module load time. The driver maintains its own, separate compute cache (default location ~/.nv) and is controlled by the CUDA JIT environment variables documented in the CUDA Programming Guide. cuEST compiles its JIT kernels directly to the target GPU architecture rather than to PTX, so the driver’s PTX JIT path is not exercised for these kernels.

What Changed#

For users upgrading from a release without JIT compilation, the affected routines behave identically in their inputs, outputs, and numerical results – only when and how their kernels are built has changed:

  • New CUDA version requirement: JIT compilation is available only in the cuEST CUDA 13 package. In the cuEST CUDA 12 package, JIT initialization is skipped at handle creation, and affected routines run their precompiled kernels instead of a JIT-specialized one.

  • First-call latency: the first invocation of an affected routine (with a cold cache) triggers a compile/link step and is slower than subsequent calls. See Compilation Times for representative figures.

  • On-disk cache: compiled kernels are cached (see below), so the cost is paid once per machine, not once per run.

  • New runtime dependencies: cuEST now calls the CUDA Toolkit’s runtime-compilation libraries, NVRTC (libnvrtc) and nvJitLink (libnvJitLink), when compiling and linking kernels. These libraries are part of the CUDA Toolkit and are not bundled with cuEST, so they must be installed and discoverable on every machine that runs the library.

  • New failure modes: because handle creation now configures the JIT compiler, cuestCreate validates the JIT settings and can return CUEST_STATUS_INVALID_ARGUMENT for an unsafe cache directory (a symlink, or a world-writable directory without the sticky bit) or a compile-thread count below 1 (see Controlling JIT Compilation). At compute time, a missing NVRTC/nvJitLink library, an unwritable cache directory, or an unsupported GPU architecture surfaces as a compute-call error rather than at handle creation.

Table 5 JIT behavior at a glance#

Aspect

Before

With JIT compilation

First call to affected routine

Runs immediately; all kernels shipped pre-built

Compiles the needed kernels on demand (cold cache), then runs; slower than later calls

Subsequent calls

Same cost every call

Load the kernel from the in-memory or on-disk cache; no recompilation

Runtime dependencies

CUDA runtime / driver only

Also requires NVRTC and nvJitLink from the CUDA Toolkit (not bundled with cuEST)

CUDA version

CUDA Toolkit 12.x or 13.x

JIT specialization requires the cuEST CUDA 13 package; with the cuEST CUDA 12 package, routines run their precompiled kernels instead

Controlling JIT Compilation in cuEST#

Two aspects of JIT compilation are configurable on the handle parameters object passed to cuestCreate. Both are set with cuestParametersConfigure on CUEST_HANDLE_PARAMETERS and take effect when the handle is created; both also support cuestParametersQuery.

Table 6 JIT handle-parameter attributes#

Attribute

Type

Meaning

CUEST_HANDLE_PARAMETERS_JIT_CACHE_DIR

const char *

Directory for the on-disk kernel cache. The empty string (the default) selects the derived default, ~/.cuest_cache/cuest-cuda<cuda_major>-v<cuest_version>/ (see Disk Caching). Must be a trusted, per-user path: an existing directory that is a symlink, or world-writable without the sticky bit, is rejected at cuestCreate with CUEST_STATUS_INVALID_ARGUMENT. A non-existent path is created under the caller’s control. A string returned by a query is allocated by cuEST with malloc and must be freed by the caller with C’s native free().

CUEST_HANDLE_PARAMETERS_JIT_COMPILE_THREADS

int32_t

Number of parallel worker threads used to compile kernels. Must be >= 1 (a smaller value is rejected at cuestCreate with CUEST_STATUS_INVALID_ARGUMENT) and is clamped to an internal upper bound of 256. Default: 16. These threads run on the host and do not touch the handle’s CUDA stream.

Note

The cache directory is security-sensitive: cuEST writes compiled cubins there and, on later runs, loads and executes them by filename match. Pointing it at a location another user can write to would let a planted file be executed, which is why an unsafe existing directory is rejected. Keep it per-user and non-world-writable; see Concurrency and Shared Storage for the shared-filesystem guidance.

See Examples below for the corresponding C and Python snippets.

Turning JIT Off Per Routine#

Unlike the two handle-wide settings above, JIT usage is switched on or off per routine, on that routine’s own compute-parameters object (not the handle parameters). Each JIT-eligible routine exposes a *_JIT_USAGE_MODE attribute of type cuestJITUsageMode_t, set with cuestParametersConfigure before the compute call – for example CUEST_POTENTIALCOMPUTE_PARAMETERS_JIT_USAGE_MODE for cuestPotentialCompute. Consult a given routine’s entry in the C API reference for its exact attribute name.

Table 7 cuestJITUsageMode_t values#

Value

Meaning

CUEST_JIT_USAGE_MODE_ON

Use the JIT-compiled, specialized kernel path.

CUEST_JIT_USAGE_MODE_OFF

Use the routine’s precompiled (AOT), non-specialized kernel path instead – the same fallback path used automatically in the cuEST CUDA 12 package.

The default is CUEST_JIT_USAGE_MODE_ON in the cuEST CUDA 13 package and CUEST_JIT_USAGE_MODE_OFF in the cuEST CUDA 12 package. Explicitly requesting CUEST_JIT_USAGE_MODE_ON with the CUDA 12 package is rejected at compute time with CUEST_STATUS_INVALID_ARGUMENT, since JIT is unavailable there (see the note at the top of this page).

Disk Caching#

cuEST caches JIT-compiled kernels to disk, so the compile cost for a given kernel is paid only once per machine rather than once per process. The first call that needs a particular specialization compiles it and writes the result to the cache; every later call – in the same run or any future run – loads the precompiled kernel from disk instead of recompiling.

Cache Location#

By default the cache lives under the user’s home cache directory, in a folder scoped by both the CUDA driver’s major version and the cuEST library version:

~/.cuest_cache/cuest-cuda<cuda_major>-v<cuest_version>/

# for example, under CUDA 13 with cuEST 2.3.4:
~/.cuest_cache/cuest-cuda13-v2.3.4/

Scoping the folder this way keeps caches from different cuEST releases and different CUDA environments physically separated. CUDA enters as the major version only (as reported by the CUDA driver), because compiled-kernel compatibility is guaranteed within a CUDA major version but not across it; cuEST enters as the full MAJOR.MINOR.PATCH, so even a patch release cannot reuse kernels compiled by a previous release.

What Is Stored#

Each cache entry is a compiled kernel for one specific specialization, GPU architecture, CUDA major version, and cuEST version. Entries are keyed so that any change to those parameters produces a distinct file. A single compiled kernel ranges from tens of kilobytes to a few megabytes depending on the specialization (high angular-momentum and gradient kernels are the largest). Because a full run spans many specializations, a fully populated cache for a given version/architecture combination can grow from roughly a hundred megabytes up to about a gigabyte.

Cache Invalidation#

The cache is self-invalidating across the dimensions that affect kernel compatibility:

  • cuEST or CUDA major upgrade – a new version writes into a new, separately named folder, so it never reads or overwrites an older version’s kernels.

  • Different GPU architecture – kernels are compiled for the real GPU architecture (not a portable intermediate form), for both load-time performance and IP protection. Running on a GPU architecture not already present in the cache results in a miss and a recompile.

To clear the cache manually, delete the relevant version-scoped folder.

Concurrency and Shared Storage#

The disk cache is safe to share across processes on one host, and across hosts on a shared POSIX filesystem. When the cache directory resides on shared storage – for example a network-mounted home directory – an entire cluster or team can amortize JIT compile time against a single cache. Concurrent writers that miss the same entry cannot corrupt one another: cache writes are committed atomically, so a reader always sees either a complete kernel or none, never a partial file. If two writers race, both produce identical output and the cache simply keeps one.

When the cache directory is on shared storage, keep the following in mind:

Setting

Guidance

Filesystem

Use a POSIX filesystem (Lustre, GPFS, BeeGFS, NFSv4+). Avoid SMB/CIFS without atomic-rename support and object stores (S3/GCS), which are not POSIX and can corrupt or duplicate cache entries.

Permissions

Make the directory group-readable/writable by all consumers, and set the sticky bit so users cannot delete one another’s files.

Quota

Plan for several hundred megabytes up to about a gigabyte per (cuEST version × CUDA version × GPU architecture) combination.

Operational guidance:

  • The first user pays the compile cost; everyone after loads from cache. For large fan-out cluster jobs, run a small representative warm-up calculation on one node first to populate the shared cache before the main job launches.

  • Do not delete the cache directory while jobs are running. A reader that hits a file mid-delete falls back to recompiling harmlessly, but bulk deletion during active runs is best avoided.

  • Do not manually copy kernels between version-scoped folders. Compiled kernels are not portable across CUDA major versions or GPU architectures; the folder scoping already handles this correctly.

Safe Operations on the Cache#

The cache is content-addressed and version-scoped, so most operations on it are safe by construction. The following summarizes what is safe to do and what to avoid.

Safe to do:

  • Run many processes and threads against the same cache directory at once, on one host or across hosts on a shared POSIX filesystem. Cache writes are committed atomically, so readers never see a partial file.

  • Let multiple cuEST versions and CUDA major versions share one root directory; they are isolated in separate version-scoped folders and cannot interfere.

  • Let multiple GPU architectures share one folder; entries are tagged by architecture and never collide.

  • Delete an entire version-scoped folder (for example to reclaim space). A process that hits a file mid-delete simply recompiles it, harmlessly.

Avoid:

  • Placing the cache on non-POSIX storage – SMB/CIFS without atomic-rename support, or object stores such as S3/GCS – which can corrupt or duplicate entries.

  • Manually copying, moving, renaming, or editing cache files, or moving them between version-scoped folders. Compiled kernels are not portable across CUDA major versions or GPU architectures; tampering with filenames or contents can cause a load failure that forces a cache clear.

  • Bulk-deleting the cache directory while jobs are actively running. It is tolerated, but best done when no jobs are using the cache.

Performance Considerations#

JIT compilation affects only the first compute call into an affected routine on a given machine; it does not change the steady-state performance of the kernels themselves. A JIT-compiled kernel is built for the real GPU architecture and is intended to match the runtime performance of the equivalent statically compiled kernel – the trade-off is purely the one-time, up-front compile cost, which is amortized by the disk cache. The cold-cache cost depends on the basis set and math mode; representative figures for the potential routine are tabulated under Compilation Times below. The compile work is parallelized across CUEST_HANDLE_PARAMETERS_JIT_COMPILE_THREADS host threads (default 16), so a larger thread count can reduce cold-cache latency on machines with spare CPU cores.

Because compilation happens on the first compute call, that call is substantially slower than later ones until the disk cache is warm. For latency-sensitive applications – and before collecting performance benchmarks – we suggest running a small, representative calculation first, one that exercises the same basis set and math mode as the production workload, to populate the cache so that the compilation cost is not attributed to later calls.

Compilation Times#

The table below reports the total cold-cache (first-call) JIT compilation time for the potential integral routine (cuestPotentialCompute), as a function of the basis set and the handle’s math mode, set with cuestSetMathMode (see Emulation and Mixed Precision in cuEST). The two modes are abbreviated below as FP64 (CUEST_NATIVE_FP64_MATH_MODE, native FP64 arithmetic everywhere) and Emulation (CUEST_DEFAULT_MATH_MODE, mixed-precision kernels chosen heuristically for performance). This is the one-time cost paid before the disk cache is warm; subsequent runs load the precompiled kernels from the cache and skip it entirely. Larger basis sets require more kernels to be compiled, and therefore take longer. The Emulation kernels compile modestly more slowly than the FP64 path (roughly 1.05–1.25x here).

The compilation cost also depends on the compute routine: the figures below are for cuestPotentialCompute, and other routines compile a different number and mix of kernels. The dominant factor is angular momentum – high-angular-momentum kernels are substantially more expensive to compile, so a basis with diffuse, high-\(L\) shells costs disproportionately more than its kernel count alone suggests. In most cases, however, the cold-cache compilation completes within about a minute with 16 JIT compile threads, and – being cached to disk – is paid only once per machine.

Note

These measurements were taken on an NVIDIA GB10 GPU (driver 580.126.09) with CUDA Toolkit 13.0 and cuEST 0.2.0, using 16 JIT compile threads. Absolute times are hardware- and toolchain-dependent; treat them as relative indicators rather than guarantees.

Table 8 Cold-cache JIT compilation time for the potential routine#

Basis

Math mode

Compilation time (ms)

sto-3g

FP64

938.0

def2-SVP

FP64

991.1

def2-TZVP

FP64

3285.9

def2-QZVP

FP64

5163.2

sto-3g

Emulation

1055.1

def2-SVP

Emulation

1097.1

def2-TZVP

Emulation

3487.2

def2-QZVP

Emulation

6457.3

Examples#

The JIT controls live on the handle parameters, so they are configured between cuestParametersCreate and cuestCreate. The snippets below set a custom cache directory and compile-thread count; any later compute call into a JIT-compiled routine then uses them. The full runnable program is the 0_context/basic_usage example shipped with the C API.

C#

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <cuest.h>
#include <helper_status.h>

int main(void)
{
    cuestHandle_t handle;
    cuestHandleParameters_t handle_parameters;

    /* Create the handle parameters (sets reasonable defaults). */
    checkCuestErrors(cuestParametersCreate(
        CUEST_HANDLE_PARAMETERS,
        &handle_parameters));

    /*
     * Configure the JIT compiler. The cache directory must be a trusted,
     * per-user, non-world-writable path; the empty string (default) selects
     * ~/.cuest_cache/cuest-cuda<N>-v<version>/. Compile threads must be >= 1.
     */
    const char *jit_cache_dir = "/tmp/cuest-jit-cache";
    int32_t jit_compile_threads = 8;
    checkCuestErrors(cuestParametersConfigure(
        CUEST_HANDLE_PARAMETERS, handle_parameters,
        CUEST_HANDLE_PARAMETERS_JIT_CACHE_DIR,
        &jit_cache_dir, sizeof(char *)));
    checkCuestErrors(cuestParametersConfigure(
        CUEST_HANDLE_PARAMETERS, handle_parameters,
        CUEST_HANDLE_PARAMETERS_JIT_COMPILE_THREADS,
        &jit_compile_threads, sizeof(int32_t)));

    /*
     * A queried string attribute is allocated by cuEST and must be freed by
     * the caller with C's native free().
     */
    char *queried_cache_dir = NULL;
    checkCuestErrors(cuestParametersQuery(
        CUEST_HANDLE_PARAMETERS, handle_parameters,
        CUEST_HANDLE_PARAMETERS_JIT_CACHE_DIR,
        &queried_cache_dir, sizeof(char *)));
    fprintf(stdout, "JIT cache directory: %s\n", queried_cache_dir);
    free(queried_cache_dir);

    /* JIT settings are applied here; an unsafe cache_dir or
     * compile_threads < 1 fails with CUEST_STATUS_INVALID_ARGUMENT. */
    checkCuestErrors(cuestCreate(handle_parameters, &handle));

    /* Parameters can be destroyed immediately after handle creation. */
    checkCuestErrors(cuestParametersDestroy(
        CUEST_HANDLE_PARAMETERS, handle_parameters));

    /* ... set up basis/plans and call JIT-compiled compute routines ... */

    checkCuestErrors(cuestDestroy(handle));
    return 0;
}

Python#

The Python bindings mirror the C API one-to-one, using small value-holder objects (data_string, data_int32_t) for the attribute value. Configure the JIT attributes on the handle parameters before cuestCreate:

import cuest.bindings as ce

params = ce.cuestHandleParameters()
ce.cuestParametersCreate(
    parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
    outParameters=params,
)

# Cache directory (empty string -> derived default; must be a trusted,
# per-user, non-world-writable path).
jit_cache_dir = ce.data_string()
jit_cache_dir.value = '/tmp/cuest-jit-cache'
ce.cuestParametersConfigure(
    parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
    parameters=params,
    attribute=ce.CuestHandleParametersAttributes.CUEST_HANDLE_PARAMETERS_JIT_CACHE_DIR,
    attributeValue=jit_cache_dir,
)

# Parallel compile threads (>= 1; default 16).
jit_compile_threads = ce.data_int32_t()
jit_compile_threads.value = 8
ce.cuestParametersConfigure(
    parametersType=ce.CuestParametersType.CUEST_HANDLE_PARAMETERS,
    parameters=params,
    attribute=ce.CuestHandleParametersAttributes.CUEST_HANDLE_PARAMETERS_JIT_COMPILE_THREADS,
    attributeValue=jit_compile_threads,
)

# JIT settings are applied at handle creation; an unsafe cache directory or a
# compile-thread count < 1 returns CUEST_STATUS_INVALID_ARGUMENT here.
handle = ce.cuestHandle()
ce.cuestCreate(parameters=params, handle=handle)

Troubleshooting#

The most common JIT-specific issues and their resolutions are listed below; see the Troubleshooting page for general problems.

  • cuEST fails to load because NVRTC / nvJitLink is not discoverable. cuEST links directly against libnvrtc and libnvJitLink, so if either is missing, loading the cuEST library itself fails (for example, a Python ImportError or a dynamic-linker error before any cuEST call is made) – this happens on every CUDA Toolkit version, not just where JIT runs. Install the matching CUDA Toolkit components and ensure they are on LD_LIBRARY_PATH.

  • Routine runs its precompiled kernel instead of a JIT-specialized one. JIT compilation is available only in the cuEST CUDA 13 package; in the cuEST CUDA 12 package, JIT initialization is skipped at handle creation, so affected routines fall back to their precompiled kernels, and explicitly requesting JIT mode returns CUEST_STATUS_INVALID_ARGUMENT. Install the cuEST CUDA 13 package (built against CUDA Toolkit 13.0.2 or newer) to enable JIT specialization.

  • cuestCreate returns CUEST_STATUS_INVALID_ARGUMENT. A JIT handle parameter was rejected: the cache directory is a symlink or is world-writable without the sticky bit, or JIT_COMPILE_THREADS is below 1. Point the cache at a trusted, per-user, non-world-writable directory and use a thread count >= 1 (see Controlling JIT Compilation).

  • Cache directory not writable. If the configured (or default) cache directory cannot be written, the compile step fails at the first affected compute call. Choose a writable location via CUEST_HANDLE_PARAMETERS_JIT_CACHE_DIR, or fix the directory permissions.

  • HOME not set. The default cache location is derived from the HOME environment variable, so handle creation fails with an error if HOME is unset (as can happen in minimal containers or batch environments) and no cache directory was configured. Either set HOME or pass an explicit CUEST_HANDLE_PARAMETERS_JIT_CACHE_DIR.

  • Unsupported compute capability / architecture. Kernels are compiled for the real GPU architecture; a device whose architecture the toolkit cannot target fails to compile. Use a CUDA Toolkit new enough for the device (see the CUDA version requirements at the top of this page).