IQ.Pilot Release Commit @ 763bad7

iqpilot/selfdrive/iqmodeld/egpu_helpers.py

@@ -57,6 +57,11 @@ def egpu_pkl_path(meta: dict) -> str:
   return os.path.join(Paths.model_root(), f"egpu_{meta['key']}_{meta['sha256'][:8]}_amd_tinygrad.pkl")
 
 
+def egpu_policy_pkl_path(meta: dict) -> str:
+  from iqpilot.system.hardware.hw import Paths
+  return os.path.join(Paths.model_root(), f"egpu_{meta['key']}_{meta['sha256'][:8]}_amd_policy.pkl")
+
+
 def onnx_cache_path(meta: dict) -> str:
   from iqpilot.system.hardware.hw import Paths
   return os.path.join(Paths.model_root(), f"{meta['model_name']}_{meta['sha256'][:8]}.onnx")
@@ -137,6 +142,15 @@ def download_onnx(meta: dict, progress_cb=None) -> str:
   return path
 
 
+def download_precompiled(meta: dict, progress_cb=None, policy: bool = False) -> str | None:
+  art = meta.get("egpu_policy_artifact" if policy else "egpu_artifact")
+  if not art or not art.get("objects"):
+    return None
+  from iqpilot.selfdrive.iqmodeld.model_bundle_downloader import download_lfs_bundle
+  dest = egpu_policy_pkl_path(meta) if policy else egpu_pkl_path(meta)
+  return download_lfs_bundle(art["objects"], dest, art["sha256"], int(art.get("size", 0)), progress_cb=progress_cb)
+
+
 def patch_tinygrad_fetch_fw() -> None:
   import pathlib
 

iqpilot/selfdrive/iqmodeld/egpu_pipeline.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 
 import numpy as np
 
+from iqpilot.selfdrive.iqmodeld.egpu_policy import PolicyRunner
 from iqpilot.selfdrive.iqmodeld.temporal_state import MODEL_INPUT_SPEC, TemporalInputState, spec_from_meta
 
 
@@ -25,13 +26,17 @@ class EgpuPipeline:
 
   def run(self, warped: np.ndarray, desire_vec: np.ndarray, traffic_convention: np.ndarray,
           action_t: np.ndarray) -> np.ndarray:
-    inputs = self.state.push_and_materialize(warped, desire_vec, traffic_convention, action_t)
-    out = np.asarray(self.infer_fn(inputs), dtype=np.float32).reshape(-1)
+    if isinstance(self.infer_fn, PolicyRunner):
+      out = np.asarray(self.infer_fn.run(warped, desire_vec, traffic_convention, action_t), dtype=np.float32).reshape(-1)
+    else:
+      inputs = self.state.push_and_materialize(warped, desire_vec, traffic_convention, action_t)
+      out = np.asarray(self.infer_fn(inputs), dtype=np.float32).reshape(-1)
     if out.shape[0] != self.output_len:
       raise EgpuPipelineError(f"eGPU output length {out.shape[0]} != {self.output_len}")
     if not np.isfinite(out).all():
       raise EgpuPipelineError("eGPU output contains non-finite values")
-    self.state.note_hidden_state(out, self.hidden_slice)
+    if not isinstance(self.infer_fn, PolicyRunner):
+      self.state.note_hidden_state(out, self.hidden_slice)
     return out
 
 

iqpilot/selfdrive/iqmodeld/egpu_policy.py

