Google TPU architecture and systolic arrays — Hinglish notes

6.5.7 · Hardware › Advanced & Emerging Architectures

TPU kyun exist karta hai?

WHY: Socho ek fully-connected layer $y = W x$ . Ise compute karne ke liye, tumhe kai multiply-accumulates karne padte hain. CPU par har MAC mein shaamil hai:

weight ko register/cache se fetch karo,
activation ko fetch karo,
multiply karo, accumulator mein add karo,
accumulator wapas likho.

Arithmetic sasta hai; data movement energy aur time par dominate karta hai. Yeh hai memory wall. Ilaaj: aisi hardware banao jahan ek value, ek baar load hone ke baad, jaane se pehle maximum reuse ho.

Systolic array KIYA hai?

Key properties (KIYA yaad rakhna hai):

Sirf local wiring → chote wires → fast clock, low energy.
Massive reuse → har operand memory se ek baar padha jaata hai, poore row/column ke liye use hota hai.
Har op ke liye instruction fetch nahi → dataflow hi program hai.

Figure — Google TPU architecture and systolic arrays

HOW karta hai yeh matrix multiply compute? (Scratch se derivation)

Hum chahte hain $C = A \times B$ jahan $A$ hai $m \times k$ aur $B$ hai $k \times n$ . Element:

$C_{ij} = \sum_{p = 1}^{k} A_{i p} B_{p j}$

Step 1 — sum ko grid par map karo. PE $(i, j)$ ko $C_{ij}$ compute karne ki jimmedaari do. Ise $k$ products accumulate karne padte hain. Toh har PE ko ek internal accumulator chahiye:

$acc_{ij} \leftarrow acc_{ij} + A_{i p} B_{p j}$

Yeh step kyun? Har output ko apna running sum chahiye, isliye accumulation cell ke liye local honi chahiye.

Step 2 — global memory ke bina operands deliver karo. Ek weight-stationary array mein (jo TPU use karta hai), weights $B_{p j}$ cells mein pehle se load karo (weight apni jagah rehta hai). Phir activations $A_{i p}$ horizontally left→right stream karo, aur partial sums vertically top→bottom stream karo.

Har clock par ek cell karta hai: $partial_{out} = partial_{in} + a_{in} \cdot w_{stored}, a_{out} = a_{in}$

Yeh step kyun? $a_{out} = a_{in}$ ka matlab hai activation aage pass ho jaati hai next cell ko — yahi hai ek-baar-padho-bahut-baar-use karne ki trick.

Step 3 — input ko skew karo (staggering). Kyunki cell $(i, j)$ tak data pahunchne mein kuch cycles lagte hain, inputs ko diagonally staggered feed kiya jaata hai (har row ek clock se delay). Yeh har cell ko sahi time par sahi operands par busy rakhta hai.

Step 4 — cost count karo. Ek $n \times n$ multiply ki naïve implementation memory se $O (n^{3})$ operands padhti hai. Systolic array $2 n^{2}$ input values mein se har ek ko ek baar padhta hai aur andar $n^{3}$ MACs karta hai — toh arithmetic intensity hai:

$Arithmetic intensity = \frac{n ^{3} MACs}{2 n ^{2} reads} = \frac{n}{2}$

Yeh kyun matter karta hai? Zyaada arithmetic intensity = compute ke liye kam memory accesses = tum memory wall ko tod dete ho. Ek $256 \times 256$ array ke liye yeh ~128 MACs per byte fetched hai.

Poora TPU block diagram (array ke aas-paas KIYA hai)

Matrix Multiply Unit (MXU) — 256×256 systolic array. Asli kaam wala.
Weight FIFO — off-chip DRAM se weights array mein feed karta hai.
Unified Buffer (UB) — bada (24 MB) on-chip SRAM jo activations rakhta hai; array ko feed karta hai aur results store karta hai. Activations on-chip rehti hain → DRAM round-trips se bachta hai.
Accumulators — neeche se aane wale column partial sums collect karte hain.
Activation unit — multiply ke baad ReLU / pooling etc. apply karta hai.
Koi cache nahi, koi branch predictor nahi, koi out-of-order nahi — jaanbujh kar simple; saari die area arithmetic ke liye jaati hai. Yeh hai domain-specific bet.

