| Accelerator | Peak bf16 (TFLOPs/s) | Peak int8 (TOPs/s) |
|-------------|----------------------|--------------------|
| TPU v6e     |                  918 |               1836 |
| TPU v5e     |                  197 |                393 |
| B200.       |                 5000 |              10000 |

And why do older chips, like TPU v4, offer identical throughput at int8 and bf16?

By the end of this article, we’ll cover:

1. What systolic arrays are

2. How the bf16 format works

3. How a bf16 processing element is implemented

4. How to extend it to support int8

5. How to achieve 2x throughput

This article reuses animations from my YouTube video covering the same topic. The article goes into much more depth, but do let me know what you think about the video format! It’s something I want to play around with more in the future.

Systolic Arrays

It would be an understatement to say that matrix multiplication is hard to optimize on traditional processors. I took a shot at CPU matmul optimizations in a previous article and it’s a good chunk of work to get even within a stone’s throw of reference implementations.

But fret not, throwing hardware at the problem is always an option and matmul is no exception. Systolic arrays are a particular type of circuit designed to achieve theoretically optimal arithmetic intensity (~do maximum amount of work per memory access) [1].

Take a look at the animation below to see a 4x4 weight-stationary systolic array in action. TPUs can have multiple systolic arrays per chip, each up to 256x256 in size.

Weight-stationary refers to how weights are loaded into the systolic array. Unlike the streamed activations, weights are loaded via double buffering. This significantly reduces the amount of data needing to be transferred between PEs every cycle.

bfloat16 format

Each cell in the systolic array, a so-called processing element, does one simple operation: a multiply-accumulate—the innermost body of the matrix multiplication code below:

for (int i = 0; i < N; ++i) {
  for (int j = 0; j < M; ++j) {
    for (int k = 0; j < K; ++k) {
      c[i][j] = a[i][k] * b[k][j] + c[i][j];
    }
  }
}

Below is an illustrated execution of the processing element.

bfloat16 format

To bridge the software-hardware gap, we need to understand how floats, particularly bf16 in our case, are encoded in bit form. I’ll assume you are familiar with how unsigned integers are encoded (if not, just nod along and trust the math).

Figure 1 illustrates the bf16 format. It has 3 parts: a sign bit, 8 exponent bits, and 7 mantissa bits. Let’s tackle these 1 by 1.

Sign

The sign bit is simple: 0 depicts a positive number, 1 a negative number.

Exponent

In the formula above in figure 1, notice how the exponent is interpreted as an unsigned integer from which a bias is subtracted and the result is used to exponentiate the number 2. While somewhat arbitrary at first glance, let’s work through this.

The range of possible exponents is between -126 to +127, using up 254 of the 256 total possible values. The remaining values are used for special states like NaN and INF. We’ll ignore the existence of these states (a luxury hardware engineers don't have), but it’s worth highlighting that they exist and the hardware must correctly handle them.

It’s also worth highlighting that single precision floats (fp32) use exactly the same exponent encoding and number of bits.

Mantissa

The mantissa represents a fractional number in the interval [1, 2). This is achieved by having 7 variable bits and an implicit 8th most-significant bit always set to 1. If we interpret this as an integer, it gives us a range from 128 to 255. To get to the fractional number, we can divide by 128.

The Mantissa Invariant: Because of the implicit leading 1, the mantissa always represents a value in the interval [1, 2). Keep this rule in mind—it is going to cause us a minor headache once we tackle multiplication.

Encoding float into bf16 format

Since the above feels a bit reversed - it tells one how to read a bf16 number not how to encode one - take a look at the animation below showing how to marshall a float into the binary format.

bfloat16 multiplication

Now that we understand how bf16 is represented, let’s see how to multiply two bf16 numbers. Figure 3 shows something interesting: algebraically we can determine the sign, exponent, and mantissa independently of each other. Let’s see how to implement this in hardware.

Sign

The resulting sign can be determined with an XOR.

| A | B | XOR |
|---|---|-----|
| 0 | 0 |   0 |
| 0 | 1 |   1 |
| 1 | 0 |   1 |
| 1 | 1 |   0 |

Exponent

Like we saw in Figure 3, multiplying two numbers of the form 2^x and 2^y is the same as adding their exponents 2^{x+y}. But remember, the 8-bit values stored in the hardware (let's call them E_A and E_B) are biased. The actual mathematical values we are multiplying are:

We need to store this new product back into the bf16 format. This means the hardware must compute a new stored exponent, E_result, which will also have a single bias of 127 implicitly subtracted from it during decoding:

If we solve for the value the hardware actually needs to store (E_result), we get:

This is great news! To compute the new exponent, the circuit just naively adds the two 8-bit unsigned integers together and subtracts 127. This can be implemented extremely efficiently using standard integer adders.

Mantissa

Earlier, we saw that the mantissa is of the form 1.M, where M is 7 bits. Logically, this is equivalent to an 8-bit integer implicitly divided by 128.

We also established the mantissa invariant: the value must always stay in the interval [1, 2).

If we multiply two mantissas together, we’ll get a 16-bit product that can fall anywhere in the range [1, 4). That breaks the invariant! Let’s see this in action using good ol’ long multiplication.

If the product happens to be 2 or greater, it breaks our invariant! To re-establish it, the hardware has to check the result. If the product overflowed into the [2, 4) range, the circuit divides the mantissa by 2 (a simple bitwise right-shift) and compensates by adding 1 to the exponent. A proper implementation would have to check if this +1 causes an exponent overflow, but we will happily sweep that under the rug to protect our sanity.

After normalization, the next thing we need to discuss is the number of bits. Since we’ve established the invariant, we have 15 bits and the MSB is guaranteed to be 1, so we can treat it as the leading implicit 1 bit. That leaves us with 14 variable bits which, if I count right, is 7 more than we fit into a bf16 mantissa. This leaves us with 2 options:

1. Round down to bf16

2. Stay in higher precision

It might be tempting to go with choice one, but let’s think back - these multiplications happen in huge systolic arrays. Rounding errors can quickly accumulate and become catastrophic at that scale.

If particularly curious, try to write a script to simulate the rounding losses. Does it match your intuition?

Instead, it’s common to keep the bf16-bf16 product in fp32. That provides enough breathing room and as seen in figure 3, the only difference between bf16 and fp32 is the number of mantissa bits, making promotion trivial.

Let’s see bf16 multiplication end-to-end in the circuit form.

Addition

So we’ve multiplied two bf16 numbers and produced fp32 product. Now we need to add it to the fp32 partial result coming from the PE above.

We don't need to dive too deep into addition, but the key takeaway is that float addition requires a lot of circuitry to align decimal places before the actual math can commence. Watch the video below and notice how the circuit might have to shift numbers twice: once during initial alignment and then later to fix mantissa invariant if leading 1s get canceled out.

These shifters are expensive! In my local implementations synthesized with yosys, the 8x8 multiplier from earlier (with support for signed and unsigned integers) costs ~680 NAND2 gate equivalents, whereas the two shifters cost combined 760 NAND2 gates.

This is a notable difference from fp32 PEs and a big part of the reason that bf16 is preferred. In fp32 the cost is dominated by the multiplier unit since its footprint scales quadratically with the number of mantissa bits. Addition, on the other hand, is dominated by the alignment shifters [2].

Supporting int8

So far we have a PE that supports bf16 - that’s very cool, but what about int8? Out of the box we don’t get support and we need to do a little extra work.

Obviously, we could just extend the PE with dedicated int8 multiplication logic, dedicated int8 addition logic, and call it a day. That would work but chips don’t grow on trees so let’s try to be a little more economical.

int8 multiplication

You might’ve already noticed earlier - we already have an 8x8 multiplier for mantissa multiplication! Mantissas are unsigned whereas for int8 support we’d like to also have signed multiplication, but this is doable. One approach (and I’m not sure if it’s the one actually used in production PEs) is to use 9x9 signed multiplier.

Let’s see why this works:

1. uint8 can be promoted to int9 by adding a leading 0 bit (positive)

2. positive int8 can be promoted identically

3. negative int8 is a little nuanced, but it just requires copying the leading bit (sign extension)

Let’s walk through the last case using -5 as an example. In two's complement, the most significant bit acts as a negative weight:

# Two's complement -5
1111_1011 = -128 + 64 + 32 + 16 + 8 + 2 + 1 = -5
# Promoting to int9
1_1111_1011 = (-256 + 128) + 64 + 32 + 16 + 8 + 2 + 1 = -5

Notice how after the promotion, the math automagically balances out! The first two bits (-256 + 128) collapse right back into the -128 we started with.

My understanding is that since we aren’t actually using the full range of possible int9 products, the synthesized multiplier can prune some logic. We also don’t need the rest of the bf16 multiplication logic, so those can be completely disabled in int8 mode as shown in the video below.

Notice how we end up with a 16-bit integer product, which we can fit into the bottom half of the fp32 register.

Addition

In int8 mode, the partial sums are stored in int32 and the product is an int16. We simply bypass the fp32 addition logic completely (no longer needing the complex shifters) and add a dedicated adder.

Adders are much cheaper and since the int8 path is completely exclusive with the bf16 path, there’s likely quite a bit of slack even in a heavily pipelined implementation. An optimized implementation would take advantage of this slack to reduce the required chip area for the adder.

Have we done it?

Have we achieved 2x throughput? Sadly, not yet. However, this is basically what older TPUs up to v4 did [3].

Let’s see why exactly this achieves identical peak throughput as bf16 and use it as an opportunity to briefly discuss three concepts:

• Clock frequency

• Critical path

• Pipelining

Clock frequency and critical path

The animation below illustrates the concepts of clock frequency and critical path. Notice how signal leaves the input registers on positive clock edge and has to propagate to the output registers before the next positive clock edge.

If the clock frequency is too high for the circuit, the signal doesn’t have time to settle into the output register. It’s common to have paths that are less complex but are required to operate at the global clock cycle speed determined by the critical path.

If we relate this back to bf16 and int8 paths, in the current implementation it’s likely the shifters that lie on the critical path. The int8 logic is much simpler, but is forced to operate at the same clock speed.

Pipelining

As we’ve seen, critical path is what limits the chip’s clock frequency. There are two ways forward:

1. Optimize critical path’s logic

2. Insert new registers into the critical path.

Optimizing the critical path might require divine intervention and even then there might be limitations. Let’s take a look at option 2, commonly known as pipelining.

In the animation below, the bf16 path uses a 4-stage pipeline, whereas the int8 path uses only 2 stages. To be honest, I’m not sure what is common in production chips - either would work with some extra accounting needed for cycle-by-cycle differences between bf16 modes and int8 modes.

While watching the video, notice even though the latencies are different, the throughput at saturation is exactly the same. This drives home the idea that int8 mode on its own doesn’t provide extra throughput.

Subword parallelism

We are on the final stretch, let’s keep pushing through.

So, how can we increase throughput to 2x? The obvious next thing to look at is that we are underutilizing our input wires by only using the bottom 8 bits in int8 mode. We can pack another int8 number there, the question remains which one?

Try to see if you can figure the answer out on your own before watching the animation below.

I suspect you see where this is going. We pack two neighboring (with respect to the dot product) int8s into a single PE. This technique—formally known as subword parallelism [5]—has actually been around since the 90s for multimedia acceleration. By applying it here, we just need the PE to do the elementwise multiplications in parallel and accumulate them. To do this, we need to add two things: a second 8x8 multiplier and a 3-way adder (to simultaneously sum product A, product B, and the incoming partial sum from the PE above).

The following animation illustrates the required changes.

This is pretty much exactly as one would expect, no surprises. Let’s talk briefly about the fact that we combine the two partial products into a single int32 accumulator. Since this operation is repeated many times, there’s a good question to ask: can it overflow?

Let’s take the worst-case: the largest possible product by the 8x8 multipliers is (-128) x (-128) = 16,384, which easily fits into int16. With subword parallelism, this doubles. Int32 has a max_val of 2^31 - 1, meaning we could repeat this operation 65k times before we have to worry about overflows.

Even if the systolic array was tiling over massive contracting dimension K, it would require K > 130k to even reach the possibility of overflow while assuming the absolute worst scenario. In practice, some products will be positive, some negative, so these largely cancel out. Suffice to say, this will never overflow, but good to sanity check it.

Cost

We’ve finally achieved 2x throughput. As we’ve seen before, the 8x8 multiplier cost around 700 NAND2 gates, so the total multiplier cost doubles.

In my functional, but not particularly optimized, SystemVerilog implementations, adding support for subword parallelism increases the cost in terms of NAND2 gates by a little over 20%. Considering that it gives us 2x throughput that seems like a no-brainer and is exactly why int8 subword parallelism is so commonly supported on modern systolic arrays.

Summary

We’ve covered a lot! We’ve seen:

• what systolic arrays are

• what processing elements do

• how bf16 format works

• how mantissa multipliers can be reused for int8 multiplication

• how clock cycle, critical path, and pipelining are used

• how subword parallelism achieves 2x throughput

And crucially, we’ve answered the question we started with: why TPU v4 doesn't offer a 2x throughput increase at int8, and why newer chips do.

References

[1] H. T. Kung, "Why Systolic Architectures?," Computer, vol. 15, no. 1, pp. 37-46, Jan. 1982.

[2] J. L. Hennessy and D. A. Patterson, Computer Architecture: A Quantitative Approach, 6th ed. Cambridge, MA, USA: Morgan Kaufmann, 2017. (Appendix J)

[3] N. P. Jouppi et al., "TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings," Proc. ACM/IEEE 50th Annu. Int. Symp. Comput. Archit. (ISCA), 2023.

[4] R. B. Lee, "Subword Parallelism with MAX-2," IEEE Micro, vol. 16, no. 4, pp. 51-59, Aug. 1996.

A Note from Michal

Hardware architecture can get pretty heavy, so my goal was to keep this as approachable as possible. I hope you learned something new about how ML chips tick and enjoyed the animations along the way! For me, understanding how the chips I work with on a daily basis work under the hood is super rewarding!

Enjoyed this article?

• Subscribe to the Substack: If you haven’t already, consider subscribing!

• Check out Animated Compute: Content for this article is entirely based on the research I did for my latest YouTube video. I’m planning to post more videos and work on improving the production quality, so I hope you’ll subscribe :)

After a well-deserved break for reading this far, perhaps check out some of my older articles.

This innovation is particularly important for LLM API providers like Google or OpenAI, where predictable performance is essential for user experience.

In this post, I’ll provide a short overview of generative LLM inference and go over DistServe’s internals.

Generating text with LLMs

Let me briefly illustrate how the generative process works. This will be important to appreciate the core insight behind DistServe.

Decoder-only models are next-token predictors, in other words, every forward pass assigns a likelihood to each token in the model’s vocabulary. Then depending on the sampling function used, one of the most-likely tokens is selected. In the example above I’m using simple greedy sampling.

I’ll be using the terms token and word interchangeably.

A naive way to implement the generative loop would be to take the token generated at step 0, append it to the user input, rerun the whole modified input through the model, and repeat until termination. As you can imagine, this would be prohibitively expensive.

Decoder-only transformers commonly use causal attention — a context mechanism where the hidden representations of input tokens are functions of the hidden representations of preceding tokens and of the current token.

This allows inference engines to cache all of the intermediate tensors that will be required in subsequent generation passes. For instance, on the first iteration of the example above, all intermediate tensors that are required for the attention layer will be cached. On the second iteration, it’s sufficient to compute the forward pass just for the token “Apples” and substitute in the cached tensors where required.

This technique is known as KV (Key-Value) caching since the intermediate tensors that need to be cached are the K, V attention tensors.

Two phases of LLM inference

You might notice that not all steps in the generative loop above are created equal. The first step, often called the prefill or context phase, is much more computationally intensive.

Similar to how website load times are measured, LLM API providers monitor the time to first token (TTFT).

It’s customary to express SLOs in percentiles. For instance, 50ms p99 TTFT would mean that the service has to generate the first token within 50ms for 99% of requests.

In contrast, the subsequent decoding phase benefits from KV caching and is more often memory-bound. A common SLO here is the inter-token latency (ITL).

Taken together, the TTFT + ITL * num_generated_tokens gives the overall request latency.

vLLM setup

Let’s walk through a reasonable inference setup with vLLM. vLLM and DistServe serve a similar purpose, so understanding vLLM will help us appreciate DistServe’s differences.

vLLM is a popular LLM inference library initially developed at UC Berkeley that was the first to introduce paged attention for improving GPU memory utilization. Since then, vLLM has incorporated many common LLM inference optimizations, for instance, quantization, prefix caching, or speculative decoding.

Let’s suppose we have a single machine with two GPUs and each GPU is large enough to fit the entire model we want to serve, say a llama-3-8B.

A pretty reasonable inference setup might look something like this:

1. Run a separate vLLM llama-3-8B instance on each GPU.

2. Set up a load balancer in front of these two model instances to distribute incoming requests.

The key limitation here is that both prefill and decode phases run on the same GPU.

vLLM implements continuous batching, where new requests can join the active batch on each generation step, significantly improving throughput. However, this creates a non-obvious issue: requests in the computationally intensive prefill stage can significantly slow down decode operations happening in the same batch.

This makes it difficult to optimize the system for both TTFT and ITL at the same time. While other mitigations have been proposed, for instance, chunked prefill, they have known limitations.

DistServe

The folks behind DistServe noticed this prefill-decode interference and came up with a clever idea to split the two phases. Disaggregating the two phases allowed them to guarantee more predictable TTFT and ITL, allowing them to maintain the same SLOs at several times higher request rates than vLLM.

Splitting these isn’t trivial — the prefill phase generates KV cache blocks required for the decode phase so they had to come up with a mechanism for moving cache blocks between instances.

DistServe is implemented on top of Ray, a popular orchestration platform for Python workloads. DistServe uses Ray actors and placement groups to assign prefill and decode worker instances to GPUs in a Ray cluster.

Let’s explore how exactly DistServe works.

Life of a request

DistServe is implemented as a single API server Python process that spawns and orchestrates Ray workers.

Requests sent to DistServe’s /generate endpoint are added to a prefill queue. The prefill-stage controller has an event loop that on each iteration tries to:

1. Create a new batch from requests in the prefill queue.

2. If a batch can be created, it reserves the memory blocks.

3. Dispatch the batch request to all prefill workers.

4. Wait for the single prefill forward pass to finish.

5. Move finished requests to a bridge queue for the decode stage to handle them.

Let me elaborate on the memory blocks. DistServe, like vLLM, implements paged attention — memory for KV cache is partitioned into blocks akin to operating system memory pages. DistServe’s block manager keeps track of all memory blocks. Information about block usage is then used by the scheduler when constructing a new batch.

The new batch is then dispatched to all prefill workers by using Ray’s remote() call. DistServe supports tensor and pipeline parallelism for each stage separately. In this walkthrough, I’ll be assuming no tensor or pipeline parallelism so there’s exactly one worker for each phase.

Being able to configure tensor and pipeline parallelism settings for each stage independently is yet another benefit of splitting prefill and decode stages.

When the prefill forward step finishes, requests are moved to DistServe’s bridge queue. DistServe has a pretty complex queue setup on the decoder side, so I’ve simplified this significantly in the diagram above.

In a nutshell, a request that has finished the prefill phase will stay waiting in a variety of queues until two conditions are met:

1. Its KV cache blocks can be moved from the prefill worker to the decode worker.

2. Once KV cache blocks are moved to the decode worker, the scheduler can try to add the request to the next continuously batched batch.

The KV cache migration from a prefill worker to a decode worker is interesting. Once DistServe determines that the decode worker has enough free GPU memory, it will send a request to the decode worker to pull relevant tensors from the prefill worker. DistServe uses a custom execution engine SwiftTransformer in which they implement support for block migration via cudaMemcpyAsync.

Interestingly, DistServe authors found that the transfer overhead even for a relatively large 175B model was only around 0.1% of the total latency. This is largely thanks to high-bandwidth InfiniBand and NVLink GPU-GPU communication links.

DistServe has to ensure that the transfer finishes before the corresponding request can be handled by the decode worker. Once the transfer finishes, the prefill worker can free those memory blocks.

The decode scheduler uses continuous batching. Since each decoding request might require a variable number of iterations to terminate, on each iteration it tries to add a new request whose KV cache blocks were already transferred to the GPU. If memory requirements of existing requests grow more than expected, most recently added requests can be temporarily swapped to main memory.

On each iteration, predicted tokens can be streamed back to users. Once a request is finished, all associated memory blocks can be freed or cached in cheaper storage. Caching can be useful to avoid recomputing everything for chat-like use-cases.

Summary

DistServe is an interesting paper that has gathered well-deserved interest from hyperscalers like Google and OpenAI. By splitting the prefill and decode phases across separate workers, it resolves a fundamental issue with traditional serving systems.

For LLM API providers, the lessons from DistServe can likely lead to a reduced need to overprovision hardware to deal with prefill × decode interference. Disaggregating prefill and decode has since been implemented in vLLM, further showing the relevancy of this finding.

I typically build interesting software From Scratch in my articles. With my background in LLM inference and Kubernetes, DistServe lies in a particularly fun problem space for me, so I made an exception and covered the paper. If there’s interest, I might do these more often :)

