ayo
7c5d3c9157
Add native K/M/RR inference engines, HIP kernels, replay and validation harnesses, and Windows DLSS setup and runtime support. Include GPU profiling tools, build and packaging scripts, platform integration patches, and implementation documentation.
128 lines
5.3 KiB
Python
128 lines
5.3 KiB
Python
#!/usr/bin/env python3
|
|
"""Compares native preset-K layer outputs with the numpy model (pwin_model.py) on sampled windows of a replay dump.
|
|
|
|
usage: pwin_check.py DUMP_DIR OUT_DIR [OUT_DIR...] [--windows N]
|
|
DUMP_DIR replay dump of one dltss_pwin_* launch (inputs)
|
|
OUT_DIR dump_runner output of a native kernel on that dump (alloc-N.bin after the launch)
|
|
For every sampled window (the four corners, border windows and random interior ones) the model output and the
|
|
kernel's stored values are compared per f16 element: exact matches, and the error in f16 ulps at the model
|
|
value's magnitude (values below 2^-10 count with the ulp of 2^-10, so near-zero sign flips do not dominate).
|
|
The kernels round as NVIDIA's f16 wmma does except in the PWIN_F32ACC layers, and gfx12 WMMA groups its f32
|
|
sums differently from gfx11, so neither build is expected to match the model exactly.
|
|
"""
|
|
import os
|
|
import random
|
|
import sys
|
|
|
|
import numpy as np
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
import pwin_model as pm # noqa: E402
|
|
|
|
F16 = np.float16
|
|
|
|
# layer -> (kind, H, C, COUT/CL, extra)
|
|
LAYERS = {
|
|
'enc0': ('enc0', 2, 32, 64, None), 'enc1': ('enc', 2, 64, 64, True), 'enc2': ('enc', 2, 64, 96, False),
|
|
'enc3': ('enc', 4, 96, 128, False), 'enc4': ('enc', 4, 128, 160, False), 'dec5': ('plain', 8, 160, 16, None),
|
|
'dec4': ('dec', 4, 128, 160, None), 'dec3': ('dec', 4, 96, 128, None), 'dec2': ('dec', 2, 64, 96, None),
|
|
'dec1': ('dec', 2, 64, 64, None), 'dec0': ('dec0', 2, 32, 64, None),
|
|
}
|
|
|
|
|
|
class Outputs:
|
|
"""the dump's pointer map over the allocations a native run saved"""
|
|
def __init__(self, dump, out_dir):
|
|
self.d = dump
|
|
self.allocs = {i: (base, np.fromfile(f'{out_dir}/alloc-{i}.bin', np.uint8)) for i, (base, _) in dump.allocs.items()}
|
|
|
|
def rows(self, off_param, index, width):
|
|
a = self.d.u64(off_param)
|
|
for base, buf in self.allocs.values():
|
|
if base <= a < base + buf.size:
|
|
o = a - base + 2 * index * width
|
|
return buf[o:o + 2 * width].view(F16)
|
|
raise ValueError(hex(a))
|
|
|
|
|
|
def ulps(model, got):
|
|
m = model.astype(np.float64)
|
|
g = got.astype(np.float64)
|
|
scale = np.maximum(np.abs(m), 2.0 ** -10)
|
|
ulp = 2.0 ** (np.floor(np.log2(scale)) - 10)
|
|
return np.abs(m - g) / ulp
|
|
|
|
|
|
def windows(grid, n, seed=1):
|
|
gx, gy = grid[0], grid[1]
|
|
w = {(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(w) < min(n, gx * gy):
|
|
w.add((r.randrange(gx), r.randrange(gy)))
|
|
return sorted(w)
|
|
|
|
|
|
def model_window(P, layer, bx, by):
|
|
kind, H, C, X, extra = LAYERS[layer]
|
|
if kind == 'enc0':
|
|
return pm.pwin_enc0(P, bx, by)
|
|
if kind == 'enc':
|
|
return pm.pwin_encoder(P, bx, by, H, C, X, posattn=extra)
|
|
if kind == 'plain':
|
|
return pm.swin_core(P, pm.Layout(H, C, C), pm.load_window(P, bx, by, C)), None
|
|
if kind == 'dec':
|
|
return pm.pwin_decoder(P, bx, by, H, C, X), None
|
|
return pm.pwin_dec0(P, bx, by)
|
|
|
|
|
|
def stored(P, O, layer, bx, by, y, extra_out):
|
|
"""(model values, kernel values) of everything the kernel stores for window (bx, by)"""
|
|
kind, H, C, X, _ = LAYERS[layer]
|
|
W, Hh = P.i32x2(0)
|
|
sx, sy = P.i32x2(56)
|
|
ms, gs = [], []
|
|
full_off = 32 if kind in ('enc0', 'enc') else 24
|
|
for t in range(64):
|
|
Y0, X0 = 8 * by - sy + (t >> 3), 8 * bx - sx + (t & 7)
|
|
if not (0 <= Y0 < Hh and 0 <= X0 < W):
|
|
continue
|
|
if kind == 'dec0':
|
|
ms.append(extra_out[t]); gs.append(O.rows(24, Y0 * W + X0, 40))
|
|
else:
|
|
ms.append(y[t]); gs.append(O.rows(full_off, Y0 * W + X0, C))
|
|
if kind in ('enc0', 'enc'):
|
|
cout = extra_out.shape[1]
|
|
for r in range(16):
|
|
Ym, Xm = int((8 * by - sy) / 2) + (r >> 2), int((8 * bx - sx) / 2) + (r & 3)
|
|
if 0 <= Ym < Hh // 2 and 0 <= Xm < W // 2:
|
|
ms.append(extra_out[r]); gs.append(O.rows(24, Ym * (W // 2) + Xm, cout))
|
|
return np.concatenate(ms), np.concatenate(gs)
|
|
|
|
|
|
def main():
|
|
args = [a for a in sys.argv[1:] if not a.startswith('--')]
|
|
n = 12
|
|
if '--windows' in sys.argv:
|
|
n = int(sys.argv[sys.argv.index('--windows') + 1])
|
|
args.remove(str(n))
|
|
dump_dir, outs = args[0], args[1:]
|
|
P = pm.Dump(dump_dir)
|
|
layer = P.kernel.replace('dltss_pwin_', '').replace('_layer', '')
|
|
wins = windows(P.grid, n)
|
|
models = {w: model_window(P, layer, *w) for w in wins}
|
|
for out in outs:
|
|
O = Outputs(P, out)
|
|
m, g = zip(*(stored(P, O, layer, bx, by, *models[(bx, by)]) for bx, by in wins))
|
|
m, g = np.concatenate(m), np.concatenate(g)
|
|
u = ulps(m, g)
|
|
exact = (m.view(np.uint16) == g.view(np.uint16)).mean()
|
|
md, gd = m.astype(np.float64), g.astype(np.float64)
|
|
psnr = 10 * np.log10(np.abs(md).max() ** 2 / max(((md - gd) ** 2).mean(), 1e-30))
|
|
print(f'{layer:5s} {os.path.basename(os.path.normpath(out)):10s} {len(wins)} windows {m.size} values: exact {100 * exact:5.1f}% '
|
|
f'<=1ulp {100 * (u <= 1).mean():5.1f}% <=4ulp {100 * (u <= 4).mean():6.2f}% max {u.max():7.1f} ulp '
|
|
f'mean {u.mean():.3f} ulp psnr {psnr:.1f} dB finite {np.isfinite(g.astype(np.float32)).mean() * 100:.0f}%')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|