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
for shape in [(1, 1, 32, 32 * 120), (1, 1, 96, 32 * 120), (1, 1, 192, 5120), (1, 1, 384, 5120)]:
    x = ttnn.from_torch(torch.randn(*shape), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1)
    nt = shape[2] * shape[3] // 1024
    a = time_op(dev, lambda: tt_kernels.gelu(x, defines=("HG_NOP",)))
    b = time_op(dev, lambda: tt_kernels.gelu(x))
    c = time_op(dev, lambda: ttnn.relu(x, memory_config=L1))
    print(f"tiles {nt:5d} ({nt/120:.1f}/core): custom nop {a:5.1f}  custom gelu {b:5.1f}  ttnn relu {c:5.1f} us", flush=True)
    ttnn.deallocate(x)
ttnn.close_device(dev)