Worked examples

Common mistakes (Steel-manned)

Flashcards

Systolic array mein "systolic" kya refer karta hai?

Har clock par data ka cell-to-cell rhythmic pumping heartbeat ki tarah (systole se).

Systolic array mein core computational unit kiya hai?

Ek multiply-accumulate (MAC) cell jo sirf apne nearest neighbours se wired hai.

Systolic array matmul ke liye CPU ko kyun beat karta hai?

Yeh har operand ek baar padhta hai aur use kai cells mein reuse karta hai, memory wall tod ke (high arithmetic intensity).

N×N MAC array ki clock f par peak throughput kiya hai?

2 N^{2} f

FLOP/s (N² cells, har clock par 2 FLOPs each).

TPUv1 ka array size aur rough throughput kiya hai?

256×256 MAC array, ~92 TOPS at ~700 MHz (8-bit).

Kiya TPUv1 training ke liye use hua ya inference ke liye?

Sirf Inference (8-bit integer). TPUs par Training TPU v2+ ke saath shuru hui (bfloat16 floating point).

Weight-stationary array mein kiya theharta hai aur kiya stream karta hai?

Weights pehle se load hote hain aur cells mein rehte hain; activations horizontally stream karti hain, partial sums vertically stream karti hain.

N×N systolic multiply ki arithmetic intensity kiya hai?

~N/2 MACs per operand read.

Systolic arrays bade batches kyun pasand karte hain?

Ek fixed ~2N−1 pipeline fill/drain cost tabhi amortize hoti hai jab M ≫ N ho.

TPU mein Unified Buffer kiya hai?

Bada on-chip SRAM (24 MB) jo activations rakhta hai taaki unhe kabhi DRAM tak round-trip na karna pade.

Teen CPU features bolo jo TPU jaanbujhkar chhod deta hai.

Caches, branch prediction, out-of-order execution (die area arithmetic ke liye jaati hai).

Weight-stationary PE mein per-cell update equation kiya hai?

partial_out = partial_in + a_in·w_stored ; a_out = a_in.

Recall Feynman: 12-saal ke bachche ko samjhao

Socho aag bujhaane wale bacchon ki ek bucket-brigade line. Har bachcha ek fixed kaam rakhta hai (ek weight). Ek bucket (ek number) line se guzarta hai; har bachcha ek chatki maarta hai aur jod leta hai, phir bucket agle bachche ko de deta hai. Koi wapas kuen par bhagne nahi jaata har chatki ke liye — bucket bas aage badhta rehta hai, aur bahut saare bachche ek saath usi par kaam karte hain. Yahi hai systolic array: numbers chhote adders ki ek grid se flow karti hain taaki bahut saari math ho sake jabki har number sirf ek baar fetch hota hai. TPU inhi chhote helpers ki ek badi grid hai jo sirf us math ke liye banayi gayi hai jis par AI nibhar karti hai.

Connections

GPU architecture and SIMT execution — doosra tarika: temporal parallelism + caches vs. TPU ka spatial dataflow.
Memory wall and arithmetic intensity — roofline wajah ki TPUs kyun exist karte hain.
Matrix multiplication algorithms — woh operation jo accelerate ho raha hai.
ASIC vs FPGA vs general-purpose processors — TPU ek domain-specific ASIC ke roop mein.
Dataflow architectures — systolic arrays ek classic dataflow design hain (Kung & Leiserson, 1978).
Quantization and 8-bit inference — kyun TPUv1 INT8 MACs use karta hai.
Roofline model — visualize karo kab tum compute- vs memory-bound ho.