| """cmm_linear (kb, slots) sweep on one shape, PCC-checked. usage: cmm_sweep.py [eth|worker] shape=fc2 cfgs=8x3,8x6,16x3 [defs=...]""" |
| 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", "fc2") |
| GX, GY, NT, KT = {"fc2": (10, 6, 4, 160), "proj": (10, 6, 4, 48), "qkv": (12, 6, 12, 40), "fc1": (10, 6, 16, 40)}[shape] |
| cfgs = [tuple(int(v) for v in c.split("x")) for c in arg("cfgs", "8x3").split(",")] |
| lay = arg("layout", "out") |
| defs = tuple(d for d in arg("defs", "").split(",") if d) |
| 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) |
| a = ttnn.from_torch(at, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.L1_MEMORY_CONFIG) |
| 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) |
| aref = torch.nn.functional.gelu(at.float()).to(torch.bfloat16).float() if "gelu" in sys.argv else at.float() |
| ref = aref.reshape(M, K) @ ttnn.to_torch(w).float() + bt.float() |
| for kb, s in cfgs: |
| wf = tk.CmmWeight(wt, dev, GX, NT, kb, ttnn.bfloat8_b) if "packed" in sys.argv else w |
| f = lambda: tk.cmm_linear(a, wf, b, omem, kb=kb, slots=s, defines=defs, layout=lay, rotate="rot" in sys.argv, gelu="gelu" in sys.argv, v1="v1" in sys.argv) |
| o = f(); ttnn.synchronize_device(dev) |
| p = pcc(ttnn.to_torch(o).float().reshape(M, N), ref) |
| print(f"[sweep] {shape} {lay} packed={'packed' in sys.argv} rot={'rot' in sys.argv} kb={kb} slots={s} defs={defs} PCC {p:.7f} {time_op(dev, f):.2f} us", flush=True) |
| if "lib" in sys.argv: |
| from hamer.vit import tuned_matmul_configs |
| 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)[shape] |
| def libf(): |
| g = tk.gelu(a, memory_config=ttnn.L1_MEMORY_CONFIG) if "gelu" in sys.argv else a |
| o = ttnn.linear(g, w, bias=b, program_config=pc2, memory_config=omem, compute_kernel_config=ck) |
| if g is not a: |
| ttnn.deallocate(g) |
| return o |
| o = libf(); ttnn.synchronize_device(dev) |
| p = pcc(ttnn.to_torch(o).float().reshape(M, N), ref) |
| print(f"[sweep] {shape} LIBRARY (gelu kernel={'gelu' in sys.argv} + 2D linear) PCC {p:.7f} {time_op(dev, libf):.2f} us", flush=True) |
| finally: |
| ttnn.close_device(dev) |
|
|
