File size: 6,934 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 | """Cooperative 2D matmul (tt_kernels.cmm_linear) bring-up / accuracy / timing / soak.
usage: cmm_test.py [eth|worker] shape=tiny|small|fc2|proj [uni] [kb=N] [slots=N] [time] [lib] [soak=N] [prof risks=...]
One call per launch unless time/soak is given."""
import sys
from pathlib import Path
from common import *
from open_device_12x10 import open_device
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from hamer import tt_kernels as tk
arg = lambda k, d: next((a.split("=", 1)[1] for a in sys.argv if a.startswith(k + "=")), d)
disp = sys.argv[1] if len(sys.argv) > 1 and sys.argv[1] in ("eth", "worker") else "eth"
shape = arg("shape", "tiny")
kb, slots = int(arg("kb", "0")), int(arg("slots", "0"))
uni = "uni" in sys.argv
GX, GY, NT, KT = {"tiny": (2, 2, 1, 8), "small": (4, 3, 2, 16), "fc2": (10, 6, 4, 160), "proj": (10, 6, 4, 48),
"fc2h": (10, 6, 4, 80), "projh": (10, 6, 4, 24),
"k2t": (2, 6, 1, 8), "k2s": (4, 6, 2, 32), # round 11: K-split (layout=k2) bring-up shapes
"k2w": (2, 6, 8, 8), "k2w2": (3, 6, 12, 16), "qkv": (12, 6, 12, 40)}[shape] # wide (NT > 4)
kb = kb or {"tiny": 2, "small": 4, "fc2": 8, "proj": 8, "k2t": 2, "k2s": 4, "k2w": 2, "k2w2": 4, "qkv": 5}[shape]
slots = slots or min(3, GX, GY)
dev = open_device(grid="12x10", dispatch=disp, trace_region_size=64 << 20)
try:
torch.manual_seed(0)
M, K, N = 32 * GY, 32 * KT, 32 * NT * GX
at = (torch.randn(1, 1, M, K) * 0.5).to(torch.bfloat16)
wt = (torch.randn(K, N) * K ** -0.5).to(torch.bfloat16)
bt = (torch.randn(N) * 0.1).to(torch.bfloat16)
adt = ttnn.bfloat8_b if "abfp8" in sys.argv else ttnn.bfloat16
a = ttnn.from_torch(at, dtype=adt, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.L1_MEMORY_CONFIG)
if "abfp8" in sys.argv:
at = ttnn.to_torch(a).to(torch.bfloat16) # the quantized activation is the reference input
gfrom = None
if "gh" in sys.argv: # round 8: GELU on the helper cores, the activation buffer is written by them
gfrom = a
at = torch.nn.functional.gelu(at.float()).to(torch.bfloat16) # reference input of the matmul
if "alocal" in sys.argv or "gh" in sys.argv: # activation in the cmm owner layout (block-sharded, one K block per core)
amem = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.BLOCK_SHARDED, ttnn.BufferType.L1, ttnn.ShardSpec(
ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(GX - 1, GY - 1))}), [32, 32 * kb * (2 if arg("layout", "out") == "k2" else 1)],
ttnn.ShardOrientation.ROW_MAJOR))
a = ttnn.to_memory_config(a, amem) if gfrom is None else ttnn.allocate_tensor_on_device(
ttnn.Shape([1, 1, M, K]), ttnn.bfloat16, ttnn.TILE_LAYOUT, dev, amem)
w = ttnn.from_torch(wt, dtype=ttnn.bfloat8_b, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG)
b = ttnn.from_torch(bt.reshape(1, 1, 1, N), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev)
omem = ttnn.create_sharded_memory_config((M, N), core_grid=ttnn.CoreGrid(y=GY, x=GX), strategy=ttnn.ShardStrategy.BLOCK,
orientation=ttnn.ShardOrientation.ROW_MAJOR)
wq = ttnn.to_torch(w).float()
aref = torch.nn.functional.gelu(at.float()).to(torch.bfloat16).float() if "gelu" in sys.argv else at.float()
ref = aref.reshape(M, K) @ wq + bt.float()
defs = tuple(d for d in arg("defs", "").split(",") if d)
_g = arg("layout", "out") == "g120"
bfn = None
if _g and "near" in sys.argv:
_nb = {(h, y): tk.near_banks(dev, 6 * h, y) for h in range(2) for y in range(10)}
bfn = lambda xw, b: _nb[(xw // 10, xw % 10)][b % 2]
wf = tk.CmmWeight(wt, dev, 20 if _g else GX, 2 if _g else NT, kb, ttnn.bfloat8_b, bank_fn=bfn) if "packed" in sys.argv else w
f = lambda: tk.cmm_linear(a, wf, b, omem, kb=kb, slots=slots, unicast=uni, defines=defs, layout=arg("layout", "out"), rotate="rot" in sys.argv, gelu="gelu" in sys.argv, v1="v1" in sys.argv, fidelity=getattr(ttnn.MathFidelity, arg("fid", "HiFi2")), gelu_from=gfrom,
out_dtype=ttnn.bfloat8_b if "obfp8" in sys.argv else None)
print(f"[cmm] {disp} shape={shape} G={GX}x{GY} NT={NT} KT={KT} kb={kb} slots={slots} uni={uni}", flush=True)
o = f()
ttnn.synchronize_device(dev)
ot = ttnn.to_torch(o).float().reshape(M, N)
print(f"[cmm] first call OK PCC {pcc(ot, ref):.7f} max|d| {float((ot - ref).abs().max()):.4f} mean|d| {float((ot - ref).abs().mean()):.5f}", flush=True)
if "lib" in sys.argv:
from hamer.vit import tuned_matmul_configs
pc = ttnn.MatmulMultiCoreReuseMultiCastProgramConfig(
compute_with_storage_grid_size=ttnn.CoreCoord(GX, GY), in0_block_w=kb, out_subblock_h=1, out_subblock_w=NT,
per_core_M=1, per_core_N=NT, transpose_mcast=False, fused_activation=None)
ck = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=ttnn.MathFidelity.HiFi2, fp32_dest_acc_en=True, packer_l1_acc=True)
g = lambda: ttnn.linear(a, w, bias=b, program_config=pc, memory_config=omem, compute_kernel_config=ck)
lt = ttnn.to_torch(g()).float().reshape(M, N)
print(f"[lib] PCC {pcc(lt, ref):.7f} mean|d| {float((lt - ref).abs().mean()):.5f} cmm==lib {bool(torch.equal(lt, ot))} max|cmm-lib| {float((lt-ot).abs().max()):.4f}", flush=True)
if "time" in sys.argv:
print(f"[lib] {time_op(dev, g):.2f} us", flush=True)
if "time" in sys.argv:
print(f"[cmm] {time_op(dev, f):.2f} us", flush=True)
if "prof" in sys.argv: # round 10: zone timeline (TT_METAL_DEVICE_PROFILER=1, defs=...,CMM_PROF)
for _ in range(4):
o = f(); ttnn.synchronize_device(dev); ttnn.deallocate(o)
ttnn.ReadDeviceProfiler(dev)
n = int(arg("soak", "0"))
if n:
reps = 50
o0 = ttnn.to_torch(f()).float()
tid = ttnn.begin_trace_capture(dev, cq_id=0)
outs = []
try:
for _ in range(reps):
outs.append(f())
finally:
ttnn.end_trace_capture(dev, tid, cq_id=0)
try:
for i in range(n // reps):
ttnn.execute_trace(dev, tid, cq_id=0, blocking=False)
if i % 40 == 39:
ttnn.synchronize_device(dev)
ttnn.synchronize_device(dev)
bad = sum(int(not torch.equal(ttnn.to_torch(x).float(), o0)) for x in outs)
print(f"[cmm] soak {n // reps * reps} replays OK, outputs differing from first: {bad}/{reps}", flush=True)
finally:
ttnn.release_trace(dev, tid)
finally:
ttnn.close_device(dev)
if "prof" in sys.argv:
tl_report(risks=tuple(arg("risks", "NCRISC,BRISC,TRISC_1").split(",")))
tl_seq(("WISS", "WDRAM", "WFWD", "XFER"), "NCRISC")
tl_seq(("RECV", "XFER"), "BRISC")
tl_seq(("CWA", "CWW"), "TRISC_0")
|