StarPilot/tinygrad_repo/test/external/external_test_usb_asm24.py

import unittest
from tinygrad.helpers import Timing, getenv, unwrap, wait_cond
from tinygrad import Tensor, Device, TinyJit, dtypes
from tinygrad.device import Buffer
from tinygrad.engine.realize import run_linear
from tinygrad.uop.ops import UOp
from tinygrad.runtime.ops_amd import AMDSDMAQueue
from tinygrad.runtime.support.hcq2 import make_submit
from tinygrad.runtime.support.usb import HALF, CHUNK, SLOT
from test.null.test_hcq2 import lower_hcq
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):
    self.dev.synchronize()
    q = unwrap(self.dev.sdma_queue(0))
    put, read = [b.host.view(fmt='Q') for b in (q.put_value, q.read_ptr)]
    ring_size = q.ring.nbytes
    # advance with NOPs, leave one dword so the next eager copy forces wrapping
    target = (put[0] // ring_size + 1) * ring_size - 4
    if nbytes := target - put[0]:
      nops = UOp.from_buffer(Buffer("CPU", nbytes, initial_value=bytes(nbytes)), dtypes.uint8)
      hq = AMDSDMAQueue(make_submit(dev=self.dev.device, queue="COPY:0", deps=(nops,)))
      run_linear(lower_hcq(hq.submit(nops)))
    wait_cond(lambda: read[0], value=target)
    self.assertEqual(put[0], target)

    a = self.rng.integers(0, 256, 64 << 20, dtype=np.uint8)
    t = Tensor(a, device="NPY").to(Device.DEFAULT).realize()
    self.assertGreater(put[0] // ring_size, target // ring_size)
    np.testing.assert_array_equal(a, t.numpy())
    self.dev.synchronize()
    wait_cond(lambda: read[0], value=put[0])

if __name__ == "__main__":
  unittest.main()