this made me cry

openpilot/sunnypilot/SConscript

@@ -1,3 +1,2 @@
 SConscript(['common/transformations/SConscript'])
-SConscript(['modeld_v2/SConscript'])
 SConscript(['selfdrive/locationd/SConscript'])

openpilot/sunnypilot/modeld_v2/SConscript

@@ -1,105 +0,0 @@
-import os
-import glob
-import sys
-import subprocess
-
-from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
-from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE
-from openpilot.common.hardware import HARDWARE, PC
-from openpilot.selfdrive.modeld.helpers import usbgpu_present
-
-Import('env', 'arch', 'release')
-lenv = env.Clone()
-tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + "/**", recursive=True, root_dir=env.Dir("#").abspath) if 'pycache' not in x]
-
-
-def get_camera_configs():
-  DEVICE_RESOLUTIONS = {
-    "tici": (_ar_ox_fisheye.width, _ar_ox_fisheye.height),
-    "tizi": (_ar_ox_fisheye.width, _ar_ox_fisheye.height),
-    "mici": (_os_fisheye.width, _os_fisheye.height),
-  }
-  if release or PC or 'CI' in os.environ:
-    return set(DEVICE_RESOLUTIONS.values())
-  return [DEVICE_RESOLUTIONS[HARDWARE.get_device_type()]]
-
-CAMERA_CONFIGS = get_camera_configs()
-
-# remove me after sync
-def probe_devices():
-  return set(subprocess.run(
-    [sys.executable, '-c', 'from tinygrad import Device\nprint("\\n".join(Device.get_available_devices()))'],
-    capture_output=True, text=True, check=True).stdout.strip().splitlines())
-
-available = probe_devices()  #remove me after sync
-if 'QCOM' in available:  # change to this after sync. if arch == 'comma_arm64':
-  tg_backend = 'QCOM'
-  tg_flags = f'DEV={tg_backend} IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1'
-else:
-  tg_backend = 'CPU'
-  tg_flags = f'DEV=CPU HOME={os.path.expanduser("~")}' if arch == 'Darwin' else 'DEV=CPU:LLVM'
-
-model_w, model_h = MEDMODEL_INPUT_SIZE
-from openpilot.selfdrive.modeld.constants import ModelConstants
-frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ
-camera_res_args = ' '.join(f'{cw}x{ch}' for cw, ch in CAMERA_CONFIGS)
-
-pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + ':' + env.Dir("#").abspath + '"'
-compile_modeld_script = File("compile_modeld.py").abspath
-upstream_compile_script = File(Dir("#openpilot/selfdrive/modeld").File("compile_modeld.py").abspath)
-script_deps = [File("compile_modeld.py"), upstream_compile_script]
-
-USBGPU = usbgpu_present()
-if USBGPU:
-  usbgpu_tg_flags = f'DEBUG=2 DEV=USB+AMD:LLVM WARP_DEV={tg_backend} FLOAT16=1 JIT_BATCH_SIZE=0 GMMU=0 TC_OPT=2'
-  usbgpu_lock = File("models/.usb_gpu.lock").abspath
-
-def compile_combined(model_type, onnx_args, output_name):
-  for usbgpu in ([False, True] if USBGPU else [False]):
-    prefix = 'big_' if usbgpu else ('big_' if os.getenv('BIG_INTO_SMALL') else '')
-    final_output_name = prefix + output_name
-    output_pkl = File(f"models/{final_output_name}").abspath
-
-    active_tg_flags = usbgpu_tg_flags if usbgpu else tg_flags
-
-    cmd = (f'{pythonpath_string} {active_tg_flags} python3 {compile_modeld_script} '
-           f'--model-type {model_type} '
-           f'--model-size {model_w}x{model_h} '
-           f'--camera-resolutions {camera_res_args} '
-           f'{onnx_args} '
-           f'--frame-skip {frame_skip} '
-           f'--output {output_pkl}')
-    onnx_files = [f for f in onnx_args.split() if f.endswith('.onnx')]
-    node = lenv.Command(output_pkl, tinygrad_files + script_deps + [File(f) for f in onnx_files if os.path.isfile(f)], cmd)
-    if usbgpu:
-      lenv.SideEffect(usbgpu_lock, node)
-
-# Vision + Policy (stock default model)
-vision_onnx = File("models/driving_vision.onnx").abspath
-policy_onnx = File("models/driving_policy.onnx").abspath
-if os.path.isfile(vision_onnx) and os.path.isfile(policy_onnx):
-  compile_combined('vision_policy',
-                   f'--vision-onnx {vision_onnx} --policy-onnx {policy_onnx}',
-                   'driving_combined_tinygrad.pkl')
-
-# Vision + Off-Policy
-off_policy_onnx = File("models/driving_off_policy.onnx").abspath
-if os.path.isfile(vision_onnx) and os.path.isfile(off_policy_onnx):
-  policy_arg = f'--policy-onnx {policy_onnx}' if os.path.isfile(policy_onnx) else ''
-  compile_combined('vision_multi_policy',
-                   f'--vision-onnx {vision_onnx} {policy_arg} --off-policy-onnx {off_policy_onnx}',
-                   'driving_combined_multi_tinygrad.pkl')
-
-# Vision + On-Policy + Off-Policy
-on_policy_onnx = File("models/driving_on_policy.onnx").abspath
-if os.path.isfile(vision_onnx) and os.path.isfile(on_policy_onnx) and os.path.isfile(off_policy_onnx):
-  compile_combined('vision_multi_policy',
-                   f'--vision-onnx {vision_onnx} --off-policy-onnx {off_policy_onnx} --on-policy-onnx {on_policy_onnx}',
-                   'driving_combined_tri_tinygrad.pkl')
-
-# Supercombo
-supercombo_onnx = File("models/supercombo.onnx").abspath
-if os.path.isfile(supercombo_onnx):
-  compile_combined('supercombo',
-                   f'--supercombo-onnx {supercombo_onnx}',
-                   'driving_combined_supercombo_tinygrad.pkl')

