woutermans
8eb7a94533
- 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
91 lines
3.8 KiB
Python
91 lines
3.8 KiB
Python
#!/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()
|