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
331 lines
14 KiB
Python
331 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
"""Exact numpy reference of the denoiser's encoder 1 (cuda_dldn_engine_swin_enc1_kernel), one 8x8 window at a time,
|
|
from its weights and real inputs only: the specification the native shader implements.
|
|
|
|
Derived phase by phase against ZLUDA replays (engine/rr/phase_dump.py), every step bit-exact:
|
|
prologue the 96 key tokens (the window plus a one-token border, corners excluded; mirrored at the image edge)
|
|
from the arena's two 32-channel planes, f16 RMS norm (NVIDIA's summation order) times the gain at W+0
|
|
attention two heads: QK and V projections (K = 64 in two k32 steps), scores = q8(QK[:64]) q8(QK)^T + bias,
|
|
the cubic-exponent softmax, P V, out-projection onto the residual (input tile + table at W+36992)
|
|
MLP RMS norm times the gain at W+37120, eight branches of 64 -> 32 (bias, ramp activation) -> 64,
|
|
accumulated onto the residual (+ table at W+41408)
|
|
outputs q8 tokens to param+48, the 2x2 merge (16 x 256 -> 96 + bias) to param+56, the 64 -> 16 head and its
|
|
3x3 kernel-prediction filter of texture 88 to surfaces 112, 120 and (half resolution) 128
|
|
|
|
usage: enc1_ref.py CAPTURE FRAME WEIGHTS.npz [--blocks X,Y ...] (verifies against the capture)
|
|
WEIGHTS.npz: engine/rr/enc1_ref.py --extract DUMP.npz CAPTURE FRAME OUT.npz builds it 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, mma_k32 # noqa: E402
|
|
|
|
H = np.float16
|
|
F = np.float32
|
|
|
|
|
|
def inter(n):
|
|
perm = []
|
|
for b in range(n // 32):
|
|
for m in range(8):
|
|
perm += [32 * b + 2 * m, 32 * b + 2 * m + 1, 32 * b + 16 + 2 * m, 32 * b + 17 + 2 * m]
|
|
return perm
|
|
|
|
|
|
INTER64 = inter(64)
|
|
# arena channel order (NVIDIA's A-fragment order) -> natural: natural channel k is arena channel ARENA[k]
|
|
ARENA = [0, 1, 4, 5, 8, 9, 12, 13, 16, 17, 20, 21, 24, 25, 28, 29, 2, 3, 6, 7, 10, 11, 14, 15, 18, 19, 22, 23, 26, 27,
|
|
30, 31]
|
|
ARENA = ARENA + [32 + c for c in ARENA]
|
|
|
|
|
|
def add(a, b):
|
|
return (a + b).astype(H)
|
|
|
|
|
|
def mm(a_codes, b_codes, c):
|
|
"""a (M x K e4m3 codes, NVIDIA k order) . b (K x N codes) + c (M x N f16 values) with ZLUDA's k32 rounding"""
|
|
return f16(mma_k32(a_codes, b_codes, np.asarray(c, H).view(np.uint16))).astype(H)
|
|
|
|
|
|
def nv_rowsum_sq(x, eps=2.0 ** -13):
|
|
"""f16(rsqrt(sum of squares + eps)) per row of a 64-column C/D tile, NVIDIA's order"""
|
|
sq = (x * x).astype(H)
|
|
|
|
def tile(j, t):
|
|
return sq[:, 8 * j + 2 * t: 8 * j + 2 * t + 2]
|
|
h = []
|
|
for t in range(4):
|
|
a, b = add(tile(0, t), tile(2, t)), add(tile(1, t), tile(3, t))
|
|
c, d = add(tile(4, t), tile(6, t)), add(tile(5, t), tile(7, t))
|
|
h.append(add(add(a, c), add(b, d)))
|
|
h1 = [add(h[t], h[t ^ 1]) for t in range(4)]
|
|
h2 = add(h1[0], h1[2])
|
|
s = add(add(h2[:, 0], h2[:, 1]), np.full(x.shape[0], H(eps)))
|
|
return (F(1) / np.sqrt(s.astype(F))).astype(H)
|
|
|
|
|
|
def input_norm(tok_arena, gain):
|
|
"""the prologue's per-token norm: tok_arena (2 planes x 32, arena order) -> f16 normalised token (natural)"""
|
|
sq = [(v * v).astype(H) for v in tok_arena]
|
|
h = []
|
|
for t in range(4):
|
|
def pair(p, w):
|
|
s = sq[p][4 * w: 4 * w + 4]
|
|
return add(s[0:2], s[2:4])
|
|
h.append(add(add(pair(0, t), pair(1, t)), add(pair(0, t + 4), pair(1, t + 4))))
|
|
h01 = [add(h[t], h[t ^ 1]) for t in range(4)]
|
|
h2 = add(h01[0], h01[2])
|
|
s = H(H(h2[0] + h2[1]) + H(2.0 ** -13))
|
|
r = H(F(1) / np.sqrt(F(s)))
|
|
x = np.concatenate(tok_arena)[ARENA]
|
|
return (x * (r * gain).astype(H)).astype(H), x
|
|
|
|
|
|
def softmax_cubic(x):
|
|
"""E5: the cubic-exponent softmax of a 64 x 96 score block (natural columns)"""
|
|
def fma16(a, b, c):
|
|
return (a.astype(np.float64) * b.astype(np.float64) + c.astype(np.float64)).astype(H)
|
|
m = np.clip((x * H(0.017181396484375)).astype(H), H(-0.55615234375), H(0.55615234375))
|
|
t = fma16(m, -m, np.full_like(m, H(0.92724609375)))
|
|
u = fma16(m, t, np.full_like(m, H(1.375)))
|
|
e = ((((u.view(np.uint16).astype(np.uint32)) << 5) & 0x7FE0).astype(np.uint16)).view(H)
|
|
|
|
def tile(j, tt, half):
|
|
return e[:, 8 * j + 2 * tt + half]
|
|
s = {}
|
|
for tt in range(4):
|
|
for half in range(2):
|
|
ev = add(add(add(tile(0, tt, half), tile(2, tt, half)), add(tile(4, tt, half), tile(6, tt, half))),
|
|
add(tile(8, tt, half), tile(10, tt, half)))
|
|
od = add(add(add(tile(1, tt, half), tile(3, tt, half)), add(tile(5, tt, half), tile(7, tt, half))),
|
|
add(tile(9, tt, half), tile(11, tt, half)))
|
|
s[tt, half] = add(ev, od)
|
|
out = [add(add(s[0, h_], s[1, h_]), add(s[2, h_], s[3, h_])) for h_ in range(2)]
|
|
rs = add(out[0], out[1])
|
|
r = (F(1) / rs.astype(F)).astype(H)
|
|
return (e * r[:, None]).astype(H)
|
|
|
|
|
|
def act(x):
|
|
y = np.clip(x, H(-2), H(2)).astype(H)
|
|
g = add(H(0.5), (y * (H(0.412109375) - (H(0.0810546875) * np.abs(y)).astype(H)).astype(H)).astype(H))
|
|
return (x * g).astype(H)
|
|
|
|
|
|
MERGE_ORDER = [(0, 0), (0, 1), (1, 0), (1, 1), (2, 0), (2, 1), (3, 0), (3, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 2),
|
|
(2, 3), (3, 2), (3, 3)]
|
|
|
|
|
|
def mirror(v, n):
|
|
v = abs(v)
|
|
return min(v, 2 * n - 2 - v)
|
|
|
|
|
|
def window(w, planes, tex, x0, y0, ex, ey, dbg=None):
|
|
"""one window: returns (tokens out [64 x 64 codes, arena order], merged [16 x 96 codes, A order],
|
|
surf112 / surf120 [64 x 3 f16], surf128 contributions)"""
|
|
border = ([(x, y0 - 1) for x in range(x0, x0 + 8)] + [(x, y0 + 8) for x in range(x0, x0 + 8)]
|
|
+ [(x0 - 1, y) for y in range(y0, y0 + 8)] + [(x0 + 8, y) for y in range(y0, y0 + 8)])
|
|
coords = [(x, y) for y in range(y0, y0 + 8) for x in range(x0, x0 + 8)] + border
|
|
norm, raw = [], []
|
|
for (x, y) in coords:
|
|
xx, yy = mirror(x, ex), mirror(y, ey)
|
|
n_, r_ = input_norm([e4m3(planes[p, yy, xx]).astype(H) for p in (0, 1)], w['gain0'])
|
|
norm.append(n_)
|
|
raw.append(r_)
|
|
X = q8(np.array(norm))[:, INTER64] # 96 x 64 codes, NVIDIA k order
|
|
res = np.array(raw[:64])
|
|
dbg = {} if dbg is None else dbg
|
|
dbg['XQ'], dbg['RES0'] = X, add(res, w['table7'][None, :])
|
|
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)[:, inter(32)]
|
|
dbg[f'QK{h}'], dbg[f'V{h}'] = Q, q8(v)[inter(96)]
|
|
S = mm(Q[:64], Q.T, w[f'bias{h}'])
|
|
P = softmax_cubic(S)
|
|
O = mm(q8(P)[:, inter(96)], q8(v)[inter(96)], np.zeros((64, 32)))
|
|
dbg[f'S{h}'], dbg[f'P{h}'], dbg[f'O{h}'] = S, q8(P)[:, inter(96)], O
|
|
C = add(res, w['table7'][None, :]) if h == 0 else R
|
|
R = mm(q8(O)[:, inter(32)], w[f'wo{h}'], C)
|
|
dbg[f'R{h}'] = R
|
|
N = q8((R * (nv_rowsum_sq(R)[:, None] * w['gain14'][None, :]).astype(H)).astype(H))[:, INTER64]
|
|
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)[:, inter(32)], w[f'w2_{b}'], acc)
|
|
dbg['ACC'] = acc
|
|
tokens = q8(acc)[:, INTER64]
|
|
Mg = np.zeros((16, 256), 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, 64 * blk: 64 * blk + 64] = tokens[(2 * my + dy) * 8 + 2 * mx + dx]
|
|
merged = q8(mm(Mg, w['wm'], np.repeat(w['bm'][None, :], 16, 0)))[:, inter(96)]
|
|
head = mm(tokens, w['wh'], np.repeat(w['bh'][None, :], 64, 0))
|
|
s112, s120 = filter_head(head, tex, x0, y0, ex, ey)
|
|
return tokens, merged, s112, s120
|
|
|
|
|
|
def filter_head(D, tex, x0, y0, ex, ey, ungated=False):
|
|
"""surfaces 112 and 120 of the encoders (the gate's split of the filtered value); ungated (the bottleneck enc5):
|
|
the filtered value itself, for surface 112"""
|
|
def sig(v):
|
|
return H(F(1) / (np.exp2(F(v) * F(-1.4426950216293335)) + F(1)))
|
|
|
|
def ex2(v):
|
|
return np.exp2(F(v) * F(1.4426950216293335))
|
|
out112, out120 = np.zeros((64, 3), H), np.zeros((64, 3), H)
|
|
for tk in range(64):
|
|
x, y = x0 + tk % 8, y0 + tk // 8
|
|
l_ = D[tk].astype(H)
|
|
m01 = np.maximum(np.maximum(np.maximum(l_[0:2], l_[2:4]), l_[4:6]), l_[6:8])
|
|
m = max(max(m01[0], m01[1]), l_[8])
|
|
e = [ex2(H(l_[i] - m)) for i in range(9)]
|
|
pair = [F(e[2 * k] + e[2 * k + 1]) for k in range(4)]
|
|
s = F(F(pair[0]) + F(0))
|
|
for k in (1, 2, 3):
|
|
s = F(s + pair[k])
|
|
s = F(s + e[8])
|
|
r = H(F(1) / s)
|
|
wgt = [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)
|
|
xm, xc, xp = cl(x - 1, ex), cl(x, ex), cl(x + 1, ex)
|
|
ym, yc, yp = cl(y - 1, ey), cl(y, ey), cl(y + 1, ey)
|
|
|
|
def T(xx, yy):
|
|
return tex[yy, xx].astype(H)
|
|
pts = [((xm, ym), (xc, ym)), ((xp, ym), (xm, yc)), ((xc, yc), (xp, yc)), ((xm, yp), (xc, yp))]
|
|
acc = np.zeros((3, 2), H)
|
|
for k, (pa, pb) in enumerate(pts):
|
|
for c in range(3):
|
|
for h_, pt in enumerate((pa, pb)):
|
|
acc[c, h_] = H(F(T(*pt)[c]) * F(wgt[2 * k + h_]) + F(acc[c, h_]))
|
|
f = [H(acc[c, 0] + acc[c, 1]) for c in range(3)]
|
|
f = [H(F(T(xp, yp)[c]) * F(wgt[8]) + F(f[c])) for c in range(3)]
|
|
if ungated:
|
|
out112[tk] = f
|
|
continue
|
|
g = [sig(l_[9]), sig(l_[10]), sig(l_[11])]
|
|
a = [H(g[c] * f[c]) for c in range(3)]
|
|
out120[tk] = a
|
|
out112[tk] = [H(f[c] - a[c]) for c in range(3)]
|
|
return out112, out120
|
|
|
|
|
|
def extract(dump, cap, frame, out):
|
|
"""the weights, from a phase dump of any block (they do not depend on the block) and the weight image"""
|
|
z = np.load(dump)
|
|
|
|
def B(p):
|
|
return z[f'p{p}_w0_B']
|
|
|
|
def C(p):
|
|
return f16(z[f'p{p}_w0_C']).astype(H)
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from rr_capture import Capture
|
|
c = Capture(os.path.expanduser(cap))
|
|
W = np.frombuffer(c.find('swin_enc1', frame).weights, np.uint8)
|
|
|
|
def table(off, n=64):
|
|
return W[off:off + 2 * n].view(np.float16)
|
|
w = {'gain0': table(0), 'gain14': table(37120), 'table7': table(36992), 'table16': table(41408)}
|
|
for h, base in ((0, 1), (1, 8)):
|
|
w[f'wk{h}a'], w[f'wv{h}a'], w[f'wk{h}b'], w[f'wv{h}b'] = B(base), B(base + 1), B(base + 2), B(base + 3)
|
|
w[f'bias{h}'] = C(base + 4)
|
|
w[f'wo{h}'] = B(base + 6)
|
|
for b in range(8):
|
|
w[f'w1_{b}'], w[f'b1_{b}'] = B(15 + 2 * b), C(15 + 2 * b)[0]
|
|
w[f'w2_{b}'] = B(16 + 2 * b)
|
|
w['wm'], w['bm'] = B(31), C(31)[0]
|
|
w['wh'], w['bh'] = B(32), C(32)[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()
|
|
if o.extract:
|
|
extract(o.extract, o.capture, o.frame, o.weights)
|
|
return
|
|
from rr_capture import Capture
|
|
c = Capture(os.path.expanduser(o.capture))
|
|
l = c.find('swin_enc1', o.frame)
|
|
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))
|
|
pre = np.frombuffer(l.pre(l.arena_alloc), np.uint8)
|
|
post = np.frombuffer(l.post_buffer(c, l.arena_alloc), np.uint8)
|
|
ptr = {p.offset: p.delta for p in l.pointers}
|
|
planes = pre[ptr[8]:ptr[8] + 2 * ex * ey * 32].reshape(2, ey, ex, 32)
|
|
outp = post[ptr[48]:ptr[48] + 2 * ex * ey * 32].reshape(2, ey, ex, 32)
|
|
mrg = post[ptr[56]:ptr[56] + 3 * (ex // 2) * (ey // 2) * 32].reshape(3, ey // 2, ex // 2, 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)
|
|
tex, s112, s120 = surf(88), surf(112), surf(120)
|
|
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, merged, o112, o120 = window(w, planes, tex, x0, y0, ex, ey)
|
|
tk_ok = 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 range(64) if 0 <= x0 + t % 8 < ex and 0 <= y0 + t // 8 < ey)
|
|
mg_ok = sum(np.array_equal(merged[r], np.concatenate([mrg[p, y0 // 2 + my, x0 // 2 + mx] for p in range(3)]))
|
|
for r, (mx, my) in enumerate(MERGE_ORDER)
|
|
if 0 <= x0 // 2 + mx < ex // 2 and 0 <= y0 // 2 + my < ey // 2)
|
|
inside = [t for t in range(64) if 0 <= x0 + t % 8 < ex and 0 <= y0 + t // 8 < ey]
|
|
s_ok = sum(np.array_equal(o112[t], s112[y0 + t // 8, x0 + t % 8, :3]) and
|
|
np.array_equal(o120[t], s120[y0 + t // 8, x0 + t % 8, :3]) for t in inside)
|
|
print(f'block {bx},{by}: tokens {tk_ok}/{len(inside)}, merged {mg_ok}, surfaces {s_ok}/{len(inside)}')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|
|
|
|
|
|
def pairs_tree(v):
|
|
"""NVIDIA's reduction over a list: (v0 + v1) + v2 for three, (v0 + v1) + (v2 + v3) for four, then the same
|
|
recursively on halves rounded up ((v0 + v1) + (v2 + v3)) + v4 for five"""
|
|
if len(v) == 1:
|
|
return v[0]
|
|
if len(v) == 2:
|
|
return add(v[0], v[1])
|
|
half = 1 << ((len(v) - 1).bit_length() - 1)
|
|
return add(pairs_tree(v[:half]), pairs_tree(v[half:]))
|
|
|
|
|
|
def rowsum_sq(x, eps=2.0 ** -13):
|
|
"""nv_rowsum_sq for any C = 8 n8-tiles: per lane word t the tiles of each parity (tile pairs 2 apart, tree),
|
|
the parities, then t ^ 1, t ^ 2 and the halves"""
|
|
sq = (x * x).astype(H)
|
|
nt = x.shape[1] // 8
|
|
|
|
def tile(j, t):
|
|
return sq[:, 8 * j + 2 * t: 8 * j + 2 * t + 2]
|
|
h = []
|
|
for t in range(4):
|
|
ev = pairs_tree([add(tile(j, t), tile(j + 2, t)) for j in range(0, nt, 4)] if nt % 4 == 0 else
|
|
[tile(j, t) for j in range(0, nt, 2)])
|
|
od = pairs_tree([add(tile(j, t), tile(j + 2, t)) for j in range(1, nt, 4)] if nt % 4 == 0 else
|
|
[tile(j, t) for j in range(1, nt, 2)])
|
|
h.append(add(ev, od))
|
|
h1 = [add(h[t], h[t ^ 1]) for t in range(4)]
|
|
h2 = add(h1[0], h1[2])
|
|
s = add(add(h2[:, 0], h2[:, 1]), np.full(x.shape[0], H(eps)))
|
|
return (F(1) / np.sqrt(s.astype(F))).astype(H)
|