@@ -0,0 +1,117 @@
+"""
+Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
+"""
+from __future__ import annotations
+
+import math
+
+import numpy as np
+
+POLICY_FORMAT = 2
+QUEUE_NAMES = ("img_q", "big_img_q", "feat_q", "desire_q")
+PACKED_ORDER = ("desire", "traffic_convention", "action_t", "prev_feat")
+
+
+def packed_layout(input_spec: dict) -> tuple[dict[str, tuple[int, ...]], list[int]]:
+  dp = input_spec["desire_pulse"][0]
+  fb = input_spec["features_buffer"][0]
+  shapes = {
+    "desire": (dp[2],),
+    "traffic_convention": tuple(input_spec["traffic_convention"][0]),
+    "action_t": tuple(input_spec["action_t"][0]),
+    "prev_feat": (fb[0], math.prod(fb[2:])),
+  }
+  return shapes, [math.prod(s) for s in shapes.values()]
+
+
+def queue_shapes(input_spec: dict, frame_skip: int) -> dict[str, tuple[tuple[int, ...], str]]:
+  img = input_spec["img"][0]
+  fb = input_spec["features_buffer"][0]
+  dp = input_spec["desire_pulse"][0]
+  n_frames = img[1] // 6
+  img_buf = (frame_skip * (n_frames - 1) + 1, 6, img[2], img[3])
+  return {
+    "img_q": (img_buf, "uint8"),
+    "big_img_q": (img_buf, "uint8"),
+    "feat_q": ((frame_skip * fb[1], fb[0], math.prod(fb[2:])), "float32"),
+    "desire_q": ((frame_skip * dp[1], dp[0], dp[2]), "float32"),
+  }
+
+
+def make_queues(input_spec: dict, frame_skip: int, device: str) -> dict:
+  from tinygrad.tensor import Tensor
+  return {name: Tensor(np.zeros(shape, dtype=dtype), device=device).contiguous().realize()
+          for name, (shape, dtype) in queue_shapes(input_spec, frame_skip).items()}
+
+
+class PackedInputs:
+  def __init__(self, input_spec: dict):
+    from tinygrad.tensor import Tensor
+    self.shapes, self.sizes = packed_layout(input_spec)
+    self.array = np.zeros(sum(self.sizes), dtype=np.float32)
+    self.views = dict(zip(self.shapes, [v.reshape(s) for s, v in zip(self.shapes.values(), np.split(self.array, np.cumsum(self.sizes[:-1])))], strict=True))
+    self.tensor = Tensor(self.array, device="NPY").realize()
+
+
+def make_run_policy(model_runner, input_spec: dict, frame_skip: int, device: str):
+  from tinygrad.tensor import Tensor
+  shapes, sizes = packed_layout(input_spec)
+  fb = input_spec["features_buffer"][0]
+
+  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):
+    return buf[::frame_skip].contiguous().flatten(0, 1).unsqueeze(0)
+
+  def sample_desire(buf):
+    return buf.reshape(-1, frame_skip, *buf.shape[1:]).max(1).flatten(0, 1).unsqueeze(0)
+
+  def run_policy(warped, img_q, big_img_q, feat_q, desire_q, packed_npy_inputs):
+    packed_npy_inputs = packed_npy_inputs.to(device)
+    warped = warped.to(device)
+    Tensor.realize(packed_npy_inputs, warped)
+    img = shift_and_sample(img_q, warped[0:1], sample_skip)
+    big_img = shift_and_sample(big_img_q, warped[1:2], sample_skip)
+    desire, traffic_convention, action_t, prev_feat = (t.reshape(s) for t, s in zip(packed_npy_inputs.split(sizes), shapes.values(), strict=True))
+    desire_buf = shift_and_sample(desire_q, desire.reshape(1, 1, -1), sample_desire)
+    feat_buf = shift_and_sample(feat_q, prev_feat.reshape(1, 1, -1), sample_skip)
+    inputs = {
+      "img": img,
+      "big_img": big_img,
+      "features_buffer": feat_buf.reshape(fb),
+      "desire_pulse": desire_buf,
+      "traffic_convention": traffic_convention,
+      "action_t": action_t,
+    }
+    out = next(iter(model_runner(inputs).values())).cast("float32")
+    return out.reshape(-1),
+
+  return run_policy
+
+
+class PolicyRunner:
+  def __init__(self, jit, input_spec: dict, frame_skip: int, hidden_slice: slice, device: str):
+    from tinygrad.tensor import Tensor
+    self._Tensor = Tensor
+    self._jit = jit
+    self._queues = make_queues(input_spec, frame_skip, device)
+    self._packed = PackedInputs(input_spec)
+    self._hidden = hidden_slice
+    self._prev_desire = np.zeros(input_spec["desire_pulse"][0][2], dtype=np.float32)
+    self._warped_shape = (2, 6, *input_spec["img"][0][2:])
+
+  def run(self, warped: np.ndarray, desire_pulse: np.ndarray, traffic_convention: np.ndarray,
+          action_t: np.ndarray) -> np.ndarray:
+    cur = desire_pulse.astype(np.float32, copy=False)
+    v = self._packed.views
+    v["desire"][:] = np.where(cur - self._prev_desire > 0.99, cur, 0)
+    self._prev_desire[:] = cur
+    v["traffic_convention"][:] = np.asarray(traffic_convention, dtype=np.float32).reshape(v["traffic_convention"].shape)
+    v["action_t"][:] = np.asarray(action_t, dtype=np.float32).reshape(v["action_t"].shape)
+    warped_t = self._Tensor(np.ascontiguousarray(warped, dtype=np.uint8).reshape(self._warped_shape), device="NPY").realize()
+    out, = self._jit(warped=warped_t, packed_npy_inputs=self._packed.tensor, **self._queues)
+    flat = out.numpy().reshape(-1)
+    v["prev_feat"][:] = flat[self._hidden].reshape(v["prev_feat"].shape)
+    return flat

