For example, the upper-left output block is A[0, 0] × B[0, 0] + A[0, 1] × B[1, 0], where the indices now select whole blocks. Each × is a matrix multiplication: every row of the A block is multiplied by every column of the B block.

We can follow that one output block. We keep 16 running totals, initially zero. The first pair of input blocks supplies the products for k = 0, 1, 2, 3. We update all 16 totals, then replace the inputs with the second pair, for k = 4, 5, 6, 7. The totals stay where they are. Only after both pairs have contributed do we store the finished output block.

To compute one output independently, we read eight values from A and eight from B. Repeating that for all 16 outputs takes 256 input loads. Working on the block, we load 16 values from A and 16 from B, keep them while updating all 16 outputs, then do the same with the next pair: 64 input loads. Each A value contributes to four output columns, and each B value contributes to four output rows.

Reusing the inputs this way leaves the arithmetic unchanged. Drawing block boundaries alone doesn’t achieve this: we have to keep the values nearby while updating the outputs. These counts describe the algorithm’s loads, not transfers from main memory; caches may satisfy some of the loads in either version.

The figure’s Keep them in cache view follows a few products from the next row of C. With whole rows, the walk through B has evicted the lines we need again. With blocks, those lines are still there. (The cache is deliberately tiny, so we can see everything it contains.) The loop swap improved reuse within a row; blocking lets us reuse the same inputs across several rows as well.

The blocked multiplication is the same formula as before, applied to a 2 × 2 matrix whose elements are 4 × 4 blocks. The + and × it uses are matrix addition and multiplication. That’s the trick: tiling is matrix multiplication at a different element type. (We can do this more than once, which is convenient when the machine has more than one level of cache.)

Scheduling for the CPU

We’ll want both 32-bit and 64-bit CPU versions, but the loop swap didn’t use any property of f32. Instead of writing the contraction again for f64, we can make its element type and operations a parameter:

cachetcontraction matmul {S : Semiring}
    (reads a: [S.Value; M, K])
    (reads b: [S.Value; K, N])
    -> out c: [S.Value; M, N] =
  c[i, j] = sum k. a[i, k] * b[k, j]

S supplies the element type S.Value, the operations, and their laws. This interface is called a semiring. PlusTimes f32 supplies ordinary floating-point arithmetic, with fma for its mul_add, so matmul @(S = PlusTimes f32) is the computation we’ve used so far. We’ll use this generic definition for the rest of the post.

tower, the CPU schedule whose first two steps split each sum into chunks, continues from there. The next two steps cut the rows and columns into blocks too, and order then moves the loop over chunks outside the output loops. We now take one block of the depth k and visit all the outputs it contributes to, giving us a chance to reuse blocks of A and B while they’re still in cache. This schedule still runs on one CPU core. The figure follows it into packed, which adds packing and register storage:

schedule tower of matmul =
  zero (i, j) c
  split k by 16 as gk ragged
  split i by 16 as gi, 4 as ui ragged
  split j by 8 as gj ragged
  order (gk, gj, gi, ui, i, j)

schedule packed of matmul @(S = PlusTimes f32)
    (mut packA: [f32; 16, 16] in PanelMajor 4 at DRAM)
    (mut packB: [f32; 16, 8] at DRAM) =
  ...
  order (gk, gj, gi, ui, i, j)
  pack a into packA along gi over (a0, a1)
  keep packA along gi within L2 into packA_L2
  ...

for i in 0 .. M do
  for j in 0 .. N do
    var acc'0: S.Value = S.zero ()
    for k in 0 .. K do
      acc'0 = S.mul_add ... acc'0
    store c[i, j] = acc'0

for zi in 0 .. M do for zj in 0 .. N do
  store c[zi, zj] = S.zero ()
for i in 0 .. M do
  for j in 0 .. N do
    for k in 0 .. K do
      let acc'0 = load c[i, j]
      store c[i, j] = S.mul_add ... acc'0

... // zero C
for i in 0 .. M do
  for j in 0 .. N do
    for gk in blocks (0 .. K) 16 do
      for k in 0 .. gk.len do
        let acc'0 = load c[i, j]
        store c[i, j] = S.mul_add ... acc'0

... // zero C
for gi in blocks (0 .. M) 16 do
  for ui in blocks (0 .. gi.len) 4 do
    for i in 0 .. ui.len do
      for j in 0 .. N do
        for gk in blocks (0 .. K) 16 do
          for k in 0 .. gk.len do
            let acc'0 = load c[gi[ui[i]], j]
            store c[gi[ui[i]], j] = S.mul_add ... acc'0

... // zero C
for gi in blocks (0 .. M) 16 do
  for ui in blocks (0 .. gi.len) 4 do
    for i in 0 .. ui.len do
      for gj in blocks (0 .. N) 8 do
        for j in 0 .. gj.len do
          for gk in blocks (0 .. K) 16 do
            for k in 0 .. gk.len do
              let acc'0 = load c[gi[ui[i]], gj[j]]
              store c[gi[ui[i]], gj[j]] = S.mul_add ... acc'0

... // zero C
for gk in blocks (0 .. K) 16 do
  for gj in blocks (0 .. N) 8 do
    for gi in blocks (0 .. M) 16 do
      for ui in blocks (0 .. gi.len) 4 do
        for i in 0 .. ui.len do
          for j in 0 .. gj.len do
            for k in 0 .. gk.len do
              let acc'0 = load c[gi[ui[i]], gj[j]]
              store c[gi[ui[i]], gj[j]] = S.mul_add ... acc'0

