"""Round 15 (D-fc1 go/no-go ceiling): the served fc1 (library 1D matmul, model config: in0 width-sharded on 8 cores,
8x10 grid, bf16 in0, bfp8 weights, HiFi2 + fp32 acc) with its weights in DRAM (served) vs L1 interleaved (= the upper
bound of any weight prefetch / streaming into L1).  Also qkv's library config for reference.  usage: fc1_l1w.py"""
import sys
from pathlib import Path
from common import *
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from hamer.vit import tuned_matmul_configs
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:
    for name, N, ws, xdt, fg in (("fc1", 5120, 8, ttnn.bfloat16, "8x10"), ("fc1", 5120, 8, ttnn.bfloat16, "12x10")):
        pc = tuned_matmul_configs(G, fp32_acc=True, in0_ws=ws, fc1_grid=fg)[name]
        sg = ttnn.num_cores_to_corerangeset(ws, G, row_wise=True)
        hmem = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.L1,
                                 ttnn.ShardSpec(sg, [192, 1280 // ws], ttnn.ShardOrientation.ROW_MAJOR))
        x = ttnn.from_torch(torch.randn(1, 1, 192, 1280), dtype=xdt, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=hmem)
        b = ttnn.from_torch(torch.randn(1, 1, 1, N) * 0.1, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev)
        wt = torch.randn(1280, N) / 1280 ** 0.5
        outs = {}
        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, bias=b, program_config=pc, memory_config=L1, compute_kernel_config=ck)
            o = fn()
            outs[wloc] = ttnn.to_torch(o).float()
            ttnn.deallocate(o)
            print(f"[fc1w] {name} grid {fg} in0 ws{ws} w={wloc:4s} {time_op(dev, fn):7.2f} us", flush=True)
            ttnn.deallocate(w)
        print(f"[fc1w] dram == l1 outputs: {bool(torch.equal(outs['dram'], outs['l1']))}", flush=True)
        ttnn.deallocate(x)
finally:
    ttnn.close_device(dev)