Lab 30: A Cranelift JIT (Milestone 14)

Background

You will compile a very small subset of Ember to native code with Cranelift: a leaf function with integer parameters, arithmetic, comparison, and a return. With guards. With a working deoptimization path.

function add(a, b) return a + b end

That is the target, and it is enough — because the interesting content is the hot counter, the guards, the deopt map, and the collector's cooperation, not the code generation.

Why This Lab Matters

• You will be able to draw the pipeline and explain every arrow, which is the stated objective.
• Deoptimization is the feature, and building even a trivial one makes every JIT paper legible.
• This is Ember's only unsafe, and it gets the full five-part treatment.

Prerequisites

• Labs 28–29 complete; the profile, the feedback, and the baseline all exist.
• JIT Architecture read.
• cargo add cranelift-jit cranelift-frontend cranelift-codegen cranelift-module --optional, behind the jit feature.

Predict First

1. fib(25) — VM versus JIT. What factor? (Be pessimistic: there is call overhead at the boundary.)
2. How many guards does function add(a, b) return a + b end need?
3. A guard fails on the 500th call. What must be reconstructed?
4. If compiled code allocated a table, what would the collector need?
5. Which is bigger: your JIT's code or its deopt/boundary machinery?

Step 1: Scope, Stated

Compile only what you can do correctly. The subset:

"No allocation in compiled code" is the scope decision that removes stack maps, and it must be stated rather than implied. A function that cannot allocate cannot trigger a collection, so the collector never needs to find roots inside a compiled frame.

#![allow(unused)]
fn main() {
fn is_jit_eligible(proto: &Proto, feedback: &TypeFeedback) -> Option<Reason> {
    // Refuse loudly and specifically. A JIT that silently declines is a JIT
    // you cannot debug.
    if proto.is_vararg               { return Some(Reason::Vararg); }
    if !proto.upvals.is_empty()      { return Some(Reason::Captures); }
    for op in proto.chunk.code() {
        match op {
            Op::Add | Op::Sub | Op::Mul | Op::Lt | Op::Le | Op::Eq
            | Op::GetLocal(_) | Op::SetLocal(_) | Op::LoadInt(_)
            | Op::Jump(_) | Op::JumpIfFalse(_) | Op::Return(_) => {}
            other => return Some(Reason::Unsupported(*other)),
        }
    }
    if !feedback.all_sites_monomorphic_int(proto) { return Some(Reason::NonIntegerFeedback); }
    None
}
}

Step 2: The Hot Counter

#![allow(unused)]
fn main() {
pub struct Proto {
    // ...
    pub call_count: Cell<u32>,
    pub compiled: Cell<Option<CompiledFn>>,
}

const JIT_THRESHOLD: u32 = 1_000;
}

#![allow(unused)]
fn main() {
// In do_call, before pushing the frame:
if let Some(f) = proto.compiled.get() {
    if let Some(result) = self.try_call_compiled(f, &args)? { return Ok(result); }
    // None ⇒ a guard failed at entry; fall through and interpret.
} else {
    let n = proto.call_count.get() + 1;
    proto.call_count.set(n);
    if n == JIT_THRESHOLD { self.jit.try_compile(proto, &self.feedback); }
}
}

Compile at exactly == THRESHOLD, not >=. Otherwise a proto that fails compilation retries on every subsequent call, and a JIT that cannot compile something spends more time trying than the interpreter spends running it.

Step 3: Bytecode to Cranelift IR

The stack-machine-to-SSA conversion, which is the pass a register VM would not need (the stack-vs-register chapter said this would land here):

#![allow(unused)]
fn main() {
fn translate(&mut self, proto: &Proto, builder: &mut FunctionBuilder) -> Result<()> {
    let mut stack: Vec<CraneliftValue> = Vec::new();   // the COMPILE-TIME operand stack

    for (ip, op) in proto.chunk.code().iter().enumerate() {
        match op {
            Op::LoadInt(n) => stack.push(builder.ins().iconst(types::I64, *n as i64)),
            Op::GetLocal(s) => stack.push(builder.use_var(Variable::from_u32(*s as u32))),
            Op::SetLocal(s) => { let v = stack.pop().unwrap();
                                 builder.def_var(Variable::from_u32(*s as u32), v); }
            Op::Add => { let b = stack.pop().unwrap(); let a = stack.pop().unwrap();
                         stack.push(builder.ins().iadd(a, b)); }
            Op::Lt  => { let b = stack.pop().unwrap(); let a = stack.pop().unwrap();
                         stack.push(builder.ins().icmp(IntCC::SignedLessThan, a, b)); }
            Op::Return(_) => { let v = stack.pop().unwrap(); builder.ins().return_(&[v]); }
            _ => unreachable!("is_jit_eligible filtered this"),
        }
    }
    Ok(())
}
}