... // zero C
for gk in blocks (0 .. K) 16 do
  for gj in blocks (0 .. N) 8 do
    for gi in blocks (0 .. M) 16 do
      for a0 in 0 .. gi.len do
        for a1 in 0 .. gk.len do
          store packA[a0, a1] = load a[gi[a0], gk[a1]]
      ... // update C from packA and B

... // zero C
for gk in blocks (0 .. K) 16 do
  for gj in blocks (0 .. N) 8 do
    for gi in blocks (0 .. M) 16 do
      for a0 in 0 .. gi.len do
        for a1 in 0 .. gk.len do
          store packA[a0, a1] = load a[gi[a0], gk[a1]]
      lease mut packA_L2 = packA at L2
      ... // update C from packA_L2 and B

... // zero C
for gk in blocks (0 .. K) 16 do
  for gj in blocks (0 .. N) 8 do
    for b0 in 0 .. gk.len do
      for b1 in 0 .. gj.len do
        store packB[b0, b1] = load b[gk[b0], gj[b1]]
    for gi in blocks (0 .. M) 16 do
      for a0 in 0 .. gi.len do
        for a1 in 0 .. gk.len do
          store packA[a0, a1] = load a[gi[a0], gk[a1]]
      lease mut packA_L2 = packA at L2
      ... // update C from packA_L2 and packB

... // zero C
for gk in blocks (0 .. K) 16 do
  for gj in blocks (0 .. N) 8 do
    for b0 in 0 .. gk.len do
      for b1 in 0 .. gj.len do
        store packB[b0, b1] = load b[gk[b0], gj[b1]]
    lease mut packB_L1 = packB at L1
    for gi in blocks (0 .. M) 16 do
      for a0 in 0 .. gi.len do
        for a1 in 0 .. gk.len do
          store packA[a0, a1] = load a[gi[a0], gk[a1]]
      lease mut packA_L2 = packA at L2
      ... // update C from packA_L2 and packB_L1

... // zero C
for gk in blocks (0 .. K) 16 do
  for gj in blocks (0 .. N) 8 do
    ... // pack B and keep it at L1
    for gi in blocks (0 .. M) 16 do
      ... // pack A and keep it at L2
      for ui in blocks (0 .. gi.len) 4 do
        var acc'0: [Vec f32 8; 4] = ...
        ... // load the current totals from C into acc'0
        for k in 0 .. gk.len do
          let lanes: Vec f32 8 = load packB_L1[...] ...
          for i in 0 .. ui.len do
            let row: Vec f32 8 = splat (load packA_L2[...])
            acc'0[i] = fma row lanes acc'0[i]
        ... // store acc'0 back to C

Each cell’s running total still carries on from one depth block to the next, so no addition is regrouped. Each step has to prove that it preserves the result, for every size of matrix. It doesn’t need to know how S implements its operations: before and after the rewrite, each output applies the same operations to the same values in the same order. So tower works for any S.

tower is actually slower than swapped. Splitting the loops into blocks hasn’t given us a fast inner kernel yet.

The first pack step in the figure starts a new schedule, packed, which continues where tower left off. Two pack steps copy blocks of A and B into temporary buffers, arranging the values for the kernel to read. These buffers become parameters of packed, and two keep steps say which cache each buffer should stay in.

The last step, registers (i, j) into c rows 4 lanes 8 through fma, holds a 4 × 8 block of C in vector registers across the depth loop. A vector instruction operates on several values at once; here, its multiply-add updates four f32 values. This ties the schedule to the machine’s vector registers, so we bind S to PlusTimes f32.

The figure writes out every step, but tower and packed share their first five. We can give those steps a name with schedule fn, and make the block sizes parameters:

cachetschedule fn block_product (i, j, k) c (MC, MR, NC, KC) =
  zero (i, j) c
  split k by KC as gk ragged
  split i by MC as gi, MR as ui ragged
  split j by NC as gj ragged
  order (gk, gj, gi, ui, i, j)

schedule tower of matmul =
  block_product (i, j, k) c (16, 4, 8, 16)

schedule packed of matmul @(S = PlusTimes f32)
    (mut packA: [f32; 16, 16] in PanelMajor 4 at DRAM)
    (mut packB: [f32; 16, 8] at DRAM) =
  block_product (i, j, k) c (16, 4, 8, 16)
  pack a into packA along gi over (a0, a1)
  keep packA along gi within L2 into packA_L2
  pack b into packB along gj over (b0, b1)
  keep packB along gj within L1 into packB_L1
  registers (i, j) into c rows 4 lanes 8 through fma

(i, j, k) c names the indices and output to rewrite, as in hoist. MC and MR choose the outer and inner row block sizes, NC the column block size, and KC the depth. Both schedules pass the same sizes; packed then adds the steps below the call.

The call to block_product expands into its steps when the compiler reads the schedule. The following steps can keep working on the loops it produces: pack copies their inputs, keep records their cache requirements, and registers holds their running totals in registers. There’s no function call left in the kernel.

Schedule functions can call other schedule functions too, so we can build up an optimisation from smaller ones. Each rewrite is checked against the program left by the preceding steps, including those from a helper. The checks compose along with the transformations: if every step preserves the result under its stated conditions, the whole sequence does too.

PanelMajor 4 lays out A in groups of four rows, putting the four values from each column in consecutive memory. The inner kernel reads four values at each depth step, and this layout keeps them together. B’s buffer keeps the ordinary row-major layout, so its eight values from that step are contiguous too. We still write packA[i, k]; only the address it refers to changes.

