File size: 1,877 Bytes
58ba6df | 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 | """device-profiler zones of hattn (TT_METAL_DEVICE_PROFILER=1)"""
import sys, os, csv
from pathlib import Path
from collections import defaultdict
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from common import *
from hamer import tt_kernels as tk
dev = open_dev()
try:
rowf = lambda v: v.reshape(1, 1, 1, -1).expand(1, 1, 32, v.numel()).contiguous()
q = ttnn.from_torch(rowf(torch.randn(512)).to(torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.L1_MEMORY_CONFIG)
kv = ttnn.from_torch(torch.randn(1, 1, 192, 6144).to(torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.L1_MEMORY_CONFIG)
for _ in range(4):
o = tk.hattn(q, kv, 2); ttnn.synchronize_device(dev); ttnn.deallocate(o)
ttnn.ReadDeviceProfiler(dev)
finally:
from open_device_12x10 import close_device
close_device(dev)
rows = list(csv.reader(open(os.environ["TT_METAL_PROFILER_DIR"] + "/.logs/profile_log_device.csv")))[2:]
ev = defaultdict(list)
for r in rows:
try:
ev[(r[7].strip(), int(r[1]), int(r[2]), r[3].strip())].append((int(r[5]), r[10].strip(), r[11].strip()))
except Exception:
pass
last = sorted({k[0] for k in ev}, key=lambda v: int(v) if v.isdigit() else 0)[-1]
t0 = min(c for k, v in ev.items() if k[0] == last for c, _, _ in v)
dur = defaultdict(list)
for k, v in ev.items():
if k[0] != last:
continue
st = {}
for c, n, t in sorted(v):
if "START" in t: st[n] = c
elif "END" in t and n in st:
dur[(k[3], n)].append(((st[n] - t0) / 1.35e3, (c - st.pop(n)) / 1.35e3))
for k in sorted(dur):
v = dur[k]; s = sorted(d for _, d in v); b0 = sorted(b for b, _ in v)
print(f"{k[0]:8s} {k[1]:20s} n={len(v):3d} start med {b0[len(b0)//2]:6.2f} dur med {s[len(s)//2]:6.2f} max {s[-1]:6.2f} us")
|