Requiem For Rancid Remmy Rental Roadster

selfdrive/modeld/modeld_v16.py

@@ -17,6 +17,7 @@ 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.dtype import dtypes
 from tinygrad.tensor import Tensor
 
 from openpilot.common.file_chunker import read_file_chunked
@@ -28,16 +29,17 @@ from openpilot.common.transformations.camera import DEVICE_CAMERAS
 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, get_curvature_from_plan, smooth_value
-from openpilot.selfdrive.modeld.compile_modeld import POLICY_INPUTS, WARP_INPUTS, make_input_queues
+from openpilot.selfdrive.modeld.compile_modeld import POLICY_INPUTS, make_input_queues
 from openpilot.selfdrive.modeld.constants import ModelConstants, Plan
 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.models.commonmodel_pyx import CLContext, DrivingModelFrame
 from openpilot.selfdrive.modeld.parse_model_outputs import Parser
+from openpilot.selfdrive.modeld.runners.tinygrad_helpers import qcom_tensor_from_opencl_address
 from openpilot.starpilot.assets.model_manager import ModelManager
 from openpilot.starpilot.common.model_versions import uses_combined_driving_artifacts
 from openpilot.starpilot.common.starpilot_variables import MODELS_PATH, get_starpilot_toggles, params_memory
 from openpilot.system import sentry
-from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
 
 
 PROCESS_NAME = "selfdrive.modeld.modeld"
@@ -159,7 +161,7 @@ class FrameMeta:
 class ModelState:
   prev_desire: np.ndarray
 
-  def __init__(self, cam_w: int, cam_h: int, usbgpu: bool):
+  def __init__(self, context: CLContext, 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)
@@ -203,48 +205,44 @@ class ModelState:
     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.frames = {name: DrivingModelFrame(context, ModelConstants.TEMPORAL_SKIP) for name in self.vision_input_names}
+    self.vision_inputs: dict[str, Tensor] = {}
     self.parser = Parser()
-    self.frame_buf_params = {key: get_nv12_info(cam_w, cam_h) for key in ("img", "big_img")}
     self.prev_desire = np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32)
     self.run_policy = jits["run_policy"]
-    self.warp_enqueue = jits[(cam_w, cam_h)]
 
   def slice_outputs(self, model_outputs: np.ndarray, output_slices: dict[str, slice]) -> dict[str, np.ndarray]:
     return {key: model_outputs[np.newaxis, value] for key, value in output_slices.items()}
 
   def run(self, bufs: dict[str, VisionBuf], transforms: dict[str, np.ndarray], inputs: dict[str, np.ndarray], prepare_only: bool) -> dict[str, np.ndarray] | None:
-    for key in bufs.keys():
-      ptr = np.frombuffer(bufs[key].data, dtype=np.uint8).ctypes.data
-      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.WARP_DEV)
-      self.full_frames[key] = self._blob_cache[cache_key]
-
     inputs[self.desire_key][0] = 0
     self.npy["desire"][:] = np.where(inputs[self.desire_key] - self.prev_desire > 0.99, inputs[self.desire_key], 0)
     self.prev_desire[:] = inputs[self.desire_key]
     self.npy["traffic_convention"][:] = inputs["traffic_convention"]
     if "action_t" in self.npy:
       self.npy["action_t"][:] = inputs["action_t"]
-    self.npy["tfm"][:, :] = transforms["img"][:, :]
-    self.npy["big_tfm"][:, :] = transforms["big_img"][:, :]
-
-    img, big_img = self.warp_enqueue(
-      **{key: self.input_queues[key] for key in WARP_INPUTS},
-      frame=self.full_frames["img"],
-      big_frame=self.full_frames["big_img"],
-    )
 
     if prepare_only:
       return None
 
+    imgs_cl = {name: self.frames[name].prepare(bufs[name], transforms[name].flatten()) for name in self.vision_input_names}
+    if TICI:
+      for key in imgs_cl:
+        if key not in self.vision_inputs:
+          self.vision_inputs[key] = qcom_tensor_from_opencl_address(
+            imgs_cl[key].mem_address,
+            self.vision_input_shapes[key],
+            dtype=dtypes.uint8,
+          )
+    else:
+      for key in imgs_cl:
+        frame_input = self.frames[key].buffer_from_cl(imgs_cl[key]).reshape(self.vision_input_shapes[key])
+        self.vision_inputs[key] = Tensor(frame_input, dtype=dtypes.uint8).realize()
+
     vision_output, policy_output, off_policy_output = self.run_policy(
       **{key: self.input_queues[key] for key in POLICY_INPUTS if key in self.input_queues},
-      img=img,
-      big_img=big_img,
+      img=self.vision_inputs["img"],
+      big_img=self.vision_inputs["big_img"],
     )
 
     vision_output = vision_output.numpy().flatten()
@@ -296,8 +294,10 @@ def main(demo=False):
     cloudlog.warning(f"connected extra cam with buffer size: {vipc_client_extra.buffer_len} ({vipc_client_extra.width} x {vipc_client_extra.height})")
 
   start_time = time.monotonic()
+  cloudlog.warning("setting up CL context")
+  cl_context = CLContext()
   cloudlog.warning("loading combined model")
-  model = ModelState(vipc_client_main.width, vipc_client_main.height, usbgpu)
+  model = ModelState(cl_context, usbgpu)
   cloudlog.warning(f"combined model loaded in {time.monotonic() - start_time:.1f}s, modeld starting")
 
   pm = messaging.PubMaster(["modelV2", "drivingModelData", "cameraOdometry", "starpilotModelV2"])