gizml:tile

GPU kernels as ordinary Gizlang over tiles, on every device.

A tile program says what one block of the output is; this module decides everything else: threads, which thread holds which element, workgroup memory, barriers, bounds masks, matrix-core fragments, layout conversions and the device. The same program runs on host tensors (the CPU, one program per grid point in turn) and on a GPU's tensors (NVIDIA PTX, AMD LLVM through comgr, both from gizhw:kernel). A tile program is ordinary Gizlang, and it reads as math: a * b + c is a.mul(b).add(c).

Tile.run([Tile.cdiv(m, 128), Tile.cdiv(n, 128)], (p: Tile.Program): void => { const i = p.id(0); const j = p.id(1); const A = p.read(a, [128, 32]); const B = p.read(b, [32, 128]); const acc = p.fold(p.int(0), A.blocks(1), p.zeros([128, 128]), (k: Tile.Index, acc: Tile.Tile): Tile.Tile => acc.add(A.at([i, k]).dot(B.at([k, j])))); p.write(c, [128, 128]).put(acc.relu(), [i, j]); });

The body runs once per launch and records the tile IR (below). Tiles are immutable values of a static shape whose sizes are powers of two; a 1-d tile of C elements is the row [1, C]. Method names are the operators' (add, sub, mul, div, lt, ...) and k lifts a literal, so operator syntax can lower onto them.

import * as Tile from "gizml:tile";

Types

Program

type Program

One grid point's program while it is traced.

Opaque: create and use it through this module's functions.

Tile

type Tile

A block of values: shape is [rows, columns], dtype the storage type its values came from (F32 after math).

• shape: Vec<i64>
• dtype: i64

Index

type Index

A uniform integer: a program id, a size, a loop counter, or math on them.

Opaque: create and use it through this module's functions.

In

type In

A tensor read in blocks.

Opaque: create and use it through this module's functions.

Out

type Out

A tensor written in blocks.

Opaque: create and use it through this module's functions.

Functions

cdiv

function cdiv(a: i64, b: i64): i64

ceil(a / b) for a launch's grid.

Effects: fail.

run

function run(grid: Vec<i64>, body: (p: Program) => void): Vec<Tensor>

Runs body once for each point of grid (1 to 3 sizes, axis 0 fastest) where its tensors are: on their GPU, or on the CPU for host tensors. Returns the tensors the body wrote, in the order of its write calls.

Effects: fail, read, write, console, io.

tune

function tune<C>(configs: Vec<C>, grid: (c: C) => Vec<i64>, body: (p: Program, c: C) => void): Vec<Tensor>

As run, for each configuration in configs (block sizes, warps: any record), keeping the fastest on this GPU for these tensor shapes (measured once, then cached); on the CPU the first. Every configuration must compute the same result, and the body must not read what it writes (each runs several times while measured).

ir

function ir(grid: Vec<i64>, storage: Vec<i64>, kind: string, arch: string, warps: i64, body: (p: Program) => void): string

The tile IR of body as text, one node a line: its index, kind, shape, storage type, operands and the layout the lowering for kind and arch gives it (free: made where used; uniform; blocked; mma WMxWN warps; slice of a layout along an axis). What to read first when a tile program is slower than expected: every "slice"/"blocked" next to an "mma" is data moved through workgroup memory.

Effects: fail, read, write, console, io.

source

function source(grid: Vec<i64>, storage: Vec<i64>, kind: string, arch: string, warps: i64, body: (p: Program) => void): string

The GPU code a launch of body runs on a GPU of kind ("cuda": PTX, "hip": AMDGPU LLVM, "metal": MSL) and arch ("sm_120", "gfx1151", "apple9"), with no GPU needed: to read what the lowering made. storage replaces the tensors' storage types, one per read or write in order (empty: theirs), so host tensors can stand for BF16 ones. warps 0 lets it choose.

Effects: fail, read, write, console, io.

Methods

Program.id

function Program.id(self: Program, axis: i64): Index

This program's position along grid axis axis (0 is the fastest).

Effects: fail, read, write, console.

Program.count

function Program.count(self: Program, axis: i64): Index

The grid's size along axis.

Effects: fail, read, write, console.

Program.int

function Program.int(self: Program, value: i64): Index

An integer constant as an index.

Effects: fail, write.

Program.k

function Program.k(self: Program, value: f64): Tile

A constant as a [1, 1] tile, which broadcasts.

Effects: fail, write, console.

Program.full

function Program.full(self: Program, shape: Vec<i64>, value: f64): Tile

A tile of SHAPE with every value VALUE.

Effects: fail, write, console.

Program.zeros

function Program.zeros(self: Program, shape: Vec<i64>): Tile

A tile of SHAPE with every value 0.

Effects: fail, write, console.

Program.iota

function Program.iota(self: Program, shape: Vec<i64>, axis: i64): Tile

Each element's row number (axis 0) or column number (axis 1) within the tile, as F32.

Effects: fail, write, console.

Program.warps

function Program.warps(self: Program, count: i64): void

