all_lessons/ai_compilers / lessons/08 · schedulinglesson 9 / 20

Part IV — Back end

Scheduling — the loop nest

Lesson 07 finished the middle end: the graph is canonicalized, fused, memory-planned, and every tensor has a physical layout chosen to coalesce and feed the tensor cores. But that graph is still declarative. A node that says C = A @ B specifies the math — every output element is a dot product over K — and nothing else. It does not say which loop is outermost, how big the working tile is, what gets staged into SRAM, or where to vectorize. Those choices do not change a single FLOP or a single output number, and yet for one fixed matmul they swing wall-clock by 10–100×. This lesson is about that gap: turning a what into a how, which is choosing a loop nest. The key move is Halide's: cleave the program into an algorithm (the math) and a schedule (the loop strategy), so you can search the second without ever touching the first.

The previous step left this broken
Layout assignment fixed the byte order of each tensor, but a tensor with a chosen layout is still just data; nothing yet says how to walk it. Take the laid-out matmul C = A·B with (M, N, K) = (1,024, 1,024, 1,024). The naive realization is the textbook triple loop — for each i, for each j, accumulate over k. That executes exactly 2·M·N·K ≈ 2.15 GFLOP and re-reads the entire A row and B column from HBM for every output element, so the same matrix data crosses the memory hierarchy hundreds of times. A tiled version computes the identical 2.15 GFLOP but loads each tile of A and B into SRAM once and reuses it across the whole tile of C — the same math, a tiny fraction of the traffic, and an order of magnitude faster. The middle end has no vocabulary for this: the graph IR records that a matmul happens, never the loop nest that runs it. That loop nest is the next decision, and it is the one that actually sets the speed.
Linear position
Forced by: a laid-out graph still only says what to compute; one fixed computation has astronomically many loop nests (orders × tile sizes × vectorization), and the choice swings performance by 10–100× at identical FLOPs.
New idea: the algorithm/schedule split (Halide) — separate what (the math) from when/where (loop order, tiling, parallelism). The schedule knobs: tile/block, reorder, vectorize, unroll, parallelize, fuse loops (compute_at), cache/stage to SRAM. Tiling for the memory hierarchy is the big one. The polyhedral model formalizes which loop transforms are legal: a transform is legal iff it preserves every data dependence.
Forces next: a chosen schedule is still a plan, not instructions. You must turn it into real machine code for a real target — that is codegen (lesson 09) — and then, because the space is too large to hand-pick per shape, search it (autotuning, lesson 10).
The plan
Five moves. (1) The matmul as a loop nest: same FLOPs, wildly different locality, priced in HBM traffic (cross-link the tiled-matmul kernel lesson). (2) Halide's split — algorithm fixed, schedule a separate program — shown as two snippets. (3) The knobs one at a time, each with what it buys (tile→locality, vectorize→SIMD/SIMT lanes, unroll→ILP, parallelize→occupancy, compute_at→fusion). (4) Tiling for the memory hierarchy as the reuse story made quantitative. (5) The polyhedral model in one tight paragraph — affine loop nests, dependences, legal = preserves dependences. Then drive the explorer until no-tiling sits far below a tiled schedule, and tiles-too-big spill.

1 · One math, exponentially many loop nests

A schedule is the loop nest that realizes an op: the order of the loops, how each loop is split into blocks, which loops run on which hardware parallel unit, and what intermediate data is held where. The math is invariant under all of it — every legal schedule computes the same numbers — but the schedule decides how many times each byte crosses each level of the memory hierarchy, and that is what sets the time.

Take C = A·B, (M, N, K) = (1,024, 1,024, 1,024), fp32. The work is fixed at 2·M·N·K ≈ 2.15 GFLOP and the inputs are 2·1,024²·4 ≈ 8.4 MB. Now look at what different loop nests cost in HBM traffic:

