"""raw multicast bandwidth on 12x10: 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 = 696 * 1024
    for noc in (0, 1):
        for size in (16384, 34816):
            for ns in (1, 2, 3, 6):
                reps = max(1, MB // size // ns)
                t = run([((x, y), (x, 0, x, 5)) for x in range(10) for y in range(ns)], size, reps, noc, lb=1)
                print(f"[mc] noc{noc} 10 cols x6, {ns} concurrent senders/col, loopback, {size}B x{reps} each: {t:.2f} us -> {ns*reps*size / (t - base) / 1e3:.1f} GB/s per col", flush=True)
            for ns in (1, 2, 5, 10):
                reps = max(1, 320 * 1024 // size // ns)
                t = run([((x, y), (0, y, 9, y)) for y in range(6) for x in range(ns)], size, reps, noc, lb=1)
                print(f"[mc] noc{noc} 6 rows x10, {ns} concurrent senders/row, loopback, {size}B x{reps} each: {t:.2f} us -> {ns*reps*size / (t - base) / 1e3:.1f} GB/s per row", flush=True)
finally:
    ttnn.close_device(dev)