"""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)