KernelContext: Explicit GPU Thread Programming Model

KernelContext gives you direct access to the GPU thread programming model from ordinary Java methods. Where the @Parallel annotation lets TornadoVM automatically infer parallelism from loop bounds, KernelContext puts you in full control: you read global and local thread IDs, allocate shared/local memory, issue barrier synchronizations, perform atomic operations, and invoke Tensor Core MMA intrinsics — all from Java source that the TornadoVM JIT transparently compiles to OpenCL C or PTX. Tasks that use KernelContext must be dispatched through a GridScheduler so the runtime knows the thread grid dimensions ahead of compilation.

KernelContext and @Parallel/@Reduce annotations are mutually exclusive programming styles within a single task. A method that uses KernelContext to read thread indices should not also annotate its loops with @Parallel — doing so produces undefined behaviour because the two parallelism models make conflicting assumptions about how the iteration space is decomposed.

Instantiation

A KernelContext object is constructed by your application code and passed as an ordinary argument to any task() method:

KernelContext ctx = new KernelContext();

TaskGraph tg = new TaskGraph("myGraph")
    .transferToDevice(DataTransferMode.EVERY_EXECUTION, a, b, c)
    .task("matmul", MyKernels::matmul, ctx, a, b, c, N)
    .transferToHost(DataTransferMode.EVERY_EXECUTION, c);

On the JVM (outside a GPU kernel), all thread-index fields return 0 and all methods execute their sequential Java fallback bodies. This makes your kernel testable without a GPU.

Thread Index Fields

All fields are public final Integer and are re-bound by the TornadoVM JIT to the appropriate hardware intrinsic when compiling for a GPU device. Their JVM-side value is always 0.

Global Thread IDs

globalIdx

Integer

The unique global thread index in the X dimension.
OpenCL: get_global_id(0) · CUDA: blockIdx.x * blockDim.x + threadIdx.x

globalIdy

Integer

The unique global thread index in the Y dimension.
OpenCL: get_global_id(1) · CUDA: blockIdx.y * blockDim.y + threadIdx.y

globalIdz

Integer

The unique global thread index in the Z dimension.
OpenCL: get_global_id(2) · CUDA: blockIdx.z * blockDim.z + threadIdx.z

Local (Workgroup) Thread IDs

localIdx

Integer

Thread index within the current workgroup (CUDA block), X dimension.
OpenCL: get_local_id(0) · CUDA: threadIdx.x

localIdy

Integer

Thread index within the current workgroup, Y dimension.
OpenCL: get_local_id(1) · CUDA: threadIdx.y

localIdz

Integer

Thread index within the current workgroup, Z dimension.
OpenCL: get_local_id(2) · CUDA: threadIdx.z

Workgroup (Block) IDs

groupIdx

Integer

Index of the current workgroup in the X dimension.
OpenCL: get_group_id(0) · CUDA: blockIdx.x

groupIdy

Integer

Index of the current workgroup in the Y dimension.
OpenCL: get_group_id(1) · CUDA: blockIdx.y

groupIdz

Integer

Index of the current workgroup in the Z dimension.
OpenCL: get_group_id(2) · CUDA: blockIdx.z

Global Grid Sizes

globalGroupSizeX

Integer

Total number of threads across all workgroups in X.
OpenCL: get_global_size(0) · CUDA: gridDim.x * blockDim.x

globalGroupSizeY

Integer

Total number of threads across all workgroups in Y.
OpenCL: get_global_size(1) · CUDA: gridDim.y * blockDim.y

globalGroupSizeZ

Integer

Total number of threads across all workgroups in Z.
OpenCL: get_global_size(2) · CUDA: gridDim.z * blockDim.z

Local Workgroup Sizes

localGroupSizeX

Integer

Number of threads per workgroup in X.
OpenCL: get_local_size(0) · CUDA: blockDim.x

localGroupSizeY

Integer

Number of threads per workgroup in Y.
OpenCL: get_local_size(1) · CUDA: blockDim.y

localGroupSizeZ

Integer

Number of threads per workgroup in Z.
OpenCL: get_local_size(2) · CUDA: blockDim.z

Local Memory Allocation

Local memory (OpenCL terminology) or shared memory (CUDA terminology) lives on-chip and is shared among all threads within a workgroup. It is much faster than global device memory but limited in size (typically 48 KB per block on modern GPUs). Call these methods inside your kernel method to allocate a local array; do not store the returned reference in a field or escape it from the kernel method.

allocateIntLocalArray(int size)

int[]

Allocates a shared-memory int array of size elements.

allocateFloatLocalArray(int size)

float[]

Allocates a shared-memory float array of size elements.

allocateDoubleLocalArray(int size)

double[]

Allocates a shared-memory double array of size elements.

allocateLongLocalArray(int size)

long[]

Allocates a shared-memory long array of size elements.

allocateByteLocalArray(int size)

