From 52c1aa6bc8081ca729de647419ce01ad846a5d04 Mon Sep 17 00:00:00 2001 From: Christoph Knittel Date: Sat, 7 Nov 2026 10:04:32 +0000 Subject: [PATCH 1/2] Replace local mutable blocks in one pass over the function MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Why: The pass analyzed and rewrote a block's whole scope once per block, so n refs used in one function cost O(n²): 1,000 refs updated in one loop allocated 151M words. One traversal now decides every block and one rewrite replaces them, so the same input allocates 17M. A value written to, or initializing, a never-read field is held back until its block is decided, as the old innermost-first order dropped such values before analyzing the blocks around them. The generated JavaScript is the same for the test suites and 1,600 generated functions. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Kbf3gK1dcjbe97yY1wRMgs Signed-off-by: Christoph Knittel --- changes/linear-sroa.polish.md | 1 + compiler/core/lam_pass_sroa.ml | 280 +++++++++++++++++++++++---------- tests/tests/src/sroa_test.mjs | 48 ++++-- tests/tests/src/sroa_test.res | 36 +++++ 4 files changed, 275 insertions(+), 90 deletions(-) create mode 100644 changes/linear-sroa.polish.md diff --git a/changes/linear-sroa.polish.md b/changes/linear-sroa.polish.md new file mode 100644 index 00000000000..9f84537c298 --- /dev/null +++ b/changes/linear-sroa.polish.md @@ -0,0 +1 @@ +- Compile functions with many local `ref`s or mutable records in linear time. diff --git a/compiler/core/lam_pass_sroa.ml b/compiler/core/lam_pass_sroa.ml index e5deff7fe9d..cffbaa06fc3 100644 --- a/compiler/core/lam_pass_sroa.ml +++ b/compiler/core/lam_pass_sroa.ml @@ -21,56 +21,47 @@ type field_use = {mutable read: bool; mutable written: bool} -let valid_field uses index = index >= 0 && index < Array.length uses +(* A local mutable block, with how its fields are used so far. [order] + numbers blocks in the order of their declarations. [fields] is set when + the block's scope has been analyzed and every occurrence was a direct, + in-range field read or write. *) +type candidate = { + order: int; + uses: field_use array; + mutable escapes: bool; + mutable fields: Ident.t array option; + mutable slots: slot list; +} -(* Does the block appear only as direct, in-range field reads and writes? While - answering, record how every field is used. [analyze] and [rewrite] below are - a matched pair: [rewrite] handles exactly the occurrences [analyze] accepts, - and asserts on the rest. Extending one without the other is a compiler crash - rather than a type error, so keep their cases in step. *) -let rec analyze block uses (lam : Lambda.t) = - match lam with - | Lvar id -> not (Ident.same id block) - | Lassign (id, value) -> - (not (Ident.same id block)) && analyze block uses value - | Lprim {primitive = Pfield (index, _); args = [Lvar id]} - when Ident.same id block -> - if valid_field uses index then ( - uses.(index).read <- true; - true) - else false - | Lprim {primitive = Psetfield (index, _); args = [Lvar id; value]} - when Ident.same id block -> - if valid_field uses index then ( - uses.(index).written <- true; - analyze block uses value) - else false - | _ -> - not - (Lambda_traverse.shallow_exists - (fun child -> not (analyze block uses child)) - lam) +(* An initializer of a field of [owner], or a value written to it, without + side occurrences. If the field is never read, the rewrite drops the value, and + the occurrences of blocks declared before [owner] in it don't count: + [occurrences] holds them until [owner] is decided. [parent] is the slot that + contains this one. *) +and slot = { + owner: candidate; + field: int; + parent: slot option; + mutable occurrences: occurrence list; +} + +and occurrence = Escape of candidate | Read of candidate * int + +let valid_field uses index = index >= 0 && index < Array.length uses -let discard_value value body = - if Lam_analysis.no_side_effects value then body else Lambda.seq value body +let is_candidate kind (init : Lambda.t) = + match (kind, init) with + | ( (Lambda.Strict | StrictOpt), + Lprim {primitive = Pmakeblock info; args = _ :: _} ) -> + not (Lambda.is_immutable_block info) + | _ -> false -let rec rewrite block fields uses (lam : Lambda.t) = - match lam with - | Lprim {primitive = Pfield (index, _); args = [Lvar id]} - when Ident.same id block -> - Lambda.var fields.(index) - | Lprim {primitive = Psetfield (index, _); args = [Lvar id; value]} - when Ident.same id block -> - let value = rewrite block fields uses value in - if not uses.(index).read then discard_value value Lambda.lambda_unit - else Lambda.assign fields.(index) value - (* Unreachable: [analyze] rejected the block for both of these, so [replace] - never reaches the rewrite. They are kept as assertions rather than dropped - so that a future occurrence form added to [analyze] but not here fails - loudly instead of silently losing the write. *) - | Lvar id when Ident.same id block -> assert false - | Lassign (id, _) when Ident.same id block -> assert false - | _ -> Lambda_traverse.shallow_map_sharing (rewrite block fields uses) lam +(* Drop [value], a value of a field that is never read, keeping it only for + its side effects, before [body]. [analyze] doesn't count the occurrences in + the values dropped here, so they are not rewritten. *) +let discard_value rewrite value body = + if Lam_analysis.no_side_effects value then body + else Lambda.seq (rewrite value) body let fields_for_block block info field_count = let fallback () = @@ -97,18 +88,133 @@ let fields_for_block block info field_count = | Some names -> Array.map (fun name -> Ident.create (Ident.name block ^ "_" ^ name)) names -let replace ~block ~info ~initializers body = - match initializers with - | [] -> None - | _ -> - let field_count = List.length initializers in - let uses = - Array.init field_count (fun _ -> {read = false; written = false}) +(* Find the blocks whose every occurrence is a direct, in-range field read or + write, recording how each field is used, in one traversal of the term. + [rewrite] below handles exactly the occurrences accepted here, and asserts + on the rest; keep their cases in step. + + The result is the one of analyzing each block's scope after rewriting the + blocks declared in it, as if blocks were replaced one at a time from the + innermost out, so pending [slot]s keep the occurrences that such a rewrite + would drop. A block is decided when its scope has been traversed, so inner + blocks are decided (and their fields named) before outer ones. *) +let analyze ~considered (lam : Lambda.t) = + let candidates = Hash_ident.create 16 in + let count = ref 0 in + let apply = function + | Escape c -> c.escapes <- true + | Read (c, index) -> c.uses.(index).read <- true + in + let rec record slot occurrence = + let target = + match occurrence with + | Escape c | Read (c, _) -> c in - if not (analyze block uses body) then None - else - let fields = fields_for_block block info field_count in - let body = rewrite block fields uses body in + match slot with + | Some slot when target.order < slot.owner.order -> + slot.occurrences <- occurrence :: slot.occurrences + | _ -> apply occurrence + and decide c = + if not c.escapes then ( + List.iter + (fun slot -> + if c.uses.(slot.field).read then + List.iter (record slot.parent) slot.occurrences) + c.slots; + true) + else ( + List.iter + (fun slot -> List.iter (record slot.parent) slot.occurrences) + c.slots; + false) + in + (* [slot]: the innermost pending slot containing [lam]. *) + let rec go slot (lam : Lambda.t) = + match lam with + | Lvar id -> ( + match Hash_ident.find_opt candidates id with + | Some c -> record slot (Escape c) + | None -> ()) + | Lassign (id, value) -> + (match Hash_ident.find_opt candidates id with + | Some c -> record slot (Escape c) + | None -> ()); + go slot value + | Lprim {primitive = Pfield (index, _); args = [Lvar id]} + when Hash_ident.mem candidates id -> + let c = Hash_ident.find candidates id in + if valid_field c.uses index then record slot (Read (c, index)) + else record slot (Escape c) + | Lprim {primitive = Psetfield (index, _); args = [Lvar id; value]} + when Hash_ident.mem candidates id -> + let c = Hash_ident.find candidates id in + if valid_field c.uses index then ( + c.uses.(index).written <- true; + go (value_slot slot c index value) value) + else ( + record slot (Escape c); + go slot value) + | Llet (kind, block, init, body) + when is_candidate kind init && considered block -> ( + match init with + | Lprim {primitive = Pmakeblock info; args = initializers} -> + let field_count = List.length initializers in + let c = + { + order = !count; + uses = + Array.init field_count (fun _ -> {read = false; written = false}); + escapes = false; + fields = None; + slots = []; + } + in + incr count; + List.iteri + (fun index init -> go (value_slot slot c index init) init) + initializers; + Hash_ident.add candidates block c; + go slot body; + if decide c then + c.fields <- Some (fields_for_block block info field_count) + | _ -> assert false) + | _ -> Lambda_traverse.iter (go slot) lam + (* The slot for a value of field [index] of [c], if the rewrite drops the + value when the field is never read. *) + and value_slot slot c index value = + if Lam_analysis.no_side_effects value then ( + let s = {owner = c; field = index; parent = slot; occurrences = []} in + c.slots <- s :: c.slots; + Some s) + else slot + in + go None lam; + candidates + +let rewrite candidates (lam : Lambda.t) = + let replaced id = + match Hash_ident.find_opt candidates id with + | Some {fields = Some fields; uses} -> Some (fields, uses) + | _ -> None + in + let rec rewrite (lam : Lambda.t) = + match lam with + | Lprim {primitive = Pfield (index, _); args = [Lvar id]} -> ( + match replaced id with + | Some (fields, _) -> Lambda.var fields.(index) + | None -> lam) + | Lprim {primitive = Psetfield (index, _); args = [Lvar id; value]} -> ( + match replaced id with + | Some (fields, uses) -> + if not uses.(index).read then + discard_value rewrite value Lambda.lambda_unit + else Lambda.assign fields.(index) (rewrite value) + | None -> Lambda_traverse.shallow_map_sharing rewrite lam) + | Llet + (_, block, Lprim {primitive = Pmakeblock _; args = initializers}, body) + when replaced block <> None -> + let fields, uses = Option.get (replaced block) in + let body = rewrite body in let rec bind_fields index initializers body = match initializers with | [] -> body @@ -119,28 +225,42 @@ let replace ~block ~info ~initializers body = retained only when they have effects. A read-only field can use a normal refined let, while a field that is both read and written still needs a mutable scalar binding. *) - if not use.read then discard_value init body + if not use.read then discard_value rewrite init body else if not use.written then - Lam_util.refine_let ~kind:Strict fields.(index) init body - else Lambda.let_ Variable fields.(index) init body + Lam_util.refine_let ~kind:Strict fields.(index) (rewrite init) body + else Lambda.let_ Variable fields.(index) (rewrite init) body in - Some (bind_fields 0 initializers body) + bind_fields 0 initializers body + (* Unreachable: [analyze] rejected the block for both of these. They are + kept as assertions rather than dropped so that a future occurrence form + added to [analyze] but not here fails loudly instead of silently losing + the write. *) + | Lvar id when replaced id <> None -> assert false + | Lassign (id, _) when replaced id <> None -> assert false + | _ -> Lambda_traverse.shallow_map_sharing rewrite lam + in + rewrite lam + +(* Analyzing and rewriting each block separately would traverse its whole + scope once per block, which is quadratic for a function with many local + references. *) +let replace ~block ~info ~initializers body = + let init = + Lambda.prim ~primitive:(Pmakeblock info) ~args:initializers Location.none + in + if not (is_candidate Strict init) then None + else + let lam = Lambda.let_ Strict block init body in + let candidates = analyze ~considered:(Ident.same block) lam in + match Hash_ident.find_opt candidates block with + | Some {fields = Some _} -> Some (rewrite candidates lam) + | _ -> None -let rec simplify (lam : Lambda.t) = - match lam with - | Llet (kind, block, init, body) -> ( - let init' = simplify init in - let body' = simplify body in - match (kind, init') with - | ( (Strict | StrictOpt), - Lambda.Lprim {primitive = Pmakeblock info; args = initializers} ) - when not (Lambda.is_immutable_block info) -> ( - match replace ~block ~info ~initializers body' with - | Some replacement -> replacement - | None -> - if init' == init && body' == body then lam - else Lambda.let_ kind block init' body') - | _ -> - if init' == init && body' == body then lam - else Lambda.let_ kind block init' body') - | _ -> Lambda_traverse.shallow_map_sharing simplify lam +let simplify (lam : Lambda.t) = + let candidates = analyze ~considered:(fun _ -> true) lam in + if + Hash_ident.fold + (fun _ c found -> found || c.fields <> None) + candidates false + then rewrite candidates lam + else lam diff --git a/tests/tests/src/sroa_test.mjs b/tests/tests/src/sroa_test.mjs index eaaa2ff8cfb..cc7df550980 100644 --- a/tests/tests/src/sroa_test.mjs +++ b/tests/tests/src/sroa_test.mjs @@ -116,21 +116,43 @@ function writeOnlyFieldStillRaises() { return cell_live; } +function storedInWriteOnly() { + let cell = 1; + return cell; +} + +function readByWriteOnly() { + let value = 1; + return value; +} + +function manyInLoop(n) { + let a = 0; + let b = 1; + let c = 2; + for (let i = 0; i <= n; ++i) { + a = a + i | 0; + b = b + a | 0; + c = c + b | 0; + } + return (a + b | 0) + c | 0; +} + Mocha.describe("Sroa_test", () => { - Mocha.test("scalarizes a local mutable record", () => Test_utils.eq("File \"sroa_test.res\", line 123, characters 53-60", 42, localPair())); - Mocha.test("shares scalar fields with closures", () => Test_utils.eq("File \"sroa_test.res\", line 124, characters 54-61", 32, capturedPair())); - Mocha.test("preserves initializer order", () => Test_utils.eq("File \"sroa_test.res\", line 125, characters 47-54", [ + Mocha.test("scalarizes a local mutable record", () => Test_utils.eq("File \"sroa_test.res\", line 156, characters 53-60", 42, localPair())); + Mocha.test("shares scalar fields with closures", () => Test_utils.eq("File \"sroa_test.res\", line 157, characters 54-61", 32, capturedPair())); + Mocha.test("preserves initializer order", () => Test_utils.eq("File \"sroa_test.res\", line 158, characters 47-54", [ 4, [ 1, 2 ] ], initializationOrder())); - Mocha.test("retains an escaping record", () => Test_utils.eq("File \"sroa_test.res\", line 126, characters 46-53", 30, consumePair({ + Mocha.test("retains an escaping record", () => Test_utils.eq("File \"sroa_test.res\", line 159, characters 46-53", 30, consumePair({ left: 10, right: 20 }))); - Mocha.test("cleans up fields according to their uses", () => Test_utils.eq("File \"sroa_test.res\", line 128, characters 7-14", [ + Mocha.test("cleans up fields according to their uses", () => Test_utils.eq("File \"sroa_test.res\", line 161, characters 7-14", [ 5, [ 2, @@ -140,20 +162,23 @@ Mocha.describe("Sroa_test", () => { 6 ] ], fieldUseCleanup())); - Mocha.test("preserves overwritten initializer effects", () => Test_utils.eq("File \"sroa_test.res\", line 131, characters 7-14", [ + Mocha.test("preserves overwritten initializer effects", () => Test_utils.eq("File \"sroa_test.res\", line 164, characters 7-14", [ 2, [ 1, 2 ] ], overwrittenBeforeRead())); - Mocha.test("removes write-only fields captured by closures", () => Test_utils.eq("File \"sroa_test.res\", line 134, characters 7-14", [ + Mocha.test("removes write-only fields captured by closures", () => Test_utils.eq("File \"sroa_test.res\", line 167, characters 7-14", [ 1, 2 ], capturedWriteOnly())); - Mocha.test("does not evaluate writes in uncalled closures", () => Test_utils.eq("File \"sroa_test.res\", line 137, characters 7-14", [1], uncalledWriteOnlyClosure())); - Mocha.test("read-only fields snapshot their initializer", () => Test_utils.eq("File \"sroa_test.res\", line 140, characters 7-14", 1, readOnlyFieldSnapshotsInitializer())); - Mocha.test("preserves exceptions from write-only field values", () => Test_utils.throws("File \"sroa_test.res\", line 143, characters 11-18", writeOnlyFieldStillRaises)); + Mocha.test("does not evaluate writes in uncalled closures", () => Test_utils.eq("File \"sroa_test.res\", line 170, characters 7-14", [1], uncalledWriteOnlyClosure())); + Mocha.test("read-only fields snapshot their initializer", () => Test_utils.eq("File \"sroa_test.res\", line 173, characters 7-14", 1, readOnlyFieldSnapshotsInitializer())); + Mocha.test("preserves exceptions from write-only field values", () => Test_utils.throws("File \"sroa_test.res\", line 176, characters 11-18", writeOnlyFieldStillRaises)); + Mocha.test("scalarizes a block stored in a write-only field", () => Test_utils.eq("File \"sroa_test.res\", line 178, characters 67-74", 1, storedInWriteOnly())); + Mocha.test("drops reads that only feed write-only fields", () => Test_utils.eq("File \"sroa_test.res\", line 179, characters 64-71", 1, readByWriteOnly())); + Mocha.test("scalarizes several blocks updated together", () => Test_utils.eq("File \"sroa_test.res\", line 180, characters 62-69", 38, manyInLoop(3))); }); export { @@ -168,5 +193,8 @@ export { uncalledWriteOnlyClosure, readOnlyFieldSnapshotsInitializer, writeOnlyFieldStillRaises, + storedInWriteOnly, + readByWriteOnly, + manyInLoop, } /* Not a pure module */ diff --git a/tests/tests/src/sroa_test.res b/tests/tests/src/sroa_test.res index 377417a3148..abe945a4ed6 100644 --- a/tests/tests/src/sroa_test.res +++ b/tests/tests/src/sroa_test.res @@ -119,6 +119,39 @@ let writeOnlyFieldStillRaises = () => { cell.live } +// [cell] only reaches [outer] through a write-only field, so [cell] still +// becomes a scalar. +let storedInWriteOnly = () => { + let cell = ref(1) + let outer = ref(cell) + outer := cell + cell.contents +} + +// [first]'s only read is a value written to the write-only [second], so +// [first] is write-only as well. +let readByWriteOnly = () => { + let value = ref(1) + let first = ref(0) + let second = ref(0) + first := value.contents + second := first.contents + value.contents +} + +// Several references updated in one loop are all replaced. +let manyInLoop = n => { + let a = ref(0) + let b = ref(1) + let c = ref(2) + for i in 0 to n { + a := a.contents + i + b := b.contents + a.contents + c := c.contents + b.contents + } + a.contents + b.contents + c.contents +} + describe(__MODULE__, () => { test("scalarizes a local mutable record", () => eq(__LOC__, 42, localPair())) test("shares scalar fields with closures", () => eq(__LOC__, 32, capturedPair())) @@ -142,4 +175,7 @@ describe(__MODULE__, () => { test("preserves exceptions from write-only field values", () => throws(__LOC__, writeOnlyFieldStillRaises) ) + test("scalarizes a block stored in a write-only field", () => eq(__LOC__, 1, storedInWriteOnly())) + test("drops reads that only feed write-only fields", () => eq(__LOC__, 1, readByWriteOnly())) + test("scalarizes several blocks updated together", () => eq(__LOC__, 38, manyInLoop(3))) }) From ba4c395e8946be74a8f6e386c0805d219a04bd60 Mon Sep 17 00:00:00 2001 From: Christoph Knittel Date: Sat, 7 Nov 2026 19:02:00 +0000 Subject: [PATCH 2/2] Analyze blocks without allocating a closure per node The benchmarks showed the single pass allocating up to 0.9% more for functions without local blocks: each node passed a new closure to the traversal. The innermost slot is now a reference, so the traversal reuses one closure, and other inputs allocate as much as before. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Kbf3gK1dcjbe97yY1wRMgs Signed-off-by: Christoph Knittel --- compiler/core/lam_pass_sroa.ml | 51 +++++++++++++++++++--------------- 1 file changed, 29 insertions(+), 22 deletions(-) diff --git a/compiler/core/lam_pass_sroa.ml b/compiler/core/lam_pass_sroa.ml index cffbaa06fc3..80fc103a123 100644 --- a/compiler/core/lam_pass_sroa.ml +++ b/compiler/core/lam_pass_sroa.ml @@ -128,32 +128,35 @@ let analyze ~considered (lam : Lambda.t) = c.slots; false) in - (* [slot]: the innermost pending slot containing [lam]. *) - let rec go slot (lam : Lambda.t) = + (* The innermost pending slot containing the term being traversed. A + reference rather than an argument of [go], so that traversing allocates + no closures. *) + let slot = ref None in + let rec go (lam : Lambda.t) = match lam with | Lvar id -> ( match Hash_ident.find_opt candidates id with - | Some c -> record slot (Escape c) + | Some c -> record !slot (Escape c) | None -> ()) | Lassign (id, value) -> (match Hash_ident.find_opt candidates id with - | Some c -> record slot (Escape c) + | Some c -> record !slot (Escape c) | None -> ()); - go slot value + go value | Lprim {primitive = Pfield (index, _); args = [Lvar id]} when Hash_ident.mem candidates id -> let c = Hash_ident.find candidates id in - if valid_field c.uses index then record slot (Read (c, index)) - else record slot (Escape c) + if valid_field c.uses index then record !slot (Read (c, index)) + else record !slot (Escape c) | Lprim {primitive = Psetfield (index, _); args = [Lvar id; value]} when Hash_ident.mem candidates id -> let c = Hash_ident.find candidates id in if valid_field c.uses index then ( c.uses.(index).written <- true; - go (value_slot slot c index value) value) + go_value c index value) else ( - record slot (Escape c); - go slot value) + record !slot (Escape c); + go value) | Llet (kind, block, init, body) when is_candidate kind init && considered block -> ( match init with @@ -170,25 +173,29 @@ let analyze ~considered (lam : Lambda.t) = } in incr count; - List.iteri - (fun index init -> go (value_slot slot c index init) init) - initializers; + List.iteri (fun index init -> go_value c index init) initializers; Hash_ident.add candidates block c; - go slot body; + go body; if decide c then c.fields <- Some (fields_for_block block info field_count) | _ -> assert false) - | _ -> Lambda_traverse.iter (go slot) lam - (* The slot for a value of field [index] of [c], if the rewrite drops the - value when the field is never read. *) - and value_slot slot c index value = + | _ -> ignore (Lambda_traverse.shallow_exists go_child lam) + and go_child lam = + go lam; + false + (* Traverse a value of field [index] of [c], in a slot of its own if the + rewrite drops the value when the field is never read. *) + and go_value c index value = if Lam_analysis.no_side_effects value then ( - let s = {owner = c; field = index; parent = slot; occurrences = []} in + let parent = !slot in + let s = {owner = c; field = index; parent; occurrences = []} in c.slots <- s :: c.slots; - Some s) - else slot + slot := Some s; + go value; + slot := parent) + else go value in - go None lam; + go lam; candidates let rewrite candidates (lam : Lambda.t) =