Kernel API: Explicit GPU Thread Programming in Java

TornadoVM’s Kernel API gives you the same degree of control over GPU execution that you would have writing OpenCL or CUDA kernels directly, but expressed entirely in Java. Instead of relying on TornadoVM to infer parallelism from @Parallel loop annotations, you query thread identifiers, manage shared memory, and insert barriers manually through a KernelContext object. This API is designed for expert GPU programmers who need fine-grained control — for example, to port existing CUDA or OpenCL kernels, implement loop-tiled matrix operations with shared memory, or exploit Tensor Core MMA intrinsics. For workloads where you don’t need that level of control, the annotation-based Loop Parallel API is simpler and equally performant.

When to Use the Kernel API

Use KernelContext when…

You need explicit access to localIdx, localIdy for shared-memory tiling
You are porting an existing OpenCL or CUDA kernel
You want to use localBarrier() to synchronise a work-group
You need local (shared) memory allocation
You want to use Tensor Core MMA intrinsics (CUDA only)

Use @Parallel when…

Your kernel is a simple loop over array elements
You want the compiler to choose the optimal thread distribution
You don’t need inter-thread communication or barriers
You’re new to GPU programming

KernelContext Thread-ID Fields

A KernelContext object is passed as the first parameter of any kernel method that uses the Kernel API. TornadoVM’s JIT compiler replaces all field reads at the GPU level with the corresponding hardware thread-ID intrinsics.

`KernelContext` field	OpenCL equivalent	CUDA equivalent
`kc.globalIdx`	`get_global_id(0)`	`blockIdx.x * blockDim.x + threadIdx.x`
`kc.globalIdy`	`get_global_id(1)`	`blockIdx.y * blockDim.y + threadIdx.y`
`kc.globalIdz`	`get_global_id(2)`	`blockIdx.z * blockDim.z + threadIdx.z`
`kc.localIdx`	`get_local_id(0)`	`threadIdx.x`
`kc.localIdy`	`get_local_id(1)`	`threadIdx.y`
`kc.localIdz`	`get_local_id(2)`	`threadIdx.z`
`kc.groupIdx`	`get_group_id(0)`	`blockIdx.x`
`kc.groupIdy`	`get_group_id(1)`	`blockIdx.y`
`kc.globalGroupSizeX`	`get_global_size(0)`	`gridDim.x * blockDim.x`
`kc.localGroupSizeX`	`get_local_size(0)`	`blockDim.x`

Local Memory Allocation

Local memory (shared memory in CUDA) is allocated through KernelContext factory methods. The allocation is per-work-group (per-block) and is not initialised — you must write before reading.

// Allocate typed local arrays inside the kernel method
int[]    intBuf    = kc.allocateIntLocalArray(size);
float[]  floatBuf  = kc.allocateFloatLocalArray(size);
double[] doubleBuf = kc.allocateDoubleLocalArray(size);
long[]   longBuf   = kc.allocateLongLocalArray(size);
byte[]   byteBuf   = kc.allocateByteLocalArray(size);

Local memory allocations are per work-group, not per thread. All threads in the same work-group share the same local array. The size parameter sets the number of elements, not bytes.

Barrier Synchronisation

Barriers synchronise all threads within a work-group (CUDA block). They are not grid-wide barriers.

// Synchronise local (shared) memory writes
kc.localBarrier();

// Synchronise global memory writes (also work-group scoped)
kc.globalBarrier();

Both localBarrier() and globalBarrier() are work-group barriers, not grid-wide fences. Do not use them expecting cross-block synchronisation — that requires splitting the computation across multiple kernel launches (separate TaskGraph tasks).

WorkerGrid: Configuring Thread Dimensions

WorkerGrid objects define the global and local thread dimensions for a kernel. TornadoVM provides 1D, 2D, and 3D variants.

WorkerGrid1D
WorkerGrid2D
WorkerGrid3D

// 1024 global threads in the X dimension
WorkerGrid1D grid = new WorkerGrid1D(1024);

// Optionally override local work size (block size)
grid.setLocalWork(64, 1, 1);

// 512 × 512 global threads
WorkerGrid2D grid = new WorkerGrid2D(512, 512);

// 16×16 local work-group (block)
grid.setLocalWork(16, 16, 1);

import uk.ac.manchester.tornado.api.WorkerGrid3D;

// 64 × 64 × 64 global threads
WorkerGrid3D grid = new WorkerGrid3D(64, 64, 64);
grid.setLocalWork(8, 8, 8);

GridScheduler: Mapping Tasks to Grids

