Use tinygrad ONNX and warp compilers for driving and DM (#38922)

* modeld: use tinygrad ONNX and warp compilers for driving and DM

* Bump tinygrad with current openpilot model compile tests

* modeld: prioritize tinygrad compiler imports in build environment

* Use compiler branch based on openpilot's pinned tinygrad

* modeld: invoke tinygrad compilers without firmware patch

.gitmodules

@@ -15,4 +15,5 @@
   url = ../../commaai/teleoprtc
 [submodule "tinygrad"]
   path = tinygrad_repo
-  url = https://github.com/tinygrad/tinygrad.git
+  url = https://github.com/commaai/tinygrad.git
+  branch = openpilot-modeld-test

openpilot/selfdrive/modeld/SConscript

@@ -1,18 +1,18 @@
 import glob
-import json
 import os
 import time
 from SCons.Script import Action, Value
 from openpilot.common.file_chunker import chunk_file, get_chunk_targets, get_existing_chunks
 from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
 from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE, DM_INPUT_SIZE
-from openpilot.selfdrive.modeld.helpers import TG_INPUT_DEVICES_PATH, chestnut_present, modeld_pkl_path
+from openpilot.selfdrive.modeld.helpers import chestnut_present, modeld_pkl_path
 from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
 
 
 Import('env', 'arch')
 chunker_file = File("#openpilot/common/file_chunker.py")
 lenv = env.Clone()
+lenv.PrependENVPath('PYTHONPATH', Dir('#tinygrad_repo').abspath)
 
 tinygrad_root = env.Dir("#").abspath
 tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + "/**", recursive=True, root_dir=tinygrad_root)
@@ -26,62 +26,27 @@ def estimate_pickle_max_size(onnx_size):
 camera_configs = [(c.width, c.height) for c in (_ar_ox_fisheye, _os_fisheye)]
 
 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'
+  tg_flags = 'DEV=QCOM IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1'
 else:
-  tg_backend = 'METAL' if arch == 'Darwin' else 'CPU'
   # JIT=2 disables graph batching, which produces incorrect outputs after buffers change.
   tg_flags = 'DEV=METAL JIT=2' if arch == 'Darwin' else 'DEV=CPU:LLVM'
 
-tg_devices = { # which device to put jit inputs to at runtime
-  'openpilot.selfdrive.modeld.dmonitoringmodeld': {
-    'default': {'DEV': tg_backend}
-  },
-}
-
 CHESTNUT = chestnut_present()
 if CHESTNUT:
   chestnut_tg_flags = 'DEBUG=1 DEV=USB+AMD:LLVM FRAME_DEV=CPU FLOAT16=1 JIT_BATCH_SIZE=0 GMMU=0 TC_OPT=2 TC_MIN_GLOBALS=32'
   # the USB+AMD GPU takes an exclusive flock; serialize all targets that touch it
   chestnut_lock = File("models/.chestnut.lock").abspath
 
-def write_tg_devices(target, source, env):
-  with open(str(target[0]), "w") as f:
-    json.dump(tg_devices, f)
-    f.write("\n")
-
-tg_devices_node = lenv.Command(
-  str(TG_INPUT_DEVICES_PATH),
-  [Value(tg_devices)],
-  write_tg_devices,
-)
-
 # tinygrad calls brew which needs a $HOME in the env
 mac_brew_string = f'HOME={os.path.expanduser("~")}' if arch == 'Darwin' else ''
 
-modeld_dir = Dir("#openpilot/selfdrive/modeld").abspath
-compile_modeld_script = [
-  File(f"{modeld_dir}/compile_modeld.py"),
-  File(f"{modeld_dir}/get_model_metadata.py"),
-  File(f"{modeld_dir}/helpers.py"),
-  File("#openpilot/common/hardware/hw.py"),
-]
-compile_warp_script = [File(f"{modeld_dir}/compile_warp.py"), File(f"{modeld_dir}/helpers.py"),
-                       File("#openpilot/system/camerad/cameras/nv12_info.py")]
-model_w, model_h = MEDMODEL_INPUT_SIZE
+warp_deps = [File("#openpilot/system/camerad/cameras/nv12_info.py")]
+compiler = 'python3 -m examples.openpilot'
+# CPU 7 is isolated with isolcpus on AGNOS, so explicitly pin the compiler to it.
+taskset = 'taskset -c 7 ' if arch == 'comma_arm64' else ''
 
