From 10a33a4bf182a7dcd96789232ccd96c51af267e7 Mon Sep 17 00:00:00 2001 From: Jason Wen Date: Sun, 26 Apr 2026 01:32:55 -0400 Subject: [PATCH] tg: reapply changes (#1817) * Reapply "bump tg (#37700)" This reverts commit 9022b4d3224a05a3509c943e50aee5a03efa8c72. * fixup! Reapply "bump tg (#37700)" * Revert "Revert "autodetect tg backend (#37778)"" This reverts commit b782958a * Reapply "autodetect tg backend: use CPU:LLVM on Linux (#37785)" This reverts commit 3fa6726f88b9afe45fd0393b81fe040a3fb4d553. * Reapply "Modeld: support uncompiled" This reverts commit 8c240cc1a4b07a2f94ccbd4dbcb9a812d7c0293a. * fixup! Reapply "bump tg (#37700)" * fixup! Reapply "autodetect tg backend: use CPU:LLVM on Linux (#37785)" * fixup! Revert "Revert "autodetect tg backend (#37778)"" * fixup! Reapply "autodetect tg backend: use CPU:LLVM on Linux (#37785)" * fixup! Reapply "autodetect tg backend: use CPU:LLVM on Linux (#37785)" * fixup! Revert "Revert "autodetect tg backend (#37778)"" * fixup! Reapply "bump tg (#37700)" * fixup! Reapply "bump tg (#37700)" --- .gitignore | 1 + selfdrive/modeld/SConscript | 46 ++++++++++++++++++++------- selfdrive/modeld/compile_warp.py | 28 ++++++---------- selfdrive/modeld/dmonitoringmodeld.py | 11 ++++--- selfdrive/modeld/modeld.py | 10 +++--- selfdrive/modeld/tinygrad_helpers.py | 12 +++++++ sunnypilot/modeld_v2/warp.py | 19 +++-------- 7 files changed, 76 insertions(+), 51 deletions(-) create mode 100644 selfdrive/modeld/tinygrad_helpers.py diff --git a/.gitignore b/.gitignore index 3434c7254..ed49682f4 100644 --- a/.gitignore +++ b/.gitignore @@ -44,6 +44,7 @@ bin/ config.json compile_commands.json compare_runtime*.html +selfdrive/modeld/models/tg_compiled_flags.json # build artifacts docs_site/ diff --git a/selfdrive/modeld/SConscript b/selfdrive/modeld/SConscript index f02e66722..544fcafc7 100644 --- a/selfdrive/modeld/SConscript +++ b/selfdrive/modeld/SConscript @@ -1,6 +1,9 @@ -import os import glob +import json +import os +from SCons.Script import Value from openpilot.common.file_chunker import chunk_file, get_chunk_paths +from tinygrad import Device Import('env', 'arch') chunker_file = File("#common/file_chunker.py") @@ -13,31 +16,52 @@ 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 -# compile warp # THREADS=0 is need to prevent bug: https://github.com/tinygrad/tinygrad/issues/14689 -tg_flags = { - 'larch64': 'DEV=QCOM FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0', - 'Darwin': f'DEV=CPU THREADS=0 HOME={os.path.expanduser("~")}', # tinygrad calls brew which needs a $HOME in the env -}.get(arch, 'DEV=CPU CPU_LLVM=1 THREADS=0') +# get fastest TG config +available = set(Device.get_available_devices()) +# FIXME-SP: reset when we bump tg +if False: # 'CUDA' in available: + tg_backend = 'CUDA' + 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' +else: + tg_backend = 'CPU' if arch == 'Darwin' else 'CPU CPU_LLVM=1' # FIXME-SP: reset when we bump tg + tg_flags = f'DEV={tg_backend} THREADS=0' + +def write_tg_compiled_flags(target, source, env): + with open(str(target[0]), "w") as f: + json.dump({"DEV": tg_backend}, f) + f.write("\n") + +compiled_flags_node = lenv.Command( + File("models/tg_compiled_flags.json").abspath, + tinygrad_files + [Value(tg_backend)], + write_tg_compiled_flags, +) + +# tinygrad calls brew which needs a $HOME in the env +mac_brew_string = f'HOME={os.path.expanduser("~")}' if arch == 'Darwin' else '' # Get model metadata for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']: fn = File(f"models/{model_name}").abspath script_files = [File(Dir("#selfdrive/modeld").File("get_model_metadata.py").abspath)] - cmd = f'{tg_flags} python3 {Dir("#selfdrive/modeld").abspath}/get_model_metadata.py {fn}.onnx' - lenv.Command(fn + "_metadata.pkl", [fn + ".onnx"] + tinygrad_files + script_files, cmd) + cmd = f'{tg_flags} {mac_brew_string} python3 {Dir("#selfdrive/modeld").abspath}/get_model_metadata.py {fn}.onnx' + lenv.Command(fn + "_metadata.pkl", [fn + ".onnx"] + tinygrad_files + script_files + [compiled_flags_node], cmd) image_flag = { 'larch64': 'IMAGE=2', }.get(arch, 'IMAGE=0') script_files = [File(Dir("#selfdrive/modeld").File("compile_warp.py").abspath)] -compile_warp_cmd = f'{tg_flags} python3 {Dir("#selfdrive/modeld").abspath}/compile_warp.py ' +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, compile_warp_cmd) +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 + '"' @@ -47,7 +71,7 @@ def tg_compile(flags, model_name): chunk_targets = get_chunk_paths(pkl, estimate_pickle_max_size(os.path.getsize(onnx_path))) compile_node = lenv.Command( pkl, - [onnx_path] + tinygrad_files + [chunker_file], + [onnx_path] + tinygrad_files + [chunker_file, compiled_flags_node], f'{pythonpath_string} {flags} {image_flag} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {pkl}', ) def do_chunk(target, source, env): diff --git a/selfdrive/modeld/compile_warp.py b/selfdrive/modeld/compile_warp.py index 75cc65f84..1144fc69c 100755 --- a/selfdrive/modeld/compile_warp.py +++ b/selfdrive/modeld/compile_warp.py @@ -97,8 +97,8 @@ def make_update_img_input(frame_prepare, model_w, model_h): def update_img_input_tinygrad(tensor, frame, M_inv): M_inv = M_inv.to(Device.DEFAULT) new_img = frame_prepare(frame, M_inv) - full_buffer = tensor[6:].cat(new_img, dim=0).contiguous() - return full_buffer, Tensor.cat(full_buffer[:6], full_buffer[-6:], dim=0).contiguous().reshape(1, 12, model_h//2, model_w//2) + tensor.assign(tensor[6:].cat(new_img, dim=0).contiguous()) + return Tensor.cat(tensor[:6], tensor[-6:], dim=0).contiguous().reshape(1, 12, model_h//2, model_w//2) return update_img_input_tinygrad @@ -107,9 +107,9 @@ def make_update_both_imgs(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_buffer, calib_img_pair = update_img(calib_img_buffer, new_img, M_inv) - calib_big_img_buffer, calib_big_img_pair = update_img(calib_big_img_buffer, new_big_img, M_inv_big) - return calib_img_buffer, calib_img_pair, calib_big_img_buffer, calib_big_img_pair + 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 @@ -136,29 +136,20 @@ def compile_modeld_warp(cam_w, cam_h): full_buffer = Tensor.zeros(IMG_BUFFER_SHAPE, dtype='uint8').contiguous().realize() big_full_buffer = Tensor.zeros(IMG_BUFFER_SHAPE, dtype='uint8').contiguous().realize() - full_buffer_np = np.zeros(IMG_BUFFER_SHAPE, dtype=np.uint8) - big_full_buffer_np = np.zeros(IMG_BUFFER_SHAPE, dtype=np.uint8) - + new_frame_np = np.random.randint(0, 256, yuv_size, dtype=np.uint8) + new_big_frame_np = np.random.randint(0, 256, yuv_size, dtype=np.uint8) for i in range(10): - new_frame_np = (32 * np.random.randn(yuv_size).astype(np.float32) + 128).clip(0, 255).astype(np.uint8) img_inputs = [full_buffer, Tensor.from_blob(new_frame_np.ctypes.data, (yuv_size,), dtype='uint8').realize(), Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')] - new_big_frame_np = (32 * np.random.randn(yuv_size).astype(np.float32) + 128).clip(0, 255).astype(np.uint8) big_img_inputs = [big_full_buffer, Tensor.from_blob(new_big_frame_np.ctypes.data, (yuv_size,), dtype='uint8').realize(), Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')] inputs = img_inputs + big_img_inputs Device.default.synchronize() - inputs_np = [x.numpy() for x in inputs] - inputs_np[0] = full_buffer_np - inputs_np[3] = big_full_buffer_np - st = time.perf_counter() - out = update_img_jit(*inputs) - full_buffer = out[0].contiguous().realize().clone() - big_full_buffer = out[2].contiguous().realize().clone() + _ = update_img_jit(*inputs) mt = time.perf_counter() Device.default.synchronize() et = time.perf_counter() @@ -182,8 +173,9 @@ def compile_dm_warp(cam_w, cam_h): warp_dm = make_warp_dm(cam_w, cam_h, dm_w, dm_h) warp_dm_jit = TinyJit(warp_dm, prune=True) + new_frame_np = np.random.randint(0, 256, yuv_size, dtype=np.uint8) for i in range(10): - inputs = [Tensor.from_blob((32 * Tensor.randn(yuv_size,) + 128).cast(dtype='uint8').realize().numpy().ctypes.data, (yuv_size,), dtype='uint8'), + inputs = [Tensor.from_blob(new_frame_np.ctypes.data, (yuv_size,), dtype='uint8').realize(), Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')] Device.default.synchronize() st = time.perf_counter() diff --git a/selfdrive/modeld/dmonitoringmodeld.py b/selfdrive/modeld/dmonitoringmodeld.py index efd8214b9..78749b773 100755 --- a/selfdrive/modeld/dmonitoringmodeld.py +++ b/selfdrive/modeld/dmonitoringmodeld.py @@ -1,12 +1,16 @@ #!/usr/bin/env python3 import os +from openpilot.selfdrive.modeld.tinygrad_helpers import MODELS_DIR, set_tinygrad_backend_from_compiled_flags +set_tinygrad_backend_from_compiled_flags() + +# FIXME-SP: remove once we bump tg from openpilot.system.hardware import TICI os.environ['DEV'] = 'QCOM' if TICI else 'CPU' + from tinygrad.tensor import Tensor import time import pickle import numpy as np -from pathlib import Path from cereal import messaging from cereal.messaging import PubMaster, SubMaster @@ -21,9 +25,8 @@ from openpilot.selfdrive.modeld.parse_model_outputs import sigmoid, safe_exp PROCESS_NAME = "selfdrive.modeld.dmonitoringmodeld" SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') -MODEL_PKL_PATH = Path(__file__).parent / 'models/dmonitoring_model_tinygrad.pkl' -METADATA_PATH = Path(__file__).parent / 'models/dmonitoring_model_metadata.pkl' -MODELS_DIR = Path(__file__).parent / 'models' +MODEL_PKL_PATH = MODELS_DIR / 'dmonitoring_model_tinygrad.pkl' +METADATA_PATH = MODELS_DIR / 'dmonitoring_model_metadata.pkl' class ModelState: inputs: dict[str, np.ndarray] diff --git a/selfdrive/modeld/modeld.py b/selfdrive/modeld/modeld.py index d774e475d..4f89058d7 100755 --- a/selfdrive/modeld/modeld.py +++ b/selfdrive/modeld/modeld.py @@ -1,7 +1,12 @@ #!/usr/bin/env python3 import os +from openpilot.selfdrive.modeld.tinygrad_helpers import MODELS_DIR, set_tinygrad_backend_from_compiled_flags +set_tinygrad_backend_from_compiled_flags() + +# FIXME-SP: remove once we bump tg from openpilot.system.hardware import TICI os.environ['DEV'] = 'QCOM' if TICI else 'CPU' + USBGPU = "USBGPU" in os.environ if USBGPU: os.environ['DEV'] = 'AMD' @@ -12,7 +17,6 @@ import pickle import numpy as np import cereal.messaging as messaging from cereal import car, log -from pathlib import Path from cereal.messaging import PubMaster, SubMaster from msgq.visionipc import VisionIpcClient, VisionStreamType, VisionBuf from opendbc.car.car_helpers import get_demo_car_params @@ -37,7 +41,6 @@ from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase PROCESS_NAME = "selfdrive.modeld.modeld" SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') -MODELS_DIR = Path(__file__).parent / 'models' 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' @@ -217,8 +220,7 @@ class ModelState(ModelStateBase): 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']) - self.img_queues['img'], self.img_queues['big_img'] = out[0].realize(), out[2].realize() - vision_inputs = {'img': out[1], 'big_img': out[3]} + vision_inputs = {'img': out[0], 'big_img': out[1]} if prepare_only: return None diff --git a/selfdrive/modeld/tinygrad_helpers.py b/selfdrive/modeld/tinygrad_helpers.py new file mode 100644 index 000000000..49a6ed616 --- /dev/null +++ b/selfdrive/modeld/tinygrad_helpers.py @@ -0,0 +1,12 @@ +import json +import os +from pathlib import Path + +MODELS_DIR = Path(__file__).parent / 'models' +COMPILED_FLAGS_PATH = MODELS_DIR / 'tg_compiled_flags.json' + + +def set_tinygrad_backend_from_compiled_flags() -> None: + if os.path.isfile(COMPILED_FLAGS_PATH): + with open(COMPILED_FLAGS_PATH) as f: + os.environ['DEV'] = str(json.load(f)['DEV']) diff --git a/sunnypilot/modeld_v2/warp.py b/sunnypilot/modeld_v2/warp.py index 829cbcca4..fd8be4683 100644 --- a/sunnypilot/modeld_v2/warp.py +++ b/sunnypilot/modeld_v2/warp.py @@ -33,29 +33,20 @@ def compile_v2_warp(cam_w, cam_h, buffer_length): full_buffer = Tensor.zeros(img_buffer_shape, dtype='uint8').contiguous().realize() big_full_buffer = Tensor.zeros(img_buffer_shape, dtype='uint8').contiguous().realize() - full_buffer_np = np.zeros(img_buffer_shape, dtype=np.uint8) - big_full_buffer_np = np.zeros(img_buffer_shape, dtype=np.uint8) - + new_frame_np = np.random.randint(0, 256, yuv_size, dtype=np.uint8) + new_big_frame_np = np.random.randint(0, 256, yuv_size, dtype=np.uint8) for i in range(10): - new_frame_np = (32 * np.random.randn(yuv_size).astype(np.float32) + 128).clip(0, 255).astype(np.uint8) img_inputs = [full_buffer, Tensor.from_blob(new_frame_np.ctypes.data, (yuv_size,), dtype='uint8').realize(), Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')] - new_big_frame_np = (32 * np.random.randn(yuv_size).astype(np.float32) + 128).clip(0, 255).astype(np.uint8) big_img_inputs = [big_full_buffer, Tensor.from_blob(new_big_frame_np.ctypes.data, (yuv_size,), dtype='uint8').realize(), Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')] inputs = img_inputs + big_img_inputs Device.default.synchronize() - inputs_np = [x.numpy() for x in inputs] - inputs_np[0] = full_buffer_np - inputs_np[3] = big_full_buffer_np - st = time.perf_counter() - out = update_img_jit(*inputs) - full_buffer = out[0].contiguous().realize().clone() - big_full_buffer = out[2].contiguous().realize().clone() + _ = update_img_jit(*inputs) mt = time.perf_counter() Device.default.synchronize() et = time.perf_counter() @@ -125,8 +116,8 @@ class Warp: self.full_buffers['img'], road_blob, self.transforms['img'], self.full_buffers['big_img'], wide_blob, self.transforms['big_img'], ) - self.full_buffers['img'], out_road = res[0].realize(), res[1].realize() - self.full_buffers['big_img'], out_wide = res[2].realize(), res[3].realize() + out_road = res[0].realize() + out_wide = res[1].realize() return {road: out_road, wide: out_wide}