The number of warps (32 threads each) of a GPU workgroup: 1, 2, 4, 8 or 16. Without it the lowering picks.

Effects: fail, write, console.

Program.stages

function Program.stages(self: Program, count: i64): void

How many blocks of a loop's matrix-core operands are in flight: 2 (the default) reads the next turn's blocks while this turn's product runs, more hide longer memory latency at the cost of workgroup memory, 1 reads each turn's in turn.

Effects: fail, write, console.

Program.read

function Program.read(self: Program, t: Tensor, block: Vec<i64>): In

t read in blocks of block (its last 1 or 2 dimensions); the dimensions before are addressed by element.

Effects: fail, write, console.

Program.write

function Program.write(self: Program, t: Tensor, block: Vec<i64>): Out

t written in blocks of block: Tile.run returns it with what the programs stored; the rest keeps its values.

Effects: fail, write, console.

Program.loop

function Program.loop(self: Program, from: Index, to: Index, inits: Vec<Tile>, body: (k: Index, carried: Vec<Tile>) => Vec<Tile>): Vec<Tile>

body(k, carried) for k from from to to - 1, each time with the values the turn before returned (inits the first time); the values of the last turn (or inits when there is none). The carried tiles keep their shapes.

Effects: fail, read, write, console, io.

Program.fold

function Program.fold(self: Program, from: Index, to: Index, init: Tile, body: (k: Index, acc: Tile) => Tile): Tile

loop with one carried tile.

Effects: fail, read, write, console, io.

Index.k

function Index.k(self: Index, value: i64): Index

A constant in this index's program.

Effects: fail, write.

Index.add

function Index.add(self: Index, other: Index): Index

SELF + OTHER.

Effects: fail, write, console.

Index.sub

function Index.sub(self: Index, other: Index): Index

SELF - OTHER.

Effects: fail, write, console.

Index.mul

function Index.mul(self: Index, other: Index): Index

SELF * OTHER.

Effects: fail, write, console.

Index.div

function Index.div(self: Index, other: Index): Index

Division and remainder truncate toward zero, as i64's.

Effects: fail, write, console.

Index.rem

function Index.rem(self: Index, other: Index): Index

The remainder of SELF / OTHER, with the sign of SELF.

Effects: fail, write, console.

Index.minimum

function Index.minimum(self: Index, other: Index): Index

The smaller of SELF and OTHER.

Effects: fail, write, console.

Index.maximum

function Index.maximum(self: Index, other: Index): Index

The larger of SELF and OTHER.

Effects: fail, write, console.

Index.value

function Index.value(self: Index): Tile

The index as a [1, 1] F32 tile (exact below 2^24).

Effects: fail, write.

In.at

function In.at(self: In, coords: Vec<Index>): Tile

The block at block coordinates coords (one per tensor dimension; leading unblocked dimensions by element). Elements past the tensor's edge read 0.

Effects: fail, read, write, console.

In.window

function In.window(self: In, offsets: Vec<Index>): Tile

The block whose first element is at element offsets offsets (any alignment; outside the tensor reads 0).

Effects: fail, read, write, console.

In.dim

function In.dim(self: In, axis: i64): Index

The tensor's size along axis.

Effects: fail, read, write, console.

Out.dim

function Out.dim(self: Out, axis: i64): Index

The output's size along AXIS.

Effects: fail, read, write, console.

In.blocks

function In.blocks(self: In, axis: i64): Index

The number of blocks along axis (its size for a leading dimension).

Effects: fail, read, write, console.

Out.blocks

function Out.blocks(self: Out, axis: i64): Index

The number of blocks of the output along AXIS (its size for a leading dimension).

Effects: fail, read, write, console.

Out.put

function Out.put(self: Out, t: Tile, coords: Vec<Index>): void