iqpilot/selfdrive/iqmodeld/iqegpumodeld.py

@@ -38,8 +38,8 @@ from iqpilot.selfdrive.iqmodeld.driving_action import (
   DESIRE_LEN, LAT_SMOOTH_SECONDS, LONG_SMOOTH_SECONDS, get_action_from_model,
 )
 from iqpilot.selfdrive.iqmodeld.egpu_helpers import (
-  download_onnx, egpu_pkl_path, egpu_present_consented, egpu_selected, local_onnx, patch_tinygrad_fetch_fw,
-  quarantine_artifact, resolve_backend, usbgpu_present,
+  download_onnx, download_precompiled, egpu_pkl_path, egpu_policy_pkl_path, egpu_present_consented, egpu_selected, local_onnx,
+  patch_tinygrad_fetch_fw, quarantine_artifact, resolve_backend, usbgpu_present,
 )
 from iqpilot.selfdrive.iqmodeld.egpu_model import resolve_egpu_model
 from iqpilot.selfdrive.iqmodeld.egpu_pipeline import EgpuPipeline, EgpuPipelineError, make_big_channel_payload
@@ -47,6 +47,7 @@ from iqpilot.selfdrive.iqmodeld.egpu_telemetry import EgpuDockTelemetry
 from iqpilot.selfdrive.iqmodeld.messaging import DrivePacketMemory, populate_drive_messages, populate_odometry_message
 from iqpilot.selfdrive.iqmodeld.metadata import Meta20hz
 from iqpilot.selfdrive.iqmodeld.model_channel import BIG_CHANNEL, ModelChannel
+from iqpilot.selfdrive.iqmodeld.egpu_policy import POLICY_FORMAT, PolicyRunner
 from iqpilot.selfdrive.iqmodeld.model_warp import FrameWarp
 from iqpilot.selfdrive.iqmodeld.parser import PhaseParser
 
@@ -89,9 +90,10 @@ def _wait_for_egpu(params: Params) -> None:
 
 def _compile_in_subprocess(meta: dict, onnx_path: str, pkl_path: str) -> None:
   cmd = [sys.executable, "-m", "iqpilot.selfdrive.iqmodeld.tools.compile_egpu_model",
-         "--model", meta["key"], "--onnx", onnx_path, "--output", pkl_path]
+         "--model", meta["key"], "--onnx", onnx_path, "--output", pkl_path,
+         "--progress-param", "UsbGpuSetupProgress", "--progress-base", "0.5", "--progress-span", "0.48"]
   compile_env = {**os.environ, "DEV": "USB+AMD:LLVM", "FLOAT16": "1",
-                 "JIT_BATCH_SIZE": "0", "GMMU": "0"}
+                 "JIT_BATCH_SIZE": "0", "GMMU": "0", "TC_OPT": "2"}
   proc = subprocess.run(cmd, timeout=COMPILE_TIMEOUT_S, capture_output=True, text=True,
                         env=compile_env, preexec_fn=lambda: os.nice(20))
   if proc.returncode != 0:
