Worked examples — Google TPU architecture and systolic arrays

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

Yeh worked-examples companion hai parent TPU note ka. Wahan humne why aur how build kiya tha ek dataflow systolic array ka. Yahan hum har case class drill karte hain jo tum meet kar sakte ho — tiny multiplies, degenerate shapes, batch limits, quantization edge cases, aur ek exam twist — har ek ground up se kaam karte hue.

Shuru karne se pehle, ek promise contract se: koi bhi symbol tab tak appear nahi hoga jab tak woh define na ho. Agar parent ne use kiya, toh hum use yahan dubara earn karte hain.

---

Scenario matrix

Is topic ke har question ko tum ek yaa zyada cells mein daal sakte ho. Neeche ke examples unke cell label ke saath hain.

#	Case class	Kya tricky hai	Example
A	Plain square multiply	"kya wiring sach mein dot product compute karti hai?" sanity check	Ex 1
B	Non-square / rectangular ($m\ne k\ne n$)	array tumhare data ke liye square-shaped nahi	Ex 2
C	Degenerate: array data se bada ($N>$ matrix side)	wasted cells, low utilization	Ex 3
D	Limiting: tiny batch $M\to 1$ vs huge batch $M\to \infty$	pipeline fill dominate karta hai vs vanish ho jata hai	Ex 4
E	Throughput / roofline (peak FLOP/s)	$N^2$ se scale karta hai, linearly nahi	Ex 5
F	Memory-traffic saving (arithmetic intensity)	reuse factor = $N$	Ex 6
G	Zero / sign / zero-point quantization mein	8-bit overflow, negatives, dequant offset	Ex 7
H	Real-world word problem	ek layer ko cycles & time mein translate karo	Ex 8
I	Exam twist: matrix array se bada (tiling)	tumhe multiply split karni padegi	Ex 9

Chalo har cell ko hit karte hain.

---

Example 1 — Case A: plain $2\times 2$ multiply (kya wiring kaam karti hai?)

Forecast: $C_{11}$ guess karo aage padhne se pehle — kya yeh $1\cdot 5+2\cdot 7$ hai, ya $1\cdot 5+3\cdot 6$?

Neeche ki figure ek saath poori picture dikhati hai: steps se pehle ise padho.

Figure kaise padhein (annotation walk-through):

• Char lavender rounded boxes hardware cells hain. Har ek par weight $B_{pj}$ label hai jo woh store karta hai: box ki vertical position row $p$ hai (top = $p=1$, bottom = $p=2$), uski horizontal position column $j$ hai (left = $j=1$, right = $j=2$).
• Left par do coral arrows $A$ ki rows hain jo cells mein rightward march kar rahi hain — ek activation ek baar read hoti hai aur cell-to-cell pass hoti hai (yahi ek-read-many-uses trick hai).
• Bottom par do mint arrows partial sums hain jo har column se downward drop kar rahi hain; column $j$ se jo number girta hai woh exactly $C_{ij}$ hai. Un do mint labels dekho — wahi jawab hai jo tum abhi derive karne wale ho.

Steps.

1. Weights load karo. Cell $(p,j)$ mein $B_{pj}$ store hota hai. Toh cells ka left column (column $j=1$) mein $B_{11}=5$ (top, row $p=1$) aur $B_{21}=7$ (bottom, row $p=2$) hai. Right column ($j=2$) mein $B_{12}=6$ aur $B_{22}=8$ hai. (Yeh figure mein char lavender boxes hain.) Yeh step kyun? Output $C_{ij}={\sum }_p{A}_{ip}{B}_{pj}$ ko $B_{pj}$ chahiye wahan jo matching $A_{ip}$ paas karegi. Column $j$ output column $j$ ka owner hai.

2. $A$ ki row $i=1$ stream karo, yaani $(A_{11},A_{12})=(1,2)$, top mein (top coral arrow) aur yeh across walk karta hai. Column 1 accumulate karta hai: $\text{acc}=A_{11}{B}_{11}+A_{12}{B}_{21}=1\cdot 5+2\cdot 7=5+14=19=C_{11}$ Yeh step kyun? Har cell karta hai ${\text{partial}}_{\text{out}}={\text{partial}}_{\text{in}}+a_{\text{in}}\cdot w_{\text{stored}}$. Partial sum neeche travel karta hai (mint arrow), har cell se ek product uthata hua — yahi hai dot product ${\sum }_p{A}_{1p}{B}_{p1}$.