If you enjoy deep dives like this, consider subscribing! I’m also somewhat active on LinkedIn, so consider connecting with me there. If you enjoyed this article, you also might enjoy some of the pinned articles below!

As we go through this, we’ll learn a thing or two about compilers, assembly, and underlying hardware.

For the sake of ever finishing this article, I’ll focus purely on single threaded optimizations and will tackle parallelization in the future.

Setup

I’m running a Ryzen 5600H processor with 6 cores and12 threads. Each core has 32 KiB of L1d and L1i cache and 512 KiB of L2 cache. All cores share 16 MiB of L3 cache.

I’ll be using clang 18.1.3 on Ubuntu 24.04 throughout this article.

Matrix multiplication review

Given two matrices, A of n rows and k columns ((n,k) from now on) and a (k,m) matrix B of, the product of AB is an (n,m) matrix C. The element C[r][c] is defined as a the dot product of row r of A with the column c of B.

In summation notation:

Let’s use the mathematical definition as the starting point.

Naive implementation

The naive implementation is a direct translation of the mathematical definition to C. For simplicity’s sake, I’ll assume that matrices are square and their size is a power of 2.

Take a look to see if you spot anything surprising.

The restrict keyword tells the compiler that no other pointer will access the object, which in turn can allow the compiler to generate more optimized code.

Let me explain what column-major order is.

Storing matrices in memory

Column-major order is a flat representation of a matrix in memory such that columns are laid out sequentially. This approach is used by Fortran and many BLAS libraries, including AMD’s BLI.

The default ordering of 2D arrays in C/C++ is row major.

Measuring performance

I’ll be measuring the time required to multiply two 4096x4096 matrices of doubles.

I’ve chosen large matrices since some optimization are more pronounced when the matrices don’t fully fit in cache. With N=4096, the 3 matrices have a total size of (3*4096*4096*8)/(1024*1024) = 384 MiB, well over my 16 MiB L3 cache.

I also want to briefly discuss how I’ll compile code. Two main options are:

• Introduce compiler flags as I introduce code optimizations

• Enable all relevant optimization flags from the outset

I’ll go with the second option as I think it’s slightly more transparent. I’ll make sure to highlight when a code change causes “unintentional” optimizations. Unless stated otherwise, I’ll be using the following flags.

clang main.c matmul.c -std=gnu11 -O3 -DNDEBUG -march=native -mfma -ffast-math -mavx2 -lrt -lblis -o matmul

Baselines

With code compiled, we can finally get first results.

The naive implementation multiplies two 4096x4096 matrices in leisurely 480 seconds,

michal@michal-lg:~/code/cpu_matmul$ ./matmul 4096
Elapsed execution time: 480.609548 sec; N: 4096, __TYPE__: double

while the much snappier bli_dgemm does it in 2.6s, around ~200x faster.

michal@michal-lg:~/code/cpu_matmul$ ./matmul 4096
Elapsed execution time: 2.564394 sec; N: 4096, __TYPE__: double

Searching for bottlenecks

I’ve run the naive implementation through perf to collect performance counters. I’m using N=1024 matrices to speed things up when profiling.

The cache miss rate stands out to me as surprisingly large. Let’s use a little tool I wrote to better visualize memory accesses.

The video below shows both the logical and memory representations of matrices. As I step through the naive multiplication, notice where the accessed elements are located in memory.

Focus on which element of A gets accessed on each iteration. While in the logical representation we are accessing sequential elements within a row, due to the column-major order, these elements are actually N elements apart in memory. This is illustrated by the memory views at the bottom.

This is a major issue for sufficiently large matrices. Each access requries reading a new cache line from main memory. By the time the loop wraps around to the first element, chances are the original cache line was already evicted - forcing us to refetch it again.

I’ve illustrated a few iterations for matrix A. For illustrative purposes, I’m assuming a very small, fully associative cache that can only store 3 lines. Each cache line is 16 bytes or 2 doubles.

In the first loop, A[0][0] is read, it’s not in the cache, which results in cache miss, so the cache line has to be fetched from main memory. Notice that we fetch a full cache line even when we want only one element from it. On modern CPUs cache lines are commonly 64 bytes.

Next loop A[0][1] is read - again a cache miss.

Shame we didn’t read A[1][0] instead that we already fetched last iteration…

This keeps repeating until we loop back to A[0][0], but unfortunately, by that point, the cache line that contained A[0][0] was already evicted.

Try to convince yourself that if we instead iterated over matrix A in column order, the cache hit rate would significantly improve.

Loop reordering

Previously, we looped over the matrices in r, c, k order. Based on the findings in the last section, we’d like to iterate in row-order in the inner loop as much as possible.

We can just reorder these loops however we want without affecting correctness.

This presents us with two compelling loop orders: c, k, r and k, c, r.

I expect c, k, r to perform better since k is used as the row index in B, but intuition is often misleading, so let’s measure.

c, k, r:

michal@michal-lg:~/code/cpu_matmul$ ./matmul 4096
Elapsed execution time: 20.784827 sec; N: 4096, __TYPE__: double

k, c, r:

michal@michal-lg:~/code/cpu_matmul$ ./matmul 4096
Elapsed execution time: 51.014005 sec; N: 4096, __TYPE__: double

Just changing the loop order improved performance by ~20x.

Re-running perf stat shows a ~50x reduction in L3 cache misses.

Besides better access pattern, reordering loops allowed the compiler to vectorize some operations! Let’s discuss that next.

Inspecting Assembly

In this section we’ll look at how the assembly generated changes with vectorization on and off.

Non-vectorized

The figure below shows assembly for the inner loop with vectorization disabled. You can see the full function at compiler explorer.

Note: Here I only use compiler flags “-O3 -fno-vectorize”

I’ve annotated the interesting lines. Hopefully with some eye squirming you can convince yourself that the C code does logically translate into the assembly code on the right.

I grew up reading left to right so I prefer the AT&T assembly syntax. Sorry my Arab friends.

As highlighted, clang applied a loop unrolling optimization with a factor of 2, i.e. on each inner loop iteration, it calculates C[r][c] and C[r+1][c].

If it didn’t do this, 3 out of 7 instructions per inner loop execution would be purely for the loop overhead: increment counter, compare with loop limit, and jump. That would be expensive so unrolling ammortizes the cost. In fact, when the number of loop iterations is known at compile time, it’s very common to see completely unrolled loops.

Note how the compiler still generated LBB0_5 in case the loop has an odd number of iterations.

Vectorized

Modern CPUs have extra hardware to support SIMD (single instruction, multiple data) instructions. This lets a single CPU instruction operate on multiple registers in parallel.

I’m using the following compiler flags “-03 -ffast-math -mavx2” so that the compiler generates vectorized code. Again, feel free to play with the compiler explorer code.

I added a promise to the compiler that the data will be aligned to 32 bytes in memory. Some vectorized instructions require memory alignment or have a faster aligned variant.

I’ve again annotated the assembly.

The general sequence is very similar to before, but now the compiler uses AVX2 SIMD instructions. Availability of instructions depends on your CPU. You can read /proc/cpuinfo or run cpuid to see what your CPU supports.

Each AVX2 register (ymm) can fit 256 bits, i.e. four doubles, so each vmulpd, vaddpd, and vmovupd processes 4 multiplications, additions, and stores in parallel.

As you can see, the compiler again did some unrolling. This time with a factor of 4, so each inner loop execution produces 16 writes to C.

The following diagram illustrates the computational dependencies.

You can try to enable the fused add-multiply option on compiler explorer (-mfma flag) to see that vmulpd and vaddpd instructions get merged into a single one.

Given the 4x increase in parallelism, one would hope for 4x speedup - Unfortunately, I’m getting a more modest 1.5x improvement.

An important corollary of the AVX2 register size is that we can trade off numerical precision for higher parallelism. This extends to GPU-based ML training/inference, where half-precision fp16 or even fp8 floats are often used.

Just to briefly recap, we are currently sitting at ~10% of bli_dgemm’s single-threaded performance. Now might be a good time to make a second coffee and stretch a bit before proceeding.

Cache utilization

In this part we’ll cover a really cool optimization that’s unfortunately super unintuitive at first, at least to me anyways.

