fix tests and lint

self = <buf real:False device:QCOM size:6504 dtype:dtypes.float offset:0>
      device = 'QCOM', size = 6504, dtype = dtypes.float, opaque = None
      options = None, initial_value = <pickle.PickleBuffer object at 0x7f21c7419d40>
      base = 1, offset = 0, preallocate = False

          def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None, options:BufferSpec|None=None,
                       initial_value:bytes|pickle.PickleBuffer|None=None, base:Buffer|None=None, offset:int=0, preallocate=False):
            assert isinstance(dtype, DType)
            self.device, self.size, self.dtype, self.offset, self.allocated_views, self._base = Device.canonicalize(device), size, dtype, offset, 0, base
            self.options = options if options is not None else BufferSpec()
            self._storage:BufferStorage|None = None
            if base is None:
              assert offset == 0, "base buffers can't have offset"
              if opaque is not None: self.allocate(opaque)
              if initial_value is not None:
                self.allocate()
                if (host:=self.get_storage().host) is not None: host[:] = memoryview(initial_value).cast('B')
                else:
                  src = Buffer("PYTHON", self.nbytes, dtypes.uint8, opaque=memoryview(bytearray(initial_value)))
                  for offset in range(0, self.nbytes, 1 << 20):
                    size = min(self.nbytes - offset, 1 << 20)
                    self.view(size, dtypes.uint8, offset).ensure_allocated().copy_from(src.view(size, dtypes.uint8, offset).ensure_allocated())
                if isinstance(initial_value, pickle.PickleBuffer): initial_value.release()
            else:
      >       assert base._base is None, "base can't have a base"
                     ^^^^^^^^^^
      E       AttributeError: 'int' object has no attribute '_base'

.github/workflows/compile-warp-kernels.yaml

@@ -0,0 +1,73 @@
+name: Compile warp for sunnypilot modeld
+
+on:
+  workflow_dispatch:
+  schedule:
+    - cron: '0 0 * * 0'
+
+jobs:
+  compile_warps:
+    runs-on: [self-hosted, chestnut]
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          submodules: recursive
+
+      - name: Set environment variables
+        run: |
+          source /etc/profile
+          export UV_PROJECT_ENVIRONMENT=${HOME}/venv
+          export UV_PYTHON_PREFERENCE=managed
+          export UV_PYTHON_INSTALL_DIR=${HOME}/uv/python
+          export VIRTUAL_ENV=$UV_PROJECT_ENVIRONMENT
+          uv sync --frozen
+          printenv >> $GITHUB_ENV
+
+      - name: Disable powersave
+        run: |
+          source ${UV_PROJECT_ENVIRONMENT}/bin/activate
+          PYTHONPATH=$PYTHONPATH:${{ github.workspace }}/ ${{ github.workspace }}/scripts/manage-powersave.py --disable
+
+      - name: Compile Warp Kernels
+        run: |
+          source ${UV_PROJECT_ENVIRONMENT}/bin/activate
+          export PYTHONPATH="${PYTHONPATH}:${{ github.workspace }}/tinygrad_repo:${{ github.workspace }}"
+          mkdir -p sunnypilot/modeld_v2/models/
+          MODEL_SIZE=$(python3 -c "from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE as s; print(f'{s[0]}x{s[1]}')")
+          CAMERA_RES=$(python3 -c "from openpilot.common.transformations.camera import _ar_ox_fisheye as a, _os_fisheye as o; print(f'{a.width}x{a.height} {o.width}x{o.height}')")
+          echo "Model size: $MODEL_SIZE"
+          echo "Camera resolutions: $CAMERA_RES"
+
+          for res in $CAMERA_RES; do
+            NV12_INFO=$(python3 -c "from openpilot.system.camerad.cameras.nv12_info import get_nv12_info; w, h = map(int, '${res}'.split('x')); print(','.join(map(str, get_nv12_info(w, h))))")
+
+            WARP_PKL="sunnypilot/modeld_v2/models/driving_warp_${res}_tinygrad.pkl"
+            echo "Compiling $WARP_PKL on QCOM"
+            taskset -c 7 env DEV=QCOM IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1 python3 tinygrad_repo/examples/openpilot/compile_warp.py \
+              --frame ${res},${NV12_INFO} --warp-to ${MODEL_SIZE} --layout yuv420 --frames 2 --output "${WARP_PKL}"
+
+            BIG_WARP_PKL="sunnypilot/modeld_v2/models/big_driving_warp_${res}_tinygrad.pkl"
+            echo "Compiling $BIG_WARP_PKL on AMD"
+            taskset -c 7 env DEBUG=1 DEV=USB+AMD:LLVM FRAME_DEV=CPU FLOAT16=1 JIT_BATCH_SIZE=0 GMMU=0 TC_OPT=2 TC_MIN_GLOBALS=32 python3 tinygrad_repo/examples/openpilot/compile_warp.py \
+              --frame ${res},${NV12_INFO} --warp-to ${MODEL_SIZE} --layout yuv420 --frames 2 --output "${BIG_WARP_PKL}"
+          done
+
+      - name: Re-enable powersave
+        if: always()
+        run: |
+          source ${UV_PROJECT_ENVIRONMENT}/bin/activate
+          PYTHONPATH=$PYTHONPATH:${{ github.workspace }}/ ${{ github.workspace }}/scripts/manage-powersave.py --enable
+
+      - name: Create Pull Request
+        uses: peter-evans/create-pull-request@9153d834b60caba6d51c9b9510b087acf9f33f83
+        with:
+          author: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
+          token: ${{ secrets.GITHUB_TOKEN }}
+          commit-message: "[bot] Update warp pkl for modeld_v2"
+          title: "[bot] Update modeld_v2 Warp"
+          branch: "auto/compile-warp-kernels"
+          base: "master"
+          delete-branch: true
+          labels: bot
+          add-paths: |
+            sunnypilot/modeld_v2/models/*.pkl

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

@@ -103,21 +103,23 @@ 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"
-            rm -f selfdrive/modeld/models/big_*.onnx selfdrive/modeld/models/dmonitoring_*.onnx
+            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
           else
-            git lfs pull -I "**/selfdrive/modeld/models/big_*.onnx" -X ""
-            find selfdrive/modeld/models -name "*.onnx" ! -name "big_*.onnx" -delete
+            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
           fi
-          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"
+          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"
             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
