models: test tinygrad concurrency (PR-2006)

.github/workflows/sunnypilot-build-model.yaml

@@ -103,23 +103,21 @@ jobs:
       - run: |
           cd ${{ github.workspace }}/openpilot/openpilot
           if [ "${{ inputs.target_hardware }}" != "chestnut" ]; then
-            git lfs pull -X "**/selfdrive/modeld/models/big_*.onnx,**/selfdrive/modeld/models/dmonitoring_*.onnx,**/selfdrive/modeld/models/big_*.pkl,**/selfdrive/modeld/models/dmonitoring_*.pkl"
-            rm -f selfdrive/modeld/models/big_*.onnx selfdrive/modeld/models/dmonitoring_*.onnx selfdrive/modeld/models/big_*.pkl selfdrive/modeld/models/dmonitoring_*.pkl
+            git lfs pull -X "**/selfdrive/modeld/models/big_*.onnx,**/selfdrive/modeld/models/dmonitoring_*.onnx"
+            rm -f selfdrive/modeld/models/big_*.onnx selfdrive/modeld/models/dmonitoring_*.onnx
           else
-            git lfs pull -I "**/selfdrive/modeld/models/big_*.onnx,**/selfdrive/modeld/models/big_*.pkl" -X ""
-            find selfdrive/modeld/models -type f \( -name "*.onnx" -o -name "*.pkl" \) ! -name "big_*.onnx" ! -name "big_*.pkl" -delete
+            git lfs pull -I "**/selfdrive/modeld/models/big_*.onnx" -X ""
+            find selfdrive/modeld/models -name "*.onnx" ! -name "big_*.onnx" -delete
           fi