Let’s inspect the access pattern again. I encourage you to step through this yourself by selecting the naive order(jki) option in my matmul visualization tool.

Notice how the calculation scans through the entire matrix A N times? Ideally we’d load each element once, do all the computations where its needed, evict it, and never load it again.

Sticking with the 4x4 example; When A[0][0] is loaded, we only use it once to calculate the partical result C[0][0] += A[0][0] * B[0][0], before moving onto C[1][0] = A[1][0] * B[0][0].

Instead of moving onto A[1][0] immediately, we’d like to use A[0][0] in as many computations as we can. It’s also used as a dependency for C[0][1], C[0][2], and C[0][3].

But if we try to reorder computation to do that, we are somewhat back to a poor strided access pattern into C. Earlier we saw how strided access can cause elements to be evicted from cache before they are needed again when the full matrix doesn’t fit in cache. Well, what if we could make it fit into cache?

If we can restructure the product of two large matrices into products of smaller matrices, then we can tune the small matrix size so that things fit nicely in cache!

Tiling

The technique I long-windedly introduced above is known as tiling. Let’s first see tiling in action, then I’ll show you the code, and then we’ll do some benchmarking.

I’m using tile size (or block size) of 2 and yet again, you can play with it in the matrix visualization tool.

In the non-tiled version, we scanned through the entire matrix A N times. You can check that every element of A is still accessed exactly the same number of times as before. So what changed? Each tile only needs to be loaded once into cache, since we can tune it to fit.

In general, we are hoping to reduce the number of cache misses by a factor proportional to the tile size.

More loops

Below you can see the tiled implementation. It’s starting to get a bit ugly with the tiling loops, but try to read through it anyways. Note that we do exactly the same amount of work as before (minus look overhead), just in a different order that better utilizes the cache.

I selected the tile_size through the best technique known to man - trying a bunch of values and picking the best.

• 4 - 45s

• 8 - 16s

• 16 - 19s

• 32 - 10s

• 64 - 7.3s

• 128 - 6.6s

• 256 - 9s

We are finally getting to pretty small numbers where sub-second precision is starting to matter. However, there can be a decent run-to-run variance. This variance is primiarly caused by other processes on my system polluting CPU cache, process context switches etc.

For the different tile sizes above, I ran the measurement a few times for each setting and picked the minimum duration for each. This seems counterintuitive, but when noise is the main source of variance, the fastest run is the one with the least noise.

6.6 seconds is a ~3x improvement, putting us at ~30% of bli_dgemm performance.

Tile size analysis

Can we explain why the 128 x 128 tile size performs the best on my system?

When I measure the L3 and L2 cache misses for N=4096 matrices, it turns out that tiled_dgemm performs significantly worse on both counts than our previous implementation. L3 misses are at ~2B compared to ~500M for naive_ordered_dgemm.

The number of L3 reads is lower by about 10B, but I’m not super positive that fully explains the better performance. Suffice to say I was pretty confused at this point. Inspecting the newly generated assembly shows very aggressive unrolling in the tiling block. My best guess is that despite the large cache miss count, the CPU is able to hide read latencies by overlaying reads and computation thanks to the aggressive unrolling.

Still, the fact that number of cache misses went up is surprising to me. Let’s explore that next.

Cache hardware

Until now we were assuming a fully associative cache - one where a cache line can be stored anywhere. Being able to store a cache line anywhere sounds convenient, but the flip side is that to check if a given cache line is stored, all entries need to be checked.

Instead, caches are commonly partitioned into sets of k cache lines. Then for a cache line, based on its address, its sufficient to search through the k lines in its corresponding set.

Let’s see how given an address, we determine where in cache it can reside. I’m basing this on my L3 cache specs: k=16-way associative cache, B=64 byte cache line, M=16 MiB cache size, and w=64 bit words.

Address resolves to a cache line using the following mapping:

• offset - lg(B=64) = least significant 6 bits

• set - lg(M/kB) = lg(16MiB/(16*64)) = 14 bits

• tag - 64 - 14 - 6 = 44 bits

Note that offset contains the least-significant bits. This intuitively makes sense - we want addresses in close proximity to belong to the same cache-line.

So, how does this relate to the increased L3 misses we saw earlier?

Let’s focus what happens when we load a single tile into memory. Suppose that tile[0][0] element is at address x. Then the element tile[1][0] is at address x+8, tile[2][0] at x+16, etc. 8 consecutive elements will belong to the same cache line. tile[8][0] will have address x+64, which flips the 7th bit, so it belongs to a different set. The last element of the first tile column will have address x+127*8 , which means that first column required ceil(127*8 / 64) = 16 lines in 16 different cache sets.

When we wrap around to the next column, tile[0][1] has address x+4096*8 or x+2^15. Recall that we are dealing with 4096 x 4096 matrices. In general, the address of tile[0][i] will be x + N*sizeof(double)*i, where N is the matrix dimension.

Notice how for my L3 cache, only the bottom 20 bits are used to map an address to the cache. When two addresses differ only in top 44 bits, those two address will map to the same cache set. Given the earlier formula, we can calculate which columns will start mapping to the same cache set by solving for i.

x + 2^20 = x + 4096*8*i → i = 2^20 - 2^15 = 2^5 = 32.

So we know that tile[0][0], tile[0][32], tile[0][64], and tile[0][96] all map to the same cache set. Below is a more dramatic visualization. Colored blocks within the tile depict cache lines that map to the same cache sets.

Is this really an issue? After all, we have 16 cache lines per set.

Well, that’s true, but we also have 2 other matrices and there are other processes running on my system. Creating unnecessary pressure on a given cache set increases likelihood of so-called conflict misses. This happens when a line is evicted because its cache set is full.

This issue would be even more pronounced in L2 and L1 caches that are smaller and typically have lower associativity.

Packing

If we lay out a 128 x 128 tile of doubles sequentially in memory, the difference between the first and last address is 2^8*2^8*2^3 = 2^19. With this we don’t run into the same cache set conflicts and the tile is uniformaly distributed across different cache sets.

This suggests an optimization strategy - copy tiles into arrays and do the tile-tile matrix multiplication over the arrays. Our hope is that the better cache hit rate will compensate for the extra copying work.

Instead of directly operating on matrices, we now copy tile of A and tile of B into helper arrays and operate on those. You can see the code and generated assembly at compiler explorer. I reran experiments for different tile sizes and 64 happens to be the optimal size.

I also extracted the inner loops into a separate function, which appears to help the compiler generate more optimized assembly on my hardware. I suspect this has to do with better register allocation. The extraction alone improves runtime from the previous ~6.6s to ~6.2s. Note that the function gets inlined so there’s no call overhead.

So how much does this improve performance?

*drum roll*

We are down to ~4.4s or around 60% of bli_dgemm. Relative to the naive implementation, the current solution is over 100x faster.

As a final sanity check, L3 cache misses did go down by about 4x.

bli_dgemm still runs circles around my implementation when it comes to L2 cache misses and instructions per cycle. That being said, I’m happy with getting within 2x of bli_dgemm.

The comparison isn’t completely fair. bli_dgemm supports arbitrary matrix sizes, while I’m only focusing on large, square, power-of-two sized matrices. If I supported arbitrary matrix sizes, I’d need to add extra instructions for boundary checks.

Next steps?

There are many optimizations we haven’t attempted yet. For instance, we could add another level of tiling for L2 cache and even for L1 cache. We could also experiment with directly using compiler instrinsics for SIMD instructions. Memory prefetching instructions are also worth experimenting with. There are also recursive algorithms that split matrices into smaller ones instead of using tiling.

While exciting, this article has gotten a lot longer than I expected, so I’ll leave those optimization for a potential followup article.

If you this domain interesting and would like to learn more, I recommend these resources:

MIT’s 6.172 on OCW - a fantastic introduction to performance engineering. I’ve been auditing it while writing this article.

Simon’s matrix multiplication article - Simon is a performance engineer at Anthropic. His article covers most of the optimizations I covered here. I like Simon’s succint style and gorgeous illustrations.

OpenBLAS repository on github - highly optimized open-source kernels.

I hope I conveyed not just what optimizations exist but also how one might go about discovering them. I find the latter a lot more satisfying.

If you enjoyed this article, consider subscribing. I’m also somewhat active on LinkedIn, so consider connecting with me there.

After a well-deserved break for reading this far, perhaps check out some of my other articles.

Safety warning

If you do decide to follow along, I’d highly recommend to setup a Linux virtual machine. I’ll be running a bunch of privileged commands and I would like to avoid unintentionally nuking my readers’ systems.

With the warning out of the way, let’s get into it!

Container filesystem

I’ll keep this section brief, for a deeper explanation of container filesystems, especially overlayFS, check out my previous post. In essence, we create a directory structure for our container, download Alpine-based minirootfs, and mount it with overlayFS.

# create folder structure in a temporary directory
mkdir -p /tmp/container-1/{lower,upper,work,merged}

cd /tmp/container-1

# download alpine minirootfs
wget https://dl-cdn.alpinelinux.org/alpine/v3.20/releases/x86_64/alpine-minirootfs-3.20.3-x86_64.tar.gz

tar -xzf alpine-minirootfs-3.20.3-x86_64.tar.gz -C lower

# mount overlayFS, our container root will be in /tmp/container-1/merged
sudo mount -t overlay overlay -o lowerdir=lower,upperdir=upper,workdir=work merged

After we run this, we should have a directory like this.

michal@michal-lg:/tmp/container-1$ ls
alpine-minirootfs-3.20.3-x86_64.tar.gz  lower  merged  upper  work

The container itself will use /tmp/container-1/merged as the root of its filesystem:

michal@michal-lg:/tmp/container-1/merged$ ls
bin  etc   lib    mnt  proc  run   srv  tmp  var
dev  home  media  opt  root  sbin  sys  usr

Control groups

Let’s restrict the resource consumption of this container to, say, 100m CPU and 500 MiB.

Setting up cgroups is super easy:

# make a new cgroup slice and a child cgroup for our container
sudo mkdir -p /sys/fs/cgroup/toydocker.slice/container-1

cd /sys/fs/cgroup/toydocker.slice/

# enable modifying cpu and memory for the child cgroup
sudo -- sh -c 'echo "+memory +cpu" > cgroup.subtree_control'

cd container-1

# set max cpu usage to 10%
sudo -- sh -c 'echo "10000 100000" > cpu.max'

# set memory limit to 500 MiB
sudo -- sh -c 'echo "500M" > memory.max'

# Disable swap
sudo -- sh -c 'echo "0" > memory.swap.max'

The cpu.max syntax is a bit unusual, but it means that out of 100 000 time units, this cgroup can consume 10 000 of those units. If we instead wanted to limit the cgroup to 2 CPUs, it would be 200 000 out of 100 000.

Interestingly, the cpu.max rule doesn’t restrict the process to use a single physical core. So on a 4 core machine, it’s fine if a process uses 2500 time units on each of cores 0, 1, 2, 3, since the total is 10 000. For limiting the number of physical cores to use, cpusets can be used.

We can see that when we created the cgroup, default rules were automatically set up.

michal@michal-lg:/sys/fs/cgroup/toydocker.slice/container-1$ ls
cgroup.controllers      cpu.pressure         memory.numa_stat
cgroup.events           cpu.stat             memory.oom.group
cgroup.freeze           cpu.stat.local       memory.peak
cgroup.kill             cpu.uclamp.max       memory.pressure
cgroup.max.depth        cpu.uclamp.min       memory.reclaim
cgroup.max.descendants  cpu.weight           memory.stat
cgroup.pressure         cpu.weight.nice      memory.swap.current
cgroup.procs            io.pressure          memory.swap.events
cgroup.stat             memory.current       memory.swap.high
cgroup.subtree_control  memory.events        memory.swap.max
cgroup.threads          memory.events.local  memory.swap.peak
cgroup.type             memory.high          memory.zswap.current
cpu.idle                memory.low           memory.zswap.max
cpu.max                 memory.max           memory.zswap.writeback
cpu.max.burst           memory.min

Let’s check that the ones we modified took effect:

michal@michal-lg:/sys/fs/cgroup/toydocker.slice/container-1$ cat cpu.max
10000 100000
michal@michal-lg:/sys/fs/cgroup/toydocker.slice/container-1$ cat memory.max
524288000
michal@michal-lg:/sys/fs/cgroup/toydocker.slice/container-1$ cat memory.swap.max
0

Looks good. Next, let’s see how we can put a process into the cgroup and further isolate it via namespaces.

Namespaces

Let’s first answer the why of namespaces, then we can see how they are used.

If cgroups are the main mechanism for restricting resource usage, namespaces are the main mechanism for isolating resources themselves.

Let’s take filesystem mounts as an example. When we mount a new filesystem in the host, it’s visible to all processes. We need to be aware of what other mounts exist on a system to avoid clashes. With mount namespace, each process can make filesystem changes as it wishes without affecting any other process outside of this namespace.

The same idea extends to other resources: networking, inter-process communication, process ids, users, etc.

With motivation out of the way, let’s see it in action.

# enter interactive root
sudo -i

# Add current process to cgroup
echo $$ > /sys/fs/cgroup/toydocker.slice/container-1/cgroup.procs

# Create new namespaces
unshare \
    --uts \
    --pid \
    --mount \
    --mount-proc \
    --net \
    --ipc \
    --cgroup \
    --fork \
    /bin/bash

This piece of code is a little arcane, mostly because I want to keep everything in a single terminal.

First, we enter an interactive root shell. This is because I need to run the next two commands from the same shell and both with root privileges.

michal@michal-lg:~$ # enter interactive root
sudo -i
[sudo] password for michal: 
root@michal-lg:~# 

The second command adds the current shell process to the cgroup we created earlier. Any children of this process will also automatically be added to the cgroup.

root@michal-lg:~# echo $$
28156
root@michal-lg:~# echo $$ > /sys/fs/cgroup/toydocker.slice/container-1/cgroup.procs

When we do this, the current shell is part of the cgroup and all the CPU and memory restrictions we set up earlier already apply.

Next, when we create the namespaces, it forks the current process and runs a bash shell. You can learn more about the unshare command in man pages.

root@michal-lg:~# unshare \
    --uts \
    --pid \
    --mount \
    --mount-proc \
    --net \
    --ipc \
    --cgroup \
    --fork \
    /bin/bash
root@michal-lg:~# 

This looks unremarkable, but we have essentially created a container through cgroup and namespace isolation. Let’s test that the UTS namespace is working correctly by changing the hostname and seeing that it doesn’t change on the host.

Container terminal:

root@michal-lg:~# hostname
michal-lg
root@michal-lg:~# hostname mycontainer
root@michal-lg:~# hostname
mycontainer
root@michal-lg:~# 