openpilot/sunnypilot/modeld_v2/compile_modeld.py

@@ -9,7 +9,7 @@ See the LICENSE.md file in the root directory for more details.
 import argparse
 import os
 import tempfile
-from collections import defaultdict
+import time
 from functools import partial
 from openpilot.selfdrive.modeld.helpers import dump_oob, load_oob
 import numpy as np
@@ -38,6 +38,9 @@ from tinygrad.engine.jit import TinyJit
 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']
+WARP_DEV = os.getenv('WARP_DEV')
 
 
 def _detect_desire_key(shapes: dict) -> str | None:
@@ -132,9 +135,35 @@ def make_supercombo_input_queues(input_shapes: dict, frame_skip: int,
   return generate_queues_and_npy(input_shapes, frame_skip, device, is_supercombo=True)
 
 
-def create_jit_runner(vision_runner, policy_runners: list, nv12: NV12Frame, model_size: tuple[int, int],
-            features_slice: slice, frame_skip: int, input_shapes: dict, prepare_only: bool):
-  frame_prepare = make_frame_prepare(nv12, *model_size)
+def make_random_images(keys, shape, device):
+  return {k: Tensor.randint(shape, low=0, high=256, dtype=dtypes.uint8, device=device).realize() for k in keys}
+
+
+def make_warp_queues(device=Device.DEFAULT):
+  npy = {
+    'tfm': np.zeros((3, 3), dtype=np.float32),
+    'big_tfm': np.zeros((3, 3), dtype=np.float32),
+  }
+  queues = {k: Tensor(v, device='NPY').realize() for k, v in npy.items()}
+  return queues, npy
+
+
+def make_warp(nv12: NV12Frame, model_w: int, model_h: int):
+  frame_prepare = make_frame_prepare(nv12, model_w, model_h)
+  WARP_DEV = os.getenv('WARP_DEV', Device.DEFAULT)
+
+  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(vision_runner, policy_runners: list, features_slice: slice, frame_skip: int, input_shapes: dict):
   sample_skip_fn = partial(sample_skip, frame_skip=frame_skip)
   sample_desire_fn = partial(sample_desire, frame_skip=frame_skip)
 
@@ -147,20 +176,14 @@ def create_jit_runner(vision_runner, policy_runners: list, nv12: NV12Frame, mode
   is_supercombo = vision_runner is None
   npy_shapes, npy_sizes = get_policy_npy_shapes(input_shapes, is_supercombo=is_supercombo)
 
-  def runner(img_q, big_img_q, feat_q, packed_npy_inputs, frame, big_frame, tfm, big_tfm, **kwargs):
+  def run_policy(warped, img_q, big_img_q, feat_q, packed_npy_inputs, **kwargs):
     desire_q = kwargs['desire_q']
-
     packed_npy_inputs_dev = packed_npy_inputs.to(Device.DEFAULT)
-    tfm_dev = tfm.to(Device.DEFAULT)
-    big_tfm_dev = big_tfm.to(Device.DEFAULT)
+    warped_dev = warped.to(Device.DEFAULT)
+    Tensor.realize(packed_npy_inputs_dev, warped_dev)
 
-    Tensor.realize(packed_npy_inputs_dev, tfm_dev, big_tfm_dev)
-
-    img = shift_and_sample(img_q, frame_prepare(frame, tfm_dev).unsqueeze(0), sample_skip_fn).realize()
-    big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm_dev).unsqueeze(0), sample_skip_fn).realize()
-
-    if prepare_only:
-      return img, big_img
+    img = shift_and_sample(img_q, warped_dev[0:1], sample_skip_fn).realize()
+    big_img = shift_and_sample(big_img_q, warped_dev[1:2], sample_skip_fn).realize()
 
     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))
