| """MANO head (M = 1 tile) ops: DRAM-sharded matmuls, width-sharded LN, vs the current minimal_matmul / 1-core LN.""" |
| from common import * |
| dev = open_dev() |
| torch.manual_seed(0) |
| L1 = ttnn.L1_MEMORY_CONFIG |
| DR = ttnn.DRAM_MEMORY_CONFIG |
| nb = dev.dram_grid_size().x |
| def rep(tag, fn, ref=None, n=None): |
| try: |
| o = fn(); ot = ttnn.to_torch(o).float() |
| if n: ot = ot[..., :n] |
| p = pcc(ot[..., :1, :], ref[..., :1, :]) if ref is not None else float('nan') |
| us = time_op(dev, fn) |
| print(f"{tag:70s} {us:8.1f} us pcc={p:.5f}", flush=True) |
| except Exception as e: |
| print(f"{tag:70s} FAIL {str(e).splitlines()[0][:200]}", flush=True) |
| ck = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=ttnn.MathFidelity.HiFi2, fp32_dest_acc_en=False, packer_l1_acc=True) |
| ck4 = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=ttnn.MathFidelity.HiFi4, fp32_dest_acc_en=True, packer_l1_acc=True) |
| def wsh(ncores, width): |
| cg = ttnn.num_cores_to_corerangeset(ncores, ttnn.CoreCoord(12, 10), row_wise=True) |
| return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.L1, ttnn.ShardSpec(cg, [32, width // ncores], ttnn.ShardOrientation.ROW_MAJOR)) |
| for K, N in [(1024, 512), (512, 1024), (1024, 1024), (1024, 128)]: |
| x_t = torch.randn(1, 1, 32, K); w_t = torch.randn(K, N) / K ** .5; b_t = torch.randn(1, 1, 1, N) * .1 |
| ref = x_t @ w_t + b_t |
| xi = ttnn.from_torch(x_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=DR) |
| wi = ttnn.from_torch(w_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev) |
| b = ttnn.from_torch(b_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev) |
| cfg = ttnn.MinimalMatmulConfig(compute_with_storage_grid_size=ttnn.CoreCoord(12, 10), M_block_size=1, K_block_size=8, N_block_size=4, subblock_h=1, subblock_w=4) |
| rep(f"K{K} N{N} minimal_matmul (current)", lambda: ttnn.experimental.minimal_matmul(xi, wi, bias_tensor=b, config=cfg, compute_kernel_config=ck), ref) |
| rep(f"K{K} N{N} linear default", lambda: ttnn.linear(xi, wi, bias=b, compute_kernel_config=ck), ref) |
| Np = -(-N // (32 * nb)) * 32 * nb |
| for wd, wn in [(ttnn.bfloat16, "bf16"), (ttnn.bfloat8_b, "bfp8")]: |
| grid = ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(nb - 1, 0))}) |
| wmc = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, ttnn.ShardSpec(grid, [K, Np // nb], ttnn.ShardOrientation.ROW_MAJOR)) |
| w = ttnn.from_torch(torch.nn.functional.pad(w_t, (0, Np - N)), dtype=wd, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=wmc) |
| bp = ttnn.from_torch(torch.nn.functional.pad(b_t, (0, Np - N)), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev) |
| for nc in [8, 16, 32]: |
| if (K // 32) % nc or (Np // 32) % nc: |
| continue |
| xs = ttnn.to_memory_config(xi, wsh(nc, K)) |
| for ibw in sorted({K // 32 // nc, max(1, K // 32 // nc // 2)}): |
| pc = ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig(in0_block_w=ibw, per_core_M=1, per_core_N=Np // 32 // nc, fused_activation=None) |
| for ckn, c in [("HiFi2", ck), ("HiFi4fp32", ck4)]: |
| rep(f"K{K} N{N} dram-sharded w={wn} cores{nc} ibw{ibw} {ckn}", lambda pc=pc, xs=xs, c=c, nc=nc: ttnn.linear(xs, w, bias=bp, program_config=pc, memory_config=wsh(nc, Np), compute_kernel_config=c), ref, n=N) |
| |
| x_t = torch.randn(1, 1, 32, 1024); g_t = 1 + .1 * torch.randn(1024); bt = .1 * torch.randn(1024) |
| ref = torch.nn.functional.layer_norm(x_t, (1024,), g_t, bt, 1e-5) |
| xi = ttnn.from_torch(x_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=DR) |
| g = ttnn.from_torch(g_t.reshape(1, 1, 1, -1), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev) |
| bb = ttnn.from_torch(bt.reshape(1, 1, 1, -1), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev) |
| rep("LN [32,1024] interleaved DRAM (current)", lambda: ttnn.layer_norm(xi, weight=g, bias=bb, epsilon=1e-5), ref) |
| 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(bt.reshape(1, 1, -1, 32), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev) |
| for nc in [4, 8, 16, 32]: |
| mc = wsh(nc, 1024) |
| xs = ttnn.to_memory_config(xi, mc) |
| pc = ttnn.LayerNormShardedMultiCoreProgramConfig(compute_with_storage_grid_size=(12, 10), subblock_w=32 // nc, block_h=1, block_w=32 // nc, inplace=False) |
| rep(f"LN width-sharded {nc} cores", lambda xs=xs, pc=pc, mc=mc: ttnn.layer_norm(xs, weight=gw, bias=bw, epsilon=1e-5, memory_config=mc, program_config=pc), ref) |
| |
| mc = wsh(8, 1024) |
| a = ttnn.to_memory_config(xi, mc) |
| rep("add width-sharded 8 cores [32,1024]", lambda: ttnn.add(a, a, memory_config=mc), x_t * 2) |
| rep("add interleaved DRAM [32,1024]", lambda: ttnn.add(xi, xi), x_t * 2) |
| rep("i2s [32,1024] -> 8 cores", lambda: ttnn.to_memory_config(xi, mc), x_t) |
| ttnn.close_device(dev) |
|
|
