| |
| """Where the e2e tt(image) time goes beyond the trace (round 11 opt tool): timestamps inside the call |
| (start -> _exec_trace entry = host prep + band uploads, -> _read_fused entry, _read_fused (band readback on |
| CQ1 + C++ tail), -> return (host finish)), medians over N calls. Device opened as served (worker, MOGE_2CQ).""" |
| import os |
| import statistics |
| import sys |
| import time |
|
|
| sys.path.insert(0, "/home/ttuser/experiments/tt-models/tools") |
| import torch |
|
|
|
|
| def main(): |
| from open_device_12x10 import open_device, close_device |
| import ttnn |
| import tt_moge |
| from tt_moge.benchmark import load_image |
| from tt_moge.reference.load_pretrained import load_moge2 |
| from tt_moge.tt.ttnn_moge import TtMoGe |
| torch.set_grad_enabled(False) |
| n = int(os.environ.get("N", 40)) |
| ref = load_moge2() |
| image = load_image(os.path.join(os.path.dirname(__file__), "..", "..", "..", "media", "source.png")).float() |
| dev = open_device(grid="12x10", dispatch=os.environ.get("MOGE_DISPATCH", "worker"), l1_small_size=32768, |
| trace_region_size=1500000000, |
| num_command_queues=2 if os.environ.get("MOGE_2CQ", "0") == "1" else 1) |
| try: |
| tt = TtMoGe(ref, dev) |
| tt(image, num_tokens=1800) |
| print(f"grouped={tt._grouped} segments={1 + len(tt._dec_segs)} scale_trace={tt._dec_tid2 is not None}") |
| marks = {} |
| ex, rd = tt._exec_trace, tt._read_fused |
|
|
| def ex_w(*a, **k): |
| marks["exec"] = time.perf_counter() |
| r = ex(*a, **k) |
| marks["exec_ret"] = time.perf_counter() |
| return r |
|
|
| def rd_w(*a, **k): |
| marks["read"] = time.perf_counter() |
| r = rd(*a, **k) |
| marks["read_ret"] = time.perf_counter() |
| return r |
| tt._exec_trace, tt._read_fused = ex_w, rd_w |
| rows = [] |
| for _ in range(n): |
| ttnn.synchronize_device(dev) |
| if tt._grouped: |
| tt._grp_times = [] |
| t0 = time.perf_counter() |
| tt(image, num_tokens=1800) |
| t1 = time.perf_counter() |
| if tt._grouped: |
| for i, t in enumerate(tt._grp_times): |
| marks[f"g{i}"] = t - marks["exec"] |
| rows.append(dict({k: v for k, v in marks.items() if k[0] == "g"},pre=marks["exec"] - t0, exec=marks["exec_ret"] - marks["exec"], |
| gap=marks["read"] - marks["exec_ret"], read=marks["read_ret"] - marks["read"], |
| finish=t1 - marks["read_ret"], e2e=t1 - t0)) |
| for k in rows[0]: |
| v = [r[k] * 1e3 for r in rows] |
| print(f"{k:8s} median {statistics.median(v):7.2f} ms min {min(v):7.2f}") |
| finally: |
| close_device(dev) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
