modeld: DEV=AMD dedupe weights across camera resolutions (#38041)

* modeld: dedupe weight accross resolutions

* cleanup

* rm compileconfig

* depends on camera targets

* dedupe doesn't work on qcom as is

selfdrive/modeld/SConscript

@@ -1,13 +1,11 @@
 import glob
 import json
 import os
-from itertools import product
 from SCons.Script import Value
 from openpilot.common.file_chunker import chunk_file, get_chunk_paths
 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.constants import ModelConstants
-from openpilot.selfdrive.modeld.helpers import CompileConfig
 from tinygrad import Device
 from openpilot.system.hardware import HARDWARE, PC
 
@@ -26,9 +24,6 @@ def get_camera_configs():
 
 CAMERA_CONFIGS = get_camera_configs()
 
-MODELD_CONFIGS = [CompileConfig(cam_w, cam_h, prepare_only, 'driving_')
-                  for (cam_w, cam_h), prepare_only in product(CAMERA_CONFIGS, [True, False])]
-DM_WARP_CONFIGS = [CompileConfig(cam_w, cam_h, True, 'dm_') for cam_w, cam_h in CAMERA_CONFIGS]
 chunker_file = File("#common/file_chunker.py")
 lenv = env.Clone()
 
@@ -80,27 +75,28 @@ driving_metadata_deps = [File(f"models/{m}_metadata.pkl").abspath for m in ['dri
 
 model_w, model_h = MEDMODEL_INPUT_SIZE
 frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ
-for cfg in MODELD_CONFIGS:
-  cmd = (f'{tg_flags} {mac_brew_string} python3 {modeld_dir}/compile_modeld.py '
-         f'--model-size {model_w}x{model_h} '
-         f'--nv12 {",".join(str(x) for x in cfg.nv12)} '
-         f'--vision-onnx {File("models/driving_vision.onnx").abspath} '
-         f'--policy-onnx {File("models/driving_policy.onnx").abspath} '
-         f'--output {cfg.pkl_path} --frame-skip {frame_skip}'
-         + (' --prepare-only' if cfg.prepare_only else ''))
-  node = lenv.Command(cfg.pkl_path, tinygrad_files + compile_modeld_script + driving_onnx_deps + driving_metadata_deps + [chunker_file, compiled_flags_node], cmd)
-  onnx_sizes_sum = sum(os.path.getsize(f) for f in driving_onnx_deps)
-  chunk_targets = get_chunk_paths(cfg.pkl_path, estimate_pickle_max_size(onnx_sizes_sum))
-  def do_chunk(target, source, env, pkl=cfg.pkl_path, chunks=chunk_targets):
-    chunk_file(pkl, chunks)
-  lenv.Command(chunk_targets, node, do_chunk)
+pkl_path = File("models/driving_tinygrad.pkl").abspath
+camera_res_args = ' '.join(f'{cw}x{ch}' for cw, ch in CAMERA_CONFIGS)
+cmd = (f'{tg_flags} {mac_brew_string} python3 {modeld_dir}/compile_modeld.py '
+       f'--model-size {model_w}x{model_h} '
+       f'--camera-resolutions {camera_res_args} '
+       f'--vision-onnx {File("models/driving_vision.onnx").abspath} '
+       f'--policy-onnx {File("models/driving_policy.onnx").abspath} '
+       f'--output {pkl_path} --frame-skip {frame_skip}')
+node = lenv.Command(pkl_path, tinygrad_files + compile_modeld_script + driving_onnx_deps + driving_metadata_deps + [Value(camera_res_args), chunker_file, compiled_flags_node], cmd)
+onnx_sizes_sum = sum(os.path.getsize(f) for f in driving_onnx_deps)
+chunk_targets = get_chunk_paths(pkl_path, estimate_pickle_max_size(onnx_sizes_sum)*2) # TODO make weight dedupe work on QCOM
+def do_chunk(target, source, env, pkl=pkl_path, chunks=chunk_targets):
+  chunk_file(pkl, chunks)
+lenv.Command(chunk_targets, node, do_chunk)
 
 dm_w, dm_h = DM_INPUT_SIZE
-for cfg in DM_WARP_CONFIGS:
+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_dm_warp.py '
-         f'--nv12 {",".join(str(x) for x in cfg.nv12)} --warp-to {dm_w}x{dm_h} '
-         f'--output {cfg.pkl_path}')
-  lenv.Command(cfg.pkl_path, tinygrad_files + compile_dm_warp_script + compile_modeld_script + [compiled_flags_node], cmd)
+         f'--camera-resolution {cam_w}x{cam_h} --warp-to {dm_w}x{dm_h} '
+         f'--output {dm_pkl_path}')
+  lenv.Command(dm_pkl_path, tinygrad_files + compile_dm_warp_script + compile_modeld_script + [compiled_flags_node], cmd)
 
 def tg_compile(flags, model_name):
   pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"'

selfdrive/modeld/compile_dm_warp.py

@@ -7,7 +7,8 @@ from tinygrad.tensor import Tensor
 from tinygrad.device import Device
 from tinygrad.engine.jit import TinyJit
 
-from openpilot.selfdrive.modeld.compile_modeld import NV12Frame, warp_perspective_tinygrad, _parse_size, _parse_nv12
+from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
+from openpilot.selfdrive.modeld.compile_modeld import NV12Frame, warp_perspective_tinygrad, _parse_size
 
 
 def make_warp_dm(nv12: NV12Frame, dm_w, dm_h):
@@ -44,11 +45,12 @@ def compile_dm_warp(nv12: NV12Frame, dm_w, dm_h, pkl_path):
 
 if __name__ == "__main__":
   p = argparse.ArgumentParser()
-  p.add_argument('--nv12', type=_parse_nv12, required=True,
-                 help=f'NV12 frame layout: {",".join(NV12Frame._fields)}')
+  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='DM input WxH')
   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))
   dm_w, dm_h = args.warp_to
