IQ.Pilot Prebuilt Release @ 27f668a
This commit is contained in:
IQ.Lvbs CI [bot]
12
iqpilot/selfdrive/iqmodeld/tools/compile_daemon.py
Normal file
12
iqpilot/selfdrive/iqmodeld/tools/compile_daemon.py
Normal file
@@ -0,0 +1,12 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from iqpilot.selfdrive.iqmodeld.tools.daemon_jit_compiler import main
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
471
iqpilot/selfdrive/iqmodeld/tools/compile_egpu_model.py
Normal file
471
iqpilot/selfdrive/iqmodeld/tools/compile_egpu_model.py
Normal file
@@ -0,0 +1,471 @@
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import os
|
||||
import pickle
|
||||
import sys
|
||||
import time
|
||||
|
||||
os.environ.setdefault("FLOAT16", "1")
|
||||
os.environ.setdefault("JIT_BATCH_SIZE", "0")
|
||||
os.environ.setdefault("GMMU", "0")
|
||||
# TC_OPT=2 lets tinygrad pick tensor-core kernels; on some models a TC kernel miscompiles and biases
|
||||
# the output (documented on Metal). A parity gate below catches it and re-compiles with TC off.
|
||||
os.environ.setdefault("TC_OPT", "0" if ("--tc-off" in sys.argv or os.environ.get("IQ_EGPU_TC_OFF")) else "2")
|
||||
|
||||
HOST = "--host" in sys.argv
|
||||
if HOST:
|
||||
from iqpilot.selfdrive.iqmodeld.tools.egpu_host_mock import DEFAULT_ARCH, activate
|
||||
activate(sys.argv[sys.argv.index("--arch") + 1] if "--arch" in sys.argv else DEFAULT_ARCH)
|
||||
os.environ.setdefault("DEV", "USB+AMD:LLVM")
|
||||
|
||||
import numpy as np
|
||||
|
||||
from iqpilot.selfdrive.iqmodeld.egpu_helpers import egpu_pkl_path, local_onnx, patch_tinygrad_fetch_fw
|
||||
from iqpilot.selfdrive.iqmodeld.egpu_model import EGPU_MODELS, get_egpu_model, resolve_egpu_model
|
||||
from iqpilot.selfdrive.iqmodeld.temporal_state import MODEL_INPUT_SPEC, spec_from_meta
|
||||
|
||||
INPUT_SPEC = dict(MODEL_INPUT_SPEC)
|
||||
|
||||
patch_tinygrad_fetch_fw()
|
||||
|
||||
SEED = 42
|
||||
|
||||
|
||||
class _ParityFail(RuntimeError):
|
||||
pass
|
||||
KERNEL_PROGRESS_SCALE = 260.0
|
||||
|
||||
|
||||
def _progress_sampler(param: str, base: float, span: float, stop) -> None:
|
||||
import math
|
||||
|
||||
from iqpilot.common.params import Params
|
||||
from tinygrad.helpers import GlobalCounters
|
||||
pm = Params()
|
||||
last = -1.0
|
||||
while not stop.wait(0.5):
|
||||
kernels = float(getattr(GlobalCounters, "kernel_count", 0))
|
||||
value = base + span * (1.0 - math.exp(-kernels / KERNEL_PROGRESS_SCALE))
|
||||
if value - last >= 0.01:
|
||||
last = value
|
||||
pm.put(param, f"{min(base + span, value):.3f}")
|
||||
|
||||
|
||||
def set_input_spec(meta: dict) -> None:
|
||||
spec = spec_from_meta(meta)
|
||||
if spec is not None:
|
||||
INPUT_SPEC.clear()
|
||||
INPUT_SPEC.update(spec)
|
||||
|
||||
|
||||
def make_run_model(model_runner):
|
||||
def run_model(**inputs):
|
||||
out = next(iter(model_runner({k: inputs[k] for k in INPUT_SPEC}).values())).cast("float32")
|
||||
return out.reshape(-1),
|
||||
return run_model
|
||||
|
||||
|
||||
def _random_inputs(seed: int):
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.tensor import Tensor
|
||||
rng = np.random.default_rng(seed)
|
||||
out = {}
|
||||
for name, (shape, dtype) in INPUT_SPEC.items():
|
||||
if dtype == "uint8":
|
||||
arr = rng.integers(0, 256, shape).astype(np.uint8)
|
||||
else:
|
||||
arr = rng.standard_normal(shape).astype(np.float32)
|
||||
out[name] = Tensor(arr, device=Device.DEFAULT).realize()
|
||||
return out
|
||||
|
||||
|
||||
def _run(fn, seed: int) -> np.ndarray:
|
||||
from tinygrad.device import Device
|
||||
st = time.perf_counter()
|
||||
outs = fn(**_random_inputs(seed))
|
||||
Device.default.synchronize()
|
||||
print(f" run(seed={seed}) {(time.perf_counter() - st) * 1e3:6.1f} ms")
|
||||
return outs[0].numpy().reshape(-1)
|
||||
|
||||
|
||||
def compile_model(meta: dict, onnx_path: str, out_path: str) -> str:
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.engine.jit import TinyJit
|
||||
from tinygrad.nn.onnx import OnnxRunner
|
||||
|
||||
if meta.get("split"):
|
||||
raise RuntimeError(f"model {meta['key']} is a split model; eGPU v1 compiles fused models only")
|
||||
|
||||
jit = TinyJit(make_run_model(OnnxRunner(onnx_path)), prune=True)
|
||||
|
||||
print("capture + replay")
|
||||
for _ in range(2):
|
||||
baseline = _run(jit, SEED)
|
||||
if baseline.shape[0] != meta["output_len"]:
|
||||
raise RuntimeError(f"model output length {baseline.shape[0]} != registry {meta['output_len']}")
|
||||
if not np.isfinite(baseline).all():
|
||||
raise RuntimeError("compiled model produced non-finite outputs")
|
||||
|
||||
bundle = {
|
||||
"run_model": jit,
|
||||
"model_key": meta["key"],
|
||||
"model_sha256": meta["sha256"],
|
||||
"output_len": int(meta["output_len"]),
|
||||
"frame_skip": int(meta["frame_skip"]),
|
||||
"input_spec": {name: (tuple(shape), dtype) for name, (shape, dtype) in INPUT_SPEC.items()},
|
||||
"input_device": Device.DEFAULT,
|
||||
}
|
||||
os.makedirs(os.path.dirname(out_path), exist_ok=True)
|
||||
tmp = out_path + ".part"
|
||||
print("serialize")
|
||||
with open(tmp, "wb") as f:
|
||||
pickle.dump(bundle, f, protocol=pickle.HIGHEST_PROTOCOL)
|
||||
|
||||
del bundle, jit
|
||||
gc.collect()
|
||||
|
||||
print("reload + validate")
|
||||
with open(tmp, "rb") as f:
|
||||
jit = pickle.load(f)["run_model"]
|
||||
if not np.array_equal(_run(jit, SEED), baseline):
|
||||
raise RuntimeError("outputs differ from baseline after pickle round trip")
|
||||
if np.array_equal(_run(jit, SEED + 1), baseline):
|
||||
raise RuntimeError("outputs insensitive to inputs after pickle round trip")
|
||||
|
||||
from tinygrad.tensor import Tensor
|
||||
zeros = {name: Tensor(np.zeros(shape, dtype=dtype), device=Device.DEFAULT).realize()
|
||||
for name, (shape, dtype) in INPUT_SPEC.items()}
|
||||
flat = jit(**zeros)[0].numpy().reshape(-1)
|
||||
from iqpilot.selfdrive.iqmodeld.parser import PhaseParser
|
||||
from iqpilot.selfdrive.iqmodeld.tools.compile_supercombo import _slice_outputs, _validate_pose_outputs
|
||||
_validate_pose_outputs(PhaseParser().parse_vision_outputs(_slice_outputs(flat, meta["output_slices"])))
|
||||
|
||||
os.replace(tmp, out_path)
|
||||
return out_path
|
||||
|
||||
|
||||
def _policy_frame(seed: int, input_spec: dict):
|
||||
from tinygrad.tensor import Tensor
|
||||
rng = np.random.default_rng(seed)
|
||||
img = input_spec["img"][0]
|
||||
warped = Tensor(rng.integers(0, 256, (2, 6, img[2], img[3])).astype(np.uint8), device="NPY").realize()
|
||||
return warped
|
||||
|
||||
|
||||
def _tc_off_reference(onnx_path: str, meta: dict, fmt: int = 2, resolutions: tuple[tuple[int, int], ...] = ()):
|
||||
"""Compile+run the model with tensor cores OFF in a child process and return the last of 3
|
||||
policy frames. This is the trusted reference: TC-off kernels are the conservative path the
|
||||
eMac gate also trusts. Used to catch a TC kernel miscompile that would bias steering."""
|
||||
import subprocess
|
||||
import tempfile
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
ref = os.path.join(td, "ref.npz" if fmt == 3 else "ref.npy")
|
||||
env = {k: v for k, v in os.environ.items() if k not in ("TC_OPT", "BEAM")}
|
||||
env["TC_OPT"] = "0"
|
||||
env["IQ_EGPU_REFERENCE"] = ref
|
||||
cmd = [sys.executable, "-m", "iqpilot.selfdrive.iqmodeld.tools.compile_egpu_model",
|
||||
"--model", meta["key"], "--onnx", onnx_path, "--tc-off", "--format", str(fmt)]
|
||||
if resolutions:
|
||||
cmd += ["--camera-resolutions", *(f"{w}x{h}" for w, h in resolutions)]
|
||||
r = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=14400)
|
||||
if r.returncode != 0 or not os.path.isfile(ref):
|
||||
raise RuntimeError(f"parity reference compile failed:\n{r.stderr[-2000:]}")
|
||||
return np.load(ref)
|
||||
|
||||
|
||||
def _parity_check(key: str, got: np.ndarray, ref: np.ndarray, label: str = "") -> None:
|
||||
rel = float(np.abs(got - ref).mean() / max(1e-3, float(np.abs(ref).mean())))
|
||||
if rel > 0.01:
|
||||
raise _ParityFail(f"PARITY FAIL: TC kernels miscompiled {key} {label}(rel={rel:.4f} vs TC-off); recompiling with tensor cores disabled")
|
||||
print(f" parity vs TC-off reference {label}: rel={rel:.6f} OK")
|
||||
|
||||
|
||||
def compile_policy_model(meta: dict, onnx_path: str, out_path: str) -> str:
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.engine.jit import TinyJit
|
||||
from tinygrad.nn.onnx import OnnxRunner
|
||||
|
||||
from iqpilot.selfdrive.iqmodeld.egpu_policy import POLICY_FORMAT, PackedInputs, dump_oob, load_bundle, make_queues, make_run_policy
|
||||
|
||||
if meta.get("split"):
|
||||
raise RuntimeError(f"model {meta['key']} is a split model; eGPU compiles fused models only")
|
||||
input_spec = {name: (tuple(shape), dtype) for name, (shape, dtype) in INPUT_SPEC.items()}
|
||||
frame_skip = int(meta["frame_skip"])
|
||||
device = Device.DEFAULT
|
||||
jit = TinyJit(make_run_policy(OnnxRunner(onnx_path), input_spec, frame_skip, device), prune=True)
|
||||
queues = make_queues(input_spec, frame_skip, device)
|
||||
packed = PackedInputs(input_spec)
|
||||
|
||||
def step(seed: int) -> np.ndarray:
|
||||
packed.views["traffic_convention"][:] = [1, 0]
|
||||
packed.views["action_t"][:] = [0.2, 0.3]
|
||||
st = time.perf_counter()
|
||||
out, = jit(warped=_policy_frame(seed, input_spec), packed_npy_inputs=packed.tensor, **queues)
|
||||
flat = out.numpy().reshape(-1)
|
||||
print(f" policy step(seed={seed}) {(time.perf_counter() - st) * 1e3:6.1f} ms")
|
||||
packed.views["prev_feat"][:] = flat[meta["output_slices"]["hidden_state"]].reshape(packed.views["prev_feat"].shape)
|
||||
return flat
|
||||
|
||||
print("capture + replay")
|
||||
for i in range(3):
|
||||
baseline = step(SEED + i)
|
||||
if baseline.shape[0] != meta["output_len"]:
|
||||
raise RuntimeError(f"model output length {baseline.shape[0]} != registry {meta['output_len']}")
|
||||
if not HOST and not np.isfinite(baseline).all():
|
||||
raise RuntimeError("compiled policy produced non-finite outputs")
|
||||
|
||||
bundle = {
|
||||
"format": POLICY_FORMAT,
|
||||
"run_policy": jit,
|
||||
"model_key": meta["key"],
|
||||
"model_sha256": meta["sha256"],
|
||||
"output_len": int(meta["output_len"]),
|
||||
"frame_skip": frame_skip,
|
||||
"input_spec": input_spec,
|
||||
"input_device": device,
|
||||
}
|
||||
os.makedirs(os.path.dirname(out_path), exist_ok=True)
|
||||
tmp = out_path + ".part"
|
||||
print("serialize (out-of-band buffers)")
|
||||
with open(tmp, "wb") as f:
|
||||
dump_oob(bundle, f)
|
||||
|
||||
del bundle, jit, queues, packed
|
||||
gc.collect()
|
||||
|
||||
print("reload + validate")
|
||||
jit = load_bundle(tmp)["run_policy"]
|
||||
queues = make_queues(input_spec, frame_skip, device)
|
||||
packed = PackedInputs(input_spec)
|
||||
outs = []
|
||||
for i in range(3):
|
||||
packed.views["traffic_convention"][:] = [1, 0]
|
||||
packed.views["action_t"][:] = [0.2, 0.3]
|
||||
out, = jit(warped=_policy_frame(SEED + i, input_spec), packed_npy_inputs=packed.tensor, **queues)
|
||||
flat = out.numpy().reshape(-1)
|
||||
packed.views["prev_feat"][:] = flat[meta["output_slices"]["hidden_state"]].reshape(packed.views["prev_feat"].shape)
|
||||
outs.append(flat)
|
||||
ref_target = os.environ.get("IQ_EGPU_REFERENCE")
|
||||
if ref_target:
|
||||
np.save(ref_target, outs[-1])
|
||||
return out_path
|
||||
if HOST:
|
||||
os.replace(tmp, out_path)
|
||||
return out_path
|
||||
if not np.array_equal(outs[-1], baseline):
|
||||
raise RuntimeError("policy outputs differ from baseline after pickle round trip")
|
||||
if np.array_equal(outs[0], outs[-1]):
|
||||
raise RuntimeError("policy outputs insensitive to inputs after pickle round trip")
|
||||
if not all(np.isfinite(o).all() for o in outs):
|
||||
raise RuntimeError("reloaded policy produced non-finite outputs")
|
||||
from iqpilot.selfdrive.iqmodeld.parser import PhaseParser
|
||||
from iqpilot.selfdrive.iqmodeld.tools.compile_supercombo import _slice_outputs, _validate_pose_outputs
|
||||
_validate_pose_outputs(PhaseParser().parse_vision_outputs(_slice_outputs(outs[-1], meta["output_slices"])))
|
||||
|
||||
if os.environ.get("TC_OPT") != "0" and not os.environ.get("IQ_EGPU_SKIP_PARITY"):
|
||||
_parity_check(meta["key"], outs[-1], _tc_off_reference(onnx_path, meta))
|
||||
|
||||
os.replace(tmp, out_path)
|
||||
return out_path
|
||||
|
||||
|
||||
DEFAULT_CAMERA_RESOLUTIONS: tuple[tuple[int, int], ...] = ((1928, 1208), (1344, 760))
|
||||
|
||||
|
||||
def camera_nv12(cam_w: int, cam_h: int) -> tuple[int, int, int, int, int]:
|
||||
from iqpilot.system.camerad.cameras.nv12_info import get_nv12_info
|
||||
stride, y_height, uv_height, _ = get_nv12_info(cam_w, cam_h)
|
||||
return (cam_w, cam_h, stride, y_height, uv_height)
|
||||
|
||||
|
||||
def _fill_model_frame(packed, seed: int, res: tuple[int, int], model_w: int, model_h: int) -> None:
|
||||
rng = np.random.default_rng(seed)
|
||||
cam_w, cam_h = res
|
||||
scale = np.array([[cam_w / model_w, 0.0, 0.0], [0.0, cam_h / model_h, 0.0], [0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
for name in ("tfm", "big_tfm"):
|
||||
packed.views[name][:, :] = scale * (1.0 + 0.02 * rng.standard_normal((3, 3))).astype(np.float32)
|
||||
for v in packed.frames.values():
|
||||
v[:] = rng.integers(0, 256, size=v.shape, dtype=np.uint8)
|
||||
packed.views["traffic_convention"][:] = [1, 0]
|
||||
packed.views["action_t"][:] = [0.2, 0.3]
|
||||
|
||||
|
||||
def compile_model_v3(meta: dict, onnx_path: str, out_path: str,
|
||||
resolutions: tuple[tuple[int, int], ...] = DEFAULT_CAMERA_RESOLUTIONS) -> str:
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.engine.jit import TinyJit
|
||||
from tinygrad.nn.onnx import OnnxRunner
|
||||
|
||||
from iqpilot.selfdrive.iqmodeld.egpu_policy import (
|
||||
MODEL_FORMAT, dump_oob, load_bundle, make_model_queues, make_run_model, make_run_policy, make_warp, model_size, nv12_copy_size,
|
||||
)
|
||||
|
||||
if meta.get("split"):
|
||||
raise RuntimeError(f"model {meta['key']} is a split model; eGPU compiles fused models only")
|
||||
input_spec = {name: (tuple(shape), dtype) for name, (shape, dtype) in INPUT_SPEC.items()}
|
||||
frame_skip = int(meta["frame_skip"])
|
||||
hidden = meta["output_slices"]["hidden_state"]
|
||||
device = Device.DEFAULT
|
||||
model_w, model_h = model_size(input_spec)
|
||||
runner = OnnxRunner(onnx_path)
|
||||
run_policy = make_run_policy(runner, input_spec, frame_skip, device)
|
||||
|
||||
def step(jit, queues, packed, seed: int, res: tuple[int, int]) -> np.ndarray:
|
||||
_fill_model_frame(packed, seed, res, model_w, model_h)
|
||||
st = time.perf_counter()
|
||||
out, = jit(**queues)
|
||||
flat = out.numpy().reshape(-1)
|
||||
print(f" model step(seed={seed}, {res[0]}x{res[1]}) {(time.perf_counter() - st) * 1e3:6.1f} ms")
|
||||
packed.views["prev_feat"][:] = flat[hidden].reshape(packed.views["prev_feat"].shape)
|
||||
return flat
|
||||
|
||||
def run_three(jit, fcs: int, res: tuple[int, int]) -> list[np.ndarray]:
|
||||
queues, packed = make_model_queues(input_spec, frame_skip, device, fcs)
|
||||
return [step(jit, queues, packed, SEED + i, res) for i in range(3)]
|
||||
|
||||
jits: dict[tuple[int, int], object] = {}
|
||||
sizes: dict[tuple[int, int], int] = {}
|
||||
nv12s: dict[tuple[int, int], tuple[int, int, int, int, int]] = {}
|
||||
baselines: dict[tuple[int, int], np.ndarray] = {}
|
||||
for res in resolutions:
|
||||
nv12 = camera_nv12(*res)
|
||||
fcs = nv12_copy_size(nv12[2], nv12[3], nv12[4])
|
||||
jit = TinyJit(make_run_model(make_warp(nv12, model_w, model_h, device), run_policy, input_spec, fcs, device), prune=True)
|
||||
print(f"capture + replay {res[0]}x{res[1]} (frame copy {fcs} B)")
|
||||
baseline = run_three(jit, fcs, res)[-1]
|
||||
if baseline.shape[0] != meta["output_len"]:
|
||||
raise RuntimeError(f"model output length {baseline.shape[0]} != registry {meta['output_len']}")
|
||||
if not HOST and not np.isfinite(baseline).all():
|
||||
raise RuntimeError("compiled model produced non-finite outputs")
|
||||
jits[res], sizes[res], nv12s[res], baselines[res] = jit, fcs, nv12, baseline
|
||||
|
||||
bundle = {
|
||||
"format": MODEL_FORMAT,
|
||||
"run_model": jits,
|
||||
"frame_copy_size": sizes,
|
||||
"nv12": nv12s,
|
||||
"model_key": meta["key"],
|
||||
"model_sha256": meta["sha256"],
|
||||
"output_len": int(meta["output_len"]),
|
||||
"frame_skip": frame_skip,
|
||||
"input_spec": input_spec,
|
||||
"input_device": device,
|
||||
}
|
||||
os.makedirs(os.path.dirname(out_path), exist_ok=True)
|
||||
tmp = out_path + ".part"
|
||||
print("serialize (out-of-band buffers)")
|
||||
with open(tmp, "wb") as f:
|
||||
dump_oob(bundle, f)
|
||||
|
||||
del bundle, jits, run_policy, runner
|
||||
gc.collect()
|
||||
|
||||
print("reload + validate")
|
||||
loaded = load_bundle(tmp)
|
||||
outs = {res: run_three(loaded["run_model"][res], loaded["frame_copy_size"][res], res) for res in resolutions}
|
||||
ref_target = os.environ.get("IQ_EGPU_REFERENCE")
|
||||
if ref_target:
|
||||
np.savez(ref_target, **{f"{w}x{h}": outs[(w, h)][-1] for (w, h) in resolutions})
|
||||
return out_path
|
||||
if HOST:
|
||||
os.replace(tmp, out_path)
|
||||
return out_path
|
||||
for res in resolutions:
|
||||
if not np.array_equal(outs[res][-1], baselines[res]):
|
||||
raise RuntimeError(f"model outputs differ from baseline after pickle round trip ({res[0]}x{res[1]})")
|
||||
if np.array_equal(outs[res][0], outs[res][-1]):
|
||||
raise RuntimeError(f"model outputs insensitive to inputs after pickle round trip ({res[0]}x{res[1]})")
|
||||
if not all(np.isfinite(o).all() for o in outs[res]):
|
||||
raise RuntimeError(f"reloaded model produced non-finite outputs ({res[0]}x{res[1]})")
|
||||
from iqpilot.selfdrive.iqmodeld.parser import PhaseParser
|
||||
from iqpilot.selfdrive.iqmodeld.tools.compile_supercombo import _slice_outputs, _validate_pose_outputs
|
||||
_validate_pose_outputs(PhaseParser().parse_vision_outputs(_slice_outputs(outs[resolutions[0]][-1], meta["output_slices"])))
|
||||
|
||||
if os.environ.get("TC_OPT") != "0" and not os.environ.get("IQ_EGPU_SKIP_PARITY"):
|
||||
ref = _tc_off_reference(onnx_path, meta, fmt=3, resolutions=resolutions)
|
||||
for (w, h) in resolutions:
|
||||
_parity_check(meta["key"], outs[(w, h)][-1], ref[f"{w}x{h}"], label=f"{w}x{h} ")
|
||||
|
||||
os.replace(tmp, out_path)
|
||||
return out_path
|
||||
|
||||
|
||||
def _parse_resolution(text: str) -> tuple[int, int]:
|
||||
w, h = text.lower().split("x")
|
||||
return int(w), int(h)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--model", default=None, help=f"registry key, one of {sorted(EGPU_MODELS)}")
|
||||
p.add_argument("--onnx", default=None)
|
||||
p.add_argument("--output", default=None)
|
||||
p.add_argument("--progress-param", default=None)
|
||||
p.add_argument("--progress-base", type=float, default=None)
|
||||
p.add_argument("--progress-span", type=float, default=0.0)
|
||||
p.add_argument("--format", type=int, default=3, choices=(1, 2, 3),
|
||||
help="3 = warp on the dock from raw NV12 (comma master); 2 = device-warped policy bundle")
|
||||
p.add_argument("--camera-resolutions", type=_parse_resolution, nargs="+", default=list(DEFAULT_CAMERA_RESOLUTIONS),
|
||||
help="WxH camera sizes bundled into a format-3 artifact")
|
||||
p.add_argument("--host", action="store_true", help="compile on a mock dock (no AMD hardware); outputs need a dock parity gate")
|
||||
p.add_argument("--arch", default=None, help="target gfx arch for --host")
|
||||
p.add_argument("--tc-off", action="store_true", help="disable tensor-core kernels (conservative; auto-set on parity failure)")
|
||||
args = p.parse_args()
|
||||
if args.host and args.format == 1:
|
||||
raise SystemExit("--host supports formats 2 and 3 only")
|
||||
|
||||
if args.model is not None:
|
||||
if args.model in EGPU_MODELS:
|
||||
meta = get_egpu_model(args.model)
|
||||
else:
|
||||
from iqpilot.common.params import Params
|
||||
meta = resolve_egpu_model(Params(), args.model)
|
||||
if meta is None:
|
||||
raise SystemExit(f"unknown model {args.model!r}: not a built-in ({sorted(EGPU_MODELS)}) and not in the synced catalog")
|
||||
else:
|
||||
meta = get_egpu_model()
|
||||
set_input_spec(meta)
|
||||
|
||||
onnx_path = args.onnx or local_onnx(meta)
|
||||
if onnx_path is None or not os.path.isfile(onnx_path):
|
||||
raise SystemExit(f"onnx not found for {meta['key']}; pass --onnx or let iqegpumodeld download it first")
|
||||
|
||||
stop = None
|
||||
sampler = None
|
||||
if args.progress_param and args.progress_base is not None:
|
||||
import threading
|
||||
stop = threading.Event()
|
||||
sampler = threading.Thread(target=_progress_sampler,
|
||||
args=(args.progress_param, args.progress_base, args.progress_span, stop),
|
||||
daemon=True)
|
||||
sampler.start()
|
||||
try:
|
||||
if args.format == 3:
|
||||
from functools import partial
|
||||
build = partial(compile_model_v3, resolutions=tuple(args.camera_resolutions))
|
||||
else:
|
||||
build = compile_policy_model if args.format == 2 else compile_model
|
||||
try:
|
||||
out = build(meta, onnx_path, args.output or egpu_pkl_path(meta))
|
||||
except _ParityFail as e:
|
||||
if os.environ.get("TC_OPT") == "0" or args.format == 1:
|
||||
raise
|
||||
print(f"{e}\nretrying compile with tensor cores disabled", flush=True)
|
||||
os.environ["TC_OPT"] = "0"
|
||||
os.environ["IQ_EGPU_TC_OFF"] = "1"
|
||||
out = build(meta, onnx_path, args.output or egpu_pkl_path(meta))
|
||||
finally:
|
||||
if stop is not None:
|
||||
stop.set()
|
||||
if sampler is not None:
|
||||
sampler.join(timeout=2)
|
||||
print(f"saved eGPU jit to {out} ({os.path.getsize(out) / 1e6:.2f} MB)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
9
iqpilot/selfdrive/iqmodeld/tools/compile_emac_warp.py
Normal file
9
iqpilot/selfdrive/iqmodeld/tools/compile_emac_warp.py
Normal file
@@ -0,0 +1,9 @@
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
"""
|
||||
from iqpilot.selfdrive.iqmodeld.tools.compile_warp import MODEL_SIZE, compile_warp, main
|
||||
|
||||
__all__ = ["MODEL_SIZE", "compile_warp", "main"]
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
125
iqpilot/selfdrive/iqmodeld/tools/compile_model.py
Normal file
125
iqpilot/selfdrive/iqmodeld/tools/compile_model.py
Normal file
@@ -0,0 +1,125 @@
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
"""
|
||||
import os
|
||||
import pickle
|
||||
import re
|
||||
import shutil
|
||||
import sys
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
|
||||
if "JIT_BATCH_SIZE" not in os.environ:
|
||||
os.environ["JIT_BATCH_SIZE"] = "0"
|
||||
|
||||
from tinygrad import Context, Device, GlobalCounters, Tensor, TinyJit, dtypes
|
||||
from tinygrad.helpers import DEBUG, getenv
|
||||
from tinygrad.nn.onnx import OnnxRunner
|
||||
from tinygrad.uop.ops import Ops
|
||||
|
||||
|
||||
def compile_model(onnx_file, output):
|
||||
run_onnx = OnnxRunner(onnx_file)
|
||||
print("loaded model")
|
||||
|
||||
input_shapes = {name: spec.shape for name, spec in run_onnx.graph_inputs.items()}
|
||||
input_types = {name: spec.dtype for name, spec in run_onnx.graph_inputs.items()}
|
||||
input_types = {key: dtypes.float32 if value is dtypes.float16 else value for key, value in input_types.items()}
|
||||
input_shapes = {key: tuple(value if isinstance(value, int) else 1 for value in shape) for key, shape in input_shapes.items()}
|
||||
|
||||
Tensor.manual_seed(100)
|
||||
inputs = {
|
||||
key: Tensor(Tensor.randn(*shape, dtype=input_types[key]).mul(8).realize().numpy(), device="NPY")
|
||||
for key, shape in sorted(input_shapes.items())
|
||||
}
|
||||
if not getenv("NPY_IMG"):
|
||||
inputs = {key: Tensor(value.numpy(), device=Device.DEFAULT).realize() if "img" in key else value for key, value in inputs.items()}
|
||||
print("created tensors")
|
||||
|
||||
run_onnx_jit = TinyJit(
|
||||
lambda **kwargs: next(iter(run_onnx({key: value.to(Device.DEFAULT) for key, value in kwargs.items()}).values())).cast("float32"),
|
||||
prune=True,
|
||||
)
|
||||
test_value = None
|
||||
for iteration in range(3):
|
||||
GlobalCounters.reset()
|
||||
print(f"run {iteration}")
|
||||
with Context(DEBUG=max(DEBUG.value, 2 if iteration == 2 else 1), OPENPILOT_HACKS=1):
|
||||
result = run_onnx_jit(**inputs).numpy()
|
||||
if iteration == 1:
|
||||
test_value = np.copy(result)
|
||||
|
||||
kernel_asts = {Ops.PROGRAM}
|
||||
kernel_calls = [
|
||||
node for node in run_onnx_jit.captured.linear.toposort(gate=lambda value: value.op not in kernel_asts)
|
||||
if node.op is Ops.CALL and node.src[0].op in kernel_asts
|
||||
]
|
||||
print(f"captured {len(kernel_calls)} kernels")
|
||||
np.testing.assert_equal(test_value, result, "JIT run failed")
|
||||
print("jit run validated")
|
||||
|
||||
kernel_count = 0
|
||||
read_image_count = 0
|
||||
gated_read_image_count = 0
|
||||
for call in kernel_calls:
|
||||
_, _, source, _ = call.src[0].src
|
||||
rendered = source.arg
|
||||
kernel_count += 1
|
||||
read_image_count += rendered.count("read_image")
|
||||
gated_read_image_count += rendered.count("?read_image")
|
||||
for value in (match.group(1) for match in re.finditer(r"(val\d+)\s*=\s*read_imagef\(", rendered)):
|
||||
if re.search(fr"[?:]{value}\.[xyzw]", rendered):
|
||||
gated_read_image_count += 1
|
||||
|
||||
print(f"{kernel_count=}, {read_image_count=}, {gated_read_image_count=}")
|
||||
expected = {
|
||||
"kernel count": (kernel_count, getenv("ALLOWED_KERNEL_COUNT", -1)),
|
||||
"read image count": (read_image_count, getenv("ALLOWED_READ_IMAGE", -1)),
|
||||
"gated read image count": (gated_read_image_count, getenv("ALLOWED_GATED_READ_IMAGE", -1)),
|
||||
}
|
||||
for name, (actual, allowed) in expected.items():
|
||||
if allowed != -1:
|
||||
assert actual == allowed, f"different {name}: {actual}, expected {allowed}"
|
||||
|
||||
with open(output, "wb") as handle:
|
||||
pickle.dump(run_onnx_jit, handle)
|
||||
print(f"model size is {os.path.getsize(onnx_file) / 1e6:.2f}M")
|
||||
print(f"pkl size is {os.path.getsize(output) / 1e6:.2f}M")
|
||||
return run_onnx_jit, inputs, test_value
|
||||
|
||||
|
||||
def test_compiled(run, inputs, test_value):
|
||||
step_times = []
|
||||
for _ in range(20):
|
||||
start = time.perf_counter()
|
||||
output = run(**inputs)
|
||||
queued = time.perf_counter()
|
||||
value = output.numpy()
|
||||
end = time.perf_counter()
|
||||
step_times.append((end - start) * 1e3)
|
||||
print(f"enqueue {(queued - start) * 1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")
|
||||
|
||||
minimum = getenv("ASSERT_MIN_STEP_TIME", 0.0)
|
||||
if minimum:
|
||||
assert min(step_times) < minimum, f"expected minimum step time below {minimum} ms, got {min(step_times)} ms"
|
||||
np.testing.assert_equal(test_value, value)
|
||||
changed_inputs = {key: Tensor(item.numpy() * 2, device=item.device) for key, item in inputs.items()}
|
||||
changed_value = run(**changed_inputs).numpy()
|
||||
np.testing.assert_raises(AssertionError, np.testing.assert_array_equal, value, changed_value)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
model_path = sys.argv[1]
|
||||
output_path = sys.argv[2]
|
||||
if stash := os.environ.get("IQPILOT_MODEL_STASH"):
|
||||
stashed_model = os.path.join(stash, os.path.basename(output_path))
|
||||
if os.path.isfile(stashed_model) and os.path.getsize(stashed_model) > 0:
|
||||
os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
|
||||
shutil.copyfile(stashed_model, output_path)
|
||||
print(f"restored device-compiled model: {output_path}")
|
||||
sys.exit(0)
|
||||
_, input_values, expected_value = compile_model(model_path, output_path)
|
||||
with open(output_path, "rb") as compiled_file:
|
||||
compiled_model = pickle.load(compiled_file)
|
||||
test_compiled(compiled_model, input_values, expected_value)
|
||||
417
iqpilot/selfdrive/iqmodeld/tools/compile_split_runtime.py
Normal file
417
iqpilot/selfdrive/iqmodeld/tools/compile_split_runtime.py
Normal file
@@ -0,0 +1,417 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import atexit
|
||||
import os
|
||||
import pickle
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def _patch_firmware_fetch() -> None:
|
||||
import hashlib
|
||||
import pathlib
|
||||
|
||||
import zstandard
|
||||
from tinygrad import helpers
|
||||
|
||||
if not hasattr(helpers, "fetch_fw"):
|
||||
return
|
||||
|
||||
original_fetch = helpers.fetch_fw
|
||||
|
||||
def fetch_fw(path, name, sha256):
|
||||
archive_path = pathlib.Path(f"/lib/firmware/{path}/{name}.zst")
|
||||
if archive_path.is_file():
|
||||
blob = zstandard.ZstdDecompressor().stream_reader(archive_path.read_bytes()).read()
|
||||
if hashlib.sha256(blob).hexdigest() == sha256:
|
||||
return blob
|
||||
return original_fetch(path, name, sha256)
|
||||
|
||||
helpers.fetch_fw = fetch_fw
|
||||
|
||||
|
||||
_patch_firmware_fetch()
|
||||
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.engine.jit import TinyJit
|
||||
from tinygrad.helpers import Context
|
||||
from tinygrad.nn.onnx import OnnxRunner
|
||||
from tinygrad.tensor import Tensor
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CameraGeometry:
|
||||
width: int
|
||||
height: int
|
||||
stride: int
|
||||
y_height: int
|
||||
uv_height: int
|
||||
size: int
|
||||
|
||||
|
||||
WARP_DEVICE = os.getenv("WARP_DEV")
|
||||
|
||||
|
||||
def _read_shared_copy(path: str) -> str:
|
||||
from iqpilot.common.file_chunker import read_file_chunked
|
||||
from iqpilot.system.hardware.hw import Paths
|
||||
|
||||
shm_path = os.path.join(Paths.shm_path(), os.path.basename(path))
|
||||
atexit.register(lambda: os.path.exists(shm_path) and os.remove(shm_path))
|
||||
with open(shm_path, "wb") as handle:
|
||||
handle.write(read_file_chunked(path))
|
||||
return shm_path
|
||||
|
||||
|
||||
def _parse_size(text: str) -> tuple[int, int]:
|
||||
width, height = text.lower().split("x")
|
||||
return int(width), int(height)
|
||||
|
||||
|
||||
def _rand_u8_inputs(keys: list[str], shape, device=None):
|
||||
return {key: Tensor.randint(shape, low=0, high=256, dtype="uint8", device=device).realize() for key in keys}
|
||||
|
||||
|
||||
def _phase_desire_key(policy_shapes: dict[str, tuple[int, ...]]) -> str:
|
||||
for key in policy_shapes:
|
||||
if key.startswith("desire"):
|
||||
return key
|
||||
raise KeyError("No desire-like key found in policy shapes")
|
||||
|
||||
|
||||
def _phase_image_keys(vision_shapes: dict[str, tuple[int, ...]]) -> tuple[str, str]:
|
||||
names = sorted(name for name in vision_shapes if "img" in name)
|
||||
road_key = next((name for name in names if "big" not in name), None)
|
||||
wide_key = next((name for name in names if "big" in name), None)
|
||||
if road_key is None or wide_key is None:
|
||||
raise ValueError(f"Unable to resolve road/wide image keys from {list(vision_shapes)}")
|
||||
return road_key, wide_key
|
||||
|
||||
|
||||
def _base_policy_keys(policy_shapes: dict[str, tuple[int, ...]]) -> set[str]:
|
||||
return {
|
||||
_phase_desire_key(policy_shapes),
|
||||
"features_buffer",
|
||||
"traffic_convention",
|
||||
"action_t",
|
||||
}
|
||||
|
||||
|
||||
def _common_policy_shapes(role_shapes: dict[str, dict[str, tuple[int, ...]]]) -> dict[str, tuple[int, ...]]:
|
||||
first_role = next(iter(role_shapes))
|
||||
baseline = role_shapes[first_role]
|
||||
for role_name, shape_map in role_shapes.items():
|
||||
if shape_map != baseline:
|
||||
raise ValueError(f"Policy input shapes differ for role {role_name}")
|
||||
return baseline
|
||||
|
||||
|
||||
def _phase_frame_skip(policy_shapes: dict[str, tuple[int, ...]]) -> int:
|
||||
feature_shape = policy_shapes.get("features_buffer")
|
||||
if feature_shape is None:
|
||||
return 1
|
||||
history_length = feature_shape[1]
|
||||
return 1 if history_length >= 99 else 4
|
||||
|
||||
|
||||
def _project_pixels(src_flat, inverse_matrix, dst_shape, src_shape, stride_pad, border_fill_val=None):
|
||||
dst_w, dst_h = dst_shape
|
||||
src_h, src_w = src_shape
|
||||
|
||||
x_coords = Tensor.arange(dst_w).to(WARP_DEVICE).reshape(1, dst_w).expand(dst_h, dst_w).reshape(-1)
|
||||
y_coords = Tensor.arange(dst_h).to(WARP_DEVICE).reshape(dst_h, 1).expand(dst_h, dst_w).reshape(-1)
|
||||
|
||||
src_x = inverse_matrix[0, 0] * x_coords + inverse_matrix[0, 1] * y_coords + inverse_matrix[0, 2]
|
||||
src_y = inverse_matrix[1, 0] * x_coords + inverse_matrix[1, 1] * y_coords + inverse_matrix[1, 2]
|
||||
scale = inverse_matrix[2, 0] * x_coords + inverse_matrix[2, 1] * y_coords + inverse_matrix[2, 2]
|
||||
|
||||
src_x = src_x / scale
|
||||
src_y = src_y / scale
|
||||
|
||||
rounded_x = Tensor.round(src_x)
|
||||
rounded_y = Tensor.round(src_y)
|
||||
gather_x = rounded_x.clip(0, src_w - 1).cast("int")
|
||||
gather_y = rounded_y.clip(0, src_h - 1).cast("int")
|
||||
gather_index = gather_y * (src_w + stride_pad) + gather_x
|
||||
sampled = src_flat[gather_index]
|
||||
|
||||
if border_fill_val is None:
|
||||
return sampled
|
||||
|
||||
inside = ((rounded_x >= 0) & (rounded_x <= src_w - 1) & (rounded_y >= 0) & (rounded_y <= src_h - 1)).cast(sampled.dtype)
|
||||
return sampled * inside + Tensor(border_fill_val, dtype=sampled.dtype) * (1 - inside)
|
||||
|
||||
|
||||
def _pack_nv12_planes(stacked_frame):
|
||||
y_height = (stacked_frame.shape[0] * 2) // 3
|
||||
frame_width = stacked_frame.shape[1]
|
||||
return Tensor.cat(
|
||||
stacked_frame[0:y_height:2, 0::2],
|
||||
stacked_frame[1:y_height:2, 0::2],
|
||||
stacked_frame[0:y_height:2, 1::2],
|
||||
stacked_frame[1:y_height:2, 1::2],
|
||||
stacked_frame[y_height:y_height + y_height // 4].reshape((y_height // 2, frame_width // 2)),
|
||||
stacked_frame[y_height + y_height // 4:y_height + y_height // 2].reshape((y_height // 2, frame_width // 2)),
|
||||
dim=0,
|
||||
).reshape((6, y_height // 2, frame_width // 2))
|
||||
|
||||
|
||||
def _warp_program(camera: CameraGeometry, model_w: int, model_h: int):
|
||||
uv_offset = camera.stride * camera.y_height
|
||||
stride_pad = camera.stride - camera.width
|
||||
|
||||
def prepare_frame(nv12_blob, inverse_matrix):
|
||||
uv_matrix = inverse_matrix * Tensor([[1.0, 1.0, 0.5], [1.0, 1.0, 0.5], [2.0, 2.0, 1.0]], device=WARP_DEVICE)
|
||||
uv_plane = nv12_blob[uv_offset:uv_offset + camera.uv_height * camera.stride].reshape(camera.uv_height, camera.stride)
|
||||
with Context(SPLIT_REDUCEOP=0):
|
||||
y_plane = _project_pixels(nv12_blob[:camera.height * camera.stride], inverse_matrix, (model_w, model_h), (camera.height, camera.width), stride_pad).realize()
|
||||
u_plane = _project_pixels(uv_plane[:camera.height // 2, :camera.width:2].flatten(), uv_matrix, (model_w // 2, model_h // 2), (camera.height // 2, camera.width // 2), 0).realize()
|
||||
v_plane = _project_pixels(uv_plane[:camera.height // 2, 1:camera.width:2].flatten(), uv_matrix, (model_w // 2, model_h // 2), (camera.height // 2, camera.width // 2), 0).realize()
|
||||
return _pack_nv12_planes(y_plane.cat(u_plane).cat(v_plane).reshape((model_h * 3 // 2, model_w)))
|
||||
|
||||
return prepare_frame
|
||||
|
||||
|
||||
def _sample_sparse(queue_tensor, frame_stride):
|
||||
return queue_tensor[::frame_stride].contiguous().flatten(0, 1).unsqueeze(0)
|
||||
|
||||
|
||||
def _sample_desire(queue_tensor, frame_stride):
|
||||
return queue_tensor.reshape(-1, frame_stride, *queue_tensor.shape[1:]).max(1).flatten(0, 1).unsqueeze(0)
|
||||
|
||||
|
||||
def _roll_queue(queue_tensor, incoming, sampler):
|
||||
queue_tensor.assign(queue_tensor[1:].cat(incoming, dim=0).contiguous())
|
||||
return sampler(queue_tensor)
|
||||
|
||||
|
||||
def _vision_queue_buffers(vision_shapes: dict[str, tuple[int, ...]], frame_stride: int, device):
|
||||
road_key, _ = _phase_image_keys(vision_shapes)
|
||||
image_shape = vision_shapes[road_key]
|
||||
frame_history = image_shape[1] // 6
|
||||
queue_depth = frame_stride * (frame_history - 1) + 1
|
||||
frame_queue_shape = (queue_depth, 6, image_shape[2], image_shape[3])
|
||||
|
||||
numpy_state = {
|
||||
"tfm": np.zeros((3, 3), dtype=np.float32),
|
||||
"big_tfm": np.zeros((3, 3), dtype=np.float32),
|
||||
}
|
||||
tensor_state = {
|
||||
"img_q": Tensor(np.zeros(frame_queue_shape, dtype=np.uint8), device=device).contiguous().realize(),
|
||||
"big_img_q": Tensor(np.zeros(frame_queue_shape, dtype=np.uint8), device=device).contiguous().realize(),
|
||||
**{name: Tensor(value, device="NPY").realize() for name, value in numpy_state.items()},
|
||||
}
|
||||
return tensor_state, numpy_state
|
||||
|
||||
|
||||
def _policy_queue_buffers(vision_shapes: dict[str, tuple[int, ...]], policy_shapes: dict[str, tuple[int, ...]], frame_stride: int, device):
|
||||
tensor_state, numpy_state = _vision_queue_buffers(vision_shapes, frame_stride, device)
|
||||
desired_key = _phase_desire_key(policy_shapes)
|
||||
feature_shape = policy_shapes["features_buffer"]
|
||||
desired_shape = policy_shapes[desired_key]
|
||||
traffic_shape = policy_shapes["traffic_convention"]
|
||||
action_shape = policy_shapes.get("action_t", traffic_shape)
|
||||
|
||||
numpy_policy = {
|
||||
"desire": np.zeros(desired_shape[2], dtype=np.float32),
|
||||
"traffic_convention": np.zeros(traffic_shape, dtype=np.float32),
|
||||
"action_t": np.zeros(action_shape, dtype=np.float32),
|
||||
}
|
||||
for key, shape in policy_shapes.items():
|
||||
if key not in _base_policy_keys(policy_shapes):
|
||||
numpy_policy[key] = np.zeros(shape, dtype=np.float32)
|
||||
|
||||
numpy_state.update(numpy_policy)
|
||||
tensor_state.update({
|
||||
"feat_q": Tensor(np.zeros((frame_stride * (feature_shape[1] - 1) + 1, feature_shape[0], feature_shape[2]), dtype=np.float32), device=device).contiguous().realize(),
|
||||
"desire_q": Tensor(np.zeros((frame_stride * desired_shape[1], desired_shape[0], desired_shape[2]), dtype=np.float32), device=device).contiguous().realize(),
|
||||
**{name: Tensor(value, device="NPY").realize() for name, value in numpy_policy.items()},
|
||||
})
|
||||
return tensor_state, numpy_state
|
||||
|
||||
|
||||
def _stage_program(camera: CameraGeometry, model_w: int, model_h: int, frame_stride: int):
|
||||
prepare_frame = _warp_program(camera, model_w, model_h)
|
||||
sparse_sampler = partial(_sample_sparse, frame_stride=frame_stride)
|
||||
|
||||
def stage_inputs(img_q, big_img_q, tfm, big_tfm, frame, big_frame):
|
||||
tfm = tfm.to(WARP_DEVICE)
|
||||
big_tfm = big_tfm.to(WARP_DEVICE)
|
||||
Tensor.realize(tfm, big_tfm)
|
||||
staged_main = prepare_frame(frame, tfm).unsqueeze(0).to(Device.DEFAULT)
|
||||
staged_wide = prepare_frame(big_frame, big_tfm).unsqueeze(0).to(Device.DEFAULT)
|
||||
return (
|
||||
_roll_queue(img_q, staged_main, sparse_sampler),
|
||||
_roll_queue(big_img_q, staged_wide, sparse_sampler),
|
||||
)
|
||||
|
||||
return stage_inputs
|
||||
|
||||
|
||||
def _role_executor(model_runners: dict[str, OnnxRunner], meta_by_role: dict[str, dict], frame_stride: int):
|
||||
desired_sampler = partial(_sample_desire, frame_stride=frame_stride)
|
||||
sparse_sampler = partial(_sample_sparse, frame_stride=frame_stride)
|
||||
vision_hidden_slice = meta_by_role["vision"]["output_slices"]["hidden_state"]
|
||||
policy_roles = [name for name in meta_by_role if name != "vision"]
|
||||
policy_shapes = _common_policy_shapes({name: meta_by_role[name]["input_shapes"] for name in policy_roles})
|
||||
desired_key = _phase_desire_key(policy_shapes)
|
||||
road_key, wide_key = _phase_image_keys(meta_by_role["vision"]["input_shapes"])
|
||||
extra_keys = [key for key in policy_shapes if key not in _base_policy_keys(policy_shapes)]
|
||||
|
||||
def execute_bundle(img, big_img, feat_q, desire_q, desire, traffic_convention, action_t, **extra):
|
||||
desired_tensor = desire.to(Device.DEFAULT)
|
||||
traffic_tensor = traffic_convention.to(Device.DEFAULT)
|
||||
action_tensor = action_t.to(Device.DEFAULT)
|
||||
extra_tensors = {key: extra[key].to(Device.DEFAULT) for key in extra_keys if key in extra}
|
||||
Tensor.realize(desired_tensor, traffic_tensor, action_tensor, *extra_tensors.values())
|
||||
|
||||
desire_buffer = _roll_queue(desire_q, desired_tensor.reshape(1, 1, -1), desired_sampler)
|
||||
vision_output = next(iter(model_runners["vision"]({road_key: img, wide_key: big_img}).values())).cast("float32")
|
||||
hidden_state = vision_output[:, vision_hidden_slice].reshape(1, -1).unsqueeze(0)
|
||||
feature_buffer = _roll_queue(feat_q, hidden_state, sparse_sampler)
|
||||
|
||||
common_inputs = {
|
||||
"features_buffer": feature_buffer,
|
||||
desired_key: desire_buffer,
|
||||
"traffic_convention": traffic_tensor,
|
||||
"action_t": action_tensor,
|
||||
**extra_tensors,
|
||||
}
|
||||
|
||||
role_outputs = []
|
||||
for role_name in policy_roles:
|
||||
role_outputs.append(next(iter(model_runners[role_name](common_inputs).values())).cast("float32"))
|
||||
return (vision_output, *role_outputs)
|
||||
|
||||
return execute_bundle
|
||||
|
||||
|
||||
def _capture_and_freeze(jit_runner, random_inputs_factory, queue_keys, queue_factory):
|
||||
seed_value = 42
|
||||
|
||||
def validate(fn, baseline_outputs=None, baseline_buffers=None, expect_match=True, replay_seed=seed_value):
|
||||
queue_tensors, numpy_values = queue_factory(Device.DEFAULT)
|
||||
np.random.seed(replay_seed)
|
||||
Tensor.manual_seed(replay_seed)
|
||||
|
||||
replay_count = 1 if (baseline_outputs is not None or baseline_buffers is not None) else 3
|
||||
for pass_index in range(replay_count):
|
||||
for value in numpy_values.values():
|
||||
value[:] = np.random.randn(*value.shape).astype(value.dtype)
|
||||
Device.default.synchronize()
|
||||
random_inputs = random_inputs_factory()
|
||||
start_time = time.perf_counter()
|
||||
outputs = fn(**{name: queue_tensors[name] for name in queue_keys}, **random_inputs)
|
||||
enqueue_time = time.perf_counter()
|
||||
Device.default.synchronize()
|
||||
total_time = time.perf_counter()
|
||||
print(f" [{pass_index + 1}/{replay_count}] enqueue {(enqueue_time - start_time) * 1e3:6.2f} ms -- total {(total_time - start_time) * 1e3:6.2f} ms")
|
||||
|
||||
if pass_index == 0:
|
||||
output_snapshot = [np.copy(value.numpy()) for value in outputs]
|
||||
buffer_snapshot = [np.copy(value.numpy().copy()) for value in queue_tensors.values()]
|
||||
|
||||
if baseline_outputs is not None:
|
||||
matches = all(np.array_equal(current, reference) for current, reference in zip(output_snapshot, baseline_outputs, strict=True))
|
||||
assert matches == expect_match, f"outputs {'differ from' if expect_match else 'match'} baseline"
|
||||
if baseline_buffers is not None:
|
||||
matches = all(np.array_equal(current, reference) for current, reference in zip(buffer_snapshot, baseline_buffers, strict=True))
|
||||
assert matches == expect_match, f"buffers {'differ from' if expect_match else 'match'} baseline"
|
||||
|
||||
return output_snapshot, buffer_snapshot
|
||||
|
||||
print("capture + replay")
|
||||
baseline_outputs, baseline_buffers = validate(jit_runner)
|
||||
print("pickle round trip")
|
||||
frozen = pickle.loads(pickle.dumps(jit_runner))
|
||||
validate(frozen, baseline_outputs, baseline_buffers, expect_match=True)
|
||||
validate(frozen, baseline_outputs, baseline_buffers, expect_match=False, replay_seed=seed_value + 1)
|
||||
return frozen
|
||||
|
||||
|
||||
def _arg_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--model-size", type=_parse_size, required=True, help="model input WxH")
|
||||
parser.add_argument("--camera-resolutions", type=_parse_size, nargs="+", required=True, help="camera resolutions WxH")
|
||||
parser.add_argument("--vision-onnx", required=True)
|
||||
parser.add_argument("--policy-onnx")
|
||||
parser.add_argument("--off-policy-onnx")
|
||||
parser.add_argument("--on-policy-onnx")
|
||||
parser.add_argument("--output", required=True)
|
||||
parser.add_argument("--frame-skip", type=int)
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
from iqpilot.selfdrive.iqmodeld.metadata import build_metadata_record
|
||||
from iqpilot.system.camerad.cameras.nv12_info import get_nv12_info
|
||||
|
||||
args = _arg_parser().parse_args(argv)
|
||||
model_w, model_h = args.model_size
|
||||
|
||||
policy_specs = [
|
||||
("policy", args.policy_onnx),
|
||||
("off_policy", args.off_policy_onnx),
|
||||
("on_policy", args.on_policy_onnx),
|
||||
]
|
||||
active_policy_specs = [(role, path) for role, path in policy_specs if path]
|
||||
if not active_policy_specs:
|
||||
raise SystemExit("At least one policy ONNX must be provided")
|
||||
|
||||
model_paths = {"vision": _read_shared_copy(args.vision_onnx)}
|
||||
for role_name, onnx_path in active_policy_specs:
|
||||
model_paths[role_name] = _read_shared_copy(onnx_path)
|
||||
|
||||
model_runners = {role_name: OnnxRunner(path) for role_name, path in model_paths.items()}
|
||||
meta_by_role = {role_name: build_metadata_record(path) for role_name, path in model_paths.items()}
|
||||
|
||||
shared_policy_shapes = _common_policy_shapes({
|
||||
role_name: meta_by_role[role_name]["input_shapes"] for role_name, _ in active_policy_specs
|
||||
})
|
||||
frame_stride = args.frame_skip if args.frame_skip is not None else _phase_frame_skip(shared_policy_shapes)
|
||||
|
||||
package: dict[Any, Any] = {
|
||||
"meta_by_role": meta_by_role,
|
||||
"roles": [role_name for role_name, _ in active_policy_specs],
|
||||
"frame_stride": frame_stride,
|
||||
}
|
||||
|
||||
executor_jit = TinyJit(_role_executor(model_runners, meta_by_role, frame_stride), prune=True)
|
||||
queue_factory = partial(_policy_queue_buffers, meta_by_role["vision"]["input_shapes"], shared_policy_shapes, frame_stride)
|
||||
image_shape = meta_by_role["vision"]["input_shapes"][_phase_image_keys(meta_by_role["vision"]["input_shapes"])[0]]
|
||||
package["execute_bundle"] = _capture_and_freeze(
|
||||
executor_jit,
|
||||
partial(_rand_u8_inputs, keys=["img", "big_img"], shape=image_shape),
|
||||
["feat_q", "desire_q", "desire", "traffic_convention", "action_t", *[k for k in shared_policy_shapes if k not in _base_policy_keys(shared_policy_shapes)]],
|
||||
queue_factory,
|
||||
)
|
||||
|
||||
for camera_width, camera_height in args.camera_resolutions:
|
||||
camera = CameraGeometry(camera_width, camera_height, *get_nv12_info(camera_width, camera_height))
|
||||
stage_jit = TinyJit(_stage_program(camera, model_w, model_h, frame_stride), prune=True)
|
||||
package[(camera_width, camera_height)] = {
|
||||
"stage_inputs": _capture_and_freeze(
|
||||
stage_jit,
|
||||
partial(_rand_u8_inputs, keys=["frame", "big_frame"], shape=camera.size, device=WARP_DEVICE),
|
||||
["img_q", "big_img_q", "tfm", "big_tfm"],
|
||||
partial(_vision_queue_buffers, meta_by_role["vision"]["input_shapes"], frame_stride),
|
||||
)
|
||||
}
|
||||
|
||||
with open(args.output, "wb") as handle:
|
||||
pickle.dump(package, handle)
|
||||
print(f"Saved combined split runtime to {args.output} ({os.path.getsize(args.output) / 1e6:.2f} MB)")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
422
iqpilot/selfdrive/iqmodeld/tools/compile_supercombo.py
Normal file
422
iqpilot/selfdrive/iqmodeld/tools/compile_supercombo.py
Normal file
@@ -0,0 +1,422 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import atexit
|
||||
import math
|
||||
import os
|
||||
import pickle
|
||||
import re
|
||||
import tempfile
|
||||
import time
|
||||
from functools import partial
|
||||
from collections import namedtuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
def _patch_tinygrad_fetch_fw():
|
||||
import hashlib
|
||||
import pathlib
|
||||
import zstandard
|
||||
from tinygrad import helpers
|
||||
_orig = helpers.fetch_fw
|
||||
def fetch_fw(path, name, sha256):
|
||||
p = pathlib.Path(f"/lib/firmware/{path}/{name}.zst")
|
||||
if p.is_file():
|
||||
blob = zstandard.ZstdDecompressor().stream_reader(p.read_bytes()).read()
|
||||
if hashlib.sha256(blob).hexdigest() == sha256:
|
||||
return blob
|
||||
return _orig(path, name, sha256)
|
||||
helpers.fetch_fw = fetch_fw
|
||||
_patch_tinygrad_fetch_fw()
|
||||
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.helpers import Context
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.engine.jit import TinyJit
|
||||
|
||||
|
||||
NV12Frame = namedtuple("NV12Frame", ['width', 'height', 'stride', 'y_height', 'uv_height', 'size'])
|
||||
WARP_INPUTS = ['tfm', 'big_tfm']
|
||||
POLICY_INPUTS = ['img_q', 'big_img_q', 'feat_q', 'desire_q', 'packed_npy_inputs']
|
||||
|
||||
UV_SCALE_MATRIX = np.array([[0.5, 0, 0], [0, 0.5, 0], [0, 0, 1]], dtype=np.float32)
|
||||
UV_SCALE_MATRIX_INV = np.linalg.inv(UV_SCALE_MATRIX)
|
||||
|
||||
WARP_DEV = os.getenv('WARP_DEV')
|
||||
|
||||
|
||||
def make_random_images(keys, shape, device=None):
|
||||
return {k: Tensor.randint(shape, low=0, high=256, dtype='uint8', device=device).realize() for k in keys}
|
||||
|
||||
|
||||
def warp_perspective_tinygrad(src_flat, M_inv, dst_shape, src_shape, stride_pad, border_fill_val=None):
|
||||
w_dst, h_dst = dst_shape
|
||||
h_src, w_src = src_shape
|
||||
|
||||
x = Tensor.arange(w_dst).reshape(1, w_dst).expand(h_dst, w_dst).reshape(-1)
|
||||
y = Tensor.arange(h_dst).reshape(h_dst, 1).expand(h_dst, w_dst).reshape(-1)
|
||||
|
||||
src_x = M_inv[0, 0] * x + M_inv[0, 1] * y + M_inv[0, 2]
|
||||
src_y = M_inv[1, 0] * x + M_inv[1, 1] * y + M_inv[1, 2]
|
||||
src_w = M_inv[2, 0] * x + M_inv[2, 1] * y + M_inv[2, 2]
|
||||
|
||||
src_x = src_x / src_w
|
||||
src_y = src_y / src_w
|
||||
|
||||
x_round = Tensor.round(src_x)
|
||||
y_round = Tensor.round(src_y)
|
||||
x_nn_clipped = x_round.clip(0, w_src - 1).cast('int')
|
||||
y_nn_clipped = y_round.clip(0, h_src - 1).cast('int')
|
||||
idx = y_nn_clipped * (w_src + stride_pad) + x_nn_clipped
|
||||
sampled = src_flat[idx]
|
||||
|
||||
if border_fill_val is None:
|
||||
return sampled
|
||||
|
||||
in_bounds = ((x_round >= 0) & (x_round <= w_src - 1) &
|
||||
(y_round >= 0) & (y_round <= h_src - 1)).cast(sampled.dtype)
|
||||
return sampled * in_bounds + Tensor(border_fill_val, dtype=sampled.dtype) * (1 - in_bounds)
|
||||
|
||||
|
||||
def frames_to_tensor(frames):
|
||||
H = (frames.shape[0] * 2) // 3
|
||||
W = frames.shape[1]
|
||||
in_img1 = Tensor.cat(frames[0:H:2, 0::2],
|
||||
frames[1:H:2, 0::2],
|
||||
frames[0:H:2, 1::2],
|
||||
frames[1:H:2, 1::2],
|
||||
frames[H:H+H//4].reshape((H//2, W//2)),
|
||||
frames[H+H//4:H+H//2].reshape((H//2, W//2)), dim=0).reshape((6, H//2, W//2))
|
||||
return in_img1
|
||||
|
||||
|
||||
def make_frame_prepare(nv12: NV12Frame, model_w, model_h):
|
||||
cam_w, cam_h, stride, y_height, uv_height, _ = nv12
|
||||
uv_offset = stride * y_height
|
||||
stride_pad = stride - cam_w
|
||||
|
||||
def frame_prepare_tinygrad(input_frame, M_inv):
|
||||
M_inv_uv = M_inv * Tensor([[1.0, 1.0, 0.5], [1.0, 1.0, 0.5], [2.0, 2.0, 1.0]], device=WARP_DEV)
|
||||
uv = input_frame[uv_offset:uv_offset + uv_height * stride].reshape(uv_height, stride)
|
||||
with Context(SPLIT_REDUCEOP=0):
|
||||
y = warp_perspective_tinygrad(input_frame[:cam_h*stride],
|
||||
M_inv, (model_w, model_h),
|
||||
(cam_h, cam_w), stride_pad).realize()
|
||||
u = warp_perspective_tinygrad(uv[:cam_h//2, :cam_w:2].flatten(),
|
||||
M_inv_uv, (model_w//2, model_h//2),
|
||||
(cam_h//2, cam_w//2), 0).realize()
|
||||
v = warp_perspective_tinygrad(uv[:cam_h//2, 1:cam_w:2].flatten(),
|
||||
M_inv_uv, (model_w//2, model_h//2),
|
||||
(cam_h//2, cam_w//2), 0).realize()
|
||||
yuv = y.cat(u).cat(v).reshape((model_h * 3 // 2, model_w))
|
||||
tensor = frames_to_tensor(yuv)
|
||||
return tensor
|
||||
return frame_prepare_tinygrad
|
||||
|
||||
|
||||
def make_warp_input_queues(vision_input_shapes, frame_skip, device):
|
||||
img = vision_input_shapes['img'] # (1, 12, 128, 256)
|
||||
n_frames = img[1] // 6
|
||||
img_buf_shape = (frame_skip * (n_frames - 1) + 1, 6, img[2], img[3])
|
||||
|
||||
npy = {
|
||||
'tfm': np.zeros((3, 3), dtype=np.float32),
|
||||
'big_tfm': np.zeros((3, 3), dtype=np.float32),
|
||||
}
|
||||
input_queues = {
|
||||
'img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(),
|
||||
'big_img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(),
|
||||
**{k: Tensor(v, device='NPY').realize() for k, v in npy.items()},
|
||||
}
|
||||
return input_queues, npy
|
||||
|
||||
|
||||
def get_policy_npy_shapes(input_shapes):
|
||||
dp = input_shapes['desire_pulse'] # (1, 25, 8)
|
||||
tc = input_shapes['traffic_convention'] # (1, 2)
|
||||
at = input_shapes['action_t'] # (1, 2)
|
||||
fb = input_shapes['features_buffer'] # (1, 24, 512)
|
||||
shapes = {'desire': (dp[2],), 'traffic_convention': tuple(tc), 'action_t': tuple(at), 'prev_feat': (fb[0], fb[2])}
|
||||
return shapes, [math.prod(s) for s in shapes.values()]
|
||||
|
||||
|
||||
def make_input_queues(input_shapes, frame_skip, device):
|
||||
input_queues, npy = make_warp_input_queues(input_shapes, frame_skip, device)
|
||||
|
||||
fb = input_shapes['features_buffer'] # (1, 24, 512), past features only; the model appends the current frame's feature
|
||||
dp = input_shapes['desire_pulse'] # (1, 25, 8)
|
||||
|
||||
shapes, sizes = get_policy_npy_shapes(input_shapes)
|
||||
packed_npy_inputs = np.zeros(sum(sizes), dtype=np.float32)
|
||||
npy.update({k: v.reshape(s) for (k, s), v in zip(shapes.items(), np.split(packed_npy_inputs, np.cumsum(sizes[:-1])), strict=True)})
|
||||
input_queues.update({
|
||||
'feat_q': Tensor(np.zeros((frame_skip * fb[1], fb[0], fb[2]), dtype=np.float32), device=device).contiguous().realize(),
|
||||
'desire_q': Tensor(np.zeros((frame_skip * dp[1], dp[0], dp[2]), dtype=np.float32), device=device).contiguous().realize(),
|
||||
'packed_npy_inputs': Tensor(packed_npy_inputs, device='NPY').realize(),
|
||||
})
|
||||
return input_queues, npy
|
||||
|
||||
|
||||
def shift_and_sample(buf, new_val, sample_fn):
|
||||
buf.assign(buf[1:].cat(new_val, dim=0).contiguous())
|
||||
return sample_fn(buf)
|
||||
|
||||
|
||||
def sample_skip(buf, frame_skip):
|
||||
return buf[::frame_skip].contiguous().flatten(0, 1).unsqueeze(0)
|
||||
|
||||
|
||||
def sample_desire(buf, frame_skip):
|
||||
return buf.reshape(-1, frame_skip, *buf.shape[1:]).max(1).flatten(0, 1).unsqueeze(0)
|
||||
|
||||
|
||||
def make_warp(nv12, model_w, model_h, frame_skip):
|
||||
frame_prepare = make_frame_prepare(nv12, model_w, model_h)
|
||||
|
||||
def warp(tfm, big_tfm, frame, big_frame):
|
||||
tfm = tfm.to(WARP_DEV)
|
||||
big_tfm = big_tfm.to(WARP_DEV)
|
||||
Tensor.realize(tfm, big_tfm)
|
||||
|
||||
warped_frame = frame_prepare(frame, tfm).unsqueeze(0)
|
||||
warped_big_frame = frame_prepare(big_frame, big_tfm).unsqueeze(0)
|
||||
return Tensor.cat(warped_frame, warped_big_frame)
|
||||
|
||||
return warp
|
||||
|
||||
|
||||
def make_run_policy(model_runner, model_metadata, frame_skip):
|
||||
sample_desire_fn = partial(sample_desire, frame_skip=frame_skip)
|
||||
sample_skip_fn = partial(sample_skip, frame_skip=frame_skip)
|
||||
npy_shapes, npy_sizes = get_policy_npy_shapes(model_metadata['input_shapes'])
|
||||
|
||||
def run_policy(warped, img_q, big_img_q, feat_q, desire_q, packed_npy_inputs):
|
||||
packed_npy_inputs = packed_npy_inputs.to(Device.DEFAULT)
|
||||
warped = warped.to(Device.DEFAULT)
|
||||
Tensor.realize(packed_npy_inputs, warped)
|
||||
|
||||
img = shift_and_sample(img_q, warped[0:1], sample_skip_fn)
|
||||
big_img = shift_and_sample(big_img_q, warped[1:2], sample_skip_fn)
|
||||
|
||||
desire, traffic_convention, action_t, prev_feat = (t.reshape(s) for t, s in zip(packed_npy_inputs.split(npy_sizes), npy_shapes.values(), strict=True))
|
||||
desire_buf = shift_and_sample(desire_q, desire.reshape(1, 1, -1), sample_desire_fn)
|
||||
feat_buf = shift_and_sample(feat_q, prev_feat.reshape(1, 1, -1), sample_skip_fn)
|
||||
|
||||
inputs = {
|
||||
'img': img,
|
||||
'big_img': big_img,
|
||||
'features_buffer': feat_buf,
|
||||
'desire_pulse': desire_buf,
|
||||
'traffic_convention': traffic_convention,
|
||||
'action_t': action_t,
|
||||
}
|
||||
out = next(iter(model_runner(inputs).values())).cast('float32')
|
||||
return out,
|
||||
return run_policy
|
||||
|
||||
|
||||
def compile_jit(jit, make_random_inputs, input_keys, make_queues):
|
||||
SEED = 42
|
||||
def random_inputs_run(fn, seed, test_val=None, test_buffers=None, expect_match=True):
|
||||
input_queues, npy = make_queues(Device.DEFAULT)
|
||||
np.random.seed(seed)
|
||||
Tensor.manual_seed(seed)
|
||||
|
||||
testing = test_val is not None or test_buffers is not None
|
||||
n_runs = 1 if testing else 3
|
||||
|
||||
for i in range(n_runs):
|
||||
for v in npy.values():
|
||||
v[:] = np.random.randn(*v.shape).astype(v.dtype)
|
||||
Device.default.synchronize()
|
||||
random_inputs = make_random_inputs()
|
||||
st = time.perf_counter()
|
||||
outs = fn(**{k: input_queues[k] for k in input_keys}, **random_inputs)
|
||||
mt = time.perf_counter()
|
||||
Device.default.synchronize()
|
||||
et = time.perf_counter()
|
||||
print(f" [{i+1}/{n_runs}] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms")
|
||||
|
||||
if i == 0:
|
||||
val = [np.copy(v.numpy()) for v in outs]
|
||||
buffers = [np.copy(v.numpy().copy()) for v in input_queues.values()]
|
||||
|
||||
if test_val is not None:
|
||||
match = all(np.array_equal(a, b) for a, b in zip(val, test_val, strict=True))
|
||||
assert match == expect_match, f"outputs {'differ from' if expect_match else 'match'} baseline (seed={seed})"
|
||||
if test_buffers is not None:
|
||||
match = all(np.array_equal(a, b) for a, b in zip(buffers, test_buffers, strict=True))
|
||||
assert match == expect_match, f"buffers {'differ from' if expect_match else 'match'} baseline (seed={seed})"
|
||||
return val, buffers
|
||||
|
||||
print('capture + replay')
|
||||
test_val, test_buffers = random_inputs_run(jit, SEED)
|
||||
print('pickle round trip')
|
||||
jit = pickle.loads(pickle.dumps(jit))
|
||||
random_inputs_run(jit, SEED, test_val, test_buffers, expect_match=True)
|
||||
random_inputs_run(jit, SEED+1, test_val, test_buffers, expect_match=False)
|
||||
return jit
|
||||
|
||||
|
||||
def _captured_devices(jit) -> set[str]:
|
||||
captured = getattr(jit, 'captured', None)
|
||||
infos = getattr(captured, 'expected_input_info', None)
|
||||
if not infos:
|
||||
return set()
|
||||
|
||||
devices: set[str] = set()
|
||||
for info in infos:
|
||||
if isinstance(info, tuple) and len(info) >= 4 and isinstance(info[3], str):
|
||||
devices.add(info[3])
|
||||
return devices
|
||||
|
||||
|
||||
def _slice_outputs(model_outputs: np.ndarray, output_slices: dict[str, slice]) -> dict[str, np.ndarray]:
|
||||
return {name: model_outputs[np.newaxis, tensor_slice] for name, tensor_slice in output_slices.items() if name != 'pad'}
|
||||
|
||||
|
||||
def _validate_pose_outputs(parsed_outputs: dict[str, np.ndarray]) -> None:
|
||||
from iqpilot.selfdrive.locationd.locationd import MIN_STD_SANITY_CHECK, ROTATION_SANITY_CHECK, TRANS_SANITY_CHECK
|
||||
|
||||
required = (
|
||||
'pose', 'pose_stds', 'wide_from_device_euler', 'wide_from_device_euler_stds',
|
||||
'road_transform', 'road_transform_stds',
|
||||
)
|
||||
missing = [name for name in required if name not in parsed_outputs]
|
||||
if missing:
|
||||
raise AssertionError(f"parsed supercombo outputs missing required odometry tensors: {missing}")
|
||||
|
||||
for name in required:
|
||||
values = parsed_outputs[name]
|
||||
if not np.isfinite(values).all():
|
||||
raise AssertionError(f"parsed supercombo output {name} contains non-finite values")
|
||||
|
||||
pose = parsed_outputs['pose'][0]
|
||||
pose_stds = parsed_outputs['pose_stds'][0]
|
||||
road_transform_stds = parsed_outputs['road_transform_stds'][0]
|
||||
wide_stds = parsed_outputs['wide_from_device_euler_stds'][0]
|
||||
|
||||
if pose_stds.min() <= MIN_STD_SANITY_CHECK:
|
||||
raise AssertionError(f"pose_stds min {pose_stds.min()} <= {MIN_STD_SANITY_CHECK}")
|
||||
if road_transform_stds.min() <= MIN_STD_SANITY_CHECK:
|
||||
raise AssertionError(f"road_transform_stds min {road_transform_stds.min()} <= {MIN_STD_SANITY_CHECK}")
|
||||
if wide_stds.min() <= MIN_STD_SANITY_CHECK:
|
||||
raise AssertionError(f"wide_from_device_euler_stds min {wide_stds.min()} <= {MIN_STD_SANITY_CHECK}")
|
||||
|
||||
if np.linalg.norm(pose[:3]) > TRANS_SANITY_CHECK:
|
||||
raise AssertionError(f"pose translation norm {np.linalg.norm(pose[:3])} exceeds {TRANS_SANITY_CHECK}")
|
||||
if np.linalg.norm(pose[3:]) > ROTATION_SANITY_CHECK:
|
||||
raise AssertionError(f"pose rotation norm {np.linalg.norm(pose[3:])} exceeds {ROTATION_SANITY_CHECK}")
|
||||
if np.linalg.norm(pose_stds[:3]) > 10 * TRANS_SANITY_CHECK:
|
||||
raise AssertionError(
|
||||
f"pose translation std norm {np.linalg.norm(pose_stds[:3])} exceeds {10 * TRANS_SANITY_CHECK}"
|
||||
)
|
||||
if np.linalg.norm(pose_stds[3:]) > 10 * ROTATION_SANITY_CHECK:
|
||||
raise AssertionError(
|
||||
f"pose rotation std norm {np.linalg.norm(pose_stds[3:])} exceeds {10 * ROTATION_SANITY_CHECK}"
|
||||
)
|
||||
|
||||
|
||||
def validate_supercombo_release(run_policy_jit, model_runner, model_metadata, frame_skip, expected_device: str) -> None:
|
||||
from iqpilot.selfdrive.iqmodeld.parser import PhaseParser
|
||||
|
||||
direct_fn = make_run_policy(model_runner, model_metadata, frame_skip)
|
||||
parser = PhaseParser()
|
||||
queue_factory = partial(make_input_queues, model_metadata['input_shapes'], frame_skip)
|
||||
image_shape = model_metadata['input_shapes']['img']
|
||||
|
||||
jit_queues, jit_npy = queue_factory(Device.DEFAULT)
|
||||
direct_queues, direct_npy = queue_factory(Device.DEFAULT)
|
||||
|
||||
for payload in (jit_npy, direct_npy):
|
||||
for name, value in payload.items():
|
||||
value[:] = 0 if value.dtype.kind in ('i', 'u') else 0.0
|
||||
|
||||
zero_inputs = {
|
||||
'warped': Tensor(np.zeros((2, 6, *image_shape[2:]), dtype=np.uint8), device=Device.DEFAULT).realize(),
|
||||
}
|
||||
|
||||
direct_outs, = direct_fn(**{k: direct_queues[k] for k in POLICY_INPUTS}, **zero_inputs)
|
||||
jit_outs, = run_policy_jit(**{k: jit_queues[k] for k in POLICY_INPUTS}, **zero_inputs)
|
||||
|
||||
direct_flat = direct_outs.numpy().astype(np.float32).reshape(-1)
|
||||
jit_flat = jit_outs.numpy().astype(np.float32).reshape(-1)
|
||||
|
||||
if not np.allclose(direct_flat, jit_flat, atol=1e-4, rtol=1e-4):
|
||||
max_delta = float(np.max(np.abs(direct_flat - jit_flat)))
|
||||
raise AssertionError(f"JIT supercombo output diverges from direct ONNX execution; max abs delta {max_delta}")
|
||||
|
||||
parsed = parser.parse_vision_outputs(_slice_outputs(jit_flat, model_metadata['output_slices']))
|
||||
_validate_pose_outputs(parsed)
|
||||
|
||||
captured_devices = _captured_devices(run_policy_jit)
|
||||
if expected_device and captured_devices and expected_device not in captured_devices:
|
||||
raise AssertionError(
|
||||
f"compiled run_policy backend mismatch: captured {sorted(captured_devices)} expected {expected_device}"
|
||||
)
|
||||
|
||||
|
||||
def _parse_size(s):
|
||||
w, h = s.lower().split('x')
|
||||
return int(w), int(h)
|
||||
|
||||
|
||||
def read_file_chunked_to_shm(path):
|
||||
from iqpilot.common.file_chunker import read_file_chunked
|
||||
from iqpilot.system.hardware.hw import Paths
|
||||
with tempfile.NamedTemporaryFile(prefix='compile_modeld_', dir=Paths.shm_path(), delete=False) as f:
|
||||
f.write(read_file_chunked(path))
|
||||
tmp_path = f.name
|
||||
atexit.register(lambda: os.path.exists(tmp_path) and os.remove(tmp_path))
|
||||
return tmp_path
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from tinygrad.nn.onnx import OnnxRunner
|
||||
from iqpilot.system.camerad.cameras.nv12_info import get_nv12_info
|
||||
from iqpilot.selfdrive.iqmodeld.metadata import build_metadata_record
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument('--model-size', type=_parse_size, required=True, help='model input WxH')
|
||||
p.add_argument('--camera-resolutions', type=_parse_size, nargs='+', required=True,
|
||||
help='camera resolutions WxH (one or more)')
|
||||
p.add_argument('--onnx', required=True)
|
||||
p.add_argument('--output', required=True)
|
||||
p.add_argument('--frame-skip', type=int, required=True)
|
||||
p.add_argument('--expected-device', default='QCOM', help='expected tinygrad backend baked into the JIT')
|
||||
args = p.parse_args()
|
||||
|
||||
model_path = read_file_chunked_to_shm(args.onnx)
|
||||
model_w, model_h = args.model_size
|
||||
|
||||
model_runner = OnnxRunner(model_path)
|
||||
out = {
|
||||
'metadata': build_metadata_record(model_path),
|
||||
'frame_skip': args.frame_skip,
|
||||
}
|
||||
|
||||
run_policy_jit = TinyJit(make_run_policy(model_runner, out['metadata'], args.frame_skip), prune=True)
|
||||
|
||||
make_policy_queues = partial(make_input_queues, out['metadata']['input_shapes'], args.frame_skip)
|
||||
make_random_model_inputs = partial(make_random_images, keys=['warped'], shape=(2, 6, *out['metadata']['input_shapes']['img'][2:]))
|
||||
out['run_policy'] = compile_jit(run_policy_jit, make_random_model_inputs, POLICY_INPUTS,
|
||||
make_policy_queues)
|
||||
validate_supercombo_release(out['run_policy'], model_runner, out['metadata'], args.frame_skip, args.expected_device)
|
||||
|
||||
for cam_w, cam_h in args.camera_resolutions:
|
||||
nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
|
||||
make_random_warp_inputs = partial(make_random_images, keys=['frame', 'big_frame'], shape=nv12.size, device=WARP_DEV)
|
||||
warp_enqueue = TinyJit(make_warp(nv12, model_w, model_h, args.frame_skip), prune=True)
|
||||
make_warp_queues = partial(make_warp_input_queues, out['metadata']['input_shapes'], args.frame_skip)
|
||||
out[(cam_w,cam_h)] = compile_jit(warp_enqueue, make_random_warp_inputs, WARP_INPUTS, make_warp_queues)
|
||||
captured_devices = _captured_devices(out[(cam_w,cam_h)])
|
||||
if args.expected_device and captured_devices and args.expected_device not in captured_devices:
|
||||
raise AssertionError(
|
||||
f"compiled warp backend mismatch for {cam_w}x{cam_h}: captured {sorted(captured_devices)} expected {args.expected_device}"
|
||||
)
|
||||
|
||||
with open(args.output, "wb") as f:
|
||||
pickle.dump(out, f)
|
||||
print(f"Saved JITs to {args.output} ({os.path.getsize(args.output) / 1e6:.2f} MB)")
|
||||
119
iqpilot/selfdrive/iqmodeld/tools/compile_warp.py
Normal file
119
iqpilot/selfdrive/iqmodeld/tools/compile_warp.py
Normal file
@@ -0,0 +1,119 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
|
||||
Compile the backend-neutral warp-only artifact: NV12 camera frames + 3x3
|
||||
transforms -> (2, 6, model_h/2, model_w/2) uint8 warped tensor, on the device
|
||||
GPU (QCOM). maciqmodeld runs this locally
|
||||
and feed the output to their backend, so the big model's image pipeline is
|
||||
bit-identical to comma's fused pkl warp stage.
|
||||
|
||||
Run ON the device (needs the QCOM backend):
|
||||
cd /data/openpilot && DEV=QCOM WARP_DEV=QCOM IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 \
|
||||
python3 iqpilot/selfdrive/iqmodeld/tools/compile_warp.py \
|
||||
--camera-resolutions 1928x1208 --output /data/models/emac_warp.pkl
|
||||
The artifact is then split per-resolution into Paths.model_root().
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import os
|
||||
import pickle
|
||||
from functools import partial
|
||||
|
||||
import numpy as np
|
||||
|
||||
SELFTEST_SEED = 20260817
|
||||
|
||||
from iqpilot.selfdrive.iqmodeld.temporal_state import DEFAULT_FRAME_SKIP, MODEL_INPUT_SPEC
|
||||
from iqpilot.selfdrive.iqmodeld.tools.compile_supercombo import (
|
||||
NV12Frame, WARP_INPUTS, compile_jit, make_random_images, make_warp, make_warp_input_queues,
|
||||
)
|
||||
|
||||
MODEL_SIZE = (MODEL_INPUT_SPEC["img"][0][3] * 2, MODEL_INPUT_SPEC["img"][0][2] * 2) # (512, 256)
|
||||
|
||||
|
||||
def _parse_size(s: str) -> tuple[int, int]:
|
||||
w, h = s.lower().split("x")
|
||||
return int(w), int(h)
|
||||
|
||||
|
||||
def compile_warp(cam_w: int, cam_h: int, out_path: str | None = None,
|
||||
frame_skip: int = DEFAULT_FRAME_SKIP) -> str:
|
||||
"""Compile the warp-only QCOM JIT for one camera resolution and write the pkl.
|
||||
Returns the artifact path. Callable from the workers so a fresh device
|
||||
self-provisions the warp instead of erroring — needs the QCOM backend."""
|
||||
# the QCOM warp env must be set before tinygrad is imported here
|
||||
os.environ.setdefault("DEV", "QCOM")
|
||||
os.environ.setdefault("WARP_DEV", "QCOM")
|
||||
os.environ.setdefault("IMAGE", "1")
|
||||
os.environ.setdefault("FLOAT16", "1")
|
||||
os.environ.setdefault("NOLOCALS", "1")
|
||||
os.environ.setdefault("JIT_BATCH_SIZE", "0")
|
||||
from tinygrad.engine.jit import TinyJit
|
||||
from iqpilot.system.camerad.cameras.nv12_info import get_nv12_info
|
||||
from iqpilot.system.hardware.hw import Paths
|
||||
|
||||
model_w, model_h = MODEL_SIZE
|
||||
input_shapes = {name: shape for name, (shape, _) in MODEL_INPUT_SPEC.items()}
|
||||
nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
|
||||
make_random_warp_inputs = partial(make_random_images, keys=["frame", "big_frame"],
|
||||
shape=nv12.size, device=os.getenv("WARP_DEV"))
|
||||
warp_jit = TinyJit(make_warp(nv12, model_w, model_h, frame_skip), prune=True)
|
||||
make_warp_queues = partial(make_warp_input_queues, input_shapes, frame_skip)
|
||||
compiled = compile_jit(warp_jit, make_random_warp_inputs, WARP_INPUTS, make_warp_queues)
|
||||
|
||||
# historical artifact name: already-provisioned devices keep their warp
|
||||
out_path = out_path or os.path.join(Paths.model_root(), f"emac_warp_{cam_w}x{cam_h}_tinygrad.pkl")
|
||||
os.makedirs(os.path.dirname(out_path), exist_ok=True)
|
||||
tmp = out_path + ".part"
|
||||
bundle = {(cam_w, cam_h): compiled, "frame_skip": frame_skip, "model_size": MODEL_SIZE}
|
||||
bundle["selftest"] = selftest_digest(compiled, cam_w, cam_h, nv12.size)
|
||||
with open(tmp, "wb") as f:
|
||||
pickle.dump(bundle, f)
|
||||
os.replace(tmp, out_path) # atomic: a reader never sees a half-written pkl
|
||||
return out_path
|
||||
|
||||
|
||||
def main() -> None:
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--camera-resolutions", type=_parse_size, nargs="+", default=[(1928, 1208)])
|
||||
p.add_argument("--output", default=None)
|
||||
p.add_argument("--frame-skip", type=int, default=DEFAULT_FRAME_SKIP)
|
||||
args = p.parse_args()
|
||||
for cam_w, cam_h in args.camera_resolutions:
|
||||
out = compile_warp(cam_w, cam_h, args.output, frame_skip=args.frame_skip)
|
||||
print(f"saved warp JIT to {out} ({os.path.getsize(out) / 1e6:.2f} MB)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
|
||||
def selftest_inputs(cam_w: int, cam_h: int, nv12_size: int):
|
||||
"""A fixed synthetic frame pair and pair of matrices. Deterministic so the
|
||||
digest is reproducible on the device that compiled the artifact."""
|
||||
rng = np.random.default_rng(SELFTEST_SEED)
|
||||
frame = rng.integers(0, 256, nv12_size, dtype=np.uint8)
|
||||
big_frame = rng.integers(0, 256, nv12_size, dtype=np.uint8)
|
||||
tfm = np.array([[0.7, 0.02, 300.0], [0.01, 0.7, 240.0], [0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
big_tfm = np.array([[0.5, 0.01, 380.0], [0.02, 0.5, 300.0], [0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
return frame, big_frame, tfm, big_tfm
|
||||
|
||||
|
||||
def selftest_digest(compiled, cam_w: int, cam_h: int, nv12_size: int) -> str:
|
||||
"""Hash the warp's output for a fixed input.
|
||||
|
||||
A warp artifact pinned to one tinygrad can still unpickle under another and
|
||||
then compute silently wrong, which reaches the model as a garbage image and
|
||||
looks like a bad model rather than a stale artifact. A version string cannot
|
||||
see that; running it can."""
|
||||
from tinygrad.tensor import Tensor
|
||||
frame, big_frame, tfm, big_tfm = selftest_inputs(cam_w, cam_h, nv12_size)
|
||||
dev = os.getenv("WARP_DEV") or "QCOM"
|
||||
out = compiled(tfm=Tensor(tfm, device="NPY").realize(),
|
||||
big_tfm=Tensor(big_tfm, device="NPY").realize(),
|
||||
frame=Tensor(frame, device=dev).realize(),
|
||||
big_frame=Tensor(big_frame, device=dev).realize())
|
||||
return hashlib.sha256(out.numpy().astype(np.uint8).tobytes()).hexdigest()
|
||||
44
iqpilot/selfdrive/iqmodeld/tools/convert_egpu_oob.py
Normal file
44
iqpilot/selfdrive/iqmodeld/tools/convert_egpu_oob.py
Normal file
@@ -0,0 +1,44 @@
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import os
|
||||
|
||||
os.environ.setdefault("DEV", "USB+AMD:LLVM")
|
||||
os.environ.setdefault("GMMU", "0")
|
||||
|
||||
from iqpilot.selfdrive.iqmodeld.egpu_helpers import patch_tinygrad_fetch_fw
|
||||
from iqpilot.selfdrive.iqmodeld.egpu_policy import dump_oob, is_oob, load_bundle
|
||||
|
||||
|
||||
def convert(src: str, dst: str) -> str:
|
||||
patch_tinygrad_fetch_fw()
|
||||
if is_oob(src):
|
||||
if src != dst:
|
||||
os.replace(src, dst)
|
||||
return dst
|
||||
bundle = load_bundle(src)
|
||||
tmp = dst + ".part"
|
||||
with open(tmp, "wb") as f:
|
||||
dump_oob(bundle, f)
|
||||
del bundle
|
||||
gc.collect()
|
||||
load_bundle(tmp)
|
||||
os.replace(tmp, dst)
|
||||
return dst
|
||||
|
||||
|
||||
def main() -> None:
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("src")
|
||||
p.add_argument("--out", default=None)
|
||||
args = p.parse_args()
|
||||
out = convert(args.src, args.out or args.src)
|
||||
print(f"converted -> {out} ({os.path.getsize(out) / 1e6:.1f} MB)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
331
iqpilot/selfdrive/iqmodeld/tools/daemon_jit_compiler.py
Normal file
331
iqpilot/selfdrive/iqmodeld/tools/daemon_jit_compiler.py
Normal file
@@ -0,0 +1,331 @@
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import atexit
|
||||
import os
|
||||
import pickle
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def _install_firmware_fetch_patch() -> None:
|
||||
import hashlib
|
||||
import pathlib
|
||||
|
||||
import zstandard
|
||||
from tinygrad import helpers
|
||||
|
||||
if not hasattr(helpers, "fetch_fw"):
|
||||
return
|
||||
|
||||
original_fetch = helpers.fetch_fw
|
||||
|
||||
def fetch_fw(path, name, sha256):
|
||||
archive_path = pathlib.Path(f"/lib/firmware/{path}/{name}.zst")
|
||||
if archive_path.is_file():
|
||||
blob = zstandard.ZstdDecompressor().stream_reader(archive_path.read_bytes()).read()
|
||||
if hashlib.sha256(blob).hexdigest() == sha256:
|
||||
return blob
|
||||
return original_fetch(path, name, sha256)
|
||||
|
||||
helpers.fetch_fw = fetch_fw
|
||||
|
||||
|
||||
_install_firmware_fetch_patch()
|
||||
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.engine.jit import TinyJit
|
||||
from tinygrad.helpers import Context
|
||||
from tinygrad.nn.onnx import OnnxRunner
|
||||
from tinygrad.tensor import Tensor
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FrameGeometry:
|
||||
width: int
|
||||
height: int
|
||||
stride: int
|
||||
y_height: int
|
||||
uv_height: int
|
||||
size: int
|
||||
|
||||
|
||||
WARP_INPUT_NAMES = ["img_q", "big_img_q", "tfm", "big_tfm"]
|
||||
POLICY_INPUT_NAMES = ["feat_q", "desire_q", "desire", "traffic_convention", "action_t"]
|
||||
WARP_DEV = os.getenv("WARP_DEV")
|
||||
|
||||
|
||||
def _random_tensor_inputs(keys: list[str], shape, device=None):
|
||||
return {key: Tensor.randint(shape, low=0, high=256, dtype="uint8", device=device).realize() for key in keys}
|
||||
|
||||
|
||||
def _project_frame(src_flat, inverse_matrix, dst_shape, src_shape, stride_pad, border_fill_val=None):
|
||||
dst_w, dst_h = dst_shape
|
||||
src_h, src_w = src_shape
|
||||
|
||||
x = Tensor.arange(dst_w, device=WARP_DEV).reshape(1, dst_w).expand(dst_h, dst_w).reshape(-1)
|
||||
y = Tensor.arange(dst_h, device=WARP_DEV).reshape(dst_h, 1).expand(dst_h, dst_w).reshape(-1)
|
||||
|
||||
src_x = inverse_matrix[0, 0] * x + inverse_matrix[0, 1] * y + inverse_matrix[0, 2]
|
||||
src_y = inverse_matrix[1, 0] * x + inverse_matrix[1, 1] * y + inverse_matrix[1, 2]
|
||||
src_w_scale = inverse_matrix[2, 0] * x + inverse_matrix[2, 1] * y + inverse_matrix[2, 2]
|
||||
|
||||
src_x = src_x / src_w_scale
|
||||
src_y = src_y / src_w_scale
|
||||
|
||||
rounded_x = Tensor.round(src_x)
|
||||
rounded_y = Tensor.round(src_y)
|
||||
clipped_x = rounded_x.clip(0, src_w - 1).cast("int")
|
||||
clipped_y = rounded_y.clip(0, src_h - 1).cast("int")
|
||||
gather_index = clipped_y * (src_w + stride_pad) + clipped_x
|
||||
sampled = src_flat[gather_index]
|
||||
|
||||
if border_fill_val is None:
|
||||
return sampled
|
||||
|
||||
inside = ((rounded_x >= 0) & (rounded_x <= src_w - 1) & (rounded_y >= 0) & (rounded_y <= src_h - 1)).cast(sampled.dtype)
|
||||
return sampled * inside + Tensor(border_fill_val, dtype=sampled.dtype) * (1 - inside)
|
||||
|
||||
|
||||
def _nv12_to_model_planes(yuv_frame):
|
||||
y_height = (yuv_frame.shape[0] * 2) // 3
|
||||
frame_width = yuv_frame.shape[1]
|
||||
return Tensor.cat(
|
||||
yuv_frame[0:y_height:2, 0::2],
|
||||
yuv_frame[1:y_height:2, 0::2],
|
||||
yuv_frame[0:y_height:2, 1::2],
|
||||
yuv_frame[1:y_height:2, 1::2],
|
||||
yuv_frame[y_height:y_height + y_height // 4].reshape((y_height // 2, frame_width // 2)),
|
||||
yuv_frame[y_height + y_height // 4:y_height + y_height // 2].reshape((y_height // 2, frame_width // 2)),
|
||||
dim=0,
|
||||
).reshape((6, y_height // 2, frame_width // 2))
|
||||
|
||||
|
||||
def _warp_kernel_factory(nv12: FrameGeometry, model_w: int, model_h: int):
|
||||
uv_offset = nv12.stride * nv12.y_height
|
||||
stride_pad = nv12.stride - nv12.width
|
||||
|
||||
def prepare_frame(nv12_blob, inverse_matrix):
|
||||
inverse_uv = inverse_matrix * Tensor([[1.0, 1.0, 0.5], [1.0, 1.0, 0.5], [2.0, 2.0, 1.0]], device=WARP_DEV)
|
||||
uv_plane = nv12_blob[uv_offset:uv_offset + nv12.uv_height * nv12.stride].reshape(nv12.uv_height, nv12.stride)
|
||||
with Context(SPLIT_REDUCEOP=0):
|
||||
y_plane = _project_frame(nv12_blob[:nv12.height * nv12.stride], inverse_matrix, (model_w, model_h), (nv12.height, nv12.width), stride_pad).realize()
|
||||
u_plane = _project_frame(uv_plane[:nv12.height // 2, :nv12.width:2].flatten(), inverse_uv, (model_w // 2, model_h // 2), (nv12.height // 2, nv12.width // 2), 0).realize()
|
||||
v_plane = _project_frame(uv_plane[:nv12.height // 2, 1:nv12.width:2].flatten(), inverse_uv, (model_w // 2, model_h // 2), (nv12.height // 2, nv12.width // 2), 0).realize()
|
||||
return _nv12_to_model_planes(y_plane.cat(u_plane).cat(v_plane).reshape((model_h * 3 // 2, model_w)))
|
||||
|
||||
return prepare_frame
|
||||
|
||||
|
||||
def _vision_queue_state(vision_shapes, frame_skip, device):
|
||||
img_shape = vision_shapes["img"]
|
||||
frame_history = img_shape[1] // 6
|
||||
queue_shape = (frame_skip * (frame_history - 1) + 1, 6, img_shape[2], img_shape[3])
|
||||
numpy_state = {
|
||||
"tfm": np.zeros((3, 3), dtype=np.float32),
|
||||
"big_tfm": np.zeros((3, 3), dtype=np.float32),
|
||||
}
|
||||
tensor_state = {
|
||||
"img_q": Tensor(np.zeros(queue_shape, dtype=np.uint8), device=device).contiguous().realize(),
|
||||
"big_img_q": Tensor(np.zeros(queue_shape, dtype=np.uint8), device=device).contiguous().realize(),
|
||||
**{name: Tensor(value, device="NPY").realize() for name, value in numpy_state.items()},
|
||||
}
|
||||
return tensor_state, numpy_state
|
||||
|
||||
|
||||
def _policy_queue_state(vision_shapes, policy_shapes, frame_skip, device):
|
||||
tensor_state, numpy_state = _vision_queue_state(vision_shapes, frame_skip, device)
|
||||
feature_shape = policy_shapes["features_buffer"]
|
||||
desire_shape = policy_shapes["desire_pulse"]
|
||||
traffic_shape = policy_shapes["traffic_convention"]
|
||||
action_shape = traffic_shape
|
||||
|
||||
policy_numpy = {
|
||||
"desire": np.zeros(desire_shape[2], dtype=np.float32),
|
||||
"traffic_convention": np.zeros(traffic_shape, dtype=np.float32),
|
||||
"action_t": np.zeros(action_shape, dtype=np.float32),
|
||||
}
|
||||
numpy_state.update(policy_numpy)
|
||||
tensor_state.update({
|
||||
"feat_q": Tensor(np.zeros((frame_skip * (feature_shape[1] - 1) + 1, feature_shape[0], feature_shape[2]), dtype=np.float32), device=device).contiguous().realize(),
|
||||
"desire_q": Tensor(np.zeros((frame_skip * desire_shape[1], desire_shape[0], desire_shape[2]), dtype=np.float32), device=device).contiguous().realize(),
|
||||
**{name: Tensor(value, device="NPY").realize() for name, value in policy_numpy.items()},
|
||||
})
|
||||
return tensor_state, numpy_state
|
||||
|
||||
|
||||
def _roll_queue(queue_tensor, incoming, sampler):
|
||||
queue_tensor.assign(queue_tensor[1:].cat(incoming, dim=0).contiguous())
|
||||
return sampler(queue_tensor)
|
||||
|
||||
|
||||
def _sample_sparse(queue_tensor, frame_skip):
|
||||
return queue_tensor[::frame_skip].contiguous().flatten(0, 1).unsqueeze(0)
|
||||
|
||||
|
||||
def _sample_desire(queue_tensor, frame_skip):
|
||||
return queue_tensor.reshape(-1, frame_skip, *queue_tensor.shape[1:]).max(1).flatten(0, 1).unsqueeze(0)
|
||||
|
||||
|
||||
def _build_warp_enqueuer(nv12: FrameGeometry, model_w: int, model_h: int, frame_skip: int):
|
||||
prepare_frame = _warp_kernel_factory(nv12, model_w, model_h)
|
||||
sparse_sampler = partial(_sample_sparse, frame_skip=frame_skip)
|
||||
|
||||
def enqueue(img_q, big_img_q, tfm, big_tfm, frame, big_frame):
|
||||
tfm = tfm.to(WARP_DEV)
|
||||
big_tfm = big_tfm.to(WARP_DEV)
|
||||
Tensor.realize(tfm, big_tfm)
|
||||
|
||||
warped_main = prepare_frame(frame, tfm).unsqueeze(0).to(Device.DEFAULT)
|
||||
warped_big = prepare_frame(big_frame, big_tfm).unsqueeze(0).to(Device.DEFAULT)
|
||||
return (
|
||||
_roll_queue(img_q, warped_main, sparse_sampler),
|
||||
_roll_queue(big_img_q, warped_big, sparse_sampler),
|
||||
)
|
||||
|
||||
return enqueue
|
||||
|
||||
|
||||
def _policy_executor(model_runners, model_metadata, frame_skip):
|
||||
desire_sampler = partial(_sample_desire, frame_skip=frame_skip)
|
||||
sparse_sampler = partial(_sample_sparse, frame_skip=frame_skip)
|
||||
hidden_slice = model_metadata["vision"]["output_slices"]["hidden_state"]
|
||||
|
||||
def execute(img, big_img, feat_q, desire_q, desire, traffic_convention, action_t):
|
||||
desire = desire.to(Device.DEFAULT)
|
||||
traffic_convention = traffic_convention.to(Device.DEFAULT)
|
||||
action_t = action_t.to(Device.DEFAULT)
|
||||
Tensor.realize(desire, traffic_convention, action_t)
|
||||
|
||||
desire_buffer = _roll_queue(desire_q, desire.reshape(1, 1, -1), desire_sampler)
|
||||
vision_output = next(iter(model_runners["vision"]({"img": img, "big_img": big_img}).values())).cast("float32")
|
||||
|
||||
hidden_state = vision_output[:, hidden_slice].reshape(1, -1).unsqueeze(0)
|
||||
feature_buffer = _roll_queue(feat_q, hidden_state, sparse_sampler)
|
||||
|
||||
on_inputs = {
|
||||
"features_buffer": feature_buffer,
|
||||
"desire_pulse": desire_buffer,
|
||||
"traffic_convention": traffic_convention,
|
||||
"action_t": action_t,
|
||||
}
|
||||
on_output = next(iter(model_runners["on_policy"](on_inputs).values())).cast("float32")
|
||||
off_output = next(iter(model_runners["off_policy"](on_inputs).values())).cast("float32")
|
||||
return vision_output, on_output, off_output
|
||||
|
||||
return execute
|
||||
|
||||
|
||||
def _replay_and_freeze(jit_runner, random_inputs_factory, queue_keys, queue_factory):
|
||||
seed = 42
|
||||
|
||||
def validate(fn, seed_value, baseline_output=None, baseline_buffers=None, expect_match=True):
|
||||
queue_tensors, numpy_values = queue_factory(Device.DEFAULT)
|
||||
np.random.seed(seed_value)
|
||||
Tensor.manual_seed(seed_value)
|
||||
|
||||
replay_count = 1 if (baseline_output is not None or baseline_buffers is not None) else 3
|
||||
for run_index in range(replay_count):
|
||||
for value in numpy_values.values():
|
||||
value[:] = np.random.randn(*value.shape).astype(value.dtype)
|
||||
Device.default.synchronize()
|
||||
random_inputs = random_inputs_factory()
|
||||
start = time.perf_counter()
|
||||
outputs = fn(**{key: queue_tensors[key] for key in queue_keys}, **random_inputs)
|
||||
enqueue_done = time.perf_counter()
|
||||
Device.default.synchronize()
|
||||
total_done = time.perf_counter()
|
||||
print(f" [{run_index + 1}/{replay_count}] enqueue {(enqueue_done - start) * 1e3:6.2f} ms -- total {(total_done - start) * 1e3:6.2f} ms")
|
||||
|
||||
if run_index == 0:
|
||||
output_snapshot = [np.copy(value.numpy()) for value in outputs]
|
||||
buffer_snapshot = [np.copy(value.numpy().copy()) for value in queue_tensors.values()]
|
||||
|
||||
if baseline_output is not None:
|
||||
matches = all(np.array_equal(current, reference) for current, reference in zip(output_snapshot, baseline_output, strict=True))
|
||||
assert matches == expect_match, f"outputs {'differ from' if expect_match else 'match'} baseline (seed={seed_value})"
|
||||
if baseline_buffers is not None:
|
||||
matches = all(np.array_equal(current, reference) for current, reference in zip(buffer_snapshot, baseline_buffers, strict=True))
|
||||
assert matches == expect_match, f"buffers {'differ from' if expect_match else 'match'} baseline (seed={seed_value})"
|
||||
return output_snapshot, buffer_snapshot
|
||||
|
||||
print("capture + replay")
|
||||
first_output, first_buffers = validate(jit_runner, seed)
|
||||
print("pickle round trip")
|
||||
frozen = pickle.loads(pickle.dumps(jit_runner))
|
||||
validate(frozen, seed, first_output, first_buffers, expect_match=True)
|
||||
validate(frozen, seed + 1, first_output, first_buffers, expect_match=False)
|
||||
return frozen
|
||||
|
||||
|
||||
def _parse_size(text: str) -> tuple[int, int]:
|
||||
width, height = text.lower().split("x")
|
||||
return int(width), int(height)
|
||||
|
||||
|
||||
def _read_file_to_shared_memory(path: str) -> str:
|
||||
from iqpilot.common.file_chunker import read_file_chunked
|
||||
from iqpilot.system.hardware.hw import Paths
|
||||
|
||||
shm_path = os.path.join(Paths.shm_path(), os.path.basename(path))
|
||||
atexit.register(lambda: os.path.exists(shm_path) and os.remove(shm_path))
|
||||
with open(shm_path, "wb") as handle:
|
||||
handle.write(read_file_chunked(path))
|
||||
return shm_path
|
||||
|
||||
|
||||
def _arg_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--model-size", type=_parse_size, required=True, help="model input WxH")
|
||||
parser.add_argument("--camera-resolutions", type=_parse_size, nargs="+", required=True, help="camera resolutions WxH (one or more)")
|
||||
parser.add_argument("--vision-onnx", required=True)
|
||||
parser.add_argument("--off-policy-onnx", required=True)
|
||||
parser.add_argument("--on-policy-onnx", required=True)
|
||||
parser.add_argument("--output", required=True)
|
||||
parser.add_argument("--frame-skip", type=int, required=True)
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
from iqpilot.selfdrive.iqmodeld.metadata import build_metadata_record
|
||||
from iqpilot.system.camerad.cameras.nv12_info import get_nv12_info
|
||||
|
||||
args = _arg_parser().parse_args(argv)
|
||||
model_w, model_h = args.model_size
|
||||
|
||||
model_paths = {
|
||||
"vision": _read_file_to_shared_memory(args.vision_onnx),
|
||||
"off_policy": _read_file_to_shared_memory(args.off_policy_onnx),
|
||||
"on_policy": _read_file_to_shared_memory(args.on_policy_onnx),
|
||||
}
|
||||
model_runners = {name: OnnxRunner(path) for name, path in model_paths.items()}
|
||||
metadata = {name: build_metadata_record(path) for name, path in model_paths.items()}
|
||||
|
||||
assert metadata["off_policy"]["input_shapes"] == metadata["on_policy"]["input_shapes"]
|
||||
|
||||
output_package: dict = {"metadata": metadata}
|
||||
policy_jit = TinyJit(_policy_executor(model_runners, metadata, args.frame_skip), prune=True)
|
||||
policy_queue_factory = partial(_policy_queue_state, metadata["vision"]["input_shapes"], metadata["on_policy"]["input_shapes"], args.frame_skip)
|
||||
random_model_inputs = partial(_random_tensor_inputs, keys=["img", "big_img"], shape=metadata["vision"]["input_shapes"]["img"])
|
||||
output_package["run_policy"] = _replay_and_freeze(policy_jit, random_model_inputs, POLICY_INPUT_NAMES, policy_queue_factory)
|
||||
|
||||
for cam_w, cam_h in args.camera_resolutions:
|
||||
nv12 = FrameGeometry(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
|
||||
warp_jit = TinyJit(_build_warp_enqueuer(nv12, model_w, model_h, args.frame_skip), prune=True)
|
||||
warp_queue_factory = partial(_vision_queue_state, metadata["vision"]["input_shapes"], args.frame_skip)
|
||||
random_warp_inputs = partial(_random_tensor_inputs, keys=["frame", "big_frame"], shape=nv12.size, device=WARP_DEV)
|
||||
output_package[(cam_w, cam_h)] = _replay_and_freeze(warp_jit, random_warp_inputs, WARP_INPUT_NAMES, warp_queue_factory)
|
||||
|
||||
output_package["frame_skip"] = args.frame_skip
|
||||
with open(args.output, "wb") as handle:
|
||||
pickle.dump(output_package, handle)
|
||||
print(f"Saved JITs to {args.output} ({os.path.getsize(args.output) / 1e6:.2f} MB)")
|
||||
return 0
|
||||
64
iqpilot/selfdrive/iqmodeld/tools/egpu_host_mock.py
Normal file
64
iqpilot/selfdrive/iqmodeld/tools/egpu_host_mock.py
Normal file
@@ -0,0 +1,64 @@
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
|
||||
DEFAULT_ARCH = "gfx1200"
|
||||
MOCK_DEV = "MOCKUSB+AMD:LLVM"
|
||||
|
||||
|
||||
def tinygrad_tree() -> str:
|
||||
override = os.environ.get("IQ_TINYGRAD_TREE")
|
||||
if override:
|
||||
return override
|
||||
here = os.path.dirname(os.path.abspath(__file__))
|
||||
root = os.path.abspath(os.path.join(here, "..", "..", "..", ".."))
|
||||
return os.path.join(root, "components", "tinygrad")
|
||||
|
||||
|
||||
def activate(arch: str = DEFAULT_ARCH, execute: bool = False) -> None:
|
||||
assert "tinygrad" not in sys.modules, "egpu_host_mock.activate must run before tinygrad is imported"
|
||||
os.environ["DEV"] = f"{MOCK_DEV}:{arch}"
|
||||
tree = tinygrad_tree()
|
||||
if tree not in sys.path:
|
||||
sys.path.insert(0, tree)
|
||||
from test.mockgpu.am import amgpu
|
||||
|
||||
# The mock dock models 512MB VRAM; big-model weights alone exceed that. Must be set before amdriver binds it.
|
||||
amgpu.VRAM_SIZE = int(os.environ.get("IQ_MOCK_VRAM_GB", "4")) << 30
|
||||
from tinygrad.runtime.autogen import libc
|
||||
if sys.platform == "darwin":
|
||||
# A Homebrew-LLVM gfx1200 kernel (no s_code_end padding) hung a real dock; ship only container-built artifacts.
|
||||
print("egpu_host_mock: native macOS LLVM output is for tests only; use scripts/iqpilot/host_egpu_compile_docker.sh for artifacts",
|
||||
file=sys.stderr)
|
||||
|
||||
def memfd_create(name, flags):
|
||||
fd, path = tempfile.mkstemp(prefix=b"iq_mock_" + bytes(name) + b"_")
|
||||
os.unlink(path)
|
||||
return fd
|
||||
libc.memfd_create = memfd_create
|
||||
if not hasattr(libc, "MFD_CLOEXEC"):
|
||||
libc.MFD_CLOEXEC = 1
|
||||
if not execute:
|
||||
import ctypes
|
||||
from test.mockgpu.amd import amdgpu
|
||||
amdgpu.remu.run_asm = lambda *args, **kwargs: 0
|
||||
pm4_wait = amdgpu.PM4Executor._exec_wait_reg_mem
|
||||
sdma_poll = amdgpu.SDMAExecutor._execute_poll_regmem
|
||||
|
||||
# Without kernel execution no memory wait carries information; a blocked wait would need a host write to re-poll it.
|
||||
def pm4_wait_passthrough(self, n):
|
||||
if not pm4_wait(self, n):
|
||||
self.rptr[0] += 7
|
||||
return True
|
||||
|
||||
def sdma_poll_passthrough(self):
|
||||
if not sdma_poll(self):
|
||||
self.rptr[0] += ctypes.sizeof(amdgpu.sdma_pkts.poll_regmem)
|
||||
return True
|
||||
amdgpu.PM4Executor._exec_wait_reg_mem = pm4_wait_passthrough
|
||||
amdgpu.SDMAExecutor._execute_poll_regmem = sdma_poll_passthrough
|
||||
128
iqpilot/selfdrive/iqmodeld/tools/install_models_pc.py
Executable file
128
iqpilot/selfdrive/iqmodeld/tools/install_models_pc.py
Executable file
@@ -0,0 +1,128 @@
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import pickle
|
||||
import shutil
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import onnx
|
||||
|
||||
from iqpilot.system.hardware.hw import Paths
|
||||
|
||||
_MODEL_STEMS = ("driving_off_policy", "driving_on_policy", "driving_policy", "driving_vision")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _ModelBundle:
|
||||
stem: str
|
||||
onnx_path: Path
|
||||
artifact_path: Path
|
||||
metadata_path: Path
|
||||
|
||||
|
||||
def _tensor_shape(value_info) -> tuple[int, ...]:
|
||||
return tuple(int(dim.dim_value) for dim in value_info.type.tensor_type.shape.dim)
|
||||
|
||||
|
||||
def _metadata_property(graph_model, key: str) -> str | None:
|
||||
for property_item in graph_model.metadata_props:
|
||||
if property_item.key == key:
|
||||
return property_item.value
|
||||
return None
|
||||
|
||||
|
||||
def _decode_output_slices(encoded_value: str):
|
||||
return pickle.loads(base64.b64decode(encoded_value.encode()))
|
||||
|
||||
|
||||
def _metadata_record(graph_model) -> dict:
|
||||
encoded_slices = _metadata_property(graph_model, "output_slices")
|
||||
if encoded_slices is None:
|
||||
raise ValueError("output_slices metadata missing")
|
||||
return {
|
||||
"model_checkpoint": _metadata_property(graph_model, "model_checkpoint"),
|
||||
"output_slices": _decode_output_slices(encoded_slices),
|
||||
"input_shapes": {item.name: _tensor_shape(item) for item in graph_model.graph.input},
|
||||
"output_shapes": {item.name: _tensor_shape(item) for item in graph_model.graph.output},
|
||||
}
|
||||
|
||||
|
||||
def generate_metadata_pkl(model_path, output_path):
|
||||
try:
|
||||
graph_model = onnx.load(str(model_path))
|
||||
metadata = _metadata_record(graph_model)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
with open(output_path, "wb") as handle:
|
||||
pickle.dump(metadata, handle)
|
||||
return True
|
||||
|
||||
|
||||
def _discover_model_bundles(model_dir: Path) -> list[_ModelBundle]:
|
||||
bundles: list[_ModelBundle] = []
|
||||
for stem in _MODEL_STEMS:
|
||||
onnx_path = model_dir / f"{stem}.onnx"
|
||||
if not onnx_path.exists():
|
||||
continue
|
||||
bundles.append(_ModelBundle(
|
||||
stem=stem,
|
||||
onnx_path=onnx_path,
|
||||
artifact_path=model_dir / f"{stem}_tinygrad.pkl",
|
||||
metadata_path=model_dir / f"{stem}_metadata.pkl",
|
||||
))
|
||||
return bundles
|
||||
|
||||
|
||||
def _prompt_short_name(found_stems: list[str]) -> str | None:
|
||||
try:
|
||||
response = input(f"Found models ({', '.join(found_stems)}). Enter model short name (e.g. wmiv4): ").strip()
|
||||
except EOFError:
|
||||
return None
|
||||
return response or None
|
||||
|
||||
|
||||
def _ensure_metadata_file(bundle: _ModelBundle) -> None:
|
||||
if bundle.metadata_path.exists():
|
||||
return
|
||||
generate_metadata_pkl(bundle.onnx_path, bundle.metadata_path)
|
||||
|
||||
|
||||
def _install_bundle(bundle: _ModelBundle, suffix: str, destination_root: Path) -> None:
|
||||
_ensure_metadata_file(bundle)
|
||||
renamed_artifact = destination_root / f"{bundle.stem}_{suffix}_tinygrad.pkl"
|
||||
renamed_metadata = destination_root / f"{bundle.stem}_{suffix}_metadata.pkl"
|
||||
if bundle.artifact_path.exists():
|
||||
shutil.move(str(bundle.artifact_path), str(renamed_artifact))
|
||||
if bundle.metadata_path.exists():
|
||||
shutil.move(str(bundle.metadata_path), str(renamed_metadata))
|
||||
|
||||
|
||||
def install_models(model_dir):
|
||||
source_root = Path(model_dir)
|
||||
bundles = _discover_model_bundles(source_root)
|
||||
if not bundles:
|
||||
return
|
||||
|
||||
short_name = _prompt_short_name([bundle.stem for bundle in bundles])
|
||||
if short_name is None:
|
||||
print("No name provided, skipping installation.")
|
||||
return
|
||||
|
||||
destination_root = Path(Paths.model_root())
|
||||
destination_root.mkdir(parents=True, exist_ok=True)
|
||||
for bundle in bundles:
|
||||
_install_bundle(bundle, short_name, destination_root)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if len(sys.argv) < 2:
|
||||
print("Usage: install_models_pc.py <model_dir>")
|
||||
sys.exit(1)
|
||||
install_models(sys.argv[1])
|
||||
64
iqpilot/selfdrive/iqmodeld/tools/oob_rewrite.py
Normal file
64
iqpilot/selfdrive/iqmodeld/tools/oob_rewrite.py
Normal file
@@ -0,0 +1,64 @@
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import pickletools
|
||||
import struct
|
||||
|
||||
from iqpilot.selfdrive.iqmodeld.egpu_policy import OOB_MAGIC
|
||||
|
||||
MIN_OOB_BYTES = 1 << 16
|
||||
NEXT_BUFFER = b"\x97"
|
||||
READONLY_BUFFER = b"\x98"
|
||||
|
||||
|
||||
def rewrite_oob(src: str, dst: str, min_bytes: int = MIN_OOB_BYTES) -> tuple[int, int]:
|
||||
# tinygrad pickles device buffers as PickleBuffers, which land in-band as BYTEARRAY8/BINBYTES8
|
||||
# without a buffer_callback; moving those opcodes out-of-band is byte-for-byte what a protocol-5
|
||||
# dump with a buffer_callback produces, so nothing has to be unpickled (no dock needed).
|
||||
with open(src, "rb") as f:
|
||||
data = f.read()
|
||||
ops = list(pickletools.genops(data))
|
||||
proto = next((arg for op, arg, _ in ops if op.name == "PROTO"), 0)
|
||||
if proto < 5:
|
||||
raise ValueError(f"{src} is pickle protocol {proto}; out-of-band buffers need protocol 5")
|
||||
moved = 0
|
||||
tmp = dst + ".part"
|
||||
with open(tmp, "wb") as out, open(tmp + ".buf", "wb") as bufs:
|
||||
ops_stream = bytearray()
|
||||
for i, (op, arg, pos) in enumerate(ops):
|
||||
end = ops[i + 1][2] if i + 1 < len(ops) else len(data)
|
||||
if op.name in ("BYTEARRAY8", "BINBYTES8", "BINBYTES") and len(arg) >= min_bytes:
|
||||
ops_stream += NEXT_BUFFER
|
||||
if op.name != "BYTEARRAY8":
|
||||
ops_stream += READONLY_BUFFER
|
||||
bufs.write(struct.pack("<q", len(arg)))
|
||||
bufs.write(arg)
|
||||
moved += 1
|
||||
else:
|
||||
ops_stream += data[pos:end]
|
||||
out.write(OOB_MAGIC)
|
||||
out.write(struct.pack("<q", len(ops_stream)))
|
||||
out.write(ops_stream)
|
||||
with open(tmp, "ab") as out, open(tmp + ".buf", "rb") as bufs:
|
||||
while chunk := bufs.read(1 << 24):
|
||||
out.write(chunk)
|
||||
os.remove(tmp + ".buf")
|
||||
os.replace(tmp, dst)
|
||||
return moved, len(ops)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("src")
|
||||
p.add_argument("dst")
|
||||
args = p.parse_args()
|
||||
moved, total = rewrite_oob(args.src, args.dst)
|
||||
print(f"{args.dst}: moved {moved} buffers out-of-band ({total} opcodes, {os.path.getsize(args.dst) / 1e6:.1f} MB)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user