Extending TornadoVM with Custom Native Library Providers

TornadoVM’s hybrid API is intentionally open: any vendor-optimized native GPU library can be plugged in as a library provider without touching a single line of the core runtime. The extension point is a standard java.util.ServiceLoader Service Provider Interface (SPI) — your module declares what it provides, the runtime discovers it automatically, and library tasks in a TaskGraph dispatch through it just like the built-in cuBLAS, cuDNN, cuFFT, cuSPARSE, and CUTLASS providers. This page walks through the complete four-step process and documents every SPI type, contract, and layout detail you need to build a production-quality provider.

Architecture Overview

When a TaskGraph reaches a .libraryTask(...) node at execution time, the TornadoVM interpreter:

Resolves each reference argument to the raw device pointer of its TornadoVM-managed buffer (skipping the 24-byte array header).
Looks up the registered TornadoLibraryProvider whose libraryName() matches the descriptor’s getLibraryName() string.
Calls prepare() (before CUDA Graph capture) to let the provider allocate any shape-dependent plans or workspaces.
Calls dispatch() at launch time, passing a LibraryInvocation with all arguments already resolved.

Library tasks ride the same CUDA stream as surrounding JIT kernels. There is no host synchronization between JIT tasks and native calls — everything executes in-order on one stream, which is also what makes CUDA Graph capture possible.

The diagram below maps the four implementation artifacts to their runtime roles:

   TaskGraph.libraryTask(...)
          │
          ▼
   LibraryTaskDescriptor          ← Factory  (Step 1)
   (.withLibrary / .withFunction / .withParameters / .withAccess)
          │  matched by libraryName()
          ▼
   TornadoLibraryProvider          ← Provider (Step 2)
   ├── createContext()  → LibraryContext  (handle, stream, plan cache)
   ├── prepare()        → allocate workspaces BEFORE graph capture
   └── dispatch()       → call into native .so via LibraryInvocation
          │  registered via
          ▼
   module-info.java  +  META-INF/services/  ← Registration (Step 3)
          │  compiled from
          ▼
   tornado-drivers/<mylib>-jni  CMake module  ← Native binding (Step 4)

Step 1 — The Factory: Build a LibraryTaskDescriptor

The factory is a plain Java class with static methods, one per supported operation. Each method constructs a LibraryTaskDescriptor that describes the call — what library handles it, what function to invoke, what the arguments are, and their access modes.

import uk.ac.manchester.tornado.api.common.Access;
import uk.ac.manchester.tornado.api.common.LibraryTaskDescriptor;
import uk.ac.manchester.tornado.api.types.arrays.FloatArray;

public final class MyLib {

    /** Unique provider ID — matched against TornadoLibraryProvider.libraryName(). */
    public static final String LIBRARY_NAME = "vendor/mylib";

    /**
     * Factory for a single-precision element-wise scale operation:
     *   out[i] = alpha * in[i]
     */
    public static LibraryTaskDescriptor myScale(float alpha, FloatArray in, FloatArray out) {
        Access[] access = {
            Access.READ_ONLY,   // alpha  (scalar — no device pointer)
            Access.READ_ONLY,   // in
            Access.WRITE_ONLY   // out
        };
        return new LibraryTaskDescriptor()
            .withLibrary(LIBRARY_NAME)
            .withFunction("myScale")
            .withParameters(new Object[] { alpha, in, out })
            .withAccess(access);
    }
}

.withLibrary()

Sets the provider ID string. Must match exactly what TornadoLibraryProvider.libraryName() returns — this is the lookup key.

.withFunction()

Names the operation. Passed verbatim to dispatch(String functionName, ...), so a single provider can multiplex many operations.

.withParameters()

The argument array in the same order as the factory method parameters. Reference types (FloatArray, HalfFloatArray, …) are resolved to device pointers; boxed primitives are passed through unchanged.

.withAccess()

