| """Cooperative 2D matmul (tt_kernels.cmm_linear) bring-up / accuracy / timing / soak. |
| usage: cmm_test.py [eth|worker] shape=tiny|small|fc2|proj [uni] [kb=N] [slots=N] [time] [lib] [soak=N] [prof risks=...] |
| One call per launch unless time/soak is given.""" |
| 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 |
|
|
| arg = lambda k, d: next((a.split("=", 1)[1] for a in sys.argv if a.startswith(k + "=")), d) |
| disp = sys.argv[1] if len(sys.argv) > 1 and sys.argv[1] in ("eth", "worker") else "eth" |
| shape = arg("shape", "tiny") |
| kb, slots = int(arg("kb", "0")), int(arg("slots", "0")) |
| uni = "uni" in sys.argv |
| GX, GY, NT, KT = {"tiny": (2, 2, 1, 8), "small": (4, 3, 2, 16), "fc2": (10, 6, 4, 160), "proj": (10, 6, 4, 48), |
| "fc2h": (10, 6, 4, 80), "projh": (10, 6, 4, 24), |
| "k2t": (2, 6, 1, 8), "k2s": (4, 6, 2, 32), |
| "k2w": (2, 6, 8, 8), "k2w2": (3, 6, 12, 16), "qkv": (12, 6, 12, 40)}[shape] |
| kb = kb or {"tiny": 2, "small": 4, "fc2": 8, "proj": 8, "k2t": 2, "k2s": 4, "k2w": 2, "k2w2": 4, "qkv": 5}[shape] |
| slots = slots or min(3, GX, GY) |
| dev = open_device(grid="12x10", dispatch=disp, trace_region_size=64 << 20) |
| try: |
| torch.manual_seed(0) |
| M, K, N = 32 * GY, 32 * KT, 32 * NT * GX |
| at = (torch.randn(1, 1, M, K) * 0.5).to(torch.bfloat16) |
| wt = (torch.randn(K, N) * K ** -0.5).to(torch.bfloat16) |
| bt = (torch.randn(N) * 0.1).to(torch.bfloat16) |
| adt = ttnn.bfloat8_b if "abfp8" in sys.argv else ttnn.bfloat16 |
| a = ttnn.from_torch(at, dtype=adt, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.L1_MEMORY_CONFIG) |
| if "abfp8" in sys.argv: |
| at = ttnn.to_torch(a).to(torch.bfloat16) |
| gfrom = None |
| if "gh" in sys.argv: |
| gfrom = a |
| at = torch.nn.functional.gelu(at.float()).to(torch.bfloat16) |
| if "alocal" in sys.argv or "gh" in sys.argv: |
| amem = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.BLOCK_SHARDED, ttnn.BufferType.L1, ttnn.ShardSpec( |
| ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(GX - 1, GY - 1))}), [32, 32 * kb * (2 if arg("layout", "out") == "k2" else 1)], |
| ttnn.ShardOrientation.ROW_MAJOR)) |
| a = ttnn.to_memory_config(a, amem) if gfrom is None else ttnn.allocate_tensor_on_device( |
| ttnn.Shape([1, 1, M, K]), ttnn.bfloat16, ttnn.TILE_LAYOUT, dev, amem) |
| 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=GY, x=GX), strategy=ttnn.ShardStrategy.BLOCK, |
| orientation=ttnn.ShardOrientation.ROW_MAJOR) |
| wq = ttnn.to_torch(w).float() |
| aref = torch.nn.functional.gelu(at.float()).to(torch.bfloat16).float() if "gelu" in sys.argv else at.float() |
| ref = aref.reshape(M, K) @ wq + bt.float() |
| defs = tuple(d for d in arg("defs", "").split(",") if d) |
| _g = arg("layout", "out") == "g120" |
| bfn = None |
| if _g and "near" in sys.argv: |
| _nb = {(h, y): tk.near_banks(dev, 6 * h, y) for h in range(2) for y in range(10)} |
| bfn = lambda xw, b: _nb[(xw // 10, xw % 10)][b % 2] |
| wf = tk.CmmWeight(wt, dev, 20 if _g else GX, 2 if _g else NT, kb, ttnn.bfloat8_b, bank_fn=bfn) if "packed" in sys.argv else w |
| f = lambda: tk.cmm_linear(a, wf, b, omem, kb=kb, slots=slots, unicast=uni, defines=defs, layout=arg("layout", "out"), rotate="rot" in sys.argv, gelu="gelu" in sys.argv, v1="v1" in sys.argv, fidelity=getattr(ttnn.MathFidelity, arg("fid", "HiFi2")), gelu_from=gfrom, |
| out_dtype=ttnn.bfloat8_b if "obfp8" in sys.argv else None) |
| print(f"[cmm] {disp} shape={shape} G={GX}x{GY} NT={NT} KT={KT} kb={kb} slots={slots} uni={uni}", flush=True) |
| o = f() |
| ttnn.synchronize_device(dev) |
| ot = ttnn.to_torch(o).float().reshape(M, N) |
| print(f"[cmm] first call OK PCC {pcc(ot, ref):.7f} max|d| {float((ot - ref).abs().max()):.4f} mean|d| {float((ot - ref).abs().mean()):.5f}", flush=True) |
| if "lib" in sys.argv: |
| from hamer.vit import tuned_matmul_configs |
| pc = ttnn.MatmulMultiCoreReuseMultiCastProgramConfig( |
| compute_with_storage_grid_size=ttnn.CoreCoord(GX, GY), in0_block_w=kb, out_subblock_h=1, out_subblock_w=NT, |
| per_core_M=1, per_core_N=NT, transpose_mcast=False, fused_activation=None) |
| ck = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=ttnn.MathFidelity.HiFi2, fp32_dest_acc_en=True, packer_l1_acc=True) |
| g = lambda: ttnn.linear(a, w, bias=b, program_config=pc, memory_config=omem, compute_kernel_config=ck) |
| lt = ttnn.to_torch(g()).float().reshape(M, N) |
| print(f"[lib] PCC {pcc(lt, ref):.7f} mean|d| {float((lt - ref).abs().mean()):.5f} cmm==lib {bool(torch.equal(lt, ot))} max|cmm-lib| {float((lt-ot).abs().max()):.4f}", flush=True) |
| if "time" in sys.argv: |
| print(f"[lib] {time_op(dev, g):.2f} us", flush=True) |
| if "time" in sys.argv: |
| print(f"[cmm] {time_op(dev, f):.2f} us", flush=True) |
| if "prof" in sys.argv: |
| for _ in range(4): |
| o = f(); ttnn.synchronize_device(dev); ttnn.deallocate(o) |
| ttnn.ReadDeviceProfiler(dev) |
| n = int(arg("soak", "0")) |
| if n: |
| reps = 50 |
| o0 = ttnn.to_torch(f()).float() |
| tid = ttnn.begin_trace_capture(dev, cq_id=0) |
| outs = [] |
| try: |
| for _ in range(reps): |
| outs.append(f()) |
| 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(), o0)) for x in outs) |
| print(f"[cmm] soak {n // reps * reps} replays OK, outputs differing from first: {bad}/{reps}", flush=True) |
| finally: |
| ttnn.release_trace(dev, tid) |
| finally: |
| ttnn.close_device(dev) |
| if "prof" in sys.argv: |
| tl_report(risks=tuple(arg("risks", "NCRISC,BRISC,TRISC_1").split(","))) |
| tl_seq(("WISS", "WDRAM", "WFWD", "XFER"), "NCRISC") |
| tl_seq(("RECV", "XFER"), "BRISC") |
| tl_seq(("CWA", "CWW"), "TRISC_0") |
|
|
