Dawn's cuTile backend

A GPU kernel in Dawn is an ordinary function under the !Dev effect. Run it once and every operation it raises is recorded as CUDA Tile IR, which NVIDIA's tileiras assembles into a cubin.

  • !Dev · Tile IR · tileiras
  • 208 golden kernels
  • 3 GPUs on record

Two layers, two effects

The host owns the buffers and the device computes over them; they meet only where the host launches a kernel by name.

the hoststd/gpu.dawn!Gpu
  1. main
  2. launch_entry3
  3. with_gpu_realwith_gpu_fake

std/gpu's Gpu effect: allocate, upload, launch, download. with_gpu_real drives the CUDA driver, with_gpu_fake runs the same program with no GPU.

the devicepackages/tileir!Dev
  1. vadd
  2. record
  3. Tile IR
  4. tileiras
  5. cubin

packages/tileir's Dev effect. One recorded run of a kernel becomes Tile IR text, or bytecode for tileiras.

Writing a kernel

A kernel says where it reads and writes once, on its parameters' markers, and states no shape it could read off its operands. Each piece of code below is cut from the golden kernels when the site is built.

Learn to write one

The markers are the addressing

vadd

In and Out cut each tensor into cells, and the Out cells are the launch grid. The body reads its block's cell and writes its block's cell; it names no offset.

fn vadd(a: Param[Float], b: Param[Float], out: Param[Float]) -> Unit !Dev =  store_cell(out, load_cell(a) + load_cell(b)) "vadd" -> {  let (p, _e) = trace3("vadd", In(F64, run_of(128)), In(F64, run_of(128)), Out(F64, run_of(128)), vadd)  p} fn run_of(tile: Int) -> Cells = cells([DYN_DIM], [tile])

kernels.dawn:48–49 · kernels.dawn:5035–5038 · kernels.dawn:4934 · vadd.mlir

Shapes come from the operands

softmax

No operation takes a shape or a format. A constant, and a reduction of a rank-1 tile, are rank-0 tiles: they widen on their own where they meet a wider operand, and nowhere else.

fn softmax(x: Param[Float], out: Param[Float]) -> Unit !Dev = {  let t = load_cell(x)  let ex = exp(t - reduce_max(t))  store_cell(out, ex / reduce_sum(ex))}

kernels.dawn:224–228 · softmax.mlir

Writes go through Out

matmul

store_cell and store_sub are the only writes an Out takes, and an accumulator starts from zeros of it. FREE_AXIS leaves a dimension for the kernel to pick, here the loop's k.

fn matmul(a: Param[Float], b: Param[Float], c: Param[Float]) -> Unit !Dev = {  var acc = zeros(c)  for k in d_range(0, MM_K / MM_TK) { acc = mma(load_at(a, [k]), load_at(b, [k]), acc) }  store_cell(c, acc)} "matmul" -> {  let (p, _e) = trace3("matmul",    In(F64, cells([MM_M, MM_K], [MM_TM, MM_TK], along: [0, FREE_AXIS])),    In(F64, cells([MM_K, MM_N], [MM_TK, MM_TN], along: [FREE_AXIS, 1])),    Out(F64, cells([MM_M, MM_N], [MM_TM, MM_TN])), matmul)  p}

kernels.dawn:389–393 · kernels.dawn:5159–5165 · matmul.mlir

Shared is the escape hatch

histogram

An atomic is not a cell write, and neither is a scatter whose addresses are data or a block that writes two regions. Those write through Shared, and the marker says so where a reader can find it.

fn histogram(input: Param[I32], hist: Param[I32]) -> Unit !Dev = {  let x = load_cell(input)  let live = tail_mask(tile_at(block_id(0), HIST_TILE), [HIST_TILE], HIST_N)  let ok = and_mask(live, in_range(x, 0, HIST_BINS))  atomic_add(hist, and_i(x, i_const(HIST_BINS - 1)), i_const(1), mask: Some(ok))} "histogram" -> {  let (p, _e) = trace2("histogram", In(I32, cells([HIST_N], [HIST_TILE])), Shared(I32), histogram)  p}

kernels.dawn:1613–1618 · kernels.dawn:5357–5360 · histogram.mlir

On sm_86, a 128 by 128 f16 matrix product wants hint_occupancy(2), which fits two blocks on each SM instead of one. How it was measured.

The package's whole surface · What changed in 0.9.0

One kernel, line by line

Every line of one kernel, a fused attention and the source of the row statistics above, sits beside the Tile IR its calls wrote, each run under the call that wrote it. The pairing comes from the recording and from Dawn's own parser, never from a hand-written table. In each run the bold line is the operation the call is for; the lines above it are addressing the lowering added. Each row starts folded to one of those bold lines, the one its outermost call wrote; click a row to open the rest. The same function beside the C and the JVM bytecode it compiles to is on the explorer page.

flash_attnone loop, two reductions inside it

