DO NOT MERGE — reduce-acc SROA regresses (defeats vec-reduce)

Reduce-accumulator scalar replacement: lower (reduce (fn [acc x] body) (->Rec
inits) coll) with a non-escaping record accumulator to a loop carrying the acc
fields as scalar vars. Correct (run-reduce-sroa.ss 7/7, make test + shakesmoke
green) but a PERF REGRESSION: ray tracer hit-all 38.4s -> ~53s.

Root cause: jolt's reduce already iterates a vector with vec-reduce (seq.ss) — a
tight index loop over the backing store, zero per-element allocation. Lowering to
a (seq)/(first)/(next) loop allocates a seq node per element, trading hit-all's
CONDITIONAL ->HitAcc allocation for a per-element seq allocation. The targeted
accumulator allocation is also conditional (only on hit improvement), not the
dominant per-sphere cost. Kept on this branch for reference; not merged.
This commit is contained in:
Yogthos 2026-06-26 11:28:34 -04:00
parent 4671e1b67e
commit e2598280fe
4 changed files with 390 additions and 132 deletions

View file

@ -4,7 +4,7 @@
# build step. `make test` is the full gate. `make remint` rebuilds the seed after a # build step. `make test` is the full gate. `make remint` rebuilds the seed after a
# source change. # source change.
.PHONY: test ci values corpus unit smoke buildsmoke selfhost sci certify ffi transient infer wp devirt fieldread numwp fieldnum directlink numeric inline shakesmoke remint .PHONY: test ci values corpus unit smoke buildsmoke selfhost sci certify ffi transient infer wp devirt fieldread numwp fieldnum reducesroa directlink numeric inline shakesmoke remint
# Full gate (dev machine). Includes the self-host byte-fixpoint, which only holds # Full gate (dev machine). Includes the self-host byte-fixpoint, which only holds
# on the same Chez that minted the seed. # on the same Chez that minted the seed.
@ -15,7 +15,7 @@ test: selfhost ci
# lockfile) — it RUNS correctly on any Chez, but `selfhost` rebuilds it and a # lockfile) — it RUNS correctly on any Chez, but `selfhost` rebuilds it and a
# different Chez version may emit byte-different (gensym/order) output, so the # different Chez version may emit byte-different (gensym/order) output, so the
# byte-fixpoint is a dev-machine check, not a CI one (jolt-8479). # byte-fixpoint is a dev-machine check, not a CI one (jolt-8479).
ci: values corpus unit smoke buildsmoke sci ffi transient infer wp devirt fieldread numwp fieldnum directlink numeric inline certify ci: values corpus unit smoke buildsmoke sci ffi transient infer wp devirt fieldread numwp fieldnum reducesroa directlink numeric inline certify
@echo "OK: CI gates passed" @echo "OK: CI gates passed"
# Self-host fixpoint: bootstrap.ss rebuild == checked-in seed. # Self-host fixpoint: bootstrap.ss rebuild == checked-in seed.
@ -92,6 +92,12 @@ numwp:
fieldnum: fieldnum:
@chez --script host/chez/run-fieldnum.ss @chez --script host/chez/run-fieldnum.ss
# Reduce-accumulator scalar replacement: a reduce over a non-escaping record
# accumulator lowers to a seq loop carrying the acc fields as scalar vars, so the
# per-step record allocation goes away; the result matches the ordinary reduce.
reducesroa:
@chez --script host/chez/run-reduce-sroa.ss
# Direct-linking emission: a closed-world build binds top-level app defs to jv$ # Direct-linking emission: a closed-world build binds top-level app defs to jv$
# Scheme bindings and routes app->app calls/refs to them, skipping var-deref + # Scheme bindings and routes app->app calls/refs to them, skipping var-deref +
# jolt-invoke; ^:dynamic/^:redef and nested defs opt out. # jolt-invoke; ^:dynamic/^:redef and nested defs opt out.

View file