-for chestnut in [False, True] if CHESTNUT else [False]:
-  target_pkl_path = File(modeld_pkl_path(chestnut)).abspath
-  file_prefix, cmd_flags = ('big_', chestnut_tg_flags) if chestnut else ('', tg_flags)
-  driving_onnx_deps = get_existing_chunks(File(f"models/{file_prefix}driving_supercombo.onnx").abspath)
-  # CPU 7 is isolated with isolcpus on AGNOS, so explicitly pin the compiler to it.
-  taskset = 'taskset -c 7 ' if arch == 'comma_arm64' else ''
-  cmd = (f'{cmd_flags} {mac_brew_string} {taskset}python3 {modeld_dir}/compile_modeld.py '
-        f'--onnx {File(f"models/{file_prefix}driving_supercombo.onnx").abspath} '
-        f'--output {target_pkl_path}')
-  onnx_sizes_sum = sum(os.path.getsize(f) for f in driving_onnx_deps)
-  chunk_targets = get_chunk_targets(target_pkl_path, estimate_pickle_max_size(onnx_sizes_sum))
-  def do_compile(target, source, env, command=cmd, pkl=target_pkl_path, chunks=chunk_targets):
+def chestnut_action(command, pkl=None, chunks=()):
+  def do_compile(target, source, env):
     from openpilot.system.hardware.chestnut.flash import link_up
     # chestnut can enumerate before its PCIe link is up due to varying 12V power behavior across cars
     for _ in range(10):
@@ -95,58 +60,49 @@ for chestnut in [False, True] if CHESTNUT else [False]:
       return ret
     if chunks:
       chunk_file(pkl, chunks)
+  return Action(do_compile, "  [CHESTNUT] $TARGET")
+
+def compile_model(onnx_path, pkl_path, flags, chestnut=False):
+  onnx_path, target_pkl_path = File(onnx_path).abspath, File(pkl_path).abspath
+  onnx_deps = get_existing_chunks(onnx_path)
+  cmd = (f'{flags} {mac_brew_string} {taskset}{compiler}.compile_onnx '
+         f'--onnx {onnx_path} --output {target_pkl_path}')
+  onnx_sizes_sum = sum(os.path.getsize(f) for f in onnx_deps)
+  chunk_targets = get_chunk_targets(target_pkl_path, estimate_pickle_max_size(onnx_sizes_sum))
   def do_chunk(target, source, env, pkl=target_pkl_path, chunks=chunk_targets):
     chunk_file(pkl, chunks)
-  actions = Action(do_compile, "  [CHESTNUT] $TARGET") if chestnut else [cmd, Action(do_chunk, "  [CHUNK] $TARGET")]
+  actions = chestnut_action(cmd, target_pkl_path, chunk_targets) if chestnut else [cmd, Action(do_chunk, "  [CHUNK] $TARGET")]
   node = lenv.Command(
     chunk_targets,
-    tinygrad_files + compile_modeld_script + driving_onnx_deps + [Value(cmd), Value(chunk_targets), chunker_file],
+    tinygrad_files + onnx_deps + [Value(cmd), Value(chunk_targets), chunker_file],
     actions,
   )
   if chestnut:
     lenv.SideEffect(chestnut_lock, node)
 
+compile_model('models/dmonitoring_model.onnx', 'models/dmonitoring_model_tinygrad.pkl', tg_flags)
+
+model_w, model_h = MEDMODEL_INPUT_SIZE
+for chestnut in [False, True] if CHESTNUT else [False]:
+  file_prefix, cmd_flags = ('big_', chestnut_tg_flags) if chestnut else ('', tg_flags)
+  compile_model(f'models/{file_prefix}driving_supercombo.onnx', modeld_pkl_path(chestnut), cmd_flags, chestnut)
   for cam_w, cam_h in camera_configs:
     warp_pkl_path = File(f"models/{file_prefix}driving_warp_{cam_w}x{cam_h}_tinygrad.pkl").abspath
     stride, y_height, uv_height, _ = get_nv12_info(cam_w, cam_h)
