Use tinygrad generic ONNX compiler artifacts (#38926)

* Use tinygrad generic ONNX compiler artifacts

* Keep compiler temporary model buffers off device tmpfs

* Keep recurrent model state updates on the GPU

* Use compiled output buffers and swap recurrent state inputs

* Note when ONNX chunk reassembly can be removed

* Bump tinygrad to rebased compilers and AMD queue fix

* Keep recurrent buffers fixed and use faster USB argument updates

* Bump tinygrad for faster model startup

.gitmodules

@@ -16,4 +16,4 @@
 [submodule "tinygrad"]
   path = tinygrad_repo
   url = https://github.com/commaai/tinygrad.git
-  branch = openpilot-modeld-test
+  branch = model-warp-compile

openpilot/selfdrive/modeld/SConscript

@@ -1,8 +1,10 @@
 import glob
 import os
+import shutil
+import tempfile
 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.file_chunker import chunk_file, get_chunk_targets, get_existing_chunks, open_file_chunked
 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 chestnut_present, modeld_pkl_path
@@ -41,7 +43,7 @@ if CHESTNUT:
 mac_brew_string = f'HOME={os.path.expanduser("~")}' if arch == 'Darwin' else ''
 
 warp_deps = [File("#openpilot/system/camerad/cameras/nv12_info.py")]
-compiler = 'python3 -m examples.openpilot'
+compiler = Dir('#tinygrad_repo/examples/openpilot').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 ''
 
@@ -65,13 +67,22 @@ def chestnut_action(command, pkl=None, chunks=()):
 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}')
+  cmd = (f'{flags} {mac_brew_string} {taskset}python3 "{compiler}/compile_onnx.py" '
+         f'"{{onnx}}" "{target_pkl_path}" --device-input "*" --out-of-band --benchmark-runs 1')
+  def do_compile(target, source, env):
+    if os.path.isfile(onnx_path):
+      return env.Execute(cmd.format(onnx=onnx_path))
+    # TODO: Remove ONNX chunk reassembly once models are precompiled.
+    with tempfile.NamedTemporaryFile(dir=os.path.dirname(onnx_path), suffix='.onnx') as tmp, open_file_chunked(onnx_path) as src:
+      shutil.copyfileobj(src, tmp)
+      tmp.flush()
+      return env.Execute(cmd.format(onnx=tmp.name))
+  compile_action = Action(do_compile, "  [ONNX] $TARGET")
   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 = chestnut_action(cmd, target_pkl_path, chunk_targets) if chestnut else [cmd, Action(do_chunk, "  [CHUNK] $TARGET")]
+  actions = chestnut_action(compile_action, target_pkl_path, chunk_targets) if chestnut else [compile_action, Action(do_chunk, "  [CHUNK] $TARGET")]
   node = lenv.Command(
     chunk_targets,
     tinygrad_files + onnx_deps + [Value(cmd), Value(chunk_targets), chunker_file],
@@ -89,9 +100,9 @@ for chestnut in [False, True] if CHESTNUT else [False]:
   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}{compiler}.compile_warp '
-           f'--camera-resolution {cam_w}x{cam_h} --warp-to {model_w}x{model_h} --layout yuv420 '
-           f'--frames 2 --stride {stride} --uv-offset {stride * y_height} --frame-size {stride * (y_height + uv_height)} '
+    cmd = (f'{cmd_flags} {mac_brew_string} {taskset}python3 "{compiler}/compile_warp.py" '
+           f'--frame {cam_w},{cam_h},{stride},{y_height},{uv_height},{stride * (y_height + uv_height)} '
+           f'--warp-to {model_w}x{model_h} --layout yuv420 --frames 2 '
            f'--output {warp_pkl_path}')
     action = chestnut_action(cmd) if chestnut else cmd
     node = lenv.Command(warp_pkl_path, tinygrad_files + warp_deps + [Value(cmd)], action)
@@ -101,8 +112,8 @@ for chestnut in [False, True] if CHESTNUT else [False]:
 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
-  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'--stride {stride} --uv-offset {stride * y_height} --frame-size {frame_size} --output {dm_pkl_path}')
+  stride, y_height, uv_height, frame_size = get_nv12_info(cam_w, cam_h)
+  cmd = (f'{tg_flags} {mac_brew_string} python3 "{compiler}/compile_warp.py" '
+         f'--frame {cam_w},{cam_h},{stride},{y_height},{uv_height},{frame_size} --warp-to {dm_w}x{dm_h} '
+         f'--layout luma --border-fill 16 --transform-device NPY --output {dm_pkl_path}')
   lenv.Command(dm_pkl_path, tinygrad_files + warp_deps + [Value(cmd)], cmd)

openpilot/selfdrive/modeld/dmonitoringmodeld.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 import os
+import base64
 from openpilot.selfdrive.modeld.helpers import MODELS_DIR, load_oob
 from tinygrad.tensor import Tensor
 import time
@@ -28,9 +29,9 @@ class ModelState:
 
   def __init__(self, cam_w: int, cam_h: int):
     jits = load_oob(open_file_chunked(MODEL_PKL_PATH))
-    self.DEV = jits['input_devices']['model']
+    self.DEV = jits['input_specs']['input_img'][2]
     self.input_shapes = jits['metadata']['input_shapes']
