Why accessing cuBLAS, cuFFT, cuDNN from within your Java application

mikepapadim2 pts0 comments

Java at the Metal: CUDA Graphs, Tensor Cores, and cuBLAS/cuDNN/cuFFT from the JVM with TornadoVM

For most of its life, TornadoVM answered a single question well: can we take ordinary Java methods, JIT-compile them to GPU kernels, and run them without the developer writing a line of CUDA? The answer was yes — you annotate a loop with @Parallel, wrap it in a TaskGraph, and the runtime generates PTX (or OpenCL, or SPIR-V) and schedules it.

That model has a ceiling. A generated kernel is only ever as good as the compiler that generated it. Real CUDA applications do not live on generated kernels alone; they live on CUDA Graphs to amortize launch overhead, on multiple streams to overlap copies with compute, on cuBLAS/cuDNN/cuFFT for the kernels NVIDIA has spent a decade tuning, and on Tensor Cores for the matrix-multiply throughput that dominates modern AI. Until recently, reaching any of these from Java meant dropping to JNI by hand and giving up TornadoVM's memory management and scheduling.

1. TornadoVM in a Nutshell

TornadoVM is a heterogeneous programming framework that JIT-compiles Java bytecode to GPU and accelerator code through its own compiler pipeline. You do not write kernels in a DSL string; you write Java, and the runtime lowers it.

Three abstractions matter for the rest of this article:

TaskGraph — a declarative description of data movement and computation: which host arrays are copied to the device, which methods run as tasks, and which results come back. Tasks are ordinary Java method references. A TaskGraph is turned into an ImmutableTaskGraph via snapshot().

TornadoExecutionPlan — the runtime handle you actually execute. It carries execution policies: device selection, profiling, warm-up, and — new here — withCUDAGraph().

KernelContext — the lower-level, explicit programming model. Instead of @Parallel loops, you get thread indices (localIdx, groupIdx), shared memory (allocateIntLocalArray), and barriers (localBarrier()). This is where Tensor Core intrinsics live, because MMA is inherently a warp-cooperative operation.

Under the hood, a TaskGraph compiles to a bytecode program for a small stack interpreter (ALLOC, TRANSFER_HOST_TO_DEVICE_*, LAUNCH, TRANSFER_DEVICE_TO_HOST_*, DEALLOC). The interpreter walks this program and issues the corresponding backend calls. This indirection is important: it is exactly the layer that CUDA Graph capture wraps, and where the new EXECUTION_GRAPH_BEGIN_CAPTURE / EXECUTION_GRAPH_LAUNCH bytecodes were added.

TornadoVM has several backends. The NVIDIA-facing ones are ptx (the compiler emits PTX ISA directly) and the newer cuda backend used throughout this article, which emits CUDA C and compiles it at runtime with NVRTC , then loads the module with the CUDA Driver API. The cuda backend is what binds the vendor libraries and Tensor Core intrinsics discussed below to a real device stream.

The keys throughout the investigation are three CLI flags: --printKernel (dump generated device code), --printBytecodes (dump the interpreter program), and --debug (verbose runtime). Every raw capture in this article came from one of them or from Nsight Systems.

2. Runtime architecture: where each feature plugs in

flowchart TD<br>subgraph JVM<br>A[Java: TaskGraph / KernelContext] --> B[TornadoVM compilersketch + IR]<br>B --> C[TornadoVM runtimegraph compiler → bytecodes]<br>C --> D[Bytecode interpreterTornadoVMInterpreter]<br>end<br>subgraph CUDA_Backend[cuda backend]<br>D --> E[CUDA C codegen+ NVRTC compile]<br>D --> F[Library SPITornadoLibraryProvider]<br>end<br>subgraph JNI[cuda-jni / cublas-jni / cudnn-jni / cufft-jni]<br>E --> G[CUDA Driver APIcuLaunchKernel, cuStream*, cuGraph*]<br>F --> H[cuBLAS / cuDNN / cuFFTcublasLtMatmul, cudnnConvolutionForward, cufftExecC2C]<br>end<br>G --> GPU[(RTX 4090)]<br>H --> GPU

The five capabilities attach at three distinct layers:

Tensor Core MMA is a codegen feature: new compiler nodes lower KernelContext.mma* calls to inline PTX mma.sync / ldmatrix inside the generated CUDA C (box E).

CUDA Streams and CUDA Graphs are interpreter/driver features: the bytecode interpreter dispatches transfers and launches onto role streams, or records them into a graph, and the JNI calls cuStream* / cuGraph* (box G).

cuBLAS / cuDNN / cuFFT are library-SPI features: a libraryTask resolves through a ServiceLoader-discovered TornadoLibraryProvider to a native library binding (boxes F/H), sharing the same device buffers and stream as the JIT kernels.

A crucial design consequence: because library tasks share TornadoVM buffers and stream, and because both streams and graphs sit below the task abstraction, the features compose — a graph can capture JIT kernels and a cuBLAS call together, on a multi-stream plan.

3. CUDA stream management

The mechanism

A native CUDA C++ programmer overlaps a host→device copy of the next batch with the kernel of the current batch by putting them on different streams and inserting an event wait. PR #800 exposes the same capability to Java behind one...

cuda java tornadovm cublas cudnn kernels

Related Articles