modeld_v2: bind a native model's inputs once per buffer set

_Jit ran tinygrad's _prepare_jit_inputs every call: graph rewrites over all
13 inputs to find the same buffers modeld binds every frame (4.8 ms on the
comma 4's little cores). Bindings are kept per set of realized input buffers
(16 at most; the camera buffers cycle through a few), not pickled into the load
cache, and dropped on warmup reset.

comma 4, visionipc buffers, copy_frames + run + readback: p50 44.4 -> 40.7 ms
with the in-place frames (46.8 before both), p95 55.3 -> 46.9 ms; outputs
unchanged.

.github/workflows/compile-warp-kernels.yaml

@@ -42,7 +42,7 @@ jobs:
             WARP_PKL="openpilot/sunnypilot/modeld_v2/models/driving_warp_${res}_tinygrad.pkl"
             echo "Compiling $WARP_PKL on QCOM"
             taskset -c 7 env DEV=QCOM IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1 python3 openpilot/sunnypilot/modeld_v2/compile_warp.py \
-              --frame ${FRAME_PARAM} --warp-to ${MODEL_SIZE} --output "${WARP_PKL}"
+              --frame ${FRAME_PARAM} --warp-to ${MODEL_SIZE} --split-frames --output "${WARP_PKL}"
 
             BIG_WARP_PKL="openpilot/sunnypilot/modeld_v2/models/big_driving_warp_${res}_tinygrad.pkl"
             echo "Compiling $BIG_WARP_PKL on AMD"

openpilot/sunnypilot/modeld_v2/compile_warp.py

@@ -17,16 +17,24 @@ def main():
   parser.add_argument("--frame", required=True, help="width,height,stride,y_height,uv_height,buffer_size")
   parser.add_argument("--warp-to", required=True, help="model input WxH")
   parser.add_argument("--output", required=True)
+  parser.add_argument("--split-frames", action="store_true", help="one input per camera (frame, big_frame), bound to the camera buffers by pointer")
   args = parser.parse_args()
 
   from examples.openpilot.compile_warp import make_frame_prepare, parse_frame, parse_size
   from tinygrad import Device, Tensor, TinyJit
   frame = parse_frame(args.frame)
   warp = make_frame_prepare(frame, *parse_size(args.warp_to))
-  run = TinyJit(lambda input_frame, M_inv: Tensor.stack(*(warp(input_frame[i], M_inv[i]) for i in range(2))), prune=True)
   rng = np.random.default_rng(42)
-  inputs = {"input_frame": Tensor(rng.integers(0, 256, (2, frame.size), dtype=np.uint8), device=Device.DEFAULT).realize(),
-            "M_inv": Tensor((rng.standard_normal((2, 3, 3)) * 8).astype(np.float32), device=Device.DEFAULT).realize()}
+  frames = rng.integers(0, 256, (2, frame.size), dtype=np.uint8)
+  m_inv = Tensor((rng.standard_normal((2, 3, 3)) * 8).astype(np.float32), device=Device.DEFAULT).realize()
+  if args.split_frames:
+    run = TinyJit(lambda frame, big_frame, M_inv: Tensor.stack(warp(frame, M_inv[0]), warp(big_frame, M_inv[1])), prune=True)
+    inputs = {"frame": Tensor(frames[0], device=Device.DEFAULT).realize(), "big_frame": Tensor(frames[1], device=Device.DEFAULT).realize(), "M_inv": m_inv}
+    specs = dict.fromkeys(("frame", "big_frame"), ((frame.size,), "|u1", Device.DEFAULT))
+  else:
+    run = TinyJit(lambda input_frame, M_inv: Tensor.stack(*(warp(input_frame[i], M_inv[i]) for i in range(2))), prune=True)
+    inputs = {"input_frame": Tensor(frames, device=Device.DEFAULT).realize(), "M_inv": m_inv}
+    specs = {"input_frame": ((2, frame.size), "|u1", Device.DEFAULT)}
   with record_kernel_linear() as recorded:
     run(**inputs)
     run(**inputs)
@@ -34,7 +42,7 @@ def main():
   assert len(recorded) == 1, f"expected one captured kernel list, got {len(recorded)}"
   part = serialize_jit(run, recorded[0], check)
   PlainJit(*part)
-  specs = {"input_frame": ((2, frame.size), "|u1", Device.DEFAULT), "M_inv": ((2, 3, 3), "<f4", Device.DEFAULT)}
+  specs |= {"M_inv": ((2, 3, 3), "<f4", Device.DEFAULT)}
   write_bundle(args.output, {"run": part}, {"input_specs": specs, "metadata": {}})
   print(f"Saved warp to {args.output} ({len(part[0]['calls'])} kernels)")
 

openpilot/sunnypilot/modeld_v2/frame_fusion.py

@@ -20,6 +20,9 @@ def pack_layout(shapes: dict[str, tuple], frame_size: int) -> tuple[dict[str, in
   return offsets | {"frames": offset}, offset + 2 * frame_size
 
 
+SPLIT_FRAMES = ("big_frame", "frame")  # a --split-frames warp: one input per camera, bound by pointer, not packed into host
+
+
 class _Blobs:  # the model's and the warp's blobs as one address space, for a part whose ranges span both files
   def __init__(self, *blobs):
     self.blobs, self.starts = blobs, list(itertools.accumulate([0] + [len(blob) for blob in blobs[:-1]]))
@@ -39,29 +42,31 @@ def frame_part(warp_path: str, parts: dict, extras: dict) -> dict:
   (model, model_blob), (warp, warp_blob) = parts["run"], read_parts(warp_path)[0]["run"]
   states = sorted(name for name in extras["input_specs"] if f"next_{name}" in extras["metadata"]["output_shapes"])
   host = {name: shape for name, (shape, _, _) in extras["input_specs"].items() if name not in states and name != "new_img"}
-  frame_size = next(s["numel"] for s in warp["inputs"] if s["name"] == "input_frame") // 2
-  offsets, packed = pack_layout(host, frame_size)
+  frame_specs = {s["name"]: s for s in warp["inputs"] if s["name"] in SPLIT_FRAMES}
+  frame_size = next(s["numel"] for s in warp["inputs"] if s["name"] in ("input_frame", "frame")) // (1 if frame_specs else 2)
+  offsets, packed = pack_layout(host, 0 if frame_specs else frame_size)
   shift, upload, u8 = len(model_blob), len(model["buffers"]) + len(warp["buffers"]), repr(_dtype("dtypes.uint8"))
 
   def view(name: str, nbytes: int, dtype: str) -> dict:
     return {"buffer": upload, "offset": offsets[name], "nbytes": nbytes, "dtype": dtype}
   outputs = sorted(s["slot"] for s in model["inputs"] if s["name"] is None)
-  names = sorted(["host", *states])
+  names = sorted(["host", *states, *frame_specs])
   slots = {s["slot"]: names.index(s["name"]) for s in model["inputs"] if s["name"] in states} | \
           {slot: len(names) + i for i, slot in enumerate(outputs)}
   returned = warp["returns"]["tensors"][0]
   model_inputs = {s["slot"]: {"input": slots[s["slot"]]} if s["slot"] in slots else
                   {k: v for k, v in returned.items() if k != "shape"} | {"buffer": returned["buffer"] + len(model["buffers"])} if s["name"] == "new_img" else
                   view(s["name"], s["numel"] * _dtype(s["dtype"]).itemsize, s["dtype"]) for s in model["inputs"]}
-  warp_inputs = {s["slot"]: view("frames" if s["name"] == "input_frame" else "tfm", s["numel"] * _dtype(s["dtype"]).itemsize, s["dtype"])
-                 for s in warp["inputs"]}
+  warp_inputs = {s["slot"]: {"input": names.index(s["name"])} if s["name"] in frame_specs else
+                 view("frames" if s["name"] == "input_frame" else "tfm", s["numel"] * _dtype(s["dtype"]).itemsize, s["dtype"]) for s in warp["inputs"]}
 
   def remap(calls: list, inputs: dict, buffers: int, lib: int) -> list:
     return [call | ({"lib": [call["lib"][0] + lib, call["lib"][1] + lib]} if "lib" in call else {}) |
             {"args": [dict(inputs[a["input"]]) if "input" in a else a | {"buffer": a["buffer"] + buffers} for a in call["args"]]} for call in calls]
   header = {"input_names": names, "returns": {"kind": "NoneType", "tensors": []}, "check": None,
             "inputs": [{"slot": names.index("host"), "name": "host", "dtype": u8, "numel": packed, "device": "NPY"}] +
-                      [s | {"slot": slots[s["slot"]]} for s in model["inputs"] if s["slot"] in slots],
+                      [s | {"slot": slots[s["slot"]]} for s in model["inputs"] if s["slot"] in slots] +
+                      [s | {"slot": names.index(name)} for name, s in frame_specs.items()],
             "buffers": model["buffers"] + [b | ({"init": [b["init"][0] + shift, b["init"][1] + shift]} if "init" in b else {}) for b in warp["buffers"]] +
                        [{"nbytes": packed, "device": model["inputs"][0]["device"], "dtype": u8}],
             "calls": [{"kind": "copy", "args": [{"buffer": upload, "offset": 0, "nbytes": packed, "dtype": u8}, {"input": names.index("host")}]}] +

openpilot/sunnypilot/modeld_v2/helpers.py

@@ -270,16 +270,42 @@ def _program_uop(call: dict, lib: bytes):
                                UOp(Ops.SOURCE, arg=""), UOp(Ops.BINARY, arg=lib)), arg=info)
 
 
+def _input_key(args: tuple, kwargs: dict) -> tuple | None:  # the realized buffers a call binds, by name; None when it can't be reused
+  from tinygrad import Tensor
+  if args:
+    return None
+  key = []
+  for name, value in sorted(kwargs.items()):
+    if not isinstance(value, (Tensor, dict)):
+      return None
+    for sub, tensor in [(None, value)] if isinstance(value, Tensor) else sorted(value.items()):
+      if not isinstance(tensor, Tensor) or not tensor.uop.is_realized:
+        return None
+      key.append((name, sub, tensor.uop))
+  return tuple(key)
+
+
 class _Jit:
   def __init__(self, captured, input_names: list, buffers: list):
     self.captured, self.input_names, self.buffers = captured, input_names, buffers
 
+  def __getstate__(self):
+    return {k: v for k, v in self.__dict__.items() if k != "_bound"}
+
   def __call__(self, *args, **kwargs):
     import tinygrad.engine.jit as jit
-    input_uops, var_vals, names, _ = jit._prepare_jit_inputs(args, kwargs)
-    if names != self.input_names:
-      raise ValueError(f"expected inputs {self.input_names}, got {names}")
-    return self.captured(input_uops, var_vals)
+    # modeld binds the same buffers every frame (the camera buffers cycle through a few): bind each set once, not per call
+    bound, key = self.__dict__.setdefault("_bound", {}), _input_key(args, kwargs)
+    if key is None or (hit := bound.get(key)) is None:
+      input_uops, var_vals, names, _ = jit._prepare_jit_inputs(args, kwargs)
+      if names != self.input_names:
+        raise ValueError(f"expected inputs {self.input_names}, got {names}")
+      hit = (input_uops, var_vals)
+      if key is not None:
+        if len(bound) >= 16:
+          bound.clear()
+        bound[key] = hit
+    return self.captured(*hit)
 
 
 class PlainJit(_Jit):

openpilot/sunnypilot/modeld_v2/model_adapters.py

@@ -11,7 +11,7 @@ import pickle
 import numpy as np
 
 from openpilot.common.basedir import BASEDIR
-from openpilot.sunnypilot.modeld_v2.frame_fusion import pack_layout
+from openpilot.sunnypilot.modeld_v2.frame_fusion import pack_layout, SPLIT_FRAMES
 from openpilot.sunnypilot.modeld_v2.helpers import load_model
 from openpilot.sunnypilot.modeld_v2.compile_modeld import (POLICY_INPUTS, derive_frame_skip,
                                                            make_split_input_queues, make_supercombo_input_queues)
@@ -169,16 +169,18 @@ class NativeTinygradAdapter(BaseModelAdapter):
     self.warp_frame_size = self.frame_copy_size
     self._init_common()
     self.run_model = self.jits['run']
+    self.split = [name for name in SPLIT_FRAMES if name in self.run_model.input_names]  # the warp reads the camera buffers in place
+    self.frame_inputs: dict[str, Tensor] = {}
 
     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 = {name: shape for name, (shape, _) in self.input_shapes.items() if name not in self.state_pairs and name != 'new_img'}
-    offsets, packed = pack_layout(shapes, self.warp_frame_size)
+    offsets, packed = pack_layout(shapes, 0 if self.split else self.warp_frame_size)
     self.packed_input = np.zeros(packed, dtype=np.uint8)
     self.host = Tensor(self.packed_input, device='NPY')
     self.numpy_inputs = {name: np.ndarray(shape, dtype=np.float32, buffer=self.packed_input, offset=offsets[name])
                          for name, shape in ({'tfm': (2, 3, 3)} | shapes).items()}
-    self.frames = self.packed_input[offsets['frames']:].reshape(2, self.warp_frame_size)
+    self.frames = None if self.split else self.packed_input[offsets['frames']:].reshape(2, self.warp_frame_size)
 
     self.outputs = {name: Tensor(np.zeros(shape, dtype=dtype), device=device).realize() for name, (shape, dtype, device) in self.jits['output_specs'].items()}
     for name, next_name in self.state_pairs.items():
@@ -187,17 +189,30 @@ class NativeTinygradAdapter(BaseModelAdapter):
 
   def copy_frames(self, bufs):
     for i, key in enumerate(self._vision_input_names):
-      if key in bufs:
+      if key not in bufs:
+        continue
+      if not self.split:
         data = bufs[key].data if hasattr(bufs[key], 'data') else bufs[key]
         np.copyto(self.frames[i], np.frombuffer(data, dtype=np.uint8, count=self.warp_frame_size))
+      elif (name := ('frame', 'big_frame')[i]) in self.split:
+        if not hasattr(bufs[key], 'data'):  # a numpy frame (warmup, tests): cached memory the GPU must not read in place, so upload it
+          self.frame_inputs[name] = Tensor(np.ascontiguousarray(bufs[key][:self.warp_frame_size]), device=self.model_device).realize()
+          continue
+        ptr = np.frombuffer(bufs[key].data, dtype=np.uint8).ctypes.data  # visionipc buffers are coherent: map once per buffer, no copy
+        if (key, ptr) not in self._blob_cache:
+          self._blob_cache[(key, ptr)] = Tensor.from_blob(ptr, (self.warp_frame_size,), dtype='uint8', device=self.model_device)
+        self.frame_inputs[name] = self._blob_cache[(key, ptr)]
 
   def reset_warmup_buffers(self) -> None:
     self.packed_input[:] = 0
+    self.frame_inputs.clear()
+    self._blob_cache.clear()
+    self.run_model.__dict__.pop('_bound', None)  # warmup's uploaded frames: drop their bindings with them
     for key in self.state_pairs:
       self.input_queues[key].assign(0).realize()
 
   def run(self):
-    self.run_model(host=self.host, **self.input_queues, output_buffers=self.outputs)
+    self.run_model(host=self.host, **self.input_queues, **self.frame_inputs, output_buffers=self.outputs)
     return self.outputs['outputs']
 
 

openpilot/sunnypilot/modeld_v2/models/driving_warp_1344x760_tinygrad.pkl

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4b22fd6bcb3a3e8a642e6866449dd9c30c12e694ff0741cad9e34e4b00d8c5fa
-size 3698799
+oid sha256:d776c72125156b5dd3eb59768409401e28c8a52a3a619cd3b9f02c10a156fda6
+size 3699072

openpilot/sunnypilot/modeld_v2/models/driving_warp_1928x1208_tinygrad.pkl

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9a17fd6acece4fab23784ef0f3f8bba28a049e9599e5e9e5ee7b36e4b1d24a6f
-size 7924847
+oid sha256:a3f13ed12e087817a86afd48b2496c66c4f61a1925065b676b2945e81fabaacd
+size 7924864

openpilot/sunnypilot/modeld_v2/tests/test_plain_artifact.py

@@ -144,6 +144,21 @@ class TestPlainArtifact(OpenpilotTestCase):
     for name in ("outputs", "next_state_feat_q"):
       assert np.array_equal(inputs["output_buffers"][name].numpy(), expected["output_buffers"][name].numpy()), name
 
+    split_path = os.path.join(self.tmp.name, "warp_split_tinygrad.pkl")
+    with mock.patch.object(sys, "argv", ["compile_warp", "--frame", FRAME, "--warp-to", f"{MODEL_W}x{MODEL_H}", "--output", split_path, "--split-frames"]):
+      compile_warp.main()
+    split = load_model(path, cache=False, extend=functools.partial(frame_part, split_path))
+    assert split["run"].input_names == ["big_frame", "frame", "host", "state_feat_q"]
+    offsets, packed = pack_layout(HOST, 0)
+    host = np.zeros(packed, np.uint8)
+    for name, shape in ({"tfm": (2, 3, 3)} | HOST).items():
+      np.ndarray(shape, np.float32, host, offsets[name])[...] = tfm.numpy() if name == "tfm" else inputs[name].numpy()
+    out = {"outputs": Tensor.zeros(OUT).contiguous().realize(), "next_state_feat_q": Tensor.zeros(STATE).contiguous().realize()}
+    split["run"](host=Tensor(host, device="NPY"), frame=Tensor(frames.numpy()[0]).realize(), big_frame=Tensor(frames.numpy()[1]).realize(),
+                 state_feat_q=expected["state_feat_q"], output_buffers=out)
+    for name in ("outputs", "next_state_feat_q"):  # the camera buffers bound in place give the packed upload's result exactly
+      assert np.array_equal(out[name].numpy(), expected["output_buffers"][name].numpy()), name
+
   def test_tinygrad_api_the_loader_reads(self):
     import tinygrad.engine.jit as jit
     from tinygrad.device import Buffer, TinyELF