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