Legion

selfdrive/modeld/helpers.py

@@ -0,0 +1,73 @@
+import json
+import os
+from pathlib import Path
+
+from tinygrad.device import Device
+
+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:
+    return env_dev.split(":", 1)[0].split("+", 1)[-1] or env_dev
+
+  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):
+  prefix = "big_" if usbgpu else ""
+  return MODELS_DIR / f"{prefix}driving_tinygrad.pkl"
+
+
+def usbgpu_present() -> bool:
+  for d in Path("/sys/bus/usb/devices").glob("*"):
+    try:
+      if int((d / "idVendor").read_text(), 16) == USBGPU_VID and \
+          int((d / "idProduct").read_text(), 16) == USBGPU_PID:
+        return True
+    except Exception:
+      pass
+  return False

selfdrive/modeld/modeld_v16.py

@@ -8,6 +8,7 @@ from pathlib import Path
 
 from openpilot.system.hardware import TICI
 
+os.environ["GMMU"] = "0"  # noop on qcom, improves load path when a USB GPU is present
 os.environ["DEV"] = "QCOM" if TICI else "LLVM"
 
 import cereal.messaging as messaging
@@ -16,7 +17,6 @@ from cereal import car, log
 from msgq.visionipc import VisionBuf, VisionIpcClient, VisionStreamType
 from opendbc.car.car_helpers import get_demo_car_params
 from setproctitle import setproctitle
-from tinygrad.device import Device
 from tinygrad.tensor import Tensor
 
 from openpilot.common.file_chunker import read_file_chunked
@@ -31,6 +31,7 @@ from openpilot.selfdrive.controls.lib.drive_helpers import smooth_value
 from openpilot.selfdrive.modeld.compile_modeld import POLICY_INPUTS, WARP_INPUTS, make_input_queues
 from openpilot.selfdrive.modeld.constants import ModelConstants
 from openpilot.selfdrive.modeld.fill_model_msg import PublishState, fill_model_msg, fill_pose_msg
+from openpilot.selfdrive.modeld.helpers import get_tg_input_devices
 from openpilot.selfdrive.modeld.parse_model_outputs import Parser
 from openpilot.starpilot.assets.model_manager import ModelManager
 from openpilot.starpilot.common.model_versions import uses_combined_driving_artifacts
@@ -142,7 +143,7 @@ class FrameMeta:
 class ModelState:
   prev_desire: np.ndarray
 
-  def __init__(self, cam_w: int, cam_h: int):
+  def __init__(self, cam_w: int, cam_h: int, usbgpu: bool):
     params = Params()
     model_id_raw = _resolve_mirrored_param(params, "Model", "DrivingModel") or BUILTIN_MODEL_KEY
     self.model_id = _canonical_model_id(model_id_raw)
@@ -176,11 +177,16 @@ class ModelState:
     self.off_policy_output_slices = off_policy_metadata["output_slices"]
     self.policy_input_shapes = on_policy_metadata["input_shapes"]
     self.policy_output_slices = on_policy_metadata["output_slices"]
-    self.desire_key = next(key for key in self.policy_input_shapes if key.startswith("desire"))
+    self.desire_key = "desire_pulse" if "desire_pulse" in self.policy_input_shapes else next(
+      key for key in self.policy_input_shapes if key.startswith("desire")
+    )
 
     self.frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ
-    self.dev = Device.DEFAULT
-    self.input_queues, self.npy = make_input_queues(self.vision_input_shapes, self.policy_input_shapes, self.frame_skip, device=self.dev)
+    input_devices = get_tg_input_devices(PROCESS_NAME, usbgpu)
+    self.WARP_DEV, self.QUEUE_DEV = input_devices["WARP_DEV"], input_devices["QUEUE_DEV"]
+    self.input_queues, self.npy = make_input_queues(
+      self.vision_input_shapes, self.policy_input_shapes, self.frame_skip, device=self.QUEUE_DEV
+    )
     self.full_frames: dict[str, Tensor] = {}
     self._blob_cache: dict[tuple[str, int], Tensor] = {}
     self.parser = Parser()
@@ -205,7 +211,7 @@ class ModelState:
       yuv_size = self.frame_buf_params[key][3]
       cache_key = (key, ptr)
       if cache_key not in self._blob_cache:
-        self._blob_cache[cache_key] = Tensor.from_blob(ptr, (yuv_size,), dtype="uint8", device=self.dev)
+        self._blob_cache[cache_key] = Tensor.from_blob(ptr, (yuv_size,), dtype="uint8", device=self.WARP_DEV)
       self.full_frames[key] = self._blob_cache[cache_key]
 
     inputs[self.desire_key][0] = 0
@@ -242,7 +248,7 @@ class ModelState:
     off_policy_output = off_policy_output.numpy().flatten()
 
     vision_outputs_dict = self.parser.parse_vision_outputs(self.slice_outputs(vision_output, self.vision_output_slices))
-    off_policy_outputs_dict = self.parser.parse_policy_outputs(self.slice_outputs(off_policy_output, self.off_policy_output_slices))
+    off_policy_outputs_dict = self.parser.parse_off_policy_outputs(self.slice_outputs(off_policy_output, self.off_policy_output_slices))
     policy_outputs_dict = self.parser.parse_policy_outputs(self.slice_outputs(policy_output, self.policy_output_slices))
     combined_outputs_dict = {**vision_outputs_dict, **off_policy_outputs_dict, **policy_outputs_dict}
 
@@ -259,6 +265,10 @@ def main(demo=False):
   setproctitle(PROCESS_NAME)
   config_realtime_process(7, 54)
 
+  # Combined downloaded models currently ship one runtime artifact, so stay on the default
+  # queue profile until a separate USBGPU artifact path exists for custom models.
+  usbgpu = False
+
   while True:
     available_streams = VisionIpcClient.available_streams("camerad", block=False)
     if available_streams:
@@ -283,7 +293,7 @@ def main(demo=False):
 
   start_time = time.monotonic()
   cloudlog.warning("loading combined model")
-  model = ModelState(vipc_client_main.width, vipc_client_main.height)
+  model = ModelState(vipc_client_main.width, vipc_client_main.height, usbgpu)
   cloudlog.warning(f"combined model loaded in {time.monotonic() - start_time:.1f}s, modeld starting")
 
   pm = messaging.PubMaster(["modelV2", "drivingModelData", "cameraOdometry", "starpilotModelV2"])

selfdrive/modeld/parse_model_outputs.py

@@ -120,6 +120,12 @@ class Parser:
     if 'lead_prob' in outs:
       self.parse_binary_crossentropy('lead_prob', outs)
 
+  def parse_off_policy_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
+    self.split_outputs(outs)
+    if 'desire_state' in outs:
+      self.parse_categorical_crossentropy('desire_state', outs, out_shape=(ModelConstants.DESIRE_PRED_WIDTH,))
+    return outs
+
   def parse_vision_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
     self.parse_mdn('pose', outs, in_N=0, out_N=0, out_shape=(ModelConstants.POSE_WIDTH,))
     self.parse_mdn('wide_from_device_euler', outs, in_N=0, out_N=0, out_shape=(ModelConstants.WIDE_FROM_DEVICE_WIDTH,))