| """custom GELU kernel: accuracy vs torch erf gelu and timing vs ttnn.gelu.""" |
| 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 |
| small = len(sys.argv) > 1 and sys.argv[1] == "small" |
| shape = (1, 1, 64, 64) if small else (1, 1, 192, 5120) |
| torch.manual_seed(0) |
| xt = torch.randn(*shape) * 3.0 |
| xt.view(-1)[:4096 if not small else 64] = torch.linspace(-8, 8, 4096 if not small else 64) |
| x = ttnn.from_torch(xt, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1) |
| xb = xt.bfloat16().float() |
| ref = torch.nn.functional.gelu(xb) |
| ref_bf = ref.bfloat16().float() |
| o = ttnn.to_torch(tt_kernels.gelu(x)).float() |
| s = ttnn.to_torch(ttnn.gelu(x, memory_config=L1)).float() |
| for name, t in [("custom", o), ("ttnn exact", s)]: |
| d = (t - ref).abs() |
| print(f"{name:11s} max|err| {float(d.max()):.3e} mean|err| {float(d.mean()):.3e} bit-equal-to-bf16(ref) {float((t == ref_bf).float().mean()):.4f} pcc {pcc(t, ref):.8f}", flush=True) |
| if not small: |
| print(f"custom gelu {time_op(dev, lambda: tt_kernels.gelu(x)):6.1f} us", flush=True) |
| print(f"custom nop {time_op(dev, lambda: tt_kernels.gelu(x, defines=('HG_NOP',))):6.1f} us", flush=True) |
| print(f"custom bf16dst {time_op(dev, lambda: tt_kernels.gelu(x, fp32_dest=False)):6.1f} us", flush=True) |
| o2 = ttnn.to_torch(tt_kernels.gelu(x, fp32_dest=False)).float(); d = (o2 - ref).abs() |
| print(f"bf16dst max|err| {float(d.max()):.3e} mean {float(d.mean()):.3e} bit-eq {float((o2 == ref_bf).float().mean()):.4f}") |
| print(f"ttnn gelu {time_op(dev, lambda: ttnn.gelu(x, memory_config=L1)):6.1f} us", flush=True) |
| ttnn.close_device(dev) |
|
|
