IQ.Pilot Release Commit @ 0798119

iqpilot/selfdrive/iqmodeld/tools/compile_daemon.py

@@ -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 openpilot.iqpilot.selfdrive.iqmodeld.tools.daemon_jit_compiler import main
+
+
+if __name__ == "__main__":
+  raise SystemExit(main())

iqpilot/selfdrive/iqmodeld/tools/compile_split_runtime.py

@@ -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 openpilot.common.file_chunker import read_file_chunked
+  from openpilot.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, device=WARP_DEVICE).reshape(1, dst_w).expand(dst_h, dst_w).reshape(-1)
+  y_coords = Tensor.arange(dst_h, device=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 openpilot.iqpilot.selfdrive.iqmodeld.metadata import build_metadata_record
+  from openpilot.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())

iqpilot/selfdrive/iqmodeld/tools/compile_supercombo.py

@@ -0,0 +1,427 @@
+#!/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)
+
+  # inline 3x3 matmul as elementwise to avoid reduce op (enables fusion with gather)
+  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):
+    # UV_SCALE @ M_inv @ UV_SCALE_INV simplifies to elementwise scaling
+    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)
+    # deinterleave NV12 UV plane (UVUV... -> separate U, V)
+    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)
+  # TODO prev_feat shouldn't exist and be handled inside the JIT, but corrupt on QCOM for now
+  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)
+  # views into the packed inputs, to be refilled at runtime
+  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 openpilot.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 openpilot.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 openpilot.common.file_chunker import read_file_chunked
+  from openpilot.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 openpilot.system.camerad.cameras.nv12_info import get_nv12_info
+  from openpilot.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)")

iqpilot/selfdrive/iqmodeld/tools/daemon_jit_compiler.py

@@ -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 openpilot.common.file_chunker import read_file_chunked
+  from openpilot.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 openpilot.iqpilot.selfdrive.iqmodeld.metadata import build_metadata_record
+  from openpilot.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

iqpilot/selfdrive/iqmodeld/tools/install_models_pc.py

@@ -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 openpilot.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])