modeld_v2: adapt to state pairs in onnx

openpilot/selfdrive/modeld/compile_modeld.py

@@ -37,7 +37,6 @@ from tinygrad.engine.jit import TinyJit
 
 
 NV12Frame = namedtuple("NV12Frame", ['width', 'height', 'stride', 'y_height', 'uv_height', 'size'])
-MODELD_INPUTS = ['img_q', 'big_img_q', 'feat_q', 'desire_q', 'packed_npy_inputs']
 
 
 def nv12_copy_size(stride: int, y_height: int, uv_height: int) -> int:
@@ -113,58 +112,26 @@ def make_frame_prepare(nv12: NV12Frame, model_w, model_h):
   return frame_prepare_tinygrad
 
 
-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, T-1, ...) e.g. (1, 24, 32, 512) with spatial features
-  feat_dim = math.prod(fb[2:])
-  # 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], feat_dim)}
+def get_npy_shapes(input_shapes, state_pairs):
+  shapes = {'tfm': (3, 3), 'big_tfm': (3, 3)} | {
+    name: shape for name, (shape, _) in input_shapes.items() if name not in state_pairs and name != 'new_img'}
   return shapes, [math.prod(s) for s in shapes.values()]
 
 
-def make_input_queues(input_shapes, frame_skip, device, frame_copy_size):
-  img = input_shapes['img']  # (1, 12, 128, 256)
-  fb = input_shapes['features_buffer']  # (1, T-1, ...), past features only; the model appends the current frame's feature
-  feat_dim = math.prod(fb[2:])
-  dp = input_shapes['desire_pulse']  # (1, 25, 8)
-  n_frames = img[1] // 6
-  img_buf_shape = (frame_skip * (n_frames - 1) + 1, 6, img[2], img[3])
-
-  policy_shapes, _ = get_policy_npy_shapes(input_shapes)
-  shapes = {'tfm': (3, 3), 'big_tfm': (3, 3)} | policy_shapes
-  sizes = [math.prod(s) for s in shapes.values()]
+def make_input_queues(input_shapes, state_pairs, device, frame_copy_size):
+  shapes, sizes = get_npy_shapes(input_shapes, state_pairs)
   packed_npy_size = sum(sizes) * np.dtype(np.float32).itemsize
   packed_input = np.zeros(packed_npy_size + 2 * frame_copy_size, dtype=np.uint8)
   packed_npy_inputs = packed_input[:packed_npy_size].view(np.float32)
   frames = packed_input[packed_npy_size:]
   frame_views = {'img': frames[:frame_copy_size], 'big_img': frames[frame_copy_size:]}
-  # views into the packed inputs, to be refilled at runtime
   npy = {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 = {
-    '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(),
-    'feat_q': Tensor(np.zeros((frame_skip * fb[1], fb[0], feat_dim), 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_input, device='NPY').realize(),
-  }
+  input_queues = {name: Tensor(np.zeros(shape, dtype=dtype.fmt), device=device).realize()
+                  for name, (shape, dtype) in input_shapes.items() if name in state_pairs}
+  input_queues['packed_npy_inputs'] = Tensor(packed_input, device='NPY').realize()
   return input_queues, npy, frame_views
 
 
-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_prepare = make_frame_prepare(nv12, model_w, model_h)
 
@@ -182,54 +149,27 @@ def make_warp(nv12, model_w, model_h):
   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'])
-  model_input_dtypes = {name: spec.dtype for name, spec in model_runner.graph_inputs.items()}
+def make_run_model(warp, model_runner, input_shapes, state_pairs, frame_copy_size):
+  shapes, sizes = get_npy_shapes(input_shapes, state_pairs)
+  packed_npy_size = sum(sizes) * np.dtype(np.float32).itemsize
 
-  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)
-    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.reshape(model_metadata['input_shapes']['features_buffer']),
-      'desire_pulse': desire_buf,
-      'traffic_convention': traffic_convention,
-      'action_t': action_t,
-    }
-    inputs = {name: value.cast(model_input_dtypes[name]) for name, value in inputs.items()}
-    out = next(iter(model_runner(inputs).values())).cast('float32')
-    return out,
-  return run_policy
-
-
-def make_run_model(warp, run_policy, model_metadata, frame_copy_size):
-  _, policy_sizes = get_policy_npy_shapes(model_metadata['input_shapes'])
-  packed_npy_size = (18 + sum(policy_sizes)) * np.dtype(np.float32).itemsize
-
-  def run_model(img_q, big_img_q, feat_q, desire_q, packed_npy_inputs):
-    packed_input = packed_npy_inputs.to(Device.DEFAULT)
-    Tensor.realize(packed_input)
+  def run_model(packed_npy_inputs, **state_inputs):
+    packed_input = packed_npy_inputs.to(Device.DEFAULT).realize()
     packed_npy_inputs = packed_input[:packed_npy_size].bitcast('float32')
+    inputs = {name: t.reshape(s) for (name, s), t in zip(shapes.items(), packed_npy_inputs.split(sizes), strict=True)}
     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)
