IQ.Pilot Release Commit @ bec7652

iqpilot/selfdrive/iqmodeld/tools/compile_daemon.py

@@ -5,7 +5,7 @@ Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed
 
 from __future__ import annotations
 
-from openpilot.iqpilot.selfdrive.iqmodeld.tools.daemon_jit_compiler import main
+from iqpilot.selfdrive.iqmodeld.tools.daemon_jit_compiler import main
 
 
 if __name__ == "__main__":

iqpilot/selfdrive/iqmodeld/tools/compile_model.py

@@ -0,0 +1,125 @@
+"""
+Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
+"""
+import os
+import pickle
+import re
+import shutil
+import sys
+import time
+
+import numpy as np
+
+if "JIT_BATCH_SIZE" not in os.environ:
+  os.environ["JIT_BATCH_SIZE"] = "0"
+
+from tinygrad import Context, Device, GlobalCounters, Tensor, TinyJit, dtypes
+from tinygrad.helpers import DEBUG, getenv
+from tinygrad.nn.onnx import OnnxRunner
+from tinygrad.uop.ops import Ops
+
+
+def compile_model(onnx_file, output):
+  run_onnx = OnnxRunner(onnx_file)
+  print("loaded model")
+
+  input_shapes = {name: spec.shape for name, spec in run_onnx.graph_inputs.items()}
+  input_types = {name: spec.dtype for name, spec in run_onnx.graph_inputs.items()}
+  input_types = {key: dtypes.float32 if value is dtypes.float16 else value for key, value in input_types.items()}
+  input_shapes = {key: tuple(value if isinstance(value, int) else 1 for value in shape) for key, shape in input_shapes.items()}
+
+  Tensor.manual_seed(100)
+  inputs = {
+    key: Tensor(Tensor.randn(*shape, dtype=input_types[key]).mul(8).realize().numpy(), device="NPY")
+    for key, shape in sorted(input_shapes.items())
+  }
+  if not getenv("NPY_IMG"):
+    inputs = {key: Tensor(value.numpy(), device=Device.DEFAULT).realize() if "img" in key else value for key, value in inputs.items()}
+  print("created tensors")
+
+  run_onnx_jit = TinyJit(
+    lambda **kwargs: next(iter(run_onnx({key: value.to(Device.DEFAULT) for key, value in kwargs.items()}).values())).cast("float32"),
+    prune=True,
+  )
+  test_value = None
+  for iteration in range(3):
+    GlobalCounters.reset()
+    print(f"run {iteration}")
+    with Context(DEBUG=max(DEBUG.value, 2 if iteration == 2 else 1), OPENPILOT_HACKS=1):
+      result = run_onnx_jit(**inputs).numpy()
+    if iteration == 1:
+      test_value = np.copy(result)
+
+  kernel_asts = {Ops.PROGRAM}
+  kernel_calls = [
+    node for node in run_onnx_jit.captured.linear.toposort(gate=lambda value: value.op not in kernel_asts)
+    if node.op is Ops.CALL and node.src[0].op in kernel_asts
+  ]
+  print(f"captured {len(kernel_calls)} kernels")
+  np.testing.assert_equal(test_value, result, "JIT run failed")
+  print("jit run validated")
+
+  kernel_count = 0
+  read_image_count = 0
+  gated_read_image_count = 0
+  for call in kernel_calls:
+    _, _, source, _ = call.src[0].src
+    rendered = source.arg
+    kernel_count += 1
+    read_image_count += rendered.count("read_image")
+    gated_read_image_count += rendered.count("?read_image")
+    for value in (match.group(1) for match in re.finditer(r"(val\d+)\s*=\s*read_imagef\(", rendered)):
+      if re.search(fr"[?:]{value}\.[xyzw]", rendered):
+        gated_read_image_count += 1
+
+  print(f"{kernel_count=}, {read_image_count=}, {gated_read_image_count=}")
+  expected = {
+    "kernel count": (kernel_count, getenv("ALLOWED_KERNEL_COUNT", -1)),
+    "read image count": (read_image_count, getenv("ALLOWED_READ_IMAGE", -1)),
+    "gated read image count": (gated_read_image_count, getenv("ALLOWED_GATED_READ_IMAGE", -1)),
+  }
+  for name, (actual, allowed) in expected.items():
+    if allowed != -1:
+      assert actual == allowed, f"different {name}: {actual}, expected {allowed}"
+
+  with open(output, "wb") as handle:
+    pickle.dump(run_onnx_jit, handle)
+  print(f"model size is {os.path.getsize(onnx_file) / 1e6:.2f}M")
+  print(f"pkl size is {os.path.getsize(output) / 1e6:.2f}M")
+  return run_onnx_jit, inputs, test_value
+
+
+def test_compiled(run, inputs, test_value):
+  step_times = []
+  for _ in range(20):
+    start = time.perf_counter()
+    output = run(**inputs)
+    queued = time.perf_counter()
+    value = output.numpy()
+    end = time.perf_counter()
+    step_times.append((end - start) * 1e3)
+    print(f"enqueue {(queued - start) * 1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")
+
+  minimum = getenv("ASSERT_MIN_STEP_TIME", 0.0)
+  if minimum:
+    assert min(step_times) < minimum, f"expected minimum step time below {minimum} ms, got {min(step_times)} ms"
+  np.testing.assert_equal(test_value, value)
+  changed_inputs = {key: Tensor(item.numpy() * 2, device=item.device) for key, item in inputs.items()}
+  changed_value = run(**changed_inputs).numpy()
+  np.testing.assert_raises(AssertionError, np.testing.assert_array_equal, value, changed_value)
+
+
+if __name__ == "__main__":
+  model_path = sys.argv[1]
+  output_path = sys.argv[2]
+  if stash := os.environ.get("IQPILOT_MODEL_STASH"):
+    stashed_model = os.path.join(stash, os.path.basename(output_path))
+    if os.path.isfile(stashed_model) and os.path.getsize(stashed_model) > 0:
+      os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
+      shutil.copyfile(stashed_model, output_path)
+      print(f"restored device-compiled model: {output_path}")
+      sys.exit(0)
+  _, input_values, expected_value = compile_model(model_path, output_path)
+  with open(output_path, "rb") as compiled_file:
+    compiled_model = pickle.load(compiled_file)
+  test_compiled(compiled_model, input_values, expected_value)

