Functional Dispatch

A bytecode interpreter is a program that executes instructions encoded as compact binary data (bytecode). A bytecode interpreter defines its own virtual instruction set and executes it in a loop. This is the approach used by languages like Python and Ruby.

The core of most bytecode interpreters is a dispatch loop.

impl Vm {
    fn run(&mut self) -> Result<Val> {
        let mut current_frame = self.stack_frames.pop().ok_or(Error::StackUnderflow)?;
        loop {
            match current_frame.next_instruction() {
                Instruction::LoadInt(x) => {
                    self.stack.push((x as i64).into());
                }
                ...
            }
        }
    }
}

For example, computing 1 + 2 compiles to three instructions. Each step either pushes a value onto the stack or pops values and pushes a result:

The size of data structures is very important to an interpreter's performance. Smaller datastructures are preferred as:

1. More data can fit in the CPU cache.
2. More data can fit in the CPU registers.
3. Less data to read and write.

Many virtual machines resort to storing both float64 and pointer values in a single 8 byte value through NaN Boxing¹. This is impressive considering that both float64 and pointers individually take 8 bytes, but we can store them in a datastructure that itself takes 8 bytes. Similarly, OCaml uses bit tagging to store either 63bit int or a pointer.

Instructions are another dimension that can be shrunk — which is why the function pointer approach is surprising: it uses a 8 byte pointer instead of a more compact enum opcode.

With the simple op code based interpreter, the representation of a function is mostly the instructions. The VM iterates and jumps around the instructions.

#[derive(Debug, PartialEq)]
struct FuncInner {
    // The number of arguments this function takes.
    arg_count: usize,
    // The instructions for the function.
    instructions: Vec<Instruction>,
    // The table of constants. Instruction references this by index.
    constants: Vec<Val>,
}

With function pointer, the Instruction enum is converted to a function and a data byte. The data byte provides extra context for the function. The data byte is used to store things for some instructions such as how many instructions to jump or which argument to fetch.

#[derive(Debug, PartialEq)]
struct FuncInner {
    arg_count: usize,
    funcs: Vec<fn(&mut Vm, StackFrame, u8) -> Result<Val>>,
    data: Vec<u8>, // data[idx] contains context for funcs[idx]
    constants: Vec<Val>,
}


// An example instruction func.
pub(crate) fn load_int_fn(vm: &mut Vm, mut frame: StackFrame, data: u8) -> Result<Val> {
    vm.stack.push((data as i8 as i64).into());
    let (fn_ptr, data) = frame.advance_instruction_fn();
    // `become` is similar to `return`, but performs a tail call which should be more optimal.
    // See https://doc.rust-lang.org/std/keyword.become.html.
    become fn_ptr(vm, frame, data);
}

This is similar to Fn Ptr, but the data byte is not passed as a function argument. Instead, instructions that need it fetch it themselves by manually reading instruction_data. Instructions that don't need a data byte (like Return) skip the fetch entirely, potentially saving a memory read. The results of the lazy fetching are hard to predict on a general virtual machine and really depend on how often the data byte is needed.

#[derive(Debug, PartialEq)]
struct FuncInner {
    arg_count: usize,
    funcs: Vec<fn(&mut Vm, StackFrame) -> Result<Val>>,
    data: Vec<u8>, // data[idx] contains context for funcs[idx]
    constants: Vec<Val>,
}


// An example instruction func.
pub(crate) fn load_int_fn(vm: &mut Vm, mut frame: StackFrame) -> Result<Val> {
    let data = frame.func.instruction_data(frame.instruction_idx);
    vm.stack.push((data as i8 as i64).into());
    frame.instruction_idx += 1;
    let fn_ptr = frame.func.instruction(frame.instruction_idx);
    become fn_ptr(vm, frame);
}

This introduces more functions based on the instruction. It can be seen as inlining some of the instructions parameters. Take the following compilation code snippet and see how common parameters have a different path.

Take the following example:

1. The general add_n_fn fetches the data byte and adds it to the top value.
2. The optimized add_n_const_fn::<N> inlines the value that will be added. There is no fetching of the data byte, just performing the addition to the top value of the stack.

For a few common values, the optimized implementation may help. For fibonacci specifically, this can be more optimal when performing the n - 2 and n - 1 operations.

// Fast path
Instruction::AddN(-2) => (add_n_const_fn::<-2> as InstructionFn, 0),
Instruction::AddN(-1) => (add_n_const_fn::<-1> as InstructionFn, 0),
Instruction::AddN(0) => (add_n_const_fn::<0> as InstructionFn, 0),
Instruction::AddN(1) => (add_n_const_fn::<1> as InstructionFn, 0),
Instruction::AddN(2) => (add_n_const_fn::<2> as InstructionFn, 0),
// Slow path
Instruction::AddN(n) => (add_n_fn as InstructionFn, n as u8),