@@ -195,50 +218,52 @@ def create_jit_runner(vision_runner, policy_runners: list, nv12: NV12Frame, mode
       shift_and_sample(feat_q, new_feat, sample_skip_fn).realize()
     return policy_out
 
-  return runner
+  return run_policy
 
 
-def compile_and_warmup(nv12: NV12Frame, model_size: tuple[int, int], prepare_only: bool, frame_skip: int, vision_runner, policy_runners: list, metadata: dict):
-  print(f"Compiling combined JIT for {nv12.width}x{nv12.height} (prepare_only={prepare_only})...")
-
-  all_shapes = {key: value for meta in metadata.values() for key, value in meta['input_shapes'].items()}
-
-  feat_meta = metadata.get('vision') or metadata.get('model') or metadata.get('policy')
-  if not feat_meta:
-    raise ValueError("Could not find vision, model, or policy metadata.")
-
-  features_slice = feat_meta['output_slices']['hidden_state']
-  WARP_DEV = os.getenv('WARP_DEV', Device.DEFAULT)
-
-  is_supercombo = vision_runner is None
-  run_func = create_jit_runner(vision_runner, policy_runners, nv12, model_size, features_slice, frame_skip, all_shapes, prepare_only)
-  run_jit = TinyJit(run_func, prune=True)
-  def run_once(seed, queues=None, npy=None):
-    if queues is None or npy is None:
-      queues, npy = generate_queues_and_npy(all_shapes, frame_skip, Device.DEFAULT, is_supercombo=is_supercombo)
+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)
     rng = np.random.default_rng(seed)
     Tensor.manual_seed(seed)
-    frame = Tensor.randint(nv12.size, low=0, high=256, dtype=dtypes.uint8, device=WARP_DEV).realize()
-    big_frame = Tensor.randint(nv12.size, low=0, high=256, dtype=dtypes.uint8, device=WARP_DEV).realize()
-    for value in npy.values():
-      value[:] = rng.standard_normal(value.shape).astype(value.dtype)
-    Device.default.synchronize()
-    outs = run_jit(**queues, frame=frame, big_frame=big_frame)
-    Device.default.synchronize()
-    return [np.copy(value.numpy()) for value in (outs if isinstance(outs, tuple) else [outs])] if outs is not None else []
 