scheduleHBM bytes read for A, B (≈)reuse per loaded byteverdict
ijk, no tilingeach C element re-reads a full A row + B col → ≈ 2·M·N·K·4 ≈ 8.6 GB≈ 1×memory-bound, ~100× off peak
ikj reorder, no tilingbetter stream order but still re-streams B → multiple GBlowhelps a little, still no reuse
tiled T = 32each tile loaded once, reused T times → ≈ 8.4 MB · (K/T) ≈ 270 MB≈ T = 32×fits SRAM, ~tens× faster
tiled T = 256 (too big)tiles spill SRAM → falls back toward untiled trafficcollapsesregister/SMEM spill, slow again

Same 2.15 GFLOP in every row; the HBM bill ranges from ≈ 8.6 GB down to ≈ 270 MB — a ≈ 30× spread — purely from the loop nest. And that is before vectorization, unrolling, and the GPU's thread mapping, each of which multiplies the spread again. The number of distinct legal nests for even this one matmul (every loop order, every tile size at every level, every split point) is astronomically large. You are not choosing whether to compute the matmul; you are choosing, out of millions of equivalent realizations, the one the chip likes. This is exactly the tiled-matmul mechanic derived kernel-side in gpu_kernels · 05 shared memory & tiled matmul — here it is one knob in a compiler's search space.

2 · The algorithm/schedule split (Halide)

The decisive idea, from Halide (Ragan-Kelley et al., 2013) and inherited by TVM, Tiramisu, and the spirit of every modern tensor compiler, is to refuse to tangle the two decisions. You write the algorithm — a pure, index-based statement of what each output is — once, and you write the schedule — the loop strategy — as a separate program over that same algorithm. Change the schedule and you change only performance; the algorithm guarantees the result is identical. That separation is what makes the schedule searchable: you can try a thousand schedules against one fixed, trusted definition of correctness.

algorithm — what (never changes)
// pure definition: each C(i,j) is a
// reduction over k. No loop order,
// no tiles, no hardware mentioned.
Func C("C");
Var  i, j;  RDom k(0, K);

C(i, j)  = 0.0f;
C(i, j) += A(i, k) * B(k, j);
//  ^ the math, and only the math.
schedule — how (free to tune)
// the loop strategy, written separately:
Var ii, jj;                  // tile vars
C.tile(i, j, ii, jj, 32, 32) // block 32×32
 .reorder(ii, jj, k)         // loop order
 .vectorize(ii, 8)           // 8-lane SIMD
 .unroll(jj, 4)              // ILP
 .parallel(j);               // across cores
// swap any line → same answer,
// different speed.

Two payoffs. First, portability of the algorithm: the same one-line math compiles to a CPU schedule (vectorize + parallelize over cores) or a GPU schedule (map tiles to thread blocks, stage to shared memory) by swapping the schedule program, never the definition. Second, safety: because the schedule is just a reordering/blocking of a fixed reduction, the compiler can prove it preserves the math (§5), so an automated search can fire off candidate schedules without anyone re-checking correctness. The schedule is, in effect, the compiler's set of knobs — and the rest of this lesson is the knobs.

3 · The knobs, and what each one buys

Every schedule is built from a small vocabulary of loop transforms. Each attacks a different bottleneck; the art is composing them so the gains stack instead of fighting.

tile / block (split)
Cut a loop of length N into an outer loop of N/T blocks and an inner loop of T. Buys locality: the inner block's working set fits a fast level (registers/SRAM), so each loaded byte is reused ≈ T times before eviction. The single highest-leverage knob.
reorder
Permute loop nesting. Buys stride-1 access and changes which tensor is reused in the inner loop — turning a strided, uncoalesced sweep into a contiguous one without moving any data.
vectorize
Map an inner loop onto SIMD lanes (AVX) or SIMT threads (a warp). Buys width: 8 / 16 / 32 elements per instruction instead of one — the difference between using one lane and all of them.
unroll
Replicate the loop body k times, removing branch/index overhead and exposing independent instructions. Buys ILP — the pipeline issues several FMAs in flight and hides latency.
parallelize
Run an outer loop's iterations on separate cores / SMs / thread blocks. Buys occupancy: it fills the machine's parallel units so compute and memory latency overlap across many independent tiles.
compute_at / fuse loops
Compute a producer's values inside the consumer's loop nest, at a chosen level, so the intermediate lives in registers/SRAM and is never written to HBM. Buys fusion — this is the schedule-level realization of operator fusion from lesson 05.

