AI Hardware: GPUs, TPUs, and Neural Engines

Two things happen at almost the same moment every time someone in Bengaluru boards a flight after unlocking a phone with their face and, thousands of kilometres away in a Google data centre, an engineer kicks off a training run for a new language model. On the phone, a chip smaller than a fingernail compares a fresh depth-map of a face against a stored template and returns "match" in well under a second, drawing a few hundred milliwatts — less power than the phone's screen backlight. In the data centre, thousands of chips wired together chew through matrix multiplications for days, drawing megawatts, to update billions of parameters. Both are "AI chips." Neither could do the other's job efficiently. A phone's face-unlock chip dropped into that data centre would take years to train the model; a data-centre training chip stuffed into a phone would drain the battery in minutes and cook the case. This chapter is about why — what a GPU, a TPU, and a Neural Engine each actually do differently at the level of transistors and data movement, not just as marketing names.

The common thread across all three is a single mathematical fact: almost everything a neural network does, at inference time and at training time, reduces to matrix multiplication. A fully connected layer computes y = Wx + b, which is a matrix-vector product. A convolution can be rewritten as a matrix multiplication via the im2col transformation. Self-attention in a transformer is built from three matrix multiplications per head (query·keyᵀ, softmax-weighted sum with value) repeated across every layer. If you could design one operation that a chip must perform blindingly fast, cheaply, and in enormous volume, it would be exactly this: take two matrices, multiply and accumulate. Every architecture in this chapter is an answer to the same design question — how do you build silicon that does multiply-accumulate (MAC) operations as fast and as cheaply as physics allows?

Why a CPU is the wrong tool for this job

A CPU core is built to run one instruction stream as fast as possible: deep pipelines, branch prediction, out-of-order execution, large caches, speculative execution — all machinery devoted to guessing what a single thread will need next and having it ready. This is exactly the right design for code with unpredictable control flow and serial data dependencies: a compiler resolving symbol tables, a database traversing a B-tree index, a linked list walk where you cannot know node k+1 until you have read node k. A desktop CPU typically has somewhere between 4 and 64 of these heavyweight cores.

Now consider computing a single output element of a matrix product, C = A · B, where C is 4096 × 4096. Each entry cij is an independent dot product of a row of A and a column of B — it does not depend on any other entry of C. There are roughly 16.8 million independent dot products to compute, each one a tight, predictable loop with no branching. This is the opposite of what a CPU is optimized for: instead of a few cores each guessing cleverly about one complicated thread, you want thousands of dumb, cheap cores each doing the same simple multiply-and-add, in lockstep, on different data. That reallocation of silicon — fewer transistors per core spent on branch prediction and speculation, many more identical cores instead — is the entire idea behind a GPU.

GPUs: SIMT and the reuse of graphics-era parallelism

A modern GPU packs thousands of small arithmetic units organized into groups (NVIDIA calls a group of 32 threads executing the same instruction a "warp"). This execution model is called SIMT — Single Instruction, Multiple Threads: one instruction is fetched and decoded once, then broadcast to dozens of execution lanes that each apply it to their own slice of data simultaneously. For our matrix multiply, thread (i, j) is assigned to compute cij, and because none of those threads depend on each other's results, thousands can run at once with no synchronization needed until the very end.

GPUs got here by accident of history, and that history still matters. They were built to rasterize triangles and shade pixels — another workload where millions of independent, identical, simple operations (colour this pixel, colour that pixel) need to happen in parallel. CUDA and similar frameworks exposed that same parallel hardware to general-purpose numerical code, and it turned out that "shade a pixel" and "compute one MAC of a matrix multiply" are structurally the same kind of problem. But a GPU still carries silicon inherited from its graphics ancestry — texture units, rasterizers, ray-tracing cores — none of which a neural network forward pass ever touches. That legacy is the seed of a very common misconception, which is worth naming explicitly.

Correcting a misconception: "a TPU is just Google's version of a GPU"

