| """fc1 (+bias) with fused GELU (accurate, in the matmul pack stage) vs linear + separate ttnn.gelu.""" |
| from common import * |
| dev = open_dev() |
| torch.manual_seed(0) |
| L1 = ttnn.L1_MEMORY_CONFIG |
| G = dev.compute_with_storage_grid_size(); cg = ttnn.CoreCoord(G.x, G.y) |
| ck = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=ttnn.MathFidelity.HiFi2, fp32_dest_acc_en=True, packer_l1_acc=True) |
| K, N = 1280, 5120 |
| x_t = torch.randn(1, 1, 192, K); w_t = torch.randn(K, N) / K ** 0.5 * 1.5; b_t = 0.3 * torch.randn(N) |
| ref = torch.nn.functional.gelu(x_t @ w_t + b_t) |
| w = ttnn.from_torch(w_t, dtype=ttnn.bfloat8_b, layout=ttnn.TILE_LAYOUT, device=dev) |
| b = ttnn.from_torch(b_t.reshape(1, 1, 1, N), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev) |
| x = ttnn.from_torch(x_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1) |
| def pc(act): |
| return ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig(compute_with_storage_grid_size=cg, in0_block_w=10, |
| out_subblock_h=2, out_subblock_w=2, per_core_M=6, per_core_N=2, fuse_batch=True, fused_activation=act, mcast_in0=True) |
| def sep(): |
| h = ttnn.linear(x, w, bias=b, program_config=pc(None), memory_config=L1, compute_kernel_config=ck) |
| g = ttnn.gelu(h, memory_config=L1); ttnn.deallocate(h); return g |
| acts = {"fused GELU(0)": ttnn.UnaryWithParam(ttnn.UnaryOpType.GELU, 0.0), |
| "fused GELU(1) fastlut": ttnn.UnaryWithParam(ttnn.UnaryOpType.GELU, 1.0)} |
| o = sep(); print(f"{'separate':28s} {time_op(dev, sep):7.1f} us pcc={pcc(ttnn.to_torch(o), ref):.6f} maxerr={float((ttnn.to_torch(o).float()-ref).abs().max()):.4f}", flush=True); ttnn.deallocate(o) |
| for k, a in acts.items(): |
| fn = lambda a=a: ttnn.linear(x, w, bias=b, program_config=pc(a), memory_config=L1, compute_kernel_config=ck) |
| try: |
| o = fn(); t = ttnn.to_torch(o).float() |
| print(f"{k:28s} {time_op(dev, fn):7.1f} us pcc={pcc(t, ref):.6f} maxerr={float((t-ref).abs().max()):.4f}", flush=True); ttnn.deallocate(o) |
| except Exception as e: |
| print(k, "FAIL", str(e).splitlines()[0][:200]) |
| ck16 = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=ttnn.MathFidelity.HiFi2, fp32_dest_acc_en=False, packer_l1_acc=True) |
| def pc16(act): |
| return ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig(compute_with_storage_grid_size=cg, in0_block_w=10, |
| out_subblock_h=2, out_subblock_w=2, per_core_M=6, per_core_N=2, fuse_batch=True, fused_activation=act, mcast_in0=True) |
| for k, a in [("bf16DST separate", None)] + list(acts.items()): |
| if a is None: |
| def fn(): |
| h = ttnn.linear(x, w, bias=b, program_config=pc16(None), memory_config=L1, compute_kernel_config=ck16) |
| g = ttnn.gelu(h, memory_config=L1); ttnn.deallocate(h); return g |
| else: |
| fn = lambda a=a: ttnn.linear(x, w, bias=b, program_config=pc16(a), memory_config=L1, compute_kernel_config=ck16) |
| o = fn(); t = ttnn.to_torch(o).float() |
| print(f"bf16DST {k:28s} {time_op(dev, fn):7.1f} us pcc={pcc(t, ref):.6f} maxerr={float((t-ref).abs().max()):.4f}", flush=True); ttnn.deallocate(o) |
| ttnn.close_device(dev) |
|
|
