mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 05:23:41 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
91 lines
4.6 KiB
Python
91 lines
4.6 KiB
Python
import unittest
|
|
from tinygrad.helpers import Timing, getenv
|
|
from tinygrad import Tensor, Device, TinyJit
|
|
from tinygrad.runtime.support.usb import HALF, CHUNK, SLOT
|
|
import numpy as np
|
|
|
|
class USBTestCase(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.sz = getenv("SIZE", 2000000)
|
|
cls.dev = Device["AMD"]
|
|
if not cls.dev.is_usb: raise unittest.SkipTest("only test this on USB devices")
|
|
cls.rng = np.random.default_rng(0)
|
|
|
|
def roundtrip(self, a:np.ndarray): # a copy in, a kernel, a copy out: the queue must order them
|
|
np.testing.assert_array_equal(a, Tensor(a, device="NPY").to(Device.DEFAULT).numpy())
|
|
np.testing.assert_array_equal(a + 1, (Tensor(a, device="NPY").to(Device.DEFAULT) + 1).numpy())
|
|
|
|
class TestDevCopySpeeds(USBTestCase):
|
|
def testCopyCPUtoDefault(self):
|
|
for _ in range(10):
|
|
t = Tensor.ones(self.sz, device="CPU", dtype='uchar').contiguous().realize()
|
|
with Timing(f"copyin of {t.nbytes()/1e6:.2f} MB: ", on_exit=lambda ns: f" @ {t.nbytes()/ns * 1e3:.2f} MB/s"): # noqa: F821
|
|
t.to(Device.DEFAULT).realize()
|
|
Device[Device.DEFAULT].synchronize()
|
|
del t
|
|
|
|
def testCopyDefaulttoCPU(self):
|
|
t = Tensor.ones(self.sz, dtype='uchar').contiguous().realize()
|
|
for _ in range(10):
|
|
with Timing(f"copyout of {t.nbytes()/1e6:.2f} MB: ", on_exit=lambda ns: f" @ {t.nbytes()/ns * 1e3:.2f} MB/s"):
|
|
t.to('CPU').realize()
|
|
|
|
class TestUSBIntegrity(USBTestCase):
|
|
def testValidateCopies(self):
|
|
t = Tensor.randn(self.sz, device="CPU", dtype='uchar').contiguous().realize()
|
|
x = t.to(Device.DEFAULT).realize()
|
|
Device[Device.DEFAULT].synchronize()
|
|
y = x.to('CPU').realize()
|
|
np.testing.assert_equal(t.numpy(), y.numpy())
|
|
|
|
def testBoundaries(self): # around the slot, the chunk and the read window
|
|
for size in (1, 3, 508, 509, SLOT - 513, SLOT - 512, SLOT - 511, CHUNK - 1, CHUNK, CHUNK + 1, 2 * CHUNK - 1, 2 * CHUNK, 2 * CHUNK + 31, HALF,
|
|
2 * HALF, 1 << 20):
|
|
with self.subTest(size=size): self.roundtrip(self.rng.integers(0, 256, size, dtype=np.uint8))
|
|
|
|
def testManyCopiesInABatch(self):
|
|
for n in (2, 7, 64, 300): # 300 chunks: the fence byte wraps
|
|
with self.subTest(n=n):
|
|
arrs = [self.rng.integers(0, 256, int(s), dtype=np.uint8) for s in self.rng.integers(1, 5000, n)]
|
|
ts = [Tensor(a, device="NPY").to(Device.DEFAULT) for a in arrs]
|
|
Tensor.realize(*ts)
|
|
for t, a in zip(ts, arrs): np.testing.assert_array_equal(a, t.numpy())
|
|
|
|
def testMixedBatch(self): # copies out and in, in one batch: runs of both directions
|
|
arrs = [self.rng.integers(0, 256, s, dtype=np.uint8) for s in (5, CHUNK + 7, 9, 2 * CHUNK + 3, 11)]
|
|
ts = [Tensor(a, device="NPY").to(Device.DEFAULT).realize() for a in arrs]
|
|
more = [self.rng.integers(0, 256, s, dtype=np.uint8) for s in (5, CHUNK + 7, 9, 2 * CHUNK + 3, 11)]
|
|
outs = [t.to("NPY") for t in ts] + [Tensor(a, device="NPY").to(Device.DEFAULT) for a in more]
|
|
Tensor.realize(*outs)
|
|
for o, a in zip(outs, arrs + more): np.testing.assert_array_equal(a, o.numpy())
|
|
|
|
def testRepeatedBatches(self): # a batch numbers its chunks from 0: the same batch again must not see what the last one left behind
|
|
a = self.rng.integers(0, 256, 2 * CHUNK + 31, dtype=np.uint8)
|
|
for _ in range(5): self.roundtrip(a)
|
|
@TinyJit
|
|
def step(x:Tensor) -> Tensor: return (x + 1).realize()
|
|
src = Tensor(a, device="NPY")
|
|
for i in range(5):
|
|
x = src.to(Device.DEFAULT)
|
|
np.testing.assert_array_equal(a + 1, step(x).numpy())
|
|
|
|
def testStaleSentinel(self): # payloads full of the tags the queue waits for, in both directions, before and around the real chunks
|
|
tags = np.array([0x51000000 | k for k in range(8)], dtype=np.uint32)
|
|
for tag in tags: # every dword of every chunk is the tag of some chunk of the copy
|
|
with self.subTest(payload=hex(tag)):
|
|
a = np.full((2 * CHUNK + 31) // 4, tag, dtype=np.uint32).view(np.uint8)
|
|
self.roundtrip(a)
|
|
with self.subTest(case="copyout residue"): # a read fills the sram with tags, then small chunks land in both halves
|
|
a = np.tile(tags, 2 * CHUNK // 32).view(np.uint8)
|
|
np.testing.assert_array_equal(a, (Tensor(a, device="NPY").to(Device.DEFAULT) * 1).numpy())
|
|
for size in (31, CHUNK + 31, 2 * CHUNK + 31): self.roundtrip(np.tile(tags, size // 32 + 1).view(np.uint8)[:size])
|
|
|
|
def testRingWrap(self): # 64MB of chunks: the sdma ring (1MB on usb) wraps within the copy
|
|
a = self.rng.integers(0, 256, 64 << 20, dtype=np.uint8)
|
|
t = Tensor(a, device="NPY").to(Device.DEFAULT).realize()
|
|
np.testing.assert_array_equal(a, t.numpy())
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|