Programming model

How Enceladus runs a kernel: programs and the launch grid, how a tile spreads across SIMD groups, types, control flow, and compile-time specialization.

Programs and the grid

You launch a kernel by indexing it with a grid and calling the result: kernel[grid](*args, BLOCK=1024, num_warps=4). The grid is a tuple of one to three non-negative integers, or a function that takes a dictionary of the arguments and returns that tuple. The GPU runs one program for each point of the grid, in no particular order.

Each program runs as one Metal threadgroup of num_warps SIMD groups, and each SIMD group has 32 threads. num_warps must be a power of two from 1 to 32, and the default is 4. The compiler spreads every tile across those threads.

Where does each element live?

The quickstart vector add with n = 98_432, BLOCK=1024, and num_warps=4. Each cell is one element of the offs tile, colored by the SIMD group that holds it. Point at a cell, or enter an index, to see the thread and register from the generated MSL.

element 032 per rowelement 1023

Tile element (0 to 1023)

offs
SIMD group
Lane
Register
mask

SIMD group 0

SIMD group 1

SIMD group 2

SIMD group 3

Masked off: offs >= n

Tiles

Every operation in a kernel works on whole tiles. x + y adds two tiles elementwise, tl.sum(x, axis=0) reduces along an axis, and tiles follow NumPy's broadcasting rules. x[:, None] inserts a dimension. Tile shapes have the following rules:

Every dimension must be a compile-time power of two from 1 to 65,536. To handle other lengths, round up with enceladus.next_power_of_2 and mask the extra elements.
A tile must fit in registers. Above 128 32-bit registers per thread, Enceladus warns that the kernel might spill. Above 256, compilation fails with a suggestion to use smaller blocks or a larger num_warps.
Tiles are immutable. To change elements, build another tile with tl.where.

A scalar, such as a kernel argument n or the result of tl.program_id, broadcasts to any tile shape.

Types

Enceladus supports int1, signed and unsigned integers from 8 to 64 bits, float16, bfloat16, and float32. It has no float64, FP8, or TF32. Arithmetic follows Triton's promotion rules:

A Python literal takes the other operand's type if it fits. A float literal otherwise becomes float32.
Mixing an integer and a float gives the float type. Mixing float16 and bfloat16 gives float32.
/ on integers gives float32, and // needs integer operands.
16-bit float operations compute in float32 and round each result back.

Control flow and helpers

A kernel can use if and else on scalars, for loops over range, tl.range, or tl.static_range, and and, or, and not on scalars. On tiles, use &, |, and ~. A condition on constexprs is decided at compile time, and the other branch isn't compiled.

A kernel can call other @enceladus.jit functions, which the compiler inlines. Helpers can return values, including tuples. while, break, continue, early return, try, with, comprehensions, lambdas, and nested functions raise enceladus.CompilationError.

gelu.py

@enceladus.jit
def gelu(x):
return 0.5 * x * (1.0 + tl.erf(x * 0.7071067811865476))


@enceladus.jit
def gelu_kernel(x_ptr, out_ptr, n, BLOCK: tl.constexpr):
offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
mask = offs < n
x = tl.load(x_ptr + offs, mask=mask)
tl.store(out_ptr + offs, gelu(x), mask=mask)

Constexprs and specialization

Enceladus compiles a separate specialization of a kernel for each combination of the following launch properties:

The value of each parameter annotated as tl.constexpr: an integer, float, boolean, tl dtype, or None.
The type of each runtime argument. A Python int is int32 if it fits in 32 bits and int64 otherwise. A Python float is float32.
Whether each integer is divisible by 16 or equals 1, and whether each array's data is 16-byte aligned.
The launch options num_warps, dot_warps, and dot_backend, and the ENCELADUS_DEBUG variable.

Each specialization takes a few milliseconds to compile and is cached in memory and on disk. If a kernel compiles more than 16 times in one process, Enceladus warns and names the argument that changed most often. For sizes that vary from launch to launch, add them to do_not_specialize:

python

@enceladus.jit(do_not_specialize=["n"])
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
# The kernel body is unchanged.
pass

@enceladus.jit also takes interpret=True, which runs every launch in the interpreter, and math_mode="fast", which lets Metal assume that no value is infinite or NaN. The default, "relaxed", keeps infinities and NaNs.