TRITON LAB #01 — From a Triton kernel to TTIR, on a laptop with no GPU

TRACK: triton | DIFFICULTY: 2 | NEEDS: ../.venv (triton 3.8), kernels.py

What this teaches

Your project compiles Triton kernels for ARM CPUs. Triton was designed for GPUs, so first understand what a Triton program is before any backend touches it, and which parts are GPU-only and must be thrown away.

A Triton kernel is a block program: you write what ONE "program" does to ONE block of data (BLOCK elements, or a BM×BN tile), and the launcher runs it for every block in a grid. tl.program_id(0) says which block you are. Everything inside is written on whole tensors (x + y adds 1024 elements), so it's already vector code in spirit.

STEP 1 — Generate the IR for three kernels

../.venv/bin/python kernels.py

vec_add: out/vec_add.ttir (48 lines), out/vec_add.ttgir (49 lines), out/vec_add.llir (155 lines)
softmax: out/softmax.ttir (73 lines), out/softmax.ttgir (74 lines), out/softmax.llir (272 lines)
matmul: out/matmul.ttir (101 lines), out/matmul.ttgir (180 lines), out/matmul.llir (633 lines)

triton.compile() needs a target description, not a device. We describe an sm_80 GPU, Triton runs its pipeline, and we keep each stage's IR. (The pip wheel bundles ptxas, so even PTX works.)

STEP 2 — Read TTIR: vec_add

$ grep -v '^#loc' out/vec_add.ttir | sed 's/ loc(#loc[0-9]*)//g; s/ loc("[a-z_]*"(#loc))//g' | sed -n '2,20p'
  tt.func public @vec_add(%x_ptr: !tt.ptr<f32>, %y_ptr: !tt.ptr<f32>, %out_ptr: !tt.ptr<f32>, %n: i32) attributes {noinline = false} {
    %c1024_i32 = arith.constant 1024 : i32
    %pid = tt.get_program_id x : i32
    %offs = arith.muli %pid, %c1024_i32 : i32
    %offs_0 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32>
    %offs_1 = tt.splat %offs : i32 -> tensor<1024xi32>
    %offs_2 = arith.addi %offs_1, %offs_0 : tensor<1024xi32>
    %mask = tt.splat %n : i32 -> tensor<1024xi32>
    %mask_3 = arith.cmpi slt, %offs_2, %mask : tensor<1024xi32>
    %x = tt.splat %x_ptr : !tt.ptr<f32> -> tensor<1024x!tt.ptr<f32>>
    %x_4 = tt.addptr %x, %offs_2 : tensor<1024x!tt.ptr<f32>>, tensor<1024xi32>
    %x_5 = tt.load %x_4, %mask_3 : tensor<1024x!tt.ptr<f32>>
    %y = tt.splat %y_ptr : !tt.ptr<f32> -> tensor<1024x!tt.ptr<f32>>
    %y_6 = tt.addptr %y, %offs_2 : tensor<1024x!tt.ptr<f32>>, tensor<1024xi32>
    %y_7 = tt.load %y_6, %mask_3 : tensor<1024x!tt.ptr<f32>>
    %0 = tt.splat %out_ptr : !tt.ptr<f32> -> tensor<1024x!tt.ptr<f32>>
    %1 = tt.addptr %0, %offs_2 : tensor<1024x!tt.ptr<f32>>, tensor<1024xi32>
    %2 = arith.addf %x_5, %y_7 : tensor<1024xf32>
    tt.store %1, %2, %mask_3 : tensor<1024x!tt.ptr<f32>>

Line by line against the Python:

The key observation: pointers are tensors of arbitrary addresses. tt.load is a gather in general. A CPU backend's first big job (triton-shared's "pointer analysis") is proving that splat(base) + arange is really base + contiguous offset, so it can become a strided memref view and a plain vector load. Lab 02 shows the result.

STEP 3 — Reductions and control flow (softmax, matmul)

$ grep -oE 'tt\.[a-z_]+|math\.[a-z]+' out/softmax.ttir | sort | uniq -c
      1 math.exp
      4 tt.addptr
      1 tt.func
      1 tt.get_program_id
      1 tt.load
      1 tt.make_range
     12 tt.ptr
      4 tt.reduce
      1 tt.return
      5 tt.splat
      1 tt.store
$ grep -oE 'tt\.[a-z_]+|scf\.[a-z]+' out/matmul.ttir | sort | uniq -c
      1 scf.for
      1 scf.yield
      8 tt.addptr
      5 tt.broadcast
      1 tt.dot
      4 tt.expand_dims
      1 tt.func
      2 tt.get_program_id
      2 tt.load
      1 tt.make_range
     30 tt.ptr
      1 tt.return
      8 tt.splat
      1 tt.store

• softmax: tt.reduce (with a region holding the combiner, maxnumf / addf) → becomes a linalg.reduce/reduction generic (lab 03_linalg/01's pattern 3), and math.exp → needs a vectorizable exp on the CPU (a real perf item: libm expf per element is slow).
• matmul: tt.dot (a block matmul → linalg.matmul, lab 04/05's world) inside an explicit scf.for over K, with tt.expand_dims / tt.broadcast building 2-D pointer blocks.

STEP 4 — What the GPU pipeline adds (and a CPU backend throws away)

$ grep -m2 -E '#blocked|ttg\.' out/vec_add.ttgir | cut -c1-150
#blocked = #ttg.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}>
module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:80", "ttg.threads-per-warp" = 32 : i32} {

TTGIR attaches a layout to every tensor: which thread of which warp holds which element (sizePerThread, threadsPerWarp = 32, warpsPerCTA = 4). A CPU has no warps. CPU backends (triton-shared and vendor CPU backends) branch off at TTIR, before any of this exists. That's why a vendor MLIR backend's compiler.py stages are ttir → ttsharedir → llir → o: no ttgir.

The mapping to remember

Triton concept              CPU compiler concept
--------------              --------------------
program / grid              outer loop (or threads) over blocks, in the LAUNCHER
BLOCK-sized tensor          a tile: linalg op on tensor<BLOCKxf32>, vectorized
tensor of pointers          (after pointer analysis) a strided memref view
mask                        a dynamic length / padded tail / vector mask
tl.dot                      linalg.matmul on the tile
tl.sum / tl.max             linalg reduction
num_warps, layouts          (ignored)

Try this yourself

1. Change BLOCK to 256 in kernels.py. What changes in the TTIR? (Only the types.)
2. Write a kernel that loads with a stride (x_ptr + offs * 2). Is the tt.addptr pattern still recognisable as base + stride?
3. Write @triton.jit def gelu(...) with tl.math.erf or a tanh approximation. Which math.* ops appear in TTIR?