diff --git a/.github/workflows/build-single-tinygrad-model.yaml b/.github/workflows/build-single-tinygrad-model.yaml index cf2d870802..13886c1aa0 100644 --- a/.github/workflows/build-single-tinygrad-model.yaml +++ b/.github/workflows/build-single-tinygrad-model.yaml @@ -34,6 +34,14 @@ on: required: false default: true type: boolean + target_hardware: + description: 'Hardware target to compile for' + required: false + type: choice + default: 'qcom' + options: + - qcom + - usbgpu workflow_dispatch: inputs: upstream_branch: @@ -81,9 +89,17 @@ on: description: 'Minimum selector version' required: false type: string + target_hardware: + description: 'Hardware target to compile for' + required: false + type: choice + default: 'qcom' + options: + - qcom + - usbgpu env: RECOMPILED_DIR: recompiled${{ inputs.recompiled_dir }} - JSON_FILE: docs/docs/driving_models_v${{ inputs.json_version }}.json + JSON_FILE: docs/docs/driving_models_${{ inputs.target_hardware == 'usbgpu' && 'usbgpu_v' || 'v' }}${{ inputs.json_version }}.json jobs: build_model: @@ -93,6 +109,7 @@ jobs: custom_name: ${{ inputs.custom_name || inputs.upstream_branch }} is_20hz: ${{ inputs.is_20hz }} artifact_suffix: ${{ inputs.artifact_suffix }} + target_hardware: ${{ inputs.target_hardware }} secrets: inherit publish_model: diff --git a/.github/workflows/sunnypilot-build-model.yaml b/.github/workflows/sunnypilot-build-model.yaml index 2c435e58a0..1aa25d4cdb 100644 --- a/.github/workflows/sunnypilot-build-model.yaml +++ b/.github/workflows/sunnypilot-build-model.yaml @@ -191,7 +191,7 @@ jobs: if [ "${{ inputs.target_hardware }}" == "usbgpu" ]; then echo "USBGPU build" export USBGPU=1 - TG_FLAGS="DEV=AMD USBGPU=1 IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1" + TG_FLAGS="DEBUG=2 DEV=USB+AMD:LLVM WARP_DEV=QCOM FLOAT16=1 JIT_BATCH_SIZE=0 GMMU=0" OUTPUT_PKL="${{ env.MODELS_DIR }}/big_driving_tinygrad.pkl" else echo "QCOM build" diff --git a/openpilot/common/params_keys.h b/openpilot/common/params_keys.h index 279d36abfd..c8527b974e 100644 --- a/openpilot/common/params_keys.h +++ b/openpilot/common/params_keys.h @@ -197,7 +197,9 @@ inline static std::unordered_map keys = { {"ModelManager_DownloadIndex", {CLEAR_ON_MANAGER_START | CLEAR_ON_ONROAD_TRANSITION, INT}}, {"ModelManager_Favs", {PERSISTENT | BACKUP, STRING}}, {"ModelManager_LastSyncTime", {CLEAR_ON_MANAGER_START | CLEAR_ON_OFFROAD_TRANSITION, INT, "0"}}, + {"ModelManager_LastSyncTime_USBGPU", {CLEAR_ON_MANAGER_START | CLEAR_ON_OFFROAD_TRANSITION, INT, "0"}}, {"ModelManager_ModelsCache", {PERSISTENT | BACKUP, JSON}}, + {"ModelManager_ModelsCache_USBGPU", {PERSISTENT | BACKUP, JSON}}, // Neural Network Lateral Control {"NeuralNetworkLateralControl", {PERSISTENT | BACKUP, BOOL, "0"}}, diff --git a/openpilot/selfdrive/modeld/compile_modeld.py b/openpilot/selfdrive/modeld/compile_modeld.py index ebc3f21a13..2d27a41496 100755 --- a/openpilot/selfdrive/modeld/compile_modeld.py +++ b/openpilot/selfdrive/modeld/compile_modeld.py @@ -3,7 +3,6 @@ import argparse import atexit import math import os -import pickle import tempfile import time import shutil @@ -13,7 +12,6 @@ from collections import namedtuple import numpy as np from openpilot.selfdrive.modeld.helpers import dump_oob, load_oob -from openpilot.selfdrive.modeld.usbgpu_link import wait_usbgpu_link def _patch_tinygrad_fetch_fw(): import hashlib @@ -31,22 +29,6 @@ def _patch_tinygrad_fetch_fw(): helpers.fetch_fw = fetch_fw _patch_tinygrad_fetch_fw() -def _patch_tinygrad_buffer_reduce(): - from tinygrad.device import Buffer - def __reduce_ex__(self, protocol): - buf = None - if self._base is not None: - return self.__class__, (self.device, self.size, self.dtype, None, None, None, 0, self.base, self.offset, self.is_allocated()) - if self.device == "NPY": - return self.__class__, (self.device, self.size, self.dtype, self._buf, self.options, None, self.uop_refcount) - if self.is_allocated(): - buf = bytearray(self.nbytes) - self.copyout(memoryview(buf)) - if protocol >= 5: - buf = pickle.PickleBuffer(buf) - return self.__class__, (self.device, self.size, self.dtype, None, self.options, buf, self.uop_refcount) - Buffer.__reduce_ex__ = __reduce_ex__ -_patch_tinygrad_buffer_reduce() from tinygrad.tensor import Tensor from tinygrad.helpers import Context @@ -312,9 +294,6 @@ if __name__ == "__main__": p.add_argument('--frame-skip', type=int, required=True) args = p.parse_args() - if 'USB+AMD' in os.environ.get('DEV', ''): - wait_usbgpu_link() - model_path = read_file_chunked_to_disk(args.onnx) model_w, model_h = args.model_size diff --git a/openpilot/selfdrive/modeld/helpers.py b/openpilot/selfdrive/modeld/helpers.py index 608bc6aa64..606ad44ae1 100644 --- a/openpilot/selfdrive/modeld/helpers.py +++ b/openpilot/selfdrive/modeld/helpers.py @@ -6,10 +6,11 @@ import struct import tempfile from pathlib import Path +from openpilot.common.file_chunker import get_manifest_path +from openpilot.common.hardware.usb import CHESTNUT_VENDOR_ID, CHESTNUT_PRODUCT_ID, USB_DEVICES_PATH + MODELS_DIR = Path(__file__).resolve().parent / 'models' TG_INPUT_DEVICES_PATH = MODELS_DIR / 'tg_input_devices.json' -USBGPU_VID = 0xADD1 -USBGPU_PID = 0x0001 def get_tg_input_devices(process_name: str, usbgpu: bool): @@ -38,22 +39,21 @@ def dump_oob(obj, f): def load_oob(f): opcodes = f.read(struct.unpack(' bool: - for d in Path("/sys/bus/usb/devices").glob("*"): + for d in USB_DEVICES_PATH.glob("*"): try: - if int((d / "idVendor").read_text(), 16) == USBGPU_VID and \ - int((d / "idProduct").read_text(), 16) == USBGPU_PID: + usb_id = (int((d / "idVendor").read_text(), 16), int((d / "idProduct").read_text(), 16)) + if usb_id == (CHESTNUT_VENDOR_ID, CHESTNUT_PRODUCT_ID): return True except Exception: pass return False + +def usbgpu_compiled() -> bool: + return Path(get_manifest_path(modeld_pkl_path(usbgpu=True))).is_file() diff --git a/openpilot/selfdrive/modeld/modeld.py b/openpilot/selfdrive/modeld/modeld.py index 1ca1a71f35..d7ddce60b9 100755 --- a/openpilot/selfdrive/modeld/modeld.py +++ b/openpilot/selfdrive/modeld/modeld.py @@ -2,6 +2,7 @@ import os os.environ['GMMU'] = '0' # for usbgpu fast loading, noop for qcom from tinygrad.tensor import Tensor +import threading import time import numpy as np import openpilot.cereal.messaging as messaging @@ -18,17 +19,13 @@ from openpilot.common.transformations.camera import DEVICE_CAMERAS from openpilot.system.camerad.cameras.nv12_info import get_nv12_info 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.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, WARP_INPUTS, POLICY_INPUTS 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, get_manifest_path +from openpilot.common.file_chunker import open_file_chunked from openpilot.selfdrive.modeld.constants import ModelConstants, Plan -from openpilot.selfdrive.modeld.helpers import usbgpu_present, modeld_pkl_path, get_tg_input_devices, load_oob -from openpilot.selfdrive.modeld.usbgpu_link import wait_usbgpu_link - -from openpilot.sunnypilot.livedelay.helpers import get_lat_delay -from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase +from openpilot.selfdrive.modeld.helpers import usbgpu_present, usbgpu_compiled, modeld_pkl_path, get_tg_input_devices, load_oob PROCESS_NAME = "openpilot.selfdrive.modeld.modeld" SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') @@ -36,16 +33,17 @@ SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') LAT_SMOOTH_SECONDS = 0.0 LONG_SMOOTH_SECONDS = 0.3 MIN_LAT_CONTROL_SPEED = 0.3 +BIG_MODEL_TIMEOUT = 60 def get_action_from_model(model_output: dict[str, np.ndarray], prev_action: log.ModelDataV2.Action, lat_action_t: float, long_action_t: float, v_ego: float) -> log.ModelDataV2.Action: if 'action' not in model_output: plan = model_output['plan'][0] - desired_accel, should_stop = get_accel_from_plan(plan[:,Plan.VELOCITY][:,0], - plan[:,Plan.ACCELERATION][:,0], - ModelConstants.T_IDXS, - action_t=long_action_t) + desired_accel = get_accel_from_plan(plan[:,Plan.VELOCITY][:,0], + plan[:,Plan.ACCELERATION][:,0], + ModelConstants.T_IDXS, + action_t=long_action_t) desired_curvature = get_curvature_from_plan(plan[:,Plan.T_FROM_CURRENT_EULER][:,2], plan[:,Plan.ORIENTATION_RATE][:,2], ModelConstants.T_IDXS, @@ -54,7 +52,7 @@ def get_action_from_model(model_output: dict[str, np.ndarray], prev_action: log. else: desired_accel = model_output['action'][0,1] desired_curvature = model_output['action'][0,0] / (max(1.0, v_ego))**2 - should_stop = (v_ego < 0.3 and desired_accel < 0.1) + stop = should_stop(v_ego, desired_accel) desired_accel = smooth_value(desired_accel, prev_action.desiredAcceleration, LONG_SMOOTH_SECONDS) if v_ego > MIN_LAT_CONTROL_SPEED: desired_curvature = smooth_value(desired_curvature, prev_action.desiredCurvature, LAT_SMOOTH_SECONDS) @@ -63,7 +61,7 @@ def get_action_from_model(model_output: dict[str, np.ndarray], prev_action: log. return log.ModelDataV2.Action(desiredCurvature=float(desired_curvature), desiredAcceleration=float(desired_accel), - shouldStop=bool(should_stop)) + shouldStop=bool(stop)) class FrameMeta: @@ -76,12 +74,10 @@ class FrameMeta: self.frame_id, self.timestamp_sof, self.timestamp_eof = vipc.frame_id, vipc.timestamp_sof, vipc.timestamp_eof -class ModelState(ModelStateBase): +class ModelState: prev_desire: np.ndarray # for tracking the rising edge of the pulse def __init__(self, cam_w: int, cam_h: int, usbgpu: bool): - ModelStateBase.__init__(self) - self.LAT_SMOOTH_SECONDS = LAT_SMOOTH_SECONDS input_devices = get_tg_input_devices(PROCESS_NAME, usbgpu) self.WARP_DEV, self.QUEUE_DEV = input_devices['WARP_DEV'], input_devices['QUEUE_DEV'] jits = load_oob(open_file_chunked(modeld_pkl_path(usbgpu))) @@ -138,16 +134,24 @@ class ModelState(ModelStateBase): outputs_dict['raw_pred'] = model_output.copy() return outputs_dict + def warmup(self) -> None: + dummy_frames = {k: np.zeros(self.frame_buf_params[k][3], dtype=np.uint8) for k in self.vision_input_names} + eye = np.eye(3, dtype=np.float32) + 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 = make_input_queues(self.input_shapes, self.frame_skip, device=self.QUEUE_DEV) + self.prev_desire[:] = 0 + self.full_frames.clear() + self._blob_cache.clear() + def main(demo=False): cloudlog.warning("modeld init") - _present = usbgpu_present() - _compiled = os.path.isfile(get_manifest_path(modeld_pkl_path(usbgpu=True))) - USBGPU = _present and _compiled + USBGPU = usbgpu_present() and usbgpu_compiled() params = Params() - params.put_bool("UsbGpuPresent", _present) - params.put_bool("UsbGpuCompiled", _compiled) + params.put_bool("UsbGpuLoading", USBGPU) + params.remove("UsbGpuActive") config_realtime_process(7, 54) @@ -174,15 +178,33 @@ 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})") - if USBGPU: - wait_usbgpu_link() st = time.monotonic() cloudlog.warning("loading model") - model = ModelState(vipc_client_main.width, vipc_client_main.height, USBGPU) + model = None + if USBGPU: + big_model = None + def load_big(): + nonlocal big_model + try: + m = ModelState(vipc_client_main.width, vipc_client_main.height, True) + m.warmup() + big_model = m + except Exception: + cloudlog.exception("big model load failed") + loader = threading.Thread(target=load_big, daemon=True) + loader.start() + loader.join(BIG_MODEL_TIMEOUT) + model = big_model + params.put_bool("UsbGpuActive", model is not None) + + small_model = ModelState(vipc_client_main.width, vipc_client_main.height, False) if model is None or USBGPU else None + if model is None: + model = small_model + params.put_bool("UsbGpuLoading", False) cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting") # messaging - pm = PubMaster(["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"]) + pm = PubMaster(["modelV2", "drivingModelData", "cameraOdometry"]) sm = SubMaster(["deviceState", "carState", "roadCameraState", "liveCalibration", "driverMonitoringState", "carControl", "liveDelay"]) publish_state = PublishState() @@ -252,9 +274,7 @@ def main(demo=False): is_rhd = sm["driverMonitoringState"].isRHD frame_id = sm["roadCameraState"].frameId v_ego = max(sm["carState"].vEgo, 0.) - if sm.frame % 60 == 0: - model.lat_delay = get_lat_delay(params, sm["liveDelay"].lateralDelay) - lat_delay = model.lat_delay + LAT_SMOOTH_SECONDS + lat_delay = sm["liveDelay"].lateralDelay + LAT_SMOOTH_SECONDS if sm.updated["liveCalibration"] and sm.seen['roadCameraState'] and sm.seen['deviceState']: device_from_calib_euler = np.array(sm["liveCalibration"].rpyCalib, dtype=np.float32) dc = DEVICE_CAMERAS[(str(sm['deviceState'].deviceType), str(sm['roadCameraState'].sensor))] @@ -293,7 +313,17 @@ def main(demo=False): } mt1 = time.perf_counter() - model_output = model.run(bufs, transforms, inputs) + try: + model_output = model.run(bufs, transforms, inputs) + except Exception: + if not params.get_bool("UsbGpuActive"): + raise + # fallback to small model + cloudlog.exception("big model failed, fall back to small") + params.put_bool("UsbGpuActive", False) + model = small_model + run_count = 0 + model_output = None mt2 = time.perf_counter() model_execution_time = mt2 - mt1 @@ -301,7 +331,6 @@ def main(demo=False): modelv2_send = messaging.new_message('modelV2') drivingdata_send = messaging.new_message('drivingModelData') posenet_send = messaging.new_message('cameraOdometry') - mdv2sp_send = messaging.new_message('modelDataV2SP') action = get_action_from_model(model_output, prev_action, lat_action_t, long_action_t, v_ego) prev_action = action @@ -316,14 +345,12 @@ def main(demo=False): DH.update(sm['carState'], sm['carControl'].latActive, lane_change_prob) modelv2_send.modelV2.meta.laneChangeState = DH.lane_change_state modelv2_send.modelV2.meta.laneChangeDirection = DH.lane_change_direction - mdv2sp_send.modelDataV2SP.laneTurnDirection = DH.lane_turn_direction fill_driving_model_data(drivingdata_send, modelv2_send) fill_pose_msg(posenet_send, model_output, meta_main.frame_id, vipc_dropped_frames, meta_main.timestamp_eof, live_calib_seen) pm.send('modelV2', modelv2_send) pm.send('drivingModelData', drivingdata_send) pm.send('cameraOdometry', posenet_send) - pm.send('modelDataV2SP', mdv2sp_send) last_vipc_frame_id = meta_main.frame_id diff --git a/openpilot/sunnypilot/modeld_v2/SConscript b/openpilot/sunnypilot/modeld_v2/SConscript index daaa199ea9..5ad2cddf48 100644 --- a/openpilot/sunnypilot/modeld_v2/SConscript +++ b/openpilot/sunnypilot/modeld_v2/SConscript @@ -1,9 +1,12 @@ import os import glob +import sys +import subprocess from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE from openpilot.common.hardware import HARDWARE, PC +from openpilot.selfdrive.modeld.helpers import usbgpu_present Import('env', 'arch', 'release') lenv = env.Clone() @@ -22,14 +25,21 @@ def get_camera_configs(): CAMERA_CONFIGS = get_camera_configs() -tg_flags = { - 'larch64': 'DEV=QCOM FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0', - 'Darwin': f'DEV=CPU HOME={os.path.expanduser("~")}', -}.get(arch, 'DEV=CPU:LLVM') +def probe_devices(): + return set(subprocess.run( + [sys.executable, '-c', 'from tinygrad import Device\nprint("\\n".join(Device.get_available_devices()))'], + capture_output=True, text=True, check=True).stdout.strip().splitlines()) -image_flag = { - 'larch64': 'IMAGE=2', -}.get(arch, 'IMAGE=0') +available = probe_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} IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1' +else: + tg_backend = 'CPU' + tg_flags = f'DEV=CPU HOME={os.path.expanduser("~")}' if arch == 'Darwin' else 'DEV=CPU:LLVM' model_w, model_h = MEDMODEL_INPUT_SIZE from openpilot.selfdrive.modeld.constants import ModelConstants @@ -41,17 +51,30 @@ compile_modeld_script = File("compile_modeld.py").abspath upstream_compile_script = File(Dir("#openpilot/selfdrive/modeld").File("compile_modeld.py").abspath) script_deps = [File("compile_modeld.py"), upstream_compile_script] +USBGPU = usbgpu_present() +if USBGPU: + usbgpu_tg_flags = f'DEBUG=2 DEV=USB+AMD:LLVM WARP_DEV={tg_backend} FLOAT16=1 JIT_BATCH_SIZE=0 GMMU=0' + usbgpu_lock = File("models/.usb_gpu.lock").abspath + def compile_combined(model_type, onnx_args, output_name): - output_pkl = File(f"models/{output_name}").abspath - cmd = (f'{pythonpath_string} {tg_flags} {image_flag} python3 {compile_modeld_script} ' - f'--model-type {model_type} ' - f'--model-size {model_w}x{model_h} ' - f'--camera-resolutions {camera_res_args} ' - f'{onnx_args} ' - f'--frame-skip {frame_skip} ' - f'--output {output_pkl}') - onnx_files = [f for f in onnx_args.split() if f.endswith('.onnx')] - return lenv.Command(output_pkl, tinygrad_files + script_deps + [File(f) for f in onnx_files if os.path.isfile(f)], cmd) + for usbgpu in ([False, True] if USBGPU else [False]): + prefix = 'big_' if usbgpu else ('big_' if os.getenv('BIG_INTO_SMALL') else '') + final_output_name = prefix + output_name + output_pkl = File(f"models/{final_output_name}").abspath + + active_tg_flags = usbgpu_tg_flags if usbgpu else tg_flags + + cmd = (f'{pythonpath_string} {active_tg_flags} python3 {compile_modeld_script} ' + f'--model-type {model_type} ' + f'--model-size {model_w}x{model_h} ' + f'--camera-resolutions {camera_res_args} ' + f'{onnx_args} ' + f'--frame-skip {frame_skip} ' + f'--output {output_pkl}') + onnx_files = [f for f in onnx_args.split() if f.endswith('.onnx')] + node = lenv.Command(output_pkl, tinygrad_files + script_deps + [File(f) for f in onnx_files if os.path.isfile(f)], cmd) + if usbgpu: + lenv.SideEffect(usbgpu_lock, node) # Vision + Policy (stock default model) vision_onnx = File("models/driving_vision.onnx").abspath diff --git a/openpilot/sunnypilot/modeld_v2/compile_modeld.py b/openpilot/sunnypilot/modeld_v2/compile_modeld.py index def54a4599..4b3d68a838 100755 --- a/openpilot/sunnypilot/modeld_v2/compile_modeld.py +++ b/openpilot/sunnypilot/modeld_v2/compile_modeld.py @@ -8,10 +8,10 @@ See the LICENSE.md file in the root directory for more details. import argparse import os -import pickle -import time +import tempfile from collections import defaultdict from functools import partial +from openpilot.selfdrive.modeld.helpers import dump_oob, load_oob import numpy as np os.environ['GMMU'] = '0' @@ -76,7 +76,7 @@ def get_policy_npy_shapes(input_shapes: dict, is_supercombo: bool = False) -> tu def generate_queues_and_npy(input_shapes: dict, frame_skip: int, device: str = Device.DEFAULT, - is_supercombo: bool = False, use_packed: bool = True) -> tuple[dict, dict]: + is_supercombo: bool = False) -> tuple[dict, dict]: road_key, _ = _detect_vision_keys(input_shapes) if not road_key: raise ValueError("Vision road key missing from input shapes.") @@ -92,69 +92,44 @@ def generate_queues_and_npy(input_shapes: dict, frame_skip: int, device: str = D desire_shape = input_shapes[desire_key] features_buffer = input_shapes.get('features_buffer') - if use_packed: # remove packed detection block after all models are recompiled - npy_arrays = { - 'tfm': np.zeros((3, 3), dtype=np.float32), - 'big_tfm': np.zeros((3, 3), dtype=np.float32) - } + npy_arrays = { + 'tfm': np.zeros((3, 3), dtype=np.float32), + 'big_tfm': np.zeros((3, 3), dtype=np.float32) + } - shapes, sizes = get_policy_npy_shapes(input_shapes, is_supercombo=is_supercombo) - packed_npy_inputs = np.zeros(sum(sizes), dtype=np.float32) + shapes, sizes = get_policy_npy_shapes(input_shapes, is_supercombo=is_supercombo) + packed_npy_inputs = np.zeros(sum(sizes), dtype=np.float32) - split_indices = np.cumsum(sizes[:-1]) if len(sizes) > 1 else [] - split_views = np.split(packed_npy_inputs, split_indices) if len(sizes) > 0 else [] - for (k, s), v in zip(shapes.items(), split_views, strict=True): - npy_arrays[k] = v.reshape(s) + split_indices = np.cumsum(sizes[:-1]) if len(sizes) > 1 else [] + split_views = np.split(packed_npy_inputs, split_indices) if len(sizes) > 0 else [] + for (k, s), v in zip(shapes.items(), split_views, strict=True): + npy_arrays[k] = v.reshape(s) - 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(), - 'desire_q': Tensor(np.zeros((frame_skip * desire_shape[1], desire_shape[0], desire_shape[2]), - dtype=np.float32), device=device).contiguous().realize(), - 'packed_npy_inputs': Tensor(packed_npy_inputs, device='NPY').realize(), - } + 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(), + 'desire_q': Tensor(np.zeros((frame_skip * desire_shape[1], desire_shape[0], desire_shape[2]), + dtype=np.float32), device=device).contiguous().realize(), + 'packed_npy_inputs': Tensor(packed_npy_inputs, device='NPY').realize(), + } - if features_buffer: - queues['feat_q'] = Tensor(np.zeros((frame_skip * (features_buffer[1] - 1) + 1, features_buffer[0], features_buffer[2]), - dtype=np.float32), device=device).contiguous().realize() + if features_buffer: + queues['feat_q'] = Tensor(np.zeros((frame_skip * (features_buffer[1] - 1) + 1, features_buffer[0], features_buffer[2]), + dtype=np.float32), device=device).contiguous().realize() - queues.update({key: Tensor(value, device='NPY').realize() for key, value in npy_arrays.items() if key in ('tfm', 'big_tfm')}) - else: - # TODO-SP: Remove legacy queuing fallback else block after all models are recompiled - npy_arrays = { - 'desire': np.zeros(desire_shape[2], dtype=np.float32), - 'tfm': np.zeros((3, 3), dtype=np.float32), - 'big_tfm': np.zeros((3, 3), dtype=np.float32) - } - - for key, shape in input_shapes.items(): - if key not in npy_arrays and 'img' not in key and key not in ('features_buffer', desire_key): - npy_arrays[key] = np.zeros(shape, dtype=np.float32) - - 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(), - 'desire_q': Tensor(np.zeros((frame_skip * desire_shape[1], desire_shape[0], desire_shape[2]), - dtype=np.float32), device=device).contiguous().realize() - } - - if features_buffer: - queues['feat_q'] = Tensor(np.zeros((frame_skip * (features_buffer[1] - 1) + 1, features_buffer[0], features_buffer[2]), - dtype=np.float32), device=device).contiguous().realize() - - queues.update({key: Tensor(value, device='NPY').realize() for key, value in npy_arrays.items()}) + queues.update({key: Tensor(value, device='NPY').realize() for key, value in npy_arrays.items() if key in ('tfm', 'big_tfm')}) return queues, npy_arrays def make_split_input_queues(vision_input_shapes: dict, policy_input_shapes: dict, - frame_skip: int, device: str = Device.DEFAULT, use_packed: bool = True) -> tuple[dict, dict]: - return generate_queues_and_npy({**vision_input_shapes, **policy_input_shapes}, frame_skip, device, is_supercombo=False, use_packed=use_packed) + frame_skip: int, device: str = Device.DEFAULT) -> tuple[dict, dict]: + return generate_queues_and_npy({**vision_input_shapes, **policy_input_shapes}, frame_skip, device, is_supercombo=False) def make_supercombo_input_queues(input_shapes: dict, frame_skip: int, - device: str = Device.DEFAULT, use_packed: bool = True) -> tuple[dict, dict]: - return generate_queues_and_npy(input_shapes, frame_skip, device, is_supercombo=True, use_packed=use_packed) + device: str = Device.DEFAULT) -> tuple[dict, dict]: + return generate_queues_and_npy(input_shapes, frame_skip, device, is_supercombo=True) def create_jit_runner(vision_runner, policy_runners: list, nv12: NV12Frame, model_size: tuple[int, int], @@ -233,29 +208,34 @@ def compile_and_warmup(nv12: NV12Frame, model_size: tuple[int, int], prepare_onl raise ValueError("Could not find vision, model, or policy metadata.") features_slice = feat_meta['output_slices']['hidden_state'] - WARP_DEV = 'CPU' if "USBGPU" in os.environ else Device.DEFAULT + WARP_DEV = os.getenv('WARP_DEV', Device.DEFAULT) is_supercombo = vision_runner is None run_func = create_jit_runner(vision_runner, policy_runners, nv12, model_size, features_slice, frame_skip, all_shapes, prepare_only) run_jit = TinyJit(run_func, prune=True) - queues, npy_arrays = generate_queues_and_npy(all_shapes, frame_skip, Device.DEFAULT, is_supercombo=is_supercombo) - - for i in range(3): - rng = np.random.default_rng(42 + i) + def run_once(seed): + queues, npy = generate_queues_and_npy(all_shapes, frame_skip, Device.DEFAULT, is_supercombo=is_supercombo) + rng = np.random.default_rng(seed) frame = Tensor.randint(nv12.size, low=0, high=256, dtype=dtypes.uint8, device=WARP_DEV).realize() big_frame = Tensor.randint(nv12.size, low=0, high=256, dtype=dtypes.uint8, device=WARP_DEV).realize() - for arr in npy_arrays.values(): - arr[:] = rng.standard_normal(arr.shape).astype(arr.dtype) - + for value in npy.values(): + value[:] = rng.standard_normal(value.shape).astype(value.dtype) Device.default.synchronize() - start_time = time.perf_counter() - run_jit(**queues, frame=frame, big_frame=big_frame) - mid_time = time.perf_counter() + outs = run_jit(**queues, frame=frame, big_frame=big_frame) Device.default.synchronize() - print(f" [{i + 1}/3] enqueue {(mid_time - start_time) * 1e3:6.2f} ms -- total {(time.perf_counter() - start_time) * 1e3:6.2f} ms") + return [np.copy(value.numpy()) for value in (outs if isinstance(outs, tuple) else [outs])] if outs is not None else [] - # TODO-SP: switch to dump_oob/load_oob on next full recompile of all models - return pickle.loads(pickle.dumps(run_jit)) if not prepare_only else run_jit + for i in range(3): + run_once(42 + i) + + if not prepare_only: + baseline = run_once(42) + with tempfile.TemporaryFile(dir=".") as f: + dump_oob(run_jit, f) + f.seek(0) + run_jit = load_oob(f) + assert all(np.array_equal(baseline, deserialized) for baseline, deserialized in zip(baseline, run_once(42), strict=True)), "OOB pickling regression" + return run_jit def _parse_size(size_str: str) -> tuple[int, int]: @@ -352,8 +332,7 @@ if __name__ == "__main__": vision_runner, policy_runners, output_data['metadata'])) with open(args.output, "wb") as file: - # TODO-SP: switch to dump_oob from openpilot/selfdrive/helpers on next full recompile of all models - pickle.dump(output_data, file) + dump_oob(output_data, file) pkl_size = os.path.getsize(args.output) print(f"Saved combined JIT to {args.output} ({pkl_size / 1e6:.2f} MB)") diff --git a/openpilot/sunnypilot/modeld_v2/modeld.py b/openpilot/sunnypilot/modeld_v2/modeld.py index b53ab18c73..153e1f375b 100755 --- a/openpilot/sunnypilot/modeld_v2/modeld.py +++ b/openpilot/sunnypilot/modeld_v2/modeld.py @@ -10,11 +10,14 @@ import os os.environ['GMMU'] = '0' from openpilot.common.hardware import TICI os.environ['DEV'] = 'QCOM' if TICI else 'CPU' -USBGPU = "USBGPU" in os.environ + +from openpilot.selfdrive.modeld.helpers import usbgpu_present, load_oob +from openpilot.selfdrive.modeld.usbgpu_link import wait_usbgpu_link + +USBGPU = usbgpu_present() if USBGPU: os.environ['DEV'] = 'AMD' os.environ['AMD_IFACE'] = 'USB' -import pickle import time import numpy as np import openpilot.cereal.messaging as messaging @@ -108,11 +111,10 @@ class ModelState(ModelStateBase): def _init_combined(self, pkl_path, cam_w, cam_h, bundle): cloudlog.warning(f"loading combined pkl: {pkl_path}") - # TODO-SP: switch to load_oob from openpilot/selfdrive/helpers on next full recompile of all models - jits = pickle.load(open_file_chunked(pkl_path)) + jits = load_oob(open_file_chunked(pkl_path)) self.DEV = Device.DEFAULT - self.WARP_DEV = 'CPU' if USBGPU else self.DEV + self.WARP_DEV = ('QCOM' if TICI else 'CPU') if USBGPU else self.DEV self.QUEUE_DEV = self.DEV metadata = jits['metadata'] @@ -120,13 +122,6 @@ class ModelState(ModelStateBase): self._run_policy = jits[(cam_w, cam_h)]['run_policy'] self._warp_enqueue = jits[(cam_w, cam_h)]['warp_enqueue'] - # TODO-SP: Remove legacy use_packed detection block after all models are recompiled - captured = getattr(self._run_policy, 'captured', None) - if captured is not None: - use_packed = 'packed_npy_inputs' in getattr(captured, 'expected_names', []) - else: - use_packed = True - if 'model' in metadata: model_metadata = metadata['model'] self.vision_output_slices = model_metadata['output_slices'] @@ -137,7 +132,7 @@ class ModelState(ModelStateBase): from openpilot.sunnypilot.modeld_v2.compile_modeld import make_supercombo_input_queues frame_skip = derive_frame_skip({}, model_metadata['input_shapes']) self.input_queues, self.numpy_inputs = make_supercombo_input_queues(model_metadata['input_shapes'], - frame_skip, device=self.QUEUE_DEV, use_packed=use_packed) + frame_skip, device=self.QUEUE_DEV) else: vision_metadata = metadata['vision'] policy_keys = [k for k in metadata if k != 'vision'] @@ -156,7 +151,7 @@ class ModelState(ModelStateBase): self._vision_input_names = [k for k in vision_input_shapes if 'img' in k] frame_skip = derive_frame_skip(vision_input_shapes, policy_input_shapes) self.input_queues, self.numpy_inputs = make_split_input_queues(vision_input_shapes, policy_input_shapes, - frame_skip, device=self.QUEUE_DEV, use_packed=use_packed) + frame_skip, device=self.QUEUE_DEV) self._desire_key = next(key for key in self.numpy_inputs if key.startswith('desire')) self._road_key = next(key for key in self._vision_input_names if 'big' not in key) @@ -189,6 +184,26 @@ class ModelState(ModelStateBase): frame=Tensor(np.zeros(yuv_size, dtype=np.uint8), device=self.WARP_DEV).contiguous().realize(), big_frame=Tensor(np.zeros(yuv_size, dtype=np.uint8), device=self.WARP_DEV).contiguous().realize()) + if USBGPU: + self.warmup() + + def warmup(self) -> None: + dummy_frames = {k: np.zeros(self.frame_buf_params[k][3], dtype=np.uint8) for k in self._vision_input_names} + transforms = {k: np.eye(3, dtype=np.float32) for k in [self._road_key, self._wide_key] if k} + + dummy_inputs = {} + for k, v in self.numpy_inputs.items(): + if k not in ['tfm', 'big_tfm', 'prev_feat']: + dummy_inputs[k] = np.zeros(v.shape, dtype=v.dtype) + + self.run(dummy_frames, transforms, dummy_inputs, prepare_only=False) + + for v in self.numpy_inputs.values(): + v[:] = 0 + self.prev_desire[:] = 0 + self.full_frames.clear() + self._blob_cache.clear() + @property def mlsim(self) -> bool: @@ -265,11 +280,6 @@ class ModelState(ModelStateBase): buf[0, :-1] = buf[0, 1:] buf[0, -1, :] = outputs['desired_curvature'][0, :] if not self.mlsim else 0 - # TODO-SP: This is a hack to prevent GPU corruption by calculating in CPU space, it can be removed on next recompile - if 'prev_feat' not in self.numpy_inputs and 'feat_q' in self.input_queues: - feat_val = self.input_queues['feat_q'].numpy() - self.input_queues['feat_q'].assign(feat_val).realize() - return outputs def get_action_from_model(self, model_output: dict[str, np.ndarray], prev_action: log.ModelDataV2.Action, @@ -306,6 +316,9 @@ def main(demo=False): setproctitle(PROCESS_NAME) config_realtime_process(7, 54) + if USBGPU: + wait_usbgpu_link() + # visionipc clients while True: available_streams = VisionIpcClient.available_streams("camerad", block=False) @@ -340,6 +353,9 @@ def main(demo=False): publish_state = PublishState() params = Params() + params.put_bool("UsbGpuPresent", USBGPU) + params.put_bool("UsbGpuCompiled", USBGPU) + # setup filter to track dropped frames frame_dropped_filter = FirstOrderFilter(0., 10., 1. / model.constants.MODEL_FREQ) frame_id = 0 diff --git a/openpilot/sunnypilot/models/fetcher.py b/openpilot/sunnypilot/models/fetcher.py index eda1117a2a..b3372eb6c7 100644 --- a/openpilot/sunnypilot/models/fetcher.py +++ b/openpilot/sunnypilot/models/fetcher.py @@ -13,6 +13,7 @@ from openpilot.common.params import Params from openpilot.common.swaglog import cloudlog from openpilot.common.hardware.hw import Paths from openpilot.sunnypilot.models.helpers import is_bundle_version_compatible +from openpilot.selfdrive.modeld.helpers import usbgpu_present from openpilot.cereal import custom @@ -103,11 +104,11 @@ class ModelParser: class ModelCache: """Handles caching of model data to avoid frequent remote fetches""" - def __init__(self, params: Params, cache_timeout: int = int(3600 * 1e9)): + def __init__(self, params: Params, cache_timeout: int = int(3600 * 1e9), suffix: str = ""): self.params = params self.cache_timeout = cache_timeout - self._LAST_SYNC_KEY = "ModelManager_LastSyncTime" - self._CACHE_KEY = "ModelManager_ModelsCache" + self._LAST_SYNC_KEY = f"ModelManager_LastSyncTime{suffix}" + self._CACHE_KEY = f"ModelManager_ModelsCache{suffix}" def _is_expired(self) -> bool: """Checks if the cache has expired""" @@ -139,24 +140,28 @@ class ModelCache: class ModelFetcher: """Handles fetching and caching of model data from remote source""" - MODEL_URL = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_v18.json" def __init__(self, params: Params): self.params = params - self.model_cache = ModelCache(params) self.model_parser = ModelParser() + if usbgpu_present(): + self.model_cache = ModelCache(params, suffix="_USBGPU") + self.model_url = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_usbgpu_v18.json" + else: + self.model_cache = ModelCache(params) + self.model_url = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_v18.json" def _fetch_and_cache_models(self) -> list[custom.ModelManagerSP.ModelBundle] | None: """Fetches fresh model data from remote and updates cache. Returns None on transport errors. Raises on 404 and other fatal HTTP errors. """ try: - response = requests.get(self.MODEL_URL, timeout=10) + response = requests.get(self.model_url, timeout=10) # Explicitly handle 404 differently if response.status_code == 404: - cloudlog.error(f"Models URL returned 404 Not Found: {self.MODEL_URL}") - raise HTTPError(f"404 Not Found: {self.MODEL_URL}", response=response) + cloudlog.error(f"Models URL returned 404 Not Found: {self.model_url}") + raise HTTPError(f"404 Not Found: {self.model_url}", response=response) # Raise for any other 4xx/5xx response.raise_for_status() diff --git a/openpilot/sunnypilot/models/tests/test_tinygrad_ref.py b/openpilot/sunnypilot/models/tests/test_tinygrad_ref.py index 3e60ab5308..d32a0e985a 100644 --- a/openpilot/sunnypilot/models/tests/test_tinygrad_ref.py +++ b/openpilot/sunnypilot/models/tests/test_tinygrad_ref.py @@ -1,11 +1,12 @@ import requests +from openpilot.common.params import Params from openpilot.sunnypilot.models.tinygrad_ref import get_tinygrad_ref from openpilot.sunnypilot.models.fetcher import ModelFetcher - def fetch_tinygrad_ref(): - response = requests.get(ModelFetcher.MODEL_URL, timeout=10) + fetcher = ModelFetcher(Params()) + response = requests.get(fetcher.model_url, timeout=10) response.raise_for_status() json_data = response.json() return json_data.get("tinygrad_ref")