"""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)