byte[]

Allocates a shared-memory byte array of size elements.

allocateHalfFloatLocalArray(int size)

HalfFloat[]

Allocates a shared-memory HalfFloat array of size elements. Useful for mixed-precision workloads.

allocateHalf2LocalArray(int size)

Half2[]

Allocates a shared-memory array of packed Half2 pairs. On backends with native packed-half2 support each element maps to a single 32-bit __half2.

Local arrays are allocated at compile time on the GPU. The size argument must be a constant known at JIT-compile time. Passing a runtime-variable size is not supported and will result in a compilation error.

Synchronization

localBarrier()

void

Synchronizes all threads within the current workgroup on local (shared) memory. All local-memory writes issued before the barrier are guaranteed to be visible to all other threads in the workgroup after it.
OpenCL: barrier(CLK_LOCAL_MEM_FENCE) · CUDA: __syncthreads()

globalBarrier()

void

Synchronizes all threads within the current workgroup on global memory. The barrier scope is the workgroup (CUDA block), not the entire device grid.
OpenCL: barrier(CLK_GLOBAL_MEM_FENCE) · CUDA: __syncthreads()

Atomic Operations

Atomic operations are essential for concurrent accumulation and histogram computation. TornadoVM provides atomics for both global (IntArray, LongArray, FloatArray, DoubleArray) and local/shared-memory (int[]) targets.

atomicAdd(IntArray array, int index, int val)

void

Atomically adds val to array[index] in global device memory.
CUDA: atomicAdd(int* address, int val)

atomicAdd(int[] array, int index, int val)

void

Atomically adds val to array[index] in local (shared) memory.

On the CUDA backend, when index is a runtime-computed (non-constant) value, a known issue causes the generated CUDA C to always update array[0] instead of array[index]. For dynamic-index local histograms, use the global IntArray overload instead.

atomicAdd(LongArray array, int index, long val)

void

Atomically adds val to array[index] in global memory (64-bit).
CUDA: atomicAdd(long long* address, long long val)

atomicAdd(FloatArray array, int index, float val)

void

Atomically adds val to array[index] in global device memory (32-bit float).
CUDA: atomicAdd(float* address, float val)

atomicAdd(DoubleArray array, int index, double val)

void

Atomically adds val to array[index] in global device memory (64-bit double).
CUDA: atomicAdd(double* address, double val)

atomicCAS(int[] array, int index, int expected, int value)

int

Compare-and-swap on local/shared memory. If array[index] == expected, atomically stores value and returns the previous value. Success if return == expected.
CUDA: atomicCAS(&array[index], expected, value)

atomicExchange(int[] array, int index, int value)

int

Atomically stores value into array[index] (local memory) and returns the previous value.
CUDA: atomicExch(&array[index], value)

atomicMin(int[] array, int index, int value)

int

Atomically replaces array[index] with min(array[index], value) (local memory) and returns the previous value.

atomicMax(int[] array, int index, int value)

int

Atomically replaces array[index] with max(array[index], value) (local memory) and returns the previous value.

SIMD / Warp Intrinsics

These intrinsics operate at the SIMD-group (warp) level and require all lanes to converge at the call site.

Show SIMD intrinsic methods

Method	Description
`simdSum(float val)`	Returns the sum of `val` across all active lanes. Metal: `simd_sum` · CUDA: butterfly reduction via `shfl.sync.down`.
`simdShuffleDown(float val, int delta)`	Returns the value held by the thread `delta` lanes ahead. Enables suffix-scan patterns.
`simdBroadcastFirst(float val)`	Broadcasts the value from lane 0 to all active lanes. Metal: `simd_broadcast_first`.
`simdAny(boolean predicate)`	Returns `true` if any active lane’s predicate is `true`. CUDA: `__any_sync`.
`simdAll(boolean predicate)`	Returns `true` if all active lanes’ predicates are `true`. CUDA: `__all_sync`.
`simdBallot(boolean predicate)`	Returns a bitmask with one bit set per lane where `predicate` is `true`. CUDA: `__ballot_sync`.

Tensor Core MMA Intrinsics

TornadoVM exposes matrix multiply-accumulate (MMA) primitives for Tensor Core workloads. These operations are warp-collective: all 32 lanes of the warp must converge at the call site.

MMA intrinsics are currently lowered on the CUDA backend only and require compute capability 8.0+ for FP16/BF16, and 8.9+ (Ada Lovelace / Hopper) for FP8. On other backends the methods execute sequential CPU fallbacks.

Show MMA fragment and load/store methods

