IQ.Pilot Release Commit @ 461be14

iqpilot/selfdrive/iqmodeld/tools/oob_rewrite.py

@@ -0,0 +1,64 @@
+"""
+Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import pickletools
+import struct
+
+from iqpilot.selfdrive.iqmodeld.egpu_policy import OOB_MAGIC
+
+MIN_OOB_BYTES = 1 << 16
+NEXT_BUFFER = b"\x97"
+READONLY_BUFFER = b"\x98"
+
+
+def rewrite_oob(src: str, dst: str, min_bytes: int = MIN_OOB_BYTES) -> tuple[int, int]:
+  # tinygrad pickles device buffers as PickleBuffers, which land in-band as BYTEARRAY8/BINBYTES8
+  # without a buffer_callback; moving those opcodes out-of-band is byte-for-byte what a protocol-5
+  # dump with a buffer_callback produces, so nothing has to be unpickled (no dock needed).
+  with open(src, "rb") as f:
+    data = f.read()
+  ops = list(pickletools.genops(data))
+  proto = next((arg for op, arg, _ in ops if op.name == "PROTO"), 0)
+  if proto < 5:
+    raise ValueError(f"{src} is pickle protocol {proto}; out-of-band buffers need protocol 5")
+  moved = 0
+  tmp = dst + ".part"
+  with open(tmp, "wb") as out, open(tmp + ".buf", "wb") as bufs:
+    ops_stream = bytearray()
+    for i, (op, arg, pos) in enumerate(ops):
+      end = ops[i + 1][2] if i + 1 < len(ops) else len(data)
+      if op.name in ("BYTEARRAY8", "BINBYTES8", "BINBYTES") and len(arg) >= min_bytes:
+        ops_stream += NEXT_BUFFER
+        if op.name != "BYTEARRAY8":
+          ops_stream += READONLY_BUFFER
+        bufs.write(struct.pack("<q", len(arg)))
+        bufs.write(arg)
+        moved += 1
+      else:
+        ops_stream += data[pos:end]
+    out.write(OOB_MAGIC)
+    out.write(struct.pack("<q", len(ops_stream)))
+    out.write(ops_stream)
+  with open(tmp, "ab") as out, open(tmp + ".buf", "rb") as bufs:
+    while chunk := bufs.read(1 << 24):
+      out.write(chunk)
+  os.remove(tmp + ".buf")
+  os.replace(tmp, dst)
+  return moved, len(ops)
+
+
+def main() -> None:
+  p = argparse.ArgumentParser()
+  p.add_argument("src")
+  p.add_argument("dst")
+  args = p.parse_args()
+  moved, total = rewrite_oob(args.src, args.dst)
+  print(f"{args.dst}: moved {moved} buffers out-of-band ({total} opcodes, {os.path.getsize(args.dst) / 1e6:.1f} MB)")
+
+
+if __name__ == "__main__":
+  main()