40. A Model of a Matrix Accelerator: Tiles, a Scratchpad, and What Loop Order Does There

What you will understand: what decides the speed of a matrix accelerator, measured on a machine you can read in full. This chapter builds GA-1, a cycle-counting model (not a chip, and not a design for one) with the three things that matter: a DMA engine that moves tiles between off-chip memory and a small on-chip scratchpad, a matrix unit that multiplies two tiles and accumulates, and a vector unit (used from Chapter 41). Then it takes the loop-order and tiling ideas of Chapters 7, 24 and 37, which on a CPU were about caches and vector registers, and asks what they do here: which schedule moves the fewest tiles, when the machine is limited by its memory and when by its arithmetic, and what double buffering buys. Every number comes from running the simulator, a simple analytic model is checked against it, and a deliberately broken simulator is shown to be caught.

What you need to know first: Chapter 24 (the ikj loop order and why ijk walks memory badly) and Chapter 37 (the idea of a core model that ignores memory). No Mountain Goat compiler is involved in this chapter: the schedules are written by hand. Chapter 41 compiles Mountain Goat programs onto GA-1.

This is a model, not a chip

Every parameter of GA-1 (tile size, scratchpad size, DMA cost, matrix-unit rate) is an assumption chosen to have plausible ratios (the matrix unit is much faster than the memory), not a measurement of any real device, and no real accelerator was consulted to calibrate it. What transfers to real chips is the shape of the arguments (reuse, double buffering, bound by memory or by arithmetic); the cycle counts do not.

Compile and run

cd docs/part40/code                               # plain Python 3, no dependencies
./sweep_orders.py > sweep_orders_out.txt          # 128 x 128 x 128 product, 14 schedules, two scratchpad sizes (8 seconds)
./timeline.py > timeline_out.txt                  # a text picture of the engines
./bandwidth_sweep.py > bandwidth_sweep_out.txt    # what a faster memory changes (8 seconds)
./check_ga.py                                     # results, traffic formulas, timing, independent hazard replay (about a second)
./sim_mutation.py > sim_mutation_out.txt          # 11 WRONG simulators the checker must catch (15 seconds)
cd ../../part15/code && ./run_lit.sh              # the whole test suite

Every listing and output on this page comes from these commands.

Why this page shows a mutation run, and why 'caught' is the wanted result

sim_mutation.py writes broken versions of the simulator and expects the checker to fail on each. "Caught" is good; "NOT CAUGHT" would be a gap. Nothing on this page is a failing test.

The machine

part	what GA-1 has	assumption (a number in ga.py)
tile	T × T words, T = 8 (64 words)	the unit of every transfer and of every matrix product
scratchpad	a few tile slots, 16 or 64 here	the only memory the units can use
DMA	one engine; one transfer at a time, DRAM ↔ one slot	16 + 64/16 = 20 cycles per tile (setup 16, then 16 words per cycle)
matrix unit	C += A @ B on three slots; accepts a new product every 8 cycles (8 × 8 × 8 = 512 multiply-adds, 64 per cycle); another engine can read a product 8 cycles after it finishes (the "fill"), but accumulating into it again needs no extra wait	mxu_cycles = 8, mxu_fill = 8
vector unit	elementwise on a tile, 16 lanes; exp on a slower path of 4 lanes (Chapter 41)

Each engine has its own in-order queue and the three run independently (as in real accelerators with decoupled access and execute); an instruction starts when its engine is free, its source slots have been written, and the slots it overwrites are no longer needed by earlier instructions. The instructions are load, store, zero, mxu and, for Chapter 41, a few vector ones. The whole simulator, including the instruction set in its header comment and the one place where the timing rule lives (_run), is short enough to read:

#!/usr/bin/env python3
"""GA-1: a MODEL of a small matrix accelerator, written in plain Python. It is not a chip and not a design for one: it is a cycle-counting simulator with the three things
that decide the speed of an inference accelerator, so that the effect of tiling and loop order on them can be measured:

  * a DMA engine that moves one TILE (T x T words) between off-chip memory ("DRAM") and an on-chip SCRATCHPAD of a few tile slots, taking setup + words / bandwidth cycles;
  * a MATRIX UNIT that multiplies two T x T tiles and adds the product to a third (C += A @ B), one tile product every `mxu_cycles` cycles;
  * a VECTOR UNIT (elementwise, with a slower special-function path for exp) and reductions, used from Chapter 41 on.

Each engine (DMA, matrix unit, vector unit) has its own in-order queue and the engines run independently of each other (decoupled access/execute, as in real accelerators); an
instruction starts when its engine is free, its operands are ready and the slots it overwrites are no longer needed (so a program that gets its loads ahead of its multiplies, into
slots that are free, overlaps the DMA with the matrix unit; a program that has no free slots to load into, waits). Every instruction really computes its result on
real numbers, in program order (so the result of a schedule is what its instructions say, whatever the timing); the TIMING rule enforces the hazards, and an optional trace of every
instruction's start and finish lets an independent checker (check_ga.py) verify that no instruction read a slot before it was written or overwrote one still being read. All the numbers below are assumptions of the model, chosen to be
plausible in ratio (the matrix unit is much faster than the DMA), not measurements of any real device."""
import math
from dataclasses import dataclass, field

@dataclass
class Machine:
    T: int = 8                     # tile edge: tiles are T x T words
    slots: int = 16                # scratchpad capacity, in tiles
    dma_setup: int = 16            # cycles of fixed cost per tile transfer
    dma_words_per_cycle: int = 16  # DMA bandwidth once started (one transfer at a time)
    mxu_cycles: int = 8            # the matrix unit accepts a new tile product every this many cycles (T x T x T MACs: T x T = 64 MACs per cycle)
    mxu_fill: int = 8              # extra cycles before a product can be READ by another engine (accumulating into it again needs no extra wait)
    vpu_lanes: int = 16            # elementwise lanes
    sfu_lanes: int = 4             # lanes of the special-function path (exp)
    @property
    def tile_words(self): return self.T * self.T
    @property
    def dma_cycles(self): return self.dma_setup + math.ceil(self.tile_words / self.dma_words_per_cycle)
    @property
    def peak_macs_per_cycle(self): return self.T ** 3 / self.mxu_cycles

