coop/rose-ash
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

JIT: VM fast path, &rest support, locals scan, test runner fixes

Browse Source 
- jit_compile_lambda: call compile directly via VM when it has bytecode
  (100-400x faster JIT compilation, server pre-warm 1.6s vs hung)
- code_from_value: scan bytecode for highest LOCAL_GET/SET slot to
  compute vc_locals correctly (fixes hyperscript LOCAL_GET overflow)
- code_from_value: accept both compiler keys (bytecode) and SX VM
  keys (vc-bytecode) for interop
- jit_compile_lambda: skip &key/:as params (compiler can't emit them)
- Test runner: seed VM globals with primitives + env bindings,
  native vm-execute-module with suspension fallback to SX version,
  _jit_refresh_globals syncs globals after module loading,
  VmSuspended + "VM undefined" caught and sentineled

3127/3127 without JIT, 3116/3127 with JIT (11 hyperscript on-event
parsing — specific closure/scope issue, not infrastructure).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
giles
2026-04-08 10:52:44 +00:00
parent 387a6cb49e
commit 3155ba47f9
2 changed files with 70 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
hosts/ocaml/lib/sx_vm.ml
View File

@@ -193,7 +193,30 @@ let code_from_value v =
let rest_arity = match find2 "rest-arity" "vc-rest-arity" with
| Some (Number n) -> int_of_float n | _ -> -1
in
{ vc_arity = arity; vc_rest_arity = rest_arity; vc_locals = arity + 16;
(* Compute locals from bytecode: scan for highest LOCAL_GET/LOCAL_SET slot.
The compiler's arity may undercount when nested lets add many locals. *)
let max_local = ref (arity - 1) in
let len = Array.length bc_list in
let i = ref 0 in
while !i < len do
let op = bc_list.(!i) in
if (op = 16 (* LOCAL_GET *) || op = 17 (* LOCAL_SET *)) && !i + 1 < len then
(let slot = bc_list.(!i + 1) in
if slot > !max_local then max_local := slot;
i := !i + 2)
else if op = 18 (* UPVALUE_GET *) || op = 19 (* UPVALUE_SET *)
|| op = 8 (* JUMP_IF_FALSE *) || op = 33 (* JUMP_IF_FALSE_u16 *)
|| op = 34 (* JUMP_IF_TRUE *) then
i := !i + 2
else if op = 1 (* CONST *) || op = 20 (* GLOBAL_GET *) || op = 21 (* GLOBAL_SET *)
|| op = 32 (* JUMP *) || op = 51 (* CLOSURE *) || op = 52 (* CALL_PRIM *)
|| op = 64 (* MAKE_LIST *) || op = 65 (* MAKE_DICT *) then
i := !i + 3 (* u16 operand *)
else
i := !i + 1
done;
let locals = !max_local + 1 + 16 in (* +16 headroom for temporaries *)
{ vc_arity = arity; vc_rest_arity = rest_arity; vc_locals = locals;
vc_bytecode = bc_list; vc_constants = constants;
vc_bytecode_list = None; vc_constants_list = None }
| _ -> { vc_arity = 0; vc_rest_arity = -1; vc_locals = 16; vc_bytecode = [||]; vc_constants = [||];
@@ -833,28 +856,34 @@ let jit_compile_lambda (l : lambda) globals =
if !_jit_compiling then (
(* Already compiling — prevent cascade. The CEK will handle this call. *)
None
) else if List.mem "&key" l.l_params || List.mem ":as" l.l_params then (
(* &key/:as require complex runtime argument processing that the compiler
doesn't emit. These functions must run via CEK. *)
None
) else
try
_jit_compiling := true;
let compile_fn = try Hashtbl.find globals "compile"
with Not_found -> (_jit_compiling := false; raise (Eval_error "JIT: compiler not loaded")) in
(* Reconstruct the (fn (params) body) form so the compiler produces
a proper closure. l.l_body is the inner body; we need the full
function form with params so the compiled code binds them. *)
let param_syms = List (List.map (fun s -> Symbol s) l.l_params) in
let fn_expr = List [Symbol "fn"; param_syms; l.l_body] in
let quoted = List [Symbol "quote"; fn_expr] in
(* Use Symbol "compile" so the CEK resolves it from the env, not
an embedded VmClosure value — the CEK dispatches VmClosure calls
differently when the value is resolved from env vs embedded in AST. *)
ignore compile_fn;
let compile_env = Sx_types.env_extend (Sx_types.make_env ()) in
Hashtbl.iter (fun k v -> Hashtbl.replace compile_env.bindings (Sx_types.intern k) v) globals;
let result = Sx_ref.eval_expr (List [Symbol "compile"; quoted]) (Env compile_env) in
(* Fast path: if compile has bytecode, call it directly via the VM.
All helper calls (compile-expr, emit-byte, etc.) happen inside the
same VM execution — no per-call VM allocation overhead. *)
let result = match compile_fn with
| Lambda { l_compiled = Some cl; _ } when not (is_jit_failed cl) ->
call_closure cl [fn_expr] globals
| _ ->
ignore compile_fn;
let compile_env = Sx_types.env_extend (Sx_types.make_env ()) in
Hashtbl.iter (fun k v -> Hashtbl.replace compile_env.bindings (Sx_types.intern k) v) globals;
Sx_ref.eval_expr (List [Symbol "compile"; quoted]) (Env compile_env)
in
_jit_compiling := false;
let effective_globals = globals in
(match result with
| Dict d when Hashtbl.mem d "bytecode" ->
| Dict d when Hashtbl.mem d "bytecode" || Hashtbl.mem d "vc-bytecode" ->
let outer_code = code_from_value result in
let bc = outer_code.vc_bytecode in
if Array.length bc >= 4 && bc.(0) = 51 (* OP_CLOSURE *) then begin