"""Round 15: the MANO-head K/V projection (library 1D matmul, [192, 1280] x [1280, 6144] bfp8, 12x10, in0_block_w 5,
per_core_N 2, HiFi2 + fp32 acc) with its weights in DRAM (served) vs L1 interleaved: is it weight-stream-bound?"""
import sys
from pathlib import Path
from common import *
dev = open_dev()
torch.manual_seed(0)
L1, DR = ttnn.L1_MEMORY_CONFIG, ttnn.DRAM_MEMORY_CONFIG
G = dev.compute_with_storage_grid_size()
ck = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=ttnn.MathFidelity.HiFi2, fp32_dest_acc_en=True,
                                            packer_l1_acc=True)
try:
    N = 6144
    x = ttnn.from_torch(torch.randn(1, 1, 192, 1280), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1)
    wt = torch.randn(1280, N) / 1280 ** 0.5
    for ibw in (5, 10, 4):
        pc = ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig(
            compute_with_storage_grid_size=ttnn.CoreCoord(G.x, G.y), in0_block_w=ibw, out_subblock_h=1, out_subblock_w=2,
            per_core_M=6, per_core_N=2, fuse_batch=True, fused_activation=None, mcast_in0=True)
        for wloc in ("dram", "l1"):
            w = ttnn.from_torch(wt, dtype=ttnn.bfloat8_b, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=DR if wloc == "dram" else L1)
            fn = lambda w=w: ttnn.linear(x, w, program_config=pc, memory_config=L1, compute_kernel_config=ck)
            print(f"[kvw] ibw {ibw} w={wloc:4s} {time_op(dev, fn):7.2f} us", flush=True)
            ttnn.deallocate(w)
finally:
    ttnn.close_device(dev)