"""DRAM-sharded matmul for the ViT shapes (M=192)."""
import sys
from common import *

dev = open_dev()
torch.manual_seed(0)
L1 = ttnn.L1_MEMORY_CONFIG
nb = dev.dram_grid_size().x
print("dram banks", dev.dram_grid_size())
SHAPES = {"qkv": (1280, 4608), "proj": (1536, 1280), "fc1": (1280, 5120), "fc2": (5120, 1280)}
for name, (K, N) in [(n, SHAPES[n]) for n in (sys.argv[1:] or SHAPES)]:
    x_t = torch.randn(1, 1, 192, K); w_t = torch.randn(K, N) / K ** 0.5
    ref = x_t @ w_t
    Np = -(-N // (32 * nb)) * 32 * nb
    w_pad = torch.nn.functional.pad(w_t, (0, Np - N))
    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(w_pad, dtype=wd, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=wmc)
        for ncores in [8, 10, 16, 20, 32, 40, 64, 80]:
            Kt = K // 32
            if Kt % ncores or (Np // 32) % ncores:
                continue
            core_grid = ttnn.num_cores_to_corerangeset(ncores, ttnn.CoreCoord(12, 10), row_wise=True)
            in0mc = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.L1,
                                      ttnn.ShardSpec(core_grid, [192, K // ncores], ttnn.ShardOrientation.ROW_MAJOR))
            x = ttnn.from_torch(x_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=in0mc)
            outmc = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.L1)
            for ibw_div in [1, 2, 4]:
                ibw = Kt // ncores // ibw_div
                if ibw < 1 or (Kt // ncores) % ibw_div:
                    continue
                pc = ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig(
                    in0_block_w=ibw, per_core_M=6, per_core_N=Np // 32 // ncores, fused_activation=None)
                for fid in ["HiFi2", "LoFi"]:
                    ck = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=getattr(ttnn.MathFidelity, fid),
                                                                fp32_dest_acc_en=False, packer_l1_acc=True)
                    fn = lambda pc=pc, ck=ck, x=x: ttnn.matmul(x, w, program_config=pc, memory_config=outmc, compute_kernel_config=ck)
                    tag = f"{name} w={wn} cores{ncores} ibw{ibw} {fid}"
                    try:
                        o = fn(); p = pcc(ttnn.to_torch(o)[..., :N], ref); ttnn.deallocate(o)
                        us = time_op(dev, fn)
                        print(f"{tag:45s} {us:8.1f} us pcc={p:.5f}", flush=True)
                    except Exception as e:
                        print(f"{tag:45s} FAIL {str(e).splitlines()[0][:150]}", flush=True)
            ttnn.deallocate(x)
        ttnn.deallocate(w)
ttnn.close_device(dev)