"""custom GELU kernel: accuracy vs torch erf gelu and timing vs ttnn.gelu."""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from common import *
from hamer import tt_kernels
dev = open_dev()
L1 = ttnn.L1_MEMORY_CONFIG
small = len(sys.argv) > 1 and sys.argv[1] == "small"
shape = (1, 1, 64, 64) if small else (1, 1, 192, 5120)
torch.manual_seed(0)
xt = torch.randn(*shape) * 3.0
xt.view(-1)[:4096 if not small else 64] = torch.linspace(-8, 8, 4096 if not small else 64)
x = ttnn.from_torch(xt, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1)
xb = xt.bfloat16().float()
ref = torch.nn.functional.gelu(xb)
ref_bf = ref.bfloat16().float()
o = ttnn.to_torch(tt_kernels.gelu(x)).float()
s = ttnn.to_torch(ttnn.gelu(x, memory_config=L1)).float()
for name, t in [("custom", o), ("ttnn exact", s)]:
    d = (t - ref).abs()
    print(f"{name:11s} max|err| {float(d.max()):.3e}  mean|err| {float(d.mean()):.3e}  bit-equal-to-bf16(ref) {float((t == ref_bf).float().mean()):.4f}  pcc {pcc(t, ref):.8f}", flush=True)
if not small:
    print(f"custom gelu {time_op(dev, lambda: tt_kernels.gelu(x)):6.1f} us", flush=True)
    print(f"custom nop  {time_op(dev, lambda: tt_kernels.gelu(x, defines=('HG_NOP',))):6.1f} us", flush=True)
    print(f"custom bf16dst {time_op(dev, lambda: tt_kernels.gelu(x, fp32_dest=False)):6.1f} us", flush=True)
    o2 = ttnn.to_torch(tt_kernels.gelu(x, fp32_dest=False)).float(); d = (o2 - ref).abs()
    print(f"bf16dst max|err| {float(d.max()):.3e} mean {float(d.mean()):.3e} bit-eq {float((o2 == ref_bf).float().mean()):.4f}")
    print(f"ttnn gelu   {time_op(dev, lambda: ttnn.gelu(x, memory_config=L1)):6.1f} us", flush=True)
ttnn.close_device(dev)