Add a Wasm JIT Using Your Knowledge from Track 2

Topic Overview

JIT for Wasm

Once you have an interpreter, compiling hot Wasm functions to native code removes dispatch overhead. Template or trace JITs map Wasm opcodes to machine instructions, reusing techniques from your bytecode VM track.

Compilation Strategy

Approaches:

• Template JIT: one native snippet per Wasm opcode, stitched in order
• Baseline + optimizing: quick code first, recompile with profiling (production engines)
• Bounds check elimination: prove memory safe, elide redundant traps

Executable memory via \`mmap\` + \`mprotect\` matches what you built for the bytecode VM JIT task.

For example, \`i32.add\` becomes x86 \`add\` after two loads from the Wasm stack slots mapped to registers or a native stack frame.

Fallback to the interpreter keeps rare opcodes simple.

Design Notes

Map Wasm locals to native stack slots or registers consistently for the whole function. Changing mapping mid-function makes debugging generated code nearly impossible without a side table.

For This Exercise

This exercise asks you to add a basic JIT to your Wasm runtime using prior JIT knowledge. You will emit native code for a short Wasm function and call it instead of interpreting opcode by opcode.

You will use the same mental model here when reading production interpreter source later in the track. Sketch one concrete input on paper, predict the outcome, then confirm with code. That discipline catches logic errors early and makes debugging far faster when you extend the implementation in follow-on tasks.

Your Task

Write a baseline JIT for straight-line WebAssembly (i32 only) that emits x86-64 assembly in two ways:

1. Naive: every Wasm value lives on the machine stack, with a push/pop for each operand.
2. Register-cached, like V8's Liftoff: the compiler keeps a virtual value stack whose entries are a constant, a not-yet-loaded local, or a register. It folds constant arithmetic, uses memory and immediate operands directly, and loads values only when an instruction needs a register.

Print both listings with their instruction counts, so the difference is visible. invoke runs the function in an interpreter as a reference.

Input

;; starts a comment.

cCopy
Loading…

At end, the function is checked first:

• no stack underflow (NAME: stack underflow at I (OP));
• at most 7 values deep (NAME: stack deeper than 7 values);
• exactly one value left (NAME: must leave exactly one i32 (leaves N));
• every i32.shl directly preceded by i32.const (NAME: i32.shl at I needs an i32.const count right before it).

Only then are the listings printed.

Frame (both versions)

push rbp, mov rbp, rsp, sub rsp, F (F = 4·(params+locals) rounded up to 16; omitted if 0). Then param i is stored with mov [rbp-4(i+1)], REG from edi esi edx ecx r8d r9d, and each local j is zeroed with mov dword [rbp-4(np+j+1)], 0. The code ends with leave and ret, and the result goes in eax.

Naive code

Register-cached code

• Registers are allocated as the lowest free one in the order eax ecx edx esi edi r8d r9d.
• Materializing an entry emits mov R, V for a constant or mov R, [rbp-S] for a local, and turns it into a register.
• local.get and i32.const only push entries, with no code.
• drop frees the entry's register.
• Binary op (right operand b popped, left a on top):
  • if both are constants, fold them (count it; 32-bit wrap);
  • otherwise materialize a into R and emit OP R, b, where b is written as V, [rbp-S] or a register. Multiplying by a constant uses imul R, R, V, and shl uses shl R, V&31;
  • then free b's register.
• local.set/tee I:
  • if the top entry is local I itself, emit nothing (set still pops);
  • otherwise, first materialize every stacked entry that is local I (bottom to top), so they keep the old value. Then materialize the top if it is a local, and emit mov dword [rbp-S], V or mov [rbp-S], R;
  • set pops the entry and frees its register.
• End: mov eax, V, mov eax, [rbp-S], mov eax, R (if R is not eax), or nothing.

Output

cCopy
Loading…

Errors:

• line N: unsupported instruction X
• line N: bad operand for OP, line N: OP takes no operand
• line N: only i32 is supported (T), line N: at most 6 params (register arguments)
• line N: outside a func
• NAME: not compiled (a function with line errors)
• invoke: no compiled function, invoke: expected N argument(s)

Example

Input:

cCopy
Loading…

Output:

cCopy
Loading…

Requirements

1. Generate code in one pass over the Wasm instructions, as a baseline compiler must (it is judged by compile speed).
2. Handle the stale-local hazard: a lazily-referenced local must be loaded before that local is overwritten.

Hidden tests cover constant folding into an immediate, local.tee with a later reuse, the swap pattern that needs the stale-local flush, a register-register operation, drop, and the error checks.