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