Host terminal:

michal@michal-lg:~$ hostname
michal-lg

Since we also used `pid` namespace, `/bin/bash` should now have process id of 1. Let’s verify from the container:

root@michal-lg:~# ps
    PID TTY          TIME CMD
      1 pts/1    00:00:00 bash
     32 pts/1    00:00:00 ps

And let’s see what the real process id is from the host’s perspective.

michal@michal-lg:~$ ps -ef | grep -i /bin/bash
root        8952    8932  0 16:10 pts/1    00:00:00 unshare --uts --pid --mount --mount-proc --net --ipc --cgroup --fork /bin/bash
root        8953    8952  0 16:10 pts/1    00:00:00 /bin/bash

There are some post-processing steps that a container runtime would do at this point before launching user’s application. Let’s go through those next.

Container-side setup

First, we isolate the container from the host filesystem by changing the root using the pivot_root command.

pivot_root is a safer equivalent to chroot /tmp/container-1/merged used by container runtimes to avoid breakout exploits. Security is not my expertise, so I’ll link to this article explaining how these exploits work and how pivot_root prevents them.

root@michal-lg:~# cd /tmp/container-1/merged
mount --make-rprivate /
mkdir old_root
pivot_root . old_root
umount -l /old_root
rm -rf /old_root
root@michal-lg:/tmp/container-1/merged# 

Making the root private prevents the container from affecting host’s mount table. This could again be used for exploits.

In my terminal, I need to run “cd ..” to refresh state after we deleted the old root. Since we removed old root, PATH variables no longer resolve correctly.

But since we are now in the `/tmp/container-1/merged` directory and this filesystem is based on Alpine minirootfs, we have basic utilities in the `bin` directory.

root@michal-lg:/tmp/container-1/merged# cd ..
root@michal-lg:/# ls
bash: /usr/bin/ls: No such file or directory
root@michal-lg:/# /bin/ls
bin    dev    etc    home   lib    media  mnt    opt    proc   root   run    sbin   srv    sys    tmp    usr    var

Let’s also setup basic devices we’ll need later and mount useful virtual filesystems:

mknod -m 666 dev/null c 1 3
mknod -m 666 dev/zero c 1 5
mknod -m 666 dev/tty c 5 0

/bin/mkdir -p dev/{pts,shm}
/bin/mount -t devpts devpts dev/pts
/bin/mount -t tmpfs tmpfs dev/shm
/bin/mount -t sysfs sysfs sys/
/bin/mount -t tmpfs tmpfs run/
/bin/mount -t proc proc proc/

For instance, if we didn’t mount `proc`, we wouldn’t have access to process information and running commands that depend on reading process info would fail:

root@michal-lg:/# top
top: no process info in /proc

After the mount, things work correctly again.

Mem: 7560280K used, 8661696K free, 161756K shrd, 135464K buff, 2364264K cached
CPU:   0% usr   0% sys   0% nic  98% idle   0% io   0% irq   0% sirq
Load average: 0.30 0.38 0.37 1/1233 64
  PID  PPID USER     STAT   VSZ %VSZ CPU %CPU COMMAND
    1     0 root     S    12896   0%   6   0% /bin/bash
   64     1 root     R     1624   0%   9   0% top

At this point, we could configure networking, export env variables, etc. For our minimal purposes, we are done and it’s time to launch user’s application!

Let’s suppose the user wanted to run a simple interactive shell. We can launch it like this:

exec /bin/busybox sh

I use busybox since it works as a minimal init script and it ships in alpine minirootfs. Using exec replaces the old process with the new one since we don’t need to keep the shell around.

root@michal-lg:/# exec /bin/busybox sh
/ # ls
bin    dev    etc    home   lib    media  mnt    opt    proc   root   run    sbin   srv    sys    tmp    usr    var
/ # 

Right now, we are roughly where we’d be if we ran the following docker command:

michal@michal-lg:~$ docker run -it --cpus="0.1" --memory="512M" --memory-swap=0 --entrypoint /bin/sh --rm alpine
/ # ls
bin    dev    etc    home   lib    media  mnt    opt    proc   root   run    sbin   srv    sys    tmp    usr    var
/ # 

Using the container

As a final step, let’s try to see that the cgroup limits we set earlier are actually working.

First, I’ll run a CPU intensive task that should use 100% of a single CPU core

/ # while true; do true; done

and open a terminal from host to see the real CPU utilization for this process. I find the process id and check consumption with top:

michal@michal-lg:~$ ps -ef | grep -i busybox
root        8953    8952  0 16:10 pts/1    00:00:07 /bin/busybox sh

And then use `top` to verify that CPU usage doesn’t exceed 10%.

michal@michal-lg:~$ top -p 8953
PID USER      PR  NI    VIRT    RES    SHR S  %CPU  %MEM     TIME+ COMMAND                                                                                                                           
   8953 root      20   0    1696   1024    896 R  10.0   0.0   0:15.25 busybox    

Similarly, for memory, I run tail to keep reading from /dev/zero. tail reads into an in-memory buffer that will shortly exceed our 500 MiB memory limit at which point the cgroup memory controller will kill the process.

/ # tail /dev/zero
Killed

We can now exit the container, and cleanup by unmounting the root directory /tmp/container-1/merged

michal@michal-lg:/tmp/container-1$ sudo umount merged

And that’s it! We’ve created a container from scratch in a terminal.

Conclusion

The main takeaway should be that containers aren’t magic. They are not virtual machines. They are an awesome feature baked into the Linux kernel for isolating processes. They achieve this isolation through cgroups and namespaces.

You can see the full-list of commands on my github.

I hope you learned something new! If you did, consider subscribing! I’m also always happy to connect with readers on LinkedIn and BlueSky

If you enjoyed the post, chances are you’ll enjoy my other writing too! All my posts are backed by a substantial deep-dive into the given problem space.

Perhaps check out my deep-dive into SQLite storage format, my implementation of MapReduce from scratch, or my introduction to CUDA programming?

I’ll start by starting a shell in a Docker container using Alpine image.

michal@michal-lg:~$ docker run -it --entrypoint /bin/sh --rm --name "alpine-container" alpine
/ # ls
bin    dev    etc    home   lib    media  mnt    opt    proc   root   run    sbin   srv    sys    tmp    usr    var
/ # hostname
a7cbf0aea1ad
/ # cd home && ls

We are in a separate filesystem and it’s pretty empty. Let’s make a file.

/ # echo -e "Hello there\nGeneral Kenobi" > /home/hello_there.txt
/ # cat /home/hello_there.txt 
Hello there
General Kenobi

What do you think, can we access this file from the host?
…

…

…
We can - let’s find it!

Docker stores everything under /var/lib/docker so we can start looking from there from a second terminal:

root@michal-lg:/var/lib/docker# find -name hello_there.txt
./overlay2/1557145fe40a1595d090eeafa72c39a7b54cca4791ae9e3ffafabff06466125c/diff/home/hello_there.txt
./overlay2/1557145fe40a1595d090eeafa72c39a7b54cca4791ae9e3ffafabff06466125c/merged/home/hello_there.txt
root@michal-lg:/var/lib/docker# 

Curiously, we found the file twice in two different directories. Let’s see what else is in the diff and merged directories.

root@michal-lg:/var/lib/docker/.../diff# ls
home  root
root@michal-lg:/var/lib/docker/.../merged# ls
bin  dev  etc  home  lib  media  mnt  opt  proc  root  run  sbin  srv  sys  tmp  usr  var

The diff directory only contains an empty root directory and a home directory with the file we created. Contents of merged exactly match the container’s filesystem.

Overlayfs

Docker uses overlayfs filesystem. Overlayfs lets us combine 2 file file trees, “lower” and “upper”, into a combined view “merged”. Docker calls the “upper” file tree “diff”, which is perhaps more fitting and I’ll be referring to it as such.

Usage of union filesystems like overlayfs comes from an interesting observation: Often we want to run multiple containers on a single host. Chances are, these containers might share the lower layer - be it Alpine, Ubuntu, or a more specialized one like golang.

By making the lower layer read-only, multiple containers can share it. Changes are only written to the upper layer. Let’s illustrate what happened when I earlier created hello_there.txt. I’m not expanding all folders to avoid clutter.

When we then create hello_there.txt in /home, it was written to diff and overlayfs constructed a combined view in merged.

What if we modify something in the lower layer? Let’s rename /bin/echo to /bin/echo.old.

/ # cd bin
/bin # mv echo echo.old
/bin # ls
...
chgrp          echo.old       gzip           ln             mount          printenv       setserial      umount
...

As mentioned, the lower layer is read-only so the only place where the file is actually modified is in diff.

root@michal-lg:/var/lib/docker/overlay2/.../diff/bin# ls -l
total 0
c--------- 1 root root 0, 0 Nov 11 23:06 echo
lrwxrwxrwx 1 root root   12 Sep  6 13:34 echo.old -> /bin/busybox

There are two files! One for the no longer existing echo and one for echo.old. Overlayfs uses special whiteout files to deal with deletion of a file. When overlayfs sees this file, it knows not to include it in the merged view.

The second file is much less interesting, it’s the renamed echo, which turns out was just a symbolic link to busybox.

Creating container filesystem

Finally, let’s see how Docker uses Overlayfs under the hood to create a new filesystem.

First, let’s create temporary directories in /tmp/ where we’ll setup a container filesystem manually.

michal@michal-lg:/tmp$ mkdir -p /tmp/container-demo/{diff,merged,work}
michal@michal-lg:/tmp$ ls container-demo/
diff  merged  work

We’ve seen diff and merged before. work is used by overlayfs as a scratchpad for and we don’t need to care about it.

Next, download Alpine minirootfs for your CPU architecture and extract it to  /tmp/container-demo/. I’ll rename the extracted folder to “alpine”.

michal@michal-lg:/tmp/container-demo$ ls
alpine  diff  merged  work

Now that everything is setup, we can mount the overlayfs filesystem.

michal@michal-lg:/tmp/container-demo$ sudo mount -t overlay  overlay  -o lowerdir=alpine,upperdir=diff,workdir=work  merged

Now if we list contents of merged, we’ll see the Alpine file system:

michal@michal-lg:/tmp/container-demo/merged$ ls
bin  etc   lib    mnt  proc  run   srv  tmp  var
dev  home  media  opt  root  sbin  sys  usr

And if we create a file in there, it will be written to the diff folder.

michal@michal-lg:/tmp/container-demo/merged$ echo hello > hello.txt
michal@michal-lg:/tmp/container-demo/merged$ ls
bin  etc        home  media  opt   root  sbin  sys  usr
dev  hello.txt  lib   mnt    proc  run   srv   tmp  var
michal@michal-lg:/tmp/container-demo/merged$ ls ../diff/
hello.txt

As a final point, we can create a new shell process and set merged as its root directory. This is how Linux containers can only see their own filesystem. Any process spawned by this shell will inherit the root so it will also be contained to this filesystem.

michal@michal-lg:/tmp/container-demo$ sudo chroot merged /bin/sh
/ # ls
bin        hello.txt  media      proc       sbin       tmp
dev        home       mnt        root       srv        usr
etc        lib        opt        run        sys        var
/ # cd ..
/ # 

A full implementation would take advantage of namespaces for additional isolation. You can learn more about those in Linux man-pages.

Conclusion

Containers are a black box for vast majority of engineers. You should now have a solid understanding of how things work under the hood!

If you are interested in a more complete implementation of Docker from scratch, consider checkout my Golang Docker clone that’s around 200 lines of code.

I hope you learned something new! If you did, consider subscribing and/or following me on LinkedIn.

You might also enjoy some of my other posts linked below. Perhaps my deep-dive into SQLite storage format, my implementation of MapReduce from scratch, or my series building an ML inference engine from scratch?

Baseline

In this post, we’ll focus on optimizing throughput: the number of inference requests processed per second. Tradeoffs between throughput and latency are pretty common. In ML applications, we often care about maximizing throughput while staying within overall latency tolerations.

First things first, we need to set up the benchmark. I’ll be using the MNIST neural network from part 1 of this series. We’ll time how long it takes to process 10000 images and calculate the throughput.

CPU baseline: 430μs per inference or ~2300 inferences per second.

As for hardware, I am using AMD Ryzen 5 5600H and an 80W mobile RTX 3060 with 6GB of VRAM.

Adding CUDA execution provider

Architecture

My engine supports only these operations: GEMM, ReLU, Add, and flatten. Our goal is to make these also executable on a GPU. I went through multiple iterations, starting with a basic approach to validate ideas

if (use_gpu) {
    gemm_cuda(...);
} else {
    gemm_cpu(...);
}

and later introducing abstractions to make it easier to work with. Eventually, I settled on an approach heavily inspired by ONNX runtime.

Shared inference functionality is encapsulated in an inference_session object. It loads the model, sets its input, and iterates over the topologically sorted computational graph. Nodes of the graph are executed via plugins called execution providers.

Execution providers implement the mathematical operations such as matmul or ReLU needed to execute each node. This way each operation’s implementation can be optimized for the hardware the provider targets.

All operations are then implemented in each provider.

Configuration is done via a yaml config file, inspired by Nvidia’s Triton Inference Server.

To see the full implementation, feel free to check out the GitHub repository.

Gotta go slow to go fast

My first attempt at adding CUDA-based matrix multiplication looked something like this. This is pretty common in CUDA samples seen online but has a huge issue!

So, how fast is this? Before we had a throughput of 2300 inferences per second on CPU. *Drum roll* 470 inferences per second: ~4 times slower.

What went wrong?

For every call, this needs to allocate memory on the GPU, transfer data from CPU to GPU, run the kernel, transfer data back to CPU, and finally free GPU memory. Those calls have significant overheads, especially for cheaper operations like ReLU and add.

Until now, I was using a batch size of 1 for inference. We can amortize some of the mentioned overheads by using larger batch sizes. The chart below shows how the throughput changes as we increase the batch size from 1 to 128 for the CPU and the naive CUDA provider.

The CPU provider experiences no scaling with increased batch sizes. The CUDA provider’s throughput is 4x slower for a batch size of 1 and grows to 11x bigger at a batch size of 128.

This is encouraging! Let’s revisit the problem with data transfer overhead and see if we can do anything about performance at smaller batch sizes.

Optimizing memory transfers

We know that all model’s weights are fixed. Once loaded from the disk, we can move them to GPU and keep them there. For inputs and outputs, we need to tolerate some memory transfers at the start and end of each inference call.

What we cannot optimize right now is the creation of intermediate tensors during the inference loop. Those are used to store outputs when we cannot reuse one of the input tensors as the output directly. Allocating new tensors on the GPU involves a call to cudaMalloc and cudaFree for each allocation and destruction respectively and those are relatively expensive as we’ve seen. We’ll see what we can do about this in the next section!

