Merge branch 'comma-pkl-converter' into reimplement-nav

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

@@ -2,8 +2,6 @@ name: Compile warp for sunnypilot modeld
 
 on:
   workflow_dispatch:
-  schedule:
-    - cron: '0 0 * * 0'
 
 jobs:
   compile_warps:
@@ -43,13 +41,13 @@ jobs:
 
             WARP_PKL="openpilot/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 ${FRAME_PARAM} --warp-to ${MODEL_SIZE} --layout yuv420 --frames 2 --output "${WARP_PKL}"
+            taskset -c 7 env DEV=QCOM IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1 python3 openpilot/sunnypilot/modeld_v2/compile_warp.py \
+              --frame ${FRAME_PARAM} --warp-to ${MODEL_SIZE} --output "${WARP_PKL}"
 
             BIG_WARP_PKL="openpilot/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 ${FRAME_PARAM} --warp-to ${MODEL_SIZE} --layout yuv420 --frames 2 --output "${BIG_WARP_PKL}"
+            taskset -c 7 env DEV=USB+AMD:LLVM FRAME_DEV=CPU FLOAT16=1 JIT_BATCH_SIZE=0 GMMU=0 TC_OPT=2 TC_MIN_GLOBALS=32 python3 openpilot/sunnypilot/modeld_v2/compile_warp.py \
+              --frame ${FRAME_PARAM} --warp-to ${MODEL_SIZE} --output "${BIG_WARP_PKL}"
           done
 
       - name: Re-enable powersave

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

@@ -102,66 +102,39 @@ jobs:
           cat $GITHUB_OUTPUT
 
       - name: Fetch Models
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: |
           cd ${{ github.workspace }}/openpilot/openpilot
+          fetch_lfs() {
+            sudo apt-get install -y -qq aria2 > /dev/null # Parallel ranges because gh is flaky
+            LFS_URL=$(git config -f "$(git rev-parse --show-toplevel)/.lfsconfig" lfs.url)
+            for f in "$@"; do
+              grep -qs "^version https://git-lfs" "$f" || continue
+              OID=$(awk -F'sha256:' '/^oid/{print $2}' "$f"); SIZE=$(awk '/^size/{print $2}' "$f")
+              HREF=$(curl -fsSL --retry 3 -X POST -H "Accept: application/vnd.git-lfs+json" -H "Content-Type: application/vnd.git-lfs+json" "$LFS_URL/objects/batch" \
+                -d "{\"operation\":\"download\",\"transfer\":[\"basic\"],\"objects\":[{\"oid\":\"$OID\",\"size\":$SIZE}]}" | jq -r '.objects[0].actions.download.href')
+              aria2c -x 16 -s 16 -k 16M --max-tries=5 --retry-wait=5 --file-allocation=none --summary-interval=30 --allow-overwrite=true \
+                -d "$(dirname "$f")" -o "$(basename "$f")" "$HREF"
+              echo "$OID  $f" | sha256sum -c -
+            done
+          }
+          UPSTREAM_PKL=selfdrive/modeld/models/big_driving_tinygrad.pkl
           if [ "${{ inputs.target_hardware }}" != "chestnut" ]; then
-            ONNX_COUNT=$(find selfdrive/modeld/models -maxdepth 1 -type f -name "*.onnx" ! -name "big_*.onnx" ! -name "dmonitoring_*.onnx" 2>/dev/null | wc -l)
+            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
+          elif compgen -G "selfdrive/modeld/models/big_*.onnx" > /dev/null; then
+            fetch_lfs selfdrive/modeld/models/big_*.onnx selfdrive/modeld/models/big_*.pkl
+            find selfdrive/modeld/models -type f \( -name "*.onnx" -o -name "*.pkl" \) ! -name "big_*.onnx" ! -name "big_*.pkl" -delete
+          elif [ -f "$UPSTREAM_PKL" ]; then
+            echo "No big model ONNX tracked: converting commaai's precompiled $UPSTREAM_PKL, with the tinygrad it was built with"
+            fetch_lfs "$UPSTREAM_PKL"
+            find selfdrive/modeld/models -type f \( -name "*.onnx" -o -name "*.pkl" \) ! -path "$UPSTREAM_PKL" -delete
+            GITMODULES="$(git rev-parse --show-toplevel)/.gitmodules"
+            TG_NAME=$(git config -f "$GITMODULES" --get-regexp '^submodule\..*\.path$' | awk '$2 == "tinygrad_repo" {sub(/^submodule\./, "", $1); sub(/\.path$/, "", $1); print $1}')
+            TG_URL=$(git config -f "$GITMODULES" "submodule.$TG_NAME.url" | sed 's|^\.\./\.\./|https://github.com/|')
+            echo "$TG_URL $(git rev-parse HEAD:tinygrad_repo) $(sha256sum "$UPSTREAM_PKL" | cut -d' ' -f1)" > selfdrive/modeld/models/upstream_pkl.txt
           else
