"""single-sender 2D-rectangle multicast bandwidth (round 6): usage mc_bench.py [eth|worker]"""
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/mc.cpp")
disp = sys.argv[1] if len(sys.argv) > 1 else "eth"
dev = open_device(grid="12x10", dispatch=disp, trace_region_size=64 << 20)
try:
    cores = ttnn.CoreRangeSet([ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(11, 9))])
    dummy = ttnn.from_torch(torch.zeros(1, 1, 32, 32), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev)
    phys = {(x, y): dev.worker_core_from_logical_core(ttnn.CoreCoord(x, y)) for x in range(12) for y in range(10)}
    def run(groups, size, reps, noc, lb=0, each=False):
        # groups: list of (sender logical, rect logical (x0,y0,x1,y1))
        args = ttnn.RuntimeArgs()
        for x in range(12):
            for y in range(10):
                args[x][y] = [0, 0, 0, 0, 0, 0, 0, 0, 0]
        for (s, (x0, y0, x1, y1)) in groups:
            p0, p1 = phys[(x0, y0)], phys[(x1, y1)]
            n = (x1 - x0 + 1) * (y1 - y0 + 1)
            inside = x0 <= s[0] <= x1 and y0 <= s[1] <= y1
            nd = n if lb else n - (1 if inside else 0)
            args[s[0]][s[1]] = [1, min(p0.x, p1.x), min(p0.y, p1.y), max(p0.x, p1.x), max(p0.y, p1.y), nd, size, reps, lb]
        cfg = ttnn.ReaderConfigDescriptor() if noc == 0 else ttnn.WriterConfigDescriptor()
        kd = ttnn.KernelDescriptor(kernel_source=K, source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH, core_ranges=cores,
                                   compile_time_args=[], defines=[("MC_BARRIER_EACH", "1")] if each else [], runtime_args=args, config=cfg)
        cb = ttnn.CBDescriptor(total_size=size, core_ranges=cores, format_descriptors=[
            ttnn.CBFormatDescriptor(buffer_index=0, data_format=ttnn.bfloat16, page_size=2048)])
        prog = ttnn.ProgramDescriptor(kernels=[kd], semaphores=[], cbs=[cb])
        f = lambda: (ttnn.generic_op([dummy, dummy], prog), None)[1]
        return time_op(dev, f, reps=10, iters=5)
    base = run([], 32768, 1, 0)
    print(f"[mc] empty {base:.2f} us", flush=True)
    MB = 480 * 1024
    for noc in (0, 1):
        for size in (8192, 49152):
            reps = MB // size
            for (s, r, nm) in (((0, 0), (0, 0, 11, 5), "12x6 from (0,0)"), ((5, 0), (0, 0, 11, 5), "12x6 from (5,0)"),
                               ((0, 0), (0, 0, 11, 9), "12x10 from (0,0)"), ((0, 0), (0, 0, 0, 5), "col x6"),
                               ((0, 0), (0, 0, 11, 0), "row x12")):
                t = run([(s, r)], size, reps, noc, lb=1)
                print(f"[mc] noc{noc} {nm} {size}B x{reps}: {t:.2f} us -> {MB / (t - base) / 1e3:.1f} GB/s", flush=True)
finally:
    ttnn.close_device(dev)