Clear windows branch; preserve previous contents in win-old
This commit is contained in:
ayo committed
1 parent
f0c881985b
commit
73994d4634
311 files changed
-57018
No files matched your search
@@ -1,113 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Compares native preset-M Swin-block outputs with the numpy model (swin_model.py) on sampled blocks of a dump.
|
||||
|
||||
usage: swin_check.py DUMP_DIR OUT_DIR [OUT_DIR...] [--blocks N]
|
||||
DUMP_DIR replay dump of one rrlite_{enc1,enc2,enc3_tube,dec1,dec2}_4x4 launch
|
||||
OUT_DIR dump_runner output of a native kernel on that dump
|
||||
The stored e4m3 outputs (block output, and the patch merge of encoders that have one) are compared code by code:
|
||||
the share of identical codes, of codes one e4m3 step apart, and the PSNR of the decoded values. The model is
|
||||
NVIDIA's arithmetic (e4m3 activations, f16 accumulation per k32 step); the production kernels keep activations
|
||||
f16 and accumulate in f32 (SWIN_NO_Q8, SWIN_F32ACC), so they are not expected to match it exactly.
|
||||
"""
|
||||
import os
|
||||
import random
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import model_enc3 as me # noqa: E402
|
||||
import swin_model as sm # noqa: E402
|
||||
|
||||
# kernel -> (C, heads, NPM, CIN)
|
||||
LAYERS = {'enc1': (64, 2, 96, 0), 'enc2': (96, 4, 128, 0), 'enc3_tube': (128, 4, 0, 0), 'dec2': (96, 4, 0, 128),
|
||||
'dec1': (64, 2, 0, 96)}
|
||||
|
||||
# Preserve the existing Linux oracle by default; Windows strict replay enables
|
||||
# the audited 310.9 f16 patch-merge bound before final e4m3 quantisation.
|
||||
CLAMP_PATCH_MERGE = False
|
||||
|
||||
|
||||
def blocks(grid, n, seed=1):
|
||||
gx, gy = grid[0], grid[1]
|
||||
b = {(0, 0), (gx - 1, 0), (0, gy - 1), (gx - 1, gy - 1), (gx // 2, 0), (0, gy // 2), (gx - 1, gy // 2), (gx // 2, gy - 1)}
|
||||
r = random.Random(seed)
|
||||
while len(b) < min(n, gx * gy):
|
||||
b.add((r.randrange(gx), r.randrange(gy)))
|
||||
return sorted(b)
|
||||
|
||||
|
||||
def model_block(P, layer, bx, by):
|
||||
C, heads, NPM, CIN = LAYERS[layer]
|
||||
pre = 4 * C * CIN + 8 * C if CIN else 0
|
||||
x0 = sm.decoder_input(P, bx, by, C, CIN, heads, P.p32) if CIN else sm.load_tokens(P, bx, by, C // 32)
|
||||
x2, out, L = sm.swin_block(P, x0, C, heads, pre)
|
||||
merged = None
|
||||
if NPM:
|
||||
A = sm.patch_merge_input(out, C)
|
||||
PM0 = L['end']
|
||||
cols = []
|
||||
for g in range(NPM // 32):
|
||||
Wg = me.E4M3[P.wbuf[sm.woff_table(PM0 + g * 128 * C, 512, 64 * C, 4 * C, 32)]]
|
||||
bias = P.f16vec(PM0 + 4 * C * NPM + 64 * g, 32)
|
||||
cols.append(me.mma_chain(A, Wg, C0=np.broadcast_to(bias, (16, 32)).copy()))
|
||||
merged = sm.pair_to_nat(np.concatenate(cols, axis=1))
|
||||
if CLAMP_PATCH_MERGE:
|
||||
limit = me.f16(2 * np.pi)
|
||||
merged = np.clip(merged, -limit, limit)
|
||||
merged = me.q8(merged)
|
||||
return out, merged
|
||||
|
||||
|
||||
def codes_at(buf, base, ptr, W, H, X, Y, C):
|
||||
"""e4m3 codes of tokens (X, Y) (in bounds) from planes of 32 channels at ptr"""
|
||||
c = np.arange(C)
|
||||
off = ptr - base + (((c[None, :] >> 5) * H + Y[:, None]) * W + X[:, None]) * 32 + (c[None, :] & 31)
|
||||
return buf[off]
|
||||
|
||||
|
||||
def main():
|
||||
args = [a for a in sys.argv[1:] if not a.startswith('--')]
|
||||
n = 8
|
||||
if '--blocks' in sys.argv:
|
||||
n = int(sys.argv[sys.argv.index('--blocks') + 1])
|
||||
args.remove(str(n))
|
||||
dump, outs = args[0], args[1:]
|
||||
P = me.Params(dump)
|
||||
layer = P.kernel.replace('rrlite_', '').replace('_4x4', '')
|
||||
C, heads, NPM, CIN = LAYERS[layer]
|
||||
bl = blocks(P.grid, n)
|
||||
models = {b: model_block(P, layer, *b) for b in bl}
|
||||
for out_dir in outs:
|
||||
got, want = [], []
|
||||
for (bx, by) in bl:
|
||||
base, _ = P.alloc_of(P.out)
|
||||
aid = next(i for i, (b0, buf) in P.allocs.items() if b0 == base)
|
||||
buf = np.fromfile(f'{out_dir}/alloc-{aid}.bin', np.uint8)
|
||||
X, Y = P.token_xy(bx, by)
|
||||
ok = (X >= 0) & (X < P.W) & (Y >= 0) & (Y < P.H)
|
||||
out, merged = models[(bx, by)]
|
||||
got.append(codes_at(buf, base, P.out, P.W, P.H, X[ok], Y[ok], C).ravel())
|
||||
want.append(out[ok].ravel())
|
||||
if NPM:
|
||||
W2, H2 = P.W // 2, P.H // 2
|
||||
r = np.arange(16)
|
||||
MX, MY = (8 * bx - P.sx) // 2 + (r & 3), (8 * by - P.sy) // 2 + (r >> 2)
|
||||
okm = (MX >= 0) & (MX < W2) & (MY >= 0) & (MY < H2)
|
||||
mbase, _ = P.alloc_of(P.p48)
|
||||
mid = next(i for i, (b0, _) in P.allocs.items() if b0 == mbase)
|
||||
mbuf = buf if mid == aid else np.fromfile(f'{out_dir}/alloc-{mid}.bin', np.uint8)
|
||||
got.append(codes_at(mbuf, mbase, P.p48, W2, H2, MX[okm], MY[okm], NPM).ravel())
|
||||
want.append(merged[okm].ravel())
|
||||
g = me.E4M3[np.concatenate(got)]
|
||||
w = np.concatenate(want).astype(np.float64)
|
||||
rank = lambda v: np.searchsorted(me._POS, np.abs(v)) * np.sign(v) # noqa: E731
|
||||
steps = np.abs(rank(g) - rank(w))
|
||||
psnr = 10 * np.log10(np.nanmax(np.abs(w)) ** 2 / max(np.nanmean((g - w) ** 2), 1e-30))
|
||||
print(f'{layer:9s} {os.path.basename(os.path.normpath(out_dir)):10s} {len(bl)} blocks {g.size} codes: '
|
||||
f'identical {100 * (g == w).mean():5.1f}% <=1 step {100 * (steps <= 1).mean():6.2f}% '
|
||||
f'max {int(steps.max())} steps psnr {psnr:.1f} dB')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in new issue
Block a user