3. Same row column $j=2$ mein: $C_{12}=1\cdot 6+2\cdot 8=6+16=22$ Yeh step kyun? Wahi activation row $(1,2)$ ek doosre column of cells ke liye reuse hoti hai bina kisi naye memory read ke — yahi woh reuse hai jo column-parallelism deta hai. Yeh second dot product ${\sum }_p{A}_{1p}{B}_{p2}$ compute karta hai.

4. $A$ ki row $i=2$ stream karo, $(3,4)$: $C_{21}=3\cdot 5+4\cdot 7=15+28=43,C_{22}=3\cdot 6+4\cdot 8=18+32=50$ Yeh step kyun? $A$ ki naya row ek naya row of $C$ produce karta hai: weights $B_{pj}$ cells mein stationary rehte hain jabke successive activation rows stream karte hain, toh ek weight-load poore batch ki service karta hai.

Forecast check: correct $C_{11}=1\cdot 5+2\cdot 7=19$ — pehla option. Agar tumne $1\cdot 5+3\cdot 6$ guess kiya toh tumne $A$ ki row ko $B$ ki row (column nahi) se pair kiya; dot product $B$ ke column $j$ se neeche chalta hai, toh doosra factor $B_{pj}$ hai jisme $p$ vary karta hai, $B_{ip}$ nahi.

Verify: by hand, $AB=\left(\begin{array}{cc}19& 22\\ 43& 50\end{array}\right)$. Array ne sum in space wire kiya loop in time ki jagah — same answer. ✓

---

Example 2 — Case B: rectangular multiply ($3\times 2$ by $2\times 4$)

Forecast: pehle $C$ ki shape aur total MAC count guess karo.

Steps.

1. Shape rule. $A$ hai $m\times k=3\times 2$, $B$ hai $k\times n=2\times 4$. Shared dimension $k=2$ (woh length jis par hum sum karte hain) match karni chahiye — karti hai. Toh $C$ hai $m\times n=3\times 4$. Yeh step kyun? $C_{ij}={\sum }_{p=1}^k{A}_{ip}{B}_{pj}$ tabhi sense karta hai jab $A$ ka column count equals $B$ ka row count ($=k$).

2. Outputs count karo. $C$ mein $m\cdot n=3\cdot 4=12$ entries hain. Yeh step kyun? Har entry $C_{ij}$ ko array mein apna खुद ka accumulator cell chahiye, toh outputs count karna batata hai kitne result cells (aur streamed activations ke columns) computation occupy karti hai.

3. MACs count karo. Har entry $k=2$ products ka sum hai, toh total MACs $=m\cdot n\cdot k=3\cdot 4\cdot 2=24$. Yeh step kyun? Matmul mein total kaam hamesha $mnk$ MACs hota hai — yahi number array ko hit karta hai.

4. Array depth needed. Partial-sum direction kam se kam $k=2$ cells deep honi chahiye taaki ek poora dot product accumulate ho sake jab woh neeche flow kare. Ek $256\times 256$ array ek pass mein $k\le 256$ handle karta hai. Yahan $k=2\ll 256$, toh array ka zyada hissa idle hai — yeh Case C ka foreshadowing hai. Yeh step kyun? Contraction length $k$ set karta hai ki partial sum ko kitni dur travel karni hai poori hone ke liye; agar $k$ array depth se zyada hoti toh tumhe tile karna padta (Example 9), isliye $k$ ko $N$ se check karna decide karta hai ki ek pass kaafi hai ya nahi.

Forecast check: $C$ hai $3\times 4$ (square nahi — output $m$ $A$ se aur $n$ $B$ se inherit karta hai), aur MAC count hai $24$. Agar tumne $2\times 2$ ya $4\times 4$ shape guess ki, toh yaad rakho contraction dimension $k$ result se disappear ho jati hai; sirf $m$ aur $n$ bachte hain.

