Branch Prediction: From CPUs to GPUs and TPUs

This is sort of random post, more for myself than anybody else, but if you’d still like to go ahead and read it, feel free :)

So essentially, in an NVIDIA GPU, the smallest unit of execution - a warp (which is made up of 32 threads that are locked together - i.e. they execute the same instruction at the same time) is responsible for the high parallelization responsible that is often talked about. Basically in things like Matrix multiplication, every operation looks identical that is, you multiply one row of the first mat with the second mat’s column and add them up. So this multiply and add operation can be run on each thread BUT on different data. And being locked together, it means that these threads will execute the instructions in the same lockstep.

But what happens when the instructions are not identical? Say for example, you have a conditional branch in your code, like an if-else statement. In this case, some threads might take the if branch, while others might take the else branch. And since all the threads in a warp must execute the same instruction at the same time, this would lead to a situation where some threads are idle while others are executing the branch. This is known as branch divergence, and it is a major cause of performance degradation in GPUs because they must execute the commands serially now since you can’t execute both if and else at the same time. This is handled in CPUs via branch predictors, which try to predict the outcome of the branch and execute the instructions speculatively.

TPUs take it a level up - dedicating even a larger portion of the on-chip computations to operations like matrix multiplication. Google’s TPUs use a systolic array architecture, which is a grid of processing elements that are connected in a way that allows data to flow through the array in a highly parallel manner. The need for fetching and writing data are eliminated here, what TPUs do is pass the tensors/arrays on which matmul must be performed from one cell to the other directly without writing to disk saving massive read/write overhead.

---

So lets go back to the CPU question.

Modern CPUs have a pipeline that executes instructions in stages. The problem with branches is that the CPU doesn’t know what to execute until it knows which way the branch goes. Branch predictors solve this by guessing.

Early predictors were simple while modern CPUs use predictors that keep a history of how branches have behaved in the past. This helps them make accurate predictions.

State-of-the-art CPUs use predictors like TAGE, which index into multiple tables using branch histories of different lengths. This gives them accuracy exceeding 95%.

Back to GPUs: how divergence is actually handled

When a warp diverges the GPU serializes the execution. It uses predication and an active mask to execute the instructions.. This can be slow.

GPU programming guides discourage branches within a warp and data layouts matter a lot. If the data is arranged such that neighboring threads always take the path divergence never happens.

TPUs: eliminating the problem entirely

TPUs eliminate the need for branch prediction by using a systolic array architecture. The computation is entirely data-driven. The control flow is fully determined at compile time.

This approach is efficient for tasks like tensor operations but limited in terms of expressiveness. A CPU can run anything while a TPU is exceptional at one thing. Matmuls.

The history of computer architecture is about identifying bottlenecks and building hardware to remove them. Branch prediction is about CPUs getting better at handling code. TPUs are about accepting a constraint to achieve extraordinary efficiency.

The right tool, for the job. Taken at the silicon level.

This post is brief. Was written as a result of the fascinating experience of going through the amazing and FREE book series Machine Learning Systems