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