File size: 5,731 Bytes
58ba6df | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """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":
# hop A: row y, head (0, y) holds shard y; hop B: column x, chain y = 0..R-1
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:
# hop A: column x (x < NSH), head (x, 0) holds shard x; hop B: row y, chain x = 0..NSH-1
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()
|