Visual walkthrough — Google TPU architecture and systolic arrays

6.5.7 · D2 · Hardware › Advanced & Emerging Architectures › Google TPU architecture and systolic arrays

Hume sirf ek running example chahiye, toh aao ise ab fix kar lete hain aur kabhi nahin badlenge:

$$A=\left(\begin{array}{cc}1& 2\\ 3& 4\end{array}\right),B=\left(\begin{array}{cc}5& 6\\ 7& 8\end{array}\right)$$

$A$ aur $B$ sirf numbers ki grids hain. $A$ mein 2 rows aur 2 columns hain; $B$ mein bhi utne hi hain. Hum unka product $C=A\times B$ chahte hain, jo bhi ek $2\times 2$ grid hai. Hamara poora mission: cells ki ek choti machine banao jo $C$ nikale bina kisi bhi number ko memory se dobara padhe.

---

Step 1 — "Do matrices multiply karna" asal mein maangta kya hai

KYA. Matrix multiplication ka rule kehta hai: $C$ ke row $i$, column $j$ mein jo number hai — use $C_{ij}$ likhte hain (padho "C-eye-jay", row $i$ aur column $j$ par entry) — yeh $A$ ki row $i$ ke saath chalte aur $B$ ke column $j$ se neeche jaate hue banta hai, jo pairs milte hain unhe multiply karo, aur add kar do.

$$C_{ij}=\underset{\text{1st pair}}{\underbrace{A_{i1}{B}_{1j}}}+\underset{\text{2nd pair}}{\underbrace{A_{i2}{B}_{2j}}}$$

Yahan $A_{i1}$ ka matlab hai "$A$ ki row $i$, column $1$ mein entry." Chote numbers sirf grid par addresses hain.

KYUN. Hardware banane se pehle, hume exactly pata hona chahiye ki har output kaunsa sum hai. Machine jo kuch bhi karti hai woh is ek formula ki physical realisation hai. Agar hum target galat le lein, toh koi bhi clever wiring hume nahi bachayegi.

PICTURE. Figure dikhata hai $C_{11}$ bante hue: $A$ ki upar wali row (orange) $B$ ke left column (teal) ke across slide karti hai; crossings woh products hain jo hum add karte hain.

Concretely, $C_{11}=1\cdot 5+2\cdot 7=5+14=19$. Yeh number yaad rakho; machine ko ise reproduce karna hai.

---

Step 2 — Woh ek cell jo ek "multiply-and-pile-up" karta hai

KYA. Hum poora array ek repeated part se banate hain: ek processing element (PE) — ek chota box jo exactly ek multiply aur ek add clock ke har tick par kar sakta hai. "Clock ka ek tick" (ek cycle) heartbeat hai: har cell ek saath ek baar karta hai, phir dobara. Cell ke andar ek stored weight $w$ rehta hai aur ek running total jise accumulator kehte hain, likha jaata hai $\text{acc}$.

Har cycle mein cell karta hai:

$${\text{acc}}_{\text{new}}=\underset{\text{pile so far}}{\underbrace{{\text{acc}}_{\text{old}}}}+\underset{\text{activation arriving}}{\underbrace{a_{\text{in}}}}\cdot \underset{\text{stored weight}}{\underbrace{w}}$$

• $a_{\text{in}}$ = woh number jo abhi left se andar aaya hai ($A$ ki ek entry).
• $w$ = woh number jo is cell ke andar parked hai ($B$ ki ek entry).
• $\text{acc}$ = growing sum jo $C$ ki ek entry banega.

KYUN. Step 1 par wapas dekho: ek output products ka sum hai. Ek "multiply-then-add-to-a-total" box sabse choti machine hai jo aisa sum grow karti hai. Hum in-place pile up karte hain taaki partial answer ko cell se bahar jaana na pade — isi tarah hum memory wall se bachte hain.

PICTURE. Ek cell, labelled: weight $w$ andar frozen, $a_{\text{in}}$ left se enter karta hua, aur accumulator tick up karta hua.

---

Step 3 — Activation ko aage paas karna (ek baar padho, kai baar use karo)

KYA. Ek cell sirf $a_{\text{in}}$ consume nahin karta; woh use apne right wale neighbour ko bhi deta hai, unchanged:

$$a_{\text{out}}=a_{\text{in}}$$

Yeh equation trivial lagti hai, lekin yahi magic hai. Yeh kehti hai: wahi activation, yahan use hone ke baad, travel karta rehta hai aur agla cell ek cycle baad use karta hai.

