IQ.Pilot Release Commit @ 7550fa9

2026-09-02 13:32:42 -05:00
parent 4efaff4cb2
commit 43ee82d228
24 changed files with 882 additions and 93 deletions
--- a/iqpilot/selfdrive/iqmodeld/egpu_policy.py
+++ b/iqpilot/selfdrive/iqmodeld/egpu_policy.py
@@ -14,9 +14,11 @@ import tempfile
 import numpy as np

 POLICY_FORMAT = 2
+MODEL_FORMAT = 3
 OOB_MAGIC = b"IQEGPUOOB1"
 QUEUE_NAMES = ("img_q", "big_img_q", "feat_q", "desire_q")
 PACKED_ORDER = ("desire", "traffic_convention", "action_t", "prev_feat")
+MODELD_INPUTS = (*QUEUE_NAMES, "packed_npy_inputs")


 def packed_layout(input_spec: dict) -> tuple[dict[str, tuple[int, ...]], list[int]]:
@@ -56,7 +58,8 @@ class PackedInputs:
    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))
+    parts = np.split(self.array, np.cumsum(self.sizes[:-1]))
+    self.views = {name: part.reshape(shape) for (name, shape), part in zip(self.shapes.items(), parts, strict=True)}
    self.tensor = Tensor(self.array, device="NPY").realize()


@@ -124,6 +127,173 @@ class PolicyRunner:
    return flat


+def nv12_copy_size(stride: int, y_height: int, uv_height: int) -> int:
+  return stride * (y_height + uv_height)
+
+
+def frame_layout(input_spec: dict) -> tuple[dict[str, tuple[int, ...]], list[int], int]:
+  policy_shapes, _ = packed_layout(input_spec)
+  shapes = {"tfm": (3, 3), "big_tfm": (3, 3)} | policy_shapes
+  sizes = [math.prod(s) for s in shapes.values()]
+  return shapes, sizes, sum(sizes) * np.dtype(np.float32).itemsize
+
+
+def model_size(input_spec: dict) -> tuple[int, int]:
+  img = input_spec["img"][0]
+  return img[3] * 2, img[2] * 2
+
+
+def warp_perspective_tinygrad(src_flat, M_inv, dst_shape, src_shape, stride_pad, border_fill_val=None):
+  from tinygrad.tensor import Tensor
+  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):
+  from tinygrad.tensor import Tensor
+  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: tuple[int, int, int, int, int], model_w: int, model_h: int, device: str):
+  from tinygrad.helpers import Context
+  from tinygrad.tensor import Tensor
+  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=device)
+    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))
+    return frames_to_tensor(yuv)
+  return frame_prepare_tinygrad
+
+
+def make_warp(nv12: tuple[int, int, int, int, int], model_w: int, model_h: int, device: str):
+  from tinygrad.tensor import Tensor
+  frame_prepare = make_frame_prepare(nv12, model_w, model_h, device)
+
+  def warp(tfm, big_tfm, frame, big_frame):
+    tfm = tfm.to(device)
+    big_tfm = big_tfm.to(device)
+    frame = frame.to(device)
+    big_frame = big_frame.to(device)
+    Tensor.realize(tfm, big_tfm, frame, big_frame)
+
+    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_model(warp, run_policy, input_spec: dict, frame_copy_size: int, device: str):
+  from tinygrad.tensor import Tensor
+  _, policy_sizes = packed_layout(input_spec)
+  _, _, packed_npy_size = frame_layout(input_spec)
+
+  def run_model(img_q, big_img_q, feat_q, desire_q, packed_npy_inputs):
+    packed_input = packed_npy_inputs.to(device)
+    Tensor.realize(packed_input)
+    packed_npy_inputs = packed_input[:packed_npy_size].bitcast("float32")
+    frame = packed_input[packed_npy_size:packed_npy_size + frame_copy_size]
+    big_frame = packed_input[packed_npy_size + frame_copy_size:]
+    tfm, big_tfm, policy_inputs = packed_npy_inputs.split([9, 9, sum(policy_sizes)])
+    warped = warp(tfm.reshape(3, 3), big_tfm.reshape(3, 3), frame, big_frame)
+    return run_policy(warped, img_q, big_img_q, feat_q, desire_q, policy_inputs)
+
+  return run_model
+
+
+class PackedFrames:
+  def __init__(self, input_spec: dict, frame_copy_size: int):
+    from tinygrad.tensor import Tensor
+    self.shapes, self.sizes, npy_bytes = frame_layout(input_spec)
+    self.frame_copy_size = frame_copy_size
+    self.array = np.zeros(npy_bytes + 2 * frame_copy_size, dtype=np.uint8)
+    npy = self.array[:npy_bytes].view(np.float32)
+    self.views = dict(zip(self.shapes, [v.reshape(s) for s, v in zip(self.shapes.values(), np.split(npy, np.cumsum(self.sizes[:-1])), strict=True)],
+                          strict=True))
+    frames = self.array[npy_bytes:]
+    self.frames = {"img": frames[:frame_copy_size], "big_img": frames[frame_copy_size:]}
+    self.tensor = Tensor(self.array, device="NPY").realize()
+
+
+def make_model_queues(input_spec: dict, frame_skip: int, device: str, frame_copy_size: int) -> tuple[dict, PackedFrames]:
+  packed = PackedFrames(input_spec, frame_copy_size)
+  return {**make_queues(input_spec, frame_skip, device), "packed_npy_inputs": packed.tensor}, packed
+
+
+class ModelRunner:
+  def __init__(self, jit, input_spec: dict, frame_skip: int, hidden_slice: slice, device: str, frame_copy_size: int):
+    self._jit = jit
+    self._queues, self._packed = make_model_queues(input_spec, frame_skip, device, frame_copy_size)
+    self._hidden = hidden_slice
+    self._prev_desire = np.zeros(input_spec["desire_pulse"][0][2], dtype=np.float32)
+    self.frame_copy_size = frame_copy_size
+
+  def run(self, main_frame, extra_frame, tfm: np.ndarray, big_tfm: np.ndarray, desire_pulse: np.ndarray,
+          traffic_convention: np.ndarray, action_t: np.ndarray) -> np.ndarray:
+    n = self.frame_copy_size
+    v = self._packed.views
+    f = self._packed.frames
+    np.copyto(f["img"], np.frombuffer(main_frame, dtype=np.uint8, count=n))
+    np.copyto(f["big_img"], np.frombuffer(extra_frame, dtype=np.uint8, count=n))
+    v["tfm"][:, :] = tfm
+    v["big_tfm"][:, :] = big_tfm
+    cur = desire_pulse.astype(np.float32, copy=False)
+    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)
+    out, = self._jit(**self._queues)
+    flat = out.numpy().reshape(-1)
+    v["prev_feat"][:] = flat[self._hidden].reshape(v["prev_feat"].shape)
+    return flat
+
+
 def dump_oob(obj, f) -> None:
  # Out-of-band pickle buffers keep the host peak at one tensor while the weights stream to the
  # dock; a plain pickle keeps every weight referenced in the memo until load() returns (~1.7GB).