# Instruction = (op, operands...). Slots are integers 0..slots-1; tensors are names in the machine's DRAM; (ti, tj) are tile coordinates.
#   ("load", slot, name, ti, tj)    DMA: DRAM tile -> slot (a tile that sticks out of the tensor is zero-padded)
#   ("store", slot, name, ti, tj)   DMA: slot -> DRAM tile (only the part inside the tensor is written)
#   ("zero", slot)                  vector unit: slot := 0
#   ("mxu", c, a, b)                matrix unit: slot c += slot a @ slot b
#   ("vop", op, dst, a, b)          vector unit, elementwise on tiles: op in add sub mul div max ; b may be None for unary op in relu neg
#   ("exp", dst, a)                 special-function path
#   ("vscale", dst, a, k)           dst := a * k   (k a number)
#   ("rowred", op, dst, a)          reduce each row of a with op in sum max, and broadcast the result across the row of dst
#   ("bcastcol", dst, a)            dst[i][j] := a[i][0]   (column 0 spread across the row)
#   ("masklen", slot, rows, cols)   vector unit: zero the part of the tile outside the first rows x cols entries (used after a padded load)
class Sim:
    def __init__(self, machine, dram):
        self.m = machine; self.dram = dram                 # dram: name -> list of rows (floats)
        T = machine.T; self.spad = [[[0.0] * T for _ in range(T)] for _ in range(machine.slots)]
        self.free = {"dma": 0, "mxu": 0, "vpu": 0}; self.busy = {"dma": 0, "mxu": 0, "vpu": 0}
        self.ready = [0] * machine.slots; self.acc_ready = [0] * machine.slots; self.last_read = [0] * machine.slots
        self.end = 0
        self.trace = []                                 # (index, engine, start, finish, reads, writes) for every instruction, for the independent hazard check
        self.dma_tiles = 0; self.dma_loads = 0; self.dma_stores = 0; self.macs = 0; self.instructions = 0
    # ---- the timing rule, in one place
    def _run(self, engine, duration, reads, writes, acc=(), fill=0):
        start = self.free[engine]                       # each engine runs its own instructions in order; engines run independently of each other
        for s in reads: start = max(start, self.ready[s])
        for s in acc: start = max(start, self.acc_ready[s], self.last_read[s])
        for s in writes:
            start = max(start, self.last_read[s], self.ready[s])
        finish = start + duration
        self.free[engine] = finish; self.busy[engine] += duration
        for s in reads: self.last_read[s] = max(self.last_read[s], finish)
        for s in acc: self.last_read[s] = max(self.last_read[s], finish)
        for s in tuple(writes) + tuple(acc): self.ready[s] = finish + fill; self.acc_ready[s] = finish
        self.end = max(self.end, finish + fill); self.instructions += 1
        self.trace.append((self.instructions - 1, engine, start, finish, tuple(reads) + tuple(acc), tuple(writes) + tuple(acc), fill))
        return start, finish
    def tile(self, name, ti, tj):
        T = self.m.T; M = self.dram[name]
        return [[M[ti * T + i][tj * T + j] if ti * T + i < len(M) and tj * T + j < len(M[0]) else 0.0 for j in range(T)] for i in range(T)]
    def run(self, program):
        m = self.m; T = m.T; S = self.spad
        for ins in program:
            op = ins[0]
            if op == "load":
                _, s, name, ti, tj = ins; self._run("dma", m.dma_cycles, (), (s,)); S[s] = self.tile(name, ti, tj); self.dma_tiles += 1; self.dma_loads += 1
            elif op == "store":
                _, s, name, ti, tj = ins; self._run("dma", m.dma_cycles, (s,), ()); M = self.dram[name]
                for i in range(T):
                    for j in range(T):
                        if ti * T + i < len(M) and tj * T + j < len(M[0]): M[ti * T + i][tj * T + j] = S[s][i][j]
                self.dma_tiles += 1; self.dma_stores += 1
            elif op == "zero":
                _, s = ins; self._run("vpu", 1, (), (s,)); S[s] = [[0.0] * T for _ in range(T)]
            elif op == "mxu":
                _, c, a, b = ins; self._run("mxu", m.mxu_cycles, (a, b), (), acc=(c,), fill=m.mxu_fill)       # c is read and written (accumulated): chained products wait for acc_ready, other readers for ready (= finish + fill)
                A, B, C = S[a], S[b], S[c]
                S[c] = [[C[i][j] + sum(A[i][k] * B[k][j] for k in range(T)) for j in range(T)] for i in range(T)]; self.macs += T ** 3
            elif op == "vop":
                _, f, d, a, b = ins; dur = math.ceil(m.tile_words / m.vpu_lanes) + 1
                self._run("vpu", dur, (a,) + ((b,) if b is not None else ()), (d,)); A = S[a]; B = S[b] if b is not None else None
                fn = {"add": lambda x, y: x + y, "sub": lambda x, y: x - y, "mul": lambda x, y: x * y, "div": lambda x, y: x / y if y != 0 else float("inf") if x > 0 else float("-inf") if x < 0 else float("nan"),
                      "max": lambda x, y: max(x, y), "relu": lambda x, y: max(x, 0.0), "neg": lambda x, y: -x}[f]
                S[d] = [[fn(A[i][j], B[i][j] if B is not None else 0.0) for j in range(T)] for i in range(T)]
            elif op == "exp":
                _, d, a = ins; self._run("vpu", math.ceil(m.tile_words / m.sfu_lanes) + 1, (a,), (d,)); S[d] = [[math.exp(x) for x in row] for row in S[a]]
            elif op == "vscale":
                _, d, a, k = ins; self._run("vpu", math.ceil(m.tile_words / m.vpu_lanes) + 1, (a,), (d,)); S[d] = [[x * k for x in row] for row in S[a]]
            elif op == "rowred":
                _, f, d, a = ins; self._run("vpu", math.ceil(m.tile_words / m.vpu_lanes) + 3, (a,), (d,)); fn = sum if f == "sum" else max
                S[d] = [[fn(row)] * T for row in S[a]]
            elif op == "bcastcol":
                _, d, a = ins; self._run("vpu", math.ceil(m.tile_words / m.vpu_lanes) + 1, (a,), (d,)); S[d] = [[row[0]] * T for row in S[a]]
            elif op == "masklen":
                _, s, r, c = ins; self._run("vpu", 1, (s,), (s,)); S[s] = [[S[s][i][j] if i < r and j < c else 0.0 for j in range(T)] for i in range(T)]
            else: raise ValueError(op)
        return self
    @property
    def cycles(self): return self.end
    def report(self):
        c = max(self.cycles, 1)
        return dict(cycles=self.cycles, dma_tiles=self.dma_tiles, loads=self.dma_loads, stores=self.dma_stores, macs=self.macs, instructions=self.instructions,
                    mxu_utilization=self.macs / (self.m.peak_macs_per_cycle * c), dma_busy=self.busy["dma"] / c, mxu_busy=self.busy["mxu"] / c, vpu_busy=self.busy["vpu"] / c)

Two things to notice. The functional result of every instruction is computed on real numbers, in program order, so a schedule's output is what its instructions say, whatever the timing. The timing rule is what enforces the hazards, and the simulator records every instruction's start and finish so that a separate checker can verify it (below).

A six-instruction program, timed by hand

Zero a slot, load two tiles, multiply them into it twice, store: ("zero", 2), ("load", 0, "A", 0, 0), ("load", 1, "B", 0, 0), ("mxu", 2, 0, 1), ("mxu", 2, 0, 1), ("store", 2, "C", 0, 0). By hand: the two loads share the one DMA engine, 20 cycles each, so the products can start at cycle 40; two chained products take 16 cycles (40 to 56, accumulating into the same slot needs no fill wait); the store must wait until the product can be read, 8 cycles after cycle 56, so it starts at 64 and takes 20: 84 cycles. The simulator says 84, and check_ga.py checks that.

Here is a text picture (timeline.py) of a real, small product (32 × 32 × 32, 4 × 4 × 4 tiles) scheduled as one block of 2 × 2 tiles of C, single and double buffered (one character is 4 cycles; D is a DMA transfer, M a matrix-unit product):

2 x 2 block of C, single buffered: 1696 cycles, 80 tiles moved, matrix-unit utilization 30%   (one character = 4 cycles)
  dma DDDDDDDDDDDDDDDDDDDD..DDDDDDDDDDDDDDDDDDDD..DDDDDDDDDDDDDDDDDDDD..DDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDD ...
  mxu ...............MM...MMMMMM...........MM...MMMMMM...........MM...MMMMMM...........MM...MMMMMM............................ ...

2 x 2 block of C, double buffered: 1600 cycles, 80 tiles moved, matrix-unit utilization 32%   (one character = 4 cycles)
  dma DDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDD ...
  mxu ...............MM...MMMMMM.........MM...MMMMMM.........MM...MMMMMM.........MM...MMMMMM.............................MM... ...