38 calls 50 operations 62 lines of Tile IR · 691 bytes

  1. 2176fn flash_attn(q: Param[Float], k: Param[Float], v: Param[Float], o: Param[Float], scale: Param[Float]) -> Unit !Dev = {
    1cuda_tile.module @m {2  entry @flash_attn(%arg0: tile<ptr<f64>>, %arg1: tile<ptr<f64>>, %arg2: tile<ptr<f64>>, %arg3: tile<ptr<f64>>, %arg4: tile<f64>) {3    %0 = make_token : token
  2. 2177  let tq = load_cell(q)
    load_cell
    4    %1 = assume div_by<16>, %arg0 : tile<ptr<f64>>5    %2 = make_tensor_view %1, shape = [64, 32], strides = [32, 1] : tensor_view<64x32xf64, strides=[32, 1]>6    %3 = make_partition_view %2 : partition_view<tile=(32x32), padding_value = zero, tensor_view<64x32xf64, strides=[32, 1]>, dim_map=[0, 1]>7    %4, %5, %6 = get_tile_block_id : tile<i32>8    %7, %8, %9 = get_tile_block_id : tile<i32>9    %10, %11 = load_view_tko weak %3[%4, %8] token=%0 : partition_view<tile=(32x32), padding_value = zero, tensor_view<64x32xf64, strides=[32, 1]>, dim_map=[0, 1]>, tile<i32> -> tile<32x32xf64>, token
  3. 2178  var m: Tile[Float] = full([FA_BQ, 1], -INFINITY)
    full
    10    %12 = constant <f64: -Infinity> : tile<32x1xf64>
  4. 2179  var l: Tile[Float] = full([FA_BQ, 1], 0.0)
    full
    11    %13 = constant <f64: 0.0> : tile<32x1xf64>
  5. 2180  var acc = zeros(o)
    zeros
    12    %14 = constant <f64: 0.0> : tile<32x32xf64>
  6. 2181  for j in d_range(0, ATT_N / FA_BK) {
    d_range
    13    %15 = constant <i32: 0> : tile<i32>14    %16 = constant <i32: 2> : tile<i32>15    %17 = constant <i32: 1> : tile<i32>16    %18, %19, %20, %21 = for %22 in (%15 to %16, step %17) : tile<i32> iter_values(%23 = %12, %24 = %13, %25 = %14, %26 = %11) -> (tile<32x1xf64>, tile<32x1xf64>, tile<32x32xf64>, token) {
  7. 2182    let s: Tile[Float] = mma(tq, load_at(k, [j]).transpose(), 0.0) * scalar(scale)
    load_at
    17      %27 = assume div_by<16>, %arg1 : tile<ptr<f64>>18      %28 = make_tensor_view %27, shape = [64, 32], strides = [32, 1] : tensor_view<64x32xf64, strides=[32, 1]>19      %29 = make_partition_view %28 : partition_view<tile=(32x32), padding_value = zero, tensor_view<64x32xf64, strides=[32, 1]>, dim_map=[0, 1]>20      %30, %31 = load_view_tko weak %29[%22, %8] token=%26 : partition_view<tile=(32x32), padding_value = zero, tensor_view<64x32xf64, strides=[32, 1]>, dim_map=[0, 1]>, tile<i32> -> tile<32x32xf64>, token
    transpose
    21      %32 = permute %30 [1, 0] : tile<32x32xf64> -> tile<32x32xf64>
    mma
    22      %33 = constant <f64: 0.0> : tile<32x32xf64>23      %34 = mmaf %10, %32, %33 : tile<32x32xf64>, tile<32x32xf64>, tile<32x32xf64>
    mul
    24      %35 = reshape %arg4 : tile<f64> -> tile<1x1xf64>25      %36 = broadcast %35 : tile<1x1xf64> -> tile<32x32xf64>26      %37 = mulf %34, %36 rounding<nearest_even> : tile<32x32xf64>
  8. 2183    let m_new = max(m, reduce_max(s, keepdims: true))
    reduce_max
    27      %38 = reduce %37 dim=1 identities=[-Infinity : f64] : tile<32x32xf64> -> tile<32xf64> (%39: tile<f64>, %40: tile<f64>) {28        %41 = maxf %39, %40 : tile<f64>29        yield %41 : tile<f64>30      }31      %42 = reshape %38 : tile<32xf64> -> tile<32x1xf64>
    max
    32      %43 = maxf %23, %42 : tile<32x1xf64>
  9. 2184    let p = exp(s - m_new)
    sub
    33      %44 = broadcast %43 : tile<32x1xf64> -> tile<32x32xf64>34      %45 = subf %37, %44 rounding<nearest_even> : tile<32x32xf64>
    exp
    35      %46 = exp %45 : tile<32x32xf64>
  10. 2185    let alpha = exp(m - m_new)
    sub
    36      %47 = subf %23, %43 rounding<nearest_even> : tile<32x1xf64>
    exp
    37      %48 = exp %47 : tile<32x1xf64>
  11. 2186    l = l * alpha + reduce_sum(p, keepdims: true)
    mul
    38      %49 = mulf %24, %48 rounding<nearest_even> : tile<32x1xf64>
    reduce_sum
    39      %50 = reduce %46 dim=1 identities=[0.0 : f64] : tile<32x32xf64> -> tile<32xf64> (%51: tile<f64>, %52: tile<f64>) {40        %53 = addf %51, %52 rounding<nearest_even> : tile<f64>41        yield %53 : tile<f64>42      }43      %54 = reshape %50 : tile<32xf64> -> tile<32x1xf64>
    add
    44      %55 = addf %49, %54 rounding<nearest_even> : tile<32x1xf64>
  12. 2187    acc = mma(p, load_at(v, [j]), acc * alpha)
    load_at
    45      %56 = assume div_by<16>, %arg2 : tile<ptr<f64>>46      %57 = make_tensor_view %56, shape = [64, 32], strides = [32, 1] : tensor_view<64x32xf64, strides=[32, 1]>47      %58 = make_partition_view %57 : partition_view<tile=(32x32), padding_value = zero, tensor_view<64x32xf64, strides=[32, 1]>, dim_map=[0, 1]>48      %59, %60 = load_view_tko weak %58[%22, %8] token=%31 : partition_view<tile=(32x32), padding_value = zero, tensor_view<64x32xf64, strides=[32, 1]>, dim_map=[0, 1]>, tile<i32> -> tile<32x32xf64>, token
    mul
    49      %61 = broadcast %48 : tile<32x1xf64> -> tile<32x32xf64>50      %62 = mulf %25, %61 rounding<nearest_even> : tile<32x32xf64>
    mma
    51      %63 = mmaf %46, %59, %62 : tile<32x32xf64>, tile<32x32xf64>, tile<32x32xf64>
  13. 2188    m = m_new
  14. 2189  }
    d_range
    52      continue %43, %55, %63, %60 : tile<32x1xf64>, tile<32x1xf64>, tile<32x32xf64>, token53    }
  15. 2190  store_cell(o, acc / l)
    div
    54    %64 = broadcast %19 : tile<32x1xf64> -> tile<32x32xf64>55    %65 = divf %20, %64 rounding<nearest_even> : tile<32x32xf64>
    store_cell
    56    %66 = assume div_by<16>, %arg3 : tile<ptr<f64>>57    %67 = make_tensor_view %66, shape = [64, 32], strides = [32, 1] : tensor_view<64x32xf64, strides=[32, 1]>58    %68 = make_partition_view %67 : partition_view<tile=(32x32), padding_value = zero, tensor_view<64x32xf64, strides=[32, 1]>, dim_map=[0, 1]>59    %69 = store_view_tko weak %65, %68[%4, %8] token=%21 : tile<32x32xf64>, partition_view<tile=(32x32), padding_value = zero, tensor_view<64x32xf64, strides=[32, 1]>, dim_map=[0, 1]>, tile<i32> -> token
  16. 2191}
    60    return61  }62}