-  warmup_queues, warmup_npy = generate_queues_and_npy(all_shapes, frame_skip, Device.DEFAULT, is_supercombo=is_supercombo)
-  for i in range(3):
-    run_once(42 + i, warmup_queues, warmup_npy)
+    testing = test_val is not None or test_buffers is not None
+    n_runs = 1 if testing else 3
 
-  if not prepare_only:
-    baseline = run_once(42)
-    with tempfile.TemporaryFile(dir=".") as f:
-      dump_oob(run_jit, f)
-      f.seek(0)
-      run_jit = load_oob(f)
-    assert all(np.array_equal(baseline, deserialized) for baseline, deserialized in zip(baseline, run_once(42), strict=True)), "OOB pickling regression"
-  return run_jit
+    for i in range(n_runs):
+      for v in npy.values():
+        v[:] = rng.standard_normal(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 if k in input_queues}, **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 if isinstance(outs, tuple) else [outs])] if outs is not None else []
+        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')
+  with tempfile.TemporaryFile(dir=".") as f:
+    dump_oob(jit, f)
+    f.seek(0)
+    deserialized_jit = load_oob(f)
+  random_inputs_run(deserialized_jit, SEED, test_val=test_val, test_buffers=test_buffers)
+  return deserialized_jit
 
 
 def _parse_size(size_str: str) -> tuple[int, int]:
@@ -260,19 +285,6 @@ def read_file_chunked_to_shm(path):
   return shm_path
 
 
-def _compile_for_resolutions(camera_resolutions: list, model_size: tuple[int, int], frame_skip: int,
-               vision_runner, policy_runners: list, metadata: dict) -> dict:
-  from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
-  return {
-    (cam_w, cam_h): {
-      name: compile_and_warmup(NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)), model_size, prepare_only,
-                   frame_skip, vision_runner, policy_runners, metadata)
-      for name, prepare_only in [('warp_enqueue', True), ('run_policy', False)]
-    }
-    for cam_w, cam_h in camera_resolutions
-  }
-
-
 def _load_policy_runners(args: argparse.Namespace) -> tuple[list, list]:
   runners, keys = [], []
   for name, onnx_arg in [('policy', args.policy_onnx), ('off_policy', args.off_policy_onnx), ('on_policy', args.on_policy_onnx)]:
@@ -284,7 +296,6 @@ def _load_policy_runners(args: argparse.Namespace) -> tuple[list, list]:
 
 if __name__ == "__main__":
   if 'USB' in os.getenv('DEV', '') or os.getenv('USBGPU'):
-    import time
     from openpilot.system.hardware.chestnut.flash import link_up
     for _ in range(10):
       if link_up():
@@ -293,7 +304,9 @@ if __name__ == "__main__":
     else:
       raise RuntimeError("Chestnut not ready, skipping big model build")
 
+  from openpilot.common.file_chunker import chunk_file, get_chunk_targets
   from openpilot.selfdrive.modeld.get_model_metadata import make_metadata_dict
+  from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
   from tinygrad.nn.onnx import OnnxRunner
 
   parser = argparse.ArgumentParser(description="Compile combined JIT pkl for sunnypilot modeld_v2")
@@ -310,7 +323,8 @@ if __name__ == "__main__":
   parser.add_argument('--supercombo-onnx', help='supercombo ONNX (for supercombo)')
 
   args = parser.parse_args()
-  output_data = defaultdict(dict)
+  model_w, model_h = args.model_size
+  output_data = {}
 
   args.vision_onnx = read_file_chunked_to_shm(args.vision_onnx)
   args.policy_onnx = read_file_chunked_to_shm(args.policy_onnx)
