From b9254f7f9667256afdb283dad3f4ba09ffb4771b Mon Sep 17 00:00:00 2001 From: woutermans Date: Sun, 18 Oct 2026 20:38:06 +0200 Subject: [PATCH] engine: Super Resolution presets E and F, the DLSS 3 CNN, on RDNA2 Presets E (5) and F (6) run on d4r's own Vulkan engine without cooperative matrices: the input and reconstruction kernels translated from the user's DLL (engine/compile_sr_e.py, rr/ptx2rr.py), the network engine/sr/e_net.comp with packed v_pk_fma_f16 accumulation (engine/cnn_sr/mma_f16.glsl, the default head; --head-accumulation f32 keeps the f32 head). Runtime runtime_e.inc (EModel/EEngine), game adapter game_e.cpp, replay d4r-e; the shim serves them from engine\e and engine\f (D4R_ENGINE_MODEL_E_DIR / _F_DIR), LDR or HDR, any render size, not with AutoExposure. build-engine-models.sh builds both. Translator fixes that apply to every translated kernel: .relu on f16/f16x2 fma/add/mul was ignored; add.rz was always taken for the roundf idiom (lost the log1pf correction term); formatted image stores expand at the call site (SPIR-V image type fix); f16 mma.sync m16n8k8/k16, .const loads, shared 16-bit stores, aligned 32-bit arena words (D4R_WORD_ARENA), unified arena. Steam Deck gfx1033: 2.55 ms at 640x360 -> 1280x720, 6.20 ms at 1080p output; Rounded Quality sizes take mid_3x3 (reconstruction 2.44 -> 1.23 ms). Ray Reconstruction CNN presets A-C (compile_rr3.py, the packed RR flavor) stay out of this commit; engine/README.md already documents their tuning. --- config/d4r.ini.default | 14 +- engine/README.md | 109 +++++++ engine/build.sh | 6 +- engine/cnn_sr/mma_f16.glsl | 70 ++++ engine/compile_e.py | 189 +++++++++++ engine/compile_sr_e.py | 207 ++++++++++++ engine/game_e.cpp | 55 ++++ engine/game_e.h | 28 ++ engine/replay_e.cpp | 196 +++++++++++ engine/rr/ptx2rr.py | 146 +++++++-- engine/runtime.cpp | 16 + engine/runtime.h | 47 +++ engine/runtime_e.inc | 490 ++++++++++++++++++++++++++++ engine/sr/e_net.comp | 408 +++++++++++++++++++++++ engine/vk_dispatch.cpp | 5 + engine/vk_dispatch.h | 10 + packaging/build_engine_models.sh | 18 +- packaging/d4r.ini | 5 +- scripts/build_d4r_nvngx_shim.sh | 2 +- scripts/stage_engine_kit.sh | 4 +- tests/test_cnn_mma_gpu.py | 93 ++++++ tests/test_cnn_shader_types.py | 52 +++ tests/test_sr_reconstruction_gpu.py | 72 ++++ tools/d4r_nvngx_shim.cpp | 180 +++++++++- vulkan/ptx2glsl.py | 20 +- 25 files changed, 2406 insertions(+), 36 deletions(-) create mode 100644 engine/cnn_sr/mma_f16.glsl create mode 100755 engine/compile_e.py create mode 100644 engine/compile_sr_e.py create mode 100644 engine/game_e.cpp create mode 100644 engine/game_e.h create mode 100644 engine/replay_e.cpp create mode 100644 engine/runtime_e.inc create mode 100644 engine/sr/e_net.comp create mode 100644 tests/test_cnn_mma_gpu.py create mode 100644 tests/test_cnn_shader_types.py create mode 100644 tests/test_sr_reconstruction_gpu.py diff --git a/config/d4r.ini.default b/config/d4r.ini.default index a621ce8..ea8b61f 100644 --- a/config/d4r.ini.default +++ b/config/d4r.ini.default @@ -25,9 +25,12 @@ GameExe = ; K | 4 | DLAA / Quality / Balanced | Default; faster alternative to M/L ; M | 4.5 | Performance | More demanding on RDNA3 ; L | 4.5 | Ultra Performance | Especially 4K; experimental -; E | 3 CNN | Quality / Balanced / Performance | Legacy faster fallback +; E | 3 CNN | Quality / Balanced / Performance | Legacy faster fallback; the engine's choice on RDNA2 ; F | 3 CNN | DLAA / Ultra Performance | Deprecated; prefer K or L ; +; On the engine ([Engine]) E and F run from engine\e and engine\f (build-engine-models.sh) without cooperative +; matrices, so they are the models for RDNA2 GPUs (Steam Deck: about 5 ms at 640x360 -> 1280x720). +; ; Select the upscaling mode separately in game/OptiScaler. These are recommendations; ; other modes are allowed. L requires the L texture build. ; Aliases: CNN = E, DLSS4 = K, DLSS4.5 = M. A/B/C/D/G/J use K in this DLL version. @@ -134,12 +137,15 @@ EvalSync = false MarkerPollUs = auto [Engine] -; Experimental. Runs DLSS 4 and 4.5 (models K, L and M) on d4r's own Vulkan inference engine instead of NVIDIA's +; Experimental. Runs DLSS 4 and 4.5 (models K, L and M) and DLSS 3's CNN (models E and F) on d4r's own Vulkan +; inference engine instead of NVIDIA's ; kernels through ZLUDA: the whole evaluation is recorded into the game's own command list, the game shows the ; result of the frame it just rendered, and FrameAge and the [Interop] hand-off settings do not apply. ; Needs VkD3DDir, a driver with VK_KHR_cooperative_matrix (Mesa RADV on RDNA3/RDNA4) and the model folders -; (RuntimeDir/engine/k, k-ldr, m, m-ldr, l, l-ldr; scripts/install_d4r_runtime.sh). A game or setting the engine -; does not cover uses the normal path, with the reason in d4r_nvngx.log: other models than K, L and M, +; (RuntimeDir/engine/k, k-ldr, m, m-ldr, l, l-ldr; scripts/install_d4r_runtime.sh). E and F (engine/e, engine/f; +; D4R_ENGINE_MODEL_E_DIR, D4R_ENGINE_MODEL_F_DIR) need no cooperative matrices: they run on RDNA2 (Steam Deck) +; too, from any render size and HDR or LDR, but not with AutoExposure. A game or setting the engine +; does not cover uses the normal path, with the reason in d4r_nvngx.log: other models than E, F, K, L and M, ; display-resolution motion vectors with L or M, or a K model built before the current engine/compile_k.py for display-resolution ; motion, non-inverted depth or exposure supplied by the game. Details and measurements: engine/README.md. ; (D4R_ENGINE) diff --git a/engine/README.md b/engine/README.md index e0454d6..a4fb3eb 100644 --- a/engine/README.md +++ b/engine/README.md @@ -1121,6 +1121,114 @@ Ultra Performance, the downsample weights and the network's summation order (M i backend in the harness as well: 1706x960 -> 2560x1440 and 640x360 -> 1920x1080, where it measured 36-41 dB with the native accuracy kernels' order; `compile_m.py` builds both presets' networks with `TEXL`). Details: `L_NOTES.md`. +## Presets E and F: DLSS 3's CNN (runtime and shim path exist; experimental, the RDNA2 route) + +Presets E (5) and F (6) are NGX 310.7's CNN upscaler, in the `nvngx_dlss.dll` the installer already fetches. Its +frame is an input kernel (`cuda_engine_input_kernel_rel_{ldr,hdr}_mvdiff_{mvhi,mvlo}`: the game's colour, depth and +motion, the previous output and motion to 16 channels per token of 4x4 output pixels, the warped history and this +frame's output-size motion), eleven `dltss_nchw8_*` network kernels (a U-Net over the token grid, the output padded +to 128 / 4, five levels) and a reconstruction kernel (`..._aniso_gaussian_dfn_{fastpath_3x3_128, mid_3x3_128, +3x3_256, fastpath_5x5_256, 5x5_256}`: the last decoder level, a 40-channel head per token and the anisotropic +Gaussian reconstruction of the output and the next history). E and F share the kernels; they differ in weights and +in two constants. + +- The network is d4r's own `sr/e_net.comp` (five modes: the first encoder, the 3x3 + pool encoders, 1x1 convs, the + upsampling 1x1, the decoders' double 1x1), packed binary16 like the RR CNN's packed flavor (`rr/cnn_pk.glsl`): a + lane owns a pixel, the weights are wave-uniform scalar loads feeding `v_pk_fma_f16`, the input of a workgroup is + staged in LDS. No cooperative matrices, so it is the engine's route on RDNA2. It writes NGX's tensor layouts + (`[C/8][H][W][8]`, the full-size skips as 2x2 phases, 8-bit activations as binary16 bits 14..7), which the + translated reconstruction kernel reads. +- The input and reconstruction kernels are translated from the user's DLL (`rr/ptx2rr.py`) in every variant; the + runtime picks a variant for the ratio (5x5 above 2.5x, `fastpath_3x3` from 1.6x, `mid` when each render axis is + floor/ceil of two thirds of its output axis, the generic kernels otherwise) and rewrites every frame-dependent + field of their argument blocks (sizes, ratios, + jitter, reset, MV scale, the HDR and MVLowRes flags; harness captures at seven sizes and ratios, 2026-10-18). +- `compile_sr_e.py REPLAY NVNGX_DLSS_DLL OUT --preset E|F` builds a model from one frame's replay dump (the weights are + packed in the DLL; NGX uploads them unpacked), `runtime_e.inc` is `EModel` / `EEngine`, `game_e.cpp` the game + adapter, `replay_e.cpp` the replay (`d4r-e`). The shim runs presets 5 and 6 with `engine\e` / `engine\f` + (`D4R_ENGINE_MODEL_E_DIR`, `_F_DIR`), LDR or HDR, either motion resolution, any render size (a new size rebuilds + the engine); not with AutoExposure. `build-engine-models.sh` builds both. + +With the f32 reconstruction head, accuracy against NGX through ZLUDA (three temporal frames): 63-68 dB at 640x360 -> 1280x720 +(LDR and HDR, both motion resolutions, E and F), 1280x720 -> 1920x1080, 1001x563 -> 1501x845 (generic kernel), +427x240 -> 1280x720 (5x5), 742x418 -> 1280x720; DLAA 57-73 dB. Network tensors 45-60 dB each (binary16 accumulation). +The translation needed two `vulkan/ptx2glsl.py` fixes that apply to every translated kernel: `.relu` on f16 / f16x2 +`fma`, `add` and `mul` was ignored, and `add.rz` was always taken for the `roundf` idiom (CUDA's `log1pf` starts with +an `add.rz`, which lost its second operand). + +Steam Deck (gfx1033, RADV, best of 12), 640x360 -> 1280x720: 4.3 ms (input 0.78, network 1.86, reconstruction 1.68); +1280x720 -> 1920x1080: 10.1 ms. The network's ideal at the Deck's packed-f16 rate is 0.75 ms at 720p; the two +translated kernels are the larger share (the reconstruction is 5400 VALU instructions as translated, its emulated +mma about 0.3 ms of it), so native input and reconstruction kernels are the next step. + +RDNA2 tuning (2026-10-18): the SR decoders now reuse the input's padded LDS rows for their hidden layer, +with a workgroup barrier before overwriting the input. This removes a second allocation and its bank conflicts. +RR's packed CNN uses smaller tiles on the stages limited by LDS occupancy; its decoder filters cache each bilinear +sample once per tile, retaining the original half-precision operation order. Rebuild the model directories with +`compile_sr_e.py` / `compile_rr3.py` to use these shaders; the existing runtime selects the packed RR flavor on RDNA2. + +Steam Deck gfx1033, RADV, GPU fixed temporarily at its stock 1600 MHz, interleaved baseline/candidate runs, +best of 12 submissions per frame (four temporal frames for RR, three for SR), 640x360 -> 1280x720: + +| Model | Before tuning, ms | After tuning, ms | +|---|---:|---:| +| SR E, complete frame | 3.13 | 3.09 | +| SR E, network | 1.24 | 1.20 | +| RR A, complete frame | 14.61 | 12.48 | +| RR A, input + network | 12.98 | 10.85 | + +These are inference timings without output readback, with no other GPU workload; automatic clocks and games can +run slower. The changes preserve the packed path's arithmetic: all outputs and histories match byte for byte at +one submission per frame and in repeated runs. Additional checks cover RR's partially filled tiles (70x66 render), +and SR's HDR, output-resolution motion, odd sizes (1001x563 -> 1501x845), 1080p output and 5x5 reconstruction. +The packed convolutions still differ from NVIDIA's matrix accumulation as described above. The translated SR +input/reconstruction passes still take most of the frame, so these results do not establish hardware-optimal +performance. `tests/test_cnn_packed_gpu.py` checks decoder reuse and partial pixel sets against a CPU reference +without proprietary weights (`D4R_RUN_GPU_TESTS=1 D4R_EVK=/path/to/evk python3 -m unittest tests.test_cnn_packed_gpu`). + +The reconstruction head now defaults to packed binary16 accumulation too (`D4R_MMA_PACKED` in +`cnn_sr/mma_f16.glsl`, enabled only for the SR model builder). Each lane computes two columns together with +`v_pk_fma_f16`, retaining the fragment layout and K order without expanding operands and accumulators to f32. +The translated stages also use aligned 32-bit arena accesses for packed words. This removes pairs of scalar +16-bit loads/stores; their binding and tensor layouts stay the same. Formatted image stores expand at the call +site, retaining the concrete SPIR-V image format instead of passing it through an unformatted function parameter. +All four input and five reconstruction variants pass `spirv-val`. + +On the same Deck at a fixed stock 1600 MHz, best of 12 submissions per frame, the complete E pass with the +previous f32 head versus the new packed head (three-frame input sequence): + +| Render -> output | f32 head, ms | Packed head, ms | +|---|---:|---:| +| 640x360 -> 1280x720, LDR | 3.09 | 2.55 | +| 640x360 -> 1280x720, HDR | 3.09 | 2.66 | +| 1001x563 -> 1501x845, generic reconstruction | 6.28 | 5.30 | +| 1280x720 -> 1920x1080 | 7.15 | 6.20 | +| 427x240 -> 1280x720, 5x5 reconstruction | 5.10 | 4.02 | + +At 640x360 -> 1280x720 the reconstruction pass itself falls from about 1.35 to 0.80 ms. The packed head rounds +after each half FMA, so its image is **not byte-identical** to the previous head. At one submission per frame, +eight tested configurations (also HDR, both motion resolutions, balanced scaling and DLAA) give 64.9-92.6 dB +RGB PSNR against the f32-head engine over three temporal frames, with finite outputs throughout. These comparisons +measure the arithmetic change, not quality parity with NGX or a long gameplay sequence. `--head-accumulation f32` +retains the previous head arithmetic, with the word-access optimization; the default is `f16` for both E and F. +`tests/test_cnn_mma_gpu.py` checks both accumulation modes, fragment mapping and independent concurrent waves +against CPU references. Rebuild model shaders to pick up the new default; old model directories continue to run. + +Rounded Quality dimensions also use the existing `mid_3x3_128` reconstruction kernel: each render axis must be +floor or ceil of two thirds of its output axis. The kernel still samples with the actual ratios and jitter. +This avoids the generic reconstruction fallback for 853x533 -> 1280x800 without changing the model arithmetic. +On the Deck at stock 1600 MHz, three interleaved before/after runs (120 temporal frames each, 30 warm-up frames +discarded, 270 measured samples per build) reduce median reconstruction from 2.44 to 1.23 ms and the complete +upscaling pass from 4.46 to 3.26 ms. This change only needs a runtime rebuild; existing model shaders work. +`tests/test_sr_reconstruction_gpu.py` compares automatic selection against forced generic reconstruction, including +mixed per-axis rounding, partial tiles, HDR/LDR, both motion resolutions, jitter, history and reset. On the Deck, +all 384 preset E test frames match bit for bit; an optional private preset F model enables the same checks for F. + +```sh +python3 engine/compile_sr_e.py "$CAP/replay" nvngx_dlss.dll build/e-model --preset E +build/native-engine/d4r-e build/e-model INPUT_DIR OUT_DIR # D4R_ENGINE_E_PROFILE=1: microseconds per pass +``` + ## Files | File | Role | @@ -1142,6 +1250,7 @@ native accuracy kernels' order; `compile_m.py` builds both presets' networks wit | `l_model.py`, `test_l.py`, `L_NOTES.md` | Preset L: weights of the input and expansion stages, their tests against a capture, the algorithms | | `enc0_features.glsl` | Presets M and L: the input stage's 32 per-pixel feature slots | | `exposure_k.comp` | Automatic exposure: log-luma block sums, then the exposure state and the stages' exposure fields | +| `sr/e_net.comp`, `compile_sr_e.py`, `runtime_e.inc`, `game_e.h`, `game_e.cpp`, `replay_e.cpp` | Presets E and F: the network shader, the model builder, the runtime, the game adapter and the replay (`d4r-e`) | | `test_frame.py`, `test_model.py` | Temporal result checks and model corruption tests | | `test_game.cpp` | Queued direct-texture, rectangle, layout, conversion and precision regression checks | | `test_k.py`, `test_in*.py`, `test_out*.py` | Isolated shader diagnosis and reference comparisons | diff --git a/engine/build.sh b/engine/build.sh index 912304e..1b880d4 100755 --- a/engine/build.sh +++ b/engine/build.sh @@ -12,11 +12,13 @@ if [[ -f "$ROCM/include/hip/hip_runtime_api.h" ]]; then HIP=(-DD4R_ENGINE_HIP "$ROOT/engine/hip_net.cpp" -I"$ROCM/include" -L"$ROCM/lib" -lamdhip64 "-Wl,-rpath,$ROCM/lib") fi "$CXX" -std=c++17 -O3 -Wall -Wextra -Wno-missing-field-initializers -fPIC -shared \ - "$ROOT/engine/runtime.cpp" "$ROOT/engine/game.cpp" "$ROOT/engine/game_m.cpp" "$ROOT/engine/game_rr.cpp" -lvulkan -o "$OUT/libd4r_engine.so" + "$ROOT/engine/runtime.cpp" "$ROOT/engine/game.cpp" "$ROOT/engine/game_m.cpp" "$ROOT/engine/game_rr.cpp" "$ROOT/engine/game_e.cpp" -lvulkan -o "$OUT/libd4r_engine.so" "$CXX" -std=c++17 -O3 -Wall -Wextra -Wno-missing-field-initializers \ "$ROOT/engine/replay.cpp" "${HIP[@]}" -L"$OUT" -ld4r_engine -lvulkan '-Wl,-rpath,$ORIGIN' -o "$OUT/d4r-k" "$CXX" -std=c++17 -O3 -Wall -Wextra -Wno-missing-field-initializers \ "$ROOT/engine/replay_m.cpp" "${HIP[@]}" -L"$OUT" -ld4r_engine -lvulkan '-Wl,-rpath,$ORIGIN' -o "$OUT/d4r-m" "$CXX" -std=c++17 -O3 -Wall -Wextra -Wno-missing-field-initializers \ "$ROOT/engine/replay_rr.cpp" -L"$OUT" -ld4r_engine -lvulkan '-Wl,-rpath,$ORIGIN' -o "$OUT/d4r-rr" -printf 'Built %s, %s and %s\n' "$OUT/d4r-k" "$OUT/d4r-m" "$OUT/d4r-rr" +"$CXX" -std=c++17 -O3 -Wall -Wextra -Wno-missing-field-initializers \ + "$ROOT/engine/replay_e.cpp" -L"$OUT" -ld4r_engine -lvulkan '-Wl,-rpath,$ORIGIN' -o "$OUT/d4r-e" +printf 'Built %s, %s, %s and %s\n' "$OUT/d4r-k" "$OUT/d4r-m" "$OUT/d4r-rr" "$OUT/d4r-e" diff --git a/engine/cnn_sr/mma_f16.glsl b/engine/cnn_sr/mma_f16.glsl new file mode 100644 index 0000000..945a9ff --- /dev/null +++ b/engine/cnn_sr/mma_f16.glsl @@ -0,0 +1,70 @@ +// Extension-free f16 MMA for wave32. PTX m16n8k8 / m16n8k16 row.col fragments: +// A: (row=g+8*(word&1), k=2*t+8*(word/2)+half). +// B: (k=2*t+8*word+half, col=g). C/D: (row=g+8*word, col=2*t+half). +// g=lane/4, t=lane%4. Shuffle sources stay within the current subgroup; +// there is no workgroup scratch or cross-wave synchronization. +// Include engine/round_half.glsl first: RADV's packHalf2x16 truncates instead of rounding to nearest. +#ifndef D4R_MMA_F16_INCLUDED +#define D4R_MMA_F16_INCLUDED + +#ifdef D4R_MMA_PACKED +// The SR reconstruction head uses the same packed-half accumulation as e_net.comp. +// Each lane owns two rows and two columns; one packed FMA updates a row's two columns. +// Keep the fragment mapping and K order, rounding after each FMA instead of each MMA. +f16vec4 mma_packed_step8(uint a0, uint a1, uint b0, f16vec4 acc) +{ + const uint lane = gl_SubgroupInvocationID, row = lane >> 2u, col = (lane & 3u) * 2u; + [[unroll]] for (uint k = 0u; k < 4u; ++k) { + const f16vec2 x0 = unpackFloat2x16(subgroupShuffle(a0, row * 4u + k)); + const f16vec2 x1 = unpackFloat2x16(subgroupShuffle(a1, row * 4u + k)); + const f16vec2 y0 = unpackFloat2x16(subgroupShuffle(b0, col * 4u + k)); + const f16vec2 y1 = unpackFloat2x16(subgroupShuffle(b0, (col + 1u) * 4u + k)); + acc.xy = fma(f16vec2(x0.x), f16vec2(y0.x, y1.x), acc.xy); + acc.xy = fma(f16vec2(x0.y), f16vec2(y0.y, y1.y), acc.xy); + acc.zw = fma(f16vec2(x1.x), f16vec2(y0.x, y1.x), acc.zw); + acc.zw = fma(f16vec2(x1.y), f16vec2(y0.y, y1.y), acc.zw); + } + return acc; +} +void mma_f16_m16n8k8(uint a0, uint a1, uint b0, uint c0, uint c1, out uint d0, out uint d1) +{ + const f16vec4 acc = mma_packed_step8(a0, a1, b0, f16vec4(unpackFloat2x16(c0), unpackFloat2x16(c1))); + d0 = packFloat2x16(acc.xy); d1 = packFloat2x16(acc.zw); +} +void mma_f16_m16n8k16(uint a0, uint a1, uint a2, uint a3, uint b0, uint b1, + uint c0, uint c1, out uint d0, out uint d1) +{ + f16vec4 acc = mma_packed_step8(a0, a1, b0, f16vec4(unpackFloat2x16(c0), unpackFloat2x16(c1))); + acc = mma_packed_step8(a2, a3, b1, acc); + d0 = packFloat2x16(acc.xy); d1 = packFloat2x16(acc.zw); +} +#else +vec4 mma_f16_step8(uint a0, uint a1, uint b0, vec4 acc) +{ + const uint lane = gl_SubgroupInvocationID, row = lane >> 2u, col = (lane & 3u) * 2u; + [[unroll]] for (uint k = 0u; k < 4u; ++k) { + const vec2 x0 = unpackHalf2x16(subgroupShuffle(a0, row * 4u + k)); + const vec2 x1 = unpackHalf2x16(subgroupShuffle(a1, row * 4u + k)); + const vec2 y0 = unpackHalf2x16(subgroupShuffle(b0, col * 4u + k)); + const vec2 y1 = unpackHalf2x16(subgroupShuffle(b0, (col + 1u) * 4u + k)); + acc.x = fma(x0.x, y0.x, acc.x); acc.x = fma(x0.y, y0.y, acc.x); + acc.y = fma(x0.x, y1.x, acc.y); acc.y = fma(x0.y, y1.y, acc.y); + acc.z = fma(x1.x, y0.x, acc.z); acc.z = fma(x1.y, y0.y, acc.z); + acc.w = fma(x1.x, y1.x, acc.w); acc.w = fma(x1.y, y1.y, acc.w); + } + return acc; +} +void mma_f16_m16n8k8(uint a0, uint a1, uint b0, uint c0, uint c1, out uint d0, out uint d1) +{ + const vec4 acc = mma_f16_step8(a0, a1, b0, vec4(unpackHalf2x16(c0), unpackHalf2x16(c1))); + d0 = packHalf2x16(roundHalfValue(acc).xy); d1 = packHalf2x16(roundHalfValue(acc).zw); +} +void mma_f16_m16n8k16(uint a0, uint a1, uint a2, uint a3, uint b0, uint b1, + uint c0, uint c1, out uint d0, out uint d1) +{ + vec4 acc = mma_f16_step8(a0, a1, b0, vec4(unpackHalf2x16(c0), unpackHalf2x16(c1))); + acc = mma_f16_step8(a2, a3, b1, acc); + d0 = packHalf2x16(roundHalfValue(acc).xy); d1 = packHalf2x16(roundHalfValue(acc).zw); +} +#endif // D4R_MMA_PACKED +#endif diff --git a/engine/compile_e.py b/engine/compile_e.py new file mode 100755 index 0000000..0dd5220 --- /dev/null +++ b/engine/compile_e.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +"""Compile an E/F CNN capture into Vulkan stages and a compact, alias-preserving tensor arena. + +Usage: compile_e.py PTX_DIR CAPTURE_DIR OUTPUT_DIR [--preset E|F] [--hdr] [--frame N] +PTX and all model payloads are derived from the user's DLL and must not be redistributed. +""" +import argparse +import json +from pathlib import Path +import re +import struct +import subprocess +import sys + +HERE = Path(__file__).resolve().parent + + +def manifest(path): + result = {'allocs': {}, 'pointers': [], 'textures': [], 'surfaces': []} + for line in path.read_text().splitlines(): + p = line.split() + if not p: + continue + if p[0] == 'kernel': + result['name'] = p[1] + elif p[0] == 'launch': + result['launch'] = list(map(int, p[1:7])) + elif p[0] == 'alloc': + result['allocs'][int(p[1])] = (int(p[2], 16), int(p[3])) + elif p[0] == 'pointer': + result['pointers'].append(tuple(map(int, p[1:4]))) + elif p[0] in ('texture', 'surface'): + result[p[0] + 's'].append(int(p[1])) + return result + + +def dump_shape(cap, stage, offset): + number = int(stage['directory'].name.split('-')[1]) + prefix = f'launch-{number:03d}-{stage["name"]}-arg{offset:03d}-' + paths = sorted((cap / 'dump').glob(prefix + '*')) + for path in paths: + m = re.search(r'-(texture|surface)-(\d+)x(\d+)x(\d+)-fmt(\d+)\.bin$', path.name) + if m: + _, w, h, channels, dtype = m.groups() + channels, dtype = int(channels), int(dtype) + if channels not in (1, 2, 4) or dtype not in (16, 32): + raise ValueError(f'unsupported captured image: {path.name}') + fmt = {1: 'r', 2: 'rg', 4: 'rgba'}[channels] + str(dtype) + 'f' + return {'width': int(w), 'height': int(h), 'format': fmt} + return None + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument('ptx_dir', type=Path) + ap.add_argument('capture_dir', type=Path) + ap.add_argument('output', type=Path) + ap.add_argument('--preset', choices=('E', 'F'), default='E') + ap.add_argument('--hdr', action='store_true') + ap.add_argument('--frame', type=int, default=1) + a = ap.parse_args() + stages = [] + for d in sorted((a.capture_dir / 'replay').glob('replay-*')): + s = manifest(d / 'manifest.txt') + s['directory'] = d + stages.append(s) + starts = [i for i, s in enumerate(stages) if s['name'].startswith('cuda_engine_input_kernel_')] + if not 1 <= a.frame <= len(starts): + ap.error(f'capture contains {len(starts)} CNN frames') + first = starts[a.frame - 1] + end = starts[a.frame] if a.frame < len(starts) else len(stages) + stages = [s for s in stages[first:end] if s['name'].startswith(('cuda_engine_input_kernel_', 'dltss_nchw8_'))] + if len(stages) != 13 or 'aniso_gaussian_dfn' not in stages[-1]['name']: + raise ValueError('capture does not contain the complete CNN input/network/reconstruction graph') + if ('_ldr_' not in stages[0]['name']) != a.hdr: + raise ValueError('--hdr does not match the captured input kernel') + color = dump_shape(a.capture_dir, stages[0], 280) + output = dump_shape(a.capture_dir, stages[-1], 464) + if color is None or output is None: + raise ValueError('capture lacks the current color or final output image') + rw, rh, ow, oh = color['width'], color['height'], output['width'], output['height'] + + # Merge only overlapping intervals, never CUDA address-space holes. A nonzero arena + # prefix preserves null-pointer tests when a valid tensor was the first allocation. + intervals = sorted({v for s in stages for v in s['allocs'].values()}) + regions = [] + for base, size in intervals: + if regions and base < regions[-1]['end']: + regions[-1]['end'] = max(regions[-1]['end'], base + size) + else: + regions.append({'base': base, 'end': base + size}) + arena_size = 256 + for r in regions: + r['offset'] = arena_size + arena_size = (arena_size + r['end'] - r['base'] + 255) & ~255 + r['mutable'] = False + def region_for(address): + return next(r for r in regions if r['base'] <= address < r['end']) + def arena_offset(address): + r = region_for(address) + return r['offset'] + address - r['base'] + + modules = {} + for p in sorted(a.ptx_dir.glob('*.ptx')): + text = p.read_text() + for name in re.findall(r'\.entry\s+(\w+)\s*\(', text): + if name not in modules or '.target sm_89' in text[:400]: + modules[name] = p + a.output.mkdir(parents=True, exist_ok=True) + cache = {} + lines = ['D4RE0001', f'size {rw} {rh} {ow} {oh}', f'preset {5 if a.preset == "E" else 6}', f'hdr {int(a.hdr)}', + f'buffer 0 {arena_size} arena-init.bin'] + diagnostics = [] + for index, s in enumerate(stages): + name, launch = s['name'], s['launch'] + if name not in cache: + generated = a.output / (name + '.generated.comp') + command = [sys.executable, str(HERE / 'rr/ptx2rr.py'), str(modules[name]), name, str(generated), + '--block', ','.join(map(str, launch[3:6])), '--weights-arg', '999', '--unified-arena'] + command += ['--buffer=' + str(off) for off, _, _ in s['pointers']] + for off in s['surfaces']: + shape = dump_shape(a.capture_dir, s, off) + if shape: + command += [f'--surface={off}={shape["format"]}'] + subprocess.run(command, check=True) + spv = a.output / (name + '.spv') + subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', str(generated), '-o', str(spv)], check=True) + info = json.loads(generated.with_suffix('.json').read_text()) + constants = generated.with_suffix('.globals.bin') + globals_name = name + '.globals.bin' + (a.output / globals_name).write_bytes(constants.read_bytes() if constants.exists() else bytes(16)) + cache[name] = (info, globals_name) + info, globals_name = cache[name] + for off, allocation, sub in s['pointers']: + if off in info['written_buffers']: + base, _ = s['allocs'][allocation] + region_for(base + sub)['mutable'] = True + params = bytearray((s['directory'] / 'args.bin').read_bytes()) + for off, allocation, sub in s['pointers']: + base, _ = s['allocs'][allocation] + struct.pack_into(' {ow}x{oh}') + + +if __name__ == '__main__': + main() diff --git a/engine/compile_sr_e.py b/engine/compile_sr_e.py new file mode 100644 index 0000000..80e9307 --- /dev/null +++ b/engine/compile_sr_e.py @@ -0,0 +1,207 @@ +#!/usr/bin/env python3 +"""Build the engine's DLSS Super Resolution CNN model directory (presets E and F) from a one-frame dump. + +usage: compile_sr_e.py REPLAY_DIR NVNGX_DLSS_DLL OUT_DIR --preset E|F [--ptx-dir DIR] + +REPLAY_DIR: the replay dump (D4R_CUDA_REPLAY_DUMP_DIR) of one frame of preset E or F through NGX, any size, any +flags: its uploads hold the network's weights, which nvngx_dlss.dll keeps packed. The frame (NGX 310.7): + input kernel cuda_engine_input_kernel_rel_{ldr,hdr}_mvdiff_{mvhi,mvlo}: the game's textures (and the warped + history) to the network's 16 input channels per token (4 x 4 output pixels) + network eleven dltss_nchw8 kernels: d4r's own engine/sr/e_net.comp (packed binary16) + reconstruction dltss_nchw8_..._aniso_gaussian_dfn_{fastpath_3x3,mid_3x3,3x3_256,fastpath_5x5,5x5_256}: the last + decoder level, the 40-channel head and the reconstruction of the output and its history +The input and reconstruction kernels are translated from the DLL's PTX (engine/rr/ptx2rr.py) in every variant the +runtime selects from (engine/runtime_e.inc). The reconstruction head defaults to packed half FMAs for RDNA2, +like the rest of the network. --head-accumulation f32 retains the reference's per-MMA rounding; all translated +variants use aligned word accesses for their packed tensors. Writes: + e.txt the manifest + weights.bin the network's weights packed for e_net.comp, then the reconstruction kernels' four + head tensors as NGX uploads them (copied into the arena by the runtime) + .spv the network stages + input-.spv / .globals.bin, final-.spv / .globals.bin the translated kernels + input.args, final.args their argument blocks from the dump (the runtime rewrites every frame-dependent field) +Everything but d4r's shaders is derived from the user's DLL: it stays on their machine and is never packaged. +""" +import argparse +import glob +import json +import os +import re +import struct +import subprocess +import sys +import tempfile +from pathlib import Path + +import numpy as np + +HERE = Path(__file__).resolve().parent +NET_SRC = HERE / 'sr' / 'e_net.comp' + +# the network: name, the kernel's index in the frame's dltss launches, mode, defines (PS: pixel sets per lane, 2 on the +# large levels), the weights to pack +# (argument offset, layout, in, out) and the biases (argument offsets). Layouts: k3 = NGX's [K/8][tap][N][8], +# k1 = [K/8][N][8], oi = [N][K]; all packed to [tap][k][n] / [k][n] binary16. +NET = [ + ('enc0', 0, 'ENC0', dict(CIN=16, C=32), [(8, 'k3', 16, 32), (40, 'k1', 32, 64)], [88, 176]), + ('enc1', 1, 'ENC3', dict(CIN=64, C=128, IN8=1), [(8, 'k3', 64, 128)], [56]), + ('enc2', 2, 'ENC3', dict(CIN=128, C=128), [(8, 'k3', 128, 128)], [56]), + ('enc3', 3, 'ENC3', dict(CIN=128, C=256), [(8, 'k3', 128, 256)], [56]), + ('mid0', 4, 'CONV1', dict(CIN=256, C=512), [(8, 'k1', 256, 512)], [24]), + ('mid1', 5, 'CONV1', dict(CIN=512, C=256), [(8, 'k1', 512, 256)], [24]), + ('up4', 6, 'UP1', dict(CIN=256, C=256), [(16, 'k1', 256, 256)], [32]), + ('red4', 7, 'CONV1', dict(CIN=256, C=128), [(8, 'k1', 256, 128)], [24]), + ('dec3', 8, 'DEC2', dict(CIN=128, C=128, C2=128), [(16, 'k1', 128, 128), (240, 'oi', 128, 128)], [24, 248]), + ('dec2', 9, 'DEC2', dict(CIN=128, C=128, C2=64, OUT8=1), [(16, 'k1', 128, 128), (240, 'oi', 128, 64)], [24, 248]), + ('dec1', 10, 'DEC2', dict(CIN=64, C=64, C2=32, LOW8=1, SKIP8=1, OUT8=1), [(16, 'k1', 64, 64), (240, 'oi', 64, 32)], + [24, 248]), +] +# the reconstruction kernels' weight arguments (the last decoder level's two convs; NGX's layouts, kept as they are) +HEAD = [(16, 2048), (24, 64), (240, 2560), (248, 80)] +INPUT_VARIANTS = ['ldr_mvdiff_mvhi', 'ldr_mvdiff_mvlo', 'hdr_mvdiff_mvhi', 'hdr_mvdiff_mvlo'] +FINAL_VARIANTS = ['fastpath_3x3_128', 'mid_3x3_128', '3x3_256', 'fastpath_5x5_256', '5x5_256'] +INPUT_KERNEL = 'cuda_engine_input_kernel_rel_{}' +FINAL_KERNEL = 'dltss_nchw8_bilinear_upsample_conv_1x1_conv_1x1_aniso_gaussian_dfn_{}_032_032_048_032_e5m3_fp16_kernel' +# the translated kernels' argument contract: pointers (arena byte offsets), surfaces and their formats +INPUT_POINTERS, INPUT_SURFACES = [376], {328: 'rg16f', 368: 'rgba16f', 384: 'rgba16f'} +FINAL_POINTERS, FINAL_SURFACES = [0, 8, 16, 24, 240, 248, 256], {488: 'rgba16f', 496: 'rgba16f'} + + +class Launch: + def __init__(self, d): + self.dir = Path(d) + self.name = self.dir.name.split('-', 2)[2] + self.allocs, self.ptr = {}, {} + for line in (self.dir / 'manifest.txt').read_text().splitlines(): + p = line.split() + if p and p[0] == 'alloc': + self.allocs[int(p[1])] = (int(p[2], 16), int(p[3])) + elif p and p[0] == 'pointer': + self.ptr[int(p[1])] = (int(p[2]), int(p[3])) + self.args = (self.dir / 'args.bin').read_bytes() + + def arg(self, off, n): + a, d = self.ptr[off] + data = (self.dir / f'alloc-{a}.bin').read_bytes() + if d + n > len(data): + raise ValueError(f'{self.name}: argument {off} is shorter than {n} bytes') + return data[d:d + n] + + +def pack(raw, layout, cin, cout): + w = np.frombuffer(raw, np.float16) + if layout == 'k3': + return w.reshape(cin // 8, 9, cout, 8).transpose(1, 0, 3, 2).reshape(9 * cin, cout) + if layout == 'k1': + return w.reshape(cin // 8, cout, 8).transpose(0, 2, 1).reshape(cin, cout) + return w.reshape(cout, cin).T + + +def glsl(source, out, defines=()): + subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *defines, str(source), '-o', str(out)], + check=True, stdout=subprocess.DEVNULL) + + +def find_module(ptx_dir, kernel): + found = [p for p in sorted(Path(ptx_dir).glob('*.ptx')) if f'.entry {kernel}(' in p.read_text(errors='replace')] + for p in found: + if '.target sm_89' in p.read_text(errors='replace')[:400]: + return p + if found: + return found[-1] + raise FileNotFoundError(f'no PTX module with the entry {kernel}') + + +def translate(ptx_dir, kernel, out, name, block, pointers, surfaces, packed=False): + comp = out / f'{name}.generated.comp' + cmd = [sys.executable, str(HERE / 'rr' / 'ptx2rr.py'), str(find_module(ptx_dir, kernel)), kernel, str(comp), + '--block', ','.join(map(str, block)), '--weights-arg', '999', '--unified-arena'] + cmd += [f'--buffer={b}' for b in pointers] + [f'--surface={o}={f}' for o, f in surfaces.items()] + subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL) + glsl(comp, out / f'{name}.spv', ['-DD4R_WORD_ARENA', *(['-DD4R_MMA_PACKED'] if packed else [])]) + g = comp.with_suffix('.globals.bin') + (out / f'{name}.globals.bin').write_bytes(g.read_bytes() if g.exists() else bytes(16)) + info = json.loads(comp.with_suffix('.json').read_text()) + for f in (comp, comp.with_suffix('.json'), g): + if f.exists(): + f.unlink() + tex = ','.join(map(str, info['textures'])) or '-' + surf = ','.join(f'{o}:{f}' for o, f in info['surfaces']) or '-' + return f'{name}.spv {name}.globals.bin {tex} {surf}' + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument('replay', type=Path) + ap.add_argument('dll') + ap.add_argument('output', type=Path) + ap.add_argument('--preset', choices=['E', 'F'], required=True, help='the preset the dump ran (the two differ in weights ' + 'and in two of the kernels\' constants)') + ap.add_argument('--ptx-dir') + ap.add_argument('--head-accumulation', choices=['f16', 'f32'], default='f16', + help='f16: packed half FMAs for the reconstruction head (RDNA2); f32: reference matrix accumulation') + a = ap.parse_args() + launches = [Launch(d) for d in sorted(glob.glob(str(a.replay / 'replay-*')))] + first = next((i for i, l in enumerate(launches) if l.name.startswith('cuda_engine_input_kernel_rel')), None) + if first is None: + raise SystemExit(f'{a.replay}: no preset E/F frame (cuda_engine_input_kernel_rel_*)') + net = [l for l in launches[first + 1:] if l.name.startswith('dltss_nchw8_')][:12] + if len(net) != 12 or 'aniso_gaussian' not in net[11].name: + raise SystemExit(f'{a.replay}: the frame does not hold the CNN network and its reconstruction') + out = a.output + out.mkdir(parents=True, exist_ok=True) + blob = bytearray() + + def put(data): + at = len(blob) + blob.extend(data) + blob.extend(bytes(-len(blob) % 256)) + return at + + # the reconstruction's per-preset constant at 584 is k / (rx ry): k from the dump's own ratios + fin = net[11].args + ow0, oh0 = struct.unpack_from(' 1 else []) + for arg, n in zip(bs, cout_of): + offs.append(put(l.arg(arg, n * 2))) + # word offsets in the order the parameter block takes them: w1 b1 w2 b2 + wb = [offs[0] // 4, offs[len(ws)] // 4, (offs[1] // 4) if len(ws) > 1 else 0, (offs[len(ws) + 1] // 4) if len(bs) > 1 else 0] + glsl(NET_SRC, out / f'{name}.spv', [f'-DMODE_{mode}', *[f'-D{k}={v}' for k, v in defines.items()]]) + lines.append(f'stage {name} {name}.spv ' + ' '.join(map(str, wb)) + f' {defines.get("PS", 1)}') + head = [put(net[11].arg(arg, n)) for arg, n in HEAD] + lines.append('head ' + ' '.join(f'{arg}:{at}:{n}' for (arg, n), at in zip(HEAD, head))) + (out / 'weights.bin').write_bytes(bytes(blob)) + with tempfile.TemporaryDirectory() as tmp: + ptx = a.ptx_dir or tmp + if not a.ptx_dir: + subprocess.run([sys.executable, str(HERE.parent / 'kernels' / 'tools' / 'extract_dlss_ptx.py'), a.dll, tmp], + check=True, stdout=subprocess.DEVNULL) + for v in INPUT_VARIANTS: + lines.append(f'input {v} ' + translate(ptx, INPUT_KERNEL.format(v), out, f'input-{v}', (128, 1, 1), + INPUT_POINTERS, INPUT_SURFACES)) + for v in FINAL_VARIANTS: + lines.append(f'final {v} ' + translate(ptx, FINAL_KERNEL.format(v), out, f'final-{v}', (256, 1, 1), + FINAL_POINTERS, FINAL_SURFACES, packed=a.head_accumulation == 'f16')) + # the argument blocks, pointers and texture handles cleared (the runtime fills them) + for name, l, size in (('input', launches[first], 408), ('final', net[11], 656)): + block = bytearray(l.args[:size].ljust(size, b'\0')) + for off in l.ptr: + block[off:off + 8] = bytes(8) + (out / f'{name}.args').write_bytes(bytes(block)) + lines += ['template input input.args', 'template final final.args'] + (out / 'e.txt').write_text('\n'.join(lines) + '\n') + print(f'Super Resolution CNN model in {out} ({len(blob)} weight bytes)') + + +if __name__ == '__main__': + main() diff --git a/engine/game_e.cpp b/engine/game_e.cpp new file mode 100644 index 0000000..8b612c7 --- /dev/null +++ b/engine/game_e.cpp @@ -0,0 +1,55 @@ +#include "game_e.h" +#include "game_internal.h" + +namespace d4r { + +struct GameUpscalerE::Impl { + EModel& model; + EEngine engine; + uint32_t width, height, outputWidth, outputHeight, motionWidth, motionHeight; + game_detail::CanonicalInputs fallback; + bool started = false; + Impl(EModel& m, VkPhysicalDevice p, VkDevice d, uint32_t ow, uint32_t oh, uint32_t rw, uint32_t rh, bool hdr, bool lowRes) + : model(m), engine(m, rw, rh, ow, oh, hdr, lowRes), width(rw), height(rh), outputWidth(ow), outputHeight(oh), + motionWidth(lowRes ? rw : ow), motionHeight(lowRes ? rh : oh), fallback(p, d, rw, rh, motionWidth, motionHeight) {} +}; +GameUpscalerE::GameUpscalerE(EModel& model, VkPhysicalDevice physical, VkDevice device, uint32_t outputWidth, uint32_t outputHeight, + uint32_t renderWidth, uint32_t renderHeight, bool hdr, bool lowResMotion) + : impl(std::make_unique(model, physical, device, outputWidth, outputHeight, renderWidth, renderHeight, hdr, lowResMotion)) {} +GameUpscalerE::~GameUpscalerE() = default; + +void GameUpscalerE::record(VkCommandBuffer cb, const GameFrameE& frame) { + auto& g = *impl; + if (!gameColorFormat(frame.color.format) || !gameMotionFormat(frame.motion.format) || + !gameDepthFormat(frame.depth.format, frame.depth.aspect) || !gameOutputFormat(frame.output.format)) + throw std::runtime_error("unsupported game texture format"); + const bool directColor = frame.color.view != VK_NULL_HANDLE && !frame.color.x && !frame.color.y; + const bool directMotion = frame.motion.view != VK_NULL_HANDLE && !frame.motion.x && !frame.motion.y; + const bool directDepth = frame.depth.view != VK_NULL_HANDLE && !frame.depth.x && !frame.depth.y; + const bool directOutput = game_detail::directOutput(frame.output, frame.color, frame.motion, frame.depth, false) && + !frame.output.x && !frame.output.y; + if (!g.started) { + if (!g.model.uploadRecorded()) g.model.recordUpload(cb); + g.engine.recordInitialize(cb); + g.started = true; + } + g.fallback.record(cb, frame.color, frame.motion, frame.depth, g.width, g.height, g.motionWidth, g.motionHeight, + directColor, directMotion, directDepth); + EFrameSettings s; + s.jitter[0] = frame.jitter[0]; s.jitter[1] = frame.jitter[1]; + s.motionScale[0] = frame.motionScale[0]; s.motionScale[1] = frame.motionScale[1]; + s.reset = frame.reset; + FrameImages images{}; + images.color = directColor ? frame.color.view : g.fallback.color.view; + images.motion = directMotion ? frame.motion.view : g.fallback.motion.view; + images.depth = directDepth ? frame.depth.view : g.fallback.depth.view; + images.colorLayout = directColor ? frame.color.layout : VK_IMAGE_LAYOUT_GENERAL; + images.motionLayout = directMotion ? frame.motion.layout : VK_IMAGE_LAYOUT_GENERAL; + images.depthLayout = directDepth ? frame.depth.layout : VK_IMAGE_LAYOUT_GENERAL; + images.exposure = frame.exposure.view; + images.exposureLayout = frame.exposure.layout; + g.engine.setFrame(images, s, directOutput ? FrameOutput{frame.output.view, frame.output.format} : FrameOutput{}); + g.engine.recordFrame(cb); + if (!directOutput) game_detail::copyOutput(cb, g.engine.outputImage(), frame.output, g.outputWidth, g.outputHeight); +} +} // namespace d4r diff --git a/engine/game_e.h b/engine/game_e.h new file mode 100644 index 0000000..1ddd007 --- /dev/null +++ b/engine/game_e.h @@ -0,0 +1,28 @@ +// Adapter between a game's own textures and the preset E/F engine (DLSS 3's CNN, engine/runtime_e.inc): input +// conversion, the frame and the copy into the game's output texture, all recorded into the game's command buffer. +#pragma once +#include "game.h" + +namespace d4r { +struct GameFrameE { + GameImage color, depth, motion, output; // as for GameUpscaler; compatible views at origin zero may be borrowed + GameImage exposure; // HDR: the game's 1x1 float exposure (a sampled view), else 1 + float jitter[2] = {0, 0}; + float motionScale[2] = {1, 1}; // NGX's MV scale + bool reset = false; +}; +class GameUpscalerE { +public: + // The model must outlive this object. Render size fixed for its lifetime (a new size needs a new adapter). + // lowResMotion: motion at the render size (NGX's MVLowRes), otherwise at the output size. + GameUpscalerE(EModel& model, VkPhysicalDevice physical, VkDevice device, uint32_t outputWidth, uint32_t outputHeight, + uint32_t renderWidth, uint32_t renderHeight, bool hdr, bool lowResMotion); + ~GameUpscalerE(); // only once every recorded frame has completed on the GPU + GameUpscalerE(const GameUpscalerE&) = delete; + GameUpscalerE& operator=(const GameUpscalerE&) = delete; + void record(VkCommandBuffer command, const GameFrameE& frame); +private: + struct Impl; + std::unique_ptr impl; +}; +} // namespace d4r diff --git a/engine/replay_e.cpp b/engine/replay_e.cpp new file mode 100644 index 0000000..8282057 --- /dev/null +++ b/engine/replay_e.cpp @@ -0,0 +1,196 @@ +// Native Super Resolution preset E/F replay: runs EEngine (engine/runtime_e.inc) on a capture's game inputs, frame +// after frame, and writes the output (and with D4R_E_DUMP_TENSORS=1 the network's tensors) after each frame. +// +// usage: d4r-e MODEL_DIR INPUT_DIR OUT_DIR +// INPUT_DIR: frames.txt ("size RW RH OW OH HDR LOWRES MW MH", then per frame "frame F jx jy mvsx mvsy reset") and +// F-color.bin (RGBA16F, render size), F-motion.bin (RG16F, MW x MH), +// F-depth.bin (R32F, render size, with LOWRES). D4R_E_TIME_REPEAT=N: best of N GPU timings per frame. +#include "runtime.h" +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { +void check(VkResult r, const char* operation) { + if (r != VK_SUCCESS) throw std::runtime_error(std::string(operation) + ": VkResult " + std::to_string(r)); +} +#define CK(call) check((call), #call) +std::vector read(const std::string& file) { + std::ifstream stream(file, std::ios::binary | std::ios::ate); + if (!stream) throw std::runtime_error("cannot read " + file); + std::vector data(static_cast(stream.tellg())); + stream.seekg(0); + stream.read(data.data(), static_cast(data.size())); + return data; +} +struct Ctx { + VkPhysicalDevice physical; + VkDevice device; + VkQueue queue; + VkCommandPool pool; + VkPhysicalDeviceMemoryProperties memory; + uint32_t type(uint32_t bits, VkMemoryPropertyFlags flags) const { + for (uint32_t i = 0; i < memory.memoryTypeCount; ++i) + if ((bits & (1u << i)) && (memory.memoryTypes[i].propertyFlags & flags) == flags) return i; + throw std::runtime_error("memory type"); + } + template void submit(F&& record) { + VkCommandBufferAllocateInfo ai{VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO}; + ai.commandPool = pool; ai.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; ai.commandBufferCount = 1; + VkCommandBuffer cb; CK(vkAllocateCommandBuffers(device, &ai, &cb)); + VkCommandBufferBeginInfo bi{VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO}; + CK(vkBeginCommandBuffer(cb, &bi)); record(cb); CK(vkEndCommandBuffer(cb)); + VkSubmitInfo si{VK_STRUCTURE_TYPE_SUBMIT_INFO}; si.commandBufferCount = 1; si.pCommandBuffers = &cb; + CK(vkQueueSubmit(queue, 1, &si, VK_NULL_HANDLE)); CK(vkQueueWaitIdle(queue)); + vkFreeCommandBuffers(device, pool, 1, &cb); + } +}; +struct HostBuffer { + VkBuffer buffer; VkDeviceMemory memory; void* mapped; + HostBuffer(Ctx& c, VkDeviceSize size) { + VkBufferCreateInfo bi{VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO}; bi.size = size; + bi.usage = VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT; + CK(vkCreateBuffer(c.device, &bi, nullptr, &buffer)); + VkMemoryRequirements r; vkGetBufferMemoryRequirements(c.device, buffer, &r); + VkMemoryAllocateInfo ai{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO}; ai.allocationSize = r.size; + ai.memoryTypeIndex = c.type(r.memoryTypeBits, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT); + CK(vkAllocateMemory(c.device, &ai, nullptr, &memory)); CK(vkBindBufferMemory(c.device, buffer, memory, 0)); + CK(vkMapMemory(c.device, memory, 0, VK_WHOLE_SIZE, 0, &mapped)); + } +}; +void copyRegion(VkBufferImageCopy& r, uint32_t w, uint32_t h) { + r = {}; r.imageSubresource = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1}; r.imageExtent = {w, h, 1}; +} +struct Input { + VkImage image; VkImageView view; VkFormat format; + Input(Ctx& c, VkFormat f, uint32_t w, uint32_t h) : format(f) { + VkImageCreateInfo ii{VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO}; + ii.imageType = VK_IMAGE_TYPE_2D; ii.format = f; ii.extent = {w, h, 1}; ii.mipLevels = ii.arrayLayers = 1; + ii.samples = VK_SAMPLE_COUNT_1_BIT; ii.usage = VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT; + CK(vkCreateImage(c.device, &ii, nullptr, &image)); + VkMemoryRequirements r; vkGetImageMemoryRequirements(c.device, image, &r); + VkMemoryAllocateInfo ai{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO}; ai.allocationSize = r.size; + ai.memoryTypeIndex = c.type(r.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT); + VkDeviceMemory m; CK(vkAllocateMemory(c.device, &ai, nullptr, &m)); CK(vkBindImageMemory(c.device, image, m, 0)); + VkImageViewCreateInfo vi{VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO}; + vi.image = image; vi.viewType = VK_IMAGE_VIEW_TYPE_2D; vi.format = f; vi.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + CK(vkCreateImageView(c.device, &vi, nullptr, &view)); + } +}; +void barrier(VkCommandBuffer cb, VkImage image, VkImageLayout from, VkImageLayout to) { + VkImageMemoryBarrier b{VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER}; + b.image = image; b.oldLayout = from; b.newLayout = to; + b.srcAccessMask = VK_ACCESS_MEMORY_WRITE_BIT; b.dstAccessMask = VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT; + b.srcQueueFamilyIndex = b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; b.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + vkCmdPipelineBarrier(cb, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, 0, 0, nullptr, 0, nullptr, 1, &b); +} +} // namespace + +int main(int argc, char** argv) try { + if (argc != 4) { std::cerr << "usage: d4r-e MODEL_DIR INPUT_DIR OUT_DIR\n"; return 2; } + const std::string modelDir = argv[1], inputDir = argv[2], outDir = argv[3]; + VkApplicationInfo app{VK_STRUCTURE_TYPE_APPLICATION_INFO}; app.apiVersion = VK_API_VERSION_1_3; + VkInstanceCreateInfo ici{VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO}; ici.pApplicationInfo = &app; + VkInstance instance; CK(vkCreateInstance(&ici, nullptr, &instance)); + uint32_t count = 8; VkPhysicalDevice devices[8]; vkEnumeratePhysicalDevices(instance, &count, devices); + Ctx c{}; + c.physical = devices[0]; + for (uint32_t i = 0; i < count; ++i) { + VkPhysicalDeviceProperties p; vkGetPhysicalDeviceProperties(devices[i], &p); + if (p.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU) c.physical = devices[i]; + } + vkGetPhysicalDeviceMemoryProperties(c.physical, &c.memory); + VkPhysicalDeviceVulkan13Features v13{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_FEATURES}; + VkPhysicalDeviceVulkan12Features v12{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES, &v13}; + VkPhysicalDeviceVulkan11Features v11{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES, &v12}; + VkPhysicalDeviceFeatures2 f2{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2, &v11}; + vkGetPhysicalDeviceFeatures2(c.physical, &f2); + uint32_t families = 8; VkQueueFamilyProperties qf[8]; vkGetPhysicalDeviceQueueFamilyProperties(c.physical, &families, qf); + uint32_t family = 0; + for (uint32_t i = 0; i < families; ++i) if (qf[i].queueFlags & VK_QUEUE_COMPUTE_BIT) { family = i; break; } + float priority = 1; + VkDeviceQueueCreateInfo qi{VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO}; qi.queueFamilyIndex = family; qi.queueCount = 1; qi.pQueuePriorities = &priority; + VkDeviceCreateInfo di{VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO, &f2}; di.queueCreateInfoCount = 1; di.pQueueCreateInfos = &qi; + CK(vkCreateDevice(c.physical, &di, nullptr, &c.device)); + vkGetDeviceQueue(c.device, family, 0, &c.queue); + VkCommandPoolCreateInfo pi{VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO}; pi.queueFamilyIndex = family; + pi.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT; + CK(vkCreateCommandPool(c.device, &pi, nullptr, &c.pool)); + + auto framesText = read(inputDir + "/frames.txt"); + std::istringstream frames(std::string(framesText.begin(), framesText.end())); + std::string key; + uint32_t rw, rh, ow, oh, hdr, lowRes, mw, mh; + frames >> key >> rw >> rh >> ow >> oh >> hdr >> lowRes >> mw >> mh; + d4r::EModel model(c.physical, c.device, modelDir); + d4r::EEngine engine(model, rw, rh, ow, oh, hdr != 0, lowRes != 0); + Input color(c, VK_FORMAT_R16G16B16A16_SFLOAT, rw, rh), motion(c, VK_FORMAT_R16G16_SFLOAT, mw, mh), depth(c, VK_FORMAT_R32_SFLOAT, rw, rh); + HostBuffer staging(c, std::max(VkDeviceSize(ow) * oh * 8, 64u << 20)); + VkQueryPoolCreateInfo qp{VK_STRUCTURE_TYPE_QUERY_POOL_CREATE_INFO}; + qp.queryType = VK_QUERY_TYPE_TIMESTAMP; qp.queryCount = 4; + VkQueryPool stamps; CK(vkCreateQueryPool(c.device, &qp, nullptr, &stamps)); + VkPhysicalDeviceProperties props; vkGetPhysicalDeviceProperties(c.physical, &props); + const int repeat = std::getenv("D4R_E_TIME_REPEAT") ? std::atoi(std::getenv("D4R_E_TIME_REPEAT")) : 1; + bool first = true; + auto upload = [&](Input& in, const std::string& name, uint32_t f, uint32_t w, uint32_t h, uint32_t texel) { + auto data = read(inputDir + "/" + std::to_string(f) + "-" + name + ".bin"); + if (data.size() != size_t(w) * h * texel) throw std::runtime_error("input size: " + name); + std::memcpy(staging.mapped, data.data(), data.size()); + c.submit([&](VkCommandBuffer cb) { + barrier(cb, in.image, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL); + VkBufferImageCopy r; copyRegion(r, w, h); + vkCmdCopyBufferToImage(cb, staging.buffer, in.image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &r); + barrier(cb, in.image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_GENERAL); + }); + }; + while (frames >> key) { + uint32_t f; int reset; + d4r::EFrameSettings s; + frames >> f >> s.jitter[0] >> s.jitter[1] >> s.motionScale[0] >> s.motionScale[1] >> reset; + s.reset = reset != 0; + upload(color, "color", f, rw, rh, 8); + upload(motion, "motion", f, mw, mh, 4); + if (lowRes) upload(depth, "depth", f, rw, rh, 4); + d4r::FrameImages images{}; + images.color = color.view; images.motion = motion.view; images.depth = depth.view; + engine.setFrame(images, s); + if (first) c.submit([&](VkCommandBuffer cb) { model.recordUpload(cb); engine.recordInitialize(cb); first = false; }); + double best = 1e30, in = 0, net = 0, fin = 0; + for (int r = 0; r < repeat; ++r) { + if (r) engine.setFrame(images, s); + c.submit([&](VkCommandBuffer cb) { + vkCmdResetQueryPool(cb, stamps, 0, 4); + engine.recordFrame(cb, stamps, 0); + }); + uint64_t t[4]; + CK(vkGetQueryPoolResults(c.device, stamps, 0, 4, sizeof(t), t, 8, VK_QUERY_RESULT_64_BIT | VK_QUERY_RESULT_WAIT_BIT)); + const double ms = props.limits.timestampPeriod * 1e-6; + if ((t[3] - t[0]) * ms < best) { best = (t[3] - t[0]) * ms; in = (t[1] - t[0]) * ms; net = (t[2] - t[1]) * ms; fin = (t[3] - t[2]) * ms; } + } + std::cout << "frame " << f << ": " << best << " ms (input " << in << ", network " << net << ", reconstruction " << fin << ")\n"; + c.submit([&](VkCommandBuffer cb) { + VkBufferImageCopy r; copyRegion(r, ow, oh); + vkCmdCopyImageToBuffer(cb, engine.outputImage(), VK_IMAGE_LAYOUT_GENERAL, staging.buffer, 1, &r); + }); + std::ofstream(outDir + "/" + std::to_string(f) + "-out.bin", std::ios::binary) + .write(static_cast(staging.mapped), static_cast(size_t(ow) * oh * 8)); + if (std::getenv("D4R_E_DUMP_TENSORS")) + for (const auto& d : engine.debugBuffers()) { + c.submit([&](VkCommandBuffer cb) { + VkBufferCopy r{d.offset, 0, d.bytes}; + vkCmdCopyBuffer(cb, d.buffer, staging.buffer, 1, &r); + }); + std::ofstream(outDir + "/" + std::to_string(f) + "-" + d.name + ".bin", std::ios::binary) + .write(static_cast(staging.mapped), static_cast(d.bytes)); + } + } + return 0; +} catch (const std::exception& e) { + std::cerr << "d4r-e: " << e.what() << "\n"; + return 1; +} diff --git a/engine/rr/ptx2rr.py b/engine/rr/ptx2rr.py index 2788659..ee7891c 100755 --- a/engine/rr/ptx2rr.py +++ b/engine/rr/ptx2rr.py @@ -67,6 +67,14 @@ uint HU2_(u16vec2 u) { return uint(u.x) | (uint(u.y) << 16); } uint H2U(f16vec2 h) { return HU2_(float16BitsToUint16(h)); } int s16(uint u) { return (int(u) << 16) >> 16; } float rcp_a(float x) { return 1.0 / x; } +// add.rz.f32: the sum rounded toward zero (round to nearest, then one ulp back when that went away from zero; the +// error term is exact, TwoSum) +float add_rz(float a, float b) { precise float s = a + b; precise float bb = s - a; precise float e = (a - (s - bb)) + (b - bb); + return (e != 0.0 && ((e < 0.0) != (s < 0.0)) && s != 0.0) ? uintBitsToFloat(floatBitsToUint(s) - 1u) : s; } +float lg2_ftz(float x) { return log2(abs(x) < 1.17549435e-38 ? 0.0 : x); } +// PTX .relu: negative results become +0, NaN stays NaN +f16vec2 relu2(f16vec2 v) { return f16vec2(v.x < 0.0hf ? 0.0hf : v.x, v.y < 0.0hf ? 0.0hf : v.y); } +float16_t relu1(float16_t v) { return v < 0.0hf ? 0.0hf : v; } vec4 fetch0(sampler2D t, ivec2 c) { const ivec2 n = textureSize(t, 0); return (c.x < 0 || c.y < 0 || c.x >= n.x || c.y >= n.y) ? vec4(0.0) : texelFetch(t, c, 0); } // x + copysign(0.5, x) rounded toward zero and then truncated: round half away from zero, exactly float round_half_away(float x) { float t = trunc(x); return abs(x - t) >= 0.5 ? t + sign(x) : t; } @@ -101,17 +109,27 @@ uint GW16(uint a) { CHECKW(a); return uint(wbuf.wh[a >> 1]); } // a 32-bit access is 4-byte aligned (PTX ISA): one word of the same buffer's word view uint GW32(uint a) { CHECKW(a + 2u); return wbuf32.ww[a >> 2]; } uint GA16(uint a) { CHECKR(a); return uint(abuf.ah[a >> 1]); } -uint GA32(uint a) { CHECKR(a + 2u); return uint(abuf.ah[a >> 1]) | (uint(abuf.ah[(a >> 1) + 1u]) << 16); } +uint GA32(uint a) { CHECKR(a + 2u); +#ifdef D4R_WORD_ARENA + return abuf32.aw[a >> 2]; +#else + return uint(abuf.ah[a >> 1]) | (uint(abuf.ah[(a >> 1) + 1u]) << 16); +#endif +} void SA16(uint a, uint v) { CHECKS(a); abuf.ah[a >> 1] = uint16_t(v & 0xffffu); } -void SA32(uint a, uint v) { CHECKS(a + 2u); abuf.ah[a >> 1] = uint16_t(v & 0xffffu); abuf.ah[(a >> 1) + 1u] = uint16_t(v >> 16); } +void SA32(uint a, uint v) { CHECKS(a + 2u); +#ifdef D4R_WORD_ARENA + abuf32.aw[a >> 2] = v; +#else + abuf.ah[a >> 1] = uint16_t(v & 0xffffu); abuf.ah[(a >> 1) + 1u] = uint16_t(v >> 16); +#endif +} // Surface stores with .zero out-of-bounds semantics. -void sust_p4(writeonly image2D s, int x, int y, vec4 v) -{ - ivec2 n = imageSize(s); - 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)); } +// Expand stores at the call site: passing a formatted image to an unformatted image2D +// parameter makes glslang emit mismatched SPIR-V image types. +#define sust_p4(S, X, Y, V) { ivec2 sn_ = imageSize(S); ivec2 sc_ = ivec2(X, Y); if (uint(sc_.x) < uint(sn_.x) && uint(sc_.y) < uint(sn_.y)) imageStore(S, sc_, V); } +#define sust_p1(S, X, Y, 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. @@ -189,14 +207,8 @@ vec4 q_unorm1010102_hw(vec4 v) // was one binary16 ulp low in half its values, which the next frame's encoder 0 amplifies): OR-ing in a zero the // compiler cannot see (the dispatch's z count is 1) keeps the conversion's rounding. #define SUST_Z (gl_NumWorkGroups.z - 1u) -void sust_b2(writeonly image2D s, int x, int y, uint a, uint b) -{ - sust_p4(s, x >> 2, y, vec4(float(H(a | SUST_Z)), float(H(b | SUST_Z)), 0.0, 0.0)); -} -void sust_b4(writeonly image2D s, int x, int y, uint a, uint b, uint c, uint d) -{ - sust_p4(s, x >> 3, y, vec4(float(H(a | SUST_Z)), float(H(b | SUST_Z)), float(H(c | SUST_Z)), float(H(d | SUST_Z)))); -} +#define sust_b2(S, X, Y, A, B) sust_p4(S, (X) >> 2, Y, vec4(float(H((A) | SUST_Z)), float(H((B) | SUST_Z)), 0.0, 0.0)) +#define sust_b4(S, X, Y, A, B, C, D) sust_p4(S, (X) >> 3, Y, vec4(float(H((A) | SUST_Z)), float(H((B) | SUST_Z)), float(H((C) | SUST_Z)), float(H((D) | SUST_Z)))) // shfl.sync.up/down within a 32-lane warp (ZLUDA's SHFL_SYNC_IMPL). in_bounds is the predicate output. uint shfl_dir(uint v, uint delta, uint opts, bool down, out bool in_bounds) @@ -507,6 +519,8 @@ class RRT(base.T): 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 + self.written_buffers = set() + self.unified_arena = False self.snorm16 = set() # texture arguments of --snorm16 bx, by, bz = block if bx * by * bz % 32: @@ -535,7 +549,7 @@ class RRT(base.T): # 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): + for gm in re.finditer(r'\.(?:global|const)\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) @@ -629,6 +643,8 @@ class RRT(base.T): def gorigin(self, reg): s = self.origin.get(reg, set()) + if self.unified_arena and s and 'w' not in s: + return 'a' if len(s) != 1: raise base.Err(f'global pointer {reg} has origin {sorted(s)} (need exactly one of w/a)') return next(iter(s)) @@ -651,6 +667,16 @@ class RRT(base.T): 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 'const' in o: + m = re.fullmatch(r'\[\s*(\w+)\s*(?:\+\s*(\d+))?\s*\]', a[1].strip()) + if not m or m.group(1) not in self.globals: + raise base.Err(f'unsupported constant address {a[1]}') + size, count = mem_shape(o) + at = self.globals[m.group(1)] + int(m.group(2) or 0) + dsts = braces(a[0]) if count > 1 else [a[0]] + for i, d in enumerate(dsts): + self.A(d, f'GW{size * 8}({at + i * size}u)', 'u16' if size == 2 else 'u32') + return if 'global' not in o: return super().ld(o, a) size, count = mem_shape(o) @@ -661,12 +687,23 @@ class RRT(base.T): self.A(d, f'{fn}{size * 8}({self.gbyte(reg, off + i * size)})', 'u16' if size == 2 else 'u32') def st(self, o, a): + if 'shared' in o and o[-1] in ('u16', 'b16', 'f16') and not {'v2', 'v4'} & set(o): + reg, off = self.addr(a[0]) + value = self.val(a[1], 'u16') + self.E('{ uint at_ = ' + self.val(reg, 'u32') + f' + {off}u;') + self.E('uint shift_ = (at_ & 2u) * 8u;') + self.E('atomicAnd(sh[at_ >> 2], ~(0xffffu << shift_));') + self.E(f'atomicOr(sh[at_ >> 2], (uint({value}) & 0xffffu) << shift_); }}') + return if 'global' not in o: return super().st(o, a) size, count = mem_shape(o) reg, off = self.gaddr(a[0]) if self.gorigin(reg) == 'w': raise base.Err(f'store through the read-only weight pointer: {a}') + for origin in self.origin.get(reg, ()): + if origin.startswith('x'): + self.written_buffers.add(int(origin[1:])) fn = self.gfn(reg, True) vals = braces(a[1]) if count > 1 else [a[1]] for i, v in enumerate(vals): @@ -799,6 +836,20 @@ class RRT(base.T): return if op.startswith('mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16'): return self.mma(a) + if op in ('mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16', + 'mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16'): + d, x, b, c = (braces(t) for t in a) + depth = 8 if '.m16n8k8.' in op else 16 + if (len(d), len(x), len(b), len(c)) != (2, depth // 4, depth // 8, 2): + raise base.Err(f'invalid fragments for {op}') + operands = ', '.join(self.val(r, 'u32') for r in x + b + c) + self.E('{ uint md0_, md1_;') + self.E(f'mma_f16_m16n8k{depth}({operands}, md0_, md1_);') + self.A(d[0], 'md0_', 'u32') + self.A(d[1], 'md1_', 'u32') + return self.E('}') + if op.startswith('mma.sync.'): + raise base.Err(f'unsupported matrix operation {op}') if op == 'mma.pair.e4m3': d, h, x, b, c, e, f = (braces(t) for t in a) args = ', '.join(self.val(r, 'u32') for r in x + b + c + e + f) @@ -875,6 +926,51 @@ class RRT(base.T): c = ('(isnan({0}) || isnan({1}))' if parts[1] == 'nan' else '!(isnan({0}) || isnan({1}))' if parts[1] == 'num' else self.cmp(set(parts), True)) return self.A(a[0], f'(({c.format(self.val(a[1], "f16"), self.val(a[2], "f16"))}) ? 1.0hf : 0.0hf)', 'f16') + if op == 'popc.b32': + return self.A(a[0], f'uint(bitCount({self.val(a[1], "u32")}))', 'u32') + if op in ('dp2a.lo.u32.u32', 'dp2a.hi.u32.u32'): + x, y, c = (self.val(t, 'u32') for t in a[1:]) + shift = 0 if parts[1] == 'lo' else 16 + return self.A(a[0], f'({c} + ({x} & 0xffffu) * (({y} >> {shift}u) & 0xffu) + ' + f'({x} >> 16u) * (({y} >> {shift + 8}u) & 0xffu))', 'u32') + if op.startswith('ldmatrix.sync.aligned.m8n8.'): + regs = braces(a[0]) + count = 4 if 'x4' in parts else 2 if 'x2' in parts else 1 + if len(regs) != count or parts[-1] != 'b16': + raise base.Err(f'unsupported matrix load {op}') + src, off = self.addr(a[1]) + addr = f'({self.val(src, "u32")} + {off}u)' + self.E('{ uint lane_ = gl_SubgroupInvocationID;') + for i, dst in enumerate(regs): + if 'trans' in parts: + self.E(f'uint lo_{i} = subgroupShuffle({addr}, {8*i}u + 2u * (lane_ & 3u)) + (lane_ >> 2u) * 2u;') + self.E(f'uint hi_{i} = subgroupShuffle({addr}, {8*i + 1}u + 2u * (lane_ & 3u)) + (lane_ >> 2u) * 2u;') + self.A(dst, f'(SH16(lo_{i}) | (SH16(hi_{i}) << 16u))', 'u32') + else: + self.E(f'uint row_{i} = subgroupShuffle({addr}, (lane_ >> 2u) + {8*i}u);') + self.A(dst, f'SH32(row_{i} + (lane_ & 3u) * 4u)', 'u32') + return self.E('}') + if op.startswith('cp.async'): + if op == 'cp.async.commit_group' or op.startswith('cp.async.wait_group'): + return # Copies below execute synchronously in program order. + if not op.startswith(('cp.async.cg.shared.global', 'cp.async.ca.shared.global')): + raise base.Err(f'unsupported asynchronous copy {op}') + dst, doff = self.addr(a[0]) + src, soff = self.gaddr(a[1]) + size = int(a[2]) + if size not in (4, 8, 16): + raise base.Err(f'unsupported asynchronous copy size {size}') + valid = self.val(a[3], 'u32') if len(a) > 3 else f'{size}u' + fn = self.gfn(src) + self.E('{ uint dst_ = ' + self.val(dst, 'u32') + f' + {doff}u;') + self.E(f'uint bytes_ = {valid};') + for i in range(size // 4): + self.E(f'uint word_{i} = 0u;') + self.E(f'if (bytes_ > {4*i}u) {{') + self.E(f'uint n_ = min(bytes_ - {4*i}u, 4u);') + self.E(f'word_{i} = {fn}32({self.gbyte(src, soff + 4*i)}) & (0xffffffffu >> ((4u - n_) * 8u)); }}') + self.E(f'sh[(dst_ + {4*i}u) >> 2] = word_{i};') + return self.E('}') return super().op(op, a) # ------------------------------------------------ stage contract @@ -882,6 +978,7 @@ class RRT(base.T): w('#version 460') w('#extension GL_EXT_shader_explicit_arithmetic_types : require') w('#extension GL_EXT_shader_16bit_storage : require') + w('#extension GL_EXT_control_flow_attributes : 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') @@ -897,6 +994,9 @@ class RRT(base.T): w('layout(set = 0, binding = 1, std430) readonly buffer WeightBuf { uint16_t wh[]; } wbuf;') w('layout(set = 0, binding = 1, std430) readonly buffer WeightBuf32 { uint ww[]; } wbuf32;') w('layout(set = 0, binding = 2, std430) buffer ArenaBuf { uint16_t ah[]; } abuf;') + w('#ifdef D4R_WORD_ARENA') + w('layout(set = 0, binding = 2, std430) buffer ArenaBuf32 { uint aw[]; } abuf32;') + w('#endif') n = 3 for off in sorted(self.textures): w(f'layout(set = 0, binding = {n}) uniform sampler2D T{off};') @@ -906,7 +1006,7 @@ class RRT(base.T): 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: + for off in (() if self.unified_arena else 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); }}') @@ -921,7 +1021,8 @@ class RRT(base.T): w(f'shared uint mms[{384 * (self.bx * self.by * self.bz // 32)}]; // mma_e4m3_f16 scratch, per subgroup') def runtime(self): - return RR_RUNTIME + return (RR_RUNTIME + (ROOT / 'engine' / 'round_half.glsl').read_text() + + (ROOT / 'engine' / 'cnn_sr' / 'mma_f16.glsl').read_text()) def main(): @@ -940,6 +1041,8 @@ def main(): 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('--unified-arena', action='store_true', + help='place all --buffer pointer arguments in binding 2; retain their write metadata') 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') @@ -956,6 +1059,7 @@ def main(): def make(): t = RRT(src, ns.kernel, block, ns.weights_arg, ns.buffer) + t.unified_arena = ns.unified_arena t.snorm16 = set(ns.snorm16) t.probes = probes t.probe_block = tuple(int(x) for x in ns.probe_block.split(',')) @@ -971,7 +1075,9 @@ def main(): # 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)], 'buffers': t.buffers, 'shared': t.shared_bytes, + 'surfaces': [[off, t.surfaces[off]] for off in sorted(t.surfaces)], + 'buffers': [] if t.unified_arena else t.buffers, 'shared': t.shared_bytes, + 'written_buffers': sorted(t.written_buffers), '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/runtime.cpp b/engine/runtime.cpp index ad9f7ba..53ce3b5 100644 --- a/engine/runtime.cpp +++ b/engine/runtime.cpp @@ -36,9 +36,24 @@ struct Device { VkDevice handle; VkPhysicalDeviceMemoryProperties memory; VkPhysicalDeviceProperties properties; + bool coopmat = false; Device(VkPhysicalDevice physical, VkDevice device) : handle(device) { vkGetPhysicalDeviceMemoryProperties(physical, &memory); vkGetPhysicalDeviceProperties(physical, &properties); +#ifndef D4R_ENGINE_VK_TABLE + uint32_t count = 0; + VKCHECK(vkEnumerateDeviceExtensionProperties(physical, nullptr, &count, nullptr)); + std::vector extensions(count); + VKCHECK(vkEnumerateDeviceExtensionProperties(physical, nullptr, &count, extensions.data())); + coopmat = std::any_of(extensions.begin(), extensions.end(), [](const auto& extension) { + return std::strcmp(extension.extensionName, VK_KHR_COOPERATIVE_MATRIX_EXTENSION_NAME) == 0; + }); + if (const char* force = std::getenv("D4R_RR_TEST_NO_COOPMAT")) + if (std::strcmp(force, "1") == 0) coopmat = false; +#endif + // The shim borrows a device whose enabled feature chain is unavailable here. + // CNN stages use the extension-free flavor instead of assuming physical support + // means VK_KHR_cooperative_matrix was enabled by the application. } VkDeviceMemory allocate(VkMemoryRequirements requirements, VkMemoryPropertyFlags flags) const { for (uint32_t i = 0; i < memory.memoryTypeCount; ++i) { @@ -804,4 +819,5 @@ VkDeviceSize Engine::networkHeadBytes(uint32_t outputWidth, uint32_t outputHeigh VkDeviceSize Engine::activationBytes() const { return impl->activations; } #include "runtime_m.inc" #include "runtime_rr.inc" +#include "runtime_e.inc" } // namespace d4r diff --git a/engine/runtime.h b/engine/runtime.h index 4ab8ed5..852561c 100644 --- a/engine/runtime.h +++ b/engine/runtime.h @@ -307,4 +307,51 @@ private: struct Impl; std::unique_ptr impl; }; +// ---------------------------------------------------------------- Super Resolution presets E and F (the DLSS 3 CNN) +// NGX 310.7's CNN upscaler: its input kernel and its reconstruction kernel translated from the user's DLL, the +// network between them d4r's own packed-f16 shaders (engine/sr/e_net.comp), so it runs on any GPU with binary16 +// arithmetic, cooperative matrices not required (engine/runtime_e.inc, built by engine/compile_sr_e.py). Same +// conventions as the classes above. The device needs shaderFloat16, storageBuffer16BitAccess, shaderInt16, +// subgroupSizeControl and computeFullSubgroups with 32-wide subgroups. +struct EFrameSettings { + float jitter[2] = {0, 0}; // render pixels, as the game passes them to NGX + float motionScale[2] = {1, 1}; // NGX's MV scale: motion texels times this are the motion's pixels + bool reset = false; +}; +class EModel { +public: + EModel(VkPhysicalDevice physical, VkDevice device, const std::string& directory, VkPipelineCache cache = VK_NULL_HANDLE); + ~EModel(); + EModel(const EModel&) = delete; + EModel& operator=(const EModel&) = delete; + void recordUpload(VkCommandBuffer command); + bool uploadRecorded() const; +private: + struct Impl; + std::unique_ptr impl; + friend class EEngine; +}; +class EEngine { +public: + // hdr: NGX's IsHDR (linear colour, the exposure textures); lowResMotion: MVLowRes (render-size motion vectors, + // in render pixels; otherwise output-size motion in output pixels). Render size at most the output size. + EEngine(EModel& model, uint32_t renderWidth, uint32_t renderHeight, uint32_t outputWidth, uint32_t outputHeight, + bool hdr, bool lowResMotion); + ~EEngine(); + EEngine(const EEngine&) = delete; + EEngine& operator=(const EEngine&) = delete; + void recordInitialize(VkCommandBuffer command); + // color (and depth with lowResMotion) at the render size, motion at the render or output size; images.exposure + // the game's 1x1 exposure (HDR; 1 when absent). A borrowed RGBA16F finalOutput (GENERAL) is written directly. + void setFrame(const FrameImages& images, const EFrameSettings& settings, FrameOutput finalOutput = {}); + // With a timestamp query pool, four stamps are written from query `first`: start, after the input kernel, after + // the network and after the reconstruction. + void recordFrame(VkCommandBuffer command, VkQueryPool stamps = VK_NULL_HANDLE, uint32_t first = 0); + VkImage outputImage() const; // RGBA16F, GENERAL, when no finalOutput was bound + struct DebugBuffer { std::string name; VkBuffer buffer; VkDeviceSize offset, bytes; }; + std::vector debugBuffers() const; // the network's tensors in the arena (engine/replay_e.cpp) +private: + struct Impl; + std::unique_ptr impl; +}; } // namespace d4r diff --git a/engine/runtime_e.inc b/engine/runtime_e.inc new file mode 100644 index 0000000..0960dfa --- /dev/null +++ b/engine/runtime_e.inc @@ -0,0 +1,490 @@ +// Super Resolution presets E and F (the CNN of nvngx_dlss.dll 310.7): the runtime around engine/sr/e_net.comp and +// the translated input and reconstruction kernels of a "D4RE0002" model (engine/compile_sr_e.py). Included by +// runtime.cpp after runtime_rr.inc (its helpers: Tex, bindRange, put32 / putF / put64). +// +// The frame: the input kernel (the game's colour, depth and motion, the previous output history and motion to the +// network's 16 channels per token, the warped history and this frame's output-size motion), the eleven network +// stages, the reconstruction kernel (the last decoder level, the head, the output and the next history). One buffer +// (the arena) holds every tensor; the translated kernels take its byte offsets in place of NGX's pointers. +// +// The two translated kernels' argument blocks are NGX's (harness captures at seven sizes and ratios, 2026-10-18): +// the dump's block as a template, every field that depends on the sizes, the frame or the flags rewritten here. +namespace { +constexpr unsigned ENet = 11; +constexpr const char* ENames[ENet] = {"enc0", "enc1", "enc2", "enc3", "mid0", "mid1", "up4", "red4", "dec3", "dec2", "dec1"}; +// the network's tensors: name, channels, level (L0 = the token grid), 8-bit +struct ETensor { const char* name; uint32_t channels, level; bool e8; }; +constexpr ETensor ETensors[] = { + {"A", 16, 0, false}, {"B", 32, 0, true}, {"S1", 64, 1, true}, {"P1", 64, 2, true}, {"K2", 128, 2, false}, + {"P2", 128, 3, false}, {"K3", 128, 3, false}, {"P3", 128, 4, false}, {"K4", 256, 4, false}, {"P4", 256, 5, false}, + {"Q5", 512, 5, false}, {"R5", 256, 5, false}, {"D4", 256, 4, false}, {"E4", 128, 4, false}, {"D3", 128, 3, false}, + {"D2", 64, 2, true}, {"D1", 32, 1, true}}; +constexpr unsigned ETensorCount = sizeof(ETensors) / sizeof(ETensors[0]); +// per stage: the level it runs at and its tensors (input, skip or second output, output, third output; -1 none) +struct EStageIO { uint32_t level; int in, skip, out, out2; }; +constexpr EStageIO EStages[ENet] = { + {0, 0, 1, 2, 3}, {2, 3, 4, 5, -1}, {3, 5, 6, 7, -1}, {4, 7, 8, 9, -1}, {5, 9, -1, 10, -1}, {5, 10, -1, 11, -1}, + {4, 11, 8, 12, -1}, {4, 12, -1, 13, -1}, {3, 13, 6, 14, -1}, {2, 14, 4, 15, -1}, {1, 15, 2, 16, -1}}; + +struct ETranslated { + std::string shader, globalsFile; + std::vector textures; + std::vector> surfaces; + std::unique_ptr pipeline; + Buffer globals; + std::vector pending; +}; +} // namespace + +struct EModel::Impl { + Device device; + VkPipelineCache cache = VK_NULL_HANDLE; + std::string directory; + std::array stages; + std::array, ENet> weightAt{}; // word offsets in weights.bin: w1 b1 w2 b2 + std::array pixelSets{}; // pixel sets per lane (e_net.comp PS) + struct Head { uint32_t arg; VkDeviceSize at, bytes; }; + std::vector head; + Buffer weights; + std::vector weightsPending; + std::map> input, final; + std::vector inputTemplate, finalTemplate; + bool uploaded = false; + uint32_t preset = 5; // 5: E, 6: F + float kspread = 0.035f; // the reconstruction's constant at 584 times rx ry (E 0.035, F 0.1) + Impl(VkPhysicalDevice physical, VkDevice dev, const std::string& dir, VkPipelineCache c) + : device(physical, dev), cache(c), directory(dir) { + const auto text = read(dir + "/e.txt"); + std::istringstream lines(std::string(text.begin(), text.end())); + std::string line; + if (!std::getline(lines, line) || line.rfind("D4RE0002", 0) != 0) throw std::runtime_error("not a preset E/F model (e.txt D4RE0002)"); + const auto b = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER; + std::array seen{}; + while (std::getline(lines, line)) { + std::istringstream f(line); + std::string key; f >> key; + if (key == "preset") f >> preset; + else if (key == "kspread") f >> kspread; + else if (key == "weights") { + std::string file; f >> file; + weightsPending = read(dir + "/" + file); + weights.create(device, weightsPending.size(), VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT | + VK_BUFFER_USAGE_TRANSFER_SRC_BIT); + } else if (key == "stage") { + std::string name, shader; f >> name >> shader; + unsigned s = 0; + while (s < ENet && name != ENames[s]) ++s; + if (s >= ENet || seen[s]) throw std::runtime_error("e.txt: invalid stage: " + line); + for (auto& w : weightAt[s]) f >> w; + if (!(f >> pixelSets[s])) pixelSets[s] = 1; + if (pixelSets[s] != 1 && pixelSets[s] != 2) throw std::runtime_error("e.txt: invalid stage: " + line); + stages[s].create(device.handle, dir + "/" + shader, {b, b, b}, 32, cache); + seen[s] = true; + } else if (key == "head") { + for (std::string item; f >> item;) { + Head h{}; + if (std::sscanf(item.c_str(), "%u:%llu:%llu", &h.arg, reinterpret_cast(&h.at), + reinterpret_cast(&h.bytes)) != 3) + throw std::runtime_error("e.txt: invalid head record"); + head.push_back(h); + } + } else if (key == "input" || key == "final") { + // "input|final VARIANT SHADER.spv GLOBALS TEXTURES SURFACES" + std::string variant, textureList, surfaceList; + auto entry = std::make_unique(); + ETranslated& k = *entry; + f >> variant >> k.shader >> k.globalsFile >> textureList >> surfaceList; + std::istringstream tl(textureList == "-" ? "" : textureList), sl(surfaceList == "-" ? "" : surfaceList); + for (std::string item; std::getline(tl, item, ',');) k.textures.push_back(static_cast(std::stoul(item))); + for (std::string item; std::getline(sl, item, ',');) { + const auto colon = item.find(':'); + k.surfaces.emplace_back(static_cast(std::stoul(item.substr(0, colon))), item.substr(colon + 1)); + } + (key == "input" ? input : final)[variant] = std::move(entry); + } else if (key == "template") { + std::string which, file; f >> which >> file; + auto bytes = read(dir + "/" + file); + (which == "input" ? inputTemplate : finalTemplate).assign(bytes.begin(), bytes.end()); + } + } + for (unsigned s = 0; s < ENet; ++s) + if (!seen[s]) throw std::runtime_error("e.txt: incomplete preset E/F model"); + if (!weights.handle || head.size() != 4 || inputTemplate.size() != 408 || finalTemplate.size() != 656 || input.empty() || final.empty()) + throw std::runtime_error("e.txt: incomplete preset E/F model"); + } + // a translated kernel, its pipeline built on first use (the engine needs one input and one reconstruction variant) + ETranslated& kernel(bool isInput, const std::string& variant) { + auto& table = isInput ? input : final; + auto found = table.find(variant); + if (found == table.end()) throw std::runtime_error("e.txt: no " + std::string(isInput ? "input" : "reconstruction") + " kernel " + variant); + auto& k = *found->second; + if (!k.pipeline) { + const auto b = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, t = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, s = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE; + std::vector types{b, b, b}; + types.insert(types.end(), k.textures.size(), t); + types.insert(types.end(), k.surfaces.size(), s); + k.pipeline = std::make_unique(); + k.pipeline->create(device.handle, directory + "/" + k.shader, types, 32, cache); + k.pending = read(directory + "/" + k.globalsFile); + k.globals.create(device, k.pending.size(), VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT); + } + return k; + } +}; +EModel::EModel(VkPhysicalDevice physical, VkDevice device, const std::string& directory, VkPipelineCache cache) + : impl(std::make_unique(physical, device, directory, cache)) {} +EModel::~EModel() = default; +void EModel::recordUpload(VkCommandBuffer cb) { + auto& m = *impl; + if (m.uploaded) throw std::runtime_error("model upload already recorded"); + recordContents(cb, m.weights, m.weightsPending); + 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_TRANSFER_READ_BIT); + m.uploaded = true; +} +bool EModel::uploadRecorded() const { return impl->uploaded; } + +struct EEngine::Impl { + EModel::Impl& model; + const Device& d; + uint32_t rw, rh, ow, oh, pw, ph, h0, w0; // render, output, the output padded to 128, the token grid + bool hdr, lowRes; + ETranslated *input = nullptr, *final = nullptr; + std::string finalVariant; + uint32_t finalGrid[2] = {0, 0}; + Buffer arena, params; + std::array at{}, bytes{}; + std::array headAt{}; + VkDeviceSize scratch = 0, arenaBytes = 0; + std::array paramAt{}; + VkDeviceSize inAt = 0, finAt = 0, paramBytes = 0; + std::vector paramStaging, inBlock, finBlock; + std::array history, motionHistory; + Tex warped, discard, zero, exposureOne, out; + bool outNew = false; + VkSampler nearest = VK_NULL_HANDLE, linear = VK_NULL_HANDLE; + VkDescriptorPool fixedPool = VK_NULL_HANDLE; + std::array netSets{}; + struct Bindings { + FrameImages images{}; + FrameOutput finalOutput{}; + VkDescriptorPool pool = VK_NULL_HANDLE; + std::array input{}, final{}; + }; + std::vector bindings; + const Bindings* bound = nullptr; + VkQueryPool profile = VK_NULL_HANDLE; // D4R_ENGINE_E_PROFILE=1: microseconds per pass, printed a frame later + bool profiled = false; + float prevJitter[2] = {0, 0}; + uint32_t frame = 0; + bool reset = true; + + uint32_t lh(uint32_t l) const { return h0 >> l; } + uint32_t lw(uint32_t l) const { return w0 >> l; } + + Impl(EModel::Impl& m, uint32_t renderWidth, uint32_t renderHeight, uint32_t outputWidth, uint32_t outputHeight, bool isHdr, bool lowResMotion) + : model(m), d(m.device), rw(renderWidth), rh(renderHeight), ow(outputWidth), oh(outputHeight), hdr(isHdr), lowRes(lowResMotion) { + if (rw < 32 || rh < 32 || rw > ow || rh > oh || ow > 16384 || oh > 16384) + throw std::runtime_error("unsupported preset E/F dimensions"); + pw = (ow + 127) & ~127u; ph = (oh + 127) & ~127u; w0 = pw / 4; h0 = ph / 4; + // The mid kernel uses the actual scale/jitter and also covers the integer-rounded 2/3 render sizes. + // Requiring exact 3:2 dimensions sends e.g. 853x533 -> 1280x800 through the much larger generic kernel. + // Restrict this extension to floor/ceil rounding on each axis; retain the generic path at other ratios. + const auto qualityAxis = [](uint32_t render, uint32_t output) { + return render == (2 * output) / 3 || render == (2 * output + 2) / 3; + }; + input = &model.kernel(true, std::string(hdr ? "hdr" : "ldr") + "_mvdiff_" + (lowRes ? "mvlo" : "mvhi")); + const double sx = double(ow) / rw, sy = double(oh) / rh, s = std::min(sx, sy); + if (s > 2.5) finalVariant = "fastpath_5x5_256"; + else if (s >= 1.6) finalVariant = "fastpath_3x3_128"; + else if (qualityAxis(rw, ow) && qualityAxis(rh, oh)) finalVariant = "mid_3x3_128"; + else finalVariant = "3x3_256"; + if (const char* v = std::getenv("D4R_ENGINE_E_FINAL")) finalVariant = v; // development: force a variant + final = &model.kernel(false, finalVariant); + const bool generic = finalVariant == "3x3_256" || finalVariant == "5x5_256"; + finalGrid[0] = (ow + 63) / 64; finalGrid[1] = generic ? (oh + 63) / 64 : (oh + 31) / 32; + // the arena: the network's tensors, the head weights, a scratch word block (NGX's unused pointer 256) + VkDeviceSize a = 0; + auto region = [&](VkDeviceSize n) { const VkDeviceSize r = a; a += (n + 255) & ~VkDeviceSize(255); return r; }; + for (unsigned t = 0; t < ETensorCount; ++t) { + const auto& e = ETensors[t]; + bytes[t] = VkDeviceSize(e.channels) * lh(e.level) * lw(e.level) * (e.e8 ? 1 : 2); + at[t] = region(bytes[t]); + } + for (unsigned i = 0; i < 4; ++i) headAt[i] = region(model.head[i].bytes); + scratch = region(256); + arenaBytes = a; + const VkBufferUsageFlags storage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT; + arena.create(d, arenaBytes, storage); + auto place = [&](size_t n) { const VkDeviceSize r = paramBytes; paramBytes += (n + 255) & ~VkDeviceSize(255); return r; }; + for (unsigned s = 0; s < ENet; ++s) paramAt[s] = place(64); + inAt = place(model.inputTemplate.size()); finAt = place(model.finalTemplate.size()); + params.create(d, paramBytes, storage); + paramStaging.assign(static_cast(paramBytes), 0); + for (unsigned s = 0; s < ENet; ++s) { + const auto& io = EStages[s]; + uint32_t p[10] = {lh(io.level), lw(io.level), uint32_t(at[io.in]), io.skip >= 0 ? uint32_t(at[io.skip]) : 0u, + uint32_t(at[io.out]), io.out2 >= 0 ? uint32_t(at[io.out2]) : 0u, + model.weightAt[s][0], model.weightAt[s][1], model.weightAt[s][2], model.weightAt[s][3]}; + std::memcpy(paramStaging.data() + paramAt[s], p, sizeof(p)); + } + const VkFormat half4 = VK_FORMAT_R16G16B16A16_SFLOAT, half2 = VK_FORMAT_R16G16_SFLOAT; + for (int i = 0; i < 2; ++i) { history[i].create(d, ow, oh, half4); motionHistory[i].create(d, ow, oh, half2); } + warped.create(d, ow, oh, half4); + discard.create(d, 1, 1, half4); + zero.create(d, 1, 1, half4); + exposureOne.create(d, 1, 1, VK_FORMAT_R32_SFLOAT); + 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)); + VkDescriptorPoolSize sizes[] = {{VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 3 * ENet}}; + VkDescriptorPoolCreateInfo pi{VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO}; + pi.maxSets = ENet; pi.poolSizeCount = 1; pi.pPoolSizes = sizes; + VKCHECK(vkCreateDescriptorPool(d.handle, &pi, nullptr, &fixedPool)); + for (unsigned s = 0; s < ENet; ++s) { + netSets[s] = descriptorSet(d.handle, fixedPool, model.stages[s].descriptors); + bindRange(d.handle, netSets[s], 0, params, paramAt[s]); + bindBuffer(d.handle, netSets[s], 1, model.weights); + bindBuffer(d.handle, netSets[s], 2, arena); + } + } + ~Impl() { + if (profile) vkDestroyQueryPool(d.handle, profile, nullptr); + for (auto& b : bindings) vkDestroyDescriptorPool(d.handle, b.pool, nullptr); + if (fixedPool) vkDestroyDescriptorPool(d.handle, fixedPool, nullptr); + if (nearest) vkDestroySampler(d.handle, nearest, nullptr); + if (linear) vkDestroySampler(d.handle, linear, nullptr); + } + struct Resource { VkImageView view = VK_NULL_HANDLE; bool linear = false; VkImageLayout layout = VK_IMAGE_LAYOUT_GENERAL; }; + // what a translated kernel's texture or surface argument is (NGX's wiring in the captures; null where NGX leaves + // it unbound). par: this frame's parity (history written this frame, the other one read) + Resource resolve(bool isInput, bool surface, uint32_t offset, uint32_t par, const FrameImages& in, VkImageView output) const { + const VkImageView exposure = hdr ? (in.exposure ? in.exposure : exposureOne.view) : VK_NULL_HANDLE; + const VkImageLayout exposureLayout = in.exposure ? in.exposureLayout : VK_IMAGE_LAYOUT_GENERAL; + if (isInput) { + if (!surface) switch (offset) { + case 280: return {in.color, true, in.colorLayout}; + case 288: return {history[par ^ 1].view, true}; + case 304: return {in.motion, false, in.motionLayout}; + case 312: return {in.depth, false, in.depthLayout}; + case 320: return {motionHistory[par ^ 1].view, true}; + case 344: case 352: case 360: return {exposure, false, exposureLayout}; + } else switch (offset) { + case 328: return {motionHistory[par].view}; + case 368: return {discard.view}; + case 384: return {warped.view}; + } + } else { + if (!surface) switch (offset) { + case 456: return {in.color, true, in.colorLayout}; + case 464: return {warped.view, false}; + case 472: case 480: return {exposure, false, exposureLayout}; + } else switch (offset) { + case 488: return {output}; + case 496: return {history[par].view}; + } + } + return {}; + } + VkDescriptorSet translatedSet(Bindings& b, bool isInput, uint32_t par, VkImageView output) { + ETranslated& k = isInput ? *input : *final; + VkDescriptorSet set = descriptorSet(d.handle, b.pool, k.pipeline->descriptors); + bindRange(d.handle, set, 0, params, isInput ? inAt : finAt); + bindBuffer(d.handle, set, 1, k.globals); + bindBuffer(d.handle, set, 2, arena); + uint32_t binding = 3; + for (uint32_t offset : k.textures) { + const auto r = resolve(isInput, false, offset, par, b.images, output); + bindImage(d.handle, set, binding++, r.view ? r.view : zero.view, r.linear ? linear : nearest, r.view ? r.layout : VK_IMAGE_LAYOUT_GENERAL); + } + for (const auto& [offset, format] : k.surfaces) { + const auto r = resolve(isInput, true, offset, par, b.images, output); + bindImage(d.handle, set, binding++, r.view ? r.view : discard.view); + } + return set; + } + void setFrame(const FrameImages& images, const EFrameSettings& settings, FrameOutput finalOutput) { + if (!images.color || !images.motion || (lowRes && !images.depth)) + throw std::runtime_error("preset E/F needs colour and motion (and depth with render-size motion)"); + if (finalOutput.view && finalOutput.format != VK_FORMAT_R16G16B16A16_SFLOAT) + throw std::runtime_error("preset E/F writes RGBA16F output only"); + const bool rst = settings.reset || frame == 0; + const float frw = float(rw), frh = float(rh), fow = float(ow), foh = float(oh); + const float rx = frw / fow, ry = frh / foh; + const float pj[2] = {rst ? 0.0f : prevJitter[0], rst ? 0.0f : prevJitter[1]}; + auto& p = inBlock; + p = model.inputTemplate; + put32(p, 0, ow); put32(p, 4, oh); put32(p, 8, pw); put32(p, 12, ph); put32(p, 16, pw / 2); put32(p, 20, ph / 2); + put32(p, 24, ow); put32(p, 28, oh); putF(p, 32, 1.0f / fow); putF(p, 36, 1.0f / foh); + put32(p, 40, rw - 1); put32(p, 44, rh - 1); + put32(p, 48, (lowRes ? rw : ow) - 1); put32(p, 52, (lowRes ? rh : oh) - 1); + putF(p, 56, settings.motionScale[0] * (lowRes ? fow / frw : 1.0f)); putF(p, 60, settings.motionScale[1] * (lowRes ? foh / frh : 1.0f)); + putF(p, 72, lowRes ? rx : 1.0f); putF(p, 76, lowRes ? ry : 1.0f); + putF(p, 80, settings.jitter[0]); putF(p, 84, settings.jitter[1]); + put32(p, 108, hdr ? 1u : 0u); put32(p, 112, rst ? 1u : 0u); + putF(p, 120, rx); putF(p, 124, ry); putF(p, 128, rx); putF(p, 132, ry); + put32(p, 136, rw); put32(p, 140, rh); putF(p, 144, rx); putF(p, 148, ry); + putF(p, 152, 1.0f / fow); putF(p, 156, 1.0f / foh); putF(p, 160, model.preset == 5 ? 500.0f / foh : 1.0f); + put32(p, 184, lowRes ? 1u : 0u); put32(p, 188, lowRes ? 1u : 0u); + put64(p, 376, at[0]); + auto& o = finBlock; + o = model.finalTemplate; + put32(o, 40, lh(1)); put32(o, 44, lw(1)); put32(o, 60, h0); put32(o, 64, w0); + put32(o, 88, lw(1) * 32); put32(o, 92, w0 * 32); put32(o, 96, w0 * 32); + put32(o, 264, h0); put32(o, 268, w0); put32(o, 284, h0); put32(o, 288, w0); + put32(o, 312, w0 * 32); put32(o, 316, w0 * 32); put32(o, 320, w0 * 40); + put32(o, 504, oh); put32(o, 508, ow); put32(o, 512, hdr ? 1u : 0u); + putF(o, 528, ry); putF(o, 532, 1.0f / ry); putF(o, 536, rx); putF(o, 540, 1.0f / rx); + putF(o, 544, -pj[0]); putF(o, 548, -pj[1]); putF(o, 552, -settings.jitter[0]); putF(o, 556, -settings.jitter[1]); + put32(o, 564, rh - 1); put32(o, 572, rw - 1); + putF(o, 584, model.kspread / (rx * ry)); + put64(o, 0, at[16]); put64(o, 8, at[1]); + for (unsigned i = 0; i < 4; ++i) put64(o, model.head[i].arg, headAt[i]); + put64(o, 256, scratch); + if (!finalOutput.view && !out.handle) { + out.create(d, ow, oh, VK_FORMAT_R16G16B16A16_SFLOAT); + outNew = true; + } + const VkImageView output = finalOutput.view ? finalOutput.view : out.view; + // texture and surface handles: 1 where something is bound, 0 where NGX leaves the argument null + for (uint32_t offset : input->textures) put64(p, offset, resolve(true, false, offset, 0, images, output).view ? 1u : 0u); + for (const auto& [offset, format] : input->surfaces) put64(p, offset, offset == 368 ? 0u : 1u); // NGX leaves 368 null + for (uint32_t offset : final->textures) put64(o, offset, resolve(false, false, offset, 0, images, output).view ? 1u : 0u); + for (const auto& [offset, format] : final->surfaces) put64(o, offset, 1u); + 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 preset E/F input view sets"); + Bindings b; + b.images = images; b.finalOutput = finalOutput; + const uint32_t tex = 2 * uint32_t(input->textures.size() + final->textures.size()); + const uint32_t surf = 2 * uint32_t(input->surfaces.size() + final->surfaces.size()); + VkDescriptorPoolSize sizes[] = {{VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 12}, {VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, tex}, + {VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, surf}}; + VkDescriptorPoolCreateInfo pi{VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO}; + pi.maxSets = 4; pi.poolSizeCount = 3; pi.pPoolSizes = sizes; + VKCHECK(vkCreateDescriptorPool(d.handle, &pi, nullptr, &b.pool)); + for (uint32_t par = 0; par < 2; ++par) { + b.input[par] = translatedSet(b, true, par, output); + b.final[par] = translatedSet(b, false, par, output); + } + bindings.push_back(std::move(b)); + found = bindings.end() - 1; + } + bound = &*found; + if (const char* dump = std::getenv("D4R_ENGINE_E_DUMP")) { // development: the frame's two argument blocks + std::ofstream(std::string(dump) + "/" + std::to_string(frame + 1) + "-input.args", std::ios::binary) + .write(reinterpret_cast(inBlock.data()), static_cast(inBlock.size())); + std::ofstream(std::string(dump) + "/" + std::to_string(frame + 1) + "-final.args", std::ios::binary) + .write(reinterpret_cast(finBlock.data()), static_cast(finBlock.size())); + } + reset = rst; + prevJitter[0] = settings.jitter[0]; prevJitter[1] = settings.jitter[1]; + } + void recordInitialize(VkCommandBuffer cb) { + VkClearColorValue clear{}; + VkImageSubresourceRange range{VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + std::vector images{warped.handle, discard.handle, zero.handle, exposureOne.handle}; + for (int i = 0; i < 2; ++i) { images.push_back(history[i].handle); images.push_back(motionHistory[i].handle); } + 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); + } + VkClearColorValue one{}; one.float32[0] = 1.0f; + vkCmdClearColorImage(cb, exposureOne.handle, VK_IMAGE_LAYOUT_GENERAL, &one, 1, &range); + vkCmdFillBuffer(cb, arena.handle, 0, VK_WHOLE_SIZE, 0); + memoryBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, + VK_ACCESS_TRANSFER_WRITE_BIT); + // the head weights into the arena (the reconstruction kernel reads them there) + for (unsigned i = 0; i < 4; ++i) { + VkBufferCopy region{model.head[i].at, headAt[i], model.head[i].bytes}; + vkCmdCopyBuffer(cb, model.weights.handle, arena.handle, 1, ®ion); + } + for (const auto* k : {input, final}) + if (!k->pending.empty()) recordContents(cb, k->globals, const_cast(k)->pending); + 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; + } + void recordFrame(VkCommandBuffer cb, VkQueryPool stamps, uint32_t first) { + if (!bound) throw std::runtime_error("EEngine::setFrame was not called"); + const auto& b = *bound; + const uint32_t cur = frame & 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 (outNew) { + clearNewImage(cb, out.handle); + outNew = false; + } + constexpr uint32_t passes = ENet + 2; + if (!profile && std::getenv("D4R_ENGINE_E_PROFILE")) { + VkQueryPoolCreateInfo qi{VK_STRUCTURE_TYPE_QUERY_POOL_CREATE_INFO}; + qi.queryType = VK_QUERY_TYPE_TIMESTAMP; qi.queryCount = passes + 1; + VKCHECK(vkCreateQueryPool(d.handle, &qi, nullptr, &profile)); + } + if (profile && profiled) { + uint64_t t[passes + 1]; + if (vkGetQueryPoolResults(d.handle, profile, 0, passes + 1, sizeof(t), t, 8, VK_QUERY_RESULT_64_BIT) == VK_SUCCESS) { + double total = 0; + std::fprintf(stderr, "sr cnn profile (us):"); + for (uint32_t i = 0; i < passes; ++i) { + const double us = (t[i + 1] - t[i]) * d.properties.limits.timestampPeriod * 1e-3; + total += us; + std::fprintf(stderr, " %s %.1f", i == 0 ? "input" : i == passes - 1 ? "final" : ENames[i - 1], us); + } + std::fprintf(stderr, " total %.1f\n", total); + } + } + uint32_t mark = 0; + auto pass = [&] { if (profile) vkCmdWriteTimestamp(cb, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, profile, mark++); }; + if (profile) vkCmdResetQueryPool(cb, profile, 0, passes + 1); + 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::memcpy(paramStaging.data() + inAt, inBlock.data(), inBlock.size()); + std::memcpy(paramStaging.data() + finAt, finBlock.data(), finBlock.size()); + for (VkDeviceSize a = 0; a < paramBytes; a += 65536) + vkCmdUpdateBuffer(cb, params.handle, a, std::min(65536, paramBytes - a), paramStaging.data() + a); + 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); + pass(); + dispatch(cb, *input->pipeline, b.input[cur], (ow + 15) / 16, (oh + 15) / 16); + step(); stamp(1); pass(); + for (unsigned s = 0; s < ENet; ++s) { + const uint32_t l = EStages[s].level; + const uint32_t ps = model.pixelSets[s]; + uint32_t gx, gy = 1; + if (s == 0) { gx = w0 / 16; gy = h0 / (8 * ps); } + else if (s <= 3) { gx = (lw(l) + 7) / 8; gy = (lh(l) + 4 * ps - 1) / (4 * ps); } + else gx = (lw(l) * lh(l) + 32 * ps - 1) / (32 * ps); + dispatch(cb, model.stages[s], netSets[s], gx, gy); + step(); pass(); + } + stamp(2); + dispatch(cb, *final->pipeline, b.final[cur], finalGrid[0], finalGrid[1]); + 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); pass(); + profiled = profile != VK_NULL_HANDLE; + ++frame; + } +}; +EEngine::EEngine(EModel& model, uint32_t renderWidth, uint32_t renderHeight, uint32_t outputWidth, uint32_t outputHeight, bool hdr, bool lowResMotion) + : impl(std::make_unique(*model.impl, renderWidth, renderHeight, outputWidth, outputHeight, hdr, lowResMotion)) { + if (!model.impl->uploaded) {} // the caller records the model's upload before the first frame +} +EEngine::~EEngine() = default; +void EEngine::recordInitialize(VkCommandBuffer cb) { impl->recordInitialize(cb); } +void EEngine::setFrame(const FrameImages& images, const EFrameSettings& settings, FrameOutput finalOutput) { impl->setFrame(images, settings, finalOutput); } +void EEngine::recordFrame(VkCommandBuffer cb, VkQueryPool stamps, uint32_t first) { impl->recordFrame(cb, stamps, first); } +VkImage EEngine::outputImage() const { return impl->out.handle; } +std::vector EEngine::debugBuffers() const { + std::vector list; + for (unsigned t = 0; t < ETensorCount; ++t) list.push_back({ETensors[t].name, impl->arena.handle, impl->at[t], impl->bytes[t]}); + return list; +} diff --git a/engine/sr/e_net.comp b/engine/sr/e_net.comp new file mode 100644 index 0000000..53a1f83 --- /dev/null +++ b/engine/sr/e_net.comp @@ -0,0 +1,408 @@ +#version 460 +// The network of DLSS Super Resolution presets E and F (nvngx_dlss.dll 310.7: the "dltss_nchw8_*" kernels between +// NGX's input kernel and its reconstruction kernel) as d4r's own packed-f16 shaders, for any GPU with binary16 +// arithmetic and no cooperative matrices required (RADV on RDNA2 included). +// +// The network is a U-Net over the token grid (L0: the output size rounded up to 128, divided by 4; level l is L0 / 2^l): +// enc0 conv 3x3 16 -> 32, ReLU -> skip0 (8-bit, L0); 2x2 max-pool, conv 1x1 32 -> 64, ReLU -> skip1 (8-bit, L1); +// its 2x2 max-pool -> pool1 (8-bit, L2) +// enc1 conv 3x3 64 -> 128 (8-bit in) -> skip2 (L2), pool -> pool2 (L3) enc2 128 -> 128 at L3, enc3 128 -> 256 +// at L4 (both binary16): MODE_ENC3 +// mid0 conv 1x1 256 -> 512, mid1 512 -> 256 at L5: MODE_CONV1 +// up4 bilinear 2x of mid1 + skip4, conv 1x1 256 -> 256 (MODE_UP1); red4 conv 1x1 256 -> 128 (MODE_CONV1) +// dec3, dec2, dec1 bilinear 2x of the level below + the skip, conv 1x1 C -> M, ReLU, conv 1x1 M -> N, ReLU +// (MODE_DEC2); dec1's output and skip0 feed NGX's reconstruction kernel (translated, engine/compile_e.py), which +// runs the last decoder level and the head itself. +// Every conv is followed by its bias and a ReLU. Tensors are NVIDIA's: [C/8][H][W][8] ("nchw8") for the pooled and +// decoder outputs, and the full-size skips as the four 2x2 phases [y&1][x&1][C/8][H/2][W/2][8] ("sub"); 8-bit +// activations are binary16 bits 14..7 (round half up, as cnn_common.glsl's enc8). +// +// Arithmetic: a lane owns a pixel, a wave 32 pixels (a 4 x 8 block in Morton order, so 2x2 pools are quads of +// lanes) and 32 output channels; the weights are wave-uniform scalar loads feeding v_pk_fma_f16 (as +// engine/rr/cnn_pk.glsl). Binary16 accumulation: not bit-exact to NGX's f16 mma (about 45-80 dB per tensor). +// +// Bindings: 0 the parameter block (below), 1 the weights (binary16, [tap][k][n] / [k][n], n fastest), 2 the arena +// (the frame's tensors; every offset in bytes). +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_control_flow_attributes : require +#extension GL_KHR_shader_subgroup_shuffle : require +#extension GL_KHR_shader_subgroup_basic : require + +layout(std430, binding = 0) readonly buffer Params { uint p[]; }; +layout(std430, binding = 1) readonly restrict buffer Weights { uint w[]; }; // binary16 pairs +layout(std430, binding = 1) readonly restrict buffer Weights4 { uvec4 w4[]; }; +layout(std430, binding = 2) buffer Arena { uint a[]; }; +layout(std430, binding = 2) buffer Arena4 { uvec4 a4[]; }; +layout(std430, binding = 2) buffer Arena2 { uvec2 a2[]; }; + +// the parameter block (words): 0 H, 1 W (this stage's level), 2 input, 3 skip (MODE_UP1 / MODE_DEC2) or second +// output, 4 output, 5 third output, 6.. weights and biases (word offsets in binding 1): 6 w1, 7 b1, 8 w2, 9 b2 +#define PH p[0] +#define PW p[1] + +#ifndef CIN +#define CIN 32 +#endif +#ifndef C +#define C 32 // the stage's (first) conv outputs +#endif +#ifndef C2 +#define C2 32 // MODE_ENC0 / MODE_DEC2: the second conv's outputs +#endif +#define NQ 32 // a wave's output channels +#ifndef LOW8 +#define LOW8 0 +#endif +#ifndef SKIP8 +#define SKIP8 0 +#endif +#ifndef IN8 +#define IN8 0 +#endif +#ifndef OUT8 +#define OUT8 0 +#endif +#if defined(MODE_DEC2) +#define WAVES ((C > C2 ? C : C2) / NQ) +#elif defined(MODE_ENC0) +#define WAVES 4 // a 8 x 16 tile of L0; its pooled 4 x 8 block of L1 is the second conv's +#else +#define WAVES (C / NQ) +#endif +layout(local_size_x = 32 * WAVES) in; + +// gl_SubgroupID, not the invocation index / 32: the compiler then knows a wave's index (and every weight address +// built from it) is uniform, and loads the weights with scalar loads +uint wave() { return gl_SubgroupID; } +uint lane() { return gl_LocalInvocationIndex & 31u; } +// a lane's pixel in its wave's 4 x 8 block: x = b0 + 2 b2 + 4 b4, y = b1 + 2 b3 +uvec2 morton(uint l) { return uvec2((l & 1u) | ((l >> 1u) & 2u) | ((l >> 2u) & 4u), ((l >> 1u) & 1u) | ((l >> 2u) & 2u)); } + +f16vec2 h2(uint u) { return unpackFloat2x16(u); } +uint u2(f16vec2 v) { return packFloat2x16(v); } +// 8-bit activations <-> binary16 +f16vec4 dec8x4(uint u) +{ + const uint lo = ((u << 7u) & 0x7f80u) | ((u << 15u) & 0x7f800000u); + const uint hi = ((u >> 9u) & 0x7f80u) | ((u >> 1u) & 0x7f800000u); + return f16vec4(h2(lo), h2(hi)); +} +uint enc8x4(f16vec2 x, f16vec2 y) +{ + const f16vec2 ex = h2(u2(x) & 0xfc00fc00u), ey = h2(u2(y) & 0xfc00fc00u); + const uint rx = u2(fma(ex, f16vec2(0.0625), x)), ry = u2(fma(ey, f16vec2(0.0625), y)); + return ((rx >> 7u) & 0xffu) | ((rx >> 15u) & 0xff00u) | ((ry << 9u) & 0xff0000u) | ((ry << 1u) & 0xff000000u); +} + +// eight channels (one nchw8 group g) of a pixel of a [CH/8][H][W][8] tensor at byte offset `at` +void load8(uint at, uint g, uint h, uint wd, uint y, uint x, bool in8, out f16vec2 v[4]) +{ + const uint e = (g * h + y) * wd + x; + if (in8) { + const uvec2 b = a2[at / 8u + e]; + const f16vec4 lo = dec8x4(b.x), hi = dec8x4(b.y); + v[0] = lo.xy; v[1] = lo.zw; v[2] = hi.xy; v[3] = hi.zw; + } else { + const uvec4 b = a4[at / 16u + e]; + v[0] = h2(b.x); v[1] = h2(b.y); v[2] = h2(b.z); v[3] = h2(b.w); + } +} +// the same from a 2x2-phase ("sub") tensor of a level of size h x wd +void loadSub(uint at, uint g, uint ch, uint h, uint wd, uint y, uint x, bool in8, out f16vec2 v[4]) +{ + const uint ph = (y & 1u) * 2u + (x & 1u); + load8(at + ph * (ch * (h / 2u) * (wd / 2u)) * (in8 ? 1u : 2u), g, h / 2u, wd / 2u, y >> 1u, x >> 1u, in8, v); +} +void store8(uint at, uint g, uint h, uint wd, uint y, uint x, bool out8, f16vec2 v[4]) +{ + const uint e = (g * h + y) * wd + x; + if (out8) a2[at / 8u + e] = uvec2(enc8x4(v[0], v[1]), enc8x4(v[2], v[3])); + else a4[at / 16u + e] = uvec4(u2(v[0]), u2(v[1]), u2(v[2]), u2(v[3])); +} +void storeSub(uint at, uint g, uint ch, uint h, uint wd, uint y, uint x, bool out8, f16vec2 v[4]) +{ + const uint ph = (y & 1u) * 2u + (x & 1u); + store8(at + ph * (ch * (h / 2u) * (wd / 2u)) * (out8 ? 1u : 2u), g, h / 2u, wd / 2u, y >> 1u, x >> 1u, out8, v); +} + +#ifndef PS +#define PS 1 // pixel sets per lane: 2 reuses every scalar-loaded weight for two pixels +#endif +// acc[i][j] (pixel set i, output pair j of the wave's 32) += the eight channels v[i] (k = k0 .. k0 + 7) times the +// weight rows k0 .. k0 + 7 of a [k][N] matrix at word offset wrow (N / 2 words per row), the wave's columns from +// word column wcol +void fma8(inout f16vec2 acc[PS][16], f16vec2 v[PS][4], uint wrow, uint rowWords, uint wcol) +{ + [[unroll]] for (uint kk = 0u; kk < 8u; ++kk) { + const uint r = wrow + kk * rowWords + wcol; + [[unroll]] for (uint j4 = 0u; j4 < 4u; ++j4) { + const uvec4 q = w4[r / 4u + j4]; + [[unroll]] for (uint i = 0u; i < uint(PS); ++i) { + const float16_t x = (kk & 1u) == 0u ? v[i][kk >> 1u].x : v[i][kk >> 1u].y; + acc[i][4u * j4] = fma(f16vec2(x), h2(q.x), acc[i][4u * j4]); + acc[i][4u * j4 + 1u] = fma(f16vec2(x), h2(q.y), acc[i][4u * j4 + 1u]); + acc[i][4u * j4 + 2u] = fma(f16vec2(x), h2(q.z), acc[i][4u * j4 + 2u]); + acc[i][4u * j4 + 3u] = fma(f16vec2(x), h2(q.w), acc[i][4u * j4 + 3u]); + } + } + } +} +void biasRelu(inout f16vec2 acc[PS][16], uint b, uint wcol) +{ + [[unroll]] for (uint j4 = 0u; j4 < 4u; ++j4) { + const uvec4 q = w4[(b + wcol) / 4u + j4]; + [[unroll]] for (uint i = 0u; i < uint(PS); ++i) { + acc[i][4u * j4] = max(acc[i][4u * j4] + h2(q.x), f16vec2(0.0)); + acc[i][4u * j4 + 1u] = max(acc[i][4u * j4 + 1u] + h2(q.y), f16vec2(0.0)); + acc[i][4u * j4 + 2u] = max(acc[i][4u * j4 + 2u] + h2(q.z), f16vec2(0.0)); + acc[i][4u * j4 + 3u] = max(acc[i][4u * j4 + 3u] + h2(q.w), f16vec2(0.0)); + } + } +} +void zero(out f16vec2 acc[PS][16]) { [[unroll]] for (uint i = 0u; i < uint(PS); ++i) [[unroll]] for (uint j = 0u; j < 16u; ++j) acc[i][j] = f16vec2(0.0); } +// the 2x2 max over a lane quad (xor 1, 2: Morton order) +void pool4(inout f16vec2 acc[16]) +{ + [[unroll]] for (uint j = 0u; j < 16u; ++j) { + acc[j] = max(acc[j], h2(subgroupShuffleXor(u2(acc[j]), 1u))); + acc[j] = max(acc[j], h2(subgroupShuffleXor(u2(acc[j]), 2u))); + } +} +// the wave's 32 channels (groups gq .. gq + 3) of a pixel +void storeAll(uint at, uint gq, uint h, uint wd, uint y, uint x, bool out8, bool sub, uint ch, f16vec2 acc[16]) +{ + [[unroll]] for (uint g = 0u; g < 4u; ++g) { + f16vec2 v[4] = f16vec2[4](acc[4u * g], acc[4u * g + 1u], acc[4u * g + 2u], acc[4u * g + 3u]); + if (sub) storeSub(at, gq + g, ch, h, wd, y, x, out8, v); else store8(at, gq + g, h, wd, y, x, out8, v); + } +} + +// the input of MODE_UP1 / MODE_DEC2 at (y, x), eight channels of group g: the level below (h/2 x wd/2) upsampled +// 2x bilinearly (half-pixel centres, edges clamped: 0.75 / 0.25 per axis) plus the skip +void upSkip(uint g, uint y, uint x, out f16vec2 v[4]) +{ + const uint h = PH, wd = PW, lh = h / 2u, lw = wd / 2u; + const uint iy = y >> 1u, ix = x >> 1u; + const uint ny = uint(clamp(int(iy) + ((y & 1u) == 0u ? -1 : 1), 0, int(lh) - 1)); + const uint nx = uint(clamp(int(ix) + ((x & 1u) == 0u ? -1 : 1), 0, int(lw) - 1)); + f16vec2 a0[4], a1[4], b0[4], b1[4], s[4]; + load8(p[2], g, lh, lw, iy, ix, LOW8 != 0, a0); load8(p[2], g, lh, lw, iy, nx, LOW8 != 0, a1); + load8(p[2], g, lh, lw, ny, ix, LOW8 != 0, b0); load8(p[2], g, lh, lw, ny, nx, LOW8 != 0, b1); + loadSub(p[3], g, uint(CIN), h, wd, y, x, SKIP8 != 0, s); + [[unroll]] for (uint i = 0u; i < 4u; ++i) { + const f16vec2 r0 = fma(f16vec2(0.25), a1[i], f16vec2(0.75) * a0[i]); + const f16vec2 r1 = fma(f16vec2(0.25), b1[i], f16vec2(0.75) * b0[i]); + v[i] = fma(f16vec2(0.25), r1, f16vec2(0.75) * r0) + s[i]; + } +} + +#if defined(MODE_ENC0) +shared f16vec2 sh[32 * PS * 18]; // pooled pixels, padded to avoid LDS bank conflicts +#endif +// the convolution's input in LDS, loaded once by the workgroup (all its loads in flight, and its waves share it): +// the tile of pixels (with the 3x3 conv's one-pixel halo) x CIN channels, binary16, rows of CINP halves (CIN / 2 + 2 +// words: 2 mod 4, so a 64-bit read's 16 lanes hit distinct bank pairs) +#if defined(MODE_DEC2) && CIN != C +#error "decoder LDS reuse requires CIN == C" +#endif +#define CINP (CIN + 4) +#if defined(MODE_ENC0) +#define TW 16 // four waves of 4 PS x 8 pixels +#define TH (8 * PS) +#elif defined(MODE_ENC3) +#define TW 8 // every wave the same 4 PS x 8 pixels +#define TH (4 * PS) +#endif +#ifdef TW +#define TILEPX ((TH + 2) * (TW + 2)) +#else +#define TILEPX (32 * PS) +#endif +#if TILEPX * CINP * 2 > 65536 +#error "CNN tile exceeds 64 KiB of shared memory; reduce PS" +#endif +shared f16vec4 tile[TILEPX * CINP / 4]; +void tileStore(uint px, uint g, f16vec2 v[4]) +{ + const uint at = (px * uint(CINP) + g * 8u) / 4u; + tile[at] = f16vec4(v[0], v[1]); tile[at + 1u] = f16vec4(v[2], v[3]); +} +void tileLoad(uint px, uint g, out f16vec2 v[4]) +{ + const uint at = (px * uint(CINP) + g * 8u) / 4u; + const f16vec4 a = tile[at], b = tile[at + 1u]; + v[0] = a.xy; v[1] = a.zw; v[2] = b.xy; v[3] = b.zw; +} +#ifdef TW +// the tile of the [CIN/8][H][W][8] input around the block at (y0, x0), zero outside +void loadTile(uint at, uint h, uint wd, int y0, int x0, bool in8) +{ + for (uint i = gl_LocalInvocationIndex; i < uint(TILEPX * CIN / 8); i += gl_WorkGroupSize.x) { + const uint px = i / uint(CIN / 8), g = i % uint(CIN / 8); + const int y = y0 - 1 + int(px / uint(TW + 2)), x = x0 - 1 + int(px % uint(TW + 2)); + f16vec2 v[4]; + if (y >= 0 && y < int(h) && x >= 0 && x < int(wd)) load8(at, g, h, wd, uint(y), uint(x), in8, v); + else { v[0] = f16vec2(0.0); v[1] = f16vec2(0.0); v[2] = f16vec2(0.0); v[3] = f16vec2(0.0); } + tileStore(px, g, v); + } + barrier(); +} +#endif + +#ifdef TW +// the 3x3 conv (CIN -> C) of the tile's pixels (ly + 4 i, lx) (block coordinates), the wave's 32 outputs +void conv3x3(inout f16vec2 acc[PS][16], uint ly, uint lx, uint w1, uint wcol) +{ + [[dont_unroll]] for (uint t = 0u; t < 9u; ++t) { + const uint px = (ly + t / 3u) * uint(TW + 2) + lx + t % 3u; + [[dont_unroll]] for (uint g = 0u; g < uint(CIN) / 8u; ++g) { + f16vec2 v[PS][4]; + [[unroll]] for (uint i = 0u; i < uint(PS); ++i) tileLoad(px + 4u * i * uint(TW + 2), g, v[i]); + fma8(acc, v, w1 + (t * uint(CIN) + g * 8u) * uint(C / 2), uint(C / 2), wcol); + } + } +} +#endif +// conv 1x1 over the tile's rows (32 PS pixels, set i the pixels 32 i + lane), K channels, the wave's 32 outputs +void conv1x1(inout f16vec2 acc[PS][16], uint lane, uint k, uint wrow, uint n, uint wcol) +{ + [[dont_unroll]] for (uint g = 0u; g < k / 8u; ++g) { + f16vec2 v[PS][4]; + [[unroll]] for (uint i = 0u; i < uint(PS); ++i) tileLoad(32u * i + lane, g, v[i]); + fma8(acc, v, wrow + g * 8u * (n / 2u), n / 2u, wcol); + } +} + +void main() +{ + const uint wv = wave(), l = lane(); + f16vec2 acc[PS][16]; + zero(acc); + const uvec2 m = morton(l); +#if defined(MODE_ENC0) + // a 8 PS x 16 tile of L0: wave wv its 4 PS x 8 part, pixel set i the 4 x 8 block i of that + const uint wx = (wv & 1u) * 8u, wy = (wv >> 1u) * 4u * uint(PS); + const uint H0 = PH, W0 = PW; + loadTile(p[2], H0, W0, int(gl_WorkGroupID.y * uint(TH)), int(gl_WorkGroupID.x * 16u), false); + conv3x3(acc, wy + m.y, wx + m.x, p[6], 0u); + biasRelu(acc, p[7], 0u); + [[unroll]] for (uint i = 0u; i < uint(PS); ++i) { + const uint x = gl_WorkGroupID.x * 16u + wx + m.x, y = gl_WorkGroupID.y * uint(TH) + wy + 4u * i + m.y; + storeAll(p[3], 0u, H0, W0, y, x, true, true, 32u, acc[i]); // skip0: 8-bit, phases + pool4(acc[i]); + // the pooled L1 pixel of every lane quad into the tile's 4 PS x 8 L1 block (4 x 8 Morton blocks again) + if ((l & 3u) == 0u) { + const uint lx = (wv & 1u) * 4u + (m.x >> 1u), ly = (wv >> 1u) * 2u * uint(PS) + 2u * i + (m.y >> 1u); + const uint b = ly >> 2u, yy = ly & 3u; + const uint q = b * 32u + ((lx & 1u) | ((yy & 1u) << 1u) | ((lx & 2u) << 1u) | ((yy & 2u) << 2u) | ((lx & 4u) << 2u)); + [[unroll]] for (uint j = 0u; j < 16u; ++j) sh[q * 18u + j] = acc[i][j]; + } + } + barrier(); + // conv 1x1 32 -> 64 at L1: wave wv the outputs 32 (wv & 1) .. +31 of the L1 block wv >> 1 (PS of them) + if (wv < 2u * uint(PS)) { + const uint b = wv >> 1u, hf = wv & 1u; + const uint x1 = gl_WorkGroupID.x * 8u + m.x, y1 = gl_WorkGroupID.y * 4u * uint(PS) + 4u * b + m.y; + f16vec2 a1[1][16]; + [[unroll]] for (uint j = 0u; j < 16u; ++j) a1[0][j] = f16vec2(0.0); + [[unroll]] for (uint g = 0u; g < 4u; ++g) { + const uint at = (b * 32u + l) * 18u + 4u * g; + f16vec2 v[1][4] = f16vec2[1][4](f16vec2[4](sh[at], sh[at + 1u], sh[at + 2u], sh[at + 3u])); +#if PS == 1 + fma8(a1, v, p[8] + g * 8u * 32u, 32u, hf * 16u); +#else + // fma8 takes PS pixel sets: the second conv runs one set per wave + [[unroll]] for (uint kk = 0u; kk < 8u; ++kk) { + const uint r = p[8] + g * 8u * 32u + kk * 32u + hf * 16u; + const float16_t xv = (kk & 1u) == 0u ? v[0][kk >> 1u].x : v[0][kk >> 1u].y; + [[unroll]] for (uint j4 = 0u; j4 < 4u; ++j4) { + const uvec4 qv = w4[r / 4u + j4]; + a1[0][4u * j4] = fma(f16vec2(xv), h2(qv.x), a1[0][4u * j4]); + a1[0][4u * j4 + 1u] = fma(f16vec2(xv), h2(qv.y), a1[0][4u * j4 + 1u]); + a1[0][4u * j4 + 2u] = fma(f16vec2(xv), h2(qv.z), a1[0][4u * j4 + 2u]); + a1[0][4u * j4 + 3u] = fma(f16vec2(xv), h2(qv.w), a1[0][4u * j4 + 3u]); + } + } +#endif + } + { + [[unroll]] for (uint j4 = 0u; j4 < 4u; ++j4) { + const uvec4 qv = w4[(p[9] + hf * 16u) / 4u + j4]; + a1[0][4u * j4] = max(a1[0][4u * j4] + h2(qv.x), f16vec2(0.0)); + a1[0][4u * j4 + 1u] = max(a1[0][4u * j4 + 1u] + h2(qv.y), f16vec2(0.0)); + a1[0][4u * j4 + 2u] = max(a1[0][4u * j4 + 2u] + h2(qv.z), f16vec2(0.0)); + a1[0][4u * j4 + 3u] = max(a1[0][4u * j4 + 3u] + h2(qv.w), f16vec2(0.0)); + } + } + storeAll(p[4], hf * 4u, H0 / 2u, W0 / 2u, y1, x1, true, true, 64u, a1[0]); // skip1: 8-bit, phases + pool4(a1[0]); + if ((l & 3u) == 0u) storeAll(p[5], hf * 4u, H0 / 4u, W0 / 4u, y1 >> 1u, x1 >> 1u, true, false, 64u, a1[0]); // pool1 + } +#elif defined(MODE_ENC3) + // a 4 PS x 8 block of the level, wave wv the outputs 32 wv .. +31; skip (binary16, phases) and its pool (nchw8) + loadTile(p[2], PH, PW, int(gl_WorkGroupID.y * uint(TH)), int(gl_WorkGroupID.x * 8u), IN8 != 0); + conv3x3(acc, m.y, m.x, p[6], wv * 16u); + biasRelu(acc, p[7], wv * 16u); + [[unroll]] for (uint i = 0u; i < uint(PS); ++i) { + const uint x = gl_WorkGroupID.x * 8u + m.x, y = gl_WorkGroupID.y * uint(TH) + 4u * i + m.y; + const bool inside = x < PW && y < PH; + if (inside) storeAll(p[3], wv * 4u, PH, PW, y, x, false, true, uint(C), acc[i]); + pool4(acc[i]); + if (inside && (l & 3u) == 0u) storeAll(p[4], wv * 4u, PH / 2u, PW / 2u, y >> 1u, x >> 1u, false, false, uint(C), acc[i]); + } +#else + // a row-major run of 32 PS pixels (set i: 32 i + lane): its input staged in LDS once for all the waves + for (uint e = gl_LocalInvocationIndex; e < uint(32 * PS * CIN / 8); e += gl_WorkGroupSize.x) { + const uint q = gl_WorkGroupID.x * 32u * uint(PS) + e / uint(CIN / 8), g = e % uint(CIN / 8), qy = q / PW, qx = q % PW; + f16vec2 v[4]; + if (qy >= PH) { v[0] = f16vec2(0.0); v[1] = f16vec2(0.0); v[2] = f16vec2(0.0); v[3] = f16vec2(0.0); } +#if defined(MODE_UP1) || defined(MODE_DEC2) + else upSkip(g, qy, qx, v); +#else + else load8(p[2], g, PH, PW, qy, qx, false, v); +#endif + tileStore(e / uint(CIN / 8), g, v); + } + barrier(); +#if defined(MODE_DEC2) + // conv 1 (C outputs over C / 32 waves) into LDS, then conv 2 (C2 outputs over C2 / 32 waves) + if (wv < uint(C) / 32u) { + conv1x1(acc, l, uint(CIN), p[6], uint(C), wv * 16u); + biasRelu(acc, p[7], wv * 16u); + } + // All waves finish reading the input before the hidden layer reuses its padded LDS rows. + barrier(); + if (wv < uint(C) / 32u) { + [[unroll]] for (uint i = 0u; i < uint(PS); ++i) + [[unroll]] for (uint g = 0u; g < 4u; ++g) { + f16vec2 v[4] = f16vec2[4](acc[i][4u * g], acc[i][4u * g + 1u], acc[i][4u * g + 2u], acc[i][4u * g + 3u]); + tileStore(32u * i + l, wv * 4u + g, v); + } + } + barrier(); + if (wv < uint(C2) / 32u) { + zero(acc); + [[dont_unroll]] for (uint g = 0u; g < uint(C) / 8u; ++g) { + f16vec2 v[PS][4]; + [[unroll]] for (uint i = 0u; i < uint(PS); ++i) { + tileLoad(32u * i + l, g, v[i]); + } + fma8(acc, v, p[8] + g * 8u * uint(C2 / 2), uint(C2 / 2), wv * 16u); + } + biasRelu(acc, p[9], wv * 16u); + [[unroll]] for (uint i = 0u; i < uint(PS); ++i) { + const uint px = gl_WorkGroupID.x * 32u * uint(PS) + 32u * i + l, y = px / PW, x = px % PW; + if (y < PH) storeAll(p[4], wv * 4u, PH, PW, y, x, OUT8 != 0, false, uint(C2), acc[i]); + } + } +#else + conv1x1(acc, l, uint(CIN), p[6], uint(C), wv * 16u); + biasRelu(acc, p[7], wv * 16u); + [[unroll]] for (uint i = 0u; i < uint(PS); ++i) { + const uint px = gl_WorkGroupID.x * 32u * uint(PS) + 32u * i + l, y = px / PW, x = px % PW; + if (y < PH) storeAll(p[4], wv * 4u, PH, PW, y, x, false, false, uint(C), acc[i]); + } +#endif +#endif +} diff --git a/engine/vk_dispatch.cpp b/engine/vk_dispatch.cpp index b404299..2e9393f 100644 --- a/engine/vk_dispatch.cpp +++ b/engine/vk_dispatch.cpp @@ -26,6 +26,10 @@ bool loadVulkan(PFN_vkGetInstanceProcAddr instanceProc, VkInstance instance, VkD ok &= (vkTable.CmdPipelineBarrier = reinterpret_cast(deviceProc(device, "vkCmdPipelineBarrier"))) != nullptr; ok &= (vkTable.CmdUpdateBuffer = reinterpret_cast(deviceProc(device, "vkCmdUpdateBuffer"))) != nullptr; ok &= (vkTable.CmdWriteTimestamp = reinterpret_cast(deviceProc(device, "vkCmdWriteTimestamp"))) != nullptr; + ok &= (vkTable.CmdResetQueryPool = reinterpret_cast(deviceProc(device, "vkCmdResetQueryPool"))) != nullptr; + ok &= (vkTable.CreateQueryPool = reinterpret_cast(deviceProc(device, "vkCreateQueryPool"))) != nullptr; + ok &= (vkTable.DestroyQueryPool = reinterpret_cast(deviceProc(device, "vkDestroyQueryPool"))) != nullptr; + ok &= (vkTable.GetQueryPoolResults = reinterpret_cast(deviceProc(device, "vkGetQueryPoolResults"))) != nullptr; ok &= (vkTable.CreateBuffer = reinterpret_cast(deviceProc(device, "vkCreateBuffer"))) != nullptr; ok &= (vkTable.CreateComputePipelines = reinterpret_cast(deviceProc(device, "vkCreateComputePipelines"))) != nullptr; ok &= (vkTable.CreateDescriptorPool = reinterpret_cast(deviceProc(device, "vkCreateDescriptorPool"))) != nullptr; @@ -46,6 +50,7 @@ bool loadVulkan(PFN_vkGetInstanceProcAddr instanceProc, VkInstance instance, VkD ok &= (vkTable.DestroyShaderModule = reinterpret_cast(deviceProc(device, "vkDestroyShaderModule"))) != nullptr; ok &= (vkTable.FreeMemory = reinterpret_cast(deviceProc(device, "vkFreeMemory"))) != nullptr; ok &= (vkTable.GetBufferMemoryRequirements = reinterpret_cast(deviceProc(device, "vkGetBufferMemoryRequirements"))) != nullptr; + ok &= (vkTable.GetBufferDeviceAddress = reinterpret_cast(deviceProc(device, "vkGetBufferDeviceAddress"))) != nullptr; ok &= (vkTable.GetImageMemoryRequirements = reinterpret_cast(deviceProc(device, "vkGetImageMemoryRequirements"))) != nullptr; ok &= (vkTable.MapMemory = reinterpret_cast(deviceProc(device, "vkMapMemory"))) != nullptr; ok &= (vkTable.UnmapMemory = reinterpret_cast(deviceProc(device, "vkUnmapMemory"))) != nullptr; diff --git a/engine/vk_dispatch.h b/engine/vk_dispatch.h index bf793c6..8a3a9a5 100644 --- a/engine/vk_dispatch.h +++ b/engine/vk_dispatch.h @@ -24,6 +24,10 @@ struct VkTable { PFN_vkCmdPipelineBarrier CmdPipelineBarrier = nullptr; PFN_vkCmdUpdateBuffer CmdUpdateBuffer = nullptr; PFN_vkCmdWriteTimestamp CmdWriteTimestamp = nullptr; + PFN_vkCmdResetQueryPool CmdResetQueryPool = nullptr; + PFN_vkCreateQueryPool CreateQueryPool = nullptr; + PFN_vkDestroyQueryPool DestroyQueryPool = nullptr; + PFN_vkGetQueryPoolResults GetQueryPoolResults = nullptr; PFN_vkCreateBuffer CreateBuffer = nullptr; PFN_vkCreateComputePipelines CreateComputePipelines = nullptr; PFN_vkCreateDescriptorPool CreateDescriptorPool = nullptr; @@ -44,6 +48,7 @@ struct VkTable { PFN_vkDestroyShaderModule DestroyShaderModule = nullptr; PFN_vkFreeMemory FreeMemory = nullptr; PFN_vkGetBufferMemoryRequirements GetBufferMemoryRequirements = nullptr; + PFN_vkGetBufferDeviceAddress GetBufferDeviceAddress = nullptr; PFN_vkGetImageMemoryRequirements GetImageMemoryRequirements = nullptr; PFN_vkMapMemory MapMemory = nullptr; PFN_vkUnmapMemory UnmapMemory = nullptr; @@ -72,6 +77,10 @@ bool loadVulkan(PFN_vkGetInstanceProcAddr instanceProc, VkInstance instance, VkD #define vkCmdPipelineBarrier ::d4r::vkTable.CmdPipelineBarrier #define vkCmdUpdateBuffer ::d4r::vkTable.CmdUpdateBuffer #define vkCmdWriteTimestamp ::d4r::vkTable.CmdWriteTimestamp +#define vkCmdResetQueryPool ::d4r::vkTable.CmdResetQueryPool +#define vkCreateQueryPool ::d4r::vkTable.CreateQueryPool +#define vkDestroyQueryPool ::d4r::vkTable.DestroyQueryPool +#define vkGetQueryPoolResults ::d4r::vkTable.GetQueryPoolResults #define vkCreateBuffer ::d4r::vkTable.CreateBuffer #define vkCreateComputePipelines ::d4r::vkTable.CreateComputePipelines #define vkCreateDescriptorPool ::d4r::vkTable.CreateDescriptorPool @@ -92,6 +101,7 @@ bool loadVulkan(PFN_vkGetInstanceProcAddr instanceProc, VkInstance instance, VkD #define vkDestroyShaderModule ::d4r::vkTable.DestroyShaderModule #define vkFreeMemory ::d4r::vkTable.FreeMemory #define vkGetBufferMemoryRequirements ::d4r::vkTable.GetBufferMemoryRequirements +#define vkGetBufferDeviceAddress ::d4r::vkTable.GetBufferDeviceAddress #define vkGetImageMemoryRequirements ::d4r::vkTable.GetImageMemoryRequirements #define vkMapMemory ::d4r::vkTable.MapMemory #define vkUnmapMemory ::d4r::vkTable.UnmapMemory diff --git a/packaging/build_engine_models.sh b/packaging/build_engine_models.sh index f5baa7b..f0e4c0a 100755 --- a/packaging/build_engine_models.sh +++ b/packaging/build_engine_models.sh @@ -4,9 +4,8 @@ # # usage: bash d4r/tools/build-engine-models.sh [--rr-only | --sr-only] [--exact] [D4R_FOLDER] # bash d4r/tools/build-engine-models.sh --find-proton prints the GE-Proton it would use (or nothing) -# # Ray Reconstruction engine/rr-d, rr-e, rr-f from d4r/nvngx_dlssd.dll 310.9.1 alone (about half a minute) -# Super Resolution engine/k, k-ldr, m, m-ldr, l, l-ldr +# Super Resolution engine/k, k-ldr, m, m-ldr, l, l-ldr, e, f # NVIDIA's DLSS keeps these weights packed in nvngx_dlss.dll; they are read where NGX unpacks # them: the bundled D3D12 test app renders one frame per model through this d4r folder under # GE-Proton, and the uploads are saved (a minute or two each; the first run also compiles @@ -121,13 +120,13 @@ if (( DO_SR )); then fail "d4r's d3d12.dll / d3d12core.dll not found next to $RUNTIME" fi winpath() { printf 'Z:%s' "$(printf '%s' "$1" | tr '/' '\\')"; } - dump() { # dump NAME PRESET: the uploads of one frame of preset PRESET in WORK/NAME/replay + dump() { # dump NAME PRESET PATTERN [FLAGS]: the uploads of one frame of preset PRESET in WORK/NAME/replay local w="$WORK/$1" mkdir -p "$w/replay" printf ' rendering one frame with model %s through d4r (GE-Proton %s)...\n' "$1" "$(basename "$PROTON")" (cd "$APP" && env STEAM_COMPAT_CLIENT_INSTALL_PATH="$STEAM_ROOT" STEAM_COMPAT_DATA_PATH="$WORK/prefix" \ PROTON_FORCE_NVAPI=1 DXVK_NVAPI_GPU_ARCH=AD100 WINEDLLOVERRIDES='d3d12=n,b;d3d12core=n,b;d4r12core=n,b' \ - D4R_ENGINE=0 D4R_DLSS_PRESET="$2" D4R_HARNESS_CREATE_FLAGS=0x4B D4R_SHIM_LOG="$(winpath "$w/shim.log")" \ + D4R_ENGINE=0 D4R_DLSS_PRESET="$2" D4R_HARNESS_CREATE_FLAGS="${4:-0x4B}" D4R_SHIM_LOG="$(winpath "$w/shim.log")" \ D4R_CUDA_REPLAY_DUMP_DIR="$w/replay" D4R_CUDA_REPLAY_DUMP_LIMIT=60 \ "$PROTON/proton" run "$APP/d3d12_dlss_harness.exe" "$(winpath "$RUNTIME/nvngx.dll")" "$(winpath "$w/out.raw")" \ 1 640 360 1280 720) > "$w/run.log" 2>&1 || true @@ -157,5 +156,16 @@ if (( DO_SR )); then done rm -rf -- "${WORK:?}/$1" done + # presets E and F (DLSS 3's CNN): one frame each, LDR with render-size motion (the model serves every flag set) + echo 'Building the Super Resolution CNN models (E and F)...' + PTX="$WORK/ptx" + run ptx.log python3 "$KIT/kernels/tools/extract_dlss_ptx.py" "$DLL" "$PTX" + for c in "e 5 E" "f 6 F"; do + set -- $c + dump "$1" "$2" cuda_engine_input_kernel_rel 0xa + run "$1.log" python3 "$KIT/engine/compile_sr_e.py" "$WORK/$1/replay" "$DLL" "$WORK/$1-out" --preset "$3" --ptx-dir "$PTX" + place "$1" "$WORK/$1-out" $(ls "$WORK/$1-out") + rm -rf -- "${WORK:?}/$1" + done fi printf 'Engine models are in %s/engine.\n' "$RUNTIME" diff --git a/packaging/d4r.ini b/packaging/d4r.ini index 4d60c1b..f9fb88c 100644 --- a/packaging/d4r.ini +++ b/packaging/d4r.ini @@ -11,9 +11,12 @@ ; K | 4 | DLAA / Quality / Balanced | Default; faster alternative to M/L ; M | 4.5 | Performance | More demanding on RDNA3 ; L | 4.5 | Ultra Performance | Especially 4K; experimental -; E | 3 CNN | Quality / Balanced / Performance | Legacy faster fallback +; E | 3 CNN | Quality / Balanced / Performance | Legacy faster fallback; the engine's choice on RDNA2 ; F | 3 CNN | DLAA / Ultra Performance | Deprecated; prefer K or L ; +; On the engine ([Engine]) E and F run from engine\e and engine\f (build-engine-models.sh) without cooperative +; matrices, so they are the models for RDNA2 GPUs (Steam Deck: about 5 ms at 640x360 -> 1280x720). +; ; Select the upscaling mode separately in game/OptiScaler. These are recommendations; ; other modes are allowed. L requires the L texture build. ; game = use the game/OptiScaler choice. Benchmarks: D4R_README.txt. diff --git a/scripts/build_d4r_nvngx_shim.sh b/scripts/build_d4r_nvngx_shim.sh index ae2dd99..e8bbf6c 100755 --- a/scripts/build_d4r_nvngx_shim.sh +++ b/scripts/build_d4r_nvngx_shim.sh @@ -24,7 +24,7 @@ PARAM_OBJECT="$ROOT/build/d4r_ngx_param_msvc.obj" # The engine (D4R_ENGINE) reaches Vulkan through a function table: MinGW has no Vulkan import library. "$MINGW_CXX" -std=c++20 -O2 -Wall -Wextra -Wno-missing-field-initializers -DD4R_ENGINE_VK_TABLE -I"$ROOT/build/vulkan-include" -I"$ROOT/build/vulkan-staging" -shared -static -static-libgcc -static-libstdc++ \ "$ROOT/tools/d4r_nvngx_shim.cpp" "$ROOT/tools/d4r_ngx_param_host.cpp" "$PARAM_OBJECT" \ - "$ROOT/engine/runtime.cpp" "$ROOT/engine/game.cpp" "$ROOT/engine/game_m.cpp" "$ROOT/engine/game_rr.cpp" "$ROOT/engine/vk_dispatch.cpp" \ + "$ROOT/engine/runtime.cpp" "$ROOT/engine/game.cpp" "$ROOT/engine/game_m.cpp" "$ROOT/engine/game_rr.cpp" "$ROOT/engine/game_e.cpp" "$ROOT/engine/vk_dispatch.cpp" \ -ldxgi -lbcrypt -Wl,--entry,d4r_dll_entry -o "$ROOT/build/d4r_nvngx.dll" printf 'Built %s\n' "$ROOT/build/d4r_nvngx.dll" cp "$ROOT/build/d4r_nvngx.dll" "$ROOT/build/d4r.asi" diff --git a/scripts/stage_engine_kit.sh b/scripts/stage_engine_kit.sh index 45a384d..8f52579 100755 --- a/scripts/stage_engine_kit.sh +++ b/scripts/stage_engine_kit.sh @@ -16,7 +16,7 @@ HARNESS="${D4R_HARNESS_EXE:-$ROOT/build/d3d12_dlss_harness.exe}" [[ -f "$HARNESS" ]] || { echo "missing $HARNESS (scripts/build_d3d12_dlss_harness.sh)" >&2; exit 2; } KIT="$D4R/tools/models" rm -rf "$KIT" -mkdir -p "$KIT/bin" "$KIT/engine/rr" "$KIT/kernels/tools" "$KIT/kernels/rr" "$KIT/vulkan" "$D4R/licenses" +mkdir -p "$KIT/bin" "$KIT/engine/rr" "$KIT/engine/sr" "$KIT/engine/cnn_sr" "$KIT/kernels/tools" "$KIT/kernels/rr" "$KIT/vulkan" "$D4R/licenses" copy() { # copy SRC_DIR DST_DIR PATTERN...: never tests local src=$1 dst=$2 pat f shift 2 @@ -30,6 +30,8 @@ copy() { # copy SRC_DIR DST_DIR PATTERN...: never tests } copy engine "$KIT/engine" '*.py' '*.comp' '*.glsl' copy engine/rr "$KIT/engine/rr" '*.py' '*.comp' '*.glsl' '*.txt' '*.npy' '*.sh' +copy engine/cnn_sr "$KIT/engine/cnn_sr" '*.glsl' +copy engine/sr "$KIT/engine/sr" '*.comp' '*.glsl' copy kernels/tools "$KIT/kernels/tools" '*.py' copy kernels/rr "$KIT/kernels/rr" '*.py' copy vulkan "$KIT/vulkan" '*.py' # ptx2glsl.py: engine/rr/ptx2rr.py builds on it diff --git a/tests/test_cnn_mma_gpu.py b/tests/test_cnn_mma_gpu.py new file mode 100644 index 0000000..4307599 --- /dev/null +++ b/tests/test_cnn_mma_gpu.py @@ -0,0 +1,93 @@ +"""Opt-in Vulkan regression for CNN f16 fragments, including independent concurrent waves. + +D4R_RUN_GPU_TESTS=1 D4R_EVK=/path/to/evk python3 -m unittest tests.test_cnn_mma_gpu +Build evk with: g++ -O2 engine/evk.cpp -lvulkan -o build/evk +""" +import os +from pathlib import Path +import subprocess +import tempfile +import unittest + +ROOT = Path(__file__).resolve().parents[1] + + +@unittest.skipUnless(os.environ.get('D4R_RUN_GPU_TESTS') == '1', 'opt-in GPU test') +class CnnMatrixGpuTests(unittest.TestCase): + def test_fragment_mapping_accumulation_and_wave_isolation(self): + self.check_fragments(packed=False) + + def test_packed_fragment_rounding_and_wave_isolation(self): + self.check_fragments(packed=True) + + def check_fragments(self, packed): + import numpy as np + rng = np.random.default_rng(412) + fragments = np.empty((8, 32, 8), dtype='; +.reg .b32 %r<4>; +.reg .f32 %f<2>; +ld.param.u64 %rd0, [args+0]; +ld.param.u64 %rd1, [args+8]; +mov.u32 %r0, 0; +mov.u32 %r1, 15360; +mov.f32 %f0, 0f3F800000; +sust.b.2d.v2.b16.zero [%rd0, {%r0, %r0}], {%r1, %r1}; +sust.p.2d.v4.b32.zero [%rd1, {%r0, %r0}], {%f0, %f0, %f0, %f0}; +ret; +} +''' + with tempfile.TemporaryDirectory() as temporary: + d = Path(temporary) + (d/'surface.ptx').write_text(ptx) + subprocess.run([sys.executable, str(ROOT/'engine/rr/ptx2rr.py'), str(d/'surface.ptx'), 'surface_test', + str(d/'surface.comp'), '--block', '32,1,1', '--weights-arg', '999', + '--surface=0=rg16f', '--surface=8=rgba16f'], check=True, capture_output=True) + for packed in (False, True): + with self.subTest(packed=packed): + subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', + *(['-DD4R_WORD_ARENA', '-DD4R_MMA_PACKED'] if packed else []), + str(d/'surface.comp'), '-o', str(d/'surface.spv')], check=True, capture_output=True) + subprocess.run(['spirv-val', '--target-env', 'vulkan1.3', str(d/'surface.spv')], + check=True, capture_output=True) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_sr_reconstruction_gpu.py b/tests/test_sr_reconstruction_gpu.py new file mode 100644 index 0000000..73d70e9 --- /dev/null +++ b/tests/test_sr_reconstruction_gpu.py @@ -0,0 +1,72 @@ +"""Opt-in rounded Quality reconstruction regression against the generic kernel. + +D4R_RUN_GPU_TESTS=1 D4R_SR_E_MODEL=/private/model/e \ + D4R_SR_REPLAY=/path/to/d4r-e python3 -m unittest tests.test_sr_reconstruction_gpu +Optionally set D4R_SR_F_MODEL to exercise preset F too. No proprietary fixtures are needed. +""" +import os +from pathlib import Path +import subprocess +import tempfile +import unittest + +ROOT = Path(__file__).resolve().parents[1] + + +@unittest.skipUnless(os.environ.get('D4R_RUN_GPU_TESTS') == '1', 'opt-in GPU test') +class SrReconstructionGpuTests(unittest.TestCase): + def test_rounded_quality_e(self): + self.check_model('D4R_SR_E_MODEL') + + def test_rounded_quality_f(self): + self.check_model('D4R_SR_F_MODEL') + + def check_model(self, variable): + import numpy as np + if not os.environ.get(variable): + self.skipTest(f'{variable} must name a private compiled model') + model = Path(os.environ[variable]).resolve() + replay = Path(os.environ.get('D4R_SR_REPLAY', ROOT / 'build/native-engine/d4r-e')).resolve() + env = dict(os.environ, D4R_E_TIME_REPEAT='1', ENABLE_VKBASALT='0') + # Test automatic selection, independent of development overrides in the shell. + env.pop('D4R_ENGINE_E_FINAL', None) + shapes = [(1280, 800), (129, 65), (128, 64), (127, 63)] + jitters = [(-.5, -.5), (.5, .5), (-.499, .499), (.499, -.499), (0, 0), (.25, -.25)] + for ow, oh in shapes: + # Independently round each render axis down/up; also exercise mixed rounding. + for upx, upy in [(0, 0), (1, 1), (0, 1), (1, 0)]: + rw, rh = (2 * ow + 2 * upx) // 3, (2 * oh + 2 * upy) // 3 + for hdr, low in [(0, 1), (1, 1), (0, 0), (1, 0)]: + with self.subTest(size=(rw, rh, ow, oh), hdr=hdr, low_motion=low): + with tempfile.TemporaryDirectory() as temporary: + case = Path(temporary) + mw, mh = (rw, rh) if low else (ow, oh) + rng = np.random.default_rng(123) + lines = [f'size {rw} {rh} {ow} {oh} {hdr} {low} {mw} {mh}'] + for frame, (jx, jy) in enumerate(jitters, 1): + color = rng.random((rh, rw, 4), dtype=np.float32) * (8 if hdr else 1) + color[..., 3] = 1 + color.astype(' engineMDynamic; // Vulkan at changing sizes; original HIP resources stay alive unsigned int engineMLastWidth = 0, engineMLastHeight = 0; unsigned int engineMRenderWidth = 0, engineMRenderHeight = 0; // first evaluation's effective extent + std::unique_ptr engineE; // presets E/F (DLSS 3's CNN): fixed render size, rebuilt on a change + unsigned int engineERenderWidth = 0, engineERenderHeight = 0; std::unique_ptr engineR; // Ray Reconstruction (engine/runtime_rr.inc) ProfileClock::time_point engineRProfileAt{}; // Preset M and preset K with the network on the native HIP layers ([Engine] Network): the buffers shared @@ -10271,6 +10274,179 @@ static bool engine_evaluate_m(Feature& feature, ID3D12GraphicsCommandList* list, return true; } +// Presets E (5) and F (6), DLSS 3's CNN (engine/runtime_e.inc, built by engine/compile_sr_e.py from +// nvngx_dlss.dll 310.7.0): engine\e and engine\f next to this DLL, or D4R_ENGINE_MODEL_E_DIR / _F_DIR. One model +// serves LDR and HDR, both motion resolutions; its network runs without cooperative matrices (RDNA2 included). +struct EngineBackendE +{ + std::once_flag once[2]; + std::unique_ptr models[2]; +}; +static EngineBackendE g_engineE; + +static void init_engine_backend_e(int slot) +{ + const char* const name = slot ? "F" : "E"; + std::string directory = env_string(slot ? "D4R_ENGINE_MODEL_F_DIR" : "D4R_ENGINE_MODEL_E_DIR"); + if (directory.empty()) + { + char folder[MAX_PATH * 3] = {}; + WideCharToMultiByte(CP_ACP, 0, (g_portable.dir + (slot ? L"\\engine\\f" : L"\\engine\\e")).c_str(), -1, folder, + sizeof(folder), nullptr, nullptr); + directory = folder; + } + if (g_vk.interop == nullptr || g_vk.lifetime == nullptr || g_vk.device == VK_NULL_HANDLE || d4r::vkTable.CreateBuffer == nullptr) + { + logf("engine backend (%s): Vulkan interop is unavailable", name); + return; + } + const std::string path = directory[0] == '/' ? "Z:" + directory : directory; + const auto start = std::chrono::steady_clock::now(); + try + { + g_engineE.models[slot] = std::make_unique(g_vk.physical, g_vk.device, path); + } + catch (const std::exception& error) + { + logf("engine backend (%s): cannot load the model from %s: %s", name, path.c_str(), error.what()); + return; + } + logf("engine backend (%s): model %s loaded in %.0f ms", name, path.c_str(), + std::chrono::duration(std::chrono::steady_clock::now() - start).count()); +} + +static bool engine_evaluate_e(Feature& feature, ID3D12GraphicsCommandList* list, ID3D12Resource* color, + ID3D12Resource* depth, ID3D12Resource* motion, ID3D12Resource* output, + ID3D12Resource* exposure, const FrameParams& p, uint32_t frame) +{ + const int slot = feature.preset == 6 ? 1 : 0; + const char* const name = slot ? "F" : "E"; + auto refuse = [&](const char* reason) { + feature.engineRefused = true; + logf("engine backend (%s): %s; feature %u uses the CUDA backend from frame %u", name, reason, feature.handle.Id, frame); + return false; + }; + std::call_once(g_engineE.once[slot], [slot] { init_engine_backend_e(slot); }); + d4r::EModel* const model = g_engineE.models[slot].get(); + if (model == nullptr) + return refuse(slot ? "no preset F model (engine\\f)" : "no preset E model (engine\\e)"); + // NVSDK_NGX_DLSS_Feature_Flags: IsHDR 1, MVLowRes 2, MVJittered 4, DepthInverted 8, AutoExposure 0x40. + const bool hdr = (feature.flags & 0x1) != 0, lowRes = (feature.flags & 0x2) != 0; + if ((feature.flags & 0x40) != 0) + return refuse("auto exposure is not implemented for preset E/F"); + if (color == output || depth == output || motion == output) + return refuse("input/output texture aliasing is not supported"); + if (p.renderWidth > feature.outWidth || p.renderHeight > feature.outHeight) + return refuse("the render size exceeds the output size"); + d4r::GameFrameE game; + using Source = EngineFrameSource; + const auto inputState = static_cast(env_uint("D4R_SHIM_INPUT_STATE", 0x40)); + const auto depthState = static_cast(env_uint("D4R_SHIM_DEPTH_STATE", 0x40)); + const auto outputState = static_cast(env_uint("D4R_SHIM_OUTPUT_STATE", 0x8)); + const uint32_t mw = lowRes ? p.renderWidth : feature.outWidth, mh = lowRes ? p.renderHeight : feature.outHeight; + // HDR: the game's exposure texture is NGX's exposure (its copy kernel moves texel 0 into the 1x1 the kernels read) + const bool gameExposure = hdr && exposure != nullptr; + Source sources[5] = { + {color, &game.color, p.colorBaseX, p.colorBaseY, p.renderWidth, p.renderHeight, inputState, D3D12_RESOURCE_STATE_COPY_SOURCE}, + {depth, &game.depth, p.depthBaseX, p.depthBaseY, p.renderWidth, p.renderHeight, depthState, D3D12_RESOURCE_STATE_COPY_SOURCE}, + {motion, &game.motion, p.mvBaseX, p.mvBaseY, mw, mh, inputState, D3D12_RESOURCE_STATE_COPY_SOURCE}, + {output, &game.output, p.outputBaseX, p.outputBaseY, feature.outWidth, feature.outHeight, outputState, + D3D12_RESOURCE_STATE_COPY_DEST}, + {exposure, &game.exposure, 0, 0, 1, 1, inputState, D3D12_RESOURCE_STATE_COPY_SOURCE}, + }; + const size_t sourceCount = gameExposure ? 5 : 4; + for (size_t index = 0; index < sourceCount; ++index) + { + Source& source = sources[index]; + UINT64 handle = 0, offset = 0; + D3D12_RESOURCE_DESC desc; + source.resource->GetDesc(&desc); + if (FAILED(g_vk.interop->GetVulkanResourceInfo1(source.resource, &handle, &offset, &source.image->format)) || + handle == 0 || desc.Dimension != D3D12_RESOURCE_DIMENSION_TEXTURE2D || desc.SampleDesc.Count != 1 || + source.x + source.width > desc.Width || source.y + source.height > desc.Height) + return refuse("a texture is not a single-sample 2D image covering its rectangle"); + source.image->image = reinterpret_cast(handle); + source.image->x = source.x; + source.image->y = source.y; + prepare_engine_view(feature, source.resource, desc, *source.image, source.image == &game.output, + source.x == 0 && source.y == 0, source.copyState); + g_vk.interop->GetVulkanImageLayout(source.resource, source.copyState, &source.image->layout); + } + if (gameExposure && (game.exposure.view == VK_NULL_HANDLE || (game.exposure.format != VK_FORMAT_R32_SFLOAT && + game.exposure.format != VK_FORMAT_R16_SFLOAT && + game.exposure.format != VK_FORMAT_R32G32B32A32_SFLOAT && + game.exposure.format != VK_FORMAT_R16G16B16A16_SFLOAT))) + return refuse("the exposure texture must be a float texture the engine samples directly"); + if (game.depth.format == VK_FORMAT_D32_SFLOAT || game.depth.format == VK_FORMAT_D32_SFLOAT_S8_UINT) + game.depth.aspect = VK_IMAGE_ASPECT_DEPTH_BIT; + if (!d4r::gameColorFormat(game.color.format) || !d4r::gameMotionFormat(game.motion.format) || + !d4r::gameDepthFormat(game.depth.format, game.depth.aspect) || !d4r::gameOutputFormat(game.output.format)) + { + logf("engine backend (%s): Vulkan formats colour %d, depth %d, motion %d, output %d", name, game.color.format, + game.depth.format, game.motion.format, game.output.format); + return refuse("a texture format is not implemented"); + } + game.jitter[0] = p.jitterX; + game.jitter[1] = p.jitterY; + game.motionScale[0] = p.mvScaleX; + game.motionScale[1] = p.mvScaleY; + const bool resized = feature.engineE != nullptr && + (p.renderWidth != feature.engineERenderWidth || p.renderHeight != feature.engineERenderHeight); + game.reset = p.reset != 0 || feature.engineE == nullptr || resized; + try + { + if (feature.engineE == nullptr || resized) + { + feature.engineE.reset(); + feature.engineE = std::make_unique(*model, g_vk.physical, g_vk.device, feature.outWidth, + feature.outHeight, p.renderWidth, p.renderHeight, hdr, lowRes); + feature.engineERenderWidth = p.renderWidth; + feature.engineERenderHeight = p.renderHeight; + logf("engine backend (%s): feature %u runs preset %s in the game's command list (%ux%u -> %ux%u, %s, %s motion)", + name, feature.handle.Id, name, p.renderWidth, p.renderHeight, feature.outWidth, feature.outHeight, + hdr ? "HDR" : "LDR", lowRes ? "render-size" : "output-size"); + } + } + catch (const std::exception& error) + { + logf("engine backend (%s): %s", name, error.what()); + return refuse("the engine could not be created"); + } + if (FAILED(g_vk.lifetime->RetainExternalResources(list, feature.resources))) + return refuse("the command list cannot retain the engine's resources"); + for (size_t index = 0; index < sourceCount; ++index) + transition(list, sources[index].resource, sources[index].state, sources[index].copyState); + VkCommandBuffer command = VK_NULL_HANDLE; + bool recorded = SUCCEEDED(g_vk.interop->BeginVkCommandBufferInterop(list, &command)); + const char* failure = recorded ? nullptr : "BeginVkCommandBufferInterop failed"; + if (recorded) + { + try + { + feature.engineE->record(command, game); + } + catch (const std::exception& error) + { + logf("engine backend (%s): %s", name, error.what()); + failure = "the frame could not be recorded"; + recorded = false; + } + if (FAILED(g_vk.interop->EndVkCommandBufferInterop(list)) && recorded) + { + failure = "EndVkCommandBufferInterop failed"; + recorded = false; + } + } + for (size_t index = 0; index < sourceCount; ++index) + transition(list, sources[index].resource, sources[index].copyState, sources[index].state); + if (!recorded) + return refuse(failure); + if (frame <= 3 || frame % 600 == 0) + logf("engine frame %u recorded, preset %s (jitter %.6f,%.6f, reset %d, mvScale %.3f,%.3f)", frame, name, p.jitterX, + p.jitterY, game.reset, p.mvScaleX, p.mvScaleY); + return true; +} + // 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 @@ -10641,6 +10817,8 @@ static bool engine_evaluate(Feature& feature, ID3D12GraphicsCommandList* list, I return engine_evaluate_rr(feature, list, color, depth, motion, output, parameters, p, frame); if (feature.preset == 13 || feature.preset == 12) return engine_evaluate_m(feature, list, color, depth, motion, output, p, frame); + if (feature.preset == 5 || feature.preset == 6) + return engine_evaluate_e(feature, list, color, depth, motion, output, exposure, p, frame); if (feature.preset == 11 && env_string("D4R_ENGINE_K") == "0" && env_string("D4R_ENGINE_NETWORK") != "hip") return refuse("[Engine] Enabled = auto keeps K on the native kernels (faster than the Vulkan network)"); // Without the HDR flag NGX runs its LDR input and output kernels: a model directory of its own. @@ -10654,7 +10832,7 @@ static bool engine_evaluate(Feature& feature, ID3D12GraphicsCommandList* list, I return refuse(ldr ? "no LDR model (engine\\k-ldr) for a game without the HDR flag" : "not available"); // NVSDK_NGX_DLSS_Feature_Flags: IsHDR 1, MVLowRes 2, MVJittered 4, DepthInverted 8. if (feature.preset != 11) - return refuse("only presets K, L and M are implemented"); + return refuse("only presets E, F, K, L and M are implemented"); // MVLowRes (0x2) clear: vectors are at output resolution (engine revision D4RO0003). const bool displayMotion = (feature.flags & 0x2) == 0; if (displayMotion && !model->hasDisplayMotion()) diff --git a/vulkan/ptx2glsl.py b/vulkan/ptx2glsl.py index 3ba673e..00c18ea 100644 --- a/vulkan/ptx2glsl.py +++ b/vulkan/ptx2glsl.py @@ -13,7 +13,7 @@ import re import sys REG_RE = re.compile(r'%([a-zA-Z]+)(\d+)$') -GT = {'f': 'float', 'r': 'uint', 'rd': 'uint64_t', 'rs': 'uint', 'p': 'bool'} +GT = {'f': 'float', 'r': 'uint', 'rd': 'uint64_t', 'rs': 'uint', 'p': 'bool', 'fd': 'double'} SREG = {'%tid.x': 'gl_LocalInvocationID.x', '%tid.y': 'gl_LocalInvocationID.y', '%tid.z': 'gl_LocalInvocationID.z', '%ctaid.x': 'gl_WorkGroupID.x', '%ctaid.y': 'gl_WorkGroupID.y', '%ctaid.z': 'gl_WorkGroupID.z', '%ntid.x': 'gl_WorkGroupSize.x', '%ntid.y': 'gl_WorkGroupSize.y', @@ -406,8 +406,10 @@ class T: A(dst, V(src, 'u32'), 'u32'); return # ------------------------------------------------ f32 if ft == 'f32' and base in ('add', 'sub', 'mul'): + if 'rz' in mods and base == 'add': + A(a[0], f'add_rz({V(a[1], "f32")}, {V(a[2], "f32")})', 'float'); return if 'rz' in mods: - A(a[0], f'round_half_away({V(a[1], "f32")})', 'float'); return # the roundf idiom: + copysign(0.5), then trunc + raise Err(f'unsupported rounding: {base}.rz') A(a[0], f'{V(a[1], "f32")} {"+-*"[["add", "sub", "mul"].index(base)]} {V(a[2], "f32")}', 'float'); return if ft == 'f32' and base == 'fma': A(a[0], f'fma({V(a[1], "f32")}, {V(a[2], "f32")}, {V(a[3], "f32")})', 'float'); return @@ -419,6 +421,8 @@ class T: A(a[0], f'abs({V(a[1], "f32")})', 'float'); return if ft == 'f32' and base in ('rcp', 'ex2', 'lg2', 'rsqrt', 'sqrt'): fn = {'rcp': 'rcp_a', 'ex2': 'exp2', 'lg2': 'log2', 'rsqrt': 'inversesqrt', 'sqrt': 'sqrt'}[base] + if base == 'lg2' and 'ftz' in mods: + fn = 'lg2_ftz' # a denormal input is zero: -inf, where log2 returns about -127 A(a[0], f'{fn}({V(a[1], "f32")})', 'float'); return if ft == 'f32' and base == 'div': A(a[0], f'{V(a[1], "f32")} * rcp_a({V(a[2], "f32")})', 'float'); return @@ -428,6 +432,8 @@ class T: e = f'fma({x[0]}, {x[1]}, {x[2]})' if base == 'fma' else f'({x[0]} {"+-*"[["add", "sub", "mul"].index(base)]} {x[1]})' if 'sat' in mods: e = f'min(max({e}, f16vec2(0.0hf)), f16vec2(1.0hf))' + if 'relu' in mods: # negative results clamp to +0 (NaN stays NaN) + e = f'relu2({e})' A(a[0], e, 'f16x2'); return if ft == 'f16x2' and base == 'abs': A(a[0], f'abs({V(a[1], "f16x2")})', 'f16x2'); return @@ -442,6 +448,8 @@ class T: e = f'fma({x[0]}, {x[1]}, {x[2]})' if base == 'fma' else f'({x[0]} {"+-*"[["add", "sub", "mul"].index(base)]} {x[1]})' if 'sat' in mods: e = f'min(max({e}, 0.0hf), 1.0hf)' + if 'relu' in mods: + e = f'relu1({e})' A(a[0], e, 'f16'); return if base == 'set' and 'f16x2' in mods and 'u32' in mods: c = self.cmp(mods, True); x, y = V(a[1], 'f16x2'), V(a[2], 'f16x2') @@ -901,6 +909,14 @@ uint HU2_(u16vec2 u) { return uint(u.x) | (uint(u.y) << 16); } uint H2U(f16vec2 h) { return HU2_(float16BitsToUint16(h)); } int s16(uint u) { return (int(u) << 16) >> 16; } float rcp_a(float x) { return 1.0 / x; } +// add.rz.f32: the sum rounded toward zero (round to nearest, then one ulp back when that went away from zero; the +// error term is exact, TwoSum) +float add_rz(float a, float b) { precise float s = a + b; precise float bb = s - a; precise float e = (a - (s - bb)) + (b - bb); + return (e != 0.0 && ((e < 0.0) != (s < 0.0)) && s != 0.0) ? uintBitsToFloat(floatBitsToUint(s) - 1u) : s; } +float lg2_ftz(float x) { return log2(abs(x) < 1.17549435e-38 ? 0.0 : x); } +// PTX .relu: negative results become +0, NaN stays NaN +f16vec2 relu2(f16vec2 v) { return f16vec2(v.x < 0.0hf ? 0.0hf : v.x, v.y < 0.0hf ? 0.0hf : v.y); } +float16_t relu1(float16_t v) { return v < 0.0hf ? 0.0hf : v; } // x + copysign(0.5, x) rounded toward zero and then truncated: round half away from zero, exactly float round_half_away(float x) { float t = trunc(x); return abs(x - t) >= 0.5 ? t + sign(x) : t; } uint PARAM32(uint o) { return param.w[o >> 2]; }