Runtime Dependencies

Extend vendored runtime helpers and diagnostics while preserving the existing compiler and model ABI.

msgq_repo/msgq/msgq.h

@@ -6,7 +6,7 @@
 #include <atomic>
 
 #define DEFAULT_SEGMENT_SIZE (1 * 1024 * 1024)
-#define NUM_READERS 15
+#define NUM_READERS 64
 #define ALIGN(n) ((n + (8 - 1)) & -8)
 
 #define UNUSED(x) (void)x

msgq_repo/msgq/msgq_tests.cc

@@ -1,3 +1,5 @@
+#include <array>
+#include <unistd.h>
 #include "msgq/test_runner.h"
 #include "msgq/msgq.h"
 
@@ -424,3 +426,57 @@ TEST_CASE("1 publisher, 2 subscribers")
     msgq_msg_close(&msg2);
   }
 }
+
+TEST_CASE("32 active subscribers retain every message across publisher restart")
+{
+  REQUIRE(NUM_READERS >= 32);
+  const std::string name = "capacity_acceptance_" + std::to_string(getpid());
+#ifdef __APPLE__
+  std::string file = "/tmp/msgq_";
+#else
+  std::string file = "/dev/shm/msgq_";
+#endif
+  const char *prefix = std::getenv("OPENPILOT_PREFIX");
+  if (prefix != nullptr) file += std::string(prefix) + "/";
+  file += name;
+  std::remove(file.c_str());
+  msgq_queue_t pub{};
+  REQUIRE(msgq_new_queue(&pub, name.c_str(), 4096) == 0);
+  msgq_init_publisher(&pub);
+  std::array<msgq_queue_t, 32> readers{};
+  for (auto &reader : readers) {
+    REQUIRE(msgq_new_queue(&reader, name.c_str(), 4096) == 0);
+    msgq_init_subscriber(&reader);
+  }
+  for (unsigned epoch = 0; epoch < 2; ++epoch) {
+    if (epoch) {
+      msgq_init_publisher(&pub);
+      for (auto &reader : readers) {
+        msgq_msg_t empty{};
+        REQUIRE(msgq_msg_recv(&empty, &reader) == 0);
+      }
+      REQUIRE(*pub.num_readers == readers.size());
+    }
+    for (uint64_t tick = 1; tick <= 200; ++tick) {
+      uint64_t sequence = epoch * 1000 + tick;
+      msgq_msg_t outgoing{};
+      msgq_msg_init_data(&outgoing, reinterpret_cast<char *>(&sequence), sizeof(sequence));
+      REQUIRE(msgq_msg_send(&outgoing, &pub) == sizeof(sequence));
+      msgq_msg_close(&outgoing);
+      for (auto &reader : readers) {
+        msgq_msg_t incoming{};
+        REQUIRE(msgq_msg_recv(&incoming, &reader) == sizeof(sequence));
+        uint64_t observed = 0;
+        std::memcpy(&observed, incoming.data, sizeof(observed));
+        REQUIRE(observed == sequence);
+        msgq_msg_close(&incoming);
+        REQUIRE(*reader.read_valids[reader.reader_id]);
+        REQUIRE(*reader.read_uids[reader.reader_id] == reader.read_uid_local);
+      }
+      REQUIRE(*pub.num_readers == readers.size());
+    }
+  }
+  for (auto &reader : readers) msgq_close_queue(&reader);
+  msgq_close_queue(&pub);
+  std::remove(file.c_str());
+}

msgq_repo/msgq/visionipc/tests/test_visionipc.py

@@ -1,5 +1,6 @@
 import struct
 import unittest
+import uuid
 from typing import Optional
 from msgq.visionipc import VisionIpcServer, VisionIpcClient
 
@@ -9,6 +10,8 @@ class TestVisionIpc(unittest.TestCase):
   client: Optional[VisionIpcClient]
 
   def setUp(self):
+    # Each test needs independent listener and publisher endpoints in parallel runs.
+    self.name = f"test_{uuid.uuid4().hex}"
     self.server = None
     self.client = None
 
@@ -16,14 +19,14 @@ class TestVisionIpc(unittest.TestCase):
     self.client = None
     self.server = None
 