-    cmd = (f'{cmd_flags} {mac_brew_string} {taskset}python3 {modeld_dir}/compile_warp.py '
+    cmd = (f'{cmd_flags} {mac_brew_string} {taskset}{compiler}.compile_warp '
            f'--camera-resolution {cam_w}x{cam_h} --warp-to {model_w}x{model_h} --layout yuv420 '
-           f'--frames 2 --frame-size {stride * (y_height + uv_height)} '
+           f'--frames 2 --stride {stride} --uv-offset {stride * y_height} --frame-size {stride * (y_height + uv_height)} '
            f'--output {warp_pkl_path}')
-    def do_compile_warp(target, source, env, command=cmd):
-      return do_compile(target, source, env, command=command, chunks=())
-    action = Action(do_compile_warp, "  [CHESTNUT] $TARGET") if chestnut else cmd
-    node = lenv.Command(warp_pkl_path, tinygrad_files + compile_warp_script + [Value(cmd)], action)
+    action = chestnut_action(cmd) if chestnut else cmd
+    node = lenv.Command(warp_pkl_path, tinygrad_files + warp_deps + [Value(cmd)], action)
     if chestnut:
       lenv.SideEffect(chestnut_lock, node)
 
-# get model metadata
-fn = File(f"models/dmonitoring_model").abspath
-script_files = [File(Dir("#openpilot/selfdrive/modeld").File("get_model_metadata.py").abspath)]
-cmd = f'{tg_flags} {mac_brew_string} python3 {Dir("#openpilot/selfdrive/modeld").abspath}/get_model_metadata.py {fn}.onnx'
-lenv.Command(fn + "_metadata.pkl", [fn + ".onnx"] + tinygrad_files + script_files + [tg_devices_node], cmd)
-
 dm_w, dm_h = DM_INPUT_SIZE
 for cam_w, cam_h in camera_configs:
   dm_pkl_path = File(f"models/dm_warp_{cam_w}x{cam_h}_tinygrad.pkl").abspath
-  cmd = (f'{tg_flags} {mac_brew_string} python3 {modeld_dir}/compile_warp.py '
+  stride, y_height, _, frame_size = get_nv12_info(cam_w, cam_h)
+  cmd = (f'{tg_flags} {mac_brew_string} {compiler}.compile_warp '
          f'--camera-resolution {cam_w}x{cam_h} --warp-to {dm_w}x{dm_h} --layout luma --border-fill 16 --transform-device NPY '
-         f'--output {dm_pkl_path}')
-  lenv.Command(dm_pkl_path, tinygrad_files + compile_warp_script + [tg_devices_node], cmd)
-
-def tg_compile(flags, model_name):
-  pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"'
-  fn = File(f"models/{model_name}").abspath
-  pkl = fn + "_tinygrad.pkl"
-  onnx_path = fn + ".onnx"
-  chunk_targets = get_chunk_targets(pkl, estimate_pickle_max_size(os.path.getsize(onnx_path)))
-  def do_chunk(target, source, env):
-    chunk_file(pkl, chunk_targets)
-  return lenv.Command(
-    chunk_targets,
-    [onnx_path] + tinygrad_files + [Value(chunk_targets), chunker_file, tg_devices_node],
-    [f'{pythonpath_string} {flags} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {pkl}',
-     Action(do_chunk, "  [CHUNK] $TARGET")],
-  )
-
-tg_compile(tg_flags, 'dmonitoring_model')
+         f'--stride {stride} --uv-offset {stride * y_height} --frame-size {frame_size} --output {dm_pkl_path}')
+  lenv.Command(dm_pkl_path, tinygrad_files + warp_deps + [Value(cmd)], cmd)

openpilot/selfdrive/modeld/compile_modeld.py

@@ -1,117 +0,0 @@
-#!/usr/bin/env python3
-import argparse
-import atexit
-import os
-import tempfile
-import time
-import shutil
-
-import numpy as np
-
-from openpilot.selfdrive.modeld.helpers import dump_oob, load_oob, patch_tinygrad_fetch_fw
-patch_tinygrad_fetch_fw()
-
-from tinygrad.tensor import Tensor
-from tinygrad.device import Device
-from tinygrad.engine.jit import TinyJit
-
-
-def make_input_queues(input_shapes, device):
-  return {name: Tensor(np.zeros(shape, dtype=dtype.fmt), device=device).realize() for name, (shape, dtype) in input_shapes.items()}
-
-
-def make_run_model(model_runner, state_pairs):
-  def run_model(**inputs):
-    outputs = {name: value.contiguous() for name, value in model_runner(inputs).items()}
-    Tensor.realize(*outputs.values())
-    if state_pairs:
-      Tensor.realize(*(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_shapes, benchmark_runs):
-  if benchmark_runs < 1:
-    raise ValueError("benchmark_runs must be at least 1")
-
-  SEED = 42
-  def random_inputs_run(fn, seed, n_runs, test_val=None, test_buffers=None, expect_match=True):
-    input_queues = make_input_queues(input_shapes, Device.DEFAULT)
-    rng = np.random.default_rng(seed)
-
-    for i in range(n_runs):
-      for value in input_queues.values():
-        values = rng.standard_normal(value.shape) if np.issubdtype(np.dtype(value.dtype.fmt), np.floating) else rng.integers(0, 256, value.shape)
-        value.assign(Tensor(values.astype(value.dtype.fmt), device=Device.DEFAULT)).realize()
-      Device.default.synchronize()
-      st = time.perf_counter()
-      outs = fn(**input_queues)
-      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 = [v.numpy() for v in outs]
-        buffers = [v.numpy() 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, 3)
-  print(f'pickle round trip ({benchmark_runs} runs per seed)')
-  with tempfile.TemporaryFile(dir=".") as f:
-    dump_oob(jit, f)
-    f.seek(0)
-    loaded_jit = load_oob(f)
-  random_inputs_run(loaded_jit, SEED, benchmark_runs, test_val, test_buffers, expect_match=True)
-  random_inputs_run(loaded_jit, SEED+1, benchmark_runs, test_val, test_buffers, expect_match=False)
-  return jit
-
-
-def read_file_chunked_to_disk(path):
-  from openpilot.common.file_chunker import open_file_chunked
-  tmp_path = f'{path}.unchunked'
-  with open(tmp_path, 'wb') as f, open_file_chunked(path) as src:
-    shutil.copyfileobj(src, f)
-  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.selfdrive.modeld.get_model_metadata import make_metadata_dict
-  p = argparse.ArgumentParser()
-  p.add_argument('--onnx', required=True)
-  p.add_argument('--output', required=True)
-  p.add_argument('--benchmark-runs', type=int, default=1,
-                 help='timed loaded-JIT runs for each correctness seed')
-  args = p.parse_args()
-
-  model_path = read_file_chunked_to_disk(args.onnx)
-
-  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 = make_run_model(model_runner, state_pairs)
-  out['run_model'] = compile_jit(TinyJit(run_model, prune=True), input_shapes, 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/compile_warp.py

@@ -1,150 +0,0 @@
-import argparse
-import pickle
-import time
-from collections import namedtuple
-
-from openpilot.selfdrive.modeld.helpers import patch_tinygrad_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
-
-from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
-
-
-NV12Frame = namedtuple("NV12Frame", ['width', 'height', 'stride', 'y_height', 'uv_height', 'size'])
-
-
-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]
-  return 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))
-
-
-def make_frame_prepare(nv12: NV12Frame, model_w, model_h, layout="yuv420", border_fill=None):
-  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):
-    if layout == "luma":
-      return warp_perspective_tinygrad(input_frame[:cam_h*stride], M_inv,
-                                      (model_w, model_h), (cam_h, cam_w), stride_pad,
-                                      border_fill_val=border_fill).reshape(-1, model_h * model_w)
-    # 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=Device.DEFAULT)
-    # 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, border_fill_val=border_fill).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, border_fill_val=border_fill).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, border_fill_val=border_fill).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, model_w, model_h, layout="yuv420", border_fill=None, frames=1):
-  frame_prepare = make_frame_prepare(nv12, model_w, model_h, layout, border_fill)
-
-  def warp(input_frames, transforms):
-    input_frames = input_frames.to(Device.DEFAULT)
-    transforms = transforms.to(Device.DEFAULT)
-    Tensor.realize(input_frames, transforms)
-    if frames == 1:
-      return frame_prepare(input_frames, transforms)
-    return Tensor.stack(*(frame_prepare(input_frames[i], transforms[i]) for i in range(frames)))
-
-  return warp
-
-
-def _parse_size(s):
-  w, h = s.lower().split('x')
-  return int(w), int(h)
-
-
-def compile_warp(nv12: NV12Frame, model_w, model_h, pkl_path, layout, border_fill=None,
-                 frames=1, transform_device=None, frame_size=None):
-  print(f"Compiling {layout} warp for {nv12.width}x{nv12.height} -> {model_w}x{model_h}...")
-
-  warp_jit = TinyJit(make_warp(nv12, model_w, model_h, layout, border_fill, frames), prune=True)
-  frame_size = nv12.size if frame_size is None else frame_size
-  frame_shape = (frame_size,) if frames == 1 else (frames, frame_size)
-  transform_shape = (3, 3) if frames == 1 else (frames, 3, 3)
-
-  for i in range(10):
-    frame = Tensor.randint(*frame_shape, low=0, high=256, dtype='uint8').realize()
-    M_inv = Tensor(Tensor.randn(*transform_shape).mul(8).realize().numpy(), device=transform_device)
-    Device.default.synchronize()
-    st = time.perf_counter()
-    warp_jit(frame, M_inv).realize()
-    mt = time.perf_counter()
-    Device.default.synchronize()
-    et = time.perf_counter()
-    print(f"  [{i+1}/10] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms")
-
-  with open(pkl_path, "wb") as f:
-    pickle.dump(warp_jit, f)
-  print(f"  Saved to {pkl_path}")
-
-
-if __name__ == "__main__":
-  p = argparse.ArgumentParser()
-  p.add_argument('--camera-resolution', type=_parse_size, required=True, help='camera resolution WxH')
-  p.add_argument('--warp-to', type=_parse_size, required=True, help='output WxH')
-  p.add_argument('--layout', choices=['luma', 'yuv420'], required=True)
-  p.add_argument('--border-fill', type=int, help='fill value outside the frame; omit to clamp coordinates')
-  p.add_argument('--frames', type=int, default=1, help='number of frames to warp together')
-  p.add_argument('--transform-device', help='device holding the input transforms; default: compute device')
-  p.add_argument('--frame-size', type=int, help='input frame size in bytes; default: full NV12 allocation')
-  p.add_argument('--output', required=True)
-  args = p.parse_args()
-
-  cam_w, cam_h = args.camera_resolution
-  nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
-  model_w, model_h = args.warp_to
-  compile_warp(nv12, model_w, model_h, args.output, args.layout, args.border_fill,
-               args.frames, args.transform_device, args.frame_size)

