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