The VM's runtime stack became a compile-time Vec. That is the whole conversion, it is fifteen lines for this subset, and it is exactly the pass V8 avoids by making Ignition register-based.

Step 4: Guards at the Boundary

For this subset, all the guards are at entry, which is what makes the deopt map trivial:

#![allow(unused)]
fn main() {
fn try_call_compiled(&mut self, f: CompiledFn, args: &[Value]) -> Result<Option<Value>> {
    // THE GUARD. Everything the compiled code assumed, checked once, cheaply.
    let mut unboxed = [0i64; MAX_JIT_PARAMS];
    for (i, a) in args.iter().enumerate() {
        match a { Value::Integer(n) => unboxed[i] = *n,
                  _ => { self.stats.deopts += 1; return Ok(None); } }   // ← bail to the interpreter
    }
    // SAFETY: see Step 5.
    let raw = unsafe { f.call(&unboxed[..args.len()]) };
    Ok(Some(Value::Integer(raw)))
}
}

Bailing at entry means the "deoptimization map" is empty: no compiled code ran, so there is no state to reconstruct. That is the honest scope of this lab, and the challenge extensions are where it gets real.

Step 5: The unsafe, With the Five-Part Treatment

Ember's only unsafe, and it needs all five parts:

#![allow(unused)]
fn main() {
/// # Safety
///
/// 1. WHY SAFE RUST IS INSUFFICIENT
///    Executing machine code produced at run time requires transmuting a data
///    pointer to a function pointer and calling it. There is no safe Rust
///    construct for this; it is the definition of what a JIT does.
///
/// 2. THE INVARIANT
///    a. `ptr` was produced by `JITModule::get_finalized_function` for a
///       function this module compiled in this process.
///    b. The signature we transmute to matches the one we declared to
///       Cranelift: `extern "C" fn(*const i64, usize) -> i64`.
///    c. `args` points to `len` initialized i64s, valid for the call's duration.
///    d. The module has been finalized (memory is executable, no longer writable)
///       and has NOT been freed — enforced by holding the JITModule alive in
///       `self.jit` for the process's lifetime.
///    e. The compiled code does not allocate, so it cannot trigger a collection
///       and cannot need a stack map.
///
/// 3. THE REGION
///    Exactly the transmute and the call below. Nothing else in this file is
///    unsafe.
///
/// 4. TESTS
///    tests/jit.rs::{compiled_matches_interpreted_on_the_corpus,
///                   guard_failure_falls_back, module_outlives_every_pointer}
///
/// 5. HOW VIOLATION CAUSES UNSOUNDNESS
///    (b) wrong ⇒ arguments read from the wrong registers ⇒ arbitrary behavior.
///    (d) wrong ⇒ jump to freed or non-executable memory ⇒ crash or worse.
///    (e) wrong ⇒ a collection with unscanned roots ⇒ use-after-free.
unsafe fn call(&self, args: &[i64]) -> i64 {
    let f: extern "C" fn(*const i64, usize) -> i64 = std::mem::transmute(self.ptr);
    f(args.as_ptr(), args.len())
}
}

Read invariant (e) again. It is the scope decision from Step 1, restated as a soundness condition — which is what makes "we do not support allocation" a safety claim rather than a limitation.

Step 6: Measure the Three-Way

cargo bench --features jit -- --baseline v0.1

Expect policy_10k to barely move, because most of it is field access and calls, which this subset cannot compile. That is the honest result and it is the interesting one: a JIT that compiles 3% of your workload gives you 3% of a speedup, and the profile told you that before you started.

The Trace

$ ember run --features jit --trace-jit -e '
local function add(a, b) return a + b end
local s = 0
for i = 1, 5000 do s = add(s, i) end
return s'