-  compile_dm_warp(args.nv12, dm_w, dm_h, args.output)
+  compile_dm_warp(nv12, dm_w, dm_h, args.output)

selfdrive/modeld/compile_modeld.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 import argparse
+import os
 import pickle
 import time
 from functools import partial
@@ -10,7 +11,6 @@ from tinygrad.tensor import Tensor
 from tinygrad.helpers import Context
 from tinygrad.device import Device
 from tinygrad.engine.jit import TinyJit
-from tinygrad.nn.onnx import OnnxRunner
 
 
 NV12Frame = namedtuple("NV12Frame", ['width', 'height', 'stride', 'y_height', 'uv_height', 'size'])
@@ -158,21 +158,10 @@ def make_run_policy(vision_runner, policy_runner, nv12: NV12Frame, model_w, mode
 
 
 def compile_modeld(nv12: NV12Frame, model_w, model_h, prepare_only, frame_skip,
-                   vision_onnx, policy_onnx, pkl_path):
-  from get_model_metadata import metadata_path_for
-
+                   vision_runner, policy_runner, vision_features_slice,
+                   vision_input_shapes, policy_input_shapes):
   print(f"Compiling combined policy JIT for {nv12.width}x{nv12.height} (prepare_only={prepare_only})...")
 
-  vision_runner = OnnxRunner(vision_onnx)
-  policy_runner = OnnxRunner(policy_onnx)
-
-  with open(metadata_path_for(vision_onnx), 'rb') as f:
-    vision_metadata = pickle.load(f)
-    vision_features_slice = vision_metadata['output_slices']['hidden_state']
-    vision_input_shapes = vision_metadata['input_shapes']
-  with open(metadata_path_for(policy_onnx), 'rb') as f:
-    policy_input_shapes = pickle.load(f)['input_shapes']
-
   _run = make_run_policy(vision_runner, policy_runner, nv12, model_w, model_h,
                          vision_features_slice, frame_skip, prepare_only)
   run_policy_jit = TinyJit(_run, prune=True)
@@ -216,13 +205,10 @@ def compile_modeld(nv12: NV12Frame, model_w, model_h, prepare_only, frame_skip,
   run_policy_jit, test_val, test_buffers = random_inputs_run_fn(run_policy_jit, SEED)
 
   print('pickle round trip')
-  with open(pkl_path, "wb") as f:
-    pickle.dump(run_policy_jit, f)
-    print(f"  Saved to {pkl_path}")
-  with open(pkl_path, "rb") as f:
-    run_policy_jit = pickle.load(f)
+  run_policy_jit = pickle.loads(pickle.dumps(run_policy_jit))
   random_inputs_run_fn(run_policy_jit, SEED, test_val, test_buffers, expect_match=True)
   random_inputs_run_fn(run_policy_jit, SEED+1, test_val, test_buffers, expect_match=False)
+  return run_policy_jit
 
 
 def _parse_size(s):
@@ -230,25 +216,42 @@ def _parse_size(s):
   return int(w), int(h)
 
 
-def _parse_nv12(s):
-  parts = s.split(',')
-  assert len(parts) == len(NV12Frame._fields), \
-    f"--nv12 expects {','.join(NV12Frame._fields)} (got {s!r})"
-  return NV12Frame(*(int(x) for x in parts))
-
-
 if __name__ == "__main__":
+  from tinygrad.nn.onnx import OnnxRunner
+  from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
   p = argparse.ArgumentParser()
   p.add_argument('--model-size', type=_parse_size, required=True, help='model input WxH')
-  p.add_argument('--nv12', type=_parse_nv12, required=True,
-                 help=f'NV12 frame layout: {",".join(NV12Frame._fields)}')
+  p.add_argument('--camera-resolutions', type=_parse_size, nargs='+', required=True,
+                 help='camera resolutions WxH (one or more)')
   p.add_argument('--vision-onnx', required=True)
   p.add_argument('--policy-onnx', required=True)
   p.add_argument('--output', required=True)
-  p.add_argument('--prepare-only', action='store_true')
   p.add_argument('--frame-skip', type=int, required=True)
   args = p.parse_args()
 
   model_w, model_h = args.model_size
-  compile_modeld(args.nv12, model_w, model_h, args.prepare_only, args.frame_skip,
-                 args.vision_onnx, args.policy_onnx, args.output)
+
+  # init runners once so weights are shared
+  from get_model_metadata import metadata_path_for
+  vision_runner = OnnxRunner(args.vision_onnx)
+  policy_runner = OnnxRunner(args.policy_onnx)
+  with open(metadata_path_for(args.vision_onnx), 'rb') as f:
+    vision_metadata = pickle.load(f)
+    vision_features_slice = vision_metadata['output_slices']['hidden_state']
+    vision_input_shapes = vision_metadata['input_shapes']
+  with open(metadata_path_for(args.policy_onnx), 'rb') as f:
+    policy_input_shapes = pickle.load(f)['input_shapes']
+
+  out = {}
+  for cam_w, cam_h in args.camera_resolutions:
+    nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
+    out[(cam_w,cam_h)] = {
+      name: compile_modeld(nv12, model_w, model_h, prepare_only, args.frame_skip,
+                           vision_runner, policy_runner, vision_features_slice,
+                           vision_input_shapes, policy_input_shapes)
+      for name, prepare_only in [('warp_enqueue', True), ('run_policy', False)]
+    }
+
+  with open(args.output, "wb") as f:
+    pickle.dump(out, f)
+  print(f"Saved combined JIT to {args.output} ({os.path.getsize(args.output) / 1e6:.2f} MB)")

selfdrive/modeld/dmonitoringmodeld.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import os
-from openpilot.selfdrive.modeld.helpers import MODELS_DIR, CompileConfig, set_tinygrad_backend_from_compiled_flags
+from openpilot.selfdrive.modeld.helpers import MODELS_DIR, set_tinygrad_backend_from_compiled_flags
 set_tinygrad_backend_from_compiled_flags()
 
 from tinygrad.tensor import Tensor
@@ -44,7 +44,7 @@ class ModelState:
     self.tensor_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items()}
     self._blob_cache : dict[int, Tensor] = {}
     self.model_run = pickle.loads(read_file_chunked(str(MODEL_PKL_PATH)))