Stores t (the block's shape) as the block at coords; elements past the tensor's edge are not written.

Effects: fail, read, write, console.

Out.putWindow

function Out.putWindow(self: Out, t: Tile, offsets: Vec<Index>): void

Stores t with its first element at element offsets offsets.

Effects: fail, read, write, console.

Tile.k

function Tile.k(self: Tile, value: f64): Tile

A constant [1, 1] tile in this tile's program.

Effects: fail, write.

Tile.add

function Tile.add(self: Tile, other: Tile): Tile

SELF + OTHER, element by element, with broadcasting.

Effects: fail, write, console.

Tile.sub

function Tile.sub(self: Tile, other: Tile): Tile

SELF - OTHER, element by element, with broadcasting.

Effects: fail, write, console.

Tile.mul

function Tile.mul(self: Tile, other: Tile): Tile

SELF * OTHER, element by element, with broadcasting.

Effects: fail, write, console.

Tile.div

function Tile.div(self: Tile, other: Tile): Tile

SELF / OTHER, element by element, with broadcasting.

Effects: fail, write, console.

Tile.maximum

function Tile.maximum(self: Tile, other: Tile): Tile

Elementwise larger and smaller value; with a NaN, the other value.

Effects: fail, write, console.

Tile.minimum

function Tile.minimum(self: Tile, other: Tile): Tile

The smaller of SELF and OTHER, element by element; with a NaN, the other value.

Effects: fail, write, console.

Tile.lt

function Tile.lt(self: Tile, other: Tile): Tile

Comparisons give masks; a NaN compares false (ne: true).

Effects: fail, write, console.

Tile.le

function Tile.le(self: Tile, other: Tile): Tile

The mask of SELF <= OTHER; a NaN compares false.

Effects: fail, write, console.

Tile.gt

function Tile.gt(self: Tile, other: Tile): Tile

The mask of SELF > OTHER; a NaN compares false.

Effects: fail, write, console.

Tile.ge

function Tile.ge(self: Tile, other: Tile): Tile

The mask of SELF >= OTHER; a NaN compares false.

Effects: fail, write, console.

Tile.eq

function Tile.eq(self: Tile, other: Tile): Tile

The mask of SELF == OTHER; a NaN compares false.

Effects: fail, write, console.

Tile.ne

function Tile.ne(self: Tile, other: Tile): Tile

The mask of SELF != OTHER; a NaN compares true.

Effects: fail, write, console.

Tile.and

function Tile.and(self: Tile, other: Tile): Tile

The mask that is true where both masks are true.

Effects: fail, write, console.

Tile.or

function Tile.or(self: Tile, other: Tile): Tile

The mask that is true where either mask is true.

Effects: fail, write, console.

Tile.not

function Tile.not(self: Tile): Tile

The mask that is true where SELF is false.

Effects: fail, write, console.

Tile.select

function Tile.select(self: Tile, yes: Tile, no: Tile): Tile

For a mask: yes where it is true, else no (all three broadcast).

Effects: fail, write, console.

Tile.neg

function Tile.neg(self: Tile): Tile

-SELF, element by element.

Effects: fail, write, console.

Tile.abs

function Tile.abs(self: Tile): Tile

The absolute value of SELF, element by element.

Effects: fail, write, console.

Tile.exp

function Tile.exp(self: Tile): Tile

e^x, ln x, tanh, sigmoid and GELU within 2-3 ulp: gizml:tensor's polynomials on every device.

Effects: fail, write, console.

Tile.log

function Tile.log(self: Tile): Tile

The natural logarithm of SELF, element by element, within 2-3 ulp.

Effects: fail, write, console.

Tile.sqrt

function Tile.sqrt(self: Tile): Tile

The square root of SELF, element by element.

Effects: fail, write, console.

Tile.rsqrt

function Tile.rsqrt(self: Tile): Tile

1 / sqrt(x), as a division of the rounded root.

Effects: fail, write, console.

Tile.sigmoid

function Tile.sigmoid(self: Tile): Tile

The sigmoid 1 / (1 + e^-x) of SELF, element by element, within 2-3 ulp.

Effects: fail, write, console.

Tile.tanh

function Tile.tanh(self: Tile): Tile

The hyperbolic tangent of SELF, element by element, within 2-3 ulp.

Effects: fail, write, console.

Tile.gelu

function Tile.gelu(self: Tile): Tile

GELU, the tanh form (gelu) and the erf form (geluErf).

Effects: fail, write, console.

Tile.geluErf

function Tile.geluErf(self: Tile): Tile

GELU of SELF in the erf form, element by element.

Effects: fail, write, console.

Tile.relu

function Tile.relu(self: Tile): Tile

max(SELF, 0), element by element.

Effects: fail, write, console.

Tile.floor

function Tile.floor(self: Tile): Tile

SELF rounded down to an integer, element by element.

Effects: fail, write, console.

Tile.to

function Tile.to(self: Tile, dtype: i64): Tile

The values rounded to storage type dtype (F32, F16, BF16 or FP8), as a store and a load of it give them.

Effects: fail, write, console.

Tile.sum

function Tile.sum(self: Tile, axis: i64): Tile

Sum, largest and smallest along axis, keeping it as size 1: [R, C].sum(1) is [R, 1].

Effects: fail, write, console.

Tile.max

function Tile.max(self: Tile, axis: i64): Tile

The largest value along AXIS, which stays as size 1.

Effects: fail, write, console.

Tile.min

function Tile.min(self: Tile, axis: i64): Tile

The smallest value along AXIS, which stays as size 1.

Effects: fail, write, console.

Tile.dot

function Tile.dot(self: Tile, other: Tile): Tile

The matrix product [M, K] [K, N] -> [M, N] in F32. With F16 or BF16 operands of one type it runs on matrix cores where the GPU has them; acc.add(a.dot(b)) accumulates into acc there.

Effects: fail, write, console.

Tile.transpose

function Tile.transpose(self: Tile): Tile

Rows become columns.

Effects: fail, write.

Constants

F32

const F32 = 0

The storage type of 32-bit floats.

F16

const F16 = 4

The storage type of IEEE half-precision floats.

BF16

const BF16 = 5

The storage type of bfloat16 floats.

FP8

const FP8 = 9

The storage type of FP8 E4M3 floats.

MASK

const MASK = -1

The type of a comparison's result: a mask, for select, and, or and not.