After a few changes to my tensor class and rewriting my CUDA operations to assume data is already on the GPU, here’s our new implementation benchmarked against the previous versions.

We got between 12x to 4x throughput increase by keeping tensors on the GPU as much as possible. The benefits are more pronounced at smaller batch sizes where the overhead from CUDA calls is larger. Nice!

Optimizing intermediate tensor allocations

I noticed intermediate tensor creation and destruction was an issue by profiling the engine with perf and seeing that a lot of the execution time is spent in cudaMalloc and cudaFree calls.

If making multiple cudaMalloc and cudaFree calls is expensive, what if we pre-allocate a large chunk of GPU memory and reuse it? Newly created tensors can get a slice of memory from this memory pool and release it back to the pool when they are destroyed. This way we can completely avoid expensive cudaMalloc and cudaFree calls in the inference loop.

To make this work, tensors accept an optional allocator argument in their constructor. The allocator can internally use a memory pool without the tensor needing to know about this. Each execution provider can implement its allocator or use the default CPU allocator - a thin wrapper around malloc and free. 

This allows us to avoid all calls to cudaMalloc and cudaFree during inference as long as we pre-allocate enough memory. Let’s see how it affects performance.

Another 2x to 1.2x increase in throughput! The optimization again has more impact on runs with smaller batch sizes.

Results

This was the last optimization I’ve written, so let’s see what the overall improvement is.

+------------+------------------+----------+
|            |     Throughput   |          |
+------------+------+-----------+ Increase |
| batch size | CPU  | CUDA_pool |          |
+============+======+===========+==========+
|          1 | 2315 |     12048 |        5 |
+------------+------+-----------+----------+
|          2 | 2345 |     23529 |       10 |
+------------+------+-----------+----------+
|          4 | 2306 |     36900 |       16 |
+------------+------+-----------+----------+
|          8 | 2360 |     45454 |       19 |
+------------+------+-----------+----------+
|         16 | 2309 |     60975 |       26 |
+------------+------+-----------+----------+
|         32 | 2328 |     90909 |       39 |
+------------+------+-----------+----------+
|         64 | 2338 |    120481 |       52 |
+------------+------+-----------+----------+
|        128 | 2329 |    144927 |       62 |
+------------+------+-----------+----------+

Pretty nice! Even at a batch size of 1, we are at 5x the CPU performance. Larger batch sizes yield higher throughput up to 60x!

Using batch sizes beyond 128 doesn’t seem to provide much benefit.

If you are curious, the current main bottleneck is the moving of results back to the CPU memory. In a future post, I’d like to explore the impact of pageable and pinned memory on this, possibly using async cuda operations where viable, and maybe adding graph optimizations to fuse kernels such as matmul followed by ReLU.

Stay tuned!

Thanks for reading! Researching and writing these articles takes a lot of time and effort. To ensure you don’t miss the next one, consider subscribing or following me on LinkedIn.

Subscribe now

If you enjoyed this article, you might enjoy some of my other work:

Baseline benchmarking

Improving performance is an empirical endeavor. I’ll be using a simple benchmark where the engine executes 500 sequential inference requests.

I’ll be compiling the binary in debug mode throughout this post so that the profiler has access to all debug symbols. In general, benchmarking should probably be done in release mode, but the extra readability will serve us well here.

We can time the execution with time.

time /home/michal/code/inference_engine/build/src/engine_exe /home/michal/code/inference_engine/models/mnist_ffn_complex.onnx 

real    0m15.194s
user    0m14.500s
sys     0m0.694s

Profiling

I’ll be using perf - a powerful profiler baked into the Linux kernel. To visualize perf’s report, I’ll use an OSS perf GUI Hotspot. This will give us access to nice flame charts.

I’m looking for major bottlenecks so a sampling analysis will work just fine for me.

perf record -o /home/michal/perf.data --call-graph dwarf --aio --sample-cpu /home/michal/code/inference_engine/build/src/engine_exe /home/michal/code/inference_engine/models/mnist_ffn_complex.onnx

The flame chart shows us where code spends the most CPU cycles. The core function in my engine is this InferenceEngine::infer call that I’ve zoomed on. It iterates over the computational graph, loads inputs, and evaluates every node.

Looking closer at the flame chart shows something concerning. Less than 40% of cycles are spent evaluating nodes. InferenceEngine::prepareNodeInputs is taking up a significant amount of time. There are many Tensor memory allocation calls and copy calls.

When I zoom in on the InferenceEngine::evaluateNode call, the story gets worse.

It only spends around 20% of its execution time doing useful work. The rest is spent on copying and destroying Tensors.

Let’s see if we can optimize this.

Optimizing

Let’s look at the InferenceEngine::prepareNodeInputs method.

Here we take elements from a map and push them into a vector. Looks innocent enough but it creates a copy of each Tensor. Let’s rework this to instead return a vector of naked pointers to Tensors. While doing so, we can also slightly optimize the map access and pre-allocate a vector of appropriate size.

In the flame chart, we saw a bunch of Tensor copies, so I’ve also gone ahead and used pointers or references where applicable in the rest of the InferenceEngine::infer method to minimize the number of unnecessary copies.

Results

To see if our previous changes made a material difference we need to benchmark again:

real    0m1.053s
user    0m1.040s
sys     0m0.013s

Wow, 15x faster just like that.

Let’s also re-run the profiler to see what proportion of time is spent on useful operations now.

As we can see, now InferenceEngine::evaluateNode constitutes the vast majority of the infer call. Most time is spent on the most complex operation my engine currently supports - generalized matrix multiplication. That’s great!

To overcome the performance limitations imposed by my current GEMM implementation, I plan to explore either offloading the computation to a GPU using CUDA or distributing the workload across multiple CPU cores.

With this in mind, I think I’ll hold off on graph optimizations as they likely won’t provide a significant improvement compared to other optimizations.

I hope you enjoyed this lighter post on performance programming! If you did, consider subscribing and/or following me on LinkedIn!

Subscribe now

You might also enjoy some of my other posts. Maybe check out one of the ones linked below?

Training a model

Before we can serve a model, we need to train one. We’ll be using the model illustrated below.

This model has some nice features: it’s easy to train, has a non-trivial topology, and only requires 4 operations: Flatten, Gemm, ReLU, and Add. Gemm stands for generalized matrix multiplication. I think the others are self-explanatory.

We want to save the trained model in ONNX format. ONNX is a standard format for saving models for interoperability between different ML frameworks. Since I probably won’t be adding support for other model formats, this is a solid default choice.

You can see all the code related to training the model on my github.

(venv) michal@michal-lg:~/code/mnist$ python nn_complex.py 
Net(
  (fc1): Linear(in_features=784, out_features=512, bias=True)
  (fc2_left): Linear(in_features=512, out_features=200, bias=True)
  (fc2_left2): Linear(in_features=200, out_features=100, bias=True)
  (fc2_right): Linear(in_features=512, out_features=100, bias=True)
  (fc3): Linear(in_features=100, out_features=10, bias=True)
)
Epoch 1 - Test loss: 0.0006, Accuracy: 84.77%
Epoch 2 - Test loss: 0.0004, Accuracy: 89.66%
Epoch 3 - Test loss: 0.0003, Accuracy: 90.73%
Epoch 4 - Test loss: 0.0003, Accuracy: 91.91%
Epoch 5 - Test loss: 0.0003, Accuracy: 92.58%
Epoch 6 - Test loss: 0.0002, Accuracy: 93.31%
Epoch 7 - Test loss: 0.0002, Accuracy: 93.59%
Epoch 8 - Test loss: 0.0002, Accuracy: 94.02%
Epoch 9 - Test loss: 0.0002, Accuracy: 94.23%
Epoch 10 - Test loss: 0.0002, Accuracy: 94.65%
Model saved as mnist_ffn_complex.onnx

With the model trained, let’s learn more about inference engines.

Why inference engines matter

Before designing the engine, let’s discuss why we even want one in the first place. Couldn’t we just reuse the same ML training framework we used to train the model?

With LLMs going mainstream, an interesting observation was made - over the lifetime of a model, serving can be more expensive than training. So it makes sense to have specialized tools optimized for inference specifically.

Inference servers are software for managing deployment, lifetime, and serving-related optimizations of already trained models. Popular inference servers include Nvidia’s Triton Inference Server or Google’s TensorFlow Serving.

Inference servers balance throughput and latency. Throughput is often optimized through dynamic batching - waiting for inference requests to accumulate before handing them off to the inference engine. This improves hardware utilization at the cost of increased latency for some requests. 

The inference engine, the subject of our discussion, then efficiently executes the model with provided inputs. To achieve high performance, inference engines employ a range of optimizations:

• Hardware acceleration

• Efficient memory management

• Graph optimizations

• Quantization to reduce numeric precision while maintaining accuracy

I want to implement graph optimizations and GPU acceleration in follow-up posts. Consider subscribing to get an email when I publish the next post.

For now, let’s stick with CPU inference.

Inference engine from scratch

Let’s outline the steps that our engine will do:

• Load the model

• Construct a graph representation of the model

• Topologically sort nodes

• Run inference with user inputs

Loading the model

Luckily for us, ONNX models are saved in Protobuf format. This means we can download the onnx-ml.proto and generate a client library for interacting with ONNX files. This will also be our only external dependency - sticking with just the standard lib from now on.

Once the model is loaded, we can extract the weights into a Tensor object. Tensor is a thin wrapper around std::vector<T> where elements are stored in row-major order.

Graph construction

In this part, we iterate over all nodes in the ONNX model, extract each into a minimal Node representation, and store them in an adjacency list in a Graph object.

Nodes simply define the operation and list the names of input and output tensors.

When we construct the graph, we add nodes 1-by-1. When a new node is added, we check if any existing nodes are the parents or children of the new node by comparing input and output tensor names.

Now that we have a graph, we have to figure out how to execute it.

Topological sorting

I’ll assume that our graphs are non-cyclic, i.e. we are dealing with DAGs. Still, we need to be careful to execute nodes in such an order that all intermediate results are ready when the node is executed.

If we didn’t pay special attention to this, we would likely run into a scenario like the one illustrated in the animation below. There we tried to execute “Add” before the results of the left branch were ready.

Instead, we want something closer to this.

If you are already thinking if we could execute those two branches in parallel, then you are on the right track! We won’t go into that in this post as I haven’t implemented that logic yet, but it’s something I’d like to explore down the line.

So how do we get this order? We can use trusty topological sort which can be concisely implemented with depth-first search. Since we assume the model is static, it’s sufficient to compute the order once when the model first loads.

For those interested in seeing this implemented, you can read it here.

Inference

We’ve loaded the model, extracted its weights, constructed a graph, and sorted the graph’s node. We are all set for inference. I’ll skip input loading as it’s not particularly interesting. For our MNIST example, every input is a Tensor<uint8>(28, 28) representing a black-and-white image.

The infer() call iterates over topologically sorted nodes, for each node, it does the following:

• Read input names used by the node.

• Read inputs from a Tensor store into a vector of inputs.

• Based on the node’s operation type, read additional input information, things as whether a matrix is transposed in Gemm.

• Pass inputs to a corresponding operator function.

• Save output to Tensor store or print if it’s the final result.

Here’s part of the method to get the main idea. You can see the full source code on github. Each case simply prepares inputs and passes them to the operator function.

The operator functions are implemented in operators.cpp. Since I decided to only support 4 operations, it’s not too bad to implement them from scratch. We are giving up some performance here, but I’d like to explore C++ profiling tooling in a follow-up post anyway. Especially around things like memory access and cache locality.

And that’s pretty much it. Let’s run the inference engine on this image of number 7.

michal@michal-lg:~/code/inference_engine$ /home/michal/code/inference_engine/build/src/engine_exe /home/michal/code/inference_engine/models/mnist_ffn_complex.onnx /home/michal/code/inference_engine/inputs/image_0.ubyte
Out: Tensor((1, 10)[[407.129, -1327.89, 827.717, 1137.59, -1497.12, -73.3868, -2284.66, 2266.74, 1.9645, 475.585]])

To determine the model’s prediction, we take the argmax of the output. Here the 7th (0 indexed) output is the largest, so the model correctly predicts the image is a 7!

Not bad for ~2000 lines of C++.

Thanks for reading! We’ve covered how inference engines work. Now that we have a minimal engine, I want to extend it. Let me know in the comments which improvements you’d like to see!

I hope you enjoyed this deep dive into inference engines! If you did, consider subscribing and/or following me on LinkedIn.

You might also enjoy some of my other posts. Maybe my deep dive into SQLite storage format or implementation of MapReduce from scratch?

C CUDA basics

C CUDA is Nvidia’s extension of ANSI C. For the most part, it is the same as C with some added syntax and built-in functions. C CUDA gives us control over what parts of our code are executed on the CPU and the GPU. We call code executed on the CPU host code and GPU code device code. Procedures that run on the GPU are for historical reasons called kernels.

Instead of focusing on CUDA itself, let’s write a simple program that blurs an image. I’ll try to fill in the details as needed.

Blurring images

We want to write code to blur an image on a GPU.

Here’s roughly what our code needs to do:

• Load the image in the host code

• Allocate memory on the GPU

• Copy over the input image to the GPU

• Blur the image with a kernel

• Copy over the output image to the CPU

• Save the output image to the disk

First, we need to know how an image is represented in memory and how to blur it. An RGB image is usually thought of as a 3-dimensional matrix of shape (channels, height, width). In memory, it’s usually represented as a flat array in row-major order. Our GPU code will assume this format.

To access the (n, row, col) pixel in a 3-channel image, we can use the following expression.

i = (row*WIDTH + column)*3 + n

To blur an image, we calculate the value of each pixel as the average of surrounding pixels and write the result into the output image.

Can we parallelize this? Of course! Each output pixel only depends on the input image, but has no dependencies on other outputs. If we had a processor with width*height cores, we could process every output pixel in parallel. Turns out, that’s pretty much what GPUs are!

Writing a kernel

Let’s finally write our kernel. It will closely follow the 2D example above, but generalize it to n-channel images. 

A kernel in execution is called a thread. Each thread will compute the RGB channels for a single pixel in the image.

To tell each thread which pixel to compute, CUDA automatically injects variables blockIdx, blockDim, and threadIdx into the kernel. We use these to determine which pixel a given thread should process.

Once we know which pixel we are processing, we iterate over the neighboring pixels and accumulate their red, green, and blue values in pixVarR, pixVarG, and pixVarB. We also count the number of pixels we’ve iterated over, to handle cases where the blur-radius reaches beyond the edges of the image.

Note that the coordinate calculation might feel unnatural since the image is flattened as discussed earlier.

