mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-29 13:16:09 +00:00
116 lines
5.7 KiB
Python
116 lines
5.7 KiB
Python
import struct, unittest
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
from tinygrad.runtime.autogen import bnxt
|
|
from extra.bnxt_driver.bnxtdev import BNXT_BACKING_STORE, BNXTDev, BNXTQP, _queue, _qwrite, ipv4_to_gid
|
|
|
|
class FakePCI:
|
|
def __init__(self): self.next_addr, self.allocations = 0x100000, []
|
|
def alloc_sysmem(self, size, contiguous=False):
|
|
pages = [self.next_addr+i*0x1000 for i in range((size+0xfff)//0x1000)]
|
|
self.next_addr += len(pages)*0x1000
|
|
self.allocations.append(mem := bytearray(size))
|
|
return mem, pages
|
|
|
|
class FakeDev:
|
|
def __init__(self): self.pci_dev, self.calls = FakePCI(), []
|
|
def hwrm(self, name, **fields):
|
|
self.calls.append((name, fields))
|
|
typ = fields.get("type", 0)
|
|
return SimpleNamespace(ctx_init_value=0x5a, ctx_init_offset=4, entry_size=16 if typ == 0 else 4,
|
|
subtype_valid_cnt=typ == 0, split_entry_0=2, instance_bit_map=5 if typ == 0 else 1, min_num_entries=0)
|
|
|
|
class FakeRCFW:
|
|
def __init__(self): self.calls, self.doorbells = [], []
|
|
def exec(self, name, **fields):
|
|
self.calls.append((name, fields))
|
|
return SimpleNamespace(xid={"create_cq":77, "create_qp":88, "register_mr":0x5678}.get(name, 0))
|
|
def doorbell(self, *args, **kwargs): self.doorbells.append((args, kwargs))
|
|
|
|
class FakeQPDev:
|
|
def __init__(self): self.pci_dev, self.fw, self.gid_id, self.nq_id = FakePCI(), FakeRCFW(), 9, 41
|
|
def rcfw(self, *args, **kwargs): return self.fw.exec(*args, **kwargs)
|
|
def doorbell(self, *args, **kwargs): self.fw.doorbell(*args, **kwargs)
|
|
|
|
class TestMemory(unittest.TestCase):
|
|
def test_cmdq_and_sq_aux(self):
|
|
dev = FakeDev()
|
|
cmdq, sq = _queue(dev), _queue(dev, aux=True)
|
|
self.assertEqual((cmdq["level"], cmdq["base"]), (0, 0x100000))
|
|
_qwrite(sq, 3, b"ABCDEFGH", aux=True)
|
|
self.assertEqual(bytes(sq["mem"][0x1018:0x1020]), b"ABCDEFGH")
|
|
|
|
def test_f320_backing_layout_and_final_marker(self):
|
|
self.assertEqual(len(BNXT_BACKING_STORE), 9)
|
|
dev = FakeDev()
|
|
small = ((0, 6), (15, 0))
|
|
with patch("extra.bnxt_driver.bnxtdev.BNXT_BACKING_STORE", small): BNXTDev.setup_backing_store(dev)
|
|
cfg = [fields for name, fields in dev.calls if name == "func_backing_store_cfg_v2"]
|
|
self.assertEqual([(x["type"], x["instance"]) for x in cfg], [(0, 0), (0, 2), (15, 0)])
|
|
self.assertTrue(all(not x["flags"] for x in cfg[:-1]))
|
|
self.assertEqual(cfg[-1]["flags"], bnxt.FUNC_BACKING_STORE_CFG_V2_REQ_FLAGS_BS_CFG_ALL_DONE)
|
|
self.assertEqual((dev.pci_dev.allocations[0][4], dev.pci_dev.allocations[0][20]), (0x5a, 0x5a))
|
|
|
|
class TestRCFW(unittest.TestCase):
|
|
def setUp(self):
|
|
patch("extra.bnxt_driver.bnxtdev.System.memory_barrier").start()
|
|
self.addCleanup(patch.stopall)
|
|
|
|
def test_doorbell_encodes_xid_type_and_index(self):
|
|
dev = BNXTDev.__new__(BNXTDev)
|
|
dev.db, dev.db_off = [0]*1024, 0x1000
|
|
dev.doorbell(0x123456, bnxt.DBC_DBC_TYPE_CQ_ARMALL, 0x456, epoch=1)
|
|
key = dev.db[0x1000//8]
|
|
self.assertEqual(key >> 32,
|
|
0x123456 & bnxt.DBC_DBC_XID_MASK | bnxt.DBC_DBC_PATH_ROCE | bnxt.DBC_DBC_TYPE_CQ_ARMALL | bnxt.BNXT_QPLIB_DBR_VALID)
|
|
self.assertEqual(key & 0xffffffff, 0x456 | 1<<bnxt.BNXT_QPLIB_DBR_EPOCH_SHIFT)
|
|
|
|
def test_command_uses_first_flag(self):
|
|
dev = BNXTDev.__new__(BNXTDev)
|
|
dev.bar0, dev.cmdq, dev.creq = [0]*1024, _queue(FakeDev()), _queue(FakeDev())
|
|
dev.rcfw_first, dev.creq_id = True, 23
|
|
dev.doorbell = lambda *args: None
|
|
_qwrite(dev.creq, 0, bytes(bnxt.struct_creq_query_version_resp(type=bnxt.CREQ_BASE_TYPE_QP_EVENT, cookie=0, v=1)))
|
|
ret = dev.rcfw("query_version")
|
|
req = bnxt.struct_cmdq_query_version.from_buffer_copy(bytes(dev.cmdq["mem"][:16]))
|
|
prod = dev.bar0[(bnxt.RCFW_COMM_BASE_OFFSET+bnxt.RCFW_PF_VF_COMM_PROD_OFFSET)//4]
|
|
self.assertEqual((req.cookie, ret.cookie, prod), (0, 0, 1 | 1<<bnxt.FIRMWARE_FIRST_FLAG))
|
|
|
|
class TestFastPath(unittest.TestCase):
|
|
def test_unified_mr(self):
|
|
dev = BNXTDev.__new__(BNXTDev)
|
|
fw = FakeRCFW()
|
|
dev.pci_dev, dev.rcfw = FakePCI(), fw.exec
|
|
self.assertEqual(dev.register_mem([0x800000, 0x900000], 0x2000), 0x5678)
|
|
mr = fw.calls[-1][1]
|
|
self.assertEqual((mr["flags"], mr["va"], mr["mr_size"], mr["log2_pg_size_lvl"]),
|
|
(bnxt.CMDQ_REGISTER_MR_FLAGS_ALLOC_MR, 0x800000, 0x2000,
|
|
1<<bnxt.CMDQ_REGISTER_MR_LVL_SFT | 12<<bnxt.CMDQ_REGISTER_MR_LOG2_PG_SIZE_SFT))
|
|
|
|
def test_qp_creation_and_connect_use_f320_layout(self):
|
|
dev = FakeQPDev()
|
|
qp = BNXTQP(dev)
|
|
create = next(fields for name, fields in dev.fw.calls if name == "create_qp")
|
|
self.assertEqual((create["sq_size"], "rq_size" in create, qp.qpn), (16, False, 88))
|
|
qp.connect(0x123, ipv4_to_gid("10.0.0.2"), 0x001122334455)
|
|
rtr, rts = dev.fw.calls[-2][1], dev.fw.calls[-1][1]
|
|
self.assertEqual((bytes(rtr["dgid"]), bytes(rtr["dest_mac"])),
|
|
(ipv4_to_gid("10.0.0.2"), bytes.fromhex("001122334455")))
|
|
self.assertEqual((rtr["modify_mask"], rts["modify_mask"]), (0x41515ad, 0xae005))
|
|
|
|
def test_rdma_write_builds_three_slots_and_host_msn(self):
|
|
qp = BNXTQP.__new__(BNXTQP)
|
|
qp.dev, qp.qpn = FakeQPDev(), 88
|
|
qp.sq, qp.sq_psn, qp.msn = _queue(FakeDev(), aux=True), 5, 0
|
|
qp._poll = lambda timeout: bytes(bnxt.struct_cq_req())
|
|
qp.rdma_write(0x1122334455667788, 0x99aa, 0x12345000, 0x55aa, 100)
|
|
hdr = bnxt.struct_sq_rdma_hdr.from_buffer_copy(bytes(qp.sq["mem"][:32]))
|
|
sge = bnxt.struct_sq_sge.from_buffer_copy(bytes(qp.sq["mem"][32:48]))
|
|
self.assertEqual((hdr.remote_va, hdr.remote_key, hdr.length, sge.va_or_pa, sge.l_key, sge.size),
|
|
(0x1122334455667788, 0x99aa, 100, 0x12345000, 0x55aa, 100))
|
|
self.assertEqual(struct.unpack_from("<Q", qp.sq["mem"], 0x1000)[0], 6<<24 | 5)
|
|
self.assertEqual(qp.dev.fw.doorbells, [((88, bnxt.DBC_DBC_TYPE_SQ, 3, 0), {})])
|
|
|
|
if __name__ == "__main__": unittest.main()
|