The Access[] array drives TornadoVM’s data-flow analysis: which buffers to transfer to/from the device and when. Length must equal withParameters() array length.

Then call the factory from a TaskGraph exactly like a built-in provider:

TaskGraph graph = new TaskGraph("scale-demo")
    .transferToDevice(DataTransferMode.EVERY_EXECUTION, input)
    .libraryTask("scale", MyLib::myScale, 2.5f, input, output)
    .transferToHost(DataTransferMode.EVERY_EXECUTION, output);

Step 2 — The Provider: Implement TornadoLibraryProvider

The provider is the bridge between the TornadoVM runtime and your native library. It manages a per-(device, executionPlan) context object and performs the actual dispatch.

import uk.ac.manchester.tornado.runtime.library.spi.*;
import uk.ac.manchester.tornado.runtime.common.TornadoXPUDevice;
import uk.ac.manchester.tornado.api.common.LibraryTaskDescriptor;

public final class MyProvider implements TornadoLibraryProvider {

    // ── Identity ────────────────────────────────────────────────────────────

    @Override
    public String libraryName() {
        return MyLib.LIBRARY_NAME;   // "vendor/mylib"
    }

    // ── Device capability check ─────────────────────────────────────────────

    @Override
    public boolean canHandle(TornadoXPUDevice device) {
        // TornadoNativeStreamSupport is implemented only by the CUDA backend device
        return device instanceof TornadoNativeStreamSupport;
    }

    // ── Context lifecycle ───────────────────────────────────────────────────

    @Override
    public LibraryContext createContext(TornadoXPUDevice device, long executionPlanId) {
        // Retrieve the backend's CUDA stream for this execution plan
        long stream = ((TornadoNativeStreamSupport) device).getNativeStream(executionPlanId);
        // Create the native handle and bind it to the stream
        long handle = MyNativeLib.createHandle(stream);
        return new MyContext(handle, stream);
    }

    // ── Pre-capture hook (idempotent) ────────────────────────────────────────

    @Override
    public void prepare(LibraryTaskDescriptor descriptor, LibraryContext context) {
        // Allocate shape-dependent device resources HERE — before CUDA Graph
        // capture begins. Implement as a plan-cache lookup so repeated calls
        // are cheap (descriptor carries shape info from withParameters()).
        MyContext ctx = (MyContext) context;
        int n = (int) descriptor.getParameters()[0];   // example shape extraction
        ctx.ensurePlan(n);  // no-op if plan already exists for this n
    }

    // ── Dispatch ─────────────────────────────────────────────────────────────

    @Override
    public void dispatch(String functionName, LibraryInvocation call) {
        MyContext ctx = (MyContext) call.getContext();
        switch (functionName) {
            case "myScale" -> {
                float alpha = (float)  call.getArg(0);        // scalar → boxed
                long  dIn   = call.getDevicePointer(1);        // reference → raw ptr
                long  dOut  = call.getDevicePointer(2);
                int   n     = ctx.currentPlanSize();
                MyNativeLib.scale(ctx.handle, alpha, dIn, dOut, n, ctx.stream);
            }
            default -> throw new UnsupportedOperationException("Unknown function: " + functionName);
        }
    }

    // ── Cleanup ───────────────────────────────────────────────────────────────

    @Override
    public void destroyContext(LibraryContext context) {
        MyContext ctx = (MyContext) context;
        ctx.freePlans();
        MyNativeLib.destroyHandle(ctx.handle);
    }
}

The `prepare()` Contract

Device memory allocations, host synchronization, and stream queries are illegal during CUDA Graph capture. Any shape-dependent workspace or plan that your dispatch() needs must be created in prepare(), which the runtime invokes in the pre-compilation pass — before capture begins. Make prepare() idempotent via a plan cache: on the second call for the same shape it must return instantly without re-allocating.

Key SPI Types Reference

TornadoLibraryProvider

The top-level interface your module implements. Registered via ServiceLoader. Methods:

Method	Called when
`libraryName()`	At dispatch-lookup time to match the descriptor’s library ID
`canHandle(device)`	During plan construction to skip unsupported devices
`createContext(device, planId)`	Once per `(device, TornadoExecutionPlan)` pair
`prepare(descriptor, context)`	Before every launch region, before CUDA Graph capture
`dispatch(functionName, invocation)`	At `LAUNCH` bytecode execution for each library task
`destroyContext(context)`	When the execution plan is closed / garbage-collected

LibraryContext

A marker interface — implement it with your own class carrying the native handle, CUDA stream reference, plan cache, and any pre-allocated workspace. The runtime stores one instance per (device, executionPlanId) pair and passes it back to prepare() and dispatch() on every invocation.

LibraryInvocation

The per-call payload delivered to dispatch():

Method	Returns
`getArg(i)`	Original Java argument at index `i` (boxed primitive or host-side array object)
`getDevicePointer(i)`	Raw GPU device pointer for reference argument `i` (past the 24-byte header)
`isReference(i)`	`true` if argument `i` is a TornadoVM off-heap array
`getContext()`	The `LibraryContext` returned by `createContext()`
`getTuning()`	Opaque object from `LibraryTaskDescriptor.withTuning(...)`, or `null`
`isCapturing()`	`true` during CUDA Graph capture — avoid allocations if `true`

TornadoNativeStreamSupport

Implemented by the CUDA backend device. Use it in canHandle() to restrict your provider to CUDA, and in createContext() to obtain the native stream:

long stream  = ((TornadoNativeStreamSupport) device).getNativeStream(planId);
long context = ((TornadoNativeStreamSupport) device).getNativeContext(planId);

Pass stream to your library’s setStream equivalent (e.g., cublasSetStream, cudnnSetStream) so all native calls run in-order with the surrounding JIT kernels.

Step 3 — Registration

Registration is two lines in two files. Both must reference the same fully-qualified provider class name. module-info.java (in your module’s src/main/java/):

open module tornado.mylib {
    requires transitive tornado.api;
    requires tornado.runtime;

    exports vendor.mylib;                              // the factory class package

    provides uk.ac.manchester.tornado.runtime.library.spi.TornadoLibraryProvider
        with vendor.mylib.provider.MyProvider;         // the provider class
}

META-INF/services/uk.ac.manchester.tornado.runtime.library.spi.TornadoLibraryProvider (in src/main/resources/):

vendor.mylib.provider.MyProvider

Both the provides directive in module-info.java and the META-INF/services file are required. The JPMS module system uses the provides clause; ServiceLoader in unnamed-module contexts uses the META-INF file as a fallback.

Step 4 — Native Binding: CMake Module

The Java side is complete — now wire up the native .so. Add a tornado-drivers/mylib-jni/ directory following the pattern of an existing JNI module:

Host library (like cudnn-jni)
Device-code library (like cutlass-jni)

Use this layout when your library is a pre-built .so (cuDNN, cuBLAS, etc.) that you link against:

tornado-drivers/mylib-jni/
├── CMakeLists.txt          ← find_library() + target_link_libraries()
└── src/main/cpp/
    └── MyLibJNI.cpp        ← JNI bridge: Java_vendor_mylib_MyNativeLib_*

# CMakeLists.txt (minimal)
find_library(MYLIB_LIB mylib HINTS $ENV{MYLIB_ROOT}/lib)
add_library(tornado-mylib SHARED src/main/cpp/MyLibJNI.cpp)
target_include_directories(tornado-mylib PRIVATE $ENV{MYLIB_ROOT}/include ${JNI_INCLUDE_DIRS})
target_link_libraries(tornado-mylib ${MYLIB_LIB} ${CUDA_LIBRARIES})

Use this layout when you compile CUDA device code (.cu files) as part of the provider:

tornado-drivers/mylib-jni/
├── CMakeLists.txt          ← enable_language(CUDA) + FetchContent
└── src/main/cpp/
    ├── MyKernel.cu         ← __global__ kernels / CUTLASS instantiations
    └── MyLibJNI.cpp        ← JNI bridge calling MyKernel launch wrappers

# CMakeLists.txt (minimal)
cmake_minimum_required(VERSION 3.18)
enable_language(CUDA)
set(CMAKE_CUDA_ARCHITECTURES $ENV{CUDA_ARCH})

add_library(tornado-mylib SHARED
    src/main/cpp/MyKernel.cu
    src/main/cpp/MyLibJNI.cpp)
target_include_directories(tornado-mylib PRIVATE ${JNI_INCLUDE_DIRS})

After adding the CMake module, wire it into four Maven locations:

tornado-drivers/pom.xml — add a <module>mylib-jni</module> entry under the cuda-backend profile.
Root pom.xml — add the module to the cuda-backend profile’s <modules> list.
tornado-assembly/assembly.xml — include the compiled .so in the distribution archive.
tornado.py launcher — add tornado.mylib to --add-modules when the CUDA backend is present.

Make your CMakeLists.txt self-guarding: wrap the entire build in a find_library check and emit a message(WARNING "mylib not found — skipping tornado-mylib-jni") if the library is absent. This way make BACKEND=cuda still succeeds on machines without your library, and the provider reports UNSUPPORTED at runtime instead of crashing.

Reference: Data Layout, Types, and Alignment

Getting data layout wrong produces silently incorrect numerical results, which is harder to debug than a crash. Use this table before writing a single line of native code.

Row-major providers

CUTLASS and cuSPARSE treat matrices as row-major — the same layout TornadoVM FloatArray and HalfFloatArray use natively. No transposition is needed; pass arrays directly.

Column-major providers

cuBLAS is column-major by convention. For row-major TornadoVM inputs either pass the transpose operation flag (e.g., CUBLAS_OP_T for SGEMV) or swap the operand order (e.g., swap A and B for SGEMM).

TornadoVM Array Types

Java type	Precision	Notes
`FloatArray`	FP32	Standard single-precision; use for cuBLAS SGEMM / SGEMV, CUTLASS SIMT
`HalfFloatArray`	FP16	Elements via `new HalfFloat(float)` / `.get(i).getFloat32()`; use for tensor-core paths
`DoubleArray`	FP64	Use for cuFFT Z2Z double-precision transforms

Device Pointer Arithmetic

device_data_ptr = buffer_base_address + 24 bytes  (TornadoVM array header)

The 24-byte header stores metadata used by the garbage collector and the runtime. When the interpreter calls getDevicePointer(i) it already adds this offset — your dispatch() receives the pointer to the first data element and should pass it directly to the native API. Alignment: All TornadoVM device buffers are guaranteed to be 8-byte aligned after the header. For FP16 CUTLASS kernels (which use 4-half vector loads), k and n must therefore be multiples of 4 — the factory should validate this and throw with a clear message for non-conforming shapes.

Existing Providers as Reference Implementations

The six built-in providers cover the full range of implementation patterns. Study the one closest to your use-case before starting:

tornado-cublas/

Dense FP32/FP16/TF32 GEMM and GEMV. Shows column-major transpose handling, beta != 0 READ_WRITE access, and the batched-strided variant.

tornado-cudnn/

Convolution, pooling, activations, softmax, and SDPA. The most complex prepare() implementation — per-shape cuDNN convolution descriptor cache + workspace allocation.

tornado-cufft/

Per-(n, batch) plan cache, real↔complex (R2C/C2R) and 2D transforms. Clean example of minimal provider state.

tornado-cutlass/

Device-code compilation via CMake FetchContent, fused epilogues (bias + ReLU/GELU), and FP16 alignment validation in the factory.

tornado-cusparse/

CSR sparse matrix products. Shows pre-allocated fixed-size workspace in prepare() and the 8 MiB workspace CUDA Graph constraint.

tornado-cufft/ (FP64)