@ -0,0 +1,84 @@
;; run-reduce-sroa.ss — reduce-accumulator scalar replacement gate.
;;
;; A (reduce (fn [acc x] body) (->Rec inits) coll) whose accumulator is a
;; non-escaping record (read only via its fields, rebuilt each step as a same-shape
;; ctor or carried forward unchanged) lowers to a seq loop that carries the acc's
;; fields as scalar loop vars and reconstructs the record once at exit — killing the
;; per-step allocation. This is the ray tracer's hit-all pattern (a HitAcc per
;; sphere test). Pinned here: the reduce call is gone (lowered to a loop), the
;; lowered result matches the generic reduce, and non-lowerable shapes (non-record
;; init, escaping acc) keep the ordinary reduce.
;;
;; chez --script host/chez/run-reduce-sroa.ss
(import (chezscheme))
(load "host/chez/rt.ss")
(set-chez-ns! "clojure.core")
(load "host/chez/seed/prelude.ss")
(load "host/chez/post-prelude.ss")
(set-chez-ns! "user")
(load "host/chez/host-contract.ss")
(load "host/chez/seed/image.ss")
(load "host/chez/compile-eval.ss")
(define analyze (var-deref "jolt.analyzer" "analyze"))
(define set-record-shapes! (var-deref "jolt.passes.types" "set-record-shapes!"))
(define set-protocol-methods! (var-deref "jolt.passes.types" "set-protocol-methods!"))
(define run-passes (var-deref "jolt.passes" "run-passes"))
(define emit (var-deref "jolt.backend-scheme" "emit"))
(define (anode src) (analyze (make-analyze-ctx "user") (jolt-ce-read src)))
(define (evals src) (jolt-compile-eval (string-append "(do " src ")") "user"))
(define (built scm-src) (eval (read (open-input-string scm-src)) (interaction-environment)))
(define (contains-sub? s sub)
(let ((n (string-length s)) (m (string-length sub)))
(let loop ((i 0))
(cond ((> (+ i m) n) #f)
((string=? (substring s i (+ i m)) sub) #t)
(else (loop (+ i 1)))))))
(define fails 0) (define total 0)
(define (check label actual expected)
(set! total (+ total 1))
(unless (equal? actual expected)
(set! fails (+ fails 1))
(printf " FAIL ~a: got ~s expected ~s\n" label actual expected)))
(evals "(defrecord Acc [sum cnt])")
(set-record-shapes! (chez-record-shapes-map))
(set-protocol-methods! (jolt-hash-map))
(set-optimize! #t)
(define (emit-opt src) (emit (run-passes (anode src) (make-analyze-ctx "user"))))
;; canonical accumulator: a same-shape ctor in one branch, the acc carried forward
;; in the other. Only the fields of acc are read; acc never escapes.
(define run-src
"(def run (fn [xs] (:sum (reduce (fn [acc x] (if (> x 0) (->Acc (+ (:sum acc) x) (+ (:cnt acc) 1)) acc)) (->Acc 0 0) xs))))")
(define run-scm (emit-opt run-src))
(check "reduce accumulator lowered (no reduce call)" (contains-sub? run-scm "reduce") #f)
(built run-scm)
(check "lowered result matches generic"
(jolt-invoke (var-deref "user" "run") (jolt-vector 1 -2 3 4))
(evals "(:sum (reduce (fn [acc x] (if (> x 0) (->Acc (+ (:sum acc) x) (+ (:cnt acc) 1)) acc)) (->Acc 0 0) [1 -2 3 4]))"))
;; a reduce over a record acc that reads BOTH fields at the end still matches.
(define cnt-src
"(def cntr (fn [xs] (:cnt (reduce (fn [acc x] (->Acc (+ (:sum acc) x) (+ (:cnt acc) 1))) (->Acc 0 0) xs))))")
(define cnt-scm (emit-opt cnt-src))
(check "second accumulator lowered" (contains-sub? cnt-scm "reduce") #f)
(built cnt-scm)
(check "count accumulator correct" (jolt-invoke (var-deref "user" "cntr") (jolt-vector 5 5 5 5)) 4)
;; empty coll returns the init (reduce semantics): (:sum (->Acc 0 0)) = 0
(check "empty coll returns init" (jolt-invoke (var-deref "user" "run") (jolt-vector)) 0)
;; --- negatives: shapes that must NOT be lowered keep the ordinary reduce --------
;; non-record init
(check "non-record reduce untouched"
(contains-sub? (emit-opt "(def sm (fn [xs] (reduce (fn [acc x] (+ acc x)) 0 xs)))") "reduce") #t)
;; acc escapes (passed whole to a fn)
(check "escaping-acc reduce untouched"
(contains-sub? (emit-opt "(def esc (fn [xs] (reduce (fn [acc x] (do (identity acc) (->Acc (+ (:sum acc) x) (:cnt acc)))) (->Acc 0 0) xs)))") "reduce") #t)
(if (= fails 0)
(begin (printf "reduce-sroa gate: ~a/~a passed\n" total total) (exit 0))
(begin (printf "reduce-sroa gate: ~a/~a passed (~a failed)\n" (- total fails) total fails) (exit 1)))

File diff suppressed because one or more lines are too long

View file

@ -546,13 +546,161 @@
(recur (inc i)))) (recur (inc i))))
node)))) node))))
;; --- reduce-accumulator scalar replacement ----------------------------------
;; (reduce (fn [acc x] body) (->Rec inits..) coll) where acc is a non-escaping
;; record — read only via its fields, returned each step as a same-shape ctor or
;; carried forward unchanged — allocates one record PER ELEMENT. Lower it to a seq
;; loop that carries acc's fields as scalar loop vars and rebuilds the record once
;; at exit, so the per-step allocation disappears. This is the ray tracer's hit-all
;; (a HitAcc per sphere test). Closed-world (--opt) only, like the rest of this pass.
(defn- local-node [nm] {:op :local :name nm})
(defn- core-call [nm cargs] {:op :invoke :fn {:op :var :ns "clojure.core" :name nm} :args (vec cargs)})
(defn- reduce3-callee?
"node is a (clojure.core/reduce f init coll) the 3-arg form with an explicit
initial value (the 2-arg form seeds from the collection, a different shape)."
[node]
(let [f (get node :fn)]
(and (= :invoke (get node :op))
(= :var (get f :op)) (= "clojure.core" (get f :ns)) (= "reduce" (get f :name))
(= 3 (count (get node :args))))))
(defn- same-shape-ctor? [node rs]
(let [cs (ctor-shape node)] (and cs (= (get cs :type) (get rs :type)))))
(defn- acc-ok?
"Is acc (local nm, record shape rs) safe to scalarize across body? It may appear
ONLY as a constant-field-read subject, or in tail position as the bare
carried-forward value. tail? marks whether node is the reduce-fn's return value.
A binding/control form that could capture or rebind acc bails (false)."
[node nm rs tail?]
(let [op (get node :op)
k (lookup-key node nm)]
(cond
;; a field read of acc: nm is consumed as the subject; a get-default (extra
;; arg) is non-tail and must not leak acc either.
k (let [args (get node :args)]
(if (> (count args) 1)
(every? (fn [a] (acc-ok? a nm rs false)) (subvec args 1 (count args)))
true))
(= op :local) (or (not= nm (get node :name)) tail?) ;; bare acc only ok in tail
(= op :const) true (= op :var) true (= op :host) true
(= op :the-var) true (= op :quote) true
(= op :if) (and (acc-ok? (get node :test) nm rs false)
(acc-ok? (get node :then) nm rs tail?)
(acc-ok? (get node :else) nm rs tail?))
(= op :do) (and (every? (fn [s] (acc-ok? s nm rs false)) (get node :statements))
(acc-ok? (get node :ret) nm rs tail?))
(= op :let) (and (every? (fn [b] (acc-ok? (nth b 1) nm rs false)) (get node :bindings))
(or (any-binding-named? (get node :bindings) nm) ;; nm shadowed below
(acc-ok? (get node :body) nm rs tail?)))
(= op :invoke) (and (acc-ok? (get node :fn) nm rs false)
(every? (fn [a] (acc-ok? a nm rs false)) (get node :args)))
(= op :throw) (acc-ok? (get node :expr) nm rs false)
(= op :vector) (every? (fn [a] (acc-ok? a nm rs false)) (get node :items))
(= op :set) (every? (fn [a] (acc-ok? a nm rs false)) (get node :items))
(= op :map) (every? (fn [p] (and (acc-ok? (nth p 0) nm rs false)
(acc-ok? (nth p 1) nm rs false))) (get node :pairs))
;; :fn/:loop/:recur/:try/:def — acc could be captured, or a nested recur would
;; collide with the loop we synthesize; conservatively bail.
:else false)))
(defn- tails-valid?
"Every tail position of body returns a same-shape ctor or the bare acc the
values the recur can carry as exploded scalars. (A `reduced` tail, or any other
shape, isn't one of these, so it keeps the ordinary reduce.)"
[node nm rs]
(let [op (get node :op)]
(cond
(= op :if) (and (tails-valid? (get node :then) nm rs) (tails-valid? (get node :else) nm rs))
(= op :do) (tails-valid? (get node :ret) nm rs)
(= op :let) (tails-valid? (get node :body) nm rs)
(and (= op :local) (= nm (get node :name))) true
:else (same-shape-ctor? node rs))))
(defn- subst-acc-fields
"Replace every (:k acc)/(get acc :k) in node with the scalar loop var for field k.
acc-ok? guarantees acc appears only as such reads (or the bare tail), so a uniform
recursion is safe."
[node nm fields acc-locals]
(let [k (lookup-key node nm)]
(if k
(nth acc-locals (field-index fields k))
(map-ir-children (fn [c] (subst-acc-fields c nm fields acc-locals)) node))))
(defn- tail->recur
"Convert each tail of body (already field-substituted) into a recur that advances
the seq and carries the next accumulator components: a same-shape ctor explodes
into its positional args, the bare acc carries the current components forward."
[node nm acc-locals snm]
(let [op (get node :op)
step (fn [t] (tail->recur t nm acc-locals snm))
recur-with (fn [comps] {:op :recur :args (reduce conj [(core-call "next" [(local-node snm)])] comps)})]
(cond
(= op :if) (assoc node :then (step (get node :then)) :else (step (get node :else)))
(= op :do) (assoc node :ret (step (get node :ret)))
(= op :let) (assoc node :body (step (get node :body)))
(and (= op :local) (= nm (get node :name))) (recur-with acc-locals)
:else (recur-with (get node :args))))) ;; a same-shape ctor: its args are the new components
(defn- lower-reduce
"Rewrite a lowerable (reduce (fn [acc x] body) (->Rec inits) coll) into a seq loop
with acc's fields carried as scalar vars. Caller has validated the shape."
[node]
(let [args (get node :args)
closure (nth args 0) init (nth args 1) coll (nth args 2)
ar (first (get closure :arities))
params (get ar :params)
nm (nth params 0) xnm (nth params 1)
body (get ar :body)
rs (ctor-shape init)
fields (get rs :fields)
snm (fresh "rseq")
accnms (mapv (fn [_] (fresh "racc")) fields)
acc-locals (mapv local-node accnms)
body' (tail->recur (subst-acc-fields body nm fields acc-locals) nm acc-locals snm)]
{:op :loop
:bindings (reduce conj [[snm (core-call "seq" [coll])]]
(mapv (fn [an ia] [an ia]) accnms (get init :args)))
:body {:op :if
:test (core-call "nil?" [(local-node snm)])
;; exhausted: rebuild the record once from the final components
:then {:op :invoke :fn (get init :fn) :args acc-locals}
:else {:op :let :bindings [[xnm (core-call "first" [(local-node snm)])]]
:body body'}}}))
(defn- try-lower-reduce
"Lower a (reduce closure (->Rec inits) coll) when the closure is a single-arity
2-param literal fn and acc is a non-escaping record; else nil."
[node]
(when (reduce3-callee? node)
(let [closure (nth (get node :args) 0)
init (nth (get node :args) 1)
ars (get closure :arities)]
(when (and (= :fn (get closure :op)) (= 1 (count ars))
(= 2 (count (get (first ars) :params)))
(not (get (first ars) :rest))
(ctor-shape init))
(let [ar (first ars)
nm (nth (get ar :params) 0)
body (get ar :body)
rs (ctor-shape init)]
(when (and (acc-ok? body nm rs true) (tails-valid? body nm rs))
(lower-reduce node)))))))
(defn scalar-replace (defn scalar-replace
"Bottom-up: scalar-replace children, then apply (a) at invokes / (b) at lets." "Bottom-up: scalar-replace children, then apply (a) at invokes / (b) at lets."
[node] [node]
(let [op (get node :op)] (let [op (get node :op)]
(cond (cond
;; (a) fold (:k <map|ctor>) at invokes, after scalar-replacing children ;; (a) at invokes: lower a reduce-over-record-accumulator to a loop, else fold
(= op :invoke) (fold-kw-literal (map-ir-children scalar-replace node)) ;; (:k <map|ctor>) — both after scalar-replacing children.
(= op :invoke)
(let [n (map-ir-children scalar-replace node)]
(if-let [low (try-lower-reduce n)]
(do (mark!) low)
(fold-kw-literal n)))
;; (b) drop a non-escaping foldable-struct let binding, after children ;; (b) drop a non-escaping foldable-struct let binding, after children
(= op :let) (elim-let-structs (map-ir-children scalar-replace node)) (= op :let) (elim-let-structs (map-ir-children scalar-replace node))
:else (map-ir-children scalar-replace node)))) :else (map-ir-children scalar-replace node))))