| from common import * |
| dev = open_dev() |
| torch.manual_seed(0) |
| x_t = torch.randn(1, 1, 192, 1280); g_t = 1 + .1 * torch.randn(1280); b_t = .1 * torch.randn(1280) |
| ref = torch.nn.functional.layer_norm(x_t, (1280,), g_t, b_t, 1e-6) |
| gw = ttnn.from_torch(g_t.reshape(1, 1, -1, 32), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev) |
| bw = ttnn.from_torch(b_t.reshape(1, 1, -1, 32), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev) |
| for gx, gy in [(10, 6), (8, 6), (5, 6), (4, 6), (2, 6), (10, 3), (10, 2), (5, 3)]: |
| mc = ttnn.create_sharded_memory_config((192, 1280), core_grid=ttnn.CoreGrid(y=gy, x=gx), strategy=ttnn.ShardStrategy.BLOCK, orientation=ttnn.ShardOrientation.ROW_MAJOR) |
| xs = ttnn.from_torch(x_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=mc) |
| bwd = 40 // gx; bh = 6 // gy |
| for sbw in sorted({d for d in range(1, min(bwd, 4) + 1) if bwd % d == 0}): |
| pc = ttnn.LayerNormShardedMultiCoreProgramConfig(compute_with_storage_grid_size=(gx, gy), subblock_w=sbw, block_h=bh, block_w=bwd, inplace=False) |
| try: |
| o = ttnn.layer_norm(xs, weight=gw, bias=bw, epsilon=1e-6, memory_config=mc, program_config=pc) |
| p = pcc(ttnn.to_torch(o), ref) |
| us = time_op(dev, lambda pc=pc, xs=xs, mc=mc: ttnn.layer_norm(xs, weight=gw, bias=bw, epsilon=1e-6, memory_config=mc, program_config=pc)) |
| print(f"LN block {gx}x{gy} sbw{sbw}: {us:.1f} us pcc={p:.5f}") |
| except Exception as e: |
| print(f"LN block {gx}x{gy} sbw{sbw}: FAIL {str(e).splitlines()[0][:120]}") |
| ttnn.close_device(dev) |
|
|
