Kernels

The kernel language contract (draft) distinguishes intended guarantees from current limitations and numerical policies still being decided. Consult it when depending on ordering, aliasing, or Python-compatible expression semantics.

Parallel Loops

The outermost for loop in a kernel is the parallel loop — each iteration maps to one GPU thread (or one chunk of work on CPU):

@tack.kernel
def fill(data, val, n):
    for i in range(n):
        data[i] = val

A kernel has exactly one parallel loop, a for statement directly in its body (not inside an if or while). Tack uses it to determine how many threads to launch. A kernel without one, or with a second, is rejected with an UnsupportedSyntaxError that gives the line and column.

Work that has to run in order, such as a recurrence in which each value needs the one before it, goes in a parallel loop of one iteration, with ordinary loops inside:

@tack.kernel
def integrate(rho, e, n, dx):
    for _ in range(1):                 # one thread; the loops below run in order
        e[0] = rho[0] * dx * 0.5
        for i in range(1, n):
            e[i] = e[i - 1] + (rho[i - 1] + rho[i]) * dx * 0.5

Code before or after the parallel loop may run any number of times per launch (once per GPU thread, once per CPU chunk), so it may only assign local variables, read arrays, and declare tack.shared or tack.local_array arrays. Array stores, atomics, barriers, block reductions and print there are rejected; move them into the loop body. Each iteration starts from the values the code before the loop assigned:

@tack.kernel
def scale(x, out, k):
    s = 2.0 * k                    # fine: a local, set up before the loop
    for i in range(x.shape[0]):
        out[i] = x[i] * s

The loop bound can come from: - A scalar argument: range(n) - An array shape: range(data.shape[0]) - len(data)

Constants

A kernel is compiled from its own source and does not see Python variables around it: a module-level dt = 0.01 could change after the kernel was compiled, and the kernel would silently keep the old value. Reading such a name raises NameError. Declare the value with tack.constant instead:

DT = tack.constant(0.01)
GRID = tack.constant(128)

@tack.func
def decay(v):
    return v * (1.0 - DT)

@tack.kernel
def step(x, v):
    for i in range(GRID):
        v[i] = decay(v[i])
        x[i] += v[i] * DT

The kernel reads DT as the literal 0.01. The constant is still an ordinary Python number, so host code uses the same name: range(GRID), np.float32(DT), tack.empty((GRID,)). Constants work from another module too (from params import DT, or params.DT), and inside device functions.

A second argument gives the constant a type, which matters in two places:

MULTIPLIER = tack.constant(747796405, tack.u32)   # x * MULTIPLIER wraps at 32 bits
TENTH = tack.constant(0.1, tack.f64)              # the exact double, not a widened f32

Without it, a u32 value times an integer literal is computed in i64.

The value is fixed where the constant is declared. Arithmetic on constants gives plain numbers, so a derived value is declared again: H = tack.constant(1.0 / (GRID - 2)). Binding the name to a different constant after a kernel has been compiled does not change that kernel. For values that change between calls, pass an argument. math.pi, math.e and math.tau can be written directly.

A tuple of numbers is a vector constant and a tuple of equal tuples a matrix constant; in a kernel they are the vector or matrix with those components, with the same typing rules per component:

SUN = tack.constant((0.5, 0.5, 0.0))
ROTATE = tack.constant(((0.0, -1.0), (1.0, 0.0)))
CELL = tack.constant((2, 3), tack.i32)

@tack.kernel
def pull(pos, grid, out, n):
    for i in range(n):
        r = pos[i] - SUN                 # a vector operand
        out[i] = r.norm() * SUN.y + grid[CELL] + (ROTATE @ [r.x, r.y]).x

On the host a vector constant indexes and iterates like a tuple of its rows and converts with np.array(SUN, np.float32); it is not a tuple, so SUN + OFFSET is an error rather than a concatenation (use np.asarray(SUN) for host arithmetic).

Shader-style helpers

tack.math holds the small vocabulary GLSL and Taichi programs lean on, as device functions that work on scalars and vectors alike: fract, mix, clamp, saturate, smoothstep, step, sign, length, distance, normalize.

from tack import math as tm

