fix ty

.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,25 @@ 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(),
+  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()}
+  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 = {"frame": ((frame.size,), "|u1", Device.DEFAULT), "big_frame": ((frame.size,), "|u1", Device.DEFAULT),
+             "M_inv": ((2, 3, 3), "<f4", 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), "M_inv": ((2, 3, 3), "<f4", Device.DEFAULT)}
   with record_kernel_linear() as recorded:
     run(**inputs)
     run(**inputs)
@@ -34,7 +43,6 @@ 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)}
   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

@@ -12,6 +12,8 @@ from functools import partial
 from openpilot.common.basedir import BASEDIR
 from openpilot.sunnypilot.modeld_v2.helpers import _dtype, read_parts
 
+SPLIT_FRAMES = ("big_frame", "frame")
+
 
 def pack_layout(shapes: dict[str, tuple], frame_size: int) -> tuple[dict[str, int], int]:
   offsets, offset = {}, 0
@@ -39,29 +41,30 @@ 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
+  frame_specs = {s["name"]: s for s in warp["inputs"] if s["name"] in SPLIT_FRAMES}
-  offsets, packed = pack_layout(host, frame_size)
+  offsets, packed = pack_layout(host, 0 if frame_specs else next(s["numel"] for s in warp["inputs"] if s["name"] == "input_frame") // 2)
   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"])
+  warp_inputs = {s["slot"]: {"input": names.index(s["name"])} if s["name"] in frame_specs else
-                 for s in warp["inputs"]}
+                 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,15 +270,32 @@ def _program_uop(call: dict, lib: bytes):
                                UOp(Ops.SOURCE, arg=""), UOp(Ops.BINARY, arg=lib)), arg=info)
 
 
+def _input_tensors(kwargs: dict) -> list:  # in _prepare_jit_inputs' order
+  from tinygrad import Tensor
+  tensors = [t for _, t in sorted(kwargs.items()) if t.__class__ is Tensor]
+  for value in kwargs.values():
+    tensors += [t for t in (value.values() if isinstance(value, dict) else value if isinstance(value, (tuple, list)) else ())
+                if t.__class__ is Tensor and not any(t is y for y in tensors)]
+  return tensors
+
+
 class _Jit:
   def __init__(self, captured, input_names: list, buffers: list):
     self.captured, self.input_names, self.buffers = captured, input_names, buffers
 
   def __call__(self, *args, **kwargs):
     import tinygrad.engine.jit as jit
+    from tinygrad.uop.ops import Ops
+    # same names, no variables: the binding is the tensors' own uops, without _prepare_jit_inputs' graph rewrites
+    if not args and self.__dict__.get("_names") == sorted(kwargs):
+      tensors = _input_tensors(kwargs)
+      if all(t.uop.is_realized for t in tensors):
+        return self.captured([t.uop.src[0] if t.uop.op is Ops.UNSHARD else t.uop for t in tensors], {})
     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}")
+    if not args and not var_vals:
+      self._names = sorted(kwargs)
     return self.captured(input_uops, var_vals)
 
 
@@ -311,7 +328,7 @@ class PlainJit(_Jit):
         calls.append(_program_uop(call, bytes(blob[slice(*call["lib"])])).call(*arguments))
       else:
         calls.append(arguments[0].store_call(arguments[1]))
-    # stand-ins for JIT inputs: copy staging to GPU is decided from the input buffers, the real ones bind per call
+    # allocate gpu buffers
     stand_ins = {slot: UOp.from_buffer(Buffer(_device(spec["device"], device_map), spec["numel"] * _dtype(spec["dtype"]).itemsize).ensure_allocated(),
                                        _dtype(spec["dtype"])) for slot, spec in inputs.items()}
     input_uops = tuple(stand_ins.get(slot, next(iter(stand_ins.values()), None)) for slot in range(max(stand_ins, default=-1) + 1))

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]
+    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 = self.packed_input[offsets['frames']:].reshape(2, -1)
 
     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,29 @@ 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 isinstance(bufs[key], np.ndarray):  # cached memory: the GPU would read stale lines in place
+          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
+        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()
     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
+oid sha256:d776c72125156b5dd3eb59768409401e28c8a52a3a619cd3b9f02c10a156fda6
-size 3698799
+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
+oid sha256:a3f13ed12e087817a86afd48b2496c66c4f61a1925065b676b2945e81fabaacd
-size 7924847
+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"):
+      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