| """gelu(prefetch) -> cmm_linear(prefetch) (fc2 with its first weight pairs prefetched during the GELU) vs the |
| model's gelu -> library 2D linear. usage: cmm_pf_test.py [pairs=P] [la=L] [time] [soak=N]""" |
| import sys |
| from pathlib import Path |
| from common import * |
| from open_device_12x10 import open_device |
| sys.path.insert(0, str(Path(__file__).resolve().parents[2])) |
| from hamer import tt_kernels as tk |
| from hamer.vit import tuned_matmul_configs |
| arg = lambda k, d: next((a.split("=", 1)[1] for a in sys.argv if a.startswith(k + "=")), d) |
| P, LA, KB = int(arg("pairs", "3")), int(arg("la", "1")), int(arg("kb", "16")) |
| dev = open_device(grid="12x10", dispatch="eth", trace_region_size=64 << 20) |
| try: |
| torch.manual_seed(0) |
| L1 = ttnn.L1_MEMORY_CONFIG |
| M, K, N = 192, 5120, 1280 |
| at = (torch.randn(1, 1, M, K) * 1.5).to(torch.bfloat16) |
| wt = (torch.randn(K, N) * K ** -0.5).to(torch.bfloat16) |
| bt = (torch.randn(N) * 0.1).to(torch.bfloat16) |
| a = ttnn.from_torch(at, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1) |
| w = ttnn.from_torch(wt, dtype=ttnn.bfloat8_b, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| b = ttnn.from_torch(bt.reshape(1, 1, 1, N), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev) |
| omem = ttnn.create_sharded_memory_config((M, N), core_grid=ttnn.CoreGrid(y=6, x=10), strategy=ttnn.ShardStrategy.BLOCK, |
| orientation=ttnn.ShardOrientation.ROW_MAJOR) |
| wpk = tk.CmmWeight(wt, dev, 10, 4, KB, ttnn.bfloat8_b) |
| pf = tk.CmmPrefetch(dev, 10, 6, wpk, P) if P else None |
| defs = ("CMM_FLUSH", f"CMM_LA={LA}", "CMM_DRAM0") |
| def fused(): |
| g = tk.gelu(a, memory_config=L1, prefetch=pf) |
| o = tk.cmm_linear(g, wpk, b, omem, kb=KB, defines=defs, prefetch=pf) |
| ttnn.deallocate(g) |
| return o |
| ck = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=ttnn.MathFidelity.HiFi2, fp32_dest_acc_en=True, packer_l1_acc=True) |
| pc2 = tuned_matmul_configs(dev.compute_with_storage_grid_size(), fp32_acc=True)["fc2"] |
| def lib(): |
| g = tk.gelu(a, memory_config=L1) |
| o = ttnn.linear(g, w, bias=b, program_config=pc2, memory_config=omem, compute_kernel_config=ck) |
| ttnn.deallocate(g) |
| return o |
| print(f"[pf] pairs={P} la={LA} kb={KB}", flush=True) |
| o = fused(); ttnn.synchronize_device(dev) |
| ot = ttnn.to_torch(o).float() |
| lt = ttnn.to_torch(lib()).float() |
| ref = torch.nn.functional.gelu(at.float()).bfloat16().float().reshape(M, K) @ ttnn.to_torch(w).float() + bt.float() |
| print(f"[pf] lib PCC {pcc(lt.reshape(M, N), ref):.7f}", flush=True) |
| print(f"[pf] first call OK PCC {pcc(ot.reshape(M, N), ref):.7f} bit-identical to library pair: {bool(torch.equal(ot, lt))} max|d| {float((ot - lt).abs().max()):.4f}", flush=True) |
| if "time" in sys.argv: |
| print(f"[pf] gelu+cmm(prefetch) {time_op(dev, fused):.2f} us", flush=True) |
| print(f"[pf] gelu+library 2D {time_op(dev, lib):.2f} us", flush=True) |
| print(f"[pf] gelu(prefetch) alone {time_op(dev, lambda: tk.gelu(a, memory_config=L1, prefetch=pf)):.2f} us gelu alone {time_op(dev, lambda: tk.gelu(a, memory_config=L1)):.2f} us", flush=True) |
| n = int(arg("soak", "0")) |
| if n: |
| reps = 50 |
| tid = ttnn.begin_trace_capture(dev, cq_id=0) |
| outs = [] |
| try: |
| for _ in range(reps): |
| outs.append(fused()) |
| finally: |
| ttnn.end_trace_capture(dev, tid, cq_id=0) |
| try: |
| for i in range(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(x).float(), ot)) for x in outs) |
| print(f"[pf] soak {n // reps * reps} calls OK, outputs differing from the first: {bad}/{reps}", flush=True) |
| finally: |
| ttnn.release_trace(dev, tid) |
| finally: |
| ttnn.close_device(dev) |
|
|
