Reduce Chestnut model startup overhead

openpilot/selfdrive/modeld/modeld.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 from collections.abc import Callable
+from typing import TYPE_CHECKING, cast
 import base64
 import ctypes
 import gc
@@ -15,6 +16,8 @@ from tinygrad.engine.realize import lower_and_compile
 from tinygrad.tensor import Tensor
 from tinygrad.helpers import all_int, round_up
 from tinygrad.uop.ops import UOp
+if TYPE_CHECKING:
+  from openpilot.starpilot.models.runtime_bridge.legacy_types import LegacyBuffer, LegacyUOpFactory
 import math
 import pickle
 import threading
@@ -158,8 +161,8 @@ class FrameMeta:
 def input_view(buffer: Buffer, shape: tuple[int, ...], dtype: DType, offset: int) -> Tensor:
   from openpilot.starpilot.models.runtime_bridge.runtime import selected_legacy
   if selected_legacy():
-    view = buffer.view(math.prod(shape), dtype, offset).ensure_allocated()
-    return Tensor(UOp.from_buffer(view)).reshape(shape)
+    view = cast("LegacyBuffer", buffer).view(math.prod(shape), dtype, offset).ensure_allocated()
+    return Tensor(cast("LegacyUOpFactory", UOp).from_buffer(view)).reshape(shape)
   view = buffer.view(math.prod(shape) * dtype.itemsize, offset).ensure_allocated()
   return Tensor(UOp.from_buffer(view, dtype)).reshape(shape)
 
@@ -237,7 +240,8 @@ class ModelState:
     self.input_device.copy_from(self.input_host)
     self.input_queues['new_img'] = self.run_warp(**self.warp_inputs)
     self.run_model(output_buffers=self.outputs, **self.input_queues)
-    if self.chestnut and not selected_legacy() and (err := Device[self.outputs['outputs'].device].error_state.host.view(fmt='q')[0]):
+    if (self.chestnut and not selected_legacy() and
+        (err := Device[cast(str, self.outputs['outputs'].device)].error_state.host.view(fmt='q')[0])):
       raise RuntimeError(f"chestnut failed with {err}")
     if after_enqueue is not None:
       after_enqueue()

openpilot/starpilot/models/compile.py

@@ -9,6 +9,7 @@ import tempfile
 import time
 from collections import namedtuple
 from functools import partial
+from typing import TYPE_CHECKING, cast
 
 if __name__ == '__main__':
   from openpilot.starpilot.models.runtime_bridge.runtime import prepare_compiler
@@ -20,6 +21,8 @@ from tinygrad.dtype import _to_np_dtype, dtypes
 from tinygrad.device import Buffer, BufferSpec
 from tinygrad.uop.ops import UOp
 from tinygrad_repo.examples.openpilot.helpers import dump_pickle, load_pickle
+if TYPE_CHECKING:
+  from openpilot.starpilot.models.runtime_bridge.legacy_types import LegacyBufferFactory, LegacyUOpFactory
 
 COMPILER_REVISION = "9d0446a4ba8a532c8b674fb6ad795af015cd9dcf"
 ARTIFACT_ABI = "tinygrad_single_v1_arena"
@@ -448,9 +451,9 @@ def make_stateful_input_queues(metadata, device):
     image_shape = tuple(metadata['input_shapes']['new_img'])
     from openpilot.starpilot.models.runtime_bridge.runtime import selected_legacy
     if selected_legacy():
-      buffer = Buffer(device, math.prod(image_shape), dtypes.uint8,
-                      options=BufferSpec(cpu_access=True), preallocate=True)
-      view = UOp.from_buffer(buffer)
+      buffer = cast("LegacyBufferFactory", Buffer)(device, math.prod(image_shape), dtypes.uint8,
+                                                   options=BufferSpec(cpu_access=True), preallocate=True)
+      view = cast("LegacyUOpFactory", UOp).from_buffer(buffer)
     else:
       buffer = Buffer(device, math.prod(image_shape) * dtypes.uint8.itemsize,
                       options=BufferSpec(cpu_access=True), preallocate=True)

openpilot/starpilot/models/runtime_bridge/guard-policy.json

