Optimizing GEMM on Apple's M5 from scratch, Part 1: Basics

Introduction

In this post, you and I will iteratively optimize a Metal matrix-multiplication kernel on Apple’s M5 chip. We want to see how fast we can compute $C = AB$, where $A$, $B$, and $C$ are large matrices stored in single-precision floating-point (fp32) format using GPU. This is a special case of single-precision general matrix multiplication (SGEMM)Show sidenoteIn full SGEMM, we compute $C = \alpha AB + \beta C$, where $\alpha$ and $\beta$ are scalars. Here, we set $\alpha=1$ and $\beta=0$.. I’ll use GEMM for short when the context is clear.

The goal of this post is served as a foundational introduction to GEMM on GPU in general and Metal in particular. Although the goal is ambitious but I will try my best. This is the result of a few months of intensive work with my buddies Claude and Codex.

I am not aiming for a state-of-the-art GEMM kernel on Apple hardware but learning how the GPU behaves and why each change makes our kernel faster — or doesn’tShow sidenoteWith the help of Codex, we can actually produce a GEMM kernel that is faster than MPS but only on a specific shape. That’s okay and honestly, agent is too good at kernel engineering.. By the end, I hope you’ll be able to analyze a kernel with the roofline model and reason about its performance: data reuse, thread ownership, and what the profiler is actually telling you.

Everything in this blog is done on my MacBook Pro with an Apple M5 (10-core GPU, 24 GB of memory)Show sidenoteIf you don’t have M5, don’t worry, the FLOPs will be smaller but we can still do the same progression to get fast kernel, the only thing we cannot do on older M-chip is the Neural Accelerator (Neural Core) part., macOS 26.7, Xcode 27.0, and MLX 0.32.1, plugged into power. Every kernel computes the same row-majorShow sidenoteRow-major is how we store matrices in memory.Figure from: The Craft of Coding. and non-transposed matrices with $M=N=K=4096$.

Here is where we’re going. Every kernel in this post fits on one chart:

We start at the bottom left, with a naive kernel starved by memory, and finish at the top right, about ten times faster and ahead of Metal Performance Shaders (MPS)Show sidenoteMPS is a library that provides high-performance matrix multiplication and other linear algebra operations for Metal. We can think this as cuBLAS but for Apple’s GPU instead. Read more about it here. and MLX, on the same plain FP32 units they use.

Before we doing anything, we should ask what the hardware allows: the peak speed we can reach.

M5 GPU Architecture

From the official source [5], the M5 GPU has 10 cores, a Neural Accelerator inside each core, and 153 GB/s of unified memory shared with the CPU.

The GPU is a single cluster of 10 identical cores, each core plays the role of an SM (Streaming Multiprocessor) on an NVIDIA GPU. On M1 and M2, each core has four schedulers, and each scheduler issues one instruction per cycle for a group of 32 threads [6]. That makes 128 FP32 lanes per core. We’ll meet these 32-thread groups, called SIMD groups, in the first kernel.

If the M5 still has 128 lanes per core as M1 and M2, then at its top two clock states, 1.578 and 1.620 GHzShow sidenoteApple does not publish GPU clocks. These come from the GPU’s performance states in the device tree (ioreg -p IODeviceTree)., the FP32 peak would be about:

The number $2$ came from the fact that our main operation is multiply-add (FMA), which is two FLOPs. From the calculation above, the peak FLOPs of our M5 chip is 4.04 TFLOP/s (use the lower bound). That is our compute ceiling for ordinary FP32 arithmetic.

For memory, Apple quotes 153 GB/s [5], [?], and our measured kernel that streams a 1 GiB buffer reaches about 118 GB/s, that is our memory ceiling. To be fair, we’ll use the measured value for our peak memory bandwidth.

Against this ceiling, MPS runs at 3.37 TFLOP/s, or 83%. MLX is right beside it. Below are all kernels we’ll build later:

This worklog has two parts. This part covers the hardware ceilings, the roofline model, and kernels 1 to 4: naive, shared memory tiling, and thread coarsening. Part 2 covers kernels 5 to 9: SIMD-group matrices and TensorOps. Our best kernel, which uses Neural Accelerator, reaches 91%.

Kernel 1: Naive Implementation

Let’s remember how we do matrix multiplication in undergraduate algorithm classes. To multiply two matrices, we visit each element of $C$ and take the dot product of the corresponding row of $A$ and column of $B$. In code, that gives us three nested loops like this:

for each row i of C:
  for each column j of C:
    sum = 0
    for k = 0, ..., K-1:
      sum += A[i, k] * B[k, j]
    C[i, j] = sum

Now look at the outer two loops: they only choose which element of $C$ to compute. Since different output elements do not depend on one another, we can compute these row-and-column pairs in parallel.

for each thread (i, j) in parallel:
  sum = 0
  for k = 0, ..., K-1:
    sum += A[i, k] * B[k, j]
  C[i, j] = sum

To achieve that in GPU programming model, we give each of those computations to a thread.

These threads are executed in parallel by the GPU and each thread has its own data, and will execute the same function so we have SPMD (Single Program Multiple Data) paradigm. Multiple threads form a threadgroup Show sidenoteIf you know CUDA, a Metal threadgroup plays the role of a CUDA block., and together multiple threadgroups cover a gridShow sidenoteWe can map thread to matrix multiplication like this: each thread will compute one element $C$, and its data is the row of $A$ and column of $B$ corresponding to the element it computes. Then all thread will run the same dot product function along $K$ dimension $\implies$ same function, different data..

Our naive kernel uses $32\times32$ threads per threadgroup, so a $4096\times4096$ output needs a $128\times128$ grid of threadgroups:

Metal gives each thread its group’s position in the grid — through threadgroup_position_in_grid — and its own position within that group — through thread_position_in_threadgroup Show sidenoteIn CUDA we use <<<gridDim, blockDim>>> to dispatch a kernel and read blockIdx and threadIdx inside it. Metal’s threadgroup_position_in_grid and thread_position_in_threadgroup play the same roles.. The kernel combines them to find the output row $i$ and column $j$ it owns. Below is the naive kernel implementation in Metal:

1
kernel void matmul_naive(
2
  device const float * A [[buffer(0)]],
3
  device const float * B [[buffer(1)]],
4
  device float * C [[buffer(2)]],
5
  uint M, uint N, uint K,
6
  uint2 block_pos [[ threadgroup_position_in_grid ]],
7
  uint2 thread_pos [[ thread_position_in_threadgroup ]],
8
  uint2 threads_per_group [[ threads_per_threadgroup ]]
9
) {
10
    // x selects the column, y selects the row
11
    uint j = block_pos.x * threads_per_group.x + thread_pos.x; // column
12
    uint i = block_pos.y * threads_per_group.y + thread_pos.y; // row
13

14
    // check if the thread is within the bounds
15
    if (i < M && j < N) {
16
      float sum = 0.f;
17
      for (uint p = 0; p < K; ++p) {
18
          sum += A[i * K + p] * B[p * N + j];
19
      }
20
      C[i * N + j] = sum;
21
    }
22
}

Let’s follow our thread through the whole kernel. We choose thread $(29, 3)$ in threadgroup $(2, 1)$ as an example. It owns $C_{35,93}$, so it walks along row 35 of $A$ and down column 93 of $B$, one element of each per loop step:

Before we go any further, one choice in the naive kernel deserves a closer look. Go back to the first two lines of the kernel. A thread at position $(x, y)$ computes $C_{y,x}$: the horizontal coordinate $x$ becomes the column $j$, and $y$ becomes the row $i$. But why?

• In memory, the column index varies fastest. $C$ is row-major, so $C_{i,j}$ lives at offset $iN + j$. Moving one column to the right moves 4 bytes; moving one row down jumps $4N$ bytes, which is 16 KiB when $N=4096$.
• Among threads, x varies fastest. A threadgroup numbers its threads linearly as $x + 32y$. The GPU executes them in SIMD groups of 32 threads that issue each instruction together and it packs consecutive linear indices into the same SIMD group Show sidenoteThe same concept in CUDA is called warp. In CUDA, we always have a fixed number of threads in a warp, which is 32. The figure is from [1]. But in Apple’s GPU, we cannot assume the same. We need to use use threadExecutionWidth to get the actual warp size. In my M5, it is 32..

Put the two statements above together. With our mapping, the 32 lanes of a SIMD group share a row $i$ and hold consecutive columns $j_0, \dots, j_0+31$. At step $p$ of the loop they all want the same $A_{i,p}$, and they want $B_{p,j_0}, \dots, B_{p,j_0+31}$, which are 32 neighbouring floats in one 128-byte span.

