| """custom GELU (tt_kernels.gelu) soak: one eager call, then N trace replays of 50 calls on the fc1-output shape; |
| every output must stay bit-identical. usage: gelu_soak.py N""" |
| import sys |
| from pathlib import Path |
| sys.path.insert(0, str(Path(__file__).resolve().parents[2])) |
| from common import * |
| from hamer import tt_kernels |
| n = int(sys.argv[1]) if len(sys.argv) > 1 else 1000 |
| dev = open_dev() |
| try: |
| L1 = ttnn.L1_MEMORY_CONFIG |
| torch.manual_seed(0) |
| xt = torch.randn(1, 1, 192, 5120) * 3.0 |
| x = ttnn.from_torch(xt, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1) |
| o0 = ttnn.to_torch(tt_kernels.gelu(x, memory_config=L1)).float() |
| ref = torch.nn.functional.gelu(xt.bfloat16().float()) |
| print(f"[gsoak] first call OK pcc {pcc(o0, ref):.8f} max|err| {float((o0 - ref).abs().max()):.3e}", flush=True) |
| reps = 50 |
| tid = ttnn.begin_trace_capture(dev, cq_id=0) |
| outs = [] |
| try: |
| for _ in range(reps): |
| outs.append(tt_kernels.gelu(x, memory_config=L1)) |
| finally: |
| ttnn.end_trace_capture(dev, tid, cq_id=0) |
| try: |
| for i in range(max(1, n // reps)): |
| ttnn.execute_trace(dev, tid, cq_id=0, blocking=False) |
| if i % 40 == 39: |
| ttnn.synchronize_device(dev) |
| ttnn.synchronize_device(dev) |
| bad = sum(int(not torch.equal(ttnn.to_torch(t).float(), o0)) for t in outs) |
| print(f"[gsoak] {max(1, n // reps) * reps} calls OK, outputs differing from first: {bad}/{reps}", flush=True) |
| finally: |
| ttnn.release_trace(dev, tid) |
| finally: |
| ttnn.close_device(dev) |
|
|
