Can the borrow checker algorithm be auto-vectorized? The answer is yes!

TL;DR — Vx's borrow checker packs the lifetimes of a reference type into four 16-bit slots of one 64-bit word. The check that one lifetime may stand in for another is a loop over those slots that stops at the first failure, and LLVM does not vectorize it. Rewritten without early exits, the same check gives the same answer on every input we could try, and LLVM turns it into SIMD code with no branches at all: four slots in one register for a single check, eight slots per instruction for a batch. In a batch it costs about 1 ns per check whatever the input, where the old loop takes 0.7 to 12.5 ns, depending on the input. The compiler now uses the new check.

A borrow checker sounds like the last thing a CPU's vector unit could help with. It walks a program, follows references and asks questions like “does this borrow outlive the value it points to?” That is pointer chasing and branching, the opposite of the long, regular loops that SIMD instructions are built for.

Part of it, though, is plain arithmetic on small integers, and Vx stores that part in a shape a vector unit likes. This post takes that one step of Vx's borrow checker and shows that LLVM vectorizes it on its own, once the code is written in a form it can vectorize. No hand-written SIMD intrinsics are involved.

Lifetimes as numbers

Vx gives every lexical block a region: its nesting depth. Region 0 is 'static, region 1 is a function's body, and each block inside adds one. A smaller region lives longer. A reference type also carries a variance, which says how its lifetime may change when the reference is assigned:

The compiler gives a type a 256-bit identifier, and for a reference type one 64-bit word of it holds up to four lifetime parameters, 16 bits each:

bit:   63      48 47      32 31      16 15       0
       [ slot 3  ] [ slot 2  ] [ slot 1  ] [ slot 0  ]

one slot:  [ variance : 4 bits | region : 12 bits ]
           slot 0 is the return slot: only 9 of its region bits are the region

A region field with every bit set means “not decided yet” and matches anything. When a reference of one type is assigned to a place of another type, the checker compares the two words slot by slot. Here is that check as it was in src/borrow.rs, in verify_subtyping_bounds, with the comments and the lookup around it taken out and a shorter name:

fn check_scalar(bits_a: u64, bits_b: u64) -> bool {
    if bits_a == bits_b {
        return true;
    }
    for i in 0..4 {
        let shift = i * 16;
        let slot_a = (bits_a >> shift) & PARAM_MASK;
        let slot_b = (bits_b >> shift) & PARAM_MASK;
        let variance_a = (slot_a & VARIANCE_MASK) >> 12;
        let variance_b = (slot_b & VARIANCE_MASK) >> 12;
        if variance_a != variance_b {
            return false;
        }
        let (region_mask, region_unset) = if i == 0 {
            (REGION_MASK_0, REGION_UNSET_0)
        } else {
            (REGION_MASK, REGION_UNSET)
        };
        let region_a = slot_a & region_mask;
        let region_b = slot_b & region_mask;
        if region_a == region_unset || region_b == region_unset {
            continue;
        }
        let valid = match variance_a {
            0x0 => region_a == region_b,
            0x1 => region_a <= region_b,
            0x2 => region_a >= region_b,
            _ => false,
        };
        if !valid {
            return false;
        }
    }
    true
}

It reads well, and it is fast when the answer comes early. But every return and continue inside the loop is a branch whose direction depends on the data. A loop that can leave in the middle is also a loop LLVM's vectorizer will not touch: to work on four slots at once, it has to know all four will be looked at.

The same check, with no way out of the loop

The fix is to stop asking questions one at a time. Work out every condition for every slot, combine them with & and |, and decide at the end:

const MASKS: [u16; 4] = [0x01FF, 0x0FFF, 0x0FFF, 0x0FFF];

// One slot. Every condition is computed, then combined with & and |.
fn slot_ok(a: u16, b: u16, mask: u16) -> bool {
    let (va, vb) = (a >> 12, b >> 12);
    let (ra, rb) = (a & mask, b & mask);
    let wildcard = (ra == mask) | (rb == mask);
    let order = ((va == 0) & (ra == rb))
              | ((va == 1) & (ra <= rb))
              | ((va == 2) & (ra >= rb));
    (va == vb) & (wildcard | order)
}

// All four slots are checked, always, and the answers are ANDed.
fn check_branchless(bits_a: u64, bits_b: u64) -> bool {
    let mut all = true;
    for i in 0..4 {
        let a = (bits_a >> (16 * i)) as u16;
        let b = (bits_b >> (16 * i)) as u16;
        all &= slot_ok(a, b, MASKS[i]);
    }
    (bits_a == bits_b) | all
}

The logic is the same, written as one formula per slot. The slot-0 mask goes in a small table instead of an if. The “not decided yet” value is the mask itself, so one comparison against the mask covers both slot kinds.

What LLVM makes of it

Compiled by rustc 1.95.0 with --release for an Apple M4, a single call to check_branchless becomes this. It is the complete function, with nothing left out:

dup.2d v0, x1
adrp x8, lCPI6_0@PAGE
ldr q1, [x8, lCPI6_0@PAGEOFF]
ushl.2d v2, v0, v1
adrp x8, lCPI6_1@PAGE
ldr q3, [x8, lCPI6_1@PAGEOFF]
ushl.2d v0, v0, v3
uzp1.4s v2, v0, v2
xtn.4h v2, v2
xtn.2s v0, v0
dup.2d v4, x0
ushl.2d v1, v4, v1
ushl.2d v3, v4, v3
uzp1.4s v1, v3, v1
xtn.4h v1, v1
ushr.4h v3, v1, #12
adrp x8, lCPI6_2@PAGE
ldr d4, [x8, lCPI6_2@PAGEOFF]
and.8b v1, v1, v4
and.8b v5, v2, v4
uzp1.4h v0, v0, v0
mov.s v0[1], v2[1]
ushr.4h v0, v0, #12
cmeq.4h v2, v1, v4
mvn.8b v2, v2
cmeq.4h v4, v5, v4
bic.8b v2, v2, v4
movi.4h v4, #2
cmeq.4h v4, v3, v4
cmhi.4h v6, v5, v1
orn.8b v4, v6, v4
cmtst.4h v6, v3, v3
cmeq.4h v7, v1, v5
orn.8b v6, v6, v7
movi.4h v7, #1
cmeq.4h v7, v3, v7
cmhi.4h v1, v1, v5
orn.8b v1, v1, v7
and.8b v1, v6, v1
and.8b v1, v4, v1
and.8b v1, v2, v1
cmeq.4h v0, v3, v0
orn.8b v0, v1, v0
umaxv.4h h0, v0
fmov w8, s0
cmp x0, x1
cset w9, eq
orn w8, w9, w8
and w0, w8, #0x1
ret

It has no conditional branches. The .4h suffix means “four 16-bit lanes”: LLVM placed the four slots of each word side by side in one NEON register, the way they already sit in the u64. ushr.4h #12 takes out all four variances at once. cmeq.4h and cmhi.4h compare all four regions at once, which covers the equal, smaller-or-equal and greater-or-equal tests. umaxv.4h folds the four answers into one. This is LLVM's SLP vectorizer at work: it found four copies of the same arithmetic in straight-line code and packed them into one register.

A compiler can check many assignments, so the next question is a loop over many pairs:

fn batch_pairs(a: &[u64], b: &[u64], out: &mut [bool]) {
    for ((x, y), o) in a.iter().zip(b).zip(out.iter_mut()) {
        *o = check_branchless(*x, *y);
    }
}

LLVM's loop vectorizer takes this one. The compiled function contains 434 vector instructions, working on eight slots per instruction (ushr.8h, cmhs.8h, cmeq.8h), with cmeq.2d testing two whole words for equality at once. The old check_scalar compiles to 98 instructions, 37 of them conditional branches, and no vector instructions.

LLVM reports both steps itself. Built with -C remark=all, it names the SLP vectorizer for the single check and the loop vectorizer for the batch, which handles 16 pairs per pass through the loop.

The form of the code matters. Our first attempt kept the branch-free slot formula but walked all the slots of all the pairs in one flat loop, looking up each slot's mask with MASKS[j % 4], and then ANDed each group of four in a second loop. It gave the same answers but took about 4.5 ns per check whatever the input. That beats the old loop only where the old loop is slowest, and loses to it on random and identical words. The version above, with the four slots written out per pair, is the fast one.

Same answers

A faster check that gives a different answer is a broken borrow checker, so we compared the two versions three ways:

The first test cannot try every pair of whole words, since there are 2128 of them. It does cover every input a single slot can see, and the result for a word is the results of its slots ANDed together, plus the whole-word equality, which the second test exercises.

Is it faster?

Sometimes. Each version checked the same 1,048,576 pairs, best of 30 runs, on one core of an Apple M4. The made-up lifetime words have variances 0 to 2 and small regions, and almost all of them fail somewhere. The passing pairs pass every slot, and none of them are identical. Every time below is per check, and each range covers three separate runs. The code, and the commands that produce every number in this post, are in benchmarks/lifetime_check.

First, a loop over a batch, which lets the loop vectorizer work on many pairs at once:

InputOld loopVectorized
Pairs that pass every slot12.1–12.5 ns1.07 ns
Made-up lifetime words4.8 ns1.07 ns
Random 64-bit words1.5–1.6 ns1.07 ns
Identical words0.70–0.74 ns1.07 ns

The vectorized check costs the same for every input, because it never branches. The old loop depends on its input. Random words almost always fail at slot 0, the branch predictor learns that, and the loop stays within about half a nanosecond of the new code. When it has to look at all four slots and the outcomes vary, mispredicted branches make it over ten times slower. When the two words are identical, its first line returns straight away and it beats the new code.

The compiler does not check assignments in batches, though. It makes one call per assignment, so the second measurement keeps each check a separate call. Here the four slots are still in one register, but nothing is shared between pairs:

InputOld loopBranch-freeEquality test, then branch-free
Pairs that pass every slot12.5–12.7 ns2.4 ns2.4 ns
Made-up lifetime words4.8–5.0 ns2.4 ns2.4 ns
Random 64-bit words1.7–1.8 ns2.3–2.4 ns2.4 ns
Identical words0.94–0.99 ns2.4 ns0.74 ns

The last column keeps the old loop's first line, if bits_a == bits_b { return true; }, and runs the branch-free check after it. That one predictable branch gives back the fast case for identical words, and the result is the fastest or close to it for every input except random words. This is the version Vx now uses. We have not measured how often each kind of pair turns up in real Vx programs, so how much it saves in practice is still open.

What this does not show

The point is the answer to the question in the title. Part of a borrow checker can be put into a form that a vector unit can run, and the compiler does the vectorizing by itself, as long as the code gives it straight lines to work with. Packing lifetimes into fixed-width fields is what makes that possible: four slots of 16 bits are already four lanes of a vector register.

The check lives in src/borrow.rs (verify_subtyping_bounds). For how Vx's borrow checker behaves from a programmer's side, see Ownership and borrowing.