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.
139 lines
6.7 KiB
Python
139 lines
6.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Exact numpy reference of the denoiser's encoder 0 (cuda_dldn_engine_swin_enc0_kernel) after its input features,
|
|
one 8x8 window at a time. The 32 FP8 features of the 96 key pixels (NVIDIA's per-pixel gather from the game's guides
|
|
and the temporal state, which engine/rr/enc0_m.comp keeps as ptx2rr.py's translation) are an input here.
|
|
|
|
embedding F (96 x 32) -> q8(relu(F W1 + b1)) -> W2 + b2: the key tokens x (32 channels, f16)
|
|
block dec0_ref's: RMS norm (gain W+6272), residual x + table W+37056, two heads, four MLP branches
|
|
outputs the tokens (argument 48, dec0's skip), the 2x2 merge 128 -> 64 (argument 56, enc1's input) and the
|
|
32 -> 16 head with enc1's filter to surfaces 112, 120 and 128
|
|
|
|
usage: enc0_ref.py CAPTURE FRAME WEIGHTS.npz --dump DUMP.npz [--block X,Y] (verifies the dump's block)
|
|
enc0_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 f16, q8 # noqa: E402
|
|
from enc1_ref import inter, mm, softmax_cubic, act, add, filter_head, MERGE_ORDER, H # noqa: E402
|
|
from dec0_ref import norm_plane # noqa: E402
|
|
|
|
I32, I64, I96 = inter(32), inter(64), inter(96)
|
|
GAIN0, TABLE0, GAIN1, TABLE1 = 6272, 37056, None, None
|
|
|
|
|
|
def block(w, feat):
|
|
"""feat: the 96 key pixels' features (codes, NVIDIA's k order). Returns (tokens 64 x 32 codes in arena order,
|
|
merged 16 x 64 codes in arena order, head 64 x 16 f16, dbg)"""
|
|
dbg = {}
|
|
e1 = mm(feat, w['w1'], np.repeat(w['b1'][None, :], 96, 0))
|
|
x = mm(q8(np.maximum(e1, H(0)))[:, I32], w['w2'], np.repeat(w['b2'][None, :], 96, 0))
|
|
dbg['x'] = x
|
|
X = np.array([norm_plane(x[k], w['gain0']) for k in range(96)])
|
|
R = None
|
|
for h in range(2):
|
|
qk = mm(X, w[f'wk{h}'], np.zeros((96, 32)))
|
|
v = mm(X, w[f'wv{h}'], np.zeros((96, 32)))
|
|
Q = q8(qk)[:, I32]
|
|
P = softmax_cubic(mm(Q[:64], Q.T, w[f'bias{h}']))
|
|
O = mm(q8(P)[:, I96], q8(v)[I96], np.zeros((64, 32)))
|
|
C = add(x[:64], w['table0'][None, :]) if h == 0 else R
|
|
R = mm(q8(O)[:, I32], w[f'wo{h}'], C)
|
|
N = np.array([norm_plane(R[t], w['gain1']) for t in range(64)])
|
|
acc = add(R, w['table1'][None, :])
|
|
for b in range(4):
|
|
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)[:, I32]
|
|
Mg = np.zeros((16, 128), np.uint8)
|
|
for r, (mx, my) in enumerate(MERGE_ORDER):
|
|
for blk, (dx, dy) in enumerate([(0, 0), (1, 0), (0, 1), (1, 1)]):
|
|
Mg[r, 32 * blk:32 * blk + 32] = tokens[(2 * my + dy) * 8 + 2 * mx + dx]
|
|
merged = q8(mm(Mg, w['wm'], np.repeat(w['bm'][None, :], 16, 0)))[:, I64]
|
|
head = mm(tokens, w['wh'], np.repeat(w['bh'][None, :], 64, 0))
|
|
return tokens, merged, head, dbg
|
|
|
|
|
|
def extract(dump, weights, out):
|
|
z = np.load(dump)
|
|
W = np.frombuffer(weights, np.uint8)
|
|
|
|
def B(p):
|
|
return z[f'p{p}_w0_B']
|
|
|
|
def C(p):
|
|
return f16(z[f'p{p}_w0_C']).astype(H)
|
|
w = {'w1': B(1), 'b1': C(1)[0], 'w2': B(2), 'b2': C(2)[0], 'gain0': W[GAIN0:GAIN0 + 64].view(H),
|
|
'table0': W[TABLE0:TABLE0 + 64].view(H)}
|
|
# the input norm's gain and the residual table sit elsewhere in preset D's (smaller) weight image: when the preset
|
|
# E offsets do not reproduce the dump, find them by matching it, like the MLP's below
|
|
x = mm(q8(np.maximum(mm(z['p1_w0_A'], w['w1'], np.repeat(w['b1'][None, :], 96, 0)), H(0)))[:, I32], w['w2'],
|
|
np.repeat(w['b2'][None, :], 96, 0))
|
|
cands = [W[o:o + 64].view(H) for o in range(0, len(W) - 63, 2)]
|
|
if not all(np.array_equal(norm_plane(x[k], w['gain0']), z['p3_w0_A'][k]) for k in (0, 9)):
|
|
w['gain0'] = next(t for t in cands if all(np.array_equal(norm_plane(x[k], t), z['p3_w0_A'][k]) for k in (0, 9)))
|
|
if not np.array_equal(add(x[:2], w['table0'][None]).view(np.uint16), C(7)[:2].view(np.uint16)):
|
|
w['table0'] = next(t for t in cands if np.array_equal(add(x[:2], t[None]).view(np.uint16), C(7)[:2].view(np.uint16)))
|
|
for h, base in ((0, 3), (1, 8)):
|
|
w[f'wk{h}'], w[f'wv{h}'] = B(base), B(base + 1)
|
|
w[f'bias{h}'] = C(base + 2)
|
|
w[f'wo{h}'] = B(base + 4)
|
|
for b in range(4):
|
|
w[f'w1_{b}'], w[f'b1_{b}'] = B(13 + 2 * b), C(13 + 2 * b)[0]
|
|
w[f'w2_{b}'] = B(14 + 2 * b)
|
|
w['wm'], w['bm'] = B(21), C(21)[0]
|
|
w['wh'], w['bh'] = B(22), C(22)[0]
|
|
# the MLP norm's gain and table, found by matching the dump
|
|
R = f16(z['p12_w0_D']).astype(H)
|
|
tables = [W[o:o + 64].view(H) for o in range(0, len(W) - 63, 2)]
|
|
w['gain1'] = next(t for t in tables if all(np.array_equal(norm_plane(R[k], t), z['p13_w0_A'][k]) for k in (0, 9)))
|
|
w['table1'] = next(t for t in tables if np.array_equal(add(R[:2], t[None]).view(np.uint16), C(14)[:2].view(np.uint16)))
|
|
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('--dump')
|
|
ap.add_argument('--block', default='5,5')
|
|
o = ap.parse_args()
|
|
from rr_capture import Capture
|
|
c = Capture(os.path.expanduser(o.capture))
|
|
l = c.find('swin_enc0', o.frame)
|
|
if o.extract:
|
|
extract(o.extract, l.weights, o.weights)
|
|
return
|
|
z = np.load(o.dump)
|
|
w = dict(np.load(o.weights))
|
|
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))
|
|
bx, by = (int(v) for v in o.block.split(','))
|
|
x0, y0 = 8 * bx - sx, 8 * by - sy
|
|
tokens, merged, head, dbg = block(w, z['p1_w0_A'])
|
|
for p, name in ((3, 'X'), ):
|
|
pass
|
|
ptr = {p.offset: p.delta for p in l.pointers}
|
|
post = np.frombuffer(l.post_buffer(c, l.arena_alloc), np.uint8)
|
|
otok = post[ptr[48]:ptr[48] + ex * ey * 32].reshape(ey, ex, 32)
|
|
omrg = post[ptr[56]:ptr[56] + 2 * (ex // 2) * (ey // 2) * 32].reshape(2, ey // 2, ex // 2, 32)
|
|
img = {i.offset: i for i in l.images}
|
|
tex = {k: np.frombuffer(i.path.read_bytes(), np.float16).reshape(i.height, i.width, 4) for k, i in img.items()
|
|
if k in (112, 120, 160)}
|
|
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], otok[y0 + t // 8, x0 + t % 8]) for t in inside)
|
|
mg = sum(np.array_equal(merged[r], np.concatenate([omrg[p, y0 // 2 + my, x0 // 2 + mx] for p in range(2)]))
|
|
for r, (mx, my) in enumerate(MERGE_ORDER))
|
|
print(f'block {bx},{by}: tokens {tk}/{len(inside)}, merged {mg}/16')
|
|
return head, x0, y0, ex, ey
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|