The scratch buffers live in main memory (DRAM); keep names the cache views the kernel will read. The cache requirements appear in the generated function’s declaration:

cachetscratch packA, packB
working_set
  @L1 { keep packB = descriptor packB }
  @L2 { keep packA = descriptor packA }

scratch requires the kernel to write to each buffer location before reading it. The caller therefore doesn’t need to clear the buffers first. The working_set rows say which buffer stays in which cache, and a caller of packed has to show that they fit. The memory hierarchy panel shows where the buffers live. The block sizes are deliberately small here.

Making it fast

packed multiplies two 1024 × 1024 matrices at 47 GFLOP/s. OpenBLAS reaches 108 on the same core with the same vector instructions.

How much data are we moving?

Copying the inputs is extra work. We expect it to pay off because the kernel can then read them repeatedly from cache, without going back to main memory. But how much copying have we actually asked for?

The compiler keeps an account of what each function does, and cachetc account prints it, one operation at a time. Here is packed’s, shortened:

packed
  operations
    loop  host  x 1
      writes c 4 * M * N
    loop  host  x ceil (K / 16) * ceil (N / 8)
      reads b 512; writes packB 512
    loop  host  x ceil (K / 16) * ceil (M / 16) * ceil (N / 8)
      reads a 1024; writes packA 1024
    loop  host  x ceil (K / 16) * ceil (M / 16) * ceil (N / 8)
      reads c 512, packA_L2 1024, packB_L1 2048; writes c 512; work fma 2048, ...

The number after each x says how many times that operation runs. The first row zeroes c, all 4 * M * N bytes of it. The next two copy a 16 × 8 block of B and a 16 × 16 block of A into their buffers. reads and writes count bytes; work fma counts multiply-adds.

The last row is the computation on one block: it reads packed A from L2 and packed B from L1, then performs 2,048 multiply-adds. The repetition counts use M, N and K, so we can see how the work grows with the matrices.

The generated function also declares bounds on this traffic, next to the working_set we saw above. These rows count the cache lines moved by each pack:

cachetmovement
  traffic packB_pull = ceil (K / 16) * ceil (N / 8) *
    lines (span_bytes packB.span) from DRAM to L1 and back
  traffic packA_pull = ceil (K / 16) * ceil (N / 8) * ceil (M / 16) *
    lines (span_bytes packA.span) from DRAM to L2 and back

A caller can put a limit on this traffic too. For two 32 × 32 matrices, the loop order above packs A once for each combination of a depth block, a column block and a row block: 2 × 4 × 2 = 16 packs. If the caller allows only eight, the call is refused. We have to change the schedule or allow the extra traffic. The full example and refusal show how the caller writes this budget.

The account is exact within the platform file’s model of the caches, but it knows nothing about prefetchers, replacement policies, or another process sharing the cache. Whether reducing these counts makes the kernel faster still has to be measured.

Tuning the kernel

Our blocks are small enough to read, but far smaller than the caches they’re kept in. The obvious next step is to make them larger:

schedule packed of matmul @(S = PlusTimes f32)
    (mut packA: [f32; 16, 16] in PanelMajor 4 at DRAM)
    (mut packB: [f32; 16, 8] at DRAM) =
  ...
  keep packB along gj within L1 into packB_L1
  registers (i, j) into c rows 4 lanes 8 through fma

schedule larger of matmul @(S = PlusTimes f32)
    (mut packA: [f32; 512, 256] in PanelMajor 4 at DRAM)
    (mut packB: [f32; 256, 8] at DRAM) =
  ...
  keep packB along gj within L1 into packB_L1
  registers (i, j) into c rows 4 lanes 8 through fma

schedule rows_first of matmul @(S = PlusTimes f32)
    (mut packA: [f32; 512, 256] in PanelMajor 4 at DRAM)
    (mut packB: [f32; 256, 8] at DRAM) =
  ...
  keep packB along gj within L1 into packB_L1
  registers (i, j) into c rows 4 lanes 8 through fma

schedule fast of matmul @(S = PlusTimes f32)
    (mut packA: [f32; 512, 256] in PanelMajor 4 at DRAM)
    (mut packB: [f32; 256, 16] at DRAM) =
  ...
  keep packB along gj within L1 into packB_L1
  registers (i, j) into c rows 4 lanes 16 through fma

for gk in blocks (0 .. K) 16 do
  for gj in blocks (0 .. N) 8 do
    ... // pack B and keep it at L1
    for gi in blocks (0 .. M) 16 do
      ... // pack A and keep it at L2
      for ui in blocks (0 .. gi.len) 4 do
        var acc'0: [Vec f32 8; 4] = ...
        ... // load C, accumulate across k, and store C

for gk in blocks (0 .. K) 256 do
  for gj in blocks (0 .. N) 8 do
    ... // pack B and keep it at L1
    for gi in blocks (0 .. M) 512 do
      ... // pack A and keep it at L2
      for ui in blocks (0 .. gi.len) 4 do
        var acc'0: [Vec f32 8; 4] = ...
        ... // load C, accumulate across k, and store C

for gk in blocks (0 .. K) 256 do
  for gi in blocks (0 .. M) 512 do
    ... // pack A and keep it at L2
    for gj in blocks (0 .. N) 8 do
      ... // pack B and keep it at L1
      for ui in blocks (0 .. gi.len) 4 do
        var acc'0: [Vec f32 8; 4] = ...
        ... // load C, accumulate across k, and store C