+    inputs['new_img'] = warp(inputs.pop('tfm'), inputs.pop('big_tfm'), frame, big_frame)
+    inputs = {name: value.cast(input_shapes[name][1]) for name, value in inputs.items()}
+    outputs = {name: value.contiguous() for name, value in model_runner(inputs | state_inputs).items()}
+    Tensor.realize(*outputs.values())
+    if state_pairs:
+      Tensor.realize(*(state_inputs[name].assign(outputs[next_name]) for name, next_name in state_pairs.items()))
+    return tuple(value for name, value in outputs.items() if name not in state_pairs.values())
   return run_model
 
 
-def compile_jit(jit, input_keys, make_queues, benchmark_runs):
+def compile_jit(jit, make_queues, benchmark_runs):
   if benchmark_runs < 1:
     raise ValueError("benchmark_runs must be at least 1")
 
@@ -245,7 +185,7 @@ def compile_jit(jit, input_keys, make_queues, benchmark_runs):
         v[:] = rng.integers(0, 256, size=v.shape, dtype=np.uint8)
       Device.default.synchronize()
       st = time.perf_counter()
-      outs = fn(**{k: input_queues[k] for k in input_keys})
+      outs = fn(**input_queues)
       mt = time.perf_counter()
       Device.default.synchronize()
       et = time.perf_counter()
@@ -300,7 +240,6 @@ if __name__ == "__main__":
                  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('--benchmark-runs', type=int, default=1,
                  help='timed loaded-JIT runs for each correctness seed')
   args = p.parse_args()
@@ -309,24 +248,28 @@ if __name__ == "__main__":
   model_w, model_h = args.model_size
 
   model_runner = OnnxRunner(model_path)
+  input_shapes = {name: (spec.shape, spec.dtype) for name, spec in model_runner.graph_inputs.items()}
+  state_pairs = {name: f'next_{name}' for name in input_shapes if f'next_{name}' in model_runner.graph_outputs}
   out = {
     'metadata': make_metadata_dict(model_path),
+    'input_shapes': input_shapes,
+    'state_pairs': state_pairs,
     'input_devices': {'model': Device.DEFAULT},
     'run_model': {},
   }
 
-  run_policy = make_run_policy(model_runner, out['metadata'], args.frame_skip)
-
   for cam_w, cam_h in args.camera_resolutions:
     nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
     frame_copy_size = nv12_copy_size(nv12.stride, nv12.y_height, nv12.uv_height)
