byteflow/jit/
dispatch.rs

1//! Execute compiled traces and map exits back to VM state.
2#![allow(unsafe_code)]
3
4use std::sync::Arc;
5
6use crate::{Chunk, Fault, Value, Vm, VmResult};
7
8use super::compiler::TraceCompiler;
9use super::exit::{ExitReason, JitReturn};
10use super::frame::{JitCallRecord, JitFrame, MAX_JIT_CALL_DEPTH};
11use super::module_local::with_jit_module;
12use super::runtime::JitRuntime;
13use super::trace::{CompiledTrace, JitContext, TraceKey, HOT_THRESHOLD};
14
15/// Result of copying VM registers into JIT shadow slots.
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum SyncSlotsResult {
18    Ok,
19    /// A register held a non-int type — interpreter must resume.
20    Deopt,
21    /// Frame unavailable.
22    Unavailable,
23}
24
25/// Attempt one native trace at the current `(function, pc)`.
26fn try_jit_step(vm: &mut Vm, ctx: &mut JitContext, budget: u32) -> Option<VmResult> {
27    if !Arc::ptr_eq(&ctx.chunk, &vm.chunk_arc()) {
28        return None;
29    }
30    let pc = vm.current_pc()?;
31    let function = vm.current_function();
32    let register_count = vm.current_num_registers()?;
33
34    let mut slots = vec![0i64; slot_capacity(vm)];
35    match sync_slots_from_vm(vm, &mut slots) {
36        SyncSlotsResult::Ok => {}
37        SyncSlotsResult::Deopt | SyncSlotsResult::Unavailable => return None,
38    }
39
40    let key = TraceKey {
41        function,
42        entry_pc: pc as u32,
43    };
44    let exit = try_run_hot(
45        ctx,
46        vm.chunk_arc().as_ref(),
47        key,
48        &mut slots,
49        budget,
50        register_count,
51        None,
52    )?;
53    match apply_exit_to_vm(vm, &slots, exit) {
54        Ok(Some(result)) => Some(result),
55        Ok(None) => None,
56        Err(fault) => Some(VmResult::Trap(fault)),
57    }
58}
59
60fn try_jit_step_runtime(
61    vm: &mut Vm,
62    runtime: &JitRuntime,
63    budget: u32,
64    metrics: Option<&crate::scheduler::RuntimeMetrics>,
65) -> Option<VmResult> {
66    if !runtime.matches_chunk(&vm.chunk_arc()) {
67        return None;
68    }
69    let pc = vm.current_pc()?;
70    let function = vm.current_function();
71    let register_count = vm.current_num_registers()?;
72
73    let mut slots = vec![0i64; slot_capacity(vm)];
74    match sync_slots_from_vm(vm, &mut slots) {
75        SyncSlotsResult::Ok => {}
76        SyncSlotsResult::Deopt => {
77            if let Some(m) = metrics {
78                crate::scheduler::RuntimeMetrics::inc(&m.jit_deopts);
79            }
80            return None;
81        }
82        SyncSlotsResult::Unavailable => return None,
83    }
84
85    let key = TraceKey {
86        function,
87        entry_pc: pc as u32,
88    };
89    let exit = try_run_hot_runtime(
90        runtime,
91        vm.chunk_arc().as_ref(),
92        key,
93        &mut slots,
94        budget,
95        register_count,
96        metrics,
97    )?;
98    match apply_exit_to_vm(vm, &slots, exit) {
99        Ok(Some(result)) => Some(result),
100        Ok(None) => None,
101        Err(fault) => Some(VmResult::Trap(fault)),
102    }
103}
104
105/// Run at most `budget` instructions, retrying compiled traces after each
106/// interpreter step when a site becomes hot.
107pub fn run_vm_with_jit(vm: &mut Vm, budget: u32, ctx: &mut JitContext) -> VmResult {
108    run_vm_with_jit_loop(vm, budget, |vm, remaining| try_jit_step(vm, ctx, remaining))
109}
110
111/// Runtime path: shared cache, per-thread compilation module.
112pub fn run_vm_with_jit_runtime(
113    vm: &mut Vm,
114    budget: u32,
115    runtime: &JitRuntime,
116    metrics: Option<&crate::scheduler::RuntimeMetrics>,
117) -> VmResult {
118    run_vm_with_jit_loop(vm, budget, |vm, remaining| {
119        try_jit_step_runtime(vm, runtime, remaining, metrics)
120    })
121}
122
123fn run_vm_with_jit_loop(
124    vm: &mut Vm,
125    budget: u32,
126    mut try_step: impl FnMut(&mut Vm, u32) -> Option<VmResult>,
127) -> VmResult {
128    let mut remaining = budget;
129    while remaining > 0 {
130        if let Some(result) = try_step(vm, remaining) {
131            return result;
132        }
133        match vm.run(1) {
134            VmResult::Yield => {
135                remaining -= 1;
136            }
137            other => return other,
138        }
139    }
140    VmResult::Yield
141}
142
143fn slot_capacity(vm: &Vm) -> usize {
144    let chunk = vm.chunk_arc();
145    let max_in_chunk = chunk
146        .functions
147        .iter()
148        .map(|f| f.num_registers as usize)
149        .max();
150    match (max_in_chunk, vm.current_num_registers()) {
151        (Some(max_regs), Some(current)) => max_regs.max(current as usize),
152        (Some(max_regs), None) => max_regs,
153        (None, Some(current)) => current as usize,
154        (None, None) => 0,
155    }
156}
157
158/// Run a compiled trace against the shadow slot table.
159pub fn run_compiled_trace(
160    entry: super::frame::JitEntry,
161    slots: &mut [i64],
162    pc: u32,
163    budget: u32,
164    function: u32,
165    register_count: u32,
166) -> JitReturn {
167    let mut call_stack = [JitCallRecord {
168        return_pc: 0,
169        return_function: 0,
170        dest_reg: 0,
171        caller_register_count: 0,
172    }; MAX_JIT_CALL_DEPTH];
173    let mut frame = JitFrame {
174        slots: slots.as_mut_ptr(),
175        register_count,
176        pc,
177        budget,
178        function,
179        exit_kind: 0,
180        return_reg: 0,
181        call_depth: 0,
182        call_stack: call_stack.as_mut_ptr(),
183    };
184    unsafe { (entry)(&mut frame) };
185    JitReturn {
186        kind: frame.exit_kind,
187        pc: frame.pc,
188        return_reg: frame.return_reg,
189    }
190}
191
192/// Run a cached trace by reference.
193pub fn run_compiled_trace_ref(
194    trace: &CompiledTrace,
195    slots: &mut [i64],
196    pc: u32,
197    budget: u32,
198    function: u32,
199    register_count: u32,
200) -> JitReturn {
201    run_compiled_trace(trace.entry, slots, pc, budget, function, register_count)
202}
203
204/// Compile `key` if needed and execute once. Returns `None` when compilation
205/// fails (caller should fall back to the interpreter).
206pub fn try_run_hot(
207    ctx: &mut JitContext,
208    chunk: &Chunk,
209    key: TraceKey,
210    slots: &mut [i64],
211    budget: u32,
212    register_count: u8,
213    metrics: Option<&crate::scheduler::RuntimeMetrics>,
214) -> Option<ExitReason> {
215    if ctx.cache.get(&key).is_none() {
216        if !ctx.hot.hit(key, ctx.hot_threshold) {
217            if let Some(m) = metrics {
218                crate::scheduler::RuntimeMetrics::inc(&m.jit_misses);
219            }
220            return None;
221        }
222        let mut compiler = TraceCompiler::new(ctx.module_mut());
223        match compiler.compile_trace(chunk, key) {
224            Ok(compiled) => {
225                ctx.cache.insert(key, compiled);
226                if let Some(m) = metrics {
227                    crate::scheduler::RuntimeMetrics::inc(&m.jit_compiles);
228                }
229            }
230            Err(_) => {
231                if let Some(m) = metrics {
232                    crate::scheduler::RuntimeMetrics::inc(&m.jit_compile_failures);
233                }
234                return None;
235            }
236        }
237    }
238    let trace = ctx.cache.get(&key)?;
239    if let Some(m) = metrics {
240        crate::scheduler::RuntimeMetrics::inc(&m.jit_executions);
241    }
242    let ret = run_compiled_trace_ref(
243        trace,
244        slots,
245        key.entry_pc,
246        budget,
247        key.function,
248        u32::from(register_count),
249    );
250    Some(ret.into_reason())
251}
252
253/// Shared-runtime variant (no global lock during execution).
254pub fn try_run_hot_runtime(
255    runtime: &JitRuntime,
256    chunk: &Chunk,
257    key: TraceKey,
258    slots: &mut [i64],
259    budget: u32,
260    register_count: u8,
261    metrics: Option<&crate::scheduler::RuntimeMetrics>,
262) -> Option<ExitReason> {
263    if let Some(copy) = runtime.get_trace(&key) {
264        if let Some(m) = metrics {
265            crate::scheduler::RuntimeMetrics::inc(&m.jit_executions);
266        }
267        let ret = run_compiled_trace(copy.entry, slots, key.entry_pc, budget, key.function, u32::from(register_count));
268        return Some(ret.into_reason());
269    }
270
271    if !runtime.record_hot_hit(key) {
272        if let Some(m) = metrics {
273            crate::scheduler::RuntimeMetrics::inc(&m.jit_misses);
274        }
275        return None;
276    }
277
278    // Double-checked compile after hot threshold.
279    if runtime.get_trace(&key).is_some() {
280        return try_run_hot_runtime(runtime, chunk, key, slots, budget, register_count, metrics);
281    }
282
283    let compiled = match with_jit_module(|module| TraceCompiler::new(module).compile_trace(chunk, key))
284    {
285        Ok(trace) => trace,
286        Err(_) => {
287            if let Some(m) = metrics {
288                crate::scheduler::RuntimeMetrics::inc(&m.jit_compile_failures);
289            }
290            return None;
291        }
292    };
293
294    runtime.insert_trace(key, compiled);
295    if let Some(m) = metrics {
296        crate::scheduler::RuntimeMetrics::inc(&m.jit_compiles);
297    }
298
299    try_run_hot_runtime(runtime, chunk, key, slots, budget, register_count, metrics)
300}
301
302/// Apply a native trace exit to a live [`Vm`], producing an optional early
303/// [`VmResult`] when the trace finished the flow.
304pub fn apply_exit_to_vm(
305    vm: &mut Vm,
306    slots: &[i64],
307    exit: ExitReason,
308) -> Result<Option<VmResult>, Fault> {
309    match exit {
310        ExitReason::Return { return_reg } => {
311            let value = Value::Int(slots[return_reg as usize]);
312            vm.set_register(return_reg, value.clone())?;
313            let pc = vm
314                .current_pc()
315                .ok_or(Fault::Invariant("empty frame stack while running"))?;
316            vm.set_pc(pc.saturating_add(1));
317            vm.return_value(value)
318        }
319        ExitReason::Trap { pc } => {
320            vm.set_pc(pc as usize);
321            Ok(Some(VmResult::Trap(Fault::DivideByZero)))
322        }
323        ExitReason::Budget { pc } => {
324            vm.set_pc(pc as usize);
325            sync_slots_to_vm(vm, slots)?;
326            Ok(None)
327        }
328        ExitReason::Continue { pc } | ExitReason::Effect { pc } | ExitReason::Deopt { pc } => {
329            vm.set_pc(pc as usize);
330            if matches!(exit, ExitReason::Continue { .. } | ExitReason::Effect { .. }) {
331                sync_slots_to_vm(vm, slots)?;
332            }
333            Ok(None)
334        }
335    }
336}
337
338/// Copy int-compatible registers from the VM into JIT shadow slots.
339pub fn sync_slots_from_vm(vm: &Vm, slots: &mut [i64]) -> SyncSlotsResult {
340    let Some(regs) = vm.top_registers() else {
341        return SyncSlotsResult::Unavailable;
342    };
343    for (slot, value) in slots.iter_mut().zip(regs.iter()) {
344        match value {
345            Value::Int(i) => *slot = *i,
346            Value::Bool(b) => *slot = i64::from(*b),
347            Value::Unit => *slot = 0,
348            _ => return SyncSlotsResult::Deopt,
349        }
350    }
351    SyncSlotsResult::Ok
352}
353
354/// Write JIT shadow slots back into the active VM register file.
355pub fn sync_slots_to_vm(vm: &mut Vm, slots: &[i64]) -> Result<(), Fault> {
356    let Some(regs) = vm.top_registers_mut() else {
357        return Err(Fault::Invariant("empty frame stack while syncing slots"));
358    };
359    for (reg, slot) in regs.iter_mut().zip(slots.iter()) {
360        *reg = Value::Int(*slot);
361    }
362    Ok(())
363}
364
365/// Threshold at which a `(function, pc)` pair becomes a compilation candidate.
366pub const fn hot_threshold() -> u32 {
367    HOT_THRESHOLD
368}
369
370/// Force-compile a trace (tests / benchmarks).
371pub fn force_compile(
372    ctx: &mut JitContext,
373    chunk: &Chunk,
374    key: TraceKey,
375) -> Result<(), super::error::CompileError> {
376    if ctx.cache.get(&key).is_none() {
377        let mut compiler = TraceCompiler::new(ctx.module_mut());
378        let compiled = compiler.compile_trace(chunk, key)?;
379        ctx.cache.insert(key, compiled);
380    }
381    Ok(())
382}
383
384#[cfg(test)]
385mod tests {
386    use std::sync::Arc;
387
388    use crate::{bytecode::builder::ChunkBuilder, FlowOutcome, NativeTable, Opcode, Runtime, Value, Vm, VmResult};
389
390    use super::*;
391    use crate::jit::trace::{JitContext, TraceKey};
392
393    type TestResult = Result<(), Box<dyn std::error::Error>>;
394
395    fn add_chunk() -> std::sync::Arc<crate::Chunk> {
396        let mut b = ChunkBuilder::new("jit-add");
397        b.begin_function("main", 0, 3);
398        b.emit_load_imm(0, 41);
399        b.emit_load_imm(1, 1);
400        b.emit_binop(Opcode::Add, 2, 0, 1);
401        b.emit_return(2);
402        Arc::new(b.finish())
403    }
404
405    #[test]
406    fn run_vm_with_jit_completes_scalar_add() -> TestResult {
407        let chunk = add_chunk();
408        let mut ctx = JitContext::new(chunk.clone())?;
409        let key = TraceKey {
410            function: 0,
411            entry_pc: 0,
412        };
413        force_compile(&mut ctx, &chunk, key)?;
414
415        let mut vm = Vm::new(chunk, NativeTable::empty(), 0, &[])?;
416        let result = run_vm_with_jit(&mut vm, 10_000, &mut ctx);
417        assert!(matches!(result, VmResult::Complete(Value::Int(42))));
418        Ok(())
419    }
420
421    #[test]
422    fn native_trace_computes_42() -> TestResult {
423        let chunk = add_chunk();
424        let mut ctx = JitContext::new(chunk.clone())?;
425        let key = TraceKey {
426            function: 0,
427            entry_pc: 0,
428        };
429        force_compile(&mut ctx, &chunk, key)?;
430        let trace = ctx
431            .cache
432            .get(&key)
433            .ok_or("trace missing after force_compile")?;
434        let mut slots = vec![0i64; 3];
435        let ret = run_compiled_trace_ref(trace, &mut slots, 0, 10_000, 0, 3);
436        let exit = ret.into_reason();
437        assert!(matches!(exit, ExitReason::Return { return_reg: 2 }));
438        assert_eq!(slots[2], 42);
439        Ok(())
440    }
441
442    #[test]
443    fn jit_matches_interpreter_on_scalar_add() -> TestResult {
444        let chunk = add_chunk();
445        let rt = Runtime::new((*chunk).clone())?;
446        let interp = rt.spawn(0, &[])?.join();
447        rt.shutdown();
448        assert!(matches!(interp, FlowOutcome::Completed(Value::Int(42))));
449
450        let mut ctx = JitContext::new(chunk.clone())?;
451        let key = TraceKey {
452            function: 0,
453            entry_pc: 0,
454        };
455        let mut slots = vec![0i64; 3];
456        force_compile(&mut ctx, &chunk, key)?;
457        let trace = ctx
458            .cache
459            .get(&key)
460            .ok_or("trace missing after force_compile")?;
461        let exit = run_compiled_trace_ref(trace, &mut slots, 0, 10_000, 0, 3).into_reason();
462        assert!(matches!(exit, ExitReason::Return { return_reg: 2 }));
463        assert_eq!(slots[2], 42);
464
465        let mut vm = Vm::new(chunk, NativeTable::empty(), 0, &[])?;
466        if let Some(result) = apply_exit_to_vm(&mut vm, &slots, exit)? {
467            assert!(matches!(result, VmResult::Complete(Value::Int(42))));
468        } else {
469            return Err("expected Completed(42)".into());
470        }
471        Ok(())
472    }
473
474    #[test]
475    fn div_by_zero_returns_trap_exit() -> TestResult {
476        let mut b = ChunkBuilder::new("jit-div0");
477        b.begin_function("main", 0, 3);
478        b.emit_load_imm(0, 1);
479        b.emit_load_imm(1, 0);
480        b.emit_binop(Opcode::Div, 2, 0, 1);
481        b.emit_return(2);
482        let chunk = Arc::new(b.finish());
483
484        let mut ctx = JitContext::new(chunk.clone())?;
485        let key = TraceKey {
486            function: 0,
487            entry_pc: 0,
488        };
489        force_compile(&mut ctx, &chunk, key)?;
490        let trace = ctx
491            .cache
492            .get(&key)
493            .ok_or("trace missing after force_compile")?;
494        let mut slots = vec![0i64; 3];
495        let ret = run_compiled_trace_ref(trace, &mut slots, 0, 10_000, 0, 3);
496        assert!(matches!(ret.into_reason(), ExitReason::Trap { .. }));
497        Ok(())
498    }
499
500    #[test]
501    fn sync_slots_deopts_on_float() -> TestResult {
502        let chunk = add_chunk();
503        let mut vm = Vm::new(chunk, NativeTable::empty(), 0, &[])?;
504        vm.set_register(0, Value::Float(1.0))?;
505        let mut slots = vec![0i64; 3];
506        assert_eq!(sync_slots_from_vm(&vm, &mut slots), SyncSlotsResult::Deopt);
507        Ok(())
508    }
509
510    #[test]
511    fn runtime_path_compiles_and_runs() -> TestResult {
512        let chunk = add_chunk();
513        let runtime = JitRuntime::new(chunk.clone(), 1);
514        let key = TraceKey {
515            function: 0,
516            entry_pc: 0,
517        };
518        let mut slots = vec![0i64; 3];
519        let exit = try_run_hot_runtime(&runtime, &chunk, key, &mut slots, 10_000, 3, None)
520            .ok_or("expected jit execution")?;
521        assert!(matches!(exit, ExitReason::Return { return_reg: 2 }));
522        assert_eq!(slots[2], 42);
523        Ok(())
524    }
525
526    #[test]
527    fn budget_exit_when_exhausted_mid_trace() -> TestResult {
528        let mut b = ChunkBuilder::new("jit-budget");
529        b.begin_function("main", 0, 1);
530        b.emit_load_imm(0, 1);
531        b.emit_load_imm(0, 2);
532        b.emit_load_imm(0, 3);
533        b.emit_return(0);
534        let chunk = Arc::new(b.finish());
535
536        let mut ctx = JitContext::new(chunk.clone())?;
537        let key = TraceKey {
538            function: 0,
539            entry_pc: 0,
540        };
541        force_compile(&mut ctx, &chunk, key)?;
542        let trace = ctx
543            .cache
544            .get(&key)
545            .ok_or("trace missing after force_compile")?;
546        let mut slots = vec![0i64; 1];
547        let ret = run_compiled_trace_ref(trace, &mut slots, 0, 2, 0, 1);
548        assert!(matches!(ret.into_reason(), ExitReason::Budget { pc: 2 }));
549        Ok(())
550    }
551
552    #[test]
553    fn interleaved_jit_compiles_mid_quantum() -> TestResult {
554        let mut b = ChunkBuilder::new("jit-interleave");
555        b.begin_function("main", 0, 2);
556        b.emit_load_imm(0, 0);
557        let head = b.new_label();
558        b.bind_label(head);
559        b.emit_load_imm(1, 1);
560        b.emit_binop(Opcode::Add, 0, 0, 1);
561        b.emit_load_imm(1, 100);
562        b.emit_binop(Opcode::Lt, 1, 0, 1);
563        let done = b.new_label();
564        b.emit_branch(1, done);
565        b.emit_jump(head);
566        b.bind_label(done);
567        b.emit_return(0);
568        let chunk = Arc::new(b.finish());
569        let head_pc = chunk.functions[0].entry + 1;
570
571        let mut ctx = JitContext::new(chunk.clone())?;
572        ctx.hot_threshold = 2;
573        let mut vm = Vm::new(chunk, NativeTable::empty(), 0, &[])?;
574        let result = run_vm_with_jit(&mut vm, 10_000, &mut ctx);
575        assert!(matches!(result, VmResult::Complete(Value::Int(100))));
576        assert!(ctx
577            .cache
578            .get(&TraceKey {
579                function: 0,
580                entry_pc: head_pc,
581            })
582            .is_some());
583        Ok(())
584    }
585
586    #[test]
587    fn reload_chunk_invalidates_jit_cache() -> TestResult {
588        let chunk = add_chunk();
589        let runtime = JitRuntime::new(chunk.clone(), 1);
590        let key = TraceKey {
591            function: 0,
592            entry_pc: 0,
593        };
594        let mut slots = vec![0i64; 3];
595        let exit = try_run_hot_runtime(&runtime, &chunk, key, &mut slots, 10_000, 3, None)
596            .ok_or("expected jit execution")?;
597        assert!(matches!(exit, ExitReason::Return { .. }));
598
599        let mut b = ChunkBuilder::new("other");
600        b.begin_function("main", 0, 1);
601        b.emit_load_imm(0, 1);
602        b.emit_return(0);
603        runtime.reload(Arc::new(b.finish()));
604        assert!(runtime.get_trace(&key).is_none());
605        Ok(())
606    }
607}
byteflow/jit/dispatch.rs

byteflow/jit/
dispatch.rs