It is tempting to assume a Tensor Processing Unit is simply a GPU with a different logo — after all, both are described as "AI chips" built for matrix math. This is wrong in a way that matters for understanding performance. A GPU is a general-purpose parallel processor: it still schedules independent threads onto ALUs, issues instructions, and can run almost any parallel workload, graphics or otherwise — general-purpose flexibility that costs silicon area and, more importantly, costs memory traffic, because every thread's ALU still needs its operands delivered to it from a register file or cache on (nearly) every cycle.

A TPU's core, the matrix multiply unit (MXU), is not a collection of independently scheduled threads at all. It is a fixed, non-programmable grid of multiply-accumulate cells wired directly to their neighbours — a systolic array. Data is loaded at the edges of the grid once, and then it flows from cell to cell like blood pumped through tissue by a heart (the name "systolic" is a direct medical analogy), being reused by several cells along the way before being discarded. There is no instruction fetch inside the array, no thread scheduler, no graphics legacy silicon at all — every transistor in the MXU exists purely to move partial sums and operands one hop at a time. The next section makes this concrete with a full worked trace, because the difference between "thousands of threads fetching from memory" and "a systolic array reusing data as it flows" is precisely why TPUs can achieve very high FLOPs-per-watt on large matrix multiplications specifically, while remaining useless for anything that doesn't decompose into that dataflow.

The memory wall, and why data reuse is the real bottleneck

Before tracing the systolic array, it helps to name the constraint it is solving. Multiplying two N × N matrices takes O(N³) multiply-accumulate operations but touches only O(N²) distinct numbers. If a chip fetched every operand fresh from memory for every MAC it performed, it would need roughly N memory fetches per unit of useful compute — and modern chips can perform an arithmetic operation far faster than they can fetch a byte from off-chip memory. This ratio, arithmetic intensity = FLOPs performed ÷ bytes moved from memory, is the central quantity in what's called the roofline model of hardware performance: a chip is either compute-bound (its ALUs are the bottleneck, memory keeps up) or memory-bound (its ALUs sit idle waiting for data to arrive). A design that lets every fetched operand be reused many times before being discarded raises arithmetic intensity without changing the FLOP count — and that is exactly what a systolic array is built to do.

Worked example: tracing a 2×2 systolic array cycle by cycle

Let A = [[1, 2], [3, 4]] and B = [[5, 6], [7, 8]], so that the true product is:

C = A · B = [[1·5+2·7, 1·6+2·8],
             [3·5+4·7, 3·6+4·8]]
  = [[19, 22],
     [43, 50]]

An output-stationary systolic array uses a grid of processing elements (PEs) where PE(i, j) permanently owns accumulator cij for the whole computation, and never moves it. Row i of A streams in from the left, one element per cycle, but row i's stream is delayed by i cycles before it starts (this staggering is what keeps every product arriving at the right cell at the right time). Column j of B streams in from the top the same way, delayed by j cycles. Every cycle, whichever operands a PE currently holds get multiplied, added into its accumulator, and then passed one hop onward — a moves right to the next PE in its row, b moves down to the next PE in its column — so a neighbour can reuse them next cycle, exactly as the memory-wall argument above requires.

Tracing all four PEs cycle by cycle for our numbers:

Every one of those four results matches the true product computed above. Notice what just happened to the value a₀₀ = 1: it was fetched from outside the array exactly once, at PE(0,0), and then reused a cycle later by PE(0,1) without ever going back to memory. The same is true of b₀₀ = 5, reused by PE(1,0). That single hop of reuse is the whole trick, scaled up: a real TPU's MXU is a 128×128 or 256×256 grid, not 2×2, so a value loaded at one edge can be reused across well over a hundred neighbouring cells before it is dropped — arithmetic intensity that a memory-fetch-per-MAC design could never approach. The total latency for our array to drain its last result was 4 cycles (t = 0 through t = 3); for an n × n array of this output-stationary design, the fill-plus-drain time scales as 3n − 2 cycles, confirmed here since 3(2) − 2 = 4.

import numpy as np

