woutermans
50fa65d87b
- engine/rr: d4r's own shaders for the denoiser frame (exposure, six encoders,
five decoders, hkpn output stage) at NGX layouts, plus PTX->GLSL/HIP packing
scripts, reference implementations and weight extraction
- runtime.h/runtime_rr.inc, game_rr.{h,cpp}: GameDenoiser wired into the engine
build (d4r-rr replay binary) and the D3D12 shim's feature-13 path
- shim: answer GetFeatureRequirements from d4r's own state only (the old
cross-binary CUDA handoff crashed OptiScaler's scan); drop the core's
CUDA GetFeatureRequirements import
- shim/harness: exact guide capture/replay of a frame's RR inputs
- kernels/rr: expanded enc0/engine prototypes; kernels/tools: ptxsim rewrite
plus rrsim and an engine test; vulkan: ptx2glsl FP8/cooperative-matrix
support, kvk inout buffers and KVK_DISABLE_OPTIMIZATION
- bridge: report the surface array in cuSurfObjectCreate traces
- packaging: RR settings and the ICU launch overrides in D4R_README
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()
|