What the backend covers

Every public opcode, type tag and attribute value of the Tile IR release has a row, a status and its evidence. Point at a card to see the rows that are not implemented.

102/105public opcodes implemented3 structuralfeatures.txt
  • entry structural · 3
  • global structural · T7
  • module structural · 3
24type tags implementedtypes.txt
52/53attribute values implemented1 unimplementedattrs.txt
  • ptr_attr.none unimplemented · T21
208golden kernels, each pinned as text and as bytecodetile-golden/

Three gates

Each layer catches what the one before it cannot, and each has to turn its own mutants red. Each dot is one of them, at the gate it was written against.

  1. 0

    Text and bytes

    10 mutants caught here

    Every golden kernel is traced twice and held to its recorded text and bytecode, byte for byte. It says whether anything changed, not whether it is right.

    scripts/tile-golden/run.sh
  2. 1

    The assembler accepts

    83 mutants caught here

    tileiras has to assemble every bytecode golden into a cubin that names its kernel. It catches encoding and type errors, not wrong answers.

    scripts/tile-golden/run.sh
  3. 2

    The device agrees

    43 mutants caught here

    On a GPU, every output buffer is held to a host reference, bit for bit or within a stated tolerance. Each run appends a ledger line, and CI refuses a tree the last line did not test.

    scripts/tile-gpu-diff/run.sh

On real hardware

The last line of each machine's ledger. CI reads the first; the sm_90 and sm_100 ledgers reach the fp8, fp4 and block-scaled rows an Ampere card cannot load.

GPUdatedrivertileirasresultbit for bitwithin tolerancetree
sm_862026-10-14616.5613.4.92pass106459fabf4af78f2
sm_902026-10-11580.95.0513.4.92pass10945b88bc2b4a1cb
sm_1002026-10-11580.159.0413.4.92pass11245b88bc2b4a1cb

Kernels compared, by tier, summed over the line's families.