-            ONNX_COUNT=$(find selfdrive/modeld/models -maxdepth 1 -type f -name "big_*.onnx" 2>/dev/null | wc -l)
-          fi
-          if [ "$ONNX_COUNT" -eq 0 ]; then
-            echo "Detected NEW stock architecture (no ONNX files tracked). Fetching from HF and skipping LFS pull."
-            COMMIT_MSG=$(git log -1 --pretty=%B)
-            HF_HASH=$(echo "$COMMIT_MSG" | grep -oE '[a-f0-9]{8}' | head -n 1)
-
-            if [ -n "$HF_HASH" ]; then
-              echo "Downloading ONNX for hash $HF_HASH from HuggingFace."
-              pip install huggingface_hub
-              rm -f selfdrive/modeld/models/*.onnx
-              python3 - <<EOF
-          import sys, os, shutil
-          from huggingface_hub import HfApi, hf_hub_download
-          short_hash = "$HF_HASH"
-          target_dir = "selfdrive/modeld/models"
-          repo = "commaai/openpilot_driving_models"
-          api = HfApi()
-          files = list(api.list_repo_tree(repo, repo_type="model", recursive=True))
-          matching_folder = next((f.path.split("/")[0] for f in files if f.path.startswith(short_hash)), None)
-          if not matching_folder:
-            print(f"Could not find HF folder for hash {short_hash}")
-            sys.exit(1)
-          os.makedirs(target_dir, exist_ok=True)
-          for f in files:
-            if f.path.startswith(matching_folder) and f.path.endswith(".onnx"):
-              p = hf_hub_download(repo_id=repo, repo_type="model", filename=f.path)
-              shutil.copy2(p, os.path.join(target_dir, os.path.basename(f.path)))
-          EOF
-              rm -f selfdrive/modeld/models/*.pkl
-              touch selfdrive/modeld/models/new_arch.txt
-            else
-              echo "Could not parse HuggingFace hash from commit message."
-              exit 1
-            fi
-          else
-            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
-            else
-              sudo apt-get install -y -qq aria2 > /dev/null # Parallel ranges because gh is flaky
-              LFS_URL=$(git config -f "$(git rev-parse --show-toplevel)/.lfsconfig" lfs.url)
-              for f in selfdrive/modeld/models/big_*.onnx selfdrive/modeld/models/big_*.pkl; do
-                grep -qs "^version https://git-lfs" "$f" || continue
-                OID=$(awk -F'sha256:' '/^oid/{print $2}' "$f"); SIZE=$(awk '/^size/{print $2}' "$f")
-                HREF=$(curl -fsSL --retry 3 -X POST -H "Accept: application/vnd.git-lfs+json" -H "Content-Type: application/vnd.git-lfs+json" "$LFS_URL/objects/batch" \
-                  -d "{\"operation\":\"download\",\"transfer\":[\"basic\"],\"objects\":[{\"oid\":\"$OID\",\"size\":$SIZE}]}" | jq -r '.objects[0].actions.download.href')
-                aria2c -x 16 -s 16 -k 16M --max-tries=5 --retry-wait=5 --file-allocation=none --summary-interval=30 --allow-overwrite=true \
-                  -d "$(dirname "$f")" -o "$(basename "$f")" "$HREF"
-                echo "$OID  $f" | sha256sum -c -
-              done
-              find selfdrive/modeld/models -type f \( -name "*.onnx" -o -name "*.pkl" \) ! -name "big_*.onnx" ! -name "big_*.pkl" -delete
-            fi
+            echo "::error::no big model ONNX or precompiled big_driving_tinygrad.pkl at this ref"
+            exit 1
           fi
 
           if grep -lIF "version https://git-lfs.github.com/spec/v1" selfdrive/modeld/models/*.onnx selfdrive/modeld/models/*.pkl 2>/dev/null; then
@@ -176,7 +149,7 @@ jobs:
           path: |
             ${{ github.workspace }}/openpilot/openpilot/selfdrive/modeld/models/*.onnx
             ${{ github.workspace }}/openpilot/openpilot/selfdrive/modeld/models/*.pkl
-            ${{ github.workspace }}/openpilot/openpilot/selfdrive/modeld/models/new_arch.txt
+            ${{ github.workspace }}/openpilot/openpilot/selfdrive/modeld/models/upstream_pkl.txt
           if-no-files-found: error
 
   build_model:
@@ -249,21 +222,20 @@ jobs:
 
           NATIVE_PKL=$(find "${{ env.MODELS_DIR }}" -maxdepth 1 -name "*.pkl" -print -quit)
 
-          if [ -f "${{ env.MODELS_DIR }}/new_arch.txt" ]; then
-            echo "Detected new arch"
-            if [ "${{ inputs.target_hardware }}" == "chestnut" ]; then
-              DOWNLOADED_ONNX=$(find "${{ env.MODELS_DIR }}" -maxdepth 1 -name "big_*driving*.onnx" -print -quit)
-            else
-              DOWNLOADED_ONNX=$(find "${{ env.MODELS_DIR }}" -maxdepth 1 -name "*driving*.onnx" ! -name "big_*" -print -quit)
-            fi
-            if [ -n "$DOWNLOADED_ONNX" ]; then
-              echo "Compiling ONNX: $DOWNLOADED_ONNX"
-              env ${TG_FLAGS} taskset -c 7 python3 openpilot/sunnypilot/modeld_v2/compile_modeld_v2.py \
-                "$DOWNLOADED_ONNX" "$OUTPUT_PKL" --device-input "*" --benchmark-runs 1
-            else
-              echo "Failed to find downloaded ONNX."
-              exit 1
-            fi
+          if [ -f "${{ env.MODELS_DIR }}/upstream_pkl.txt" ]; then
+            read -r TG_URL TG_SHA _ < "${{ env.MODELS_DIR }}/upstream_pkl.txt"
+            echo "Converting commaai's precompiled model with tinygrad $TG_SHA"
+            UPSTREAM_TG="${{ runner.temp }}/upstream_tinygrad"
+            rm -rf "$UPSTREAM_TG" && git init -q "$UPSTREAM_TG" && git -C "$UPSTREAM_TG" fetch -q --depth 1 "$TG_URL" "$TG_SHA" && git -C "$UPSTREAM_TG" checkout -q FETCH_HEAD
+            CONVERTED="${{ runner.temp }}/big_driving_tinygrad_plain.pkl"
+            env DEV=USB+AMD PYTHONPATH="$UPSTREAM_TG:${{ github.workspace }}" taskset -c 7 python3 openpilot/sunnypilot/modeld_v2/convert_upstream_pkl.py \
+              "${{ env.MODELS_DIR }}/big_driving_tinygrad.pkl" "$CONVERTED"
+            mv "$CONVERTED" "$OUTPUT_PKL"
+          elif NEW_ARCH_ONNX=$(find "${{ env.MODELS_DIR }}" -maxdepth 1 -name "${{ inputs.target_hardware == 'chestnut' && 'big_' || '' }}*driving*.onnx" -print -quit) &&
+               [ -n "$NEW_ARCH_ONNX" ] && grep -qa new_img "$NEW_ARCH_ONNX"; then # the new architecture: one ONNX taking the warped frames as new_img
+            echo "Compiling new architecture ONNX: $NEW_ARCH_ONNX"
+            env ${TG_FLAGS} taskset -c 7 python3 openpilot/sunnypilot/modeld_v2/compile_modeld_v2.py \
+              "$NEW_ARCH_ONNX" "$OUTPUT_PKL" --device-input "*" --benchmark-runs 1
           else
             echo "Detected onnx architecture. Routing to native compile_modeld.py."
             # Detect model type and build compile args
@@ -340,6 +312,7 @@ jobs:
             --output-dir "${{ env.OUTPUT_DIR }}" \
             --custom-name "${{ env.MODEL_NAME }}" \
             --upstream-branch "${{ inputs.upstream_branch }}" \
+            $([ -f "${{ env.MODELS_DIR }}/upstream_pkl.txt" ] && echo "--onnx-sha256 $(cut -d' ' -f3 "${{ env.MODELS_DIR }}/upstream_pkl.txt")") \
             ${{ inputs.is_20hz && '--is-20hz' || '' }}
 
       - name: Write artifact name to file

openpilot/selfdrive/modeld/dmonitoringmodeld.py

@@ -16,7 +16,6 @@ from openpilot.common.realtime import config_realtime_process
 from openpilot.common.transformations.model import dmonitoringmodel_intrinsics
 from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
 from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
-from openpilot.common.file_chunker import open_file_chunked
 from openpilot.selfdrive.modeld.parse_model_outputs import sigmoid, safe_exp
 
 SEND_RAW_PRED = os.getenv('SEND_RAW_PRED')
@@ -28,7 +27,7 @@ class ModelState:
   output: np.ndarray
 
   def __init__(self, cam_w: int, cam_h: int):
-    jits = load_oob(open_file_chunked(MODEL_PKL_PATH))
+    jits = load_oob(MODEL_PKL_PATH)
     self.DEV = jits['input_specs']['input_img'][2]
     self.input_shapes = jits['metadata']['input_shapes']
     self.output_slices = pickle.loads(base64.b64decode(jits['metadata']['metadata']['output_slices']))

openpilot/selfdrive/modeld/helpers.py

@@ -1,10 +1,11 @@
 import io
+import os
 import pickle
 import struct
 import sys
 from pathlib import Path
 
-from openpilot.common.file_chunker import get_manifest_path
+from openpilot.common.file_chunker import get_chunk_name, get_manifest_path, open_file_chunked
 from openpilot.common.hardware import AGNOS
 from openpilot.common.hardware.usb import CHESTNUT_USB_PRODUCT, USB_DEVICES_PATH, is_chestnut_usb_id
 
@@ -15,13 +16,18 @@ def modeld_pkl_path(chestnut: bool):
   prefix = 'big_' if chestnut else ''
   return MODELS_DIR / f'{prefix}driving_tinygrad.pkl'
 
-def load_oob(f, chestnut=False):
-  import numpy as np
-  from tinygrad import Context, Tensor
+def load_oob(path, chestnut=False):
+  from tinygrad import Context, Device
+  from tinygrad.device import Buffer
   device = 'USB+AMD:LLVM' if chestnut else 'QCOM' if AGNOS else 'METAL' if sys.platform == 'darwin' else 'CPU:LLVM'
-  with Context(DEV=device):
+  chunks = 0 if os.path.isfile(path) else int(Path(get_manifest_path(path)).read_text())
+  paths = [get_chunk_name(path, i, chunks) for i in range(chunks)] or [path]
+  with Context(DEV=device), open_file_chunked(path) as f:
     opcodes = f.read(struct.unpack('<q', f.read(8))[0])
-    buffers = Tensor(np.frombuffer(f.read(), np.uint8)).realize().uop.buffer
+    buffers, stage = Buffer(Device.DEFAULT, sum(map(os.path.getsize, paths)) - 8 - len(opcodes)).ensure_allocated(), bytearray(64 << 20)
+    for offset in range(0, buffers.nbytes, len(stage)):
+      n = f.readinto(piece := memoryview(stage)[:min(len(stage), buffers.nbytes - offset)])
+      buffers.view(n, offset).ensure_allocated().copy_from(Buffer("PYTHON", n, initial_value=piece[:n]))
     return type("_Unpickler", (pickle.Unpickler,), {"persistent_load": lambda self, pid: buffers.view(*pid)})(io.BytesIO(opcodes)).load()
 
 def chestnut_present() -> bool:

openpilot/selfdrive/modeld/modeld.py

@@ -6,6 +6,7 @@ from functools import cached_property
 import os
 os.environ['GMMU'] = '0' # for chestnut fast loading, noop for qcom
 from tinygrad.device import Buffer, Device
+from tinygrad.engine.realize import lower_and_compile
 from tinygrad.dtype import DType, dtypes
 from tinygrad.tensor import Tensor
 from tinygrad.helpers import round_up
@@ -34,7 +35,6 @@ 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.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.selfdrive.modeld.constants import ModelConstants, Plan
 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
@@ -143,7 +143,7 @@ class ModelState(ModelStateBase):
 
   def __init__(self, cam_w: int, cam_h: int, chestnut: bool):
     ModelStateBase.__init__(self)
-    jits = load_oob(open_file_chunked(modeld_pkl_path(chestnut)), chestnut)
+    jits = load_oob(modeld_pkl_path(chestnut), chestnut)
     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']}
@@ -158,7 +158,9 @@ class ModelState(ModelStateBase):
     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_warp.captured._linear = lower_and_compile(self.run_warp.captured._linear)
     self.run_model = jits['run']
+    self.run_model.captured._linear = lower_and_compile(self.run_model.captured._linear)
     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]

openpilot/selfdrive/navigation/helpers.py

@@ -0,0 +1,43 @@
+import time
+from typing import Any
+
+import numpy as np
+
+
+def latest_gps_source(sm) -> str | None:
+  """The service with the freshest usable GPS fix, or None."""
+  now = time.monotonic()
+  sources = [s for s in ('gpsLocation', 'gpsLocationExternal')
+             if sm.valid[s] and sm[s].hasFix and 0 <= now - sm.recv_time[s] < 2]
+  return max(sources, key=lambda s: sm.logMonoTime[s]) if sources else None
+
+
+# convert mapbox steps to valhalla enums the model was trained on
+VALHALLA_COMMAND_COUNT = 46
+VALHALLA_SIDED_TYPES = {  # (straight, right, left)
+  'depart': (1, 2, 3), 'arrive': (4, 5, 6),
+  'on ramp': (17, 18, 19), 'off ramp': (20, 20, 21), 'fork': (22, 23, 24),
+  'roundabout': (26, 26, 26), 'rotary': (26, 26, 26), 'roundabout turn': (26, 26, 26),
+  'exit roundabout': (27, 27, 27), 'exit rotary': (27, 27, 27), 'notification': (28, 28, 28),
+}
+VALHALLA_TURN_TYPES = {'straight': 8, 'slight right': 9, 'right': 10, 'sharp right': 11, 'uturn': 13,
+                       'sharp left': 14, 'left': 15, 'slight left': 16}  # turn, continue, new name, end of road
+
+
+def valhalla_type(maneuver_type: str, modifier: str) -> int | None:
+  if maneuver_type == 'merge':
+    return None  # Valhalla gives no maneuver at a merge, so the model never saw one
+  if maneuver_type in VALHALLA_SIDED_TYPES:
+    straight, right, left = VALHALLA_SIDED_TYPES[maneuver_type]
+    return right if 'right' in modifier else left if 'left' in modifier else straight
+  return VALHALLA_TURN_TYPES.get(modifier, 8)
+
+
+def navigation_input(maneuvers: Any, traveled: float) -> np.ndarray:
+  navigation = np.zeros(VALHALLA_COMMAND_COUNT + 1, dtype=np.float32)
+  for maneuver in maneuvers:
+    command, distance = valhalla_type(maneuver.type, maneuver.modifier), maneuver.distance - traveled
+    if command is not None and command > 6 and distance >= 0:
+      navigation[command], navigation[-1] = 1, distance / 100
+      break
+  return navigation

openpilot/sunnypilot/modeld_v2/compile_modeld.py

@@ -20,7 +20,7 @@ from typing import Any
 
 import numpy as np
 
-from openpilot.sunnypilot.modeld_v2.helpers import PlainJit, read_parts, write_bundle
+from openpilot.sunnypilot.modeld_v2.helpers import PlainJit, patch_fetch_fw, read_parts, write_bundle
 from openpilot.sunnypilot.modeld_v2.plain_artifact_compiler import capture_frame, record_kernel_linear, serialize_jit
 import openpilot.sunnypilot.modeld_v2.stock_dependencies as stock
 from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
@@ -34,24 +34,7 @@ from tinygrad.tensor import Tensor
 
 os.environ['GMMU'] = '0'
 
-def _patch_tinygrad_fetch_fw():
-  try:
-    import hashlib
-    import pathlib
-    import zstandard
-    from tinygrad import helpers
-    _orig_fetch_fw = helpers.fetch_fw
-    def fetch_fw(path, name, sha256):
-      p = pathlib.Path(f"/lib/firmware/{path}/{name}.zst")
-      if p.is_file():
-        blob = zstandard.ZstdDecompressor().stream_reader(p.read_bytes()).read()
-        if hashlib.sha256(blob).hexdigest() == sha256:
-          return blob
-      return _orig_fetch_fw(path, name, sha256)
-    helpers.fetch_fw = fetch_fw
-  except ImportError:
-    pass
-_patch_tinygrad_fetch_fw()
+patch_fetch_fw()
 
 MODEL_TYPES = ('vision_policy', 'supercombo', 'vision_multi_policy')
 WARP_INPUTS = ['tfm', 'big_tfm']

openpilot/sunnypilot/modeld_v2/compile_warp.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+"""
+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 argparse
+import numpy as np
+
+from openpilot.sunnypilot.modeld_v2.helpers import PlainJit, write_bundle
+from openpilot.sunnypilot.modeld_v2.plain_artifact_compiler import capture_frame, record_kernel_linear, serialize_jit
+
+
+def main():
+  parser = argparse.ArgumentParser(description="Compile the two-frame yuv420 camera warp into a plain tinygrad artifact")
+  parser.add_argument("--frame", required=True, help="width,height,stride,y_height,uv_height,buffer_size")
+  parser.add_argument("--warp-to", required=True, help="model input WxH")
+  parser.add_argument("--output", required=True)
+  args = parser.parse_args()
+
+  from examples.openpilot.compile_warp import make_frame_prepare, parse_frame, parse_size
+  from tinygrad import Device, Tensor, TinyJit
+  frame = parse_frame(args.frame)
+  warp = make_frame_prepare(frame, *parse_size(args.warp_to))
+  run = TinyJit(lambda input_frame, M_inv: Tensor.stack(*(warp(input_frame[i], M_inv[i]) for i in range(2))), prune=True)
+  rng = np.random.default_rng(42)
+  inputs = {"input_frame": Tensor(rng.integers(0, 256, (2, frame.size), dtype=np.uint8), device=Device.DEFAULT).realize(),
+            "M_inv": Tensor((rng.standard_normal((2, 3, 3)) * 8).astype(np.float32), device=Device.DEFAULT).realize()}
+  with record_kernel_linear() as recorded:
+    run(**inputs)
+    run(**inputs)
+    check, _ = capture_frame(run, inputs)
+  assert len(recorded) == 1, f"expected one captured kernel list, got {len(recorded)}"
+  part = serialize_jit(run, recorded[0], check)
+  PlainJit(*part)
+  specs = {"input_frame": ((2, frame.size), "|u1", Device.DEFAULT), "M_inv": ((2, 3, 3), "<f4", Device.DEFAULT)}
+  write_bundle(args.output, {"run": part}, {"input_specs": specs, "metadata": {}})
+  print(f"Saved warp to {args.output} ({len(part[0]['calls'])} kernels)")
+
+
+if __name__ == "__main__":
+  main()

openpilot/sunnypilot/modeld_v2/convert_upstream_pkl.py

@@ -0,0 +1,214 @@
+#!/usr/bin/env python3
+"""
+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 argparse
+import dataclasses
+import itertools
+import mmap
+import re
+import struct
+import tempfile
+import numpy as np
+
+from openpilot.sunnypilot.modeld_v2.helpers import ALIGNMENT, Blob, PlainJit, flatten, patch_fetch_fw, write_bundle
+from openpilot.sunnypilot.modeld_v2.plain_artifact_compiler import capture_frame
+
+
+def _root(u):
+  from tinygrad.uop.ops import Ops
+  while u.op in (Ops.AFTER, Ops.BITCAST, Ops.SHRINK, Ops.INDEX):
+    u = u.src[0]
+  return u
+
+
+def _view(u) -> tuple[int, int | None]:
+  from tinygrad.uop.ops import Ops
+  start, size = 0, None
+  while u.op in (Ops.AFTER, Ops.BITCAST, Ops.SHRINK):
+    if u.op is Ops.SHRINK:
+      start, size = start + u.src[1].arg, u.src[2].arg if size is None else size
+    u = u.src[0]
+  return start, size
+
+
+def _address(v):
+  from tinygrad.uop.ops import Ops
+  add = 0
+  while v.op is not Ops.GETADDR:
+    add += v.src[1].arg if v.op is Ops.ADD else 0
+    v = v.src[0]
+  start, size = _view(v.src[0])
+  return _root(v.src[0]), add + start, size
+
+
+def _patched(buf, topo) -> tuple[bytearray, dict]:
+  from tinygrad.uop.ops import Ops
+  blob, rows = bytearray(buf.max_numel() * buf.dtype.itemsize), {}
+  for u in topo:
+    if u.op is not Ops.STORE or _root(u.src[0]) is not buf:
+      continue
+    if (value := u.src[1]).op is Ops.BINARY or value.src[0].op is Ops.BINARY:
+      data = (value if value.op is Ops.BINARY else value.src[0]).arg
+      blob[:len(data)] = data
+    else:
+      view, offsets = u.src[0].src
+      rows |= {_view(view)[0] + o.arg * view.dtype.itemsize: w for o, w in zip(offsets.src, value.src, strict=True)}
+  return blob, rows
+
+
+def _evaluate(u, env):
+  from tinygrad.uop.ops import GroupOp, Ops, exec_alu
+  if u.op in (Ops.AFTER, Ops.CAST, Ops.BITCAST):
+    return _evaluate(u.src[0], env)
+  if u.op is Ops.CONST:
+    return u.arg
+  if u.op in (Ops.PARAM, Ops.RANGE):
+    return env.get(u.arg.slot if u.op is Ops.PARAM else u)
+  if u.op is Ops.LOAD:
+    table, index = _root(u.src[0].src[0]), _evaluate(u.src[0].src[1], env)
+    return struct.unpack_from(u.dtype.fmt, table.arg, index * u.dtype.itemsize)[0] if table.op is Ops.BINARY and index is not None else None
+  if u.op in GroupOp.ALU and None not in (values := [_evaluate(x, env) for x in u.src]):
+    return exec_alu(u.op, u.dtype, values, False)
+  return None
+
+
+def _input_pokes(submit, kernargs) -> dict[int, tuple[int, int]]:
+  from tinygrad.uop.ops import Ops
+  outer = dict(enumerate(submit.src[1:]))
+  kernargs_params = {i for i, a in outer.items() if a.op is Ops.GETADDR and _root(a.src[0]) is kernargs}
+  inputs = {i: _root(a.src[0]).arg.slot for i, a in outer.items() if a.op is Ops.GETADDR and _root(a.src[0]).op is Ops.PARAM}
+  pokes = {}
+  for call in (u for u in submit.src[0].src[0].toposort() if u.op is Ops.CALL and len(u.src) > 3):
+    _, address, value, *_ = call.src[1:]
+    used = list({x.arg.slot for x in value.toposort() if x.op is Ops.PARAM} & inputs.keys())
+    if len(used) != 1 or not kernargs_params & {x.arg.slot for x in address.toposort() if x.op is Ops.PARAM}:
+      continue
+    ranges = list(dict.fromkeys(x for x in (*address.toposort(), *value.toposort()) if x.op is Ops.RANGE))
+    for trip in itertools.product(*[range(int(r.vmax) + 1) for r in ranges]):
+      env = dict.fromkeys(outer, 0) | dict(zip(ranges, trip, strict=True))
+      if (offset := _evaluate(address, env)) % 8 == 0:
+        pokes[offset] = (inputs[used[0]], _evaluate(value, env))
+  return pokes
+
+
+def _elf(image: bytes, rodata: int) -> bytes:
+  names, offset, data, headers = b"\0.text\0.rodata\0.shstrtab\0", 64, b"", []
+  for name, kind, address, content in [(0, 0, 0, b""), (1, 1, ALIGNMENT, image[:rodata]), (7, 1, ALIGNMENT + rodata, image[rodata:]), (15, 3, 0, names)]:
+    headers.append(struct.pack("<IIQQQQIIQQ", name, kind, 2 if kind == 1 else 0, address, offset, len(content), 0, 0, 1, 0))
+    data, offset = data + content, offset + len(content)
+  return b"\x7fELF\x02\x01\x01" + bytes(9) + struct.pack("<HHIQQQIHHHHHH", 1, 0xe0, 1, 0, 0, offset, 0, 64, 0, 0, 64, 4, 3) + data + b"".join(headers)
+
+
+def convert(jit, kwargs: dict, frame: dict, target) -> tuple[dict, memoryview]:
+  import tinygrad.runtime.autogen.am.pm4_nv as pm4
+  import tinygrad.runtime.autogen.amd_gpu as amd_gpu
+  from tinygrad.dtype import dtypes
+  from tinygrad.engine.jit import _prepare_jit_inputs
+  from tinygrad.helpers import unwrap
+  from tinygrad.uop.ops import Ops
+  submit = jit.captured._linear.src[0].without_after
+  kernels, topo = submit.arg.aux.kernels, list(jit.captured._linear.toposort(enter_calls=True))
+  allocs = {str(u.tag): u for u in topo if u.op is Ops.ALLOC and u.tag is not None}
+  (commands, patched), kernargs = _patched(allocs["cmdbuf_compute_0"], topo), allocs["kernargs_compute_0"]
+  argument_words, pokes = _patched(kernargs, topo)[1], _input_pokes(submit, kernargs)
+  sh = {name: getattr(amd_gpu, f"reg{name}") + amd_gpu.GC_BASE__INST0_SEG0 - pm4.PACKET3_SET_SH_REG_START
+        for name in ("COMPUTE_PGM_LO", "COMPUTE_USER_DATA_0", "COMPUTE_NUM_THREAD_X")} if hasattr(amd_gpu, "regCOMPUTE_PGM_LO") else \
+       {"COMPUTE_PGM_LO": 0x20c, "COMPUTE_USER_DATA_0": 0x240, "COMPUTE_NUM_THREAD_X": 0x207}
+
+  dispatches, registers, offset = [], {}, 0
+  while offset < len(commands):
+    word = lambda o: patched.get(o, struct.unpack_from("<I", commands, o)[0])  # noqa: E731
+    if not isinstance(header := word(offset), int) or header >> 30 != 3:
+      offset += 4
+      continue
+    count, opcode = ((header >> 16) & 0x3fff) + 1, (header >> 8) & 0xff
+    body = [word(offset + 4 + 4 * i) for i in range(count)]
+    if opcode == pm4.PACKET3_SET_SH_REG:
+      registers |= {body[0] + i: v for i, v in enumerate(body[1:])}
+    elif opcode == pm4.PACKET3_DISPATCH_DIRECT:
+      local = [registers[sh["COMPUTE_NUM_THREAD_X"] + i] for i in range(3)]
+      dispatches.append((registers[sh["COMPUTE_PGM_LO"]], registers[sh["COMPUTE_USER_DATA_0"]], local, body[:3]))
+    offset += 4 + 4 * count
+  assert len(dispatches) == len(kernels), f"decoded {len(dispatches)} dispatches, the JIT has {len(kernels)} kernels"
+
+  sink = tempfile.TemporaryFile(dir=".")
+  blob, buffers, programs = Blob(sink), {}, {}
+  u8 = repr(dtypes.uint8)
+
+  def store(array) -> dict:
+    array = np.ascontiguousarray(array)
+    return {"data": blob.add(array.tobytes()), "dtype": array.dtype.str, "shape": list(array.shape)}
+
+  calls, written, needs_init = [], set(), set()
+  for (program, arguments, local, grid), kernel in zip(dispatches, kernels, strict=True):
+    image_alloc, entry, _ = _address(program)
+    if image_alloc not in programs:
+      image = bytes(_patched(image_alloc, topo)[0])
+      rodata = next(r for r in range(0, len(image) - 24, 64) if r + struct.unpack_from("<q", image, r + 16)[0] == entry)
+      programs[image_alloc] = blob.add(_elf(image, rodata))
+    start, size = _view(arguments.src[0])
+    args = []
+    for slot in range(start, start + unwrap(size), 8):
+      if slot in pokes:
+        assert pokes[slot][1] == 0, "the plain format passes whole inputs, this kernel takes a view into one"
+        args.append({"input": pokes[slot][0]})
+      elif slot in argument_words:
+        base, at, nbytes = _address(argument_words[slot])
+        args.append({"buffer": buffers.setdefault(base, len(buffers)), "offset": at, "nbytes": nbytes or base.arg.buffer.nbytes, "dtype": u8})
+    outs, ins = kernel[6]
+    calls.append({"kind": "program", "name": re.sub(r"\x1b\[[0-9;]*m", "", kernel[1]), "lib": programs[image_alloc], "target": dataclasses.asdict(target),
+                  "signature": [[f"data{i}", i, u8, [a.get("nbytes", 1)]] for i, a in enumerate(args)], "global_size": grid, "local_size": local,
+                  "globals": list(range(len(args))), "outs": list(outs), "ins": list(ins), "args": args})
+    needs_init |= {args[i]["buffer"] for i in ins if "buffer" in args[i] and args[i]["buffer"] not in written}
+    written |= {args[i]["buffer"] for i in outs if "buffer" in args[i]}
+
+  names = list(jit.captured.expected_names)
+  inputs = [{"slot": slot, "name": names[slot] if slot < len(names) else None, "dtype": repr(u.dtype), "numel": int(u.max_numel()), "device": u.device}
+            for slot, u in enumerate(_prepare_jit_inputs((), kwargs)[0])]
+  header = {"input_names": names, "calls": calls, "returns": {"kind": "NoneType", "tensors": []}, "inputs": inputs,
+            "buffers": [{"nbytes": b.arg.buffer.nbytes, "device": b.arg.buffer.device, "dtype": u8}
+                        | ({"init": blob.add(b.arg.buffer.ensure_allocated().as_memoryview())} if i in needs_init else {}) for b, i in buffers.items()],
+            "check": {"inputs": [[path, store(array), device] for path, array, device in frame["inputs"]],
+                      "outputs": [store(array) for array in frame["outputs"]], "output_container": frame["output_container"]}}
+  with sink:
+    sink.flush()
+    return header, memoryview(mmap.mmap(sink.fileno(), 0, access=mmap.ACCESS_READ))
+
+
+def main():
+  parser = argparse.ArgumentParser(description="Convert commaai's precompiled chestnut driving model into a plain tinygrad artifact")
+  parser.add_argument("pkl", help="big_driving_tinygrad.pkl")
+  parser.add_argument("output")
+  args = parser.parse_args()
+
+  patch_fetch_fw()
+  from examples.openpilot.helpers import allocate_inputs, load_pickle
+  from tinygrad import Device, Tensor
+  from tinygrad.engine.realize import lower_and_compile
+  artifact = load_pickle(args.pkl, out_of_band=True)
+  artifact["run"].captured._linear = lower_and_compile(artifact["run"].captured._linear)
+  rng = np.random.default_rng(7)
+
+  def initialize(arrays: dict):
+    for value in arrays.values():
+      value[...] = rng.standard_normal(value.shape) if value.dtype.kind == "f" else rng.integers(0, 2 if value.dtype == np.bool_ else 16, value.shape)
+  kwargs = allocate_inputs(artifact["input_specs"], initialize) | {"output_buffers": allocate_inputs(artifact["output_specs"])}
+  frame, _ = capture_frame(artifact["run"], kwargs, "output_buffers")
+  part = convert(artifact["run"], kwargs, frame, Device[Device.DEFAULT].renderer.target)
+
+  plain = PlainJit(*part)  # replays the frame, within tolerance
+  replay = {path[0]: Tensor(array, device=device).realize() for path, array, device in frame["inputs"] if len(path) == 1}
+  replay["output_buffers"] = {path[1]: Tensor(array, device=device).realize() for path, array, device in frame["inputs"] if len(path) == 2}
+  flatten(plain(**replay))
+  for index, (got, expected) in enumerate(zip([t.numpy() for t in replay["output_buffers"].values()], frame["outputs"], strict=True)):
+    assert np.array_equal(got, expected), f"output {index} of the plain artifact differs from the pickle"
+  write_bundle(args.output, {"run": part}, {key: artifact[key] for key in ("metadata", "input_specs", "output_specs")})
+  print(f"converted {len(part[0]['calls'])} kernels, bit-exact with the pickle: {args.output}")
+
+
+if __name__ == "__main__":
+  main()

openpilot/sunnypilot/modeld_v2/frame_fusion.py

@@ -0,0 +1,71 @@
+"""
+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 bisect
+import itertools
+import math
+from functools import partial
+
+import numpy as np
+
+from openpilot.common.basedir import BASEDIR
+from openpilot.sunnypilot.modeld_v2.helpers import _dtype, read_parts
+
+
+def pack_layout(specs: dict[str, tuple], frame_size: int) -> tuple[dict[str, int], int]:
+  offsets, offset = {}, 0
+  for name, (shape, dtype) in ({"tfm": ((2, 3, 3), np.float32)} | specs).items():
+    offsets[name], offset = offset, offset + -(-math.prod(shape) * np.dtype(dtype).itemsize // 128) * 128
+  return offsets | {"frames": offset}, offset + 2 * frame_size
+
+
+class _Blobs:  # the model's and the warp's blobs as one address space, for a part whose ranges span both files
+  def __init__(self, *blobs):
+    self.blobs, self.starts = blobs, list(itertools.accumulate([0] + [len(blob) for blob in blobs[:-1]]))
+
+  def __getitem__(self, span: slice):
+    index = bisect.bisect_right(self.starts, span.start) - 1
+    return self.blobs[index][span.start - self.starts[index]:span.stop - self.starts[index]]
+
+
+def fused_with_warp(cam_w: int, cam_h: int, chestnut: bool):  # load_model's extend: a native model runs upload, warp and model as one submit
+  return partial(frame_part, f"{BASEDIR}/openpilot/sunnypilot/modeld_v2/models/{'big_' if chestnut else ''}driving_warp_{cam_w}x{cam_h}_tinygrad.pkl")
+
+
+def frame_part(warp_path: str, parts: dict, extras: dict) -> dict:
+  if "input_specs" not in extras:  # a legacy model warps in its own parts
+    return {}
+  (model, model_blob), (warp, warp_blob) = parts["run"], read_parts(warp_path)[0]["run"]
+  states = sorted(name for name in extras["input_specs"] if f"next_{name}" in extras["metadata"]["output_shapes"])
+  host = {name: (shape, dtype) for name, (shape, dtype, _) in extras["input_specs"].items() if name not in states and name != "new_img"}
+  frame_size = next(s["numel"] for s in warp["inputs"] if s["name"] == "input_frame") // 2
+  offsets, packed = pack_layout(host, frame_size)
+  shift, upload, u8 = len(model_blob), len(model["buffers"]) + len(warp["buffers"]), repr(_dtype("dtypes.uint8"))
+
+  def view(name: str, nbytes: int, dtype: str) -> dict:
+    return {"buffer": upload, "offset": offsets[name], "nbytes": nbytes, "dtype": dtype}
+  outputs = sorted(s["slot"] for s in model["inputs"] if s["name"] is None)
+  names = sorted(["host", *states])
+  slots = {s["slot"]: names.index(s["name"]) for s in model["inputs"] if s["name"] in states} | \
+          {slot: len(names) + i for i, slot in enumerate(outputs)}
+  returned = warp["returns"]["tensors"][0]
+  model_inputs = {s["slot"]: {"input": slots[s["slot"]]} if s["slot"] in slots else
+                  {k: v for k, v in returned.items() if k != "shape"} | {"buffer": returned["buffer"] + len(model["buffers"])} if s["name"] == "new_img" else
+                  view(s["name"], s["numel"] * _dtype(s["dtype"]).itemsize, s["dtype"]) for s in model["inputs"]}
+  warp_inputs = {s["slot"]: view("frames" if s["name"] == "input_frame" else "tfm", s["numel"] * _dtype(s["dtype"]).itemsize, s["dtype"])
+                 for s in warp["inputs"]}
+
+  def remap(calls: list, inputs: dict, buffers: int, lib: int) -> list:
+    return [call | ({"lib": [call["lib"][0] + lib, call["lib"][1] + lib]} if "lib" in call else {}) |
+            {"args": [dict(inputs[a["input"]]) if "input" in a else a | {"buffer": a["buffer"] + buffers} for a in call["args"]]} for call in calls]
+  header = {"input_names": names, "returns": {"kind": "NoneType", "tensors": []}, "check": None,
+            "inputs": [{"slot": names.index("host"), "name": "host", "dtype": u8, "numel": packed, "device": "NPY"}] +
+                      [s | {"slot": slots[s["slot"]]} for s in model["inputs"] if s["slot"] in slots],
+            "buffers": model["buffers"] + [b | ({"init": [b["init"][0] + shift, b["init"][1] + shift]} if "init" in b else {}) for b in warp["buffers"]] +
+                       [{"nbytes": packed, "device": model["inputs"][0]["device"], "dtype": u8}],
+            "calls": [{"kind": "copy", "args": [{"buffer": upload, "offset": 0, "nbytes": packed, "dtype": u8}, {"input": names.index("host")}]}] +
+                     remap(warp["calls"], warp_inputs, len(model["buffers"]), shift) + remap(model["calls"], model_inputs, 0, 0)}
+  return {"run": (header, _Blobs(model_blob, warp_blob))}  # replaces the model's run: one set of buffers, no separate warp

openpilot/sunnypilot/modeld_v2/helpers.py

@@ -5,12 +5,14 @@ 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 bisect
 import builtins
 import contextlib
 import dataclasses
 import functools
 import hashlib
 import io
+import itertools
 import json
 import math
 import mmap
@@ -42,6 +44,14 @@ def _patch_system_flock_acquire():
 _patch_system_flock_acquire()
 
 
+def patch_fetch_fw() -> None:
+  try:
+    import zstandard
+  except ImportError:
+    return
+  zstandard.decompress = lambda data, *_: zstandard.ZstdDecompressor().stream_reader(data).read()  # ty: ignore[invalid-assignment]
+
+
 class ArtifactMismatch(RuntimeError):
   pass
 
@@ -92,7 +102,7 @@ def write_bundle(path: str, parts: dict, extras: dict | None = None):
 
 def read_parts(path: str) -> tuple[dict, dict]:
   with open(path, "rb") as file:
-    data = memoryview(mmap.mmap(file.fileno(), 0, access=mmap.ACCESS_READ))
+    data = memoryview(mmap.mmap(file.fileno(), 0, access=mmap.ACCESS_COPY))
   if bytes(data[:len(MAGIC)]) != MAGIC:
     raise ValueError(f"{path} is not a plain tinygrad model, recompile it with compile_modeld.py")
   start = len(MAGIC) + 8 + struct.unpack_from("<Q", data, len(MAGIC))[0]
@@ -101,14 +111,20 @@ def read_parts(path: str) -> tuple[dict, dict]:
   return parts, _MetadataUnpickler(io.BytesIO(bytes(blob[slice(*header["extras"])]))).load()
 
 
-def load_model(path: str, device_map: dict | None = None, verify: bool = True, cache: bool = True) -> dict:
+def _extended(path: str, extend) -> tuple[dict, dict, str]:
   parts, extras = read_parts(path)
-  if cache and (cached := _read_cache(path, parts, device_map)) is not None:
+  added = extend(parts, extras) if extend else {}
+  return parts | added, extras, _parts_digest(added) if added else ""
+
+
+def load_model(path: str, device_map: dict | None = None, verify: bool = True, cache: bool = True, extend=None) -> dict:
+  parts, extras, extra = _extended(path, extend)
+  if cache and (cached := _read_cache(path, parts, device_map, extra)) is not None:
     return extras | cached
   shared: dict = {}
   jits = {key: PlainJit(header, blob, device_map, verify, shared) for key, (header, blob) in parts.items()}
   if cache:
-    _write_cache(path, jits, device_map)
+    _write_cache(path, jits, device_map, extra)
   return extras | jits
 
 
@@ -118,19 +134,30 @@ def tinygrad_identity() -> str:
   root, digest = os.path.dirname(tinygrad.__file__), hashlib.sha256()
   for directory, _, files in sorted(os.walk(root)):
     for name in sorted(file for file in files if file.endswith(".py")):
-      stat = os.stat(path := os.path.join(directory, name))
-      digest.update(f"{os.path.relpath(path, root)}:{stat.st_size}:{stat.st_mtime_ns};".encode())
+      with open(path := os.path.join(directory, name), "rb") as file:
+        digest.update(f"{os.path.relpath(path, root)}:".encode() + file.read())
   return digest.hexdigest()
 
 
-def _cache_key(path: str, device_map: dict | None) -> bytes:
+def _parts_digest(parts: dict) -> str:
+  digest = hashlib.sha256()
+  for _, (header, blob) in sorted(parts.items(), key=lambda item: repr(item[0])):
+    digest.update(json.dumps(header, sort_keys=True).encode())
+    for call in header["calls"]:
+      if call["kind"] == "program":
+        digest.update(blob[slice(*call["lib"])])
+  return digest.hexdigest()
+
+
+def _cache_key(path: str, device_map: dict | None, extra: str = "") -> bytes:
   stat = os.stat(path)
   return json.dumps({"format": MAGIC.decode(), "tinygrad": tinygrad_identity(), "size": stat.st_size, "mtime_ns": stat.st_mtime_ns,
-                     "device_map": sorted((device_map or {}).items()), "env": [os.environ.get(name, "") for name in RUNTIME_ENVIRONMENT]}).encode()
+                     "extra": extra, "device_map": sorted((device_map or {}).items()),
+                     "env": [os.environ.get(name, "") for name in RUNTIME_ENVIRONMENT]}).encode()
 
 
-def _cache_header_matches(file, path: str, device_map: dict | None) -> bool:
-  return file.read(struct.unpack("<Q", file.read(8))[0]) == _cache_key(path, device_map)
+def _cache_header_matches(file, path: str, device_map: dict | None, extra: str = "") -> bool:
+  return file.read(struct.unpack("<Q", file.read(8))[0]) == _cache_key(path, device_map, extra)
 
 
 def cache_is_fresh(path: str, device_map: dict | None = None) -> bool: # entry point for warmup proc
@@ -141,10 +168,10 @@ def cache_is_fresh(path: str, device_map: dict | None = None) -> bool: # entry p
     return False
 
 
-def _read_cache(path: str, parts: dict, device_map: dict | None) -> dict | None:
+def _read_cache(path: str, parts: dict, device_map: dict | None, extra: str = "") -> dict | None:
   try:
     with open(f"{path}.cache", "rb") as file:
-      if not _cache_header_matches(file, path, device_map):
+      if not _cache_header_matches(file, path, device_map, extra):
         return None
       shared: dict = {}
       buffers = {repr(key): _allocate_buffers(header, blob, device_map, shared) for key, (header, blob) in parts.items()}
@@ -153,7 +180,7 @@ def _read_cache(path: str, parts: dict, device_map: dict | None) -> dict | None:
     return None
 
 
-def _write_cache(path: str, jits: dict, device_map: dict | None) -> None:
+def _write_cache(path: str, jits: dict, device_map: dict | None, extra: str = "") -> None:
   from tinygrad.device import Buffer
   owners = {id(buffer): (repr(key), index) for key, jit in jits.items() for index, buffer in enumerate(jit.buffers)}
 
@@ -161,7 +188,7 @@ def _write_cache(path: str, jits: dict, device_map: dict | None) -> None:
     def persistent_id(self, obj):
       return owners.get(id(obj)) if isinstance(obj, Buffer) else None
 
-  key, tmp_path = _cache_key(path, device_map), f"{path}.cache.{os.getpid()}.tmp"
+  key, tmp_path = _cache_key(path, device_map, extra), f"{path}.cache.{os.getpid()}.tmp"
   with contextlib.suppress(Exception):
     with open(tmp_path, "wb") as file:
       file.write(struct.pack("<Q", len(key)) + key)
@@ -189,16 +216,34 @@ def _device(recorded: str, device_map: dict | None) -> str:
   return Device.canonicalize(device_map.get(recorded, device_map.get(recorded.split(":", 1)[0], recorded)))
 
 
+def _upload(buffer, data: memoryview, chunk: int = 16 << 20) -> None:
+  from tinygrad.device import Buffer
+  source, chunk = Buffer("NPY", len(data), opaque=memoryview(bytearray(data)) if data.readonly else data), min(chunk, len(data))
+  for offset in sorted({*range(0, len(data) - chunk, chunk), len(data) - chunk}):
+    buffer.view(chunk, offset).ensure_allocated().copy_from(source.view(chunk, offset).ensure_allocated())
+
+
 def _allocate_buffers(header: dict, blob: memoryview, device_map: dict | None, shared: dict) -> list:
   import numpy as np
   from tinygrad.device import Buffer
   written = {call["args"][i]["buffer"] for call in header["calls"] for i in (call["outs"] if call["kind"] == "program" else (0,))
              if "buffer" in call["args"][i]}
+  entries = [(_device(entry["device"], device_map), blob[slice(*entry["init"])] if "init" in entry else None, entry) for entry in header["buffers"]]
+  keys = [(device, entry["dtype"], np.frombuffer(data, np.uint8).ctypes.data, len(data)) if data and index not in written else object()
+          for index, (device, data, entry) in enumerate(entries)]
+  for device in {key[0] for key in keys if isinstance(key, tuple) and key not in shared}:
+    if shared.setdefault(("mapped", device), Buffer(device, 1).ensure_allocated().get_storage().host is not None):
+      continue
+    spans = [(entry["init"], entry["nbytes"], key) for (_, _, entry), key in zip(entries, keys, strict=True)
+             if isinstance(key, tuple) and key[0] == device and key not in shared]
+    files = getattr(blob, "starts", [0])  # a part read from several files (frame_fusion's) is uploaded a block per file
+    for _, group in itertools.groupby(sorted(spans, key=lambda span: span[0][0]), key=lambda span: bisect.bisect_right(files, span[0][0])):
+      group = list(group)
+      start, end = min(span[0] for span, _, _ in group), max(span[1] for span, _, _ in group)
+      _upload(block := Buffer(device, end - start).ensure_allocated(), blob[start:end])
+      shared |= {key: block.view(nbytes, span[0] - start).ensure_allocated() for span, nbytes, key in group}
   buffers = []
-  for index, entry in enumerate(header["buffers"]):
-    device, data = _device(entry["device"], device_map), blob[slice(*entry["init"])] if "init" in entry else None
-    # read-only data the parts share in the file (weights common to every camera resolution) is uploaded once to the nut
-    key = (device, entry["dtype"], np.frombuffer(data, np.uint8).ctypes.data, len(data)) if data and index not in written else object()
+  for (device, data, entry), key in zip(entries, keys, strict=True):
     if key not in shared:
       shared[key] = Buffer(device, entry["nbytes"], initial_value=pickle.PickleBuffer(data) if data else None).ensure_allocated()
     buffers.append(shared[key])

openpilot/sunnypilot/modeld_v2/model_adapters.py

@@ -11,13 +11,14 @@ import pickle
 import numpy as np
 
 from openpilot.common.basedir import BASEDIR
+from openpilot.sunnypilot.modeld_v2.frame_fusion import pack_layout
+from openpilot.sunnypilot.modeld_v2.helpers import load_model
 from openpilot.sunnypilot.modeld_v2.compile_modeld import (POLICY_INPUTS, derive_frame_skip,
                                                            make_split_input_queues, make_supercombo_input_queues)
 from openpilot.sunnypilot.modeld_v2.stock_dependencies import make_input_queues as stock_make_input_queues, nv12_copy_size
 from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
 from tinygrad.device import Buffer
-from tinygrad.dtype import DType, dtypes
-from tinygrad.helpers import round_up
+from tinygrad.dtype import DType
 from tinygrad.tensor import Tensor
 from tinygrad.uop.ops import UOp
 
@@ -64,16 +65,10 @@ class BaseModelAdapter:
       self.warp_frame_size = self.nv12_info[3]
       return self.jits[(self.cam_w, self.cam_h)]
 
-    warp_dir = Path(BASEDIR) / "openpilot/sunnypilot/modeld_v2/models"
     warp_name = f'{"big_" if self.chestnut else ""}driving_warp_{self.cam_w}x{self.cam_h}_tinygrad.pkl'
-    with open(warp_dir / warp_name, 'rb') as f:
-      warp_data = pickle.load(f)
-      run_warp = warp_data['run']
-      if 'input_specs' in warp_data and 'input_frame' in warp_data['input_specs']:
-        self.warp_frame_size = warp_data['input_specs']['input_frame'][0][1]
-      else:
-        self.warp_frame_size = self.nv12_info[3] if self.chestnut else self.frame_copy_size
-      return run_warp
+    warp = load_model(str(Path(BASEDIR) / "openpilot/sunnypilot/modeld_v2/models" / warp_name), cache=False)
+    self.warp_frame_size = warp['input_specs']['input_frame'][0][1]
+    return warp['run']
 
 
 class LegacyModelAdapter(BaseModelAdapter):
@@ -171,27 +166,19 @@ class NativeTinygradAdapter(BaseModelAdapter):
     self._vision_input_names = ['img', 'big_img']
     self.vision_output_slices = pickle.loads(codecs.decode(self.jits['metadata']['metadata']['output_slices'].encode(), 'base64'))
 
-    self.run_warp = self._load_warp()
+    self.warp_frame_size = self.frame_copy_size
     self._init_common()
     self.run_model = self.jits['run']
 
-    warp_frame_size = getattr(self, 'warp_frame_size', self.nv12_info[3])
     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 * warp_frame_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, warp_frame_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')}
+    specs = {name: spec for name, spec in self.input_shapes.items() if name not in self.state_pairs and name != 'new_img'}
+    offsets, packed = pack_layout(specs, self.warp_frame_size)
+    self.packed_input = np.zeros(packed, dtype=np.uint8)
+    self.host = Tensor(self.packed_input, device='NPY')
+    self.numpy_inputs = {name: np.ndarray(shape, dtype=dtype, buffer=self.packed_input, offset=offsets[name])
+                         for name, (shape, dtype) in ({'tfm': ((2, 3, 3), np.dtype(np.float32))} | specs).items()}
+    self.frames = self.packed_input[offsets['frames']:].reshape(2, self.warp_frame_size)
 
     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():
@@ -210,9 +197,7 @@ class NativeTinygradAdapter(BaseModelAdapter):
       self.input_queues[key].assign(0).realize()
 
   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)
+    self.run_model(host=self.host, **self.input_queues, output_buffers=self.outputs)
     return self.outputs['outputs']
 
 

openpilot/sunnypilot/modeld_v2/modeld.py

@@ -32,6 +32,7 @@ from openpilot.common.transformations.camera import DEVICE_CAMERAS
 from openpilot.common.transformations.model import get_warp_matrix
 from openpilot.selfdrive.controls.lib.desire_helper import DesireHelper
 from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_from_plan, smooth_value
+from openpilot.selfdrive.navigation.helpers import navigation_input
 from openpilot.selfdrive.modeld.modeld import ChestnutGpuState
 
 from openpilot.sunnypilot.modeld_v2.fill_model_msg import fill_model_msg, fill_pose_msg, PublishState, get_curvature_from_output
@@ -42,6 +43,7 @@ from openpilot.sunnypilot.modeld_v2.camera_offset_helper import CameraOffsetHelp
 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_model
+from openpilot.sunnypilot.modeld_v2.frame_fusion import fused_with_warp
 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
@@ -106,7 +108,7 @@ class ModelState(ModelStateBase):
 
   def _init_combined(self, pkl_path, cam_w, cam_h, bundle):
     cloudlog.warning(f"loading combined pkl: {pkl_path}")
-    jits = load_model(pkl_path)
+    jits = load_model(pkl_path, extend=fused_with_warp(cam_w, cam_h, self.chestnut))
 
     metadata = jits['metadata']
     self.WARP_DEV = metadata.get('warp_dev', 'QCOM') if COMMA_HARDWARE else 'CPU'
@@ -168,6 +170,9 @@ class ModelState(ModelStateBase):
     for key in ('traffic_convention', 'lateral_control_params', 'action_t'):
       if key in self.numpy_inputs and key in inputs:
         self.numpy_inputs[key][:] = inputs[key]
+    for name in ('nav_map', 'navigation'):
+      if name in self.numpy_inputs:
+        self.numpy_inputs[name][:] = inputs.get(name, 0)
 
     if self.adapter.is_native:
       for i, key in enumerate(self._vision_input_names):
@@ -247,6 +252,10 @@ class ModelState(ModelStateBase):
     return log.ModelDataV2.Action(desiredCurvature=float(desired_curvature), desiredAcceleration=float(desired_accel), shouldStop=bool(stop))
 
 
+def takes_nav(*models) -> bool:
+  return any(name in m.numpy_inputs for m in models if m is not None for name in ('nav_map', 'navigation'))
+
+
 def main(demo=False):
   cloudlog.warning("modeld init")
 
@@ -275,6 +284,7 @@ def main(demo=False):
   vipc_client_main_stream = VisionStreamType.VISION_STREAM_WIDE_ROAD if main_wide_camera else VisionStreamType.VISION_STREAM_NARROW_ROAD
   vipc_client_main = VisionIpcClient("camerad", vipc_client_main_stream, True)
   vipc_client_extra = VisionIpcClient("camerad", VisionStreamType.VISION_STREAM_WIDE_ROAD, False)
+  vipc_client_map = VisionIpcClient("mapsd", VisionStreamType.VISION_STREAM_MAP, True)
   cloudlog.warning(f"vision stream set up, main_wide_camera: {main_wide_camera}, use_extra_client: {use_extra_client}")
 
   while not vipc_client_main.connect(False):
@@ -316,7 +326,9 @@ def main(demo=False):
   # messaging
   pub_socks = ["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"] + (["chestnutGpuState"] if CHESTNUT else [])
   pm = PubMaster(pub_socks)
-  sm = SubMaster(["deviceState", "carState", "narrowRoadCameraState", "extrinsicsCalibration", "driverMonitoringState", "carControl", "lateralDelay"])
+  nav = takes_nav(model, small_model)
+  sm = SubMaster(["deviceState", "carState", "narrowRoadCameraState", "extrinsicsCalibration", "driverMonitoringState", "carControl", "lateralDelay"] +
+                 (["navInstruction"] if nav else []))
 
   publish_state = PublishState()
   chestnut_state = ChestnutGpuState(pm, model.chestnut) if CHESTNUT else None
@@ -349,6 +361,7 @@ def main(demo=False):
   DH = DesireHelper()
   meta_constants = load_meta_constants()
   RELC = RoadEdgeLaneChangeController()
+  nav_map, nav_map_time = None, 0.
 
   while True:
     # Keep receiving frames until we are at least 1 frame ahead of previous extra frame
@@ -438,6 +451,16 @@ def main(demo=False):
     if 'action_t' in model.numpy_inputs:
       inputs['action_t'] = np.array([lat_action_t, long_action_t], dtype=np.float32)
 
+    now = time.monotonic()
+    if nav and (vipc_client_map.is_connected() or vipc_client_map.connect(False)):
+      if (buf_map := vipc_client_map.recv(0)) is not None:
+        nav_map = np.frombuffer(buf_map.data, dtype=np.uint8, count=buf_map.width * buf_map.height).reshape(buf_map.height, buf_map.width).copy()
+        nav_map_time = now
+    if nav_map is not None and now - nav_map_time < 0.5: # map timeout
+      inputs['nav_map'] = nav_map
+    if nav and sm.valid['navInstruction'] and now - sm.recv_time['navInstruction'] < 2:
+      inputs['navigation'] = navigation_input(sm['navInstruction'].allManeuvers, v_ego * (now - sm.recv_time['navInstruction']))
+
     mt1 = time.perf_counter()
     try:
       send_chestnut = (chestnut_state is not None and

openpilot/sunnypilot/modeld_v2/models/big_driving_warp_1344x760_tinygrad.pkl

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ebac006f5e35ddb9c7da8ab38f9af87d6d0a87a3299b07cd118f88581f0f561f
-size 872573
+oid sha256:f1689356ff8e516f87e3d75f6b53e66953a11ee5a97e859b4808c9a16bfe671b
+size 3677806

openpilot/sunnypilot/modeld_v2/models/big_driving_warp_1928x1208_tinygrad.pkl

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d8eab79e53ebe1bbc3c633acc2a6bee80612cfe7fb5246afff6e0b7259c3a0d3
-size 872317
+oid sha256:56cc9690e0e9fb719323d2c6646c1305e8bc6bbdfdc807724a605a0d0a37554b
+size 7904366

openpilot/sunnypilot/modeld_v2/models/driving_warp_1344x760_tinygrad.pkl

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c48c9150afd01bd7071641415a248308c9a7b3cb47e7095862ef66c9f0429e43
-size 838183
+oid sha256:4b22fd6bcb3a3e8a642e6866449dd9c30c12e694ff0741cad9e34e4b00d8c5fa
+size 3698799

openpilot/sunnypilot/modeld_v2/models/driving_warp_1928x1208_tinygrad.pkl

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:59564a7d214fe13dbb30cdd068437c4bf35ec071789c922288046636f711cdfa
-size 837991
+oid sha256:9a17fd6acece4fab23784ef0f3f8bba28a049e9599e5e9e5ee7b36e4b1d24a6f
+size 7924847

openpilot/sunnypilot/modeld_v2/plain_artifact_compiler.py

@@ -67,7 +67,7 @@ def serialize_jit(jit, kernel_linear, check: dict | None = None) -> tuple[dict,
                     "global_size": [int(x) for x in global_size], "local_size": [int(x) for x in local_size],
                     "globals": list(info.globals), "outs": list(info.outs), "ins": list(info.ins), "args": arguments})
       outs, ins = info.outs, info.ins
-    elif ast.op is Ops.COPY:
+    elif ast.op is Ops.STORE:  # a copy: STORE(PARAM dst, PARAM src)
       calls.append({"kind": "copy", "args": arguments})
       outs, ins = (0,), (1,)
     else:

openpilot/sunnypilot/modeld_v2/tests/helpers.py

@@ -209,7 +209,7 @@ def patch_modeld(monkeypatch):
   def _patch(bundle):
     monkeypatch.setattr(helpers, 'get_active_bundle', lambda params=None, *, chestnut=None: bundle)
     monkeypatch.setattr(modeld_module, 'get_active_bundle', lambda params=None, *, chestnut=None: bundle)
-    monkeypatch.setattr(modeld_module, 'load_model', lambda path: pickle.loads(pathlib.Path(path).read_bytes()))
+    monkeypatch.setattr(modeld_module, 'load_model', lambda path, **_: pickle.loads(pathlib.Path(path).read_bytes()))
 
   return _patch
 

openpilot/sunnypilot/modeld_v2/tests/test_fallback.py

@@ -42,7 +42,7 @@ class TestFallback(OpenpilotTestCase):
     small_bundle = ModelParser.parse_models(small_json)[-1]
 
     oob_bytes = pickle.dumps(tests_helpers.make_pkl_data(tests_helpers.ARCHETYPES['supercombo_non20hz']))
-    monkeypatch.setattr(modeld_module, 'load_model', lambda path: pickle.loads(open(path, 'rb').read()))
+    monkeypatch.setattr(modeld_module, 'load_model', lambda path, **_: pickle.loads(open(path, 'rb').read()))
 
     for bundle in (big_bundle, small_bundle):
       (tmp_path / bundle.models[0].artifact.fileName).write_bytes(oob_bytes)

openpilot/sunnypilot/modeld_v2/warmup.py

@@ -13,28 +13,28 @@ import time
 from openpilot.common.hardware.hw import Paths
 from openpilot.common.params import Params
 from openpilot.common.swaglog import cloudlog
-from openpilot.selfdrive.modeld.helpers import chestnut_present
-from openpilot.sunnypilot.modeld_v2.helpers import cache_is_fresh, load_model
+from openpilot.sunnypilot.modeld_v2.helpers import cache_is_fresh, load_model, read_parts
 from openpilot.sunnypilot.models.helpers import get_selected_bundle
 
 
 def main():
   """
   The main entry point for offroad cache creation.
-  Runs every 10 seconds and checks for stale model caches.
-  Chestnut cache needs the gpu to build, so that slot waits until it is plugged in, or simply runs during modeld startup.
+  Runs every 10 seconds and checks for a stale qcom model cache.
   """
-  params, failed, spawn = Params(), set(), multiprocessing.get_context("spawn")
+  params, failed, native, spawn = Params(), set(), set(), multiprocessing.get_context("spawn")
   while True:
-    for bundle in [get_selected_bundle(params, source) for source in (("qcom", "chestnut") if chestnut_present() else ("qcom",))]:
-      path = os.path.join(Paths.model_root(), bundle.models[0].artifact.fileName) if bundle and bundle.models else ""
-      if not os.path.isfile(path) or cache_is_fresh(path) or (attempt := (path, os.stat(path).st_mtime_ns)) in failed:
-        continue
-      (worker := spawn.Process(target=load_model, args=(path,), daemon=True)).start()
-      worker.join()
-      if not cache_is_fresh(path):
-        cloudlog.error(f"model cache warmup failed for {path} with exit code {worker.exitcode}")
-        failed.add(attempt)
+    bundle = get_selected_bundle(params, "qcom")
+    path = os.path.join(Paths.model_root(), bundle.models[0].artifact.fileName) if bundle and bundle.models else ""
+    if os.path.isfile(path) and not cache_is_fresh(path) and (attempt := (path, os.stat(path).st_mtime_ns)) not in failed | native:
+      if "input_specs" in read_parts(path)[1]:  # native models are fused
+        native.add(attempt)
+      else:
+        (worker := spawn.Process(target=load_model, args=(path,), daemon=True)).start()
+        worker.join()
+        if not cache_is_fresh(path):
+          cloudlog.error(f"model cache warmup failed for {path} with exit code {worker.exitcode}")
+          failed.add(attempt)
     time.sleep(10)