From 8b254da432452faa42697e4e31b3f3707a1fd981 Mon Sep 17 00:00:00 2001 From: firestar5683 <168790843+firestar5683@users.noreply.github.com> Date: Fri, 25 Sep 2026 14:53:21 -0500 Subject: [PATCH] Runtime Dependencies Extend vendored runtime helpers and diagnostics while preserving the existing compiler and model ABI. --- msgq_repo/msgq/msgq.h | 2 +- msgq_repo/msgq/msgq_tests.cc | 56 ++++++++ .../msgq/visionipc/tests/test_visionipc.py | 23 ++-- .../models/tests/test_helper_import.py | 1 + .../models/tests/test_startup_trace.py | 122 ++++++++++++++++++ tinygrad_repo/LOCAL_PATCHES.md | 9 ++ tinygrad_repo/examples/openpilot/helpers.py | 17 +-- tinygrad_repo/test/device/test_qcom.py | 31 +++++ tinygrad_repo/test/unit/test_helpers.py | 28 +++- tinygrad_repo/tinygrad/helpers.py | 11 +- tinygrad_repo/tinygrad/runtime/ops_qcom.py | 9 +- .../tinygrad/runtime/support/am/amdev.py | 23 +++- .../tinygrad/runtime/support/am/ip.py | 6 +- .../runtime/support/am/startup_trace.py | 92 +++++++++++++ tinygrad_repo/tinygrad/runtime/support/usb.py | 2 + 15 files changed, 398 insertions(+), 34 deletions(-) create mode 100644 openpilot/starpilot/models/tests/test_startup_trace.py create mode 100644 tinygrad_repo/LOCAL_PATCHES.md create mode 100644 tinygrad_repo/tinygrad/runtime/support/am/startup_trace.py diff --git a/msgq_repo/msgq/msgq.h b/msgq_repo/msgq/msgq.h index 1160f66805..6f64406c8b 100644 --- a/msgq_repo/msgq/msgq.h +++ b/msgq_repo/msgq/msgq.h @@ -6,7 +6,7 @@ #include #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 diff --git a/msgq_repo/msgq/msgq_tests.cc b/msgq_repo/msgq/msgq_tests.cc index b8d6824da2..056b941d3b 100644 --- a/msgq_repo/msgq/msgq_tests.cc +++ b/msgq_repo/msgq/msgq_tests.cc @@ -1,3 +1,5 @@ +#include +#include #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 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(&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()); +} diff --git a/msgq_repo/msgq/visionipc/tests/test_visionipc.py b/msgq_repo/msgq/visionipc/tests/test_visionipc.py index b938df4472..40966b4f27 100644 --- a/msgq_repo/msgq/visionipc/tests/test_visionipc.py +++ b/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 diff --git a/openpilot/starpilot/models/tests/test_helper_import.py b/openpilot/starpilot/models/tests/test_helper_import.py index e029ceab50..367b85b563 100644 --- a/openpilot/starpilot/models/tests/test_helper_import.py +++ b/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. diff --git a/openpilot/starpilot/models/tests/test_startup_trace.py b/openpilot/starpilot/models/tests/test_startup_trace.py new file mode 100644 index 0000000000..86fd95edde --- /dev/null +++ b/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() diff --git a/tinygrad_repo/LOCAL_PATCHES.md b/tinygrad_repo/LOCAL_PATCHES.md new file mode 100644 index 0000000000..cff3648c2a --- /dev/null +++ b/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. diff --git a/tinygrad_repo/examples/openpilot/helpers.py b/tinygrad_repo/examples/openpilot/helpers.py index ae34b8ffc2..7432b34f9a 100644 --- a/tinygrad_repo/examples/openpilot/helpers.py +++ b/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) diff --git a/tinygrad_repo/test/device/test_qcom.py b/tinygrad_repo/test/device/test_qcom.py index 381f1c1402..7560815563 100644 --- a/tinygrad_repo/test/device/test_qcom.py +++ b/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") diff --git a/tinygrad_repo/test/unit/test_helpers.py b/tinygrad_repo/test/unit/test_helpers.py index 2b93834958..080c54cc7c 100644 --- a/tinygrad_repo/test/unit/test_helpers.py +++ b/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() diff --git a/tinygrad_repo/tinygrad/helpers.py b/tinygrad_repo/tinygrad/helpers.py index a4f864dc59..541a37ee91 100644 --- a/tinygrad_repo/tinygrad/helpers.py +++ b/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() diff --git a/tinygrad_repo/tinygrad/runtime/ops_qcom.py b/tinygrad_repo/tinygrad/runtime/ops_qcom.py index 671e041a10..be88c16e0b 100644 --- a/tinygrad_repo/tinygrad/runtime/ops_qcom.py +++ b/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)) diff --git a/tinygrad_repo/tinygrad/runtime/support/am/amdev.py b/tinygrad_repo/tinygrad/runtime/support/am/amdev.py index 3105d5ba02..0245a4d32b 100644 --- a/tinygrad_repo/tinygrad/runtime/support/am/amdev.py +++ b/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): diff --git a/tinygrad_repo/tinygrad/runtime/support/am/ip.py b/tinygrad_repo/tinygrad/runtime/support/am/ip.py index d42821d363..cff46fd4ba 100644 --- a/tinygrad_repo/tinygrad/runtime/support/am/ip.py +++ b/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)) diff --git a/tinygrad_repo/tinygrad/runtime/support/am/startup_trace.py b/tinygrad_repo/tinygrad/runtime/support/am/startup_trace.py new file mode 100644 index 0000000000..ebc0edce4d --- /dev/null +++ b/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 diff --git a/tinygrad_repo/tinygrad/runtime/support/usb.py b/tinygrad_repo/tinygrad/runtime/support/usb.py index a2b222d943..9c0af13295 100644 --- a/tinygrad_repo/tinygrad/runtime/support/usb.py +++ b/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")