Beating the compiler

In modern times, everyone knows that writing assembly is a fool's errand: compilers are the result of literal engineer-centuries of work, and they know the processor much better than you do. And yet – one hears rumors. Written in ancient tomes, muttered in quiet watering holes, scrawled on the walls of bygone temples, hinted at by mysterious texts; the rumors paint a specific picture: Compilers are bad at generating code for interpreters, and it's possible to outperform them by writing your interpreter in assembly. I recently wrote a fast interpreter for the Uxn CPU, a stack-based architecture with 256 opcodes. The interpreter is a simple loop which reads a byte from RAM then selects the appropriate instruction: impl Uxn { /// Runs the VM starting at the given address until it terminates #[inline] pub fn run<D: Device>(&mut self, dev: &mut D, mut pc: u16) { loop { let op = self.ram[usize::from(pc)]; pc = pc.wrapping_add(1); let Some(next) = self.op(op, dev, pc) else { break; }; pc = next; } } /// Executes a single operation #[inline] fn op<D: Device>(&mut self, op: u8, dev: &mut D, pc: u16) -> Option<u16> { match op { 0x00 => op::brk(self, dev, pc), 0x01 => op::inc::<0b000>(self, dev, pc), 0x02 => op::pop::<0b000>(self, dev, pc), 0x03 => op::nip::<0b000>(self, dev, pc), 0x04 => op::swp::<0b000>(self, dev, pc), 0x05 => op::rot::<0b000>(self, dev, pc), 0x06 => op::dup::<0b000>(self, dev, pc), 0x07 => op::ovr::<0b000>(self, dev, pc), 0x08 => op::equ::<0b000>(self, dev, pc), 0x09 => op::neq::<0b000>(self, dev, pc), 0x0a => op::gth::<0b000>(self, dev, pc), 0x0b => op::lth::<0b000>(self, dev, pc), 0x0c => op::jmp::<0b000>(self, dev, pc), 0x0d => op::jcn::<0b000>(self, dev, pc), 0x0e => op::jsr::<0b000>(self, dev, pc), // ... etc } } } All of the opcode implementations end up monomorphized and inlined into the body of Uxn::run(..), and the compiler is smart enough to keep key values in registers. This makes it relatively fast; I see 10-20% speedup over the reference implementation. Let's look at the assembly and see what the compiler is doing – and whether we can do any better. For context, the Uxn CPU has four different memories: The data stack, which is a [u8; 256] along with a u8 index The return stack, which has the same format RAM, which is a [u8; 65535] Device memory, which we'll ignore for the moment (along with the D: Device argument) During evaluation, we also track the program counter pc, which is a u16 used to index into the RAM. In each cycle, we load a byte from RAM, then call the appropriate opcode. Some opcodes can also read and write to RAM, so self-modifying code is possible! By examining the assembly, we can reverse-engineer which values are stored where. Consider the INC operation, which loads a value from the top of the data stack and increments it: ; INC 0x100002d4c: ldrb w8, [x25] ; read the current data stack index 0x100002d50: ldrb w9, [x24, x8] ; read a byte from the data stack 0x100002d54: add w9, w9, #1 ; increment that byte 0x100002d58: strb w9, [x24, x8] ; write that byte back to the stack 0x100002d5c: b 0x100002d1c ; jump back to the dispatch loop From this assembly, we learn the following: x25 is the address of the data stack index (not its value!) x24 is the address of the data stack array w9 is used as a temporary register Similarly, INCr – increment the top value in the return stack – teaches us that x22 and x23 are the return stack's data and index addresses. JMP shows that our program counter is stored in w27: ; JMP 0x100002eac: ldrb w8, [x25] ; read the current data stack index 0x100002eb0: ldrsb w9, [x24, x8] ; read a signed jump offset from the data stack 0x100002eb4: sub w8, w8, #1 ; decrement the data stack index 0x100002eb8: strb w8, [x25] ; write back the data stack index 0x100002ebc: add w27, w27, w9 ; apply the jump to our program counter 0x100002ec0: b 0x100002d1c ; jump back to the dispatch loop Finally, the dispatch loop itself is worth examining: 0x100002d1c: and x10, x27, #0xffff ; mask pc to a u16 0x100002d20: ldr x8, [x20, #256] ; load RAM base from *mut Uxn 0x100002d24: ldrb w10, [x8, x10] ; load opcode byte from RAM 0x100002d28: add w27, w27, #1 ; increment pc 0x100002d2c: adr x11, #-96 ; load base for jump 0x100002d30: ldrh w12, [x27, x10, lsl #1] ; load per-opcode jump amount 0x100002d34: add x11, x11, x12, lsl #2 ; compute jump location 0x100002d38: br x11 ; jump into opcode implementation The compiler has generated a jump table of 256 offsets (each a 2-byte value, indicated by lsl #1). It reads an opcode-specific value from this table to compute a jump target, then performs an indirect branch to jump into the opcode's implementation. We can run this in a debugger and dump the actual jump table: (lldb) disas -p -c3 raven-cli`raven_uxn::Uxn::run::had9dba0d7d1b5105: -> 0x100002d30 <+236>: ldrh w12, [x27, x10, lsl #1] 0x100002d34 <+240>: add x11, x11, x12, lsl #2 0x100002d38 <+244>: br x11 (lldb) reg read x27 x27 = 0x0000000100170b10 (lldb) memory read -s2 -fu -c256 0x0000000100170b10 0x100170b10: 2923 0x100170b12: 31 0x100170b14: 36 0x100170b16: 40 0x100170b18: 44 0x100170b1a: 52 0x100170b1c: 28 0x100170b1e: 64 0x100170b20: 70 0x100170b22: 78 0x100170b24: 86 0x100170b26: 94 0x100170b28: 119 0x100170b2a: 102 0x100170b2c: 110 0x100170b2e: 125 ; etc... (indeed, this is how I generated the per-opcode instruction listing) Having looked at the assembly, there are two things that stick out as possible inefficiencies: Some critical values (stack indices, the base address of RAM) are kept in memory instead of registers; for example, INC has an extra load operation to get the current data stack index. The dispatch loop takes a single indirect branch to the opcode-specific implementation. This means that the branch will be nigh unpredictable! Profiling the code, the hottest instructions are all in the dispatch loop; the ldrh takes over 1/3 of the total runtime! (I'm not confident that the profiler is attributing the time to the correct specific instruction here, but the vibes definitely indicate that dispatch is expensive) LuaJIT is the fast interpreter par excellence, and it's written in assembly. Mike Pall specifically calls out keeping state in registers and indirect threading as two contributors to its speed, which can only be accomplished reliably in assembly. Since persuading our compiler to generate extremely specific patterns is hard, let's get started writing some assembly of our own. My home machine is an M1 Macbook, so all of the assembly will be AArch64-flavored. The implementation uses general-purpose registers; be aware that w* and x* refer to 32-bit and 64-bit views of the same register. Register assignment Our first optimization is to store all important data in registers, to avoid superfluous loads and stores. My implementation ends up using 9 registers (x0-x8), along with a handful of scratch registers: ; x0 - stack pointer (&mut [u8; 256]) ; x1 - stack index (u8) ; x2 - return stack pointer (&mut [u8; 256]) ; x3 - return stack index (u8) ; x4 - RAM pointer (&mut [u8; 65536]) ; x5 - program counter (u16), offset of the next value in RAM ; x6 - VM pointer (&mut Uxn) ; x7 - Device handle pointer (&DeviceHandle) ; x8 - Jump table pointer ; x9-15 - scratch registers The AArch64 calling convention only gives you 8 input arguments, so we can't call a function directly with all of these values in registers; we'll need a C ABI-flavored entry point (discussed below). Indirect threading Our second optimization is using threaded code to eliminate the dispatch loop. Each opcode's implementation will end with a jump to the next opcode's implementation. Opcodes are stored as single bytes in VM RAM, with a base address of x4. I'll build a separate jump table of function pointers, then pass its address in register x8. On the Rust side, here's what that table looks like: extern "C" { fn BRK(); fn INC(); fn POP(); fn NIP(); fn SWP(); fn ROT(); fn DUP(); fn OVR(); fn EQU(); // ...etc } const JUMP_TABLE: [unsafe extern "C" fn(); 256] = [ (BRK as unsafe extern "C" fn()), (INC as unsafe extern "C" fn()), (POP as unsafe extern "C" fn()), (NIP as unsafe extern "C" fn()), (SWP as unsafe extern "C" fn()), (ROT as unsafe extern "C" fn()), (DUP as unsafe extern "C" fn()), (OVR as unsafe extern "C" fn()), (EQU as unsafe extern "C" fn()), (NEQ as unsafe extern "C" fn()), // ... etc ]; In assembly, we want to read the current byte from VM RAM (x4), use it to pick an address in the jump table (x8), then jump to that address. I defined a macro to do this dispatch: .macro next ldrb w9, [x4, x5] ; load the byte from RAM add x5, x5, #1 ; increment the program counter and x5, x5, #0xffff ; wrap the program counter ldr x10, [x8, x9, lsl #3] ; load the opcode implementation address br x10 ; jump to the opcode's implementation .endm Notice that this is a macro, not a function; we'll add next to the end of each opcode, which will expand into this text. For example, here's INC: .global _INC _INC: ldrb w9, [x0, x1] ; read the byte from the top of the stack add w9, w9, #1 ; increment it strb w9, [x0, x1] ; write it back next ; jump to the next opcode Unlike LuaJIT, there's no decoding step for instructions; there are no register arguments, and the single-byte opcode uniquely defines program behavior. Implementation Implementing the other 255 opcodes is mostly just turning the crank; there's nothing particularly exotic here, just good honest assembly. In many cases, I'll use helper macros to generate code for a group of instructions: .macro binary_op op ldrb w10, [x0, x1] ; read the top value from the data stack pop ; decrement the data stack index (this is a macro!) ldrb w11, [x0, x1] ; read the next value from the data stack