Two clarifications worth pinning down. Vectorize and parallelize are not the same axis: vectorize widens a single instruction across lanes; parallelize spreads independent iterations across execution units. On a GPU you use both — the warp's 32 threads are the "vector," the grid of blocks is the "parallel." compute_at is the knob that connects scheduling back to fusion: where you place a producer's computation in the consumer's loops decides whether its output is a materialized HBM buffer (compute it at the root) or a transient register value recomputed per tile (compute it deep inside). That choice is exactly the fusion-vs-materialize trade of lesson 05, now expressed as a loop position. And every knob has a cliff: tile too big and the working set spills (§4); unroll too far and you blow the instruction cache or register file; vectorize a non-contiguous axis and you get gathers instead of wide loads.

4 · Tiling for the memory hierarchy

Tiling is worth its own section because it is the knob that makes the others matter, and because it is the same reuse story the whole track keeps telling — now made quantitative. The memory hierarchy is a stack of ever-smaller, ever-faster levels: HBM (≈3.3 TB/s, hundreds of GB), L2 (tens of MB), shared memory / L1 (≈100–200 KB per SM at ≈20× HBM bandwidth), registers (KB per thread, fastest of all). A byte read from HBM is "expensive"; the same byte reused from SRAM is nearly free. Tiling structures the loops so that once a block of data is hauled up a level, it is used as many times as possible before being evicted.

Quantify it for the matmul. A naive nest re-reads inputs for every output element: arithmetic intensity ≈ 1 FLOP per byte, far below the GPU's ridge point (~tens to hundreds of FLOP/byte), so it is hopelessly memory-bound — the tensor cores starve while the kernel waits on HBM (the roofline from lesson 01 and gpu_kernels · 02 memory hierarchy). Tile with block size T: each T×T tile of A and B is loaded into SRAM once and reused across the T output columns/rows it touches, so the bytes-per-FLOP drops by ≈ T and arithmetic intensity rises to ≈ T. Push T from 1 to 32 and you climb from far below the ridge to above it — the op crosses from memory-bound to compute-bound, which is the entire point.

But there is a ceiling. The tiles for A and B must fit the fast level simultaneously: two fp32 tiles cost 2·T²·4 bytes of SRAM, and there is only ~100–200 KB per SM. At T = 64 that is 2·64²·4 = 32 KB — already a large slice — and a thread-per-element mapping wants T² = 4,096 threads, past the 1,024-per-block limit. Push T bigger and the working set spills: data evicts back to L2/HBM mid-computation, occupancy craters because each block hoards too much SRAM, and your reuse advantage evaporates back toward the untiled traffic. So tiling has a sweet spot: large enough to amortize the load, small enough to fit. Real kernels tile at multiple levels at once (a register tile inside an SRAM tile inside an L2 tile), each sized to its level — which multiplies the schedule space yet again and is precisely why nobody hand-picks it per shape.

arithmetic intensity vs tile size T  (matmul, fp32)

 FLOP/byte
   ^                                  spill: tiles exceed SRAM,
   |                         ____        traffic falls back to HBM
   |                    ____/    \____
   |               ____/              \________   ← reuse collapses
   |          ____/   ridge point .........
   |     ____/   (memory-bound below, compute-bound above)
   |  __/
   | /  naive, T=1: ~1 FLOP/byte, deeply memory-bound
   +----+----+----+----+----+----+----+----+---->  T
       4    8   16   32   48   64   96  128       tile size
                     ^ sweet spot: fits SRAM, max reuse

5 · The polyhedral model: which transforms are legal