-    make_model_queues = partial(make_input_queues, out['metadata']['input_shapes'], args.frame_skip,
+    make_model_queues = partial(make_input_queues, input_shapes, state_pairs,
                                 frame_copy_size=frame_copy_size)
     warp = make_warp(nv12, model_w, model_h)
-    run_model_jit = TinyJit(make_run_model(warp, run_policy, out['metadata'], frame_copy_size), prune=True)
-    out['run_model'][(cam_w,cam_h)] = compile_jit(run_model_jit, MODELD_INPUTS, make_model_queues,
-                                                  args.benchmark_runs)
+    run_model_jit = TinyJit(make_run_model(warp, model_runner, input_shapes, state_pairs, frame_copy_size), prune=True)
+    out['run_model'][(cam_w,cam_h)] = compile_jit(run_model_jit, make_model_queues, args.benchmark_runs)
 
   with open(args.output, "wb") as f:
     dump_oob(out, f)
+  with open(args.output, "rb") as f:
+    load_oob(f)
+    assert not f.read(1), "unexpected model buffer data"
   print(f"Saved JITs to {args.output} ({os.path.getsize(args.output) / 1e6:.2f} MB)")

openpilot/selfdrive/modeld/modeld.py

@@ -5,8 +5,6 @@ from functools import cached_property
 import os
 os.environ['GMMU'] = '0' # for chestnut fast loading, noop for qcom
 from tinygrad.device import Device
-import usb1
-import struct
 import threading
 import time
 import numpy as np
@@ -28,12 +26,11 @@ from openpilot.common.transformations.model import get_warp_matrix
 from openpilot.selfdrive.controls.lib.desire_helper import DesireHelper
 from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_from_plan, should_stop, smooth_value, get_curvature_from_plan
 from openpilot.selfdrive.modeld.parse_model_outputs import Parser
-from openpilot.selfdrive.modeld.compile_modeld import make_input_queues, nv12_copy_size, MODELD_INPUTS
+from openpilot.selfdrive.modeld.compile_modeld import make_input_queues, nv12_copy_size
 from openpilot.selfdrive.modeld.fill_model_msg import fill_model_msg, fill_driving_model_data, fill_pose_msg, PublishState
 from openpilot.common.file_chunker import open_file_chunked
-from openpilot.common.hardware.usb import CHESTNUT_USB_IDS
 from openpilot.selfdrive.modeld.constants import ModelConstants, Plan
-from openpilot.selfdrive.modeld.helpers import chestnut_present, chestnut_compiled, chestnut_ready, modeld_pkl_path, load_oob
+from openpilot.selfdrive.modeld.helpers import chestnut_present, chestnut_compiled, modeld_pkl_path, load_oob
 
 from openpilot.sunnypilot.livedelay.helpers import get_lat_delay
 from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase
@@ -75,45 +72,14 @@ def get_action_from_model(model_output: dict[str, np.ndarray], prev_action: log.
                                 shouldStop=bool(stop))
 
 
-class ChestnutState:
-  # only modeld can access chestnut
+class ChestnutGpuState:
+  # GPU metrics require modeld's GPU context
   def __init__(self, pm: PubMaster, big: bool):
     self.pm = pm
     self.big = big
     self.valid = True
     self.sends = 0
     self.metrics = {}
-    self._asm_usb = None
-
-  def _close_asm_usb(self) -> None:
-    if self._asm_usb is not None:
-      self._asm_usb.close()
-      self._asm_usb = None
-
-  def _open_asm_usb(self):
-    context = usb1.USBContext()
-    for vendor_id, product_id in CHESTNUT_USB_IDS:
-      if (handle := context.openByVendorIDAndProductID(vendor_id, product_id, skip_on_error=True)) is not None:
-        return handle
-    context.close()
-
-  def _read_ina(self) -> tuple[int, int, bool]:
-    if "AMD" in Device._opened_devices and self._asm_usb is None:
-      try:
-        raw = Device["AMD"].iface.pci_dev.usb.usb.control_read(0xC0, 5)
-        return struct.unpack('<Hh?', bytes(raw))
-      except Exception:
-        pass
-    if self._asm_usb is None:
-      self._asm_usb = self._open_asm_usb()
-    if self._asm_usb is None:
-      raise usb1.USBErrorNoDevice
-    try:
-      raw = self._asm_usb.controlRead(0xC0, 0xC0, 0, 0, 5, timeout=100)
-    except usb1.USBError:
-      self._close_asm_usb()
-      raise
-    return struct.unpack('<Hh?', bytes(raw))
 
   @cached_property
   def power_limit(self) -> int:
@@ -121,8 +87,8 @@ class ChestnutState:
     return smu._send_msg(smu.smu_mod.PPSMC_MSG_GetPptLimit, 0, read_back_arg=True, timeout=100)
 
   def send(self) -> None:
-    msg = messaging.new_message('chestnutState')
-    state = msg.chestnutState
+    msg = messaging.new_message('chestnutGpuState')
+    state = msg.chestnutGpuState
     self.sends += 1
     if self.big and "AMD" in Device._opened_devices and self.sends % 100 == 1:
       try:
@@ -148,21 +114,8 @@ class ChestnutState:
       for k, v in self.metrics.items():
         setattr(state, k, v)
 
-    asm_valid = False
-    try:
-      # ASM runs on USB-C power, these still read without a gpu
-      state.supplyVoltage, state.supplyCurrent, state.supplyFault = self._read_ina()
-      asm_valid = True
-    except Exception:
-      pass
-    if "AMD" in Device._opened_devices:
-      try:
-        state.pcieLtssm = Device["AMD"].iface.pci_dev.usb.read(0xB450, 1)[0]
-      except Exception:
-        pass
-
-    msg.valid = asm_valid and (not self.big or self.valid)
-    self.pm.send('chestnutState', msg)
+    msg.valid = not self.big or (self.valid and bool(self.metrics))
+    self.pm.send('chestnutGpuState', msg)
 
 
 class FrameMeta:
@@ -184,17 +137,17 @@ class ModelState(ModelStateBase):
     input_devices = jits['input_devices']
     self.model_device = input_devices['model']
     metadata = jits['metadata']
-    self.input_shapes = metadata['input_shapes']
-    self.vision_input_names = [k for k in self.input_shapes if 'img' in k]
+    self.input_shapes = jits['input_shapes']
+    self.state_pairs = jits['state_pairs']
+    self.vision_input_names = ('img', 'big_img')
     self.output_slices = metadata['output_slices']
 
     self.prev_desire = np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32)
     self.chestnut = chestnut
 
-    self.frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ
     self.frame_copy_size = nv12_copy_size(*get_nv12_info(cam_w, cam_h)[:3])
     self.input_queues, self.npy, self.frame_views = make_input_queues(
-      self.input_shapes, self.frame_skip, device=self.model_device, frame_copy_size=self.frame_copy_size)
+      self.input_shapes, self.state_pairs, device=self.model_device, frame_copy_size=self.frame_copy_size)
     self.parser = Parser()
     self.run_model = jits['run_model'][(cam_w,cam_h)]
 
@@ -216,14 +169,13 @@ class ModelState(ModelStateBase):
     self.npy['tfm'][:,:] = transforms['img'][:,:]
     self.npy['big_tfm'][:,:] = transforms['big_img'][:,:]
 
-    outs, = self.run_model(**{k: self.input_queues[k] for k in MODELD_INPUTS})
+    outs, = self.run_model(**self.input_queues)
     if after_enqueue is not None:
       after_enqueue()
     model_output = outs.numpy()[0]
     if self.chestnut and not np.all(np.isfinite(model_output)):
       raise RuntimeError("model output not finite")
     outputs_dict = self.parser.parse_outputs(self.slice_outputs(model_output, self.output_slices))
