| """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): |
| |
| 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) |
|
|