Flip the mapping and the lanes share a column instead. Then they all want the same $B_{p,j}$ but 32 different rows of $A$, each 16 KiB from the next.

From our measurement, the naive kernel reached 365 GFLOP/s, which is 9% of the FP32 ceiling and 11% of MPS. So where does the other $91\%$ go? We will move to another important concept in GPU programming: Roofline Analysis.

Roofline Analysis

To answer that, we need a way to say what a kernel could reach before we measure it. The roofline model [7] gives us one, and it needs only the two ceilings from the setup section. We will follow Part 1 of the JAX Scaling Book [8] closely (Really recommend this book to anyone who wants to be good at inference).

Compute Time and Memory Time

A kernel does two kinds of work: it computes (multiplying and adding numbers as in GEMM), and it moves data between memory and the coresShow sidenoteThe memory here can be refered to registers, threadgroup memory (or shared memory in CUDA), devices memory (or global memory in CUDA) and even cache too.. Both take time, and our two ceilings tell us the least time each one could take.

Let’s $F$ be the number of floating-point operations (FLOPs) that kernel performs and $Q$ be the number of bytes moved to or from device memoryShow sidenoteMetal calls the GPU’s main memory device memory, after the device address space our buffers live in. It plays the role of CUDA’s global memory. For now, “memory” in this section means device memory and its 118 GB/s ceiling. We’ll see soon why that qualifier matters.. Write $P = 4.04$ TFLOP/s for the compute ceiling and $\beta = 118$ GB/s for the memory bandwidth.

If the GPU did nothing but arithmetic at full speed, the arithmetic would take:

If it did nothing but move bytes at full bandwidth, the traffic would take:

A GPU can do both at once: while some SIMD groups wait for their loads, others keep the arithmetic units busy. At best, the slower of the two hides the faster one completely, so:

Whichever term is larger names the kernel’s bottleneckShow sidenoteWith no overlap at all, the two times add up instead, so a real kernel lands somewhere between the max and the sum. So we have this inequality: $\max(T_\text{compute}, T_\text{memory}) \leq T \leq T_\text{compute} + T_\text{memory}$..

• If $T_\text{compute}$ is larger, the kernel is compute-bound: the arithmetic units are the busy part, and memory keeps up.
• If $T_\text{memory}$ is larger, the kernel is memory-bound: the arithmetic units sit idle, waiting for data.

Arithmetic Intensity

If a kernel is compute-bound, then $T_\text{compute} > T_\text{memory}$, or $F/P > Q/\beta$. Then we can rearrange the inequality to:

The left side describes how much arithmetic the kernel does for the bytes it transfers and it also has its own name.

On the hardware side, $P/\beta$ is the intensity at which the compute time and memory time are equal. On our M5:

We can understand it this way: in the time the M5 moves one byte from device memory, it can do about 34 FLOPs.

Let’s recall the inequality:

Dividing $F$ by both sides of the bound on $T$, we get:

If we plot this bound against $I$ on log-log axes, we will get the shape that names the model: roofline model. It is a slanted line $I\beta$ that rises until it meets the flat line $P$, this is the ridge point.

• A kernel left of the ridge sits under the slanted part of the roof. Raising its intensity raises the bandwidth ceiling. If it already reaches that ceiling, it is memory-bound.
• A kernel right of the ridge sits under the flat part. The FP32 ceiling sets the bound, and more intensity no longer raises it.

The Naive Kernel on the Roofline

Now we go back to the naive kernel’s inner loop (this is where the kernel reads and computes the data):

1
sum += A[i * K + p] * B[p * N + j];

Each iteration does one multiply-add, 2 FLOPs, and asks for two floats, 8 bytes. If every one of those requests went to device memory, the kernel’s intensity would be:

That is only $(0.25 / 683) \times 100\% \approx 1 / 2700 \approx 0.036\%$ of the perfect kernel’s intensity. With the same arithmetic, naive kernel is spending more than 2700 times on data movement. So “on paper”, our naive kernel is memory-bound.

Put that on the roofline, at 0.25 FLOP/byte, the naive kernel sits far left of the ridge, and its roof is:

But in our benchmark, we measured 365 GFLOP/s, twelve times higher than that roof! A kernel cannot run above its roof, so one of our inputs must be wrong.