for gk in blocks (0 .. K) 256 do
  for gi in blocks (0 .. M) 512 do
    ... // pack A and keep it at L2
    for gj in blocks (0 .. N) 16 do
      ... // pack B and keep it at L1
      for ui in blocks (0 .. gi.len) 4 do
        var acc'0: [Vec f32 16; 4] = ...
        ... // load C, accumulate across k, and store C

The first step, Larger blocks, mostly changes sizes. The blocks of A are now 512 rows by 256 columns, half a megabyte, kept in L2, and the blocks of B are 256 by 8, kept in L1. These tall rectangular blocks are often called panels. Surprisingly, it’s barely faster: 48 GFLOP/s. (The correctness proof is perfectly happy with this.)

The traffic account shows a lot of repeated copying. The factor ceil (N / 8) in A’s traffic row comes from gj, which is still outside gi: we pack each panel of A again for every panel of B. At 1024 × 1024, that’s 1,024 copies of half a megabyte for one multiplication. A traffic contract written for these sizes can refuse this order.

Rows first puts gi before gj in the order, so the rows are outside the columns, and the factor ceil (N / 8) moves to B’s row, whose panels are 64 times smaller. That gets us to 66 GFLOP/s.

That’s still only about 61% of OpenBLAS. Another problem is that we’re waiting for our own results. A vector fma on the M4 Max takes four cycles before its result can be used by the next fma in the same chain. The core can start four independent vector fmas per cycle, so it needs 16 independent chains to keep doing that while earlier results are on their way.

Each vector holds four f32 values, so our 4 × 8 tile occupies only eight vectors. After updating those eight, we have to wait before we can update them again. Widening the tile gives the core more independent totals to work on in the meantime:

The platform file gives the number of values per vector as lanes, the number of fmas it can start per cycle as fma_pipes, and the delay before using a result as fma_latency. We can use them to write a requirement next to the schedule, whose last step becomes registers (i, j) into c rows MR lanes NR through fma:

cachetconst MR : Nat = 4
const NR : Nat = 8

proof enough_chains =
  assert MR * ceil (NR / platform.lanes @Vec @f32) >=
    platform.fma_pipes @Vec * platform.fma_latency @Vec

With a 4 × 8 tile, it’s refused:

could not prove proof enough_chains: MR * ceil (NR / platform.lanes Vec f32) >= platform.fma_pipes @Vec * platform.fma_latency @Vec

A reader who has been counting independent chains will already know what to try: NR = 16. Each of the four rows now occupies four vectors, giving us sixteen independent chains, and the proof goes through. fast now runs at 116 GFLOP/s, as fast as OpenBLAS:

fast still produces exactly the same result as matmul. If someone later narrows the tile to save registers, the enough_chains check fails, pointing out the dependency we’d reintroduce. This check relies on the platform file’s numbers: Apple doesn’t publish the latency or throughput, so the file uses published measurements.

We could also have asked Cachet to find that width. With MR still fixed at four, we can replace const NR : Nat = 16 with:

cachetsynth const NR : Nat =
  minimize n in (n): n where
    n > 0, n <= 64,
    n % platform.lanes @Vec @f32 == 0,
    MR * ceil (n / platform.lanes @Vec @f32) >=
      platform.fma_pipes @Vec * platform.fma_latency @Vec

This asks for the smallest width from 1 to 64 that fills whole vectors and supplies enough independent chains. For the platform above, that is 16. So the condition that refused our 4 × 8 tile can also pick the width for us. synth const chooses the number at compile time, and the rest of the schedule is checked with that number in place, just as if we’d written it ourselves.

Several parameters can be chosen together too, and the candidates ranked by a cost expression, such as the estimated load cost per multiply-add. We’d still benchmark the promising candidates: a low cost in the model needn’t mean a short running time.

Using the other cores

We now have a fast kernel for one core. To use the other cores, we can divide the output rows between them. In bands, parallel gi turns the loop over blocks of 64 rows into a grid of independent iterations. It runs on the host, so CPU threads take the bands of rows.

Each iteration has a write window: the part of the output it’s allowed to update. The compiler infers this from the rows it stores, and checks that the windows don’t overlap. Threads can read the same inputs, but can’t write each other’s outputs.

parallel is a library step too, and it’s short enough to show in full:

cachetstep fn parallel t =
  from
    for t in $lo .. $hi do
      $body
  to
    grid t in $lo .. $hi do
      $body

There’s no relate proof recipe this time. The grid it produces has to prove its windows don’t overlap anyway.

The threads share the packed inputs: A is packed once per block of depth, for every row, and B is packed in groups of 16 columns and kept in L2:

cachetschedule bands of matmul @(S = PlusTimes f32)
    (mut packA: [f32; M, 256] in PanelMajor 4)
    (mut packB: [f32; 256, 1024] in ColumnPanels 16 at DRAM) =
  zero (i, j) c

  // Block and reorder the loops.
  split k by 256 as gk ragged
  split j by 1024 as gj, 16 as uj ragged
  split i by 64 as gi, 4 as ui
  order (gk, gj, gi, ui, uj, i, j)

  // Give each CPU thread a band of output rows.
  parallel gi

  // Pack the inputs for reuse across the output bands.
  pack a into packA along gk over (a0, a1)
  strip a0 4 (pa, ra)
  interchange (ra, a1)
  parallel pa

  pack_lanes b into packB along gj over (b0, b1, bt) lanes 16
  keep packB along gj within L2 into packB_L2
  parallel b0

  // Accumulate a 4 × 16 output block in vector registers.
  registers (i, j) into c rows 4 lanes 16 through fma