@@ -99,12 +101,43 @@ def _compile_in_subprocess(meta: dict, onnx_path: str, pkl_path: str) -> None:
     raise RuntimeError(f"eGPU model compile failed (rc={proc.returncode}): {tail}")
 
 
+_precompiled_tried = False
+
+
 def _ensure_artifact(params: Params, meta: dict) -> str:
-  pkl_path = egpu_pkl_path(meta)
-  if os.path.isfile(pkl_path):
-    return pkl_path
+  global _precompiled_tried
+  policy_path = egpu_policy_pkl_path(meta)
+  if os.path.isfile(policy_path):
+    return policy_path
+  legacy_path = egpu_pkl_path(meta)
 
   params.put_bool("UsbGpuCompiled", False)
+  params.put_bool("UsbGpuReady", False)
+
+  if meta.get("egpu_policy_artifact") and not _precompiled_tried:
+    _precompiled_tried = True
+    params.put("UsbGpuSetupProgress", "0.0")
+    dl_last = [-1.0]
+
+    def _dl_prog(p: float) -> None:
+      if p - dl_last[0] >= 0.02 or p >= 1.0:
+        dl_last[0] = p
+        params.put("UsbGpuSetupProgress", f"{p:.3f}")
+
+    try:
+      cloudlog.warning(f"iqegpumodeld downloading precompiled {meta['key']} policy "
+                       f"({int(meta['egpu_policy_artifact'].get('size', 0)) / 1e6:.0f}MB)")
+      precompiled = download_precompiled(meta, progress_cb=_dl_prog, policy=True)
+      if precompiled is not None:
+        cloudlog.warning(f"iqegpumodeld precompiled ready -> {precompiled}")
+        return precompiled
+    except Exception as e:
+      cloudlog.warning(f"iqegpumodeld precompiled policy unavailable ({e}); falling back")
+
+  if os.path.isfile(legacy_path):
+    cloudlog.warning(f"iqegpumodeld using legacy per-tensor artifact {legacy_path}; policy artifact not hosted yet")
+    return legacy_path
+
   onnx_path = local_onnx(meta)
   if onnx_path is None:
     params.put("UsbGpuSetupProgress", "0.0")
@@ -114,14 +147,14 @@ def _ensure_artifact(params: Params, meta: dict) -> str:
     def _prog(p: float) -> None:
       if p - last[0] >= 0.02 or p >= 1.0:
         last[0] = p
-        params.put("UsbGpuSetupProgress", f"{p:.3f}")
+        params.put("UsbGpuSetupProgress", f"{p * 0.5:.3f}")
 
     onnx_path = download_onnx(meta, progress_cb=_prog)
 
   cloudlog.warning(f"iqegpumodeld compiling {meta['key']} for USB-AMD (one-time, can take minutes)")
-  _compile_in_subprocess(meta, onnx_path, pkl_path)
-  cloudlog.warning(f"iqegpumodeld compiled -> {pkl_path}")
-  return pkl_path
+  _compile_in_subprocess(meta, onnx_path, policy_path)
+  cloudlog.warning(f"iqegpumodeld compiled -> {policy_path}")
+  return policy_path
 
 
 def _load_infer_fn(pkl_path: str, meta: dict):
@@ -136,6 +169,10 @@ def _load_infer_fn(pkl_path: str, meta: dict):
   if int(bundle.get("output_len", -1)) != int(meta["output_len"]):
     quarantine_artifact(pkl_path, "pkl output_len mismatch")
     raise RuntimeError(f"artifact output_len {bundle.get('output_len')} != {meta['output_len']}")