Let’s follow the bytes to find out which one. At 365 GFLOP/s, the kernel finishes its 137.4 GFLOP in about 376 ms. In that time, device memory can deliver at most:

But the kernel asks for 550 GB from device memory insteadShow sidenoteWe need two floats per multiply-add so the total bytes requested is $2MNK \times 4$ bytes or $2 \times 4 \times 4096^3 \approx 550$ GB.. So at least 92% of the loads never reached device memoryShow sidenoteEven if we use Apple’s 153 GB/s spec, it only raises the 44 GB to about 58 GB, still barely a tenth of what the kernel asks for. The conclusion doesn’t depend on our measured bandwidth.. It must be the case that almost data has been cached by the GPUShow sidenoteThis is an inference from counting loads against measured throughput, not a measured cache hit rate.. So the wrong input was $Q$: we counted the bytes the code requests, not the bytes that actually travel from device memory.

Our 0.25 FLOP/byte was a no-cache estimate of arithmetic intensity: it assumed every buffer load reached device memory. With the caches absorbing most of those requests, the intensity is at least $137.4\ \text{GFLOP} / 44\ \text{GB} \approx 3.1$ FLOP/byte. We still count at device memory, the only thing that changed is our estimate of $Q$.

Profiling the Naive Kernel

Besides analyzing and predicting the kernel performance, we can also ask the GPU. Like Nsight Compute on NVIDIA GPUs, Apple’s Instruments can read the GPU’s performance counters while our kernel runsShow sidenoteprobes/profile_kernel.sh in the repository records a Metal System Trace with Instruments’ “Performance Limiters” counter set and averages each counter over our kernel’s GPU time. xctrace has no public option for the counter set, so the script patches Xcode’s template. The counters are sampled for the whole GPU, not per kernel, so keep the machine idle while profiling. The script also times the GPU alone, so its GFLOP/s are a little higher than our tables, which come from the harness and include encoding and submitting the work.. Here is the report for the naive kernel:

A utilization counter is the share of a unit’s peak work that was actually done and a limiter counts the time the unit tried to work but was stalledShow sidenoteThese are Apple’s definitions, which Instruments prints with each counter. For example, “F32 Utilization: Measures the time during which F32 work is executed as a percentage of peak F32 performance”.. For now, we can see three lines:

• F32 Utilization is 9.1%: the FP32 units do 9% of the work they could, the same 9% of the FP32 ceiling our stopwatch gave.
• Buffer L1 Miss Rate is 18.6%: the L1 cache inside each core answers about four out of five buffer reads.
• GPU Read Bandwidth is 120 GB/s: what gets past the GPU’s caches is read from memory outside the GPUShow sidenoteApple describes this counter as reads “from a memory external to the GPU (potentially device memory)”, so it may include the system-level cache that the GPU shares with the CPU. at our device-memory ceiling of 118 GB/sShow sidenoteAs we mentioned before, 118 GB/s is our estimated from measurements, not the real value and it isn’t different much from the real value (120 GB/s)..

Over one 366 ms run, 120 GB/s is about 44 GB, this is the most requests that device memory can deliver. So the dashed circle in our naive kernel roofline figure is where the naive kernel really sits. The caches absorb over 90% of its loads, and the rest saturate device memory.

Areas of Improvement

Every kernel in the rest of this post performs the same 137.4 GFLOP. What we can change is how many bytes each FLOP costs: we can raise the kernel’s intensity toward the 683 that GEMM allows. The perfect kernel gets there by loading each input value once and reusing it thousands of times from somewhere fast. Every optimization in this post is a way of getting closer to that reuse.

Back to the question the first result raised: can threads working on nearby outputs share input values on purpose instead of hoping the cache does it for them? Our next kernel answers it with tiling.

Kernels 2 and 3: Shared Memory Tiling

Look again at one $32\times32$ threadgroup of the naive kernel. In order to compute that threadgroup’s final output, we need a a block of $32 \times K$ from $A$ and a block of $K \times 32$ from $B$. So the threadgroup needs:

Now let’s look how our naive kernel acquires these elements. The 32 threads in a row of the threadgroup share the same row $i$, so they all walk along the same row of $A$. The 32 threads in a column share the same column $j$, so they all walk down the same column of $B$. Over the whole loop, the threadgroup asks for:

So although we only need 262K different elements, our naive kernel asks for 8.4M loads, that means for each element, it is read 32 times. If the threadgroup has size $T \times T$, then this is $T$ times and it is increasing linearly with the threadgroup size.

Our tiled kernel still keeps everything from the naive kernel: the same $32\times32$ threadgroup, and each thread still owns one element of $C$, so thread $(29, 3)$ of threadgroup $(2, 1)$ still computes $C_{35,93}$. The only thing we change is where the inputs come from so we can control it. Each value is:

• Loaded from a device buffer once per threadgroup.
• Then kept on the core.
• Finally read by all the threads that need it.

To keep values on the core, we need a place that every thread in the threadgroup can write to and read from. Metal calls it threadgroup memoryShow sidenoteCUDA calls it shared memory and declares it with __shared__. Metal marks it with the threadgroup address space, the same way device marks our buffers. Besides being a place that all threads in the same block can work with, it is also faster to access than global memory., and we declare it inside the kernel:

1
threadgroup float A_shmem[32 * 32];

Every thread in the threadgroup sees the same A_shmem, and it lives only as long as the threadgroup does. Threads in other threadgroups can’t see it.

But why do we need threadgroup memory? It gives us control: we decide what stays on the core and for how long, instead of hoping the cache decides the same wayShow sidenoteIn CUDA, without any help from caching, the pain point of the naive kernel is that it reads the same element many times, and each read is a trip to global memory, which is expensive. So we change the destination of the trip from global memory to the core. Shared memory has lower latency and higher bandwidth than global memory, which makes each trip cheaper. This is the same idea in PMPP [10] and Simon’s blog [1]..

Tiling the Loop over K

Let’s go back to the loops, as we did for the naive kernel. Each thread still computes one dot product of length $K$, but now we cut that dot product into phases of $T$ steps. In each phase, the threadgroup copies a $T\times T$ tile of $A$ and a $T\times T$ tile of $B$ into threadgroup memory, and every thread does $T$ multiply-adds from the tiles:

for each threadgroup (bCol, bRow) in parallel:
  for each thread (tCol, tRow) in parallel:
    # the element of C this thread owns
    i = bRow*T + tRow
    j = bCol*T + tCol
    sum = 0
    # one phase per pair of tiles
    for ph = 0, T, 2T, ..., K-T:
      # each thread loads one element of A
      A_shmem[tRow, tCol] = A[i, ph + tCol]
      # and one element of B
      B_shmem[tRow, tCol] = B[ph + tRow, j]
      wait until the whole threadgroup has loaded
      for k = 0, ..., T-1:
        # = A[i, ph+k] * B[ph+k, j]
        sum += A_shmem[tRow, k] * B_shmem[k, tCol]
      wait until the whole threadgroup has finished reading
    C[i, j] = sum

Across all phases, $ph + k$ visits every index from $0$ to $K-1$ exactly once, so sum ends up as the same dot product as in the naive kernel. The difference is who loads what. With $T\times T$ threads and two $T\times T$ tiles, each thread loads exactly one element of each tile, and then reads $T$ elements that other threads loaded.

Keep that $T/4$ in mind. Going from $16\times16$ to $32\times32$ tiles halves the bytes requested from the device buffers. If those requests all reached device memory and its bandwidth still set the time, the $32\times32$ kernel should run close to twice as fast. We’ll check that prediction at the end of this section.

Implementation

1
template <uint BLOCK_M, uint BLOCK_N>
2
kernel void matmul_tiling(device const float* A [[buffer(0)]],
3
                          device const float* B [[buffer(1)]],
4
                          device float* C       [[buffer(2)]],
5
                          uint M, uint N, uint K,
6
                          uint2 block_pos  [[ threadgroup_position_in_grid ]],
7
                          uint2 thread_pos [[ thread_position_in_threadgroup ]])