openpilot/selfdrive/modeld/dmonitoringmodeld.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import os
-from openpilot.selfdrive.modeld.helpers import MODELS_DIR, get_tg_input_devices
+from openpilot.selfdrive.modeld.helpers import MODELS_DIR, load_oob
 from tinygrad.tensor import Tensor
 import time
 import pickle
@@ -18,10 +18,8 @@ from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
 from openpilot.common.file_chunker import open_file_chunked
 from openpilot.selfdrive.modeld.parse_model_outputs import sigmoid, safe_exp
 
-PROCESS_NAME = "openpilot.selfdrive.modeld.dmonitoringmodeld"
 SEND_RAW_PRED = os.getenv('SEND_RAW_PRED')
 MODEL_PKL_PATH = MODELS_DIR / 'dmonitoring_model_tinygrad.pkl'
-METADATA_PATH = MODELS_DIR / 'dmonitoring_model_metadata.pkl'
 
 
 class ModelState:
@@ -29,11 +27,10 @@ class ModelState:
   output: np.ndarray
 
   def __init__(self, cam_w: int, cam_h: int):
-    self.DEV = get_tg_input_devices(PROCESS_NAME, chestnut=False)['DEV']
-    with open(METADATA_PATH, 'rb') as f:
-      model_metadata = pickle.load(f)
-      self.input_shapes = model_metadata['input_shapes']
-      self.output_slices = model_metadata['output_slices']
+    jits = load_oob(open_file_chunked(MODEL_PKL_PATH))
+    self.DEV = jits['input_devices']['model']
+    self.input_shapes = jits['metadata']['input_shapes']
+    self.output_slices = jits['metadata']['output_slices']
 
     self.numpy_inputs = {
       'calib': np.zeros(self.input_shapes['calib'], dtype=np.float32),
@@ -42,9 +39,10 @@ class ModelState:
     self.warp_inputs_np = {'transform': np.zeros((3,3), dtype=np.float32)}
     self.warp_inputs = {k: Tensor(v, device='NPY') for k,v in self.warp_inputs_np.items()}
     self.frame_buf_params = get_nv12_info(cam_w, cam_h)
-    self.tensor_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items()}
+    self.tensor_inputs = {k: Tensor(v, device=self.DEV).realize() for k,v in self.numpy_inputs.items()}
+    self.calib_host = Tensor(self.numpy_inputs['calib'], device='NPY')._buffer()
     self._blob_cache : dict[int, Tensor] = {}
