mirror of
https://gitlvb.teallvbs.xyz/IQ.Lvbs/IQ.Pilot.git
synced 2026-07-16 04:42:05 +08:00
IQ.Lvbs CI [bot]
62 lines
2.5 KiB
Python
62 lines
2.5 KiB
Python
"""On-device finalize for a host-compiled supercombo pkl.
|
|
|
|
Host compiles with local-size tuning disabled (workgroup sizes can only be
|
|
chosen by timing real kernels). This rebakes them by running tinygrad's
|
|
pm_optimize_local_size pass on the real GPU, then re-pickles. Takes well
|
|
under a minute; no scheduling or kernel compilation happens here.
|
|
|
|
sudo env ... DEV=QCOM WARP_DEV=QCOM IMAGE=2 FLOAT16=1 NOLOCALS=1 \
|
|
JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1 python finalize_pkl.py <in.pkl> <out.pkl>
|
|
"""
|
|
import dataclasses, pickle, sys, time
|
|
|
|
from tinygrad.engine.realize import optimize_local_size, run_linear
|
|
from tinygrad.uop.ops import Ops
|
|
|
|
pkl_in, pkl_out = sys.argv[1], sys.argv[2]
|
|
with open(pkl_in, "rb") as f:
|
|
out = pickle.load(f)
|
|
|
|
# replay the sliced eager weight-prep chains with device codegen first; the
|
|
# host ran them with container clang/tinygrad codegen (rounding drift)
|
|
st = time.perf_counter()
|
|
pre = out.pop("_pretime", [])
|
|
for lin, vv in pre:
|
|
run_linear(lin, vv)
|
|
if pre:
|
|
print(f"pretime: ran {sum(len(l.src) for l, _ in pre)} calls in {time.perf_counter() - st:.1f}s")
|
|
|
|
# execute the pruned one-time kernels (IMAGE weight packing etc.) on the real
|
|
# GPU; on the NULL host they never ran, so their baked outputs are zeros
|
|
for k, linears in out.pop("_onetime", {}).items():
|
|
st = time.perf_counter()
|
|
for lin in linears:
|
|
run_linear(lin, {})
|
|
print(f"{k}: ran {len(linears)} one-time linear(s) in {time.perf_counter() - st:.1f}s")
|
|
|
|
for k, v in out.items():
|
|
cap = getattr(v, "captured", None)
|
|
if cap is None:
|
|
continue
|
|
st = time.perf_counter()
|
|
# CALL(PROGRAM) nodes live inside batched CUSTOM_FUNCTION "graph" nodes,
|
|
# which graph_rewrite/substitute treat as opaque; rebuild bottom-up instead.
|
|
changed: dict = {}
|
|
tuned = 0
|
|
for u in cap.linear.toposort():
|
|
nu = u.replace(src=ns) if (ns := tuple(changed.get(s, s) for s in u.src)) != u.src else u
|
|
if nu.op is Ops.CALL and nu.src and nu.src[0].op is Ops.PROGRAM and nu.src[0].arg.local_size is None:
|
|
if (r := optimize_local_size(nu, nu.src[0])) is not None:
|
|
nu, tuned = r, tuned + 1
|
|
if nu is not u:
|
|
changed[u] = nu
|
|
new_linear = changed.get(cap.linear, cap.linear)
|
|
left = sum(1 for u in new_linear.toposort() if u.op is Ops.PROGRAM and u.arg.local_size is None)
|
|
assert left == 0, f"{k}: {left} kernels still unbaked"
|
|
v.captured = dataclasses.replace(cap, linear=new_linear)
|
|
print(f"{k}: rebaked {tuned} calls in {time.perf_counter() - st:.1f}s")
|
|
|
|
with open(pkl_out, "wb") as f:
|
|
pickle.dump(out, f)
|
|
print(f"finalized -> {pkl_out}")
|