Verify: $mnk=3\cdot 4\cdot 2=24$ MACs; $C$ shape $3\times 4$ mein $12$ entries hain; $24/12=2=k$. Consistent. ✓

---

Example 3 — Case C: array data se bada (utilization)

Forecast: guess karo — kya yeh $100\%$ ke near hai, $50\%$ ke near, ya tiny?

Steps.

1. Useful cells. Ek $2\times 2\times 2$ multiply zyada se zyada $2\times 2=4$ output cells busy rakhti hai (aur depth $2$). Yeh step kyun? Sirf woh cells jo real $C_{ij}$ own karte hain aur real operands dekhte hain kuch compute karte hain.

2. Total cells. $256\times 256=65536$. Yeh step kyun? Array ek fixed physical grid hai; uska cell count hardware ka constant hai, isliye yeh woh denominator hai jiske against kisi bhi problem ka useful work measure hota hai.

3. Utilization: $U=\frac{4}{65536}=6.1\times 10^{-5}\approx 0.0061\%$ Yeh step kyun? Yeh degenerate case hai — array ek fixed-size grid hai; chhote problems unused silicon mein dub jaate hain.

Forecast check: answer tiny hai ($0.0061\%$), $50\%$ ya $100\%$ ke near nahi. Intuition trap yeh hai ki "$2\times 2$ fair chunk use karta hai"; actually utilization $(\text{data side}/N{)}^2$ hai, jo quickly collapse ho jaati hai jab data $N$ se bahut chhota ho.

Verify: $4/65536=0.00006103\dots$, yaani $0.0061\%$. Yahi reason hai ki parent note kehta hai TPUs bade matrices ke liye hain — chhote ops catastrophically wasteful hain. ✓

---

Example 4 — Case D: batch size $M$ ki do limits

Forecast: $\eta (1)$ nearest percent tak guess karo, aur huge $M$ ke liye $\eta$ kya approach karta hai.

Steps.

1. Formula kahan se aata hai. Total cycles $=\underset{\text{fill+drain}}{\underbrace{(2N-1)}}+\underset{\text{stream}}{\underbrace{M}}$. Sirf $M$ cycles results emit karte hain; fill overhead hai. Toh $\eta =M/(M+2N-1)$. Yeh step kyun? Pehla result tabhi aata hai jab data diagonally sab $N$ rows aur $N$ columns cross kar le — woh diagonal delay $2N-1$ clocks hai.

2. Limit $M=1$ with $2N-1=511$: $\eta (1)=\frac{1}{1+511}=\frac{1}{512}\approx 0.00195=0.195\%$ Yeh step kyun? Ek tiny op ek hi result ke liye poora fill cost pay karta hai — worst case. (Yeh parent ka "latency of one small matmul" trap hai.)

3. Limit $M\to \infty$: ${\lim }_{M\to \infty }\frac{M}{M+511}=1=100\%$ Yeh step kyun? Jaise batch badhta hai, fixed 511-cycle fill negligible ho jaata hai — array full efficiency par run karta hai. TPUs throughput machines hain.

4. $\eta \ge 0.9$ solve karo, yeh dhyan rakhte hue ki $M$ puri sankhya hai. Pehle equality set karo: $\frac{M}{M+511}=0.9\Rightarrow M=0.9M+0.9\cdot 511\Rightarrow 0.1M=459.9\Rightarrow M=4599$ Yeh step kyun? $M$ activation rows count karta hai, isliye yeh ek non-negative integer hona chahiye — fractional batch meaningless hai. $\eta (M)$ $M$ mein increasing hai, isliye hume sabse chhota integer $M$ chahiye jisme $\eta (M)\ge 0.9$ ho. Equality boundary exactly integer $4599$ par land karti hai: check karo $\eta (4599)=4599/5110=0.9000$ (exactly $\ge 0.9$), jabki $\eta (4598)=4598/5109=0.89998<0.9$. Agar boundary fractional hoti toh hum ceiling $\lceil M\rceil$ lete; yahan yeh already integral hai, toh $M=4599$ answer hai.

Forecast check: $\eta (1)\approx 0.2\%$ (buri tarah chhota — ek op poori pipeline fill waste karta hai), aur huge-batch limit exactly $100\%$ hai. Agar tumne $\eta (1)$ $50\%$ ke near guess kiya, tumne underestimate kiya ki fixed $511$-cycle fill ek single streamed row ke relative kitna bada hai.