-          if grep -lIF "version https://git-lfs.github.com/spec/v1" selfdrive/modeld/models/*.onnx selfdrive/modeld/models/*.pkl 2>/dev/null; then
-            echo "::error::the ONNX or PKL files above are still LFS pointers, not real models"
+          if grep -lIF "version https://git-lfs.github.com/spec/v1" selfdrive/modeld/models/*.onnx; then
+            echo "::error::the ONNX files above are still LFS pointers, not real models"
             exit 1
           fi
       - name: 'Upload Artifact'
         uses: actions/upload-artifact@v4
         with:
           name: models-${{ env.REF }}${{ inputs.artifact_suffix  }}
-          path: |
-            ${{ github.workspace }}/openpilot/openpilot/selfdrive/modeld/models/*.onnx
-            ${{ github.workspace }}/openpilot/openpilot/selfdrive/modeld/models/*.pkl
+          path: ${{ github.workspace }}/openpilot/openpilot/selfdrive/modeld/models/*.onnx
           if-no-files-found: error
 
   build_model:
@@ -198,75 +196,64 @@ jobs:
             OUTPUT_PKL="${{ env.MODELS_DIR }}/driving_tinygrad.pkl"
           fi
 
-          NATIVE_PKL=$(find "${{ env.MODELS_DIR }}" -maxdepth 1 -name "*.pkl" -print -quit)
+          # Generate metadata for all ONNX files
+          find "${{ env.MODELS_DIR }}" -maxdepth 1 -name '*.onnx' | while IFS= read -r onnx_file; do
+            echo "Generating metadata: $onnx_file"
+            env ${TG_FLAGS_QCOM} python3 "${{ env.MODELS_DIR }}/../get_model_metadata.py" "$onnx_file" || true
+          done
 
-          if [ -n "$NATIVE_PKL" ]; then
-            echo "Found native precompiled pkl: $NATIVE_PKL"
-            if [ "$NATIVE_PKL" != "$OUTPUT_PKL" ]; then
-              mv "$NATIVE_PKL" "$OUTPUT_PKL"
+          # Detect model type and build compile args
+          VISION_ONNX=""
+          for f in "${{ env.MODELS_DIR }}/driving_vision.onnx" "${{ env.MODELS_DIR }}/big_driving_vision.onnx"; do
+            [ -f "$f" ] && VISION_ONNX="$f" && break
+          done
+
+          POLICY_ONNX=""
+          for f in "${{ env.MODELS_DIR }}/driving_policy.onnx" "${{ env.MODELS_DIR }}/big_driving_policy.onnx"; do
+            [ -f "$f" ] && POLICY_ONNX="$f" && break
+          done
+
+          OFF_POLICY_ONNX=""
+          for f in "${{ env.MODELS_DIR }}/driving_off_policy.onnx" "${{ env.MODELS_DIR }}/big_driving_off_policy.onnx"; do
+            [ -f "$f" ] && OFF_POLICY_ONNX="$f" && break
+          done
+
+          ON_POLICY_ONNX=""
+          for f in "${{ env.MODELS_DIR }}/driving_on_policy.onnx" "${{ env.MODELS_DIR }}/big_driving_on_policy.onnx"; do
+            [ -f "$f" ] && ON_POLICY_ONNX="$f" && break
+          done
+
+          SUPERCOMBO_ONNX=""
+          for f in "${{ env.MODELS_DIR }}/supercombo.onnx" "${{ env.MODELS_DIR }}/driving_supercombo.onnx" "${{ env.MODELS_DIR }}/big_supercombo.onnx" "${{ env.MODELS_DIR }}/big_driving_supercombo.onnx"; do
+            [ -f "$f" ] && SUPERCOMBO_ONNX="$f" && break
+          done
+
+          MODEL_TYPE="" ONNX_ARGS="" OUTPUT_NAME=""
+          if [ -f "$VISION_ONNX" ]; then
+            ONNX_ARGS="--vision-onnx $VISION_ONNX"
+            if [ -f "$ON_POLICY_ONNX" ] && [ -f "$OFF_POLICY_ONNX" ]; then
+              MODEL_TYPE=vision_multi_policy
+              ONNX_ARGS="$ONNX_ARGS --off-policy-onnx $OFF_POLICY_ONNX --on-policy-onnx $ON_POLICY_ONNX"
+            elif [ -f "$OFF_POLICY_ONNX" ] && [ -f "$POLICY_ONNX" ]; then
+              MODEL_TYPE=vision_multi_policy
+              ONNX_ARGS="$ONNX_ARGS --policy-onnx $POLICY_ONNX --off-policy-onnx $OFF_POLICY_ONNX"
+            elif [ -f "$POLICY_ONNX" ]; then
+              MODEL_TYPE=vision_policy
+              ONNX_ARGS="$ONNX_ARGS --policy-onnx $POLICY_ONNX"
             fi
-            echo "Chunking pkl"
-            python3 -c "from openpilot.common.file_chunker import chunk_file, get_chunk_targets; import os; p='$OUTPUT_PKL'; chunk_file(p, get_chunk_targets(p, os.path.getsize(p)))"
-          else
-            # Generate metadata for all ONNX files
-            find "${{ env.MODELS_DIR }}" -maxdepth 1 -name '*.onnx' | while IFS= read -r onnx_file; do
-              echo "Generating metadata: $onnx_file"
-              env ${TG_FLAGS_QCOM} python3 "${{ env.MODELS_DIR }}/../get_model_metadata.py" "$onnx_file" || true
-            done
+          elif [ -f "$SUPERCOMBO_ONNX" ]; then
+            MODEL_TYPE=supercombo
+            ONNX_ARGS="--supercombo-onnx $SUPERCOMBO_ONNX"
+          fi
 
-            # Detect model type and build compile args
-            VISION_ONNX=""
-            for f in "${{ env.MODELS_DIR }}/driving_vision.onnx" "${{ env.MODELS_DIR }}/big_driving_vision.onnx"; do
-              [ -f "$f" ] && VISION_ONNX="$f" && break
-            done
-
-            POLICY_ONNX=""
-            for f in "${{ env.MODELS_DIR }}/driving_policy.onnx" "${{ env.MODELS_DIR }}/big_driving_policy.onnx"; do
-              [ -f "$f" ] && POLICY_ONNX="$f" && break
-            done
-
-            OFF_POLICY_ONNX=""
-            for f in "${{ env.MODELS_DIR }}/driving_off_policy.onnx" "${{ env.MODELS_DIR }}/big_driving_off_policy.onnx"; do
-              [ -f "$f" ] && OFF_POLICY_ONNX="$f" && break
-            done
-
-            ON_POLICY_ONNX=""
-            for f in "${{ env.MODELS_DIR }}/driving_on_policy.onnx" "${{ env.MODELS_DIR }}/big_driving_on_policy.onnx"; do
-              [ -f "$f" ] && ON_POLICY_ONNX="$f" && break
-            done
-
-            SUPERCOMBO_ONNX=""
-            for f in "${{ env.MODELS_DIR }}/supercombo.onnx" "${{ env.MODELS_DIR }}/driving_supercombo.onnx" "${{ env.MODELS_DIR }}/big_supercombo.onnx" "${{ env.MODELS_DIR }}/big_driving_supercombo.onnx"; do
-              [ -f "$f" ] && SUPERCOMBO_ONNX="$f" && break
-            done
-
-            MODEL_TYPE="" ONNX_ARGS="" OUTPUT_NAME=""
-            if [ -f "$VISION_ONNX" ]; then
-              ONNX_ARGS="--vision-onnx $VISION_ONNX"
-              if [ -f "$ON_POLICY_ONNX" ] && [ -f "$OFF_POLICY_ONNX" ]; then
-                MODEL_TYPE=vision_multi_policy
-                ONNX_ARGS="$ONNX_ARGS --off-policy-onnx $OFF_POLICY_ONNX --on-policy-onnx $ON_POLICY_ONNX"
-              elif [ -f "$OFF_POLICY_ONNX" ] && [ -f "$POLICY_ONNX" ]; then
-                MODEL_TYPE=vision_multi_policy
-                ONNX_ARGS="$ONNX_ARGS --policy-onnx $POLICY_ONNX --off-policy-onnx $OFF_POLICY_ONNX"
-              elif [ -f "$POLICY_ONNX" ]; then
-                MODEL_TYPE=vision_policy
-                ONNX_ARGS="$ONNX_ARGS --policy-onnx $POLICY_ONNX"
-              fi
-            elif [ -f "$SUPERCOMBO_ONNX" ]; then
-              MODEL_TYPE=supercombo
-              ONNX_ARGS="--supercombo-onnx $SUPERCOMBO_ONNX"
-            fi
-
-            if [ -n "$MODEL_TYPE" ]; then
-              echo "Detected: $MODEL_TYPE -> $OUTPUT_PKL"
-              env ${TG_FLAGS} python3 "$COMPILE_MODELD" \
-                --model-type $MODEL_TYPE \
-                --model-size $MODEL_SIZE \
-                --camera-resolutions $CAMERA_RES \
-                $ONNX_ARGS \
-                --output "$OUTPUT_PKL"
-            fi
+          if [ -n "$MODEL_TYPE" ]; then
+            echo "Detected: $MODEL_TYPE -> $OUTPUT_PKL"
+            env ${TG_FLAGS} python3 "$COMPILE_MODELD" \
+              --model-type $MODEL_TYPE \
+              --model-size $MODEL_SIZE \
+              --camera-resolutions $CAMERA_RES \
+              $ONNX_ARGS \
+              --output "$OUTPUT_PKL"
           fi
 
       - name: Prepare Output

.github/workflows/test_models.yml

@@ -1,5 +1,8 @@
 name: Test Models Compatibility With Tinygrad Changes
 on:
+  push:
+    paths:
+      - 'tinygrad_repo'
   pull_request:
     paths:
       - 'tinygrad_repo'

openpilot/sunnypilot/modeld_v2/compile_modeld.py

@@ -32,7 +32,7 @@ def _patch_tinygrad_fetch_fw():
   helpers.fetch_fw = fetch_fw
 _patch_tinygrad_fetch_fw()
 
-import openpilot.sunnypilot.modeld_v2.stock_dependencies as stock
+import openpilot.selfdrive.modeld.compile_modeld as stock
 from tinygrad import dtypes
 from tinygrad.device import Device
 from tinygrad.engine.jit import TinyJit

openpilot/sunnypilot/modeld_v2/model_adapters.py

@@ -1,177 +0,0 @@
-"""
-Copyright (c) 2021-, Haibin Wen, sunnypilot, and a number of other contributors.
-
-This file is part of sunnypilot and is licensed under the MIT License.
-See the LICENSE.md file in the root directory for more details.
-"""
-import numpy as np
-from tinygrad.tensor import Tensor
-
-from openpilot.sunnypilot.modeld_v2.stock_dependencies import MODELD_INPUTS, make_input_queues as make_stock_input_queues
-from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
-from openpilot.sunnypilot.modeld_v2.compile_modeld import (derive_frame_skip, make_split_input_queues, make_supercombo_input_queues,
-                                                           WARP_INPUTS, POLICY_INPUTS, nv12_copy_size)
-
-
-class BaseModelAdapter:
-  def __init__(self, jits, cam_w, cam_h, model_device, queue_device, warp_device, chestnut=False):
-    self.jits = jits
-    self.DEV = model_device
-    self.QUEUE_DEV = queue_device
-    self.WARP_DEV = warp_device
-    self.chestnut = chestnut
-    self.cam_w = cam_w
-    self.cam_h = cam_h
-    self._combined_model_type = 'supercombo'
-    self._policy_slices_list = []
-    self._has_on_policy = False
-    self.policy_output_slices = {}
-    self._policy_keys = []
-    self.full_frames = {}
-    self._blob_cache = {}
-    self.frame_buffers = {}
-    self.frame_views = {}
-    self.nv12_info = get_nv12_info(cam_w, cam_h)
-
-  def _init_common(self):
-    self._desire_key = next((key for key in getattr(self, 'numpy_inputs', {}) if key.startswith('desire')), 'desire')
-    self._road_key = next((key for key in getattr(self, '_vision_input_names', []) if 'big' not in key), 'img')
-    self._wide_key = next((key for key in getattr(self, '_vision_input_names', []) if 'big' in key), 'big_img')
-    self.frame_buf_params = dict.fromkeys(getattr(self, '_vision_input_names', ['img', 'big_img']), self.nv12_info)
-
-  def get_dummy_inputs(self):
-    dummy_size = getattr(self, 'frame_copy_size', self.frame_buf_params[self._road_key][3])
-    if getattr(self, 'is_run_model', True) is False:
-      dummy_size = self.frame_buf_params[self._road_key][3]
-
-    dummy_frames = {k: np.zeros(dummy_size, 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 = {k: np.zeros(v.shape, dtype=v.dtype) for k, v in self.numpy_inputs.items() if k not in ['tfm', 'big_tfm', 'prev_feat']}
-    return dummy_frames, transforms, dummy_inputs
-
-
-class LegacyModelAdapter(BaseModelAdapter):
-  def __init__(self, *args, **kwargs):
-    super().__init__(*args, **kwargs)
-
-    metadata = self.jits['metadata']
-    self.is_run_model = 'run_model' in self.jits
-    self.frame_copy_size = nv12_copy_size(*self.nv12_info[:3])
-
-    if self.is_run_model or 'model' in metadata:
-      model_metadata = metadata.get('model', metadata)
-      self.input_shapes = model_metadata['input_shapes']
-      self.vision_output_slices = model_metadata['output_slices']
-      self._vision_input_names = [key for key in self.input_shapes if 'img' in key]
-      self.frame_skip = derive_frame_skip({}, self.input_shapes)
-
-      if self.is_run_model:
-        self.input_queues, self.numpy_inputs, self.frame_buffers = make_stock_input_queues(
-          self.input_shapes, self.frame_skip, device=self.DEV, frame_copy_size=self.frame_copy_size)
-        self.frame_views, self.npy = self.frame_buffers, self.numpy_inputs
-        self.run_model, self.run_policy, self.warp = self.jits['run_model'][(self.cam_w, self.cam_h)], None, None
-      else:
-        self.input_queues, self.numpy_inputs = make_supercombo_input_queues(self.input_shapes, self.frame_skip, device=self.QUEUE_DEV)
-        self.run_model, self.run_policy, self.warp = None, self.jits['run_policy'], self.jits[(self.cam_w, self.cam_h)]
-    else:
-      self.run_model, self.run_policy, self.warp = None, self.jits['run_policy'], self.jits[(self.cam_w, self.cam_h)]
-      vision_metadata = metadata['vision']
-      policy_keys = [k for k in metadata if k not in ('vision', 'warp_dev')]
-      self._combined_model_type = 'split' if policy_keys == ['policy'] else 'multi_policy'
-      self.vision_output_slices = vision_metadata['output_slices']
-      self._policy_keys = policy_keys
-      self._policy_slices_list = [metadata[k]['output_slices'] for k in policy_keys]
-      self.policy_output_slices = self._policy_slices_list[0]
-      self._has_on_policy = any('on' in k.lower() for k in policy_keys)
-      self._vision_input_names = [key for key in vision_metadata['input_shapes'] if 'img' in key]
-      first_policy_meta = metadata[policy_keys[0]]
-      self.frame_skip = derive_frame_skip(vision_metadata['input_shapes'], first_policy_meta['input_shapes'])
-      self.input_queues, self.numpy_inputs = make_split_input_queues(vision_metadata['input_shapes'],
-                                                                     first_policy_meta['input_shapes'],
-                                                                     self.frame_skip, device=self.QUEUE_DEV)
-
-    self._init_common()
-    if self.warp is not None:
-      self.full_frames = {k: Tensor(np.zeros(self.nv12_info[3], dtype=np.uint8), device=self.WARP_DEV).contiguous().realize() for k in self._vision_input_names}
-      self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=self.full_frames[self._road_key], big_frame=self.full_frames[self._wide_key])
-
-  def copy_frames(self, bufs):
-    if getattr(self, 'is_run_model', True):
-      for key, buf in bufs.items():
-        data = buf.data if hasattr(buf, 'data') else buf
-        np.copyto(self.frame_buffers[key], np.frombuffer(data, dtype=np.uint8, count=self.frame_copy_size))
-    else:
-      for key, buf in bufs.items():
-        ptr = np.frombuffer(buf.data, dtype=np.uint8).ctypes.data
-        cache_key = (key, ptr)
-        if cache_key not in self._blob_cache:
-          self._blob_cache[cache_key] = Tensor.from_blob(ptr, (self.frame_buf_params[key][3],), dtype='uint8', device=self.WARP_DEV)
-        self.full_frames[key] = self._blob_cache[cache_key]
-
-  def reset_warmup_buffers(self):
-    if getattr(self, 'is_run_model', True):
-      self.input_queues, self.numpy_inputs, self.frame_buffers = make_stock_input_queues(
-        self.input_shapes, self.frame_skip, device=self.DEV, frame_copy_size=self.frame_copy_size)
-      self.frame_views = self.frame_buffers
-      self.npy = self.numpy_inputs
-    else:
-      for v in self.numpy_inputs.values():
-        v[:] = 0
-      self.full_frames.clear()
-      self._blob_cache.clear()
-
-  def run(self):
-    if self.run_model is not None:
-      outs, = self.run_model(**{k: self.input_queues[k] for k in MODELD_INPUTS})
-      return outs
-    else:
-      warped = self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=self.full_frames[self._road_key], big_frame=self.full_frames[self._wide_key])
-      raw_outputs = self.run_policy(**{k: self.input_queues[k] for k in POLICY_INPUTS if k in self.input_queues}, warped=warped)
-      return raw_outputs
-
-
-class NativeTinygradAdapter(BaseModelAdapter):
-  def __init__(self, *args, **kwargs):
-    super().__init__(*args, **kwargs)
-
-    self.input_specs = self.jits['input_specs'][(self.cam_w, self.cam_h)]
-    self.npy_shapes = self.jits['npy_shapes']
-    self.run_model = self.jits['run_model'][(self.cam_w, self.cam_h)]
-    self.vision_output_slices = self.jits['metadata']['model']['output_slices']
-
-    self._vision_input_names = ['img', 'big_img']
-
-    self.reset_warmup_buffers()
-    self._init_common()
-
-  def copy_frames(self, bufs):
-    for key, buf in bufs.items():
-      data = buf.data if hasattr(buf, 'data') else buf
-      if key in self.frame_views:
-        np.copyto(self.frame_views[key], np.frombuffer(data, dtype=np.uint8, count=self.frame_copy_size))
-
-  def reset_warmup_buffers(self) -> None:
-    buffers = {name: np.zeros(shape, dtype=dtype) for name, (shape, dtype, _) in self.input_specs.items()}
-    self.input_queues = {name: Tensor(buffers[name], device=device).realize() for name, (_, _, device) in self.input_specs.items()}
-
-    sizes = [int(np.prod(shape)) for shape in self.npy_shapes.values()]
-    packed = buffers['packed_npy_inputs']
-    npy_size = sum(sizes) * np.dtype(np.float32).itemsize
-
-    self.numpy_inputs = {name: v.reshape(shape) for (name, shape), v in
-                zip(self.npy_shapes.items(), np.split(packed[:npy_size].view(np.float32), np.cumsum(sizes[:-1])), strict=True)}
-
-    self.frame_copy_size = (packed.size - npy_size) // 2
-    self.frame_views = {'img': packed[npy_size:npy_size+self.frame_copy_size],
-                        'big_img': packed[npy_size+self.frame_copy_size:]}
-
-  def run(self):
-    outs, = self.run_model(**self.input_queues)
-    return outs
-
-
-def get_model_adapter(jits, cam_w, cam_h, model_device, queue_device, warp_device, chestnut=False):
-  if 'input_specs' in jits:
-    return NativeTinygradAdapter(jits, cam_w, cam_h, model_device, queue_device, warp_device, chestnut=chestnut)
-  else:
-    return LegacyModelAdapter(jits, cam_w, cam_h, model_device, queue_device, warp_device, chestnut=chestnut)

openpilot/sunnypilot/modeld_v2/modeld.py

@@ -38,15 +38,21 @@ from openpilot.selfdrive.controls.lib.desire_helper import DesireHelper
 from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_from_plan, smooth_value
 from openpilot.selfdrive.modeld.modeld import ChestnutState
 
+from openpilot.selfdrive.modeld.compile_modeld import (
+  MODELD_INPUTS,
+  make_input_queues as make_stock_input_queues,
+)
 from openpilot.sunnypilot.modeld_v2.fill_model_msg import fill_model_msg, fill_pose_msg, PublishState, get_curvature_from_output
 from openpilot.sunnypilot.modeld_v2.parse_model_outputs import Parser
 from openpilot.sunnypilot.modeld_v2.constants import ModelConstants, Plan
 from openpilot.sunnypilot.modeld_v2.meta_helper import load_meta_constants
 from openpilot.sunnypilot.modeld_v2.camera_offset_helper import CameraOffsetHelper
+from openpilot.sunnypilot.modeld_v2.compile_modeld import (derive_frame_skip, make_split_input_queues,
+                                                           make_supercombo_input_queues, nv12_copy_size,
+                                                           WARP_INPUTS, POLICY_INPUTS)
 from openpilot.sunnypilot.livedelay.helpers import get_lat_delay
 from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase
 from openpilot.sunnypilot.modeld_v2.helpers import load_oob
-from openpilot.sunnypilot.modeld_v2.model_adapters import get_model_adapter
 from openpilot.sunnypilot.models.helpers import get_active_bundle
 from openpilot.sunnypilot.selfdrive.controls.lib.relc import RoadEdgeLaneChangeController
 
@@ -117,19 +123,52 @@ class ModelState(ModelStateBase):
     self.WARP_DEV = metadata.get('warp_dev', 'QCOM') if COMMA_HARDWARE else 'CPU'
     self.DEV = ('AMD' if self.chestnut else 'QCOM') if COMMA_HARDWARE else 'CPU'
     self.QUEUE_DEV = self.DEV
-    self.adapter = get_model_adapter(jits, cam_w, cam_h, self.DEV, self.QUEUE_DEV, self.WARP_DEV, self.chestnut)
-    self.vision_output_slices = self.adapter.vision_output_slices
-    self.policy_output_slices = self.adapter.policy_output_slices
-    self._policy_slices_list = self.adapter._policy_slices_list
-    self._combined_model_type = self.adapter._combined_model_type
-    self._vision_input_names = self.adapter._vision_input_names
-    self.numpy_inputs = self.adapter.numpy_inputs
-    self._policy_keys = self.adapter._policy_keys
-    self._has_on_policy = self.adapter._has_on_policy
-    self._desire_key = self.adapter._desire_key
-    self._road_key = self.adapter._road_key
-    self._wide_key = self.adapter._wide_key
-    self.frame_buf_params = self.adapter.frame_buf_params
+    self.is_run_model = 'run_model' in jits
+
+    nv12_info = get_nv12_info(cam_w, cam_h)
+    self.frame_copy_size = nv12_copy_size(*nv12_info[:3])
+    self.full_frames: dict = {}
+    self._blob_cache: dict = {}
+    self.frame_buffers: dict = {}
+
+    if self.is_run_model or 'model' in metadata:
+      model_metadata = metadata.get('model', metadata)
+      self.input_shapes = model_metadata['input_shapes']
+      self.vision_output_slices = model_metadata['output_slices']
+      self.policy_output_slices = {}
+      self._policy_slices_list = []
+      self._combined_model_type = 'supercombo'
+      self._vision_input_names = [key for key in self.input_shapes if 'img' in key]
+      self.frame_skip = derive_frame_skip({}, self.input_shapes)
+      if self.is_run_model:
+        self.input_queues, self.numpy_inputs, self.frame_buffers = make_stock_input_queues(
+          self.input_shapes, self.frame_skip, device=self.DEV, frame_copy_size=self.frame_copy_size)
+        self.frame_views, self.npy = self.frame_buffers, self.numpy_inputs
+        self.run_model, self.run_policy, self.warp = jits['run_model'][(cam_w, cam_h)], None, None
+      else:
+        self.input_queues, self.numpy_inputs = make_supercombo_input_queues(self.input_shapes, self.frame_skip, device=self.QUEUE_DEV)
+        self.run_model, self.run_policy, self.warp = None, jits['run_policy'], jits[(cam_w, cam_h)]
+    else:
+      self.run_model, self.run_policy, self.warp = None, jits['run_policy'], jits[(cam_w, cam_h)]
+      vision_metadata = metadata['vision']
+      policy_keys = [k for k in metadata if k not in ('vision', 'warp_dev')]
+      self._combined_model_type = 'split' if policy_keys == ['policy'] else 'multi_policy'
+      self.vision_output_slices = vision_metadata['output_slices']
+      self._policy_keys = policy_keys
+      self._policy_slices_list = [metadata[k]['output_slices'] for k in policy_keys]
+      self.policy_output_slices = self._policy_slices_list[0]
+      self._has_on_policy = any('on' in k.lower() for k in policy_keys)
+      self._vision_input_names = [key for key in vision_metadata['input_shapes'] if 'img' in key]
+      first_policy_meta = metadata[policy_keys[0]]
+      frame_skip = derive_frame_skip(vision_metadata['input_shapes'], first_policy_meta['input_shapes'])
+      self.input_queues, self.numpy_inputs = make_split_input_queues(vision_metadata['input_shapes'],
+                                                                     first_policy_meta['input_shapes'],
+                                                                     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)
+    self._wide_key = next(key for key in self._vision_input_names if 'big' in key)
+    self.frame_buf_params = dict.fromkeys(self._vision_input_names, nv12_info)
 
     is_20hz = bundle.is20hz if bundle else self._combined_model_type in ('split', 'multi_policy')
     if is_20hz:
@@ -141,10 +180,26 @@ class ModelState(ModelStateBase):
     self.parser = Parser()
     self.prev_desire = np.zeros(self.constants.DESIRE_LEN, dtype=np.float32)
 
+    if self.warp is not None:
+      self.full_frames = {k: Tensor(np.zeros(nv12_info[3], dtype=np.uint8), device=self.WARP_DEV).contiguous().realize() for k in self._vision_input_names}
+      self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=self.full_frames[self._road_key], big_frame=self.full_frames[self._wide_key])
+
   def warmup(self) -> None:
-    dummy_frames, transforms, dummy_inputs = self.adapter.get_dummy_inputs()
+    dummy_size = self.frame_copy_size if self.is_run_model else self.frame_buf_params[self._road_key][3]
+    dummy_frames = {k: np.zeros(dummy_size, 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 = {k: np.zeros(v.shape, dtype=v.dtype) for k, v in self.numpy_inputs.items() if k not in ['tfm', 'big_tfm', 'prev_feat']}
     self.run(dummy_frames, transforms, dummy_inputs)
-    self.adapter.reset_warmup_buffers()
+    if self.is_run_model:
+      self.input_queues, self.numpy_inputs, self.frame_buffers = make_stock_input_queues(
+        self.input_shapes, self.frame_skip, device=self.DEV, frame_copy_size=self.frame_copy_size)
+      self.frame_views = self.frame_buffers
+      self.npy = self.numpy_inputs
+    else:
+      for v in self.numpy_inputs.values():
+        v[:] = 0
+      self.full_frames.clear()
+      self._blob_cache.clear()
     self.prev_desire[:] = 0
 
   @property
@@ -162,7 +217,17 @@ class ModelState(ModelStateBase):
   def run(self, bufs: dict[str, VisionBuf], transforms: dict[str, np.ndarray],
           inputs: dict[str, np.ndarray],
           after_enqueue: Callable[[], None] | None = None) -> dict[str, np.ndarray] | None:
-    self.adapter.copy_frames(bufs)
+    if self.is_run_model:
+      for key, buf in bufs.items():
+        data = buf.data if hasattr(buf, 'data') else buf
+        np.copyto(self.frame_buffers[key], np.frombuffer(data, dtype=np.uint8, count=self.frame_copy_size))
+    else:
+      for key, buf in bufs.items():
+        ptr = np.frombuffer(buf.data, dtype=np.uint8).ctypes.data
+        cache_key = (key, ptr)
+        if cache_key not in self._blob_cache:
+          self._blob_cache[cache_key] = Tensor.from_blob(ptr, (self.frame_buf_params[key][3],), dtype='uint8', device=self.WARP_DEV)
+        self.full_frames[key] = self._blob_cache[cache_key]
 
     desire_key = self.desire_key
     inputs[desire_key][0] = 0
@@ -176,7 +241,13 @@ class ModelState(ModelStateBase):
     self.numpy_inputs['tfm'][:, :] = transforms[self._road_key].reshape(3, 3)
     self.numpy_inputs['big_tfm'][:, :] = transforms[self._wide_key].reshape(3, 3)
 
-    raw_outputs = self.adapter.run()
+    if self.run_model is not None:
+      outs, = self.run_model(**{k: self.input_queues[k] for k in MODELD_INPUTS})
+      raw_outputs = outs
+    else:
+      assert self.warp is not None and self.run_policy is not None
+      warped = self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=self.full_frames[self._road_key], big_frame=self.full_frames[self._wide_key])
+      raw_outputs = self.run_policy(**{k: self.input_queues[k] for k in POLICY_INPUTS if k in self.input_queues}, warped=warped)
 
     if after_enqueue is not None:
       after_enqueue()

openpilot/sunnypilot/modeld_v2/stock_dependencies.py

@@ -1,167 +0,0 @@
-"""
-Copyright (c) 2021-, Haibin Wen, sunnypilot, and a number of other contributors.
-
-This file is part of sunnypilot and is licensed under the MIT License.
-See the LICENSE.md file in the root directory for more details.
-"""
-import math
-import numpy as np
-from collections import namedtuple
-from tinygrad.tensor import Tensor
-from tinygrad.helpers import Context
-from tinygrad.device import Device
-
-"""
-Frozen in time compile_modeld dependencies to support all models prior to transition to tinygrad compilation.
-This file is not meant to be modified.
-"""
-NV12Frame = namedtuple("NV12Frame", ['width', 'height', 'stride', 'y_height', 'uv_height', 'size'])
-MODELD_INPUTS = ['img_q', 'big_img_q', 'feat_q', 'desire_q', 'packed_npy_inputs']
-
-
-def nv12_copy_size(stride: int, y_height: int, uv_height: int) -> int:
-  return stride * (y_height + uv_height)
-
-
-def warp_perspective_tinygrad(src_flat, M_inv, dst_shape, src_shape, stride_pad, border_fill_val=None):
-  w_dst, h_dst = dst_shape
-  h_src, w_src = src_shape
-
-  x = Tensor.arange(w_dst).reshape(1, w_dst).expand(h_dst, w_dst).reshape(-1)
-  y = Tensor.arange(h_dst).reshape(h_dst, 1).expand(h_dst, w_dst).reshape(-1)
-
-  src_x = M_inv[0, 0] * x + M_inv[0, 1] * y + M_inv[0, 2]
-  src_y = M_inv[1, 0] * x + M_inv[1, 1] * y + M_inv[1, 2]
-  src_w = M_inv[2, 0] * x + M_inv[2, 1] * y + M_inv[2, 2]
-
-  src_x = src_x / src_w
-  src_y = src_y / src_w
-
-  x_round = Tensor.round(src_x)
-  y_round = Tensor.round(src_y)
-  x_nn_clipped = x_round.clip(0, w_src - 1).cast('int')
-  y_nn_clipped = y_round.clip(0, h_src - 1).cast('int')
-  idx = y_nn_clipped * (w_src + stride_pad) + x_nn_clipped
-  sampled = src_flat[idx]
-
-  if border_fill_val is None:
-    return sampled
-
-  in_bounds = ((x_round >= 0) & (x_round <= w_src - 1) &
-               (y_round >= 0) & (y_round <= h_src - 1)).cast(sampled.dtype)
-  return sampled * in_bounds + Tensor(border_fill_val, dtype=sampled.dtype) * (1 - in_bounds)
-
-
-def frames_to_tensor(frames):
-  H = (frames.shape[0] * 2) // 3
-  W = frames.shape[1]
-  in_img1 = Tensor.cat(frames[0:H:2, 0::2],
-                       frames[1:H:2, 0::2],
-                       frames[0:H:2, 1::2],
-                       frames[1:H:2, 1::2],
-                       frames[H:H+H//4].reshape((H//2, W//2)),
-                       frames[H+H//4:H+H//2].reshape((H//2, W//2)), dim=0).reshape((6, H//2, W//2))
-  return in_img1
-
-
-def make_frame_prepare(nv12: NV12Frame, model_w, model_h):
-  cam_w, cam_h, stride, y_height, uv_height, _ = nv12
-  uv_offset = stride * y_height
-  stride_pad = stride - cam_w
-
-  def frame_prepare_tinygrad(input_frame, M_inv):
-    M_inv_uv = M_inv * Tensor([[1.0, 1.0, 0.5], [1.0, 1.0, 0.5], [2.0, 2.0, 1.0]], device=Device.DEFAULT)
-    uv = input_frame[uv_offset:uv_offset + uv_height * stride].reshape(uv_height, stride)
-    with Context(SPLIT_REDUCEOP=0):
-      y = warp_perspective_tinygrad(input_frame[:cam_h*stride],
-                                    M_inv, (model_w, model_h),
-                                    (cam_h, cam_w), stride_pad).realize()
-      u = warp_perspective_tinygrad(uv[:cam_h//2, :cam_w:2].flatten(),
-                                    M_inv_uv, (model_w//2, model_h//2),
-                                    (cam_h//2, cam_w//2), 0).realize()
-      v = warp_perspective_tinygrad(uv[:cam_h//2, 1:cam_w:2].flatten(),
-                                    M_inv_uv, (model_w//2, model_h//2),
-                                    (cam_h//2, cam_w//2), 0).realize()
-    yuv = y.cat(u).cat(v).reshape((model_h * 3 // 2, model_w))
-    tensor = frames_to_tensor(yuv)
-    return tensor
-  return frame_prepare_tinygrad
-
-
-def get_policy_npy_shapes(input_shapes):
-  dp = input_shapes['desire_pulse']
-  tc = input_shapes['traffic_convention']
-  at = input_shapes['action_t']
-  fb = input_shapes['features_buffer']
-  feat_dim = math.prod(fb[2:])
-  shapes = {'desire': (dp[2],), 'traffic_convention': tuple(tc), 'action_t': tuple(at), 'prev_feat': (fb[0], feat_dim)}
-  return shapes, [math.prod(s) for s in shapes.values()]
-
-
-def make_input_queues(input_shapes, frame_skip, device, frame_copy_size):
-  img = input_shapes['img']
-  fb = input_shapes['features_buffer']
-  feat_dim = math.prod(fb[2:])
-  dp = input_shapes['desire_pulse']
-  n_frames = img[1] // 6
-  img_buf_shape = (frame_skip * (n_frames - 1) + 1, 6, img[2], img[3])
-
-  policy_shapes, _ = get_policy_npy_shapes(input_shapes)
-  shapes = {'tfm': (3, 3), 'big_tfm': (3, 3)} | policy_shapes
-  sizes = [math.prod(s) for s in shapes.values()]
-  packed_npy_size = sum(sizes) * np.dtype(np.float32).itemsize
-  packed_input = np.zeros(packed_npy_size + 2 * frame_copy_size, dtype=np.uint8)
-  packed_npy_inputs = packed_input[:packed_npy_size].view(np.float32)
-  frames = packed_input[packed_npy_size:]
-  frame_views = {'img': frames[:frame_copy_size], 'big_img': frames[frame_copy_size:]}
-  npy = {k: v.reshape(s) for (k, s), v in zip(shapes.items(), np.split(packed_npy_inputs, np.cumsum(sizes[:-1])), strict=True)}
-  input_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(),
-    'feat_q': Tensor(np.zeros((frame_skip * fb[1], fb[0], feat_dim), dtype=np.float32), device=device).contiguous().realize(),
-    'desire_q': Tensor(np.zeros((frame_skip * dp[1], dp[0], dp[2]), dtype=np.float32), device=device).contiguous().realize(),
-    'packed_npy_inputs': Tensor(packed_input, device='NPY').realize(),
-  }
-  return input_queues, npy, frame_views
-
-
-def shift_and_sample(buf, new_val, sample_fn):
-  buf.assign(buf[1:].cat(new_val, dim=0).contiguous())
-  return sample_fn(buf)
-
-
-def sample_skip(buf, frame_skip):
-  return buf[::frame_skip].contiguous().flatten(0, 1).unsqueeze(0)
-
-
-def sample_desire(buf, frame_skip):
-  return buf.reshape(-1, frame_skip, *buf.shape[1:]).max(1).flatten(0, 1).unsqueeze(0)
-
-
-def make_warp(nv12, model_w, model_h):
-  frame_prepare = make_frame_prepare(nv12, model_w, model_h)
-  def warp(tfm, big_tfm, frame, big_frame):
-    tfm = tfm.to(Device.DEFAULT)
-    big_tfm = big_tfm.to(Device.DEFAULT)
-    frame = frame.to(Device.DEFAULT)
-    big_frame = big_frame.to(Device.DEFAULT)
-    Tensor.realize(tfm, big_tfm, frame, big_frame)
-    warped_frame = frame_prepare(frame, tfm).unsqueeze(0)
-    warped_big_frame = frame_prepare(big_frame, big_tfm).unsqueeze(0)
-    return Tensor.cat(warped_frame, warped_big_frame)
-  return warp
-
-
-def make_run_model(warp, run_policy, model_metadata, frame_copy_size):
-  _, policy_sizes = get_policy_npy_shapes(model_metadata['input_shapes'])
-  packed_npy_size = (18 + sum(policy_sizes)) * np.dtype(np.float32).itemsize
-  def run_model(img_q, big_img_q, feat_q, desire_q, packed_npy_inputs):
-    packed_input = packed_npy_inputs.to(Device.DEFAULT)
-    Tensor.realize(packed_input)
-    packed_npy_inputs = packed_input[:packed_npy_size].bitcast('float32')
-    frame = packed_input[packed_npy_size:packed_npy_size + frame_copy_size]
-    big_frame = packed_input[packed_npy_size + frame_copy_size:]
-    tfm, big_tfm, policy_inputs = packed_npy_inputs.split([9, 9, sum(policy_sizes)])
-    warped = warp(tfm.reshape(3, 3), big_tfm.reshape(3, 3), frame, big_frame)
-    return run_policy(warped, img_q, big_img_q, feat_q, desire_q, policy_inputs)
-  return run_model