mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 13:33:42 +08:00
firestar5683
8b254da432
Extend vendored runtime helpers and diagnostics while preserving the existing compiler and model ABI.
123 lines
4.4 KiB
Python
123 lines
4.4 KiB
Python
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()
|