| """Round 8: zone timeline of the fused cmm + add+LN epilogue program (proj shape, LN output width-sharded on 10 cores). |
| usage: ADDLN_EXP=prof TT_METAL_DEVICE_PROFILER=1 fused_tl.py [shape=proj|fc2]""" |
| import sys, os, csv |
| from pathlib import Path |
| from collections import defaultdict |
| from common import * |
| sys.path.insert(0, str(Path(__file__).resolve().parents[2])) |
| from hamer import tt_kernels as tk |
| arg = lambda k, d: next((a.split("=", 1)[1] for a in sys.argv if a.startswith(k + "=")), d) |
| shape = arg("shape", "proj") |
| GX, GY, NT, KT, kb = {"proj": (10, 6, 4, 48, 8), "fc2": (10, 6, 4, 160, 16)}[shape] |
| dev = open_dev() |
| M, K, N = 32 * GY, 32 * KT, 32 * NT * GX |
| G = dev.compute_with_storage_grid_size() |
| res = ttnn.create_sharded_memory_config((M, N), core_grid=ttnn.CoreGrid(y=GY, x=GX), strategy=ttnn.ShardStrategy.BLOCK, |
| orientation=ttnn.ShardOrientation.ROW_MAJOR) |
| a = ttnn.from_torch(torch.randn(1, 1, M, K) * 0.5, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.L1_MEMORY_CONFIG) |
| wf = tk.CmmWeight((torch.randn(K, N) * K ** -0.5).to(torch.bfloat16), dev, GX, NT, kb, ttnn.bfloat8_b) |
| b = ttnn.from_torch(torch.zeros(1, 1, 1, N), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev) |
| x = ttnn.from_torch(torch.randn(1, 1, M, N), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=res) |
| C = tk.addln_constants(torch.ones(N), torch.zeros(N), dev) |
| if "noaff" in sys.argv: |
| C["noaff"] = True |
| if "nobeta" in sys.argv: |
| C["nobeta"] = True |
| sg = ttnn.num_cores_to_corerangeset(10, G, row_wise=True) |
| hmem = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.L1, ttnn.ShardSpec(sg, [M, N // 10], ttnn.ShardOrientation.ROW_MAJOR)) |
| defs = ("CMM_FLUSH", "CMM_LA=1", "CMM_NOSTAGE", "CMM_PROF") |
| for _ in range(4): |
| s, h = tk.cmm_linear(a, wf, b, res, kb=kb, slots=3, defines=defs, v1=True, addln=dict(x=x, consts=C, eps=1e-6, out_mem=hmem)) |
| ttnn.synchronize_device(dev); ttnn.deallocate(s); ttnn.deallocate(h) |
| ttnn.ReadDeviceProfiler(dev) |
| ttnn.close_device(dev) |
| f = os.environ["TT_METAL_PROFILER_DIR"] + "/.logs/profile_log_device.csv" |
| rows = list(csv.reader(open(f)))[2:] |
| ev = defaultdict(list) |
| for r in rows: |
| try: |
| cx, cy, risc, cyc, name, typ, run = int(r[1]), int(r[2]), r[3].strip(), int(r[5]), r[10].strip(), r[11].strip(), r[7].strip() |
| except Exception: |
| continue |
| ev[(run, cx, cy)].append((cyc, risc, name, typ)) |
| last = sorted({k[0] for k in ev}, key=lambda v: int(v) if v.isdigit() else 0)[-1] |
| t00 = min(c for (run, _, _), v in ev.items() if run == last for c, rk, n, t in v if "KERNEL" in n) |
| tl = defaultdict(lambda: defaultdict(list)) |
| for (run, cx, cy), v in ev.items(): |
| if run != last: |
| continue |
| st = {} |
| for c, rk, n, t in sorted(v): |
| if "START" in t: |
| st[(rk, n)] = c |
| elif "END" in t and (rk, n) in st: |
| tl[rk][n].append(((st.pop((rk, n)) - t00) / 1.35e3, (c - t00) / 1.35e3)) |
| for rk in tuple(os.environ.get("TL_RISCS", "TRISC_1,BRISC,NCRISC").split(",")): |
| for n, v in sorted(tl[rk].items(), key=lambda kv: sorted(kv[1])[len(kv[1]) // 2][0]): |
| if n in ("RECV", "XFER", "XDATA", "XFL1", "WRDY", "WISS", "WDRAM"): |
| continue |
| s_ = sorted(a_ for a_, b_ in v); e_ = sorted(b_ for a_, b_ in v) |
| print(f"{rk:8s} {n:22s} start med {s_[len(v) // 2]:6.2f} max {s_[-1]:6.2f} end med {e_[len(v) // 2]:6.2f} max {e_[-1]:6.2f}") |
|
|
