Merge pull request #181 from jolt-lang/spike/perf-levers

Perf levers: IR inlining + loop-counter typing

Makefile

@@ -4,7 +4,7 @@
 # build step. `make test` is the full gate. `make remint` rebuilds the seed after a
 # 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
 # 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
 # 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).
-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"
 
 # Self-host fixpoint: bootstrap.ss rebuild == checked-in seed.
@@ -72,6 +72,11 @@ directlink:
 numeric:
 	@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.
 certify:
 	@if command -v clojure >/dev/null 2>&1; then \

host/chez/host-contract.ss

@@ -300,7 +300,15 @@
 (define hc-optimize? #f)
 (define (set-optimize! on) (set! hc-optimize? on))
 (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 ------
 ;; (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-shapes" hc-record-shapes)
   (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!)

jolt-core/jolt/backend_scheme.clj

@@ -571,6 +571,11 @@
                        (mapcat (fn [p] [(emit (nth p 0)) (emit (nth p 1))]) (:pairs node)))
     :quote (emit-quoted (:form 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)
     ;; regex literal #"…" -> a jolt-regex value (regex.ss, vendored irregex).
     :regex (str "(jolt-regex " (chez-str-lit (:source node)) ")")

jolt-core/jolt/ir.clj

@@ -66,6 +66,13 @@
 (defn quote-node [form] {:op :quote :form form})
 (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.
 ;;
@@ -91,6 +98,7 @@
       (= op :do)     (assoc node :statements (mapv f (get node :statements))
                                  :ret (f (get node :ret)))
       (= 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-field) (assoc node :obj (f (get node :obj)) :val (f (get node :val)))
       (= 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 :do) (f (reduce f acc (get node :statements)) (get node :ret))
       (= op :throw) (f acc (get node :expr))
+      (= op :coerce) (f acc (get node :expr))
       (= op :set-var) (f acc (get node :val))
       (= op :set-field) (f (f acc (get node :obj)) (get node :val))
       (= op :defmacro) (f acc (get node :fn))

jolt-core/jolt/passes.clj

@@ -13,7 +13,7 @@
   :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? record-shapes]]
+  (:require [jolt.host :refer [inline-enabled? record-shapes stash-inline!]]
             [jolt.passes.fold :refer [const-fold]]
             [jolt.passes.numeric :as numeric]
             [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.
 (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
   "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),
@@ -40,6 +51,10 @@
   numeric/annotate runs last in both branches (hint-directed fl*/fx* arithmetic);
   it benefits open builds too, so it is not gated on inlining."
   [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
     (if (inline-enabled? ctx)
       (let [_ (set-rec-shapes! (record-shapes ctx))   ;; record ctor fold

jolt-core/jolt/passes/inline.clj

@@ -4,7 +4,7 @@
   share the alpha-rename invariant (every spliced binder is made globally fresh)
   and the `dirty` fixpoint flag. Portable Clojure (compiler-tier)."
   (: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?]]))
 
 ;; ---------------------------------------------------------------------------
@@ -47,7 +47,7 @@
   ;; is rejected by body-size below and never inlined or alpha-renamed.
   (or (= op :const) (= op :local) (= op :var) (= op :host) (= op :the-var)
       (= 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)
 
@@ -141,6 +141,8 @@
         (if stash
           (let [params (get stash :params)
                 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)]
             (if (and (= (count params) (count args))
                      (<= (body-size body) inline-budget)
@@ -149,19 +151,25 @@
                     ;; trivial args (local/const) substitute straight in (copy
                     ;; propagation); the rest get a fresh local bound once in a
                     ;; 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 []]
                           (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)
+                                :else (let [f (fresh p)]
+                                        (recur (inc i) (assoc env p {:op :local :name f})
                                                (conj binds [f a])))))
                             [env binds]))
                     env (nth res 0)
                     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!)
                 (if (= 0 (count binds))
                   rbody

jolt-core/jolt/passes/numeric.clj

@@ -66,6 +66,41 @@
       (= op :do) (recur-kinds (get node :ret) tenv)
       :else [])))
 
+;; The recur-arg NODE lists for the recurs at THIS loop level (structural, no env),
+;; parallel to recur-kinds. Used to recognise a counter.
+(defn- recur-arg-lists [node]
+  (let [op (get node :op)]
+    (cond
+      (= op :recur) [(get node :args)]
+      (= op :let) (recur-arg-lists (get node :body))
+      (= op :if) (concat (recur-arg-lists (get node :then)) (recur-arg-lists (get node :else)))
+      (= op :do) (recur-arg-lists (get node :ret))
+      :else [])))
+
+;; Is `arg` an increment-style step of loop var `vname`: the var unchanged, or
+;; inc/dec/unchecked-inc/dec, or (+/- var <int-literal>)? Bounded growth that a
+;; fixnum-range counter can sustain for any realistic loop — unlike (* acc x), which
+;; overflows fast, so a multiplicative accumulator never qualifies and stays
+;; arbitrary-precision.
+(defn- counter-step? [arg vname]
+  (cond
+    (and (= :local (get arg :op)) (= vname (get arg :name))) true
+    (= :invoke (get arg :op))
+    (let [f (get arg :fn) as (get arg :args)]
+      (and (= :var (get f :op)) (= "clojure.core" (get f :ns))
+           (let [nm (get f :name)
+                 v? (fn [n] (and (= :local (get n :op)) (= vname (get n :name))))]
+             (cond
+               (and (contains? #{"inc" "dec" "unchecked-inc" "unchecked-dec"} nm) (= 1 (count as)))
+               (v? (nth as 0))
+               (and (contains? #{"+" "unchecked-add"} nm) (= 2 (count as)))
+               (or (and (v? (nth as 0)) (int-lit? (nth as 1)))
+                   (and (v? (nth as 1)) (int-lit? (nth as 0))))
+               (and (contains? #{"-" "unchecked-subtract"} nm) (= 2 (count as)))
+               (and (v? (nth as 0)) (int-lit? (nth as 1)))
+               :else false))))
+    :else false))
+
 ;; Loop-var kinds by bounded fixpoint. A var keeps its init kind (:double or :long)
 ;; only if every recur arg in that slot is the same kind (under the current
 ;; assumption) — a monotone demotion that stops at a fixpoint, bounded by the var
@@ -73,8 +108,8 @@
 ;; keeps arbitrary precision (no :long from a bare literal). A typed loop var's init
 ;; and recur args are all flonums/fixnums (a :long init flows from a coerced ^long
 ;; value or an fx op), so no entry coercion is needed here, unlike a fn param.
-(defn- loop-kinds [names ik body tenv]
+(defn- loop-kinds [names seed body tenv]
-  (loop [cur ik iter 0]
+  (loop [cur seed iter 0]
     (if (> iter (count names))
       cur
       (let [te (reduce (fn [t i] (assoc t (nth names i) (nth cur i))) tenv (range (count names)))
@@ -137,6 +172,7 @@
     (cond
       (= op :const) [(if (float-lit? node) :double nil) 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 :let)
       (let [res (reduce (fn [acc b]
@@ -151,7 +187,24 @@
       (let [binds (get node :bindings)
             names (mapv (fn [b] (nth b 0)) binds)
             ik (mapv (fn [b] (nth (an (nth b 1) tenv) 0)) binds)
-            lk (loop-kinds names ik (get node :body) tenv)
+            rlists (recur-arg-lists (get node :body))
+            ;; seed each var: an already-typed init keeps its kind; an integer-literal
+            ;; init whose recur args are all counter steps is a fixnum counter (:long).
+            seed (mapv (fn [j]
+                         (let [k (nth ik j) b (nth binds j)]
+                           (cond
+                             k k
+                             ;; an int-literal var is a fixnum counter only in a real
+                             ;; iterating loop (>= 1 recur) whose every step is bounded.
+                             ;; A recur-less loop is a let — its int literal stays
+                             ;; generic (arbitrary precision), like a let binding.
+                             (and (seq rlists)
+                                  (int-lit? (nth b 1))
+                                  (every? (fn [args] (counter-step? (nth args j) (nth b 0))) rlists))
+                             :long
+                             :else nil)))
+                       (range (count names)))
+            lk (loop-kinds names seed (get node :body) tenv)
             te (reduce (fn [t i] (assoc t (nth names i) (nth lk i))) tenv (range (count names)))]
         [nil (assoc node
                     :bindings (mapv (fn [b] [(nth b 0) (nth (an (nth b 1) tenv) 1)]) binds)

test/chez/inline-test.ss

@@ -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))

test/chez/numeric-test.ss

@@ -30,6 +30,7 @@
   (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"))
 
 ;; --- emission: ^double -> fl-ops, ^long -> fx-ops ---
 (let ((e (emitf "u" "(fn* ([^double a ^double b] (+ (* a a) (* b b))))")))
@@ -64,6 +65,25 @@
 ;; precision (no fx* overflow).
 (let ((e (emitf "u" "(fn* ([] (loop [acc 1 i 1] (if (< i 25) (recur (* acc i) (inc i)) acc))))")))
   (ok "loop integer accumulator is NOT fx-specialized" (not (has? e "(fx*"))))
+;; a literal-init increment counter types as a fixnum (fx1+), even with no hint.
+(let ((e (emitf "u" "(fn* ([] (loop [i 0] (if (< i 5) (recur (inc i)) i))))")))
+  (ok "literal-init increment counter lowers to fx1+" (has? e "(fx1+")))
+;; but a multiplicative accumulator in the SAME loop stays generic (bignum-safe);
+;; only the counter types.
+(let ((e (emitf "u" "(fn* ([] (loop [acc 1 i 0] (if (< i 100) (recur (* acc i) (inc i)) acc))))")))
+  (ok "counter beside a * accumulator: counter is fx1+" (has? e "(fx1+"))
+  (ok "the * accumulator is NOT fx-specialized (bignum-safe)" (not (has? e "(fx*"))))
+(ok "counter+bignum-accumulator stays exact (1*2*...*99 is a bignum)"
+    (jolt-truthy? (ev "(< 1000000000000000000000 ((fn* ([] (loop [acc 1 i 1] (if (< i 100) (recur (* acc i) (inc i)) acc))))))")))
+(ok "increment counter runtime: counts to 1000"
+    (= 1000 (jnum->exact (ev "((fn* ([] (loop [i 0] (if (< i 1000) (recur (inc i)) i)))))"))))
+;; a recur-less loop is a let: its int-literal binding stays generic (no fx), so
+;; arbitrary precision is preserved (matches (let [i 5] ...)).
+(let ((e (emitf "u" "(fn* ([] (loop [i 5] (+ i 9223372036854775807))))")))
+  (ok "recur-less loop int binding is NOT fx-typed" (not (has? e "(fx+"))))
+(ok "recur-less loop with a big add stays exact (bignum)"
+    (jolt-truthy? (ev "(< 9223372036854775807 ((fn* ([] (loop [i 5] (+ i 9223372036854775807))))))")))
+
 ;; a ^long-seeded loop accumulator IS fx-typed (the hint is a fixnum promise, and
 ;; the value flows from a coerced ^long param).
 (let ((e (emitf "u" "(fn* ([^long start] (loop [acc start] (if (< acc 100) (recur (inc acc)) acc))))")))
@@ -85,7 +105,6 @@
   (ok "long division is NOT specialized (stays generic /)" (not (has? e "(fx"))))
 
 ;; --- runtime values match the generic result ---
-(define (ev s) (jolt-compile-eval s "u"))
 (ok "double dot: 3^2+4^2 = 25" (= 25 (jnum->exact (ev "((fn* ([^double a ^double b] (+ (* a a) (* b b)))) 3.0 4.0)"))))
 (ok "long sum: 2+3 = 5"        (= 5  (jnum->exact (ev "((fn* ([^long a ^long b] (+ a b))) 2 3)"))))
 (ok "double compare true"      (jolt-truthy? (ev "((fn* ([^double x] (< x 5.0))) 3.0)")))