"""custom GELU timing vs tiles per core (grid size)."""
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()
try:
    L1 = ttnn.L1_MEMORY_CONFIG
    x = ttnn.from_torch(torch.randn(1, 1, 192, 5120), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1)
    for gx, gy in ((12, 10), (10, 6), (6, 5), (4, 3)):
        g = ttnn.CoreCoord(gx, gy)
        t = time_op(dev, lambda: tt_kernels.gelu(x, memory_config=L1, grid=g))
        tn = time_op(dev, lambda: tt_kernels.gelu(x, memory_config=L1, grid=g, defines=("HG_NOP",)))
        n = 960 / (gx * gy)
        print(f"[ggrid] {gx}x{gy}: {n:.0f} tiles/core  gelu {t:.2f} us  nop {tn:.2f} us  -> {(t - tn) / n:.3f} us/tile", flush=True)
finally:
    ttnn.close_device(dev)