-    self.npy['prev_feat'][:] = model_output[self.output_slices['hidden_state']]
 
     if SEND_RAW_PRED:
       outputs_dict['raw_pred'] = model_output.copy()
@@ -235,32 +187,19 @@ class ModelState(ModelStateBase):
     dims = {'desire_pulse': ModelConstants.DESIRE_LEN, 'traffic_convention': 2, 'action_t': 2}
     self.run(dummy_frames, dict.fromkeys(self.vision_input_names, eye), {k: np.zeros(v, dtype=np.float32) for k, v in dims.items()})
     self.input_queues, self.npy, self.frame_views = make_input_queues(
-      self.input_shapes, self.frame_skip, device=self.model_device, frame_copy_size=self.frame_copy_size)
+      self.input_shapes, self.state_pairs, device=self.model_device, frame_copy_size=self.frame_copy_size)
     self.prev_desire[:] = 0
 
 
 def main(demo=False):
   cloudlog.warning("modeld init")
 
-  chestnut_available = chestnut_present() and chestnut_compiled()
-  CHESTNUT = False
-  if chestnut_available:
-    poller = messaging.Poller()
-    sock = messaging.sub_sock("chestnutState", poller=poller, conflate=True)
-    deadline = time.monotonic() + 4. / SERVICE_LIST['deviceState'].frequency
-    while not CHESTNUT and (remaining := deadline - time.monotonic()) > 0.:
-      if not poller.poll(round(remaining * 1000)):
-        break
-      msg = messaging.recv_one_or_none(sock)
-      CHESTNUT = msg is not None and msg.valid and chestnut_ready(msg.chestnutState)
+  CHESTNUT = chestnut_present() and chestnut_compiled()
   if CHESTNUT:
     os.environ['HCQDEV_WAIT_TIMEOUT_MS'] = '3000'
   params = Params()
   params.put_bool("ChestnutLoading", CHESTNUT)
-  if chestnut_available and not CHESTNUT:
-    params.put_bool("ChestnutActive", False)
-  else:
-    params.remove("ChestnutActive")
+  params.remove("ChestnutActive")
 
   config_realtime_process(7, 54)
 
@@ -304,11 +243,7 @@ def main(demo=False):
     loader.start()
     loader.join(BIG_MODEL_TIMEOUT)
     model = big_model
-    if model is None:
-      params.put_bool("ChestnutModelError", True)
     params.put_bool("ChestnutActive", model is not None)
-    if model is not None:
-      params.remove("ChestnutModelError")
 
   small_model = ModelState(vipc_client_main.width, vipc_client_main.height, False) if model is None or CHESTNUT else None
   if model is None:
@@ -318,13 +253,13 @@ def main(demo=False):
   cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting")
 
   # messaging
-  pub_socks = ["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"] + (["chestnutState"] if CHESTNUT else [])
+  pub_socks = ["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"] + (["chestnutGpuState"] if CHESTNUT else [])
   pm = PubMaster(pub_socks)
   sm = SubMaster(["deviceState", "carState", "narrowRoadCameraState", "extrinsicsCalibration", "driverMonitoringState", "carControl", "lateralDelay"])
 
   publish_state = PublishState()
   params = Params()
-  chestnut_state = ChestnutState(pm, model.chestnut) if CHESTNUT else None
+  chestnut_state = ChestnutGpuState(pm, model.chestnut) if CHESTNUT else None
 
   # setup filter to track dropped frames
   frame_dropped_filter = FirstOrderFilter(0., 10., 1. / ModelConstants.MODEL_RUN_FREQ)
@@ -435,14 +370,13 @@ def main(demo=False):
     mt1 = time.perf_counter()
     try:
       send_chestnut = (chestnut_state is not None and
-                       run_count % round(ModelConstants.MODEL_RUN_FREQ / SERVICE_LIST['chestnutState'].frequency) == 0)
+                       run_count % round(ModelConstants.MODEL_RUN_FREQ / SERVICE_LIST['chestnutGpuState'].frequency) == 0)
       model_output = model.run(bufs, transforms, inputs, chestnut_state.send if send_chestnut else None)
     except Exception:
       if not params.get_bool("ChestnutActive"):
         raise
       # fallback to small model
       cloudlog.exception("big model failed, fall back to small")
-      params.put_bool("ChestnutModelError", True)
       params.put_bool("ChestnutActive", False)
       assert small_model is not None
       model = small_model

openpilot/sunnypilot/modeld_v2/compile_modeld.py

@@ -33,6 +33,7 @@ def _patch_tinygrad_fetch_fw():
 _patch_tinygrad_fetch_fw()
 
 import openpilot.selfdrive.modeld.compile_modeld as stock
+import openpilot.sunnypilot.modeld_v2.stock_dependencies as legacy
 from tinygrad import dtypes
 from tinygrad.device import Device
 from tinygrad.engine.jit import TinyJit
@@ -41,7 +42,7 @@ from tinygrad.tensor import Tensor
 MODEL_TYPES = ('vision_policy', 'supercombo', 'vision_multi_policy')
 WARP_INPUTS = ['tfm', 'big_tfm']
 POLICY_INPUTS = ['img_q', 'big_img_q', 'feat_q', 'desire_q', 'packed_npy_inputs']
