code/scripts/micro/dram_bench.py · changh95/hamer-p150 at main

hamer-p150 / code /scripts /micro /dram_bench.py

Optimized build (2026-10-03): model call 3.49 ms, device 3.36 ms

58ba6df verified about 1 month ago

7.67 kB

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

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