wouter/d4r
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

engine/rr: native preset F (DLSS 4.5 Ray Reconstruction) and the reference's fused roundings

Browse Source 
- engine/compile_rr2.py, runtime_rr2.inc, engine/rr/rr2_m.comp, exposure_f.comp,
  gen_rr2_enc0.py, rr2_pack.py, rr2_ref.py: the rr2_* model (nvngx_dlssd.dll 310.9's
  third network) as d4r's own shaders - exposure, enc0 gather/block split, enc1-5 and
  dec4-0 through one matrix shader, translated output passes; weights packed from the
  DLL, argument blocks from the capture
- ptx2rr: match the driver's fused roundings - opaque-zero OR around cvt.rn.f16.f32
  against v_fma_mix, half-away-from-zero formatted stores, UNORM 10:10:10:2 as RGBA32F
  of q/1023, OCML tanhF with the AMDGPU exp lowering, binary16 fma and fetch roundings;
  --probe for register bisection against ZLUDA
- dec0_m/swin_m/enc0_*: RR_MIX build option (D4R_RR_MIX=1), one v_fma_mix rounding per
  32-deep step, ~5% faster, not bit-exact; default build stays exact
- runtime_rr.inc: RModel/REngine dispatch to the D4RR0002 model; replay_rr: RGBA32F and
  A2B10G10R1010 texel sizes, roughnessPacked bit, larger debug staging
- shim: engine\rr-f model slot, 310.9 version probe so preset F is only claimed when the
  CUDA path has it; d4r.ini / d4r_config.py Model = F; run_rr_engine_compare.sh D4R_RR_DLL
- kernels/rr/swin_native.h + rr_swin_native_test.cpp: HIP port of encoder 1 for the
  FP8-product and rounding cost measurements
- engine/rr/README: preset F's block, options, capture-by-capture results, timings and
  the open camera-matrix difference
This commit is contained in:
woutermans committed 2026-10-13 13:59:17 +02:00
1 parent 065fe2c2a1
commit 8eb7a94533
31 files changed
+3545 -126
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

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