Verify: $1/512=0.001953$; $\lim =1$; $4599/(4599+511)=4599/5110=0.9000$; aur $4598/5109=0.89998<0.9$ confirm karta hai ki $4599$ sabse chhota integer qualify karne wala hai. ✓

---

Example 5 — Case E: peak throughput $N^2$ ke roop mein scale karta hai

Forecast: TOPS number guess karo, phir guess karo ki $N$ double karne par yeh kitna multiply hota hai (×2 ya ×4?).

Steps.

1. Formula. $N^2$ cells, har ek steady state mein 1 MAC/clock, 2 ops per MAC: $\text{Peak}=2N^{2}f\text{ops/s}$ Yeh step kyun? $N^2$ cells mein se har ek har clock fire karta hai — parallelism spatial hai, isliye throughput $\propto$ area $\propto N^2$.

2. $N=256$ plug karo: $2\cdot 256^2\cdot 0.7\times 10^9=2\cdot 65536\cdot 0.7\times 10^9=9.175\times 10^{13}\text{ops/s}\approx 92\text{TOPS}$

3. Side double karo, $N=512$: $2\cdot 512^2\cdot 0.7\times 10^9=3.67\times 10^{14}=4\times (9.175\times 10^{13})$ Yeh step kyun? $N$ double karna throughput quadruple karta hai — yeh "CPU ki tarah clock×cores nahi" wali lesson hai.

Forecast check: answer $\approx 92$ TOPS hai, aur $N$ double karna throughput ×4 multiply karta hai, ×2 nahi. Agar tumne ×2 guess kiya toh tum linearly soch rahe the (jaise CPU cores add karna); array area mein grow karta hai, isliye factor squared hai.

Verify: $2\cdot 256^2\cdot 0.7\times 10^9=9.1750\times 10^{13}$ ops/s; $(512/256{)}^2=4$; ratio $N^2$ scaling confirm karta hai. Yeh peak kaise flat "compute-bound" ceiling ban jaati hai dekho Roofline model mein. ✓

---

Example 6 — Case F: memory-traffic saving (arithmetic intensity)

Forecast: ratio guess karo — kya yeh ×16, ×256, ya ×65536 ke aas paas hai?

Steps.

1. Naïve reads. Har MAC apne do operands dobara fetch karta hai. MACs ki count $=mnk=256^3$, har ek ko 2 reads chahiye: $2\cdot 256^3=2\cdot 16,777,216=3.355\times 10^7\text{reads}$ Yeh step kyun? Yeh memory-wall worst case hai — compute sasta hai, data move karna nahi.

2. Systolic reads. $A$ aur $B$ mein se har ek ($256^2$ entries each) ek baar read hota hai: $2\cdot 256^2=2\cdot 65536=1.31\times 10^5\text{reads}$

3. Ratio: $\frac{2N^3}{2N^2}=N=256$ Yeh step kyun? Reuse factor array side $N$ ke equal hai: fetch kiya hua har byte cells ki ek poori row/column ko feed karta hai.

4. Arithmetic intensity (MACs per input read): $\text{AI}=\frac{N^3}{2N^2}=\frac{N}{2}=128\text{MACs/read}$

Forecast check: ratio ×256 hai (equals $N$), ×16 ya ×65536 nahi. Agar tumne ×65536 guess kiya toh tumne ise cell count $N^2$ se confuse kiya; reuse factor array side $N$ hai, kyunki har input cells ki ek row/column feed karta hai, poori grid ko nahi.

Verify: $256^3=16777216$; ratio $2N^3/2N^2=256$; AI $=256/2=128$. ✓

---

Example 7 — Case G: 8-bit MACs mein zero, sign, zero-point, aur overflow

Forecast: padhne se pehle — ek 8-bit×8-bit product ki sabse badi magnitude guess karo, aur guess karo kyun hume sirf scale ki jagah ek offset (zero-point) chahiye.

Steps.