jit: proto 'add' call 1000/1000 → eligible, compiling
jit: translate  7 bytecode ops → 4 IR instructions
jit:   v0 = iconst.i64  (param a)
jit:   v1 = iconst.i64  (param b)
jit:   v2 = iadd v0, v1
jit:   return v2
jit: compiled 'add' → 24 bytes at 0x7f3a...  (compile time: 84 us)
jit: call 1001..5000 → compiled, 0 deopts
--- jit report ---
  protos compiled: 1        protos rejected: 0
  compiled calls: 4,000     deopts: 0
  compile time: 84 us       payback after ~<n> calls
12502500

payback after ~n calls is the number that matters. Compilation costs 84 µs; if a compiled call saves 20 ns, you need ~4,200 calls to break even — which is more than this loop makes. Print it, and let the reader discover that a JIT can be a net loss on a short run. That is a real and under-taught fact.

Now the deopt path:

$ ember run --features jit --trace-jit -e '
local function add(a, b) return a + b end
local s = 0
for i = 1, 2000 do s = add(s, i) end
s = add(s, 1.5)                            -- a FLOAT: the guard fails
for i = 1, 2000 do s = add(s, i) end       -- integers again
return s'

jit: proto 'add' → compiled
jit: call 1001..2000 → compiled
jit: DEOPT at entry guard: param 1 is Float, expected Integer
jit:   → interpreting this call
jit: call 2002..4000 → compiled, guard holds again
--- jit report ---
  compiled calls: 3,999     deopts: 1

The deopt cost one interpreted call and nothing else. Because the guard is at entry, no compiled code ran and no state needed reconstructing — which is exactly why this scope was chosen, and exactly what makes a real JIT hard.

And rejection, which must be loud:

$ ember run --features jit --trace-jit -e '
local function f(t) return t.x end
for i = 1, 2000 do f({x = i}) end
return 1'

jit: proto 'f' call 1000/1000 → REJECTED: unsupported opcode GET_FIELD
jit: proto 'f' will not be retried

"Will not be retried" is the == THRESHOLD decision, visible.

Expected Output

$ cargo test --features jit --test jit
test compiled_matches_interpreted_on_the_corpus ... ok
test guard_failure_falls_back_to_the_interpreter ... ok
test ineligible_protos_are_rejected_once ... ok
test module_outlives_every_function_pointer ... ok

$ cargo test                      # no jit feature: forbid(unsafe_code) holds
$ cargo clippy --features jit -- -D warnings

Debugging Steps

Cranelift verifier errors

Usually a block without a terminator, or a Variable used before def_var. Cranelift's verifier is strict and its messages are good — read them before guessing.

The compiled function returns garbage

The signature you declared does not match the one you transmuted to. Invariant (b). Print both.

A segfault after many calls

The JITModule was dropped while a pointer was still live. Invariant (d): hold it for the process's lifetime, or reference-count it.

The JIT is slower than the VM

Entirely possible on short runs — check the payback number. Also check that you are benchmarking a release build with the compilation cost excluded from the steady-state measurement (or included deliberately, and labelled).

Deopt fires on every call

The feedback said monomorphic-int but the actual arguments are not. Your feedback is being collected at the wrong site or is stale.

cargo test without --features jit fails to compile

The JIT module is not properly #[cfg(feature = "jit")]-gated. The default build must have zero unsafe and zero Cranelift.

Experiment

CLAIM. A JIT is a net loss below a certain call count, and the break-even point is computable from the compile time and the per-call saving.

METHOD. For loop counts of 100, 1,000, 10,000, 100,000, and 1,000,000, run the add benchmark with the JIT enabled and disabled. Plot or tabulate total time.

PREDICTION. Where do the lines cross? Does it match the payback after ~n calls estimate the trace printed?

RESULT. Record it in docs/learning/15-jit.md. This is why real JITs have tiers: HotSpot's C1 compiles fast and badly, C2 compiles slowly and well, and the interpreter handles everything below C1's threshold. You have just derived the reason for tiered compilation from your own measurement.

Test

