Accelerating the Zig Tokenizer (with AVX-512)

# Accelerating the Zig Tokenizer (with AVX-512) **TL;DR: Niles Salter explains how to use SIMD data parallelism and bitwise primitives (CTZ, etc.) to accelerate the Zig tokenizer, achieving throughput of 820 MB/s, and explores the potential for further speedups with AVX-512.** Speaker: Niles Salter --- ## Project Background This project to accelerate the Zig parser started a long time ago. The first talk was early last year, and the project's goals kept changing as it progressed. For those who thought this was a project that would be "ready for everyone to use as soon as it was finished," the speaker apologizes—this is actually a project that helps him keep growing as a developer and truly learn how instructions interact, involving optimizations for AVX-512, ARM, and RISC-V. Because of this project, he has also filed quite a few issues with LLVM about missed optimizations. Talk outline: 1. What SIMD is 2. Reading assembly code 3. What a tokenizer is and how it traditionally works 4. How a tokenizer can be accelerated 5. Why accelerate it --- ## What Is SIMD? The basic idea of SIMD is **data parallelism**, a form of parallelism provided by modern instruction sets and a major source of speed. On the left, you can see four addition instructions that are independent of one another, with eight input operands in total. On the right is an equivalent SIMD instruction that actually performs the same number of additions. Here is how to read it: `Y` indicates a 256-bit vector, and `Q` indicates that the elements are quadwords, or 64 bits. So this single instruction performs four 64-bit additions. This is helpful because it: - **Reduces the number of instructions**, relieving pressure on the CPU front end, including the decoder and instruction cache. The ITLB and some other components also benefit. - **Reduces power consumption**. SIMD briefly increases power consumption while running, but the key idea is that using SIMD lets you finish the work faster. - **Does more work per cycle**. A CPU can execute only a limited number of micro-operations per cycle—typically four, depending on the type of operation. The four additions on the left can finish in one cycle; the `VPADDD` instruction on the right also takes only one cycle, and theoretically, four more could be performed. With AVX-512, the width of addition instructions can be doubled, allowing significantly more work to be done at once. --- ## Which Workloads Can Be Optimized? If you search for SIMD, you will usually see explanations involving something like SAXPY (single-precision `a * x + y`): lots of mathematical operations performed in any order, which can easily be parallelized. Here, a spectrum is drawn from "optimizable" to "not optimizable." Easily parallelized code like SAXPY is straightforward to handle—take a bunch of numbers, do a bunch of multiplications, do a bunch of additions, and write them out. The basic idea is: **however many multipliers you can give me, I can use them all.** There was a time when people really did consider tokenization impossible to optimize. Now we tend to place it on the optimizable side, but it is worth exploring in depth why that changed. On one hand, you could say LLVM's optimization passes took care of it; on the other, you could say things like simdjson showed that we can actually tokenize or parse several GB per second. But the question is: how many people can really explain why this happens? --- ## Reading Assembly Code First, let's look at some assembly code without explaining its relevance yet. There is a `foo` block whose input is `k0`. The first instruction is `KXNORW`: - `K` means it uses K mask registers. - `XNOR` is like the NOT of `XOR`. - XNORing a value with itself produces all 1 bits. - Shifting it right by 15 bits produces a bit string with just one bit set. Next comes a loop label. Inside the loop: 1. **The `TEST` instruction**—it actually performs two operations: AND on the two operands, and AND NOT on the two operands. Here, AND NOT is used. It sets the carry flag based on whether the AND NOT result is all zeros. In this case, the result is not all zeros, so the carry flag is 0. 2. **`XOR`**—flips the lowest bit. 3. **Left shift**. 4. **`JNC` (jump on not carry)**—checks the carry bit: if it is 0, it jumps back up; if it is 1, it does not jump and falls through. (In the slide's diagram, registers with the same name are aligned vertically: K0 with K0, K1 with K1, helping show how the state changes across iterations.) On the second iteration, K1 has already shifted left by one bit, and `KXOR`, the left shift, and `JNC` are performed again. After shifting left by one bit on the final iteration, execution falls through to `return`. ### Hardware Acceleration: TZCNT This loop can be turned into a single hardware-accelerated instruction—**`TZCNT`**, which was added in 2017. It lets the work just shown on the left be completed significantly faster. Why is this useful? We will find out later. We usually think of an instruction as "delegating a task to something in the hardware." Here, we need to adjust that intuition: think of it as a **fixed-size, hardware-accelerated loop**. What we just saw was a loop—as long as the bits are 0, flip them; when a 1 is encountered, flip it too, then stop. --- ## What Is a Tokenizer? Simply put, it splits source code into chunks. Each chunk is a **token**, and each token has a **kind**—for example, it might be an identifier (essentially "a word" in a programming language), or it might be a multi-character symbol. We also want to know where it starts, where it ends, and how long it is. We do not necessarily have to store all of this information, but we need it at some level. The pieces being split out are **semantically irreducible**. ### Pseudocode for a Traditional Tokenizer `` while (not at end of file) { switch (first character at cursor) { case alphabetic character: token start = current position while (character is alphanumeric or underscore) { advance cursor } token end = current position output token optionally check a hash table to determine whether it is a keyword (keywords are a separate kind in Zig) ... } } `` If it starts with a letter, keep matching alphanumeric characters and underscores until something else is encountered. Then we know the token has ended and can output it. ### The Intuition Behind Parallelization Here we get our first intuition: we are branching on **whether each character is a letter, digit, or underscore**. This information can actually be generated for all files at once, entirely in parallel. This does not mean we should parallelize it without limits, but it does mean we can separate the part that **generates these information bits** from the loop that **keeps checking whether those information bits are true or false**. ### Equivalence to the TZCNT Loop The earlier code is strikingly similar: examine each bit, test it, and see whether we eventually reach a bit that is not 0. We are looking for a sequence of 0s (or a sequence of 1s) and checking where the change occurs. We can generate this mask: we see a sequence of 1s, then a 0—where is the transition? This is exactly what the previous loop was doing. This means we can actually **reduce the branching factor by up to 64 times**—there is a certain equivalence between the two loops, so we can use this technique to make it 64 times faster (not overall, but we are on our way). ### Example of the Scanning Process Look at the first character and jump to the end; look at the first character of the next token and jump to the end. In this example, we have not really gained anything in terms of "skipping multiple characters." But we have gained something in **branch prediction**: the CPU is actually trying to guess how long the tokenizer's next token will be—based on the length of the previous token. This method eliminates most of that prediction, as long as the token fits within a block of bits. --- ## Using Subtraction to Jump Ahead Here is a close-up of the specific implementation for jumping ahead: 1. **Right shift**: shift `x` right by the number of bits corresponding to the current position (11 in this example). Visually, it moves to the left, meaning we have reached the position of that vertical line. 2. **Take NOT**: the main reason is that when shifting left or right, the incoming bits default to 0. Taking NOT gives us 1s at the tail or head. 3. **Subtract one**: perform subtraction using the elementary-school algorithm—subtracting 1 from 0 requires borrowing from the preceding 0; then borrowing from the preceding 0; then borrowing from the preceding 0... Continue borrowing until the first 1 is encountered, and then we are done. The remaining bits stay unchanged. This is exactly the bit-level algorithm for decrementing. 4. **AND with NOT**: capture the bits that changed from 0 to 1—exactly our trailing zeros—and then count them. Of course, as everyone knows, this operation can be turned into **CTZ (count trailing zeros)**. ### Independent of Any Particular Instruction The reason for not showing CTZ earlier was to demonstrate that this is a general fact. When implementing the basic idea of this tokenizer on a single-board computer, **it was still more than twice as fast as a traditional tokenizer, even without hardware CTZ support**. So this is not about having that particular instruction; it is about having primitives that can propagate information bits from the least significant bit toward the most significant bit. --- ## Implementing the Hot-Trace Tokenizer The core idea is: 1. There is a `while` loop. 2. On each iteration, generate some bit strings—the first relates to control characters, the second to spaces, and the last to alphanumeric characters and underscores. 3. At some position in the file, take the **low 6 bits** of that position, **shift the selected bitmap left** by that amount, take **NOT**, and calculate **CTZ**. This lets us skip an entire token at once—unless we run out of bits, in which case we go back and get more. 4. Choose which bit string to query based on the token's first character (select a particular bitmap). > There is a tokenizer bug in the demo that causes the colors to be wrong—but that is the JavaScript tokenizer, so it is fine. --- ## Results This was the speed the hot-trace version of the tokenizer achieved at the time—**820 MB/s**. This was measured before Zig's tokenizer became faster; faster speeds will be shown later. But the question on everyone's mind is: **can we go faster?** --- Source: Accelerating the Zig Tokenizer (with AVX-512) — youtube.com via rcalixte (https://www.youtube.com/watch?v=E6dWIwPLx3s)