1. (a) Zero weight. Agar $w=0$ toh kisi bhi activation $a$ ke liye $a\cdot w=0$. Cell phir bhi $a_{\text{out}}=a_{\text{in}}$ along pass karta hai. Yeh step kyun? Zero hardware mein special nahi hai — yeh kisi bhi number ki tarah flow karta hai; isliye zeros dense array par automatically time save nahi karte.

2. (b) Most-negative aur most-positive product (sign corners). Extreme positive signed product do negatives se hai: $(-128)\cdot (-128)=16384$ aur sabse negative single product ek negative aur sabse bade positive se: $(-128)\cdot (127)=-16256$ Dono 16 bits mein easily fit ho jaate hain. Yeh step kyun? Sign corner cover karta hai — do negatives sabse bada positive dete hain, mixed signs sabse zyada negative dete hain, toh humne dikhaya ki dono sign quadrants multiplier par overflow ke bina handle hote hain.

3. (c) Koi accumulator overflow nahi. $k=256$ worst-case products ki magnitude $16384$ sum karo: $256\cdot 16384=4,194,304$ Ek signed 32-bit register $2^{31}-1=2,147,483,647$ tak hold karta hai. Kyunki $4.19\times 10^6\ll 2.15\times 10^9$, overflow nahi. Yeh step kyun? Explain karta hai kyun accumulator (32-bit) deliberately operands (8-bit) se wider hai: bahut saare products sum karna running total badhata hai, toh use headroom chahiye jo 8-bit inputs mein nahi hai.

4. (d) Ek concrete value quantize phir dequantize karo. ReLU activation range $[0,6]$ lo jo unsigned codes $[0,255]$ par map hoti hai. Toh scale hai $s=6/255=0.023529\dots$ aur zero-point hai $z=0$ (real $0$ code $0$ par map hota hai). $r=4.5$ quantize karo: $q=\text{round}\left(\frac{4.5}{0.023529}\right)+0=\text{round}(191.25)=191$ Wapas dequantize karo: $\hat{r}=s(q-z)=0.023529\cdot (191-0)=4.4941\dots$ Round-trip error hai $|4.5-4.4941|=0.0059$, lagbhag aadha scale step $s/2$ — 8 bits ki cost. Yeh step kyun? Poora quantize → compute → dequantize loop dikhata hai jo TPU actually run karta hai: integers MAC array se flow karte hain (cheap, dense), aur results baad mein $s$ se rescale hote hain. Zero-point $z$ wahi hai jo lopsided (all-positive) range fit karta hai.

Forecast check: sabse badi magnitude wala single product $16384$ hai ($(-128{)}^2$ se), do negatives se milta hai, do large positives se nahi (jinका max $127^2=16129$ hai). Aur zero-point isliye chahiye kyunki akela scale sirf ek range stretch kar sakta hai jo $0$ ke aas paas symmetric ho; offset $z$ range slide karta hai taaki all-positive interval jaise $[0,6]$ saare 256 codes use kar sake.

Verify: $(-128)(-128)=16384$; $(-128)(127)=-16256$; $256\cdot 16384=4194304<2147483647$; $\text{round}(4.5\cdot 255/6)=191$; aur $6/255\cdot 191=4.4941$, error $\approx 0.006$. Quantization and 8-bit inference dekho. ✓

---

Example 8 — Case H: real-world word problem (ek dense layer ka runtime)

Forecast: guess karo ki kya fill cost ($\sim 511$ cycles) answer mein noticeably fark karti hai.

Steps.

1. Cycles. $\text{cycles}=(2N-1)+M=511+8192=8703$. Yeh step kyun? Fill/drain ek baar 511 cycles, phir $M=8192$ rows ke liye ek result-row per clock.

2. Efficiency check. $\eta =8192/8703=0.941=94.1\%$ — ek fat batch fill amortize karta hai, exactly jaise Case D ne predict kiya.

3. Time. Ek cycle $=1/(0.7\times 10^9)\text{s}=1.4286\times 10^{-9}\text{s}$. Toh $8703\times 1.4286\times 10^{-9}\text{s}=1.243\times 10^{-5}\text{s}=12.43\mu \text{s}$ Yeh step kyun? Abstract cycles ko wall-clock number mein convert karta hai — unit ($\mu$s) wahi hai jo ek engineer actually report karta hai.