The packing loops are parallel too. strip a0 4 (pa, ra) divides A’s rows into groups of four, with pa selecting a group and ra a row within it. interchange (ra, a1) moves the column loop outside those four rows, so the copy writes neighbouring values in PanelMajor 4 order. parallel pa lets separate threads pack separate groups.

For B, ColumnPanels 16 stores groups of 16 columns, with each group’s rows one after another. pack_lanes copies 16 values at a time into that layout, and parallel b0 distributes the rows. On NEON, 16 f32 values fit in four vectors.

On all cores, bands reaches 757 GFLOP/s, where OpenBLAS reaches 863. No sum is split between threads, only rows, so the result is still identical to matmul’s.

CPU matrix units

fast and bands use NEON, the CPU’s vector instructions. The M4 also has a matrix unit, accessed through SME, Arm’s Scalable Matrix Extension. We can use it for f64 by choosing another multiply-add, outer_fma_primitive, which updates 8 × 8 blocks with outer products. An outer product multiplies each of eight values from A by each of eight values from B, contributing one product to each of the block’s 64 totals. sme holds a 32 × 16 tile of C as eight such blocks, and lays out the rest of the work for the M4 Max’s two matrix units, one in each cluster of performance cores:

schedule sme of matmul @(S = PlusTimes f64)
    (mut packA: [f64; M, 3072] in PanelMajor 32)
    (mut packB: [f64; 3072, 256] in ColumnPanels 16 at DRAM) =

  split k by 3072 as gk ragged
  split j by 256 as gj, 16 as uj
  ...

for i in 0 .. M do
  for j in 0 .. N do
    var acc'0: f64 = -0.0
    for k in 0 .. K do
      acc'0 = fma ... acc'0
    store c[i, j] = acc'0

for i in 0 .. M do
  for j in 0 .. N do
    for gk in blocks (0 .. K) 3072 do
      var acc'0: f64 = if gk.lo == 0 then -0.0 else load c[i, j]
      for k in 0 .. gk.len do
        acc'0 = fma ... acc'0
      store c[i, j] = acc'0

for i in 0 .. M do
  for gj in 0 .. N / 256 do
    for uj in blocks (0 .. 256) 16 do
      for j in 0 .. uj.len do
        for gk in blocks (0 .. K) 3072 do
          var acc'0: f64 = if gk.lo == 0 then -0.0 else load c[...]
          for k in 0 .. gk.len do
            acc'0 = fma ... acc'0
          store c[...] = acc'0

for gi in 0 .. M / 1536 do
  for ui in blocks (0 .. 1536) 32 do
    for i in 0 .. ui.len do
      for gj in 0 .. N / 256 do
        for uj in blocks (0 .. 256) 16 do
          for j in 0 .. uj.len do
            for gk in blocks (0 .. K) 3072 do
              var acc'0: f64 = if gk.lo == 0 then -0.0 else load c[...]
              for k in 0 .. gk.len do
                acc'0 = fma ... acc'0
              store c[...] = acc'0

for gk in blocks (0 .. K) 3072 do
  for gj in 0 .. N / 256 do
    for gi in 0 .. M / 1536 do
      for ui in blocks (0 .. 1536) 32 do
        for uj in blocks (0 .. 256) 16 do
          for i in 0 .. ui.len do
            for j in 0 .. uj.len do
              var acc'0: f64 = if gk.lo == 0 then -0.0 else load c[...]
              for k in 0 .. gk.len do
                acc'0 = fma ... acc'0
              store c[...] = acc'0

for gk in blocks (0 .. K) 3072 do
  for gj in 0 .. N / 256 do
    grid gi in 0 .. M / 1536
        writes c[...] do
      for ui in blocks (0 .. 1536) 32 do
        for uj in blocks (0 .. 256) 16 do
          for i in 0 .. ui.len do
            for j in 0 .. uj.len do
              var acc'0: f64 = if gk.lo == 0 then -0.0 else load c[...]
              for k in 0 .. gk.len do
                acc'0 = fma ... acc'0
              store c[...] = acc'0

for gk in blocks (0 .. K) 3072 do
  for a0 in 0 .. M do
    for a1 in 0 .. gk.len do
      store packA[a0, a1] = load a[a0, gk[a1]]
  for gj in 0 .. N / 256 do
    grid gi in 0 .. M / 1536
        writes c[...] do
      ... // update C from packA and B

for gk in blocks (0 .. K) 3072 do
  for pa in 0 .. M / 32 do
    for ra in 0 .. 32 do
      for a1 in 0 .. gk.len do
        store packA[0 + pa * 32 + ra, a1] = load a[...]
  for gj in 0 .. N / 256 do
    grid gi in 0 .. M / 1536
        writes c[...] do
      ... // update C from packA and B

for gk in blocks (0 .. K) 3072 do
  grid pa in 0 .. M / 32 writes packA[...] do
    for ra in 0 .. 32 do
      for a1 in 0 .. gk.len do
        store packA[0 + pa * 32 + ra, a1] = load a[...]
  for gj in 0 .. N / 256 do
    grid gi in 0 .. M / 1536
        writes c[...] do
      ... // update C from packA and B

