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