Skip to main content

Documentation Index

Fetch the complete documentation index at: https://mintlify.com/Deepak-Sangle/TornadoVM/llms.txt

Use this file to discover all available pages before exploring further.

CUDA Graphs let NVIDIA GPUs execute a pre-recorded sequence of GPU commands — kernels, memory copies, library calls — with a single cuGraphLaunch instead of the individual dispatch overhead of each iteration. For compute loops that repeat the same pipeline hundreds or thousands of times, the per-iteration host-side launch cost of setting up kernel arguments, submitting commands to the CUDA driver, and synchronizing stream dependencies can become a measurable bottleneck. TornadoVM exposes CUDA Graph capture through a single method on TornadoExecutionPlan: call plan.withCUDAGraph() and the runtime records iteration zero, then replays the captured graph for every subsequent execution.
CUDA Graph support requires the CUDA backend. Build with make BACKEND=cuda. Both JIT-compiled Java kernels and library tasks (cuBLAS, cuFFT, cuDNN, etc.) are captured together in the same graph.

How CUDA Graph Capture Works in TornadoVM

1

Call plan.withCUDAGraph()

Enable graph capture mode on the execution plan before the iteration loop. This is the only code change required.
try (TornadoExecutionPlan plan = new TornadoExecutionPlan(graph.snapshot())) {
    plan.withCUDAGraph();   // arm capture
    for (int i = 0; i < iterations; i++) {
        plan.execute();     // iteration 0: capture; i > 0: replay
    }
}
2

Iteration 0: Graph Capture

On the first call to plan.execute(), TornadoVM runs a pre-compilation pass (the prepare() hook on all library tasks) to create native handles, cuFFT plans, cuDNN descriptors, and workspace allocations. Allocation is illegal inside a CUDA capture region — prepare() guarantees it runs before capture begins. The runtime then opens a capture on the backend CUDA stream, executes the full pipeline once (transfers, JIT kernels, library calls), and closes the capture to produce a CUgraph. Per-call profiler timing is disabled during this iteration.
3

Iterations 1+: Graph Replay

All subsequent plan.execute() calls dispatch the captured graph with a single cuGraphLaunch. The GPU executes the exact same command sequence as iteration 0 — no CPU-side argument marshalling, no per-task CUDA API calls, no stream-dependency setup. The host returns as soon as the launch is submitted.

The prepare() Hook: Why It Matters

The most important constraint for CUDA Graph capture is that no device memory allocations may happen inside the capture region. Library contexts (cuBLAS handles, cuFFT plan workspaces, cuDNN descriptors) all require device allocations. TornadoVM handles this automatically through the prepare() hook:

What prepare() Does

Before capture starts, the runtime calls prepare() on every library task in the graph. Each provider uses this hook to create and cache per-shape resources:
  • cuBLAS: handle creation and optional workspace
  • cuFFT: plan creation and device work area allocation
  • cuDNN: convolution descriptors, algorithm selection, grow-only workspace
  • cuBLASLt: matmul plan, 32 MiB device workspace
prepare() is idempotent — repeated calls with the same shape hit the cache and return immediately.

What dispatch() Does

Inside the capture region, dispatch() is called for each library task. Because prepare() has already allocated everything, dispatch() makes no device allocations — it only issues the native library call with the pre-computed handles, plans, and pointers. This makes dispatch() safe to call inside a CUDA stream capture.
Changing the data size or device buffer addresses between replays invalidates the captured graph (device pointers are baked in at capture time). If your problem size changes, close the current TornadoExecutionPlan, create a new one, and let iteration 0 re-capture at the new shape.

Complete Example: Matrix Multiply Loop

The following example runs a cuBLAS SGEMM inside a CUDA Graph-captured execution plan for 500 iterations. Iteration 0 captures; iterations 1–499 replay the captured graph.
import uk.ac.manchester.tornado.api.*;
import uk.ac.manchester.tornado.api.enums.DataTransferMode;
import uk.ac.manchester.tornado.api.types.arrays.FloatArray;
import uk.ac.manchester.tornado.cublas.CuBlas;
import uk.ac.manchester.tornado.cublas.enums.CuBlasOperation;

int size = 2048;
FloatArray matA   = new FloatArray(size * size);
FloatArray matB   = new FloatArray(size * size);
FloatArray output = new FloatArray(size * size);

// ... initialize matA and matB ...

TaskGraph graph = new TaskGraph("sgemm_loop")
    .transferToDevice(DataTransferMode.EVERY_EXECUTION, matA, matB)
    .task("preprocess", MyKernels::scale, matA, 0.5f)      // JIT kernel
    .libraryTask("sgemm", CuBlas::cublasSgemm,             // cuBLAS library task
            CuBlasOperation.CUBLAS_OP_N.operation(),
            CuBlasOperation.CUBLAS_OP_N.operation(),
            size, size, size,
            1.0f, matB, size, matA, size,
            0.0f, output, size)
    .task("postprocess", MyKernels::relu, output)          // JIT kernel
    .transferToHost(DataTransferMode.EVERY_EXECUTION, output);