for gk in blocks (0 .. K) 3072 do
  grid pa in 0 .. M / 32 writes packA[...] do
    ... // pack A
  for gj in 0 .. N / 256 do
    for b0 in 0 .. gk.len do
      for b1 in 0 .. 256 do
        store packB[b0, b1] = load b[gk[b0], gj * 256 + b1]
    grid gi in 0 .. M / 1536
        writes c[...] do
      ... // update C from packA and packB

for gk in blocks (0 .. K) 3072 do
  grid pa in 0 .. M / 32 writes packA[...] do
    ... // pack A
  for gj in 0 .. N / 256 do
    for b0 in 0 .. gk.len do
      for pb in 0 .. 256 / 16 do
        for rb in 0 .. 16 do
          store packB[b0, 0 + pb * 16 + rb] = load b[...]
    grid gi in 0 .. M / 1536
        writes c[...] do
      ... // update C from packA and packB

for gk in blocks (0 .. K) 3072 do
  grid pa in 0 .. M / 32 writes packA[...] do
    ... // pack A
  for gj in 0 .. N / 256 do
    for pb in 0 .. 16 do
      for b0 in 0 .. gk.len do
        for rb in 0 .. 16 do
          store packB[b0, 0 + pb * 16 + rb] = load b[...]
    grid gi in 0 .. M / 1536
        writes c[...] do
      ... // update C from packA and packB

for gk in blocks (0 .. K) 3072 do
  grid pa in 0 .. M / 32 writes packA[...] do
    ... // pack A
  for gj in 0 .. N / 256 do
    grid pb in 0 .. 16 writes packB[...] do
      for b0 in 0 .. gk.len do
        for rb in 0 .. 16 do
          store packB[b0, 0 + pb * 16 + rb] = load b[...]
    grid gi in 0 .. M / 1536
        writes c[...] do
      ... // update C from packA and packB

... // pack inputs, and distribute the output bands
for ui in blocks (0 .. 1536) 32 do
  for uj in blocks (0 .. 256) 16 do
    var acc'0: [Vec f64 16; 32] = ...
    ... // the 32 × 16 output tile: −0.0 on the first block, else c's
    for k in 0 .. gk.len do
      let lanes: Vec f64 16 = ...
      for i in 0 .. ui.len do
        let row: Vec f64 16 = ...
        acc'0[i] = fma row lanes acc'0[i]
    ... // store the output tile

... // pack inputs, and distribute the output bands
for ui in blocks (0 .. 1536) 32 do
  for uj in blocks (0 .. 256) 16 do
    var blocks_2: [Vec f64 (8 * 8); 4 * 2] = ...
    ... // the same tile as eight 8 × 8 blocks: −0.0, else c's
    for k in 0 .. gk.len do
      ... // row and column masks
      let column_all: Vec f64 32 = load packA[...]
      let row_all: Vec f64 16 = load packB[...]
      for q in 0 .. ui.len / 8 * 2 do
        let column = slice_lanes column_all (floor (q / 2) * 8)
        let row = slice_lanes row_all (q % 2 * 8)
        blocks_2[q] =
          outer_fma_primitive column row blocks_2[q] rows cols
    ... // store the eight blocks

Like bands, sme divides the rows into bands, but uses larger blocks. Each band of 1536 rows runs on a thread of its own, so at 3072 × 3072 each matrix unit works on one band. Packing the inputs is just copying, so that work can still run across all the cores.

Within a band, the order in step 4 visits the tiles along C’s rows. We work across the same 32 rows before moving down to the next 32, keeping the accesses to C on the same memory pages. Visiting the tiles down the columns would start every tile on 32 new rows.

At this matrix size, one depth block covers the whole sum: 3072 steps. Each tile’s running totals can stay in the matrix unit until the sum is finished, and we store the result once. With blocks of 768, we’d write each output tile four times, reading it back between blocks. Those extra transfers cost about a tenth of the running time.

There’s also no zero step here. The first depth block starts the running totals at -0.0 in the matrix unit, so we don’t need to clear C and then read it back. At this size, C is only ever written. Two library steps and a few new checker rules make this possible while keeping the result identical to matmul’s.

For each step along the depth, we load a column of A and a row of B once, using SME2’s multi-vector load instructions. All eight 8 × 8 blocks then reuse those values.

For the checker, outer_fma_primitive is a scalar fma on each active cell of the accumulator. Here the LLVM backend selects the matrix instructions. We rely on the compiler and hardware to implement the operation correctly, without supplying a template of our own.

On 3072 × 3072 matrices of f64, sme reaches 894 GFLOP/s, slightly faster than Accelerate, Apple’s BLAS, at 871 on the same matrix units. Like bands, it splits no sums between threads, so it produces the same result as matmul.

Scheduling for the GPU

The GPU schedules start again from matmul. The examples below bind S to PlusTimes f32, and their performance bars compare with MLX on 4096 × 4096 matrices.

Independent outputs

We gave CPU threads separate bands of output rows. On the GPU, the unit of work is one output element.

A GPU runs the same small program, called a kernel, many times. Each invocation is a thread, and its indices tell it which piece of work to do. Thread (i, j) computes c[i, j]: it reads row i of A and column j of B and keeps one running total. The GPU assigns threads to hardware as space becomes available; they needn’t all run at once.

We use a grid again, this time on the GPU:

schedule rows of matmul @(S = PlusTimes f32) on M4Gpu =
  grid over (i, j)

for i in 0 .. M do
  for j in 0 .. N do
    var acc'0: f32 = -0.0
    for k in 0 .. K do
      acc'0 = fma ... acc'0
    store c[i, j] = acc'0