+  if bundle.get("format") == POLICY_FORMAT:
+    runner = PolicyRunner(bundle["run_policy"], bundle["input_spec"], int(bundle["frame_skip"]),
+                          meta["output_slices"]["hidden_state"], bundle.get("input_device", "AMD"))
+    return runner, bundle["input_spec"]
   jit = bundle["run_model"]
   input_dev = bundle.get("input_device", "AMD")
   input_spec = bundle["input_spec"]
@@ -152,7 +189,12 @@ def _load_infer_fn(pkl_path: str, meta: dict):
 def _warmup(infer_fn, input_spec: dict, output_len: int) -> float:
   zeros = {name: np.zeros(shape, dtype=dtype) for name, (shape, dtype) in input_spec.items()}
   t0 = time.perf_counter()
-  out = infer_fn(zeros)
+  if isinstance(infer_fn, PolicyRunner):
+    img = input_spec["img"][0]
+    out = infer_fn.run(np.zeros((2, 6, img[2], img[3]), dtype=np.uint8), np.zeros(input_spec["desire_pulse"][0][2], dtype=np.float32),
+                       np.zeros(2, dtype=np.float32), np.zeros(2, dtype=np.float32))
+  else:
+    out = infer_fn(zeros)
   dt = time.perf_counter() - t0
   if out.shape[0] != output_len or not np.isfinite(out).all():
     raise RuntimeError(f"warmup produced invalid output (len={out.shape[0]})")
@@ -207,6 +249,7 @@ def main(demo: bool = False) -> None:
 
   params.put_bool("UsbGpuLoading", False)
   params.put_bool("UsbGpuCompiled", True)
+  params.put_bool("UsbGpuReady", True)
   params.put("UsbGpuSetupProgress", "1.0")
   cloudlog.warning(f"iqegpumodeld model: {meta['key']} ({meta['model_name']})")
   cloudlog.warning(f"iqegpumodeld model up (warmup {warm_s * 1e3:.0f}ms)")

iqpilot/selfdrive/iqmodeld/model_bundle_downloader.py

