File size: 5,350 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 | """hgemv (kernels/hgemv, tt_kernels.hgemv) bring-up: correctness vs torch fp32 and timing.
usage: hgemv_test.py [K=1024] [N=1024] [ln=1] [gelu=0] [res=1] [xf32=1] [time=1]"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from common import *
from hamer import tt_kernels as tk
a = dict(s.split("=") for s in sys.argv[1:])
K, N = int(a.get("K", 1024)), int(a.get("N", 1024))
ln, gelu, res, xf32 = a.get("ln", "1") == "1", a.get("gelu", "0") == "1", a.get("res", "0") == "1", a.get("xf32", "1") == "1"
dev = open_dev()
try:
torch.manual_seed(0)
xv = torch.randn(K) * 2 + 0.5
W = torch.randn(K, N) / K ** 0.5
bv = torch.randn(N) * 0.1
rv = torch.randn(N)
xdt = ttnn.float32 if xf32 else ttnn.bfloat16
rowf = lambda v: v.reshape(1, 1, 1, -1).expand(1, 1, 32, v.numel()).contiguous()
x = ttnn.from_torch(rowf(xv), dtype=xdt, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.L1_MEMORY_CONFIG)
D = int(a.get("din", 0))
if D: # dense replica input: D tiles, each = x reshaped 32x32 row-major
x = ttnn.from_torch(tk.hg_dense_layout(xv).repeat(D, 1).reshape(1, 1, 32 * D, 32), dtype=xdt, layout=ttnn.TILE_LAYOUT,
device=dev, memory_config=ttnn.L1_MEMORY_CONFIG)
b = ttnn.from_torch(rowf(bv).to(torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev,
memory_config=ttnn.L1_MEMORY_CONFIG) # as in the model (_hg_rowf)
r = ttnn.from_torch(rowf(rv), dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.L1_MEMORY_CONFIG) if res else None
w = tk.CmmWeight(W, dev, N // 32, 1, K // 32, ttnn.bfloat16)
xin = xv.to(torch.float32 if xf32 else torch.bfloat16).double()
h = torch.nn.functional.layer_norm(xin, (K,), eps=1e-5) if ln else xin
ref = h @ W.to(torch.bfloat16).double() + bv.to(torch.bfloat16).double()
if gelu:
ref = torch.nn.functional.gelu(ref)
if res:
ref = ref + rv.double()
DO = int(a.get("dout", 0))
o = tk.hgemv(x, w, b, ln_eps=1e-5 if ln else None, gelu=gelu, res=r, out_dtype=ttnn.float32, dense_out=DO)
ttnn.synchronize_device(dev)
if DO:
o, od = o
dd = ttnn.to_torch(od).double().reshape(DO, 32, 32)
got0 = ttnn.to_torch(o).double()[0, 0, 0]
print(f"[hg] dense out: replicas equal row-form: {all(torch.equal(tk.hg_dense_vector(dd[c]), got0) for c in range(DO))}", flush=True)
got = ttnn.to_torch(o).double()[0, 0, 0]
print(f"[hg] K{K} N{N} ln{int(ln)} gelu{int(gelu)} res{int(res)} xf32{int(xf32)}: pcc {pcc(got, ref):.7f} "
f"max|d| {float((got - ref).abs().max()):.5f} rel {float((got - ref).abs().max() / ref.abs().max()):.2e}", flush=True)
if a.get("pf", "0") == "1": # round 10: cross-op weight prefetch -- GEMV 1 on rect A prefetches GEMV 2's weights
rA = ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(7, 3))})
rB = ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 4), ttnn.CoreCoord(7, 7))})
bufB = tk.hg_weight_buffer(dev, rB)
kw = dict(ln_eps=1e-5 if ln else None, gelu=gelu, res=r, out_dtype=ttnn.float32, dense_out=DO)
def pair(pf):
o1 = tk.hgemv(x, w, b, cores=rA, prefetch=(w, bufB, rB) if pf else None, **kw)
o2 = tk.hgemv(x, w, b, cores=rB, w_l1=bufB if pf else None, **kw)
return o1, o2
p1, p2 = pair(True)
ttnn.synchronize_device(dev)
g = lambda o: ttnn.to_torch(o[0] if DO else o).double()[0, 0, 0]
print(f"[hg] prefetch pair: out1 == out2 {bool(torch.equal(g(p1), g(p2)))} == plain {bool(torch.equal(g(p2), got))}", flush=True)
if a.get("time", "1") == "1":
bufA = tk.hg_weight_buffer(dev, rA)
def chain(pf, n=8): # steady state: every program prefetches the next one's weights
rs, bs, outs = (rA, rB), (bufA, bufB), []
for i in range(n):
nxt = (w, bs[(i + 1) % 2], rs[(i + 1) % 2]) if pf and i + 1 < n else None
outs.append(tk.hgemv(x, w, b, cores=rs[i % 2], w_l1=bs[i % 2] if pf and i > 0 else None,
prefetch=nxt, **kw))
return outs
print(f"[hg] chain of 8: plain {time_op(dev, lambda: chain(False), reps=4) / 8:.2f} us/gemv prefetched "
f"{time_op(dev, lambda: chain(True), reps=4) / 8:.2f} us/gemv", flush=True)
print(f"[hg] pair plain {time_op(dev, lambda: pair(False)):.2f} us prefetched {time_op(dev, lambda: pair(True)):.2f} us", flush=True)
if a.get("time", "1") == "1":
print(f"[hg] {time_op(dev, lambda: tk.hgemv(x, w, b, ln_eps=1e-5 if ln else None, gelu=gelu, res=r, out_dtype=ttnn.float32, dense_out=DO)):.2f} us/call", flush=True)
if a.get("prof", "0") == "1": # round 10: zone timeline (HG_PROF=1 TT_METAL_DEVICE_PROFILER=1)
for _ in range(4):
o2 = tk.hgemv(x, w, b, ln_eps=1e-5 if ln else None, gelu=gelu, res=r, out_dtype=ttnn.float32, dense_out=DO)
ttnn.synchronize_device(dev)
ttnn.ReadDeviceProfiler(dev)
finally:
from open_device_12x10 import close_device
close_device(dev)
if a.get("prof", "0") == "1":
tl_report(risks=("NCRISC", "BRISC", "TRISC_0", "TRISC_1", "TRISC_2"))
|