iqpilot/selfdrive/iqmodeld/tools/compile_split_runtime.py

@@ -62,8 +62,8 @@ WARP_DEVICE = os.getenv("WARP_DEV")
 
 
 def _read_shared_copy(path: str) -> str:
-  from openpilot.common.file_chunker import read_file_chunked
-  from openpilot.system.hardware.hw import Paths
+  from iqpilot.common.file_chunker import read_file_chunked
+  from iqpilot.system.hardware.hw import Paths
 
   shm_path = os.path.join(Paths.shm_path(), os.path.basename(path))
   atexit.register(lambda: os.path.exists(shm_path) and os.remove(shm_path))
@@ -127,8 +127,8 @@ def _project_pixels(src_flat, inverse_matrix, dst_shape, src_shape, stride_pad,
   dst_w, dst_h = dst_shape
   src_h, src_w = src_shape
 
-  x_coords = Tensor.arange(dst_w, device=WARP_DEVICE).reshape(1, dst_w).expand(dst_h, dst_w).reshape(-1)
-  y_coords = Tensor.arange(dst_h, device=WARP_DEVICE).reshape(dst_h, 1).expand(dst_h, dst_w).reshape(-1)
+  x_coords = Tensor.arange(dst_w).to(WARP_DEVICE).reshape(1, dst_w).expand(dst_h, dst_w).reshape(-1)
+  y_coords = Tensor.arange(dst_h).to(WARP_DEVICE).reshape(dst_h, 1).expand(dst_h, dst_w).reshape(-1)
 
   src_x = inverse_matrix[0, 0] * x_coords + inverse_matrix[0, 1] * y_coords + inverse_matrix[0, 2]
   src_y = inverse_matrix[1, 0] * x_coords + inverse_matrix[1, 1] * y_coords + inverse_matrix[1, 2]
@@ -352,8 +352,8 @@ def _arg_parser() -> argparse.ArgumentParser:
 
 
 def main(argv: list[str] | None = None) -> int:
-  from openpilot.iqpilot.selfdrive.iqmodeld.metadata import build_metadata_record
-  from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
+  from iqpilot.selfdrive.iqmodeld.metadata import build_metadata_record
+  from iqpilot.system.camerad.cameras.nv12_info import get_nv12_info
 
   args = _arg_parser().parse_args(argv)
   model_w, model_h = args.model_size

iqpilot/selfdrive/iqmodeld/tools/compile_supercombo.py

@@ -59,7 +59,6 @@ def warp_perspective_tinygrad(src_flat, M_inv, dst_shape, src_shape, stride_pad,
   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)
 
-  # inline 3x3 matmul as elementwise to avoid reduce op (enables fusion with gather)
   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]
@@ -100,9 +99,7 @@ def make_frame_prepare(nv12: NV12Frame, model_w, model_h):
   stride_pad = stride - cam_w
 
   def frame_prepare_tinygrad(input_frame, M_inv):
-    # UV_SCALE @ M_inv @ UV_SCALE_INV simplifies to elementwise scaling
     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=WARP_DEV)
-    # deinterleave NV12 UV plane (UVUV... -> separate U, V)
     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],
@@ -142,7 +139,6 @@ def get_policy_npy_shapes(input_shapes):
   tc = input_shapes['traffic_convention']  # (1, 2)
   at = input_shapes['action_t']  # (1, 2)
   fb = input_shapes['features_buffer']  # (1, 24, 512)