#![allow(unused)]
fn main() {
#[test]
fn compiled_matches_interpreted_on_the_corpus() {
    // Every JIT-eligible function in the corpus, both ways. The differential
    // discipline, extended to a THIRD backend — and the harness needed zero
    // changes, which is Lab 8's design paying off one last time.
    for case in corpus().iter().filter(|c| c.has_jit_eligible_functions) {
        assert_eq!(run_with_jit(&case.src), run_without_jit(&case.src), "{}", case.name);
    }
}

#[test]
fn guard_failure_falls_back_to_the_interpreter() {
    let out = run_with_jit_stats(
        "local function add(a,b) return a+b end
         local s = 0
         for i = 1, 2000 do s = add(s, i) end
         s = add(s, 1.5)
         for i = 1, 2000 do s = add(s, i) end
         return s");
    assert_eq!(out.deopts, 1);
    assert_eq!(out.value, run_without_jit_value(/* same source */));
}

#[test]
fn ineligible_protos_are_rejected_once() {
    let out = run_with_jit_stats("local function f(t) return t.x end
                                  for i = 1, 5000 do f({x=i}) end return 1");
    assert_eq!(out.compile_attempts, 1, "a rejected proto was retried");
}

#[test]
fn module_outlives_every_function_pointer() {
    // Invariant (d). Compile many functions, drop the protos, keep calling.
    let mut e = Engine::with_jit();
    let ptrs = compile_many(&mut e, 100);
    drop_protos(&mut e);
    for p in ptrs { assert!(call_compiled(p).is_ok()); }
}

#[test]
fn the_default_build_has_no_unsafe_and_no_cranelift() {
    // `#![cfg_attr(not(feature = "jit"), forbid(unsafe_code))]` is the claim;
    // this checks the dependency half.
    let out = std::process::Command::new("cargo").args(["tree", "--depth", "1"]).output().unwrap();
    assert!(!String::from_utf8_lossy(&out.stdout).contains("cranelift"));
}
}

Challenge Extensions

1. Real deoptimization. Support a guard inside compiled code (an overflow check on ADD), which requires a map from the guard site to interpreter state: which registers hold which slots, and what ip to resume at. This is the actual lesson, and it is a weekend.
2. Loops. Compile a whole loop body, with an on-stack-replacement entry so a running loop can transfer into compiled code. OSR is how HotSpot handles a long-running loop in an interpreted method.
3. Stack maps. Allow allocation in compiled code, and use Cranelift's stack-map support so the collector can find roots in a compiled frame. Then verify under --gc-stress.
4. A second tier. Compile at 1,000 calls with minimal optimization and at 100,000 with full Cranelift optimization. Measure whether tiering helps — you predicted the reason for it in the experiment.
5. Compare with luajit. Same benchmark, both runtimes. Write down the factor honestly. It will be large. Then read lj_record.c and identify three things LuaJIT does that Ember does not.

Deliverables

• The jit feature, off by default; #![cfg_attr(not(feature = "jit"), forbid(unsafe_code))].
• is_jit_eligible with specific rejection reasons, and no retry after rejection.
• A hot counter compiling at == THRESHOLD.
• Bytecode → Cranelift IR for the stated subset, with the compile-time operand stack.
• Entry guards; a guard failure falls back to the interpreter and is counted.
• The unsafe block with all five parts of the treatment, including invariant (e).
• The three-way benchmark table, with policy_10k included and its honest result.
• --trace-jit showing translation, compilation, deopts, rejections, and the payback estimate.
• The break-even experiment, and the tiered-compilation conclusion drawn from it.
• Corpus tests run with and without the JIT, agreeing.
• docs/adr/ADR-014-cranelift.md and docs/learning/15-jit.md written.
• Milestone 14 complete.

Validation / Self-check

1. State the compiled subset and the one scope decision that removes the need for stack maps.
2. Why compile at == THRESHOLD rather than >=?
3. What does the bytecode-to-IR pass do, and which VM design would not need it?
4. Where are this lab's guards, and why does that make the deopt map empty?
5. Give all five parts of the unsafe treatment for call. Which invariant is a restatement of the scope decision?
6. What is the payback point, and what did your experiment say about where the lines cross?
7. What does that imply about why real JITs have tiers?
8. Why did policy_10k barely move, and why is that the expected result rather than a failure?
9. How many changes did the test harness need to accommodate a third backend?

Next: The Capstone — The Recommendation Policy Engine.