From b751b04cdad7dde9c738ede3cab79fa673d15541 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Harald=20Sch=C3=A4fer?= Date: Tue, 15 Sep 2026 11:21:47 -0700 Subject: [PATCH] Use tinygrad ONNX and warp compilers for driving and DM (#38922) * modeld: use tinygrad ONNX and warp compilers for driving and DM * Bump tinygrad with current openpilot model compile tests * modeld: prioritize tinygrad compiler imports in build environment * Use compiler branch based on openpilot's pinned tinygrad * modeld: invoke tinygrad compilers without firmware patch --- .gitmodules | 3 +- openpilot/selfdrive/modeld/SConscript | 112 ++++--------- openpilot/selfdrive/modeld/compile_modeld.py | 117 -------------- openpilot/selfdrive/modeld/compile_warp.py | 150 ------------------ .../selfdrive/modeld/dmonitoringmodeld.py | 22 +-- .../selfdrive/modeld/get_model_metadata.py | 55 ------- openpilot/selfdrive/modeld/helpers.py | 38 ----- openpilot/selfdrive/modeld/modeld.py | 4 +- tinygrad_repo | 2 +- 9 files changed, 50 insertions(+), 453 deletions(-) delete mode 100755 openpilot/selfdrive/modeld/compile_modeld.py delete mode 100644 openpilot/selfdrive/modeld/compile_warp.py delete mode 100755 openpilot/selfdrive/modeld/get_model_metadata.py diff --git a/.gitmodules b/.gitmodules index ad6530de9a..7e37596b0a 100644 --- a/.gitmodules +++ b/.gitmodules @@ -15,4 +15,5 @@ url = ../../commaai/teleoprtc [submodule "tinygrad"] path = tinygrad_repo - url = https://github.com/tinygrad/tinygrad.git + url = https://github.com/commaai/tinygrad.git + branch = openpilot-modeld-test diff --git a/openpilot/selfdrive/modeld/SConscript b/openpilot/selfdrive/modeld/SConscript index 6af7f09092..6ccf711c4c 100644 --- a/openpilot/selfdrive/modeld/SConscript +++ b/openpilot/selfdrive/modeld/SConscript @@ -1,18 +1,18 @@ import glob -import json import os import time from SCons.Script import Action, Value from openpilot.common.file_chunker import chunk_file, get_chunk_targets, get_existing_chunks from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE, DM_INPUT_SIZE -from openpilot.selfdrive.modeld.helpers import TG_INPUT_DEVICES_PATH, chestnut_present, modeld_pkl_path +from openpilot.selfdrive.modeld.helpers import chestnut_present, modeld_pkl_path from openpilot.system.camerad.cameras.nv12_info import get_nv12_info Import('env', 'arch') chunker_file = File("#openpilot/common/file_chunker.py") lenv = env.Clone() +lenv.PrependENVPath('PYTHONPATH', Dir('#tinygrad_repo').abspath) tinygrad_root = env.Dir("#").abspath tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + "/**", recursive=True, root_dir=tinygrad_root) @@ -26,62 +26,27 @@ def estimate_pickle_max_size(onnx_size): camera_configs = [(c.width, c.height) for c in (_ar_ox_fisheye, _os_fisheye)] if arch == 'comma_arm64': - tg_backend = 'QCOM' - tg_flags = f'DEV={tg_backend} IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1' + tg_flags = 'DEV=QCOM IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1' else: - tg_backend = 'METAL' if arch == 'Darwin' else 'CPU' # JIT=2 disables graph batching, which produces incorrect outputs after buffers change. tg_flags = 'DEV=METAL JIT=2' if arch == 'Darwin' else 'DEV=CPU:LLVM' -tg_devices = { # which device to put jit inputs to at runtime - 'openpilot.selfdrive.modeld.dmonitoringmodeld': { - 'default': {'DEV': tg_backend} - }, -} - CHESTNUT = chestnut_present() if CHESTNUT: chestnut_tg_flags = 'DEBUG=1 DEV=USB+AMD:LLVM FRAME_DEV=CPU FLOAT16=1 JIT_BATCH_SIZE=0 GMMU=0 TC_OPT=2 TC_MIN_GLOBALS=32' # the USB+AMD GPU takes an exclusive flock; serialize all targets that touch it chestnut_lock = File("models/.chestnut.lock").abspath -def write_tg_devices(target, source, env): - with open(str(target[0]), "w") as f: - json.dump(tg_devices, f) - f.write("\n") - -tg_devices_node = lenv.Command( - str(TG_INPUT_DEVICES_PATH), - [Value(tg_devices)], - write_tg_devices, -) - # tinygrad calls brew which needs a $HOME in the env mac_brew_string = f'HOME={os.path.expanduser("~")}' if arch == 'Darwin' else '' -modeld_dir = Dir("#openpilot/selfdrive/modeld").abspath -compile_modeld_script = [ - File(f"{modeld_dir}/compile_modeld.py"), - File(f"{modeld_dir}/get_model_metadata.py"), - File(f"{modeld_dir}/helpers.py"), - File("#openpilot/common/hardware/hw.py"), -] -compile_warp_script = [File(f"{modeld_dir}/compile_warp.py"), File(f"{modeld_dir}/helpers.py"), - File("#openpilot/system/camerad/cameras/nv12_info.py")] -model_w, model_h = MEDMODEL_INPUT_SIZE +warp_deps = [File("#openpilot/system/camerad/cameras/nv12_info.py")] +compiler = 'python3 -m examples.openpilot' +# CPU 7 is isolated with isolcpus on AGNOS, so explicitly pin the compiler to it. +taskset = 'taskset -c 7 ' if arch == 'comma_arm64' else '' -for chestnut in [False, True] if CHESTNUT else [False]: - target_pkl_path = File(modeld_pkl_path(chestnut)).abspath - file_prefix, cmd_flags = ('big_', chestnut_tg_flags) if chestnut else ('', tg_flags) - driving_onnx_deps = get_existing_chunks(File(f"models/{file_prefix}driving_supercombo.onnx").abspath) - # CPU 7 is isolated with isolcpus on AGNOS, so explicitly pin the compiler to it. - taskset = 'taskset -c 7 ' if arch == 'comma_arm64' else '' - cmd = (f'{cmd_flags} {mac_brew_string} {taskset}python3 {modeld_dir}/compile_modeld.py ' - f'--onnx {File(f"models/{file_prefix}driving_supercombo.onnx").abspath} ' - f'--output {target_pkl_path}') - onnx_sizes_sum = sum(os.path.getsize(f) for f in driving_onnx_deps) - chunk_targets = get_chunk_targets(target_pkl_path, estimate_pickle_max_size(onnx_sizes_sum)) - def do_compile(target, source, env, command=cmd, pkl=target_pkl_path, chunks=chunk_targets): +def chestnut_action(command, pkl=None, chunks=()): + def do_compile(target, source, env): from openpilot.system.hardware.chestnut.flash import link_up # chestnut can enumerate before its PCIe link is up due to varying 12V power behavior across cars for _ in range(10): @@ -95,58 +60,49 @@ for chestnut in [False, True] if CHESTNUT else [False]: return ret if chunks: chunk_file(pkl, chunks) + return Action(do_compile, " [CHESTNUT] $TARGET") + +def compile_model(onnx_path, pkl_path, flags, chestnut=False): + onnx_path, target_pkl_path = File(onnx_path).abspath, File(pkl_path).abspath + onnx_deps = get_existing_chunks(onnx_path) + cmd = (f'{flags} {mac_brew_string} {taskset}{compiler}.compile_onnx ' + f'--onnx {onnx_path} --output {target_pkl_path}') + onnx_sizes_sum = sum(os.path.getsize(f) for f in onnx_deps) + chunk_targets = get_chunk_targets(target_pkl_path, estimate_pickle_max_size(onnx_sizes_sum)) def do_chunk(target, source, env, pkl=target_pkl_path, chunks=chunk_targets): chunk_file(pkl, chunks) - actions = Action(do_compile, " [CHESTNUT] $TARGET") if chestnut else [cmd, Action(do_chunk, " [CHUNK] $TARGET")] + actions = chestnut_action(cmd, target_pkl_path, chunk_targets) if chestnut else [cmd, Action(do_chunk, " [CHUNK] $TARGET")] node = lenv.Command( chunk_targets, - tinygrad_files + compile_modeld_script + driving_onnx_deps + [Value(cmd), Value(chunk_targets), chunker_file], + tinygrad_files + onnx_deps + [Value(cmd), Value(chunk_targets), chunker_file], actions, ) if chestnut: lenv.SideEffect(chestnut_lock, node) +compile_model('models/dmonitoring_model.onnx', 'models/dmonitoring_model_tinygrad.pkl', tg_flags) + +model_w, model_h = MEDMODEL_INPUT_SIZE +for chestnut in [False, True] if CHESTNUT else [False]: + file_prefix, cmd_flags = ('big_', chestnut_tg_flags) if chestnut else ('', tg_flags) + compile_model(f'models/{file_prefix}driving_supercombo.onnx', modeld_pkl_path(chestnut), cmd_flags, chestnut) for cam_w, cam_h in camera_configs: warp_pkl_path = File(f"models/{file_prefix}driving_warp_{cam_w}x{cam_h}_tinygrad.pkl").abspath stride, y_height, uv_height, _ = get_nv12_info(cam_w, cam_h) - cmd = (f'{cmd_flags} {mac_brew_string} {taskset}python3 {modeld_dir}/compile_warp.py ' + cmd = (f'{cmd_flags} {mac_brew_string} {taskset}{compiler}.compile_warp ' f'--camera-resolution {cam_w}x{cam_h} --warp-to {model_w}x{model_h} --layout yuv420 ' - f'--frames 2 --frame-size {stride * (y_height + uv_height)} ' + f'--frames 2 --stride {stride} --uv-offset {stride * y_height} --frame-size {stride * (y_height + uv_height)} ' f'--output {warp_pkl_path}') - def do_compile_warp(target, source, env, command=cmd): - return do_compile(target, source, env, command=command, chunks=()) - action = Action(do_compile_warp, " [CHESTNUT] $TARGET") if chestnut else cmd - node = lenv.Command(warp_pkl_path, tinygrad_files + compile_warp_script + [Value(cmd)], action) + action = chestnut_action(cmd) if chestnut else cmd + node = lenv.Command(warp_pkl_path, tinygrad_files + warp_deps + [Value(cmd)], action) if chestnut: lenv.SideEffect(chestnut_lock, node) -# get model metadata -fn = File(f"models/dmonitoring_model").abspath -script_files = [File(Dir("#openpilot/selfdrive/modeld").File("get_model_metadata.py").abspath)] -cmd = f'{tg_flags} {mac_brew_string} python3 {Dir("#openpilot/selfdrive/modeld").abspath}/get_model_metadata.py {fn}.onnx' -lenv.Command(fn + "_metadata.pkl", [fn + ".onnx"] + tinygrad_files + script_files + [tg_devices_node], cmd) - dm_w, dm_h = DM_INPUT_SIZE for cam_w, cam_h in camera_configs: dm_pkl_path = File(f"models/dm_warp_{cam_w}x{cam_h}_tinygrad.pkl").abspath - cmd = (f'{tg_flags} {mac_brew_string} python3 {modeld_dir}/compile_warp.py ' + stride, y_height, _, frame_size = get_nv12_info(cam_w, cam_h) + cmd = (f'{tg_flags} {mac_brew_string} {compiler}.compile_warp ' f'--camera-resolution {cam_w}x{cam_h} --warp-to {dm_w}x{dm_h} --layout luma --border-fill 16 --transform-device NPY ' - f'--output {dm_pkl_path}') - lenv.Command(dm_pkl_path, tinygrad_files + compile_warp_script + [tg_devices_node], cmd) - -def tg_compile(flags, model_name): - pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"' - fn = File(f"models/{model_name}").abspath - pkl = fn + "_tinygrad.pkl" - onnx_path = fn + ".onnx" - chunk_targets = get_chunk_targets(pkl, estimate_pickle_max_size(os.path.getsize(onnx_path))) - def do_chunk(target, source, env): - chunk_file(pkl, chunk_targets) - return lenv.Command( - chunk_targets, - [onnx_path] + tinygrad_files + [Value(chunk_targets), chunker_file, tg_devices_node], - [f'{pythonpath_string} {flags} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {pkl}', - Action(do_chunk, " [CHUNK] $TARGET")], - ) - -tg_compile(tg_flags, 'dmonitoring_model') + f'--stride {stride} --uv-offset {stride * y_height} --frame-size {frame_size} --output {dm_pkl_path}') + lenv.Command(dm_pkl_path, tinygrad_files + warp_deps + [Value(cmd)], cmd) diff --git a/openpilot/selfdrive/modeld/compile_modeld.py b/openpilot/selfdrive/modeld/compile_modeld.py deleted file mode 100755 index 463e697519..0000000000 --- a/openpilot/selfdrive/modeld/compile_modeld.py +++ /dev/null @@ -1,117 +0,0 @@ -#!/usr/bin/env python3 -import argparse -import atexit -import os -import tempfile -import time -import shutil - -import numpy as np - -from openpilot.selfdrive.modeld.helpers import dump_oob, load_oob, patch_tinygrad_fetch_fw -patch_tinygrad_fetch_fw() - -from tinygrad.tensor import Tensor -from tinygrad.device import Device -from tinygrad.engine.jit import TinyJit - - -def make_input_queues(input_shapes, device): - return {name: Tensor(np.zeros(shape, dtype=dtype.fmt), device=device).realize() for name, (shape, dtype) in input_shapes.items()} - - -def make_run_model(model_runner, state_pairs): - def run_model(**inputs): - outputs = {name: value.contiguous() for name, value in model_runner(inputs).items()} - Tensor.realize(*outputs.values()) - if state_pairs: - Tensor.realize(*(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_shapes, benchmark_runs): - if benchmark_runs < 1: - raise ValueError("benchmark_runs must be at least 1") - - SEED = 42 - def random_inputs_run(fn, seed, n_runs, test_val=None, test_buffers=None, expect_match=True): - input_queues = make_input_queues(input_shapes, Device.DEFAULT) - rng = np.random.default_rng(seed) - - for i in range(n_runs): - for value in input_queues.values(): - values = rng.standard_normal(value.shape) if np.issubdtype(np.dtype(value.dtype.fmt), np.floating) else rng.integers(0, 256, value.shape) - value.assign(Tensor(values.astype(value.dtype.fmt), device=Device.DEFAULT)).realize() - Device.default.synchronize() - st = time.perf_counter() - outs = fn(**input_queues) - mt = time.perf_counter() - Device.default.synchronize() - et = time.perf_counter() - print(f" [{i+1}/{n_runs}] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms") - - if i == 0: - val = [v.numpy() for v in outs] - buffers = [v.numpy() for v in input_queues.values()] - - if test_val is not None: - match = all(np.array_equal(a, b) for a, b in zip(val, test_val, strict=True)) - assert match == expect_match, f"outputs {'differ from' if expect_match else 'match'} baseline (seed={seed})" - if test_buffers is not None: - match = all(np.array_equal(a, b) for a, b in zip(buffers, test_buffers, strict=True)) - assert match == expect_match, f"buffers {'differ from' if expect_match else 'match'} baseline (seed={seed})" - return val, buffers - - print('capture + replay') - test_val, test_buffers = random_inputs_run(jit, SEED, 3) - print(f'pickle round trip ({benchmark_runs} runs per seed)') - with tempfile.TemporaryFile(dir=".") as f: - dump_oob(jit, f) - f.seek(0) - loaded_jit = load_oob(f) - random_inputs_run(loaded_jit, SEED, benchmark_runs, test_val, test_buffers, expect_match=True) - random_inputs_run(loaded_jit, SEED+1, benchmark_runs, test_val, test_buffers, expect_match=False) - return jit - - -def read_file_chunked_to_disk(path): - from openpilot.common.file_chunker import open_file_chunked - tmp_path = f'{path}.unchunked' - with open(tmp_path, 'wb') as f, open_file_chunked(path) as src: - shutil.copyfileobj(src, f) - atexit.register(lambda: os.path.exists(tmp_path) and os.remove(tmp_path)) - return tmp_path - - -if __name__ == "__main__": - from tinygrad.nn.onnx import OnnxRunner - from openpilot.selfdrive.modeld.get_model_metadata import make_metadata_dict - p = argparse.ArgumentParser() - p.add_argument('--onnx', required=True) - p.add_argument('--output', required=True) - p.add_argument('--benchmark-runs', type=int, default=1, - help='timed loaded-JIT runs for each correctness seed') - args = p.parse_args() - - model_path = read_file_chunked_to_disk(args.onnx) - - 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 = make_run_model(model_runner, state_pairs) - out['run_model'] = compile_jit(TinyJit(run_model, prune=True), input_shapes, 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/compile_warp.py b/openpilot/selfdrive/modeld/compile_warp.py deleted file mode 100644 index 1b2c085537..0000000000 --- a/openpilot/selfdrive/modeld/compile_warp.py +++ /dev/null @@ -1,150 +0,0 @@ -import argparse -import pickle -import time -from collections import namedtuple - -from openpilot.selfdrive.modeld.helpers import patch_tinygrad_fetch_fw -patch_tinygrad_fetch_fw() - -from tinygrad.tensor import Tensor -from tinygrad.helpers import Context -from tinygrad.device import Device -from tinygrad.engine.jit import TinyJit - -from openpilot.system.camerad.cameras.nv12_info import get_nv12_info - - -NV12Frame = namedtuple("NV12Frame", ['width', 'height', 'stride', 'y_height', 'uv_height', 'size']) - - -def warp_perspective_tinygrad(src_flat, M_inv, dst_shape, src_shape, stride_pad, border_fill_val=None): - w_dst, h_dst = dst_shape - h_src, w_src = src_shape - - x = Tensor.arange(w_dst).reshape(1, w_dst).expand(h_dst, w_dst).reshape(-1) - y = Tensor.arange(h_dst).reshape(h_dst, 1).expand(h_dst, w_dst).reshape(-1) - - # inline 3x3 matmul as elementwise to avoid reduce op (enables fusion with gather) - src_x = M_inv[0, 0] * x + M_inv[0, 1] * y + M_inv[0, 2] - src_y = M_inv[1, 0] * x + M_inv[1, 1] * y + M_inv[1, 2] - src_w = M_inv[2, 0] * x + M_inv[2, 1] * y + M_inv[2, 2] - - src_x = src_x / src_w - src_y = src_y / src_w - - x_round = Tensor.round(src_x) - y_round = Tensor.round(src_y) - x_nn_clipped = x_round.clip(0, w_src - 1).cast('int') - y_nn_clipped = y_round.clip(0, h_src - 1).cast('int') - idx = y_nn_clipped * (w_src + stride_pad) + x_nn_clipped - sampled = src_flat[idx] - - if border_fill_val is None: - return sampled - - in_bounds = ((x_round >= 0) & (x_round <= w_src - 1) & - (y_round >= 0) & (y_round <= h_src - 1)).cast(sampled.dtype) - return sampled * in_bounds + Tensor(border_fill_val, dtype=sampled.dtype) * (1 - in_bounds) - - -def frames_to_tensor(frames): - H = (frames.shape[0] * 2) // 3 - W = frames.shape[1] - return Tensor.cat(frames[0:H:2, 0::2], - frames[1:H:2, 0::2], - frames[0:H:2, 1::2], - frames[1:H:2, 1::2], - frames[H:H+H//4].reshape((H//2, W//2)), - frames[H+H//4:H+H//2].reshape((H//2, W//2)), dim=0).reshape((6, H//2, W//2)) - - -def make_frame_prepare(nv12: NV12Frame, model_w, model_h, layout="yuv420", border_fill=None): - cam_w, cam_h, stride, y_height, uv_height, _ = nv12 - uv_offset = stride * y_height - stride_pad = stride - cam_w - - def frame_prepare_tinygrad(input_frame, M_inv): - if layout == "luma": - return warp_perspective_tinygrad(input_frame[:cam_h*stride], M_inv, - (model_w, model_h), (cam_h, cam_w), stride_pad, - border_fill_val=border_fill).reshape(-1, model_h * model_w) - # UV_SCALE @ M_inv @ UV_SCALE_INV simplifies to elementwise scaling - M_inv_uv = M_inv * Tensor([[1.0, 1.0, 0.5], [1.0, 1.0, 0.5], [2.0, 2.0, 1.0]], device=Device.DEFAULT) - # deinterleave NV12 UV plane (UVUV... -> separate U, V) - uv = input_frame[uv_offset:uv_offset + uv_height * stride].reshape(uv_height, stride) - with Context(SPLIT_REDUCEOP=0): - y = warp_perspective_tinygrad(input_frame[:cam_h*stride], - M_inv, (model_w, model_h), - (cam_h, cam_w), stride_pad, border_fill_val=border_fill).realize() - u = warp_perspective_tinygrad(uv[:cam_h//2, :cam_w:2].flatten(), - M_inv_uv, (model_w//2, model_h//2), - (cam_h//2, cam_w//2), 0, border_fill_val=border_fill).realize() - v = warp_perspective_tinygrad(uv[:cam_h//2, 1:cam_w:2].flatten(), - M_inv_uv, (model_w//2, model_h//2), - (cam_h//2, cam_w//2), 0, border_fill_val=border_fill).realize() - yuv = y.cat(u).cat(v).reshape((model_h * 3 // 2, model_w)) - return frames_to_tensor(yuv) - return frame_prepare_tinygrad - - -def make_warp(nv12, model_w, model_h, layout="yuv420", border_fill=None, frames=1): - frame_prepare = make_frame_prepare(nv12, model_w, model_h, layout, border_fill) - - def warp(input_frames, transforms): - input_frames = input_frames.to(Device.DEFAULT) - transforms = transforms.to(Device.DEFAULT) - Tensor.realize(input_frames, transforms) - if frames == 1: - return frame_prepare(input_frames, transforms) - return Tensor.stack(*(frame_prepare(input_frames[i], transforms[i]) for i in range(frames))) - - return warp - - -def _parse_size(s): - w, h = s.lower().split('x') - return int(w), int(h) - - -def compile_warp(nv12: NV12Frame, model_w, model_h, pkl_path, layout, border_fill=None, - frames=1, transform_device=None, frame_size=None): - print(f"Compiling {layout} warp for {nv12.width}x{nv12.height} -> {model_w}x{model_h}...") - - warp_jit = TinyJit(make_warp(nv12, model_w, model_h, layout, border_fill, frames), prune=True) - frame_size = nv12.size if frame_size is None else frame_size - frame_shape = (frame_size,) if frames == 1 else (frames, frame_size) - transform_shape = (3, 3) if frames == 1 else (frames, 3, 3) - - for i in range(10): - frame = Tensor.randint(*frame_shape, low=0, high=256, dtype='uint8').realize() - M_inv = Tensor(Tensor.randn(*transform_shape).mul(8).realize().numpy(), device=transform_device) - Device.default.synchronize() - st = time.perf_counter() - warp_jit(frame, M_inv).realize() - mt = time.perf_counter() - Device.default.synchronize() - et = time.perf_counter() - print(f" [{i+1}/10] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms") - - with open(pkl_path, "wb") as f: - pickle.dump(warp_jit, f) - print(f" Saved to {pkl_path}") - - -if __name__ == "__main__": - p = argparse.ArgumentParser() - p.add_argument('--camera-resolution', type=_parse_size, required=True, help='camera resolution WxH') - p.add_argument('--warp-to', type=_parse_size, required=True, help='output WxH') - p.add_argument('--layout', choices=['luma', 'yuv420'], required=True) - p.add_argument('--border-fill', type=int, help='fill value outside the frame; omit to clamp coordinates') - p.add_argument('--frames', type=int, default=1, help='number of frames to warp together') - p.add_argument('--transform-device', help='device holding the input transforms; default: compute device') - p.add_argument('--frame-size', type=int, help='input frame size in bytes; default: full NV12 allocation') - p.add_argument('--output', required=True) - args = p.parse_args() - - cam_w, cam_h = args.camera_resolution - nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)) - model_w, model_h = args.warp_to - compile_warp(nv12, model_w, model_h, args.output, args.layout, args.border_fill, - args.frames, args.transform_device, args.frame_size) diff --git a/openpilot/selfdrive/modeld/dmonitoringmodeld.py b/openpilot/selfdrive/modeld/dmonitoringmodeld.py index 4010725b89..4d84aa0cd6 100755 --- a/openpilot/selfdrive/modeld/dmonitoringmodeld.py +++ b/openpilot/selfdrive/modeld/dmonitoringmodeld.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 import os -from openpilot.selfdrive.modeld.helpers import MODELS_DIR, get_tg_input_devices +from openpilot.selfdrive.modeld.helpers import MODELS_DIR, load_oob from tinygrad.tensor import Tensor import time import pickle @@ -18,10 +18,8 @@ from openpilot.system.camerad.cameras.nv12_info import get_nv12_info from openpilot.common.file_chunker import open_file_chunked from openpilot.selfdrive.modeld.parse_model_outputs import sigmoid, safe_exp -PROCESS_NAME = "openpilot.selfdrive.modeld.dmonitoringmodeld" SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') MODEL_PKL_PATH = MODELS_DIR / 'dmonitoring_model_tinygrad.pkl' -METADATA_PATH = MODELS_DIR / 'dmonitoring_model_metadata.pkl' class ModelState: @@ -29,11 +27,10 @@ class ModelState: output: np.ndarray def __init__(self, cam_w: int, cam_h: int): - self.DEV = get_tg_input_devices(PROCESS_NAME, chestnut=False)['DEV'] - with open(METADATA_PATH, 'rb') as f: - model_metadata = pickle.load(f) - self.input_shapes = model_metadata['input_shapes'] - self.output_slices = model_metadata['output_slices'] + jits = load_oob(open_file_chunked(MODEL_PKL_PATH)) + self.DEV = jits['input_devices']['model'] + self.input_shapes = jits['metadata']['input_shapes'] + self.output_slices = jits['metadata']['output_slices'] self.numpy_inputs = { 'calib': np.zeros(self.input_shapes['calib'], dtype=np.float32), @@ -42,9 +39,10 @@ class ModelState: self.warp_inputs_np = {'transform': np.zeros((3,3), dtype=np.float32)} self.warp_inputs = {k: Tensor(v, device='NPY') for k,v in self.warp_inputs_np.items()} self.frame_buf_params = get_nv12_info(cam_w, cam_h) - self.tensor_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items()} + self.tensor_inputs = {k: Tensor(v, device=self.DEV).realize() for k,v in self.numpy_inputs.items()} + self.calib_host = Tensor(self.numpy_inputs['calib'], device='NPY')._buffer() self._blob_cache : dict[int, Tensor] = {} - self.model_run = pickle.load(open_file_chunked(str(MODEL_PKL_PATH))) + self.model_run = jits['run_model'] with open(MODELS_DIR / f'dm_warp_{cam_w}x{cam_h}_tinygrad.pkl', "rb") as f: self.image_warp = pickle.load(f) @@ -52,6 +50,7 @@ class ModelState: self.numpy_inputs['calib'][0,:] = calib t1 = time.perf_counter() + self.tensor_inputs['calib']._buffer().copy_from(self.calib_host) ptr = np.frombuffer(buf.data, dtype=np.uint8).ctypes.data # There is a ringbuffer of imgs, just cache tensors pointing to all of them @@ -61,7 +60,8 @@ class ModelState: self.warp_inputs_np['transform'][:] = transform[:] self.tensor_inputs['input_img'] = self.image_warp(self._blob_cache[ptr], self.warp_inputs['transform']) - output = self.model_run(**self.tensor_inputs).numpy().flatten() + output, = self.model_run(**self.tensor_inputs) + output = output.numpy().astype(np.float32).reshape(-1) t2 = time.perf_counter() return output, t2 - t1 diff --git a/openpilot/selfdrive/modeld/get_model_metadata.py b/openpilot/selfdrive/modeld/get_model_metadata.py deleted file mode 100755 index e4c173957a..0000000000 --- a/openpilot/selfdrive/modeld/get_model_metadata.py +++ /dev/null @@ -1,55 +0,0 @@ -#!/usr/bin/env python3 -import sys -import pathlib -import codecs -import pickle -from typing import Any - -from tinygrad.nn.onnx import OnnxPBParser - - -class MetadataOnnxPBParser(OnnxPBParser): - def _parse_ModelProto(self) -> dict: - obj: dict[str, Any] = {"graph": {"input": [], "output": []}, "metadata_props": []} - for fid, wire_type in self._parse_message(self.reader.len): - match fid: - case 7: - obj["graph"] = self._parse_GraphProto() - case 14: - obj["metadata_props"].append(self._parse_StringStringEntryProto()) - case _: - self.reader.skip_field(wire_type) - return obj - - -def get_name_and_shape(value_info: dict[str, Any]) -> tuple[str, tuple[int, ...]]: - shape = tuple(int(dim) if isinstance(dim, int) else 0 for dim in value_info["parsed_type"].shape) - name = value_info["name"] - return name, shape - - -def get_metadata_value_by_name(model: dict[str, Any], name: str) -> str | Any: - for prop in model["metadata_props"]: - if prop["key"] == name: - return prop["value"] - return None - - -def make_metadata_dict(model_path): - model = MetadataOnnxPBParser(model_path).parse() - output_slices = get_metadata_value_by_name(model, 'output_slices') - assert output_slices is not None, 'output_slices not found in metadata' - return { - 'model_checkpoint': get_metadata_value_by_name(model, 'model_checkpoint'), - 'output_slices': pickle.loads(codecs.decode(output_slices.encode(), "base64")), - 'input_shapes': dict(get_name_and_shape(x) for x in model["graph"]["input"]), - 'output_shapes': dict(get_name_and_shape(x) for x in model["graph"]["output"]), - } - - -if __name__ == "__main__": - model_path = pathlib.Path(sys.argv[1]) - metadata_path = model_path.parent / (model_path.stem + '_metadata.pkl') - with open(metadata_path, 'wb') as f: - pickle.dump(make_metadata_dict(model_path), f) - print(f'saved metadata to {metadata_path}') diff --git a/openpilot/selfdrive/modeld/helpers.py b/openpilot/selfdrive/modeld/helpers.py index a5b64d1b25..307cb7c163 100644 --- a/openpilot/selfdrive/modeld/helpers.py +++ b/openpilot/selfdrive/modeld/helpers.py @@ -1,56 +1,18 @@ import io -import json import pickle -import shutil import struct -import tempfile from pathlib import Path from openpilot.common.file_chunker import get_manifest_path from openpilot.common.hardware.usb import CHESTNUT_USB_PRODUCT, USB_DEVICES_PATH, is_chestnut_usb_id MODELS_DIR = Path(__file__).resolve().parent / 'models' -TG_INPUT_DEVICES_PATH = MODELS_DIR / 'tg_input_devices.json' -def patch_tinygrad_fetch_fw(): - import hashlib - import zstandard - from tinygrad import helpers - original_fetch_fw = helpers.fetch_fw - def fetch_fw(path, name, sha256): - p = Path(f"/lib/firmware/{path}/{name}.zst") - if p.is_file(): - blob = zstandard.ZstdDecompressor().stream_reader(p.read_bytes()).read() - if hashlib.sha256(blob).hexdigest() == sha256: - return blob - return original_fetch_fw(path, name, sha256) - helpers.fetch_fw = fetch_fw - - -def get_tg_input_devices(process_name: str, chestnut: bool): - with open(TG_INPUT_DEVICES_PATH) as f: - return json.load(f)[process_name]['default' if not chestnut else 'chestnut'] - def modeld_pkl_path(chestnut: bool): prefix = 'big_' if chestnut else '' return MODELS_DIR / f'{prefix}driving_tinygrad.pkl' -def dump_oob(obj, f): - with tempfile.TemporaryFile(dir=".") as tmp: - def buffer_callback(pb: pickle.PickleBuffer): - m = pb.raw() - tmp.write(struct.pack(' None: # Pack host inputs into one upload to reduce USB transfer overhead for the eGPU. - self.input_queues = make_input_queues({name: self.input_shapes[name] for name in self.state_pairs}, self.model_device) + self.input_queues = {name: Tensor(np.zeros(shape, dtype=dtype.fmt), device=self.model_device).realize() + for name, (shape, dtype) in self.input_shapes.items() if name in self.state_pairs} shapes = {'tfm': (2, 3, 3)} | {name: shape for name, (shape, _) in self.input_shapes.items() if name not in self.state_pairs and name != 'new_img'} npy_size = sum(round_up(math.prod(shape) * 4, 128) for shape in shapes.values()) diff --git a/tinygrad_repo b/tinygrad_repo index f6fc4e3f2c..953a7f36cf 160000 --- a/tinygrad_repo +++ b/tinygrad_repo @@ -1 +1 @@ -Subproject commit f6fc4e3f2c3db5fae1e19cbfbc3ad9fc579a12ae +Subproject commit 953a7f36cfda50db17ee94505143474061c12e9d