File size: 3,594 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 | """Round 12: raw DRAM read ceilings for the xmm fc1 weight-sender placements (kernels/mcbench/dram_rd.cpp, all reads issued
at once, one barrier). usage: dram_fc1.py"""
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/dram_rd.cpp")
dev = open_device(grid="12x10", dispatch="eth", trace_region_size=64 << 20)
try:
nb = dev.dram_grid_size().x
SH = 4 << 20
cg = ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(nb - 1, 0))})
mem = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.BufferType.DRAM,
ttnn.ShardSpec(cg, [SH // 64, 32], ttnn.ShardOrientation.ROW_MAJOR))
src = ttnn.from_torch(torch.zeros(nb * SH // 64, 32), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=mem)
base_addr = src.buffer_address()
cores = ttnn.CoreRangeSet([ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(11, 9))])
LS = 256 * 1024
def run(rd0, rd1):
kds = []
tot = 0
for noc, readers in ((0, rd0), (1, rd1)):
args = ttnn.RuntimeArgs()
for x in range(12):
for y in range(10):
args[x][y] = [0] * 8
for (c, b0, chunk, reps) in readers:
args[c[0]][c[1]] = [1, b0, nb, base_addr, chunk, reps, chunk, LS]
tot += chunk * reps
cfg = ttnn.ReaderConfigDescriptor() if noc == 0 else ttnn.WriterConfigDescriptor()
kds.append(ttnn.KernelDescriptor(kernel_source=K, source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
core_ranges=cores, compile_time_args=[], defines=[], runtime_args=args, config=cfg))
cb = ttnn.CBDescriptor(total_size=2 * LS, core_ranges=cores, format_descriptors=[
ttnn.CBFormatDescriptor(buffer_index=0, data_format=ttnn.bfloat16, page_size=2048)])
prog = ttnn.ProgramDescriptor(kernels=kds, semaphores=[], cbs=[cb])
f = lambda: (ttnn.generic_op([src, src], prog), None)[1]
return time_op(dev, f, reps=10, iters=5), tot
base, _ = run([], [])
print(f"[dram] empty {base:.2f} us", flush=True)
def rep(name, r0, r1=()):
t, tot = run(list(r0), list(r1))
print(f"[dram] {name}: {t:.2f} us {tot / 1e3 / max(t - base, 1e-3):.0f} GB/s (tot {tot >> 10} KB)", flush=True)
B = 17408
top = [(x, 3 * s) for s in range(3) for x in range(12)] + [(3 * g, 9) for g in range(4)]
bot = [(x, 3 * s + 2) for s in range(3) for x in range(12)] + [(3 * g + 2, 9) for g in range(4)]
rep("40 senders NOC0 10 x 17KB", [(c, i % nb, B, 10) for i, c in enumerate(top)])
rep("40 top NOC0 5 + 40 bottom NOC1 5", [(c, i % nb, B, 5) for i, c in enumerate(top)], [(c, (i + 3) % nb, B, 5) for i, c in enumerate(bot)])
rep("80 (top+bottom) NOC0 5 each", [(c, i % nb, B, 5) for i, c in enumerate(top + bot)])
g120 = [(x, y) for y in range(10) for x in range(12)]
rep("120 NOC0 1/3 each", [(c, i % nb, B, 10) for i, c in enumerate(g120[:120])][:120] and [(c, i % nb, 5803, 10) for i, c in enumerate(g120)])
rep("60 NOC0 + 60 NOC1", [(c, i % nb, 5803 * 2, 5) for i, c in enumerate(g120[:60])], [(c, i % nb, 5803 * 2, 5) for i, c in enumerate(g120[60:])])
rep("120: 60 NOC0 even rows-ish + 60 NOC1", [(c, i % nb, 5803, 10) for i, c in enumerate(g120) if c[1] % 2 == 0] + [],
[(c, i % nb, 5803 * 1, 10) for i, c in enumerate(g120) if c[1] % 2 == 1])
finally:
ttnn.close_device(dev)
|