GridScheduler associates task IDs (in "graphName.taskName" format) with their WorkerGrid.

// Single task
GridScheduler scheduler = new GridScheduler("s0.t0", workerGrid);

// Multiple tasks — use addWorkerGrid()
GridScheduler scheduler = new GridScheduler();
scheduler.addWorkerGrid("s0.t0", grid1D);
scheduler.addWorkerGrid("s0.t1", grid2D);

The GridScheduler is then attached to the execution plan via .withGridScheduler(...):

try (TornadoExecutionPlan plan = new TornadoExecutionPlan(itg)) {
    plan.withGridScheduler(scheduler).execute();
}

Any task that references a KernelContext must have a corresponding entry in a GridScheduler. Running a KernelContext task without a grid attachment causes a runtime error.

Complete Example: Tiled Matrix Multiplication

The following example combines KernelContext, local memory, barriers, and WorkerGrid2D to implement a shared-memory tiled matrix multiplication — the classic GPU optimisation that avoids redundant global memory reads.

Define the kernel method

The KernelContext is the first parameter. Local memory is allocated inside the method body.

import uk.ac.manchester.tornado.api.KernelContext;
import uk.ac.manchester.tornado.api.types.arrays.FloatArray;

public class TiledMxM {

    private static final int TS = 32; // tile size = local work-group size

    public static void matMulTiled(KernelContext kc,
                                   FloatArray A,
                                   FloatArray B,
                                   FloatArray C,
                                   int size) {
        // Thread indices within the work-group
        int row = kc.localIdx;
        int col = kc.localIdy;

        // Global position of this thread
        int globalRow = TS * kc.groupIdx + row;
        int globalCol = TS * kc.groupIdy + col;

        // Allocate shared memory tiles
        float[] aSub = kc.allocateFloatLocalArray(TS * TS);
        float[] bSub = kc.allocateFloatLocalArray(TS * TS);

        float sum = 0.0f;

        // Iterate over all tiles in the K dimension
        int numTiles = size / TS;
        for (int t = 0; t < numTiles; t++) {

            // Load one tile of A and B into shared memory
            int tiledRow = TS * t + row;
            int tiledCol = TS * t + col;
            aSub[col * TS + row] = A.get(tiledCol * size + globalRow);
            bSub[col * TS + row] = B.get(globalCol * size + tiledRow);

            // Wait until all threads have filled the tile
            kc.localBarrier();

            // Multiply within the tile
            for (int k = 0; k < TS; k++) {
                sum += aSub[k * TS + row] * bSub[col * TS + k];
            }

            // Wait before loading the next tile
            kc.localBarrier();
        }

        C.set(globalCol * size + globalRow, sum);
    }
}

Set up the WorkerGrid and GridScheduler

// Grid matches the matrix dimensions; local work-group = tile size
WorkerGrid2D workerGrid = new WorkerGrid2D(size, size);
workerGrid.setLocalWork(TiledMxM.TS, TiledMxM.TS, 1);

GridScheduler scheduler = new GridScheduler("mxm.t0", workerGrid);
KernelContext context    = new KernelContext();

Build the TaskGraph

import uk.ac.manchester.tornado.api.*;
import uk.ac.manchester.tornado.api.enums.DataTransferMode;

TaskGraph tg = new TaskGraph("mxm")
    .transferToDevice(DataTransferMode.FIRST_EXECUTION, matA, matB)
    .task("t0", TiledMxM::matMulTiled, context, matA, matB, matC, size)
    .transferToHost(DataTransferMode.EVERY_EXECUTION, matC);

ImmutableTaskGraph itg = tg.snapshot();

Execute with the GridScheduler

try (TornadoExecutionPlan plan = new TornadoExecutionPlan(itg)) {
    plan.withGridScheduler(scheduler).execute();
}

Combining KernelContext and @Parallel Tasks

A single TaskGraph can mix Kernel API tasks and Loop Parallel API tasks. Tasks that do not use KernelContext do not need a WorkerGrid entry — TornadoVM schedules them automatically.

WorkerGrid1D grid = new WorkerGrid1D(size);
GridScheduler scheduler = new GridScheduler();
scheduler.addWorkerGrid("s0.t0", grid);   // t0 uses KernelContext
scheduler.addWorkerGrid("s0.t1", grid);   // t1 uses KernelContext
// t2 uses @Parallel — no entry needed in the scheduler

KernelContext kc = new KernelContext();

