Introducing Lightbeam: An Optimising Streaming WebAssembly Compiler

Lightbeam is a new streaming compiler for WebAssembly, designed to produce the best possible assembly while still being fast enough to produce assembly faster than the WebAssembly is received over the wire ¹ .

WebAssembly was designed for streaming compilation, and even from its first public release there was a streaming implementation - Firefox has had its own optimising streaming compiler for a long time and V8 has LiftOff.

Lightbeam is similar in concept, but has a different internal mechanism which leads to surprisingly high-quality code considering the constraints, and I’ll explain how it works in this article.

Lightbeam §

Lightbeam is intended for use as the initial compiler in Wasmtime and as the main compilation engine for Substrate‘s smart contract subsystem.

In relation to the latter, you might have some questions: Why do we need a compiler for our smart contracts? What even are smart contracts? We really want to put a compiler in a blockchain client? For answers to these questions and more you can check out my article about WebAssembly on the blockchain, but it’s not particularly important for the purposes of this article.

What is important is how WebAssembly works and how we can work around its limitations to produce high-quality code in a streaming compiler.

Streaming compilation §

So I’m not going to go too deep into this since Lin Clark has already covered this much more thoroughly than I can, but here’s the pitch. If you want to have fast startup times for a program, you have a few options, each with their own pros and cons:

• Distribute the program as machine code

  • Pros: Very fast, predictable performance
  • Cons: Massively insecure, can’t use microarchitecture-specific opcodes without extra complexity and runtime checks, must distribute separate versions of the program for each target platform

• Distribute bytecode and use an interpreter

  • Pros: Easy to implement, easy to debug, easy to add extra functionality, easy to ensure correctness, easy to write a convenient API, 100% thread- and memory-safe by default, portable to any platform that the host language is
  • Cons: Slow (really, really slow)

• Distribute bytecode and use a JIT

  • Pros: Very fast startup and execution, can lazily compile functions (so unused functions are never compiled), can take advantage of host microarchitecture, can do incredibly powerful optimisation techniques that are impossible with ahead-of-time compilation
  • Cons: Easily the most complex option, no thread- or memory-safety by default, good average-case startup but without extremely careful implementation can have very bad worst-case startup time, must implement a different backend for each target architecture

• Distribute bytecode and use streaming compilation

  • Pros: Can give very good performance with good bytecode design, relatively easy to ensure worst-case startup time is low, can take advantage of host microarchitecture
  • Cons: Must implement a different backend for each target architecture, no thread- or memory-safety by default, need to design bytecode around constraints of streaming compilation

The main reason we decided to use a streaming compiler at Parity is that the strict bounds on compilation time are necessary in the context of the blockchain, currently we use our own in-house WebAssembly interpreter.

Any one of these can be combined with a slower but more-powerful compiler that works in the background as the more immediate method of execution is running the program, with execution switching to the output of the heavyweight compiler when it’s ready.

WebAssembly primer §

So to explain what Lightbeam does internally, I need to explain a few things about WebAssembly. WebAssembly is a weird middle-ground between a high-level language and a more traditional bytecode. It is a bytecode in that it’s a stack machine, canonically represented as a series of opcodes instead of using a text-based format, but it has many high-level features that are rarely seen in bytecodes, like having hierarchical blocks using block...end that you can break out of, a separate block type for blocks that you can jump to the end of and blocks that you can jump to the start of, and an if..else..end construct instead of a using more traditional “test-and-jump” instruction (it also has a test-and-jump instruction, but you currently can’t use it to emulate if statements). Plus, it has “locals”, a concept similar to variables in most high-level languages. These locals are actually more trouble than they’re worth, however, they’re more an artifact of WebAssembly’s history than a useful feature.

So although WebAssembly is designed to be easily compiled and easily compiled with a streaming compiler, it’s actually pretty non-trivial to generate optimal code from WebAssembly directly. That’s why Lightbeam doesn’t.

Gotta go fast §

So, how does Lightbeam get such good code? Well, there’s lots of small optimisations that make it possible, but we start by converting to an SSA, CFG intermediate representation like a “normal”, non-streaming compiler would. This is comparable to LLVM IR or Cranelift IR - a simpler version of the input code that allows you to perform optimisations without having to deal with as many cases. I’ll get into precisely how this differs from WebAssembly later. The difference between this and LLVM IR is that Lightbeam’s IR conversion is streaming - we can convert a stream of WebAssembly to a stream of IR and the backend can convert a stream of IR to a stream of native code. In practice this means that we can generate native code with only one-opcode lookahead.

