Enable IR inlining: splice small defns at call sites (lever 1/4)

jolt.passes.inline was fully written but dormant — it fetched bodies via the
inline-ir host hook, which was a stub returning nil. Wire it up: run-passes stashes
each inline-eligible defn (single fixed arity) as its form is optimized, and
inline-ir hands the body back at call sites under --opt.

The catch was the ^double/^long coercion: an inlined fn drops its param-entry and
return coercion, so (work 3 4) on a ^double fn would return 25 instead of 25.0. New
:coerce IR node carries the coercion inside the spliced body — the inline pass wraps
a hinted param's arg and the return in :coerce, the back end lowers it
(exact->inexact / jolt->fx), and jolt.passes.numeric reads its :kind. So an inlined
call matches the called one and the body's fl*/fx* fast path still fires.

Only under --opt (closed world); the seed mint and -e don't inline, so selfhost and
the corpus are unaffected. test/chez/inline-test.ss 12/12 (make inline); full make
test green, 0 new corpus divergences.

Bench (hot loop, body is a ^double helper call): direct-link 500ms -> --opt
(inlined) 184ms = 2.7x, by eliminating the call + coercion wrappers and letting Chez
fuse the fl-ops unboxed. ~26x over the default dispatched build.
This commit is contained in:
Yogthos 2026-06-23 17:43:13 -04:00
parent 9927fd7f74
commit 79fa22eeab
9 changed files with 365 additions and 240 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 directlink numeric remint .PHONY: test ci values corpus unit smoke buildsmoke selfhost sci certify ffi transient infer directlink numeric inline 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 directlink numeric certify ci: values corpus unit smoke buildsmoke sci ffi transient infer 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.
@ -72,6 +72,11 @@ directlink:
numeric: numeric:
@chez --script test/chez/numeric-test.ss @chez --script test/chez/numeric-test.ss
# IR inlining: a small single-arity defn is spliced at call sites (under optimize),
# with ^double/^long entry/return coercions carried through via :coerce nodes.
inline:
@chez --script test/chez/inline-test.ss
# JVM oracle: certify the corpus against reference Clojure. Skips if clojure absent. # JVM oracle: certify the corpus against reference Clojure. Skips if clojure absent.
certify: certify:
@if command -v clojure >/dev/null 2>&1; then \ @if command -v clojure >/dev/null 2>&1; then \

View file

