From 21538e5a099f525c42b6f5cca440de2c9a22b069 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Harald=20Sch=C3=A4fer?= Date: Wed, 8 Apr 2026 13:55:00 -0700 Subject: [PATCH] autodetect tg backend (#37778) * pick fastest * save config * fix * Ignore generated tg_compiled_flags file * helper * cleaner * whitespace not needed * no shebang * whitespace --- .gitignore | 1 + selfdrive/modeld/SConscript | 45 ++++++++++++++++++++------- selfdrive/modeld/dmonitoringmodeld.py | 11 +++---- selfdrive/modeld/modeld.py | 7 ++--- selfdrive/modeld/tinygrad_helpers.py | 11 +++++++ 5 files changed, 54 insertions(+), 21 deletions(-) create mode 100644 selfdrive/modeld/tinygrad_helpers.py diff --git a/.gitignore b/.gitignore index 1f58a371e0..c4022a8653 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 selfdrive/pandad/pandad diff --git a/selfdrive/modeld/SConscript b/selfdrive/modeld/SConscript index bad1cdd500..f45dd4e7c6 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,51 @@ 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:LLVM THREADS=0') +# get fastest TG config +available = set(Device.get_available_devices()) +if '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' + 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_off_policy', 'driving_on_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 +70,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/dmonitoringmodeld.py b/selfdrive/modeld/dmonitoringmodeld.py index efd8214b9f..d5901e8943 100755 --- a/selfdrive/modeld/dmonitoringmodeld.py +++ b/selfdrive/modeld/dmonitoringmodeld.py @@ -1,12 +1,12 @@ #!/usr/bin/env python3 import os -from openpilot.system.hardware import TICI -os.environ['DEV'] = 'QCOM' if TICI else 'CPU' +from openpilot.selfdrive.modeld.tinygrad_helpers import MODELS_DIR, set_tinygrad_backend_from_compiled_flags +set_tinygrad_backend_from_compiled_flags() + 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 +21,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 82e750cf8b..c61e417e1a 100755 --- a/selfdrive/modeld/modeld.py +++ b/selfdrive/modeld/modeld.py @@ -1,7 +1,8 @@ #!/usr/bin/env python3 import os -from openpilot.system.hardware import TICI -os.environ['DEV'] = 'QCOM' if TICI else 'CPU' +from openpilot.selfdrive.modeld.tinygrad_helpers import MODELS_DIR, set_tinygrad_backend_from_compiled_flags +set_tinygrad_backend_from_compiled_flags() + USBGPU = "USBGPU" in os.environ if USBGPU: os.environ['DEV'] = 'AMD' @@ -12,7 +13,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 @@ -34,7 +34,6 @@ from openpilot.selfdrive.modeld.constants import ModelConstants, Plan 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' ON_POLICY_PKL_PATH = MODELS_DIR / 'driving_on_policy_tinygrad.pkl' diff --git a/selfdrive/modeld/tinygrad_helpers.py b/selfdrive/modeld/tinygrad_helpers.py new file mode 100644 index 0000000000..e833c20968 --- /dev/null +++ b/selfdrive/modeld/tinygrad_helpers.py @@ -0,0 +1,11 @@ +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: + with open(COMPILED_FLAGS_PATH) as f: + os.environ['DEV'] = str(json.load(f)['DEV'])