-nv12_copy_size = stock.nv12_copy_size
+
 
 def _detect_desire_key(shapes: dict) -> str | None:
   return next((key for key in shapes if key.startswith('desire')), None)
@@ -152,8 +153,8 @@ def make_warp_queues(device=Device.DEFAULT):
 
 
 def make_run_policy(vision_runner, policy_runners: list, features_slice: slice, frame_skip: int, input_shapes: dict):
-  sample_skip_fn = partial(stock.sample_skip, frame_skip=frame_skip)
-  sample_desire_fn = partial(stock.sample_desire, frame_skip=frame_skip)
+  sample_skip_fn = partial(legacy.sample_skip, frame_skip=frame_skip)
+  sample_desire_fn = partial(legacy.sample_desire, frame_skip=frame_skip)
 
   desire_key = _detect_desire_key(input_shapes)
   road_key, wide_key = _detect_vision_keys(input_shapes)
@@ -170,14 +171,14 @@ def make_run_policy(vision_runner, policy_runners: list, features_slice: slice,
     warped_dev = warped.to(Device.DEFAULT)
     Tensor.realize(packed_npy_inputs_dev, warped_dev)
 
-    img = stock.shift_and_sample(img_q, warped_dev[0:1], sample_skip_fn)
-    big_img = stock.shift_and_sample(big_img_q, warped_dev[1:2], sample_skip_fn)
+    img = legacy.shift_and_sample(img_q, warped_dev[0:1], sample_skip_fn)
+    big_img = legacy.shift_and_sample(big_img_q, warped_dev[1:2], sample_skip_fn)
 
     unpacked_tensors = [tensor.reshape(shape) for tensor, shape in zip(packed_npy_inputs_dev.split(npy_sizes), npy_shapes.values(), strict=True)]
     unpacked_dict = dict(zip(npy_shapes.keys(), unpacked_tensors, strict=True))
 
     desire_dev = unpacked_dict['desire']
-    desire_buf = stock.shift_and_sample(desire_q, desire_dev.reshape(1, 1, -1), sample_desire_fn)
+    desire_buf = legacy.shift_and_sample(desire_q, desire_dev.reshape(1, 1, -1), sample_desire_fn)
 
     inputs = {desire_key: desire_buf}
     for key, tensor_val in unpacked_dict.items():
@@ -186,13 +187,13 @@ def make_run_policy(vision_runner, policy_runners: list, features_slice: slice,
 
     if 'prev_feat' in unpacked_dict:
       prev_feat_dev = unpacked_dict['prev_feat']
-      inputs['features_buffer'] = stock.shift_and_sample(feat_q, prev_feat_dev.reshape(1, 1, -1), sample_skip_fn).reshape(input_shapes['features_buffer'])
+      inputs['features_buffer'] = legacy.shift_and_sample(feat_q, prev_feat_dev.reshape(1, 1, -1), sample_skip_fn).reshape(input_shapes['features_buffer'])
 
     if vision_runner:
       vision_out_cast = next(iter(vision_runner({road_key: img, wide_key: big_img}).values())).cast('float32').realize()
       if 'features_buffer' not in inputs:
         new_feat = vision_out_cast[:, features_slice].reshape(1, -1).unsqueeze(0)
-        inputs['features_buffer'] = stock.shift_and_sample(feat_q, new_feat, sample_skip_fn).realize()
+        inputs['features_buffer'] = legacy.shift_and_sample(feat_q, new_feat, sample_skip_fn).realize()
       policy_outs = [next(iter(pol_runner(inputs).values())).cast('float32').realize() for pol_runner in policy_runners]
       return (vision_out_cast, *policy_outs) if len(policy_outs) > 1 else (vision_out_cast, policy_outs[0])
 
@@ -203,7 +204,7 @@ def make_run_policy(vision_runner, policy_runners: list, features_slice: slice,
     policy_out = next(iter(policy_runners[0](inputs).values())).cast('float32').realize()
     if 'features_buffer' not in inputs and features_slice is not None:
       new_feat = policy_out[:, features_slice].reshape(1, -1).unsqueeze(0)
-      stock.shift_and_sample(feat_q, new_feat, sample_skip_fn).realize()
+      legacy.shift_and_sample(feat_q, new_feat, sample_skip_fn).realize()
     return policy_out
 
   return run_policy
@@ -330,16 +331,32 @@ if __name__ == "__main__":
     output_data['run_model'] = {}
     derived_frame_skip = args.frame_skip or derive_frame_skip({}, model_metadata['input_shapes'])
     model_runner = OnnxRunner(args.supercombo_onnx)
-    run_policy = stock.make_run_policy(model_runner, model_metadata, derived_frame_skip)
-    for cam_w, cam_h in args.camera_resolutions:
-      print(f"Compiling unified run_model JIT for {cam_w}x{cam_h}...")
-      nv12 = stock.NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
-      frame_copy_size = stock.nv12_copy_size(nv12.stride, nv12.y_height, nv12.uv_height)
-      make_model_queues = partial(stock.make_input_queues, model_metadata['input_shapes'], derived_frame_skip,
-                                  frame_copy_size=frame_copy_size)
-      warp = stock.make_warp(nv12, model_w, model_h)
-      run_model_jit = TinyJit(stock.make_run_model(warp, run_policy, model_metadata, frame_copy_size), prune=True)
-      output_data['run_model'][(cam_w, cam_h)] = compile_jit(run_model_jit, stock.MODELD_INPUTS, make_model_queues, benchmark_runs=args.benchmark_runs)
+    new_img_model = 'new_img' in model_runner.graph_inputs
+
+    if new_img_model:
+      input_shapes = {name: (spec.shape, spec.dtype) for name, spec in model_runner.graph_inputs.items()}
+      state_pairs = {name: f'next_{name}' for name in input_shapes if f'next_{name}' in model_runner.graph_outputs}
+      output_data['metadata'] = {'model': model_metadata, **model_metadata, 'input_shapes': input_shapes, 'state_pairs': state_pairs}
+      for cam_w, cam_h in args.camera_resolutions:
+        print(f"Compiling unified run_model JIT for {cam_w}x{cam_h} (new architecture)...")
+        nv12 = stock.NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
+        frame_copy_size = stock.nv12_copy_size(nv12.stride, nv12.y_height, nv12.uv_height)
+        make_model_queues = partial(stock.make_input_queues, input_shapes, state_pairs, frame_copy_size=frame_copy_size)
+        warp = stock.make_warp(nv12, model_w, model_h)
+        run_model_jit = TinyJit(stock.make_run_model(warp, model_runner, input_shapes, state_pairs, frame_copy_size), prune=True)
+        output_data['run_model'][(cam_w, cam_h)] = compile_jit(run_model_jit, list(state_pairs.keys()) + ['packed_npy_inputs'], make_model_queues,
+                                                               benchmark_runs=args.benchmark_runs)
+    else:
+      run_policy = legacy.make_legacy_run_policy(model_runner, model_metadata, derived_frame_skip)
+      for cam_w, cam_h in args.camera_resolutions:
+        print(f"Compiling unified run_model JIT for {cam_w}x{cam_h}...")
+        nv12 = stock.NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
+        frame_copy_size = stock.nv12_copy_size(nv12.stride, nv12.y_height, nv12.uv_height)
+        make_model_queues = partial(stock.make_input_queues, model_metadata['input_shapes'], derived_frame_skip,
+                                    frame_copy_size=frame_copy_size)
+        warp = stock.make_warp(nv12, model_w, model_h)
+        run_model_jit = TinyJit(legacy.make_legacy_run_model(warp, run_policy, model_metadata, frame_copy_size), prune=True)
+        output_data['run_model'][(cam_w, cam_h)] = compile_jit(run_model_jit, POLICY_INPUTS, make_model_queues, benchmark_runs=args.benchmark_runs)
   else:
     vision_runner = OnnxRunner(args.vision_onnx) if args.vision_onnx else None
     if args.model_type == 'vision_policy':
@@ -355,13 +372,12 @@ if __name__ == "__main__":
         output_data['metadata'][name] = make_metadata_dict(runner_arg)
 
     policy_keys = [key for key in output_data['metadata'].keys() if key != 'vision']
-    first_policy_meta = output_data['metadata'][policy_keys[0]] if policy_keys else {}
-    vision_meta = output_data['metadata'].get('vision', {})
+    first_policy_meta: dict = output_data['metadata'][policy_keys[0]] if policy_keys else {}
+    vision_meta: dict = output_data['metadata'].get('vision', {})
 
     derived_frame_skip = args.frame_skip or derive_frame_skip(vision_meta.get('input_shapes', {}), first_policy_meta.get('input_shapes', {}))
     all_shapes = {key: value for meta in output_data['metadata'].values() for key, value in meta['input_shapes'].items()}
-    feat_meta = output_data['metadata'].get('vision') or output_data['metadata'].get('policy')
-    assert feat_meta is not None
+    feat_meta: dict = vision_meta or first_policy_meta
     features_slice = feat_meta['output_slices']['hidden_state']
 
     print(f"Compiling run_policy JIT (model_size={model_w}x{model_h}, frame_skip={derived_frame_skip})...")

openpilot/sunnypilot/modeld_v2/modeld.py

@@ -36,12 +36,9 @@ from openpilot.system import sentry
 from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
 from openpilot.selfdrive.controls.lib.desire_helper import DesireHelper
 from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_from_plan, smooth_value
-from openpilot.selfdrive.modeld.modeld import ChestnutState
+from openpilot.selfdrive.modeld.modeld import ChestnutGpuState
 
-from openpilot.selfdrive.modeld.compile_modeld import (
-  MODELD_INPUTS,
-  make_input_queues as make_stock_input_queues,
-)
+from openpilot.selfdrive.modeld.compile_modeld import make_input_queues
 from openpilot.sunnypilot.modeld_v2.fill_model_msg import fill_model_msg, fill_pose_msg, PublishState, get_curvature_from_output
 from openpilot.sunnypilot.modeld_v2.parse_model_outputs import Parser
 from openpilot.sunnypilot.modeld_v2.constants import ModelConstants, Plan
@@ -50,6 +47,7 @@ from openpilot.sunnypilot.modeld_v2.camera_offset_helper import CameraOffsetHelp
 from openpilot.sunnypilot.modeld_v2.compile_modeld import (derive_frame_skip, make_split_input_queues,
                                                            make_supercombo_input_queues, nv12_copy_size,
                                                            WARP_INPUTS, POLICY_INPUTS)
+from openpilot.sunnypilot.modeld_v2.stock_dependencies import make_legacy_stock_input_queues
 from openpilot.sunnypilot.livedelay.helpers import get_lat_delay
 from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase
 from openpilot.sunnypilot.modeld_v2.helpers import load_oob
@@ -141,8 +139,14 @@ class ModelState(ModelStateBase):
       self._vision_input_names = [key for key in self.input_shapes if 'img' in key]
       self.frame_skip = derive_frame_skip({}, self.input_shapes)
       if self.is_run_model:
-        self.input_queues, self.numpy_inputs, self.frame_buffers = make_stock_input_queues(
-          self.input_shapes, self.frame_skip, device=self.DEV, frame_copy_size=self.frame_copy_size)
+        self.state_pairs = model_metadata.get('state_pairs', {})
+        self.is_new_model = len(self.state_pairs) > 0
+        if self.is_new_model:
+          self.input_queues, self.numpy_inputs, self.frame_buffers = make_input_queues(self.input_shapes, self.state_pairs,
+                                                                                       device=self.DEV, frame_copy_size=self.frame_copy_size)
+        else:
+          self.input_queues, self.numpy_inputs, self.frame_buffers = make_legacy_stock_input_queues(self.input_shapes, self.frame_skip, device=self.DEV,
+                                                                                                    frame_copy_size=self.frame_copy_size)
         self.frame_views, self.npy = self.frame_buffers, self.numpy_inputs
         self.run_model, self.run_policy, self.warp = jits['run_model'][(cam_w, cam_h)], None, None
       else:
@@ -191,8 +195,12 @@ class ModelState(ModelStateBase):
     dummy_inputs = {k: np.zeros(v.shape, dtype=v.dtype) for k, v in self.numpy_inputs.items() if k not in ['tfm', 'big_tfm', 'prev_feat']}
     self.run(dummy_frames, transforms, dummy_inputs)
     if self.is_run_model:
-      self.input_queues, self.numpy_inputs, self.frame_buffers = make_stock_input_queues(
-        self.input_shapes, self.frame_skip, device=self.DEV, frame_copy_size=self.frame_copy_size)
+      if self.is_new_model:
+        self.input_queues, self.numpy_inputs, self.frame_buffers = make_input_queues(self.input_shapes, self.state_pairs, device=self.DEV,
+                                                                                     frame_copy_size=self.frame_copy_size)
+      else:
+        self.input_queues, self.numpy_inputs, self.frame_buffers = make_legacy_stock_input_queues(self.input_shapes, self.frame_skip, device=self.DEV,
+                                                                                                  frame_copy_size=self.frame_copy_size)
       self.frame_views = self.frame_buffers
       self.npy = self.numpy_inputs
     else:
@@ -242,7 +250,10 @@ class ModelState(ModelStateBase):
     self.numpy_inputs['big_tfm'][:, :] = transforms[self._wide_key].reshape(3, 3)
 
     if self.run_model is not None:
-      outs, = self.run_model(**{k: self.input_queues[k] for k in MODELD_INPUTS})
+      if self.is_new_model:
+        outs, = self.run_model(**self.input_queues)
+      else:
+        outs, = self.run_model(**{k: self.input_queues[k] for k in POLICY_INPUTS})
       raw_outputs = outs
     else:
       assert self.warp is not None and self.run_policy is not None
@@ -373,11 +384,7 @@ def main(demo=False):
     loader.start()
     loader.join(BIG_MODEL_TIMEOUT)
     model = big_model
-    if model is None:
-      params.put_bool("ChestnutModelError", True)
     params.put_bool("ChestnutActive", model is not None)
-    if model is not None:
-      params.remove("ChestnutModelError")
 
   small_model = ModelState(cam_w=vipc_client_main.width, cam_h=vipc_client_main.height, chestnut=False) if model is None or CHESTNUT else None
   if model is None:
@@ -387,12 +394,12 @@ def main(demo=False):
   cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting")
 
   # messaging
-  pub_socks = ["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"] + (["chestnutState"] if CHESTNUT else [])
+  pub_socks = ["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"] + (["chestnutGpuState"] if CHESTNUT else [])
   pm = PubMaster(pub_socks)
   sm = SubMaster(["deviceState", "carState", "narrowRoadCameraState", "extrinsicsCalibration", "driverMonitoringState", "carControl", "lateralDelay"])
 
   publish_state = PublishState()
-  chestnut_state = ChestnutState(pm, model.chestnut) if CHESTNUT else None
+  chestnut_state = ChestnutGpuState(pm, model.chestnut) if CHESTNUT else None
 
   # setup filter to track dropped frames
   frame_dropped_filter = FirstOrderFilter(0., 10., 1. / model.constants.MODEL_FREQ)
@@ -514,13 +521,12 @@ def main(demo=False):
     mt1 = time.perf_counter()
     try:
       send_chestnut = (chestnut_state is not None and
-                       run_count % round(model.constants.MODEL_FREQ / SERVICE_LIST['chestnutState'].frequency) == 0)
+                       run_count % round(model.constants.MODEL_FREQ / SERVICE_LIST['chestnutGpuState'].frequency) == 0)
       model_output = model.run(bufs, transforms, inputs, chestnut_state.send if send_chestnut else None)
     except Exception:
       if not params.get_bool("ChestnutActive"):
         raise
       cloudlog.exception("chestnut failed, falling back to small")
-      params.put_bool("ChestnutModelError", True)
       params.put_bool("ChestnutActive", False)
       assert small_model is not None
       model = small_model

openpilot/sunnypilot/modeld_v2/stock_dependencies.py

@@ -0,0 +1,119 @@
+"""
+Copyright (c) 2021-, Haibin Wen, sunnypilot, and a number of other contributors.
+
+This file is part of sunnypilot and is licensed under the MIT License.
+See the LICENSE.md file in the root directory for more details.
+"""
+
+import math
+import numpy as np
+from functools import partial
+from tinygrad import dtypes
+from tinygrad.device import Device
+from tinygrad.tensor import Tensor
+
+# The old openpilot/selfdrive/modeld/compile_modeld.py functions needed for legacy models
+# We freeze them here so they aren't lost.
+
+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 _detect_desire_key(shapes: dict) -> str | None:
+  return next((key for key in shapes if key.startswith('desire')), None)
+
+def get_policy_npy_shapes(input_shapes: dict, is_supercombo: bool = False) -> tuple[dict, list[int]]:
+  desire_key = _detect_desire_key(input_shapes)
+  shapes = {}
+  if desire_key:
+    shapes['desire'] = (input_shapes[desire_key][2],)
+
+  for key, shape in input_shapes.items():
+    if key not in (desire_key, 'features_buffer') and 'img' not in key:
+      shapes[key] = tuple(shape)
+
+  if is_supercombo and 'features_buffer' in input_shapes:
+    fb = input_shapes['features_buffer']
+    feat_dim = math.prod(fb[2:])
+    shapes['prev_feat'] = (fb[0], feat_dim)
+
+  sizes = [int(np.prod(size)) for size in shapes.values()]
+  return shapes, sizes
+
+def make_legacy_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'], is_supercombo=True)
+  model_input_dtypes = {name: spec.dtype for name, spec in model_runner.graph_inputs.items()}
+
+  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)
+    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.reshape(model_metadata['input_shapes']['features_buffer']),
+      'desire_pulse': desire_buf,
+      'traffic_convention': traffic_convention,
+      'action_t': action_t,
+    }
+    inputs = {name: value.cast(model_input_dtypes.get(name, dtypes.float32)) for name, value in inputs.items()}
+    out = next(iter(model_runner(inputs).values())).cast('float32')
+    return out,
+  return run_policy
+
+def make_legacy_run_model(warp, run_policy, model_metadata, frame_copy_size):
+  _, policy_sizes = get_policy_npy_shapes(model_metadata['input_shapes'], is_supercombo=True)
+  packed_npy_size = (18 + sum(policy_sizes)) * np.dtype(np.float32).itemsize
+
+  def run_model(img_q, big_img_q, feat_q, desire_q, packed_npy_inputs):
+    packed_input = packed_npy_inputs.to(Device.DEFAULT)
+    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
+
+def make_legacy_stock_input_queues(input_shapes, frame_skip, device, frame_copy_size):
+  img = input_shapes['img']  # (1, 12, 128, 256)
+  fb = input_shapes['features_buffer']  # (1, T-1, ...), past features only; the model appends the current frame's feature
+  feat_dim = math.prod(fb[2:])
+  dp = input_shapes['desire_pulse']  # (1, 25, 8)
+  n_frames = img[1] // 6
+  img_buf_shape = (frame_skip * (n_frames - 1) + 1, 6, img[2], img[3])
+
+  policy_shapes, _ = get_policy_npy_shapes(input_shapes, is_supercombo=True)
+  shapes = {'tfm': (3, 3), 'big_tfm': (3, 3)} | policy_shapes
+  sizes = [math.prod(s) for s in shapes.values()]
+  packed_npy_size = sum(sizes) * np.dtype(np.float32).itemsize
+  packed_input = np.zeros(packed_npy_size + 2 * frame_copy_size, dtype=np.uint8)
+  packed_npy_inputs = packed_input[:packed_npy_size].view(np.float32)
+  frames = packed_input[packed_npy_size:]
+  frame_views = {'img': frames[:frame_copy_size], 'big_img': frames[frame_copy_size:]}
+  # views into the packed inputs, to be refilled at runtime
+  npy = {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 = {
+    '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(),
+    'feat_q': Tensor(np.zeros((frame_skip * fb[1], fb[0], feat_dim), 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_input, device='NPY').realize(),
+  }
+  return input_queues, npy, frame_views