b5075b73be
Adds two IR passes to jolt.passes that run when a unit opts into direct-linking (JOLT_DIRECT_LINK=1, off by default). The inline pass splices small direct-linked fns at their call sites, copy-propagating trivial args so that scalar replacement can then see map literals across the call boundary. Scalar replacement is AOT escape analysis: a map allocation whose only use is constant-keyword lookup is dropped and each (:k m) is replaced with the value at :k, both for a literal lookup subject and for a non-escaping let-bound map. Inlining and scalar replacement iterate to a capped fixpoint, since inlining exposes literals that scalar replacement then collapses. The back end stashes the body IR of each single-fixed-arity defn on its var cell (inline-stash!), and the portable pass reads it through two new jolt.host contract fns (inline-enabled?, inline-ir). Inlining is gated on :inline?, which is off for all of init so core and the self-hosted compiler compile exactly as before (const-fold only); api/init and main re-read JOLT_DIRECT_LINK so the flag works both for a freshly built context and for the build-time-baked one in the shipped binary. Only inline-safe targets are spliced: a single fixed arity, no recur/loop/fn/ try crossing the boundary, within a size budget, a closed body (no free locals beyond the params, so a self-recursive fn's name reference can't dangle), and not ^:redef / ^:dynamic. Bodies are fully alpha-renamed so no spliced name can collide with a caller local. On the ray tracer this is 15.3s -> 13.0s (1.18x). The ceiling is honest: that workload's cost is dominated by lookups on maps that genuinely escape (rays, hits, materials) and by dynamic dispatch (the reduce closure, the :scatter fn), which escape analysis cannot remove. On allocation-bound code where the temporaries are local it is far larger: a vec3 reflect+dot loop goes 9.3s -> 0.38s (25x), with the loop body reduced to pure arithmetic. Verified: full jpm test passes (inline off, no regression); conformance 335/335 in all three modes and the clojure-test-suite both pass with inline on; new inline-sra-test pins the transform and its semantics.
64 lines
3.7 KiB
Text
64 lines
3.7 KiB
Text
# Inline + scalar-replacement passes (jolt-87f, Route 1 AOT escape analysis).
|
|
# When a unit opts into direct-linking (:inline?, JOLT_DIRECT_LINK=1), the IR
|
|
# pipeline inlines small direct-linked fns and then scalar-replaces the now-
|
|
# exposed non-escaping map allocations: (:r {:r a ..}) -> a. This pins the
|
|
# transform (allocations actually vanish) AND that it stays semantics-preserving.
|
|
(import ../../src/jolt/api :as api)
|
|
(import ../../src/jolt/backend :as backend)
|
|
(import ../../src/jolt/reader :as reader)
|
|
|
|
(print "Inline + scalar replacement (jolt-87f)...")
|
|
|
|
# A ctx with inlining ON (independent of the build-time JOLT_DIRECT_LINK).
|
|
(def ctx (api/init {:compile? true}))
|
|
(put (ctx :env) :direct-linking? true)
|
|
(put (ctx :env) :inline? true)
|
|
(api/eval-string ctx "(ns rt)")
|
|
(each s ["(defn v3 [r g b] {:r r :g g :b b})"
|
|
"(defn scale [l n] {:r (* (:r l) n) :g (* (:g l) n) :b (* (:b l) n)})"
|
|
"(defn add [l r] {:r (+ (:r l) (:r r)) :g (+ (:g l) (:g r)) :b (+ (:b l) (:b r))})"
|
|
"(defn dot [l r] (+ (+ (* (:r l) (:r r)) (* (:g l) (:g r))) (* (:b l) (:b r))))"
|
|
"(defn sub [l r] {:r (- (:r l) (:r r)) :g (- (:g l) (:g r)) :b (- (:b l) (:b r))})"
|
|
"(defn reflect [v n] (sub v (scale n (* 2.0 (dot v n)))))"
|
|
# self-recursive: must NOT be inlined into callers (its body has a free
|
|
# local — the fn-name self-reference — that would dangle when spliced).
|
|
"(defn countdown [n] (if (< n 1) :done (countdown (- n 1))))"]
|
|
(api/eval-string ctx s))
|
|
|
|
(defn alloc-count [src]
|
|
# struct / build-map literal occurrences in the emitted Janet = surviving map
|
|
# allocations (jolt builds a struct, falling back to build-map-literal).
|
|
(def code (string/format "%p" (backend/emit-ir ctx (backend/analyze-form ctx (reader/parse-string src)))))
|
|
[(length (string/find-all "struct " code))
|
|
(length (string/find-all "build-map" code))])
|
|
|
|
# A vec3 chain whose intermediates never escape collapses to ONE result map.
|
|
(let [[s b] (alloc-count "(fn [v n] (reflect v n))")]
|
|
(assert (= 1 s) (string "reflect keeps exactly one alloc, got " s " struct"))
|
|
(assert (= 1 b) (string "reflect keeps exactly one build-map fallback, got " b)))
|
|
|
|
# A fully consumed chain (result not returned as a map) allocates NOTHING.
|
|
(let [[s b] (alloc-count "(fn [v n] (dot (reflect v n) (reflect v n)))")]
|
|
(assert (= 0 s) (string "fully-consumed chain allocates no struct, got " s))
|
|
(assert (= 0 b) (string "fully-consumed chain has no build-map fallback, got " b)))
|
|
|
|
# Loop bodies optimize too (recur is not a blanket escape).
|
|
(let [[s _] (alloc-count "(fn [k] (loop [i 0 acc 0.0] (if (< i k) (recur (inc i) (+ acc (dot (v3 1.0 2.0 3.0) (v3 0.1 0.2 0.3)))) acc)))")]
|
|
(assert (= 0 s) (string "loop body allocates no struct, got " s)))
|
|
|
|
# Correctness: inlined results match the obvious computation.
|
|
(assert (= 32.0 (api/eval-string ctx "(dot (v3 1.0 2.0 3.0) (v3 4.0 5.0 6.0))")) "dot value")
|
|
(assert (= 9.0 (api/eval-string ctx "(:r (add (v3 1.0 0.0 0.0) (scale (v3 4.0 0.0 0.0) 2.0)))")) "add+scale value")
|
|
# the self-recursive fn still runs (the closed-body guard kept it un-inlined)
|
|
(assert (= :done (api/eval-string ctx "(countdown 5)")) "recursive fn still works")
|
|
|
|
# A redefinable (^:redef) callee must NOT be inlined — it stays a live var call.
|
|
(api/eval-string ctx "(defn ^:redef wobble [x] {:v x})")
|
|
(let [[s _] (alloc-count "(fn [] (:v (wobble 1)))")]
|
|
# wobble is not inlined, so its map isn't visible to scalar replacement: the
|
|
# lookup stays a call, and the (:v ...) result is whatever wobble returns.
|
|
(assert (= 7 (do (api/eval-string ctx "(defn ^:redef wobble [x] {:v (+ x 6)})")
|
|
(api/eval-string ctx "(:v (wobble 1))")))
|
|
"redef callee stays live (redefinition is visible)"))
|
|
|
|
(print "Inline + scalar replacement passed!")
|