grid on M4Gpu i in 0 .. M, j in 0 .. N
    writes c[i, j] as ct do
  ... // the same running total and k loop
  store ct = acc'0

grid over (i, j) gives each pair of indices its own grid iteration. In this schedule, each iteration runs as one GPU thread. on M4Gpu selects the device, from the description of the target machine named at the top of the file. We call the resulting function with rows a b c.

A grid is an ordinary Cachet construct, so we could have written rows by hand, and it would have been checked in the same way. In particular, the grid’s header says what each iteration may write:

cachetgrid on M4Gpu i in 0 .. M, j in 0 .. N writes c[i, j] as ct do

writes c[i, j] gives the iteration’s write window, referred to as ct in the grid’s body. Cachet checks that different iterations’ windows don’t overlap, and that each iteration writes only within its own.⁴ For example, suppose we wrote this grid by hand for 32 × 32 matrices, but stored each result one column to the right, into c[i, (j + 1) % 32]. The checker would find a thread that writes outside its window:

could not prove a grid iteration reaches only its own window: i == i && (j + 1) % 32 == j
  counterexample: j = 0

rows performs the same multiply-adds in the same order as matmul, but it runs them on the GPU. A GPU may handle some floating-point values differently from the CPU. The M4 Max’s GPU, for example, flushes subnormals (very small numbers near zero) to zero. The guarantee is therefore that rows produces exactly the same result as matmul evaluated using that GPU’s arithmetic.

The grid layout is currently a built-in step. Its side conditions are checked, but its rewrite law is part of the compiler we trust.

Tiles

rows gives each grid iteration one output number. We can give it a whole output tile instead, with several threads cooperating on it. In Cachet, that tile is an ordinary local matrix value: an accumulator, or a block loaded from an input. The grid divides the independent work; the tiles hold the values each iteration works with.

The illustration uses 32 × 32 matrices, with 16 × 16 output tiles, advancing eight steps along k at a time.

The M4’s GPU executes threads in groups of 32, called SIMD groups. A group can collectively perform an 8 × 8 matrix operation. Our 16 × 16 output tile has four 8 × 8 parts, so we’ll give each part to one SIMD group. Those four groups make up a threadgroup, which cooperates on one tile. Separate threadgroups compute the four independent output tiles:

Here’s the schedule:

schedule tiled of matmul @(S = PlusTimes f32) on M4Gpu =
  grid over (i, j)
  tiles (i, j, k) (16, 16) (8, 8)

for i in 0 .. M do
  for j in 0 .. N do
    var acc'0: f32 = -0.0
    for k in 0 .. K do
      acc'0 = fma ... acc'0
    store c[i, j] = acc'0

grid on M4Gpu i in 0 .. M, j in 0 .. N writes c[i, j] as ct do
  var acc'0: f32 = -0.0
  for k in 0 .. K do
    acc'0 = fma ... acc'0
  store ct = acc'0

grid on M4Gpu gi in 0 .. M / 16, gj in 0 .. N / 16
    by tile (si, sj) parts [8, 8]
    writes c[...] as ct_2 do
  var t: [f32; 16, 16] = repeat (-0.0)
  for k in 0 .. K / 8 do
    let x: [f32; 16, 8] = load a[...]
    let y: [f32; 8, 16] = load b[...]
    t = mma x y t
  store ct_2 = t

grid over (i, j) starts with the independent outputs, as in rows. The tiles step groups them into 16 × 16 tiles and gives each SIMD group an 8 × 8 part. These parts don’t overlap, so each output still has one owner.

For k, tiles keeps the same running total across chunks, just as tower does. The GPU matrix instruction we use is called simdgroup_mma8. Its reference body adds each block’s eight products in ascending order of k, so each output receives the same chain of multiply-adds, now eight at a time.

t holds the output tile’s running totals, starting from zero, and x and y are the next pair of input blocks. mma x y t is matrix multiply-accumulate: multiply the two input blocks and add their product to t. The totals stay in t throughout the depth loop, and are stored at the end.

In effect, tiles makes each element an 8 × 8 matrix, with simdgroup_mma8 as its multiply-add. It’s still our matmul, just at a different element type. Select a level to see one element of its type:

Inputs are now reused within each SIMD group. Different groups still load their inputs separately; to share those loads too, we’d need threadgroup memory, a small memory that all the groups in a threadgroup can access. We could also give each group more outputs to work on. That would reuse each input more times, but would need room for more running totals. We’ll try larger tiles next.

The compiler reasons about a reference implementation of the matrix multiply-add and uses a code template to replace it with a GPU instruction. We test that instruction against the reference, but don’t prove that they always produce identical results, so the template remains a trusted assumption.

tiles itself is a library step: its relation is checked in the same way as hoist’s.

Larger tiles

tiled gives each SIMD group only one 8 × 8 block of the accumulator. We can give it more outputs, so that each input is reused more times. With 64 × 32 output tiles and 32 × 16 parts, each SIMD group holds eight accumulator blocks:

cachetschedule big_tiles of matmul @(S = PlusTimes f32) on M4Gpu =
  grid over (i, j)
  tiles (i, j, k) (64, 32) (32, 16)

At 4096 × 4096, big_tiles reaches 12,407 GFLOP/s: about 2.3 times tiled’s speed, and 95% of MLX’s. The result is still identical to tiled’s.

Comparing the results