-    self.output_slices = jits['metadata']['output_slices']
+    self.output_slices = pickle.loads(base64.b64decode(jits['metadata']['metadata']['output_slices']))
 
     self.numpy_inputs = {
       'calib': np.zeros(self.input_shapes['calib'], dtype=np.float32),
@@ -42,9 +43,10 @@ class ModelState:
     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 = jits['run_model']
+    self.model_run = jits['run']
+    self.outputs = {name: Tensor(np.zeros(shape, dtype=dtype), device=device).realize() for name, (shape, dtype, device) in jits['output_specs'].items()}
     with open(MODELS_DIR / f'dm_warp_{cam_w}x{cam_h}_tinygrad.pkl', "rb") as f:
-      self.image_warp = pickle.load(f)
+      self.image_warp = pickle.load(f)['run']
 
   def run(self, buf: VisionBuf, calib: np.ndarray, transform: np.ndarray) -> tuple[np.ndarray, float]:
     self.numpy_inputs['calib'][0,:] = calib
@@ -58,10 +60,10 @@ class ModelState:
       self._blob_cache[ptr] = Tensor.from_blob(ptr, (self.frame_buf_params[3],), dtype='uint8', device=self.DEV)
 
     self.warp_inputs_np['transform'][:] = transform[:]
-    self.tensor_inputs['input_img'] = self.image_warp(self._blob_cache[ptr], self.warp_inputs['transform'])
+    self.tensor_inputs['input_img'] = self.image_warp(input_frame=self._blob_cache[ptr], M_inv=self.warp_inputs['transform'])
 
-    output, = self.model_run(**self.tensor_inputs)
-    output = output.numpy().astype(np.float32).reshape(-1)
+    self.model_run(output_buffers=self.outputs, **self.tensor_inputs)
+    output = self.outputs['outputs'].numpy().astype(np.float32).reshape(-1)
 
     t2 = time.perf_counter()
     return output, t2 - t1

openpilot/selfdrive/modeld/modeld.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 from collections.abc import Callable
+import base64
 import ctypes
 from functools import cached_property
 import os
@@ -139,11 +140,11 @@ class ModelState:
 
   def __init__(self, cam_w: int, cam_h: int, chestnut: bool):
     jits = load_oob(open_file_chunked(modeld_pkl_path(chestnut)))
-    self.model_device = jits['input_devices']['model']
-    self.input_shapes = jits['input_shapes']
-    self.state_pairs = jits['state_pairs']
+    self.model_device = jits['input_specs']['new_img'][2]
+    self.input_shapes = {name: (shape, np.dtype(dtype)) for name, (shape, dtype, _) in jits['input_specs'].items()}
+    self.state_pairs = {name: f'next_{name}' for name in self.input_shapes if f'next_{name}' in jits['metadata']['output_shapes']}
     self.vision_input_names = ('img', 'big_img')
-    self.output_slices = jits['metadata']['output_slices']
+    self.output_slices = pickle.loads(base64.b64decode(jits['metadata']['metadata']['output_slices']))
 
     self.prev_desire = np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32)
     self.chestnut = chestnut
@@ -152,13 +153,17 @@ class ModelState:
     self.frame_copy_size = stride * (y_height + uv_height)
     self.pack_inputs()
     with open(MODELS_DIR / f'{"big_" if chestnut else ""}driving_warp_{cam_w}x{cam_h}_tinygrad.pkl', 'rb') as f:
-      self.run_warp = pickle.load(f)
-    self.run_model = jits['run_model']
+      self.run_warp = pickle.load(f)['run']
+    self.run_model = jits['run']
+    self.outputs = {name: Tensor(np.zeros(shape, dtype=dtype), device=device).realize() for name, (shape, dtype, device) in jits['output_specs'].items()}
+    for name, next_name in self.state_pairs.items():
+      state = self.input_queues[name]
+      self.outputs[next_name] = input_view(state._buffer(), state.shape, state.dtype, 0)
     self.parser = Parser()
 
   def pack_inputs(self) -> None:
     # Pack host inputs into one upload to reduce USB transfer overhead for the eGPU.
-    self.input_queues = {name: Tensor(np.zeros(shape, dtype=dtype.fmt), device=self.model_device).realize()
+    self.input_queues = {name: Tensor(np.zeros(shape, dtype=dtype), 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'}
@@ -173,7 +178,7 @@ class ModelState:
       self.input_queues[name] = input_view(self.input_device, shape, dtypes.float32, offset)
       offset += round_up(self.npy[name].nbytes, 128)
     self.frames = self.packed_input[npy_size:].reshape(2, self.frame_copy_size)
-    self.warp_inputs = (input_view(self.input_device, self.frames.shape, dtypes.uint8, npy_size), self.input_queues.pop('tfm'))
+    self.warp_inputs = {'input_frame': input_view(self.input_device, self.frames.shape, dtypes.uint8, npy_size), 'M_inv': self.input_queues.pop('tfm')}
 
   def slice_outputs(self, model_outputs: np.ndarray, output_slices: dict[str, slice]) -> dict[str, np.ndarray]:
     return {k: model_outputs[np.newaxis, v] for k,v in output_slices.items()}
@@ -192,11 +197,11 @@ class ModelState:
     self.npy['action_t'][:] = inputs['action_t']
 
     self.input_device.copy_from(self.input_host)
-    self.input_queues['new_img'] = self.run_warp(*self.warp_inputs)
-    outs, = self.run_model(**self.input_queues)
+    self.input_queues['new_img'] = self.run_warp(**self.warp_inputs)
+    self.run_model(output_buffers=self.outputs, **self.input_queues)
     if after_enqueue is not None:
       after_enqueue()
-    model_output = outs.numpy()[0]
+    model_output = self.outputs['outputs'].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))