From bd764bbf2e2df4a49695b1d4268fafd6e9791399 Mon Sep 17 00:00:00 2001 From: firestar5683 <168790843+firestar5683@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:52:19 -0500 Subject: [PATCH] Reduce Chestnut model startup overhead --- openpilot/selfdrive/modeld/modeld.py | 10 +- openpilot/starpilot/models/compile.py | 9 +- .../models/runtime_bridge/guard-policy.json | 2 +- .../models/runtime_bridge/legacy_types.py | 40 ++++ .../starpilot/models/runtime_bridge/qcom.py | 14 +- .../models/runtime_bridge/runtime.py | 50 +++- .../starpilot/models/runtime_bridge/usb.py | 27 ++- .../models/tests/runtime_bridge_fixtures.py | 10 +- .../models/tests/test_chestnut_reader.py | 215 ++++++++++++++++++ .../tests/test_stock_model_error_state.py | 5 +- prebuilt.json | 2 +- 11 files changed, 350 insertions(+), 34 deletions(-) create mode 100644 openpilot/starpilot/models/runtime_bridge/legacy_types.py create mode 100644 openpilot/starpilot/models/tests/test_chestnut_reader.py diff --git a/openpilot/selfdrive/modeld/modeld.py b/openpilot/selfdrive/modeld/modeld.py index 9ce6e512e5..f7a4bda3b5 100755 --- a/openpilot/selfdrive/modeld/modeld.py +++ b/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() diff --git a/openpilot/starpilot/models/compile.py b/openpilot/starpilot/models/compile.py index cf86f0b287..c00e06181a 100644 --- a/openpilot/starpilot/models/compile.py +++ b/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) diff --git a/openpilot/starpilot/models/runtime_bridge/guard-policy.json b/openpilot/starpilot/models/runtime_bridge/guard-policy.json index 728c27dd7a..821fd52d84 100644 --- a/openpilot/starpilot/models/runtime_bridge/guard-policy.json +++ b/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"} diff --git a/openpilot/starpilot/models/runtime_bridge/legacy_types.py b/openpilot/starpilot/models/runtime_bridge/legacy_types.py new file mode 100644 index 0000000000..5ff8bef04b --- /dev/null +++ b/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] diff --git a/openpilot/starpilot/models/runtime_bridge/qcom.py b/openpilot/starpilot/models/runtime_bridge/qcom.py index ee97f07cb8..8808bd7636 100644 --- a/openpilot/starpilot/models/runtime_bridge/qcom.py +++ b/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: diff --git a/openpilot/starpilot/models/runtime_bridge/runtime.py b/openpilot/starpilot/models/runtime_bridge/runtime.py index da2f52bb0c..d7d6f02a86 100644 --- a/openpilot/starpilot/models/runtime_bridge/runtime.py +++ b/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(' "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': diff --git a/openpilot/starpilot/models/tests/runtime_bridge_fixtures.py b/openpilot/starpilot/models/tests/runtime_bridge_fixtures.py index ac9771d77e..08d97b949d 100644 --- a/openpilot/starpilot/models/tests/runtime_bridge_fixtures.py +++ b/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, } diff --git a/openpilot/starpilot/models/tests/test_chestnut_reader.py b/openpilot/starpilot/models/tests/test_chestnut_reader.py new file mode 100644 index 0000000000..19b5e82240 --- /dev/null +++ b/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('