import unittest, struct, contextlib, tempfile, pathlib, json, time, atexit, random from tinygrad import Device, Tensor, dtypes, TinyJit from tinygrad.helpers import CI, getenv, Context, ProfileLogger from tinygrad.device import Buffer, BufferOptions, HCQCompiled from tinygrad.engine.schedule import create_schedule from tinygrad.engine.realize import get_runner MOCKGPU = getenv("MOCKGPU") @contextlib.contextmanager def helper_collect_profile(*devs, random_setup_delay=False): ProfileLogger.mjson, ProfileLogger.actors = [], {} if random_setup_delay: devs = list(devs) for dev in devs: dev.synchronize() random.shuffle(devs) for dev in devs: dev._prof_setup() time.sleep(random.randint(1, 1000) / 1000) else: for dev in devs: dev._prof_setup() profile_dict = {} _, tmp = tempfile.mkstemp() with Context(PROFILE=1, PROFILEPATH=tmp): try: yield profile_dict finally: for dev in devs: dev.synchronize() dev._prof_finalize() atexit.unregister(dev._prof_finalize) for k,v in json.loads(pathlib.Path(tmp).read_text()).items(): profile_dict[k] = v pathlib.Path(tmp).unlink() def helper_profile_filter_node(profile, **kwargs): assert len(profile) > 0, "Empty profile" assert 'traceEvents' in profile, "traceEvents should present" return [x for x in profile['traceEvents'] if all(x.get(k, None) == v for k,v in kwargs.items())] def helper_profile_parse_pids(profile): pids, tids = {}, {} procs = helper_profile_filter_node(profile, name='process_name') for proc in procs: pids[proc['pid']] = proc['args']['name'] threads = helper_profile_filter_node(profile, name='thread_name') for th in threads: tids[th['tid']] = th['args']['name'] return pids, tids def helper_profile_parse_deps(profile): deps = [] for s in helper_profile_filter_node(profile, ph="s"): f = helper_profile_filter_node(profile, ph="f", id=s['id'])[0] starts, ends = [], [] for x in helper_profile_filter_node(profile, ph="X"): if s['pid'] == x['pid'] and s['tid'] == x['tid'] and x['ts'] <= s['ts'] <= x['ts'] + x['dur']: starts.append(x) if f['pid'] == x['pid'] and f['tid'] == x['tid'] and x['ts'] <= f['ts'] <= x['ts'] + x['dur']: ends.append(x) assert len(starts) == 1 and len(ends) == 1, "more than one start and end possible, valid?" deps.append((s, f, starts[0], ends[0])) return deps def helper_validate_node(node, duration_s=10, ts_age_s=30, profile=None, pid_name=None, tid_name=None): pids, tids = helper_profile_parse_pids(profile) assert abs(node['ts'] - time.perf_counter_ns() / 1e3) < ts_age_s * 1e6, "timestimp is not in 30s range" assert 0 < node['dur'] < duration_s * 1e6, "duration is not in 10s range" assert pid_name is None or pids[node['pid']] == pid_name assert tid_name is None or tids[node['tid']] == tid_name @unittest.skipUnless(issubclass(type(Device[Device.DEFAULT]), HCQCompiled), "HCQ device required to run") class TestProfiler(unittest.TestCase): @classmethod def setUpClass(self): TestProfiler.d0 = Device[Device.DEFAULT] TestProfiler.a = Tensor([0.,1.], device=Device.DEFAULT).realize() TestProfiler.b = self.a + 1 si = create_schedule([self.b.lazydata])[-1] TestProfiler.runner = get_runner(TestProfiler.d0.dname, si.ast) TestProfiler.b.lazydata.buffer.allocate() TestProfiler.kernargs_ba_ptr = TestProfiler.runner.clprg.fill_kernargs([TestProfiler.b.lazydata.buffer._buf, TestProfiler.a.lazydata.buffer._buf]) TestProfiler.kernargs_ab_ptr = TestProfiler.runner.clprg.fill_kernargs([TestProfiler.a.lazydata.buffer._buf, TestProfiler.b.lazydata.buffer._buf]) def test_profile_kernel_run(self): runner_name = TestProfiler.runner.clprg.name with helper_collect_profile(TestProfiler.d0) as profile: TestProfiler.runner([TestProfiler.b.lazydata.buffer, TestProfiler.a.lazydata.buffer], var_vals={}) kernel_node = helper_profile_filter_node(profile, name=runner_name)[0] helper_validate_node(kernel_node, profile=profile, pid_name=Device.DEFAULT, tid_name="COMPUTE") def test_profile_copyin(self): buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferOptions(nolru=True)).ensure_allocated() with helper_collect_profile(TestProfiler.d0) as profile: buf1.copyin(memoryview(bytearray(struct.pack("ff", 0, 1)))) copyin_node = helper_profile_filter_node(profile, name=f"CPU -> {Device.DEFAULT}")[0] helper_validate_node(copyin_node, profile=profile, pid_name=Device.DEFAULT, tid_name="DMA") def test_profile_multiops(self): runner_name = TestProfiler.runner.clprg.name buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferOptions(nolru=True)).ensure_allocated() with helper_collect_profile(TestProfiler.d0) as profile: buf1.copyin(memoryview(bytearray(struct.pack("ff", 0, 1)))) TestProfiler.runner([buf1, TestProfiler.a.lazydata.buffer], var_vals={}) buf1.as_buffer() copyin_node = helper_profile_filter_node(profile, name=f"CPU -> {Device.DEFAULT}")[0] helper_validate_node(copyin_node, profile=profile, pid_name=Device.DEFAULT, tid_name="DMA") kernel_node = helper_profile_filter_node(profile, name=runner_name)[0] helper_validate_node(kernel_node, profile=profile, pid_name=Device.DEFAULT, tid_name="COMPUTE") copyout_node = helper_profile_filter_node(profile, name=f"{Device.DEFAULT} -> CPU")[0] helper_validate_node(copyout_node, profile=profile, pid_name=Device.DEFAULT, tid_name="DMA") assert copyin_node['ts'] + copyin_node['dur'] < kernel_node['ts'], "timestamp not aranged" assert kernel_node['ts'] + kernel_node['dur'] < copyout_node['ts'], "timestamp not aranged" def test_profile_multidev_copyin(self): d1 = Device[f"{Device.DEFAULT}:1"] buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferOptions(nolru=True)).ensure_allocated() buf2 = Buffer(f"{Device.DEFAULT}:1", 2, dtypes.float, options=BufferOptions(nolru=True)).ensure_allocated() with helper_collect_profile(TestProfiler.d0, d1) as profile: buf1.copyin(memoryview(bytearray(struct.pack("ff", 0, 1)))) buf2.copyin(memoryview(bytearray(struct.pack("ff", 0, 1)))) copyin_node_1 = helper_profile_filter_node(profile, name=f"CPU -> {Device.DEFAULT}")[0] helper_validate_node(copyin_node_1, profile=profile, pid_name=Device.DEFAULT, tid_name="DMA") copyin_node_2 = helper_profile_filter_node(profile, name=f"CPU -> {Device.DEFAULT}:1")[0] helper_validate_node(copyin_node_2, profile=profile, pid_name=f"{Device.DEFAULT}:1", tid_name="DMA") def test_profile_multidev_transfer(self): d1 = Device[f"{Device.DEFAULT}:1"] a = Tensor.randn(1 << 20, device=Device.DEFAULT).realize() with helper_collect_profile(TestProfiler.d0, d1) as profile: y = a.to(f"{Device.DEFAULT}:1") y.realize() transfer_node_1 = helper_profile_filter_node(profile, name=f"{Device.DEFAULT} -> {Device.DEFAULT}:1")[0] helper_validate_node(transfer_node_1, profile=profile, pid_name=Device.DEFAULT, tid_name="DMA") assert 80 < transfer_node_1['dur'] < (5000 if CI else 1400), f"Duration is not in the range: {transfer_node_1['dur']}" @unittest.skipIf(MOCKGPU and Device.DEFAULT == "AMD", "AMD mockgpu with indirect buffers does not support queue wait interrupts") def test_profile_deps(self): d1 = Device[f"{Device.DEFAULT}:1"] def f(a): x = (a + 1).realize() return x, x.to(d1.dname).realize() a = Tensor.randn(10, 10, device=TestProfiler.d0.dname).realize() with helper_collect_profile(TestProfiler.d0, d1) as profile: jf = TinyJit(f) for _ in range(3): jf(a) del jf deps = helper_profile_parse_deps(profile) assert len(deps) == 1, "one dep is expected, one launch" _, _, l, r = deps[0] assert l['name'].find("->") == -1, "should be kernel" assert r['name'] == f"{Device.DEFAULT} -> {Device.DEFAULT}:1", "should be copy" @unittest.skipIf(MOCKGPU and Device.DEFAULT == "AMD", "AMD mockgpu with indirect buffers does not support queue wait interrupts") def test_profile_copy_args(self): d1 = Device[f"{Device.DEFAULT}:1"] def f(a): x = (a + 1).realize() return x, x.to(d1.dname).realize() a = Tensor.randn(10, 10, device=TestProfiler.d0.dname).realize() with helper_collect_profile(TestProfiler.d0, d1) as profile: jf = TinyJit(f) for _ in range(3): TestProfiler.d0.raw_prof_records, TestProfiler.d0.sig_prof_records = [], [] # reset to collect only graph logs d1.raw_prof_records, d1.sig_prof_records = [], [] jf(a) del jf node = helper_profile_filter_node(profile, name=f"{Device.DEFAULT} -> {Device.DEFAULT}:1")[-1] assert node['args']['Size'] == "400.00 B" assert abs(float(node['args']['GB/S']) - ((10 * 10 * 4) / 1e3) / (node['dur'])) < 0.01 @unittest.skipIf(CI, "skip CI") def test_profile_sync(self): mv = memoryview(bytearray(struct.pack("ff", 0, 1))) expected_diff = 100000 # sleep in us devs = [Device[f"{Device.DEFAULT}:{i}"] for i in range(6)] bufs = [Buffer(f"{Device.DEFAULT}:{i}", 2, dtypes.float, options=BufferOptions(nolru=True)).ensure_allocated() for i in range(6)] # enqueue ops on different queues to check the timer sync cpu_time = [] with helper_collect_profile(*devs, random_setup_delay=True) as profile: for i in range(6): x = time.perf_counter_ns() time.sleep(expected_diff / 1e6) bufs[i].copyin(mv) cpu_time.append(((time.perf_counter_ns() - x) / 1000) - expected_diff) nodes = [helper_profile_filter_node(profile, name=f"CPU -> {Device.canonicalize(f'{Device.DEFAULT}:{i}')}")[-1] for i in range(6)] avg_diff = [] for i in range(1, 6): diff = nodes[i]['ts'] - nodes[i-1]['ts'] - cpu_time[i] avg_diff.append(diff - expected_diff) assert expected_diff * 0.998 < diff < expected_diff * 1.002, "more that 0.2% diff" print(f"total avg delay is {sum(avg_diff) / len(avg_diff)} us") if __name__ == "__main__": unittest.main()