From 8eb7a94533eb8dd7b61770e23c28a11d6742a96f Mon Sep 17 00:00:00 2001 From: woutermans Date: Tue, 13 Oct 2026 13:49:46 +0200 Subject: [PATCH] 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 --- engine/compile_rr2.py | 203 ++++++++ engine/replay_rr.cpp | 7 +- engine/rr/README.md | 164 +++++++ engine/rr/build_native_model.sh | 2 + engine/rr/dec0_m.comp | 34 +- engine/rr/enc0_m_decl.glsl | 16 + engine/rr/enc0_m_tail.glsl | 32 +- engine/rr/exposure_f.comp | 93 ++++ engine/rr/gen_rr2_enc0.py | 298 ++++++++++++ engine/rr/phase_dump.py | 9 +- engine/rr/ptx2rr.py | 188 +++++++- engine/rr/rr2_m.comp | 530 +++++++++++++++++++++ engine/rr/rr2_pack.py | 90 ++++ engine/rr/rr2_ref.py | 362 +++++++++++++++ engine/rr/rr_attn.glsl | 7 +- engine/rr/rr_capture.py | 12 +- engine/rr/rr_inputs.py | 41 ++ engine/rr/swin_m.comp | 40 +- engine/rr/test_layer.py | 30 +- engine/rr/zluda_replay.py | 30 +- engine/runtime_rr.inc | 30 +- engine/runtime_rr2.inc | 643 ++++++++++++++++++++++++++ kernels/rr/swin_native.h | 491 ++++++++++++++++++++ kernels/tools/extract_dlss_ptx.py | 22 +- kernels/tools/rr_swin_native_test.cpp | 121 +++++ packaging/d4r.ini | 6 +- scripts/d4r_config.py | 6 +- scripts/run_rr_engine_compare.sh | 8 +- tests/test_ray_reconstruction.py | 7 +- tools/d3d12_dlss_harness.cpp | 4 +- tools/d4r_nvngx_shim.cpp | 149 ++++-- 31 files changed, 3547 insertions(+), 128 deletions(-) create mode 100644 engine/compile_rr2.py create mode 100644 engine/rr/exposure_f.comp create mode 100644 engine/rr/gen_rr2_enc0.py create mode 100644 engine/rr/rr2_m.comp create mode 100644 engine/rr/rr2_pack.py create mode 100644 engine/rr/rr2_ref.py create mode 100644 engine/runtime_rr2.inc create mode 100644 kernels/rr/swin_native.h create mode 100644 kernels/tools/rr_swin_native_test.cpp diff --git a/engine/compile_rr2.py b/engine/compile_rr2.py new file mode 100644 index 0000000..a679969 --- /dev/null +++ b/engine/compile_rr2.py @@ -0,0 +1,203 @@ +#!/usr/bin/env python3 +"""Build the engine's Ray Reconstruction model directory for preset F (DLSS 4.5, the rr2_* kernels of +nvngx_dlssd.dll 310.9) from a capture of the denoiser. + +usage: compile_rr2.py CAPTURE_DIR OUT_DIR [--ptx-dir DIR] [--frame N] [--native NVNGX_DLSSD_DLL] + +CAPTURE_DIR is a bridge capture of a few preset F frames (engine/rr/rr_capture.py, taken with +D4R_CUDA_REPLAY_DUMP_FILTER=rr2 D4R_CUDA_LAUNCH_STATS_FILTER=rr2 and preset hint 6). Writes: + .spv one shader per kernel, generated from its PTX by engine/rr/ptx2rr.py + weights-.bin the stage's weight image as NGX uploads it (argument 24); the output filter's is the constant + table of its PTX module; enc0-aux.bin is encoder 0's second table (argument 104) + .args the argument block of the capture's frame, pointers and handles cleared: the runtime keeps its + constants and fills what depends on the frame or the geometry (engine/runtime_rr2.inc) + exposure_f0.spv, exposure_f1.spv engine/rr/exposure_f.comp + rr.txt the manifest RModel reads ("D4RR0002") +With --native, the stages engine/rr/rr2_m.comp implements (enc1 ... enc5, dec4 ... dec1, dec0's first kernel) are that +shader with weights read from the DLL itself (engine/rr/rr2_pack.py) instead of the translation: bit-identical to +NVIDIA's kernels on the reference captures and about three times faster. D4R_RR_MIX=1 builds the one-rounding variant. +Everything but the exposure shader is derived from NVIDIA's DLL: it stays on the user's machine and is never packaged. +""" +import argparse +import json +import os +import subprocess +import sys +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE / 'rr')) +from rr_capture import Capture # noqa: E402 +from test_layer import stage_params # noqa: E402 + +NETWORK = ['enc0', 'enc1', 'enc2', 'enc3', 'enc4', 'enc5'] + [f'dec{n}_{k}' for n in (4, 3, 2, 1, 0) for k in (0, 1)] +STAGES = NETWORK + ['post', 'down', 'pooled'] +KERNEL = {s: f'rr2_{s}_kernel' for s in NETWORK[:6]} +KERNEL.update({f'dec{n}_{k}': f'rr2_dec{n}_kernel_{k}' for n in range(5) for k in (0, 1)}) +KERNEL.update(post='rr2_post_kernel', down='rr2_downsample_kernel_static_hdr_no_transparency', + pooled='rr2_pooled_history2_kernel') +BLOCK = {'enc0': '32,1,1', 'enc1': '32,1,1', 'enc2': '32,1,4', 'enc3': '32,1,4', 'enc4': '32,1,8', 'enc5': '32,1,8', + 'post': '16,16,1', 'down': '16,16,1', 'pooled': '16,16,1'} +for n, b in ((4, '32,1,8'), (3, '32,1,4'), (2, '32,1,4'), (1, '32,1,2'), (0, '32,1,1')): + BLOCK[f'dec{n}_0'] = BLOCK[f'dec{n}_1'] = b +# pointer arguments that are neither the weights (24) nor the plane arena: the exposure values (16 bytes), encoder +# 0's second table, and the output passes' progress counters (the runtime leaves those pointers null) +BUFFERS = {'enc0': [104, 256], 'post': [16, 256, 264], 'down': [72, 168, 176]} +# surfaces whose image format is not the store's own guess. The two pooled histories are UNORM 10:10:10:2 arrays for +# NVIDIA; the reference this engine is checked against (ZLUDA) backs them with RGBA32F, and so does the engine. +NATIVE = ['enc1', 'enc2', 'enc3', 'enc4', 'enc5', 'dec4_0', 'dec4_1', 'dec3_0', 'dec3_1', 'dec2_0', 'dec2_1', 'dec1_0', + 'dec1_1', 'dec0_0'] +# --particles: the downsample kernel with the particle layer +KERNEL['downp'] = 'rr2_downsample_kernel_static_hdr' +BLOCK['downp'] = '16,16,1' +BUFFERS['downp'] = [72, 168, 176] +# with --native: the translation of the per-pixel code around the native block (engine/rr/gen_rr2_enc0.py) +HYBRID = {'enc0': [], 'dec0_1': ['--dec0_1']} +# windows per workgroup of the hybrid stages (gen_rr2_enc0.py --group; D4R_RR2_GROUP to experiment) +GROUP = int(os.environ.get('D4R_RR2_GROUP', '1')) +# encoder 0's gather as a pass of its own (D4R_RR2_SPLIT=0: one shader, the gather on one wave of the block's four) +SPLIT = os.environ.get('D4R_RR2_SPLIT', '1') != '0' +# The pooled histories as RGBA32F holding q / 1023, as ZLUDA backs them (16 bytes a pixel). D4R_RR2_POOLED=rgb10_a2 +# makes them real UNORM 10:10:10:2 images (0.1 ms faster at 4K), but the hardware's bilinear filter then rounds +# differently: identical at 3x and 2x scaling, 44 - 51 dB on temporal frames at 1.5x (capFQ, 2026-10-13) +POOLED = os.environ.get('D4R_RR2_POOLED', 'rgba32f') +SURFACES = {'dec0_1': {176: 'rgba8_snorm', 184: 'rgba8', 192: 'rgba8'}, 'post': {120: POOLED}, 'pooled': {8: POOLED}} + + +def find_module(ptx_dir, kernel): + for path in sorted(Path(ptx_dir).glob('*.ptx')): + if f'.entry {kernel}(' in path.read_text(errors='replace'): + return path + raise FileNotFoundError(f'no PTX module with the entry {kernel} in {ptx_dir}') + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument('capture') + ap.add_argument('output', type=Path) + ap.add_argument('--ptx-dir', default=os.path.expanduser('~/.cache/d4r-rr-corpus-f')) + ap.add_argument('--frame', type=int, default=2, help='which captured frame supplies the templates (1-based)') + ap.add_argument('--particles', metavar='CAPTURE_DIR', + help='a capture taken with DLSSD.ColorBeforeParticles: adds the downsample kernel NGX runs then (downp)') + ap.add_argument('--native', metavar='DLL', help='nvngx_dlssd.dll 310.9: build the native stages from its weights') + a = ap.parse_args() + cap = Capture(a.capture) + out = a.output + out.mkdir(parents=True, exist_ok=True) + capped = ['systemd-run', '--user', '--scope', '-q', '-p', 'MemoryMax=24G', '-p', 'MemorySwapMax=0'] + lines = ['D4RR0002', 'preset 6'] + pcap = Capture(a.particles) if a.particles else None + for stage in STAGES + (['downp'] if pcap else []): + kernel = KERNEL[stage] + if stage == 'downp': + cap = pcap + launches = [l for l in cap.launches if l.kernel == kernel] + per_frame = len(launches) // len([l for l in cap.launches if l.kernel == KERNEL['enc0']]) + launch = launches[(a.frame - 1) * per_frame] + network = stage in NETWORK + if a.native and stage in NATIVE: + work = out / 'native-work' / stage + work.mkdir(parents=True, exist_ok=True) + weights = f'weights-{stage}m.bin' + subprocess.run([sys.executable, str(HERE / 'rr' / 'rr2_pack.py'), stage, a.native, str(out / weights), + str(work / 'rr2_stage.h')], check=True, stdout=subprocess.DEVNULL) + defs = ['-DRR_EXACT_CONVERT'] + (['-DRR_MIX'] if os.environ.get('D4R_RR_MIX') == '1' else []) + subprocess.run(capped + ['glslangValidator', '-V', '--target-env', 'vulkan1.3', *defs, f'-I{work}', + f'-I{HERE / "rr"}', str(HERE / 'rr' / 'rr2_m.comp'), '-o', str(out / f'{stage}_m.spv')], + check=True, stdout=subprocess.DEVNULL) + subprocess.run([sys.executable, str(HERE / 'rr' / 'spv_exact_convert.py'), str(out / f'{stage}_m.spv')], + check=True, stdout=subprocess.DEVNULL) + (out / f'{stage}.args').write_bytes(stage_params(launch, 24)[:len(launch.args)]) + arena = launch.arena_alloc + unit = launch.allocs[arena][1] / 3152 + ptrs = [f'{p.offset}:{p.delta / unit:g}' for p in launch.pointers if p.alloc == arena and p.offset != 24] + lines.append(f'stage {stage} {stage}_m.spv 32 {weights} {stage}.args - - - {",".join(ptrs)}') + print(f'{stage}: native (engine/rr/rr2_m.comp)') + continue + generated = out / f'{stage}.generated.comp' + cmd = [sys.executable, str(HERE / 'rr' / 'ptx2rr.py'), str(find_module(a.ptx_dir, kernel)), kernel, str(generated), + '--block', BLOCK[stage], '--weights-arg', '24' if network else '9999'] + cmd += [f'--buffer={off}' for off in BUFFERS.get(stage, [])] + cmd += [f'--surface={off}={fmt}' for off, fmt in SURFACES.get(stage, {}).items()] + subprocess.run(cmd, check=True, stderr=subprocess.DEVNULL) + info = json.loads(generated.with_suffix('.json').read_text()) + weights = '-' + hybrid = a.native and stage in HYBRID + if hybrid: + work = out / 'native-work' / stage + work.mkdir(parents=True, exist_ok=True) + weights = f'weights-{stage}m.bin' + subprocess.run([sys.executable, str(HERE / 'rr' / 'rr2_pack.py'), stage, a.native, str(out / weights), + str(work / 'rr2_stage.h')], check=True, stdout=subprocess.DEVNULL) + defs = ['-DRR_EXACT_CONVERT'] + (['-DRR_MIX'] if os.environ.get('D4R_RR_MIX') == '1' else []) + # encoder 0 in two passes (gen_rr2_enc0.py --split): its gather as the stage enc0g, whose record's address + # is a word appended to the argument block + split = stage == 'enc0' and SPLIT and GROUP == 1 + how = ['--split', 'block', '--xf-word', str((len(launch.args) + 7) // 8 * 2)] if split else ['--group', str(GROUP)] + if split: + subprocess.run([sys.executable, str(HERE / 'rr' / 'gen_rr2_enc0.py'), str(find_module(a.ptx_dir, kernel)), + str(generated), str(work / 'enc0_g.comp'), '--split', 'gather', *how[2:]], + check=True, stdout=subprocess.DEVNULL) + subprocess.run(capped + ['glslangValidator', '-V', '--target-env', 'vulkan1.3', *defs, f'-I{work}', + f'-I{HERE / "rr"}', str(work / 'enc0_g.comp'), '-o', str(out / 'enc0g.spv')], + check=True, stdout=subprocess.DEVNULL) + subprocess.run([sys.executable, str(HERE / 'rr' / 'spv_exact_convert.py'), str(out / 'enc0g.spv')], + check=True, stdout=subprocess.DEVNULL) + record = int((work / 'enc0_g.record').read_text()) + subprocess.run([sys.executable, str(HERE / 'rr' / 'gen_rr2_enc0.py'), str(find_module(a.ptx_dir, kernel)), + str(generated), str(work / f'{stage}_m.comp'), *HYBRID[stage], *how], + check=True, stdout=subprocess.DEVNULL) + subprocess.run(capped + ['glslangValidator', '-V', '--target-env', 'vulkan1.3', *defs, f'-I{work}', + f'-I{HERE / "rr"}', str(work / f'{stage}_m.comp'), '-o', str(out / f'{stage}.spv')], + check=True, stdout=subprocess.DEVNULL) + subprocess.run([sys.executable, str(HERE / 'rr' / 'spv_exact_convert.py'), str(out / f'{stage}.spv')], + check=True, stdout=subprocess.DEVNULL) + else: + subprocess.run(capped + ['glslangValidator', '-V', '--target-env', 'vulkan1.3', str(generated), '-o', + str(out / f'{stage}.spv')], check=True, stdout=subprocess.DEVNULL) + if hybrid: + pass + elif network: + weights = f'weights-{stage}.bin' + (out / weights).write_bytes(launch.weights_at(24)) + elif generated.with_suffix('.globals.bin').exists(): + weights = f'weights-{stage}.bin' + (out / weights).write_bytes(generated.with_suffix('.globals.bin').read_bytes()) + if stage == 'enc0': + (out / 'enc0-aux.bin').write_bytes(launch.pre(launch.pointer(104).alloc)) + args = stage_params(launch, 24 if network else 9999)[:len(launch.args)] + if hybrid and split: + args = args + bytes(-len(args) % 8) + bytes(8) # the gather record's address (the runtime's) + (out / f'{stage}.args').write_bytes(args) + # the stage's arena regions, in units of 1/3152 of NGX's arena (W0 * H0 / 32 bytes: the regions scale with it) + arena = launch.arena_alloc + unit = launch.allocs[arena][1] / 3152 if arena is not None else 1 + ptrs = [f'{p.offset}:{p.delta / unit:g}' for p in launch.pointers + if p.alloc == arena and p.offset not in BUFFERS.get(stage, []) and not (network and p.offset == 24)] + if not network: + ptrs = [] # the output passes' regions are the runtime's own + tex = ','.join(map(str, info['textures'])) or '-' + surf = ','.join(f'{off}:{fmt}' for off, fmt in info['surfaces']) or '-' + buf = ','.join(map(str, info['buffers'])) or '-' + # the output passes have no 32-lane matrix products: 64-wide waves (their warp-level shuffles address the + # lane's own half, engine/rr/ptx2rr.py) + wave = 32 if network else 64 + lines.append(f'stage {stage} {stage}.spv {wave} {weights} {stage}.args {tex} {surf} {buf} {",".join(ptrs) or "-"}' + + (f' x{GROUP}' if hybrid and GROUP > 1 else '')) # windows per workgroup (engine/rr/gen_rr2_enc0.py) + if hybrid and split: # the same block, bindings and regions; "rN": N words per window record + lines.append(f'stage enc0g enc0g.spv 32 - {stage}.args {tex} {surf} {buf} {",".join(ptrs) or "-"} r{record}') + print(f'{stage}: {kernel}, {len(launch.args)} argument bytes' + (' (native block)' if hybrid else '')) + for p in (0, 1): + subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', f'-DPASS={p}', + str(HERE / 'rr' / 'exposure_f.comp'), '-o', str(out / f'exposure_f{p}.spv')], + check=True, stdout=subprocess.DEVNULL) + # the driver folds float(float16_t(x)) away unless the conversions are marked (engine/rr/spv_exact_convert.py): + # the kernel truncates the exposure to binary16 + subprocess.run([sys.executable, str(HERE / 'rr' / 'spv_exact_convert.py'), str(out / f'exposure_f{p}.spv')], + check=True, stdout=subprocess.DEVNULL) + (out / 'rr.txt').write_text('\n'.join(lines) + '\n') + print(f'Ray Reconstruction preset F model in {out}') + + +if __name__ == '__main__': + main() diff --git a/engine/replay_rr.cpp b/engine/replay_rr.cpp index 23999ab..cd2d377 100644 --- a/engine/replay_rr.cpp +++ b/engine/replay_rr.cpp @@ -32,8 +32,9 @@ std::vector read(const std::string& file) { } uint32_t texelBytes(VkFormat f) { switch (f) { + case VK_FORMAT_R32G32B32A32_SFLOAT: return 16; case VK_FORMAT_R16G16B16A16_SFLOAT: return 8; - case VK_FORMAT_R16G16_SFLOAT: case VK_FORMAT_R32_SFLOAT: case VK_FORMAT_R8G8B8A8_UNORM: case VK_FORMAT_R8G8B8A8_SNORM: return 4; + case VK_FORMAT_A2B10G10R10_UNORM_PACK32: case VK_FORMAT_R16G16_SFLOAT: case VK_FORMAT_R32_SFLOAT: case VK_FORMAT_R8G8B8A8_UNORM: case VK_FORMAT_R8G8B8A8_SNORM: return 4; case VK_FORMAT_R16_SFLOAT: return 2; default: throw std::runtime_error("texelBytes: format"); } @@ -163,7 +164,8 @@ int main(int argc, char** argv) try { const VkFormat optionalFormats[] = {VK_FORMAT_R32_SFLOAT, VK_FORMAT_R16_SFLOAT, VK_FORMAT_R16_SFLOAT, VK_FORMAT_R16G16B16A16_SFLOAT}; std::vector optional; for (VkFormat f : optionalFormats) optional.emplace_back(c, f, rw, rh); - HostBuffer staging(c, VkDeviceSize(ow) * oh * 8); + // large enough for every debug image: preset F's intermediate history is 1.5 times the output size on both axes + HostBuffer staging(c, VkDeviceSize(ow) * oh * 8 * 4); // GPU time of every frame: REngine's four stamps (start, after exposure, after the network, after the output) VkQueryPoolCreateInfo qp{VK_STRUCTURE_TYPE_QUERY_POOL_CREATE_INFO}; qp.queryType = VK_QUERY_TYPE_TIMESTAMP; qp.queryCount = 4; @@ -179,6 +181,7 @@ int main(int argc, char** argv) try { for (float& v : s.worldToView) frames >> v; for (float& v : s.viewToClip) frames >> v; s.reset = reset != 0; s.depthInverted = (inverted & 1) != 0; s.hardwareDepth = (inverted & 2) != 0; // bit 1: DLSS.Use.HW.Depth + s.roughnessPacked = (inverted & 4) != 0; // bit 2: DLSS.Roughness.Mode (preset F inputs) for (size_t i = 0; i < inputs.size(); ++i) { auto data = read(inputDir + "/" + std::to_string(f) + "-" + names[i] + ".bin"); if (data.size() != size_t(rw) * rh * texelBytes(formats[i])) throw std::runtime_error(std::string("input size: ") + names[i]); diff --git a/engine/rr/README.md b/engine/rr/README.md index 36efe29..5ef023c 100644 --- a/engine/rr/README.md +++ b/engine/rr/README.md @@ -172,6 +172,25 @@ single-use multiply-subtract pairs into fma, merging the two constant multiplica `precise` on the inlined exp2, NoContraction on its conversions. The differing values come in spatial clusters, one colour channel at a time, always one ulp, mostly just above 0.5 in surface 224. +## HIP against Vulkan, and the fused rounding (2026-10-12) + +Encoder 1 ported line for line to a standalone HIP kernel (`kernels/rr/swin_native.h`, runner +`kernels/tools/rr_swin_native_test.cpp`; capD1080 frame 2, 5 launches per round): 835 us against the Vulkan stage's +770-788 us, the arena byte-identical. A HIP port of the network is not worth making. What the port did show: + +- The FP8 product alone costs 0.61 us per product per wave on this grid, 2.07 us when each is followed by the + f32 -> binary16 -> f32 rounding of its eight accumulator registers. Products and that rounding are about half of + a stage on either backend. +- Keeping the accumulator as binary16 and taking f16(P + acc) with `v_fma_mix` (as `kernels/m` does) halves the + rounding's instructions: HIP 780 us. RADV emits the same instruction for `float16_t(p + float(acc))`, so the + Vulkan stages have it as a build option: `D4R_RR_MIX=1 build_native_model.sh ...` (`RR_MIX`, the shaders' K32). + +`RR_MIX` is one rounding of the exact sum where the reference has two, so it is not bit-exact by construction. +Preset D: every network stage still identical to the capture on all 21 frames of capDT12 / capD4K / capD1080 / capDX +/ capDXcp / capDXhw (the same figures as the default build), 4.22 ms at 1280x720 -> 4K and 8.3 ms at 1920x1080 -> 4K +(default 4.45 / 8.7). Preset E (capE3): frames 1 and 3 have a value of encoder 0 one step off (127 dB), which the +later stages amplify to 41 dB at the output; frame 2 is identical. The default build stays the exact one. + ## Cyberpunk 2077: more inputs and settings (2026-10-12) Found on harness captures (`D4R_HARNESS_RR_EXTRA=particles,sss`, `D4R_HARNESS_RR_HW_DEPTH=1`, `D4R_RR_PRESET=N`): @@ -189,3 +208,148 @@ byte-identical to the CUDA path through the shim on the frames the harness can c with RE Requiem's set the network is identical on every frame (output 107-109 dB). At 1280x720 and 1920x1080 -> 4K without these inputs the output still falls to 43-53 dB after the first frame (the output stage's one-ulp values); a correctly rounded log2 there matches the capture worse than the GPU instruction does. + +## The output stage's one-step differences: fused multiply-convert (2026-10-13) + +The rare one-binary16-step values of the translated output stage (above: ~0.004% of surfaces 224 / 240, "cause not +found") were the driver fusing an f32 product (or sum) with the conversion that follows it into `v_fma_mix`, which +rounds the exact result once. The kernel rounds twice, to f32 and then to binary16, and the two differ when the f32 +result is exactly a binary16 midpoint, as `log2(x) * c1 * c2` of the output stage's tone curve sometimes is. +`ptx2rr.py` now ORs an opaque zero into the f32 bits before every `cvt.rn.f16.f32` (and into the binary16 bits +after, against the fold of `float(float16_t(x))`). With it presets D and E are byte-identical to the CUDA path on +every frame of every capture (capDT12, capD4K, capD1080, capDX, capDXcp, capDXhw, capE3; history and output), at ++0.05 ms: 4.27 / 4.45 ms at 1280x720 -> 4K. Found with `build/rr`-style bisection: probe every register of a block +through ZLUDA (debug stores added to the PTX, `zluda_replay.py`) and through the translated shader +(`ptx2rr.py --probe`) and list the first that differ. + +Other reference behaviours the translation now follows, all found the same way on preset F's kernels: + +- formatted stores to normalized surfaces round half away from zero in f32 (ZLUDA's `surface_formatted_store_bits`); + UNORM 10:10:10:2 arrays are RGBA32F images holding the quantised values `q / 1023` (the correctly rounded quotient, + from a table: the driver divides by multiplying with a reciprocal); +- `tanh.approx.f32` is OCML's `tanhF` with LLVM's AMDGPU lowering of `exp` (`lowerFEXP`, fast-FMA path): the same + operations, the fused ones through f64 (the driver's f32 `fma` is a product and a sum whenever the product also + exists on its own); `tanh.approx.f16` is the correctly rounded value (a table); +- the driver's GLSL `fma` on binary16 is a true fused operation, as the reference's is: no emulation there; +- texture fetches with a binary16 result keep that rounding. + +## Preset F: DLSS 4.5 Ray Reconstruction (2026-10-13) + +`nvngx_dlssd.dll` 310.9 has a third model with its own kernels, `rr2_*`, and runs it for every preset hint but 4 (D) +and 5 (E): none, 6, Cyberpunk 2077's 0x80000002. The engine runs it from a `D4RR0002` model: + +```sh +python3 kernels/tools/extract_dlss_ptx.py NVNGX_DLSSD_DLL ~/.cache/d4r-rr-corpus-f # its PTX (zstd in 310.9) +python3 engine/compile_rr2.py CAPTURE OUT_DIR --native NVNGX_DLSSD_DLL [--particles CAPTURE_WITH_PARTICLES] +``` + +The runtime is `engine/runtime_rr2.inc` (behind `RModel` / `REngine`); the shim uses `engine\rr-f` (or +`D4R_ENGINE_MODEL_RR_F_DIR`) and, when that model exists, NGX 310.9's hint mapping. + +A frame, as NGX launches it: + +| pass | NVIDIA kernel | engine | +|---|---|---| +| exposure | `rr2_histogram_auto_exposure_basic_kernel` | `exposure_f.comp` (two passes) | +| enc0 | `rr2_enc0_kernel` | translated feature gather and guide tail around the native block (`gen_rr2_enc0.py`) | +| enc1 - enc5 | `rr2_enc1_kernel` ... | `rr2_m.comp` | +| dec4 - dec1 | `rr2_decN_kernel_0`, `rr2_decN_kernel_1` | `rr2_m.comp` | +| dec0 | `rr2_dec0_kernel_0`, `rr2_dec0_kernel_1` | `rr2_m.comp`; the second: native block, translated output tail | +| output | `rr2_post_kernel`, `rr2_downsample_kernel_static_hdr[_no_transparency]` per tile, `rr2_pooled_history2_kernel` | translated (`ptx2rr.py`) | + +The network (`rr2_ref.py` is its exact numpy reference and documents the block): D/E's levels, window grids and +arithmetic (FP8 operands, binary16 sums rounded per 32-deep step, the cubic softmax, the MLP activation), with a +different block. Each kernel is one block on 8x8-token windows that alternate their shift per frame: heads that +share one projection for queries and keys (so the scores are symmetric; bias `(T + T^T) sqrt(2)` of a 64 x 64 +relative position table), all reading the block's input and summing their out-projections; MLP branches likewise; +both sums enter the residual times a binary16 scalar from the argument block. Encoders merge 2x2 for the next level, +a decoder's first kernel expands the level below and adds the skip times a third scalar. Next to the tokens every +level carries a three-channel guide image: encoders gate it (`sigmoid(head) * s0`) and average it 2x2 for the next +level, decoders upsample it bilinearly, gate it and blend it with the encoder's. All stages exchange data through +the arena (3152 units of W0 * H0 / 32 bytes; NGX's region offsets, which scale with the render size), tokens as FP8 +planes of 32 channels, guides as four binary16 per token. + +Weights: the DLL holds the model as plain binary16 tensors in natural `[out][in]` order, kernel after kernel +(`rr2_ref.dll_weights`; found by the sixteen kernels' relative position tables, 2 2 4 4 8 8 8 8 4 4 4 4 2 2 2 2 +heads). `rr2_pack.py` packs them for `rr2_m.comp`; nothing of the network comes from the capture. The capture still +supplies every kernel's argument block (the runtime keeps their constants and fills the geometry, the frame's +values and the pointers) and encoder 0's second table. + +The output stage filters an intermediate image of 1.5 times the output size in tiles (columns of at most 720, rows +of at most 1620 intermediate pixels, 4 pixels of overlap inside), reduces each tile to the output, and keeps two +histories for the next frame: the intermediate image, and a half-size copy whose quarter-size reduction +(`rr2_pooled_history2_kernel`) encoder 0 also reads (UNORM 10:10:10:2 arrays, see above). NGX runs filter and +reduction concurrently with progress counters; the kernels skip that when the counter pointers are null, which is how +the engine runs them, in order. They have no matrix products and run with 64-wide waves (their warp-level shuffles +address the lane's own half). + +Options (harness captures capFX*): encoder 0 byte 466 packed roughness, 568 inverted depth, 569 / texture 272 +`DLSSD.SpecularHitDistance`, 571 / texture 336 `DLSS.DisocclusionMask`, texture 288 the subsurface guide, 1028 +`DLSS.Use.HW.Depth`; `DLSSD.ColorBeforeParticles` selects the other reduction kernel (model stage `downp`, texture +208, an output-size history 256 <-> 264, flag bytes 196-198). Encoder 0's surfaces 376 / 384 are the output filter's +textures 64 / 72. Pre-exposure does not reach these kernels. + +Checked (`d4r-rr` on `rr_inputs.py` inputs; every stage surface and the presented image against the CUDA path): +byte-identical on every frame of capF (640x360 -> 1080p, 3 frames), capF900 (800x450 -> 1600x900), capF4K +(1280x720 -> 4K, 16 tiles) and of the option captures (motion, a real camera, inverted and hardware depth, packed +roughness, hit distance, disocclusion mask, subsurface guide, particles); through the shim (D3D12 harness under +Proton) the engine's frames are byte-identical to the CUDA backend's too (`D4R_RR_PRESET=6 D4R_RR_DLL=<310.9 DLL> +scripts/run_rr_engine_compare.sh`; that script still compares frame N with frame N, and the harness presents the +CUDA result one frame late in the jitter scene). + +Speed, RX 9070 XT, 1280x720 -> 3840x2160: 6.2 / 6.6 ms per frame (exposure 0.03, network 4.1 - 4.2, output 2.1 - +2.4), from 14.5 ms with every kernel translated; preset D is 4.3 / 4.5. Per stage (us): enc0 1076 (its block 521), +dec0_1 775, dec0_0 630, enc1 340, dec1_0 316, dec1_1 269, the rest under 180; filter and reduction about 130 per +tile, 16 tiles. Not done: native output passes and feature gather. + +### Where the translated parts' time went (2026-10-13, later) + +Now about 5.8 / 6.2 ms (5.7 / 6.1 with the 10-bit histories that were later withdrawn), still byte-identical on every capture and through the +shim. Profiling the translated code (cutting it at points behind an opaque condition, `D4R_ENGINE_RR_SKIP=post,down, +pooled`) showed its arithmetic is spread evenly - no hot spot a rewrite would remove - and that the losses were in +how the passes were arranged: + +- **Encoder 0 in two passes** (-0.2 ms). One shader ran NVIDIA's 32-lane gather on one of the block's four waves + while the other three waited: 0.40 ms for the gather, 0.24 ms as a pass of its own in 32-thread groups. + `gen_rr2_enc0.py --split gather | block` writes both; the gather pass (`stage enc0g`, the same bindings and + parameter block as enc0) stores each window's record in the arena after the tile slots: the first product's A + fragments (1024 words) and, per lane, the values the tail reads from the gather (six binary16 on the reference + build; the generator finds them and recomputes the tail's index arithmetic). The record's address is a word + compile_rr2.py appends to enc0's argument block. Cost: 76 bytes per render pixel of arena. + `D4R_RR2_SPLIT=0` builds the single shader. +- **Tile passes in overlapping groups** (-0.15 ms). The 16 tiles' post and down passes ran one after another, each + behind a barrier, and the gaps between 32 short dispatches added up. Each tile now has a slot (its two regions; + eight slots, 58 MB at 4K) and a group's post passes run without barriers between them, as do its down passes + together with the next group's post passes on the other half of the slots - what NGX does with two slots and its + progress counters. `D4R_ENGINE_RR_SKIP=serial` restores the old order. +- **Pooled histories as real 10:10:10:2 images: withdrawn.** 4 bytes a pixel instead of RGBA32F's 16 saved + 0.1 ms and was identical at 3x and 2x scaling, but the hardware's bilinear filter rounds a UNORM image differently + from a float one: at 1.5x (capFQ) temporal frames came out at 44 - 51 dB. The default is RGBA32F again; + `D4R_RR2_POOLED=rgb10_a2` builds the other form. + +Tried without gain: 32-bit arena accesses for the translated code's paired 16-bit ones (0.01 ms), dropping the +subgroup shuffles or the depth search (they cost nothing measurable). What is left is NVIDIA's arithmetic: gather +0.24 - 0.33 ms, decoder 0's output tail 0.18, post 1.1 - 1.3, down 0.7. + +### Other scaling ratios, and what a game-like configuration showed (2026-10-13, later) + +Asked whether instability seen in RE Requiem came from this implementation, the engine was compared with the CUDA +path through the shim in configurations the earlier captures did not cover (`D4R_RR_PRESET=6 D4R_RR_DLL=<310.9> +scripts/run_rr_engine_compare.sh`, frames aligned by one): + +- **Render widths that are not a multiple of 32 were wrong** (1114x626 -> 1080p: 15 dB): the post pass's argument 8 + is the row stride of the network's output, the padded width W0, not the render width. Fixed; identical now. +- **Below 4/3 scaling NGX runs another output filter**, `rr2_post_kernel_3_2` (1456x820 -> 1080p and larger render + sizes, DLAA included; 1440x810 and smaller use `rr2_post_kernel`). The model has no such stage, so the shim hands + those features to the CUDA backend. `rr2_histogram_auto_exposure_histogram_kernel` was never launched in any run. +- **DLSS.Roughness.Mode 1 with a roughness texture bound**: NGX follows the mode, the shim followed the texture. + Fixed for preset F (RE Requiem binds none, so it was not affected). +- 1.5x, 1.72x, 2x and 4/3 scaling, and RE Requiem's inputs together (inverted and hardware depth, hit distance, + mask, packed roughness, scaled motion vectors, a moving scene): identical. +- **Open: the camera matrices are not NGX's bit for bit.** NGX computes them in single precision: argument 644 is + the f32 product view x clip, 836 its inverse by cofactors times 1 / det (glm::inverse's operation order; both + reproduced bit for bit in build/rr/f/matfit.py), 900 is an inverse of the pixel matrix (it carries that inverse's + residues, and 0.9999999 for 1 at 1114x626) and 964 was not reproduced. The engine computes them in double and + rounds, a last-bit difference in a few elements. With a non-identity camera, hit distance and mask bound and a + moving scene that showed as 40 - 50 dB on later frames in a synthetic scene whose affected feature is rounding + noise around zero in both paths; no other configuration was affected. diff --git a/engine/rr/build_native_model.sh b/engine/rr/build_native_model.sh index ecd4da5..0250ec9 100755 --- a/engine/rr/build_native_model.sh +++ b/engine/rr/build_native_model.sh @@ -42,6 +42,8 @@ PRESET="$(awk '$1 == "preset" { print $2 }' "$OUT/rr.txt")" DEFS=(-DRR_EXACT_CONVERT); [[ "$PRESET" == 4 ]] && DEFS+=(-DRR_PRESET_D) # D4R_RR_RELAXED=1: one rounding where NVIDIA has two (engine/rr/rr_exact.glsl): faster, not bit-exact [[ -n "${D4R_RR_RELAXED:-}" ]] && DEFS+=("-DRR_RELAXED=$D4R_RR_RELAXED") +# D4R_RR_MIX=1: the 32-deep steps round once through v_fma_mix (the shaders' K32): ~5% faster, not provably bit-exact +[[ "${D4R_RR_MIX:-0}" == 1 ]] && DEFS+=(-DRR_MIX) glsl() { # glsl OUT.spv SOURCE INCLUDE_DIR... local out="$1" src="$2"; shift 2 local dir inc=(); for dir in "$@"; do inc+=("-I$dir"); done diff --git a/engine/rr/dec0_m.comp b/engine/rr/dec0_m.comp index 85e0472..603d741 100644 --- a/engine/rr/dec0_m.comp +++ b/engine/rr/dec0_m.comp @@ -48,7 +48,23 @@ layout(set = 0, binding = 7, rgba8_snorm) uniform writeonly image2D S144; #define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM) #define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]); #define Q8(dst, src) Q8PAIRS(dst, src) +// One 32-deep step onto an accumulator that holds binary16 values. The reference (and the default here) adds the +// accumulator inside the f32 products and rounds the result to binary16. RR_MIX (build_native_model.sh with +// D4R_RR_MIX=1) keeps the accumulator as binary16, sums the products from zero and takes f16(P + acc), which the +// driver compiles to one v_fma_mix per element in place of two conversions: about 5% faster, but a single rounding +// where the reference has two, so a value can come out one binary16 step away (engine/rr/README.md). +// RAC is the accumulator type, RV(x) a binary16 value as its element; K32H is the step on a binary16 matrix. +#ifdef RR_MIX +#define RAC HAC +#define RV(x) (x) +#define K32(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(0.0))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_] + float((acc)[i_])); } +#define K32H(acc, a0, b0, a1, b1) K32(acc, a0, b0, a1, b1) +#else +#define RAC ACC +#define RV(x) float(x) #define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) } +#define K32H(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(acc))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_]); } +#endif const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0; // shared memory, aliased by phase: expand and prologue [LOW | EX], then the attention [K | V^T] @@ -111,8 +127,8 @@ void main() [[unroll]] for (uint kt = 0u; kt < 4u; ++kt) coopMatLoad(lb[kt], LOW, n * 256u + 4u * kt, 16u, CM); [[unroll]] for (uint mm = 0u; mm < 2u; ++mm) { const uint m = 2u * wv + mm; - ACC a; - [[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_BE / 2u + 16u * m + 8u * hl + uint(i)]); + RAC a; + [[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_BE / 2u + 16u * m + 8u * hl + uint(i)]); K32(a, we[mm][0], lb[0], we[mm][1], lb[1]) K32(a, we[mm][2], lb[2], we[mm][3], lb[3]) B8 q; @@ -167,13 +183,13 @@ void main() barrier(); // ---------------------------------------------------------------- the residual: window tokens + table 0 - ACC R[2]; + RAC R[2]; [[unroll]] for (uint m = 0u; m < 2u; ++m) { HAC x; coopMatLoad(x, XR, wv * 512u + 16u * m, 32u, CM); [[unroll]] for (int i = 0; i < 8; ++i) { precise float16_t c = float16_t(x[i]) + w16.w[W_TABLE0 / 2u + 16u * m + 8u * hl + uint(i)]; - R[m][i] = float(c); + R[m][i] = RV(c); } } @@ -190,12 +206,12 @@ void main() LDA(wk[kt], W_WK0 + HEAD * hd, 32u, c, kt); LDB(wvb[kt], W_WV0 + HEAD * hd, 32u, kt, c); } - ACC a = ACC(0.0); + RAC a = RAC(RV(float16_t(0.0))); K32(a, wk[0], xb[0], wk[1], xb[1]) A8 ka; Q8(ka, a) coopMatStore(ka, POOL, n * 128u + 4u * c, 8u, RM); - ACC v = ACC(0.0); + RAC v = RAC(RV(float16_t(0.0))); K32(v, xa[0], wvb[0], xa[1], wvb[1]) A8 va; Q8(va, v) @@ -241,15 +257,15 @@ void main() precise float16_t n = rv[m][i] * g; Nb[m][i] = q8h(n); precise float16_t a = rv[m][i] + w16.w[W_TABLE1 / 2u + k]; - R[m][i] = float(a); + R[m][i] = RV(a); } } for (uint br = 0u; br < 4u; ++br) { const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br; B8 Hb[2]; [[unroll]] for (uint c = 0u; c < 2u; ++c) { - ACC h; - [[unroll]] for (int i = 0; i < 8; ++i) h[i] = float(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]); + RAC h; + [[unroll]] for (int i = 0; i < 8; ++i) h[i] = RV(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]); A8 a0, a1; LDA(a0, w1, 32u, c, 0u); LDA(a1, w1, 32u, c, 1u); K32(h, a0, Nb[0], a1, Nb[1]) diff --git a/engine/rr/enc0_m_decl.glsl b/engine/rr/enc0_m_decl.glsl index b338440..48cf822 100644 --- a/engine/rr/enc0_m_decl.glsl +++ b/engine/rr/enc0_m_decl.glsl @@ -22,7 +22,23 @@ layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena; #define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM) #define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]); #define Q8(dst, src) Q8PAIRS(dst, src) +// One 32-deep step onto an accumulator that holds binary16 values. The reference (and the default here) adds the +// accumulator inside the f32 products and rounds the result to binary16. RR_MIX (build_native_model.sh with +// D4R_RR_MIX=1) keeps the accumulator as binary16, sums the products from zero and takes f16(P + acc), which the +// driver compiles to one v_fma_mix per element in place of two conversions: about 5% faster, but a single rounding +// where the reference has two, so a value can come out one binary16 step away (engine/rr/README.md). +// RAC is the accumulator type, RV(x) a binary16 value as its element; K32H is the step on a binary16 matrix. +#ifdef RR_MIX +#define RAC HAC +#define RV(x) (x) +#define K32(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(0.0))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_] + float((acc)[i_])); } +#define K32H(acc, a0, b0, a1, b1) K32(acc, a0, b0, a1, b1) +#else +#define RAC ACC +#define RV(x) float(x) #define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) } +#define K32H(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(acc))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_]); } +#endif const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0; shared uint KVS[2 * 768]; // the current head's K, then V^T; before the attention the key pixels' features #define XF KVS // the key pixels' 32 FP8 features (plane 0 | plane 1), slot-major diff --git a/engine/rr/enc0_m_tail.glsl b/engine/rr/enc0_m_tail.glsl index 043f9a3..47d0166 100644 --- a/engine/rr/enc0_m_tail.glsl +++ b/engine/rr/enc0_m_tail.glsl @@ -20,23 +20,23 @@ // the embedding x = W2 q8(relu(W1 f + b1)) + b2 for token tiles n = wv (this wave's queries: x is also their // residual) and n = wv + 4 (the border keys, waves 0 and 1), then each key's norm in registers (dec0_ref.norm_plane // on the accumulator layout, as dec0_m.comp's MLP norm) to XT - ACC R[2]; + RAC R[2]; for (uint n = wv; n < 6u; n += 4u) { B8 fb0, fb1, hb_[2]; coopMatLoad(fb0, XF, n * 128u, 8u, CM); coopMatLoad(fb1, XF, n * 128u + 4u, 8u, CM); [[unroll]] for (uint m = 0u; m < 2u; ++m) { - ACC a; - [[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_B1 / 2u + 16u * m + 8u * hl + uint(i)]); + RAC a; + [[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_B1 / 2u + 16u * m + 8u * hl + uint(i)]); A8 a0, a1; LDA(a0, W_W1, 32u, m, 0u); LDA(a1, W_W1, 32u, m, 1u); K32(a, a0, fb0, a1, fb1) - [[unroll]] for (int i = 0; i < 8; ++i) hb_[m][i] = q8s(max(a[i], 0.0)); + [[unroll]] for (int i = 0; i < 8; ++i) hb_[m][i] = q8s(max(float(a[i]), 0.0)); } float16_t xv[2][8]; [[unroll]] for (uint m = 0u; m < 2u; ++m) { - ACC a; - [[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_B2 / 2u + 16u * m + 8u * hl + uint(i)]); + RAC a; + [[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_B2 / 2u + 16u * m + 8u * hl + uint(i)]); A8 a0, a1; LDA(a0, W_W2, 32u, m, 0u); LDA(a1, W_W2, 32u, m, 1u); K32(a, a0, hb_[0], a1, hb_[1]) @@ -44,7 +44,7 @@ if (n == wv) { // the residual: x + table 0 [[unroll]] for (int i = 0; i < 8; ++i) { precise float16_t c = xv[m][i] + w16.w[W_TABLE0 / 2u + 16u * m + 8u * hl + uint(i)]; - R[m][i] = float(c); + R[m][i] = RV(c); } } } @@ -90,12 +90,12 @@ LDA(wk[kt], W_WK0 + HEAD * hd, 32u, c, kt); LDB(wvb[kt], W_WV0 + HEAD * hd, 32u, kt, c); } - ACC a = ACC(0.0); + RAC a = RAC(RV(float16_t(0.0))); K32(a, wk[0], xb[0], wk[1], xb[1]) A8 ka; Q8(ka, a) coopMatStore(ka, KVS, n * 128u + 4u * c, 8u, RM); - ACC vv = ACC(0.0); + RAC vv = RAC(RV(float16_t(0.0))); K32(vv, xa[0], wvb[0], xa[1], wvb[1]) A8 va; Q8(va, vv) @@ -139,15 +139,15 @@ precise float16_t nn = rv[m][i] * g; Nb[m][i] = q8h(nn); precise float16_t a = rv[m][i] + w16.w[W_TABLE1 / 2u + k]; - R[m][i] = float(a); + R[m][i] = RV(a); } } for (uint br = 0u; br < 4u; ++br) { const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br; B8 Hb[2]; [[unroll]] for (uint c = 0u; c < 2u; ++c) { - ACC h; - [[unroll]] for (int i = 0; i < 8; ++i) h[i] = float(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]); + RAC h; + [[unroll]] for (int i = 0; i < 8; ++i) h[i] = RV(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]); A8 a0, a1; LDA(a0, w1, 32u, c, 0u); LDA(a1, w1, 32u, c, 1u); K32(h, a0, Nb[0], a1, Nb[1]) @@ -180,8 +180,8 @@ } } { // the 32 -> 16 head and encoder 0's filter of the noise grid (two f16 fma chains, taps even / odd) - ACC hv; - [[unroll]] for (int i = 0; i < 8; ++i) hv[i] = float(w16.w[W_BH / 2u + 8u * hl + uint(i)]); + RAC hv; + [[unroll]] for (int i = 0; i < 8; ++i) hv[i] = RV(w16.w[W_BH / 2u + 8u * hl + uint(i)]); A8 a0, a1; LDA(a0, W_WH, 32u, 0u, 0u); LDA(a1, W_WH, 32u, 0u, 1u); K32(hv, a0, Tb[0], a1, Tb[1]) @@ -228,8 +228,8 @@ const uvec2 mxy = mergeXY(col); const int X = x0 / 2 + int(mxy.x), Y = y0 / 2 + int(mxy.y), hx = ex / 2, hy = ey / 2; const uint m = wv; - ACC a; - [[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_BM / 2u + 16u * m + 8u * hl + uint(i)]); + RAC a; + [[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_BM / 2u + 16u * m + 8u * hl + uint(i)]); [[unroll]] for (uint ks = 0u; ks < 4u; ++ks) { B8 g0, g1; A8 w0, w1; coopMatLoad(g0, MG, 8u * ks, 32u, CM); diff --git a/engine/rr/exposure_f.comp b/engine/rr/exposure_f.comp new file mode 100644 index 0000000..fb4e308 --- /dev/null +++ b/engine/rr/exposure_f.comp @@ -0,0 +1,93 @@ +#version 460 +// Ray Reconstruction preset F (rr2_*): the auto exposure of rr2_histogram_auto_exposure_basic_kernel, in two passes. +// +// NVIDIA's kernel sums a clamped luma per 16x16 block (a pairwise tree over the block's threads), converts each block +// sum to fixed point (x 2^32 / 10000, round to nearest) and adds them with a 64-bit atomic; the last block to finish +// turns the total into the exposure. The integer total does not depend on the order of the blocks, so PASS 0 writes +// every block's fixed-point sum and PASS 1 (one workgroup) adds them and does what the last block does: +// exposure = 0.23547088 / max(mean luma, 0.005) -> the 1x1 exposure image of this frame +// values = { e, p, 1 / e, clamp(log2(e / p), -6, 6) / 6 } -> the 16-byte buffer encoder 0 and the output passes read +// with e the exposure truncated to binary16 and p the previous frame's (1 where either is zero). +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_control_flow_attributes : require +#ifndef PASS +#define PASS 0 +#endif +layout(set = 0, binding = 0, std430) readonly buffer Params { + ivec2 last; // render size - 1 + uint blocksX, blocks; + float invCount; // 10000 / (2^32 * render pixels) + uint reset; // the kernel's argument 32: no log ratio when set +} p; +#if PASS == 0 +// One thread per column of the 16x16 block. NVIDIA's tree over the block's 256 threads (index 16 y + x: i += i + 128, +// then 64, 32, ... 1) is, per column, ((r0 + r8) + (r4 + r12)) + ((r2 + r10) + (r6 + r14)) of the even rows plus the +// same of the odd rows, and then the columns (c + c8, + 4, + 2, + 1). +layout(local_size_x = 16) in; +layout(set = 0, binding = 1) uniform sampler2D color; +layout(set = 0, binding = 2, std430) writeonly buffer Sums { uint64_t s[]; } sums; +shared float col[16]; +void main() +{ + const uint tx = gl_LocalInvocationID.x; + const int x = int(gl_GlobalInvocationID.x), y0 = int(gl_WorkGroupID.y) * 16; + float r[16]; + [[unroll]] for (int k = 0; k < 16; ++k) { + float v = 0.0; + if (x <= p.last.x && y0 + k <= p.last.y) { + const vec3 t = texelFetch(color, ivec2(x, y0 + k), 0).rgb; + precise float l = fma(t.b, 0.25, fma(t.r, 0.25, t.g * 0.5)); + v = min(max(l, 0.0), 10000.0); + } + r[k] = v; + } + precise float u[8]; + [[unroll]] for (int k = 0; k < 8; ++k) u[k] = r[k + 8] + r[k]; + precise float v4[4]; + [[unroll]] for (int k = 0; k < 4; ++k) v4[k] = u[k + 4] + u[k]; + precise float w0 = v4[2] + v4[0], w1 = v4[3] + v4[1]; + precise float c = w1 + w0; + col[tx] = c; + barrier(); + if (tx == 0u) { + precise float d[8]; + [[unroll]] for (int k = 0; k < 8; ++k) d[k] = col[k + 8] + col[k]; + precise float e0 = d[4] + d[0], e1 = d[5] + d[1], e2 = d[6] + d[2], e3 = d[7] + d[3]; + precise float f0 = e2 + e0, f1 = e3 + e1; + precise float g = f1 + f0; + sums.s[gl_WorkGroupID.y * p.blocksX + gl_WorkGroupID.x] = uint64_t(roundEven(g * 429496.71875)); + } +} +#else +layout(local_size_x = 64) in; +shared uint64_t part[64]; +layout(set = 0, binding = 1, std430) readonly buffer Sums { uint64_t s[]; } sums; +layout(set = 0, binding = 2, std430) writeonly buffer Values { vec4 v; } values; +layout(set = 0, binding = 3) uniform sampler2D previous; +layout(set = 0, binding = 4, r16f) uniform writeonly image2D current; +void main() +{ + // the block sums are integers: any order. 64 strided partial sums, then thread 0 finishes + const uint tid = gl_LocalInvocationID.x; + uint64_t mine = 0ul; + for (uint b = tid; b < p.blocks; b += 64u) mine += sums.s[b]; + part[tid] = mine; + barrier(); + if (tid != 0u) return; + uint64_t total = 0ul; + for (uint b = 0u; b < 64u; ++b) total += part[b]; + precise float mean = p.invCount * float(total); + precise float e = 0.23547088 / max(mean, 0.005); + imageStore(current, ivec2(0), vec4(e)); + // cvt.rz.f16.f32: toward zero (e > 0) + float16_t h = float16_t(e); + if (float(h) > e) h = uint16BitsToFloat16(float16BitsToUint16(h) - uint16_t(1)); + float cur = float(h); + if (cur == 0.0) cur = 1.0; + float prev = texelFetch(previous, ivec2(0), 0).r; + if (prev == 0.0) prev = 1.0; + float ratio = 0.0; + if (p.reset == 0u) ratio = clamp(log2(max(cur, 1e-8) / max(prev, 1e-8)), -6.0, 6.0) * 0.16666667; + values.v = vec4(cur, prev, 1.0 / cur, ratio); +} +#endif diff --git a/engine/rr/gen_rr2_enc0.py b/engine/rr/gen_rr2_enc0.py new file mode 100644 index 0000000..efe8b59 --- /dev/null +++ b/engine/rr/gen_rr2_enc0.py @@ -0,0 +1,298 @@ +#!/usr/bin/env python3 +"""Generate preset F's encoder 0 shader: ptx2rr.py's translation of NVIDIA's per-pixel feature gather and of its guide +tail, around the native block (engine/rr/rr2_m.comp, KIND 4). + +usage: gen_rr2_enc0.py ENC0.ptx ENC0.generated.comp OUT.comp + gen_rr2_enc0.py DEC0_1.ptx DEC0_1.generated.comp OUT.comp --dec0_1 + +NVIDIA's kernel runs one 32-lane warp per 8x8-token window and keeps everything in registers. Its matrix products +(the embedding, the attention block and the guide head) are one contiguous run of the translated main, which is cut +out: + wave w the translation up to the first product: the gather; the first product's A fragments (the window's 64 + feature slots per token) go to the shared array XF + 4 waves rr2_block(): embedding, attention, MLP, the tokens to the arena, the guide head to the shared array HD + wave w the last product's D fragments are loaded from HD and the translation continues after it: the guide tail +With --group N (2 to 4) a workgroup holds N windows along x: waves 0 ... N - 1 run the translated code of one window +each and the block then runs once per window; the stage's dispatch is 1 / N as wide ("xN" in rr.txt). Measured on +an RX 9070 XT (2026-10-13) that is slower than one window per workgroup (network +0.4 ms at 1280x720 with N = 4). +With --dec0_1 the same for dec0's second kernel (rr2_m.comp KIND 5): nothing before the block (it reads its tokens +from the arena); after it the last product's D fragments (the MLP sum) and the registers holding r are loaded from +MSH / RSH, and the translated tail forms y = r + MLP * s2 and computes the kernel's outputs from it. +The translated parts hold code derived from the user's DLL: the output is model data, do not commit it. +""" +import re +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from phase_dump import mma_statements, phases_of # noqa: E402 + + +def tiles(mm): + """per mma of a phase: (m-tile, n-tile, k-step), as phase_dump.py orders them""" + dset = {tuple(m[2]): n for n, m in enumerate(mm)} + prev = {n: dset.get(tuple(m[5])) for n, m in enumerate(mm)} + step = {} + for n in range(len(mm)): + s, q = 0, n + while prev[q] is not None: + q = prev[q] + s += 1 + step[n] = s + a_tile, b_tile = {}, {} + for s in range(max(step.values()) + 1): + for t, key in enumerate(dict.fromkeys(tuple(m[3]) for n, m in enumerate(mm) if step[n] == s)): + a_tile[(s, key)] = t + for t, key in enumerate(dict.fromkeys(tuple(m[4]) for n, m in enumerate(mm) if step[n] == s)): + b_tile[(s, key)] = t + return [(a_tile[(step[n], tuple(m[3]))], b_tile[(step[n], tuple(m[4]))], step[n]) for n, m in enumerate(mm)] + + +def f16vars_of(decls): + names = set() + for l in decls: + if re.match(r'\s*(precise )?float16_t ', l): + names |= set(re.findall(r'\b(rs\d+)\b', l)) + return names + + +REG = r'\b(?:rs|rd|r|f|p|l)\d+(?:_low|_high)?\b' +# what makes a statement of the gather more than index arithmetic on the parameter block +DATA = re.compile(r'fetch0|texture|shfl_|GX\d|GA\d|GW\d|SH\d|sust_|movm_|mma_') + + +def tail_inputs(pre, tail): + """The gather's values the tail reads: (recomputed, saved). Recomputed are the statements, in order, of values + that are index arithmetic on the parameter block (assigned once, outside any block); saved are the names of the + others, which the gather pass writes to the window's record.""" + assigned, depth = {}, 0 + for i, l in enumerate(pre): + m = re.match(r'\s*(\w+) = (.*)', l) + if m: + assigned.setdefault(m.group(1), []).append((i, depth, m.group(2))) + depth += l.count('{') - l.count('}') + used, own = [], set() + for l in tail: + m = re.match(r'\s*(\w+) = (.*)', l) + for v in re.findall(REG, m.group(2) if m else l): + if v not in own and v in assigned and v not in used: + used.append(v) + if m: + own.add(m.group(1)) + + def closure(v, seen): + """statement indices v needs, or None when it is not plain index arithmetic""" + if v in seen: + return seen[v] + seen[v] = None + defs = assigned.get(v) + if defs is None: + seen[v] = set() # never assigned in the gather: a declaration's zero + return seen[v] + if len(defs) != 1 or defs[0][1] != 0 or DATA.search(defs[0][2]): + return None + need = {defs[0][0]} + for u in re.findall(REG, defs[0][2]): + c = closure(u, seen) + if c is None: + return None + need |= c + seen[v] = need + return need + + lines, saved, seen = set(), [], {} + for v in used: + c = closure(v, seen) + if c is None: + saved.append(v) + else: + lines |= c + return [pre[i] for i in sorted(lines)], saved + + +def split_layout(xf_word, pre, tail, f16vars): + """(recomputed statements, saved names, the record's words, the GLSL locating a window's record)""" + recomputed, saved = tail_inputs(pre, tail) + for v in saved: + if not (v in f16vars or re.fullmatch(r'r\d+', v)): + raise SystemExit(f'the tail reads {v} from the gather: not a 16- or 32-bit register') + record = 1024 + 32 * len(saved) + where = (f' const uint rr2_gx = (param.w[0] + (param.w[20] & 0xffffu) + 7u) / 8u;\n' + f' const uint rr2_rec = param.w[{xf_word}u] / 4u + {record}u * (gl_WorkGroupID.y * (param.w[0] / 8u + 1u) + gl_WorkGroupID.x);\n') + return recomputed, saved, record, where + + +def gather_pass(xf_word, head, decls, pre, tail, first, f16vars): + """Encoder 0 as two passes (--split gather | block): the translated gather in NVIDIA's own 32-thread groups, which + writes each window's record to the arena, and the native block with the translated tail in 128-thread groups. + One shader holding both runs the gather on one of the block's four waves while the other three wait, which + costs more than the record's trip through memory (RX 9070 XT, 1280x720, 2026-10-13: gather 0.40 -> 0.24 ms). + The record: 1024 words of A fragments ([token slot][64 FP8 feature slots], as the block's XF), then per lane the + gather's values the tail reads. Its address: parameter word XF_WORD (bytes) + 4 * RECORD * (y * (W / 8 + 1) + x).""" + recomputed, saved, record, where = split_layout(xf_word, pre, tail, f16vars) + var = lambda r: r.lstrip('%') + store = [' if (rr2_valid) { // the first product\'s A fragments: lane 4g + t, register q -> row, 4 k', + ' const uint l_ = gl_SubgroupInvocationID, g_ = l_ >> 2u, t_ = l_ & 3u;'] + seen = set() + for mm_, (i, j, s) in zip(first, tiles(first)): + if (i, s) in seen: + continue + seen.add((i, s)) + for q, r in enumerate(mm_[3]): + store.append(f' rr2_arena.w[rr2_rec + rr2_uninter({16 * i + 8 * (q & 1)}u + g_) * 16u + {(32 * s + 16 * (q >> 1)) // 4}u + t_] = {var(r)};') + for n, v in enumerate(saved): + bits = f'uint(float16BitsToUint16({v}))' if v in f16vars else v + store.append(f' rr2_arena.w[rr2_rec + 1024u + {len(saved)}u * l_ + {n}u] = {bits};') + store.append(' }') + print(f'gather pass: record of {record} words, the tail reads {", ".join(saved) or "nothing"} from it') + return (head + '\nlayout(set = 0, binding = 2, std430) buffer Rr2Arena { uint w[]; } rr2_arena;\n' + 'uint rr2_uninter(uint c) { return (c & ~31u) | (((c & 15u) >> 1u) << 2u) | (c & 1u) | (((c >> 4u) & 1u) << 1u); }\n\n' + 'void main()\n{\n' + '\n'.join(decls) + '\n' + where + + ' const bool rr2_valid = gl_WorkGroupID.x < rr2_gx;\n' + ' if (rr2_valid) { // NVIDIA\'s warp: the translation before the first product\n' + '\n'.join(pre) + '\n }\n' + + '\n'.join(store) + '\n}\n'), record + + +def main(): + ptx, src, out = sys.argv[1:4] + dec = '--dec0_1' in sys.argv + # --group N: windows per workgroup (1 to 4; wave w runs the translated code of window w) + G = int(sys.argv[sys.argv.index('--group') + 1]) if '--group' in sys.argv else 1 + kernel = 'rr2_dec0_kernel_1' if dec else 'rr2_enc0_kernel' + lines = Path(ptx).read_text().replace('\r', '').split('\n') + k = next(i for i, l in enumerate(lines) if f'.entry {kernel}(' in l) + phases = phases_of(mma_statements(lines, k, len(lines)), lines) + first, last = phases[0], phases[-1] + var = lambda r: r.lstrip('%') + + g = Path(src).read_text().split('\n') + m = next(i for i, l in enumerate(g) if l.startswith('void main()')) + head = '\n'.join(g[:m]).replace('layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in;', + 'layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in;') + if dec: # the tail's tanh from rr2_pack.py's table (ptx2rr.py's RR_TANH_TABLE; its offset from OUT's rr2_stage.h) + v = head.index('\n') + 1 + table = re.search(r'#define W_TANH (\d+)u', (Path(out).parent / 'rr2_stage.h').read_text()).group(1) + head = head[:v] + f'#define RR_TANH_TABLE {table}u\n' + head[v:] + body = g[m + 2:] + while body and body[-1].strip() in ('', '}'): + last_line = body.pop() + if last_line.strip() == '}': + break + nd = next(i for i, l in enumerate(body) if not re.match(r' (precise |bool |uint64_t |float |uint |int |float16_t |f16vec2 )', l)) + calls = [i for i, l in enumerate(body) if 'mma_e4m3_f16' in l] + a = calls[0] - 1 # '{ uint md0_, ...;' + b = next(i for i in range(calls[-1], len(body)) if body[i].strip() == '}') + if not body[a].strip().startswith('{ uint md0_'): + raise SystemExit('unexpected statement before the first product: ' + body[a]) + pre, cut, tail = body[nd:a], body[a:b + 1], body[b + 1:] + if sum(l.count('{') - l.count('}') for l in pre): + raise SystemExit('the first product is inside a block of the translated main') + + # values the tail takes from the cut region other than the last product's results + assigned = lambda ls: {mm_.group(1) for l in ls for mm_ in [re.match(r'\s*(\w+) = ', l)] if mm_} + restored = {var(r) for mm_ in last for r in mm_[2]} + rmap = {} # dec0_1: the register holding r for (row offset, column offset): y = r + (D half) * s2 in the tail + if dec: + place = {} + for mm_, (i, j, s_) in zip(last, tiles(last)): + for q, r in enumerate(mm_[2]): + place[r] = (16 * i + 8 * q, 8 * j) + text_tail = '\n'.join(lines[last[-1][1]:]) + half = {} + for mo in re.finditer(r'mov\.b32\s+\{(%rs\d+),\s*(%rs\d+)\},\s*(%r\d+);', text_tail): + if mo.group(3) in place: + half[mo.group(1)] = (*place[mo.group(3)], 0) + half[mo.group(2)] = (*place[mo.group(3)], 1) + prod = {} + for mo in re.finditer(r'mul\.f16\s+(%rs\d+),\s*(%rs\d+),\s*(%rs\d+);', text_tail): + if mo.group(2) in half: + prod[mo.group(1)] = half[mo.group(2)] + for mo in re.finditer(r'add\.f16\s+(%rs\d+),\s*(%rs\d+),\s*(%rs\d+);', text_tail): + if mo.group(2) in prod: + rmap[var(mo.group(3))] = prod[mo.group(2)] + if len(rmap) != 2 * 2 * len(last): + raise SystemExit(f'found {len(rmap)} residual registers, expected {4 * len(last)}') + restored |= set(rmap) + used = set(re.findall(r'\b[a-z]+\d+\b', '\n'.join(tail))) + lost = sorted((assigned(cut) - assigned(pre) - restored) & used) + # (values assigned again in the tail before their use are fine; the kernel has none on the reference build) + keep = [] + cut_set = assigned(cut) + for name in lost: # thread indices and the like: their assignment is carried over when it depends on nothing cut + line = next(l for l in reversed(cut) if re.match(rf'\s*{name} = ', l)) + if set(re.findall(r'\b[a-z]+\d+\b', line.split('=', 1)[1])) & cut_set: + raise SystemExit(f'the tail needs {name}, which the cut region computes: {line.strip()}') + keep.append(' ' + line) + + split = sys.argv[sys.argv.index('--split') + 1] if '--split' in sys.argv else None + xf_word = int(sys.argv[sys.argv.index('--xf-word') + 1]) if split else 0 + if split == 'gather': + text, record = gather_pass(xf_word, '\n'.join(g[:m]), body[:nd], pre, tail, first, f16vars_of(body[:nd])) + Path(out).write_text(text) + Path(out).with_suffix('.record').write_text(str(record)) + return + if dec: + first = [] + store = [' if (rr2_valid) { // the first product\'s A fragments: lane 4g + t, register q -> row, 4 k', + ' const uint l_ = gl_SubgroupInvocationID, g_ = l_ >> 2u, t_ = l_ & 3u;'] + seen = set() + for mm_, (i, j, s) in zip(first, tiles(first) if first else []): + if (i, s) in seen: + continue + seen.add((i, s)) + for q, r in enumerate(mm_[3]): + store.append(f' XF[gl_SubgroupID * 1024u + uninter({16 * i + 8 * (q & 1)}u + g_) * 16u + {(32 * s + 16 * (q >> 1)) // 4}u + t_] = {var(r)};') + store.append(' }') + load = [' { // the guide head as the last product\'s D fragments (rows g, g + 8; columns 2t, 2t + 1)', + ' const uint l_ = gl_SubgroupInvocationID, g_ = l_ >> 2u, t_ = l_ & 3u;'] + f16vars = set() + for l in body[:nd]: + if re.match(r'\s*(precise )?float16_t ', l): + f16vars |= set(re.findall(r'\b(rs\d+)\b', l)) + for mm_, (i, j, s) in zip(last, tiles(last)): + for q, r in enumerate(mm_[2]): + row = f'({16 * i + 8 * q}u + g_)' + load.append(f' {var(r)} = HDV({row}, {8 * j}u + 2u * t_) | (HDV({row}, {8 * j + 1}u + 2u * t_) << 16u);') + for name, (ro, co, hf) in rmap.items(): + v = f'RSV({ro}u + g_, {co + hf}u + 2u * t_)' + load.append(f' {name} = ' + (f'uint16BitsToFloat16(uint16_t({v}));' if name in f16vars else f'{v};')) + load.append(' }') + # the native block's source, without its #version line (this shader has its own) + block = (Path(__file__).resolve().parent / 'rr2_m.comp').read_text().split('\n', 1)[1] + if split == 'block': # the second pass: the block on the gather pass's record, then the tail on wave 0 + recomputed, saved, record, where = split_layout(xf_word, pre, tail, f16vars) + restore = [' ' + (f'{v} = uint16BitsToFloat16(uint16_t(arena.w[rr2_rec + 1024u + {len(saved)}u * l_ + {n}u]));' if v in f16vars + else f'{v} = arena.w[rr2_rec + 1024u + {len(saved)}u * l_ + {n}u];') for n, v in enumerate(saved)] + load[2:2] = restore + Path(out).write_text( + head + '\n#extension GL_GOOGLE_include_directive : require\n#define RR2_EMBED 1\n#define RR2_GROUP 1\n#define RR2_XF_ARENA 1\n' + + block + '\nuint HDV(uint row, uint col) { return col < 4u ? uint(float16BitsToUint16(HD[gl_SubgroupID * 256u + row * 4u + col])) : 0u; }\n\n' + 'void main()\n{\n' + '\n'.join(body[:nd]) + '\n' + where + + ' const bool rr2_valid = gl_SubgroupID < 1u && gl_WorkGroupID.x < rr2_gx;\n' + ' rr2_xf = rr2_rec;\n rr2_block();\n barrier();\n if (rr2_valid) {\n' + + ' // the gather\'s index arithmetic the tail reads\n' + '\n'.join(recomputed) + '\n' + + '\n'.join(load + keep) + '\n' + '\n'.join(tail) + '\n }\n}\n') + return + wx = lambda t: t.replace('gl_WorkGroupID.x', f'({G}u * gl_WorkGroupID.x + gl_SubgroupID)') + text = (head + f'\n#extension GL_GOOGLE_include_directive : require\n#define RR2_EMBED 1\n#define RR2_GROUP {G}\n' + block + '\n' + + ('uint HDV(uint row, uint col) { return uint(float16BitsToUint16(MSH[gl_SubgroupID * 2048u + row * 32u + col])); }\n' + 'uint RSV(uint row, uint col) { return uint(float16BitsToUint16(RSH[gl_SubgroupID * 2048u + row * 32u + col])); }\n\n' if dec else + 'uint HDV(uint row, uint col) { return col < 4u ? uint(float16BitsToUint16(HD[gl_SubgroupID * 256u + row * 4u + col])) : 0u; }\n\n') + + 'void main()\n{\n' + '\n'.join(body[:nd]) + + '\n // this wave\'s window, when the level has it\n' + ' const uint rr2_gx = (param.w[0] + (param.w[20] & 0xffffu) + 7u) / 8u;\n' + f' const bool rr2_valid = gl_SubgroupID < {G}u && {G}u * gl_WorkGroupID.x + gl_SubgroupID < rr2_gx;\n' + ' if (rr2_valid) { // NVIDIA\'s warp: the translation before the first product\n' + wx('\n'.join(pre)) + '\n }\n' + + ('' if dec else '\n'.join(store)) + + # (the block runs unconditionally, also for a window past the level's last one: it writes nothing there, and + # inside a branch the driver compiles it much slower) + ('\n barrier();\n rr2_block();\n barrier();\n' if G == 1 else + f'\n barrier();\n for (uint w_ = 0u; w_ < {G}u; ++w_) {{\n rr2_win = w_;\n' + ' rr2_block();\n barrier();\n }\n') + ' if (rr2_valid) {\n' + + '\n'.join(load + keep) + '\n' + wx('\n'.join(tail)) + '\n }\n}\n') + Path(out).write_text(text) + print(f'{out}: {len(pre)} gather lines, {len(cut)} cut, {len(tail)} tail lines, {len(seen)} A fragment sets') + + +if __name__ == '__main__': + main() diff --git a/engine/rr/phase_dump.py b/engine/rr/phase_dump.py index 2eca0b6..553c787 100644 --- a/engine/rr/phase_dump.py +++ b/engine/rr/phase_dump.py @@ -29,7 +29,8 @@ REG = r'%r\d+' def kernel_module(kernel): - for p in sorted(Path('~/.cache/d4r-rr-corpus').expanduser().glob('*.ptx')): + corpus = '~/.cache/d4r-rr-corpus-f' if kernel.startswith('rr2_') else '~/.cache/d4r-rr-corpus' + for p in sorted(Path(corpus).expanduser().glob('*.ptx')): if f'.entry {kernel}(' in p.read_text(errors='replace'): return p raise SystemExit(f'no corpus module for {kernel}') @@ -80,7 +81,9 @@ def main(): o = ap.parse_args() # the capture decides the kernel: a preset D capture runs cuda_dldn_engine_swin_{enc0,dec0}_diamond_wallaby_kernel cap = Capture(os.path.expanduser(o.capture)) - launch = [l for l in cap.launches if f'swin_{o.stage}' in l.kernel][o.frame - 1] + # preset F: STAGE is the kernel's full name (rr2_enc1_kernel, rr2_dec1_kernel_0) + rr2 = o.stage.startswith('rr2_') + launch = [l for l in cap.launches if (l.kernel == o.stage if rr2 else f'swin_{o.stage}' in l.kernel)][o.frame - 1] kernel = launch.kernel module = kernel_module(kernel) lines = module.read_text().replace('\r', '').split('\n') @@ -126,7 +129,7 @@ def main(): cmd = ['systemd-run', '--user', '--scope', '-q', '-p', 'MemoryMax=24G', '-p', 'MemorySwapMax=0'] cmd += [f'--setenv={k_}={v_}' for k_, v_ in env.items() if k_.startswith('D4R_ZLUDA')] r = subprocess.run(cmd + [sys.executable, str(HERE / 'zluda_replay.py'), os.path.expanduser(o.capture), - f'swin_{o.stage}' + ('_diamond' if 'diamond' in kernel else ''), str(o.frame), str(work), + kernel if rr2 else f'swin_{o.stage}' + ('_diamond' if 'diamond' in kernel else ''), str(o.frame), str(work), '--ptx', str(work / 'probe.ptx'), '--pad', str(pad)], capture_output=True, text=True) if 'failed' in r.stdout + r.stderr: diff --git a/engine/rr/ptx2rr.py b/engine/rr/ptx2rr.py index 741e6cb..5d4223c 100755 --- a/engine/rr/ptx2rr.py +++ b/engine/rr/ptx2rr.py @@ -14,6 +14,7 @@ model/build directory only; do not commit or redistribute it. import argparse import importlib.util import json +import os import re import struct import sys @@ -110,6 +111,78 @@ void sust_p4(writeonly image2D s, int x, int y, vec4 v) if (uint(x) < uint(n.x) && uint(y) < uint(n.y)) imageStore(s, ivec2(x, y), v); } void sust_p1(writeonly image2D s, int x, int y, float v) { sust_p4(s, x, y, vec4(v, 0.0, 0.0, 0.0)); } +// Formatted stores to normalized surfaces, as the reference converts them (ZLUDA's surface_formatted_store_bits): NaN +// to zero, clamp, round half away from zero in f32. The quotient is what the image store then converts back exactly. +// UNORM 10:10:10:2 surfaces (the pooled histories) are RGBA32F images holding the quantised values, like ZLUDA's. +// tanh.approx.f32 as the reference computes it: ROCm's OCML tanh (device-libs ocml/src/tanhF.cl) with LLVM's AMDGPU +// lowering of exp (lowerFEXP, fast-FMA path): the same fused operations, hardware exp2, ldexp and reciprocal. +// (a fused f32 a * b + c through f64: the product is exact there and the sum is rounded to 53 bits before the f32 +// rounding. The driver's own f32 fma is a product and a sum whenever its product also exists on its own, as in +// exp's ph / fma(x, c, -ph) pair.) +float fmad(float a, float b, float c) { precise double s = double(a) * double(b) + double(c); return float(s); } +float tanh_ocml(float x) +{ + const float y = abs(x); + precise float z; + if (y < 0.625) { + precise float y2 = y * y; + precise float p = fmad(y2, fmad(y2, fmad(y2, fmad(y2, -0.005700020585209131, 0.02063407190144062), -0.053737930953502655), 0.13331416249275208), -0.3333328068256378); + precise float yp = y * p; + z = fmad(y2, yp, y); + } else { + const float w = 2.0 * y; + precise float ph = w * 1.4426950216293335; + precise float f0 = fmad(w, 1.4426950216293335, -ph); + precise float pl = fmad(w, 1.925962855864327e-08, f0); + const float e = roundEven(ph); + precise float d = ph - e; + precise float a = d + pl; + float t = ldexp(exp2(a), int(e)); + if (w > 88.72283172607422) t = uintBitsToFloat(0x7f800000u); + precise float s1 = t + 1.0; + z = fmad(-2.0, rcp_a(s1), 1.0); + } + return x < 0.0 || floatBitsToUint(x) == 0x80000000u ? -z : z; +} +// tanh.approx.f16: from a table of the correctly rounded values indexed by the argument's bits (RR_TANH_TABLE: its +// byte offset in the weight buffer; engine/rr/rr2_pack.py) where a stage has one, else through f32 +#ifdef RR_TANH_TABLE +float16_t tanh16(float16_t x) { return uint16BitsToFloat16(wbuf.wh[RR_TANH_TABLE / 2u + uint(float16BitsToUint16(x))]); } +#else +float16_t tanh16(float16_t x) { return float16_t(tanh(float(x))); } +#endif +vec4 q_snorm8(vec4 v) +{ + v = clamp(mix(v, vec4(0.0), isnan(v)), vec4(-1.0), vec4(1.0)); + precise vec4 s = v * 127.0; + precise vec4 a = abs(s) + 0.5; + return sign(s) * floor(a) / 127.0; +} +vec4 q_unorm8(vec4 v) +{ + v = clamp(mix(v, vec4(0.0), isnan(v)), vec4(0.0), vec4(1.0)); + precise vec4 a = v * 255.0 + 0.5; + return floor(a) / 255.0; +} +// (the quotient q / 1023 must be the correctly rounded f32 the reference stores: the driver divides by multiplying +// with a reciprocal, one step off for some q, so it comes from a table) +const float Q1023[1024] = float[1024](Q1023_VALUES); +vec4 q_unorm1010102(vec4 v) +{ + v = clamp(mix(v, vec4(0.0), isnan(v)), vec4(0.0), vec4(1.0)); + precise vec4 a = v * vec4(1023.0, 1023.0, 1023.0, 3.0) + 0.5; + const ivec4 q = ivec4(floor(a)); + const float third[4] = float[4](0.0, 0.3333333432674408, 0.6666666865348816, 1.0); + return vec4(Q1023[q.x], Q1023[q.y], Q1023[q.z], third[q.w]); +} +// for a real UNORM 10:10:10:2 image: the reference's rounding (half away from zero) here, and any value within half a +// step of q / 1023 for the store, which the hardware quantises back to q +vec4 q_unorm1010102_hw(vec4 v) +{ + v = clamp(mix(v, vec4(0.0), isnan(v)), vec4(0.0), vec4(1.0)); + precise vec4 a = v * vec4(1023.0, 1023.0, 1023.0, 3.0) + 0.5; + return floor(a) * vec4(1.0 / 1023.0, 1.0 / 1023.0, 1.0 / 1023.0, 1.0 / 3.0); +} // Binary stores write the halves' bits. RADV folds float(float16_t(x)) back to x when the half came from a // cvt.rn.f16.f32 in a register, and the image store then truncates x instead (the output stage's history surface // was one binary16 ulp low in half its values, which the next frame's encoder 0 amplifies): OR-ing in a zero the @@ -135,7 +208,7 @@ uint shfl_dir(uint v, uint delta, uint opts, bool down, out bool in_bounds) bool oob = down ? idx > sub_end : idx < sub_end; if (oob) idx = self; in_bounds = !oob; - return subgroupShuffle(v, uint(idx) & 31u); + return subgroupShuffle(v, (gl_SubgroupInvocationID & 32u) | (uint(idx) & 31u)); // the warp is this half of a 64-wide wave } // shfl.sync.idx (PTX ISA): j = minLane | (b & ~segmask), in bounds while j <= maxLane, otherwise the lane's own value @@ -144,7 +217,7 @@ uint shfl_idx(uint v, uint b, uint opts, out bool in_bounds) uint mask = (opts >> 8) & 31u, lane = gl_SubgroupInvocationID & 31u; uint maxLane = (lane & mask) | (opts & 31u & ~mask), j = (lane & mask) | (b & 31u & ~mask); in_bounds = j <= maxLane; - return subgroupShuffle(v, in_bounds ? j : lane); + return subgroupShuffle(v, (gl_SubgroupInvocationID & 32u) | (in_bounds ? j : lane)); } // movmatrix.sync.trans.aligned.m8n8.b16 on one b32 fragment register per lane: lane (g = lane>>2, t = lane&3) // holds row g, columns 2t and 2t+1. The transpose gathers (row 2t, col g) and (row 2t+1, col g). @@ -337,12 +410,16 @@ def base_helper(name): return '\n'.join(lines[start:end + 1]) + '\n' +RR_RUNTIME = RR_RUNTIME.replace('Q1023_VALUES', ', '.join(repr(float(np.float32(q) / np.float32(1023))) for q in range(1024))) RR_RUNTIME += base_helper('prmt') + base_helper('shfl_bfly').replace( - 'return subgroupShuffle(v, (gl_SubgroupInvocationID & 32u) | uint(idx));', 'return subgroupShuffle(v, uint(idx) & 31u);') + 'return subgroupShuffle(v, (gl_SubgroupInvocationID & 32u) | uint(idx));', + 'return subgroupShuffle(v, (gl_SubgroupInvocationID & 32u) | (uint(idx) & 31u));') class RRT(base.T): - def __init__(self, ptx, kernel, block): + def __init__(self, ptx, kernel, block, weights_arg=40, buffers=()): + self.weights_arg = weights_arg + self.buffers = sorted(buffers) # argument offsets of further buffers, each its own binding bx, by, bz = block if bx * by * bz % 32: raise base.Err('block size must be a multiple of the 32-lane subgroup') @@ -367,10 +444,18 @@ class RRT(base.T): # without it the compiler may contract a product and a sum across them depending on scheduling (results then # changed with unrelated code and with the device's robustness features) self.precise_all = True + # module globals with an initialiser (the output filter's bicubic phase table): constants, laid out one after + # another in a blob that is the stage's weight buffer (binding 1); `mov.u64 %rd, NAME` is its offset there + self.globals, self.globals_blob = {}, b'' + for gm in re.finditer(r'\.global\s+\.align\s+\d+\s+\.b8\s+(\w+)\[(\d+)\]\s*=\s*\{([^}]*)\}', ptx): + self.globals[gm.group(1)] = len(self.globals_blob) + data = bytes(int(x) for x in gm.group(3).split(',')) + self.globals_blob += data + bytes(-len(data) % 16) self.stmt_lines = statement_lines(ptx, kernel) self.stmt_index = 0 super().__init__(ptx, kernel) - self.pair_mma() + if not os.environ.get('D4R_PTX2RR_NO_PAIR'): + self.pair_mma() self.origins() def stmt(self, s): @@ -438,9 +523,12 @@ class RRT(base.T): if not a or not re.fullmatch(r'%rd\d+', a[0]) or op.startswith(('ld.', 'st.', 'bra')): if op != 'ld.param.u64': continue - if op == 'ld.param.u64': + if op.startswith('mov') and len(a) > 1 and a[1].strip() in self.globals: + new = {'w'} + elif op == 'ld.param.u64': # the argument's true offset (the load may go through a register based inside the block) - new = {'w' if self.param_off(a[1]) == 40 else 'a'} + off_ = self.param_off(a[1]) + new = {'w' if off_ == self.weights_arg else f'x{off_}' if off_ in self.buffers else 'a'} else: new = set() for x in a[1:]: @@ -457,21 +545,29 @@ class RRT(base.T): raise base.Err(f'global pointer {reg} has origin {sorted(s)} (need exactly one of w/a)') return next(iter(s)) + def gfn(self, reg, store=False): + o = self.gorigin(reg) + if o == 'w': + return 'GW' + if o == 'a': + return 'SA' if store else 'GA' + return ('SX' if store else 'GX') + o[1:] + '_' + def gaddr(self, tok): - m = re.fullmatch(r'\[\s*(%rd\d+)\s*(?:\+\s*(\d+))?\s*\]', tok.strip()) + m = re.fullmatch(r'\[\s*(%rd\d+)\s*(?:\+\s*(-?\d+))?\s*\]', tok.strip()) if not m: raise base.Err(f'unsupported global address {tok}') return m.group(1), int(m.group(2) or 0) def gbyte(self, reg, off): - return f'uint({self.val(reg, "u64")}) + {off}u' + return f'uint({self.val(reg, "u64")}) + {off}u' if off >= 0 else f'uint({self.val(reg, "u64")}) - {-off}u' def ld(self, o, a): if 'global' not in o: return super().ld(o, a) size, count = mem_shape(o) reg, off = self.gaddr(a[1]) - fn = 'GW' if self.gorigin(reg) == 'w' else 'GA' + fn = self.gfn(reg) dsts = braces(a[0]) if count > 1 else [a[0]] for i, d in enumerate(dsts): self.A(d, f'{fn}{size * 8}({self.gbyte(reg, off + i * size)})', 'u16' if size == 2 else 'u32') @@ -481,14 +577,15 @@ class RRT(base.T): return super().st(o, a) size, count = mem_shape(o) reg, off = self.gaddr(a[0]) - if self.gorigin(reg) != 'a': - raise base.Err('store through the read-only weight pointer') + if self.gorigin(reg) == 'w': + raise base.Err(f'store through the read-only weight pointer: {a}') + fn = self.gfn(reg, True) vals = braces(a[1]) if count > 1 else [a[1]] for i, v in enumerate(vals): if size == 2: - self.E(f'SA16({self.gbyte(reg, off + i * size)}, {self.val(v, "u16")});') + self.E(f'{fn}16({self.gbyte(reg, off + i * size)}, {self.val(v, "u16")});') else: - self.E(f'SA32({self.gbyte(reg, off + i * size)}, {self.val(v, "u32")});') + self.E(f'{fn}32({self.gbyte(reg, off + i * size)}, {self.val(v, "u32")});') # ------------------------------------------------ textures and surfaces def handle_off(self, h): @@ -515,7 +612,7 @@ class RRT(base.T): # (zluda_ptx_impl.cpp fetch_2D). The output stage's specular reprojection reads such neighbours at the image # border: with clamp its surfaces were 99.92% identical to the capture, with zero 99.997% (capDX, 2026-10-12) fetch = f'fetch0(T{off}, ivec2({self.val(c[0], "s32")}, {self.val(c[1], "s32")}))' - elif 'base' in o and o[-1] == 'f32': + elif ('base' in o or o[0] == '2d') and o[-1] == 'f32': # tex.2d: no derivatives in a kernel, level 0 fetch = f'textureLod(T{off}, vec2({self.val(c[0], "f32")}, {self.val(c[1], "f32")}), 0.0)' elif 'level' in o and o[-1] == 'f32': fetch = (f'textureLod(T{off}, vec2({self.val(c[0], "f32")}, {self.val(c[1], "f32")}), ' @@ -525,7 +622,9 @@ class RRT(base.T): f16 = 'f16' in o self.E(f'{{ vec4 t_ = {fetch};') for i, d in enumerate(dst): - self.A(d, f'float16_t(t_[{i}])' if f16 else f't_[{i}]', 'f16' if f16 else 'float') + # a binary16 result is the texel rounded to binary16: through H(HU(...) | SUST_Z) so that the driver cannot + # fold the rounding away when the value is widened again (RR_RUNTIME's note on sust_b) + self.A(d, f'H(HU(float16_t(t_[{i}])) | SUST_Z)' if f16 else f't_[{i}]', 'f16' if f16 else 'float') self.E('}') def sust(self, op, a): @@ -534,9 +633,10 @@ class RRT(base.T): v = braces(a[1]) x, y = self.val(c[0], 's32'), self.val(c[1], 's32') S = f'S{off}' + op = op.replace('.cs.', '.') # the streaming cache policy: a hint, the store is the same if op == 'sust.p.2d.v4.b32.zero': self.set_surface(off, 'rgba16f') - self.E(f'sust_p4({S}, {x}, {y}, vec4({", ".join(self.val(q, "f32") for q in v)}));') + self.E(f'sust_p4({S}, {x}, {y}, SUSTQ_{off}(vec4({", ".join(self.val(q, "f32") for q in v)})));') elif op == 'sust.p.2d.b32.zero': self.set_surface(off, 'r16f') self.E(f'sust_p1({S}, {x}, {y}, {self.val(v[0], "f32")});') @@ -579,6 +679,13 @@ class RRT(base.T): types = [t for t in o if re.match(r'^(f|s|u|b)\d+(x2)?$', t)] if types[:2] in (['u32', 'u64'], ['s32', 's64'], ['u32', 's64'], ['s32', 'u64']): return self.A(a[0], f'uint({self.val(a[1], "u64")})', 'u32') + if types[:2] == ['f16', 'f32'] and 'rn' in o and 'satfinite' not in o: + # Two things the driver must not do. It folds float(float16_t(x)) to x when the result is widened again + # (RR_RUNTIME's note on sust_b): the OR on the binary16 bits. And it fuses a product or sum with this + # conversion into v_fma_mix, one rounding of the exact value, where the kernel rounds to f32 first: an f32 + # result that is exactly a binary16 midpoint then rounds the other way (the output stages' rare one-step + # differences, e.g. log2(x) * c1 * c2; found 2026-10-13): the OR on the f32 bits keeps the f32 value. + return self.A(a[0], f'H(HU(float16_t(uintBitsToFloat(floatBitsToUint({self.val(a[1], "f32")}) | SUST_Z))) | SUST_Z)', 'f16') return super().cvt(op, o, a) def f64_op(self, op, a): @@ -610,6 +717,22 @@ class RRT(base.T): return self.E('}') if op == 'mma.paired.noop': return + if parts[0] == 'mov' and len(a) > 1 and a[1].strip() in self.globals: + return self.A(a[0], f'uint64_t({self.globals[a[1].strip()]}u)', 'u64') + # cross-kernel progress counters (NGX runs the output passes concurrently and has each wait for the rows it + # needs): plain loads and adds here. The engine runs the passes one after another and leaves the counter + # pointers null, which the kernels test, so these never execute. + if parts[0] == 'ld' and ('relaxed' in parts or 'acquire' in parts) and 'global' in parts: + return self.ld([q for q in parts[1:] if q not in ('relaxed', 'acquire', 'gpu') and '::' not in q and not q.startswith('L1')], a) + if parts[0] == 'red' and parts[-2] == 'add': + reg, off = self.gaddr(a[0]) + fn, sfn = self.gfn(reg), self.gfn(reg, True) + return self.E(f'{sfn}32({self.gbyte(reg, off)}, {fn}32({self.gbyte(reg, off)}) + {self.val(a[1], "u32")});') + if op == 'membar.gl': + return + if op == 'activemask.b32': + # the warp's lanes: this half of a 64-wide wave + return self.A(a[0], '((gl_SubgroupInvocationID & 32u) != 0u ? subgroupBallot(true).y : subgroupBallot(true).x)', 'u32') if op == 'movmatrix.sync.trans.aligned.m8n8.b16': return self.A(a[0], f'movm_trans({self.val(a[1], "u32")})', 'u32') # PTX clamps shift amounts at the register width (shl/shr.u give 0, shr.s the sign); GLSL leaves them undefined @@ -635,6 +758,16 @@ class RRT(base.T): return self.tex(op, a) if parts[0] == 'cvt' and [t for t in parts[1:] if t in ('f16x2', 'e4m3x2')] == ['e4m3x2', 'f16x2']: return self.cvt_to_e4m3(op, a) + if op == 'copysign.f32': # PTX: the sign of the first source on the value of the second + x, y = self.val(a[1], 'f32'), self.val(a[2], 'f32') + return self.A(a[0], f'uintBitsToFloat((floatBitsToUint({y}) & 0x7fffffffu) | (floatBitsToUint({x}) & 0x80000000u))', 'float') + if op == 'tanh.approx.f16': + return self.A(a[0], f'tanh16({self.val(a[1], "f16")})', 'f16') + if op == 'tanh.approx.f32': + return self.A(a[0], f'tanh_ocml({self.val(a[1], "f32")})', 'float') + if op == 'ex2.approx.f16': + # v_exp_f32 of the widened value, rounded to binary16; the ORs keep that rounding (see cvt) + return self.A(a[0], f'H(HU(float16_t(uintBitsToFloat(floatBitsToUint(exp2(float({self.val(a[1], "f16")}))) | SUST_Z))) | SUST_Z)', 'f16') if op == 'neg.f16': return self.A(a[0], f'-{self.val(a[1], "f16")}', 'f16') if parts[0] in ('min', 'max') and 'NaN' not in parts and any(t in parts for t in ('f16', 'f16x2', 'f32')): @@ -659,6 +792,7 @@ class RRT(base.T): w('#extension GL_EXT_shader_16bit_storage : require') w('#extension GL_KHR_shader_subgroup_basic : require') w('#extension GL_KHR_shader_subgroup_shuffle : require') + w('#extension GL_KHR_shader_subgroup_ballot : require') w('#extension GL_KHR_cooperative_matrix : require') w('#extension GL_KHR_memory_scope_semantics : require') w('#extension GL_EXT_float_e4m3 : require') @@ -673,6 +807,15 @@ class RRT(base.T): n += 1 for off in sorted(self.surfaces): w(f'layout(set = 0, binding = {n}, {self.surfaces[off]}) uniform writeonly image2D S{off};') + q = {'rgba8_snorm': 'q_snorm8', 'rgba8': 'q_unorm8', 'rgba32f': 'q_unorm1010102', 'rgb10_a2': 'q_unorm1010102_hw'}.get(self.surfaces[off], '') + w(f'#define SUSTQ_{off}(v) {q}(v)') + n += 1 + for off in self.buffers: + w(f'layout(set = 0, binding = {n}, std430) buffer XBuf{off} {{ uint16_t h[]; }} xb{off};') + w(f'uint GX{off}_16(uint a) {{ return uint(xb{off}.h[a >> 1]); }}') + w(f'uint GX{off}_32(uint a) {{ return uint(xb{off}.h[a >> 1]) | (uint(xb{off}.h[(a >> 1) + 1u]) << 16); }}') + w(f'void SX{off}_16(uint a, uint v) {{ xb{off}.h[a >> 1] = uint16_t(v & 0xffffu); }}') + w(f'void SX{off}_32(uint a, uint v) {{ xb{off}.h[a >> 1] = uint16_t(v & 0xffffu); xb{off}.h[(a >> 1) + 1u] = uint16_t(v >> 16); }}') n += 1 w(f'shared uint sh[{max(1, (self.shared_bytes + 3) // 4)}];') if self.probe_slots or self.oob: @@ -694,6 +837,11 @@ def main(): ap.add_argument('--surface', action='append', default=[], metavar='OFFSET=FORMAT', help="the GLSL image format of a surface whose stores are formatted (sust.p), when it is not the " "store's own guess, e.g. 144=rgba8_snorm (repeatable)") + ap.add_argument('--weights-arg', type=int, default=40, + help='argument offset of the weight pointer (40 for the DLSS-D kernels, 24 for rr2_*)') + ap.add_argument('--buffer', action='append', type=int, default=[], metavar='OFFSET', + help='a pointer argument that is neither the weights nor the arena: its own storage buffer, bound ' + 'after the surfaces in ascending argument offset (repeatable)') ap.add_argument('--probe', action='append', default=[], metavar='LINE:%REG,...', help='debugging: store the registers after the statement at that module line executes (last ' 'execution) for workgroup --probe-block, into a debug buffer bound after the surfaces') @@ -709,7 +857,7 @@ def main(): probes[int(line)] = [r.strip() for r in regs.split(',')] def make(): - t = RRT(src, ns.kernel, block) + t = RRT(src, ns.kernel, block, ns.weights_arg, ns.buffer) t.probes = probes t.probe_block = tuple(int(x) for x in ns.probe_block.split(',')) t.oob = ns.oob_report @@ -719,10 +867,12 @@ def main(): if off in t.surfaces: t.surfaces[off] = fmt open(ns.out, 'w').write(t.emit({})) + if t.globals_blob: + Path(ns.out).with_suffix('.globals.bin').write_bytes(t.globals_blob) # the binding order (stage contract): textures, then surfaces, in ascending argument offset Path(ns.out).with_suffix('.json').write_text(json.dumps( {'kernel': ns.kernel, 'block': block, 'textures': sorted(t.textures), - 'surfaces': [[off, t.surfaces[off]] for off in sorted(t.surfaces)], 'shared': t.shared_bytes, + 'surfaces': [[off, t.surfaces[off]] for off in sorted(t.surfaces)], 'buffers': t.buffers, 'shared': t.shared_bytes, 'probes': [[line, reg] for line in sorted(probes) for reg in probes[line]], 'oob': ns.oob_report}) + '\n') sys.stderr.write(f'block={block} textures={sorted(t.textures)} surfaces={t.surfaces} shared={t.shared_bytes}\n') diff --git a/engine/rr/rr2_m.comp b/engine/rr/rr2_m.comp new file mode 100644 index 0000000..4b902a9 --- /dev/null +++ b/engine/rr/rr2_m.comp @@ -0,0 +1,530 @@ +#version 460 +// Native stage of the DLSS 4.5 Ray Reconstruction network (preset F, the rr2_* kernels of nvngx_dlssd.dll 310.9): +// one shader for enc1 ... enc5 and both kernels of dec4 ... dec1 and dec0's first. The arithmetic is +// engine/rr/rr2_ref.py's (which documents the block); the shape and the weight offsets come from the include +// generated by engine/rr/rr2_pack.py: +// KIND 0 encoder (tokens from the arena), 1 enc1 (first merges encoder 0's tokens 2x2), 2 a decoder's first kernel +// (expands the CL channels one level down and adds the skip), 3 a decoder's second kernel (with the guide blend), +// 4 encoder 0 after its per-pixel feature gather: included (RR2_EMBED) by the shader engine/rr/gen_rr2_enc0.py +// generates around ptx2rr.py's translation of the gather; it embeds the 64 feature slots of XF and leaves +// the guide head in HD for the translated tail; 5 dec0's second kernel, included likewise by the shader +// generated around the translation of its output tail: it leaves r and the MLP sum in RSH / MSH +// C channels (MT = C / 16 tiles), NH heads of 32 channels, NB MLP branches, CM merged channels for the next +// encoder (0: none), GUIDE 1 when the stage gates its guide image (KIND 0 and 1) +// Like engine/rr/swin_m.comp, every product runs transposed (weights x activations), so results stay in registers as +// [channel tile][token] accumulators: one workgroup per 8x8-token window, four waves, wave w owning the token slots +// 16w ... 16w + 15; token and channel slot j holds the natural index inter(j) (NVIDIA's k order, which is also the +// arena's byte order within a 32-channel plane). +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_float_e4m3 : require +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_shuffle : require +#extension GL_EXT_control_flow_attributes : require +#extension GL_GOOGLE_include_directive : require +#ifndef RR2_EMBED +layout(local_size_x = 128) in; +#endif +#include "rr2_stage.h" +#include "rr_exact.glsl" + +#define MT (C / 16) + +#ifdef RR2_EMBED +#define PRM(i) param.w[i] +// the generated shader's workgroup holds RR2_GROUP (1 to 4) windows along x (wave w runs the translated code of +// window w, then the block runs once per window): RR2_WIN is the one this call works on +#if RR2_GROUP == 1 +#define RR2_WIN 0u +#define RR2_WGX gl_WorkGroupID.x +#else +uint rr2_win; +#define RR2_WIN rr2_win +#define RR2_WGX (uint(RR2_GROUP) * gl_WorkGroupID.x + rr2_win) +#endif +#else +#define RR2_GROUP 1 +#define RR2_WIN 0u +#define RR2_WGX gl_WorkGroupID.x +layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param; +#define PRM(i) param.p[i] +#endif +layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8; +layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16; +layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u; +layout(set = 0, binding = 1, std430) readonly buffer W8U { uint8_t w[]; } w8u; +layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena; + +#define ACC coopmat +#define HAC coopmat +#define A8 coopmat +#define B8 coopmat +#define RM gl_CooperativeMatrixLayoutRowMajor +#define CMJ gl_CooperativeMatrixLayoutColumnMajor +#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM) +#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM) +#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]); +#define Q8(dst, src) Q8PAIRS(dst, src) +// One 32-deep step onto an accumulator of binary16 values: the reference adds it inside the f32 products and rounds +// once (swin_m.comp explains RR_MIX, the one-rounding variant that is not bit-exact by construction) +#ifdef RR_MIX +#define RAC HAC +#define RV(x) (x) +#define K32(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(0.0))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_] + float((acc)[i_])); } +#define K32H(acc, a0, b0, a1, b1) K32(acc, a0, b0, a1, b1) +#else +#define RAC ACC +#define RV(x) float(x) +#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) } +#define K32H(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(acc))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_]); } +#endif +const uint HEAD = W_WQK1 - W_WQK0, BRANCH = W_WD1 - W_WD0; + +shared uint XT[16 * C]; // the window's tokens: [token slot][C FP8 channels in slot order]; later the merge rows +shared uint KVS[2 * 512]; // the current head's QK ([token slot][32]), then V^T ([32][token slot]) +#if KIND == 1 +shared uint GT[64 * 32]; // the merge input: [token slot][4 sub-positions x 32 FP8 channels] +#endif +#if KIND == 2 +shared uint LOW[4 * CL]; // the 4 x 4 tokens one level down: [row-major][CL FP8 channels] +shared uint EX[16 * C]; // their expansion: [low token][sub-position][C FP8 channels in slot order] +shared float16_t XE[64 * C]; // the block's input before its FP8 conversion: [token slot][channel slot] +#endif +#if KIND == 4 +#ifdef RR2_XF_ARENA +// the features come from the gather pass (engine/rr/gen_rr2_enc0.py --split): the window's record in the arena, as +// the index of its first 32-bit word, set by the stage's main +uint rr2_xf; +#define XF_SRC arena.w +#define XF_AT rr2_xf +#else +shared uint XF[RR2_GROUP * 64 * 16]; // per window, the features from the gather: [token slot][64 FP8 feature slots] +#define XF_SRC XF +#define XF_AT (RR2_WIN * 1024u) +#endif +shared float16_t HD[RR2_GROUP * 64 * 4]; // per window, the guide head per token (raster order), for the translated tail +const uint EBRANCH = W_EWD1 - W_EWD0; +#endif +#if KIND == 5 +shared float16_t RSH[RR2_GROUP * 64 * 32], MSH[RR2_GROUP * 64 * 32]; // per window, r and the MLP sum per token (raster order) and natural channel +#endif +#if GUIDE || KIND == 3 +shared float16_t OUTG[64 * 4]; // the guide output per window token (raster order) +#endif + +uint pack4(floate4m3_t a, floate4m3_t b, floate4m3_t c, floate4m3_t d) +{ + return code(a) | (code(b) << 8u) | (code(c) << 16u) | (code(d) << 24u); +} +float16_t e4(uint w, uint b) { return float16_t(uintBitsToFloate4m3EXT(uint8_t((w >> (8u * b)) & 0xffu))); } +uvec2 mergeXY(uint r) { return r < 8u ? uvec2(r >> 1u, r & 1u) : uvec2((r - 8u) >> 1u, 2u + (r & 1u)); } +// four binary16 values of a guide image at token (x, y) of a level w wide +f16vec4 guideAt(uint off, int x, int y, int w) +{ + const uint a = (off + uint(y * w + x) * 8u) >> 2u; + const uint u0 = arena.w[a], u1 = arena.w[a + 1u]; + return f16vec4(hb(u0 & 0xffffu), hb(u0 >> 16u), hb(u1 & 0xffffu), hb(u1 >> 16u)); +} + +#ifdef RR2_EMBED +void rr2_block() +#else +void main() +#endif +{ + const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID; + const uint hl = lane >> 4u, col = lane & 15u; + RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u); + const int ex = int(PRM(0)), ey = int(PRM(1)), hx = ex / 2, hy = ey / 2; + const uint inOff = PRM(2), skOff = PRM(4), outOff = PRM(8), mrgOff = PRM(10), shift = PRM(20); + const uint gInOff = PRM(12), gEncOff = PRM(14), gOutOff = PRM(16), gHalfOff = PRM(18); + const float16_t s0 = hb(PRM(23) & 0xffffu), s1 = hb(PRM(23) >> 16u), s2 = hb(PRM(24) & 0xffffu), s3 = hb(PRM(24) >> 16u); + const int x0 = 8 * int(RR2_WGX) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u); + const uint plane = uint(ex * ey) * 32u, lplane = uint(hx * hy) * 32u; + // this lane's own token (slot 16 wv + col) + const uint t = inter(16u * wv + col); + const int tx = x0 + int(t & 7u), ty = y0 + int(t >> 3u); + const bool inside = tx >= 0 && ty >= 0 && tx < ex && ty < ey; + + // ---------------------------------------------------------------- the block's input x: one token per thread +#if KIND == 0 || KIND == 3 || KIND == 5 + if (tid < 64u) { + const uint k = inter(tid); + const int xx = mirror(x0 + int(k & 7u), ex), yy = mirror(y0 + int(k >> 3u), ey); + const uint tb = (inOff + uint(yy * ex + xx) * 32u) / 4u; + [[unroll]] for (uint wd = 0u; wd < C / 4; ++wd) XT[tid * (C / 4u) + wd] = arena.w[tb + (wd >> 3u) * (plane >> 2u) + (wd & 7u)]; + } + barrier(); +#elif KIND == 4 + // encoder 0: x = bu + sum over two branches of Wu . act(Wd . features + bd) + { + B8 fb[4]; + [[unroll]] for (uint kt = 0u; kt < 4u; ++kt) coopMatLoad(fb[kt], XF_SRC, XF_AT + wv * 256u + 4u * kt, 16u, CMJ); + RAC X[MT]; + [[unroll]] for (uint m = 0u; m < MT; ++m) + [[unroll]] for (int i = 0; i < 8; ++i) X[m][i] = RV(w16.w[W_EBU / 2u + 16u * m + 8u * hl + uint(i)]); + [[unroll]] for (uint br = 0u; br < 2u; ++br) { + B8 Hb[2]; + [[unroll]] for (uint c = 0u; c < 2u; ++c) { + RAC h; + [[unroll]] for (int i = 0; i < 8; ++i) h[i] = RV(w16.w[(W_EBD0 + EBRANCH * br) / 2u + 16u * c + 8u * hl + uint(i)]); + [[unroll]] for (uint kt = 0u; kt < 4u; kt += 2u) { + A8 a0, a1; + LDA(a0, W_EWD0 + EBRANCH * br, 64u, c, kt); LDA(a1, W_EWD0 + EBRANCH * br, 64u, c, kt + 1u); + K32(h, a0, fb[kt], a1, fb[kt + 1u]) + } + [[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = ACTQ(float16_t(h[i])); + } + [[unroll]] for (uint m = 0u; m < MT; ++m) { + A8 a0, a1; + LDA(a0, W_EWU0 + EBRANCH * br, 32u, m, 0u); LDA(a1, W_EWU0 + EBRANCH * br, 32u, m, 1u); + K32(X[m], a0, Hb[0], a1, Hb[1]) + } + } + [[unroll]] for (uint m = 0u; m < MT; ++m) { + B8 q; + Q8(q, X[m]) + coopMatStore(q, XT, wv * 4u * C + 4u * m, C / 4u, CMJ); + } + } + barrier(); +#elif KIND == 1 + // enc1: the 2x2 group of encoder 0's 32-channel tokens, then x = W1 . group + b1 + if (tid < 64u) { + const uint k = inter(tid); + const int xx = mirror(x0 + int(k & 7u), ex), yy = mirror(y0 + int(k >> 3u), ey); + [[unroll]] for (uint q = 0u; q < 4u; ++q) { + const uint tb = (inOff + uint((2 * yy + int(q >> 1u)) * (2 * ex) + 2 * xx + int(q & 1u)) * 32u) / 4u; + [[unroll]] for (uint wd = 0u; wd < 8u; ++wd) GT[tid * 32u + q * 8u + wd] = arena.w[tb + wd]; + } + } + barrier(); + { + B8 gb[8]; + [[unroll]] for (uint kt = 0u; kt < 8u; ++kt) coopMatLoad(gb[kt], GT, wv * 512u + 4u * kt, 32u, CMJ); + [[unroll]] for (uint m = 0u; m < MT; ++m) { + RAC a; + [[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_B1 / 2u + 16u * m + 8u * hl + uint(i)]); + [[unroll]] for (uint kt = 0u; kt < 8u; kt += 2u) { + A8 w0, w1; + LDA(w0, W_W1, 128u, m, kt); LDA(w1, W_W1, 128u, m, kt + 1u); + K32(a, w0, gb[kt], w1, gb[kt + 1u]) + } + B8 q; + Q8(q, a) + coopMatStore(q, XT, wv * 4u * C + 4u * m, C / 4u, CMJ); + } + } + barrier(); +#else + // a decoder's first kernel: wave w expands the 16 low tokens to sub-position w, then x = dq(e) + dq(skip) * s3 + { + const int X0 = x0 >> 1, Y0 = y0 >> 1; + for (uint i = tid; i < 4u * CL; i += 128u) { + const uint r = i / (CL / 4u), wd = i % (CL / 4u); + const int X = mirror(X0 + int(r & 3u), hx), Y = mirror(Y0 + int(r >> 2u), hy); + LOW[i] = arena.w[(inOff + (wd >> 3u) * lplane + uint(Y * hx + X) * 32u) / 4u + (wd & 7u)]; + } + barrier(); + B8 lb[CL / 16]; + [[unroll]] for (uint kt = 0u; kt < CL / 16; ++kt) coopMatLoad(lb[kt], LOW, 4u * kt, CL / 4u, CMJ); + for (uint mm = 0u; mm < MT; ++mm) { + const uint m = MT * wv + mm; + RAC a; + [[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_BE / 2u + 16u * m + 8u * hl + uint(i)]); + [[unroll]] for (uint kt = 0u; kt < CL / 16; kt += 2u) { + A8 w0, w1; + LDA(w0, W_WE, CL, m, kt); LDA(w1, W_WE, CL, m, kt + 1u); + K32(a, w0, lb[kt], w1, lb[kt + 1u]) + } + B8 q; + Q8(q, a) + coopMatStore(q, EX, 4u * m, C, CMJ); + } + barrier(); + if (tid < 64u) { + const uint k = inter(tid); + const int kx = x0 + int(k & 7u), ky = y0 + int(k >> 3u); + const uint tb = (skOff + uint(mirror(ky, ey) * ex + mirror(kx, ex)) * 32u) / 4u; + const uint r = uint(((ky >> 1) - Y0) * 4 + ((kx >> 1) - X0)), s = uint(2 * (ky & 1) + (kx & 1)); + [[unroll]] for (uint wd = 0u; wd < C / 4; ++wd) { + const uint sk = arena.w[tb + (wd >> 3u) * (plane >> 2u) + (wd & 7u)]; + const uint e = EX[r * C + s * (C / 4u) + wd]; + floate4m3_t c[4]; + [[unroll]] for (uint b = 0u; b < 4u; ++b) { + precise float16_t ms = e4(sk, b) * s3; + precise float16_t x = e4(e, b) + ms; + XE[tid * C + 4u * wd + b] = x; + c[b] = q8h(x); + } + XT[tid * (C / 4u) + wd] = pack4(c[0], c[1], c[2], c[3]); + } + } + barrier(); + } +#endif + + // ---------------------------------------------------------------- attention: the heads' out-projections sum up + RAC AT[MT]; + [[unroll]] for (uint m = 0u; m < MT; ++m) + [[unroll]] for (int i = 0; i < 8; ++i) AT[m][i] = RV(w16.w[W_BO / 2u + 16u * m + 8u * hl + uint(i)]); + [[dont_unroll]] for (uint hd = 0u; hd < NH; ++hd) { + barrier(); // every wave is done with the previous head's QK / V + for (uint j = 0u; j < 2u; ++j) { // 8 jobs (token tile n, channel tile c) + const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u; + RAC a = RAC(RV(float16_t(0.0))), v = RAC(RV(float16_t(0.0))); + [[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) { + B8 xb0, xb1, wv0, wv1; + A8 xa0, xa1, wk0, wk1; + coopMatLoad(xb0, XT, n * 4u * C + 4u * kt, C / 4u, CMJ); + coopMatLoad(xb1, XT, n * 4u * C + 4u * kt + 4u, C / 4u, CMJ); + coopMatLoad(xa0, XT, n * 4u * C + 4u * kt, C / 4u, RM); + coopMatLoad(xa1, XT, n * 4u * C + 4u * kt + 4u, C / 4u, RM); + LDA(wk0, W_WQK0 + HEAD * hd, C, c, kt); LDA(wk1, W_WQK0 + HEAD * hd, C, c, kt + 1u); + LDB(wv0, W_WV0 + HEAD * hd, 32u, kt, c); LDB(wv1, W_WV0 + HEAD * hd, 32u, kt + 1u, c); + K32(a, wk0, xb0, wk1, xb1) + K32(v, xa0, wv0, xa1, wv1) + } + A8 ka, va; + Q8(ka, a) + Q8(va, v) + coopMatStore(ka, KVS, n * 128u + 4u * c, 8u, RM); + coopMatStore(va, KVS, 512u + c * 256u + 4u * n, 16u, RM); + } + barrier(); + // scores^T (keys x this wave's queries) = QK . QK^T + bias, the cubic softmax down each column + B8 Qb[2]; + [[unroll]] for (uint c = 0u; c < 2u; ++c) coopMatLoad(Qb[c], KVS, wv * 128u + 4u * c, 8u, CMJ); + B8 Pb[4]; + { + float16_t e[4][8]; + [[unroll]] for (uint kt = 0u; kt < 4u; ++kt) { + HAC c; + coopMatLoad(c, w16.w, (W_BIAS0 + HEAD * hd) / 2u + kt * 1024u + 16u * wv, 64u, RM); + A8 k0, k1; + coopMatLoad(k0, KVS, kt * 128u, 8u, RM); + coopMatLoad(k1, KVS, kt * 128u + 4u, 8u, RM); + K32H(c, k0, Qb[0], k1, Qb[1]) + [[unroll]] for (int i = 0; i < 8; ++i) e[kt][i] = hb(uint(w16u.w[W_CUBIC / 2u + bh(c[i])])); + } + // key slot 16kt + 8hl + i is NVIDIA's column 8j + 2tt + half with j = 4(kt >> 1) + (kt & 1) + 2((i >> 1) & 1), + // tt = 2hl + (i >> 2), half = i & 1 (rr2_ref.softmax): pairs (e[2q], e[2q + 1]) are the two halves + f16vec2 ep[4][4]; + [[unroll]] for (uint kt = 0u; kt < 4u; ++kt) + [[unroll]] for (uint q = 0u; q < 4u; ++q) ep[kt][q] = f16vec2(e[kt][2u * q], e[kt][2u * q + 1u]); + precise f16vec2 st[2]; + [[unroll]] for (uint tq = 0u; tq < 2u; ++tq) { + const uint q0 = 2u * tq, q1 = q0 + 1u; + precise f16vec2 ev = (ep[0][q0] + ep[0][q1]) + (ep[2][q0] + ep[2][q1]); + precise f16vec2 od = (ep[1][q0] + ep[1][q1]) + (ep[3][q0] + ep[3][q1]); + st[tq] = ev + od; + } + precise f16vec2 loc = st[0] + st[1]; + const uint other = subgroupShuffleXor(packFloat2x16(loc), 16u); + // hl = 0 holds tt 0, 1 and hl = 1 tt 2, 3: the sum is (tt0 + tt1) + (tt2 + tt3) on both halves + precise f16vec2 o = hl == 0u ? loc + unpackFloat2x16(other) : unpackFloat2x16(other) + loc; + precise float16_t rs = o.x + o.y; + const float16_t r = float16_t(1.0 / float(rs)); + const f16vec2 r2 = f16vec2(r, r); + [[unroll]] for (uint kt = 0u; kt < 4u; ++kt) + [[unroll]] for (uint q = 0u; q < 4u; ++q) { + precise f16vec2 pv = ep[kt][q] * r2; + { const fe4m3vec2 t_ = fe4m3vec2(pv); Pb[kt][2u * q] = t_.x; Pb[kt][2u * q + 1u] = t_.y; } + } + } + // O^T = V^T P^T, then the out-projection + B8 Ob[2]; + [[unroll]] for (uint c = 0u; c < 2u; ++c) { + RAC o = RAC(RV(float16_t(0.0))); + [[unroll]] for (uint ks = 0u; ks < 2u; ++ks) { + A8 v0, v1; + coopMatLoad(v0, KVS, 512u + c * 256u + 8u * ks, 16u, RM); + coopMatLoad(v1, KVS, 512u + c * 256u + 8u * ks + 4u, 16u, RM); + K32(o, v0, Pb[2u * ks], v1, Pb[2u * ks + 1u]) + } + Q8(Ob[c], o) + } + [[unroll]] for (uint m = 0u; m < MT; ++m) { + A8 a0, a1; + LDA(a0, W_WO0 + HEAD * hd, 32u, m, 0u); LDA(a1, W_WO0 + HEAD * hd, 32u, m, 1u); + K32(AT[m], a0, Ob[0], a1, Ob[1]) + } + } + + // ---------------------------------------------------------------- r = x + attention * s1; the MLP's input q8(r) + float16_t rv[MT][8]; + B8 Nb[MT]; + [[unroll]] for (uint m = 0u; m < MT; ++m) { +#if KIND != 2 + const uint xa = (16u * wv + col) * (C / 4u) + 4u * m + 2u * hl; + const uvec2 xw = uvec2(XT[xa], XT[xa + 1u]); +#endif + [[unroll]] for (int i = 0; i < 8; ++i) { +#if KIND == 2 + const float16_t xb = XE[(16u * wv + col) * C + 16u * m + 8u * hl + uint(i)]; +#else + const float16_t xb = e4(xw[i >> 2], uint(i) & 3u); +#endif + precise float16_t sa = float16_t(AT[m][i]) * s1; + precise float16_t r = xb + sa; + rv[m][i] = r; + Nb[m][i] = q8h(r); + } + } + RAC ML[MT]; + [[unroll]] for (uint m = 0u; m < MT; ++m) + [[unroll]] for (int i = 0; i < 8; ++i) ML[m][i] = RV(w16.w[W_BU / 2u + 16u * m + 8u * hl + uint(i)]); + [[dont_unroll]] for (uint br = 0u; br < NB; ++br) { + const uint w1 = W_WD0 + BRANCH * br, b1 = W_BD0 + BRANCH * br, w2 = W_WU0 + BRANCH * br; + B8 Hb[2]; + [[unroll]] for (uint c = 0u; c < 2u; ++c) { + RAC h; + [[unroll]] for (int i = 0; i < 8; ++i) h[i] = RV(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]); + [[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) { + A8 a0, a1; + LDA(a0, w1, C, c, kt); LDA(a1, w1, C, c, kt + 1u); + K32(h, a0, Nb[kt], a1, Nb[kt + 1u]) + } + [[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = ACTQ(float16_t(h[i])); + } + [[unroll]] for (uint m = 0u; m < MT; ++m) { + A8 a0, a1; + LDA(a0, w2, 32u, m, 0u); LDA(a1, w2, 32u, m, 1u); + K32(ML[m], a0, Hb[0], a1, Hb[1]) + } + } + +#if KIND == 5 + // dec0's second kernel: the translated tail forms y = r + MLP * s2 itself and computes the outputs from it + [[unroll]] for (uint m = 0u; m < MT; ++m) + [[unroll]] for (int i = 0; i < 8; ++i) { + const uint at = RR2_WIN * 2048u + t * 32u + inter(16u * m + 8u * hl + uint(i)); + RSH[at] = rv[m][i]; + MSH[at] = float16_t(ML[m][i]); + } +#else + // ---------------------------------------------------------------- y = r + MLP * s2: the stage's tokens + B8 Tb[MT]; + { + const uint o = (outOff + uint(ty * ex + tx) * 32u) >> 2u; +#if CM + const uint ty2 = t >> 3u, tx2 = t & 7u; + const uint mr = (ty2 >> 1u) < 2u ? 2u * (tx2 >> 1u) + (ty2 >> 1u) : 8u + 2u * (tx2 >> 1u) + (ty2 >> 1u) - 2u; + const uint blk = 2u * (ty2 & 1u) + (tx2 & 1u); + barrier(); // every wave is done with XT (the merge rows reuse it) +#endif + [[unroll]] for (uint m = 0u; m < MT; ++m) { + [[unroll]] for (int i = 0; i < 8; ++i) { + precise float16_t sm = float16_t(ML[m][i]) * s2; + precise float16_t y = rv[m][i] + sm; + Tb[m][i] = q8h(y); + } + const uint w0 = pack4(Tb[m][0], Tb[m][1], Tb[m][2], Tb[m][3]), w1 = pack4(Tb[m][4], Tb[m][5], Tb[m][6], Tb[m][7]); + const uint at = 4u * (m & 1u) + 2u * hl; + if (inside) { arena.w[o + (m >> 1u) * (plane >> 2u) + at] = w0; arena.w[o + (m >> 1u) * (plane >> 2u) + at + 1u] = w1; } +#if CM + const uint g = mr * C + blk * (C / 4u) + 4u * m + 2u * hl; + XT[g] = w0; XT[g + 1u] = w1; +#endif + } + } +#endif +#if GUIDE || KIND == 3 || KIND == 4 + { // the guide head (rows 0-3 on lanes 0-15) + RAC hv; + [[unroll]] for (int i = 0; i < 8; ++i) hv[i] = RV(w16.w[W_BH / 2u + 8u * hl + uint(i)]); + [[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) { + A8 a0, a1; + LDA(a0, W_WH, C, 0u, kt); LDA(a1, W_WH, C, 0u, kt + 1u); + K32(hv, a0, Tb[kt], a1, Tb[kt + 1u]) + } +#if KIND == 4 + if (hl == 0u) [[unroll]] for (uint c = 0u; c < 4u; ++c) HD[RR2_WIN * 256u + t * 4u + c] = float16_t(hv[c]); +#else + if (hl == 0u) { + f16vec4 o4 = f16vec4(0.0); + if (inside) { + precise f16vec3 gate = f16vec3(sig(float16_t(hv[0])), sig(float16_t(hv[1])), sig(float16_t(hv[2]))) * s0; +#if KIND == 3 + // the guide one level down, bilinear (half-pixel centres, edge clamped), every product and sum rounded + const int bx = tx == 0 ? 0 : (tx - 1) >> 1, by = ty == 0 ? 0 : (ty - 1) >> 1; + const float16_t wx = tx == 0 ? float16_t(0.0) : ((tx & 1) != 0 ? float16_t(0.25) : float16_t(0.75)); + const float16_t wy = ty == 0 ? float16_t(0.0) : ((ty & 1) != 0 ? float16_t(0.25) : float16_t(0.75)); + const int x1 = min(bx + 1, hx - 1), y1 = min(by + 1, hy - 1); + precise float16_t ux = float16_t(1.0) - wx, uy = float16_t(1.0) - wy; + precise f16vec4 p00 = guideAt(gInOff, bx, by, hx) * ux, p10 = guideAt(gInOff, x1, by, hx) * wx; + precise f16vec4 p01 = guideAt(gInOff, bx, y1, hx) * ux, p11 = guideAt(gInOff, x1, y1, hx) * wx; + precise f16vec4 a00 = p00 * uy, a10 = p10 * uy, a01 = p01 * wy, a11 = p11 * wy; + precise f16vec4 u0 = a00 + a10; + precise f16vec4 u1 = u0 + a01; + precise f16vec4 up = u1 + a11; + const float16_t al = sig(float16_t(hv[3])); + precise f16vec4 gt = up * f16vec4(gate, float16_t(0.0)); + precise float16_t ia = float16_t(1.0) - al; + precise f16vec4 pe = guideAt(gEncOff, tx, ty, ex) * ia; + precise f16vec4 pg = gt * al; + precise f16vec4 mx = pe + pg; + o4 = min(mx, f16vec4(65504.0)); +#else + precise f16vec4 gg = guideAt(gInOff, tx, ty, ex) * f16vec4(gate, float16_t(0.0)); + o4 = min(gg, f16vec4(65504.0)); +#endif + const uint a = (gOutOff + uint(ty * ex + tx) * 8u) >> 2u; + arena.w[a] = bh(o4.x) | (bh(o4.y) << 16u); + arena.w[a + 1u] = bh(o4.z) | (bh(o4.w) << 16u); + } + [[unroll]] for (uint c = 0u; c < 4u; ++c) OUTG[t * 4u + c] = o4[c]; + } +#endif + } +#endif +#if CM + barrier(); + { // the 2x2 merge for the next encoder: merged^T (CM x 16 rows), channel tiles m = wv, wv + 4, ... + const uvec2 mxy = mergeXY(col); + const int X = x0 / 2 + int(mxy.x), Y = y0 / 2 + int(mxy.y); + const uint o = (mrgOff + uint(Y * hx + X) * 32u) >> 2u; + for (uint m = wv; m < CM / 16u; m += 4u) { + RAC a; + [[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_BM / 2u + 16u * m + 8u * hl + uint(i)]); + for (uint ks = 0u; ks < C / 8u; ++ks) { + B8 g0, g1; A8 w0, w1; + coopMatLoad(g0, XT, 8u * ks, C, CMJ); + coopMatLoad(g1, XT, 8u * ks + 4u, C, CMJ); + LDA(w0, W_WM, 4u * C, m, 2u * ks); LDA(w1, W_WM, 4u * C, m, 2u * ks + 1u); + K32(a, w0, g0, w1, g1) + } + B8 q; + Q8(q, a) + if (X >= 0 && Y >= 0 && X < hx && Y < hy) { + const uint at = o + (m >> 1u) * (lplane >> 2u) + 4u * (m & 1u) + 2u * hl; + arena.w[at] = pack4(q[0], q[1], q[2], q[3]); + arena.w[at + 1u] = pack4(q[4], q[5], q[6], q[7]); + } + } +#if GUIDE + if (wv == 3u && lane < 16u) { // the guide for the next encoder: 2x2 mean, rightwards first, then downwards + const uint mx = lane & 3u, my = lane >> 2u; + const int x = x0 + 2 * int(mx), y = y0 + 2 * int(my); + if (x >= 0 && y >= 0 && x < ex && y < ey) { + const uint t00 = (2u * my) * 8u + 2u * mx; + f16vec4 q4[4]; + [[unroll]] for (uint d = 0u; d < 4u; ++d) { + const uint tt = t00 + (d & 1u) + 8u * (d >> 1u); + precise f16vec4 qv = f16vec4(OUTG[tt * 4u], OUTG[tt * 4u + 1u], OUTG[tt * 4u + 2u], OUTG[tt * 4u + 3u]) * float16_t(0.25); + q4[d] = qv; + } + precise f16vec4 h0 = q4[0] + q4[1], h1 = q4[2] + q4[3]; + precise f16vec4 hm = h0 + h1; + const uint a = (gHalfOff + uint((y / 2) * hx + x / 2) * 8u) >> 2u; + arena.w[a] = bh(hm.x) | (bh(hm.y) << 16u); + arena.w[a + 1u] = bh(hm.z) | (bh(hm.w) << 16u); + } + } +#endif + } +#endif +} diff --git a/engine/rr/rr2_pack.py b/engine/rr/rr2_pack.py new file mode 100644 index 0000000..f953b7c --- /dev/null +++ b/engine/rr/rr2_pack.py @@ -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() diff --git a/engine/rr/rr2_ref.py b/engine/rr/rr2_ref.py new file mode 100644 index 0000000..d2024f1 --- /dev/null +++ b/engine/rr/rr2_ref.py @@ -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() diff --git a/engine/rr/rr_attn.glsl b/engine/rr/rr_attn.glsl index 99797ec..e235510 100644 --- a/engine/rr/rr_attn.glsl +++ b/engine/rr/rr_attn.glsl @@ -21,13 +21,12 @@ [[unroll]] for (uint kt = 0u; kt < 6u; ++kt) { HAC c; coopMatLoad(c, w16.w, (W_BIAS0 + HEAD * hd) / 2u + kt * 1024u + 16u * wv, 64u, RM); - ACC s = ACC(c); A8 k0, k1; coopMatLoad(k0, RR_KS, RR_KS0 + kt * 128u, 8u, RM); coopMatLoad(k1, RR_KS, RR_KS0 + kt * 128u + 4u, 8u, RM); - s = coopMatMulAdd(k1, Qb[1], coopMatMulAdd(k0, Qb[0], s)); // rounded once, by the conversion below + K32H(c, k0, Qb[0], k1, Qb[1]) // the cubic numerator from enc1_pack.cubic_table, indexed by the score's binary16 bits - [[unroll]] for (int i = 0; i < 8; ++i) e[kt][i] = hb(uint(w16u.w[W_CUBIC / 2u + bh(float16_t(s[i]))])); + [[unroll]] for (int i = 0; i < 8; ++i) e[kt][i] = hb(uint(w16u.w[W_CUBIC / 2u + bh(c[i])])); } // key slot 16kt + 8hl + i is NVIDIA's column 8j + 2tt + half with j = 4(kt >> 1) + (kt & 1) + 2((i >> 1) & 1), // tt = 2hl + (i >> 2), half = i & 1 (enc1_ref.softmax_cubic's tree) @@ -60,7 +59,7 @@ // O^T = V^T P^T, then the out-projection onto the residual B8 Ob[2]; [[unroll]] for (uint c = 0u; c < 2u; ++c) { - ACC o = ACC(0.0); + RAC o = RAC(RV(float16_t(0.0))); [[unroll]] for (uint ks = 0u; ks < 3u; ++ks) { A8 v0, v1; coopMatLoad(v0, RR_VS, RR_VS0 + c * 384u + 8u * ks, 24u, RM); diff --git a/engine/rr/rr_capture.py b/engine/rr/rr_capture.py index b0d827a..da147ec 100644 --- a/engine/rr/rr_capture.py +++ b/engine/rr/rr_capture.py @@ -70,16 +70,20 @@ class Launch: def image(self, offset): return next(i for i in self.images if i.offset == offset) + def weights_at(self, offset): + """the weight image (the pointer at argument OFFSET) as NGX uploaded it; 16 zero bytes for a stage without one""" + pointer = next((p for p in self.pointers if p.offset == offset), None) + return self.pre(pointer.alloc) if pointer is not None else bytes(16) + @property def weights(self): - """the weight image (argument 40) as NGX uploaded it; 16 zero bytes for a stage without one (the output stage)""" - pointer = next((p for p in self.pointers if p.offset == 40), None) - return self.pre(pointer.alloc) if pointer is not None else bytes(16) + """the DLSS-D kernels' weight image (argument 40)""" + return self.weights_at(40) @property def arena_alloc(self): """the plane arena: the allocation every arena pointer of the layer points into""" - candidates = [p.alloc for p in self.pointers if p.offset != 40] + candidates = [p.alloc for p in self.pointers] return max(set(candidates), key=lambda a: self.allocs[a][1]) if candidates else None def post_buffer(self, capture, alloc): diff --git a/engine/rr/rr_inputs.py b/engine/rr/rr_inputs.py index d48c209..8e48759 100644 --- a/engine/rr/rr_inputs.py +++ b/engine/rr/rr_inputs.py @@ -25,6 +25,45 @@ OPTIONAL_OUTPUT = {'particles': 416} # output stage INPUTS = {'color': 160, 'motion': 168, 'normals': 184, 'roughness': 192, 'diffuse': 232, 'specular': 240, 'depth': 512} +# preset F (rr2_enc0_kernel): its textures and argument offsets +INPUTS_F = {'color': 144, 'motion': 192, 'normals': 224, 'roughness': 240, 'diffuse': 160, 'specular': 176, 'depth': 552} + + +def preset_f(cap, o): + lines = [] + for f in range(1, o.frames + 1): + enc0 = cap.find('rr2_enc0', f) + a = enc0.args + u32 = lambda off: struct.unpack_from('> 3u); @@ -235,7 +251,7 @@ void main() const float16_t x = e4(t[i >> 2], uint(i) & 3u); #endif precise float16_t c = x + w16.w[W_TABLE7 / 2u + 16u * m + 8u * hl + uint(i)]; - R[m][i] = float(c); + R[m][i] = RV(c); } } } @@ -245,7 +261,7 @@ void main() barrier(); // every wave is done with the previous head's K / V (and, for the first, the expand's LOW) for (uint j = 0u; j < 3u; ++j) { // 12 jobs (token tile n, channel tile c) const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u; - ACC a = ACC(0.0), v = ACC(0.0); + RAC a = RAC(RV(float16_t(0.0))), v = RAC(RV(float16_t(0.0))); [[unroll]] for (uint kt = 0u; kt < 2u * KC; kt += 2u) { B8 xb0, xb1, wv0, wv1; A8 xa0, xa1, wk0, wk1; @@ -312,15 +328,15 @@ void main() precise float16_t n = rv[m][i] * g; Nb[m][i] = q8h(n); precise float16_t a = rv[m][i] + w16.w[W_TABLE16 / 2u + k]; - R[m][i] = float(a); + R[m][i] = RV(a); } } [[dont_unroll]] for (uint br = 0u; br < NB; ++br) { const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br; B8 Hb[2]; [[unroll]] for (uint c = 0u; c < 2u; ++c) { - ACC h; - [[unroll]] for (int i = 0; i < 8; ++i) h[i] = float(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]); + RAC h; + [[unroll]] for (int i = 0; i < 8; ++i) h[i] = RV(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]); [[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) { A8 a0, a1; LDA(a0, w1, C, c, kt); LDA(a1, w1, C, c, kt + 1u); @@ -360,8 +376,8 @@ void main() } } // the C -> 16 head - ACC hv; - [[unroll]] for (int i = 0; i < 8; ++i) hv[i] = float(w16.w[W_BH / 2u + 8u * hl + uint(i)]); + RAC hv; + [[unroll]] for (int i = 0; i < 8; ++i) hv[i] = RV(w16.w[W_BH / 2u + 8u * hl + uint(i)]); [[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) { A8 a0, a1; LDA(a0, W_WH, C, 0u, kt); LDA(a1, W_WH, C, 0u, kt + 1u); @@ -427,8 +443,8 @@ void main() const int X = x0 / 2 + int(mxy.x), Y = y0 / 2 + int(mxy.y); const uint o = (mrgOff + uint(Y * hx + X) * 32u) >> 2u; for (uint m = wv; m < CX / 16u; m += 4u) { - ACC a; - [[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_BM / 2u + 16u * m + 8u * hl + uint(i)]); + RAC a; + [[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_BM / 2u + 16u * m + 8u * hl + uint(i)]); for (uint ks = 0u; ks < C / 8u; ++ks) { B8 g0, g1; A8 w0, w1; coopMatLoad(g0, MG, 8u * ks, C, CM); diff --git a/engine/rr/test_layer.py b/engine/rr/test_layer.py index 197200d..442767e 100755 --- a/engine/rr/test_layer.py +++ b/engine/rr/test_layer.py @@ -34,11 +34,11 @@ from rr_capture import Capture # noqa: E402 SURFACE_FORMATS = {'rgba16f': (4, 16), 'rg16f': (2, 16), 'r16f': (1, 16), 'rgba8_snorm': (4, 9), 'rgba8': (4, 8)} -def stage_params(launch): +def stage_params(launch, weights_arg=40): args = bytearray(launch.args) arena = launch.arena_alloc for p in launch.pointers: - value = p.delta if p.alloc == arena and p.offset != 40 else 0 + value = p.delta if p.alloc == arena and p.offset != weights_arg else 0 struct.pack_into('& p, uint32_t offset, const Mat4& m) { for (int i = 0; i < 16; ++i) putF(p, offset + 4 * i, static_cast(m[i])); } } // namespace +#include "runtime_rr2.inc" struct RModel::Impl { Device device; uint32_t preset = 0; + std::unique_ptr v2; // preset F ("D4RR0002"): engine/runtime_rr2.inc // per stage, the shader's bindings after the three buffers (engine/rr/README.md): texture argument offsets, then // surface argument offsets with their GLSL image format, both ascending std::array, RStages> textures; @@ -95,7 +97,14 @@ RModel::RModel(VkPhysicalDevice physical, VkDevice device, const std::string& di auto text = read(directory + "/rr.txt"); std::istringstream lines(std::string(text.begin(), text.end())); std::string line; - if (!std::getline(lines, line) || line.rfind("D4RR0001", 0) != 0) throw std::runtime_error("not a Ray Reconstruction model (rr.txt D4RR0001)"); + if (!std::getline(lines, line)) throw std::runtime_error("not a Ray Reconstruction model (rr.txt)"); + if (line.rfind("D4RR0002", 0) == 0) { + m.preset = 6; + m.v2 = std::make_unique(); + m.v2->load(m.device, directory, lines, cache); + return; + } + if (line.rfind("D4RR0001", 0) != 0) throw std::runtime_error("not a Ray Reconstruction model (rr.txt D4RR0001)"); unsigned stage = 0; const auto b = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, t = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, s = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE; auto split = [](const std::string& list) { @@ -153,6 +162,12 @@ uint32_t RModel::preset() const { return impl->preset; } void RModel::recordUpload(VkCommandBuffer cb) { auto& m = *impl; if (m.uploaded) throw std::runtime_error("model upload already recorded"); + if (m.v2) { + m.v2->recordUpload(cb); + memoryBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT); + m.uploaded = true; + return; + } for (unsigned i = 0; i < RStages; ++i) { VkBufferCopy copy{0, 0, m.uploads[i].size}; vkCmdCopyBuffer(cb, m.uploads[i].handle, m.weights[i].handle, 1, ©); @@ -163,6 +178,7 @@ void RModel::recordUpload(VkCommandBuffer cb) { struct REngine::Impl { RModel::Impl& model; + std::unique_ptr v2; // preset F uint32_t rw, rh, ow, oh, w0, h0; Buffer arena, expP, expSums, expState; std::array params; @@ -266,6 +282,12 @@ REngine::REngine(RModel& model, uint32_t renderWidth, uint32_t renderHeight, uin : impl(std::make_unique(*model.impl)) { auto& e = *impl; const Device& d = e.model.device; + if (e.model.v2) { + if (renderWidth < 64 || renderHeight < 64 || renderWidth > outputWidth || renderHeight > outputHeight || outputWidth > 16384 || outputHeight > 16384) + throw std::runtime_error("unsupported Ray Reconstruction dimensions"); + e.v2 = std::make_unique(d, *e.model.v2, renderWidth, renderHeight, outputWidth, outputHeight); + return; + } if (renderWidth < 64 || renderHeight < 64 || renderWidth > outputWidth || renderHeight > outputHeight || outputWidth > 16384 || outputHeight > 16384) throw std::runtime_error("unsupported Ray Reconstruction dimensions"); e.rw = renderWidth; e.rh = renderHeight; e.ow = outputWidth; e.oh = outputHeight; @@ -326,9 +348,10 @@ REngine::REngine(RModel& model, uint32_t renderWidth, uint32_t renderHeight, uin } } REngine::~REngine() = default; -VkImage REngine::outputImage() const { return impl->out.handle; } +VkImage REngine::outputImage() const { return impl->v2 ? impl->v2->out.handle : impl->out.handle; } std::vector REngine::debugImages() const { const auto& e = *impl; + if (e.v2) return e.v2->debugImages(); std::vector list; const VkFormat half4 = VK_FORMAT_R16G16B16A16_SFLOAT; for (int l = 0; l < 6; ++l) { @@ -350,6 +373,7 @@ std::vector REngine::debugImages() const { void REngine::recordInitialize(VkCommandBuffer cb) { auto& e = *impl; + if (e.v2) { e.v2->recordInitialize(cb); return; } std::vector images; for (auto* group : {&e.s112, &e.s120}) for (auto& t : *group) images.push_back(t.handle); for (auto* group : {&e.s128, &e.d112, &e.d120}) for (auto& t : *group) images.push_back(t.handle); @@ -373,6 +397,7 @@ void REngine::recordInitialize(VkCommandBuffer cb) { void REngine::setFrame(const RFrameImages& images, const RFrameSettings& settings, FrameOutput finalOutput) { auto& e = *impl; + if (e.v2) { e.v2->setFrame(images, settings, finalOutput); return; } const Device& d = e.model.device; if (!images.color || !images.motion || !images.depth || !images.normals || !images.roughness || !images.diffuseAlbedo || !images.specularAlbedo) throw std::runtime_error("Ray Reconstruction needs colour, motion, depth and the four guides"); @@ -524,6 +549,7 @@ void REngine::setFrame(const RFrameImages& images, const RFrameSettings& setting void REngine::recordFrame(VkCommandBuffer cb, VkQueryPool stamps, uint32_t first) { auto& e = *impl; + if (e.v2) { e.v2->recordFrame(cb, stamps, first); return; } if (!e.bound) throw std::runtime_error("REngine::setFrame was not called"); const auto& b = *e.bound; const uint32_t f = e.frame, cur = f & 1u; diff --git a/engine/runtime_rr2.inc b/engine/runtime_rr2.inc new file mode 100644 index 0000000..f3579ff --- /dev/null +++ b/engine/runtime_rr2.inc @@ -0,0 +1,643 @@ +// Ray Reconstruction preset F (DLSS 4.5, the rr2_* kernels of nvngx_dlssd.dll 310.9): the runtime around the stage +// shaders of a "D4RR0002" model (engine/compile_rr2.py). Included by runtime_rr.inc after its helpers. +// +// NGX's frame: the exposure, encoder 0 to 5, five decoders of two kernels each, then the output stage, which filters +// an intermediate image of 1.5 times the output size in tiles (post) and reduces each tile to the output (down), and +// a quarter-size copy of the half-size history for the next frame's encoder 0 (pooled). Every stage's parameter +// block is NVIDIA's argument block: the model carries one captured block per kernel and this file fills the fields +// that depend on the frame or the geometry (offsets as observed on harness captures at three sizes, 2026-10-13). +// The network's activations live in one arena whose regions are NGX's, in units of W0 * H0 / 32 bytes (they scale +// with the render size); the output passes' two tile regions follow it. NGX runs post and down concurrently and +// has them wait on progress counters; here they run in order and the counter pointers stay null, which the kernels +// test. +namespace { +// R2DownP: the downsample kernel NGX runs instead of R2Down when the game passes DLSSD.ColorBeforeParticles +// (rr2_downsample_kernel_static_hdr); optional in a model +// R2Gather: encoder 0's feature gather as a pass of its own, before R2Enc0 (engine/rr/gen_rr2_enc0.py --split); optional +constexpr unsigned R2Network = 16, R2Post = 16, R2Down = 17, R2Pooled = 18, R2DownP = 19, R2Gather = 20, R2Stages = 21; +constexpr unsigned R2Enc0 = 0, R2Enc5 = 5, R2Dec0B = 15; +constexpr const char* R2Names[R2Stages] = {"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", "post", "down", "pooled", "downp", "enc0g"}; +constexpr int R2Level[R2Network] = {0, 1, 2, 3, 4, 5, 4, 4, 3, 3, 2, 2, 1, 1, 0, 0}; +constexpr uint32_t R2ArenaUnits = 3152; +// The window shift of a network stage (x | y << 16, argument 80): it alternates every frame +uint32_t r2Shift(unsigned stage, uint32_t frame) { + const bool even = (frame & 1u) == 0; + if (stage < R2Enc5) return even ? 0x00060002u : 0x00020006u; + const bool first = stage == R2Enc5 || ((stage - R2Enc5) & 1u) == 1; // enc5 and the decoders' kernel 0 + return (even == first) ? 0x00000004u : 0x00040000u; +} +struct R2Stage { + Pipeline pipeline; + Buffer weights, upload; + std::vector args; + std::vector textures, buffers; + std::vector> surfaces; + std::vector> pointers; // argument offset, arena offset in units + uint32_t group = 1; // "xN": a workgroup holds N windows along x (the stages generated around translated code) + uint32_t record = 0; // R2Gather, "rN": 32-bit words per window in the arena +}; +struct R2Model { + std::array stages; + Buffer aux, auxUpload; // encoder 0's second table (argument 104) + std::array exposure; + bool particles = false; // the model has R2DownP + bool gather = false; // the model has R2Gather + void load(const Device& d, const std::string& directory, std::istringstream& lines, VkPipelineCache cache) { + const auto b = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, t = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, s = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE; + auto split = [](const std::string& list) { + std::vector items; + if (list == "-") return items; + std::istringstream in(list); + for (std::string item; std::getline(in, item, ',');) items.push_back(item); + return items; + }; + const VkBufferUsageFlags usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT; + auto blob = [&](Buffer& buffer, Buffer& upload, const std::vector& bytes) { + buffer.create(d, bytes.size(), usage); + upload.create(d, bytes.size(), VK_BUFFER_USAGE_TRANSFER_SRC_BIT, true); + std::memcpy(upload.mapped, bytes.data(), bytes.size()); + }; + std::array seen{}; + std::string line; + while (std::getline(lines, line)) { + std::istringstream f(line); + std::string key; f >> key; + if (key != "stage") continue; + // "stage NAME SHADER.spv SUBGROUP WEIGHTS|- ARGS TEXTURES SURFACES BUFFERS POINTERS": the lists are + // comma-separated argument offsets ("-": none), surfaces OFFSET:FORMAT, pointers OFFSET:UNITS; an optional "xN" + std::string name, shader, weightFile, argsFile, textureList, surfaceList, bufferList, pointerList, option; uint32_t subgroup = 32; + f >> name >> shader >> subgroup >> weightFile >> argsFile >> textureList >> surfaceList >> bufferList >> pointerList >> option; + unsigned stage = 0; + while (stage < R2Stages && name != R2Names[stage]) ++stage; + if (stage >= R2Stages || seen[stage] || pointerList.empty()) + throw std::runtime_error("rr.txt: unknown, repeated or invalid stage: " + line); + seen[stage] = true; + auto& st = stages[stage]; + if (option.size() == 2 && option[0] == 'x' && option[1] >= '1' && option[1] <= '4') st.group = static_cast(option[1] - '0'); + if (option.size() > 1 && option[0] == 'r') st.record = static_cast(std::stoul(option.substr(1))); + for (const auto& item : split(textureList)) st.textures.push_back(static_cast(std::stoul(item))); + for (const auto& item : split(bufferList)) st.buffers.push_back(static_cast(std::stoul(item))); + for (const auto& item : split(surfaceList)) { + const auto colon = item.find(':'); + st.surfaces.emplace_back(static_cast(std::stoul(item.substr(0, colon))), item.substr(colon + 1)); + } + for (const auto& item : split(pointerList)) { + const auto colon = item.find(':'); + st.pointers.emplace_back(static_cast(std::stoul(item.substr(0, colon))), std::stod(item.substr(colon + 1))); + } + std::vector types{b, b, b}; + types.insert(types.end(), st.textures.size(), t); + types.insert(types.end(), st.surfaces.size(), s); + types.insert(types.end(), st.buffers.size(), b); + st.pipeline.create(d.handle, directory + "/" + shader, types, subgroup, cache); + { const auto bytes = read(directory + "/" + argsFile); st.args.assign(bytes.begin(), bytes.end()); } + blob(st.weights, st.upload, weightFile == "-" ? std::vector(16) : read(directory + "/" + weightFile)); + } + for (unsigned stage = 0; stage < R2DownP; ++stage) + if (!seen[stage]) throw std::runtime_error("rr.txt: incomplete Ray Reconstruction preset F model"); + particles = seen[R2DownP]; + gather = seen[R2Gather] && stages[R2Gather].record != 0 && stages[R2Enc0].args.size() >= 8; + for (unsigned stage : {R2DownP, R2Gather}) + if (!seen[stage]) blob(stages[stage].weights, stages[stage].upload, std::vector(16)); + blob(aux, auxUpload, read(directory + "/enc0-aux.bin")); + exposure[0].create(d.handle, directory + "/exposure_f0.spv", {b, t, b}, 64, cache); + exposure[1].create(d.handle, directory + "/exposure_f1.spv", {b, b, b, t, s}, 64, cache); + } + void recordUpload(VkCommandBuffer cb) { + for (auto& st : stages) { VkBufferCopy copy{0, 0, st.upload.size}; vkCmdCopyBuffer(cb, st.upload.handle, st.weights.handle, 1, ©); } + VkBufferCopy copy{0, 0, auxUpload.size}; + vkCmdCopyBuffer(cb, auxUpload.handle, aux.handle, 1, ©); + } +}; +struct R2ExposureParams { int32_t last[2]; uint32_t blocksX, blocks; float invCount; uint32_t reset; }; + +struct R2Engine { + const Device& d; + R2Model& model; + uint32_t rw, rh, ow, oh, w0, h0; + uint32_t iw, ih, cw, ch, qw, qh; // the intermediate (1.5x output), its half and quarter + uint32_t nx, ny, maxW, maxH; // output tiles and the largest tile rectangle (with its overlap) + uint64_t unit, tileA, tileB; // arena unit; the first tile's two regions; a tile slot's bytes + uint64_t tileBytes = 0; + // Tile slots (a pair of regions each) and tiles per group. The post passes of a group have a slot each and run + // without waiting for one another, as do the down passes that read them; a group's down passes also overlap the + // next group's post passes, which write the other half of the slots (NGX overlaps one post with one down the same + // way, on two slots). Serialising all 32 passes of a 4K frame cost 0.17 ms in the gaps (RX 9070 XT, 2026-10-13). + static constexpr uint32_t R2TileSlots = 8; + uint32_t tileGroup = 1; + uint64_t gatherAt = 0; // the gather pass's records (R2Gather): after the tile regions + struct Rect { uint32_t x0, y0, x1, y1; }; + std::vector src, dst; // per tile: intermediate rectangle (with overlap), output rectangle + std::vector downPAt; + std::vector> downPBlocks; + // the particle layer's output-size history (R2DownP's surface 264, read back as texture 256): created on first use + Tex pb[2]; + bool pbNew = false; + Buffer arena, values, sums, dummy; + // every parameter block of the frame in one buffer (one upload per frame): the network stages', per tile the two + // output passes', the pooling pass's and the exposure's, each at a 256-byte boundary + Buffer paramAll; + std::array paramAt{}; + std::vector postAt, downAt; + VkDeviceSize pooledAt = 0, expAt = 0, paramBytes = 0; + std::array, R2Network> blocks; + std::vector> postBlocks, downBlocks; + std::vector pooledBlock; + R2ExposureParams exposure{}; + // enc0's feature (surface 368; next frame's texture 264), the intermediate history (post's surface 96 / texture + // 48) and the exposure, ping-pong; the rest is written and read within a frame or carried as is: + // mv[0] / mv[1]: enc0's surfaces 376 / 384 (post's textures 64 / 72; capFXmot); rec: dec0_1's 176 (enc0's 528 next frame); + // ka / kb / kc: dec0_1's 184 / 192 / 200 (post's 24 / 32 / 40); half: post's 120 (enc0's 536, pooled's input); + // quarter: pooled's output (enc0's 544) + std::array feat, hist, e, mv; + Tex rec, ka, kb, kc, half, quarter, out, zero; + std::map spare; + VkSampler nearest = VK_NULL_HANDLE, linear = VK_NULL_HANDLE; + struct Bindings { + RFrameImages images{}; + FrameOutput finalOutput{}; + VkDescriptorPool pool = VK_NULL_HANDLE; + std::array, R2Network> network{}; + std::array gather{}; + std::vector> post; + std::vector down; + std::vector> downP; + VkDescriptorSet pooled = VK_NULL_HANDLE, exposure0 = VK_NULL_HANDLE; + std::array exposure1{}; + }; + std::vector bindings; + const Bindings* bound = nullptr; + float prevJitter[2] = {0, 0}; + uint32_t frame = 0; + bool reset = true; + + static VkFormat surfaceFormat(const std::string& f) { + if (f == "rgba16f") return VK_FORMAT_R16G16B16A16_SFLOAT; + if (f == "rgba32f") return VK_FORMAT_R32G32B32A32_SFLOAT; + if (f == "rgb10_a2") return VK_FORMAT_A2B10G10R10_UNORM_PACK32; + if (f == "rg16f") return VK_FORMAT_R16G16_SFLOAT; + if (f == "r16f") return VK_FORMAT_R16_SFLOAT; + if (f == "rgba8_snorm") return VK_FORMAT_R8G8B8A8_SNORM; + if (f == "rgba8") return VK_FORMAT_R8G8B8A8_UNORM; + throw std::runtime_error("Ray Reconstruction stage surface format " + f + " is not implemented"); + } + // the pooled histories: the model's format of the post pass's surface 120 (UNORM 10:10:10:2, or RGBA32F holding + // the same q / 1023 in models built with D4R_RR2_POOLED=rgba32f) + VkFormat pooledFormat() const { + for (const auto& [offset, format] : model.stages[R2Post].surfaces) + if (offset == 120) return surfaceFormat(format); + return VK_FORMAT_R32G32B32A32_SFLOAT; + } + // tile i's slot: groups alternate between the two halves of the slots (one group when every tile has a slot) + uint32_t tileSlot(size_t i) const { + const uint32_t g = static_cast(i) / tileGroup, k = static_cast(i) % tileGroup; + return (g & 1u) * tileGroup + k; + } + uint32_t lw(int level) const { return w0 >> level; } + uint32_t lh(int level) const { return h0 >> level; } + struct Resource { VkImageView view = VK_NULL_HANDLE; bool linear = false; bool game = false; }; + // what a stage's texture or surface at an argument offset is, for frame parity `par` (the capture's wiring; + // filter modes are NGX's texture objects'); null view: NGX leaves it unbound + Resource resolve(unsigned stage, bool surface, uint32_t offset, uint32_t par, const RFrameImages& in, VkImageView output) const { + auto T = [](const Tex& t, bool lin = false) { return Resource{t.view, lin, false}; }; + auto G = [](VkImageView v, bool lin = false) { return Resource{v, lin, true}; }; + if (stage == R2Enc0) { + if (!surface) switch (offset) { + case 144: return G(in.color, true); case 160: return G(in.diffuseAlbedo); case 176: return G(in.specularAlbedo); + case 192: return G(in.motion); case 224: return G(in.normals); case 240: return G(in.roughness); + case 264: return T(feat[par ^ 1], true); case 528: return T(rec, true); case 536: return T(half, true); + case 544: return T(quarter, true); case 552: return G(in.depth); + // the optional inputs (harness captures capFX*, 2026-10-13): unbound without them, like NGX + case 272: return in.specularHitDistance ? G(in.specularHitDistance) : Resource{}; + case 336: return in.disocclusionMask ? G(in.disocclusionMask) : Resource{}; + case 288: return in.subsurfaceGuide ? G(in.subsurfaceGuide) : Resource{}; + } else switch (offset) { + case 368: return T(feat[par]); case 376: return T(mv[0]); case 384: return T(mv[1]); + } + } else if (stage == R2Dec0B) { + if (surface) switch (offset) { + case 176: return T(rec); case 184: return T(ka); case 192: return T(kb); case 200: return T(kc); + } + } else if (stage == R2Post) { + if (!surface) switch (offset) { + case 24: return T(ka); case 32: return T(kb, true); case 40: return T(kc); case 48: return T(hist[par ^ 1], true); + case 56: return T(feat[par]); case 64: return T(mv[0]); case 72: return T(mv[1], true); case 80: return G(in.color, true); + } else switch (offset) { + case 96: return T(hist[par]); case 120: return T(half); + } + } else if (stage == R2Down) { + if (surface && offset == 40) return Resource{output, false, false}; + } else if (stage == R2DownP) { + if (surface && offset == 40) return Resource{output, false, false}; + if (surface && offset == 264) return T(pb[par]); + if (!surface && offset == 208) return in.colorBeforeParticles ? G(in.colorBeforeParticles) : Resource{}; + if (!surface && offset == 256) return T(pb[par ^ 1], true); + } else if (stage == R2Pooled) { + return surface ? T(quarter) : T(half); + } + return {}; + } + + R2Engine(const Device& device, R2Model& m, uint32_t renderWidth, uint32_t renderHeight, uint32_t outputWidth, uint32_t outputHeight) + : d(device), model(m), rw(renderWidth), rh(renderHeight), ow(outputWidth), oh(outputHeight) { + w0 = (rw + 31) & ~31u; h0 = (rh + 31) & ~31u; + iw = (ow * 3 + 1) / 2; ih = (oh * 3 + 1) / 2; + cw = iw / 2; ch = ih / 2; qw = (cw + 1) / 2; qh = (ch + 1) / 2; + // tiles: columns of at most 720 and rows of at most 1620 intermediate pixels, 4 pixels of overlap inside + nx = (iw + 719) / 720; ny = (ih + 1619) / 1620; + for (uint32_t ty = 0; ty < ny; ++ty) + for (uint32_t tx = 0; tx < nx; ++tx) { + const uint32_t x0 = tx * iw / nx, x1 = (tx + 1) * iw / nx, y0 = ty * ih / ny, y1 = (ty + 1) * ih / ny; + src.push_back({x0 >= 4 ? x0 - 4 : 0, y0 >= 4 ? y0 - 4 : 0, std::min(iw, x1 + 4), std::min(ih, y1 + 4)}); + dst.push_back({tx * ow / nx, ty * oh / ny, (tx + 1) * ow / nx, (ty + 1) * oh / ny}); + } + maxW = maxH = 0; + for (const auto& r : src) { maxW = std::max(maxW, r.x1 - r.x0); maxH = std::max(maxH, r.y1 - r.y0); } + unit = static_cast(w0) * h0 / 32; + const uint64_t tilePixels = static_cast(maxW + 8) * (maxH + 8); + tileA = unit * R2ArenaUnits; tileB = tileA + tilePixels * 4; + tileBytes = tilePixels * 6; + const uint32_t tileCount = nx * ny; + const uint32_t slots = std::min(tileCount, R2TileSlots); + // every tile its own slot when they fit; otherwise groups of half the slots + tileGroup = tileCount <= R2TileSlots ? tileCount : R2TileSlots / 2; + const VkBufferUsageFlags storage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT; + // debugging: D4R_ENGINE_RR_ARENA=DIR keeps the arena and the exposure values host-visible and writes them at the + // next setFrame (DIR/F-arena.bin, F-values.bin), when the caller has waited for frame F + const bool debug = std::getenv("D4R_ENGINE_RR_ARENA") != nullptr; + // the gather pass's records: one per window of level 0, (W0 / 8 + 1) to a row (a shifted level has one more) + gatherAt = (tileA + tileBytes * slots + 255) & ~uint64_t(255); + const uint64_t gatherBytes = model.gather ? uint64_t(w0 / 8 + 1) * (h0 / 8 + 1) * model.stages[R2Gather].record * 4 : 0; + if (gatherAt + gatherBytes > 0xffffffffull) throw std::runtime_error("Ray Reconstruction: the render size is too large for the arena"); + arena.create(d, gatherAt + gatherBytes, storage, debug); + values.create(d, 16, storage, debug); + dummy.create(d, 4096, storage); + exposure.blocksX = (rw + 15) / 16; exposure.blocks = exposure.blocksX * ((rh + 15) / 16); + exposure.last[0] = static_cast(rw - 1); exposure.last[1] = static_cast(rh - 1); + exposure.invCount = static_cast(10000.0 / (4294967296.0 * rw * rh)); + sums.create(d, exposure.blocks * 8, storage); + const size_t tiles = src.size(); + auto place = [&](size_t bytes) { const VkDeviceSize at = paramBytes; paramBytes += (bytes + 255) & ~VkDeviceSize(255); return at; }; + for (unsigned i = 0; i < R2Network; ++i) { blocks[i] = model.stages[i].args; paramAt[i] = place(blocks[i].size()); } + postBlocks.assign(tiles, model.stages[R2Post].args); downBlocks.assign(tiles, model.stages[R2Down].args); + if (model.particles) downPBlocks.assign(tiles, model.stages[R2DownP].args); + for (size_t i = 0; i < tiles; ++i) { + postAt.push_back(place(postBlocks[i].size())); downAt.push_back(place(downBlocks[i].size())); + if (model.particles) downPAt.push_back(place(downPBlocks[i].size())); + } + pooledBlock = model.stages[R2Pooled].args; + pooledAt = place(pooledBlock.size()); + expAt = place(sizeof(R2ExposureParams)); + paramAll.create(d, paramBytes, storage); + const VkFormat half4 = VK_FORMAT_R16G16B16A16_SFLOAT; + for (int i = 0; i < 2; ++i) { + feat[i].create(d, rw, rh, half4); hist[i].create(d, iw, ih, half4); + e[i].create(d, 1, 1, VK_FORMAT_R16_SFLOAT); mv[i].create(d, rw, rh, VK_FORMAT_R16G16_SFLOAT); + } + rec.create(d, rw, rh, VK_FORMAT_R8G8B8A8_SNORM); + ka.create(d, rw, rh, VK_FORMAT_R8G8B8A8_UNORM); kb.create(d, rw, rh, VK_FORMAT_R8G8B8A8_UNORM); kc.create(d, rw, rh, half4); + half.create(d, cw, ch, pooledFormat()); quarter.create(d, qw, qh, pooledFormat()); + out.create(d, ow, oh, half4); + zero.create(d, 1, 1, half4); + for (const auto& st : model.stages) + for (const auto& [offset, format] : st.surfaces) + if (!spare.count(format)) spare[format].create(d, 1, 1, surfaceFormat(format)); + VkSamplerCreateInfo si{VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO}; + si.addressModeU = si.addressModeV = si.addressModeW = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE; si.maxLod = 0; + VKCHECK(vkCreateSampler(d.handle, &si, nullptr, &nearest)); + si.magFilter = si.minFilter = VK_FILTER_LINEAR; + VKCHECK(vkCreateSampler(d.handle, &si, nullptr, &linear)); + // the fixed fields. Network stages: size, arena regions, 1 / size + for (unsigned stage = 0; stage < R2Network; ++stage) { + auto& p = blocks[stage]; + const int l = R2Level[stage]; + put32(p, 0, lw(l)); put32(p, 4, lh(l)); + putF(p, 84, 1.0f / static_cast(lw(l))); putF(p, 88, 1.0f / static_cast(lh(l))); + for (const auto& [offset, units] : model.stages[stage].pointers) + put64(p, offset, static_cast(std::llround(units * static_cast(unit)))); + } + const float fr = static_cast(rw), fh = static_cast(rh), fiw = static_cast(iw), fih = static_cast(ih); + const float fow = static_cast(ow), foh = static_cast(oh); + auto& p = blocks[R2Enc0]; + if (model.gather) put64(p, p.size() - 8, gatherAt); // the word compile_rr2.py appends to the block + put32(p, 128, rw); put32(p, 132, rh); put32(p, 136, w0); put32(p, 140, h0); + putF(p, 408, fr / fiw); putF(p, 412, fh / fih); putF(p, 416, fiw / fr); putF(p, 420, fih / fh); + putF(p, 424, 1.0f / fr); putF(p, 428, 1.0f / fh); putF(p, 432, 1.0f / fiw); putF(p, 436, 1.0f / fih); + put32(p, 468, rw - 1); put32(p, 472, rh - 1); put32(p, 488, iw); put32(p, 492, ih); put32(p, 496, cw); put32(p, 500, ch); + putF(p, 504, 1.0f / fiw); putF(p, 508, 1.0f / fih); + putF(p, 512, 1.0f / static_cast(2 * qw)); putF(p, 516, 1.0f / static_cast(2 * qh)); + put32(p, 1076, rw); put32(p, 1080, rh); + // the network's output for the output stage: dec0_1's argument 64 + uint64_t netOut = 0; + for (const auto& [offset, units] : model.stages[R2Dec0B].pointers) + if (offset == 64) netOut = static_cast(std::llround(units * static_cast(unit))); + const uint32_t postCount = ((maxW + 15) / 16) * ((maxH + 31) / 32); + for (size_t i = 0; i < tiles; ++i) { + auto& o = postBlocks[i]; + put64(o, 0, netOut); put32(o, 8, w0); // the row stride of the network's output: the padded width + put64(o, 104, tileA + tileBytes * tileSlot(i)); put64(o, 112, tileB + tileBytes * tileSlot(i)); + putF(o, 132, fr / fiw); putF(o, 136, fh / fih); put32(o, 148, iw); put32(o, 152, ih); + putF(o, 160, 1.0f / fiw); putF(o, 164, 1.0f / fih); put32(o, 168, rw); put32(o, 172, rh); + put32(o, 176, src[i].x0); put32(o, 180, src[i].y0); put32(o, 184, src[i].x1); put32(o, 188, src[i].y1); + put32(o, 192, maxW); put32(o, 196, maxH); + putF(o, 200, 1.0f / fr); putF(o, 204, 1.0f / fh); + put32(o, 272, 0xffffffffu); put32(o, 276, 0); put32(o, 280, static_cast(i)); + for (auto* wb : {&downBlocks[i], model.particles ? &downPBlocks[i] : nullptr}) { + if (!wb) continue; + auto& w = *wb; + put32(w, 0, iw); put32(w, 4, ih); put32(w, 8, ow); put32(w, 12, oh); put32(w, 16, ow); put32(w, 20, oh); + put64(w, 24, tileA + tileBytes * tileSlot(i)); put64(w, 32, tileB + tileBytes * tileSlot(i)); + put32(w, 96, rw - 1); put32(w, 100, rh - 1); + putF(w, 104, fr / fow); putF(w, 108, fh / foh); putF(w, 112, fow / fiw); putF(w, 116, foh / fih); + put32(w, 128, dst[i].x0); put32(w, 132, dst[i].y0); put32(w, 136, dst[i].x1); put32(w, 140, dst[i].y1); + put32(w, 144, maxW); put32(w, 148, maxH); + put32(w, 152, src[i].x0); put32(w, 156, src[i].y0); put32(w, 160, src[i].x1); put32(w, 164, src[i].y1); + put32(w, 184, static_cast(i)); put32(w, 188, postCount); put32(w, 192, static_cast(i) + 2); + } + } + put32(pooledBlock, 16, cw); put32(pooledBlock, 20, ch); put32(pooledBlock, 24, qw); put32(pooledBlock, 28, qh); + } + ~R2Engine() { + for (auto& b : bindings) vkDestroyDescriptorPool(d.handle, b.pool, nullptr); + if (nearest) vkDestroySampler(d.handle, nearest, nullptr); + if (linear) vkDestroySampler(d.handle, linear, nullptr); + } + std::vector debugImages() const { + std::vector list; + const VkFormat half4 = VK_FORMAT_R16G16B16A16_SFLOAT; + for (int i = 0; i < 2; ++i) { + const std::string n = std::to_string(i); + list.push_back({"feat" + n, feat[i].handle, half4, rw, rh}); + list.push_back({"hist" + n, hist[i].handle, half4, iw, ih}); + list.push_back({"e" + n, e[i].handle, VK_FORMAT_R16_SFLOAT, 1, 1}); + list.push_back({"mv" + n, mv[i].handle, VK_FORMAT_R16G16_SFLOAT, rw, rh}); + } + list.push_back({"rec", rec.handle, VK_FORMAT_R8G8B8A8_SNORM, rw, rh}); + list.push_back({"ka", ka.handle, VK_FORMAT_R8G8B8A8_UNORM, rw, rh}); + list.push_back({"kb", kb.handle, VK_FORMAT_R8G8B8A8_UNORM, rw, rh}); + list.push_back({"kc", kc.handle, half4, rw, rh}); + list.push_back({"half", half.handle, pooledFormat(), cw, ch}); + list.push_back({"quarter", quarter.handle, pooledFormat(), qw, qh}); + list.push_back({"out", out.handle, half4, ow, oh}); + return list; + } + void recordInitialize(VkCommandBuffer cb) { + std::vector images; + for (auto* group : {&feat, &hist, &e, &mv}) for (auto& t : *group) images.push_back(t.handle); + for (auto* t : {&rec, &ka, &kb, &kc, &half, &quarter, &out, &zero}) images.push_back(t->handle); + for (auto& [format, t] : spare) images.push_back(t.handle); + VkClearColorValue clear{}; + VkImageSubresourceRange range{VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + for (VkImage image : images) { + imageBarrier(cb, image, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_GENERAL, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, + VK_PIPELINE_STAGE_TRANSFER_BIT, 0, VK_ACCESS_TRANSFER_WRITE_BIT); + vkCmdClearColorImage(cb, image, VK_IMAGE_LAYOUT_GENERAL, &clear, 1, &range); + } + for (auto* b : {&arena, &values, &sums, &dummy}) vkCmdFillBuffer(cb, b->handle, 0, VK_WHOLE_SIZE, 0); + memoryBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, + VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT | VK_ACCESS_TRANSFER_WRITE_BIT); + frame = 0; reset = true; + } + // a stage's descriptor set: parameter block, weights, arena, textures, surfaces, further buffers + VkDescriptorSet stageSet(Bindings& b, unsigned pass, VkDeviceSize param, uint32_t par, VkImageView output) { + const auto& st = model.stages[pass]; + const unsigned stage = pass == R2Gather ? R2Enc0 : pass; // the gather pass has encoder 0's bindings + VkDescriptorSet set = descriptorSet(d.handle, b.pool, st.pipeline.descriptors); + bindRange(d.handle, set, 0, paramAll, param); + bindBuffer(d.handle, set, 1, st.weights); + bindBuffer(d.handle, set, 2, arena); + uint32_t binding = 3; + for (uint32_t offset : st.textures) { + const auto res = resolve(stage, false, offset, par, b.images, output); + bindImage(d.handle, set, binding++, res.view ? res.view : zero.view, res.linear ? linear : nearest, + res.game ? b.images.layout : VK_IMAGE_LAYOUT_GENERAL); + } + for (const auto& [offset, format] : st.surfaces) { + const auto res = resolve(stage, true, offset, par, b.images, output); + bindImage(d.handle, set, binding++, res.view ? res.view : spare.at(format).view); + } + for (uint32_t offset : st.buffers) { + // the exposure values (enc0 256, post 16, down 72), enc0's second table (104); the progress counters + // are never accessed (their pointers are null) + const bool vals = (stage == R2Enc0 && offset == 256) || (stage == R2Post && offset == 16) || (stage == R2Down && offset == 72); + bindBuffer(d.handle, set, binding++, vals ? values : (stage == R2Enc0 && offset == 104) ? model.aux : dummy); + } + return set; + } + void setFrame(const RFrameImages& images, const RFrameSettings& settings, FrameOutput finalOutput) { + if (!images.color || !images.motion || !images.depth || !images.normals || !images.roughness || !images.diffuseAlbedo || !images.specularAlbedo) + throw std::runtime_error("Ray Reconstruction needs colour, motion, depth and the four guides"); + if (finalOutput.view && finalOutput.format != VK_FORMAT_R16G16B16A16_SFLOAT) + throw std::runtime_error("Ray Reconstruction writes RGBA16F output only"); + const bool rst = settings.reset || frame == 0; + const uint32_t f = frame; + if (const char* dir = std::getenv("D4R_ENGINE_RR_ARENA"); dir && f > 0 && arena.mapped) { + std::ofstream(std::string(dir) + "/" + std::to_string(f) + "-arena.bin", std::ios::binary) + .write(static_cast(arena.mapped), static_cast(unit * R2ArenaUnits)); + std::ofstream(std::string(dir) + "/" + std::to_string(f) + "-values.bin", std::ios::binary) + .write(static_cast(values.mapped), 16); + } + const float fr = static_cast(rw), fh = static_cast(rh), fiw = static_cast(iw), fih = static_cast(ih); + const float fow = static_cast(ow), foh = static_cast(oh); + const float pj[2] = {rst ? 0.0f : prevJitter[0], rst ? 0.0f : prevJitter[1]}; + for (unsigned stage = 0; stage < R2Network; ++stage) put32(blocks[stage], 80, r2Shift(stage, f)); + auto& p = blocks[R2Enc0]; + putF(p, 392, pj[0]); putF(p, 396, pj[1]); putF(p, 400, settings.jitter[0]); putF(p, 404, settings.jitter[1]); + putF(p, 440, pj[0] / fr); putF(p, 444, pj[1] / fh); + putF(p, 448, settings.motionScale[0] * (fiw / fr)); putF(p, 452, settings.motionScale[1] * (fih / fh)); + // bytes 465 / 466: the reset flag and DLSS.Roughness.Mode packed (the roughness is the normals' w); 568: inverted + // depth; 569 / 571: DLSSD.SpecularHitDistance / DLSS.DisocclusionMask are bound (textures 272 / 336); 1028: + // DLSS.Use.HW.Depth (harness captures capFXrough, capFXinv, capFXhit, capFXmask, capFXhw, 2026-10-13) + put32(p, 464, (rst ? 256u : 0u) | (settings.roughnessPacked ? 65536u : 0u)); + put32(p, 568, (settings.depthInverted ? 1u : 0u) | (images.specularHitDistance ? 256u : 0u) | (images.disocclusionMask ? 0x1000000u : 0u)); + p[1028] = settings.hardwareDepth ? 1u : 0u; + // the camera, as presets D and E (runtime_rr.inc), 48 bytes further on + Mat4 view, clip; + for (int i = 0; i < 16; ++i) { view[i] = settings.worldToView[i]; clip[i] = settings.viewToClip[i]; } + view[12] = view[13] = view[14] = 0.0; + const Mat4 rc = mul(view, clip), inv = inverse(rc); + const Mat4 toNdc{2.0 / rw, 0, 0, 0, 0, -2.0 / rh, 0, 0, 0, 0, 1, 0, -1, 1, 0, 1}; + const Mat4 toPixels{rw / 2.0, 0, 0, 0, 0, -(rh / 2.0), 0, 0, 0, 0, 1, 0, rw / 2.0, rh / 2.0, 0, 1}; + putMat(p, 644, rc); putMat(p, 708, clip); putMat(p, 772, clip); putMat(p, 836, inv); + putMat(p, 900, mul(toNdc, inv)); putMat(p, 964, mul(rc, toPixels)); + putF(p, 804, settings.jitter[0] * (2.0f / fr)); putF(p, 808, -settings.jitter[1] * (2.0f / fh)); + put32(p, 1032, f); + for (size_t i = 0; i < postBlocks.size(); ++i) { + auto& o = postBlocks[i]; + putF(o, 140, settings.jitter[0]); putF(o, 144, settings.jitter[1]); + put32(o, 208, rst ? 1u : 0u); + for (auto* wb : {&downBlocks[i], model.particles ? &downPBlocks[i] : nullptr}) { + if (!wb) continue; + auto& w = *wb; + putF(w, 120, settings.jitter[0]); putF(w, 124, settings.jitter[1]); + putF(w, 288, settings.jitter[0] * (fow / fr)); putF(w, 292, settings.jitter[1] * (foh / fh)); + // byte 196: the reset flag; with the particle layer also byte 197, and byte 198 set (capFXparticles) + put32(w, 196, wb == &downBlocks[i] ? (rst ? 1u : 0u) : (rst ? 0x010101u : 0x010000u)); + } + } + if (images.colorBeforeParticles) { + if (!model.particles) throw std::runtime_error("this preset F model has no particle-layer downsample stage (downp)"); + if (!pb[0].handle) { + for (auto& t : pb) t.create(d, ow, oh, VK_FORMAT_R16G16B16A16_SFLOAT); + pbNew = true; + } + } + exposure.reset = rst ? 1u : 0u; + reset = rst; + auto same = [&](const Bindings& b) { + return !std::memcmp(&b.images, &images, sizeof(images)) && b.finalOutput.view == finalOutput.view; + }; + auto found = std::find_if(bindings.begin(), bindings.end(), same); + if (found == bindings.end()) { + if (bindings.size() >= 16) throw std::runtime_error("too many distinct Ray Reconstruction input view sets"); + Bindings b; + b.images = images; b.finalOutput = finalOutput; + const uint32_t tiles = static_cast(src.size()); + auto count = [&](unsigned stage, uint32_t sets, uint32_t& bufs, uint32_t& tex, uint32_t& surf) { + const auto& st = model.stages[stage]; + bufs += sets * (3 + static_cast(st.buffers.size())); + tex += sets * static_cast(st.textures.size()); surf += sets * static_cast(st.surfaces.size()); + }; + uint32_t bufs = 8, tex = 4, surf = 4, sets = 3; + for (unsigned stage = 0; stage < R2Network; ++stage) { count(stage, 2, bufs, tex, surf); sets += 2; } + if (model.gather) { count(R2Gather, 2, bufs, tex, surf); sets += 2; } + count(R2Post, 2 * tiles, bufs, tex, surf); count(R2Down, tiles, bufs, tex, surf); count(R2Pooled, 1, bufs, tex, surf); + if (images.colorBeforeParticles) count(R2DownP, 2 * tiles, bufs, tex, surf); + sets += 5 * tiles + 1; + VkDescriptorPoolSize sizes[] = {{VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, bufs}, {VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, tex}, + {VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, surf}}; + VkDescriptorPoolCreateInfo pi{VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO}; + pi.maxSets = sets; pi.poolSizeCount = 3; pi.pPoolSizes = sizes; + VKCHECK(vkCreateDescriptorPool(d.handle, &pi, nullptr, &b.pool)); + const VkImageView output = finalOutput.view ? finalOutput.view : out.view; + for (unsigned stage = 0; stage < R2Network; ++stage) + for (uint32_t par = 0; par < 2; ++par) b.network[stage][par] = stageSet(b, stage, paramAt[stage], par, output); + if (model.gather) + for (uint32_t par = 0; par < 2; ++par) b.gather[par] = stageSet(b, R2Gather, paramAt[R2Enc0], par, output); + for (uint32_t i = 0; i < tiles; ++i) { + b.post.push_back({stageSet(b, R2Post, postAt[i], 0, output), stageSet(b, R2Post, postAt[i], 1, output)}); + b.down.push_back(stageSet(b, R2Down, downAt[i], 0, output)); + if (images.colorBeforeParticles) + b.downP.push_back({stageSet(b, R2DownP, downPAt[i], 0, output), stageSet(b, R2DownP, downPAt[i], 1, output)}); + } + b.pooled = stageSet(b, R2Pooled, pooledAt, 0, output); + b.exposure0 = descriptorSet(d.handle, b.pool, model.exposure[0].descriptors); + bindRange(d.handle, b.exposure0, 0, paramAll, expAt); + bindImage(d.handle, b.exposure0, 1, images.color, nearest, images.layout); + bindBuffer(d.handle, b.exposure0, 2, sums); + for (uint32_t par = 0; par < 2; ++par) { + VkDescriptorSet set = b.exposure1[par] = descriptorSet(d.handle, b.pool, model.exposure[1].descriptors); + bindRange(d.handle, set, 0, paramAll, expAt); + bindBuffer(d.handle, set, 1, sums); + bindBuffer(d.handle, set, 2, values); + bindImage(d.handle, set, 3, e[par ^ 1].view, nearest); + bindImage(d.handle, set, 4, e[par].view); + } + bindings.push_back(std::move(b)); + found = bindings.end() - 1; + } + bound = &*found; + // texture and surface handles: 1 where something is bound, 0 where NGX leaves the argument null (the kernels + // test optional inputs' handles); buffer pointers other than the arena's are zero + auto handles = [&](unsigned stage, std::vector& block) { + const auto& st = model.stages[stage]; + for (uint32_t offset : st.textures) put64(block, offset, resolve(stage, false, offset, 0, images, out.view).view ? 1u : 0u); + for (const auto& [offset, format] : st.surfaces) put64(block, offset, resolve(stage, true, offset, 0, images, out.view).view ? 1u : 0u); + for (uint32_t offset : st.buffers) put64(block, offset, 0); + }; + for (unsigned stage = 0; stage < R2Network; ++stage) handles(stage, blocks[stage]); + for (auto& o : postBlocks) handles(R2Post, o); + for (auto& w : downBlocks) handles(R2Down, w); + for (auto& w : downPBlocks) handles(R2DownP, w); + handles(R2Pooled, pooledBlock); + if (const char* dump = std::getenv("D4R_ENGINE_RR_DUMP")) { + auto write = [&](const std::string& name, const std::vector& block) { + std::ofstream(std::string(dump) + "/" + std::to_string(f + 1) + "-" + name + ".bin", std::ios::binary) + .write(reinterpret_cast(block.data()), static_cast(block.size())); + }; + for (unsigned stage = 0; stage < R2Network; ++stage) write(R2Names[stage], blocks[stage]); + for (size_t i = 0; i < postBlocks.size(); ++i) { write("post" + std::to_string(i), postBlocks[i]); write("down" + std::to_string(i), downBlocks[i]); } + } + prevJitter[0] = settings.jitter[0]; prevJitter[1] = settings.jitter[1]; + } + void recordFrame(VkCommandBuffer cb, VkQueryPool stamps, uint32_t first) { + if (!bound) throw std::runtime_error("REngine::setFrame was not called"); + const auto& b = *bound; + const uint32_t f = frame, cur = f & 1u; + auto stamp = [&](uint32_t n) { if (stamps) vkCmdWriteTimestamp(cb, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, stamps, first + n); }; + auto step = [&] { + memoryBarrier(cb, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, + VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT); + }; + if (pbNew) { // the particle history starts from zero, like the rest of the temporal state + VkClearColorValue clear{}; + VkImageSubresourceRange range{VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + for (auto& t : pb) { + imageBarrier(cb, t.handle, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_GENERAL, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, + VK_PIPELINE_STAGE_TRANSFER_BIT, 0, VK_ACCESS_TRANSFER_WRITE_BIT); + vkCmdClearColorImage(cb, t.handle, VK_IMAGE_LAYOUT_GENERAL, &clear, 1, &range); + } + memoryBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, + VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT); + pbNew = false; + } + stamp(0); + memoryBarrier(cb, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_SHADER_READ_BIT, VK_ACCESS_TRANSFER_WRITE_BIT); + std::vector all(paramBytes); + auto put = [&](VkDeviceSize at, const void* data, size_t bytes) { std::memcpy(all.data() + at, data, bytes); }; + for (unsigned i = 0; i < R2Network; ++i) put(paramAt[i], blocks[i].data(), blocks[i].size()); + for (size_t i = 0; i < postBlocks.size(); ++i) { put(postAt[i], postBlocks[i].data(), postBlocks[i].size()); put(downAt[i], downBlocks[i].data(), downBlocks[i].size()); } + for (size_t i = 0; i < downPBlocks.size(); ++i) put(downPAt[i], downPBlocks[i].data(), downPBlocks[i].size()); + put(pooledAt, pooledBlock.data(), pooledBlock.size()); + put(expAt, &exposure, sizeof(exposure)); + for (VkDeviceSize at = 0; at < paramBytes; at += 65536) // vkCmdUpdateBuffer takes at most 65536 bytes + vkCmdUpdateBuffer(cb, paramAll.handle, at, std::min(65536, paramBytes - at), all.data() + at); + memoryBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, + VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT); + dispatch(cb, model.exposure[0], b.exposure0, exposure.blocksX, exposure.blocks / exposure.blocksX); + step(); + dispatch(cb, model.exposure[1], b.exposure1[cur], 1, 1); + step(); + stamp(1); + for (unsigned stage = 0; stage < R2Network; ++stage) { + const int l = R2Level[stage]; + const uint32_t shift = r2Shift(stage, f), sx = shift & 0xffffu, sy = shift >> 16; + const uint32_t gx = (lw(l) + sx + 7) / 8; + if (stage == R2Enc0 && model.gather) { + dispatch(cb, model.stages[R2Gather].pipeline, b.gather[cur], gx, (lh(l) + sy + 7) / 8); + step(); + } + dispatch(cb, model.stages[stage].pipeline, b.network[stage][cur], (gx + model.stages[stage].group - 1) / model.stages[stage].group, (lh(l) + sy + 7) / 8); + step(); + } + stamp(2); + // timing only: D4R_ENGINE_RR_SKIP=post,down,pooled leaves those passes out (the picture is wrong then) + static const std::string skip = std::getenv("D4R_ENGINE_RR_SKIP") ? std::getenv("D4R_ENGINE_RR_SKIP") : ""; + static const bool noPost = skip.find("post") != skip.npos, noDown = skip.find("down") != skip.npos, noPooled = skip.find("pooled") != skip.npos; + static const bool serial = skip.find("serial") != skip.npos; // a barrier after every output pass, as before + // group g: its post passes, then (after one barrier) its down passes together with group g + 1's post passes + const bool particles = b.images.colorBeforeParticles != VK_NULL_HANDLE; + const size_t tiles = src.size(), groups = (tiles + tileGroup - 1) / tileGroup; + for (size_t g = 0; g <= groups; ++g) { + if (g > 0) + for (size_t i = (g - 1) * tileGroup; i < std::min(tiles, g * tileGroup); ++i) { + if (!noDown) dispatch(cb, model.stages[particles ? R2DownP : R2Down].pipeline, particles ? b.downP[i][cur] : b.down[i], + (dst[i].x1 - dst[i].x0 + 15) / 16, (dst[i].y1 - dst[i].y0 + 15) / 16); + if (serial) step(); + } + for (size_t i = g * tileGroup; i < std::min(tiles, (g + 1) * tileGroup); ++i) { + if (!noPost) dispatch(cb, model.stages[R2Post].pipeline, b.post[i][cur], (maxW + 15) / 16, (maxH + 31) / 32); + if (serial) step(); + } + if (g < groups) step(); + } + // (the pooling pass reads what the post passes wrote: the last group's down passes do not touch it) + if (!noPooled) dispatch(cb, model.stages[R2Pooled].pipeline, b.pooled, (qw + 15) / 16, (qh + 15) / 16); + memoryBarrier(cb, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, + VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT | VK_ACCESS_TRANSFER_READ_BIT); + stamp(3); + ++frame; + } +}; +} // namespace diff --git a/kernels/rr/swin_native.h b/kernels/rr/swin_native.h new file mode 100644 index 0000000..6b216b3 --- /dev/null +++ b/kernels/rr/swin_native.h @@ -0,0 +1,491 @@ +// A complete Ray Reconstruction Swin encoder stage as one native HIP kernel (gfx12, wave32, native FP8 WMMA): the port +// of engine/rr/swin_m.comp's KIND 0 path (tokens from the arena, attention, MLP, kernel-prediction filter, 2x2 merge), +// with nothing of NVIDIA's PTX in it. The arithmetic, slot orders and weight layout are that shader's (see its header +// and engine/rr/swin_pack.py); a line here corresponds to a line there. +// +// Differences from the shader's inputs, all made by the host (kernels/tools/rr_swin_native_test.cpp): +// - the V weights' tiles are stored [n16][k16] and the attention bias [query][key], so every operand of every +// product is 8 contiguous bytes per lane (the shader's row-major B loads are strided gathers); +// - texture 88 and the three surfaces are plain RGBA16F buffers. +// +// The including file defines the stage's shape and weight offsets first (the swin_stage.h that swin_pack.py writes). +// Block (32, 1, 4): one 8x8-token window, wave w = threadIdx.z owns query slots 16w .. 16w + 15. +#pragma once +#include +#include + +#if KIND != 0 || CX == 0 +#error "swin_native.h covers the encoder stages with a merge (enc1 - enc4)" +#endif +#define SN_KC (C / 32) +#define SN_MT (C / 16) +#define SN_NB (C / 8) +#define SN_HEAD (W_WK1 - W_WK0) +#define SN_BRANCH (W_W1_1 - W_W1_0) +// profiling only: fewer attention heads / MLP branches than the stage has (the output is then wrong) +#ifndef SN_HEADS +#define SN_HEADS NH +#endif +#ifndef SN_BRANCHES +#define SN_BRANCHES SN_NB +#endif + +#pragma clang fp contract(off) + +namespace swin_native +{ +typedef _Float16 hf; +typedef float f8v __attribute__((ext_vector_type(8))); +typedef int i2v __attribute__((ext_vector_type(2))); +struct h2 { hf x, y; }; +typedef _Float16 hv2 __attribute__((ext_vector_type(2))); +#define SN_DEV __device__ __forceinline__ + +SN_DEV h2 operator+(h2 a, h2 b) { return {(hf)(a.x + b.x), (hf)(a.y + b.y)}; } +SN_DEV h2 operator*(h2 a, h2 b) { return {(hf)(a.x * b.x), (hf)(a.y * b.y)}; } +SN_DEV uint32_t inter(uint32_t j) { return (j & ~31u) | (((j & 31u) >> 2) << 1) | (j & 1u) | (((j >> 1) & 1u) << 4); } +SN_DEV int mirror(int v, int n) { v = v < 0 ? -v : v; const int m = 2 * n - 2 - v; return v < m ? v : m; } +SN_DEV hf hb(uint32_t b) { return __builtin_bit_cast(hf, (uint16_t)b); } +SN_DEV uint32_t bh(hf h) { return __builtin_bit_cast(uint16_t, h); } +SN_DEV float r16(float x) { return (float)(hf)x; } +SN_DEV uint32_t other_half(uint32_t v) { return __builtin_amdgcn_permlanex16(v, v, 0x76543210u, 0xfedcba98u, false, false); } +// byte b of w as an e4m3 value +SN_DEV hf e4(uint32_t w, int b) { return (hf)__builtin_amdgcn_cvt_f32_fp8((int)(w >> (8 * b)), 0); } +SN_DEV float sat(float x) { return __builtin_fminf(__builtin_fmaxf(x, -448.0f), 448.0f); } +// four values -> four e4m3 codes in one word (no saturation: the caller clamps where the value can exceed 448) +SN_DEV uint32_t pk4(float a, float b, float c, float d) +{ + const int w = __builtin_amdgcn_cvt_pk_fp8_f32(a, b, 0, false); + return (uint32_t)__builtin_amdgcn_cvt_pk_fp8_f32(c, d, w, true); +} +SN_DEV uint32_t q4(float a, float b, float c, float d) { return pk4(sat(a), sat(b), sat(c), sat(d)); } +// an accumulator (rows 8h + i of this lane's column) as the operand of the next product (the rows become K) +// A binary16 accumulator: rows 8h + i of this lane's column, two per register +struct A16 +{ + hv2 p[4]; + SN_DEV hf get(int i) const { return p[i >> 1][i & 1]; } + SN_DEV void set(int i, hf v) { p[i >> 1][i & 1] = v; } +}; +SN_DEV A16 zero16() { A16 r; r.p[0] = r.p[1] = r.p[2] = r.p[3] = (hv2){0, 0}; return r; } +SN_DEV A16 load16(const uint8_t* q) { return *(const A16*)q; } // 8 contiguous binary16 values +SN_DEV i2v q8(const A16& a) +{ + return (i2v){(int)q4((float)a.get(0), (float)a.get(1), (float)a.get(2), (float)a.get(3)), + (int)q4((float)a.get(4), (float)a.get(5), (float)a.get(6), (float)a.get(7))}; +} +// c = f16(P + c) with one rounding (v_fma_mix: P is f32, the addend and the result binary16), in place: mixlo writes +// only the low half and mixhi reads only the high half of the addend. The reference adds c inside the f32 product +// and rounds that to binary16 (two roundings); kernels/m/swin_common.h found the two equal on its launches, and the +// stage test compares every byte against the capture. +SN_DEV void mix(A16& c, f8v P) +{ +#pragma unroll + for (int j = 0; j < 4; ++j) { + asm("v_fma_mixlo_f16 %0, %1, 1.0, %0 op_sel_hi:[0,0,1]" : "+v"(c.p[j]) : "v"(P[2 * j])); + asm("v_fma_mixhi_f16 %0, %1, 1.0, %0 op_sel:[0,0,1] op_sel_hi:[0,0,1]" : "+v"(c.p[j]) : "v"(P[2 * j + 1])); + } +} +SN_DEV void k32(A16& c, i2v a0, i2v b0, i2v a1, i2v b1) +{ + f8v P = __builtin_amdgcn_wmma_f32_16x16x16_fp8_fp8_w32_gfx12(a0, b0, (f8v){0, 0, 0, 0, 0, 0, 0, 0}); + P = __builtin_amdgcn_wmma_f32_16x16x16_fp8_fp8_w32_gfx12(a1, b1, P); + mix(c, P); +} +// this lane's 8 K values of row / column `col` of a 16-byte-per-row image +SN_DEV i2v op(const uint8_t* p) { return *(const i2v*)p; } +SN_DEV i2v tile(const uint8_t* w, uint32_t base, uint32_t t, uint32_t col, uint32_t hl) { return op(w + base + 256u * t + 16u * col + 8u * hl); } +SN_DEV hf h16at(const uint8_t* w, uint32_t byte) { return *(const hf*)(w + byte); } + +SN_DEV int cmp1(float m, float x) +{ + const float t = m * x; + const float e = __builtin_fmaf(m, x, -t); + return t > 1.0f ? 1 : (t < 1.0f ? -1 : (e > 0.0f ? 1 : (e < 0.0f ? -1 : 0))); +} +// binary16 rounding of 1 / sqrt(s) (s > 0), as engine/rr/rr_exact.glsl +SN_DEV hf rsq16(hf s16) +{ + const float s = (float)s16; + hf q = (hf)__builtin_amdgcn_rsqf(s); +#pragma unroll + for (int i = 0; i < 2; ++i) { + const uint32_t b = bh(q); + const float lo = ((float)q + (float)hb(b - 1u)) * 0.5f, hi = ((float)q + (float)hb(b + 1u)) * 0.5f; + const int cl = cmp1(lo * lo, s), ch = cmp1(hi * hi, s); + if (cl > 0 || (cl == 0 && (b & 1u) != 0u)) q = hb(b - 1u); + else if (ch < 0 || (ch == 0 && (b & 1u) != 0u)) q = hb(b + 1u); + } + return q; +} +SN_DEV hf rcp16(float x) { return (hf)__builtin_amdgcn_rcpf(x); } +SN_DEV hf sig(hf v) { return rcp16(__builtin_amdgcn_exp2f((float)v * -1.4426950216293335f) + 1.0f); } +template SN_DEV T tree(const T* v) +{ +#if SN_KC == 1 + return v[0]; +#elif SN_KC == 2 + return v[0] + v[1]; +#elif SN_KC == 3 + return (v[0] + v[1]) + v[2]; +#elif SN_KC == 4 + return (v[0] + v[1]) + (v[2] + v[3]); +#else + return ((v[0] + v[1]) + (v[2] + v[3])) + v[4]; +#endif +} +SN_DEV void key_xy(uint32_t t, int x0, int y0, int& x, int& y) +{ + if (t < 64u) { x = x0 + (int)(t & 7u); y = y0 + (int)(t >> 3); return; } + const int i = (int)(t & 7u), side = (int)((t - 64u) >> 3); + x = side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8); + y = side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i); +} +SN_DEV const hf* texel(const hf* t88, int ex, int x, int y) { return t88 + 4 * (y * ex + x); } + +// param: the stage's parameter block (engine/rr/test_layer.stage_params); t88 and s112 / s120 are ex x ey, s128 half +__global__ void swin_stage(const uint32_t* __restrict__ param, const uint8_t* __restrict__ w, uint8_t* __restrict__ arena_b, + const hf* __restrict__ t88, hf* __restrict__ s112, hf* __restrict__ s120, hf* __restrict__ s128) +{ + __shared__ __attribute__((aligned(16))) uint8_t XT[96 * C]; // the normalised key tokens, then the merge rows + __shared__ __attribute__((aligned(16))) uint8_t KVS[2 * 3072]; // the current head's K, then V^T + __shared__ hf O112[64 * 4]; + uint8_t* const KS = KVS; + uint8_t* const VS = KVS + 3072; + uint32_t* const arena = (uint32_t*)arena_b; + + const uint32_t lane = threadIdx.x, wv = threadIdx.z, tid = 32u * wv + lane; + const uint32_t hl = lane >> 4, col = lane & 15u; + const int ex = (int)param[0], ey = (int)param[1], hx = ex / 2, hy = ey / 2; + const uint32_t inOff = param[2], outOff = param[12], mrgOff = param[14], shift = param[20]; + const int x0 = 8 * (int)blockIdx.x - (int)(shift & 0xffffu), y0 = 8 * (int)blockIdx.y - (int)(shift >> 16); + const uint32_t plane = (uint32_t)(ex * ey) * 32u, lplane = (uint32_t)(hx * hy) * 32u; + + // ---------------------------------------------------------------- prologue: one key token per thread + if (tid < 96u) { + int kx, ky; + key_xy(inter(tid), x0, y0, kx, ky); + const int xx = mirror(kx, ex), yy = mirror(ky, ey); + const uint32_t tb = (inOff + (uint32_t)(yy * ex + xx) * 32u) / 4u; + hf v[C]; +#pragma unroll + for (uint32_t wd = 0; wd < C / 4; ++wd) { + const uint32_t t = arena[tb + (wd >> 3) * (plane >> 2) + (wd & 7u)]; +#pragma unroll + for (int b = 0; b < 4; ++b) v[4u * wd + b] = e4(t, b); + } + h2 h[4]; +#pragma unroll + for (uint32_t q = 0; q < 4; ++q) { + h2 hh[2]; +#pragma unroll + for (uint32_t half_ = 0; half_ < 2; ++half_) { + h2 pr[SN_KC]; +#pragma unroll + for (uint32_t p = 0; p < SN_KC; ++p) { + const uint32_t o = 32u * p + 4u * (q + 4u * half_); + const h2 a = h2{v[o], v[o + 1]} * h2{v[o], v[o + 1]}; + const h2 c = h2{v[o + 2], v[o + 3]} * h2{v[o + 2], v[o + 3]}; + pr[p] = a + c; + } + hh[half_] = tree(pr); + } + h[q] = hh[0] + hh[1]; + } + const h2 h0 = h[0] + h[1], h2_ = h[2] + h[3]; + const h2 hs = h0 + h2_; + const hf sm = hs.x + hs.y; + const hf se = sm + (hf)0.0001220703125f; + const hf rr = rsq16(se); +#pragma unroll + for (uint32_t wd = 0; wd < C / 4; ++wd) { + float n[4]; +#pragma unroll + for (uint32_t b = 0; b < 4; ++b) { + const uint32_t j = 4u * wd + b; + const hf g = rr * h16at(w, W_GAIN0 + 2u * inter(j)); + n[b] = (float)(hf)(v[j] * g); + } + *(uint32_t*)(XT + tid * C + 4u * wd) = q4(n[0], n[1], n[2], n[3]); + } + } + __syncthreads(); + + // ---------------------------------------------------------------- the residual: the window token + table 7 + A16 R[SN_MT]; + { + const uint32_t tq = inter(16u * wv + col); + const int qx = x0 + (int)(tq & 7u), qy = y0 + (int)(tq >> 3); + const uint32_t tb = (inOff + (uint32_t)(mirror(qy, ey) * ex + mirror(qx, ex)) * 32u) / 4u; +#pragma unroll + for (uint32_t m = 0; m < SN_MT; ++m) { + const uint32_t ta = tb + (m >> 1) * (plane >> 2) + 4u * (m & 1u) + 2u * hl; + const uint32_t t[2] = {arena[ta], arena[ta + 1]}; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const hf c = e4(t[i >> 2], i & 3) + h16at(w, W_TABLE7 + 2u * (16u * m + 8u * hl + (uint32_t)i)); + R[m].set(i, c); + } + } + } + + // ---------------------------------------------------------------- attention, NH heads +#pragma nounroll + for (uint32_t hd = 0; hd < SN_HEADS; ++hd) { + __syncthreads(); + const uint32_t wk = W_WK0 + SN_HEAD * hd, wvb = W_WV0 + SN_HEAD * hd, wo = W_WO0 + SN_HEAD * hd; + for (uint32_t j = 0; j < 3; ++j) { // 12 jobs (token tile n, channel tile c) + const uint32_t job = wv + 4u * j, n = job >> 1, c = job & 1u; + A16 a = zero16(), v = zero16(); +#pragma unroll + for (uint32_t kt = 0; kt < 2 * SN_KC; kt += 2) { + // token tile n's rows are the A operand of V and the B operand of K: the same 8 bytes + const i2v x0_ = op(XT + (16u * n + col) * C + 16u * kt + 8u * hl), x1_ = op(XT + (16u * n + col) * C + 16u * kt + 16u + 8u * hl); + k32(a, tile(w, wk, c * (C / 16) + kt, col, hl), x0_, tile(w, wk, c * (C / 16) + kt + 1, col, hl), x1_); + k32(v, x0_, tile(w, wvb, kt * 2 + c, col, hl), x1_, tile(w, wvb, (kt + 1) * 2 + c, col, hl)); + } + *(i2v*)(KS + (16u * n + col) * 32u + 16u * c + 8u * hl) = q8(a); + *(i2v*)(VS + (16u * c + col) * 96u + 16u * n + 8u * hl) = q8(v); + } + __syncthreads(); + i2v Qb[2], Pb[6]; +#pragma unroll + for (uint32_t c = 0; c < 2; ++c) Qb[c] = op(KS + (16u * wv + col) * 32u + 16u * c + 8u * hl); + { + hf e[6][8]; + const uint8_t* bias = w + W_BIAS0 + SN_HEAD * hd + (16u * wv + col) * 192u + 16u * hl; // [query][key] f16 +#pragma unroll + for (uint32_t kt = 0; kt < 6; ++kt) { + A16 s = load16(bias + 32u * kt); + k32(s, op(KS + (16u * kt + col) * 32u + 8u * hl), Qb[0], op(KS + (16u * kt + col) * 32u + 16u + 8u * hl), Qb[1]); +#pragma unroll + for (int i = 0; i < 8; ++i) e[kt][i] = hb(*(const uint16_t*)(w + W_CUBIC + 2u * bh(s.get(i)))); + } + h2 ep[6][4]; +#pragma unroll + for (uint32_t kt = 0; kt < 6; ++kt) +#pragma unroll + for (uint32_t q = 0; q < 4; ++q) ep[kt][q] = h2{e[kt][2 * q], e[kt][2 * q + 1]}; + h2 st[2]; +#pragma unroll + for (uint32_t tq = 0; tq < 2; ++tq) { + const uint32_t q0 = 2u * tq, q1 = q0 + 1u; + const h2 ev = ((ep[0][q0] + ep[0][q1]) + (ep[2][q0] + ep[2][q1])) + (ep[4][q0] + ep[4][q1]); + const h2 od = ((ep[1][q0] + ep[1][q1]) + (ep[3][q0] + ep[3][q1])) + (ep[5][q0] + ep[5][q1]); + st[tq] = ev + od; + } + const h2 loc = st[0] + st[1]; + const uint32_t other = other_half(bh(loc.x) | (bh(loc.y) << 16)); + const h2 o = loc + h2{hb(other & 0xffffu), hb(other >> 16)}; + const hf rs = o.x + o.y; + const hf r = (hf)__builtin_amdgcn_rcpf((float)rs); +#pragma unroll + for (uint32_t kt = 0; kt < 6; ++kt) { + float p[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) p[i] = (float)(hf)(e[kt][i] * r); + Pb[kt] = (i2v){(int)pk4(p[0], p[1], p[2], p[3]), (int)pk4(p[4], p[5], p[6], p[7])}; + } + } + i2v Ob[2]; +#pragma unroll + for (uint32_t c = 0; c < 2; ++c) { + A16 o = zero16(); +#pragma unroll + for (uint32_t ks = 0; ks < 3; ++ks) { + const uint8_t* vp = VS + (16u * c + col) * 96u + 32u * ks + 8u * hl; + k32(o, op(vp), Pb[2 * ks], op(vp + 16), Pb[2 * ks + 1]); + } + Ob[c] = q8(o); + } +#pragma unroll + for (uint32_t m = 0; m < SN_MT; ++m) + k32(R[m], tile(w, wo, 2 * m, col, hl), Ob[0], tile(w, wo, 2 * m + 1, col, hl), Ob[1]); + } + + // ---------------------------------------------------------------- MLP + i2v Nb[SN_MT]; + { + hf rv[SN_MT][8]; +#pragma unroll + for (uint32_t m = 0; m < SN_MT; ++m) +#pragma unroll + for (int i = 0; i < 8; ++i) rv[m][i] = R[m].get(i); + hf h1[2]; +#pragma unroll + for (uint32_t c = 0; c < 2; ++c) { + hf ht[2]; +#pragma unroll + for (uint32_t tq = 0; tq < 2; ++tq) { + const uint32_t i0 = 4u * tq + c, i1 = i0 + 2u; + hf par[2]; +#pragma unroll + for (uint32_t pa = 0; pa < 2; ++pa) { + hf pr[SN_KC]; +#pragma unroll + for (uint32_t b = 0; b < SN_KC; ++b) { + const hf s0 = rv[2 * b + pa][i0] * rv[2 * b + pa][i0], s1 = rv[2 * b + pa][i1] * rv[2 * b + pa][i1]; + pr[b] = s0 + s1; + } + par[pa] = tree(pr); + } + ht[tq] = par[0] + par[1]; + } + h1[c] = ht[0] + ht[1]; + } + const uint32_t other = other_half(bh(h1[0]) | (bh(h1[1]) << 16)); + const hf g0 = h1[0] + hb(other & 0xffffu), g1 = h1[1] + hb(other >> 16); + const hf s = g0 + g1; + const hf se = s + (hf)0.0001220703125f; + const hf r = rsq16(se); +#pragma unroll + for (uint32_t m = 0; m < SN_MT; ++m) { + float n[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const uint32_t k = 16u * m + 8u * hl + (uint32_t)i; + const hf g = r * h16at(w, W_GAIN14 + 2u * k); + n[i] = (float)(hf)(rv[m][i] * g); + const hf a = rv[m][i] + h16at(w, W_TABLE16 + 2u * k); + R[m].set(i, a); + } + Nb[m] = (i2v){(int)q4(n[0], n[1], n[2], n[3]), (int)q4(n[4], n[5], n[6], n[7])}; + } + } +#pragma nounroll + for (uint32_t br = 0; br < SN_BRANCHES; ++br) { + const uint32_t w1 = W_W1_0 + SN_BRANCH * br, b1 = W_B1_0 + SN_BRANCH * br, w2 = W_W2_0 + SN_BRANCH * br; + i2v Hb[2]; +#pragma unroll + for (uint32_t c = 0; c < 2; ++c) { + A16 h = load16(w + b1 + 2u * (16u * c + 8u * hl)); +#pragma unroll + for (uint32_t kt = 0; kt < SN_MT; kt += 2) + k32(h, tile(w, w1, c * (C / 16) + kt, col, hl), Nb[kt], tile(w, w1, c * (C / 16) + kt + 1, col, hl), Nb[kt + 1]); + uint32_t b[2] = {0, 0}; +#pragma unroll + for (int i = 0; i < 8; ++i) b[i >> 2] |= (uint32_t)w[W_ACT + bh(h.get(i))] << (8 * (i & 3)); + Hb[c] = (i2v){(int)b[0], (int)b[1]}; + } +#pragma unroll + for (uint32_t m = 0; m < SN_MT; ++m) + k32(R[m], tile(w, w2, 2 * m, col, hl), Hb[0], tile(w, w2, 2 * m + 1, col, hl), Hb[1]); + } + + // ---------------------------------------------------------------- outputs + const uint32_t t = inter(16u * wv + col); + const int tx = x0 + (int)(t & 7u), ty = y0 + (int)(t >> 3); + const bool inside = tx >= 0 && ty >= 0 && tx < ex && ty < ey; + i2v Tb[SN_MT]; + { + const uint32_t o = (outOff + (uint32_t)(ty * ex + tx) * 32u) >> 2; + const uint32_t ty2 = t >> 3, tx2 = t & 7u; + const uint32_t mr = (ty2 >> 1) < 2u ? 2u * (tx2 >> 1) + (ty2 >> 1) : 8u + 2u * (tx2 >> 1) + (ty2 >> 1) - 2u; + const uint32_t blk = 2u * (ty2 & 1u) + (tx2 & 1u); + __syncthreads(); // every wave is done with XT (the merge rows reuse it) +#pragma unroll + for (uint32_t m = 0; m < SN_MT; ++m) { + Tb[m] = q8(R[m]); + const uint32_t at = 4u * (m & 1u) + 2u * hl; + if (inside) { arena[o + (m >> 1) * (plane >> 2) + at] = (uint32_t)Tb[m][0]; arena[o + (m >> 1) * (plane >> 2) + at + 1u] = (uint32_t)Tb[m][1]; } + *(i2v*)(XT + mr * 4u * C + blk * C + 16u * m + 8u * hl) = Tb[m]; + } + } + // the C -> 16 head + A16 hv = load16(w + W_BH + 16u * hl); +#pragma unroll + for (uint32_t kt = 0; kt < SN_MT; kt += 2) + k32(hv, tile(w, W_WH, kt, col, hl), Tb[kt], tile(w, W_WH, kt + 1, col, hl), Tb[kt + 1]); + { // the kernel-prediction filter (lanes 0-15; 16-31 hold head channels 8-15) + hf l[12]; +#pragma unroll + for (int i = 0; i < 8; ++i) l[i] = hv.get(i); + const uint32_t u0 = other_half(bh(l[0]) | (bh(l[1]) << 16)), u1 = other_half(bh(l[2]) | (bh(l[3]) << 16)); + l[8] = hb(u0 & 0xffffu); l[9] = hb(u0 >> 16); l[10] = hb(u1 & 0xffffu); l[11] = hb(u1 >> 16); +#define SN_MAX(a, b) __builtin_fmaxf16((a), (b)) + const hf ma = SN_MAX(SN_MAX(SN_MAX(l[0], l[2]), l[4]), l[6]), mb = SN_MAX(SN_MAX(SN_MAX(l[1], l[3]), l[5]), l[7]); + const hf mx = SN_MAX(SN_MAX(ma, mb), l[8]); + float e[9]; +#pragma unroll + for (int i = 0; i < 9; ++i) { const hf d = l[i] - mx; e[i] = __builtin_amdgcn_exp2f((float)d * 1.4426950216293335f); } + float s = (e[0] + e[1]) + 0.0f; + s = s + (e[2] + e[3]); + s = s + (e[4] + e[5]); + s = s + (e[6] + e[7]); + s = s + e[8]; + const hf rr = rcp16(s); + float wg[9]; +#pragma unroll + for (int i = 0; i < 9; ++i) wg[i] = r16(r16(e[i]) * (float)rr); + // a token outside the image stores nothing: its taps only have to stay inside the texture buffer + const int x = inside ? tx : 0, y = inside ? ty : 0; + const int xm = x - 1 < 0 ? 1 : x - 1, xc = x < 0 ? 1 : (x >= ex ? ex - 2 : x), xp = x + 1 >= ex ? ex - 2 : x + 1; + const int ym = y - 1 < 0 ? 1 : y - 1, yc = y < 0 ? 1 : (y >= ey ? ey - 2 : y), yp = y + 1 >= ey ? ey - 2 : y + 1; + const int px[9] = {xm, xc, xp, xm, xc, xp, xm, xc, xp}, py[9] = {ym, ym, ym, yc, yc, yc, yp, yp, yp}; + float acc[3][2] = {{0.0f, 0.0f}, {0.0f, 0.0f}, {0.0f, 0.0f}}; +#pragma unroll + for (int k = 0; k < 8; ++k) { + const hf* tv = texel(t88, ex, px[k], py[k]); +#pragma unroll + for (int c = 0; c < 3; ++c) { const float f = (float)tv[c] * wg[k] + acc[c][k & 1]; acc[c][k & 1] = r16(f); } + } + const hf* t8 = texel(t88, ex, px[8], py[8]); + float o112[3], o120[3]; +#pragma unroll + for (int c = 0; c < 3; ++c) { + const float f0 = acc[c][0] + acc[c][1]; + const float f1 = (float)t8[c] * wg[8] + r16(f0); + const float f = r16(f1); + const float a = (float)sig(l[9 + c]) * f; + const float ar = r16(a); + const float b = f - ar; + o120[c] = ar; o112[c] = r16(b); + } + if (hl == 0u && inside) { + hf* p = s112 + 4 * (y * ex + x); + hf* q = s120 + 4 * (y * ex + x); +#pragma unroll + for (int c = 0; c < 3; ++c) { p[c] = (hf)o112[c]; q[c] = (hf)o120[c]; } + p[3] = (hf)0.0f; q[3] = (hf)0.0f; + } + if (hl == 0u) +#pragma unroll + for (int c = 0; c < 3; ++c) O112[t * 4u + c] = (hf)o112[c]; + } + __syncthreads(); + { // the 2x2 merge: merged^T (CX x 16 rows), channel tiles m = wv, wv + 4, ... + const uint32_t mxx = col < 8u ? col >> 1 : (col - 8u) >> 1, mxy = col < 8u ? col & 1u : 2u + (col & 1u); + const int X = x0 / 2 + (int)mxx, Y = y0 / 2 + (int)mxy; + const uint32_t o = (mrgOff + (uint32_t)(Y * hx + X) * 32u) >> 2; + for (uint32_t m = wv; m < CX / 16u; m += 4u) { + A16 a = load16(w + W_BM + 2u * (16u * m + 8u * hl)); + for (uint32_t ks = 0; ks < C / 8u; ++ks) { + const uint8_t* g = XT + col * 4u * C + 32u * ks + 8u * hl; + k32(a, tile(w, W_WM, m * (4u * C / 16u) + 2u * ks, col, hl), op(g), tile(w, W_WM, m * (4u * C / 16u) + 2u * ks + 1u, col, hl), op(g + 16)); + } + const i2v q = q8(a); + if (X >= 0 && Y >= 0 && X < hx && Y < hy) { + const uint32_t at = o + (m >> 1) * (lplane >> 2) + 4u * (m & 1u) + 2u * hl; + arena[at] = (uint32_t)q[0]; + arena[at + 1u] = (uint32_t)q[1]; + } + } + if (wv == 3u && lane < 16u) { // surface 128: 2x2 sums of 0.25 * surface 112, rightwards first, then downwards + const uint32_t mx_ = lane & 3u, my_ = lane >> 2; + const int x = x0 + 2 * (int)mx_, y = y0 + 2 * (int)my_; + if (x >= 0 && y >= 0 && x < ex && y < ey) { + const uint32_t t00 = (2u * my_) * 8u + 2u * mx_; + hf* p = s128 + 4 * ((y >> 1) * hx + (x >> 1)); +#pragma unroll + for (uint32_t c = 0; c < 3; ++c) { + const hf a00 = O112[t00 * 4u + c] * (hf)0.25f, a01 = O112[(t00 + 1u) * 4u + c] * (hf)0.25f; + const hf a10 = O112[(t00 + 8u) * 4u + c] * (hf)0.25f, a11 = O112[(t00 + 9u) * 4u + c] * (hf)0.25f; + const hf b0 = a00 + a01, b1 = a10 + a11; + p[c] = b0 + b1; + } + p[3] = (hf)0.0f; + } + } + } +} +} // namespace swin_native diff --git a/kernels/tools/extract_dlss_ptx.py b/kernels/tools/extract_dlss_ptx.py index 9461f0d..fea3fd5 100644 --- a/kernels/tools/extract_dlss_ptx.py +++ b/kernels/tools/extract_dlss_ptx.py @@ -48,6 +48,23 @@ def lz4_block(src: bytes, size: int) -> bytes: return bytes(dst[:size]) +ZSTD_MAGIC = b'\x28\xb5\x2f\xfd' + + +def zstd_frame(src: bytes) -> bytes: + """One zstd frame, through the zstandard module when it is installed and the zstd tool otherwise.""" + try: + import zstandard + return zstandard.ZstdDecompressor().decompressobj().decompress(src) + except ImportError: + import subprocess + # the entry is padded: the tool decodes the frame and then complains about the bytes after it + out = subprocess.run(['zstd', '-dqc'], input=src, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL).stdout + if not out: + raise RuntimeError('zstd could not decode a PTX entry') + return out + + def fatbins(data: bytes): """Yields (offset, bytes) of every fatbin container (magic, version 1) in file order.""" magic = struct.pack('= 64 else 0 payload = fatbin[offset + entry_header:offset + entry_header + entry_size] if kind == 1: # PTX - yield lz4_block(payload, decompressed) if (flags & 0x2000 and decompressed) else payload + if payload[:4] == ZSTD_MAGIC: # nvngx_dlssd.dll 310.9 stores its rr2_* modules as zstd frames + yield zstd_frame(payload) + else: + yield lz4_block(payload, decompressed) if (flags & 0x2000 and decompressed) else payload offset += entry_header + entry_size diff --git a/kernels/tools/rr_swin_native_test.cpp b/kernels/tools/rr_swin_native_test.cpp new file mode 100644 index 0000000..f311aa2 --- /dev/null +++ b/kernels/tools/rr_swin_native_test.cpp @@ -0,0 +1,121 @@ +// Runs kernels/rr/swin_native.h (one Ray Reconstruction encoder stage as a native HIP kernel) on a stage work +// directory written by engine/rr/test_layer.py, compares it with that directory's reference outputs and times it. +// +// hipcc --offload-arch=gfx1201 -O3 -I STAGE_DIR -o rr_swin_native_test kernels/tools/rr_swin_native_test.cpp +// rr_swin_native_test WORK WEIGHTS.bin TEXTURE88.bin [REPEATS] +// +// STAGE_DIR holds the stage's swin_stage.h (engine/rr/swin_pack.py), WEIGHTS.bin is the blob the same run wrote, WORK +// is test_layer.py's (params.bin, arena.bin; arena.out.bin and surf112 / 120 / 128.bin from a run of the Vulkan stage +// that was identical to the capture). Channel 3 of the surfaces is left out, as test_layer.py --ignore does. +#include "swin_stage.h" +#include "../rr/swin_native.h" +#include +#include +#include +#include +#include +#include +#include +#include + +#define CHECK(expr) \ + do { \ + hipError_t e_ = (expr); \ + if (e_ != hipSuccess) { std::printf("HIP error at line %d: %s -> %s\n", __LINE__, #expr, hipGetErrorString(e_)); std::exit(2); } \ + } while (0) + +static std::vector slurp(const std::string& path) +{ + std::ifstream f(path, std::ios::binary); + if (!f) { std::printf("cannot open %s\n", path.c_str()); std::exit(2); } + return std::vector((std::istreambuf_iterator(f)), std::istreambuf_iterator()); +} +template static T* upload(const std::vector& host, size_t bytes = 0) +{ + void* p = nullptr; + bytes = std::max(bytes, host.size()); + CHECK(hipMalloc(&p, bytes)); + CHECK(hipMemset(p, 0, bytes)); + if (!host.empty()) CHECK(hipMemcpy(p, host.data(), host.size(), hipMemcpyHostToDevice)); + return (T*)p; +} +static size_t differing(const std::vector& a, const std::vector& b, bool skipAlpha) +{ + size_t n = a.size() != b.size(); + for (size_t i = 0; i < std::min(a.size(), b.size()); ++i) + if (a[i] != b[i] && !(skipAlpha && (i & 7) >= 6)) ++n; + return n; +} + +int main(int argc, char** argv) +{ + if (argc < 4) { std::printf("usage: %s WORK WEIGHTS.bin TEXTURE88.bin [REPEATS]\n", argv[0]); return 2; } + const std::string work = argv[1]; + const int repeats = argc > 4 ? atoi(argv[4]) : 20; + const auto params = slurp(work + "/params.bin"), arena0 = slurp(work + "/arena.bin"), tex = slurp(argv[3]); + auto weights = slurp(argv[2]); + const uint32_t* p = (const uint32_t*)params.data(); + const int ex = (int)p[0], ey = (int)p[1], hx = ex / 2, hy = ey / 2; + const uint32_t shift = p[20]; + const int gx = (ex + (int)(shift & 0xffffu) + 7) / 8, gy = (ey + (int)(shift >> 16) + 7) / 8; + if (tex.size() != (size_t)ex * ey * 8) { std::printf("texture 88 is not %dx%d RGBA16F\n", ex, ey); return 2; } + + // the layout changes swin_native.h documents: V tiles [k16][n16] -> [n16][k16], bias [key][query] -> [query][key] + for (uint32_t hd = 0; hd < NH; ++hd) { + uint8_t* wv = weights.data() + W_WV0 + SN_HEAD * hd; + for (uint32_t t = 0; t < (C / 16) * 2; ++t) + for (int k = 0; k < 16; ++k) + for (int n = 0; n < k; ++n) std::swap(wv[256 * t + 16 * k + n], wv[256 * t + 16 * n + k]); + uint16_t* b = (uint16_t*)(weights.data() + W_BIAS0 + SN_HEAD * hd); + std::vector o(96 * 64); + for (int key = 0; key < 96; ++key) + for (int q = 0; q < 64; ++q) o[q * 96 + key] = b[key * 64 + q]; + memcpy(b, o.data(), o.size() * 2); + } + + auto* dParam = upload(params); + auto* dW = upload(weights); + auto* dArena = upload(arena0); + auto* dTex = upload(tex); + const size_t fullBytes = (size_t)ex * ey * 8, halfBytes = (size_t)hx * hy * 8; + auto* d112 = upload({}, fullBytes); + auto* d120 = upload({}, fullBytes); + auto* d128 = upload({}, halfBytes); + const dim3 grid(gx, gy), block(32, 1, 4); + std::printf("grid %d x %d, %d x %d tokens, C %d\n", gx, gy, ex, ey, C); + + swin_native::swin_stage<<>>(dParam, dW, dArena, dTex, d112, d120, d128); + CHECK(hipGetLastError()); + CHECK(hipDeviceSynchronize()); + + bool same = true; + { + std::vector got(arena0.size()); + CHECK(hipMemcpy(got.data(), dArena, got.size(), hipMemcpyDeviceToHost)); + const size_t d = differing(got, slurp(work + "/arena.out.bin"), false); + std::printf(" arena: %zu bytes differ\n", d); + same &= d == 0; + } + const struct { const char* name; swin_native::hf* dev; size_t bytes; } surfaces[] = { + {"surf112.bin", d112, fullBytes}, {"surf120.bin", d120, fullBytes}, {"surf128.bin", d128, halfBytes}}; + for (const auto& s : surfaces) { + std::vector got(s.bytes); + CHECK(hipMemcpy(got.data(), s.dev, got.size(), hipMemcpyDeviceToHost)); + const size_t d = differing(got, slurp(work + "/" + s.name), true); + std::printf(" %s: %zu bytes differ\n", s.name, d); + same &= d == 0; + } + std::printf("%s\n", same ? "IDENTICAL" : "DIFFERENT"); + + // like vulkan/kvk: the median over submits of REPEATS launches each + std::vector t; + for (int it = 0; it < 12; ++it) { + const auto t0 = std::chrono::steady_clock::now(); + for (int r = 0; r < repeats; ++r) swin_native::swin_stage<<>>(dParam, dW, dArena, dTex, d112, d120, d128); + CHECK(hipDeviceSynchronize()); + if (it >= 3) t.push_back(std::chrono::duration(std::chrono::steady_clock::now() - t0).count() / repeats); + } + std::sort(t.begin(), t.end()); + std::printf("%.1f us per launch (median of %zu rounds of %d, min %.1f)\n", t[t.size() / 2], t.size(), repeats, t[0]); + return same ? 0 : 1; +} diff --git a/packaging/d4r.ini b/packaging/d4r.ini index b715539..c6a13b6 100644 --- a/packaging/d4r.ini +++ b/packaging/d4r.ini @@ -27,8 +27,10 @@ ShowWatermark = false ; auto uses the denoiser when available; false disables the route. ; RR D/E temporal validation on gfx1201 uses [Kernels] NativeFp8 = on; scalar FP8 RR is not validated. Enable = auto -; Independent denoiser presets: auto (game choice), default, D or E (transformer). -; E requires a denoiser DLL containing that model. +; Independent denoiser presets: auto (game choice), default, D, E (transformer) or F (DLSS 4.5). +; E and F through NVIDIA's library need a denoiser DLL containing that model (F: 310.9 or later). On the engine +; ([Engine] Enabled = true) F runs from the model in d4r\engine\rr-f; frames the library serves instead show +; preset D when the DLL is older (d4r_nvngx.log says so). Model = auto [Latency] diff --git a/scripts/d4r_config.py b/scripts/d4r_config.py index d850488..f5c0028 100755 --- a/scripts/d4r_config.py +++ b/scripts/d4r_config.py @@ -92,14 +92,14 @@ def main(argv): env["D4R_RR_ENABLE"] = "1" if rr_enabled else "0" rr_model = get("RayReconstruction", "Model") if rr_model is not None: - rr_presets = {"DEFAULT": 0, "D": 4, "E": 5} + rr_presets = {"DEFAULT": 0, "D": 4, "E": 5, "F": 6} # F: DLSS 4.5, nvngx_dlssd.dll 310.9 or later rr_model = rr_model.upper() if rr_model in rr_presets: env["D4R_RR_PRESET"] = str(rr_presets[rr_model]) - elif rr_model in ("0", "4", "5"): + elif rr_model in ("0", "4", "5", "6"): env["D4R_RR_PRESET"] = rr_model else: - fail(f"[RayReconstruction] Model {rr_model!r}: use auto, default, D or E") + fail(f"[RayReconstruction] Model {rr_model!r}: use auto, default, D, E or F") age = get("Latency", "FrameAge") if age is not None: diff --git a/scripts/run_rr_engine_compare.sh b/scripts/run_rr_engine_compare.sh index 3e37f36..f54e403 100755 --- a/scripts/run_rr_engine_compare.sh +++ b/scripts/run_rr_engine_compare.sh @@ -6,7 +6,7 @@ set -euo pipefail # # usage: run_rr_engine_compare.sh OUTPUT_DIR MODEL_DIR # environment: as scripts/run_rr_guide_compare.sh (D4R_PROTON_DIR, D4R_RUNTIME_DIR, D4R_VKD3D_DIR, D4R_ZLUDA_DIR, -# D4R_HARNESS_SIZE, D4R_HARNESS_FRAMES); D4R_RR_PRESET 5 = E (default; the model must match) +# D4R_HARNESS_SIZE, D4R_HARNESS_FRAMES); D4R_RR_PRESET 5 = E (default; the model must match), 6 = F with D4R_RR_DLL=<310.9 nvngx_dlssd.dll> ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" OUTPUT_DIR="$(realpath -m "${1:?usage: run_rr_engine_compare.sh OUTPUT_DIR MODEL_DIR}")" MODEL_DIR="$(realpath "${2:?usage: run_rr_engine_compare.sh OUTPUT_DIR MODEL_DIR}")" @@ -22,6 +22,8 @@ mkdir -p "$APP_DIR" cp -f "$ROOT/build/d3d12_dlss_harness.exe" "$ROOT/build/d4r_nvngx.dll" \ "$RUNTIME_DIR/bin/nvngx_dlssd.dll" "$RUNTIME_DIR/bin/nvngx_dlss.dll" \ "$VKD3D_DIR/d3d12.dll" "$VKD3D_DIR/d3d12core.dll" "$NVAPI_DIR/nvapi64.dll" "$APP_DIR/" +# D4R_RR_DLL: another nvngx_dlssd.dll than the runtime's (preset F needs 310.9: D4R_RR_PRESET=6 with a compile_rr2.py model) +[[ -z "${D4R_RR_DLL:-}" ]] || cp -f "$D4R_RR_DLL" "$APP_DIR/nvngx_dlssd.dll" export D4R_HARNESS_SKIP_BUILD=1 D4R_HARNESS_APP_DIR="$APP_DIR" D4R_HARNESS_NGX_DLL="$APP_DIR/d4r_nvngx.dll" export D4R_RUNTIME_DIR="$RUNTIME_DIR" D4R_ZLUDA_DIR="$ZLUDA_DIR" @@ -32,9 +34,11 @@ export D4R_HARNESS_FEATURE_ID=13 D4R_HARNESS_SAVE_FRAMES=1 export D4R_HARNESS_GUIDE_FORMATS="${D4R_HARNESS_GUIDE_FORMATS:-rgba16f,rgba16f,rgba8unorm,rgba8unorm}" export WINEDLLOVERRIDES="d3d12=n,b;d3d12core=n,b" +# a preset F model also changes the shim's mapping of the hints (every one but 4 and 5 runs F then) +F_DIR="$OUTPUT_DIR/no-preset-f-model"; grep -q "^preset 6" "$MODEL_DIR/rr.txt" && F_DIR="$MODEL_DIR" run_one() { # run_one LABEL ENGINE printf '\n== %s\n' "$1" >&2 - D4R_ENGINE="$2" D4R_ENGINE_RR_SYNC=1 D4R_ENGINE_MODEL_RR_E_DIR="$MODEL_DIR" D4R_ENGINE_MODEL_RR_D_DIR="$MODEL_DIR" \ + D4R_ENGINE="$2" D4R_ENGINE_RR_SYNC=1 D4R_ENGINE_MODEL_RR_E_DIR="$MODEL_DIR" D4R_ENGINE_MODEL_RR_D_DIR="$MODEL_DIR" D4R_ENGINE_MODEL_RR_F_DIR="$F_DIR" \ "$ROOT/scripts/run_d3d12_dlss_harness_proton.sh" "$PROTON_DIR" "$OUTPUT_DIR/$1.raw" \ "$FRAMES" "$IN_W" "$IN_H" "$OUT_W" "$OUT_H" > "$OUTPUT_DIR/$1.log" 2>&1 || true grep -E 'engine backend|engine frame|Ray Reconstruction|output read back|rror' "$OUTPUT_DIR/$1.log" | head -20 || true diff --git a/tests/test_ray_reconstruction.py b/tests/test_ray_reconstruction.py index cfe4a39..bb330ce 100644 --- a/tests/test_ray_reconstruction.py +++ b/tests/test_ray_reconstruction.py @@ -44,8 +44,13 @@ class RayReconstructionConfigTests(unittest.TestCase): self.assertEqual(decided["D4R_RR_PRESET"], "5") # Ray Reconstruction preset E self.assertEqual(decided["D4R_DLSS_PRESET"], "11") # Super Sampling preset K + def test_dlss_45_ray_reconstruction_is_model_f(self): + code, decided = run_config("[RayReconstruction]\nEnable = true\nModel = F\n") + self.assertEqual(code, 0) + self.assertEqual(decided["D4R_RR_PRESET"], "6") # preset F: nvngx_dlssd.dll 310.9 or later + def test_invalid_rr_settings_are_rejected_rather_than_guessed(self): - for text in ("[RayReconstruction]\nModel = F\n", + for text in ("[RayReconstruction]\nModel = G\n", "[RayReconstruction]\nModel = K\n", "[RayReconstruction]\nEnable = maybe\n"): with self.subTest(ini=text): diff --git a/tools/d3d12_dlss_harness.cpp b/tools/d3d12_dlss_harness.cpp index c76af3e..982e91b 100644 --- a/tools/d3d12_dlss_harness.cpp +++ b/tools/d3d12_dlss_harness.cpp @@ -1546,7 +1546,9 @@ int main(int argc, char** argv) if (reconstruction) { d4r_ngx_set_int(parameters, "DLSS.Denoise.Mode", 1); - d4r_ngx_set_int(parameters, "DLSS.Roughness.Mode", 0); + // D4R_HARNESS_RR_ROUGHNESS_MODE=1: the roughness is packed in the normals' w (RE Requiem) + d4r_ngx_set_int(parameters, "DLSS.Roughness.Mode", + std::getenv("D4R_HARNESS_RR_ROUGHNESS_MODE") != nullptr ? std::atoi(std::getenv("D4R_HARNESS_RR_ROUGHNESS_MODE")) : 0); // D4R_HARNESS_RR_HW_DEPTH=1: the depth is the hardware (nonlinear) depth buffer, as Cyberpunk 2077 declares it d4r_ngx_set_int(parameters, "DLSS.Use.HW.Depth", std::getenv("D4R_HARNESS_RR_HW_DEPTH") != nullptr ? std::atoi(std::getenv("D4R_HARNESS_RR_HW_DEPTH")) : 0); diff --git a/tools/d4r_nvngx_shim.cpp b/tools/d4r_nvngx_shim.cpp index f7348ff..4642084 100644 --- a/tools/d4r_nvngx_shim.cpp +++ b/tools/d4r_nvngx_shim.cpp @@ -612,10 +612,12 @@ static void load_portable_config() portable_set("D4R_RR_PRESET", "4"); else if (upper == "E") portable_set("D4R_RR_PRESET", "5"); - else if (upper == "0" || upper == "4" || upper == "5") + else if (upper == "F") // DLSS 4.5 Ray Reconstruction: nvngx_dlssd.dll 310.9 or later + portable_set("D4R_RR_PRESET", "6"); + else if (upper == "0" || upper == "4" || upper == "5" || upper == "6") portable_set("D4R_RR_PRESET", upper); else - g_portable.notes.push_back("d4r.ini: [RayReconstruction] Model must be auto, default, D or E, not '" + + g_portable.notes.push_back("d4r.ini: [RayReconstruction] Model must be auto, default, D, E or F, not '" + rrModel + "'; leaving the model to the game"); } @@ -1344,9 +1346,17 @@ struct Global DenoiserProbe denoiser; bool denoiserAttempted = false; std::wstring denoiserPath; + // nvngx_dlssd.dll's file version (major, minor, build, revision); all zero when it has no version resource + unsigned denoiserVersion[4] = {}; }; static Global g; + +// nvngx_dlssd.dll 310.9 and newer hold preset F (DLSS 4.5 Ray Reconstruction); unknown versions count as without. +static bool denoiser_has_preset_f() +{ + return g.denoiserVersion[0] > 310 || (g.denoiserVersion[0] == 310 && g.denoiserVersion[1] >= 9); +} static std::mutex g_featureApiMutex; template static bool load_export(HMODULE module, const char* name, Function& function) @@ -1608,6 +1618,35 @@ static bool load_denoiser_library() logf("found the DLSS-Denoiser library %ls (mapped as %ls): %s", g.denoiserPath.c_str(), mapped, g.denoiser.cudaCapable ? "CUDA entry points present" : "no CUDA entry points - not a CUDA-capable denoiser build"); + // The file version, from the module's version resource (VS_FIXEDFILEINFO follows its signature 0xFEEF04BD). + // Preset F (DLSS 4.5, the rr2_* kernels) exists from 310.9 on: an older library runs preset D for that hint. + if (HRSRC found = FindResourceW(g.denoiserModule, MAKEINTRESOURCEW(1), MAKEINTRESOURCEW(16))) + { + const DWORD size = SizeofResource(g.denoiserModule, found); + HGLOBAL loaded = LoadResource(g.denoiserModule, found); + const unsigned char* bytes = loaded != nullptr ? static_cast(LockResource(loaded)) : nullptr; + for (DWORD at = 0; bytes != nullptr && at + 16 <= size; at += 4) + { + uint32_t words[4]; + std::memcpy(words, bytes + at, sizeof(words)); + if (words[0] != 0xFEEF04BDu) + continue; + g.denoiserVersion[0] = words[2] >> 16; + g.denoiserVersion[1] = words[2] & 0xffffu; + g.denoiserVersion[2] = words[3] >> 16; + g.denoiserVersion[3] = words[3] & 0xffffu; + break; + } + } + if (g.denoiserVersion[0] == 0) + logf("Ray Reconstruction: the DLSS-Denoiser library has no file version; whether its CUDA path has preset F " + "(DLSS 4.5, 310.9 or newer) is unknown"); + else + logf("Ray Reconstruction: DLSS-Denoiser library version %u.%u.%u.%u: its CUDA path %s", g.denoiserVersion[0], + g.denoiserVersion[1], g.denoiserVersion[2], g.denoiserVersion[3], + denoiser_has_preset_f() ? "has preset F (DLSS 4.5)" + : "has NO preset F (DLSS 4.5 needs 310.9 or newer) - preset F is shown only by " + "frames the engine serves; frames the CUDA backend serves show preset D"); return g.denoiser.cudaCapable; } @@ -2561,6 +2600,7 @@ struct Feature double engineNetWaitEmaUs = 0.0; // worker only: smoothed launch-to-completion wait, for adaptive polling void* engineNet = nullptr; bool engineRefused = false; + bool rrCompileWarned = false; // the "CUDA backend has no preset F" warning was logged for the compile frames float engineJitter[2] = {0, 0}; // of the previous engine frame // Serializes picking the slot to present (game thread) against picking // the slot to overwrite (worker). @@ -8248,54 +8288,62 @@ static bool engine_evaluate_m(Feature& feature, ID3D12GraphicsCommandList* list, return true; } -// Ray Reconstruction models (engine/compile_rr.py): one per preset, engine\rr-d and engine\rr-e next to this DLL, or -// D4R_ENGINE_MODEL_RR_D_DIR / D4R_ENGINE_MODEL_RR_E_DIR. +// Ray Reconstruction models (engine/compile_rr.py; preset F: engine/compile_rr2.py): one per preset, engine\rr-d, +// engine\rr-e and engine\rr-f next to this DLL, or D4R_ENGINE_MODEL_RR_D_DIR / _E_DIR / _F_DIR. Slots: 0 D, 1 E, 2 F. struct EngineBackendR { // The model's pipelines take minutes to compile on a cold shader cache, so each preset's model loads on a thread // of its own; until it is ready the CUDA backend serves the frames. state: 0 not started, 1 loading, 2 ready, // 3 failed. - std::once_flag once[2]; // [preset E] - std::atomic state[2] = {}; - std::unique_ptr models[2]; + std::once_flag once[3]; + std::atomic state[3] = {}; + std::unique_ptr models[3]; }; static EngineBackendR g_engineR; -static void init_engine_backend_rr(bool presetE) +static const char g_rrLetter[3] = {'D', 'E', 'F'}; +static const uint32_t g_rrPreset[3] = {4, 5, 6}; +static std::string rr_model_directory(int slot) { - std::string directory = env_string(presetE ? "D4R_ENGINE_MODEL_RR_E_DIR" : "D4R_ENGINE_MODEL_RR_D_DIR"); + static const char* const variables[3] = {"D4R_ENGINE_MODEL_RR_D_DIR", "D4R_ENGINE_MODEL_RR_E_DIR", "D4R_ENGINE_MODEL_RR_F_DIR"}; + static const wchar_t* const folders[3] = {L"\\engine\\rr-d", L"\\engine\\rr-e", L"\\engine\\rr-f"}; + std::string directory = env_string(variables[slot]); if (directory.empty()) { char folder[MAX_PATH * 3] = {}; - WideCharToMultiByte(CP_ACP, 0, (g_portable.dir + (presetE ? L"\\engine\\rr-e" : L"\\engine\\rr-d")).c_str(), -1, - folder, sizeof(folder), nullptr, nullptr); + WideCharToMultiByte(CP_ACP, 0, (g_portable.dir + folders[slot]).c_str(), -1, folder, sizeof(folder), nullptr, nullptr); directory = folder; } + return directory[0] == '/' ? "Z:" + directory : directory; +} + +static void init_engine_backend_rr(int slot) +{ + const std::string path = rr_model_directory(slot); if (g_vk.interop == nullptr || g_vk.lifetime == nullptr || g_vk.device == VK_NULL_HANDLE || d4r::vkTable.CreateBuffer == nullptr) { logf("engine backend (Ray Reconstruction): Vulkan interop is unavailable"); - g_engineR.state[presetE] = 3; + g_engineR.state[slot] = 3; return; } - const std::string path = directory[0] == '/' ? "Z:" + directory : directory; logf("engine backend (Ray Reconstruction %c): loading the model from %s and compiling its pipelines; the CUDA " - "backend serves frames until it is ready", presetE ? 'E' : 'D', path.c_str()); + "backend serves frames until it is ready", g_rrLetter[slot], path.c_str()); const auto start = std::chrono::steady_clock::now(); try { auto model = std::make_unique(g_vk.physical, g_vk.device, path); - if (model->preset() != (presetE ? 5u : 4u)) - throw std::runtime_error(presetE ? "the directory holds a preset D model" : "the directory holds a preset E model"); - g_engineR.models[presetE] = std::move(model); + if (model->preset() != g_rrPreset[slot]) + throw std::runtime_error("the directory holds a model of another preset"); + g_engineR.models[slot] = std::move(model); } catch (const std::exception& error) { logf("engine backend (Ray Reconstruction): cannot load the model from %s: %s", path.c_str(), error.what()); - g_engineR.state[presetE] = 3; + g_engineR.state[slot] = 3; return; } - g_engineR.state[presetE] = 2; - logf("engine backend (Ray Reconstruction %c): model %s loaded and pipelines compiled in %.0f ms", presetE ? 'E' : 'D', + g_engineR.state[slot] = 2; + logf("engine backend (Ray Reconstruction %c): model %s loaded and pipelines compiled in %.0f ms", g_rrLetter[slot], path.c_str(), std::chrono::duration(std::chrono::steady_clock::now() - start).count()); } @@ -8305,30 +8353,55 @@ static bool engine_evaluate_rr(Feature& feature, ID3D12GraphicsCommandList* list ID3D12Resource* depth, ID3D12Resource* motion, ID3D12Resource* output, void* parameters, const FrameParams& p, uint32_t frame) { + // set below once the preset is known: preset F was asked for but nvngx_dlssd.dll (the CUDA backend) has none + bool cudaLacksF = false; + auto warn_not_f = [&](const char* when) { + if (!cudaLacksF) + return; + logf("WARNING: Ray Reconstruction preset F (DLSS 4.5) was selected, but %s feature %u is served by the CUDA " + "backend, whose nvngx_dlssd.dll %u.%u.%u.%u has no preset F: those frames show preset D, NOT F", + when, feature.handle.Id, g.denoiserVersion[0], g.denoiserVersion[1], g.denoiserVersion[2], + g.denoiserVersion[3]); + }; auto refuse = [&](const char* reason) { feature.engineRefused = true; logf("engine backend (Ray Reconstruction): %s; feature %u uses the CUDA backend from frame %u", reason, feature.handle.Id, frame); + warn_not_f("from this frame on"); return false; }; // NGX 310.7 runs preset E's kernels for the hint 5 only and preset D's (diamond_wallaby) for every other value: // none, the retired A / B / C (1-3), 6, and values with the top bit set such as Cyberpunk 2077's 0x80000002 or // 0x80000005 (harness captures with D4R_RR_PRESET, 2026-10-12). - const bool presetE = feature.preset == 5; - std::call_once(g_engineR.once[presetE], [presetE] { - g_engineR.state[presetE] = 1; + // NGX 310.9 has preset F (DLSS 4.5, the rr2_* kernels) and runs it for every hint but 4 (D) and 5 (E): none, + // 6, Cyberpunk 2077's 0x80000002 (harness runs with D4R_RR_PRESET, 2026-10-13). The shim cannot see which DLL + // version the CUDA backend loaded, so that mapping applies when the user built a preset F model (engine\rr-f), + // which only a 310.9 DLL yields. + static const bool haveF = GetFileAttributesA((rr_model_directory(2) + "\\rr.txt").c_str()) != INVALID_FILE_ATTRIBUTES; + const int slot = feature.preset == 5 ? 1 : (haveF && feature.preset != 4) ? 2 : 0; + std::call_once(g_engineR.once[slot], [slot] { + g_engineR.state[slot] = 1; // D4R_ENGINE_RR_SYNC=1 (harness comparisons): load in place so the engine serves the first frame if (env_uint("D4R_ENGINE_RR_SYNC", 0) != 0) - init_engine_backend_rr(presetE); + init_engine_backend_rr(slot); else - std::thread([presetE] { init_engine_backend_rr(presetE); }).detach(); + std::thread([slot] { init_engine_backend_rr(slot); }).detach(); }); - const int state = g_engineR.state[presetE].load(); + cudaLacksF = slot == 2 && !denoiser_has_preset_f(); + const int state = g_engineR.state[slot].load(); if (state == 1) - return false; // still compiling: the CUDA backend serves this frame, the engine is tried again next frame - d4r::RModel* const model = g_engineR.models[presetE].get(); + { + // still compiling: the CUDA backend serves this frame, the engine is tried again next frame + if (!feature.rrCompileWarned) + { + feature.rrCompileWarned = true; + warn_not_f("while the preset F model compiles,"); + } + return false; + } + d4r::RModel* const model = g_engineR.models[slot].get(); if (state != 2 || model == nullptr) - return refuse(presetE ? "no preset E model (engine\\rr-e)" : "no preset D model (engine\\rr-d)"); + return refuse(slot == 1 ? "no preset E model (engine\\rr-e)" : slot == 2 ? "no preset F model (engine\\rr-f)" : "no preset D model (engine\\rr-d)"); // NVSDK_NGX_DLSS_Feature_Flags: IsHDR 1, MVLowRes 2, MVJittered 4, DepthInverted 8. if ((feature.flags & 0x3) != 0x3) return refuse("needs the HDR flag and render-resolution motion vectors"); @@ -8336,6 +8409,11 @@ static bool engine_evaluate_rr(Feature& feature, ID3D12GraphicsCommandList* list return refuse("jittered motion vectors are not implemented"); if (!p.auxMatrixSet[0] || !p.auxMatrixSet[1]) return refuse("the camera matrices are missing"); + // Preset F below 4/3 scaling (DLAA and the like): NGX runs another output filter there, rr2_post_kernel_3_2, + // which the model does not have (harness runs at 1920x1080 output, 2026-10-13: 1440x810 and smaller render sizes + // use rr2_post_kernel, 1456x820 and larger rr2_post_kernel_3_2) + if (slot == 2 && (3ull * feature.outWidth < 4ull * p.renderWidth || 3ull * feature.outHeight < 4ull * p.renderHeight)) + return refuse("preset F below 4/3 scaling uses an output filter the engine does not have (rr2_post_kernel_3_2)"); if (feature.engineR != nullptr && (p.renderWidth != feature.engineMRenderWidth || p.renderHeight != feature.engineMRenderHeight)) return refuse("changing the effective render resolution after initialization is not implemented"); @@ -8343,15 +8421,16 @@ static bool engine_evaluate_rr(Feature& feature, ID3D12GraphicsCommandList* list {"DLSS.Input.DiffuseAlbedo", "GBuffer.DiffuseAlbedo"}, {"DLSS.Input.SpecularAlbedo", "GBuffer.SpecularAlbedo"}}; ID3D12Resource* guides[4] = {}; - // DLSS.Roughness.Mode 1 (packed): no roughness texture, the roughness is the normals' w - const bool roughnessPacked = get_resource(parameters, "GBuffer.Roughness") == nullptr && - get_uint_or(parameters, "DLSS.Roughness.Mode", 0) == 1; + // DLSS.Roughness.Mode 1 (packed): the roughness is the normals' w. Preset F follows the mode also when a roughness + // texture is bound (harness run against the CUDA path, 2026-10-13); D and E are only known without one (RE Requiem) + const bool roughnessPacked = get_uint_or(parameters, "DLSS.Roughness.Mode", 0) == 1 && + (slot == 2 || get_resource(parameters, "GBuffer.Roughness") == nullptr); for (int i = 0; i < 4; ++i) { guides[i] = get_resource(parameters, guideNames[i][0]); if (guides[i] == nullptr && guideNames[i][1] != nullptr) guides[i] = get_resource(parameters, guideNames[i][1]); - if (guides[i] == nullptr && i == 1 && roughnessPacked) + if (i == 1 && roughnessPacked) guides[i] = guides[0]; if (guides[i] == nullptr) return refuse("a guide the engine needs (normals, roughness, diffuse or specular albedo) is missing"); @@ -8474,7 +8553,7 @@ static bool engine_evaluate_rr(Feature& feature, ID3D12GraphicsCommandList* list feature.engineMRenderWidth = rw; feature.engineMRenderHeight = rh; logf("engine backend (Ray Reconstruction %c): feature %u runs in the game's command list (%ux%u -> %ux%u)", - presetE ? 'E' : 'D', feature.handle.Id, rw, rh, feature.outWidth, feature.outHeight); + g_rrLetter[slot], feature.handle.Id, rw, rh, feature.outWidth, feature.outHeight); } } catch (const std::exception& error) @@ -8515,7 +8594,7 @@ static bool engine_evaluate_rr(Feature& feature, ID3D12GraphicsCommandList* list return refuse(failure); if (frame <= 3 || frame % 600 == 0) logf("engine frame %u recorded, Ray Reconstruction %c (jitter %.6f,%.6f, reset %d, mvScale %.3f,%.3f)", frame, - presetE ? 'E' : 'D', p.jitterX, p.jitterY, s.reset, p.mvScaleX, p.mvScaleY); + g_rrLetter[slot], p.jitterX, p.jitterY, s.reset, p.mvScaleX, p.mvScaleY); return true; }