simon/jolt
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Propagate declared record hints through the whole-program fixpoint (jolt-3ko) (#134)

Browse source 
A ^Record param hint was applied only at the final re-emit (reinfer-def), not
during the inter-procedural fixpoint. So a hinted param with no callers stayed
:any while inference ran, and a field read off it (e.g. (:origin ^Ray r)) never
told a non-inlined callee that its arg is a Vec3 — the callee's params stayed
unproven and its field reads kept the dynamic guard.

Seed declared hints as a param-type floor in the fixpoint: phint-seed (passes/
types) resolves an arity's :phints to positional record types via the
record-shapes registry, and infer-unit! initializes each fn's fresh param slots
from them instead of nil. A fixed declared type can't poison the least-fixpoint
the way an early-iteration :any would, and a hinted param now propagates its
(and its field reads') types to its callees during inference.

Scope: this closes the hinted-propagation gap. It does NOT help the ray tracer,
which uses zero ^-hinted params (only hinted fields) — its remaining type gap is
unhinted record-param inference on recursive/non-inlined hot fns, and per the
jolt-15jq A/B it's allocation-bound regardless (jolt-8flj). Tracked on the bead.

Co-authored-by: Yogthos <yogthos@gmail.com>
This commit is contained in:
Dmitri Sotnikov 2026-06-16 03:38:03 +00:00 • committed by GitHub
parent 30a12f39ff
commit bd9e542c78
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 116 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

2
jolt-core/jolt/passes.clj
View file

@ -17,7 +17,7 @@
[jolt.passes.fold :refer [const-fold]]
[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
check-form infer-body reinfer-def phint-seed
set-rtenv! set-vtypes! join-types
set-record-shapes! set-map-shapes! set-protocol-methods!
reset-escapes! collected-escapes

18
jolt-core/jolt/passes/types.clj
View file

@ -791,6 +791,24 @@
(get fnode :arities))))
def-node)))
(defn phint-seed
"Positional declared-hint type seeds for a fn arity. Given the param-name
vector and the arity's :phints (a seq of [name ctor-key] pairs), return a
vector parallel to params whose slot i is the resolved record TYPE of that
param's ^Record hint (via the record-shapes registry), or nil. The
whole-program fixpoint seeds these as a param-type FLOOR so a declared hint
propagates to a fn's callees DURING inference — not only at the final re-emit
(reinfer-def). Without it a hinted param with no callers stays :any through the
fixpoint, so a field read off it (e.g. (:origin ^Ray r)) never tells a shared
callee its arg is a Vec3 (jolt-3ko)."
[params phints]
(let [m (reduce (fn [acc pr] (assoc acc (nth pr 0) (nth pr 1))) {} phints)]
(mapv (fn [nm]
(let [ck (get m nm)
e (and ck (get @record-shapes-box ck))]
(when e (record-type-from-entry e type-depth))))
params)))
;; Piggyback checking (jolt audit). In direct-link mode infer-top already runs
;; one inference pass for specialization; turning checking? on during it makes
;; the success checker nearly free there (no extra traversal — just the

23
src/jolt/backend.janet
View file

@ -1023,6 +1023,7 @@
(def f-set-rshapes (and pns (ns-find pns "set-record-shapes!"))) # jolt-t34
(def f-set-mshapes (and pns (ns-find pns "set-map-shapes!"))) # jolt-t34
(def f-set-pmethods (and pns (ns-find pns "set-protocol-methods!"))) # jolt-41m
(def f-phint-seed (and pns (ns-find pns "phint-seed"))) # jolt-3ko
(def f-get-esc (and pns (ns-find pns "collected-escapes")))
(def report @{})
(when (and f-set-rtenv f-set-vtypes f-join f-infer-body f-reinfer f-reset-esc f-get-esc)
@ -1048,7 +1049,8 @@
(unless (ar :rest)
(def pv (vview (ar :params)))
(def rec @{:key key :cell v :def d :params (ar :params) :body (ar :body)
:np (length pv) :pt (array/new-filled (length pv)) :ret nil})
:phints (ar :phints) :np (length pv)
:pt (array/new-filled (length pv)) :ret nil})
(array/push fns rec)
(put by-key key rec)
# a fn value is non-nil -> :truthy (sealed root in opt mode)
@ -1060,6 +1062,13 @@
# jolt-t34: feed record-ctor shapes + the map-shaping flag to the inference
(when f-set-rshapes ((var-get f-set-rshapes) (or (get (ctx :env) :record-shapes) @{})))
(when f-set-mshapes ((var-get f-set-mshapes) (get (ctx :env) :map-shapes?)))
# jolt-3ko: resolve each fn's declared ^Record param hints to positional
# type seeds (needs the registry above). Seeded as a param-type FLOOR in the
# fixpoint so a hinted param propagates its (field-read) types to callees
# during inference, not only at the final re-emit.
(when f-phint-seed
(each f fns
(put f :phint-types (vview ((var-get f-phint-seed) (f :params) (or (f :phints) []))))))
# jolt-41m: feed the protocol-method registry for devirtualization
(when f-set-pmethods ((var-get f-set-pmethods) (or (get (ctx :env) :protocol-methods) @{})))
# --- param/return/value-type fixpoint (chaotic iteration to LEAST fixpoint) ---
@ -1083,9 +1092,17 @@
# infer each def's VALUE type from its init
(each dv defs
(put dv :tvt (in (vview ((var-get f-infer-body) (dv :init) @{})) 0)))
# recompute param types FRESH (start at bottom = nil) from this round's calls
# recompute param types FRESH (start at bottom) from this round's calls.
# Bottom is nil, EXCEPT a declared ^Record hint seeds its slot as a floor
# (jolt-3ko) — a fixed declared type can't poison the fixpoint the way an
# early-iteration :any would, and it lets a hinted param propagate to its
# callees even when it has no callers of its own.
(def newpt @{})
(each f fns (put newpt (f :key) (array/new-filled (f :np))))
(each f fns
(def npa (array/new-filled (f :np)))
(def phts (f :phint-types))
(when phts (for i 0 (f :np) (when (in phts i) (put npa i (in phts i)))))
(put newpt (f :key) npa))
(each f fns
(each c (vview (f :tcalls))
(def cv (vview c))

77
test/integration/phint-propagation-test.janet Normal file
View file

@ -0,0 +1,77 @@
# Declared record param hints propagate THROUGH the whole-program fixpoint
# (jolt-3ko). A ^Record param's field reads type to the field's record type, and
# when such a field-read value is passed to a SHARED helper (no hints, inferred
# from call sites), the helper's params must pick up that record type so its own
# field reads bare-index. Before the fix, ^-hints were applied only at the final
# re-emit (reinfer-def), not during the fixpoint, so a hinted param with no
# callers stayed :any during inference and never propagated to its callees —
# exactly why the ray tracer's vec ops (called with (:origin ray) etc.) stayed
# unproven even under whole-program optimization.
(import ../../src/jolt/api :as api)
(import ../../src/jolt/backend :as backend)
(import ../../src/jolt/types :as types)
(print "phint propagation through the fixpoint (jolt-3ko)...")
(def dir (string (or (os/getenv "TMPDIR") "/tmp") "/jolt-phint-prop"))
(os/mkdir dir)
(spit (string dir "/pp.clj")
(string
"(ns pp)\n"
"(defrecord Vec3 [r g b])\n"
"(defrecord Ray [^Vec3 origin ^Vec3 direction])\n"
# shared ops, NO hints — must be inferred from the call sites below. A
# `loop` makes them inline-INELIGIBLE so they stay real calls (a small fn
# would be spliced into the caller, where the caller's hint already proves
# it); the point here is propagation to a separate callee's PARAMS.
"(defn dot [a b]\n"
" (loop [i 0 s 0.0]\n"
" (if (< i 1) (recur (inc i) (+ (+ (* (:r a) (:r b)) (* (:g a) (:g b))) (* (:b a) (:b b)))) s)))\n"
"(defn add [a b]\n"
" (loop [i 0 s nil]\n"
" (if (< i 1) (recur (inc i) (->Vec3 (+ (:r a) (:r b)) (+ (:g a) (:g b)) (+ (:b a) (:b b)))) s)))\n"
# callers pass field-read Vec3s from a ^Ray param into the shared ops
"(defn use-dot [^Ray r] (dot (:origin r) (:direction r)))\n"
"(defn use-add [^Ray r] (add (:origin r) (:direction r)))\n"))
(os/setenv "JOLT_DIRECT_LINK" "1")
(os/setenv "JOLT_WHOLE_PROGRAM" "1")
(os/setenv "JOLT_PATH" dir)
(def ctx (api/init {:compile? true}))
(api/eval-string ctx "(require '[pp])")
(def report (backend/infer-program! ctx))
(def pns (types/ctx-find-ns ctx "jolt.passes"))
(def reinfer (types/var-get (types/ns-find pns "reinfer-def")))
(def ppns (types/ctx-find-ns ctx "pp"))
(defn ptmap-for [fname params]
(def pts (get report (string "pp/" fname)))
(def m @{}) (when pts (var i 0) (each p params (put m p (get pts i)) (++ i))) m)
(defn guards [fname params]
(def cell (get (ppns :mappings) fname))
(length (string/find-all ":jolt/type"
(string/format "%p" (backend/emit-ir ctx (reinfer (get cell :infer-ir) (ptmap-for fname params)))))))
(var fails 0)
(defn check [label got expected]
(if (= got expected) (print " ok " label)
(do (++ fails) (printf " FAIL %s: want %p got %p" label expected got))))
# the shared ops' params get the Vec3 record type from the field-read call args,
# so all their field reads bare-index (no :jolt/type guard)
(check "dot params typed Vec3 -> reads bare" (guards "dot" ["a" "b"]) 0)
(check "add params typed Vec3 -> reads bare" (guards "add" ["a" "b"]) 0)
# the hinted callers were already fine (re-emit applies the phint)
(check "use-dot hinted reads bare" (guards "use-dot" ["r"]) 0)
# the report shows the shared ops' params as a Vec3 struct (has :shape / :type)
(def dot-a (get (get report "pp/dot") 0))
(check "dot param a is a struct type" (truthy? (and dot-a (or (get dot-a :shape) (get dot-a :struct)))) true)
# correctness: results are unchanged
(check "dot computes" (api/eval-string ctx "(pp/dot (pp/->Vec3 1 2 3) (pp/->Vec3 4 5 6))") 32)
(check "use-dot computes"
(api/eval-string ctx "(pp/use-dot (pp/->Ray (pp/->Vec3 1 2 3) (pp/->Vec3 4 5 6)))") 32)
(if (> fails 0) (do (printf "phint-propagation: %d FAILED" fails) (os/exit 1))
(print "phint propagation (jolt-3ko) passed!"))