A = np.array([[1, 2],
              [3, 4]])
B = np.array([[5, 6],
              [7, 8]])

C = A @ B
print(C)
# [[19 22]
#  [43 50]]

def macs_for_matmul(M, K, N):
    """Multiply-accumulate ops for an (M x K) . (K x N) matmul."""
    return M * K * N

print(macs_for_matmul(2, 2, 2))   # 8 -- four dot products of length 2

def systolic_latency(n):
    """Fill+drain cycles for an n x n output-stationary
    systolic array multiplying two n x n matrices."""
    return 3 * n - 2

print(systolic_latency(2))   # 4, matching the hand trace t = 0..3

Scale the same idea up to a realistic layer size. A transformer feed-forward expansion from d_model = 2048 to d_ff = 8192 performs 2048 × 8192 = 16,777,216 MACs per token, which is 2 × 16,777,216 = 33,554,432 FLOPs (one multiply and one add per MAC). A single CPU core running at a few billion scalar FLOPs per second would need on the order of tens of milliseconds just for this one layer, for one token — and a transformer has dozens of such layers, applied to every token in a sequence, for every token generated. A systolic array that keeps operands resident and reuses them across a 128-wide or 256-wide grid, doing tens of thousands of MACs per cycle instead of one, is what turns "tens of milliseconds per layer" into microseconds — the difference between a chatbot that answers in real time and one that doesn't.

Neural Engines: inference at the milliwatt scale

A GPU's SIMT design and a TPU's systolic array both assume you want to process many inputs — large batches, long training runs — where keeping thousands of ALUs continuously fed is worth the engineering. A phone doing face-unlock has the opposite workload: exactly one inference, on exactly one image, right now, with a battery and a thermal budget measured in single-digit watts, not the hundreds of watts a training GPU or TPU chip draws. This is the job of what Apple calls the Neural Engine and other vendors call an NPU (neural processing unit) — Qualcomm's Hexagon and MediaTek's APU are the same category of chip under different names.

Three design choices distinguish a Neural Engine from a training-oriented accelerator, and all three follow directly from "batch size 1, milliwatts, no cloud connection required." First, precision: training needs enough numerical range and precision (typically FP16, BF16, or mixed with FP32 accumulation) to keep gradient updates stable across millions of steps, but a trained, frozen model being used only for inference can usually be quantized to INT8 or even lower without meaningfully hurting accuracy — and integer multiply-accumulate hardware is dramatically smaller and lower-power per operation than floating-point hardware. Second, the matrix units on a Neural Engine are small and fixed-function rather than a huge programmable systolic grid, because the matrices involved (a single face embedding, one small vision-model layer) are nowhere near the scale that justifies a 256×256 array. Third, and most important, the entire chip is optimized for latency at batch size one, not throughput at large batch size — a training chip's efficiency comes from feeding one huge array with a continuous stream of independent work, which is precisely the thing a phone doing one face-unlock at a time cannot supply.

Comparing the four architectures

The pattern across the row is not "which chip is best" but "which chip matches the shape of the workload": how much independent parallel work is available at once, how much of a power and thermal budget exists, and how much precision the numbers genuinely need. Training a large model produces an enormous, batchable, matmul-heavy workload with a data-centre power budget behind it, which is exactly what a TPU pod or a cluster of GPUs is built for. Running that same trained model once, on one input, on a device with a battery, is a completely different shape of problem, and a Neural Engine is the chip shaped to fit it.

Active recall

Attempt these before reading the answers.