TaskGraph tg = new TaskGraph("s0")
    .transferToDevice(DataTransferMode.EVERY_EXECUTION, a, b)
    .task("t0", MyKernels::vectorAddKernel,  kc, a, b, c)     // KernelContext
    .task("t1", MyKernels::vectorMulKernel,  kc, c, b, c)     // KernelContext
    .task("t2", MyKernels::vectorSubAnnot,   c, b)            // @Parallel
    .transferToHost(DataTransferMode.EVERY_EXECUTION, c);

try (TornadoExecutionPlan plan = new TornadoExecutionPlan(tg.snapshot())) {
    plan.withGridScheduler(scheduler).execute();
}

Backend Support Matrix

Not every KernelContext feature is available on all backends. Unsupported operations are rejected at compile time — they are never silently replaced with a Java fallback.

Advanced operation support by backend

Operation	CUDA	OpenCL	Metal
`atomicAdd` (int / long arrays)	✅	✅	✅
`atomicAdd` (float / double arrays)	✅	❌	✅
`atomicCAS`, `atomicMin`, `atomicMax`	local only	❌	❌
`simdSum`, `simdShuffleDown`	✅	❌	✅
`allocateHalf2LocalArray`	✅	❌	❌
MMA Tensor Core (`mma*`)	✅	❌	❌
`asyncCopyToLocal` (Ampere+)	✅	❌	❌
`simdgroupMatrix*`, `matrixMultiply8x8`	❌	❌	✅

Tensor Core MMA Intrinsics (CUDA Only)

TornadoVM exposes NVIDIA Tensor Core operations through KernelContext for FP16, BF16, INT8, and FP8 matrix multiply-accumulate workflows. These are low-level intrinsics intended for advanced CUDA programmers implementing custom mixed-precision linear-algebra kernels.

// Load fragments: shared-memory tiles are int[] with fp16 elements packed 2-per-int
// mmaLoadA(int[] aTile, int wmmaK) -> HalfFloat[]
// mmaLoadB(int[] bTile, int wmmaK) -> HalfFloat[]
// mmaFragment(float initVal)       -> float[]       (initialise accumulator)
// mma(HalfFloat[] fragA, HalfFloat[] fragB, float[] fragC, MMAShape) -> float[]
// mmaStore(float[] fragD, FloatArray c, int tileRow, int tileCol, int dimN)

HalfFloat[] fragA = kc.mmaLoadA(aTile, 16);
HalfFloat[] fragB = kc.mmaLoadB(bTile, 16);
float[]     fragC = kc.mmaFragment(0.0f);
float[]     fragD = kc.mma(fragA, fragB, fragC, MMAShape.M16N8K16);
kc.mmaStore(fragD, outputArray, tileRow, tileCol, dimN);

MMA intrinsics require the CUDA backend and a GPU with Tensor Core support (Volta architecture or newer). They are rejected at compile time on OpenCL and Metal backends.

For most use cases, the @Parallel Loop API produces competitive performance without the complexity of explicit thread management. Reach for the Kernel API when you have profiling evidence that shared-memory tiling or warp-level primitives would yield meaningful speedups.

Get Started

Programming Guide

Backends

NVIDIA Integration

Guides

Contributing

Kernel API: Explicit GPU Thread Programming in Java

When to Use the Kernel API

Use KernelContext when…

Use @Parallel when…

KernelContext Thread-ID Fields

Local Memory Allocation

Barrier Synchronisation

WorkerGrid: Configuring Thread Dimensions

GridScheduler: Mapping Tasks to Grids

Complete Example: Tiled Matrix Multiplication

Combining KernelContext and @Parallel Tasks

Backend Support Matrix

Tensor Core MMA Intrinsics (CUDA Only)

Build docs developers (and LLMs) love

Get Started

Programming Guide

Backends

NVIDIA Integration

Guides

Contributing

Documentation Index

​When to Use the Kernel API

Use KernelContext when…

Use @Parallel when…

​KernelContext Thread-ID Fields

​Local Memory Allocation

​Barrier Synchronisation

​WorkerGrid: Configuring Thread Dimensions

​GridScheduler: Mapping Tasks to Grids

​Complete Example: Tiled Matrix Multiplication

​Combining KernelContext and @Parallel Tasks

​Backend Support Matrix

​Tensor Core MMA Intrinsics (CUDA Only)

Build docs developers (and LLMs) love

When to Use the Kernel API

KernelContext Thread-ID Fields

Local Memory Allocation

Barrier Synchronisation

WorkerGrid: Configuring Thread Dimensions

GridScheduler: Mapping Tasks to Grids

Complete Example: Tiled Matrix Multiplication

Combining KernelContext and @Parallel Tasks

Backend Support Matrix

Tensor Core MMA Intrinsics (CUDA Only)