@@ -1 +1 @@
-{"compiler_revision":"9d0446a4ba8a532c8b674fb6ad795af015cd9dcf","files":{"bootstrap.py":"dcb4a3469f4e3cf93b7f2384a47e72a5559ed859d3a94c6f1b78ba72d5efa86c","bounded_wait.c":"a9e1a82e50ebefd983b3dbc6f784adc465cbd06848debed11dfff80cc8a832fe","capture.py":"834cf5d83725729f30aec68a7ba5ca4897701dd4aa34147f0b49d8b0a9b9e81a","libbounded_wait.so":"24ed559861c8636c1026cad9a244854bb2236b5ee8b91bf21c5efbf688fa4a69","libusb_guards.so":"66a3dfe48ac17899862ab80160cf6dac2c9b9707234527ba0d808f3b68404114","qcom.py":"1721afdcb9f6e7d413e36438390f84cfbb9ca33e646b84fdd321bd652af99611","usb.py":"9cef2454861fe90d84d083ab335e2a5d29a902b1dd0571ca4bc2cd73d5e4f5e2","usb_guards.c":"8b3c207de74f92b74980da35428cff6f4a198ce97f9f8b6e64aded74e0110316"},"native_abi":"linux-aarch64-64bit","runtime_manifest_sha256":"20ca9e77010c1ab788f2adb824103db3d595b0da8c4069d6e32b1266f7977b1f"}
+{"compiler_revision":"9d0446a4ba8a532c8b674fb6ad795af015cd9dcf","files":{"bootstrap.py":"dcb4a3469f4e3cf93b7f2384a47e72a5559ed859d3a94c6f1b78ba72d5efa86c","bounded_wait.c":"a9e1a82e50ebefd983b3dbc6f784adc465cbd06848debed11dfff80cc8a832fe","capture.py":"834cf5d83725729f30aec68a7ba5ca4897701dd4aa34147f0b49d8b0a9b9e81a","libbounded_wait.so":"24ed559861c8636c1026cad9a244854bb2236b5ee8b91bf21c5efbf688fa4a69","libusb_guards.so":"66a3dfe48ac17899862ab80160cf6dac2c9b9707234527ba0d808f3b68404114","qcom.py":"5f35de06bae649342953bfef48c54267f0c7971756537b8758e1d0d2ec6fa7b0","usb.py":"a2e84aa6d6be4cd7b8db7bf9dae85f67a6008de5c6e93ecb288fefede2b46e23","usb_guards.c":"8b3c207de74f92b74980da35428cff6f4a198ce97f9f8b6e64aded74e0110316"},"native_abi":"linux-aarch64-64bit","runtime_manifest_sha256":"20ca9e77010c1ab788f2adb824103db3d595b0da8c4069d6e32b1266f7977b1f"}

openpilot/starpilot/models/runtime_bridge/legacy_types.py

@@ -0,0 +1,40 @@
+"""Type-only facets of the selected, pinned 9d runtime's legacy APIs."""
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Protocol
+
+if TYPE_CHECKING:
+  from pickle import PickleBuffer
+  from tinygrad.device import BufferSpec
+  from tinygrad.dtype import DType
+  from tinygrad.uop.ops import UOp
+
+
+class LegacyBuffer(Protocol):
+  @property
+  def nbytes(self) -> int: ...
+  def copy_from(self, src: LegacyBuffer) -> LegacyBuffer: ...
+  def view(self, size: int, dtype: DType, offset: int) -> LegacyBuffer: ...
+  def ensure_allocated(self) -> LegacyBuffer: ...
+
+
+class LegacyBufferFactory(Protocol):
+  def __call__(self, device: str, size: int, dtype: DType, *, options: BufferSpec | None = None,
+               initial_value: bytes | PickleBuffer | None = None, preallocate: bool = False) -> LegacyBuffer: ...
+
+
+class LegacyUOpFactory(Protocol):
+  def from_buffer(self, opaque: LegacyBuffer, device: str | tuple[str, ...] | None = None) -> UOp: ...
+
+
+class LegacyCall(Protocol):
+  def call(self, *srcs: UOp, ret_dtype: DType | None = None) -> UOp: ...
+
+
+class LegacyHCQModule(Protocol):
+  lower_and_compile: Callable[[UOp], UOp]
+
+
+class LegacyPersistentUnpickler(Protocol):
+  persistent_load: Callable[[tuple[int, DType, int]], LegacyBuffer]

openpilot/starpilot/models/runtime_bridge/qcom.py

@@ -4,6 +4,7 @@ from dataclasses import replace
 import ctypes
 import hashlib
 from pathlib import Path
+from typing import TYPE_CHECKING, cast
 
 from tinygrad.device import Buffer, HCQ_RUNTIME_DEV
 from tinygrad.dtype import AddrSpace, dtypes
@@ -11,6 +12,9 @@ from tinygrad.engine.realize import lower_and_compile
 from tinygrad.uop.ops import AxisType, Ops, UOp
 from tinygrad.runtime.support.hcq2 import HCQInfo
 
