"""Round 12: raw DRAM read ceilings for the xmm fc1 weight-sender placements (kernels/mcbench/dram_rd.cpp, all reads issued
at once, one barrier).  usage: dram_fc1.py"""
import sys
from pathlib import Path
from common import *
from open_device_12x10 import open_device
K = str(Path(__file__).resolve().parents[2] / "hamer/kernels/mcbench/dram_rd.cpp")
dev = open_device(grid="12x10", dispatch="eth", trace_region_size=64 << 20)
try:
    nb = dev.dram_grid_size().x
    SH = 4 << 20
    cg = ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(nb - 1, 0))})
    mem = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.BufferType.DRAM,
                            ttnn.ShardSpec(cg, [SH // 64, 32], ttnn.ShardOrientation.ROW_MAJOR))
    src = ttnn.from_torch(torch.zeros(nb * SH // 64, 32), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=mem)
    base_addr = src.buffer_address()
    cores = ttnn.CoreRangeSet([ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(11, 9))])
    LS = 256 * 1024
    def run(rd0, rd1):
        kds = []
        tot = 0
        for noc, readers in ((0, rd0), (1, rd1)):
            args = ttnn.RuntimeArgs()
            for x in range(12):
                for y in range(10):
                    args[x][y] = [0] * 8
            for (c, b0, chunk, reps) in readers:
                args[c[0]][c[1]] = [1, b0, nb, base_addr, chunk, reps, chunk, LS]
                tot += chunk * reps
            cfg = ttnn.ReaderConfigDescriptor() if noc == 0 else ttnn.WriterConfigDescriptor()
            kds.append(ttnn.KernelDescriptor(kernel_source=K, source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
                                             core_ranges=cores, compile_time_args=[], defines=[], runtime_args=args, config=cfg))
        cb = ttnn.CBDescriptor(total_size=2 * LS, core_ranges=cores, format_descriptors=[
            ttnn.CBFormatDescriptor(buffer_index=0, data_format=ttnn.bfloat16, page_size=2048)])
        prog = ttnn.ProgramDescriptor(kernels=kds, semaphores=[], cbs=[cb])
        f = lambda: (ttnn.generic_op([src, src], prog), None)[1]
        return time_op(dev, f, reps=10, iters=5), tot
    base, _ = run([], [])
    print(f"[dram] empty {base:.2f} us", flush=True)
    def rep(name, r0, r1=()):
        t, tot = run(list(r0), list(r1))
        print(f"[dram] {name}: {t:.2f} us  {tot / 1e3 / max(t - base, 1e-3):.0f} GB/s (tot {tot >> 10} KB)", flush=True)
    B = 17408
    top = [(x, 3 * s) for s in range(3) for x in range(12)] + [(3 * g, 9) for g in range(4)]
    bot = [(x, 3 * s + 2) for s in range(3) for x in range(12)] + [(3 * g + 2, 9) for g in range(4)]
    rep("40 senders NOC0 10 x 17KB", [(c, i % nb, B, 10) for i, c in enumerate(top)])
    rep("40 top NOC0 5 + 40 bottom NOC1 5", [(c, i % nb, B, 5) for i, c in enumerate(top)], [(c, (i + 3) % nb, B, 5) for i, c in enumerate(bot)])
    rep("80 (top+bottom) NOC0 5 each", [(c, i % nb, B, 5) for i, c in enumerate(top + bot)])
    g120 = [(x, y) for y in range(10) for x in range(12)]
    rep("120 NOC0 1/3 each", [(c, i % nb, B, 10) for i, c in enumerate(g120[:120])][:120] and [(c, i % nb, 5803, 10) for i, c in enumerate(g120)])
    rep("60 NOC0 + 60 NOC1", [(c, i % nb, 5803 * 2, 5) for i, c in enumerate(g120[:60])], [(c, i % nb, 5803 * 2, 5) for i, c in enumerate(g120[60:])])
    rep("120: 60 NOC0 even rows-ish + 60 NOC1", [(c, i % nb, 5803, 10) for i, c in enumerate(g120) if c[1] % 2 == 0] + [],
        [(c, i % nb, 5803 * 1, 10) for i, c in enumerate(g120) if c[1] % 2 == 1])
finally:
    ttnn.close_device(dev)