"""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)
# LayerNorm [32,1024]
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)
# residual add width-sharded
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)