We’ve now implemented the same matmul using four different multiply-add operations: scalar fma, vector fma, the CPU’s outer_fma_primitive, and the GPU’s simdgroup_mma8. In each case, the schedule chooses a block of the output to update at once, and the instruction that updates it.

Here are all the matrix multiplication schedules together, grouped by the hardware they use. Each group compares Cachet with a library on the same matrix size and precision:

For these sizes, fast is about level with OpenBLAS on one core, and sme is slightly faster than Accelerate. bands reaches 88% of OpenBLAS’s speed across all cores, and big_tiles reaches 95% of MLX’s on the GPU. Each kernel must produce exactly the same result as matmul using its target device’s arithmetic. The libraries don’t promise this; we checked each kernel against its reference before timing it.

Attention

So far, every schedule has kept the original order of additions. Attention gives us a reason to relax that requirement.

Attention lets each position in a sequence mix information from other positions, giving more weight to the ones relevant to it. Each position has a query to compare with the others’ keys; the resulting weights say how much of each value to include. The Transformer in Attention Is All You Need writes this as:

Attention(Q, K, V) = softmax(QK^T / √d_k)V

There are two matrix multiplications, with a softmax between them. The first computes scores; softmax turns each row into weights that add up to one; the second mixes the values. We divide the scores by the square root of the query width, d_k. We’ll leave that scale factor out of the example below.

Here it is as a contraction. q holds the queries, kt the transposed keys, and v the values. Online is a library implementation of the weighted sum, which we’ll unpack in a moment:

cachetcontraction attend
    (reads q: [f32; 64, 16])
    (reads kt: [f32; 16, 32])
    (reads v: [f32; 32, 16])
    -> out y: [f32; 64, 16] =
  s[i, j] = sum c. q[i, c] * kt[c, j]
  y[i, d] = normalised (reduce[reassoc] Online j. key (s[i, j]) (v[j, d]))

The score matrix has one entry for every query–key pair. For self-attention, doubling the sequence length quadruples that matrix. Avoiding the cost of storing and rereading it is the motivation behind FlashAttention.

We’d like to compute a block of scores, use it, and discard it. Softmax makes this awkward: it exponentiates each score and divides by the sum of all the exponentials in the row, including scores we haven’t computed yet. To prevent overflow, it also subtracts the row’s largest score before exponentiating. We don’t know that maximum yet either.

Online keeps a running summary that can be rescaled when a larger score arrives. key s v creates a summary for one key, and reduce[reassoc] Online j. combines these summaries over j, much as sum combines numbers. The summary keeps the total weight and the weighted sum of values. At the end, normalised divides the weighted sum by the total weight. The side panel shows the implementation. The [reassoc] annotation is permission to regroup the reduction.

The contraction combines keys one at a time. The schedule will instead find the maximum across the old summary and a whole block of scores, rescale the summary once, then add the block’s contributions. Rescaling both the sum of weights and the weighted sum of values by the same factor leaves their ratio unchanged over the reals.

Combining a block at a time regroups the reduction, so it can change the answer in floating point. Here are six f32 values, added one at a time on the left, and in pairs on the right:

Both expressions equal one over the reals. In f32, adding one to minus 2²⁵ rounds back to minus 2²⁵, so the paired version loses the one and returns zero. Neither order is more accurate in general: with other values, the chain could just as well be the one that loses precision. But they’re different, so Cachet only regroups a reduction that the contraction allows to be regrouped. That’s what [reassoc] in attend is for: it permits a different rounding. Without it, the schedule below is refused:

a fused block of `j` combines a block's keys at once, regrouping the reduction, which it permits with `[reassoc]`, `reduce[reassoc] Online j.`

With that permission, we can use a library law that combines a whole block of keys with the running summary. The tile operations do the rest: they can find a maximum or sum along a row, exponentiate each score, and use the resulting weights in another matrix multiplication. The scores stay local throughout, so we never need to write out the full score matrix. As with a ReLU after a product, the contraction’s intermediate values only get buffers of their own if the schedule asks for them.

We’ll call the scheduled version flash. Because we’ve allowed it to change the rounding, we can no longer require it to produce exactly the same result as attend. We need a reference that follows the new reduction order. Cachet constructs this shadow by applying only the order-changing steps to the contraction, leaving plain loops with no grids, tiles or threads. flash must produce exactly the same result as this reference, which appears as flash_shadow in the claim:

schedule flash of attend on M4Gpu =
  split i by 32, j by 8, c by 16
  grid over i.outer
  compute s at j.outer
  reduce j.inner by (max: halving, sum: order)
  parts [8, 16]
  element [f32; 8, 8]

for i in 0 .. 64 do
  for d in 0 .. 16 do
    var acc'1: Online.Value = Online.empty ()
    for j in 0 .. 32 do
      ... // compute this key's score in acc'0
      acc'1 = Online.combine acc'1 (key acc'0 (load v[j, d]))
    store y[i, d] = normalised acc'1

grid on M4Gpu gi in 0 .. 2
    by tile (si, sj) parts [8, 16]
    writes y[...] as yt do
  ... // initialise the running summary: m, l and o
  for gj in 0 .. 4 do
    ... // compute a 32 × 8 score tile s with mma
    let m2 = max m (reduce Max f32 over cols by halving do s)
    let e = exp (s - m2)
    let alpha = rescale m m2
    l = l * alpha + (reduce Sum f32 over cols by order do e)
    let w: [f32; 8, 16] = load v[...]
    o = o * alpha
    o = mma e w o
    m = m2
  store yt = o / l