You might notice the special __global__ identifier before the kernel name. This is how we specify that a procedure is a kernel and should be compiled to run on the GPU. It’s also how the c cuda compiler (NVCC) knows to inject the blockIdx, blockDim, and threadIdx variables.

Now that we have the kernel, let’s briefly write the main function to set things up and run it. The main function closely follows the setup steps outlined earlier. The cuda-prefixed functions are automatically included by the NVCC compiler. These mostly copy built-in C functions to provide similar functionality but for GPUs.

The interesting part is when we call the kernel. 

We designed our kernel to process all channels for a given pixel with a single thread.

To make it concrete, our lenna.png is of shape 512 * 512, so we need that many threads to process the whole image. My GPU, however, only has 1920 CUDA cores. That’s fine - GPUs are much faster at context switching than CPUs, so having more threads than physical cores is desirable to maximize throughput.

To do this, we use the <<<dimGrid, dimBlock>>> syntax. Each argument is a struct with 3 fields {x, y, z}. The configuration generally follows the shape of the input data. Since each thread corresponds to a single pixel, a natural division is using 2 dimensions. We can group threads into 16x16 blocks, meaning each block will process a 16x16 patch of the image. Since we are not using multiple threads per the z-axis, we leave it at 1.

The grid dimension tells us how many blocks per dimension to create. Since we need to cover the whole image using 16x16 patches, we need width/16 blocks in the x direction and height/16 blocks in the y direction. In case our image dimensions don’t divide evenly by 16, we round up. 

This rounding means that we might need to spawn some extra blocks where only some threads are utilized. To make sure these unused threads behave correctly, we added the conditional check in our kernel. Only threads that have a corresponding pixel will do some work!

Oof, that’s a lot of low-level details!

So, why 16x16 blocks? In our case, it is pretty arbitrary. We could’ve used 8x8 blocks or 32x32. There’s a hard limit on the number of threads per block, which is usually 1024 on recent cards. As far as I understand, properly organizing threads into blocks and threads can improve performance thanks to memory locality.

Finally, let’s compile and run our code. I’m using BLUR_SIZE=31 for an extra blurry effect.

We can compile this with NVCC to yield the blurred image.

You might wonder if we can run code directly on the GPU without any intervention of the CPU. As far as I know, that’s not possible. Our executable makes calls to cuda runtime API, which in turn communicates with the GPU drivers. However, it’s possible to chain kernels to keep as much of the computation on the GPU without CPU intervention.

If you would like to learn more about this area, consider checking out the book Programming Massively Parallel Processors and the official CUDA C++ Programming guide from Nvidia.

Thanks for reading! If you enjoyed this write-up, you might enjoy my previous one where I explain how MapReduce works by building it from scratch!

Researching and writing these articles takes a lot of time and effort. To ensure you don’t miss the next one, consider subscribing or following me on LinkedIn.

This will be a long article: we’ll understand the need for distributed computing, rediscover why MapReduce is a natural way to model many problems, build our own version, understand how individual parts fit together, and solve a real problem with it!

Motivating the problem

Suppose we want to count the word frequencies in a massive dataset. We tried to process it on a single machine but found it would take over a month. What can we do?

The first thought should be to get a faster machine, but it either might not exist or be too expensive. Instead, let’s see how we can distribute the problem across N commodity machines. We want an easy way to look up the frequency of any word with a simple query, i.e. grep over a file.

Let’s start by splitting the dataset into N partitions and computing word frequencies for each subset using a different machine. This already gives us Nx speedup minus some fixed overhead. To combine the results, we can add a final machine that takes these N partial result files and sums the frequencies of corresponding words.

This is not too far from the core idea behind MapReduce. The above process has two main steps: first, we map words to their frequency in a subset of input data and then reduce the intermediate results to obtain the final answer. 

Note that this is pretty general, imagine we have a large dataset of photos that we want to classify: we could do the image classification task as a map operation and then group images with the same class in the reduce phase.

Another observation is that the map part is usually the more expensive phase of the two, so generally, we will have more mappers than reducers.

Having hopefully convinced you that MapReduce is a reasonable idea, let’s see how the MapReduce paper solves the word frequency problem.

Using MapReduce

Below, I copy-pasted the WordCounter MapReduce program from the paper. Let’s see how it works. Later, when we implement our version, we’ll aim to keep the usage semantics the same.

#include "mapreduce/mapreduce.h"

// User’s map function
class WordCounter : public Mapper
{
public:
    virtual void Map(const MapInput &input)
    {
        const string &text = input.value();
        const int n = text.size();
        for (int i = 0; i < n;)
        {
            // Skip past leading whitespace
            while ((i < n) && isspace(text[i]))
                i++;
            // Find word end
            int start = i;
            while ((i < n) && !isspace(text[i]))
                i++;
            if (start < i)
                Emit(text.substr(start, i - start), "1");
        }
    }
};
REGISTER_MAPPER(WordCounter);

// User’s reduce function
class Adder : public Reducer
{
    virtual void Reduce(ReduceInput *input)
    {
        // Iterate over all entries with the
        // same key and add the values
        int64 value = 0;
        while (!input->done())
        {
            value += StringToInt(input->value());
            input->NextValue();
        }
        // Emit sum for input->key()
        Emit(IntToString(value));
    }
};
REGISTER_REDUCER(Adder);

int main(int argc, char **argv)
{
    ParseCommandLineFlags(argc, argv);
    MapReduceSpecification spec;
    
    // Store list of input files into "spec"
    for (int i = 1; i < argc; i++)
    {
        MapReduceInput *input = spec.add_input();
        input->set_format("text");
        input->set_filepattern(argv[i]);
        input->set_mapper_class("WordCounter");
    }
    
    MapReduceOutput *out = spec.output();
    out->set_filebase("/gfs/test/freq");
    out->set_num_tasks(100);
    out->set_format("text");
    out->set_reducer_class("Adder");
    
    // Tuning parameters: use at most 2000
    // machines and 100 MB of memory per task
    spec.set_machines(2000);
    spec.set_map_megabytes(100);
    spec.set_reduce_megabytes(100);
    
    // Now run it
    MapReduceResult result;
    if (!MapReduce(spec, &result))
        abort();
    // Done: ’result’ structure contains info
    // about counters, time taken, number of
    // machines used, etc.
    return 0;
}

The user program has 3 parts: map function, reduce function, and configuration. Most of the heavy lifting is handled by the imported mapreduce library.

The map function splits the input text into words and emits key-value pairs. For instance, “the quick brown fox jumps over the lazy dog” is mapped to these key-value pairs: [“the”:1, “quick”:1, “brown”:1, “fox”:1, “jumps”: 1, “over”:1, “the”:1, “lazy”:1, “dog”:1]. Notice that the pair (“the”: 1) is emitted twice.

The reduce function operates per key. For instance, given the input [“the”:1, “the”:1”, “the”:1] it would emit (“the”:3).

The config handles input-output, formats, and the number of resources available for the MapReduce job.

In less than 100 lines of code, we can solve the word counting problem by utilizing 1000s of machines! That’s pretty neat and shows how powerful of an abstraction MapReduce is.

Gathering requirements

With the paper as our inspiration, let’s start outlining the requirements for our implementation:

• The user imports our MapReduce library, implements Map and Reduce functions, and provides configuration.

• When a user executes the binary, it is copied to multiple machines where it executes in either Mapper or Reducer mode.

• Those machines must have access to input data.

• Outputs of mappers must be available to reducers.

• The user can access the final results.

Infrastructure

When I started working on this, these requirements presented two main unknowns: how to distribute the binary to other machines and how to make input data available to them.

After some research, I settled on hosting my input data on a networked storage server - I chose NFS. We can mount the networked directory on every machine and allow the machines to read and write to it. It won’t be as performant or scalable as Google File System or HDFS, but it will do for our purposes.

To distribute the binary to multiple machines, execute it, and monitor its status, we could probably get quite far by copying the binary via SCP, SSHing into machines, executing the binary, and then waiting for it to finish. We could even implement health-pinging and progress-reporting.

However, after twiddling my thumbs for a while, I realized that this sounds awfully like cluster orchestration and decided to leverage Kubernetes to make my life easier. We can build the binary as a docker image and execute it as Kubernetes Jobs. Integrating NFS with Kubernetes is easy enough via Persistent Volumes and Persistent Volume Claims. Don’t worry, you don’t need to understand these concepts, but I mention them for completeness.

Since I’m developing this for educational purposes, I use minikube to spin up a local virtual cluster on my laptop. I run my NFS server as a docker container outside the cluster and connect it to the cluster via docker networking. The advantage of this setup is that it can be easily migrated to several physical or virtual machines.

With infrastructure sorted out, let’s start writing our MapReduce framework!

Using my MapReduce

First, we’ll explore how to solve the word-counting problem using my MapReduce implementation. Then, we’ll deep dive into the implementation to understand how individual components function.

As promised, I tried to keep the usage semantics similar to the original paper. Compare this MapReduce program using my Go implementation against the C++ example we saw before!

// main.go

import (
  …

  "github.com/MichalPitr/map_reduce/pkg/config"
  "github.com/MichalPitr/map_reduce/pkg/interfaces"
  "github.com/MichalPitr/map_reduce/pkg/mapreduce"
)

type WordCounter struct {
  wordRegex *regexp.Regexp
}

func (wc *WordCounter) Map(input interfaces.MapInput, emit func(key, value string)) {
  text := input.Value()
  text = strings.ToLower(text)
  words := wc.wordRegex.FindAllString(text, -1)
  for _, word := range words {
    emit(word, "1")
  }
}

type Adder struct{}

func (a *Adder) Reduce(input interfaces.ReducerInput, emit func(value string)) {
  val := 0
  for !input.Done() {
    num, err := strconv.Atoi(input.Value())
    if err != nil {
      log.Printf("Failed converting input to integer, skipping: %q", input.Value())
      input.NextValue()
      continue
    }
  val += num
  input.NextValue()
  }
  emit(strconv.Itoa(val))
}

func main() {
  cfg := config.SetupJobConfig()
  log.Printf("cfg: %+v", cfg)
  cfg.NumReducers = 2
  cfg.NumMappers = 4


  cfg.Mapper = &WordCounter{wordRegex: regexp.MustCompile(`\b\w+\b`)}
  cfg.Reducer = &Adder{}

  mapreduce.Execute(cfg)
}

Let’s spend some time understanding how my solution works behind the scenes. We can reference a handy diagram from the paper illustrating this.

The user’s program is executed in three modes, Master, Mapper, and Reducer, on different machines. At a high level, the master handles the entire job orchestration, mappers do the expensive map operation over input files, and reducers combine intermediate results from mappers. In my solution, the mode is determined by a CLI flag, e.g. --mode=master for master mode.

Master

Master mode splits input files into subsets, prepares NFS directories, launches mapper jobs with assigned files, and waits for them to finish. Then it repeats the process for reducers. The main master function is simple enough that I can list it here:

// /pkg/master/master.go

func Run(cfg *config.Config) {
  clientset := createKubernetesClient()
  numNodes := getNumberOfNodes(clientset)
  mustValidateConfig(cfg, numNodes)

  jobId := fmt.Sprintf("job-%s", time.Now().Format("2006-01-02-15-04-05"))
  log.Printf("Running master: %s", jobId)
	
  mustCreateJobDir(cfg.NfsPath, jobId)
  fileRanges := partitionInputFiles(cfg.InputDir, cfg.NumMappers)

  t0 := time.Now()
  launchMappers(cfg, clientset, jobId, fileRanges)
  waitForJobsToComplete(clientset, jobId, "mapper")
  log.Printf("Mappers took %v to finish", time.Since(t0))

  t1 := time.Now()
  launchReducers(cfg, clientset, jobId)
  waitForJobsToComplete(clientset, jobId, "reducer")
  log.Printf("Reducers took %v to finish", time.Since(t1))
  log.Printf("Total runtime: %v", time.Since(t0))
}

Let’s understand the launchMappers function. It creates a Kubernetes job for each mapper. The job spec specifies:

• Docker image containing our binary.

• CLI arguments necessary for the mapper: mapper mode, input/output dirs, and files to process.

• How to mount the NFS storage.

Wait…Docker image? We never created one! I’ll show this in more detail when we finally run our solution, but Kubernetes revolves around containers so we dockerize our binary so that Kubernetes pods can download the binary.

After launchMappers creates a mapper job, Kubernetes assigns jobs to nodes within the cluster, starts the containers, and monitors their status.

Now that we can run mappers, it’s time to learn how they work!

Mappers

In mapper mode, the binary scans input files line-by-line and passes the line to the user’s map function. Map processes the text and emits key-value pairs using the injected emit function.

In my implementation, the emit function is super simple - it stores the key-value pairs in a map of dynamic lists. A more complete implementation would do some buffering and then flush the buffer to the disk.

When the mapper finishes processing all inputs, it saves sorted key-value pairs to intermediate files in the NFS storage from where reducers will later read them for the final processing. Let’s zoom out here and see how this mapping from intermediate files to reducers works.

We want to partition the intermediate files by keys such that all the same keys are processed by one reducer. For instance, “brown” and “the” appear in Mapper 1 and Mapper 2 intermediate files. We want to ensure each ends up being processed by the same reducer. Let’s see an example.

To achieve this, when we save intermediate results from mappers, we partition the keys by the number of reducers R using the formula

For instance, using FNV hash and R = 2, we get

(maths note: this reads as “1 is congruent to FNV(brown) mod 2”.)

Based on this, we assign “brown” to the 2nd intermediate file, which is to be processed by reducer 2. Notice how in the example above, all “the”s landed in the blue files, while all the “brown”s landed in the red files!

This concludes the discussion on mappers - next, let’s see how reducers work.

Reducers

As highlighted previously, the reducer’s job is to read key-value pairs from the assigned intermediate files and process them using the user-defined reduce function.

We can be certain of two things: the keys in intermediate files are sorted by key and if a key A is present in one of these intermediate files, we are guaranteed that key A is not present in any file assigned to other reducers.

I tried to be smarter about my reducer implementation and avoid loading all intermediate files into memory. This comes with an interesting algorithmic problem:

Suppose we are trying to process 3 intermediate files, key-value pair at a time without loading everything into memory.

We can merge the key-value pairs on the fly using a min-heap! We load the first key-value pair from each file into the heap. Whenever we pop from the heap, we read the next line from the corresponding file and push it onto the heap. This gives us a memory-efficient way to read a stream of key-value pairs! You can find the implementation here.

After processing all intermediate files, the reducer saves results to the NFS storage.

For instance, the above example would yield: [“abby”: 1, “alice”: 3, “bob”: 2, “car”:2, “dog”:1, …].

