File size: 2,449 Bytes
cffc837 | 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 | """Check a hand-off bundle with nothing but numpy: python verify.py [dir]"""
import json
import sys
import numpy as np
root = sys.argv[1] if len(sys.argv) > 1 else "."
manifest = json.load(open(f"{root}/manifest.json"))
npz = np.load(f"{root}/weights.npz")
print(f"{manifest['model']} pattern={manifest['pattern']} "
f"axis={manifest['group_axis']} N_inference={manifest['N_inference']}")
print(f"{'matrix':<20}{'M':>6}{'K':>6}{'sparsity':>10}{'pattern':>10}{'ref':>10}")
print("-" * 62)
bad = 0
for m in manifest["matrices"]:
name, M, K = m["name"], m["M"], m["K"]
w = npz[f"{name}.weight"]
mask = npz[f"{name}.mask"]
assert w.shape == (M, K), f"{name}: shape {w.shape} != {(M, K)}"
assert np.all(w[mask == 0] == 0), f"{name}: nonzero weight under a zero mask"
assert int((mask != 0).sum()) == m["nonzero"], f"{name}: nonzero count"
# Structural check: for an N:M pattern every full group of M columns along
# K must hold at most N nonzeros. This is the contract a generated kernel
# relies on, so verify it against the shipped array, not the manifest.
ok = "n/a"
pat = m["pattern"]
if pat.startswith("1x"):
# Block pattern: every full block of B columns along K must be kept or
# dropped whole. Partial blocks would break a block-packed kernel.
block = int(pat[2:].split(":")[0])
nb = K // block
if nb:
counts = (w[:, :nb * block] != 0).reshape(M, nb, block).sum(-1)
viol = int(((counts != 0) & (counts != block)).sum())
bad += viol
ok = "OK" if viol == 0 else f"{viol} BAD"
elif ":" in pat and not pat.startswith("unstructured"):
n, g = (int(v) for v in pat.split(":"))
ng = K // g
if ng:
counts = (w[:, :ng * g] != 0).reshape(M, ng, g).sum(-1)
viol = int((counts > n).sum())
bad += viol
ok = "OK" if viol == 0 else f"{viol} BAD"
ref = "-"
if f"{name}.ref_x" in npz.files:
y = w @ npz[f"{name}.ref_x"]
if f"{name}.bias" in npz.files:
y = y + npz[f"{name}.bias"]
err = np.max(np.abs(y - npz[f"{name}.ref_y"]))
ref = f"{err:.2e}"
if err > 1e-3:
bad += 1
ref += " BAD"
print(f"{name:<20}{M:>6}{K:>6}{m['sparsity']:>10.4f}{ok:>10}{ref:>10}")
print("-" * 62)
print("FAILED" if bad else "all checks passed")
sys.exit(1 if bad else 0)
|