8
{
9
    // one element of A and one of B per thread only works for a square tile
10
    static_assert(BLOCK_M == BLOCK_N, "BLOCK_M must equal BLOCK_N");
11

12
    uint bRow = block_pos.y;
13
    uint bCol = block_pos.x;
14
    uint tRow = thread_pos.y;
15
    uint tCol = thread_pos.x;
16

17
    // 1. allocate threadgroup memory and move the pointers
18
    threadgroup float A_shmem[BLOCK_M * BLOCK_N];
19
    threadgroup float B_shmem[BLOCK_M * BLOCK_N];
20

21
    A += bRow * BLOCK_M * K;                   // A now starts at row bRow * BLOCK_M
22
    B += bCol * BLOCK_N;                       // B now starts at column bCol * BLOCK_N
23
    C += bRow * BLOCK_M * N + bCol * BLOCK_N;  // C now starts at our tile's top-left corner
24

25
    float sum = 0.f;
26
    for (uint ph = 0; ph < K; ph += BLOCK_N) {
27
        // 2. collaboratively load one tile of A and one tile of B
28
        A_shmem[tRow * BLOCK_N + tCol] = A[tRow * K + tCol];
29
        B_shmem[tRow * BLOCK_N + tCol] = B[tRow * N + tCol];
30
        threadgroup_barrier(mem_flags::mem_threadgroup);  // wait until both tiles are full
31

32
        // 3. collaboratively compute from the tiles
33
        for (uint k = 0; k < BLOCK_N; ++k) {
34
            sum += A_shmem[tRow * BLOCK_N + k] * B_shmem[k * BLOCK_N + tCol];
35
        }
36
        threadgroup_barrier(mem_flags::mem_threadgroup);  // wait until everyone is done reading
37

38
        // 4. slide both tiles along K
39
        A += BLOCK_N;
40
        B += BLOCK_N * N;
41
    }
42

43
    // 5. write the result
44
    C[tRow * N + tCol] = sum;
45
}

Barrier Synchronization

Each thread needs a way to wait for each other. In Metal it is one function call:

1
threadgroup_barrier(mem_flags::mem_threadgroup);

When a thread calls it, the thread is held at that line until every thread in its threadgroup has reached the same line, and then all of them continueShow sidenoteLet’s think we have a group of friends who drive to a mall together that each will shop in a different store, but the car leaves only when everyone is back. Without the barrier, someone gets left at the mall. Below is PMPP’s illustration for threads synchronization.. The mem_flags flag adds a guarantee about memory: every write to threadgroup memory made before the barrier is visible to every thread in the threadgroup after itShow sidenoteCUDA’s __syncthreads() does both at once. In Metal, the flag chooses which memory the barrier covers: mem_threadgroup for threadgroup memory, mem_device for device buffers, or mem_none for an execution barrier only. See the Metal Shading Language Specification..

A barrier only waits for the threads of one threadgroup. There is no barrier between threadgroups inside a kernel. That is enough for tiling, because the threads that share a tile all belong to the same threadgroup.

Our kernel calls the barrier twice per phase, and each call protects a different dependence between threads.

• The barrier after section 2 protects a read-after-write dependence: no thread reads a tile before every thread has written its element into it.
• The barrier after section 3 protects a write-after-read dependence: no thread overwrites a tile with the next phase’s values before every thread has finished reading the current ones. Without it, a SIMD group that finishes section 3 early would start section 2 of the next phase while a slower SIMD group is still reading.

Results

Tiling actually works! Both tiled kernels are about 1.8x faster than the naive one. But our no-cache intensity estimate suggested that $32\times32$ tiles could be nearly twice as fast as $16\times16$ ones, since the larger tile halves the requested buffer bytes. In reality, they’re only 2.5% faster. What did we miss?

First, let’s see the profiler results:

The memory counters show what tiling bought us. The naive kernel reads about 120 GB/s from outside the GPU, right at our measured bandwidth ceiling. The tiled kernels run faster while reading only 81 and 40 GB/s. Loading each input once per threadgroup has relieved that demand on global memory. But the $32\times32$ kernel uses half the bandwidth of the $16\times16$ kernel and runs at almost the same speed, we need to know why?

First, let’s check our intensity estimates. The 4 and 8 FLOP/byte we calculated before assumed that every requested byte came from device memory. Even though tiling makes sharing within a threadgroup explicit, the caches can still share inputs between threadgroupsShow sidenoteThreadgroups in the same row of the grid read the same rows of $A$, and those in the same column read the same columns of $B$. A value one threadgroup brought in may still be in the cache when its neighbour asks for it.. The profiler’s $16\times16$ run takes about 206 ms, so its global memory reads is:

That is about half the 34.4 GB requested by the kernel. Using this traffic to estimate intensity gives us $137.4 / 16.6 \approx 8.3$ FLOP/byte. The same calculation for $32\times32$ gives about 8.0 GB of reads and 17 FLOP/byte. Caching changed our estimate of the bytes moved, not what arithmetic intensity means.