+          path: |
+            ${{ github.workspace }}/openpilot/openpilot/selfdrive/modeld/models/*.onnx
+            ${{ github.workspace }}/openpilot/openpilot/selfdrive/modeld/models/*.pkl
           if-no-files-found: error
 
   build_model:
@@ -196,64 +198,75 @@ jobs:
             OUTPUT_PKL="${{ env.MODELS_DIR }}/driving_tinygrad.pkl"
           fi
 
-          # 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
+          NATIVE_PKL=$(find "${{ env.MODELS_DIR }}" -maxdepth 1 -name "*.pkl" -print -quit)
 
-          # 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"
+          if [ -n "$NATIVE_PKL" ]; then
+            echo "Found native precompiled pkl: $NATIVE_PKL"
+            if [ "$NATIVE_PKL" != "$OUTPUT_PKL" ]; then
+              mv "$NATIVE_PKL" "$OUTPUT_PKL"
             fi
-          elif [ -f "$SUPERCOMBO_ONNX" ]; then
-            MODEL_TYPE=supercombo
-            ONNX_ARGS="--supercombo-onnx $SUPERCOMBO_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
 
-          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"
+            # 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
           fi
 
       - name: Prepare Output

.github/workflows/test_models.yml

@@ -48,6 +48,8 @@ jobs:
           submodules: true
 
       - name: Download Model Chunks in Parallel
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: |
           mkdir -p /tmp/model_chunks
           echo '${{ toJson(matrix.artifact.chunks) }}' > chunks.json
@@ -63,12 +65,15 @@ jobs:
           with open(manifest_path, "w") as f:
             f.write(str(len(chunks)))
           base_dir = os.environ["BASE_DIR"]
+          hf_token = os.environ.get("HF_TOKEN", "")
           with open("/tmp/curl_config.txt", "w") as f:
             for c in chunks:
               fn = c["file_name"]
               f.write(f"url = \"{base_dir}/{fn}\"\noutput = \"/tmp/model_chunks/{fn}\"\n")
+              if hf_token:
+                f.write(f"header = \"Authorization: Bearer {hf_token}\"\n")
           '
-          curl -Z --parallel-immediate --parallel-max 16 -s -S -f -L -K /tmp/curl_config.txt
+          curl -Z --parallel-immediate --parallel-max 12 --retry 3 --retry-all-errors -s -S -f -L -K /tmp/curl_config.txt
 
       - name: Run Model Compatibility Test
         env:

openpilot/selfdrive/modeld/SConscript

@@ -1,18 +1,20 @@
 import glob
-import json
 import os
+import shutil
+import tempfile
 import time
 from SCons.Script import Action, Value
-from openpilot.common.file_chunker import chunk_file, get_chunk_targets, get_existing_chunks
+from openpilot.common.file_chunker import chunk_file, get_chunk_targets, get_existing_chunks, open_file_chunked
 from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
 from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE, DM_INPUT_SIZE
-from openpilot.selfdrive.modeld.constants import ModelConstants
-from openpilot.selfdrive.modeld.helpers import TG_INPUT_DEVICES_PATH, chestnut_present, modeld_pkl_path
+from openpilot.selfdrive.modeld.helpers import chestnut_present
+from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
 
 
 Import('env', 'arch')
 chunker_file = File("#openpilot/common/file_chunker.py")
 lenv = env.Clone()
+lenv.PrependENVPath('PYTHONPATH', Dir('#tinygrad_repo').abspath)
 
 tinygrad_root = env.Dir("#").abspath
 tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + "/**", recursive=True, root_dir=tinygrad_root)
@@ -23,120 +25,89 @@ def estimate_pickle_max_size(onnx_size):
   # the ONNX size. Overestimating only adds an empty trailing chunk.
   return 2.0 * onnx_size + 10 * 1024 * 1024
 
-if arch == 'comma_arm64':
-  from openpilot.common.hardware import HARDWARE
-  camera = _os_fisheye if HARDWARE.get_device_type() == "mici" else _ar_ox_fisheye
-  camera_configs = [(camera.width, camera.height)]
-  tg_backend = 'QCOM'
-  tg_flags = f'DEV={tg_backend} IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1'
-else:
-  camera_configs = [(c.width, c.height) for c in (_ar_ox_fisheye, _os_fisheye)]
-  tg_backend = 'CPU'
-  tg_flags = f'DEV=CPU' if arch == 'Darwin' else 'DEV=CPU:LLVM'
+camera_configs = [(c.width, c.height) for c in (_ar_ox_fisheye, _os_fisheye)]
 
-tg_devices = { # which device to put jit inputs to at runtime
-  'openpilot.selfdrive.modeld.dmonitoringmodeld': {
-    'default': {'DEV': tg_backend}
-  },
-}
+if arch == 'comma_arm64':
+  tg_flags = 'DEV=QCOM IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1'
+else:
+  # JIT=2 disables graph batching, which produces incorrect outputs after buffers change.
+  tg_flags = 'DEV=METAL JIT=2' if arch == 'Darwin' else 'DEV=CPU:LLVM'
 
 CHESTNUT = chestnut_present()
 if CHESTNUT:
-  chestnut_tg_flags = 'DEBUG=1 DEV=USB+AMD:LLVM FRAME_DEV=CPU FLOAT16=1 JIT_BATCH_SIZE=0 GMMU=0 TC_OPT=2 TC_OCCUPANCY_OPT=1'
+  chestnut_tg_flags = 'DEBUG=1 DEV=USB+AMD:LLVM FRAME_DEV=CPU FLOAT16=1 JIT_BATCH_SIZE=0 GMMU=0 TC_OPT=2 TC_MIN_GLOBALS=32'
   # the USB+AMD GPU takes an exclusive flock; serialize all targets that touch it
   chestnut_lock = File("models/.chestnut.lock").abspath
 
-def write_tg_devices(target, source, env):
-  with open(str(target[0]), "w") as f:
-    json.dump(tg_devices, f)
-    f.write("\n")
-
-tg_devices_node = lenv.Command(
-  str(TG_INPUT_DEVICES_PATH),
-  [Value(tg_devices)],
-  write_tg_devices,
-)
-
 # tinygrad calls brew which needs a $HOME in the env
 mac_brew_string = f'HOME={os.path.expanduser("~")}' if arch == 'Darwin' else ''
 
-modeld_dir = Dir("#openpilot/selfdrive/modeld").abspath
-compile_modeld_script = [
-  File(f"{modeld_dir}/compile_modeld.py"),
-  File(f"{modeld_dir}/get_model_metadata.py"),
-  File("#openpilot/system/camerad/cameras/nv12_info.py"),
-  File("#openpilot/common/hardware/hw.py"),
-]
-model_w, model_h = MEDMODEL_INPUT_SIZE
-frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ
+warp_deps = [File("#openpilot/system/camerad/cameras/nv12_info.py")]
+compiler = Dir('#tinygrad_repo/examples/openpilot').abspath
+# CPU 7 is isolated with isolcpus on AGNOS, so explicitly pin the compiler to it.
+taskset = 'taskset -c 7 ' if arch == 'comma_arm64' else ''
 
-if not os.getenv('SKIP_TINYGRAD_COMPILE'):
-  for chestnut in [False, True] if CHESTNUT else [False]:
-    target_pkl_path = File(modeld_pkl_path(chestnut)).abspath
-    file_prefix, cmd_flags = ('big_', chestnut_tg_flags) if chestnut else ('', tg_flags)
-    driving_onnx_deps = get_existing_chunks(File(f"models/{file_prefix}driving_supercombo.onnx").abspath)
-    camera_res_args = ' '.join(f'{cw}x{ch}' for cw, ch in camera_configs)
-    # CPU 7 is isolated with isolcpus on AGNOS, so explicitly pin the compiler to it.
-    taskset = 'taskset -c 7 ' if arch == 'comma_arm64' else ''
-    cmd = (f'{cmd_flags} {mac_brew_string} {taskset}python3 {modeld_dir}/compile_modeld.py '
-          f'--model-size {model_w}x{model_h} '
-          f'--camera-resolutions {camera_res_args} '
-          f'--onnx {File(f"models/{file_prefix}driving_supercombo.onnx").abspath} '
-          f'--output {target_pkl_path} --frame-skip {frame_skip}')
-    onnx_sizes_sum = sum(os.path.getsize(f) for f in driving_onnx_deps)
-    chunk_targets = get_chunk_targets(target_pkl_path, estimate_pickle_max_size(onnx_sizes_sum))
-    def do_compile(target, source, env, command=cmd, pkl=target_pkl_path, chunks=chunk_targets):
-      from openpilot.system.hardware.chestnut.flash import link_up
-      # chestnut can enumerate before its PCIe link is up due to varying 12V power behavior across cars
-      for _ in range(10):
-        if link_up():
-          break
-        time.sleep(1)
-      else:
-        print("Chestnut not ready, skipping big model build")
-        return
-      if ret := env.Execute(command):
-        return ret
-      chunk_file(pkl, chunks)
-    def do_chunk(target, source, env, pkl=target_pkl_path, chunks=chunk_targets):
-      chunk_file(pkl, chunks)
-    actions = Action(do_compile, "  [CHESTNUT] $TARGET") if chestnut else [cmd, Action(do_chunk, "  [CHUNK] $TARGET")]
-    node = lenv.Command(
-      chunk_targets,
-      tinygrad_files + compile_modeld_script + driving_onnx_deps + [Value(camera_res_args), Value(chunk_targets), chunker_file],
-      actions,
-    )
+def chestnut_action(command):
+  def do_compile(target, source, env):
+    from openpilot.system.hardware.chestnut.flash import link_up
+    # chestnut can enumerate before its PCIe link is up due to varying 12V power behavior across cars
+    for _ in range(10):
+      if link_up():
+        break
+      time.sleep(1)
+    else:
+      print("Chestnut not ready, skipping warp build")
+      return
+    return env.Execute(command)
+  return Action(do_compile, "  [CHESTNUT] $TARGET")
+
+def compile_model(onnx_path, pkl_path, flags):
+  onnx_path, target_pkl_path = File(onnx_path).abspath, File(pkl_path).abspath
+  onnx_deps = get_existing_chunks(onnx_path)
+  cmd = (f'{flags} {mac_brew_string} {taskset}python3 "{compiler}/compile_onnx.py" '
+         f'"{{onnx}}" "{target_pkl_path}" --device-input "*" --out-of-band --benchmark-runs 1')
+  def do_compile(target, source, env):
+    if os.path.isfile(onnx_path):
+      return env.Execute(cmd.format(onnx=onnx_path))
+    # TODO: Remove ONNX chunk reassembly once models are precompiled.
+    with tempfile.NamedTemporaryFile(dir=os.path.dirname(onnx_path), suffix='.onnx') as tmp, open_file_chunked(onnx_path) as src:
+      shutil.copyfileobj(src, tmp)
+      tmp.flush()
+      return env.Execute(cmd.format(onnx=tmp.name))
+  compile_action = Action(do_compile, "  [ONNX] $TARGET")
+  onnx_sizes_sum = sum(os.path.getsize(f) for f in onnx_deps)
+  chunk_targets = get_chunk_targets(target_pkl_path, estimate_pickle_max_size(onnx_sizes_sum))
+  def do_chunk(target, source, env, pkl=target_pkl_path, chunks=chunk_targets):
+    chunk_file(pkl, chunks)
+  lenv.Command(
+    chunk_targets,
+    tinygrad_files + onnx_deps + [Value(cmd), Value(chunk_targets), chunker_file],
+    [compile_action, Action(do_chunk, "  [CHUNK] $TARGET")],
+  )
+
+compile_model('models/dmonitoring_model.onnx', 'models/dmonitoring_model_tinygrad.pkl', tg_flags)
+compile_model('models/driving_supercombo.onnx', 'models/driving_tinygrad.pkl', tg_flags)
+
+model_w, model_h = MEDMODEL_INPUT_SIZE
+for chestnut in [False, True] if CHESTNUT else [False]:
+  file_prefix, cmd_flags = ('big_', chestnut_tg_flags) if chestnut else ('', tg_flags)
+  for cam_w, cam_h in camera_configs:
+    warp_pkl_path = File(f"models/{file_prefix}driving_warp_{cam_w}x{cam_h}_tinygrad.pkl").abspath
+    stride, y_height, uv_height, _ = get_nv12_info(cam_w, cam_h)
+    cmd = (f'{cmd_flags} {mac_brew_string} {taskset}python3 "{compiler}/compile_warp.py" '
+           f'--frame {cam_w},{cam_h},{stride},{y_height},{uv_height},{stride * (y_height + uv_height)} '
+           f'--warp-to {model_w}x{model_h} --layout yuv420 --frames 2 '
+           f'--output {warp_pkl_path}')
+    action = chestnut_action(cmd) if chestnut else cmd
+    node = lenv.Command(warp_pkl_path, tinygrad_files + warp_deps + [Value(cmd)], action)
     if chestnut:
       lenv.SideEffect(chestnut_lock, node)
 
-# get model metadata
-fn = File(f"models/dmonitoring_model").abspath
-script_files = [File(Dir("#openpilot/selfdrive/modeld").File("get_model_metadata.py").abspath)]
-cmd = f'{tg_flags} {mac_brew_string} python3 {Dir("#openpilot/selfdrive/modeld").abspath}/get_model_metadata.py {fn}.onnx'
-lenv.Command(fn + "_metadata.pkl", [fn + ".onnx"] + tinygrad_files + script_files + [tg_devices_node], cmd)
-
 dm_w, dm_h = DM_INPUT_SIZE
-compile_dm_warp_script = [File(f"{modeld_dir}/compile_dm_warp.py")]
 for cam_w, cam_h in camera_configs:
   dm_pkl_path = File(f"models/dm_warp_{cam_w}x{cam_h}_tinygrad.pkl").abspath
-  cmd = (f'{tg_flags} {mac_brew_string} python3 {modeld_dir}/compile_dm_warp.py '
-         f'--camera-resolution {cam_w}x{cam_h} --warp-to {dm_w}x{dm_h} '
-         f'--output {dm_pkl_path}')
-  lenv.Command(dm_pkl_path, tinygrad_files + compile_dm_warp_script + compile_modeld_script + [tg_devices_node], cmd)
-
-def tg_compile(flags, model_name):
-  pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"'
-  fn = File(f"models/{model_name}").abspath
-  pkl = fn + "_tinygrad.pkl"
-  onnx_path = fn + ".onnx"
-  chunk_targets = get_chunk_targets(pkl, estimate_pickle_max_size(os.path.getsize(onnx_path)))
-  def do_chunk(target, source, env):
-    chunk_file(pkl, chunk_targets)
-  return lenv.Command(
-    chunk_targets,
-    [onnx_path] + tinygrad_files + [Value(chunk_targets), chunker_file, tg_devices_node],
-    [f'{pythonpath_string} {flags} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {pkl}',
-     Action(do_chunk, "  [CHUNK] $TARGET")],
-  )
-
-tg_compile(tg_flags, 'dmonitoring_model')
+  stride, y_height, uv_height, frame_size = get_nv12_info(cam_w, cam_h)
+  cmd = (f'{tg_flags} {mac_brew_string} python3 "{compiler}/compile_warp.py" '
+         f'--frame {cam_w},{cam_h},{stride},{y_height},{uv_height},{frame_size} --warp-to {dm_w}x{dm_h} '
+         f'--layout luma --border-fill 16 --transform-device NPY --output {dm_pkl_path}')
+  lenv.Command(dm_pkl_path, tinygrad_files + warp_deps + [Value(cmd)], cmd)

openpilot/selfdrive/modeld/modeld.py

@@ -1,12 +1,17 @@
 #!/usr/bin/env python3
 from collections.abc import Callable
+import base64
 import ctypes
 from functools import cached_property
 import os
 os.environ['GMMU'] = '0' # for chestnut fast loading, noop for qcom
-from tinygrad.device import Device
-import usb1
-import struct
+from tinygrad.device import Buffer, Device
+from tinygrad.dtype import DType, dtypes
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import round_up
+from tinygrad.uop.ops import UOp
+import math
+import pickle
 import threading
 import time
 import numpy as np
@@ -28,12 +33,10 @@ 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, 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, nv12_copy_size, MODELD_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
-from openpilot.common.hardware.usb import CHESTNUT_USB_IDS
 from openpilot.selfdrive.modeld.constants import ModelConstants, Plan
-from openpilot.selfdrive.modeld.helpers import chestnut_present, chestnut_compiled, chestnut_ready, modeld_pkl_path, load_oob
+from openpilot.selfdrive.modeld.helpers import  MODELS_DIR, chestnut_present, chestnut_compiled, modeld_pkl_path, load_oob
 
 from openpilot.sunnypilot.livedelay.helpers import get_lat_delay
 from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase
@@ -75,45 +78,14 @@ def get_action_from_model(model_output: dict[str, np.ndarray], prev_action: log.
                                 shouldStop=bool(stop))
 
 
-class ChestnutState:
-  # only modeld can access chestnut
+class ChestnutGpuState:
+  # GPU metrics require modeld's GPU context
   def __init__(self, pm: PubMaster, big: bool):
     self.pm = pm
     self.big = big
     self.valid = True
     self.sends = 0
     self.metrics = {}
-    self._asm_usb = None
-
-  def _close_asm_usb(self) -> None:
-    if self._asm_usb is not None:
-      self._asm_usb.close()
-      self._asm_usb = None
-
-  def _open_asm_usb(self):
-    context = usb1.USBContext()
-    for vendor_id, product_id in CHESTNUT_USB_IDS:
-      if (handle := context.openByVendorIDAndProductID(vendor_id, product_id, skip_on_error=True)) is not None:
-        return handle
-    context.close()
-
-  def _read_ina(self) -> tuple[int, int, bool]:
-    if "AMD" in Device._opened_devices and self._asm_usb is None:
-      try:
-        raw = Device["AMD"].iface.pci_dev.usb.usb.control_read(0xC0, 5)
-        return struct.unpack('<Hh?', bytes(raw))
-      except Exception:
-        pass
-    if self._asm_usb is None:
-      self._asm_usb = self._open_asm_usb()
-    if self._asm_usb is None:
-      raise usb1.USBErrorNoDevice
-    try:
-      raw = self._asm_usb.controlRead(0xC0, 0xC0, 0, 0, 5, timeout=100)
-    except usb1.USBError:
-      self._close_asm_usb()
-      raise
-    return struct.unpack('<Hh?', bytes(raw))
 
   @cached_property
   def power_limit(self) -> int:
@@ -121,8 +93,8 @@ class ChestnutState:
     return smu._send_msg(smu.smu_mod.PPSMC_MSG_GetPptLimit, 0, read_back_arg=True, timeout=100)
 
   def send(self) -> None:
-    msg = messaging.new_message('chestnutState')
-    state = msg.chestnutState
+    msg = messaging.new_message('chestnutGpuState')
+    state = msg.chestnutGpuState
     self.sends += 1
     if self.big and "AMD" in Device._opened_devices and self.sends % 100 == 1:
       try:
@@ -148,21 +120,8 @@ class ChestnutState:
       for k, v in self.metrics.items():
         setattr(state, k, v)
 
-    asm_valid = False
-    try:
-      # ASM runs on USB-C power, these still read without a gpu
-      state.supplyVoltage, state.supplyCurrent, state.supplyFault = self._read_ina()
-      asm_valid = True
-    except Exception:
-      pass
-    if "AMD" in Device._opened_devices:
-      try:
-        state.pcieLtssm = Device["AMD"].iface.pci_dev.usb.read(0xB450, 1)[0]
-      except Exception:
-        pass
-
-    msg.valid = asm_valid and (not self.big or self.valid)
-    self.pm.send('chestnutState', msg)
+    msg.valid = not self.big or (self.valid and bool(self.metrics))
+    self.pm.send('chestnutGpuState', msg)
 
 
 class FrameMeta:
@@ -175,37 +134,64 @@ class FrameMeta:
       self.frame_id, self.timestamp_sof, self.timestamp_eof = vipc.frame_id, vipc.timestamp_sof, vipc.timestamp_eof
 
 
+def input_view(buffer: Buffer, shape: tuple[int, ...], dtype: DType, offset: int) -> Tensor:
+  view = buffer.view(math.prod(shape), dtype, offset).ensure_allocated()
+  return Tensor(UOp.from_buffer(view)).reshape(shape)
+
+
 class ModelState(ModelStateBase):
   prev_desire: np.ndarray  # for tracking the rising edge of the pulse
 
   def __init__(self, cam_w: int, cam_h: int, chestnut: bool):
-    ModelStateBase.__init__(self)
     jits = load_oob(open_file_chunked(modeld_pkl_path(chestnut)))
-    input_devices = jits['input_devices']
-    self.model_device = input_devices['model']
-    metadata = jits['metadata']
-    self.input_shapes = metadata['input_shapes']
-    self.vision_input_names = [k for k in self.input_shapes if 'img' in k]
-    self.output_slices = metadata['output_slices']
+    self.model_device = jits['input_specs']['new_img'][2]
+    self.input_shapes = {name: (shape, np.dtype(dtype)) for name, (shape, dtype, _) in jits['input_specs'].items()}
+    self.state_pairs = {name: f'next_{name}' for name in self.input_shapes if f'next_{name}' in jits['metadata']['output_shapes']}
+    self.vision_input_names = ('img', 'big_img')
+    self.output_slices = pickle.loads(base64.b64decode(jits['metadata']['metadata']['output_slices']))
 
     self.prev_desire = np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32)
     self.chestnut = chestnut
 
-    self.frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ
-    self.frame_copy_size = nv12_copy_size(*get_nv12_info(cam_w, cam_h)[:3])
-    self.input_queues, self.npy, self.frame_views = make_input_queues(
-      self.input_shapes, self.frame_skip, device=self.model_device, frame_copy_size=self.frame_copy_size)
+    stride, y_height, uv_height, _ = get_nv12_info(cam_w, cam_h)
+    self.frame_copy_size = stride * (y_height + uv_height)
+    self.pack_inputs()
+    with open(MODELS_DIR / f'{"big_" if chestnut else ""}driving_warp_{cam_w}x{cam_h}_tinygrad.pkl', 'rb') as f:
+      self.run_warp = pickle.load(f)['run']
+    self.run_model = jits['run']
+    self.outputs = {name: Tensor(np.zeros(shape, dtype=dtype), device=device).realize() for name, (shape, dtype, device) in jits['output_specs'].items()}
+    for name, next_name in self.state_pairs.items():
+      state = self.input_queues[name]
+      self.outputs[next_name] = input_view(state._buffer(), state.shape, state.dtype, 0)
     self.parser = Parser()
-    self.run_model = jits['run_model'][(cam_w,cam_h)]
+
+  def pack_inputs(self) -> None:
+    # Pack host inputs into one upload to reduce USB transfer overhead for the eGPU.
+    self.input_queues = {name: Tensor(np.zeros(shape, dtype=dtype), device=self.model_device).realize()
+                         for name, (shape, dtype) in self.input_shapes.items() if name in self.state_pairs}
+    shapes = {'tfm': (2, 3, 3)} | {name: shape for name, (shape, _) in self.input_shapes.items()
+                                   if name not in self.state_pairs and name != 'new_img'}
+    npy_size = sum(round_up(math.prod(shape) * 4, 128) for shape in shapes.values())
+    self.packed_input = np.zeros(npy_size + 2 * self.frame_copy_size, dtype=np.uint8)
+    self.input_host = Tensor(self.packed_input, device='NPY')._buffer()
+    self.input_device = Tensor(self.packed_input, device=self.model_device)._buffer()
+    self.npy = {}
+    offset = 0
+    for name, shape in shapes.items():
+      self.npy[name] = np.ndarray(shape, dtype=np.float32, buffer=self.packed_input, offset=offset)
+      self.input_queues[name] = input_view(self.input_device, shape, dtypes.float32, offset)
+      offset += round_up(self.npy[name].nbytes, 128)
+    self.frames = self.packed_input[npy_size:].reshape(2, self.frame_copy_size)
+    self.warp_inputs = {'input_frame': input_view(self.input_device, self.frames.shape, dtypes.uint8, npy_size), 'M_inv': self.input_queues.pop('tfm')}
 
   def slice_outputs(self, model_outputs: np.ndarray, output_slices: dict[str, slice]) -> dict[str, np.ndarray]:
-    parsed_model_outputs = {k: model_outputs[np.newaxis, v] for k,v in output_slices.items()}
-    return parsed_model_outputs
+    return {k: model_outputs[np.newaxis, v] for k,v in output_slices.items()}
 
   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]:
-    for key, buf in bufs.items():
-      np.copyto(self.frame_views[key], np.frombuffer(buf.data, dtype=np.uint8, count=self.frame_copy_size))
+    for i, key in enumerate(self.vision_input_names):
+      np.copyto(self.frames[i], np.frombuffer(bufs[key].data, dtype=np.uint8, count=self.frame_copy_size))
+      self.npy['tfm'][i] = transforms[key]
 
     # Model decides when action is completed, so desire input is just a pulse triggered on rising edge
     inputs['desire_pulse'][0] = 0
@@ -213,17 +199,16 @@ class ModelState(ModelStateBase):
     self.prev_desire[:] = inputs['desire_pulse']
     self.npy['traffic_convention'][:] = inputs['traffic_convention']
     self.npy['action_t'][:] = inputs['action_t']
-    self.npy['tfm'][:,:] = transforms['img'][:,:]
-    self.npy['big_tfm'][:,:] = transforms['big_img'][:,:]
 
-    outs, = self.run_model(**{k: self.input_queues[k] for k in MODELD_INPUTS})
+    self.input_device.copy_from(self.input_host)
+    self.input_queues['new_img'] = self.run_warp(**self.warp_inputs)
+    self.run_model(output_buffers=self.outputs, **self.input_queues)
     if after_enqueue is not None:
       after_enqueue()
-    model_output = outs.numpy()[0]
+    model_output = self.outputs['outputs'].numpy()[0]
     if self.chestnut and not np.all(np.isfinite(model_output)):
       raise RuntimeError("model output not finite")
     outputs_dict = self.parser.parse_outputs(self.slice_outputs(model_output, self.output_slices))
-    self.npy['prev_feat'][:] = model_output[self.output_slices['hidden_state']]
 
     if SEND_RAW_PRED:
       outputs_dict['raw_pred'] = model_output.copy()
@@ -234,33 +219,21 @@ class ModelState(ModelStateBase):
     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, self.frame_views = make_input_queues(
-      self.input_shapes, self.frame_skip, device=self.model_device, frame_copy_size=self.frame_copy_size)
+    self.packed_input[:] = 0
+    for key in self.state_pairs:
+      self.input_queues[key].assign(0).realize()
     self.prev_desire[:] = 0
 
 
 def main(demo=False):
   cloudlog.warning("modeld init")
 
-  chestnut_available = chestnut_present() and chestnut_compiled()
-  CHESTNUT = False
-  if chestnut_available:
-    poller = messaging.Poller()
-    sock = messaging.sub_sock("chestnutState", poller=poller, conflate=True)
-    deadline = time.monotonic() + 4. / SERVICE_LIST['deviceState'].frequency
-    while not CHESTNUT and (remaining := deadline - time.monotonic()) > 0.:
-      if not poller.poll(round(remaining * 1000)):
-        break
-      msg = messaging.recv_one_or_none(sock)
-      CHESTNUT = msg is not None and msg.valid and chestnut_ready(msg.chestnutState)
+  CHESTNUT = chestnut_present() and chestnut_compiled()
   if CHESTNUT:
     os.environ['HCQDEV_WAIT_TIMEOUT_MS'] = '3000'
   params = Params()
   params.put_bool("ChestnutLoading", CHESTNUT)
-  if chestnut_available and not CHESTNUT:
-    params.put_bool("ChestnutActive", False)
-  else:
-    params.remove("ChestnutActive")
+  params.remove("ChestnutActive")
 
   config_realtime_process(7, 54)
 
@@ -304,11 +277,7 @@ def main(demo=False):
     loader.start()
     loader.join(BIG_MODEL_TIMEOUT)
     model = big_model
-    if model is None:
-      params.put_bool("ChestnutModelError", True)
     params.put_bool("ChestnutActive", model is not None)
-    if model is not None:
-      params.remove("ChestnutModelError")
 
   small_model = ModelState(vipc_client_main.width, vipc_client_main.height, False) if model is None or CHESTNUT else None
   if model is None:
@@ -318,13 +287,13 @@ def main(demo=False):
   cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting")
 
   # messaging
-  pub_socks = ["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"] + (["chestnutState"] if CHESTNUT else [])
+  pub_socks = ["modelV2", "drivingModelData", "cameraOdometry"] + (["chestnutGpuState"] if CHESTNUT else [])
   pm = PubMaster(pub_socks)
   sm = SubMaster(["deviceState", "carState", "narrowRoadCameraState", "extrinsicsCalibration", "driverMonitoringState", "carControl", "lateralDelay"])
 
   publish_state = PublishState()
   params = Params()
-  chestnut_state = ChestnutState(pm, model.chestnut) if CHESTNUT else None
+  chestnut_state = ChestnutGpuState(pm, model.chestnut) if CHESTNUT else None
 
   # setup filter to track dropped frames
   frame_dropped_filter = FirstOrderFilter(0., 10., 1. / ModelConstants.MODEL_RUN_FREQ)
@@ -435,14 +404,13 @@ def main(demo=False):
     mt1 = time.perf_counter()
     try:
       send_chestnut = (chestnut_state is not None and
-                       run_count % round(ModelConstants.MODEL_RUN_FREQ / SERVICE_LIST['chestnutState'].frequency) == 0)
+                       run_count % round(ModelConstants.MODEL_RUN_FREQ / SERVICE_LIST['chestnutGpuState'].frequency) == 0)
       model_output = model.run(bufs, transforms, inputs, chestnut_state.send if send_chestnut else None)
     except Exception:
       if not params.get_bool("ChestnutActive"):
         raise
       # fallback to small model
       cloudlog.exception("big model failed, fall back to small")
-      params.put_bool("ChestnutModelError", True)
       params.put_bool("ChestnutActive", False)
       assert small_model is not None
       model = small_model

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.selfdrive.modeld.compile_modeld as stock
+import openpilot.sunnypilot.modeld_v2.stock_dependencies as stock
 from tinygrad import dtypes
 from tinygrad.device import Device
 from tinygrad.engine.jit import TinyJit

openpilot/sunnypilot/modeld_v2/helpers.py

@@ -55,6 +55,11 @@ def _dynamic_factory(real_class):
     return _enum_factory(real_class)
 
   def factory(*args, **kwargs):
+    if real_class.__name__ == 'Buffer':
+      # Tinygrad, before compile_modeld moved to tg, serialized lb_refcount at index 6. It has been removed.
+      if len(args) >= 7 and isinstance(args[6], int):
+        args = tuple(list(args[:6]) + list(args[7:]))
+
     try:
       return real_class(*args, **kwargs)
     except TypeError:

openpilot/sunnypilot/modeld_v2/model_adapters.py

@@ -0,0 +1,224 @@
+"""
+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 codecs
+import pickle
+import numpy as np
+from tinygrad.tensor import Tensor
+from tinygrad.dtype import DType, dtypes
+from tinygrad.device import Buffer
+from tinygrad.uop.ops import UOp
+from tinygrad.helpers import round_up
+import math
+from openpilot.common.basedir import BASEDIR
+from pathlib import Path
+from openpilot.sunnypilot.modeld_v2.helpers import DynamicTinygradUnpickler
+
+def input_view(buffer: Buffer, shape: tuple[int, ...], dtype: DType, offset: int) -> Tensor:
+  view = buffer.view(math.prod(shape), dtype, offset).ensure_allocated()
+  return Tensor(UOp.from_buffer(view)).reshape(shape)
+
+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)
+    self.is_native = False
+
+  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(self.frame_views[k].size, dtype=np.uint8) if self.is_native else 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:
+      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)
+      return raw_outputs
+
+
+class NativeTinygradAdapter(BaseModelAdapter):
+  def __init__(self, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+    self.is_native = True
+    self.input_specs = self.jits['input_specs']
+    self.packed_specs = self.jits['packed_specs']
+
+    self.model_device = self.input_specs['new_img'][2]
+    self.input_shapes = {name: (shape, np.dtype(dtype)) for name, (shape, dtype, _) in self.input_specs.items()}
+    self.state_pairs = {name: f'next_{name}' for name in self.input_shapes if f'next_{name}' in self.jits['metadata']['output_shapes']}
+
+    stride, y_height, uv_height, _ = get_nv12_info(self.cam_w, self.cam_h)
+    self.frame_copy_size = stride * (y_height + uv_height)
+
+    self.input_shapes_orig = self.jits['metadata']['input_shapes']
+    self._vision_input_names = [k for k in self.input_shapes_orig if 'img' in k]
+    self.vision_output_slices = pickle.loads(codecs.decode(self.jits['metadata']['metadata']['output_slices'].encode(), 'base64'))
+
+    self.reset_warmup_buffers()
+    self._init_common()
+
+    warp_dir = Path(BASEDIR) / "openpilot/sunnypilot/modeld_v2/models"
+    with open(warp_dir / f'{"big_" if self.chestnut else ""}driving_warp_{self.cam_w}x{self.cam_h}_tinygrad.pkl', 'rb') as f:
+      self.run_warp = DynamicTinygradUnpickler(f).load()['run']
+    self.run_model = self.jits['run']
+
+    self.outputs = {name: Tensor(np.zeros(shape, dtype=dtype), device=device).realize() for name, (shape, dtype, device) in self.jits['output_specs'].items()}
+    for name, next_name in self.state_pairs.items():
+      state = self.input_queues[name]
+      self.outputs[next_name] = input_view(state._buffer(), state.shape, state.dtype, 0)
+
+  def copy_frames(self, bufs):
+    for i, key in enumerate(self._vision_input_names):
+      if key in bufs:
+        data = bufs[key].data if hasattr(bufs[key], 'data') else bufs[key]
+        np.copyto(self.frames[i], np.frombuffer(data, dtype=np.uint8, count=self.frame_copy_size))
+
+  def reset_warmup_buffers(self) -> None:
+    self.input_queues = {name: Tensor(np.zeros(shape, dtype=dtype), device=self.model_device).realize()
+                         for name, (shape, dtype) in self.input_shapes.items() if name in self.state_pairs}
+    shapes = {'tfm': (2, 3, 3)} | {name: shape for name, (shape, _) in self.input_shapes.items()
+                                   if name not in self.state_pairs and name != 'new_img'}
+    npy_size = sum(round_up(math.prod(shape) * 4, 128) for shape in shapes.values())
+    self.packed_input = np.zeros(npy_size + 2 * self.frame_copy_size, dtype=np.uint8)
+    self.input_host = Tensor(self.packed_input, device='NPY')._buffer()
+    self.input_device = Tensor(self.packed_input, device=self.model_device)._buffer()
+    self.numpy_inputs = {}
+    offset = 0
+    for name, shape in shapes.items():
+      self.numpy_inputs[name] = np.ndarray(shape, dtype=np.float32, buffer=self.packed_input, offset=offset)
+      self.input_queues[name] = input_view(self.input_device, shape, dtypes.float32, offset)
+      offset += round_up(self.numpy_inputs[name].nbytes, 128)
+    self.frames = self.packed_input[npy_size:].reshape(2, self.frame_copy_size)
+    self.warp_inputs = {'input_frame': input_view(self.input_device, self.frames.shape, dtypes.uint8, npy_size), 'M_inv': self.input_queues.pop('tfm')}
+
+    # Split tfm into tfm and big_tfm for modeld compatibility
+    if 'tfm' in self.numpy_inputs and self.numpy_inputs['tfm'].shape == (2, 3, 3):
+      real_tfm = self.numpy_inputs.pop('tfm')
+      self.numpy_inputs['tfm'] = real_tfm[0]
+      self.numpy_inputs['big_tfm'] = real_tfm[1]
+
+  def run(self):
+    self.input_device.copy_from(self.input_host)
+    self.input_queues['new_img'] = self.run_warp(**self.warp_inputs)
+    self.run_model(output_buffers=self.outputs, **self.input_queues)
+    return self.outputs['outputs']
+
+
+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

@@ -13,7 +13,6 @@ import numpy as np
 import threading
 import time
 from setproctitle import setproctitle
-from tinygrad.tensor import Tensor
 
 import openpilot.cereal.messaging as messaging
 from openpilot.common.hardware import COMMA_HARDWARE
@@ -33,26 +32,19 @@ from openpilot.common.realtime import config_realtime_process, DT_MDL
 from openpilot.common.transformations.camera import DEVICE_CAMERAS
 from openpilot.common.transformations.model import get_warp_matrix
 from openpilot.system import sentry
-from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
 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.modeld import ChestnutGpuState
 
-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
 
@@ -123,52 +115,19 @@ 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.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)
+    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
 
     is_20hz = bundle.is20hz if bundle else self._combined_model_type in ('split', 'multi_policy')
     if is_20hz:
@@ -180,26 +139,10 @@ 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_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']}
+    dummy_frames, transforms, dummy_inputs = self.adapter.get_dummy_inputs()
     self.run(dummy_frames, transforms, dummy_inputs)
-    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.adapter.reset_warmup_buffers()
     self.prev_desire[:] = 0
 
   @property
@@ -217,21 +160,12 @@ 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:
-    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]
+    self.adapter.copy_frames(bufs)
 
     desire_key = self.desire_key
     inputs[desire_key][0] = 0
-    self.numpy_inputs[desire_key][:] = np.where(inputs[desire_key] - self.prev_desire > .99, inputs[desire_key], 0)
+    (self.numpy_inputs[desire_key].flat if self.adapter.is_native else self.numpy_inputs[desire_key])[:] = \
+      np.where(inputs[desire_key] - self.prev_desire > .99, inputs[desire_key], 0)
     self.prev_desire[:] = inputs[desire_key]
 
     for key in ('traffic_convention', 'lateral_control_params', 'action_t'):
@@ -241,13 +175,7 @@ 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)
 
-    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)
+    raw_outputs = self.adapter.run()
 
     if after_enqueue is not None:
       after_enqueue()
@@ -259,14 +187,16 @@ class ModelState(ModelStateBase):
       sliced = {k: model_output[np.newaxis, v] for k, v in self.vision_output_slices.items()}
       outputs = self.parser.parse_outputs(sliced)
       if 'prev_feat' in self.numpy_inputs and 'hidden_state' in self.vision_output_slices:
-        self.numpy_inputs['prev_feat'][:] = model_output[self.vision_output_slices['hidden_state']]
+        (self.numpy_inputs['prev_feat'].flat if self.adapter.is_native else self.numpy_inputs['prev_feat'])[:] = \
+          model_output[self.vision_output_slices['hidden_state']]
     else:
       vision_output = raw_outputs[0].numpy().flatten()
       vision_sliced = {k: vision_output[np.newaxis, v] for k, v in self.vision_output_slices.items()}
       outputs = self.parser.parse_vision_outputs(vision_sliced)
 
       if 'prev_feat' in self.numpy_inputs and 'hidden_state' in self.vision_output_slices:
-        self.numpy_inputs['prev_feat'][:] = vision_output[self.vision_output_slices['hidden_state']]
+        (self.numpy_inputs['prev_feat'].flat if self.adapter.is_native else self.numpy_inputs['prev_feat'])[:] = \
+          vision_output[self.vision_output_slices['hidden_state']]
 
       for i, policy_slices in enumerate(self._policy_slices_list):
         policy_output = raw_outputs[i + 1].numpy().flatten()
@@ -373,11 +303,7 @@ def main(demo=False):
     loader.start()
     loader.join(BIG_MODEL_TIMEOUT)
     model = big_model
-    if model is None:
-      params.put_bool("ChestnutModelError", True)
     params.put_bool("ChestnutActive", model is not None)
-    if model is not None:
-      params.remove("ChestnutModelError")
 
   small_model = ModelState(cam_w=vipc_client_main.width, cam_h=vipc_client_main.height, chestnut=False) if model is None or CHESTNUT else None
   if model is None:
@@ -387,12 +313,12 @@ def main(demo=False):
   cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting")
 
   # messaging
-  pub_socks = ["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"] + (["chestnutState"] if CHESTNUT else [])
+  pub_socks = ["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"] + (["chestnutGpuState"] if CHESTNUT else [])
   pm = PubMaster(pub_socks)
   sm = SubMaster(["deviceState", "carState", "narrowRoadCameraState", "extrinsicsCalibration", "driverMonitoringState", "carControl", "lateralDelay"])
 
   publish_state = PublishState()
-  chestnut_state = ChestnutState(pm, model.chestnut) if CHESTNUT else None
+  chestnut_state = ChestnutGpuState(pm, model.chestnut) if CHESTNUT else None
 
   # setup filter to track dropped frames
   frame_dropped_filter = FirstOrderFilter(0., 10., 1. / model.constants.MODEL_FREQ)
@@ -514,13 +440,12 @@ def main(demo=False):
     mt1 = time.perf_counter()
     try:
       send_chestnut = (chestnut_state is not None and
-                       run_count % round(model.constants.MODEL_FREQ / SERVICE_LIST['chestnutState'].frequency) == 0)
+                       run_count % round(model.constants.MODEL_FREQ / SERVICE_LIST['chestnutGpuState'].frequency) == 0)
       model_output = model.run(bufs, transforms, inputs, chestnut_state.send if send_chestnut else None)
     except Exception:
       if not params.get_bool("ChestnutActive"):
         raise
-      cloudlog.exception("chestnut failed, falling back to small")
-      params.put_bool("ChestnutModelError", True)
+      cloudlog.exception("big model failed, fall back to small")
       params.put_bool("ChestnutActive", False)
       assert small_model is not None
       model = small_model

openpilot/sunnypilot/modeld_v2/stock_dependencies.py

@@ -0,0 +1,171 @@
+"""
+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 _detect_desire_key(shapes: dict) -> str | None:
+  return next((key for key in shapes if key.startswith('desire')), None)
+
+
+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[_detect_desire_key(input_shapes)]
+  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

openpilot/sunnypilot/modeld_v2/tests/helpers.py

@@ -7,6 +7,8 @@ See the LICENSE.md file in the root directory for more details.
 
 import pathlib
 import tempfile
+import codecs
+import pickle
 
 import openpilot.sunnypilot.models.helpers as helpers
 import openpilot.sunnypilot.modeld_v2.modeld as modeld_module
@@ -161,6 +163,23 @@ ARCHETYPES = {
 
 
 def make_pkl_data(archetype):
+  if archetype.expected_model_type == 'supercombo':
+    slices_b64 = codecs.encode(pickle.dumps(archetype.metadata_structure['model']['output_slices']), 'base64').decode()
+    return {
+      'metadata': {
+        'metadata': {'output_slices': slices_b64},
+        'model': {
+          'input_shapes': archetype.metadata_structure['model']['input_shapes'],
+          'output_slices': archetype.metadata_structure['model']['output_slices']
+        },
+        'input_shapes': archetype.metadata_structure['model']['input_shapes'],
+        'output_slices': archetype.metadata_structure['model']['output_slices'],
+        'output_shapes': {}
+      },
+      'run_policy': _noop_jit,
+      (CAM_W, CAM_H): _noop_jit
+    }
+
   return {
     'metadata': archetype.metadata_structure,
     'run_policy': _noop_jit,

openpilot/sunnypilot/modeld_v2/tests/test_combined_pkl_loader.py

@@ -12,7 +12,7 @@ import openpilot.sunnypilot.modeld_v2.modeld as modeld_module
 from openpilot.sunnypilot.modeld_v2.modeld import _find_driving_pkl
 from openpilot.sunnypilot.modeld_v2.tests import helpers as tests_helpers
 from openpilot.sunnypilot.modeld_v2.tests.helpers import DummyModel, DummyBundle, ARCHETYPES, CAM_W, CAM_H, \
-  SPLIT_VISION_INPUT_SHAPES, SPLIT_POLICY_INPUT_SHAPES
+  SPLIT_VISION_INPUT_SHAPES
 from openpilot.common.test import OpenpilotTestCase
 
 # resolved by name from this module when a test asks for them
@@ -66,21 +66,6 @@ class TestModelStateCombinedInit(OpenpilotTestCase):
 
 
 class TestStockEquivalence(OpenpilotTestCase):
-
-  def test_split_queue_keys_match_stock(self, model_state_factory):
-    from openpilot.selfdrive.modeld.compile_modeld import make_input_queues
-    from openpilot.sunnypilot.modeld_v2.compile_modeld import derive_frame_skip
-
-    state = model_state_factory(ARCHETYPES['vision_policy_split'])
-
-    frame_skip = derive_frame_skip(SPLIT_VISION_INPUT_SHAPES, SPLIT_POLICY_INPUT_SHAPES)
-    stock_shapes = {**SPLIT_VISION_INPUT_SHAPES, **SPLIT_POLICY_INPUT_SHAPES, 'action_t': (1, 2)}
-    stock_queues, stock_npy, _frame_views = make_input_queues(stock_shapes, frame_skip, device='NPY', frame_copy_size=49152)
-
-    # sunnypilot split pipeline has tfm/big_tfm as queues (stock has them in npy only)
-    assert set(stock_queues.keys()) <= set(state.input_queues.keys())
-    assert {'desire', 'traffic_convention'} <= set(state.numpy_inputs.keys())
-
   def test_split_queue_keys_work_with_desire_key(self, model_state_factory):
     from openpilot.sunnypilot.modeld_v2.compile_modeld import derive_frame_skip, make_split_input_queues
 
@@ -108,17 +93,23 @@ class TestStockEquivalence(OpenpilotTestCase):
     from openpilot.selfdrive.modeld.helpers import dump_oob
     shapes = {'img': (1, 12, 128, 256), 'big_img': (1, 12, 128, 256), 'features_buffer': (1, 24, 32, 512),
               'desire_pulse': (1, 25, 8), 'traffic_convention': (1, 2), 'action_t': (1, 2)}
-    pkl_data = {'metadata': {'model': {'input_shapes': shapes, 'output_slices': {}}},
-                'run_model': {(CAM_W, CAM_H): tests_helpers._noop_jit}}
+    import codecs
+    import pickle
+    slices_b64 = codecs.encode(pickle.dumps({}), 'base64').decode()
+    pkl_data = {
+      'metadata': {'model': {'input_shapes': shapes, 'output_slices': {}}, 'metadata': {'output_slices': slices_b64},
+                   'input_shapes': shapes, 'output_slices': {}, 'output_shapes': {}},
+      'run_model': {(CAM_W, CAM_H): tests_helpers._noop_jit}
+    }
     with open(tmp_path / 'driving_test_tinygrad.pkl', 'wb') as f:
       dump_oob(pkl_data, f)
     bundle = DummyBundle(models=[DummyModel('supercombo', 'driving_test_tinygrad.pkl')])
     patch_modeld(bundle)
     monkeypatch.setattr(hw.Paths, 'model_root', staticmethod(lambda: str(tmp_path)))
     state = ModelState(cam_w=CAM_W, cam_h=CAM_H)
-    assert state.is_run_model and state.run_model is not None
-    assert state.run_policy is None and state.warp is None
-    assert 'img' in state.frame_views and 'big_img' in state.frame_views
+    assert state.adapter.is_run_model and state.adapter.run_model is not None
+    assert state.adapter.run_policy is None and state.adapter.warp is None
+    assert 'img' in state.adapter.frame_views and 'big_img' in state.adapter.frame_views
 
 
 ARCHETYPE_NAMES = list(ARCHETYPES.keys())
@@ -193,7 +184,7 @@ class TestInputQueueCreation(OpenpilotTestCase):
   def test_queues_not_empty(self, archetype_name, model_state_factory):
     arch = ARCHETYPES[archetype_name]
     state = model_state_factory(arch)
-    assert len(state.input_queues) > 0, f"{arch.name}: input_queues empty"
+    assert len(state.adapter.input_queues) > 0, f"{arch.name}: input_queues empty"
 
   @parameterized.expand(ARCHETYPE_NAMES, names=["archetype_name"])
   def test_npy_contains_transforms(self, archetype_name, model_state_factory):