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:
1 parent
065fe2c2a1
commit
8eb7a94533
31 files changed
+3547
-128
No files matched your search
@@ -0,0 +1,90 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Pack a preset F stage's weights, read from the user's nvngx_dlssd.dll (engine/rr/rr2_ref.dll_weights), for
|
||||
engine/rr/rr2_m.comp.
|
||||
|
||||
usage: rr2_pack.py STAGE NVNGX_DLSSD_DLL OUT.bin OUT.h
|
||||
|
||||
swin_pack.py's layout: transposed products as 16 x 16 tiles, every tensor's outputs in slot order (slot j holds the
|
||||
natural index inter(j)); a merge's or the expand's inputs are sub-positions of C channels each. OUT.h also defines
|
||||
the stage's shape for the shader (KIND, C, NH, NB, CM, CL, GUIDE).
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from enc1_ref import inter # noqa: E402
|
||||
from enc1_pack import a_tiles, b_tiles, cubic_table, act_table # noqa: E402
|
||||
from rr2_ref import STAGES, EXPAND, dll_weights # noqa: E402
|
||||
|
||||
|
||||
def main():
|
||||
stage, dll, out_bin, out_h = sys.argv[1:5]
|
||||
w = dll_weights(open(dll, 'rb').read(), stage)
|
||||
if stage in EXPAND:
|
||||
C, CL, NH, NB = EXPAND[stage]
|
||||
kind, CM, guide = 2, 0, 0
|
||||
else:
|
||||
C, c_up, CM, NH, NB = STAGES[stage]
|
||||
kind, CL = (4 if stage == 'enc0' else 5 if stage == 'dec0_1' else 3 if stage.startswith('dec') else 1 if c_up else 0), 0
|
||||
guide = 0 if kind >= 3 else 1
|
||||
I32, I64, IC = inter(32), inter(64), inter(C)
|
||||
e = {}
|
||||
if kind == 1:
|
||||
e['w1'], e['b1'] = ('A', w['w1'][:, IC]), ('H', w['b1'][IC])
|
||||
if kind == 4: # encoder 0's embedding branches
|
||||
for b in range(2):
|
||||
e[f'ewd{b}'] = ('A', w['ewd'][b][:, I32])
|
||||
e[f'ebd{b}'] = ('H', w['ebd'][b][I32])
|
||||
e[f'ewu{b}'] = ('A', w['ewu'][b][:, IC])
|
||||
e['ebu'] = ('H', w['ebu'][IC])
|
||||
if kind == 2:
|
||||
I4C = [C * (j // C) + IC[j % C] for j in range(4 * C)]
|
||||
e['we'], e['be'] = ('A', w['we'][:, I4C]), ('H', w['be'][I4C])
|
||||
for h in range(NH):
|
||||
e[f'wqk{h}'] = ('A', w['wqk'][h][:, I32])
|
||||
e[f'wv{h}'] = ('B', w['wv'][h][:, I32])
|
||||
e[f'wo{h}'] = ('A', w['wo'][h][:, IC])
|
||||
e[f'bias{h}'] = ('H', w['bias'][h][np.ix_(I64, I64)]) # [key slot][query slot]; symmetric
|
||||
e['bo'] = ('H', w['bo'][IC])
|
||||
for b in range(NB):
|
||||
e[f'wd{b}'] = ('A', w['wd'][b][:, I32])
|
||||
e[f'bd{b}'] = ('H', w['bd'][b][I32])
|
||||
e[f'wu{b}'] = ('A', w['wu'][b][:, IC])
|
||||
e['bu'] = ('H', w['bu'][IC])
|
||||
if CM:
|
||||
e['wm'], e['bm'] = ('A', w['wm'][:, inter(CM)]), ('H', w['bm'][inter(CM)])
|
||||
if 'wh' in w:
|
||||
e['wh'], e['bh'] = ('A', w['wh']), ('H', w['bh'])
|
||||
if kind == 5: # dec0's output tail: tanh of every binary16 value, correctly rounded, by its bits
|
||||
with np.errstate(all='ignore'):
|
||||
x = np.arange(65536, dtype=np.uint32).astype(np.uint16).view(np.float16).astype(np.float64)
|
||||
e['tanh'] = ('U', np.tanh(x).astype(np.float16).view(np.uint16))
|
||||
e['cubic'] = ('U', cubic_table())
|
||||
e['act'] = ('Y', act_table())
|
||||
blob, defines = bytearray(), []
|
||||
for name, (kd, a) in e.items():
|
||||
while len(blob) % 256:
|
||||
blob.append(0)
|
||||
off = len(blob)
|
||||
if kd == 'A':
|
||||
blob += a_tiles(a).tobytes()
|
||||
elif kd == 'B':
|
||||
blob += b_tiles(a).tobytes()
|
||||
elif kd == 'Y':
|
||||
blob += np.ascontiguousarray(a, np.uint8).tobytes()
|
||||
elif kd == 'U':
|
||||
blob += np.ascontiguousarray(a, np.uint16).tobytes()
|
||||
else:
|
||||
blob += np.ascontiguousarray(a, np.float16).tobytes()
|
||||
defines.append(f'#define W_{name.upper()} {off}u // {kd} {a.shape}')
|
||||
shape = [f'#define KIND {kind}', f'#define C {C}', f'#define NH {NH}', f'#define NB {NB}', f'#define CM {CM}',
|
||||
f'#define CL {CL}', f'#define GUIDE {guide}']
|
||||
open(out_bin, 'wb').write(bytes(blob))
|
||||
open(out_h, 'w').write(f'// generated by engine/rr/rr2_pack.py for {stage}\n' + '\n'.join(shape + defines) + '\n')
|
||||
print(f'{stage}: {len(blob)} bytes, {len(e)} entries')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in new issue
Block a user