-    self.model_run = pickle.load(open_file_chunked(str(MODEL_PKL_PATH)))
+    self.model_run = jits['run_model']
     with open(MODELS_DIR / f'dm_warp_{cam_w}x{cam_h}_tinygrad.pkl', "rb") as f:
       self.image_warp = pickle.load(f)
 
@@ -52,6 +50,7 @@ class ModelState:
     self.numpy_inputs['calib'][0,:] = calib
 
     t1 = time.perf_counter()
+    self.tensor_inputs['calib']._buffer().copy_from(self.calib_host)
 
     ptr = np.frombuffer(buf.data, dtype=np.uint8).ctypes.data
     # There is a ringbuffer of imgs, just cache tensors pointing to all of them
@@ -61,7 +60,8 @@ class ModelState:
     self.warp_inputs_np['transform'][:] = transform[:]
     self.tensor_inputs['input_img'] = self.image_warp(self._blob_cache[ptr], self.warp_inputs['transform'])
 
-    output = self.model_run(**self.tensor_inputs).numpy().flatten()
+    output, = self.model_run(**self.tensor_inputs)
+    output = output.numpy().astype(np.float32).reshape(-1)
 
     t2 = time.perf_counter()
     return output, t2 - t1

openpilot/selfdrive/modeld/get_model_metadata.py