Forecast check: fill yahan barely matter karti hai — $8192$ par $511$ cycles sirf $\sim 6\%$ overhead hai, toh $\eta \approx 94\%$. Bade batch ke saath fixed fill almost invisible hai, yahi reason hai ki TPUs aggressively batch karte hain.

Verify: cycles $=511+8192=8703$; $\eta =8192/8703=0.9413$; time $=8703/(0.7\times 10^9)=1.2433\times 10^{-5}$ s $=12.43\mu$s. ✓

---

Example 9 — Case I: exam twist — matrix array se bada (tiling)

Forecast: pehle $256\times 256$ tiles ki count guess karo, aur kya total MAC count untiled $512^3$ se different hoga.

Steps.

1. Har dimension split karo. $512=2\times 256$, toh $A$, $B$, $C$ ko $256\times 256$ blocks mein todo. Kyunki teen matrix dimensions mein se har ek ($m,n,k$) $512/256=2$ chunks mein split hota hai, $C$ output tiles ki ek $2\times 2$ grid ban jaati hai $=4$ output tiles. Yeh step kyun? Array ek fixed $256\times 256$ window hai; tum ise bade problem par slide karte ho — yeh tiling hai, standard fix jab matrix hardware se bada ho.

2. Output tile per passes. Har output block $C_{IJ}={\sum }_P{A}_{IP}{B}_{PJ}$ contraction tile index $P$ par sum karta hai, aur $k=512$ bhi $512/256=2$ pieces mein split hota hai. Toh 4 output tiles mein se har ek ko $2$ accumulate-passes chahiye. $\text{tile-passes}=(\text{output tiles})\times (\text{contraction tiles})=4\times 2=8$ Yeh step kyun? Ek pass ek $256\times 256$ weight block load karta hai aur ek $256$-column slab stream karta hai; tumhe ek pass per (output-tile, contraction-tile) pair chahiye, unke partial sums add karte hue.

3. Total MACs (untiled count ke equal hona chahiye). Har pass ek poori $256\times 256\times 256$ multiply drive karta hai, yaani $256^3=16,777,216$ MACs. Saare 8 passes mein: $8\times 256^3=8\times 16,777,216=134,217,728$ Ek untiled $512\times 512$ multiply se compare karo, jo $mnk=512^3=134,217,728$ MACs karta hai — identical. Yeh step kyun? Tiling kaam ko space/time mein rearrange karta hai lekin arithmetic total kabhi nahi badalta — woh equality ($8\cdot 256^3=512^3$) correctness invariant hai jo tum exam mein check karte ho.

Forecast check: 4 output tiles aur 8 total passes hain, aur MAC count $512^3$ ke identical hai. Agar tumne guess kiya ki total MACs shrink honge, yaad rakho tiling sirf kaam reorder karta hai — arithmetic $mnk$ conserved hai.

Verify: output tiles $=(512/256{)}^2=4$; contraction tiles $=512/256=2$; passes $=8$; $8\cdot 256^3=134217728=512^3$. ✓

---

Recap

Recall 256×256 array par 90% efficiency kaunsa batch size deta hai?

$M\ge (2N-1)\cdot \frac{0.9}{0.1}=511\cdot 9=4599$ rows. ::: ~4600 rows — TPUs ko bade batches chahiye.

Recall Array side double karne par throughput kyun quadruple hoti hai?

Throughput $=2N^{2}f$; yeh area $N^2$ ke saath scale karta hai kyunki saare $N^2$ cells har clock fire karte hain. ::: Yeh spatial (area) parallelism hai, linear nahi.

Recall Kya tiling total MAC count badalta hai?

Nahi — $8\times 256^3=512^3$. Tiling sirf same arithmetic reorder karta hai. ::: Same total kaam, passes mein split.

Recall 8-bit quantization mein zero-point kya karta hai?

Yeh woh integer code hai jo real value $0$ represent karta hai, jo ek lopsided (all-positive, post-ReLU) range ko bina waste ke 8-bit codes mein fit karne deta hai. ::: Woh offset jo real 0 ko integer space mein anchor karta hai.

Related depth: GPU architecture and SIMT execution · ASIC vs FPGA vs general-purpose processors · Roofline model.