166. Optimizing the Rust VM

Twelve and a half percent

Chapter 165’s profile of fib put 12.5% of all instructions in slice/index.rs: indexing a slice or a Vec, bounds check included. clox reads bytecode through a raw pointer and checks nothing. It is tempting to do the same. Before touching anything, two rules from chapters 24 and 142: know where the work is, and know how you’ll tell whether a change helped.

What unchecked would take

The standard library has the unchecked read: xs.get_unchecked(i). It is an unsafe fn. Calling it with i out of bounds is undefined behaviour, “even if the resulting reference is not used”^[1]. In read_byte:

// SAFETY: ip < code.len(), because the
// compiler ends every chunk with a return
// and every jump lands inside the chunk.
let byte = unsafe {
    *self.chunk.code.get_unchecked(self.ip)
};

Read that SAFETY comment again. It isn’t about this line at all: it’s a promise about compiler.rs, hundreds of lines away. One wrong jump offset, and instead of a panic with a message, the VM reads whatever is past the chunk, with no telling what happens next^[3].

And drip-rs refuses it. In a scratch copy, rustc 1.97 says (trimmed):

error: usage of an `unsafe` block
   --> src/vm.rs:433:20
note: the lint level is defined here
  9 | #![forbid(unsafe_code)]

#![forbid(unsafe_code)]

fn main() {
    let code = vec![1u8, 2, 3];
    let ip = 1;
    // SAFETY: ip < code.len()
    let b = unsafe {
        *code.get_unchecked(ip)
    };
    println!("{b}");
}

What it would buy

So measure, in a scratch copy with the forbid deleted. Chapter 141 did this with wall times. Here I added callgrind’s counts. Four builds of drip-rs, each passing the 246 corpus programs and the 95 stress-GC programs: the committed one; one with read_byte unchecked; one with read_constant, peek and the local-slot reads unchecked too; and, for scale, one without chapter 150’s registers, reading every byte through frame, closure, function and chunk. Release builds, rustc 1.97, the book’s shared 4-vCPU VM; 14 fresh runs of fib.drip each, the four builds alternating:

The unsafe builds ran 6.2% and 7.9% fewer instructions. In time, their runs (1266–1571 and 1246–1786 ms) landed among the committed build’s (1331–1677 ms). On this VM that is noise. The safe design change, the registers, was worth 1.95× in instructions and about 2× in time, close to chapter 150’s measurement. (drip-rs’s README says “about 1.4×”; its starting point may not have been this one.)

The lesson: the walk to the byte cost far more than the check on it.

Checks the compiler already moved

drip-rs’s hash table probes with keys[index], where index = hash & mask and mask = keys.len() - 1. Here is the release build’s machine code for Table::get, trimmed:

sub  $0x1,%rbx      # mask = len - 1
jb   <panic_bounds_check>
loop:
inc  %r8
and  %rbx,%r8       # i = (i + 1) & mask
cmpb $0x0,(%r11,%r15,4)  # keys[i]

One test before the loop, that the length isn’t zero. Inside the loop, none: i & (len - 1) can’t reach len. The optimizer hoisted the check itself. I tried to help by reslicing values to keys.len() before the loop. zoo.drip then ran 7,540,512,526 instructions instead of 7,528,846,449, 0.15% more. The line only added a check.

Simple loops get the same treatment. For xs[i] += ys[i] over 1000 elements, a plain indexed loop, one with ys resliced first, and a zip each ran about 2,770 instructions per call, all vectorized.

Checks that can’t fail

Some indexes can’t be proven in range: a byte read from the bytecode, used to pick a constant. But a u8 is at most 255. Index a [T; 256], whose length is in its type, and there is nothing left to check. A slice becomes an array reference with one check up front, try_into^[2]:

// A u8 index into a [T; 256] array can
// never be out of bounds.
fn sum(code: &[u8], k: &[f64; 256])
    -> f64 {
    let mut s = 0.0;
    for &b in code {
        s += k[b as usize];
    }
    s
}

fn main() {
    let consts = vec![0.5; 256];
    // The one check: is the length 256?
    let k: &[f64; 256] = consts
        .as_slice()
        .try_into()
        .expect("256 constants");
    let code: Vec<u8> =
        (0..=255).collect();
    println!("{}", sum(&code, k));

    let short = vec![0.5; 10];
    let r: Result<&[f64; 256], _> =
        short.as_slice().try_into();
    println!("len 10: {}", r.is_ok());
}

I timed that loop against the same one taking &[f64] (rustc 1.97 -C opt-level=3, 4096 bytes, 20,000 times). callgrind counted 573,560,000 instructions for the slice version and 194,800,000 for the array: 2.9× fewer. Wall time: 53 to 57 ms for both, three runs each. Each s += waits for the one before, so the extra compares ran in the gaps.

drip-rs could use this. Constant operands are u8 and a chunk holds at most 256 constants, so a frozen chunk could keep Box<[Value; 256]>. That is design, not code, and the unsafe build bounds the prize: all four unchecked reads together saved 8% of instructions.

Iterators are the third way. Table::iter zips the keys and values columns, and the collector’s Trace for [T] is a plain for t in self. No index, no check^[4].

Your turn

#[derive(Debug, PartialEq)]
enum Verdict {
    Faster,
    Slower,
    Noise,
}

fn min(xs: &[f64]) -> f64 {
    todo!()
}

fn max(xs: &[f64]) -> f64 {
    todo!()
}

fn speedup(
    before: &[f64],
    after: &[f64],
) -> f64 {
    todo!()
}

fn verdict(
    before: &[f64],
    after: &[f64],
) -> Verdict {
    todo!()
}

What you learned

• get_unchecked is unchecked indexing behind unsafe: its safety rests on promises about distant code, and drip-rs forbids it.
• Four unchecked reads saved about 8% of fib’s instructions and no time this VM could tell from noise. The safe registers saved half.
• The optimizer already hoists many checks, such as the hash table’s & mask probe. Adding your own can cost.
• A u8 into a [T; 256] needs no check; try_into makes the array reference with one check. Iterators carry no index at all.
• Fewer instructions isn’t less time. Claim a speed-up only when the runs don’t overlap.