@@ -1,55 +0,0 @@
-#!/usr/bin/env python3
-import sys
-import pathlib
-import codecs
-import pickle
-from typing import Any
-
-from tinygrad.nn.onnx import OnnxPBParser
-
-
-class MetadataOnnxPBParser(OnnxPBParser):
-  def _parse_ModelProto(self) -> dict:
-    obj: dict[str, Any] = {"graph": {"input": [], "output": []}, "metadata_props": []}
-    for fid, wire_type in self._parse_message(self.reader.len):
-      match fid:
-        case 7:
-          obj["graph"] = self._parse_GraphProto()
-        case 14:
-          obj["metadata_props"].append(self._parse_StringStringEntryProto())
-        case _:
-          self.reader.skip_field(wire_type)
-    return obj
-
-
-def get_name_and_shape(value_info: dict[str, Any]) -> tuple[str, tuple[int, ...]]:
-  shape = tuple(int(dim) if isinstance(dim, int) else 0 for dim in value_info["parsed_type"].shape)
-  name = value_info["name"]
-  return name, shape
-
-
-def get_metadata_value_by_name(model: dict[str, Any], name: str) -> str | Any:
-  for prop in model["metadata_props"]:
-    if prop["key"] == name:
-      return prop["value"]
-  return None
-
-
-def make_metadata_dict(model_path):
-  model = MetadataOnnxPBParser(model_path).parse()
-  output_slices = get_metadata_value_by_name(model, 'output_slices')
-  assert output_slices is not None, 'output_slices not found in metadata'
-  return {
-    'model_checkpoint': get_metadata_value_by_name(model, 'model_checkpoint'),
-    'output_slices': pickle.loads(codecs.decode(output_slices.encode(), "base64")),
-    'input_shapes': dict(get_name_and_shape(x) for x in model["graph"]["input"]),
-    'output_shapes': dict(get_name_and_shape(x) for x in model["graph"]["output"]),
-  }
-
-
-if __name__ == "__main__":
-  model_path = pathlib.Path(sys.argv[1])
-  metadata_path = model_path.parent / (model_path.stem + '_metadata.pkl')
-  with open(metadata_path, 'wb') as f:
-    pickle.dump(make_metadata_dict(model_path), f)
-  print(f'saved metadata to {metadata_path}')

