File size: 2,837 Bytes
13b4736 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | # SPDX-License-Identifier: Apache-2.0
"""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 # noqa: F401
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()
|