-  # TODO prev_feat shouldn't exist and be handled inside the JIT, but corrupt on QCOM for now
   shapes = {'desire': (dp[2],), 'traffic_convention': tuple(tc), 'action_t': tuple(at), 'prev_feat': (fb[0], fb[2])}
   return shapes, [math.prod(s) for s in shapes.values()]
 
@@ -155,7 +151,6 @@ def make_input_queues(input_shapes, frame_skip, device):
 
   shapes, sizes = get_policy_npy_shapes(input_shapes)
   packed_npy_inputs = np.zeros(sum(sizes), dtype=np.float32)
-  # views into the packed inputs, to be refilled at runtime
   npy.update({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.update({
     'feat_q': Tensor(np.zeros((frame_skip * fb[1], fb[0], fb[2]), dtype=np.float32), device=device).contiguous().realize(),
@@ -284,7 +279,7 @@ def _slice_outputs(model_outputs: np.ndarray, output_slices: dict[str, slice]) -
 
 
 def _validate_pose_outputs(parsed_outputs: dict[str, np.ndarray]) -> None:
-  from openpilot.selfdrive.locationd.locationd import MIN_STD_SANITY_CHECK, ROTATION_SANITY_CHECK, TRANS_SANITY_CHECK
+  from iqpilot.selfdrive.locationd.locationd import MIN_STD_SANITY_CHECK, ROTATION_SANITY_CHECK, TRANS_SANITY_CHECK
 
   required = (
     'pose', 'pose_stds', 'wide_from_device_euler', 'wide_from_device_euler_stds',
@@ -326,7 +321,7 @@ def _validate_pose_outputs(parsed_outputs: dict[str, np.ndarray]) -> None:
 
 
 def validate_supercombo_release(run_policy_jit, model_runner, model_metadata, frame_skip, expected_device: str) -> None:
-  from openpilot.iqpilot.selfdrive.iqmodeld.parser import PhaseParser
+  from iqpilot.selfdrive.iqmodeld.parser import PhaseParser
 
   direct_fn = make_run_policy(model_runner, model_metadata, frame_skip)
   parser = PhaseParser()
@@ -370,8 +365,8 @@ def _parse_size(s):
 
 
 def read_file_chunked_to_shm(path):
-  from openpilot.common.file_chunker import read_file_chunked
-  from openpilot.system.hardware.hw import Paths
+  from iqpilot.common.file_chunker import read_file_chunked
+  from iqpilot.system.hardware.hw import Paths
   with tempfile.NamedTemporaryFile(prefix='compile_modeld_', dir=Paths.shm_path(), delete=False) as f:
     f.write(read_file_chunked(path))
     tmp_path = f.name
@@ -381,8 +376,8 @@ def read_file_chunked_to_shm(path):
 
 if __name__ == "__main__":
   from tinygrad.nn.onnx import OnnxRunner
-  from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
-  from openpilot.iqpilot.selfdrive.iqmodeld.metadata import build_metadata_record
+  from iqpilot.system.camerad.cameras.nv12_info import get_nv12_info
+  from iqpilot.selfdrive.iqmodeld.metadata import build_metadata_record
   p = argparse.ArgumentParser()
   p.add_argument('--model-size', type=_parse_size, required=True, help='model input WxH')
   p.add_argument('--camera-resolutions', type=_parse_size, nargs='+', required=True,

iqpilot/selfdrive/iqmodeld/tools/daemon_jit_compiler.py

@@ -272,8 +272,8 @@ def _parse_size(text: str) -> tuple[int, int]:
 
 
 def _read_file_to_shared_memory(path: str) -> str:
-  from openpilot.common.file_chunker import read_file_chunked
-  from openpilot.system.hardware.hw import Paths
+  from iqpilot.common.file_chunker import read_file_chunked
+  from iqpilot.system.hardware.hw import Paths
 
   shm_path = os.path.join(Paths.shm_path(), os.path.basename(path))
   atexit.register(lambda: os.path.exists(shm_path) and os.remove(shm_path))
@@ -295,8 +295,8 @@ def _arg_parser() -> argparse.ArgumentParser:
 
 
 def main(argv: list[str] | None = None) -> int:
-  from openpilot.iqpilot.selfdrive.iqmodeld.metadata import build_metadata_record
-  from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
+  from iqpilot.selfdrive.iqmodeld.metadata import build_metadata_record
+  from iqpilot.system.camerad.cameras.nv12_info import get_nv12_info
 
   args = _arg_parser().parse_args(argv)
   model_w, model_h = args.model_size

iqpilot/selfdrive/iqmodeld/tools/install_models_pc.py

@@ -13,7 +13,7 @@ from pathlib import Path
 
 import onnx
 
-from openpilot.system.hardware.hw import Paths
+from iqpilot.system.hardware.hw import Paths
 
 _MODEL_STEMS = ("driving_off_policy", "driving_on_policy", "driving_policy", "driving_vision")