mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 13:33:42 +08:00
Daniel Koepping
224 lines
12 KiB
Python
224 lines
12 KiB
Python
import unittest
|
|
from typing import cast
|
|
from types import SimpleNamespace
|
|
from collections import defaultdict
|
|
from unittest.mock import patch
|
|
from tinygrad import Device, Tensor, TinyJit, dtypes
|
|
from tinygrad.device import Buffer, Compiled, ProfileGraphEvent
|
|
from tinygrad.helpers import Context, unwrap, to_tuple
|
|
from tinygrad.dtype import AddrSpace
|
|
from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, KernelInfo
|
|
from tinygrad.engine.realize import compile_linear, link_linear, get_call_arg_uops, lower_and_compile, run_linear
|
|
import tinygrad.runtime.support.hcq2 as hcq2
|
|
from tinygrad.runtime.support.hcq2 import HCQInfo
|
|
|
|
def lower_hcq(*body:UOp) -> UOp: # body through lower_call, as a linear
|
|
return UOp(Ops.LINEAR, src=(unwrap(hcq2.lower_call(UOp.sink(*body, arg=KernelInfo("test"), tag=1).call(aux=HCQInfo(("CPU",))))),))
|
|
|
|
def chain(x:Tensor, n:int) -> Tensor:
|
|
for _ in range(n): x = (x + 1).contiguous()
|
|
return x
|
|
|
|
def chain_input(value:int=2, device="NULL") -> Tensor: return Tensor.full((4,), value, dtype=dtypes.int32, device=device).contiguous().realize()
|
|
|
|
def compiled_chain(n:int, jit=False, device="NULL") -> tuple[Tensor, UOp, list[UOp]]:
|
|
x, inputs = chain_input(device=device), []
|
|
if jit:
|
|
f = TinyJit(lambda a: chain(a, n).realize())
|
|
f(x)
|
|
return f(x), f.captured._linear, [x.uop.base]
|
|
out = chain(x, n)
|
|
return out, compile_linear(out.schedule_linear(), input_uops=inputs, cache=True), inputs
|
|
|
|
# NULL never runs a batch, so the scheduler is tested on the commands it hands each queue, run by a small executor with symbolic
|
|
# signals and timelines. the fence and the ffi are tested by running them on CPU, see test/runtime/test_hcq2.py.
|
|
|
|
def scheduled(*ts:Tensor, **kwargs) -> list[UOp]:
|
|
batches, orig = list[UOp](), hcq2.sched_batches
|
|
def track(l, profile):
|
|
lin = orig(l, profile)
|
|
batches.extend(c for c in lin.src if c.op is Ops.CALL and isinstance(c.arg.aux, HCQInfo))
|
|
return lin
|
|
with patch.object(hcq2, "sched_batches", track): compile_linear(ts[0].schedule_linear(*ts[1:]), **kwargs)
|
|
return batches
|
|
|
|
def queues(batch:UOp) -> dict[tuple[str, str], list[UOp]]:
|
|
return {(lin.arg[0][0], lin.arg[1]): list(lin.src) for lin in (s.src[1].without_after for s in batch.body.src)}
|
|
def calls(batch:UOp) -> list[UOp]: return [c for cmds in queues(batch).values() for c in cmds if c.op is Ops.CALL]
|
|
def devices_of(call:UOp) -> set[str]: return {to_tuple(a.device)[0] for a in get_call_arg_uops(call)}
|
|
|
|
def word(u:UOp) -> tuple[UOp, int]:
|
|
if u.op is Ops.INDEX: return (v:=hcq2.unwrap_view(u.src[0]))[0], v[1] + u.src[1].val * u.dtype.itemsize
|
|
return hcq2.unwrap_view(u)
|
|
|
|
T = 5 # the timeline value the last submitted batch of every device signals
|
|
def run(batch:UOp, done:dict[str, int]|None=None, prio:list|None=None) -> tuple[list, dict, dict]:
|
|
qs, cs, mem = queues(batch), calls(batch), defaultdict(int)
|
|
for d in batch.arg.aux.device:
|
|
mem[word(hcq2.timeline((d,)))] = (done or {}).get(d, T)
|
|
mem[word(hcq2.timeline((d,)).index(1))] = T
|
|
def val(u:UOp) -> int:
|
|
if u.op is Ops.CONST: return u.val
|
|
if u.op is Ops.LOAD: return mem[word(u.src[0])]
|
|
return sum(val(s) for s in u.src)
|
|
log:list[int|str] = [] # the calls that ran, and a device each time its timeline bumps
|
|
while True:
|
|
for q in (list(qs) if prio is None else prio):
|
|
if not (cmds:=qs[q]): continue
|
|
c = cmds[0]
|
|
if c.op is Ops.INS and c.arg[0].startswith("wait"):
|
|
sig, target = mem[word(c.src[0])], val(c.src[1])
|
|
if sig != target if c.arg[0] == "wait_eq" else sig < target: continue
|
|
elif c.op is Ops.INS and c.arg[0] == "store":
|
|
mem[w:=word(c.src[0])] = val(c.src[1])
|
|
if w[0].tag == "timeline": log.append(w[0].device)
|
|
elif c.op is Ops.CALL: log.append(cs.index(c))
|
|
cmds.pop(0)
|
|
break
|
|
else: return log, {q: len(cmds) for q, cmds in qs.items()}, {d: mem[word(hcq2.timeline((d,)))] for d in batch.arg.aux.device}
|
|
|
|
def rotations(batch:UOp) -> list[list]: return [(qs:=list(queues(batch)))[i:] + qs[:i] for i in range(len(queues(batch)))]
|
|
def orders(batch:UOp) -> set[tuple[int, ...]]: return {tuple(x for x in run(batch, prio=p)[0] if isinstance(x, int)) for p in rotations(batch)}
|
|
|
|
class TestHCQ2Deps(unittest.TestCase):
|
|
def test_dependencies_through_selected_slices(self):
|
|
b = UOp.param(0, dtypes.float32, 64, device=("NULL", "NULL:1"))
|
|
for view in [b.mselect(0).shrink(((8, 16),)), b.shrink(((8, 16),)).mselect(0), b.shrink(((4, 32),)).mselect(0).shrink(((4, 12),))]:
|
|
tracker = hcq2.DepsTracker()
|
|
tracker.access_resources([view], [0], 0)
|
|
self.assertEqual(tracker.access_resources([b.mselect(1)], [], 1), [])
|
|
self.assertEqual(tracker.access_resources([b.mselect(0).shrink(((16, 24),))], [], 2), [])
|
|
self.assertEqual(tracker.access_resources([b.mselect(0).shrink(((12, 20),))], [], 3), [0])
|
|
|
|
def test_disjoint_write_preserves_dependencies(self):
|
|
b = UOp.param(0, dtypes.uint8, 16, device="NULL")
|
|
for write in ([], [0]):
|
|
tracker = hcq2.DepsTracker()
|
|
tracker.access_resources([b.shrink(((0, 4),))], write, 0)
|
|
self.assertEqual(tracker.access_resources([b.shrink(((4, 8),))], [0], 1), [])
|
|
self.assertEqual(tracker.access_resources([b.shrink(((0, 4),))], [0], 2), [0])
|
|
|
|
def test_partial_write_preserves_dependencies(self):
|
|
b = UOp.param(0, dtypes.uint8, 16, device="NULL")
|
|
for write in ([], [0]):
|
|
tracker = hcq2.DepsTracker()
|
|
tracker.access_resources([b], write, 0)
|
|
self.assertEqual(tracker.access_resources([b.shrink(((4, 12),))], [0], 1), [0])
|
|
self.assertEqual(tracker.access_resources([b.shrink(((0, 4),))], [0], 2), [0])
|
|
self.assertEqual(tracker.access_resources([b.shrink(((12, 16),))], [0], 3), [0])
|
|
self.assertEqual(tracker.access_resources([b.shrink(((4, 12),))], [], 4), [1])
|
|
|
|
def test_write_waits_for_all_readers(self):
|
|
b = UOp.param(0, dtypes.uint8, 16, device="NULL")
|
|
tracker = hcq2.DepsTracker()
|
|
self.assertEqual(tracker.access_resources([b], [], 0), [])
|
|
self.assertEqual(tracker.access_resources([b], [], 1), [])
|
|
self.assertEqual(tracker.access_resources([b], [0], 2), [0, 1])
|
|
self.assertEqual(tracker.access_resources([b], [], 3), [2])
|
|
|
|
def test_aliases_do_not_wait_on_themselves(self):
|
|
b = UOp.param(0, dtypes.uint8, 16, device="NULL")
|
|
for write in ([0], [1], [0, 1]):
|
|
tracker = hcq2.DepsTracker()
|
|
tracker.access_resources([b], [0], 0)
|
|
self.assertEqual(tracker.access_resources([b, b.shrink(((4, 12),))], write, 1), [0])
|
|
self.assertEqual(tracker.access_resources([b.shrink(((4, 12),))], [0], 2), [1])
|
|
|
|
class TestHCQ2Schedule(unittest.TestCase):
|
|
def setUp(self):
|
|
self.enterContext(Context(DEV="NULL"))
|
|
self.x = Tensor.ones(4).contiguous().realize()
|
|
|
|
def check(self, batch:UOp) -> UOp: # what every batch promises
|
|
cs = calls(batch)
|
|
for prio in rotations(batch):
|
|
log, left, timelines = run(batch, prio=prio)
|
|
self.assertFalse(any(left.values()), f"deadlock, left {left}")
|
|
for d in batch.arg.aux.device:
|
|
bumps = [i for i, x in enumerate(log) if x == d]
|
|
self.assertEqual((len(bumps), timelines[d]), (1, T + 1), f"{d} must bump its timeline once")
|
|
self.assertTrue(all(i < bumps[0] for i, x in enumerate(log) if isinstance(x, int) and d in devices_of(cs[x])), f"{d} bumps too early")
|
|
for d in batch.arg.aux.device:
|
|
log = run(batch, done={d: T - 1})[0]
|
|
self.assertFalse([x for x in log if isinstance(x, int) and d in devices_of(cs[x])], f"{d} runs before its previous batch is done")
|
|
return batch
|
|
|
|
def scheduled(self, *ts:Tensor) -> list[UOp]: return [self.check(b) for b in scheduled(*ts)]
|
|
def batch(self, *ts:Tensor) -> UOp: return unwrap(self.scheduled(*ts)[0])
|
|
|
|
def test_kernels_run_in_order(self): self.assertEqual(orders(self.batch(chain(self.x, 3))), {(0, 1, 2)})
|
|
|
|
def test_a_peer_kernel_runs_after_the_copy_that_feeds_it(self):
|
|
self.assertEqual(orders(self.batch((self.x.to("NULL:1") + 1).contiguous())), {(0, 1)})
|
|
|
|
def test_lanes_of_a_sharded_kernel_do_not_wait_for_each_other(self):
|
|
s = Tensor.ones(8).contiguous().realize().shard(("NULL", "NULL:1"), axis=0).contiguous().realize()
|
|
b = self.batch((s + 1).contiguous())
|
|
self.assertEqual(orders(b), {(0, 1), (1, 0)})
|
|
self.assertEqual(len([x for x in run(b, done={"NULL": T - 1})[0] if isinstance(x, int)]), 1)
|
|
|
|
def test_a_device_without_a_copy_queue_copies_with_a_kernel(self):
|
|
with patch.object(type(Device["NULL"]), "has_copy_queue", property(lambda _: False)):
|
|
b = self.batch((self.x.to("NULL:1") + 1).contiguous())
|
|
self.assertTrue(all(c.body.op is Ops.PROGRAM for c in calls(b)))
|
|
self.assertEqual(orders(b), {(0, 1)})
|
|
|
|
def test_a_host_kernel_splits_the_batch(self):
|
|
self.assertEqual(len(self.scheduled(((self.x + 1).contiguous().to("PYTHON") + 2).contiguous().to("NULL") + 3)), 2)
|
|
|
|
def test_batches_of_real_workloads_are_well_formed(self):
|
|
t = Tensor.ones(6).contiguous().realize().shard(("NULL", "NULL:1", "NULL:2"), axis=0)
|
|
self.scheduled((t + 1).sum(0).contiguous())
|
|
self.scheduled((self.x.to("NULL:1") + 1).to("NULL:2").contiguous().to("NULL") + 1)
|
|
|
|
class TestHCQ2Profile(unittest.TestCase):
|
|
def setUp(self): self.enterContext(Context(DEV="NULL"))
|
|
|
|
def test_profiling_reports_a_range_per_kernel(self, n=2):
|
|
x = Tensor.ones(4).contiguous().realize()
|
|
with Context(PROFILE=1):
|
|
seen = len(Compiled.profile_events)
|
|
chain(x, n).realize()
|
|
Device["NULL"].synchronize()
|
|
(ev,) = [e for e in Compiled.profile_events[seen:] if isinstance(e, ProfileGraphEvent)]
|
|
ranges = [(ev.sigs[e.st_id], ev.sigs[e.en_id]) for e in ev.ents]
|
|
self.assertEqual([e.device for e in ev.ents], ["NULL"] * n)
|
|
self.assertEqual([en - st for st, en in ranges], [1] * n, "NULL emulates 1us per kernel")
|
|
self.assertEqual(ranges, sorted(ranges))
|
|
|
|
def test_slots_addressed_by_the_device(self):
|
|
pm = PatternMatcher([(UPat((Ops.LOAD, Ops.STORE), src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat())),), allow_any_len=True),
|
|
lambda buf: buf.getaddr(Device["NULL"].host) if hcq2.unwrap_view(buf)[0].tag == "slots" else None)])
|
|
with patch.object(Device["NULL"], "pm_lower", pm): self.test_profiling_reports_a_range_per_kernel(n=3)
|
|
|
|
class TestHCQ2Link(unittest.TestCase):
|
|
def setUp(self): self.enterContext(Context(DEV="NULL"))
|
|
|
|
def test_links_serve_any_input(self):
|
|
a, inputs = chain_input(), list[UOp]()
|
|
linear = compile_linear(chain(a, 2).schedule_linear(), input_uops=inputs, cache=True)
|
|
linked = link_linear(linear, input_uops=inputs)
|
|
self.assertIs(link_linear(linear, input_uops=[chain_input(3).uop.base, *inputs[1:]]), linked)
|
|
bufs = [cast(Buffer, u.buffer) for u in linked.toposort() if u.op is Ops.BUFFER and u.addrspace is AddrSpace.GLOBAL]
|
|
self.assertNotIn(a.uop.base.buffer, bufs)
|
|
words = [w for b in bufs if b.options.external_ptr and b.nbytes % 8 == 0 for w in b.host.view(fmt='Q')[:]]
|
|
self.assertNotIn(cast(Buffer, a.uop.base.buffer)._buf, words)
|
|
|
|
def test_eager_templates_compile_once(self): self.assertIs(compiled_chain(3)[1], compiled_chain(3)[1])
|
|
|
|
def test_repeated_word_loops(self): # for (i..10) cmdbuf[off[i]] = var
|
|
var, offs = UOp.alloc((1,), dtypes.uint32, device="CPU").rtag("var"), [4 * i * i for i in range(10)]
|
|
var = hcq2.patch(var, [], bytes(4)) # initialize at link: the allocator may return a reused buffer
|
|
hq = SimpleNamespace(blob=bytearray(offs[-1] + 4), patches=[(o, var.index(0).load()) for o in offs], devs=("CPU",), queue="COPY:0", deps=())
|
|
lowered = lower_hcq(var.after(hcq2.encode_cmdbuf(hq, device="CPU")).index(0).store(var.index(0).load() + 1))
|
|
self.assertEqual(len([u for u in lowered.src[0].without_after.src[0].toposort() if u.op is Ops.RANGE]), 1)
|
|
linked = hcq2.hcq_link(lower_and_compile(lowered), allow_cache=False)
|
|
args = zip(lowered.src[0].without_after.src[1:], linked.src[0].without_after.src[1:])
|
|
cmdbuf = next(b.buffer for p, b in args if p.without_after.tag == "cmdbuf_copy_0")
|
|
for step in range(3):
|
|
run_linear(linked, jit=True)
|
|
self.assertEqual([cmdbuf.host.view(fmt="I")[o // 4] for o in offs], [step] * 10)
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|