@ -300,7 +300,15 @@
(define hc-optimize? #f) (define hc-optimize? #f)
(define (set-optimize! on) (set! hc-optimize? on)) (define (set-optimize! on) (set! hc-optimize? on))
(define (hc-inline-enabled? ctx) hc-optimize?) (define (hc-inline-enabled? ctx) hc-optimize?)
(define (hc-inline-ir ctx ns-name nm) jolt-nil) ;; Inline-body registry: jolt.passes stashes an inline-eligible defn's
;; {:params :body :nhints :ret} here (keyed ns/name) as its form is optimized;
;; jolt.passes.inline fetches it to splice the body at a call site. The stash is an
;; opaque jolt value to the host — IR maps round-tripping through the table.
(define inline-stash-table (make-hashtable string-hash string=?))
(define (hc-stash-inline! ctx ns-name nm m)
(hashtable-set! inline-stash-table (string-append ns-name "/" nm) m) jolt-nil)
(define (hc-inline-ir ctx ns-name nm)
(or (hashtable-ref inline-stash-table (string-append ns-name "/" nm) #f) jolt-nil))
;; --- declare the hot clojure.core primitives so resolve-global sees them ------ ;; --- declare the hot clojure.core primitives so resolve-global sees them ------
;; (mirrors backend_scheme.clj native-ops keys — the emitter lowers these inline, ;; (mirrors backend_scheme.clj native-ops keys — the emitter lowers these inline,
@ -354,6 +362,7 @@
(def-var! "jolt.host" "record-ctor-key" hc-record-ctor-key) (def-var! "jolt.host" "record-ctor-key" hc-record-ctor-key)
(def-var! "jolt.host" "record-shapes" hc-record-shapes) (def-var! "jolt.host" "record-shapes" hc-record-shapes)
(def-var! "jolt.host" "inline-enabled?" hc-inline-enabled?) (def-var! "jolt.host" "inline-enabled?" hc-inline-enabled?)
(def-var! "jolt.host" "inline-ir" hc-inline-ir)) (def-var! "jolt.host" "inline-ir" hc-inline-ir)
(def-var! "jolt.host" "stash-inline!" hc-stash-inline!))
(hc-install!) (hc-install!)

File diff suppressed because one or more lines are too long

View file

@ -571,6 +571,11 @@
(mapcat (fn [p] [(emit (nth p 0)) (emit (nth p 1))]) (:pairs node))) (mapcat (fn [p] [(emit (nth p 0)) (emit (nth p 1))]) (:pairs node)))
:quote (emit-quoted (:form node)) :quote (emit-quoted (:form node))
:throw (str "(jolt-throw " (emit (:expr node)) ")") :throw (str "(jolt-throw " (emit (:expr node)) ")")
;; numeric coercion (from an inlined ^double/^long param or return).
:coerce (let [e (emit (:expr node))]
(cond (= :double (:kind node)) (str "(exact->inexact " e ")")
(= :long (:kind node)) (str "(jolt->fx " e ")")
:else e))
:try (emit-try node) :try (emit-try node)
;; regex literal #"…" -> a jolt-regex value (regex.ss, vendored irregex). ;; regex literal #"…" -> a jolt-regex value (regex.ss, vendored irregex).
:regex (str "(jolt-regex " (chez-str-lit (:source node)) ")") :regex (str "(jolt-regex " (chez-str-lit (:source node)) ")")

View file

@ -66,6 +66,13 @@
(defn quote-node [form] {:op :quote :form form}) (defn quote-node [form] {:op :quote :form form})
(defn throw-node [expr] {:op :throw :expr expr}) (defn throw-node [expr] {:op :throw :expr expr})
;; Numeric coercion of a value to a primitive kind (:double / :long), the way a JVM
;; ^double/^long parameter or return coerces. The back end lowers it (exact->inexact
;; / jolt->fx) and jolt.passes.numeric reads its :kind as the value's numeric kind.
;; Carrying coercion as an IR node (rather than a back-end string wrap) lets it
;; travel with inlining and keeps the typed-arithmetic fast path sound.
(defn coerce-node [kind expr] {:op :coerce :kind kind :expr expr})
;; --------------------------------------------------------------------------- ;; ---------------------------------------------------------------------------
;; Structural recursion over IR child nodes. ;; Structural recursion over IR child nodes.
;; ;;
@ -91,6 +98,7 @@
(= op :do) (assoc node :statements (mapv f (get node :statements)) (= op :do) (assoc node :statements (mapv f (get node :statements))
:ret (f (get node :ret))) :ret (f (get node :ret)))
(= op :throw) (assoc node :expr (f (get node :expr))) (= op :throw) (assoc node :expr (f (get node :expr)))
(= op :coerce) (assoc node :expr (f (get node :expr)))
(= op :set-var) (assoc node :val (f (get node :val))) (= op :set-var) (assoc node :val (f (get node :val)))
(= op :set-field) (assoc node :obj (f (get node :obj)) :val (f (get node :val))) (= op :set-field) (assoc node :obj (f (get node :obj)) :val (f (get node :val)))
(= op :defmacro) (assoc node :fn (f (get node :fn))) (= op :defmacro) (assoc node :fn (f (get node :fn)))
@ -138,6 +146,7 @@
(= op :if) (f (f (f acc (get node :test)) (get node :then)) (get node :else)) (= op :if) (f (f (f acc (get node :test)) (get node :then)) (get node :else))
(= op :do) (f (reduce f acc (get node :statements)) (get node :ret)) (= op :do) (f (reduce f acc (get node :statements)) (get node :ret))
(= op :throw) (f acc (get node :expr)) (= op :throw) (f acc (get node :expr))
(= op :coerce) (f acc (get node :expr))
(= op :set-var) (f acc (get node :val)) (= op :set-var) (f acc (get node :val))
(= op :set-field) (f (f acc (get node :obj)) (get node :val)) (= op :set-field) (f (f acc (get node :obj)) (get node :val))
(= op :defmacro) (f acc (get node :fn)) (= op :defmacro) (f acc (get node :fn))

View file

@ -13,7 +13,7 @@
:refer, so jolt.passes stays the only namespace the back end imports. :refer, so jolt.passes stays the only namespace the back end imports.
Portable Clojure: kernel-tier fns + seed primitives only." Portable Clojure: kernel-tier fns + seed primitives only."
(:require [jolt.host :refer [inline-enabled? record-shapes]] (:require [jolt.host :refer [inline-enabled? record-shapes stash-inline!]]
[jolt.passes.fold :refer [const-fold]] [jolt.passes.fold :refer [const-fold]]
[jolt.passes.numeric :as numeric] [jolt.passes.numeric :as numeric]
[jolt.passes.inline :refer [inline-node flatten-lets scalar-replace dirty set-rec-shapes!]] [jolt.passes.inline :refer [inline-node flatten-lets scalar-replace dirty set-rec-shapes!]]
@ -28,6 +28,17 @@
;; sets `dirty` when it rewrote something; the loop stops at a clean pass or here. ;; sets `dirty` when it rewrote something; the loop stops at a clean pass or here.
(def ^:private inline-fixpoint-cap 8) (def ^:private inline-fixpoint-cap 8)
;; A top-level defn the inline pass may splice: a single fixed arity (no rest). The
;; pass itself checks body size + closedness, so any such fn is stashable.
(defn- inline-eligible? [node]
(and (= :def (:op node)) (:init node) (= :fn (:op (:init node)))
(= 1 (count (:arities (:init node))))
(not (:rest (first (:arities (:init node)))))))
(defn- stash-of [node]
(let [a (first (:arities (:init node)))]
{:params (:params a) :body (:body a) :nhints (:nhints a) :ret (:ret-nhint a)}))
(defn run-passes (defn run-passes
"All passes, in order. The back end applies this to every analyzed form. When "All passes, in order. The back end applies this to every analyzed form. When
inlining is enabled for the unit (user code under direct-linking), inlining is enabled for the unit (user code under direct-linking),
@ -40,6 +51,10 @@
numeric/annotate runs last in both branches (hint-directed fl*/fx* arithmetic); numeric/annotate runs last in both branches (hint-directed fl*/fx* arithmetic);
it benefits open builds too, so it is not gated on inlining." it benefits open builds too, so it is not gated on inlining."
[node ctx] [node ctx]
;; stash an inline-eligible defn so later call sites can splice it (closed-world
;; optimization only). Done before optimizing, from the analyzed node.
(when (and (inline-enabled? ctx) (inline-eligible? node))
(stash-inline! ctx (:ns node) (:name node) (stash-of node)))
(numeric/annotate (numeric/annotate
(if (inline-enabled? ctx) (if (inline-enabled? ctx)
(let [_ (set-rec-shapes! (record-shapes ctx)) ;; record ctor fold (let [_ (set-rec-shapes! (record-shapes ctx)) ;; record ctor fold

View file

@ -4,7 +4,7 @@
share the alpha-rename invariant (every spliced binder is made globally fresh) share the alpha-rename invariant (every spliced binder is made globally fresh)
and the `dirty` fixpoint flag. Portable Clojure (compiler-tier)." and the `dirty` fixpoint flag. Portable Clojure (compiler-tier)."
(:require [jolt.host :refer [inline-ir]] (:require [jolt.host :refer [inline-ir]]
[jolt.ir :refer [map-ir-children reduce-ir-children]] [jolt.ir :refer [map-ir-children reduce-ir-children coerce-node]]
[jolt.passes.fold :refer [scalar-const?]])) [jolt.passes.fold :refer [scalar-const?]]))
;; --------------------------------------------------------------------------- ;; ---------------------------------------------------------------------------
@ -47,7 +47,7 @@
;; is rejected by body-size below and never inlined or alpha-renamed. ;; is rejected by body-size below and never inlined or alpha-renamed.
(or (= op :const) (= op :local) (= op :var) (= op :host) (= op :the-var) (or (= op :const) (= op :local) (= op :var) (= op :host) (= op :the-var)
(= op :quote) (= op :if) (= op :do) (= op :let) (= op :invoke) (= op :quote) (= op :if) (= op :do) (= op :let) (= op :invoke)
(= op :map) (= op :vector) (= op :set) (= op :throw))) (= op :map) (= op :vector) (= op :set) (= op :throw) (= op :coerce)))
(def ^:private inline-budget 120) (def ^:private inline-budget 120)
@ -141,6 +141,8 @@
(if stash (if stash
(let [params (get stash :params) (let [params (get stash :params)
body (get stash :body) body (get stash :body)
nh (reduce (fn [m pr] (assoc m (nth pr 0) (nth pr 1))) {} (get stash :nhints))
ret (get stash :ret)
args (get node :args)] args (get node :args)]
(if (and (= (count params) (count args)) (if (and (= (count params) (count args))
(<= (body-size body) inline-budget) (<= (body-size body) inline-budget)
@ -149,19 +151,25 @@
;; trivial args (local/const) substitute straight in (copy ;; trivial args (local/const) substitute straight in (copy
;; propagation); the rest get a fresh local bound once in a ;; propagation); the rest get a fresh local bound once in a
;; wrapping let, so they evaluate exactly once in source order. ;; wrapping let, so they evaluate exactly once in source order.
;; A ^double/^long param always binds (no copy-prop) so its
;; entry coercion runs — preserving the called fn's semantics.
res (loop [i 0 env {} binds []] res (loop [i 0 env {} binds []]
(if (< i n) (if (< i n)
(let [p (nth params i) a (nth args i)] (let [p (nth params i) a (nth args i) k (get nh p)]
(if (trivial-arg? a) (cond
(recur (inc i) (assoc env p a) binds) k (let [f (fresh p)]
(let [f (fresh p)] (recur (inc i) (assoc env p {:op :local :name f})
(recur (inc i) (conj binds [f (coerce-node k a)])))
(assoc env p {:op :local :name f}) (trivial-arg? a) (recur (inc i) (assoc env p a) binds)
(conj binds [f a]))))) :else (let [f (fresh p)]
(recur (inc i) (assoc env p {:op :local :name f})
(conj binds [f a])))))
[env binds])) [env binds]))
env (nth res 0) env (nth res 0)
binds (nth res 1) binds (nth res 1)
rbody (subst body env)] rbody0 (subst body env)
;; preserve the fn's ^double/^long return coercion.
rbody (if ret (coerce-node ret rbody0) rbody0)]
(mark!) (mark!)
(if (= 0 (count binds)) (if (= 0 (count binds))
rbody rbody

View file

@ -137,6 +137,7 @@
(cond (cond
(= op :const) [(if (float-lit? node) :double nil) node] (= op :const) [(if (float-lit? node) :double nil) node]
(= op :local) [(get tenv (get node :name)) node] (= op :local) [(get tenv (get node :name)) node]
(= op :coerce) [(get node :kind) (assoc node :expr (nth (an (get node :expr) tenv) 1))]
(= op :invoke) (an-invoke node tenv) (= op :invoke) (an-invoke node tenv)
(= op :let) (= op :let)
(let [res (reduce (fn [acc b] (let [res (reduce (fn [acc b]

67
test/chez/inline-test.ss Normal file
View file

@ -0,0 +1,67 @@
;; IR inlining (jolt.passes.inline), enabled under optimization. A small
;; single-arity defn is stashed and spliced at its call sites, removing the call.
;; A ^double/^long fn's param-entry and return coercions travel with the splice
;; (via :coerce nodes) so an inlined call matches the called one — incl. coercing a
;; non-double arg — and the body's fl*/fx* fast path still fires. Run:
;; chez --script test/chez/inline-test.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 total 0) (define fails 0)
(define (ok name pred) (set! total (+ total 1)) (unless pred (set! fails (+ fails 1)) (printf "FAIL: ~a\n" name)))
(define (has? s sub)
(let ((ns (string-length s)) (nsub (string-length sub)))
(let loop ((i 0))
(cond ((> (+ i nsub) ns) #f)
((string=? (substring s i (+ i nsub)) sub) #t)
(else (loop (+ i 1)))))))
(define (emitf ns str) ; analyze + run-passes (optimize on) + emit
(let-values (((f j) (rdr-read-form str 0 (string-length str))))
(let ((ctx (make-analyze-ctx ns)))
(jolt-ce-emit (jolt-ce-run-passes (jolt-ce-analyze ctx f) ctx)))))
(define (ev s) (jolt-compile-eval s "u"))
;; inlining is a closed-world optimization — only under optimize.
(set-optimize! #t)
;; a small plain fn is spliced; the call to it disappears.
(ev "(def add1 (fn* ([x] (+ x 1))))")
(let ((e (emitf "u" "(fn* ([y] (add1 y)))")))
(ok "plain fn is inlined (call to add1 gone)" (not (has? e "add1")))
(ok "inlined body present (+ ... 1)" (has? e "(+")))
(ok "inlined plain fn runtime: (add1 41) = 42" (= 42 (jnum->exact (ev "((fn* ([y] (add1 y))) 41)"))))
;; a ^double fn: body fl-ops fire after inlining, and the call is gone.
(ev "(def ^double dwork (fn* ([^double a ^double b] (+ (* a a) (* b b)))))")
(let ((e (emitf "u" "(fn* ([] (dwork 3.0 4.0)))")))
(ok "inlined ^double fn body uses fl*" (has? e "(fl*"))
(ok "inlined ^double fn call to dwork is gone" (not (has? e "dwork"))))
(ok "inlined ^double call: 3^2+4^2 = 25" (= 25 (jnum->exact (ev "((fn* ([] (dwork 3.0 4.0))))"))))
;; coercion travels with the splice: int args become doubles, so the result is a
;; flonum 25.0 — matching the called fn, not an exact 25.
(ok "inlined ^double with int args still returns a flonum" (flonum? (ev "((fn* ([] (dwork 3 4))))")))
;; a ^long fn inlines with fixnum coercion + fx ops.
(ev "(def ^long lsum (fn* ([^long a ^long b] (+ a b))))")
(let ((e (emitf "u" "(fn* ([] (lsum 3 4)))")))
(ok "inlined ^long fn body uses fx+" (has? e "(fx+")))
(ok "inlined ^long call: 3+4 = 7 (fixnum)" (let ((r (ev "((fn* ([] (lsum 3 4))))"))) (and (fixnum? r) (= r 7))))
;; an accumulator over an inlined ^double call: the whole loop body fuses to fl-ops.
(ev "(def ^double sq (fn* ([^double x] (* x x))))")
(let ((e (emitf "u" "(fn* ([] (loop [acc 0.0 i 0] (if (< i 3) (recur (+ acc (sq 2.0)) (inc i)) acc))))")))
(ok "accumulator over inlined ^double call lowers to fl+" (has? e "(fl+"))
(ok "the sq call is inlined away" (not (has? e "sq"))))
(ok "accumulator over inlined ^double call: 3*4.0 = 12" (= 12 (jnum->exact (ev "((fn* ([] (loop [acc 0.0 i 0] (if (< i 3) (recur (+ acc (sq 2.0)) (inc i)) acc)))))"))))
(set-optimize! #f)
(printf "~a/~a passed~n" (- total fails) total)
(exit (if (zero? fails) 0 1))