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