+if TYPE_CHECKING:
+  from openpilot.starpilot.models.runtime_bridge.legacy_types import LegacyBufferFactory, LegacyCall, LegacyUOpFactory
+
 
 def byte_inventory(root):
   return Counter(hashlib.sha256(u.arg).hexdigest() for u in root.toposort()
@@ -25,7 +29,7 @@ class Protection:
     self.function.restype = ctypes.c_uint64
     address = ctypes.cast(self.function, ctypes.c_void_p).value
     assert address is not None
-    self.pointer = Buffer(HCQ_RUNTIME_DEV.value, 1, dtypes.uint64,
+    self.pointer = cast("LegacyBufferFactory", Buffer)(cast(str, HCQ_RUNTIME_DEV.value), 1, dtypes.uint64,
                           initial_value=address.to_bytes(8, 'little'))
     self.records = []
 
@@ -45,7 +49,7 @@ class Protection:
     params = {u: u.replace(arg=replace(u.arg, slot=u.arg.slot + 1)) for u in sink.toposort()
               if u.op is Ops.PARAM and u.arg.addrspace is AddrSpace.ALU}
     sink = sink.substitute(params, enter_calls=True)
-    pointer = UOp.param(info.nargs, dtypes.uint64, 1, HCQ_RUNTIME_DEV.value, name='bounded_wait_pointer')
+    pointer = UOp.param(info.nargs, dtypes.uint64, 1, cast(str, HCQ_RUNTIME_DEV.value), name='bounded_wait_pointer')
     callee = UOp.custom_function('starpilot_bounded_wait', pointer.index(0).load())
     replacements, seen_loops = {}, set()
     for end in sink.toposort():
@@ -70,7 +74,7 @@ class Protection:
       if loop in target.toposort():
         raise ValueError('Fence target depends on its polling loop')
       signal_index = index.substitute({loop: UOp(Ops.NOOP)}, enter_calls=True)
-      replacements[end] = callee.call(signal_index, target, ret_dtype=dtypes.uint64)
+      replacements[end] = cast("LegacyCall", callee).call(signal_index, target, ret_dtype=dtypes.uint64)
       seen_loops.add(loop)
     weak = {u for u in sink.toposort() if u.op is Ops.RANGE and u.src and u.src[0].op is Ops.NOOP}
     if len(replacements) != 1 or weak != seen_loops:
@@ -83,9 +87,9 @@ class Protection:
     if any(u.op is Ops.CALL and u.body.op in (Ops.SINK, Ops.PROGRAM) for u in rewritten.toposort()):
       raise ValueError('Nested compilation request in host source')
     # The only compilation request is this CPU SINK.
-    new_call = call.replace(src=(rewritten, *call.src[1:], UOp.from_buffer(self.pointer)),
+    new_call = call.replace(src=(rewritten, *call.src[1:], cast("LegacyUOpFactory", UOp).from_buffer(self.pointer)),
                             arg=replace(call.arg, aux=replace(info, nargs=info.nargs + 1)))
-    if new_call.device != HCQ_RUNTIME_DEV.value or HCQ_RUNTIME_DEV.value.split(':')[0] != 'CPU':
+    if new_call.device != cast(str, HCQ_RUNTIME_DEV.value) or cast(str, HCQ_RUNTIME_DEV.value).split(':')[0] != 'CPU':
       raise ValueError('Refusing any non-CPU host compilation')
     compiled = lower_and_compile(UOp(Ops.LINEAR, src=(new_call,))).src[0]
     if compiled.body.arg.target.device != 'CPU' or byte_inventory(compiled.body.src[0]) != before:

openpilot/starpilot/models/runtime_bridge/runtime.py

@@ -1,16 +1,21 @@
 """Load existing bank and stock artifacts with one explicitly selected runtime."""
 import hashlib
 from pathlib import Path
+from typing import TYPE_CHECKING, cast
 
 from openpilot.starpilot.models.runtime_bridge.bootstrap import LEGACY_REVISION, activate
 
+if TYPE_CHECKING:
+  from tinygrad.uop.ops import UOp
+  from openpilot.starpilot.models.runtime_bridge.legacy_types import LegacyBuffer, LegacyBufferFactory, LegacyHCQModule, LegacyPersistentUnpickler
+
 BUNDLE = Path(__file__).resolve().parents[4] / 'tinygrad_model_runtime/legacy9d'
 MANIFEST_SHA256 = "20ca9e77010c1ab788f2adb824103db3d595b0da8c4069d6e32b1266f7977b1f"
 _GUARD_HASHES = {
   'libbounded_wait.so': '24ed559861c8636c1026cad9a244854bb2236b5ee8b91bf21c5efbf688fa4a69',
   'libusb_guards.so': '66a3dfe48ac17899862ab80160cf6dac2c9b9707234527ba0d808f3b68404114',
 }
-_POLICY_SHA256 = '39c3387f44b570895aab440c6afa52eeb1996c99510205c92ea774565d881b66'
+_POLICY_SHA256 = 'c7eae9e0fab99e8c5fa56854b5a1e1ca8c36c81ef73668dd18ede27de61e3f4b'
 _selection = None
 _protection = None
 
@@ -102,14 +107,15 @@ def _guard(chestnut):
       from tinygrad.runtime.support import hcq2
       from tinygrad.uop.ops import Ops, ProgramInfo
       candidate = Protection(_library('libbounded_wait.so'))
-      original = hcq2.lower_and_compile
-      def checked(linear):
+      legacy_hcq = cast("LegacyHCQModule", hcq2)
+      original = legacy_hcq.lower_and_compile
+      def checked(linear: "UOp") -> "UOp":
         if any(u.op is Ops.CALL and (u.body.op is Ops.SINK or
                (u.body.op is Ops.PROGRAM and not (isinstance(u.body.arg, ProgramInfo) and
                 u.body.src and u.body.src[-1].op is Ops.BINARY))) for u in linear.toposort()):
           raise ValueError('Unexpected fresh monitoring submission compilation')
         return original(linear)
-      hcq2.lower_and_compile = checked
+      legacy_hcq.lower_and_compile = checked
     _protection = candidate
   return _protection
 
@@ -159,12 +165,37 @@ def _load_monitoring(data):
   end = 8 + opcode_size
   if end == len(data):
     return pickle.loads(data[8:end])
-  arena = Buffer(Device.DEFAULT, len(data) - end, dtypes.uchar, initial_value=data[end:])
+  arena = cast("LegacyBufferFactory", Buffer)(Device.DEFAULT, len(data) - end, dtypes.uchar, initial_value=data[end:])
   unpickler = pickle.Unpickler(io.BytesIO(data[8:end]))
-  unpickler.persistent_load = lambda record: arena.view(*record)
+  cast("LegacyPersistentUnpickler", unpickler).persistent_load = lambda record: arena.view(*record)
   return unpickler.load()
 
 
+def _load_chestnut(path):
+  import io
+  import pickle
+  import struct
+  from tinygrad import Device, Tensor
+  from tinygrad.device import Buffer
+  from tinygrad.dtype import dtypes
+  with open(path, "rb") as f:
+    opcodes = f.read(struct.unpack('<q', f.read(8))[0])
+    if f.tell() == Path(path).stat().st_size:
+      return pickle.loads(opcodes)
+    buffers = cast("LegacyBuffer", Tensor(Path(path))[f.tell():].uop.buffer).ensure_allocated()
+
+  arena, chunk_size = cast("LegacyBufferFactory", Buffer)(Device.DEFAULT, buffers.nbytes, dtypes.uchar, preallocate=True), 8 << 20
+  for off in range(0, buffers.nbytes, chunk_size):
+    size = min(chunk_size, buffers.nbytes-off)
+    arena.view(size, dtypes.uchar, off).ensure_allocated().copy_from(buffers.view(size, dtypes.uchar, off).ensure_allocated())
+
+  def persistent_load(pid): return arena.view(*pid)
+
+  u = pickle.Unpickler(io.BytesIO(opcodes))
+  cast("LegacyPersistentUnpickler", u).persistent_load = persistent_load
+  return u.load()
+
+
 def load_oob(path, chestnut=False):
   if _selection is None:
     raise RuntimeError('A device model loader requires the selected runtime')
@@ -179,7 +210,12 @@ def load_oob(path, chestnut=False):
   from tinygrad_repo.examples.openpilot.helpers import load_pickle
   protection = _guard(chestnut)
   with Context(DEV='USB+AMD:LLVM' if chestnut else 'QCOM'):
-    artifact = _load_monitoring(authenticated) if monitoring else load_pickle(path, out_of_band=True)
+    if monitoring:
+      artifact = _load_monitoring(authenticated)
+    elif chestnut:
+      artifact = _load_chestnut(path)
+    else:
+      artifact = load_pickle(path, out_of_band=True)
   if isinstance(artifact, dict) and artifact.get('artifact_abi') == 'tinygrad_single_v1_arena':
     if chestnut:
       return protection.artifact(artifact)

openpilot/starpilot/models/runtime_bridge/usb.py

@@ -3,6 +3,7 @@ from collections import Counter
 from dataclasses import replace
 import ctypes
 import hashlib
+from typing import TYPE_CHECKING, cast
 
 from tinygrad import Device
 from tinygrad.device import Buffer, HCQ_RUNTIME_DEV
@@ -15,6 +16,9 @@ from tinygrad.runtime.support.hcq2 import HCQInfo
 from tinygrad.uop.ops import AxisType, Ops, UOp, ProgramInfo
 from openpilot.starpilot.models.runtime_bridge.qcom import Protection, byte_inventory
 
+if TYPE_CHECKING:
+  from openpilot.starpilot.models.runtime_bridge.legacy_types import LegacyBufferFactory, LegacyCall, LegacyHCQModule, LegacyUOpFactory
+
 
 class USBProtection(Protection):
   ffi_names = {'libusb_control_transfer': 'sp_control', 'libusb_bulk_transfer': 'sp_bulk',
@@ -33,7 +37,7 @@ class USBProtection(Protection):
     addresses = [ctypes.cast(getattr(self.usb_library, name), ctypes.c_void_p).value for name in self.names]
     if any(address is None for address in addresses):
       raise ValueError('Missing native submission guard')
-    self.usb_pointer = Buffer(HCQ_RUNTIME_DEV.value, len(addresses), dtypes.uint64,
+    self.usb_pointer = cast("LegacyBufferFactory", Buffer)(cast(str, HCQ_RUNTIME_DEV.value), len(addresses), dtypes.uint64,
                               initial_value=b''.join(address.to_bytes(8, 'little') for address in addresses))
     self.fresh_records = []
 
@@ -60,7 +64,7 @@ class USBProtection(Protection):
     params = {u: u.replace(arg=replace(u.arg, slot=u.arg.slot + 1)) for u in sink.toposort()
               if u.op is Ops.PARAM and u.arg.addrspace is AddrSpace.ALU}
     sink = sink.substitute(params, enter_calls=True)
-    pointer = UOp.param(info.nargs, dtypes.uint64, len(self.names), HCQ_RUNTIME_DEV.value, name='usb_guard_pointer')
+    pointer = UOp.param(info.nargs, dtypes.uint64, len(self.names), cast(str, HCQ_RUNTIME_DEV.value), name='usb_guard_pointer')
     def callee(name):
       return UOp.custom_function(name, pointer.index(self.names.index(name)).load())
     weak = [u for u in sink.toposort() if u.op is Ops.RANGE and u.src and u.src[0].op is Ops.NOOP]
@@ -104,11 +108,11 @@ class USBProtection(Protection):
           token = next((a for a in u.src[1:] if loop in a.ranges), None)
           if token is None:
             raise ValueError('Polling call has no original open-loop argument')
-          touches.append(callee('sp_loop_touch').call(UOp.const(slot, dtypes.uint64),
+          touches.append(cast("LegacyCall", callee('sp_loop_touch')).call(UOp.const(slot, dtypes.uint64),
                            UOp.const(timeout, dtypes.uint64), scalar64(token), ret_dtype=dtypes.uint64))
           touched.add(loop.arg)
       marker = sum(touches[1:], touches[0]) if touches else UOp.const(0, dtypes.uint64)
-      substitutions[u] = callee(self.ffi_names[u.body.arg]).call(
+      substitutions[u] = cast("LegacyCall", callee(self.ffi_names[u.body.arg])).call(
         *[scalar64(a) for a in (u.body.src[0], *u.src[1:])], marker, ret_dtype=dtypes.int32)
     if touched != set(specifications) or not substitutions:
       raise ValueError('A native polling loop has no entry guard')
@@ -116,7 +120,7 @@ class USBProtection(Protection):
     guarded_ranges = [u for u in sink.toposort() if u.op is Ops.RANGE and u.src and u.src[0].op is Ops.NOOP]
     if len(guarded_ranges) != len(weak) or {r.arg for r in guarded_ranges} != set(specifications):
       raise ValueError('Native polling ranges changed')
-    ends = {u: u.replace(src=(*u.src[:2], callee('sp_loop_check').call(
+    ends = {u: u.replace(src=(*u.src[:2], cast("LegacyCall", callee('sp_loop_check')).call(
             u.src[2].cast(dtypes.uint64), UOp.const(specifications[u.src[1].arg][0], dtypes.uint64),
             ret_dtype=dtypes.uint64).cast(dtypes.bool)))
             for u in sink.toposort() if u.op is Ops.END and len(u.src) == 3 and u.src[1].op is Ops.RANGE and u.src[1].arg in specifications}
@@ -132,12 +136,12 @@ class USBProtection(Protection):
       raise ValueError('Opaque device payload changed')
     if any(u.op is Ops.CALL and u.body.op in (Ops.SINK, Ops.PROGRAM) for u in rewritten.toposort()):
       raise ValueError('Nested device compilation in a CPU host')
-    updated = call.replace(src=(rewritten, *call.src[1:], UOp.from_buffer(self.usb_pointer)),
+    updated = call.replace(src=(rewritten, *call.src[1:], cast("LegacyUOpFactory", UOp).from_buffer(self.usb_pointer)),
                            arg=replace(call.arg, aux=replace(info, nargs=info.nargs + 1)))
-    if updated.device != HCQ_RUNTIME_DEV.value or HCQ_RUNTIME_DEV.value.split(':')[0] != 'CPU':
+    if updated.device != cast(str, HCQ_RUNTIME_DEV.value) or cast(str, HCQ_RUNTIME_DEV.value).split(':')[0] != 'CPU':
       raise ValueError('Refusing non-CPU compilation')
     compiled = lower_and_compile(UOp(Ops.LINEAR, src=(updated,))).src[0]
-    host_renderer = Device[HCQ_RUNTIME_DEV.value].renderer
+    host_renderer = Device[cast(str, HCQ_RUNTIME_DEV.value)].renderer
     if (compiled.body.arg.target.device != 'CPU' or not isinstance(host_renderer, ClangRenderer) or
         compiled.body.arg.target != host_renderer.target or byte_inventory(compiled.body.src[0]) != before):
       raise ValueError({'failure':'USB host-only compilation contract failed','target':repr(compiled.body.arg.target),
@@ -148,10 +152,11 @@ class USBProtection(Protection):
     return compiled
 
   def install_fresh_guard(self):
-    original = hcq2.lower_and_compile
+    legacy_hcq = cast("LegacyHCQModule", hcq2)
+    original = legacy_hcq.lower_and_compile
     if original is not lower_and_compile:
       raise ValueError('Host compiler was already replaced')
-    def protect(linear):
+    def protect(linear: "UOp") -> "UOp":
       replacements = {u: self.protect_call(u) for u in linear.toposort()
                       if u.op is Ops.CALL and isinstance(u.arg.aux, HCQInfo) and u.body.op is Ops.SINK}
       self.fresh_records.append(len(replacements))
@@ -161,7 +166,7 @@ class USBProtection(Protection):
               u.body.src[-1].op is Ops.BINARY)) for u in prepared.toposort()):
         raise ValueError('Refusing residual unqualified compilation during native weight loading')
       return original(prepared)
-    hcq2.lower_and_compile = protect
+    legacy_hcq.lower_and_compile = protect
 
   def artifact(self, artifact):
     if artifact.get('execution_device') != 'AMD' or artifact.get('warp_device') != 'QCOM':

openpilot/starpilot/models/tests/runtime_bridge_fixtures.py

@@ -6,11 +6,19 @@ import os
 from pathlib import Path
 import sys
 import tempfile
+from typing import TypedDict
 import unittest
 
 from openpilot.starpilot.models.runtime_bridge import bootstrap
 
 
+class RuntimeManifest(TypedDict):
+  compiler_revision: str
+  adapter: str
+  roles: dict[str, list[str]]
+  files: dict[str, str]
+
+
 class ProcessCustody(unittest.TestCase):
   """Restore state even when fixture construction or an assertion fails.
 
@@ -80,7 +88,7 @@ class BootstrapFixture(ProcessCustody):
 
   def seal(self):
     files = {name: hashlib.sha256((self.bundle / name).read_bytes()).hexdigest() for name in self.sources}
-    self.manifest = {
+    self.manifest: RuntimeManifest = {
       'compiler_revision': bootstrap.LEGACY_REVISION, 'adapter': 'bounded-host-v2',
       'roles': {'driving': ['QCOM', 'USB+AMD:LLVM'], 'monitoring': ['QCOM']}, 'files': files,
     }

openpilot/starpilot/models/tests/test_chestnut_reader.py

@@ -0,0 +1,215 @@
+"""Actual reader ASTs and real pickle envelopes; fake host storage, no devices."""
+import ast
+import io
+from pathlib import Path
+import pickle
+import struct
+import sys
+import tempfile
+from types import ModuleType, SimpleNamespace
+from typing import cast
+import unittest
+from unittest.mock import patch
+
+
+class DType:
+  def __init__(self, name, itemsize):
+    self.name, self.itemsize = name, itemsize
+
+
+UCHAR, UINT32 = DType('uchar', 1), DType('uint32', 4)
+
+
+class Storage:
+  def __init__(self):
+    self.arenas, self.copies, self.source_views = [], [], []
+
+  def buffer(self, device, size, dtype, *, preallocate=False):
+    value = View(bytearray(size * dtype.itemsize), size, dtype, 0, self, 'arena')
+    self.arenas.append((device, size, dtype, preallocate, value))
+    return value
+
+  def tensor(self, path):
+    owner = self
+
+    class FileTensor:
+      def __getitem__(self, key):
+        if key.stop is not None or key.step is not None:
+          raise AssertionError('Unexpected file slice')
+        raw = bytearray(path.read_bytes()[key.start:])
+        buffer = View(raw, len(raw), UCHAR, 0, owner, 'source')
+        return SimpleNamespace(uop=SimpleNamespace(buffer=buffer))
+
+    return FileTensor()
+
+
+class View:
+  def __init__(self, storage, size, dtype, offset, owner, kind):
+    self.storage, self.size, self.dtype, self.offset = storage, size, dtype, offset
+    self.owner, self.kind = owner, kind
+    self.nbytes = size * dtype.itemsize
+
+  def ensure_allocated(self):
+    return self
+
+  def view(self, size, dtype, offset):
+    if offset < 0 or offset + size * dtype.itemsize > self.nbytes:
+      raise AssertionError('Out of bounds fixture view')
+    result = View(self.storage, size, dtype, self.offset + offset, self.owner, self.kind)
+    if self.kind == 'source':
+      self.owner.source_views.append((size, dtype.name, result.offset))
+    return result
+
+  def copy_from(self, source):
+    if self.nbytes != source.nbytes:
+      raise AssertionError('Copy size mismatch')
+    self.owner.copies.append((self.offset, source.offset, self.nbytes, self.dtype.name, source.dtype.name))
+    self.storage[self.offset:self.offset + self.nbytes] = source.storage[source.offset:source.offset + source.nbytes]
+    return self
+
+
+def reader(path, name):
+  node = next(item for item in ast.parse(path.read_text()).body if isinstance(item, ast.FunctionDef) and item.name == name)
+  before = ast.dump(node, include_attributes=False)
+  namespace = {'Path': Path, 'pathlib': SimpleNamespace(Path=Path), 'io': io, 'pickle': pickle, 'struct': struct, 'cast': cast}
+  exec(compile(ast.Module(body=[node], type_ignores=[]), str(path), 'exec'), namespace)
+  assert ast.dump(node, include_attributes=False) == before
+  return namespace[name]
+
+
+def envelope(path, payload):
+  markers = (object(), object(), object())
+  records = ((32, UINT32, (8 << 20) - 48), (19, UCHAR, len(payload) - 19), (32, UINT32, (8 << 20) - 48))
+
+  class OOBPickler(pickle.Pickler):
+    def persistent_id(self, value):
+      for marker, record in zip(markers, records, strict=True):
+        if value is marker:
+          return record
+      return None
+
+  stream = io.BytesIO()
+  OOBPickler(stream).dump({'views': markers, 'metadata': {'abi': 'fixture', 'shape': (3, 7)}})
+  opcodes = stream.getvalue()
+  path.write_bytes(struct.pack('<q', len(opcodes)) + opcodes + payload)
+  return opcodes, records
+
+
+def invoke(function, path, legacy=False):
+  owner = Storage()
+  tinygrad, device, dtype = (ModuleType(name) for name in ('tinygrad', 'tinygrad.device', 'tinygrad.dtype'))
+  tinygrad.Device, tinygrad.Tensor = SimpleNamespace(DEFAULT='CPU:LLVM'), owner.tensor
+  device.Buffer, dtype.dtypes = owner.buffer, SimpleNamespace(uchar=UCHAR)
+  saved = dict(sys.modules)
+  with patch.dict(sys.modules, {'tinygrad': tinygrad, 'tinygrad.device': device, 'tinygrad.dtype': dtype}):
+    # Legacy loader imports these names at module level; supply only the same
+    # fake metadata globals to its unchanged extracted function.
+    globals_ = function.__globals__
+    with patch.dict(globals_, {'Device': tinygrad.Device, 'Tensor': tinygrad.Tensor,
+                              'Buffer': owner.buffer, 'dtypes': dtype.dtypes}):
+      result = function(path, out_of_band=True) if legacy else function(path)
+  if saved != dict(sys.modules):
+    raise AssertionError('Fake reader module custody changed')
+  return result, owner
+
+
+class TestChestnutReader(unittest.TestCase):
+  def sources(self):
+    root = Path(__file__).resolve().parents[4]
+    current = root / 'openpilot/starpilot/models/runtime_bridge/runtime.py'
+    legacy = root / 'tinygrad_model_runtime/legacy9d/tinygrad_repo/examples/openpilot/helpers.py'
+    return reader(current, '_load_chestnut'), reader(legacy, 'load_pickle')
+
+  def test_chunked_bytes_offsets_dtype_alias_and_tail_match_legacy(self):
+    current, legacy = self.sources()
+    size = 2 * (8 << 20) + 173
+    pattern = bytes((index * 37 + 11) % 256 for index in range(251))
+    payload = (pattern * (size // len(pattern) + 1))[:size]
+    with tempfile.TemporaryDirectory() as temporary:
+      path = Path(temporary) / 'captured.pkl'
+      opcodes, records = envelope(path, payload)
+      old, old_store = invoke(legacy, path, legacy=True)
+      new, store = invoke(current, path)
+      self.assertEqual(path.read_bytes()[8:8 + len(opcodes)], opcodes)
+      self.assertEqual(new['metadata'], old['metadata'])
+      self.assertEqual(new['metadata'], {'abi': 'fixture', 'shape': (3, 7)})
+      self.assertEqual(store.copies, [(0, 0, 8 << 20, 'uchar', 'uchar'),
+                                    (8 << 20, 8 << 20, 8 << 20, 'uchar', 'uchar'),
+                                    (16 << 20, 16 << 20, 173, 'uchar', 'uchar')])
+      self.assertEqual(old_store.copies, [(0, 0, size, 'uchar', 'uchar')])
+      self.assertEqual(sum(row[2] for row in store.copies), size)
+      self.assertEqual(store.source_views, [(row[2], 'uchar', row[1]) for row in store.copies])
+      self.assertEqual(len(store.arenas), 1)
+      device, count, dtype, preallocate, arena = store.arenas[0]
+      self.assertEqual((device, count, dtype.name, preallocate), ('CPU:LLVM', size, 'uchar', True))
+      self.assertEqual(bytes(arena.storage), payload)
+      self.assertEqual(bytes(old_store.arenas[0][4].storage), payload)
+      for expected, old_view, view in zip(records, old['views'], new['views'], strict=True):
+        count, dtype, offset = expected
+        self.assertEqual((view.size, view.dtype.name, view.offset), (count, dtype.name, offset))
+        self.assertEqual((old_view.size, old_view.dtype.name, old_view.offset), (count, dtype.name, offset))
+        self.assertIs(view.storage, arena.storage)
+        self.assertEqual(bytes(view.storage[offset:offset + view.nbytes]), payload[offset:offset + view.nbytes])
+      # Two independently unpickled persistent records alias the same arena.
+      first, _, alias = new['views']
+      first.storage[first.offset] ^= 0xff
+      self.assertEqual(alias.storage[alias.offset], payload[first.offset] ^ 0xff)
+      self.assertEqual(bytes(new['views'][1].storage[-19:]), payload[-19:])
+
+  def test_opcode_only_pickle_does_not_allocate_an_arena(self):
+    current, legacy = self.sources()
+    expected = {'metadata': (1, 'no arena'), 'values': [3, 5, 8]}
+    opcodes = pickle.dumps(expected)
+    with tempfile.TemporaryDirectory() as temporary:
+      path = Path(temporary) / 'empty-arena.pkl'
+      path.write_bytes(struct.pack('<q', len(opcodes)) + opcodes)
+      for function, selected_legacy in ((current, False), (legacy, True)):
+        with self.subTest(legacy=selected_legacy):
+          result, store = invoke(function, path, legacy=selected_legacy)
+          self.assertEqual(result, expected)
+          self.assertFalse(store.arenas)
+          self.assertFalse(store.copies)
+          self.assertFalse(store.source_views)
+
+  def test_zero_single_byte_and_chunk_boundary_lengths(self):
+    current, legacy = self.sources()
+    chunk = 8 << 20
+    opcodes = pickle.dumps({'boundary': 'same opcode envelope'})
+    pattern = bytes(range(251))
+    with tempfile.TemporaryDirectory() as temporary:
+      path = Path(temporary) / 'boundary.pkl'
+      for size in (0, 1, chunk - 1, chunk, chunk + 1):
+        with self.subTest(size=size):
+          payload = (pattern * (size // len(pattern) + 1))[:size]
+          path.write_bytes(struct.pack('<q', len(opcodes)) + opcodes + payload)
+          old, old_store = invoke(legacy, path, legacy=True)
+          new, store = invoke(current, path)
+          self.assertEqual(new, old)
+          self.assertEqual(new, {'boundary': 'same opcode envelope'})
+          expected = [(offset, offset, min(chunk, size - offset), 'uchar', 'uchar')
+                      for offset in range(0, size, chunk)]
+          self.assertEqual(store.copies, expected)
+          self.assertEqual(store.source_views, [(row[2], 'uchar', row[1]) for row in expected])
+          self.assertEqual(sum(row[2] for row in store.copies), size)
+          if size:
+            self.assertEqual(old_store.copies, [(0, 0, size, 'uchar', 'uchar')])
+            self.assertEqual(bytes(store.arenas[0][4].storage), payload)
+            self.assertEqual(bytes(old_store.arenas[0][4].storage), payload)
+            self.assertEqual(store.arenas[0][4].storage[-1], payload[-1])
+          else:
+            self.assertFalse(store.arenas)
+            self.assertFalse(old_store.arenas)
+            self.assertFalse(old_store.copies)
+
+  def test_truncated_envelope_header_rejects_before_copy(self):
+    current, legacy = self.sources()
+    with tempfile.TemporaryDirectory() as temporary:
+      path = Path(temporary) / 'truncated.pkl'
+      path.write_bytes(b'\x01\x00')
+      for function, selected_legacy in ((current, False), (legacy, True)):
+        with self.subTest(legacy=selected_legacy), self.assertRaises(struct.error):
+          invoke(function, path, legacy=selected_legacy)
+
+
+if __name__ == '__main__':
+  unittest.main()

openpilot/starpilot/models/tests/test_stock_model_error_state.py

@@ -8,6 +8,7 @@ import __future__
 import ast
 from pathlib import Path
 from types import SimpleNamespace
+from typing import cast
 import unittest
 from unittest.mock import patch
 
@@ -24,8 +25,8 @@ def load_subject(device):
   methods = [node for node in cls.body if isinstance(node, ast.FunctionDef) and node.name in ('run', 'slice_outputs')]
   assert {node.name for node in methods} == {'run', 'slice_outputs'}
   original = {node.name: ast.dump(node, include_attributes=False) for node in methods}
-  namespace = {'np': np, 'Device': device, 'selected_legacy': runtime.selected_legacy, 'SEND_RAW_PRED': False}
-  module = ast.Module(body=list(methods), type_ignores=[])
+  namespace = {'np': np, 'Device': device, 'selected_legacy': runtime.selected_legacy, 'SEND_RAW_PRED': False, 'cast': cast}
+  module = ast.Module(body=[*methods], type_ignores=[])
   exec(compile(module, str(path), 'exec', flags=__future__.annotations.compiler_flag), namespace)
   # AST nodes are consumed unchanged; no rewritten expression or surrogate run.
   assert original == {node.name: ast.dump(node, include_attributes=False) for node in methods}

prebuilt.json

@@ -1 +1 @@
-{"protected_inventory_sha256": "e9af39a992a51f2ebd7998a963134538d01d2e54aa92d91f924cb4be745e172a", "schema_version": 1}
+{"protected_inventory_sha256": "82bba6649f143d9f5c68183db87c01f68fd8539a93b8abab02bb104682f8c8c6", "schema_version": 1}