At these intensities, the bandwidth roof allows roughly 1 and 2 TFLOP/s. The profiler measures only 667 and 683 GFLOP/s so both kernels sit below their roofs. As we can see, the roofline still helps us see the gap but it doesn’t tell us which part of the kernel prevents us from closing it.

Now, let’s go back to the inner loop:

1
sum += A_shmem[tRow * BLOCK_N + k] * B_shmem[k * BLOCK_N + tCol];

Each thread reads one value from A_shmem and one from B_shmem, then does one multiply-add. Tiling kernel moved those reads onto the core, but it didn’t remove them. With either tile size, every multiply-add still asks for two values from threadgroup memory, along with the address calculations and loop control. Even though a larger tile lets more threads share each buffer load, each of those threads still has to read the values from the shared tile.

The roofline analysis in kernel 1 asked for reuse “from somewhere fast”. Threadgroup memory gave us that reuse between threads. Now we want reuse within a thread. If one thread computes several outputs that need the same input, it can load that input into a register once and use it in several multiply-adds, without another load for each use. That is the next kernel: thread coarsening.

Kernel 4: 1D Thread Coarsening

Take two outputs in the same column, $C_{i,j}$ and $C_{i+1,j}$. At step $k$, they need different values of $A$, but the same value $B_{k,j}$. In the tiled kernel, two threads load that value separately from threadgroup memory. What if one thread computed both outputs? It could load $B_{k,j}$ once into a register and use it in both multiply-adds. We would need three loads instead of four for the same arithmetic.

Thread coarsening [10] gives a thread the work of several threads. Here, each thread computes $T_M$ outputs stacked in one column:

At each step $k$, the thread loads $B_{k,j}$ once and reads $T_M$ different values of $A$, one for each output. That gives $T_M + 1$ threadgroup-memory loads for $T_M$ multiply-adds, or:

With $T_M = 1$, we have the tiled kernel’s two loads per multiply-add. Increasing $T_M$ spreads the one load of $B$ over more outputs. Since those outputs form a strip in one direction, we call this 1D coarsening.

Implementation

1
template <uint BM, uint BN, uint BK, uint TM>
2
kernel void matmul_1D_coarsening(device const float* A [[buffer(0)]],
3
                                 device const float* B [[buffer(1)]],
4
                                 device float* C       [[buffer(2)]],
5
                                 uint M, uint N, uint K,
6
                                 uint2 block_pos  [[ threadgroup_position_in_grid ]],
7
                                 uint2 thread_pos [[ thread_position_in_threadgroup ]])
8
{
9
    uint bRow = block_pos.y;
10
    uint bCol = block_pos.x;
11
    uint tid  = thread_pos.x;  // 1D threadgroup: tid = 0, ..., BM * BN / TM - 1
12

13
    // this thread owns rows tRow * TM, ..., tRow * TM + TM - 1 of column tCol
14
    uint tRow = tid / BN;
15
    uint tCol = tid % BN;
16

17
    // and copies one element of each tile
18
    uint tileRowA = tid / BK, tileColA = tid % BK;
19
    uint tileRowB = tid / BN, tileColB = tid % BN;
20

21
    // 1. allocate threadgroup memory and move the pointers
22
    threadgroup float A_shmem[BM * BK];
23
    threadgroup float B_shmem[BK * BN];
24

25
    A += bRow * BM * K;
26
    B += bCol * BN;
27
    C += bRow * BM * N + bCol * BN;
28

29
    float sum[TM] = {0.f};  // TM results, kept in registers
30
    for (uint ph = 0; ph < K; ph += BK) {
31
        // 2. collaboratively load one tile of A and one tile of B
32
        A_shmem[tileRowA * BK + tileColA] = A[tileRowA * K + tileColA];
33
        B_shmem[tileRowB * BN + tileColB] = B[tileRowB * N + tileColB];
34
        threadgroup_barrier(mem_flags::mem_threadgroup);
35

36
        // 3. compute TM results from the tiles
37
        for (uint k = 0; k < BK; ++k) {
38
            float Btmp = B_shmem[k * BN + tCol];  // one load, reused TM times
39
            for (uint i = 0; i < TM; ++i) {
40
                sum[i] += A_shmem[(tRow * TM + i) * BK + k] * Btmp;
41
            }
42
        }
43
        threadgroup_barrier(mem_flags::mem_threadgroup);
44

45
        // 4. slide both tiles along K
46
        A += BK;
47
        B += BK * N;
48
    }
49

50
    // 5. write the TM results
51
    for (uint i = 0; i < TM; ++i) {
52
        C[(tRow * TM + i) * N + tCol] = sum[i];
53
    }
54
}

