modeld: move history queues into ONNX (#38916)

* modeld: move history queues into ONNX * Move ONNX history export into xx exporter * Rename ONNX image input to new_img * Combine ONNX model history state * Revert model history README addition * Keep float32 input packing from master * Rename model input specs to input shapes
2026-09-18 10:23:43 +08:00 · 2026-09-14 20:44:33 -07:00
parent 7b469af438
commit 64f9b47b6e
5 changed files with 42 additions and 105 deletions
@@ -6,7 +6,6 @@ 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.constants import ModelConstants
 from openpilot.selfdrive.modeld.helpers import TG_INPUT_DEVICES_PATH, chestnut_present, modeld_pkl_path


@@ -68,7 +67,6 @@ compile_modeld_script = [
  File("#openpilot/common/hardware/hw.py"),
 ]
 model_w, model_h = MEDMODEL_INPUT_SIZE
-frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ

 for chestnut in [False, True] if CHESTNUT else [False]:
  target_pkl_path = File(modeld_pkl_path(chestnut)).abspath
@@ -81,7 +79,7 @@ for chestnut in [False, True] if CHESTNUT else [False]:
        f'--model-size {model_w}x{model_h} '
        f'--camera-resolutions {camera_res_args} '
        f'--onnx {File(f"models/{file_prefix}driving_supercombo.onnx").abspath} '
-        f'--output {target_pkl_path} --frame-skip {frame_skip}')
+        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 * len(camera_configs)))
  def do_compile(target, source, env, command=cmd, pkl=target_pkl_path, chunks=chunk_targets):
@@ -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,23 +248,24 @@ 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)
@@ -26,7 +26,7 @@ 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.selfdrive.modeld.constants import ModelConstants, Plan
@@ -132,17 +132,17 @@ class ModelState:
    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)]

@@ -164,14 +164,13 @@ class ModelState:
    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()
@@ -183,7 +182,7 @@ class ModelState:
    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


@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:09d080f36965bb2a0790500452bd328aa03c484d0222aa79d1ad9f021a522aec
-size 766040736
+oid sha256:6fee5937923c74848df4a63f6239eb6331c6274dd4bdb7a5d6ec0388a8b543d5
+size 766018462
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:659727c4d4839adc4992a254409a54259a8756a743f2d567bf5fdc6579f8009b
-size 60881999
+oid sha256:65a08adc31d5c456219687d99b7bf5e44d61dae2d49ea67850e76105c7248cce
+size 60918562