try (TornadoExecutionPlan plan = new TornadoExecutionPlan(graph.snapshot())) {
    plan.withCUDAGraph();   // arm graph capture

    long startMs = System.currentTimeMillis();
    for (int iter = 0; iter < 500; iter++) {
        // Modify host-side matA here if needed (transferred on EVERY_EXECUTION)
        plan.execute();
        // iter == 0: graph captured (includes EVERY_EXECUTION transfer + JIT + cuBLAS)
        // iter > 0:  single cuGraphLaunch replays the entire recorded pipeline
    }
    long elapsedMs = System.currentTimeMillis() - startMs;
    System.out.printf("500 iterations: %d ms (%.2f ms/iter)%n",
            elapsedMs, elapsedMs / 500.0);
}
Run the built-in CUDA Graph test for cuBLAS with pre/post JIT tasks:
tornado -m tornado.cublas/uk.ac.manchester.tornado.cublas.tests.TestCuBlasSgemvWithTasksCudaGraph

What Gets Captured

Everything that runs on the backend CUDA stream is recorded into the graph:

Captured

  • EVERY_EXECUTION host-to-device transfers (recorded as cuMemcpyHtoDAsync nodes)
  • JIT-compiled Java kernels (@Parallel, KernelContext)
  • Library task dispatches (cuBLAS, cuFFT, cuDNN, cuSPARSE, CUTLASS)
  • EVERY_EXECUTION device-to-host transfers
  • Stream dependencies between all of the above

Not Captured

  • FIRST_EXECUTION transfers (these run before iteration 0, outside the capture region)
  • prepare() allocations (run in the pre-compilation pass, also before capture)
  • Per-call profiler timing (requires stream synchronization, which invalidates capture)
  • JVM-side Java code between plan.execute() calls

Profiling a CUDA Graph Workload

Because per-call profiler timing is disabled during capture, --enableProfiler console reports timing only for iteration 0 (the capture iteration). To measure steady-state throughput, time the loop externally as shown in the example above. For detailed per-kernel analysis of the captured graph, use nsys:
tornado --enableProfiler console \
  -m tornado.cublas/uk.ac.manchester.tornado.cublas.tests.TestCuBlasSgemvWithTasksCudaGraph
In the nsys CUDA GPU trace, JIT kernels and cuBLAS/cuFFT/cuDNN library kernels all appear on the same stream, confirming that the in-order CUDA stream guarantee is preserved inside the captured graph.

cuBLAS-Specific Graph Notes

When using host pointer mode (the default), the alpha and beta scalars passed to cuBLAS factory methods are baked into the captured graph by value. Replaying the graph always uses the values from iteration 0. If alpha or beta must change between replays, create a new TornadoExecutionPlan to re-capture with the new values.
The cuBLAS handle (and any user-specified workspace) is created in prepare(), before capture begins. Handle creation allocates device memory, which is illegal inside a capture region. TornadoVM ensures this happens automatically — you do not need to manage handle lifetime.
cuBLASLt matmul plans (descriptors + heuristic algorithm + 32 MiB workspace) are also created in prepare(). The cuBLASLt cublasLtMatmul call inside dispatch() is therefore allocation-free and safe to record into the graph.

CUDA Graph with cuFFT and cuDNN

CUDA Graph capture works identically for cuFFT and cuDNN library tasks.
// FFT filter pipeline captured into a CUDA Graph
TaskGraph graph = new TaskGraph("filter")
    .transferToDevice(DataTransferMode.EVERY_EXECUTION, signal)
    .libraryTask("fwd",      CuFft::cufftForwardC2C, signal, freq, n, 1)
    .task("lowpass",         Filters::lowPass, freq, cutoff)
    .libraryTask("inv",      CuFft::cufftInverseC2C, freq, out, n, 1)
    .task("normalize",       Filters::scaleBy, out, 1.0f / n)
    .transferToHost(DataTransferMode.EVERY_EXECUTION, out);

try (TornadoExecutionPlan plan = new TornadoExecutionPlan(graph.snapshot())) {
    plan.withCUDAGraph();
    for (int i = 0; i < 1000; i++) plan.execute();
}
Verify with:
tornado-test -V uk.ac.manchester.tornado.unittests.cufft.TestCuFft#testRoundTripWithCudaGraph

Limitations and Constraints

ConstraintDetail
Fixed data sizesDevice pointers are baked into the captured graph. Changing the size of any buffer invalidates the graph and requires re-capture via a new TornadoExecutionPlan.
Fixed alpha/betacuBLAS scalar arguments are captured by value (host pointer mode). Re-create the execution plan to change them.
No per-call profiling during replay--enableProfiler console only times iteration 0 (capture). Use nsys for replay timing.
cuSPARSE workspacecuSPARSE pre-allocates an 8 MiB workspace in prepare(). If the sparse matrix requires more, the graph must run once without withCUDAGraph() first to grow the workspace, then re-capture.
withBatch unsupportedTornadoExecutionPlan.withBatch() is not supported with library tasks — use the library’s own batched entry point (e.g. cublasSgemmStridedBatched).

Build docs developers (and LLMs) love