@tack.kernel
def grid_lines(uv, out, n):
    for i in range(n):
        st = tm.fract(uv[i] * 8.0) - 0.5                  # a 2-vector
        edge = tm.smoothstep(0.0, 0.05, 0.5 - abs(st).min())
        out[i] = tm.mix(0.2, 1.0, edge)

The definitions are GLSL's (mix(x, y, a) is x * (1 - a) + y * a, and smoothstep with its edges reversed reverses the step). They are only callable from kernels and device functions.

Array Dimensions

array.shape[k] and len(array) work anywhere in a kernel — not only as the loop bound:

@tack.kernel
def reverse(x, out):
    for i in range(x.shape[0]):
        out[i] = x[x.shape[0] - 1 - i]

@tack.kernel
def normalize_rows(a, out):
    for i, j in tack.ndindex(a.shape[0], a.shape[1]):
        out[i, j] = a[i, j] / float(a.shape[1])

The dimension index must be a literal: x.shape[0] is fine, x.shape[d] with a runtime d is not.

One thing to know about compilation. A dimension used in the loop bound is passed to the launch, so one compiled kernel serves every array length. A dimension used inside the kernel becomes part of the generated code, so Tack compiles a separate version per shape. That is usually what you want — it lets the compiler fold and vectorize around a known size. If you would rather not specialize, pass the length as a scalar argument:

@tack.kernel
def reverse(x, out, n):     # one compiled kernel for every n
    for i in range(n):
        out[i] = x[n - 1 - i]

Multi-Dimensional Iteration

Use tack.ndindex for 2D or 3D parallel iteration:

@tack.kernel
def fill_2d(grid, height, width):
    for row, column in tack.ndindex(height, width):
        grid[row, column] = float(row + column)

This launches width * height threads. Each thread gets its (i, j) pair via index decomposition.

An argument can also be a (start, end) pair, which is how a stencil visits the interior of a grid and leaves the boundary alone:

@tack.kernel
def smooth(u, out, n, m):
    for i, j in tack.ndindex((1, n - 1), (1, m - 1)):
        out[i, j] = 0.25 * (u[i - 1, j] + u[i + 1, j] + u[i, j - 1] + u[i, j + 1])

Sizes and pairs can be mixed. An empty or reversed pair makes the loop run no iterations.

Sequential Inner Loops

Loops nested inside the parallel loop run sequentially per thread:

@tack.kernel
def reduction_per_row(data, out, width, height):
    for row in range(height):          # parallel (one thread per row)
        total = 0.0
        for col in range(width):       # sequential (per thread)
            total = total + data[row * width + col]
        out[row] = total

Sequential loops support: - range(n) with runtime bounds (including array loads) - range(start, end) - range(start, end, step)

Variable-Length Inner Loops

The loop bound can come from an array load, enabling variable-length iteration patterns like Viskores-style connectivity:

@tack.kernel
def sum_segments(offsets, data, output, n_cells):
    for c in range(n_cells):
        total = 0.0
        for i in range(offsets[c], offsets[c + 1]):
            total = total + data[i]
        output[c] = total

Returning a Value

A kernel whose last statement is return expr, after the parallel loop, hands that value back to the caller. The return annotation names the type: -> tack.f32, -> float (f32) or -> int (i32).

@tack.kernel
def count_live(cells, count, n) -> int:
    for i in range(n):
        if cells[i] > 0:
            tack.atomic_add(count, 0, 1)
    return count[0]

count.fill(0)                            # reset before counting each board
live = count_live(cells, count, n)        # a Python int

The expression is evaluated once, after the launch, by a one-thread epilogue kernel that stores it into a hidden array the call reads back. So it may use arrays, arguments, constants and locals bound before the loop, but not a local the loop assigns, which has a value per thread. This replaces writing a one-element array and reading it with numpy() or f[0]; the cost is one extra, tiny launch and one element's transfer.

Return the length of an atomic append

An atomic's result and a kernel's result have different roles. Inside the loop, atomic_add returns the counter's old value, giving that iteration a unique output slot. After the loop, the kernel returns the completed counter to Python:

@tack.kernel
def compact_positive(values, out, count, n) -> int:
    for i in range(n):
        if values[i] > 0.0:
            slot = tack.atomic_add(count, 0, 1)
            out[slot] = values[i]
    return count[0]

Allocate enough capacity and reset the counter before each call:

    values = tack.array(np.array([-2, 1, -0.5, 3, 0, 2], dtype=np.float32))
    out = tack.empty(values.shape)  # capacity for every input
    count = tack.empty((1,), tack.i32)
    count.fill(0)
    length = compact_positive(values, out, count, values.size)
    result = out.numpy()[:length]
    # Arrival order is unspecified: compare the values after sorting.
    assert length == 3 and isinstance(length, int)
    np.testing.assert_array_equal(np.sort(result), np.array([1, 2, 3], dtype=np.float32))

The returned length is deterministic; the order of the appended values is not. Only the first length output elements are valid. A scalar return does not turn the loop's per-iteration locals into a global reduction. For ordered output, use count → scan → write.

Download the complete program, or run:

uv run python docs/examples/kernel_results.py --arch cpu --check
uv run python docs/examples/kernel_results.py --arch metal --check

Kernel Caching

Kernels are compiled on first call. A compiled variant is reused by later calls that agree on everything the generated code depends on: argument dtypes, which arguments are arrays, scalars or textures, vector widths, texture extents, @tack.device_class class constants and layout, any array dimensions the kernel bakes in (for example x.shape[1] used as a row stride), and on CPU whether the arrays overlap. Anything else is a runtime parameter.

So changing scalar values (passing alpha=2.5, then alpha=3.0) does not recompile, and neither does changing the length of a 1-D array used only as the parallel loop bound. Changing a dtype, a vector width or a baked-in dimension does. If a kernel recompiles for every array size, pass the size as a scalar argument instead of reading it from shape inside the body. Specialization and Caching lists exactly what goes into the key, with worked examples.

Compiled variants are also kept on disk, so a later process that calls the same kernel the same way starts warm; editing the kernel, a device function it calls, a constant or a template class compiles it again. See The Kernel Cache.

Inspecting Generated Code

tack.inspect() shows the generated code for a kernel without executing it. This is useful for debugging, learning, and understanding what the compiler produces:

@tack.kernel
def saxpy(x, y, out, a):
    for i in range(len(x)):
        out[i] = a * x[i] + y[i]

n = 1024
x = tack.empty((n,))
y = tack.empty((n,))
out = tack.empty((n,))

# Tack intermediate representation
print(tack.inspect(saxpy, x, y, out, 2.0, mode="ir"))

# Backend-specific source (MSL, CUDA C, LLVM IR, etc.)
print(tack.inspect(saxpy, x, y, out, 2.0, mode="source"))

# Post-optimization LLVM IR (CPU backend only — runs LLVM O3 passes)
print(tack.inspect(saxpy, x, y, out, 2.0, mode="optimized"))

The three modes are:

Mode	Output
"ir"	Tack intermediate representation
"source"	Backend source code: LLVM IR (CPU), MSL (Metal), CUDA C, HIP C, OpenCL C
"optimized"	Post-optimization LLVM IR (loop vectorization, unrolling, etc.). CPU only; other backends raise ValueError

You must pass the same arguments the kernel would receive at runtime, since type inference, dimension resolution, and template expansion all depend on them. Inspection applies the same checks as a call, so arguments a call would reject (an f64 array on Metal, for example) raise here too, and on GPUs without texture hardware the source shows the software sampling the call would use. Templates and scalar arguments work as expected:

@tack.device_class
class Sim:
    GRAVITY = -9.81

    def __init__(self, n):
        self.pos = tack.empty((n,))
        self.vel = tack.empty((n,))
        self.dt = 0.01

    @tack.func
    def step(self, i):
        self.vel[i] = self.vel[i] + self.GRAVITY * self.dt
        self.pos[i] = self.pos[i] + self.vel[i] * self.dt

@tack.kernel
def update(sim):
    for i in range(len(sim.pos)):
        sim.step(i)

sim = Sim(1024)
print(tack.inspect(update, sim))  # shows template expansion + inlined methods

For complete applications that combine these constructs, see Heat diffusion, Physarum and the tutorial gallery.