The MapReduce paper introduces a couple of additional optimizations that I skipped in my implementation. Astute readers can probably come up with some already - for instance, couldn’t we optionally do some reduction already on the mapper?

Congratulations, you now have a pretty complete understanding of MapReduce! The hard part is over! Let’s put it together and count the word frequencies in the top 100 most popular books on Project Gutenberg using our MapReduce!

MapReduce in action

I downloaded the 100 most popular books on Project Gutenberg to /mnt/nfs/input/ with this script. The book files are labeled book-0 to book-99. I also created a 4-node minikube cluster with access to my NFS storage.

michal@michal-ThinkPad-T490s:/mnt/nfs$ kubectl get nodes
NAME            STATUS   ROLES           AGE    VERSION
multinode       Ready    control-plane   27d    v1.28.3
multinode-m02   Ready    <none>          7d2h   v1.28.3
multinode-m04   Ready    <none>          7d1h   v1.28.3
multinode-m05   Ready    <none>          7d1h   v1.28.3

I will be reusing the word-frequencies Go main file I showed earlier. We must prepare the docker image using this Dockerfile and push it to my registry. The mapper and reducer nodes will pull this image to run our workloads.

docker build -t michalpitr/mapreduce:latest .
docker push michalpitr/mapreduce:latest

Now all we need to do is run the master locally!

go build .
./map_reduce --mode=master --image=michalpitr/mapreduce:latest --input-dir /mnt/input/ --nfs-path /mnt/nfs/

Let’s inspect the execution logs.

michal@michal-ThinkPad-T490s:~/code/map_reduce$ ./map_reduce --mode=ma
ster --image=michalpitr/mapreduce:latest --input-dir /mnt/nfs/input/ -
-nfs-path /mnt/nfs/
2024/04/28 17:13:53 Running master job-2024-04-28-17-13-53
2024/04/28 17:13:53 Creating mapper-0 for book-0-30
2024/04/28 17:13:53 Creating mapper-1 for book-31-53
2024/04/28 17:13:53 Creating mapper-2 for book-54-76
2024/04/28 17:13:53 Creating mapper-3 for book-77-99
2024/04/28 17:13:53 Waiting for jobs to finish.
2024/04/28 17:14:03 Waiting for jobs to finish.
2024/04/28 17:14:13 Waiting for jobs to finish.
2024/04/28 17:14:23 All jobs completed.
2024/04/28 17:14:23 Mappers took 30.1843218s to finish
2024/04/28 17:14:23 Creating reducer-0
2024/04/28 17:14:23 Creating reducer-1
2024/04/28 17:14:23 Waiting for jobs to finish.
2024/04/28 17:14:33 All jobs completed.
2024/04/28 17:14:33 Reducers took 10.106014703s to finish
2024/04/28 17:14:33 Total runtime: 40.290376921s

The whole job took 40s. We aren’t expecting anything crazy here, after all, I’m running this on an underpowered laptop and not a real cluster.

Let’s walk through this a little. The master creates 4 mappers, each with an assigned book range.

Using kubectl, we can see that the cluster is creating 4 containers. This is when the cluster pulls the image we pushed earlier!

michal@michal-ThinkPad-T490s:/mnt/nfs$ kubectl get pods
NAME             READY   STATUS              RESTARTS   AGE
mapper-0-zgz24   0/1     ContainerCreating   0          11s
mapper-1-m824k   0/1     ContainerCreating   0          11s
mapper-2-ctqqt   0/1     ContainerCreating   0          11s
mapper-3-l2ddh   0/1     ContainerCreating   0          11s

These don’t take very long to start running and to finish.

michal@michal-ThinkPad-T490s:/mnt/nfs$ kubectl get pods
NAME             READY   STATUS      RESTARTS   AGE
mapper-0-zgz24   1/1     Running     0          18s
mapper-1-m824k   0/1     Completed   0          18s
mapper-2-ctqqt   0/1     Completed   0          18s
mapper-3-l2ddh   0/1     Completed   0          18s

The master periodically polls the Kubernetes API server for the status of these jobs. Once they are all finished, it launches the reducers.

Before we go to reducers, let’s see the intermediate files produced by one of the mappers:

michal@michal-ThinkPad-T490s:/mnt/nfs/job-2024-04-28-17-13-53/mapper-3$ ls
partition-0  partition-1

Inspecting the end of partition-1 shows key-value pairs of a character from Dostoevsky’s Notes from Underground. So far so good!

michal@michal-ThinkPad-T490s:/mnt/nfs/job-2024-04-28-17-13-53/mapper-3$ tail partition-1
zverkov,1
zverkov,1
zverkov,1
zverkov,1
zverkov,1
…

Next, the master launches reducers,

michal@michal-ThinkPad-T490s:/mnt/nfs$ kubectl get pods
NAME              READY   STATUS              RESTARTS   AGE
…
reducer-0-4xt58   0/1     ContainerCreating   0          2s
reducer-1-xvhfk   0/1     ContainerCreating   0          2s

which swiftly finish.

michal@michal-ThinkPad-T490s:/mnt/nfs$ kubectl get pods
NAME              READY   STATUS      RESTARTS   AGE
…
reducer-0-4xt58   0/1     Completed   0          7s
reducer-1-xvhfk   0/1     Completed   0          7s

That’s it - we just computed word frequencies using our very own MapReduce framework! Let’s see what we can learn from the results!

I can see the output files in the root of the job folder:

michal@michal-ThinkPad-T490s:/mnt/nfs/job-2024-04-28-17-13-53$ ls
…  reducer-0  reducer-1

Let’s use grep to find the frequencies of some words.

michal@michal-ThinkPad-T490s:/mnt/nfs/job-2024-04-28-17-13-53$ grep -w zverkov reducer-0 reducer-1
reducer-1:zverkov,112

michal@michal-ThinkPad-T490s:/mnt/nfs/job-2024-04-28-17-13-53$ grep -w the reducer-0 reducer-1
reducer-1:the,822175

michal@michal-ThinkPad-T490s:/mnt/nfs/job-2024-04-28-17-13-53$ grep -w mapreduce reducer-0 reducer-1
michal@michal-ThinkPad-T490s:/mnt/nfs/job-2024-04-28-17-13-53$ 

michal@michal-ThinkPad-T490s:/mnt/nfs/job-2024-04-28-17-13-53$ grep -w map reducer-0 reducer-1
reducer-1:map,188

michal@michal-ThinkPad-T490s:/mnt/nfs/job-2024-04-28-17-13-53$ grep -w reduce reducer-0 reducer-1
reducer-0:reduce,123

Unfortunately, there are no mentions of MapReduce in the top 100 most popular books on Project Gutenberg… maybe one day?

One final note, notice how these output files partition the result space by key. When we look for a word, it’s only present in one file! It’s almost as if we did something right!

If you made it this far, you might as well check out the GitHub repo. You’ll already be familiar with most ideas!

Congratulations, you made it! If you enjoyed this deep dive, you might enjoy my previous one into SQLite storage format internals. Researching and writing these articles takes a lot of time and effort. To ensure you don’t miss the next one, consider subscribing or following me on LinkedIn.

Let’s keep it hands-on and create a new SQLite database, a Users table, and insert a single user. 

michal@michal-ThinkPad-T490s:$ sqlite3 mydb.db
SQLite version 3.37.2 2022-01-06 13:25:41
Enter ".help" for usage hints.
sqlite> CREATE TABLE Users (
  Id INT PRIMARY KEY,
  Username VARCHAR(32),  
  Email VARCHAR(255));
sqlite> INSERT INTO Users (Id, Username, Email) VALUES (1, "michal", "michal@example.com");
sqlite> SELECT * FROM Users WHERE Id = 1;
1|michal|michal@example.com
sqlite> 

As expected, everything works just fine. Let’s exit and reopen the database.

sqlite> .exit
$sqlite3 mydb.db
sqlite> SELECT * FROM Users WHERE Id = 1;
1|michal|michal@example.com

We got our user back. At some point, SQLite wrote our changes to the mydb.db file - SQLite is unusual because it stores all data in a single file. Let’s inspect this file with a hex editor to see if we can find our data. The first 16 bytes of the file inform us of the SQLite format used.

A bit deeper, we find the schema for the Users table.

And finally, the row we inserted!

Let’s see what else we can learn about this file by running stat mydb.db.

michal@michal-ThinkPad-T490s:$ stat mydb.db 
  File: mydb.db
  Size: 12288           Blocks: 24         IO Block: 4096   regular file
Device: 10302h/66306d   Inode: 8938166     Links: 1

The file size is 12288 bytes, we’ll come back to this. We might also notice that the IO Block is 4096. This also turns out to be the default page size on my machine, but in general, it depends on the OS and file system.

SQLite stores the database configuration in the first 100 bytes of the root page. The page size is stored at bytes 16-17 and it is 4096 bytes. The number of pages used by the DB is stored in bytes 28-31 and it is 3. Multiply those two together and we get back the file size 4096 * 3 = 12288, exactly what we found with stat earlier!

Let’s recap. We can create a table, insert rows, and store it on disk. Then we can retrieve it from the disk.

Why do we split the database into pages?

I’ve gone ahead and populated the Users table with 100k entries using the following format: sprintf(“%d, user%, user%d@example.com”, i, i, i)

sqlite> SELECT COUNT(Id) FROM Users;
100000

To store those 100k entries, SQLite uses 1361 pages taking up 5.4MB. SQLite internally uses a B+ tree data structure to store data. B+ trees are balanced N-ary trees with a (usually) large number of children per internal node. You can think of them like balanced binary search trees, but with more than 2 children per node, stored values only in leaf nodes, and information on how to traverse the tree in the interior nodes. 

Each node in the tree corresponds to a physical page stored on disk. Inserted rows are stored in leaf nodes only. Index information is stored in interior nodes. I might mix calling them pages and nodes, but they mean the same thing.

Leaf pages

Let’s start with the simpler leaf pages. Both leaf and interior nodes share the same header format visualized in the diagram below. I’ve provided a legend for the fields that we care about. 

The 0th byte of the page determines the page type - 0x0d indicates a leaf page. Bytes 3-4 store the number of cells stored in the node. In the leaf page shown below it’s 123 cells. The entries are stored in reverse order, so we can see that user123 is stored first. The next two bytes store the offset at which the first entry starts - 254. Bytes 8-11 are only used by interior pages.

Once we get to the first entry, we see an array of rows. When I was implementing my version of SQLite from scratch, I also stored a pointer to the parent page in the leaf page to make the implementation a little simpler, since sometimes you need to access the parent node from a leaf.

Next, let’s take a look at the more exciting interior nodes.

Interior pages

So far, we know how to store data. What’s missing is a way to efficiently retrieve it. That is precisely the role of interior pages in a B+ tree. Let’s see how they achieve this. 

Interior pages have the 0th byte set to 0x05.

The header format is the same as for leaf nodes, so the number of cells is in bytes 3-4, here it’s 2 cells. The start of the cell array is indicated by bytes 5-6 as offset from the start of the page: 4082. The right child pointer is stored at bytes 8-11: 1216.

Let’s look at the array of cells and try to understand the cell format for interior pages.

The array starts at 0x1ff2. Each cell is composed of an uint32 page pointer followed by a varint max primary key. 

Interpreting the big-endian uint32 is straightforward: 0x00000267 = 615.

Varint is an encoding that can represent up to 64-bit unsigned ints using 1 to 9 bytes but is more memory efficient for smaller integers.

To decode a varint, we need to look at each byte, starting with the most significant one, and write it in binary:

0x8493370
x84 = 10000100
0x93 = 10010011
0x37 = 00110111

When we reach a byte with the most significant bit set to 0, we have hit the last byte of the varint, which is how I knew that we only needed these 3 bytes. The stored value is then calculated by concatenating the 7 least significant bits:

0000100 | 0010011 | 0110111 = 68023 in decimal

Similarly for the second cell:

0x00000266 = 614
0x82962C = 0000010 | 0010110 | 0101100 = 35628

The values stored as varints indicate that all rows in the subtree pointed to by the page pointer have a primary key lower or equal to this key. Think of it again like a generalization of binary search trees.

So what does this mean? It means that if we are looking for a row with a primary key below or equal to 35628, we should visit page 614. If the key is greater than that but less than or equal to 68023, we should visit page 615. For anything larger, we should follow the right-most pointer to page 1216.

It’s giving us directions on how to find the row we are looking for!

Looking for user4242

Now that we understand interior and leaf pages, let’s look for the user with Id = 4242 to see how this works end to end.

In the previous section, we looked at the root interior node and found that page 614 has entries <= 35628 so let’s look there.

We need to go to page 614, which is at byte offset (614-1) * 4096 (0-indexed) in the mydb.db file. It’s again an internal node as indicated by the 0th byte. I’ll just list the cell entries but notice that I can do some skipping. It turns out that SQLite keeps the cells in sorted order - we can use binary search!

Page_num, max_key
461, 0000010|0010101|1001101 = 35533
…
79,0110000|0111011 = 6203
…
60, 0100100|0100011 = 4643
…
58, 0100011|0111011 = 4539
57, 0100010|1010011 = 4435
56,0100001|1101011 = 4331
55, 0100001|0000011 = 4227
…
4, 123

User4242 should be on page 56. Let’s inspect it and woohoo - it’s a leaf node! We should be close!

Notice how the largest entered user is indeed user4331 as indicated by the max key stored in the parent interior node.

Let’s go find our user4242. If we scroll a bit down, we finally see the row we wanted!

This is amazing - we were able to find our row by just visiting 3 pages: the root page, internal page 614, and leaf page 56. This is important because it means SQLite only had to read those 3 pages from the disk to find the row! Reading a specific page from a file can be implemented efficiently with fseek and reading only PAGE_SIZE bytes.

Disk IO tends to be the slowest part of a database, especially if using physical hard drives. If we didn’t have an index (interior nodes of the B+ tree), we would have to read all 1361 pages. Here we used a tiny database - just 100k rows, but this logarithmic scaling property becomes crucial as the database grows.

If you are wondering how I learned this, I’ve been implementing a subset of SQLite from scratch in Go. I started by following cstack’s SQLite from scratch blog series. For SQLite implementational details, the official SQLite documentation is fantastic.

I hope you enjoyed this deep dive into database internals! If you did, consider subscribing and/or following me on LinkedIn.

I’ve been learning about compilers lately and about how call stacks work. This got me wondering how functional programming languages, where everything is handled recursively deal with maximum call stack size.

Recursion in Python

Let’s start with something familiar and write a valid albeit impractical Python function to sum numbers from 1 to n.

def recursive_sum(n: int) -> int:
    if n <= 0:
        return 0
    return n + recursive_sum(n-1)

