engine/rr: native Ray Reconstruction stages with capture/replay and shim integration
- 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
This commit is contained in:
woutermans committed
1 parent
30472ae195
commit
50fa65d87b
67 files changed
+11926
-504
No files matched your search
@@ -0,0 +1,183 @@
|
||||
#!/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()
|
||||
Reference in new issue
Block a user