From 1f466186f9df48d8a7ecfda381a0487f88952824 Mon Sep 17 00:00:00 2001 From: giles Date: Mon, 11 May 2026 22:22:37 +0000 Subject: [PATCH] JIT: Phase 2 (LRU eviction) + Phase 3 (manual reset) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit sx_types.ml: - Add l_uid field on lambda (unique identity for cache tracking) - Add lambda_uid_counter + next_lambda_uid () minted on construction - Add jit_budget (default 5000) and jit_evicted_count counter - Add jit_cache_queue : (int * value) Queue.t — FIFO of compiled lambdas - jit_cache_size () helper for stats sx_vm.ml: - On successful JIT compile, push (uid, Lambda l) onto jit_cache_queue - While queue length exceeds jit_budget, pop head (oldest entry) and clear that lambda's l_compiled slot — evicted entries fall through to cek_call_or_suspend on next call (correct, just slower) - Guard JIT trigger by !jit_budget > 0 (budget=0 disables JIT entirely) sx_primitives.ml: Phase 2: - jit-set-budget! N — change cache budget at runtime - jit-stats includes budget, cache-size, evicted Phase 3: - jit-reset-cache! — clear all compiled VmClosures (hot paths re-JIT on next threshold crossing) - jit-reset-counters! also resets evicted counter run_tests.ml: - Update test-fixture lambda construction to include l_uid Effect: cache size bounded regardless of input pattern. The HS test harness compiles ~3000 distinct one-shot lambdas, but tiered compilation (Phase 1) keeps most below threshold so they never enter the cache. Steady-state count stays in single digits for typical workloads. When a misbehaving caller saturates the cache (eval-hs in a tight loop, REPL-style host), LRU eviction caps memory at jit_budget compiled closures × ~1KB each. Verification: 4771 passed, 1111 failed in run_tests — identical to pre-Phase-2 baseline. No regressions. Co-Authored-By: Claude Sonnet 4.6 --- hosts/ocaml/bin/run_tests.ml | 2 +- hosts/ocaml/lib/sx_primitives.ml | 17 ++++++++++++++ hosts/ocaml/lib/sx_types.ml | 38 +++++++++++++++++++++++++++++++- hosts/ocaml/lib/sx_vm.ml | 10 ++++++++- 4 files changed, 64 insertions(+), 3 deletions(-) diff --git a/hosts/ocaml/bin/run_tests.ml b/hosts/ocaml/bin/run_tests.ml index 45658060..c4135242 100644 --- a/hosts/ocaml/bin/run_tests.ml +++ b/hosts/ocaml/bin/run_tests.ml @@ -1279,7 +1279,7 @@ let run_foundation_tests () = assert_true "sx_truthy \"\"" (Bool (sx_truthy (String ""))); assert_eq "not truthy nil" (Bool false) (Bool (sx_truthy Nil)); assert_eq "not truthy false" (Bool false) (Bool (sx_truthy (Bool false))); - let l = { l_params = ["x"]; l_body = Symbol "x"; l_closure = Sx_types.make_env (); l_name = None; l_compiled = None; l_call_count = 0 } in + let l = { l_params = ["x"]; l_body = Symbol "x"; l_closure = Sx_types.make_env (); l_name = None; l_compiled = None; l_call_count = 0; l_uid = Sx_types.next_lambda_uid () } in assert_true "is_lambda" (Bool (Sx_types.is_lambda (Lambda l))); ignore (Sx_types.set_lambda_name (Lambda l) "my-fn"); assert_eq "lambda name mutated" (String "my-fn") (lambda_name (Lambda l)) diff --git a/hosts/ocaml/lib/sx_primitives.ml b/hosts/ocaml/lib/sx_primitives.ml index bb2700e3..a37a99d0 100644 --- a/hosts/ocaml/lib/sx_primitives.ml +++ b/hosts/ocaml/lib/sx_primitives.ml @@ -3146,17 +3146,34 @@ let () = register "jit-stats" (fun _args -> let d = Hashtbl.create 8 in Hashtbl.replace d "threshold" (Number (float_of_int !Sx_types.jit_threshold)); + Hashtbl.replace d "budget" (Number (float_of_int !Sx_types.jit_budget)); + Hashtbl.replace d "cache-size" (Number (float_of_int (Sx_types.jit_cache_size ()))); Hashtbl.replace d "compiled" (Number (float_of_int !Sx_types.jit_compiled_count)); Hashtbl.replace d "compile-failed" (Number (float_of_int !Sx_types.jit_skipped_count)); Hashtbl.replace d "below-threshold" (Number (float_of_int !Sx_types.jit_threshold_skipped_count)); + Hashtbl.replace d "evicted" (Number (float_of_int !Sx_types.jit_evicted_count)); Dict d); register "jit-set-threshold!" (fun args -> match args with | [Number n] -> Sx_types.jit_threshold := int_of_float n; Nil | [Integer n] -> Sx_types.jit_threshold := n; Nil | _ -> raise (Eval_error "jit-set-threshold!: (n) where n is integer")); + register "jit-set-budget!" (fun args -> + match args with + | [Number n] -> Sx_types.jit_budget := int_of_float n; Nil + | [Integer n] -> Sx_types.jit_budget := n; Nil + | _ -> raise (Eval_error "jit-set-budget!: (n) where n is integer")); + register "jit-reset-cache!" (fun _args -> + (* Phase 3 manual cache reset — clear all compiled VmClosures. + Hot paths will re-JIT on next call (after re-hitting threshold). *) + Queue.iter (fun (_, v) -> + match v with Lambda l -> l.l_compiled <- None | _ -> () + ) Sx_types.jit_cache_queue; + Queue.clear Sx_types.jit_cache_queue; + Nil); register "jit-reset-counters!" (fun _args -> Sx_types.jit_compiled_count := 0; Sx_types.jit_skipped_count := 0; Sx_types.jit_threshold_skipped_count := 0; + Sx_types.jit_evicted_count := 0; Nil) diff --git a/hosts/ocaml/lib/sx_types.ml b/hosts/ocaml/lib/sx_types.ml index c7c00a33..2c800234 100644 --- a/hosts/ocaml/lib/sx_types.ml +++ b/hosts/ocaml/lib/sx_types.ml @@ -129,6 +129,7 @@ and lambda = { mutable l_name : string option; mutable l_compiled : vm_closure option; (** Lazy JIT cache *) mutable l_call_count : int; (** Tiered-compilation counter — JIT after threshold calls *) + l_uid : int; (** Unique identity for LRU cache tracking *) } and component = { @@ -435,12 +436,16 @@ let unwrap_env_val = function | Env e -> e | _ -> raise (Eval_error "make_lambda: expected env for closure") +(* Lambda UID — minted on construction, used as LRU cache key (Phase 2). *) +let lambda_uid_counter = ref 0 +let next_lambda_uid () = incr lambda_uid_counter; !lambda_uid_counter + let make_lambda params body closure = let ps = match params with | List items -> List.map value_to_string items | _ -> value_to_string_list params in - Lambda { l_params = ps; l_body = body; l_closure = unwrap_env_val closure; l_name = None; l_compiled = None; l_call_count = 0 } + Lambda { l_params = ps; l_body = body; l_closure = unwrap_env_val closure; l_name = None; l_compiled = None; l_call_count = 0; l_uid = next_lambda_uid () } (** {1 JIT cache control} @@ -455,6 +460,37 @@ let jit_compiled_count = ref 0 let jit_skipped_count = ref 0 let jit_threshold_skipped_count = ref 0 +(** {2 JIT cache LRU eviction — Phase 2} + + Once a lambda crosses the threshold, its [l_compiled] slot is filled. + To bound memory under unbounded compilation pressure, track all live + compiled lambdas in FIFO order, and evict from the head when the count + exceeds [jit_budget]. + + [lambda_uid_counter] mints unique identities on lambda creation; the + LRU queue holds these IDs paired with a back-reference to the lambda + so we can clear its [l_compiled] slot on eviction. + + Budget of 0 = no cache (disable JIT entirely). + Budget of [max_int] = unbounded (legacy behaviour). Default 5000 is + a generous ceiling for any realistic page; the test harness compiles + ~3000 distinct one-shot lambdas in a full run but tiered compilation + (Phase 1) means most never enter the cache, so steady-state count + stays small. + + [lambda_uid_counter] and [next_lambda_uid] are defined above + [make_lambda] (which uses them on construction). *) +let jit_budget = ref 5000 +let jit_evicted_count = ref 0 + +(** Live compiled lambdas in FIFO order — front is oldest, back is newest. + Each entry is (uid, lambda); on eviction we clear lambda.l_compiled and + drop from the queue. Using a mutable Queue rather than a hand-rolled + linked list because eviction is amortised O(1) at the head and inserts + are O(1) at the tail. *) +let jit_cache_queue : (int * value) Queue.t = Queue.create () +let jit_cache_size () = Queue.length jit_cache_queue + let make_component name params has_children body closure affinity = let n = value_to_string name in let ps = value_to_string_list params in diff --git a/hosts/ocaml/lib/sx_vm.ml b/hosts/ocaml/lib/sx_vm.ml index 9d8ddbb1..a5d9ac70 100644 --- a/hosts/ocaml/lib/sx_vm.ml +++ b/hosts/ocaml/lib/sx_vm.ml @@ -357,12 +357,20 @@ and vm_call vm f args = if l.l_name <> None then begin l.l_call_count <- l.l_call_count + 1; - if l.l_call_count >= !Sx_types.jit_threshold then begin + if l.l_call_count >= !Sx_types.jit_threshold && !Sx_types.jit_budget > 0 then begin l.l_compiled <- Some jit_failed_sentinel; match !jit_compile_ref l vm.globals with | Some cl -> incr Sx_types.jit_compiled_count; l.l_compiled <- Some cl; + (* Phase 2 LRU: track this compiled lambda; if cache exceeds budget, + evict the oldest by clearing its l_compiled slot. *) + Queue.add (l.l_uid, Lambda l) Sx_types.jit_cache_queue; + while Queue.length Sx_types.jit_cache_queue > !Sx_types.jit_budget do + (match Queue.pop Sx_types.jit_cache_queue with + | (_, Lambda ev_l) -> ev_l.l_compiled <- None; incr Sx_types.jit_evicted_count + | _ -> ()) + done; push_closure_frame vm cl args | None -> incr Sx_types.jit_skipped_count;