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

feat: add d4r inference engines and Windows parity tooling

Browse Source 
Add native K/M/RR inference engines, HIP kernels, replay and validation
harnesses, and Windows DLSS setup and runtime support. Include GPU
profiling tools, build and packaging scripts, platform integration
patches, and implementation documentation.
This commit is contained in:
ayo committed 2026-10-16 08:54:47 -04:00
1 parent 25493b3198
commit 7c5d3c9157
430 files changed
+76191
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

engine/rr/README.md
+558
View File
@@ -0,0 +1,558 @@
# Ray Reconstruction stages (DLSS-D, NGX feature 13)
d4r's own Vulkan compute shaders for the denoiser in `nvngx_dlssd.dll` 310.7 (presets D and E). One shader per
NVIDIA kernel of the frame; the engine records them in NVIDIA's order:
| stage | NVIDIA kernel | shader | grid (640x360 render, 1280x720 output) |
|---|---|---|---|
| exposure | `cuda_dldn_engine_luma_convert_kernel`, `reduce_sum` x2-3, `auto_exposure_copy_kernel` | `engine/exposure_k.comp` (MODE 1) | |
| enc0 | `cuda_dldn_engine_swin_enc0_kernel` (D: `..._enc0_diamond_wallaby_kernel`) | `enc0.comp` | 81x49, 32 threads |
| enc1-enc5 | `cuda_dldn_engine_swin_enc1_kernel` ... `enc5` | `enc1.comp` ... `enc5.comp` | |
| dec4-dec1 | `cuda_dldn_engine_swin_dec4_kernel` ... `dec1` | `dec4.comp` ... `dec1.comp` | |
| dec0 | `cuda_dldn_engine_swin_dec0_kernel` (D: `..._dec0_diamond_wallaby_kernel`) | `dec0.comp` | 81x49, 32 threads |
| output | `cuda_dldn_engine_hkpn_output_kernel_transformer` (D: `..._diamond_wallaby`) | `hkpn.comp` | 80x45, 16x16 threads |
The D variants of enc0, dec0 and the output kernel are different code; the other nine kernels are shared and only
their weights differ between presets.
## RDNA3 exact preset D (2026-10-14)
Build with `D4R_RR_GFX11=1 D4R_RR_RELAXED=0 engine/rr/build_native_model.sh CAPTURE OUT DLL`.
`rr_gfx11.glsl` reproduces ZLUDA's gfx11 FP8 widening, K positions, scaled f32 accumulators and final binary16
rounding using f16 cooperative matrices. The weight files include a widened copy. Relaxed and mix arithmetic are
unavailable in this build; `VK_EXT_shader_float8` is unnecessary. `kvk`, `evk` and `d4r-rr` enable float8 only when
the driver advertises the extension, because RADV gfx11 may report its feature bits without the extension.
`arithmetic gfx11` in `rr.txt` also selects NGX 310.7's CPU camera operation order. Its inversions, relative camera
translation and intermediate reprojection products leave small residues that affect specular motion history;
double-precision equivalents were insufficient for temporal equality. Other model builds keep their existing path.
RX 7700 XT (gfx1101), RADV, preset D, AW2 town at 853×480 → 2560×1440: enc0 and hkpn are byte-identical on all three
captured launches, joining the ten previously verified stages. Encoder surface 112's undefined alpha is excluded
from the isolated stage check. The complete presented RGBA16F output, including alpha, is byte-identical over all
24 replay frames in two engine runs versus the deterministic CUDA control. Engine GPU time is about 11.4 ms versus
27.9 ms for the current control; warmed samples varied by about 0.1 ms within each engine run. Older CUDA runs
included timing outliers, so their lower 26.7 ms figure is historical, not the paired control. No frame-age or
resolution change contributes to the result. Other RDNA3 GPUs, preset E, and other games require validation.
AW2 live at the same setting: CUDA / engine denoiser GPU medians 27.3 / 12.0 ms; frame intervals 42.6 / 24.4 ms
in the faster sessions and 53.4 / 36.2 ms in the slower sessions. These are separate session states, one launch per
backend in each state; the engine does not resolve the pre-existing slower game state. Frame intervals use the
frame-count/time slope over the final 30 seconds, with sampled instantaneous spreads of about 1 ms for CUDA and
3-4 ms for the engine. The same-frame engine is opt-in; configuration was restored after testing.
Put the local model under `d4r/engine/rr-d`, or set `D4R_ENGINE_MODEL_RR_D_DIR` to its Windows path. The shim runs it
with `[RayReconstruction] Backend = dlss` and `[Engine] Enabled = true`. `D4R_ENGINE_RR_PROFILE` log lines give
game-queue GPU time around the complete engine evaluation; `gpu_frame` identifies the earlier frame whose query
has completed, while `game_interval` is the current call interval. Query readback does not delay the image.
## Stage contract
Every stage shader reads the same data NVIDIA's kernel reads, in the same layouts, so each one can be checked on
its own against a captured launch (`test_layer.py`) and the stages chain exactly like NGX's kernels:
- binding 0: the parameter block, `std430` readonly. It is NVIDIA's argument block byte for byte, except that every
pointer into the plane arena holds its **byte offset in the arena** (a 64-bit value), the weight pointer is zero,
and a texture or surface handle is 1 where NGX bound one and 0 where it left the argument null (the kernels test
optional inputs' handles for null). Fields are read at NVIDIA's argument offsets.
- binding 1: the layer's weight image exactly as NGX uploads it (argument 40's allocation), `std430` readonly.
- binding 2: the plane arena, `std430` read/write.
- bindings 3, 4, ...: the textures in ascending argument offset (combined image samplers, clamp-to-edge,
normalized coordinates, linear or nearest as NVIDIA's texture object; 8-bit unsigned arrays are sampled as UNORM,
`SNORM_INT8X4` as SNORM), then the surfaces in ascending argument offset (storage images, `rgba16f` or
`rgba8_snorm`).
- The dispatch grid is NVIDIA's launch grid. The workgroup size is the shader's choice.
## Checking a stage
Captures (bridge launch and replay dumps of a harness run, see `rr_capture.py`) live outside the repository.
```sh
glslangValidator -V --target-env vulkan1.3 engine/rr/enc1.comp -o build/rr/enc1.spv
python3 engine/rr/test_layer.py CAPTURE swin_enc1 2 build/rr/enc1.spv build/rr/work-enc1 [--time 50]
```
It starts from the arena as it was before NVIDIA's launch, with zeroed surfaces, and compares every surface and
the whole arena with what NVIDIA's launch left. `build/rr/kvk` is `vulkan/kvk.cpp`
(`g++ -O2 -std=c++17 -o build/rr/kvk vulkan/kvk.cpp -lvulkan`).
## Native stages (preset E)
`build_native.py CAPTURE MODEL_DIR` turns a `compile_rr.py` model directory into one where every Swin stage is
d4r's own cooperative-matrix shader instead of translated PTX. Each was derived phase by phase from ZLUDA replays
(`phase_dump.py`: every mma's operands of one window) into an exact numpy reference, then written against it:
| stages | reference | shader | weights |
|---|---|---|---|
| enc1-enc5, dec4-dec1 | `swin_ref.py` (one generic model: C, heads, merge / expand) | `swin_m.comp` | `swin_pack.py` |
| dec0 | `dec0_ref.py` | `dec0_m.comp` | `dec0_pack.py` |
| enc0 | `enc0_ref.py` (after the per-pixel features) | `enc0_m.comp` from `gen_enc0_m.py` | `enc0_pack.py` |
`enc1_m.comp` / `dec1_m.comp` are the hand-written C = 64 instances `swin_m.comp` generalises; `enc1_ref.py` holds
the shared numerics (inter, mm, the cubic softmax, the norm trees). enc0 keeps ptx2rr's translation of NVIDIA's feature
gather (waves 0-2 run its pixel loop) and is native from the embedding on.
Binding 1 differs from the stage contract above: it holds the packed weights (enc0: NGX's image, then the packed
blob). Everything else is the contract, so `test_layer.py --weights FILE` checks a native stage the same way. Against
capE3 every native stage leaves the arena byte-identical to NVIDIA's; surfaces differ in a few values by one binary16
ulp (`exp2` approximation, and decoder texture 88 where NVIDIA's bilinear unit rounds differently), and the whole
network's output is ~96 dB from NVIDIA's.
The shaders run every product transposed (weights x activations) so results feed the next product in registers;
`enc1_m.comp`'s header explains the layouts and `rr_exact.glsl` the binary16 rounding RADV needs (it folds
`float(float16_t(x))` and does not fuse f16 `fma`).
## Native models and preset D
`engine/rr/build_native_model.sh CAPTURE OUT_DIR [NVNGX_DLSSD_DLL]` builds a model whose eleven network stages are
the native shaders (`swin_m.comp`, `dec0_m.comp`, the generated `enc0_m.comp`) for the preset of the capture, with the
translated output stage.
The weights are read from the user's DLL itself (`dll_weights.py`), not from the capture. The DLL holds each model's
tensors as plain binary16 arrays in natural `[out][in]` order, stage after stage; nothing names them, so they are
found by structure: every attention head has a 100 x 100 relative position bias table whose entries depend only on
the offset between two window positions, and a model is eleven consecutive stages with 2 2 4 4 8 8 8 4 4 2 2 heads.
What NGX uploads is the same data rearranged: FP8 codes in the kernels' fragment order, the norm gains times
sqrt(C), the bias times 2 sqrt(2) for the kernels' 64 x 96 query / key order, and (preset E) the 4096-byte noise
table encoder 0 dithers with. `dll_weights.py DLL OUT --model N --check DIR` compares a reading with a phase-dump
extraction: for 310.7 every tensor of both presets is identical, and so are the packed weight files. The capture
only selects which of the DLL's models it ran (and supplies the two argument-block templates); a DLL whose network
kernels are not the ones the shaders were written against (`native_kernels.txt`, SHA-256 of the PTX modules) is
refused unless `D4R_RR_ANY_KERNELS=1`.
`compile_rr.py` uses the translated output kernel for both presets: `engine/rr/hkpn.comp` has the reset path only
(`--reset-only` selects it for captures that reset every frame).
Preset D (`..._diamond_wallaby` kernels) differs from E in three places:
- **enc0**: another feature gather (kept as the translation, like E's) with 30 input features for E's 32, and no
noise table (its weight image is 4096 bytes smaller).
- **dec0**: the eight bytes per pixel for the output stage are channels 12, 13, 0, 18, 19, 20, 21, 0
(`-DRR_PRESET_D`; E: 12, 13, 18-23). Everything else is E's kernel.
- **output stage**: its own kernel (closest-depth motion from texture 296, a 4x4 bicubic history fetch). It is the
translated stage; it has no subgroup operations and runs with 64-wide waves (`rr.txt`).
The reset flag (enc0 byte 365, output stage byte 440) follows `Reset`, not the jitter. With `Reset = 0` the output
stage reprojects and blends its history: check stages against a temporal capture (the harness with
`D4R_HARNESS_JITTER_SCENE=1`, and `D4R_CUDA_REPLAY_DUMP_LIMIT=400` so every launch is saved), not only against
captures that reset every frame. `engine/rr/hkpn.comp` implements the reset path only.
Preset D, RX 9070 XT, RADV 26.2.4, 2026-10-12, 1280x720 -> 3840x2160 (`d4r-rr`, `D4R_RR_TIME_REPEAT=20`): 5.0-5.2 ms
per frame (network 4.5, output 0.55-0.65), from 17.5 ms with every stage translated. Every network stage reproduces
the capture's arena byte for byte on reset and temporal frames. The presented image against the CUDA path over
temporal frames: 96-106 dB on the reset frame, then 41-52 dB (one-ulp differences in the filtered surfaces are
amplified by the recurrent feature; twelve frames stay at 41-45 dB).
RE Requiem (2026-10-12, 1280x720 -> 3840x2160, R9G9B9E5 colour, roughness packed in the normals' w): with
`hkpn.comp` the picture was 3-pixel blocks that shook with the jitter (no history); with the translated output stage
presets D and E both give a stable, reconstructed image on the engine. Not compared with the CUDA path in that game.
## Matching the reference bit for bit (2026-10-12)
The reference is the CUDA path through ZLUDA on the same GPU, so its `rcp.approx`, `ex2.approx` and bilinear texture
fetches are the GPU's own instructions and texture unit, not correctly rounded arithmetic. Two places in the native
filters had been "corrected" to exact arithmetic and were one ulp off in ~0.01-0.2% of the filtered surfaces:
- `rcp16` (the filters' softmax normalisation and gate) is now the plain hardware reciprocal;
- the decoders' texture 88 goes through the bilinear sampler at NVIDIA's coordinates (`rr_decfilter.glsl`,
`dec0_m.comp`) instead of a hand-computed 1/4, 3/4 blend.
With both, every network stage (arena and surfaces 112 / 120 / 128 / 144 / 304) is identical to the capture on every
frame of capDT12, capD4K, capD1080 (preset D) and capE3 (preset E). What remains is the translated output stage:
~0.004% of surface 224 / 240 one ulp off (single channels, values just above 0.5; cause not found). The output is
107-112 dB on frames where the network's inputs are still exact. At 640x360 -> 1080p that holds for all twelve
frames of capDT12 (was 96 dB, then 41-45); at 1280x720 and 1920x1080 -> 4K the history's differing values reach
encoder 0 through texture 496 and the output falls to 52 / 45 / 43 dB (capD4K), 53 / 49 dB (capD1080) as before.
Speed: the compiled stages spend about half their instructions on float conversions. `spv_exact_convert.py` marks
the conversions NoContraction so RADV keeps a binary16 rounding without the opaque OR (`RR_EXACT_CONVERT`), which is
bit-exact and ~8% faster: preset D 4.8 ms at 1280x720 -> 4K, 9.5 ms at 1920x1080 -> 4K (was 5.1 / 10.3).
`D4R_RR_RELAXED=1|2` builds drop NVIDIA's double rounding instead (8.2 / 7.6 ms at 1080p) but are not exact: about 1%
of a layer's FP8 codes change and the output is 39-44 dB from the reference on every frame.
Profile of a full-resolution stage at 1080p (dec0, 3.0 ms): expand 0.31, key norm 0.16, attention 1.37 (per head:
projections 0.15, scores + softmax 0.26, P V 0.17, out-projection 0.10), MLP 0.69, output filter 0.50. The 136
matrix products themselves are 0.45-0.5 ms; weight and table memory traffic is under 8%. enc0's translated feature
gather is 0.28 of its 3.0 ms.
## Optional inputs: specular hit distance and disocclusion mask (2026-10-12)
RE Requiem passes `DLSSD.SpecularHitDistance` (R32_FLOAT) and `DLSS.DisocclusionMask` (R16_FLOAT). Harness captures
with synthetic ones (`D4R_HARNESS_RR_EXTRA=hit,mask`; capDX, capDXhit, capDXmask) show the wiring:
| input | encoder 0 texture | encoder 0 flag byte | output stage |
|---|---|---|---|
| specular hit distance | argument 248 | 536 | flag byte 160 |
| disocclusion mask | argument 296 | 539 | - |
Encoder 0 also tests the handle at argument 248 for null before it takes the hit distance path (it writes a
reprojected specular motion into channels 2 and 3 of surface 304), which is why the parameter block now carries the
handles' presence. With both bound the network is identical to the capture on every frame (capDX). The output
stage then differs from the CUDA path on the image's border pixels only (63-67 dB on temporal frames): NVIDIA's
kernel fetches neighbours outside the image there, CUDA clamps such integer fetches to the edge (every denoiser
texture's address mode is clamp; the translation and rrsim do the same), and ZLUDA's image load returns zero. With
zero outside the image the stage matches the capture like before (99.997%); the engine keeps CUDA's clamp.
Exact speed-ups since the first native build, all bit-identical per stage: conversions kept by NoContraction instead of
an opaque OR, the MLP activation and its FP8 conversion from a 64 KiB table (`enc1_pack.act_table`), FP8 conversions in
pairs, no saturation clamp on attention probabilities. Preset D: 4.4 ms at 1280x720 -> 4K, 8.6 ms at 1920x1080 -> 4K.
Output stage, hypotheses tested for its remaining ~0.004% one-ulp values and rejected (2026-10-12): correctly rounded
exp2 / reciprocal (rrsim's exact arithmetic is further from the capture than the GPU), fusing the kernel's 34
single-use multiply-subtract pairs into fma, merging the two constant multiplications after the log transform,
`precise` on the inlined exp2, NoContraction on its conversions. The differing values come in spatial clusters, one
colour channel at a time, always one ulp, mostly just above 0.5 in surface 224.
## HIP against Vulkan, and the fused rounding (2026-10-12)
Encoder 1 ported line for line to a standalone HIP kernel (`kernels/rr/swin_native.h`, runner
`kernels/tools/rr_swin_native_test.cpp`; capD1080 frame 2, 5 launches per round): 835 us against the Vulkan stage's
770-788 us, the arena byte-identical. A HIP port of the network is not worth making. What the port did show:
- The FP8 product alone costs 0.61 us per product per wave on this grid, 2.07 us when each is followed by the
f32 -> binary16 -> f32 rounding of its eight accumulator registers. Products and that rounding are about half of
a stage on either backend.
- Keeping the accumulator as binary16 and taking f16(P + acc) with `v_fma_mix` (as `kernels/m` does) halves the
rounding's instructions: HIP 780 us. RADV emits the same instruction for `float16_t(p + float(acc))`, so the
Vulkan stages have it as a build option: `D4R_RR_MIX=1 build_native_model.sh ...` (`RR_MIX`, the shaders' K32).
`RR_MIX` is one rounding of the exact sum where the reference has two, so it is not bit-exact by construction.
Preset D: every network stage still identical to the capture on all 21 frames of capDT12 / capD4K / capD1080 / capDX
/ capDXcp / capDXhw (the same figures as the default build), 4.22 ms at 1280x720 -> 4K and 8.3 ms at 1920x1080 -> 4K
(default 4.45 / 8.7). Preset E (capE3): frames 1 and 3 have a value of encoder 0 one step off (127 dB), which the
later stages amplify to 41 dB at the output; frame 2 is identical. The default build stays the exact one.
## Cyberpunk 2077: more inputs and settings (2026-10-12)
Found on harness captures (`D4R_HARNESS_RR_EXTRA=particles,sss`, `D4R_HARNESS_RR_HW_DEPTH=1`, `D4R_RR_PRESET=N`):
| what the game passes | where NGX puts it |
|---|---|
| `DLSSD.ScreenSpaceSubsurfaceScatteringGuide` (R16F) | encoder 0 texture 264 (its handle is the only switch) |
| `DLSSD.ColorBeforeParticles` (RGBA16F) | output stage texture 416, flag byte 432, and a second output-size history: surface 384, read back next frame as texture 392 (linear) |
| `DLSS.Use.HW.Depth = 1` | encoder 0 byte 980, output stage byte 443 |
| preset hint | 5 runs preset E; every other value (0-4, 6, 0x80000002, 0x80000005) runs preset D |
Integer texture fetches outside the image return zero (`fetch0`), which is what the reference does; CUDA would clamp.
With Cyberpunk's set (hit distance, particles, subsurface guide, hardware depth) the engine's presented image is
byte-identical to the CUDA path through the shim on the frames the harness can compare (640x360 -> 1920x1080), and
with RE Requiem's set the network is identical on every frame (output 107-109 dB). At 1280x720 and 1920x1080 -> 4K
without these inputs the output still falls to 43-53 dB after the first frame (the output stage's one-ulp values);
a correctly rounded log2 there matches the capture worse than the GPU instruction does.
## The output stage's one-step differences: fused multiply-convert (2026-10-13)
The rare one-binary16-step values of the translated output stage (above: ~0.004% of surfaces 224 / 240, "cause not
found") were the driver fusing an f32 product (or sum) with the conversion that follows it into `v_fma_mix`, which
rounds the exact result once. The kernel rounds twice, to f32 and then to binary16, and the two differ when the f32
result is exactly a binary16 midpoint, as `log2(x) * c1 * c2` of the output stage's tone curve sometimes is.
`ptx2rr.py` now ORs an opaque zero into the f32 bits before every `cvt.rn.f16.f32` (and into the binary16 bits
after, against the fold of `float(float16_t(x))`). With it presets D and E are byte-identical to the CUDA path on
every frame of every capture (capDT12, capD4K, capD1080, capDX, capDXcp, capDXhw, capE3; history and output), at
+0.05 ms: 4.27 / 4.45 ms at 1280x720 -> 4K. Found with `build/rr`-style bisection: probe every register of a block
through ZLUDA (debug stores added to the PTX, `zluda_replay.py`) and through the translated shader
(`ptx2rr.py --probe`) and list the first that differ.
Other reference behaviours the translation now follows, all found the same way on preset F's kernels:
- formatted stores to normalized surfaces round half away from zero in f32 (ZLUDA's `surface_formatted_store_bits`);
UNORM 10:10:10:2 arrays are RGBA32F images holding the quantised values `q / 1023` (the correctly rounded quotient,
from a table: the driver divides by multiplying with a reciprocal);
- `tanh.approx.f32` is OCML's `tanhF` with LLVM's AMDGPU lowering of `exp` (`lowerFEXP`, fast-FMA path): the same
operations, the fused ones through f64 (the driver's f32 `fma` is a product and a sum whenever the product also
exists on its own); `tanh.approx.f16` is the correctly rounded value (a table);
- the driver's GLSL `fma` on binary16 is a true fused operation, as the reference's is: no emulation there;
- texture fetches with a binary16 result keep that rounding.
## Preset F: DLSS 4.5 Ray Reconstruction (2026-10-13)
`nvngx_dlssd.dll` 310.9 has a third model with its own kernels, `rr2_*`, and runs it for every preset hint but 4 (D)
and 5 (E): none, 6, Cyberpunk 2077's 0x80000002. The engine runs it from a `D4RR0002` model:
```sh
python3 kernels/tools/extract_dlss_ptx.py NVNGX_DLSSD_DLL ~/.cache/d4r-rr-corpus-f # its PTX (zstd in 310.9)
python3 engine/compile_rr2.py CAPTURE OUT_DIR --native NVNGX_DLSSD_DLL [--particles CAPTURE_WITH_PARTICLES]
```
The runtime is `engine/runtime_rr2.inc` (behind `RModel` / `REngine`); the shim uses `engine\rr-f` (or
`D4R_ENGINE_MODEL_RR_F_DIR`) and, when that model exists, NGX 310.9's hint mapping.
A frame, as NGX launches it:
| pass | NVIDIA kernel | engine |
|---|---|---|
| exposure | `rr2_histogram_auto_exposure_basic_kernel` | `exposure_f.comp` (two passes) |
| enc0 | `rr2_enc0_kernel` | translated feature gather and guide tail around the native block (`gen_rr2_enc0.py`) |
| enc1 - enc5 | `rr2_enc1_kernel` ... | `rr2_m.comp` |
| dec4 - dec1 | `rr2_decN_kernel_0`, `rr2_decN_kernel_1` | `rr2_m.comp` |
| dec0 | `rr2_dec0_kernel_0`, `rr2_dec0_kernel_1` | `rr2_m.comp`; the second: native block, translated output tail |
| output | `rr2_post_kernel`, `rr2_downsample_kernel_static_hdr[_no_transparency]` per tile, `rr2_pooled_history2_kernel` | translated (`ptx2rr.py`) |
The network (`rr2_ref.py` is its exact numpy reference and documents the block): D/E's levels, window grids and
arithmetic (FP8 operands, binary16 sums rounded per 32-deep step, the cubic softmax, the MLP activation), with a
different block. Each kernel is one block on 8x8-token windows that alternate their shift per frame: heads that
share one projection for queries and keys (so the scores are symmetric; bias `(T + T^T) sqrt(2)` of a 64 x 64
relative position table), all reading the block's input and summing their out-projections; MLP branches likewise;
both sums enter the residual times a binary16 scalar from the argument block. Encoders merge 2x2 for the next level,
a decoder's first kernel expands the level below and adds the skip times a third scalar. Next to the tokens every
level carries a three-channel guide image: encoders gate it (`sigmoid(head) * s0`) and average it 2x2 for the next
level, decoders upsample it bilinearly, gate it and blend it with the encoder's. All stages exchange data through
the arena (3152 units of W0 * H0 / 32 bytes; NGX's region offsets, which scale with the render size), tokens as FP8
planes of 32 channels, guides as four binary16 per token.
Weights: the DLL holds the model as plain binary16 tensors in natural `[out][in]` order, kernel after kernel
(`rr2_ref.dll_weights`; found by the sixteen kernels' relative position tables, 2 2 4 4 8 8 8 8 4 4 4 4 2 2 2 2
heads). `rr2_pack.py` packs them for `rr2_m.comp`; nothing of the network comes from the capture. The capture still
supplies every kernel's argument block (the runtime keeps their constants and fills the geometry, the frame's
values and the pointers) and encoder 0's second table.
The output stage filters an intermediate image of 1.5 times the output size in tiles (columns of at most 720, rows
of at most 1620 intermediate pixels, 4 pixels of overlap inside), reduces each tile to the output, and keeps two
histories for the next frame: the intermediate image, and a half-size copy whose quarter-size reduction
(`rr2_pooled_history2_kernel`) encoder 0 also reads (UNORM 10:10:10:2 arrays, see above). NGX runs filter and
reduction concurrently with progress counters; the kernels skip that when the counter pointers are null, which is how
the engine runs them, in order. They have no matrix products and run with 64-wide waves (their warp-level shuffles
address the lane's own half).
Options (harness captures capFX*): encoder 0 byte 466 packed roughness, 568 inverted depth, 569 / texture 272
`DLSSD.SpecularHitDistance`, 571 / texture 336 `DLSS.DisocclusionMask`, texture 288 the subsurface guide, 1028
`DLSS.Use.HW.Depth`; `DLSSD.ColorBeforeParticles` selects the other reduction kernel (model stage `downp`, texture
208, an output-size history 256 <-> 264, flag bytes 196-198). Encoder 0's surfaces 376 / 384 are the output filter's
textures 64 / 72. Pre-exposure does not reach these kernels.
Checked (`d4r-rr` on `rr_inputs.py` inputs; every stage surface and the presented image against the CUDA path):
byte-identical on every frame of capF (640x360 -> 1080p, 3 frames), capF900 (800x450 -> 1600x900), capF4K
(1280x720 -> 4K, 16 tiles) and of the option captures (motion, a real camera, inverted and hardware depth, packed
roughness, hit distance, disocclusion mask, subsurface guide, particles); through the shim (D3D12 harness under
Proton) the engine's frames are byte-identical to the CUDA backend's too (`D4R_RR_PRESET=6 D4R_RR_DLL=<310.9 DLL>
scripts/run_rr_engine_compare.sh`; that script still compares frame N with frame N, and the harness presents the
CUDA result one frame late in the jitter scene).
Speed, RX 9070 XT, 1280x720 -> 3840x2160: 6.2 / 6.6 ms per frame (exposure 0.03, network 4.1 - 4.2, output 2.1 -
2.4), from 14.5 ms with every kernel translated; preset D is 4.3 / 4.5. Per stage (us): enc0 1076 (its block 521),
dec0_1 775, dec0_0 630, enc1 340, dec1_0 316, dec1_1 269, the rest under 180; filter and reduction about 130 per
tile, 16 tiles. Not done: native output passes and feature gather.
### Where the translated parts' time went (2026-10-13, later)
Now about 5.8 / 6.2 ms (5.7 / 6.1 with the 10-bit histories that were later withdrawn), still byte-identical on every capture and through the
shim. Profiling the translated code (cutting it at points behind an opaque condition, `D4R_ENGINE_RR_SKIP=post,down,
pooled`) showed its arithmetic is spread evenly - no hot spot a rewrite would remove - and that the losses were in
how the passes were arranged:
- **Encoder 0 in two passes** (-0.2 ms). One shader ran NVIDIA's 32-lane gather on one of the block's four waves
while the other three waited: 0.40 ms for the gather, 0.24 ms as a pass of its own in 32-thread groups.
`gen_rr2_enc0.py --split gather | block` writes both; the gather pass (`stage enc0g`, the same bindings and
parameter block as enc0) stores each window's record in the arena buffer (see Memory below): the first product's A
fragments (1024 words) and, per lane, the values the tail reads from the gather (six binary16 on the reference
build; the generator finds them and recomputes the tail's index arithmetic). The record's address is a word
compile_rr2.py appends to enc0's argument block. Cost: 76 bytes per render pixel of arena.
`D4R_RR2_SPLIT=0` builds the single shader.
- **Tile passes in overlapping groups** (-0.15 ms). The 16 tiles' post and down passes ran one after another, each
behind a barrier, and the gaps between 32 short dispatches added up. Each tile now has a slot (its two regions;
eight slots, 58 MB at 4K) and a group's post passes run without barriers between them, as do its down passes
together with the next group's post passes on the other half of the slots - what NGX does with two slots and its
progress counters. `D4R_ENGINE_RR_SKIP=serial` restores the old order.
- **Pooled histories as real 10:10:10:2 images: withdrawn.** 4 bytes a pixel instead of RGBA32F's 16 saved
0.1 ms and was identical at 3x and 2x scaling, but the hardware's bilinear filter rounds a UNORM image differently
from a float one: at 1.5x (capFQ) temporal frames came out at 44 - 51 dB. The default is RGBA32F again;
`D4R_RR2_POOLED=rgb10_a2` builds the other form.
Tried without gain: 32-bit arena accesses for the translated code's paired 16-bit ones (0.01 ms), dropping the
subgroup shuffles or the depth search (they cost nothing measurable). What is left is NVIDIA's arithmetic: gather
0.24 - 0.33 ms, decoder 0's output tail 0.18, post 1.1 - 1.3, down 0.7.
### Other scaling ratios, and what a game-like configuration showed (2026-10-13, later)
Asked whether instability seen in RE Requiem came from this implementation, the engine was compared with the CUDA
path through the shim in configurations the earlier captures did not cover (`D4R_RR_PRESET=6 D4R_RR_DLL=<310.9>
scripts/run_rr_engine_compare.sh`, frames aligned by one):
- **Render widths that are not a multiple of 32 were wrong** (1114x626 -> 1080p: 15 dB): the post pass's argument 8
is the row stride of the network's output, the padded width W0, not the render width. Fixed; identical now.
- **Below 4/3 scaling NGX runs another output filter**, `rr2_post_kernel_3_2` (1456x820 -> 1080p and larger render
sizes, DLAA included; 1440x810 and smaller use `rr2_post_kernel`). The model has no such stage, so the shim hands
those features to the CUDA backend. `rr2_histogram_auto_exposure_histogram_kernel` was never launched in any run.
- **DLSS.Roughness.Mode 1 with a roughness texture bound**: NGX follows the mode, the shim followed the texture.
Fixed for preset F (RE Requiem binds none, so it was not affected).
- 1.5x, 1.72x, 2x and 4/3 scaling, and RE Requiem's inputs together (inverted and hardware depth, hit distance,
mask, packed roughness, scaled motion vectors, a moving scene): identical.
- **Open: the camera matrices are not NGX's bit for bit.** NGX computes them in single precision: argument 644 is
the f32 product view x clip, 836 its inverse by cofactors times 1 / det (glm::inverse's operation order; both
reproduced bit for bit in build/rr/f/matfit.py), 900 is an inverse of the pixel matrix (it carries that inverse's
residues, and 0.9999999 for 1 at 1114x626) and 964 was not reproduced. The engine computes them in double and
rounds, a last-bit difference in a few elements. With a non-identity camera, hit distance and mask bound and a
moving scene that showed as 40 - 50 dB on later frames in a synthetic scene whose affected feature is rounding
noise around zero in both paths; no other configuration was affected.
### Memory (2026-10-14)
At 1280x720 -> 4K the engine needs about 600 MB for preset F; the two 1.5x output histories (2 x 149 MB, NGX has
the same) are half of it. NVIDIA's DLL, its allocations traced through the CUDA path and replayed on an RTX PRO 4500
Blackwell, takes 526 MiB of feature state, 106 MiB of input / output arrays and 2 MiB of code, plus 262 MiB of CUDA
context. Three cuts, every capture still byte-identical (`verify.sh`, `verify_x.sh`) and the shim's frames
byte-identical to the previous build, same speed:
- **The tile slots share the arena with the network's regions.** Every arena region is written and read within one
frame; in the output stage only the network's output (dec0_1's argument 64) and something after it are read, so
the slots that fit before the network's output go there (all 8 at 720p -> 4K) and the rest after the arena. A
first try that put the slots right after the network's output broke 360p -> 1080p's output stage.
- **The gather records' rows of windows fill the gaps between encoder 0's regions**: the two largest (units
1216 - 2240 and 2304 - 2560) and then from the end of its last region (unit 2816) on, instead of one piece after
the tile slots. The records (76 bytes a render pixel, about 2450 units) are larger than all the gaps, so 46 of 93
rows at 720p still reach past the arena (36 MB). Models from compile_rr2.py mark this with "s3" on the enc0g line
and append six words to enc0's block (three ranges' addresses, the first two's rows; `gen_rr2_enc0.py`); older
models keep the one-piece layout. Same speed as one piece (four alternating runs at 720p -> 4K).
- **The engine's own output image is created only for frames that do not write the game's texture**, and the
weights go into the command buffer with `vkCmdUpdateBuffer` instead of through a host-visible staging copy that
stayed allocated.
Through the shim (direct output) the harness process peaks 159 MiB lower (1100 -> 941 MiB). Not exact, so not
done: the pooled histories as RGB10A2 (`D4R_RR2_POOLED=rgb10_a2`, -70 MB, see above); they hold q / 1023, which
neither binary16 nor UNORM16 represents exactly. `D4R_ENGINE_RR_ARENA`
keeps the old separate layout, so that arena dumps show the network's regions as NGX leaves them.
## Against a real NVIDIA card: `D4R_RR_FUSE` (2026-10-13)
Every reference above is the CUDA path through ZLUDA. NVIDIA's own kernels (310.9.1's `rr2_*` PTX, replayed through
plain CUDA on an RTX PRO 4500 Blackwell, driver 580; tools and logs in `build/rr/f/nv/`) give an image 39-42 dB from
that reference on every frame of capF, capF4K and capFQ. Register bisection between the two (`probe_pack.py`,
`probe_cmp.py`: a debug store after every instruction, run on both) found why:
- **Fused multiply-adds.** NVIDIA's compiler contracts a `mul.f16` / `mul.f16x2` whose product only feeds an
`add` / `sub` into one fused multiply-add; the kernels give those instructions no rounding modifier. ZLUDA runs them
as written. `ptx_contract.py` is the rule (single-use product, same block, the first operand when both are
products): with it ZLUDA's result of every network kernel matches the card's (0 of ~5900 4 KiB chunks for most
launches, a handful where an `ex2.approx` / `div.approx` differs by one f32 step). The FP8 products, the e4m3
conversions and every unfiltered texture fetch were already identical. The sites are the residual steps
`out = in + branch * s`, the MLP activation `x (0.5 + y (0.412 - 0.081 |y|))` and the guide's bilinear blend.
Probing a product stops the card from fusing it: compare unprobed results, or probe everything.
- **Filtered fetches of 8-bit textures.** A filtered fetch of the SNORM8 normals returns this card's value (8-bit
weights) rounded to a 16-bit step: `f32(roundEven(v * 32767) / 32767)`, exact on 256 of 256 probed values. The
output filter's UNORM8 fetches differ too; no model for those yet.
- **Approximate instructions.** `lg2.approx` near 1 differs by ~8e-8 absolute, `ex2.approx` / `rcp.approx` /
`div.approx` by one f32 step. Not emulated.
`D4R_RR_FUSE=1 compile_rr2.py CAPTURE OUT --ptx-dir <build/rr/f/nv/mkptx.sh output> --native DLL` builds the model
with the first two: `rr2_m.comp -DRR_FUSE` (`MADD` / `MADD4`), the activation table from `enc1_ref.act` with
`FUSE`, the translated parts from the rewritten modules, encoder 0's normals through `ptx2rr.py --snorm16 528`. All
14 native stages are byte-identical to ZLUDA running the rewritten modules (`build/rr/f/nv/nfuse.sh`), and so is
the whole model on capF's three frames. Presented image against the NVIDIA card, 640x360 -> 1080p:
| frame | default build | `D4R_RR_FUSE` |
|---|---|---|
| 1 (reset) | 40.9 dB, 27.3% of values equal | 85.6 dB, 99.5% equal |
| 2 | 39.2 dB | 43.0 dB |
| 3 | 39.3 dB | 41.8 dB |
Temporal frames stay near 42 dB: the remaining one-step differences (approximate instructions, the UNORM8 fetches)
enter encoder 0's reprojection and the recurrent history amplifies them, as the one-step values of the output stage
did before the fused-convert fix. 1280x720 -> 4K: 5.68 / 6.24 ms (default 5.78 / 6.35). The default build is
unchanged and still byte-identical to the captures.
### This GPU's own arithmetic for preset F: `D4R_RR_RELAXED` (2026-10-13)
Two NVIDIA setups differ from each other as much as a relaxed build differs from the reference: NVIDIA's kernels on
an RTX 4000 Ada (driver 595) and on the Blackwell cards (driver 580) give images 41.0 / 39.0 / 38.9 dB apart on
capF's three frames, although their approximate instructions and filtered texture fetches are bit-identical
(`build/rr/f/nv/math/`); the `D4R_RR_FUSE` build matches the Blackwell result only. So there is no single NVIDIA
image to reproduce, and `compile_rr2.py` now takes `D4R_RR_RELAXED=1|2` like `build_native_model.sh` does for D / E
(`rr_exact.glsl`: f32 accumulators across the 32-deep steps, one rounding into FP8; 2 also leaves the conversions to
the driver). Preset F, 1280x720 -> 4K: 5.40 / 5.99 ms against 5.78 / 6.35 (network 3.45 / 3.77 against 3.83 / 4.14);
40.8 / 39.1 / 39.2 dB from the exact build and 39-41 dB from either NVIDIA card. The default stays exact.
Why only 6%: with the rounding gone the products are about a sixth of a full-resolution stage. dec0's first kernel
at 720p, relaxed (`build/rr/f/nv/exp/t.sh`, cutting parts out): 626 us, of which the two attention heads ~370, the
four MLP branches ~100-160, expansion / residual / store ~120-160, an empty dispatch 10; inside the attention the
softmax table ~30, the bias tiles ~40, weight loads ~12. Its 88 products alone would take ~108 us at the measured
198 T MAC/s. The compiled stage still has 338 binary16 <-> f32 conversions, 176 FP8 conversions and 128 clamps for
44 product instructions (biases and tables are binary16, the activation and softmax tables are indexed by binary16
bits, FP8 conversions saturate): what is left is format traffic and the four waves' shared-memory hand-offs, not
arithmetic. NVIDIA's kernel runs the same stage in 203 us with one 32-thread warp per window.
`D4R_RR_RELAXED=3` adds `RR_NATIVE` (rr2_m.comp): the residual steps and the MLP activation in f32, straight into
FP8, with no binary16 in between and no activation table. 1280x720 -> 4K: 5.17 / 5.72 ms (network 3.24 / 3.48), still
39-41 dB from the exact build and from both NVIDIA cards. dec0's first kernel at that level: 588 us; the barriers
cost nothing; reading the constants (weight tiles, bias tiles, the softmax table) costs 170 us; the products about
108; the rest is the tokens' way through shared memory and the arena, conversions and the expansion
(`build/rr/f/nv/exp/run1.sh`).
What bounds a stage (2026-10-13; `build/rr/f/nv/exp/`). dec0's first kernel at level 3 takes 581 us at 720p and a
wave executes about 2500 instructions (`dyn.py`): 96 products, 356 binary16 <-> f32 conversions, 296 FP8
conversions, 224 clamps of the saturating FP8 conversions, 142 buffer loads, 75 shared-memory operations, about 400
wait / delay markers. What was tried on top of level 3, none of it a gain:
- the clamps removed: -30 us (2.2 ps per instruction and wave), but not usable: without the clamp the hardware
conversion returns NaN above 464 where the clamp gives 448 (`satprobe.comp`);
- f32 copies of the biases: 176 fewer conversions and 78 more loads per wave, slightly slower (587-595 us): a load
costs about three arithmetic instructions;
- the softmax cubic computed, not read from its table: the same time;
- no barriers: the same time;
- one wave per window walking the four token slices, every weight, key and value tile read once
(`rr2_w.comp`, byte-identical to the four-wave build): 27% fewer instructions per window, 30% fewer buffer loads,
60% fewer shared-memory loads, and the same time (600 us); it needs all 256 vector registers (four waves: 192).
So time here is not a function of the instruction count alone, and there is no model yet that predicts both the
clamp result and the one-wave result. The output stage is compute-bound too: 100 / 300 extra integer steps per
thread in the filter add 0.45 / 1.33 ms (`dz.sh`). NVIDIA's kernel runs dec0's first kernel in 203 us.
Shared memory throttles this GPU (2026-10-13, `build/rr/f/nv/cost/occ.py`, `occ2.py`). A dependent chain of f32
multiply-adds runs at full rate while the workgroup declares up to about 4 KiB of `shared` storage per wave; above
that the rate falls in proportion (one wave per workgroup: 6 KiB 1.24x slower, 8 KiB 1.55x, 13 KiB 3.1x, 62 KiB 12x;
four waves: 13 KiB full rate, 26 KiB 1.55x, 52 KiB 3.1x). The register count does not matter (0 to 224 live values:
the same rate). The decoders' first kernels held 21 to 46 KiB (LOW, EX, XE): `rr2_m.comp` now keeps the 4 x 4 low
tokens in KVS and the block's unquantised input in each lane's registers (`xr`), byte-identical on every stage
checked, and the network is 0.10 - 0.13 ms faster at 1280x720 -> 4K: exact build 5.68 / 6.25 ms (was 5.78 / 6.35),
`D4R_RR_RELAXED=3` 5.07 / 5.56 ms (was 5.17 / 5.72). It is also why the first one-wave prototype (13 KiB in one
wave) showed no gain: with 4 KiB (`exp/rr2_w.comp`: the low tokens borrow XT, the expansion borrows KVS, V^T is
stored, read into registers and overwritten by QK) it takes 534 us against the four-wave shader's 581, still
byte-identical to it. Extra arithmetic injected into either shader costs one instruction per cycle on about 128
wave slots (3.8e11 wave-instructions per second); NVIDIA's card measures 7.5e11 (`build/rr/f/nv/sass/rate.cpp`), and
its compiled kernel for this stage is 3,729 instructions per window (`sass/nv/dec0_0.sass`) against about 10,000
executed here. What an instruction costs here (`cost/cost.py`, `cost2.py`): an f32 multiply-add is the unit; a
32-bit integer multiply about seven; converting a pair to FP8 with the clamp about eight; a buffer read about 0.1 ps
per byte and lane (16 bytes: some forty multiply-adds), nearly twice that from a 1 MiB table.
`engine/rr/rr2_w.comp` (2026-10-13) is decoder 0's first kernel for `D4R_RR_RELAXED=3` with one wave per window:
every weight, key and value tile read once per window, the expansion plus skip computed once and kept in registers,
the scores' cross quadrant computed once and reused transposed (368 products for 384), FP8 values converted in
pairs, 4 KiB of shared memory. `compile_rr2.py` uses it for dec0_0 at level 3 (`ONE_WAVE`). The stage alone:
453 us against 581 us at 1280x720, arena byte-identical to `rr2_m.comp` at that level
(`build/rr/f/nv/invest-20261013/REPORT.md`: 20 variants, each change's instruction count and time). In the engine
every image of capF's three frames is identical to the four-wave level 3 model's, and 1280x720 -> 4K is about
5.00 / 5.55 ms (four-wave level 3: 5.10 / 5.61; network 3.05 / 3.31 against 3.15 / 3.39). The shader covers C = 32
only; the other stages still use `rr2_m.comp`.
`rr2_w.comp` also covers encoder 0's block now (KIND 4, embedded by `gen_rr2_enc0.py --one-wave` in NVIDIA's own
32-thread group, the translated guide tail on the same wave; `compile_rr2.py` `ONE_WAVE = {'dec0_0', 'enc0'}` at
`D4R_RR_RELAXED=3`). All 45 images of capF's three frames are identical to the four-wave level 3 model's.
1280x720 -> 4K, four interleaved runs each: four-wave 5.10 / 5.68 ms, dec0_0 on one wave 5.04 / 5.62, encoder 0 as
well 4.98 / 5.59. Decoder 0's second kernel is not converted: it hands r and the MLP sum (8 KiB of binary16) to its
translated tail through shared memory, which a single wave cannot hold without losing the 4 KiB rate.
### Presets D and E against a real NVIDIA card, and level 3 for them (2026-10-14)
NVIDIA's own kernels for presets D and E were replayed on an RTX PRO 4500 Blackwell (drivers 595.91 and 580.178: the
two give identical images; `build/rr/f/nv/pack.py --corpus ~/.cache/d4r-rr-corpus --first luma_convert`, captures
capD / capE3 at 640x360 -> 1280x720, every frame a reset frame). `D4R_RR_RELAXED=3` now also exists for
`build_native_model.sh`: level 2 plus `RR_NATIVE` (`rr_exact.glsl` `ACTQ8`: the MLP activation in f32, two elements
per saturating conversion, no activation table), and from level 2 on the modules are no longer passed through
`spv_exact_convert.py`. Presented image against the NVIDIA card, and the time at 1280x720 -> 4K (reset / next frame):
| build | D, dB | D, ms | E, dB | E, ms |
|---|---|---|---|---|
| exact (identical to the CUDA path through ZLUDA) | 38.5 / 39.0 / 38.5 | 4.4 / 4.7 | 40.0 / 39.9 / 40.0 | 4.34 / 4.54 |
| `D4R_RR_RELAXED=1` | 36.8 / 37.4 / 36.8 | 3.86 / 4.08 | 38.5 / 38.3 / 38.3 | 3.74 / 3.98 |
| `D4R_RR_RELAXED=2` | 36.8 / 37.4 / 36.8 | 3.76 / 3.99 | 38.5 / 38.3 / 38.3 | 3.65 / 3.82 |
| `D4R_RR_RELAXED=3` | 36.7 / 37.4 / 36.7 | 3.65 / 3.95 | 38.5 / 38.3 / 38.3 | 3.51 / 3.70 |
The NVIDIA card runs D at 1280x720 -> 4K in 2.19 / 2.39 ms. Unlike preset F, relaxing costs about 1.6 dB here. What
was measured and not built for D / E (`build/rr/f/nv/dstage.sh`, `dnl.sh`, level 3, 720p):
- shared memory: encoder 1 with 4 / 8 / 16 KiB of extra shared storage takes 287 -> ~305 / ~320 / ~343 us, so the
stages above the limit (decoder 1 at 21 KiB, decoder 0 at 17 KiB) would gain a few per cent each, under 1% of a
frame;
- one wave per window: with every weight tile replaced by one tile loaded once (wrong values, the ceiling of sharing
loads between a window's waves) decoder 0 takes the same time (803-818 us), encoder 1 the same, decoder 1 10% less.
These windows also have 96 key tokens (3 KiB each for the tokens, the keys and the values at C = 32), which does
not fit the 4 KiB a single wave can hold at full rate.
engine/rr/build_native.py
+110
View File
@@ -0,0 +1,110 @@
#!/usr/bin/env python3
"""Turn a preset-E Ray Reconstruction model directory (engine/compile_rr.py) into the native one: every Swin stage
from d4r's own shaders instead of translated PTX.
usage: build_native.py CAPTURE MODEL_DIR [--frame N] [--work DIR]
enc0 engine/rr/enc0_m.comp, generated (gen_enc0_m.py) from the model's translated enc0.generated.comp,
whose per-pixel feature gather it keeps; weights: NGX's image + enc0_pack.py's blob
enc1-enc5, engine/rr/swin_m.comp with swin_pack.py's per-stage header and weights
dec4-dec1
dec0 engine/rr/dec0_m.comp with dec0_pack.py's weights
The weights are extracted from ZLUDA replays of the capture's launches (engine/rr/phase_dump.py: every mma's operands
of one window) and the stages' weight images; extraction checks itself against the dump. Like the model directory,
everything written is derived from the user's DLL and stays on the machine. Runs one replay or compile at a time
(phase_dump.py caps its replay at 24 GB).
"""
import argparse
import os
import shutil
import subprocess
import sys
from pathlib import Path
import numpy as np
HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
from rr_capture import Capture # noqa: E402
import swin_ref # noqa: E402
import dec0_ref # noqa: E402
import enc0_ref # noqa: E402
SWIN = ['enc1', 'enc2', 'enc3', 'enc4', 'enc5', 'dec4', 'dec3', 'dec2', 'dec1']
def run(*cmd, **kw):
subprocess.run([str(c) for c in cmd], check=True, **kw)
def glsl(source, output, *includes):
run('glslangValidator', '-V', '--target-env', 'vulkan1.3', *[f'-I{i}' for i in includes], source, '-o', output,
stdout=subprocess.DEVNULL)
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('model')
ap.add_argument('--frame', type=int, default=2)
ap.add_argument('--work', help='dumps and build files (default MODEL/native-work)')
o = ap.parse_args()
model = Path(o.model)
work = Path(o.work) if o.work else model / 'native-work'
work.mkdir(parents=True, exist_ok=True)
cap = Capture(os.path.expanduser(o.capture))
manifest = (model / 'rr.txt').read_text().split('\n')
def dump(stage):
out = work / f'{stage}_all.npz'
launch = cap.find(f'swin_{stage}', o.frame)
block = '5,5' if launch.grid[0] > 5 and launch.grid[1] > 5 else '1,1'
if not out.exists():
run(sys.executable, HERE / 'phase_dump.py', o.capture, stage, o.frame, block, out)
return out, block
def stage_line(stage, shader, weights):
for i, line in enumerate(manifest):
f = line.split()
if len(f) > 4 and f[0] == 'stage' and f[1] == stage:
f[2], f[4] = shader, weights
manifest[i] = ' '.join(x for x in f if x != 'noopt')
for name in ('rr_exact.glsl', 'rr_attn.glsl', 'rr_decfilter.glsl', 'enc0_m_decl.glsl', 'enc0_m_tail.glsl'):
shutil.copy(HERE / name, work / name)
for stage in SWIN:
d, block = dump(stage)
print(f'{stage}: extracting', flush=True)
l, ex, ey, sx, sy, inp, outs, tex = swin_ref.load(o.capture, stage, o.frame)
bx, by = (int(v) for v in block.split(','))
swin_ref.extract(stage, d, l.weights, work / f'{stage}_w.npz', inp, 8 * bx - sx, 8 * by - sy, ex, ey)
sd = work / stage
sd.mkdir(exist_ok=True)
run(sys.executable, HERE / 'swin_pack.py', stage, work / f'{stage}_w.npz', model / f'weights-{stage}m.bin',
sd / 'swin_stage.h')
glsl(HERE / 'swin_m.comp', model / f'{stage}_m.spv', sd, work)
stage_line(stage, f'{stage}_m.spv', f'weights-{stage}m.bin')
# dec0
d, _ = dump('dec0')
print('dec0: extracting', flush=True)
dec0_ref.extract(d, cap.find('swin_dec0', o.frame).weights, work / 'dec0_w.npz')
run(sys.executable, HERE / 'dec0_pack.py', work / 'dec0_w.npz', model / 'weights-dec0m.bin', work / 'dec0_m_w.h')
glsl(HERE / 'dec0_m.comp', model / 'dec0_m.spv', work)
stage_line('dec0', 'dec0_m.spv', 'weights-dec0m.bin')
# enc0: NGX's weight image (the translated feature gather reads it) followed by the native blob
d, _ = dump('enc0')
print('enc0: extracting', flush=True)
nv = cap.find('swin_enc0', o.frame).weights
enc0_ref.extract(d, nv, work / 'enc0_w.npz')
base = (len(nv) + 255) // 256 * 256
run(sys.executable, HERE / 'enc0_pack.py', work / 'enc0_w.npz', work / 'enc0_native.bin', work / 'enc0_m_w.h', base)
(model / 'weights-enc0m.bin').write_bytes(nv + bytes(base - len(nv)) + (work / 'enc0_native.bin').read_bytes())
run(sys.executable, HERE / 'gen_enc0_m.py', model / 'enc0.generated.comp', work / 'enc0_m.comp')
glsl(work / 'enc0_m.comp', model / 'enc0_m.spv', work)
stage_line('enc0', 'enc0_m.spv', 'weights-enc0m.bin')
(model / 'rr.txt').write_text('\n'.join(manifest))
print(f'{model / "rr.txt"}: native stages written')
if __name__ == '__main__':
main()
engine/rr/build_native_model.sh
+158
View File
@@ -0,0 +1,158 @@
#!/usr/bin/env bash
set -euo pipefail
# Builds a Ray Reconstruction model directory whose network stages are the native shaders (engine/rr/*_m.comp),
# for the preset the capture was taken with (D or E).
#
# usage: build_native_model.sh CAPTURE_DIR OUT_DIR [NVNGX_DLSSD_DLL]
# build_native_model.sh --dll NVNGX_DLSSD_DLL D|E OUT_DIR
#
# With --dll no capture is needed: everything comes from the DLL (its PTX: kernels/tools/extract_dlss_ptx.py; the
# argument-block templates: engine/rr/rr_templates.py; the weights: dll_weights.py, model 0 = D, model 1 = E in 310.7's
# file order) and step 3's checks against a captured frame are skipped. This is what install.sh runs.
#
# CAPTURE_DIR is a bridge capture of a few denoiser frames (engine/rr/rr_capture.py); the DLL is the one it was taken
# with (default: CAPTURE_DIR/run/app/nvngx_dlssd.dll, then $D4R_RR_DLL). The steps:
# 1. engine/compile_rr.py: the translated model (the argument-block templates, the exposure passes, the translated
# output stage and the translated encoder 0 whose feature gather the native one keeps)
# 2. engine/rr/dll_weights.py: the network's weights, read from the DLL itself. The capture only selects which of
# the DLL's models it ran, and a DLL whose kernels are not the ones the native shaders were written against is
# refused (D4R_RR_ANY_KERNELS=1 builds anyway)
# 3. per stage: the weights checked against the capture's frame (STAGE_ref.py), packed (STAGE_pack.py) and the
# native shader compiled against that packing's header
# 4. rr.txt rewritten to the native stages; the weight images NGX uploaded in the capture are removed
# Preset D differs from E in three kernels: encoder 0 (its feature gather and 30 input features for E's 32), decoder 0
# (the eight bytes it hands the output stage: -DRR_PRESET_D) and the output stage, which is the translated kernel
# (it has no subgroup operations, so it runs with 64-wide waves).
# Everything written is derived from the user's own nvngx_dlssd.dll and stays on their machine.
#
# GPU jobs run one at a time, each in a memory-capped scope (a translated stage's pipeline compile can take tens of GB).
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
USAGE="usage: build_native_model.sh CAPTURE_DIR OUT_DIR [NVNGX_DLSSD_DLL] | --dll NVNGX_DLSSD_DLL D|E OUT_DIR"
CAP=""
if [[ "${1:-}" == --dll ]]; then
DLL="$(realpath "${2:?$USAGE}")"; WANT="${3:?$USAGE}"; OUT="$(realpath -m "${4:?$USAGE}")"
[[ "$WANT" == D || "$WANT" == E ]] || { echo "$USAGE" >&2; exit 2; }
else
CAP="$(realpath "${1:?$USAGE}")"
OUT="$(realpath -m "${2:?$USAGE}")"
DLL="${3:-${D4R_RR_DLL:-$CAP/run/app/nvngx_dlssd.dll}}"
fi
RR="$ROOT/engine/rr"
WORK="$OUT/native-work"
[[ -f "$DLL" ]] || { echo "no nvngx_dlssd.dll at $DLL: pass the DLL the capture was taken with" >&2; exit 1; }
FRAME="${D4R_RR_MODEL_FRAME:-2}"
# each GPU / compiler job in a memory-capped scope where systemd provides one (D4R_NO_SCOPE=1: none)
CAPPED=()
if [[ "${D4R_NO_SCOPE:-0}" != 1 ]] && command -v systemd-run > /dev/null && systemd-run --user --scope -q true 2> /dev/null; then
CAPPED=(systemd-run --user --scope -q -p MemoryMax=24G -p MemorySwapMax=0)
fi
mkdir -p "$OUT" "$WORK"
cd "$ROOT"
if [[ -n "$CAP" ]]; then
python3 engine/compile_rr.py "$CAP" "$OUT" > "$WORK/compile_rr.log" 2>&1 || { tail -5 "$WORK/compile_rr.log"; exit 1; }
python3 "$RR/dll_weights.py" "$DLL" "$WORK/dll" --capture "$CAP" ${D4R_RR_ANY_KERNELS:+--any-kernels}
else
python3 kernels/tools/extract_dlss_ptx.py "$DLL" "$WORK/ptx" > /dev/null
python3 engine/compile_rr.py --dll "$DLL" --preset "$WANT" "$OUT" --ptx-dir "$WORK/ptx" --network-bindings-only \
> "$WORK/compile_rr.log" 2>&1 || { tail -5 "$WORK/compile_rr.log"; exit 1; }
python3 "$RR/dll_weights.py" "$DLL" "$WORK/dll" --model "$([[ "$WANT" == D ]] && echo 0 || echo 1)" \
${D4R_RR_ANY_KERNELS:+--any-kernels} > /dev/null
fi
PRESET="$(awk '$1 == "preset" { print $2 }' "$OUT/rr.txt")"
# RR_EXACT_CONVERT: the compiled module's float conversions are decorated NoContraction (spv_exact_convert.py), which
# keeps the driver from folding a binary16 rounding away without the opaque OR the shaders otherwise need
DEFS=(-DRR_EXACT_CONVERT); [[ "$PRESET" == 4 ]] && DEFS+=(-DRR_PRESET_D)
# D4R_RR_RELAXED=1: one rounding where NVIDIA has two (engine/rr/rr_exact.glsl): faster, not bit-exact
# D4R_RR_RELAXED=3 adds RR_NATIVE: the MLP activation in f32 straight into FP8, no activation table
# Level 3 is the default (the fastest build; about 1.6 dB further from an NVIDIA card for D / E than the exact one);
# D4R_RR_RELAXED=0 builds NVIDIA's roundings exactly (bit-identical to the CUDA path through ZLUDA).
D4R_RR_RELAXED="${D4R_RR_RELAXED:-3}"
if (( D4R_RR_RELAXED > 0 )); then
DEFS+=("-DRR_RELAXED=$(( D4R_RR_RELAXED > 2 ? 2 : D4R_RR_RELAXED ))"); (( D4R_RR_RELAXED >= 3 )) && DEFS+=(-DRR_NATIVE)
fi
# D4R_RR_MIX=1: the 32-deep steps round once through v_fma_mix (the shaders' K32): ~5% faster, not provably bit-exact
[[ "${D4R_RR_MIX:-0}" == 1 ]] && DEFS+=(-DRR_MIX)
# D4R_RR_GFX11=1: RDNA3 (no FP8 matrices, no VK_EXT_shader_float8): the stages reproduce ZLUDA's gfx11 lowering bit for
# bit (engine/rr/rr_gfx11.glsl); exact arithmetic only, and each native weight file carries a widened copy (W_WIDE)
GFX11=0
if [[ "${D4R_RR_GFX11:-0}" == 1 ]]; then
GFX11=1; DEFS=(-DRR_EXACT_CONVERT -DRR_GFX11); [[ "$PRESET" == 4 ]] && DEFS+=(-DRR_PRESET_D)
(( D4R_RR_RELAXED == 0 )) && [[ "${D4R_RR_MIX:-0}" != 1 ]] ||
echo "D4R_RR_GFX11: relaxed / mix builds do not exist for RDNA3; building the exact model" >&2
D4R_RR_RELAXED=0
fi
# The gfx11 runtime also reproduces NGX's CPU camera arithmetic, whose residues feed temporal specular motion.
(( GFX11 == 0 )) || sed -i '/^preset /a arithmetic gfx11' "$OUT/rr.txt"
# widen WEIGHTS: on RDNA3, append the widened copy and print -DW_WIDE=...; nothing otherwise
widen() { (( GFX11 )) && printf -- '-DW_WIDE=%su' "$(python3 "$RR/widen_weights.py" "$1")" || true; }
glsl() { # glsl OUT.spv SOURCE INCLUDE_DIR... [-DEXTRA]
local out="$1" src="$2"; shift 2
local dir inc=() extra=(); for dir in "$@"; do [[ "$dir" == -D* ]] && extra+=("$dir") || inc+=("-I$dir"); done
"${CAPPED[@]}" glslangValidator --target-env vulkan1.3 -V "${DEFS[@]}" "${extra[@]}" "${inc[@]}" -o "$out" "$src" > "$WORK/glsl.log" 2>&1 ||
{ grep -E 'ERROR' "$WORK/glsl.log" | head -5; exit 1; }
# from level 2 on the conversions are left to the driver (as compile_rr2.py does for preset F)
(( D4R_RR_RELAXED >= 2 )) || python3 "$RR/spv_exact_convert.py" "$out" > /dev/null
}
native() { # native STAGE WEIGHT_FILE: point rr.txt's stage line at STAGE_m.spv
sed -i -E "s/^stage $1 $1\.spv 32 weights-$1\.bin (.*)/stage $1 $1_m.spv 32 $2 \1/; /^stage $1 /s/ noopt\$//" "$OUT/rr.txt"
grep -q "^stage $1 $1_m.spv " "$OUT/rr.txt" || { echo "rr.txt: no stage line for $1" >&2; exit 1; }
}
# the Swin stages shared by both presets (engine/rr/swin_m.comp)
for s in enc1 enc2 enc3 enc4 enc5 dec4 dec3 dec2 dec1; do
b=1,1; case $s in enc1|enc2|dec1|dec2) b=5,5 ;; esac
d="$WORK/$s"; mkdir -p "$d"
cp "$WORK/dll/$s.npz" "$d/weights.npz"
[[ -z "$CAP" ]] || python3 "$RR/swin_ref.py" "$CAP" "$s" "$FRAME" "$d/weights.npz" --blocks "$b" 0,0 | tail -2 | sed "s/^/$s: /"
python3 "$RR/swin_pack.py" "$s" "$d/weights.npz" "$OUT/weights-${s}m.bin" "$d/swin_stage.h" > /dev/null
glsl "$OUT/${s}_m.spv" "$RR/swin_m.comp" "$d" "$RR" $(widen "$OUT/weights-${s}m.bin")
native "$s" "weights-${s}m.bin"
done
# decoder 0
d="$WORK/dec0"; mkdir -p "$d"
cp "$WORK/dll/dec0.npz" "$d/weights.npz"
[[ -z "$CAP" ]] || python3 "$RR/dec0_ref.py" "$CAP" "$FRAME" "$d/weights.npz" --blocks 5,5 0,0 | tail -2 | sed 's/^/dec0: /'
python3 "$RR/dec0_pack.py" "$d/weights.npz" "$OUT/weights-dec0m.bin" "$d/dec0_m_w.h" > /dev/null
glsl "$OUT/dec0_m.spv" "$RR/dec0_m.comp" "$d" "$RR" $(widen "$OUT/weights-dec0m.bin")
native dec0 weights-dec0m.bin
# encoder 0: the translated feature gather + the native rest
d="$WORK/enc0"; mkdir -p "$d"
cp "$WORK/dll/enc0.npz" "$d/weights.npz"
# preset E's gather dithers with the noise table that starts NGX's weight image: the table, then the packing
BASE=0; [[ "$PRESET" == 4 ]] || BASE=4096
python3 "$RR/enc0_pack.py" "$d/weights.npz" "$d/enc0_m.bin" "$d/enc0_m_w.h" "$BASE" > /dev/null
if [[ "$BASE" != 0 ]]; then
[[ "$(stat -c %s "$WORK/dll/noise.bin" 2>/dev/null)" == 4096 ]] || { echo "no noise table found in $DLL" >&2; exit 1; }
cat "$WORK/dll/noise.bin" "$d/enc0_m.bin" > "$OUT/weights-enc0m.bin"
else
cp "$d/enc0_m.bin" "$OUT/weights-enc0m.bin"
fi
python3 "$RR/gen_enc0_m.py" "$OUT/enc0.generated.comp" "$d/enc0_m.comp" > /dev/null
glsl "$OUT/enc0_m.spv" "$d/enc0_m.comp" "$d" "$RR" $(widen "$OUT/weights-enc0m.bin")
native enc0 weights-enc0m.bin
# the translated output stage (preset D; preset E's when engine/rr/hkpn.comp is not used) has no subgroup operations
if grep -q '"shader": "generated from the DLL"' <(python3 -c "import json,sys; print(json.dumps(json.load(open('$OUT/model.json'))['stages']['hkpn']))"); then
sed -i -E 's/^stage hkpn hkpn\.spv 32 /stage hkpn hkpn.spv 64 /' "$OUT/rr.txt"
fi
for s in enc0 enc1 enc2 enc3 enc4 enc5 dec4 dec3 dec2 dec1 dec0; do rm -f "$OUT/weights-$s.bin"; done
# test_layer.py must bind optional images as well as the ones present in a capture. Keep a sidecar for every used
# shader, including native stages, so an absent input never shifts the shader's descriptor bindings.
python3 - "$OUT" <<'PY'
import json, sys
from pathlib import Path
out = Path(sys.argv[1])
for line in (out / 'rr.txt').read_text().splitlines():
fields = line.split()
if not fields or fields[0] != 'stage':
continue
textures = [] if fields[5] == '-' else list(map(int, fields[5].split(',')))
surfaces = [] if fields[6] == '-' else [[int(s.split(':')[0]), s.split(':')[1]] for s in fields[6].split(',')]
(out / fields[2]).with_suffix('.json').write_text(json.dumps({'textures': textures, 'surfaces': surfaces}) + '\n')
PY
printf 'Native Ray Reconstruction model (preset %s) in %s\n' "$([[ "$PRESET" == 4 ]] && echo D || echo E)" "$OUT"
grep '^stage' "$OUT/rr.txt" | awk '{ printf " %s %s (%s-wide)\n", $2, $3, $4 }'
engine/rr/dec0_m.comp
+390
View File
@@ -0,0 +1,390 @@
#version 460
// Native decoder 0 of the DLSS-D (preset E) denoiser, register-resident: cuda_dldn_engine_swin_dec0_kernel computed
// like engine/rr/enc1_m.comp (whose header explains the transposed products and the slot orders). The arithmetic is
// engine/rr/dec0_ref.py's (bit-exact to NVIDIA except surface 112's bilinear taps, which NVIDIA's texture unit rounds
// differently in a few percent of the values by one ulp); binding 1 holds engine/rr/dec0_pack.py's packing.
//
// One workgroup per 8x8-token window at the render resolution, four waves:
// expand the 6 x 6 low-resolution tokens under the window and its border (dec1's 64 channels) -> 4 sub-positions
// x 32 channels; wave w computes sub-position w, FP8 into EX in slot order
// prologue one key token per thread: expanded + skip (enc0's tokens), RMS norm -> XT; window tokens -> XR (f16)
// block two 32-channel heads and four MLP branches of 32 -> 32 -> 32, as enc1_m with C = 32
// outputs one token per thread: 8 FP8 channels for the output stage (argument 48, [y][x][8]), tanh of four
// channels to surface 144, and the kernel-prediction filter of texture 88 (bilinear from half resolution)
// blended with textures 96 / 104 to surface 112
#extension GL_KHR_cooperative_matrix : require
#extension GL_KHR_memory_scope_semantics : require
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_8bit_storage : require
#ifndef RR_GFX11
#extension GL_EXT_float_e4m3 : require
#endif
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_EXT_control_flow_attributes : require
#extension GL_GOOGLE_include_directive : require
layout(local_size_x = 128) in;
#include "dec0_m_w.h"
#include "rr_exact.glsl"
layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param;
#ifndef RR_GFX11
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
#endif
layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16;
layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u;
layout(set = 0, binding = 1, std430) readonly buffer W8U { uint8_t w[]; } w8u;
layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
layout(set = 0, binding = 3) uniform sampler2D T88;
layout(set = 0, binding = 4) uniform sampler2D T96;
layout(set = 0, binding = 5) uniform sampler2D T104;
layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S112;
layout(set = 0, binding = 7, rgba8_snorm) uniform writeonly image2D S144;
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define HAC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#ifndef RR_GFX11
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#endif
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CM gl_CooperativeMatrixLayoutColumnMajor
#ifdef RR_GFX11
#include "rr_gfx11_k32.glsl"
#else
#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM)
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
#define Q8(dst, src) Q8PAIRS(dst, src)
// One 32-deep step onto an accumulator that holds binary16 values. The reference (and the default here) adds the
// accumulator inside the f32 products and rounds the result to binary16. RR_MIX (build_native_model.sh with
// D4R_RR_MIX=1) keeps the accumulator as binary16, sums the products from zero and takes f16(P + acc), which the
// driver compiles to one v_fma_mix per element in place of two conversions: about 5% faster, but a single rounding
// where the reference has two, so a value can come out one binary16 step away (engine/rr/README.md).
// RAC is the accumulator type, RV(x) a binary16 value as its element; K32H is the step on a binary16 matrix.
#ifdef RR_MIX
#define RAC HAC
#define RV(x) (x)
#define K32(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(0.0))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_] + float((acc)[i_])); }
#define K32H(acc, a0, b0, a1, b1) K32(acc, a0, b0, a1, b1)
#else
#define RAC ACC
#define RV(x) float(x)
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) }
#define K32H(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(acc))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_]); }
#endif
#endif
const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0;
// shared memory, aliased by phase: expand and prologue [LOW | EX], then the attention [K | V^T]
shared uint POOL[48 * 16 + 48 * 32];
#define LOW POOL // the low-resolution tokens (6 x 6, padded to 48): 64 FP8 channels, arena order
#define EX_AT 768u // their expansion: 4 sub-positions x 32 FP8 channels in slot order (32 words a row)
#define RR_KS POOL // the current head's K
#define RR_KS0 0u
#define RR_VS POOL // the current head's V^T
#define RR_VS0 768u
shared uint XT[96 * 8]; // the normalised key tokens: slot-major, 32 FP8 channels in slot order
shared float16_t XR[64 * 32]; // the window tokens before the norm (the residual), slot-major
shared float16_t TAP[100 * 4]; // texture 88 bilinear at the window's 10 x 10 filter tap positions
#define OUTF XR // after the block: the final tokens, token-major, natural channels (f16)
// the window's key tokens: the 8x8 window, then a border row above, below, a column left and right
ivec2 keyXY(uint t, int x0, int y0)
{
if (t < 64u) return ivec2(x0 + int(t & 7u), y0 + int(t >> 3u));
const int i = int(t & 7u), side = int((t - 64u) >> 3u);
return ivec2(side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8), side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i));
}
int clampTap(int v, int n) { return v < 0 ? 1 : (v >= n ? n - 2 : v); }
float tanhNV(float t)
{
const float m = abs(t), s = t * t;
const float big = m >= 9.0109138 ? 1.0 : fma(1.0 / (exp2(m * 2.885390043258667) + 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);
}
void main()
{
const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID;
const uint hl = lane >> 4u, col = lane & 15u;
RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u);
const int ex = int(param.p[0]), ey = int(param.p[1]), hx = ex / 2, hy = ey / 2;
const uint loOff = param.p[2], skOff = param.p[6], outOff = param.p[12], shift = param.p[20];
const int x0 = 8 * int(gl_WorkGroupID.x) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u);
const int X0 = (x0 >> 1) - 1, Y0 = (y0 >> 1) - 1;
const uint lplane = uint(hx * hy) * 32u;
// ---------------------------------------------------------------- expand
for (uint i = tid; i < 48u * 16u; i += 128u) {
const uint r = i >> 4u, wd = i & 15u;
uint u = 0u;
if (r < 36u) {
const int X = mirror(X0 + int(r % 6u), hx), Y = mirror(Y0 + int(r / 6u), hy);
u = arena.w[(loOff + (wd >> 3u) * lplane + uint(Y * hx + X) * 32u) / 4u + (wd & 7u)];
}
LOW[i] = u;
}
barrier();
{
A8 we[2][4];
[[unroll]] for (uint mm = 0u; mm < 2u; ++mm)
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) LDA(we[mm][kt], W_WE, 64u, 2u * wv + mm, kt);
[[unroll]] for (uint n = 0u; n < 3u; ++n) {
B8 lb[4];
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) LDS8B(lb[kt], LOW, n * 256u + 4u * kt, 16u);
[[unroll]] for (uint mm = 0u; mm < 2u; ++mm) {
const uint m = 2u * wv + mm;
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_BE / 2u + 16u * m + 8u * hl + uint(i)]);
K32(a, we[mm][0], lb[0], we[mm][1], lb[1])
K32(a, we[mm][2], lb[2], we[mm][3], lb[3])
Q8STB(a, POOL, EX_AT + n * 512u + 4u * m, 32u)
}
}
}
barrier();
// ---------------------------------------------------------------- prologue: one key token per thread
if (tid < 96u) {
const ivec2 k = keyXY(inter(tid), x0, y0);
const uint r = uint(((k.y >> 1) - Y0) * 6 + ((k.x >> 1) - X0)), s = uint(2 * (k.y & 1) + (k.x & 1));
const int xx = mirror(k.x, ex), yy = mirror(k.y, ey);
const uint sb = (skOff + uint(yy * ex + xx) * 32u) / 4u;
float16_t v[32];
[[unroll]] for (uint wd = 0u; wd < 8u; ++wd) {
const uint e = POOL[EX_AT + r * 32u + 8u * s + wd], sk = arena.w[sb + wd];
[[unroll]] for (uint b = 0u; b < 4u; ++b) {
#ifdef RR_GFX11
precise float16_t x = g11_e4h((e >> (8u * b)) & 0xffu) + g11_e4h((sk >> (8u * b)) & 0xffu);
#else
precise float16_t x = float16_t(uintBitsToFloate4m3EXT(uint8_t((e >> (8u * b)) & 0xffu)))
+ float16_t(uintBitsToFloate4m3EXT(uint8_t((sk >> (8u * b)) & 0xffu)));
#endif
v[4u * wd + b] = x;
}
}
// dec0_ref.norm_plane: per word t the pair sums, words t and t + 4, then t ^ 1, t ^ 2, the halves
precise f16vec2 h[4];
[[unroll]] for (uint t = 0u; t < 4u; ++t) {
precise f16vec2 p0 = f16vec2(v[4u * t], v[4u * t + 1u]) * f16vec2(v[4u * t], v[4u * t + 1u])
+ f16vec2(v[4u * t + 2u], v[4u * t + 3u]) * f16vec2(v[4u * t + 2u], v[4u * t + 3u]);
precise f16vec2 p1 = f16vec2(v[16u + 4u * t], v[17u + 4u * t]) * f16vec2(v[16u + 4u * t], v[17u + 4u * t])
+ f16vec2(v[18u + 4u * t], v[19u + 4u * t]) * f16vec2(v[18u + 4u * t], v[19u + 4u * t]);
h[t] = p0 + p1;
}
precise f16vec2 h0 = h[0] + h[1], h2_ = h[2] + h[3];
precise f16vec2 hs = h0 + h2_;
precise float16_t sm = hs.x + hs.y;
precise float16_t se = sm + float16_t(0.0001220703125);
const float16_t rr = rsq16(se);
[[unroll]] for (uint wd = 0u; wd < 8u; ++wd) {
F8 c[4];
[[unroll]] for (uint b = 0u; b < 4u; ++b) {
const uint j = 4u * wd + b;
precise float16_t g = rr * w16.w[W_GAIN0 / 2u + inter(j)];
precise float16_t n = v[j] * g;
c[b] = q8h(n);
}
XT[tid * 8u + wd] = code(c[0]) | (code(c[1]) << 8u) | (code(c[2]) << 16u) | (code(c[3]) << 24u);
}
if (tid < 64u) [[unroll]] for (uint j = 0u; j < 32u; ++j) XR[tid * 32u + j] = v[j];
}
barrier();
// ---------------------------------------------------------------- the residual: window tokens + table 0
RAC R[2];
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
HAC x;
#ifdef RR_GFX11
[[unroll]] for (int i = 0; i < 8; ++i) x[i] = XR[wv * 512u + 16u * m + col * 32u + 8u * hl + uint(i)];
#else
coopMatLoad(x, XR, wv * 512u + 16u * m, 32u, CM);
#endif
[[unroll]] for (int i = 0; i < 8; ++i) {
precise float16_t c = float16_t(x[i]) + w16.w[W_TABLE0 / 2u + 16u * m + 8u * hl + uint(i)];
R[m][i] = RV(c);
}
}
// ---------------------------------------------------------------- attention, two heads
[[unroll]] for (uint hd = 0u; hd < 2u; ++hd) {
if (hd == 1u) barrier();
for (uint j = 0u; j < 3u; ++j) { // 12 jobs (token tile n, channel tile c)
const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u;
A8 wk[2], xa[2];
B8 wvb[2], xb[2];
[[unroll]] for (uint kt = 0u; kt < 2u; ++kt) {
LDS8B(xb[kt], XT, n * 128u + 4u * kt, 8u);
LDS8A(xa[kt], XT, n * 128u + 4u * kt, 8u);
LDA(wk[kt], W_WK0 + HEAD * hd, 32u, c, kt);
LDB(wvb[kt], W_WV0 + HEAD * hd, 32u, kt, c);
}
RAC a = RAC(RV(float16_t(0.0)));
K32(a, wk[0], xb[0], wk[1], xb[1])
Q8STA(a, POOL, n * 128u + 4u * c, 8u)
RAC v = RAC(RV(float16_t(0.0)));
K32(v, xa[0], wvb[0], xa[1], wvb[1])
Q8STA(v, POOL, 768u + c * 384u + 4u * n, 24u)
}
barrier();
#include "rr_attn.glsl"
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
A8 a0, a1;
LDA(a0, W_WO0 + HEAD * hd, 32u, m, 0u); LDA(a1, W_WO0 + HEAD * hd, 32u, m, 1u);
K32(R[m], a0, Ob[0], a1, Ob[1])
}
}
// ---------------------------------------------------------------- MLP
B8 Nb[2];
#ifdef RR_GFX11
F8V NbC[2] = F8V[2](uvec2(0u), uvec2(0u));
#endif
{
float16_t rv[2][8];
[[unroll]] for (uint m = 0u; m < 2u; ++m) [[unroll]] for (int i = 0; i < 8; ++i) rv[m][i] = float16_t(R[m][i]);
// dec0_ref.norm_plane on slot order: lane half hl holds words 2hl, 2hl + 1 (tile 0) and 4 + 2hl, 5 + 2hl (tile 1)
float16_t h1[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
float16_t ht[2];
[[unroll]] for (uint tq = 0u; tq < 2u; ++tq) {
const uint i0 = 4u * tq + c, i1 = i0 + 2u;
precise float16_t p0 = rv[0][i0] * rv[0][i0] + rv[0][i1] * rv[0][i1];
precise float16_t p1 = rv[1][i0] * rv[1][i0] + rv[1][i1] * rv[1][i1];
precise float16_t hv = p0 + p1;
ht[tq] = hv;
}
precise float16_t hh = ht[0] + ht[1];
h1[c] = hh;
}
const uint other = subgroupShuffleXor(bh(h1[0]) | (bh(h1[1]) << 16u), 16u);
precise float16_t g0 = h1[0] + hb(other & 0xffffu), g1 = h1[1] + hb(other >> 16u);
precise float16_t s = g0 + g1;
precise float16_t se = s + float16_t(0.0001220703125);
const float16_t r = rsq16(se);
[[unroll]] for (uint m = 0u; m < 2u; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) {
const uint k = 16u * m + 8u * hl + uint(i);
precise float16_t g = r * w16.w[W_GAIN1 / 2u + k];
precise float16_t n = rv[m][i] * g;
#ifdef RR_GFX11
NbC[m][i >> 2] |= q8h(n) << (8u * uint(i & 3));
#else
Nb[m][i] = q8h(n);
#endif
precise float16_t a = rv[m][i] + w16.w[W_TABLE1 / 2u + k];
R[m][i] = RV(a);
}
#ifdef RR_GFX11
[[unroll]] for (uint m = 0u; m < 2u; ++m) G11_B(Nb[m], NbC[m])
#endif
}
for (uint br = 0u; br < 4u; ++br) {
const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br;
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
RAC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = RV(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]);
A8 a0, a1;
LDA(a0, w1, 32u, c, 0u); LDA(a1, w1, 32u, c, 1u);
K32(h, a0, Nb[0], a1, Nb[1])
#ifdef RR_GFX11
{ F8V hc_; ACTQ8C(hc_, h) G11_B(Hb[c], hc_) }
#else
ACTQ8(Hb[c], h)
#endif
}
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
A8 a0, a1;
LDA(a0, w2, 32u, m, 0u); LDA(a1, w2, 32u, m, 1u);
K32(R[m], a0, Hb[0], a1, Hb[1])
}
}
// ---------------------------------------------------------------- outputs, one token per thread
{
const uint t = inter(16u * wv + col);
[[unroll]] for (uint m = 0u; m < 2u; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) OUTF[t * 32u + inter(16u * m + 8u * hl + uint(i))] = float16_t(R[m][i]);
}
barrier();
// texture 88 (one level down) at the window's 10 x 10 full-resolution tap positions (clamped like the filter's
// coordinates), through the bilinear sampler at NVIDIA's coordinates: the reference's texture unit does not round
// like f32 arithmetic (computing the 1/4, 3/4 blend by hand left ~0.2% of surface 112 one ulp off, capDT 2026-10-12)
if (tid < 100u) {
const int xi = clampTap(x0 - 1 + int(tid % 10u), ex), yi = clampTap(y0 - 1 + int(tid / 10u), ey);
const float iw = uintBitsToFloat(param.p[34]), ih = uintBitsToFloat(param.p[35]);
precise float cu = fma(iw, float(xi), iw * 0.5), cv = fma(ih, float(yi), ih * 0.5);
const vec4 bl = textureLod(T88, vec2(cu, cv), 0.0);
[[unroll]] for (uint c = 0u; c < 3u; ++c) TAP[tid * 4u + c] = float16_t(bl[c]);
}
barrier();
// two threads per token: half 0 filters channels 0 and 1, half 1 channel 2 and writes the output stage's codes
// and surface 144
{
const uint t = tid >> 1u, hf = tid & 1u;
const int x = x0 + int(t & 7u), y = y0 + int(t >> 3u);
if (x < 0 || y < 0 || x >= ex || y >= ey) return;
float16_t a[10];
[[unroll]] for (uint c = 0u; c < 10u; ++c) a[c] = OUTF[t * 32u + c];
if (hf == 1u) {
float16_t b[12];
[[unroll]] for (uint c = 0u; c < 12u; ++c) b[c] = OUTF[t * 32u + 12u + c];
const uint o = (outOff + uint(y * ex + x) * 8u) / 4u;
#ifdef RR_PRESET_D
// cuda_dldn_engine_swin_dec0_diamond_wallaby_kernel: channels 12, 13, 0, 18, 19, 20, 21, 0 (its only
// difference from the preset E kernel; capDT, 2026-10-12)
arena.w[o] = code(q8h(b[0])) | (code(q8h(b[1])) << 8u) | (code(q8h(b[6])) << 24u);
arena.w[o + 1u] = code(q8h(b[7])) | (code(q8h(b[8])) << 8u) | (code(q8h(b[9])) << 16u);
#else
arena.w[o] = code(q8h(b[0])) | (code(q8h(b[1])) << 8u) | (code(q8h(b[6])) << 16u) | (code(q8h(b[7])) << 24u);
arena.w[o + 1u] = code(q8h(b[8])) | (code(q8h(b[9])) << 8u) | (code(q8h(b[10])) << 16u) | (code(q8h(b[11])) << 24u);
#endif
imageStore(S144, ivec2(x, y), vec4(tanhNV(float(b[2])), tanhNV(float(b[3])), tanhNV(float(b[4])), tanhNV(float(b[5]))));
}
// surface 112: softmax of channels 0-8 over the 3x3 taps, one f16 fma chain per channel; the gate (channel 9)
// blends it with texture 96; + texture 104
const float16_t ma = max(max(max(a[0], a[2]), a[4]), a[6]), mb = max(max(max(a[1], a[3]), a[5]), a[7]);
const float16_t mx = max(max(ma, mb), a[8]);
float e[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) { precise float16_t d = a[i] - mx; e[i] = exp2(float(d) * 1.4426950216293335); }
precise float s = (e[0] + e[1]) + 0.0;
s = s + (e[2] + e[3]);
s = s + (e[4] + e[5]);
s = s + (e[6] + e[7]);
s = s + e[8];
const float16_t rr = rcp16(s);
float16_t wg[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) wg[i] = h16(r16(e[i]) * float(rr));
const uint tx = uint(x - x0), ty = uint(y - y0), c0 = 2u * hf, nc = 2u - hf;
float16_t f[2] = float16_t[2](0.0hf, 0.0hf);
bool mid = false;
[[unroll]] for (uint k = 0u; k < 9u; ++k) {
const uint at = ((ty + k / 3u) * 10u + tx + k % 3u) * 4u + c0;
[[unroll]] for (uint c = 0u; c < 2u; ++c) f[c] = fma16q(TAP[at + c], wg[k], f[c], mid);
}
if (mid) { // an f32 result on a binary16 midpoint: redo the chain with the exact fma (rare)
f[0] = 0.0hf; f[1] = 0.0hf;
for (uint k = 0u; k < 9u; ++k) {
const uint at = ((ty + k / 3u) * 10u + tx + k % 3u) * 4u + c0;
for (uint c = 0u; c < 2u; ++c) f[c] = fma16(TAP[at + c], wg[k], f[c]);
}
}
const float16_t g = sig(a[9]);
precise float16_t om = float16_t(1.0) - g;
const vec4 t96 = texelFetch(T96, ivec2(x, y), 0), t104 = texelFetch(T104, ivec2(x, y), 0);
float16_t o[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
precise float16_t fg = f[c] * g;
float16_t v = fma16(float16_t(t96[c0 + c]), om, fg);
const float16_t t4 = float16_t(t104[c0 + c]);
if (!isnan(t4)) { precise float16_t vv = v + t4; v = vv; }
o[c] = v;
}
const uint other = subgroupShuffleXor(bh(o[0]), 1u);
if (hf == 0u) imageStore(S112, ivec2(x, y), vec4(o[0], o[1], hb(other), 0.0));
}
}
engine/rr/dec0_pack.py
+63
View File
@@ -0,0 +1,63 @@
#!/usr/bin/env python3
"""Pack decoder 0's extracted weights (engine/rr/dec0_ref.py --extract) for engine/rr/dec0_m.comp.
usage: dec0_pack.py WEIGHTS.npz OUT.bin OUT.h
The layout rules are enc1_pack.py's: transposed products, every weight's columns in the k order of the product that
consumes its result (C = 32: slot j holds natural channel inter(32)[j]). The expand's 128 columns are four
sub-positions of 32 channels, each permuted the same way, so a key token's expanded half is already in slot order.
"""
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from enc1_ref import inter # noqa: E402
from enc1_pack import a_tiles, b_tiles, cubic_table, act_table # noqa: E402
I32, I96 = inter(32), inter(96)
I128 = [32 * (j // 32) + I32[j % 32] for j in range(128)]
def main():
src, out_bin, out_h = sys.argv[1:4]
w = dict(np.load(src))
e = {'we': ('A', w['we'][:, I128]), 'be': ('H', w['be'][I128])}
for h in range(2):
e[f'wk{h}'] = ('A', w[f'wk{h}'][:, I32])
e[f'wv{h}'] = ('B', w[f'wv{h}'][:, I32])
e[f'wo{h}'] = ('A', w[f'wo{h}'][:, I32])
e[f'bias{h}'] = ('H', w[f'bias{h}'][np.ix_(inter(64), I96)].T)
for b in range(4):
e[f'w1_{b}'] = ('A', w[f'w1_{b}'][:, I32])
e[f'b1_{b}'] = ('H', w[f'b1_{b}'][I32])
e[f'w2_{b}'] = ('A', w[f'w2_{b}'][:, I32])
e['gain0'] = ('H', w['gain0'])
for k in ('table0', 'gain1', 'table1'):
e[k] = ('H', w[k][I32])
e['cubic'] = ('U', cubic_table())
e['act'] = ('Y', act_table())
blob, defines = bytearray(), []
for name, (kind, a) in e.items():
while len(blob) % 256:
blob.append(0)
off = len(blob)
if kind == 'A':
blob += a_tiles(a).tobytes()
elif kind == 'B':
blob += b_tiles(a).tobytes()
elif kind == 'Y':
blob += np.ascontiguousarray(a, np.uint8).tobytes()
elif kind == 'U':
blob += np.ascontiguousarray(a, np.uint16).tobytes()
else:
blob += np.ascontiguousarray(a, np.float16).tobytes()
defines.append(f'#define W_{name.upper()} {off}u // {kind} {a.shape}')
open(out_bin, 'wb').write(bytes(blob))
open(out_h, 'w').write('// generated by engine/rr/dec0_pack.py\n' + '\n'.join(defines) + '\n')
print(f'{len(blob)} bytes, {len(e)} entries')
if __name__ == '__main__':
main()
engine/rr/dec0_ref.py
+258
View File
@@ -0,0 +1,258 @@
#!/usr/bin/env python3
"""Exact numpy reference of the denoiser's decoder 0 (cuda_dldn_engine_swin_dec0_kernel), one 8x8 window at a time,
from its weights and real inputs only: the specification the native shader implements. Derived like
engine/rr/enc1_ref.py, phase by phase against ZLUDA replays (engine/rr/phase_dump.py).
expand the low-resolution tokens under the window and its border (6 x 6 of dec1's 64-channel output, arena order)
-> 128 = 4 sub-positions x 32 channels (+ bias), FP8; sub-position 2 (y & 1) + (x & 1)
prologue key token = expanded + skip token (enc0's 32-channel output at argument 24), f16; RMS norm (enc1's per-word
tree on one plane) times the gain at W+8448; the residual is the token + the table at W+39232
attention two heads of 32 channels (K = 32, one k32 step), as enc1
MLP RMS norm, four branches of 32 -> 32 (+ bias, ramp activation) -> 32 onto the residual (+ table)
usage: dec0_ref.py CAPTURE FRAME WEIGHTS.npz [--blocks X,Y ...] (verifies the phases of a dump, see --dump)
dec0_ref.py CAPTURE FRAME OUT.npz --extract DUMP.npz (weights from a phase dump)
"""
import argparse
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rrmath import e4m3, f16, q8 # noqa: E402
from enc1_ref import inter, mm, softmax_cubic, act, mirror, add, H, F # noqa: E402
I32 = inter(32)
UN32 = [I32.index(i) for i in range(32)]
I96 = inter(96)
def norm_plane(xn, gain):
"""RMS norm of one 32-channel token (natural order) in NVIDIA's tree over the arena order; returns codes in k order"""
xa = xn[I32]
sq = (xa * xa).astype(H)
h = []
for t in range(4):
def pair(w):
return add(sq[4 * w:4 * w + 2], sq[4 * w + 2:4 * w + 4])
h.append(add(pair(t), pair(t + 4)))
h01 = [add(h[t], h[t ^ 1]) for t in range(4)]
h2 = add(h01[0], h01[2])
s = H(H(h2[0] + h2[1]) + H(2.0 ** -13))
r = H(F(1) / np.sqrt(F(s)))
return q8((xn * (r * gain).astype(H)).astype(H))[I32]
def keys(x0, y0):
border = ([(x, y0 - 1) for x in range(x0, x0 + 8)] + [(x, y0 + 8) for x in range(x0, x0 + 8)]
+ [(x0 - 1, y) for y in range(y0, y0 + 8)] + [(x0 + 8, y) for y in range(y0, y0 + 8)])
return [(x, y) for y in range(y0, y0 + 8) for x in range(x0, x0 + 8)] + border
def window(w, lo, skip, x0, y0, ex, ey, dbg=None):
"""lo: dec1's output [2][ey/2][ex/2][32] codes; skip: enc0's [ey][ex][32] codes. Returns the residual stream after
the block (64 x 32 f16, natural channels)"""
dbg = {} if dbg is None else dbg
hx, hy = ex // 2, ey // 2
X0, Y0 = x0 // 2 - 1, y0 // 2 - 1
low = np.zeros((48, 64), np.uint8)
for r in range(36):
X, Y = mirror(X0 + r % 6, hx), mirror(Y0 + r // 6, hy)
low[r] = np.concatenate([lo[0, Y, X], lo[1, Y, X]])
Ex = q8(mm(low, w['we'], np.repeat(w['be'][None, :], 48, 0)))
dbg['p1_A'] = low
xs, ns = [], []
for (x, y) in keys(x0, y0):
xx, yy = mirror(x, ex), mirror(y, ey)
r = (yy // 2 - Y0) * 6 + (xx // 2 - X0) if False else (y // 2 - Y0) * 6 + (x // 2 - X0)
s = 2 * (y & 1) + (x & 1)
xn = (e4m3(Ex[r, 32 * s:32 * s + 32]).astype(H) + e4m3(skip[yy, xx][UN32]).astype(H)).astype(H)
xs.append(xn)
ns.append(norm_plane(xn, w['gain0']))
X = np.array(ns)
dbg['p2_A'] = X
R = None
for h in range(2):
qk = mm(X, w[f'wk{h}'], np.zeros((96, 32)))
v = mm(X, w[f'wv{h}'], np.zeros((96, 32)))
Q = q8(qk)[:, I32]
S = mm(Q[:64], Q.T, w[f'bias{h}'])
P = softmax_cubic(S)
O = mm(q8(P)[:, I96], q8(v)[I96], np.zeros((64, 32)))
C = add(np.array(xs[:64]), w['table0'][None, :]) if h == 0 else R
R = mm(q8(O)[:, I32], w[f'wo{h}'], C)
dbg[f'p{4 + 5 * h}_A'], dbg[f'p{5 + 5 * h}_A'], dbg[f'p{6 + 5 * h}_A'] = Q[:64], q8(P)[:, I96], q8(O)[:, I32]
dbg[f'p{6 + 5 * h}_C'], dbg[f'p{6 + 5 * h}_D'] = C, R
N = np.array([norm_plane(R[t], w['gain1']) for t in range(64)])
acc = add(R, w['table1'][None, :])
for b in range(4):
hid = act(mm(N, w[f'w1_{b}'], np.repeat(w[f'b1_{b}'][None, :], 64, 0)))
dbg[f'p{12 + 2 * b}_A'] = N
dbg[f'p{13 + 2 * b}_A'], dbg[f'p{13 + 2 * b}_C'] = q8(hid)[:, I32], acc
acc = mm(q8(hid)[:, I32], w[f'w2_{b}'], acc)
dbg[f'p{13 + 2 * b}_D'] = acc
return acc
OUT_CH = [12, 13, 18, 19, 20, 21, 22, 23] # the 8 FP8 channels of the output stage's input (argument 48)
OUT_CH_D = [12, 13, None, 18, 19, 20, 21, None] # preset D (..._dec0_diamond_wallaby_kernel): bytes 2 and 7 are zero
TANH_CH = [14, 15, 16, 17] # surface 144 (rgba8 snorm)
def fma16(a, b, c):
return H(np.float64(a) * np.float64(b) + np.float64(c))
def tanh_nv(t):
"""NVIDIA's tanh of the surface-144 channels: an odd polynomial below 0.6, 1 - 2 / (e^2x + 1) above, f32"""
t = F(t)
m = abs(t)
big = F(1) if m >= F(9.0109138) else F(np.float64(F(1) / F(np.exp2(F(m * F(2.8853900432586670))) + F(1))) * -2 + 1)
s = F(t * t)
p = F(0.01573968306183815)
for c in (-0.05230396240949631, 0.13315297663211823, -0.33332768082618713, 0.0):
p = F(np.float64(p) * np.float64(s) + np.float64(F(c)))
small = F(np.float64(p) * np.float64(t) + np.float64(t))
return np.copysign(big, t) if m >= F(0.6) else small
def bilinear88(t88, xi, yi):
"""texture 88 (half resolution, linear) at the full-resolution texel centre (xi, yi): weights 1/4 and 3/4,
clamped at the edge; f32 sum rounded to binary16 (NVIDIA's texture unit rounds a few of these differently)"""
def axis(i):
k = i // 2
return ((k - 1, 0.25), (k, 0.75)) if i % 2 == 0 else ((k, 0.75), (k + 1, 0.25))
h, w = t88.shape[:2]
v = np.zeros(4)
for (yy, wy) in axis(yi):
for (xx, wx) in axis(xi):
v += wx * wy * t88[min(max(yy, 0), h - 1), min(max(xx, 0), w - 1)].astype(np.float64)
return v.astype(H)
def outputs(acc, t88, t96, t104, x0, y0, ex, ey):
"""per window token: the 8 output codes, surface 144 (4 floats before the snorm store), surface 112 (3 f16)"""
codes, s144, s112 = [], [], []
for t in range(64):
x, y = x0 + t % 8, y0 + t // 8
a = acc[t].astype(H)
codes.append(q8(a)[OUT_CH])
s144.append([tanh_nv(a[c]) for c in TANH_CH])
lg = a[:9]
m01 = np.maximum(np.maximum(np.maximum(lg[0:2], lg[2:4]), lg[4:6]), lg[6:8])
m = max(max(m01[0], m01[1]), lg[8])
e = [np.exp2(F(H(lg[i] - m)) * F(1.4426950216293335)) for i in range(9)]
pair = [F(e[2 * k] + e[2 * k + 1]) for k in range(4)]
sm = F(pair[0] + F(0))
for k in (1, 2, 3):
sm = F(sm + pair[k])
sm = F(sm + e[8])
r = H(F(1) / sm)
wg = [H(H(e[i]) * r) for i in range(9)]
def cl(v, n):
return 1 if v < 0 else (n - 2 if v >= n else v)
xs, ys = [cl(x - 1, ex), cl(x, ex), cl(x + 1, ex)], [cl(y - 1, ey), cl(y, ey), cl(y + 1, ey)]
f = np.zeros(3, H)
k = 0
for yy in ys:
for xx in xs:
T = bilinear88(t88, xx, yy)
for c in range(3):
f[c] = fma16(T[c], wg[k], f[c])
k += 1
g = H(F(1) / (np.exp2(F(a[9]) * F(-1.4426950216293335)) + F(1)))
om = H(H(1) - g)
o = np.zeros(3, H)
for c in range(3):
if 0 <= x < ex and 0 <= y < ey:
o[c] = fma16(t96[y, x, c], om, H(f[c] * g))
if not np.isnan(t104[y, x, c]):
o[c] = H(o[c] + t104[y, x, c])
s112.append(o)
return np.array(codes), np.array(s144), np.array(s112)
def extract(dump, weights, out):
z = np.load(dump)
W = np.frombuffer(weights, np.uint8)
def B(p):
return z[f'p{p}_w0_B']
def C(p):
return f16(z[f'p{p}_w0_C']).astype(H)
w = {'we': B(1), 'be': C(1)[0], 'gain0': W[8448:8448 + 64].view(H), 'table0': W[39232:39232 + 64].view(H),
'gain1': W[39296:39296 + 64].view(H), 'table1': W[41472:41472 + 64].view(H)}
for h, base in ((0, 2), (1, 7)):
w[f'wk{h}'], w[f'wv{h}'] = B(base), B(base + 1)
w[f'bias{h}'] = C(base + 2)
w[f'wo{h}'] = B(base + 4)
for b in range(4):
w[f'w1_{b}'], w[f'b1_{b}'] = B(12 + 2 * b), C(12 + 2 * b)[0]
w[f'w2_{b}'] = B(13 + 2 * b)
np.savez_compressed(out, **w)
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('frame', type=int)
ap.add_argument('weights')
ap.add_argument('--extract', metavar='DUMP')
ap.add_argument('--dump', help='phase dump of --blocks[0] to verify phase by phase')
ap.add_argument('--blocks', nargs='*', default=['5,5'])
o = ap.parse_args()
from rr_capture import Capture
c = Capture(os.path.expanduser(o.capture))
l = c.find('swin_dec0', o.frame)
if o.extract:
extract(o.extract, l.weights, o.weights)
return
ex, ey = (int(v) for v in np.frombuffer(l.args[:8], np.uint32))
sx, sy = (int(v) for v in np.frombuffer(l.args[80:84], np.uint16))
ptr = {p.offset: p.delta for p in l.pointers}
pre = np.frombuffer(l.pre(l.arena_alloc), np.uint8)
lo = pre[ptr[8]:ptr[8] + 2 * (ex // 2) * (ey // 2) * 32].reshape(2, ey // 2, ex // 2, 32)
skip = pre[ptr[24]:ptr[24] + ex * ey * 32].reshape(ey, ex, 32)
w = dict(np.load(o.weights))
img = {i.offset: i for i in l.images}
def surf(off, dt=np.float16):
i = img[off]
return np.frombuffer(i.path.read_bytes(), dt).reshape(i.height, i.width, i.channels)
t88, t96, t104, r112, r144 = surf(88), surf(96), surf(104), surf(112), surf(144, np.int8)
post = np.frombuffer(l.post_buffer(c, l.arena_alloc), np.uint8)
rout = post[ptr[48]:ptr[48] + ex * ey * 8].reshape(ey, ex, 8)
for blk in o.blocks:
bx, by = (int(v) for v in blk.split(','))
d = {}
x0, y0 = 8 * bx - sx, 8 * by - sy
acc = window(w, lo, skip, x0, y0, ex, ey, d)
codes, s144, s112 = outputs(acc, t88, t96, t104, x0, y0, ex, ey)
if 'diamond' in l.kernel:
codes = np.array([[0 if ch is None else q8(acc[t].astype(H))[ch] for ch in OUT_CH_D] for t in range(64)])
n = ok_c = ok_s = ok_t = 0
for t in range(64):
x, y = x0 + t % 8, y0 + t // 8
if not (0 <= x < ex and 0 <= y < r144.shape[0]):
continue
n += 1
ok_c += np.array_equal(codes[t], rout[y, x])
ok_t += np.array_equal(np.clip(np.round(np.array(s144[t], np.float64) * 127), -127, 127).astype(np.int8), r144[y, x])
ok_s += np.array_equal(s112[t], r112[y, x, :3])
print(f'block {bx},{by}: output codes {ok_c}/{n}, surface 144 {ok_t}/{n}, surface 112 {ok_s}/{n}')
if o.dump:
z = np.load(o.dump)
for k in sorted(d, key=lambda s: (int(s[1:].split('_')[0]), s)):
ref = z[f'{k.split("_")[0]}_w0_{k.split("_")[1]}']
ours = np.asarray(d[k])
ours = ours.view(np.uint16) if ours.dtype == H else ours
n = min(len(ref), len(ours))
print(f'{k}: rows equal {int((ref[:n] == ours[:n]).all(1).sum())}/{n}')
o.dump = None
if __name__ == '__main__':
main()
engine/rr/dec1_m.comp
+295
View File
@@ -0,0 +1,295 @@
#version 460
// Native decoder 1 of the DLSS-D (preset E) denoiser, register-resident: cuda_dldn_engine_swin_dec1_kernel.
// engine/rr/dec1_ref.py's arithmetic (bit-exact to NVIDIA except surface 112's bilinear taps, see dec0_m.comp);
// binding 1 holds engine/rr/dec1_pack.py's packing. The expand and prologue are dec0_m.comp's (64 channels from dec2's
// 96), the block is enc1_m.comp's (whose header explains the layouts), the outputs are the tokens (argument 48) and
// dec0_m.comp's filter on the 64 -> 16 head.
#extension GL_KHR_cooperative_matrix : require
#extension GL_KHR_memory_scope_semantics : require
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_8bit_storage : require
#extension GL_EXT_float_e4m3 : require
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_EXT_control_flow_attributes : require
#extension GL_GOOGLE_include_directive : require
layout(local_size_x = 128) in;
#include "dec1_m_w.h"
#include "rr_exact.glsl"
layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param;
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16;
layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u;
layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
layout(set = 0, binding = 3) uniform sampler2D T88;
layout(set = 0, binding = 4) uniform sampler2D T96;
layout(set = 0, binding = 5) uniform sampler2D T104;
layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S112;
#define SKIP 0
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define HAC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CM gl_CooperativeMatrixLayoutColumnMajor
#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM)
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
#define Q8(dst, src) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (dst)[i_] = q8s((src)[i_]);
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) }
const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0;
// shared memory, aliased by phase: expand [LOW | EX], attention [KS | VS], outputs [TAP | OUTF]
shared uint POOL[48 * 24 + 48 * 64];
#define LOW POOL // 48 low-resolution tokens (6 x 6, padded): 96 FP8 channels, arena order (24 words)
#define EX_AT (48u * 24u) // their expansion: 4 sub-positions x 64 FP8 channels in slot order (64 words)
#define KS POOL // the current head's K: token slot x 32 FP8 channels
#define VS_AT 768u // the current head's V^T: 32 channels x 96 token slots
#define RR_KS POOL
#define RR_KS0 0u
#define RR_VS POOL
#define RR_VS0 VS_AT
shared uint XT[96 * 16]; // the normalised key tokens: slot-major, 64 FP8 channels in slot order
shared float16_t TAP[100 * 4]; // texture 88 bilinear at the window's 10 x 10 filter tap positions
shared float16_t OUTF[64 * 16]; // the head per window token (raster order)
ivec2 keyXY(uint t, int x0, int y0)
{
if (t < 64u) return ivec2(x0 + int(t & 7u), y0 + int(t >> 3u));
const int i = int(t & 7u), side = int((t - 64u) >> 3u);
return ivec2(side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8), side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i));
}
int clampTap(int v, int n) { return v < 0 ? 1 : (v >= n ? n - 2 : v); }
uint pack4(floate4m3_t a, floate4m3_t b, floate4m3_t c, floate4m3_t d)
{
return code(a) | (code(b) << 8u) | (code(c) << 16u) | (code(d) << 24u);
}
float16_t e4(uint w, uint b) { return float16_t(uintBitsToFloate4m3EXT(uint8_t((w >> (8u * b)) & 0xffu))); }
void main()
{
const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID;
const uint hl = lane >> 4u, col = lane & 15u;
RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u);
const int ex = int(param.p[0]), ey = int(param.p[1]), hx = ex / 2, hy = ey / 2;
const uint loOff = param.p[2], skOff = param.p[6], outOff = param.p[12], shift = param.p[20];
const int x0 = 8 * int(gl_WorkGroupID.x) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u);
const int X0 = (x0 >> 1) - 1, Y0 = (y0 >> 1) - 1;
const uint lplane = uint(hx * hy) * 32u, plane = uint(ex * ey) * 32u;
// ---------------------------------------------------------------- expand: wave w computes sub-position w
for (uint i = tid; i < 48u * 24u; i += 128u) {
const uint r = i / 24u, wd = i % 24u;
uint u = 0u;
if (r < 36u) {
const int X = mirror(X0 + int(r % 6u), hx), Y = mirror(Y0 + int(r / 6u), hy);
u = arena.w[(loOff + (wd >> 3u) * lplane + uint(Y * hx + X) * 32u) / 4u + (wd & 7u)];
}
LOW[i] = u;
}
barrier();
[[unroll]] for (uint n = 0u; n < 3u; ++n) {
B8 lb[6];
[[unroll]] for (uint kt = 0u; kt < 6u; ++kt) coopMatLoad(lb[kt], POOL, n * 384u + 4u * kt, 24u, CM);
[[unroll]] for (uint mm = 0u; mm < 4u; ++mm) {
const uint m = 4u * wv + mm;
ACC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_BE / 2u + 16u * m + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < 6u; kt += 2u) {
A8 w0, w1;
LDA(w0, W_WE, 96u, m, kt); LDA(w1, W_WE, 96u, m, kt + 1u);
K32(a, w0, lb[kt], w1, lb[kt + 1u])
}
B8 q;
Q8(q, a)
coopMatStore(q, POOL, EX_AT + n * 1024u + 4u * m, 64u, CM);
}
}
barrier();
// ---------------------------------------------------------------- prologue: one key token per thread
if (tid < 96u) {
const ivec2 k = keyXY(inter(tid), x0, y0);
const uint r = uint(((k.y >> 1) - Y0) * 6 + ((k.x >> 1) - X0)), s = uint(2 * (k.y & 1) + (k.x & 1));
const int xx = mirror(k.x, ex), yy = mirror(k.y, ey);
const uint sb = (skOff + uint(yy * ex + xx) * 32u) / 4u;
float16_t v[64];
[[unroll]] for (uint wd = 0u; wd < 16u; ++wd) {
const uint e = POOL[EX_AT + r * 64u + 16u * s + wd], sk = arena.w[sb + (wd >> 3u) * (plane >> 2u) + (wd & 7u)];
[[unroll]] for (uint b = 0u; b < 4u; ++b) { precise float16_t x = e4(e, b) + e4(sk, b); v[4u * wd + b] = x; }
}
// enc1's input norm: per word t (and t + 4) the f16x2 pair sums of both planes, then lanes ^1, ^2, the halves
precise f16vec2 h[4];
[[unroll]] for (uint q = 0u; q < 4u; ++q) {
precise f16vec2 p[4];
[[unroll]] for (uint kk = 0u; kk < 4u; ++kk) {
const uint o = 32u * (kk & 1u) + 4u * (q + 4u * (kk >> 1u));
precise f16vec2 a = f16vec2(v[o], v[o + 1u]) * f16vec2(v[o], v[o + 1u]);
precise f16vec2 c = f16vec2(v[o + 2u], v[o + 3u]) * f16vec2(v[o + 2u], v[o + 3u]);
p[kk] = a + c;
}
precise f16vec2 s0 = p[0] + p[1], s1 = p[2] + p[3];
h[q] = s0 + s1;
}
precise f16vec2 h0 = h[0] + h[1], h2_ = h[2] + h[3];
precise f16vec2 hs = h0 + h2_;
precise float16_t sm = hs.x + hs.y;
precise float16_t se = sm + float16_t(0.0001220703125);
const float16_t rr = rsq16(se);
[[unroll]] for (uint wd = 0u; wd < 16u; ++wd) {
floate4m3_t c[4];
[[unroll]] for (uint b = 0u; b < 4u; ++b) {
const uint j = 4u * wd + b;
precise float16_t g = rr * w16.w[W_GAIN0 / 2u + inter(j)];
precise float16_t n = v[j] * g;
c[b] = q8h(n);
}
XT[tid * 16u + wd] = pack4(c[0], c[1], c[2], c[3]);
}
}
barrier();
// ---------------------------------------------------------------- the residual: expanded + skip + table 7
ACC R[4];
{
const uint tq = inter(16u * wv + col);
const int qx = x0 + int(tq & 7u), qy = y0 + int(tq >> 3u);
const uint r = uint(((qy >> 1) - Y0) * 6 + ((qx >> 1) - X0)), s = uint(2 * (qy & 1) + (qx & 1));
const uint sb = (skOff + uint(mirror(qy, ey) * ex + mirror(qx, ex)) * 32u) / 4u;
[[unroll]] for (uint m = 0u; m < 4u; ++m) {
const uint ea = EX_AT + r * 64u + 16u * s + 4u * m + 2u * hl;
const uint sa = sb + (m >> 1u) * (plane >> 2u) + 4u * (m & 1u) + 2u * hl;
const uvec2 e = uvec2(POOL[ea], POOL[ea + 1u]), sk = uvec2(arena.w[sa], arena.w[sa + 1u]);
[[unroll]] for (int i = 0; i < 8; ++i) {
precise float16_t x = e4(e[i >> 2], uint(i) & 3u) + e4(sk[i >> 2], uint(i) & 3u);
precise float16_t c = x + w16.w[W_TABLE7 / 2u + 16u * m + 8u * hl + uint(i)];
R[m][i] = float(c);
}
}
}
barrier(); // EX is read; the attention reuses its memory
// ---------------------------------------------------------------- attention, two heads
[[unroll]] for (uint hd = 0u; hd < ((SKIP & 64) != 0 ? 0u : 2u); ++hd) {
// this head's Q/K and V, shared: 12 jobs (token tile n, channel tile c), three per wave. K = (QK^T)^T rows
// to KS (token slot x 32 channels, also Q^T read column-major), V^T to VS (32 channels x 96 token slots)
if (hd == 1u) barrier(); // every wave is done with head 0's KS / VS
for (uint j = 0u; j < ((SKIP & 128) != 0 ? 0u : 3u); ++j) {
const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u;
A8 wk[4];
B8 wvb[4], xb[4];
A8 xa[4];
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) {
coopMatLoad(xb[kt], XT, n * 256u + 4u * kt, 16u, CM);
coopMatLoad(xa[kt], XT, n * 256u + 4u * kt, 16u, RM);
LDA(wk[kt], W_WK0 + HEAD * hd, 64u, c, kt);
LDB(wvb[kt], W_WV0 + HEAD * hd, 32u, kt, c);
}
ACC a = ACC(0.0);
K32(a, wk[0], xb[0], wk[1], xb[1])
K32(a, wk[2], xb[2], wk[3], xb[3])
A8 ka;
Q8(ka, a)
coopMatStore(ka, KS, n * 128u + 4u * c, 8u, RM);
ACC v = ACC(0.0);
K32(v, xa[0], wvb[0], xa[1], wvb[1])
K32(v, xa[2], wvb[2], xa[3], wvb[3])
A8 va;
Q8(va, v)
coopMatStore(va, POOL, VS_AT + c * 384u + 4u * n, 24u, RM);
}
barrier();
#include "rr_attn.glsl"
[[unroll]] for (uint m = 0u; m < 4u; ++m) {
A8 a0, a1;
LDA(a0, W_WO0 + HEAD * hd, 32u, m, 0u); LDA(a1, W_WO0 + HEAD * hd, 32u, m, 1u);
K32(R[m], a0, Ob[0], a1, Ob[1])
}
}
// ---------------------------------------------------------------- MLP
B8 Nb[4];
{
float16_t rv[4][8];
[[unroll]] for (uint m = 0u; m < 4u; ++m) [[unroll]] for (int i = 0; i < 8; ++i) rv[m][i] = float16_t(R[m][i]);
// slot 16m + 8hl + i is NVIDIA's channel 8j + 2t + c with j = 4(m >> 1) + (m & 1) + 2((i >> 1) & 1),
// t = 2hl + (i >> 2), c = i & 1 (enc1_ref.nv_rowsum_sq's tree)
float16_t h1[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
float16_t ht[2];
[[unroll]] for (uint tq = 0u; tq < 2u; ++tq) {
const uint i0 = 4u * tq + c, i1 = i0 + 2u;
precise float16_t a = rv[0][i0] * rv[0][i0] + rv[0][i1] * rv[0][i1];
precise float16_t b = rv[1][i0] * rv[1][i0] + rv[1][i1] * rv[1][i1];
precise float16_t cc = rv[2][i0] * rv[2][i0] + rv[2][i1] * rv[2][i1];
precise float16_t d = rv[3][i0] * rv[3][i0] + rv[3][i1] * rv[3][i1];
precise float16_t hv = (a + cc) + (b + d);
ht[tq] = hv;
}
precise float16_t hh = ht[0] + ht[1];
h1[c] = hh;
}
const uint other = subgroupShuffleXor(bh(h1[0]) | (bh(h1[1]) << 16u), 16u);
precise float16_t g0 = h1[0] + hb(other & 0xffffu), g1 = h1[1] + hb(other >> 16u);
precise float16_t s = g0 + g1;
precise float16_t se = s + float16_t(0.0001220703125);
const float16_t r = rsq16(se);
[[unroll]] for (uint m = 0u; m < 4u; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) {
const uint k = 16u * m + 8u * hl + uint(i);
precise float16_t g = r * w16.w[W_GAIN14 / 2u + k];
precise float16_t n = rv[m][i] * g;
Nb[m][i] = q8h(n);
precise float16_t a = rv[m][i] + w16.w[W_TABLE16 / 2u + k];
R[m][i] = float(a);
}
}
for (uint br = 0u; br < ((SKIP & 4) != 0 ? 1u : 8u); ++br) {
const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br;
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
ACC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = float(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]);
A8 a0, a1, a2, a3;
LDA(a0, w1, 64u, c, 0u); LDA(a1, w1, 64u, c, 1u); LDA(a2, w1, 64u, c, 2u); LDA(a3, w1, 64u, c, 3u);
K32(h, a0, Nb[0], a1, Nb[1])
K32(h, a2, Nb[2], a3, Nb[3])
[[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = q8h(act(float16_t(h[i])));
}
[[unroll]] for (uint m = 0u; m < 4u; ++m) {
A8 a0, a1;
LDA(a0, w2, 32u, m, 0u); LDA(a1, w2, 32u, m, 1u);
K32(R[m], a0, Hb[0], a1, Hb[1])
}
}
// ---------------------------------------------------------------- outputs
{
const uint t = inter(16u * wv + col);
const int tx = x0 + int(t & 7u), ty = y0 + int(t >> 3u);
const bool inside = tx >= 0 && ty >= 0 && tx < ex && ty < ey;
B8 Tb[4];
const uint o = (outOff + uint(ty * ex + tx) * 32u) >> 2u;
[[unroll]] for (uint m = 0u; m < 4u; ++m) {
Q8(Tb[m], R[m])
const uint at = 4u * (m & 1u) + 2u * hl;
if (inside) {
arena.w[o + (m >> 1u) * (plane >> 2u) + at] = pack4(Tb[m][0], Tb[m][1], Tb[m][2], Tb[m][3]);
arena.w[o + (m >> 1u) * (plane >> 2u) + at + 1u] = pack4(Tb[m][4], Tb[m][5], Tb[m][6], Tb[m][7]);
}
}
// the 64 -> 16 head, per token to OUTF
ACC hv;
[[unroll]] for (int i = 0; i < 8; ++i) hv[i] = float(w16.w[W_BH / 2u + 8u * hl + uint(i)]);
A8 a0, a1, a2, a3;
LDA(a0, W_WH, 64u, 0u, 0u); LDA(a1, W_WH, 64u, 0u, 1u); LDA(a2, W_WH, 64u, 0u, 2u); LDA(a3, W_WH, 64u, 0u, 3u);
K32(hv, a0, Tb[0], a1, Tb[1])
K32(hv, a2, Tb[2], a3, Tb[3])
[[unroll]] for (int i = 0; i < 8; ++i) OUTF[t * 16u + 8u * hl + uint(i)] = float16_t(hv[i]);
}
barrier();
#include "rr_decfilter.glsl"
}
engine/rr/dec1_pack.py
+60
View File
@@ -0,0 +1,60 @@
#!/usr/bin/env python3
"""Pack decoder 1's extracted weights (engine/rr/dec1_ref.py --extract) for engine/rr/dec1_m.comp.
usage: dec1_pack.py WEIGHTS.npz OUT.bin OUT.h
enc1_pack.py's layout (transposed products, columns in the consumer's k order) plus the expand (dec0_pack.py): 256
columns, four sub-positions of 64 channels each permuted to slot order.
"""
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from enc1_ref import inter # noqa: E402
from enc1_pack import a_tiles, b_tiles, cubic_table # noqa: E402
I32, I64, I96 = inter(32), inter(64), inter(96)
I256 = [64 * (j // 64) + I64[j % 64] for j in range(256)]
def main():
src, out_bin, out_h = sys.argv[1:4]
w = dict(np.load(src))
e = {'we': ('A', w['we'][:, I256]), 'be': ('H', w['be'][I256])}
for h in range(2):
e[f'wk{h}'] = ('A', np.vstack([w[f'wk{h}a'], w[f'wk{h}b']])[:, I32])
e[f'wv{h}'] = ('B', np.vstack([w[f'wv{h}a'], w[f'wv{h}b']])[:, I32])
e[f'wo{h}'] = ('A', w[f'wo{h}'][:, I64])
e[f'bias{h}'] = ('H', w[f'bias{h}'][np.ix_(I64, I96)].T)
for b in range(8):
e[f'w1_{b}'] = ('A', w[f'w1_{b}'][:, I32])
e[f'b1_{b}'] = ('H', w[f'b1_{b}'][I32])
e[f'w2_{b}'] = ('A', w[f'w2_{b}'][:, I64])
e['wh'], e['bh'] = ('A', w['wh']), ('H', w['bh'])
e['gain0'] = ('H', w['gain0'])
for k in ('table7', 'gain14', 'table16'):
e[k] = ('H', w[k][I64])
e['cubic'] = ('U', cubic_table())
blob, defines = bytearray(), []
for name, (kind, a) in e.items():
while len(blob) % 256:
blob.append(0)
off = len(blob)
if kind == 'A':
blob += a_tiles(a).tobytes()
elif kind == 'B':
blob += b_tiles(a).tobytes()
elif kind == 'U':
blob += np.ascontiguousarray(a, np.uint16).tobytes()
else:
blob += np.ascontiguousarray(a, np.float16).tobytes()
defines.append(f'#define W_{name.upper()} {off}u // {kind} {a.shape}')
open(out_bin, 'wb').write(bytes(blob))
open(out_h, 'w').write('// generated by engine/rr/dec1_pack.py\n' + '\n'.join(defines) + '\n')
print(f'{len(blob)} bytes, {len(e)} entries')
if __name__ == '__main__':
main()
engine/rr/dec1_ref.py
+183
View File
@@ -0,0 +1,183 @@
#!/usr/bin/env python3
"""Exact numpy reference of the denoiser's decoder 1 (cuda_dldn_engine_swin_dec1_kernel), one 8x8 window at a time.
Derived like engine/rr/enc1_ref.py, phase by phase against ZLUDA replays (engine/rr/phase_dump.py); NVIDIA runs it with
two warps, each one head and half of the expand's columns.
expand the 6 x 6 level-2 tokens under the window and its border (dec2's 96 channels, arena order) -> 256 =
4 sub-positions x 64 channels (+ bias), FP8; sub-position 2 (y & 1) + (x & 1) (dec0_ref)
prologue key token = expanded + skip token (enc1's 64-channel output at argument 24), f16; enc1's input norm (two
planes) times the gain at W+25088; the residual is the token + the table at W+62080
block enc1's: two 32-channel heads, eight MLP branches 64 -> 32 -> 64 (gain W+62208, table W+66496)
outputs the tokens (argument 48); the 64 -> 16 head and dec0's filter (bilinear texture 88 from level 2, gate,
textures 96 / 104) to surface 112
usage: dec1_ref.py CAPTURE FRAME WEIGHTS.npz [--blocks X,Y ...] (verifies against the capture)
dec1_ref.py CAPTURE FRAME OUT.npz --extract DUMP.npz (weights from a phase dump)
"""
import argparse
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rrmath import e4m3, f16, q8 # noqa: E402
from enc1_ref import inter, mm, softmax_cubic, act, mirror, add, input_norm, nv_rowsum_sq, H, F # noqa: E402
from dec0_ref import keys, fma16, bilinear88 # noqa: E402
I32, I64, I96 = inter(32), inter(64), inter(96)
UN64 = [I64.index(i) for i in range(64)]
GAIN0, TABLE7, GAIN14, TABLE16 = 25088, 62080, 62208, 66496
def window(w, lo, skip, x0, y0, ex, ey):
"""lo: dec2's output [3][ey/2][ex/2][32] codes; skip: enc1's [2][ey][ex][32]. Returns (tokens 64 x 64 codes in arena
order, head 64 x 16 f16)"""
hx, hy = ex // 2, ey // 2
X0, Y0 = x0 // 2 - 1, y0 // 2 - 1
low = np.zeros((48, 96), np.uint8)
for r in range(36):
X, Y = mirror(X0 + r % 6, hx), mirror(Y0 + r // 6, hy)
low[r] = np.concatenate([lo[p, Y, X] for p in range(3)])
Ex = q8(mm(low, w['we'], np.repeat(w['be'][None, :], 48, 0)))
xs, ns = [], []
for (x, y) in keys(x0, y0):
xx, yy = mirror(x, ex), mirror(y, ey)
r = (y // 2 - Y0) * 6 + (x // 2 - X0)
s = 2 * (y & 1) + (x & 1)
sk = np.concatenate([skip[0, yy, xx], skip[1, yy, xx]])[UN64]
xn = (e4m3(Ex[r, 64 * s:64 * s + 64]).astype(H) + e4m3(sk).astype(H)).astype(H)
xs.append(xn)
xa = xn[I64]
n, _ = input_norm([xa[:32], xa[32:]], w['gain0'])
ns.append(q8(n)[I64])
X = np.array(ns)
res = np.array(xs[:64])
R = None
for h in range(2):
qk = mm(X[:, 32:], w[f'wk{h}b'], mm(X[:, :32], w[f'wk{h}a'], np.zeros((96, 32))))
v = mm(X[:, 32:], w[f'wv{h}b'], mm(X[:, :32], w[f'wv{h}a'], np.zeros((96, 32))))
Q = q8(qk)[:, I32]
S = mm(Q[:64], Q.T, w[f'bias{h}'])
P = softmax_cubic(S)
O = mm(q8(P)[:, I96], q8(v)[I96], np.zeros((64, 32)))
C = add(res, w['table7'][None, :]) if h == 0 else R
R = mm(q8(O)[:, I32], w[f'wo{h}'], C)
N = q8((R * (nv_rowsum_sq(R)[:, None] * w['gain14'][None, :]).astype(H)).astype(H))[:, I64]
acc = add(R, w['table16'][None, :])
for b in range(8):
hid = act(mm(N, w[f'w1_{b}'], np.repeat(w[f'b1_{b}'][None, :], 64, 0)))
acc = mm(q8(hid)[:, I32], w[f'w2_{b}'], acc)
tokens = q8(acc)[:, I64]
head = mm(tokens, w['wh'], np.repeat(w['bh'][None, :], 64, 0))
return tokens, head
def filter112(head, t88, t96, t104, x0, y0, ex, ey):
"""dec0_ref.outputs' surface 112 on the head's channels 0-8 (logits) and 9 (gate)"""
out = []
for t in range(64):
x, y = x0 + t % 8, y0 + t // 8
a = head[t].astype(H)
lg = a[:9]
m01 = np.maximum(np.maximum(np.maximum(lg[0:2], lg[2:4]), lg[4:6]), lg[6:8])
m = max(max(m01[0], m01[1]), lg[8])
e = [np.exp2(F(H(lg[i] - m)) * F(1.4426950216293335)) for i in range(9)]
pair = [F(e[2 * k] + e[2 * k + 1]) for k in range(4)]
sm = F(pair[0] + F(0))
for k in (1, 2, 3):
sm = F(sm + pair[k])
sm = F(sm + e[8])
r = H(F(1) / sm)
wg = [H(H(e[i]) * r) for i in range(9)]
def cl(v, n):
return 1 if v < 0 else (n - 2 if v >= n else v)
f = np.zeros(3, H)
k = 0
for yy in (cl(y - 1, ey), cl(y, ey), cl(y + 1, ey)):
for xx in (cl(x - 1, ex), cl(x, ex), cl(x + 1, ex)):
T = bilinear88(t88, xx, yy)
for c in range(3):
f[c] = fma16(T[c], wg[k], f[c])
k += 1
g = H(F(1) / (np.exp2(F(a[9]) * F(-1.4426950216293335)) + F(1)))
om = H(H(1) - g)
o = np.zeros(3, H)
if 0 <= x < ex and 0 <= y < ey:
for c in range(3):
o[c] = fma16(t96[y, x, c], om, H(f[c] * g))
if not np.isnan(t104[y, x, c]):
o[c] = H(o[c] + t104[y, x, c])
out.append(o)
return np.array(out)
def extract(dump, weights, out):
z = np.load(dump)
W = np.frombuffer(weights, np.uint8)
def B(p, wv=0):
return z[f'p{p}_w{wv}_B']
def C(p, wv=0):
return f16(z[f'p{p}_w{wv}_C']).astype(H)
def table(off, n=64):
return W[off:off + 2 * n].view(H)
w = {'we': np.hstack([B(1, 0), B(1, 1)]), 'be': np.concatenate([C(1, 0)[0], C(1, 1)[0]]),
'gain0': table(GAIN0), 'table7': table(TABLE7), 'gain14': table(GAIN14), 'table16': table(TABLE16)}
for h in range(2):
w[f'wk{h}a'], w[f'wv{h}a'], w[f'wk{h}b'], w[f'wv{h}b'] = B(2, h), B(3, h), B(4, h), B(5, h)
w[f'bias{h}'] = C(6, h)
w[f'wo{h}'] = B(8 + h)
for b in range(8):
w[f'w1_{b}'], w[f'b1_{b}'] = B(10 + 2 * b), C(10 + 2 * b)[0]
w[f'w2_{b}'] = B(11 + 2 * b)
w['wh'], w['bh'] = B(26), C(26)[0]
np.savez_compressed(out, **w)
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('frame', type=int)
ap.add_argument('weights')
ap.add_argument('--extract', metavar='DUMP')
ap.add_argument('--blocks', nargs='*', default=['5,5'])
o = ap.parse_args()
from rr_capture import Capture
c = Capture(os.path.expanduser(o.capture))
l = c.find('swin_dec1', o.frame)
if o.extract:
extract(o.extract, l.weights, o.weights)
return
ex, ey = (int(v) for v in np.frombuffer(l.args[:8], np.uint32))
sx, sy = (int(v) for v in np.frombuffer(l.args[80:84], np.uint16))
ptr = {p.offset: p.delta for p in l.pointers}
pre = np.frombuffer(l.pre(l.arena_alloc), np.uint8)
post = np.frombuffer(l.post_buffer(c, l.arena_alloc), np.uint8)
lo = pre[ptr[8]:ptr[8] + 3 * (ex // 2) * (ey // 2) * 32].reshape(3, ey // 2, ex // 2, 32)
skip = pre[ptr[24]:ptr[24] + 2 * ex * ey * 32].reshape(2, ey, ex, 32)
outp = post[ptr[48]:ptr[48] + 2 * ex * ey * 32].reshape(2, ey, ex, 32)
img = {i.offset: i for i in l.images}
def surf(off):
i = img[off]
return np.frombuffer(i.path.read_bytes(), np.float16).reshape(i.height, i.width, 4)
t88, t96, t104, r112 = surf(88), surf(96), surf(104), surf(112)
w = dict(np.load(o.weights))
for blk in o.blocks:
bx, by = (int(v) for v in blk.split(','))
x0, y0 = 8 * bx - sx, 8 * by - sy
tokens, head = window(w, lo, skip, x0, y0, ex, ey)
s112 = filter112(head, t88, t96, t104, x0, y0, ex, ey)
inside = [t for t in range(64) if 0 <= x0 + t % 8 < ex and 0 <= y0 + t // 8 < ey]
tk = sum(np.array_equal(tokens[t], np.concatenate([outp[0, y0 + t // 8, x0 + t % 8], outp[1, y0 + t // 8, x0 + t % 8]]))
for t in inside)
sf = sum(np.array_equal(s112[t], r112[y0 + t // 8, x0 + t % 8, :3]) for t in inside)
print(f'block {bx},{by}: tokens {tk}/{len(inside)}, surface 112 {sf}/{len(inside)}')
if __name__ == '__main__':
main()
engine/rr/dll_weights.py
+320
View File
@@ -0,0 +1,320 @@
#!/usr/bin/env python3
"""Read the Ray Reconstruction network's weights straight from the user's nvngx_dlssd.dll.
usage: dll_weights.py DLL --list
dll_weights.py DLL --kernels
dll_weights.py DLL OUT_DIR (--model N | --capture CAPTURE_DIR) [--check EXTRACTED_DIR] [--any-kernels]
The DLL holds every model's tensors as plain binary16 arrays in natural ([out][in]) order, one stage after the other
(encoder 0 ... decoder 0), with nothing naming them. They are found by structure alone: every attention head has a
relative position bias table of 100 x 100 values (the 10 x 10 window with its border) whose entries depend only on the
offset between the two positions, and a model is eleven consecutive stages whose head counts are 2 2 4 4 8 8 8 4 4 2 2.
The other tensors of a stage sit at fixed distances from its tables (LAYOUT below).
OUT_DIR/STAGE.npz gets the tensors in the form engine/rr/{swin,enc0,dec0}_ref.py --extract produced from a ZLUDA
phase dump, which is what the packers (swin_pack.py, enc0_pack.py, dec0_pack.py) take:
FP8 matrices e4m3 codes, [k][n] with k in NVIDIA's fragment order (natural input channel inter(K)[k])
gains times sqrt(32), as NGX uploads them
bias [64 queries][96 keys] in the kernels' key order, times 2 sqrt(2), as NGX uploads it
--model picks the N-th model of the DLL (--list); --capture picks the one a bridge capture of the denoiser ran, by
looking up each stage's vectors in the weight image NGX uploaded (and fails if the capture is of another DLL).
--check compares every tensor with an extracted one (EXTRACTED_DIR/STAGE/weights.npz).
The native shaders were written against the kernels of one DLL version. Before anything is written the DLL's network
kernels (the PTX modules, engine/rr/native_kernels.txt: their SHA-256) are compared with those; another version is
refused unless --any-kernels is given. --kernels prints the digests of a DLL.
Everything read and written is NVIDIA's: it stays on the user's machine and is never packaged.
"""
import argparse
import os
import sys
from pathlib import Path
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rrmath import e4m3, q8 # noqa: E402
from enc1_ref import inter # noqa: E402
from dec0_ref import keys # noqa: E402
H = np.float16
# stage: (kind, C, C_LO (decoders: the channels one level down) or C_M (encoders: the merged channels), heads)
STAGES = {
'enc0': ('enc0', 32, 64, 2), 'enc1': ('enc', 64, 96, 2), 'enc2': ('enc', 96, 128, 4), 'enc3': ('enc', 128, 160, 4),
'enc4': ('enc', 160, 160, 8), 'enc5': ('enc', 160, 0, 8), 'dec4': ('dec', 160, 160, 8), 'dec3': ('dec', 128, 160, 4),
'dec2': ('dec', 96, 128, 4), 'dec1': ('dec', 64, 96, 2), 'dec0': ('dec0', 32, 64, 2),
}
ORDER = list(STAGES)
HEADS = [STAGES[s][3] for s in ORDER]
TABLE = 100 * 100 # one head's relative position bias, in values
# the 96 keys of a window (its 64 tokens, then the border without the corners) in the 10 x 10 table
POS = [(y + 1) * 10 + x + 1 for (x, y) in keys(0, 0)]
BIAS_SCALE = np.float32(2.0 * np.sqrt(2.0))
def bias_tables(u):
"""offsets (in values) of every relative position bias table in the binary16 view u of the DLL"""
# entry (i, j) equals entry (i + 1, j + 1) and (i + 10, j + 10): one step right / down for both positions
e1 = u[:-1111] == u[101:-1010]
e10 = u[:-1111] == u[1010:-101]
e11 = u[:-1111] == u[1111:]
first = e1[:-18] & e1[1:-17] & e1[2:-16] & e10[:-18] & e10[1:-17] & e11[:-18] & e11[1:-17] & (u[:len(e1) - 18] != u[1:len(e1) - 17])
found = []
dy, dx = np.divmod(np.arange(100), 10)
rel = ((dy[None, :] - dy[:, None] + 9) * 19 + dx[None, :] - dx[:, None] + 9).reshape(-1)
for p in np.nonzero(first)[0]:
if p + TABLE > len(u) or (found and p < found[-1] + TABLE):
continue
t = u[p:p + TABLE]
r = np.zeros(361, np.uint16)
r[rel] = t
if np.array_equal(r[rel], t) and len(np.unique(t)) > 64:
found.append(int(p))
return found
def models(dll):
"""[[bias offset (bytes) of each of the eleven stages] per model], in file order"""
u = np.frombuffer(dll[:len(dll) // 2 * 2], np.uint16)
tables = bias_tables(u)
runs = [] # (first table, count) of consecutive tables
for p in tables:
if runs and p == runs[-1][0] + TABLE * runs[-1][1]:
runs[-1][1] += 1
else:
runs.append([p, 1])
out, i = [], 0
while i + len(HEADS) <= len(runs):
if [n for _, n in runs[i:i + len(HEADS)]] == HEADS:
out.append([2 * p for p, _ in runs[i:i + len(HEADS)]])
i += len(HEADS)
else:
i += 1
return out
def kernel_digests(dll):
"""{(kernel, SHA-256 of its PTX module)} of the DLL's network kernels"""
import hashlib
import re
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..', 'kernels', 'tools'))
from extract_dlss_ptx import fatbins, ptx_entries
out = set()
for _, fatbin in fatbins(dll):
for text in ptx_entries(fatbin):
for name in re.findall(rb'\.entry (cuda_dldn_engine_swin_\w+)\(', text):
out.add((name.decode(), hashlib.sha256(text).hexdigest()))
return out
def check_kernels(dll, path):
"""the kernels the native shaders were written against must all be in the DLL, and no other version of them"""
known = {tuple(l.split()) for l in open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'native_kernels.txt'))
if l.strip() and not l.startswith('#')}
have = kernel_digests(dll)
missing, other = sorted(known - have), sorted(have - known)
if missing or other:
raise SystemExit(f'{path}: its network kernels are not the ones the native shaders were written against '
f'({len(missing)} missing, {len(other)} unknown; e.g. {(missing + other)[0][0]}). The native '
'stages are unverified for this DLL version: use the translated model (engine/compile_rr.py), or '
'--any-kernels to build anyway.')
def noise_table(dll):
"""preset E's encoder 0 dithers its features with 8 x 16 x 16 16-bit values, the first 4096 bytes of the weight
image NGX uploads. In the DLL they are the array of 2048 values whose low five bits are all clear or all set."""
u = np.frombuffer(dll[:len(dll) // 2 * 2], np.uint16)
low = u & 31
ok = np.concatenate([[False], (low == 0) | (low == 31), [False]])
edge = np.nonzero(ok[1:] != ok[:-1])[0]
found = []
for a, b in zip(edge[::2], edge[1::2]):
if b - a >= 2048:
t = u[b - 2048:b]
if len(np.unique(t)) > 1024 and (t & 31 == 31).sum() > 256:
found.append((2 * (b - 2048), t.tobytes()))
# 310.9.1 also holds a second, unreferenced array of this kind at an odd address; the table the code uses is aligned
if len({t for _, t in found}) > 1:
found = [(at, t) for at, t in found if at % 16 == 0]
return found[0][1] if len({t for _, t in found}) == 1 else None # some DLLs hold the same table twice
class Reader:
def __init__(self, dll, at):
self.dll, self.at = dll, at
def back(self, *shape):
self.at -= 2 * int(np.prod(shape))
return np.frombuffer(self.dll, H, int(np.prod(shape)), self.at).reshape(shape)
def take(self, *shape):
a = np.frombuffer(self.dll, H, int(np.prod(shape)), self.at).reshape(shape)
self.at += 2 * int(np.prod(shape))
return a
def fp8(nat):
"""a natural [out][in] binary16 matrix -> NVIDIA's B operand: codes [k][n], k in fragment order"""
k = nat.shape[1]
return np.ascontiguousarray(q8(nat)[:, inter(k) if k % 32 == 0 else np.arange(k)].T)
def gain(g, c):
return (g.astype(np.float32) * np.float32(np.sqrt(np.float32(c)))).astype(H)
def read_stage(dll, stage, bias_at):
"""the stage's tensors (see the module text); returns (dict, first byte, one past the last byte)"""
kind, C, CX, NH = STAGES[stage]
w = {}
r = Reader(dll, bias_at)
table_a = r.back(C)
wo = r.back(C, 32 * NH)
wv = r.back(32 * NH, C)
wk = r.back(32 * NH, C)
gain_a = r.back(C)
if kind == 'enc0': # the embedding of the 32 input features
w['b2'] = r.back(32).copy()
w['w2'] = fp8(r.back(32, 32))
w['b1'] = r.back(32).copy()
# [32][inputs]: preset E has 32 input features, preset D 30 (NGX pads the image's matrix with zeros). The
# matrix is the widest one that starts with a weight and holds FP8 values only
end = r.at
for inputs in range(32, 15, -1):
m = np.frombuffer(dll, H, 32 * inputs, end - 64 * inputs).reshape(32, inputs)
if m[0, :4].any() and np.array_equal(e4m3(q8(m)).astype(H).view(np.uint16), m.view(np.uint16)):
break
else:
raise SystemExit('enc0: no input embedding before its bias')
r.at = end - 64 * inputs
w['w1'] = fp8(np.hstack([m, np.zeros((32, 32 - inputs), H)]))
first = r.at
r = Reader(dll, bias_at)
bias = r.take(NH, 100, 100)
gain_b = r.take(C)
w1 = r.take(4 * C, C)
b1 = r.take(4 * C)
w2 = r.take(C, 4 * C)
table_b = r.take(C)
for h in range(NH):
w[f'wk{h}'] = fp8(wk[32 * h:32 * h + 32])
w[f'wv{h}'] = fp8(wv[32 * h:32 * h + 32])
w[f'wo{h}'] = fp8(wo[:, 32 * h:32 * h + 32])
w[f'bias{h}'] = (bias[h][np.ix_(POS[:64], POS)].astype(np.float32) * BIAS_SCALE).astype(H)
for b in range(C // 8):
w[f'w1_{b}'] = fp8(w1[32 * b:32 * b + 32])
w[f'b1_{b}'] = b1[32 * b:32 * b + 32].copy()
w[f'w2_{b}'] = fp8(w2[:, 32 * b:32 * b + 32])
names = ('gain0', 'table0', 'gain1', 'table1') if kind in ('enc0', 'dec0') else ('gain0', 'table7', 'gain14', 'table16')
for name, a in zip(names, (gain(gain_a, C), table_a.copy(), gain(gain_b, C), table_b.copy())):
w[name] = a
if kind in ('enc', 'enc0') and CX: # the 2x2 merge
w['wm'] = fp8(r.take(CX, 4 * C))
w['bm'] = r.take(CX).copy()
if kind in ('dec', 'dec0'): # the expand of the level below
w['we'] = fp8(r.take(4 * C, CX))
w['be'] = r.take(4 * C).copy()
if kind != 'dec0': # the 16-channel head of the stage's filter
w['wh'] = fp8(r.take(16, C))
w['bh'] = r.take(16).copy()
return w, first, r.at
def read_model(dll, bias_offsets):
"""{stage: tensors}; checks that the stages lie back to back in the DLL, which the layout implies"""
out, end = {}, None
for stage, at in zip(ORDER, bias_offsets):
out[stage], first, last = read_stage(dll, stage, at)
if end is not None and first != end:
raise SystemExit(f'{stage}: starts at {first}, the stage before ends at {end}: not the layout this reader knows')
end = last
return out
def in_capture(capture, model):
"""whether the capture's launches uploaded this model: every stage's vectors must be in its weight image"""
from rr_capture import Capture
cap = Capture(os.path.expanduser(capture))
for stage, w in model.items():
image = next(l for l in cap.launches if f'swin_{stage}' in l.kernel).weights
for name, a in w.items():
if a.dtype == H and a.ndim == 1 and bytes(image).find(a.tobytes()) < 0:
return False
return True
def main():
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument('dll')
ap.add_argument('output', nargs='?')
ap.add_argument('--list', action='store_true')
ap.add_argument('--model', type=int)
ap.add_argument('--capture')
ap.add_argument('--check', metavar='EXTRACTED_DIR')
ap.add_argument('--kernels', action='store_true')
ap.add_argument('--any-kernels', action='store_true')
o = ap.parse_args()
dll = Path(o.dll).read_bytes()
if o.kernels:
for name, digest in sorted(kernel_digests(dll)):
print(name, digest)
return
if o.output and not o.any_kernels:
check_kernels(dll, o.dll)
found = models(dll)
if not found:
raise SystemExit(f'{o.dll}: no Ray Reconstruction model with the layout this reader knows')
if o.list or not o.output:
for i, m in enumerate(found):
print(f'model {i}: bias tables at ' + ' '.join(map(str, m)))
return
if o.model is not None:
model = read_model(dll, found[o.model])
index = o.model
else:
if not o.capture:
raise SystemExit('--model N or --capture CAPTURE_DIR')
index, model = next(((i, m) for i, m in ((i, read_model(dll, f)) for i, f in enumerate(found))
if in_capture(o.capture, m)), (None, None))
if model is None:
raise SystemExit(f'none of the {len(found)} models in {o.dll} is the one the capture ran: the capture was taken '
'with another nvngx_dlssd.dll')
out = Path(o.output)
out.mkdir(parents=True, exist_ok=True)
for stage, w in model.items():
np.savez_compressed(out / f'{stage}.npz', **w)
noise = noise_table(dll)
if noise is not None:
if o.capture: # preset E's image starts with it; preset D's has none
from rr_capture import Capture
image = bytes(next(l for l in Capture(os.path.expanduser(o.capture)).launches if 'swin_enc0' in l.kernel).weights)
if 'diamond' not in next(l for l in Capture(os.path.expanduser(o.capture)).launches if 'swin_enc0' in l.kernel).kernel \
and image[:4096] != noise:
raise SystemExit('the noise table found in the DLL is not the one the capture uploaded')
(out / 'noise.bin').write_bytes(noise)
print(f'{o.dll}: model {index} of {len(found)}, {len(model)} stages in {out}')
if o.check:
bad = 0
for stage, w in model.items():
ref = dict(np.load(Path(o.check) / stage / 'weights.npz'))
for name in sorted(set(ref) | set(w)):
a, b = w.get(name), ref.get(name)
if a is None or b is None:
print(f'{stage}.{name}: only in the {"extraction" if a is None else "DLL reading"}')
bad += 1
continue
if name in ('we', 'be'): # a multi-warp extraction repeats the columns
b = b[..., :a.shape[-1]]
same = a.shape == b.shape and np.array_equal(a.view(np.uint8) if a.dtype == np.uint8 else a.astype(H).view(np.uint16),
b.view(np.uint8) if b.dtype == np.uint8 else b.astype(H).view(np.uint16))
if not same:
n = int((a != b).sum()) if a.shape == b.shape else -1
print(f'{stage}.{name}: differs ({a.shape} vs {b.shape}, {n} values)')
bad += 1
print('check: every tensor equals the extraction' if not bad else f'check: {bad} tensors differ')
sys.exit(1 if bad else 0)
if __name__ == '__main__':
main()
engine/rr/enc0_m_decl.glsl
+79
View File
@@ -0,0 +1,79 @@
// Native part of engine/rr/enc0_m.comp (generated by engine/rr/gen_enc0_m.py): declarations, after the translated
// header. Binding 1 holds NGX's weight image followed by engine/rr/enc0_pack.py's blob (offsets in enc0_m_w.h already
// include its base); binding 2 is the arena.
#extension GL_EXT_shader_8bit_storage : require
#extension GL_KHR_shader_subgroup_vote : require
#extension GL_EXT_control_flow_attributes : require
#extension GL_GOOGLE_include_directive : require
#include "enc0_m_w.h"
#include "rr_exact.glsl"
#ifndef RR_GFX11
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
#endif
layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16;
layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u;
layout(set = 0, binding = 1, std430) readonly buffer W8U { uint8_t w[]; } w8u;
layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define HAC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#ifndef RR_GFX11
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#endif
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CM gl_CooperativeMatrixLayoutColumnMajor
#ifdef RR_GFX11
#include "rr_gfx11_k32.glsl"
#else
#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM)
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
#define Q8(dst, src) Q8PAIRS(dst, src)
// One 32-deep step onto an accumulator that holds binary16 values. The reference (and the default here) adds the
// accumulator inside the f32 products and rounds the result to binary16. RR_MIX (build_native_model.sh with
// D4R_RR_MIX=1) keeps the accumulator as binary16, sums the products from zero and takes f16(P + acc), which the
// driver compiles to one v_fma_mix per element in place of two conversions: about 5% faster, but a single rounding
// where the reference has two, so a value can come out one binary16 step away (engine/rr/README.md).
// RAC is the accumulator type, RV(x) a binary16 value as its element; K32H is the step on a binary16 matrix.
#ifdef RR_MIX
#define RAC HAC
#define RV(x) (x)
#define K32(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(0.0))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_] + float((acc)[i_])); }
#define K32H(acc, a0, b0, a1, b1) K32(acc, a0, b0, a1, b1)
#else
#define RAC ACC
#define RV(x) float(x)
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) }
#define K32H(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(acc))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_]); }
#endif
#endif
const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0;
shared uint KVS[2 * 768]; // the current head's K, then V^T; before the attention the key pixels' features
#define XF KVS // the key pixels' 32 FP8 features (plane 0 | plane 1), slot-major
shared uint XT[96 * 8]; // the normalised key tokens
#define RR_KS KVS
#define RR_KS0 0u
#define RR_VS KVS
#define RR_VS0 768u
#define MG XT // the merge rows (16 x 128 FP8) reuse XT
shared float16_t O112[64 * 4]; // surface 112 per token, for surface 128's 2x2 sums
ivec2 keyXY(uint t, int x0, int y0)
{
if (t < 64u) return ivec2(x0 + int(t & 7u), y0 + int(t >> 3u));
const int i = int(t & 7u), side = int((t - 64u) >> 3u);
return ivec2(side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8), side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i));
}
#ifndef RR_GFX11
uint pack4(floate4m3_t a, floate4m3_t b, floate4m3_t c, floate4m3_t d)
{
return code(a) | (code(b) << 8u) | (code(c) << 16u) | (code(d) << 24u);
}
#endif
uvec2 mergeXY(uint r) { return r < 8u ? uvec2(r >> 1u, r & 1u) : uvec2((r - 8u) >> 1u, 2u + (r & 1u)); }
// the translated prologue's shared buffer sh: input_noise (10 x 10 pixels, 4 f16) at byte 0, input_tensor (two planes
// of 10 x 10 pixels x 16 FP8 features) at byte 800
float16_t noiseAt(uint gx, uint gy, uint c)
{
const uint b = 80u * gy + 8u * gx + 2u * c;
return hb((sh[b >> 2] >> (8u * (b & 2u))) & 0xffffu);
}
engine/rr/enc0_m_tail.glsl
+307
View File
@@ -0,0 +1,307 @@
// ======================================================================== native part (engine/rr/enc0_ref.py)
memoryBarrierShared(); barrier(); // wave 0's translated prologue filled sh (features, noise) and wrote 304 / 1016
const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID;
const uint hl = lane >> 4u, col = lane & 15u;
RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u);
const int ex = int(param.w[0]), ey = int(param.w[1]);
const uint outOff = param.w[12], mrgOff = param.w[14], shift = param.w[20];
const int x0 = 8 * int(gl_WorkGroupID.x) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u);
const uint lplane = uint((ex / 2) * (ey / 2)) * 32u;
// the key slots' features, from the 10 x 10 grid (origin x0 - 1, y0 - 1)
if (tid < 96u) {
const ivec2 k = keyXY(inter(tid), x0, y0);
const uint g = uint((k.y - y0 + 1) * 10 + (k.x - x0 + 1));
[[unroll]] for (uint wd = 0u; wd < 4u; ++wd) {
XF[tid * 8u + wd] = sh[(800u + g * 16u) / 4u + wd];
XF[tid * 8u + 4u + wd] = sh[(2400u + g * 16u) / 4u + wd];
}
}
barrier();
// the embedding x = W2 q8(relu(W1 f + b1)) + b2 for token tiles n = wv (this wave's queries: x is also their
// residual) and n = wv + 4 (the border keys, waves 0 and 1), then each key's norm in registers (dec0_ref.norm_plane
// on the accumulator layout, as dec0_m.comp's MLP norm) to XT
RAC R[2];
for (uint n = wv; n < 6u; n += 4u) {
B8 fb0, fb1, hb_[2];
LDS8B(fb0, XF, n * 128u, 8u);
LDS8B(fb1, XF, n * 128u + 4u, 8u);
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_B1 / 2u + 16u * m + 8u * hl + uint(i)]);
A8 a0, a1;
LDA(a0, W_W1, 32u, m, 0u); LDA(a1, W_W1, 32u, m, 1u);
K32(a, a0, fb0, a1, fb1)
#ifdef RR_GFX11
{ F8V hc_ = uvec2(0u); [[unroll]] for (int i = 0; i < 8; ++i) hc_[i >> 2] |= q8s(max(float(a[i]), 0.0)) << (8u * uint(i & 3));
G11_B(hb_[m], hc_) }
#else
[[unroll]] for (int i = 0; i < 8; ++i) hb_[m][i] = q8s(max(float(a[i]), 0.0));
#endif
}
float16_t xv[2][8];
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_B2 / 2u + 16u * m + 8u * hl + uint(i)]);
A8 a0, a1;
LDA(a0, W_W2, 32u, m, 0u); LDA(a1, W_W2, 32u, m, 1u);
K32(a, a0, hb_[0], a1, hb_[1])
[[unroll]] for (int i = 0; i < 8; ++i) xv[m][i] = float16_t(a[i]);
if (n == wv) { // the residual: x + table 0
[[unroll]] for (int i = 0; i < 8; ++i) {
precise float16_t c = xv[m][i] + w16.w[W_TABLE0 / 2u + 16u * m + 8u * hl + uint(i)];
R[m][i] = RV(c);
}
}
}
float16_t h1[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
float16_t ht[2];
[[unroll]] for (uint tq = 0u; tq < 2u; ++tq) {
const uint i0 = 4u * tq + c, i1 = i0 + 2u;
precise float16_t p0 = xv[0][i0] * xv[0][i0] + xv[0][i1] * xv[0][i1];
precise float16_t p1 = xv[1][i0] * xv[1][i0] + xv[1][i1] * xv[1][i1];
precise float16_t hv = p0 + p1;
ht[tq] = hv;
}
precise float16_t hh = ht[0] + ht[1];
h1[c] = hh;
}
const uint other = subgroupShuffleXor(bh(h1[0]) | (bh(h1[1]) << 16u), 16u);
precise float16_t g0 = h1[0] + hb(other & 0xffffu), g1 = h1[1] + hb(other >> 16u);
precise float16_t s = g0 + g1;
precise float16_t se = s + float16_t(0.0001220703125);
const float16_t r = rsq16(se);
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
#ifdef RR_GFX11
F8V xq = uvec2(0u);
#else
B8 xq;
#endif
[[unroll]] for (int i = 0; i < 8; ++i) {
precise float16_t g = r * w16.w[W_GAIN0 / 2u + inter(16u * m + 8u * hl + uint(i))];
precise float16_t nn = xv[m][i] * g;
#ifdef RR_GFX11
xq[i >> 2] |= q8h(nn) << (8u * uint(i & 3));
#else
xq[i] = q8h(nn);
#endif
}
#ifdef RR_GFX11
G11_ST(xq, XT, n * 128u + 4u * m, 8u)
#else
coopMatStore(xq, XT, n * 128u + 4u * m, 8u, CM);
#endif
}
}
barrier(); // XT complete; XF is read (the attention reuses its memory)
// two heads
[[unroll]] for (uint hd = 0u; hd < 2u; ++hd) {
if (hd == 1u) barrier();
for (uint j = 0u; j < 3u; ++j) {
const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u;
A8 wk[2], xa[2];
B8 wvb[2], xb[2];
[[unroll]] for (uint kt = 0u; kt < 2u; ++kt) {
LDS8B(xb[kt], XT, n * 128u + 4u * kt, 8u);
LDS8A(xa[kt], XT, n * 128u + 4u * kt, 8u);
LDA(wk[kt], W_WK0 + HEAD * hd, 32u, c, kt);
LDB(wvb[kt], W_WV0 + HEAD * hd, 32u, kt, c);
}
RAC a = RAC(RV(float16_t(0.0)));
K32(a, wk[0], xb[0], wk[1], xb[1])
Q8STA(a, KVS, n * 128u + 4u * c, 8u)
RAC vv = RAC(RV(float16_t(0.0)));
K32(vv, xa[0], wvb[0], xa[1], wvb[1])
Q8STA(vv, KVS, 768u + c * 384u + 4u * n, 24u)
}
barrier();
#include "rr_attn.glsl"
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
A8 a0, a1;
LDA(a0, W_WO0 + HEAD * hd, 32u, m, 0u); LDA(a1, W_WO0 + HEAD * hd, 32u, m, 1u);
K32(R[m], a0, Ob[0], a1, Ob[1])
}
}
// MLP (dec0_m.comp's)
B8 Nb[2];
#ifdef RR_GFX11
F8V NbC[2] = F8V[2](uvec2(0u), uvec2(0u));
#endif
{
float16_t rv[2][8];
[[unroll]] for (uint m = 0u; m < 2u; ++m) [[unroll]] for (int i = 0; i < 8; ++i) rv[m][i] = float16_t(R[m][i]);
float16_t h1[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
float16_t ht[2];
[[unroll]] for (uint tq = 0u; tq < 2u; ++tq) {
const uint i0 = 4u * tq + c, i1 = i0 + 2u;
precise float16_t p0 = rv[0][i0] * rv[0][i0] + rv[0][i1] * rv[0][i1];
precise float16_t p1 = rv[1][i0] * rv[1][i0] + rv[1][i1] * rv[1][i1];
precise float16_t hv = p0 + p1;
ht[tq] = hv;
}
precise float16_t hh = ht[0] + ht[1];
h1[c] = hh;
}
const uint other = subgroupShuffleXor(bh(h1[0]) | (bh(h1[1]) << 16u), 16u);
precise float16_t g0 = h1[0] + hb(other & 0xffffu), g1 = h1[1] + hb(other >> 16u);
precise float16_t s = g0 + g1;
precise float16_t se = s + float16_t(0.0001220703125);
const float16_t r = rsq16(se);
[[unroll]] for (uint m = 0u; m < 2u; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) {
const uint k = 16u * m + 8u * hl + uint(i);
precise float16_t g = r * w16.w[W_GAIN1 / 2u + k];
precise float16_t nn = rv[m][i] * g;
#ifdef RR_GFX11
NbC[m][i >> 2] |= q8h(nn) << (8u * uint(i & 3));
#else
Nb[m][i] = q8h(nn);
#endif
precise float16_t a = rv[m][i] + w16.w[W_TABLE1 / 2u + k];
R[m][i] = RV(a);
}
#ifdef RR_GFX11
[[unroll]] for (uint m = 0u; m < 2u; ++m) G11_B(Nb[m], NbC[m])
#endif
}
for (uint br = 0u; br < 4u; ++br) {
const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br;
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
RAC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = RV(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]);
A8 a0, a1;
LDA(a0, w1, 32u, c, 0u); LDA(a1, w1, 32u, c, 1u);
K32(h, a0, Nb[0], a1, Nb[1])
#ifdef RR_GFX11
{ F8V hc_; ACTQ8C(hc_, h) G11_B(Hb[c], hc_) }
#else
ACTQ8(Hb[c], h)
#endif
}
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
A8 a0, a1;
LDA(a0, w2, 32u, m, 0u); LDA(a1, w2, 32u, m, 1u);
K32(R[m], a0, Hb[0], a1, Hb[1])
}
}
// outputs: the tokens (one plane, dec0's skip) and the merge rows
const uint t = inter(16u * wv + col);
const int tx = x0 + int(t & 7u), ty = y0 + int(t >> 3u);
const bool inside = tx >= 0 && ty >= 0 && tx < ex && ty < ey;
B8 Tb[2];
{
const uint ty2 = t >> 3u, tx2 = t & 7u;
const uint mr = (ty2 >> 1u) < 2u ? 2u * (tx2 >> 1u) + (ty2 >> 1u) : 8u + 2u * (tx2 >> 1u) + (ty2 >> 1u) - 2u;
const uint blk = 2u * (ty2 & 1u) + (tx2 & 1u);
const uint o = (outOff + uint(ty * ex + tx) * 32u) >> 2u;
barrier(); // XT is read; the merge rows reuse it
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
#ifdef RR_GFX11
F8V tc_;
G11_CODES(tc_, R[m])
G11_B(Tb[m], tc_)
const uint q0 = tc_.x, q1 = tc_.y;
#else
Q8(Tb[m], R[m])
const uint q0 = pack4(Tb[m][0], Tb[m][1], Tb[m][2], Tb[m][3]), q1 = pack4(Tb[m][4], Tb[m][5], Tb[m][6], Tb[m][7]);
#endif
const uint at = 4u * m + 2u * hl;
if (inside) { arena.w[o + at] = q0; arena.w[o + at + 1u] = q1; }
const uint gm = mr * 32u + blk * 8u + 4u * m + 2u * hl;
MG[gm] = q0; MG[gm + 1u] = q1;
}
}
{ // the 32 -> 16 head and encoder 0's filter of the noise grid (two f16 fma chains, taps even / odd)
RAC hv;
[[unroll]] for (int i = 0; i < 8; ++i) hv[i] = RV(w16.w[W_BH / 2u + 8u * hl + uint(i)]);
A8 a0, a1;
LDA(a0, W_WH, 32u, 0u, 0u); LDA(a1, W_WH, 32u, 0u, 1u);
K32(hv, a0, Tb[0], a1, Tb[1])
float16_t l[12];
[[unroll]] for (int i = 0; i < 8; ++i) l[i] = float16_t(hv[i]);
const uint u0 = subgroupShuffleXor(bh(l[0]) | (bh(l[1]) << 16u), 16u), u1 = subgroupShuffleXor(bh(l[2]) | (bh(l[3]) << 16u), 16u);
l[8] = hb(u0 & 0xffffu); l[9] = hb(u0 >> 16u); l[10] = hb(u1 & 0xffffu); l[11] = hb(u1 >> 16u);
const float16_t ma = max(max(max(l[0], l[2]), l[4]), l[6]), mb = max(max(max(l[1], l[3]), l[5]), l[7]);
const float16_t mx = max(max(ma, mb), l[8]);
float e[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) { precise float16_t d = l[i] - mx; e[i] = exp2(float(d) * 1.4426950216293335); }
precise float s = (e[0] + e[1]) + 0.0;
s = s + (e[2] + e[3]);
s = s + (e[4] + e[5]);
s = s + (e[6] + e[7]);
s = s + e[8];
const float16_t rr = rcp16(s);
float16_t wg[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) wg[i] = h16(r16(e[i]) * float(rr));
const uint gx = t & 7u, gy = t >> 3u; // the token's 3 x 3 taps in the grid: (gx .. gx + 2, gy .. gy + 2)
float16_t o112[3], o120[3];
[[unroll]] for (uint c = 0u; c < 3u; ++c) {
float16_t lo = 0.0hf, hi = 0.0hf;
[[unroll]] for (uint k = 0u; k < 8u; k += 2u) {
lo = fma16(noiseAt(gx + k % 3u, gy + k / 3u, c), wg[k], lo);
hi = fma16(noiseAt(gx + (k + 1u) % 3u, gy + (k + 1u) / 3u, c), wg[k + 1u], hi);
}
precise float16_t sum = lo + hi;
const float16_t f = fma16(noiseAt(gx + 2u, gy + 2u, c), wg[8], sum);
precise float16_t a = sig(l[9u + c]) * f;
precise float16_t b = f - a;
o120[c] = a; o112[c] = b;
}
if (hl == 0u) {
if (inside) {
imageStore(S112, ivec2(tx, ty), vec4(o112[0], o112[1], o112[2], 0.0));
imageStore(S120, ivec2(tx, ty), vec4(o120[0], o120[1], o120[2], 0.0));
}
[[unroll]] for (uint c = 0u; c < 3u; ++c) O112[t * 4u + c] = o112[c];
}
}
barrier();
{ // the 2x2 merge (128 -> 64): wave w computes channel tile w
const uvec2 mxy = mergeXY(col);
const int X = x0 / 2 + int(mxy.x), Y = y0 / 2 + int(mxy.y), hx = ex / 2, hy = ey / 2;
const uint m = wv;
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_BM / 2u + 16u * m + 8u * hl + uint(i)]);
[[unroll]] for (uint ks = 0u; ks < 4u; ++ks) {
B8 g0, g1; A8 w0, w1;
LDS8B(g0, MG, 8u * ks, 32u);
LDS8B(g1, MG, 8u * ks + 4u, 32u);
LDA(w0, W_WM, 128u, m, 2u * ks); LDA(w1, W_WM, 128u, m, 2u * ks + 1u);
K32(a, w0, g0, w1, g1)
}
#ifdef RR_GFX11
F8V q;
G11_CODES(q, a)
if (X >= 0 && Y >= 0 && X < hx && Y < hy) {
const uint at = (mrgOff + (m >> 1u) * lplane + uint(Y * hx + X) * 32u) / 4u + 4u * (m & 1u) + 2u * hl;
arena.w[at] = q.x;
arena.w[at + 1u] = q.y;
}
#else
B8 q;
Q8(q, a)
if (X >= 0 && Y >= 0 && X < hx && Y < hy) {
const uint at = (mrgOff + (m >> 1u) * lplane + uint(Y * hx + X) * 32u) / 4u + 4u * (m & 1u) + 2u * hl;
arena.w[at] = pack4(q[0], q[1], q[2], q[3]);
arena.w[at + 1u] = pack4(q[4], q[5], q[6], q[7]);
}
#endif
}
if (wv == 3u && lane < 16u) { // surface 128: 2x2 sums of 0.25 * surface 112, rightwards first, then downwards
const uint mx = lane & 3u, my = lane >> 2u;
const int x = x0 + 2 * int(mx), y = y0 + 2 * int(my);
if (x >= 0 && y >= 0 && x < ex && y < ey) {
const uint t00 = (2u * my) * 8u + 2u * mx;
vec4 out_ = vec4(0.0);
[[unroll]] for (uint c = 0u; c < 3u; ++c) {
precise float16_t a00 = O112[t00 * 4u + c] * float16_t(0.25), a01 = O112[(t00 + 1u) * 4u + c] * float16_t(0.25);
precise float16_t a10 = O112[(t00 + 8u) * 4u + c] * float16_t(0.25), a11 = O112[(t00 + 9u) * 4u + c] * float16_t(0.25);
precise float16_t b0 = a00 + a01, b1 = a10 + a11;
precise float16_t cc = b0 + b1;
out_[c] = float(cc);
}
imageStore(S128, ivec2(x >> 1, y >> 1), out_);
}
}
engine/rr/enc0_pack.py
+65
View File
@@ -0,0 +1,65 @@
#!/usr/bin/env python3
"""Pack encoder 0's extracted weights (engine/rr/enc0_ref.py --extract) for engine/rr/enc0_m.comp's native part.
usage: enc0_pack.py WEIGHTS.npz OUT.bin OUT.h [BASE]
dec0_pack.py's layout (C = 32) plus the embedding (W1, W2: 32 -> 32 -> 32, columns in slot order), the 2x2 merge
(128 -> 64, columns in arena order) and the 32 -> 16 head. BASE (default 0) is added to every offset in OUT.h: the
shader binds NGX's weight image first (its translated feature prologue reads it) and this blob after it.
"""
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from enc1_ref import inter # noqa: E402
from enc1_pack import a_tiles, b_tiles, cubic_table, act_table # noqa: E402
I32, I64, I96 = inter(32), inter(64), inter(96)
def main():
src, out_bin, out_h = sys.argv[1:4]
base = int(sys.argv[4]) if len(sys.argv) > 4 else 0
w = dict(np.load(src))
e = {'w1': ('A', w['w1'][:, I32]), 'b1': ('H', w['b1'][I32]), 'w2': ('A', w['w2'][:, I32]), 'b2': ('H', w['b2'][I32])}
for h in range(2):
e[f'wk{h}'] = ('A', w[f'wk{h}'][:, I32])
e[f'wv{h}'] = ('B', w[f'wv{h}'][:, I32])
e[f'wo{h}'] = ('A', w[f'wo{h}'][:, I32])
e[f'bias{h}'] = ('H', w[f'bias{h}'][np.ix_(inter(64), I96)].T)
for b in range(4):
e[f'w1_{b}'] = ('A', w[f'w1_{b}'][:, I32])
e[f'b1_{b}'] = ('H', w[f'b1_{b}'][I32])
e[f'w2_{b}'] = ('A', w[f'w2_{b}'][:, I32])
e['wm'], e['bm'] = ('A', w['wm'][:, I64]), ('H', w['bm'][I64])
e['wh'], e['bh'] = ('A', w['wh']), ('H', w['bh'])
e['gain0'] = ('H', w['gain0'])
for k in ('table0', 'gain1', 'table1'):
e[k] = ('H', w[k][I32])
e['cubic'] = ('U', cubic_table())
e['act'] = ('Y', act_table())
blob, defines = bytearray(), []
for name, (kind, a) in e.items():
while len(blob) % 256:
blob.append(0)
off = len(blob)
if kind == 'A':
blob += a_tiles(a).tobytes()
elif kind == 'B':
blob += b_tiles(a).tobytes()
elif kind == 'Y':
blob += np.ascontiguousarray(a, np.uint8).tobytes()
elif kind == 'U':
blob += np.ascontiguousarray(a, np.uint16).tobytes()
else:
blob += np.ascontiguousarray(a, np.float16).tobytes()
defines.append(f'#define W_{name.upper()} {base + off}u // {kind} {a.shape}')
open(out_bin, 'wb').write(bytes(blob))
open(out_h, 'w').write('// generated by engine/rr/enc0_pack.py\n' + '\n'.join(defines) + '\n')
print(f'{len(blob)} bytes, {len(e)} entries, base {base}')
if __name__ == '__main__':
main()
engine/rr/enc0_ref.py
+138
View File
@@ -0,0 +1,138 @@
#!/usr/bin/env python3
"""Exact numpy reference of the denoiser's encoder 0 (cuda_dldn_engine_swin_enc0_kernel) after its input features,
one 8x8 window at a time. The 32 FP8 features of the 96 key pixels (NVIDIA's per-pixel gather from the game's guides
and the temporal state, which engine/rr/enc0_m.comp keeps as ptx2rr.py's translation) are an input here.
embedding F (96 x 32) -> q8(relu(F W1 + b1)) -> W2 + b2: the key tokens x (32 channels, f16)
block dec0_ref's: RMS norm (gain W+6272), residual x + table W+37056, two heads, four MLP branches
outputs the tokens (argument 48, dec0's skip), the 2x2 merge 128 -> 64 (argument 56, enc1's input) and the
32 -> 16 head with enc1's filter to surfaces 112, 120 and 128
usage: enc0_ref.py CAPTURE FRAME WEIGHTS.npz --dump DUMP.npz [--block X,Y] (verifies the dump's block)
enc0_ref.py CAPTURE FRAME OUT.npz --extract DUMP.npz (weights from a phase dump)
"""
import argparse
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rrmath import f16, q8 # noqa: E402
from enc1_ref import inter, mm, softmax_cubic, act, add, filter_head, MERGE_ORDER, H # noqa: E402
from dec0_ref import norm_plane # noqa: E402
I32, I64, I96 = inter(32), inter(64), inter(96)
GAIN0, TABLE0, GAIN1, TABLE1 = 6272, 37056, None, None
def block(w, feat):
"""feat: the 96 key pixels' features (codes, NVIDIA's k order). Returns (tokens 64 x 32 codes in arena order,
merged 16 x 64 codes in arena order, head 64 x 16 f16, dbg)"""
dbg = {}
e1 = mm(feat, w['w1'], np.repeat(w['b1'][None, :], 96, 0))
x = mm(q8(np.maximum(e1, H(0)))[:, I32], w['w2'], np.repeat(w['b2'][None, :], 96, 0))
dbg['x'] = x
X = np.array([norm_plane(x[k], w['gain0']) for k in range(96)])
R = None
for h in range(2):
qk = mm(X, w[f'wk{h}'], np.zeros((96, 32)))
v = mm(X, w[f'wv{h}'], np.zeros((96, 32)))
Q = q8(qk)[:, I32]
P = softmax_cubic(mm(Q[:64], Q.T, w[f'bias{h}']))
O = mm(q8(P)[:, I96], q8(v)[I96], np.zeros((64, 32)))
C = add(x[:64], w['table0'][None, :]) if h == 0 else R
R = mm(q8(O)[:, I32], w[f'wo{h}'], C)
N = np.array([norm_plane(R[t], w['gain1']) for t in range(64)])
acc = add(R, w['table1'][None, :])
for b in range(4):
hid = act(mm(N, w[f'w1_{b}'], np.repeat(w[f'b1_{b}'][None, :], 64, 0)))
acc = mm(q8(hid)[:, I32], w[f'w2_{b}'], acc)
tokens = q8(acc)[:, I32]
Mg = np.zeros((16, 128), np.uint8)
for r, (mx, my) in enumerate(MERGE_ORDER):
for blk, (dx, dy) in enumerate([(0, 0), (1, 0), (0, 1), (1, 1)]):
Mg[r, 32 * blk:32 * blk + 32] = tokens[(2 * my + dy) * 8 + 2 * mx + dx]
merged = q8(mm(Mg, w['wm'], np.repeat(w['bm'][None, :], 16, 0)))[:, I64]
head = mm(tokens, w['wh'], np.repeat(w['bh'][None, :], 64, 0))
return tokens, merged, head, dbg
def extract(dump, weights, out):
z = np.load(dump)
W = np.frombuffer(weights, np.uint8)
def B(p):
return z[f'p{p}_w0_B']
def C(p):
return f16(z[f'p{p}_w0_C']).astype(H)
w = {'w1': B(1), 'b1': C(1)[0], 'w2': B(2), 'b2': C(2)[0], 'gain0': W[GAIN0:GAIN0 + 64].view(H),
'table0': W[TABLE0:TABLE0 + 64].view(H)}
# the input norm's gain and the residual table sit elsewhere in preset D's (smaller) weight image: when the preset
# E offsets do not reproduce the dump, find them by matching it, like the MLP's below
x = mm(q8(np.maximum(mm(z['p1_w0_A'], w['w1'], np.repeat(w['b1'][None, :], 96, 0)), H(0)))[:, I32], w['w2'],
np.repeat(w['b2'][None, :], 96, 0))
cands = [W[o:o + 64].view(H) for o in range(0, len(W) - 63, 2)]
if not all(np.array_equal(norm_plane(x[k], w['gain0']), z['p3_w0_A'][k]) for k in (0, 9)):
w['gain0'] = next(t for t in cands if all(np.array_equal(norm_plane(x[k], t), z['p3_w0_A'][k]) for k in (0, 9)))
if not np.array_equal(add(x[:2], w['table0'][None]).view(np.uint16), C(7)[:2].view(np.uint16)):
w['table0'] = next(t for t in cands if np.array_equal(add(x[:2], t[None]).view(np.uint16), C(7)[:2].view(np.uint16)))
for h, base in ((0, 3), (1, 8)):
w[f'wk{h}'], w[f'wv{h}'] = B(base), B(base + 1)
w[f'bias{h}'] = C(base + 2)
w[f'wo{h}'] = B(base + 4)
for b in range(4):
w[f'w1_{b}'], w[f'b1_{b}'] = B(13 + 2 * b), C(13 + 2 * b)[0]
w[f'w2_{b}'] = B(14 + 2 * b)
w['wm'], w['bm'] = B(21), C(21)[0]
w['wh'], w['bh'] = B(22), C(22)[0]
# the MLP norm's gain and table, found by matching the dump
R = f16(z['p12_w0_D']).astype(H)
tables = [W[o:o + 64].view(H) for o in range(0, len(W) - 63, 2)]
w['gain1'] = next(t for t in tables if all(np.array_equal(norm_plane(R[k], t), z['p13_w0_A'][k]) for k in (0, 9)))
w['table1'] = next(t for t in tables if np.array_equal(add(R[:2], t[None]).view(np.uint16), C(14)[:2].view(np.uint16)))
np.savez_compressed(out, **w)
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('frame', type=int)
ap.add_argument('weights')
ap.add_argument('--extract', metavar='DUMP')
ap.add_argument('--dump')
ap.add_argument('--block', default='5,5')
o = ap.parse_args()
from rr_capture import Capture
c = Capture(os.path.expanduser(o.capture))
l = c.find('swin_enc0', o.frame)
if o.extract:
extract(o.extract, l.weights, o.weights)
return
z = np.load(o.dump)
w = dict(np.load(o.weights))
ex, ey = (int(v) for v in np.frombuffer(l.args[:8], np.uint32))
sx, sy = (int(v) for v in np.frombuffer(l.args[80:84], np.uint16))
bx, by = (int(v) for v in o.block.split(','))
x0, y0 = 8 * bx - sx, 8 * by - sy
tokens, merged, head, dbg = block(w, z['p1_w0_A'])
for p, name in ((3, 'X'), ):
pass
ptr = {p.offset: p.delta for p in l.pointers}
post = np.frombuffer(l.post_buffer(c, l.arena_alloc), np.uint8)
otok = post[ptr[48]:ptr[48] + ex * ey * 32].reshape(ey, ex, 32)
omrg = post[ptr[56]:ptr[56] + 2 * (ex // 2) * (ey // 2) * 32].reshape(2, ey // 2, ex // 2, 32)
img = {i.offset: i for i in l.images}
tex = {k: np.frombuffer(i.path.read_bytes(), np.float16).reshape(i.height, i.width, 4) for k, i in img.items()
if k in (112, 120, 160)}
inside = [t for t in range(64) if 0 <= x0 + t % 8 < ex and 0 <= y0 + t // 8 < ey]
tk = sum(np.array_equal(tokens[t], otok[y0 + t // 8, x0 + t % 8]) for t in inside)
mg = sum(np.array_equal(merged[r], np.concatenate([omrg[p, y0 // 2 + my, x0 // 2 + mx] for p in range(2)]))
for r, (mx, my) in enumerate(MERGE_ORDER))
print(f'block {bx},{by}: tokens {tk}/{len(inside)}, merged {mg}/16')
return head, x0, y0, ex, ey
if __name__ == '__main__':
main()
engine/rr/enc1_m.comp
+357
View File
@@ -0,0 +1,357 @@
#version 460
// Native encoder 1 of the DLSS-D (preset E) denoiser, register-resident: cuda_dldn_engine_swin_enc1_kernel
// computed like engine/layer_m.comp. The arithmetic is engine/rr/enc1_ref.py's (bit-exact to NVIDIA); the stage
// contract is the translated stage's, except binding 1, which holds engine/rr/enc1_pack.py's packing.
//
// Every product runs transposed, D^T = W^T X^T: weights are the A operand, activations the B operand, and a result's
// rows are channels, its columns tokens. On RADV gfx1201 an accumulator's element i in lane l is (row 8 * (l >> 4)
// + i, column l & 15), a B operand's the same and an A operand's the transpose (layer_m.comp): so a result is the
// next product's B operand element for element, and a result tile is also the A operand of its own transpose (the
// keys of the scores, the V^T of P V). enc1_pack.py orders every weight's columns so that the rows come out in the
// k order NVIDIA's next product contracts them in; keys and queries are token slots (slot s = token inter(s)).
// Accumulators stay f32 and are rounded to binary16 values in place (r16) after each logical k32 step.
//
// One workgroup per 8x8-token window, four waves; wave w owns query slots 16w..16w+15 and computes the window's
// keys and values itself (no barrier between the prologue and the outputs). Shared memory holds only the
// normalised tokens, the raw tokens and the merge rows.
#extension GL_KHR_cooperative_matrix : require
#extension GL_KHR_memory_scope_semantics : require
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_8bit_storage : require
#extension GL_EXT_float_e4m3 : require
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_KHR_shader_subgroup_vote : require
#extension GL_EXT_control_flow_attributes : require
#extension GL_GOOGLE_include_directive : require
layout(local_size_x = 128) in;
#include "enc1_m_w.h"
#include "rr_exact.glsl"
layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param;
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16;
layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u;
layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
layout(set = 0, binding = 3) uniform sampler2D T88;
layout(set = 0, binding = 4, rgba16f) uniform writeonly image2D S112;
layout(set = 0, binding = 5, rgba16f) uniform writeonly image2D S120;
layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S128;
#ifndef SKIP
#define SKIP 0
#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>
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CM gl_CooperativeMatrixLayoutColumnMajor
// W^T tile (mt, kt) of a weight with K rows (enc1_pack.a_tiles); W tile (kt, nt) of one with N columns (b_tiles)
#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM)
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
// binary16 rounding in place: a real conversion, kept by OR-ing in a zero the compiler cannot see (RADV folds a
// plain float(float16_t(x)) back to x); about 2 instructions per element against r16's 6
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
#define Q8(dst, src) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (dst)[i_] = q8s((src)[i_]);
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) }
const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0;
shared uint XT[96 * 16]; // the normalised key tokens: slot-major, 64 FP8 channels in arena (= k) order
// the merge rows (16 x 256 FP8, enc1_ref.MERGE_ORDER, 4 tokens of 64 channels) reuse XT, dead after the projections
#define MG XT
shared float16_t O112[64 * 4]; // surface 112 per token, for surface 128's 2x2 sums
shared uint KS[96 * 8]; // the current head's K, token slot x 32 FP8 channels (k order)
shared uint VS[32 * 24]; // the current head's V^T, 32 channels x 96 token slots
uvec2 mergeXY(uint r) { return r < 8u ? uvec2(r >> 1u, r & 1u) : uvec2((r - 8u) >> 1u, 2u + (r & 1u)); }
uint pack4(floate4m3_t a, floate4m3_t b, floate4m3_t c, floate4m3_t d)
{
return code(a) | (code(b) << 8u) | (code(c) << 16u) | (code(d) << 24u);
}
void main()
{
const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID;
const uint hl = lane >> 4u, col = lane & 15u;
RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u);
const int ex = int(param.p[0]), ey = int(param.p[1]);
const uint inOff = param.p[2], outOff = param.p[12], mrgOff = param.p[14], shift = param.p[20];
const int x0 = 8 * int(gl_WorkGroupID.x) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u);
const uint plane = uint(ex * ey) * 32u;
// ---------------------------------------------------------------- prologue: one key token per thread
if ((SKIP & 16) == 0 && tid < 96u) {
const uint t = inter(tid);
int x, y;
if (t < 64u) {
x = x0 + int(t & 7u); y = y0 + int(t >> 3u);
} else {
const int i = int(t & 7u), side = int((t - 64u) >> 3u);
x = side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8);
y = side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i);
}
x = mirror(x, ex); y = mirror(y, ey);
const uint base = (inOff + uint(y * ex + x) * 32u) >> 2u;
uint raw[16];
float16_t v[64];
[[unroll]] for (uint wd = 0u; wd < 16u; ++wd) {
raw[wd] = arena.w[base + (wd >> 3u) * (plane >> 2u) + (wd & 7u)];
[[unroll]] for (uint b = 0u; b < 4u; ++b)
v[4u * wd + b] = float16_t(uintBitsToFloate4m3EXT(uint8_t((raw[wd] >> (8u * b)) & 0xffu)));
}
// NVIDIA's order: per word t (and t + 4) the f16x2 pair sums of both planes, then lanes ^1, ^2, the halves
precise f16vec2 h[4];
[[unroll]] for (uint q = 0u; q < 4u; ++q) {
precise f16vec2 p[4];
[[unroll]] for (uint k = 0u; k < 4u; ++k) {
const uint o = 32u * (k & 1u) + 4u * (q + 4u * (k >> 1u));
precise f16vec2 a = f16vec2(v[o], v[o + 1u]) * f16vec2(v[o], v[o + 1u]);
precise f16vec2 c = f16vec2(v[o + 2u], v[o + 3u]) * f16vec2(v[o + 2u], v[o + 3u]);
p[k] = a + c;
}
precise f16vec2 s0 = p[0] + p[1], s1 = p[2] + p[3];
h[q] = s0 + s1;
}
precise f16vec2 h0 = h[0] + h[1], h2_ = h[2] + h[3];
precise f16vec2 hs = h0 + h2_;
precise float16_t s = hs.x + hs.y;
precise float16_t se = s + float16_t(0.0001220703125);
const float16_t r = rsq16(se);
[[unroll]] for (uint wd = 0u; wd < 16u; ++wd) {
floate4m3_t c[4];
[[unroll]] for (uint b = 0u; b < 4u; ++b) {
const uint j = 4u * wd + b;
precise float16_t g = r * w16.w[W_GAIN0 / 2u + inter(j)];
precise float16_t n = v[j] * g;
c[b] = q8h(n);
}
XT[tid * 16u + wd] = pack4(c[0], c[1], c[2], c[3]);
}
}
barrier();
// ---------------------------------------------------------------- the residual: raw tokens + table 7
// (each lane's token, mirrored like the prologue's; its channels 16m + 8hl .. + 7 are 8 bytes of plane m >> 1)
ACC R[4];
{
const uint tq = inter(16u * wv + col);
const int qx = mirror(x0 + int(tq & 7u), ex), qy = mirror(y0 + int(tq >> 3u), ey);
const uint qb = (inOff + uint(qy * ex + qx) * 32u) >> 2u;
[[unroll]] for (uint m = 0u; m < 4u; ++m) {
const uint at = qb + (m >> 1u) * (plane >> 2u) + 4u * (m & 1u) + 2u * hl;
const uvec2 u = uvec2(arena.w[at], arena.w[at + 1u]);
[[unroll]] for (int i = 0; i < 8; ++i) {
const float16_t x = float16_t(uintBitsToFloate4m3EXT(uint8_t((u[i >> 2] >> (8u * (uint(i) & 3u))) & 0xffu)));
precise float16_t c = x + w16.w[W_TABLE7 / 2u + 16u * m + 8u * hl + uint(i)];
R[m][i] = float(c);
}
}
}
// ---------------------------------------------------------------- attention, two heads
[[unroll]] for (uint hd = 0u; hd < ((SKIP & 64) != 0 ? 0u : 2u); ++hd) {
// this head's Q/K and V, shared: 12 jobs (token tile n, channel tile c), three per wave. K = (QK^T)^T rows
// to KS (token slot x 32 channels, also Q^T read column-major), V^T to VS (32 channels x 96 token slots)
if (hd == 1u) barrier(); // every wave is done with head 0's KS / VS
for (uint j = 0u; j < ((SKIP & 128) != 0 ? 0u : 3u); ++j) {
const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u;
A8 wk[4];
B8 wvb[4], xb[4];
A8 xa[4];
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) {
coopMatLoad(xb[kt], XT, n * 256u + 4u * kt, 16u, CM);
coopMatLoad(xa[kt], XT, n * 256u + 4u * kt, 16u, RM);
LDA(wk[kt], W_WK0 + HEAD * hd, 64u, c, kt);
LDB(wvb[kt], W_WV0 + HEAD * hd, 32u, kt, c);
}
ACC a = ACC(0.0);
K32(a, wk[0], xb[0], wk[1], xb[1])
K32(a, wk[2], xb[2], wk[3], xb[3])
A8 ka;
Q8(ka, a)
coopMatStore(ka, KS, n * 128u + 4u * c, 8u, RM);
ACC v = ACC(0.0);
K32(v, xa[0], wvb[0], xa[1], wvb[1])
K32(v, xa[2], wvb[2], xa[3], wvb[3])
A8 va;
Q8(va, v)
coopMatStore(va, VS, c * 384u + 4u * n, 24u, RM);
}
barrier();
#include "rr_attn.glsl"
[[unroll]] for (uint m = 0u; m < 4u; ++m) {
A8 a0, a1;
LDA(a0, W_WO0 + HEAD * hd, 32u, m, 0u); LDA(a1, W_WO0 + HEAD * hd, 32u, m, 1u);
K32(R[m], a0, Ob[0], a1, Ob[1])
}
}
// ---------------------------------------------------------------- MLP
B8 Nb[4];
{
float16_t rv[4][8];
[[unroll]] for (uint m = 0u; m < 4u; ++m) [[unroll]] for (int i = 0; i < 8; ++i) rv[m][i] = float16_t(R[m][i]);
// slot 16m + 8hl + i is NVIDIA's channel 8j + 2t + c with j = 4(m >> 1) + (m & 1) + 2((i >> 1) & 1),
// t = 2hl + (i >> 2), c = i & 1 (enc1_ref.nv_rowsum_sq's tree)
float16_t h1[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
float16_t ht[2];
[[unroll]] for (uint tq = 0u; tq < 2u; ++tq) {
const uint i0 = 4u * tq + c, i1 = i0 + 2u;
precise float16_t a = rv[0][i0] * rv[0][i0] + rv[0][i1] * rv[0][i1];
precise float16_t b = rv[1][i0] * rv[1][i0] + rv[1][i1] * rv[1][i1];
precise float16_t cc = rv[2][i0] * rv[2][i0] + rv[2][i1] * rv[2][i1];
precise float16_t d = rv[3][i0] * rv[3][i0] + rv[3][i1] * rv[3][i1];
precise float16_t hv = (a + cc) + (b + d);
ht[tq] = hv;
}
precise float16_t hh = ht[0] + ht[1];
h1[c] = hh;
}
const uint other = subgroupShuffleXor(bh(h1[0]) | (bh(h1[1]) << 16u), 16u);
precise float16_t g0 = h1[0] + hb(other & 0xffffu), g1 = h1[1] + hb(other >> 16u);
precise float16_t s = g0 + g1;
precise float16_t se = s + float16_t(0.0001220703125);
const float16_t r = rsq16(se);
[[unroll]] for (uint m = 0u; m < 4u; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) {
const uint k = 16u * m + 8u * hl + uint(i);
precise float16_t g = r * w16.w[W_GAIN14 / 2u + k];
precise float16_t n = rv[m][i] * g;
Nb[m][i] = q8h(n);
precise float16_t a = rv[m][i] + w16.w[W_TABLE16 / 2u + k];
R[m][i] = float(a);
}
}
for (uint br = 0u; br < ((SKIP & 4) != 0 ? 1u : 8u); ++br) {
const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br;
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
ACC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = float(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]);
A8 a0, a1, a2, a3;
LDA(a0, w1, 64u, c, 0u); LDA(a1, w1, 64u, c, 1u); LDA(a2, w1, 64u, c, 2u); LDA(a3, w1, 64u, c, 3u);
K32(h, a0, Nb[0], a1, Nb[1])
K32(h, a2, Nb[2], a3, Nb[3])
[[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = q8h(act(float16_t(h[i])));
}
[[unroll]] for (uint m = 0u; m < 4u; ++m) {
A8 a0, a1;
LDA(a0, w2, 32u, m, 0u); LDA(a1, w2, 32u, m, 1u);
K32(R[m], a0, Hb[0], a1, Hb[1])
}
}
// ---------------------------------------------------------------- outputs
const uint t = inter(16u * wv + col);
const int tx = x0 + int(t & 7u), ty = y0 + int(t >> 3u);
const bool inside = tx >= 0 && ty >= 0 && tx < ex && ty < ey;
B8 Tb[4];
{
const uint ty2 = t >> 3u, tx2 = t & 7u;
const uint r = (ty2 >> 1u) < 2u ? 2u * (tx2 >> 1u) + (ty2 >> 1u) : 8u + 2u * (tx2 >> 1u) + (ty2 >> 1u) - 2u;
const uint blk = 2u * (ty2 & 1u) + (tx2 & 1u);
const uint o = (outOff + uint(ty * ex + tx) * 32u) >> 2u;
[[unroll]] for (uint m = 0u; m < 4u; ++m) {
Q8(Tb[m], R[m])
const uint w0 = pack4(Tb[m][0], Tb[m][1], Tb[m][2], Tb[m][3]), w1 = pack4(Tb[m][4], Tb[m][5], Tb[m][6], Tb[m][7]);
const uint at = 4u * (m & 1u) + 2u * hl;
if (inside) { arena.w[o + (m >> 1u) * (plane >> 2u) + at] = w0; arena.w[o + (m >> 1u) * (plane >> 2u) + at + 1u] = w1; }
const uint g = r * 64u + 16u * blk + 4u * m + 2u * hl;
MG[g] = w0; MG[g + 1u] = w1;
}
}
if ((SKIP & 32) == 0) { // the 64 -> 16 head and its kernel-prediction filter (lanes 0-15; 16-31 hold head channels 8-15)
ACC hv;
[[unroll]] for (int i = 0; i < 8; ++i) hv[i] = float(w16.w[W_BH / 2u + 8u * hl + uint(i)]);
A8 a0, a1, a2, a3;
LDA(a0, W_WH, 64u, 0u, 0u); LDA(a1, W_WH, 64u, 0u, 1u); LDA(a2, W_WH, 64u, 0u, 2u); LDA(a3, W_WH, 64u, 0u, 3u);
K32(hv, a0, Tb[0], a1, Tb[1])
K32(hv, a2, Tb[2], a3, Tb[3])
float16_t l[12];
[[unroll]] for (int i = 0; i < 8; ++i) l[i] = float16_t(hv[i]);
const uint u0 = subgroupShuffleXor(bh(l[0]) | (bh(l[1]) << 16u), 16u), u1 = subgroupShuffleXor(bh(l[2]) | (bh(l[3]) << 16u), 16u);
l[8] = hb(u0 & 0xffffu); l[9] = hb(u0 >> 16u); l[10] = hb(u1 & 0xffffu); l[11] = hb(u1 >> 16u);
const float16_t ma = max(max(max(l[0], l[2]), l[4]), l[6]), mb = max(max(max(l[1], l[3]), l[5]), l[7]);
const float16_t mx = max(max(ma, mb), l[8]);
float e[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) { precise float16_t d = l[i] - mx; e[i] = exp2(float(d) * 1.4426950216293335); }
precise float s = (e[0] + e[1]) + 0.0;
s = s + (e[2] + e[3]);
s = s + (e[4] + e[5]);
s = s + (e[6] + e[7]);
s = s + e[8];
const float16_t rr = rcp16(s);
float wg[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) wg[i] = r16(r16(e[i]) * float(rr));
const int x = tx, y = ty;
const int xm = x - 1 < 0 ? 1 : x - 1, xc = x < 0 ? 1 : (x >= ex ? ex - 2 : x), xp = x + 1 >= ex ? ex - 2 : x + 1;
const int ym = y - 1 < 0 ? 1 : y - 1, yc = y < 0 ? 1 : (y >= ey ? ey - 2 : y), yp = y + 1 >= ey ? ey - 2 : y + 1;
const ivec2 pts[9] = ivec2[9](ivec2(xm, ym), ivec2(xc, ym), ivec2(xp, ym), ivec2(xm, yc), ivec2(xc, yc),
ivec2(xp, yc), ivec2(xm, yp), ivec2(xc, yp), ivec2(xp, yp));
float acc[3][2] = float[3][2](float[2](0.0, 0.0), float[2](0.0, 0.0), float[2](0.0, 0.0));
[[unroll]] for (uint k = 0u; k < 8u; ++k) {
const vec4 tv = texelFetch(T88, pts[k], 0);
[[unroll]] for (uint c = 0u; c < 3u; ++c) { precise float f = tv[c] * wg[k] + acc[c][k & 1u]; acc[c][k & 1u] = r16(f); }
}
const vec4 t8 = texelFetch(T88, pts[8], 0);
float o112[3], o120[3];
[[unroll]] for (uint c = 0u; c < 3u; ++c) {
precise float f0 = acc[c][0] + acc[c][1];
precise float f1 = t8[c] * wg[8] + r16(f0);
const float f = r16(f1);
precise float a = float(sig(l[9u + c])) * f;
const float ar = r16(a);
precise float b = f - ar;
o120[c] = ar; o112[c] = r16(b);
}
if (hl == 0u) {
if (inside) {
imageStore(S112, ivec2(x, y), vec4(o112[0], o112[1], o112[2], 0.0));
imageStore(S120, ivec2(x, y), vec4(o120[0], o120[1], o120[2], 0.0));
}
[[unroll]] for (uint c = 0u; c < 3u; ++c) O112[t * 4u + c] = float16_t(o112[c]);
}
}
barrier();
if ((SKIP & 8) == 0 && wv < 3u) { // the 2x2 merge: merged^T (96 x 16 rows), channel tiles 2w and 2w + 1 = arena plane w
const uvec2 mxy = mergeXY(col);
const int X = x0 / 2 + int(mxy.x), Y = y0 / 2 + int(mxy.y), hx = ex / 2, hy = ey / 2;
const uint o = (mrgOff + wv * uint(hx * hy) * 32u + uint(Y * hx + X) * 32u) >> 2u;
[[unroll]] for (uint m2 = 0u; m2 < 2u; ++m2) {
const uint m = 2u * wv + m2;
ACC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_BM / 2u + 16u * m + 8u * hl + uint(i)]);
for (uint ks = 0u; ks < 8u; ++ks) {
B8 g0, g1; A8 w0, w1;
coopMatLoad(g0, MG, 8u * ks, 64u, CM);
coopMatLoad(g1, MG, 8u * ks + 4u, 64u, CM);
LDA(w0, W_WM, 256u, m, 2u * ks); LDA(w1, W_WM, 256u, m, 2u * ks + 1u);
K32(a, w0, g0, w1, g1)
}
B8 q;
Q8(q, a)
if (X >= 0 && Y >= 0 && X < hx && Y < hy) {
arena.w[o + 4u * m2 + 2u * hl] = pack4(q[0], q[1], q[2], q[3]);
arena.w[o + 4u * m2 + 2u * hl + 1u] = pack4(q[4], q[5], q[6], q[7]);
}
}
} else if (lane < 16u) { // surface 128: 2x2 sums of 0.25 * surface 112, rightwards first, then downwards
const uint mx = lane & 3u, my = lane >> 2u;
const int x = x0 + 2 * int(mx), y = y0 + 2 * int(my);
if (x >= 0 && y >= 0 && x < ex && y < ey) {
const uint t00 = (2u * my) * 8u + 2u * mx;
vec4 out_ = vec4(0.0);
[[unroll]] for (uint c = 0u; c < 3u; ++c) {
precise float16_t a00 = O112[t00 * 4u + c] * float16_t(0.25), a01 = O112[(t00 + 1u) * 4u + c] * float16_t(0.25);
precise float16_t a10 = O112[(t00 + 8u) * 4u + c] * float16_t(0.25), a11 = O112[(t00 + 9u) * 4u + c] * float16_t(0.25);
precise float16_t b0 = a00 + a01, b1 = a10 + a11;
precise float16_t cc = b0 + b1;
out_[c] = float(cc);
}
imageStore(S128, ivec2(x >> 1, y >> 1), out_);
}
}
}
engine/rr/enc1_native.comp
+559
View File
@@ -0,0 +1,559 @@
#version 460
// Native encoder 1 of the DLSS-D (preset E) denoiser: cuda_dldn_engine_swin_enc1_kernel rewritten for cooperative
// matrices, bit-exact to engine/rr/enc1_ref.py (which is bit-exact to NVIDIA's launch). Same stage contract as the
// translated stage (engine/rr/README.md) except binding 1, which holds engine/rr/native_pack.py's packing of the
// extracted weights (offsets in the generated enc1_w.h).
//
// One workgroup is one 8x8-token window, four waves. Products run tokens x weights (A = activations from shared
// memory, B = weight tiles), one 16x16 output tile at a time, with NVIDIA's arithmetic: per logical k32 step the
// running value enters the first of two 16x16x16 FP8 products as C, the sum stays f32 and is rounded to binary16
// once, through shared memory (RADV folds a register-only f32 -> f16 -> f32). Every FP8 operand is stored in NVIDIA's
// k order: within each 32-block, slot j holds natural channel inter(j). The reductions (norms, softmax row sums) run
// in NVIDIA's exact order.
//
// 1 prologue the 96 key tokens (window, then a border row above, below, a column left and right; mirrored at the
// image edge), RMS norm (gain W_GAIN0) -> XQ; the residual + W_TABLE7 -> RES
// 2 QK and V per head, 96 x 64 -> 32 each: 24 tile jobs spread over the waves -> KV
// 3 attention per wave 16 query rows, both heads: scores + bias, cubic softmax, P V, out-projection onto RES
// 4 MLP per wave: RMS norm (W_GAIN14), RES += W_TABLE16, eight branches 64 -> 32 -> 64 onto RES
// 5 outputs the tokens (argument 48), the 2x2 merge to 96 channels (argument 56, waves 0-2), the 64 -> 16 head
// and its 3x3 kernel-prediction filter of texture 88 to surfaces 112, 120 and (half resolution) 128
#extension GL_KHR_cooperative_matrix : require
#extension GL_KHR_memory_scope_semantics : require
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_8bit_storage : require
#extension GL_EXT_float_e4m3 : require
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_GOOGLE_include_directive : require
layout(local_size_x = 128) in;
#include "enc1_w.h"
layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param;
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16;
layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
layout(set = 0, binding = 3) uniform sampler2D T88;
layout(set = 0, binding = 4, rgba16f) uniform writeonly image2D S112;
layout(set = 0, binding = 5, rgba16f) uniform writeonly image2D S120;
layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S128;
#ifdef DBG
// debugging: intermediates of window (DBX, DBY) as raw words: XQ, KV, RES after each head and after the MLP
layout(set = 0, binding = 7, std430) buffer Dbg { uint d[]; } dbg;
#define DBGWIN (gl_WorkGroupID.x == DBX && gl_WorkGroupID.y == DBY)
#define DUMP8(arr, n, at) if (DBGWIN) { barrier(); for (uint i_ = tid; i_ < (n); i_ += 128u) dbg.d[(at) + i_] = code(arr[i_]); barrier(); }
#define DUMP16(arr, n, at) if (DBGWIN) { barrier(); for (uint i_ = tid; i_ < (n); i_ += 128u) dbg.d[(at) + i_] = bh(arr[i_]); barrier(); }
#else
#define DUMP8(arr, n, at)
#define DUMP16(arr, n, at)
#endif
#ifndef SKIP
#define SKIP 0
#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>
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CM gl_CooperativeMatrixLayoutColumnMajor
#define SYNC { memoryBarrierShared(); subgroupBarrier(); }
shared floate4m3_t XQ[96 * 64]; // the normalised key tokens; after the MLP the output tokens (64 rows)
shared float16_t RES[64 * 64]; // the residual stream, natural channel order
shared floate4m3_t KV[4 * 96 * 32]; // Q/K head 0, V head 0, Q/K head 1, V head 1; later the merge rows (16 x 256)
shared float16_t SW[4 * 1536]; // per wave: f16 results (scores 16 x 96), rounding scratch at +1024
shared floate4m3_t PW[4 * 1536]; // per wave: FP8 operands (P 16 x 96; O 16 x 32; N 16 x 64 + hidden 16 x 32)
// a weight tile: matrix W (K x N) packed by native_pack.py, k-tile kt, n-tile nt
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
// one binary16 rounding of an accumulator, through the wave's scratch
#define RND(acc) { \
coopMatStore(HAC(acc), SW, scr, 16u, RM); SYNC; \
HAC r_; coopMatLoad(r_, SW, scr, 16u, RM); acc = ACC(r_); SYNC; }
// NVIDIA's k order within a 32-block: slot j holds natural channel inter(j); uninter is the inverse (and also the
// arena's channel order: arena slot uninter(k) holds natural channel k)
uint inter(uint j) { return (j & ~31u) | (((j & 31u) >> 2u) << 1u) | (j & 1u) | (((j >> 1u) & 1u) << 4u); }
uint uninter(uint c) { return (c & ~31u) | (((c & 15u) >> 1u) << 2u) | (c & 1u) | (((c >> 4u) & 1u) << 1u); }
int mirror(int v, int n) { v = abs(v); return min(v, 2 * n - 2 - v); }
// cvt.rn.satfinite.e4m3 (NaN -> 0x7f)
floate4m3_t q8(float16_t h)
{
return isnan(h) ? uintBitsToFloate4m3EXT(uint8_t(0x7fu)) : floate4m3_t(clamp(h, float16_t(-448.0), float16_t(448.0)));
}
uint code(floate4m3_t v) { return uint(floate4m3BitsToUintEXT(v)); }
float16_t hb(uint b) { return uint16BitsToFloat16(uint16_t(b)); }
uint bh(float16_t h) { return uint(float16BitsToUint16(h)); }
// x rounded to binary16 precision (nearest even), staying f32. RADV folds float(float16_t(x)) back to x, so every
// f32 -> f16 rounding whose result is used as f32 again goes through this: adding and subtracting 1.5 * 2^(E + 13)
// (E = x's binary16 exponent) leaves x rounded to a multiple of its binary16 ulp; precise keeps the pair.
float r16(float x)
{
const uint e = max(floatBitsToUint(x) & 0x7f800000u, 0x38800000u);
const float c = uintBitsToFloat(e + (13u << 23u)) * 1.5;
precise float t = x + c;
precise float r = t - c;
return r;
}
float16_t h16(float x) { return float16_t(r16(x)); }
// m * x compared with 1, exactly (the fma recovers the product's rounding error): -1, 0 or 1
int cmp1(float m, float x)
{
precise float t = m * x;
precise float e = fma(m, x, -t);
return t > 1.0 ? 1 : (t < 1.0 ? -1 : (e > 0.0 ? 1 : (e < 0.0 ? -1 : 0)));
}
// binary16 rounding of 1 / x (x > 0) to nearest even, from the hardware approximation and two midpoint checks
float16_t rcp16(float x)
{
float16_t q = h16(1.0 / x);
for (int i = 0; i < 2; ++i) {
const uint b = bh(q);
const float lo = (float(q) + float(hb(b - 1u))) * 0.5, hi = (float(q) + float(hb(b + 1u))) * 0.5;
const int cl = cmp1(lo, x), ch = cmp1(hi, x);
if (cl > 0 || (cl == 0 && (b & 1u) != 0u)) q = hb(b - 1u);
else if (ch < 0 || (ch == 0 && (b & 1u) != 0u)) q = hb(b + 1u);
}
return q;
}
// binary16 rounding of 1 / sqrt(s) (s > 0)
float16_t rsq16(float16_t s16)
{
const float s = float(s16);
float16_t q = h16(inversesqrt(s));
for (int i = 0; i < 2; ++i) {
const uint b = bh(q);
const float lo = (float(q) + float(hb(b - 1u))) * 0.5, hi = (float(q) + float(hb(b + 1u))) * 0.5;
const int cl = cmp1(lo * lo, s), ch = cmp1(hi * hi, s);
if (cl > 0 || (cl == 0 && (b & 1u) != 0u)) q = hb(b - 1u);
else if (ch < 0 || (ch == 0 && (b & 1u) != 0u)) q = hb(b + 1u);
}
return q;
}
// fma of binary16 operands with one binary16 rounding, exactly: the product is exact in f32, the sum's f32 rounding
// error comes from TwoSum and decides the f32 results that land on a binary16 midpoint
float16_t fma16(float16_t a, float16_t b, float16_t c)
{
precise float p = float(a) * float(b);
precise float s = p + float(c);
precise float bb = s - p;
precise float err = (p - (s - bb)) + (float(c) - bb);
const uint u = floatBitsToUint(s);
if ((u & 0x1fffu) == 0x1000u && err != 0.0 && abs(s) >= 6.103515625e-05)
return h16(uintBitsToFloat(err * s > 0.0 ? u + 1u : u - 1u));
return h16(s);
}
float16_t act(float16_t x)
{
precise float16_t y = clamp(x, float16_t(-2.0), float16_t(2.0));
precise float16_t a = float16_t(0.0810546875) * abs(y);
precise float16_t b = float16_t(0.412109375) - a;
precise float16_t c = y * b;
precise float16_t g = float16_t(0.5) + c;
precise float16_t r = x * g;
return r;
}
float16_t sig(float16_t v) { return rcp16(exp2(float(v) * -1.4426950216293335) + 1.0); }
uvec2 mergeXY(uint r) { return r < 8u ? uvec2(r >> 1u, r & 1u) : uvec2((r - 8u) >> 1u, 2u + (r & 1u)); }
void main()
{
const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID;
const int ex = int(param.p[0]), ey = int(param.p[1]);
const uint inOff = param.p[2], outOff = param.p[12], mrgOff = param.p[14], shift = param.p[20];
const int x0 = 8 * int(gl_WorkGroupID.x) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u);
const uint plane = uint(ex * ey) * 32u;
const uint sw = wv * 1536u, scr = sw + 1024u, r0 = wv * 16u;
// ---------------------------------------------------------------- 1 prologue
#if !(SKIP & 1)
if (tid < 96u) {
int x, y;
if (tid < 64u) {
x = x0 + int(tid & 7u); y = y0 + int(tid >> 3u);
} else {
const int i = int(tid & 7u), side = int((tid - 64u) >> 3u);
x = side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8);
y = side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i);
}
x = mirror(x, ex); y = mirror(y, ey);
const uint base = (inOff + uint(y * ex + x) * 32u) >> 2u;
float16_t v[64];
for (uint pl = 0u; pl < 2u; ++pl)
for (uint wd = 0u; wd < 8u; ++wd) {
const uint u = arena.w[base + pl * (plane >> 2u) + wd];
for (uint b = 0u; b < 4u; ++b)
v[32u * pl + 4u * wd + b] = float16_t(uintBitsToFloate4m3EXT(uint8_t((u >> (8u * b)) & 0xffu)));
}
// NVIDIA's order: per word t (and t + 4) the f16x2 pair sums of both planes, then lanes ^1, ^2, the halves
precise f16vec2 h[4];
for (uint t = 0u; t < 4u; ++t) {
precise f16vec2 p[4];
for (uint k = 0u; k < 4u; ++k) {
const uint o = 32u * (k & 1u) + 4u * (t + 4u * (k >> 1u));
precise f16vec2 a = f16vec2(v[o], v[o + 1u]) * f16vec2(v[o], v[o + 1u]);
precise f16vec2 c = f16vec2(v[o + 2u], v[o + 3u]) * f16vec2(v[o + 2u], v[o + 3u]);
p[k] = a + c;
}
precise f16vec2 s0 = p[0] + p[1], s1 = p[2] + p[3];
h[t] = s0 + s1;
}
precise f16vec2 h0 = h[0] + h[1], h2_ = h[2] + h[3];
precise f16vec2 hs = h0 + h2_;
precise float16_t s = hs.x + hs.y;
precise float16_t se = s + float16_t(0.0001220703125);
const float16_t r = rsq16(se);
for (uint j = 0u; j < 64u; ++j) {
precise float16_t g = r * w16.w[W_GAIN0 / 2u + inter(j)];
precise float16_t n = v[j] * g;
XQ[tid * 64u + j] = q8(n);
}
if (tid < 64u)
for (uint k = 0u; k < 64u; ++k) {
precise float16_t c = v[uninter(k)] + w16.w[W_TABLE7 / 2u + k];
RES[tid * 64u + k] = c;
}
}
#endif
barrier();
DUMP8(XQ, 6144u, 0u)
DUMP16(RES, 4096u, 6144u)
// ---------------------------------------------------------------- 2 QK and V projections
#if !(SKIP & 2)
for (uint job = wv; job < 24u; job += 4u) {
const uint mt = job % 6u, mat = job / 6u, hd = mat >> 1u;
const bool isV = (mat & 1u) != 0u;
const uint wa = (isV ? W_WV0A : W_WK0A) + 2048u * hd, wb = wa + 1024u;
A8 a0, a1, a2, a3;
coopMatLoad(a0, XQ, mt * 1024u, 64u, RM);
coopMatLoad(a1, XQ, mt * 1024u + 16u, 64u, RM);
coopMatLoad(a2, XQ, mt * 1024u + 32u, 64u, RM);
coopMatLoad(a3, XQ, mt * 1024u + 48u, 64u, RM);
for (uint nt = 0u; nt < 2u; ++nt) {
B8 b0, b1;
LDB(b0, wa, 32u, 0u, nt); LDB(b1, wa, 32u, 1u, nt);
ACC acc = coopMatMulAdd(a1, b1, coopMatMulAdd(a0, b0, ACC(0.0)));
RND(acc);
LDB(b0, wb, 32u, 0u, nt); LDB(b1, wb, 32u, 1u, nt);
acc = coopMatMulAdd(a3, b1, coopMatMulAdd(a2, b0, acc));
coopMatStore(HAC(acc), SW, sw + nt * 16u, 32u, RM);
}
SYNC;
const uint row = lane >> 1u, cb = (lane & 1u) * 16u, kv = mat * 3072u;
if (!isV) {
for (uint j = cb; j < cb + 16u; ++j) KV[kv + (mt * 16u + row) * 32u + j] = q8(SW[sw + row * 32u + inter(j)]);
} else {
const uint dr = uninter(mt * 16u + row);
for (uint c = cb; c < cb + 16u; ++c) KV[kv + dr * 32u + c] = q8(SW[sw + row * 32u + c]);
}
SYNC;
}
#endif
barrier();
DUMP8(KV, 12288u, 10240u)
// ---------------------------------------------------------------- 3 attention, 16 query rows per wave
#if !(SKIP & 4)
for (uint hd = 0u; hd < 2u; ++hd) {
const uint qk = 6144u * hd, vv = qk + 3072u;
{
A8 q0, q1;
coopMatLoad(q0, KV, qk + r0 * 32u, 32u, RM);
coopMatLoad(q1, KV, qk + r0 * 32u + 16u, 32u, RM);
for (uint nt = 0u; nt < 6u; ++nt) {
HAC c;
coopMatLoad(c, w16.w, (W_BIAS0 + 12288u * hd) / 2u + r0 * 96u + nt * 16u, 96u, RM);
B8 k0, k1;
coopMatLoad(k0, KV, qk + nt * 512u, 32u, CM);
coopMatLoad(k1, KV, qk + nt * 512u + 16u, 32u, CM);
ACC acc = coopMatMulAdd(q1, k1, coopMatMulAdd(q0, k0, ACC(c)));
coopMatStore(HAC(acc), SW, sw + nt * 16u, 96u, RM);
}
}
SYNC;
{ // the cubic-exponent softmax: lane = (row, column parity)
const uint row = lane >> 1u, par = lane & 1u;
float16_t e[48];
for (uint j = 0u; j < 12u; ++j)
for (uint tt = 0u; tt < 4u; ++tt) {
const float16_t x = SW[sw + row * 96u + 8u * j + 2u * tt + par];
precise float16_t m = clamp(x * float16_t(0.017181396484375), float16_t(-0.55615234375), float16_t(0.55615234375));
const float16_t t = fma16(m, -m, float16_t(0.92724609375));
const float16_t u = fma16(m, t, float16_t(1.375));
e[4u * j + tt] = hb((bh(u) << 5u) & 0x7FE0u);
}
precise float16_t sm[4];
for (uint tt = 0u; tt < 4u; ++tt) {
precise float16_t ev = ((e[tt] + e[8u + tt]) + (e[16u + tt] + e[24u + tt])) + (e[32u + tt] + e[40u + tt]);
precise float16_t od = ((e[4u + tt] + e[12u + tt]) + (e[20u + tt] + e[28u + tt])) + (e[36u + tt] + e[44u + tt]);
sm[tt] = ev + od;
}
precise float16_t out_ = (sm[0] + sm[1]) + (sm[2] + sm[3]);
const float16_t other = hb(subgroupShuffleXor(bh(out_), 1u));
precise float16_t rs = par == 0u ? out_ + other : other + out_;
const float16_t r = rcp16(float(rs));
for (uint j = 0u; j < 12u; ++j)
for (uint tt = 0u; tt < 4u; ++tt) {
precise float16_t pv = e[4u * j + tt] * r;
PW[sw + row * 96u + uninter(8u * j + 2u * tt + par)] = q8(pv);
}
}
SYNC;
if (hd == 0u) { DUMP16(SW, 6144u, 34816u) DUMP8(PW, 6144u, 40960u) }
for (uint nt = 0u; nt < 2u; ++nt) { // P V
ACC acc = ACC(0.0);
for (uint ks = 0u; ks < 3u; ++ks) {
A8 p0, p1; B8 v0, v1;
coopMatLoad(p0, PW, sw + ks * 32u, 96u, RM);
coopMatLoad(p1, PW, sw + ks * 32u + 16u, 96u, RM);
coopMatLoad(v0, KV, vv + ks * 1024u + nt * 16u, 32u, RM);
coopMatLoad(v1, KV, vv + ks * 1024u + 512u + nt * 16u, 32u, RM);
acc = coopMatMulAdd(p1, v1, coopMatMulAdd(p0, v0, acc));
if (ks < 2u) RND(acc);
}
coopMatStore(HAC(acc), SW, sw + nt * 16u, 32u, RM);
}
SYNC;
{
const uint row = lane >> 1u, cb = (lane & 1u) * 16u;
for (uint j = cb; j < cb + 16u; ++j) PW[sw + row * 32u + j] = q8(SW[sw + row * 32u + inter(j)]);
}
SYNC;
{ // out-projection onto the residual
A8 o0, o1;
coopMatLoad(o0, PW, sw, 32u, RM);
coopMatLoad(o1, PW, sw + 16u, 32u, RM);
for (uint nt = 0u; nt < 4u; ++nt) {
HAC c;
coopMatLoad(c, RES, r0 * 64u + nt * 16u, 64u, RM);
B8 b0, b1;
LDB(b0, W_WO0 + 2048u * hd, 64u, 0u, nt); LDB(b1, W_WO0 + 2048u * hd, 64u, 1u, nt);
ACC acc = coopMatMulAdd(o1, b1, coopMatMulAdd(o0, b0, ACC(c)));
coopMatStore(HAC(acc), RES, r0 * 64u + nt * 16u, 64u, RM);
}
}
SYNC;
DUMP16(RES, 4096u, 22528u + 4096u * hd)
}
#endif
// ---------------------------------------------------------------- 4 MLP
#if !(SKIP & 8)
if (lane < 16u) {
const uint base = (r0 + lane) * 64u;
float16_t R[64];
for (uint k = 0u; k < 64u; ++k) R[k] = RES[base + k];
precise f16vec2 h[4];
for (uint t = 0u; t < 4u; ++t) {
precise f16vec2 q[8];
for (uint j = 0u; j < 8u; ++j) {
const f16vec2 x = f16vec2(R[8u * j + 2u * t], R[8u * j + 2u * t + 1u]);
q[j] = x * x;
}
precise f16vec2 a = q[0] + q[2], b = q[1] + q[3], c = q[4] + q[6], d = q[5] + q[7];
precise f16vec2 ac = a + c, bd = b + d;
h[t] = ac + bd;
}
precise f16vec2 h0 = h[0] + h[1], h2_ = h[2] + h[3];
precise f16vec2 hs = h0 + h2_;
precise float16_t s = hs.x + hs.y;
precise float16_t se = s + float16_t(0.0001220703125);
const float16_t r = rsq16(se);
for (uint j = 0u; j < 64u; ++j) {
const uint k = inter(j);
precise float16_t g = r * w16.w[W_GAIN14 / 2u + k];
precise float16_t n = R[k] * g;
PW[sw + lane * 96u + j] = q8(n);
}
for (uint k = 0u; k < 64u; ++k) {
precise float16_t c = R[k] + w16.w[W_TABLE16 / 2u + k];
RES[base + k] = c;
}
}
SYNC;
{
A8 n0, n1, n2, n3;
coopMatLoad(n0, PW, sw, 96u, RM);
coopMatLoad(n1, PW, sw + 16u, 96u, RM);
coopMatLoad(n2, PW, sw + 32u, 96u, RM);
coopMatLoad(n3, PW, sw + 48u, 96u, RM);
for (uint br = 0u; br < 8u; ++br) {
const uint w1 = W_W1_0 + 2048u * br, w2 = W_W2_0 + 2048u * br;
for (uint nt = 0u; nt < 2u; ++nt) {
HAC c;
coopMatLoad(c, w16.w, (W_B1_0 + 1024u * br) / 2u + nt * 16u, 32u, RM);
B8 b0, b1;
LDB(b0, w1, 32u, 0u, nt); LDB(b1, w1, 32u, 1u, nt);
ACC acc = coopMatMulAdd(n1, b1, coopMatMulAdd(n0, b0, ACC(c)));
RND(acc);
LDB(b0, w1, 32u, 2u, nt); LDB(b1, w1, 32u, 3u, nt);
acc = coopMatMulAdd(n3, b1, coopMatMulAdd(n2, b0, acc));
coopMatStore(HAC(acc), SW, sw + nt * 16u, 32u, RM);
}
SYNC;
{
const uint row = lane >> 1u, cb = (lane & 1u) * 16u;
for (uint j = cb; j < cb + 16u; ++j) PW[sw + row * 96u + 64u + j] = q8(act(SW[sw + row * 32u + inter(j)]));
}
SYNC;
A8 g0, g1;
coopMatLoad(g0, PW, sw + 64u, 96u, RM);
coopMatLoad(g1, PW, sw + 80u, 96u, RM);
for (uint nt = 0u; nt < 4u; ++nt) {
HAC c;
coopMatLoad(c, RES, r0 * 64u + nt * 16u, 64u, RM);
B8 b0, b1;
LDB(b0, w2, 64u, 0u, nt); LDB(b1, w2, 64u, 1u, nt);
ACC acc = coopMatMulAdd(g1, b1, coopMatMulAdd(g0, b0, ACC(c)));
coopMatStore(HAC(acc), RES, r0 * 64u + nt * 16u, 64u, RM);
}
SYNC;
}
}
#endif
DUMP16(RES, 4096u, 30720u)
// ---------------------------------------------------------------- 5 outputs
{ // the tokens, in arena order (= NVIDIA's k order)
const uint row = lane >> 1u, pl = lane & 1u, t = r0 + row;
const int x = x0 + int(t & 7u), y = y0 + int(t >> 3u);
uint words[8];
for (uint wd = 0u; wd < 8u; ++wd) {
uint u = 0u;
for (uint b = 0u; b < 4u; ++b) {
const uint j = 32u * pl + 4u * wd + b;
const floate4m3_t c = q8(RES[t * 64u + inter(j)]);
XQ[t * 64u + j] = c;
u |= code(c) << (8u * b);
}
words[wd] = u;
}
if (x >= 0 && y >= 0 && x < ex && y < ey) {
const uint o = (outOff + pl * plane + uint(y * ex + x) * 32u) >> 2u;
for (uint wd = 0u; wd < 8u; ++wd) arena.w[o + wd] = words[wd];
}
}
barrier();
#if !(SKIP & 16)
for (uint i = 0u; i < 32u; ++i) { // the merge rows: (mx, my) in NVIDIA's order, 4 tokens of 64 channels
const uint idx = tid * 32u + i, r = idx >> 8u, k = idx & 255u, blk = k >> 6u;
const uvec2 m = mergeXY(r);
KV[idx] = XQ[((2u * m.y + (blk >> 1u)) * 8u + 2u * m.x + (blk & 1u)) * 64u + (k & 63u)];
}
barrier();
if (wv < 3u) { // the 2x2 merge, one 32-channel plane per wave
for (uint n2 = 0u; n2 < 2u; ++n2) {
const uint nt = 2u * wv + n2;
HAC c;
coopMatLoad(c, w16.w, W_BM / 2u + nt * 16u, 96u, RM);
ACC acc = ACC(c);
for (uint ks = 0u; ks < 8u; ++ks) {
A8 m0, m1; B8 b0, b1;
coopMatLoad(m0, KV, ks * 32u, 256u, RM);
coopMatLoad(m1, KV, ks * 32u + 16u, 256u, RM);
LDB(b0, W_WM, 96u, 2u * ks, nt); LDB(b1, W_WM, 96u, 2u * ks + 1u, nt);
acc = coopMatMulAdd(m1, b1, coopMatMulAdd(m0, b0, acc));
if (ks < 7u) RND(acc);
}
coopMatStore(HAC(acc), SW, sw + n2 * 16u, 32u, RM);
}
SYNC;
const uint row = lane >> 1u, hf = lane & 1u;
const uvec2 m = mergeXY(row);
const int X = x0 / 2 + int(m.x), Y = y0 / 2 + int(m.y), hx = ex / 2, hy = ey / 2;
uint words[4];
for (uint wd = 0u; wd < 4u; ++wd) {
uint u = 0u;
for (uint b = 0u; b < 4u; ++b) u |= code(q8(SW[sw + row * 32u + inter(16u * hf + 4u * wd + b)])) << (8u * b);
words[wd] = u;
}
if (X >= 0 && Y >= 0 && X < hx && Y < hy) {
const uint o = (mrgOff + wv * uint(hx * hy) * 32u + uint(Y * hx + X) * 32u) >> 2u;
for (uint wd = 0u; wd < 4u; ++wd) arena.w[o + 4u * hf + wd] = words[wd];
}
SYNC;
}
#endif
{ // the 64 -> 16 head
HAC c;
coopMatLoad(c, w16.w, W_BH / 2u, 16u, RM);
A8 t0, t1, t2, t3; B8 b0, b1;
coopMatLoad(t0, XQ, r0 * 64u, 64u, RM);
coopMatLoad(t1, XQ, r0 * 64u + 16u, 64u, RM);
coopMatLoad(t2, XQ, r0 * 64u + 32u, 64u, RM);
coopMatLoad(t3, XQ, r0 * 64u + 48u, 64u, RM);
LDB(b0, W_WH, 16u, 0u, 0u); LDB(b1, W_WH, 16u, 1u, 0u);
ACC acc = coopMatMulAdd(t1, b1, coopMatMulAdd(t0, b0, ACC(c)));
RND(acc);
LDB(b0, W_WH, 16u, 2u, 0u); LDB(b1, W_WH, 16u, 3u, 0u);
acc = coopMatMulAdd(t3, b1, coopMatMulAdd(t2, b0, acc));
coopMatStore(HAC(acc), SW, sw, 16u, RM);
}
SYNC;
{ // the kernel-prediction filter; lanes 16-31 repeat 0-15 so that the shuffles below stay uniform
const uint tl = lane & 15u, tk = r0 + tl;
const int x = x0 + int(tk & 7u), y = y0 + int(tk >> 3u);
float16_t l[12];
for (uint i = 0u; i < 12u; ++i) l[i] = SW[sw + tl * 16u + i];
const float16_t ma = max(max(max(l[0], l[2]), l[4]), l[6]), mb = max(max(max(l[1], l[3]), l[5]), l[7]);
const float16_t mx = max(max(ma, mb), l[8]);
float e[9];
for (uint i = 0u; i < 9u; ++i) {
precise float16_t d = l[i] - mx;
e[i] = exp2(float(d) * 1.4426950216293335);
}
precise float s = (e[0] + e[1]) + 0.0;
s = s + (e[2] + e[3]);
s = s + (e[4] + e[5]);
s = s + (e[6] + e[7]);
s = s + e[8];
const float16_t r = rcp16(s);
float wg[9];
for (uint i = 0u; i < 9u; ++i) wg[i] = r16(r16(e[i]) * float(r));
const int xm = x - 1 < 0 ? 1 : x - 1, xc = x < 0 ? 1 : (x >= ex ? ex - 2 : x), xp = x + 1 >= ex ? ex - 2 : x + 1;
const int ym = y - 1 < 0 ? 1 : y - 1, yc = y < 0 ? 1 : (y >= ey ? ey - 2 : y), yp = y + 1 >= ey ? ey - 2 : y + 1;
const ivec2 pts[9] = ivec2[9](ivec2(xm, ym), ivec2(xc, ym), ivec2(xp, ym), ivec2(xm, yc), ivec2(xc, yc),
ivec2(xp, yc), ivec2(xm, yp), ivec2(xc, yp), ivec2(xp, yp));
// f32 values that are binary16 numbers throughout (r16 at every NVIDIA f16 rounding)
float o112[3], o120[3];
float acc[3][2] = float[3][2](float[2](0.0, 0.0), float[2](0.0, 0.0), float[2](0.0, 0.0));
for (uint k = 0u; k < 8u; ++k) {
const vec4 tv = texelFetch(T88, pts[k], 0);
for (uint c = 0u; c < 3u; ++c) {
precise float f = tv[c] * wg[k] + acc[c][k & 1u];
acc[c][k & 1u] = r16(f);
}
}
const vec4 t8 = texelFetch(T88, pts[8], 0);
for (uint c = 0u; c < 3u; ++c) {
precise float f0 = acc[c][0] + acc[c][1];
precise float f1 = t8[c] * wg[8] + r16(f0);
const float f = r16(f1);
precise float a = float(sig(l[9u + c])) * f;
const float ar = r16(a);
precise float b = f - ar;
o120[c] = ar; o112[c] = r16(b);
}
// surface 128: the 2x2 sums of 0.25 * surface 112, rightwards first, then downwards
float q[3], cs[3];
const uint nb = min(lane + 1u, 31u), dn = min(lane + 8u, 31u);
for (uint c = 0u; c < 3u; ++c) {
q[c] = o112[c] * 0.25;
precise float b = q[c] + subgroupShuffle(q[c], nb);
const float br = r16(b);
precise float d = br + subgroupShuffle(br, dn);
cs[c] = r16(d);
}
if (lane < 16u && x >= 0 && y >= 0 && x < ex && y < ey) {
imageStore(S112, ivec2(x, y), vec4(o112[0], o112[1], o112[2], 0.0));
imageStore(S120, ivec2(x, y), vec4(o120[0], o120[1], o120[2], 0.0));
if (((x | y) & 1) == 0) imageStore(S128, ivec2(x >> 1, y >> 1), vec4(cs[0], cs[1], cs[2], 0.0));
}
}
}
engine/rr/enc1_pack.py
+106
View File
@@ -0,0 +1,106 @@
#!/usr/bin/env python3
"""Pack encoder 1's extracted weights (engine/rr/enc1_ref.py --extract) for engine/rr/enc1_m.comp.
usage: enc1_pack.py WEIGHTS.npz OUT.bin OUT.h
The shader runs every product transposed (D^T = W^T X^T: weights are the A operand, activations the B operand), so a
result's rows are output channels and its columns tokens, and the result is the next product's B operand in place.
The rows come out in the order of W's columns, so every weight's columns are permuted here into the k order NVIDIA's
next product contracts them in (engine/rr/enc1_ref.py: q8(x)[:, inter(n)]), and the f16 vectors that act on those
rows likewise. Keys and queries are token slots: slot s is token inter(96)[s] (the prologue's key list), so the
attention bias is stored transposed in slot order.
A tiles W (K x N) -> W^T 16x16 tiles [mt][kt], row-major (m, k): 256 * (mt * K/16 + kt) bytes
B tiles the V projection runs X Wv (A = tokens): tiles [kt][nt], row-major (k, n)
f16 vectors as is; the attention bias as 96 x 64 [key slot][query slot]
"""
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from enc1_ref import inter # noqa: E402
I32, I64, I96 = inter(32), inter(64), inter(96)
def a_tiles(w):
k, n = w.shape
return np.ascontiguousarray(w.T).reshape(n // 16, 16, k // 16, 16).transpose(0, 2, 1, 3).reshape(-1)
def b_tiles(w):
k, n = w.shape
return w.reshape(k // 16, 16, n // 16, 16).transpose(0, 2, 1, 3).reshape(-1)
CUBIC_MAX = 0x3873 # binary16 bits of 0.55615234375, the clamp of the scaled score
def act_table():
"""the MLP's hidden activation and its FP8 conversion, q8(act(h)), for every binary16 h, indexed by its bits: one
byte load in place of the six binary16 operations and the conversion (the shaders' ACTQ)"""
from enc1_ref import act
from rrmath import q8
h = np.arange(65536, dtype=np.uint32).astype(np.uint16).view(np.float16)
with np.errstate(all='ignore'):
return np.asarray(q8(act(h).astype(np.float64)), np.uint8)
def cubic_table():
"""the cubic softmax numerator's bits (enc1_ref.softmax_cubic) for every binary16 score, indexed by its bits (the
scaling and the clamp included; NaN scores take NVIDIA's min / max, which return the other operand)"""
x = np.arange(65536, dtype=np.uint32).astype(np.uint16).view(np.float16)
with np.errstate(invalid='ignore', over='ignore'):
m = (x * np.float16(0.017181396484375)).astype(np.float16)
lim = np.float16(0.55615234375)
m = np.where(np.isnan(m), lim, np.clip(m, -lim, lim)).astype(np.float16)
def fma16(a, b, c):
return (a.astype(np.float64) * b.astype(np.float64) + c.astype(np.float64)).astype(np.float16)
t = fma16(m, -m, np.full_like(m, np.float16(0.92724609375)))
u = fma16(m, t, np.full_like(m, np.float16(1.375)))
return ((u.view(np.uint16).astype(np.uint32) << 5) & 0x7FE0).astype(np.uint16)
def main():
src, out_bin, out_h = sys.argv[1:4]
w = dict(np.load(src))
e = {}
for h in range(2):
e[f'wk{h}'] = ('A', np.vstack([w[f'wk{h}a'], w[f'wk{h}b']])[:, I32])
e[f'wv{h}'] = ('B', np.vstack([w[f'wv{h}a'], w[f'wv{h}b']])[:, I32])
e[f'wo{h}'] = ('A', w[f'wo{h}'][:, I64])
e[f'bias{h}'] = ('H', w[f'bias{h}'][np.ix_(I64, I96)].T)
for b in range(8):
e[f'w1_{b}'] = ('A', w[f'w1_{b}'][:, I32])
e[f'b1_{b}'] = ('H', w[f'b1_{b}'][I32])
e[f'w2_{b}'] = ('A', w[f'w2_{b}'][:, I64])
e['wm'], e['bm'] = ('A', w['wm'][:, I96]), ('H', w['bm'][I96])
e['wh'], e['bh'] = ('A', w['wh']), ('H', w['bh'])
e['gain0'] = ('H', w['gain0'])
for k in ('table7', 'gain14', 'table16'):
e[k] = ('H', w[k][I64])
e['cubic'] = ('U', cubic_table())
blob, defines = bytearray(), []
for name, (kind, a) in e.items():
while len(blob) % 256:
blob.append(0)
off = len(blob)
if kind == 'A':
blob += a_tiles(a).tobytes()
elif kind == 'B':
blob += b_tiles(a).tobytes()
elif kind == 'U':
blob += np.ascontiguousarray(a, np.uint16).tobytes()
else:
blob += np.ascontiguousarray(a, np.float16).tobytes()
defines.append(f'#define W_{name.upper()} {off}u // {kind} {a.shape}')
open(out_bin, 'wb').write(bytes(blob))
open(out_h, 'w').write('// generated by engine/rr/enc1_pack.py\n' + '\n'.join(defines) + '\n')
print(f'{len(blob)} bytes, {len(e)} entries')
if __name__ == '__main__':
main()
engine/rr/enc1_ref.py
+345
View File
@@ -0,0 +1,345 @@
#!/usr/bin/env python3
"""Exact numpy reference of the denoiser's encoder 1 (cuda_dldn_engine_swin_enc1_kernel), one 8x8 window at a time,
from its weights and real inputs only: the specification the native shader implements.
Derived phase by phase against ZLUDA replays (engine/rr/phase_dump.py), every step bit-exact:
prologue the 96 key tokens (the window plus a one-token border, corners excluded; mirrored at the image edge)
from the arena's two 32-channel planes, f16 RMS norm (NVIDIA's summation order) times the gain at W+0
attention two heads: QK and V projections (K = 64 in two k32 steps), scores = q8(QK[:64]) q8(QK)^T + bias,
the cubic-exponent softmax, P V, out-projection onto the residual (input tile + table at W+36992)
MLP RMS norm times the gain at W+37120, eight branches of 64 -> 32 (bias, ramp activation) -> 64,
accumulated onto the residual (+ table at W+41408)
outputs q8 tokens to param+48, the 2x2 merge (16 x 256 -> 96 + bias) to param+56, the 64 -> 16 head and its
3x3 kernel-prediction filter of texture 88 to surfaces 112, 120 and (half resolution) 128
usage: enc1_ref.py CAPTURE FRAME WEIGHTS.npz [--blocks X,Y ...] (verifies against the capture)
WEIGHTS.npz: engine/rr/enc1_ref.py --extract DUMP.npz CAPTURE FRAME OUT.npz builds it from a phase dump
"""
import argparse
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rrmath import e4m3, f16, q8, mma_k32 # noqa: E402
H = np.float16
F = np.float32
def inter(n):
perm = []
for b in range(n // 32):
for m in range(8):
perm += [32 * b + 2 * m, 32 * b + 2 * m + 1, 32 * b + 16 + 2 * m, 32 * b + 17 + 2 * m]
return perm
INTER64 = inter(64)
# arena channel order (NVIDIA's A-fragment order) -> natural: natural channel k is arena channel ARENA[k]
ARENA = [0, 1, 4, 5, 8, 9, 12, 13, 16, 17, 20, 21, 24, 25, 28, 29, 2, 3, 6, 7, 10, 11, 14, 15, 18, 19, 22, 23, 26, 27,
30, 31]
ARENA = ARENA + [32 + c for c in ARENA]
def add(a, b):
return (a + b).astype(H)
def mm(a_codes, b_codes, c):
"""a (M x K e4m3 codes, NVIDIA k order) . b (K x N codes) + c (M x N f16 values) with ZLUDA's k32 rounding"""
return f16(mma_k32(a_codes, b_codes, np.asarray(c, H).view(np.uint16))).astype(H)
def nv_rowsum_sq(x, eps=2.0 ** -13):
"""f16(rsqrt(sum of squares + eps)) per row of a 64-column C/D tile, NVIDIA's order"""
sq = (x * x).astype(H)
def tile(j, t):
return sq[:, 8 * j + 2 * t: 8 * j + 2 * t + 2]
h = []
for t in range(4):
a, b = add(tile(0, t), tile(2, t)), add(tile(1, t), tile(3, t))
c, d = add(tile(4, t), tile(6, t)), add(tile(5, t), tile(7, t))
h.append(add(add(a, c), add(b, d)))
h1 = [add(h[t], h[t ^ 1]) for t in range(4)]
h2 = add(h1[0], h1[2])
s = add(add(h2[:, 0], h2[:, 1]), np.full(x.shape[0], H(eps)))
return (F(1) / np.sqrt(s.astype(F))).astype(H)
def input_norm(tok_arena, gain):
"""the prologue's per-token norm: tok_arena (2 planes x 32, arena order) -> f16 normalised token (natural)"""
sq = [(v * v).astype(H) for v in tok_arena]
h = []
for t in range(4):
def pair(p, w):
s = sq[p][4 * w: 4 * w + 4]
return add(s[0:2], s[2:4])
h.append(add(add(pair(0, t), pair(1, t)), add(pair(0, t + 4), pair(1, t + 4))))
h01 = [add(h[t], h[t ^ 1]) for t in range(4)]
h2 = add(h01[0], h01[2])
s = H(H(h2[0] + h2[1]) + H(2.0 ** -13))
r = H(F(1) / np.sqrt(F(s)))
x = np.concatenate(tok_arena)[ARENA]
return (x * (r * gain).astype(H)).astype(H), x
def softmax_cubic(x):
"""E5: the cubic-exponent softmax of a 64 x 96 score block (natural columns)"""
def fma16(a, b, c):
return (a.astype(np.float64) * b.astype(np.float64) + c.astype(np.float64)).astype(H)
m = np.clip((x * H(0.017181396484375)).astype(H), H(-0.55615234375), H(0.55615234375))
t = fma16(m, -m, np.full_like(m, H(0.92724609375)))
u = fma16(m, t, np.full_like(m, H(1.375)))
e = ((((u.view(np.uint16).astype(np.uint32)) << 5) & 0x7FE0).astype(np.uint16)).view(H)
def tile(j, tt, half):
return e[:, 8 * j + 2 * tt + half]
s = {}
for tt in range(4):
for half in range(2):
ev = add(add(add(tile(0, tt, half), tile(2, tt, half)), add(tile(4, tt, half), tile(6, tt, half))),
add(tile(8, tt, half), tile(10, tt, half)))
od = add(add(add(tile(1, tt, half), tile(3, tt, half)), add(tile(5, tt, half), tile(7, tt, half))),
add(tile(9, tt, half), tile(11, tt, half)))
s[tt, half] = add(ev, od)
out = [add(add(s[0, h_], s[1, h_]), add(s[2, h_], s[3, h_])) for h_ in range(2)]
rs = add(out[0], out[1])
r = (F(1) / rs.astype(F)).astype(H)
return (e * r[:, None]).astype(H)
# D4R_RR_FUSE=1: the arithmetic as an NVIDIA card runs it. NVIDIA's compiler contracts a binary16 multiply whose
# product only feeds an add / sub into a fused multiply-add (the kernels give those no rounding modifier); ZLUDA, and
# so the captures, round the product first. build/rr/f/nv has the measurements (ptx_contract.py is the rule).
FUSE = os.environ.get('D4R_RR_FUSE') == '1'
def fma16(a, b, c):
"""binary16 a * b + c with one rounding (the exact value fits a double)"""
with np.errstate(all='ignore'):
return (np.asarray(a, np.float64) * np.asarray(b, np.float64) + np.asarray(c, np.float64)).astype(H)
def act(x):
y = np.clip(x, H(-2), H(2)).astype(H)
if FUSE:
g = fma16(y, fma16(-np.abs(y), H(0.0810546875), H(0.412109375)), H(0.5))
return (x * g).astype(H)
g = add(H(0.5), (y * (H(0.412109375) - (H(0.0810546875) * np.abs(y)).astype(H)).astype(H)).astype(H))
return (x * g).astype(H)
MERGE_ORDER = [(0, 0), (0, 1), (1, 0), (1, 1), (2, 0), (2, 1), (3, 0), (3, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 2),
(2, 3), (3, 2), (3, 3)]
def mirror(v, n):
v = abs(v)
return min(v, 2 * n - 2 - v)
def window(w, planes, tex, x0, y0, ex, ey, dbg=None):
"""one window: returns (tokens out [64 x 64 codes, arena order], merged [16 x 96 codes, A order],
surf112 / surf120 [64 x 3 f16], surf128 contributions)"""
border = ([(x, y0 - 1) for x in range(x0, x0 + 8)] + [(x, y0 + 8) for x in range(x0, x0 + 8)]
+ [(x0 - 1, y) for y in range(y0, y0 + 8)] + [(x0 + 8, y) for y in range(y0, y0 + 8)])
coords = [(x, y) for y in range(y0, y0 + 8) for x in range(x0, x0 + 8)] + border
norm, raw = [], []
for (x, y) in coords:
xx, yy = mirror(x, ex), mirror(y, ey)
n_, r_ = input_norm([e4m3(planes[p, yy, xx]).astype(H) for p in (0, 1)], w['gain0'])
norm.append(n_)
raw.append(r_)
X = q8(np.array(norm))[:, INTER64] # 96 x 64 codes, NVIDIA k order
res = np.array(raw[:64])
dbg = {} if dbg is None else dbg
dbg['XQ'], dbg['RES0'] = X, add(res, w['table7'][None, :])
R = None
for h in range(2):
qk = mm(X[:, 32:], w[f'wk{h}b'], mm(X[:, :32], w[f'wk{h}a'], np.zeros((96, 32))))
v = mm(X[:, 32:], w[f'wv{h}b'], mm(X[:, :32], w[f'wv{h}a'], np.zeros((96, 32))))
Q = q8(qk)[:, inter(32)]
dbg[f'QK{h}'], dbg[f'V{h}'] = Q, q8(v)[inter(96)]
S = mm(Q[:64], Q.T, w[f'bias{h}'])
P = softmax_cubic(S)
O = mm(q8(P)[:, inter(96)], q8(v)[inter(96)], np.zeros((64, 32)))
dbg[f'S{h}'], dbg[f'P{h}'], dbg[f'O{h}'] = S, q8(P)[:, inter(96)], O
C = add(res, w['table7'][None, :]) if h == 0 else R
R = mm(q8(O)[:, inter(32)], w[f'wo{h}'], C)
dbg[f'R{h}'] = R
N = q8((R * (nv_rowsum_sq(R)[:, None] * w['gain14'][None, :]).astype(H)).astype(H))[:, INTER64]
acc = add(R, w['table16'][None, :])
for b in range(8):
hid = act(mm(N, w[f'w1_{b}'], np.repeat(w[f'b1_{b}'][None, :], 64, 0)))
acc = mm(q8(hid)[:, inter(32)], w[f'w2_{b}'], acc)
dbg['ACC'] = acc
tokens = q8(acc)[:, INTER64]
Mg = np.zeros((16, 256), np.uint8)
for r, (mx, my) in enumerate(MERGE_ORDER):
for blk, (dx, dy) in enumerate([(0, 0), (1, 0), (0, 1), (1, 1)]):
Mg[r, 64 * blk: 64 * blk + 64] = tokens[(2 * my + dy) * 8 + 2 * mx + dx]
merged = q8(mm(Mg, w['wm'], np.repeat(w['bm'][None, :], 16, 0)))[:, inter(96)]
head = mm(tokens, w['wh'], np.repeat(w['bh'][None, :], 64, 0))
s112, s120 = filter_head(head, tex, x0, y0, ex, ey)
return tokens, merged, s112, s120
def filter_head(D, tex, x0, y0, ex, ey, ungated=False):
"""surfaces 112 and 120 of the encoders (the gate's split of the filtered value); ungated (the bottleneck enc5):
the filtered value itself, for surface 112"""
def sig(v):
return H(F(1) / (np.exp2(F(v) * F(-1.4426950216293335)) + F(1)))
def ex2(v):
return np.exp2(F(v) * F(1.4426950216293335))
out112, out120 = np.zeros((64, 3), H), np.zeros((64, 3), H)
for tk in range(64):
x, y = x0 + tk % 8, y0 + tk // 8
l_ = D[tk].astype(H)
m01 = np.maximum(np.maximum(np.maximum(l_[0:2], l_[2:4]), l_[4:6]), l_[6:8])
m = max(max(m01[0], m01[1]), l_[8])
e = [ex2(H(l_[i] - m)) for i in range(9)]
pair = [F(e[2 * k] + e[2 * k + 1]) for k in range(4)]
s = F(F(pair[0]) + F(0))
for k in (1, 2, 3):
s = F(s + pair[k])
s = F(s + e[8])
r = H(F(1) / s)
wgt = [H(H(e[i]) * r) for i in range(9)]
def cl(v, n):
return 1 if v < 0 else (n - 2 if v >= n else v)
xm, xc, xp = cl(x - 1, ex), cl(x, ex), cl(x + 1, ex)
ym, yc, yp = cl(y - 1, ey), cl(y, ey), cl(y + 1, ey)
def T(xx, yy):
return tex[yy, xx].astype(H)
pts = [((xm, ym), (xc, ym)), ((xp, ym), (xm, yc)), ((xc, yc), (xp, yc)), ((xm, yp), (xc, yp))]
acc = np.zeros((3, 2), H)
for k, (pa, pb) in enumerate(pts):
for c in range(3):
for h_, pt in enumerate((pa, pb)):
acc[c, h_] = H(F(T(*pt)[c]) * F(wgt[2 * k + h_]) + F(acc[c, h_]))
f = [H(acc[c, 0] + acc[c, 1]) for c in range(3)]
f = [H(F(T(xp, yp)[c]) * F(wgt[8]) + F(f[c])) for c in range(3)]
if ungated:
out112[tk] = f
continue
g = [sig(l_[9]), sig(l_[10]), sig(l_[11])]
a = [H(g[c] * f[c]) for c in range(3)]
out120[tk] = a
out112[tk] = [H(f[c] - a[c]) for c in range(3)]
return out112, out120
def extract(dump, cap, frame, out):
"""the weights, from a phase dump of any block (they do not depend on the block) and the weight image"""
z = np.load(dump)
def B(p):
return z[f'p{p}_w0_B']
def C(p):
return f16(z[f'p{p}_w0_C']).astype(H)
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rr_capture import Capture
c = Capture(os.path.expanduser(cap))
W = np.frombuffer(c.find('swin_enc1', frame).weights, np.uint8)
def table(off, n=64):
return W[off:off + 2 * n].view(np.float16)
w = {'gain0': table(0), 'gain14': table(37120), 'table7': table(36992), 'table16': table(41408)}
for h, base in ((0, 1), (1, 8)):
w[f'wk{h}a'], w[f'wv{h}a'], w[f'wk{h}b'], w[f'wv{h}b'] = B(base), B(base + 1), B(base + 2), B(base + 3)
w[f'bias{h}'] = C(base + 4)
w[f'wo{h}'] = B(base + 6)
for b in range(8):
w[f'w1_{b}'], w[f'b1_{b}'] = B(15 + 2 * b), C(15 + 2 * b)[0]
w[f'w2_{b}'] = B(16 + 2 * b)
w['wm'], w['bm'] = B(31), C(31)[0]
w['wh'], w['bh'] = B(32), C(32)[0]
np.savez_compressed(out, **w)
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('frame', type=int)
ap.add_argument('weights')
ap.add_argument('--extract', metavar='DUMP')
ap.add_argument('--blocks', nargs='*', default=['5,5'])
o = ap.parse_args()
if o.extract:
extract(o.extract, o.capture, o.frame, o.weights)
return
from rr_capture import Capture
c = Capture(os.path.expanduser(o.capture))
l = c.find('swin_enc1', o.frame)
ex, ey = (int(v) for v in np.frombuffer(l.args[:8], np.uint32))
sx, sy = (int(v) for v in np.frombuffer(l.args[80:84], np.uint16))
pre = np.frombuffer(l.pre(l.arena_alloc), np.uint8)
post = np.frombuffer(l.post_buffer(c, l.arena_alloc), np.uint8)
ptr = {p.offset: p.delta for p in l.pointers}
planes = pre[ptr[8]:ptr[8] + 2 * ex * ey * 32].reshape(2, ey, ex, 32)
outp = post[ptr[48]:ptr[48] + 2 * ex * ey * 32].reshape(2, ey, ex, 32)
mrg = post[ptr[56]:ptr[56] + 3 * (ex // 2) * (ey // 2) * 32].reshape(3, ey // 2, ex // 2, 32)
img = {i.offset: i for i in l.images}
def surf(off):
i = img[off]
return np.frombuffer(i.path.read_bytes(), np.float16).reshape(i.height, i.width, 4)
tex, s112, s120 = surf(88), surf(112), surf(120)
w = dict(np.load(o.weights))
for blk in o.blocks:
bx, by = (int(v) for v in blk.split(','))
x0, y0 = 8 * bx - sx, 8 * by - sy
tokens, merged, o112, o120 = window(w, planes, tex, x0, y0, ex, ey)
tk_ok = sum(np.array_equal(tokens[t], np.concatenate([outp[0, y0 + t // 8, x0 + t % 8],
outp[1, y0 + t // 8, x0 + t % 8]]))
for t in range(64) if 0 <= x0 + t % 8 < ex and 0 <= y0 + t // 8 < ey)
mg_ok = sum(np.array_equal(merged[r], np.concatenate([mrg[p, y0 // 2 + my, x0 // 2 + mx] for p in range(3)]))
for r, (mx, my) in enumerate(MERGE_ORDER)
if 0 <= x0 // 2 + mx < ex // 2 and 0 <= y0 // 2 + my < ey // 2)
inside = [t for t in range(64) if 0 <= x0 + t % 8 < ex and 0 <= y0 + t // 8 < ey]
s_ok = sum(np.array_equal(o112[t], s112[y0 + t // 8, x0 + t % 8, :3]) and
np.array_equal(o120[t], s120[y0 + t // 8, x0 + t % 8, :3]) for t in inside)
print(f'block {bx},{by}: tokens {tk_ok}/{len(inside)}, merged {mg_ok}, surfaces {s_ok}/{len(inside)}')
if __name__ == '__main__':
main()
def pairs_tree(v):
"""NVIDIA's reduction over a list: (v0 + v1) + v2 for three, (v0 + v1) + (v2 + v3) for four, then the same
recursively on halves rounded up ((v0 + v1) + (v2 + v3)) + v4 for five"""
if len(v) == 1:
return v[0]
if len(v) == 2:
return add(v[0], v[1])
half = 1 << ((len(v) - 1).bit_length() - 1)
return add(pairs_tree(v[:half]), pairs_tree(v[half:]))
def rowsum_sq(x, eps=2.0 ** -13):
"""nv_rowsum_sq for any C = 8 n8-tiles: per lane word t the tiles of each parity (tile pairs 2 apart, tree),
the parities, then t ^ 1, t ^ 2 and the halves"""
sq = (x * x).astype(H)
nt = x.shape[1] // 8
def tile(j, t):
return sq[:, 8 * j + 2 * t: 8 * j + 2 * t + 2]
h = []
for t in range(4):
ev = pairs_tree([add(tile(j, t), tile(j + 2, t)) for j in range(0, nt, 4)] if nt % 4 == 0 else
[tile(j, t) for j in range(0, nt, 2)])
od = pairs_tree([add(tile(j, t), tile(j + 2, t)) for j in range(1, nt, 4)] if nt % 4 == 0 else
[tile(j, t) for j in range(1, nt, 2)])
h.append(add(ev, od))
h1 = [add(h[t], h[t ^ 1]) for t in range(4)]
h2 = add(h1[0], h1[2])
s = add(add(h2[:, 0], h2[:, 1]), np.full(x.shape[0], H(eps)))
return (F(1) / np.sqrt(s.astype(F))).astype(H)
engine/rr/exposure_f.comp
+93
View File
@@ -0,0 +1,93 @@
#version 460
// Ray Reconstruction preset F (rr2_*): the auto exposure of rr2_histogram_auto_exposure_basic_kernel, in two passes.
//
// NVIDIA's kernel sums a clamped luma per 16x16 block (a pairwise tree over the block's threads), converts each block
// sum to fixed point (x 2^32 / 10000, round to nearest) and adds them with a 64-bit atomic; the last block to finish
// turns the total into the exposure. The integer total does not depend on the order of the blocks, so PASS 0 writes
// every block's fixed-point sum and PASS 1 (one workgroup) adds them and does what the last block does:
// exposure = 0.23547088 / max(mean luma, 0.005) -> the 1x1 exposure image of this frame
// values = { e, p, 1 / e, clamp(log2(e / p), -6, 6) / 6 } -> the 16-byte buffer encoder 0 and the output passes read
// with e the exposure truncated to binary16 and p the previous frame's (1 where either is zero).
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_EXT_control_flow_attributes : require
#ifndef PASS
#define PASS 0
#endif
layout(set = 0, binding = 0, std430) readonly buffer Params {
ivec2 last; // render size - 1
uint blocksX, blocks;
float invCount; // 10000 / (2^32 * render pixels)
uint reset; // the kernel's argument 32: no log ratio when set
} p;
#if PASS == 0
// One thread per column of the 16x16 block. NVIDIA's tree over the block's 256 threads (index 16 y + x: i += i + 128,
// then 64, 32, ... 1) is, per column, ((r0 + r8) + (r4 + r12)) + ((r2 + r10) + (r6 + r14)) of the even rows plus the
// same of the odd rows, and then the columns (c + c8, + 4, + 2, + 1).
layout(local_size_x = 16) in;
layout(set = 0, binding = 1) uniform sampler2D color;
layout(set = 0, binding = 2, std430) writeonly buffer Sums { uint64_t s[]; } sums;
shared float col[16];
void main()
{
const uint tx = gl_LocalInvocationID.x;
const int x = int(gl_GlobalInvocationID.x), y0 = int(gl_WorkGroupID.y) * 16;
float r[16];
[[unroll]] for (int k = 0; k < 16; ++k) {
float v = 0.0;
if (x <= p.last.x && y0 + k <= p.last.y) {
const vec3 t = texelFetch(color, ivec2(x, y0 + k), 0).rgb;
precise float l = fma(t.b, 0.25, fma(t.r, 0.25, t.g * 0.5));
v = min(max(l, 0.0), 10000.0);
}
r[k] = v;
}
precise float u[8];
[[unroll]] for (int k = 0; k < 8; ++k) u[k] = r[k + 8] + r[k];
precise float v4[4];
[[unroll]] for (int k = 0; k < 4; ++k) v4[k] = u[k + 4] + u[k];
precise float w0 = v4[2] + v4[0], w1 = v4[3] + v4[1];
precise float c = w1 + w0;
col[tx] = c;
barrier();
if (tx == 0u) {
precise float d[8];
[[unroll]] for (int k = 0; k < 8; ++k) d[k] = col[k + 8] + col[k];
precise float e0 = d[4] + d[0], e1 = d[5] + d[1], e2 = d[6] + d[2], e3 = d[7] + d[3];
precise float f0 = e2 + e0, f1 = e3 + e1;
precise float g = f1 + f0;
sums.s[gl_WorkGroupID.y * p.blocksX + gl_WorkGroupID.x] = uint64_t(roundEven(g * 429496.71875));
}
}
#else
layout(local_size_x = 64) in;
shared uint64_t part[64];
layout(set = 0, binding = 1, std430) readonly buffer Sums { uint64_t s[]; } sums;
layout(set = 0, binding = 2, std430) writeonly buffer Values { vec4 v; } values;
layout(set = 0, binding = 3) uniform sampler2D previous;
layout(set = 0, binding = 4, r16f) uniform writeonly image2D current;
void main()
{
// the block sums are integers: any order. 64 strided partial sums, then thread 0 finishes
const uint tid = gl_LocalInvocationID.x;
uint64_t mine = 0ul;
for (uint b = tid; b < p.blocks; b += 64u) mine += sums.s[b];
part[tid] = mine;
barrier();
if (tid != 0u) return;
uint64_t total = 0ul;
for (uint b = 0u; b < 64u; ++b) total += part[b];
precise float mean = p.invCount * float(total);
precise float e = 0.23547088 / max(mean, 0.005);
imageStore(current, ivec2(0), vec4(e));
// cvt.rz.f16.f32: toward zero (e > 0)
float16_t h = float16_t(e);
if (float(h) > e) h = uint16BitsToFloat16(float16BitsToUint16(h) - uint16_t(1));
float cur = float(h);
if (cur == 0.0) cur = 1.0;
float prev = texelFetch(previous, ivec2(0), 0).r;
if (prev == 0.0) prev = 1.0;
float ratio = 0.0;
if (p.reset == 0u) ratio = clamp(log2(max(cur, 1e-8) / max(prev, 1e-8)), -6.0, 6.0) * 0.16666667;
values.v = vec4(cur, prev, 1.0 / cur, ratio);
}
#endif
engine/rr/gen_enc0_m.py
+45
View File
@@ -0,0 +1,45 @@
#!/usr/bin/env python3
"""Generate engine/rr/enc0_m.comp's source: encoder 0 with ptx2rr.py's translation of NVIDIA's per-pixel feature gather
(run by wave 0, as NVIDIA runs it in one warp) and the native rest (engine/rr/enc0_m_decl.glsl, enc0_m_tail.glsl).
usage: gen_enc0_m.py ENC0.generated.comp OUT.comp
The translation's main is cut at its first workgroup barrier: by then the 10 x 10 pixels' features and the noise grid
are in its shared buffer sh and the state surfaces 304 / 1016 are written. Model data derived from the user's DLL is in
the translated part; do not commit the output.
"""
import re
import sys
def main():
src, out = sys.argv[1:3]
lines = open(src).read().split('\n')
m = next(i for i, l in enumerate(lines) if l.startswith('void main()'))
head = '\n'.join(lines[:m]).replace('layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in;',
'layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in;')
body = lines[m + 2:]
first = next(i for i, l in enumerate(body) if not re.match(r' (precise |bool |uint64_t |float |uint |int )', l))
bar = next(i for i, l in enumerate(body) if l.strip() == 'memoryBarrierShared(); barrier();')
decls, prologue = body[:first], body[first:bar]
# NVIDIA's warp gathers the 96 key pixels in a loop of three 32-pixel steps, then the four corners: the steps go to
# waves 0-2 (the setup before the loop runs in every wave, the corners in wave 0)
do_ = next(i for i, l in enumerate(prologue) if l.strip() == 'do')
end = next(i for i, l in enumerate(prologue) if l.strip().startswith('} while (again_'))
loop = '\n'.join(prologue[do_ - 1:end + 1])
inc = re.search(r'(r\d+) = uint\((r\d+) \+ 32u\);\n\s+(p\d+) = \(\1 < 96u\);', loop)
counter = inc.group(1)
loop = loop.replace(inc.group(0), inc.group(0).replace('+ 32u', '+ 96u'))
pre = [l.replace(f'{counter} = uint(0u);', f'{counter} = uint(32u * gl_SubgroupID);') for l in prologue[:do_ - 1]]
post = prologue[end + 1:]
head = head.replace('shared uint mms[384];', 'shared uint mms[384 * 4];')
text = (head + '\n#extension GL_GOOGLE_include_directive : require\n#include "enc0_m_decl.glsl"\n\nvoid main()\n{\n' +
'\n'.join(decls) + '\n // the translated feature gather (NVIDIA\'s one warp, its pixel loop split over the waves)\n' +
'\n'.join(pre) + '\n if (gl_SubgroupID < 3u) {\n' + loop + '\n }\n if (gl_SubgroupID == 0u) {\n' +
'\n'.join(post) + '\n }\n#include "enc0_m_tail.glsl"\n}\n')
open(out, 'w').write(text)
print(f'{out}: {len(decls)} declaration lines, {len(prologue)} prologue lines')
if __name__ == '__main__':
main()
engine/rr/gen_rr2_enc0.py
+314
View File
@@ -0,0 +1,314 @@
#!/usr/bin/env python3
"""Generate preset F's encoder 0 shader: ptx2rr.py's translation of NVIDIA's per-pixel feature gather and of its guide
tail, around the native block (engine/rr/rr2_m.comp, KIND 4).
usage: gen_rr2_enc0.py ENC0.ptx ENC0.generated.comp OUT.comp
gen_rr2_enc0.py DEC0_1.ptx DEC0_1.generated.comp OUT.comp --dec0_1
NVIDIA's kernel runs one 32-lane warp per 8x8-token window and keeps everything in registers. Its matrix products
(the embedding, the attention block and the guide head) are one contiguous run of the translated main, which is cut
out:
wave w the translation up to the first product: the gather; the first product's A fragments (the window's 64
feature slots per token) go to the shared array XF
4 waves rr2_block(): embedding, attention, MLP, the tokens to the arena, the guide head to the shared array HD
wave w the last product's D fragments are loaded from HD and the translation continues after it: the guide tail
With --group N (2 to 4) a workgroup holds N windows along x: waves 0 ... N - 1 run the translated code of one window
each and the block then runs once per window; the stage's dispatch is 1 / N as wide ("xN" in rr.txt). Measured on
an RX 9070 XT (2026-10-13) that is slower than one window per workgroup (network +0.4 ms at 1280x720 with N = 4).
With --dec0_1 the same for dec0's second kernel (rr2_m.comp KIND 5): nothing before the block (it reads its tokens
from the arena); after it the last product's D fragments (the MLP sum) and the registers holding r are loaded from
MSH / RSH, and the translated tail forms y = r + MLP * s2 and computes the kernel's outputs from it.
The translated parts hold code derived from the user's DLL: the output is model data, do not commit it.
"""
import re
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent))
from phase_dump import mma_statements, phases_of # noqa: E402
def tiles(mm):
"""per mma of a phase: (m-tile, n-tile, k-step), as phase_dump.py orders them"""
dset = {tuple(m[2]): n for n, m in enumerate(mm)}
prev = {n: dset.get(tuple(m[5])) for n, m in enumerate(mm)}
step = {}
for n in range(len(mm)):
s, q = 0, n
while prev[q] is not None:
q = prev[q]
s += 1
step[n] = s
a_tile, b_tile = {}, {}
for s in range(max(step.values()) + 1):
for t, key in enumerate(dict.fromkeys(tuple(m[3]) for n, m in enumerate(mm) if step[n] == s)):
a_tile[(s, key)] = t
for t, key in enumerate(dict.fromkeys(tuple(m[4]) for n, m in enumerate(mm) if step[n] == s)):
b_tile[(s, key)] = t
return [(a_tile[(step[n], tuple(m[3]))], b_tile[(step[n], tuple(m[4]))], step[n]) for n, m in enumerate(mm)]
def f16vars_of(decls):
names = set()
for l in decls:
if re.match(r'\s*(precise )?float16_t ', l):
names |= set(re.findall(r'\b(rs\d+)\b', l))
return names
REG = r'\b(?:rs|rd|r|f|p|l)\d+(?:_low|_high)?\b'
# what makes a statement of the gather more than index arithmetic on the parameter block
DATA = re.compile(r'fetch0|texture|shfl_|GX\d|GA\d|GW\d|SH\d|sust_|movm_|mma_')
def tail_inputs(pre, tail):
"""The gather's values the tail reads: (recomputed, saved). Recomputed are the statements, in order, of values
that are index arithmetic on the parameter block (assigned once, outside any block); saved are the names of the
others, which the gather pass writes to the window's record."""
assigned, depth = {}, 0
for i, l in enumerate(pre):
m = re.match(r'\s*(\w+) = (.*)', l)
if m:
assigned.setdefault(m.group(1), []).append((i, depth, m.group(2)))
depth += l.count('{') - l.count('}')
used, own = [], set()
for l in tail:
m = re.match(r'\s*(\w+) = (.*)', l)
for v in re.findall(REG, m.group(2) if m else l):
if v not in own and v in assigned and v not in used:
used.append(v)
if m:
own.add(m.group(1))
def closure(v, seen):
"""statement indices v needs, or None when it is not plain index arithmetic"""
if v in seen:
return seen[v]
seen[v] = None
defs = assigned.get(v)
if defs is None:
seen[v] = set() # never assigned in the gather: a declaration's zero
return seen[v]
if len(defs) != 1 or defs[0][1] != 0 or DATA.search(defs[0][2]):
return None
need = {defs[0][0]}
for u in re.findall(REG, defs[0][2]):
c = closure(u, seen)
if c is None:
return None
need |= c
seen[v] = need
return need
lines, saved, seen = set(), [], {}
for v in used:
c = closure(v, seen)
if c is None:
saved.append(v)
else:
lines |= c
return [pre[i] for i in sorted(lines)], saved
def split_layout(xf_word, pre, tail, f16vars):
"""(recomputed statements, saved names, the record's words, the GLSL locating a window's record)"""
recomputed, saved = tail_inputs(pre, tail)
for v in saved:
if not (v in f16vars or re.fullmatch(r'r\d+', v)):
raise SystemExit(f'the tail reads {v} from the gather: not a 16- or 32-bit register')
record = 1024 + 32 * len(saved)
# the records' rows of windows lie in up to three ranges of the arena (the runtime's choice, where the network's
# regions are dead during encoder 0): words XF_WORD, +2, +4 their byte addresses, +1 and +3 the first two's rows
where = (f' const uint rr2_gx = (param.w[0] + (param.w[20] & 0xffffu) + 7u) / 8u;\n'
f' uint rr2_y = gl_WorkGroupID.y, rr2_base = param.w[{xf_word}u];\n'
f' if (rr2_y >= param.w[{xf_word + 1}u]) {{\n'
f' rr2_y -= param.w[{xf_word + 1}u]; rr2_base = param.w[{xf_word + 2}u];\n'
f' if (rr2_y >= param.w[{xf_word + 3}u]) {{ rr2_y -= param.w[{xf_word + 3}u]; rr2_base = param.w[{xf_word + 4}u]; }}\n'
f' }}\n'
f' const uint rr2_rec = rr2_base / 4u + {record}u * (rr2_y * (param.w[0] / 8u + 1u) + gl_WorkGroupID.x);\n')
return recomputed, saved, record, where
def gather_pass(xf_word, head, decls, pre, tail, first, f16vars):
"""Encoder 0 as two passes (--split gather | block): the translated gather in NVIDIA's own 32-thread groups, which
writes each window's record to the arena, and the native block with the translated tail in 128-thread groups.
One shader holding both runs the gather on one of the block's four waves while the other three wait, which
costs more than the record's trip through memory (RX 9070 XT, 1280x720, 2026-10-13: gather 0.40 -> 0.24 ms).
The record: 1024 words of A fragments ([token slot][64 FP8 feature slots], as the block's XF), then per lane the
gather's values the tail reads. Its address: the byte address of row y's range (parameter words XF_WORD to
XF_WORD + 4) + 4 * RECORD * (y' * (W / 8 + 1) + x), y' the row within its range."""
recomputed, saved, record, where = split_layout(xf_word, pre, tail, f16vars)
var = lambda r: r.lstrip('%')
store = [' if (rr2_valid) { // the first product\'s A fragments: lane 4g + t, register q -> row, 4 k',
' const uint l_ = gl_SubgroupInvocationID, g_ = l_ >> 2u, t_ = l_ & 3u;']
seen = set()
for mm_, (i, j, s) in zip(first, tiles(first)):
if (i, s) in seen:
continue
seen.add((i, s))
for q, r in enumerate(mm_[3]):
store.append(f' rr2_arena.w[rr2_rec + rr2_uninter({16 * i + 8 * (q & 1)}u + g_) * 16u + {(32 * s + 16 * (q >> 1)) // 4}u + t_] = {var(r)};')
for n, v in enumerate(saved):
bits = f'uint(float16BitsToUint16({v}))' if v in f16vars else v
store.append(f' rr2_arena.w[rr2_rec + 1024u + {len(saved)}u * l_ + {n}u] = {bits};')
store.append(' }')
print(f'gather pass: record of {record} words, the tail reads {", ".join(saved) or "nothing"} from it')
return (head + '\nlayout(set = 0, binding = 2, std430) buffer Rr2Arena { uint w[]; } rr2_arena;\n'
'uint rr2_uninter(uint c) { return (c & ~31u) | (((c & 15u) >> 1u) << 2u) | (c & 1u) | (((c >> 4u) & 1u) << 1u); }\n\n'
'void main()\n{\n' + '\n'.join(decls) + '\n' + where +
' const bool rr2_valid = gl_WorkGroupID.x < rr2_gx;\n'
' if (rr2_valid) { // NVIDIA\'s warp: the translation before the first product\n' + '\n'.join(pre) + '\n }\n' +
'\n'.join(store) + '\n}\n'), record
def main():
ptx, src, out = sys.argv[1:4]
dec = '--dec0_1' in sys.argv
# --group N: windows per workgroup (1 to 4; wave w runs the translated code of window w)
G = int(sys.argv[sys.argv.index('--group') + 1]) if '--group' in sys.argv else 1
kernel = 'rr2_dec0_kernel_1' if dec else 'rr2_enc0_kernel'
lines = Path(ptx).read_text().replace('\r', '').split('\n')
k = next(i for i, l in enumerate(lines) if f'.entry {kernel}(' in l)
phases = phases_of(mma_statements(lines, k, len(lines)), lines)
first, last = phases[0], phases[-1]
var = lambda r: r.lstrip('%')
g = Path(src).read_text().split('\n')
m = next(i for i, l in enumerate(g) if l.startswith('void main()'))
# --one-wave (D4R_RR_RELAXED=3): the block as rr2_w.comp, in NVIDIA's own 32-thread group, the tail on the same wave
one_wave = '--one-wave' in sys.argv
head = '\n'.join(g[:m])
if not one_wave:
head = head.replace('layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in;',
'layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in;')
if dec: # the tail's tanh from rr2_pack.py's table (ptx2rr.py's RR_TANH_TABLE; its offset from OUT's rr2_stage.h)
v = head.index('\n') + 1
table = re.search(r'#define W_TANH (\d+)u', (Path(out).parent / 'rr2_stage.h').read_text()).group(1)
head = head[:v] + f'#define RR_TANH_TABLE {table}u\n' + head[v:]
body = g[m + 2:]
while body and body[-1].strip() in ('', '}'):
last_line = body.pop()
if last_line.strip() == '}':
break
nd = next(i for i, l in enumerate(body) if not re.match(r' (precise |bool |uint64_t |float |uint |int |float16_t |f16vec2 )', l))
calls = [i for i, l in enumerate(body) if 'mma_e4m3_f16' in l]
a = calls[0] - 1 # '{ uint md0_, ...;'
b = next(i for i in range(calls[-1], len(body)) if body[i].strip() == '}')
if not body[a].strip().startswith('{ uint md0_'):
raise SystemExit('unexpected statement before the first product: ' + body[a])
pre, cut, tail = body[nd:a], body[a:b + 1], body[b + 1:]
if sum(l.count('{') - l.count('}') for l in pre):
raise SystemExit('the first product is inside a block of the translated main')
# values the tail takes from the cut region other than the last product's results
assigned = lambda ls: {mm_.group(1) for l in ls for mm_ in [re.match(r'\s*(\w+) = ', l)] if mm_}
restored = {var(r) for mm_ in last for r in mm_[2]}
rmap = {} # dec0_1: the register holding r for (row offset, column offset): y = r + (D half) * s2 in the tail
if dec:
place = {}
for mm_, (i, j, s_) in zip(last, tiles(last)):
for q, r in enumerate(mm_[2]):
place[r] = (16 * i + 8 * q, 8 * j)
text_tail = '\n'.join(lines[last[-1][1]:])
half = {}
for mo in re.finditer(r'mov\.b32\s+\{(%rs\d+),\s*(%rs\d+)\},\s*(%r\d+);', text_tail):
if mo.group(3) in place:
half[mo.group(1)] = (*place[mo.group(3)], 0)
half[mo.group(2)] = (*place[mo.group(3)], 1)
prod = {}
for mo in re.finditer(r'mul\.f16\s+(%rs\d+),\s*(%rs\d+),\s*(%rs\d+);', text_tail):
if mo.group(2) in half:
prod[mo.group(1)] = half[mo.group(2)]
for mo in re.finditer(r'add\.f16\s+(%rs\d+),\s*(%rs\d+),\s*(%rs\d+);', text_tail):
if mo.group(2) in prod:
rmap[var(mo.group(3))] = prod[mo.group(2)]
# the same step in a module rewritten by build/rr/f/nv/ptx_contract.py (D4R_RR_FUSE builds): one fused instruction
for mo in re.finditer(r'fma\.rn\.f16\s+(%rs\d+),\s*(%rs\d+),\s*(%rs\d+),\s*(%rs\d+);', text_tail):
if mo.group(2) in half:
rmap[var(mo.group(4))] = half[mo.group(2)]
if len(rmap) != 2 * 2 * len(last):
raise SystemExit(f'found {len(rmap)} residual registers, expected {4 * len(last)}')
restored |= set(rmap)
used = set(re.findall(r'\b[a-z]+\d+\b', '\n'.join(tail)))
lost = sorted((assigned(cut) - assigned(pre) - restored) & used)
# (values assigned again in the tail before their use are fine; the kernel has none on the reference build)
keep = []
cut_set = assigned(cut)
for name in lost: # thread indices and the like: their assignment is carried over when it depends on nothing cut
line = next(l for l in reversed(cut) if re.match(rf'\s*{name} = ', l))
if set(re.findall(r'\b[a-z]+\d+\b', line.split('=', 1)[1])) & cut_set:
raise SystemExit(f'the tail needs {name}, which the cut region computes: {line.strip()}')
keep.append(' ' + line)
split = sys.argv[sys.argv.index('--split') + 1] if '--split' in sys.argv else None
xf_word = int(sys.argv[sys.argv.index('--xf-word') + 1]) if split else 0
if split == 'gather':
text, record = gather_pass(xf_word, '\n'.join(g[:m]), body[:nd], pre, tail, first, f16vars_of(body[:nd]))
Path(out).write_text(text)
Path(out).with_suffix('.record').write_text(str(record))
return
if dec:
first = []
store = [' if (rr2_valid) { // the first product\'s A fragments: lane 4g + t, register q -> row, 4 k',
' const uint l_ = gl_SubgroupInvocationID, g_ = l_ >> 2u, t_ = l_ & 3u;']
seen = set()
for mm_, (i, j, s) in zip(first, tiles(first) if first else []):
if (i, s) in seen:
continue
seen.add((i, s))
for q, r in enumerate(mm_[3]):
store.append(f' XF[gl_SubgroupID * 1024u + uninter({16 * i + 8 * (q & 1)}u + g_) * 16u + {(32 * s + 16 * (q >> 1)) // 4}u + t_] = {var(r)};')
store.append(' }')
load = [' { // the guide head as the last product\'s D fragments (rows g, g + 8; columns 2t, 2t + 1)',
' const uint l_ = gl_SubgroupInvocationID, g_ = l_ >> 2u, t_ = l_ & 3u;']
f16vars = set()
for l in body[:nd]:
if re.match(r'\s*(precise )?float16_t ', l):
f16vars |= set(re.findall(r'\b(rs\d+)\b', l))
for mm_, (i, j, s) in zip(last, tiles(last)):
for q, r in enumerate(mm_[2]):
row = f'({16 * i + 8 * q}u + g_)'
load.append(f' {var(r)} = HDV({row}, {8 * j}u + 2u * t_) | (HDV({row}, {8 * j + 1}u + 2u * t_) << 16u);')
for name, (ro, co, hf) in rmap.items():
v = f'RSV({ro}u + g_, {co + hf}u + 2u * t_)'
load.append(f' {name} = ' + (f'uint16BitsToFloat16(uint16_t({v}));' if name in f16vars else f'{v};'))
load.append(' }')
# the native block's source, without its #version line (this shader has its own)
block = (Path(__file__).resolve().parent / ('rr2_w.comp' if one_wave else 'rr2_m.comp')).read_text().split('\n', 1)[1]
if split == 'block': # the second pass: the block on the gather pass's record, then the tail on wave 0
recomputed, saved, record, where = split_layout(xf_word, pre, tail, f16vars)
restore = [' ' + (f'{v} = uint16BitsToFloat16(uint16_t(arena.w[rr2_rec + 1024u + {len(saved)}u * l_ + {n}u]));' if v in f16vars
else f'{v} = arena.w[rr2_rec + 1024u + {len(saved)}u * l_ + {n}u];') for n, v in enumerate(saved)]
load[2:2] = restore
Path(out).write_text(
head + '\n#extension GL_GOOGLE_include_directive : require\n#define RR2_EMBED 1\n#define RR2_GROUP 1\n#define RR2_XF_ARENA 1\n' +
block + '\nuint HDV(uint row, uint col) { return col < 4u ? uint(float16BitsToUint16(HD[gl_SubgroupID * 256u + row * 4u + col])) : 0u; }\n\n'
'void main()\n{\n' + '\n'.join(body[:nd]) + '\n' + where +
' const bool rr2_valid = gl_SubgroupID < 1u && gl_WorkGroupID.x < rr2_gx;\n'
' rr2_xf = rr2_rec;\n rr2_block();\n barrier();\n if (rr2_valid) {\n' +
' // the gather\'s index arithmetic the tail reads\n' + '\n'.join(recomputed) + '\n' +
'\n'.join(load + keep) + '\n' + '\n'.join(tail) + '\n }\n}\n')
return
wx = lambda t: t.replace('gl_WorkGroupID.x', f'({G}u * gl_WorkGroupID.x + gl_SubgroupID)')
text = (head + f'\n#extension GL_GOOGLE_include_directive : require\n#define RR2_EMBED 1\n#define RR2_GROUP {G}\n' + block + '\n' +
('uint HDV(uint row, uint col) { return uint(float16BitsToUint16(MSH[gl_SubgroupID * 2048u + row * 32u + col])); }\n'
'uint RSV(uint row, uint col) { return uint(float16BitsToUint16(RSH[gl_SubgroupID * 2048u + row * 32u + col])); }\n\n' if dec else
'uint HDV(uint row, uint col) { return col < 4u ? uint(float16BitsToUint16(HD[gl_SubgroupID * 256u + row * 4u + col])) : 0u; }\n\n') +
'void main()\n{\n' + '\n'.join(body[:nd]) +
'\n // this wave\'s window, when the level has it\n'
' const uint rr2_gx = (param.w[0] + (param.w[20] & 0xffffu) + 7u) / 8u;\n'
f' const bool rr2_valid = gl_SubgroupID < {G}u && {G}u * gl_WorkGroupID.x + gl_SubgroupID < rr2_gx;\n'
' if (rr2_valid) { // NVIDIA\'s warp: the translation before the first product\n' + wx('\n'.join(pre)) + '\n }\n' +
('' if dec else '\n'.join(store)) +
# (the block runs unconditionally, also for a window past the level's last one: it writes nothing there, and
# inside a branch the driver compiles it much slower)
('\n barrier();\n rr2_block();\n barrier();\n' if G == 1 else
f'\n barrier();\n for (uint w_ = 0u; w_ < {G}u; ++w_) {{\n rr2_win = w_;\n'
' rr2_block();\n barrier();\n }\n') + ' if (rr2_valid) {\n' +
'\n'.join(load + keep) + '\n' + wx('\n'.join(tail)) + '\n }\n}\n')
Path(out).write_text(text)
print(f'{out}: {len(pre)} gather lines, {len(cut)} cut, {len(tail)} tail lines, {len(seen)} A fragment sets')
if __name__ == '__main__':
main()
engine/rr/hkpn.comp
+298
View File
@@ -0,0 +1,298 @@
#version 460
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
#extension GL_EXT_shader_16bit_storage : require
// Ray Reconstruction hkpn output stage, preset E (no diamond-wallaby branch).
// One workgroup per 16x16 output tile. Stage 1 stages the render-space neighbourhood of
// the tile into shared memory (colour log, per-texel weights, anisotropy); stage 2 gathers
// a 5x5 elliptical kernel per output pixel and writes the two output surfaces.
// Parameter block (binding 0) holds NVIDIA's kernel argument block; word index = byte / 4.
layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in;
layout(std430, binding = 0) readonly buffer Params { uint P[]; };
layout(std430, binding = 2) buffer Arena { uint A[]; };
layout(binding = 3) uniform sampler2D colourTex; // arg 176, 640x360 rgba16f
layout(binding = 5) uniform sampler2D exposureTex; // arg 200, 1x1 rgba16f
layout(binding = 8) uniform sampler2D albedoA; // arg 280, 640x360 rgba8unorm
layout(binding = 9) uniform sampler2D albedoB; // arg 288, 640x360 rgba8unorm
layout(binding = 11) uniform sampler2D decTex; // arg 448, 640x384 rgba16f
layout(rgba16f, binding = 12) uniform writeonly image2D outLog; // arg 224
layout(rgba16f, binding = 13) uniform writeonly image2D outLin; // arg 240
#define F_LOG2E uintBitsToFloat(0x3FB8AA3Bu)
#define F_2LOG2E uintBitsToFloat(0x4038AA3Bu)
#define F_LOG10_2 uintBitsToFloat(0x3E9A209Bu)
#define F_0_36 uintBitsToFloat(0x3EB851ECu)
#define F_100 uintBitsToFloat(0x42C80000u)
#define F_1E4 uintBitsToFloat(0x38D1B717u)
#define F_0_005 uintBitsToFloat(0x3BA3D70Au)
#define F_0_1 uintBitsToFloat(0x3DCCCCCDu)
#define F_1_7 uintBitsToFloat(0x3FD9999Au)
#define F_0_17 uintBitsToFloat(0x3E2E147Bu)
#define F_0_95 uintBitsToFloat(0x3F733333u)
#define F_0_6 uintBitsToFloat(0x3F19999Au)
#define F_9 uintBitsToFloat(0x41102CB4u)
// shared tile: 21 entries per row, up to 21x21 entries (render scale <= 1)
#define TILE_STRIDE 21
// binary16 numbers, stored as such: half the shared-memory traffic of the 5x5 gather
shared f16vec4 tLog[TILE_STRIDE * TILE_STRIDE]; // (L_R, L_G, L_B, alpha = 0) log-domain colour
shared f16vec4 tSig[TILE_STRIDE * TILE_STRIDE]; // (s0, s1, s2, s3) sigmoid of arena bytes 0..3
shared vec4 tAn[TILE_STRIDE * TILE_STRIDE]; // (p0, p1, p2, p3) anisotropy parameters
float pf(uint idx) { return uintBitsToFloat(P[idx]); }
// e4m3 byte -> float (exact); e4m3fn NaN for 0x7F/0xFF.
float e4m3(uint b) {
uint e = (b >> 3u) & 15u;
uint m = b & 7u;
if (e == 15u && m == 7u) return uintBitsToFloat(0x7FC00000u);
float v = (e == 0u) ? float(m) * uintBitsToFloat(0x3B000000u)
: (1.0 + float(m) * 0.125) * exp2(float(e) - 7.0);
return ((b & 128u) != 0u) ? -v : v;
}
// round to nearest even binary16, returned as f32: the hardware conversion (RNE), kept from being folded back to x by
// OR-ing an opaque zero into its bits (RADV folds a register-only f32 -> f16 -> f32)
// Beyond the binary16 range NVIDIA's arithmetic keeps a finite rounded value: that case takes the explicit rounding.
uint16_t zero16;
float f16rnSlow(float x) {
if (x == 0.0 || isnan(x) || isinf(x)) return x;
int e;
frexp(x, e);
int ee = max(e - 11, -24);
float q = uintBitsToFloat(uint(ee + 127) << 23);
return roundEven(x / q) * q;
}
float f16rn(float x) {
if (abs(x) >= 65504.0) return f16rnSlow(x);
return float(uint16BitsToFloat16(float16BitsToUint16(float16_t(x)) | zero16));
}
// Sigmoid as evaluated in binary16 on this lowering: r = f16(2^(-x*log2e)), den = f16(1+r),
// s = f16(1/den), with one Newton step for small s and the observed fix-ups of r.
float sig16(float e) {
float x = f16rn(e);
float r = f16rn(exp2(-x * F_LOG2E));
uint hb = uint(float16BitsToUint16(float16_t(r)));
if (hb == 0x1F79u || hb == 0x25CFu) {
r = f16rn(r - 0.0009765625);
} else if (hb == 0xC13Bu) {
r = f16rn(r + 0.00006103515625);
} else if (hb == 0xC1EFu) {
r = f16rn(r + 0.000030517578125);
}
float d32 = f16rn(1.0 + r);
float rc = 1.0 / d32;
float s16 = f16rn(rc);
uint sb = uint(float16BitsToUint16(float16_t(s16))) & 0x7FFFu;
if (sb != 0u && sb < 0x8Fu) {
float ee = fma(-d32, s16, 1.0);
s16 = f16rn(fma(rc, ee, s16));
}
return s16;
}
// tanh as evaluated: polynomial for small |x|, exponential form otherwise.
float tanhApprox(float x) {
float ax = abs(x);
float x2 = x * x;
float q = fma(uintBitsToFloat(0x3C80F082u), x2, uintBitsToFloat(0xBD563CAEu));
q = fma(q, x2, uintBitsToFloat(0x3E085941u));
q = fma(q, x2, uintBitsToFloat(0xBEAAA9EDu));
q = fma(q, x2, 0.0);
float small = fma(q, x, x);
float ex = exp2(ax * F_2LOG2E);
float big = fma(1.0 / (ex + 1.0), -2.0, 1.0);
if (ax >= F_9) big = 1.0;
float bigSigned = uintBitsToFloat((floatBitsToUint(x) & 0x80000000u) | floatBitsToUint(big));
return (ax >= F_0_6) ? bigSigned : small;
}
// round half away from zero
int roundAway(float v) {
return int(trunc(v + ((v < 0.0) ? -0.5 : 0.5)));
}
void main() {
zero16 = uint16_t(P[0] >> 24u); // the output width is below 2^24
// ---- parameters (byte offsets / 4) ----
uint outW = P[0];
uint outH = P[1];
float sx = pf(6u); // 24: render/out scale x
float sy = pf(7u); // 28
float ox = pf(10u); // 40: jitter offset x
float oy = pf(11u); // 44
float expScale = pf(37u); // 148
uint decMaxX = P[14] - 1u; // 56: dec0 width
uint decMaxY = P[15] - 1u; // 60: dec0 height
uint colMaxX = P[35]; // 140: colour/albedo max x
uint colMaxY = P[36]; // 144: colour/albedo max y
uint arenaBase = P[64]; // 256: arena offset
uint arenaPitch = P[4]; // 16: render width (640)
// ---- tile origin / size in render space ----
uvec2 wg = gl_WorkGroupID.xy;
float tileX = ox + fma(float(wg.x * 16u) + 0.5, sx, -0.5);
float tileY = oy + fma(float(wg.y * 16u) + 0.5, sy, -0.5);
int c0 = int(floor(tileX)) - 2;
int r0 = int(floor(tileY)) - 2;
int tw = int(ceil(16.0 * sx)) + 5;
int th = int(ceil(16.0 * sy)) + 5;
// ---- stage 1: tile load ----
uint total = uint(tw * th);
for (uint idx = gl_LocalInvocationIndex; idx < total; idx += 256u) {
int col = int(idx) % tw;
int row = int(idx) / tw;
int ix = clamp(c0 + col, 0, int(decMaxX));
int iy = clamp(r0 + row, 0, int(decMaxY));
vec3 d = texelFetch(decTex, ivec2(ix, iy), 0).rgb;
int ax = min(ix, int(colMaxX));
int ay = min(iy, int(colMaxY));
vec3 alb = texelFetch(albedoA, ivec2(ax, ay), 0).rgb
+ texelFetch(albedoB, ivec2(ax, ay), 0).rgb;
float albSum = (alb.r + alb.g) + alb.b;
if (albSum != 0.0) {
d.r *= max(alb.r, F_0_005);
d.g *= max(alb.g, F_0_005);
d.b *= max(alb.b, F_0_005);
}
vec3 lg;
lg.r = f16rn(F_0_36 * (log2(max(fma(d.r, F_100, 1.0), F_1E4)) * F_LOG10_2));
lg.g = f16rn(F_0_36 * (log2(max(fma(d.g, F_100, 1.0), F_1E4)) * F_LOG10_2));
lg.b = f16rn(F_0_36 * (log2(max(fma(d.b, F_100, 1.0), F_1E4)) * F_LOG10_2));
uint wordIdx = (arenaBase >> 2u) + uint(arenaPitch * uint(iy) + uint(ix)) * 2u;
uint lo = A[wordIdx];
uint hi = A[wordIdx + 1u];
vec4 sg = vec4(sig16(e4m3(lo & 0xFFu)),
sig16(e4m3((lo >> 8u) & 0xFFu)),
sig16(e4m3((lo >> 16u) & 0xFFu)),
sig16(e4m3(lo >> 24u)));
vec4 an = vec4(e4m3(hi & 0xFFu),
e4m3((hi >> 8u) & 0xFFu),
e4m3((hi >> 16u) & 0xFFu),
e4m3(hi >> 24u));
uint se = uint(row * TILE_STRIDE + col);
tLog[se] = f16vec4(lg, 0.0);
tSig[se] = f16vec4(sg);
tAn[se] = an;
}
barrier();
// ---- stage 2: per output pixel ----
uvec2 pix = gl_GlobalInvocationID.xy;
if (pix.x >= outW || pix.y >= outH) return;
float expo = texelFetch(exposureTex, ivec2(0, 0), 0).r;
float e41 = ((expo == 0.0) ? 1.0 : expo) * expScale;
float fx = (float(pix.x) + 0.5) * sx;
float fy = (float(pix.y) + 0.5) * sy;
float cx = fx + ox;
float cy = fy + oy;
// colour fetch: log of exposure-scaled colour at the nearest render texel
int ux = clamp(roundAway((fx - 0.5) + ox), 0, int(colMaxX));
int uy = clamp(roundAway((fy - 0.5) + oy), 0, int(colMaxY));
vec3 col = texelFetch(colourTex, ivec2(ux, uy), 0).rgb;
vec3 gcol;
gcol.r = F_0_36 * (log2(max(fma(col.r * e41, F_100, 1.0), F_1E4)) * F_LOG10_2);
gcol.g = F_0_36 * (log2(max(fma(col.g * e41, F_100, 1.0), F_1E4)) * F_LOG10_2);
gcol.b = F_0_36 * (log2(max(fma(col.b * e41, F_100, 1.0), F_1E4)) * F_LOG10_2);
// centre entry, anisotropy
float flx = floor(cx);
float fly = floor(cy);
int ix0 = int(flx) - c0;
int iy0 = int(fly) - r0;
float dx0 = (flx + 0.5) - cx;
float dy0 = (fly + 0.5) - cy;
float ks = (sy >= 0.5) ? 1.0 : (0.5 / sy);
vec4 an = tAn[iy0 * TILE_STRIDE + ix0];
float sa = fma(1.0 / (1.0 + exp2(-F_LOG2E * an.x)), F_1_7, F_0_17);
float sb = fma(1.0 / (1.0 + exp2(-F_LOG2E * an.y)), F_1_7, F_0_17);
float rho = F_0_95 * tanhApprox(an.z);
float sa2 = sa * sa;
float sb2 = sb * sb;
float denom = (sa2 * sb2) * (1.0 - rho * rho);
float invDenom = -(1.0 / denom);
float cross2 = rho * (2.0 * sqrt(sa2 * sb2));
// 5x5 gather
float sumw = 0.0;
vec3 accL = vec3(0.0);
vec4 accS = vec4(0.0);
for (int c = -2; c <= 2; c++) {
for (int r = -2; r <= 2; r++) {
float dx = ks * (dx0 + float(c));
float dy = ks * (dy0 + float(r));
uint se = uint((iy0 + r) * TILE_STRIDE + (ix0 + c));
precise float t_c = (cross2 * dx) * dy;
precise float t_b = dx * (sb2 * dx);
precise float t_a = dy * (sa2 * dy);
precise float num = (t_c - t_b) - t_a;
precise float Q = invDenom * num;
float w = exp2((Q * -0.5) * F_LOG2E);
sumw += w;
vec4 lv = vec4(tLog[se]);
vec4 sv = vec4(tSig[se]);
accL = fma(vec3(w), lv.rgb, accL);
accS = fma(vec4(w), sv, accS);
}
}
float inv = 1.0 / sumw;
vec3 Lm = accL * inv;
float S0 = accS.x * inv;
float S1 = accS.y * inv;
float S2 = accS.z * inv;
// blend weight
float q = F_0_1 / (sx * sy);
float t1 = sumw * q;
float t2 = S0 * t1;
float t3 = (sx * sy) * t2;
float t4 = 4.0 * t3;
float t5 = (sy < 0.5) ? t4 : t2;
float wgt = (t5 < 0.0) ? 0.0 : ((t5 > 1.0) ? 1.0 : t5);
float fA = fma(1.0 - S1, wgt, S1);
float fB = fma(1.0 - S2, wgt, S2);
// surface 224 (E): fA * D + (1 - fA) * gcol; D = Lm (f16 tile log channels),
// gcol = log of the exposure-scaled colour texel (PTX linear texel fetch path); D.a = 0 so aA = 1 - fA
vec3 cA = fma(Lm, vec3(fA), gcol * (1.0 - fA));
float aA = 1.0 - fA;
float sumA = cA.r + cA.g + cA.b + aA;
if (isnan(sumA) || isinf(sumA)) {
cA = Lm;
aA = 0.0;
}
imageStore(outLog, ivec2(pix), vec4(f16rn(cA.r), f16rn(cA.g), f16rn(cA.b), f16rn(aA)));
// surface 240 (E, PTX BB0_83): {R,G,B,1} of lin = (10^(clamp(c)/0.36)-1)/100/e41,
// c = fma(Lm, fB, gcol*(1-fB)); c = Lm when the sum is not finite; stored RNE f16
vec3 cB = fma(Lm, vec3(fB), gcol * (1.0 - fB));
float sumB = cB.r + cB.g + cB.b;
if (isnan(sumB) || isinf(sumB)) cB = Lm;
cB = clamp(cB, vec3(F_1E4), vec3(3.0));
// PTX divides by constants and takes lg2(10) with the hardware approximations at run time; keep the
// constants opaque so the compiler cannot fold them with different rounding.
uint zero = P[0] >> 24u;
float c036 = uintBitsToFloat(0x3EB851ECu | zero);
float c100 = uintBitsToFloat(0x42C80000u | zero);
float lg10 = log2(uintBitsToFloat(0x41200000u | zero));
float rc036 = 1.0 / c036;
float rc100 = 1.0 / c100;
float rce41 = 1.0 / e41;
precise vec3 lin = (exp2((cB * rc036) * lg10) - 1.0) * rc100 * rce41;
imageStore(outLin, ivec2(pix), vec4(float(float16_t(lin.r)), float(float16_t(lin.g)), float(float16_t(lin.b)), 1.0));
}
engine/rr/native_kernels.txt
+15
View File
@@ -0,0 +1,15 @@
# SHA-256 of the PTX modules of the network kernels engine/rr's native shaders (swin_m.comp, dec0_m.comp, enc0_m)
# were written against and verified with: nvngx_dlssd.dll 310.7. engine/rr/dll_weights.py DLL --kernels prints a DLL's.
cuda_dldn_engine_swin_dec0_diamond_wallaby_kernel f9420970b6e29c1cd100057ba9eb56a0465a95a4aaeaa21dd355605b75160e5b
cuda_dldn_engine_swin_dec0_kernel 5e533151fd1d469bb344548b113e67b5fc798db807cebaf627864ce877f01242
cuda_dldn_engine_swin_dec1_kernel ac293a9826ea7c8be44168698e79f122d0bb2cb23dbf55999a094876e04ce32f
cuda_dldn_engine_swin_dec2_kernel 282d8f093bf343c7a9080c7379ff933443943d503c7dea647535444711494528
cuda_dldn_engine_swin_dec3_kernel f2325cc78cac7c3f20d4b3b5d6604b8a350774dd83642da10ab525bd3fd1e5b0
cuda_dldn_engine_swin_dec4_kernel 9bb7cfeba4380b023f2a1b2db5fa311edd40d4e4fdce62af2fe9227a83c52213
cuda_dldn_engine_swin_enc0_diamond_wallaby_kernel 35a2b205de3feb52dacfd54ebdcb2eda4dc75fe78a181a28687876f9a1d2a17d
cuda_dldn_engine_swin_enc0_kernel e0927a196eeffc93996a0c53337fe6fecc8a869145d832a02c59d9e12741db9d
cuda_dldn_engine_swin_enc1_kernel c9df416b891e4dddbeb6311ce110b8362c1c7f1d9e9efb2595b9b4ccae155fce
cuda_dldn_engine_swin_enc2_kernel 58702ea68b478da919a31cb9052fdb82eba3d7d1ec0d94ef2bb2827686a3b36a
cuda_dldn_engine_swin_enc3_kernel 7f07b4cf84adcdc6bfb367c00767b1f124bbdb1956f704a5dbc72f9e1d3fb9ac
cuda_dldn_engine_swin_enc4_kernel d1466bfa00c702d1a67820d1ad12e351fc4185d90695b21286c1d9523e772d6e
cuda_dldn_engine_swin_enc5_kernel 720c531da5552ff9a0304c8628aea46d9eaeae3c60f9ee0cdb4f94bebfe1a6e0
engine/rr/native_pack.py
+46
View File
@@ -0,0 +1,46 @@
#!/usr/bin/env python3
"""Pack a denoiser stage's weights (engine/rr/enc1_ref.py --extract) for the native stage shaders.
usage: native_pack.py WEIGHTS.npz OUT.bin OUT.h
Every FP8 matrix W (K x N, rows in NVIDIA's k order) becomes 16x16 tiles, tile (kt, nt) at
256 * (kt * N/16 + nt) bytes, row-major inside the tile (stride 16): a B operand of the cooperative matrix loads.
Every f16 matrix is stored as is (row-major); an f16 bias vector (a 1-D entry whose name starts with 'b') as 16
equal rows, so that it loads as an accumulator tile. OUT.h lists the byte offset of every entry as a #define,
for the shader's include.
"""
import sys
import numpy as np
def tiles(w):
k, n = w.shape
assert k % 16 == 0 and n % 16 == 0, w.shape
return w.reshape(k // 16, 16, n // 16, 16).transpose(0, 2, 1, 3).reshape(-1)
def main():
src, out_bin, out_h = sys.argv[1:4]
w = dict(np.load(src))
blob = bytearray()
defines = []
for name in sorted(w):
a = w[name]
while len(blob) % 256:
blob.append(0)
off = len(blob)
if a.dtype == np.uint8:
blob += tiles(a).tobytes()
elif a.ndim == 1 and name.startswith('b'):
blob += np.tile(np.asarray(a, np.float16), (16, 1)).tobytes()
else:
blob += np.asarray(a, np.float16).tobytes()
defines.append(f'#define W_{name.upper()} {off}u // {a.dtype} {a.shape}')
open(out_bin, 'wb').write(bytes(blob))
open(out_h, 'w').write('// generated by engine/rr/native_pack.py\n' + '\n'.join(defines) + '\n')
print(f'{len(blob)} bytes, {len(w)} entries')
if __name__ == '__main__':
main()
engine/rr/phase_dump.py
+206
View File
@@ -0,0 +1,206 @@
#!/usr/bin/env python3
"""Dump every mma of a denoiser layer as matrices, from a ZLUDA replay of a captured launch: the oracle for porting
a layer to native shaders.
usage: phase_dump.py CAPTURE STAGE FRAME BX,BY OUT.npz [--phases A-B]
For one thread block, debug stores added to the PTX record every mma's A, B, C and D registers (each mma's last
execution) into padding behind the plane arena; the replay runs through ZLUDA (engine/rr/zluda_replay.py). The
fragments are then put back together with the m16n8k32 layouts (A row-major 16x32 e4m3, B column-major 32x8 e4m3,
C/D 16x8 f16) into one matrix per phase: phase P's mma are grouped into m-tiles by A fragment and into n-tiles by
B fragment, k-steps follow the C <- D chains. OUT.npz holds, per phase p and wave w: p{p}_w{w}_A (M x K, e4m3 codes),
_B (K x N codes), _C (M x N f16 of the first k-step), _D (M x N f16 of the last).
"""
import argparse
import os
import re
import subprocess
import sys
from pathlib import Path
import numpy as np
HERE = Path(__file__).resolve().parent
ROOT = HERE.parents[1]
sys.path.insert(0, str(HERE))
from rr_capture import Capture # noqa: E402
REG = r'%r\d+'
def kernel_module(kernel):
corpus = '~/.cache/d4r-rr-corpus-f' if kernel.startswith('rr2_') else '~/.cache/d4r-rr-corpus'
for p in sorted(Path(corpus).expanduser().glob('*.ptx')):
if f'.entry {kernel}(' in p.read_text(errors='replace'):
return p
raise SystemExit(f'no corpus module for {kernel}')
def mma_statements(lines, start, end):
"""(first line, last line, D, A, B, C) of every mma in [start, end)"""
out, i = [], start
while i < end:
if 'mma.sync.aligned.m16n8k32' in lines[i]:
j, text = i, ''
while True:
text += lines[j]
if lines[j].rstrip().endswith(';'):
break
j += 1
groups = [re.findall(REG, g) for g in re.findall(r'\{([^}]*)\}', text)]
out.append((i + 1, j + 1, *groups[:4]))
i = j + 1
else:
i += 1
return out
def phases_of(mmas, lines):
"""maximal runs of mma separated only by blank or comment lines"""
phases, cur = [], []
for m in mmas:
if cur:
gap = [l.strip() for l in lines[cur[-1][1]:m[0] - 1]]
if any(g and not g.startswith('//') for g in gap):
phases.append(cur)
cur = []
cur.append(m)
if cur:
phases.append(cur)
return phases
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('stage')
ap.add_argument('frame', type=int)
ap.add_argument('block')
ap.add_argument('out')
ap.add_argument('--phases', help='first-last (1-based), default all')
o = ap.parse_args()
# the capture decides the kernel: a preset D capture runs cuda_dldn_engine_swin_{enc0,dec0}_diamond_wallaby_kernel
cap = Capture(os.path.expanduser(o.capture))
# preset F: STAGE is the kernel's full name (rr2_enc1_kernel, rr2_dec1_kernel_0)
rr2 = o.stage.startswith('rr2_')
launch = [l for l in cap.launches if (l.kernel == o.stage if rr2 else f'swin_{o.stage}' in l.kernel)][o.frame - 1]
kernel = launch.kernel
module = kernel_module(kernel)
lines = module.read_text().replace('\r', '').split('\n')
k = next(i for i, l in enumerate(lines) if f'.entry {kernel}(' in l)
e = next((i for i in range(k + 1, len(lines)) if '.entry ' in lines[i]), len(lines))
phases = phases_of(mma_statements(lines, k, e), lines)
first, last = (int(x) for x in o.phases.split('-')) if o.phases else (1, len(phases))
T = launch.block[0] * launch.block[1] * launch.block[2]
arena_bytes = len(launch.pre(launch.arena_alloc))
# any argument pointing into the arena locates its base (enc0 has none at 8)
aptr = next(p for p in launch.pointers if p.alloc == launch.arena_alloc)
delta8, off8 = aptr.delta, aptr.offset
bx, by = (int(v) for v in o.block.split(','))
# one probe slot per register per mma: after the mma, all of A, B, C, D (C read after the mma is still C
# unless C and D share registers, which the k-step chains never do for the first step)
slots, inserts = [], {}
pname = f'{kernel}_param_0'
for p in range(first, last + 1):
for (l0, l1, D, A, B, C) in phases[p - 1]:
code = []
for role, regs in (('D', D), ('A', A), ('B', B), ('C', C)):
for r in regs:
s = len(slots)
slots.append((p, l0, role, r))
code.append(f'mov.b32 %dv, {r}; mad.wide.u32 %da, %dt, 4, %db; '
f'add.s64 %da, %da, {4 * T * s}; @%dq2 st.global.u32 [%da], %dv;')
inserts[l1] = ('{ .reg .b64 %da, %db; .reg .b32 %dv, %dt, %dx, %dy; .reg .pred %dq, %dq2;\n'
f'mov.u32 %dx, %ctaid.x; mov.u32 %dy, %ctaid.y; setp.eq.u32 %dq, %dx, {bx}; '
f'setp.eq.and.u32 %dq2, %dy, {by}, %dq;\n'
f'mov.u32 %dt, %tid.z; mov.u32 %dx, %tid.y; mov.u32 %dy, %tid.x; '
f'mad.lo.u32 %dt, %dt, {launch.block[1]}, %dx; mad.lo.u32 %dt, %dt, {launch.block[0]}, %dy;\n'
f'ld.param.u64 %db, [{pname}+{off8}]; cvta.to.global.u64 %db, %db; '
f'add.s64 %db, %db, {arena_bytes - delta8};\n' + '\n'.join(code) + '\n}')
out = list(lines)
for l1 in sorted(inserts, reverse=True):
out[l1:l1] = [inserts[l1]]
work = ROOT / 'build/rr' / f'phase-{o.stage}'
work.mkdir(parents=True, exist_ok=True)
(work / 'probe.ptx').write_text('\n'.join(out))
pad = 4 * T * len(slots)
env = dict(os.environ, D4R_ZLUDA_WMMA='1', D4R_ZLUDA_WMMA_FP8='1', D4R_ZLUDA_WMMA_FP8_NATIVE='1',
D4R_ZLUDA_IMPLICIT_MAX_BLOCK='256')
cmd = ['systemd-run', '--user', '--scope', '-q', '-p', 'MemoryMax=24G', '-p', 'MemorySwapMax=0']
cmd += [f'--setenv={k_}={v_}' for k_, v_ in env.items() if k_.startswith('D4R_ZLUDA')]
r = subprocess.run(cmd + [sys.executable, str(HERE / 'zluda_replay.py'), os.path.expanduser(o.capture),
kernel if rr2 else f'swin_{o.stage}' + ('_diamond' if 'diamond' in kernel else ''), str(o.frame), str(work),
'--ptx', str(work / 'probe.ptx'),
'--pad', str(pad)], capture_output=True, text=True)
if 'failed' in r.stdout + r.stderr:
sys.exit(r.stdout[-2000:] + r.stderr[-3000:])
raw = np.frombuffer((work / 'arena.out.bin').read_bytes()[arena_bytes:], np.uint32).reshape(len(slots), T)
val = {(p, l0, role, r): raw[s] for s, (p, l0, role, r) in enumerate(slots)}
result = {}
waves = T // 32
for p in range(first, last + 1):
mm = phases[p - 1]
# k-step chains: C of a later mma = D of an earlier one
dset = {tuple(m[2]): n for n, m in enumerate(mm)}
prev = {n: dset.get(tuple(m[5])) for n, m in enumerate(mm)}
step = {}
for n in range(len(mm)):
s_, q = 0, n
while prev[q] is not None:
q = prev[q]
s_ += 1
step[n] = s_
ksteps = max(step.values()) + 1
a_keys = list(dict.fromkeys(tuple(m[3]) for m in mm))
b_keys = list(dict.fromkeys(tuple(m[4]) for m in mm))
# m-tile of an A fragment: its index among the A fragments of the same k-step
a_tile = {}
for s_ in range(ksteps):
for t_, key in enumerate(dict.fromkeys(tuple(m[3]) for n, m in enumerate(mm) if step[n] == s_)):
a_tile[(s_, key)] = t_
b_tile = {}
for s_ in range(ksteps):
for t_, key in enumerate(dict.fromkeys(tuple(m[4]) for n, m in enumerate(mm) if step[n] == s_)):
b_tile[(s_, key)] = t_
mt = max(t for (s_, _), t in a_tile.items()) + 1
nt = max(t for (s_, _), t in b_tile.items()) + 1
for w in range(waves):
lanes = np.arange(32) + 32 * w
A = np.zeros((16 * mt, 32 * ksteps), np.uint8)
B = np.zeros((32 * ksteps, 8 * nt), np.uint8)
C = np.zeros((16 * mt, 8 * nt), np.uint16)
D = np.zeros((16 * mt, 8 * nt), np.uint16)
for n, (l0, l1, Dr, Ar, Br, Cr) in enumerate(mm):
s_ = step[n]
i, j = a_tile[(s_, tuple(Ar))], b_tile[(s_, tuple(Br))]
for lane in range(32):
g, t = lane >> 2, lane & 3
av = [int(val[(p, l0, 'A', r)][lanes[lane]]) for r in Ar]
bv = [int(val[(p, l0, 'B', r)][lanes[lane]]) for r in Br]
for q in range(4): # A regs: (g, k 4t..), (g+8, 4t..), (g, 16+4t..), (g+8, 16+4t..)
row = 16 * i + g + 8 * (q & 1)
col = 32 * s_ + 4 * t + 16 * (q >> 1)
for byte in range(4):
A[row, col + byte] = (av[q] >> (8 * byte)) & 0xff
for q in range(2): # B regs: column g, k 4t.. and 16+4t..
for byte in range(4):
B[32 * s_ + 4 * t + 16 * q + byte, 8 * j + g] = (bv[q] >> (8 * byte)) & 0xff
for q in range(2): # C/D regs: row g (+8), columns 2t, 2t+1
row = 16 * i + g + 8 * q
if s_ == 0:
cv = int(val[(p, l0, 'C', Cr[q])][lanes[lane]])
C[row, 8 * j + 2 * t] = cv & 0xffff
C[row, 8 * j + 2 * t + 1] = cv >> 16
if s_ == ksteps - 1:
dv = int(val[(p, l0, 'D', Dr[q])][lanes[lane]])
D[row, 8 * j + 2 * t] = dv & 0xffff
D[row, 8 * j + 2 * t + 1] = dv >> 16
for name, mat in (('A', A), ('B', B), ('C', C), ('D', D)):
result[f'p{p}_w{w}_{name}'] = mat
print(f'phase {p}: {len(mm)} mma, {mt} m-tiles x {nt} n-tiles x {ksteps} k-steps -> '
f'{16 * mt}x{8 * nt}x{32 * ksteps}, lines {mm[0][0]}-{mm[-1][1]}')
np.savez_compressed(o.out, **result)
if __name__ == '__main__':
main()
engine/rr/ptx2rr.py
+980
View File
@@ -0,0 +1,980 @@
#!/usr/bin/env python3
"""Generate a GLSL compute shader for one Ray Reconstruction (DLSS-D) PTX kernel.
python3 engine/rr/ptx2rr.py MODULE.ptx KERNEL OUT.comp [--block X,Y,Z]
The translation core is vulkan/ptx2glsl.py. This module subclasses its translator and
implements the stage contract of engine/rr/README.md: binding 0 = parameter block,
binding 1 = weight buffer (argument 40), binding 2 = plane arena, then textures and
surfaces in ascending argument offset. One PTX block is one workgroup of X*Y*Z invocations.
The generated shader is derived from the user's own DLL and is model data: write it to the
model/build directory only; do not commit or redistribute it.
"""
import argparse
import importlib.util
import json
import os
import re
import struct
import sys
from pathlib import Path
import numpy as np
ROOT = Path(__file__).resolve().parents[2]
_spec = importlib.util.spec_from_file_location('ptx2glsl', ROOT / 'vulkan' / 'ptx2glsl.py')
base = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(base)
# Workgroup shapes of the 15 Swin/HKPN kernels, from the captured launch manifests.
BLOCKS = {
'enc0': (32, 1, 1), 'enc1': (32, 1, 1), 'enc2': (32, 1, 4), 'enc3': (32, 1, 4), 'enc4': (32, 1, 8),
'enc5': (32, 1, 8), 'dec0': (32, 1, 1), 'dec1': (32, 1, 2), 'dec2': (32, 1, 4), 'dec3': (32, 1, 4),
'dec4': (32, 1, 8), 'hkpn': (16, 16, 1),
}
def default_block(kernel):
m = re.search(r'swin_(enc|dec)(\d)', kernel)
key = f'{m.group(1)}{m.group(2)}' if m else ('hkpn' if 'hkpn_output' in kernel else None)
if key not in BLOCKS:
raise SystemExit(f'no default block shape for {kernel}; pass --block X,Y,Z')
return BLOCKS[key]
def braces(tok):
s = tok.strip()
if not (s.startswith('{') and s.endswith('}')):
raise base.Err(f'expected a {{...}} vector operand, got {tok}')
return [x.strip() for x in s[1:-1].split(',')]
def mem_shape(o):
"""(bytes per element, element count) of ld/st.global modifiers o."""
count = 2 if 'v2' in o else 4 if 'v4' in o else 1
size = int(re.sub(r'\D', '', o[-1])) // 8
if size not in (2, 4):
raise base.Err(f'unsupported global access width in {".".join(o)}')
return size, count
RR_RUNTIME = r'''
float16_t H(uint u) { return uint16BitsToFloat16(uint16_t(u)); }
uint HU(float16_t h) { return uint(float16BitsToUint16(h)); }
f16vec2 H2(uint u) { return uint16BitsToFloat16(u16vec2(uint16_t(u & 0xffffu), uint16_t(u >> 16))); }
uint HU2_(u16vec2 u) { return uint(u.x) | (uint(u.y) << 16); }
uint H2U(f16vec2 h) { return HU2_(float16BitsToUint16(h)); }
int s16(uint u) { return (int(u) << 16) >> 16; }
float rcp_a(float x) { return 1.0 / x; }
vec4 fetch0(sampler2D t, ivec2 c) { const ivec2 n = textureSize(t, 0); return (c.x < 0 || c.y < 0 || c.x >= n.x || c.y >= n.y) ? vec4(0.0) : texelFetch(t, c, 0); }
// x + copysign(0.5, x) rounded toward zero and then truncated: round half away from zero, exactly
float round_half_away(float x) { float t = trunc(x); return abs(x - t) >= 0.5 ? t + sign(x) : t; }
uint PARAM32(uint o) { return param.w[o >> 2]; }
uint64_t PARAM64(uint o) { return uint64_t(param.w[o >> 2]) | (uint64_t(param.w[(o >> 2) + 1u]) << 32); }
uint PARAM16(uint o) { return (param.w[o >> 2] >> ((o & 2u) * 8u)) & 0xffffu; }
uint PARAM8(uint o) { return (param.w[o >> 2] >> ((o & 3u) * 8u)) & 0xffu; }
uint SH32(uint a) { return sh[a >> 2]; }
uint SH16(uint a) { return (sh[a >> 2] >> ((a & 2u) * 8u)) & 0xffffu; }
uint bfi32(uint ins, uint base, uint pos, uint len)
{
pos &= 0xffu; len &= 0xffu;
if (pos >= 32u) return base;
uint mask = len >= 32u ? (0xffffffffu << pos) : (((1u << len) - 1u) << pos);
return (~mask & base) | (mask & (ins << pos));
}
// Global memory: weight buffer (binding 1, read-only) and plane arena (binding 2). Addresses are byte offsets.
#ifdef OOB_REPORT
// debugging (ptx2rr.py --oob-report): per buffer (0 weights read, 1 arena read, 2 arena write) a count and the lowest
// and highest out-of-range byte address, in dbg.d[3k .. 3k+2]
void OOB(uint k, uint a, uint size) { if (a + 2u > size) { atomicAdd(dbg.d[3u * k], 1u); atomicMin(dbg.d[3u * k + 1u], a); atomicMax(dbg.d[3u * k + 2u], a); } }
#define CHECKW(a) OOB(0u, a, uint(wbuf.wh.length()) * 2u)
#define CHECKR(a) OOB(1u, a, uint(abuf.ah.length()) * 2u)
#define CHECKS(a) OOB(2u, a, uint(abuf.ah.length()) * 2u)
#else
#define CHECKW(a)
#define CHECKR(a)
#define CHECKS(a)
#endif
uint GW16(uint a) { CHECKW(a); return uint(wbuf.wh[a >> 1]); }
// a 32-bit access is 4-byte aligned (PTX ISA): one word of the same buffer's word view
uint GW32(uint a) { CHECKW(a + 2u); return wbuf32.ww[a >> 2]; }
uint GA16(uint a) { CHECKR(a); return uint(abuf.ah[a >> 1]); }
uint GA32(uint a) { CHECKR(a + 2u); return uint(abuf.ah[a >> 1]) | (uint(abuf.ah[(a >> 1) + 1u]) << 16); }
void SA16(uint a, uint v) { CHECKS(a); abuf.ah[a >> 1] = uint16_t(v & 0xffffu); }
void SA32(uint a, uint v) { CHECKS(a + 2u); abuf.ah[a >> 1] = uint16_t(v & 0xffffu); abuf.ah[(a >> 1) + 1u] = uint16_t(v >> 16); }
// Surface stores with .zero out-of-bounds semantics.
void sust_p4(writeonly image2D s, int x, int y, vec4 v)
{
ivec2 n = imageSize(s);
if (uint(x) < uint(n.x) && uint(y) < uint(n.y)) imageStore(s, ivec2(x, y), v);
}
void sust_p1(writeonly image2D s, int x, int y, float v) { sust_p4(s, x, y, vec4(v, 0.0, 0.0, 0.0)); }
// Formatted stores to normalized surfaces, as the reference converts them (ZLUDA's surface_formatted_store_bits): NaN
// to zero, clamp, round half away from zero in f32. The quotient is what the image store then converts back exactly.
// UNORM 10:10:10:2 surfaces (the pooled histories) are RGBA32F images holding the quantised values, like ZLUDA's.
// tanh.approx.f32 as the reference computes it: ROCm's OCML tanh (device-libs ocml/src/tanhF.cl) with LLVM's AMDGPU
// lowering of exp (lowerFEXP, fast-FMA path): the same fused operations, hardware exp2, ldexp and reciprocal.
// (a fused f32 a * b + c through f64: the product is exact there and the sum is rounded to 53 bits before the f32
// rounding. The driver's own f32 fma is a product and a sum whenever its product also exists on its own, as in
// exp's ph / fma(x, c, -ph) pair.)
float fmad(float a, float b, float c) { precise double s = double(a) * double(b) + double(c); return float(s); }
float tanh_ocml(float x)
{
const float y = abs(x);
precise float z;
if (y < 0.625) {
precise float y2 = y * y;
precise float p = fmad(y2, fmad(y2, fmad(y2, fmad(y2, -0.005700020585209131, 0.02063407190144062), -0.053737930953502655), 0.13331416249275208), -0.3333328068256378);
precise float yp = y * p;
z = fmad(y2, yp, y);
} else {
const float w = 2.0 * y;
precise float ph = w * 1.4426950216293335;
precise float f0 = fmad(w, 1.4426950216293335, -ph);
precise float pl = fmad(w, 1.925962855864327e-08, f0);
const float e = roundEven(ph);
precise float d = ph - e;
precise float a = d + pl;
float t = ldexp(exp2(a), int(e));
if (w > 88.72283172607422) t = uintBitsToFloat(0x7f800000u);
precise float s1 = t + 1.0;
z = fmad(-2.0, rcp_a(s1), 1.0);
}
return x < 0.0 || floatBitsToUint(x) == 0x80000000u ? -z : z;
}
// tanh.approx.f16: from a table of the correctly rounded values indexed by the argument's bits (RR_TANH_TABLE: its
// byte offset in the weight buffer; engine/rr/rr2_pack.py) where a stage has one, else through f32
#ifdef RR_TANH_TABLE
float16_t tanh16(float16_t x) { return uint16BitsToFloat16(wbuf.wh[RR_TANH_TABLE / 2u + uint(float16BitsToUint16(x))]); }
#else
float16_t tanh16(float16_t x) { return float16_t(tanh(float(x))); }
#endif
vec4 q_snorm8(vec4 v)
{
v = clamp(mix(v, vec4(0.0), isnan(v)), vec4(-1.0), vec4(1.0));
precise vec4 s = v * 127.0;
precise vec4 a = abs(s) + 0.5;
return sign(s) * floor(a) / 127.0;
}
vec4 q_unorm8(vec4 v)
{
v = clamp(mix(v, vec4(0.0), isnan(v)), vec4(0.0), vec4(1.0));
precise vec4 a = v * 255.0 + 0.5;
return floor(a) / 255.0;
}
// (the quotient q / 1023 must be the correctly rounded f32 the reference stores: the driver divides by multiplying
// with a reciprocal, one step off for some q, so it comes from a table)
const float Q1023[1024] = float[1024](Q1023_VALUES);
vec4 q_unorm1010102(vec4 v)
{
v = clamp(mix(v, vec4(0.0), isnan(v)), vec4(0.0), vec4(1.0));
precise vec4 a = v * vec4(1023.0, 1023.0, 1023.0, 3.0) + 0.5;
const ivec4 q = ivec4(floor(a));
const float third[4] = float[4](0.0, 0.3333333432674408, 0.6666666865348816, 1.0);
return vec4(Q1023[q.x], Q1023[q.y], Q1023[q.z], third[q.w]);
}
// for a real UNORM 10:10:10:2 image: the reference's rounding (half away from zero) here, and any value within half a
// step of q / 1023 for the store, which the hardware quantises back to q
vec4 q_unorm1010102_hw(vec4 v)
{
v = clamp(mix(v, vec4(0.0), isnan(v)), vec4(0.0), vec4(1.0));
precise vec4 a = v * vec4(1023.0, 1023.0, 1023.0, 3.0) + 0.5;
return floor(a) * vec4(1.0 / 1023.0, 1.0 / 1023.0, 1.0 / 1023.0, 1.0 / 3.0);
}
// Binary stores write the halves' bits. RADV folds float(float16_t(x)) back to x when the half came from a
// cvt.rn.f16.f32 in a register, and the image store then truncates x instead (the output stage's history surface
// was one binary16 ulp low in half its values, which the next frame's encoder 0 amplifies): OR-ing in a zero the
// compiler cannot see (the dispatch's z count is 1) keeps the conversion's rounding.
#define SUST_Z (gl_NumWorkGroups.z - 1u)
void sust_b2(writeonly image2D s, int x, int y, uint a, uint b)
{
sust_p4(s, x >> 2, y, vec4(float(H(a | SUST_Z)), float(H(b | SUST_Z)), 0.0, 0.0));
}
void sust_b4(writeonly image2D s, int x, int y, uint a, uint b, uint c, uint d)
{
sust_p4(s, x >> 3, y, vec4(float(H(a | SUST_Z)), float(H(b | SUST_Z)), float(H(c | SUST_Z)), float(H(d | SUST_Z))));
}
// shfl.sync.up/down within a 32-lane warp (ZLUDA's SHFL_SYNC_IMPL). in_bounds is the predicate output.
uint shfl_dir(uint v, uint delta, uint opts, bool down, out bool in_bounds)
{
int section_mask = int((opts >> 8) & 31u), warp_end = int(opts & 31u);
int self = int(gl_SubgroupInvocationID & 31u);
int d = int(delta & 31u);
int idx = down ? self + d : self - d;
int sub_end = (section_mask & self) | (~section_mask & (down ? warp_end : 0));
bool oob = down ? idx > sub_end : idx < sub_end;
if (oob) idx = self;
in_bounds = !oob;
return subgroupShuffle(v, (gl_SubgroupInvocationID & 32u) | (uint(idx) & 31u)); // the warp is this half of a 64-wide wave
}
// shfl.sync.idx (PTX ISA): j = minLane | (b & ~segmask), in bounds while j <= maxLane, otherwise the lane's own value
uint shfl_idx(uint v, uint b, uint opts, out bool in_bounds)
{
uint mask = (opts >> 8) & 31u, lane = gl_SubgroupInvocationID & 31u;
uint maxLane = (lane & mask) | (opts & 31u & ~mask), j = (lane & mask) | (b & 31u & ~mask);
in_bounds = j <= maxLane;
return subgroupShuffle(v, (gl_SubgroupInvocationID & 32u) | (in_bounds ? j : lane));
}
// movmatrix.sync.trans.aligned.m8n8.b16 on one b32 fragment register per lane: lane (g = lane>>2, t = lane&3)
// holds row g, columns 2t and 2t+1. The transpose gathers (row 2t, col g) and (row 2t+1, col g).
uint movm_trans(uint src)
{
// through the subgroup's mma scratch rather than two shuffles: RADV miscompiled the shuffle form in the large
// stages (enc1's movmatrix results were wrong unless a store followed them; ZLUDA replay, 2026-10-11)
const uint base = gl_SubgroupID * 384u + 256u;
uint lane = gl_SubgroupInvocationID & 31u;
uint g = lane >> 2u, t = lane & 3u;
// order this operation's writes after every lane's reads of the previous one's results (the same scratch)
memoryBarrierShared(); subgroupBarrier();
mms[base + lane] = src;
memoryBarrierShared(); subgroupBarrier();
uint w1 = mms[base + 8u * t + (g >> 1u)];
uint w2 = mms[base + 8u * t + 4u + (g >> 1u)];
uint s = 16u * (g & 1u);
return ((w1 >> s) & 0xffffu) | (((w2 >> s) & 0xffffu) << 16);
}
// PTX min/max: a NaN operand yields the other one. GLSL's min/max leave NaN undefined (the result depended on
// codegen); SPIR-V's NMin/NMax (GLSL.std.450 79/80) are defined that way, and compile to the bare v_min/v_max
// (checked on RADV / RDNA4 for f16, f16x2 and f32 with NaN, infinity and signed zero operands, 2026-10-15)
spirv_instruction(set = "GLSL.std.450", id = 79) float16_t pmin(float16_t a, float16_t b);
spirv_instruction(set = "GLSL.std.450", id = 80) float16_t pmax(float16_t a, float16_t b);
spirv_instruction(set = "GLSL.std.450", id = 79) f16vec2 pmin2(f16vec2 a, f16vec2 b);
spirv_instruction(set = "GLSL.std.450", id = 80) f16vec2 pmax2(f16vec2 a, f16vec2 b);
spirv_instruction(set = "GLSL.std.450", id = 79) float pminf(float a, float b);
spirv_instruction(set = "GLSL.std.450", id = 80) float pmaxf(float a, float b);
// e4m3 <-> f16 (PTX cvt semantics; RNE, satfinite; NaN -> 0x7f e4m3 / 0x7e00 f16).
uint e4m3_to_h(uint b)
{
uint e = (b >> 3) & 15u, m = b & 7u, s = (b & 0x80u) << 8;
if (e == 15u && m == 7u) return s | 0x7e00u;
float f = e == 0u ? ldexp(float(m), -9) : ldexp(float(8u + m), int(max(e, 1u)) - 10);
return s | HU(float16_t(f));
}
#if defined(RR_GFX11)
// RDNA3 (engine/rr/rr_gfx11.glsl): ZLUDA's gfx11 conversions exactly. e4m3 -> f16: NaN -> 0x7fff with the sign
uint g11_e4m3_to_h(uint c)
{
const uint m = c & 0x7fu;
uint h = m >= 8u ? (m << 7u) + 0x2000u : uint(float16BitsToUint16(float16_t(float(m & 7u) * 0.001953125)));
h = m == 0x7fu ? 0x7fffu : h;
return ((c & 0x80u) << 8u) | h;
}
f16vec2 e4m3x2_to_f16x2(uint v) { return H2(g11_e4m3_to_h(v & 0xffu) | (g11_e4m3_to_h((v >> 8) & 0xffu) << 16)); }
#elif defined(SOFTWARE_E4M3)
f16vec2 e4m3x2_to_f16x2(uint v) { return H2(e4m3_to_h(v & 0xffu) | (e4m3_to_h((v >> 8) & 0xffu) << 16)); }
#else
// the hardware conversion: exact for every non-NaN code (2026-10-12); NaN keeps e4m3_to_h's sign | 0x7e00
uint e4m3_to_h_hw(uint b)
{
if ((b & 0x7fu) == 0x7fu) return ((b & 0x80u) << 8) | 0x7e00u;
return uint(float16BitsToUint16(float16_t(uintBitsToFloate4m3EXT(uint8_t(b)))));
}
f16vec2 e4m3x2_to_f16x2(uint v) { return H2(e4m3_to_h_hw(v & 0xffu) | (e4m3_to_h_hw((v >> 8) & 0xffu) << 16)); }
#endif
uint h2e4m3(float f, bool relu)
{
if (isnan(f)) return 0x7fu;
if (relu && !(f > 0.0)) return 0u;
uint s = f < 0.0 ? 0x80u : 0u;
float a = abs(f);
if (a >= 448.0) return s | 0x7eu;
if (a == 0.0) return s;
int ex;
frexp(a, ex);
int qe = max(ex - 4, -9);
uint n = uint(roundEven(ldexp(a, -qe)));
uint bits;
if (n < 8u) bits = n;
else { int p = findMSB(n); bits = (uint(qe + p + 7) << 3) | ((n >> uint(p - 3)) & 7u); }
return s | min(bits, 0x7eu);
}
#if defined(RR_GFX11)
// f16 -> e4m3 as ZLUDA computes it on gfx11 (HIP's satfinite conversion, infinities clamped first): nearest even,
// |x| > 448 -> 0x7e, NaN -> 0x7f, the sign kept; .relu first zeroes negative non-NaN halves
uint g11_h2e4m3(uint h)
{
const uint s = (h >> 8u) & 0x80u, a = h & 0x7fffu;
if (a > 0x7c00u) return s | 0x7fu;
uint c;
if (a >= 0x2400u) {
const uint t = a - 0x2000u;
c = min((t + 0x3fu + ((t >> 7u) & 1u)) >> 7u, 0x7eu);
} else {
c = uint(roundEven(float(uint16BitsToFloat16(uint16_t(a))) * 512.0));
}
return s | c;
}
#define G11_HAVE_H2E4M3
uint g11_relu(uint h) { return (h & 0x8000u) != 0u && (h & 0x7fffu) <= 0x7c00u ? 0u : h; }
uint f16x2_to_e4m3x2(f16vec2 h, bool relu)
{
uint x = HU(h.x), y = HU(h.y);
if (relu) { x = g11_relu(x); y = g11_relu(y); }
return g11_h2e4m3(x) | (g11_h2e4m3(y) << 8);
}
#elif defined(SOFTWARE_E4M3)
uint f16x2_to_e4m3x2(f16vec2 h, bool relu) { return h2e4m3(float(h.x), relu) | (h2e4m3(float(h.y), relu) << 8); }
#else
// cvt.rn.satfinite(.relu).e4m3x2.f16x2 on the hardware conversion: round to nearest even after clamping to +-448
// (satfinite), NaN to 0x7f. Equal to the software h2e4m3 for all 65536 f16 inputs (2026-10-12) and about two thirds
// of a stage's time faster (dec0: 2.85 -> 1.03 ms).
uint h2e4m3_hw(float16_t h, bool relu)
{
if (isnan(h)) return 0x7fu;
if (relu) h = h > float16_t(0.0) ? h : float16_t(0.0);
return uint(floate4m3BitsToUintEXT(floate4m3_t(clamp(h, float16_t(-448.0), float16_t(448.0)))));
}
uint f16x2_to_e4m3x2(f16vec2 h, bool relu) { return h2e4m3_hw(h.x, relu) | (h2e4m3_hw(h.y, relu) << 8); }
#endif
// mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 on the FP8 cooperative matrices, with ZLUDA's lowering of
// it (engine/layer_m.comp's TEXL): C enters the first of two 16x16x16 products, the sum stays f32 and is rounded to
// binary16 once. The operands pass through the subgroup's own scratch (MMS): NVIDIA's fragments are written in
// place (A row-major 16x32, B column-major 32x8 in a 16-column tile, C/D row-major f16 16x16), then loaded as
// whole matrices. Lane l = 4g + t holds A rows g and g+8 (k 4t.. and 16+4t..), B column g (same k), C/D row g
// and g+8, columns 2t and 2t+1.
#define MMF coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define MMH coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#ifdef RR_GFX11
// RDNA3: f16 operands holding ZLUDA's widening (value * 2^-8, k positions 2q + j <- 4q + j, 8 + 2q + j <- 4q + 2 + j
// in each K half), C scaled by 2^-16 into the two products and back (engine/rr/rr_gfx11.glsl)
#define MMA coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define MMB coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
uvec2 g11_sc4(uint packed)
{
uint mag = packed & 0x7f7f7f7fu;
mag |= (mag + 0x01010101u) & 0x80808080u;
const uint sign = packed & 0x80808080u;
return uvec2(((mag & 0xffu) << 7u) | ((mag & 0xff00u) << 15u) | ((sign & 0x80u) << 8u) | ((sign & 0x8000u) << 16u),
(((mag >> 16u) & 0xffu) << 7u) | ((mag >> 24u) << 23u) | ((sign & 0x800000u) >> 8u) | (sign & 0x80000000u));
}
#define G11_FILL(m, w0, w1, w2, w3) { \
const uvec2 s0_ = g11_sc4(w0), s1_ = g11_sc4(w1), s2_ = g11_sc4(w2), s3_ = g11_sc4(w3); f16vec2 l_; \
l_ = unpackFloat2x16(s0_.x); (m)[0] = l_.x; (m)[1] = l_.y; l_ = unpackFloat2x16(s0_.y); (m)[8] = l_.x; (m)[9] = l_.y; \
l_ = unpackFloat2x16(s1_.x); (m)[2] = l_.x; (m)[3] = l_.y; l_ = unpackFloat2x16(s1_.y); (m)[10] = l_.x; (m)[11] = l_.y; \
l_ = unpackFloat2x16(s2_.x); (m)[4] = l_.x; (m)[5] = l_.y; l_ = unpackFloat2x16(s2_.y); (m)[12] = l_.x; (m)[13] = l_.y; \
l_ = unpackFloat2x16(s3_.x); (m)[6] = l_.x; (m)[7] = l_.y; l_ = unpackFloat2x16(s3_.y); (m)[14] = l_.x; (m)[15] = l_.y; }
#define G11_HAVE_FILL
#define coopMatLoad8(m, arr, off, stride) { const uint b_ = (off) + (gl_SubgroupInvocationID & 15u) * (stride); \
G11_FILL(m, arr[b_], arr[b_ + 1u], arr[b_ + 2u], arr[b_ + 3u]) }
#define MM8(x1, y1, x0, y0, c) MMF acc_ = MMF(c); for (int i_ = 0; i_ < acc_.length(); ++i_) acc_[i_] *= 1.52587890625e-05; \
acc_ = coopMatMulAdd(x1, y1, coopMatMulAdd(x0, y0, acc_)); for (int i_ = 0; i_ < acc_.length(); ++i_) acc_[i_] *= 65536.0;
#else
#define MMA coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define MMB coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#define coopMatLoad8(m, arr, off, stride) coopMatLoad(m, arr, off, stride, gl_CooperativeMatrixLayoutRowMajor)
#define MM8(x1, y1, x0, y0, c) MMF acc_ = coopMatMulAdd(x1, y1, coopMatMulAdd(x0, y0, MMF(c)));
#endif
void mma_e4m3_f16(uint a0, uint a1, uint a2, uint a3, uint b0, uint b1, uint c0, uint c1, out uint d0, out uint d1)
{
const uint base = gl_SubgroupID * 384u, A = base, B = base + 128u, C = base + 256u;
const uint lane = gl_SubgroupInvocationID & 31u, g = lane >> 2u, t = lane & 3u;
// order this operation's writes after every lane's reads of the previous one's results (the same scratch)
memoryBarrierShared(); subgroupBarrier();
mms[A + g * 8u + t] = a0; mms[A + (g + 8u) * 8u + t] = a1;
mms[A + g * 8u + 4u + t] = a2; mms[A + (g + 8u) * 8u + 4u + t] = a3;
mms[B + g * 8u + t] = b0; mms[B + g * 8u + 4u + t] = b1;
mms[B + (g + 8u) * 8u + t] = 0u; mms[B + (g + 8u) * 8u + 4u + t] = 0u;
mms[C + g * 8u + t] = c0; mms[C + (g + 8u) * 8u + t] = c1;
mms[C + g * 8u + 4u + t] = 0u; mms[C + (g + 8u) * 8u + 4u + t] = 0u;
memoryBarrierShared(); subgroupBarrier();
MMA x0, x1; MMB y0, y1; MMH c;
#ifdef RR_GFX11
coopMatLoad8(x0, mms, A, 8u);
coopMatLoad8(x1, mms, A + 4u, 8u);
coopMatLoad8(y0, mms, B, 8u);
coopMatLoad8(y1, mms, B + 4u, 8u);
coopMatLoad(c, mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor);
MM8(x1, y1, x0, y0, c)
memoryBarrierShared(); subgroupBarrier();
coopMatStore(MMH(acc_), mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor);
#else
coopMatLoad(x0, mms, A, 8u, gl_CooperativeMatrixLayoutRowMajor);
coopMatLoad(x1, mms, A + 4u, 8u, gl_CooperativeMatrixLayoutRowMajor);
coopMatLoad(y0, mms, B, 8u, gl_CooperativeMatrixLayoutColumnMajor);
coopMatLoad(y1, mms, B + 4u, 8u, gl_CooperativeMatrixLayoutColumnMajor);
coopMatLoad(c, mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor);
MMF acc = coopMatMulAdd(x1, y1, coopMatMulAdd(x0, y0, MMF(c)));
memoryBarrierShared(); subgroupBarrier();
coopMatStore(MMH(acc), mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor);
#endif
memoryBarrierShared(); subgroupBarrier();
d0 = mms[C + g * 8u + t]; d1 = mms[C + (g + 8u) * 8u + t];
}
// Two m16n8k32 products with the same A fragment (NVIDIA issues every one of them in such adjacent pairs): one 16x16
// product, the first's B, C and D in columns 0-7, the second's in columns 8-15.
void mma_e4m3_f16_pair(uint a0, uint a1, uint a2, uint a3, uint b0, uint b1, uint c0, uint c1,
uint e0, uint e1, uint f0, uint f1, out uint d0, out uint d1, out uint h0, out uint h1)
{
const uint base = gl_SubgroupID * 384u, A = base, B = base + 128u, C = base + 256u;
const uint lane = gl_SubgroupInvocationID & 31u, g = lane >> 2u, t = lane & 3u;
memoryBarrierShared(); subgroupBarrier();
mms[A + g * 8u + t] = a0; mms[A + (g + 8u) * 8u + t] = a1;
mms[A + g * 8u + 4u + t] = a2; mms[A + (g + 8u) * 8u + 4u + t] = a3;
mms[B + g * 8u + t] = b0; mms[B + g * 8u + 4u + t] = b1;
mms[B + (g + 8u) * 8u + t] = e0; mms[B + (g + 8u) * 8u + 4u + t] = e1;
mms[C + g * 8u + t] = c0; mms[C + (g + 8u) * 8u + t] = c1;
mms[C + g * 8u + 4u + t] = f0; mms[C + (g + 8u) * 8u + 4u + t] = f1;
memoryBarrierShared(); subgroupBarrier();
MMA x0, x1; MMB y0, y1; MMH c;
#ifdef RR_GFX11
coopMatLoad8(x0, mms, A, 8u);
coopMatLoad8(x1, mms, A + 4u, 8u);
coopMatLoad8(y0, mms, B, 8u);
coopMatLoad8(y1, mms, B + 4u, 8u);
coopMatLoad(c, mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor);
MM8(x1, y1, x0, y0, c)
memoryBarrierShared(); subgroupBarrier();
coopMatStore(MMH(acc_), mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor);
#else
coopMatLoad(x0, mms, A, 8u, gl_CooperativeMatrixLayoutRowMajor);
coopMatLoad(x1, mms, A + 4u, 8u, gl_CooperativeMatrixLayoutRowMajor);
coopMatLoad(y0, mms, B, 8u, gl_CooperativeMatrixLayoutColumnMajor);
coopMatLoad(y1, mms, B + 4u, 8u, gl_CooperativeMatrixLayoutColumnMajor);
coopMatLoad(c, mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor);
MMF acc = coopMatMulAdd(x1, y1, coopMatMulAdd(x0, y0, MMF(c)));
memoryBarrierShared(); subgroupBarrier();
coopMatStore(MMH(acc), mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor);
#endif
memoryBarrierShared(); subgroupBarrier();
d0 = mms[C + g * 8u + t]; d1 = mms[C + (g + 8u) * 8u + t];
h0 = mms[C + g * 8u + 4u + t]; h1 = mms[C + (g + 8u) * 8u + 4u + t];
}
'''
def statement_lines(ptx, kernel):
"""the source line of every statement of KERNEL, in the order vulkan/ptx2glsl.py's tokenizer yields them"""
text = re.sub(r'//[^\n]*', '', ptx.replace('\r', ''))
m = re.search(r'\.entry\s+' + re.escape(kernel) + r'\s*\((.*?)\)', text, re.S)
start = text.index('{', m.end())
depth, k = 0, start
while True:
if text[k] == '{':
depth += 1
elif text[k] == '}':
depth -= 1
if depth == 0:
break
k += 1
lines, i, line = [], start + 1, text.count('\n', 0, start + 1) + 1
while i < k:
c = text[i]
if c.isspace():
line += c == '\n'
i += 1
continue
if c in '{}':
i += 1
continue
lm = re.match(r'(\$?[A-Za-z_][\w$]*):', text[i:i + 200])
if lm:
i += lm.end()
continue
j = text.index(';', i)
lines.append(line)
line += text.count('\n', i, j)
i = j + 1
return lines
def base_helper(name):
"""one helper function of vulkan/ptx2glsl.py's runtime, with its leading comment"""
lines = base.RUNTIME.split('\n')
start = next(i for i, l in enumerate(lines) if re.match(r'\w+ ' + name + r'\(', l))
while lines[start - 1].startswith('//'):
start -= 1
end = start
if not lines[start].rstrip().endswith('}'): # one-line helpers end on their own line
while lines[end] != '}':
end += 1
return '\n'.join(lines[start:end + 1]) + '\n'
RR_RUNTIME = RR_RUNTIME.replace('Q1023_VALUES', ', '.join(repr(float(np.float32(q) / np.float32(1023))) for q in range(1024)))
RR_RUNTIME += base_helper('prmt') + base_helper('shfl_bfly').replace(
'return subgroupShuffle(v, (gl_SubgroupInvocationID & 32u) | uint(idx));',
'return subgroupShuffle(v, (gl_SubgroupInvocationID & 32u) | (uint(idx) & 31u));')
class RRT(base.T):
def __init__(self, ptx, kernel, block, weights_arg=40, buffers=()):
self.weights_arg = weights_arg
self.buffers = sorted(buffers) # argument offsets of further buffers, each its own binding
self.snorm16 = set() # texture arguments of --snorm16
bx, by, bz = block
if bx * by * bz % 32:
raise base.Err('block size must be a multiple of the 32-lane subgroup')
self.bx, self.by, self.bz = bx, by, bz
self.sreg = {
'%tid.x': f'(gl_LocalInvocationIndex % {bx}u)',
'%tid.y': f'((gl_LocalInvocationIndex / {bx}u) % {by}u)',
'%tid.z': f'(gl_LocalInvocationIndex / {bx * by}u)',
'%ntid.x': f'{bx}u', '%ntid.y': f'{by}u', '%ntid.z': f'{bz}u',
'%ctaid.x': 'gl_WorkGroupID.x', '%ctaid.y': 'gl_WorkGroupID.y', '%ctaid.z': 'gl_WorkGroupID.z',
'%nctaid.x': 'gl_NumWorkGroups.x', '%nctaid.y': 'gl_NumWorkGroups.y', '%nctaid.z': 'gl_NumWorkGroups.z',
'%laneid': 'gl_SubgroupInvocationID',
'WARP_SZ': '32u',
}
self.origin = {}
self.f64 = {} # %fd registers: the kernels only load f64 immediates into them, to round them to f16
self.probes = {} # debugging: PTX line -> registers stored to the debug buffer after it executes
self.probe_block = (0, 0)
self.probe_slots = 0
self.oob = False
# every register precise, the integer ones too: they carry f16x2 and f32 bit patterns between operations, and
# without it the compiler may contract a product and a sum across them depending on scheduling (results then
# changed with unrelated code and with the device's robustness features)
self.precise_all = True
# module globals with an initialiser (the output filter's bicubic phase table): constants, laid out one after
# another in a blob that is the stage's weight buffer (binding 1); `mov.u64 %rd, NAME` is its offset there
self.globals, self.globals_blob = {}, b''
for gm in re.finditer(r'\.global\s+\.align\s+\d+\s+\.b8\s+(\w+)\[(\d+)\]\s*=\s*\{([^}]*)\}', ptx):
self.globals[gm.group(1)] = len(self.globals_blob)
data = bytes(int(x) for x in gm.group(3).split(','))
self.globals_blob += data + bytes(-len(data) % 16)
self.stmt_lines = statement_lines(ptx, kernel)
self.stmt_index = 0
super().__init__(ptx, kernel)
if not os.environ.get('D4R_PTX2RR_NO_PAIR'):
self.pair_mma()
self.origins()
def stmt(self, s):
line = self.stmt_lines[self.stmt_index] if self.stmt_index < len(self.stmt_lines) else None
self.stmt_index += 1
super().stmt(s)
for reg in self.probes.get(line, ()):
# the debug buffer: 32-bit words [slot][local invocation], the last execution of the statement
n = self.reg(reg)
g = self.regs[n]
v = (f'floatBitsToUint({n})' if g == 'float' else f'uint({n})' if g in ('uint', 'uint64_t', 'bool') else
f'HU({n})' if g == 'float16_t' else f'H2U({n})')
bx, by = self.probe_block
self.E(f'if (gl_WorkGroupID.xy == uvec2({bx}u, {by}u)) dbg.d[{self.probe_slots}u * (gl_WorkGroupSize.x * gl_WorkGroupSize.y) + '
f'gl_LocalInvocationIndex] = {v};')
self.probe_slots += 1
def pair_mma(self):
"""merge adjacent mma statements with the same A fragment and independent operands into one paired product;
the second statement becomes a no-op so statement numbering (probes) is unchanged"""
stmts = [k for k, (kind, _) in enumerate(self.toks) if kind == 'stmt']
n = 0
while n + 1 < len(stmts):
i, j = stmts[n], stmts[n + 1]
x, y = self.toks[i][1], self.toks[j][1]
if x.startswith('mma.sync.aligned.m16n8k32') and y.startswith('mma.sync.aligned.m16n8k32'):
gx = [braces(t) for t in base.split_operands(x.split(None, 1)[1])]
gy = [braces(t) for t in base.split_operands(y.split(None, 1)[1])]
if gx[1] == gy[1] and not set(gx[0]) & set(gy[2] + gy[3] + gy[1]):
ops = ', '.join('{' + ', '.join(g) + '}' for g in (gx[0], gy[0], gx[1], gx[2], gx[3], gy[2], gy[3]))
self.toks[i] = ('stmt', 'mma.pair.e4m3 ' + ops)
self.toks[j] = ('stmt', 'mma.paired.noop')
n += 2
continue
n += 1
# ------------------------------------------------ SREG remap onto the flat workgroup
def is_reg(self, tok):
return tok.strip() not in self.sreg and super().is_reg(tok)
def val(self, tok, t):
tok = tok.strip()
if re.fullmatch(r'(0[xX][0-9a-fA-F]+|-?\d+)[uU]', tok):
tok = tok[:-1]
tok = re.sub(r'^0X', '0x', tok)
if tok in self.sreg:
e = self.sreg[tok]
return f'int({e})' if t == 's32' else f'uint64_t({e})' if t in ('u64', 'b64') else f'uint({e})'
return super().val(tok, t)
# ------------------------------------------------ global memory origin: weight buffer or arena
def origins(self):
stmts = []
for kind, s in self.toks:
if kind != 'stmt' or s.startswith('.'):
continue
s = re.sub(r'^@!?%p\d+\s+', '', s)
parts = s.split(None, 1)
stmts.append((parts[0], base.split_operands(parts[1]) if len(parts) > 1 else []))
org = self.origin
changed = True
while changed:
changed = False
for op, a in stmts:
if not a or not re.fullmatch(r'%rd\d+', a[0]) or op.startswith(('ld.', 'st.', 'bra')):
if op != 'ld.param.u64':
continue
if op.startswith('mov') and len(a) > 1 and a[1].strip() in self.globals:
new = {'w'}
elif op == 'ld.param.u64':
# the argument's true offset (the load may go through a register based inside the block)
off_ = self.param_off(a[1])
new = {'w' if off_ == self.weights_arg else f'x{off_}' if off_ in self.buffers else 'a'}
else:
new = set()
for x in a[1:]:
if re.fullmatch(r'%rd\d+', x):
new |= org.get(x, set())
cur = org.setdefault(a[0], set())
if new - cur:
cur |= new
changed = True
def gorigin(self, reg):
s = self.origin.get(reg, set())
if len(s) != 1:
raise base.Err(f'global pointer {reg} has origin {sorted(s)} (need exactly one of w/a)')
return next(iter(s))
def gfn(self, reg, store=False):
o = self.gorigin(reg)
if o == 'w':
return 'GW'
if o == 'a':
return 'SA' if store else 'GA'
return ('SX' if store else 'GX') + o[1:] + '_'
def gaddr(self, tok):
m = re.fullmatch(r'\[\s*(%rd\d+)\s*(?:\+\s*(-?\d+))?\s*\]', tok.strip())
if not m:
raise base.Err(f'unsupported global address {tok}')
return m.group(1), int(m.group(2) or 0)
def gbyte(self, reg, off):
return f'uint({self.val(reg, "u64")}) + {off}u' if off >= 0 else f'uint({self.val(reg, "u64")}) - {-off}u'
def ld(self, o, a):
if 'global' not in o:
return super().ld(o, a)
size, count = mem_shape(o)
reg, off = self.gaddr(a[1])
fn = self.gfn(reg)
dsts = braces(a[0]) if count > 1 else [a[0]]
for i, d in enumerate(dsts):
self.A(d, f'{fn}{size * 8}({self.gbyte(reg, off + i * size)})', 'u16' if size == 2 else 'u32')
def st(self, o, a):
if 'global' not in o:
return super().st(o, a)
size, count = mem_shape(o)
reg, off = self.gaddr(a[0])
if self.gorigin(reg) == 'w':
raise base.Err(f'store through the read-only weight pointer: {a}')
fn = self.gfn(reg, True)
vals = braces(a[1]) if count > 1 else [a[1]]
for i, v in enumerate(vals):
if size == 2:
self.E(f'{fn}16({self.gbyte(reg, off + i * size)}, {self.val(v, "u16")});')
else:
self.E(f'{fn}32({self.gbyte(reg, off + i * size)}, {self.val(v, "u32")});')
# ------------------------------------------------ textures and surfaces
def handle_off(self, h):
off = self.handle.get(h)
if off is None:
raise base.Err(f'{h} is not a texture or surface handle loaded from the parameter block')
return off
def set_surface(self, off, fmt):
old = self.surfaces.get(off)
if old is not None and old != fmt:
raise base.Err(f'surface {off} used with formats {old} and {fmt}')
self.surfaces[off] = fmt
def tex(self, op, a):
o = op.split('.')[1:]
dst = braces(a[0])
h, c = self.texref(a[1])
off = self.handle_off(h)
self.textures.add(off)
if 'base' in o and o[-1] == 's32':
# an integer fetch outside the image: CUDA applies the texture's address mode (clamp for every denoiser
# texture), but the reference this engine is checked against is ZLUDA, whose image load returns zero there
# (zluda_ptx_impl.cpp fetch_2D). The output stage's specular reprojection reads such neighbours at the image
# border: with clamp its surfaces were 99.92% identical to the capture, with zero 99.997% (capDX, 2026-10-12)
fetch = f'fetch0(T{off}, ivec2({self.val(c[0], "s32")}, {self.val(c[1], "s32")}))'
elif ('base' in o or o[0] == '2d') and o[-1] == 'f32': # tex.2d: no derivatives in a kernel, level 0
fetch = f'textureLod(T{off}, vec2({self.val(c[0], "f32")}, {self.val(c[1], "f32")}), 0.0)'
elif 'level' in o and o[-1] == 'f32':
fetch = (f'textureLod(T{off}, vec2({self.val(c[0], "f32")}, {self.val(c[1], "f32")}), '
f'{self.val(a[2], "f32")})')
else:
raise base.Err(f'unsupported texture op {op}')
f16 = 'f16' in o
self.E(f'{{ vec4 t_ = {fetch};')
if off in self.snorm16 and o[-1] == 'f32':
# --snorm16: a filtered fetch of an 8-bit SNORM texture on an NVIDIA card is this card's value (8-bit weights)
# rounded to a 16-bit step (measured on enc0's reprojected normals, build/rr/f/nv; D4R_RR_FUSE models)
self.E('{ precise vec4 q_ = roundEven(t_ * 32767.0); precise vec4 d_ = q_ / 32767.0; t_ = d_; }')
for i, d in enumerate(dst):
# a binary16 result is the texel rounded to binary16: through H(HU(...) | SUST_Z) so that the driver cannot
# fold the rounding away when the value is widened again (RR_RUNTIME's note on sust_b)
self.A(d, f'H(HU(float16_t(t_[{i}])) | SUST_Z)' if f16 else f't_[{i}]', 'f16' if f16 else 'float')
self.E('}')
def sust(self, op, a):
h, c = self.texref(a[0])
off = self.handle_off(h)
v = braces(a[1])
x, y = self.val(c[0], 's32'), self.val(c[1], 's32')
S = f'S{off}'
op = op.replace('.cs.', '.') # the streaming cache policy: a hint, the store is the same
if op == 'sust.p.2d.v4.b32.zero':
self.set_surface(off, 'rgba16f')
self.E(f'sust_p4({S}, {x}, {y}, SUSTQ_{off}(vec4({", ".join(self.val(q, "f32") for q in v)})));')
elif op == 'sust.p.2d.b32.zero':
self.set_surface(off, 'r16f')
self.E(f'sust_p1({S}, {x}, {y}, {self.val(v[0], "f32")});')
elif op == 'sust.b.2d.v2.b16.zero':
self.set_surface(off, 'rg16f')
self.E(f'sust_b2({S}, {x}, {y}, {", ".join(self.val(q, "u16") for q in v)});')
elif op == 'sust.b.2d.v4.b16.zero':
self.set_surface(off, 'rgba16f')
self.E(f'sust_b4({S}, {x}, {y}, {", ".join(self.val(q, "u16") for q in v)});')
else:
raise base.Err(f'unsupported surface op {op}')
# ------------------------------------------------ warp-level ops
def mma(self, a):
d, x, b, c = (braces(t) for t in a)
if not (len(d) == 2 and len(x) == 4 and len(b) == 2 and len(c) == 2):
raise base.Err('mma operand shape')
args = ', '.join(self.val(r, 'u32') for r in x + b + c)
self.E('{ uint md0_, md1_;')
self.E(f'mma_e4m3_f16({args}, md0_, md1_);')
self.A(d[0], 'md0_', 'u32')
self.A(d[1], 'md1_', 'u32')
self.E('}')
def shfl_dir(self, op, a):
down = op.split('.')[2] == 'down'
d = a[0].split('|')
self.E('{ bool ib_;')
self.A(d[0], f'shfl_dir({self.val(a[1], "u32")}, {self.val(a[2], "u32")}, {self.val(a[3], "u32")}, '
f'{"true" if down else "false"}, ib_)', 'u32')
if len(d) > 1:
self.A(d[1], 'ib_', 'bool')
self.E('}')
def cvt_to_e4m3(self, op, a):
relu = 'relu' in op.split('.')
self.A(a[0], f'f16x2_to_e4m3x2({self.val(a[1], "f16x2")}, {"true" if relu else "false"})', 'u16')
def cvt(self, op, o, a):
types = [t for t in o if re.match(r'^(f|s|u|b)\d+(x2)?$', t)]
if types[:2] in (['u32', 'u64'], ['s32', 's64'], ['u32', 's64'], ['s32', 'u64']):
return self.A(a[0], f'uint({self.val(a[1], "u64")})', 'u32')
if types[:2] == ['f16', 'f32'] and 'rn' in o and 'satfinite' not in o:
# Two things the driver must not do. It folds float(float16_t(x)) to x when the result is widened again
# (RR_RUNTIME's note on sust_b): the OR on the binary16 bits. And it fuses a product or sum with this
# conversion into v_fma_mix, one rounding of the exact value, where the kernel rounds to f32 first: an f32
# result that is exactly a binary16 midpoint then rounds the other way (the output stages' rare one-step
# differences, e.g. log2(x) * c1 * c2; found 2026-10-13): the OR on the f32 bits keeps the f32 value.
return self.A(a[0], f'H(HU(float16_t(uintBitsToFloat(floatBitsToUint({self.val(a[1], "f32")}) | SUST_Z))) | SUST_Z)', 'f16')
return super().cvt(op, o, a)
def f64_op(self, op, a):
if op == 'mov.f64' and re.fullmatch(r'0[dD][0-9A-Fa-f]{16}', a[1].strip()):
self.f64[a[0].strip()] = struct.unpack('<d', int(a[1].strip()[2:], 16).to_bytes(8, 'little'))[0]
return True
if op == 'cvt.rn.f16.f64':
v = self.f64.get(a[1].strip())
if v is None:
raise base.Err(f'{a[1]} is not a known f64 constant')
bits = int(np.float16(v).view(np.uint16))
self.A(a[0], f'{bits}u', 'u16')
return True
return False
def op(self, op, a):
parts = op.split('.')
if '%fd' in ','.join(a) and self.f64_op(op, a):
return
if op.startswith('mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16'):
return self.mma(a)
if op == 'mma.pair.e4m3':
d, h, x, b, c, e, f = (braces(t) for t in a)
args = ', '.join(self.val(r, 'u32') for r in x + b + c + e + f)
self.E('{ uint md0_, md1_, mh0_, mh1_;')
self.E(f'mma_e4m3_f16_pair({args}, md0_, md1_, mh0_, mh1_);')
for reg, v in zip(d + h, ('md0_', 'md1_', 'mh0_', 'mh1_')):
self.A(reg, v, 'u32')
return self.E('}')
if op == 'mma.paired.noop':
return
if parts[0] == 'mov' and len(a) > 1 and a[1].strip() in self.globals:
return self.A(a[0], f'uint64_t({self.globals[a[1].strip()]}u)', 'u64')
# cross-kernel progress counters (NGX runs the output passes concurrently and has each wait for the rows it
# needs): plain loads and adds here. The engine runs the passes one after another and leaves the counter
# pointers null, which the kernels test, so these never execute.
if parts[0] == 'ld' and ('relaxed' in parts or 'acquire' in parts) and 'global' in parts:
return self.ld([q for q in parts[1:] if q not in ('relaxed', 'acquire', 'gpu') and '::' not in q and not q.startswith('L1')], a)
if parts[0] == 'red' and parts[-2] == 'add':
reg, off = self.gaddr(a[0])
fn, sfn = self.gfn(reg), self.gfn(reg, True)
return self.E(f'{sfn}32({self.gbyte(reg, off)}, {fn}32({self.gbyte(reg, off)}) + {self.val(a[1], "u32")});')
if op == 'membar.gl':
return
if op == 'activemask.b32':
# the warp's lanes: this half of a 64-wide wave
return self.A(a[0], '((gl_SubgroupInvocationID & 32u) != 0u ? subgroupBallot(true).y : subgroupBallot(true).x)', 'u32')
if op == 'movmatrix.sync.trans.aligned.m8n8.b16':
return self.A(a[0], f'movm_trans({self.val(a[1], "u32")})', 'u32')
# PTX clamps shift amounts at the register width (shl/shr.u give 0, shr.s the sign); GLSL leaves them undefined
if parts[0] in ('shl', 'shr') and len(parts) == 2 and parts[1] in ('b32', 'u32', 's32'):
x, n = self.val(a[1], 's32' if parts[1] == 's32' else 'u32'), self.val(a[2], 'u32')
if parts[0] == 'shl':
return self.A(a[0], f'(({n}) >= 32u ? 0u : ({x}) << ({n}))', 'u32')
if parts[1] == 's32':
return self.A(a[0], f'uint(({x}) >> min({n}, 31u))', 'u32')
return self.A(a[0], f'(({n}) >= 32u ? 0u : ({x}) >> ({n}))', 'u32')
if op == 'shfl.sync.idx.b32':
d = a[0].split('|')
self.E('{ bool ib_;')
self.A(d[0], f'shfl_idx({self.val(a[1], "u32")}, {self.val(a[2], "u32")}, {self.val(a[3], "u32")}, ib_)', 'u32')
if len(d) > 1:
self.A(d[1], 'ib_', 'bool')
return self.E('}')
if op in ('shfl.sync.down.b32', 'shfl.sync.up.b32'):
return self.shfl_dir(op, a)
if parts[0] == 'sust':
return self.sust(op, a)
if parts[0] == 'tex':
return self.tex(op, a)
if parts[0] == 'cvt' and [t for t in parts[1:] if t in ('f16x2', 'e4m3x2')] == ['e4m3x2', 'f16x2']:
return self.cvt_to_e4m3(op, a)
if op == 'copysign.f32': # PTX: the sign of the first source on the value of the second
x, y = self.val(a[1], 'f32'), self.val(a[2], 'f32')
return self.A(a[0], f'uintBitsToFloat((floatBitsToUint({y}) & 0x7fffffffu) | (floatBitsToUint({x}) & 0x80000000u))', 'float')
if op == 'tanh.approx.f16':
return self.A(a[0], f'tanh16({self.val(a[1], "f16")})', 'f16')
if op == 'tanh.approx.f32':
return self.A(a[0], f'tanh_ocml({self.val(a[1], "f32")})', 'float')
if op == 'ex2.approx.f16':
# v_exp_f32 of the widened value, rounded to binary16; the ORs keep that rounding (see cvt)
return self.A(a[0], f'H(HU(float16_t(uintBitsToFloat(floatBitsToUint(exp2(float({self.val(a[1], "f16")}))) | SUST_Z))) | SUST_Z)', 'f16')
if op == 'neg.f16':
return self.A(a[0], f'-{self.val(a[1], "f16")}', 'f16')
if parts[0] in ('min', 'max') and 'NaN' not in parts and any(t in parts for t in ('f16', 'f16x2', 'f32')):
t = 'f16x2' if 'f16x2' in parts else 'f16' if 'f16' in parts else 'f32'
fn = {'f16x2': 'p{}2', 'f16': 'p{}', 'f32': 'p{}f'}[t].format(parts[0])
x, y = self.val(a[1], t), self.val(a[2], t)
return self.A(a[0], f'{fn}({x}, {y})', 'float' if t == 'f32' else t)
if op == 'neg.f16x2':
return self.A(a[0], f'-{self.val(a[1], "f16x2")}', 'f16x2')
if op == 'abs.f16x2':
return self.A(a[0], f'abs({self.val(a[1], "f16x2")})', 'f16x2')
if parts[0] == 'set' and parts[2] == 'f16':
c = ('(isnan({0}) || isnan({1}))' if parts[1] == 'nan' else
'!(isnan({0}) || isnan({1}))' if parts[1] == 'num' else self.cmp(set(parts), True))
return self.A(a[0], f'(({c.format(self.val(a[1], "f16"), self.val(a[2], "f16"))}) ? 1.0hf : 0.0hf)', 'f16')
return super().op(op, a)
# ------------------------------------------------ stage contract
def preamble(self, w, tex_filter):
w('#version 460')
w('#extension GL_EXT_shader_explicit_arithmetic_types : require')
w('#extension GL_EXT_shader_16bit_storage : require')
w('#extension GL_KHR_shader_subgroup_basic : require')
w('#extension GL_KHR_shader_subgroup_shuffle : require')
w('#extension GL_KHR_shader_subgroup_ballot : require')
w('#extension GL_KHR_cooperative_matrix : require')
w('#extension GL_KHR_memory_scope_semantics : require')
w('#ifndef RR_GFX11')
w('#extension GL_EXT_float_e4m3 : require')
w('#endif')
w('#extension GL_EXT_spirv_intrinsics : require')
w(f'// Generated by engine/rr/ptx2rr.py from {self.kernel}. Model data derived from the user\'s DLL; do not commit.')
w(f'layout(local_size_x = {self.bx * self.by * self.bz}, local_size_y = 1, local_size_z = 1) in;')
w('layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint w[]; } param;')
w('layout(set = 0, binding = 1, std430) readonly buffer WeightBuf { uint16_t wh[]; } wbuf;')
w('layout(set = 0, binding = 1, std430) readonly buffer WeightBuf32 { uint ww[]; } wbuf32;')
w('layout(set = 0, binding = 2, std430) buffer ArenaBuf { uint16_t ah[]; } abuf;')
n = 3
for off in sorted(self.textures):
w(f'layout(set = 0, binding = {n}) uniform sampler2D T{off};')
n += 1
for off in sorted(self.surfaces):
w(f'layout(set = 0, binding = {n}, {self.surfaces[off]}) uniform writeonly image2D S{off};')
q = {'rgba8_snorm': 'q_snorm8', 'rgba8': 'q_unorm8', 'rgba32f': 'q_unorm1010102', 'rgb10_a2': 'q_unorm1010102_hw'}.get(self.surfaces[off], '')
w(f'#define SUSTQ_{off}(v) {q}(v)')
n += 1
for off in self.buffers:
w(f'layout(set = 0, binding = {n}, std430) buffer XBuf{off} {{ uint16_t h[]; }} xb{off};')
w(f'uint GX{off}_16(uint a) {{ return uint(xb{off}.h[a >> 1]); }}')
w(f'uint GX{off}_32(uint a) {{ return uint(xb{off}.h[a >> 1]) | (uint(xb{off}.h[(a >> 1) + 1u]) << 16); }}')
w(f'void SX{off}_16(uint a, uint v) {{ xb{off}.h[a >> 1] = uint16_t(v & 0xffffu); }}')
w(f'void SX{off}_32(uint a, uint v) {{ xb{off}.h[a >> 1] = uint16_t(v & 0xffffu); xb{off}.h[(a >> 1) + 1u] = uint16_t(v >> 16); }}')
n += 1
w(f'shared uint sh[{max(1, (self.shared_bytes + 3) // 4)}];')
if self.probe_slots or self.oob:
w(f'layout(set = 0, binding = {n}, std430) buffer Debug {{ uint d[]; }} dbg;')
if self.oob:
w('#define OOB_REPORT 1')
w(f'shared uint mms[{384 * (self.bx * self.by * self.bz // 32)}]; // mma_e4m3_f16 scratch, per subgroup')
def runtime(self):
return RR_RUNTIME
def main():
ap = argparse.ArgumentParser(description='PTX kernel -> GLSL compute shader (Ray Reconstruction stage contract)')
ap.add_argument('module')
ap.add_argument('kernel')
ap.add_argument('out')
ap.add_argument('--block', help='workgroup shape X,Y,Z (default: from the kernel table)')
ap.add_argument('--surface', action='append', default=[], metavar='OFFSET=FORMAT',
help="the GLSL image format of a surface whose stores are formatted (sust.p), when it is not the "
"store's own guess, e.g. 144=rgba8_snorm (repeatable)")
ap.add_argument('--snorm16', action='append', type=int, default=[], metavar='OFFSET',
help='texture argument whose filtered fetches are rounded to 16-bit steps, as an NVIDIA card returns them')
ap.add_argument('--weights-arg', type=int, default=40,
help='argument offset of the weight pointer (40 for the DLSS-D kernels, 24 for rr2_*)')
ap.add_argument('--buffer', action='append', type=int, default=[], metavar='OFFSET',
help='a pointer argument that is neither the weights nor the arena: its own storage buffer, bound '
'after the surfaces in ascending argument offset (repeatable)')
ap.add_argument('--probe', action='append', default=[], metavar='LINE:%REG,...',
help='debugging: store the registers after the statement at that module line executes (last '
'execution) for workgroup --probe-block, into a debug buffer bound after the surfaces')
ap.add_argument('--probe-block', default='0,0')
ap.add_argument('--oob-report', action='store_true', help='debugging: count out-of-range buffer accesses')
ns = ap.parse_args()
block = tuple(int(x) for x in ns.block.split(',')) if ns.block else default_block(ns.kernel)
formats = {int(k): v for k, v in (x.split('=') for x in ns.surface)}
src = open(ns.module).read()
probes = {}
for spec in ns.probe:
line, regs = spec.split(':', 1)
probes[int(line)] = [r.strip() for r in regs.split(',')]
def make():
t = RRT(src, ns.kernel, block, ns.weights_arg, ns.buffer)
t.snorm16 = set(ns.snorm16)
t.probes = probes
t.probe_block = tuple(int(x) for x in ns.probe_block.split(','))
t.oob = ns.oob_report
return t
t = base.build(make, False, False, False)
for off, fmt in formats.items():
if off in t.surfaces:
t.surfaces[off] = fmt
open(ns.out, 'w').write(t.emit({}))
if t.globals_blob:
Path(ns.out).with_suffix('.globals.bin').write_bytes(t.globals_blob)
# the binding order (stage contract): textures, then surfaces, in ascending argument offset
Path(ns.out).with_suffix('.json').write_text(json.dumps(
{'kernel': ns.kernel, 'block': block, 'textures': sorted(t.textures),
'surfaces': [[off, t.surfaces[off]] for off in sorted(t.surfaces)], 'buffers': t.buffers, 'shared': t.shared_bytes,
'probes': [[line, reg] for line in sorted(probes) for reg in probes[line]], 'oob': ns.oob_report}) + '\n')
sys.stderr.write(f'block={block} textures={sorted(t.textures)} surfaces={t.surfaces} shared={t.shared_bytes}\n')
if __name__ == '__main__':
main()
engine/rr/ptx_expr.py
+108
View File
@@ -0,0 +1,108 @@
#!/usr/bin/env python3
"""Print the expression tree that defines a PTX register, followed back through arithmetic, pair packs/unpacks,
conversions and warp shuffles down to loads, constants and registers named as leaves.
usage: ptx_expr.py STAGE LINE %REG [--depth N] [--stop %r1,%r2]
LINE is a line at or after the register's definition (the search goes backwards from it). Brace-scoped temporaries
(`{ .reg .b16 hl, hu; ... }`) are resolved within their block. Shuffles print as `shfl.bfly^k(expr)`: the same
expression on the lane `laneid ^ k`.
"""
import argparse
import re
from pathlib import Path
def module_of(kernel):
for p in sorted(Path('~/.cache/d4r-rr-corpus').expanduser().glob('*.ptx')):
if f'.entry {kernel}(' in p.read_text(errors='replace'):
return p
class Tracer:
def __init__(self, lines, start):
self.lines = lines
self.start = start
self.consts = {}
self.cache = {}
def find_def(self, reg, before, block_lo):
"""the statement (line index) that last writes REG before line BEFORE"""
pat = re.compile(r'(^|[\s{,|])' + re.escape(reg) + r'([\s,|}]|$)')
lo = block_lo if not reg.startswith('%') else self.start
for i in range(before - 1, lo - 1, -1):
s = re.sub(r'//.*', '', self.lines[i]).strip().lstrip('{').strip()
s = re.sub(r'^@!?%p\d+\s+', '', s)
if not s or s.startswith(('.', '$', '}')):
continue
parts = s.split(None, 1)
if len(parts) < 2:
continue
rest = parts[1]
dsts = ([x.strip() for x in rest[1:rest.index('}')].split(',')] if rest.startswith('{')
else [x.strip() for x in rest.split(',')[0].split('|')])
if reg in dsts and not parts[0].startswith(('st.', 'sust', 'bra')):
return i, parts[0], rest.rstrip(';').strip(), dsts
return None
def block_start(self, line):
depth = 0
for i in range(line, self.start, -1):
depth += self.lines[i].count('}') - self.lines[i].count('{')
if depth < 0:
return i
return self.start
def expr(self, reg, before, depth, stop, half=None):
reg = reg.strip()
if re.fullmatch(r'-?\d+|0[fFdD][0-9A-Fa-f]+|0x[0-9a-fA-F]+', reg):
return reg
if reg in stop or depth <= 0:
return reg + (f'.{half}' if half is not None else '')
d = self.find_def(reg, before, self.block_start(before - 1))
if d is None:
return reg + '?'
i, op, rest, dsts = d
args = [a.strip() for a in re.split(r',(?![^{]*\})', rest)]
srcs = args[1:] if not rest.startswith('{') else args[1:]
e = lambda r, h=None: self.expr(r, i, depth - 1, stop, h)
if op.startswith('mov.b32') and rest.startswith('{'): # unpack: {lo, hi} <- r
which = 'lo' if dsts.index(reg) == 0 else 'hi'
return e(srcs[0], which)
if op.startswith('mov.b32') and srcs and srcs[0].startswith('{'): # pack r <- {lo, hi}
lo, hi = [x.strip() for x in srcs[0][1:-1].split(',')]
if half == 'lo':
return e(lo)
if half == 'hi':
return e(hi)
return f'pack({e(lo)}, {e(hi)})'
if op.startswith(('mov.', 'cvt.')) and len(srcs) == 1:
inner = e(srcs[0], half)
return inner if op.startswith('mov.') else f'{op.split(".")[0]}.{op.split(".")[-2]}({inner})'
if op.startswith('shfl.sync.bfly'):
return f'shfl^{srcs[1]}({e(srcs[0], half)})'
if op.startswith('ld.'):
return f'load@L{i + 1}[{srcs[0]}]' + (f'.{half}' if half else '')
name = op.split('.')[0]
if op.endswith('x2') and half:
return f'{name}(' + ', '.join(e(s, half) for s in srcs) + ')'
return f'{name}(' + ', '.join(e(s, half) for s in srcs) + ')'
def main():
ap = argparse.ArgumentParser()
ap.add_argument('stage')
ap.add_argument('line', type=int)
ap.add_argument('reg')
ap.add_argument('--depth', type=int, default=30)
ap.add_argument('--stop', default='')
o = ap.parse_args()
kernel = f'cuda_dldn_engine_swin_{o.stage}_kernel'
lines = module_of(kernel).read_text().replace('\r', '').split('\n')
k = next(i for i, l in enumerate(lines) if f'.entry {kernel}(' in l)
t = Tracer(lines, k)
print(t.expr(o.reg, o.line, o.depth, set(filter(None, o.stop.split(',')))))
if __name__ == '__main__':
main()
engine/rr/relations.py
+55
View File
@@ -0,0 +1,55 @@
#!/usr/bin/env python3
"""Which earlier result each phase operand of a dumped layer (engine/rr/phase_dump.py) comes from.
usage: relations.py DUMP.npz [WAVE]"""
import sys
import numpy as np
sys.path.insert(0, __import__('os').path.dirname(__file__))
from rrmath import f16, q8
z = np.load(sys.argv[1]); w = int(sys.argv[2]) if len(sys.argv) > 2 else 0
phases = sorted({int(k.split('_')[0][1:]) for k in z.files})
P = lambda p, n: z[f'p{p}_w{w}_{n}']
def inter(n):
"""the C/D -> A reinterpretation within each 32 block: column 4m+j <- 2m+j, 4m+2+j <- 16+2m+j"""
perm = []
for b in range(n // 32):
for m in range(8):
perm += [32 * b + 2 * m, 32 * b + 2 * m + 1, 32 * b + 16 + 2 * m, 32 * b + 16 + 2 * m + 1]
return perm
def rel(op, M, p):
found = []
if not M.any():
return ['zero']
for q in range(p - 1, 0, -1):
D = P(q, 'D')
if op == 'C':
for name, X in (('D', D),):
if X.shape == M.shape and np.array_equal(X, M):
found.append(f'{name}{q}')
continue
Q = q8(f16(D))
cands = []
if Q.shape[1] % 32 == 0:
cands += [('q8(D%d)[:, inter]' % q, Q[:, inter(Q.shape[1])])]
cands += [('q8(D%d)' % q, Q)]
for name, X in cands:
for tname, Y in ((name, X), (name + '.T', X.T)):
if Y.shape == M.shape and np.array_equal(Y, M):
found.append(tname)
elif Y.shape[1] == M.shape[1] and Y.shape[0] > M.shape[0] and np.array_equal(Y[:M.shape[0]], M):
found.append(tname + f'[:{M.shape[0]}]')
if op == 'B' and Y.shape[0] == M.shape[0] and Y.shape[1] > M.shape[1] and np.array_equal(Y[:, :M.shape[1]], M):
found.append(tname + f'[:, :{M.shape[1]}]')
if found:
break
return found or ['?']
for p in phases:
A, B, C = P(p, 'A'), P(p, 'B'), P(p, 'C')
print(f'phase {p:2d} {A.shape[0]}x{B.shape[1]}x{A.shape[1]}: A <- {rel("A", A, p)} B <- {rel("B", B, p)} '
f'C <- {rel("C", C, p)}')
engine/rr/rr2_m.comp
+578
View File
@@ -0,0 +1,578 @@
#version 460
// Native stage of the DLSS 4.5 Ray Reconstruction network (preset F, the rr2_* kernels of nvngx_dlssd.dll 310.9):
// one shader for enc1 ... enc5 and both kernels of dec4 ... dec1 and dec0's first. The arithmetic is
// engine/rr/rr2_ref.py's (which documents the block); the shape and the weight offsets come from the include
// generated by engine/rr/rr2_pack.py:
// KIND 0 encoder (tokens from the arena), 1 enc1 (first merges encoder 0's tokens 2x2), 2 a decoder's first kernel
// (expands the CL channels one level down and adds the skip), 3 a decoder's second kernel (with the guide blend),
// 4 encoder 0 after its per-pixel feature gather: included (RR2_EMBED) by the shader engine/rr/gen_rr2_enc0.py
// generates around ptx2rr.py's translation of the gather; it embeds the 64 feature slots of XF and leaves
// the guide head in HD for the translated tail; 5 dec0's second kernel, included likewise by the shader
// generated around the translation of its output tail: it leaves r and the MLP sum in RSH / MSH
// C channels (MT = C / 16 tiles), NH heads of 32 channels, NB MLP branches, CM merged channels for the next
// encoder (0: none), GUIDE 1 when the stage gates its guide image (KIND 0 and 1)
// Like engine/rr/swin_m.comp, every product runs transposed (weights x activations), so results stay in registers as
// [channel tile][token] accumulators: one workgroup per 8x8-token window, four waves, wave w owning the token slots
// 16w ... 16w + 15; token and channel slot j holds the natural index inter(j) (NVIDIA's k order, which is also the
// arena's byte order within a 32-channel plane).
#extension GL_KHR_cooperative_matrix : require
#extension GL_KHR_memory_scope_semantics : require
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_8bit_storage : require
#extension GL_EXT_float_e4m3 : require
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_EXT_control_flow_attributes : require
#extension GL_GOOGLE_include_directive : require
#ifndef RR2_EMBED
layout(local_size_x = 128) in;
#endif
#include "rr2_stage.h"
#include "rr_exact.glsl"
#define MT (C / 16)
#ifdef RR2_EMBED
#define PRM(i) param.w[i]
// the generated shader's workgroup holds RR2_GROUP (1 to 4) windows along x (wave w runs the translated code of
// window w, then the block runs once per window): RR2_WIN is the one this call works on
#if RR2_GROUP == 1
#define RR2_WIN 0u
#define RR2_WGX gl_WorkGroupID.x
#else
uint rr2_win;
#define RR2_WIN rr2_win
#define RR2_WGX (uint(RR2_GROUP) * gl_WorkGroupID.x + rr2_win)
#endif
#else
#define RR2_GROUP 1
#define RR2_WIN 0u
#define RR2_WGX gl_WorkGroupID.x
layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param;
#define PRM(i) param.p[i]
#endif
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16;
layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u;
layout(set = 0, binding = 1, std430) readonly buffer W8U { uint8_t w[]; } w8u;
layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define HAC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CMJ gl_CooperativeMatrixLayoutColumnMajor
#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM)
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
#define Q8(dst, src) Q8PAIRS(dst, src)
// One 32-deep step onto an accumulator of binary16 values: the reference adds it inside the f32 products and rounds
// RR_FUSE (D4R_RR_FUSE=1): out = in + branch * s and the guide's blend as one fused multiply-add each, which is what
// NVIDIA's compiler makes of those multiply / add pairs on an NVIDIA card; without it the product is rounded first,
// as ZLUDA (the captures) does. The activation table is built to match (enc1_ref.act).
#ifdef RR_FUSE
#define MADD(d, a, b, c) precise float16_t d = fma((a), (b), (c));
#define MADD4(d, a, b, c) precise f16vec4 d = fma((a), (b), (c));
#else
#define MADD(d, a, b, c) precise float16_t d##_p = (a) * (b); precise float16_t d = (c) + d##_p;
#define MADD4(d, a, b, c) precise f16vec4 d##_p = (a) * (b); precise f16vec4 d = d##_p + (c);
#endif
// RR_NATIVE (D4R_RR_RELAXED=3): the residual steps and the MLP activation in f32, straight into FP8 - no binary16 in
// between and no activation table. Not the reference's roundings (about 40 dB from it, like RR_RELAXED).
// (The biases stay binary16: f32 copies were tried for RR_NATIVE and were slower - half the conversions, but one more
// load per accumulator, and a load costs about three arithmetic instructions here.)
#define BV(X, i) RV(w16.w[W_##X / 2u + (i)])
#define BVS(X, stride, n, i) RV(w16.w[(W_##X + (stride) * (n)) / 2u + (i)])
#ifdef RR_NATIVE
#define RVT float
#define E4V(w, b) float(uintBitsToFloate4m3EXT(uint8_t(((w) >> (8u * (b))) & 0xffu)))
#define RESID(d, a, s, c) const float d = fma(float(a), float(s), float(c));
#define Q8R(x) q8s(x)
#define ACT8(dst, h) [[unroll]] for (int i_ = 0; i_ < 8; i_ += 2) { const vec2 x_ = vec2(float((h)[i_]), float((h)[i_ + 1])); const vec2 y_ = clamp(x_, vec2(-2.0), vec2(2.0)); const vec2 g_ = fma(y_, fma(-abs(y_), vec2(0.0810546875), vec2(0.412109375)), vec2(0.5)); fe4m3vec2 t_; saturatedConvertEXT(t_, x_ * g_); (dst)[i_] = t_.x; (dst)[i_ + 1] = t_.y; }
#else
#define RVT float16_t
#define E4V(w, b) e4((w), (b))
#define RESID(d, a, s, c) MADD(d, float16_t(a), s, c)
#define Q8R(x) q8h(x)
#define ACT8(dst, h) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (dst)[i_] = ACTQ(float16_t((h)[i_]));
#endif
// once (swin_m.comp explains RR_MIX, the one-rounding variant that is not bit-exact by construction)
#ifdef RR_MIX
#define RAC HAC
#define RV(x) (x)
#define K32(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(0.0))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_] + float((acc)[i_])); }
#define K32H(acc, a0, b0, a1, b1) K32(acc, a0, b0, a1, b1)
#else
#define RAC ACC
#define RV(x) float(x)
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) }
#define K32H(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(acc))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_]); }
#endif
const uint HEAD = W_WQK1 - W_WQK0, BRANCH = W_WD1 - W_WD0;
#define BRANCHB BRANCH
shared uint XT[16 * C]; // the window's tokens: [token slot][C FP8 channels in slot order]; later the merge rows
shared uint KVS[2 * 512]; // the current head's QK ([token slot][32]), then V^T ([32][token slot])
#if KIND == 1
shared uint GT[64 * 32]; // the merge input: [token slot][4 sub-positions x 32 FP8 channels]
#endif
#if KIND == 2
// Shared memory is what limits how many waves this GPU keeps running: above about 4 KiB per wave (16 KiB for this
// workgroup) the rate falls in proportion (measured: build/rr/f/nv/cost/occ2.py). So the 4 x 4 tokens one level down
// ([row-major][CL FP8 channels]) borrow KVS, which the attention only writes later, and the block's input before its
// FP8 conversion stays in each lane's registers (xr), not in a [token slot][channel slot] array.
#if 4 * CL > 2 * 512
#error rr2_m.comp: LOW no longer fits in KVS
#endif
#define LOW KVS
shared uint EX[16 * C]; // their expansion: [low token][sub-position][C FP8 channels in slot order]
#endif
#if KIND == 4
#ifdef RR2_XF_ARENA
// the features come from the gather pass (engine/rr/gen_rr2_enc0.py --split): the window's record in the arena, as
// the index of its first 32-bit word, set by the stage's main
uint rr2_xf;
#define XF_SRC arena.w
#define XF_AT rr2_xf
#else
shared uint XF[RR2_GROUP * 64 * 16]; // per window, the features from the gather: [token slot][64 FP8 feature slots]
#define XF_SRC XF
#define XF_AT (RR2_WIN * 1024u)
#endif
shared float16_t HD[RR2_GROUP * 64 * 4]; // per window, the guide head per token (raster order), for the translated tail
const uint EBRANCH = W_EWD1 - W_EWD0;
#define EBRANCHB EBRANCH
#endif
#if KIND == 5
shared float16_t RSH[RR2_GROUP * 64 * 32], MSH[RR2_GROUP * 64 * 32]; // per window, r and the MLP sum per token (raster order) and natural channel
#endif
#if GUIDE || KIND == 3
shared float16_t OUTG[64 * 4]; // the guide output per window token (raster order)
#endif
uint pack4(floate4m3_t a, floate4m3_t b, floate4m3_t c, floate4m3_t d)
{
return code(a) | (code(b) << 8u) | (code(c) << 16u) | (code(d) << 24u);
}
float16_t e4(uint w, uint b) { return float16_t(uintBitsToFloate4m3EXT(uint8_t((w >> (8u * b)) & 0xffu))); }
uvec2 mergeXY(uint r) { return r < 8u ? uvec2(r >> 1u, r & 1u) : uvec2((r - 8u) >> 1u, 2u + (r & 1u)); }
// four binary16 values of a guide image at token (x, y) of a level w wide
f16vec4 guideAt(uint off, int x, int y, int w)
{
const uint a = (off + uint(y * w + x) * 8u) >> 2u;
const uint u0 = arena.w[a], u1 = arena.w[a + 1u];
return f16vec4(hb(u0 & 0xffffu), hb(u0 >> 16u), hb(u1 & 0xffffu), hb(u1 >> 16u));
}
#ifdef RR2_EMBED
void rr2_block()
#else
void main()
#endif
{
const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID;
const uint hl = lane >> 4u, col = lane & 15u;
RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u);
const int ex = int(PRM(0)), ey = int(PRM(1)), hx = ex / 2, hy = ey / 2;
const uint inOff = PRM(2), skOff = PRM(4), outOff = PRM(8), mrgOff = PRM(10), shift = PRM(20);
const uint gInOff = PRM(12), gEncOff = PRM(14), gOutOff = PRM(16), gHalfOff = PRM(18);
const float16_t s0 = hb(PRM(23) & 0xffffu), s1 = hb(PRM(23) >> 16u), s2 = hb(PRM(24) & 0xffffu), s3 = hb(PRM(24) >> 16u);
const int x0 = 8 * int(RR2_WGX) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u);
const uint plane = uint(ex * ey) * 32u, lplane = uint(hx * hy) * 32u;
// this lane's own token (slot 16 wv + col)
const uint t = inter(16u * wv + col);
const int tx = x0 + int(t & 7u), ty = y0 + int(t >> 3u);
const bool inside = tx >= 0 && ty >= 0 && tx < ex && ty < ey;
// ---------------------------------------------------------------- the block's input x: one token per thread
#if KIND == 0 || KIND == 3 || KIND == 5
if (tid < 64u) {
const uint k = inter(tid);
const int xx = mirror(x0 + int(k & 7u), ex), yy = mirror(y0 + int(k >> 3u), ey);
const uint tb = (inOff + uint(yy * ex + xx) * 32u) / 4u;
[[unroll]] for (uint wd = 0u; wd < C / 4; ++wd) XT[tid * (C / 4u) + wd] = arena.w[tb + (wd >> 3u) * (plane >> 2u) + (wd & 7u)];
}
barrier();
#elif KIND == 4
// encoder 0: x = bu + sum over two branches of Wu . act(Wd . features + bd)
{
B8 fb[4];
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) coopMatLoad(fb[kt], XF_SRC, XF_AT + wv * 256u + 4u * kt, 16u, CMJ);
RAC X[MT];
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) X[m][i] = BV(EBU, 16u * m + 8u * hl + uint(i));
[[unroll]] for (uint br = 0u; br < 2u; ++br) {
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
RAC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = BVS(EBD0, EBRANCHB, br, 16u * c + 8u * hl + uint(i));
[[unroll]] for (uint kt = 0u; kt < 4u; kt += 2u) {
A8 a0, a1;
LDA(a0, W_EWD0 + EBRANCH * br, 64u, c, kt); LDA(a1, W_EWD0 + EBRANCH * br, 64u, c, kt + 1u);
K32(h, a0, fb[kt], a1, fb[kt + 1u])
}
ACT8(Hb[c], h)
}
[[unroll]] for (uint m = 0u; m < MT; ++m) {
A8 a0, a1;
LDA(a0, W_EWU0 + EBRANCH * br, 32u, m, 0u); LDA(a1, W_EWU0 + EBRANCH * br, 32u, m, 1u);
K32(X[m], a0, Hb[0], a1, Hb[1])
}
}
[[unroll]] for (uint m = 0u; m < MT; ++m) {
B8 q;
Q8(q, X[m])
coopMatStore(q, XT, wv * 4u * C + 4u * m, C / 4u, CMJ);
}
}
barrier();
#elif KIND == 1
// enc1: the 2x2 group of encoder 0's 32-channel tokens, then x = W1 . group + b1
if (tid < 64u) {
const uint k = inter(tid);
const int xx = mirror(x0 + int(k & 7u), ex), yy = mirror(y0 + int(k >> 3u), ey);
[[unroll]] for (uint q = 0u; q < 4u; ++q) {
const uint tb = (inOff + uint((2 * yy + int(q >> 1u)) * (2 * ex) + 2 * xx + int(q & 1u)) * 32u) / 4u;
[[unroll]] for (uint wd = 0u; wd < 8u; ++wd) GT[tid * 32u + q * 8u + wd] = arena.w[tb + wd];
}
}
barrier();
{
B8 gb[8];
[[unroll]] for (uint kt = 0u; kt < 8u; ++kt) coopMatLoad(gb[kt], GT, wv * 512u + 4u * kt, 32u, CMJ);
[[unroll]] for (uint m = 0u; m < MT; ++m) {
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = BV(B1, 16u * m + 8u * hl + uint(i));
[[unroll]] for (uint kt = 0u; kt < 8u; kt += 2u) {
A8 w0, w1;
LDA(w0, W_W1, 128u, m, kt); LDA(w1, W_W1, 128u, m, kt + 1u);
K32(a, w0, gb[kt], w1, gb[kt + 1u])
}
B8 q;
Q8(q, a)
coopMatStore(q, XT, wv * 4u * C + 4u * m, C / 4u, CMJ);
}
}
barrier();
#else
// a decoder's first kernel: wave w expands the 16 low tokens to sub-position w, then x = dq(e) + dq(skip) * s3
float16_t xr[MT][8]; // x of this lane's token (slot 16 wv + col), channel slots 16 m + 8 hl + i
{
const int X0 = x0 >> 1, Y0 = y0 >> 1;
for (uint i = tid; i < 4u * CL; i += 128u) {
const uint r = i / (CL / 4u), wd = i % (CL / 4u);
const int X = mirror(X0 + int(r & 3u), hx), Y = mirror(Y0 + int(r >> 2u), hy);
LOW[i] = arena.w[(inOff + (wd >> 3u) * lplane + uint(Y * hx + X) * 32u) / 4u + (wd & 7u)];
}
barrier();
B8 lb[CL / 16];
[[unroll]] for (uint kt = 0u; kt < CL / 16; ++kt) coopMatLoad(lb[kt], LOW, 4u * kt, CL / 4u, CMJ);
for (uint mm = 0u; mm < MT; ++mm) {
const uint m = MT * wv + mm;
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = BV(BE, 16u * m + 8u * hl + uint(i));
[[unroll]] for (uint kt = 0u; kt < CL / 16; kt += 2u) {
A8 w0, w1;
LDA(w0, W_WE, CL, m, kt); LDA(w1, W_WE, CL, m, kt + 1u);
K32(a, w0, lb[kt], w1, lb[kt + 1u])
}
B8 q;
Q8(q, a)
coopMatStore(q, EX, 4u * m, C, CMJ);
}
barrier();
if (tid < 64u) {
const uint k = inter(tid);
const int kx = x0 + int(k & 7u), ky = y0 + int(k >> 3u);
const uint tb = (skOff + uint(mirror(ky, ey) * ex + mirror(kx, ex)) * 32u) / 4u;
const uint r = uint(((ky >> 1) - Y0) * 4 + ((kx >> 1) - X0)), s = uint(2 * (ky & 1) + (kx & 1));
[[unroll]] for (uint wd = 0u; wd < C / 4; ++wd) {
const uint sk = arena.w[tb + (wd >> 3u) * (plane >> 2u) + (wd & 7u)];
const uint e = EX[r * C + s * (C / 4u) + wd];
floate4m3_t c[4];
[[unroll]] for (uint b = 0u; b < 4u; ++b) {
MADD(x, e4(sk, b), s3, e4(e, b))
c[b] = q8h(x);
}
XT[tid * (C / 4u) + wd] = pack4(c[0], c[1], c[2], c[3]);
}
}
{ // the same x for the residual step: the bytes of the words 4 m + 2 hl and 4 m + 2 hl + 1 of this lane's token
const int kx = x0 + int(t & 7u), ky = y0 + int(t >> 3u);
const uint tb = (skOff + uint(mirror(ky, ey) * ex + mirror(kx, ex)) * 32u) / 4u;
const uint r = uint(((ky >> 1) - Y0) * 4 + ((kx >> 1) - X0)), s = uint(2 * (ky & 1) + (kx & 1));
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (uint h2 = 0u; h2 < 2u; ++h2) {
const uint wd = 4u * m + 2u * hl + h2;
const uint sk = arena.w[tb + (wd >> 3u) * (plane >> 2u) + (wd & 7u)];
const uint e = EX[r * C + s * (C / 4u) + wd];
[[unroll]] for (uint b = 0u; b < 4u; ++b) {
MADD(x, e4(sk, b), s3, e4(e, b))
xr[m][4u * h2 + b] = x;
}
}
}
barrier();
}
#endif
// ---------------------------------------------------------------- attention: the heads' out-projections sum up
RAC AT[MT];
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) AT[m][i] = BV(BO, 16u * m + 8u * hl + uint(i));
[[dont_unroll]] for (uint hd = 0u; hd < NH; ++hd) {
barrier(); // every wave is done with the previous head's QK / V
for (uint j = 0u; j < 2u; ++j) { // 8 jobs (token tile n, channel tile c)
const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u;
RAC a = RAC(RV(float16_t(0.0))), v = RAC(RV(float16_t(0.0)));
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) {
B8 xb0, xb1, wv0, wv1;
A8 xa0, xa1, wk0, wk1;
coopMatLoad(xb0, XT, n * 4u * C + 4u * kt, C / 4u, CMJ);
coopMatLoad(xb1, XT, n * 4u * C + 4u * kt + 4u, C / 4u, CMJ);
coopMatLoad(xa0, XT, n * 4u * C + 4u * kt, C / 4u, RM);
coopMatLoad(xa1, XT, n * 4u * C + 4u * kt + 4u, C / 4u, RM);
LDA(wk0, W_WQK0 + HEAD * hd, C, c, kt); LDA(wk1, W_WQK0 + HEAD * hd, C, c, kt + 1u);
LDB(wv0, W_WV0 + HEAD * hd, 32u, kt, c); LDB(wv1, W_WV0 + HEAD * hd, 32u, kt + 1u, c);
K32(a, wk0, xb0, wk1, xb1)
K32(v, xa0, wv0, xa1, wv1)
}
A8 ka, va;
Q8(ka, a)
Q8(va, v)
coopMatStore(ka, KVS, n * 128u + 4u * c, 8u, RM);
coopMatStore(va, KVS, 512u + c * 256u + 4u * n, 16u, RM);
}
barrier();
// scores^T (keys x this wave's queries) = QK . QK^T + bias, the cubic softmax down each column
B8 Qb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) coopMatLoad(Qb[c], KVS, wv * 128u + 4u * c, 8u, CMJ);
B8 Pb[4];
{
float16_t e[4][8];
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) {
A8 k0, k1;
coopMatLoad(k0, KVS, kt * 128u, 8u, RM);
coopMatLoad(k1, KVS, kt * 128u + 4u, 8u, RM);
HAC c;
coopMatLoad(c, w16.w, (W_BIAS0 + HEAD * hd) / 2u + kt * 1024u + 16u * wv, 64u, RM);
K32H(c, k0, Qb[0], k1, Qb[1])
[[unroll]] for (int i = 0; i < 8; ++i) e[kt][i] = hb(uint(w16u.w[W_CUBIC / 2u + bh(c[i])]));
}
// key slot 16kt + 8hl + i is NVIDIA's column 8j + 2tt + half with j = 4(kt >> 1) + (kt & 1) + 2((i >> 1) & 1),
// tt = 2hl + (i >> 2), half = i & 1 (rr2_ref.softmax): pairs (e[2q], e[2q + 1]) are the two halves
f16vec2 ep[4][4];
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt)
[[unroll]] for (uint q = 0u; q < 4u; ++q) ep[kt][q] = f16vec2(e[kt][2u * q], e[kt][2u * q + 1u]);
precise f16vec2 st[2];
[[unroll]] for (uint tq = 0u; tq < 2u; ++tq) {
const uint q0 = 2u * tq, q1 = q0 + 1u;
precise f16vec2 ev = (ep[0][q0] + ep[0][q1]) + (ep[2][q0] + ep[2][q1]);
precise f16vec2 od = (ep[1][q0] + ep[1][q1]) + (ep[3][q0] + ep[3][q1]);
st[tq] = ev + od;
}
precise f16vec2 loc = st[0] + st[1];
const uint other = subgroupShuffleXor(packFloat2x16(loc), 16u);
// hl = 0 holds tt 0, 1 and hl = 1 tt 2, 3: the sum is (tt0 + tt1) + (tt2 + tt3) on both halves
precise f16vec2 o = hl == 0u ? loc + unpackFloat2x16(other) : unpackFloat2x16(other) + loc;
precise float16_t rs = o.x + o.y;
const float16_t r = float16_t(1.0 / float(rs));
const f16vec2 r2 = f16vec2(r, r);
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt)
[[unroll]] for (uint q = 0u; q < 4u; ++q) {
precise f16vec2 pv = ep[kt][q] * r2;
{ const fe4m3vec2 t_ = fe4m3vec2(pv); Pb[kt][2u * q] = t_.x; Pb[kt][2u * q + 1u] = t_.y; }
}
}
// O^T = V^T P^T, then the out-projection
B8 Ob[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
RAC o = RAC(RV(float16_t(0.0)));
[[unroll]] for (uint ks = 0u; ks < 2u; ++ks) {
A8 v0, v1;
coopMatLoad(v0, KVS, 512u + c * 256u + 8u * ks, 16u, RM);
coopMatLoad(v1, KVS, 512u + c * 256u + 8u * ks + 4u, 16u, RM);
K32(o, v0, Pb[2u * ks], v1, Pb[2u * ks + 1u])
}
Q8(Ob[c], o)
}
[[unroll]] for (uint m = 0u; m < MT; ++m) {
A8 a0, a1;
LDA(a0, W_WO0 + HEAD * hd, 32u, m, 0u); LDA(a1, W_WO0 + HEAD * hd, 32u, m, 1u);
K32(AT[m], a0, Ob[0], a1, Ob[1])
}
}
// ---------------------------------------------------------------- r = x + attention * s1; the MLP's input q8(r)
RVT rv[MT][8];
B8 Nb[MT];
[[unroll]] for (uint m = 0u; m < MT; ++m) {
#if KIND != 2
const uint xa = (16u * wv + col) * (C / 4u) + 4u * m + 2u * hl;
const uvec2 xw = uvec2(XT[xa], XT[xa + 1u]);
#endif
[[unroll]] for (int i = 0; i < 8; ++i) {
#if KIND == 2
const float16_t xb = xr[m][i];
#else
const RVT xb = E4V(xw[i >> 2], uint(i) & 3u);
#endif
RESID(r, AT[m][i], s1, xb)
rv[m][i] = r;
Nb[m][i] = Q8R(r);
}
}
RAC ML[MT];
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) ML[m][i] = BV(BU, 16u * m + 8u * hl + uint(i));
[[dont_unroll]] for (uint br = 0u; br < NB; ++br) {
const uint w1 = W_WD0 + BRANCH * br, b1 = W_BD0 + BRANCH * br, w2 = W_WU0 + BRANCH * br;
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
RAC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = BVS(BD0, BRANCHB, br, 16u * c + 8u * hl + uint(i));
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) {
A8 a0, a1;
LDA(a0, w1, C, c, kt); LDA(a1, w1, C, c, kt + 1u);
K32(h, a0, Nb[kt], a1, Nb[kt + 1u])
}
ACT8(Hb[c], h)
}
[[unroll]] for (uint m = 0u; m < MT; ++m) {
A8 a0, a1;
LDA(a0, w2, 32u, m, 0u); LDA(a1, w2, 32u, m, 1u);
K32(ML[m], a0, Hb[0], a1, Hb[1])
}
}
#if KIND == 5
// dec0's second kernel: the translated tail forms y = r + MLP * s2 itself and computes the outputs from it
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) {
const uint at = RR2_WIN * 2048u + t * 32u + inter(16u * m + 8u * hl + uint(i));
RSH[at] = float16_t(rv[m][i]);
MSH[at] = float16_t(ML[m][i]);
}
#else
// ---------------------------------------------------------------- y = r + MLP * s2: the stage's tokens
B8 Tb[MT];
{
const uint o = (outOff + uint(ty * ex + tx) * 32u) >> 2u;
#if CM
const uint ty2 = t >> 3u, tx2 = t & 7u;
const uint mr = (ty2 >> 1u) < 2u ? 2u * (tx2 >> 1u) + (ty2 >> 1u) : 8u + 2u * (tx2 >> 1u) + (ty2 >> 1u) - 2u;
const uint blk = 2u * (ty2 & 1u) + (tx2 & 1u);
barrier(); // every wave is done with XT (the merge rows reuse it)
#endif
[[unroll]] for (uint m = 0u; m < MT; ++m) {
[[unroll]] for (int i = 0; i < 8; ++i) {
RESID(y, ML[m][i], s2, rv[m][i])
Tb[m][i] = Q8R(y);
}
const uint w0 = pack4(Tb[m][0], Tb[m][1], Tb[m][2], Tb[m][3]), w1 = pack4(Tb[m][4], Tb[m][5], Tb[m][6], Tb[m][7]);
const uint at = 4u * (m & 1u) + 2u * hl;
if (inside) { arena.w[o + (m >> 1u) * (plane >> 2u) + at] = w0; arena.w[o + (m >> 1u) * (plane >> 2u) + at + 1u] = w1; }
#if CM
const uint g = mr * C + blk * (C / 4u) + 4u * m + 2u * hl;
XT[g] = w0; XT[g + 1u] = w1;
#endif
}
}
#endif
#if GUIDE || KIND == 3 || KIND == 4
{ // the guide head (rows 0-3 on lanes 0-15)
RAC hv;
[[unroll]] for (int i = 0; i < 8; ++i) hv[i] = BV(BH, 8u * hl + uint(i));
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) {
A8 a0, a1;
LDA(a0, W_WH, C, 0u, kt); LDA(a1, W_WH, C, 0u, kt + 1u);
K32(hv, a0, Tb[kt], a1, Tb[kt + 1u])
}
#if KIND == 4
if (hl == 0u) [[unroll]] for (uint c = 0u; c < 4u; ++c) HD[RR2_WIN * 256u + t * 4u + c] = float16_t(hv[c]);
#else
if (hl == 0u) {
f16vec4 o4 = f16vec4(0.0);
if (inside) {
precise f16vec3 gate = f16vec3(sig(float16_t(hv[0])), sig(float16_t(hv[1])), sig(float16_t(hv[2]))) * s0;
#if KIND == 3
// the guide one level down, bilinear (half-pixel centres, edge clamped), every product and sum rounded
const int bx = tx == 0 ? 0 : (tx - 1) >> 1, by = ty == 0 ? 0 : (ty - 1) >> 1;
const float16_t wx = tx == 0 ? float16_t(0.0) : ((tx & 1) != 0 ? float16_t(0.25) : float16_t(0.75));
const float16_t wy = ty == 0 ? float16_t(0.0) : ((ty & 1) != 0 ? float16_t(0.25) : float16_t(0.75));
const int x1 = min(bx + 1, hx - 1), y1 = min(by + 1, hy - 1);
precise float16_t ux = float16_t(1.0) - wx, uy = float16_t(1.0) - wy;
precise f16vec4 p00 = guideAt(gInOff, bx, by, hx) * ux, p10 = guideAt(gInOff, x1, by, hx) * wx;
precise f16vec4 p01 = guideAt(gInOff, bx, y1, hx) * ux, p11 = guideAt(gInOff, x1, y1, hx) * wx;
precise f16vec4 a10 = p10 * uy;
MADD4(u0, p00, f16vec4(uy), a10)
MADD4(u1, p01, f16vec4(wy), u0)
MADD4(up, p11, f16vec4(wy), u1)
const float16_t al = sig(float16_t(hv[3]));
precise f16vec4 gt = up * f16vec4(gate, float16_t(0.0));
precise float16_t ia = float16_t(1.0) - al;
precise f16vec4 pg = gt * al;
MADD4(mx, guideAt(gEncOff, tx, ty, ex), f16vec4(ia), pg)
o4 = min(mx, f16vec4(65504.0));
#else
precise f16vec4 gg = guideAt(gInOff, tx, ty, ex) * f16vec4(gate, float16_t(0.0));
o4 = min(gg, f16vec4(65504.0));
#endif
const uint a = (gOutOff + uint(ty * ex + tx) * 8u) >> 2u;
arena.w[a] = bh(o4.x) | (bh(o4.y) << 16u);
arena.w[a + 1u] = bh(o4.z) | (bh(o4.w) << 16u);
}
[[unroll]] for (uint c = 0u; c < 4u; ++c) OUTG[t * 4u + c] = o4[c];
}
#endif
}
#endif
#if CM
barrier();
{ // the 2x2 merge for the next encoder: merged^T (CM x 16 rows), channel tiles m = wv, wv + 4, ...
const uvec2 mxy = mergeXY(col);
const int X = x0 / 2 + int(mxy.x), Y = y0 / 2 + int(mxy.y);
const uint o = (mrgOff + uint(Y * hx + X) * 32u) >> 2u;
for (uint m = wv; m < CM / 16u; m += 4u) {
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = BV(BM, 16u * m + 8u * hl + uint(i));
for (uint ks = 0u; ks < C / 8u; ++ks) {
B8 g0, g1; A8 w0, w1;
coopMatLoad(g0, XT, 8u * ks, C, CMJ);
coopMatLoad(g1, XT, 8u * ks + 4u, C, CMJ);
LDA(w0, W_WM, 4u * C, m, 2u * ks); LDA(w1, W_WM, 4u * C, m, 2u * ks + 1u);
K32(a, w0, g0, w1, g1)
}
B8 q;
Q8(q, a)
if (X >= 0 && Y >= 0 && X < hx && Y < hy) {
const uint at = o + (m >> 1u) * (lplane >> 2u) + 4u * (m & 1u) + 2u * hl;
arena.w[at] = pack4(q[0], q[1], q[2], q[3]);
arena.w[at + 1u] = pack4(q[4], q[5], q[6], q[7]);
}
}
#if GUIDE
if (wv == 3u && lane < 16u) { // the guide for the next encoder: 2x2 mean, rightwards first, then downwards
const uint mx = lane & 3u, my = lane >> 2u;
const int x = x0 + 2 * int(mx), y = y0 + 2 * int(my);
if (x >= 0 && y >= 0 && x < ex && y < ey) {
const uint t00 = (2u * my) * 8u + 2u * mx;
f16vec4 q4[4];
[[unroll]] for (uint d = 0u; d < 4u; ++d) {
const uint tt = t00 + (d & 1u) + 8u * (d >> 1u);
precise f16vec4 qv = f16vec4(OUTG[tt * 4u], OUTG[tt * 4u + 1u], OUTG[tt * 4u + 2u], OUTG[tt * 4u + 3u]) * float16_t(0.25);
q4[d] = qv;
}
precise f16vec4 h0 = q4[0] + q4[1], h1 = q4[2] + q4[3];
precise f16vec4 hm = h0 + h1;
const uint a = (gHalfOff + uint((y / 2) * hx + x / 2) * 8u) >> 2u;
arena.w[a] = bh(hm.x) | (bh(hm.y) << 16u);
arena.w[a + 1u] = bh(hm.z) | (bh(hm.w) << 16u);
}
}
#endif
}
#endif
}
engine/rr/rr2_pack.py
+90
View File
@@ -0,0 +1,90 @@
#!/usr/bin/env python3
"""Pack a preset F stage's weights, read from the user's nvngx_dlssd.dll (engine/rr/rr2_ref.dll_weights), for
engine/rr/rr2_m.comp.
usage: rr2_pack.py STAGE NVNGX_DLSSD_DLL OUT.bin OUT.h
swin_pack.py's layout: transposed products as 16 x 16 tiles, every tensor's outputs in slot order (slot j holds the
natural index inter(j)); a merge's or the expand's inputs are sub-positions of C channels each. OUT.h also defines
the stage's shape for the shader (KIND, C, NH, NB, CM, CL, GUIDE).
"""
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from enc1_ref import inter # noqa: E402
from enc1_pack import a_tiles, b_tiles, cubic_table, act_table # noqa: E402
from rr2_ref import STAGES, EXPAND, dll_weights # noqa: E402
def main():
stage, dll, out_bin, out_h = sys.argv[1:5]
w = dll_weights(open(dll, 'rb').read(), stage)
if stage in EXPAND:
C, CL, NH, NB = EXPAND[stage]
kind, CM, guide = 2, 0, 0
else:
C, c_up, CM, NH, NB = STAGES[stage]
kind, CL = (4 if stage == 'enc0' else 5 if stage == 'dec0_1' else 3 if stage.startswith('dec') else 1 if c_up else 0), 0
guide = 0 if kind >= 3 else 1
I32, I64, IC = inter(32), inter(64), inter(C)
e = {}
if kind == 1:
e['w1'], e['b1'] = ('A', w['w1'][:, IC]), ('H', w['b1'][IC])
if kind == 4: # encoder 0's embedding branches
for b in range(2):
e[f'ewd{b}'] = ('A', w['ewd'][b][:, I32])
e[f'ebd{b}'] = ('H', w['ebd'][b][I32])
e[f'ewu{b}'] = ('A', w['ewu'][b][:, IC])
e['ebu'] = ('H', w['ebu'][IC])
if kind == 2:
I4C = [C * (j // C) + IC[j % C] for j in range(4 * C)]
e['we'], e['be'] = ('A', w['we'][:, I4C]), ('H', w['be'][I4C])
for h in range(NH):
e[f'wqk{h}'] = ('A', w['wqk'][h][:, I32])
e[f'wv{h}'] = ('B', w['wv'][h][:, I32])
e[f'wo{h}'] = ('A', w['wo'][h][:, IC])
e[f'bias{h}'] = ('H', w['bias'][h][np.ix_(I64, I64)]) # [key slot][query slot]; symmetric
e['bo'] = ('H', w['bo'][IC])
for b in range(NB):
e[f'wd{b}'] = ('A', w['wd'][b][:, I32])
e[f'bd{b}'] = ('H', w['bd'][b][I32])
e[f'wu{b}'] = ('A', w['wu'][b][:, IC])
e['bu'] = ('H', w['bu'][IC])
if CM:
e['wm'], e['bm'] = ('A', w['wm'][:, inter(CM)]), ('H', w['bm'][inter(CM)])
if 'wh' in w:
e['wh'], e['bh'] = ('A', w['wh']), ('H', w['bh'])
if kind == 5: # dec0's output tail: tanh of every binary16 value, correctly rounded, by its bits
with np.errstate(all='ignore'):
x = np.arange(65536, dtype=np.uint32).astype(np.uint16).view(np.float16).astype(np.float64)
e['tanh'] = ('U', np.tanh(x).astype(np.float16).view(np.uint16))
e['cubic'] = ('U', cubic_table())
e['act'] = ('Y', act_table())
blob, defines = bytearray(), []
for name, (kd, a) in e.items():
while len(blob) % 256:
blob.append(0)
off = len(blob)
if kd == 'A':
blob += a_tiles(a).tobytes()
elif kd == 'B':
blob += b_tiles(a).tobytes()
elif kd == 'Y':
blob += np.ascontiguousarray(a, np.uint8).tobytes()
elif kd == 'U':
blob += np.ascontiguousarray(a, np.uint16).tobytes()
else:
blob += np.ascontiguousarray(a, np.float16).tobytes()
defines.append(f'#define W_{name.upper()} {off}u // {kd} {a.shape}')
shape = [f'#define KIND {kind}', f'#define C {C}', f'#define NH {NH}', f'#define NB {NB}', f'#define CM {CM}',
f'#define CL {CL}', f'#define GUIDE {guide}']
open(out_bin, 'wb').write(bytes(blob))
open(out_h, 'w').write(f'// generated by engine/rr/rr2_pack.py for {stage}\n' + '\n'.join(shape + defines) + '\n')
print(f'{stage}: {len(blob)} bytes, {len(e)} entries')
if __name__ == '__main__':
main()
engine/rr/rr2_ref.py
+362
View File
@@ -0,0 +1,362 @@
#!/usr/bin/env python3
"""Exact numpy reference of preset F's encoder stages (rr2_enc1_kernel ... rr2_enc4_kernel of nvngx_dlssd.dll
310.9), one 8x8 window at a time, from weights and real inputs only; checked against a capture (main).
A stage has C channels and works on windows of 64 tokens (origin 8 * block - shift). Arithmetic is binary16 with FP8
(e4m3) operands in the products, exactly as in presets D and E (enc1_ref.py: mm, the cubic softmax, the activation):
merge x = W1 . [the four tokens of the 2x2 group one level up, (0,0) (1,0) (0,1) (1,1)] + b1; window tokens
outside the level are the mirrored tokens
heads per head, all from q8(x): QK = Wqk . x and V = Wv . x (32 channels each, no bias). The scores are
QK . QK^T + bias: symmetric, so NVIDIA computes the blocks (queries 0-31, keys 0-31), (0-31, 32-63) and
(32-63, 32-63) and transposes the second for (32-63, 0-31). Cubic softmax over the 64 keys, P . V, and the
out-projections of all heads accumulate into one sum a (the first one's C is the bias bo)
residual r = dq(q8(x)) + a * s1 (dq: the FP8 code back as binary16; s1: binary16 at argument 94)
MLP NB branches from q8(r): C -> 32 (+ bias) -> act -> C, accumulated into m (first C: the bias bu)
residual y = r + m * s2 (s2: argument 96); the stage's tokens are q8(y)
outputs merged = Wm . [y of the 2x2 group] + bm for the next encoder (q8), and a 16-channel head h = Wh . y + bh
whose first three channels gate the stage's guide image (argument 48, three binary16 channels per token):
out = min(guide * (sigmoid(h) * s0), 65504) (s0: argument 92) at argument 64, and its 2x2 mean
((a + b) + (c + d) of the quarters) at argument 72 for the next encoder
Arena layouts: tokens are FP8 codes in planes of 32 channels ([plane][y][x][32], NVIDIA's k order within a plane);
the guide images are four binary16 per token.
usage: rr2_ref.py CAPTURE STAGE FRAME WEIGHTS.npz [--blocks X,Y ...] (verifies against the capture)
rr2_ref.py CAPTURE STAGE FRAME OUT.npz --extract DUMP.npz (weights from a phase dump, phase_dump.py)
"""
import argparse
import os
import struct
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rrmath import e4m3, q8 # noqa: E402
from enc1_ref import inter, mm, act, add, pairs_tree, mirror, H, F # noqa: E402
from rr_capture import Capture # noqa: E402
# stage: C, channels one level up (merge input), merged channels for the next encoder, heads, MLP branches
STAGES = {'enc0': (32, 0, 0, 2, 4), 'enc1': (64, 32, 96, 2, 8), 'enc2': (96, 0, 128, 4, 12), 'enc3': (128, 0, 160, 4, 16),
'enc4': (160, 0, 160, 8, 20), 'enc5': (160, 0, 0, 8, 20),
'dec0_1': (32, 0, 0, 2, 4), 'dec4_1': (160, 0, 0, 8, 20), 'dec3_1': (128, 0, 0, 4, 16), 'dec2_1': (96, 0, 0, 4, 12), 'dec1_1': (64, 0, 0, 2, 8)}
# the decoders' first kernels (decN_0): C, channels one level down, heads, MLP branches. They expand the 4 x 4 tokens
# one level down to the window, e = We . low + be (4C outputs: sub-position (y & 1) * 2 + (x & 1), then channel), and
# start from x = dq(q8(e)) + dq(skip) * s3 (s3: binary16 at argument 98; skip: the encoder's tokens of this level,
# argument 16); then the heads and the MLP as above, the first residual on x itself (not dq(q8(x))). No merge, no guide.
EXPAND = {'dec4_0': (160, 160, 8, 20), 'dec3_0': (128, 160, 4, 16), 'dec2_0': (96, 128, 4, 12), 'dec1_0': (64, 96, 2, 8),
'dec0_0': (32, 64, 2, 4)}
# the network's sixteen kernels in the DLL's order, with their head counts
ORDER = ['enc0', 'enc1', 'enc2', 'enc3', 'enc4', 'enc5', 'dec4_0', 'dec4_1', 'dec3_0', 'dec3_1', 'dec2_0', 'dec2_1',
'dec1_0', 'dec1_1', 'dec0_0', 'dec0_1']
HEADS = [2, 2, 4, 4, 8, 8, 8, 8, 4, 4, 4, 4, 2, 2, 2, 2]
def find_tables(dll):
"""byte offsets of the sixteen kernels' relative position tables in nvngx_dlssd.dll: runs of 64 x 64 binary16
tables (one per head) with a constant diagonal and at most 225 distinct values, in ORDER"""
a = np.frombuffer(dll[:len(dll) & ~1], np.uint16)
n = len(a) - 4096
m = (a[:n] & 0x7fff) != 0
for k in range(1, 8):
m &= a[:n] == a[65 * k:65 * k + n]
m &= a[1:n + 1] == a[66:66 + n]
groups = []
for i in np.nonzero(m)[0]:
t = a[i:i + 4096].reshape(64, 64)
if not (np.all(t.diagonal() == t[0, 0]) and 50 < len(np.unique(t)) <= 225):
continue
if groups and 2 * int(i) - groups[-1][-1] == 8192:
groups[-1].append(2 * int(i))
else:
groups.append([2 * int(i)])
for start in range(len(groups) - 15):
if [len(g) for g in groups[start:start + 16]] == HEADS:
return {name: g[0] for name, g in zip(ORDER, groups[start:start + 16])}
raise SystemExit('no preset F model found in the DLL (its sixteen relative position table groups)')
def dll_weights(dll, stage):
"""an encoder stage's tensors from the DLL: plain binary16 arrays in natural [out][in] order, one after another.
Before the tables: {s0, 0, 0, 0} and (enc1) the merge of encoder 0's tokens, W1 then b1. After them, per tensor
kind all heads / branches together: Wqk, Wv (heads x 32 x C), Wo (C x heads x 32), bo, Wd (branches x 32 x C), bd,
Wu (C x branches x 32), bu, the merge
for the next encoder Wm (C_M x 4C) and bm, the guide head Wh (3 x C) and bh. A merge's input index is
channel * 4 + sub-position; the kernels' bias is (T + T^T) * sqrt(2)."""
if stage in EXPAND:
c, c_lo, nh, nb = EXPAND[stage]
c_up = c_m = 0
else:
c, c_up, c_m, nh, nb = STAGES[stage]
g = find_tables(dll)[stage]
f = lambda off, n: np.frombuffer(dll[off:off + 2 * n], H)
def B(mat, merge=0):
"""[out][in] binary16 -> the product's B operand (in x out FP8 codes, the inputs in NVIDIA's k order)"""
k = mat.shape[1]
if merge: # channel * 4 + q -> q * channels + channel
mat = mat.reshape(mat.shape[0], merge, 4).transpose(0, 2, 1).reshape(mat.shape[0], k)
return q8(mat).T[inter(k)]
w = {}
if c_up:
o = g - 8 - 2 * (3 * c_up + 4) - 2 * c - 2 * c * 4 * c_up # before encoder 0's guide head (3 x C_UP, bias, 0)
w['w1'] = B(f(o, c * 4 * c_up).reshape(c, 4 * c_up), c_up)
w['b1'] = f(o + 2 * c * 4 * c_up, c)
if stage in EXPAND: # We (C x 4 x C_LO) and be (C x 4) end where the tables start
o = g - 8 * c - 8 * c * c_lo
we = f(o, 4 * c * c_lo).reshape(c, 4, c_lo).transpose(1, 0, 2).reshape(4 * c, c_lo)
w['we'] = q8(we).T[inter(c_lo)]
w['be'] = f(g - 8 * c, 4 * c).reshape(c, 4).T.reshape(-1)
if stage == 'enc0':
# encoder 0's embedding of the 43 per-pixel features (the kernel's 64 feature slots, the rest unused): two MLP
# branches 43 -> 32 -> act -> 32, x = bu + sum: Wd (2 x 32 x 43), bd, Wu (32 x 2 x 32), bu, then 16 bytes
o = g - 16 - 64 - 4096 - 128 - 2 * 2 * 32 * 43
wd = np.zeros((2, 32, 64), H); wd[:, :, :43] = f(o, 2 * 32 * 43).reshape(2, 32, 43); o += 2 * 2 * 32 * 43
w['ewd'] = np.array([B(x) for x in wd])
w['ebd'] = f(o, 64).reshape(2, 32); o += 128
w['ewu'] = np.array([B(x) for x in f(o, 2048).reshape(32, 2, 32).transpose(1, 0, 2)]); o += 4096
w['ebu'] = f(o, 32)
o = g
t = f(o, nh * 4096).reshape(nh, 64, 64).astype(np.float64); o += nh * 8192
w['bias'] = ((t + t.transpose(0, 2, 1)) * np.sqrt(2)).astype(H)
w['wqk'] = np.array([B(x) for x in f(o, nh * 32 * c).reshape(nh, 32, c)]); o += nh * 64 * c
w['wv'] = np.array([B(x) for x in f(o, nh * 32 * c).reshape(nh, 32, c)]); o += nh * 64 * c
w['wo'] = np.array([B(x) for x in f(o, nh * 32 * c).reshape(c, nh, 32).transpose(1, 0, 2)]); o += nh * 64 * c
w['bo'] = f(o, c); o += 2 * c
w['wd'] = np.array([B(x) for x in f(o, nb * 32 * c).reshape(nb, 32, c)]); o += nb * 64 * c
w['bd'] = f(o, nb * 32).reshape(nb, 32); o += nb * 64
w['wu'] = np.array([B(x) for x in f(o, nb * 32 * c).reshape(c, nb, 32).transpose(1, 0, 2)]); o += nb * 64 * c
w['bu'] = f(o, c); o += 2 * c
if c_m:
w['wm'] = B(f(o, c_m * 4 * c).reshape(c_m, 4 * c), c); o += 2 * c_m * 4 * c
w['bm'] = f(o, c_m); o += 2 * c_m
if stage in EXPAND:
return w
if stage == 'enc0': # the merge of its tokens (enc1's W1, b1) comes before its guide head
o += 2 * 64 * 128 + 2 * 64
nhd = 4 if stage.startswith('dec') else 3 # the guide head: three gates; decoders also a blend factor
wh = np.zeros((16, c), H); wh[:nhd] = f(o, nhd * c).reshape(nhd, c); o += 2 * nhd * c
bh = np.zeros(16, H); bh[:nhd] = f(o, nhd)
w['wh'], w['bh'] = B(wh), bh
return w
def dq(x):
"""binary16 values through FP8 and back"""
return e4m3(q8(x)).astype(H)
def A(x):
"""binary16 rows -> the product's A operand: FP8 codes in NVIDIA's k order"""
return q8(x)[:, inter(x.shape[1])]
def softmax(s):
"""the cubic-exponent softmax of presets D and E over the 64 keys of a row"""
def fma16(a, b, c):
return (a.astype(np.float64) * b.astype(np.float64) + c.astype(np.float64)).astype(H)
m = np.clip((s * H(0.017181396484375)).astype(H), H(-0.55615234375), H(0.55615234375))
t = fma16(m, -m, np.full_like(m, H(0.92724609375)))
u = fma16(m, t, np.full_like(m, H(1.375)))
e = ((((u.view(np.uint16).astype(np.uint32)) << 5) & 0x7FE0).astype(np.uint16)).view(H)
nt = e.shape[1] // 8
part = {}
for tt in range(4):
for half in range(2):
ev = pairs_tree([e[:, 8 * j + 2 * tt + half] for j in range(0, nt, 2)])
od = pairs_tree([e[:, 8 * j + 2 * tt + half] for j in range(1, nt, 2)])
part[tt, half] = add(ev, od)
out = [add(add(part[0, h], part[1, h]), add(part[2, h], part[3, h])) for h in range(2)]
r = (F(1) / add(out[0], out[1]).astype(F)).astype(H)
return (e * r[:, None]).astype(H)
def sigmoid(d):
"""1 / (1 + 2^(-d log2 e)) in f32, rounded to binary16 (the GPU's exp2 and division; see rr_exact.glsl sig)"""
f = F(1) / (np.exp2(d.astype(F) * F(-1.4426950216293335)).astype(F) + F(1))
return f.astype(H)
def upsample(low, x, y):
"""the decoders' 2x bilinear fetch of the guide one level down (half-pixel centres, edge clamped), with NVIDIA's
binary16 rounding: every weight product and the three sums in tap order"""
def axis(v, n):
if v == 0:
return 0, H(0)
b = (v - 1) // 2
return b, H(0.25) if v & 1 else H(0.75)
hh, hw = low.shape[:2]
bx, wx = axis(x, hw)
by, wy = axis(y, hh)
x1, y1 = min(bx + 1, hw - 1), min(by + 1, hh - 1)
ux, uy = H(1) - wx, H(1) - wy
m = lambda v, a, b: ((v * a).astype(H) * b).astype(H)
return add(add(add(m(low[by, bx], ux, uy), m(low[by, x1], wx, uy)), m(low[y1, bx], ux, wy)), m(low[y1, x1], wx, wy))
def window(w, scal, xin, x=None):
"""xin: 64 x 4*C_UP FP8 codes (the merge input, already in k order). Returns y (64 x C binary16), the merged
tokens (16 x C_M binary16) and the head (64 x 16 binary16)."""
s0, s1, s2 = scal
base = None
if x is not None: # a decoder's first kernel: the expanded tokens plus the skip; the residual keeps it unquantised
xa, base = A(x), x
elif 'w1' in w:
x = mm(xin, w['w1'], np.tile(w['b1'], (64, 1)))
xa = A(x)
else: # the tokens themselves (already FP8)
x, xa = e4m3(xin)[:, np.argsort(inter(xin.shape[1]))].astype(H), xin
acc = np.tile(w['bo'], (64, 1))
for h in range(len(w['wqk'])):
qk = mm(xa, w['wqk'][h], np.zeros((64, 32), H))
v = mm(xa, w['wv'][h], np.zeros((64, 32), H))
qa = A(qk)
s = mm(qa, qa.T, w['bias'][h])
# NVIDIA computes three 32 x 32 blocks and transposes one; the full product is the same values
p = softmax(s)
o = mm(A(p), A(v.T).T if False else q8(v)[inter(64), :], np.zeros((64, 32), H))
acc = mm(A(o), w['wo'][h], acc)
r = add(dq(x) if base is None else base, (acc * s1).astype(H))
ra = A(r)
m = np.tile(w['bu'], (64, 1))
for j in range(len(w['wd'])):
d = mm(ra, w['wd'][j], np.tile(w['bd'][j], (64, 1)))
m = mm(A(act(d)), w['wu'][j], m)
y = add(r, (m * s2).astype(H))
ya = A(y)
# the 2x2 groups of the window: token order (0,0) (1,0) (0,1) (1,1), 16 groups row-major
g = np.zeros((16, 4 * y.shape[1]), np.uint8)
for n in range(16):
gx, gy = n % 4, n // 4
for q, (dx, dy) in enumerate(((0, 0), (1, 0), (0, 1), (1, 1))):
g[n, q * y.shape[1]:(q + 1) * y.shape[1]] = ya[(2 * gy + dy) * 8 + 2 * gx + dx]
merged = mm(g, w['wm'], np.tile(w['bm'], (16, 1))) if 'wm' in w else None
head = mm(ya, w['wh'], np.tile(w['bh'], (64, 1))) if 'wh' in w else None
return y, merged, head
def extract(z, out):
"""weights of enc1 from a phase dump (phases as phase_dump.py numbers them)"""
P = lambda p, n: z[f'p{p}_w0_{n}']
Ch = lambda p: P(p, 'C').view(H)
w = {'w1': P(1, 'B'), 'b1': Ch(1)[0], 'wqk': [], 'wv': [], 'bias': [], 'wo': [], 'bo': Ch(11)[0]}
for h in range(2):
b = 2 + 10 * h
w['wqk'].append(np.vstack([P(b, 'B'), P(b + 2, 'B')]))
w['wv'].append(np.vstack([P(b + 1, 'B'), P(b + 3, 'B')]))
tl, tr, br = Ch(b + 4), Ch(b + 5), Ch(b + 7)
w['bias'].append(np.block([[tl, tr], [tr.T, br]]))
w['wo'].append(P(b + 9, 'B'))
w['wd'] = [P(22 + 2 * j, 'B') for j in range(8)]
w['bd'] = [Ch(22 + 2 * j)[0] for j in range(8)]
w['wu'] = [P(23 + 2 * j, 'B') for j in range(8)]
w['bu'] = Ch(23)[0]
w['wm'], w['bm'], w['wh'], w['bh'] = P(38, 'B'), Ch(38)[0], P(39, 'B'), Ch(39)[0]
np.savez_compressed(out, **{k: np.array(v) for k, v in w.items()})
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('stage')
ap.add_argument('frame', type=int)
ap.add_argument('weights')
ap.add_argument('--extract')
ap.add_argument('--blocks', nargs='*', default=['5,5'])
o = ap.parse_args()
if o.extract:
return extract(np.load(o.extract), o.weights)
expand = o.stage in EXPAND
if expand:
c, c_lo, heads, nb = EXPAND[o.stage]
c_up = c_m = 0
else:
c, c_up, c_m, heads, nb = STAGES[o.stage]
w = dll_weights(open(o.weights, 'rb').read(), o.stage) if o.weights.lower().endswith('.dll') else dict(np.load(o.weights))
cap = Capture(os.path.expanduser(o.capture))
kernel = f'rr2_{o.stage}_kernel' if o.stage.startswith('enc') else f'rr2_{o.stage[:4]}_kernel_{o.stage[5]}'
l = [l for l in cap.launches if l.kernel == kernel][o.frame - 1]
pre = np.frombuffer(l.pre(l.arena_alloc), np.uint8)
post = np.frombuffer(l.post_buffer(cap, l.arena_alloc), np.uint8)
ptr = {p.offset: p.delta for p in l.pointers if p.alloc == l.arena_alloc}
wd, ht = struct.unpack_from('<2I', l.args, 0)
sx, sy = struct.unpack_from('<2H', l.args, 80)
scal = np.frombuffer(l.args[92:98], H)
s3 = np.frombuffer(l.args[98:100], H)[0]
planes = lambda buf, off, w_, h_, ch: buf[off:off + w_ * h_ * ch].reshape(ch // 32, h_, w_, 32)
up = (planes(pre, ptr[8], wd // 2, ht // 2, c_lo) if expand else planes(pre, ptr[8], 2 * wd, 2 * ht, c_up) if c_up else
planes(pre, ptr[8], wd, ht, c))
skip = planes(pre, ptr[16], wd, ht, c) if expand else None
guided = 48 in ptr and o.stage.startswith('enc')
dec_guide = o.stage.startswith('dec') and 56 in ptr
tok_ref = planes(post, ptr[32], wd, ht, c)
mer_ref = planes(post, ptr[40], wd // 2, ht // 2, c_m) if c_m else None
img = lambda buf, off, w_, h_: np.frombuffer(bytes(buf[off:off + w_ * h_ * 8]) + bytes(8), H)[:w_ * h_ * 4].reshape(h_, w_, 4)
guide, out_ref = (img(pre, ptr[48], wd, ht), img(post, ptr[64], wd, ht)) if guided else (None, None)
half_ref = img(post, ptr[72], wd // 2, ht // 2) if c_m else None
if dec_guide: # the guide one level down (argument 48), the encoder's of this level (56), the output (64)
g_low, g_enc, out_ref = img(pre, ptr[48], wd // 2, ht // 2), img(pre, ptr[56], wd, ht), img(post, ptr[64], wd, ht)
ok = True
for blk in o.blocks:
bx, by = (int(v) for v in blk.split(','))
x0, y0 = 8 * bx - sx, 8 * by - sy
inside = np.array([0 <= x0 + t % 8 < wd and 0 <= y0 + t // 8 < ht for t in range(64)])
xin = np.zeros((64, 4 * c_up if c_up else c), np.uint8)
xe = None
if expand:
lows = np.array([np.concatenate([up[p, mirror(y0 // 2 + n // 4, ht // 2), mirror(x0 // 2 + n % 4, wd // 2)]
for p in range(c_lo // 32)]) for n in range(16)])
e = mm(lows, w['we'], np.tile(w['be'], (16, 1)))
un = np.argsort(inter(c))
xe = np.zeros((64, c), H)
for t in range(64):
tx, ty = t % 8, t // 8
sub = (ty % 2) * 2 + tx % 2
sk = np.concatenate([skip[p, mirror(y0 + ty, ht), mirror(x0 + tx, wd)] for p in range(c // 32)])
xe[t] = add(dq(e[(ty // 2) * 4 + tx // 2, sub * c:(sub + 1) * c]), (e4m3(sk)[un].astype(H) * s3).astype(H))
for t in range(64 if not expand else 0): # tokens outside the level are the mirrored ones (no edge repeat), as in presets D and E
X, Y = mirror(x0 + t % 8, wd), mirror(y0 + t // 8, ht)
if not c_up:
xin[t] = np.concatenate([up[p, Y, X] for p in range(c // 32)])
continue
for q, (dx, dy) in enumerate(((0, 0), (1, 0), (0, 1), (1, 1))):
xin[t, q * c_up:(q + 1) * c_up] = np.concatenate([up[p, 2 * Y + dy, 2 * X + dx] for p in range(c_up // 32)])
y, merged, head = window(w, scal, xin, xe)
ya, ma = A(y), (A(merged) if c_m else None)
bad_t = bad_m = bad_o = bad_h = 0
out = np.zeros((64, 4), H)
for t in np.nonzero(inside)[0]:
X, Y = x0 + t % 8, y0 + t // 8
bad_t += int(np.sum(np.concatenate([tok_ref[p, Y, X] for p in range(c // 32)]) != ya[t]))
if dec_guide:
# out = min(enc * (1 - a) + (upsampled * sigmoid(h[0:3]) * s0) * a, 65504), a = sigmoid(h[3])
sg = sigmoid(head[t, :4])
gate = np.array([(sg[0] * scal[0]).astype(H), (sg[1] * scal[0]).astype(H), (sg[2] * scal[0]).astype(H), H(0)])
gt = (upsample(g_low, X, Y) * gate).astype(H)
gt[3] = H(0)
mix = add((g_enc[Y, X] * (H(1) - sg[3])).astype(H), (gt * sg[3]).astype(H))
out[t] = np.minimum(mix, H(65504))
bad_o += int(np.sum(out[t].view(np.uint16) != out_ref[Y, X].view(np.uint16)))
if not guided:
continue
gate = (sigmoid(head[t, :3]) * scal[0]).astype(H)
out[t, :3] = np.minimum((guide[Y, X, :3] * gate).astype(H), H(65504))
bad_o += int(np.sum(out[t].view(np.uint16) != out_ref[Y, X].view(np.uint16)))
for n in range(16):
X, Y = x0 // 2 + n % 4, y0 // 2 + n // 4
t0 = (2 * (n // 4)) * 8 + 2 * (n % 4)
if not c_m or not (inside[t0] and x0 % 2 == 0 and y0 % 2 == 0):
continue
bad_m += int(np.sum(np.concatenate([mer_ref[p, Y, X] for p in range(c_m // 32)]) != ma[n]))
qv = [(out[t0 + d] * H(0.25)).astype(H) for d in (0, 1, 8, 9)]
bad_h += int(np.sum(add(add(qv[0], qv[1]), add(qv[2], qv[3])).view(np.uint16) != half_ref[Y, X].view(np.uint16)))
print(f'block {blk}: {int(inside.sum())} tokens inside; differing tokens bytes {bad_t}, merged bytes {bad_m}, '
f'guide values {bad_o}, half guide values {bad_h}')
ok &= bad_t + bad_m + bad_o + bad_h == 0
print('IDENTICAL' if ok else 'DIFFERENT')
sys.exit(0 if ok else 1)
if __name__ == '__main__':
main()
engine/rr/rr2_templates.py
+263
View File
@@ -0,0 +1,263 @@
#!/usr/bin/env python3
"""What a preset F model takes from a capture, read from nvngx_dlssd.dll 310.9 instead: the argument-block templates
(<stage>.args), encoder 0's second table (enc0-aux.bin, argument 104) and the stages' plane-arena regions.
usage: rr2_templates.py NVNGX_DLSSD_DLL OUT_DIR write the templates and enc0-aux.bin
rr2_templates.py NVNGX_DLSSD_DLL --check MODEL_DIR compare with a model compile_rr2.py built from a capture
engine/runtime_rr2.inc fills every field of the blocks that depends on the frame, the geometry or the bound inputs.
Across the preset F captures (2026-10-14: capF, capF900, capF1440, capF4K, capFQ and the capFX* input variants) what
is left is per stage:
network stages binary16 s0 (byte 92), s1 (94), s2 (96) and, for the decoders, s3 (98) (engine/rr/rr2_ref.py)
enc0 two float32 (572, 576), two binary16 widened to float32 (580, 584), 1.0 (588, 1072), 1 (1036)
dec0_1 8.0 (208) and exp(t) (212); post: exp(t') (128); down / downp: 0.25 (204), 1.0 (296, 300)
NGX's model description is built by one routine of plain stores (mov [rdx + d], imm32 / ax / a pointer); describe()
replays it. It holds every tensor's pointer and size, s1 / s2 of each stage at 0x220 + 0x238 * n (a decoder level's
two kernels side by side, its s3 behind them), encoder 0's float32 pair at 0x230 and the binary16 pair at 0x185a, and
three one-value tensors (t at 0x1838, t' at 0x1848). s0 sits in the weights, before each encoder's relative position
tables (encoder 0: 16 bytes before; a decoder level uses its encoder's).
enc0-aux.bin is preset E's noise table (engine/rr/dll_weights.py) and then encoder 0's embedding, both branches: Wd
(32 x 64 FP8 codes, 43 inputs used), bd, Wu (32 x 32 FP8), bu (once, after the first branch), the FP8 matrices in the
kernel's fragment order: a fixed permutation of the index bits of the natural [in][out] order (FRAG_WD, FRAG_WU).
Everything written is derived from the user's DLL and stays on their machine.
"""
import math
import struct
import sys
from pathlib import Path
import numpy as np
sys.path.insert(0, str(Path(__file__).resolve().parent))
from dll_weights import noise_table # noqa: E402
from enc1_ref import inter # noqa: E402
from rr2_ref import ORDER, dll_weights, find_tables # noqa: E402
H = np.float16
NETWORK = ORDER # enc0 ... enc5, dec4_0, dec4_1, ... dec0_1
ARG_BYTES = dict({s: 104 for s in NETWORK}, enc0=1088, dec0_1=216, post=320, down=416, downp=416, pooled=32)
# the plane-arena regions of the network stages: argument offset -> region start in units of 1/3152 of NGX's arena
# (W0 * H0 / 32 bytes; the same on every capture)
ARENA = {
'enc0': '32:0,40:1024,64:2560,72:2240', 'enc1': '8:0,32:1536,40:2048,48:2240,64:3008,72:2304',
'enc2': '8:2048,32:1024,40:1216,48:2304,64:1472,72:1364', 'enc3': '8:1216,32:1280,40:1344,48:1364,64:1368,72:1372',
'enc4': '8:1344,32:1236,40:1216,48:1372,64:1261,72:1221', 'enc5': '8:1216,32:1256,48:1221,64:1262',
'dec4_0': '8:1256,16:1236,32:1216', 'dec4_1': '8:1216,32:1344,48:1262,56:1261,64:1364',
'dec3_0': '8:1344,16:1280,32:1216', 'dec3_1': '8:1216,32:1408,48:1364,56:1368,64:1488',
'dec2_0': '8:1408,16:1024,32:1216', 'dec2_1': '8:1216,32:2816,48:1488,56:1472,64:3136',
'dec1_0': '8:2816,16:1536,32:1024', 'dec1_1': '8:1024,32:2048,48:3136,56:3008,64:3072',
'dec0_0': '8:2048,16:0,32:1024', 'dec0_1': '8:1024,48:3072,56:2560,64:2816',
}
# fragment order of encoder 0's FP8 matrices in enc0-aux.bin: bit j of the table index is bit FRAG[j] of the natural
# index in * 32 + out (found on 310.9.1 and checked on every entry of both branches)
FRAG_WD = [5, 9, 8, 3, 6, 7, 0, 1, 2, 10, 4]
FRAG_WU = [5, 9, 8, 3, 6, 7, 0, 1, 2, 4]
def sections(dll):
pe = struct.unpack_from('<I', dll, 0x3c)[0]
count, opt = struct.unpack_from('<HH', dll, pe + 6)[0], struct.unpack_from('<H', dll, pe + 20)[0]
base = struct.unpack_from('<Q', dll, pe + 24 + 24)[0]
out = []
for i in range(count):
at = pe + 24 + opt + 40 * i
vsize, va, rsize, raw = struct.unpack_from('<IIII', dll, at + 8)
out.append((base + va, max(vsize, rsize), raw))
return out
def file_offset(dll, va):
for start, size, raw in sections(dll):
if start <= va < start + size:
return va - start + raw
raise ValueError(f'address {va:#x} is in no section')
def modrm(b, i):
m = b[i]
mod, reg, rm = m >> 6, (m >> 3) & 7, m & 7
if rm == 4 or (mod == 0 and rm == 5) or mod == 3:
raise ValueError
base = {2: 'rdx', 1: 'rcx'}.get(rm)
if mod == 0:
return base, reg, 0, 1
if mod == 1:
return base, reg, struct.unpack_from('<b', b, i + 1)[0], 2
return base, reg, struct.unpack_from('<i', b, i + 1)[0], 5
def replay(dll, i, image_base):
"""the stores of the routine starting at file offset i, up to its ret: ({offset: bytes}, {offset: pointer})"""
stores, pointers = {}, {}
eax, r8, rax_ptr = 0, 0, None
text = next((s, sz, raw) for s, sz, raw in sections(dll) if raw <= i < raw + sz)
for _ in range(4000):
b = dll[i:i + 16]
p66 = b[0] == 0x66
j = 1 if p66 else 0
rex = b[j] if 0x40 <= b[j] <= 0x4f else 0
j += 1 if rex else 0
op, w, rexr = b[j], bool(rex & 8), bool(rex & 4)
size = 8 if w else 2 if p66 else 4
if op == 0xc3:
return stores, pointers
if op == 0x33 and rex == 0x45 and b[j + 1] == 0xc0: # xor r8d, r8d
r8 = 0; i += j + 2; continue
if op == 0xb8 and not rex: # mov eax, imm32
eax = struct.unpack_from('<I', b, j + 1)[0]; i += j + 5; continue
if op == 0x8d and w and b[j + 1] == 0x05: # lea rax, [rip + d]
rel = struct.unpack_from('<i', b, j + 2)[0]
rax_ptr = text[0] + (i + j + 6 - text[2]) + rel; i += j + 6; continue
if op == 0x8b: # mov rax, [rip + d] / mov rax, rcx
rax_ptr = None; i += j + (2 if b[j + 1] >> 6 == 3 else 6); continue
if op == 0xc7: # mov [rdx + d], imm32
base, _, d, n = modrm(b, j + 1)
if base == 'rdx':
stores[d] = struct.pack('<q' if w else '<i', struct.unpack_from('<i', b, j + 1 + n)[0])
i += j + 1 + n + 4; continue
if op == 0x89: # mov [rdx + d], eax / ax / rax / r8
base, reg, d, n = modrm(b, j + 1)
if base == 'rdx' and reg == 0:
if w and not rexr:
pointers[d] = rax_ptr
else:
v = r8 if rexr else eax
stores[d] = (v & ((1 << (8 * size)) - 1)).to_bytes(size, 'little')
i += j + 1 + n; continue
raise ValueError(f'instruction at file offset {i:#x} ({b[:8].hex()})')
raise ValueError('no ret')
def describe(dll):
"""replay the routine that builds the model description: the one that stores encoder 0's tensor sizes 2752
(Wd, 2 x 32 x 43) at offset 0 and 64 (bd) at 0x10"""
marker = b'\xc7\x02\xc0\x0a\x00\x00\xc7\x42\x10\x40\x00\x00\x00'
found = []
at = dll.find(marker)
while at >= 0:
start = dll.rfind(b'\xcc\xcc', 0, at) + 2
try:
stores, pointers = replay(dll, start, None)
if 0x220 in stores and 0x1858 in stores:
found.append((stores, pointers))
except ValueError:
pass
at = dll.find(marker, at + 1)
if len(found) != 1:
raise SystemExit(f'found {len(found)} preset F model descriptions in the DLL (expected one): not a layout this '
'builder knows')
return found[0]
def h16(b):
return np.frombuffer(b, H)[0]
def values(dll):
"""{stage: {argument offset: bytes}} of every constant the runtime keeps"""
stores, pointers = describe(dll)
image = bytearray(max(d + len(v) for d, v in stores.items()))
for d in sorted(stores):
image[d:d + len(stores[d])] = stores[d]
word = lambda d: bytes(image[d:d + 2])
tables = find_tables(dll)
s0 = {}
for n, stage in enumerate(['enc0', 'enc1', 'enc2', 'enc3', 'enc4', 'enc5']):
at = tables[stage] - (16 if stage == 'enc0' else 8)
s0[stage] = dll[at:at + 2]
out = {}
for n, stage in enumerate(NETWORK):
v = {}
if stage.startswith('enc'):
k = int(stage[3])
v[92] = s0[stage]
d = 0x220 + 0x238 * k
v[94] = bytes(image[d:d + 4])
else:
level, half = int(stage[3]), int(stage[5])
v[92] = s0[f'enc{level + 1}']
d = 0x220 + 0x238 * (10 - level)
v[94] = bytes(image[d + 4 * half:d + 4 * half + 4])
v[98] = word(d + 8)
out[stage] = v
f16f32 = lambda b: struct.pack('<f', float(h16(b)))
one = struct.pack('<f', 1.0)
out['enc0'].update({572: bytes(image[0x230:0x234]), 576: bytes(image[0x234:0x238]), 580: f16f32(word(0x185c)), 584: f16f32(word(0x185a)),
588: one, 1036: struct.pack('<I', 1), 1072: one})
tensor = lambda d: dll[file_offset(dll, pointers[d]):file_offset(dll, pointers[d]) + 2]
expf = lambda b: struct.pack('<f', np.float32(math.exp(float(h16(b)))))
out['dec0_1'].update({208: struct.pack('<f', 8.0), 212: expf(tensor(0x1838))})
out['post'] = {128: expf(tensor(0x1848))}
for s in ('down', 'downp'):
out[s] = {204: struct.pack('<f', 0.25), 296: one, 300: one}
out['pooled'] = {}
return out
def templates(dll):
blocks = {}
for stage, v in values(dll).items():
b = bytearray(ARG_BYTES[stage])
for off, x in v.items():
b[off:off + len(x)] = x
blocks[stage] = bytes(b)
return blocks
def frag(natural, bits):
"""natural [in][out] FP8 codes -> the table's order: bit j of the table index is bit bits[j] of the natural index"""
flat = natural.ravel()
idx = np.arange(flat.size)
src = np.zeros_like(idx)
for j, b in enumerate(bits):
src |= ((idx >> j) & 1) << b
return flat[src].tobytes()
def aux_table(dll):
noise = noise_table(dll)
if noise is None:
raise SystemExit('no noise table found in the DLL')
w = dll_weights(dll, 'enc0')
def natural(b_operand): # B = q8(mat).T[inter(k)] (k x n) -> [in][out]
n = np.empty_like(b_operand)
n[np.asarray(inter(b_operand.shape[0]))] = b_operand
return n
out = [noise]
for b in range(2):
out += [frag(natural(np.asarray(w['ewd'][b])), FRAG_WD), np.asarray(w['ebd'][b], H).tobytes(),
frag(natural(np.asarray(w['ewu'][b])), FRAG_WU)]
if b == 0:
out.append(np.asarray(w['ebu'], H).tobytes())
return b''.join(out)
def main():
if len(sys.argv) < 3:
raise SystemExit(__doc__)
dll = Path(sys.argv[1]).read_bytes()
blocks, aux = templates(dll), aux_table(dll)
if sys.argv[2] == '--check':
model, bad = Path(sys.argv[3]), 0
theirs = (model / 'enc0-aux.bin').read_bytes()
print('enc0-aux.bin:', 'same' if theirs == aux else 'DIFFERENT'); bad += theirs != aux
for stage, ours in blocks.items():
path = model / f'{stage}.args'
if not path.exists():
continue
t = path.read_bytes()[:len(ours)]
diff = [i for i in range(len(ours)) if ours[i] != t[i]]
words = sorted({i // 4 * 4 for i in diff})
print(f'{stage}: ' + ('same' if not diff else f'{len(words)} words differ: {words[:40]}'))
sys.exit(1 if bad else 0)
out = Path(sys.argv[2])
out.mkdir(parents=True, exist_ok=True)
for stage, b in blocks.items():
(out / f'{stage}.args').write_bytes(b)
(out / 'enc0-aux.bin').write_bytes(aux)
if __name__ == '__main__':
main()
engine/rr/rr2_w.comp
+370
View File
@@ -0,0 +1,370 @@
#version 460
// Decoder 0's first kernel of preset F (rr2_m.comp KIND 2 with C = 32) for D4R_RR_RELAXED=3, with ONE wave per
// 8x8-token window where rr2_m.comp uses four. The wave walks the four 16-token slices itself, so every weight, key
// and value tile is read once per window; the expansion plus skip is computed once and kept in registers; the scores'
// cross quadrant is computed once and reused transposed (the scores are symmetric); FP8 values are converted in pairs.
// Shared memory stays at 4 KiB, above which this GPU runs fewer waves (engine/rr/README.md). The arithmetic is
// rr2_m.comp's with -DRR_RELAXED=2 -DRR_NATIVE and the arena comes out byte-identical to that build
// (build/rr/f/nv/invest-20261013/REPORT.md: 453 us against 581 us at 1280x720). compile_rr2.py uses it for dec0_0.
// KIND 4 is encoder 0's block the same way (the embedding of the gather pass's record, the tokens, the guide head to
// HD), included with RR2_EMBED and RR2_XF_ARENA by the shader engine/rr/gen_rr2_enc0.py --one-wave generates.
#extension GL_KHR_cooperative_matrix : require
#extension GL_KHR_memory_scope_semantics : require
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_8bit_storage : require
#extension GL_EXT_float_e4m3 : require
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_EXT_control_flow_attributes : require
#extension GL_GOOGLE_include_directive : require
#ifndef RR2_EMBED
layout(local_size_x = 32) in;
#endif
#include "rr2_stage.h"
#include "rr_exact.glsl"
#if !(KIND == 2 || KIND == 4) || CM != 0 || GUIDE != 0 || C != 32
#error rr2_w.comp: a 32-channel decoder's first kernel, or encoder 0's block
#endif
#if KIND == 4 && !(defined(RR2_EMBED) && defined(RR2_XF_ARENA))
#error rr2_w.comp: encoder 0's block reads the gather pass's record (gen_rr2_enc0.py --split block --one-wave)
#endif
#define MT (C / 16)
#ifdef RR2_EMBED
#define PRM(i) param.w[i]
#define RR2_WGX gl_WorkGroupID.x
#else
layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param;
#define PRM(i) param.p[i]
#define RR2_WGX gl_WorkGroupID.x
#endif
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16;
layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u;
layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define HAC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CMJ gl_CooperativeMatrixLayoutColumnMajor
#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM)
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); }
#define Q8(dst, src) Q8PAIRS(dst, src)
#define ACT8(dst, h) [[unroll]] for (int i_ = 0; i_ < 8; i_ += 2) { const vec2 x_ = vec2((h)[i_], (h)[i_ + 1]); const vec2 y_ = clamp(x_, vec2(-2.0), vec2(2.0)); const vec2 g_ = fma(y_, fma(-abs(y_), vec2(0.0810546875), vec2(0.412109375)), vec2(0.5)); fe4m3vec2 t_; saturatedConvertEXT(t_, x_ * g_); (dst)[i_] = t_.x; (dst)[i_ + 1] = t_.y; }
const uint HEAD = W_WQK1 - W_WQK0, BRANCH = W_WD1 - W_WD0;
shared uint XT[16 * C];
shared uint KVS[512]; // the expansion (input phase), then per head V^T, then QK
// Shared memory is what limits how many waves this GPU keeps running: above about 4 KiB per wave the rate falls in
// proportion (build/rr/f/nv/cost/occ2.py). So the input phase borrows KVS, which the attention only needs later
// (LOW at word 0, EX at word 512), and the block's unquantised input stays in registers (xr), not in memory.
// (4 KiB in all: the low tokens borrow XT, which is only written after they are in registers; the expansion borrows
// KVS; V^T is stored, read into registers, and then overwritten by QK.)
#if KIND == 2 && (4 * CL > 16 * C || 16 * C > 512)
#error rr2_w.comp: the input phase no longer fits
#endif
#define EX_AT 0u
#if KIND == 4
uint rr2_xf; // the window's record from the gather pass: the index of its first arena word
shared float16_t HD[64 * 4]; // the guide head per token (raster order), for the translated tail
const uint EBRANCH = W_EWD1 - W_EWD0;
#endif
uint pack4(floate4m3_t a, floate4m3_t b, floate4m3_t c, floate4m3_t d)
{
return code(a) | (code(b) << 8u) | (code(c) << 16u) | (code(d) << 24u);
}
f16vec2 e4pair(uint w, uint h) { return f16vec2(uintBitsToFloate4m3EXT(u8vec2(uint8_t(w >> (16u*h)), uint8_t(w >> (16u*h+8u))))); }
float16_t e4(uint w, uint b) { return float16_t(uintBitsToFloate4m3EXT(uint8_t((w >> (8u * b)) & 0xffu))); }
#ifdef RR2_EMBED
void rr2_block()
#else
void main()
#endif
{
const uint lane = gl_SubgroupInvocationID;
const uint hl = lane >> 4u, col = lane & 15u;
RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u);
const int ex = int(PRM(0)), ey = int(PRM(1)), hx = ex / 2, hy = ey / 2;
const uint inOff = PRM(2), skOff = PRM(4), outOff = PRM(8), shift = PRM(20);
const float16_t s1 = hb(PRM(23) >> 16u), s2 = hb(PRM(24) & 0xffffu), s3 = hb(PRM(24) >> 16u);
const int x0 = 8 * int(RR2_WGX) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u);
const uint plane = uint(ex * ey) * 32u, lplane = uint(hx * hy) * 32u;
#if KIND == 4
// ---------------------------------------------------------------- the block's input: x = bu + sum over two
// branches of Wu . act(Wd . features + bd); the features are the gather pass's record, read once per slice
B8 qx[4][MT]; // x as FP8, this lane's part per token slice (the residual step decodes it again)
{
B8 fb[4][4];
ACC X[4][MT];
[[unroll]] for (uint wv = 0u; wv < 4u; ++wv) {
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) coopMatLoad(fb[wv][kt], arena.w, rr2_xf + wv * 256u + 4u * kt, 16u, CMJ);
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) X[wv][m][i] = float(w16.w[W_EBU / 2u + 16u * m + 8u * hl + uint(i)]);
}
[[unroll]] for (uint br = 0u; br < 2u; ++br) {
A8 ewd[2][4], ewu0[MT], ewu1[MT];
[[unroll]] for (uint c = 0u; c < 2u; ++c)
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) LDA(ewd[c][kt], W_EWD0 + EBRANCH * br, 64u, c, kt);
[[unroll]] for (uint m = 0u; m < MT; ++m) {
LDA(ewu0[m], W_EWU0 + EBRANCH * br, 32u, m, 0u); LDA(ewu1[m], W_EWU0 + EBRANCH * br, 32u, m, 1u);
}
[[unroll]] for (uint wv = 0u; wv < 4u; ++wv) {
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
ACC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = float(w16.w[(W_EBD0 + EBRANCH * br) / 2u + 16u * c + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < 4u; kt += 2u) K32(h, ewd[c][kt], fb[wv][kt], ewd[c][kt + 1u], fb[wv][kt + 1u])
ACT8(Hb[c], h)
}
[[unroll]] for (uint m = 0u; m < MT; ++m) K32(X[wv][m], ewu0[m], Hb[0], ewu1[m], Hb[1])
}
}
[[unroll]] for (uint wv = 0u; wv < 4u; ++wv)
[[unroll]] for (uint m = 0u; m < MT; ++m) {
Q8(qx[wv][m], X[wv][m])
coopMatStore(qx[wv][m], XT, wv * 4u * C + 4u * m, C / 4u, CMJ);
}
}
#else
// ---------------------------------------------------------------- the block's input: expansion plus skip
float16_t xr[4][MT][8]; // this lane's part of x for each token slice: [slice][channel tile][element]
{
const int X0 = x0 >> 1, Y0 = y0 >> 1;
for (uint i = lane; i < 4u * CL; i += 32u) {
const uint r = i / (CL / 4u), wd = i % (CL / 4u);
const int X = mirror(X0 + int(r & 3u), hx), Y = mirror(Y0 + int(r >> 2u), hy);
XT[i] = arena.w[(inOff + (wd >> 3u) * lplane + uint(Y * hx + X) * 32u) / 4u + (wd & 7u)];
}
B8 lb[CL / 16];
[[unroll]] for (uint kt = 0u; kt < CL / 16; ++kt) coopMatLoad(lb[kt], XT, 4u * kt, CL / 4u, CMJ);
for (uint m = 0u; m < 4u * MT; ++m) {
ACC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_BE / 2u + 16u * m + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < CL / 16; kt += 2u) {
A8 w0, w1;
LDA(w0, W_WE, CL, m, kt); LDA(w1, W_WE, CL, m, kt + 1u);
K32(a, w0, lb[kt], w1, lb[kt + 1u])
}
B8 q;
Q8(q, a)
coopMatStore(q, KVS, EX_AT + 4u * m, C, CMJ);
}
// the same x for the residual step, in this lane's own layout (token slot 16 wv + col, channel slots
// 16 m + 8 hl + i = bytes of the words 4 m + 2 hl and 4 m + 2 hl + 1)
[[unroll]] for (uint wv = 0u; wv < 4u; ++wv) {
const uint k = inter(16u * wv + col);
const int kx = x0 + int(k & 7u), ky = y0 + int(k >> 3u);
const uint tb = (skOff + uint(mirror(ky, ey) * ex + mirror(kx, ex)) * 32u) / 4u;
const uint r = uint(((ky >> 1) - Y0) * 4 + ((kx >> 1) - X0)), s = uint(2 * (ky & 1) + (kx & 1));
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (uint h2 = 0u; h2 < 2u; ++h2) {
const uint wd = 4u * m + 2u * hl + h2;
const uint sk = arena.w[tb + (wd >> 3u) * (plane >> 2u) + (wd & 7u)];
const uint e = KVS[EX_AT + r * C + s * (C / 4u) + wd];
[[unroll]] for (uint pair = 0u; pair < 2u; ++pair) {
precise f16vec2 x_p = e4pair(sk, pair) * f16vec2(s3);
precise f16vec2 x = e4pair(e, pair) + x_p;
xr[wv][m][4u * h2 + 2u*pair] = x.x;
xr[wv][m][4u * h2 + 2u*pair+1u] = x.y;
}
}
[[unroll]] for (uint m = 0u; m < MT; ++m) {
B8 q;
Q8(q, xr[wv][m])
coopMatStore(q, XT, wv * 4u * C + 4u * m, C / 4u, CMJ);
}
}
}
#endif
// ---------------------------------------------------------------- attention
ACC AT[4][MT];
[[unroll]] for (uint wv = 0u; wv < 4u; ++wv)
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) AT[wv][m][i] = float(w16.w[W_BO / 2u + 16u * m + 8u * hl + uint(i)]);
[[unroll]] for (uint hd = 0u; hd < NH; ++hd) {
// V^T of every token tile first: stored, read back as the product's A operands, then QK takes its place
A8 kt0[4], kt1[4], v0[2][2], v1[2][2], wo0[MT], wo1[MT];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
B8 wvt[MT];
[[unroll]] for (uint kt = 0u; kt < MT; ++kt) LDB(wvt[kt], W_WV0 + HEAD * hd, 32u, kt, c);
[[unroll]] for (uint n = 0u; n < 4u; ++n) {
ACC v = ACC(0.0);
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) {
A8 xa0, xa1;
coopMatLoad(xa0, XT, n * 4u * C + 4u * kt, C / 4u, RM);
coopMatLoad(xa1, XT, n * 4u * C + 4u * kt + 4u, C / 4u, RM);
K32(v, xa0, wvt[kt], xa1, wvt[kt + 1u])
}
A8 va;
Q8(va, v)
coopMatStore(va, KVS, c * 256u + 4u * n, 16u, RM);
}
}
[[unroll]] for (uint c = 0u; c < 2u; ++c)
[[unroll]] for (uint ks = 0u; ks < 2u; ++ks) {
coopMatLoad(v0[c][ks], KVS, c * 256u + 8u * ks, 16u, RM);
coopMatLoad(v1[c][ks], KVS, c * 256u + 8u * ks + 4u, 16u, RM);
}
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
A8 wk[MT];
[[unroll]] for (uint kt = 0u; kt < MT; ++kt) LDA(wk[kt], W_WQK0 + HEAD * hd, C, c, kt);
[[unroll]] for (uint n = 0u; n < 4u; ++n) {
ACC a = ACC(0.0);
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) {
B8 xb0, xb1;
coopMatLoad(xb0, XT, n * 4u * C + 4u * kt, C / 4u, CMJ);
coopMatLoad(xb1, XT, n * 4u * C + 4u * kt + 4u, C / 4u, CMJ);
K32(a, wk[kt], xb0, wk[kt + 1u], xb1)
}
A8 ka;
Q8(ka, a)
coopMatStore(ka, KVS, n * 128u + 4u * c, 8u, RM);
}
}
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) {
coopMatLoad(kt0[kt], KVS, kt * 128u, 8u, RM);
coopMatLoad(kt1[kt], KVS, kt * 128u + 4u, 8u, RM);
}
[[unroll]] for (uint m = 0u; m < MT; ++m) {
LDA(wo0[m], W_WO0 + HEAD * hd, 32u, m, 0u); LDA(wo1[m], W_WO0 + HEAD * hd, 32u, m, 1u);
}
// Preserve all query operands before KVS becomes a 32x32 f16 numerator cache.
B8 allQ[4][2];
[[unroll]] for (uint wv = 0u; wv < 4u; ++wv)
[[unroll]] for (uint c = 0u; c < 2u; ++c) coopMatLoad(allQ[wv][c], KVS, wv * 128u + 4u * c, 8u, CMJ);
[[unroll]] for (uint wv = 0u; wv < 4u; ++wv) {
if (wv == 2u) memoryBarrierShared();
B8 Pb[4];
{
float16_t e[4][8];
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) {
if (wv >= 2u && kt < 2u) {
HAC cached;
// Transpose keys/query roles. Stride is measured in uint words.
coopMatLoad(cached, KVS, (wv - 2u) * 256u + kt * 8u, 16u, CMJ);
[[unroll]] for (int i = 0; i < 8; ++i) e[kt][i] = cached[i];
} else {
HAC c;
coopMatLoad(c, w16.w, (W_BIAS0 + HEAD * hd) / 2u + kt * 1024u + 16u * wv, 64u, RM);
ACC p_ = coopMatMulAdd(kt1[kt], allQ[wv][1], coopMatMulAdd(kt0[kt], allQ[wv][0], ACC(c)));
[[unroll]] for (int i = 0; i < 8; ++i) e[kt][i] = hb(uint(w16u.w[W_CUBIC / 2u + bh(float16_t(p_[i]))]));
if (wv < 2u && kt >= 2u) {
HAC cached;
[[unroll]] for (int i = 0; i < 8; ++i) cached[i] = e[kt][i];
coopMatStore(cached, KVS, (kt - 2u) * 256u + wv * 8u, 16u, RM);
}
}
}
f16vec2 ep[4][4];
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt)
[[unroll]] for (uint q = 0u; q < 4u; ++q) ep[kt][q] = f16vec2(e[kt][2u * q], e[kt][2u * q + 1u]);
precise f16vec2 st[2];
[[unroll]] for (uint tq = 0u; tq < 2u; ++tq) {
const uint q0 = 2u * tq, q1 = q0 + 1u;
precise f16vec2 ev = (ep[0][q0] + ep[0][q1]) + (ep[2][q0] + ep[2][q1]);
precise f16vec2 od = (ep[1][q0] + ep[1][q1]) + (ep[3][q0] + ep[3][q1]);
st[tq] = ev + od;
}
precise f16vec2 loc = st[0] + st[1];
const uint other = subgroupShuffleXor(packFloat2x16(loc), 16u);
precise f16vec2 o = hl == 0u ? loc + unpackFloat2x16(other) : unpackFloat2x16(other) + loc;
precise float16_t rs = o.x + o.y;
const float16_t r = float16_t(1.0 / float(rs));
const f16vec2 r2 = f16vec2(r, r);
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt)
[[unroll]] for (uint q = 0u; q < 4u; ++q) {
precise f16vec2 pv = ep[kt][q] * r2;
{ const fe4m3vec2 t_ = fe4m3vec2(pv); Pb[kt][2u * q] = t_.x; Pb[kt][2u * q + 1u] = t_.y; }
}
}
B8 Ob[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
ACC o = ACC(0.0);
[[unroll]] for (uint ks = 0u; ks < 2u; ++ks) K32(o, v0[c][ks], Pb[2u * ks], v1[c][ks], Pb[2u * ks + 1u])
Q8(Ob[c], o)
}
[[unroll]] for (uint m = 0u; m < MT; ++m) K32(AT[wv][m], wo0[m], Ob[0], wo1[m], Ob[1])
}
}
// ---------------------------------------------------------------- r = x + attention * s1; the MLP
float rv[4][MT][8];
B8 Nb[4][MT];
ACC ML[4][MT];
const float f1 = float(s1), f2 = float(s2);
[[unroll]] for (uint wv = 0u; wv < 4u; ++wv)
[[unroll]] for (uint m = 0u; m < MT; ++m) {
[[unroll]] for (int i = 0; i < 8; ++i) {
#if KIND == 4
const float xb = float(qx[wv][m][i]);
#else
const float xb = float(xr[wv][m][i]);
#endif
const float r = fma(AT[wv][m][i], f1, xb);
rv[wv][m][i] = r;
ML[wv][m][i] = float(w16.w[W_BU / 2u + 16u * m + 8u * hl + uint(i)]);
}
Q8(Nb[wv][m], rv[wv][m])
}
[[unroll]] for (uint br = 0u; br < NB; ++br) {
const uint w1 = W_WD0 + BRANCH * br, b1 = W_BD0 + BRANCH * br, w2 = W_WU0 + BRANCH * br;
A8 wd[2][MT], wu0[MT], wu1[MT];
[[unroll]] for (uint c = 0u; c < 2u; ++c)
[[unroll]] for (uint kt = 0u; kt < MT; ++kt) LDA(wd[c][kt], w1, C, c, kt);
[[unroll]] for (uint m = 0u; m < MT; ++m) { LDA(wu0[m], w2, 32u, m, 0u); LDA(wu1[m], w2, 32u, m, 1u); }
[[unroll]] for (uint wv = 0u; wv < 4u; ++wv) {
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
ACC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = float(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) K32(h, wd[c][kt], Nb[wv][kt], wd[c][kt + 1u], Nb[wv][kt + 1u])
ACT8(Hb[c], h)
}
[[unroll]] for (uint m = 0u; m < MT; ++m) K32(ML[wv][m], wu0[m], Hb[0], wu1[m], Hb[1])
}
}
// ---------------------------------------------------------------- y = r + MLP * s2: the stage's tokens
#if KIND == 4
A8 wh[MT];
[[unroll]] for (uint kt = 0u; kt < MT; ++kt) LDA(wh[kt], W_WH, C, 0u, kt);
#endif
[[unroll]] for (uint wv = 0u; wv < 4u; ++wv) {
const uint t = inter(16u * wv + col);
const int tx = x0 + int(t & 7u), ty = y0 + int(t >> 3u);
const bool inside = tx >= 0 && ty >= 0 && tx < ex && ty < ey;
const uint o = (outOff + uint(ty * ex + tx) * 32u) >> 2u;
#if KIND == 4
B8 Tm[MT];
#endif
[[unroll]] for (uint m = 0u; m < MT; ++m) {
floate4m3_t Tb[8];
[[unroll]] for (int i = 0; i < 8; i += 2) {
fe4m3vec2 q; saturatedConvertEXT(q, vec2(fma(ML[wv][m][i], f2, rv[wv][m][i]), fma(ML[wv][m][i+1], f2, rv[wv][m][i+1])));
Tb[i] = q.x; Tb[i+1] = q.y;
}
#if KIND == 4
[[unroll]] for (int i = 0; i < 8; ++i) Tm[m][i] = Tb[i];
#endif
const uint w0 = pack4(Tb[0], Tb[1], Tb[2], Tb[3]), w1 = pack4(Tb[4], Tb[5], Tb[6], Tb[7]);
const uint at = 4u * (m & 1u) + 2u * hl;
if (inside) { arena.w[o + (m >> 1u) * (plane >> 2u) + at] = w0; arena.w[o + (m >> 1u) * (plane >> 2u) + at + 1u] = w1; }
}
#if KIND == 4
{ // the guide head (rows 0-3 on lanes 0-15), for the translated tail
ACC hv;
[[unroll]] for (int i = 0; i < 8; ++i) hv[i] = float(w16.w[W_BH / 2u + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) K32(hv, wh[kt], Tm[kt], wh[kt + 1u], Tm[kt + 1u])
if (hl == 0u) [[unroll]] for (uint c = 0u; c < 4u; ++c) HD[t * 4u + c] = float16_t(hv[c]);
}
#endif
}
}
engine/rr/rr_attn.glsl
+91
View File
@@ -0,0 +1,91 @@
// The attention core shared by the 32-channel-head stages (engine/rr/enc1_m.comp, dec0_m.comp), included inside the
// head loop: this wave's 16 query slots against the window's 96 key slots. Expects KS (K, token slot x 32 FP8
// channels in k order) and VS (V^T, 32 channels x 96 token slots) filled for head hd, the bias at W_BIAS0 + HEAD * hd
// (enc1_pack: 96 x 64, [key slot][query slot]) and the cubic table at W_CUBIC. Leaves Ob[2]: O^T (32 channels x
// this wave's queries) as the out-projection's B operand. RR_KS / RR_VS (+ word offsets RR_KS0 / RR_VS0) name
// the arrays when they are not KS and VS.
#ifndef RR_KS
#define RR_KS KS
#define RR_KS0 0u
#endif
#ifndef RR_VS
#define RR_VS VS
#define RR_VS0 0u
#endif
B8 Qb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) LDS8B(Qb[c], RR_KS, RR_KS0 + wv * 128u + 4u * c, 8u);
// scores S^T (keys x this wave's queries) + bias, then the cubic softmax down each column
B8 Pb[6];
{
float16_t e[6][8];
[[unroll]] for (uint kt = 0u; kt < 6u; ++kt) {
HAC c;
#ifdef RR_GFX11
// the accumulator layout the stages are written for (rr_gfx11.glsl): element i is row i + 8 hl
[[unroll]] for (int i = 0; i < 8; ++i)
c[i] = w16.w[(W_BIAS0 + HEAD * hd) / 2u + kt * 1024u + 16u * wv + (8u * hl + uint(i)) * 64u + col];
#else
coopMatLoad(c, w16.w, (W_BIAS0 + HEAD * hd) / 2u + kt * 1024u + 16u * wv, 64u, RM);
#endif
A8 k0, k1;
LDS8A(k0, RR_KS, RR_KS0 + kt * 128u, 8u);
LDS8A(k1, RR_KS, RR_KS0 + kt * 128u + 4u, 8u);
K32H(c, k0, Qb[0], k1, Qb[1])
// the cubic numerator from enc1_pack.cubic_table, indexed by the score's binary16 bits
[[unroll]] for (int i = 0; i < 8; ++i) e[kt][i] = hb(uint(w16u.w[W_CUBIC / 2u + bh(c[i])]));
}
// key slot 16kt + 8hl + i is NVIDIA's column 8j + 2tt + half with j = 4(kt >> 1) + (kt & 1) + 2((i >> 1) & 1),
// tt = 2hl + (i >> 2), half = i & 1 (enc1_ref.softmax_cubic's tree)
// packed: component h of a pair is column parity h (elements i = 4tq + h, i + 2), so the tree runs on f16x2
f16vec2 ep[6][4];
[[unroll]] for (uint kt = 0u; kt < 6u; ++kt)
[[unroll]] for (uint q = 0u; q < 4u; ++q) ep[kt][q] = f16vec2(e[kt][2u * q], e[kt][2u * q + 1u]);
precise f16vec2 st[2];
[[unroll]] for (uint tq = 0u; tq < 2u; ++tq) { // pairs q = 2tq (i0) and 2tq + 1 (i1 = i0 + 2)
const uint q0 = 2u * tq, q1 = q0 + 1u;
precise f16vec2 ev = ((ep[0][q0] + ep[0][q1]) + (ep[2][q0] + ep[2][q1])) + (ep[4][q0] + ep[4][q1]);
precise f16vec2 od = ((ep[1][q0] + ep[1][q1]) + (ep[3][q0] + ep[3][q1])) + (ep[5][q0] + ep[5][q1]);
st[tq] = ev + od;
}
precise f16vec2 loc = st[0] + st[1];
const uint other = subgroupShuffleXor(packFloat2x16(loc), 16u);
precise f16vec2 o = loc + unpackFloat2x16(other);
precise float16_t rs = o.x + o.y;
// f16(1 / rs) for a binary16 rs: x * m = 1 + d with |d| >= 2^-22 for every binary16 midpoint m, so the f32
// reciprocal (within an ulp) already rounds correctly
const float16_t r = float16_t(1.0 / float(rs));
const f16vec2 r2 = f16vec2(r, r);
[[unroll]] for (uint kt = 0u; kt < 6u; ++kt) {
#ifdef RR_GFX11
F8V pc_ = uvec2(0u);
#endif
[[unroll]] for (uint q = 0u; q < 4u; ++q) {
precise f16vec2 pv = ep[kt][q] * r2;
#ifdef RR_GFX11
pc_[q >> 1u] |= (q8h(pv.x) | (q8h(pv.y) << 8u)) << (16u * (q & 1u));
#else
// a probability is in [0, 1]: the plain conversion, no saturation clamp, both in one instruction
{ const fe4m3vec2 t_ = fe4m3vec2(pv); Pb[kt][2u * q] = t_.x; Pb[kt][2u * q + 1u] = t_.y; }
#endif
}
#ifdef RR_GFX11
G11_B(Pb[kt], pc_)
#endif
}
}
// O^T = V^T P^T, then the out-projection onto the residual
B8 Ob[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
RAC o = RAC(RV(float16_t(0.0)));
[[unroll]] for (uint ks = 0u; ks < 3u; ++ks) {
A8 v0, v1;
LDS8A(v0, RR_VS, RR_VS0 + c * 384u + 8u * ks, 24u);
LDS8A(v1, RR_VS, RR_VS0 + c * 384u + 8u * ks + 4u, 24u);
K32(o, v0, Pb[2u * ks], v1, Pb[2u * ks + 1u])
}
#ifdef RR_GFX11
{ F8V oc_; G11_CODES(oc_, o) G11_B(Ob[c], oc_) }
#else
Q8(Ob[c], o)
#endif
}
engine/rr/rr_capture.py
+160
View File
@@ -0,0 +1,160 @@
"""Ray Reconstruction captures: the launches of one denoiser frame as the CUDA bridge recorded them.
A capture directory holds what the bridge writes with
D4R_CUDA_VERBOSE=1 D4R_CUDA_REPLAY_DUMP_DIR=CAP/replay D4R_CUDA_REPLAY_DUMP_FILTER=cuda_dldn
D4R_CUDA_LAUNCH_STATS=1 D4R_CUDA_LAUNCH_STATS_FILTER=cuda_dldn D4R_CUDA_LAUNCH_DUMP_DIR=CAP/dump
plus the run's log as CAP/run/run.log (texture sampler settings and array shapes come from the verbose
cuTexObjectCreate / cuSurfObjectCreate / cuArrayCreate lines).
replay/replay-NNNNNN-KERNEL/ args.bin, manifest.txt and every buffer the launch points into, *before* it ran
dump/launch-NNN-KERNEL-argOOO-{buffer,texture,surface}-... every resource *after* the launch ran
Captures are reference data only: nothing at runtime reads them.
"""
from dataclasses import dataclass, field
from pathlib import Path
import re
import struct
# CUarray_format codes the denoiser uses, and how a Vulkan image of the same texels is described
# (kvk's BITS column: 8 = UNORM bytes, 9 = SNORM bytes, 16 = binary16, 32 = binary32)
ARRAY_FORMATS = {16: 16, 32: 32, 1: 8, 200: 9, 194: 8}
@dataclass
class Image:
offset: int
kind: str # 'texture' or 'surface'
width: int
height: int
channels: int
format: int # CUarray_format
path: Path # post-launch contents
linear: bool = False # textures: bilinear filtering
@property
def bits(self):
return ARRAY_FORMATS[self.format]
@property
def texel_bytes(self):
return self.channels * (1 if self.bits in (8, 9) else self.bits // 8)
@dataclass
class Pointer:
offset: int
alloc: int
delta: int
@dataclass
class Launch:
sequence: int
kernel: str
directory: Path
grid: tuple
block: tuple
args: bytes
allocs: dict = field(default_factory=dict) # index -> (base, size)
pointers: list = field(default_factory=list)
images: list = field(default_factory=list)
def pre(self, alloc):
"""an allocation's contents before the launch"""
return (self.directory / f'alloc-{alloc}.bin').read_bytes()
def pointer(self, offset):
return next(p for p in self.pointers if p.offset == offset)
def image(self, offset):
return next(i for i in self.images if i.offset == offset)
def weights_at(self, offset):
"""the weight image (the pointer at argument OFFSET) as NGX uploaded it; 16 zero bytes for a stage without one"""
pointer = next((p for p in self.pointers if p.offset == offset), None)
return self.pre(pointer.alloc) if pointer is not None else bytes(16)
@property
def weights(self):
"""the DLSS-D kernels' weight image (argument 40)"""
return self.weights_at(40)
@property
def arena_alloc(self):
"""the plane arena: the allocation every arena pointer of the layer points into"""
candidates = [p.alloc for p in self.pointers]
return max(set(candidates), key=lambda a: self.allocs[a][1]) if candidates else None
def post_buffer(self, capture, alloc):
base = self.allocs[alloc][0]
for path in (capture.root / 'dump').glob(f'launch-{self.sequence:03d}-{self.kernel}-arg*-buffer-*'):
m = re.search(r'-base0x([0-9a-f]+)\+', path.name)
if m and int(m.group(1), 16) == base:
return path.read_bytes()
raise FileNotFoundError(f'no post-launch dump of allocation {alloc} for launch {self.sequence}')
def u32(self, offset):
return struct.unpack_from('<I', self.args, offset)[0]
def f32(self, offset):
return struct.unpack_from('<f', self.args, offset)[0]
class Capture:
def __init__(self, root):
self.root = Path(root)
log_path = self.root / 'run' / 'run.log'
self.log = log_path.read_text(errors='replace') if log_path.exists() else ''
# (log position, texture handle, filter mode) of every texture object creation, in order
self._tex_events = [(m.start(), int(m.group(2), 16), int(m.group(1))) for m in re.finditer(
r'cuTexObjectCreate .*? filter=(\d+) .*? object=(0x[0-9a-f]+)', self.log)]
self.launches = []
for d in sorted((self.root / 'replay').iterdir()):
self.launches.append(self._launch(d))
def filter_at(self, handle, position):
value = None
for pos, h, f in self._tex_events:
if pos > position:
break
if h == handle:
value = f
return value
def _launch(self, directory):
manifest = (directory / 'manifest.txt').read_text().split('\n')
kernel = manifest[0].split()[1]
launch = [int(x) for x in manifest[1].split()[1:]]
sequence = int(re.search(r'replay-(\d+)-', directory.name).group(1))
args = (directory / 'args.bin').read_bytes()
result = Launch(sequence, kernel, directory, tuple(launch[0:3]), tuple(launch[3:6]), args)
for line in manifest[2:]:
f = line.split()
if not f:
continue
if f[0] == 'alloc':
result.allocs[int(f[1])] = (int(f[2], 16), int(f[3]))
elif f[0] == 'pointer':
result.pointers.append(Pointer(int(f[1]), int(f[2]), int(f[3])))
# the log position of this launch, so a reused texture handle resolves to its settings at the time
position = self.log.find(f'launch[{sequence}] {kernel}')
if position < 0:
position = len(self.log)
for path in sorted((self.root / 'dump').glob(f'launch-{sequence:03d}-{kernel}-arg*')):
m = re.search(r'-arg(\d+)-(texture|surface)-(\d+)x(\d+)x(\d+)-fmt(-?\d+)\.bin$', path.name)
if not m:
continue
offset = int(m.group(1))
image = Image(offset, m.group(2), int(m.group(3)), int(m.group(4)), int(m.group(5)), int(m.group(6)), path)
if image.kind == 'texture':
handle = struct.unpack_from('<Q', args, offset)[0]
image.linear = self.filter_at(handle, position) == 1
result.images.append(image)
result.images.sort(key=lambda i: (i.kind != 'texture', i.offset))
return result
def find(self, kernel_part, frame=None):
"""the launches whose kernel name contains KERNEL_PART, in order; FRAME (1-based) selects one"""
found = [l for l in self.launches if kernel_part in l.kernel]
return found if frame is None else found[frame - 1]
engine/rr/rr_decfilter.glsl
+65
View File
@@ -0,0 +1,65 @@
// The decoders' output filter (engine/rr/dec0_ref.outputs' surface 112), shared by dec1_m.comp and swin_m.comp: included
// after the head has been written to OUTF (16 f16 per window token, raster order) and a barrier. Texture 88 (one level
// down, bilinear) at the window's 10 x 10 tap positions into TAP; then two threads per token: the softmax of channels
// 0-8 over the 3x3 taps, one f16 fma chain per colour channel, blended with texture 96 by the gate (channel 9) plus
// texture 104, to surface 112.
// texture 88 (one level down) at the window's 10 x 10 full-resolution tap positions (clamped like the filter's
// coordinates), through the bilinear sampler at NVIDIA's coordinates: the reference's texture unit does not round
// like f32 arithmetic (computing the 1/4, 3/4 blend by hand left ~0.2% of surface 112 one ulp off, capDT 2026-10-12)
if (tid < 100u) {
const int xi = clampTap(x0 - 1 + int(tid % 10u), ex), yi = clampTap(y0 - 1 + int(tid / 10u), ey);
const float iw = uintBitsToFloat(param.p[34]), ih = uintBitsToFloat(param.p[35]);
precise float cu = fma(iw, float(xi), iw * 0.5), cv = fma(ih, float(yi), ih * 0.5);
const vec4 bl = textureLod(T88, vec2(cu, cv), 0.0);
[[unroll]] for (uint c = 0u; c < 3u; ++c) TAP[tid * 4u + c] = float16_t(bl[c]);
}
barrier();
// two threads per token: half 0 filters channels 0 and 1, half 1 channel 2
{
const uint t = tid >> 1u, hf = tid & 1u;
const int x = x0 + int(t & 7u), y = y0 + int(t >> 3u);
if (x < 0 || y < 0 || x >= ex || y >= ey) return;
float16_t a[10];
[[unroll]] for (uint c = 0u; c < 10u; ++c) a[c] = OUTF[t * 16u + c];
// surface 112: softmax of channels 0-8 over the 3x3 taps, one f16 fma chain per channel; the gate (channel 9)
// blends it with texture 96; + texture 104
const float16_t ma = max(max(max(a[0], a[2]), a[4]), a[6]), mb = max(max(max(a[1], a[3]), a[5]), a[7]);
const float16_t mx = max(max(ma, mb), a[8]);
float e[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) { precise float16_t d = a[i] - mx; e[i] = exp2(float(d) * 1.4426950216293335); }
precise float s = (e[0] + e[1]) + 0.0;
s = s + (e[2] + e[3]);
s = s + (e[4] + e[5]);
s = s + (e[6] + e[7]);
s = s + e[8];
const float16_t rr = rcp16(s);
float16_t wg[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) wg[i] = h16(r16(e[i]) * float(rr));
const uint tx = uint(x - x0), ty = uint(y - y0), c0 = 2u * hf, nc = 2u - hf;
float16_t f[2] = float16_t[2](0.0hf, 0.0hf);
bool mid = false;
[[unroll]] for (uint k = 0u; k < 9u; ++k) {
const uint at = ((ty + k / 3u) * 10u + tx + k % 3u) * 4u + c0;
[[unroll]] for (uint c = 0u; c < 2u; ++c) f[c] = fma16q(TAP[at + c], wg[k], f[c], mid);
}
if (mid) { // an f32 result on a binary16 midpoint: redo the chain with the exact fma (rare)
f[0] = 0.0hf; f[1] = 0.0hf;
for (uint k = 0u; k < 9u; ++k) {
const uint at = ((ty + k / 3u) * 10u + tx + k % 3u) * 4u + c0;
for (uint c = 0u; c < 2u; ++c) f[c] = fma16(TAP[at + c], wg[k], f[c]);
}
}
const float16_t g = sig(a[9]);
precise float16_t om = float16_t(1.0) - g;
const vec4 t96 = texelFetch(T96, ivec2(x, y), 0), t104 = texelFetch(T104, ivec2(x, y), 0);
float16_t o[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
precise float16_t fg = f[c] * g;
float16_t v = fma16(float16_t(t96[c0 + c]), om, fg);
const float16_t t4 = float16_t(t104[c0 + c]);
if (!isnan(t4)) { precise float16_t vv = v + t4; v = vv; }
o[c] = v;
}
const uint other = subgroupShuffleXor(bh(o[0]), 1u);
if (hf == 0u) imageStore(S112, ivec2(x, y), vec4(o[0], o[1], hb(other), 0.0));
}
engine/rr/rr_exact.glsl
+185
View File
@@ -0,0 +1,185 @@
// Exact NVIDIA arithmetic helpers for the native Ray Reconstruction stages (engine/rr/*_m.comp, enc1_native.comp).
// Requires GL_EXT_shader_explicit_arithmetic_types, GL_EXT_float_e4m3 (RR_GFX11: GL_KHR_shader_subgroup_shuffle instead;
// FP8 values are then codes, rr_gfx11.glsl).
//
// RADV folds a register-only float(float16_t(x)) back to x, so every binary16 rounding whose result is used again goes
// through a conversion with an opaque zero (RR_ZERO16, set in main from something the compiler cannot evaluate, e.g.
// gl_NumWorkGroups.z - 1) OR-ed into its bits: one instruction more than the plain conversion.
uint16_t RR_ZERO16;
#if defined(RR_RELAXED)
// Relaxed rounding (build_native_model.sh with D4R_RR_RELAXED=1): the accumulators stay f32 between 32-deep steps and
// are converted to FP8 in one rounding instead of NVIDIA's two (binary16, then FP8). About 1% of the FP8 codes of a
// layer then differ by one step, so a stage no longer reproduces the reference bit for bit; it saves the conversions
// that are a quarter of a stage's instructions.
float rnd16(float x) { return x; }
#elif defined(RR_EXACT_CONVERT)
// the module's conversions are decorated NoContraction (engine/rr/spv_exact_convert.py): the driver keeps them
float rnd16(float x) { return float(float16_t(x)); }
#else
float rnd16(float x) { return float(uint16BitsToFloat16(float16BitsToUint16(float16_t(x)) | RR_ZERO16)); }
#endif
// x rounded to binary16 (nearest even), as f32 / as f16, through a conversion the compiler keeps (see rnd16)
float r16(float x) { return rnd16(x); }
#if (defined(RR_RELAXED) && RR_RELAXED >= 2) || defined(RR_EXACT_CONVERT)
float16_t h16(float x) { return float16_t(x); }
#else
float16_t h16(float x) { return uint16BitsToFloat16(float16BitsToUint16(float16_t(x)) | RR_ZERO16); }
#endif
// NVIDIA's k order within a 32-block: slot j holds natural channel inter(j); uninter is the inverse (and also the
// arena's channel order: arena slot uninter(k) holds natural channel k)
uint inter(uint j) { return (j & ~31u) | (((j & 31u) >> 2u) << 1u) | (j & 1u) | (((j >> 1u) & 1u) << 4u); }
uint uninter(uint c) { return (c & ~31u) | (((c & 15u) >> 1u) << 2u) | (c & 1u) | (((c >> 4u) & 1u) << 1u); }
int mirror(int v, int n) { v = abs(v); return min(v, 2 * n - 2 - v); }
#ifdef RR_GFX11
#include "rr_gfx11.glsl"
#define F8 uint
uint q8h(float16_t x) { return g11_code(x); }
uint q8s(float x) { return g11_codef(x); }
uint code(uint c) { return c; }
float16_t e4h(uint c) { return g11_e4h(c); }
// q8(act(h)) from the packed table (enc1_pack.act_table; the weight buffer's W8U view at W_ACT)
#define ACTQ(h) uint(w8u.w[W_ACT + bh(h)])
// an accumulator's activation codes (accumulator layout)
#define ACTQ8C(dst, h) { (dst) = uvec2(0u); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (dst)[i_ >> 2] |= ACTQ(float16_t((h)[i_])) << (8u * uint(i_ & 3)); }
#else
#define F8 floate4m3_t
float16_t e4h(floate4m3_t v) { return float16_t(v); }
// cvt.rn.satfinite.e4m3 (NaN -> 0x7f)
floate4m3_t q8(float16_t h)
{
return isnan(h) ? uintBitsToFloate4m3EXT(uint8_t(0x7fu)) : floate4m3_t(clamp(h, float16_t(-448.0), float16_t(448.0)));
}
// the same from f32 / f16 values in one saturating hardware conversion (satfinite; for an f32 value that is a binary16
// number the result equals the f16 conversion's)
floate4m3_t q8s(float x) { floate4m3_t r; saturatedConvertEXT(r, x); return r; }
floate4m3_t q8h(float16_t x) { floate4m3_t r; saturatedConvertEXT(r, x); return r; }
uint code(floate4m3_t v) { return uint(floate4m3BitsToUintEXT(v)); }
#endif
float16_t hb(uint b) { return uint16BitsToFloat16(uint16_t(b)); }
uint bh(float16_t h) { return uint(float16BitsToUint16(h)); }
// m * x compared with 1, exactly (the fma recovers the product's rounding error): -1, 0 or 1
int cmp1(float m, float x)
{
precise float t = m * x;
precise float e = fma(m, x, -t);
return t > 1.0 ? 1 : (t < 1.0 ? -1 : (e > 0.0 ? 1 : (e < 0.0 ? -1 : 0)));
}
// 1 / x of the filters' softmax and gate, rounded to binary16. The reference is the GPU's own reciprocal instruction
// (rcp.approx; ZLUDA lowers it to llvm.amdgcn.rcp.f32, which is what 1.0 / x compiles to here), not the correctly
// rounded quotient: correcting it left ~0.01% of the filtered surfaces one ulp off (capDT, 2026-10-12).
float16_t rcp16(float x) { return h16(1.0 / x); }
// binary16 rounding of 1 / sqrt(s) (s > 0)
float16_t rsq16(float16_t s16)
{
const float s = float(s16);
float16_t q = h16(inversesqrt(s));
#if defined(RR_RELAXED) && RR_RELAXED >= 2
return q;
#endif
for (int i = 0; i < 2; ++i) {
const uint b = bh(q);
const float lo = (float(q) + float(hb(b - 1u))) * 0.5, hi = (float(q) + float(hb(b + 1u))) * 0.5;
const int cl = cmp1(lo * lo, s), ch = cmp1(hi * hi, s);
if (cl > 0 || (cl == 0 && (b & 1u) != 0u)) q = hb(b - 1u);
else if (ch < 0 || (ch == 0 && (b & 1u) != 0u)) q = hb(b + 1u);
}
return q;
}
// fma of binary16 operands with one binary16 rounding, exactly. The product is exact in f32, so an f32 fma rounds the
// exact result once; converting that to binary16 rounds twice, which differs only when the f32 value lands exactly
// on a binary16 midpoint: then the sum's rounding error (TwoSum, the rare branch) decides the direction.
float16_t fma16(float16_t a, float16_t b, float16_t c)
{
const float fa = float(a), fb = float(b), fc = float(c);
precise float s = fma(fa, fb, fc);
#if defined(RR_RELAXED) && RR_RELAXED >= 2
return float16_t(s);
#endif
const uint u = floatBitsToUint(s);
if ((u & 0x1fffu) == 0x1000u && abs(s) >= 6.103515625e-05) {
precise float p = fa * fb;
precise float bb = s - p;
precise float err = (p - (s - bb)) + (fc - bb);
if (err != 0.0) return h16(uintBitsToFloat(err * s > 0.0 ? u + 1u : u - 1u));
}
return h16(s);
}
// the same in one f32 fma, rounded once more to binary16: mid is set when the f32 result lies on a binary16 midpoint,
// where that double rounding can differ (the caller then redoes the computation with fma16)
float16_t fma16q(float16_t a, float16_t b, float16_t c, inout bool mid)
{
precise float s = fma(float(a), float(b), float(c));
#if defined(RR_RELAXED) && RR_RELAXED >= 2
return float16_t(s);
#endif
mid = mid || (floatBitsToUint(s) & 0x1fffu) == 0x1000u;
#ifdef RR_EXACT_CONVERT
return float16_t(s);
#else
return hb(bh(float16_t(s)) | uint(RR_ZERO16));
#endif
}
#ifndef RR_GFX11
// The saturating FP8 conversion of a whole accumulator, two elements per conversion instruction
#define Q8PAIRS(dst, src) [[unroll]] for (int i_ = 0; i_ < 8; i_ += 2) { fe4m3vec2 t_; saturatedConvertEXT(t_, vec2((src)[i_], (src)[i_ + 1])); (dst)[i_] = t_.x; (dst)[i_ + 1] = t_.y; }
// q8(act(h)) from the packed table (enc1_pack.act_table; the weight buffer's W8U view at W_ACT)
#define ACTQ(h) uintBitsToFloate4m3EXT(w8u.w[W_ACT + bh(h)])
// the activation of a whole accumulator into FP8. RR_NATIVE (D4R_RR_RELAXED=3): in f32, two elements per
// instruction, straight into FP8 - no binary16 in between and no table (not the reference's roundings)
#ifdef RR_NATIVE
#define ACTQ8(dst, h) [[unroll]] for (int i_ = 0; i_ < 8; i_ += 2) { const vec2 x_ = vec2((h)[i_], (h)[i_ + 1]); const vec2 y_ = clamp(x_, vec2(-2.0), vec2(2.0)); const vec2 g_ = fma(y_, fma(-abs(y_), vec2(0.0810546875), vec2(0.412109375)), vec2(0.5)); fe4m3vec2 t_; saturatedConvertEXT(t_, x_ * g_); (dst)[i_] = t_.x; (dst)[i_ + 1] = t_.y; }
#else
#define ACTQ8(dst, h) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (dst)[i_] = ACTQ(float16_t((h)[i_]));
#endif
#endif
float16_t act(float16_t x)
{
precise float16_t y = clamp(x, float16_t(-2.0), float16_t(2.0));
precise float16_t a = float16_t(0.0810546875) * abs(y);
precise float16_t b = float16_t(0.412109375) - a;
precise float16_t c = y * b;
precise float16_t g = float16_t(0.5) + c;
precise float16_t r = x * g;
return r;
}
float16_t sig(float16_t v) { return rcp16(exp2(float(v) * -1.4426950216293335) + 1.0); }
// the cubic-exponent softmax numerator (enc1_ref.softmax_cubic): a cubic in the clamped scaled score whose binary16
// bits, shifted into the exponent field, are the exponential
// the same in plain f32 fmas (double rounding); bad is set when an f32 result lies exactly on a binary16 midpoint,
// the only case where that differs: the caller then recomputes with cubic_e (subgroupAny, rare)
float16_t cubic_e_fast(float16_t x, inout bool bad)
{
precise float16_t m = clamp(x * float16_t(0.017181396484375), float16_t(-0.55615234375), float16_t(0.55615234375));
const float fm = float(m);
precise float t = fma(fm, -fm, 0.92724609375);
const float tr = rnd16(t);
precise float u = fma(fm, tr, 1.375);
// t - c and u - 1.375 are exact (Sterbenz), so one more fma gives the rounding error's sign: only an inexact
// midpoint is a double rounding
precise float et = fma(fm, fm, t - 0.92724609375), eu = fma(fm, -tr, u - 1.375);
bad = bad || ((floatBitsToUint(t) & 0x1fffu) == 0x1000u && et != 0.0) || ((floatBitsToUint(u) & 0x1fffu) == 0x1000u && eu != 0.0);
return hb((bh(float16_t(u)) << 5u) & 0x7FE0u);
}
float16_t cubic_e(float16_t x)
{
precise float16_t m = clamp(x * float16_t(0.017181396484375), float16_t(-0.55615234375), float16_t(0.55615234375));
const float16_t t = fma16(m, -m, float16_t(0.92724609375));
const float16_t u = fma16(m, t, float16_t(1.375));
return hb((bh(u) << 5u) & 0x7FE0u);
}
// FP8 operands from a shared byte image (word offset and stride): an A row-major, a B column-major; and an accumulator
// quantized and stored as the transposed A row-major / as a B column-major (the same bytes)
#ifdef RR_GFX11
#define LDS8A(m, ARR, off, stride) LD8A(m, ARR, off, stride)
#define LDS8B(m, ARR, off, stride) LD8B(m, ARR, off, stride)
#define Q8STA(src, ARR, off, stride) { F8V c_; G11_CODES(c_, src) G11_ST(c_, ARR, off, stride) }
#define Q8STB(src, ARR, off, stride) Q8STA(src, ARR, off, stride)
#else
#define LDS8A(m, ARR, off, stride) coopMatLoad(m, ARR, off, stride, gl_CooperativeMatrixLayoutRowMajor)
#define LDS8B(m, ARR, off, stride) coopMatLoad(m, ARR, off, stride, gl_CooperativeMatrixLayoutColumnMajor)
#define Q8STA(src, ARR, off, stride) { A8 q; Q8(q, src) coopMatStore(q, ARR, off, stride, gl_CooperativeMatrixLayoutRowMajor); }
#define Q8STB(src, ARR, off, stride) { B8 q; Q8(q, src) coopMatStore(q, ARR, off, stride, gl_CooperativeMatrixLayoutColumnMajor); }
#endif
engine/rr/rr_gfx11.glsl
+95
View File
@@ -0,0 +1,95 @@
// RDNA3 (gfx11, RR_GFX11) forms of the native Ray Reconstruction stages' FP8 operations, included by rr_exact.glsl.
// gfx11 has no FP8 matrix instruction and RADV no VK_EXT_shader_float8, so the reference is ZLUDA's gfx11 lowering of
// mma.m16n8k32.f16.e4m3.e4m3.f16 (D4R_ZLUDA_WMMA_FP8, patches/zluda 0004), which these reproduce bit for bit:
// - each e4m3 byte is widened on its bits to the binary16 of value * 2^-8 (magnitude at bits 7..14, the sign at 15,
// the NaN code to a NaN), so the products carry 2^-16;
// - the accumulator enters scaled by 2^-16, the two K halves (k 0..15, 16..31) run as two chained
// v_wmma_f32_16x16x16_f16, the sum is scaled back and rounded to binary16 once;
// - inside each K half the operand positions hold k in ZLUDA's order: position 2q + j holds k 4q + j and position
// 8 + 2q + j holds k 4q + 2 + j (q < 4, j < 2). The instruction's f32 sum depends on that order (tests on an
// RX 7700 XT: 29% of the outputs change when the same products are permuted), so it is part of the arithmetic.
// FP8 values in registers are their codes (uint); an operand matrix is coopmat<float16_t> holding the widened values.
// RADV's gfx11 layouts: an A or B operand lane l holds row / column l & 15, element i = k i (both half-waves alike);
// an accumulator lane holds column l & 15, element i = row 2i + (l >> 4). The stages' elementwise code is written for
// gfx12's accumulator layout (element i = row i + 8 (l >> 4)), so every A operand is loaded with its rows permuted:
// lane l takes row g11_row(l & 15), and a product's physical row 2i + h is then row i + 8h, i.e. the gfx12 layout.
// That only relabels the outputs (each is its own dot product), so it is exact. Accumulators loaded or built
// elementwise therefore use the gfx12 positions too.
#define A8 coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B8 coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
// cvt.rn.satfinite.e4m3 of binary16 bits as ZLUDA computes it on gfx11 (HIP's conversion after clamping infinities
// to 448): nearest even, |x| > 448 -> 0x7e, NaN -> 0x7f, the sign kept. Equal to it for all 65536 inputs.
#ifndef G11_HAVE_H2E4M3 // ptx2rr.py's preamble (a translated head before this file) defines the same
uint g11_h2e4m3(uint h)
{
const uint s = (h >> 8u) & 0x80u, a = h & 0x7fffu;
if (a > 0x7c00u) return s | 0x7fu;
uint c;
if (a >= 0x2400u) { // |x| >= 2^-6: an e4m3 normal; rebias 15 -> 7 and round 10 -> 3 mantissa bits
const uint t = a - 0x2000u;
c = min((t + 0x3fu + ((t >> 7u) & 1u)) >> 7u, 0x7eu);
} else { // e4m3 subnormals: steps of 2^-9 (x * 2^9 is exact in f32)
c = uint(roundEven(float(uint16BitsToFloat16(uint16_t(a))) * 512.0));
}
return s | c;
}
#endif
uint g11_code(float16_t h) { return g11_h2e4m3(uint(float16BitsToUint16(h))); }
// an f32 that holds a binary16 value
uint g11_codef(float x) { return g11_code(float16_t(x)); }
// the exact binary16 value of a code (ZLUDA's e4m3_to_f16_bits; NaN -> 0x7fff with the sign)
float16_t g11_e4h(uint c)
{
const uint m = c & 0x7fu;
uint h = m >= 8u ? (m << 7u) + 0x2000u : uint(float16BitsToUint16(float16_t(float(m & 7u) * 0.001953125)));
h = m == 0x7fu ? 0x7fffu : h;
return uint16BitsToFloat16(uint16_t(((c & 0x80u) << 8u) | h));
}
#ifndef G11_HAVE_FILL
// four codes -> the scaled binary16 pairs (bytes 0, 1) and (bytes 2, 3), ZLUDA's e4m3x4_to_scaled_f16x2x2
uvec2 g11_sc4(uint packed)
{
uint mag = packed & 0x7f7f7f7fu;
mag |= (mag + 0x01010101u) & 0x80808080u;
const uint sign = packed & 0x80808080u;
const uint lo = ((mag & 0xffu) << 7u) | ((mag & 0xff00u) << 15u) | ((sign & 0x80u) << 8u) | ((sign & 0x8000u) << 16u);
const uint hi = (((mag >> 16u) & 0xffu) << 7u) | ((mag >> 24u) << 23u) | ((sign & 0x800000u) >> 8u) | (sign & 0x80000000u);
return uvec2(lo, hi);
}
// an operand from its lane's 16 codes in k order (four words): positions in ZLUDA's order
#define G11_FILL(m, w0, w1, w2, w3) { \
const uvec2 s0_ = g11_sc4(w0), s1_ = g11_sc4(w1), s2_ = g11_sc4(w2), s3_ = g11_sc4(w3); f16vec2 l_; \
l_ = unpackFloat2x16(s0_.x); (m)[0] = l_.x; (m)[1] = l_.y; l_ = unpackFloat2x16(s0_.y); (m)[8] = l_.x; (m)[9] = l_.y; \
l_ = unpackFloat2x16(s1_.x); (m)[2] = l_.x; (m)[3] = l_.y; l_ = unpackFloat2x16(s1_.y); (m)[10] = l_.x; (m)[11] = l_.y; \
l_ = unpackFloat2x16(s2_.x); (m)[4] = l_.x; (m)[5] = l_.y; l_ = unpackFloat2x16(s2_.y); (m)[12] = l_.x; (m)[13] = l_.y; \
l_ = unpackFloat2x16(s3_.x); (m)[6] = l_.x; (m)[7] = l_.y; l_ = unpackFloat2x16(s3_.y); (m)[14] = l_.x; (m)[15] = l_.y; }
#endif
// the A operand row a lane holds: physical row r carries row (r >> 1) + 8 (r & 1)
uint g11_row(uint r) { return (r >> 1u) | ((r & 1u) << 3u); }
// an A operand row-major (rows permuted) or a B operand column-major from a shared byte image (word offset and
// stride, as coopMatLoad takes them for an FP8 matrix): the lane's row / column is 16 consecutive bytes
#define LD8A(m, ARR, off, stride) { const uint b_ = (off) + g11_row(gl_SubgroupInvocationID & 15u) * (stride); \
G11_FILL(m, ARR[b_], ARR[b_ + 1u], ARR[b_ + 2u], ARR[b_ + 3u]) }
#define LD8B(m, ARR, off, stride) { const uint b_ = (off) + (gl_SubgroupInvocationID & 15u) * (stride); \
G11_FILL(m, ARR[b_], ARR[b_ + 1u], ARR[b_ + 2u], ARR[b_ + 3u]) }
// natural k order (a matrix loaded from the widened weights) -> ZLUDA's positions
#define G11_PERM(d, s) { (d)[0] = (s)[0]; (d)[1] = (s)[1]; (d)[2] = (s)[4]; (d)[3] = (s)[5]; (d)[4] = (s)[8]; \
(d)[5] = (s)[9]; (d)[6] = (s)[12]; (d)[7] = (s)[13]; (d)[8] = (s)[2]; (d)[9] = (s)[3]; (d)[10] = (s)[6]; \
(d)[11] = (s)[7]; (d)[12] = (s)[10]; (d)[13] = (s)[11]; (d)[14] = (s)[14]; (d)[15] = (s)[15]; }
// FP8 vectors in the accumulator layout: element i of a lane (row i + 8 half) as byte i & 3 of word i >> 2
#define F8V uvec2
// the codes of an accumulator that holds binary16 values
#define G11_CODES(dst, src) { (dst) = uvec2(0u); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) \
(dst)[i_ >> 2] |= g11_codef((src)[i_]) << (8u * uint(i_ & 3)); }
// the B operand that is the accumulator-layout codes c (the gfx12 elementwise conversion of an accumulator into B):
// rows 0..7 are half-wave 0's codes, 8..15 half-wave 1's
#define G11_B(dst, c) { const uvec2 o_ = uvec2(subgroupShuffleXor((c).x, 16u), subgroupShuffleXor((c).y, 16u)); \
const bool h_ = gl_SubgroupInvocationID >= 16u; const uvec2 e0_ = h_ ? o_ : (c), e1_ = h_ ? (c) : o_; \
G11_FILL(dst, e0_.x, e0_.y, e1_.x, e1_.y) }
// store the codes c column-major as bytes (column j at word off + j * stride): what an FP8 B stored column-major or
// an FP8 A (the accumulator's transpose) stored row-major leaves; each half-wave writes its 8 rows
#define G11_ST(c, ARR, off, stride) { const uint b_ = (off) + (gl_SubgroupInvocationID & 15u) * (stride) + \
(gl_SubgroupInvocationID >= 16u ? 2u : 0u); ARR[b_] = (c).x; ARR[b_ + 1u] = (c).y; }
engine/rr/rr_gfx11_k32.glsl
+29
View File
@@ -0,0 +1,29 @@
// RR_GFX11: the weight loads and the 32-deep step of the native stages (rr_gfx11.glsl explains the arithmetic).
// Included after ACC / HAC / RM / CM and the binding-1 views. W_WIDE (a -D of the build) is the f16 element at which
// the weight file carries its widened copy: element W_WIDE + j is byte j of the image widened like an FP8 operand
// (value * 2^-8), so a stage's byte offsets address it unchanged.
layout(set = 0, binding = 1, std430) readonly buffer W128 { uvec4 w[]; } w128;
#ifndef W_WIDE
#error RR_GFX11 needs W_WIDE (the widened weight copy, engine/rr/widen_weights.py)
#endif
// A: the lane's (permuted) row is 16 consecutive f16, two 16-byte loads (base and tiles are multiples of 16 bytes)
#define LDA(dst, base, k, mt, kt) { const uint e_ = W_WIDE + (base) + 256u * ((mt) * ((k) / 16u) + (kt)) + \
16u * g11_row(gl_SubgroupInvocationID & 15u); const uvec4 u0_ = w128.w[e_ >> 3u], u1_ = w128.w[(e_ >> 3u) + 1u]; \
f16vec2 l_; l_ = unpackFloat2x16(u0_.x); (dst)[0] = l_.x; (dst)[1] = l_.y; l_ = unpackFloat2x16(u0_.y); (dst)[8] = l_.x; (dst)[9] = l_.y; \
l_ = unpackFloat2x16(u0_.z); (dst)[2] = l_.x; (dst)[3] = l_.y; l_ = unpackFloat2x16(u0_.w); (dst)[10] = l_.x; (dst)[11] = l_.y; \
l_ = unpackFloat2x16(u1_.x); (dst)[4] = l_.x; (dst)[5] = l_.y; l_ = unpackFloat2x16(u1_.y); (dst)[12] = l_.x; (dst)[13] = l_.y; \
l_ = unpackFloat2x16(u1_.z); (dst)[6] = l_.x; (dst)[7] = l_.y; l_ = unpackFloat2x16(u1_.w); (dst)[14] = l_.x; (dst)[15] = l_.y; }
#define LDB(dst, base, n, kt, nt) { B8 t_; coopMatLoad(t_, w16.w, W_WIDE + (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM); G11_PERM(dst, t_) }
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
#if defined(RR_MIX) || defined(RR_RELAXED)
#error RR_GFX11 builds the exact arithmetic only
#endif
#define RAC ACC
#define RV(x) float(x)
// C * 2^-16 (exact), the two halves, * 2^16 (exact), one rounding to binary16
#define K32(acc, a0, b0, a1, b1) { [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] *= 1.52587890625e-05; \
(acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); \
[[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = rnd16((acc)[i_] * 65536.0); }
#define K32H(acc, a0, b0, a1, b1) { ACC p_ = ACC(acc); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) p_[i_] *= 1.52587890625e-05; \
p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), p_)); \
[[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_] * 65536.0); }
engine/rr/rr_inputs.py
+111
View File
@@ -0,0 +1,111 @@
#!/usr/bin/env python3
"""Write the game inputs of a capture's denoiser frames for engine/replay_rr.cpp (d4r-rr).
usage: rr_inputs.py CAPTURE OUT_DIR [--frames N] [--no-reset]
The inputs are enc0's textures (colour 160, motion 168, normals 184, roughness 192, the albedos 232 and 240, depth
512); the frame settings come from enc0's argument block: jitter (320), motion scale (344, in output pixels per
render pixel), depth inversion (528) and the camera, whose rotation-times-projection (596) and projection (660) give
back a world-to-view rotation and the view-to-clip matrix that reproduce the block. The harness passes reset on every
frame (scripts/run_rr_engine_compare.sh); --no-reset clears it after the first frame.
"""
import argparse
import os
import struct
import sys
from pathlib import Path
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rr_capture import Capture # noqa: E402
OPTIONAL = {'hitdistance': 248, 'disocclusion': 296, 'sss': 264} # encoder 0 textures
OPTIONAL_OUTPUT = {'particles': 416} # output stage textures
INPUTS = {'color': 160, 'motion': 168, 'normals': 184, 'roughness': 192, 'diffuse': 232, 'specular': 240, 'depth': 512}
# preset F (rr2_enc0_kernel): its textures and argument offsets
INPUTS_F = {'color': 144, 'motion': 192, 'normals': 224, 'roughness': 240, 'diffuse': 160, 'specular': 176, 'depth': 552}
def preset_f(cap, o):
lines = []
for f in range(1, o.frames + 1):
enc0 = cap.find('rr2_enc0', f)
a = enc0.args
u32 = lambda off: struct.unpack_from('<I', a, off)[0]
f32 = lambda off: struct.unpack_from('<f', a, off)[0]
mat = lambda off: np.array(struct.unpack_from('<16f', a, off), np.float64).reshape(4, 4)
down = cap.find('rr2_downsample', 1).args # either variant: the sizes are at the same offsets
rw, rh, iw, ih = u32(128), u32(132), u32(488), u32(492)
ow, oh = struct.unpack_from('<2I', down, 8)
if f == 1:
lines.append(f'size {rw} {rh} {ow} {oh}')
images = {i.offset: i for i in enc0.images}
for name, off in INPUTS_F.items():
(o.out / f'{f}-{name}.bin').write_bytes(images[off].path.read_bytes())
# the optional inputs: encoder 0's textures 272 / 336 / 288, the particle downsample kernel's texture 208
for name, off in (('hitdistance', 272), ('disocclusion', 336), ('sss', 288)):
if off in images:
(o.out / f'{f}-{name}.bin').write_bytes(images[off].path.read_bytes())
part = [l for l in cap.launches if l.kernel == 'rr2_downsample_kernel_static_hdr']
if part:
n = len(part) // len(cap.find('rr2_enc0'))
(o.out / f'{f}-particles.bin').write_bytes(part[(f - 1) * n].image(208).path.read_bytes())
clip = mat(708)
view = mat(644) @ np.linalg.inv(clip)
view[3, :3] = 0.0
reset = 1 if f == 1 or not o.no_reset else 0
values = ' '.join(f'{v:.9g}' for v in list(view.flatten()) + list(clip.flatten()))
flags = a[568] | (a[1028] << 1) | (a[466] << 2) # inverted depth, hardware depth, packed roughness
lines.append(f'frame {f} {f32(400):.9g} {f32(404):.9g} {f32(448) * rw / iw:.9g} {f32(452) * rh / ih:.9g} 1 {reset} {flags} {values}')
(o.out / 'frames.txt').write_text('\n'.join(lines) + '\n')
print(f'{o.frames} preset F frames of {o.capture} in {o.out}')
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('out', type=Path)
ap.add_argument('--frames', type=int, default=3)
ap.add_argument('--no-reset', action='store_true')
o = ap.parse_args()
cap = Capture(o.capture)
o.out.mkdir(parents=True, exist_ok=True)
lines = []
if any(l.kernel == 'rr2_enc0_kernel' for l in cap.launches):
return preset_f(cap, o)
for f in range(1, o.frames + 1):
enc0 = cap.find('swin_enc0', f)
a = enc0.args
u32 = lambda off: struct.unpack_from('<I', a, off)[0]
f32 = lambda off: struct.unpack_from('<f', a, off)[0]
mat = lambda off: np.array(struct.unpack_from('<16f', a, off), np.float64).reshape(4, 4)
rw, rh, ow, oh = u32(144), u32(148), u32(472), u32(476)
if f == 1:
lines.append(f'size {rw} {rh} {ow} {oh}')
images = {i.offset: i for i in enc0.images}
for name, off in INPUTS.items():
(o.out / f'{f}-{name}.bin').write_bytes(images[off].path.read_bytes())
for name, off in OPTIONAL.items(): # DLSSD.SpecularHitDistance, DLSS.DisocclusionMask when the game passed them
if off in images:
(o.out / f'{f}-{name}.bin').write_bytes(images[off].path.read_bytes())
out_images = {i.offset: i for i in cap.find('hkpn_output', f).images}
for name, off in OPTIONAL_OUTPUT.items(): # DLSSD.ColorBeforeParticles
if off in out_images:
(o.out / f'{f}-{name}.bin').write_bytes(out_images[off].path.read_bytes())
clip = mat(660)
view = mat(596) @ np.linalg.inv(clip)
view[3, :3] = 0.0
jx, jy = f32(320), f32(324)
mvx, mvy = f32(344) * rw / ow, f32(348) * rh / oh
reset = 1 if f == 1 or not o.no_reset else 0
values = ' '.join(f'{v:.9g}' for v in list(view.flatten()) + list(clip.flatten()))
lines.append(f'frame {f} {jx:.9g} {jy:.9g} {mvx:.9g} {mvy:.9g} 1 {reset} {u32(528) | (a[980] << 1)} {values}')
(o.out / 'frames.txt').write_text('\n'.join(lines) + '\n')
print(f'{o.frames} frames of {o.capture} in {o.out}')
if __name__ == '__main__':
main()
engine/rr/rr_templates.py
+103
View File
@@ -0,0 +1,103 @@
#!/usr/bin/env python3
"""The two argument-block templates of a preset D / E model (enc0.args, hkpn.args) without a capture.
usage: rr_templates.py NVNGX_DLSSD_DLL D|E OUT_DIR
rr_templates.py NVNGX_DLSSD_DLL D|E --check MODEL_DIR... (compare with templates compile_rr.py took from captures)
engine/runtime_rr.inc fills every field of the two blocks that depends on the frame, the geometry or the bound inputs.
Across captures of different sizes, inputs and frames (2026-10-14: capD, capDT, capD4K, capDX*, capE, capE3) what is
left is the same few constants for both presets, plus two floats of encoder 0 (bytes 540 and 544) that differ between
the presets. NGX sets those two in code: `sub ecx, 4; je D; cmp ecx, 1; jne ...; mov [rbx + X], E0; mov [rbx + X + 4],
E1; jmp ...; D: mov [rbx + X], D0; mov [rbx + X + 4], D1` (ecx = the preset), which find_floats() looks up in the DLL.
The templates are written on the user's machine and are never packaged.
"""
import re
import struct
import sys
from pathlib import Path
ENC0_BYTES, HKPN_BYTES = 1032, 480
# (offset, packed value) of the constants the runtime keeps
ENC0_CONST = [(360, struct.pack('<f', 1.0)), (532, struct.pack('<I', 2)), (988, struct.pack('<I', 1))]
HKPN_CONST = [(148, struct.pack('<f', 1.0)), (436, struct.pack('<f', 0.25)), (456, struct.pack('<I', 1))]
# the byte ranges engine/runtime_rr.inc writes every frame (setFrame), for --check
ENC0_WRITTEN = [(0, 8), (48, 64), (80, 84), (136, 160), (312, 352), (364, 368), (432, 440), (472, 480), (528, 532),
(536, 537), (539, 540), (596, 980), (980, 981), (984, 988), (1024, 1032)]
HKPN_WRITTEN = [(0, 48), (56, 64), (140, 148), (152, 156), (160, 161), (256, 264), (336, 344), (432, 433), (440, 444)]
# mov [rbx + disp32], imm32 is c7 83 disp32 imm32
PATTERN = re.compile(rb'\x83\xe9\x04\x74.\x83\xf9\x01\x75.'
rb'\xc7\x83(....)(....)\xc7\x83(....)(....)\xeb.'
rb'\xc7\x83(....)(....)\xc7\x83(....)(....)', re.S)
def find_floats(dll):
"""{'D': (f0, f1), 'E': (f0, f1)}: encoder 0's bytes 540 / 544 as NGX's code sets them"""
found = set()
for m in PATTERN.finditer(dll):
e0, e0v, e1, e1v, d0, d0v, d1, d1v = m.groups()
x = struct.unpack('<I', e0)[0]
if struct.unpack('<I', e1)[0] == x + 4 and e0 == d0 and e1 == d1:
found.add((d0v + d1v, e0v + e1v))
if len(found) != 1:
raise SystemExit(f'found {len(found)} candidates for the preset constants of encoder 0 (expected one): '
'not a DLL layout this builder knows')
d, e = found.pop()
return {'D': d, 'E': e}
def templates(dll, preset):
enc0, hkpn = bytearray(ENC0_BYTES), bytearray(HKPN_BYTES)
for off, v in ENC0_CONST:
enc0[off:off + len(v)] = v
for off, v in HKPN_CONST:
hkpn[off:off + len(v)] = v
enc0[540:548] = find_floats(dll)[preset]
return bytes(enc0), bytes(hkpn)
def masked(b, written, handles):
b = bytearray(b)
for lo, hi in written:
b[lo:hi] = bytes(hi - lo)
for off in handles:
b[off:off + 8] = bytes(8)
return bytes(b)
def handles(rr_txt, stage):
for line in Path(rr_txt).read_text().splitlines():
f = line.split()
if f[:2] == ['stage', stage]:
tex = [] if f[5] == '-' else [int(x) for x in f[5].split(',')]
surf = [] if f[6] == '-' else [int(x.split(':')[0]) for x in f[6].split(',')]
return tex + surf
return []
def main():
if len(sys.argv) < 4:
raise SystemExit(__doc__)
dll, preset = Path(sys.argv[1]).read_bytes(), sys.argv[2].upper()
if preset not in ('D', 'E'):
raise SystemExit('preset D or E')
enc0, hkpn = templates(dll, preset)
if sys.argv[3] == '--check':
bad = 0
for model in sys.argv[4:]:
for name, ours, written in (('enc0', enc0, ENC0_WRITTEN), ('hkpn', hkpn, HKPN_WRITTEN)):
theirs = (Path(model) / f'{name}.args').read_bytes()
h = handles(Path(model) / 'rr.txt', name)
a, b = masked(ours, written, h), masked(theirs, written, h)
diff = [i for i in range(len(a)) if a[i] != b[i]]
bad += bool(diff)
print(f'{model} {name}: ' + ('same' if not diff else f'{len(diff)} bytes differ, first at {diff[:8]}'))
sys.exit(1 if bad else 0)
out = Path(sys.argv[3])
out.mkdir(parents=True, exist_ok=True)
(out / 'enc0.args').write_bytes(enc0)
(out / 'hkpn.args').write_bytes(hkpn)
if __name__ == '__main__':
main()
engine/rr/rrmath.py
+55
View File
@@ -0,0 +1,55 @@
"""Shared numeric helpers for the Ray Reconstruction port: e4m3 codes, f16, and ZLUDA's m16n8k32 accumulation."""
import numpy as np
_E4M3 = np.zeros(256)
for _c in range(256):
_s, _e, _m = _c >> 7, (_c >> 3) & 15, _c & 7
_v = np.nan if (_e == 15 and _m == 7) else (_m * 2.0 ** -9 if _e == 0 else (1 + _m / 8) * 2.0 ** (_e - 7))
_E4M3[_c] = -_v if _s else _v
def e4m3(codes):
return _E4M3[np.asarray(codes, np.uint8)]
def f16(bits):
return np.asarray(bits, np.uint16).view(np.float16).astype(np.float64)
def f16bits(x):
return np.asarray(x, np.float64).astype(np.float16).view(np.uint16)
def mma_k32(a_codes, b_codes, c_f16bits, ksteps=None):
"""D = A.B + C with ZLUDA's lowering: per 32-deep k-step, f32 sum of the exact products and C, one f16 rounding;
the next k-step takes that f16 value as C. (f64 stands in for f32: the products and short sums are exact.)"""
a, b = e4m3(a_codes), e4m3(b_codes)
acc = f16(c_f16bits)
for k in range(0, a.shape[1], 32):
acc = (a[:, k:k + 32] @ b[k:k + 32, :] + acc).astype(np.float32).astype(np.float16).astype(np.float64)
return f16bits(acc)
# cvt.rn.satfinite.e4m3 of every f16 bit pattern. Generate this deterministic
# format table locally; never ship an extracted or generated binary table.
def _f16_to_e4m3_table():
h = np.arange(65536, dtype=np.uint32)
mag = h & 0x7fff
normal = (mag - 0x2000 + 0x3f + ((mag >> 7) & 1)) >> 7
exponent = mag >> 10
significand = (mag & 0x3ff) | np.where(exponent != 0, 0x400, 0)
clamped_exponent = np.clip(exponent, 1, 8)
shift = 16 - clamped_exponent
subnormal = (significand + (np.left_shift(1, shift - 1)) - 1 + ((significand >> shift) & 1)) >> shift
code = np.where(mag >= 0x2400, normal, subnormal)
code = np.where(mag > 0x5f00, 0x7e, code)
result = ((h >> 8) & 0x80) | code
result = np.where(mag > 0x7c00, 0x7f, result) # All NaNs map to the canonical positive NaN.
return result.astype(np.uint8)
_Q8 = _f16_to_e4m3_table()
def q8(x_f16):
"""f16 values (float array) -> e4m3 codes"""
return _Q8[np.asarray(x_f16, np.float64).astype(np.float16).view(np.uint16)]
engine/rr/spv_exact_convert.py
+41
View File
@@ -0,0 +1,41 @@
#!/usr/bin/env python3
"""Decorate every float conversion of a SPIR-V module NoContraction, so the driver keeps it.
usage: spv_exact_convert.py IN.spv [OUT.spv]
RADV folds float(float16_t(x)) back to x unless the conversions are marked exact; glslang's `precise` does not reach
conversions. With the decoration a binary16 rounding is two instructions (engine/rr/rr_exact.glsl, RR_EXACT_CONVERT)
instead of three with an opaque zero OR-ed into the bits.
"""
import struct
import sys
OP_DECORATE, OP_FCONVERT, NO_CONTRACTION = 71, 115, 42
def main():
src = sys.argv[1]
dst = sys.argv[2] if len(sys.argv) > 2 else src
data = open(src, 'rb').read()
words = list(struct.unpack(f'<{len(data) // 4}I', data))
if words[0] != 0x07230203:
sys.exit(f'{src}: not SPIR-V')
insts, i = [], 5
while i < len(words):
count = words[i] >> 16
insts.append((words[i] & 0xffff, words[i:i + count]))
i += count
ids = [w[2] for op, w in insts if op == OP_FCONVERT]
first_type = next(k for k, (op, _) in enumerate(insts) if 19 <= op <= 39) # annotations end where types begin
out = list(words[:5])
for k, (op, w) in enumerate(insts):
if k == first_type:
for result in ids:
out += [3 << 16 | OP_DECORATE, result, NO_CONTRACTION]
out += w
open(dst, 'wb').write(struct.pack(f'<{len(out)}I', *out))
print(f'{dst}: {len(ids)} conversions marked exact')
if __name__ == '__main__':
main()
engine/rr/spv_float_controls.py
+71
View File
@@ -0,0 +1,71 @@
#!/usr/bin/env python3
"""Give a compute SPIR-V module PTX's floating-point environment: RoundingModeRTE for 16- and 32-bit results and
DenormPreserve for 16-bit values (SPV_KHR_float_controls, core in SPIR-V 1.4).
usage: spv_float_controls.py IN.spv [OUT.spv]
Without these execution modes the rounding of float conversions is the driver's choice, where every cvt.rn of
NVIDIA's kernels rounds to nearest even. glslang keeps only one execution mode per mode id, so the modes are inserted
into the compiled module instead.
"""
import struct
import sys
OP_EXTENSION, OP_CAPABILITY, OP_ENTRY_POINT, OP_EXECUTION_MODE = 10, 17, 15, 16
CAP_DENORM_PRESERVE, CAP_ROUNDING_RTE = 4464, 4467
MODES = ((4462, 16), (4462, 32), (4459, 16)) # RoundingModeRTE 16/32, DenormPreserve 16
EXTENSION = 'SPV_KHR_float_controls'
def instructions(words):
i = 5
while i < len(words):
count = words[i] >> 16
yield words[i] & 0xffff, words[i:i + count]
i += count
def string_words(text):
data = text.encode() + b'\0'
data += b'\0' * (-len(data) % 4)
return list(struct.unpack(f'<{len(data) // 4}I', data))
def patch(words):
insts = list(instructions(words))
caps = {w[1] for op, w in insts if op == OP_CAPABILITY}
has_ext = any(op == OP_EXTENSION and EXTENSION.encode() in struct.pack(f'<{len(w) - 1}I', *w[1:]) for op, w in insts)
entry = next(w[2] for op, w in insts if op == OP_ENTRY_POINT)
modes = {(w[2], w[3]) for op, w in insts if op == OP_EXECUTION_MODE and w[1] == entry and len(w) > 3}
last_cap = max(i for i, (op, _) in enumerate(insts) if op == OP_CAPABILITY)
first_mode = min(i for i, (op, _) in enumerate(insts) if op == OP_EXECUTION_MODE)
out = list(words[:5])
for i, (op, w) in enumerate(insts):
if i == first_mode:
for mode, width in MODES:
if (mode, width) not in modes:
out += [4 << 16 | OP_EXECUTION_MODE, entry, mode, width]
out += w
if i == last_cap:
for cap in (CAP_DENORM_PRESERVE, CAP_ROUNDING_RTE):
if cap not in caps:
out += [2 << 16 | OP_CAPABILITY, cap]
if not has_ext:
name = string_words(EXTENSION)
out += [(1 + len(name)) << 16 | OP_EXTENSION, *name]
return out
def main():
src = sys.argv[1]
dst = sys.argv[2] if len(sys.argv) > 2 else src
data = open(src, 'rb').read()
words = list(struct.unpack(f'<{len(data) // 4}I', data))
if words[0] != 0x07230203:
sys.exit(f'{src}: not SPIR-V')
out = patch(words)
open(dst, 'wb').write(struct.pack(f'<{len(out)}I', *out))
if __name__ == '__main__':
main()
engine/rr/swin_m.comp
+537
View File
@@ -0,0 +1,537 @@
#version 460
// Native Swin stage of the DLSS-D (preset E) denoiser, register-resident, for any stage of engine/rr/swin_ref.STAGES
// (enc2-enc5, dec2-dec4; enc1_m.comp / dec1_m.comp are the hand-written C = 64 instances it generalises). The stage's
// shape and weight offsets come from the include generated by engine/rr/swin_pack.py:
// KIND 0 encoder (tokens from the arena, 2x2 merge to CX channels, enc1's gated filter; CX = 0: the bottleneck,
// ungated filter, no merge), 1 decoder (expand of the CX channels one level down + skip, dec0's filter)
// C channels (KC = C / 32 planes, MT = C / 16 tiles), NH heads of 32 channels, C / 8 MLP branches
// The arithmetic is engine/rr/swin_ref.py's; enc1_m.comp's header explains the transposed products and slot orders.
// One workgroup per 8x8-token window, four waves, wave w owning query slots 16w..16w+15.
#extension GL_KHR_cooperative_matrix : require
#extension GL_KHR_memory_scope_semantics : require
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_8bit_storage : require
#ifndef RR_GFX11
#extension GL_EXT_float_e4m3 : require
#endif
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_EXT_control_flow_attributes : require
#extension GL_GOOGLE_include_directive : require
layout(local_size_x = 128) in;
#include "swin_stage.h"
#include "rr_exact.glsl"
#define KC (C / 32)
#define MT (C / 16)
#define NB (C / 8)
layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param;
#ifndef RR_GFX11
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
#endif
layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16;
layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u;
layout(set = 0, binding = 1, std430) readonly buffer W8U { uint8_t w[]; } w8u;
layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
layout(set = 0, binding = 3) uniform sampler2D T88;
#if KIND == 1
layout(set = 0, binding = 4) uniform sampler2D T96;
layout(set = 0, binding = 5) uniform sampler2D T104;
layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S112;
#else
layout(set = 0, binding = 4, rgba16f) uniform writeonly image2D S112;
#if CX
layout(set = 0, binding = 5, rgba16f) uniform writeonly image2D S120;
layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S128;
#endif
#endif
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define HAC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#ifndef RR_GFX11
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#endif
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CM gl_CooperativeMatrixLayoutColumnMajor
#ifdef RR_GFX11
#include "rr_gfx11_k32.glsl"
#else
#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM)
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
#define Q8(dst, src) Q8PAIRS(dst, src)
// One 32-deep step onto an accumulator that holds binary16 values. The reference (and the default here) adds the
// accumulator inside the f32 products and rounds the result to binary16. RR_MIX (build_native_model.sh with
// D4R_RR_MIX=1) keeps the accumulator as binary16, sums the products from zero and takes f16(P + acc), which the
// driver compiles to one v_fma_mix per element in place of two conversions: about 5% faster, but a single rounding
// where the reference has two, so a value can come out one binary16 step away (engine/rr/README.md).
// RAC is the accumulator type, RV(x) a binary16 value as its element; K32H is the step on a binary16 matrix.
#ifdef RR_MIX
#define RAC HAC
#define RV(x) (x)
#define K32(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(0.0))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_] + float((acc)[i_])); }
#define K32H(acc, a0, b0, a1, b1) K32(acc, a0, b0, a1, b1)
#else
#define RAC ACC
#define RV(x) float(x)
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) }
#define K32H(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(acc))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_]); }
#endif
#endif
const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0;
// NVIDIA's reduction over a short list (swin_ref.pairs_tree)
#if KC == 1
#define TREE(v) (v[0])
#elif KC == 2
#define TREE(v) (v[0] + v[1])
#elif KC == 3
#define TREE(v) ((v[0] + v[1]) + v[2])
#elif KC == 4
#define TREE(v) ((v[0] + v[1]) + (v[2] + v[3]))
#else
#define TREE(v) (((v[0] + v[1]) + (v[2] + v[3])) + v[4])
#endif
#if KIND == 1
// one pool, by phase: [EX | LOW] for the expand; the prologue reads EX and writes XT over LOW (dead) and beyond;
// the attention's K and V^T reuse EX (dead after the residual)
#define EX_W (48u * C)
#define POOL_W (EX_W + (12u * CX > 24u * C ? 12u * CX : 24u * C))
shared uint POOL[POOL_W];
#define EX POOL // the expansion: 48 rows x 4 sub-positions x C FP8 channels in slot order
#define LOW_AT EX_W // the low-resolution tokens (6 x 6, padded): CX FP8 channels, arena order
#define XT_AT EX_W // the normalised key tokens: slot-major, C FP8 channels in slot order
#define XTS POOL
#define RR_KS POOL
#define RR_KS0 0u
#define RR_VS POOL
#define RR_VS0 768u
shared float16_t TAP[100 * 4]; // texture 88 bilinear at the window's 10 x 10 filter tap positions
shared float16_t OUTF[64 * 16]; // the head per window token (raster order)
#else
shared uint XT[96 * C / 4]; // the normalised key tokens: slot-major, C FP8 channels in slot order
#define XT_AT 0u
#define XTS XT
shared uint KVS[2 * 768]; // the current head's K, then V^T
#define RR_KS KVS
#define RR_KS0 0u
#define RR_VS KVS
#define RR_VS0 768u
#define MG XT // the merge rows (16 x 4C FP8) reuse XT, dead after the projections
#if CX
shared float16_t O112[64 * 4]; // surface 112 per token, for surface 128's 2x2 sums
#endif
#endif
ivec2 keyXY(uint t, int x0, int y0)
{
if (t < 64u) return ivec2(x0 + int(t & 7u), y0 + int(t >> 3u));
const int i = int(t & 7u), side = int((t - 64u) >> 3u);
return ivec2(side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8), side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i));
}
int clampTap(int v, int n) { return v < 0 ? 1 : (v >= n ? n - 2 : v); }
#ifdef RR_GFX11
uint pack4(uint a, uint b, uint c, uint d) { return a | (b << 8u) | (c << 16u) | (d << 24u); }
float16_t e4(uint w, uint b) { return g11_e4h((w >> (8u * b)) & 0xffu); }
#else
uint pack4(floate4m3_t a, floate4m3_t b, floate4m3_t c, floate4m3_t d)
{
return code(a) | (code(b) << 8u) | (code(c) << 16u) | (code(d) << 24u);
}
float16_t e4(uint w, uint b) { return float16_t(uintBitsToFloate4m3EXT(uint8_t((w >> (8u * b)) & 0xffu))); }
#endif
uvec2 mergeXY(uint r) { return r < 8u ? uvec2(r >> 1u, r & 1u) : uvec2((r - 8u) >> 1u, 2u + (r & 1u)); }
void main()
{
const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID;
const uint hl = lane >> 4u, col = lane & 15u;
RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u);
const int ex = int(param.p[0]), ey = int(param.p[1]), hx = ex / 2, hy = ey / 2;
const uint inOff = param.p[2], skOff = param.p[6], outOff = param.p[12], mrgOff = param.p[14], shift = param.p[20];
const int x0 = 8 * int(gl_WorkGroupID.x) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u);
const uint plane = uint(ex * ey) * 32u, lplane = uint(hx * hy) * 32u;
#if KIND == 1
// ---------------------------------------------------------------- expand: wave w computes sub-position w
const int X0 = (x0 >> 1) - 1, Y0 = (y0 >> 1) - 1;
for (uint i = tid; i < 48u * (CX / 4u); i += 128u) {
const uint r = i / (CX / 4u), wd = i % (CX / 4u);
uint u = 0u;
if (r < 36u) {
const int X = mirror(X0 + int(r % 6u), hx), Y = mirror(Y0 + int(r / 6u), hy);
u = arena.w[(inOff + (wd >> 3u) * lplane + uint(Y * hx + X) * 32u) / 4u + (wd & 7u)];
}
POOL[LOW_AT + i] = u;
}
barrier();
[[unroll]] for (uint n = 0u; n < 3u; ++n) {
B8 lb[CX / 16];
[[unroll]] for (uint kt = 0u; kt < CX / 16; ++kt) LDS8B(lb[kt], POOL, LOW_AT + n * 4u * CX + 4u * kt, CX / 4u);
for (uint mm = 0u; mm < MT; ++mm) {
const uint m = MT * wv + mm;
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_BE / 2u + 16u * m + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < CX / 16; kt += 2u) {
A8 w0, w1;
LDA(w0, W_WE, CX, m, kt); LDA(w1, W_WE, CX, m, kt + 1u);
K32(a, w0, lb[kt], w1, lb[kt + 1u])
}
Q8STB(a, EX, n * 16u * C + 4u * m, C)
}
}
barrier();
#endif
// ---------------------------------------------------------------- prologue: one key token per thread
if (tid < 96u) {
const ivec2 k = keyXY(inter(tid), x0, y0);
const int xx = mirror(k.x, ex), yy = mirror(k.y, ey);
const uint tb = ((KIND == 1 ? skOff : inOff) + uint(yy * ex + xx) * 32u) / 4u;
#if KIND == 1
const uint r = uint(((k.y >> 1) - Y0) * 6 + ((k.x >> 1) - X0)), s = uint(2 * (k.y & 1) + (k.x & 1));
#endif
float16_t v[C];
[[unroll]] for (uint wd = 0u; wd < C / 4; ++wd) {
const uint t = arena.w[tb + (wd >> 3u) * (plane >> 2u) + (wd & 7u)];
#if KIND == 1
const uint e = EX[r * C + s * (C / 4u) + wd];
[[unroll]] for (uint b = 0u; b < 4u; ++b) { precise float16_t x = e4(e, b) + e4(t, b); v[4u * wd + b] = x; }
#else
[[unroll]] for (uint b = 0u; b < 4u; ++b) v[4u * wd + b] = e4(t, b);
#endif
}
// swin_ref.input_norm: per lane word q (and q + 4) the pair sums of the planes (TREE), the halves, then
// q ^ 1, q ^ 2, the components
precise f16vec2 h[4];
[[unroll]] for (uint q = 0u; q < 4u; ++q) {
precise f16vec2 hh[2];
[[unroll]] for (uint half_ = 0u; half_ < 2u; ++half_) {
precise f16vec2 pr[KC];
[[unroll]] for (uint p = 0u; p < KC; ++p) {
const uint o = 32u * p + 4u * (q + 4u * half_);
precise f16vec2 a = f16vec2(v[o], v[o + 1u]) * f16vec2(v[o], v[o + 1u]);
precise f16vec2 c = f16vec2(v[o + 2u], v[o + 3u]) * f16vec2(v[o + 2u], v[o + 3u]);
pr[p] = a + c;
}
precise f16vec2 tr = TREE(pr);
hh[half_] = tr;
}
h[q] = hh[0] + hh[1];
}
precise f16vec2 h0 = h[0] + h[1], h2_ = h[2] + h[3];
precise f16vec2 hs = h0 + h2_;
precise float16_t sm = hs.x + hs.y;
precise float16_t se = sm + float16_t(0.0001220703125);
const float16_t rr = rsq16(se);
[[unroll]] for (uint wd = 0u; wd < C / 4; ++wd) {
F8 c[4];
[[unroll]] for (uint b = 0u; b < 4u; ++b) {
const uint j = 4u * wd + b;
precise float16_t g = rr * w16.w[W_GAIN0 / 2u + inter(j)];
precise float16_t n = v[j] * g;
c[b] = q8h(n);
}
XTS[XT_AT + tid * (C / 4u) + wd] = pack4(c[0], c[1], c[2], c[3]);
}
}
barrier();
// ---------------------------------------------------------------- the residual: the window token + table 7
RAC R[MT];
{
const uint tq = inter(16u * wv + col);
const int qx = x0 + int(tq & 7u), qy = y0 + int(tq >> 3u);
const uint tb = ((KIND == 1 ? skOff : inOff) + uint(mirror(qy, ey) * ex + mirror(qx, ex)) * 32u) / 4u;
#if KIND == 1
const uint r = uint(((qy >> 1) - Y0) * 6 + ((qx >> 1) - X0)), s = uint(2 * (qy & 1) + (qx & 1));
#endif
[[unroll]] for (uint m = 0u; m < MT; ++m) {
const uint ta = tb + (m >> 1u) * (plane >> 2u) + 4u * (m & 1u) + 2u * hl;
const uvec2 t = uvec2(arena.w[ta], arena.w[ta + 1u]);
#if KIND == 1
const uint ea = r * C + s * (C / 4u) + 4u * m + 2u * hl;
const uvec2 e = uvec2(EX[ea], EX[ea + 1u]);
#endif
[[unroll]] for (int i = 0; i < 8; ++i) {
#if KIND == 1
precise float16_t x = e4(e[i >> 2], uint(i) & 3u) + e4(t[i >> 2], uint(i) & 3u);
#else
const float16_t x = e4(t[i >> 2], uint(i) & 3u);
#endif
precise float16_t c = x + w16.w[W_TABLE7 / 2u + 16u * m + 8u * hl + uint(i)];
R[m][i] = RV(c);
}
}
}
// ---------------------------------------------------------------- attention, NH heads
[[dont_unroll]] for (uint hd = 0u; hd < NH; ++hd) {
barrier(); // every wave is done with the previous head's K / V (and, for the first, the expand's LOW)
for (uint j = 0u; j < 3u; ++j) { // 12 jobs (token tile n, channel tile c)
const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u;
RAC a = RAC(RV(float16_t(0.0))), v = RAC(RV(float16_t(0.0)));
[[unroll]] for (uint kt = 0u; kt < 2u * KC; kt += 2u) {
B8 xb0, xb1, wv0, wv1;
A8 xa0, xa1, wk0, wk1;
LDS8B(xb0, XTS, XT_AT + n * 4u * C + 4u * kt, C / 4u);
LDS8B(xb1, XTS, XT_AT + n * 4u * C + 4u * kt + 4u, C / 4u);
LDS8A(xa0, XTS, XT_AT + n * 4u * C + 4u * kt, C / 4u);
LDS8A(xa1, XTS, XT_AT + n * 4u * C + 4u * kt + 4u, C / 4u);
LDA(wk0, W_WK0 + HEAD * hd, C, c, kt); LDA(wk1, W_WK0 + HEAD * hd, C, c, kt + 1u);
LDB(wv0, W_WV0 + HEAD * hd, 32u, kt, c); LDB(wv1, W_WV0 + HEAD * hd, 32u, kt + 1u, c);
K32(a, wk0, xb0, wk1, xb1)
K32(v, xa0, wv0, xa1, wv1)
}
#ifdef RR_GFX11
Q8STA(a, RR_KS, RR_KS0 + n * 128u + 4u * c, 8u)
Q8STA(v, RR_VS, RR_VS0 + c * 384u + 4u * n, 24u)
#else
A8 ka, va;
Q8(ka, a)
Q8(va, v)
coopMatStore(ka, RR_KS, RR_KS0 + n * 128u + 4u * c, 8u, RM);
coopMatStore(va, RR_VS, RR_VS0 + c * 384u + 4u * n, 24u, RM);
#endif
}
barrier();
#include "rr_attn.glsl"
[[unroll]] for (uint m = 0u; m < MT; ++m) {
A8 a0, a1;
LDA(a0, W_WO0 + HEAD * hd, 32u, m, 0u); LDA(a1, W_WO0 + HEAD * hd, 32u, m, 1u);
K32(R[m], a0, Ob[0], a1, Ob[1])
}
}
// ---------------------------------------------------------------- MLP
B8 Nb[MT];
#ifdef RR_GFX11
F8V NbC[MT];
[[unroll]] for (uint m = 0u; m < MT; ++m) NbC[m] = uvec2(0u);
#endif
{
float16_t rv[MT][8];
[[unroll]] for (uint m = 0u; m < MT; ++m) [[unroll]] for (int i = 0; i < 8; ++i) rv[m][i] = float16_t(R[m][i]);
// swin_ref.rowsum_sq: slot 16m + 8hl + i is NVIDIA's channel 8j + 2t + c with j = 4(m >> 1) + (m & 1) +
// 2((i >> 1) & 1), t = 2hl + (i >> 2), c = i & 1: per parity (m & 1) the tile pairs (i, i + 2) of every
// 32-channel block, TREE over the blocks; the parities; t ^ 1 (local), t ^ 2 (the other lane half)
float16_t h1[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
float16_t ht[2];
[[unroll]] for (uint tq = 0u; tq < 2u; ++tq) {
const uint i0 = 4u * tq + c, i1 = i0 + 2u;
precise float16_t par[2];
[[unroll]] for (uint pa = 0u; pa < 2u; ++pa) {
precise float16_t pr[KC];
[[unroll]] for (uint b = 0u; b < KC; ++b)
pr[b] = rv[2u * b + pa][i0] * rv[2u * b + pa][i0] + rv[2u * b + pa][i1] * rv[2u * b + pa][i1];
precise float16_t tr = TREE(pr);
par[pa] = tr;
}
precise float16_t hv = par[0] + par[1];
ht[tq] = hv;
}
precise float16_t hh = ht[0] + ht[1];
h1[c] = hh;
}
const uint other = subgroupShuffleXor(bh(h1[0]) | (bh(h1[1]) << 16u), 16u);
precise float16_t g0 = h1[0] + hb(other & 0xffffu), g1 = h1[1] + hb(other >> 16u);
precise float16_t s = g0 + g1;
precise float16_t se = s + float16_t(0.0001220703125);
const float16_t r = rsq16(se);
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) {
const uint k = 16u * m + 8u * hl + uint(i);
precise float16_t g = r * w16.w[W_GAIN14 / 2u + k];
precise float16_t n = rv[m][i] * g;
#ifdef RR_GFX11
NbC[m][i >> 2] |= q8h(n) << (8u * uint(i & 3));
#else
Nb[m][i] = q8h(n);
#endif
precise float16_t a = rv[m][i] + w16.w[W_TABLE16 / 2u + k];
R[m][i] = RV(a);
}
#ifdef RR_GFX11
[[unroll]] for (uint m = 0u; m < MT; ++m) G11_B(Nb[m], NbC[m])
#endif
}
[[dont_unroll]] for (uint br = 0u; br < NB; ++br) {
const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br;
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
RAC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = RV(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) {
A8 a0, a1;
LDA(a0, w1, C, c, kt); LDA(a1, w1, C, c, kt + 1u);
K32(h, a0, Nb[kt], a1, Nb[kt + 1u])
}
#ifdef RR_GFX11
{ F8V hc_; ACTQ8C(hc_, h) G11_B(Hb[c], hc_) }
#else
ACTQ8(Hb[c], h)
#endif
}
[[unroll]] for (uint m = 0u; m < MT; ++m) {
A8 a0, a1;
LDA(a0, w2, 32u, m, 0u); LDA(a1, w2, 32u, m, 1u);
K32(R[m], a0, Hb[0], a1, Hb[1])
}
}
// ---------------------------------------------------------------- outputs
const uint t = inter(16u * wv + col);
const int tx = x0 + int(t & 7u), ty = y0 + int(t >> 3u);
const bool inside = tx >= 0 && ty >= 0 && tx < ex && ty < ey;
B8 Tb[MT];
{
const uint o = (outOff + uint(ty * ex + tx) * 32u) >> 2u;
#if KIND == 0 && CX
const uint ty2 = t >> 3u, tx2 = t & 7u;
const uint mr = (ty2 >> 1u) < 2u ? 2u * (tx2 >> 1u) + (ty2 >> 1u) : 8u + 2u * (tx2 >> 1u) + (ty2 >> 1u) - 2u;
const uint blk = 2u * (ty2 & 1u) + (tx2 & 1u);
barrier(); // every wave is done with XT (the merge rows reuse it)
#endif
[[unroll]] for (uint m = 0u; m < MT; ++m) {
#ifdef RR_GFX11
F8V tc_;
G11_CODES(tc_, R[m])
G11_B(Tb[m], tc_)
const uint w0 = tc_.x, w1 = tc_.y;
#else
Q8(Tb[m], R[m])
const uint w0 = pack4(Tb[m][0], Tb[m][1], Tb[m][2], Tb[m][3]), w1 = pack4(Tb[m][4], Tb[m][5], Tb[m][6], Tb[m][7]);
#endif
const uint at = 4u * (m & 1u) + 2u * hl;
if (inside) { arena.w[o + (m >> 1u) * (plane >> 2u) + at] = w0; arena.w[o + (m >> 1u) * (plane >> 2u) + at + 1u] = w1; }
#if KIND == 0 && CX
const uint g = mr * C + blk * (C / 4u) + 4u * m + 2u * hl;
MG[g] = w0; MG[g + 1u] = w1;
#endif
}
}
// the C -> 16 head
RAC hv;
[[unroll]] for (int i = 0; i < 8; ++i) hv[i] = RV(w16.w[W_BH / 2u + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) {
A8 a0, a1;
LDA(a0, W_WH, C, 0u, kt); LDA(a1, W_WH, C, 0u, kt + 1u);
K32(hv, a0, Tb[kt], a1, Tb[kt + 1u])
}
#if KIND == 0
{ // enc1's kernel-prediction filter (lanes 0-15; 16-31 hold head channels 8-15)
float16_t l[12];
[[unroll]] for (int i = 0; i < 8; ++i) l[i] = float16_t(hv[i]);
const uint u0 = subgroupShuffleXor(bh(l[0]) | (bh(l[1]) << 16u), 16u), u1 = subgroupShuffleXor(bh(l[2]) | (bh(l[3]) << 16u), 16u);
l[8] = hb(u0 & 0xffffu); l[9] = hb(u0 >> 16u); l[10] = hb(u1 & 0xffffu); l[11] = hb(u1 >> 16u);
const float16_t ma = max(max(max(l[0], l[2]), l[4]), l[6]), mb = max(max(max(l[1], l[3]), l[5]), l[7]);
const float16_t mx = max(max(ma, mb), l[8]);
float e[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) { precise float16_t d = l[i] - mx; e[i] = exp2(float(d) * 1.4426950216293335); }
precise float s = (e[0] + e[1]) + 0.0;
s = s + (e[2] + e[3]);
s = s + (e[4] + e[5]);
s = s + (e[6] + e[7]);
s = s + e[8];
const float16_t rr = rcp16(s);
float wg[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) wg[i] = r16(r16(e[i]) * float(rr));
const int x = tx, y = ty;
const int xm = x - 1 < 0 ? 1 : x - 1, xc = x < 0 ? 1 : (x >= ex ? ex - 2 : x), xp = x + 1 >= ex ? ex - 2 : x + 1;
const int ym = y - 1 < 0 ? 1 : y - 1, yc = y < 0 ? 1 : (y >= ey ? ey - 2 : y), yp = y + 1 >= ey ? ey - 2 : y + 1;
const ivec2 pts[9] = ivec2[9](ivec2(xm, ym), ivec2(xc, ym), ivec2(xp, ym), ivec2(xm, yc), ivec2(xc, yc),
ivec2(xp, yc), ivec2(xm, yp), ivec2(xc, yp), ivec2(xp, yp));
float acc[3][2] = float[3][2](float[2](0.0, 0.0), float[2](0.0, 0.0), float[2](0.0, 0.0));
[[unroll]] for (uint k = 0u; k < 8u; ++k) {
const vec4 tv = texelFetch(T88, pts[k], 0);
[[unroll]] for (uint c = 0u; c < 3u; ++c) { precise float f = tv[c] * wg[k] + acc[c][k & 1u]; acc[c][k & 1u] = r16(f); }
}
const vec4 t8 = texelFetch(T88, pts[8], 0);
float o112[3], o120[3];
[[unroll]] for (uint c = 0u; c < 3u; ++c) {
precise float f0 = acc[c][0] + acc[c][1];
precise float f1 = t8[c] * wg[8] + r16(f0);
const float f = r16(f1);
#if CX
precise float a = float(sig(l[9u + c])) * f;
const float ar = r16(a);
precise float b = f - ar;
o120[c] = ar; o112[c] = r16(b);
#else
o112[c] = f; // the bottleneck: the filtered value itself
#endif
}
if (hl == 0u && inside) {
imageStore(S112, ivec2(x, y), vec4(o112[0], o112[1], o112[2], 0.0));
#if CX
imageStore(S120, ivec2(x, y), vec4(o120[0], o120[1], o120[2], 0.0));
#endif
}
#if CX
if (hl == 0u) [[unroll]] for (uint c = 0u; c < 3u; ++c) O112[t * 4u + c] = float16_t(o112[c]);
#endif
}
#if CX
barrier();
{ // the 2x2 merge: merged^T (CX x 16 rows), channel tiles m = wv, wv + 4, ...
const uvec2 mxy = mergeXY(col);
const int X = x0 / 2 + int(mxy.x), Y = y0 / 2 + int(mxy.y);
const uint o = (mrgOff + uint(Y * hx + X) * 32u) >> 2u;
for (uint m = wv; m < CX / 16u; m += 4u) {
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_BM / 2u + 16u * m + 8u * hl + uint(i)]);
for (uint ks = 0u; ks < C / 8u; ++ks) {
B8 g0, g1; A8 w0, w1;
LDS8B(g0, MG, 8u * ks, C);
LDS8B(g1, MG, 8u * ks + 4u, C);
LDA(w0, W_WM, 4u * C, m, 2u * ks); LDA(w1, W_WM, 4u * C, m, 2u * ks + 1u);
K32(a, w0, g0, w1, g1)
}
#ifdef RR_GFX11
F8V q;
G11_CODES(q, a)
if (X >= 0 && Y >= 0 && X < hx && Y < hy) {
const uint at = o + (m >> 1u) * (lplane >> 2u) + 4u * (m & 1u) + 2u * hl;
arena.w[at] = q.x;
arena.w[at + 1u] = q.y;
}
#else
B8 q;
Q8(q, a)
if (X >= 0 && Y >= 0 && X < hx && Y < hy) {
const uint at = o + (m >> 1u) * (lplane >> 2u) + 4u * (m & 1u) + 2u * hl;
arena.w[at] = pack4(q[0], q[1], q[2], q[3]);
arena.w[at + 1u] = pack4(q[4], q[5], q[6], q[7]);
}
#endif
}
if (wv == 3u && lane < 16u) { // surface 128: 2x2 sums of 0.25 * surface 112, rightwards first, then downwards
const uint mx = lane & 3u, my = lane >> 2u;
const int x = x0 + 2 * int(mx), y = y0 + 2 * int(my);
if (x >= 0 && y >= 0 && x < ex && y < ey) {
const uint t00 = (2u * my) * 8u + 2u * mx;
vec4 out_ = vec4(0.0);
[[unroll]] for (uint c = 0u; c < 3u; ++c) {
precise float16_t a00 = O112[t00 * 4u + c] * float16_t(0.25), a01 = O112[(t00 + 1u) * 4u + c] * float16_t(0.25);
precise float16_t a10 = O112[(t00 + 8u) * 4u + c] * float16_t(0.25), a11 = O112[(t00 + 9u) * 4u + c] * float16_t(0.25);
precise float16_t b0 = a00 + a01, b1 = a10 + a11;
precise float16_t cc = b0 + b1;
out_[c] = float(cc);
}
imageStore(S128, ivec2(x >> 1, y >> 1), out_);
}
}
}
#endif
#else
[[unroll]] for (int i = 0; i < 8; ++i) OUTF[t * 16u + 8u * hl + uint(i)] = float16_t(hv[i]);
barrier();
#define DEC_FILTER
#include "rr_decfilter.glsl"
#endif
}
engine/rr/swin_pack.py
+75
View File
@@ -0,0 +1,75 @@
#!/usr/bin/env python3
"""Pack a Swin stage's extracted weights (engine/rr/swin_ref.py --extract) for engine/rr/swin_m.comp.
usage: swin_pack.py STAGE WEIGHTS.npz OUT.bin OUT.h
enc1_pack.py's layout for any stage of engine/rr/swin_ref.STAGES: transposed products, every weight's columns in the
k order of the product that consumes its result (slot j holds natural channel inter(n)[j]); the expand's 4C columns
are four sub-positions of C channels, each permuted the same way. OUT.h also defines the stage's shape for the shader
(KIND, C, CX, NH).
"""
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from enc1_ref import inter # noqa: E402
from enc1_pack import a_tiles, b_tiles, cubic_table, act_table # noqa: E402
from swin_ref import STAGES # noqa: E402
def main():
stage, src, out_bin, out_h = sys.argv[1:5]
kind, C, CX, NH, _ = STAGES[stage]
w = dict(np.load(src))
I32, I96, IC = inter(32), inter(96), inter(C)
e = {}
if kind == 'dec':
I4C = [C * (j // C) + IC[j % C] for j in range(4 * C)]
e['we'], e['be'] = ('A', w['we'][:, I4C]), ('H', w['be'][I4C])
for h in range(NH):
e[f'wk{h}'] = ('A', w[f'wk{h}'][:, I32])
e[f'wv{h}'] = ('B', w[f'wv{h}'][:, I32])
e[f'wo{h}'] = ('A', w[f'wo{h}'][:, IC])
e[f'bias{h}'] = ('H', w[f'bias{h}'][np.ix_(inter(64), I96)].T)
for b in range(C // 8):
e[f'w1_{b}'] = ('A', w[f'w1_{b}'][:, I32])
e[f'b1_{b}'] = ('H', w[f'b1_{b}'][I32])
e[f'w2_{b}'] = ('A', w[f'w2_{b}'][:, IC])
if kind == 'enc' and CX:
e['wm'], e['bm'] = ('A', w['wm'][:, inter(CX)]), ('H', w['bm'][inter(CX)])
e['wh'], e['bh'] = ('A', w['wh']), ('H', w['bh'])
e['gain0'] = ('H', w['gain0'])
for k in ('table7', 'gain14', 'table16'):
e[k] = ('H', w[k][IC])
e['cubic'] = ('U', cubic_table())
e['act'] = ('Y', act_table())
blob, defines = bytearray(), []
for name, (kd, a) in e.items():
while len(blob) % 256:
blob.append(0)
off = len(blob)
if kd == 'A':
blob += a_tiles(a).tobytes()
elif kd == 'B':
blob += b_tiles(a).tobytes()
elif kd == 'Y':
blob += np.ascontiguousarray(a, np.uint8).tobytes()
elif kd == 'U':
blob += np.ascontiguousarray(a, np.uint16).tobytes()
else:
blob += np.ascontiguousarray(a, np.float16).tobytes()
defines.append(f'#define W_{name.upper()} {off}u // {kd} {a.shape}')
shape = [f'#define KIND {0 if kind == "enc" else 1}', f'#define C {C}', f'#define CX {CX}', f'#define NH {NH}']
if NH < 2:
defines.append('#define W_WK1 (W_WK0 + 1u)')
if C // 8 < 2:
defines.append('#define W_W1_1 (W_W1_0 + 1u)')
open(out_bin, 'wb').write(bytes(blob))
open(out_h, 'w').write(f'// generated by engine/rr/swin_pack.py for {stage}\n' + '\n'.join(shape + defines) + '\n')
print(f'{stage}: {len(blob)} bytes, {len(e)} entries')
if __name__ == '__main__':
main()
engine/rr/swin_ref.py
+334
View File
@@ -0,0 +1,334 @@
#!/usr/bin/env python3
"""Exact numpy reference of the denoiser's Swin stages enc1-enc5 and dec1-dec4 (cuda_dldn_engine_swin_*_kernel), one
8x8 window at a time, from weights and real inputs only. Generalises engine/rr/enc1_ref.py (encoders) and
engine/rr/dec1_ref.py (decoders); every stage was checked against its capture (main).
A stage has C channels (32 * KC), NH heads of 32 channels, NB = C / 8 MLP branches of C -> 32 -> C:
prologue encoders: the 96 key tokens from the arena (C / 32 planes); decoders: the expand of the 6 x 6 tokens
one level down (C_LO channels -> 4 sub-positions x C) + the skip token. RMS norm (input_norm: per lane word
the planes summed pairwise (pairs_tree), the two word halves, then lanes ^1, ^2 and the halves) times gain0; the
residual is the token + table7
heads Q/K and V projections over KC k32 steps, scores + bias, the cubic softmax, P V, out-projection onto the
residual, head after head
MLP RMS norm (rowsum_sq) times gain14, residual + table16, NB branches
outputs the tokens; encoders: the 2x2 merge (4C -> C_M) and the 64 -> 16 head with enc1's filter (surfaces 112,
120, 128); decoders: the head with dec0's filter (surface 112)
usage: swin_ref.py CAPTURE STAGE FRAME WEIGHTS.npz [--blocks X,Y ...] (verifies against the capture)
swin_ref.py CAPTURE STAGE FRAME OUT.npz --extract DUMP.npz (weights from a phase dump)
"""
import argparse
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rrmath import e4m3, f16, q8 # noqa: E402
from enc1_ref import (inter, mm, softmax_cubic, act, mirror, add, rowsum_sq, pairs_tree, filter_head, # noqa: E402
MERGE_ORDER, H, F)
from dec0_ref import keys # noqa: E402
from dec1_ref import filter112 # noqa: E402
# kind, C, C_LO (decoders: channels one level down) or C_M (encoders: merged channels), heads, warps
STAGES = {
'enc1': ('enc', 64, 96, 2, 1), 'enc2': ('enc', 96, 128, 4, 4), 'enc3': ('enc', 128, 160, 4, 4),
'enc4': ('enc', 160, 160, 8, 8), 'enc5': ('enc', 160, 0, 8, 8),
'dec4': ('dec', 160, 160, 8, 8), 'dec3': ('dec', 128, 160, 4, 4), 'dec2': ('dec', 96, 128, 4, 4),
'dec1': ('dec', 64, 96, 2, 2),
}
def input_norm(xn, gain):
"""xn: one token, natural channels (f16). NVIDIA's tree over the arena order; returns the normalised token
(natural, f16)"""
c = len(xn)
xa = xn[inter(c)]
sq = (xa * xa).astype(H)
def pair(p, w):
return add(sq[32 * p + 4 * w:32 * p + 4 * w + 2], sq[32 * p + 4 * w + 2:32 * p + 4 * w + 4])
fold = pairs_tree # (p0 + p1) + p2, (p0 + p1) + (p2 + p3), ((p0 + p1) + (p2 + p3)) + p4
h = [add(fold([pair(p, t) for p in range(c // 32)]), fold([pair(p, t + 4) for p in range(c // 32)]))
for t in range(4)]
h01 = [add(h[t], h[t ^ 1]) for t in range(4)]
h2 = add(h01[0], h01[2])
s = H(H(h2[0] + h2[1]) + H(2.0 ** -13))
r = H(F(1) / np.sqrt(F(s)))
return (xn * (r * gain).astype(H)).astype(H)
def block(stage, w, x0, y0, ex, ey, inp):
"""the window: returns (tokens 64 x C codes in arena order, merged 16 x C_M codes or None, head 64 x 16 f16).
inp: 'tok' [C/32][ey][ex][32] (encoders) or 'lo' [C_LO/32][ey/2][ex/2][32] and 'skip' [C/32][ey][ex][32]"""
kind, C, CX, NH, _ = STAGES[stage]
IC = inter(C)
UN = [IC.index(i) for i in range(C)]
xs = []
if kind == 'dec':
hx, hy = ex // 2, ey // 2
X0, Y0 = x0 // 2 - 1, y0 // 2 - 1
low = np.zeros((48, CX), np.uint8)
for r in range(36):
X, Y = mirror(X0 + r % 6, hx), mirror(Y0 + r // 6, hy)
low[r] = np.concatenate([inp['lo'][p, Y, X] for p in range(CX // 32)])
Ex = q8(mm(low, w['we'], np.repeat(w['be'][None, :], 48, 0)))
for (x, y) in keys(x0, y0):
xx, yy = mirror(x, ex), mirror(y, ey)
r, s = (y // 2 - Y0) * 6 + (x // 2 - X0), 2 * (y & 1) + (x & 1)
sk = np.concatenate([inp['skip'][p, yy, xx] for p in range(C // 32)])[UN]
xs.append((e4m3(Ex[r, C * s:C * s + C]).astype(H) + e4m3(sk).astype(H)).astype(H))
else:
for (x, y) in keys(x0, y0):
xx, yy = mirror(x, ex), mirror(y, ey)
xs.append(e4m3(np.concatenate([inp['tok'][p, yy, xx] for p in range(C // 32)])[UN]).astype(H))
X = np.array([q8(input_norm(x, w['gain0']))[IC] for x in xs])
R = add(np.array(xs[:64]), w['table7'][None, :])
I32, I96 = inter(32), inter(96)
for h in range(NH):
qk, v = np.zeros((96, 32)), np.zeros((96, 32))
for c in range(C // 32):
qk = mm(X[:, 32 * c:32 * c + 32], w[f'wk{h}'][32 * c:32 * c + 32], qk)
v = mm(X[:, 32 * c:32 * c + 32], w[f'wv{h}'][32 * c:32 * c + 32], v)
Q = q8(qk)[:, I32]
P = softmax_cubic(mm(Q[:64], Q.T, w[f'bias{h}']))
O = mm(q8(P)[:, I96], q8(v)[I96], np.zeros((64, 32)))
R = mm(q8(O)[:, I32], w[f'wo{h}'], R)
N = q8((R * (rowsum_sq(R)[:, None] * w['gain14'][None, :]).astype(H)).astype(H))[:, IC]
acc = add(R, w['table16'][None, :])
for b in range(C // 8):
hid = act(mm(N, w[f'w1_{b}'], np.repeat(w[f'b1_{b}'][None, :], 64, 0)))
acc = mm(q8(hid)[:, I32], w[f'w2_{b}'], acc)
tokens = q8(acc)[:, IC]
merged = None
if kind == 'enc' and CX:
Mg = np.zeros((16, 4 * C), np.uint8)
for r, (mx, my) in enumerate(MERGE_ORDER):
for blk, (dx, dy) in enumerate([(0, 0), (1, 0), (0, 1), (1, 1)]):
Mg[r, C * blk:C * blk + C] = tokens[(2 * my + dy) * 8 + 2 * mx + dx]
merged = q8(mm(Mg, w['wm'], np.repeat(w['bm'][None, :], 16, 0)))[:, inter(CX)]
head = mm(tokens, w['wh'], np.repeat(w['bh'][None, :], 64, 0))
return tokens, merged, head
def phase_map(stage):
"""dump phase numbers: projections [h][c] = (QK, V), bias[h] (scores phase, wave), out[h], mlp up / down, merge,
head, expand"""
kind, C, CX, NH, waves = STAGES[stage]
KC = C // 32
m = {'proj': {}, 'bias': {}, 'out': {}}
if waves == 1: # heads in sequence: per head 2 KC projections, scores, P V, out-projection
p = 1
for h in range(NH):
for c in range(KC):
m['proj'][h, c] = (p, p + 1, 0)
p += 2
m['bias'][h] = (p, 0)
m['out'][h] = p + 2
p += 3
else: # one head per warp
p = 2 if kind == 'dec' else 1
for h in range(NH):
for c in range(KC):
m['proj'][h, c] = (p + 2 * c, p + 2 * c + 1, h)
p += 2 * KC
for h in range(NH):
m['bias'][h] = (p, h)
m['out'][h] = p + 2 + h
p += 2 + NH
m['mlp'] = [(p + 2 * b, p + 2 * b + 1) for b in range(C // 8)]
p += 2 * (C // 8)
if kind == 'enc' and CX:
m['merge'] = p
p += 1
m['head'] = p
return m
def read_b(W, base, K):
"""an FP8 B matrix (K x 32) as NVIDIA stores it in the weight image: 4-byte k groups, offset 512 (kg >> 3) +
16 (kg & 3) + 4 ((kg >> 2) & 1) + 64 (n & 7) + 8 ((n >> 3) & 1) + 16 K (n >> 4)"""
b = np.zeros((K, 32), np.uint8)
for kg in range(K // 4):
for n in range(32):
o = base + 512 * (kg >> 3) + 16 * (kg & 3) + 4 * ((kg >> 2) & 1) + 64 * (n & 7) + 8 * ((n >> 3) & 1) + 16 * K * (n >> 4)
b[4 * kg:4 * kg + 4, n] = W[o:o + 4]
return b
def merge_weights(z, W, p, waves, C, CX):
"""the merge (4C -> CX) of a multi-warp encoder: warp w computes the 32-column chunks w, w + warps, ...; the dump
holds each warp's last chunk, so the chunks come from the weight image (located through warp 1's, chunk 1) and the
dumped ones check the layout"""
K, size = 4 * C, 4 * C * 32
W32 = W[:len(W) // 4 * 4].view(np.uint32)
first = z[f'p{p}_w1_B'][0:4, 0].copy().view(np.uint32)[0]
pos1 = int(next(i for i in np.nonzero(W32 == first)[0] * 4 if np.array_equal(read_b(W, int(i), K), z[f'p{p}_w1_B'])))
base0 = pos1 - size
chunks = [read_b(W, base0 + size * j, K) for j in range(CX // 32)]
for v in range(waves):
last = max(j for j in range(CX // 32) if j % waves == v) if any(j % waves == v for j in range(CX // 32)) else None
if last is not None:
assert np.array_equal(chunks[last], z[f'p{p}_w{v}_B']), (v, last)
b1 = f16(z[f'p{p}_w1_C']).astype(H)[0].tobytes()
bpos1 = W.tobytes().find(b1)
bias = np.concatenate([W[bpos1 + 64 * (j - 1):bpos1 + 64 * j].view(H) for j in range(CX // 32)])
for v in range(waves):
last = max((j for j in range(CX // 32) if j % waves == v), default=None)
if last is not None:
assert np.array_equal(bias[32 * last:32 * last + 32].view(np.uint16), z[f'p{p}_w{v}_C'][0]), ('bias', v)
return np.hstack(chunks), bias
def extract(stage, dump, weights, out, inp, x0, y0, ex, ey):
"""weights from the phase dump; gains and tables are found in the weight image by matching the dump"""
kind, C, CX, NH, waves = STAGES[stage]
z = np.load(dump)
W = np.frombuffer(weights, np.uint8)
def B(p, wv=0):
return z[f'p{p}_w{wv}_B']
def Cm(p, wv=0):
return f16(z[f'p{p}_w{wv}_C']).astype(H)
pm = phase_map(stage)
w = {}
if kind == 'dec':
w['we'] = np.hstack([B(1, v) for v in range(waves)])
w['be'] = np.concatenate([Cm(1, v)[0] for v in range(waves)])
for h in range(NH):
w[f'wk{h}'] = np.vstack([B(pm['proj'][h, c][0], pm['proj'][h, c][2]) for c in range(C // 32)])
w[f'wv{h}'] = np.vstack([B(pm['proj'][h, c][1], pm['proj'][h, c][2]) for c in range(C // 32)])
w[f'bias{h}'] = Cm(*pm['bias'][h])
w[f'wo{h}'] = B(pm['out'][h])
for b, (up, dn) in enumerate(pm['mlp']):
w[f'w1_{b}'], w[f'b1_{b}'], w[f'w2_{b}'] = B(up), Cm(up)[0], B(dn)
if 'merge' in pm:
mw = min(waves, 4) # the merge runs on (at most) four warps
if waves > 1 and CX // 32 > mw: # a warp computes several chunks: read them from the weight image
w['wm'], w['bm'] = merge_weights(z, W, pm['merge'], mw, C, CX)
elif waves > 1:
w['wm'] = np.hstack([B(pm['merge'], v) for v in range(CX // 32)])
w['bm'] = np.concatenate([Cm(pm['merge'], v)[0] for v in range(CX // 32)])
else:
w['wm'], w['bm'] = B(pm['merge']), Cm(pm['merge'])[0]
w['wh'], w['bh'] = B(pm['head']), Cm(pm['head'])[0]
# the four f16 vectors: candidates at every even offset, checked against the first token rows of the dump
tables = [W[o:o + 2 * C].view(H) for o in range(0, len(W) - 2 * C + 1, 2)]
w0 = {k: v for k, v in w.items()}
# gain0: the key tokens' norm against the first projection's A (window row 0..)
xs = key_tokens(stage, w0, x0, y0, ex, ey, inp)
IC = inter(C)
A1 = np.hstack([z[f'p{pm["proj"][0, c][0]}_w{pm["proj"][0, c][2]}_A'] for c in range(C // 32)])
w['gain0'] = next(t for t in tables if np.array_equal(q8(input_norm(xs[0], t))[IC], A1[0]) and
np.array_equal(q8(input_norm(xs[5], t))[IC], A1[5]))
p_out = pm['out'][0]
C0 = Cm(p_out, 0)
rows = np.array(xs[:64])
w['table7'] = next(t for t in tables if np.array_equal(add(rows[:len(C0)][:2], t[None]).view(np.uint16), C0[:2].view(np.uint16)))
up, dn = pm['mlp'][0]
Rt = f16(z[f'p{pm["out"][NH - 1]}_w0_D']).astype(H)
Aup, Cdn = z[f'p{up}_w0_A'], Cm(dn)
w['gain14'] = next(t for t in tables if np.array_equal(
q8((Rt[:2] * (rowsum_sq(Rt[:2])[:, None] * t[None]).astype(H)).astype(H))[:, IC], Aup[:2]))
w['table16'] = next(t for t in tables if np.array_equal(add(Rt[:2], t[None]).view(np.uint16), Cdn[:2].view(np.uint16)))
np.savez_compressed(out, **w)
def key_tokens(stage, w, x0, y0, ex, ey, inp):
kind, C, CX, NH, _ = STAGES[stage]
IC = inter(C)
UN = [IC.index(i) for i in range(C)]
xs = []
if kind == 'dec':
hx, hy = ex // 2, ey // 2
X0, Y0 = x0 // 2 - 1, y0 // 2 - 1
low = np.zeros((48, CX), np.uint8)
for r in range(36):
X, Y = mirror(X0 + r % 6, hx), mirror(Y0 + r // 6, hy)
low[r] = np.concatenate([inp['lo'][p, Y, X] for p in range(CX // 32)])
Ex = q8(mm(low, w['we'], np.repeat(w['be'][None, :], 48, 0)))
for (x, y) in keys(x0, y0):
xx, yy = mirror(x, ex), mirror(y, ey)
if kind == 'dec':
r, s = (y // 2 - Y0) * 6 + (x // 2 - X0), 2 * (y & 1) + (x & 1)
sk = np.concatenate([inp['skip'][p, yy, xx] for p in range(C // 32)])[UN]
xs.append((e4m3(Ex[r, C * s:C * s + C]).astype(H) + e4m3(sk).astype(H)).astype(H))
else:
xs.append(e4m3(np.concatenate([inp['tok'][p, yy, xx] for p in range(C // 32)])[UN]).astype(H))
return xs
def load(cap, stage, frame):
from rr_capture import Capture
c = Capture(os.path.expanduser(cap))
l = c.find(f'swin_{stage}', frame)
kind, C, CX, NH, _ = STAGES[stage]
ex, ey = (int(v) for v in np.frombuffer(l.args[:8], np.uint32))
sx, sy = (int(v) for v in np.frombuffer(l.args[80:84], np.uint16))
ptr = {p.offset: p.delta for p in l.pointers}
pre = np.frombuffer(l.pre(l.arena_alloc), np.uint8)
post = np.frombuffer(l.post_buffer(c, l.arena_alloc), np.uint8)
inp = {}
if kind == 'dec':
inp['lo'] = pre[ptr[8]:ptr[8] + CX * (ex // 2) * (ey // 2)].reshape(CX // 32, ey // 2, ex // 2, 32)
inp['skip'] = pre[ptr[24]:ptr[24] + C * ex * ey].reshape(C // 32, ey, ex, 32)
else:
inp['tok'] = pre[ptr[8]:ptr[8] + C * ex * ey].reshape(C // 32, ey, ex, 32)
outs = {'tok': post[ptr[48]:ptr[48] + C * ex * ey].reshape(C // 32, ey, ex, 32)}
if kind == 'enc' and CX:
outs['mrg'] = post[ptr[56]:ptr[56] + CX * (ex // 2) * (ey // 2)].reshape(CX // 32, ey // 2, ex // 2, 32)
img = {i.offset: i for i in l.images}
def surf(off):
i = img[off]
return np.frombuffer(i.path.read_bytes(), np.float16).reshape(i.height, i.width, 4)
tex = {k: surf(k) for k in img}
return l, ex, ey, sx, sy, inp, outs, tex
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('stage')
ap.add_argument('frame', type=int)
ap.add_argument('weights')
ap.add_argument('--extract', metavar='DUMP')
ap.add_argument('--dump-block', default='1,1', help='the block the dump was taken of (for --extract)')
ap.add_argument('--blocks', nargs='*', default=['1,1'])
o = ap.parse_args()
l, ex, ey, sx, sy, inp, outs, tex = load(o.capture, o.stage, o.frame)
kind, C, CX, NH, _ = STAGES[o.stage]
if o.extract:
bx, by = (int(v) for v in o.dump_block.split(','))
extract(o.stage, o.extract, l.weights, o.weights, inp, 8 * bx - sx, 8 * by - sy, ex, ey)
return
w = dict(np.load(o.weights))
for blk in o.blocks:
bx, by = (int(v) for v in blk.split(','))
x0, y0 = 8 * bx - sx, 8 * by - sy
tokens, merged, head = block(o.stage, w, x0, y0, ex, ey, inp)
inside = [t for t in range(64) if 0 <= x0 + t % 8 < ex and 0 <= y0 + t // 8 < ey]
tk = sum(np.array_equal(tokens[t], np.concatenate([outs['tok'][p, y0 + t // 8, x0 + t % 8] for p in range(C // 32)]))
for t in inside)
msg = f'block {bx},{by}: tokens {tk}/{len(inside)}'
if merged is not None:
mg = sum(np.array_equal(merged[r], np.concatenate([outs['mrg'][p, y0 // 2 + my, x0 // 2 + mx] for p in range(CX // 32)]))
for r, (mx, my) in enumerate(MERGE_ORDER) if 0 <= x0 // 2 + mx < ex // 2 and 0 <= y0 // 2 + my < ey // 2)
msg += f', merged {mg}'
if kind == 'enc' and not CX: # the bottleneck: the ungated filter to surface 112 only
o112, _ = filter_head(head, tex[88], x0, y0, ex, ey, ungated=True)
ok = sum(np.array_equal(o112[t], tex[112][y0 + t // 8, x0 + t % 8, :3]) for t in inside)
elif kind == 'enc':
o112, o120 = filter_head(head, tex[88], x0, y0, ex, ey)
ok = sum(np.array_equal(o112[t], tex[112][y0 + t // 8, x0 + t % 8, :3]) and
np.array_equal(o120[t], tex[120][y0 + t // 8, x0 + t % 8, :3]) for t in inside)
else:
s112 = filter112(head, tex[88], tex[96], tex[104], x0, y0, ex, ey)
ok = sum(np.array_equal(s112[t], tex[112][y0 + t // 8, x0 + t % 8, :3]) for t in inside)
print(msg + f', surfaces {ok}/{len(inside)}')
if __name__ == '__main__':
main()
engine/rr/test_layer.py
+223
View File
@@ -0,0 +1,223 @@
#!/usr/bin/env python3
"""Run one Ray Reconstruction stage shader on a captured launch and compare it with NVIDIA's result.
usage: test_layer.py CAPTURE KERNEL FRAME SHADER.spv WORK [--kvk build/rr/kvk] [--time N] [--subgroup 32|64]
KERNEL is a substring of the denoiser kernel's name (e.g. swin_enc1, hkpn_output_kernel_transformer), FRAME the
1-based occurrence in the capture. The shader follows the RR stage contract (engine/rr/README.md):
binding 0 the parameter block: NVIDIA's argument block with every pointer into the plane arena
replaced by its byte offset in the arena, the weight pointer and every texture or surface
handle by zero (stage_params below; the runtime builds the same block)
binding 1 the weight image exactly as NGX uploads it (argument 40), read-only
binding 2 the plane arena, read/write
bindings 3, 4, ... the textures in ascending argument offset, then the surfaces in ascending argument offset
Every surface starts zeroed and the arena starts as it was before NVIDIA's launch; afterwards each surface and the
whole arena are compared with what NVIDIA's launch left. Exit status 0 only when everything is byte-identical.
"""
import argparse
import json
import os
import struct
import subprocess
import sys
from pathlib import Path
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rr_capture import Capture # noqa: E402
# GLSL image formats of the stage shaders' surfaces -> (channels, kvk bits)
SURFACE_FORMATS = {'rgba16f': (4, 16), 'rg16f': (2, 16), 'r16f': (1, 16), 'rgba8_snorm': (4, 9), 'rgba8': (4, 8)}
def stage_params(launch, weights_arg=40):
args = bytearray(launch.args)
arena = launch.arena_alloc
for p in launch.pointers:
value = p.delta if p.alloc == arena and p.offset != weights_arg else 0
struct.pack_into('<Q', args, p.offset, value)
# texture and surface handles: 1 where NGX bound one (the kernels test optional inputs' handles for null, e.g.
# enc0's specular hit distance at argument 248), 0 elsewhere
for image in launch.images:
struct.pack_into('<Q', args, image.offset, 1)
return bytes(args)
def e4m3_table():
t = np.zeros(256)
for c in range(256):
s, e, m = c >> 7, (c >> 3) & 15, c & 7
v = np.nan if (e == 15 and m == 7) else (m * 2.0 ** -9 if e == 0 else (1 + m / 8) * 2.0 ** (e - 7))
t[c] = -v if s else v
return t
def psnr(a, b):
finite = np.isfinite(a) & np.isfinite(b)
if not finite.any():
return float('nan')
err = np.mean((a[finite] - b[finite]) ** 2)
peak = max(np.max(np.abs(b[finite])), 1e-6)
return float('inf') if err == 0 else 10 * np.log10(peak * peak / err)
def compare_image(name, image, got, ref, ignore_channels=()):
"""IGNORE_CHANNELS: channels NVIDIA's kernel fills from an undefined register (left out of the comparison)"""
if image.bits == 16:
a, b = np.frombuffer(got, np.float16).astype(np.float64), np.frombuffer(ref, np.float16).astype(np.float64)
elif image.bits == 32:
a, b = np.frombuffer(got, np.float32).astype(np.float64), np.frombuffer(ref, np.float32).astype(np.float64)
else:
dt = np.int8 if image.bits == 9 else np.uint8
a, b = np.frombuffer(got, dt).astype(np.float64), np.frombuffer(ref, dt).astype(np.float64)
keep = np.ones(len(a), bool)
for c in ignore_channels:
keep[c::image.channels] = False
a, b = a[keep], b[keep]
element_bytes = len(got) // len(keep)
same = (np.frombuffer(got, np.uint8).reshape(-1, element_bytes) == np.frombuffer(ref, np.uint8).reshape(-1, element_bytes))[keep]
elements = len(a)
equal = np.sum(a == b) + np.sum(np.isnan(a) & np.isnan(b))
worst = np.nanmax(np.abs(a - b)) if elements else 0
print(f' {name}: {image.width}x{image.height}x{image.channels} fmt{image.format}'
f'{" (channels " + ",".join(map(str, ignore_channels)) + " ignored)" if ignore_channels else ""}: '
f'{100.0 * equal / elements:.4f}% identical values, max |diff| {worst:.6g}, PSNR {psnr(a, b):.2f} dB, '
f'nonfinite ours/ref {int(np.sum(~np.isfinite(a)))}/{int(np.sum(~np.isfinite(b)))}')
return bool(same.all())
def compare_arena(got, before, ref):
g, p, r = (np.frombuffer(x, np.uint8) for x in (got, before, ref))
written = (r != p) | (g != p)
differ = g != r
print(f' arena: {int(written.sum())} bytes written by either side, {int(differ.sum())} differ', end='')
if differ.any():
idx = np.nonzero(differ)[0]
t = e4m3_table()
dv = np.abs(t[g[idx]] - t[r[idx]])
print(f' (first at byte {idx[0]:#x}, last at {idx[-1]:#x}; as e4m3 values: max |diff| {np.nanmax(dv):.6g}, '
f'{int(np.sum(dv <= 0.0625 * np.abs(t[r[idx]]) + 1e-12))} within one e4m3 step)')
else:
print()
return not differ.any()
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('kernel')
ap.add_argument('frame', type=int)
ap.add_argument('shader')
ap.add_argument('work')
ap.add_argument('--kvk', default=str(Path(__file__).resolve().parents[2] / 'build/rr/kvk'))
ap.add_argument('--time', type=int, default=1)
ap.add_argument('--subgroup', type=int, default=32)
ap.add_argument('--weights', help='bind this file at binding 1 instead of NGX\'s weight image (native stages)')
ap.add_argument('--weights-arg', type=int, default=40, help='argument offset of the weight pointer (rr2_*: 24)')
ap.add_argument('--grid', help='X,Y workgroups (default: the launch grid)')
ap.add_argument('--ignore', action='append', default=[], metavar='SURFACE:CHANNEL',
help='leave a channel out of the comparison, e.g. 112:3 for a value NVIDIA stores from an '
'undefined register (repeatable)')
o = ap.parse_args()
cap = Capture(o.capture)
launch = [l for l in cap.launches if o.kernel in l.kernel and 'diamond' not in l.kernel or
('diamond' in o.kernel and o.kernel in l.kernel)][o.frame - 1]
work = Path(o.work)
work.mkdir(parents=True, exist_ok=True)
print(f'{launch.kernel} (launch {launch.sequence}, grid {launch.grid}, block {launch.block})')
(work / 'params.bin').write_bytes(stage_params(launch, o.weights_arg))
(work / 'weights.bin').write_bytes(Path(o.weights).read_bytes() if o.weights else launch.weights_at(o.weights_arg))
arena = launch.arena_alloc
arena_before = launch.pre(arena) if arena is not None else bytes(16)
(work / 'arena.bin').write_bytes(arena_before)
gx, gy = (int(v) for v in o.grid.split(',')) if o.grid else launch.grid[:2]
lines = [f'param {work}/params.bin', f'grid {gx} {gy}']
# Bind in the shader's order (its ptx2rr.py sidecar, SHADER.json: textures then surfaces by argument offset).
# A handle the shader reads but NGX left null gets a zeroed 1x1 image.
sidecar = Path(o.shader).with_suffix('.json')
if sidecar.exists():
info = json.loads(sidecar.read_text())
order = [('texture', off, None) for off in info['textures']] + [('surface', off, f) for off, f in info['surfaces']]
else:
order = [(i.kind, i.offset, None) for i in launch.images]
by_key = {(i.kind, i.offset): i for i in launch.images}
(work / 'zero.bin').write_bytes(bytes(64))
for kind, offset, fmt in order:
image = by_key.get((kind, offset))
if image is None:
ch, bits = SURFACE_FORMATS.get(fmt, (4, 16))
lines.append(f'tex {offset} {work}/zero.bin 1 1 4 16 nearest' if kind == 'texture' else
f'surf {offset} 1 1 {ch} {bits} - {work}/dummy{offset}.bin')
elif kind == 'texture':
lines.append(f'tex {image.offset} {image.path} {image.width} {image.height} {image.channels} {image.bits} '
f'{"linear" if image.linear else "nearest"}')
else:
lines.append(f'surf {image.offset} {image.width} {image.height} {image.channels} {image.bits} - '
f'{work}/surf{image.offset}.bin')
n_images = len(order)
lines.append(f'lut {work}/weights.bin')
lines.append(f'inout {work}/arena.bin {work}/arena.out.bin')
# further buffers of the stage (ptx2rr.py --buffer), bound after the surfaces
extra = info.get('buffers', []) if sidecar.exists() else []
for off in extra:
present = any(p.offset == off for p in launch.pointers) # NGX may leave the pointer null
(work / f'buf{off}.bin').write_bytes(launch.pre(launch.pointer(off).alloc) if present else bytes(16))
lines.append(f'inout {work}/buf{off}.bin {work}/buf{off}.out.bin')
probes = info.get('probes', []) if sidecar.exists() else []
oob = info.get('oob', False) if sidecar.exists() else False
if probes: # ptx2rr.py --probe: the debug buffer after the surfaces
# D4R_PROBE_THREADS: the shader's workgroup size when it is not the launch's block (the hybrid stages)
threads = int(os.environ.get('D4R_PROBE_THREADS', launch.block[0] * launch.block[1] * launch.block[2]))
lines.append(f'outbuf {work}/debug.bin {4 * threads * len(probes)}')
elif oob: # ptx2rr.py --oob-report: count, lowest and highest address per buffer
(work / 'oob.init').write_bytes(struct.pack('<9I', *([0, 0xffffffff, 0] * 3)))
lines.append(f'inout {work}/oob.init {work}/oob.bin')
lines.append('bindings 0 ' + ' '.join(str(3 + i) for i in range(n_images)) + ' 1 2' +
''.join(f' {3 + n_images + i}' for i in range(len(extra))) +
(f' {3 + n_images + len(extra)}' if probes or oob else ''))
(work / 'manifest').write_text('\n'.join(lines) + '\n')
r = subprocess.run([o.kvk, o.shader, str(work / 'manifest'), str(o.time), str(o.subgroup)],
capture_output=True, text=True)
sys.stdout.write(r.stdout)
if r.returncode != 0:
sys.stderr.write(r.stderr)
sys.exit(f'kvk failed ({r.returncode})')
identical = True
bound = {(kind, offset) for kind, offset, _ in order}
for image in launch.images:
if image.kind != 'surface':
continue
if ('surface', image.offset) not in bound:
print(f' surface {image.offset}: not written by this kernel')
continue
got = (work / f'surf{image.offset}.bin').read_bytes()
waived = [int(s.split(':')[1]) for s in o.ignore if int(s.split(':')[0]) == image.offset]
identical &= compare_image(f'surface {image.offset}', image, got, image.path.read_bytes(), waived)
if arena is not None:
identical &= compare_arena((work / 'arena.out.bin').read_bytes(), arena_before,
launch.post_buffer(cap, arena))
for off in extra:
if not any(p.offset == off for p in launch.pointers):
continue
got, ref = (work / f'buf{off}.out.bin').read_bytes(), launch.post_buffer(cap, launch.pointer(off).alloc)
same = got[:len(ref)] == ref
print(f' buffer {off}: {len(ref)} bytes, ' + ('identical' if same else f'{sum(x != y for x, y in zip(got, ref))} differ'))
identical &= same
if oob:
r = struct.unpack('<9I', (work / 'oob.bin').read_bytes()[:36])
for k, name in enumerate(('weight reads', 'arena reads', 'arena writes')):
print(f' out of range {name}: {r[3 * k]}' + (f' (bytes {r[3 * k + 1]:#x}..{r[3 * k + 2]:#x}, size '
f'{len(launch.weights) if k == 0 else len(arena_before):#x})' if r[3 * k] else ''))
print('IDENTICAL' if identical else 'DIFFERENT')
sys.exit(0 if identical else 1)
if __name__ == '__main__':
main()
engine/rr/widen_weights.py
+33
View File
@@ -0,0 +1,33 @@
#!/usr/bin/env python3
"""Append the RDNA3 (RR_GFX11) widened copy of a native stage's weight file and print its f16 element offset (W_WIDE).
usage: widen_weights.py WEIGHTS.bin [OUT.bin] (default: in place)
Element W_WIDE + j of the file's f16 view is byte j of the original image as ZLUDA widens an FP8 operand on gfx11
(rr_gfx11.glsl): the magnitude at bits 7..14, the sign at bit 15, the NaN code with bit 14 set too, i.e. the binary16
of value * 2^-8. Bytes that are not FP8 (tables, binary16 data) are widened as well and never read. W_WIDE is the
original size rounded up to 512 bytes, halved.
"""
import sys
from pathlib import Path
import numpy as np
def widen(image):
b = np.frombuffer(image, np.uint8).astype(np.uint16)
mag = b & 0x7f
mag = mag | ((mag + 1) & 0x80)
return ((mag << 7) | ((b & 0x80) << 8)).astype('<u2').tobytes()
def main():
src = Path(sys.argv[1]); dst = Path(sys.argv[2]) if len(sys.argv) > 2 else src
image = src.read_bytes()
base = (len(image) + 511) // 512 * 512
dst.write_bytes(image + bytes(base - len(image)) + widen(image))
print(base // 2)
if __name__ == '__main__':
main()
engine/rr/zluda_replay.py
+104
View File
@@ -0,0 +1,104 @@
#!/usr/bin/env python3
"""Replay one captured denoiser launch through ZLUDA on Linux (tools/replay_rr_launch.cpp) and compare it with the
capture, like test_layer.py does for a stage shader.
usage: zluda_replay.py CAPTURE KERNEL FRAME WORK [--ptx MODULE] [--zluda DIR]
KERNEL and FRAME select the launch as in test_layer.py. MODULE defaults to the corpus module holding the kernel
(~/.cache/d4r-rr-corpus); pass an edited copy to add debug stores. ZLUDA's own settings (D4R_ZLUDA_WMMA,
D4R_ZLUDA_WMMA_FP8, D4R_ZLUDA_WMMA_FP8_NATIVE, ...) are taken from the environment. Outputs land in WORK
(surfOFFSET.bin, arena.out.bin) so the two replays can be compared with each other too.
"""
import argparse
import os
import subprocess
import sys
from pathlib import Path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rr_capture import Capture # noqa: E402
from test_layer import compare_arena, compare_image # noqa: E402
ROOT = Path(__file__).resolve().parents[2]
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('kernel')
ap.add_argument('frame', type=int)
ap.add_argument('work')
ap.add_argument('--ptx')
ap.add_argument('--zluda', default=os.path.expanduser('~/.cache/d4r-zluda-rr'))
ap.add_argument('--rocm', default='/opt/rocm/lib')
ap.add_argument('--tex', action='append', default=[], metavar='OFFSET=FILE:W:H',
help='an RGBA32F texture for a handle the capture has no dump of (preset F pooled histories)')
ap.add_argument('--pad', type=int, default=0, help='bytes of zeros appended to the arena (debug stores)')
o = ap.parse_args()
cap = Capture(o.capture)
exact = [l for l in cap.launches if l.kernel == o.kernel]
launch = (exact or [l for l in cap.launches if o.kernel in l.kernel and 'diamond' not in l.kernel or
('diamond' in o.kernel and o.kernel in l.kernel)])[o.frame - 1]
work = Path(o.work)
work.mkdir(parents=True, exist_ok=True)
corpus = '~/.cache/d4r-rr-corpus-f' if launch.kernel.startswith('rr2_') else '~/.cache/d4r-rr-corpus'
module = o.ptx or next(str(p) for p in sorted(Path(corpus).expanduser().glob('*.ptx'))
if f'.entry {launch.kernel}(' in p.read_text(errors='replace'))
(work / 'args.bin').write_bytes(launch.args)
arena = launch.arena_alloc
lines = [f'args {work}/args.bin', 'grid ' + ' '.join(map(str, launch.grid)),
'block ' + ' '.join(map(str, launch.block))]
if o.pad:
(work / 'arena-pad.bin').write_bytes(launch.pre(arena) + bytes(o.pad))
for p in launch.pointers:
out = f' {work}/arena.out.bin' if p.alloc == arena else ''
source = f'{work}/arena-pad.bin' if o.pad and p.alloc == arena else f'{launch.directory}/alloc-{p.alloc}.bin'
lines.append(f'buf {p.offset} {source} {p.delta}{out}')
# textures the launch names but the bridge could not read back (preset F's 10:10:10:2 pooled histories, which
# ZLUDA backs with RGBA32F): zeros, which is what they hold on a first frame
have = {i.offset for i in launch.images}
for line in (launch.directory / 'manifest.txt').read_text().split('\n'):
f = line.split()
given = {int(t.split('=')[0]): t.split('=')[1].split(':') for t in o.tex}
if len(f) == 2 and f[0] == 'texture' and int(f[1]) in given:
path, w_, h_ = given[int(f[1])]
lines.append(f'tex {f[1]} {path} {w_} {h_} 4 32 linear')
elif len(f) == 2 and f[0] == 'texture' and int(f[1]) not in have:
(work / 'zero16.bin').write_bytes(bytes(16 * 16 * 16))
lines.append(f'tex {f[1]} {work}/zero16.bin 16 16 4 32 linear')
if len(f) == 2 and f[0] == 'surface' and int(f[1]) not in have: # half the largest surface: the pooled history
big = max((i for i in launch.images if i.kind == 'surface'), key=lambda i: i.width, default=None)
if big is not None:
lines.append(f'surf {f[1]} {big.width // 2} {big.height // 2} 4 32 {work}/surf{f[1]}.bin')
for image in launch.images:
if image.kind == 'texture':
# SNORM_INT8X4 (200): a SIGNED_INT8 array read as normalized floats gives the same values
fmt = 8 if image.format == 200 else 1 if image.format == 194 else image.format
lines.append(f'tex {image.offset} {image.path} {image.width} {image.height} {image.channels} '
f'{fmt} {"linear" if image.linear else "nearest"}')
else:
# ZLUDA has no SNORM_INT8X4 arrays (format 200); surface stores write the same bytes to a SIGNED_INT8 one
fmt = 8 if image.format == 200 else 1 if image.format == 194 else image.format # 194: UNORM_INT8X4
lines.append(f'surf {image.offset} {image.width} {image.height} {image.channels} {fmt} '
f'{work}/surf{image.offset}.bin')
(work / 'manifest').write_text('\n'.join(lines) + '\n')
env = dict(os.environ, LD_LIBRARY_PATH=f'{o.zluda}:{o.rocm}:' + os.environ.get('LD_LIBRARY_PATH', ''))
print(f'{launch.kernel} (launch {launch.sequence}) through ZLUDA {o.zluda}')
r = subprocess.run([str(ROOT / 'build/replay_rr_launch'), module, launch.kernel, str(work / 'manifest')],
env=env, capture_output=True, text=True)
sys.stdout.write(r.stdout)
if r.returncode:
sys.stderr.write(r.stderr[-4000:])
sys.exit('replay failed')
same = True
for image in launch.images:
if image.kind == 'surface':
same &= compare_image(f'surface {image.offset}', image, (work / f'surf{image.offset}.bin').read_bytes(),
image.path.read_bytes())
if arena is not None and not o.pad:
same &= compare_arena((work / 'arena.out.bin').read_bytes(), launch.pre(arena), launch.post_buffer(cap, arena))
print('IDENTICAL' if same else 'DIFFERENT')
if __name__ == '__main__':
main()