From d4a2036bf7388e622c5bdf2989c2c8d665715897 Mon Sep 17 00:00:00 2001 From: discountchubbs Date: Mon, 14 Sep 2026 21:48:21 -0700 Subject: [PATCH] modeld_v2: adapt to state pairs in onnx --- openpilot/selfdrive/modeld/compile_modeld.py | 123 +++++------------- openpilot/selfdrive/modeld/modeld.py | 104 +++------------ .../sunnypilot/modeld_v2/compile_modeld.py | 62 +++++---- openpilot/sunnypilot/modeld_v2/modeld.py | 42 +++--- .../modeld_v2/stock_dependencies.py | 119 +++++++++++++++++ 5 files changed, 234 insertions(+), 216 deletions(-) create mode 100644 openpilot/sunnypilot/modeld_v2/stock_dependencies.py diff --git a/openpilot/selfdrive/modeld/compile_modeld.py b/openpilot/selfdrive/modeld/compile_modeld.py index d851c3cc86..c2fd116f1a 100755 --- a/openpilot/selfdrive/modeld/compile_modeld.py +++ b/openpilot/selfdrive/modeld/compile_modeld.py @@ -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,24 +248,28 @@ 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) + with open(args.output, "rb") as f: + load_oob(f) + assert not f.read(1), "unexpected model buffer data" print(f"Saved JITs to {args.output} ({os.path.getsize(args.output) / 1e6:.2f} MB)") diff --git a/openpilot/selfdrive/modeld/modeld.py b/openpilot/selfdrive/modeld/modeld.py index b98b4d0f58..af2e60a0d2 100755 --- a/openpilot/selfdrive/modeld/modeld.py +++ b/openpilot/selfdrive/modeld/modeld.py @@ -5,8 +5,6 @@ from functools import cached_property import os os.environ['GMMU'] = '0' # for chestnut fast loading, noop for qcom from tinygrad.device import Device -import usb1 -import struct import threading import time import numpy as np @@ -28,12 +26,11 @@ 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.common.hardware.usb import CHESTNUT_USB_IDS from openpilot.selfdrive.modeld.constants import ModelConstants, Plan -from openpilot.selfdrive.modeld.helpers import chestnut_present, chestnut_compiled, chestnut_ready, modeld_pkl_path, load_oob +from openpilot.selfdrive.modeld.helpers import chestnut_present, chestnut_compiled, modeld_pkl_path, load_oob from openpilot.sunnypilot.livedelay.helpers import get_lat_delay from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase @@ -75,45 +72,14 @@ def get_action_from_model(model_output: dict[str, np.ndarray], prev_action: log. shouldStop=bool(stop)) -class ChestnutState: - # only modeld can access chestnut +class ChestnutGpuState: + # GPU metrics require modeld's GPU context def __init__(self, pm: PubMaster, big: bool): self.pm = pm self.big = big self.valid = True self.sends = 0 self.metrics = {} - self._asm_usb = None - - def _close_asm_usb(self) -> None: - if self._asm_usb is not None: - self._asm_usb.close() - self._asm_usb = None - - def _open_asm_usb(self): - context = usb1.USBContext() - for vendor_id, product_id in CHESTNUT_USB_IDS: - if (handle := context.openByVendorIDAndProductID(vendor_id, product_id, skip_on_error=True)) is not None: - return handle - context.close() - - def _read_ina(self) -> tuple[int, int, bool]: - if "AMD" in Device._opened_devices and self._asm_usb is None: - try: - raw = Device["AMD"].iface.pci_dev.usb.usb.control_read(0xC0, 5) - return struct.unpack(' int: @@ -121,8 +87,8 @@ class ChestnutState: return smu._send_msg(smu.smu_mod.PPSMC_MSG_GetPptLimit, 0, read_back_arg=True, timeout=100) def send(self) -> None: - msg = messaging.new_message('chestnutState') - state = msg.chestnutState + msg = messaging.new_message('chestnutGpuState') + state = msg.chestnutGpuState self.sends += 1 if self.big and "AMD" in Device._opened_devices and self.sends % 100 == 1: try: @@ -148,21 +114,8 @@ class ChestnutState: for k, v in self.metrics.items(): setattr(state, k, v) - asm_valid = False - try: - # ASM runs on USB-C power, these still read without a gpu - state.supplyVoltage, state.supplyCurrent, state.supplyFault = self._read_ina() - asm_valid = True - except Exception: - pass - if "AMD" in Device._opened_devices: - try: - state.pcieLtssm = Device["AMD"].iface.pci_dev.usb.read(0xB450, 1)[0] - except Exception: - pass - - msg.valid = asm_valid and (not self.big or self.valid) - self.pm.send('chestnutState', msg) + msg.valid = not self.big or (self.valid and bool(self.metrics)) + self.pm.send('chestnutGpuState', msg) class FrameMeta: @@ -184,17 +137,17 @@ class ModelState(ModelStateBase): 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)] @@ -216,14 +169,13 @@ class ModelState(ModelStateBase): 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() @@ -235,32 +187,19 @@ class ModelState(ModelStateBase): 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 def main(demo=False): cloudlog.warning("modeld init") - chestnut_available = chestnut_present() and chestnut_compiled() - CHESTNUT = False - if chestnut_available: - poller = messaging.Poller() - sock = messaging.sub_sock("chestnutState", poller=poller, conflate=True) - deadline = time.monotonic() + 4. / SERVICE_LIST['deviceState'].frequency - while not CHESTNUT and (remaining := deadline - time.monotonic()) > 0.: - if not poller.poll(round(remaining * 1000)): - break - msg = messaging.recv_one_or_none(sock) - CHESTNUT = msg is not None and msg.valid and chestnut_ready(msg.chestnutState) + CHESTNUT = chestnut_present() and chestnut_compiled() if CHESTNUT: os.environ['HCQDEV_WAIT_TIMEOUT_MS'] = '3000' params = Params() params.put_bool("ChestnutLoading", CHESTNUT) - if chestnut_available and not CHESTNUT: - params.put_bool("ChestnutActive", False) - else: - params.remove("ChestnutActive") + params.remove("ChestnutActive") config_realtime_process(7, 54) @@ -304,11 +243,7 @@ def main(demo=False): loader.start() loader.join(BIG_MODEL_TIMEOUT) model = big_model - if model is None: - params.put_bool("ChestnutModelError", True) params.put_bool("ChestnutActive", model is not None) - if model is not None: - params.remove("ChestnutModelError") small_model = ModelState(vipc_client_main.width, vipc_client_main.height, False) if model is None or CHESTNUT else None if model is None: @@ -318,13 +253,13 @@ def main(demo=False): cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting") # messaging - pub_socks = ["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"] + (["chestnutState"] if CHESTNUT else []) + pub_socks = ["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"] + (["chestnutGpuState"] if CHESTNUT else []) pm = PubMaster(pub_socks) sm = SubMaster(["deviceState", "carState", "narrowRoadCameraState", "extrinsicsCalibration", "driverMonitoringState", "carControl", "lateralDelay"]) publish_state = PublishState() params = Params() - chestnut_state = ChestnutState(pm, model.chestnut) if CHESTNUT else None + chestnut_state = ChestnutGpuState(pm, model.chestnut) if CHESTNUT else None # setup filter to track dropped frames frame_dropped_filter = FirstOrderFilter(0., 10., 1. / ModelConstants.MODEL_RUN_FREQ) @@ -435,14 +370,13 @@ def main(demo=False): mt1 = time.perf_counter() try: send_chestnut = (chestnut_state is not None and - run_count % round(ModelConstants.MODEL_RUN_FREQ / SERVICE_LIST['chestnutState'].frequency) == 0) + run_count % round(ModelConstants.MODEL_RUN_FREQ / SERVICE_LIST['chestnutGpuState'].frequency) == 0) model_output = model.run(bufs, transforms, inputs, chestnut_state.send if send_chestnut else None) except Exception: if not params.get_bool("ChestnutActive"): raise # fallback to small model cloudlog.exception("big model failed, fall back to small") - params.put_bool("ChestnutModelError", True) params.put_bool("ChestnutActive", False) assert small_model is not None model = small_model diff --git a/openpilot/sunnypilot/modeld_v2/compile_modeld.py b/openpilot/sunnypilot/modeld_v2/compile_modeld.py index 1e88769928..71a90eb720 100755 --- a/openpilot/sunnypilot/modeld_v2/compile_modeld.py +++ b/openpilot/sunnypilot/modeld_v2/compile_modeld.py @@ -33,6 +33,7 @@ def _patch_tinygrad_fetch_fw(): _patch_tinygrad_fetch_fw() import openpilot.selfdrive.modeld.compile_modeld as stock +import openpilot.sunnypilot.modeld_v2.stock_dependencies as legacy from tinygrad import dtypes from tinygrad.device import Device from tinygrad.engine.jit import TinyJit @@ -41,7 +42,7 @@ from tinygrad.tensor import Tensor MODEL_TYPES = ('vision_policy', 'supercombo', 'vision_multi_policy') WARP_INPUTS = ['tfm', 'big_tfm'] POLICY_INPUTS = ['img_q', 'big_img_q', 'feat_q', 'desire_q', 'packed_npy_inputs'] -nv12_copy_size = stock.nv12_copy_size + def _detect_desire_key(shapes: dict) -> str | None: return next((key for key in shapes if key.startswith('desire')), None) @@ -152,8 +153,8 @@ def make_warp_queues(device=Device.DEFAULT): def make_run_policy(vision_runner, policy_runners: list, features_slice: slice, frame_skip: int, input_shapes: dict): - sample_skip_fn = partial(stock.sample_skip, frame_skip=frame_skip) - sample_desire_fn = partial(stock.sample_desire, frame_skip=frame_skip) + sample_skip_fn = partial(legacy.sample_skip, frame_skip=frame_skip) + sample_desire_fn = partial(legacy.sample_desire, frame_skip=frame_skip) desire_key = _detect_desire_key(input_shapes) road_key, wide_key = _detect_vision_keys(input_shapes) @@ -170,14 +171,14 @@ def make_run_policy(vision_runner, policy_runners: list, features_slice: slice, warped_dev = warped.to(Device.DEFAULT) Tensor.realize(packed_npy_inputs_dev, warped_dev) - img = stock.shift_and_sample(img_q, warped_dev[0:1], sample_skip_fn) - big_img = stock.shift_and_sample(big_img_q, warped_dev[1:2], sample_skip_fn) + img = legacy.shift_and_sample(img_q, warped_dev[0:1], sample_skip_fn) + big_img = legacy.shift_and_sample(big_img_q, warped_dev[1:2], sample_skip_fn) unpacked_tensors = [tensor.reshape(shape) for tensor, shape in zip(packed_npy_inputs_dev.split(npy_sizes), npy_shapes.values(), strict=True)] unpacked_dict = dict(zip(npy_shapes.keys(), unpacked_tensors, strict=True)) desire_dev = unpacked_dict['desire'] - desire_buf = stock.shift_and_sample(desire_q, desire_dev.reshape(1, 1, -1), sample_desire_fn) + desire_buf = legacy.shift_and_sample(desire_q, desire_dev.reshape(1, 1, -1), sample_desire_fn) inputs = {desire_key: desire_buf} for key, tensor_val in unpacked_dict.items(): @@ -186,13 +187,13 @@ def make_run_policy(vision_runner, policy_runners: list, features_slice: slice, if 'prev_feat' in unpacked_dict: prev_feat_dev = unpacked_dict['prev_feat'] - inputs['features_buffer'] = stock.shift_and_sample(feat_q, prev_feat_dev.reshape(1, 1, -1), sample_skip_fn).reshape(input_shapes['features_buffer']) + inputs['features_buffer'] = legacy.shift_and_sample(feat_q, prev_feat_dev.reshape(1, 1, -1), sample_skip_fn).reshape(input_shapes['features_buffer']) if vision_runner: vision_out_cast = next(iter(vision_runner({road_key: img, wide_key: big_img}).values())).cast('float32').realize() if 'features_buffer' not in inputs: new_feat = vision_out_cast[:, features_slice].reshape(1, -1).unsqueeze(0) - inputs['features_buffer'] = stock.shift_and_sample(feat_q, new_feat, sample_skip_fn).realize() + inputs['features_buffer'] = legacy.shift_and_sample(feat_q, new_feat, sample_skip_fn).realize() policy_outs = [next(iter(pol_runner(inputs).values())).cast('float32').realize() for pol_runner in policy_runners] return (vision_out_cast, *policy_outs) if len(policy_outs) > 1 else (vision_out_cast, policy_outs[0]) @@ -203,7 +204,7 @@ def make_run_policy(vision_runner, policy_runners: list, features_slice: slice, policy_out = next(iter(policy_runners[0](inputs).values())).cast('float32').realize() if 'features_buffer' not in inputs and features_slice is not None: new_feat = policy_out[:, features_slice].reshape(1, -1).unsqueeze(0) - stock.shift_and_sample(feat_q, new_feat, sample_skip_fn).realize() + legacy.shift_and_sample(feat_q, new_feat, sample_skip_fn).realize() return policy_out return run_policy @@ -330,16 +331,32 @@ if __name__ == "__main__": output_data['run_model'] = {} derived_frame_skip = args.frame_skip or derive_frame_skip({}, model_metadata['input_shapes']) model_runner = OnnxRunner(args.supercombo_onnx) - run_policy = stock.make_run_policy(model_runner, model_metadata, derived_frame_skip) - for cam_w, cam_h in args.camera_resolutions: - print(f"Compiling unified run_model JIT for {cam_w}x{cam_h}...") - nv12 = stock.NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)) - frame_copy_size = stock.nv12_copy_size(nv12.stride, nv12.y_height, nv12.uv_height) - make_model_queues = partial(stock.make_input_queues, model_metadata['input_shapes'], derived_frame_skip, - frame_copy_size=frame_copy_size) - warp = stock.make_warp(nv12, model_w, model_h) - run_model_jit = TinyJit(stock.make_run_model(warp, run_policy, model_metadata, frame_copy_size), prune=True) - output_data['run_model'][(cam_w, cam_h)] = compile_jit(run_model_jit, stock.MODELD_INPUTS, make_model_queues, benchmark_runs=args.benchmark_runs) + new_img_model = 'new_img' in model_runner.graph_inputs + + if new_img_model: + 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} + output_data['metadata'] = {'model': model_metadata, **model_metadata, 'input_shapes': input_shapes, 'state_pairs': state_pairs} + for cam_w, cam_h in args.camera_resolutions: + print(f"Compiling unified run_model JIT for {cam_w}x{cam_h} (new architecture)...") + nv12 = stock.NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)) + frame_copy_size = stock.nv12_copy_size(nv12.stride, nv12.y_height, nv12.uv_height) + make_model_queues = partial(stock.make_input_queues, input_shapes, state_pairs, frame_copy_size=frame_copy_size) + warp = stock.make_warp(nv12, model_w, model_h) + run_model_jit = TinyJit(stock.make_run_model(warp, model_runner, input_shapes, state_pairs, frame_copy_size), prune=True) + output_data['run_model'][(cam_w, cam_h)] = compile_jit(run_model_jit, list(state_pairs.keys()) + ['packed_npy_inputs'], make_model_queues, + benchmark_runs=args.benchmark_runs) + else: + run_policy = legacy.make_legacy_run_policy(model_runner, model_metadata, derived_frame_skip) + for cam_w, cam_h in args.camera_resolutions: + print(f"Compiling unified run_model JIT for {cam_w}x{cam_h}...") + nv12 = stock.NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)) + frame_copy_size = stock.nv12_copy_size(nv12.stride, nv12.y_height, nv12.uv_height) + make_model_queues = partial(stock.make_input_queues, model_metadata['input_shapes'], derived_frame_skip, + frame_copy_size=frame_copy_size) + warp = stock.make_warp(nv12, model_w, model_h) + run_model_jit = TinyJit(legacy.make_legacy_run_model(warp, run_policy, model_metadata, frame_copy_size), prune=True) + output_data['run_model'][(cam_w, cam_h)] = compile_jit(run_model_jit, POLICY_INPUTS, make_model_queues, benchmark_runs=args.benchmark_runs) else: vision_runner = OnnxRunner(args.vision_onnx) if args.vision_onnx else None if args.model_type == 'vision_policy': @@ -355,13 +372,12 @@ if __name__ == "__main__": output_data['metadata'][name] = make_metadata_dict(runner_arg) policy_keys = [key for key in output_data['metadata'].keys() if key != 'vision'] - first_policy_meta = output_data['metadata'][policy_keys[0]] if policy_keys else {} - vision_meta = output_data['metadata'].get('vision', {}) + first_policy_meta: dict = output_data['metadata'][policy_keys[0]] if policy_keys else {} + vision_meta: dict = output_data['metadata'].get('vision', {}) derived_frame_skip = args.frame_skip or derive_frame_skip(vision_meta.get('input_shapes', {}), first_policy_meta.get('input_shapes', {})) all_shapes = {key: value for meta in output_data['metadata'].values() for key, value in meta['input_shapes'].items()} - feat_meta = output_data['metadata'].get('vision') or output_data['metadata'].get('policy') - assert feat_meta is not None + feat_meta: dict = vision_meta or first_policy_meta features_slice = feat_meta['output_slices']['hidden_state'] print(f"Compiling run_policy JIT (model_size={model_w}x{model_h}, frame_skip={derived_frame_skip})...") diff --git a/openpilot/sunnypilot/modeld_v2/modeld.py b/openpilot/sunnypilot/modeld_v2/modeld.py index 6e31045b89..05d07f20d6 100755 --- a/openpilot/sunnypilot/modeld_v2/modeld.py +++ b/openpilot/sunnypilot/modeld_v2/modeld.py @@ -36,12 +36,9 @@ from openpilot.system import sentry from openpilot.system.camerad.cameras.nv12_info import get_nv12_info from openpilot.selfdrive.controls.lib.desire_helper import DesireHelper from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_from_plan, smooth_value -from openpilot.selfdrive.modeld.modeld import ChestnutState +from openpilot.selfdrive.modeld.modeld import ChestnutGpuState -from openpilot.selfdrive.modeld.compile_modeld import ( - MODELD_INPUTS, - make_input_queues as make_stock_input_queues, -) +from openpilot.selfdrive.modeld.compile_modeld import make_input_queues from openpilot.sunnypilot.modeld_v2.fill_model_msg import fill_model_msg, fill_pose_msg, PublishState, get_curvature_from_output from openpilot.sunnypilot.modeld_v2.parse_model_outputs import Parser from openpilot.sunnypilot.modeld_v2.constants import ModelConstants, Plan @@ -50,6 +47,7 @@ from openpilot.sunnypilot.modeld_v2.camera_offset_helper import CameraOffsetHelp from openpilot.sunnypilot.modeld_v2.compile_modeld import (derive_frame_skip, make_split_input_queues, make_supercombo_input_queues, nv12_copy_size, WARP_INPUTS, POLICY_INPUTS) +from openpilot.sunnypilot.modeld_v2.stock_dependencies import make_legacy_stock_input_queues from openpilot.sunnypilot.livedelay.helpers import get_lat_delay from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase from openpilot.sunnypilot.modeld_v2.helpers import load_oob @@ -141,8 +139,14 @@ class ModelState(ModelStateBase): self._vision_input_names = [key for key in self.input_shapes if 'img' in key] self.frame_skip = derive_frame_skip({}, self.input_shapes) if self.is_run_model: - self.input_queues, self.numpy_inputs, self.frame_buffers = make_stock_input_queues( - self.input_shapes, self.frame_skip, device=self.DEV, frame_copy_size=self.frame_copy_size) + self.state_pairs = model_metadata.get('state_pairs', {}) + self.is_new_model = len(self.state_pairs) > 0 + if self.is_new_model: + self.input_queues, self.numpy_inputs, self.frame_buffers = make_input_queues(self.input_shapes, self.state_pairs, + device=self.DEV, frame_copy_size=self.frame_copy_size) + else: + self.input_queues, self.numpy_inputs, self.frame_buffers = make_legacy_stock_input_queues(self.input_shapes, self.frame_skip, device=self.DEV, + frame_copy_size=self.frame_copy_size) self.frame_views, self.npy = self.frame_buffers, self.numpy_inputs self.run_model, self.run_policy, self.warp = jits['run_model'][(cam_w, cam_h)], None, None else: @@ -191,8 +195,12 @@ class ModelState(ModelStateBase): dummy_inputs = {k: np.zeros(v.shape, dtype=v.dtype) for k, v in self.numpy_inputs.items() if k not in ['tfm', 'big_tfm', 'prev_feat']} self.run(dummy_frames, transforms, dummy_inputs) if self.is_run_model: - self.input_queues, self.numpy_inputs, self.frame_buffers = make_stock_input_queues( - self.input_shapes, self.frame_skip, device=self.DEV, frame_copy_size=self.frame_copy_size) + if self.is_new_model: + self.input_queues, self.numpy_inputs, self.frame_buffers = make_input_queues(self.input_shapes, self.state_pairs, device=self.DEV, + frame_copy_size=self.frame_copy_size) + else: + self.input_queues, self.numpy_inputs, self.frame_buffers = make_legacy_stock_input_queues(self.input_shapes, self.frame_skip, device=self.DEV, + frame_copy_size=self.frame_copy_size) self.frame_views = self.frame_buffers self.npy = self.numpy_inputs else: @@ -242,7 +250,10 @@ class ModelState(ModelStateBase): self.numpy_inputs['big_tfm'][:, :] = transforms[self._wide_key].reshape(3, 3) if self.run_model is not None: - outs, = self.run_model(**{k: self.input_queues[k] for k in MODELD_INPUTS}) + if self.is_new_model: + outs, = self.run_model(**self.input_queues) + else: + outs, = self.run_model(**{k: self.input_queues[k] for k in POLICY_INPUTS}) raw_outputs = outs else: assert self.warp is not None and self.run_policy is not None @@ -373,11 +384,7 @@ def main(demo=False): loader.start() loader.join(BIG_MODEL_TIMEOUT) model = big_model - if model is None: - params.put_bool("ChestnutModelError", True) params.put_bool("ChestnutActive", model is not None) - if model is not None: - params.remove("ChestnutModelError") small_model = ModelState(cam_w=vipc_client_main.width, cam_h=vipc_client_main.height, chestnut=False) if model is None or CHESTNUT else None if model is None: @@ -387,12 +394,12 @@ def main(demo=False): cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting") # messaging - pub_socks = ["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"] + (["chestnutState"] if CHESTNUT else []) + pub_socks = ["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"] + (["chestnutGpuState"] if CHESTNUT else []) pm = PubMaster(pub_socks) sm = SubMaster(["deviceState", "carState", "narrowRoadCameraState", "extrinsicsCalibration", "driverMonitoringState", "carControl", "lateralDelay"]) publish_state = PublishState() - chestnut_state = ChestnutState(pm, model.chestnut) if CHESTNUT else None + chestnut_state = ChestnutGpuState(pm, model.chestnut) if CHESTNUT else None # setup filter to track dropped frames frame_dropped_filter = FirstOrderFilter(0., 10., 1. / model.constants.MODEL_FREQ) @@ -514,13 +521,12 @@ def main(demo=False): mt1 = time.perf_counter() try: send_chestnut = (chestnut_state is not None and - run_count % round(model.constants.MODEL_FREQ / SERVICE_LIST['chestnutState'].frequency) == 0) + run_count % round(model.constants.MODEL_FREQ / SERVICE_LIST['chestnutGpuState'].frequency) == 0) model_output = model.run(bufs, transforms, inputs, chestnut_state.send if send_chestnut else None) except Exception: if not params.get_bool("ChestnutActive"): raise cloudlog.exception("chestnut failed, falling back to small") - params.put_bool("ChestnutModelError", True) params.put_bool("ChestnutActive", False) assert small_model is not None model = small_model diff --git a/openpilot/sunnypilot/modeld_v2/stock_dependencies.py b/openpilot/sunnypilot/modeld_v2/stock_dependencies.py new file mode 100644 index 0000000000..84aa0f63b6 --- /dev/null +++ b/openpilot/sunnypilot/modeld_v2/stock_dependencies.py @@ -0,0 +1,119 @@ +""" +Copyright (c) 2021-, Haibin Wen, sunnypilot, and a number of other contributors. + +This file is part of sunnypilot and is licensed under the MIT License. +See the LICENSE.md file in the root directory for more details. +""" + +import math +import numpy as np +from functools import partial +from tinygrad import dtypes +from tinygrad.device import Device +from tinygrad.tensor import Tensor + +# The old openpilot/selfdrive/modeld/compile_modeld.py functions needed for legacy models +# We freeze them here so they aren't lost. + +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 _detect_desire_key(shapes: dict) -> str | None: + return next((key for key in shapes if key.startswith('desire')), None) + +def get_policy_npy_shapes(input_shapes: dict, is_supercombo: bool = False) -> tuple[dict, list[int]]: + desire_key = _detect_desire_key(input_shapes) + shapes = {} + if desire_key: + shapes['desire'] = (input_shapes[desire_key][2],) + + for key, shape in input_shapes.items(): + if key not in (desire_key, 'features_buffer') and 'img' not in key: + shapes[key] = tuple(shape) + + if is_supercombo and 'features_buffer' in input_shapes: + fb = input_shapes['features_buffer'] + feat_dim = math.prod(fb[2:]) + shapes['prev_feat'] = (fb[0], feat_dim) + + sizes = [int(np.prod(size)) for size in shapes.values()] + return shapes, sizes + +def make_legacy_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'], is_supercombo=True) + model_input_dtypes = {name: spec.dtype for name, spec in model_runner.graph_inputs.items()} + + 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.get(name, dtypes.float32)) for name, value in inputs.items()} + out = next(iter(model_runner(inputs).values())).cast('float32') + return out, + return run_policy + +def make_legacy_run_model(warp, run_policy, model_metadata, frame_copy_size): + _, policy_sizes = get_policy_npy_shapes(model_metadata['input_shapes'], is_supercombo=True) + 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) + packed_npy_inputs = packed_input[:packed_npy_size].bitcast('float32') + 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) + return run_model + +def make_legacy_stock_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, is_supercombo=True) + shapes = {'tfm': (3, 3), 'big_tfm': (3, 3)} | policy_shapes + sizes = [math.prod(s) for s in shapes.values()] + 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(), + } + return input_queues, npy, frame_views