It feels weird but it’s pretty much exactly how we’d write this in a functional language like Haskell or OCaml. When I run this with sufficiently large input, 1000 is enough on my machine, I get an error:

RecursionError: maximum recursion depth exceeded.

Every time a function call is made, CPython pushes a new stack frame to the call stack. Among other things, each frame contains function arguments and local variables. When the function returns, the stack frame can be popped from the call stack.

Here’s what the call stack looks like for n = 5.

recursive_sum(5)
5 + recursive_sum(4)
5 + (4 + recursive_sum(3))
5 + (4 + (3 + recursive_sum(2)))
5 + (4 + (3 + (2 + recursive_sum(1))))
5 + (4 + (3 + (2 + (1 + recursive_sum(0))))
5 + (4 + (3 + (2 + (1))))
5 + (4 + (3 + (3)))
5 + (4 + (6))
5 + (10)
15

So why do we get a RecursionError?

Python assumes that a function that exceeds the recursion limit is incorrect — maybe the user forgot to add a base case. That’s a good thing and lets users handle the exception. But in our case, we know that recursive_sum is valid and would terminate if CPython didn’t stop it.

The other part is memory usage. Each stack frame takes up memory. It’s more graceful to let the user handle the exception than for the program to unexpectedly crash because it ran out of memory.

Turns out, there is a way to get pseudo-tail recursion in Python by leveraging exceptions to unwind the stack. Thanks to Chris Penner’s article for explaining this well.

Tail recursion

Functional programming languages do something pretty smart for recursive functions that meet a certain form— the current stack frame is reused for the next call.

First, let’s rewrite our Python function into a tail-recursive form.

def tail_recursive_sum(n: int, accumulator: int) -> int:
    if n <= 0:
        return accumulator
    return tail_recursive_sum(n-1, accumulator + n)

All I did was introduce an accumulator to the function parameters. Then whenever it makes a recursive call, the accumulator is incremented.

Let’s illustrate how this works.

tail_recursive_sum(5, 0)
| tail_recursive_sum(4, 0 + 5)
| | tail_recursive_sum(3, 5 + 4)
| | | tail_recursive_sum(2, 9 + 3)
| | | | tail_recursive_sum(1, 12 + 2)
| | | | | tail_recursive_sum(0, 14 + 1)
| | | | | 15
| | | | 15
| | | 15
| | 15
| 15
15

Notice that once the function reaches the base case, it immediately has the result. However, Python does not take advantage of this. If we run the function above with n = 1000, we get the familiar maximum recursion depth exceeded.

What we can notice though, is that previous stack frames don’t serve any purpose. Compilers that optimize tail recursion take advantage of this fact to reuse the existing stack frame! The call stack then looks something like this:

tail_recursive_sum(5, 0)
tail_recursive_sum(4, 0 + 5)
tail_recursive_sum(3, 5 + 4)
tail_recursive_sum(2, 9 + 3)
tail_recursive_sum(1, 12 + 2)
tail_recursive_sum(0, 14 + 1)
15

Next, let’s see how this works in a language that supports it!

Tail recursion in OCaml

Here’s our tail_recursive_sum written in OCaml.

To explain the syntax, we define a function with the tail_recursive_sum name with two parameters n and accumulator. We use the rec keyword to tell the compiler that it’s a recursive function. Note that there’s no explicit return statement but the last expression is implicitly returned.

let rec tail_recursive_sum n accumulator =
    if n <= 0 then accumulator
    else tail_recursive_sum (n-1) (accumulator + n)

For good measure, let’s test it with10^9,which correctly returns 500000000500000000

Woohoo! That’s pretty remarkable. It would work for larger values as well, but we’d soon need to start worrying about integer overflow.

There’s one last loose end to tie up. To verify that this behavior is indeed due to tail recursion optimization and not a fluke, let’s try using the non-tail recursive form that we used before.

let rec sum n =
    if n <= 1 then 1
    else n + sum (n-1)

As we would hope, we start getting Fatal error: exception Stack_overflow errors starting at around n = 10^6. Great!

If you enjoyed this shorter write-up, consider subscribing. I write technical deep dives, often about implementing complex software from scratch to properly understand its inner workings.

Maybe take a look at one of my other posts?

I’ve included relevant code listings below. You can see the full repository here: https://github.com/MichalPitr/hashmap_vs_trie_c

Here’s the setup. We have a sliding window scanner that reads the source code and emits Tokens which are easier to work with. To illustrate, the source code

var name = “Michael”;

would be converted to the following tokens

TOKEN_VAR, TOKEN_IDENTIFIER, TOKEN_EQUAL, TOKEN_STRING, TOKEN_SEMICOLON

Each token would have associated metadata, like start and length so that we can retrieve the literal string from the source code.

When the scanner sees a lexeme (a sequence of characters), it decides what Token it is. { maps to TOKEN_LEFT_BRACE, == maps to TOKEN_EQUAL_EQUAL.

How does it go about differentiating keywords and identifiers, such as variable names or function names? The rest of this write-up compares two approaches: Hash tables and tries.

Hash Tables

Let’s start with hash tables (HashMap in Java, dictionary in Python, unordered_map in C++) since they are the intuitive go-to solution. I use simple chaining hash tables. When the scanner starts, we can create a hash table where keywords map to their TokenType. Then whenever the scanner finds a lexeme that is either a keyword or an identifier, it looks it up in the hash table.

The listing below shows a function TokenType identifierTypeUsingHashMap() that returns the TokenType of the current lexeme pointed to by scanner.start. If the lexeme isn’t a reserved keyword, it returns TOKEN_IDENTIFIER.

TokenType identifierTypeUsingHashMap() {
    int length = scanner.current - scanner.start;
    TokenType type = hashMapGet(&scanner.keywords, scanner.start, length);
    return type;
}

Let’s have a look at TokenType hashMapGet() to understand how it works:

1. It hashes the key, which requires a full traversal of the string

2. Looks up the corresponding bucket in the hash table

3. Due to possible collision, it checks if the key matches the bucket’s key. This again requires a full traversal of the string. If there’s a match, it returns the keyword’s TokenType.

4. If this is a collision, it checks the next value in the bucket’s linked list.

5. If it reaches the end of the linked list, we know that the key is an identifier, not a keyword.

typedef struct HashNode {
    char* key;
    int keyLength;
    TokenType value;
    struct HashNode* next;
} HashNode;

typedef struct {
    HashNode* table[HASHMAP_CAPACITY];
} HashMap;

TokenType hashMapGet(HashMap* map, const char* key, int length) {
    uint32_t index = hash(key, length);
    HashNode* node = map->table[index];
    while (node) {
        if (node->keyLength == length && strncmp(node->key, key, length) == 0) {
            return node->value; 
        }
        node = node->next;
    }
    return TOKEN_IDENTIFIER;
}

This seems reasonably fast. Especially if the hash function distributes keys uniformly in the hash table to avoid expensive chaining. We can control this through the load factor (entries/capacity).

We are traversing the string multiple times, but when we consider the lexeme “formula”, once the scanner reads “form” it already knows that there is no keyword with the prefix “form”. Can we leverage it? That’s exactly the approach Crafting Interpreters takes.

Tries

Tries are a type of n-ary tree that can be used to efficiently represent a set of words. Any path from the root to the leaf constitutes a valid word in the vocabulary. Here’s an illustration for the following keywords: “false”, “for”, and “fun”.

Let’s represent all keywords with a trie and implement it with a switch-case statement. It’s pretty simple:

1. Check if the first letter of the current lexeme is also the first letter of any keyword. This is equivalent to starting from “Start” in the tree diagram above. If not, it immediately knows that the lexeme is an identifier without any extra work.

2. If the first character belongs to some keyword, it further checks which keyword it could be or compares the rest of the string if there’s only one option left.

static TokenType checkKeyword(int start, int length, const char* rest, TokenType type) {
    if (scanner.current - scanner.start == start + length && 
        memcmp(scanner.start + start, rest, length) == 0) {
        return type;
    } 
    return TOKEN_IDENTIFIER;
}

static TokenType identifierType() {
    switch (scanner.start[0]) {
        case 'a': return checkKeyword(1, 2, "nd", TOKEN_AND);
        case 'c': return checkKeyword(1, 4, "lass", TOKEN_CLASS);
        case 'e': return checkKeyword(1, 3, "lse", TOKEN_ELSE);
        case 'f':
            if (scanner.current - scanner.start > 1) {
                switch (scanner.start[1]) {
                    case 'a': return checkKeyword(2, 3, "lse", TOKEN_FALSE);
                    case 'o': return checkKeyword(2, 1, "r", TOKEN_FOR);
                    case 'u': return checkKeyword(2, 1, "n", TOKEN_FUN);
                }
            }
            break;
        case 'i': return checkKeyword(1, 1, "f", TOKEN_IF);
        case 'n': return checkKeyword(1, 2, "il", TOKEN_NIL);
        case 'o': return checkKeyword(1, 1, "r", TOKEN_OR);
        case 'p': return checkKeyword(1, 4, "rint", TOKEN_PRINT);
        case 'r': return checkKeyword(1, 5, "eturn", TOKEN_RETURN);
        case 's': return checkKeyword(1, 4, "uper", TOKEN_SUPER);
        case 't':
            if (scanner.current - scanner.start > 1) {
                switch(scanner.start[1]) {
                    case 'h': return checkKeyword(2, 2, "is", TOKEN_THIS);
                    case 'r': return checkKeyword(2, 2, "ue", TOKEN_TRUE);
                }
            }
        case 'v': return checkKeyword(1, 2, "ar", TOKEN_VAR);
        case 'w': return checkKeyword(1, 4, "hile", TOKEN_WHILE);
    }

    return TOKEN_IDENTIFIER;
}

TokenType checkKeyword() compares the rest of the two strings with a call to memcmp.

This has several nice properties. For many identifiers, it can determine that it is not a keyword by looking at just the first character. In comparison, the hash table approach had to hash the lexeme every time.

Glue

The driving code doesn’t do anything useful, it simply loads a source file and scans it. Normally it would hand off the Tokens to a parser/compiler, but here I just wanted to focus on the scanner’s performance.

It reads the source file to memory and hands it off to the scanner. We time how long it takes the scanner to finish scanning.

static void runFile(const char* path) {
    char* source = readFile(path);
    initScanner(source);
    // Prime the scanner.
    Token token;
    
    struct timespec start, end;
    clock_gettime(CLOCK_MONOTONIC, &start);
    
    // Scans the source file.
    while (token.type != TOKEN_EOF) {
        advance(&token);
    }

    clock_gettime(CLOCK_MONOTONIC, &end);
    double elapsed = (end.tv_sec - start.tv_sec) + 
                 ((end.tv_nsec - start.tv_nsec)/1e9);
    printf("%f\n", elapsed);
    free(source);
}

int main(int argc, const char* argv[]) {
    if (argc == 2) {
        runFile(argv[1]);
    } else {
        fprintf(stderr, "Usage: clox [path]\n");
        exit(64);
    }
    
    return 0;
}

For completeness, the advance() function just scans the next token. This scanToken() is another switch-case statement that given a character (with some lookahead) decides what token it is. identifier() is what decides if a lexeme is a keyword or an identifier!

Token scanToken() {
    skipWhitespace();
    scanner.start = scanner.current;

    if (isAtEnd()) return makeToken(TOKEN_EOF);

    char c = advance();
    // identifier() uses either the hash map or trie to determine the type
    if (isAlpha(c)) return identifier();
    if (isDigit(c)) return number();

    switch (c) {
        case '(': return makeToken(TOKEN_LEFT_PAREN);
        case ')': return makeToken(TOKEN_RIGHT_PAREN);
        case '{': return makeToken(TOKEN_LEFT_BRACE);
        case '}': return makeToken(TOKEN_RIGHT_BRACE);
        case ';': return makeToken(TOKEN_SEMICOLON);
        case ',': return makeToken(TOKEN_COMMA);
        case '.': return makeToken(TOKEN_DOT);
        case '-': return makeToken(TOKEN_MINUS);
        case '+': return makeToken(TOKEN_PLUS);
        case '*': return makeToken(TOKEN_STAR);
        case '/': return makeToken(TOKEN_SLASH);
        case '!':
            return makeToken(
                match('=') ? TOKEN_BANG_EQUAL : TOKEN_BANG);
        case '=':
            return makeToken(
                match('=') ? TOKEN_EQUAL_EQUAL : TOKEN_EQUAL);
        case '<':
            return makeToken(
                match('=') ? TOKEN_LESS_EQUAL : TOKEN_LESS);
        case '>':
            return makeToken(
                match('=') ? TOKEN_GREATER_EQUAL : TOKEN_GREATER);
        case '"': return string();
    }

    return errorToken("Unexpected character.");
}

static void advance(Token* token) {
    for (;;) {
        *token = scanToken();
        if (token->type != TOKEN_ERROR) break;
    }
}

Benchmark

I compile the scanner with gcc -O3 to enable optimizations. I then run the scanner against source files of varying sizes multiple times (500) to measure its performance. I re-run and recompile with the hashmap implementation and the trie implementation.

Since we are just scanning, I composed a ~1k line source file from different unit tests. I then created a couple of versions of this source file where I copied it consecutively size times, i.e. size 128 corresponds to ~128k lines of code.

def execute_command(command):
    result = subprocess.run(command, stdout=subprocess.PIPE)
    return float(result.stdout.strip())

def main():
    results = defaultdict(list)

    for size in [1, 2, 4, 8, 16, 32, 64, 128]:
        for _ in range(500):
            command = ["./a.out", f"code{size}.clox"]
            output = execute_command(command)
            results[size].append(output)

        print(f"size {size}")
        mean = statistics.mean(results[size])
        std_dev = statistics.stdev(results[size])
        print(f"Mean: {mean}")
        print(f"Standard Deviation: {std_dev}")
        print("")

    with open("results_hashmap.pkl", "wb") as file:
        pickle.dump(results, file)

if __name__ == "__main__":
    main()p

Results

Plotting the results yields the plot below. Vertical bars indicate standard deviation.

As expected, both approaches scale linearly with the size of the source file. 
Tries are on average faster by about 30%. That’s a pretty sizeable difference, especially since modern IDEs re-run the scanner and parser whenever the user makes any changes to provide syntax highlighting and type hints.

In the test source file, I used pretty standard-length identifier names. If used very long ones, the Trie implementation would show even better results, since it can avoid unnecessarily traversing the full string.

Since the trie implementation is so simple, there’s not much of a reason not to use it.

Thanks for reading! If you liked this, consider subscribing and/or following me on LinkedIn. I like to understand the inner workings of software systems. Whenever I learn something interesting, I compile it into one of these blog posts!

Want to read another one? Consider checking out one of my other posts!