| |
| |
| """SE-K probe (OPT_PLAN SE-K, ``tt/sek.py``), per the hang protocol: the one-program SE MLP vs the stock two |
| ``layers.Linear`` (fc1 + SiLU, fc2 + sigmoid; two weight terms, two activation terms, EPI) on the 7 distinct SE shapes |
| of SceneSeg (C, squeeze), random fp32 rows: eager x2 identical, a 2-call trace, the error of both vs float64, and the |
| time of each in a trace. |
| |
| bin/devrun -t 600 -- env TT_METAL_WATCHER=2 python -u code/scripts/diag/se_probe.py --smallest |
| bin/devrun -t 900 -- python -u code/scripts/diag/se_probe.py |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| import time |
|
|
| import numpy as np |
|
|
| SHAPES = [(32, 8), (96, 4), (144, 6), (240, 10), (480, 20), (672, 28), (1152, 48)] |
|
|
|
|
| def main() -> int: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--smallest", action="store_true") |
| a = ap.parse_args() |
|
|
| import torch |
| import ttnn |
|
|
| from tt_sceneseg.device import close_device, open_device |
| from tt_sceneseg.tt.layers import Linear |
| from tt_sceneseg.tt.sek import SEKernel |
| from tt_sceneseg.ttaw.tensors import to_numpy |
|
|
| dev = open_device() |
| rng = np.random.default_rng(2) |
| rc = 0 |
|
|
| def tm(fn): |
| r = fn() |
| ttnn.deallocate(r) |
| tid = ttnn.begin_trace_capture(dev, cq_id=0) |
| rs = [fn() for _ in range(10)] |
| ttnn.end_trace_capture(dev, tid, cq_id=0) |
| ttnn.execute_trace(dev, tid, cq_id=0, blocking=True) |
| t1 = time.perf_counter() |
| for _ in range(5): |
| ttnn.execute_trace(dev, tid, cq_id=0, blocking=False) |
| ttnn.synchronize_device(dev) |
| us = (time.perf_counter() - t1) / 50 * 1e6 |
| ttnn.release_trace(dev, tid) |
| for r in rs: |
| ttnn.deallocate(r) |
| return us |
|
|
| for c, s in (SHAPES[:2] if a.smallest else SHAPES): |
| w1 = (rng.standard_normal((s, c)) / np.sqrt(c) / 50).astype(np.float32) |
| b1 = (rng.standard_normal(s) * 0.5).astype(np.float32) |
| w2 = (rng.standard_normal((c, s)) / np.sqrt(s)).astype(np.float32) |
| b2 = (rng.standard_normal(c) * 0.5).astype(np.float32) |
| x = (rng.standard_normal((1, 1, 1, c)) * 30).astype(np.float32) |
| fc1 = Linear(w1, b1, activation="silu", output_dtype="float32", weight_terms=2, act_terms=2, name="fc1", |
| epi=True) |
| fc2 = Linear(w2, b2, activation="sigmoid", output_dtype="float32", weight_terms=2, act_terms=2, name="fc2", |
| epi=True) |
| sek = SEKernel(fc1, fc2) |
| xt = ttnn.from_torch(torch.from_numpy(x), dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, device=dev, |
| memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| x64 = x.reshape(-1).astype(np.float64) |
| h = x64 @ w1.T.astype(np.float64) + b1 |
| h = h / (1 + np.exp(-h)) |
| ref = 1 / (1 + np.exp(-(h @ w2.T.astype(np.float64) + b2))) |
| st = fc2(fc1(xt)) |
| sn = to_numpy(st).reshape(-1)[:c].astype(np.float64) |
| ttnn.deallocate(st) |
| res = [] |
| for _ in range(2): |
| o = sek(xt) |
| res.append(to_numpy(o).reshape(-1)[:c]) |
| ttnn.deallocate(o) |
| same = np.array_equal(res[0].view(np.uint32), res[1].view(np.uint32)) |
| tid = ttnn.begin_trace_capture(dev, cq_id=0) |
| o1 = sek(xt) |
| o2 = sek(xt) |
| ttnn.end_trace_capture(dev, tid, cq_id=0) |
| ttnn.execute_trace(dev, tid, cq_id=0, blocking=True) |
| tr = np.array_equal(to_numpy(o2).reshape(-1)[:c].view(np.uint32), res[0].view(np.uint32)) |
| ttnn.release_trace(dev, tid) |
| ttnn.deallocate(o1) |
| ttnn.deallocate(o2) |
| ek = float(np.abs(res[0] - ref).max()) |
| es = float(np.abs(sn - ref).max()) |
| t_k = tm(lambda: sek(xt)) |
| t_s = tm(lambda: fc2(fc1(xt))) |
| ok = same and tr and ek <= max(2 * es, 1e-6) |
| rc |= 0 if ok else 1 |
| print(f"SE C {c:5d} S {s:3d}: eager x2 identical {same}, trace identical {tr}, max abs err vs float64: " |
| f"kernel {ek:.2e}, stock {es:.2e} | kernel {t_k:7.1f} us, stock {t_s:7.1f} us {'OK' if ok else 'FAIL'}", |
| flush=True) |
| sek.release() |
| fc1.release() |
| fc2.release() |
| ttnn.deallocate(xt) |
| close_device(dev) |
| print("ALL OK" if rc == 0 else "FAIL") |
| return rc |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