Let’s keep following $C_{35,93}$. It lies in threadgroup $(1, 0)$, at row 35 and column 29 of that group’s output tile. Row 35 falls in the band covering rows 32 through 39, so tRow = 4 and tCol = 29. Its owner is thread $4 \times 64 + 29 = 285$. That thread computes $C_{32,93}$ through $C_{39,93}$; our $C_{35,93}$ is its fourth output, accumulated in sum[3].

Before thread 285 can touch any of those values, the 512 threads have to fill the two tiles together, just like in kernels 2 and 3. Each thread copies one element of each tile per phase, so you can read the tile shapes straight from the thread count:

Once the barrier says both tiles are full, thread 285 only looks at its own eight rows of A_shmem and its own column of B_shmem. Now we can count what one phase costs it:

Keeping $B$ in a register gives us fewer loads to issue and fewer addresses to calculate for the same arithmetic. That is the overhead the tiled kernels’ profiles pointed us toward. We now need eight accumulators per thread, though, and the smaller BK means more phases and barriers. The load count gives us a reason to try this kernel; the benchmark will tell us whether the changes pay off.

Here is that reuse at work in two steps of thread 285’s inner loop. At each step, one load of B goes into a register and feeds all eight multiply-adds.

Results

The 1D kernel reaches 1222 GFLOP/s, 1.83x the $32\times32$ tiled kernel’s rate. Giving each thread eight outputs has paid off: we now reach 30% of the FP32 ceiling and 36% of MPS. Let’s check what changed inside the GPU:

The FP32 units do about 30% of their peak work, up from 16.5%, matching the increase in throughput. At the same time, integer and conditional utilization falls from 62% to 35%, and threadgroup-memory read bandwidth falls from 1312 to 435 GiB/s. The kernel is doing the same arithmetic faster while putting less demand on those parts of the GPU. That supports the reason why we tried coarsening: less work around each multiply-addShow sidenoteThe nine loads per step are a count of read requests in the source code. The L1 bandwidth counter measures hardware traffic, which need not equal four bytes for every such request. Also, this comparison includes the larger output tile, the smaller BK, and the new thread mapping; it does not isolate the effect of reusing $B$ in a register..

The external read bandwidth actually rises, from about 40 to 46 GB/s. Did we make device-memory traffic worse? Remember that bandwidth is a rate. The profiled kernel finishes in about 112 ms instead of 201 ms, so its total external reads fall from 8.0 GB to:

As in kernels 2 and 3, we use those external reads to estimate device-memory traffic. The intensity is then $137.4 / 5.1 \approx 27$ FLOP/byte, above our no-cache estimate of 16. Its bandwidth roof is about $27 \times 118 \approx 3.2$ TFLOP/s, while the profiler measures 1.23 TFLOP/s. We have raised the intensity and the performance, but there is still a large gap below the roof.

Instruction Throughput Limiter remains high at 78%. It includes both work and stalls, so it does not identify a precise issue-rate limit. The counter still points us toward instruction work and stalls, and the inner loop still has an input it doesn’t reuse within the thread: every multiply-add reads a different value of $A$.

Our count $1 + 1/T_M$ shows the limit of this particular reuse. Increasing $T_M$ spreads the load of $B$ more thinly, but the one load of $A$ per multiply-add remains. A thread could reuse both inputs by computing a rectangle of outputs instead and that would be 2D coarsening but we will stop coarsening hereShow sidenoteLet’s have some fun to try 2D coarsening yourself, 2D coarsening is really interesting! Remember to read [1] for more details.

Our next kernels will move to modern Metal features: simdgroup_matrix, it let the 32 threads of a SIMD group hold $8\times8$ tiles together and multiply them with one call.

Next

Kernel 4 reaches 1222 GFLOP/s, 36% of MPS, and every multiply-add is still a separate scalar instruction. Part 2 will continue with kernels 5 to 9: SIMD-group matrices, sharing inputs through threadgroup memory, and TensorOps, ending ahead of MPS.

References