1. By hand, trace a single accumulator computing the dot product [2, 3, 1] · [4, 0, 5] one multiply-add per cycle. What is the value after each of the three cycles, and what is the final result?
2. Why does a GPU parallelize matrix multiplication well but perform poorly, relative to a CPU, on an in-order linked-list traversal or a recursive quicksort partition step?
3. Name the best-fit hardware (CPU / GPU / TPU / Neural Engine) for each, with a one-line reason: (a) training a 70-billion-parameter language model from scratch; (b) face-unlock on a phone; (c) compiling a large C++ codebase; (d) serving millions of concurrent recommendation-model inferences in a data centre.
4. Using the idea of arithmetic intensity, explain why a systolic array achieves better hardware utilization than a design where every ALU independently fetches its operands from memory every cycle.
5. A dense layer has input dimension 2048 and output dimension 8192. (a) How many MACs does one forward pass through this layer perform, per token? (b) How many FLOPs is that? (c) If a chip sustains 4×10¹¹ FLOPs/second, what is the minimum time for this one layer, for one token?
6. True or false, with justification: "Since a TPU cannot render graphics, it must be slower than a GPU at AI workloads."

Answers

1. Cycle 1: 2 × 4 = 8, accumulator = 8. Cycle 2: 3 × 0 = 0, accumulator stays 8. Cycle 3: 1 × 5 = 5, accumulator = 13. Final result: 13 (check: 2·4 + 3·0 + 1·5 = 8 + 0 + 5 = 13). A single accumulator with no neighbouring PEs takes exactly 3 cycles for a length-3 dot product with no possibility of overlap — this is why real speed comes from many PEs working on independent outputs simultaneously, not from making one accumulator faster.
2. Each output element cij of a matrix product is computed independently of every other output element, so thousands of GPU threads can compute different entries at once with zero data dependency between them (embarrassingly parallel, SIMT-friendly). A linked-list traversal or quicksort partition has a strict serial dependency — you cannot process node k+1 until you've read node k's pointer, and the partition boundary only exists after each comparison — so parallel threads would have nothing independent to do, and GPUs additionally penalize the branching that this kind of code needs (threads in the same warp taking different paths serialize instead of running in parallel). A CPU's few, fast, branch-predicting cores handle this dependency-chained, branchy pattern far better.
3. (a) GPU/TPU — training is a massive, batchable, matmul-dominated workload needing huge parallel throughput and high memory bandwidth across many chips. (b) Neural Engine/NPU — a single low-latency inference on a milliwatt power budget with a quantized model. (c) CPU — compilation is sequential, branch-heavy, dependency-chained work with little exploitable data parallelism. (d) TPU (or GPU) — large-batch inference at data-centre scale benefits from the same matmul-optimized, high-throughput-per-watt design used for training.
4. Arithmetic intensity is FLOPs performed divided by bytes moved from memory. If every MAC fetches its own operands independently, an O(N³)-FLOP matrix multiply demands close to O(N³) memory accesses too, which saturates memory bandwidth long before the ALUs saturate — the memory-bound region of the roofline model, where compute sits idle waiting for data. A systolic array loads each operand once and passes it to a neighbouring PE for reuse (as in the trace above, where a₀₀ was used by both PE(0,0) and PE(0,1) after a single fetch), raising the FLOPs-per-byte-moved ratio without changing the FLOP count, which pushes the workload into the compute-bound region where all those ALUs can actually stay busy.
5. (a) 2048 × 8192 = 16,777,216 MACs. (b) 2 × 16,777,216 = 33,554,432 FLOPs (≈ 33.55 million). (c) 33,554,432 ÷ (4 × 10¹¹) ≈ 8.39 × 10⁻⁵ seconds ≈ 83.9 microseconds.
6. False. The inability to render graphics reflects that a TPU discarded all graphics-specific silicon — rasterizers, texture units, ray-tracing cores — and devoted every transistor to its systolic matrix-multiply array and supporting memory system instead. For the narrow task of large matrix multiplication specifically, that specialization typically makes a TPU more efficient in FLOPs per watt and per unit of silicon area than a general-purpose GPU on a comparable process node, precisely because it spends nothing on capabilities that neural-network training or inference never uses.

Think About It

Think about this: How would you explain ai hardware: gpus, tpus, and neural engines to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where ai hardware: gpus, tpus, and neural engines is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting ai hardware: gpus, tpus, and neural engines to at least 3 other topics you have studied.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind ai hardware: gpus, tpus, and neural engines, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.