| """ETH-dispatch config cost probe: trace of reps x [short op S, op X] vs reps x X alone. |
| pair - (X alone) - S alone ~= exposed dispatch cost of X after a short op. usage: disp_cost.py eth|worker""" |
| import sys |
| from common import * |
| from open_device_12x10 import open_device |
| disp = sys.argv[1] if len(sys.argv) > 1 else "eth" |
| dev = open_device(grid="12x10", dispatch=disp, trace_region_size=256 << 20) |
| torch.manual_seed(0) |
| L1 = ttnn.L1_MEMORY_CONFIG |
| G = dev.compute_with_storage_grid_size(); cg = ttnn.CoreCoord(G.x, G.y) |
| ck = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=ttnn.MathFidelity.HiFi2, fp32_dest_acc_en=True, packer_l1_acc=True) |
| res = ttnn.create_sharded_memory_config((192, 1280), core_grid=ttnn.CoreGrid(y=6, x=10), strategy=ttnn.ShardStrategy.BLOCK, |
| orientation=ttnn.ShardOrientation.ROW_MAJOR) |
| xs = ttnn.from_torch(torch.randn(1, 1, 192, 1280), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=res) |
| xi = ttnn.from_torch(torch.randn(1, 1, 192, 1280), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1) |
| S = lambda: ttnn.sharded_to_interleaved(xs, L1) |
| cases = {} |
| for pcn in (2, 3, 4, 6, 8): |
| N = 4608 |
| w = ttnn.from_torch(torch.randn(1280, N) / 36, dtype=ttnn.bfloat8_b, layout=ttnn.TILE_LAYOUT, device=dev) |
| b = ttnn.from_torch(torch.randn(1, 1, 1, N), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev) |
| sw = max(d for d in (1, 2, 4) if pcn % d == 0) |
| pc = ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig(compute_with_storage_grid_size=cg, in0_block_w=10, |
| out_subblock_h=1, out_subblock_w=sw, per_core_M=6, per_core_N=pcn, fuse_batch=True, fused_activation=None, mcast_in0=True) |
| cases[f"qkv1d pcn{pcn} cores{-(-144//pcn)}"] = (lambda w=w, b=b, pc=pc: ttnn.linear(xi, w, bias=b, program_config=pc, memory_config=L1, compute_kernel_config=ck)) |
| cases[f"qkv1d nobias pcn{pcn}"] = (lambda w=w, pc=pc: ttnn.linear(xi, w, program_config=pc, memory_config=L1, compute_kernel_config=ck)) |
| ga = ttnn.from_torch(torch.ones(1, 1, 1, 1280), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev) |
| ga = ttnn.from_torch(torch.ones(1280).reshape(1, 1, -1, 32), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev) |
| lnpc = ttnn.LayerNormShardedMultiCoreProgramConfig(compute_with_storage_grid_size=(10, 6), subblock_w=4, block_h=1, block_w=4, inplace=False) |
| cases["ln sharded"] = lambda: ttnn.layer_norm(xs, weight=ga, bias=ga, epsilon=1e-6, memory_config=res, program_config=lnpc) |
| cases["add sharded"] = lambda: ttnn.add(xs, xs, memory_config=res) |
| cases["s2i"] = S |
| cases["gelu 192x5120"] = (lambda g=ttnn.from_torch(torch.randn(1, 1, 192, 5120), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1): ttnn.gelu(g, memory_config=L1)) |
| for name, X in cases.items(): |
| try: |
| tx = time_op(dev, X); ts = time_op(dev, S) |
| tp = time_op(dev, lambda: [S(), X()]) |
| print(f"{disp:6s} {name:28s} X={tx:6.2f} S={ts:5.2f} pair={tp:6.2f} exposed={tp - tx - ts:6.2f} us", flush=True) |
| except Exception as e: |
| print(name, "FAIL", str(e).splitlines()[0][:200], flush=True) |
| ttnn.close_device(dev) |
|
|
