engine/rr: native preset F (DLSS 4.5 Ray Reconstruction) and the reference's fused roundings
- engine/compile_rr2.py, runtime_rr2.inc, engine/rr/rr2_m.comp, exposure_f.comp, gen_rr2_enc0.py, rr2_pack.py, rr2_ref.py: the rr2_* model (nvngx_dlssd.dll 310.9's third network) as d4r's own shaders - exposure, enc0 gather/block split, enc1-5 and dec4-0 through one matrix shader, translated output passes; weights packed from the DLL, argument blocks from the capture - ptx2rr: match the driver's fused roundings - opaque-zero OR around cvt.rn.f16.f32 against v_fma_mix, half-away-from-zero formatted stores, UNORM 10:10:10:2 as RGBA32F of q/1023, OCML tanhF with the AMDGPU exp lowering, binary16 fma and fetch roundings; --probe for register bisection against ZLUDA - dec0_m/swin_m/enc0_*: RR_MIX build option (D4R_RR_MIX=1), one v_fma_mix rounding per 32-deep step, ~5% faster, not bit-exact; default build stays exact - runtime_rr.inc: RModel/REngine dispatch to the D4RR0002 model; replay_rr: RGBA32F and A2B10G10R1010 texel sizes, roughnessPacked bit, larger debug staging - shim: engine\rr-f model slot, 310.9 version probe so preset F is only claimed when the CUDA path has it; d4r.ini / d4r_config.py Model = F; run_rr_engine_compare.sh D4R_RR_DLL - kernels/rr/swin_native.h + rr_swin_native_test.cpp: HIP port of encoder 1 for the FP8-product and rounding cost measurements - engine/rr/README: preset F's block, options, capture-by-capture results, timings and the open camera-matrix difference
This commit is contained in:
woutermans committed
1 parent
065fe2c2a1
commit
8eb7a94533
31 files changed
+3547
-128
No files matched your search
@@ -0,0 +1,362 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Exact numpy reference of preset F's encoder stages (rr2_enc1_kernel ... rr2_enc4_kernel of nvngx_dlssd.dll
|
||||
310.9), one 8x8 window at a time, from weights and real inputs only; checked against a capture (main).
|
||||
|
||||
A stage has C channels and works on windows of 64 tokens (origin 8 * block - shift). Arithmetic is binary16 with FP8
|
||||
(e4m3) operands in the products, exactly as in presets D and E (enc1_ref.py: mm, the cubic softmax, the activation):
|
||||
|
||||
merge x = W1 . [the four tokens of the 2x2 group one level up, (0,0) (1,0) (0,1) (1,1)] + b1; window tokens
|
||||
outside the level are the mirrored tokens
|
||||
heads per head, all from q8(x): QK = Wqk . x and V = Wv . x (32 channels each, no bias). The scores are
|
||||
QK . QK^T + bias: symmetric, so NVIDIA computes the blocks (queries 0-31, keys 0-31), (0-31, 32-63) and
|
||||
(32-63, 32-63) and transposes the second for (32-63, 0-31). Cubic softmax over the 64 keys, P . V, and the
|
||||
out-projections of all heads accumulate into one sum a (the first one's C is the bias bo)
|
||||
residual r = dq(q8(x)) + a * s1 (dq: the FP8 code back as binary16; s1: binary16 at argument 94)
|
||||
MLP NB branches from q8(r): C -> 32 (+ bias) -> act -> C, accumulated into m (first C: the bias bu)
|
||||
residual y = r + m * s2 (s2: argument 96); the stage's tokens are q8(y)
|
||||
outputs merged = Wm . [y of the 2x2 group] + bm for the next encoder (q8), and a 16-channel head h = Wh . y + bh
|
||||
whose first three channels gate the stage's guide image (argument 48, three binary16 channels per token):
|
||||
out = min(guide * (sigmoid(h) * s0), 65504) (s0: argument 92) at argument 64, and its 2x2 mean
|
||||
((a + b) + (c + d) of the quarters) at argument 72 for the next encoder
|
||||
|
||||
Arena layouts: tokens are FP8 codes in planes of 32 channels ([plane][y][x][32], NVIDIA's k order within a plane);
|
||||
the guide images are four binary16 per token.
|
||||
|
||||
usage: rr2_ref.py CAPTURE STAGE FRAME WEIGHTS.npz [--blocks X,Y ...] (verifies against the capture)
|
||||
rr2_ref.py CAPTURE STAGE FRAME OUT.npz --extract DUMP.npz (weights from a phase dump, phase_dump.py)
|
||||
"""
|
||||
import argparse
|
||||
import os
|
||||
import struct
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from rrmath import e4m3, q8 # noqa: E402
|
||||
from enc1_ref import inter, mm, act, add, pairs_tree, mirror, H, F # noqa: E402
|
||||
from rr_capture import Capture # noqa: E402
|
||||
|
||||
# stage: C, channels one level up (merge input), merged channels for the next encoder, heads, MLP branches
|
||||
STAGES = {'enc0': (32, 0, 0, 2, 4), 'enc1': (64, 32, 96, 2, 8), 'enc2': (96, 0, 128, 4, 12), 'enc3': (128, 0, 160, 4, 16),
|
||||
'enc4': (160, 0, 160, 8, 20), 'enc5': (160, 0, 0, 8, 20),
|
||||
'dec0_1': (32, 0, 0, 2, 4), 'dec4_1': (160, 0, 0, 8, 20), 'dec3_1': (128, 0, 0, 4, 16), 'dec2_1': (96, 0, 0, 4, 12), 'dec1_1': (64, 0, 0, 2, 8)}
|
||||
# the decoders' first kernels (decN_0): C, channels one level down, heads, MLP branches. They expand the 4 x 4 tokens
|
||||
# one level down to the window, e = We . low + be (4C outputs: sub-position (y & 1) * 2 + (x & 1), then channel), and
|
||||
# start from x = dq(q8(e)) + dq(skip) * s3 (s3: binary16 at argument 98; skip: the encoder's tokens of this level,
|
||||
# argument 16); then the heads and the MLP as above, the first residual on x itself (not dq(q8(x))). No merge, no guide.
|
||||
EXPAND = {'dec4_0': (160, 160, 8, 20), 'dec3_0': (128, 160, 4, 16), 'dec2_0': (96, 128, 4, 12), 'dec1_0': (64, 96, 2, 8),
|
||||
'dec0_0': (32, 64, 2, 4)}
|
||||
# the network's sixteen kernels in the DLL's order, with their head counts
|
||||
ORDER = ['enc0', 'enc1', 'enc2', 'enc3', 'enc4', 'enc5', 'dec4_0', 'dec4_1', 'dec3_0', 'dec3_1', 'dec2_0', 'dec2_1',
|
||||
'dec1_0', 'dec1_1', 'dec0_0', 'dec0_1']
|
||||
HEADS = [2, 2, 4, 4, 8, 8, 8, 8, 4, 4, 4, 4, 2, 2, 2, 2]
|
||||
|
||||
|
||||
def find_tables(dll):
|
||||
"""byte offsets of the sixteen kernels' relative position tables in nvngx_dlssd.dll: runs of 64 x 64 binary16
|
||||
tables (one per head) with a constant diagonal and at most 225 distinct values, in ORDER"""
|
||||
a = np.frombuffer(dll[:len(dll) & ~1], np.uint16)
|
||||
n = len(a) - 4096
|
||||
m = (a[:n] & 0x7fff) != 0
|
||||
for k in range(1, 8):
|
||||
m &= a[:n] == a[65 * k:65 * k + n]
|
||||
m &= a[1:n + 1] == a[66:66 + n]
|
||||
groups = []
|
||||
for i in np.nonzero(m)[0]:
|
||||
t = a[i:i + 4096].reshape(64, 64)
|
||||
if not (np.all(t.diagonal() == t[0, 0]) and 50 < len(np.unique(t)) <= 225):
|
||||
continue
|
||||
if groups and 2 * int(i) - groups[-1][-1] == 8192:
|
||||
groups[-1].append(2 * int(i))
|
||||
else:
|
||||
groups.append([2 * int(i)])
|
||||
for start in range(len(groups) - 15):
|
||||
if [len(g) for g in groups[start:start + 16]] == HEADS:
|
||||
return {name: g[0] for name, g in zip(ORDER, groups[start:start + 16])}
|
||||
raise SystemExit('no preset F model found in the DLL (its sixteen relative position table groups)')
|
||||
|
||||
|
||||
def dll_weights(dll, stage):
|
||||
"""an encoder stage's tensors from the DLL: plain binary16 arrays in natural [out][in] order, one after another.
|
||||
Before the tables: {s0, 0, 0, 0} and (enc1) the merge of encoder 0's tokens, W1 then b1. After them, per tensor
|
||||
kind all heads / branches together: Wqk, Wv (heads x 32 x C), Wo (C x heads x 32), bo, Wd (branches x 32 x C), bd,
|
||||
Wu (C x branches x 32), bu, the merge
|
||||
for the next encoder Wm (C_M x 4C) and bm, the guide head Wh (3 x C) and bh. A merge's input index is
|
||||
channel * 4 + sub-position; the kernels' bias is (T + T^T) * sqrt(2)."""
|
||||
if stage in EXPAND:
|
||||
c, c_lo, nh, nb = EXPAND[stage]
|
||||
c_up = c_m = 0
|
||||
else:
|
||||
c, c_up, c_m, nh, nb = STAGES[stage]
|
||||
g = find_tables(dll)[stage]
|
||||
f = lambda off, n: np.frombuffer(dll[off:off + 2 * n], H)
|
||||
|
||||
def B(mat, merge=0):
|
||||
"""[out][in] binary16 -> the product's B operand (in x out FP8 codes, the inputs in NVIDIA's k order)"""
|
||||
k = mat.shape[1]
|
||||
if merge: # channel * 4 + q -> q * channels + channel
|
||||
mat = mat.reshape(mat.shape[0], merge, 4).transpose(0, 2, 1).reshape(mat.shape[0], k)
|
||||
return q8(mat).T[inter(k)]
|
||||
w = {}
|
||||
if c_up:
|
||||
o = g - 8 - 2 * (3 * c_up + 4) - 2 * c - 2 * c * 4 * c_up # before encoder 0's guide head (3 x C_UP, bias, 0)
|
||||
w['w1'] = B(f(o, c * 4 * c_up).reshape(c, 4 * c_up), c_up)
|
||||
w['b1'] = f(o + 2 * c * 4 * c_up, c)
|
||||
if stage in EXPAND: # We (C x 4 x C_LO) and be (C x 4) end where the tables start
|
||||
o = g - 8 * c - 8 * c * c_lo
|
||||
we = f(o, 4 * c * c_lo).reshape(c, 4, c_lo).transpose(1, 0, 2).reshape(4 * c, c_lo)
|
||||
w['we'] = q8(we).T[inter(c_lo)]
|
||||
w['be'] = f(g - 8 * c, 4 * c).reshape(c, 4).T.reshape(-1)
|
||||
if stage == 'enc0':
|
||||
# encoder 0's embedding of the 43 per-pixel features (the kernel's 64 feature slots, the rest unused): two MLP
|
||||
# branches 43 -> 32 -> act -> 32, x = bu + sum: Wd (2 x 32 x 43), bd, Wu (32 x 2 x 32), bu, then 16 bytes
|
||||
o = g - 16 - 64 - 4096 - 128 - 2 * 2 * 32 * 43
|
||||
wd = np.zeros((2, 32, 64), H); wd[:, :, :43] = f(o, 2 * 32 * 43).reshape(2, 32, 43); o += 2 * 2 * 32 * 43
|
||||
w['ewd'] = np.array([B(x) for x in wd])
|
||||
w['ebd'] = f(o, 64).reshape(2, 32); o += 128
|
||||
w['ewu'] = np.array([B(x) for x in f(o, 2048).reshape(32, 2, 32).transpose(1, 0, 2)]); o += 4096
|
||||
w['ebu'] = f(o, 32)
|
||||
o = g
|
||||
t = f(o, nh * 4096).reshape(nh, 64, 64).astype(np.float64); o += nh * 8192
|
||||
w['bias'] = ((t + t.transpose(0, 2, 1)) * np.sqrt(2)).astype(H)
|
||||
w['wqk'] = np.array([B(x) for x in f(o, nh * 32 * c).reshape(nh, 32, c)]); o += nh * 64 * c
|
||||
w['wv'] = np.array([B(x) for x in f(o, nh * 32 * c).reshape(nh, 32, c)]); o += nh * 64 * c
|
||||
w['wo'] = np.array([B(x) for x in f(o, nh * 32 * c).reshape(c, nh, 32).transpose(1, 0, 2)]); o += nh * 64 * c
|
||||
w['bo'] = f(o, c); o += 2 * c
|
||||
w['wd'] = np.array([B(x) for x in f(o, nb * 32 * c).reshape(nb, 32, c)]); o += nb * 64 * c
|
||||
w['bd'] = f(o, nb * 32).reshape(nb, 32); o += nb * 64
|
||||
w['wu'] = np.array([B(x) for x in f(o, nb * 32 * c).reshape(c, nb, 32).transpose(1, 0, 2)]); o += nb * 64 * c
|
||||
w['bu'] = f(o, c); o += 2 * c
|
||||
if c_m:
|
||||
w['wm'] = B(f(o, c_m * 4 * c).reshape(c_m, 4 * c), c); o += 2 * c_m * 4 * c
|
||||
w['bm'] = f(o, c_m); o += 2 * c_m
|
||||
if stage in EXPAND:
|
||||
return w
|
||||
if stage == 'enc0': # the merge of its tokens (enc1's W1, b1) comes before its guide head
|
||||
o += 2 * 64 * 128 + 2 * 64
|
||||
nhd = 4 if stage.startswith('dec') else 3 # the guide head: three gates; decoders also a blend factor
|
||||
wh = np.zeros((16, c), H); wh[:nhd] = f(o, nhd * c).reshape(nhd, c); o += 2 * nhd * c
|
||||
bh = np.zeros(16, H); bh[:nhd] = f(o, nhd)
|
||||
w['wh'], w['bh'] = B(wh), bh
|
||||
return w
|
||||
|
||||
|
||||
def dq(x):
|
||||
"""binary16 values through FP8 and back"""
|
||||
return e4m3(q8(x)).astype(H)
|
||||
|
||||
|
||||
def A(x):
|
||||
"""binary16 rows -> the product's A operand: FP8 codes in NVIDIA's k order"""
|
||||
return q8(x)[:, inter(x.shape[1])]
|
||||
|
||||
|
||||
def softmax(s):
|
||||
"""the cubic-exponent softmax of presets D and E over the 64 keys of a row"""
|
||||
def fma16(a, b, c):
|
||||
return (a.astype(np.float64) * b.astype(np.float64) + c.astype(np.float64)).astype(H)
|
||||
m = np.clip((s * 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)
|
||||
nt = e.shape[1] // 8
|
||||
part = {}
|
||||
for tt in range(4):
|
||||
for half in range(2):
|
||||
ev = pairs_tree([e[:, 8 * j + 2 * tt + half] for j in range(0, nt, 2)])
|
||||
od = pairs_tree([e[:, 8 * j + 2 * tt + half] for j in range(1, nt, 2)])
|
||||
part[tt, half] = add(ev, od)
|
||||
out = [add(add(part[0, h], part[1, h]), add(part[2, h], part[3, h])) for h in range(2)]
|
||||
r = (F(1) / add(out[0], out[1]).astype(F)).astype(H)
|
||||
return (e * r[:, None]).astype(H)
|
||||
|
||||
|
||||
def sigmoid(d):
|
||||
"""1 / (1 + 2^(-d log2 e)) in f32, rounded to binary16 (the GPU's exp2 and division; see rr_exact.glsl sig)"""
|
||||
f = F(1) / (np.exp2(d.astype(F) * F(-1.4426950216293335)).astype(F) + F(1))
|
||||
return f.astype(H)
|
||||
|
||||
|
||||
def upsample(low, x, y):
|
||||
"""the decoders' 2x bilinear fetch of the guide one level down (half-pixel centres, edge clamped), with NVIDIA's
|
||||
binary16 rounding: every weight product and the three sums in tap order"""
|
||||
def axis(v, n):
|
||||
if v == 0:
|
||||
return 0, H(0)
|
||||
b = (v - 1) // 2
|
||||
return b, H(0.25) if v & 1 else H(0.75)
|
||||
hh, hw = low.shape[:2]
|
||||
bx, wx = axis(x, hw)
|
||||
by, wy = axis(y, hh)
|
||||
x1, y1 = min(bx + 1, hw - 1), min(by + 1, hh - 1)
|
||||
ux, uy = H(1) - wx, H(1) - wy
|
||||
m = lambda v, a, b: ((v * a).astype(H) * b).astype(H)
|
||||
return add(add(add(m(low[by, bx], ux, uy), m(low[by, x1], wx, uy)), m(low[y1, bx], ux, wy)), m(low[y1, x1], wx, wy))
|
||||
|
||||
|
||||
def window(w, scal, xin, x=None):
|
||||
"""xin: 64 x 4*C_UP FP8 codes (the merge input, already in k order). Returns y (64 x C binary16), the merged
|
||||
tokens (16 x C_M binary16) and the head (64 x 16 binary16)."""
|
||||
s0, s1, s2 = scal
|
||||
base = None
|
||||
if x is not None: # a decoder's first kernel: the expanded tokens plus the skip; the residual keeps it unquantised
|
||||
xa, base = A(x), x
|
||||
elif 'w1' in w:
|
||||
x = mm(xin, w['w1'], np.tile(w['b1'], (64, 1)))
|
||||
xa = A(x)
|
||||
else: # the tokens themselves (already FP8)
|
||||
x, xa = e4m3(xin)[:, np.argsort(inter(xin.shape[1]))].astype(H), xin
|
||||
acc = np.tile(w['bo'], (64, 1))
|
||||
for h in range(len(w['wqk'])):
|
||||
qk = mm(xa, w['wqk'][h], np.zeros((64, 32), H))
|
||||
v = mm(xa, w['wv'][h], np.zeros((64, 32), H))
|
||||
qa = A(qk)
|
||||
s = mm(qa, qa.T, w['bias'][h])
|
||||
# NVIDIA computes three 32 x 32 blocks and transposes one; the full product is the same values
|
||||
p = softmax(s)
|
||||
o = mm(A(p), A(v.T).T if False else q8(v)[inter(64), :], np.zeros((64, 32), H))
|
||||
acc = mm(A(o), w['wo'][h], acc)
|
||||
r = add(dq(x) if base is None else base, (acc * s1).astype(H))
|
||||
ra = A(r)
|
||||
m = np.tile(w['bu'], (64, 1))
|
||||
for j in range(len(w['wd'])):
|
||||
d = mm(ra, w['wd'][j], np.tile(w['bd'][j], (64, 1)))
|
||||
m = mm(A(act(d)), w['wu'][j], m)
|
||||
y = add(r, (m * s2).astype(H))
|
||||
ya = A(y)
|
||||
# the 2x2 groups of the window: token order (0,0) (1,0) (0,1) (1,1), 16 groups row-major
|
||||
g = np.zeros((16, 4 * y.shape[1]), np.uint8)
|
||||
for n in range(16):
|
||||
gx, gy = n % 4, n // 4
|
||||
for q, (dx, dy) in enumerate(((0, 0), (1, 0), (0, 1), (1, 1))):
|
||||
g[n, q * y.shape[1]:(q + 1) * y.shape[1]] = ya[(2 * gy + dy) * 8 + 2 * gx + dx]
|
||||
merged = mm(g, w['wm'], np.tile(w['bm'], (16, 1))) if 'wm' in w else None
|
||||
head = mm(ya, w['wh'], np.tile(w['bh'], (64, 1))) if 'wh' in w else None
|
||||
return y, merged, head
|
||||
|
||||
|
||||
def extract(z, out):
|
||||
"""weights of enc1 from a phase dump (phases as phase_dump.py numbers them)"""
|
||||
P = lambda p, n: z[f'p{p}_w0_{n}']
|
||||
Ch = lambda p: P(p, 'C').view(H)
|
||||
w = {'w1': P(1, 'B'), 'b1': Ch(1)[0], 'wqk': [], 'wv': [], 'bias': [], 'wo': [], 'bo': Ch(11)[0]}
|
||||
for h in range(2):
|
||||
b = 2 + 10 * h
|
||||
w['wqk'].append(np.vstack([P(b, 'B'), P(b + 2, 'B')]))
|
||||
w['wv'].append(np.vstack([P(b + 1, 'B'), P(b + 3, 'B')]))
|
||||
tl, tr, br = Ch(b + 4), Ch(b + 5), Ch(b + 7)
|
||||
w['bias'].append(np.block([[tl, tr], [tr.T, br]]))
|
||||
w['wo'].append(P(b + 9, 'B'))
|
||||
w['wd'] = [P(22 + 2 * j, 'B') for j in range(8)]
|
||||
w['bd'] = [Ch(22 + 2 * j)[0] for j in range(8)]
|
||||
w['wu'] = [P(23 + 2 * j, 'B') for j in range(8)]
|
||||
w['bu'] = Ch(23)[0]
|
||||
w['wm'], w['bm'], w['wh'], w['bh'] = P(38, 'B'), Ch(38)[0], P(39, 'B'), Ch(39)[0]
|
||||
np.savez_compressed(out, **{k: np.array(v) for k, v in w.items()})
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument('capture')
|
||||
ap.add_argument('stage')
|
||||
ap.add_argument('frame', type=int)
|
||||
ap.add_argument('weights')
|
||||
ap.add_argument('--extract')
|
||||
ap.add_argument('--blocks', nargs='*', default=['5,5'])
|
||||
o = ap.parse_args()
|
||||
if o.extract:
|
||||
return extract(np.load(o.extract), o.weights)
|
||||
expand = o.stage in EXPAND
|
||||
if expand:
|
||||
c, c_lo, heads, nb = EXPAND[o.stage]
|
||||
c_up = c_m = 0
|
||||
else:
|
||||
c, c_up, c_m, heads, nb = STAGES[o.stage]
|
||||
w = dll_weights(open(o.weights, 'rb').read(), o.stage) if o.weights.lower().endswith('.dll') else dict(np.load(o.weights))
|
||||
cap = Capture(os.path.expanduser(o.capture))
|
||||
kernel = f'rr2_{o.stage}_kernel' if o.stage.startswith('enc') else f'rr2_{o.stage[:4]}_kernel_{o.stage[5]}'
|
||||
l = [l for l in cap.launches if l.kernel == kernel][o.frame - 1]
|
||||
pre = np.frombuffer(l.pre(l.arena_alloc), np.uint8)
|
||||
post = np.frombuffer(l.post_buffer(cap, l.arena_alloc), np.uint8)
|
||||
ptr = {p.offset: p.delta for p in l.pointers if p.alloc == l.arena_alloc}
|
||||
wd, ht = struct.unpack_from('<2I', l.args, 0)
|
||||
sx, sy = struct.unpack_from('<2H', l.args, 80)
|
||||
scal = np.frombuffer(l.args[92:98], H)
|
||||
s3 = np.frombuffer(l.args[98:100], H)[0]
|
||||
planes = lambda buf, off, w_, h_, ch: buf[off:off + w_ * h_ * ch].reshape(ch // 32, h_, w_, 32)
|
||||
up = (planes(pre, ptr[8], wd // 2, ht // 2, c_lo) if expand else planes(pre, ptr[8], 2 * wd, 2 * ht, c_up) if c_up else
|
||||
planes(pre, ptr[8], wd, ht, c))
|
||||
skip = planes(pre, ptr[16], wd, ht, c) if expand else None
|
||||
guided = 48 in ptr and o.stage.startswith('enc')
|
||||
dec_guide = o.stage.startswith('dec') and 56 in ptr
|
||||
tok_ref = planes(post, ptr[32], wd, ht, c)
|
||||
mer_ref = planes(post, ptr[40], wd // 2, ht // 2, c_m) if c_m else None
|
||||
img = lambda buf, off, w_, h_: np.frombuffer(bytes(buf[off:off + w_ * h_ * 8]) + bytes(8), H)[:w_ * h_ * 4].reshape(h_, w_, 4)
|
||||
guide, out_ref = (img(pre, ptr[48], wd, ht), img(post, ptr[64], wd, ht)) if guided else (None, None)
|
||||
half_ref = img(post, ptr[72], wd // 2, ht // 2) if c_m else None
|
||||
if dec_guide: # the guide one level down (argument 48), the encoder's of this level (56), the output (64)
|
||||
g_low, g_enc, out_ref = img(pre, ptr[48], wd // 2, ht // 2), img(pre, ptr[56], wd, ht), img(post, ptr[64], wd, ht)
|
||||
ok = True
|
||||
for blk in o.blocks:
|
||||
bx, by = (int(v) for v in blk.split(','))
|
||||
x0, y0 = 8 * bx - sx, 8 * by - sy
|
||||
inside = np.array([0 <= x0 + t % 8 < wd and 0 <= y0 + t // 8 < ht for t in range(64)])
|
||||
xin = np.zeros((64, 4 * c_up if c_up else c), np.uint8)
|
||||
xe = None
|
||||
if expand:
|
||||
lows = np.array([np.concatenate([up[p, mirror(y0 // 2 + n // 4, ht // 2), mirror(x0 // 2 + n % 4, wd // 2)]
|
||||
for p in range(c_lo // 32)]) for n in range(16)])
|
||||
e = mm(lows, w['we'], np.tile(w['be'], (16, 1)))
|
||||
un = np.argsort(inter(c))
|
||||
xe = np.zeros((64, c), H)
|
||||
for t in range(64):
|
||||
tx, ty = t % 8, t // 8
|
||||
sub = (ty % 2) * 2 + tx % 2
|
||||
sk = np.concatenate([skip[p, mirror(y0 + ty, ht), mirror(x0 + tx, wd)] for p in range(c // 32)])
|
||||
xe[t] = add(dq(e[(ty // 2) * 4 + tx // 2, sub * c:(sub + 1) * c]), (e4m3(sk)[un].astype(H) * s3).astype(H))
|
||||
for t in range(64 if not expand else 0): # tokens outside the level are the mirrored ones (no edge repeat), as in presets D and E
|
||||
X, Y = mirror(x0 + t % 8, wd), mirror(y0 + t // 8, ht)
|
||||
if not c_up:
|
||||
xin[t] = np.concatenate([up[p, Y, X] for p in range(c // 32)])
|
||||
continue
|
||||
for q, (dx, dy) in enumerate(((0, 0), (1, 0), (0, 1), (1, 1))):
|
||||
xin[t, q * c_up:(q + 1) * c_up] = np.concatenate([up[p, 2 * Y + dy, 2 * X + dx] for p in range(c_up // 32)])
|
||||
y, merged, head = window(w, scal, xin, xe)
|
||||
ya, ma = A(y), (A(merged) if c_m else None)
|
||||
bad_t = bad_m = bad_o = bad_h = 0
|
||||
out = np.zeros((64, 4), H)
|
||||
for t in np.nonzero(inside)[0]:
|
||||
X, Y = x0 + t % 8, y0 + t // 8
|
||||
bad_t += int(np.sum(np.concatenate([tok_ref[p, Y, X] for p in range(c // 32)]) != ya[t]))
|
||||
if dec_guide:
|
||||
# out = min(enc * (1 - a) + (upsampled * sigmoid(h[0:3]) * s0) * a, 65504), a = sigmoid(h[3])
|
||||
sg = sigmoid(head[t, :4])
|
||||
gate = np.array([(sg[0] * scal[0]).astype(H), (sg[1] * scal[0]).astype(H), (sg[2] * scal[0]).astype(H), H(0)])
|
||||
gt = (upsample(g_low, X, Y) * gate).astype(H)
|
||||
gt[3] = H(0)
|
||||
mix = add((g_enc[Y, X] * (H(1) - sg[3])).astype(H), (gt * sg[3]).astype(H))
|
||||
out[t] = np.minimum(mix, H(65504))
|
||||
bad_o += int(np.sum(out[t].view(np.uint16) != out_ref[Y, X].view(np.uint16)))
|
||||
if not guided:
|
||||
continue
|
||||
gate = (sigmoid(head[t, :3]) * scal[0]).astype(H)
|
||||
out[t, :3] = np.minimum((guide[Y, X, :3] * gate).astype(H), H(65504))
|
||||
bad_o += int(np.sum(out[t].view(np.uint16) != out_ref[Y, X].view(np.uint16)))
|
||||
for n in range(16):
|
||||
X, Y = x0 // 2 + n % 4, y0 // 2 + n // 4
|
||||
t0 = (2 * (n // 4)) * 8 + 2 * (n % 4)
|
||||
if not c_m or not (inside[t0] and x0 % 2 == 0 and y0 % 2 == 0):
|
||||
continue
|
||||
bad_m += int(np.sum(np.concatenate([mer_ref[p, Y, X] for p in range(c_m // 32)]) != ma[n]))
|
||||
qv = [(out[t0 + d] * H(0.25)).astype(H) for d in (0, 1, 8, 9)]
|
||||
bad_h += int(np.sum(add(add(qv[0], qv[1]), add(qv[2], qv[3])).view(np.uint16) != half_ref[Y, X].view(np.uint16)))
|
||||
print(f'block {blk}: {int(inside.sum())} tokens inside; differing tokens bytes {bad_t}, merged bytes {bad_m}, '
|
||||
f'guide values {bad_o}, half guide values {bad_h}')
|
||||
ok &= bad_t + bad_m + bad_o + bad_h == 0
|
||||
print('IDENTICAL' if ok else 'DIFFERENT')
|
||||
sys.exit(0 if ok else 1)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in new issue
Block a user