From 8745cd0f385371488f7b0a84f4e072a829c4bbda Mon Sep 17 00:00:00 2001 From: Jason Wen Date: Sun, 26 Apr 2026 02:24:19 -0400 Subject: [PATCH] Revert "modeld: single jit (#37758)" This reverts commit d81d66193fc9d46291f03fb16557bb32166e4091. --- selfdrive/modeld/SConscript | 34 ++-- selfdrive/modeld/compile_modeld.py | 289 --------------------------- selfdrive/modeld/compile_warp.py | 198 ++++++++++++++++++ selfdrive/modeld/modeld.py | 150 +++++++++++--- selfdrive/modeld/tinygrad_helpers.py | 2 +- 5 files changed, 330 insertions(+), 343 deletions(-) delete mode 100755 selfdrive/modeld/compile_modeld.py create mode 100755 selfdrive/modeld/compile_warp.py diff --git a/selfdrive/modeld/SConscript b/selfdrive/modeld/SConscript index 053eb157b6..86b62d2c17 100644 --- a/selfdrive/modeld/SConscript +++ b/selfdrive/modeld/SConscript @@ -16,6 +16,7 @@ tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + " def estimate_pickle_max_size(onnx_size): return 1.2 * onnx_size + 10 * 1024 * 1024 # 20% + 10MB is plenty +# THREADS=0 is need to prevent bug: https://github.com/tinygrad/tinygrad/issues/14689 # get fastest TG config available = set(Device.get_available_devices()) if 'CUDA' in available: @@ -23,10 +24,9 @@ if 'CUDA' in available: tg_flags = f'DEV={tg_backend}' elif 'QCOM' in available: tg_backend = 'QCOM' - tg_flags = f'DEV={tg_backend} FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1' + tg_flags = f'DEV={tg_backend} FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0' else: tg_backend = 'CPU' if arch == 'Darwin' else 'CPU:LLVM' - # THREADS=0 is need to prevent bug: https://github.com/tinygrad/tinygrad/issues/14689 tg_flags = f'DEV={tg_backend} THREADS=0' def write_tg_compiled_flags(target, source, env): @@ -53,23 +53,14 @@ for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']: image_flag = { 'larch64': 'IMAGE=2', }.get(arch, 'IMAGE=0') -script_files = [File(Dir("#selfdrive/modeld").File("compile_modeld.py").abspath)] -compile_modeld_cmd = f'{tg_flags} {mac_brew_string} {image_flag} python3 {Dir("#selfdrive/modeld").abspath}/compile_modeld.py ' -driving_onnx_deps = [File(f"models/{m}.onnx").abspath for m in ['driving_vision', 'driving_policy']] -driving_metadata_deps = [File(f"models/{m}_metadata.pkl").abspath for m in ['driving_vision', 'driving_policy']] - -from openpilot.selfdrive.modeld.compile_modeld import MODELD_CONFIGS, DM_WARP_CONFIGS -policy_pkls = [File(cfg.pkl_path).abspath for cfg in MODELD_CONFIGS] -modeld_targets = policy_pkls + [File(cfg.pkl_path).abspath for cfg in DM_WARP_CONFIGS] -compile_node = lenv.Command(modeld_targets, tinygrad_files + script_files + driving_onnx_deps + driving_metadata_deps + [chunker_file, compiled_flags_node], compile_modeld_cmd) - -# chunk the combined policy pkls -for policy_pkl in policy_pkls: - onnx_sizes_sum = sum(os.path.getsize(f) for f in driving_onnx_deps) - chunk_targets = get_chunk_paths(policy_pkl, estimate_pickle_max_size(onnx_sizes_sum)) - def do_chunk(target, source, env, pkl=policy_pkl, chunks=chunk_targets): - chunk_file(pkl, chunks) - lenv.Command(chunk_targets, compile_node, do_chunk) +script_files = [File(Dir("#selfdrive/modeld").File("compile_warp.py").abspath)] +compile_warp_cmd = f'{tg_flags} {mac_brew_string} python3 {Dir("#selfdrive/modeld").abspath}/compile_warp.py ' +from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye +warp_targets = [] +for cam in [_ar_ox_fisheye, _os_fisheye]: + w, h = cam.width, cam.height + warp_targets += [File(f"models/warp_{w}x{h}_tinygrad.pkl").abspath, File(f"models/dm_warp_{w}x{h}_tinygrad.pkl").abspath] +lenv.Command(warp_targets, tinygrad_files + script_files + [compiled_flags_node], compile_warp_cmd) def tg_compile(flags, model_name): pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"' @@ -90,4 +81,7 @@ def tg_compile(flags, model_name): do_chunk, ) -tg_compile(tg_flags, 'dmonitoring_model') +# Compile small models +for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']: + tg_compile(tg_flags, model_name) + diff --git a/selfdrive/modeld/compile_modeld.py b/selfdrive/modeld/compile_modeld.py deleted file mode 100755 index 469110a857..0000000000 --- a/selfdrive/modeld/compile_modeld.py +++ /dev/null @@ -1,289 +0,0 @@ -#!/usr/bin/env python3 -import time -import pickle -from dataclasses import dataclass -from itertools import product -from functools import partial - -import numpy as np -from tinygrad.tensor import Tensor -from tinygrad.helpers import Context -from tinygrad.device import Device -from tinygrad.engine.jit import TinyJit - -from openpilot.selfdrive.modeld.tinygrad_helpers import MODELS_DIR -from openpilot.system.camerad.cameras.nv12_info import get_nv12_info -from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE, DM_INPUT_SIZE -from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye - -# https://github.com/tinygrad/tinygrad/issues/15682 -from tinygrad.uop.ops import UOp, Ops -_orig = UOp.__reduce__ -UOp.__reduce__ = lambda self: (UOp.unique, ()) if self.op is Ops.UNIQUE else _orig(self) - - -@dataclass -class CompileConfig: - cam_w: int - cam_h: int - prepare_only: bool - prefix: str - - @property - def pkl_path(self): - return str(MODELS_DIR / f'{self.prefix}{"warp_" if self.prepare_only else ""}{self.cam_w}x{self.cam_h}_tinygrad.pkl') - -CAMERA_CONFIGS = [ - (_ar_ox_fisheye.width, _ar_ox_fisheye.height), # tici: 1928x1208 - (_os_fisheye.width, _os_fisheye.height), # mici: 1344x760 -] -MODELD_CONFIGS = [CompileConfig(cam_w, cam_h, prepare_only, 'driving_') for (cam_w, cam_h), prepare_only in product(CAMERA_CONFIGS, [True, False])] -DM_WARP_CONFIGS = [CompileConfig(cam_w, cam_h, True, 'dm_') for cam_w, cam_h in CAMERA_CONFIGS] - - -UV_SCALE_MATRIX = np.array([[0.5, 0, 0], [0, 0.5, 0], [0, 0, 1]], dtype=np.float32) -UV_SCALE_MATRIX_INV = np.linalg.inv(UV_SCALE_MATRIX) - - -def warp_perspective_tinygrad(src_flat, M_inv, dst_shape, src_shape, stride_pad): - 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_nn_clipped = Tensor.round(src_x).clip(0, w_src - 1).cast('int') - y_nn_clipped = Tensor.round(src_y).clip(0, h_src - 1).cast('int') - idx = y_nn_clipped * (w_src + stride_pad) + x_nn_clipped - - return src_flat[idx] - - -def frames_to_tensor(frames): - H = (frames.shape[0] * 2) // 3 - W = frames.shape[1] - in_img1 = 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)) - return in_img1 - - -def make_frame_prepare(cam_w, cam_h, model_w, model_h): - stride, y_height, uv_height, _ = get_nv12_info(cam_w, cam_h) - uv_offset = stride * y_height - stride_pad = stride - cam_w - - def frame_prepare_tinygrad(input_frame, M_inv): - # 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]]) - # 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).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).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).realize() - yuv = y.cat(u).cat(v).reshape((model_h * 3 // 2, model_w)) - tensor = frames_to_tensor(yuv) - return tensor - return frame_prepare_tinygrad - - -def make_input_queues(vision_input_shapes, policy_input_shapes, frame_skip): - img = vision_input_shapes['img'] # (1, 12, 128, 256) - n_frames = img[1] // 6 - img_buf_shape = (frame_skip * (n_frames - 1) + 1, 6, img[2], img[3]) - - fb = policy_input_shapes['features_buffer'] # (1, 25, 512) - dp = policy_input_shapes['desire_pulse'] # (1, 25, 8) - tc = policy_input_shapes['traffic_convention'] # (1, 2) - - npy = { - 'desire': np.zeros(dp[2], dtype=np.float32), - 'traffic_convention': np.zeros(tc, dtype=np.float32), - 'tfm': np.zeros((3, 3), dtype=np.float32), - 'big_tfm': np.zeros((3, 3), dtype=np.float32), - } - input_queues = { - 'img_q': Tensor.zeros(img_buf_shape, dtype='uint8').contiguous().realize(), - 'big_img_q': Tensor.zeros(img_buf_shape, dtype='uint8').contiguous().realize(), - 'feat_q': Tensor.zeros(frame_skip * (fb[1] - 1) + 1, fb[0], fb[2]).contiguous().realize(), - 'desire_q': Tensor.zeros(frame_skip * dp[1], dp[0], dp[2]).contiguous().realize(), - **{k: Tensor(v, device='NPY').realize() for k, v in npy.items()}, - } - return input_queues, npy - - -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_dm(cam_w, cam_h, dm_w, dm_h): - stride, y_height, _, _ = get_nv12_info(cam_w, cam_h) - stride_pad = stride - cam_w - - def warp_dm(input_frame, M_inv): - M_inv = M_inv.to(Device.DEFAULT) - result = warp_perspective_tinygrad(input_frame[:cam_h*stride], M_inv, (dm_w, dm_h), (cam_h, cam_w), stride_pad).reshape(-1, dm_h * dm_w) - return result - return warp_dm - - -def make_run_policy(vision_runner, policy_runner, cam_w, cam_h, - vision_features_slice, frame_skip, prepare_only=False): - model_w, model_h = MEDMODEL_INPUT_SIZE - frame_prepare = make_frame_prepare(cam_w, cam_h, model_w, model_h) - sample_skip_fn = partial(sample_skip, frame_skip=frame_skip) - sample_desire_fn = partial(sample_desire, frame_skip=frame_skip) - - def run_policy(img_q, big_img_q, feat_q, desire_q, desire, traffic_convention, tfm, big_tfm, frame, big_frame): - img = shift_and_sample(img_q, frame_prepare(frame, tfm.to(Device.DEFAULT)).unsqueeze(0), sample_skip_fn) - big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm.to(Device.DEFAULT)).unsqueeze(0), sample_skip_fn) - - if prepare_only: - return img, big_img - - vision_out = next(iter(vision_runner({'img': img, 'big_img': big_img}).values())).cast('float32') - - new_feat = vision_out[:, vision_features_slice].reshape(1, -1).unsqueeze(0) - feat_buf = shift_and_sample(feat_q, new_feat, sample_skip_fn) - desire_buf = shift_and_sample(desire_q, desire.to(Device.DEFAULT).reshape(1, 1, -1), sample_desire_fn) - - inputs = {'features_buffer': feat_buf, 'desire_pulse': desire_buf, 'traffic_convention': traffic_convention.to(Device.DEFAULT)} - policy_out = next(iter(policy_runner(inputs).values())).cast('float32') - - return vision_out, policy_out - return run_policy - - -def compile_modeld(cam_w, cam_h, prepare_only, pkl_path): - from tinygrad.nn.onnx import OnnxRunner - from openpilot.selfdrive.modeld.constants import ModelConstants - - _, _, _, yuv_size = get_nv12_info(cam_w, cam_h) - print(f"Compiling combined policy JIT for {cam_w}x{cam_h}...") - - vision_runner = OnnxRunner(MODELS_DIR / 'driving_vision.onnx') - policy_runner = OnnxRunner(MODELS_DIR / 'driving_policy.onnx') - - with open(MODELS_DIR / 'driving_vision_metadata.pkl', 'rb') as f: - vision_metadata = pickle.load(f) - vision_features_slice = vision_metadata['output_slices']['hidden_state'] - vision_input_shapes = vision_metadata['input_shapes'] - with open(MODELS_DIR / 'driving_policy_metadata.pkl', 'rb') as f: - policy_input_shapes = pickle.load(f)['input_shapes'] - - frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ - - _run = make_run_policy(vision_runner, policy_runner, - cam_w, cam_h, vision_features_slice, frame_skip, prepare_only) - run_policy_jit = TinyJit(_run, prune=True) - - N_RUNS = 3 - SEED = 42 - - def random_inputs_run_fn(fn, seed, test_val=None, test_buffers=None, expect_match=True): - input_queues, npy = make_input_queues(vision_input_shapes, policy_input_shapes, frame_skip) - np.random.seed(seed) - - for i in range(N_RUNS): - frame = Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize() - big_frame = Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize() - for v in npy.values(): - v[:] = np.random.randn(*v.shape).astype(v.dtype) - Device.default.synchronize() - st = time.perf_counter() - outs = fn(**input_queues, frame=frame, big_frame=big_frame) - mt = time.perf_counter() - for o in outs: - # .realize() not needed once jitted, but needed for unjitted fn - o.realize() - 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") - - val = [np.copy(v.numpy()) for v in outs] - buffers = [np.copy(v.numpy().copy()) 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 fn, val, buffers - - print('run unjitted') - _, test_val, test_buffers = random_inputs_run_fn(_run, seed=SEED) - print('capture + replay') - run_policy_jit, _, _ = random_inputs_run_fn(run_policy_jit, SEED, test_val, test_buffers) - - print('pickle round trip') - with open(pkl_path, "wb") as f: - pickle.dump(run_policy_jit, f) - print(f" Saved to {pkl_path}") - with open(pkl_path, "rb") as f: - run_policy_jit = pickle.load(f) - random_inputs_run_fn(run_policy_jit, SEED, test_val, test_buffers, expect_match=True) - random_inputs_run_fn(run_policy_jit, SEED+1, test_val, test_buffers, expect_match=False) - - -def compile_dm_warp(cam_w, cam_h, pkl_path): - dm_w, dm_h = DM_INPUT_SIZE - _, _, _, yuv_size = get_nv12_info(cam_w, cam_h) - - print(f"Compiling DM warp for {cam_w}x{cam_h}...") - - warp_dm = make_warp_dm(cam_w, cam_h, dm_w, dm_h) - warp_dm_jit = TinyJit(warp_dm, prune=True) - - for i in range(10): - inputs = [Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize(), - Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')] - Device.default.synchronize() - st = time.perf_counter() - warp_dm_jit(*inputs) - 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_dm_jit, f) - print(f" Saved to {pkl_path}") - - -def run_and_save_pickle(): - for cfg in MODELD_CONFIGS: - compile_modeld(cfg.cam_w, cfg.cam_h, cfg.prepare_only, cfg.pkl_path) - for cfg in DM_WARP_CONFIGS: - compile_dm_warp(cfg.cam_w, cfg.cam_h, cfg.pkl_path) - - -if __name__ == "__main__": - run_and_save_pickle() diff --git a/selfdrive/modeld/compile_warp.py b/selfdrive/modeld/compile_warp.py new file mode 100755 index 0000000000..47511f2a2b --- /dev/null +++ b/selfdrive/modeld/compile_warp.py @@ -0,0 +1,198 @@ +#!/usr/bin/env python3 +import time +import pickle +import numpy as np +from pathlib import Path +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 +from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE, DM_INPUT_SIZE +from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye + +MODELS_DIR = Path(__file__).parent / 'models' + +CAMERA_CONFIGS = [ + (_ar_ox_fisheye.width, _ar_ox_fisheye.height), # tici: 1928x1208 + (_os_fisheye.width, _os_fisheye.height), # mici: 1344x760 +] + +UV_SCALE_MATRIX = np.array([[0.5, 0, 0], [0, 0.5, 0], [0, 0, 1]], dtype=np.float32) +UV_SCALE_MATRIX_INV = np.linalg.inv(UV_SCALE_MATRIX) + +IMG_BUFFER_SHAPE = (30, MEDMODEL_INPUT_SIZE[1] // 2, MEDMODEL_INPUT_SIZE[0] // 2) + + +def warp_pkl_path(w, h): + return MODELS_DIR / f'warp_{w}x{h}_tinygrad.pkl' + + +def dm_warp_pkl_path(w, h): + return MODELS_DIR / f'dm_warp_{w}x{h}_tinygrad.pkl' + + +def warp_perspective_tinygrad(src_flat, M_inv, dst_shape, src_shape, stride_pad): + 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_nn_clipped = Tensor.round(src_x).clip(0, w_src - 1).cast('int') + y_nn_clipped = Tensor.round(src_y).clip(0, h_src - 1).cast('int') + idx = y_nn_clipped * (w_src + stride_pad) + x_nn_clipped + + return src_flat[idx] + + +def frames_to_tensor(frames, model_w, model_h): + H = (frames.shape[0] * 2) // 3 + W = frames.shape[1] + in_img1 = 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)) + return in_img1 + + +def make_frame_prepare(cam_w, cam_h, model_w, model_h): + stride, y_height, uv_height, _ = get_nv12_info(cam_w, cam_h) + uv_offset = stride * y_height + stride_pad = stride - cam_w + + def frame_prepare_tinygrad(input_frame, M_inv): + # 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]]) + # 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).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).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).realize() + yuv = y.cat(u).cat(v).reshape((model_h * 3 // 2, model_w)) + tensor = frames_to_tensor(yuv, model_w, model_h) + return tensor + return frame_prepare_tinygrad + + +def make_update_img_input(frame_prepare, model_w, model_h): + def update_img_input_tinygrad(frame_buffer, frame, M_inv): + M_inv = M_inv.to(Device.DEFAULT) + new_img = frame_prepare(frame, M_inv) + frame_buffer.assign(frame_buffer[6:].cat(new_img, dim=0).contiguous()) + return Tensor.cat(frame_buffer[:6], frame_buffer[-6:], dim=0).contiguous().reshape(1, 12, model_h//2, model_w//2) + return update_img_input_tinygrad + + +def make_update_both_imgs(frame_prepare, model_w, model_h): + update_img = make_update_img_input(frame_prepare, model_w, model_h) + + def update_both_imgs_tinygrad(calib_img_buffer, new_img, M_inv, + calib_big_img_buffer, new_big_img, M_inv_big): + calib_img_pair = update_img(calib_img_buffer, new_img, M_inv) + calib_big_img_pair = update_img(calib_big_img_buffer, new_big_img, M_inv_big) + return calib_img_pair, calib_big_img_pair + return update_both_imgs_tinygrad + + +def make_warp_dm(cam_w, cam_h, dm_w, dm_h): + stride, y_height, _, _ = get_nv12_info(cam_w, cam_h) + stride_pad = stride - cam_w + + def warp_dm(input_frame, M_inv): + M_inv = M_inv.to(Device.DEFAULT) + result = warp_perspective_tinygrad(input_frame[:cam_h*stride], M_inv, (dm_w, dm_h), (cam_h, cam_w), stride_pad).reshape(-1, dm_h * dm_w) + return result + return warp_dm + + +def compile_modeld_warp(cam_w, cam_h): + model_w, model_h = MEDMODEL_INPUT_SIZE + _, _, _, yuv_size = get_nv12_info(cam_w, cam_h) + + print(f"Compiling modeld warp for {cam_w}x{cam_h}...") + + frame_prepare = make_frame_prepare(cam_w, cam_h, model_w, model_h) + update_both_imgs = make_update_both_imgs(frame_prepare, model_w, model_h) + update_img_jit = TinyJit(update_both_imgs, prune=True) + + full_buffer = Tensor.zeros(IMG_BUFFER_SHAPE, dtype='uint8').contiguous().realize() + big_full_buffer = Tensor.zeros(IMG_BUFFER_SHAPE, dtype='uint8').contiguous().realize() + for i in range(10): + img_inputs = [full_buffer, + Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize(), + Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')] + big_img_inputs = [big_full_buffer, + Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize(), + Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')] + inputs = img_inputs + big_img_inputs + Device.default.synchronize() + + st = time.perf_counter() + _ = update_img_jit(*inputs) + 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") + + pkl_path = warp_pkl_path(cam_w, cam_h) + with open(pkl_path, "wb") as f: + pickle.dump(update_img_jit, f) + print(f" Saved to {pkl_path}") + + jit = pickle.load(open(pkl_path, "rb")) + jit(*inputs) + + +def compile_dm_warp(cam_w, cam_h): + dm_w, dm_h = DM_INPUT_SIZE + _, _, _, yuv_size = get_nv12_info(cam_w, cam_h) + + print(f"Compiling DM warp for {cam_w}x{cam_h}...") + + warp_dm = make_warp_dm(cam_w, cam_h, dm_w, dm_h) + warp_dm_jit = TinyJit(warp_dm, prune=True) + + for i in range(10): + inputs = [Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize(), + Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')] + Device.default.synchronize() + st = time.perf_counter() + warp_dm_jit(*inputs) + 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") + + pkl_path = dm_warp_pkl_path(cam_w, cam_h) + with open(pkl_path, "wb") as f: + pickle.dump(warp_dm_jit, f) + print(f" Saved to {pkl_path}") + + +def run_and_save_pickle(): + for cam_w, cam_h in CAMERA_CONFIGS: + compile_modeld_warp(cam_w, cam_h) + compile_dm_warp(cam_w, cam_h) + + +if __name__ == "__main__": + run_and_save_pickle() diff --git a/selfdrive/modeld/modeld.py b/selfdrive/modeld/modeld.py index 7a8a442981..0c50d8beb5 100755 --- a/selfdrive/modeld/modeld.py +++ b/selfdrive/modeld/modeld.py @@ -26,7 +26,6 @@ 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, smooth_value, get_curvature_from_plan from openpilot.selfdrive.modeld.parse_model_outputs import Parser -from openpilot.selfdrive.modeld.compile_modeld import CompileConfig, make_input_queues from openpilot.selfdrive.modeld.fill_model_msg import fill_model_msg, fill_pose_msg, PublishState from openpilot.common.file_chunker import read_file_chunked from openpilot.selfdrive.modeld.constants import ModelConstants, Plan @@ -35,13 +34,17 @@ from openpilot.selfdrive.modeld.constants import ModelConstants, Plan PROCESS_NAME = "selfdrive.modeld.modeld" SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') +VISION_PKL_PATH = MODELS_DIR / 'driving_vision_tinygrad.pkl' VISION_METADATA_PATH = MODELS_DIR / 'driving_vision_metadata.pkl' +POLICY_PKL_PATH = MODELS_DIR / 'driving_policy_tinygrad.pkl' POLICY_METADATA_PATH = MODELS_DIR / 'driving_policy_metadata.pkl' LAT_SMOOTH_SECONDS = 0.0 LONG_SMOOTH_SECONDS = 0.3 MIN_LAT_CONTROL_SPEED = 0.3 +IMG_QUEUE_SHAPE = (6*(ModelConstants.MODEL_RUN_FREQ//ModelConstants.MODEL_CONTEXT_FREQ + 1), 128, 256) +assert IMG_QUEUE_SHAPE[0] == 30 def get_action_from_model(model_output: dict[str, np.ndarray], prev_action: log.ModelDataV2.Action, @@ -76,36 +79,106 @@ class FrameMeta: if vipc is not None: self.frame_id, self.timestamp_sof, self.timestamp_eof = vipc.frame_id, vipc.timestamp_sof, vipc.timestamp_eof +class InputQueues: + def __init__ (self, model_fps, env_fps, n_frames_input): + assert env_fps % model_fps == 0 + assert env_fps >= model_fps + self.model_fps = model_fps + self.env_fps = env_fps + self.n_frames_input = n_frames_input + + self.dtypes = {} + self.shapes = {} + self.q = {} + + def update_dtypes_and_shapes(self, input_dtypes, input_shapes) -> None: + self.dtypes.update(input_dtypes) + if self.env_fps == self.model_fps: + self.shapes.update(input_shapes) + else: + for k in input_shapes: + shape = list(input_shapes[k]) + if 'img' in k: + n_channels = shape[1] // self.n_frames_input + shape[1] = (self.env_fps // self.model_fps + (self.n_frames_input - 1)) * n_channels + else: + shape[1] = (self.env_fps // self.model_fps) * shape[1] + self.shapes[k] = tuple(shape) + + def reset(self) -> None: + self.q = {k: np.zeros(self.shapes[k], dtype=self.dtypes[k]) for k in self.dtypes.keys()} + + def enqueue(self, inputs:dict[str, np.ndarray]) -> None: + for k in inputs.keys(): + if inputs[k].dtype != self.dtypes[k]: + raise ValueError(f'supplied input <{k}({inputs[k].dtype})> has wrong dtype, expected {self.dtypes[k]}') + input_shape = list(self.shapes[k]) + input_shape[1] = -1 + single_input = inputs[k].reshape(tuple(input_shape)) + sz = single_input.shape[1] + self.q[k][:,:-sz] = self.q[k][:,sz:] + self.q[k][:,-sz:] = single_input + + def get(self, *names) -> dict[str, np.ndarray]: + if self.env_fps == self.model_fps: + return {k: self.q[k] for k in names} + else: + out = {} + for k in names: + shape = self.shapes[k] + if 'img' in k: + n_channels = shape[1] // (self.env_fps // self.model_fps + (self.n_frames_input - 1)) + out[k] = np.concatenate([self.q[k][:, s:s+n_channels] for s in np.linspace(0, shape[1] - n_channels, self.n_frames_input, dtype=int)], axis=1) + elif 'pulse' in k: + # any pulse within interval counts + out[k] = self.q[k].reshape((shape[0], shape[1] * self.model_fps // self.env_fps, self.env_fps // self.model_fps, -1)).max(axis=2) + else: + idxs = np.arange(-1, -shape[1], -self.env_fps // self.model_fps)[::-1] + out[k] = self.q[k][:, idxs] + return out class ModelState: + inputs: dict[str, np.ndarray] + output: np.ndarray prev_desire: np.ndarray # for tracking the rising edge of the pulse - def __init__(self, cam_w: int, cam_h: int): + def __init__(self): with open(VISION_METADATA_PATH, 'rb') as f: vision_metadata = pickle.load(f) self.vision_input_shapes = vision_metadata['input_shapes'] self.vision_input_names = list(self.vision_input_shapes.keys()) self.vision_output_slices = vision_metadata['output_slices'] + vision_output_size = vision_metadata['output_shapes']['outputs'][1] with open(POLICY_METADATA_PATH, 'rb') as f: policy_metadata = pickle.load(f) self.policy_input_shapes = policy_metadata['input_shapes'] self.policy_output_slices = policy_metadata['output_slices'] + policy_output_size = policy_metadata['output_shapes']['outputs'][1] self.prev_desire = np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32) - self.frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ - self.input_queues, self.npy = make_input_queues(self.vision_input_shapes, self.policy_input_shapes, self.frame_skip) + # policy inputs + self.numpy_inputs = {k: np.zeros(self.policy_input_shapes[k], dtype=np.float32) for k in self.policy_input_shapes} + self.full_input_queues = InputQueues(ModelConstants.MODEL_CONTEXT_FREQ, ModelConstants.MODEL_RUN_FREQ, ModelConstants.N_FRAMES) + for k in ['desire_pulse', 'features_buffer']: + self.full_input_queues.update_dtypes_and_shapes({k: self.numpy_inputs[k].dtype}, {k: self.numpy_inputs[k].shape}) + self.full_input_queues.reset() + + self.img_queues = {'img': Tensor.zeros(IMG_QUEUE_SHAPE, dtype='uint8').contiguous().realize(), + 'big_img': Tensor.zeros(IMG_QUEUE_SHAPE, dtype='uint8').contiguous().realize()} self.full_frames : dict[str, Tensor] = {} self._blob_cache : dict[int, Tensor] = {} + self.transforms_np = {k: np.zeros((3,3), dtype=np.float32) for k in self.img_queues} + self.transforms = {k: Tensor(v, device='NPY').realize() for k, v in self.transforms_np.items()} + self.vision_output = np.zeros(vision_output_size, dtype=np.float32) + self.policy_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items()} + self.policy_output = np.zeros(policy_output_size, dtype=np.float32) self.parser = Parser() - self.frame_buf_params = {k: get_nv12_info(cam_w, cam_h) for k in ('img', 'big_img')} - self.run_policy = pickle.loads(read_file_chunked(CompileConfig(cam_w, cam_h, prefix='driving_', prepare_only=False).pkl_path)) - self.warp_enqueue = pickle.loads(read_file_chunked(CompileConfig(cam_w, cam_h, prefix='driving_', prepare_only=True).pkl_path)) - self.warp_enqueue( - **self.input_queues, - frame=Tensor.zeros(self.frame_buf_params['img'][3], dtype='uint8').contiguous().realize(), - big_frame=Tensor.zeros(self.frame_buf_params['big_img'][3], dtype='uint8').contiguous().realize()) + self.frame_buf_params : dict[str, tuple[int, int, int, int]] = {} + self.update_imgs = None + self.vision_run = pickle.loads(read_file_chunked(str(VISION_PKL_PATH))) + self.policy_run = pickle.loads(read_file_chunked(str(POLICY_PKL_PATH))) def slice_outputs(self, model_outputs: np.ndarray, output_slices: dict[str, slice]) -> dict[str, np.ndarray]: parsed_model_outputs = {k: model_outputs[np.newaxis, v] for k,v in output_slices.items()} @@ -113,6 +186,18 @@ class ModelState: def run(self, bufs: dict[str, VisionBuf], transforms: dict[str, np.ndarray], inputs: dict[str, np.ndarray], prepare_only: bool) -> dict[str, np.ndarray] | None: + # Model decides when action is completed, so desire input is just a pulse triggered on rising edge + inputs['desire_pulse'][0] = 0 + new_desire = np.where(inputs['desire_pulse'] - self.prev_desire > .99, inputs['desire_pulse'], 0) + self.prev_desire[:] = inputs['desire_pulse'] + if self.update_imgs is None: + for key in bufs.keys(): + w, h = bufs[key].width, bufs[key].height + self.frame_buf_params[key] = get_nv12_info(w, h) + warp_path = MODELS_DIR / f'warp_{w}x{h}_tinygrad.pkl' + with open(warp_path, "rb") as f: + self.update_imgs = pickle.load(f) + for key in bufs.keys(): ptr = bufs[key].data.ctypes.data yuv_size = self.frame_buf_params[key][3] @@ -121,31 +206,30 @@ class ModelState: if cache_key not in self._blob_cache: self._blob_cache[cache_key] = Tensor.from_blob(ptr, (yuv_size,), dtype='uint8') self.full_frames[key] = self._blob_cache[cache_key] + for key in bufs.keys(): + self.transforms_np[key][:,:] = transforms[key][:,:] - # Model decides when action is completed, so desire input is just a pulse triggered on rising edge - inputs['desire_pulse'][0] = 0 - self.npy['desire'][:] = np.where(inputs['desire_pulse'] - self.prev_desire > .99, inputs['desire_pulse'], 0) - self.prev_desire[:] = inputs['desire_pulse'] - self.npy['traffic_convention'][:] = inputs['traffic_convention'] - self.npy['tfm'][:,:] = transforms['img'][:,:] - self.npy['big_tfm'][:,:] = transforms['big_img'][:,:] + out = self.update_imgs(self.img_queues['img'], self.full_frames['img'], self.transforms['img'], + self.img_queues['big_img'], self.full_frames['big_img'], self.transforms['big_img']) + vision_inputs = {'img': out[0], 'big_img': out[1]} if prepare_only: - self.warp_enqueue(**self.input_queues, frame=self.full_frames['img'], big_frame=self.full_frames['big_img']) return None - vision_output, policy_output = self.run_policy( - **self.input_queues, frame=self.full_frames['img'], big_frame=self.full_frames['big_img'] - ) + self.vision_output = self.vision_run(**vision_inputs).contiguous().realize().uop.base.buffer.numpy().flatten() + vision_outputs_dict = self.parser.parse_vision_outputs(self.slice_outputs(self.vision_output, self.vision_output_slices)) - vision_output = vision_output.numpy().flatten() - policy_output = policy_output.numpy().flatten() - vision_outputs_dict = self.parser.parse_vision_outputs(self.slice_outputs(vision_output, self.vision_output_slices)) - policy_outputs_dict = self.parser.parse_policy_outputs(self.slice_outputs(policy_output, self.policy_output_slices)) + self.full_input_queues.enqueue({'features_buffer': vision_outputs_dict['hidden_state'], 'desire_pulse': new_desire}) + for k in ['desire_pulse', 'features_buffer']: + self.numpy_inputs[k][:] = self.full_input_queues.get(k)[k] + self.numpy_inputs['traffic_convention'][:] = inputs['traffic_convention'] + + self.policy_output = self.policy_run(**self.policy_inputs).contiguous().realize().uop.base.buffer.numpy().flatten() + policy_outputs_dict = self.parser.parse_policy_outputs(self.slice_outputs(self.policy_output, self.policy_output_slices)) combined_outputs_dict = {**vision_outputs_dict, **policy_outputs_dict} - if SEND_RAW_PRED: - combined_outputs_dict['raw_pred'] = np.concatenate([vision_output.copy(), policy_output.copy()]) + combined_outputs_dict['raw_pred'] = np.concatenate([self.vision_output.copy(), self.policy_output.copy()]) + return combined_outputs_dict @@ -157,6 +241,11 @@ def main(demo=False): # also need to move the aux USB interrupts for good timings config_realtime_process(7, 54) + st = time.monotonic() + cloudlog.warning("loading model") + model = ModelState() + cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting") + # visionipc clients while True: available_streams = VisionIpcClient.available_streams("camerad", block=False) @@ -180,11 +269,6 @@ def main(demo=False): if use_extra_client: cloudlog.warning(f"connected extra cam with buffer size: {vipc_client_extra.buffer_len} ({vipc_client_extra.width} x {vipc_client_extra.height})") - st = time.monotonic() - cloudlog.warning("loading model") - model = ModelState(vipc_client_main.width, vipc_client_main.height) - cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting") - # messaging pm = PubMaster(["modelV2", "drivingModelData", "cameraOdometry"]) sm = SubMaster(["deviceState", "carState", "roadCameraState", "liveCalibration", "driverMonitoringState", "carControl", "liveDelay"]) diff --git a/selfdrive/modeld/tinygrad_helpers.py b/selfdrive/modeld/tinygrad_helpers.py index 9c1a0d4689..49a6ed6161 100644 --- a/selfdrive/modeld/tinygrad_helpers.py +++ b/selfdrive/modeld/tinygrad_helpers.py @@ -2,7 +2,7 @@ import json import os from pathlib import Path -MODELS_DIR = Path(__file__).resolve().parent / 'models' +MODELS_DIR = Path(__file__).parent / 'models' COMPILED_FLAGS_PATH = MODELS_DIR / 'tg_compiled_flags.json'