The DMA row is almost always busy; the matrix unit is idle most of the time, waiting for tiles (the Ms come in bursts after each step's loads). That is the whole chapter in one picture: with these ratios the memory, not the arithmetic, is the bottleneck, and the schedule decides how bad it is.

The schedules

Two families of hand-written programs for C = A @ B (schedules.py):

#!/usr/bin/env python3
"""Hand-written GA-1 programs for C = A @ B, in two families, and the helpers that run and check them.

  blocked(M, K, N, bi, bj, double_buffer): "output-stationary" schedules. A block of bi x bj tiles of C stays in the scratchpad while the K dimension is streamed through it:
        for each block of C:  zero it;  for each k: load bi tiles of A and bj tiles of B, do bi*bj tile products;  store the block.
      bi = bj = 1 is the plain (i, j, k) loop order; bi = 1, bj = (all tiles of a row) is the (i, k, j) order of Chapter 24 (a row of C stays, a tile of A is reused across the row);
      bigger blocks reuse both. With double_buffer the loads for step k+1 are issued BEFORE the products of step k, into a second set of slots.
  kouter(M, K, N, double_buffer): the (k, i, j) loop order: every tile of C is read from DRAM, updated by one tile product and written back for each k.
Every program is checked against a plain triple loop by check_ga.py."""
import math
from ga import Machine, Sim

def tiles(n, T): return math.ceil(n / T)

def blocked(M, K, N, m, bi, bj, double_buffer=True, a="A", b="B", c="C"):
    Mt, Kt, Nt = tiles(M, m.T), tiles(K, m.T), tiles(N, m.T); nbuf = 2 if double_buffer else 1
    need = bi * bj + nbuf * (bi + bj)
    if need > m.slots: raise ValueError(f"blocked({bi},{bj}, double_buffer={double_buffer}) needs {need} slots, the scratchpad has {m.slots}")
    Cs = lambda x, y: x * bj + y                                        # slot of C tile (x, y) of the block
    As = lambda buf, x: bi * bj + buf * (bi + bj) + x
    Bs = lambda buf, y: bi * bj + buf * (bi + bj) + bi + y
    prog = []
    for I in range(0, Mt, bi):
        for J in range(0, Nt, bj):
            ii = [i for i in range(I, min(I + bi, Mt))]; jj = [j for j in range(J, min(J + bj, Nt))]
            for x in range(len(ii)):
                for y in range(len(jj)): prog.append(("zero", Cs(x, y)))
            def loads(k, buf):
                for x, i in enumerate(ii): prog.append(("load", As(buf, x), a, i, k))
                for y, j in enumerate(jj): prog.append(("load", Bs(buf, y), b, k, j))
            if double_buffer: loads(0, 0)
            for k in range(Kt):
                if double_buffer:
                    if k + 1 < Kt: loads(k + 1, (k + 1) % 2)
                    buf = k % 2
                else: loads(k, 0); buf = 0
                for x in range(len(ii)):
                    for y in range(len(jj)): prog.append(("mxu", Cs(x, y), As(buf, x), Bs(buf, y)))
            for x, i in enumerate(ii):
                for y, j in enumerate(jj): prog.append(("store", Cs(x, y), c, i, j))
    return prog

def kouter(M, K, N, m, double_buffer=True, a="A", b="B", c="C"):
    Mt, Kt, Nt = tiles(M, m.T), tiles(K, m.T), tiles(N, m.T); groups = 2 if double_buffer else 1
    if 3 * groups > m.slots: raise ValueError("scratchpad too small")
    prog = []; g = 0
    for k in range(Kt):
        for i in range(Mt):
            for j in range(Nt):
                sc, sa, sb = 3 * g, 3 * g + 1, 3 * g + 2; g = (g + 1) % groups
                if k == 0: prog.append(("zero", sc))
                else: prog.append(("load", sc, c, i, j))               # the partial sum comes back from DRAM
                prog += [("load", sa, a, i, k), ("load", sb, b, k, j), ("mxu", sc, sa, sb), ("store", sc, c, i, j)]
    return prog

def matrices(M, K, N, seed=1):
    import random; r = random.Random(seed)
    A = [[r.randint(-3, 3) for _ in range(K)] for _ in range(M)]; B = [[r.randint(-3, 3) for _ in range(N)] for _ in range(K)]
    return A, B

def reference(A, B):
    return [[sum(A[i][k] * B[k][j] for k in range(len(B))) for j in range(len(B[0]))] for i in range(len(A))]

def run(program, A, B, m):
    dram = {"A": A, "B": B, "C": [[0.0] * len(B[0]) for _ in range(len(A))]}
    sim = Sim(m, dram).run(program); return sim, dram["C"]

def lower_bounds(M, K, N, m):
    """The two simple bounds: the matrix unit alone, and the DMA alone moving the tiles that MUST move at least once (A, B in; C out)."""
    Mt, Kt, Nt = tiles(M, m.T), tiles(K, m.T), tiles(N, m.T)
    return dict(compute=Mt * Nt * Kt * m.mxu_cycles, dma=(Mt * Kt + Kt * Nt + Mt * Nt) * m.dma_cycles)

def expected_blocked_tiles(M, K, N, m, bi, bj):
    """Tiles moved by blocked(): for every block, bi_eff + bj_eff loads per k, and one store per C tile."""
    Mt, Kt, Nt = tiles(M, m.T), tiles(K, m.T), tiles(N, m.T); loads = 0
    for I in range(0, Mt, bi):
        for J in range(0, Nt, bj): loads += Kt * (min(bi, Mt - I) + min(bj, Nt - J))
    return loads, Mt * Nt

def predicted_double_buffered_cycles(M, K, N, m, bi, bj):
    """A simple model of a double-buffered blocked schedule: per k step the DMA needs (bi + bj) tile loads and the matrix unit bi * bj products, and with double buffering the two
    overlap, so a step takes the LONGER of the two; each block also pays its first loads before the first product and its stores after the last (nothing overlaps them)."""
    Mt, Kt, Nt = tiles(M, m.T), tiles(K, m.T), tiles(N, m.T); blocks = math.ceil(Mt / bi) * math.ceil(Nt / bj)
    load = (bi + bj) * m.dma_cycles; comp = bi * bj * m.mxu_cycles
    return blocks * (Kt * max(load, comp) + load + bi * bj * m.dma_cycles)

• blocked(bi, bj): keep a block of bi × bj tiles of C in the scratchpad and stream K through it. Each step loads bi tiles of A and bj tiles of B and does bi · bj products, so the loads per product are (bi + bj) / (bi · bj): 2 for 1 × 1 (the plain i, j, k order), 1.06 for a whole row of 1 × 16 (the i, k, j order of Chapter 24), 1 for 2 × 2, 0.5 for 4 × 4, 0.375 for 4 × 8. This is reuse: bigger blocks do more work per loaded tile, and cost scratchpad slots.
• double buffering: with two sets of A/B slots, the loads for step k + 1 are issued before the products of step k and land in the other set, so the DMA and the matrix unit overlap. It costs bi + bj more slots.
• kouter: the k, i, j order. Every tile of C is loaded (the partial sum), updated by one product and stored again, for every k.

Results

sweep_orders.py runs all of them on a 128 × 128 × 128 product (16 × 16 × 16 tiles, 4,096 tile products), checking each result against a plain triple loop first. The matrix unit alone would need 32,768 cycles; the DMA alone, moving every tile exactly once, 15,360.

scratchpad of 64 tile slots (4096 words); matrix unit alone needs 32768 cycles, DMA alone (each tile moved once) 15360 cycles
schedule                           tiles moved   cycles  MXU util  DMA busy  vs best  model, 2 buf
blocked 1x1 (i,j,k), 1 buf                8448   203776      16%       83%    5.44x
blocked 1x1 (i,j,k), 2 buf                8448   173056      19%       98%    4.62x   179200 (1.04)
blocked 1x16 (row: i,k,j), 1 buf          4608    94080      35%       98%    2.51x
blocked 1x16 (row: i,k,j), 2 buf          4608    92160      36%      100%    2.46x    97600 (1.06)
blocked 2x2, 1 buf                        4352    94720      35%       92%    2.53x
blocked 2x2, 2 buf                        4352    87040      38%      100%    2.32x    92160 (1.06)
blocked 2x4, 1 buf                        3328    76160      43%       87%    2.03x
blocked 2x4, 2 buf                        3328    66560      49%      100%    1.78x    70400 (1.06)
blocked 2x8, 1 buf                        2816    68800      48%       82%    1.84x
blocked 2x8, 2 buf                        2816    56320      58%      100%    1.50x    59520 (1.06)
blocked 4x4, 1 buf                        2304    56640      58%       81%    1.51x
blocked 4x4, 2 buf                        2304    46080      71%      100%    1.23x    48640 (1.06)
blocked 4x8, 1 buf                        1792    52640      62%       68%    1.41x
blocked 4x8, 2 buf                        1792    37440      88%       96%    1.00x    39808 (1.06)
k outermost (k,i,j), 1 buf               16128   388096       8%       83%   10.37x
k outermost (k,i,j), 2 buf               16128   388096       8%       83%   10.37x

scratchpad of 16 tile slots (1024 words); matrix unit alone needs 32768 cycles, DMA alone (each tile moved once) 15360 cycles
schedule                           tiles moved   cycles  MXU util  DMA busy  vs best  model, 2 buf
blocked 1x1 (i,j,k), 1 buf                8448   203776      16%       83%    2.68x
blocked 1x1 (i,j,k), 2 buf                8448   173056      19%       98%    2.27x   179200 (1.04)
blocked 1x16 (row: i,k,j), single buffered           does not fit in the scratchpad
blocked 1x16 (row: i,k,j), double buffered           does not fit in the scratchpad
blocked 2x2, 1 buf                        4352    94720      35%       92%    1.24x
blocked 2x2, 2 buf                        4352    87040      38%      100%    1.14x    92160 (1.06)
blocked 2x4, 1 buf                        3328    76160      43%       87%    1.00x
blocked 2x4, double buffered                 does not fit in the scratchpad
blocked 2x8, single buffered                 does not fit in the scratchpad
blocked 2x8, double buffered                 does not fit in the scratchpad
blocked 4x4, single buffered                 does not fit in the scratchpad
blocked 4x4, double buffered                 does not fit in the scratchpad
blocked 4x8, single buffered                 does not fit in the scratchpad
blocked 4x8, double buffered                 does not fit in the scratchpad
k outermost (k,i,j), 1 buf               16128   388096       8%       83%    5.10x
k outermost (k,i,j), 2 buf               16128   388096       8%       83%    5.10x

Read the 64-slot table first:

• The plain i, j, k order (1 × 1) is the worst of the blocked ones: 8,448 tiles moved, and even with double buffering 173,056 cycles and 19% of the matrix unit's peak. The DMA is busy 98% of the time: the machine is DMA-bound.
• Reuse is everything. Tiles moved fall from 8,448 (1 × 1) to 4,608 (a row), 4,352 (2 × 2), 3,328 (2 × 4), 2,816 (2 × 8), 2,304 (4 × 4) and 1,792 (4 × 8). Cycles follow the tiles when the DMA is the bottleneck: 4 × 4 double buffered takes exactly 46,080 = 2,304 × 20 cycles (the DMA 100% busy; the matrix unit at 71%, which is just 32,768 / 46,080).
• At 4 × 8 the bottleneck flips. A step has 12 loads (240 cycles) and 32 products (256 cycles): now the matrix unit is the longer, utilization reaches 88% and the cycles (37,440) are close to the arithmetic bound of 32,768. The condition for being compute-bound is 8 · bi · bj ≥ 20 · (bi + bj); with these numbers it first holds at 4 × 8 among the shapes tried (it does not hold at 4 × 4: 128 against 160).
• Double buffering helps exactly when something is left to overlap. It gives 1.18× at 1 × 1, 1.09× at 2 × 2, 1.27× at 4 × 4, 1.41× at 4 × 8, and almost nothing (even slightly negative) for a whole row, where the single-buffered version is already almost DMA-saturated (98%).
• The k outermost order is catastrophic here: 16,128 tiles moved (every partial sum goes out to DRAM and back for every k), 388,096 cycles, 10.4 times the best, and the matrix unit at 8%. This is the accelerator version of the CPU result of Chapter 24: the loop order that keeps the accumulator close wins.
• A simple model predicts the double-buffered rows. S.predicted_double_buffered_cycles says a step takes the longer of its loads and its products, plus the unoverlapped first loads and last stores. It is within 4% to 6% above the simulator for every double-buffered row in the table (the last column; it is always a little pessimistic) and within 12% in the checker's smaller runs. It is a model of a model: it explains the simulator, it does not validate it against anything real.

The 16-slot table shows what a small scratchpad does: the large blocks and the double-buffered versions of the medium ones do not fit, so the best feasible schedule is 2 × 4 single buffered (76,160 cycles), and the "best" there is a worse machine than the 64-slot one's best. Capacity is a trade between blocking (reuse) and double buffering (overlap): 2 × 2 with two buffers needs 16 slots and gives 87,040 cycles, 2 × 4 with one buffer needs 16 and gives 76,160.

What if the memory were faster?

bandwidth_sweep.py runs three double-buffered schedules (1 × 1, 2 × 4, 4 × 8) on machines whose DMA moves 8 to 64 words per cycle, with the setup cost of a transfer either 16 cycles or 4:

128 x 128 x 128, double buffered, 64 tile slots; cycles (matrix-unit utilization); the matrix unit alone needs 32768 cycles
DMA words/cycle, setup   tile cycles                1x1                2x4                4x8   1x1 / 4x8
       8 words, setup 16              24       206848 (16%)        79872 (41%)        43008 (76%)       4.81x
      16 words, setup 16              20       173056 (19%)        66560 (49%)        37440 (88%)       4.62x
      32 words, setup 16              18       156160 (21%)        59904 (55%)        36736 (89%)       4.25x
      64 words, setup 16              17       147712 (22%)        56576 (58%)        36384 (90%)       4.06x
       8 words, setup 4               12       105472 (31%)        39936 (82%)        34624 (95%)       3.05x
      16 words, setup 4                8        71680 (46%)        34048 (96%)        33216 (99%)       2.16x
      32 words, setup 4                6        54784 (60%)        33792 (97%)        33120 (99%)       1.65x
      64 words, setup 4                5        46336 (71%)        33664 (97%)        33072 (99%)       1.40x

Two lessons, both about the model's assumptions: raising the bandwidth barely helps while the per-transfer setup stays at 16 cycles (the gap between the worst and best schedule goes from 4.8× only down to 4.1× as bandwidth grows eightfold, because a tile takes 24 cycles at 8 words per cycle and still 17 at 64: the setup dominates), whereas cutting the setup to 4 cycles helps a great deal (at 16 words per cycle the gap falls to 2.2×, and 2 × 4 already reaches 96% utilization). So for this machine, larger or fewer transfers is a better lever than a faster bus. Whether a real device behaves this way depends on its real setup cost, which this chapter does not know.

Tests

One new lit file group test/ga40/ (the suite is now 152 tests; two files: check and timeline). check_ga.py checks, in order:

1. Results: 70 combinations (7 shapes, including sizes that are not multiples of the tile and a single row and column; 5 schedules; single and double buffering) each give exactly the plain triple loop's C.
2. Traffic: the tiles loaded and stored equal the formula in all 70 runs (blocked: per block, per k, bi + bj loads and one store per tile of C; kouter: two loads per product, a partial-sum load for every k after the first, a store per k).
3. Timing, from first principles: the six-instruction program takes 84 cycles; 4 × 4 double buffered on 64 × 64 × 64 takes exactly (tiles moved × 20) cycles; no schedule beats the matrix-unit bound or its own DMA bound; more reuse always moves fewer tiles; double buffering is faster for 1 × 1 and 4 × 4; the simple model is within 12%.
4. Hazards, independently: the simulator records when every instruction started and finished; the checker replays that trace, written separately from the timing rule, and verifies for every slot that no instruction read a value before the write that produced it (plus fill) was done and that no write began before the earlier reads of that slot had finished.

#!/usr/bin/env python3
"""Checks the Chapter 40 simulator and schedules.

  1. RESULTS: every schedule (blocked with several block shapes, with and without double buffering, and k-outermost) gives exactly the plain triple loop's C, for shapes that
     are and are not multiples of the tile size (including a single row and a single column);
  2. TRAFFIC: the number of tiles moved equals the formula (blocked: per block, per k, bi + bj loads, plus one store per tile of C; k-outermost: 2 loads per product, a load of the
     partial C for every k after the first, and a store for every k);
  3. TIMING, from first principles: a six-instruction program whose cycle count is worked out by hand; the 4 x 4 double-buffered schedule is DMA-bound (cycles = tiles moved x
     cycles per tile, exactly); no schedule beats the matrix-unit bound or the DMA bound; more reuse never moves more tiles;
  4. HAZARDS, independently: the simulator records when every instruction started and finished; this check replays the trace and verifies, for every slot, that no instruction read
     a value before the write that produced it had finished (plus the fill time for a non-accumulating reader), and that no write began before the earlier reads of that slot ended.
Usage: check_ga.py     Exit status 0 only if every check passes."""
import itertools, sys, random
from ga import Machine, Sim
import schedules as S
failures = 0
def report(ok, name, detail=""):
    global failures; failures += not ok
    print(f"  {'ok  ' if ok else 'FAIL'} {name}" + ("" if ok else f"\n       {detail}"))
m = Machine(slots=64)

def hazards_general(sim):
    """The same idea, written the straightforward way: for each slot, the sequence of its uses in program order; every use must respect the previous ones."""
    bad = []; last_write = {}; reads_since = {}
    for idx, eng, start, finish, rd, wr, fill in sim.trace:
        accs = set(rd) & set(wr)                          # accumulating (read-modify-write) slots
        for s in rd:
            if s in last_write:
                wf, wfill, was_acc = last_write[s]
                need = wf if (s in accs and was_acc) else wf + wfill      # an accumulating read after an accumulating write needs no fill
                if start < need: bad.append((idx, "read too early", s, start, need))
        for s in wr:
            for rf in reads_since.get(s, []):
                if start < rf: bad.append((idx, "write too early", s, start, rf))
        for s in rd: reads_since.setdefault(s, []).append(finish)
        for s in wr:
            last_write[s] = (finish, fill, s in accs); reads_since[s] = [finish] if s in accs else []
    return bad

print("-- 1. results")
shapes = [(8, 8, 8), (24, 16, 32), (20, 13, 30), (1, 17, 9), (9, 17, 1), (40, 40, 40), (64, 64, 64)]
schedules = [("blocked 1x1", lambda M, K, N, db: S.blocked(M, K, N, m, 1, 1, db)), ("blocked 2x2", lambda M, K, N, db: S.blocked(M, K, N, m, 2, 2, db)),
             ("blocked 1x4", lambda M, K, N, db: S.blocked(M, K, N, m, 1, 4, db)), ("blocked 3x2", lambda M, K, N, db: S.blocked(M, K, N, m, 3, 2, db)),
             ("k outermost", lambda M, K, N, db: S.kouter(M, K, N, m, db))]
bad_results = []; bad_traffic = []; bad_haz = []; count = 0; sims = {}
for (M, K, N), (name, build), db in itertools.product(shapes, schedules, (False, True)):
    A, B = S.matrices(M, K, N, seed=M + K + N); prog = build(M, K, N, db); sim, C = S.run(prog, A, B, m); count += 1; sims[(M, K, N, name, db)] = sim
    if C != S.reference(A, B): bad_results.append((M, K, N, name, db))
    Mt, Kt, Nt = (-(-M // 8), -(-K // 8), -(-N // 8))
    if name.startswith("blocked"):
        bi, bj = map(int, name.split()[1].split("x")); loads, stores = S.expected_blocked_tiles(M, K, N, m, bi, bj)
    else: loads, stores = 2 * Mt * Nt * Kt + Mt * Nt * (Kt - 1), Mt * Nt * Kt
    if (sim.dma_loads, sim.dma_stores) != (loads, stores): bad_traffic.append((M, K, N, name, db, (sim.dma_loads, sim.dma_stores), (loads, stores)))
    hz = hazards_general(sim)
    if hz: bad_haz.append((M, K, N, name, db, hz[:2]))
report(not bad_results, f"all {count} schedule/shape/buffering combinations give exactly the triple loop's result", str(bad_results[:3]))
print("-- 2. traffic")
report(not bad_traffic, f"tiles loaded and stored equal the formula in all {count} runs", str(bad_traffic[:3]))
print("-- 3. timing")
A = [[1.0] * 8 for _ in range(8)]; d = {"A": A, "B": A, "C": [[0.0] * 8 for _ in range(8)]}
t = Sim(m, d).run([("zero", 2), ("load", 0, "A", 0, 0), ("load", 1, "B", 0, 0), ("mxu", 2, 0, 1), ("mxu", 2, 0, 1), ("store", 2, "C", 0, 0)]).cycles
report(t == 84, f"a six-instruction program takes 84 cycles (2 loads of {m.dma_cycles} = 40, two chained products = 16, fill {m.mxu_fill}, store {m.dma_cycles}: 40 + 8 + 8 + 8 + 20)", str(t))
A128, B128 = S.matrices(64, 64, 64); res = {}
for bi, bj, db in [(1, 1, False), (1, 1, True), (2, 2, True), (4, 4, False), (4, 4, True)]:
    sim, C = S.run(S.blocked(64, 64, 64, m, bi, bj, db), A128, B128, m); res[(bi, bj, db)] = sim
sim44 = res[(4, 4, True)]
report(sim44.cycles == sim44.dma_tiles * m.dma_cycles, f"4x4 double buffered on 64x64x64 is DMA-bound: {sim44.cycles} cycles = {sim44.dma_tiles} tiles x {m.dma_cycles}", f"{sim44.cycles} vs {sim44.dma_tiles * m.dma_cycles}")
lb = S.lower_bounds(64, 64, 64, m)
report(all(s.cycles >= lb["compute"] and s.cycles >= s.dma_tiles * m.dma_cycles for s in res.values()), f"no schedule beats the matrix-unit bound ({lb['compute']}) or its own DMA bound", str({k: v.cycles for k, v in res.items()}))
tiles_by_block = [S.expected_blocked_tiles(64, 64, 64, m, bi, bj) for bi, bj in [(1, 1), (1, 2), (2, 2), (2, 4), (4, 4)]]
report(all(x[0] > y[0] for x, y in zip(tiles_by_block, tiles_by_block[1:])), "more reuse moves fewer tiles: loads fall at every step from 1x1 to 1x2, 2x2, 2x4, 4x4: " + " > ".join(str(x[0]) for x in tiles_by_block), str(tiles_by_block))
report(res[(1, 1, True)].cycles < res[(1, 1, False)].cycles and res[(4, 4, True)].cycles < res[(4, 4, False)].cycles, "double buffering makes 1x1 and 4x4 faster", str({k: v.cycles for k, v in res.items()}))
errs = {}
for bi, bj in [(1, 1), (1, 8), (2, 2), (2, 4), (4, 4), (2, 8)]:
    sim, _ = S.run(S.blocked(64, 64, 64, m, bi, bj, True), A128, B128, m); errs[(bi, bj)] = S.predicted_double_buffered_cycles(64, 64, 64, m, bi, bj) / sim.cycles
report(all(1.0 <= e <= 1.12 for e in errs.values()), "the simple model (a step takes the longer of its loads and its products; first loads and stores do not overlap) is always a little ABOVE the simulator and within 12% for the double-buffered cycles: " + ", ".join(f"{k[0]}x{k[1]}: {v:.2f}" for k, v in errs.items()), str(errs))
print("-- 4. hazards (independent replay of the trace)")
report(not bad_haz, f"no read-before-write or write-before-read hazard in the traces of all {count} runs", str(bad_haz[:2]))
print("all checks pass" if failures == 0 else f"{failures} CHECK(S) FAILED"); sys.exit(1 if failures else 0)

(The listing is the checker itself.) Its output:

-- 1. results
  ok   all 70 schedule/shape/buffering combinations give exactly the triple loop's result
-- 2. traffic
  ok   tiles loaded and stored equal the formula in all 70 runs
-- 3. timing
  ok   a six-instruction program takes 84 cycles (2 loads of 20 = 40, two chained products = 16, fill 8, store 20: 40 + 8 + 8 + 8 + 20)
  ok   4x4 double buffered on 64x64x64 is DMA-bound: 6400 cycles = 320 tiles x 20
  ok   no schedule beats the matrix-unit bound (4096) or its own DMA bound
  ok   more reuse moves fewer tiles: loads fall at every step from 1x1 to 1x2, 2x2, 2x4, 4x4: 1024 > 768 > 512 > 384 > 256
  ok   double buffering makes 1x1 and 4x4 faster
  ok   the simple model (a step takes the longer of its loads and its products; first loads and stores do not overlap) is always a little ABOVE the simulator and within 12% for the double-buffered cycles: 1x1: 1.07, 1x8: 1.11, 2x2: 1.11, 2x4: 1.11, 4x4: 1.10, 2x8: 1.10
-- 4. hazards (independent replay of the trace)
  ok   no read-before-write or write-before-read hazard in the traces of all 70 runs
all checks pass

Are the checks good enough? Eleven wrong simulators.

sim_mutation.py breaks ga.py eleven ways and runs the checker on each:

NOTE: this script deliberately breaks the simulator. 'caught' lines are EXPECTED: they show the checker can detect the mistake.
baseline (nothing broken): checker passes (expected)
a write does not wait for the earlier READS of its slot (write-after-read hazard dropped): caught by 2 failing check(s) (double buffering, the independent hazard replay)
an instruction waits for only the FIRST of its source slots (the second operand is not waited for): caught by 3 failing check(s) (the hand-timed program, double buffering, the independent hazard replay)
a DMA transfer has no setup cost: caught by 3 failing check(s) (the hand-timed program, the DMA-bound schedule, other)
chained accumulating products wait for the fill time as well: caught by 1 failing check(s) (the hand-timed program)
the product's fill time is zero: caught by 1 failing check(s) (the hand-timed program)
each store is counted as two tile transfers: caught by 1 failing check(s) (traffic)
a tile that sticks out of the matrix wraps around instead of being padded with zeros: caught by 1 failing check(s) (the results)
a store writes the padding too (writes past the edge): caught (the checker did not finish: IndexError: list assignment index out of range)
the matrix unit multiplies with A transposed: caught by 1 failing check(s) (the results)
all engines run in lock step (an instruction waits for every engine to be free): caught by 4 failing check(s) (the hand-timed program, the DMA-bound schedule, double buffering, other)
a load does not wait for the slot to become free (no WAW either): caught by 2 failing check(s) (double buffering, the independent hazard replay)

restored: ga.py is back to its original

All eleven are caught. Which check caught which is informative: the independent hazard replay catches exactly the hazard bugs (a write that does not wait for earlier reads; an instruction that waits for only one of its operands; a load that ignores the slot's earlier use), and only it and the double-buffering check notice them, because the results are computed in program order and do not depend on timing; the hand-timed program catches the cost bugs (no DMA setup, a wrong fill, accumulation waiting for the fill); the traffic formula catches a double-counted store; the results catch the padding and the transposed-operand bugs; and the DMA-bound equality catches the lock-step engines.

The full suite:

-- Testing: 152 tests, 4 workers --
PASS: IR Forge / Mountain Goat :: lm36/against-reference.mlir (1 of 152)
PASS: IR Forge / Mountain Goat :: profile39/claims.mlir (2 of 152)
PASS: IR Forge / Mountain Goat :: lm35/against-reference.mlir (3 of 152)
PASS: IR Forge / Mountain Goat :: lm36/examples.mlir (4 of 152)
PASS: IR Forge / Mountain Goat :: generalize32/against-reference.mlir (5 of 152)
PASS: IR Forge / Mountain Goat :: attention33/against-reference.mlir (6 of 152)
PASS: IR Forge / Mountain Goat :: lm35/examples.mlir (7 of 152)
PASS: IR Forge / Mountain Goat :: attention33/examples.mlir (8 of 152)
PASS: IR Forge / Mountain Goat :: transformer30/against-reference.mlir (9 of 152)
PASS: IR Forge / Mountain Goat :: contraction28/against-definition.mlir (10 of 152)
PASS: IR Forge / Mountain Goat :: contraction/tensor-contraction.mlir (11 of 152)
PASS: IR Forge / Mountain Goat :: ffn34/examples.mlir (12 of 152)
PASS: IR Forge / Mountain Goat :: generalize32/examples.mlir (13 of 152)
PASS: IR Forge / Mountain Goat :: outline38/results.mlir (14 of 152)
PASS: IR Forge / Mountain Goat :: llvm37/claims.mlir (15 of 152)
PASS: IR Forge / Mountain Goat :: attention29/against-reference.mlir (16 of 152)
PASS: IR Forge / Mountain Goat :: tour/gallery.mlir (17 of 152)
PASS: IR Forge / Mountain Goat :: learning/gradient-descent.mlir (18 of 152)
PASS: IR Forge / Mountain Goat :: bigram31/against-reference.mlir (19 of 152)
PASS: IR Forge / Mountain Goat :: attention29/examples.mlir (20 of 152)
PASS: IR Forge / Mountain Goat :: contraction28/examples.mlir (21 of 152)
PASS: IR Forge / Mountain Goat :: broadcast/values.mlir (22 of 152)
PASS: IR Forge / Mountain Goat :: ga40/check.mlir (23 of 152)
PASS: IR Forge / Mountain Goat :: tour/tour-examples.mlir (24 of 152)
PASS: IR Forge / Mountain Goat :: perf/options-same-results.mlir (25 of 152)
PASS: IR Forge / Mountain Goat :: loop-order/results-identical.mlir (26 of 152)
PASS: IR Forge / Mountain Goat :: transformer30/examples.mlir (27 of 152)
PASS: IR Forge / Mountain Goat :: tour/more-examples.mlir (28 of 152)
PASS: IR Forge / Mountain Goat :: bigram31/examples.mlir (29 of 152)
PASS: IR Forge / Mountain Goat :: ops/values.mlir (30 of 152)
PASS: IR Forge / Mountain Goat :: loop-order/bitwise-identical.mlir (31 of 152)
PASS: IR Forge / Mountain Goat :: contraction/mismatch-trap.mlir (32 of 152)
PASS: IR Forge / Mountain Goat :: outline38/structure.mlir (33 of 152)
PASS: IR Forge / Mountain Goat :: broadcast/cpp-interop-reductions.mlir (34 of 152)
PASS: IR Forge / Mountain Goat :: ops/cpp-interop-ops.mlir (35 of 152)
PASS: IR Forge / Mountain Goat :: llvm37/examples.mlir (36 of 152)
PASS: IR Forge / Mountain Goat :: execution/gpu-host-with-stub-runtime.mlir (37 of 152)
PASS: IR Forge / Mountain Goat :: loop-order/cpp-ikj.mlir (38 of 152)
PASS: IR Forge / Mountain Goat :: chapters/ch09-execution-engine.mlir (39 of 152)
PASS: IR Forge / Mountain Goat :: ops/elementwise-mismatch-aborts.mlir (40 of 152)
PASS: IR Forge / Mountain Goat :: frontend/more-examples.mlir (41 of 152)
PASS: IR Forge / Mountain Goat :: broadcast/identities.mlir (42 of 152)
PASS: IR Forge / Mountain Goat :: frontend/cpp-interop.mlir (43 of 152)
PASS: IR Forge / Mountain Goat :: frontend/runtime-mismatch-aborts.mlir (44 of 152)
PASS: IR Forge / Mountain Goat :: perf/flags-reach-clang.mlir (45 of 152)
PASS: IR Forge / Mountain Goat :: ops/negate-and-dynamic-divide.mlir (46 of 152)
PASS: IR Forge / Mountain Goat :: perf/tiling-structure.mlir (47 of 152)
PASS: IR Forge / Mountain Goat :: assert-stderr/message-on-stderr.mlir (48 of 152)
PASS: IR Forge / Mountain Goat :: frontend/hello-and-static.mlir (49 of 152)
PASS: IR Forge / Mountain Goat :: loop-order/flag-reaches-pass.mlir (50 of 152)
PASS: IR Forge / Mountain Goat :: broadcast/networks.mlir (51 of 152)
PASS: IR Forge / Mountain Goat :: frontend/dynamic-runs.mlir (52 of 152)
PASS: IR Forge / Mountain Goat :: tour/scientific-literals.mlir (53 of 152)
PASS: IR Forge / Mountain Goat :: ops/dynamic-ops.mlir (54 of 152)
PASS: IR Forge / Mountain Goat :: ops/precedence-and-identities.mlir (55 of 152)
PASS: IR Forge / Mountain Goat :: ops/composed-programs.mlir (56 of 152)
PASS: IR Forge / Mountain Goat :: broadcast/dynamic-reductions.mlir (57 of 152)
PASS: IR Forge / Mountain Goat :: contraction28/errors.mlir (58 of 152)
PASS: IR Forge / Mountain Goat :: execution/mismatch-columns-aborts.mlir (59 of 152)
PASS: IR Forge / Mountain Goat :: execution/mismatch-aborts.mlir (60 of 152)
PASS: IR Forge / Mountain Goat :: tour/tour-errors.mlir (61 of 152)
PASS: IR Forge / Mountain Goat :: ops/ptx-new-ops.mlir (62 of 152)
PASS: IR Forge / Mountain Goat :: transformer30/errors.mlir (63 of 152)
PASS: IR Forge / Mountain Goat :: execution/static-add-runs.mlir (64 of 152)
PASS: IR Forge / Mountain Goat :: execution/dynamic-shapes-run.mlir (65 of 152)
PASS: IR Forge / Mountain Goat :: ops/matmul-mismatch-aborts.mlir (66 of 152)
PASS: IR Forge / Mountain Goat :: dynamic-loops/tiling-runs.mlir (67 of 152)
PASS: IR Forge / Mountain Goat :: dynamic-loops/runtime-check-survives-transforms.mlir (68 of 152)
PASS: IR Forge / Mountain Goat :: chapters/ch01-pipeline.mlir (69 of 152)
PASS: IR Forge / Mountain Goat :: execution/mixed-static-dynamic.mlir (70 of 152)
PASS: IR Forge / Mountain Goat :: dynamic-loops/unroll-runs.mlir (71 of 152)
PASS: IR Forge / Mountain Goat :: ops/division-ieee.mlir (72 of 152)
PASS: IR Forge / Mountain Goat :: attention29/errors.mlir (73 of 152)
PASS: IR Forge / Mountain Goat :: attention33/errors.mlir (74 of 152)
PASS: IR Forge / Mountain Goat :: ffn34/errors.mlir (75 of 152)
PASS: IR Forge / Mountain Goat :: frontend/ptx-from-source.mlir (76 of 152)
PASS: IR Forge / Mountain Goat :: assert-stderr/nested-assert-runs.mlir (77 of 152)
PASS: IR Forge / Mountain Goat :: execution/strided-view.mlir (78 of 152)
PASS: IR Forge / Mountain Goat :: lm36/errors.mlir (79 of 152)
PASS: IR Forge / Mountain Goat :: assert-stderr/valid-shapes-still-pass.mlir (80 of 152)
PASS: IR Forge / Mountain Goat :: dynamic-loops/fusion-after-tiling.mlir (81 of 152)
PASS: IR Forge / Mountain Goat :: lm35/errors.mlir (82 of 152)
PASS: IR Forge / Mountain Goat :: assert-stderr/message-bytes-exact.mlir (83 of 152)
PASS: IR Forge / Mountain Goat :: generalize32/errors.mlir (84 of 152)
PASS: IR Forge / Mountain Goat :: contraction28/ir.mlir (85 of 152)
PASS: IR Forge / Mountain Goat :: bigram31/errors.mlir (86 of 152)
PASS: IR Forge / Mountain Goat :: gpu/pass-order-crash.mlir (87 of 152)
PASS: IR Forge / Mountain Goat :: dynamic-loops/fusion-needs-constant-trip-counts.mlir (88 of 152)
PASS: IR Forge / Mountain Goat :: broadcast/ptx-layer.mlir (89 of 152)
PASS: IR Forge / Mountain Goat :: gpu/ptx-bare-pointers.mlir (90 of 152)
PASS: IR Forge / Mountain Goat :: gpu/ptx-default.mlir (91 of 152)
PASS: IR Forge / Mountain Goat :: gpu/host-runtime-calls-bare.mlir (92 of 152)
PASS: IR Forge / Mountain Goat :: frontend/front-end-errors.mlir (93 of 152)
PASS: IR Forge / Mountain Goat :: gpu/ptx-index32.mlir (94 of 152)
PASS: IR Forge / Mountain Goat :: ga40/timeline.mlir (95 of 152)
PASS: IR Forge / Mountain Goat :: ops/front-end-errors.mlir (96 of 152)
PASS: IR Forge / Mountain Goat :: gpu/host-runtime-calls.mlir (97 of 152)
PASS: IR Forge / Mountain Goat :: dynamic-loops/unroll-full-needs-constant-trip-count.mlir (98 of 152)
PASS: IR Forge / Mountain Goat :: broadcast/front-end-errors.mlir (99 of 152)
PASS: IR Forge / Mountain Goat :: generalize32/ir.mlir (100 of 152)
PASS: IR Forge / Mountain Goat :: loop-order/structure.mlir (101 of 152)
PASS: IR Forge / Mountain Goat :: gpu/nvvm-lowering.mlir (102 of 152)
PASS: IR Forge / Mountain Goat :: bigram31/ir.mlir (103 of 152)
PASS: IR Forge / Mountain Goat :: chapters/ch02-example-programs.mlir (104 of 152)
PASS: IR Forge / Mountain Goat :: broadcast/roundtrip-and-lowering.mlir (105 of 152)
PASS: IR Forge / Mountain Goat :: outline38/pass.mlir (106 of 152)
PASS: IR Forge / Mountain Goat :: assert-stderr/guard-nested-assert.mlir (107 of 152)
PASS: IR Forge / Mountain Goat :: dynamic-loops/fusion-fires-when-static.mlir (108 of 152)
PASS: IR Forge / Mountain Goat :: ops/matmul-lowering-structure.mlir (109 of 152)
PASS: IR Forge / Mountain Goat :: dynamic-loops/fusion-declines.mlir (110 of 152)
PASS: IR Forge / Mountain Goat :: gpu/kernel-outlining.mlir (111 of 152)
PASS: IR Forge / Mountain Goat :: frontend/emit-mlir.mlir (112 of 152)
PASS: IR Forge / Mountain Goat :: ops/elementwise-static-mismatch.mlir (113 of 152)
PASS: IR Forge / Mountain Goat :: transformer30/ir.mlir (114 of 152)
PASS: IR Forge / Mountain Goat :: lowering/dynamic-transpose-no-check.mlir (115 of 152)
PASS: IR Forge / Mountain Goat :: assert-stderr/static-unchanged.mlir (116 of 152)
PASS: IR Forge / Mountain Goat :: gpu/cuda-comparison.mlir (117 of 152)
PASS: IR Forge / Mountain Goat :: dynamic-loops/unroll-epilogue.mlir (118 of 152)
PASS: IR Forge / Mountain Goat :: attention29/ir.mlir (119 of 152)
PASS: IR Forge / Mountain Goat :: broadcast/verifier-reduce.mlir (120 of 152)
PASS: IR Forge / Mountain Goat :: dynamic-loops/tiling-structure.mlir (121 of 152)
PASS: IR Forge / Mountain Goat :: assert-stderr/rewrite-structure.mlir (122 of 152)
PASS: IR Forge / Mountain Goat :: assert-stderr/run-twice-no-collision.mlir (123 of 152)
PASS: IR Forge / Mountain Goat :: loops/unroll-full.mlir (124 of 152)
PASS: IR Forge / Mountain Goat :: ops/matmul-inner-mismatch.mlir (125 of 152)
PASS: IR Forge / Mountain Goat :: ops/matmul-bad-result-shape.mlir (126 of 152)
PASS: IR Forge / Mountain Goat :: broadcast/verifier-broadcast.mlir (127 of 152)
PASS: IR Forge / Mountain Goat :: loops/tiling.mlir (128 of 152)
PASS: IR Forge / Mountain Goat :: verifier/transpose-rank3.mlir (129 of 152)
PASS: IR Forge / Mountain Goat :: contraction28/verifiers.mlir (130 of 152)
PASS: IR Forge / Mountain Goat :: canonicalize/fold-constant-add.mlir (131 of 152)
PASS: IR Forge / Mountain Goat :: lowering/mixed-add-one-check.mlir (132 of 152)
PASS: IR Forge / Mountain Goat :: lowering/static-add-affine.mlir (133 of 152)
PASS: IR Forge / Mountain Goat :: verifier/add-static-mismatch.mlir (134 of 152)
PASS: IR Forge / Mountain Goat :: loops/fusion.mlir (135 of 152)
PASS: IR Forge / Mountain Goat :: ops/scalar-bad-op.mlir (136 of 152)
PASS: IR Forge / Mountain Goat :: verifier/transpose-bad-shape.mlir (137 of 152)
PASS: IR Forge / Mountain Goat :: canonicalize/transpose-twice-dynamic-exact.mlir (138 of 152)
PASS: IR Forge / Mountain Goat :: canonicalize/transpose-twice-dynamic-mixed.mlir (139 of 152)
PASS: IR Forge / Mountain Goat :: lowering/bufferize-dynamic-add.mlir (140 of 152)
PASS: IR Forge / Mountain Goat :: verifier/unranked-rejected.mlir (141 of 152)
PASS: IR Forge / Mountain Goat :: lowering/dynamic-add-affine.mlir (142 of 152)
PASS: IR Forge / Mountain Goat :: attention29/verifiers.mlir (143 of 152)
PASS: IR Forge / Mountain Goat :: loop-order/bad-option.mlir (144 of 152)
PASS: IR Forge / Mountain Goat :: verifier/dynamic-accepted.mlir (145 of 152)
PASS: IR Forge / Mountain Goat :: ops/dynamic-accepted.mlir (146 of 152)
PASS: IR Forge / Mountain Goat :: bigram31/verifiers.mlir (147 of 152)
PASS: IR Forge / Mountain Goat :: canonicalize/transpose-twice.mlir (148 of 152)
PASS: IR Forge / Mountain Goat :: lowering/bufferize-static-add.mlir (149 of 152)
PASS: IR Forge / Mountain Goat :: generalize32/verifiers.mlir (150 of 152)
PASS: IR Forge / Mountain Goat :: transformer30/verifiers.mlir (151 of 152)
PASS: IR Forge / Mountain Goat :: ffn34/against-reference.mlir (152 of 152)

Testing Time: 113.03s

Total Discovered Tests: 152
  Passed: 152 (100.00%)

Limits and what is not established

• A model of a machine that does not exist. No real device calibrated the parameters. The matrix unit's rate, the DMA's setup and bandwidth, the scratchpad sizes and the tile size are inventions with plausible ratios; the conclusions that depend on the ratios (for example that 4 × 8 is the first compute-bound block, or that setup cost dominates bandwidth) hold for GA-1 and may not for any real chip.
• What is not modelled: instruction fetch and decode, a real memory hierarchy (DRAM is a flat store with a constant cost per transfer: no banks, no contention, no refresh, no row buffers), more than one DMA channel, tiles that are not square or not T × T, energy and area, numerical formats (everything is a Python number), sparsity, and any cost for the host or for launching a program.
• Partial tiles cost a full tile. A tile that sticks out of the matrix is padded with zeros and costs a whole transfer and a whole product, so the model charges padding waste (Chapter 41 shows what that does to small matrices).
• The schedules are hand-written and not searched. Only block shapes 1 × 1 to 4 × 8 and two buffer counts were tried; there is no claim that any is optimal, and the inter-block overlap is not optimized (each block's stores and the next block's loads do not overlap).
• The analytic model is only checked against the simulator, not against hardware.
• Pure Python, small sizes. The largest product is 128 × 128 × 128.

Reproducing

cd docs/part40/code
./sweep_orders.py > sweep_orders_out.txt && ./timeline.py > timeline_out.txt && ./bandwidth_sweep.py > bandwidth_sweep_out.txt
./check_ga.py && ./sim_mutation.py > sim_mutation_out.txt
cd ../../part15/code && ./run_lit.sh          # 152 tests

Chapter summary

• GA-1 is a small, readable model of a matrix accelerator: DMA, scratchpad, matrix unit, independent in-order queues, and a timing rule that enforces hazards. Every instruction computes real values and every run is checked against a plain triple loop.
• With ratios in which the matrix unit is much faster than the memory, the plain i, j, k order is DMA-bound (19% of peak); blocking cuts the tiles moved from 8,448 to 1,792, and a 4 × 8 block with double buffering reaches 88% (it becomes compute-bound); the k-outermost order is 10.4× slower than the best.
• Double buffering helps exactly when there is a compute phase to hide the loads behind; scratchpad capacity forces a choice between blocking and double buffering.
• In this model per-transfer setup cost matters more than bandwidth.
• A simple analytic model of a double-buffered step (the longer of loads and products) matches the simulator within 4% to 6% on the table; an independent replay of the instruction trace verifies the hazards; eleven broken simulators are caught.
• It is a model: its numbers describe GA-1, not any real chip.

Self-check questions

Each answer is collapsed; try the question first.

1. Why is the plain 1 × 1 schedule so much worse than 4 × 4, with the same arithmetic?

   Answer

   The arithmetic is identical (4,096 tile products) but 1 × 1 loads two tiles per product (8,448 tiles moved in all) while 4 × 4 loads eight tiles for sixteen products (2,304 moved). With the DMA at 20 cycles per tile, the plain schedule keeps the DMA busy for 98% of 173,056 cycles while the matrix unit idles at 19%; 4 × 4 reuses each loaded tile four times.

2. When does a blocked schedule stop being memory-bound on GA-1?

   Answer

   When a step's products take at least as long as its loads: 8 · bi · bj ≥ 20 · (bi + bj). That holds first at 4 × 8 among the shapes tried (256 against 240 cycles per step); at 4 × 4 it fails (128 against 160), which is why 4 × 4 is DMA-bound at exactly tiles × 20 cycles and 4 × 8 reaches 88% utilization. The thresholds follow from the invented parameters; other ratios would move them.

3. What does double buffering need, and what does it buy?

   Answer

   A second set of A and B slots (bi + bj more), so that the DMA can load step k + 1 while the matrix unit works on step k. It buys overlap: 18% at 1 × 1, 27% at 4 × 4, 41% at 4 × 8. It buys almost nothing when the DMA is already saturated by a single buffer, and it can make a schedule not fit: with 16 slots the best choice was 2 × 4 single buffered.

4. Why is the k-outermost order 10 times slower?

   Answer

   The accumulator does not stay on chip. For every k each tile of C is loaded from DRAM, updated by one product and stored again: 16,128 tiles moved against 1,792 for 4 × 8, and the DMA queue, being in order, also stalls behind each store that waits for its product. It is the same lesson as Chapter 24's ijk against ikj, here with the accumulator in off-chip memory instead of a cache line.

5. The checker's hazard replay and the results check are separate. What does each catch that the other cannot?

   Answer

   Results are computed in program order, so they cannot reveal a timing hazard: a simulator that lets a load overwrite a slot too early still produces the right numbers. Only the replay of the recorded start and finish times (written independently of the timing rule) sees it. Conversely, the replay says nothing about whether a tile was padded or multiplied correctly; the results check does.

6. What does the chapter's bandwidth sweep say, and why must it be read with care?

   Answer

   That in GA-1 the fixed per-transfer setup (16 cycles) dominates a 64-word tile, so a faster bus changes little (the gap between plain and blocked stays about 4×) but a smaller setup helps a lot. It must be read with care because setup and bandwidth are invented numbers: it shows how to ask the question and what shape the answer takes, not what any real device does.

7. Name three things GA-1 does not model that would matter on a real accelerator.

   Answer

   Any of: a real memory hierarchy (banks, contention, row buffers, more than one DMA channel), instruction fetch and decode, energy and area, numerical formats such as 8-bit integers or 16-bit floats, sparsity, non-square or differently sized tiles, and launch overhead from a host. The chapter lists these as limits.