Method	Description
`mmaFragment(float v)`	Allocates a 4-element `float[]` accumulator fragment initialised to `v`. Maps to register allocation on GPU.
`mmaLoadA(int[] aTile, int wmmaK)`	Loads the A-operand fragment from a shared-memory int-packed tile for `m16n8k16`.
`mmaLoadB(int[] bTile, int wmmaK)`	Loads the B-operand fragment from a shared-memory tile.
`mmaLoadBSwizzled(HalfFloat[] bTile, int wmmaK)`	Loads B from a bank-conflict-free swizzled tile. Includes XOR permutation to avoid shared-memory bank conflicts.
`mma(HalfFloat[] fragA, HalfFloat[] fragB, float[] fragC, MMAShape shape)`	Warp-collective FP16 MMA: `D = A * B + C`. PTX: `mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32`.
`mmaBF16(HalfFloat[] fragA, HalfFloat[] fragB, float[] fragC, MMAShape shape)`	BF16 variant: `D = A * B + C` with bf16 inputs. Compute capability 8.0+ required.
`mmaStore(float[] fragD, FloatArray c, int tileRow, int tileCol, int dimN)`	Stores the D fragment to global memory.
`mmaInt8(byte[] fragA, byte[] fragB, int[] fragC, MMAShape shape)`	INT8 MMA: `D = A * B + C` with int8 inputs, int32 accumulator. PTX: `mma.sync.m16n8k32`.
`mmaFP8E4M3(byte[] fragA, byte[] fragB, float[] fragC, MMAShape shape)`	FP8 E4M3 MMA. Requires compute capability 8.9+.
`mmaFP8E5M2(byte[] fragA, byte[] fragB, float[] fragC, MMAShape shape)`	FP8 E5M2 MMA. Requires compute capability 8.9+.

Matrix-Vector Multiplication Example

The following example demonstrates KernelContext usage for a 2D matrix-vector multiplication kernel with local memory tiling.

Always call localBarrier() after cooperatively loading data into local memory, and again after the compute phase before writing results. Missing a barrier produces race conditions that manifest as non-deterministic output.

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

public class MatVecMul {

    public static void matvec(KernelContext ctx,
                               float[] matrix, float[] vector,
                               float[] result, int N) {
        int row = ctx.globalIdx;
        if (row >= N) return;

        // Allocate a local tile for the vector
        float[] tile = ctx.allocateFloatLocalArray(256);
        int localId = ctx.localIdx;

        // Cooperatively load a tile of the vector into local memory
        if (localId < N) {
            tile[localId] = vector[localId];
        }
        ctx.localBarrier();

        // Each thread computes one row of the result
        float sum = 0.0f;
        for (int col = 0; col < N; col++) {
            sum += matrix[row * N + col] * tile[col];
        }
        result[row] = sum;
    }

    public static void main(String[] args) {
        int N = 256;
        float[] matrix = new float[N * N];
        float[] vector = new float[N];
        float[] result = new float[N];
        // ... fill matrix and vector ...

        KernelContext ctx = new KernelContext();

        WorkerGrid1D grid = new WorkerGrid1D(N);
        grid.setLocalWork(256, 1, 1);
        GridScheduler sched = new GridScheduler("mv.matvec", grid);

        TaskGraph tg = new TaskGraph("mv")
            .transferToDevice(DataTransferMode.EVERY_EXECUTION, matrix, vector)
            .task("matvec", MatVecMul::matvec, ctx, matrix, vector, result, N)
            .transferToHost(DataTransferMode.EVERY_EXECUTION, result);

        try (TornadoExecutionPlan plan = new TornadoExecutionPlan(tg.snapshot())) {
            plan.withGridScheduler(sched).execute();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}

Worker Grid

Configure global and local work dimensions that KernelContext tasks require.

Core API

Data Types

Library Tasks

KernelContext: Explicit GPU Thread Programming Model

Instantiation

Thread Index Fields

Global Thread IDs

Local (Workgroup) Thread IDs

Workgroup (Block) IDs

Global Grid Sizes

Local Workgroup Sizes

Local Memory Allocation

Synchronization

Atomic Operations

SIMD / Warp Intrinsics

Tensor Core MMA Intrinsics

Matrix-Vector Multiplication Example

Worker Grid

Task Graph

Build docs developers (and LLMs) love

Core API

Data Types

Library Tasks

Documentation Index

​Instantiation

​Thread Index Fields

​Global Thread IDs

​Local (Workgroup) Thread IDs

​Workgroup (Block) IDs

​Global Grid Sizes

​Local Workgroup Sizes

​Local Memory Allocation

​Synchronization

​Atomic Operations

​SIMD / Warp Intrinsics

​Tensor Core MMA Intrinsics

​Matrix-Vector Multiplication Example

Worker Grid

Task Graph

Build docs developers (and LLMs) love

Instantiation

Thread Index Fields

Global Thread IDs

Local (Workgroup) Thread IDs

Workgroup (Block) IDs

Global Grid Sizes

Local Workgroup Sizes

Local Memory Allocation

Synchronization

Atomic Operations

SIMD / Warp Intrinsics

Tensor Core MMA Intrinsics

Matrix-Vector Multiplication Example