Reordering and tiling loops is only safe if it does not change the answer, and the answer is fixed by the program's dependences: if statement instance S(i,k) writes a value that S(i,k+1) reads, the second must still run after the first in any reordering. The polyhedral model is the framework that makes "preserve the dependences" checkable. It represents a loop nest's iteration space as an integer polyhedron — the set of integer points (i, j, k) satisfying the loops' affine bounds, e.g. 0 ≤ i < M, 0 ≤ k < K — and represents each loop transform (tile, skew, reorder, fuse) as an affine map on those points. A dependence is a pair of iterations that touch the same memory with at least one write; a transform is legal iff its affine map preserves the relative order of every dependent pair (the transformed schedule never moves a read before the write it depends on). Because all the constraints are affine over integers, legality and even optimality (minimize a cost like reuse distance) become integer-linear problems a solver can decide — no guessing, no per-case correctness review. This is the machinery behind Tiramisu, LLVM's Polly, parts of XLA's fusion/tiling, and MLIR's affine dialect; it is the formal floor under the "swap any schedule line, same answer" promise of §2, and it is the same preserve-semantics invariant introduced in lessons 03–04, now stated geometrically.

6 · Drive it: the schedule explorer

The widget below is one matmul with (M, N, K) you set. Slide the tile size and toggle vectorize and unroll; the FLOPs never change, but the schedule does. Watch the KPIs: data reuse and arithmetic intensity climb with tiling, on-chip vs off-chip traffic shift, and a toy GFLOP/s estimate tracks how close the schedule gets to peak. Start at tile = 1 (no tiling) and see the throughput floor — memory-bound, a fraction of peak. Climb to the sweet spot and watch it cross toward compute-bound. Push the tile past what SRAM holds and watch it spill: reuse collapses and throughput falls back down. Vectorize and unroll add multipliers on top, until the tile is so big they cannot rescue it.

Schedule explorer — same FLOPs, different loop nest
The matmul's FLOPs are fixed; only the schedule changes. The knob that proves it: with tile = 1 the kernel is deeply memory-bound (reuse ≈ 1, throughput on the floor); raise the tile to the sweet spot and arithmetic intensity climbs past the ridge into compute-bound territory; push the tile until two tiles no longer fit SRAM and it spills — reuse and throughput collapse back down.
data reuse / loaded byte
arithmetic intensity
est. throughput
SRAM for 2 tiles
off-chip HBM
on-chip reuse
% of peak

The shape of the throughput curve is the whole lesson: flat and low at tile = 1 (no reuse, memory-bound), a broad peak where the tile fits SRAM and the chip runs compute-bound, then a cliff where the tile spills and you are back near the floor — all at identical FLOPs. Vectorize and unroll lift the peak but cannot save a spilling tile. That curve, swept across every shape and every hardware, is what lesson 10's autotuner is searching.

Failure modes & checklist