Z2Z double-precision FFT. Minimal delta from the FP32 provider — good template for adding a new precision variant to an existing library.

Unit Test Structure

TornadoVM’s unit tests double as worked integration examples. Locate the test class for the provider you’re modelling, then mirror that structure for your own:

tornado-unittests/src/main/java/uk/ac/manchester/tornado/unittests/
├── cublas/
│   ├── TestCuBlas.java          ← SGEMM, SGEMV, batched
│   └── TestCuBlasLt.java        ← FP16 / FP8 matmul, fused epilogues
├── cufft/
│   └── TestCuFft.java           ← C2C, R2C, 2D transforms
├── cudnn/
│   └── TestCuDnn.java           ← conv2d, pooling, SDPA
├── cusparse/
│   └── TestCusparse.java        ← SpMV, SpMM
└── cutlass/
    └── TestCutlass.java         ← SGEMM, HGEMM, fused bias+ReLU/GELU

Run the full library test suite for a single provider:

# Run all cuBLAS tests
tornado-test -V uk.ac.manchester.tornado.unittests.cublas.TestCuBlas

# Run with the profiler to see per-task kernel timings
tornado --enableProfiler console \
  -V uk.ac.manchester.tornado.unittests.cutlass.TestCutlass

Add your provider’s test class to the same package convention (uk.ac.manchester.tornado.unittests.<yourlib>/) and guard CUDA-only tests with a @Before method that checks getTornadoRuntime().getDefaultDevice().getTornadoVMBackend() and throws a typed TornadoVMCUDANotSupported exception when the backend is not CUDA. The test runner counts that as [UNSUPPORTED] rather than a failure — follow the pattern in TestCuBlas.cuBlasMustBeAvailable().

Get Started

Programming Guide

Backends

NVIDIA Integration

Guides

Contributing

Extending TornadoVM with Custom Native Library Providers

Architecture Overview

Step 1 — The Factory: Build a LibraryTaskDescriptor

.withLibrary()

.withFunction()

.withParameters()

.withAccess()

Step 2 — The Provider: Implement TornadoLibraryProvider

The `prepare()` Contract

Key SPI Types Reference

Step 3 — Registration

Step 4 — Native Binding: CMake Module

Reference: Data Layout, Types, and Alignment

Row-major providers

Column-major providers

TornadoVM Array Types

Device Pointer Arithmetic

Existing Providers as Reference Implementations

tornado-cublas/

tornado-cudnn/

tornado-cufft/

tornado-cutlass/

tornado-cusparse/

tornado-cufft/ (FP64)

Unit Test Structure

Build docs developers (and LLMs) love

Get Started

Programming Guide

Backends

NVIDIA Integration

Guides

Contributing

Documentation Index

​Architecture Overview

​Step 1 — The Factory: Build a LibraryTaskDescriptor

.withLibrary()

.withFunction()

.withParameters()

.withAccess()

​Step 2 — The Provider: Implement TornadoLibraryProvider

​The prepare() Contract

​Key SPI Types Reference

​Step 3 — Registration

​Step 4 — Native Binding: CMake Module

​Reference: Data Layout, Types, and Alignment

Row-major providers

Column-major providers

​TornadoVM Array Types

​Device Pointer Arithmetic

​Existing Providers as Reference Implementations

tornado-cublas/

tornado-cudnn/

tornado-cufft/

tornado-cutlass/

tornado-cusparse/

tornado-cufft/ (FP64)

​Unit Test Structure

Build docs developers (and LLMs) love

Architecture Overview

Step 1 — The Factory: Build a LibraryTaskDescriptor

Step 2 — The Provider: Implement TornadoLibraryProvider

The `prepare()` Contract

Key SPI Types Reference

Step 3 — Registration

Step 4 — Native Binding: CMake Module

Reference: Data Layout, Types, and Alignment

TornadoVM Array Types

Device Pointer Arithmetic

Existing Providers as Reference Implementations

Unit Test Structure