@@ -0,0 +1,87 @@
+"""
+Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+
+MODELS_BASE_URLS = (
+  "https://git.konn3kt.com/teal/IQModels/raw/branch/main",
+  "https://gitlvb.teallvbs.xyz/teal/IQModels/raw/branch/main",
+)
+CHUNK = 4 * 1024 * 1024
+HTTP_TIMEOUT_S = 60.0
+STREAM_RETRIES = 6
+
+
+def _requests_auth():
+  import importlib
+  for mod in ("iqpilot_private.models.git_auth", "iqpilot.models_private_src.git_auth",
+              "iqpilot.selfdrive.iqmodeld.models.git_auth"):
+    try:
+      return importlib.import_module(mod).get_requests_auth()
+    except Exception:
+      continue
+  return None
+
+
+def _lfs_endpoint(base_url: str) -> str:
+  return base_url.split("/raw/", 1)[0] + ".git/info/lfs"
+
+
+def _resolve_oid(session, base_url: str, oid: str, size: int, auth):
+  import requests
+  batch = session.post(f"{_lfs_endpoint(base_url)}/objects/batch",
+                       data=json.dumps({"operation": "download", "transfers": ["basic"],
+                                        "objects": [{"oid": oid, "size": size}]}),
+                       headers={"Content-Type": "application/vnd.git-lfs+json",
+                                "Accept": "application/vnd.git-lfs+json"},
+                       auth=auth, timeout=HTTP_TIMEOUT_S)
+  batch.raise_for_status()
+  entry = batch.json()["objects"][0]
+  if "actions" not in entry:
+    raise requests.RequestException(f"LFS object unavailable: {entry.get('error', oid)}")
+  action = entry["actions"]["download"]
+  return action["href"], action.get("header", {})
+
+
+def download_lfs_bundle(objects: list, dst: str, sha256: str, size: int, progress_cb=None) -> str:
+  import requests
+  auth = _requests_auth()
+  session = requests.Session()
+  os.makedirs(os.path.dirname(dst), exist_ok=True)
+  tmp = dst + ".part"
+  total = int(size) or sum(int(o["size"]) for o in objects)
+  last_error: Exception | None = None
+  for base_url in MODELS_BASE_URLS:
+    for attempt in range(STREAM_RETRIES):
+      try:
+        digest = hashlib.sha256()
+        got = 0
+        with open(tmp, "wb") as f:
+          for obj in objects:
+            href, headers = _resolve_oid(session, base_url, obj["oid"], int(obj["size"]), auth)
+            obj_auth = None if headers.get("Authorization") else auth
+            with session.get(href, headers=headers, stream=True, timeout=120, auth=obj_auth) as r:
+              r.raise_for_status()
+              for chunk in r.iter_content(CHUNK):
+                f.write(chunk)
+                digest.update(chunk)
+                got += len(chunk)
+                if progress_cb is not None and total:
+                  progress_cb(min(1.0, got / total))
+        if total and got != total:
+          raise RuntimeError(f"size mismatch: {got}/{total} bytes")
+        if sha256 and digest.hexdigest() != sha256:
+          raise RuntimeError("sha256 mismatch")
+        os.replace(tmp, dst)
+        return dst
+      except Exception as e:
+        last_error = e
+        try:
+          os.remove(tmp)
+        except OSError:
+          pass
+  raise RuntimeError(f"model bundle download failed: {last_error}")

iqpilot/selfdrive/iqmodeld/tests/test_egpu_policy.py

@@ -0,0 +1,82 @@
+"""
+Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
+"""
+import os
+
+import numpy as np
+
+os.environ["DEV"] = "CPU"
+
+from iqpilot.selfdrive.iqmodeld.egpu_policy import PolicyRunner, make_run_policy, packed_layout, queue_shapes
+from iqpilot.selfdrive.iqmodeld.temporal_state import TemporalInputState
+
+SPEC = {
+  "img": ((1, 12, 8, 16), "uint8"),
+  "big_img": ((1, 12, 8, 16), "uint8"),
+  "desire_pulse": ((1, 25, 8), "float32"),
+  "traffic_convention": ((1, 2), "float32"),
+  "action_t": ((1, 2), "float32"),
+  "features_buffer": ((1, 24, 512), "float32"),
+}
+FS = 4
+OUT_LEN = 2580
+HIDDEN = slice(1064, 1576)
+
+
+def _pack(inputs):
+  from tinygrad.tensor import Tensor
+  parts = [inputs[k].cast("float32").reshape(-1) for k in ("img", "big_img", "features_buffer", "desire_pulse", "traffic_convention", "action_t")]
+  flat = Tensor.cat(*parts)
+  hidden = (flat[:512] * 0.001).reshape(1, 512)
+  return flat, hidden
+
+
+def _fake_model(inputs):
+  from tinygrad.tensor import Tensor
+  flat, hidden = _pack(inputs)
+  n = flat.shape[0]
+  head = flat[:min(n, HIDDEN.start)]
+  out = Tensor.cat(head.pad((0, HIDDEN.start - head.shape[0])), hidden.reshape(-1), Tensor.zeros(OUT_LEN - HIDDEN.stop, device="CPU"))
+  return {"outputs": out.reshape(1, -1)}
+
+
+class _Reference:
+  def __init__(self):
+    self.state = TemporalInputState(FS, SPEC)
+
+  def run(self, warped, desire, traffic, action_t):
+    inputs = self.state.push_and_materialize(warped, desire, traffic, action_t)
+    from tinygrad.tensor import Tensor
+    t = {k: Tensor(np.ascontiguousarray(v), device="CPU") for k, v in inputs.items()}
+    out = _fake_model(t)["outputs"].numpy().reshape(-1)
+    self.state.note_hidden_state(out, HIDDEN)
+    return out
+
+
+def test_policy_queues_match_temporal_state():
+  from tinygrad.engine.jit import TinyJit
+  jit = TinyJit(make_run_policy(_fake_model, SPEC, FS, "CPU"), prune=True)
+  runner = PolicyRunner(jit, SPEC, FS, HIDDEN, "CPU")
+  ref = _Reference()
+  rng = np.random.default_rng(3)
+  desire = np.zeros(8, dtype=np.float32)
+  for i in range(14):
+    warped = rng.integers(0, 256, (2, 6, 8, 16), dtype=np.int64).astype(np.uint8)
+    if i in (2, 3, 9):
+      desire[:] = 0
+      desire[1 + (i % 3)] = 1
+    elif i == 5:
+      desire[:] = 0
+    traffic = np.array([1.0, 0.0], dtype=np.float32) if i % 2 else np.array([0.0, 1.0], dtype=np.float32)
+    action_t = np.array([0.1 * i, 0.2], dtype=np.float32)
+    got = runner.run(warped, desire, traffic, action_t)
+    want = ref.run(warped, desire, traffic, action_t)
+    np.testing.assert_array_equal(got, want, err_msg=f"frame {i}")
+
+
+def test_layouts():
+  shapes, sizes = packed_layout(SPEC)
+  assert list(shapes) == ["desire", "traffic_convention", "action_t", "prev_feat"]
+  assert sum(sizes) == 8 + 2 + 2 + 512
+  q = queue_shapes(SPEC, FS)
+  assert q["img_q"][0] == (5, 6, 8, 16) and q["feat_q"][0] == (96, 1, 512) and q["desire_q"][0] == (100, 1, 8)

iqpilot/selfdrive/iqmodeld/tools/compile_egpu_model.py

@@ -4,6 +4,7 @@ Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed
 from __future__ import annotations
 
 import argparse
+import gc
 import os
 import pickle
 import time
@@ -12,6 +13,7 @@ os.environ.setdefault("DEV", "USB+AMD:LLVM")
 os.environ.setdefault("FLOAT16", "1")
 os.environ.setdefault("JIT_BATCH_SIZE", "0")
 os.environ.setdefault("GMMU", "0")
+os.environ.setdefault("TC_OPT", "2")
 
 import numpy as np
 
@@ -24,6 +26,22 @@ INPUT_SPEC = dict(MODEL_INPUT_SPEC)
 patch_tinygrad_fetch_fw()
 
 SEED = 42
+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:
@@ -81,8 +99,27 @@ def compile_model(meta: dict, onnx_path: str, out_path: str) -> str:
   if not np.isfinite(baseline).all():
     raise RuntimeError("compiled model produced non-finite outputs")
 
-  print("pickle round trip")
-  jit = pickle.loads(pickle.dumps(jit))
+  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):
@@ -96,19 +133,94 @@ def compile_model(meta: dict, onnx_path: str, out_path: str) -> str:
   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 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, 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 np.isfinite(baseline).all():
+    raise RuntimeError("compiled policy produced non-finite outputs")
+
   bundle = {
-    "run_model": jit,
+    "format": POLICY_FORMAT,
+    "run_policy": 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,
+    "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")
   with open(tmp, "wb") as f:
     pickle.dump(bundle, f, protocol=pickle.HIGHEST_PROTOCOL)
+
+  del bundle, jit, queues, packed
+  gc.collect()
+
+  print("reload + validate")
+  with open(tmp, "rb") as f:
+    jit = pickle.load(f)["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)
+  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"])))
+
   os.replace(tmp, out_path)
   return out_path
 
@@ -118,6 +230,10 @@ def main() -> None:
   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=2, choices=(1, 2))
   args = p.parse_args()
 
   if args.model is not None:
@@ -136,7 +252,23 @@ def main() -> None:
   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")
 
-  out = compile_model(meta, onnx_path, args.output or egpu_pkl_path(meta))
+  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:
+    build = compile_policy_model if args.format == 2 else compile_model
+    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)")