mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-07-26 03:32:06 +08:00
firestar5683
119 lines
3.5 KiB
Python
119 lines
3.5 KiB
Python
import io
|
|
import json
|
|
import os
|
|
import pickle
|
|
import shutil
|
|
import struct
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
from tinygrad.device import Device
|
|
from openpilot.system.hardware.usb import chestnut_present
|
|
|
|
MODELS_DIR = Path(__file__).resolve().parent / "models"
|
|
TG_INPUT_DEVICES_PATH = MODELS_DIR / "tg_input_devices.json"
|
|
USBGPU_VID = 0xADD1
|
|
USBGPU_PID = 0x0001
|
|
|
|
|
|
def _default_tinygrad_backend() -> str:
|
|
env_dev = os.getenv("DEV", "").strip()
|
|
if env_dev:
|
|
default_target = env_dev.split(";", 1)[0]
|
|
return default_target.split(":", 1)[0].split("+", 1)[-1] or default_target
|
|
|
|
try:
|
|
default = Device.DEFAULT
|
|
except Exception:
|
|
default = ""
|
|
|
|
if isinstance(default, str) and default:
|
|
return default.split(":", 1)[0].split("+", 1)[-1] or default
|
|
return "QCOM" if Path("/dev/kgsl-3d0").exists() else "CPU"
|
|
|
|
|
|
def _load_tg_devices() -> dict:
|
|
if not TG_INPUT_DEVICES_PATH.is_file():
|
|
return {}
|
|
with open(TG_INPUT_DEVICES_PATH) as f:
|
|
return json.load(f)
|
|
|
|
|
|
def _fallback_tg_devices(process_name: str, usbgpu: bool) -> dict[str, str]:
|
|
backend = _default_tinygrad_backend()
|
|
if process_name == "selfdrive.modeld.dmonitoringmodeld":
|
|
return {"DEV": backend}
|
|
|
|
queue_dev = backend
|
|
if usbgpu:
|
|
try:
|
|
available = {name.split(":", 1)[0] for name in Device.get_available_devices()}
|
|
except Exception:
|
|
available = set()
|
|
if "AMD" in available:
|
|
queue_dev = "AMD"
|
|
return {"WARP_DEV": backend, "QUEUE_DEV": queue_dev}
|
|
|
|
|
|
def get_tg_input_devices(process_name: str, usbgpu: bool) -> dict[str, str]:
|
|
tg_devices = _load_tg_devices()
|
|
process_devices = tg_devices.get(process_name, {})
|
|
profile = "usbgpu" if usbgpu else "default"
|
|
if profile in process_devices:
|
|
return process_devices[profile]
|
|
return _fallback_tg_devices(process_name, usbgpu)
|
|
|
|
|
|
def modeld_pkl_path(usbgpu: bool) -> Path:
|
|
prefix = "big_" if usbgpu else ""
|
|
return MODELS_DIR / f"{prefix}driving_tinygrad.pkl"
|
|
|
|
|
|
def usbgpu_present() -> bool:
|
|
return chestnut_present()
|
|
|
|
|
|
def tinygrad_dev_config(usbgpu: bool, tici: bool) -> str:
|
|
default = "QCOM" if tici else ("CPU:LLVM" if usbgpu else "LLVM")
|
|
return f"{default};USB+AMD:LLVM" if usbgpu else default
|
|
|
|
|
|
def dump_oob(obj, output) -> None:
|
|
"""Pickle large buffers out-of-band so model weights are never duplicated in RAM."""
|
|
with tempfile.TemporaryFile(dir=".") as buffers:
|
|
def buffer_callback(pickle_buffer: pickle.PickleBuffer):
|
|
view = pickle_buffer.raw()
|
|
buffers.write(struct.pack("<q", view.nbytes))
|
|
buffers.write(view)
|
|
pickle_buffer.release()
|
|
|
|
opcodes = io.BytesIO()
|
|
pickle.Pickler(opcodes, protocol=5, buffer_callback=buffer_callback).dump(obj)
|
|
opcode_data = opcodes.getvalue()
|
|
output.write(struct.pack("<q", len(opcode_data)))
|
|
output.write(opcode_data)
|
|
buffers.seek(0)
|
|
shutil.copyfileobj(buffers, output)
|
|
|
|
|
|
def load_oob(source):
|
|
header = source.read(8)
|
|
if len(header) != 8:
|
|
raise EOFError("truncated out-of-band pickle header")
|
|
opcodes = source.read(struct.unpack("<q", header)[0])
|
|
|
|
def buffers():
|
|
previous = None
|
|
while header := source.read(8):
|
|
if len(header) != 8:
|
|
raise EOFError("truncated out-of-band buffer header")
|
|
if previous is not None:
|
|
previous.release()
|
|
buffer = bytearray(struct.unpack("<q", header)[0])
|
|
if source.readinto(buffer) != len(buffer):
|
|
raise EOFError("truncated out-of-band buffer")
|
|
previous = pickle.PickleBuffer(buffer)
|
|
yield previous
|
|
|
|
return pickle.load(io.BytesIO(opcodes), buffers=buffers())
|