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

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

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

No files matched your search

engine/rr/README.md
+164
View File
@@ -172,6 +172,25 @@ single-use multiply-subtract pairs into fma, merging the two constant multiplica
`precise` on the inlined exp2, NoContraction on its conversions. The differing values come in spatial clusters, one
colour channel at a time, always one ulp, mostly just above 0.5 in surface 224.
## HIP against Vulkan, and the fused rounding (2026-10-12)
Encoder 1 ported line for line to a standalone HIP kernel (`kernels/rr/swin_native.h`, runner
`kernels/tools/rr_swin_native_test.cpp`; capD1080 frame 2, 5 launches per round): 835 us against the Vulkan stage's
770-788 us, the arena byte-identical. A HIP port of the network is not worth making. What the port did show:
- The FP8 product alone costs 0.61 us per product per wave on this grid, 2.07 us when each is followed by the
f32 -> binary16 -> f32 rounding of its eight accumulator registers. Products and that rounding are about half of
a stage on either backend.
- Keeping the accumulator as binary16 and taking f16(P + acc) with `v_fma_mix` (as `kernels/m` does) halves the
rounding's instructions: HIP 780 us. RADV emits the same instruction for `float16_t(p + float(acc))`, so the
Vulkan stages have it as a build option: `D4R_RR_MIX=1 build_native_model.sh ...` (`RR_MIX`, the shaders' K32).
`RR_MIX` is one rounding of the exact sum where the reference has two, so it is not bit-exact by construction.
Preset D: every network stage still identical to the capture on all 21 frames of capDT12 / capD4K / capD1080 / capDX
/ capDXcp / capDXhw (the same figures as the default build), 4.22 ms at 1280x720 -> 4K and 8.3 ms at 1920x1080 -> 4K
(default 4.45 / 8.7). Preset E (capE3): frames 1 and 3 have a value of encoder 0 one step off (127 dB), which the
later stages amplify to 41 dB at the output; frame 2 is identical. The default build stays the exact one.
## Cyberpunk 2077: more inputs and settings (2026-10-12)
Found on harness captures (`D4R_HARNESS_RR_EXTRA=particles,sss`, `D4R_HARNESS_RR_HW_DEPTH=1`, `D4R_RR_PRESET=N`):
@@ -189,3 +208,148 @@ byte-identical to the CUDA path through the shim on the frames the harness can c
with RE Requiem's set the network is identical on every frame (output 107-109 dB). At 1280x720 and 1920x1080 -> 4K
without these inputs the output still falls to 43-53 dB after the first frame (the output stage's one-ulp values);
a correctly rounded log2 there matches the capture worse than the GPU instruction does.
## The output stage's one-step differences: fused multiply-convert (2026-10-13)
The rare one-binary16-step values of the translated output stage (above: ~0.004% of surfaces 224 / 240, "cause not
found") were the driver fusing an f32 product (or sum) with the conversion that follows it into `v_fma_mix`, which
rounds the exact result once. The kernel rounds twice, to f32 and then to binary16, and the two differ when the f32
result is exactly a binary16 midpoint, as `log2(x) * c1 * c2` of the output stage's tone curve sometimes is.
`ptx2rr.py` now ORs an opaque zero into the f32 bits before every `cvt.rn.f16.f32` (and into the binary16 bits
after, against the fold of `float(float16_t(x))`). With it presets D and E are byte-identical to the CUDA path on
every frame of every capture (capDT12, capD4K, capD1080, capDX, capDXcp, capDXhw, capE3; history and output), at
+0.05 ms: 4.27 / 4.45 ms at 1280x720 -> 4K. Found with `build/rr`-style bisection: probe every register of a block
through ZLUDA (debug stores added to the PTX, `zluda_replay.py`) and through the translated shader
(`ptx2rr.py --probe`) and list the first that differ.
Other reference behaviours the translation now follows, all found the same way on preset F's kernels:
- formatted stores to normalized surfaces round half away from zero in f32 (ZLUDA's `surface_formatted_store_bits`);
UNORM 10:10:10:2 arrays are RGBA32F images holding the quantised values `q / 1023` (the correctly rounded quotient,
from a table: the driver divides by multiplying with a reciprocal);
- `tanh.approx.f32` is OCML's `tanhF` with LLVM's AMDGPU lowering of `exp` (`lowerFEXP`, fast-FMA path): the same
operations, the fused ones through f64 (the driver's f32 `fma` is a product and a sum whenever the product also
exists on its own); `tanh.approx.f16` is the correctly rounded value (a table);
- the driver's GLSL `fma` on binary16 is a true fused operation, as the reference's is: no emulation there;
- texture fetches with a binary16 result keep that rounding.
## Preset F: DLSS 4.5 Ray Reconstruction (2026-10-13)
`nvngx_dlssd.dll` 310.9 has a third model with its own kernels, `rr2_*`, and runs it for every preset hint but 4 (D)
and 5 (E): none, 6, Cyberpunk 2077's 0x80000002. The engine runs it from a `D4RR0002` model:
```sh
python3 kernels/tools/extract_dlss_ptx.py NVNGX_DLSSD_DLL ~/.cache/d4r-rr-corpus-f # its PTX (zstd in 310.9)
python3 engine/compile_rr2.py CAPTURE OUT_DIR --native NVNGX_DLSSD_DLL [--particles CAPTURE_WITH_PARTICLES]
```
The runtime is `engine/runtime_rr2.inc` (behind `RModel` / `REngine`); the shim uses `engine\rr-f` (or
`D4R_ENGINE_MODEL_RR_F_DIR`) and, when that model exists, NGX 310.9's hint mapping.
A frame, as NGX launches it:
| pass | NVIDIA kernel | engine |
|---|---|---|
| exposure | `rr2_histogram_auto_exposure_basic_kernel` | `exposure_f.comp` (two passes) |
| enc0 | `rr2_enc0_kernel` | translated feature gather and guide tail around the native block (`gen_rr2_enc0.py`) |
| enc1 - enc5 | `rr2_enc1_kernel` ... | `rr2_m.comp` |
| dec4 - dec1 | `rr2_decN_kernel_0`, `rr2_decN_kernel_1` | `rr2_m.comp` |
| dec0 | `rr2_dec0_kernel_0`, `rr2_dec0_kernel_1` | `rr2_m.comp`; the second: native block, translated output tail |
| output | `rr2_post_kernel`, `rr2_downsample_kernel_static_hdr[_no_transparency]` per tile, `rr2_pooled_history2_kernel` | translated (`ptx2rr.py`) |
The network (`rr2_ref.py` is its exact numpy reference and documents the block): D/E's levels, window grids and
arithmetic (FP8 operands, binary16 sums rounded per 32-deep step, the cubic softmax, the MLP activation), with a
different block. Each kernel is one block on 8x8-token windows that alternate their shift per frame: heads that
share one projection for queries and keys (so the scores are symmetric; bias `(T + T^T) sqrt(2)` of a 64 x 64
relative position table), all reading the block's input and summing their out-projections; MLP branches likewise;
both sums enter the residual times a binary16 scalar from the argument block. Encoders merge 2x2 for the next level,
a decoder's first kernel expands the level below and adds the skip times a third scalar. Next to the tokens every
level carries a three-channel guide image: encoders gate it (`sigmoid(head) * s0`) and average it 2x2 for the next
level, decoders upsample it bilinearly, gate it and blend it with the encoder's. All stages exchange data through
the arena (3152 units of W0 * H0 / 32 bytes; NGX's region offsets, which scale with the render size), tokens as FP8
planes of 32 channels, guides as four binary16 per token.
Weights: the DLL holds the model as plain binary16 tensors in natural `[out][in]` order, kernel after kernel
(`rr2_ref.dll_weights`; found by the sixteen kernels' relative position tables, 2 2 4 4 8 8 8 8 4 4 4 4 2 2 2 2
heads). `rr2_pack.py` packs them for `rr2_m.comp`; nothing of the network comes from the capture. The capture still
supplies every kernel's argument block (the runtime keeps their constants and fills the geometry, the frame's
values and the pointers) and encoder 0's second table.
The output stage filters an intermediate image of 1.5 times the output size in tiles (columns of at most 720, rows
of at most 1620 intermediate pixels, 4 pixels of overlap inside), reduces each tile to the output, and keeps two
histories for the next frame: the intermediate image, and a half-size copy whose quarter-size reduction
(`rr2_pooled_history2_kernel`) encoder 0 also reads (UNORM 10:10:10:2 arrays, see above). NGX runs filter and
reduction concurrently with progress counters; the kernels skip that when the counter pointers are null, which is how
the engine runs them, in order. They have no matrix products and run with 64-wide waves (their warp-level shuffles
address the lane's own half).
Options (harness captures capFX*): encoder 0 byte 466 packed roughness, 568 inverted depth, 569 / texture 272
`DLSSD.SpecularHitDistance`, 571 / texture 336 `DLSS.DisocclusionMask`, texture 288 the subsurface guide, 1028
`DLSS.Use.HW.Depth`; `DLSSD.ColorBeforeParticles` selects the other reduction kernel (model stage `downp`, texture
208, an output-size history 256 <-> 264, flag bytes 196-198). Encoder 0's surfaces 376 / 384 are the output filter's
textures 64 / 72. Pre-exposure does not reach these kernels.
Checked (`d4r-rr` on `rr_inputs.py` inputs; every stage surface and the presented image against the CUDA path):
byte-identical on every frame of capF (640x360 -> 1080p, 3 frames), capF900 (800x450 -> 1600x900), capF4K
(1280x720 -> 4K, 16 tiles) and of the option captures (motion, a real camera, inverted and hardware depth, packed
roughness, hit distance, disocclusion mask, subsurface guide, particles); through the shim (D3D12 harness under
Proton) the engine's frames are byte-identical to the CUDA backend's too (`D4R_RR_PRESET=6 D4R_RR_DLL=<310.9 DLL>
scripts/run_rr_engine_compare.sh`; that script still compares frame N with frame N, and the harness presents the
CUDA result one frame late in the jitter scene).
Speed, RX 9070 XT, 1280x720 -> 3840x2160: 6.2 / 6.6 ms per frame (exposure 0.03, network 4.1 - 4.2, output 2.1 -
2.4), from 14.5 ms with every kernel translated; preset D is 4.3 / 4.5. Per stage (us): enc0 1076 (its block 521),
dec0_1 775, dec0_0 630, enc1 340, dec1_0 316, dec1_1 269, the rest under 180; filter and reduction about 130 per
tile, 16 tiles. Not done: native output passes and feature gather.
### Where the translated parts' time went (2026-10-13, later)
Now about 5.8 / 6.2 ms (5.7 / 6.1 with the 10-bit histories that were later withdrawn), still byte-identical on every capture and through the
shim. Profiling the translated code (cutting it at points behind an opaque condition, `D4R_ENGINE_RR_SKIP=post,down,
pooled`) showed its arithmetic is spread evenly - no hot spot a rewrite would remove - and that the losses were in
how the passes were arranged:
- **Encoder 0 in two passes** (-0.2 ms). One shader ran NVIDIA's 32-lane gather on one of the block's four waves
while the other three waited: 0.40 ms for the gather, 0.24 ms as a pass of its own in 32-thread groups.
`gen_rr2_enc0.py --split gather | block` writes both; the gather pass (`stage enc0g`, the same bindings and
parameter block as enc0) stores each window's record in the arena after the tile slots: the first product's A
fragments (1024 words) and, per lane, the values the tail reads from the gather (six binary16 on the reference
build; the generator finds them and recomputes the tail's index arithmetic). The record's address is a word
compile_rr2.py appends to enc0's argument block. Cost: 76 bytes per render pixel of arena.
`D4R_RR2_SPLIT=0` builds the single shader.
- **Tile passes in overlapping groups** (-0.15 ms). The 16 tiles' post and down passes ran one after another, each
behind a barrier, and the gaps between 32 short dispatches added up. Each tile now has a slot (its two regions;
eight slots, 58 MB at 4K) and a group's post passes run without barriers between them, as do its down passes
together with the next group's post passes on the other half of the slots - what NGX does with two slots and its
progress counters. `D4R_ENGINE_RR_SKIP=serial` restores the old order.
- **Pooled histories as real 10:10:10:2 images: withdrawn.** 4 bytes a pixel instead of RGBA32F's 16 saved
0.1 ms and was identical at 3x and 2x scaling, but the hardware's bilinear filter rounds a UNORM image differently
from a float one: at 1.5x (capFQ) temporal frames came out at 44 - 51 dB. The default is RGBA32F again;
`D4R_RR2_POOLED=rgb10_a2` builds the other form.
Tried without gain: 32-bit arena accesses for the translated code's paired 16-bit ones (0.01 ms), dropping the
subgroup shuffles or the depth search (they cost nothing measurable). What is left is NVIDIA's arithmetic: gather
0.24 - 0.33 ms, decoder 0's output tail 0.18, post 1.1 - 1.3, down 0.7.
### Other scaling ratios, and what a game-like configuration showed (2026-10-13, later)
Asked whether instability seen in RE Requiem came from this implementation, the engine was compared with the CUDA
path through the shim in configurations the earlier captures did not cover (`D4R_RR_PRESET=6 D4R_RR_DLL=<310.9>
scripts/run_rr_engine_compare.sh`, frames aligned by one):
- **Render widths that are not a multiple of 32 were wrong** (1114x626 -> 1080p: 15 dB): the post pass's argument 8
is the row stride of the network's output, the padded width W0, not the render width. Fixed; identical now.
- **Below 4/3 scaling NGX runs another output filter**, `rr2_post_kernel_3_2` (1456x820 -> 1080p and larger render
sizes, DLAA included; 1440x810 and smaller use `rr2_post_kernel`). The model has no such stage, so the shim hands
those features to the CUDA backend. `rr2_histogram_auto_exposure_histogram_kernel` was never launched in any run.
- **DLSS.Roughness.Mode 1 with a roughness texture bound**: NGX follows the mode, the shim followed the texture.
Fixed for preset F (RE Requiem binds none, so it was not affected).
- 1.5x, 1.72x, 2x and 4/3 scaling, and RE Requiem's inputs together (inverted and hardware depth, hit distance,
mask, packed roughness, scaled motion vectors, a moving scene): identical.
- **Open: the camera matrices are not NGX's bit for bit.** NGX computes them in single precision: argument 644 is
the f32 product view x clip, 836 its inverse by cofactors times 1 / det (glm::inverse's operation order; both
reproduced bit for bit in build/rr/f/matfit.py), 900 is an inverse of the pixel matrix (it carries that inverse's
residues, and 0.9999999 for 1 at 1114x626) and 964 was not reproduced. The engine computes them in double and
rounds, a last-bit difference in a few elements. With a non-identity camera, hit distance and mask bound and a
moving scene that showed as 40 - 50 dB on later frames in a synthetic scene whose affected feature is rounding
noise around zero in both paths; no other configuration was affected.
engine/rr/build_native_model.sh
+2
View File
@@ -42,6 +42,8 @@ PRESET="$(awk '$1 == "preset" { print $2 }' "$OUT/rr.txt")"
DEFS=(-DRR_EXACT_CONVERT); [[ "$PRESET" == 4 ]] && DEFS+=(-DRR_PRESET_D)
# D4R_RR_RELAXED=1: one rounding where NVIDIA has two (engine/rr/rr_exact.glsl): faster, not bit-exact
[[ -n "${D4R_RR_RELAXED:-}" ]] && DEFS+=("-DRR_RELAXED=$D4R_RR_RELAXED")
# D4R_RR_MIX=1: the 32-deep steps round once through v_fma_mix (the shaders' K32): ~5% faster, not provably bit-exact
[[ "${D4R_RR_MIX:-0}" == 1 ]] && DEFS+=(-DRR_MIX)
glsl() { # glsl OUT.spv SOURCE INCLUDE_DIR...
local out="$1" src="$2"; shift 2
local dir inc=(); for dir in "$@"; do inc+=("-I$dir"); done
engine/rr/dec0_m.comp
+25 -9
View File
@@ -48,7 +48,23 @@ layout(set = 0, binding = 7, rgba8_snorm) uniform writeonly image2D S144;
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
#define Q8(dst, src) Q8PAIRS(dst, src)
// One 32-deep step onto an accumulator that holds binary16 values. The reference (and the default here) adds the
// accumulator inside the f32 products and rounds the result to binary16. RR_MIX (build_native_model.sh with
// D4R_RR_MIX=1) keeps the accumulator as binary16, sums the products from zero and takes f16(P + acc), which the
// driver compiles to one v_fma_mix per element in place of two conversions: about 5% faster, but a single rounding
// where the reference has two, so a value can come out one binary16 step away (engine/rr/README.md).
// RAC is the accumulator type, RV(x) a binary16 value as its element; K32H is the step on a binary16 matrix.
#ifdef RR_MIX
#define RAC HAC
#define RV(x) (x)
#define K32(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(0.0))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_] + float((acc)[i_])); }
#define K32H(acc, a0, b0, a1, b1) K32(acc, a0, b0, a1, b1)
#else
#define RAC ACC
#define RV(x) float(x)
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) }
#define K32H(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(acc))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_]); }
#endif
const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0;
// shared memory, aliased by phase: expand and prologue [LOW | EX], then the attention [K | V^T]
@@ -111,8 +127,8 @@ void main()
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) coopMatLoad(lb[kt], LOW, n * 256u + 4u * kt, 16u, CM);
[[unroll]] for (uint mm = 0u; mm < 2u; ++mm) {
const uint m = 2u * wv + mm;
ACC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_BE / 2u + 16u * m + 8u * hl + uint(i)]);
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_BE / 2u + 16u * m + 8u * hl + uint(i)]);
K32(a, we[mm][0], lb[0], we[mm][1], lb[1])
K32(a, we[mm][2], lb[2], we[mm][3], lb[3])
B8 q;
@@ -167,13 +183,13 @@ void main()
barrier();
// ---------------------------------------------------------------- the residual: window tokens + table 0
ACC R[2];
RAC R[2];
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
HAC x;
coopMatLoad(x, XR, wv * 512u + 16u * m, 32u, CM);
[[unroll]] for (int i = 0; i < 8; ++i) {
precise float16_t c = float16_t(x[i]) + w16.w[W_TABLE0 / 2u + 16u * m + 8u * hl + uint(i)];
R[m][i] = float(c);
R[m][i] = RV(c);
}
}
@@ -190,12 +206,12 @@ void main()
LDA(wk[kt], W_WK0 + HEAD * hd, 32u, c, kt);
LDB(wvb[kt], W_WV0 + HEAD * hd, 32u, kt, c);
}
ACC a = ACC(0.0);
RAC a = RAC(RV(float16_t(0.0)));
K32(a, wk[0], xb[0], wk[1], xb[1])
A8 ka;
Q8(ka, a)
coopMatStore(ka, POOL, n * 128u + 4u * c, 8u, RM);
ACC v = ACC(0.0);
RAC v = RAC(RV(float16_t(0.0)));
K32(v, xa[0], wvb[0], xa[1], wvb[1])
A8 va;
Q8(va, v)
@@ -241,15 +257,15 @@ void main()
precise float16_t n = rv[m][i] * g;
Nb[m][i] = q8h(n);
precise float16_t a = rv[m][i] + w16.w[W_TABLE1 / 2u + k];
R[m][i] = float(a);
R[m][i] = RV(a);
}
}
for (uint br = 0u; br < 4u; ++br) {
const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br;
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
ACC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = float(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]);
RAC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = RV(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]);
A8 a0, a1;
LDA(a0, w1, 32u, c, 0u); LDA(a1, w1, 32u, c, 1u);
K32(h, a0, Nb[0], a1, Nb[1])
engine/rr/enc0_m_decl.glsl
+16
View File
@@ -22,7 +22,23 @@ layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
#define Q8(dst, src) Q8PAIRS(dst, src)
// One 32-deep step onto an accumulator that holds binary16 values. The reference (and the default here) adds the
// accumulator inside the f32 products and rounds the result to binary16. RR_MIX (build_native_model.sh with
// D4R_RR_MIX=1) keeps the accumulator as binary16, sums the products from zero and takes f16(P + acc), which the
// driver compiles to one v_fma_mix per element in place of two conversions: about 5% faster, but a single rounding
// where the reference has two, so a value can come out one binary16 step away (engine/rr/README.md).
// RAC is the accumulator type, RV(x) a binary16 value as its element; K32H is the step on a binary16 matrix.
#ifdef RR_MIX
#define RAC HAC
#define RV(x) (x)
#define K32(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(0.0))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_] + float((acc)[i_])); }
#define K32H(acc, a0, b0, a1, b1) K32(acc, a0, b0, a1, b1)
#else
#define RAC ACC
#define RV(x) float(x)
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) }
#define K32H(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(acc))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_]); }
#endif
const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0;
shared uint KVS[2 * 768]; // the current head's K, then V^T; before the attention the key pixels' features
#define XF KVS // the key pixels' 32 FP8 features (plane 0 | plane 1), slot-major
engine/rr/enc0_m_tail.glsl
+16 -16
View File
@@ -20,23 +20,23 @@
// the embedding x = W2 q8(relu(W1 f + b1)) + b2 for token tiles n = wv (this wave's queries: x is also their
// residual) and n = wv + 4 (the border keys, waves 0 and 1), then each key's norm in registers (dec0_ref.norm_plane
// on the accumulator layout, as dec0_m.comp's MLP norm) to XT
ACC R[2];
RAC R[2];
for (uint n = wv; n < 6u; n += 4u) {
B8 fb0, fb1, hb_[2];
coopMatLoad(fb0, XF, n * 128u, 8u, CM);
coopMatLoad(fb1, XF, n * 128u + 4u, 8u, CM);
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
ACC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_B1 / 2u + 16u * m + 8u * hl + uint(i)]);
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_B1 / 2u + 16u * m + 8u * hl + uint(i)]);
A8 a0, a1;
LDA(a0, W_W1, 32u, m, 0u); LDA(a1, W_W1, 32u, m, 1u);
K32(a, a0, fb0, a1, fb1)
[[unroll]] for (int i = 0; i < 8; ++i) hb_[m][i] = q8s(max(a[i], 0.0));
[[unroll]] for (int i = 0; i < 8; ++i) hb_[m][i] = q8s(max(float(a[i]), 0.0));
}
float16_t xv[2][8];
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
ACC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_B2 / 2u + 16u * m + 8u * hl + uint(i)]);
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_B2 / 2u + 16u * m + 8u * hl + uint(i)]);
A8 a0, a1;
LDA(a0, W_W2, 32u, m, 0u); LDA(a1, W_W2, 32u, m, 1u);
K32(a, a0, hb_[0], a1, hb_[1])
@@ -44,7 +44,7 @@
if (n == wv) { // the residual: x + table 0
[[unroll]] for (int i = 0; i < 8; ++i) {
precise float16_t c = xv[m][i] + w16.w[W_TABLE0 / 2u + 16u * m + 8u * hl + uint(i)];
R[m][i] = float(c);
R[m][i] = RV(c);
}
}
}
@@ -90,12 +90,12 @@
LDA(wk[kt], W_WK0 + HEAD * hd, 32u, c, kt);
LDB(wvb[kt], W_WV0 + HEAD * hd, 32u, kt, c);
}
ACC a = ACC(0.0);
RAC a = RAC(RV(float16_t(0.0)));
K32(a, wk[0], xb[0], wk[1], xb[1])
A8 ka;
Q8(ka, a)
coopMatStore(ka, KVS, n * 128u + 4u * c, 8u, RM);
ACC vv = ACC(0.0);
RAC vv = RAC(RV(float16_t(0.0)));
K32(vv, xa[0], wvb[0], xa[1], wvb[1])
A8 va;
Q8(va, vv)
@@ -139,15 +139,15 @@
precise float16_t nn = rv[m][i] * g;
Nb[m][i] = q8h(nn);
precise float16_t a = rv[m][i] + w16.w[W_TABLE1 / 2u + k];
R[m][i] = float(a);
R[m][i] = RV(a);
}
}
for (uint br = 0u; br < 4u; ++br) {
const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br;
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
ACC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = float(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]);
RAC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = RV(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]);
A8 a0, a1;
LDA(a0, w1, 32u, c, 0u); LDA(a1, w1, 32u, c, 1u);
K32(h, a0, Nb[0], a1, Nb[1])
@@ -180,8 +180,8 @@
}
}
{ // the 32 -> 16 head and encoder 0's filter of the noise grid (two f16 fma chains, taps even / odd)
ACC hv;
[[unroll]] for (int i = 0; i < 8; ++i) hv[i] = float(w16.w[W_BH / 2u + 8u * hl + uint(i)]);
RAC hv;
[[unroll]] for (int i = 0; i < 8; ++i) hv[i] = RV(w16.w[W_BH / 2u + 8u * hl + uint(i)]);
A8 a0, a1;
LDA(a0, W_WH, 32u, 0u, 0u); LDA(a1, W_WH, 32u, 0u, 1u);
K32(hv, a0, Tb[0], a1, Tb[1])
@@ -228,8 +228,8 @@
const uvec2 mxy = mergeXY(col);
const int X = x0 / 2 + int(mxy.x), Y = y0 / 2 + int(mxy.y), hx = ex / 2, hy = ey / 2;
const uint m = wv;
ACC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_BM / 2u + 16u * m + 8u * hl + uint(i)]);
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_BM / 2u + 16u * m + 8u * hl + uint(i)]);
[[unroll]] for (uint ks = 0u; ks < 4u; ++ks) {
B8 g0, g1; A8 w0, w1;
coopMatLoad(g0, MG, 8u * ks, 32u, CM);
engine/rr/exposure_f.comp
+93
View File
@@ -0,0 +1,93 @@
#version 460
// Ray Reconstruction preset F (rr2_*): the auto exposure of rr2_histogram_auto_exposure_basic_kernel, in two passes.
//
// NVIDIA's kernel sums a clamped luma per 16x16 block (a pairwise tree over the block's threads), converts each block
// sum to fixed point (x 2^32 / 10000, round to nearest) and adds them with a 64-bit atomic; the last block to finish
// turns the total into the exposure. The integer total does not depend on the order of the blocks, so PASS 0 writes
// every block's fixed-point sum and PASS 1 (one workgroup) adds them and does what the last block does:
// exposure = 0.23547088 / max(mean luma, 0.005) -> the 1x1 exposure image of this frame
// values = { e, p, 1 / e, clamp(log2(e / p), -6, 6) / 6 } -> the 16-byte buffer encoder 0 and the output passes read
// with e the exposure truncated to binary16 and p the previous frame's (1 where either is zero).
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_EXT_control_flow_attributes : require
#ifndef PASS
#define PASS 0
#endif
layout(set = 0, binding = 0, std430) readonly buffer Params {
ivec2 last; // render size - 1
uint blocksX, blocks;
float invCount; // 10000 / (2^32 * render pixels)
uint reset; // the kernel's argument 32: no log ratio when set
} p;
#if PASS == 0
// One thread per column of the 16x16 block. NVIDIA's tree over the block's 256 threads (index 16 y + x: i += i + 128,
// then 64, 32, ... 1) is, per column, ((r0 + r8) + (r4 + r12)) + ((r2 + r10) + (r6 + r14)) of the even rows plus the
// same of the odd rows, and then the columns (c + c8, + 4, + 2, + 1).
layout(local_size_x = 16) in;
layout(set = 0, binding = 1) uniform sampler2D color;
layout(set = 0, binding = 2, std430) writeonly buffer Sums { uint64_t s[]; } sums;
shared float col[16];
void main()
{
const uint tx = gl_LocalInvocationID.x;
const int x = int(gl_GlobalInvocationID.x), y0 = int(gl_WorkGroupID.y) * 16;
float r[16];
[[unroll]] for (int k = 0; k < 16; ++k) {
float v = 0.0;
if (x <= p.last.x && y0 + k <= p.last.y) {
const vec3 t = texelFetch(color, ivec2(x, y0 + k), 0).rgb;
precise float l = fma(t.b, 0.25, fma(t.r, 0.25, t.g * 0.5));
v = min(max(l, 0.0), 10000.0);
}
r[k] = v;
}
precise float u[8];
[[unroll]] for (int k = 0; k < 8; ++k) u[k] = r[k + 8] + r[k];
precise float v4[4];
[[unroll]] for (int k = 0; k < 4; ++k) v4[k] = u[k + 4] + u[k];
precise float w0 = v4[2] + v4[0], w1 = v4[3] + v4[1];
precise float c = w1 + w0;
col[tx] = c;
barrier();
if (tx == 0u) {
precise float d[8];
[[unroll]] for (int k = 0; k < 8; ++k) d[k] = col[k + 8] + col[k];
precise float e0 = d[4] + d[0], e1 = d[5] + d[1], e2 = d[6] + d[2], e3 = d[7] + d[3];
precise float f0 = e2 + e0, f1 = e3 + e1;
precise float g = f1 + f0;
sums.s[gl_WorkGroupID.y * p.blocksX + gl_WorkGroupID.x] = uint64_t(roundEven(g * 429496.71875));
}
}
#else
layout(local_size_x = 64) in;
shared uint64_t part[64];
layout(set = 0, binding = 1, std430) readonly buffer Sums { uint64_t s[]; } sums;
layout(set = 0, binding = 2, std430) writeonly buffer Values { vec4 v; } values;
layout(set = 0, binding = 3) uniform sampler2D previous;
layout(set = 0, binding = 4, r16f) uniform writeonly image2D current;
void main()
{
// the block sums are integers: any order. 64 strided partial sums, then thread 0 finishes
const uint tid = gl_LocalInvocationID.x;
uint64_t mine = 0ul;
for (uint b = tid; b < p.blocks; b += 64u) mine += sums.s[b];
part[tid] = mine;
barrier();
if (tid != 0u) return;
uint64_t total = 0ul;
for (uint b = 0u; b < 64u; ++b) total += part[b];
precise float mean = p.invCount * float(total);
precise float e = 0.23547088 / max(mean, 0.005);
imageStore(current, ivec2(0), vec4(e));
// cvt.rz.f16.f32: toward zero (e > 0)
float16_t h = float16_t(e);
if (float(h) > e) h = uint16BitsToFloat16(float16BitsToUint16(h) - uint16_t(1));
float cur = float(h);
if (cur == 0.0) cur = 1.0;
float prev = texelFetch(previous, ivec2(0), 0).r;
if (prev == 0.0) prev = 1.0;
float ratio = 0.0;
if (p.reset == 0u) ratio = clamp(log2(max(cur, 1e-8) / max(prev, 1e-8)), -6.0, 6.0) * 0.16666667;
values.v = vec4(cur, prev, 1.0 / cur, ratio);
}
#endif
engine/rr/gen_rr2_enc0.py
+298
View File
@@ -0,0 +1,298 @@
#!/usr/bin/env python3
"""Generate preset F's encoder 0 shader: ptx2rr.py's translation of NVIDIA's per-pixel feature gather and of its guide
tail, around the native block (engine/rr/rr2_m.comp, KIND 4).
usage: gen_rr2_enc0.py ENC0.ptx ENC0.generated.comp OUT.comp
gen_rr2_enc0.py DEC0_1.ptx DEC0_1.generated.comp OUT.comp --dec0_1
NVIDIA's kernel runs one 32-lane warp per 8x8-token window and keeps everything in registers. Its matrix products
(the embedding, the attention block and the guide head) are one contiguous run of the translated main, which is cut
out:
wave w the translation up to the first product: the gather; the first product's A fragments (the window's 64
feature slots per token) go to the shared array XF
4 waves rr2_block(): embedding, attention, MLP, the tokens to the arena, the guide head to the shared array HD
wave w the last product's D fragments are loaded from HD and the translation continues after it: the guide tail
With --group N (2 to 4) a workgroup holds N windows along x: waves 0 ... N - 1 run the translated code of one window
each and the block then runs once per window; the stage's dispatch is 1 / N as wide ("xN" in rr.txt). Measured on
an RX 9070 XT (2026-10-13) that is slower than one window per workgroup (network +0.4 ms at 1280x720 with N = 4).
With --dec0_1 the same for dec0's second kernel (rr2_m.comp KIND 5): nothing before the block (it reads its tokens
from the arena); after it the last product's D fragments (the MLP sum) and the registers holding r are loaded from
MSH / RSH, and the translated tail forms y = r + MLP * s2 and computes the kernel's outputs from it.
The translated parts hold code derived from the user's DLL: the output is model data, do not commit it.
"""
import re
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent))
from phase_dump import mma_statements, phases_of # noqa: E402
def tiles(mm):
"""per mma of a phase: (m-tile, n-tile, k-step), as phase_dump.py orders them"""
dset = {tuple(m[2]): n for n, m in enumerate(mm)}
prev = {n: dset.get(tuple(m[5])) for n, m in enumerate(mm)}
step = {}
for n in range(len(mm)):
s, q = 0, n
while prev[q] is not None:
q = prev[q]
s += 1
step[n] = s
a_tile, b_tile = {}, {}
for s in range(max(step.values()) + 1):
for t, key in enumerate(dict.fromkeys(tuple(m[3]) for n, m in enumerate(mm) if step[n] == s)):
a_tile[(s, key)] = t
for t, key in enumerate(dict.fromkeys(tuple(m[4]) for n, m in enumerate(mm) if step[n] == s)):
b_tile[(s, key)] = t
return [(a_tile[(step[n], tuple(m[3]))], b_tile[(step[n], tuple(m[4]))], step[n]) for n, m in enumerate(mm)]
def f16vars_of(decls):
names = set()
for l in decls:
if re.match(r'\s*(precise )?float16_t ', l):
names |= set(re.findall(r'\b(rs\d+)\b', l))
return names
REG = r'\b(?:rs|rd|r|f|p|l)\d+(?:_low|_high)?\b'
# what makes a statement of the gather more than index arithmetic on the parameter block
DATA = re.compile(r'fetch0|texture|shfl_|GX\d|GA\d|GW\d|SH\d|sust_|movm_|mma_')
def tail_inputs(pre, tail):
"""The gather's values the tail reads: (recomputed, saved). Recomputed are the statements, in order, of values
that are index arithmetic on the parameter block (assigned once, outside any block); saved are the names of the
others, which the gather pass writes to the window's record."""
assigned, depth = {}, 0
for i, l in enumerate(pre):
m = re.match(r'\s*(\w+) = (.*)', l)
if m:
assigned.setdefault(m.group(1), []).append((i, depth, m.group(2)))
depth += l.count('{') - l.count('}')
used, own = [], set()
for l in tail:
m = re.match(r'\s*(\w+) = (.*)', l)
for v in re.findall(REG, m.group(2) if m else l):
if v not in own and v in assigned and v not in used:
used.append(v)
if m:
own.add(m.group(1))
def closure(v, seen):
"""statement indices v needs, or None when it is not plain index arithmetic"""
if v in seen:
return seen[v]
seen[v] = None
defs = assigned.get(v)
if defs is None:
seen[v] = set() # never assigned in the gather: a declaration's zero
return seen[v]
if len(defs) != 1 or defs[0][1] != 0 or DATA.search(defs[0][2]):
return None
need = {defs[0][0]}
for u in re.findall(REG, defs[0][2]):
c = closure(u, seen)
if c is None:
return None
need |= c
seen[v] = need
return need
lines, saved, seen = set(), [], {}
for v in used:
c = closure(v, seen)
if c is None:
saved.append(v)
else:
lines |= c
return [pre[i] for i in sorted(lines)], saved
def split_layout(xf_word, pre, tail, f16vars):
"""(recomputed statements, saved names, the record's words, the GLSL locating a window's record)"""
recomputed, saved = tail_inputs(pre, tail)
for v in saved:
if not (v in f16vars or re.fullmatch(r'r\d+', v)):
raise SystemExit(f'the tail reads {v} from the gather: not a 16- or 32-bit register')
record = 1024 + 32 * len(saved)
where = (f' const uint rr2_gx = (param.w[0] + (param.w[20] & 0xffffu) + 7u) / 8u;\n'
f' const uint rr2_rec = param.w[{xf_word}u] / 4u + {record}u * (gl_WorkGroupID.y * (param.w[0] / 8u + 1u) + gl_WorkGroupID.x);\n')
return recomputed, saved, record, where
def gather_pass(xf_word, head, decls, pre, tail, first, f16vars):
"""Encoder 0 as two passes (--split gather | block): the translated gather in NVIDIA's own 32-thread groups, which
writes each window's record to the arena, and the native block with the translated tail in 128-thread groups.
One shader holding both runs the gather on one of the block's four waves while the other three wait, which
costs more than the record's trip through memory (RX 9070 XT, 1280x720, 2026-10-13: gather 0.40 -> 0.24 ms).
The record: 1024 words of A fragments ([token slot][64 FP8 feature slots], as the block's XF), then per lane the
gather's values the tail reads. Its address: parameter word XF_WORD (bytes) + 4 * RECORD * (y * (W / 8 + 1) + x)."""
recomputed, saved, record, where = split_layout(xf_word, pre, tail, f16vars)
var = lambda r: r.lstrip('%')
store = [' if (rr2_valid) { // the first product\'s A fragments: lane 4g + t, register q -> row, 4 k',
' const uint l_ = gl_SubgroupInvocationID, g_ = l_ >> 2u, t_ = l_ & 3u;']
seen = set()
for mm_, (i, j, s) in zip(first, tiles(first)):
if (i, s) in seen:
continue
seen.add((i, s))
for q, r in enumerate(mm_[3]):
store.append(f' rr2_arena.w[rr2_rec + rr2_uninter({16 * i + 8 * (q & 1)}u + g_) * 16u + {(32 * s + 16 * (q >> 1)) // 4}u + t_] = {var(r)};')
for n, v in enumerate(saved):
bits = f'uint(float16BitsToUint16({v}))' if v in f16vars else v
store.append(f' rr2_arena.w[rr2_rec + 1024u + {len(saved)}u * l_ + {n}u] = {bits};')
store.append(' }')
print(f'gather pass: record of {record} words, the tail reads {", ".join(saved) or "nothing"} from it')
return (head + '\nlayout(set = 0, binding = 2, std430) buffer Rr2Arena { uint w[]; } rr2_arena;\n'
'uint rr2_uninter(uint c) { return (c & ~31u) | (((c & 15u) >> 1u) << 2u) | (c & 1u) | (((c >> 4u) & 1u) << 1u); }\n\n'
'void main()\n{\n' + '\n'.join(decls) + '\n' + where +
' const bool rr2_valid = gl_WorkGroupID.x < rr2_gx;\n'
' if (rr2_valid) { // NVIDIA\'s warp: the translation before the first product\n' + '\n'.join(pre) + '\n }\n' +
'\n'.join(store) + '\n}\n'), record
def main():
ptx, src, out = sys.argv[1:4]
dec = '--dec0_1' in sys.argv
# --group N: windows per workgroup (1 to 4; wave w runs the translated code of window w)
G = int(sys.argv[sys.argv.index('--group') + 1]) if '--group' in sys.argv else 1
kernel = 'rr2_dec0_kernel_1' if dec else 'rr2_enc0_kernel'
lines = Path(ptx).read_text().replace('\r', '').split('\n')
k = next(i for i, l in enumerate(lines) if f'.entry {kernel}(' in l)
phases = phases_of(mma_statements(lines, k, len(lines)), lines)
first, last = phases[0], phases[-1]
var = lambda r: r.lstrip('%')
g = Path(src).read_text().split('\n')
m = next(i for i, l in enumerate(g) if l.startswith('void main()'))
head = '\n'.join(g[:m]).replace('layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in;',
'layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in;')
if dec: # the tail's tanh from rr2_pack.py's table (ptx2rr.py's RR_TANH_TABLE; its offset from OUT's rr2_stage.h)
v = head.index('\n') + 1
table = re.search(r'#define W_TANH (\d+)u', (Path(out).parent / 'rr2_stage.h').read_text()).group(1)
head = head[:v] + f'#define RR_TANH_TABLE {table}u\n' + head[v:]
body = g[m + 2:]
while body and body[-1].strip() in ('', '}'):
last_line = body.pop()
if last_line.strip() == '}':
break
nd = next(i for i, l in enumerate(body) if not re.match(r' (precise |bool |uint64_t |float |uint |int |float16_t |f16vec2 )', l))
calls = [i for i, l in enumerate(body) if 'mma_e4m3_f16' in l]
a = calls[0] - 1 # '{ uint md0_, ...;'
b = next(i for i in range(calls[-1], len(body)) if body[i].strip() == '}')
if not body[a].strip().startswith('{ uint md0_'):
raise SystemExit('unexpected statement before the first product: ' + body[a])
pre, cut, tail = body[nd:a], body[a:b + 1], body[b + 1:]
if sum(l.count('{') - l.count('}') for l in pre):
raise SystemExit('the first product is inside a block of the translated main')
# values the tail takes from the cut region other than the last product's results
assigned = lambda ls: {mm_.group(1) for l in ls for mm_ in [re.match(r'\s*(\w+) = ', l)] if mm_}
restored = {var(r) for mm_ in last for r in mm_[2]}
rmap = {} # dec0_1: the register holding r for (row offset, column offset): y = r + (D half) * s2 in the tail
if dec:
place = {}
for mm_, (i, j, s_) in zip(last, tiles(last)):
for q, r in enumerate(mm_[2]):
place[r] = (16 * i + 8 * q, 8 * j)
text_tail = '\n'.join(lines[last[-1][1]:])
half = {}
for mo in re.finditer(r'mov\.b32\s+\{(%rs\d+),\s*(%rs\d+)\},\s*(%r\d+);', text_tail):
if mo.group(3) in place:
half[mo.group(1)] = (*place[mo.group(3)], 0)
half[mo.group(2)] = (*place[mo.group(3)], 1)
prod = {}
for mo in re.finditer(r'mul\.f16\s+(%rs\d+),\s*(%rs\d+),\s*(%rs\d+);', text_tail):
if mo.group(2) in half:
prod[mo.group(1)] = half[mo.group(2)]
for mo in re.finditer(r'add\.f16\s+(%rs\d+),\s*(%rs\d+),\s*(%rs\d+);', text_tail):
if mo.group(2) in prod:
rmap[var(mo.group(3))] = prod[mo.group(2)]
if len(rmap) != 2 * 2 * len(last):
raise SystemExit(f'found {len(rmap)} residual registers, expected {4 * len(last)}')
restored |= set(rmap)
used = set(re.findall(r'\b[a-z]+\d+\b', '\n'.join(tail)))
lost = sorted((assigned(cut) - assigned(pre) - restored) & used)
# (values assigned again in the tail before their use are fine; the kernel has none on the reference build)
keep = []
cut_set = assigned(cut)
for name in lost: # thread indices and the like: their assignment is carried over when it depends on nothing cut
line = next(l for l in reversed(cut) if re.match(rf'\s*{name} = ', l))
if set(re.findall(r'\b[a-z]+\d+\b', line.split('=', 1)[1])) & cut_set:
raise SystemExit(f'the tail needs {name}, which the cut region computes: {line.strip()}')
keep.append(' ' + line)
split = sys.argv[sys.argv.index('--split') + 1] if '--split' in sys.argv else None
xf_word = int(sys.argv[sys.argv.index('--xf-word') + 1]) if split else 0
if split == 'gather':
text, record = gather_pass(xf_word, '\n'.join(g[:m]), body[:nd], pre, tail, first, f16vars_of(body[:nd]))
Path(out).write_text(text)
Path(out).with_suffix('.record').write_text(str(record))
return
if dec:
first = []
store = [' if (rr2_valid) { // the first product\'s A fragments: lane 4g + t, register q -> row, 4 k',
' const uint l_ = gl_SubgroupInvocationID, g_ = l_ >> 2u, t_ = l_ & 3u;']
seen = set()
for mm_, (i, j, s) in zip(first, tiles(first) if first else []):
if (i, s) in seen:
continue
seen.add((i, s))
for q, r in enumerate(mm_[3]):
store.append(f' XF[gl_SubgroupID * 1024u + uninter({16 * i + 8 * (q & 1)}u + g_) * 16u + {(32 * s + 16 * (q >> 1)) // 4}u + t_] = {var(r)};')
store.append(' }')
load = [' { // the guide head as the last product\'s D fragments (rows g, g + 8; columns 2t, 2t + 1)',
' const uint l_ = gl_SubgroupInvocationID, g_ = l_ >> 2u, t_ = l_ & 3u;']
f16vars = set()
for l in body[:nd]:
if re.match(r'\s*(precise )?float16_t ', l):
f16vars |= set(re.findall(r'\b(rs\d+)\b', l))
for mm_, (i, j, s) in zip(last, tiles(last)):
for q, r in enumerate(mm_[2]):
row = f'({16 * i + 8 * q}u + g_)'
load.append(f' {var(r)} = HDV({row}, {8 * j}u + 2u * t_) | (HDV({row}, {8 * j + 1}u + 2u * t_) << 16u);')
for name, (ro, co, hf) in rmap.items():
v = f'RSV({ro}u + g_, {co + hf}u + 2u * t_)'
load.append(f' {name} = ' + (f'uint16BitsToFloat16(uint16_t({v}));' if name in f16vars else f'{v};'))
load.append(' }')
# the native block's source, without its #version line (this shader has its own)
block = (Path(__file__).resolve().parent / 'rr2_m.comp').read_text().split('\n', 1)[1]
if split == 'block': # the second pass: the block on the gather pass's record, then the tail on wave 0
recomputed, saved, record, where = split_layout(xf_word, pre, tail, f16vars)
restore = [' ' + (f'{v} = uint16BitsToFloat16(uint16_t(arena.w[rr2_rec + 1024u + {len(saved)}u * l_ + {n}u]));' if v in f16vars
else f'{v} = arena.w[rr2_rec + 1024u + {len(saved)}u * l_ + {n}u];') for n, v in enumerate(saved)]
load[2:2] = restore
Path(out).write_text(
head + '\n#extension GL_GOOGLE_include_directive : require\n#define RR2_EMBED 1\n#define RR2_GROUP 1\n#define RR2_XF_ARENA 1\n' +
block + '\nuint HDV(uint row, uint col) { return col < 4u ? uint(float16BitsToUint16(HD[gl_SubgroupID * 256u + row * 4u + col])) : 0u; }\n\n'
'void main()\n{\n' + '\n'.join(body[:nd]) + '\n' + where +
' const bool rr2_valid = gl_SubgroupID < 1u && gl_WorkGroupID.x < rr2_gx;\n'
' rr2_xf = rr2_rec;\n rr2_block();\n barrier();\n if (rr2_valid) {\n' +
' // the gather\'s index arithmetic the tail reads\n' + '\n'.join(recomputed) + '\n' +
'\n'.join(load + keep) + '\n' + '\n'.join(tail) + '\n }\n}\n')
return
wx = lambda t: t.replace('gl_WorkGroupID.x', f'({G}u * gl_WorkGroupID.x + gl_SubgroupID)')
text = (head + f'\n#extension GL_GOOGLE_include_directive : require\n#define RR2_EMBED 1\n#define RR2_GROUP {G}\n' + block + '\n' +
('uint HDV(uint row, uint col) { return uint(float16BitsToUint16(MSH[gl_SubgroupID * 2048u + row * 32u + col])); }\n'
'uint RSV(uint row, uint col) { return uint(float16BitsToUint16(RSH[gl_SubgroupID * 2048u + row * 32u + col])); }\n\n' if dec else
'uint HDV(uint row, uint col) { return col < 4u ? uint(float16BitsToUint16(HD[gl_SubgroupID * 256u + row * 4u + col])) : 0u; }\n\n') +
'void main()\n{\n' + '\n'.join(body[:nd]) +
'\n // this wave\'s window, when the level has it\n'
' const uint rr2_gx = (param.w[0] + (param.w[20] & 0xffffu) + 7u) / 8u;\n'
f' const bool rr2_valid = gl_SubgroupID < {G}u && {G}u * gl_WorkGroupID.x + gl_SubgroupID < rr2_gx;\n'
' if (rr2_valid) { // NVIDIA\'s warp: the translation before the first product\n' + wx('\n'.join(pre)) + '\n }\n' +
('' if dec else '\n'.join(store)) +
# (the block runs unconditionally, also for a window past the level's last one: it writes nothing there, and
# inside a branch the driver compiles it much slower)
('\n barrier();\n rr2_block();\n barrier();\n' if G == 1 else
f'\n barrier();\n for (uint w_ = 0u; w_ < {G}u; ++w_) {{\n rr2_win = w_;\n'
' rr2_block();\n barrier();\n }\n') + ' if (rr2_valid) {\n' +
'\n'.join(load + keep) + '\n' + wx('\n'.join(tail)) + '\n }\n}\n')
Path(out).write_text(text)
print(f'{out}: {len(pre)} gather lines, {len(cut)} cut, {len(tail)} tail lines, {len(seen)} A fragment sets')
if __name__ == '__main__':
main()
engine/rr/phase_dump.py
+6 -3
View File
@@ -29,7 +29,8 @@ REG = r'%r\d+'
def kernel_module(kernel):
for p in sorted(Path('~/.cache/d4r-rr-corpus').expanduser().glob('*.ptx')):
corpus = '~/.cache/d4r-rr-corpus-f' if kernel.startswith('rr2_') else '~/.cache/d4r-rr-corpus'
for p in sorted(Path(corpus).expanduser().glob('*.ptx')):
if f'.entry {kernel}(' in p.read_text(errors='replace'):
return p
raise SystemExit(f'no corpus module for {kernel}')
@@ -80,7 +81,9 @@ def main():
o = ap.parse_args()
# the capture decides the kernel: a preset D capture runs cuda_dldn_engine_swin_{enc0,dec0}_diamond_wallaby_kernel
cap = Capture(os.path.expanduser(o.capture))
launch = [l for l in cap.launches if f'swin_{o.stage}' in l.kernel][o.frame - 1]
# preset F: STAGE is the kernel's full name (rr2_enc1_kernel, rr2_dec1_kernel_0)
rr2 = o.stage.startswith('rr2_')
launch = [l for l in cap.launches if (l.kernel == o.stage if rr2 else f'swin_{o.stage}' in l.kernel)][o.frame - 1]
kernel = launch.kernel
module = kernel_module(kernel)
lines = module.read_text().replace('\r', '').split('\n')
@@ -126,7 +129,7 @@ def main():
cmd = ['systemd-run', '--user', '--scope', '-q', '-p', 'MemoryMax=24G', '-p', 'MemorySwapMax=0']
cmd += [f'--setenv={k_}={v_}' for k_, v_ in env.items() if k_.startswith('D4R_ZLUDA')]
r = subprocess.run(cmd + [sys.executable, str(HERE / 'zluda_replay.py'), os.path.expanduser(o.capture),
f'swin_{o.stage}' + ('_diamond' if 'diamond' in kernel else ''), str(o.frame), str(work),
kernel if rr2 else f'swin_{o.stage}' + ('_diamond' if 'diamond' in kernel else ''), str(o.frame), str(work),
'--ptx', str(work / 'probe.ptx'),
'--pad', str(pad)], capture_output=True, text=True)
if 'failed' in r.stdout + r.stderr:
engine/rr/ptx2rr.py
+169 -19
View File
@@ -14,6 +14,7 @@ model/build directory only; do not commit or redistribute it.
import argparse
import importlib.util
import json
import os
import re
import struct
import sys
@@ -110,6 +111,78 @@ void sust_p4(writeonly image2D s, int x, int y, vec4 v)
if (uint(x) < uint(n.x) && uint(y) < uint(n.y)) imageStore(s, ivec2(x, y), v);
}
void sust_p1(writeonly image2D s, int x, int y, float v) { sust_p4(s, x, y, vec4(v, 0.0, 0.0, 0.0)); }
// Formatted stores to normalized surfaces, as the reference converts them (ZLUDA's surface_formatted_store_bits): NaN
// to zero, clamp, round half away from zero in f32. The quotient is what the image store then converts back exactly.
// UNORM 10:10:10:2 surfaces (the pooled histories) are RGBA32F images holding the quantised values, like ZLUDA's.
// tanh.approx.f32 as the reference computes it: ROCm's OCML tanh (device-libs ocml/src/tanhF.cl) with LLVM's AMDGPU
// lowering of exp (lowerFEXP, fast-FMA path): the same fused operations, hardware exp2, ldexp and reciprocal.
// (a fused f32 a * b + c through f64: the product is exact there and the sum is rounded to 53 bits before the f32
// rounding. The driver's own f32 fma is a product and a sum whenever its product also exists on its own, as in
// exp's ph / fma(x, c, -ph) pair.)
float fmad(float a, float b, float c) { precise double s = double(a) * double(b) + double(c); return float(s); }
float tanh_ocml(float x)
{
const float y = abs(x);
precise float z;
if (y < 0.625) {
precise float y2 = y * y;
precise float p = fmad(y2, fmad(y2, fmad(y2, fmad(y2, -0.005700020585209131, 0.02063407190144062), -0.053737930953502655), 0.13331416249275208), -0.3333328068256378);
precise float yp = y * p;
z = fmad(y2, yp, y);
} else {
const float w = 2.0 * y;
precise float ph = w * 1.4426950216293335;
precise float f0 = fmad(w, 1.4426950216293335, -ph);
precise float pl = fmad(w, 1.925962855864327e-08, f0);
const float e = roundEven(ph);
precise float d = ph - e;
precise float a = d + pl;
float t = ldexp(exp2(a), int(e));
if (w > 88.72283172607422) t = uintBitsToFloat(0x7f800000u);
precise float s1 = t + 1.0;
z = fmad(-2.0, rcp_a(s1), 1.0);
}
return x < 0.0 || floatBitsToUint(x) == 0x80000000u ? -z : z;
}
// tanh.approx.f16: from a table of the correctly rounded values indexed by the argument's bits (RR_TANH_TABLE: its
// byte offset in the weight buffer; engine/rr/rr2_pack.py) where a stage has one, else through f32
#ifdef RR_TANH_TABLE
float16_t tanh16(float16_t x) { return uint16BitsToFloat16(wbuf.wh[RR_TANH_TABLE / 2u + uint(float16BitsToUint16(x))]); }
#else
float16_t tanh16(float16_t x) { return float16_t(tanh(float(x))); }
#endif
vec4 q_snorm8(vec4 v)
{
v = clamp(mix(v, vec4(0.0), isnan(v)), vec4(-1.0), vec4(1.0));
precise vec4 s = v * 127.0;
precise vec4 a = abs(s) + 0.5;
return sign(s) * floor(a) / 127.0;
}
vec4 q_unorm8(vec4 v)
{
v = clamp(mix(v, vec4(0.0), isnan(v)), vec4(0.0), vec4(1.0));
precise vec4 a = v * 255.0 + 0.5;
return floor(a) / 255.0;
}
// (the quotient q / 1023 must be the correctly rounded f32 the reference stores: the driver divides by multiplying
// with a reciprocal, one step off for some q, so it comes from a table)
const float Q1023[1024] = float[1024](Q1023_VALUES);
vec4 q_unorm1010102(vec4 v)
{
v = clamp(mix(v, vec4(0.0), isnan(v)), vec4(0.0), vec4(1.0));
precise vec4 a = v * vec4(1023.0, 1023.0, 1023.0, 3.0) + 0.5;
const ivec4 q = ivec4(floor(a));
const float third[4] = float[4](0.0, 0.3333333432674408, 0.6666666865348816, 1.0);
return vec4(Q1023[q.x], Q1023[q.y], Q1023[q.z], third[q.w]);
}
// for a real UNORM 10:10:10:2 image: the reference's rounding (half away from zero) here, and any value within half a
// step of q / 1023 for the store, which the hardware quantises back to q
vec4 q_unorm1010102_hw(vec4 v)
{
v = clamp(mix(v, vec4(0.0), isnan(v)), vec4(0.0), vec4(1.0));
precise vec4 a = v * vec4(1023.0, 1023.0, 1023.0, 3.0) + 0.5;
return floor(a) * vec4(1.0 / 1023.0, 1.0 / 1023.0, 1.0 / 1023.0, 1.0 / 3.0);
}
// Binary stores write the halves' bits. RADV folds float(float16_t(x)) back to x when the half came from a
// cvt.rn.f16.f32 in a register, and the image store then truncates x instead (the output stage's history surface
// was one binary16 ulp low in half its values, which the next frame's encoder 0 amplifies): OR-ing in a zero the
@@ -135,7 +208,7 @@ uint shfl_dir(uint v, uint delta, uint opts, bool down, out bool in_bounds)
bool oob = down ? idx > sub_end : idx < sub_end;
if (oob) idx = self;
in_bounds = !oob;
return subgroupShuffle(v, uint(idx) & 31u);
return subgroupShuffle(v, (gl_SubgroupInvocationID & 32u) | (uint(idx) & 31u)); // the warp is this half of a 64-wide wave
}
// shfl.sync.idx (PTX ISA): j = minLane | (b & ~segmask), in bounds while j <= maxLane, otherwise the lane's own value
@@ -144,7 +217,7 @@ uint shfl_idx(uint v, uint b, uint opts, out bool in_bounds)
uint mask = (opts >> 8) & 31u, lane = gl_SubgroupInvocationID & 31u;
uint maxLane = (lane & mask) | (opts & 31u & ~mask), j = (lane & mask) | (b & 31u & ~mask);
in_bounds = j <= maxLane;
return subgroupShuffle(v, in_bounds ? j : lane);
return subgroupShuffle(v, (gl_SubgroupInvocationID & 32u) | (in_bounds ? j : lane));
}
// movmatrix.sync.trans.aligned.m8n8.b16 on one b32 fragment register per lane: lane (g = lane>>2, t = lane&3)
// holds row g, columns 2t and 2t+1. The transpose gathers (row 2t, col g) and (row 2t+1, col g).
@@ -337,12 +410,16 @@ def base_helper(name):
return '\n'.join(lines[start:end + 1]) + '\n'
RR_RUNTIME = RR_RUNTIME.replace('Q1023_VALUES', ', '.join(repr(float(np.float32(q) / np.float32(1023))) for q in range(1024)))
RR_RUNTIME += base_helper('prmt') + base_helper('shfl_bfly').replace(
'return subgroupShuffle(v, (gl_SubgroupInvocationID & 32u) | uint(idx));', 'return subgroupShuffle(v, uint(idx) & 31u);')
'return subgroupShuffle(v, (gl_SubgroupInvocationID & 32u) | uint(idx));',
'return subgroupShuffle(v, (gl_SubgroupInvocationID & 32u) | (uint(idx) & 31u));')
class RRT(base.T):
def __init__(self, ptx, kernel, block):
def __init__(self, ptx, kernel, block, weights_arg=40, buffers=()):
self.weights_arg = weights_arg
self.buffers = sorted(buffers) # argument offsets of further buffers, each its own binding
bx, by, bz = block
if bx * by * bz % 32:
raise base.Err('block size must be a multiple of the 32-lane subgroup')
@@ -367,10 +444,18 @@ class RRT(base.T):
# without it the compiler may contract a product and a sum across them depending on scheduling (results then
# changed with unrelated code and with the device's robustness features)
self.precise_all = True
# module globals with an initialiser (the output filter's bicubic phase table): constants, laid out one after
# another in a blob that is the stage's weight buffer (binding 1); `mov.u64 %rd, NAME` is its offset there
self.globals, self.globals_blob = {}, b''
for gm in re.finditer(r'\.global\s+\.align\s+\d+\s+\.b8\s+(\w+)\[(\d+)\]\s*=\s*\{([^}]*)\}', ptx):
self.globals[gm.group(1)] = len(self.globals_blob)
data = bytes(int(x) for x in gm.group(3).split(','))
self.globals_blob += data + bytes(-len(data) % 16)
self.stmt_lines = statement_lines(ptx, kernel)
self.stmt_index = 0
super().__init__(ptx, kernel)
self.pair_mma()
if not os.environ.get('D4R_PTX2RR_NO_PAIR'):
self.pair_mma()
self.origins()
def stmt(self, s):
@@ -438,9 +523,12 @@ class RRT(base.T):
if not a or not re.fullmatch(r'%rd\d+', a[0]) or op.startswith(('ld.', 'st.', 'bra')):
if op != 'ld.param.u64':
continue
if op == 'ld.param.u64':
if op.startswith('mov') and len(a) > 1 and a[1].strip() in self.globals:
new = {'w'}
elif op == 'ld.param.u64':
# the argument's true offset (the load may go through a register based inside the block)
new = {'w' if self.param_off(a[1]) == 40 else 'a'}
off_ = self.param_off(a[1])
new = {'w' if off_ == self.weights_arg else f'x{off_}' if off_ in self.buffers else 'a'}
else:
new = set()
for x in a[1:]:
@@ -457,21 +545,29 @@ class RRT(base.T):
raise base.Err(f'global pointer {reg} has origin {sorted(s)} (need exactly one of w/a)')
return next(iter(s))
def gfn(self, reg, store=False):
o = self.gorigin(reg)
if o == 'w':
return 'GW'
if o == 'a':
return 'SA' if store else 'GA'
return ('SX' if store else 'GX') + o[1:] + '_'
def gaddr(self, tok):
m = re.fullmatch(r'\[\s*(%rd\d+)\s*(?:\+\s*(\d+))?\s*\]', tok.strip())
m = re.fullmatch(r'\[\s*(%rd\d+)\s*(?:\+\s*(-?\d+))?\s*\]', tok.strip())
if not m:
raise base.Err(f'unsupported global address {tok}')
return m.group(1), int(m.group(2) or 0)
def gbyte(self, reg, off):
return f'uint({self.val(reg, "u64")}) + {off}u'
return f'uint({self.val(reg, "u64")}) + {off}u' if off >= 0 else f'uint({self.val(reg, "u64")}) - {-off}u'
def ld(self, o, a):
if 'global' not in o:
return super().ld(o, a)
size, count = mem_shape(o)
reg, off = self.gaddr(a[1])
fn = 'GW' if self.gorigin(reg) == 'w' else 'GA'
fn = self.gfn(reg)
dsts = braces(a[0]) if count > 1 else [a[0]]
for i, d in enumerate(dsts):
self.A(d, f'{fn}{size * 8}({self.gbyte(reg, off + i * size)})', 'u16' if size == 2 else 'u32')
@@ -481,14 +577,15 @@ class RRT(base.T):
return super().st(o, a)
size, count = mem_shape(o)
reg, off = self.gaddr(a[0])
if self.gorigin(reg) != 'a':
raise base.Err('store through the read-only weight pointer')
if self.gorigin(reg) == 'w':
raise base.Err(f'store through the read-only weight pointer: {a}')
fn = self.gfn(reg, True)
vals = braces(a[1]) if count > 1 else [a[1]]
for i, v in enumerate(vals):
if size == 2:
self.E(f'SA16({self.gbyte(reg, off + i * size)}, {self.val(v, "u16")});')
self.E(f'{fn}16({self.gbyte(reg, off + i * size)}, {self.val(v, "u16")});')
else:
self.E(f'SA32({self.gbyte(reg, off + i * size)}, {self.val(v, "u32")});')
self.E(f'{fn}32({self.gbyte(reg, off + i * size)}, {self.val(v, "u32")});')
# ------------------------------------------------ textures and surfaces
def handle_off(self, h):
@@ -515,7 +612,7 @@ class RRT(base.T):
# (zluda_ptx_impl.cpp fetch_2D). The output stage's specular reprojection reads such neighbours at the image
# border: with clamp its surfaces were 99.92% identical to the capture, with zero 99.997% (capDX, 2026-10-12)
fetch = f'fetch0(T{off}, ivec2({self.val(c[0], "s32")}, {self.val(c[1], "s32")}))'
elif 'base' in o and o[-1] == 'f32':
elif ('base' in o or o[0] == '2d') and o[-1] == 'f32': # tex.2d: no derivatives in a kernel, level 0
fetch = f'textureLod(T{off}, vec2({self.val(c[0], "f32")}, {self.val(c[1], "f32")}), 0.0)'
elif 'level' in o and o[-1] == 'f32':
fetch = (f'textureLod(T{off}, vec2({self.val(c[0], "f32")}, {self.val(c[1], "f32")}), '
@@ -525,7 +622,9 @@ class RRT(base.T):
f16 = 'f16' in o
self.E(f'{{ vec4 t_ = {fetch};')
for i, d in enumerate(dst):
self.A(d, f'float16_t(t_[{i}])' if f16 else f't_[{i}]', 'f16' if f16 else 'float')
# a binary16 result is the texel rounded to binary16: through H(HU(...) | SUST_Z) so that the driver cannot
# fold the rounding away when the value is widened again (RR_RUNTIME's note on sust_b)
self.A(d, f'H(HU(float16_t(t_[{i}])) | SUST_Z)' if f16 else f't_[{i}]', 'f16' if f16 else 'float')
self.E('}')
def sust(self, op, a):
@@ -534,9 +633,10 @@ class RRT(base.T):
v = braces(a[1])
x, y = self.val(c[0], 's32'), self.val(c[1], 's32')
S = f'S{off}'
op = op.replace('.cs.', '.') # the streaming cache policy: a hint, the store is the same
if op == 'sust.p.2d.v4.b32.zero':
self.set_surface(off, 'rgba16f')
self.E(f'sust_p4({S}, {x}, {y}, vec4({", ".join(self.val(q, "f32") for q in v)}));')
self.E(f'sust_p4({S}, {x}, {y}, SUSTQ_{off}(vec4({", ".join(self.val(q, "f32") for q in v)})));')
elif op == 'sust.p.2d.b32.zero':
self.set_surface(off, 'r16f')
self.E(f'sust_p1({S}, {x}, {y}, {self.val(v[0], "f32")});')
@@ -579,6 +679,13 @@ class RRT(base.T):
types = [t for t in o if re.match(r'^(f|s|u|b)\d+(x2)?$', t)]
if types[:2] in (['u32', 'u64'], ['s32', 's64'], ['u32', 's64'], ['s32', 'u64']):
return self.A(a[0], f'uint({self.val(a[1], "u64")})', 'u32')
if types[:2] == ['f16', 'f32'] and 'rn' in o and 'satfinite' not in o:
# Two things the driver must not do. It folds float(float16_t(x)) to x when the result is widened again
# (RR_RUNTIME's note on sust_b): the OR on the binary16 bits. And it fuses a product or sum with this
# conversion into v_fma_mix, one rounding of the exact value, where the kernel rounds to f32 first: an f32
# result that is exactly a binary16 midpoint then rounds the other way (the output stages' rare one-step
# differences, e.g. log2(x) * c1 * c2; found 2026-10-13): the OR on the f32 bits keeps the f32 value.
return self.A(a[0], f'H(HU(float16_t(uintBitsToFloat(floatBitsToUint({self.val(a[1], "f32")}) | SUST_Z))) | SUST_Z)', 'f16')
return super().cvt(op, o, a)
def f64_op(self, op, a):
@@ -610,6 +717,22 @@ class RRT(base.T):
return self.E('}')
if op == 'mma.paired.noop':
return
if parts[0] == 'mov' and len(a) > 1 and a[1].strip() in self.globals:
return self.A(a[0], f'uint64_t({self.globals[a[1].strip()]}u)', 'u64')
# cross-kernel progress counters (NGX runs the output passes concurrently and has each wait for the rows it
# needs): plain loads and adds here. The engine runs the passes one after another and leaves the counter
# pointers null, which the kernels test, so these never execute.
if parts[0] == 'ld' and ('relaxed' in parts or 'acquire' in parts) and 'global' in parts:
return self.ld([q for q in parts[1:] if q not in ('relaxed', 'acquire', 'gpu') and '::' not in q and not q.startswith('L1')], a)
if parts[0] == 'red' and parts[-2] == 'add':
reg, off = self.gaddr(a[0])
fn, sfn = self.gfn(reg), self.gfn(reg, True)
return self.E(f'{sfn}32({self.gbyte(reg, off)}, {fn}32({self.gbyte(reg, off)}) + {self.val(a[1], "u32")});')
if op == 'membar.gl':
return
if op == 'activemask.b32':
# the warp's lanes: this half of a 64-wide wave
return self.A(a[0], '((gl_SubgroupInvocationID & 32u) != 0u ? subgroupBallot(true).y : subgroupBallot(true).x)', 'u32')
if op == 'movmatrix.sync.trans.aligned.m8n8.b16':
return self.A(a[0], f'movm_trans({self.val(a[1], "u32")})', 'u32')
# PTX clamps shift amounts at the register width (shl/shr.u give 0, shr.s the sign); GLSL leaves them undefined
@@ -635,6 +758,16 @@ class RRT(base.T):
return self.tex(op, a)
if parts[0] == 'cvt' and [t for t in parts[1:] if t in ('f16x2', 'e4m3x2')] == ['e4m3x2', 'f16x2']:
return self.cvt_to_e4m3(op, a)
if op == 'copysign.f32': # PTX: the sign of the first source on the value of the second
x, y = self.val(a[1], 'f32'), self.val(a[2], 'f32')
return self.A(a[0], f'uintBitsToFloat((floatBitsToUint({y}) & 0x7fffffffu) | (floatBitsToUint({x}) & 0x80000000u))', 'float')
if op == 'tanh.approx.f16':
return self.A(a[0], f'tanh16({self.val(a[1], "f16")})', 'f16')
if op == 'tanh.approx.f32':
return self.A(a[0], f'tanh_ocml({self.val(a[1], "f32")})', 'float')
if op == 'ex2.approx.f16':
# v_exp_f32 of the widened value, rounded to binary16; the ORs keep that rounding (see cvt)
return self.A(a[0], f'H(HU(float16_t(uintBitsToFloat(floatBitsToUint(exp2(float({self.val(a[1], "f16")}))) | SUST_Z))) | SUST_Z)', 'f16')
if op == 'neg.f16':
return self.A(a[0], f'-{self.val(a[1], "f16")}', 'f16')
if parts[0] in ('min', 'max') and 'NaN' not in parts and any(t in parts for t in ('f16', 'f16x2', 'f32')):
@@ -659,6 +792,7 @@ class RRT(base.T):
w('#extension GL_EXT_shader_16bit_storage : require')
w('#extension GL_KHR_shader_subgroup_basic : require')
w('#extension GL_KHR_shader_subgroup_shuffle : require')
w('#extension GL_KHR_shader_subgroup_ballot : require')
w('#extension GL_KHR_cooperative_matrix : require')
w('#extension GL_KHR_memory_scope_semantics : require')
w('#extension GL_EXT_float_e4m3 : require')
@@ -673,6 +807,15 @@ class RRT(base.T):
n += 1
for off in sorted(self.surfaces):
w(f'layout(set = 0, binding = {n}, {self.surfaces[off]}) uniform writeonly image2D S{off};')
q = {'rgba8_snorm': 'q_snorm8', 'rgba8': 'q_unorm8', 'rgba32f': 'q_unorm1010102', 'rgb10_a2': 'q_unorm1010102_hw'}.get(self.surfaces[off], '')
w(f'#define SUSTQ_{off}(v) {q}(v)')
n += 1
for off in self.buffers:
w(f'layout(set = 0, binding = {n}, std430) buffer XBuf{off} {{ uint16_t h[]; }} xb{off};')
w(f'uint GX{off}_16(uint a) {{ return uint(xb{off}.h[a >> 1]); }}')
w(f'uint GX{off}_32(uint a) {{ return uint(xb{off}.h[a >> 1]) | (uint(xb{off}.h[(a >> 1) + 1u]) << 16); }}')
w(f'void SX{off}_16(uint a, uint v) {{ xb{off}.h[a >> 1] = uint16_t(v & 0xffffu); }}')
w(f'void SX{off}_32(uint a, uint v) {{ xb{off}.h[a >> 1] = uint16_t(v & 0xffffu); xb{off}.h[(a >> 1) + 1u] = uint16_t(v >> 16); }}')
n += 1
w(f'shared uint sh[{max(1, (self.shared_bytes + 3) // 4)}];')
if self.probe_slots or self.oob:
@@ -694,6 +837,11 @@ def main():
ap.add_argument('--surface', action='append', default=[], metavar='OFFSET=FORMAT',
help="the GLSL image format of a surface whose stores are formatted (sust.p), when it is not the "
"store's own guess, e.g. 144=rgba8_snorm (repeatable)")
ap.add_argument('--weights-arg', type=int, default=40,
help='argument offset of the weight pointer (40 for the DLSS-D kernels, 24 for rr2_*)')
ap.add_argument('--buffer', action='append', type=int, default=[], metavar='OFFSET',
help='a pointer argument that is neither the weights nor the arena: its own storage buffer, bound '
'after the surfaces in ascending argument offset (repeatable)')
ap.add_argument('--probe', action='append', default=[], metavar='LINE:%REG,...',
help='debugging: store the registers after the statement at that module line executes (last '
'execution) for workgroup --probe-block, into a debug buffer bound after the surfaces')
@@ -709,7 +857,7 @@ def main():
probes[int(line)] = [r.strip() for r in regs.split(',')]
def make():
t = RRT(src, ns.kernel, block)
t = RRT(src, ns.kernel, block, ns.weights_arg, ns.buffer)
t.probes = probes
t.probe_block = tuple(int(x) for x in ns.probe_block.split(','))
t.oob = ns.oob_report
@@ -719,10 +867,12 @@ def main():
if off in t.surfaces:
t.surfaces[off] = fmt
open(ns.out, 'w').write(t.emit({}))
if t.globals_blob:
Path(ns.out).with_suffix('.globals.bin').write_bytes(t.globals_blob)
# the binding order (stage contract): textures, then surfaces, in ascending argument offset
Path(ns.out).with_suffix('.json').write_text(json.dumps(
{'kernel': ns.kernel, 'block': block, 'textures': sorted(t.textures),
'surfaces': [[off, t.surfaces[off]] for off in sorted(t.surfaces)], 'shared': t.shared_bytes,
'surfaces': [[off, t.surfaces[off]] for off in sorted(t.surfaces)], 'buffers': t.buffers, 'shared': t.shared_bytes,
'probes': [[line, reg] for line in sorted(probes) for reg in probes[line]], 'oob': ns.oob_report}) + '\n')
sys.stderr.write(f'block={block} textures={sorted(t.textures)} surfaces={t.surfaces} shared={t.shared_bytes}\n')
engine/rr/rr2_m.comp
+530
View File
@@ -0,0 +1,530 @@
#version 460
// Native stage of the DLSS 4.5 Ray Reconstruction network (preset F, the rr2_* kernels of nvngx_dlssd.dll 310.9):
// one shader for enc1 ... enc5 and both kernels of dec4 ... dec1 and dec0's first. The arithmetic is
// engine/rr/rr2_ref.py's (which documents the block); the shape and the weight offsets come from the include
// generated by engine/rr/rr2_pack.py:
// KIND 0 encoder (tokens from the arena), 1 enc1 (first merges encoder 0's tokens 2x2), 2 a decoder's first kernel
// (expands the CL channels one level down and adds the skip), 3 a decoder's second kernel (with the guide blend),
// 4 encoder 0 after its per-pixel feature gather: included (RR2_EMBED) by the shader engine/rr/gen_rr2_enc0.py
// generates around ptx2rr.py's translation of the gather; it embeds the 64 feature slots of XF and leaves
// the guide head in HD for the translated tail; 5 dec0's second kernel, included likewise by the shader
// generated around the translation of its output tail: it leaves r and the MLP sum in RSH / MSH
// C channels (MT = C / 16 tiles), NH heads of 32 channels, NB MLP branches, CM merged channels for the next
// encoder (0: none), GUIDE 1 when the stage gates its guide image (KIND 0 and 1)
// Like engine/rr/swin_m.comp, every product runs transposed (weights x activations), so results stay in registers as
// [channel tile][token] accumulators: one workgroup per 8x8-token window, four waves, wave w owning the token slots
// 16w ... 16w + 15; token and channel slot j holds the natural index inter(j) (NVIDIA's k order, which is also the
// arena's byte order within a 32-channel plane).
#extension GL_KHR_cooperative_matrix : require
#extension GL_KHR_memory_scope_semantics : require
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_8bit_storage : require
#extension GL_EXT_float_e4m3 : require
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_EXT_control_flow_attributes : require
#extension GL_GOOGLE_include_directive : require
#ifndef RR2_EMBED
layout(local_size_x = 128) in;
#endif
#include "rr2_stage.h"
#include "rr_exact.glsl"
#define MT (C / 16)
#ifdef RR2_EMBED
#define PRM(i) param.w[i]
// the generated shader's workgroup holds RR2_GROUP (1 to 4) windows along x (wave w runs the translated code of
// window w, then the block runs once per window): RR2_WIN is the one this call works on
#if RR2_GROUP == 1
#define RR2_WIN 0u
#define RR2_WGX gl_WorkGroupID.x
#else
uint rr2_win;
#define RR2_WIN rr2_win
#define RR2_WGX (uint(RR2_GROUP) * gl_WorkGroupID.x + rr2_win)
#endif
#else
#define RR2_GROUP 1
#define RR2_WIN 0u
#define RR2_WGX gl_WorkGroupID.x
layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param;
#define PRM(i) param.p[i]
#endif
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16;
layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u;
layout(set = 0, binding = 1, std430) readonly buffer W8U { uint8_t w[]; } w8u;
layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define HAC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CMJ gl_CooperativeMatrixLayoutColumnMajor
#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM)
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
#define Q8(dst, src) Q8PAIRS(dst, src)
// One 32-deep step onto an accumulator of binary16 values: the reference adds it inside the f32 products and rounds
// once (swin_m.comp explains RR_MIX, the one-rounding variant that is not bit-exact by construction)
#ifdef RR_MIX
#define RAC HAC
#define RV(x) (x)
#define K32(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(0.0))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_] + float((acc)[i_])); }
#define K32H(acc, a0, b0, a1, b1) K32(acc, a0, b0, a1, b1)
#else
#define RAC ACC
#define RV(x) float(x)
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) }
#define K32H(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(acc))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_]); }
#endif
const uint HEAD = W_WQK1 - W_WQK0, BRANCH = W_WD1 - W_WD0;
shared uint XT[16 * C]; // the window's tokens: [token slot][C FP8 channels in slot order]; later the merge rows
shared uint KVS[2 * 512]; // the current head's QK ([token slot][32]), then V^T ([32][token slot])
#if KIND == 1
shared uint GT[64 * 32]; // the merge input: [token slot][4 sub-positions x 32 FP8 channels]
#endif
#if KIND == 2
shared uint LOW[4 * CL]; // the 4 x 4 tokens one level down: [row-major][CL FP8 channels]
shared uint EX[16 * C]; // their expansion: [low token][sub-position][C FP8 channels in slot order]
shared float16_t XE[64 * C]; // the block's input before its FP8 conversion: [token slot][channel slot]
#endif
#if KIND == 4
#ifdef RR2_XF_ARENA
// the features come from the gather pass (engine/rr/gen_rr2_enc0.py --split): the window's record in the arena, as
// the index of its first 32-bit word, set by the stage's main
uint rr2_xf;
#define XF_SRC arena.w
#define XF_AT rr2_xf
#else
shared uint XF[RR2_GROUP * 64 * 16]; // per window, the features from the gather: [token slot][64 FP8 feature slots]
#define XF_SRC XF
#define XF_AT (RR2_WIN * 1024u)
#endif
shared float16_t HD[RR2_GROUP * 64 * 4]; // per window, the guide head per token (raster order), for the translated tail
const uint EBRANCH = W_EWD1 - W_EWD0;
#endif
#if KIND == 5
shared float16_t RSH[RR2_GROUP * 64 * 32], MSH[RR2_GROUP * 64 * 32]; // per window, r and the MLP sum per token (raster order) and natural channel
#endif
#if GUIDE || KIND == 3
shared float16_t OUTG[64 * 4]; // the guide output per window token (raster order)
#endif
uint pack4(floate4m3_t a, floate4m3_t b, floate4m3_t c, floate4m3_t d)
{
return code(a) | (code(b) << 8u) | (code(c) << 16u) | (code(d) << 24u);
}
float16_t e4(uint w, uint b) { return float16_t(uintBitsToFloate4m3EXT(uint8_t((w >> (8u * b)) & 0xffu))); }
uvec2 mergeXY(uint r) { return r < 8u ? uvec2(r >> 1u, r & 1u) : uvec2((r - 8u) >> 1u, 2u + (r & 1u)); }
// four binary16 values of a guide image at token (x, y) of a level w wide
f16vec4 guideAt(uint off, int x, int y, int w)
{
const uint a = (off + uint(y * w + x) * 8u) >> 2u;
const uint u0 = arena.w[a], u1 = arena.w[a + 1u];
return f16vec4(hb(u0 & 0xffffu), hb(u0 >> 16u), hb(u1 & 0xffffu), hb(u1 >> 16u));
}
#ifdef RR2_EMBED
void rr2_block()
#else
void main()
#endif
{
const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID;
const uint hl = lane >> 4u, col = lane & 15u;
RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u);
const int ex = int(PRM(0)), ey = int(PRM(1)), hx = ex / 2, hy = ey / 2;
const uint inOff = PRM(2), skOff = PRM(4), outOff = PRM(8), mrgOff = PRM(10), shift = PRM(20);
const uint gInOff = PRM(12), gEncOff = PRM(14), gOutOff = PRM(16), gHalfOff = PRM(18);
const float16_t s0 = hb(PRM(23) & 0xffffu), s1 = hb(PRM(23) >> 16u), s2 = hb(PRM(24) & 0xffffu), s3 = hb(PRM(24) >> 16u);
const int x0 = 8 * int(RR2_WGX) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u);
const uint plane = uint(ex * ey) * 32u, lplane = uint(hx * hy) * 32u;
// this lane's own token (slot 16 wv + col)
const uint t = inter(16u * wv + col);
const int tx = x0 + int(t & 7u), ty = y0 + int(t >> 3u);
const bool inside = tx >= 0 && ty >= 0 && tx < ex && ty < ey;
// ---------------------------------------------------------------- the block's input x: one token per thread
#if KIND == 0 || KIND == 3 || KIND == 5
if (tid < 64u) {
const uint k = inter(tid);
const int xx = mirror(x0 + int(k & 7u), ex), yy = mirror(y0 + int(k >> 3u), ey);
const uint tb = (inOff + uint(yy * ex + xx) * 32u) / 4u;
[[unroll]] for (uint wd = 0u; wd < C / 4; ++wd) XT[tid * (C / 4u) + wd] = arena.w[tb + (wd >> 3u) * (plane >> 2u) + (wd & 7u)];
}
barrier();
#elif KIND == 4
// encoder 0: x = bu + sum over two branches of Wu . act(Wd . features + bd)
{
B8 fb[4];
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) coopMatLoad(fb[kt], XF_SRC, XF_AT + wv * 256u + 4u * kt, 16u, CMJ);
RAC X[MT];
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) X[m][i] = RV(w16.w[W_EBU / 2u + 16u * m + 8u * hl + uint(i)]);
[[unroll]] for (uint br = 0u; br < 2u; ++br) {
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
RAC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = RV(w16.w[(W_EBD0 + EBRANCH * br) / 2u + 16u * c + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < 4u; kt += 2u) {
A8 a0, a1;
LDA(a0, W_EWD0 + EBRANCH * br, 64u, c, kt); LDA(a1, W_EWD0 + EBRANCH * br, 64u, c, kt + 1u);
K32(h, a0, fb[kt], a1, fb[kt + 1u])
}
[[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = ACTQ(float16_t(h[i]));
}
[[unroll]] for (uint m = 0u; m < MT; ++m) {
A8 a0, a1;
LDA(a0, W_EWU0 + EBRANCH * br, 32u, m, 0u); LDA(a1, W_EWU0 + EBRANCH * br, 32u, m, 1u);
K32(X[m], a0, Hb[0], a1, Hb[1])
}
}
[[unroll]] for (uint m = 0u; m < MT; ++m) {
B8 q;
Q8(q, X[m])
coopMatStore(q, XT, wv * 4u * C + 4u * m, C / 4u, CMJ);
}
}
barrier();
#elif KIND == 1
// enc1: the 2x2 group of encoder 0's 32-channel tokens, then x = W1 . group + b1
if (tid < 64u) {
const uint k = inter(tid);
const int xx = mirror(x0 + int(k & 7u), ex), yy = mirror(y0 + int(k >> 3u), ey);
[[unroll]] for (uint q = 0u; q < 4u; ++q) {
const uint tb = (inOff + uint((2 * yy + int(q >> 1u)) * (2 * ex) + 2 * xx + int(q & 1u)) * 32u) / 4u;
[[unroll]] for (uint wd = 0u; wd < 8u; ++wd) GT[tid * 32u + q * 8u + wd] = arena.w[tb + wd];
}
}
barrier();
{
B8 gb[8];
[[unroll]] for (uint kt = 0u; kt < 8u; ++kt) coopMatLoad(gb[kt], GT, wv * 512u + 4u * kt, 32u, CMJ);
[[unroll]] for (uint m = 0u; m < MT; ++m) {
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_B1 / 2u + 16u * m + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < 8u; kt += 2u) {
A8 w0, w1;
LDA(w0, W_W1, 128u, m, kt); LDA(w1, W_W1, 128u, m, kt + 1u);
K32(a, w0, gb[kt], w1, gb[kt + 1u])
}
B8 q;
Q8(q, a)
coopMatStore(q, XT, wv * 4u * C + 4u * m, C / 4u, CMJ);
}
}
barrier();
#else
// a decoder's first kernel: wave w expands the 16 low tokens to sub-position w, then x = dq(e) + dq(skip) * s3
{
const int X0 = x0 >> 1, Y0 = y0 >> 1;
for (uint i = tid; i < 4u * CL; i += 128u) {
const uint r = i / (CL / 4u), wd = i % (CL / 4u);
const int X = mirror(X0 + int(r & 3u), hx), Y = mirror(Y0 + int(r >> 2u), hy);
LOW[i] = arena.w[(inOff + (wd >> 3u) * lplane + uint(Y * hx + X) * 32u) / 4u + (wd & 7u)];
}
barrier();
B8 lb[CL / 16];
[[unroll]] for (uint kt = 0u; kt < CL / 16; ++kt) coopMatLoad(lb[kt], LOW, 4u * kt, CL / 4u, CMJ);
for (uint mm = 0u; mm < MT; ++mm) {
const uint m = MT * wv + mm;
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_BE / 2u + 16u * m + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < CL / 16; kt += 2u) {
A8 w0, w1;
LDA(w0, W_WE, CL, m, kt); LDA(w1, W_WE, CL, m, kt + 1u);
K32(a, w0, lb[kt], w1, lb[kt + 1u])
}
B8 q;
Q8(q, a)
coopMatStore(q, EX, 4u * m, C, CMJ);
}
barrier();
if (tid < 64u) {
const uint k = inter(tid);
const int kx = x0 + int(k & 7u), ky = y0 + int(k >> 3u);
const uint tb = (skOff + uint(mirror(ky, ey) * ex + mirror(kx, ex)) * 32u) / 4u;
const uint r = uint(((ky >> 1) - Y0) * 4 + ((kx >> 1) - X0)), s = uint(2 * (ky & 1) + (kx & 1));
[[unroll]] for (uint wd = 0u; wd < C / 4; ++wd) {
const uint sk = arena.w[tb + (wd >> 3u) * (plane >> 2u) + (wd & 7u)];
const uint e = EX[r * C + s * (C / 4u) + wd];
floate4m3_t c[4];
[[unroll]] for (uint b = 0u; b < 4u; ++b) {
precise float16_t ms = e4(sk, b) * s3;
precise float16_t x = e4(e, b) + ms;
XE[tid * C + 4u * wd + b] = x;
c[b] = q8h(x);
}
XT[tid * (C / 4u) + wd] = pack4(c[0], c[1], c[2], c[3]);
}
}
barrier();
}
#endif
// ---------------------------------------------------------------- attention: the heads' out-projections sum up
RAC AT[MT];
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) AT[m][i] = RV(w16.w[W_BO / 2u + 16u * m + 8u * hl + uint(i)]);
[[dont_unroll]] for (uint hd = 0u; hd < NH; ++hd) {
barrier(); // every wave is done with the previous head's QK / V
for (uint j = 0u; j < 2u; ++j) { // 8 jobs (token tile n, channel tile c)
const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u;
RAC a = RAC(RV(float16_t(0.0))), v = RAC(RV(float16_t(0.0)));
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) {
B8 xb0, xb1, wv0, wv1;
A8 xa0, xa1, wk0, wk1;
coopMatLoad(xb0, XT, n * 4u * C + 4u * kt, C / 4u, CMJ);
coopMatLoad(xb1, XT, n * 4u * C + 4u * kt + 4u, C / 4u, CMJ);
coopMatLoad(xa0, XT, n * 4u * C + 4u * kt, C / 4u, RM);
coopMatLoad(xa1, XT, n * 4u * C + 4u * kt + 4u, C / 4u, RM);
LDA(wk0, W_WQK0 + HEAD * hd, C, c, kt); LDA(wk1, W_WQK0 + HEAD * hd, C, c, kt + 1u);
LDB(wv0, W_WV0 + HEAD * hd, 32u, kt, c); LDB(wv1, W_WV0 + HEAD * hd, 32u, kt + 1u, c);
K32(a, wk0, xb0, wk1, xb1)
K32(v, xa0, wv0, xa1, wv1)
}
A8 ka, va;
Q8(ka, a)
Q8(va, v)
coopMatStore(ka, KVS, n * 128u + 4u * c, 8u, RM);
coopMatStore(va, KVS, 512u + c * 256u + 4u * n, 16u, RM);
}
barrier();
// scores^T (keys x this wave's queries) = QK . QK^T + bias, the cubic softmax down each column
B8 Qb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) coopMatLoad(Qb[c], KVS, wv * 128u + 4u * c, 8u, CMJ);
B8 Pb[4];
{
float16_t e[4][8];
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) {
HAC c;
coopMatLoad(c, w16.w, (W_BIAS0 + HEAD * hd) / 2u + kt * 1024u + 16u * wv, 64u, RM);
A8 k0, k1;
coopMatLoad(k0, KVS, kt * 128u, 8u, RM);
coopMatLoad(k1, KVS, kt * 128u + 4u, 8u, RM);
K32H(c, k0, Qb[0], k1, Qb[1])
[[unroll]] for (int i = 0; i < 8; ++i) e[kt][i] = hb(uint(w16u.w[W_CUBIC / 2u + bh(c[i])]));
}
// key slot 16kt + 8hl + i is NVIDIA's column 8j + 2tt + half with j = 4(kt >> 1) + (kt & 1) + 2((i >> 1) & 1),
// tt = 2hl + (i >> 2), half = i & 1 (rr2_ref.softmax): pairs (e[2q], e[2q + 1]) are the two halves
f16vec2 ep[4][4];
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt)
[[unroll]] for (uint q = 0u; q < 4u; ++q) ep[kt][q] = f16vec2(e[kt][2u * q], e[kt][2u * q + 1u]);
precise f16vec2 st[2];
[[unroll]] for (uint tq = 0u; tq < 2u; ++tq) {
const uint q0 = 2u * tq, q1 = q0 + 1u;
precise f16vec2 ev = (ep[0][q0] + ep[0][q1]) + (ep[2][q0] + ep[2][q1]);
precise f16vec2 od = (ep[1][q0] + ep[1][q1]) + (ep[3][q0] + ep[3][q1]);
st[tq] = ev + od;
}
precise f16vec2 loc = st[0] + st[1];
const uint other = subgroupShuffleXor(packFloat2x16(loc), 16u);
// hl = 0 holds tt 0, 1 and hl = 1 tt 2, 3: the sum is (tt0 + tt1) + (tt2 + tt3) on both halves
precise f16vec2 o = hl == 0u ? loc + unpackFloat2x16(other) : unpackFloat2x16(other) + loc;
precise float16_t rs = o.x + o.y;
const float16_t r = float16_t(1.0 / float(rs));
const f16vec2 r2 = f16vec2(r, r);
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt)
[[unroll]] for (uint q = 0u; q < 4u; ++q) {
precise f16vec2 pv = ep[kt][q] * r2;
{ const fe4m3vec2 t_ = fe4m3vec2(pv); Pb[kt][2u * q] = t_.x; Pb[kt][2u * q + 1u] = t_.y; }
}
}
// O^T = V^T P^T, then the out-projection
B8 Ob[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
RAC o = RAC(RV(float16_t(0.0)));
[[unroll]] for (uint ks = 0u; ks < 2u; ++ks) {
A8 v0, v1;
coopMatLoad(v0, KVS, 512u + c * 256u + 8u * ks, 16u, RM);
coopMatLoad(v1, KVS, 512u + c * 256u + 8u * ks + 4u, 16u, RM);
K32(o, v0, Pb[2u * ks], v1, Pb[2u * ks + 1u])
}
Q8(Ob[c], o)
}
[[unroll]] for (uint m = 0u; m < MT; ++m) {
A8 a0, a1;
LDA(a0, W_WO0 + HEAD * hd, 32u, m, 0u); LDA(a1, W_WO0 + HEAD * hd, 32u, m, 1u);
K32(AT[m], a0, Ob[0], a1, Ob[1])
}
}
// ---------------------------------------------------------------- r = x + attention * s1; the MLP's input q8(r)
float16_t rv[MT][8];
B8 Nb[MT];
[[unroll]] for (uint m = 0u; m < MT; ++m) {
#if KIND != 2
const uint xa = (16u * wv + col) * (C / 4u) + 4u * m + 2u * hl;
const uvec2 xw = uvec2(XT[xa], XT[xa + 1u]);
#endif
[[unroll]] for (int i = 0; i < 8; ++i) {
#if KIND == 2
const float16_t xb = XE[(16u * wv + col) * C + 16u * m + 8u * hl + uint(i)];
#else
const float16_t xb = e4(xw[i >> 2], uint(i) & 3u);
#endif
precise float16_t sa = float16_t(AT[m][i]) * s1;
precise float16_t r = xb + sa;
rv[m][i] = r;
Nb[m][i] = q8h(r);
}
}
RAC ML[MT];
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) ML[m][i] = RV(w16.w[W_BU / 2u + 16u * m + 8u * hl + uint(i)]);
[[dont_unroll]] for (uint br = 0u; br < NB; ++br) {
const uint w1 = W_WD0 + BRANCH * br, b1 = W_BD0 + BRANCH * br, w2 = W_WU0 + BRANCH * br;
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
RAC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = RV(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) {
A8 a0, a1;
LDA(a0, w1, C, c, kt); LDA(a1, w1, C, c, kt + 1u);
K32(h, a0, Nb[kt], a1, Nb[kt + 1u])
}
[[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = ACTQ(float16_t(h[i]));
}
[[unroll]] for (uint m = 0u; m < MT; ++m) {
A8 a0, a1;
LDA(a0, w2, 32u, m, 0u); LDA(a1, w2, 32u, m, 1u);
K32(ML[m], a0, Hb[0], a1, Hb[1])
}
}
#if KIND == 5
// dec0's second kernel: the translated tail forms y = r + MLP * s2 itself and computes the outputs from it
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) {
const uint at = RR2_WIN * 2048u + t * 32u + inter(16u * m + 8u * hl + uint(i));
RSH[at] = rv[m][i];
MSH[at] = float16_t(ML[m][i]);
}
#else
// ---------------------------------------------------------------- y = r + MLP * s2: the stage's tokens
B8 Tb[MT];
{
const uint o = (outOff + uint(ty * ex + tx) * 32u) >> 2u;
#if CM
const uint ty2 = t >> 3u, tx2 = t & 7u;
const uint mr = (ty2 >> 1u) < 2u ? 2u * (tx2 >> 1u) + (ty2 >> 1u) : 8u + 2u * (tx2 >> 1u) + (ty2 >> 1u) - 2u;
const uint blk = 2u * (ty2 & 1u) + (tx2 & 1u);
barrier(); // every wave is done with XT (the merge rows reuse it)
#endif
[[unroll]] for (uint m = 0u; m < MT; ++m) {
[[unroll]] for (int i = 0; i < 8; ++i) {
precise float16_t sm = float16_t(ML[m][i]) * s2;
precise float16_t y = rv[m][i] + sm;
Tb[m][i] = q8h(y);
}
const uint w0 = pack4(Tb[m][0], Tb[m][1], Tb[m][2], Tb[m][3]), w1 = pack4(Tb[m][4], Tb[m][5], Tb[m][6], Tb[m][7]);
const uint at = 4u * (m & 1u) + 2u * hl;
if (inside) { arena.w[o + (m >> 1u) * (plane >> 2u) + at] = w0; arena.w[o + (m >> 1u) * (plane >> 2u) + at + 1u] = w1; }
#if CM
const uint g = mr * C + blk * (C / 4u) + 4u * m + 2u * hl;
XT[g] = w0; XT[g + 1u] = w1;
#endif
}
}
#endif
#if GUIDE || KIND == 3 || KIND == 4
{ // the guide head (rows 0-3 on lanes 0-15)
RAC hv;
[[unroll]] for (int i = 0; i < 8; ++i) hv[i] = RV(w16.w[W_BH / 2u + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) {
A8 a0, a1;
LDA(a0, W_WH, C, 0u, kt); LDA(a1, W_WH, C, 0u, kt + 1u);
K32(hv, a0, Tb[kt], a1, Tb[kt + 1u])
}
#if KIND == 4
if (hl == 0u) [[unroll]] for (uint c = 0u; c < 4u; ++c) HD[RR2_WIN * 256u + t * 4u + c] = float16_t(hv[c]);
#else
if (hl == 0u) {
f16vec4 o4 = f16vec4(0.0);
if (inside) {
precise f16vec3 gate = f16vec3(sig(float16_t(hv[0])), sig(float16_t(hv[1])), sig(float16_t(hv[2]))) * s0;
#if KIND == 3
// the guide one level down, bilinear (half-pixel centres, edge clamped), every product and sum rounded
const int bx = tx == 0 ? 0 : (tx - 1) >> 1, by = ty == 0 ? 0 : (ty - 1) >> 1;
const float16_t wx = tx == 0 ? float16_t(0.0) : ((tx & 1) != 0 ? float16_t(0.25) : float16_t(0.75));
const float16_t wy = ty == 0 ? float16_t(0.0) : ((ty & 1) != 0 ? float16_t(0.25) : float16_t(0.75));
const int x1 = min(bx + 1, hx - 1), y1 = min(by + 1, hy - 1);
precise float16_t ux = float16_t(1.0) - wx, uy = float16_t(1.0) - wy;
precise f16vec4 p00 = guideAt(gInOff, bx, by, hx) * ux, p10 = guideAt(gInOff, x1, by, hx) * wx;
precise f16vec4 p01 = guideAt(gInOff, bx, y1, hx) * ux, p11 = guideAt(gInOff, x1, y1, hx) * wx;
precise f16vec4 a00 = p00 * uy, a10 = p10 * uy, a01 = p01 * wy, a11 = p11 * wy;
precise f16vec4 u0 = a00 + a10;
precise f16vec4 u1 = u0 + a01;
precise f16vec4 up = u1 + a11;
const float16_t al = sig(float16_t(hv[3]));
precise f16vec4 gt = up * f16vec4(gate, float16_t(0.0));
precise float16_t ia = float16_t(1.0) - al;
precise f16vec4 pe = guideAt(gEncOff, tx, ty, ex) * ia;
precise f16vec4 pg = gt * al;
precise f16vec4 mx = pe + pg;
o4 = min(mx, f16vec4(65504.0));
#else
precise f16vec4 gg = guideAt(gInOff, tx, ty, ex) * f16vec4(gate, float16_t(0.0));
o4 = min(gg, f16vec4(65504.0));
#endif
const uint a = (gOutOff + uint(ty * ex + tx) * 8u) >> 2u;
arena.w[a] = bh(o4.x) | (bh(o4.y) << 16u);
arena.w[a + 1u] = bh(o4.z) | (bh(o4.w) << 16u);
}
[[unroll]] for (uint c = 0u; c < 4u; ++c) OUTG[t * 4u + c] = o4[c];
}
#endif
}
#endif
#if CM
barrier();
{ // the 2x2 merge for the next encoder: merged^T (CM x 16 rows), channel tiles m = wv, wv + 4, ...
const uvec2 mxy = mergeXY(col);
const int X = x0 / 2 + int(mxy.x), Y = y0 / 2 + int(mxy.y);
const uint o = (mrgOff + uint(Y * hx + X) * 32u) >> 2u;
for (uint m = wv; m < CM / 16u; m += 4u) {
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_BM / 2u + 16u * m + 8u * hl + uint(i)]);
for (uint ks = 0u; ks < C / 8u; ++ks) {
B8 g0, g1; A8 w0, w1;
coopMatLoad(g0, XT, 8u * ks, C, CMJ);
coopMatLoad(g1, XT, 8u * ks + 4u, C, CMJ);
LDA(w0, W_WM, 4u * C, m, 2u * ks); LDA(w1, W_WM, 4u * C, m, 2u * ks + 1u);
K32(a, w0, g0, w1, g1)
}
B8 q;
Q8(q, a)
if (X >= 0 && Y >= 0 && X < hx && Y < hy) {
const uint at = o + (m >> 1u) * (lplane >> 2u) + 4u * (m & 1u) + 2u * hl;
arena.w[at] = pack4(q[0], q[1], q[2], q[3]);
arena.w[at + 1u] = pack4(q[4], q[5], q[6], q[7]);
}
}
#if GUIDE
if (wv == 3u && lane < 16u) { // the guide for the next encoder: 2x2 mean, rightwards first, then downwards
const uint mx = lane & 3u, my = lane >> 2u;
const int x = x0 + 2 * int(mx), y = y0 + 2 * int(my);
if (x >= 0 && y >= 0 && x < ex && y < ey) {
const uint t00 = (2u * my) * 8u + 2u * mx;
f16vec4 q4[4];
[[unroll]] for (uint d = 0u; d < 4u; ++d) {
const uint tt = t00 + (d & 1u) + 8u * (d >> 1u);
precise f16vec4 qv = f16vec4(OUTG[tt * 4u], OUTG[tt * 4u + 1u], OUTG[tt * 4u + 2u], OUTG[tt * 4u + 3u]) * float16_t(0.25);
q4[d] = qv;
}
precise f16vec4 h0 = q4[0] + q4[1], h1 = q4[2] + q4[3];
precise f16vec4 hm = h0 + h1;
const uint a = (gHalfOff + uint((y / 2) * hx + x / 2) * 8u) >> 2u;
arena.w[a] = bh(hm.x) | (bh(hm.y) << 16u);
arena.w[a + 1u] = bh(hm.z) | (bh(hm.w) << 16u);
}
}
#endif
}
#endif
}
engine/rr/rr2_pack.py
+90
View File
@@ -0,0 +1,90 @@
#!/usr/bin/env python3
"""Pack a preset F stage's weights, read from the user's nvngx_dlssd.dll (engine/rr/rr2_ref.dll_weights), for
engine/rr/rr2_m.comp.
usage: rr2_pack.py STAGE NVNGX_DLSSD_DLL OUT.bin OUT.h
swin_pack.py's layout: transposed products as 16 x 16 tiles, every tensor's outputs in slot order (slot j holds the
natural index inter(j)); a merge's or the expand's inputs are sub-positions of C channels each. OUT.h also defines
the stage's shape for the shader (KIND, C, NH, NB, CM, CL, GUIDE).
"""
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from enc1_ref import inter # noqa: E402
from enc1_pack import a_tiles, b_tiles, cubic_table, act_table # noqa: E402
from rr2_ref import STAGES, EXPAND, dll_weights # noqa: E402
def main():
stage, dll, out_bin, out_h = sys.argv[1:5]
w = dll_weights(open(dll, 'rb').read(), stage)
if stage in EXPAND:
C, CL, NH, NB = EXPAND[stage]
kind, CM, guide = 2, 0, 0
else:
C, c_up, CM, NH, NB = STAGES[stage]
kind, CL = (4 if stage == 'enc0' else 5 if stage == 'dec0_1' else 3 if stage.startswith('dec') else 1 if c_up else 0), 0
guide = 0 if kind >= 3 else 1
I32, I64, IC = inter(32), inter(64), inter(C)
e = {}
if kind == 1:
e['w1'], e['b1'] = ('A', w['w1'][:, IC]), ('H', w['b1'][IC])
if kind == 4: # encoder 0's embedding branches
for b in range(2):
e[f'ewd{b}'] = ('A', w['ewd'][b][:, I32])
e[f'ebd{b}'] = ('H', w['ebd'][b][I32])
e[f'ewu{b}'] = ('A', w['ewu'][b][:, IC])
e['ebu'] = ('H', w['ebu'][IC])
if kind == 2:
I4C = [C * (j // C) + IC[j % C] for j in range(4 * C)]
e['we'], e['be'] = ('A', w['we'][:, I4C]), ('H', w['be'][I4C])
for h in range(NH):
e[f'wqk{h}'] = ('A', w['wqk'][h][:, I32])
e[f'wv{h}'] = ('B', w['wv'][h][:, I32])
e[f'wo{h}'] = ('A', w['wo'][h][:, IC])
e[f'bias{h}'] = ('H', w['bias'][h][np.ix_(I64, I64)]) # [key slot][query slot]; symmetric
e['bo'] = ('H', w['bo'][IC])
for b in range(NB):
e[f'wd{b}'] = ('A', w['wd'][b][:, I32])
e[f'bd{b}'] = ('H', w['bd'][b][I32])
e[f'wu{b}'] = ('A', w['wu'][b][:, IC])
e['bu'] = ('H', w['bu'][IC])
if CM:
e['wm'], e['bm'] = ('A', w['wm'][:, inter(CM)]), ('H', w['bm'][inter(CM)])
if 'wh' in w:
e['wh'], e['bh'] = ('A', w['wh']), ('H', w['bh'])
if kind == 5: # dec0's output tail: tanh of every binary16 value, correctly rounded, by its bits
with np.errstate(all='ignore'):
x = np.arange(65536, dtype=np.uint32).astype(np.uint16).view(np.float16).astype(np.float64)
e['tanh'] = ('U', np.tanh(x).astype(np.float16).view(np.uint16))
e['cubic'] = ('U', cubic_table())
e['act'] = ('Y', act_table())
blob, defines = bytearray(), []
for name, (kd, a) in e.items():
while len(blob) % 256:
blob.append(0)
off = len(blob)
if kd == 'A':
blob += a_tiles(a).tobytes()
elif kd == 'B':
blob += b_tiles(a).tobytes()
elif kd == 'Y':
blob += np.ascontiguousarray(a, np.uint8).tobytes()
elif kd == 'U':
blob += np.ascontiguousarray(a, np.uint16).tobytes()
else:
blob += np.ascontiguousarray(a, np.float16).tobytes()
defines.append(f'#define W_{name.upper()} {off}u // {kd} {a.shape}')
shape = [f'#define KIND {kind}', f'#define C {C}', f'#define NH {NH}', f'#define NB {NB}', f'#define CM {CM}',
f'#define CL {CL}', f'#define GUIDE {guide}']
open(out_bin, 'wb').write(bytes(blob))
open(out_h, 'w').write(f'// generated by engine/rr/rr2_pack.py for {stage}\n' + '\n'.join(shape + defines) + '\n')
print(f'{stage}: {len(blob)} bytes, {len(e)} entries')
if __name__ == '__main__':
main()
engine/rr/rr2_ref.py
+362
View File
@@ -0,0 +1,362 @@
#!/usr/bin/env python3
"""Exact numpy reference of preset F's encoder stages (rr2_enc1_kernel ... rr2_enc4_kernel of nvngx_dlssd.dll
310.9), one 8x8 window at a time, from weights and real inputs only; checked against a capture (main).
A stage has C channels and works on windows of 64 tokens (origin 8 * block - shift). Arithmetic is binary16 with FP8
(e4m3) operands in the products, exactly as in presets D and E (enc1_ref.py: mm, the cubic softmax, the activation):
merge x = W1 . [the four tokens of the 2x2 group one level up, (0,0) (1,0) (0,1) (1,1)] + b1; window tokens
outside the level are the mirrored tokens
heads per head, all from q8(x): QK = Wqk . x and V = Wv . x (32 channels each, no bias). The scores are
QK . QK^T + bias: symmetric, so NVIDIA computes the blocks (queries 0-31, keys 0-31), (0-31, 32-63) and
(32-63, 32-63) and transposes the second for (32-63, 0-31). Cubic softmax over the 64 keys, P . V, and the
out-projections of all heads accumulate into one sum a (the first one's C is the bias bo)
residual r = dq(q8(x)) + a * s1 (dq: the FP8 code back as binary16; s1: binary16 at argument 94)
MLP NB branches from q8(r): C -> 32 (+ bias) -> act -> C, accumulated into m (first C: the bias bu)
residual y = r + m * s2 (s2: argument 96); the stage's tokens are q8(y)
outputs merged = Wm . [y of the 2x2 group] + bm for the next encoder (q8), and a 16-channel head h = Wh . y + bh
whose first three channels gate the stage's guide image (argument 48, three binary16 channels per token):
out = min(guide * (sigmoid(h) * s0), 65504) (s0: argument 92) at argument 64, and its 2x2 mean
((a + b) + (c + d) of the quarters) at argument 72 for the next encoder
Arena layouts: tokens are FP8 codes in planes of 32 channels ([plane][y][x][32], NVIDIA's k order within a plane);
the guide images are four binary16 per token.
usage: rr2_ref.py CAPTURE STAGE FRAME WEIGHTS.npz [--blocks X,Y ...] (verifies against the capture)
rr2_ref.py CAPTURE STAGE FRAME OUT.npz --extract DUMP.npz (weights from a phase dump, phase_dump.py)
"""
import argparse
import os
import struct
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rrmath import e4m3, q8 # noqa: E402
from enc1_ref import inter, mm, act, add, pairs_tree, mirror, H, F # noqa: E402
from rr_capture import Capture # noqa: E402
# stage: C, channels one level up (merge input), merged channels for the next encoder, heads, MLP branches
STAGES = {'enc0': (32, 0, 0, 2, 4), 'enc1': (64, 32, 96, 2, 8), 'enc2': (96, 0, 128, 4, 12), 'enc3': (128, 0, 160, 4, 16),
'enc4': (160, 0, 160, 8, 20), 'enc5': (160, 0, 0, 8, 20),
'dec0_1': (32, 0, 0, 2, 4), 'dec4_1': (160, 0, 0, 8, 20), 'dec3_1': (128, 0, 0, 4, 16), 'dec2_1': (96, 0, 0, 4, 12), 'dec1_1': (64, 0, 0, 2, 8)}
# the decoders' first kernels (decN_0): C, channels one level down, heads, MLP branches. They expand the 4 x 4 tokens
# one level down to the window, e = We . low + be (4C outputs: sub-position (y & 1) * 2 + (x & 1), then channel), and
# start from x = dq(q8(e)) + dq(skip) * s3 (s3: binary16 at argument 98; skip: the encoder's tokens of this level,
# argument 16); then the heads and the MLP as above, the first residual on x itself (not dq(q8(x))). No merge, no guide.
EXPAND = {'dec4_0': (160, 160, 8, 20), 'dec3_0': (128, 160, 4, 16), 'dec2_0': (96, 128, 4, 12), 'dec1_0': (64, 96, 2, 8),
'dec0_0': (32, 64, 2, 4)}
# the network's sixteen kernels in the DLL's order, with their head counts
ORDER = ['enc0', 'enc1', 'enc2', 'enc3', 'enc4', 'enc5', 'dec4_0', 'dec4_1', 'dec3_0', 'dec3_1', 'dec2_0', 'dec2_1',
'dec1_0', 'dec1_1', 'dec0_0', 'dec0_1']
HEADS = [2, 2, 4, 4, 8, 8, 8, 8, 4, 4, 4, 4, 2, 2, 2, 2]
def find_tables(dll):
"""byte offsets of the sixteen kernels' relative position tables in nvngx_dlssd.dll: runs of 64 x 64 binary16
tables (one per head) with a constant diagonal and at most 225 distinct values, in ORDER"""
a = np.frombuffer(dll[:len(dll) & ~1], np.uint16)
n = len(a) - 4096
m = (a[:n] & 0x7fff) != 0
for k in range(1, 8):
m &= a[:n] == a[65 * k:65 * k + n]
m &= a[1:n + 1] == a[66:66 + n]
groups = []
for i in np.nonzero(m)[0]:
t = a[i:i + 4096].reshape(64, 64)
if not (np.all(t.diagonal() == t[0, 0]) and 50 < len(np.unique(t)) <= 225):
continue
if groups and 2 * int(i) - groups[-1][-1] == 8192:
groups[-1].append(2 * int(i))
else:
groups.append([2 * int(i)])
for start in range(len(groups) - 15):
if [len(g) for g in groups[start:start + 16]] == HEADS:
return {name: g[0] for name, g in zip(ORDER, groups[start:start + 16])}
raise SystemExit('no preset F model found in the DLL (its sixteen relative position table groups)')
def dll_weights(dll, stage):
"""an encoder stage's tensors from the DLL: plain binary16 arrays in natural [out][in] order, one after another.
Before the tables: {s0, 0, 0, 0} and (enc1) the merge of encoder 0's tokens, W1 then b1. After them, per tensor
kind all heads / branches together: Wqk, Wv (heads x 32 x C), Wo (C x heads x 32), bo, Wd (branches x 32 x C), bd,
Wu (C x branches x 32), bu, the merge
for the next encoder Wm (C_M x 4C) and bm, the guide head Wh (3 x C) and bh. A merge's input index is
channel * 4 + sub-position; the kernels' bias is (T + T^T) * sqrt(2)."""
if stage in EXPAND:
c, c_lo, nh, nb = EXPAND[stage]
c_up = c_m = 0
else:
c, c_up, c_m, nh, nb = STAGES[stage]
g = find_tables(dll)[stage]
f = lambda off, n: np.frombuffer(dll[off:off + 2 * n], H)
def B(mat, merge=0):
"""[out][in] binary16 -> the product's B operand (in x out FP8 codes, the inputs in NVIDIA's k order)"""
k = mat.shape[1]
if merge: # channel * 4 + q -> q * channels + channel
mat = mat.reshape(mat.shape[0], merge, 4).transpose(0, 2, 1).reshape(mat.shape[0], k)
return q8(mat).T[inter(k)]
w = {}
if c_up:
o = g - 8 - 2 * (3 * c_up + 4) - 2 * c - 2 * c * 4 * c_up # before encoder 0's guide head (3 x C_UP, bias, 0)
w['w1'] = B(f(o, c * 4 * c_up).reshape(c, 4 * c_up), c_up)
w['b1'] = f(o + 2 * c * 4 * c_up, c)
if stage in EXPAND: # We (C x 4 x C_LO) and be (C x 4) end where the tables start
o = g - 8 * c - 8 * c * c_lo
we = f(o, 4 * c * c_lo).reshape(c, 4, c_lo).transpose(1, 0, 2).reshape(4 * c, c_lo)
w['we'] = q8(we).T[inter(c_lo)]
w['be'] = f(g - 8 * c, 4 * c).reshape(c, 4).T.reshape(-1)
if stage == 'enc0':
# encoder 0's embedding of the 43 per-pixel features (the kernel's 64 feature slots, the rest unused): two MLP
# branches 43 -> 32 -> act -> 32, x = bu + sum: Wd (2 x 32 x 43), bd, Wu (32 x 2 x 32), bu, then 16 bytes
o = g - 16 - 64 - 4096 - 128 - 2 * 2 * 32 * 43
wd = np.zeros((2, 32, 64), H); wd[:, :, :43] = f(o, 2 * 32 * 43).reshape(2, 32, 43); o += 2 * 2 * 32 * 43
w['ewd'] = np.array([B(x) for x in wd])
w['ebd'] = f(o, 64).reshape(2, 32); o += 128
w['ewu'] = np.array([B(x) for x in f(o, 2048).reshape(32, 2, 32).transpose(1, 0, 2)]); o += 4096
w['ebu'] = f(o, 32)
o = g
t = f(o, nh * 4096).reshape(nh, 64, 64).astype(np.float64); o += nh * 8192
w['bias'] = ((t + t.transpose(0, 2, 1)) * np.sqrt(2)).astype(H)
w['wqk'] = np.array([B(x) for x in f(o, nh * 32 * c).reshape(nh, 32, c)]); o += nh * 64 * c
w['wv'] = np.array([B(x) for x in f(o, nh * 32 * c).reshape(nh, 32, c)]); o += nh * 64 * c
w['wo'] = np.array([B(x) for x in f(o, nh * 32 * c).reshape(c, nh, 32).transpose(1, 0, 2)]); o += nh * 64 * c
w['bo'] = f(o, c); o += 2 * c
w['wd'] = np.array([B(x) for x in f(o, nb * 32 * c).reshape(nb, 32, c)]); o += nb * 64 * c
w['bd'] = f(o, nb * 32).reshape(nb, 32); o += nb * 64
w['wu'] = np.array([B(x) for x in f(o, nb * 32 * c).reshape(c, nb, 32).transpose(1, 0, 2)]); o += nb * 64 * c
w['bu'] = f(o, c); o += 2 * c
if c_m:
w['wm'] = B(f(o, c_m * 4 * c).reshape(c_m, 4 * c), c); o += 2 * c_m * 4 * c
w['bm'] = f(o, c_m); o += 2 * c_m
if stage in EXPAND:
return w
if stage == 'enc0': # the merge of its tokens (enc1's W1, b1) comes before its guide head
o += 2 * 64 * 128 + 2 * 64
nhd = 4 if stage.startswith('dec') else 3 # the guide head: three gates; decoders also a blend factor
wh = np.zeros((16, c), H); wh[:nhd] = f(o, nhd * c).reshape(nhd, c); o += 2 * nhd * c
bh = np.zeros(16, H); bh[:nhd] = f(o, nhd)
w['wh'], w['bh'] = B(wh), bh
return w
def dq(x):
"""binary16 values through FP8 and back"""
return e4m3(q8(x)).astype(H)
def A(x):
"""binary16 rows -> the product's A operand: FP8 codes in NVIDIA's k order"""
return q8(x)[:, inter(x.shape[1])]
def softmax(s):
"""the cubic-exponent softmax of presets D and E over the 64 keys of a row"""
def fma16(a, b, c):
return (a.astype(np.float64) * b.astype(np.float64) + c.astype(np.float64)).astype(H)
m = np.clip((s * H(0.017181396484375)).astype(H), H(-0.55615234375), H(0.55615234375))
t = fma16(m, -m, np.full_like(m, H(0.92724609375)))
u = fma16(m, t, np.full_like(m, H(1.375)))
e = ((((u.view(np.uint16).astype(np.uint32)) << 5) & 0x7FE0).astype(np.uint16)).view(H)
nt = e.shape[1] // 8
part = {}
for tt in range(4):
for half in range(2):
ev = pairs_tree([e[:, 8 * j + 2 * tt + half] for j in range(0, nt, 2)])
od = pairs_tree([e[:, 8 * j + 2 * tt + half] for j in range(1, nt, 2)])
part[tt, half] = add(ev, od)
out = [add(add(part[0, h], part[1, h]), add(part[2, h], part[3, h])) for h in range(2)]
r = (F(1) / add(out[0], out[1]).astype(F)).astype(H)
return (e * r[:, None]).astype(H)
def sigmoid(d):
"""1 / (1 + 2^(-d log2 e)) in f32, rounded to binary16 (the GPU's exp2 and division; see rr_exact.glsl sig)"""
f = F(1) / (np.exp2(d.astype(F) * F(-1.4426950216293335)).astype(F) + F(1))
return f.astype(H)
def upsample(low, x, y):
"""the decoders' 2x bilinear fetch of the guide one level down (half-pixel centres, edge clamped), with NVIDIA's
binary16 rounding: every weight product and the three sums in tap order"""
def axis(v, n):
if v == 0:
return 0, H(0)
b = (v - 1) // 2
return b, H(0.25) if v & 1 else H(0.75)
hh, hw = low.shape[:2]
bx, wx = axis(x, hw)
by, wy = axis(y, hh)
x1, y1 = min(bx + 1, hw - 1), min(by + 1, hh - 1)
ux, uy = H(1) - wx, H(1) - wy
m = lambda v, a, b: ((v * a).astype(H) * b).astype(H)
return add(add(add(m(low[by, bx], ux, uy), m(low[by, x1], wx, uy)), m(low[y1, bx], ux, wy)), m(low[y1, x1], wx, wy))
def window(w, scal, xin, x=None):
"""xin: 64 x 4*C_UP FP8 codes (the merge input, already in k order). Returns y (64 x C binary16), the merged
tokens (16 x C_M binary16) and the head (64 x 16 binary16)."""
s0, s1, s2 = scal
base = None
if x is not None: # a decoder's first kernel: the expanded tokens plus the skip; the residual keeps it unquantised
xa, base = A(x), x
elif 'w1' in w:
x = mm(xin, w['w1'], np.tile(w['b1'], (64, 1)))
xa = A(x)
else: # the tokens themselves (already FP8)
x, xa = e4m3(xin)[:, np.argsort(inter(xin.shape[1]))].astype(H), xin
acc = np.tile(w['bo'], (64, 1))
for h in range(len(w['wqk'])):
qk = mm(xa, w['wqk'][h], np.zeros((64, 32), H))
v = mm(xa, w['wv'][h], np.zeros((64, 32), H))
qa = A(qk)
s = mm(qa, qa.T, w['bias'][h])
# NVIDIA computes three 32 x 32 blocks and transposes one; the full product is the same values
p = softmax(s)
o = mm(A(p), A(v.T).T if False else q8(v)[inter(64), :], np.zeros((64, 32), H))
acc = mm(A(o), w['wo'][h], acc)
r = add(dq(x) if base is None else base, (acc * s1).astype(H))
ra = A(r)
m = np.tile(w['bu'], (64, 1))
for j in range(len(w['wd'])):
d = mm(ra, w['wd'][j], np.tile(w['bd'][j], (64, 1)))
m = mm(A(act(d)), w['wu'][j], m)
y = add(r, (m * s2).astype(H))
ya = A(y)
# the 2x2 groups of the window: token order (0,0) (1,0) (0,1) (1,1), 16 groups row-major
g = np.zeros((16, 4 * y.shape[1]), np.uint8)
for n in range(16):
gx, gy = n % 4, n // 4
for q, (dx, dy) in enumerate(((0, 0), (1, 0), (0, 1), (1, 1))):
g[n, q * y.shape[1]:(q + 1) * y.shape[1]] = ya[(2 * gy + dy) * 8 + 2 * gx + dx]
merged = mm(g, w['wm'], np.tile(w['bm'], (16, 1))) if 'wm' in w else None
head = mm(ya, w['wh'], np.tile(w['bh'], (64, 1))) if 'wh' in w else None
return y, merged, head
def extract(z, out):
"""weights of enc1 from a phase dump (phases as phase_dump.py numbers them)"""
P = lambda p, n: z[f'p{p}_w0_{n}']
Ch = lambda p: P(p, 'C').view(H)
w = {'w1': P(1, 'B'), 'b1': Ch(1)[0], 'wqk': [], 'wv': [], 'bias': [], 'wo': [], 'bo': Ch(11)[0]}
for h in range(2):
b = 2 + 10 * h
w['wqk'].append(np.vstack([P(b, 'B'), P(b + 2, 'B')]))
w['wv'].append(np.vstack([P(b + 1, 'B'), P(b + 3, 'B')]))
tl, tr, br = Ch(b + 4), Ch(b + 5), Ch(b + 7)
w['bias'].append(np.block([[tl, tr], [tr.T, br]]))
w['wo'].append(P(b + 9, 'B'))
w['wd'] = [P(22 + 2 * j, 'B') for j in range(8)]
w['bd'] = [Ch(22 + 2 * j)[0] for j in range(8)]
w['wu'] = [P(23 + 2 * j, 'B') for j in range(8)]
w['bu'] = Ch(23)[0]
w['wm'], w['bm'], w['wh'], w['bh'] = P(38, 'B'), Ch(38)[0], P(39, 'B'), Ch(39)[0]
np.savez_compressed(out, **{k: np.array(v) for k, v in w.items()})
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('stage')
ap.add_argument('frame', type=int)
ap.add_argument('weights')
ap.add_argument('--extract')
ap.add_argument('--blocks', nargs='*', default=['5,5'])
o = ap.parse_args()
if o.extract:
return extract(np.load(o.extract), o.weights)
expand = o.stage in EXPAND
if expand:
c, c_lo, heads, nb = EXPAND[o.stage]
c_up = c_m = 0
else:
c, c_up, c_m, heads, nb = STAGES[o.stage]
w = dll_weights(open(o.weights, 'rb').read(), o.stage) if o.weights.lower().endswith('.dll') else dict(np.load(o.weights))
cap = Capture(os.path.expanduser(o.capture))
kernel = f'rr2_{o.stage}_kernel' if o.stage.startswith('enc') else f'rr2_{o.stage[:4]}_kernel_{o.stage[5]}'
l = [l for l in cap.launches if l.kernel == kernel][o.frame - 1]
pre = np.frombuffer(l.pre(l.arena_alloc), np.uint8)
post = np.frombuffer(l.post_buffer(cap, l.arena_alloc), np.uint8)
ptr = {p.offset: p.delta for p in l.pointers if p.alloc == l.arena_alloc}
wd, ht = struct.unpack_from('<2I', l.args, 0)
sx, sy = struct.unpack_from('<2H', l.args, 80)
scal = np.frombuffer(l.args[92:98], H)
s3 = np.frombuffer(l.args[98:100], H)[0]
planes = lambda buf, off, w_, h_, ch: buf[off:off + w_ * h_ * ch].reshape(ch // 32, h_, w_, 32)
up = (planes(pre, ptr[8], wd // 2, ht // 2, c_lo) if expand else planes(pre, ptr[8], 2 * wd, 2 * ht, c_up) if c_up else
planes(pre, ptr[8], wd, ht, c))
skip = planes(pre, ptr[16], wd, ht, c) if expand else None
guided = 48 in ptr and o.stage.startswith('enc')
dec_guide = o.stage.startswith('dec') and 56 in ptr
tok_ref = planes(post, ptr[32], wd, ht, c)
mer_ref = planes(post, ptr[40], wd // 2, ht // 2, c_m) if c_m else None
img = lambda buf, off, w_, h_: np.frombuffer(bytes(buf[off:off + w_ * h_ * 8]) + bytes(8), H)[:w_ * h_ * 4].reshape(h_, w_, 4)
guide, out_ref = (img(pre, ptr[48], wd, ht), img(post, ptr[64], wd, ht)) if guided else (None, None)
half_ref = img(post, ptr[72], wd // 2, ht // 2) if c_m else None
if dec_guide: # the guide one level down (argument 48), the encoder's of this level (56), the output (64)
g_low, g_enc, out_ref = img(pre, ptr[48], wd // 2, ht // 2), img(pre, ptr[56], wd, ht), img(post, ptr[64], wd, ht)
ok = True
for blk in o.blocks:
bx, by = (int(v) for v in blk.split(','))
x0, y0 = 8 * bx - sx, 8 * by - sy
inside = np.array([0 <= x0 + t % 8 < wd and 0 <= y0 + t // 8 < ht for t in range(64)])
xin = np.zeros((64, 4 * c_up if c_up else c), np.uint8)
xe = None
if expand:
lows = np.array([np.concatenate([up[p, mirror(y0 // 2 + n // 4, ht // 2), mirror(x0 // 2 + n % 4, wd // 2)]
for p in range(c_lo // 32)]) for n in range(16)])
e = mm(lows, w['we'], np.tile(w['be'], (16, 1)))
un = np.argsort(inter(c))
xe = np.zeros((64, c), H)
for t in range(64):
tx, ty = t % 8, t // 8
sub = (ty % 2) * 2 + tx % 2
sk = np.concatenate([skip[p, mirror(y0 + ty, ht), mirror(x0 + tx, wd)] for p in range(c // 32)])
xe[t] = add(dq(e[(ty // 2) * 4 + tx // 2, sub * c:(sub + 1) * c]), (e4m3(sk)[un].astype(H) * s3).astype(H))
for t in range(64 if not expand else 0): # tokens outside the level are the mirrored ones (no edge repeat), as in presets D and E
X, Y = mirror(x0 + t % 8, wd), mirror(y0 + t // 8, ht)
if not c_up:
xin[t] = np.concatenate([up[p, Y, X] for p in range(c // 32)])
continue
for q, (dx, dy) in enumerate(((0, 0), (1, 0), (0, 1), (1, 1))):
xin[t, q * c_up:(q + 1) * c_up] = np.concatenate([up[p, 2 * Y + dy, 2 * X + dx] for p in range(c_up // 32)])
y, merged, head = window(w, scal, xin, xe)
ya, ma = A(y), (A(merged) if c_m else None)
bad_t = bad_m = bad_o = bad_h = 0
out = np.zeros((64, 4), H)
for t in np.nonzero(inside)[0]:
X, Y = x0 + t % 8, y0 + t // 8
bad_t += int(np.sum(np.concatenate([tok_ref[p, Y, X] for p in range(c // 32)]) != ya[t]))
if dec_guide:
# out = min(enc * (1 - a) + (upsampled * sigmoid(h[0:3]) * s0) * a, 65504), a = sigmoid(h[3])
sg = sigmoid(head[t, :4])
gate = np.array([(sg[0] * scal[0]).astype(H), (sg[1] * scal[0]).astype(H), (sg[2] * scal[0]).astype(H), H(0)])
gt = (upsample(g_low, X, Y) * gate).astype(H)
gt[3] = H(0)
mix = add((g_enc[Y, X] * (H(1) - sg[3])).astype(H), (gt * sg[3]).astype(H))
out[t] = np.minimum(mix, H(65504))
bad_o += int(np.sum(out[t].view(np.uint16) != out_ref[Y, X].view(np.uint16)))
if not guided:
continue
gate = (sigmoid(head[t, :3]) * scal[0]).astype(H)
out[t, :3] = np.minimum((guide[Y, X, :3] * gate).astype(H), H(65504))
bad_o += int(np.sum(out[t].view(np.uint16) != out_ref[Y, X].view(np.uint16)))
for n in range(16):
X, Y = x0 // 2 + n % 4, y0 // 2 + n // 4
t0 = (2 * (n // 4)) * 8 + 2 * (n % 4)
if not c_m or not (inside[t0] and x0 % 2 == 0 and y0 % 2 == 0):
continue
bad_m += int(np.sum(np.concatenate([mer_ref[p, Y, X] for p in range(c_m // 32)]) != ma[n]))
qv = [(out[t0 + d] * H(0.25)).astype(H) for d in (0, 1, 8, 9)]
bad_h += int(np.sum(add(add(qv[0], qv[1]), add(qv[2], qv[3])).view(np.uint16) != half_ref[Y, X].view(np.uint16)))
print(f'block {blk}: {int(inside.sum())} tokens inside; differing tokens bytes {bad_t}, merged bytes {bad_m}, '
f'guide values {bad_o}, half guide values {bad_h}')
ok &= bad_t + bad_m + bad_o + bad_h == 0
print('IDENTICAL' if ok else 'DIFFERENT')
sys.exit(0 if ok else 1)
if __name__ == '__main__':
main()
engine/rr/rr_attn.glsl
+3 -4
View File
@@ -21,13 +21,12 @@
[[unroll]] for (uint kt = 0u; kt < 6u; ++kt) {
HAC c;
coopMatLoad(c, w16.w, (W_BIAS0 + HEAD * hd) / 2u + kt * 1024u + 16u * wv, 64u, RM);
ACC s = ACC(c);
A8 k0, k1;
coopMatLoad(k0, RR_KS, RR_KS0 + kt * 128u, 8u, RM);
coopMatLoad(k1, RR_KS, RR_KS0 + kt * 128u + 4u, 8u, RM);
s = coopMatMulAdd(k1, Qb[1], coopMatMulAdd(k0, Qb[0], s)); // rounded once, by the conversion below
K32H(c, k0, Qb[0], k1, Qb[1])
// the cubic numerator from enc1_pack.cubic_table, indexed by the score's binary16 bits
[[unroll]] for (int i = 0; i < 8; ++i) e[kt][i] = hb(uint(w16u.w[W_CUBIC / 2u + bh(float16_t(s[i]))]));
[[unroll]] for (int i = 0; i < 8; ++i) e[kt][i] = hb(uint(w16u.w[W_CUBIC / 2u + bh(c[i])]));
}
// key slot 16kt + 8hl + i is NVIDIA's column 8j + 2tt + half with j = 4(kt >> 1) + (kt & 1) + 2((i >> 1) & 1),
// tt = 2hl + (i >> 2), half = i & 1 (enc1_ref.softmax_cubic's tree)
@@ -60,7 +59,7 @@
// O^T = V^T P^T, then the out-projection onto the residual
B8 Ob[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
ACC o = ACC(0.0);
RAC o = RAC(RV(float16_t(0.0)));
[[unroll]] for (uint ks = 0u; ks < 3u; ++ks) {
A8 v0, v1;
coopMatLoad(v0, RR_VS, RR_VS0 + c * 384u + 8u * ks, 24u, RM);
engine/rr/rr_capture.py
+8 -4
View File
@@ -70,16 +70,20 @@ class Launch:
def image(self, offset):
return next(i for i in self.images if i.offset == offset)
def weights_at(self, offset):
"""the weight image (the pointer at argument OFFSET) as NGX uploaded it; 16 zero bytes for a stage without one"""
pointer = next((p for p in self.pointers if p.offset == offset), None)
return self.pre(pointer.alloc) if pointer is not None else bytes(16)
@property
def weights(self):
"""the weight image (argument 40) as NGX uploaded it; 16 zero bytes for a stage without one (the output stage)"""
pointer = next((p for p in self.pointers if p.offset == 40), None)
return self.pre(pointer.alloc) if pointer is not None else bytes(16)
"""the DLSS-D kernels' weight image (argument 40)"""
return self.weights_at(40)
@property
def arena_alloc(self):
"""the plane arena: the allocation every arena pointer of the layer points into"""
candidates = [p.alloc for p in self.pointers if p.offset != 40]
candidates = [p.alloc for p in self.pointers]
return max(set(candidates), key=lambda a: self.allocs[a][1]) if candidates else None
def post_buffer(self, capture, alloc):
engine/rr/rr_inputs.py
+41
View File
@@ -25,6 +25,45 @@ OPTIONAL_OUTPUT = {'particles': 416} # output stage
INPUTS = {'color': 160, 'motion': 168, 'normals': 184, 'roughness': 192, 'diffuse': 232, 'specular': 240, 'depth': 512}
# preset F (rr2_enc0_kernel): its textures and argument offsets
INPUTS_F = {'color': 144, 'motion': 192, 'normals': 224, 'roughness': 240, 'diffuse': 160, 'specular': 176, 'depth': 552}
def preset_f(cap, o):
lines = []
for f in range(1, o.frames + 1):
enc0 = cap.find('rr2_enc0', f)
a = enc0.args
u32 = lambda off: struct.unpack_from('<I', a, off)[0]
f32 = lambda off: struct.unpack_from('<f', a, off)[0]
mat = lambda off: np.array(struct.unpack_from('<16f', a, off), np.float64).reshape(4, 4)
down = cap.find('rr2_downsample', 1).args # either variant: the sizes are at the same offsets
rw, rh, iw, ih = u32(128), u32(132), u32(488), u32(492)
ow, oh = struct.unpack_from('<2I', down, 8)
if f == 1:
lines.append(f'size {rw} {rh} {ow} {oh}')
images = {i.offset: i for i in enc0.images}
for name, off in INPUTS_F.items():
(o.out / f'{f}-{name}.bin').write_bytes(images[off].path.read_bytes())
# the optional inputs: encoder 0's textures 272 / 336 / 288, the particle downsample kernel's texture 208
for name, off in (('hitdistance', 272), ('disocclusion', 336), ('sss', 288)):
if off in images:
(o.out / f'{f}-{name}.bin').write_bytes(images[off].path.read_bytes())
part = [l for l in cap.launches if l.kernel == 'rr2_downsample_kernel_static_hdr']
if part:
n = len(part) // len(cap.find('rr2_enc0'))
(o.out / f'{f}-particles.bin').write_bytes(part[(f - 1) * n].image(208).path.read_bytes())
clip = mat(708)
view = mat(644) @ np.linalg.inv(clip)
view[3, :3] = 0.0
reset = 1 if f == 1 or not o.no_reset else 0
values = ' '.join(f'{v:.9g}' for v in list(view.flatten()) + list(clip.flatten()))
flags = a[568] | (a[1028] << 1) | (a[466] << 2) # inverted depth, hardware depth, packed roughness
lines.append(f'frame {f} {f32(400):.9g} {f32(404):.9g} {f32(448) * rw / iw:.9g} {f32(452) * rh / ih:.9g} 1 {reset} {flags} {values}')
(o.out / 'frames.txt').write_text('\n'.join(lines) + '\n')
print(f'{o.frames} preset F frames of {o.capture} in {o.out}')
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
@@ -35,6 +74,8 @@ def main():
cap = Capture(o.capture)
o.out.mkdir(parents=True, exist_ok=True)
lines = []
if any(l.kernel == 'rr2_enc0_kernel' for l in cap.launches):
return preset_f(cap, o)
for f in range(1, o.frames + 1):
enc0 = cap.find('swin_enc0', f)
a = enc0.args
engine/rr/swin_m.comp
+28 -12
View File
@@ -54,7 +54,23 @@ layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S128;
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
#define Q8(dst, src) Q8PAIRS(dst, src)
// One 32-deep step onto an accumulator that holds binary16 values. The reference (and the default here) adds the
// accumulator inside the f32 products and rounds the result to binary16. RR_MIX (build_native_model.sh with
// D4R_RR_MIX=1) keeps the accumulator as binary16, sums the products from zero and takes f16(P + acc), which the
// driver compiles to one v_fma_mix per element in place of two conversions: about 5% faster, but a single rounding
// where the reference has two, so a value can come out one binary16 step away (engine/rr/README.md).
// RAC is the accumulator type, RV(x) a binary16 value as its element; K32H is the step on a binary16 matrix.
#ifdef RR_MIX
#define RAC HAC
#define RV(x) (x)
#define K32(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(0.0))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_] + float((acc)[i_])); }
#define K32H(acc, a0, b0, a1, b1) K32(acc, a0, b0, a1, b1)
#else
#define RAC ACC
#define RV(x) float(x)
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) }
#define K32H(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(acc))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_]); }
#endif
const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0;
// NVIDIA's reduction over a short list (swin_ref.pairs_tree)
@@ -143,8 +159,8 @@ void main()
[[unroll]] for (uint kt = 0u; kt < CX / 16; ++kt) coopMatLoad(lb[kt], POOL, LOW_AT + n * 4u * CX + 4u * kt, CX / 4u, CM);
for (uint mm = 0u; mm < MT; ++mm) {
const uint m = MT * wv + mm;
ACC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_BE / 2u + 16u * m + 8u * hl + uint(i)]);
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_BE / 2u + 16u * m + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < CX / 16; kt += 2u) {
A8 w0, w1;
LDA(w0, W_WE, CX, m, kt); LDA(w1, W_WE, CX, m, kt + 1u);
@@ -213,7 +229,7 @@ void main()
barrier();
// ---------------------------------------------------------------- the residual: the window token + table 7
ACC R[MT];
RAC R[MT];
{
const uint tq = inter(16u * wv + col);
const int qx = x0 + int(tq & 7u), qy = y0 + int(tq >> 3u);
@@ -235,7 +251,7 @@ void main()
const float16_t x = e4(t[i >> 2], uint(i) & 3u);
#endif
precise float16_t c = x + w16.w[W_TABLE7 / 2u + 16u * m + 8u * hl + uint(i)];
R[m][i] = float(c);
R[m][i] = RV(c);
}
}
}
@@ -245,7 +261,7 @@ void main()
barrier(); // every wave is done with the previous head's K / V (and, for the first, the expand's LOW)
for (uint j = 0u; j < 3u; ++j) { // 12 jobs (token tile n, channel tile c)
const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u;
ACC a = ACC(0.0), v = ACC(0.0);
RAC a = RAC(RV(float16_t(0.0))), v = RAC(RV(float16_t(0.0)));
[[unroll]] for (uint kt = 0u; kt < 2u * KC; kt += 2u) {
B8 xb0, xb1, wv0, wv1;
A8 xa0, xa1, wk0, wk1;
@@ -312,15 +328,15 @@ void main()
precise float16_t n = rv[m][i] * g;
Nb[m][i] = q8h(n);
precise float16_t a = rv[m][i] + w16.w[W_TABLE16 / 2u + k];
R[m][i] = float(a);
R[m][i] = RV(a);
}
}
[[dont_unroll]] for (uint br = 0u; br < NB; ++br) {
const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br;
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
ACC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = float(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]);
RAC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = RV(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) {
A8 a0, a1;
LDA(a0, w1, C, c, kt); LDA(a1, w1, C, c, kt + 1u);
@@ -360,8 +376,8 @@ void main()
}
}
// the C -> 16 head
ACC hv;
[[unroll]] for (int i = 0; i < 8; ++i) hv[i] = float(w16.w[W_BH / 2u + 8u * hl + uint(i)]);
RAC hv;
[[unroll]] for (int i = 0; i < 8; ++i) hv[i] = RV(w16.w[W_BH / 2u + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) {
A8 a0, a1;
LDA(a0, W_WH, C, 0u, kt); LDA(a1, W_WH, C, 0u, kt + 1u);
@@ -427,8 +443,8 @@ void main()
const int X = x0 / 2 + int(mxy.x), Y = y0 / 2 + int(mxy.y);
const uint o = (mrgOff + uint(Y * hx + X) * 32u) >> 2u;
for (uint m = wv; m < CX / 16u; m += 4u) {
ACC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_BM / 2u + 16u * m + 8u * hl + uint(i)]);
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_BM / 2u + 16u * m + 8u * hl + uint(i)]);
for (uint ks = 0u; ks < C / 8u; ++ks) {
B8 g0, g1; A8 w0, w1;
coopMatLoad(g0, MG, 8u * ks, C, CM);
engine/rr/test_layer.py
+23 -7
View File
@@ -34,11 +34,11 @@ from rr_capture import Capture # noqa: E402
SURFACE_FORMATS = {'rgba16f': (4, 16), 'rg16f': (2, 16), 'r16f': (1, 16), 'rgba8_snorm': (4, 9), 'rgba8': (4, 8)}
def stage_params(launch):
def stage_params(launch, weights_arg=40):
args = bytearray(launch.args)
arena = launch.arena_alloc
for p in launch.pointers:
value = p.delta if p.alloc == arena and p.offset != 40 else 0
value = p.delta if p.alloc == arena and p.offset != weights_arg else 0
struct.pack_into('<Q', args, p.offset, value)
# texture and surface handles: 1 where NGX bound one (the kernels test optional inputs' handles for null, e.g.
# enc0's specular hit distance at argument 248), 0 elsewhere
@@ -117,6 +117,7 @@ def main():
ap.add_argument('--time', type=int, default=1)
ap.add_argument('--subgroup', type=int, default=32)
ap.add_argument('--weights', help='bind this file at binding 1 instead of NGX\'s weight image (native stages)')
ap.add_argument('--weights-arg', type=int, default=40, help='argument offset of the weight pointer (rr2_*: 24)')
ap.add_argument('--grid', help='X,Y workgroups (default: the launch grid)')
ap.add_argument('--ignore', action='append', default=[], metavar='SURFACE:CHANNEL',
help='leave a channel out of the comparison, e.g. 112:3 for a value NVIDIA stores from an '
@@ -130,10 +131,10 @@ def main():
work.mkdir(parents=True, exist_ok=True)
print(f'{launch.kernel} (launch {launch.sequence}, grid {launch.grid}, block {launch.block})')
(work / 'params.bin').write_bytes(stage_params(launch))
(work / 'weights.bin').write_bytes(Path(o.weights).read_bytes() if o.weights else launch.weights)
(work / 'params.bin').write_bytes(stage_params(launch, o.weights_arg))
(work / 'weights.bin').write_bytes(Path(o.weights).read_bytes() if o.weights else launch.weights_at(o.weights_arg))
arena = launch.arena_alloc
arena_before = launch.pre(arena)
arena_before = launch.pre(arena) if arena is not None else bytes(16)
(work / 'arena.bin').write_bytes(arena_before)
gx, gy = (int(v) for v in o.grid.split(',')) if o.grid else launch.grid[:2]
lines = [f'param {work}/params.bin', f'grid {gx} {gy}']
@@ -162,16 +163,24 @@ def main():
n_images = len(order)
lines.append(f'lut {work}/weights.bin')
lines.append(f'inout {work}/arena.bin {work}/arena.out.bin')
# further buffers of the stage (ptx2rr.py --buffer), bound after the surfaces
extra = info.get('buffers', []) if sidecar.exists() else []
for off in extra:
present = any(p.offset == off for p in launch.pointers) # NGX may leave the pointer null
(work / f'buf{off}.bin').write_bytes(launch.pre(launch.pointer(off).alloc) if present else bytes(16))
lines.append(f'inout {work}/buf{off}.bin {work}/buf{off}.out.bin')
probes = info.get('probes', []) if sidecar.exists() else []
oob = info.get('oob', False) if sidecar.exists() else False
if probes: # ptx2rr.py --probe: the debug buffer after the surfaces
threads = launch.block[0] * launch.block[1] * launch.block[2]
# D4R_PROBE_THREADS: the shader's workgroup size when it is not the launch's block (the hybrid stages)
threads = int(os.environ.get('D4R_PROBE_THREADS', launch.block[0] * launch.block[1] * launch.block[2]))
lines.append(f'outbuf {work}/debug.bin {4 * threads * len(probes)}')
elif oob: # ptx2rr.py --oob-report: count, lowest and highest address per buffer
(work / 'oob.init').write_bytes(struct.pack('<9I', *([0, 0xffffffff, 0] * 3)))
lines.append(f'inout {work}/oob.init {work}/oob.bin')
lines.append('bindings 0 ' + ' '.join(str(3 + i) for i in range(n_images)) + ' 1 2' +
(f' {3 + n_images}' if probes or oob else ''))
''.join(f' {3 + n_images + i}' for i in range(len(extra))) +
(f' {3 + n_images + len(extra)}' if probes or oob else ''))
(work / 'manifest').write_text('\n'.join(lines) + '\n')
r = subprocess.run([o.kvk, o.shader, str(work / 'manifest'), str(o.time), str(o.subgroup)],
capture_output=True, text=True)
@@ -194,6 +203,13 @@ def main():
if arena is not None:
identical &= compare_arena((work / 'arena.out.bin').read_bytes(), arena_before,
launch.post_buffer(cap, arena))
for off in extra:
if not any(p.offset == off for p in launch.pointers):
continue
got, ref = (work / f'buf{off}.out.bin').read_bytes(), launch.post_buffer(cap, launch.pointer(off).alloc)
same = got[:len(ref)] == ref
print(f' buffer {off}: {len(ref)} bytes, ' + ('identical' if same else f'{sum(x != y for x, y in zip(got, ref))} differ'))
identical &= same
if oob:
r = struct.unpack('<9I', (work / 'oob.bin').read_bytes()[:36])
for k, name in enumerate(('weight reads', 'arena reads', 'arena writes')):
engine/rr/zluda_replay.py
+25 -5
View File
@@ -31,14 +31,18 @@ def main():
ap.add_argument('--ptx')
ap.add_argument('--zluda', default=os.path.expanduser('~/.cache/d4r-zluda-rr'))
ap.add_argument('--rocm', default='/opt/rocm/lib')
ap.add_argument('--tex', action='append', default=[], metavar='OFFSET=FILE:W:H',
help='an RGBA32F texture for a handle the capture has no dump of (preset F pooled histories)')
ap.add_argument('--pad', type=int, default=0, help='bytes of zeros appended to the arena (debug stores)')
o = ap.parse_args()
cap = Capture(o.capture)
launch = [l for l in cap.launches if o.kernel in l.kernel and 'diamond' not in l.kernel or
('diamond' in o.kernel and o.kernel in l.kernel)][o.frame - 1]
exact = [l for l in cap.launches if l.kernel == o.kernel]
launch = (exact or [l for l in cap.launches if o.kernel in l.kernel and 'diamond' not in l.kernel or
('diamond' in o.kernel and o.kernel in l.kernel)])[o.frame - 1]
work = Path(o.work)
work.mkdir(parents=True, exist_ok=True)
module = o.ptx or next(str(p) for p in sorted(Path('~/.cache/d4r-rr-corpus').expanduser().glob('*.ptx'))
corpus = '~/.cache/d4r-rr-corpus-f' if launch.kernel.startswith('rr2_') else '~/.cache/d4r-rr-corpus'
module = o.ptx or next(str(p) for p in sorted(Path(corpus).expanduser().glob('*.ptx'))
if f'.entry {launch.kernel}(' in p.read_text(errors='replace'))
(work / 'args.bin').write_bytes(launch.args)
arena = launch.arena_alloc
@@ -50,15 +54,31 @@ def main():
out = f' {work}/arena.out.bin' if p.alloc == arena else ''
source = f'{work}/arena-pad.bin' if o.pad and p.alloc == arena else f'{launch.directory}/alloc-{p.alloc}.bin'
lines.append(f'buf {p.offset} {source} {p.delta}{out}')
# textures the launch names but the bridge could not read back (preset F's 10:10:10:2 pooled histories, which
# ZLUDA backs with RGBA32F): zeros, which is what they hold on a first frame
have = {i.offset for i in launch.images}
for line in (launch.directory / 'manifest.txt').read_text().split('\n'):
f = line.split()
given = {int(t.split('=')[0]): t.split('=')[1].split(':') for t in o.tex}
if len(f) == 2 and f[0] == 'texture' and int(f[1]) in given:
path, w_, h_ = given[int(f[1])]
lines.append(f'tex {f[1]} {path} {w_} {h_} 4 32 linear')
elif len(f) == 2 and f[0] == 'texture' and int(f[1]) not in have:
(work / 'zero16.bin').write_bytes(bytes(16 * 16 * 16))
lines.append(f'tex {f[1]} {work}/zero16.bin 16 16 4 32 linear')
if len(f) == 2 and f[0] == 'surface' and int(f[1]) not in have: # half the largest surface: the pooled history
big = max((i for i in launch.images if i.kind == 'surface'), key=lambda i: i.width, default=None)
if big is not None:
lines.append(f'surf {f[1]} {big.width // 2} {big.height // 2} 4 32 {work}/surf{f[1]}.bin')
for image in launch.images:
if image.kind == 'texture':
# SNORM_INT8X4 (200): a SIGNED_INT8 array read as normalized floats gives the same values
fmt = 8 if image.format == 200 else image.format
fmt = 8 if image.format == 200 else 1 if image.format == 194 else image.format
lines.append(f'tex {image.offset} {image.path} {image.width} {image.height} {image.channels} '
f'{fmt} {"linear" if image.linear else "nearest"}')
else:
# ZLUDA has no SNORM_INT8X4 arrays (format 200); surface stores write the same bytes to a SIGNED_INT8 one
fmt = 8 if image.format == 200 else image.format
fmt = 8 if image.format == 200 else 1 if image.format == 194 else image.format # 194: UNORM_INT8X4
lines.append(f'surf {image.offset} {image.width} {image.height} {image.channels} {fmt} '
f'{work}/surf{image.offset}.bin')
(work / 'manifest').write_text('\n'.join(lines) + '\n')