In the backend we keep track of where values are stored, so that we don’t need to move items around unless absolutely necessary. The upshot of this is that values stay in registers as much as possible and often do not need to be spilled to memory - and when they are spilled to memory they’re spilled in order of least-recently-used. Currently we have 4 kinds of value locations: a value can be stored in memory, in a register, in a constant or as a condition code. The latter means that where many streaming compilers would emit a comparison, a conversion of the comparison’s result to an integer, a comparison of that integer to 0, and then a jump, in Lightbeam we just emit a comparison and then a jump on that condition code, just as a non-streaming compiler would. We also get constant folding for free with this method, all without sacrificing streaming.

In our case we simplify if..end, if..else..end, br_if, br_table, block..end and loop..end constructs to a single, flat form. Instead of nested blocks like WebAssembly, we have a flat list of basic blocks which must end with br, br_if or br_table. Instead of having many ways to switch control flow, we only have 3 ² . Entering a block becomes a no-op, but if, else, loop and end need to be converted into this form. This is luckily extremely simple in practice.

One very useful fact is that there’s no way for a block to jump without also ending the block, which far simplifies value lifetimes. It isn’t possible to conditionally drop a value, you need to conditionally jump to a block which then unconditionally drops the value. Unfortunately, WebAssembly does allow that, so we have some complications. Consider the following WebAssembly:

(block (result i32)
    i32.const 1
    i32.const 2
    get_local 0
    br_if 0
    i32.add
)

(block (result i32)
    i32.const 1
    i32.const 2
    get_local 0
    br_if 0
    i32.add
)

So what this does is returns the 2 constant from the function if local 0 is non-zero, and otherwise returns the 1 and 2 constants added together. That is, the 1 is discarded only if the branch is taken. In the Lightbeam IR we explicitly state which elements are dropped for each branch target, which means that it is impossible to forget to handle this case or to handle it incorrectly - a branch where nothing is discarded and a branch where something is discarded are treated the same. This makes the conversion from WebAssembly to Lightbeam IR more complicated, but that’s the point. Lightbeam IR is far easier to verify for correctness than x86_64 assembly, we want the complex parts of the compilation process to happen during conversion to the IR.

We also annotate each block header with the number of calls and whether it may have so-called backwards callers, i.e. whether there could be branches to this block that the compiler has not yet seen ³ . Because of Wasm’s approach to control flow, we can’t always know ahead of time how many callers a block will have, so as a backup we keep track of the number of callers ourselves. However, blocks generated from Wasm loop instructions may have callers that we don’t yet know about when beginning compilation of that block, since br-family instructions branch to the loop’s header. Because of this, we explicitly mark any block that could have backwards callers. Making this information explicit turns what used to be a complicated set of implicit assumptions into a single boolean check.

The meat of the compiler comes in the form of the virtual stack. This is a LIFO queue of value locations, which can be either:

• Stack: A location on the physical stack (i.e. a memory location that is an offset from rsp). Since rsp can change within a function, this is stored as an offset from what rsp was at the start of the function and the real offset is recalculated each time the value is accessed.
• Register: A value in a register. There are some subtleties about register use that are covered below.
• Immediate: A constant value, which allows us to fold constants and do optimisations like using the immediate-operand version of instructions like add instead of spilling the constant to a register first.
• Condition code: A value that represents one of the bits in the FLAGS register. This means that a comparison operation followed by a select or br_if can be compiled to one of the cc forms of the associated instructions. For example, i32.lt_u + select compiles to cmp + cmovb, the cmp instruction sets some flags on the CPU and the cmovb reads the flags and sets the corresponding register to a certain value if the b flag is set. Since there’s only one FLAGS register and many things overwrite it, we spill this kind of value to a register if anything gets pushed onto the stack on top of it ⁴ .

One important thing is that other than immediates, other values are refcounted. To see why that’s important, we have to talk about locals.

