From 048de0200db47da47ff2daf759b1edeff9f9547e Mon Sep 17 00:00:00 2001 From: Dmitri Sotnikov Date: Tue, 16 Jun 2026 01:21:36 +0000 Subject: [PATCH] Scalar-replace short-lived record allocations (jolt-15jq) (#132) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit scalar-replace already folds non-escaping const-key map literals ((:k {:k a ..}) -> a, and drops a let-bound map that doesn't escape). Extend the same fold to record constructors: a (->Rec a b c) is a positional struct whose declared field order lives in the record-shapes registry, so a field read on a non-escaping ctor folds to the matching positional arg and the allocation disappears. Direct form (:field (->Rec ..)) and the let-bound form both handled, threaded through run-passes via a per-unit shape registry (new jolt.host/record-shapes accessor). Soundness: ctor args must be pure (duplicated/discarded like map vals), arg count must equal the field count, and only declared-field reads fold — a record answers the virtual :jolt/deftype key with its type tag and any other key with nil, neither of which is a positional arg, so those keep the allocation. pure? now treats a record ctor of pure args as pure, so nested records (a Ray holding a Vec3) fold bottom-up. Allocation-bound microbench (non-escaping record built + field-read in a hot loop): 69.6s -> 2.4s, landing on the no-record arithmetic baseline. The ray tracer is unchanged — its vec3 results escape (returned/stored each op), so they genuinely allocate; that's a separate problem. Co-authored-by: Yogthos --- jolt-core/jolt/passes.clj | 7 +- jolt-core/jolt/passes/inline.clj | 155 +++++++++++++++--- src/jolt/host_iface.janet | 9 +- .../scalar-replace-record-test.janet | 61 +++++++ 4 files changed, 204 insertions(+), 28 deletions(-) create mode 100644 test/integration/scalar-replace-record-test.janet diff --git a/jolt-core/jolt/passes.clj b/jolt-core/jolt/passes.clj index a1e4ba3..075a0d3 100644 --- a/jolt-core/jolt/passes.clj +++ b/jolt-core/jolt/passes.clj @@ -13,9 +13,9 @@ :refer, so jolt.passes stays the only namespace the back end imports. Portable Clojure: kernel-tier fns + seed primitives only." - (:require [jolt.host :refer [inline-enabled?]] + (:require [jolt.host :refer [inline-enabled? record-shapes]] [jolt.passes.fold :refer [const-fold]] - [jolt.passes.inline :refer [inline-node flatten-lets scalar-replace dirty]] + [jolt.passes.inline :refer [inline-node flatten-lets scalar-replace dirty set-rec-shapes!]] [jolt.passes.types :refer [run-inference check-form infer-body reinfer-def set-rtenv! set-vtypes! join-types @@ -33,7 +33,8 @@ type is proven. Otherwise (core + bootstrap) just const-fold, as before." [node ctx] (if (inline-enabled? ctx) - (let [opt (loop [i 0 n (const-fold node)] + (let [_ (set-rec-shapes! (record-shapes ctx)) ;; record ctor fold (jolt-15jq) + opt (loop [i 0 n (const-fold node)] (reset! dirty false) (let [n2 (const-fold (scalar-replace (flatten-lets (inline-node n ctx))))] (if (and @dirty (< i 8)) diff --git a/jolt-core/jolt/passes/inline.clj b/jolt-core/jolt/passes/inline.clj index b1ed6a5..622fa77 100644 --- a/jolt-core/jolt/passes/inline.clj +++ b/jolt-core/jolt/passes/inline.clj @@ -14,6 +14,16 @@ (def dirty (atom false)) ;; read/reset by the run-passes fixpoint (jolt.passes) (defn- mark! [] (reset! dirty true)) +;; Record-ctor shape registry ("ns/->Name" -> {:fields (:k ..) :type tag}), fed +;; per unit by run-passes (set-rec-shapes!) before the fixpoint so scalar-replace +;; can recognize a (->Rec ..) call and map its positional args to declared fields +;; — the record analogue of the inline keys a map literal already carries in the +;; IR (jolt-15jq). +(def ^:private rec-shapes (atom {})) +(defn set-rec-shapes! + "Install the record-ctor shape registry the record fold consults (jolt-15jq)." + [m] (reset! rec-shapes (or m {}))) + (def ^:private fresh-counter (atom 0)) (defn- fresh [base] (let [n @fresh-counter] @@ -244,10 +254,16 @@ (= op :host) (contains? pure-fns (get f :name)) :else false))) +;; forward ref: a record ctor (allocating an immutable struct from its args) is +;; side-effect-free, so pure? treats (->Rec pure-args..) as pure — which lets a +;; nested record (a Ray holding a Vec3) fold bottom-up (jolt-15jq). +(declare ctor-shape) + (defn- pure? "Conservative: true only for expressions with no side effects that are safe to - duplicate or discard. A var/host ref is a pure read; an invoke is pure only - for a known-pure fn (arithmetic, comparison, keyword lookup, get)." + duplicate or discard. A var/host ref is a pure read; an invoke is pure for a + known-pure fn (arithmetic, comparison, keyword lookup, get) or a record + constructor (an immutable struct alloc) whose args are themselves pure." [node] (let [op (get node :op)] (cond @@ -263,7 +279,8 @@ (= op :vector) (every? pure? (get node :items)) (= op :set) (every? pure? (get node :items)) (= op :map) (every? (fn [pr] (and (pure? (nth pr 0)) (pure? (nth pr 1)))) (get node :pairs)) - (= op :invoke) (and (pure-fn? (get node :fn)) (every? pure? (get node :args))) + (= op :invoke) (and (or (pure-fn? (get node :fn)) (ctor-shape node)) + (every? pure? (get node :args))) :else false))) (defn- const-key-map? [node] @@ -410,46 +427,136 @@ (= op :def) (local-escapes? (get node :init) nm) :else true))) -(defn- subst-lookup - "Replace every (:k nm)/(get nm :k) in node with the map value at k. The caller - guarantees (via local-escapes?) that nm is never rebound here and appears only - as a lookup subject, so no shadowing logic is needed." - [node nm mapnode] +;; --- record constructors as foldable struct sources (jolt-15jq) ------------- +;; A record ctor (->Rec a b ..) is a positional struct: the registry maps its +;; ctor key ("ns/->Name", exactly how the IR names the call head) to the DECLARED +;; field order. A field read on a non-escaping ctor folds to the matching arg, +;; just as (:k {:k a ..}) folds to a. Two soundness differences from maps: +;; - the ctor's args are duplicated/discarded, so they must be pure (like map +;; vals), and the arg count must equal the field count (a positional call); +;; - a record answers the virtual :jolt/deftype key with its type tag and any +;; other non-field key with nil — neither is a positional arg, so we only +;; fold DECLARED-field reads and keep the allocation otherwise. + +(defn- ctor-shape + "If node is a record-constructor :invoke (its :fn a :var whose ns/name is a + registered ctor key, with arg count matching the declared field count), return + that record's shape entry; else nil." + [node] + (if (= :invoke (get node :op)) + (let [f (get node :fn)] + (if (= :var (get f :op)) + (let [rs (get @rec-shapes (str (get f :ns) "/" (get f :name)))] + (if (and rs (= (count (get rs :fields)) (count (get node :args)))) + rs + nil)) + nil)) + nil)) + +(defn- ctor-all-args-pure? [node] (every? pure? (get node :args))) + +(defn- field-index + "Index of scalar key k in the declared field tuple fields, or nil." + [fields k] + (let [n (count fields)] + (loop [i 0] + (if (< i n) + (if (= (nth fields i) k) i (recur (inc i))) + nil)))) + +(defn- ctor-val + "The positional arg IR at declared field k of record ctor node (shape rs). Only + called for a key known to be a field, so the index is always present." + [ctor rs k] + (nth (get ctor :args) (field-index (get rs :fields) k))) + +(defn- collect-keys! + "Accumulate (into atom acc) every constant-keyword lookup key applied to local + nm in node. The caller has proven (via local-escapes?) that nm appears only as + a lookup subject and is never rebound, so a uniform recursion suffices: at a + lookup of nm we record the key and stop (its subject is nm itself); elsewhere + we recurse into children." + [node nm acc] (let [k (lookup-key node nm)] (if k - (map-val mapnode k) + (swap! acc conj k) + (map-ir-children (fn [c] (collect-keys! c nm acc) c) node)))) + +(defn- lookups-all-fields? + "True if every lookup of nm across nodes uses a declared field in fields — the + record-only guard that keeps a :jolt/deftype/unknown-key read (not a positional + arg) from being folded to the wrong value." + [nodes nm fields] + (every? (fn [node] + (let [acc (atom #{})] + (collect-keys! node nm acc) + (every? (fn [k] (field-index fields k)) @acc))) + nodes)) + +(defn- src-val + "Field value at k from a foldable struct source — a const-key map (absent key + -> nil, struct-map semantics) or a record ctor (k is always a declared field + here, guaranteed by lookups-all-fields?)." + [src k] + (if (= :map (get src :op)) + (map-val src k) + (ctor-val src (ctor-shape src) k))) + +(defn- subst-lookup + "Replace every (:k nm)/(get nm :k) in node with the source value at k. The + caller guarantees (via local-escapes?) that nm is never rebound here and + appears only as a lookup subject, so no shadowing logic is needed." + [node nm src] + (let [k (lookup-key node nm)] + (if k + (src-val src k) ;; the caller's escape check guarantees nm is never rebound below, so we ;; recurse uniformly into every child — leaving any lookup of nm ;; un-substituted would dangle. - (map-ir-children (fn [c] (subst-lookup c nm mapnode)) node)))) + (map-ir-children (fn [c] (subst-lookup c nm src)) node)))) (defn- fold-kw-literal - "(a) (:k {:k a ..}) -> a (siblings pure)." + "(a) (:k ) -> the value at k. is a const-key pure map + ((:k {:k a ..}) -> a) or a pure record ctor ((:k (->Rec a ..)) -> the arg for + field k). Siblings are duplicated/discarded, so all must be pure; a record + lookup folds only for a declared field." [node] (let [f (get node :fn) args (get node :args)] (if (and (= :const (get f :op)) (keyword? (get f :val)) (= 1 (count args))) - (let [m (nth args 0)] + (let [m (nth args 0) k (get f :val)] (if (and (= :map (get m :op)) (const-key-map? m) (all-vals-pure? m)) - (do (mark!) (map-val m (get f :val))) - node)) + (do (mark!) (map-val m k)) + (let [rs (ctor-shape m)] + (if (and rs (ctor-all-args-pure? m) (field-index (get rs :fields) k)) + (do (mark!) (ctor-val m rs k)) + node)))) node))) -(defn- elim-let-maps - "(b) Drop the first non-escaping let binding whose init is a pure const-key map - literal, substituting its field lookups into the remaining bindings and body. - Fixpoint re-runs us for the rest, so one elimination per call keeps it simple." +(defn- elim-let-structs + "(b) Drop the first non-escaping let binding whose init is a foldable struct + source — a pure const-key map literal or a pure record ctor — substituting its + field reads into the remaining bindings and body. Fixpoint re-runs us for the + rest, so one elimination per call keeps it simple. For a record every lookup + of the binding must hit a declared field, else we keep the allocation." [node] (let [binds (get node :bindings) n (count binds) body (get node :body)] (loop [i 0] (if (< i n) - (let [b (nth binds i) nm (nth b 0) init (nth b 1)] - (if (and (= :map (get init :op)) (const-key-map? init) (all-vals-pure? init) + (let [b (nth binds i) nm (nth b 0) init (nth b 1) + ismap (and (= :map (get init :op)) (const-key-map? init) (all-vals-pure? init)) + rs (when (not ismap) (ctor-shape init)) + isrec (and rs (ctor-all-args-pure? init))] + (if (and (or ismap isrec) (not (any-binding-named? (subvec binds (inc i) n) nm)) (not (loop [j (inc i)] (if (< j n) (if (local-escapes? (nth (nth binds j) 1) nm) true (recur (inc j))) false))) - (not (local-escapes? body nm))) + (not (local-escapes? body nm)) + (or ismap + (lookups-all-fields? + (conj (mapv (fn [bb] (nth bb 1)) (subvec binds (inc i) n)) body) + nm (get rs :fields)))) (let [head (subvec binds 0 i) tail (mapv (fn [bb] [(nth bb 0) (subst-lookup (nth bb 1) nm init)]) (subvec binds (inc i) n)) @@ -467,8 +574,8 @@ [node] (let [op (get node :op)] (cond - ;; (a) fold (:k {:k a ..}) at invokes, after scalar-replacing children + ;; (a) fold (:k ) at invokes, after scalar-replacing children (= op :invoke) (fold-kw-literal (map-ir-children scalar-replace node)) - ;; (b) drop a non-escaping const-key-map let binding, after children - (= op :let) (elim-let-maps (map-ir-children scalar-replace node)) + ;; (b) drop a non-escaping foldable-struct let binding, after children + (= op :let) (elim-let-structs (map-ir-children scalar-replace node)) :else (map-ir-children scalar-replace node)))) diff --git a/src/jolt/host_iface.janet b/src/jolt/host_iface.janet index 9c1ca01..767821a 100644 --- a/src/jolt/host_iface.janet +++ b/src/jolt/host_iface.janet @@ -239,6 +239,13 @@ # not just one whose ->Name happens to be interned in the current ns). (if (record-hint-ctor-key ctx name) true false)) +# The record-ctor shape registry ("ns/->Name" -> {:fields (:k ..) :type tag ..}), +# or an empty table. Lets scalar-replace (jolt-15jq) recognize a (->Rec ..) call +# and map its positional args to the declared field order — the record analogue +# of the inline map-literal keys the map fold already sees in the IR. +(defn h-record-shapes [ctx] + (or (get (ctx :env) :record-shapes) @{})) + (def- exports {"form-sym?" h-sym? "form-sym-name" h-sym-name "form-sym-ns" h-sym-ns "ref-put!" h-ref-put! @@ -255,7 +262,7 @@ "host-intern!" h-intern! "inline-enabled?" h-inline-enabled? "inline-ir" h-inline-ir "record-type?" h-record-type? "form-position" h-form-position - "record-ctor-key" h-record-ctor-key}) + "record-ctor-key" h-record-ctor-key "record-shapes" h-record-shapes}) (defn install! [ctx] (def ns (ctx-find-ns ctx "jolt.host")) diff --git a/test/integration/scalar-replace-record-test.janet b/test/integration/scalar-replace-record-test.janet new file mode 100644 index 0000000..b84b556 --- /dev/null +++ b/test/integration/scalar-replace-record-test.janet @@ -0,0 +1,61 @@ +# Scalar-replacement of short-lived RECORD allocations (jolt-15jq). The pass +# already folds the const-key MAP-literal form ((:k {:k a ..}) -> a and drops a +# non-escaping let-bound map); this extends it to record CONSTRUCTORS. A record +# ctor (->Rec a b ..) is a positional struct whose declared field order lives in +# the record-shapes registry, so a field read on a non-escaping ctor result folds +# to the corresponding positional arg and the allocation disappears. +# +# Probe: count occurrences of the ctor var "->V3" in the analyzed IR. Folded => +# the ctor is gone (0). Mirrors type-infer-test's guard-counting harness. +(import ../../src/jolt/api :as api) +(import ../../src/jolt/backend :as backend) +(import ../../src/jolt/reader :as reader) + +(print "Scalar-replace of records (jolt-15jq)...") + +(os/setenv "JOLT_DIRECT_LINK" "1") +(def ctx (api/init {:compile? true})) +(api/eval-string ctx "(ns srr)") +(api/eval-string ctx "(defrecord V3 [r g b])") + +(defn ctors [src] + (length (string/find-all "->V3" + (string/format "%p" (backend/analyze-form ctx (reader/parse-string src)))))) +(defn ev [src] (api/eval-string ctx src)) + +# --- direct form: (:field (->V3 a b c)) -> the positional arg ---------------- +(assert (= 0 (ctors "(fn [] (:r (->V3 1 2 3)))")) "direct record lookup :r -> arg, ctor gone") +(assert (= 0 (ctors "(fn [] (:g (->V3 1 2 3)))")) "direct record lookup :g -> arg, ctor gone") +(assert (= 0 (ctors "(fn [] (:b (->V3 1 2 3)))")) "direct record lookup :b -> arg, ctor gone") +# pure non-constant args fold too (each discarded sibling is pure) +(assert (= 0 (ctors "(fn [a b] (:r (->V3 (+ a 1) (* b 2) 7)))")) "direct fold with pure arith args") + +# --- let form: non-escaping let-bound record, field reads -> args ------------ +(assert (= 0 (ctors "(fn [a b c] (let [v (->V3 a b c)] (+ (:r v) (:g v) (:b v))))")) "let-bound record, all field reads folded") +(assert (= 0 (ctors "(fn [a b c] (let [v (->V3 a b c)] (:r v)))")) "let-bound record, single field read folded (siblings discarded, pure)") + +# --- sound fallbacks: keep the allocation ------------------------------------ +# escaping record (passed to a sink) must NOT be folded +(assert (>= (ctors "(fn [sink a b c] (let [v (->V3 a b c)] (sink v) (:r v)))") 1) "escaping record keeps the allocation") +# returning the record escapes it +(assert (>= (ctors "(fn [a b c] (->V3 a b c))") 1) "returned record keeps the allocation") +# a non-field key read (:jolt/deftype is a virtual key) -> not folded, keep alloc +(assert (>= (ctors "(fn [a b c] (let [v (->V3 a b c)] (:jolt/deftype v)))") 1) ":jolt/deftype lookup keeps the allocation") + +# --- correctness: folded path evaluates identically ------------------------- +(assert (= 1 (ev "((fn [] (:r (->V3 1 2 3))))")) "direct :r value") +(assert (= 2 (ev "((fn [] (:g (->V3 1 2 3))))")) "direct :g value") +(assert (= 3 (ev "((fn [] (:b (->V3 1 2 3))))")) "direct :b value") +(assert (= 6 (ev "((fn [a b c] (let [v (->V3 a b c)] (+ (:r v) (:g v) (:b v)))) 1 2 3)")) "let-bound sum value") +(assert (= 10 (ev "((fn [a b] (:r (->V3 (+ a 1) (* b 2) 7))) 9 5)")) "arith-arg direct value") +# correctness of the kept-allocation fallbacks +(assert (= 1 (ev "((fn [sink a b c] (let [v (->V3 a b c)] (sink v) (:r v))) (fn [_] nil) 1 2 3)")) "escaping record reads correctly") +(assert (= "srr.V3" (ev "((fn [a b c] (let [v (->V3 a b c)] (:jolt/deftype v))) 1 2 3)")) ":jolt/deftype reads the type tag") + +# --- nested records fold compositionally (bottom-up) ------------------------- +(api/eval-string ctx "(defrecord Ray [orig dir])") +# (:r (:orig (->Ray (->V3 a b c) d))): inner ctors both fold away +(assert (= 0 (ctors "(fn [a b c d] (:r (:orig (->Ray (->V3 a b c) d))))")) "nested record reads fold both ctors") +(assert (= 7 (ev "((fn [a b c d] (:r (:orig (->Ray (->V3 a b c) d)))) 7 8 9 0)")) "nested record fold value") + +(print "Scalar-replace of records passed!")