11. Capra, the Compiler: From a Tensor Graph to GA-2 Instructions¶
About the animals in this picture and where they are standing
- Alpine ibex (Capra ibex) lives in the European Alps, usually above the tree line on steep rock, and is a famously sure-footed climber. Males carry horns that can reach about a metre. Hunting left only a small remnant, around the Gran Paradiso in Italy, by the early 1800s; reintroductions have since restored it across the Alps.
- Siberian ibex (Capra sibirica) lives in the high mountains of Central Asia, from the Altai and Tian Shan to the Himalaya. It is one of the largest ibexes, and the males' ridged horns can exceed a metre in length.
- Markhor (Capra falconeri) lives in the mountains of Pakistan, Afghanistan and neighbouring regions. The males' horns twist in a corkscrew; the name is often explained as Persian for 'snake-eater'. It is Pakistan's national animal, and community-based conservation has helped some populations recover.
- West Caucasian tur (Capra caucasica) lives only in the western Caucasus, on steep slopes above the forest line. Its horns are short and heavy, curving outward and back, and it moves up and down cliffs with ease.
- Bezoar ibex (wild goat) (Capra aegagrus) lives from Turkey and the Caucasus through Iran to Pakistan. It is the wild ancestor of the domestic goat, which people began keeping some ten thousand years ago.
- Iberian ibex (Capra pyrenaica) lives in the mountains of Spain and Portugal. Its Pyrenean subspecies, the bucardo, died out in 2000. In 2003 scientists cloned one; the kid lived only minutes, making it the first extinct animal to be cloned.
The scene is imaginary: it is set on Pluto, a dwarf planet with a heart-shaped plain of nitrogen ice; its largest moon, Charon, is about half its width. The animals are stylised drawings, shown where they could not live, for effect.
What you will understand: how Capra (named for Capra, the genus of the wild goats and ibexes) turns a model, written as a graph of tensor operations in floating point, becomes a program for the chip of Chapter 7 that no human wrote: how the compiler chooses an int8 scale for every tensor, cuts big matrix products into the 4 x 4 tiles the hardware can do, places tensors in the 4,096-word scratchpad and reuses the space as soon as they are dead, and emits instructions. You will also see how a compiler is tested: against a plain-Python interpreter, against the meaning of each number, against its own allocator, and against graphs it must refuse.
What you need to know first: Chapters 3 (quantization), 7 (the ISA) and 8-9 (programs written by hand, which is what this chapter automates). Appendix D (ML linear algebra) covers matrices and matrix products if they are unfamiliar.
What this chapter builds: model/capra.py (the compiler: graph IR, calibration, scale planning, tiling, allocator, emitter, two interpreters), model/capra_tests.py (the test battery), tools/ch11_run.py, tools/mut_ch11.py, and for the running examples tools/ch11_example_a.py (compile one graph pass by pass) and tools/ch11_example_b.py (write your own graph).
Why a compiler¶
Chapters 7 to 9 showed what programming this chip by hand is like. For the 10-instruction network of Chapter 7 it was tolerable. For the attention step of Chapter 8 it took several screens of Python that chose addresses, strides, tile loops, and eight pairs of requantizer constants; a single wrong address or shift gave a plausible, wrong answer. A real model has hundreds of such steps. The work is mechanical, error-prone and exactly the kind a program should do.
A compiler is that program. The user writes what to compute (a graph of tensor operations, in ordinary floating point) and the compiler decides how the chip does it: which int8 scale each tensor gets, how a product that is too big for the array is cut up, where each tensor lives in the scratchpad and when that space can be reused, and which instructions to emit. The name Capra is the genus of the wild goats and ibexes: surefooted animals that find a route up a cliff, a fair description of what a compiler does with a graph and a very small memory.
The compiler is also where the whole book's methods meet. It must be correct for every graph (not just for the ones its author tried), so it is tested against an interpreter, against the real-number meaning of each tensor, against itself with the allocator switched off, and against graphs it must refuse.
Where the compiler sits¶
graph of tensor ops (float)
| calibrate: run on sample data, record each tensor's range
v
scale planning -> an int8 scale per tensor; a multiplier + shift per requantizer
v
tiling -> each matmul becomes 4 x 4 output tiles (MM instructions)
v
allocation -> a scratchpad address for every tensor, reusing dead space
v
emission -> LD / MM / RQ / SM / VADD / AMAX / ST ... HALT
Figure 11.1: the pipeline. Each stage hands a more concrete description to the next.
The input is a small graph IR: tensors are 2-D, and the operations are input, weight, matmul (optionally with the second operand transposed and a constant factor), softmax (per row), add, relu, concat_rows, argmax (per row) and output. That is exactly what a transformer decode step needs; Chapter 12 uses it for one.
Capra: model/capra.py¶
#!/usr/bin/env python3
"""Chapter 11: CAPRA, a compiler from a tensor graph to GA-2 programs (named for the goat genus, Capra).
PIPELINE. graph IR -> float evaluation (calibration) -> scale planning -> tiling -> scratchpad allocation -> instruction emission -> a Program that runs on the reference simulator or the RTL.
THE IR. Tensors are 2-D (rows, cols). Nodes: input, weight, matmul (optionally with B transposed and a constant factor), softmax (per row), add, relu, concat_rows, slice_rows (a view, Chapter 16), argmax (per row), output.
WHAT THE COMPILER DECIDES FOR THE USER: every tensor's int8 scale (from calibration data), every requantizer's multiplier and shift, how a big matrix product is cut into 4x4 output tiles, where each tensor lives in the scratchpad
(buffers are reused once their last consumer has run), and where inputs, weights and outputs live in external memory. Limits of the target: a product's inner dimension K <= 64, 4096 words of scratchpad, int8 tensors.
"""
import math, os, sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
os.environ.setdefault("GA2_EXT", "8192")
import ga2_isa as I
from quant import quantize, mantissa_shift, requant, matmul_int8
from softmax_gold import softmax_fixed
B = I.build
class CompileError(Exception): pass
# ------------------------------------------------------------------ the IR
class Node:
def __init__(self, id, op, ins, shape, **attrs): self.id, self.op, self.ins, self.shape, self.attrs = id, op, ins, shape, attrs
class Graph:
def __init__(self): self.nodes = []
def _add(self, op, ins, shape, **a): n = Node(len(self.nodes), op, ins, shape, **a); self.nodes.append(n); return n.id
def input(self, name, shape): return self._add("input", [], shape, name=name)
def weight(self, name, array, bits=8):
"""A constant matrix. bits=8: stored as int8, one value per word. bits=4 (Chapter 13): stored as signed 4-bit values packed eight to a 32-bit word in external memory and expanded by an UNPACK instruction after the load."""
if bits not in (4, 8): raise CompileError(f"weight {name}: bits must be 4 or 8")
return self._add("weight", [], (len(array), len(array[0])), name=name, array=array, bits=bits)
def matmul(self, a, b, transpose_b=False, scale=1.0):
(m, k), (k2, n) = self.nodes[a].shape, self.nodes[b].shape[::-1] if transpose_b else self.nodes[b].shape
if k != k2: raise CompileError(f"matmul shapes {self.nodes[a].shape} x {self.nodes[b].shape}{'^T' if transpose_b else ''}")
return self._add("matmul", [a, b], (m, n), transpose_b=transpose_b, scale=scale)
def softmax(self, a): return self._add("softmax", [a], self.nodes[a].shape)
def add(self, a, b):
if self.nodes[a].shape != self.nodes[b].shape: raise CompileError("add: shapes differ")
return self._add("add", [a, b], self.nodes[a].shape)
def relu(self, a): return self._add("relu", [a], self.nodes[a].shape)
def concat_rows(self, a, b):
if self.nodes[a].shape[1] != self.nodes[b].shape[1]: raise CompileError("concat_rows: widths differ")
return self._add("concat", [a, b], (self.nodes[a].shape[0] + self.nodes[b].shape[0], self.nodes[a].shape[1]))
def slice_rows(self, a, start, count):
"""Rows start .. start+count-1 of a tensor, as a VIEW: no instruction, no copy (Chapter 16). The slice shares the source's scale and lives inside the source's buffer."""
r = self.nodes[a].shape[0]
if count < 1 or start < 0 or start + count > r: raise CompileError(f"slice_rows: rows {start}..{start + count - 1} do not lie inside a tensor of {r} rows")
return self._add("slice", [a], (count, self.nodes[a].shape[1]), start=start, count=count)
def argmax(self, a): return self._add("argmax", [a], (self.nodes[a].shape[0], 1))
def output(self, a, name): return self._add("output", [a], self.nodes[a].shape, name=name)
def tag(self, a, name): self.nodes[a].attrs["tag"] = name; return a # a name for a tensor, so that its range can be fixed from outside (see compile_graph(ranges=...))
def evaluate(g, inputs):
"""Floating-point evaluation of the graph. inputs: {name: 2-D list}. Returns {node id: 2-D list}."""
v = {}
for n in g.nodes:
if n.op == "input": v[n.id] = inputs[n.attrs["name"]]
elif n.op == "weight": v[n.id] = n.attrs["array"]
elif n.op == "matmul":
a, b = v[n.ins[0]], v[n.ins[1]]; b = [list(r) for r in zip(*b)] if n.attrs["transpose_b"] else b
v[n.id] = [[n.attrs["scale"] * sum(a[i][t] * b[t][j] for t in range(len(b))) for j in range(len(b[0]))] for i in range(len(a))]
elif n.op == "softmax":
out = []
for r in v[n.ins[0]]: mx = max(r); e = [math.exp(x - mx) for x in r]; z = sum(e); out.append([x / z for x in e])
v[n.id] = out
elif n.op == "add": v[n.id] = [[x + y for x, y in zip(r, s)] for r, s in zip(v[n.ins[0]], v[n.ins[1]])]
elif n.op == "relu": v[n.id] = [[max(0.0, x) for x in r] for r in v[n.ins[0]]]
elif n.op == "concat": v[n.id] = [list(r) for r in v[n.ins[0]]] + [list(r) for r in v[n.ins[1]]]
elif n.op == "slice": v[n.id] = [list(r) for r in v[n.ins[0]][n.attrs["start"]:n.attrs["start"] + n.attrs["count"]]]
elif n.op == "argmax": v[n.id] = [[float(r.index(max(r)))] for r in v[n.ins[0]]]
elif n.op == "output": v[n.id] = v[n.ins[0]]
return v
# ------------------------------------------------------------------ scale planning
def plan_scales(g, calib, ranges=None):
"""Assign one int8 scale to every tensor. Tensors that must share a scale (the parts and the result of a concatenation, input and result of relu) form a group; an add does NOT force one: the compiler rescales an operand whose scale differs from the sum's and get the group's largest calibrated range;
the input of a softmax is pinned to 1/16 (the softmax unit's Q4.4 input format) and its output to 1/127. Returns (scale per node id, notes)."""
parent = list(range(len(g.nodes)))
def find(x):
while parent[x] != x: parent[x] = parent[parent[x]]; x = parent[x]
return x
def union(a, b): parent[find(a)] = find(b)
for n in g.nodes:
if n.op == "concat": union(n.ins[0], n.id); union(n.ins[1], n.id)
elif n.op == "relu": union(n.ins[0], n.id)
elif n.op == "slice": union(n.ins[0], n.id)
elif n.op == "output": union(n.ins[0], n.id)
rng = {}
for sample in calib:
vals = evaluate(g, sample)
for nid, m in vals.items(): rng[nid] = max(rng.get(nid, 0.0), max(abs(x) for r in m for x in r))
for n in g.nodes: # ranges fixed from outside (e.g. so that a KV cache keeps one scale from step to step)
if ranges and n.attrs.get("tag") in ranges: rng[n.id] = ranges[n.attrs["tag"]]
grp_max = {}
for n in g.nodes:
if n.op != "argmax": grp_max[find(n.id)] = max(grp_max.get(find(n.id), 0.0), rng[n.id])
if n.op == "add": # the sum's scale must also hold both operands once rescaled to it (large operands may cancel)
for i in n.ins: grp_max[find(n.id)] = max(grp_max[find(n.id)], rng[i])
pinned = {}; notes = []
for n in g.nodes:
if n.op == "softmax":
pinned[find(n.ins[0])] = 1 / 16
if grp_max[find(n.ins[0])] > 127 / 16: notes.append(f"node {n.ins[0]}: scores reach {grp_max[find(n.ins[0])]:.2f}, above the softmax input range of {127/16:.2f}: they will saturate")
pinned[find(n.id)] = 1 / 127
scale = {}
for n in g.nodes:
if n.op == "argmax": continue
r = find(n.id); scale[n.id] = pinned.get(r, (grp_max[r] if grp_max[r] > 0 else 1.0) / 127)
if r in pinned and rng[n.id] > pinned[r] * 127 * 1.001 and n.op in ("concat", "relu", "input", "weight") and n.id != r:
notes.append(f"node {n.id} ({n.op}) reaches {rng[n.id]:.3f}, above the {pinned[r]*127:.3f} that its pinned scale can hold: it will saturate")
for n in g.nodes: # int4 weights use the range -7..7: the scale is range / 7 (Chapter 13)
if n.op == "weight" and n.attrs.get("bits", 8) == 4: scale[n.id] = (rng[n.id] if rng[n.id] > 0 else 1.0) / 7
return scale, notes
# ------------------------------------------------------------------ memory allocation
class Allocator:
"""First-fit allocator for the scratchpad (4096 words) with coalescing free."""
def __init__(self, size=I.SPAD): self.free_list = [(0, size)]; self.high = 0 # high: the highest scratchpad address ever handed out + 1
def alloc(self, n):
for i, (s, sz) in enumerate(self.free_list):
if sz >= n:
self.free_list[i:i + 1] = [(s + n, sz - n)] if sz > n else []; self.high = max(self.high, s + n); return s
raise CompileError(f"scratchpad is full: cannot place a buffer of {n} words (free: {sum(sz for _, sz in self.free_list)} words in {len(self.free_list)} pieces)")
def release(self, s, n):
self.free_list.append((s, n)); self.free_list.sort(); merged = []
for a, b in self.free_list:
if merged and merged[-1][0] + merged[-1][1] == a: merged[-1] = (merged[-1][0], merged[-1][1] + b)
else: merged.append((a, b))
self.free_list = merged
# ------------------------------------------------------------------ the compiler
class Program:
pass
def compile_graph(g, calib, reuse=True, ranges=None):
"""Compile `g` to GA-2 instructions. `calib` is a list of input dicts for calibration; `ranges` optionally fixes the range of tagged tensors. reuse=False keeps every buffer alive (for debugging and for differential testing of the allocator)."""
P = Program(); P.graph = g; P.reuse = reuse; scale, P.notes = plan_scales(g, calib, ranges); P.scale = scale; N = g.nodes
cons = {n.id: [] for n in N}
for n in N:
for i in n.ins: cons[i].append(n.id)
# relu fusion
P.relu_fused = {}
for n in N:
if n.op == "relu":
src = N[n.ins[0]]
if src.op != "matmul" or len(cons[src.id]) != 1: raise CompileError(f"relu on node {src.id} ({src.op}) cannot be fused into a requantizer")
P.relu_fused[src.id] = True
# buffers: where does each tensor live?
loc = {} # node id -> (buffer id, row offset)
bufs = [] # buffer id -> dict(size, members)
def new_buf(size): bufs.append({"size": size, "members": [], "addr": None}); return len(bufs) - 1
for n in reversed(N): # consumers first: a concat's output is placed before its inputs, which become views of it; a fused relu's input shares the relu's place
if n.op in ("output", "argmax", "slice"): continue
if n.id not in loc: loc[n.id] = (new_buf(((n.shape[0] * n.shape[1] + 7) // 8) * 8 if n.op == "weight" and n.attrs.get("bits", 8) == 4 else n.shape[0] * n.shape[1]), 0)
if n.op == "relu":
if n.ins[0] in loc: raise CompileError(f"node {n.ins[0]} feeds two consumers that need it in different places")
loc[n.ins[0]] = loc[n.id]
if n.op == "concat":
b0, off = loc[n.id]
for k, i in enumerate(n.ins):
if i in loc: raise CompileError(f"node {i} feeds two concatenations")
loc[i] = (b0, off + (0 if k == 0 else N[n.ins[0]].shape[0]))
for n in N:
if n.op == "slice": # a view into its source's place (Chapter 16)
if any(N[c].op in ("concat", "relu") for c in cons[n.id]): raise CompileError(f"node {n.id}: a slice cannot feed a concatenation or a relu (it would need a copy)")
b, off = loc[n.ins[0]]; loc[n.id] = (b, off + n.attrs["start"])
for n in N:
if n.op == "argmax": loc[n.id] = (new_buf(n.shape[0]), 0)
for nid, (b, off) in loc.items(): bufs[b]["members"].append(nid)
# lifetimes in steps (= node order)
first = {}; last = {}
for n in N:
if n.op in ("input", "weight"): continue
if n.id in loc: first.setdefault(loc[n.id][0], n.id); last[loc[n.id][0]] = max(last.get(loc[n.id][0], 0), n.id)
for i in n.ins:
if i in loc: b = loc[i][0]; first.setdefault(b, n.id); last[b] = max(last.get(b, 0), n.id)
P.ext = {}; ext_top = 0
# external layout: weights, inputs, outputs
for n in N:
if n.op in ("input", "weight"):
words = (n.shape[0] * n.shape[1] + 7) // 8 if n.op == "weight" and n.attrs.get("bits", 8) == 4 else n.shape[0] * n.shape[1] # int4 weights are packed eight to a word
P.ext[n.attrs["name"]] = {"addr": ext_top, "shape": n.shape, "node": n.id, "kind": n.op, "words": words}; ext_top += words
for n in N:
if n.op == "output": P.ext[n.attrs["name"]] = {"addr": ext_top, "shape": n.shape, "node": n.id, "kind": "output"}; ext_top += n.shape[0] * n.shape[1]
if ext_top > I.EXT: raise CompileError(f"external memory is full ({ext_top} words)")
alloc = Allocator(); code = []; loaded = set()
def addr_of(nid):
b, off = loc[nid]
if bufs[b]["addr"] is None: bufs[b]["addr"] = alloc.alloc(bufs[b]["size"])
return bufs[b]["addr"] + off * N[nid].shape[1]
def ensure(nid):
n = N[nid]
if n.op in ("input", "weight") and nid not in loaded:
a = addr_of(nid); e = P.ext[n.attrs["name"]]["addr"]; loaded.add(nid)
if n.op == "weight" and n.attrs.get("bits", 8) == 4: # load the packed words, then expand them into the weight's buffer
pw = P.ext[n.attrs["name"]]["words"]; tmp = alloc.alloc(pw); code.append(B("LD", dst=tmp, src=e, len=pw)); code.append(B("UNPACK", dst=a, src=tmp, len=pw)); alloc.release(tmp, pw)
else: code.append(B("LD", dst=a, src=e, len=n.shape[0] * n.shape[1]))
return addr_of(nid)
P.params = {}
for n in N:
if n.op == "input" or n.op == "weight": continue
if n.op == "slice": ensure(n.ins[0]); addr_of(n.id); continue
if n.op == "concat" or n.op == "relu": ensure(n.ins[0]); ensure(n.ins[1]) if n.op == "concat" else None; addr_of(n.id); continue
ins = [ensure(i) for i in n.ins]
if n.op == "matmul":
(m, k) = N[n.ins[0]].shape; nn = n.shape[1]; tb = n.attrs["transpose_b"]
if k > 64: raise CompileError(f"matmul node {n.id}: inner dimension {k} exceeds the matrix unit's limit of 64")
out = addr_of(n.id); acc = alloc.alloc(m * nn)
for i0 in range(0, m, 4):
for j0 in range(0, nn, 4):
mt, nt = min(4, m - i0), min(4, nn - j0)
code.append(B("MM", dst=acc + i0 * nn + j0, A=ins[0] + i0 * k, B=ins[1] + (j0 * k if tb else j0), M=mt, K=k, N=nt, tb=int(tb), lda=k, ldb=(k if tb else nn), ldc=nn))
ms = mantissa_shift(n.attrs["scale"] * scale[n.ins[0]] * scale[n.ins[1]] / scale[n.id]); P.params[n.id] = ms
code.append(B("RQ", dst=out, src=acc, len=m * nn, m=ms[0], s=ms[1], relu=int(P.relu_fused.get(n.id, False)))); alloc.release(acc, m * nn)
elif n.op == "softmax":
m, nn = n.shape; out = addr_of(n.id); tmp = alloc.alloc(m * nn)
for r in range(m): code.append(B("SM", dst=tmp + r * nn, src=ins[0] + r * nn, len=nn))
ms = mantissa_shift(127 / 65536); P.params[n.id] = ms; code.append(B("RQ", dst=out, src=tmp, len=m * nn, m=ms[0], s=ms[1])); alloc.release(tmp, m * nn)
elif n.op == "add":
size = n.shape[0] * n.shape[1]; srcs = []; temps = []
for k, i in enumerate(n.ins):
ratio = scale[i] / scale[n.id]
if abs(ratio - 1) < 1e-9: srcs.append(ins[k]); continue
t = alloc.alloc(size); temps.append(t); ms = mantissa_shift(ratio); P.params[(n.id, k)] = ms
code.append(B("RQ", dst=t, src=ins[k], len=size, m=ms[0], s=ms[1])); srcs.append(t) # int8 -> int8 at the sum's scale
code.append(B("VADD", dst=addr_of(n.id), src1=srcs[0], src2=srcs[1], len=size))
for t in temps: alloc.release(t, size)
elif n.op == "argmax":
out = addr_of(n.id); (m, nn) = N[n.ins[0]].shape
for r in range(m): code.append(B("AMAX", dst=out + r, src=ins[0] + r * nn, len=nn))
elif n.op == "output":
e = P.ext[n.attrs["name"]]; code.append(B("ST", src=ins[0], dst=e["addr"], len=n.shape[0] * n.shape[1]))
# free what is no longer needed after this step
if reuse:
for b, l in list(last.items()):
if l == n.id and bufs[b]["addr"] is not None: alloc.release(bufs[b]["addr"], bufs[b]["size"]); bufs[b]["addr"] = None; last[b] = -1
P.spad_high = alloc.high; P.loc = loc; P.bufs = bufs; P.code = code + [B("HALT")]; P.ext_words = ext_top
return P
# ------------------------------------------------------------------ running a compiled program
def _q(vals, sc): return [[quantize(v, sc) for v in r] for r in vals]
def ext_image(P, inputs):
"""External memory for a run: weights and the given inputs quantized with the compiler's scales."""
img = [0] * I.EXT; G = P.graph
for name, e in P.ext.items():
if e["kind"] == "output": continue
n = G.nodes[e["node"]]; data = n.attrs["array"] if e["kind"] == "weight" else inputs[name]
if e["kind"] == "weight" and n.attrs.get("bits", 8) == 4: # pack: element 8i + j goes to bits 4j+3..4j of word i
flat = [v for row in _q(data, P.scale[n.id]) for v in row]; flat += [0] * (-len(flat) % 8)
for i in range(len(flat) // 8): img[e["addr"] + i] = sum((flat[8 * i + j] & 15) << (4 * j) for j in range(8))
continue
for r, row in enumerate(_q(data, P.scale[n.id])):
for c, v in enumerate(row): img[e["addr"] + r * n.shape[1] + c] = v & 0xFFFFFFFF
return img
def run(P, inputs, machine=False):
"""Run on the reference simulator. Returns {output name: (ints, floats)} (and the machine if asked)."""
m = I.Machine(ext_image(P, inputs)).run(P.code); out = {}; G = P.graph
for name, e in P.ext.items():
if e["kind"] != "output": continue
n = G.nodes[e["node"]]; src = G.nodes[n.ins[0]]; r, c = n.shape; w = [m.ext[e["addr"] + i] for i in range(r * c)]
ints = [[(I.s32(w[i * c + j]) if src.op == "argmax" else I.s8(w[i * c + j])) for j in range(c)] for i in range(r)]
sc = 1.0 if src.op == "argmax" else P.scale[n.id]; out[name] = (ints, [[v * sc for v in row] for row in ints])
return (out, m) if machine else out
def read_tensor(P, m, nid):
"""The int values of a tensor, read from the scratchpad after a run (valid for every tensor only when compiled with reuse=False)."""
n = P.graph.nodes[nid]; b, off = P.loc[nid]; base = P.bufs[b]["addr"]
if base is None: raise ValueError("buffer already released")
a = base + off * n.shape[1]; r, c = n.shape
return [[(I.s32(m.spad[a + i * c + j]) if n.op == "argmax" else I.s8(m.spad[a + i * c + j])) for j in range(c)] for i in range(r)]
def interpret(P, inputs):
"""The compiled program's meaning in plain Python integers, node by node: no instructions, no tiles, no addresses. Shares only the arithmetic helpers and the planned scales and requantizers with the compiler."""
G = P.graph; v = {}
for n in G.nodes:
if n.op == "input": v[n.id] = _q(inputs[n.attrs["name"]], P.scale[n.id])
elif n.op == "weight": v[n.id] = _q(n.attrs["array"], P.scale[n.id])
elif n.op == "matmul":
b = v[n.ins[1]]; b = [list(r) for r in zip(*b)] if n.attrs["transpose_b"] else b; acc = matmul_int8(v[n.ins[0]], b); m_, s_ = P.params[n.id]
v[n.id] = [[requant(a, m_, s_, bool(P.relu_fused.get(n.id))) for a in r] for r in acc]
elif n.op == "softmax": m_, s_ = P.params[n.id]; v[n.id] = [[requant(p, m_, s_) for p in softmax_fixed(r)] for r in v[n.ins[0]]]
elif n.op == "add":
ops = [v[i] if (n.id, k) not in P.params else [[requant(x, *P.params[(n.id, k)]) for x in r] for r in v[i]] for k, i in enumerate(n.ins)]
v[n.id] = [[max(-127, min(127, x + y)) for x, y in zip(r, s)] for r, s in zip(ops[0], ops[1])]
elif n.op in ("relu", "output"): v[n.id] = v[n.ins[0]]
elif n.op == "slice": v[n.id] = [list(r) for r in v[n.ins[0]][n.attrs["start"]:n.attrs["start"] + n.attrs["count"]]]
elif n.op == "concat": v[n.id] = [list(r) for r in v[n.ins[0]]] + [list(r) for r in v[n.ins[1]]]
elif n.op == "argmax": v[n.id] = [[r.index(max(r))] for r in v[n.ins[0]]]
return v
def stagewise_problems(P, ints):
"""Every tensor must agree, within one level, with the REAL-NUMBER meaning of its operation applied to the dequantized values of its inputs. Uses only the planned scales, not the requantizer parameters."""
G = P.graph; probs = []; sc = P.scale
deq = lambda nid: [[x * sc[nid] for x in r] for r in ints[nid]]
rnd = lambda x: max(-127, min(127, int(math.floor(abs(x) + 0.5)) * (1 if x >= 0 else -1)))
for n in G.nodes:
if n.op == "matmul":
a, b = deq(n.ins[0]), deq(n.ins[1]); b = [list(r) for r in zip(*b)] if n.attrs["transpose_b"] else b
want = [[rnd(n.attrs["scale"] * sum(a[i][t] * b[t][j] for t in range(len(b))) / sc[n.id]) for j in range(len(b[0]))] for i in range(len(a))]
if P.relu_fused.get(n.id): want = [[max(0, x) for x in r] for r in want]
elif n.op == "softmax":
want = []
for r in deq(n.ins[0]): mx = max(r); e = [math.exp(x - mx) for x in r]; z = sum(e); want.append([rnd(x / z / sc[n.id]) for x in e])
elif n.op == "add": want = [[rnd((x + y) / sc[n.id]) for x, y in zip(r, s)] for r, s in zip(deq(n.ins[0]), deq(n.ins[1]))]
else: continue
tol = 2 if n.op == "add" else 1 # an add may rescale both operands (half a level each) before it rounds
bad = sum(1 for r, s in zip(ints[n.id], want) for x, y in zip(r, s) if abs(x - y) > tol)
if bad: probs.append(f"node {n.id} ({n.op}): {bad} values more than {tol} level(s) from the real-number result")
return probs
Reading the code. Four things are worth looking for. Graph and Node are the whole IR: a graph is a list of nodes, each with an operation, the ids of its inputs and a shape; the methods (matmul, softmax, concat_rows...) check shapes at construction, so a bad graph fails when built. plan_scales uses a union-find to group tensors that must share a scale, then looks at calibration ranges. Allocator is a first-fit allocator with coalescing: the classic malloc algorithm on a 4,096-word arena. compile_graph is one pass over the nodes in order, emitting instructions as it goes; the tiling loops are the two nested for loops over i0 and j0 in the matmul branch.
What it decides, in order:
- Calibration (
plan_scales). The graph is evaluated in floating point on sample inputs; each tensor's largest absolute value gives its int8 scale (max / 127). Tensors that must share a scale form a group: the parts of a concatenation and its result, or the input and output of arelu. Two scales are pinned by the hardware: a softmax's input must be in Q4.4 (1/16) and its output is a probability (1/127). Anadddoes not force a shared scale: the compiler rescales an operand whose scale differs from the sum's, with anRQfrom int8 to int8. - Requantizers. For
matmulthe multiplier isfactor x scale_a x scale_b / scale_out, turned into the 24-bit mantissa and shift of Chapter 3. - Tiling. A product of an
M x Kand aK x Nmatrix becomesceil(M/4) x ceil(N/4)MMinstructions, each with the right strides (lda,ldb,ldc) and the transpose flag. The matrix unit's limitK <= 64is a compile error, not a surprise. - Allocation. Each tensor gets a buffer; a concatenation's inputs become views of its output buffer (so the concatenation costs no instructions: the producers write straight into place), and a fused
relushares the buffer of its producer. Buffers are allocated first-fit when first needed and released after their last consumer, with coalescing. - Emission. Weights and inputs are loaded just before first use, outputs stored at the end, softmax and argmax emitted row by row.
Two interpreters accompany it (interpret and stagewise_problems), written for testing and discussed below.
An example, compiled¶
#!/usr/bin/env python3
"""Chapter 11: the compiler end to end. A readable example, twelve random graphs run on the RTL in both simulators, and a network that fits the scratchpad only because the allocator reuses buffers. Usage: ch11_run.py"""
import math, os, random, re, sys
os.environ["GA2_EXT"] = "8192"
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))); sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "model"))
import hw, ga2_isa as I, ga2_progs as G, capra as C, capra_tests as T
F = ["rtl/sram.v", "rtl/dma.v", "rtl/systolic.v", "rtl/requant.v", "rtl/exp_lut.v", "rtl/divu.v", "rtl/rowmem.v", "rtl/ga2_vec.v", "rtl/ga2_sm.v", "rtl/ga2_mm.v", "rtl/ga2.v", "tb/extmem_rw.v", "tb/ga2_tb.v"]
def line(out): return ([l for l in out.splitlines() if l.startswith(("PASS", "FAIL", "MISMATCH"))] or out.strip().splitlines()[-2:])[0]
def dump(g):
for n in g.nodes:
extra = n.attrs.get("name", ""); extra += (" ^T" if n.attrs.get("transpose_b") else "") + (f" x{n.attrs['scale']:.3g}" if n.op == "matmul" and n.attrs.get("scale") != 1.0 else "")
print(f" %{n.id:<2d} = {n.op:8s} {', '.join('%' + str(i) for i in n.ins):10s} -> {n.shape} {extra}")
def err(P, out, test):
fl = C.evaluate(P.graph, test); w = 0.0
for nm, e in P.ext.items():
if e["kind"] != "output" or P.graph.nodes[P.graph.nodes[e["node"]].ins[0]].op == "argmax": continue
ref = fl[e["node"]]; got = out[nm][1]; num = math.sqrt(sum((a - b) ** 2 for r, t in zip(got, ref) for a, b in zip(r, t))); den = math.sqrt(sum(b * b for t in ref for b in t)) or 1.0; w = max(w, num / den)
return w
print("== 1. one graph, compiled: a two-layer network with a softmax head and a residual connection")
R = random.Random(11); g = C.Graph()
x = g.input("x", (5, 12)); w1 = g.weight("w1", [[R.gauss(0, 0.3) for _ in range(10)] for _ in range(12)]); h = g.relu(g.matmul(x, w1))
w2 = g.weight("w2", [[R.gauss(0, 0.3) for _ in range(10)] for _ in range(10)]); y = g.add(g.matmul(h, w2), h)
w3 = g.weight("w3", [[R.gauss(0, 0.3) for _ in range(3)] for _ in range(10)]); p = g.softmax(g.matmul(y, w3, scale=0.5)); g.output(p, "probs"); g.output(g.argmax(p), "class")
print("the graph:"); dump(g)
calib = [{"x": [[R.gauss(0, 1) for _ in range(12)] for _ in range(5)]} for _ in range(30)]; test = {"x": [[R.gauss(0, 1) for _ in range(12)] for _ in range(5)]}
P = C.compile_graph(g, calib); print(f"\ncompiled: {len(P.code)} instructions, scratchpad high-water mark {P.spad_high} of {I.SPAD} words, external memory {P.ext_words} words; notes: {P.notes or 'none'}")
print("scales chosen by the compiler (real value per int8 step):", {f"%{k}": round(v, 5) for k, v in P.scale.items()})
print("the program:"); [print(f" {k:3d} {I.disassemble(i)}") for k, i in enumerate(P.code)]
o = C.run(P, test); ints = C.interpret(P, test); fl = C.evaluate(g, test)
print("\nclass (circuit-model):", [r[0] for r in o["class"][0]], " class (floating point):", [int(r[0]) for r in fl[g.nodes[-1].ins[0]]], f" probabilities: relative error {err(P, o, test)*100:.2f}%")
print("\n== 2. twelve random graphs (seeds 100-111, not in the test battery) on the RTL")
progs = []; meta = []
for seed in range(100, 112):
g, names = T.random_graph(seed); Rr = random.Random(seed); cal = [T.sample_inputs(names, Rr) for _ in range(30)]; tst = T.sample_inputs(names, Rr)
P = C.compile_graph(g, cal); P0 = C.compile_graph(g, cal, reuse=False); out = C.run(P, tst)
ext = C.ext_image(P, tst); progs.append((P.code, ext)); meta.append((seed, len(g.nodes), P, P0.spad_high, err(P, out, tst)))
words = [len(progs)]
for prog, ext in progs: words += G.record(prog, ext)
open(os.path.join(hw.ROOT, "out", "ga2_programs.hex"), "w").write("\n".join("%08x" % (w & 0xFFFFFFFF) for w in words) + "\n")
res = {}
for name, fn in (("Icarus", hw.sim_icarus), ("Verilator", hw.sim_verilator)):
rc, out = fn(F, "ga2_tb", defines=("EXTW=8192", "DUMP")); print(f" {name:10s}{line(out)} (exit {rc})")
for m in re.finditer(r"COUNTERS (\d+) (\d+) (\d+) (\d+) (\d+) (\d+)", out): res[int(m.group(1))] = [int(q) for q in m.groups()[1:]]
print(f"\n{'seed':>5s} {'nodes':>6s} {'instr':>6s} {'scratchpad words (reuse)':>25s} {'(no reuse)':>11s} {'cycles':>8s} {'error vs float':>15s}")
for i, (seed, nn, P, hi0, e) in enumerate(meta): print(f"{seed:5d} {nn:6d} {len(P.code):6d} {P.spad_high:25d} {hi0:11d} {res[i][0]:8d} {e*100:14.2f}%")
print("\n== 3. a network that fits only because buffers are reused: 6 layers of (8 x 32) times (32 x 32)")
g = C.Graph(); Rr = random.Random(5); cur = g.input("x", (8, 32))
for k in range(6): cur = g.relu(g.matmul(cur, g.weight(f"w{k}", [[Rr.gauss(0, 1 / math.sqrt(32)) for _ in range(32)] for _ in range(32)])))
g.output(cur, "y"); cal = [{"x": [[Rr.gauss(0, 1) for _ in range(32)] for _ in range(8)]} for _ in range(10)]
P = C.compile_graph(g, cal); print(f"with reuse: {len(P.code)} instructions, scratchpad high-water mark {P.spad_high} words")
try: C.compile_graph(g, cal, reuse=False); print("without reuse: compiled")
except C.CompileError as e: print("without reuse:", e)
Output (cloud sandbox -- live-executed; Icarus Verilog 12.0, Verilator 5.020)
== 1. one graph, compiled: a two-layer network with a softmax head and a residual connection
the graph:
%0 = input -> (5, 12) x
%1 = weight -> (12, 10) w1
%2 = matmul %0, %1 -> (5, 10)
%3 = relu %2 -> (5, 10)
%4 = weight -> (10, 10) w2
%5 = matmul %3, %4 -> (5, 10)
%6 = add %5, %3 -> (5, 10)
%7 = weight -> (10, 3) w3
%8 = matmul %6, %7 -> (5, 3) x0.5
%9 = softmax %8 -> (5, 3)
%10 = output %9 -> (5, 3) probs
%11 = argmax %9 -> (5, 1)
%12 = output %11 -> (5, 1) class
compiled: 38 instructions, scratchpad high-water mark 280 of 4096 words, external memory 330 words; notes: none
scales chosen by the compiler (real value per int8 step): {'%0': 0.02763, '%1': 0.0078, '%2': 0.03246, '%3': 0.03246, '%4': 0.00822, '%5': 0.03763, '%6': 0.04381, '%7': 0.00561, '%8': 0.0625, '%9': 0.00787, '%10': 0.00787, '%12': 0.01575}
the program:
0 LD dst=0 src=0 len=60
1 LD dst=60 src=60 len=120
2 MM dst=230 A=0 B=60 M=4 K=12 N=4 tb=0 lda=12 ldb=10 ldc=10
3 MM dst=234 A=0 B=64 M=4 K=12 N=4 tb=0 lda=12 ldb=10 ldc=10
4 MM dst=238 A=0 B=68 M=4 K=12 N=2 tb=0 lda=12 ldb=10 ldc=10
5 MM dst=270 A=48 B=60 M=1 K=12 N=4 tb=0 lda=12 ldb=10 ldc=10
6 MM dst=274 A=48 B=64 M=1 K=12 N=4 tb=0 lda=12 ldb=10 ldc=10
7 MM dst=278 A=48 B=68 M=1 K=12 N=2 tb=0 lda=12 ldb=10 ldc=10
8 RQ dst=180 src=230 len=50 m=14251070 s=31 relu=1
9 LD dst=0 src=180 len=100
10 MM dst=230 A=180 B=0 M=4 K=10 N=4 tb=0 lda=10 ldb=10 ldc=10
11 MM dst=234 A=180 B=4 M=4 K=10 N=4 tb=0 lda=10 ldb=10 ldc=10
12 MM dst=238 A=180 B=8 M=4 K=10 N=2 tb=0 lda=10 ldb=10 ldc=10
13 MM dst=270 A=220 B=0 M=1 K=10 N=4 tb=0 lda=10 ldb=10 ldc=10
14 MM dst=274 A=220 B=4 M=1 K=10 N=4 tb=0 lda=10 ldb=10 ldc=10
15 MM dst=278 A=220 B=8 M=1 K=10 N=2 tb=0 lda=10 ldb=10 ldc=10
16 RQ dst=100 src=230 len=50 m=15229501 s=31 relu=0
17 RQ dst=0 src=100 len=50 m=14409190 s=24 relu=0
18 RQ dst=50 src=180 len=50 m=12429377 s=24 relu=0
19 VADD dst=230 src1=0 src2=50 len=50
20 LD dst=0 src=280 len=30
21 MM dst=45 A=230 B=0 M=4 K=10 N=3 tb=0 lda=10 ldb=3 ldc=3
22 MM dst=57 A=270 B=0 M=1 K=10 N=3 tb=0 lda=10 ldb=3 ldc=3
23 RQ dst=30 src=45 len=15 m=8446838 s=32 relu=0
24 SM dst=15 src=30 len=3
25 SM dst=18 src=33 len=3
26 SM dst=21 src=36 len=3
27 SM dst=24 src=39 len=3
28 SM dst=27 src=42 len=3
29 RQ dst=0 src=15 len=15 m=16646144 s=33 relu=0
30 ST src=0 dst=310 len=15
31 AMAX dst=15 src=0 len=3
32 AMAX dst=16 src=3 len=3
33 AMAX dst=17 src=6 len=3
34 AMAX dst=18 src=9 len=3
35 AMAX dst=19 src=12 len=3
36 ST src=15 dst=325 len=5
37 HALT
class (circuit-model): [2, 2, 1, 1, 1] class (floating point): [2, 2, 1, 1, 1] probabilities: relative error 2.00%
== 2. twelve random graphs (seeds 100-111, not in the test battery) on the RTL
Icarus PASS: 12 programs match the Python reference (memory and all five cycle counters), 17022 cycles in total (exit status 0) (exit 0)
Verilator PASS: 12 programs match the Python reference (memory and all five cycle counters), 17022 cycles in total (exit status 0) (exit 0)
seed nodes instr scratchpad words (reuse) (no reuse) cycles error vs float
100 15 20 240 282 830 2.92%
101 21 40 184 325 1497 2.62%
102 13 24 145 329 1124 2.35%
103 11 32 480 684 2659 3.21%
104 10 15 38 47 237 0.00%
105 18 46 203 471 2200 3.71%
106 12 41 276 528 2600 2.37%
107 15 26 221 418 1503 1.52%
108 18 33 175 385 1183 1.62%
109 13 24 64 138 770 1.07%
110 11 21 188 328 1080 2.48%
111 16 27 198 430 1339 1.68%
== 3. a network that fits only because buffers are reused: 6 layers of (8 x 32) times (32 x 32)
with reuse: 111 instructions, scratchpad high-water mark 1792 words
without reuse: scratchpad is full: cannot place a buffer of 256 words (free: 0 words in 0 pieces)
Reading it¶
- Section 1 shows the whole pipeline on a two-layer network with a residual connection and a softmax head. The compiler chose a scale for every tensor, fused the
reluinto the first matmul's requantizer, inserted two rescales (instructions 17 and 18) for the residualadd, whose operands (scales 0.0376 and 0.0325) differ from the sum's (0.0438), emitted 38 instructions in all, and needed 280 of the 4,096 scratchpad words. The class decision agrees with floating point on all five rows, and the probabilities are 2.0% off. - Section 2 compiles twelve random graphs the test battery has never seen, puts them on the RTL, and compares with the reference simulator: identical memory and cycle counters, in both simulators. Their error against floating point is 0-3.7%.
- Allocation matters. On those graphs, reuse cuts the scratchpad footprint to 43-85% of what keeping every buffer would need. In section 3 a six-layer network of 32 x 32 layers compiles to 111 instructions in 1,792 words with reuse, and fails without it ("scratchpad is full"): its weights alone (6 x 1,024 words) exceed the memory.
Running example A: compile one attention step, pass by pass¶
The point of this example: to see every decision the compiler makes, in order, on a graph small enough to follow. The graph is a single decode-step attention head (Chapter 8's computation): a query against a key cache plus a new key, a softmax, and a weighted sum of a value cache plus a new value. The script prints what each stage decided, records the allocator's events as it runs, and checks the result.
#!/usr/bin/env python3
"""Chapter 11, running example A: compile one attention step pass by pass. The graph is a single decode-step attention (scores of a query against a key cache plus the new key, softmax, weighted sum of the value cache plus the new value). The script prints what every stage of Capra decided: the graph, the calibrated ranges, the int8 scales (and which are pinned), the requantizer parameters, the tiles, the scratchpad allocation over time (recorded from the allocator itself), the emitted program, and then checks the compiled program against the integer interpreter and floating point. Writes out/ch11_example_a.json. Usage: ch11_example_a.py"""
import json, math, os, random, sys
os.environ["GA2_EXT"] = "8192"
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))); import hw
sys.path.insert(0, os.path.join(hw.ROOT, "model")); import ga2_isa as I, capra as C
D, L = 8, 5; R = random.Random(3)
g = C.Graph(); q = g.input("q", (1, D)); kc = g.input("kc", (L, D)); kn = g.input("kn", (1, D)); K = g.concat_rows(kc, kn)
s = g.matmul(q, K, transpose_b=True, scale=1 / math.sqrt(D)); p = g.softmax(s); vc = g.input("vc", (L, D)); vn = g.input("vn", (1, D)); V = g.concat_rows(vc, vn)
o = g.matmul(p, V); g.output(o, "out")
def sample(): return {"q": [[R.gauss(0, 1) for _ in range(D)]], "kc": [[R.gauss(0, 1) for _ in range(D)] for _ in range(L)], "kn": [[R.gauss(0, 1) for _ in range(D)]], "vc": [[R.gauss(0, 1) for _ in range(D)] for _ in range(L)], "vn": [[R.gauss(0, 1) for _ in range(D)]]}
calib = [sample() for _ in range(20)]; test = sample()
print("== PASS 0: the graph (what the user writes)")
names = {n.id: n.attrs.get("name", "") for n in g.nodes}
for n in g.nodes: print(f" %{n.id:<2d} = {n.op:8s} {', '.join('%' + str(i) for i in n.ins):8s} -> {n.shape} {names[n.id]}" + (" ^T" if n.attrs.get("transpose_b") else "") + (f" x{n.attrs['scale']:.4f}" if n.op == "matmul" and n.attrs["scale"] != 1 else ""))
print("\n== PASS 1: calibration: run the graph in floating point on 20 sample sets; each tensor's range is the largest |value| seen")
rng = {}
for c in calib:
for nid, m in C.evaluate(g, c).items(): rng[nid] = max(rng.get(nid, 0.0), max(abs(x) for r in m for x in r))
for n in g.nodes: print(f" %{n.id:<2d} {n.op:8s} range {rng[n.id]:8.4f}")
print("\n== PASS 2: scale planning: one int8 scale per tensor (range / 127), with groups and pinned scales")
P = C.compile_graph(g, calib)
for n in g.nodes:
if n.op == "argmax": continue
pin = " <- pinned: softmax input must be Q4.4 (1/16)" if n.id == s else " <- pinned: softmax output is a probability, 1/127" if n.id == p else " <- shares a scale with its concat result (same buffer)" if n.id in (kc, kn, vc, vn) else ""
print(f" %{n.id:<2d} {n.op:8s} scale {P.scale[n.id]:.6f} (range/127 = {rng[n.id]/127:.6f}){pin}")
print("\n== PASS 3: requantizer parameters: M = factor x scale_a x scale_b / scale_out, written as mantissa / 2^shift (Chapter 3)")
for k, (m_, s_) in P.params.items(): print(f" node {k}: M ~ {m_} / 2^{s_} = {m_ / 2**s_:.8f}")
print("\n== PASS 4: tiling: each matmul is cut into 4x4 output tiles, one MM instruction each")
for n in g.nodes:
if n.op != "matmul": continue
(m, k) = g.nodes[n.ins[0]].shape; nn = n.shape[1]; tiles = [(i, j) for i in range(0, m, 4) for j in range(0, nn, 4)]
print(f" %{n.id}: ({m} x {k}) x ({k} x {nn}) -> {math.ceil(m/4)} x {math.ceil(nn/4)} = {len(tiles)} tile(s):", ", ".join(f"[rows {i}-{min(i+3,m-1)}, cols {j}-{min(j+3,nn-1)}]" for i, j in tiles))
print("\n== PASS 5: allocation: scratchpad words over time, recorded from the allocator itself, numbered in the order the events happen")
events = []; code_len = lambda: None
orig_alloc, orig_rel = C.Allocator.alloc, C.Allocator.release
cur = {"code": None}
def alloc(self, n): a = orig_alloc(self, n); events.append(("alloc", a, n, len(cur["code"]) if cur["code"] is not None else 0)); return a
def rel(self, s_, n): orig_rel(self, s_, n); events.append(("free", s_, n, len(cur["code"]) if cur["code"] is not None else 0))
# the compiler builds `code` as a local list; to timestamp events we count them in order instead
C.Allocator.alloc = alloc; C.Allocator.release = rel
P1 = C.compile_graph(g, calib); C.Allocator.alloc, C.Allocator.release = orig_alloc, orig_rel
print(" events in order (the allocator is first-fit with coalescing; a buffer is released after its last consumer):")
live = {}; blocks = []
for k, (kind, a, n, _) in enumerate(events):
if kind == "alloc": live[(a, n)] = k; print(f" {k:2d}: alloc {n:4d} words at {a:4d} .. {a+n-1:4d}")
else: print(f" {k:2d}: free {n:4d} words at {a:4d}"); blocks.append((a, n, live.pop((a, n), 0), k))
for (a, n), k0 in live.items(): blocks.append((a, n, k0, len(events)))
print(f" scratchpad high-water mark with reuse: {P1.spad_high} words; without reuse: {C.compile_graph(g, calib, reuse=False).spad_high} words (of {I.SPAD})")
print(" note: kc and kn are VIEWS of one buffer (the concat result), as are vc and vn: the cache plus the new row is assembled in place, so the concatenation costs no instruction.")
print("\n== PASS 6: emission: the program")
for k, ins in enumerate(P.code): print(f" {k:3d} {I.disassemble(ins)}")
print(f"\n{len(P.code)} instructions (including HALT); cycles on GA-2 (model): {I.cycle_counts(P.code)[0]}")
print("\n== checks: the compiled program against the integer interpreter, and against floating point")
out = C.run(P, test); ints = C.interpret(P, test); fl = C.evaluate(g, test); got = out["out"][1]; ref = fl[o]
num = math.sqrt(sum((a - b) ** 2 for r, t in zip(got, ref) for a, b in zip(r, t))); den = math.sqrt(sum(b * b for t in ref for b in t))
exact = out["out"][0] == ints[g.nodes[o].id] if False else out["out"][0] == ints[[n.id for n in g.nodes if n.op == "output"][0]]
print(f" simulator output equals the interpreter's integers: {exact}; relative error against floating point: {100*num/den:.2f}%")
print(" stage-wise real-number check (every tensor within one level of the meaning of its operation):", C.stagewise_problems(P, ints) or "no problems")
print(" output (float):", [round(v, 3) for v in got[0]]); print(" reference :", [round(v, 3) for v in ref[0]])
json.dump({"events": events, "blocks": blocks, "high": P1.spad_high, "code": [I.disassemble(i) for i in P.code], "ops": [i["op"] for i in P.code], "scales": {str(k): v for k, v in P.scale.items()}, "nodes": [(n.id, n.op, names[n.id]) for n in g.nodes]}, open(os.path.join(hw.ROOT, "out", "ch11_example_a.json"), "w"))
To compile and run: python3 tools/ch11_example_a.py (pure Python; instant).
Output (cloud sandbox -- live-executed)
== PASS 0: the graph (what the user writes)
%0 = input -> (1, 8) q
%1 = input -> (5, 8) kc
%2 = input -> (1, 8) kn
%3 = concat %1, %2 -> (6, 8)
%4 = matmul %0, %3 -> (1, 6) ^T x0.3536
%5 = softmax %4 -> (1, 6)
%6 = input -> (5, 8) vc
%7 = input -> (1, 8) vn
%8 = concat %6, %7 -> (6, 8)
%9 = matmul %5, %8 -> (1, 8)
%10 = output %9 -> (1, 8) out
== PASS 1: calibration: run the graph in floating point on 20 sample sets; each tensor's range is the largest |value| seen
%0 input range 2.5978
%1 input range 3.1547
%2 input range 2.8494
%3 concat range 3.1547
%4 matmul range 2.5511
%5 softmax range 0.7433
%6 input range 3.6058
%7 input range 3.0773
%8 concat range 3.6058
%9 matmul range 2.2089
%10 output range 2.2089
== PASS 2: scale planning: one int8 scale per tensor (range / 127), with groups and pinned scales
%0 input scale 0.020455 (range/127 = 0.020455)
%1 input scale 0.024840 (range/127 = 0.024840) <- shares a scale with its concat result (same buffer)
%2 input scale 0.024840 (range/127 = 0.022436) <- shares a scale with its concat result (same buffer)
%3 concat scale 0.024840 (range/127 = 0.024840)
%4 matmul scale 0.062500 (range/127 = 0.020087) <- pinned: softmax input must be Q4.4 (1/16)
%5 softmax scale 0.007874 (range/127 = 0.005853) <- pinned: softmax output is a probability, 1/127
%6 input scale 0.028392 (range/127 = 0.028392) <- shares a scale with its concat result (same buffer)
%7 input scale 0.028392 (range/127 = 0.024231) <- shares a scale with its concat result (same buffer)
%8 concat scale 0.028392 (range/127 = 0.028392)
%9 matmul scale 0.017393 (range/127 = 0.017393)
%10 output scale 0.017393 (range/127 = 0.017393)
== PASS 3: requantizer parameters: M = factor x scale_a x scale_b / scale_out, written as mantissa / 2^shift (Chapter 3)
node 4: M ~ 12345120 / 2^32 = 0.00287432
node 5: M ~ 16646144 / 2^33 = 0.00193787
node 9: M ~ 13800986 / 2^30 = 0.01285317
== PASS 4: tiling: each matmul is cut into 4x4 output tiles, one MM instruction each
%4: (1 x 8) x (8 x 6) -> 1 x 2 = 2 tile(s): [rows 0-0, cols 0-3], [rows 0-0, cols 4-5]
%9: (1 x 6) x (6 x 8) -> 1 x 2 = 2 tile(s): [rows 0-0, cols 0-3], [rows 0-0, cols 4-7]
== PASS 5: allocation: scratchpad words over time, recorded from the allocator itself, numbered in the order the events happen
events in order (the allocator is first-fit with coalescing; a buffer is released after its last consumer):
0: alloc 48 words at 0 .. 47
1: alloc 8 words at 48 .. 55
2: alloc 6 words at 56 .. 61
3: alloc 6 words at 62 .. 67
4: free 6 words at 62
5: free 48 words at 0
6: free 8 words at 48
7: alloc 6 words at 0 .. 5
8: alloc 6 words at 6 .. 11
9: free 6 words at 6
10: free 6 words at 56
11: alloc 48 words at 6 .. 53
12: alloc 8 words at 54 .. 61
13: alloc 8 words at 62 .. 69
14: free 8 words at 62
15: free 6 words at 0
16: free 48 words at 6
17: free 8 words at 54
scratchpad high-water mark with reuse: 70 words; without reuse: 132 words (of 4096)
note: kc and kn are VIEWS of one buffer (the concat result), as are vc and vn: the cache plus the new row is assembled in place, so the concatenation costs no instruction.
== PASS 6: emission: the program
0 LD dst=0 src=8 len=40
1 LD dst=40 src=48 len=8
2 LD dst=48 src=0 len=8
3 MM dst=62 A=48 B=0 M=1 K=8 N=4 tb=1 lda=8 ldb=8 ldc=6
4 MM dst=66 A=48 B=32 M=1 K=8 N=2 tb=1 lda=8 ldb=8 ldc=6
5 RQ dst=56 src=62 len=6 m=12345120 s=32 relu=0
6 SM dst=6 src=56 len=6
7 RQ dst=0 src=6 len=6 m=16646144 s=33 relu=0
8 LD dst=6 src=56 len=40
9 LD dst=46 src=96 len=8
10 MM dst=62 A=0 B=6 M=1 K=6 N=4 tb=0 lda=6 ldb=8 ldc=8
11 MM dst=66 A=0 B=10 M=1 K=6 N=4 tb=0 lda=6 ldb=8 ldc=8
12 RQ dst=54 src=62 len=8 m=13800986 s=30 relu=0
13 ST src=54 dst=104 len=8
14 HALT
15 instructions (including HALT); cycles on GA-2 (model): 484
== checks: the compiled program against the integer interpreter, and against floating point
simulator output equals the interpreter's integers: True; relative error against floating point: 2.36%
stage-wise real-number check (every tensor within one level of the meaning of its operation): no problems
output (float): [0.244, 0.209, 0.052, 0.244, -0.974, 1.078, 0.122, 0.417]
reference : [0.239, 0.21, 0.044, 0.22, -0.998, 1.092, 0.125, 0.413]
Figure 11.2: the graph the user writes in ten lines.
Walkthrough, pass by pass.
- Pass 0, the graph. Eleven nodes. Node %4 is a matmul with
^T(the key matrix is used transposed) and a constant factor 0.3536 = 1/sqrt(8), the attention scaling. The twoconcatnodes append the new key and value to the caches. - Pass 1, calibration. The graph is evaluated in floating point on 20 sample input sets, and each tensor's largest absolute value is recorded: the scores reach 2.55, the softmax output 0.74, the final output 2.21.
- Pass 2, scales. Normally scale = range / 127. Three things are different here. The scores (%4) are pinned to 1/16 because the softmax unit's input format is Q4.4: that scale is 0.0625, not 0.0201. The softmax output (%5) is pinned to 1/127 (a probability). And
kc,knandkc's concatenation share one scale (0.02484, the largest of the three ranges), as dovc,vnand their concat: a buffer holds one scale. That is whyknshows range/127 = 0.0224 but scale 0.0248. - Pass 3, requantizer parameters. For the first matmul
M = 0.3536 x scale_q x scale_K / 0.0625 = 0.002874, stored as the integer pair (12345120, 32), i.e. 12345120 / 2^32. Each requantizing node has such a pair. - Pass 4, tiling. Both matmuls have one row, so each becomes two
MMinstructions: columns 0-3 and the remaining columns (2 or 4 wide). A real model's matmuls would have many more tiles (Figure 11.4). - Pass 5, allocation. The allocator events show first-fit at work. The first four allocations (48, 8, 6 and 6 words) are the key buffer (cache plus new key, one buffer), the query, the scores, and the matmul's int32 accumulator. Once the scores are computed the key buffer and the query are released, and later the value buffer (48 words) is placed at address 6, on addresses the key buffer had used. The peak is 70 words instead of 132 without reuse: a 47% cut on a graph with no waste to speak of.
- Pass 6, emission. 15 instructions.
LDbrings in the cache and the new key (the cache and the new row are loaded into one buffer, as the figure below shows), twoMMtiles withtb=1compute the scores, oneRQbrings them to Q4.4,SMis the softmax,RQmakes int8 weights, a secondLDbrings the value cache, two moreMMtiles do the weighted sum, anRQand anSTfinish. On GA-2 it takes 484 cycles. - Checks. The simulator's output is exactly the integer interpreter's; the stage-wise real-number check finds no problems; the error against floating point is 2.4%.
Figure 11.3: the scales. The two pinned tensors (the scores at 1/16 = 0.0625 and the softmax output at 1/127 = 0.0079) stand out; the others follow their ranges.
Figure 11.4: tiling a larger product (the five-row product of Section 1's graph). Each outlined tile is one MM instruction; the edge tiles are smaller.
Figure 11.5: a concatenation is a view. Both parts are written into the same buffer at the right row offsets, so the concatenation costs no instruction.
Later chapters extend the IR. Chapter 13 adds weight(..., bits=4) (packed int4 weights) and Chapter 16 adds slice_rows, the reverse of concat_rows: a view of some rows of a tensor, with no instruction and no copy, which is how a causal mask is built without a mask operation.
Figure 11.6: the allocator's events drawn as a memory map. The key-side buffers (left) are freed before the value-side buffers (right) are placed, partly on top of the same addresses.
Running example B: write your own graph¶
The point of this example: the compiler is only useful if you can give it your own network. tools/ch11_example_b.py is a template: the function build makes a residual network with a softmax head; edit it to describe yours. The script compiles it, runs it on the reference simulator and on the RTL in both simulators, compares with floating point, sweeps the hidden width to see how footprint, cycles and accuracy move, and shows the exact messages for the graphs Capra refuses.
#!/usr/bin/env python3
"""Chapter 11, running example B: write your own graph. A template for a small residual network (the function `build` is the part to edit): it is compiled with Capra, run on the reference simulator AND on the RTL in both simulators, compared with floating point; then the hidden width is swept to see footprint, cycles and accuracy; and the five kinds of graph the compiler REFUSES are shown with the exact error messages and the fix for each. Writes out/ch11_example_b.json. Usage: ch11_example_b.py"""
import json, math, os, random, re, sys
os.environ["GA2_EXT"] = "8192"
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))); import hw
sys.path.insert(0, os.path.join(hw.ROOT, "model")); import ga2_isa as I, ga2_progs as G, capra as C
F = ["rtl/sram.v", "rtl/dma.v", "rtl/systolic.v", "rtl/requant.v", "rtl/exp_lut.v", "rtl/divu.v", "rtl/rowmem.v", "rtl/ga2_vec.v", "rtl/ga2_sm.v", "rtl/ga2_mm.v", "rtl/ga2.v", "tb/extmem_rw.v", "tb/ga2_tb.v"]
def build(batch=4, d_in=16, hidden=32, d_out=8, seed=1):
"""THE PART TO EDIT. A two-layer network with a residual connection and a softmax head; returns (graph, names of inputs)."""
R = random.Random(seed); w = lambda r, c: [[R.gauss(0, 1 / math.sqrt(r)) for _ in range(c)] for _ in range(r)]
g = C.Graph(); x = g.input("x", (batch, d_in)); h = g.relu(g.matmul(x, g.weight("w1", w(d_in, hidden))))
r = g.add(g.matmul(h, g.weight("w2", w(hidden, hidden))), h); p = g.softmax(g.matmul(r, g.weight("w3", w(hidden, d_out)), scale=0.5))
g.output(p, "probs"); g.output(g.argmax(p), "class"); return g, {"x": (batch, d_in)}
def inputs(shapes, R): return {k: [[R.gauss(0, 1) for _ in range(c)] for _ in range(r)] for k, (r, c) in shapes.items()}
def relerr(P, out, test):
fl = C.evaluate(P.graph, test); e = [e for e in P.ext.values() if e["kind"] == "output" and P.graph.nodes[P.graph.nodes[e["node"]].ins[0]].op != "argmax"][0]
ref = fl[e["node"]]; got = out["probs"][1]; return math.sqrt(sum((a - b) ** 2 for r, t in zip(got, ref) for a, b in zip(r, t))) / math.sqrt(sum(b * b for t in ref for b in t))
print("== 1. the default graph (batch 4, 16 inputs, 32 hidden, 8 classes): compile, run, compare")
g, shapes = build(); R = random.Random(9); cal = [inputs(shapes, R) for _ in range(30)]; test = inputs(shapes, R)
P = C.compile_graph(g, cal); out = C.run(P, test); fl = C.evaluate(g, test); cls_f = [int(r[0]) for r in fl[g.nodes[-1].ins[0]]]
print(f" {len(P.code)} instructions, scratchpad high-water {P.spad_high} words, {I.cycle_counts(P.code)[0]} cycles on GA-2 (model)")
print(" class (compiled):", [r[0] for r in out["class"][0]], " class (floating point):", cls_f, f" probabilities relative error {100*relerr(P, out, test):.2f}%")
words = [1] + G.record(P.code, C.ext_image(P, test)); open(os.path.join(hw.ROOT, "out", "ga2_programs.hex"), "w").write("\n".join("%08x" % (w & 0xFFFFFFFF) for w in words) + "\n")
for name, fn in (("Icarus", hw.sim_icarus), ("Verilator", hw.sim_verilator)):
rc, o = fn(F, "ga2_tb", defines=("EXTW=8192",)); print(f" RTL in {name:10s}", ([l for l in o.splitlines() if l.startswith(("PASS", "FAIL", "MISMATCH"))] or [o[-120:]])[0], f"(exit {rc})")
print("\n== 2. sweep the hidden width: footprint, cycles and accuracy (reference simulator + cycle model; the K <= 64 limit shows at 64 and above)")
print(" hidden instr spad words (reuse / no reuse) cycles error vs float class agreement")
res = {"hidden": [], "instr": [], "high": [], "high0": [], "cycles": [], "err": []}
notes = []
for hdn in (8, 16, 32, 48, 64):
g2, sh = build(hidden=hdn); Rr = random.Random(hdn); cal2 = [inputs(sh, Rr) for _ in range(30)]; t2 = inputs(sh, Rr)
try: P2 = C.compile_graph(g2, cal2)
except C.CompileError as e:
print(f" {hdn:6d} does not compile: {e}"); res["hidden"].append(hdn); res["instr"].append(None); res["high"].append(None); res["high0"].append(None); res["cycles"].append(None); res["err"].append(None); continue
try: h0 = C.compile_graph(g2, cal2, reuse=False).spad_high
except C.CompileError: h0 = None
o2 = C.run(P2, t2); f2 = C.evaluate(g2, t2); pf = f2[g2.nodes[-3].id] if False else None
cf = [int(r[0]) for r in f2[g2.nodes[-1].ins[0]]]; cc = [r[0] for r in o2["class"][0]]; ok = cf == cc
cyc = I.cycle_counts(P2.code)[0]; e = relerr(P2, o2, t2); res["hidden"].append(hdn); res["instr"].append(len(P2.code)); res["high"].append(P2.spad_high); res["high0"].append(h0); res["cycles"].append(cyc); res["err"].append(e)
if not ok:
probs = [n for n in g2.nodes if n.op == "softmax"][0].id; i = [k for k in range(len(cf)) if cf[k] != cc[k]][0]; row = f2[probs][i]; top = sorted(row, reverse=True)[:2]
notes.append(f"hidden {hdn}: sample {i} gets class {cc[i]} instead of {cf[i]}; the two largest floating-point probabilities are {top[0]:.3f} and {top[1]:.3f}, a gap of {top[0]-top[1]:.3f}")
print(f" {hdn:6d} {len(P2.code):5d} {P2.spad_high:12d} / {('full' if h0 is None else h0):>6} {cyc:6d} {100*e:9.2f}% {'all agree' if ok else 'DIFFER (see below)'}")
for n_ in notes: print(" note:", n_)
if notes: print(" (a class flips only where two probabilities are nearly tied: a 2% error in the probabilities can reorder a 0.001 gap. The probabilities themselves stay within about 2-3% in every row.)")
print("\n== 3. what Capra refuses, with the exact message and the fix")
def tryit(title, fix, f):
try: f(); print(f" {title}: compiled (unexpected)")
except C.CompileError as e: print(f" {title}\n message: {e}\n fix: {fix}")
def g_k65():
g = C.Graph(); x = g.input("x", (2, 65)); g.output(g.matmul(x, g.weight("w", [[0.1] * 4 for _ in range(65)])), "y"); C.compile_graph(g, [{"x": [[0.5] * 65 for _ in range(2)]}])
def g_relu():
g = C.Graph(); x = g.input("x", (2, 4)); y = g.relu(g.add(x, x)); g.output(y, "y"); C.compile_graph(g, [{"x": [[0.5] * 4 for _ in range(2)]}])
def g_mismatch(): g = C.Graph(); a = g.input("a", (2, 3)); b = g.weight("b", [[0.1] * 4 for _ in range(4)]); g.matmul(a, b)
def g_big():
g = C.Graph(); x = g.input("x", (8, 64)); g.output(g.matmul(x, g.weight("w", [[0.01] * 64 for _ in range(64)])), "y"); C.compile_graph(g, [{"x": [[0.5] * 64 for _ in range(8)]}])
tryit("a product with inner dimension 65", "split the contraction into pieces of at most 64 (the hardware has no int32 accumulate instruction), or reduce the width", g_k65)
tryit("a relu that is not directly after a matmul", "put the relu directly after a matmul that has no other consumer (the compiler fuses it into the requantizer)", g_relu)
tryit("matrix shapes that do not match", "check the shapes; transpose_b=True multiplies by the transpose of the second operand", g_mismatch)
tryit("tensors too big for the 4,096-word scratchpad (a 64 x 64 weight is the whole of it)", "use smaller tensors or fewer live at once (the allocator already reuses dead buffers); a bigger model must be split into pieces that fit", g_big)
json.dump(res, open(os.path.join(hw.ROOT, "out", "ch11_example_b.json"), "w"))
To compile and run: python3 tools/ch11_example_b.py (about a minute; it runs the RTL twice).
Output (cloud sandbox -- live-executed; Icarus Verilog 12.0, Verilator 5.020)
== 1. the default graph (batch 4, 16 inputs, 32 hidden, 8 classes): compile, run, compare
40 instructions, scratchpad high-water 1728 words, 7740 cycles on GA-2 (model)
class (compiled): [2, 2, 1, 2] class (floating point): [2, 2, 1, 2] probabilities relative error 2.92%
RTL in Icarus PASS: 1 programs match the Python reference (memory and all five cycle counters), 7740 cycles in total (exit status 0) (exit 0)
RTL in Verilator PASS: 1 programs match the Python reference (memory and all five cycle counters), 7740 cycles in total (exit status 0) (exit 0)
== 2. sweep the hidden width: footprint, cycles and accuracy (reference simulator + cycle model; the K <= 64 limit shows at 64 and above)
hidden instr spad words (reuse / no reuse) cycles error vs float class agreement
8 28 256 / 512 1800 1.91% DIFFER (see below)
16 32 448 / 992 3364 2.63% all agree
32 40 1728 / 2336 7740 1.88% all agree
48 48 3328 / full 13780 2.44% all agree
64 does not compile: scratchpad is full: cannot place a buffer of 4096 words (free: 3840 words in 2 pieces)
note: hidden 8: sample 2 gets class 1 instead of 6; the two largest floating-point probabilities are 0.166 and 0.165, a gap of 0.001
(a class flips only where two probabilities are nearly tied: a 2% error in the probabilities can reorder a 0.001 gap. The probabilities themselves stay within about 2-3% in every row.)
== 3. what Capra refuses, with the exact message and the fix
a product with inner dimension 65
message: matmul node 2: inner dimension 65 exceeds the matrix unit's limit of 64
fix: split the contraction into pieces of at most 64 (the hardware has no int32 accumulate instruction), or reduce the width
a relu that is not directly after a matmul
message: relu on node 1 (add) cannot be fused into a requantizer
fix: put the relu directly after a matmul that has no other consumer (the compiler fuses it into the requantizer)
matrix shapes that do not match
message: matmul shapes (2, 3) x (4, 4)
fix: check the shapes; transpose_b=True multiplies by the transpose of the second operand
tensors too big for the 4,096-word scratchpad (a 64 x 64 weight is the whole of it)
message: scratchpad is full: cannot place a buffer of 4096 words (free: 3584 words in 1 pieces)
fix: use smaller tensors or fewer live at once (the allocator already reuses dead buffers); a bigger model must be split into pieces that fit
Figure 11.7 (measured): footprint against width. At width 48 the program needs 3,328 words with reuse; without reuse it does not fit.
Figure 11.8 (cycle model, equal to the circuit's counters): cycles grow faster than linearly with width, since the weight matrices grow with its square.
Walkthrough.
- The default graph. 40 instructions, 1,728 scratchpad words, 7,740 cycles. The compiled class decisions agree with floating point, the probabilities are 2.9% off, and the RTL in both simulators matches the reference simulator's memory and cycle counters exactly (7,740 cycles). The program that no human wrote runs on the chip.
- The width sweep. With hidden width 8, 16, 32 and 48 the program grows from 28 to 48 instructions, and the scratchpad need from 256 to 3,328 words with reuse. Without reuse the 48-wide network (4,096 words or more) does not fit at all. At width 64 even reuse does not help: a single 64 x 64 weight is the whole scratchpad, and Capra reports it.
- A class that flipped. At hidden width 8 one sample gets class 1 instead of 6. The two largest floating-point probabilities for that sample are 0.166 and 0.165, a gap of 0.001: with an error of about 2% in the probabilities such a near-tie can reorder. The decision is unstable in the floating-point model too; nothing is wrong with the compiler. This is why accuracy is judged on probabilities, not on argmax, in check (e).
- Refusals. Four of the five things the compiler will not do are shown with the exact messages: an inner dimension of 65 (the hardware has no instruction to accumulate across matrix-unit passes), a relu that cannot be fused into a requantizer, mismatched shapes, and a graph whose tensors do not fit the scratchpad. Each message names the cause, and each has a fix. A compiler that refuses clearly is much better than one that emits a wrong program.
How a compiler is tested¶
Figure 11.9: the battery. Each check sees a different kind of bug; the table of mutants shows what each one catches.
A compiler's output is a program, and "does the program run" is not the question; the question is whether it computes the right thing for every graph. The battery in model/capra_tests.py generates random graphs (random shapes, matmul chains, fused relus, self-products, softmax heads, residual adds, concatenations, argmax outputs), compiles each twice (with and without buffer reuse), and applies five checks:
#!/usr/bin/env python3
"""Chapter 11: Capra's test battery. A generator of random graphs and the checks applied to each:
(a) ALLOCATOR: the program compiled with buffer reuse and the one compiled with every buffer kept alive give identical outputs;
(b) MEANING: the outputs of the reference simulator equal the integer interpreter's, exactly;
(c) EVERY TENSOR: with reuse off, every tensor read back from the scratchpad equals the interpreter's;
(d) SCALES: every tensor is within one level of the real-number meaning of its operation (stagewise_problems);
(e) ACCURACY: the dequantized outputs stay within a bound of the floating-point result on inputs NOT used for calibration.
Usage: capra_tests.py [nseeds]"""
import math, os, random, sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import capra as C
def random_graph(seed, wbits=8):
R = random.Random(seed); g = C.Graph(); rows = R.randrange(1, 7); cols = R.randrange(2, 20); names = {}
def new_input(shape, sd=1.0): nm = f"in{len(names)}"; names[nm] = (shape, sd); return g.input(nm, shape)
def std_of(t): # typical magnitude of a tensor, measured by running the graph so far; a partner input in an add or concat is generated at the same magnitude
v = C.evaluate(g, sample_inputs(names, random.Random(0)))[t]; flat = [x for r in v for x in r]; return max(1e-3, math.sqrt(sum(x * x for x in flat) / len(flat)))
def weight(r, c): return g.weight(f"w{len(g.nodes)}", [[R.gauss(0, 1 / math.sqrt(r)) for _ in range(c)] for _ in range(r)], bits=wbits) # wbits=4: Chapter 13's int4 weights
cur = new_input((rows, cols)); pool = [cur]; in_concat = set(); outputs = []; used_scores = False
for _ in range(R.randrange(3, 8)):
shape = g.nodes[cur].shape; kind = R.choice(["matmul", "matmul", "matmul_relu", "scores", "softmax", "add", "concat"])
if kind in ("matmul", "matmul_relu"):
cur = g.matmul(cur, weight(shape[1], R.randrange(1, 11)))
if kind == "matmul_relu": cur = g.relu(cur)
elif kind == "scores" and shape[1] <= 64 and not used_scores: # at most one self-product: repeated squaring spans orders of magnitude, which no static per-tensor scale can serve
used_scores = True; cur = g.matmul(cur, cur, transpose_b=True, scale=1 / math.sqrt(shape[1]))
elif kind == "softmax":
cur = g.softmax(g.matmul(cur, weight(shape[1], R.randrange(2, 9)), scale=0.5))
elif kind == "add":
same = [p for p in pool if g.nodes[p].shape == shape and p != cur]; other = R.choice(same) if same else new_input(shape, std_of(cur)); cur = g.add(cur, other)
elif kind == "concat" and shape[0] <= 9 and cur not in in_concat:
extra = new_input((R.randrange(1, 4), shape[1]), std_of(cur)); in_concat.update([cur, extra]); cur = g.concat_rows(cur, extra)
pool.append(cur)
if R.random() < 0.3: g.output(g.argmax(cur), "tok")
g.output(cur, "out")
if len(pool) > 2 and R.random() < 0.3 and pool[1] not in in_concat: g.output(pool[1], "mid")
return g, names
def sample_inputs(names, R): return {nm: [[R.gauss(0, sd) for _ in range(c)] for _ in range(r)] for nm, ((r, c), sd) in names.items()}
def check_graph(seed, bound=0.16, wbits=8, build=None):
problems = []; g, names = (build or random_graph)(seed, wbits); R = random.Random(seed * 7 + 1)
calib = [sample_inputs(names, R) for _ in range(30)]; test = sample_inputs(names, R)
try: P1 = C.compile_graph(g, calib, reuse=True); P0 = C.compile_graph(g, calib, reuse=False)
except C.CompileError as e: return [f"seed {seed}: compile error: {e}"], None
o1 = C.run(P1, test); o0, m0 = C.run(P0, test, machine=True); ints = C.interpret(P0, test)
if {k: v[0] for k, v in o1.items()} != {k: v[0] for k, v in o0.items()}: problems.append(f"seed {seed}: (a) reuse and no-reuse programs give different outputs")
for nm, e in P0.ext.items():
if e["kind"] == "output":
if o0[nm][0] != ints[P0.graph.nodes[e["node"]].id]: problems.append(f"seed {seed}: (b) output {nm} differs from the interpreter")
for n in g.nodes:
if n.op in ("output",): continue
try:
if C.read_tensor(P0, m0, n.id) != ints[n.id]: problems.append(f"seed {seed}: (c) tensor {n.id} ({n.op}) read back from the scratchpad differs from the interpreter")
except ValueError: pass
problems += [f"seed {seed}: (d) {p}" for p in C.stagewise_problems(P0, ints)]
fl = C.evaluate(g, test); worst = 0.0
for nm, e in P0.ext.items():
if e["kind"] != "output" or g.nodes[g.nodes[e["node"]].ins[0]].op == "argmax": continue
ref = fl[e["node"]]; got = o0[nm][1]; num = math.sqrt(sum((a - b) ** 2 for r, s in zip(got, ref) for a, b in zip(r, s))); den = math.sqrt(sum(b * b for s in ref for b in s)) or 1.0
worst = max(worst, num / den)
if worst > bound: problems.append(f"seed {seed}: (e) error {worst:.3f} over the bound {bound}")
return problems, worst
def random_slice_graph(seed, wbits=8):
"""Chapter 16: graphs built around slice_rows: slices of a product, of a concatenation and of a slice, feeding matmuls, softmaxes, adds and argmax."""
R = random.Random(seed); g = C.Graph(); rows = R.randrange(4, 10); cols = R.randrange(3, 14); names = {}
def new_input(shape): nm = f"in{len(names)}"; names[nm] = (shape, 1.0); return g.input(nm, shape)
def weight(r, c): return g.weight(f"w{len(g.nodes)}", [[R.gauss(0, 1 / math.sqrt(r)) for _ in range(c)] for _ in range(r)], bits=wbits)
x = new_input((rows, cols)); y = g.matmul(x, weight(cols, R.randrange(3, 9)))
extra = new_input((R.randrange(1, 4), g.nodes[y].shape[1])); cat = g.concat_rows(y, extra); total = g.nodes[cat].shape[0]; wsl = weight(R.randrange(3, 8), R.randrange(2, 6)); sources = [(cat, total), (y, rows), (x, rows), (wsl, g.nodes[wsl].shape[0])]; outs = []
for k in range(R.randrange(2, 5)):
src, r = R.choice(sources); start = R.randrange(0, r); count = R.randrange(1, r - start + 1); sl = g.slice_rows(src, start, count)
if R.random() < 0.3 and count > 1: sl2 = g.slice_rows(sl, R.randrange(0, count), 1) if False else g.slice_rows(sl, 0, R.randrange(1, count + 1)); sl = sl2
c = g.nodes[sl].shape[1]; kind = R.choice(["matmul", "softmax", "scores", "argmax"])
if kind == "matmul": o = g.matmul(sl, weight(c, R.randrange(2, 7)))
elif kind == "softmax": o = g.softmax(g.matmul(sl, weight(c, R.randrange(2, 7)), scale=0.5))
elif kind == "scores": o = g.matmul(sl, g.slice_rows(src, 0, r), transpose_b=True, scale=1 / math.sqrt(c))
else: o = g.argmax(g.matmul(sl, weight(c, R.randrange(2, 7))))
g.output(o, f"o{k}")
g.output(cat, "cat"); return g, names
def check_slices(seeds, bound=0.40, wbits=8):
probs = []
for s in seeds: p, _ = check_graph(s, bound, wbits, build=random_slice_graph); probs += p
return probs
def check_slice_errors():
"""slice_rows must REFUSE a range outside the tensor and a slice that feeds a concatenation or a relu. Returns the cases wrongly accepted."""
wrong = []
def attempt(name, build):
try: build(); wrong.append(f"accepted: {name}")
except C.CompileError: pass
def out_of_range(): g = C.Graph(); x = g.input("x", (3, 4)); g.slice_rows(x, 2, 2)
def empty(): g = C.Graph(); x = g.input("x", (3, 4)); g.slice_rows(x, 1, 0)
def into_concat():
g = C.Graph(); x = g.input("x", (3, 4)); y = g.input("y", (1, 4)); g.output(g.concat_rows(g.slice_rows(x, 0, 2), y), "o"); C.compile_graph(g, [{"x": [[1.0] * 4] * 3, "y": [[1.0] * 4]}])
attempt("a slice past the end", out_of_range); attempt("an empty slice", empty); attempt("a slice feeding a concatenation", into_concat)
return wrong
def check_many(seeds, bound=0.16, wbits=8):
probs = []; worst = 0.0
for s in seeds:
p, w = check_graph(s, bound, wbits); probs += p; worst = max(worst, w or 0.0)
return probs, worst
def check_cancellation():
"""Directed case: add(x, y) with y close to -x, so the operands are large and their sum is tiny. The sum's scale must still hold the operands (they are rescaled to it); a scale chosen from the sum alone clips them.
int8 cannot represent the tiny sum itself, so no accuracy bound is applied here: only checks (a)-(d) (reuse, interpreter, every tensor, real-number meaning)."""
R = random.Random(1); g = C.Graph(); x = g.input("x", (3, 6)); y = g.input("y", (3, 6)); sm = g.add(x, y)
w = g.weight("w", [[R.gauss(0, 0.4) for _ in range(4)] for _ in range(6)]); g.output(g.matmul(sm, w), "o")
def smp(): a = [[R.gauss(0, 3) for _ in range(6)] for _ in range(3)]; return {"x": a, "y": [[-v + R.gauss(0, 0.05) for v in r] for r in a]}
calib = [smp() for _ in range(20)]; test = smp(); P1 = C.compile_graph(g, calib, reuse=True); P0 = C.compile_graph(g, calib, reuse=False)
o1 = C.run(P1, test); o0, m0 = C.run(P0, test, machine=True); ints = C.interpret(P0, test); problems = []
if o1["o"][0] != o0["o"][0]: problems.append("cancellation: (a) reuse and no-reuse differ")
if o0["o"][0] != ints[g.nodes[-1].id]: problems.append("cancellation: (b) output differs from the interpreter")
for n in g.nodes:
if n.op != "output" and C.read_tensor(P0, m0, n.id) != ints[n.id]: problems.append(f"cancellation: (c) tensor {n.id} ({n.op}) differs from the interpreter")
return problems + ["cancellation: (d) " + q for q in C.stagewise_problems(P0, ints)]
def check_errors():
"""Things Capra must REFUSE. A compiler that quietly accepts them produces wrong code. Returns the list of cases it wrongly accepted."""
R = random.Random(0); wrong = []
def W(r, c): return [[R.gauss(0, 0.1) for _ in range(c)] for _ in range(r)]
def attempt(name, build, inputs):
try: C.compile_graph(build(), [inputs]); wrong.append(f"accepted: {name}")
except C.CompileError: pass
def inner_too_big():
g = C.Graph(); x = g.input("x", (2, 65)); g.output(g.matmul(x, g.weight("w", W(65, 3))), "o"); return g
def relu_with_second_consumer():
g = C.Graph(); x = g.input("x", (2, 4)); y = g.matmul(x, g.weight("w", W(4, 4))); g.output(g.add(g.relu(y), y), "o"); return g
def relu_on_input():
g = C.Graph(); x = g.input("x", (2, 4)); g.output(g.relu(x), "o"); return g
def too_big_for_scratchpad():
g = C.Graph(); x = g.input("x", (64, 64)); g.output(g.matmul(x, g.weight("w", W(64, 64))), "o"); return g
def feeds_two_concats():
g = C.Graph(); x = g.input("x", (1, 3)); y = g.input("y", (1, 3)); a = g.concat_rows(x, y); b = g.concat_rows(x, y); g.output(g.add(a, b), "o"); return g
mk = lambda r, c: [[R.gauss(0, 1) for _ in range(c)] for _ in range(r)]
attempt("inner dimension 65", inner_too_big, {"x": mk(2, 65)}); attempt("relu on a tensor with a second consumer", relu_with_second_consumer, {"x": mk(2, 4)})
attempt("relu that cannot be fused", relu_on_input, {"x": mk(2, 4)}); attempt("buffers larger than the scratchpad", too_big_for_scratchpad, {"x": mk(64, 64)})
attempt("a tensor feeding two concatenations", feeds_two_concats, {"x": mk(1, 3), "y": mk(1, 3)})
return wrong
if __name__ == "__main__":
n = int(sys.argv[1]) if len(sys.argv) > 1 else 60
p, w = check_many(range(n)); print(f"{n} random graphs: worst relative error {w:.4f}; problems: {p[:8] if p else 'none'}"); print("refusals:", check_errors() or "all five bad graphs rejected"); print("cancellation:", check_cancellation() or "ok")
- (a) The allocator. The program with buffer reuse and the one that keeps every buffer alive must give identical outputs. A buffer freed too early, or two live buffers given overlapping space, makes them differ.
- (b) Meaning. The simulator's outputs must equal the integer interpreter's, which runs the graph in plain Python with no instructions, tiles or addresses.
- (c) Every tensor. With reuse off, every tensor read back from the scratchpad must equal the interpreter's, so a bug in the middle of the graph cannot hide behind a correct-looking output.
- (d) Scales. Every tensor must be within one level of the real-number meaning of its operation applied to the dequantized inputs (
stagewise_problems). This check uses the planned scales but not the requantizer parameters, so it sees a wrong multiplier that (b) cannot (the interpreter shares the parameters with the compiler). - (e) Accuracy. The dequantized outputs must stay within 0.16 of the floating-point result (relative error) on inputs not used for calibration; 0.16 is twice the worst error the correct compiler has on the 60 battery graphs (0.079).
Beyond random graphs, check_errors lists five graphs the compiler must refuse (inner dimension 65, a relu that cannot be fused, a tensor feeding two concatenations, buffers larger than the scratchpad, a relu on a tensor with another consumer) and check_cancellation is a directed case described below.
What the tests found while the compiler was being written¶
None of these is an invented example; each was a failing check:
- A fused relu lost its place. In a graph that concatenated the output of a relu with another tensor, the placement code gave the relu's output a view inside the concatenation buffer and then overwrote that placement with its producer's separate buffer. Check (a) failed, (b) failed, and (c) pointed at the exact tensor. The fix: the matmul under a relu takes the relu's location.
- Forced scale sharing saturates. The first version made both operands and the result of an
addshare one scale. When one operand was a softmax output (pinned to 1/127, a range of 1.0) and the other an ordinary tensor reaching 3, the ordinary tensor clipped at 1.0 and the error was 50-100%. The principled fix is the rescale instruction described above. A second fix followed: the sum's scale must hold both operands (not just the sum) because two large operands that cancel would otherwise be clipped when rescaled. - The generator was wrong too. Concatenating a unit-variance input onto a tensor of magnitude 0.01 forces one scale on both and ruins the small one. That is not a compiler bug: real concatenations (a key cache and a new key) have matching magnitudes. The generator now creates partner inputs at the magnitude of what they join.
- A limit of per-tensor static scales. Repeated squaring (
x x^T, three times) spans orders of magnitude between samples; no single scale serves it. The generator allows one self-product per graph. This is a property of the technique, not a defect to fix. - Cancellation. In
add(x, y)withyclose to-x, int8 cannot hold the tiny sum however the compiler behaves: the correct compiler gets this wrong by 100% in relative terms. The accuracy check therefore cannot judge it. But a compiler bug that chooses the sum's scale from the sum alone (ignoring the operands) makes the rescaled operands clip, which the real-number check (d) sees. So the battery includes this one graph, with checks (a)-(d) only. Random graphs almost never contain a cancellation, and an experiment showed the mutant surviving until this directed case was added.
Testing the tests¶
#!/usr/bin/env python3
"""Chapter 11: test the tests of the COMPILER. Break model/capra.py one line at a time and run the battery: 60 random graphs under checks (a)-(e), the five graphs the compiler must refuse, and a directed cancellation case. 'caught' = some check reports a problem or the compiler crashes."""
import concurrent.futures, os, shutil, subprocess, sys, tempfile
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))); SRC = open(os.path.join(ROOT, "model", "capra.py")).read()
M = [
("tiling: edge tiles are always 4 x 4", "mt, nt = min(4, m - i0), min(4, nn - j0)", "mt, nt = 4, 4"),
("tiling: the A tile offset ignores the row length", "A=ins[0] + i0 * k,", "A=ins[0] + i0,"),
("tiling: the transposed B tile offset ignores the row length", "B=ins[1] + (j0 * k if tb else j0)", "B=ins[1] + j0"),
("tiling: the stride of B is the output width even when transposed", "ldb=(k if tb else nn)", "ldb=nn"),
("tiling: the transpose flag is dropped", "tb=int(tb), lda=k", "tb=0, lda=k"),
("tiling: results are stored with row stride 4", "ldc=nn))", "ldc=4))"),
("quantization: the constant factor of a matmul is ignored", 'ms = mantissa_shift(n.attrs["scale"] * scale[n.ins[0]] * scale[n.ins[1]] / scale[n.id])', 'ms = mantissa_shift(scale[n.ins[0]] * scale[n.ins[1]] / scale[n.id])'),
("quantization: the second operand's scale is replaced by the first's", 'scale[n.ins[0]] * scale[n.ins[1]] / scale[n.id])', 'scale[n.ins[0]] * scale[n.ins[0]] / scale[n.id])'),
("quantization: the fused relu flag is dropped", 'relu=int(P.relu_fused.get(n.id, False)))', 'relu=0)'),
("quantization: softmax probabilities are requantized with 1/65536", "ms = mantissa_shift(127 / 65536)", "ms = mantissa_shift(1 / 65536)"),
("planning: the softmax input scale is pinned to 1/8", "pinned[find(n.ins[0])] = 1 / 16", "pinned[find(n.ins[0])] = 1 / 8"),
("planning: calibration keeps the smallest range, not the largest", "rng[nid] = max(rng.get(nid, 0.0), max(abs(x) for r in m for x in r))", "rng[nid] = min(rng.get(nid, 1e9), max(abs(x) for r in m for x in r))"),
("planning: concatenated parts get independent scales", 'if n.op == "concat": union(n.ins[0], n.id); union(n.ins[1], n.id)', 'if n.op == "concat": union(n.ins[0], n.id)'),
("planning: the sum's scale ignores the operands' ranges", " for i in n.ins: grp_max[find(n.id)] = max(grp_max[find(n.id)], rng[i])", " pass"),
("lowering: softmax rows all read the first row", 'B("SM", dst=tmp + r * nn, src=ins[0] + r * nn, len=nn)', 'B("SM", dst=tmp + r * nn, src=ins[0], len=nn)'),
("lowering: argmax of every row is written to the first slot", "dst=out + r, src=ins[0] + r * nn, len=nn", "dst=out, src=ins[0] + r * nn, len=nn"),
("lowering: an add rescales with the inverse ratio", "ratio = scale[i] / scale[n.id]", "ratio = scale[n.id] / scale[i]"),
("lowering: an add never rescales", "if abs(ratio - 1) < 1e-9: srcs.append(ins[k]); continue", "srcs.append(ins[k]); continue"),
("lowering: concatenated parts are placed on top of each other", "loc[i] = (b0, off + (0 if k == 0 else N[n.ins[0]].shape[0]))", "loc[i] = (b0, off)"),
("lowering: a load copies only the first row", "code.append(B(\"LD\", dst=a, src=e, len=n.shape[0] * n.shape[1]))", "code.append(B(\"LD\", dst=a, src=e, len=n.shape[1]))"),
("lowering: an output is stored without its last row", "code.append(B(\"ST\", src=ins[0], dst=e[\"addr\"], len=n.shape[0] * n.shape[1]))", "code.append(B(\"ST\", src=ins[0], dst=e[\"addr\"], len=max(1, n.shape[0] * n.shape[1] - n.shape[1])))"),
("lowering: external regions of inputs and weights overlap", "\"kind\": n.op, \"words\": words}; ext_top += words", "\"kind\": n.op, \"words\": words}; ext_top += n.shape[0]"),
("allocator: buffers are released one step too early", "if l == n.id and bufs[b][\"addr\"] is not None", "if l <= n.id + 1 and bufs[b][\"addr\"] is not None"),
("allocator: free pieces are merged across a gap", "if merged and merged[-1][0] + merged[-1][1] == a:", "if merged and merged[-1][0] + merged[-1][1] <= a:"),
("checks: the inner-dimension limit is 640", "if k > 64:", "if k > 640:"),
("checks: relu may follow a matmul with other consumers", "or len(cons[src.id]) != 1", ""),
("checks: a tensor may feed two concatenations", 'if i in loc: raise CompileError(f"node {i} feeds two concatenations")', "pass"),
]
EQUIV = [("EQUIVALENT in effect: the allocator skips an exact-fit hole (it uses the next larger one, or fails only when memory is nearly full)", "if sz >= n:", "if sz > n:")]
def check(label, old, new):
if SRC.count(old) != 1: return label, "BAD ANCHOR (%d occurrences)" % SRC.count(old)
d = tempfile.mkdtemp(prefix="mut11_"); shutil.copytree(os.path.join(ROOT, "model"), os.path.join(d, "model")); open(os.path.join(d, "model", "capra.py"), "w").write(SRC.replace(old, new))
code = "import sys; sys.path.insert(0,'model'); import capra_tests as T; p,w=T.check_many(range(60)); e=T.check_errors(); c=T.check_cancellation(); sys.exit(1 if (p or e or c) else 0)"
try: r = subprocess.run([sys.executable, "-c", code], cwd=d, capture_output=True, text=True, timeout=900).returncode != 0
except subprocess.TimeoutExpired: r = True
shutil.rmtree(d, ignore_errors=True); return label, r
print("NOTE: this script deliberately breaks copies of the compiler. 'caught' is EXPECTED: it shows the battery notices the mistake.")
with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool: rows = list(pool.map(lambda m: check(*m), M)); eq = list(pool.map(lambda m: check(*m), EQUIV))
n = 0
for label, r in rows: print(f"{label}: {'caught' if r is True else r if isinstance(r, str) else 'NOT CAUGHT'}"); n += r is True
print(f"\ncompiler: {n} of {len(M)} broken compilers caught")
print("\nA change that looks like a bug and is not:")
for label, r in eq: print(f"{label}: {'caught' if r is True else 'NOT CAUGHT'}")
sys.exit(0 if n == len(M) else 1)
Output (cloud sandbox -- live-executed)
NOTE: this script deliberately breaks copies of the compiler. 'caught' is EXPECTED: it shows the battery notices the mistake.
tiling: edge tiles are always 4 x 4: caught
tiling: the A tile offset ignores the row length: caught
tiling: the transposed B tile offset ignores the row length: caught
tiling: the stride of B is the output width even when transposed: caught
tiling: the transpose flag is dropped: caught
tiling: results are stored with row stride 4: caught
quantization: the constant factor of a matmul is ignored: caught
quantization: the second operand's scale is replaced by the first's: caught
quantization: the fused relu flag is dropped: caught
quantization: softmax probabilities are requantized with 1/65536: caught
planning: the softmax input scale is pinned to 1/8: caught
planning: calibration keeps the smallest range, not the largest: caught
planning: concatenated parts get independent scales: caught
planning: the sum's scale ignores the operands' ranges: caught
lowering: softmax rows all read the first row: caught
lowering: argmax of every row is written to the first slot: caught
lowering: an add rescales with the inverse ratio: caught
lowering: an add never rescales: caught
lowering: concatenated parts are placed on top of each other: caught
lowering: a load copies only the first row: caught
lowering: an output is stored without its last row: caught
lowering: external regions of inputs and weights overlap: caught
allocator: buffers are released one step too early: caught
allocator: free pieces are merged across a gap: caught
checks: the inner-dimension limit is 640: caught
checks: relu may follow a matmul with other consumers: caught
checks: a tensor may feed two concatenations: caught
compiler: 27 of 27 broken compilers caught
A change that looks like a bug and is not:
EQUIVALENT in effect: the allocator skips an exact-fit hole (it uses the next larger one, or fails only when memory is nearly full): NOT CAUGHT
All 27 broken compilers are caught: tiling mistakes (edge tiles, offsets, strides, the transpose flag), quantization mistakes (a dropped constant factor, a wrong operand's scale, a lost relu flag, a wrong probability scale, a wrong pinned scale, calibration keeping the smallest range), lowering mistakes (rows reading the wrong row, wrong destinations, a missing rescale, concatenation parts overlapping, a truncated load or store, overlapping external regions), allocator mistakes (releasing a buffer a step early, merging free blocks across a gap) and removed safety checks (the inner-dimension limit, relu fusion with another consumer, a tensor feeding two concatenations). The removed checks are caught only because the battery contains graphs the compiler must refuse; without check_errors those mutants survive.
One change survives and is equivalent in effect: an allocator that takes the next larger hole instead of an exact fit only differs when memory is nearly full.
Common mistakes¶
- Calibrating on too little data. A range seen in 20 samples is not the range of the real inputs; values beyond it clip at +-127. Chapter 12 fixes the ranges of the KV cache from outside for this reason (
ranges=). - Assuming per-tensor scales suit every graph. Repeated squaring (
x x^Tthree times) spans orders of magnitude and no single scale serves it. The technique has limits; the generator avoids that case and the page says so. - Forgetting that a relu must be fused. In this IR a relu is not a separate instruction: it is a flag on the requantizer of the matmul before it. A relu after an
addis refused. - Letting a tensor feed two places. A tensor that is part of two concatenations cannot be a view of both. The compiler refuses; restructure the graph.
- Trusting an accuracy number from a graph with cancellation.
add(x, y)with y near -x loses its significant bits in int8 however good the compiler is. The directed cancellation test checks the compiler's behaviour, not the accuracy. - Testing a compiler on the graphs its author thought of. The random graphs and the five refusal graphs found bugs the author's examples did not.
- Reading the compiled program as if it were hand-written. It reuses scratchpad addresses on purpose: an address that looks overwritten is a buffer whose last reader has run.
What this chapter does and does not establish¶
- Verified: on 60 battery graphs and 12 fresh ones, the compiled programs match the interpreter exactly, every tensor is consistent with its real-number meaning, the allocator is correct under reuse, and the programs run on the RTL with exactly the reference's memory and cycle counters.
- Not covered: a matrix product with
K > 64(it needs an int32 accumulate instruction the ISA does not have; the compiler refuses it); scratchpad spilling; scheduling that overlaps loads with compute (the machine of Chapter 7 runs one instruction at a time); other operations (layer normalization, GELU); dynamic shapes; and anything about calibration data beyond random Gaussians. The accuracy numbers say nothing about a trained model.
Chapter summary¶
The compiler turns a float graph into a GA-2 program in five steps: calibrate, plan scales, tile, allocate with reuse, emit. It was tested against an interpreter and the real-number meaning of every tensor, against itself with and without reuse, and against graphs it must refuse; 27 of 27 broken compilers were caught, and the process found a placement bug, a scale-sharing flaw in add, and a blind spot (cancellation) that needed a directed case.
Self-check questions¶
- Why does the compiler reject a matmul with
K = 65, and what would the hardware need to support it? - Why was forcing both operands of an
addto share a scale wrong, and what does the compiler do instead? - Why is check (c) run with buffer reuse switched off?
- Which check catches "buffers are released one step too early", and why do the others not?
- Why is the cancellation graph tested without an accuracy bound, and what bug does it exist to catch?
- In Running example A, why does the scores tensor have scale 0.0625 although its range divided by 127 is 0.0201? What would go wrong with 0.0201?
- In the 5 x 12 times 12 x 10 product of Figure 11.4, how many
MMinstructions are emitted, and what are their (M, N) values? What arelda,ldbandldc? - At hidden width 48 the compiled network needs 3,328 scratchpad words with reuse and does not fit without it. At width 64 it fails even with reuse. Why?
Exercises¶
- Your own graph. Edit
buildinch11_example_b.pyso that the second layer ismatmul(h, w2, transpose_b=True)with a square weight, and compile it. Does the output change? Why does it not change the footprint? - Add a layer. Add a third hidden layer to
build. Predict whether it still fits the scratchpad at width 32, then check. Which of the checks would you rerun first after such a change? - A refusal. Make a graph that applies
reluto the output of anaddand read the error. Rewrite it so that the relu follows a matmul and compiles. - Pin a scale. Add
ranges={"tag": 4.0}to a compile of a graph whose tensors you tag (g.tag(node, "tag")). What happens to that tensor's scale, and to the neighbours that share it? - Read the allocator. Modify Example A's graph so that the value cache is 20 rows instead of 5. Predict the new high-water mark before running, then check. Which buffer defines it?