| """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): |
| |
| 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 |
| if "row" in sys.argv: |
| for noc in (0, 1): |
| for sx_ in (0, 4, 9): |
| for lb in (0, 1): |
| size, reps = 32768, MB // 32768 |
| t = run([((sx_, 0), (0, 0, 9, 0))], size, reps, noc, lb=lb) |
| print(f"[mc] noc{noc} row0 x0-9 sender x{sx_} lb{lb}, {size}B x{reps}: {t:.2f} us -> {MB / (t - base) / 1e3:.1f} GB/s", flush=True) |
| t = run([((sx_, 0), (0, 0, 9, 0))], size, reps, noc, lb=lb, each=True) |
| print(f"[mc] noc{noc} row0 x0-9 sender x{sx_} lb{lb} barrier-each: {t:.2f} us -> {MB / (t - base) / 1e3:.1f} GB/s", flush=True) |
| sys.exit(0) |
| for noc in (0, 1): |
| for size in (16384, 32768): |
| reps = MB // size |
| |
| t = run([((x, 0), (x, 0, x, 5)) for x in range(10)], size, reps, noc) |
| print(f"[mc] noc{noc} 10 cols x6, sender top, {size}B x{reps}: {t:.2f} us -> {MB / (t - base) / 1e3:.1f} GB/s per group", flush=True) |
| t = run([((x, 3), (x, 0, x, 5)) for x in range(10)], size, reps, noc) |
| print(f"[mc] noc{noc} 10 cols x6, sender mid, {size}B x{reps}: {t:.2f} us -> {MB / (t - base) / 1e3:.1f} GB/s", flush=True) |
| t = run([((0, y), (0, y, 9, y)) for y in range(6)], size, reps, noc) |
| print(f"[mc] noc{noc} 6 rows x10, sender left, {size}B x{reps}: {t:.2f} us -> {MB / (t - base) / 1e3:.1f} GB/s", flush=True) |
| t = run([((x, 0), (x, 0, x, 9)) for x in range(12)], size, reps, noc) |
| print(f"[mc] noc{noc} 12 cols x10, sender top, {size}B x{reps}: {t:.2f} us -> {MB / (t - base) / 1e3:.1f} GB/s", flush=True) |
| t = run([((x, 0), (x, 0, x, 5)) for x in range(10)], size, reps, noc, lb=1) |
| print(f"[mc] noc{noc} 10 cols x6 loopback, {size}B x{reps}: {t:.2f} us -> {MB / (t - base) / 1e3:.1f} GB/s", flush=True) |
| t = run([((0, y), (0, y, 2, y)) for y in range(10)] + [((3, y), (3, y, 5, y)) for y in range(10)], size, reps, noc) |
| print(f"[mc] noc{noc} 20 row segs x3, {size}B x{reps}: {t:.2f} us -> {MB / (t - base) / 1e3:.1f} GB/s", flush=True) |
| t = run([((0, 0), (0, 0, 0, 5))], size, reps, noc) |
| print(f"[mc] noc{noc} single col x6 alone, {size}B x{reps}: {t:.2f} us -> {MB / (t - base) / 1e3:.1f} GB/s", flush=True) |
| finally: |
| ttnn.close_device(dev) |
|
|
