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

Engine: native preset L (rrlite unfolded input, expansion and reconstruction stages)

Browse Source
This commit is contained in:
woutermans committed 2026-10-10 23:58:18 +02:00
1 parent de37ca5869
commit 653460f7b2
22 files changed
+1109 -237
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

engine/L_NOTES.md
+83
View File
@@ -0,0 +1,83 @@
# Preset L (rrlite, unfolded) - notes
Preset L (DLSS 4.5, render preset 12, NVIDIA's model for Ultra Performance) is preset M's pipeline (M_NOTES.md) with
three stages of its own. NGX runs the kernels without the `_folded` suffix: `rrlite_enc0_4x4_mv{hi,lo}_{hdr,ldr}`,
`rrlite_dec0_4x4` and `rrlite_post_3_{1,2}_mv{hi,lo}_{hdr,ldr}`. The exposure kernels, the ten Swin layers (with L's
own weights) and the downsample are M's. Sources: harness captures (`build/engine-work/lcap`, 640x360 -> 1920x1080;
`lcapq`, 1280x720 -> 1920x1080; `lcapldr`, LDR), the translated kernels (`vulkan/ptx2glsl.py` + `condense.py`, for
reading only) and d4r's native HIP parts of L's texture kernels (`kernels/tex/enc0l_tail.hip`, `dec0l_block.hip`).
## Input stage (`rrlite_enc0_4x4_*`; layer_m.comp KIND 3)
Grid 8x8-pixel blocks, block origin 8 b - (sx, sy); the shift alternates per frame like the layers': (2, 2), (0, 2).
1. The 32 FP8 feature slots per pixel, exactly as M's input stage (`enc0_features.glsl`, M_NOTES.md), including the
mirrored pixels outside the image (their dither uses the unmirrored |position|).
2. Embedding, the shape of an MLP chunk: z = W1 f + b1 (K 32 -> 32), q8(relu(z)), x0 = W2 q8(relu(z)) + b2
(32 -> 32, f16). Weights at 4096 of the allocation (after the 4096-byte dither table): W1 `woff_table(4096, 512,
512, 32, 32)`, b1 at 5120, W2 at 5184 (rows in pair order), b2 at 6208 (pair order).
3. A Swin block of M's shape at pixel resolution: C 32, 2 heads (32 attention channels each), 4x4-pixel windows,
4 MLP chunks; weights at 6272 laid out as a network layer (`swin_model.layout(32, 2)`).
4. Its FP8 output at every pixel of the 32-aligned pixel grid (P16 = render size rounded up to 32): the expansion's
skip input, `[gh][gw][32]`.
5. Patch merge of 2x2 pixels (K 128) to 64 channels, clamped to +-2 pi as f16 before the FP8 encoding: the network's
input tokens, `[gh/2][gw/2][64]`.
## Expansion stage (`rrlite_dec0_4x4`; layer_m.comp KIND 4)
Shift (0, 0) / (2, 0) alternating. The patch expand of the network's 64-channel output (64 -> 4 x 32), x0 =
f16(q8(expand) + skip), the same C 32 block, then the per-pixel head on the block's f16 output (no FP8 rounding).
The head reads NGX's pair-order columns j = 0, 1, 6 (gates), 2-5 (feature), 7-9 (mix), 10-12 (covariance a, b,
rho), which are natural channels 0, 1, 12; 4, 5, 8, 9; 13, 16, 17; 20, 21, 24. These are M's columns without M's
four feature gates:
- feature (SNORM8x4): tanh(v) - the poly below |t| 0.6, 1 - 2 / (e^2t + 1) above, as M's rho - with **no history**:
L feeds the previous feature back only through the input stage's slots.
- gates (UNORM8x4): f16(sigmoid(v)), alpha 0. mix: softmax of three. covariance: a = sigma f16(sigmoid(clamp(v, +-20))),
b likewise, (a^2, b^2, a b tanh(rho v)).
The f16 rounding of the sigmoid has to go through round_half.glsl: RADV folds `float(float16_t(x))`.
## Reconstruction (`rrlite_post_3_*`; post_m.comp PRESET_L)
M's body with two differences:
- the anisotropic Gaussians cover 4x4 render texels, offsets -1 .. +2 around floor(position - 0.5), instead of 3x3.
The engine's tiles are one texel wider per tier: 11 (reconstruction/render >= 3), 12 (compact), 16 (post_3_2);
- the history is blended in linear light: hist = e1 (b1 h1 + b2 h2), each candidate h = expand(bicubic sample) / e2
(when e2 > 0), and a candidate off screen or after a reset takes the colour texel itself (M: compress(e2 colour),
blended compressed, expanded after the blend). Both candidates use the same fallback.
## Constants (identical in every captured mode)
| | M | L |
|---|---:|---:|
| input motion-change scale (enc0 P268) | 9.045066 | 5.137703 |
| input motion-change gain (enc0 P272) | 0.223568 | 0.082259 |
| expansion sigma, rho (dec0 P160, P164) | 8, 0.078633 | 8, 0.097860 |
| first gate gain (post P120) | 2.753011 | 2.979623 |
compile_m.py stores L's five in offsets.bin (`D4RL0001`), read from the capture.
## Found on the way (shared with M)
- **Merge clamp.** NVIDIA's encoders (enc1, enc2, L's enc0) clamp the patch merge to +-2 pi before the FP8 encoding
(kernels/m `merge_clamp`); layer_m.comp did not. Without it enc1's merged output matched 98.6% and enc2's 88.3%.
- **Exposure grid.** NGX measures the luma on an endpoint-aligned grid of half the output, except where the output is
more than twice the render size on an axis: there half the render size, at least 256 (640x360 -> 1920x1080:
320x256; 854x480 -> 2560x1440: 427x256; 1280x720 -> 3840x2160: 640x360; presets M and L alike, eleven harness
geometries). The engine used half the output everywhere, which changed the exposure at Ultra Performance and
with it every value downstream (39 dB).
- **Downsample weights.** glslang folds `float16_t(<literal>)` toward zero, so the first weight was -0.041321 instead
of NVIDIA's -0.041351 (cvt.rn of its f32 constant). As bit patterns the downsample is byte-identical (it was
93.8% / 71 dB on every capture, M's included).
- **Network arithmetic.** With TEXL (ZLUDA's lowering: C enters the first WMMA of a k32 step, natural-order norm
trees, ZLUDA's P V key slots, hardware rsq) the ten layers chain byte-identically against NVIDIA's translated
kernels; the native accuracy kernels' K32_MIX order drifts to about 40 dB at dec1 after the tube.
## Results (RX 9070 XT, RADV, DLSS 310.7)
Stage tests (`test_l.py enc0|dec0`, `test_post_m.py` with `POST_KERNEL`, `test_down_m.py`, `test_m.py --chain`) on
the three captures, every captured frame: 0 differing values. Whole frames:
- `d4r-m` on the three captured frames (640x360 -> 1920x1080): byte-identical to NGX's downsample output.
- D3D12 harness through the shim, engine (Vulkan network) against the CUDA backend running NVIDIA's translated
kernels, five frames each (the CUDA path presents one frame later): byte-identical at 640x360 -> 1920x1080 HDR and
LDR. At 1280x720 -> 1920x1080, 960x540 -> 1920x1080 and 1280x720 -> 3840x2160 one frame each differs in 2,615,
6,079 and 4,207 values (96-122 dB) and the next frame is identical again: an output-only knife edge (the packed
11-11-10 store or the downsample), not state. Both paths are deterministic run to run.
- With d4r's native HIP layers (`compile_m.py --hip`, accuracy build) the network is K32_MIX: 48, 42, 40 dB on the
three captured frames.
GPU time, `d4r-m`, 1280x720 -> 3840x2160 (stand-in frames), median of 60: 3.79 ms with the Vulkan network (exposure
0.01, input 0.67, network 1.44, expansion 0.54, reconstruction 0.92, downsample 0.22), 3.57 ms with the HIP layers
(network 1.12). The CUDA backend's native L tails measure 0.63 ms (enc0) and 0.54 ms (dec0) at the same render size.
engine/README.md
+44 -11
View File
@@ -1,14 +1,15 @@
# Native weight-based inference engine
This directory implements standalone preset K and M inference engines. NGX is used **once to obtain weights
and reference data**. K loads extracted f16 matrices and a reconstruction table; M loads FP8 matrices and
f16 vectors. Both run d4r's own Vulkan shaders. M can run its network on d4r's native HIP kernels instead;
neither engine translates PTX or launches NVIDIA kernels.
This directory implements standalone preset K, M and L inference engines. NGX is used **once to obtain weights
and reference data**. K loads extracted f16 matrices and a reconstruction table; M and L load FP8 matrices and
f16 vectors. All run d4r's own Vulkan shaders. M and L can run their network on d4r's native HIP kernels instead;
no engine translates PTX or launches NVIDIA kernels.
The reusable library is `libd4r_engine.so`; `d4r-k` and `d4r-m` provide full-frame replay and benchmarking.
The D3D12 game shim supports both as opt-in engine backends (see [In a game](#in-a-game-d3d12-shim)).
The CUDA/ZLUDA backend remains the default and fallback. Presets L/E and RTX/game-quality parity are not
implemented or established here. Exposure is supplied by the application or measured on the GPU.
The CUDA/ZLUDA backend remains the default and fallback. Preset E and RTX/game-quality parity are not
implemented or established here (L and M are byte-identical to the CUDA backend running NVIDIA's translated kernels
in the harness; see [Preset L](#preset-l-runtime-and-shim-path-exist-experimental)). Exposure is supplied by the application or measured on the GPU.
## Architecture
@@ -284,9 +285,11 @@ The shim loads each model on first use from a folder next to it, or from a varia
| K, no HDR flag | `engine\k-ldr` | `D4R_ENGINE_MODEL_LDR_DIR` | `compile_k.py --ldr` |
| M, HDR flag set | `engine\m` | `D4R_ENGINE_MODEL_M_DIR` | `compile_m.py` |
| M, no HDR flag | `engine\m-ldr` | `D4R_ENGINE_MODEL_M_LDR_DIR` | `compile_m.py --ldr` |
| L, HDR flag set | `engine\l` | `D4R_ENGINE_MODEL_L_DIR` | `compile_m.py` on a capture of preset L |
| L, no HDR flag | `engine\l-ldr` | `D4R_ENGINE_MODEL_L_LDR_DIR` | `compile_m.py --ldr` on a capture of preset L |
`scripts/install_d4r_runtime.sh` stages them from `D4R_ENGINE_MODEL`, `D4R_ENGINE_MODEL_LDR`, `D4R_ENGINE_MODEL_M`
and `D4R_ENGINE_MODEL_M_LDR`, copying each folder whole (with `hipnet.bin` and `hip/` from `--hip`). A model whose
`scripts/install_d4r_runtime.sh` stages them from `D4R_ENGINE_MODEL`, `D4R_ENGINE_MODEL_LDR`, `D4R_ENGINE_MODEL_M`,
`D4R_ENGINE_MODEL_M_LDR`, `D4R_ENGINE_MODEL_L` and `D4R_ENGINE_MODEL_L_LDR`, copying each folder whole (with `hipnet.bin` and `hip/` from `--hip`). A model whose
`hipnet.bin` belongs to the other preset keeps the Vulkan network. `D4R_ENGINE_NETWORK=vulkan` (`[Engine] Network`)
forces the Vulkan network for both presets; `D4R_SHIM_ENGINE_GPU_WAIT=0` makes the HIP network wait for its inputs on
the CPU-visible marker instead of a GPU-side wait, for comparisons.
@@ -299,8 +302,8 @@ evaluation, CUDA launch, worker thread, frame marker, shared-memory import or sp
the result of the frame it just rendered. The shim log says `engine backend: feature N runs preset K in the game's
command list`.
The shim hands a feature to the CUDA backend, with a log line giving the reason, when: the preset is not K or M;
the feature is preset M with display-resolution motion vectors (MVLowRes clear), which M does not implement yet; the
The shim hands a feature to the CUDA backend, with a log line giving the reason, when: the preset is not K, L or M;
the feature is preset L or M with display-resolution motion vectors (MVLowRes clear), which they do not implement yet; the
motion vectors are display-resolution and the K package predates `D4RO0003` (K with a current package takes them, and
jittered vectors are handled by adding the jitter change, as NGX does); the depth is regular (not inverted) and the K
package predates `D4RO0002`; HDR K runs without AutoExposure and either has no float exposure texture the engine can
@@ -631,7 +634,7 @@ against captured launches:
| network | `layer_m.comp` | per layer on the captured inputs, all ten layers and all three captured frames: byte-identical to the capture (2026-10-09) |
| expansion | `dec0_m.comp` | 96-99.6% of the 8-bit values identical; covariance 70-73 dB |
| reconstruction | `post_m.comp` | 61.5 dB on the reset frame, 76.9 dB with history; packed output 94-98% identical |
| downsample | `down_m.comp` | 70-71 dB, 92% identical |
| downsample | `down_m.comp` | byte-identical (2026-10-11; was 70-71 dB, 92%: its weights, see L_NOTES.md) |
`test_frame_m.py` chains all of them over captured frames with the engine's own intermediate data and history.
Final image, three frames of the panning stress scene:
@@ -848,6 +851,34 @@ flag; they are not current accuracy results. Feature release and re-creation at
clean. The current game timing and longer accuracy comparison are reported above. The kernels are built per
GPU architecture (gfx1201 with FP8 here); another GPU falls back to Vulkan when they do not load.
## Preset L (runtime and shim path exist; experimental)
Preset L (render preset 12, NVIDIA's model for Ultra Performance) is M's pipeline with its own input, expansion and
reconstruction stages: a C 32 Swin block at render-pixel resolution behind an embedding (input) and in front of the
per-pixel head (expansion), where M has single folded products, and a reconstruction with 4x4-texel Gaussians and a
history blended in linear light. `L_NOTES.md` has the algorithms. The network is M's ten layers with L's weights.
The runtime is M's (`MModel`, `MEngine`, `GameUpscalerM`, `d4r-m`): `compile_m.py` recognises a capture of preset L by
its unfolded kernels and writes an L model directory (`offsets.bin` `D4RL0001`: the input and expansion stages as
layer blocks of `layer_m.comp` KIND 3 and 4, and L's five constants), which `MModel` loads as preset L
(`isPresetL()`). The shim runs preset 12 with the model in `engine\l` / `engine\l-ldr`; the log says `engine backend
(L): feature N runs preset L in the game's command list`.
```sh
# capture one frame of preset L with the CUDA backend (D4R_DLSS_PRESET=12; replay filter rrlite_), then
python3 engine/compile_m.py "$CAP" build/l-model # --ldr for a capture without the HDR flag
python3 engine/test_l.py enc0 "$CAP" 0 build/l-work vulkan/kvk # also: dec0; test_post_m.py with POST_KERNEL
```
Every stage is byte-identical to NVIDIA's kernels as ZLUDA translates them on every captured frame (640x360 ->
1920x1080 HDR and LDR, 1280x720 -> 1920x1080); `d4r-m` on the captured frames and the D3D12 harness through the shim
are byte-identical to the CUDA backend except for isolated single-frame knife edges in three of 20 compared frames
(96-122 dB, gone the next frame). 1280x720 -> 3840x2160 takes 3.8 ms on an RX 9070 XT (3.6 ms with `--hip`).
Building L showed four engine issues that M shared, all fixed: the encoders' merge clamp, the exposure grid at
Ultra Performance, the downsample weights and the network's summation order (M is now byte-identical to the CUDA
backend in the harness as well: 1706x960 -> 2560x1440 and 640x360 -> 1920x1080, where it measured 36-41 dB with the
native accuracy kernels' order; `compile_m.py` builds both presets' networks with `TEXL`). Details: `L_NOTES.md`.
## Files
| File | Role |
@@ -866,6 +897,8 @@ GPU architecture (gfx1201 with FP8 here); another GPU falls back to Vulkan when
| `activation_m.glsl`, `gen_activation_m.py` | Preset M: sigmoid/tanh/gate table of the 256 E4M3 codes and the script that computes it (`--check` verifies the table) |
| `test_enc0_m.py`, `test_dec0_m.py`, `test_post_m.py`, `test_down_m.py`, `test_frame_m.py` | Preset M: per-stage and whole-frame tests |
| `M_NOTES.md` | Preset M: the algorithms as reverse-engineered |
| `l_model.py`, `test_l.py`, `L_NOTES.md` | Preset L: weights of the input and expansion stages, their tests against a capture, the algorithms |
| `enc0_features.glsl` | Presets M and L: the input stage's 32 per-pixel feature slots |
| `exposure_k.comp` | Automatic exposure: log-luma block sums, then the exposure state and the stages' exposure fields |
| `test_frame.py`, `test_model.py` | Temporal result checks and model corruption tests |
| `test_game.cpp` | Queued direct-texture, rectangle, layout, conversion and precision regression checks |
engine/compile_m.py
+90 -18
View File
@@ -1,9 +1,12 @@
#!/usr/bin/env python3
"""Build the engine's preset M model directory from a capture of one frame.
"""Build the engine's preset M or preset L model directory from a capture of one frame.
usage: compile_m.py CAPTURE_DIR OUT_DIR [--frame N]
CAPTURE_DIR/replay must hold the frame's rrlite_enc0 and rrlite_dec0 launches (the `_folded` kernels) and its
ten Swin layer launches, with their weight allocations (D4R_CUDA_REPLAY_DUMP_DIR, filter `rrlite_`).
CAPTURE_DIR/replay must hold the frame's rrlite_enc0 and rrlite_dec0 launches and its ten Swin layer launches,
with their weight allocations (D4R_CUDA_REPLAY_DUMP_DIR, filter `rrlite_`). A capture of preset M has the `_folded`
input and expansion kernels; one of preset L (DLSS preset 12) has the unfolded ones, and gets an L model directory
(offsets.bin `D4RL0001`, L_NOTES.md): the same network, layer_m.comp's L input and expansion stages, post_m.comp's
L reconstruction.
Writes weights8.bin (the layers' FP8 matrices in the blocked fixed-16 tile layout of test_m.pack8, with the front-end
enc0/dec0 matrices plain and NGX-ordered as their own stages read them), weights16.bin (f16 vectors and the dither
table), offsets.bin and the compiled shaders. The weights are values read from NVIDIA's DLSS and are not redistributable
@@ -17,6 +20,7 @@ import numpy as np
HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
import m_model as mm
import l_model as lm
import model_enc3 as me
import swin_model as sm
from test_m import pack_layer, defines
@@ -49,6 +53,11 @@ def main():
if len(d) <= a.frame: raise ValueError(f'no launch of {kernel} for frame {a.frame}')
return d[a.frame]
R = 'ldr' if a.ldr else 'hdr'
preset_l = bool(glob.glob(f'{a.capture}/replay/replay-*-rrlite_enc0_4x4_mvlo_{R}'))
if preset_l == bool(glob.glob(f'{a.capture}/replay/replay-*-rrlite_enc0_4x4_mvlo_{R}_folded')):
raise ValueError(f'the capture must hold the input stage of exactly one of preset M or L ({R})')
if preset_l:
return compile_l(a, one, R)
w8, w16 = [], []
n8 = n16 = 0
def add8(x):
@@ -82,7 +91,10 @@ def main():
(a.output / 'offsets.bin').write_bytes(head)
def glsl(src, out, flags=()):
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *flags, str(HERE / src), '-o', str(a.output / out)], check=True, stdout=subprocess.DEVNULL)
for name, out in SHADERS.items(): glsl('layer_m.comp', out + '.spv', defines(name))
# the network in the arithmetic of ZLUDA's lowering of NVIDIA's PTX (TEXL): byte-identical to the CUDA backend
# with NVIDIA's translated kernels in the harness (2026-10-11), where the native accuracy kernels' K32_MIX order
# drifted to 36-41 dB through the FP8 tube
for name, out in SHADERS.items(): glsl('layer_m.comp', out + '.spv', defines(name) + ['-DTEXL=1'])
for s in ('enc0_m', 'dec0_m', 'post_m', 'down_m'): glsl(s + '.comp', s + '.spv', ['-DLDR=1'] if a.ldr and s != 'dec0_m' else [])
glsl('down_m.comp', 'down_m_rgb10.spv', ['-DOUTPUT_RGB10=1'] + (['-DLDR=1'] if a.ldr else []))
glsl('post_m.comp', 'post_m_q.spv', ['-DTILE_W=15', '-DTILE_H=15'] + (['-DLDR=1'] if a.ldr else []))
@@ -90,20 +102,7 @@ def main():
# the same stages with the token buffers in NGX's plane layout, for a network that runs outside the engine
for s in ('enc0_m', 'dec0_m'): glsl(s + '.comp', s + '_planes.spv', ['-DPLANES=1'] + (['-DLDR=1'] if a.ldr and s != 'dec0_m' else []))
if a.hip:
# hipnet.bin: "D4RMHIP1", layer count, then per layer {kernel name[32], shift x, y, token grid divisor, weight
# offset, weight bytes} and the weight allocations as NGX uploads them (the native kernels read that layout)
scale = {'enc1': 1, 'enc2': 2, 'enc3_tube': 4, 'dec2': 2, 'dec1': 1}
blobs, head, at = [], struct.pack('<8sI4x', b'D4RMHIP1', len(Ps)), 16 + 64 * len(Ps)
for P, name in zip(Ps, mm.ORDER):
w = np.ascontiguousarray(P.wbuf); pad = (-w.size) % 64
head += struct.pack('<32siiIIQQ', f'rrlite_{name}_4x4'.encode(), P.sx, P.sy, scale[name], 0, at, w.size)
blobs += [w, np.zeros(pad, np.uint8)]; at += w.size + pad
(a.output / 'hipnet.bin').write_bytes(head + np.concatenate(blobs).tobytes())
(a.output / 'hip').mkdir(exist_ok=True)
for name in sorted(set(mm.ORDER)):
src = a.hip / f'rrlite_{name}_4x4.hsaco'
if not src.exists(): raise ValueError(f'{src} is missing')
(a.output / 'hip' / src.name).write_bytes(src.read_bytes())
package_hip(a, Ps)
marker = a.output / 'ldr' # tells the runtime to skip exposure
if a.ldr: marker.write_text('1\n')
elif marker.exists(): marker.unlink()
@@ -113,5 +112,78 @@ def main():
print(f'Compiled native preset M model in {a.output} ({blob8.size} FP8 bytes, {blob16.size} bytes of f16 data)')
def compile_l(a, one, R):
"""preset L: the network's layers plus the input (KIND 3) and expansion (KIND 4) stages, each packed as a layer"""
w8, w16 = [], []
n8 = n16 = 0
def add(lst, x, n):
x = np.ascontiguousarray(x).view(np.uint8).ravel(); pad = (-x.size) % 64
lst.append(x); lst.append(np.zeros(pad + 64, np.uint8)); return x.size + pad + 64
blocks = []
def layer(t, name):
nonlocal n8, n16
o, b8, b16 = pack_layer(t, name)
blocks.append((n8, n16, o)); n8 += add(w8, b8, n8); n16 += add(w16, b16, n16)
wb = blob_of(one(f'rrlite_enc0_4x4_mvlo_{R}'))
dither = n16; n16 += add(w16, wb[:4096], n16)
layer(lm.enc0_layer(wb), 'enc0l')
Ps, dirs = mm.launches(a.capture, a.frame)
for P, name in zip(Ps, mm.ORDER): layer(mm.plain_layer(P, name), name)
layer(lm.dec0_layer(blob_of(one('rrlite_dec0_4x4'))), 'dec0l')
blob8, blob16 = np.concatenate(w8), np.concatenate(w16)
blob8.tofile(a.output / 'weights8.bin'); blob16.tofile(a.output / 'weights16.bin')
# the stages' constants, as NGX passes them (they differ from preset M's): the input stage's motion-change scale and
# gain, the expansion's covariance sigma and rho, the reconstruction's first-gate gain
arg = lambda kernel, off: struct.unpack_from('<f', open(f'{one(kernel)}/args.bin', 'rb').read(), off)[0]
consts = (arg(f'rrlite_enc0_4x4_mvlo_{R}', 268), arg(f'rrlite_enc0_4x4_mvlo_{R}', 272), arg('rrlite_dec0_4x4', 160),
arg('rrlite_dec0_4x4', 164), arg(next(k for k in (f'rrlite_post_3_1_mvlo_{R}', f'rrlite_post_3_2_mvlo_{R}')
if glob.glob(f'{a.capture}/replay/replay-*-{k}')), 120))
head = struct.pack('<8sIII5f', b'D4RL0001', blob8.size, blob16.size, dither, *consts)
for o8, o16, o in blocks: head += struct.pack('<2I72I', o8, o16, *o)
(a.output / 'offsets.bin').write_bytes(head)
def glsl(src, out, flags=()):
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *flags, str(HERE / src), '-o', str(a.output / out)], check=True, stdout=subprocess.DEVNULL)
ldr = ['-DLDR=1'] if a.ldr else []
# the network in the arithmetic of ZLUDA's lowering of NVIDIA's PTX (TEXL), as for preset M (L_NOTES.md)
for name, out in SHADERS.items(): glsl('layer_m.comp', out + '.spv', defines(name) + ['-DTEXL=1'])
for planes, suffix in ((0, ''), (1, '_planes')):
glsl('layer_m.comp', f'enc0_l{suffix}.spv', defines('enc0l') + ldr + [f'-DPLANES={planes}'])
glsl('layer_m.comp', f'dec0_l{suffix}.spv', defines('dec0l') + [f'-DPLANES={planes}'])
# the reconstruction one render texel wider per tier than M's (4x4 taps): 11 (reconstruction/render >= 3),
# 16 (the whole post_3_2 range), 12 (the compact tier: render/reconstruction <= 0.45, jitter within half a texel)
for out, tile in (('post_l', 11), ('post_l_q', 16), ('post_l_c', 12)):
glsl('post_m.comp', out + '.spv', ['-DPRESET_L=1', f'-DTILE_W={tile}', f'-DTILE_H={tile}'] + ldr)
glsl('down_m.comp', 'down_m.spv', ldr)
glsl('down_m.comp', 'down_m_rgb10.spv', ['-DOUTPUT_RGB10=1'] + ldr)
if a.hip:
package_hip(a, Ps)
marker = a.output / 'ldr'
if a.ldr: marker.write_text('1\n')
elif marker.exists(): marker.unlink()
for p in '01': glsl('exposure_k.comp', f'exposure_m{p}.spv', ['-DMODE=1', '-DPASS=' + p])
for stale in ('enc0_m.spv', 'dec0_m.spv', 'post_m.spv', 'post_m_q.spv', 'post_m_c.spv', 'enc0_m_planes.spv', 'dec0_m_planes.spv'):
(a.output / stale).unlink(missing_ok=True)
files = sorted(a.output.glob('*.bin')) + sorted(a.output.glob('*.spv')) + sorted(a.output.glob('hip/*.hsaco'))
(a.output / 'model.json').write_text(json.dumps({'format': 'D4RL0001', 'preset': 'L', 'files': {str(f.relative_to(a.output)): hashlib.sha256(f.read_bytes()).hexdigest() for f in files}}, indent=1))
print(f'Compiled native preset L model in {a.output} ({blob8.size} FP8 bytes, {blob16.size} bytes of f16 data)')
def package_hip(a, Ps):
"""hipnet.bin: "D4RMHIP1", layer count, then per layer {kernel name[32], shift x, y, token grid divisor, weight
offset, weight bytes} and the weight allocations as NGX uploads them (the native kernels read that layout)"""
scale = {'enc1': 1, 'enc2': 2, 'enc3_tube': 4, 'dec2': 2, 'dec1': 1}
blobs, head, at = [], struct.pack('<8sI4x', b'D4RMHIP1', len(Ps)), 16 + 64 * len(Ps)
for P, name in zip(Ps, mm.ORDER):
w = np.ascontiguousarray(P.wbuf); pad = (-w.size) % 64
head += struct.pack('<32siiIIQQ', f'rrlite_{name}_4x4'.encode(), P.sx, P.sy, scale[name], 0, at, w.size)
blobs += [w, np.zeros(pad, np.uint8)]; at += w.size + pad
(a.output / 'hipnet.bin').write_bytes(head + np.concatenate(blobs).tobytes())
(a.output / 'hip').mkdir(exist_ok=True)
for name in sorted(set(mm.ORDER)):
src = a.hip / f'rrlite_{name}_4x4.hsaco'
if not src.exists(): raise ValueError(f'{src} is missing')
(a.output / 'hip' / src.name).write_bytes(src.read_bytes())
if __name__ == '__main__':
main()
engine/down_m.comp
+8 -2
View File
@@ -30,8 +30,14 @@ layout(set = 0, binding = 3, std430) readonly buffer Src { uint w[]; } src;
shared f16vec4 tile[TILE_W * TILE_H]; // source texels around the block
shared f16vec4 rows[TILE_H * BLOCK_X]; // after the horizontal pass
// weights of taps -2 .. +2 for an even output coordinate; an odd one uses them mirrored
const float16_t W0 = float16_t(-0.041347626596689224), W1 = float16_t(0.08728941529989243), W2 = float16_t(0.6202144026756287),
W3 = float16_t(0.3721286356449127), W4 = float16_t(-0.03828485310077667);
// as binary16 bit patterns: NGX rounds its f32 constants to the nearest binary16 (cvt.rn), while glslang folds
// float16_t(<literal>) toward zero (W0 became -0.041321 instead of -0.041351: 6% of the captured frame's outputs
// one code or more off). The values are -0.041347627, 0.087289415, 0.6202144, 0.37212864, -0.038284853.
#define W0 uint16BitsToFloat16(uint16_t(0xA94Bu))
#define W1 uint16BitsToFloat16(uint16_t(0x2D96u))
#define W2 uint16BitsToFloat16(uint16_t(0x38F6u))
#define W3 uint16BitsToFloat16(uint16_t(0x35F4u))
#define W4 uint16BitsToFloat16(uint16_t(0xA8E7u))
f16vec3 filter5(f16vec3 a, f16vec3 b, f16vec3 c, f16vec3 d, f16vec3 e, bool odd)
{
const f16vec3 w0 = f16vec3(odd ? W4 : W0), w1 = f16vec3(odd ? W3 : W1), w3 = f16vec3(odd ? W1 : W3), w4 = f16vec3(odd ? W0 : W4);
engine/enc0_features.glsl
+119
View File
@@ -0,0 +1,119 @@
// d4r engine, presets M and L: the input stage's per-pixel features, shared by enc0_m.comp (M) and layer_m.comp
// KIND 3 (L, whose input stage gathers the same 32 feature slots). The including shader declares the parameter
// block (EP) and the bindings colorTex, motionTex, depthTex, prevMvTex, halfTex, featTex, mvAOut, mvBOut, mvCOut and
// noise. EP names the parameter block (default p).
#ifndef EP
#define EP p
#endif
const float L10 = 3.3219281;
// LDR (NGX's *_ldr_* kernels, the game did not set the HDR flag): display-referred colour, no exposure and no
// range compression anywhere.
#ifndef LDR
#define LDR 0
#endif
#if LDR
float compress1(float v) { return v; }
#else
float compress1(float v) {
precise float mapped = fma(v, 100.0, 1.0);
precise float lg = log2(max(mapped, 1e-4));
precise float base10 = lg * 0.30103;
precise float result = base10 * 0.36;
return result;
}
#endif
vec3 compress(vec3 v) { return vec3(compress1(v.x), compress1(v.y), compress1(v.z)); }
#if LDR
float colour1(float v) { return isinf(v) || isnan(v) ? 0.0 : min(v, 32752.0); }
#else
float colour1(float v) { return compress1(isinf(v) || isnan(v) ? 0.0 : min(v, 32752.0)); }
#endif
ivec2 texel(ivec2 t) { return ivec2(t.x > 0 ? min(t.x, EP.size.x - 1) : 0, t.y > 0 ? min(t.y, EP.size.y - 1) : 0); }
vec2 motion(ivec2 t) { return EP.mvScale * (EP.mvOffs + texelFetch(motionTex, texel(t), 0).xy); }
// NGX's dithered truncation of an f16 value to the FP8 grid
floate4m3_t dither(float v, uint n)
{
const float16_t h = float16_t(isinf(v) || isnan(v) ? 0.0 : min(v, 1024.0));
const uint e = uint(float16BitsToUint16(h)) & 0xFC00u;
precise float16_t added = h + uint16BitsToFloat16(uint16_t(e | n));
precise float16_t t = added - uint16BitsToFloat16(uint16_t(e));
const float16_t r = uint16BitsToFloat16(uint16_t(uint(float16BitsToUint16(t)) & 0xFF80u));
return floate4m3_t(clamp(r, float16_t(-448.0), float16_t(448.0)));
}
// The 32 dithered FP8 feature slots of block pixel pi (16 per plane) and the motion candidates of the pixel.
void gatherFeatures(uint pi, ivec2 block, float e1, float e2, bool fusedFourth, out floate4m3_t fe[32])
{
const uint px = pi & 7u, py = pi >> 3;
const ivec2 raw = block + ivec2(px, py), a = abs(raw);
const ivec2 m = ivec2(a.x < EP.size.x ? a.x : 2 * EP.size.x - a.x - 2, a.y < EP.size.y ? a.y : 2 * EP.size.y - a.y - 2);
const bool outside = raw.x < 0 || raw.y < 0 || a.x >= EP.size.x || a.y >= EP.size.y;
// the neighbours with the smallest and the largest depth
ivec2 lo = m, hi = m;
float dlo = texelFetch(depthTex, texel(m), 0).r, dhi = dlo;
const ivec2 offs[8] = ivec2[8](ivec2(-1, -1), ivec2(0, -1), ivec2(1, -1), ivec2(-1, 0), ivec2(1, 0), ivec2(-1, 1), ivec2(0, 1), ivec2(1, 1));
for (int k = 0; k < 8; ++k)
{
const float d = texelFetch(depthTex, texel(m + offs[k]), 0).r;
if (d < dlo) { dlo = d; lo = m + offs[k]; }
if (d > dhi) { dhi = d; hi = m + offs[k]; }
}
const vec2 mvA = motion(EP.depthFlag != 0u ? lo : hi), mvB = motion(EP.depthFlag != 0u ? hi : lo), mvC = motion(m);
if (!outside)
{
imageStore(mvAOut, m, vec4(mvA, 0.0, 0.0)); imageStore(mvBOut, m, vec4(mvB, 0.0, 0.0)); imageStore(mvCOut, m, vec4(mvC, 0.0, 0.0));
}
const vec3 own = texelFetch(colorTex, texel(m), 0).rgb;
const vec2 centre = vec2(m) + 0.5, f = centre - EP.jitter, inv = 1.0 / vec2(EP.size);
// how much the motion changed since the previous frame
const vec2 prev = textureLod(prevMvTex, clamp(((f + EP.prevJitter) + mvA * EP.hi2rnd) * vec2(EP.renderInvX, EP.renderInvY), vec2(0.0), vec2(EP.size) * vec2(EP.renderInvX, EP.renderInvY)), 0.0).xy;
const vec2 dm = (EP.changeScale * 540.0 / float(EP.size.y)) * (mvA * EP.hi2rnd - EP.hi2rnd * prev);
const float d2 = fma(dm.x, dm.x, dm.y * dm.y);
const float change = EP.changeGain * (log2(d2 > 0.0 ? sqrt(d2) + 1.0 : 1.0) * 0.6931472);
// the previous result along the candidates
precise vec2 q[4];
q[0] = mvA + f * EP.rnd2hi; q[1] = mvC + f * EP.rnd2hi; q[2] = mvB + f * EP.rnd2hi;
// The captured PTX fuses the first pixel's fourth coordinate, but reuses a rounded
// centre*scale product for the second pixel. A single FMA for both changes FP8 dither.
q[3] = fusedFourth ? fma(centre, EP.rnd2hi, mvC) : centre * EP.rnd2hi + mvC;
vec3 h[4];
for (int k = 0; k < 4; ++k)
{
precise vec3 s = own;
if (!(q[k].x < 0.0 || q[k].y < 0.0 || q[k].x > EP.hiSize.x || q[k].y > EP.hiSize.y || EP.reset != 0u))
#if LDR
s = textureLod(halfTex, q[k] * vec2(EP.halfInvX, EP.halfInvY), 0.0).rgb;
#else
{
precise vec3 packedHistory = textureLod(halfTex, q[k] * vec2(EP.halfInvX, EP.halfInvY), 0.0).rgb * 2.45;
precise vec3 scaled = clamp(packedHistory, 1e-4, 3.0) * (1.0 / 0.36);
precise vec3 exponent = scaled * L10;
precise vec3 linear = (exp2(exponent) - 1.0) * (1.0 / 100.0);
s = linear * (1.0 / e2);
}
#endif
h[k] = compress(e1 * s);
}
vec4 feat = vec4(0.0);
float mf = 0.0;
if (EP.reset == 0u)
{
mf = change;
if (!(q[0].x < 0.0 || q[0].y < 0.0 || q[0].x > EP.hiSize.x || q[0].y > EP.hiSize.y))
{
precise vec2 normalized = fma(q[0], 1.0 / EP.hiSize, EP.prevJitter * inv);
precise vec2 texels = normalized * vec2(EP.size);
precise vec2 uv = texels * vec2(EP.renderInvX, EP.renderInvY);
feat = textureLod(featTex, uv, 0.0);
}
}
const vec3 c = e1 * own;
const vec3 col = vec3(colour1(c.r), colour1(c.g), colour1(c.b));
const uint n = uint(noise.n[(EP.frame << 8) | uint((a.y & 15) << 4) | uint(a.x & 15)]) >> 9;
// 32 slots: 16 per plane
float v[32];
for (int k = 0; k < 32; ++k) v[k] = 0.0;
v[0] = col.r; v[1] = col.g; v[2] = feat.x; v[3] = feat.y; v[4] = col.b; v[5] = h[0].r; v[6] = feat.z; v[7] = feat.w;
v[8] = h[0].g; v[9] = h[0].b; v[12] = h[1].r; v[13] = h[1].g;
v[16] = h[1].b; v[17] = h[2].r; v[20] = h[2].g; v[21] = h[2].b; v[24] = h[3].r; v[25] = h[3].g; v[28] = h[3].b; v[29] = mf;
for (uint k = 0u; k < 32u; ++k) fe[k] = dither(v[k], n);
}
engine/enc0_m.comp
+4 -107
View File
@@ -58,121 +58,18 @@ layout(set = 0, binding = 13, std430) writeonly buffer Out { floate4m3_t a[]; }
#endif
shared floate4m3_t FE[16 * 136];
shared float16_t RS[gl_WorkGroupSize.x * 8u];
const float L10 = 3.3219281;
// LDR (NGX's *_ldr_* kernels, the game did not set the HDR flag): display-referred colour, no exposure and no
// range compression anywhere.
#ifndef LDR
#define LDR 0
#endif
// PLANES: the tokens in NGX's layout, two planes of 32 channels ([2][tok.y][tok.x][32]), for the HIP network
#ifndef PLANES
#define PLANES 0
#endif
#if LDR
float compress1(float v) { return v; }
#else
float compress1(float v) {
precise float mapped = fma(v, 100.0, 1.0);
precise float lg = log2(max(mapped, 1e-4));
precise float base10 = lg * 0.30103;
precise float result = base10 * 0.36;
return result;
}
#endif
vec3 compress(vec3 v) { return vec3(compress1(v.x), compress1(v.y), compress1(v.z)); }
#if LDR
float colour1(float v) { return isinf(v) || isnan(v) ? 0.0 : min(v, 32752.0); }
#else
float colour1(float v) { return compress1(isinf(v) || isnan(v) ? 0.0 : min(v, 32752.0)); }
#endif
ivec2 texel(ivec2 t) { return ivec2(t.x > 0 ? min(t.x, p.size.x - 1) : 0, t.y > 0 ? min(t.y, p.size.y - 1) : 0); }
vec2 motion(ivec2 t) { return p.mvScale * (p.mvOffs + texelFetch(motionTex, texel(t), 0).xy); }
// NGX's dithered truncation of an f16 value to the FP8 grid
floate4m3_t dither(float v, uint n)
{
const float16_t h = float16_t(isinf(v) || isnan(v) ? 0.0 : min(v, 1024.0));
const uint e = uint(float16BitsToUint16(h)) & 0xFC00u;
precise float16_t added = h + uint16BitsToFloat16(uint16_t(e | n));
precise float16_t t = added - uint16BitsToFloat16(uint16_t(e));
const float16_t r = uint16BitsToFloat16(uint16_t(uint(float16BitsToUint16(t)) & 0xFF80u));
return floate4m3_t(clamp(r, float16_t(-448.0), float16_t(448.0)));
}
#include "enc0_features.glsl"
void gatherPixel(uint pi, ivec2 block, float e1, float e2, bool fusedFourth)
{
floate4m3_t fe[32];
gatherFeatures(pi, block, e1, e2, fusedFourth, fe);
const uint px = pi & 7u, py = pi >> 3;
const ivec2 raw = block + ivec2(px, py), a = abs(raw);
const ivec2 m = ivec2(a.x < p.size.x ? a.x : 2 * p.size.x - a.x - 2, a.y < p.size.y ? a.y : 2 * p.size.y - a.y - 2);
const bool outside = raw.x < 0 || raw.y < 0 || a.x >= p.size.x || a.y >= p.size.y;
// the neighbours with the smallest and the largest depth
ivec2 lo = m, hi = m;
float dlo = texelFetch(depthTex, texel(m), 0).r, dhi = dlo;
const ivec2 offs[8] = ivec2[8](ivec2(-1, -1), ivec2(0, -1), ivec2(1, -1), ivec2(-1, 0), ivec2(1, 0), ivec2(-1, 1), ivec2(0, 1), ivec2(1, 1));
for (int k = 0; k < 8; ++k)
{
const float d = texelFetch(depthTex, texel(m + offs[k]), 0).r;
if (d < dlo) { dlo = d; lo = m + offs[k]; }
if (d > dhi) { dhi = d; hi = m + offs[k]; }
}
const vec2 mvA = motion(p.depthFlag != 0u ? lo : hi), mvB = motion(p.depthFlag != 0u ? hi : lo), mvC = motion(m);
if (!outside)
{
imageStore(mvAOut, m, vec4(mvA, 0.0, 0.0)); imageStore(mvBOut, m, vec4(mvB, 0.0, 0.0)); imageStore(mvCOut, m, vec4(mvC, 0.0, 0.0));
}
const vec3 own = texelFetch(colorTex, texel(m), 0).rgb;
const vec2 centre = vec2(m) + 0.5, f = centre - p.jitter, inv = 1.0 / vec2(p.size);
// how much the motion changed since the previous frame
const vec2 prev = textureLod(prevMvTex, clamp(((f + p.prevJitter) + mvA * p.hi2rnd) * vec2(p.renderInvX, p.renderInvY), vec2(0.0), vec2(p.size) * vec2(p.renderInvX, p.renderInvY)), 0.0).xy;
const vec2 dm = (p.changeScale * 540.0 / float(p.size.y)) * (mvA * p.hi2rnd - p.hi2rnd * prev);
const float d2 = fma(dm.x, dm.x, dm.y * dm.y);
const float change = p.changeGain * (log2(d2 > 0.0 ? sqrt(d2) + 1.0 : 1.0) * 0.6931472);
// the previous result along the candidates
precise vec2 q[4];
q[0] = mvA + f * p.rnd2hi; q[1] = mvC + f * p.rnd2hi; q[2] = mvB + f * p.rnd2hi;
// The captured PTX fuses the first pixel's fourth coordinate, but reuses a rounded
// centre*scale product for the second pixel. A single FMA for both changes FP8 dither.
q[3] = fusedFourth ? fma(centre, p.rnd2hi, mvC) : centre * p.rnd2hi + mvC;
vec3 h[4];
for (int k = 0; k < 4; ++k)
{
precise vec3 s = own;
if (!(q[k].x < 0.0 || q[k].y < 0.0 || q[k].x > p.hiSize.x || q[k].y > p.hiSize.y || p.reset != 0u))
#if LDR
s = textureLod(halfTex, q[k] * vec2(p.halfInvX, p.halfInvY), 0.0).rgb;
#else
{
precise vec3 packedHistory = textureLod(halfTex, q[k] * vec2(p.halfInvX, p.halfInvY), 0.0).rgb * 2.45;
precise vec3 scaled = clamp(packedHistory, 1e-4, 3.0) * (1.0 / 0.36);
precise vec3 exponent = scaled * L10;
precise vec3 linear = (exp2(exponent) - 1.0) * (1.0 / 100.0);
s = linear * (1.0 / e2);
}
#endif
h[k] = compress(e1 * s);
}
vec4 feat = vec4(0.0);
float mf = 0.0;
if (p.reset == 0u)
{
mf = change;
if (!(q[0].x < 0.0 || q[0].y < 0.0 || q[0].x > p.hiSize.x || q[0].y > p.hiSize.y))
{
precise vec2 normalized = fma(q[0], 1.0 / p.hiSize, p.prevJitter * inv);
precise vec2 texels = normalized * vec2(p.size);
precise vec2 uv = texels * vec2(p.renderInvX, p.renderInvY);
feat = textureLod(featTex, uv, 0.0);
}
}
const vec3 c = e1 * own;
const vec3 col = vec3(colour1(c.r), colour1(c.g), colour1(c.b));
const uint n = uint(noise.n[(p.frame << 8) | uint((a.y & 15) << 4) | uint(a.x & 15)]) >> 9;
// 32 slots: 16 per plane
float v[32];
for (int k = 0; k < 32; ++k) v[k] = 0.0;
v[0] = col.r; v[1] = col.g; v[2] = feat.x; v[3] = feat.y; v[4] = col.b; v[5] = h[0].r; v[6] = feat.z; v[7] = feat.w;
v[8] = h[0].g; v[9] = h[0].b; v[12] = h[1].r; v[13] = h[1].g;
v[16] = h[1].b; v[17] = h[2].r; v[20] = h[2].g; v[21] = h[2].b; v[24] = h[3].r; v[25] = h[3].g; v[28] = h[3].b; v[29] = mf;
const uint at = ((py >> 1) * 4u + (px >> 1)) * 136u + ((py & 1u) * 2u + (px & 1u)) * 32u;
for (uint k = 0u; k < 32u; ++k) FE[at + k] = dither(v[k], n);
for (uint k = 0u; k < 32u; ++k) FE[at + k] = fe[k];
}
void main()
{
engine/l_model.py
+50
View File
@@ -0,0 +1,50 @@
#!/usr/bin/env python3
"""Preset L (rrlite, unfolded): plain weights of the input stage (rrlite_enc0_4x4_*) and the expansion stage
(rrlite_dec0_4x4), the two stages L does not share with preset M. See L_NOTES.md.
Both are a Swin block of preset M's shape at render-pixel resolution (C 32, 2 heads, 4x4-pixel windows); their weight
allocations are laid out like a network layer's, so m_model.plain_layer reads them:
enc0: the dither table (4096 bytes), the embedding (W1 [32][32] FP8, b1, W2 [32][32] FP8, b2) at 4096, then the
block and its patch merge (128 -> 64 channels) at 6272
dec0: the patch expand (64 -> 4 x 32 channels) and its bias, then the block (m_model.plain_layer with CIN 64)
"""
import os, sys
import numpy as np
HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import m_model as mm
import model_enc3 as me
import swin_model as sm
E_EMB, E_BLOCK = 4096, 6272
mm.LAYERS.update({'enc0l': (32, 2, 64, 0), 'dec0l': (32, 2, 0, 64)})
class Blob:
"""the part of a weight allocation that m_model.plain_layer reads (model_enc3.Params' wbuf and f16vec)"""
def __init__(self, wbuf):
self.wbuf = wbuf
def f16vec(self, off, n):
return self.wbuf[off:off + 2 * n].view(np.float16).copy()
def enc0_layer(wbuf):
"""plain tensors of the input stage: the block, its merge and the embedding (emb_*), natural channel order"""
t = mm.plain_layer(Blob(wbuf[E_BLOCK:]), 'enc0l')
fp8 = lambda T: me.E4M3[wbuf[T]]
b = Blob(wbuf)
g = mm.gp(32)
# the embedding has the layout of an MLP chunk of a C 32 block: fc1's columns in K-slot order, fc2's rows in pair order
t['emb_w1'] = np.asarray(fp8(sm.woff_table(E_EMB, 512, 16 * 32, 32, 32)), np.float16)
t['emb_b1'] = b.f16vec(E_EMB + 32 * 32, 32)
raw = fp8(sm.woff_table(E_EMB + 32 * 32 + 64, 0, 512, 32, 32))
w2 = np.zeros((32, 32)); w2[g] = raw[:, g]
t['emb_w2'] = np.asarray(w2, np.float16)
t['emb_b2'] = b.f16vec(E_EMB + 2 * 32 * 32 + 64, 32)[g]
return t
def dec0_layer(wbuf):
"""plain tensors of the expansion stage: the patch expand and the block"""
return mm.plain_layer(Blob(wbuf), 'dec0l')
engine/layer_m.comp
+290 -19
View File
@@ -31,16 +31,32 @@
// accumulator registers with a pair-aligned half-wave exchange. Without this the ten layers reach only 41-46 dB
// against the capture even though every product, weight map and bias is already the native one.
//
// Build-time shape: C, NH, KIND (0 plain, 1 encoder with 2x2 patch merge, 2 decoder with patch expand + skip),
// XC (merged channels or low-resolution channels). One workgroup is an 8x8-token block: four waves, one per
// 4x4-token window. Bindings and the parameter block are unchanged; the FP8 weight matrices are the packed
// fixed-16 tiles of test_m.pack8 (the f16 vectors stay row-major).
// Build-time shape: C, NH, KIND (0 plain, 1 encoder with 2x2 patch merge, 2 decoder with patch expand + skip,
// 3 preset L's input stage, 4 preset L's expansion stage), XC (merged channels or low-resolution channels). One
// workgroup is an 8x8-token block: four waves, one per 4x4-token window. Bindings 0-6 and the parameter block are
// the same for every kind; the FP8 weight matrices are the packed fixed-16 tiles of test_m.pack8 (the f16 vectors
// stay row-major).
//
// Preset L (L_NOTES.md) runs a C 32, 2-head block of this shape at render-pixel resolution in its input and
// expansion stages, where M has single folded matrix products:
// KIND 3 (rrlite_enc0_4x4_*): enc0_features.glsl's 32 FP8 feature slots per pixel, an embedding
// x0 = W2 q8(relu(W1 f + b1)) + b2 (32 -> 32 -> 32), the block, its FP8 output (dec0's skip input) and a
// patch merge to 64 channels at half resolution (the network's input). Weights o[60..63]: W1, b1, W2, b2.
// KIND 4 (rrlite_dec0_4x4): the patch expand of the network's 64-channel output plus that skip input, the
// block, and the reconstruction's per-pixel inputs from its f16 output (no FP8 rounding, no history).
// TEXL selects the arithmetic of ZLUDA's lowering of NVIDIA's PTX (kernels/tex/swin_tex.h): the bias or previous
// accumulator enters the first of a k32 step's two products as C, the norms sum in natural channel order, the P V
// product's keys take ZLUDA's K slots, and the reciprocal square root is the hardware approximation. compile_m.py
// builds every layer of M and L with it: the chained network and the whole frame are then byte-identical to the CUDA
// backend running NVIDIA's translated kernels (L_NOTES.md). Without it the layers reproduce d4r's native accuracy
// kernels (K32_MIX) instead, which drift from that reference through the FP8 tube.
#extension GL_KHR_cooperative_matrix : require
#extension GL_KHR_memory_scope_semantics : require
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_EXT_float_e4m3 : require
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_GOOGLE_include_directive : require
layout(local_size_x = 128) in;
// o: 0 g1, 1 bo, 2 g2, 3 b2 (f16); per head h at 4 + 4h: q, v (FP8), bias (f16), wo (FP8);
// per MLP chunk c at 20 + 3c: w1 (FP8), b1 (f16), w2 (FP8); 68, 69: merge or expand weights (FP8) and bias (f16)
@@ -49,8 +65,65 @@ layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
layout(set = 0, binding = 2, std430) readonly buffer W16 { f16vec4 w[]; } w16;
layout(set = 0, binding = 3, std430) readonly buffer Ib { floate4m3_t a[]; } inb; // [H][W][C]; decoders: [H/2][W/2][XC]
layout(set = 0, binding = 4, std430) readonly buffer Sb { floate4m3_t a[]; } skb; // decoders: skip [H][W][C]
#if KIND == 4 && defined(L_DUMP)
layout(set = 0, binding = 5, std430) buffer O0 { float16_t a[]; } ldump; // debug: the head's inputs [H][W][C]
#else
layout(set = 0, binding = 5, std430) buffer O0 { floate4m3_t a[]; } out0; // [H][W][C]
#endif
layout(set = 0, binding = 6, std430) buffer O1 { floate4m3_t a[]; } out1; // encoders: merged [H/2][W/2][XC]
#ifndef TEXL
#define TEXL 0
#endif
// PLANES (KIND 3 and 4): the half-resolution tokens in NGX's layout, planes of 32 channels ([XC/32][H/2][W/2][32]),
// for a network that runs outside the engine (the native HIP layers)
#ifndef PLANES
#define PLANES 0
#endif
#if KIND == 3
// the input stage's parameter block (MEnc0Params; W, Hh above are the token pixel grid)
layout(set = 0, binding = 7, std430) readonly buffer Ep
{
ivec2 size, shift, tok;
vec2 jitter, prevJitter;
vec2 mvScale, mvOffs;
vec2 rnd2hi, hi2rnd;
vec2 hiSize;
float changeScale, changeGain;
uint depthFlag, reset, frame;
float expo, expoPrev;
float halfInvX, halfInvY, renderInvX, renderInvY;
} ep;
layout(set = 0, binding = 8) uniform sampler2D colorTex;
layout(set = 0, binding = 9) uniform sampler2D motionTex;
layout(set = 0, binding = 10) uniform sampler2D depthTex;
layout(set = 0, binding = 11) uniform sampler2D prevMvTex; // candidate A of the previous frame, bilinear
layout(set = 0, binding = 12) uniform sampler2D halfTex; // half-resolution previous result, bilinear
layout(set = 0, binding = 13) uniform sampler2D featTex; // recurrent feature, bilinear
layout(set = 0, binding = 14, rg16f) uniform writeonly image2D mvAOut;
layout(set = 0, binding = 15, rg16f) uniform writeonly image2D mvBOut;
layout(set = 0, binding = 16, rg16f) uniform writeonly image2D mvCOut;
layout(set = 0, binding = 17, std430) readonly buffer Noise { uint16_t n[]; } noise; // 8 x 16 x 16
#define EP ep
#include "enc0_features.glsl"
#undef EP
#endif
#if KIND == 4
#include "normalized_store.glsl"
#include "round_half.glsl"
// the expansion stage's parameter block (MDec0Params): render size, the covariance's sigma and rho
layout(set = 0, binding = 7, std430) readonly buffer Dp
{
ivec2 size;
int tokW; uint reset;
vec2 rndOffs, rnd2hi, hiSize, histOffs;
float sigma, rho;
vec2 renderInv;
} dp;
layout(set = 0, binding = 8, rgba8_snorm) uniform writeonly image2D featOut;
layout(set = 0, binding = 9, rgba8) uniform writeonly image2D gateOut;
layout(set = 0, binding = 10, rgba8) uniform writeonly image2D mixOut;
layout(set = 0, binding = 11, rgba16f) uniform writeonly image2D covOut;
#endif
#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>
@@ -66,6 +139,7 @@ const uint VWD = 16u; // row width of the per-wave V trans
int mirror(int v, int n) { v = abs(v); return min(v, 2 * n - 2 - v); }
#define Q16(v) float16_t(v)
#define CL8(v) clamp(v, float16_t(-448.0), float16_t(448.0))
const float16_t MCL = float16_t(6.2831855);
// NVIDIA's accumulator rounding: one binary16 rounding of the whole tile after every logical 32 products.
// RADV folds a register-only f32->f16->f32 conversion away (measured: half of all values and every tie
// midpoint unchanged), so the rounding has to pass through binary16 memory. Each invocation owns 8 halves of
@@ -85,6 +159,18 @@ shared f16vec4 RS[128 * 2];
// operand values as the old plain W[K][N] ColumnMajor stride-N load, but a plain (non-transposed) load.
#define LDA8(dst, base, ks, stride, moff) coopMatLoad((dst), w8.w, (base) + 256u * (uint(ks) * (uint(stride) / 16u) + (uint(moff) / 16u)), 16u, RM)
#if TEXL
// ZLUDA's lowering of m16n8k32 (K32_NO_MIX): the accumulator enters the first product as C, f32 throughout the step,
// one binary16 rounding at its end
#define K32H(dst, a0, b0, a1, b1) { \
ACC P_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(dst))); \
RS[gl_LocalInvocationID.x * 2u + 0u] = f16vec4(float16_t(P_[0]), float16_t(P_[1]), float16_t(P_[2]), float16_t(P_[3])); \
RS[gl_LocalInvocationID.x * 2u + 1u] = f16vec4(float16_t(P_[4]), float16_t(P_[5]), float16_t(P_[6]), float16_t(P_[7])); \
subgroupMemoryBarrierShared(); \
const f16vec4 u0_ = RS[gl_LocalInvocationID.x * 2u + 0u], u1_ = RS[gl_LocalInvocationID.x * 2u + 1u]; \
(dst)[0] = u0_.x; (dst)[1] = u0_.y; (dst)[2] = u0_.z; (dst)[3] = u0_.w; \
(dst)[4] = u1_.x; (dst)[5] = u1_.y; (dst)[6] = u1_.z; (dst)[7] = u1_.w; }
#else
#define K32H(dst, a0, b0, a1, b1) { \
ACC P_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(0.0))); \
RS[gl_LocalInvocationID.x * 2u + 0u] = f16vec4( \
@@ -97,6 +183,7 @@ shared f16vec4 RS[128 * 2];
const f16vec4 u0_ = RS[gl_LocalInvocationID.x * 2u + 0u], u1_ = RS[gl_LocalInvocationID.x * 2u + 1u]; \
(dst)[0] = u0_.x; (dst)[1] = u0_.y; (dst)[2] = u0_.z; (dst)[3] = u0_.w; \
(dst)[4] = u1_.x; (dst)[5] = u1_.y; (dst)[6] = u1_.z; (dst)[7] = u1_.w; }
#endif
// a tiled f16 vector (16 equal rows) as the initial value of an accumulator whose rows are its channels
HAC bias16(uint off, uint stride) { HAC b; coopMatLoad(b, w16.w, off / 4u, stride / 4u, CM); return b; }
#define MK8_F8(dst, a) { \
@@ -123,7 +210,10 @@ float16_t gelu1(float16_t x)
precise float16_t outer = float16_t(0.5) + c * inner;
return x * outer;
}
#if C == 64
#if C == 32
#define EACH_T(X) X(0) X(1)
#define EACH_T16(X) X(0, 1)
#elif C == 64
#define EACH_T(X) X(0) X(1) X(2) X(3)
#define EACH_T16(X) X(0, 1) X(2, 3)
#elif C == 96
@@ -145,7 +235,41 @@ float16_t gelu1(float16_t x)
// #pragma clang fp contract(off)).
#define SH16(v) subgroupShuffleXor(packHalf2x16(vec2(float((v).x), float((v).y))), 16u)
#define UN16(p) f16vec2(float16_t(unpackHalf2x16(p).x), float16_t(unpackHalf2x16(p).y))
#if C == 64
#if TEXL
// rsqrt.approx as ZLUDA lowers it (v_rsq_f32)
#define RSQ16(S) float16_t(inversesqrt(float(float16_t((S) + float16_t(0.0001220703125)))))
#else
#define RSQ16(S) float16_t(1.0 / sqrt(float(float16_t((S) + float16_t(0.0001220703125)))))
#endif
#if TEXL
// SWIN_TOK_NATURAL's tok_rsq: the lane's own channels 16T + 8hf + I; per g the tiles of each parity (WS), the two
// parities, then g, then the halves
#if C == 32
#define WS(P, ww) P[ww]
#elif C == 64
#define WS(P, ww) (P[ww] + P[2 + ww])
#elif C == 96
#define WS(P, ww) ((P[ww] + P[2 + ww]) + P[4 + ww])
#else
#define WS(P, ww) ((P[ww] + P[2 + ww]) + (P[4 + ww] + P[6 + ww]))
#endif
#define RSQ_TREE(OUT) { \
f16vec2 P0_[NT], P1_[NT]; \
EACH_T(RSQ_TILE) \
precise f16vec2 W00_ = WS(P0_, 0), W10_ = WS(P0_, 1), W01_ = WS(P1_, 0), W11_ = WS(P1_, 1); \
precise f16vec2 U0_ = W00_ + W10_, U1_ = W01_ + W11_; \
precise f16vec2 A_ = U0_ + U1_; \
precise f16vec2 S_ = A_ + UN16(SH16(A_)); \
OUT = RSQ16(S_.x + S_.y); }
#elif C == 32
#define RSQ_TREE(OUT) { \
f16vec2 P0_[NT], P1_[NT]; \
EACH_T(RSQ_TILE) \
precise f16vec2 A0_ = P0_[0] + P0_[1], A1_ = P1_[0] + P1_[1]; \
precise f16vec2 U0_ = A0_ + UN16(SH16(A0_)), U1_ = A1_ + UN16(SH16(A1_)); \
precise f16vec2 S_ = U0_ + U1_; \
OUT = RSQ16(S_.x + S_.y); }
#elif C == 64
#define RSQ_TREE(OUT) { \
f16vec2 P0_[NT], P1_[NT]; \
EACH_T(RSQ_TILE) \
@@ -154,7 +278,7 @@ float16_t gelu1(float16_t x)
precise f16vec2 A0_ = W00_ + W10_, A1_ = W01_ + W11_; \
precise f16vec2 U0_ = A0_ + UN16(SH16(A0_)), U1_ = A1_ + UN16(SH16(A1_)); \
precise f16vec2 S_ = U0_ + U1_; \
OUT = float16_t(1.0 / sqrt(float(float16_t((S_.x + S_.y) + float16_t(0.0001220703125))))); }
OUT = RSQ16(S_.x + S_.y); }
#elif C == 96
#define RSQ_TREE(OUT) { \
f16vec2 P0_[NT], P1_[NT]; \
@@ -164,7 +288,7 @@ float16_t gelu1(float16_t x)
precise f16vec2 A0_ = W00_ + W10_, A1_ = W01_ + W11_; \
precise f16vec2 U0_ = A0_ + UN16(SH16(A0_)), U1_ = A1_ + UN16(SH16(A1_)); \
precise f16vec2 S_ = U0_ + U1_; \
OUT = float16_t(1.0 / sqrt(float(float16_t((S_.x + S_.y) + float16_t(0.0001220703125))))); }
OUT = RSQ16(S_.x + S_.y); }
#else
#define RSQ_TREE(OUT) { \
f16vec2 P0_[NT], P1_[NT]; \
@@ -174,22 +298,70 @@ float16_t gelu1(float16_t x)
precise f16vec2 A0_ = W00_ + W10_, A1_ = W01_ + W11_; \
precise f16vec2 U0_ = A0_ + UN16(SH16(A0_)), U1_ = A1_ + UN16(SH16(A1_)); \
precise f16vec2 S_ = U0_ + U1_; \
OUT = float16_t(1.0 / sqrt(float(float16_t((S_.x + S_.y) + float16_t(0.0001220703125))))); }
OUT = RSQ16(S_.x + S_.y); }
#endif
#define ZK(m) { \
const bool up_ = hf != 0u; \
const f16vec2 w0_ = f16vec2((m)[0], (m)[1]), w1_ = f16vec2((m)[2], (m)[3]); \
const f16vec2 w2_ = f16vec2((m)[4], (m)[5]), w3_ = f16vec2((m)[6], (m)[7]); \
const f16vec2 r0_ = UN16(SH16(up_ ? w0_ : w2_)), r1_ = UN16(SH16(up_ ? w1_ : w3_)); \
const f16vec2 a_ = up_ ? r0_ : w0_, b_ = up_ ? w2_ : r0_, c_ = up_ ? r1_ : w1_, d_ = up_ ? w3_ : r1_; \
(m)[0] = a_.x; (m)[1] = a_.y; (m)[2] = b_.x; (m)[3] = b_.y; (m)[4] = c_.x; (m)[5] = c_.y; (m)[6] = d_.x; (m)[7] = d_.y; }
#define RSQ_TILE(T) { \
precise f16vec2 A_ = f16vec2(EV(T, 0), EV(T, 1)), B_ = f16vec2(EV(T, 2), EV(T, 3)); \
precise f16vec2 C_ = f16vec2(EV(T, 4), EV(T, 5)), D_ = f16vec2(EV(T, 6), EV(T, 7)); \
P0_[T] = (A_ * A_) + (B_ * B_); P1_[T] = (C_ * C_) + (D_ * D_); }
shared f16vec4 VSH[4 * 2 * 16 * VWD / 4]; // per wave and head: the V tiles, transposed through shared memory
#if KIND == 1
#define MERGE (KIND == 1 || KIND == 3)
#define EXPAND (KIND == 2 || KIND == 4)
#if MERGE
shared floate4m3_t NB[32 * S4]; // 16 merged rows, preceded by 16 rows of slack for the gather
#endif
#if KIND == 2
#if EXPAND
shared floate4m3_t NLI[16 * SLO]; // the block's 16 low-resolution tokens
shared floate4m3_t NEX[16 * S4]; // their expanded 4C columns
#endif
#if KIND == 3
const uint FS = 40u, XS = 40u; // row strides of the feature and x0 images (32 + slack)
shared floate4m3_t FEW[64 * FS]; // the block's 64 pixels' feature slots, window-major (wave, token)
shared float16_t X0S[64 * XS]; // the embedding's output x0, the block's input
#endif
#if KIND == 4
const uint HS = 40u;
shared float16_t HO[64 * HS]; // the block's f16 output, window-major, for the per-pixel head
#endif
#define SYNC memoryBarrierShared(); barrier();
#if KIND == 4
// NVIDIA's f32 activations of the head (rrlite_dec0_4x4's output code): a sigmoid rounded to binary16, and a
// tanh that is a polynomial below 0.6 and 1 - 2 / (e^2x + 1) above
// (RADV folds a register f32 -> f16 -> f32 conversion away: round_half.glsl's integer rounding keeps it)
float sigmoidL(float v) { return roundHalfValue(1.0 / (exp2(v * -1.442695) + 1.0)); }
float tanhL(float t)
{
const float m = abs(t), s = t * t;
const float big = m >= 9.0109138 ? 1.0 : fma(1.0 / (exp2(m * 2.88539) + 1.0), -2.0, 1.0);
return m >= 0.6 ? (t < 0.0 ? -big : big) :
fma(fma(fma(fma(fma(0.01573968306183815, s, -0.05230396240949631), s, 0.13315297663211823), s, -0.33332768082618713), s, 0.0), t, t);
}
// the block output channel of each head input: NGX's pair-order columns 0, 1, 6 (gates), 2-5 (feature), 7-9 (mix),
// 10-12 (covariance a, b, rho) - preset M's columns without its four feature gates - in natural channel order
#ifndef LH_GATE0
#define LH_GATE0 0
#define LH_GATE1 1
#define LH_GATE2 12
#define LH_MIX0 13
#define LH_MIX1 16
#define LH_MIX2 17
#define LH_COVA 20
#define LH_COVB 21
#define LH_RHO 24
#define LH_FEAT0 4
#define LH_FEAT1 5
#define LH_FEAT2 8
#define LH_FEAT3 9
#endif
#endif
void main()
{
const uint mt = gl_LocalInvocationID.x >> 5, lane = gl_LocalInvocationID.x & 31u;
@@ -202,11 +374,15 @@ void main()
HAC res[NT]; // the residual stream: rows are channels, columns are tokens
B8 h1[NT]; // the normalised activation, the K operand of every product
// ------------------------------------------------------------ patch expand: x0 = q8(E) + skip
#if KIND == 2
#if EXPAND
{
const uint W2 = uint(p.W) / 2u, H2 = uint(p.Hh) / 2u, part = UXC / 8u, c0 = (mt * 2u + hf) * part;
const uint lo = uint(mirror((Y0 >> 1) + int(row >> 2), int(H2)) * int(W2) + mirror((X0 >> 1) + int(row & 3u), p.W / 2)) * UXC;
for (uint c = c0; c < c0 + part; ++c) NLI[row * SLO + c] = inb.a[lo + c];
const uint lt = uint(mirror((Y0 >> 1) + int(row >> 2), int(H2)) * int(W2) + mirror((X0 >> 1) + int(row & 3u), p.W / 2));
#if PLANES
for (uint c = c0; c < c0 + part; ++c) NLI[row * SLO + c] = inb.a[((c >> 5) * H2 * W2 + lt) * 32u + (c & 31u)];
#else
for (uint c = c0; c < c0 + part; ++c) NLI[row * SLO + c] = inb.a[lt * UXC + c];
#endif
SYNC
// this wave computes sub-position mt of every low-resolution token's 4C columns
HAC ef[NT];
@@ -227,12 +403,55 @@ void main()
EACH_T(EF_STORE)
SYNC
}
#endif
#if KIND == 3
// ------------------------------------------------------------ L's input: features, embedding -> x0
{
const float e1 = ep.expo == 0.0 ? 1.0 : ep.expo, e2 = ep.expoPrev == 0.0 ? 1.0 : ep.expoPrev;
const uint li = gl_LocalInvocationID.x;
if (li < 64u)
{
// NGX stages two pixels per lane of one wave (rows y and y + 4); the first of them fuses its fourth
// history coordinate (enc0_features.glsl)
floate4m3_t fe[32];
gatherFeatures(li, ivec2(X0, Y0), e1, e2, li < 32u, fe);
const uint px = li & 7u, py = li >> 3;
const uint at = (((py >> 2) * 2u + (px >> 2)) * 16u + (py & 3u) * 4u + (px & 3u)) * FS;
for (uint k = 0u; k < 32u; ++k) FEW[at + k] = fe[k];
}
SYNC
// the embedding has the shape of an MLP chunk: 32 features -> 32 hidden (ReLU, FP8) -> C
B8 f0, f1;
coopMatLoad(f0, FEW, mt * 16u * FS, FS, CM);
coopMatLoad(f1, FEW, mt * 16u * FS + 16u, FS, CM);
HAC z0 = bias16(p.o[61], 32u), z1 = bias16(p.o[61] + 16u, 32u);
{ A8 a0, a1; LDA8(a0, p.o[60], 0u, 32u, 0u); LDA8(a1, p.o[60], 1u, 32u, 0u); K32H(z0, a0, f0, a1, f1) }
{ A8 a0, a1; LDA8(a0, p.o[60], 0u, 32u, 16u); LDA8(a1, p.o[60], 1u, 32u, 16u); K32H(z1, a0, f0, a1, f1) }
B8 g0, g1;
for (int i = 0; i < 8; ++i)
{
g0[i] = floate4m3_t(CL8(max(float16_t(z0[i]), float16_t(0.0))));
g1[i] = floate4m3_t(CL8(max(float16_t(z1[i]), float16_t(0.0))));
}
const uint xb = (mt * 16u + row) * XS + ch8;
#define EMB_N(N) { \
HAC x_ = bias16(p.o[63] + 16u * (N), UC); \
A8 a0, a1; LDA8(a0, p.o[62], 0u, UC, 16u * (N)); LDA8(a1, p.o[62], 1u, UC, 16u * (N)); \
K32H(x_, a0, g0, a1, g1) \
for (int i = 0; i < 8; ++i) X0S[xb + 16u * (N) + uint(i)] = float16_t(x_[i]); }
EACH_T(EMB_N)
#undef EMB_N
SYNC
}
#endif
// ------------------------------------------------------------ norm 1: sum of squares, residual and h1
#if KIND == 2
#if EXPAND
const uint lb0 = ((byr >> 1) * 4u + (bxr >> 1)) * S4 + ((byr & 1u) * 2u + (bxr & 1u)) * UC + ch8;
const uint sb0 = (iy * uint(p.W) + ix) * UC + ch8;
#define XVAL(C) (float16_t(NEX[lb0 + (C)]) + float16_t(skb.a[sb0 + (C)]))
#elif KIND == 3
const uint xb0 = (mt * 16u + row) * XS + ch8;
#define XVAL(C) X0S[xb0 + (C)]
#else
const uint ib0 = (iy * uint(p.W) + ix) * UC + ch8;
#define XVAL(C) float16_t(inb.a[ib0 + (C)])
@@ -241,12 +460,18 @@ void main()
// 16T + 8*((I>>2)&1) + 4hf + (I&3) (kernels/m/swin_tok.h tok_rsq on the register layout the input plane implies).
// This shader's own accumulators and A rows keep natural channels, so h1 and the residual use XVAL (the lane's
// own half); the norm tree, however, groups per the native element order, so it reads the absolute channel.
#if KIND == 2
#if EXPAND
#define XVALN(C) (float16_t(NEX[(lb0 - ch8) + (C)]) + float16_t(skb.a[(sb0 - ch8) + (C)]))
#elif KIND == 3
#define XVALN(C) X0S[(xb0 - ch8) + (C)]
#else
#define XVALN(C) float16_t(inb.a[(iy * uint(p.W) + ix) * UC + (C)])
#endif
#if TEXL
#define EV(T, I) XVAL(16u * (T) + uint(I)) // natural order: the lane's own channels
#else
#define EV(T, I) XVALN(16u * (T) + 8u * ((uint(I) >> 2) & 1u) + 4u * hf + (uint(I) & 3u))
#endif
float16_t rn; RSQ_TREE(rn)
#undef EV
#define NM1(N) { \
@@ -320,6 +545,12 @@ void main()
A16 vt0, vt1;
coopMatLoad(vt0, VSH, (mt * 2u + 0u) * 64u, VWD / 4u, RM);
coopMatLoad(vt1, VSH, (mt * 2u + 1u) * 64u, VWD / 4u, RM);
#if TEXL
// P V is the block's one product of two f16 operands, and its f32 sum over the 16 keys is not exact: ZLUDA's
// lowering of m16n8k16 puts key pair (j & 1) * 4 + 2 h + (j >> 1) in K dword j of half h (swin_tok.h
// SWIN_PV_ZLUDA_ORDER), not 4 h + j. Both operands are permuted alike; it changes about one token in 30,000.
ZK(pb); ZK(vt0); ZK(vt1);
#endif
o[0] = coopMatMulAdd(vt0, pb, ACC(0.0));
o[1] = coopMatMulAdd(vt1, pb, ACC(0.0));
R16(o[0]); R16(o[1]);
@@ -337,6 +568,9 @@ void main()
// the native element (T, 8h+i) holds the natural channel tok_xnat(T, 8h+i) = 16T + 4h + (i&3) + 8*((i>>2)&1),
// which lives in my natural-order accumulator at the half ((i>>2)&1) and the element 4h+(i&3). The tree needs
// the pairs in that order, so fetch the (pair 2h, pair 2h+1) of the half ((i>>2)&1) and select by this lane's half.
#if TEXL
#define EV(T, I) float16_t(res[T][I]) // natural order: the lane's own elements
#else
f16vec2 XA_[NT], XB_[NT], XC_[NT], XD_[NT];
for (uint t_ = 0u; t_ < NT; ++t_)
{
@@ -350,6 +584,7 @@ void main()
}
#define EV(T, I) ((uint(I) < 2u) ? XA_[T][uint(I) & 1u] : (uint(I) < 4u) ? XB_[T][uint(I) & 1u] : \
(uint(I) < 6u) ? XC_[T][uint(I) & 1u] : XD_[T][uint(I) & 1u])
#endif
float16_t rn2; RSQ_TREE(rn2)
#undef EV
#define N2(N) { \
@@ -398,6 +633,33 @@ void main()
K32H(res[N], a0, g0, a1, g1) }
EACH_T(MLP2_N)
}
#if KIND == 4
// ------------------------------------------------------------ L's expansion head: per-pixel reconstruction inputs
{
const uint hb = (mt * 16u + row) * HS + ch8;
for (uint nt = 0u; nt < NT; ++nt)
for (int i = 0; i < 8; ++i) HO[hb + 16u * nt + uint(i)] = float16_t(res[nt][i]);
SYNC
const uint li = gl_LocalInvocationID.x;
const uint px = li & 7u, py = li >> 3;
const ivec2 at = ivec2(X0 + int(px), Y0 + int(py));
if (li < 64u && at.x >= 0 && at.y >= 0 && at.x < dp.size.x && at.y < dp.size.y)
{
const uint hp = (((py >> 2) * 2u + (px >> 2)) * 16u + (py & 3u) * 4u + (px & 3u)) * HS;
#ifdef L_DUMP
for (uint c = 0u; c < UC; ++c) ldump.a[(uint(at.y) * uint(p.W) + uint(at.x)) * UC + c] = HO[hp + c];
#endif
#define V(c) float(HO[hp + uint(c)])
imageStore(gateOut, at, unorm8Value(vec4(sigmoidL(V(LH_GATE0)), sigmoidL(V(LH_GATE1)), sigmoidL(V(LH_GATE2)), 0.0)));
const vec3 m = vec3(V(LH_MIX0), V(LH_MIX1), V(LH_MIX2)), ex = exp2((m - max(max(m.x, m.y), m.z)) * 1.442695);
imageStore(mixOut, at, unorm8Value(vec4(ex * (1.0 / (ex.z + (ex.x + ex.y))), 0.0)));
const float a = dp.sigma * sigmoidL(clamp(V(LH_COVA), -20.0, 20.0)), b = dp.sigma * sigmoidL(clamp(V(LH_COVB), -20.0, 20.0));
imageStore(covOut, at, vec4(a * a, b * b, a * b * tanhL(V(LH_RHO) * dp.rho), 0.0));
imageStore(featOut, at, snorm8Value(vec4(tanhL(V(LH_FEAT0)), tanhL(V(LH_FEAT1)), tanhL(V(LH_FEAT2)), tanhL(V(LH_FEAT3)))));
#undef V
}
}
#else
// ------------------------------------------------------------ the block output, FP8
const uint dst = uint(ty * p.W + tx) * UC + ch8;
if (inside)
@@ -407,8 +669,9 @@ void main()
OUT8(0) OUT8(1) OUT8(2) OUT8(3) OUT8(4) OUT8(5) OUT8(6) OUT8(7)
#undef OUT8
}
#endif
// ------------------------------------------------------------ patch merge: 2x2 tokens (4C) -> XC, half resolution
#if KIND == 1
#if MERGE
{
// the 16 merged rows are gathered from the four waves: the row is the merged token of the lane's pixel
const uint mrow = (byr >> 1) * 4u + (bxr >> 1), qq = (byr & 1u) * 2u + (bxr & 1u);
@@ -418,7 +681,13 @@ void main()
const uint MT = UXC / 16u, NI = (MT + 3u) / 4u;
const int ym = (Y0 >> 1) + int(row >> 2), xm = (X0 >> 1) + int(row & 3u);
const bool mok = ym >= 0 && ym < p.Hh / 2 && xm >= 0 && xm < p.W / 2;
const uint md = uint(ym * (p.W / 2) + xm) * UXC + ch8;
#if PLANES
const uint mt2 = uint(ym * (p.W / 2) + xm), mplane = uint(p.Hh / 2) * uint(p.W / 2) * 32u;
#define MADDR(C) ((C) >> 5) * mplane + mt2 * 32u + ((C) & 31u)
#else
const uint md = uint(ym * (p.W / 2) + xm) * UXC;
#define MADDR(C) md + (C)
#endif
for (uint i = 0u; i < NI; ++i)
{
const uint u = mt + 4u * i;
@@ -435,9 +704,11 @@ void main()
}
if (mok)
{
#define MGO(I) out1.a[md + 16u * u + uint(I)] = floate4m3_t(CL8(y[I]));
// NVIDIA's encoders clamp the merged result to +-2 pi (as f16) before its e4m3 encoding (kernels/m merge_clamp)
#define MGO(I) out1.a[MADDR(16u * u + ch8 + uint(I))] = floate4m3_t(clamp(float16_t(y[I]), -MCL, MCL));
MGO(0) MGO(1) MGO(2) MGO(3) MGO(4) MGO(5) MGO(6) MGO(7)
#undef MGO
#undef MADDR
}
}
}
engine/post_m.comp
+41 -2
View File
@@ -9,6 +9,12 @@
// worst case of its range, 1.5x). TILE_W/TILE_H select the tile for the engine's own 16x16-pixel block: 10x10
// covers reconstruction/render >= 3; 11x11 covers render/reconstruction <= 0.45 with jitter in [-0.5, 0.5];
// 15x15 covers the whole post_3_2 range (render/reconstruction <= 2/3).
//
// PRESET_L: preset L's rrlite_post_3_* (unfolded). The same body with two differences (L_NOTES.md): the Gaussians
// cover the 4x4 render texels from -1 to +2 around floor(position - 0.5), so every tile tier is one texel wider
// (11, 12, 16), and the history is blended in linear light: each candidate's bicubic sample is expanded and divided by
// the previous exposure on its own, a candidate off screen (or after a reset) takes the linear colour texel, and the
// blend is scaled by the current exposure.
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_GOOGLE_include_directive : require
#include "normalized_store.glsl"
@@ -42,6 +48,14 @@ layout(set = 0, binding = 11, std430) writeonly buffer Packed { uint w[]; } pack
#ifndef TILE_H
#define TILE_H 10
#endif
#ifndef PRESET_L
#define PRESET_L 0
#endif
#if PRESET_L
const int TAP_LO = -1, TAP_HI = 2;
#else
const int TAP_LO = -1, TAP_HI = 1;
#endif
shared f16vec4 tile0[TILE_W * TILE_H], tile1[TILE_W * TILE_H], tile2[TILE_W * TILE_H];
shared vec3 colSum[128]; // vertical pair first, then the neighbouring column
const float L10 = 3.3219281; // log2(10)
@@ -88,6 +102,17 @@ vec3 history(vec2 q)
textureLod(histTex, vec2(lo.x, mid.y), 0.0).rgb * a))));
return s * (1.0 / (e + (d + (c + (b + a)))));
}
#if PRESET_L
vec3 historyLinear(vec2 q, float e2)
{
const vec3 h = history(q);
#if LDR
return h;
#else
return e2 > 0.0 ? expand(h) * (1.0 / e2) : h;
#endif
}
#endif
void main()
{
const ivec2 lid = ivec2(gl_LocalInvocationID.xy), gid = ivec2(gl_GlobalInvocationID.xy);
@@ -154,10 +179,10 @@ void main()
f16vec4 a0 = f16vec4(0.0), a1 = f16vec4(0.0);
float16_t a2 = float16_t(0.0);
float W = 0.0;
for (int i = -1; i <= 1; ++i)
for (int i = TAP_LO; i <= TAP_HI; ++i)
{
const float16_t dx = d0.x + float16_t(i), sq = dx * dx;
for (int j = -1; j <= 1; ++j)
for (int j = TAP_LO; j <= TAP_HI; ++j)
{
const int k = (base.y - origin.y + j) * TILE_W + (base.x - origin.x + i);
const f16vec4 t2 = tile2[k];
@@ -186,6 +211,19 @@ void main()
// the fallback colour is only fetched where the history sample leaves the image or the history was reset:
// the value is the same either way, so computing it in the branch saves the fetch on ordinary pixels
const bool out1 = q1.x < 0.0 || q1.y < 0.0 || q1.x > size.x || q1.y > size.y || p.reset != 0u;
#if PRESET_L
// linear light: a sample's expansion (when the previous exposure is positive) or the colour texel
const vec3 h1 = out1 ? texelFetch(colorTex, colourTexel(near), 0).rgb : historyLinear(q1, p.expoPrev == 0.0 ? 1.0 : p.expoPrev);
vec3 h2 = vec3(0.0);
if (b2 > 0.0)
{
const vec2 q2 = centre + second.xy;
if (q2.x < 0.0 || q2.y < 0.0 || q2.x > size.x || q2.y > size.y || p.reset != 0u)
h2 = out1 ? h1 : texelFetch(colorTex, colourTexel(near), 0).rgb;
else h2 = historyLinear(q2, p.expoPrev == 0.0 ? 1.0 : p.expoPrev);
}
const vec3 hist = e1 * fma(vec3(b1), h1, b2 * h2);
#else
const vec3 h1 = out1 ? compress(e2 * texelFetch(colorTex, colourTexel(near), 0).rgb) : history(q1);
vec3 h2 = vec3(0.0);
if (b2 > 0.0)
@@ -196,6 +234,7 @@ void main()
else h2 = history(q2);
}
const vec3 hist = e1 * (expand(fma(vec3(b1), h1, b2 * h2)) * (1.0 / e2));
#endif
const vec3 lin = expand(cur);
const float g0 = clamp(p.gain * W * g.x, 0.0, 1.0);
const float ga = fma(g0, 1.0 - g.y, g.y), gb = fma(g0, 1.0 - g.z, g.z);
engine/replay_m.cpp
+14
View File
@@ -185,6 +185,20 @@ int main(int argc, char** argv) try {
if (hip) d[2] = networkUs; // measured HIP network time; the complete span includes host handoffs
const double us = double(t[6] - t[0]) * tick;
if (pass) for (int k = 0; k < 6; ++k) stage[k].push_back(d[k]);
if (pass == 0 && std::getenv("D4R_M_DUMP")) {
// development: every intermediate of this frame, to PREFIX-NAME-FRAME.bin
for (const auto& d : engine.debugResources()) {
Staging st = staging(c, d.bytes);
begin();
if (d.buffer) { VkBufferCopy bc{0, 0, d.bytes}; vkCmdCopyBuffer(cb, d.buffer, st.buffer, 1, &bc); }
else { VkBufferImageCopy ic{}; ic.imageSubresource = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1}; ic.imageExtent = {d.width, d.height, 1};
vkCmdCopyImageToBuffer(cb, d.image, VK_IMAGE_LAYOUT_GENERAL, st.buffer, 1, &ic); }
submit();
std::ofstream(std::string(std::getenv("D4R_M_DUMP")) + "-" + d.name + "-" + std::to_string(f) + ".bin", std::ios::binary)
.write(static_cast<const char*>(st.map), std::streamsize(d.bytes));
vkDestroyBuffer(c.device, st.buffer, nullptr); vkFreeMemory(c.device, st.memory, nullptr);
}
}
if (pass == 0) {
std::ofstream(prefix + "-" + std::to_string(f) + ".rgba16f", std::ios::binary).write(static_cast<const char*>(out.map), std::streamsize(ow) * oh * 8);
std::printf("frame %d: %.1f us GPU\n", f, us);
engine/runtime.h
+6
View File
@@ -4,6 +4,7 @@
#include <array>
#include <cstdint>
#include <memory>
#include <vector>
#include <string>
namespace d4r {
@@ -202,6 +203,7 @@ public:
bool uploadRecorded() const;
bool hasExternalNetworkStages() const; // the model directory has the plane-layout input and expansion stages
bool hasRgb10Output() const;
bool isPresetL() const; // a preset L model directory (offsets.bin D4RL0001): MEngine then runs preset L
private:
struct Impl;
std::unique_ptr<Impl> impl;
@@ -230,6 +232,10 @@ public:
void recordFront(VkCommandBuffer command, VkQueryPool stamps = VK_NULL_HANDLE, uint32_t first = 0);
void recordBack(VkCommandBuffer command, VkQueryPool stamps = VK_NULL_HANDLE, uint32_t first = 0);
VkImage outputImage() const; // RGBA16F, GENERAL
// Development: the engine's intermediate buffers and render-resolution images (GENERAL), for replays that compare
// them with a capture. Valid while the engine lives; read them only after a recorded frame has completed.
struct DebugResource { const char* name; VkBuffer buffer; VkImage image; VkDeviceSize bytes; uint32_t width, height; };
std::vector<DebugResource> debugResources() const;
private:
struct Impl;
std::unique_ptr<Impl> impl;
engine/runtime_m.inc
+113 -20
View File
@@ -19,6 +19,11 @@ struct MDec0Params { int32_t size[2], tokW; uint32_t reset; float rndOffs[2], rn
struct MPostParams { int32_t colOrigin[2], hiSize[2], rndSize[2]; float hi2rnd[2], rndOffs[2], invRnd[2], gain; uint32_t reset; float expo, expoPrev; };
struct MDownParams { int32_t srcSize[2], outOrigin[2], colOrigin[2], colSize[2]; float out2rnd[2], rndOffs[2], expo; };
struct MLayerParams { int32_t w, h, sx, sy; uint32_t o[72]; };
// The grid NGX measures the exposure on (cuda_dldn_engine_luma_convert_kernel's launch, presets M and L alike,
// harness captures 2026-10-11): half the output, unless the output is more than twice the render size on that axis
// (Ultra Performance), then half the render size but at least 256 (640x360 -> 1920x1080: 320x256, 854x480 ->
// 2560x1440: 427x256, 1280x720 -> 3840x2160: 640x360).
uint32_t exposureGrid(uint32_t render, uint32_t output) { return output / 2 <= render ? output / 2 : std::max(render / 2, 256u); }
static_assert(sizeof(MEnc0Params) == 124 && offsetof(MEnc0Params, expo) == 100 && offsetof(MEnc0Params, halfInv) == 108, "enc0_m.comp and exposure_k.comp (words 25, 26)");
static_assert(sizeof(MDec0Params) == 64, "dec0_m.comp");
static_assert(sizeof(MPostParams) == 64 && offsetof(MPostParams, expo) == 56, "post_m.comp and exposure_k.comp (words 14, 15)");
@@ -72,6 +77,12 @@ struct MModel::Impl {
bool planes = false;
std::array<Pipeline, 2> exposure;
bool uploaded = false, ldr = false;
// preset L (offsets.bin D4RL0001, L_NOTES.md): the input and expansion stages are layer_m.comp's KIND 3 and 4,
// with their own layer offset blocks; constants: change scale and gain, sigma, rho, the first gate's gain
bool presetL = false;
std::array<uint32_t, 74> lEnc0{}, lDec0{};
uint32_t dither16 = 0;
float lConst[5]{};
Impl(VkPhysicalDevice physical, VkDevice dev) : device(physical, dev) {}
};
MModel::MModel(VkPhysicalDevice physical, VkDevice device, const std::string& directory, VkPipelineCache cache)
@@ -79,16 +90,27 @@ MModel::MModel(VkPhysicalDevice physical, VkDevice device, const std::string& di
auto& m = *impl;
m.ldr = exists(directory + "/ldr"); // the stages were compiled for a game without the HDR flag: no exposure
auto head = read(directory + "/offsets.bin");
const size_t expect = 16 + 5 * 4 + MLayers * 74 * 4;
if (head.size() != expect || std::memcmp(head.data(), "D4RM0003", 8)) throw std::runtime_error("incompatible M model metadata");
m.presetL = head.size() >= 8 && !std::memcmp(head.data(), "D4RL0001", 8);
const size_t expect = m.presetL ? 40 + (MLayers + 2) * 74 * 4 : 16 + 5 * 4 + MLayers * 74 * 4;
if (head.size() != expect || (!m.presetL && std::memcmp(head.data(), "D4RM0003", 8))) throw std::runtime_error("incompatible M/L model metadata");
uint32_t n8 = 0, n16 = 0;
std::memcpy(&n8, head.data() + 8, 4); std::memcpy(&n16, head.data() + 12, 4);
auto b8 = read(directory + "/weights8.bin"), b16 = read(directory + "/weights16.bin");
if (b8.size() != n8 || b16.size() != n16 || n16 % 64) throw std::runtime_error("invalid M model sizes");
std::memcpy(m.enc0, head.data() + 16, 12); std::memcpy(m.dec0, head.data() + 28, 8);
std::memcpy(m.layers.data(), head.data() + 36, sizeof(m.layers));
auto inside = [&](uint32_t off8, uint32_t off16) { if (off8 >= n8 || off16 >= n16 || off8 % 64 || off16 % 64) throw std::runtime_error("invalid M model offsets"); };
inside(m.enc0[0], m.enc0[1]); inside(m.dec0[0], m.dec0[1]); inside(0, m.enc0[2]);
if (m.presetL) {
// the input stage's block, the ten layers, the expansion stage's block
std::memcpy(&m.dither16, head.data() + 16, 4); std::memcpy(m.lConst, head.data() + 20, sizeof(m.lConst));
std::memcpy(m.lEnc0.data(), head.data() + 40, sizeof(m.lEnc0));
std::memcpy(m.layers.data(), head.data() + 40 + sizeof(m.lEnc0), sizeof(m.layers));
std::memcpy(m.lDec0.data(), head.data() + 40 + sizeof(m.lEnc0) + sizeof(m.layers), sizeof(m.lDec0));
inside(m.lEnc0[0], m.lEnc0[1]); inside(m.lDec0[0], m.lDec0[1]); inside(0, m.dither16);
for (float c : m.lConst) if (!(c > 0.f && c < 100.f)) throw std::runtime_error("invalid L model constants");
} else {
std::memcpy(m.enc0, head.data() + 16, 12); std::memcpy(m.dec0, head.data() + 28, 8);
std::memcpy(m.layers.data(), head.data() + 36, sizeof(m.layers));
inside(m.enc0[0], m.enc0[1]); inside(m.dec0[0], m.dec0[1]); inside(0, m.enc0[2]);
}
for (auto& l : m.layers) inside(l[0], l[1]);
const VkBufferUsageFlags usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT;
m.w8.create(m.device, b8.size(), usage); m.w16.create(m.device, b16.size(), usage);
@@ -96,16 +118,21 @@ MModel::MModel(VkPhysicalDevice physical, VkDevice device, const std::string& di
std::memcpy(m.up8.mapped, b8.data(), b8.size()); std::memcpy(m.up16.mapped, b16.data(), b16.size());
const auto b = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, t = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, s = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE;
for (unsigned i = 0; i < MLayers; ++i) m.network[i].create(device, directory + "/" + MShader[i] + ".spv", {b,b,b,b,b,b,b}, 32, cache);
m.input.create(device, directory + "/enc0_m.spv", {b,t,t,t,t,t,t,s,s,s,b,b,b,b}, 32, cache);
m.expand.create(device, directory + "/dec0_m.spv", {b,t,t,s,s,s,s,b,b,b}, 32, cache);
if ((m.planes = exists(directory + "/enc0_m_planes.spv") && exists(directory + "/dec0_m_planes.spv"))) {
m.inputPlanes.create(device, directory + "/enc0_m_planes.spv", {b,t,t,t,t,t,t,s,s,s,b,b,b,b}, 32, cache);
m.expandPlanes.create(device, directory + "/dec0_m_planes.spv", {b,t,t,s,s,s,s,b,b,b}, 32, cache);
// layer_m.comp's L stages: bindings 0-6 of a layer, then the stage's own (layer_m.comp KIND 3 and 4)
const std::vector<VkDescriptorType> lIn{b,b,b,b,b,b,b,b,t,t,t,t,t,t,s,s,s,b}, lEx{b,b,b,b,b,b,b,b,s,s,s,s};
const std::string in = m.presetL ? "enc0_l" : "enc0_m", ex = m.presetL ? "dec0_l" : "dec0_m", po = m.presetL ? "post_l" : "post_m";
const auto inTypes = m.presetL ? lIn : std::vector<VkDescriptorType>{b,t,t,t,t,t,t,s,s,s,b,b,b,b};
const auto exTypes = m.presetL ? lEx : std::vector<VkDescriptorType>{b,t,t,s,s,s,s,b,b,b};
m.input.create(device, directory + "/" + in + ".spv", inTypes, 32, cache);
m.expand.create(device, directory + "/" + ex + ".spv", exTypes, 32, cache);
if ((m.planes = exists(directory + "/" + in + "_planes.spv") && exists(directory + "/" + ex + "_planes.spv"))) {
m.inputPlanes.create(device, directory + "/" + in + "_planes.spv", inTypes, 32, cache);
m.expandPlanes.create(device, directory + "/" + ex + "_planes.spv", exTypes, 32, cache);
}
m.post.create(device, directory + "/post_m.spv", {b,t,t,t,t,t,t,t,t,s,s,b}, 64, cache);
m.postQuality.create(device, directory + "/post_m_q.spv", {b,t,t,t,t,t,t,t,t,s,s,b}, 64, cache);
if (exists(directory + "/post_m_c.spv"))
m.postCompact.create(device, directory + "/post_m_c.spv", {b,t,t,t,t,t,t,t,t,s,s,b}, 64, cache);
m.post.create(device, directory + "/" + po + ".spv", {b,t,t,t,t,t,t,t,t,s,s,b}, 64, cache);
m.postQuality.create(device, directory + "/" + po + "_q.spv", {b,t,t,t,t,t,t,t,t,s,s,b}, 64, cache);
if (exists(directory + "/" + po + "_c.spv"))
m.postCompact.create(device, directory + "/" + po + "_c.spv", {b,t,t,t,t,t,t,t,t,s,s,b}, 64, cache);
m.down.create(device, directory + "/down_m.spv", {b,t,s,b}, 64, cache);
if (exists(directory + "/down_m_rgb10.spv"))
m.downRgb10.create(device, directory + "/down_m_rgb10.spv", {b,t,s,b}, 64, cache);
@@ -116,6 +143,7 @@ MModel::~MModel() = default;
bool MModel::uploadRecorded() const { return impl->uploaded; }
bool MModel::hasExternalNetworkStages() const { return impl->planes; }
bool MModel::hasRgb10Output() const { return impl->downRgb10.handle != VK_NULL_HANDLE; }
bool MModel::isPresetL() const { return impl->presetL; }
void MModel::recordUpload(VkCommandBuffer cb) {
auto& m = *impl;
if (m.uploaded) throw std::runtime_error("model upload already recorded");
@@ -131,6 +159,7 @@ struct MEngine::Impl {
Buffer enc0P, dec0P, postP, downP, expP, expSums, expState, dummy, packed;
std::array<Buffer, MLayers> layerP;
Buffer tokens, skip1, merged1, skip2, merged2, tubeA, tubeB, dec2Out, dec1Out;
Buffer skipL, enc0LP, dec0LP; // preset L: the input stage's pixel output and the two stages' layer blocks
std::array<Tex, 2> mvA, feat, hist, half;
Tex mvB, mvC, gates, mix, cov, output;
VkDescriptorPool pool = VK_NULL_HANDLE;
@@ -183,7 +212,7 @@ MEngine::MEngine(MModel& model, uint32_t renderWidth, uint32_t renderHeight, uin
const VkBufferUsageFlags sb = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, sbt = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT;
e.enc0P.create(d, sizeof(MEnc0Params), sbt); e.dec0P.create(d, sizeof(MDec0Params), sbt); e.postP.create(d, sizeof(MPostParams), sbt);
e.downP.create(d, sizeof(MDownParams), sbt); e.expP.create(d, sizeof(ExposureParams), sbt);
e.expSums.create(d, VkDeviceSize((outputWidth / 2 + 31) / 32) * ((outputHeight / 2 + 31) / 32) * 4, sb);
e.expSums.create(d, VkDeviceSize((exposureGrid(renderWidth, outputWidth) + 31) / 32) * ((exposureGrid(renderHeight, outputHeight) + 31) / 32) * 4, sb);
e.expState.create(d, 8, sbt); e.dummy.create(d, 256, sb);
e.packed.create(d, VkDeviceSize(e.hw) * e.hh * 4, sb);
auto tok = [&](Buffer& b, unsigned scale, unsigned channels) { b.create(d, VkDeviceSize(tw / scale) * (th / scale) * channels + 64, sb); };
@@ -227,8 +256,31 @@ MEngine::MEngine(MModel& model, uint32_t renderWidth, uint32_t renderHeight, uin
bindBuffer(dev, set, 0, e.layerP[i]); bindRange(dev, set, 1, m.w8, m.layers[i][0]); bindRange(dev, set, 2, m.w16, m.layers[i][1]);
bindBuffer(dev, set, 3, *in[i]); bindBuffer(dev, set, 4, *skip[i]); bindBuffer(dev, set, 5, *out[i]); bindBuffer(dev, set, 6, *merged[i]);
}
if (m.presetL) {
// the input stage's pixel-resolution block output (the expansion's skip input) and the two stages' layer
// parameter blocks (token pixel grid, window shift, offsets; the shift alternates per frame like the layers')
e.skipL.create(d, VkDeviceSize(e.gw) * e.gh * 32 + 64, sb);
e.enc0LP.create(d, sizeof(MLayerParams), sbt, true); e.dec0LP.create(d, sizeof(MLayerParams), sbt, true);
MLayerParams lp{int32_t(e.gw), int32_t(e.gh), 2, 2, {}};
std::memcpy(lp.o, m.lEnc0.data() + 2, sizeof(lp.o)); std::memcpy(e.enc0LP.mapped, &lp, sizeof(lp));
lp.sx = lp.sy = 0; std::memcpy(lp.o, m.lDec0.data() + 2, sizeof(lp.o)); std::memcpy(e.dec0LP.mapped, &lp, sizeof(lp));
}
for (unsigned i = 0; i < 2; ++i) { // i: the index of this frame's history outputs
const unsigned prev = 1 - i;
if (m.presetL) {
auto s0 = e.enc0Sets[i] = descriptorSet(dev, e.pool, (external ? m.inputPlanes : m.input).descriptors);
bindBuffer(dev, s0, 0, e.enc0LP); bindRange(dev, s0, 1, m.w8, m.lEnc0[0]); bindRange(dev, s0, 2, m.w16, m.lEnc0[1]);
bindBuffer(dev, s0, 3, e.dummy); bindBuffer(dev, s0, 4, e.dummy); bindBuffer(dev, s0, 5, e.skipL); bindBuffer(dev, s0, 6, e.tokens);
bindBuffer(dev, s0, 7, e.enc0P);
bindImage(dev, s0, 11, e.mvA[prev].view, e.linear); bindImage(dev, s0, 12, e.half[prev].view, e.linear); bindImage(dev, s0, 13, e.feat[prev].view, e.linear);
bindImage(dev, s0, 14, e.mvA[i].view); bindImage(dev, s0, 15, e.mvB.view); bindImage(dev, s0, 16, e.mvC.view);
bindRange(dev, s0, 17, m.w16, m.dither16);
auto s1 = e.dec0Sets[i] = descriptorSet(dev, e.pool, (external ? m.expandPlanes : m.expand).descriptors);
bindBuffer(dev, s1, 0, e.dec0LP); bindRange(dev, s1, 1, m.w8, m.lDec0[0]); bindRange(dev, s1, 2, m.w16, m.lDec0[1]);
bindBuffer(dev, s1, 3, e.dec1Out); bindBuffer(dev, s1, 4, e.skipL); bindBuffer(dev, s1, 5, e.dummy); bindBuffer(dev, s1, 6, e.dummy);
bindBuffer(dev, s1, 7, e.dec0P);
bindImage(dev, s1, 8, e.feat[i].view); bindImage(dev, s1, 9, e.gates.view); bindImage(dev, s1, 10, e.mix.view); bindImage(dev, s1, 11, e.cov.view);
} else {
auto s0 = e.enc0Sets[i] = descriptorSet(dev, e.pool, (external ? m.inputPlanes : m.input).descriptors);
bindBuffer(dev, s0, 0, e.enc0P);
bindImage(dev, s0, 4, e.mvA[prev].view, e.linear); bindImage(dev, s0, 5, e.half[prev].view, e.linear); bindImage(dev, s0, 6, e.feat[prev].view, e.linear);
@@ -238,6 +290,7 @@ MEngine::MEngine(MModel& model, uint32_t renderWidth, uint32_t renderHeight, uin
bindBuffer(dev, s1, 0, e.dec0P); bindImage(dev, s1, 1, e.mvA[i].view, e.nearest); bindImage(dev, s1, 2, e.feat[prev].view, e.linear);
bindImage(dev, s1, 3, e.feat[i].view); bindImage(dev, s1, 4, e.gates.view); bindImage(dev, s1, 5, e.mix.view); bindImage(dev, s1, 6, e.cov.view);
bindBuffer(dev, s1, 7, e.dec1Out); bindRange(dev, s1, 8, m.w8, m.dec0[0]); bindRange(dev, s1, 9, m.w16, m.dec0[1]);
}
auto s2 = e.postSets[i] = descriptorSet(dev, e.pool, m.post.descriptors);
bindBuffer(dev, s2, 0, e.postP);
bindImage(dev, s2, 2, e.mix.view, e.nearest); bindImage(dev, s2, 3, e.gates.view, e.nearest); bindImage(dev, s2, 4, e.cov.view, e.nearest);
@@ -273,16 +326,18 @@ void MEngine::setFrame(const FrameImages& images, const MFrameSettings& s, Frame
(finalOutput.format != VK_FORMAT_A2B10G10R10_UNORM_PACK32 || !e.model.downRgb10.handle))
throw std::runtime_error("unsupported borrowed output format or missing RGB10A2 shader");
const bool reset = s.reset || e.frame == 0;
const bool L = e.model.presetL;
const float fx = float(e.hw) / float(e.rw), fy = float(e.hh) / float(e.rh);
e.enc0 = {{int32_t(e.rw), int32_t(e.rh)}, {2, 2}, {int32_t(e.gw / 2), int32_t(e.gh / 2)}, {s.jitter[0], s.jitter[1]}, {e.prevJitter[0], e.prevJitter[1]},
{s.motionScale[0] * fx, s.motionScale[1] * fy}, {s.motionOffset[0], s.motionOffset[1]}, {fx, fy}, {1.f / fx, 1.f / fy}, {float(e.hw), float(e.hh)}, 9.045065879821777f, 0.22356846928596497f,
{s.motionScale[0] * fx, s.motionScale[1] * fy}, {s.motionOffset[0], s.motionOffset[1]}, {fx, fy}, {1.f / fx, 1.f / fy}, {float(e.hw), float(e.hh)},
L ? e.model.lConst[0] : 9.045065879821777f, L ? e.model.lConst[1] : 0.22356846928596497f,
s.depthInverted ? 1u : 0u, reset ? 1u : 0u, e.frame & 7u, 1.f, 0.f, {1.f / float(e.hw), 1.f / float(e.hh)}, {1.f / float(e.rw), 1.f / float(e.rh)}};
e.dec0 = {{int32_t(e.rw), int32_t(e.rh)}, int32_t(e.gw / 2), reset ? 1u : 0u, {s.jitter[0], s.jitter[1]}, {fx, fy}, {float(e.hw), float(e.hh)},
{e.prevJitter[0], e.prevJitter[1]}, 8.f, 0.07863261550664902f, {1.f / float(e.rw), 1.f / float(e.rh)}};
{e.prevJitter[0], e.prevJitter[1]}, L ? e.model.lConst[2] : 8.f, L ? e.model.lConst[3] : 0.07863261550664902f, {1.f / float(e.rw), 1.f / float(e.rh)}};
e.post = {{0, 0}, {int32_t(e.hw), int32_t(e.hh)}, {int32_t(e.rw), int32_t(e.rh)}, {1.f / fx, 1.f / fy}, {s.jitter[0], s.jitter[1]},
{1.f / float(e.rw), 1.f / float(e.rh)}, 2.7530112f, reset ? 1u : 0u, 1.f, 0.f};
{1.f / float(e.rw), 1.f / float(e.rh)}, L ? e.model.lConst[4] : 2.7530112f, reset ? 1u : 0u, 1.f, 0.f};
e.down = {{int32_t(e.hw), int32_t(e.hh)}, {0, 0}, {0, 0}, {int32_t(e.rw), int32_t(e.rh)}, {float(e.rw) / float(e.ow), float(e.rh) / float(e.oh)}, {s.jitter[0], s.jitter[1]}, 1.f};
const uint32_t sw = e.ow / 2, sh = e.oh / 2;
const uint32_t sw = exposureGrid(e.rw, e.ow), sh = exposureGrid(e.rh, e.oh);
const uint32_t bx = (sw + 31) / 32, by = (sh + 31) / 32;
e.exposure = {{0, 0}, {int32_t(e.rw), int32_t(e.rh)}, 1.f / (float(sw) * float(sh)), 1.f, 0.23547087609767914f, 1u | (reset ? 2u : 0u), bx * by, bx, {int32_t(sw), int32_t(sh)}};
e.prevJitter[0] = s.jitter[0]; e.prevJitter[1] = s.jitter[1];
@@ -321,6 +376,18 @@ void MEngine::setFrame(const FrameImages& images, const MFrameSettings& s, Frame
b.post[i] = descriptorSet(dev, b.pool, m.post.descriptors);
}
auto in = b.input[i], post = b.post[i];
if (L) {
bindBuffer(dev, in, 0, e.enc0LP); bindRange(dev, in, 1, m.w8, m.lEnc0[0]); bindRange(dev, in, 2, m.w16, m.lEnc0[1]);
bindBuffer(dev, in, 3, e.dummy); bindBuffer(dev, in, 4, e.dummy); bindBuffer(dev, in, 5, e.skipL); bindBuffer(dev, in, 6, e.tokens);
bindBuffer(dev, in, 7, e.enc0P);
bindImage(dev, in, 8, images.color, e.nearest, images.colorLayout);
bindImage(dev, in, 9, images.motion, e.nearest, images.motionLayout);
bindImage(dev, in, 10, images.depth, e.nearest, images.depthLayout);
bindImage(dev, in, 11, e.mvA[1 - i].view, e.linear); bindImage(dev, in, 12, e.half[1 - i].view, e.linear);
bindImage(dev, in, 13, e.feat[1 - i].view, e.linear);
bindImage(dev, in, 14, e.mvA[i].view); bindImage(dev, in, 15, e.mvB.view); bindImage(dev, in, 16, e.mvC.view);
bindRange(dev, in, 17, m.w16, m.dither16);
} else {
bindBuffer(dev, in, 0, e.enc0P);
bindImage(dev, in, 1, images.color, e.nearest, images.colorLayout);
bindImage(dev, in, 2, images.motion, e.nearest, images.motionLayout);
@@ -330,6 +397,7 @@ void MEngine::setFrame(const FrameImages& images, const MFrameSettings& s, Frame
bindImage(dev, in, 7, e.mvA[i].view); bindImage(dev, in, 8, e.mvB.view); bindImage(dev, in, 9, e.mvC.view);
bindRange(dev, in, 10, m.w8, m.enc0[0]); bindRange(dev, in, 11, m.w16, m.enc0[1]);
bindRange(dev, in, 12, m.w16, m.enc0[2]); bindBuffer(dev, in, 13, e.tokens);
}
bindBuffer(dev, post, 0, e.postP); bindImage(dev, post, 1, images.color, e.nearest, images.colorLayout);
bindImage(dev, post, 2, e.mix.view, e.nearest); bindImage(dev, post, 3, e.gates.view, e.nearest);
bindImage(dev, post, 4, e.cov.view, e.nearest); bindImage(dev, post, 5, e.hist[1 - i].view, e.linear);
@@ -370,6 +438,12 @@ void MEngine::recordFront(VkCommandBuffer cb, VkQueryPool stamps, uint32_t first
const int32_t shifts[2] = {MShift[l] ^ int32_t((e.frame & 1u) * 2u), MShift[l]};
vkCmdUpdateBuffer(cb, e.layerP[l].handle, offsetof(MLayerParams, sx), sizeof(shifts), shifts);
}
if (m.presetL) {
// NGX alternates the horizontal window shift every frame: input stage (2, 2) / (0, 2), expansion (0, 0) / (2, 0)
const int32_t in[2] = {2 ^ int32_t((e.frame & 1u) * 2u), 2}, ex[2] = {int32_t((e.frame & 1u) * 2u), 0};
vkCmdUpdateBuffer(cb, e.enc0LP.handle, offsetof(MLayerParams, sx), sizeof(in), in);
vkCmdUpdateBuffer(cb, e.dec0LP.handle, offsetof(MLayerParams, sx), sizeof(ex), ex);
}
memoryBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, compute, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT);
uint32_t query = first;
auto stamp = [&] { if (stamps) vkCmdWriteTimestamp(cb, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, stamps, query++); };
@@ -379,6 +453,7 @@ void MEngine::recordFront(VkCommandBuffer cb, VkQueryPool stamps, uint32_t first
dispatch(cb, m.exposure[1], e.expSets[1], 1, 1); nextStage(cb);
}
stamp();
// (preset L's input stage covers the same blocks: its shift is at most 2 in x and 2 in y)
dispatch(cb, e.external ? m.inputPlanes : m.input, e.enc0Sets[e.current], (e.gw + 2 + 7) / 8, (e.gh + 2 + 7) / 8);
stamp();
if (e.external) closeStages(cb); else nextStage(cb);
@@ -393,7 +468,11 @@ void MEngine::recordBack(VkCommandBuffer cb, VkQueryPool stamps, uint32_t first)
uint32_t query = first;
auto stamp = [&] { if (stamps) vkCmdWriteTimestamp(cb, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, stamps, query++); };
stamp();
dispatch(cb, e.external ? m.expandPlanes : m.expand, e.dec0Sets[i], e.gw / 8, e.gh / 8); nextStage(cb);
if (m.presetL) // one block more in x on frames whose expansion shift is 2
dispatch(cb, e.external ? m.expandPlanes : m.expand, e.dec0Sets[i], (e.gw + (e.frame & 1u) * 2u + 7) / 8, e.gh / 8);
else
dispatch(cb, e.external ? m.expandPlanes : m.expand, e.dec0Sets[i], e.gw / 8, e.gh / 8);
nextStage(cb);
stamp();
const bool compact = e.compact && e.post.rndOffs[0] >= -0.5f && e.post.rndOffs[0] <= 0.5f &&
e.post.rndOffs[1] >= -0.5f && e.post.rndOffs[1] <= 0.5f;
@@ -417,3 +496,17 @@ void MEngine::recordFrame(VkCommandBuffer cb, VkQueryPool stamps, uint32_t first
recordBack(cb, stamps, stamps ? first + 3 : 0);
}
VkImage MEngine::outputImage() const { return impl->output.handle; }
std::vector<MEngine::DebugResource> MEngine::debugResources() const {
auto& e = *impl;
const uint32_t w = e.rw, h = e.rh;
std::vector<DebugResource> r{{"enc0-params", e.enc0P.handle, VK_NULL_HANDLE, e.enc0P.size, 0, 0},
{"tokens", e.tokens.handle, VK_NULL_HANDLE, e.tokens.size, 0, 0},
{"head", e.dec1Out.handle, VK_NULL_HANDLE, e.dec1Out.size, 0, 0},
{"gates", VK_NULL_HANDLE, e.gates.handle, VkDeviceSize(w) * h * 4, w, h},
{"mix", VK_NULL_HANDLE, e.mix.handle, VkDeviceSize(w) * h * 4, w, h},
{"cov", VK_NULL_HANDLE, e.cov.handle, VkDeviceSize(w) * h * 8, w, h},
{"mvA", VK_NULL_HANDLE, e.mvA[1 - e.current].handle, VkDeviceSize(w) * h * 4, w, h},
{"feat", VK_NULL_HANDLE, e.feat[1 - e.current].handle, VkDeviceSize(w) * h * 4, w, h}};
if (e.model.presetL) r.push_back({"skip", e.skipL.handle, VK_NULL_HANDLE, e.skipL.size, 0, 0});
return r;
}
engine/test_l.py
+145
View File
@@ -0,0 +1,145 @@
#!/usr/bin/env python3
"""Run the engine's own preset L stages on a captured launch and compare them with its outputs.
usage: test_l.py enc0|dec0 CAPTURE_DIR FRAME WORK_DIR KVK [repeats]
enc0: layer_m.comp KIND 3 (features, embedding, block, merge) against rrlite_enc0_4x4_mv*_*: the block output
(dec0's skip input), the merged tokens (the network's input) and the three motion surfaces
dec0: layer_m.comp KIND 4 (expand, block, head) against rrlite_dec0_4x4: its four surfaces
CAPTURE_DIR as for test_enc0_m.py (replay/ before each launch, dump/ after it, filter `rrlite_`). KVK is
vulkan/kvk built with its `bindings` manifest line. The half-resolution previous result (texture 152) is not in the
dumps: enc0 uses WORK_DIR/half-prev.bin if it exists (zeros otherwise), which only matters after the first frame.
"""
import glob, os, re, struct, subprocess, sys
import numpy as np
HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import l_model as lm
import model_enc3 as me
from test_m import pack_layer, defines, token_codes, report
def launch(cap, kernel_re, frame):
"""(sequence number, replay dir, dumped resources by parameter offset) of the frame's launch of the kernel"""
names = sorted({re.match(r'launch-(\d+)-(.*)-arg', os.path.basename(f)).groups() for f in glob.glob(f'{cap}/dump/launch-*-arg*')})
hits = [(int(s), k) for s, k in names if re.fullmatch(kernel_re, k)]
if len(hits) <= frame:
raise SystemExit(f'no launch of {kernel_re} for frame {frame}')
seq, k = hits[frame]
res = {int(re.search(r'-arg(\d+)-', f).group(1)): f for f in glob.glob(f'{cap}/dump/launch-{seq:03d}-{k}-arg*')}
return seq, k, f'{cap}/replay/replay-{seq:06d}-{k}', res
class Replay:
def __init__(self, d):
self.P = open(d + '/args.bin', 'rb').read()
self.allocs = {}
for line in open(d + '/manifest.txt'):
q = line.split()
if q and q[0] == 'alloc':
self.allocs[int(q[1])] = (int(q[2], 16), np.fromfile(f'{d}/alloc-{q[1]}.bin', np.uint8))
def at(self, ptr):
"""the replayed allocation bytes from address ptr on"""
a, b = next((a, b) for a, b in self.allocs.values() if a <= ptr < a + b.size)
return b[ptr - a:]
def q(self, o): return struct.unpack_from('<Q', self.P, o)[0]
def i2(self, o): return struct.unpack_from('<2i', self.P, o)
def f2(self, o): return struct.unpack_from('<2f', self.P, o)
def f1(self, o): return struct.unpack_from('<f', self.P, o)[0]
def compile_stage(name, out, extra=()):
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *defines(name), *extra, f'{HERE}/layer_m.comp', '-o', out],
check=True, stdout=subprocess.DEVNULL)
def run(kvk, spv, manifest, reps):
r = subprocess.run([kvk, spv, manifest, reps, '32'], capture_output=True, text=True)
if r.returncode:
print('FAILED', r.stderr[-800:]); sys.exit(1)
for l in r.stdout.split('\n'):
if 'us per' in l: print(l)
def enc0(cap, frame, work, kvk, reps):
seq, k, d, res = launch(cap, r'rrlite_enc0_4x4_mv(hi|lo)_(hdr|ldr)', frame)
R = Replay(d)
wbuf = R.at(R.q(0))
t = lm.enc0_layer(wbuf)
o, b8, b16 = pack_layer(t, 'enc0l')
rw, rh = R.i2(56); gw, gh = R.i2(16); sx, sy = R.i2(8)
b8.tofile(f'{work}/w8.bin'); b16.tofile(f'{work}/w16.bin'); wbuf[:4096].tofile(f'{work}/noise.bin')
np.zeros(64, np.uint8).tofile(f'{work}/dummy.bin')
open(f'{work}/lp.bin', 'wb').write(struct.pack('<4i72I', gw, gh, sx, sy, *o))
expo = lambda off: float(np.fromfile(res[off], np.float16)[0])
hw, hh = R.i2(248)
open(f'{work}/ep.bin', 'wb').write(struct.pack('<2i2i2i2f2f2f2f2f2f2fff3Iffffff', rw, rh, sx, sy, gw // 2, gh // 2, *R.f2(168), *R.f2(176),
*R.f2(200), *R.f2(208), *R.f2(192), *R.f2(184), float(hw), float(hh), R.f1(268), R.f1(272),
R.P[216], R.P[316], R.i2(312)[0], expo(112), expo(120), *R.f2(256), *R.f2(64)))
half = f'{work}/half-prev.bin'
if not os.path.exists(half):
np.zeros((hh // 2) * (hw // 2) * 4, np.float32).tofile(half)
lines = [f'param {work}/lp.bin', f'grid {(gw + sx + 7) // 8} {(gh + sy + 7) // 8}',
f'tex 1 {res[80]} {rw} {rh} 4 16 nearest', f'tex 2 {res[96]} {rw} {rh} 2 16 nearest', f'tex 3 {res[136]} {rw} {rh} 1 32 nearest',
f'tex 4 {res[128]} {rw} {rh} 2 16 linear', f'tex 5 {half} {hw // 2} {hh // 2} 4 32 linear', f'tex 6 {res[160]} {rw} {rh} 4 9 linear']
for i, off in enumerate((280, 288, 296)): lines.append(f'surf {7 + i} {rw} {rh} 2 16 - {work}/mv-{off}.bin')
lines += [f'lut {work}/w8.bin', f'lut {work}/w16.bin', f'lut {work}/dummy.bin', f'lut {work}/dummy.bin',
f'outbuf {work}/skip.bin {gw * gh * 32}', f'outbuf {work}/merged.bin {(gw // 2) * (gh // 2) * 64}',
f'lut {work}/ep.bin', f'lut {work}/noise.bin',
'bindings 0 8 9 10 11 12 13 14 15 16 1 2 3 4 5 6 7 17']
open(f'{work}/manifest', 'w').write('\n'.join(lines) + '\n')
spv = f'{work}/enc0_l.spv'
compile_stage('enc0l', spv, ['-DLDR=1'] if k.endswith('_ldr') else [])
run(kvk, spv, f'{work}/manifest', reps)
print(f'frame {frame} ({k}, launch {seq}) {rw}x{rh}, pixel grid {gw}x{gh}, shift {sx},{sy}, reset {R.P[316]}')
for off, name in ((280, 'A'), (288, 'B'), (296, 'C')):
a = np.fromfile(res[off], np.float16).astype(np.float64); b = np.fromfile(f'{work}/mv-{off}.bin', np.float16).astype(np.float64)
print(f' motion {name}: identical {100 * (a == b).mean():5.1f}% max |diff| {np.abs(a - b).max():.4g}')
out, merged = R.q(40), R.q(48)
after = np.fromfile(res[40], np.uint8); base = out
ref = token_codes(after, 0, gw, gh, 32)
got = np.fromfile(f'{work}/skip.bin', np.uint8).reshape(gh, gw, 32)
print(' ' + report('block output', got, ref))
print(' ' + report('block output, rows inside the image', got[:rh], ref[:rh]))
mref = token_codes(after, merged - base, gw // 2, gh // 2, 64)
mgot = np.fromfile(f'{work}/merged.bin', np.uint8).reshape(gh // 2, gw // 2, 64)
print(' ' + report('merged', mgot, mref))
print(' ' + report('merged, rows inside the image', mgot[:rh // 2], mref[:rh // 2]))
def dec0(cap, frame, work, kvk, reps):
seq, k, d, res = launch(cap, r'rrlite_dec0_4x4', frame)
R = Replay(d)
t = lm.dec0_layer(R.at(R.q(0)))
o, b8, b16 = pack_layer(t, 'dec0l')
sx, sy = R.i2(8); gw, gh = R.i2(16)
rw, rh = [int(x) for x in re.search(r'surface-(\d+)x(\d+)x', res[128]).groups()]
b8.tofile(f'{work}/dw8.bin'); b16.tofile(f'{work}/dw16.bin')
np.zeros(64, np.uint8).tofile(f'{work}/dummy.bin')
token_codes(R.at(R.q(24)), 0, gw // 2, gh // 2, 64).tofile(f'{work}/din.bin')
token_codes(R.at(R.q(32)), 0, gw, gh, 32).tofile(f'{work}/dskip.bin')
open(f'{work}/dlp.bin', 'wb').write(struct.pack('<4i72I', gw, gh, sx, sy, *o))
open(f'{work}/dp.bin', 'wb').write(struct.pack('<2iiI2f2f2f2fff2f', rw, rh, gw // 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, R.f1(160), R.f1(164), 0, 0))
lines = [f'param {work}/dlp.bin', f'grid {(gw + sx + 7) // 8} {(gh + sy + 7) // 8}']
fmt = {128: (4, 9), 136: (4, 8), 144: (4, 8), 152: (4, 16)}
for off in (128, 136, 144, 152): lines.append(f'surf {off} {rw} {rh} {fmt[off][0]} {fmt[off][1]} - {work}/s{off}.bin')
lines += [f'lut {work}/dw8.bin', f'lut {work}/dw16.bin', f'lut {work}/din.bin', f'lut {work}/dskip.bin',
f'outbuf {work}/dump.bin {gw * gh * 64}', f'lut {work}/dummy.bin', f'lut {work}/dp.bin', 'bindings 0 8 9 10 11 1 2 3 4 5 6 7']
open(f'{work}/dmanifest', 'w').write('\n'.join(lines) + '\n')
spv = f'{work}/dec0_l.spv'
compile_stage('dec0l', spv, [a for a in os.environ.get('L_DEFS', '').split()])
run(kvk, spv, f'{work}/dmanifest', reps)
print(f'frame {frame} (launch {seq}) {rw}x{rh}, pixel grid {gw}x{gh}, shift {sx},{sy}, sigma {R.f1(160):.4g} rho {R.f1(164):.4g}')
for off, name, dt in ((128, 'feature', np.int8), (136, 'gates', np.uint8), (144, 'mix', np.uint8), (152, 'covariance', np.float16)):
a = np.fromfile(res[off], dt).reshape(rh, rw, 4).astype(np.float64); b = np.fromfile(f'{work}/s{off}.bin', dt).reshape(rh, rw, 4).astype(np.float64)
same = (a == b) | (np.isnan(a) & np.isnan(b))
print(f' {name:10s}: identical {100 * same.mean():6.2f}% (per channel ' + ' '.join(f'{100 * same[..., c].mean():6.2f}' for c in range(4)) +
f'), max |diff| {np.nanmax(np.abs(a - b)):.4g}')
if __name__ == '__main__':
stage, cap, frame, work, kvk = sys.argv[1], sys.argv[2], int(sys.argv[3]), sys.argv[4], sys.argv[5]
reps = sys.argv[6] if len(sys.argv) > 6 else '1'
os.makedirs(work, exist_ok=True)
{'enc0': enc0, 'dec0': dec0}[stage](cap, frame, work, kvk, reps)
engine/test_m.py
+9 -1
View File
@@ -76,12 +76,20 @@ def pack_layer(t, name):
if CIN:
assert t['expand_w'].shape == (CIN, 4 * C), t['expand_w'].shape
o[68], o[69] = b8.add(pack8(t['expand_w'])), b16.add(tile(t['expand_b']))
if 'emb_w1' in t:
# preset L's input embedding (layer_m.comp KIND 3): 32 features -> 32 hidden -> C
assert t['emb_w1'].shape == (32, 32) and t['emb_w2'].shape == (32, C), (t['emb_w1'].shape, t['emb_w2'].shape)
o[60], o[61] = b8.add(pack8(t['emb_w1'])), b16.add(tile(t['emb_b1']))
o[62], o[63] = b8.add(pack8(t['emb_w2'])), b16.add(tile(t['emb_b2']))
return o, b8.data(), b16.data()
def defines(name):
C, heads, NPM, CIN = mm.LAYERS[name]
return [f'-DC={C}', f'-DNH={heads}', f'-DKIND={1 if NPM else 2 if CIN else 0}', f'-DXC={NPM or CIN or 16}'] + os.environ.get('M_DEFS', '').split()
kind = {'enc0l': 3, 'dec0l': 4}.get(name, 1 if NPM else 2 if CIN else 0)
# compile_m.py's arithmetic (ZLUDA's lowering); M_TEXL=0 gives the native accuracy kernels' for the network layers
texl = ['-DTEXL=1'] if kind >= 3 or os.environ.get('M_TEXL', '1') != '0' else []
return [f'-DC={C}', f'-DNH={heads}', f'-DKIND={kind}', f'-DXC={NPM or CIN or 16}'] + texl + os.environ.get('M_DEFS', '').split()
def token_codes(buf, off, W, H, C):
engine/test_post_m.py
+5
View File
@@ -10,6 +10,11 @@ cap, frame, work, kvk = sys.argv[1], int(sys.argv[2]), sys.argv[3], sys.argv[4]
K = os.environ.get('POST_KERNEL') or next((k for k in ('rrlite_post_3_1_mvlo_hdr_folded', 'rrlite_post_3_2_mvlo_hdr_folded')
if glob.glob(f'{cap}/replay/replay-*-{k}')), 'rrlite_post_3_1_mvlo_hdr_folded')
TILE = ['-DTILE_W=15', '-DTILE_H=15'] if '3_2' in K else []
# preset L's unfolded kernels (no _folded suffix) run post_m.comp's PRESET_L variant, one tap wider per tier
if '_ldr' in K:
TILE = TILE + ['-DLDR=1']
if not K.endswith('_folded'):
TILE = [d for d in TILE if d == '-DLDR=1'] + ['-DPRESET_L=1'] + (['-DTILE_W=16', '-DTILE_H=16'] if '3_2' in K else ['-DTILE_W=11', '-DTILE_H=11'])
reps = sys.argv[5] if len(sys.argv) > 5 else '1'
HERE = os.path.dirname(os.path.abspath(__file__))
os.makedirs(work, exist_ok=True)