KYUN. Matmul rule mein, number $A_{i1}$ $C$ ki row $i$ ki har entry mein aata hai ($C_{i1},C_{i2},\dots$). Toh $A_{i1}$ ko memory se ek baar padha jaana chahiye aur phir physically un sab cells mein chalna chahiye jise iska zaroorat hai. $a_{\text{out}}=a_{\text{in}}$ wahi walk hai. Ek baar padhna aur $N$ baar reuse karna exactly wahi hai jo arithmetic intensity badhata hai.

PICTURE. Ek akela activation (orange dot) cells ki ek row mein left→right march karta hua, har ek mein ek alag parked weight se multiply hota hua, jabki partial sums neeche girte hain.

---

Step 4 — Weights park karo: $B$ ko grid mein rakhna

KYA. $2\times 2$ cells ka grid banao. Koi bhi activation flow hone se pehle, hum $B$ ko cells mein load karte hain, ek weight per cell, position ke hisaab se matched: row $p$, column $j$ wala cell $B_{pj}$ store karta hai.

$$\begin{array}{ccc}& \text{col }1& \text{col }2\\ \text{row }1& B_{11}=5& B_{12}=6\\ \text{row }2& B_{21}=7& B_{22}=8\end{array}$$

Layout dhyan se padho: array ka column $j$, $B$ ka column $j$ hold karta hai. Yeh deliberate hai — array ka column $j$, $C$ ka column $j$ produce karega.

KYUN. Har output $C_{ij}$ ko pair $(B_{1j},B_{2j})$ chahiye. Array-column $j$ mein $B$ ka column $j$ vertically stack karke, us column mein stream karta hua activation exactly woh weights milta hai jine use multiply karna hai. Geometry hamare liye bookkeeping kar deti hai.

PICTURE. $2\times 2$ grid mein charon cells mein $5,6,7,8$ frozen, colour-coded by ki woh $C$ ke kis column ki serve karte hain.

---

Step 5 — Activations stream karo aur dekho $C_{11}$ kaise grow karta hai

KYA. Ab grid mein $A$ ki ek row push karo. $A$ ki row 1 lo, jo hai $(A_{11},A_{12})=(1,2)$, aur ise array-column 1 ke cells mein feed karo (jo $5$ aur $7$ hold karte hain). Us column ka accumulator collect karta hai:

$$\text{acc}=A_{11}{B}_{11}+A_{12}{B}_{21}=1\cdot 5+2\cdot 7=5+14=19=C_{11}✓$$

Wahi row 1 ko array-column 2 mein feed karo (jo $6$ aur $8$ hold karta hai):

$$\text{acc}=A_{11}{B}_{12}+A_{12}{B}_{22}=1\cdot 6+2\cdot 8=6+16=22=C_{12}✓$$

Har partial sum (${\text{partial}}_{\text{in}}$) ek cell mein upar se enter karta hai, ek product add hota hai, aur neeche ${\text{partial}}_{\text{out}}$ ke roop mein nikalta hai:

$${\text{partial}}_{\text{out}}={\text{partial}}_{\text{in}}+a_{\text{in}}\cdot w$$

KYUN. Step 1 ke target $C_{11}=1\cdot 5+2\cdot 7$ se compare karo. Array ne ise bina ek bhi instruction fetch ke reproduce kiya — wiring hi program hai. Woh sum jo pehle time mein unfold hota tha (CPU par loop iterations) ab space mein unfold hota hai (cells ke column mein neeche).

PICTURE. Row $(1,2)$ column 1 mein enter karti hui; partial sum $5$ accumulate karta hua, phir $5+14=19$ hota hua jaise woh do cells se neeche girta hai.

---

Step 6 — Skew: kyun inputs diagonally enter karte hain (timing edge case)

KYA. Ek subtlety hai jo hum gloss over kar gaye: cells ek saath har cycle mein kaam karte hain, lekin data ko travel karne mein time lagta hai. Agar hum har activation ko ek hi cycle mein push kar dein, toh upar se aate partial sums galat activation se milenge. Fix yeh hai ki inputs ko stagger (skew) karo: data ki row 2, row 1 se ek cycle baad enter karti hai, row 3 ek aur baad, ek slanted diagonal wavefront banate hue.

$$\text{delay of row }r=(r-1)\text{cycles}$$

KYUN. Kisi product ke correct hone ke liye, activation aur partial sum ko same cell mein same cycle par rendezvous karna chahiye. Kyunki partial sum ko neeche jaane mein har cycle ek hop lagta hai, aur ek activation ko cross karne mein har cycle ek hop, hume unke launch times offset karne chahiye taaki unke paths sahi instant par cross karein. Har row ke liye ek cycle skewing se poora grid synchronise ho jaata hai.

PICTURE. Inputs ka ek slanted "wavefront": data ki top row already array ke andar deep jabki agla row abhi sirf enter ho raha hai — woh tell-tale diagonal staircase.

