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.
346 lines
15 KiB
Python
346 lines
15 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)
|
|
|
|
|
|
# D4R_RR_FUSE=1: the arithmetic as an NVIDIA card runs it. NVIDIA's compiler contracts a binary16 multiply whose
|
|
# product only feeds an add / sub into a fused multiply-add (the kernels give those no rounding modifier); ZLUDA, and
|
|
# so the captures, round the product first. build/rr/f/nv has the measurements (ptx_contract.py is the rule).
|
|
FUSE = os.environ.get('D4R_RR_FUSE') == '1'
|
|
|
|
|
|
def fma16(a, b, c):
|
|
"""binary16 a * b + c with one rounding (the exact value fits a double)"""
|
|
with np.errstate(all='ignore'):
|
|
return (np.asarray(a, np.float64) * np.asarray(b, np.float64) + np.asarray(c, np.float64)).astype(H)
|
|
|
|
|
|
def act(x):
|
|
y = np.clip(x, H(-2), H(2)).astype(H)
|
|
if FUSE:
|
|
g = fma16(y, fma16(-np.abs(y), H(0.0810546875), H(0.412109375)), H(0.5))
|
|
return (x * g).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)
|