openpilot/selfdrive/modeld/helpers.py

@@ -1,56 +1,18 @@
 import io
-import json
 import pickle
-import shutil
 import struct
-import tempfile
 from pathlib import Path
 
 from openpilot.common.file_chunker import get_manifest_path
 from openpilot.common.hardware.usb import CHESTNUT_USB_PRODUCT, USB_DEVICES_PATH, is_chestnut_usb_id
 
 MODELS_DIR = Path(__file__).resolve().parent / 'models'
-TG_INPUT_DEVICES_PATH = MODELS_DIR / 'tg_input_devices.json'
 
 
-def patch_tinygrad_fetch_fw():
-  import hashlib
-  import zstandard
-  from tinygrad import helpers
-  original_fetch_fw = helpers.fetch_fw
-  def fetch_fw(path, name, sha256):
-    p = 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 original_fetch_fw(path, name, sha256)
-  helpers.fetch_fw = fetch_fw
-
-
-def get_tg_input_devices(process_name: str, chestnut: bool):
-  with open(TG_INPUT_DEVICES_PATH) as f:
-    return json.load(f)[process_name]['default' if not chestnut else 'chestnut']
-
 def modeld_pkl_path(chestnut: bool):
   prefix = 'big_' if chestnut else ''
   return MODELS_DIR / f'{prefix}driving_tinygrad.pkl'
 
-def dump_oob(obj, f):
-  with tempfile.TemporaryFile(dir=".") as tmp:
-    def buffer_callback(pb: pickle.PickleBuffer):
-      m = pb.raw()
-      tmp.write(struct.pack('<q', m.nbytes))
-      tmp.write(m)
-      pb.release() # keep peak ram at ~1 buffer
-    stream = io.BytesIO()
-    pickle.Pickler(stream, protocol=5, buffer_callback=buffer_callback).dump(obj)
-    opcodes = stream.getvalue()
-    f.write(struct.pack('<q', len(opcodes)))
-    f.write(opcodes)
-    tmp.seek(0)
-    shutil.copyfileobj(tmp, f)
-
 def load_oob(f):
   opcodes = f.read(struct.unpack('<q', f.read(8))[0])
   def buffers():

openpilot/selfdrive/modeld/modeld.py

@@ -32,7 +32,6 @@ 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
 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.selfdrive.modeld.constants import ModelConstants, Plan
@@ -159,7 +158,8 @@ class ModelState:
 
   def pack_inputs(self) -> None:
     # Pack host inputs into one upload to reduce USB transfer overhead for the eGPU.
-    self.input_queues = make_input_queues({name: self.input_shapes[name] for name in self.state_pairs}, self.model_device)
+    self.input_queues = {name: Tensor(np.zeros(shape, dtype=dtype.fmt), device=self.model_device).realize()
+                         for name, (shape, dtype) in self.input_shapes.items() if name in self.state_pairs}
     shapes = {'tfm': (2, 3, 3)} | {name: shape for name, (shape, _) in self.input_shapes.items()
                                    if name not in self.state_pairs and name != 'new_img'}
     npy_size = sum(round_up(math.prod(shape) * 4, 128) for shape in shapes.values())