// The normal implementation.
pub(crate) fn add_n_fn(vm: &mut Vm, mut frame: StackFrame) -> Result<Val> {
    let data = frame.func.instruction_data(frame.instruction_idx);
    let top = vm.stack.last_mut().ok_or(Error::StackUnderflow)?;
    match top {
        Val::Int(x) => *x += data as i8 as i64,
        _ => {
            return Err(Error::WrongType {
                expected: "Int",
                got: top.type_name(),
            });
        }
    }
    frame.instruction_idx += 1;
    let fn_ptr = frame.func.instruction(frame.instruction_idx);
    become fn_ptr(vm, frame);
}

// A fast path implementation.
pub(crate) fn add_n_const_fn<const N: i64>(vm: &mut Vm, mut frame: StackFrame) -> Result<Val> {
    let top = vm.stack.last_mut().ok_or(Error::StackUnderflow)?;
    match top {
        Val::Int(x) => *x += N,
        _ => {
            return Err(Error::WrongType {
                expected: "Int",
                got: top.type_name(),
            });
        }
    }
    frame.instruction_idx += 1;
    let fn_ptr = frame.func.instruction(frame.instruction_idx);
    become fn_ptr(vm, frame);
}

This adds two extra specializations. For the adder, triop_add_fn fuses two consecutive Binop::Add instructions into a single three-operand add. This one is rather simple compared to the fibonacci instruction…

For fibonacci, jump_if_local0_lt2_fn performs a jump if register 0 is less than 2. The enum -> function converter automatically detects these patterns and uses the specialized functions when appropriate.

jump_if_local0_lt2_fn alone handles the first 3 operations in each fibonacci call. This is significant as each fibonacci call executes 5 instructions for the base case and 11 instructions for the recursive call case.

pub(crate) fn jump_if_local0_lt2_fn(vm: &mut Vm, mut frame: StackFrame) -> Result<Val> {
    let data = frame.func.instruction_data(frame.instruction_idx);
    let local0 = match &vm.stack[frame.stack_start] {
        Val::Int(x) => *x,
        top => {
            return Err(Error::WrongType {
                expected: "Int",
                got: top.type_name(),
            });
        }
    };
    let idx = if local0 < 2 {
        (1 + frame.instruction_idx as isize + data as i8 as isize) as usize
    } else {
        frame.instruction_idx + 1
    };
    frame.instruction_idx = idx;
    let fn_ptr = frame.func.instruction(idx);
    become fn_ptr(vm, frame);
}

Function pointers are dramatically slower, but heavy specialization can beat enum dispatch. The simple adder function takes 15% longer with fibonacci taking 54% longer!

In the Fn Ptr Slim, instead of fetching (next_function, databyte), we only fetch next_function. The function itself will have to manually get the databyte if needed. This one was better for fibonacci, but worse for add. This is somewhat expected and depends on the code. if the databyte is usually used, then prefetching it and storing it in the proper register is good. If not, then copying it around just adds to our overhead.

Specialization greatly improves the cost of fibonacci while the add is around the same level. As a refresher, specialization allowed us to replace some special instructions with more optimized versions. The main optimization is that the functions have constants built-in instead of relying on the value in the databyte.

// Used in both fib and adder
Instruction::LoadLocal(0) => (load_local_const_fn::<0> as InstructionFn, 0),

// Used only in fib
Instruction::AddN(-2) => (add_n_const_fn::<-2> as InstructionFn, 0),
Instruction::AddN(-1) => (add_n_const_fn::<-1> as InstructionFn, 0),
Instruction::LessThan(2) => (less_than_const_fn::<2> as InstructionFn, 0),
Instruction::EvalRecursive(1) => (eval_recursive_const_fn::<1> as InstructionFn, 0),

// Used only in adder
Instruction::LoadLocal(1) => (load_local_const_fn::<1> as InstructionFn, 0),
Instruction::LoadLocal(2) => (load_local_const_fn::<2> as InstructionFn, 0),
Instruction::LoadLocal(3) => (load_local_const_fn::<3> as InstructionFn, 0),

For hyper specialization, I added the silly jump_if_local0_lt2_fn function which greatly improves performance for fibonacci and this pushes it ahead of the original instruction enum based implementation. The adder benchmark only gets the triop_add_fn specialization which was a tiny win; We could probably muck around with better specializations to optimize add further, but it's not a very good exercise.

The takeaway here is that with the added flexibility, we can hack up some functions to improve our benchmark scores. Two microbenchmarks aren't enough to characterize a dispatch mechanism. Results may look very different with a more diverse or realistic workload.

Note: As the interpreter evolves, the performance implications may change. As we add more complexity and branches, function pointers give the branch predictor a stable target per instruction type. A large match has a single indirect branch with many possible targets.