| """Two-hop activation broadcast (kernels/mcbench/bcast2.cpp), round 7. |
| usage: bcast_test.py mode=rc|cr [cols=12] [rows=10] [sb=49152] [noc=0|1] [verify=1] [time=1] [nodata=0] |
| rc: shard r starts at (0, r); hop A = row multicast, hop B = chain of column multicasts (12 columns in parallel) |
| cr: shard c starts at (c, 0) (c < NSH); hop A = column multicast, hop B = chain of row multicasts (rows in parallel)""" |
| import sys |
| from pathlib import Path |
| sys.path.insert(0, str(Path(__file__).resolve().parents[2])) |
| from common import * |
|
|
| K = str(Path(__file__).resolve().parents[2] / "hamer/kernels/mcbench/bcast2.cpp") |
| args = dict(a.split("=") for a in sys.argv[1:]) |
| mode = args.get("mode", "rc") |
| C, R = int(args.get("cols", 12)), int(args.get("rows", 10)) |
| SB = int(args.get("sb", 49152)) |
| noc = int(args.get("noc", 0)) |
| verify = args.get("verify", "1") == "1" |
| NSH = R if mode == "rc" else int(args.get("nsh", 10)) |
| assert NSH <= 14 |
|
|
| dev = open_dev() |
| try: |
| phys = {(x, y): dev.worker_core_from_logical_core(ttnn.CoreCoord(x, y)) for x in range(C) for y in range(R)} |
| cores = ttnn.CoreRangeSet([ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(C - 1, R - 1))]) |
|
|
| def rect(x0, y0, x1, y1): |
| p0, p1 = phys[(x0, y0)], phys[(x1, y1)] |
| return [min(p0.x, p1.x), min(p0.y, p1.y), max(p0.x, p1.x), max(p0.y, p1.y)] |
|
|
| out = ttnn.from_torch(torch.zeros(C * R, 8, dtype=torch.int32), dtype=ttnn.uint32, layout=ttnn.ROW_MAJOR_LAYOUT, |
| device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
|
|
| def build(empty=False): |
| rt = ttnn.RuntimeArgs() |
| for x in range(C): |
| for y in range(R): |
| a = [0xff, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff, 0, 0, 0, 0, 0, 0, 0, 0, 0, SB, out.buffer_address(), |
| y * C + x] |
| if not empty: |
| if mode == "rc": |
| |
| p = phys[(0, y)] |
| a[0] = y |
| a[1] = 1 if x == 0 else 0 |
| a[2:6] = rect(0, y, C - 1, y); a[6] = C |
| a[7:11] = rect(x, 0, x, R - 1); a[11] = R |
| a[12] = y; a[13] = y - 1 if y else 0 |
| a[14] = 0 if x == 0 else 1; a[15:17] = [p.x, p.y] |
| q = phys[(x, 0)] |
| a[17] = 0 if y == 0 else 1; a[18:20] = [q.x, q.y]; a[20] = 1 if y == 0 else 0 |
| else: |
| |
| if x < NSH: |
| p = phys[(x, 0)] |
| a[0] = x |
| a[1] = 1 if y == 0 else 0 |
| a[2:6] = rect(x, 0, x, R - 1); a[6] = R |
| a[12] = x; a[13] = x - 1 if x else 0 |
| a[14] = 0 if y == 0 else 1; a[15:17] = [p.x, p.y] |
| a[7:11] = rect(0, y, C - 1, y); a[11] = C |
| q = phys[(0, y)] |
| a[17] = 0 if x == 0 else 1; a[18:20] = [q.x, q.y]; a[20] = 1 if x == 0 else 0 |
| a[21] = NSH |
| rt[x][y] = a |
| cfg = ttnn.ReaderConfigDescriptor() if noc == 0 else ttnn.WriterConfigDescriptor() |
| defs = ([("BC_VERIFY", "1")] if verify else []) + ([("BC_NO_DATA", "1")] if args.get("nodata") == "1" else []) \ |
| + [(d, "1") for d in filter(None, args.get("defs", "").split(","))] |
| kd = ttnn.KernelDescriptor(kernel_source=K, source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH, core_ranges=cores, |
| compile_time_args=list(ttnn.TensorAccessorArgs(out).get_compile_time_args()), |
| defines=defs, runtime_args=rt, config=cfg) |
| cbs = [ttnn.CBDescriptor(total_size=NSH * SB if not empty else 2048, core_ranges=cores, format_descriptors=[ |
| ttnn.CBFormatDescriptor(buffer_index=0, data_format=ttnn.bfloat16, page_size=2048)]), |
| ttnn.CBDescriptor(total_size=256, core_ranges=cores, format_descriptors=[ |
| ttnn.CBFormatDescriptor(buffer_index=1, data_format=ttnn.uint32, page_size=256)])] |
| sems = [ttnn.SemaphoreDescriptor(id=i, core_ranges=cores, initial_value=0) for i in range(2 + NSH)] |
| return ttnn.ProgramDescriptor(kernels=[kd], semaphores=sems, cbs=cbs) |
|
|
| prog = build() |
| ttnn.generic_op([out, out], prog) |
| ttnn.synchronize_device(dev) |
| if verify: |
| r = ttnn.to_torch(out).reshape(C * R, 8) |
| bad = int(r[:, 0].sum()); tags = int(((r[:, 1].to(torch.int64) & 0xFFFF0000) == 0xC0DE0000).sum()) |
| print(f"[bc] {mode} verify: cores reporting {tags}/{C*R}, mismatched words {bad}", flush=True) |
| if args.get("prof") == "1": |
| pe = build(empty=True) |
| prof_run(dev, "empty", lambda: (ttnn.generic_op([out, out], pe), None)[1]) |
| prof_run(dev, "bcast", lambda: (ttnn.generic_op([out, out], prog), None)[1]) |
| if args.get("time", "1") == "1": |
| reps = int(args.get("reps", 20)) |
| pe = build(empty=True) |
| te = time_op(dev, lambda: (ttnn.generic_op([out, out], pe), None)[1], reps=reps) |
| tb = time_op(dev, lambda: (ttnn.generic_op([out, out], prog), None)[1], reps=reps) |
| tot = NSH * SB |
| print(f"[bc] {mode} noc{noc} {C}x{R} NSH={NSH} SB={SB}: empty {te:.2f} us, bcast {tb:.2f} us -> {tb - te:.2f} us" |
| f" ({tot / max(tb - te, 1e-3) / 1e3:.1f} GB/s per receiver)", flush=True) |
| finally: |
| from open_device_12x10 import close_device |
| close_device(dev) |
| if args.get("prof") == "1": |
| prof_report() |
|
|