@@ -341,16 +355,31 @@ if __name__ == "__main__":
   vision_meta = 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', {}))
-  output_data.update(_compile_for_resolutions(args.camera_resolutions, args.model_size, derived_frame_skip,
-                        vision_runner, policy_runners, output_data['metadata']))
+  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('model') or output_data['metadata'].get('policy')
+  assert feat_meta is not None
+  features_slice = feat_meta['output_slices']['hidden_state']
+  is_supercombo = vision_runner is None
+
+  print(f"Compiling run_policy JIT (model_size={model_w}x{model_h}, frame_skip={derived_frame_skip})...")
+  run_policy_func = make_run_policy(vision_runner, policy_runners, features_slice, derived_frame_skip, all_shapes)
+  run_policy_jit = TinyJit(run_policy_func, prune=True)
+  make_policy_queues = partial(generate_queues_and_npy, all_shapes, derived_frame_skip, is_supercombo=is_supercombo)
+  make_random_model_inputs = partial(make_random_images, keys=['warped'], shape=(2, 6, model_h // 2, model_w // 2), device=WARP_DEV)
+  output_data['run_policy'] = compile_jit(run_policy_jit, make_random_model_inputs, POLICY_INPUTS, make_policy_queues)
+
+  for cam_w, cam_h in args.camera_resolutions:
+    print(f"Compiling warp JIT for {cam_w}x{cam_h}...")
+    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 = TinyJit(make_warp(nv12, model_w, model_h), prune=True)
+    output_data[(cam_w, cam_h)] = compile_jit(warp, make_random_warp_inputs, WARP_INPUTS, make_warp_queues)
 
   with open(args.output, "wb") as file:
     dump_oob(output_data, file)
 
   pkl_size = os.path.getsize(args.output)
   print(f"Saved combined JIT to {args.output} ({pkl_size / 1e6:.2f} MB)")
-
-  from openpilot.common.file_chunker import chunk_file, get_chunk_targets
   chunk_targets = get_chunk_targets(args.output, pkl_size)
   chunk_file(args.output, chunk_targets)
   print(f"Chunked into {len(chunk_targets) - 1} file(s)")

openpilot/sunnypilot/modeld_v2/modeld.py

@@ -41,7 +41,7 @@ from openpilot.sunnypilot.modeld_v2.fill_model_msg import fill_model_msg, fill_p
 from openpilot.sunnypilot.modeld_v2.constants import Plan
 from openpilot.sunnypilot.modeld_v2.meta_helper import load_meta_constants
 from openpilot.sunnypilot.modeld_v2.camera_offset_helper import CameraOffsetHelper
-from openpilot.sunnypilot.modeld_v2.compile_modeld import derive_frame_skip, make_split_input_queues
+from openpilot.sunnypilot.modeld_v2.compile_modeld import derive_frame_skip, make_split_input_queues, make_supercombo_input_queues, WARP_INPUTS, POLICY_INPUTS
 
 from openpilot.sunnypilot.livedelay.helpers import get_lat_delay
 from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase
@@ -112,11 +112,10 @@ class ModelState(ModelStateBase):
     self.WARP_DEV = 'QCOM' if COMMA_HARDWARE else 'CPU'
     self.DEV = 'AMD' if self.usbgpu else self.WARP_DEV
     self.QUEUE_DEV = self.DEV
-
     metadata = jits['metadata']
 
-    self._run_policy = jits[(cam_w, cam_h)]['run_policy']
-    self._warp_enqueue = jits[(cam_w, cam_h)]['warp_enqueue']
+    self.run_policy = jits['run_policy']
+    self.warp = jits[(cam_w, cam_h)]
 
     if 'model' in metadata:
       model_metadata = metadata['model']
@@ -125,7 +124,6 @@ class ModelState(ModelStateBase):
       self._policy_slices_list = []
       self._combined_model_type = 'supercombo'
       self._vision_input_names = [key for key in model_metadata['input_shapes'] if 'img' in key]
-      from openpilot.sunnypilot.modeld_v2.compile_modeld import make_supercombo_input_queues
       frame_skip = derive_frame_skip({}, model_metadata['input_shapes'])
       self.input_queues, self.numpy_inputs = make_supercombo_input_queues(model_metadata['input_shapes'],
                                                                           frame_skip, device=self.QUEUE_DEV)
@@ -141,12 +139,11 @@ class ModelState(ModelStateBase):
       self._policy_slices_list = [metadata[k]['output_slices'] for k in policy_keys]
       self.policy_output_slices = self._policy_slices_list[0]
       self._has_on_policy = any('on' in k.lower() for k in policy_keys)
-      first_policy_metadata = metadata[policy_keys[0]]
-      vision_input_shapes = vision_metadata['input_shapes']
-      policy_input_shapes = first_policy_metadata['input_shapes']
-      self._vision_input_names = [k for k in vision_input_shapes if 'img' in k]
-      frame_skip = derive_frame_skip(vision_input_shapes, policy_input_shapes)
-      self.input_queues, self.numpy_inputs = make_split_input_queues(vision_input_shapes, policy_input_shapes,
+      self._vision_input_names = [key for key in vision_metadata['input_shapes'] if 'img' in key]
+      first_policy_meta = metadata[policy_keys[0]]
+      frame_skip = derive_frame_skip(vision_metadata['input_shapes'], first_policy_meta['input_shapes'])
+      self.input_queues, self.numpy_inputs = make_split_input_queues(vision_metadata['input_shapes'],
+                                                                     first_policy_meta['input_shapes'],
                                                                      frame_skip, device=self.QUEUE_DEV)
 
     self._desire_key = next(key for key in self.numpy_inputs if key.startswith('desire'))
@@ -175,8 +172,9 @@ class ModelState(ModelStateBase):
     self.frame_buf_params = dict.fromkeys(self._vision_input_names, nv12_info)
 
     yuv_size = self.frame_buf_params[self._road_key][3]
-    self._warp_enqueue(
-      **self.input_queues,
+    self.warp(
+      tfm=self.input_queues['tfm'],
+      big_tfm=self.input_queues['big_tfm'],
       frame=Tensor(np.zeros(yuv_size, dtype=np.uint8), device=self.WARP_DEV).contiguous().realize(),
       big_frame=Tensor(np.zeros(yuv_size, dtype=np.uint8), device=self.WARP_DEV).contiguous().realize())
 
@@ -237,10 +235,10 @@ class ModelState(ModelStateBase):
     self.numpy_inputs['big_tfm'][:, :] = transforms[wide_key].reshape(3, 3)
 
     if prepare_only:
-      self._warp_enqueue(**self.input_queues, frame=self.full_frames[road_key], big_frame=self.full_frames[wide_key])
+      self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=self.full_frames[road_key], big_frame=self.full_frames[wide_key])
       return None
-
-    raw_outputs = self._run_policy(**self.input_queues, frame=self.full_frames[road_key], big_frame=self.full_frames[wide_key])
+    warped = self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=self.full_frames[road_key], big_frame=self.full_frames[wide_key])
+    raw_outputs = self.run_policy(**{k: self.input_queues[k] for k in POLICY_INPUTS if k in self.input_queues}, warped=warped)
 
     if self._combined_model_type == 'supercombo':
       model_output = raw_outputs.numpy().flatten()

openpilot/sunnypilot/modeld_v2/tests/helpers.py

@@ -163,7 +163,8 @@ ARCHETYPES = {
 def make_pkl_data(archetype):
   return {
     'metadata': archetype.metadata_structure,
-    (CAM_W, CAM_H): {'run_policy': _noop_jit, 'warp_enqueue': _noop_jit},
+    'run_policy': _noop_jit,
+    (CAM_W, CAM_H): _noop_jit,
   }