File size: 7,672 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 102 103 104 105 106 107 108 109 | """raw DRAM read bandwidth patterns on 12x10 (kernels/mcbench/dram_rd.cpp): usage dram_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/dram_rd.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:
nb = dev.dram_grid_size().x
print("[dram] banks", nb, flush=True)
SH = 4 << 20 # 4 MB per bank
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(readers, noc, defs=()):
# readers: list of ((x, y), bank0, nrot, chunk, reps, stride)
args = ttnn.RuntimeArgs()
for x in range(12):
for y in range(10):
args[x][y] = [0] * 8
tot = 0
for rr in readers:
(c, b0, nrot, chunk, reps, stride), ad = rr[:6], (rr[6] if len(rr) > 6 else base_addr)
args[c[0]][c[1]] = [1, b0, nrot, ad, chunk, reps, stride, LS]
tot += chunk * reps
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=list(defs), runtime_args=args, config=cfg)
cb = ttnn.CBDescriptor(total_size=LS, 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([src, src], prog), None)[1]
return time_op(dev, f, reps=10, iters=5), tot
base, _ = run([], 0)
print(f"[dram] empty {base:.2f} us", flush=True)
TOT = 6960 * 1024 # fc2 weight bytes (bfp8)
def rep(name, readers, noc, defs=()):
t, tot = run(readers, noc, defs)
print(f"[dram] noc{noc} {name}: {t:.2f} us {tot / 1e3 / max(t - base, 1e-3):.0f} GB/s (tot {tot >> 10} KB)", flush=True)
grid60 = [(x, y) for y in range(6) for x in range(10)]
grid120 = [(x, y) for y in range(10) for x in range(12)]
if "l1" in sys.argv: # round 10: L1 -> L1 pulls (the cmm forwarder pattern): reader (x, ry) <- core (x, sy)
# source L1 address: the scratch CB (same address on every core); data content irrelevant
L1A = 0x100000
def pc(x, y):
p = dev.worker_core_from_logical_core(ttnn.CoreCoord(x, y))
return (p.x << 8) | p.y
for ry, sy in ((0, 6), (0, 1), (0, 5), (0, 9), (3, 6)):
rep(f"10 pulls (x,{sy}) -> (x,{ry}), 10 x 69632B", [((x, ry), pc(x, sy), 0, 69632, 10, 0, L1A) for x in range(10)], 0, (("L1_SRC", "1"),))
rep("1 pull (0,6) -> (0,0), 10 x 69632B", [((0, 0), pc(0, 6), 0, 69632, 10, 0, L1A)], 0, (("L1_SRC", "1"),))
sys.exit(0)
if "cmm" in sys.argv: # round 10: the fc2 cmm weight feed (10 columns x 10 blocks of 69632 B, bank (x + b) % 8)
BLK = 69632
def cmm_readers(cells, nper):
# cells: reader cores; reader i reads nper consecutive blocks rotating banks from bank0
return [(c, i % nb, nb, BLK, nper, BLK) for i, c in enumerate(cells)]
for y in range(6):
rep(f"10 readers row {y}, 10 x 69632B rotating", cmm_readers([(x, y) for x in range(10)], 10), 0)
rep("10 readers diag x%6, 10 x 69632B", cmm_readers([(x, x % 6) for x in range(10)], 10), 0)
rep("10 readers rows 6-9 / cols 10-11 (idle cores), 10 x 69632B", cmm_readers([(x, 6 + x % 4) for x in range(10)], 10), 0)
rep("20 readers rows 0-1, 5 x 69632B", cmm_readers([(x, y) for y in range(2) for x in range(10)], 5), 0)
rep("20 readers rows 0,3, 5 x 69632B", cmm_readers([(x, y) for y in (0, 3) for x in range(10)], 5), 0)
rep("60 readers 10x6, 2 x 69632B (approx 1/6)", cmm_readers([(x, y) for y in range(6) for x in range(10)], 2), 0)
rep("120 readers, 1 x 69632B", cmm_readers([(x, y) for y in range(10) for x in range(12)], 1), 0)
rep("8 readers optimal per bank (row 0 x 0..7), 12 x 69632B", cmm_readers([(x, 0) for x in range(8)], 12), 0)
# round 11: split-K layouts (2 / 3 sender rows per column)
rep("20 readers rows 0,5, 5 x 69632B", cmm_readers([(x, y) for y in (0, 5) for x in range(10)], 5), 0)
rep("30 readers rows 0,3,6, 3 x 69632B (90 blocks)", cmm_readers([(x, y) for y in (0, 3, 6) for x in range(10)], 3), 0)
rep("30 readers rows 0,3,6, 4 x 69632B (120 blocks)", cmm_readers([(x, y) for y in (0, 3, 6) for x in range(10)], 4), 0)
rep("10 readers row 0, 10 x 69632B (today)", cmm_readers([(x, 0) for x in range(10)], 10), 0)
sys.exit(0)
if "row" in sys.argv:
for noc in (0,):
for rows in (1, 2, 3):
rd = [(x, y) for y in range(rows) for x in range(10)]
rep(f"{len(rd)} readers rows 0..{rows-1}, 2 x 69632B single bank", [(c, i % nb, 0, 69632, 2, 69632) for i, c in enumerate(rd)], noc)
rep(f"{len(rd)} readers rows 0..{rows-1}, 16 x 8704B rotating banks", [(c, i % nb, nb, 8704, 16, 8704) for i, c in enumerate(rd)], noc)
rep(f"{len(rd)} readers rows 0..{rows-1}, 128 x 1088B rotating banks", [(c, i % nb, nb, 1088, 128, 1088) for i, c in enumerate(rd)], noc)
rd = [(x, y) for x in range(10) for y in range(6) if y == x % 6]
rep(f"{len(rd)} readers diagonal, 2 x 69632B single bank", [(c, i % nb, 0, 69632, 2, 69632) for i, c in enumerate(rd)], noc)
rep(f"{len(rd)} readers diagonal, 16 x 8704B rotating", [(c, i % nb, nb, 8704, 16, 8704) for i, c in enumerate(rd)], noc)
sys.exit(0) if False else None
for noc in ((0, 1) if "row" not in sys.argv else ()):
# (a) per-tile interleaved-like: each of 60 cores reads 1088 B pieces rotating over banks
n = TOT // 60 // 1088
rep("60 cores, 1088B rotating banks", [(c, i % nb, nb, 1088, n, 1088) for i, c in enumerate(grid60)], noc)
n = TOT // 120 // 1088
rep("120 cores, 1088B rotating banks", [(c, i % nb, nb, 1088, n, 1088) for i, c in enumerate(grid120)], noc)
# (b) big contiguous reads from ONE bank per reader
for chunk in (8704, 69632):
n = max(1, TOT // 60 // chunk)
rep(f"60 cores, {chunk}B single bank each (i%nb)", [(c, i % nb, 0, chunk, n, chunk) for i, c in enumerate(grid60)], noc)
rep(f"60 cores, {chunk}B rotating banks", [(c, i % nb, nb, chunk, n, chunk) for i, c in enumerate(grid60)], noc)
# (c) 8 readers, one per bank
for chunk in (8704, 69632):
n = TOT // nb // chunk
rep(f"{nb} readers, own bank, {chunk}B", [((i, 0), i, 0, chunk, n, chunk) for i in range(nb)], noc)
# (d) one bank, 1 reader
rep("1 reader, bank 0, 69632B x10", [((0, 0), 0, 0, 69632, 10, 69632)], noc)
for q in range(nb):
rep(f"6 readers bank {q} 69632B x2", [((x, 0), q, 0, 69632, 2, 69632) for x in range(6)], noc)
finally:
ttnn.close_device(dev)
|