| """GELU variants: separate op vs fused into the fc1 1D-mcast matmul.""" |
| from common import * |
| dev = open_dev() |
| torch.manual_seed(0) |
| L1 = ttnn.L1_MEMORY_CONFIG |
| def T(t, dt=ttnn.bfloat16, mem=L1): |
| return ttnn.from_torch(t, dtype=dt, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=mem) |
| def rep(tag, fn, ref=None): |
| try: |
| o = fn(); p = pcc(ttnn.to_torch(o), ref) if ref is not None else float('nan') |
| mx = (ttnn.to_torch(o).float() - ref).abs().max().item() if ref is not None else float('nan') |
| us = time_op(dev, fn) |
| print(f"{tag:60s} {us:8.1f} us pcc={p:.6f} maxabs={mx:.4f}", flush=True) |
| except Exception as e: |
| print(f"{tag:60s} FAIL {str(e).splitlines()[0][:200]}", flush=True) |
| h_t = torch.randn(1, 1, 192, 5120) * 1.5; h = T(h_t) |
| ref = torch.nn.functional.gelu(h_t) |
| for v in ["Accurate", "Tanh", "FastLut"]: |
| rep(f"gelu {v}", lambda v=v: ttnn.gelu(h, variant=getattr(ttnn.GeluVariant, v), memory_config=L1), ref) |
| x_t = torch.randn(1, 1, 192, 1280); x = T(x_t) |
| w_t = torch.randn(1280, 5120) / 1280 ** 0.5; bb_t = 0.1 * torch.randn(1, 1, 1, 5120) |
| w = ttnn.from_torch(w_t, dtype=ttnn.bfloat8_b, layout=ttnn.TILE_LAYOUT, device=dev); bb = T(bb_t, mem=None) |
| ck = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=ttnn.MathFidelity.HiFi2, fp32_dest_acc_en=False, packer_l1_acc=True) |
| ref = torch.nn.functional.gelu(x_t @ w_t + bb_t) |
| for act in [None, ttnn.UnaryWithParam(ttnn.UnaryOpType.GELU_TANH), ttnn.UnaryWithParam(ttnn.UnaryOpType.GELU, 0.0)]: |
| pc = ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig(compute_with_storage_grid_size=ttnn.CoreCoord(12, 10), in0_block_w=10, out_subblock_h=3, out_subblock_w=2, per_core_M=6, per_core_N=2, fuse_batch=True, fused_activation=act, mcast_in0=True) |
| if act is None: |
| rep("fc1 lin1d + gelu Accurate", lambda pc=pc: ttnn.gelu(ttnn.linear(x, w, bias=bb, program_config=pc, memory_config=L1, compute_kernel_config=ck), memory_config=L1), ref) |
| rep("fc1 lin1d + gelu Tanh", lambda pc=pc: ttnn.gelu(ttnn.linear(x, w, bias=bb, program_config=pc, memory_config=L1, compute_kernel_config=ck), variant=ttnn.GeluVariant.Tanh, memory_config=L1), ref) |
| rep("fc1 lin1d only", lambda pc=pc: ttnn.linear(x, w, bias=bb, program_config=pc, memory_config=L1, compute_kernel_config=ck), x_t @ w_t + bb_t) |
| else: |
| rep(f"fc1 lin1d fused {act.op_type}", lambda pc=pc: ttnn.linear(x, w, bias=bb, program_config=pc, memory_config=L1, compute_kernel_config=ck), ref) |
| ttnn.close_device(dev) |
|
|