-  def setup_vipc(self, name, *stream_types, num_buffers=1, width=100, height=100, conflate=False):
-    self.server = VisionIpcServer(name)
+  def setup_vipc(self, *stream_types, num_buffers=1, width=100, height=100, conflate=False):
+    self.server = VisionIpcServer(self.name)
     for stream_type in stream_types:
       self.server.create_buffers(stream_type, num_buffers, width, height)
     self.server.start_listener()
 
     if len(stream_types):
-      self.client = VisionIpcClient(name, stream_types[0], conflate)
+      self.client = VisionIpcClient(self.name, stream_types[0], conflate)
       assert self.client.connect(True)
     else:
       self.client = None
@@ -31,19 +34,19 @@ class TestVisionIpc(unittest.TestCase):
     return self.server, self.client
 
   def test_connect(self):
-    self.setup_vipc("camerad", 0)
+    self.setup_vipc(0)
     assert self.client is not None
     assert self.client.is_connected()
 
   def test_available_streams(self):
     stream_types = (0, 2)
-    self.setup_vipc("camerad", *stream_types)
-    available_streams = VisionIpcClient.available_streams("camerad", True)
+    self.setup_vipc(*stream_types)
+    available_streams = VisionIpcClient.available_streams(self.name, True)
     assert available_streams == set(stream_types)
 
   def test_buffers(self):
     width, height, num_buffers = 100, 200, 5
-    self.setup_vipc("camerad", 0, num_buffers=num_buffers, width=width, height=height)
+    self.setup_vipc(0, num_buffers=num_buffers, width=width, height=height)
     assert self.client is not None
     assert self.client.width == width
     assert self.client.height == height
@@ -51,7 +54,7 @@ class TestVisionIpc(unittest.TestCase):
     assert self.client.num_buffers == num_buffers
 
   def test_send_single_buffer(self):
-    self.setup_vipc("camerad", 0)
+    self.setup_vipc(0)
     assert self.server is not None
     assert self.client is not None
     assert self.client.buffer_len is not None
@@ -70,7 +73,7 @@ class TestVisionIpc(unittest.TestCase):
     assert recv_buf.frame_id == 1337
 
   def test_no_conflate(self):
-    self.setup_vipc("camerad", 0)
+    self.setup_vipc(0)
     assert self.server is not None
     assert self.client is not None
     assert self.client.buffer_len is not None
@@ -87,7 +90,7 @@ class TestVisionIpc(unittest.TestCase):
     assert self.client.frame_id == 2
 
   def test_conflate(self):
-    self.setup_vipc("camerad", 0, conflate=True)
+    self.setup_vipc(0, conflate=True)
     assert self.server is not None
     assert self.client is not None
     assert self.client.buffer_len is not None

openpilot/starpilot/models/tests/test_helper_import.py