Failure modes

  • Tiles too big for the fast level. Picking T so two tiles exceed SRAM / the register file. Signal: register-spill or low-occupancy warnings, and a "tiled" kernel no faster than the naive nest — the working set is bouncing off HBM.
  • No tiling at all. Shipping the textbook triple loop. Signal: arithmetic intensity ≈ 1, the op pinned at a small fraction of peak, profiler says memory-bound on a clearly compute-heavy GEMM.
  • Vectorizing a strided axis. Mapping SIMD lanes onto a non-contiguous loop. Signal: the vector load lowers to a gather; bandwidth craters instead of widening.
  • Over-unrolling. Unrolling so far you overflow the instruction cache or run out of registers. Signal: code size explodes, occupancy drops, and throughput falls below the modestly-unrolled version.
  • An illegal reorder. Swapping loops across a real dependence (a reduction's accumulation order, an in-place update). Signal: wrong numbers, often only at certain sizes — the schedule violated a dependence the polyhedral check would have caught.
  • Hand-picking one schedule for all shapes. Tuning the tile for one (M,N,K) and reusing it everywhere. Signal: great on the benchmark shape, mediocre on production shapes — the sweet spot moved.

Checklist

  • Separate algorithm from schedule. State the math once; treat every loop decision as a tunable, swappable schedule.
  • Tile first. It is the highest-leverage knob; size each tile to the level it targets (registers, then SRAM, then L2).
  • Check the working set fits. 2·T²·dtype ≤ SRAM per block; if it spills, shrink the tile or tile at more levels.
  • Vectorize the contiguous axis only. Reorder so the innermost loop is stride-1 before mapping it to lanes.
  • Watch arithmetic intensity cross the ridge. Tile until FLOP/byte clears the roofline ridge point; that is the memory-bound→compute-bound transition.
  • Confirm legality. Any reorder/tile must preserve dependences; lean on the polyhedral check rather than eyeballing it.

Checkpoint

Try it
Open the explorer at M=N=K=1,024 and set tile = 1 with vectorize and unroll off: note the throughput floor and that reuse ≈ 1 (deeply memory-bound). Now raise the tile to 32: reuse jumps to ≈ 32×, arithmetic intensity climbs past the ridge, and throughput leaps — same 2.15 GFLOP, far faster. Compute the SRAM by hand at T = 32 fp32: 2·32²·4 = 8 KB, comfortably on-chip. Now push the tile to 128: 2·128²·4 = 128 KB exceeds the ~96 KB SRAM budget, the KPI turns red, and throughput falls back down — that is the spill. Finally, with a good tile, toggle vectorize and unroll on and off and watch them add a multiplier on top of the tiling win, but never rescue the spilling tile. Predict before you slide: at T = 64, does it still fit? (2·64²·4 = 32 KB — yes, comfortably; the spill only appears between T = 96 at 72 KB and T = 128 at 128 KB.)

Where this points next

You can now describe a good loop nest: tile to the memory hierarchy, reorder for stride-1, vectorize the lanes, unroll for ILP, parallelize for occupancy, compute_at to fuse — all guaranteed legal by the polyhedral dependence check. But a schedule is still only a plan. It says "tile by 32, vectorize the inner loop, map blocks to SMs" in the compiler's own vocabulary; nothing in it is an instruction the GPU can run. Something must lower that plan to a real target — emit LLVM IR that becomes PTX and SASS, or emit a tile DSL like Triton and let its compiler map tiles to warps. That translation, with all its target-specific choices, is code generation. And the moment codegen can emit any schedule, a new problem surfaces: which tile size, which unroll factor, which num_warps — out of the millions this lesson opened up? The space is too large and too shape-dependent to hand-pick, which forces an automated search (lesson 10). Next: lesson 09, Code generation.

Takeaway
A schedule is the loop nest that realizes an op — loop order, tile sizes, vectorization, unrolling, parallel mapping, and where intermediates live — and it is a no-op on the math while swinging wall-clock by 10–100×, because it decides how many times each byte crosses the memory hierarchy. The organizing idea is Halide's algorithm/schedule split: write the math once, write the loop strategy as a separate, swappable program, so the schedule is searchable against a fixed definition of correctness. The knobs each buy a specific thing — tile buys locality (reuse ≈ T per loaded byte, raising arithmetic intensity across the roofline ridge), reorder buys stride-1 access, vectorize buys SIMD/SIMT lane width, unroll buys instruction-level parallelism, parallelize buys occupancy, and compute_at buys fusion by placing a producer inside the consumer's loops. Tiling is the headline: size the tile so two tiles fit SRAM (2·T²·dtype) and the op turns compute-bound — too big and it spills back to HBM-bound. The polyhedral model makes the legality of every reorder/tile decidable: represent the iteration space as an integer polyhedron and each transform as an affine map; a transform is legal iff it preserves every data dependence — the same preserve-semantics invariant from lessons 03–04, now geometric. The schedule is still only a plan; turning it into machine code is codegen, and searching the enormous schedule space is autotuning.

Interview prompts