-    with open(CompileConfig(cam_w, cam_h, prefix='dm_', prepare_only=True).pkl_path, "rb") as f:
+    with open(MODELS_DIR / f'dm_warp_{cam_w}x{cam_h}_tinygrad.pkl', "rb") as f:
       self.image_warp = pickle.load(f)
 
   def run(self, buf: VisionBuf, calib: np.ndarray, transform: np.ndarray) -> tuple[np.ndarray, float]:

selfdrive/modeld/helpers.py

@@ -1,10 +1,7 @@
 import json
 import os
-from dataclasses import dataclass
 from pathlib import Path
 
-from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
-
 MODELS_DIR = Path(__file__).resolve().parent / 'models'
 COMPILED_FLAGS_PATH = MODELS_DIR / 'tg_compiled_flags.json'
 
@@ -13,19 +10,3 @@ def set_tinygrad_backend_from_compiled_flags() -> None:
   if os.path.isfile(COMPILED_FLAGS_PATH):
     with open(COMPILED_FLAGS_PATH) as f:
       os.environ['DEV'] = str(json.load(f)['DEV'])
-
-
-@dataclass
-class CompileConfig:
-  cam_w: int
-  cam_h: int
-  prepare_only: bool
-  prefix: str
-
-  @property
-  def pkl_path(self):
-    return str(MODELS_DIR / f'{self.prefix}{"warp_" if self.prepare_only else ""}{self.cam_w}x{self.cam_h}_tinygrad.pkl')
-
-  @property
-  def nv12(self):
-    return (self.cam_w, self.cam_h, *get_nv12_info(self.cam_w, self.cam_h))

selfdrive/modeld/modeld.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import os
-from openpilot.selfdrive.modeld.helpers import MODELS_DIR, CompileConfig, set_tinygrad_backend_from_compiled_flags
+from openpilot.selfdrive.modeld.helpers import MODELS_DIR, set_tinygrad_backend_from_compiled_flags
 set_tinygrad_backend_from_compiled_flags()
 
 USBGPU = "USBGPU" in os.environ
@@ -100,8 +100,9 @@ class ModelState:
     self._blob_cache : dict[int, Tensor] = {}
     self.parser = Parser()
     self.frame_buf_params = {k: get_nv12_info(cam_w, cam_h) for k in ('img', 'big_img')}
-    self.run_policy = pickle.loads(read_file_chunked(CompileConfig(cam_w, cam_h, prefix='driving_', prepare_only=False).pkl_path))
-    self.warp_enqueue = pickle.loads(read_file_chunked(CompileConfig(cam_w, cam_h, prefix='driving_', prepare_only=True).pkl_path))
+    jits = pickle.loads(read_file_chunked(MODELS_DIR / 'driving_tinygrad.pkl'))[(cam_w,cam_h)]
+    self.run_policy = jits['run_policy']
+    self.warp_enqueue = jits['warp_enqueue']
     self.warp_enqueue(
       **self.input_queues,
       frame=Tensor.zeros(self.frame_buf_params['img'][3], dtype='uint8').contiguous().realize(),