---

Step 7 — Cost count karo: fill, drain, aur kyun arrays bade batches chahte hain

KYA. Ek $N\times N$ array ke liye (yahan $N=2$; TPUv1 mein $N=256$), pehla result tab tak nahin aa sakta jab tak data poore grid mein diagonally propagate na ho jaaye. Woh warm-up fill hai; pipeline afterward khaali karna drain hai. Saath mein unka cost lagbhag $2N-1$ cycles hai. Pipeline full hone ke baad, results ki ek nayi row har cycle mein nikalt hai. Toh $M$ rows input push karne mein lagta hai:

$$\text{cycles}\approx \underset{\text{fill \& drain, fixed}}{\underbrace{(2N-1)}}+\underset{\text{one row per cycle}}{\underbrace{M}}$$

Efficiency — useful streaming karne wale cycles ka fraction — hai:

$$\text{efficiency}=\frac{M}{M+(2N-1)}$$

KYUN. Us fraction ki shape dekho. Agar $M$ tiny hai, toh fixed $2N-1$ dominate karta hai aur zyaadatar cycles warm up mein waste hote hain. Agar $M$ bahut bada hai, toh fixed cost average mein vanish ho jaata hai. Yahi precisely reason hai ki TPU bade batches ke liye throughput machine hai, low-latency machine nahin — ek idea jise Roofline model aur arithmetic-intensity view quantitative banata hai.

Numbers. $N=256$ ke liye: fill/drain $\approx 2\cdot 256-1=511$. $90\%$ efficiency achieve karne ke liye:

$$\frac{M}{M+511}\ge 0.9\Rightarrow M\ge 0.9M+459.9\Rightarrow 0.1M\ge 459.9\Rightarrow M\ge 4599$$

toh lagbhag 4600 rows (parent ka "$\ge 4608$" rounder $2N=512$ use karte hue).

PICTURE. Efficiency curve near-zero se $1$ ki taraf badhti hui jaise batch size $M$ grow karta hai, $90\%$ line marked ke saath.

---

Ek-picture summary

Upar ki sab cheez ek single frame mein collapse ho jaati hai: grid mein parked weights, left se flow karte activations ka skewed diagonal, neeche cascade karte partial sums, aur $C$ ki finished entries bottom se drop karti hain.

Recall Feynman retelling — plain words mein wapas bolo

Hum do grids of numbers multiply karna chahte the. Har answer-box sirf ek row-times-column "dot product" hai — matching pairs multiply karo aur pile up karo. Toh humne ek chota box banaya jo ek multiply-and-add har heartbeat par karta hai, aur humne unka poora grid glue kiya. Humne doosri matrix ke numbers boxes ke andar park kiye (har ek apni jagah rehta hai — weight-stationary). Phir humne pehli matrix ko left se feed kiya, aur yahan trick hai: har box, ek incoming number use karne ke baad, use apne neighbour ko unchanged deta hai — toh har number memory se ek baar padha jaata hai lekin ek row mein sab jagah reuse hota hai. Partial sums har column mein neeche slide karte hain, badhte hue jab tak woh ek answer-box ke barabar na ho jaayein, phir bottom se gir jaate hain. Incoming number aur girte sum ko sahi box mein sahi beat par milane ke liye, hum input rows ko ek slanted staircase mein release karte hain (har row ke liye ek beat baad). Pipeline warm up karna aur drain karna ek fixed number of beats cost karta hai, toh chote jobs warm up mein time waste karte hain — isliye TPU bade batches pyaar karta hai, jahan har ek beat par ek answer-row niklti hai.

Recall Quick self-test

Activation ko right neighbour ko unchanged kyun pass kiya jaata hai? ::: Taaki ise memory se ek baar padha jaaye lekin row mein har cell reuse kare — woh reuse trick jo memory wall ko beat karti hai ($a_{\text{out}}=a_{\text{in}}$). Inputs diagonally staggered kyun enter karne chahiye? ::: Taaki har activation aur uska descending partial sum same cycle mein same cell mein pahunche; har row ke liye ek cycle skewing se rendezvous synchronise hota hai. Ek $256\times 256$ array ke liye, roughly kitna batch $M$ $90\%$ efficiency deta hai? ::: Lagbhag $4600$ rows (fill/drain $\approx 511$, solve $M/(M+511)\ge 0.9$). Throughput $N$ ki jagah $N^2$ kyun scale karta hai? ::: $N^2$ cells hain, har ek har cycle mein ek MAC karta hai — parallelism spatial hai, toh peak $=2N^{2}f$ FLOP/s.

---

Related: ASIC vs FPGA vs general-purpose processors · Quantization and 8-bit inference · GPU architecture and SIMT execution