Skew, fill and drain

Why the inputs arrive as a staircase, how long the array takes to fill and empty, and what that does to utilisation as the matrices grow or stop fitting.

Loading the animation…

Concept

The skew that chapter 2 needed has a price. Row $m$ of $A$ reaches PE $(k, n)$ at cycle $m + k + n$ , so in any one cycle the busy PEs are those on a band of anti-diagonals: the wavefront. It enters at the top-left corner, takes $K + N - 2$ cycles to reach the far corner (the fill), and after the last row of $A$ has entered it takes as long again to leave (the drain). Only in between, and only if the stream is long enough, is every PE busy.

With the default 10 rows of $A$ through a 4 × 4 array, the fill takes 6 cycles, all 16 PEs are busy only briefly, and the whole run (with the weight load) takes 20 cycles at 50.0% utilisation. Shorten the stream below $K + N - 1$ rows and the array never fills at all.

The skew itself costs hardware too: something has to delay row $k$ by $k$ cycles. The RTL challenge on the model page uses a chain of $k$ registers per row, $\tfrac12 K(K-1)$ for the left edge (6 registers for 4 rows); some designs skew in the memory system instead.

Loading the animation…

Concept

Utilisation is the fraction of the array's MAC slots that did useful work. The fill, drain and weight load are a fixed overhead per pass, so the way to amortise them is a long stream. On an 8 × 8 weight-stationary array with its weights all used, one row of $A$ gives 4.3%, 8 rows give 26.7%, and 64 rows give 74.4%; reaching 90% takes 198 rows. For a 256 × 256 array, as in the first TPU, the overhead is 766 cycles and 90% takes 6,894 rows. This is why such arrays want big batches (more rows of activations per loaded weight tile), and why the TPU double-buffers its weights: a weight load hidden behind the previous tile's stream costs nothing (chapter 5).

Switch to matrix width to see the other loss. Here $M =$ 32 and $K =$ 8 fit, and $N$ grows past the array's 8 columns. At $N = 8$ the tile fills the array (59.3%); at $N = 9$ a second tile is needed for one column, and utilisation drops to 35.6%. It climbs back as that tile fills, and drops again at 17 (43.9%): a matrix dimension just past a multiple of the array size wastes most of a tile. It is one reason model dimensions are usually multiples of a power of two.

Maths

For weight-stationary on a $K \times N$ block, PE $(k, n)$ is busy at compute cycle $\tau$ when it holds a row of $A$ : $0 \le \tau - k - n < M$ . The busy count at $\tau$ is the number of $(k, n)$ in the block with $\tau - M < k + n \le \tau$ ; it rises over the fill, is $KN$ while $K + N - 2 \le \tau \le M - 1$ (if $M \ge K + N - 1$ ), and falls over the drain. Summed over all cycles it is $MKN$ , every MAC once.

Including the $K$ load cycles, a pass takes $T = K + M + K + N - 2$ , so with the block filling an $R \times C$ array ( $K = R$ , $N = C$ )

U = \frac{MRC}{RC\,(M + 2R + C - 2)} = \frac{M}{M + 2R + C - 2},

and $U \ge u$ needs $M \ge u\,(2R + C - 2)/(1 - u)$ : for $u = 0.9$ and $R = C = 8$ , $M \ge$ 198. When a dimension does not fit, the problem is cut into tiles of at most $R$ rows of $K$ and $C$ columns of $N$ , run one after another; tile $j$ has $k_j \le R$ and $n_j \le C$ and costs $T_j = k_j + M + k_j + n_j - 2$ , and

U = \frac{MKN}{RC \sum_j T_j}.

The model runs every tile cycle by cycle and checks the sum of their cycles against this closed form.

Code

The closed form for one pass, cut from src/lib/sa/model.ts (tests check it equals the number of frames the cycle-accurate model produces, for every dataflow and shape):

export function cycles(df: Dataflow, m: number, k: number, n: number): number {
  if (df === "ws") return k + (m + k + n - 2);
  if (df === "os") return k + m + n - 2 + m;
  return k + (n + k + m - 2);
}