@@ -38,6 +38,7 @@ with Context(DEV=""), patch.object(type(Device), "_select_device", property(unex
   def load_helper():
     path = Path(__file__).resolve().parents[4] / "tinygrad_repo/examples/openpilot/helpers.py"
     spec = importlib.util.spec_from_file_location("benchmark_import_regression", path)
+    assert spec is not None and spec.loader is not None
     helper = importlib.util.module_from_spec(spec)
     # Import under QCOM, then call under another device to catch settings that
     # were incorrectly captured by an import-time decorator.

openpilot/starpilot/models/tests/test_startup_trace.py

@@ -0,0 +1,122 @@
+from types import SimpleNamespace
+import unittest
+from unittest.mock import patch
+
+from tinygrad.runtime.support.am import startup_trace as trace
+
+
+class TestStartupTrace(unittest.TestCase):
+  def setUp(self):
+    self.packets = []
+    for name, replacement in (("ENABLED", True), ("REPORT", self.packets.append), ("_short_reads", 0)):
+      override = patch.object(trace, name, replacement)
+      override.start()
+      self.addCleanup(override.stop)
+
+  def test_entries_and_poll_samples_are_bounded(self):
+    dev = SimpleNamespace()
+    for result in range(10000):
+      self.assertEqual(trace.value(dev, "cp_status", result), result)
+    for index in range(100):
+      trace.note(dev, "stage", index=index)
+      trace.value(dev, f"extra_{index}", index)
+    trace.emit(dev, "complete")
+    entries = self.packets[0]["entries"]
+    self.assertEqual(len(entries), 64)
+    self.assertEqual((entries[0]["first"], entries[0]["last"], entries[0]["reads"]), (0, 9999, 10000))
+    self.assertNotIn("_startup_trace_entries", dev.__dict__)
+
+  def test_ack_mask_read_count_and_first_gc_only(self):
+    dev = SimpleNamespace()
+    reads = []
+    values = iter((0x80, 0x82, 0x84, 0x84))
+    returned = object()
+
+    def read_ack():
+      raw = next(values)
+      reads.append(raw)
+      return raw
+
+    def wait(read, **kwargs):
+      self.assertEqual(kwargs, {"value": 4, "msg": "flush_tlb timeout"})
+      while read() != kwargs["value"]:
+        pass
+      return returned
+
+    for _ in range(2):
+      self.assertIs(trace.wait_tlb(dev, "GC", 1, 123, 2, wait, read_ack), returned)
+    self.assertEqual(reads, [0x80, 0x82, 0x84, 0x84])
+    self.assertEqual(len(self.packets), 1)
+    self.assertEqual(self.packets[0]["outcome"], "first_gc_acknowledged")
+    sample = self.packets[0]["entries"][-1]
+    self.assertEqual((sample["reads"], sample["first"], sample["last"]), (3, 0x80, 0x84))
+
+  def test_non_gc_and_disabled_trace_preserve_wait(self):
+    for ip, enabled in (("MM", True), ("GC", False)):
+      with self.subTest(ip=ip, enabled=enabled), patch.object(trace, "ENABLED", enabled):
+        dev = SimpleNamespace()
+        reads = []
+
+        def read_ack(reads=reads):
+          reads.append(9)
+          return 9
+
+        def wait(read, **kwargs):
+          self.assertEqual(kwargs, {"value": 1, "msg": "flush_tlb timeout"})
+          return read()
+
+        self.assertEqual(trace.wait_tlb(dev, ip, 0, 123, 0, wait, read_ack), 1)
+        self.assertEqual(reads, [9])
+        self.assertEqual(dev.__dict__, {})
+    self.assertEqual(self.packets, [])
+
+  def test_wait_failure_preserves_exception_and_observed_reads(self):
+    error = TimeoutError("original timeout")
+
+    def wait(read, **kwargs):
+      self.assertEqual(read(), 0)
+      self.assertEqual(read(), 0)
+      raise error
+
+    with self.assertRaises(TimeoutError) as caught:
+      trace.wait_tlb(SimpleNamespace(), "GC", 0, 123, 0, wait, lambda: 0x80)
+    self.assertIs(caught.exception, error)
+    self.assertEqual(self.packets[0]["outcome"], "first_gc_failure")
+    sample = self.packets[0]["entries"][-1]
+    self.assertEqual((sample["reads"], sample["first"], sample["last"]), (2, 0x80, 0x80))
+
+  def test_callback_failure_does_not_replace_result_or_exception(self):
+    error = RuntimeError("original initialization failure")
+    returned = object()
+
+    @trace.initialization
+    def initialize(dev, fail):
+      if fail:
+        raise error
+      return returned
+
+    def broken_report(packet):
+      raise OSError("logger unavailable")
+
+    with patch.object(trace, "REPORT", broken_report):
+      self.assertIs(initialize(SimpleNamespace(), False), returned)
+      with self.assertRaises(RuntimeError) as caught:
+        initialize(SimpleNamespace(), True)
+      self.assertIs(caught.exception, error)
+      self.assertEqual(trace.wait_tlb(SimpleNamespace(), "GC", 0, 123, 0, lambda read, **kwargs: read(), lambda: 1), 1)
+      trace.short_read(0x1000, 4, 0)
+
+  def test_short_read_reports_are_bounded(self):
+    for _ in range(100):
+      trace.short_read(0x1000, 4, 2)
+    self.assertEqual(len(self.packets), 8)
+    self.assertEqual([packet["sample"] for packet in self.packets], list(range(1, 9)))
+    self.assertTrue(all((packet["address"], packet["expected"], packet["actual"]) == (0x1000, 4, 2)
+                        for packet in self.packets))
+    with patch.object(trace, "ENABLED", False), patch.object(trace, "_short_reads", 0):
+      trace.short_read(0x1000, 4, 0)
+    self.assertEqual(len(self.packets), 8)
+
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/LOCAL_PATCHES.md

@@ -0,0 +1,9 @@
+# Local compatibility patch
+
+`tinygrad/helpers.py::fetch_fw` reads installed `.zst` firmware through the
+`zstandard` streaming decoder on Python before 3.14. AGNOS 19.6.20's
+`amdgpu/psp_14_0_2_sos.bin.zst` has no frame content-size field, so the former
+one-shot decoder raised before Tinygrad could compare the decompressed bytes
+with its pinned firmware SHA-256. Streaming retains the same hash check and
+the pinned upstream fallback for a missing or mismatched local file. It does
+not change device firmware or the AMD runtime protocol.

tinygrad_repo/examples/openpilot/helpers.py

@@ -50,15 +50,16 @@ def load_pickle(path, *, out_of_band=False):
   u.persistent_load = persistent_load
   return u.load()
 
-@Context(OPENPILOT_HACKS=1, **{'AMD': {'TC_OPT': 2, 'TC_MIN_GLOBALS': 32}}.get(Device.DEFAULT, {}))
 def benchmark(fxn:Callable, cb=None, **kwargs):
-  Device.default.synchronize()
-  start = time.perf_counter()
-  if (output := fxn(**kwargs)) is not None: output.realize()
-  Device.default.synchronize()
-  end = time.perf_counter()
-  if cb: cb(end-start)
-  return [t.numpy().copy() for t in get_parameters(kwargs.get('output_buffers', output))]
+  # Select benchmark options when called, without probing devices on import.
+  with Context(OPENPILOT_HACKS=1, **{'AMD': {'TC_OPT': 2, 'TC_MIN_GLOBALS': 32}}.get(Device.DEFAULT, {})):
+    Device.default.synchronize()
+    start = time.perf_counter()
+    if (output := fxn(**kwargs)) is not None: output.realize()
+    Device.default.synchronize()
+    end = time.perf_counter()
+    if cb: cb(end-start)
+    return [t.numpy().copy() for t in get_parameters(kwargs.get('output_buffers', output))]
 
 pm_retargetable = PatternMatcher([
   (UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat()), name="p"), lambda p: p.replace(src=p.src[:-1]) if p.arg.target.device == "CPU" else None)

tinygrad_repo/test/device/test_qcom.py

@@ -1,8 +1,39 @@
 import ctypes, platform, unittest
+from types import SimpleNamespace
+from unittest.mock import patch
 from tinygrad import Device
+from tinygrad.device import TinyELF
+from tinygrad.dtype import dtypes
+from tinygrad.helpers import Target
 from tinygrad.renderer.cstyle import ClangRenderer
 
 class TestQCOM(unittest.TestCase):
+  def test_program_cache_preserves_signature(self):
+    from tinygrad.runtime import ops_qcom
+
+    target = Target(device='QCOM')
+    common = (b'compiled-binary', 'same_kernel', target)
+    first = TinyELF(*common, ((None, 0, dtypes.float, (4,)),))
+    second = TinyELF(*common, ((None, 0, dtypes.half, (8,)),))
+
+    class Program:
+      src = (None, None, None, SimpleNamespace(arg=common[0]))
+      def __init__(self, elf): self.elf = elf
+      def to_elf(self): return self.elf
+
+    ops_qcom._qcom_program_cache.clear()
+    try:
+      with patch.object(ops_qcom, 'QCOMProgramData', side_effect=lambda _dev, elf: SimpleNamespace(image=b'code', signature=elf.signature)) as parsed:
+        with patch.object(ops_qcom, 'patch', side_effect=lambda _buf, _offsets, _image: object()):
+          a = ops_qcom.qcom_build_program(object(), Program(first), ('QCOM',))
+          b = ops_qcom.qcom_build_program(object(), Program(second), ('QCOM',))
+          again = ops_qcom.qcom_build_program(object(), Program(first), ('QCOM',))
+      self.assertIsNot(a, b)
+      self.assertIs(a, again)
+      self.assertEqual(parsed.call_count, 2)
+    finally:
+      ops_qcom._qcom_program_cache.clear()
+
   # although part of the QCOM runtime, this tests flushing the CPU's dcache
   @unittest.skipUnless(isinstance(Device["CPU"].renderer, ClangRenderer) and platform.machine().lower() in {"arm64", "aarch64"},
                        "dcache_flush's inline asm needs ClangRenderer, and runs on arm64")

tinygrad_repo/test/unit/test_helpers.py

@@ -1,6 +1,7 @@
-import unittest, gc
+import unittest, gc, hashlib, io, pathlib
+from unittest.mock import patch
 import numpy as np
-from tinygrad.helpers import polyN, disable_gc
+from tinygrad.helpers import polyN, disable_gc, fetch_fw
 from tinygrad.tensor import Tensor, is_numpy_ndarray
 
 class TestPolyN(unittest.TestCase):
@@ -25,5 +26,28 @@ class TestDisableGC(unittest.TestCase):
     finally:
       (gc.enable if was_enabled else gc.disable)()
 
+class TestFirmwareFetch(unittest.TestCase):
+  def test_local_zstd_unknown_content_size_and_hash(self):
+    from zstandard import ZstdCompressor
+    firmware = b"verified AMD firmware fixture" * 20
+    digest = hashlib.sha256(firmware).hexdigest()
+    for write_content_size in (False, True):
+      compressed = ZstdCompressor(write_content_size=write_content_size).compress(firmware)
+      with self.subTest(write_content_size=write_content_size), patch.object(pathlib.Path, "is_file", return_value=True), \
+           patch.object(pathlib.Path, "open", return_value=io.BytesIO(compressed)), \
+           patch("tinygrad.helpers.fetch", side_effect=AssertionError("network fallback")):
+        self.assertEqual(fetch_fw("amdgpu", "fixture.bin", digest), firmware)
+
+  def test_local_zstd_wrong_hash_uses_pinned_fallback(self):
+    from zstandard import ZstdCompressor
+    compressed = ZstdCompressor(write_content_size=False).compress(b"wrong firmware")
+    with patch.object(pathlib.Path, "is_file", return_value=True), patch.object(pathlib.Path, "open", return_value=io.BytesIO(compressed)), \
+         patch("tinygrad.helpers.fetch", side_effect=RuntimeError("hash mismatch fallback")) as fallback:
+      with self.assertRaisesRegex(RuntimeError, "hash mismatch fallback"):
+        fetch_fw("amdgpu", "fixture.bin", hashlib.sha256(b"expected firmware").hexdigest())
+    fallback.assert_called_once_with(
+      "https://gitlab.com/kernel-firmware/linux-firmware/-/raw/0a6871b19abf5d6e024b5d208b101ae53e7fa0de/amdgpu/fixture.bin",
+      subdir="fw", sha256=hashlib.sha256(b"expected firmware").hexdigest())
+
 if __name__ == '__main__':
   unittest.main()

tinygrad_repo/tinygrad/helpers.py

@@ -508,9 +508,14 @@ def fetch(url:str, name:pathlib.Path|str|None=None, subdir:str|None=None, gunzip
 def fetch_fw(path:str, name:str, sha256:str) -> bytes:
   if (p:=pathlib.Path(f"/lib/firmware/{path}/{name}.zst")).is_file():
     try:
-      if sys.version_info >= (3,14): from compression.zstd import decompress
-      else: from zstandard import decompress
-      if hashlib.sha256(b:=decompress(p.read_bytes())).hexdigest() == sha256: return b
+      if sys.version_info >= (3,14):
+        from compression.zstd import decompress
+        b = decompress(p.read_bytes())
+      else:
+        # Distro firmware frames can omit content size, which decompress(bytes) rejects.
+        from zstandard import ZstdDecompressor
+        with p.open("rb") as f, ZstdDecompressor().stream_reader(f) as reader: b = reader.read()
+      if hashlib.sha256(b).hexdigest() == sha256: return b
     except ImportError: pass
   return fetch(f"https://gitlab.com/kernel-firmware/linux-firmware/-/raw/0a6871b19abf5d6e024b5d208b101ae53e7fa0de/{path}/{name}",
                subdir="fw", sha256=sha256).read_bytes()

tinygrad_repo/tinygrad/runtime/ops_qcom.py

@@ -281,10 +281,13 @@ class QCOMProgramData:
     reg_desc_off = _read_lib(lib, 0x34)
     self.fregs, self.hregs = _read_lib(lib, reg_desc_off + 0x14), _read_lib(lib, reg_desc_off + 0x18)
 
-_qcom_program_cache:dict[tuple[bytes, tuple[str, ...]], tuple[QCOMProgramData, UOp]] = {}
+_qcom_program_cache:dict[tuple[Any, ...], tuple[QCOMProgramData, UOp]] = {}
 def qcom_build_program(dev:QCOMDevice, prg:UOp, devs:tuple[str, ...]) -> tuple[QCOMProgramData, UOp]:
-  if (cached:=_qcom_program_cache.get(key:=(prg.src[3].arg, devs))) is None:
-    data = QCOMProgramData(dev, prg.to_elf())
+  elf = prg.to_elf()
+  # The same binary can be called with different argument signatures. Keep
+  # each descriptor layout separate while reusing identical programs.
+  if (cached:=_qcom_program_cache.get(key:=(elf.lib, devs, elf.name, elf.target, elf.signature))) is None:
+    data = QCOMProgramData(dev, elf)
     image = bytes(data.image).ljust(round_up(len(data.image), 4), b"\x00")
     buf = UOp.placeholder((len(image),), dtypes.uint8, next(UOp.unique_num), device=devs).rtag("program")
     cached = _qcom_program_cache[key] = (data, patch(buf, [], image))

tinygrad_repo/tinygrad/runtime/support/am/amdev.py

@@ -7,6 +7,7 @@ from tinygrad.runtime.support.amd import AMDReg, import_module, import_asic_regs
 from tinygrad.runtime.support.memory import MemoryManager, AddrSpace
 from tinygrad.runtime.support.system import PCIDevice
 from tinygrad.runtime.support.am.ip import AM_IP, AM_SOC, AM_GMC, AM_IH, AM_PSP, AM_SMU, AM_GFX, AM_SDMA
+from tinygrad.runtime.support.am.startup_trace import initialization, note, value
 
 AM_DEBUG = getenv("AM_DEBUG", 0)
 
@@ -153,6 +154,7 @@ class AMDev:
       cap = self.pci_dev.read_config(cap + 1, 1) & 0xfc
     if cap and cap not in seen: self.pci_dev.write_config_flush(cap + 0x10, self.pci_dev.read_config(cap + 0x10, 2) & ~3, 2) # PCIe cap lnkctl
 
+  @initialization
   def __init__(self, pci_dev:PCIDevice, reset_mode=False):
     self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
     self._disable_aspm()
@@ -166,6 +168,7 @@ class AMDev:
 
     self._run_discovery()
     self._build_regs()
+    note(self, "discovery", gc=self.ip_ver[am.GC_HWIP], mp0=self.ip_ver[am.MP0_HWIP], mp1=self.ip_ver[am.MP1_HWIP], is_vf=self.is_vf)
 
     # AM boot Process:
     # The GPU being passed can be in one of several states: 1. Not initialized. 2. Initialized by amdgpu. 3. Initialized by AM.
@@ -180,25 +183,33 @@ class AMDev:
     self.is_booting = True # During boot only boot memory can be allocated. This flag is to validate this.
     self.init_sw(smi_dev=False)
 
-    self.partial_boot = (self.reg("regSCRATCH_REG7").read() == AMDev.Version) and (getenv("AM_RESET", 0) != 1)
-    if self.partial_boot and (self.reg("regSCRATCH_REG6").read() != 0 or self.reg(self.gmc.pf_status_reg("GC")).read() != 0):
+    self.partial_boot = (value(self, "scratch7", self.reg("regSCRATCH_REG7").read()) == AMDev.Version) and (getenv("AM_RESET", 0) != 1)
+    if self.partial_boot and (value(self, "scratch6", self.reg("regSCRATCH_REG6").read()) != 0 or
+                              value(self, "gc_fault", self.reg(self.gmc.pf_status_reg("GC")).read()) != 0):
       if DEBUG >= 2: print(f"am {self.devfmt}: Malformed state. Issuing a full reset.")
       self.partial_boot = False
 
     # aqua (gc 9.5.0): full boot over live state can kill the fabric (power cycle recovers); partial boot+reset_mec is the deepest safe reset
-    if self.ip_ver[am.GC_HWIP] == (9,5,0) and self.reg("regSCRATCH_REG7").read() == AMDev.Version: self.partial_boot = True
+    if self.ip_ver[am.GC_HWIP] == (9,5,0) and value(self, "scratch7", self.reg("regSCRATCH_REG7").read()) == AMDev.Version: self.partial_boot = True
 
+    note(self, "boot_choice", partial_boot=self.partial_boot)
     # Init hw for IP blocks where it is needed
     if not self.partial_boot:
-      if not self.is_vf and self.psp.is_sos_alive() and self.smu.is_smu_alive(): # skip in vf mode, these are pf funcs.
+      if (not self.is_vf and value(self, "psp_alive", self.psp.is_sos_alive()) and
+          value(self, "smu_alive", self.smu.is_smu_alive())): # skip in vf mode, these are pf funcs.
         self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) & ~pci.PCI_COMMAND_MASTER, 2)
         if self.is_hive():
           if reset_mode: return # in reset mode, do not raise
           raise RuntimeError("Malformed state. Use extra/amdpci/hive_reset.py to reset the hive")
+        note(self, "mode1_reset_enter")
         self.smu.mode1_reset()
+        note(self, "mode1_reset_exit")
       self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
       self.init_hw(self.soc, self.gmc, self.ih, *(() if self.is_vf else (self.psp, self.smu)))
-    elif not self.is_vf: self.psp._tmr_init()
+    elif not self.is_vf:
+      note(self, "partial_tmr_enter")
+      self.psp._tmr_init()
+      note(self, "partial_tmr_exit")
 
     # Booting done
     self.is_booting = False
@@ -241,7 +252,9 @@ class AMDev:
 
   def init_hw(self, *blocks:AM_IP):
     for ip in blocks:
+      note(self, "ip_enter", block=ip.__class__.__name__)
       ip.init_hw()
+      note(self, "ip_exit", block=ip.__class__.__name__)
       if DEBUG >= 2: print(f"am {self.devfmt}: {ip.__class__.__name__} initialized")
 
   def fini(self):

tinygrad_repo/tinygrad/runtime/support/am/ip.py

@@ -5,6 +5,8 @@ from tinygrad.runtime.autogen.am import am, pm4_soc15 as pm4
 from tinygrad.runtime.support.amd import import_soc
 from tinygrad.runtime.support.memory import AddrSpace
 
+from tinygrad.runtime.support.am.startup_trace import value, wait_tlb
+
 class AM_IP:
   def __init__(self, adev): self.adev = adev
   def init_sw(self): pass # Prepare sw/allocations for this IP
@@ -118,7 +120,7 @@ class AM_GMC(AM_IP):
 
       self.adev.reg(f"reg{ip}VM_INVALIDATE_ENG17_REQ").write(req, inst=inst)
 
-      wait_cond(lambda: self.adev.reg(f"reg{ip}VM_INVALIDATE_ENG17_ACK").read(inst=inst) & (1 << vmid), value=(1 << vmid), msg="flush_tlb timeout")
+      wait_tlb(self.adev, ip, inst, req, vmid, wait_cond, lambda: self.adev.reg(f"reg{ip}VM_INVALIDATE_ENG17_ACK").read(inst=inst))
 
       if ip == "MM": self.adev.regMMVM_INVALIDATE_ENG17_SEM.write(0x0, inst=inst)
       if self.adev.ip_ver[am.GC_HWIP] >= (11,0,0) and ip == "MM":
@@ -275,7 +277,7 @@ class AM_GFX(AM_IP):
 
   def init_hw(self):
     # Wait for RLC autoload to complete
-    wait_cond(lambda: self.adev.regCP_STAT.read() == 0 or self.adev.regRLC_RLCS_BOOTLOAD_STATUS.read_bitfields()['bootload_complete'] == 0,
+    wait_cond(lambda: value(self.adev, "cp_stat", self.adev.regCP_STAT.read()) == 0 or value(self.adev, "bootload_complete", self.adev.regRLC_RLCS_BOOTLOAD_STATUS.read_bitfields()['bootload_complete']) == 0,
               value=True, msg="RLC autoload timeout")
 
     self.adev.gmc.init_hub("GC", insts=range(self.xccs))

tinygrad_repo/tinygrad/runtime/support/am/startup_trace.py

@@ -0,0 +1,92 @@
+import functools, json, os, time
+
+ENABLED = os.environ.get("STARPILOT_GPU_STARTUP_TRACE") == "1"
+REPORT = None
+
+def report(packet):
+  if REPORT is not None: REPORT(packet)
+  else: os.write(2, (json.dumps(packet, separators=(",", ":")) + "\n").encode())
+
+def note(dev, event, **fields):
+  if not ENABLED: return
+  try:
+    entries = dev.__dict__.setdefault("_startup_trace_entries", [])
+    if len(entries) < 64:
+      entries.append({"mono_ns": time.monotonic_ns(), "event": event, **fields})
+  except Exception:
+    pass
+
+def value(dev, name, result):
+  if ENABLED:
+    try:
+      values = dev.__dict__.setdefault("_startup_trace_values", {})
+      entries = dev.__dict__.setdefault("_startup_trace_entries", [])
+      if name not in values and len(entries) < 64:
+        values[name] = {"mono_ns": time.monotonic_ns(), "event": name, "first": result, "last": result, "reads": 0}
+        entries.append(values[name])
+      if name in values:
+        values[name]["last"] = result
+        values[name]["reads"] += 1
+    except Exception:
+      pass
+  return result
+
+def emit(dev, outcome):
+  if not ENABLED: return
+  try:
+    entries = dev.__dict__.pop("_startup_trace_entries", [])
+    packet = {"kind": "starpilot_gpu_startup", "pid": os.getpid(), "device": getattr(dev, "devfmt", None),
+              "outcome": outcome, "entries": entries}
+    report(packet)
+  except Exception:
+    pass
+
+_short_reads = 0
+
+def short_read(address, expected, actual):
+  global _short_reads
+  if not ENABLED or _short_reads >= 8: return
+  _short_reads += 1
+  try:
+    packet = {"kind": "starpilot_gpu_short_mmio", "pid": os.getpid(), "mono_ns": time.monotonic_ns(),
+              "address": address, "expected": expected, "actual": actual, "sample": _short_reads}
+    report(packet)
+  except Exception:
+    pass
+
+def initialization(fn):
+  @functools.wraps(fn)
+  def call(dev, *args, **kwargs):
+    note(dev, "initialization_enter")
+    try:
+      result = fn(dev, *args, **kwargs)
+    except BaseException as exc:
+      note(dev, "initialization_failed", exception=type(exc).__name__)
+      emit(dev, "initialization_failed")
+      raise
+    note(dev, "initialization_complete")
+    return result
+  return call
+
+def wait_tlb(dev, ip, inst, req, vmid, wait, read_ack):
+  mask = 1 << vmid
+  if not ENABLED or ip != "GC" or dev.__dict__.get("_startup_trace_gc_seen", False):
+    return wait(lambda: read_ack() & mask, value=mask, msg="flush_tlb timeout")
+  dev._startup_trace_gc_seen = True
+  note(dev, "first_gc_request", instance=inst, request=req, vmid=vmid)
+  sample = {"reads": 0, "first": None, "last": None}
+  def observe():
+    raw = read_ack()
+    sample["reads"] += 1
+    if sample["first"] is None: sample["first"] = raw
+    sample["last"] = raw
+    return raw & mask
+  try:
+    result = wait(observe, value=mask, msg="flush_tlb timeout")
+  except BaseException as exc:
+    note(dev, "first_gc_failure", exception=type(exc).__name__, **sample)
+    emit(dev, "first_gc_failure")
+    raise
+  note(dev, "first_gc_acknowledged", **sample)
+  emit(dev, "first_gc_acknowledged")
+  return result

tinygrad_repo/tinygrad/runtime/support/usb.py

@@ -9,6 +9,7 @@ from tinygrad.device import Buffer, BufferSpec
 from tinygrad.runtime.support.hcq2 import HCQ_RUNTIME_DEV, HCQ_DEVS, ccall, cfield, patch, rt_addr, unwrap_view, all_devices_in
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support import c
+from tinygrad.runtime.support.am.startup_trace import short_read
 
 def alloc_cbuffer(sz:int) -> tuple[ctypes.Array, memoryview]: return (buf:=(ctypes.c_ubyte * sz)()), to_mv(ctypes.addressof(buf), sz)
 def checked(fn, msg=None):
@@ -232,6 +233,7 @@ class USBMMIOInterface(MMIOInterface):
     if self.pcimem:
       assert sz % 4 == 0 and off % 4 == 0, f"pcie_mem_read requires 4-byte aligned access, got off={off}, sz={sz}"
       data = self.usb.pcie_mem_read(self.addr + off, sz)
+      if len(data) != sz: short_read(self.addr + off, sz, len(data))
     else: data = self.usb.scsi_read(sz) if self.addr == 0xf000 else self.usb.read(self.addr + off, sz)
     return data if isinstance(index, slice) else int.from_bytes(data, "little")