File size: 2,761 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 | """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,):
for size in (34816, 69632):
reps = MB // size
for each in (False, True):
for lb in (0, 1):
for sy in (0, 2, 5):
t = run([((x, sy), (x, 0, x, 5)) for x in range(10)], size, reps, noc, lb=lb, each=each)
print(f"[mc] noc{noc} 10 cols x6 sender y={sy} lb={lb} barrier_each={each} {size}B x{reps}: {t:.2f} us -> {(t-base)/reps:.2f} us/xfer {reps*size / (t - base) / 1e3:.1f} GB/s", flush=True)
finally:
ttnn.close_device(dev)
|