"""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)