As mentioned above, in WebAssembly there are two kinds of place that a value can be: on the virtual stack or in a local. You can move a value from the stack to a local with set_local and from a local to the stack with get_local. There’s also tee_local, which is equivalent to set_local followed by get_local ⁵ . Anyway, for our purposes the important thing is that get_local causes a value to be duplicated. This also applies to tee_local, but for simplicity we’ll ignore that for now. In Lightbeam IR, we only have the virtual stack, so all values are treated the same. Operations on locals are converted to peek and poke instructions, which get/set values further up the stack. Refcounting means that get_local becomes a peek, and Lightbeam doesn’t need to emit anything when encountering it. It just increases the refcount on that value. set_local becomes poke, which just overwrites the value in the compiler’s virtual stack, again without emitting anything. A register or physical stack slot can be reused if the value stored in it is no longer in the virtual stack, meaning that if an instruction reduces a register’s refcount to zero then the emitted instruction can operate in-place.

For example, on x86 you don’t really have a + operator ⁶ , you only have the equivalent of +=. Normally you’d need to allocate a new register, copy the LHS of the operator into it and add the RHS in-place, but if the refcount on that register is precisely 1 you can operate on it in-place. Another use of refcounting in a streaming compiler is that although in straight-line code you can have aliasing values, when you start a loop (and in some other circumstances) you have to have all your locals and stack elements in unique places. To illustrate why this is, let’s use an example. Here’s a simple loop in WebAssembly:

(func $bad_factorial (param $param i32)
  (local $counter i32)
  (local $output i32)
  (set_local $counter (get_local $param))
  (set_local $output (get_local $counter))
  (loop
    (set_local $counter
      (i32.sub (get_local $counter)
               (i32.const 1)))
    (set_local $output
      (i32.mul (get_local $output)
               (get_local $counter)))
    (br_if 0 (i32.ne (get_local $counter) (i32.const 0)))
  )
  (return (get_local $output))
)

(func $bad_factorial (param $param i32)
  (local $counter i32)
  (local $output i32)
  (set_local $counter (get_local $param))
  (set_local $output (get_local $counter))
  (loop
    (set_local $counter
      (i32.sub (get_local $counter)
               (i32.const 1)))
    (set_local $output
      (i32.mul (get_local $output)
               (get_local $counter)))
    (br_if 0 (i32.ne (get_local $counter) (i32.const 0)))
  )
  (return (get_local $output))
)

So in this function, when the loop starts both locals point to the same place. This is bad, since after the first iteration they will have non-aliasing locations and if we generate code that looks in the same place for both variables we’ll end up returning the wrong result. We could choose arbitrary but static and non-aliasing locations for values used by blocks (in the same way that functions have a specific calling convention called SystemV) but we want to avoid moving values around if possible. So, the first time a block is called we try to set the block’s calling convention to whatever the locations of values happen to be at the first call. When we set the calling convention for the first time, we allocate the minimum number of new locations that will remove any aliasing values. This is why it’s important that we know if a block has precisely one caller: if it only has one caller then we can always keep values where they are. If there are aliasing values then that’s fine, since we know that it’s impossible for this block be called with non-aliasing values ⁷ .

Other baseline compilers will assign consistent locations for locals, usually on the physical stack, in order to keep the compiler simple. By converting to this SSA-form stack machine, Lightbeam can compile many functions without ever spilling values to the stack, and make far better use of registers in general. In fact, at time of writing, even SpiderMonkey (FireFox’s multi-pass optimising compiler) unconditionally spills locals to the stack for the vast majority of functions. Combined with how, in WebAssembly, locals are necessary to persist values across even the most basic of control flow, this means that Lightbeam produces significantly faster code for many small functions. Of course, SpiderMonkey produces better assembly than Lightbeam for most larger functions

For a deeper explanation of Lightbeam’s IR check out the article series that lead me to develop it.

What does this mean for me? §

If you’re just someone who interacts with WebAssembly as a consumer - i.e. you write code that compiles to WebAssembly and the Wasm runtime is just a black box - you might be asking what this means for you. Well, this means that Wasmtime can have vastly better startup performance, and will almost certainly have more consistent startup performance. Other than that, it’s probably something that will be mostly transparent for you. It’s the same as V8′s LiftOff or FireFox’s baseline compiler, but for Wasmtime. However, I hope this article gives an amount of insight into how this unconventional form of compiler can work.