File size: 853 Bytes
58ba6df | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 | """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)
|