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.
184 lines
7.9 KiB
Python
184 lines
7.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Exact numpy reference of the denoiser's decoder 1 (cuda_dldn_engine_swin_dec1_kernel), one 8x8 window at a time.
|
|
Derived like engine/rr/enc1_ref.py, phase by phase against ZLUDA replays (engine/rr/phase_dump.py); NVIDIA runs it with
|
|
two warps, each one head and half of the expand's columns.
|
|
|
|
expand the 6 x 6 level-2 tokens under the window and its border (dec2's 96 channels, arena order) -> 256 =
|
|
4 sub-positions x 64 channels (+ bias), FP8; sub-position 2 (y & 1) + (x & 1) (dec0_ref)
|
|
prologue key token = expanded + skip token (enc1's 64-channel output at argument 24), f16; enc1's input norm (two
|
|
planes) times the gain at W+25088; the residual is the token + the table at W+62080
|
|
block enc1's: two 32-channel heads, eight MLP branches 64 -> 32 -> 64 (gain W+62208, table W+66496)
|
|
outputs the tokens (argument 48); the 64 -> 16 head and dec0's filter (bilinear texture 88 from level 2, gate,
|
|
textures 96 / 104) to surface 112
|
|
|
|
usage: dec1_ref.py CAPTURE FRAME WEIGHTS.npz [--blocks X,Y ...] (verifies against the capture)
|
|
dec1_ref.py CAPTURE FRAME OUT.npz --extract DUMP.npz (weights from a phase dump)
|
|
"""
|
|
import argparse
|
|
import os
|
|
import sys
|
|
|
|
import numpy as np
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from rrmath import e4m3, f16, q8 # noqa: E402
|
|
from enc1_ref import inter, mm, softmax_cubic, act, mirror, add, input_norm, nv_rowsum_sq, H, F # noqa: E402
|
|
from dec0_ref import keys, fma16, bilinear88 # noqa: E402
|
|
|
|
I32, I64, I96 = inter(32), inter(64), inter(96)
|
|
UN64 = [I64.index(i) for i in range(64)]
|
|
GAIN0, TABLE7, GAIN14, TABLE16 = 25088, 62080, 62208, 66496
|
|
|
|
|
|
def window(w, lo, skip, x0, y0, ex, ey):
|
|
"""lo: dec2's output [3][ey/2][ex/2][32] codes; skip: enc1's [2][ey][ex][32]. Returns (tokens 64 x 64 codes in arena
|
|
order, head 64 x 16 f16)"""
|
|
hx, hy = ex // 2, ey // 2
|
|
X0, Y0 = x0 // 2 - 1, y0 // 2 - 1
|
|
low = np.zeros((48, 96), np.uint8)
|
|
for r in range(36):
|
|
X, Y = mirror(X0 + r % 6, hx), mirror(Y0 + r // 6, hy)
|
|
low[r] = np.concatenate([lo[p, Y, X] for p in range(3)])
|
|
Ex = q8(mm(low, w['we'], np.repeat(w['be'][None, :], 48, 0)))
|
|
xs, ns = [], []
|
|
for (x, y) in keys(x0, y0):
|
|
xx, yy = mirror(x, ex), mirror(y, ey)
|
|
r = (y // 2 - Y0) * 6 + (x // 2 - X0)
|
|
s = 2 * (y & 1) + (x & 1)
|
|
sk = np.concatenate([skip[0, yy, xx], skip[1, yy, xx]])[UN64]
|
|
xn = (e4m3(Ex[r, 64 * s:64 * s + 64]).astype(H) + e4m3(sk).astype(H)).astype(H)
|
|
xs.append(xn)
|
|
xa = xn[I64]
|
|
n, _ = input_norm([xa[:32], xa[32:]], w['gain0'])
|
|
ns.append(q8(n)[I64])
|
|
X = np.array(ns)
|
|
res = np.array(xs[:64])
|
|
R = None
|
|
for h in range(2):
|
|
qk = mm(X[:, 32:], w[f'wk{h}b'], mm(X[:, :32], w[f'wk{h}a'], np.zeros((96, 32))))
|
|
v = mm(X[:, 32:], w[f'wv{h}b'], mm(X[:, :32], w[f'wv{h}a'], np.zeros((96, 32))))
|
|
Q = q8(qk)[:, I32]
|
|
S = mm(Q[:64], Q.T, w[f'bias{h}'])
|
|
P = softmax_cubic(S)
|
|
O = mm(q8(P)[:, I96], q8(v)[I96], np.zeros((64, 32)))
|
|
C = add(res, w['table7'][None, :]) if h == 0 else R
|
|
R = mm(q8(O)[:, I32], w[f'wo{h}'], C)
|
|
N = q8((R * (nv_rowsum_sq(R)[:, None] * w['gain14'][None, :]).astype(H)).astype(H))[:, I64]
|
|
acc = add(R, w['table16'][None, :])
|
|
for b in range(8):
|
|
hid = act(mm(N, w[f'w1_{b}'], np.repeat(w[f'b1_{b}'][None, :], 64, 0)))
|
|
acc = mm(q8(hid)[:, I32], w[f'w2_{b}'], acc)
|
|
tokens = q8(acc)[:, I64]
|
|
head = mm(tokens, w['wh'], np.repeat(w['bh'][None, :], 64, 0))
|
|
return tokens, head
|
|
|
|
|
|
def filter112(head, t88, t96, t104, x0, y0, ex, ey):
|
|
"""dec0_ref.outputs' surface 112 on the head's channels 0-8 (logits) and 9 (gate)"""
|
|
out = []
|
|
for t in range(64):
|
|
x, y = x0 + t % 8, y0 + t // 8
|
|
a = head[t].astype(H)
|
|
lg = a[:9]
|
|
m01 = np.maximum(np.maximum(np.maximum(lg[0:2], lg[2:4]), lg[4:6]), lg[6:8])
|
|
m = max(max(m01[0], m01[1]), lg[8])
|
|
e = [np.exp2(F(H(lg[i] - m)) * F(1.4426950216293335)) for i in range(9)]
|
|
pair = [F(e[2 * k] + e[2 * k + 1]) for k in range(4)]
|
|
sm = F(pair[0] + F(0))
|
|
for k in (1, 2, 3):
|
|
sm = F(sm + pair[k])
|
|
sm = F(sm + e[8])
|
|
r = H(F(1) / sm)
|
|
wg = [H(H(e[i]) * r) for i in range(9)]
|
|
|
|
def cl(v, n):
|
|
return 1 if v < 0 else (n - 2 if v >= n else v)
|
|
f = np.zeros(3, H)
|
|
k = 0
|
|
for yy in (cl(y - 1, ey), cl(y, ey), cl(y + 1, ey)):
|
|
for xx in (cl(x - 1, ex), cl(x, ex), cl(x + 1, ex)):
|
|
T = bilinear88(t88, xx, yy)
|
|
for c in range(3):
|
|
f[c] = fma16(T[c], wg[k], f[c])
|
|
k += 1
|
|
g = H(F(1) / (np.exp2(F(a[9]) * F(-1.4426950216293335)) + F(1)))
|
|
om = H(H(1) - g)
|
|
o = np.zeros(3, H)
|
|
if 0 <= x < ex and 0 <= y < ey:
|
|
for c in range(3):
|
|
o[c] = fma16(t96[y, x, c], om, H(f[c] * g))
|
|
if not np.isnan(t104[y, x, c]):
|
|
o[c] = H(o[c] + t104[y, x, c])
|
|
out.append(o)
|
|
return np.array(out)
|
|
|
|
|
|
def extract(dump, weights, out):
|
|
z = np.load(dump)
|
|
W = np.frombuffer(weights, np.uint8)
|
|
|
|
def B(p, wv=0):
|
|
return z[f'p{p}_w{wv}_B']
|
|
|
|
def C(p, wv=0):
|
|
return f16(z[f'p{p}_w{wv}_C']).astype(H)
|
|
|
|
def table(off, n=64):
|
|
return W[off:off + 2 * n].view(H)
|
|
w = {'we': np.hstack([B(1, 0), B(1, 1)]), 'be': np.concatenate([C(1, 0)[0], C(1, 1)[0]]),
|
|
'gain0': table(GAIN0), 'table7': table(TABLE7), 'gain14': table(GAIN14), 'table16': table(TABLE16)}
|
|
for h in range(2):
|
|
w[f'wk{h}a'], w[f'wv{h}a'], w[f'wk{h}b'], w[f'wv{h}b'] = B(2, h), B(3, h), B(4, h), B(5, h)
|
|
w[f'bias{h}'] = C(6, h)
|
|
w[f'wo{h}'] = B(8 + h)
|
|
for b in range(8):
|
|
w[f'w1_{b}'], w[f'b1_{b}'] = B(10 + 2 * b), C(10 + 2 * b)[0]
|
|
w[f'w2_{b}'] = B(11 + 2 * b)
|
|
w['wh'], w['bh'] = B(26), C(26)[0]
|
|
np.savez_compressed(out, **w)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument('capture')
|
|
ap.add_argument('frame', type=int)
|
|
ap.add_argument('weights')
|
|
ap.add_argument('--extract', metavar='DUMP')
|
|
ap.add_argument('--blocks', nargs='*', default=['5,5'])
|
|
o = ap.parse_args()
|
|
from rr_capture import Capture
|
|
c = Capture(os.path.expanduser(o.capture))
|
|
l = c.find('swin_dec1', o.frame)
|
|
if o.extract:
|
|
extract(o.extract, l.weights, o.weights)
|
|
return
|
|
ex, ey = (int(v) for v in np.frombuffer(l.args[:8], np.uint32))
|
|
sx, sy = (int(v) for v in np.frombuffer(l.args[80:84], np.uint16))
|
|
ptr = {p.offset: p.delta for p in l.pointers}
|
|
pre = np.frombuffer(l.pre(l.arena_alloc), np.uint8)
|
|
post = np.frombuffer(l.post_buffer(c, l.arena_alloc), np.uint8)
|
|
lo = pre[ptr[8]:ptr[8] + 3 * (ex // 2) * (ey // 2) * 32].reshape(3, ey // 2, ex // 2, 32)
|
|
skip = pre[ptr[24]:ptr[24] + 2 * ex * ey * 32].reshape(2, ey, ex, 32)
|
|
outp = post[ptr[48]:ptr[48] + 2 * ex * ey * 32].reshape(2, ey, ex, 32)
|
|
img = {i.offset: i for i in l.images}
|
|
|
|
def surf(off):
|
|
i = img[off]
|
|
return np.frombuffer(i.path.read_bytes(), np.float16).reshape(i.height, i.width, 4)
|
|
t88, t96, t104, r112 = surf(88), surf(96), surf(104), surf(112)
|
|
w = dict(np.load(o.weights))
|
|
for blk in o.blocks:
|
|
bx, by = (int(v) for v in blk.split(','))
|
|
x0, y0 = 8 * bx - sx, 8 * by - sy
|
|
tokens, head = window(w, lo, skip, x0, y0, ex, ey)
|
|
s112 = filter112(head, t88, t96, t104, x0, y0, ex, ey)
|
|
inside = [t for t in range(64) if 0 <= x0 + t % 8 < ex and 0 <= y0 + t // 8 < ey]
|
|
tk = sum(np.array_equal(tokens[t], np.concatenate([outp[0, y0 + t // 8, x0 + t % 8], outp[1, y0 + t // 8, x0 + t % 8]]))
|
|
for t in inside)
|
|
sf = sum(np.array_equal(s112[t], r112[y0 + t // 8, x0 + t % 8, :3]) for t in inside)
|
|
print(f'block {bx},{by}: tokens {tk}/{len(inside)}, surface 112 {sf}/{len(inside)}')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|