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

engine/rr: native Ray Reconstruction stages with capture/replay and shim integration

Browse Source 
- engine/rr: d4r's own shaders for the denoiser frame (exposure, six encoders,
  five decoders, hkpn output stage) at NGX layouts, plus PTX->GLSL/HIP packing
  scripts, reference implementations and weight extraction
- runtime.h/runtime_rr.inc, game_rr.{h,cpp}: GameDenoiser wired into the engine
  build (d4r-rr replay binary) and the D3D12 shim's feature-13 path
- shim: answer GetFeatureRequirements from d4r's own state only (the old
  cross-binary CUDA handoff crashed OptiScaler's scan); drop the core's
  CUDA GetFeatureRequirements import
- shim/harness: exact guide capture/replay of a frame's RR inputs
- kernels/rr: expanded enc0/engine prototypes; kernels/tools: ptxsim rewrite
  plus rrsim and an engine test; vulkan: ptx2glsl FP8/cooperative-matrix
  support, kvk inout buffers and KVK_DISABLE_OPTIMIZATION
- bridge: report the surface array in cuSurfObjectCreate traces
- packaging: RR settings and the ICU launch overrides in D4R_README
This commit is contained in:
woutermans committed 2026-10-12 12:47:53 +02:00
1 parent 30472ae195
commit 50fa65d87b
67 files changed
+11926 -504
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

engine/rr/README.md
+173
View File
@@ -0,0 +1,173 @@
# Ray Reconstruction stages (DLSS-D, NGX feature 13)
d4r's own Vulkan compute shaders for the denoiser in `nvngx_dlssd.dll` 310.7 (presets D and E). One shader per
NVIDIA kernel of the frame; the engine records them in NVIDIA's order:
| stage | NVIDIA kernel | shader | grid (640x360 render, 1280x720 output) |
|---|---|---|---|
| exposure | `cuda_dldn_engine_luma_convert_kernel`, `reduce_sum` x2-3, `auto_exposure_copy_kernel` | `engine/exposure_k.comp` (MODE 1) | |
| enc0 | `cuda_dldn_engine_swin_enc0_kernel` (D: `..._enc0_diamond_wallaby_kernel`) | `enc0.comp` | 81x49, 32 threads |
| enc1-enc5 | `cuda_dldn_engine_swin_enc1_kernel` ... `enc5` | `enc1.comp` ... `enc5.comp` | |
| dec4-dec1 | `cuda_dldn_engine_swin_dec4_kernel` ... `dec1` | `dec4.comp` ... `dec1.comp` | |
| dec0 | `cuda_dldn_engine_swin_dec0_kernel` (D: `..._dec0_diamond_wallaby_kernel`) | `dec0.comp` | 81x49, 32 threads |
| output | `cuda_dldn_engine_hkpn_output_kernel_transformer` (D: `..._diamond_wallaby`) | `hkpn.comp` | 80x45, 16x16 threads |
The D variants of enc0, dec0 and the output kernel are different code; the other nine kernels are shared and only
their weights differ between presets.
## Stage contract
Every stage shader reads the same data NVIDIA's kernel reads, in the same layouts, so each one can be checked on
its own against a captured launch (`test_layer.py`) and the stages chain exactly like NGX's kernels:
- binding 0: the parameter block, `std430` readonly. It is NVIDIA's argument block byte for byte, except that every
pointer into the plane arena holds its **byte offset in the arena** (a 64-bit value), the weight pointer is zero,
and a texture or surface handle is 1 where NGX bound one and 0 where it left the argument null (the kernels test
optional inputs' handles for null). Fields are read at NVIDIA's argument offsets.
- binding 1: the layer's weight image exactly as NGX uploads it (argument 40's allocation), `std430` readonly.
- binding 2: the plane arena, `std430` read/write.
- bindings 3, 4, ...: the textures in ascending argument offset (combined image samplers, clamp-to-edge,
normalized coordinates, linear or nearest as NVIDIA's texture object; 8-bit unsigned arrays are sampled as UNORM,
`SNORM_INT8X4` as SNORM), then the surfaces in ascending argument offset (storage images, `rgba16f` or
`rgba8_snorm`).
- The dispatch grid is NVIDIA's launch grid. The workgroup size is the shader's choice.
## Checking a stage
Captures (bridge launch and replay dumps of a harness run, see `rr_capture.py`) live outside the repository.
```sh
glslangValidator -V --target-env vulkan1.3 engine/rr/enc1.comp -o build/rr/enc1.spv
python3 engine/rr/test_layer.py CAPTURE swin_enc1 2 build/rr/enc1.spv build/rr/work-enc1 [--time 50]
```
It starts from the arena as it was before NVIDIA's launch, with zeroed surfaces, and compares every surface and
the whole arena with what NVIDIA's launch left. `build/rr/kvk` is `vulkan/kvk.cpp`
(`g++ -O2 -std=c++17 -o build/rr/kvk vulkan/kvk.cpp -lvulkan`).
## Native stages (preset E)
`build_native.py CAPTURE MODEL_DIR` turns a `compile_rr.py` model directory into one where every Swin stage is
d4r's own cooperative-matrix shader instead of translated PTX. Each was derived phase by phase from ZLUDA replays
(`phase_dump.py`: every mma's operands of one window) into an exact numpy reference, then written against it:
| stages | reference | shader | weights |
|---|---|---|---|
| enc1-enc5, dec4-dec1 | `swin_ref.py` (one generic model: C, heads, merge / expand) | `swin_m.comp` | `swin_pack.py` |
| dec0 | `dec0_ref.py` | `dec0_m.comp` | `dec0_pack.py` |
| enc0 | `enc0_ref.py` (after the per-pixel features) | `enc0_m.comp` from `gen_enc0_m.py` | `enc0_pack.py` |
`enc1_m.comp` / `dec1_m.comp` are the hand-written C = 64 instances `swin_m.comp` generalises; `enc1_ref.py` holds
the shared numerics (inter, mm, the cubic softmax, the norm trees). enc0 keeps ptx2rr's translation of NVIDIA's feature
gather (waves 0-2 run its pixel loop) and is native from the embedding on.
Binding 1 differs from the stage contract above: it holds the packed weights (enc0: NGX's image, then the packed
blob). Everything else is the contract, so `test_layer.py --weights FILE` checks a native stage the same way. Against
capE3 every native stage leaves the arena byte-identical to NVIDIA's; surfaces differ in a few values by one binary16
ulp (`exp2` approximation, and decoder texture 88 where NVIDIA's bilinear unit rounds differently), and the whole
network's output is ~96 dB from NVIDIA's.
The shaders run every product transposed (weights x activations) so results feed the next product in registers;
`enc1_m.comp`'s header explains the layouts and `rr_exact.glsl` the binary16 rounding RADV needs (it folds
`float(float16_t(x))` and does not fuse f16 `fma`).
## Native models and preset D
`engine/rr/build_native_model.sh CAPTURE OUT_DIR [NVNGX_DLSSD_DLL]` builds a model whose eleven network stages are
the native shaders (`swin_m.comp`, `dec0_m.comp`, the generated `enc0_m.comp`) for the preset of the capture, with the
translated output stage.
The weights are read from the user's DLL itself (`dll_weights.py`), not from the capture. The DLL holds each model's
tensors as plain binary16 arrays in natural `[out][in]` order, stage after stage; nothing names them, so they are
found by structure: every attention head has a 100 x 100 relative position bias table whose entries depend only on
the offset between two window positions, and a model is eleven consecutive stages with 2 2 4 4 8 8 8 4 4 2 2 heads.
What NGX uploads is the same data rearranged: FP8 codes in the kernels' fragment order, the norm gains times
sqrt(C), the bias times 2 sqrt(2) for the kernels' 64 x 96 query / key order, and (preset E) the 4096-byte noise
table encoder 0 dithers with. `dll_weights.py DLL OUT --model N --check DIR` compares a reading with a phase-dump
extraction: for 310.7 every tensor of both presets is identical, and so are the packed weight files. The capture
only selects which of the DLL's models it ran (and supplies the two argument-block templates); a DLL whose network
kernels are not the ones the shaders were written against (`native_kernels.txt`, SHA-256 of the PTX modules) is
refused unless `D4R_RR_ANY_KERNELS=1`.
`compile_rr.py` uses the translated output kernel for both presets: `engine/rr/hkpn.comp` has the reset path only
(`--reset-only` selects it for captures that reset every frame).
Preset D (`..._diamond_wallaby` kernels) differs from E in three places:
- **enc0**: another feature gather (kept as the translation, like E's) with 30 input features for E's 32, and no
noise table (its weight image is 4096 bytes smaller).
- **dec0**: the eight bytes per pixel for the output stage are channels 12, 13, 0, 18, 19, 20, 21, 0
(`-DRR_PRESET_D`; E: 12, 13, 18-23). Everything else is E's kernel.
- **output stage**: its own kernel (closest-depth motion from texture 296, a 4x4 bicubic history fetch). It is the
translated stage; it has no subgroup operations and runs with 64-wide waves (`rr.txt`).
The reset flag (enc0 byte 365, output stage byte 440) follows `Reset`, not the jitter. With `Reset = 0` the output
stage reprojects and blends its history: check stages against a temporal capture (the harness with
`D4R_HARNESS_JITTER_SCENE=1`, and `D4R_CUDA_REPLAY_DUMP_LIMIT=400` so every launch is saved), not only against
captures that reset every frame. `engine/rr/hkpn.comp` implements the reset path only.
Preset D, RX 9070 XT, RADV 26.2.4, 2026-10-12, 1280x720 -> 3840x2160 (`d4r-rr`, `D4R_RR_TIME_REPEAT=20`): 5.0-5.2 ms
per frame (network 4.5, output 0.55-0.65), from 17.5 ms with every stage translated. Every network stage reproduces
the capture's arena byte for byte on reset and temporal frames. The presented image against the CUDA path over
temporal frames: 96-106 dB on the reset frame, then 41-52 dB (one-ulp differences in the filtered surfaces are
amplified by the recurrent feature; twelve frames stay at 41-45 dB).
RE Requiem (2026-10-12, 1280x720 -> 3840x2160, R9G9B9E5 colour, roughness packed in the normals' w): with
`hkpn.comp` the picture was 3-pixel blocks that shook with the jitter (no history); with the translated output stage
presets D and E both give a stable, reconstructed image on the engine. Not compared with the CUDA path in that game.
## Matching the reference bit for bit (2026-10-12)
The reference is the CUDA path through ZLUDA on the same GPU, so its `rcp.approx`, `ex2.approx` and bilinear texture
fetches are the GPU's own instructions and texture unit, not correctly rounded arithmetic. Two places in the native
filters had been "corrected" to exact arithmetic and were one ulp off in ~0.01-0.2% of the filtered surfaces:
- `rcp16` (the filters' softmax normalisation and gate) is now the plain hardware reciprocal;
- the decoders' texture 88 goes through the bilinear sampler at NVIDIA's coordinates (`rr_decfilter.glsl`,
`dec0_m.comp`) instead of a hand-computed 1/4, 3/4 blend.
With both, every network stage (arena and surfaces 112 / 120 / 128 / 144 / 304) is identical to the capture on every
frame of capDT12, capD4K, capD1080 (preset D) and capE3 (preset E). What remains is the translated output stage:
~0.004% of surface 224 / 240 one ulp off (single channels, values just above 0.5; cause not found). The output is
107-112 dB on frames where the network's inputs are still exact. At 640x360 -> 1080p that holds for all twelve
frames of capDT12 (was 96 dB, then 41-45); at 1280x720 and 1920x1080 -> 4K the history's differing values reach
encoder 0 through texture 496 and the output falls to 52 / 45 / 43 dB (capD4K), 53 / 49 dB (capD1080) as before.
Speed: the compiled stages spend about half their instructions on float conversions. `spv_exact_convert.py` marks
the conversions NoContraction so RADV keeps a binary16 rounding without the opaque OR (`RR_EXACT_CONVERT`), which is
bit-exact and ~8% faster: preset D 4.8 ms at 1280x720 -> 4K, 9.5 ms at 1920x1080 -> 4K (was 5.1 / 10.3).
`D4R_RR_RELAXED=1|2` builds drop NVIDIA's double rounding instead (8.2 / 7.6 ms at 1080p) but are not exact: about 1%
of a layer's FP8 codes change and the output is 39-44 dB from the reference on every frame.
Profile of a full-resolution stage at 1080p (dec0, 3.0 ms): expand 0.31, key norm 0.16, attention 1.37 (per head:
projections 0.15, scores + softmax 0.26, P V 0.17, out-projection 0.10), MLP 0.69, output filter 0.50. The 136
matrix products themselves are 0.45-0.5 ms; weight and table memory traffic is under 8%. enc0's translated feature
gather is 0.28 of its 3.0 ms.
## Optional inputs: specular hit distance and disocclusion mask (2026-10-12)
RE Requiem passes `DLSSD.SpecularHitDistance` (R32_FLOAT) and `DLSS.DisocclusionMask` (R16_FLOAT). Harness captures
with synthetic ones (`D4R_HARNESS_RR_EXTRA=hit,mask`; capDX, capDXhit, capDXmask) show the wiring:
| input | encoder 0 texture | encoder 0 flag byte | output stage |
|---|---|---|---|
| specular hit distance | argument 248 | 536 | flag byte 160 |
| disocclusion mask | argument 296 | 539 | - |
Encoder 0 also tests the handle at argument 248 for null before it takes the hit distance path (it writes a
reprojected specular motion into channels 2 and 3 of surface 304), which is why the parameter block now carries the
handles' presence. With both bound the network is identical to the capture on every frame (capDX). The output
stage then differs from the CUDA path on the image's border pixels only (63-67 dB on temporal frames): NVIDIA's
kernel fetches neighbours outside the image there, CUDA clamps such integer fetches to the edge (every denoiser
texture's address mode is clamp; the translation and rrsim do the same), and ZLUDA's image load returns zero. With
zero outside the image the stage matches the capture like before (99.997%); the engine keeps CUDA's clamp.
Exact speed-ups since the first native build, all bit-identical per stage: conversions kept by NoContraction instead of
an opaque OR, the MLP activation and its FP8 conversion from a 64 KiB table (`enc1_pack.act_table`), FP8 conversions in
pairs, no saturation clamp on attention probabilities. Preset D: 4.4 ms at 1280x720 -> 4K, 8.6 ms at 1920x1080 -> 4K.
Output stage, hypotheses tested for its remaining ~0.004% one-ulp values and rejected (2026-10-12): correctly rounded
exp2 / reciprocal (rrsim's exact arithmetic is further from the capture than the GPU), fusing the kernel's 34
single-use multiply-subtract pairs into fma, merging the two constant multiplications after the log transform,
`precise` on the inlined exp2, NoContraction on its conversions. The differing values come in spatial clusters, one
colour channel at a time, always one ulp, mostly just above 0.5 in surface 224.
engine/rr/build_native.py
+110
View File
@@ -0,0 +1,110 @@
#!/usr/bin/env python3
"""Turn a preset-E Ray Reconstruction model directory (engine/compile_rr.py) into the native one: every Swin stage
from d4r's own shaders instead of translated PTX.
usage: build_native.py CAPTURE MODEL_DIR [--frame N] [--work DIR]
enc0 engine/rr/enc0_m.comp, generated (gen_enc0_m.py) from the model's translated enc0.generated.comp,
whose per-pixel feature gather it keeps; weights: NGX's image + enc0_pack.py's blob
enc1-enc5, engine/rr/swin_m.comp with swin_pack.py's per-stage header and weights
dec4-dec1
dec0 engine/rr/dec0_m.comp with dec0_pack.py's weights
The weights are extracted from ZLUDA replays of the capture's launches (engine/rr/phase_dump.py: every mma's operands
of one window) and the stages' weight images; extraction checks itself against the dump. Like the model directory,
everything written is derived from the user's DLL and stays on the machine. Runs one replay or compile at a time
(phase_dump.py caps its replay at 24 GB).
"""
import argparse
import os
import shutil
import subprocess
import sys
from pathlib import Path
import numpy as np
HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
from rr_capture import Capture # noqa: E402
import swin_ref # noqa: E402
import dec0_ref # noqa: E402
import enc0_ref # noqa: E402
SWIN = ['enc1', 'enc2', 'enc3', 'enc4', 'enc5', 'dec4', 'dec3', 'dec2', 'dec1']
def run(*cmd, **kw):
subprocess.run([str(c) for c in cmd], check=True, **kw)
def glsl(source, output, *includes):
run('glslangValidator', '-V', '--target-env', 'vulkan1.3', *[f'-I{i}' for i in includes], source, '-o', output,
stdout=subprocess.DEVNULL)
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('model')
ap.add_argument('--frame', type=int, default=2)
ap.add_argument('--work', help='dumps and build files (default MODEL/native-work)')
o = ap.parse_args()
model = Path(o.model)
work = Path(o.work) if o.work else model / 'native-work'
work.mkdir(parents=True, exist_ok=True)
cap = Capture(os.path.expanduser(o.capture))
manifest = (model / 'rr.txt').read_text().split('\n')
def dump(stage):
out = work / f'{stage}_all.npz'
launch = cap.find(f'swin_{stage}', o.frame)
block = '5,5' if launch.grid[0] > 5 and launch.grid[1] > 5 else '1,1'
if not out.exists():
run(sys.executable, HERE / 'phase_dump.py', o.capture, stage, o.frame, block, out)
return out, block
def stage_line(stage, shader, weights):
for i, line in enumerate(manifest):
f = line.split()
if len(f) > 4 and f[0] == 'stage' and f[1] == stage:
f[2], f[4] = shader, weights
manifest[i] = ' '.join(x for x in f if x != 'noopt')
for name in ('rr_exact.glsl', 'rr_attn.glsl', 'rr_decfilter.glsl', 'enc0_m_decl.glsl', 'enc0_m_tail.glsl'):
shutil.copy(HERE / name, work / name)
for stage in SWIN:
d, block = dump(stage)
print(f'{stage}: extracting', flush=True)
l, ex, ey, sx, sy, inp, outs, tex = swin_ref.load(o.capture, stage, o.frame)
bx, by = (int(v) for v in block.split(','))
swin_ref.extract(stage, d, l.weights, work / f'{stage}_w.npz', inp, 8 * bx - sx, 8 * by - sy, ex, ey)
sd = work / stage
sd.mkdir(exist_ok=True)
run(sys.executable, HERE / 'swin_pack.py', stage, work / f'{stage}_w.npz', model / f'weights-{stage}m.bin',
sd / 'swin_stage.h')
glsl(HERE / 'swin_m.comp', model / f'{stage}_m.spv', sd, work)
stage_line(stage, f'{stage}_m.spv', f'weights-{stage}m.bin')
# dec0
d, _ = dump('dec0')
print('dec0: extracting', flush=True)
dec0_ref.extract(d, cap.find('swin_dec0', o.frame).weights, work / 'dec0_w.npz')
run(sys.executable, HERE / 'dec0_pack.py', work / 'dec0_w.npz', model / 'weights-dec0m.bin', work / 'dec0_m_w.h')
glsl(HERE / 'dec0_m.comp', model / 'dec0_m.spv', work)
stage_line('dec0', 'dec0_m.spv', 'weights-dec0m.bin')
# enc0: NGX's weight image (the translated feature gather reads it) followed by the native blob
d, _ = dump('enc0')
print('enc0: extracting', flush=True)
nv = cap.find('swin_enc0', o.frame).weights
enc0_ref.extract(d, nv, work / 'enc0_w.npz')
base = (len(nv) + 255) // 256 * 256
run(sys.executable, HERE / 'enc0_pack.py', work / 'enc0_w.npz', work / 'enc0_native.bin', work / 'enc0_m_w.h', base)
(model / 'weights-enc0m.bin').write_bytes(nv + bytes(base - len(nv)) + (work / 'enc0_native.bin').read_bytes())
run(sys.executable, HERE / 'gen_enc0_m.py', model / 'enc0.generated.comp', work / 'enc0_m.comp')
glsl(work / 'enc0_m.comp', model / 'enc0_m.spv', work)
stage_line('enc0', 'enc0_m.spv', 'weights-enc0m.bin')
(model / 'rr.txt').write_text('\n'.join(manifest))
print(f'{model / "rr.txt"}: native stages written')
if __name__ == '__main__':
main()
engine/rr/build_native_model.sh
Executable
+98
View File
@@ -0,0 +1,98 @@
#!/usr/bin/env bash
set -euo pipefail
# Builds a Ray Reconstruction model directory whose network stages are the native shaders (engine/rr/*_m.comp),
# for the preset the capture was taken with (D or E).
#
# usage: build_native_model.sh CAPTURE_DIR OUT_DIR [NVNGX_DLSSD_DLL]
#
# CAPTURE_DIR is a bridge capture of a few denoiser frames (engine/rr/rr_capture.py); the DLL is the one it was taken
# with (default: CAPTURE_DIR/run/app/nvngx_dlssd.dll, then $D4R_RR_DLL). The steps:
# 1. engine/compile_rr.py: the translated model (the argument-block templates, the exposure passes, the translated
# output stage and the translated encoder 0 whose feature gather the native one keeps)
# 2. engine/rr/dll_weights.py: the network's weights, read from the DLL itself. The capture only selects which of
# the DLL's models it ran, and a DLL whose kernels are not the ones the native shaders were written against is
# refused (D4R_RR_ANY_KERNELS=1 builds anyway)
# 3. per stage: the weights checked against the capture's frame (STAGE_ref.py), packed (STAGE_pack.py) and the
# native shader compiled against that packing's header
# 4. rr.txt rewritten to the native stages; the weight images NGX uploaded in the capture are removed
# Preset D differs from E in three kernels: encoder 0 (its feature gather and 30 input features for E's 32), decoder 0
# (the eight bytes it hands the output stage: -DRR_PRESET_D) and the output stage, which is the translated kernel
# (it has no subgroup operations, so it runs with 64-wide waves).
# Everything written is derived from the user's own nvngx_dlssd.dll and stays on their machine.
#
# GPU jobs run one at a time, each in a memory-capped scope (a translated stage's pipeline compile can take tens of GB).
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
CAP="$(realpath "${1:?usage: build_native_model.sh CAPTURE_DIR OUT_DIR}")"
OUT="$(realpath -m "${2:?usage: build_native_model.sh CAPTURE_DIR OUT_DIR}")"
RR="$ROOT/engine/rr"
WORK="$OUT/native-work"
DLL="${3:-${D4R_RR_DLL:-$CAP/run/app/nvngx_dlssd.dll}}"
[[ -f "$DLL" ]] || { echo "no nvngx_dlssd.dll at $DLL: pass the DLL the capture was taken with" >&2; exit 1; }
FRAME="${D4R_RR_MODEL_FRAME:-2}"
CAPPED=(systemd-run --user --scope -q -p MemoryMax=24G -p MemorySwapMax=0)
mkdir -p "$OUT" "$WORK"
cd "$ROOT"
python3 engine/compile_rr.py "$CAP" "$OUT" > "$WORK/compile_rr.log" 2>&1 || { tail -5 "$WORK/compile_rr.log"; exit 1; }
python3 "$RR/dll_weights.py" "$DLL" "$WORK/dll" --capture "$CAP" ${D4R_RR_ANY_KERNELS:+--any-kernels}
PRESET="$(awk '$1 == "preset" { print $2 }' "$OUT/rr.txt")"
# RR_EXACT_CONVERT: the compiled module's float conversions are decorated NoContraction (spv_exact_convert.py), which
# keeps the driver from folding a binary16 rounding away without the opaque OR the shaders otherwise need
DEFS=(-DRR_EXACT_CONVERT); [[ "$PRESET" == 4 ]] && DEFS+=(-DRR_PRESET_D)
# D4R_RR_RELAXED=1: one rounding where NVIDIA has two (engine/rr/rr_exact.glsl): faster, not bit-exact
[[ -n "${D4R_RR_RELAXED:-}" ]] && DEFS+=("-DRR_RELAXED=$D4R_RR_RELAXED")
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
"${CAPPED[@]}" glslangValidator --target-env vulkan1.3 -V "${DEFS[@]}" "${inc[@]}" -o "$out" "$src" > "$WORK/glsl.log" 2>&1 ||
{ grep -E 'ERROR' "$WORK/glsl.log" | head -5; exit 1; }
python3 "$RR/spv_exact_convert.py" "$out" > /dev/null
}
native() { # native STAGE WEIGHT_FILE: point rr.txt's stage line at STAGE_m.spv
sed -i -E "s/^stage $1 $1\.spv 32 weights-$1\.bin (.*)/stage $1 $1_m.spv 32 $2 \1/; /^stage $1 /s/ noopt\$//" "$OUT/rr.txt"
grep -q "^stage $1 $1_m.spv " "$OUT/rr.txt" || { echo "rr.txt: no stage line for $1" >&2; exit 1; }
}
# the Swin stages shared by both presets (engine/rr/swin_m.comp)
for s in enc1 enc2 enc3 enc4 enc5 dec4 dec3 dec2 dec1; do
b=1,1; case $s in enc1|enc2|dec1|dec2) b=5,5 ;; esac
d="$WORK/$s"; mkdir -p "$d"
cp "$WORK/dll/$s.npz" "$d/weights.npz"
python3 "$RR/swin_ref.py" "$CAP" "$s" "$FRAME" "$d/weights.npz" --blocks "$b" 0,0 | tail -2 | sed "s/^/$s: /"
python3 "$RR/swin_pack.py" "$s" "$d/weights.npz" "$OUT/weights-${s}m.bin" "$d/swin_stage.h" > /dev/null
glsl "$OUT/${s}_m.spv" "$RR/swin_m.comp" "$d" "$RR"
native "$s" "weights-${s}m.bin"
done
# decoder 0
d="$WORK/dec0"; mkdir -p "$d"
cp "$WORK/dll/dec0.npz" "$d/weights.npz"
python3 "$RR/dec0_ref.py" "$CAP" "$FRAME" "$d/weights.npz" --blocks 5,5 0,0 | tail -2 | sed 's/^/dec0: /'
python3 "$RR/dec0_pack.py" "$d/weights.npz" "$OUT/weights-dec0m.bin" "$d/dec0_m_w.h" > /dev/null
glsl "$OUT/dec0_m.spv" "$RR/dec0_m.comp" "$d" "$RR"
native dec0 weights-dec0m.bin
# encoder 0: the translated feature gather + the native rest
d="$WORK/enc0"; mkdir -p "$d"
cp "$WORK/dll/enc0.npz" "$d/weights.npz"
# preset E's gather dithers with the noise table that starts NGX's weight image: the table, then the packing
BASE=0; [[ "$PRESET" == 4 ]] || BASE=4096
python3 "$RR/enc0_pack.py" "$d/weights.npz" "$d/enc0_m.bin" "$d/enc0_m_w.h" "$BASE" > /dev/null
if [[ "$BASE" != 0 ]]; then
[[ "$(stat -c %s "$WORK/dll/noise.bin" 2>/dev/null)" == 4096 ]] || { echo "no noise table found in $DLL" >&2; exit 1; }
cat "$WORK/dll/noise.bin" "$d/enc0_m.bin" > "$OUT/weights-enc0m.bin"
else
cp "$d/enc0_m.bin" "$OUT/weights-enc0m.bin"
fi
python3 "$RR/gen_enc0_m.py" "$OUT/enc0.generated.comp" "$d/enc0_m.comp" > /dev/null
glsl "$OUT/enc0_m.spv" "$d/enc0_m.comp" "$d" "$RR"
native enc0 weights-enc0m.bin
# the translated output stage (preset D; preset E's when engine/rr/hkpn.comp is not used) has no subgroup operations
if grep -q '"shader": "generated from the DLL"' <(python3 -c "import json,sys; print(json.dumps(json.load(open('$OUT/model.json'))['stages']['hkpn']))"); then
sed -i -E 's/^stage hkpn hkpn\.spv 32 /stage hkpn hkpn.spv 64 /' "$OUT/rr.txt"
fi
for s in enc0 enc1 enc2 enc3 enc4 enc5 dec4 dec3 dec2 dec1 dec0; do rm -f "$OUT/weights-$s.bin"; done
printf 'Native Ray Reconstruction model (preset %s) in %s\n' "$([[ "$PRESET" == 4 ]] && echo D || echo E)" "$OUT"
grep '^stage' "$OUT/rr.txt" | awk '{ printf " %s %s (%s-wide)\n", $2, $3, $4 }'
engine/rr/dec0_m.comp
+348
View File
@@ -0,0 +1,348 @@
#version 460
// Native decoder 0 of the DLSS-D (preset E) denoiser, register-resident: cuda_dldn_engine_swin_dec0_kernel computed
// like engine/rr/enc1_m.comp (whose header explains the transposed products and the slot orders). The arithmetic is
// engine/rr/dec0_ref.py's (bit-exact to NVIDIA except surface 112's bilinear taps, which NVIDIA's texture unit rounds
// differently in a few percent of the values by one ulp); binding 1 holds engine/rr/dec0_pack.py's packing.
//
// One workgroup per 8x8-token window at the render resolution, four waves:
// expand the 6 x 6 low-resolution tokens under the window and its border (dec1's 64 channels) -> 4 sub-positions
// x 32 channels; wave w computes sub-position w, FP8 into EX in slot order
// prologue one key token per thread: expanded + skip (enc0's tokens), RMS norm -> XT; window tokens -> XR (f16)
// block two 32-channel heads and four MLP branches of 32 -> 32 -> 32, as enc1_m with C = 32
// outputs one token per thread: 8 FP8 channels for the output stage (argument 48, [y][x][8]), tanh of four
// channels to surface 144, and the kernel-prediction filter of texture 88 (bilinear from half resolution)
// blended with textures 96 / 104 to surface 112
#extension GL_KHR_cooperative_matrix : require
#extension GL_KHR_memory_scope_semantics : require
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_8bit_storage : require
#extension GL_EXT_float_e4m3 : require
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_EXT_control_flow_attributes : require
#extension GL_GOOGLE_include_directive : require
layout(local_size_x = 128) in;
#include "dec0_m_w.h"
#include "rr_exact.glsl"
layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param;
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16;
layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u;
layout(set = 0, binding = 1, std430) readonly buffer W8U { uint8_t w[]; } w8u;
layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
layout(set = 0, binding = 3) uniform sampler2D T88;
layout(set = 0, binding = 4) uniform sampler2D T96;
layout(set = 0, binding = 5) uniform sampler2D T104;
layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S112;
layout(set = 0, binding = 7, rgba8_snorm) uniform writeonly image2D S144;
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define HAC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CM gl_CooperativeMatrixLayoutColumnMajor
#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM)
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
#define Q8(dst, src) Q8PAIRS(dst, src)
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) }
const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0;
// shared memory, aliased by phase: expand and prologue [LOW | EX], then the attention [K | V^T]
shared uint POOL[48 * 16 + 48 * 32];
#define LOW POOL // the low-resolution tokens (6 x 6, padded to 48): 64 FP8 channels, arena order
#define EX_AT 768u // their expansion: 4 sub-positions x 32 FP8 channels in slot order (32 words a row)
#define RR_KS POOL // the current head's K
#define RR_KS0 0u
#define RR_VS POOL // the current head's V^T
#define RR_VS0 768u
shared uint XT[96 * 8]; // the normalised key tokens: slot-major, 32 FP8 channels in slot order
shared float16_t XR[64 * 32]; // the window tokens before the norm (the residual), slot-major
shared float16_t TAP[100 * 4]; // texture 88 bilinear at the window's 10 x 10 filter tap positions
#define OUTF XR // after the block: the final tokens, token-major, natural channels (f16)
// the window's key tokens: the 8x8 window, then a border row above, below, a column left and right
ivec2 keyXY(uint t, int x0, int y0)
{
if (t < 64u) return ivec2(x0 + int(t & 7u), y0 + int(t >> 3u));
const int i = int(t & 7u), side = int((t - 64u) >> 3u);
return ivec2(side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8), side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i));
}
int clampTap(int v, int n) { return v < 0 ? 1 : (v >= n ? n - 2 : v); }
float tanhNV(float t)
{
const float m = abs(t), s = t * t;
const float big = m >= 9.0109138 ? 1.0 : fma(1.0 / (exp2(m * 2.885390043258667) + 1.0), -2.0, 1.0);
return m >= 0.6 ? (t < 0.0 ? -big : big) :
fma(fma(fma(fma(fma(0.01573968306183815, s, -0.05230396240949631), s, 0.13315297663211823), s, -0.33332768082618713), s, 0.0), t, t);
}
void main()
{
const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID;
const uint hl = lane >> 4u, col = lane & 15u;
RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u);
const int ex = int(param.p[0]), ey = int(param.p[1]), hx = ex / 2, hy = ey / 2;
const uint loOff = param.p[2], skOff = param.p[6], outOff = param.p[12], shift = param.p[20];
const int x0 = 8 * int(gl_WorkGroupID.x) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u);
const int X0 = (x0 >> 1) - 1, Y0 = (y0 >> 1) - 1;
const uint lplane = uint(hx * hy) * 32u;
// ---------------------------------------------------------------- expand
for (uint i = tid; i < 48u * 16u; i += 128u) {
const uint r = i >> 4u, wd = i & 15u;
uint u = 0u;
if (r < 36u) {
const int X = mirror(X0 + int(r % 6u), hx), Y = mirror(Y0 + int(r / 6u), hy);
u = arena.w[(loOff + (wd >> 3u) * lplane + uint(Y * hx + X) * 32u) / 4u + (wd & 7u)];
}
LOW[i] = u;
}
barrier();
{
A8 we[2][4];
[[unroll]] for (uint mm = 0u; mm < 2u; ++mm)
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) LDA(we[mm][kt], W_WE, 64u, 2u * wv + mm, kt);
[[unroll]] for (uint n = 0u; n < 3u; ++n) {
B8 lb[4];
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) 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)]);
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;
Q8(q, a)
coopMatStore(q, POOL, EX_AT + n * 512u + 4u * m, 32u, CM);
}
}
}
barrier();
// ---------------------------------------------------------------- prologue: one key token per thread
if (tid < 96u) {
const ivec2 k = keyXY(inter(tid), x0, y0);
const uint r = uint(((k.y >> 1) - Y0) * 6 + ((k.x >> 1) - X0)), s = uint(2 * (k.y & 1) + (k.x & 1));
const int xx = mirror(k.x, ex), yy = mirror(k.y, ey);
const uint sb = (skOff + uint(yy * ex + xx) * 32u) / 4u;
float16_t v[32];
[[unroll]] for (uint wd = 0u; wd < 8u; ++wd) {
const uint e = POOL[EX_AT + r * 32u + 8u * s + wd], sk = arena.w[sb + wd];
[[unroll]] for (uint b = 0u; b < 4u; ++b) {
precise float16_t x = float16_t(uintBitsToFloate4m3EXT(uint8_t((e >> (8u * b)) & 0xffu)))
+ float16_t(uintBitsToFloate4m3EXT(uint8_t((sk >> (8u * b)) & 0xffu)));
v[4u * wd + b] = x;
}
}
// dec0_ref.norm_plane: per word t the pair sums, words t and t + 4, then t ^ 1, t ^ 2, the halves
precise f16vec2 h[4];
[[unroll]] for (uint t = 0u; t < 4u; ++t) {
precise f16vec2 p0 = f16vec2(v[4u * t], v[4u * t + 1u]) * f16vec2(v[4u * t], v[4u * t + 1u])
+ f16vec2(v[4u * t + 2u], v[4u * t + 3u]) * f16vec2(v[4u * t + 2u], v[4u * t + 3u]);
precise f16vec2 p1 = f16vec2(v[16u + 4u * t], v[17u + 4u * t]) * f16vec2(v[16u + 4u * t], v[17u + 4u * t])
+ f16vec2(v[18u + 4u * t], v[19u + 4u * t]) * f16vec2(v[18u + 4u * t], v[19u + 4u * t]);
h[t] = p0 + p1;
}
precise f16vec2 h0 = h[0] + h[1], h2_ = h[2] + h[3];
precise f16vec2 hs = h0 + h2_;
precise float16_t sm = hs.x + hs.y;
precise float16_t se = sm + float16_t(0.0001220703125);
const float16_t rr = rsq16(se);
[[unroll]] for (uint wd = 0u; wd < 8u; ++wd) {
floate4m3_t c[4];
[[unroll]] for (uint b = 0u; b < 4u; ++b) {
const uint j = 4u * wd + b;
precise float16_t g = rr * w16.w[W_GAIN0 / 2u + inter(j)];
precise float16_t n = v[j] * g;
c[b] = q8h(n);
}
XT[tid * 8u + wd] = code(c[0]) | (code(c[1]) << 8u) | (code(c[2]) << 16u) | (code(c[3]) << 24u);
}
if (tid < 64u) [[unroll]] for (uint j = 0u; j < 32u; ++j) XR[tid * 32u + j] = v[j];
}
barrier();
// ---------------------------------------------------------------- the residual: window tokens + table 0
ACC 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);
}
}
// ---------------------------------------------------------------- attention, two heads
[[unroll]] for (uint hd = 0u; hd < 2u; ++hd) {
if (hd == 1u) barrier();
for (uint j = 0u; j < 3u; ++j) { // 12 jobs (token tile n, channel tile c)
const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u;
A8 wk[2], xa[2];
B8 wvb[2], xb[2];
[[unroll]] for (uint kt = 0u; kt < 2u; ++kt) {
coopMatLoad(xb[kt], XT, n * 128u + 4u * kt, 8u, CM);
coopMatLoad(xa[kt], XT, n * 128u + 4u * kt, 8u, RM);
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);
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);
K32(v, xa[0], wvb[0], xa[1], wvb[1])
A8 va;
Q8(va, v)
coopMatStore(va, POOL, 768u + c * 384u + 4u * n, 24u, RM);
}
barrier();
#include "rr_attn.glsl"
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
A8 a0, a1;
LDA(a0, W_WO0 + HEAD * hd, 32u, m, 0u); LDA(a1, W_WO0 + HEAD * hd, 32u, m, 1u);
K32(R[m], a0, Ob[0], a1, Ob[1])
}
}
// ---------------------------------------------------------------- MLP
B8 Nb[2];
{
float16_t rv[2][8];
[[unroll]] for (uint m = 0u; m < 2u; ++m) [[unroll]] for (int i = 0; i < 8; ++i) rv[m][i] = float16_t(R[m][i]);
// dec0_ref.norm_plane on slot order: lane half hl holds words 2hl, 2hl + 1 (tile 0) and 4 + 2hl, 5 + 2hl (tile 1)
float16_t h1[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
float16_t ht[2];
[[unroll]] for (uint tq = 0u; tq < 2u; ++tq) {
const uint i0 = 4u * tq + c, i1 = i0 + 2u;
precise float16_t p0 = rv[0][i0] * rv[0][i0] + rv[0][i1] * rv[0][i1];
precise float16_t p1 = rv[1][i0] * rv[1][i0] + rv[1][i1] * rv[1][i1];
precise float16_t hv = p0 + p1;
ht[tq] = hv;
}
precise float16_t hh = ht[0] + ht[1];
h1[c] = hh;
}
const uint other = subgroupShuffleXor(bh(h1[0]) | (bh(h1[1]) << 16u), 16u);
precise float16_t g0 = h1[0] + hb(other & 0xffffu), g1 = h1[1] + hb(other >> 16u);
precise float16_t s = g0 + g1;
precise float16_t se = s + float16_t(0.0001220703125);
const float16_t r = rsq16(se);
[[unroll]] for (uint m = 0u; m < 2u; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) {
const uint k = 16u * m + 8u * hl + uint(i);
precise float16_t g = r * w16.w[W_GAIN1 / 2u + k];
precise float16_t n = rv[m][i] * g;
Nb[m][i] = q8h(n);
precise float16_t a = rv[m][i] + w16.w[W_TABLE1 / 2u + k];
R[m][i] = float(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)]);
A8 a0, a1;
LDA(a0, w1, 32u, c, 0u); LDA(a1, w1, 32u, c, 1u);
K32(h, a0, Nb[0], a1, Nb[1])
[[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = ACTQ(float16_t(h[i]));
}
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
A8 a0, a1;
LDA(a0, w2, 32u, m, 0u); LDA(a1, w2, 32u, m, 1u);
K32(R[m], a0, Hb[0], a1, Hb[1])
}
}
// ---------------------------------------------------------------- outputs, one token per thread
{
const uint t = inter(16u * wv + col);
[[unroll]] for (uint m = 0u; m < 2u; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) OUTF[t * 32u + inter(16u * m + 8u * hl + uint(i))] = float16_t(R[m][i]);
}
barrier();
// texture 88 (one level down) at the window's 10 x 10 full-resolution tap positions (clamped like the filter's
// coordinates), through the bilinear sampler at NVIDIA's coordinates: the reference's texture unit does not round
// like f32 arithmetic (computing the 1/4, 3/4 blend by hand left ~0.2% of surface 112 one ulp off, capDT 2026-10-12)
if (tid < 100u) {
const int xi = clampTap(x0 - 1 + int(tid % 10u), ex), yi = clampTap(y0 - 1 + int(tid / 10u), ey);
const float iw = uintBitsToFloat(param.p[34]), ih = uintBitsToFloat(param.p[35]);
precise float cu = fma(iw, float(xi), iw * 0.5), cv = fma(ih, float(yi), ih * 0.5);
const vec4 bl = textureLod(T88, vec2(cu, cv), 0.0);
[[unroll]] for (uint c = 0u; c < 3u; ++c) TAP[tid * 4u + c] = float16_t(bl[c]);
}
barrier();
// two threads per token: half 0 filters channels 0 and 1, half 1 channel 2 and writes the output stage's codes
// and surface 144
{
const uint t = tid >> 1u, hf = tid & 1u;
const int x = x0 + int(t & 7u), y = y0 + int(t >> 3u);
if (x < 0 || y < 0 || x >= ex || y >= ey) return;
float16_t a[10];
[[unroll]] for (uint c = 0u; c < 10u; ++c) a[c] = OUTF[t * 32u + c];
if (hf == 1u) {
float16_t b[12];
[[unroll]] for (uint c = 0u; c < 12u; ++c) b[c] = OUTF[t * 32u + 12u + c];
const uint o = (outOff + uint(y * ex + x) * 8u) / 4u;
#ifdef RR_PRESET_D
// cuda_dldn_engine_swin_dec0_diamond_wallaby_kernel: channels 12, 13, 0, 18, 19, 20, 21, 0 (its only
// difference from the preset E kernel; capDT, 2026-10-12)
arena.w[o] = code(q8h(b[0])) | (code(q8h(b[1])) << 8u) | (code(q8h(b[6])) << 24u);
arena.w[o + 1u] = code(q8h(b[7])) | (code(q8h(b[8])) << 8u) | (code(q8h(b[9])) << 16u);
#else
arena.w[o] = code(q8h(b[0])) | (code(q8h(b[1])) << 8u) | (code(q8h(b[6])) << 16u) | (code(q8h(b[7])) << 24u);
arena.w[o + 1u] = code(q8h(b[8])) | (code(q8h(b[9])) << 8u) | (code(q8h(b[10])) << 16u) | (code(q8h(b[11])) << 24u);
#endif
imageStore(S144, ivec2(x, y), vec4(tanhNV(float(b[2])), tanhNV(float(b[3])), tanhNV(float(b[4])), tanhNV(float(b[5]))));
}
// surface 112: softmax of channels 0-8 over the 3x3 taps, one f16 fma chain per channel; the gate (channel 9)
// blends it with texture 96; + texture 104
const float16_t ma = max(max(max(a[0], a[2]), a[4]), a[6]), mb = max(max(max(a[1], a[3]), a[5]), a[7]);
const float16_t mx = max(max(ma, mb), a[8]);
float e[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) { precise float16_t d = a[i] - mx; e[i] = exp2(float(d) * 1.4426950216293335); }
precise float s = (e[0] + e[1]) + 0.0;
s = s + (e[2] + e[3]);
s = s + (e[4] + e[5]);
s = s + (e[6] + e[7]);
s = s + e[8];
const float16_t rr = rcp16(s);
float16_t wg[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) wg[i] = h16(r16(e[i]) * float(rr));
const uint tx = uint(x - x0), ty = uint(y - y0), c0 = 2u * hf, nc = 2u - hf;
float16_t f[2] = float16_t[2](0.0hf, 0.0hf);
bool mid = false;
[[unroll]] for (uint k = 0u; k < 9u; ++k) {
const uint at = ((ty + k / 3u) * 10u + tx + k % 3u) * 4u + c0;
[[unroll]] for (uint c = 0u; c < 2u; ++c) f[c] = fma16q(TAP[at + c], wg[k], f[c], mid);
}
if (mid) { // an f32 result on a binary16 midpoint: redo the chain with the exact fma (rare)
f[0] = 0.0hf; f[1] = 0.0hf;
for (uint k = 0u; k < 9u; ++k) {
const uint at = ((ty + k / 3u) * 10u + tx + k % 3u) * 4u + c0;
for (uint c = 0u; c < 2u; ++c) f[c] = fma16(TAP[at + c], wg[k], f[c]);
}
}
const float16_t g = sig(a[9]);
precise float16_t om = float16_t(1.0) - g;
const vec4 t96 = texelFetch(T96, ivec2(x, y), 0), t104 = texelFetch(T104, ivec2(x, y), 0);
float16_t o[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
precise float16_t fg = f[c] * g;
float16_t v = fma16(float16_t(t96[c0 + c]), om, fg);
const float16_t t4 = float16_t(t104[c0 + c]);
if (!isnan(t4)) { precise float16_t vv = v + t4; v = vv; }
o[c] = v;
}
const uint other = subgroupShuffleXor(bh(o[0]), 1u);
if (hf == 0u) imageStore(S112, ivec2(x, y), vec4(o[0], o[1], hb(other), 0.0));
}
}
engine/rr/dec0_pack.py
+63
View File
@@ -0,0 +1,63 @@
#!/usr/bin/env python3
"""Pack decoder 0's extracted weights (engine/rr/dec0_ref.py --extract) for engine/rr/dec0_m.comp.
usage: dec0_pack.py WEIGHTS.npz OUT.bin OUT.h
The layout rules are enc1_pack.py's: transposed products, every weight's columns in the k order of the product that
consumes its result (C = 32: slot j holds natural channel inter(32)[j]). The expand's 128 columns are four
sub-positions of 32 channels, each permuted the same way, so a key token's expanded half is already in slot order.
"""
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from enc1_ref import inter # noqa: E402
from enc1_pack import a_tiles, b_tiles, cubic_table, act_table # noqa: E402
I32, I96 = inter(32), inter(96)
I128 = [32 * (j // 32) + I32[j % 32] for j in range(128)]
def main():
src, out_bin, out_h = sys.argv[1:4]
w = dict(np.load(src))
e = {'we': ('A', w['we'][:, I128]), 'be': ('H', w['be'][I128])}
for h in range(2):
e[f'wk{h}'] = ('A', w[f'wk{h}'][:, I32])
e[f'wv{h}'] = ('B', w[f'wv{h}'][:, I32])
e[f'wo{h}'] = ('A', w[f'wo{h}'][:, I32])
e[f'bias{h}'] = ('H', w[f'bias{h}'][np.ix_(inter(64), I96)].T)
for b in range(4):
e[f'w1_{b}'] = ('A', w[f'w1_{b}'][:, I32])
e[f'b1_{b}'] = ('H', w[f'b1_{b}'][I32])
e[f'w2_{b}'] = ('A', w[f'w2_{b}'][:, I32])
e['gain0'] = ('H', w['gain0'])
for k in ('table0', 'gain1', 'table1'):
e[k] = ('H', w[k][I32])
e['cubic'] = ('U', cubic_table())
e['act'] = ('Y', act_table())
blob, defines = bytearray(), []
for name, (kind, a) in e.items():
while len(blob) % 256:
blob.append(0)
off = len(blob)
if kind == 'A':
blob += a_tiles(a).tobytes()
elif kind == 'B':
blob += b_tiles(a).tobytes()
elif kind == 'Y':
blob += np.ascontiguousarray(a, np.uint8).tobytes()
elif kind == 'U':
blob += np.ascontiguousarray(a, np.uint16).tobytes()
else:
blob += np.ascontiguousarray(a, np.float16).tobytes()
defines.append(f'#define W_{name.upper()} {off}u // {kind} {a.shape}')
open(out_bin, 'wb').write(bytes(blob))
open(out_h, 'w').write('// generated by engine/rr/dec0_pack.py\n' + '\n'.join(defines) + '\n')
print(f'{len(blob)} bytes, {len(e)} entries')
if __name__ == '__main__':
main()
engine/rr/dec0_ref.py
+258
View File
@@ -0,0 +1,258 @@
#!/usr/bin/env python3
"""Exact numpy reference of the denoiser's decoder 0 (cuda_dldn_engine_swin_dec0_kernel), one 8x8 window at a time,
from its weights and real inputs only: the specification the native shader implements. Derived like
engine/rr/enc1_ref.py, phase by phase against ZLUDA replays (engine/rr/phase_dump.py).
expand the low-resolution tokens under the window and its border (6 x 6 of dec1's 64-channel output, arena order)
-> 128 = 4 sub-positions x 32 channels (+ bias), FP8; sub-position 2 (y & 1) + (x & 1)
prologue key token = expanded + skip token (enc0's 32-channel output at argument 24), f16; RMS norm (enc1's per-word
tree on one plane) times the gain at W+8448; the residual is the token + the table at W+39232
attention two heads of 32 channels (K = 32, one k32 step), as enc1
MLP RMS norm, four branches of 32 -> 32 (+ bias, ramp activation) -> 32 onto the residual (+ table)
usage: dec0_ref.py CAPTURE FRAME WEIGHTS.npz [--blocks X,Y ...] (verifies the phases of a dump, see --dump)
dec0_ref.py CAPTURE FRAME OUT.npz --extract DUMP.npz (weights from a phase dump)
"""
import argparse
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rrmath import e4m3, f16, q8 # noqa: E402
from enc1_ref import inter, mm, softmax_cubic, act, mirror, add, H, F # noqa: E402
I32 = inter(32)
UN32 = [I32.index(i) for i in range(32)]
I96 = inter(96)
def norm_plane(xn, gain):
"""RMS norm of one 32-channel token (natural order) in NVIDIA's tree over the arena order; returns codes in k order"""
xa = xn[I32]
sq = (xa * xa).astype(H)
h = []
for t in range(4):
def pair(w):
return add(sq[4 * w:4 * w + 2], sq[4 * w + 2:4 * w + 4])
h.append(add(pair(t), pair(t + 4)))
h01 = [add(h[t], h[t ^ 1]) for t in range(4)]
h2 = add(h01[0], h01[2])
s = H(H(h2[0] + h2[1]) + H(2.0 ** -13))
r = H(F(1) / np.sqrt(F(s)))
return q8((xn * (r * gain).astype(H)).astype(H))[I32]
def keys(x0, y0):
border = ([(x, y0 - 1) for x in range(x0, x0 + 8)] + [(x, y0 + 8) for x in range(x0, x0 + 8)]
+ [(x0 - 1, y) for y in range(y0, y0 + 8)] + [(x0 + 8, y) for y in range(y0, y0 + 8)])
return [(x, y) for y in range(y0, y0 + 8) for x in range(x0, x0 + 8)] + border
def window(w, lo, skip, x0, y0, ex, ey, dbg=None):
"""lo: dec1's output [2][ey/2][ex/2][32] codes; skip: enc0's [ey][ex][32] codes. Returns the residual stream after
the block (64 x 32 f16, natural channels)"""
dbg = {} if dbg is None else dbg
hx, hy = ex // 2, ey // 2
X0, Y0 = x0 // 2 - 1, y0 // 2 - 1
low = np.zeros((48, 64), np.uint8)
for r in range(36):
X, Y = mirror(X0 + r % 6, hx), mirror(Y0 + r // 6, hy)
low[r] = np.concatenate([lo[0, Y, X], lo[1, Y, X]])
Ex = q8(mm(low, w['we'], np.repeat(w['be'][None, :], 48, 0)))
dbg['p1_A'] = low
xs, ns = [], []
for (x, y) in keys(x0, y0):
xx, yy = mirror(x, ex), mirror(y, ey)
r = (yy // 2 - Y0) * 6 + (xx // 2 - X0) if False else (y // 2 - Y0) * 6 + (x // 2 - X0)
s = 2 * (y & 1) + (x & 1)
xn = (e4m3(Ex[r, 32 * s:32 * s + 32]).astype(H) + e4m3(skip[yy, xx][UN32]).astype(H)).astype(H)
xs.append(xn)
ns.append(norm_plane(xn, w['gain0']))
X = np.array(ns)
dbg['p2_A'] = X
R = None
for h in range(2):
qk = mm(X, w[f'wk{h}'], np.zeros((96, 32)))
v = mm(X, w[f'wv{h}'], np.zeros((96, 32)))
Q = q8(qk)[:, I32]
S = mm(Q[:64], Q.T, w[f'bias{h}'])
P = softmax_cubic(S)
O = mm(q8(P)[:, I96], q8(v)[I96], np.zeros((64, 32)))
C = add(np.array(xs[:64]), w['table0'][None, :]) if h == 0 else R
R = mm(q8(O)[:, I32], w[f'wo{h}'], C)
dbg[f'p{4 + 5 * h}_A'], dbg[f'p{5 + 5 * h}_A'], dbg[f'p{6 + 5 * h}_A'] = Q[:64], q8(P)[:, I96], q8(O)[:, I32]
dbg[f'p{6 + 5 * h}_C'], dbg[f'p{6 + 5 * h}_D'] = C, R
N = np.array([norm_plane(R[t], w['gain1']) for t in range(64)])
acc = add(R, w['table1'][None, :])
for b in range(4):
hid = act(mm(N, w[f'w1_{b}'], np.repeat(w[f'b1_{b}'][None, :], 64, 0)))
dbg[f'p{12 + 2 * b}_A'] = N
dbg[f'p{13 + 2 * b}_A'], dbg[f'p{13 + 2 * b}_C'] = q8(hid)[:, I32], acc
acc = mm(q8(hid)[:, I32], w[f'w2_{b}'], acc)
dbg[f'p{13 + 2 * b}_D'] = acc
return acc
OUT_CH = [12, 13, 18, 19, 20, 21, 22, 23] # the 8 FP8 channels of the output stage's input (argument 48)
OUT_CH_D = [12, 13, None, 18, 19, 20, 21, None] # preset D (..._dec0_diamond_wallaby_kernel): bytes 2 and 7 are zero
TANH_CH = [14, 15, 16, 17] # surface 144 (rgba8 snorm)
def fma16(a, b, c):
return H(np.float64(a) * np.float64(b) + np.float64(c))
def tanh_nv(t):
"""NVIDIA's tanh of the surface-144 channels: an odd polynomial below 0.6, 1 - 2 / (e^2x + 1) above, f32"""
t = F(t)
m = abs(t)
big = F(1) if m >= F(9.0109138) else F(np.float64(F(1) / F(np.exp2(F(m * F(2.8853900432586670))) + F(1))) * -2 + 1)
s = F(t * t)
p = F(0.01573968306183815)
for c in (-0.05230396240949631, 0.13315297663211823, -0.33332768082618713, 0.0):
p = F(np.float64(p) * np.float64(s) + np.float64(F(c)))
small = F(np.float64(p) * np.float64(t) + np.float64(t))
return np.copysign(big, t) if m >= F(0.6) else small
def bilinear88(t88, xi, yi):
"""texture 88 (half resolution, linear) at the full-resolution texel centre (xi, yi): weights 1/4 and 3/4,
clamped at the edge; f32 sum rounded to binary16 (NVIDIA's texture unit rounds a few of these differently)"""
def axis(i):
k = i // 2
return ((k - 1, 0.25), (k, 0.75)) if i % 2 == 0 else ((k, 0.75), (k + 1, 0.25))
h, w = t88.shape[:2]
v = np.zeros(4)
for (yy, wy) in axis(yi):
for (xx, wx) in axis(xi):
v += wx * wy * t88[min(max(yy, 0), h - 1), min(max(xx, 0), w - 1)].astype(np.float64)
return v.astype(H)
def outputs(acc, t88, t96, t104, x0, y0, ex, ey):
"""per window token: the 8 output codes, surface 144 (4 floats before the snorm store), surface 112 (3 f16)"""
codes, s144, s112 = [], [], []
for t in range(64):
x, y = x0 + t % 8, y0 + t // 8
a = acc[t].astype(H)
codes.append(q8(a)[OUT_CH])
s144.append([tanh_nv(a[c]) for c in TANH_CH])
lg = a[:9]
m01 = np.maximum(np.maximum(np.maximum(lg[0:2], lg[2:4]), lg[4:6]), lg[6:8])
m = max(max(m01[0], m01[1]), lg[8])
e = [np.exp2(F(H(lg[i] - m)) * F(1.4426950216293335)) for i in range(9)]
pair = [F(e[2 * k] + e[2 * k + 1]) for k in range(4)]
sm = F(pair[0] + F(0))
for k in (1, 2, 3):
sm = F(sm + pair[k])
sm = F(sm + e[8])
r = H(F(1) / sm)
wg = [H(H(e[i]) * r) for i in range(9)]
def cl(v, n):
return 1 if v < 0 else (n - 2 if v >= n else v)
xs, ys = [cl(x - 1, ex), cl(x, ex), cl(x + 1, ex)], [cl(y - 1, ey), cl(y, ey), cl(y + 1, ey)]
f = np.zeros(3, H)
k = 0
for yy in ys:
for xx in xs:
T = bilinear88(t88, xx, yy)
for c in range(3):
f[c] = fma16(T[c], wg[k], f[c])
k += 1
g = H(F(1) / (np.exp2(F(a[9]) * F(-1.4426950216293335)) + F(1)))
om = H(H(1) - g)
o = np.zeros(3, H)
for c in range(3):
if 0 <= x < ex and 0 <= y < ey:
o[c] = fma16(t96[y, x, c], om, H(f[c] * g))
if not np.isnan(t104[y, x, c]):
o[c] = H(o[c] + t104[y, x, c])
s112.append(o)
return np.array(codes), np.array(s144), np.array(s112)
def extract(dump, weights, out):
z = np.load(dump)
W = np.frombuffer(weights, np.uint8)
def B(p):
return z[f'p{p}_w0_B']
def C(p):
return f16(z[f'p{p}_w0_C']).astype(H)
w = {'we': B(1), 'be': C(1)[0], 'gain0': W[8448:8448 + 64].view(H), 'table0': W[39232:39232 + 64].view(H),
'gain1': W[39296:39296 + 64].view(H), 'table1': W[41472:41472 + 64].view(H)}
for h, base in ((0, 2), (1, 7)):
w[f'wk{h}'], w[f'wv{h}'] = B(base), B(base + 1)
w[f'bias{h}'] = C(base + 2)
w[f'wo{h}'] = B(base + 4)
for b in range(4):
w[f'w1_{b}'], w[f'b1_{b}'] = B(12 + 2 * b), C(12 + 2 * b)[0]
w[f'w2_{b}'] = B(13 + 2 * b)
np.savez_compressed(out, **w)
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('frame', type=int)
ap.add_argument('weights')
ap.add_argument('--extract', metavar='DUMP')
ap.add_argument('--dump', help='phase dump of --blocks[0] to verify phase by phase')
ap.add_argument('--blocks', nargs='*', default=['5,5'])
o = ap.parse_args()
from rr_capture import Capture
c = Capture(os.path.expanduser(o.capture))
l = c.find('swin_dec0', o.frame)
if o.extract:
extract(o.extract, l.weights, o.weights)
return
ex, ey = (int(v) for v in np.frombuffer(l.args[:8], np.uint32))
sx, sy = (int(v) for v in np.frombuffer(l.args[80:84], np.uint16))
ptr = {p.offset: p.delta for p in l.pointers}
pre = np.frombuffer(l.pre(l.arena_alloc), np.uint8)
lo = pre[ptr[8]:ptr[8] + 2 * (ex // 2) * (ey // 2) * 32].reshape(2, ey // 2, ex // 2, 32)
skip = pre[ptr[24]:ptr[24] + ex * ey * 32].reshape(ey, ex, 32)
w = dict(np.load(o.weights))
img = {i.offset: i for i in l.images}
def surf(off, dt=np.float16):
i = img[off]
return np.frombuffer(i.path.read_bytes(), dt).reshape(i.height, i.width, i.channels)
t88, t96, t104, r112, r144 = surf(88), surf(96), surf(104), surf(112), surf(144, np.int8)
post = np.frombuffer(l.post_buffer(c, l.arena_alloc), np.uint8)
rout = post[ptr[48]:ptr[48] + ex * ey * 8].reshape(ey, ex, 8)
for blk in o.blocks:
bx, by = (int(v) for v in blk.split(','))
d = {}
x0, y0 = 8 * bx - sx, 8 * by - sy
acc = window(w, lo, skip, x0, y0, ex, ey, d)
codes, s144, s112 = outputs(acc, t88, t96, t104, x0, y0, ex, ey)
if 'diamond' in l.kernel:
codes = np.array([[0 if ch is None else q8(acc[t].astype(H))[ch] for ch in OUT_CH_D] for t in range(64)])
n = ok_c = ok_s = ok_t = 0
for t in range(64):
x, y = x0 + t % 8, y0 + t // 8
if not (0 <= x < ex and 0 <= y < r144.shape[0]):
continue
n += 1
ok_c += np.array_equal(codes[t], rout[y, x])
ok_t += np.array_equal(np.clip(np.round(np.array(s144[t], np.float64) * 127), -127, 127).astype(np.int8), r144[y, x])
ok_s += np.array_equal(s112[t], r112[y, x, :3])
print(f'block {bx},{by}: output codes {ok_c}/{n}, surface 144 {ok_t}/{n}, surface 112 {ok_s}/{n}')
if o.dump:
z = np.load(o.dump)
for k in sorted(d, key=lambda s: (int(s[1:].split('_')[0]), s)):
ref = z[f'{k.split("_")[0]}_w0_{k.split("_")[1]}']
ours = np.asarray(d[k])
ours = ours.view(np.uint16) if ours.dtype == H else ours
n = min(len(ref), len(ours))
print(f'{k}: rows equal {int((ref[:n] == ours[:n]).all(1).sum())}/{n}')
o.dump = None
if __name__ == '__main__':
main()
engine/rr/dec1_m.comp
+295
View File
@@ -0,0 +1,295 @@
#version 460
// Native decoder 1 of the DLSS-D (preset E) denoiser, register-resident: cuda_dldn_engine_swin_dec1_kernel.
// engine/rr/dec1_ref.py's arithmetic (bit-exact to NVIDIA except surface 112's bilinear taps, see dec0_m.comp);
// binding 1 holds engine/rr/dec1_pack.py's packing. The expand and prologue are dec0_m.comp's (64 channels from dec2's
// 96), the block is enc1_m.comp's (whose header explains the layouts), the outputs are the tokens (argument 48) and
// dec0_m.comp's filter on the 64 -> 16 head.
#extension GL_KHR_cooperative_matrix : require
#extension GL_KHR_memory_scope_semantics : require
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_8bit_storage : require
#extension GL_EXT_float_e4m3 : require
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_EXT_control_flow_attributes : require
#extension GL_GOOGLE_include_directive : require
layout(local_size_x = 128) in;
#include "dec1_m_w.h"
#include "rr_exact.glsl"
layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param;
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16;
layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u;
layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
layout(set = 0, binding = 3) uniform sampler2D T88;
layout(set = 0, binding = 4) uniform sampler2D T96;
layout(set = 0, binding = 5) uniform sampler2D T104;
layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S112;
#define SKIP 0
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define HAC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CM gl_CooperativeMatrixLayoutColumnMajor
#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM)
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
#define Q8(dst, src) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (dst)[i_] = q8s((src)[i_]);
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) }
const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0;
// shared memory, aliased by phase: expand [LOW | EX], attention [KS | VS], outputs [TAP | OUTF]
shared uint POOL[48 * 24 + 48 * 64];
#define LOW POOL // 48 low-resolution tokens (6 x 6, padded): 96 FP8 channels, arena order (24 words)
#define EX_AT (48u * 24u) // their expansion: 4 sub-positions x 64 FP8 channels in slot order (64 words)
#define KS POOL // the current head's K: token slot x 32 FP8 channels
#define VS_AT 768u // the current head's V^T: 32 channels x 96 token slots
#define RR_KS POOL
#define RR_KS0 0u
#define RR_VS POOL
#define RR_VS0 VS_AT
shared uint XT[96 * 16]; // the normalised key tokens: slot-major, 64 FP8 channels in slot order
shared float16_t TAP[100 * 4]; // texture 88 bilinear at the window's 10 x 10 filter tap positions
shared float16_t OUTF[64 * 16]; // the head per window token (raster order)
ivec2 keyXY(uint t, int x0, int y0)
{
if (t < 64u) return ivec2(x0 + int(t & 7u), y0 + int(t >> 3u));
const int i = int(t & 7u), side = int((t - 64u) >> 3u);
return ivec2(side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8), side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i));
}
int clampTap(int v, int n) { return v < 0 ? 1 : (v >= n ? n - 2 : v); }
uint pack4(floate4m3_t a, floate4m3_t b, floate4m3_t c, floate4m3_t d)
{
return code(a) | (code(b) << 8u) | (code(c) << 16u) | (code(d) << 24u);
}
float16_t e4(uint w, uint b) { return float16_t(uintBitsToFloate4m3EXT(uint8_t((w >> (8u * b)) & 0xffu))); }
void main()
{
const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID;
const uint hl = lane >> 4u, col = lane & 15u;
RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u);
const int ex = int(param.p[0]), ey = int(param.p[1]), hx = ex / 2, hy = ey / 2;
const uint loOff = param.p[2], skOff = param.p[6], outOff = param.p[12], shift = param.p[20];
const int x0 = 8 * int(gl_WorkGroupID.x) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u);
const int X0 = (x0 >> 1) - 1, Y0 = (y0 >> 1) - 1;
const uint lplane = uint(hx * hy) * 32u, plane = uint(ex * ey) * 32u;
// ---------------------------------------------------------------- expand: wave w computes sub-position w
for (uint i = tid; i < 48u * 24u; i += 128u) {
const uint r = i / 24u, wd = i % 24u;
uint u = 0u;
if (r < 36u) {
const int X = mirror(X0 + int(r % 6u), hx), Y = mirror(Y0 + int(r / 6u), hy);
u = arena.w[(loOff + (wd >> 3u) * lplane + uint(Y * hx + X) * 32u) / 4u + (wd & 7u)];
}
LOW[i] = u;
}
barrier();
[[unroll]] for (uint n = 0u; n < 3u; ++n) {
B8 lb[6];
[[unroll]] for (uint kt = 0u; kt < 6u; ++kt) coopMatLoad(lb[kt], POOL, n * 384u + 4u * kt, 24u, CM);
[[unroll]] for (uint mm = 0u; mm < 4u; ++mm) {
const uint m = 4u * wv + mm;
ACC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_BE / 2u + 16u * m + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < 6u; kt += 2u) {
A8 w0, w1;
LDA(w0, W_WE, 96u, m, kt); LDA(w1, W_WE, 96u, m, kt + 1u);
K32(a, w0, lb[kt], w1, lb[kt + 1u])
}
B8 q;
Q8(q, a)
coopMatStore(q, POOL, EX_AT + n * 1024u + 4u * m, 64u, CM);
}
}
barrier();
// ---------------------------------------------------------------- prologue: one key token per thread
if (tid < 96u) {
const ivec2 k = keyXY(inter(tid), x0, y0);
const uint r = uint(((k.y >> 1) - Y0) * 6 + ((k.x >> 1) - X0)), s = uint(2 * (k.y & 1) + (k.x & 1));
const int xx = mirror(k.x, ex), yy = mirror(k.y, ey);
const uint sb = (skOff + uint(yy * ex + xx) * 32u) / 4u;
float16_t v[64];
[[unroll]] for (uint wd = 0u; wd < 16u; ++wd) {
const uint e = POOL[EX_AT + r * 64u + 16u * s + wd], sk = arena.w[sb + (wd >> 3u) * (plane >> 2u) + (wd & 7u)];
[[unroll]] for (uint b = 0u; b < 4u; ++b) { precise float16_t x = e4(e, b) + e4(sk, b); v[4u * wd + b] = x; }
}
// enc1's input norm: per word t (and t + 4) the f16x2 pair sums of both planes, then lanes ^1, ^2, the halves
precise f16vec2 h[4];
[[unroll]] for (uint q = 0u; q < 4u; ++q) {
precise f16vec2 p[4];
[[unroll]] for (uint kk = 0u; kk < 4u; ++kk) {
const uint o = 32u * (kk & 1u) + 4u * (q + 4u * (kk >> 1u));
precise f16vec2 a = f16vec2(v[o], v[o + 1u]) * f16vec2(v[o], v[o + 1u]);
precise f16vec2 c = f16vec2(v[o + 2u], v[o + 3u]) * f16vec2(v[o + 2u], v[o + 3u]);
p[kk] = a + c;
}
precise f16vec2 s0 = p[0] + p[1], s1 = p[2] + p[3];
h[q] = s0 + s1;
}
precise f16vec2 h0 = h[0] + h[1], h2_ = h[2] + h[3];
precise f16vec2 hs = h0 + h2_;
precise float16_t sm = hs.x + hs.y;
precise float16_t se = sm + float16_t(0.0001220703125);
const float16_t rr = rsq16(se);
[[unroll]] for (uint wd = 0u; wd < 16u; ++wd) {
floate4m3_t c[4];
[[unroll]] for (uint b = 0u; b < 4u; ++b) {
const uint j = 4u * wd + b;
precise float16_t g = rr * w16.w[W_GAIN0 / 2u + inter(j)];
precise float16_t n = v[j] * g;
c[b] = q8h(n);
}
XT[tid * 16u + wd] = pack4(c[0], c[1], c[2], c[3]);
}
}
barrier();
// ---------------------------------------------------------------- the residual: expanded + skip + table 7
ACC R[4];
{
const uint tq = inter(16u * wv + col);
const int qx = x0 + int(tq & 7u), qy = y0 + int(tq >> 3u);
const uint r = uint(((qy >> 1) - Y0) * 6 + ((qx >> 1) - X0)), s = uint(2 * (qy & 1) + (qx & 1));
const uint sb = (skOff + uint(mirror(qy, ey) * ex + mirror(qx, ex)) * 32u) / 4u;
[[unroll]] for (uint m = 0u; m < 4u; ++m) {
const uint ea = EX_AT + r * 64u + 16u * s + 4u * m + 2u * hl;
const uint sa = sb + (m >> 1u) * (plane >> 2u) + 4u * (m & 1u) + 2u * hl;
const uvec2 e = uvec2(POOL[ea], POOL[ea + 1u]), sk = uvec2(arena.w[sa], arena.w[sa + 1u]);
[[unroll]] for (int i = 0; i < 8; ++i) {
precise float16_t x = e4(e[i >> 2], uint(i) & 3u) + e4(sk[i >> 2], uint(i) & 3u);
precise float16_t c = x + w16.w[W_TABLE7 / 2u + 16u * m + 8u * hl + uint(i)];
R[m][i] = float(c);
}
}
}
barrier(); // EX is read; the attention reuses its memory
// ---------------------------------------------------------------- attention, two heads
[[unroll]] for (uint hd = 0u; hd < ((SKIP & 64) != 0 ? 0u : 2u); ++hd) {
// this head's Q/K and V, shared: 12 jobs (token tile n, channel tile c), three per wave. K = (QK^T)^T rows
// to KS (token slot x 32 channels, also Q^T read column-major), V^T to VS (32 channels x 96 token slots)
if (hd == 1u) barrier(); // every wave is done with head 0's KS / VS
for (uint j = 0u; j < ((SKIP & 128) != 0 ? 0u : 3u); ++j) {
const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u;
A8 wk[4];
B8 wvb[4], xb[4];
A8 xa[4];
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) {
coopMatLoad(xb[kt], XT, n * 256u + 4u * kt, 16u, CM);
coopMatLoad(xa[kt], XT, n * 256u + 4u * kt, 16u, RM);
LDA(wk[kt], W_WK0 + HEAD * hd, 64u, c, kt);
LDB(wvb[kt], W_WV0 + HEAD * hd, 32u, kt, c);
}
ACC a = ACC(0.0);
K32(a, wk[0], xb[0], wk[1], xb[1])
K32(a, wk[2], xb[2], wk[3], xb[3])
A8 ka;
Q8(ka, a)
coopMatStore(ka, KS, n * 128u + 4u * c, 8u, RM);
ACC v = ACC(0.0);
K32(v, xa[0], wvb[0], xa[1], wvb[1])
K32(v, xa[2], wvb[2], xa[3], wvb[3])
A8 va;
Q8(va, v)
coopMatStore(va, POOL, VS_AT + c * 384u + 4u * n, 24u, RM);
}
barrier();
#include "rr_attn.glsl"
[[unroll]] for (uint m = 0u; m < 4u; ++m) {
A8 a0, a1;
LDA(a0, W_WO0 + HEAD * hd, 32u, m, 0u); LDA(a1, W_WO0 + HEAD * hd, 32u, m, 1u);
K32(R[m], a0, Ob[0], a1, Ob[1])
}
}
// ---------------------------------------------------------------- MLP
B8 Nb[4];
{
float16_t rv[4][8];
[[unroll]] for (uint m = 0u; m < 4u; ++m) [[unroll]] for (int i = 0; i < 8; ++i) rv[m][i] = float16_t(R[m][i]);
// slot 16m + 8hl + i is NVIDIA's channel 8j + 2t + c with j = 4(m >> 1) + (m & 1) + 2((i >> 1) & 1),
// t = 2hl + (i >> 2), c = i & 1 (enc1_ref.nv_rowsum_sq's tree)
float16_t h1[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
float16_t ht[2];
[[unroll]] for (uint tq = 0u; tq < 2u; ++tq) {
const uint i0 = 4u * tq + c, i1 = i0 + 2u;
precise float16_t a = rv[0][i0] * rv[0][i0] + rv[0][i1] * rv[0][i1];
precise float16_t b = rv[1][i0] * rv[1][i0] + rv[1][i1] * rv[1][i1];
precise float16_t cc = rv[2][i0] * rv[2][i0] + rv[2][i1] * rv[2][i1];
precise float16_t d = rv[3][i0] * rv[3][i0] + rv[3][i1] * rv[3][i1];
precise float16_t hv = (a + cc) + (b + d);
ht[tq] = hv;
}
precise float16_t hh = ht[0] + ht[1];
h1[c] = hh;
}
const uint other = subgroupShuffleXor(bh(h1[0]) | (bh(h1[1]) << 16u), 16u);
precise float16_t g0 = h1[0] + hb(other & 0xffffu), g1 = h1[1] + hb(other >> 16u);
precise float16_t s = g0 + g1;
precise float16_t se = s + float16_t(0.0001220703125);
const float16_t r = rsq16(se);
[[unroll]] for (uint m = 0u; m < 4u; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) {
const uint k = 16u * m + 8u * hl + uint(i);
precise float16_t g = r * w16.w[W_GAIN14 / 2u + k];
precise float16_t n = rv[m][i] * g;
Nb[m][i] = q8h(n);
precise float16_t a = rv[m][i] + w16.w[W_TABLE16 / 2u + k];
R[m][i] = float(a);
}
}
for (uint br = 0u; br < ((SKIP & 4) != 0 ? 1u : 8u); ++br) {
const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br;
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
ACC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = float(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]);
A8 a0, a1, a2, a3;
LDA(a0, w1, 64u, c, 0u); LDA(a1, w1, 64u, c, 1u); LDA(a2, w1, 64u, c, 2u); LDA(a3, w1, 64u, c, 3u);
K32(h, a0, Nb[0], a1, Nb[1])
K32(h, a2, Nb[2], a3, Nb[3])
[[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = q8h(act(float16_t(h[i])));
}
[[unroll]] for (uint m = 0u; m < 4u; ++m) {
A8 a0, a1;
LDA(a0, w2, 32u, m, 0u); LDA(a1, w2, 32u, m, 1u);
K32(R[m], a0, Hb[0], a1, Hb[1])
}
}
// ---------------------------------------------------------------- outputs
{
const uint t = inter(16u * wv + col);
const int tx = x0 + int(t & 7u), ty = y0 + int(t >> 3u);
const bool inside = tx >= 0 && ty >= 0 && tx < ex && ty < ey;
B8 Tb[4];
const uint o = (outOff + uint(ty * ex + tx) * 32u) >> 2u;
[[unroll]] for (uint m = 0u; m < 4u; ++m) {
Q8(Tb[m], R[m])
const uint at = 4u * (m & 1u) + 2u * hl;
if (inside) {
arena.w[o + (m >> 1u) * (plane >> 2u) + at] = pack4(Tb[m][0], Tb[m][1], Tb[m][2], Tb[m][3]);
arena.w[o + (m >> 1u) * (plane >> 2u) + at + 1u] = pack4(Tb[m][4], Tb[m][5], Tb[m][6], Tb[m][7]);
}
}
// the 64 -> 16 head, per token to OUTF
ACC hv;
[[unroll]] for (int i = 0; i < 8; ++i) hv[i] = float(w16.w[W_BH / 2u + 8u * hl + uint(i)]);
A8 a0, a1, a2, a3;
LDA(a0, W_WH, 64u, 0u, 0u); LDA(a1, W_WH, 64u, 0u, 1u); LDA(a2, W_WH, 64u, 0u, 2u); LDA(a3, W_WH, 64u, 0u, 3u);
K32(hv, a0, Tb[0], a1, Tb[1])
K32(hv, a2, Tb[2], a3, Tb[3])
[[unroll]] for (int i = 0; i < 8; ++i) OUTF[t * 16u + 8u * hl + uint(i)] = float16_t(hv[i]);
}
barrier();
#include "rr_decfilter.glsl"
}
engine/rr/dec1_pack.py
+60
View File
@@ -0,0 +1,60 @@
#!/usr/bin/env python3
"""Pack decoder 1's extracted weights (engine/rr/dec1_ref.py --extract) for engine/rr/dec1_m.comp.
usage: dec1_pack.py WEIGHTS.npz OUT.bin OUT.h
enc1_pack.py's layout (transposed products, columns in the consumer's k order) plus the expand (dec0_pack.py): 256
columns, four sub-positions of 64 channels each permuted to slot order.
"""
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from enc1_ref import inter # noqa: E402
from enc1_pack import a_tiles, b_tiles, cubic_table # noqa: E402
I32, I64, I96 = inter(32), inter(64), inter(96)
I256 = [64 * (j // 64) + I64[j % 64] for j in range(256)]
def main():
src, out_bin, out_h = sys.argv[1:4]
w = dict(np.load(src))
e = {'we': ('A', w['we'][:, I256]), 'be': ('H', w['be'][I256])}
for h in range(2):
e[f'wk{h}'] = ('A', np.vstack([w[f'wk{h}a'], w[f'wk{h}b']])[:, I32])
e[f'wv{h}'] = ('B', np.vstack([w[f'wv{h}a'], w[f'wv{h}b']])[:, I32])
e[f'wo{h}'] = ('A', w[f'wo{h}'][:, I64])
e[f'bias{h}'] = ('H', w[f'bias{h}'][np.ix_(I64, I96)].T)
for b in range(8):
e[f'w1_{b}'] = ('A', w[f'w1_{b}'][:, I32])
e[f'b1_{b}'] = ('H', w[f'b1_{b}'][I32])
e[f'w2_{b}'] = ('A', w[f'w2_{b}'][:, I64])
e['wh'], e['bh'] = ('A', w['wh']), ('H', w['bh'])
e['gain0'] = ('H', w['gain0'])
for k in ('table7', 'gain14', 'table16'):
e[k] = ('H', w[k][I64])
e['cubic'] = ('U', cubic_table())
blob, defines = bytearray(), []
for name, (kind, a) in e.items():
while len(blob) % 256:
blob.append(0)
off = len(blob)
if kind == 'A':
blob += a_tiles(a).tobytes()
elif kind == 'B':
blob += b_tiles(a).tobytes()
elif kind == 'U':
blob += np.ascontiguousarray(a, np.uint16).tobytes()
else:
blob += np.ascontiguousarray(a, np.float16).tobytes()
defines.append(f'#define W_{name.upper()} {off}u // {kind} {a.shape}')
open(out_bin, 'wb').write(bytes(blob))
open(out_h, 'w').write('// generated by engine/rr/dec1_pack.py\n' + '\n'.join(defines) + '\n')
print(f'{len(blob)} bytes, {len(e)} entries')
if __name__ == '__main__':
main()
engine/rr/dec1_ref.py
+183
View File
@@ -0,0 +1,183 @@
#!/usr/bin/env python3
"""Exact numpy reference of the denoiser's decoder 1 (cuda_dldn_engine_swin_dec1_kernel), one 8x8 window at a time.
Derived like engine/rr/enc1_ref.py, phase by phase against ZLUDA replays (engine/rr/phase_dump.py); NVIDIA runs it with
two warps, each one head and half of the expand's columns.
expand the 6 x 6 level-2 tokens under the window and its border (dec2's 96 channels, arena order) -> 256 =
4 sub-positions x 64 channels (+ bias), FP8; sub-position 2 (y & 1) + (x & 1) (dec0_ref)
prologue key token = expanded + skip token (enc1's 64-channel output at argument 24), f16; enc1's input norm (two
planes) times the gain at W+25088; the residual is the token + the table at W+62080
block enc1's: two 32-channel heads, eight MLP branches 64 -> 32 -> 64 (gain W+62208, table W+66496)
outputs the tokens (argument 48); the 64 -> 16 head and dec0's filter (bilinear texture 88 from level 2, gate,
textures 96 / 104) to surface 112
usage: dec1_ref.py CAPTURE FRAME WEIGHTS.npz [--blocks X,Y ...] (verifies against the capture)
dec1_ref.py CAPTURE FRAME OUT.npz --extract DUMP.npz (weights from a phase dump)
"""
import argparse
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rrmath import e4m3, f16, q8 # noqa: E402
from enc1_ref import inter, mm, softmax_cubic, act, mirror, add, input_norm, nv_rowsum_sq, H, F # noqa: E402
from dec0_ref import keys, fma16, bilinear88 # noqa: E402
I32, I64, I96 = inter(32), inter(64), inter(96)
UN64 = [I64.index(i) for i in range(64)]
GAIN0, TABLE7, GAIN14, TABLE16 = 25088, 62080, 62208, 66496
def window(w, lo, skip, x0, y0, ex, ey):
"""lo: dec2's output [3][ey/2][ex/2][32] codes; skip: enc1's [2][ey][ex][32]. Returns (tokens 64 x 64 codes in arena
order, head 64 x 16 f16)"""
hx, hy = ex // 2, ey // 2
X0, Y0 = x0 // 2 - 1, y0 // 2 - 1
low = np.zeros((48, 96), np.uint8)
for r in range(36):
X, Y = mirror(X0 + r % 6, hx), mirror(Y0 + r // 6, hy)
low[r] = np.concatenate([lo[p, Y, X] for p in range(3)])
Ex = q8(mm(low, w['we'], np.repeat(w['be'][None, :], 48, 0)))
xs, ns = [], []
for (x, y) in keys(x0, y0):
xx, yy = mirror(x, ex), mirror(y, ey)
r = (y // 2 - Y0) * 6 + (x // 2 - X0)
s = 2 * (y & 1) + (x & 1)
sk = np.concatenate([skip[0, yy, xx], skip[1, yy, xx]])[UN64]
xn = (e4m3(Ex[r, 64 * s:64 * s + 64]).astype(H) + e4m3(sk).astype(H)).astype(H)
xs.append(xn)
xa = xn[I64]
n, _ = input_norm([xa[:32], xa[32:]], w['gain0'])
ns.append(q8(n)[I64])
X = np.array(ns)
res = np.array(xs[:64])
R = None
for h in range(2):
qk = mm(X[:, 32:], w[f'wk{h}b'], mm(X[:, :32], w[f'wk{h}a'], np.zeros((96, 32))))
v = mm(X[:, 32:], w[f'wv{h}b'], mm(X[:, :32], w[f'wv{h}a'], np.zeros((96, 32))))
Q = q8(qk)[:, I32]
S = mm(Q[:64], Q.T, w[f'bias{h}'])
P = softmax_cubic(S)
O = mm(q8(P)[:, I96], q8(v)[I96], np.zeros((64, 32)))
C = add(res, w['table7'][None, :]) if h == 0 else R
R = mm(q8(O)[:, I32], w[f'wo{h}'], C)
N = q8((R * (nv_rowsum_sq(R)[:, None] * w['gain14'][None, :]).astype(H)).astype(H))[:, I64]
acc = add(R, w['table16'][None, :])
for b in range(8):
hid = act(mm(N, w[f'w1_{b}'], np.repeat(w[f'b1_{b}'][None, :], 64, 0)))
acc = mm(q8(hid)[:, I32], w[f'w2_{b}'], acc)
tokens = q8(acc)[:, I64]
head = mm(tokens, w['wh'], np.repeat(w['bh'][None, :], 64, 0))
return tokens, head
def filter112(head, t88, t96, t104, x0, y0, ex, ey):
"""dec0_ref.outputs' surface 112 on the head's channels 0-8 (logits) and 9 (gate)"""
out = []
for t in range(64):
x, y = x0 + t % 8, y0 + t // 8
a = head[t].astype(H)
lg = a[:9]
m01 = np.maximum(np.maximum(np.maximum(lg[0:2], lg[2:4]), lg[4:6]), lg[6:8])
m = max(max(m01[0], m01[1]), lg[8])
e = [np.exp2(F(H(lg[i] - m)) * F(1.4426950216293335)) for i in range(9)]
pair = [F(e[2 * k] + e[2 * k + 1]) for k in range(4)]
sm = F(pair[0] + F(0))
for k in (1, 2, 3):
sm = F(sm + pair[k])
sm = F(sm + e[8])
r = H(F(1) / sm)
wg = [H(H(e[i]) * r) for i in range(9)]
def cl(v, n):
return 1 if v < 0 else (n - 2 if v >= n else v)
f = np.zeros(3, H)
k = 0
for yy in (cl(y - 1, ey), cl(y, ey), cl(y + 1, ey)):
for xx in (cl(x - 1, ex), cl(x, ex), cl(x + 1, ex)):
T = bilinear88(t88, xx, yy)
for c in range(3):
f[c] = fma16(T[c], wg[k], f[c])
k += 1
g = H(F(1) / (np.exp2(F(a[9]) * F(-1.4426950216293335)) + F(1)))
om = H(H(1) - g)
o = np.zeros(3, H)
if 0 <= x < ex and 0 <= y < ey:
for c in range(3):
o[c] = fma16(t96[y, x, c], om, H(f[c] * g))
if not np.isnan(t104[y, x, c]):
o[c] = H(o[c] + t104[y, x, c])
out.append(o)
return np.array(out)
def extract(dump, weights, out):
z = np.load(dump)
W = np.frombuffer(weights, np.uint8)
def B(p, wv=0):
return z[f'p{p}_w{wv}_B']
def C(p, wv=0):
return f16(z[f'p{p}_w{wv}_C']).astype(H)
def table(off, n=64):
return W[off:off + 2 * n].view(H)
w = {'we': np.hstack([B(1, 0), B(1, 1)]), 'be': np.concatenate([C(1, 0)[0], C(1, 1)[0]]),
'gain0': table(GAIN0), 'table7': table(TABLE7), 'gain14': table(GAIN14), 'table16': table(TABLE16)}
for h in range(2):
w[f'wk{h}a'], w[f'wv{h}a'], w[f'wk{h}b'], w[f'wv{h}b'] = B(2, h), B(3, h), B(4, h), B(5, h)
w[f'bias{h}'] = C(6, h)
w[f'wo{h}'] = B(8 + h)
for b in range(8):
w[f'w1_{b}'], w[f'b1_{b}'] = B(10 + 2 * b), C(10 + 2 * b)[0]
w[f'w2_{b}'] = B(11 + 2 * b)
w['wh'], w['bh'] = B(26), C(26)[0]
np.savez_compressed(out, **w)
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('frame', type=int)
ap.add_argument('weights')
ap.add_argument('--extract', metavar='DUMP')
ap.add_argument('--blocks', nargs='*', default=['5,5'])
o = ap.parse_args()
from rr_capture import Capture
c = Capture(os.path.expanduser(o.capture))
l = c.find('swin_dec1', o.frame)
if o.extract:
extract(o.extract, l.weights, o.weights)
return
ex, ey = (int(v) for v in np.frombuffer(l.args[:8], np.uint32))
sx, sy = (int(v) for v in np.frombuffer(l.args[80:84], np.uint16))
ptr = {p.offset: p.delta for p in l.pointers}
pre = np.frombuffer(l.pre(l.arena_alloc), np.uint8)
post = np.frombuffer(l.post_buffer(c, l.arena_alloc), np.uint8)
lo = pre[ptr[8]:ptr[8] + 3 * (ex // 2) * (ey // 2) * 32].reshape(3, ey // 2, ex // 2, 32)
skip = pre[ptr[24]:ptr[24] + 2 * ex * ey * 32].reshape(2, ey, ex, 32)
outp = post[ptr[48]:ptr[48] + 2 * ex * ey * 32].reshape(2, ey, ex, 32)
img = {i.offset: i for i in l.images}
def surf(off):
i = img[off]
return np.frombuffer(i.path.read_bytes(), np.float16).reshape(i.height, i.width, 4)
t88, t96, t104, r112 = surf(88), surf(96), surf(104), surf(112)
w = dict(np.load(o.weights))
for blk in o.blocks:
bx, by = (int(v) for v in blk.split(','))
x0, y0 = 8 * bx - sx, 8 * by - sy
tokens, head = window(w, lo, skip, x0, y0, ex, ey)
s112 = filter112(head, t88, t96, t104, x0, y0, ex, ey)
inside = [t for t in range(64) if 0 <= x0 + t % 8 < ex and 0 <= y0 + t // 8 < ey]
tk = sum(np.array_equal(tokens[t], np.concatenate([outp[0, y0 + t // 8, x0 + t % 8], outp[1, y0 + t // 8, x0 + t % 8]]))
for t in inside)
sf = sum(np.array_equal(s112[t], r112[y0 + t // 8, x0 + t % 8, :3]) for t in inside)
print(f'block {bx},{by}: tokens {tk}/{len(inside)}, surface 112 {sf}/{len(inside)}')
if __name__ == '__main__':
main()
engine/rr/dll_weights.py
+317
View File
@@ -0,0 +1,317 @@
#!/usr/bin/env python3
"""Read the Ray Reconstruction network's weights straight from the user's nvngx_dlssd.dll.
usage: dll_weights.py DLL --list
dll_weights.py DLL --kernels
dll_weights.py DLL OUT_DIR (--model N | --capture CAPTURE_DIR) [--check EXTRACTED_DIR] [--any-kernels]
The DLL holds every model's tensors as plain binary16 arrays in natural ([out][in]) order, one stage after the other
(encoder 0 ... decoder 0), with nothing naming them. They are found by structure alone: every attention head has a
relative position bias table of 100 x 100 values (the 10 x 10 window with its border) whose entries depend only on the
offset between the two positions, and a model is eleven consecutive stages whose head counts are 2 2 4 4 8 8 8 4 4 2 2.
The other tensors of a stage sit at fixed distances from its tables (LAYOUT below).
OUT_DIR/STAGE.npz gets the tensors in the form engine/rr/{swin,enc0,dec0}_ref.py --extract produced from a ZLUDA
phase dump, which is what the packers (swin_pack.py, enc0_pack.py, dec0_pack.py) take:
FP8 matrices e4m3 codes, [k][n] with k in NVIDIA's fragment order (natural input channel inter(K)[k])
gains times sqrt(32), as NGX uploads them
bias [64 queries][96 keys] in the kernels' key order, times 2 sqrt(2), as NGX uploads it
--model picks the N-th model of the DLL (--list); --capture picks the one a bridge capture of the denoiser ran, by
looking up each stage's vectors in the weight image NGX uploaded (and fails if the capture is of another DLL).
--check compares every tensor with an extracted one (EXTRACTED_DIR/STAGE/weights.npz).
The native shaders were written against the kernels of one DLL version. Before anything is written the DLL's network
kernels (the PTX modules, engine/rr/native_kernels.txt: their SHA-256) are compared with those; another version is
refused unless --any-kernels is given. --kernels prints the digests of a DLL.
Everything read and written is NVIDIA's: it stays on the user's machine and is never packaged.
"""
import argparse
import os
import sys
from pathlib import Path
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rrmath import e4m3, q8 # noqa: E402
from enc1_ref import inter # noqa: E402
from dec0_ref import keys # noqa: E402
H = np.float16
# stage: (kind, C, C_LO (decoders: the channels one level down) or C_M (encoders: the merged channels), heads)
STAGES = {
'enc0': ('enc0', 32, 64, 2), 'enc1': ('enc', 64, 96, 2), 'enc2': ('enc', 96, 128, 4), 'enc3': ('enc', 128, 160, 4),
'enc4': ('enc', 160, 160, 8), 'enc5': ('enc', 160, 0, 8), 'dec4': ('dec', 160, 160, 8), 'dec3': ('dec', 128, 160, 4),
'dec2': ('dec', 96, 128, 4), 'dec1': ('dec', 64, 96, 2), 'dec0': ('dec0', 32, 64, 2),
}
ORDER = list(STAGES)
HEADS = [STAGES[s][3] for s in ORDER]
TABLE = 100 * 100 # one head's relative position bias, in values
# the 96 keys of a window (its 64 tokens, then the border without the corners) in the 10 x 10 table
POS = [(y + 1) * 10 + x + 1 for (x, y) in keys(0, 0)]
BIAS_SCALE = np.float32(2.0 * np.sqrt(2.0))
def bias_tables(u):
"""offsets (in values) of every relative position bias table in the binary16 view u of the DLL"""
# entry (i, j) equals entry (i + 1, j + 1) and (i + 10, j + 10): one step right / down for both positions
e1 = u[:-1111] == u[101:-1010]
e10 = u[:-1111] == u[1010:-101]
e11 = u[:-1111] == u[1111:]
first = e1[:-18] & e1[1:-17] & e1[2:-16] & e10[:-18] & e10[1:-17] & e11[:-18] & e11[1:-17] & (u[:len(e1) - 18] != u[1:len(e1) - 17])
found = []
dy, dx = np.divmod(np.arange(100), 10)
rel = ((dy[None, :] - dy[:, None] + 9) * 19 + dx[None, :] - dx[:, None] + 9).reshape(-1)
for p in np.nonzero(first)[0]:
if p + TABLE > len(u) or (found and p < found[-1] + TABLE):
continue
t = u[p:p + TABLE]
r = np.zeros(361, np.uint16)
r[rel] = t
if np.array_equal(r[rel], t) and len(np.unique(t)) > 64:
found.append(int(p))
return found
def models(dll):
"""[[bias offset (bytes) of each of the eleven stages] per model], in file order"""
u = np.frombuffer(dll[:len(dll) // 2 * 2], np.uint16)
tables = bias_tables(u)
runs = [] # (first table, count) of consecutive tables
for p in tables:
if runs and p == runs[-1][0] + TABLE * runs[-1][1]:
runs[-1][1] += 1
else:
runs.append([p, 1])
out, i = [], 0
while i + len(HEADS) <= len(runs):
if [n for _, n in runs[i:i + len(HEADS)]] == HEADS:
out.append([2 * p for p, _ in runs[i:i + len(HEADS)]])
i += len(HEADS)
else:
i += 1
return out
def kernel_digests(dll):
"""{(kernel, SHA-256 of its PTX module)} of the DLL's network kernels"""
import hashlib
import re
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..', 'kernels', 'tools'))
from extract_dlss_ptx import fatbins, ptx_entries
out = set()
for _, fatbin in fatbins(dll):
for text in ptx_entries(fatbin):
for name in re.findall(rb'\.entry (cuda_dldn_engine_swin_\w+)\(', text):
out.add((name.decode(), hashlib.sha256(text).hexdigest()))
return out
def check_kernels(dll, path):
"""the kernels the native shaders were written against must all be in the DLL, and no other version of them"""
known = {tuple(l.split()) for l in open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'native_kernels.txt'))
if l.strip() and not l.startswith('#')}
have = kernel_digests(dll)
missing, other = sorted(known - have), sorted(have - known)
if missing or other:
raise SystemExit(f'{path}: its network kernels are not the ones the native shaders were written against '
f'({len(missing)} missing, {len(other)} unknown; e.g. {(missing + other)[0][0]}). The native '
'stages are unverified for this DLL version: use the translated model (engine/compile_rr.py), or '
'--any-kernels to build anyway.')
def noise_table(dll):
"""preset E's encoder 0 dithers its features with 8 x 16 x 16 16-bit values, the first 4096 bytes of the weight
image NGX uploads. In the DLL they are the array of 2048 values whose low five bits are all clear or all set."""
u = np.frombuffer(dll[:len(dll) // 2 * 2], np.uint16)
low = u & 31
ok = np.concatenate([[False], (low == 0) | (low == 31), [False]])
edge = np.nonzero(ok[1:] != ok[:-1])[0]
found = []
for a, b in zip(edge[::2], edge[1::2]):
if b - a >= 2048:
t = u[b - 2048:b]
if len(np.unique(t)) > 1024 and (t & 31 == 31).sum() > 256:
found.append(t.tobytes())
return found[0] if len(set(found)) == 1 else None # some DLLs hold the same table twice
class Reader:
def __init__(self, dll, at):
self.dll, self.at = dll, at
def back(self, *shape):
self.at -= 2 * int(np.prod(shape))
return np.frombuffer(self.dll, H, int(np.prod(shape)), self.at).reshape(shape)
def take(self, *shape):
a = np.frombuffer(self.dll, H, int(np.prod(shape)), self.at).reshape(shape)
self.at += 2 * int(np.prod(shape))
return a
def fp8(nat):
"""a natural [out][in] binary16 matrix -> NVIDIA's B operand: codes [k][n], k in fragment order"""
k = nat.shape[1]
return np.ascontiguousarray(q8(nat)[:, inter(k) if k % 32 == 0 else np.arange(k)].T)
def gain(g, c):
return (g.astype(np.float32) * np.float32(np.sqrt(np.float32(c)))).astype(H)
def read_stage(dll, stage, bias_at):
"""the stage's tensors (see the module text); returns (dict, first byte, one past the last byte)"""
kind, C, CX, NH = STAGES[stage]
w = {}
r = Reader(dll, bias_at)
table_a = r.back(C)
wo = r.back(C, 32 * NH)
wv = r.back(32 * NH, C)
wk = r.back(32 * NH, C)
gain_a = r.back(C)
if kind == 'enc0': # the embedding of the 32 input features
w['b2'] = r.back(32).copy()
w['w2'] = fp8(r.back(32, 32))
w['b1'] = r.back(32).copy()
# [32][inputs]: preset E has 32 input features, preset D 30 (NGX pads the image's matrix with zeros). The
# matrix is the widest one that starts with a weight and holds FP8 values only
end = r.at
for inputs in range(32, 15, -1):
m = np.frombuffer(dll, H, 32 * inputs, end - 64 * inputs).reshape(32, inputs)
if m[0, :4].any() and np.array_equal(e4m3(q8(m)).astype(H).view(np.uint16), m.view(np.uint16)):
break
else:
raise SystemExit('enc0: no input embedding before its bias')
r.at = end - 64 * inputs
w['w1'] = fp8(np.hstack([m, np.zeros((32, 32 - inputs), H)]))
first = r.at
r = Reader(dll, bias_at)
bias = r.take(NH, 100, 100)
gain_b = r.take(C)
w1 = r.take(4 * C, C)
b1 = r.take(4 * C)
w2 = r.take(C, 4 * C)
table_b = r.take(C)
for h in range(NH):
w[f'wk{h}'] = fp8(wk[32 * h:32 * h + 32])
w[f'wv{h}'] = fp8(wv[32 * h:32 * h + 32])
w[f'wo{h}'] = fp8(wo[:, 32 * h:32 * h + 32])
w[f'bias{h}'] = (bias[h][np.ix_(POS[:64], POS)].astype(np.float32) * BIAS_SCALE).astype(H)
for b in range(C // 8):
w[f'w1_{b}'] = fp8(w1[32 * b:32 * b + 32])
w[f'b1_{b}'] = b1[32 * b:32 * b + 32].copy()
w[f'w2_{b}'] = fp8(w2[:, 32 * b:32 * b + 32])
names = ('gain0', 'table0', 'gain1', 'table1') if kind in ('enc0', 'dec0') else ('gain0', 'table7', 'gain14', 'table16')
for name, a in zip(names, (gain(gain_a, C), table_a.copy(), gain(gain_b, C), table_b.copy())):
w[name] = a
if kind in ('enc', 'enc0') and CX: # the 2x2 merge
w['wm'] = fp8(r.take(CX, 4 * C))
w['bm'] = r.take(CX).copy()
if kind in ('dec', 'dec0'): # the expand of the level below
w['we'] = fp8(r.take(4 * C, CX))
w['be'] = r.take(4 * C).copy()
if kind != 'dec0': # the 16-channel head of the stage's filter
w['wh'] = fp8(r.take(16, C))
w['bh'] = r.take(16).copy()
return w, first, r.at
def read_model(dll, bias_offsets):
"""{stage: tensors}; checks that the stages lie back to back in the DLL, which the layout implies"""
out, end = {}, None
for stage, at in zip(ORDER, bias_offsets):
out[stage], first, last = read_stage(dll, stage, at)
if end is not None and first != end:
raise SystemExit(f'{stage}: starts at {first}, the stage before ends at {end}: not the layout this reader knows')
end = last
return out
def in_capture(capture, model):
"""whether the capture's launches uploaded this model: every stage's vectors must be in its weight image"""
from rr_capture import Capture
cap = Capture(os.path.expanduser(capture))
for stage, w in model.items():
image = next(l for l in cap.launches if f'swin_{stage}' in l.kernel).weights
for name, a in w.items():
if a.dtype == H and a.ndim == 1 and bytes(image).find(a.tobytes()) < 0:
return False
return True
def main():
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument('dll')
ap.add_argument('output', nargs='?')
ap.add_argument('--list', action='store_true')
ap.add_argument('--model', type=int)
ap.add_argument('--capture')
ap.add_argument('--check', metavar='EXTRACTED_DIR')
ap.add_argument('--kernels', action='store_true')
ap.add_argument('--any-kernels', action='store_true')
o = ap.parse_args()
dll = Path(o.dll).read_bytes()
if o.kernels:
for name, digest in sorted(kernel_digests(dll)):
print(name, digest)
return
if o.output and not o.any_kernels:
check_kernels(dll, o.dll)
found = models(dll)
if not found:
raise SystemExit(f'{o.dll}: no Ray Reconstruction model with the layout this reader knows')
if o.list or not o.output:
for i, m in enumerate(found):
print(f'model {i}: bias tables at ' + ' '.join(map(str, m)))
return
if o.model is not None:
model = read_model(dll, found[o.model])
index = o.model
else:
if not o.capture:
raise SystemExit('--model N or --capture CAPTURE_DIR')
index, model = next(((i, m) for i, m in ((i, read_model(dll, f)) for i, f in enumerate(found))
if in_capture(o.capture, m)), (None, None))
if model is None:
raise SystemExit(f'none of the {len(found)} models in {o.dll} is the one the capture ran: the capture was taken '
'with another nvngx_dlssd.dll')
out = Path(o.output)
out.mkdir(parents=True, exist_ok=True)
for stage, w in model.items():
np.savez_compressed(out / f'{stage}.npz', **w)
noise = noise_table(dll)
if noise is not None:
if o.capture: # preset E's image starts with it; preset D's has none
from rr_capture import Capture
image = bytes(next(l for l in Capture(os.path.expanduser(o.capture)).launches if 'swin_enc0' in l.kernel).weights)
if 'diamond' not in next(l for l in Capture(os.path.expanduser(o.capture)).launches if 'swin_enc0' in l.kernel).kernel \
and image[:4096] != noise:
raise SystemExit('the noise table found in the DLL is not the one the capture uploaded')
(out / 'noise.bin').write_bytes(noise)
print(f'{o.dll}: model {index} of {len(found)}, {len(model)} stages in {out}')
if o.check:
bad = 0
for stage, w in model.items():
ref = dict(np.load(Path(o.check) / stage / 'weights.npz'))
for name in sorted(set(ref) | set(w)):
a, b = w.get(name), ref.get(name)
if a is None or b is None:
print(f'{stage}.{name}: only in the {"extraction" if a is None else "DLL reading"}')
bad += 1
continue
if name in ('we', 'be'): # a multi-warp extraction repeats the columns
b = b[..., :a.shape[-1]]
same = a.shape == b.shape and np.array_equal(a.view(np.uint8) if a.dtype == np.uint8 else a.astype(H).view(np.uint16),
b.view(np.uint8) if b.dtype == np.uint8 else b.astype(H).view(np.uint16))
if not same:
n = int((a != b).sum()) if a.shape == b.shape else -1
print(f'{stage}.{name}: differs ({a.shape} vs {b.shape}, {n} values)')
bad += 1
print('check: every tensor equals the extraction' if not bad else f'check: {bad} tensors differ')
sys.exit(1 if bad else 0)
if __name__ == '__main__':
main()
engine/rr/enc0_m_decl.glsl
+53
View File
@@ -0,0 +1,53 @@
// Native part of engine/rr/enc0_m.comp (generated by engine/rr/gen_enc0_m.py): declarations, after the translated
// header. Binding 1 holds NGX's weight image followed by engine/rr/enc0_pack.py's blob (offsets in enc0_m_w.h already
// include its base); binding 2 is the arena.
#extension GL_EXT_shader_8bit_storage : require
#extension GL_KHR_shader_subgroup_vote : require
#extension GL_EXT_control_flow_attributes : require
#extension GL_GOOGLE_include_directive : require
#include "enc0_m_w.h"
#include "rr_exact.glsl"
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 CM gl_CooperativeMatrixLayoutColumnMajor
#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM)
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
#define Q8(dst, src) Q8PAIRS(dst, src)
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) }
const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0;
shared uint KVS[2 * 768]; // the current head's K, then V^T; before the attention the key pixels' features
#define XF KVS // the key pixels' 32 FP8 features (plane 0 | plane 1), slot-major
shared uint XT[96 * 8]; // the normalised key tokens
#define RR_KS KVS
#define RR_KS0 0u
#define RR_VS KVS
#define RR_VS0 768u
#define MG XT // the merge rows (16 x 128 FP8) reuse XT
shared float16_t O112[64 * 4]; // surface 112 per token, for surface 128's 2x2 sums
ivec2 keyXY(uint t, int x0, int y0)
{
if (t < 64u) return ivec2(x0 + int(t & 7u), y0 + int(t >> 3u));
const int i = int(t & 7u), side = int((t - 64u) >> 3u);
return ivec2(side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8), side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i));
}
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);
}
uvec2 mergeXY(uint r) { return r < 8u ? uvec2(r >> 1u, r & 1u) : uvec2((r - 8u) >> 1u, 2u + (r & 1u)); }
// the translated prologue's shared buffer sh: input_noise (10 x 10 pixels, 4 f16) at byte 0, input_tensor (two planes
// of 10 x 10 pixels x 16 FP8 features) at byte 800
float16_t noiseAt(uint gx, uint gy, uint c)
{
const uint b = 80u * gy + 8u * gx + 2u * c;
return hb((sh[b >> 2] >> (8u * (b & 2u))) & 0xffffu);
}
engine/rr/enc0_m_tail.glsl
+263
View File
@@ -0,0 +1,263 @@
// ======================================================================== native part (engine/rr/enc0_ref.py)
memoryBarrierShared(); barrier(); // wave 0's translated prologue filled sh (features, noise) and wrote 304 / 1016
const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID;
const uint hl = lane >> 4u, col = lane & 15u;
RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u);
const int ex = int(param.w[0]), ey = int(param.w[1]);
const uint outOff = param.w[12], mrgOff = param.w[14], shift = param.w[20];
const int x0 = 8 * int(gl_WorkGroupID.x) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u);
const uint lplane = uint((ex / 2) * (ey / 2)) * 32u;
// the key slots' features, from the 10 x 10 grid (origin x0 - 1, y0 - 1)
if (tid < 96u) {
const ivec2 k = keyXY(inter(tid), x0, y0);
const uint g = uint((k.y - y0 + 1) * 10 + (k.x - x0 + 1));
[[unroll]] for (uint wd = 0u; wd < 4u; ++wd) {
XF[tid * 8u + wd] = sh[(800u + g * 16u) / 4u + wd];
XF[tid * 8u + 4u + wd] = sh[(2400u + g * 16u) / 4u + wd];
}
}
barrier();
// the embedding x = W2 q8(relu(W1 f + b1)) + b2 for token tiles n = wv (this wave's queries: x is also their
// residual) and n = wv + 4 (the border keys, waves 0 and 1), then each key's norm in registers (dec0_ref.norm_plane
// on the accumulator layout, as dec0_m.comp's MLP norm) to XT
ACC 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)]);
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));
}
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)]);
A8 a0, a1;
LDA(a0, W_W2, 32u, m, 0u); LDA(a1, W_W2, 32u, m, 1u);
K32(a, a0, hb_[0], a1, hb_[1])
[[unroll]] for (int i = 0; i < 8; ++i) xv[m][i] = float16_t(a[i]);
if (n == wv) { // the residual: x + table 0
[[unroll]] for (int i = 0; i < 8; ++i) {
precise float16_t c = xv[m][i] + w16.w[W_TABLE0 / 2u + 16u * m + 8u * hl + uint(i)];
R[m][i] = float(c);
}
}
}
float16_t h1[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
float16_t ht[2];
[[unroll]] for (uint tq = 0u; tq < 2u; ++tq) {
const uint i0 = 4u * tq + c, i1 = i0 + 2u;
precise float16_t p0 = xv[0][i0] * xv[0][i0] + xv[0][i1] * xv[0][i1];
precise float16_t p1 = xv[1][i0] * xv[1][i0] + xv[1][i1] * xv[1][i1];
precise float16_t hv = p0 + p1;
ht[tq] = hv;
}
precise float16_t hh = ht[0] + ht[1];
h1[c] = hh;
}
const uint other = subgroupShuffleXor(bh(h1[0]) | (bh(h1[1]) << 16u), 16u);
precise float16_t g0 = h1[0] + hb(other & 0xffffu), g1 = h1[1] + hb(other >> 16u);
precise float16_t s = g0 + g1;
precise float16_t se = s + float16_t(0.0001220703125);
const float16_t r = rsq16(se);
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
B8 xq;
[[unroll]] for (int i = 0; i < 8; ++i) {
precise float16_t g = r * w16.w[W_GAIN0 / 2u + inter(16u * m + 8u * hl + uint(i))];
precise float16_t nn = xv[m][i] * g;
xq[i] = q8h(nn);
}
coopMatStore(xq, XT, n * 128u + 4u * m, 8u, CM);
}
}
barrier(); // XT complete; XF is read (the attention reuses its memory)
// two heads
[[unroll]] for (uint hd = 0u; hd < 2u; ++hd) {
if (hd == 1u) barrier();
for (uint j = 0u; j < 3u; ++j) {
const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u;
A8 wk[2], xa[2];
B8 wvb[2], xb[2];
[[unroll]] for (uint kt = 0u; kt < 2u; ++kt) {
coopMatLoad(xb[kt], XT, n * 128u + 4u * kt, 8u, CM);
coopMatLoad(xa[kt], XT, n * 128u + 4u * kt, 8u, RM);
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);
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);
K32(vv, xa[0], wvb[0], xa[1], wvb[1])
A8 va;
Q8(va, vv)
coopMatStore(va, KVS, 768u + c * 384u + 4u * n, 24u, RM);
}
barrier();
#include "rr_attn.glsl"
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
A8 a0, a1;
LDA(a0, W_WO0 + HEAD * hd, 32u, m, 0u); LDA(a1, W_WO0 + HEAD * hd, 32u, m, 1u);
K32(R[m], a0, Ob[0], a1, Ob[1])
}
}
// MLP (dec0_m.comp's)
B8 Nb[2];
{
float16_t rv[2][8];
[[unroll]] for (uint m = 0u; m < 2u; ++m) [[unroll]] for (int i = 0; i < 8; ++i) rv[m][i] = float16_t(R[m][i]);
float16_t h1[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
float16_t ht[2];
[[unroll]] for (uint tq = 0u; tq < 2u; ++tq) {
const uint i0 = 4u * tq + c, i1 = i0 + 2u;
precise float16_t p0 = rv[0][i0] * rv[0][i0] + rv[0][i1] * rv[0][i1];
precise float16_t p1 = rv[1][i0] * rv[1][i0] + rv[1][i1] * rv[1][i1];
precise float16_t hv = p0 + p1;
ht[tq] = hv;
}
precise float16_t hh = ht[0] + ht[1];
h1[c] = hh;
}
const uint other = subgroupShuffleXor(bh(h1[0]) | (bh(h1[1]) << 16u), 16u);
precise float16_t g0 = h1[0] + hb(other & 0xffffu), g1 = h1[1] + hb(other >> 16u);
precise float16_t s = g0 + g1;
precise float16_t se = s + float16_t(0.0001220703125);
const float16_t r = rsq16(se);
[[unroll]] for (uint m = 0u; m < 2u; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) {
const uint k = 16u * m + 8u * hl + uint(i);
precise float16_t g = r * w16.w[W_GAIN1 / 2u + k];
precise float16_t nn = rv[m][i] * g;
Nb[m][i] = q8h(nn);
precise float16_t a = rv[m][i] + w16.w[W_TABLE1 / 2u + k];
R[m][i] = float(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)]);
A8 a0, a1;
LDA(a0, w1, 32u, c, 0u); LDA(a1, w1, 32u, c, 1u);
K32(h, a0, Nb[0], a1, Nb[1])
[[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = ACTQ(float16_t(h[i]));
}
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
A8 a0, a1;
LDA(a0, w2, 32u, m, 0u); LDA(a1, w2, 32u, m, 1u);
K32(R[m], a0, Hb[0], a1, Hb[1])
}
}
// outputs: the tokens (one plane, dec0's skip) and the merge rows
const uint t = inter(16u * wv + col);
const int tx = x0 + int(t & 7u), ty = y0 + int(t >> 3u);
const bool inside = tx >= 0 && ty >= 0 && tx < ex && ty < ey;
B8 Tb[2];
{
const uint ty2 = t >> 3u, tx2 = t & 7u;
const uint mr = (ty2 >> 1u) < 2u ? 2u * (tx2 >> 1u) + (ty2 >> 1u) : 8u + 2u * (tx2 >> 1u) + (ty2 >> 1u) - 2u;
const uint blk = 2u * (ty2 & 1u) + (tx2 & 1u);
const uint o = (outOff + uint(ty * ex + tx) * 32u) >> 2u;
barrier(); // XT is read; the merge rows reuse it
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
Q8(Tb[m], R[m])
const uint q0 = pack4(Tb[m][0], Tb[m][1], Tb[m][2], Tb[m][3]), q1 = pack4(Tb[m][4], Tb[m][5], Tb[m][6], Tb[m][7]);
const uint at = 4u * m + 2u * hl;
if (inside) { arena.w[o + at] = q0; arena.w[o + at + 1u] = q1; }
const uint gm = mr * 32u + blk * 8u + 4u * m + 2u * hl;
MG[gm] = q0; MG[gm + 1u] = q1;
}
}
{ // the 32 -> 16 head and encoder 0's filter of the noise grid (two f16 fma chains, taps even / odd)
ACC hv;
[[unroll]] for (int i = 0; i < 8; ++i) hv[i] = float(w16.w[W_BH / 2u + 8u * hl + uint(i)]);
A8 a0, a1;
LDA(a0, W_WH, 32u, 0u, 0u); LDA(a1, W_WH, 32u, 0u, 1u);
K32(hv, a0, Tb[0], a1, Tb[1])
float16_t l[12];
[[unroll]] for (int i = 0; i < 8; ++i) l[i] = float16_t(hv[i]);
const uint u0 = subgroupShuffleXor(bh(l[0]) | (bh(l[1]) << 16u), 16u), u1 = subgroupShuffleXor(bh(l[2]) | (bh(l[3]) << 16u), 16u);
l[8] = hb(u0 & 0xffffu); l[9] = hb(u0 >> 16u); l[10] = hb(u1 & 0xffffu); l[11] = hb(u1 >> 16u);
const float16_t ma = max(max(max(l[0], l[2]), l[4]), l[6]), mb = max(max(max(l[1], l[3]), l[5]), l[7]);
const float16_t mx = max(max(ma, mb), l[8]);
float e[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) { precise float16_t d = l[i] - mx; e[i] = exp2(float(d) * 1.4426950216293335); }
precise float s = (e[0] + e[1]) + 0.0;
s = s + (e[2] + e[3]);
s = s + (e[4] + e[5]);
s = s + (e[6] + e[7]);
s = s + e[8];
const float16_t rr = rcp16(s);
float16_t wg[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) wg[i] = h16(r16(e[i]) * float(rr));
const uint gx = t & 7u, gy = t >> 3u; // the token's 3 x 3 taps in the grid: (gx .. gx + 2, gy .. gy + 2)
float16_t o112[3], o120[3];
[[unroll]] for (uint c = 0u; c < 3u; ++c) {
float16_t lo = 0.0hf, hi = 0.0hf;
[[unroll]] for (uint k = 0u; k < 8u; k += 2u) {
lo = fma16(noiseAt(gx + k % 3u, gy + k / 3u, c), wg[k], lo);
hi = fma16(noiseAt(gx + (k + 1u) % 3u, gy + (k + 1u) / 3u, c), wg[k + 1u], hi);
}
precise float16_t sum = lo + hi;
const float16_t f = fma16(noiseAt(gx + 2u, gy + 2u, c), wg[8], sum);
precise float16_t a = sig(l[9u + c]) * f;
precise float16_t b = f - a;
o120[c] = a; o112[c] = b;
}
if (hl == 0u) {
if (inside) {
imageStore(S112, ivec2(tx, ty), vec4(o112[0], o112[1], o112[2], 0.0));
imageStore(S120, ivec2(tx, ty), vec4(o120[0], o120[1], o120[2], 0.0));
}
[[unroll]] for (uint c = 0u; c < 3u; ++c) O112[t * 4u + c] = o112[c];
}
}
barrier();
{ // the 2x2 merge (128 -> 64): wave w computes channel tile w
const uvec2 mxy = mergeXY(col);
const int X = x0 / 2 + int(mxy.x), Y = y0 / 2 + int(mxy.y), hx = ex / 2, hy = ey / 2;
const uint m = wv;
ACC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(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);
coopMatLoad(g1, MG, 8u * ks + 4u, 32u, CM);
LDA(w0, W_WM, 128u, m, 2u * ks); LDA(w1, W_WM, 128u, 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 = (mrgOff + (m >> 1u) * lplane + uint(Y * hx + X) * 32u) / 4u + 4u * (m & 1u) + 2u * hl;
arena.w[at] = pack4(q[0], q[1], q[2], q[3]);
arena.w[at + 1u] = pack4(q[4], q[5], q[6], q[7]);
}
}
if (wv == 3u && lane < 16u) { // surface 128: 2x2 sums of 0.25 * surface 112, rightwards first, then downwards
const uint mx = lane & 3u, my = lane >> 2u;
const int x = x0 + 2 * int(mx), y = y0 + 2 * int(my);
if (x >= 0 && y >= 0 && x < ex && y < ey) {
const uint t00 = (2u * my) * 8u + 2u * mx;
vec4 out_ = vec4(0.0);
[[unroll]] for (uint c = 0u; c < 3u; ++c) {
precise float16_t a00 = O112[t00 * 4u + c] * float16_t(0.25), a01 = O112[(t00 + 1u) * 4u + c] * float16_t(0.25);
precise float16_t a10 = O112[(t00 + 8u) * 4u + c] * float16_t(0.25), a11 = O112[(t00 + 9u) * 4u + c] * float16_t(0.25);
precise float16_t b0 = a00 + a01, b1 = a10 + a11;
precise float16_t cc = b0 + b1;
out_[c] = float(cc);
}
imageStore(S128, ivec2(x >> 1, y >> 1), out_);
}
}
engine/rr/enc0_pack.py
+65
View File
@@ -0,0 +1,65 @@
#!/usr/bin/env python3
"""Pack encoder 0's extracted weights (engine/rr/enc0_ref.py --extract) for engine/rr/enc0_m.comp's native part.
usage: enc0_pack.py WEIGHTS.npz OUT.bin OUT.h [BASE]
dec0_pack.py's layout (C = 32) plus the embedding (W1, W2: 32 -> 32 -> 32, columns in slot order), the 2x2 merge
(128 -> 64, columns in arena order) and the 32 -> 16 head. BASE (default 0) is added to every offset in OUT.h: the
shader binds NGX's weight image first (its translated feature prologue reads it) and this blob after it.
"""
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from enc1_ref import inter # noqa: E402
from enc1_pack import a_tiles, b_tiles, cubic_table, act_table # noqa: E402
I32, I64, I96 = inter(32), inter(64), inter(96)
def main():
src, out_bin, out_h = sys.argv[1:4]
base = int(sys.argv[4]) if len(sys.argv) > 4 else 0
w = dict(np.load(src))
e = {'w1': ('A', w['w1'][:, I32]), 'b1': ('H', w['b1'][I32]), 'w2': ('A', w['w2'][:, I32]), 'b2': ('H', w['b2'][I32])}
for h in range(2):
e[f'wk{h}'] = ('A', w[f'wk{h}'][:, I32])
e[f'wv{h}'] = ('B', w[f'wv{h}'][:, I32])
e[f'wo{h}'] = ('A', w[f'wo{h}'][:, I32])
e[f'bias{h}'] = ('H', w[f'bias{h}'][np.ix_(inter(64), I96)].T)
for b in range(4):
e[f'w1_{b}'] = ('A', w[f'w1_{b}'][:, I32])
e[f'b1_{b}'] = ('H', w[f'b1_{b}'][I32])
e[f'w2_{b}'] = ('A', w[f'w2_{b}'][:, I32])
e['wm'], e['bm'] = ('A', w['wm'][:, I64]), ('H', w['bm'][I64])
e['wh'], e['bh'] = ('A', w['wh']), ('H', w['bh'])
e['gain0'] = ('H', w['gain0'])
for k in ('table0', 'gain1', 'table1'):
e[k] = ('H', w[k][I32])
e['cubic'] = ('U', cubic_table())
e['act'] = ('Y', act_table())
blob, defines = bytearray(), []
for name, (kind, a) in e.items():
while len(blob) % 256:
blob.append(0)
off = len(blob)
if kind == 'A':
blob += a_tiles(a).tobytes()
elif kind == 'B':
blob += b_tiles(a).tobytes()
elif kind == 'Y':
blob += np.ascontiguousarray(a, np.uint8).tobytes()
elif kind == 'U':
blob += np.ascontiguousarray(a, np.uint16).tobytes()
else:
blob += np.ascontiguousarray(a, np.float16).tobytes()
defines.append(f'#define W_{name.upper()} {base + off}u // {kind} {a.shape}')
open(out_bin, 'wb').write(bytes(blob))
open(out_h, 'w').write('// generated by engine/rr/enc0_pack.py\n' + '\n'.join(defines) + '\n')
print(f'{len(blob)} bytes, {len(e)} entries, base {base}')
if __name__ == '__main__':
main()
engine/rr/enc0_ref.py
+138
View File
@@ -0,0 +1,138 @@
#!/usr/bin/env python3
"""Exact numpy reference of the denoiser's encoder 0 (cuda_dldn_engine_swin_enc0_kernel) after its input features,
one 8x8 window at a time. The 32 FP8 features of the 96 key pixels (NVIDIA's per-pixel gather from the game's guides
and the temporal state, which engine/rr/enc0_m.comp keeps as ptx2rr.py's translation) are an input here.
embedding F (96 x 32) -> q8(relu(F W1 + b1)) -> W2 + b2: the key tokens x (32 channels, f16)
block dec0_ref's: RMS norm (gain W+6272), residual x + table W+37056, two heads, four MLP branches
outputs the tokens (argument 48, dec0's skip), the 2x2 merge 128 -> 64 (argument 56, enc1's input) and the
32 -> 16 head with enc1's filter to surfaces 112, 120 and 128
usage: enc0_ref.py CAPTURE FRAME WEIGHTS.npz --dump DUMP.npz [--block X,Y] (verifies the dump's block)
enc0_ref.py CAPTURE FRAME OUT.npz --extract DUMP.npz (weights from a phase dump)
"""
import argparse
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rrmath import f16, q8 # noqa: E402
from enc1_ref import inter, mm, softmax_cubic, act, add, filter_head, MERGE_ORDER, H # noqa: E402
from dec0_ref import norm_plane # noqa: E402
I32, I64, I96 = inter(32), inter(64), inter(96)
GAIN0, TABLE0, GAIN1, TABLE1 = 6272, 37056, None, None
def block(w, feat):
"""feat: the 96 key pixels' features (codes, NVIDIA's k order). Returns (tokens 64 x 32 codes in arena order,
merged 16 x 64 codes in arena order, head 64 x 16 f16, dbg)"""
dbg = {}
e1 = mm(feat, w['w1'], np.repeat(w['b1'][None, :], 96, 0))
x = mm(q8(np.maximum(e1, H(0)))[:, I32], w['w2'], np.repeat(w['b2'][None, :], 96, 0))
dbg['x'] = x
X = np.array([norm_plane(x[k], w['gain0']) for k in range(96)])
R = None
for h in range(2):
qk = mm(X, w[f'wk{h}'], np.zeros((96, 32)))
v = mm(X, w[f'wv{h}'], np.zeros((96, 32)))
Q = q8(qk)[:, I32]
P = softmax_cubic(mm(Q[:64], Q.T, w[f'bias{h}']))
O = mm(q8(P)[:, I96], q8(v)[I96], np.zeros((64, 32)))
C = add(x[:64], w['table0'][None, :]) if h == 0 else R
R = mm(q8(O)[:, I32], w[f'wo{h}'], C)
N = np.array([norm_plane(R[t], w['gain1']) for t in range(64)])
acc = add(R, w['table1'][None, :])
for b in range(4):
hid = act(mm(N, w[f'w1_{b}'], np.repeat(w[f'b1_{b}'][None, :], 64, 0)))
acc = mm(q8(hid)[:, I32], w[f'w2_{b}'], acc)
tokens = q8(acc)[:, I32]
Mg = np.zeros((16, 128), np.uint8)
for r, (mx, my) in enumerate(MERGE_ORDER):
for blk, (dx, dy) in enumerate([(0, 0), (1, 0), (0, 1), (1, 1)]):
Mg[r, 32 * blk:32 * blk + 32] = tokens[(2 * my + dy) * 8 + 2 * mx + dx]
merged = q8(mm(Mg, w['wm'], np.repeat(w['bm'][None, :], 16, 0)))[:, I64]
head = mm(tokens, w['wh'], np.repeat(w['bh'][None, :], 64, 0))
return tokens, merged, head, dbg
def extract(dump, weights, out):
z = np.load(dump)
W = np.frombuffer(weights, np.uint8)
def B(p):
return z[f'p{p}_w0_B']
def C(p):
return f16(z[f'p{p}_w0_C']).astype(H)
w = {'w1': B(1), 'b1': C(1)[0], 'w2': B(2), 'b2': C(2)[0], 'gain0': W[GAIN0:GAIN0 + 64].view(H),
'table0': W[TABLE0:TABLE0 + 64].view(H)}
# the input norm's gain and the residual table sit elsewhere in preset D's (smaller) weight image: when the preset
# E offsets do not reproduce the dump, find them by matching it, like the MLP's below
x = mm(q8(np.maximum(mm(z['p1_w0_A'], w['w1'], np.repeat(w['b1'][None, :], 96, 0)), H(0)))[:, I32], w['w2'],
np.repeat(w['b2'][None, :], 96, 0))
cands = [W[o:o + 64].view(H) for o in range(0, len(W) - 63, 2)]
if not all(np.array_equal(norm_plane(x[k], w['gain0']), z['p3_w0_A'][k]) for k in (0, 9)):
w['gain0'] = next(t for t in cands if all(np.array_equal(norm_plane(x[k], t), z['p3_w0_A'][k]) for k in (0, 9)))
if not np.array_equal(add(x[:2], w['table0'][None]).view(np.uint16), C(7)[:2].view(np.uint16)):
w['table0'] = next(t for t in cands if np.array_equal(add(x[:2], t[None]).view(np.uint16), C(7)[:2].view(np.uint16)))
for h, base in ((0, 3), (1, 8)):
w[f'wk{h}'], w[f'wv{h}'] = B(base), B(base + 1)
w[f'bias{h}'] = C(base + 2)
w[f'wo{h}'] = B(base + 4)
for b in range(4):
w[f'w1_{b}'], w[f'b1_{b}'] = B(13 + 2 * b), C(13 + 2 * b)[0]
w[f'w2_{b}'] = B(14 + 2 * b)
w['wm'], w['bm'] = B(21), C(21)[0]
w['wh'], w['bh'] = B(22), C(22)[0]
# the MLP norm's gain and table, found by matching the dump
R = f16(z['p12_w0_D']).astype(H)
tables = [W[o:o + 64].view(H) for o in range(0, len(W) - 63, 2)]
w['gain1'] = next(t for t in tables if all(np.array_equal(norm_plane(R[k], t), z['p13_w0_A'][k]) for k in (0, 9)))
w['table1'] = next(t for t in tables if np.array_equal(add(R[:2], t[None]).view(np.uint16), C(14)[:2].view(np.uint16)))
np.savez_compressed(out, **w)
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('frame', type=int)
ap.add_argument('weights')
ap.add_argument('--extract', metavar='DUMP')
ap.add_argument('--dump')
ap.add_argument('--block', default='5,5')
o = ap.parse_args()
from rr_capture import Capture
c = Capture(os.path.expanduser(o.capture))
l = c.find('swin_enc0', o.frame)
if o.extract:
extract(o.extract, l.weights, o.weights)
return
z = np.load(o.dump)
w = dict(np.load(o.weights))
ex, ey = (int(v) for v in np.frombuffer(l.args[:8], np.uint32))
sx, sy = (int(v) for v in np.frombuffer(l.args[80:84], np.uint16))
bx, by = (int(v) for v in o.block.split(','))
x0, y0 = 8 * bx - sx, 8 * by - sy
tokens, merged, head, dbg = block(w, z['p1_w0_A'])
for p, name in ((3, 'X'), ):
pass
ptr = {p.offset: p.delta for p in l.pointers}
post = np.frombuffer(l.post_buffer(c, l.arena_alloc), np.uint8)
otok = post[ptr[48]:ptr[48] + ex * ey * 32].reshape(ey, ex, 32)
omrg = post[ptr[56]:ptr[56] + 2 * (ex // 2) * (ey // 2) * 32].reshape(2, ey // 2, ex // 2, 32)
img = {i.offset: i for i in l.images}
tex = {k: np.frombuffer(i.path.read_bytes(), np.float16).reshape(i.height, i.width, 4) for k, i in img.items()
if k in (112, 120, 160)}
inside = [t for t in range(64) if 0 <= x0 + t % 8 < ex and 0 <= y0 + t // 8 < ey]
tk = sum(np.array_equal(tokens[t], otok[y0 + t // 8, x0 + t % 8]) for t in inside)
mg = sum(np.array_equal(merged[r], np.concatenate([omrg[p, y0 // 2 + my, x0 // 2 + mx] for p in range(2)]))
for r, (mx, my) in enumerate(MERGE_ORDER))
print(f'block {bx},{by}: tokens {tk}/{len(inside)}, merged {mg}/16')
return head, x0, y0, ex, ey
if __name__ == '__main__':
main()
engine/rr/enc1_m.comp
+357
View File
@@ -0,0 +1,357 @@
#version 460
// Native encoder 1 of the DLSS-D (preset E) denoiser, register-resident: cuda_dldn_engine_swin_enc1_kernel
// computed like engine/layer_m.comp. The arithmetic is engine/rr/enc1_ref.py's (bit-exact to NVIDIA); the stage
// contract is the translated stage's, except binding 1, which holds engine/rr/enc1_pack.py's packing.
//
// Every product runs transposed, D^T = W^T X^T: weights are the A operand, activations the B operand, and a result's
// rows are channels, its columns tokens. On RADV gfx1201 an accumulator's element i in lane l is (row 8 * (l >> 4)
// + i, column l & 15), a B operand's the same and an A operand's the transpose (layer_m.comp): so a result is the
// next product's B operand element for element, and a result tile is also the A operand of its own transpose (the
// keys of the scores, the V^T of P V). enc1_pack.py orders every weight's columns so that the rows come out in the
// k order NVIDIA's next product contracts them in; keys and queries are token slots (slot s = token inter(s)).
// Accumulators stay f32 and are rounded to binary16 values in place (r16) after each logical k32 step.
//
// One workgroup per 8x8-token window, four waves; wave w owns query slots 16w..16w+15 and computes the window's
// keys and values itself (no barrier between the prologue and the outputs). Shared memory holds only the
// normalised tokens, the raw tokens and the merge rows.
#extension GL_KHR_cooperative_matrix : require
#extension GL_KHR_memory_scope_semantics : require
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_8bit_storage : require
#extension GL_EXT_float_e4m3 : require
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_KHR_shader_subgroup_vote : require
#extension GL_EXT_control_flow_attributes : require
#extension GL_GOOGLE_include_directive : require
layout(local_size_x = 128) in;
#include "enc1_m_w.h"
#include "rr_exact.glsl"
layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param;
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16;
layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u;
layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
layout(set = 0, binding = 3) uniform sampler2D T88;
layout(set = 0, binding = 4, rgba16f) uniform writeonly image2D S112;
layout(set = 0, binding = 5, rgba16f) uniform writeonly image2D S120;
layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S128;
#ifndef SKIP
#define SKIP 0
#endif
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define HAC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CM gl_CooperativeMatrixLayoutColumnMajor
// W^T tile (mt, kt) of a weight with K rows (enc1_pack.a_tiles); W tile (kt, nt) of one with N columns (b_tiles)
#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM)
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
// binary16 rounding in place: a real conversion, kept by OR-ing in a zero the compiler cannot see (RADV folds a
// plain float(float16_t(x)) back to x); about 2 instructions per element against r16's 6
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
#define Q8(dst, src) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (dst)[i_] = q8s((src)[i_]);
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) }
const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0;
shared uint XT[96 * 16]; // the normalised key tokens: slot-major, 64 FP8 channels in arena (= k) order
// the merge rows (16 x 256 FP8, enc1_ref.MERGE_ORDER, 4 tokens of 64 channels) reuse XT, dead after the projections
#define MG XT
shared float16_t O112[64 * 4]; // surface 112 per token, for surface 128's 2x2 sums
shared uint KS[96 * 8]; // the current head's K, token slot x 32 FP8 channels (k order)
shared uint VS[32 * 24]; // the current head's V^T, 32 channels x 96 token slots
uvec2 mergeXY(uint r) { return r < 8u ? uvec2(r >> 1u, r & 1u) : uvec2((r - 8u) >> 1u, 2u + (r & 1u)); }
uint pack4(floate4m3_t a, floate4m3_t b, floate4m3_t c, floate4m3_t d)
{
return code(a) | (code(b) << 8u) | (code(c) << 16u) | (code(d) << 24u);
}
void main()
{
const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID;
const uint hl = lane >> 4u, col = lane & 15u;
RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u);
const int ex = int(param.p[0]), ey = int(param.p[1]);
const uint inOff = param.p[2], outOff = param.p[12], mrgOff = param.p[14], shift = param.p[20];
const int x0 = 8 * int(gl_WorkGroupID.x) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u);
const uint plane = uint(ex * ey) * 32u;
// ---------------------------------------------------------------- prologue: one key token per thread
if ((SKIP & 16) == 0 && tid < 96u) {
const uint t = inter(tid);
int x, y;
if (t < 64u) {
x = x0 + int(t & 7u); y = y0 + int(t >> 3u);
} else {
const int i = int(t & 7u), side = int((t - 64u) >> 3u);
x = side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8);
y = side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i);
}
x = mirror(x, ex); y = mirror(y, ey);
const uint base = (inOff + uint(y * ex + x) * 32u) >> 2u;
uint raw[16];
float16_t v[64];
[[unroll]] for (uint wd = 0u; wd < 16u; ++wd) {
raw[wd] = arena.w[base + (wd >> 3u) * (plane >> 2u) + (wd & 7u)];
[[unroll]] for (uint b = 0u; b < 4u; ++b)
v[4u * wd + b] = float16_t(uintBitsToFloate4m3EXT(uint8_t((raw[wd] >> (8u * b)) & 0xffu)));
}
// NVIDIA's order: per word t (and t + 4) the f16x2 pair sums of both planes, then lanes ^1, ^2, the halves
precise f16vec2 h[4];
[[unroll]] for (uint q = 0u; q < 4u; ++q) {
precise f16vec2 p[4];
[[unroll]] for (uint k = 0u; k < 4u; ++k) {
const uint o = 32u * (k & 1u) + 4u * (q + 4u * (k >> 1u));
precise f16vec2 a = f16vec2(v[o], v[o + 1u]) * f16vec2(v[o], v[o + 1u]);
precise f16vec2 c = f16vec2(v[o + 2u], v[o + 3u]) * f16vec2(v[o + 2u], v[o + 3u]);
p[k] = a + c;
}
precise f16vec2 s0 = p[0] + p[1], s1 = p[2] + p[3];
h[q] = s0 + s1;
}
precise f16vec2 h0 = h[0] + h[1], h2_ = h[2] + h[3];
precise f16vec2 hs = h0 + h2_;
precise float16_t s = hs.x + hs.y;
precise float16_t se = s + float16_t(0.0001220703125);
const float16_t r = rsq16(se);
[[unroll]] for (uint wd = 0u; wd < 16u; ++wd) {
floate4m3_t c[4];
[[unroll]] for (uint b = 0u; b < 4u; ++b) {
const uint j = 4u * wd + b;
precise float16_t g = r * w16.w[W_GAIN0 / 2u + inter(j)];
precise float16_t n = v[j] * g;
c[b] = q8h(n);
}
XT[tid * 16u + wd] = pack4(c[0], c[1], c[2], c[3]);
}
}
barrier();
// ---------------------------------------------------------------- the residual: raw tokens + table 7
// (each lane's token, mirrored like the prologue's; its channels 16m + 8hl .. + 7 are 8 bytes of plane m >> 1)
ACC R[4];
{
const uint tq = inter(16u * wv + col);
const int qx = mirror(x0 + int(tq & 7u), ex), qy = mirror(y0 + int(tq >> 3u), ey);
const uint qb = (inOff + uint(qy * ex + qx) * 32u) >> 2u;
[[unroll]] for (uint m = 0u; m < 4u; ++m) {
const uint at = qb + (m >> 1u) * (plane >> 2u) + 4u * (m & 1u) + 2u * hl;
const uvec2 u = uvec2(arena.w[at], arena.w[at + 1u]);
[[unroll]] for (int i = 0; i < 8; ++i) {
const float16_t x = float16_t(uintBitsToFloate4m3EXT(uint8_t((u[i >> 2] >> (8u * (uint(i) & 3u))) & 0xffu)));
precise float16_t c = x + w16.w[W_TABLE7 / 2u + 16u * m + 8u * hl + uint(i)];
R[m][i] = float(c);
}
}
}
// ---------------------------------------------------------------- attention, two heads
[[unroll]] for (uint hd = 0u; hd < ((SKIP & 64) != 0 ? 0u : 2u); ++hd) {
// this head's Q/K and V, shared: 12 jobs (token tile n, channel tile c), three per wave. K = (QK^T)^T rows
// to KS (token slot x 32 channels, also Q^T read column-major), V^T to VS (32 channels x 96 token slots)
if (hd == 1u) barrier(); // every wave is done with head 0's KS / VS
for (uint j = 0u; j < ((SKIP & 128) != 0 ? 0u : 3u); ++j) {
const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u;
A8 wk[4];
B8 wvb[4], xb[4];
A8 xa[4];
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) {
coopMatLoad(xb[kt], XT, n * 256u + 4u * kt, 16u, CM);
coopMatLoad(xa[kt], XT, n * 256u + 4u * kt, 16u, RM);
LDA(wk[kt], W_WK0 + HEAD * hd, 64u, c, kt);
LDB(wvb[kt], W_WV0 + HEAD * hd, 32u, kt, c);
}
ACC a = ACC(0.0);
K32(a, wk[0], xb[0], wk[1], xb[1])
K32(a, wk[2], xb[2], wk[3], xb[3])
A8 ka;
Q8(ka, a)
coopMatStore(ka, KS, n * 128u + 4u * c, 8u, RM);
ACC v = ACC(0.0);
K32(v, xa[0], wvb[0], xa[1], wvb[1])
K32(v, xa[2], wvb[2], xa[3], wvb[3])
A8 va;
Q8(va, v)
coopMatStore(va, VS, c * 384u + 4u * n, 24u, RM);
}
barrier();
#include "rr_attn.glsl"
[[unroll]] for (uint m = 0u; m < 4u; ++m) {
A8 a0, a1;
LDA(a0, W_WO0 + HEAD * hd, 32u, m, 0u); LDA(a1, W_WO0 + HEAD * hd, 32u, m, 1u);
K32(R[m], a0, Ob[0], a1, Ob[1])
}
}
// ---------------------------------------------------------------- MLP
B8 Nb[4];
{
float16_t rv[4][8];
[[unroll]] for (uint m = 0u; m < 4u; ++m) [[unroll]] for (int i = 0; i < 8; ++i) rv[m][i] = float16_t(R[m][i]);
// slot 16m + 8hl + i is NVIDIA's channel 8j + 2t + c with j = 4(m >> 1) + (m & 1) + 2((i >> 1) & 1),
// t = 2hl + (i >> 2), c = i & 1 (enc1_ref.nv_rowsum_sq's tree)
float16_t h1[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
float16_t ht[2];
[[unroll]] for (uint tq = 0u; tq < 2u; ++tq) {
const uint i0 = 4u * tq + c, i1 = i0 + 2u;
precise float16_t a = rv[0][i0] * rv[0][i0] + rv[0][i1] * rv[0][i1];
precise float16_t b = rv[1][i0] * rv[1][i0] + rv[1][i1] * rv[1][i1];
precise float16_t cc = rv[2][i0] * rv[2][i0] + rv[2][i1] * rv[2][i1];
precise float16_t d = rv[3][i0] * rv[3][i0] + rv[3][i1] * rv[3][i1];
precise float16_t hv = (a + cc) + (b + d);
ht[tq] = hv;
}
precise float16_t hh = ht[0] + ht[1];
h1[c] = hh;
}
const uint other = subgroupShuffleXor(bh(h1[0]) | (bh(h1[1]) << 16u), 16u);
precise float16_t g0 = h1[0] + hb(other & 0xffffu), g1 = h1[1] + hb(other >> 16u);
precise float16_t s = g0 + g1;
precise float16_t se = s + float16_t(0.0001220703125);
const float16_t r = rsq16(se);
[[unroll]] for (uint m = 0u; m < 4u; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) {
const uint k = 16u * m + 8u * hl + uint(i);
precise float16_t g = r * w16.w[W_GAIN14 / 2u + k];
precise float16_t n = rv[m][i] * g;
Nb[m][i] = q8h(n);
precise float16_t a = rv[m][i] + w16.w[W_TABLE16 / 2u + k];
R[m][i] = float(a);
}
}
for (uint br = 0u; br < ((SKIP & 4) != 0 ? 1u : 8u); ++br) {
const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br;
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
ACC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = float(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]);
A8 a0, a1, a2, a3;
LDA(a0, w1, 64u, c, 0u); LDA(a1, w1, 64u, c, 1u); LDA(a2, w1, 64u, c, 2u); LDA(a3, w1, 64u, c, 3u);
K32(h, a0, Nb[0], a1, Nb[1])
K32(h, a2, Nb[2], a3, Nb[3])
[[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = q8h(act(float16_t(h[i])));
}
[[unroll]] for (uint m = 0u; m < 4u; ++m) {
A8 a0, a1;
LDA(a0, w2, 32u, m, 0u); LDA(a1, w2, 32u, m, 1u);
K32(R[m], a0, Hb[0], a1, Hb[1])
}
}
// ---------------------------------------------------------------- outputs
const uint t = inter(16u * wv + col);
const int tx = x0 + int(t & 7u), ty = y0 + int(t >> 3u);
const bool inside = tx >= 0 && ty >= 0 && tx < ex && ty < ey;
B8 Tb[4];
{
const uint ty2 = t >> 3u, tx2 = t & 7u;
const uint r = (ty2 >> 1u) < 2u ? 2u * (tx2 >> 1u) + (ty2 >> 1u) : 8u + 2u * (tx2 >> 1u) + (ty2 >> 1u) - 2u;
const uint blk = 2u * (ty2 & 1u) + (tx2 & 1u);
const uint o = (outOff + uint(ty * ex + tx) * 32u) >> 2u;
[[unroll]] for (uint m = 0u; m < 4u; ++m) {
Q8(Tb[m], R[m])
const uint w0 = pack4(Tb[m][0], Tb[m][1], Tb[m][2], Tb[m][3]), w1 = pack4(Tb[m][4], Tb[m][5], Tb[m][6], Tb[m][7]);
const uint at = 4u * (m & 1u) + 2u * hl;
if (inside) { arena.w[o + (m >> 1u) * (plane >> 2u) + at] = w0; arena.w[o + (m >> 1u) * (plane >> 2u) + at + 1u] = w1; }
const uint g = r * 64u + 16u * blk + 4u * m + 2u * hl;
MG[g] = w0; MG[g + 1u] = w1;
}
}
if ((SKIP & 32) == 0) { // the 64 -> 16 head and its kernel-prediction filter (lanes 0-15; 16-31 hold head channels 8-15)
ACC hv;
[[unroll]] for (int i = 0; i < 8; ++i) hv[i] = float(w16.w[W_BH / 2u + 8u * hl + uint(i)]);
A8 a0, a1, a2, a3;
LDA(a0, W_WH, 64u, 0u, 0u); LDA(a1, W_WH, 64u, 0u, 1u); LDA(a2, W_WH, 64u, 0u, 2u); LDA(a3, W_WH, 64u, 0u, 3u);
K32(hv, a0, Tb[0], a1, Tb[1])
K32(hv, a2, Tb[2], a3, Tb[3])
float16_t l[12];
[[unroll]] for (int i = 0; i < 8; ++i) l[i] = float16_t(hv[i]);
const uint u0 = subgroupShuffleXor(bh(l[0]) | (bh(l[1]) << 16u), 16u), u1 = subgroupShuffleXor(bh(l[2]) | (bh(l[3]) << 16u), 16u);
l[8] = hb(u0 & 0xffffu); l[9] = hb(u0 >> 16u); l[10] = hb(u1 & 0xffffu); l[11] = hb(u1 >> 16u);
const float16_t ma = max(max(max(l[0], l[2]), l[4]), l[6]), mb = max(max(max(l[1], l[3]), l[5]), l[7]);
const float16_t mx = max(max(ma, mb), l[8]);
float e[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) { precise float16_t d = l[i] - mx; e[i] = exp2(float(d) * 1.4426950216293335); }
precise float s = (e[0] + e[1]) + 0.0;
s = s + (e[2] + e[3]);
s = s + (e[4] + e[5]);
s = s + (e[6] + e[7]);
s = s + e[8];
const float16_t rr = rcp16(s);
float wg[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) wg[i] = r16(r16(e[i]) * float(rr));
const int x = tx, y = ty;
const int xm = x - 1 < 0 ? 1 : x - 1, xc = x < 0 ? 1 : (x >= ex ? ex - 2 : x), xp = x + 1 >= ex ? ex - 2 : x + 1;
const int ym = y - 1 < 0 ? 1 : y - 1, yc = y < 0 ? 1 : (y >= ey ? ey - 2 : y), yp = y + 1 >= ey ? ey - 2 : y + 1;
const ivec2 pts[9] = ivec2[9](ivec2(xm, ym), ivec2(xc, ym), ivec2(xp, ym), ivec2(xm, yc), ivec2(xc, yc),
ivec2(xp, yc), ivec2(xm, yp), ivec2(xc, yp), ivec2(xp, yp));
float acc[3][2] = float[3][2](float[2](0.0, 0.0), float[2](0.0, 0.0), float[2](0.0, 0.0));
[[unroll]] for (uint k = 0u; k < 8u; ++k) {
const vec4 tv = texelFetch(T88, pts[k], 0);
[[unroll]] for (uint c = 0u; c < 3u; ++c) { precise float f = tv[c] * wg[k] + acc[c][k & 1u]; acc[c][k & 1u] = r16(f); }
}
const vec4 t8 = texelFetch(T88, pts[8], 0);
float o112[3], o120[3];
[[unroll]] for (uint c = 0u; c < 3u; ++c) {
precise float f0 = acc[c][0] + acc[c][1];
precise float f1 = t8[c] * wg[8] + r16(f0);
const float f = r16(f1);
precise float a = float(sig(l[9u + c])) * f;
const float ar = r16(a);
precise float b = f - ar;
o120[c] = ar; o112[c] = r16(b);
}
if (hl == 0u) {
if (inside) {
imageStore(S112, ivec2(x, y), vec4(o112[0], o112[1], o112[2], 0.0));
imageStore(S120, ivec2(x, y), vec4(o120[0], o120[1], o120[2], 0.0));
}
[[unroll]] for (uint c = 0u; c < 3u; ++c) O112[t * 4u + c] = float16_t(o112[c]);
}
}
barrier();
if ((SKIP & 8) == 0 && wv < 3u) { // the 2x2 merge: merged^T (96 x 16 rows), channel tiles 2w and 2w + 1 = arena plane w
const uvec2 mxy = mergeXY(col);
const int X = x0 / 2 + int(mxy.x), Y = y0 / 2 + int(mxy.y), hx = ex / 2, hy = ey / 2;
const uint o = (mrgOff + wv * uint(hx * hy) * 32u + uint(Y * hx + X) * 32u) >> 2u;
[[unroll]] for (uint m2 = 0u; m2 < 2u; ++m2) {
const uint m = 2u * wv + m2;
ACC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_BM / 2u + 16u * m + 8u * hl + uint(i)]);
for (uint ks = 0u; ks < 8u; ++ks) {
B8 g0, g1; A8 w0, w1;
coopMatLoad(g0, MG, 8u * ks, 64u, CM);
coopMatLoad(g1, MG, 8u * ks + 4u, 64u, CM);
LDA(w0, W_WM, 256u, m, 2u * ks); LDA(w1, W_WM, 256u, m, 2u * ks + 1u);
K32(a, w0, g0, w1, g1)
}
B8 q;
Q8(q, a)
if (X >= 0 && Y >= 0 && X < hx && Y < hy) {
arena.w[o + 4u * m2 + 2u * hl] = pack4(q[0], q[1], q[2], q[3]);
arena.w[o + 4u * m2 + 2u * hl + 1u] = pack4(q[4], q[5], q[6], q[7]);
}
}
} else if (lane < 16u) { // surface 128: 2x2 sums of 0.25 * surface 112, rightwards first, then downwards
const uint mx = lane & 3u, my = lane >> 2u;
const int x = x0 + 2 * int(mx), y = y0 + 2 * int(my);
if (x >= 0 && y >= 0 && x < ex && y < ey) {
const uint t00 = (2u * my) * 8u + 2u * mx;
vec4 out_ = vec4(0.0);
[[unroll]] for (uint c = 0u; c < 3u; ++c) {
precise float16_t a00 = O112[t00 * 4u + c] * float16_t(0.25), a01 = O112[(t00 + 1u) * 4u + c] * float16_t(0.25);
precise float16_t a10 = O112[(t00 + 8u) * 4u + c] * float16_t(0.25), a11 = O112[(t00 + 9u) * 4u + c] * float16_t(0.25);
precise float16_t b0 = a00 + a01, b1 = a10 + a11;
precise float16_t cc = b0 + b1;
out_[c] = float(cc);
}
imageStore(S128, ivec2(x >> 1, y >> 1), out_);
}
}
}
engine/rr/enc1_native.comp
+559
View File
@@ -0,0 +1,559 @@
#version 460
// Native encoder 1 of the DLSS-D (preset E) denoiser: cuda_dldn_engine_swin_enc1_kernel rewritten for cooperative
// matrices, bit-exact to engine/rr/enc1_ref.py (which is bit-exact to NVIDIA's launch). Same stage contract as the
// translated stage (engine/rr/README.md) except binding 1, which holds engine/rr/native_pack.py's packing of the
// extracted weights (offsets in the generated enc1_w.h).
//
// One workgroup is one 8x8-token window, four waves. Products run tokens x weights (A = activations from shared
// memory, B = weight tiles), one 16x16 output tile at a time, with NVIDIA's arithmetic: per logical k32 step the
// running value enters the first of two 16x16x16 FP8 products as C, the sum stays f32 and is rounded to binary16
// once, through shared memory (RADV folds a register-only f32 -> f16 -> f32). Every FP8 operand is stored in NVIDIA's
// k order: within each 32-block, slot j holds natural channel inter(j). The reductions (norms, softmax row sums) run
// in NVIDIA's exact order.
//
// 1 prologue the 96 key tokens (window, then a border row above, below, a column left and right; mirrored at the
// image edge), RMS norm (gain W_GAIN0) -> XQ; the residual + W_TABLE7 -> RES
// 2 QK and V per head, 96 x 64 -> 32 each: 24 tile jobs spread over the waves -> KV
// 3 attention per wave 16 query rows, both heads: scores + bias, cubic softmax, P V, out-projection onto RES
// 4 MLP per wave: RMS norm (W_GAIN14), RES += W_TABLE16, eight branches 64 -> 32 -> 64 onto RES
// 5 outputs the tokens (argument 48), the 2x2 merge to 96 channels (argument 56, waves 0-2), the 64 -> 16 head
// and its 3x3 kernel-prediction filter of texture 88 to surfaces 112, 120 and (half resolution) 128
#extension GL_KHR_cooperative_matrix : require
#extension GL_KHR_memory_scope_semantics : require
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_8bit_storage : require
#extension GL_EXT_float_e4m3 : require
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_GOOGLE_include_directive : require
layout(local_size_x = 128) in;
#include "enc1_w.h"
layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param;
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16;
layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
layout(set = 0, binding = 3) uniform sampler2D T88;
layout(set = 0, binding = 4, rgba16f) uniform writeonly image2D S112;
layout(set = 0, binding = 5, rgba16f) uniform writeonly image2D S120;
layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S128;
#ifdef DBG
// debugging: intermediates of window (DBX, DBY) as raw words: XQ, KV, RES after each head and after the MLP
layout(set = 0, binding = 7, std430) buffer Dbg { uint d[]; } dbg;
#define DBGWIN (gl_WorkGroupID.x == DBX && gl_WorkGroupID.y == DBY)
#define DUMP8(arr, n, at) if (DBGWIN) { barrier(); for (uint i_ = tid; i_ < (n); i_ += 128u) dbg.d[(at) + i_] = code(arr[i_]); barrier(); }
#define DUMP16(arr, n, at) if (DBGWIN) { barrier(); for (uint i_ = tid; i_ < (n); i_ += 128u) dbg.d[(at) + i_] = bh(arr[i_]); barrier(); }
#else
#define DUMP8(arr, n, at)
#define DUMP16(arr, n, at)
#endif
#ifndef SKIP
#define SKIP 0
#endif
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define HAC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CM gl_CooperativeMatrixLayoutColumnMajor
#define SYNC { memoryBarrierShared(); subgroupBarrier(); }
shared floate4m3_t XQ[96 * 64]; // the normalised key tokens; after the MLP the output tokens (64 rows)
shared float16_t RES[64 * 64]; // the residual stream, natural channel order
shared floate4m3_t KV[4 * 96 * 32]; // Q/K head 0, V head 0, Q/K head 1, V head 1; later the merge rows (16 x 256)
shared float16_t SW[4 * 1536]; // per wave: f16 results (scores 16 x 96), rounding scratch at +1024
shared floate4m3_t PW[4 * 1536]; // per wave: FP8 operands (P 16 x 96; O 16 x 32; N 16 x 64 + hidden 16 x 32)
// a weight tile: matrix W (K x N) packed by native_pack.py, k-tile kt, n-tile nt
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
// one binary16 rounding of an accumulator, through the wave's scratch
#define RND(acc) { \
coopMatStore(HAC(acc), SW, scr, 16u, RM); SYNC; \
HAC r_; coopMatLoad(r_, SW, scr, 16u, RM); acc = ACC(r_); SYNC; }
// NVIDIA's k order within a 32-block: slot j holds natural channel inter(j); uninter is the inverse (and also the
// arena's channel order: arena slot uninter(k) holds natural channel k)
uint inter(uint j) { return (j & ~31u) | (((j & 31u) >> 2u) << 1u) | (j & 1u) | (((j >> 1u) & 1u) << 4u); }
uint uninter(uint c) { return (c & ~31u) | (((c & 15u) >> 1u) << 2u) | (c & 1u) | (((c >> 4u) & 1u) << 1u); }
int mirror(int v, int n) { v = abs(v); return min(v, 2 * n - 2 - v); }
// cvt.rn.satfinite.e4m3 (NaN -> 0x7f)
floate4m3_t q8(float16_t h)
{
return isnan(h) ? uintBitsToFloate4m3EXT(uint8_t(0x7fu)) : floate4m3_t(clamp(h, float16_t(-448.0), float16_t(448.0)));
}
uint code(floate4m3_t v) { return uint(floate4m3BitsToUintEXT(v)); }
float16_t hb(uint b) { return uint16BitsToFloat16(uint16_t(b)); }
uint bh(float16_t h) { return uint(float16BitsToUint16(h)); }
// x rounded to binary16 precision (nearest even), staying f32. RADV folds float(float16_t(x)) back to x, so every
// f32 -> f16 rounding whose result is used as f32 again goes through this: adding and subtracting 1.5 * 2^(E + 13)
// (E = x's binary16 exponent) leaves x rounded to a multiple of its binary16 ulp; precise keeps the pair.
float r16(float x)
{
const uint e = max(floatBitsToUint(x) & 0x7f800000u, 0x38800000u);
const float c = uintBitsToFloat(e + (13u << 23u)) * 1.5;
precise float t = x + c;
precise float r = t - c;
return r;
}
float16_t h16(float x) { return float16_t(r16(x)); }
// m * x compared with 1, exactly (the fma recovers the product's rounding error): -1, 0 or 1
int cmp1(float m, float x)
{
precise float t = m * x;
precise float e = fma(m, x, -t);
return t > 1.0 ? 1 : (t < 1.0 ? -1 : (e > 0.0 ? 1 : (e < 0.0 ? -1 : 0)));
}
// binary16 rounding of 1 / x (x > 0) to nearest even, from the hardware approximation and two midpoint checks
float16_t rcp16(float x)
{
float16_t q = h16(1.0 / x);
for (int i = 0; i < 2; ++i) {
const uint b = bh(q);
const float lo = (float(q) + float(hb(b - 1u))) * 0.5, hi = (float(q) + float(hb(b + 1u))) * 0.5;
const int cl = cmp1(lo, x), ch = cmp1(hi, x);
if (cl > 0 || (cl == 0 && (b & 1u) != 0u)) q = hb(b - 1u);
else if (ch < 0 || (ch == 0 && (b & 1u) != 0u)) q = hb(b + 1u);
}
return q;
}
// binary16 rounding of 1 / sqrt(s) (s > 0)
float16_t rsq16(float16_t s16)
{
const float s = float(s16);
float16_t q = h16(inversesqrt(s));
for (int i = 0; i < 2; ++i) {
const uint b = bh(q);
const float lo = (float(q) + float(hb(b - 1u))) * 0.5, hi = (float(q) + float(hb(b + 1u))) * 0.5;
const int cl = cmp1(lo * lo, s), ch = cmp1(hi * hi, s);
if (cl > 0 || (cl == 0 && (b & 1u) != 0u)) q = hb(b - 1u);
else if (ch < 0 || (ch == 0 && (b & 1u) != 0u)) q = hb(b + 1u);
}
return q;
}
// fma of binary16 operands with one binary16 rounding, exactly: the product is exact in f32, the sum's f32 rounding
// error comes from TwoSum and decides the f32 results that land on a binary16 midpoint
float16_t fma16(float16_t a, float16_t b, float16_t c)
{
precise float p = float(a) * float(b);
precise float s = p + float(c);
precise float bb = s - p;
precise float err = (p - (s - bb)) + (float(c) - bb);
const uint u = floatBitsToUint(s);
if ((u & 0x1fffu) == 0x1000u && err != 0.0 && abs(s) >= 6.103515625e-05)
return h16(uintBitsToFloat(err * s > 0.0 ? u + 1u : u - 1u));
return h16(s);
}
float16_t act(float16_t x)
{
precise float16_t y = clamp(x, float16_t(-2.0), float16_t(2.0));
precise float16_t a = float16_t(0.0810546875) * abs(y);
precise float16_t b = float16_t(0.412109375) - a;
precise float16_t c = y * b;
precise float16_t g = float16_t(0.5) + c;
precise float16_t r = x * g;
return r;
}
float16_t sig(float16_t v) { return rcp16(exp2(float(v) * -1.4426950216293335) + 1.0); }
uvec2 mergeXY(uint r) { return r < 8u ? uvec2(r >> 1u, r & 1u) : uvec2((r - 8u) >> 1u, 2u + (r & 1u)); }
void main()
{
const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID;
const int ex = int(param.p[0]), ey = int(param.p[1]);
const uint inOff = param.p[2], outOff = param.p[12], mrgOff = param.p[14], shift = param.p[20];
const int x0 = 8 * int(gl_WorkGroupID.x) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u);
const uint plane = uint(ex * ey) * 32u;
const uint sw = wv * 1536u, scr = sw + 1024u, r0 = wv * 16u;
// ---------------------------------------------------------------- 1 prologue
#if !(SKIP & 1)
if (tid < 96u) {
int x, y;
if (tid < 64u) {
x = x0 + int(tid & 7u); y = y0 + int(tid >> 3u);
} else {
const int i = int(tid & 7u), side = int((tid - 64u) >> 3u);
x = side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8);
y = side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i);
}
x = mirror(x, ex); y = mirror(y, ey);
const uint base = (inOff + uint(y * ex + x) * 32u) >> 2u;
float16_t v[64];
for (uint pl = 0u; pl < 2u; ++pl)
for (uint wd = 0u; wd < 8u; ++wd) {
const uint u = arena.w[base + pl * (plane >> 2u) + wd];
for (uint b = 0u; b < 4u; ++b)
v[32u * pl + 4u * wd + b] = float16_t(uintBitsToFloate4m3EXT(uint8_t((u >> (8u * b)) & 0xffu)));
}
// NVIDIA's order: per word t (and t + 4) the f16x2 pair sums of both planes, then lanes ^1, ^2, the halves
precise f16vec2 h[4];
for (uint t = 0u; t < 4u; ++t) {
precise f16vec2 p[4];
for (uint k = 0u; k < 4u; ++k) {
const uint o = 32u * (k & 1u) + 4u * (t + 4u * (k >> 1u));
precise f16vec2 a = f16vec2(v[o], v[o + 1u]) * f16vec2(v[o], v[o + 1u]);
precise f16vec2 c = f16vec2(v[o + 2u], v[o + 3u]) * f16vec2(v[o + 2u], v[o + 3u]);
p[k] = a + c;
}
precise f16vec2 s0 = p[0] + p[1], s1 = p[2] + p[3];
h[t] = s0 + s1;
}
precise f16vec2 h0 = h[0] + h[1], h2_ = h[2] + h[3];
precise f16vec2 hs = h0 + h2_;
precise float16_t s = hs.x + hs.y;
precise float16_t se = s + float16_t(0.0001220703125);
const float16_t r = rsq16(se);
for (uint j = 0u; j < 64u; ++j) {
precise float16_t g = r * w16.w[W_GAIN0 / 2u + inter(j)];
precise float16_t n = v[j] * g;
XQ[tid * 64u + j] = q8(n);
}
if (tid < 64u)
for (uint k = 0u; k < 64u; ++k) {
precise float16_t c = v[uninter(k)] + w16.w[W_TABLE7 / 2u + k];
RES[tid * 64u + k] = c;
}
}
#endif
barrier();
DUMP8(XQ, 6144u, 0u)
DUMP16(RES, 4096u, 6144u)
// ---------------------------------------------------------------- 2 QK and V projections
#if !(SKIP & 2)
for (uint job = wv; job < 24u; job += 4u) {
const uint mt = job % 6u, mat = job / 6u, hd = mat >> 1u;
const bool isV = (mat & 1u) != 0u;
const uint wa = (isV ? W_WV0A : W_WK0A) + 2048u * hd, wb = wa + 1024u;
A8 a0, a1, a2, a3;
coopMatLoad(a0, XQ, mt * 1024u, 64u, RM);
coopMatLoad(a1, XQ, mt * 1024u + 16u, 64u, RM);
coopMatLoad(a2, XQ, mt * 1024u + 32u, 64u, RM);
coopMatLoad(a3, XQ, mt * 1024u + 48u, 64u, RM);
for (uint nt = 0u; nt < 2u; ++nt) {
B8 b0, b1;
LDB(b0, wa, 32u, 0u, nt); LDB(b1, wa, 32u, 1u, nt);
ACC acc = coopMatMulAdd(a1, b1, coopMatMulAdd(a0, b0, ACC(0.0)));
RND(acc);
LDB(b0, wb, 32u, 0u, nt); LDB(b1, wb, 32u, 1u, nt);
acc = coopMatMulAdd(a3, b1, coopMatMulAdd(a2, b0, acc));
coopMatStore(HAC(acc), SW, sw + nt * 16u, 32u, RM);
}
SYNC;
const uint row = lane >> 1u, cb = (lane & 1u) * 16u, kv = mat * 3072u;
if (!isV) {
for (uint j = cb; j < cb + 16u; ++j) KV[kv + (mt * 16u + row) * 32u + j] = q8(SW[sw + row * 32u + inter(j)]);
} else {
const uint dr = uninter(mt * 16u + row);
for (uint c = cb; c < cb + 16u; ++c) KV[kv + dr * 32u + c] = q8(SW[sw + row * 32u + c]);
}
SYNC;
}
#endif
barrier();
DUMP8(KV, 12288u, 10240u)
// ---------------------------------------------------------------- 3 attention, 16 query rows per wave
#if !(SKIP & 4)
for (uint hd = 0u; hd < 2u; ++hd) {
const uint qk = 6144u * hd, vv = qk + 3072u;
{
A8 q0, q1;
coopMatLoad(q0, KV, qk + r0 * 32u, 32u, RM);
coopMatLoad(q1, KV, qk + r0 * 32u + 16u, 32u, RM);
for (uint nt = 0u; nt < 6u; ++nt) {
HAC c;
coopMatLoad(c, w16.w, (W_BIAS0 + 12288u * hd) / 2u + r0 * 96u + nt * 16u, 96u, RM);
B8 k0, k1;
coopMatLoad(k0, KV, qk + nt * 512u, 32u, CM);
coopMatLoad(k1, KV, qk + nt * 512u + 16u, 32u, CM);
ACC acc = coopMatMulAdd(q1, k1, coopMatMulAdd(q0, k0, ACC(c)));
coopMatStore(HAC(acc), SW, sw + nt * 16u, 96u, RM);
}
}
SYNC;
{ // the cubic-exponent softmax: lane = (row, column parity)
const uint row = lane >> 1u, par = lane & 1u;
float16_t e[48];
for (uint j = 0u; j < 12u; ++j)
for (uint tt = 0u; tt < 4u; ++tt) {
const float16_t x = SW[sw + row * 96u + 8u * j + 2u * tt + par];
precise float16_t m = clamp(x * float16_t(0.017181396484375), float16_t(-0.55615234375), float16_t(0.55615234375));
const float16_t t = fma16(m, -m, float16_t(0.92724609375));
const float16_t u = fma16(m, t, float16_t(1.375));
e[4u * j + tt] = hb((bh(u) << 5u) & 0x7FE0u);
}
precise float16_t sm[4];
for (uint tt = 0u; tt < 4u; ++tt) {
precise float16_t ev = ((e[tt] + e[8u + tt]) + (e[16u + tt] + e[24u + tt])) + (e[32u + tt] + e[40u + tt]);
precise float16_t od = ((e[4u + tt] + e[12u + tt]) + (e[20u + tt] + e[28u + tt])) + (e[36u + tt] + e[44u + tt]);
sm[tt] = ev + od;
}
precise float16_t out_ = (sm[0] + sm[1]) + (sm[2] + sm[3]);
const float16_t other = hb(subgroupShuffleXor(bh(out_), 1u));
precise float16_t rs = par == 0u ? out_ + other : other + out_;
const float16_t r = rcp16(float(rs));
for (uint j = 0u; j < 12u; ++j)
for (uint tt = 0u; tt < 4u; ++tt) {
precise float16_t pv = e[4u * j + tt] * r;
PW[sw + row * 96u + uninter(8u * j + 2u * tt + par)] = q8(pv);
}
}
SYNC;
if (hd == 0u) { DUMP16(SW, 6144u, 34816u) DUMP8(PW, 6144u, 40960u) }
for (uint nt = 0u; nt < 2u; ++nt) { // P V
ACC acc = ACC(0.0);
for (uint ks = 0u; ks < 3u; ++ks) {
A8 p0, p1; B8 v0, v1;
coopMatLoad(p0, PW, sw + ks * 32u, 96u, RM);
coopMatLoad(p1, PW, sw + ks * 32u + 16u, 96u, RM);
coopMatLoad(v0, KV, vv + ks * 1024u + nt * 16u, 32u, RM);
coopMatLoad(v1, KV, vv + ks * 1024u + 512u + nt * 16u, 32u, RM);
acc = coopMatMulAdd(p1, v1, coopMatMulAdd(p0, v0, acc));
if (ks < 2u) RND(acc);
}
coopMatStore(HAC(acc), SW, sw + nt * 16u, 32u, RM);
}
SYNC;
{
const uint row = lane >> 1u, cb = (lane & 1u) * 16u;
for (uint j = cb; j < cb + 16u; ++j) PW[sw + row * 32u + j] = q8(SW[sw + row * 32u + inter(j)]);
}
SYNC;
{ // out-projection onto the residual
A8 o0, o1;
coopMatLoad(o0, PW, sw, 32u, RM);
coopMatLoad(o1, PW, sw + 16u, 32u, RM);
for (uint nt = 0u; nt < 4u; ++nt) {
HAC c;
coopMatLoad(c, RES, r0 * 64u + nt * 16u, 64u, RM);
B8 b0, b1;
LDB(b0, W_WO0 + 2048u * hd, 64u, 0u, nt); LDB(b1, W_WO0 + 2048u * hd, 64u, 1u, nt);
ACC acc = coopMatMulAdd(o1, b1, coopMatMulAdd(o0, b0, ACC(c)));
coopMatStore(HAC(acc), RES, r0 * 64u + nt * 16u, 64u, RM);
}
}
SYNC;
DUMP16(RES, 4096u, 22528u + 4096u * hd)
}
#endif
// ---------------------------------------------------------------- 4 MLP
#if !(SKIP & 8)
if (lane < 16u) {
const uint base = (r0 + lane) * 64u;
float16_t R[64];
for (uint k = 0u; k < 64u; ++k) R[k] = RES[base + k];
precise f16vec2 h[4];
for (uint t = 0u; t < 4u; ++t) {
precise f16vec2 q[8];
for (uint j = 0u; j < 8u; ++j) {
const f16vec2 x = f16vec2(R[8u * j + 2u * t], R[8u * j + 2u * t + 1u]);
q[j] = x * x;
}
precise f16vec2 a = q[0] + q[2], b = q[1] + q[3], c = q[4] + q[6], d = q[5] + q[7];
precise f16vec2 ac = a + c, bd = b + d;
h[t] = ac + bd;
}
precise f16vec2 h0 = h[0] + h[1], h2_ = h[2] + h[3];
precise f16vec2 hs = h0 + h2_;
precise float16_t s = hs.x + hs.y;
precise float16_t se = s + float16_t(0.0001220703125);
const float16_t r = rsq16(se);
for (uint j = 0u; j < 64u; ++j) {
const uint k = inter(j);
precise float16_t g = r * w16.w[W_GAIN14 / 2u + k];
precise float16_t n = R[k] * g;
PW[sw + lane * 96u + j] = q8(n);
}
for (uint k = 0u; k < 64u; ++k) {
precise float16_t c = R[k] + w16.w[W_TABLE16 / 2u + k];
RES[base + k] = c;
}
}
SYNC;
{
A8 n0, n1, n2, n3;
coopMatLoad(n0, PW, sw, 96u, RM);
coopMatLoad(n1, PW, sw + 16u, 96u, RM);
coopMatLoad(n2, PW, sw + 32u, 96u, RM);
coopMatLoad(n3, PW, sw + 48u, 96u, RM);
for (uint br = 0u; br < 8u; ++br) {
const uint w1 = W_W1_0 + 2048u * br, w2 = W_W2_0 + 2048u * br;
for (uint nt = 0u; nt < 2u; ++nt) {
HAC c;
coopMatLoad(c, w16.w, (W_B1_0 + 1024u * br) / 2u + nt * 16u, 32u, RM);
B8 b0, b1;
LDB(b0, w1, 32u, 0u, nt); LDB(b1, w1, 32u, 1u, nt);
ACC acc = coopMatMulAdd(n1, b1, coopMatMulAdd(n0, b0, ACC(c)));
RND(acc);
LDB(b0, w1, 32u, 2u, nt); LDB(b1, w1, 32u, 3u, nt);
acc = coopMatMulAdd(n3, b1, coopMatMulAdd(n2, b0, acc));
coopMatStore(HAC(acc), SW, sw + nt * 16u, 32u, RM);
}
SYNC;
{
const uint row = lane >> 1u, cb = (lane & 1u) * 16u;
for (uint j = cb; j < cb + 16u; ++j) PW[sw + row * 96u + 64u + j] = q8(act(SW[sw + row * 32u + inter(j)]));
}
SYNC;
A8 g0, g1;
coopMatLoad(g0, PW, sw + 64u, 96u, RM);
coopMatLoad(g1, PW, sw + 80u, 96u, RM);
for (uint nt = 0u; nt < 4u; ++nt) {
HAC c;
coopMatLoad(c, RES, r0 * 64u + nt * 16u, 64u, RM);
B8 b0, b1;
LDB(b0, w2, 64u, 0u, nt); LDB(b1, w2, 64u, 1u, nt);
ACC acc = coopMatMulAdd(g1, b1, coopMatMulAdd(g0, b0, ACC(c)));
coopMatStore(HAC(acc), RES, r0 * 64u + nt * 16u, 64u, RM);
}
SYNC;
}
}
#endif
DUMP16(RES, 4096u, 30720u)
// ---------------------------------------------------------------- 5 outputs
{ // the tokens, in arena order (= NVIDIA's k order)
const uint row = lane >> 1u, pl = lane & 1u, t = r0 + row;
const int x = x0 + int(t & 7u), y = y0 + int(t >> 3u);
uint words[8];
for (uint wd = 0u; wd < 8u; ++wd) {
uint u = 0u;
for (uint b = 0u; b < 4u; ++b) {
const uint j = 32u * pl + 4u * wd + b;
const floate4m3_t c = q8(RES[t * 64u + inter(j)]);
XQ[t * 64u + j] = c;
u |= code(c) << (8u * b);
}
words[wd] = u;
}
if (x >= 0 && y >= 0 && x < ex && y < ey) {
const uint o = (outOff + pl * plane + uint(y * ex + x) * 32u) >> 2u;
for (uint wd = 0u; wd < 8u; ++wd) arena.w[o + wd] = words[wd];
}
}
barrier();
#if !(SKIP & 16)
for (uint i = 0u; i < 32u; ++i) { // the merge rows: (mx, my) in NVIDIA's order, 4 tokens of 64 channels
const uint idx = tid * 32u + i, r = idx >> 8u, k = idx & 255u, blk = k >> 6u;
const uvec2 m = mergeXY(r);
KV[idx] = XQ[((2u * m.y + (blk >> 1u)) * 8u + 2u * m.x + (blk & 1u)) * 64u + (k & 63u)];
}
barrier();
if (wv < 3u) { // the 2x2 merge, one 32-channel plane per wave
for (uint n2 = 0u; n2 < 2u; ++n2) {
const uint nt = 2u * wv + n2;
HAC c;
coopMatLoad(c, w16.w, W_BM / 2u + nt * 16u, 96u, RM);
ACC acc = ACC(c);
for (uint ks = 0u; ks < 8u; ++ks) {
A8 m0, m1; B8 b0, b1;
coopMatLoad(m0, KV, ks * 32u, 256u, RM);
coopMatLoad(m1, KV, ks * 32u + 16u, 256u, RM);
LDB(b0, W_WM, 96u, 2u * ks, nt); LDB(b1, W_WM, 96u, 2u * ks + 1u, nt);
acc = coopMatMulAdd(m1, b1, coopMatMulAdd(m0, b0, acc));
if (ks < 7u) RND(acc);
}
coopMatStore(HAC(acc), SW, sw + n2 * 16u, 32u, RM);
}
SYNC;
const uint row = lane >> 1u, hf = lane & 1u;
const uvec2 m = mergeXY(row);
const int X = x0 / 2 + int(m.x), Y = y0 / 2 + int(m.y), hx = ex / 2, hy = ey / 2;
uint words[4];
for (uint wd = 0u; wd < 4u; ++wd) {
uint u = 0u;
for (uint b = 0u; b < 4u; ++b) u |= code(q8(SW[sw + row * 32u + inter(16u * hf + 4u * wd + b)])) << (8u * b);
words[wd] = u;
}
if (X >= 0 && Y >= 0 && X < hx && Y < hy) {
const uint o = (mrgOff + wv * uint(hx * hy) * 32u + uint(Y * hx + X) * 32u) >> 2u;
for (uint wd = 0u; wd < 4u; ++wd) arena.w[o + 4u * hf + wd] = words[wd];
}
SYNC;
}
#endif
{ // the 64 -> 16 head
HAC c;
coopMatLoad(c, w16.w, W_BH / 2u, 16u, RM);
A8 t0, t1, t2, t3; B8 b0, b1;
coopMatLoad(t0, XQ, r0 * 64u, 64u, RM);
coopMatLoad(t1, XQ, r0 * 64u + 16u, 64u, RM);
coopMatLoad(t2, XQ, r0 * 64u + 32u, 64u, RM);
coopMatLoad(t3, XQ, r0 * 64u + 48u, 64u, RM);
LDB(b0, W_WH, 16u, 0u, 0u); LDB(b1, W_WH, 16u, 1u, 0u);
ACC acc = coopMatMulAdd(t1, b1, coopMatMulAdd(t0, b0, ACC(c)));
RND(acc);
LDB(b0, W_WH, 16u, 2u, 0u); LDB(b1, W_WH, 16u, 3u, 0u);
acc = coopMatMulAdd(t3, b1, coopMatMulAdd(t2, b0, acc));
coopMatStore(HAC(acc), SW, sw, 16u, RM);
}
SYNC;
{ // the kernel-prediction filter; lanes 16-31 repeat 0-15 so that the shuffles below stay uniform
const uint tl = lane & 15u, tk = r0 + tl;
const int x = x0 + int(tk & 7u), y = y0 + int(tk >> 3u);
float16_t l[12];
for (uint i = 0u; i < 12u; ++i) l[i] = SW[sw + tl * 16u + i];
const float16_t ma = max(max(max(l[0], l[2]), l[4]), l[6]), mb = max(max(max(l[1], l[3]), l[5]), l[7]);
const float16_t mx = max(max(ma, mb), l[8]);
float e[9];
for (uint i = 0u; i < 9u; ++i) {
precise float16_t d = l[i] - mx;
e[i] = exp2(float(d) * 1.4426950216293335);
}
precise float s = (e[0] + e[1]) + 0.0;
s = s + (e[2] + e[3]);
s = s + (e[4] + e[5]);
s = s + (e[6] + e[7]);
s = s + e[8];
const float16_t r = rcp16(s);
float wg[9];
for (uint i = 0u; i < 9u; ++i) wg[i] = r16(r16(e[i]) * float(r));
const int xm = x - 1 < 0 ? 1 : x - 1, xc = x < 0 ? 1 : (x >= ex ? ex - 2 : x), xp = x + 1 >= ex ? ex - 2 : x + 1;
const int ym = y - 1 < 0 ? 1 : y - 1, yc = y < 0 ? 1 : (y >= ey ? ey - 2 : y), yp = y + 1 >= ey ? ey - 2 : y + 1;
const ivec2 pts[9] = ivec2[9](ivec2(xm, ym), ivec2(xc, ym), ivec2(xp, ym), ivec2(xm, yc), ivec2(xc, yc),
ivec2(xp, yc), ivec2(xm, yp), ivec2(xc, yp), ivec2(xp, yp));
// f32 values that are binary16 numbers throughout (r16 at every NVIDIA f16 rounding)
float o112[3], o120[3];
float acc[3][2] = float[3][2](float[2](0.0, 0.0), float[2](0.0, 0.0), float[2](0.0, 0.0));
for (uint k = 0u; k < 8u; ++k) {
const vec4 tv = texelFetch(T88, pts[k], 0);
for (uint c = 0u; c < 3u; ++c) {
precise float f = tv[c] * wg[k] + acc[c][k & 1u];
acc[c][k & 1u] = r16(f);
}
}
const vec4 t8 = texelFetch(T88, pts[8], 0);
for (uint c = 0u; c < 3u; ++c) {
precise float f0 = acc[c][0] + acc[c][1];
precise float f1 = t8[c] * wg[8] + r16(f0);
const float f = r16(f1);
precise float a = float(sig(l[9u + c])) * f;
const float ar = r16(a);
precise float b = f - ar;
o120[c] = ar; o112[c] = r16(b);
}
// surface 128: the 2x2 sums of 0.25 * surface 112, rightwards first, then downwards
float q[3], cs[3];
const uint nb = min(lane + 1u, 31u), dn = min(lane + 8u, 31u);
for (uint c = 0u; c < 3u; ++c) {
q[c] = o112[c] * 0.25;
precise float b = q[c] + subgroupShuffle(q[c], nb);
const float br = r16(b);
precise float d = br + subgroupShuffle(br, dn);
cs[c] = r16(d);
}
if (lane < 16u && x >= 0 && y >= 0 && x < ex && y < ey) {
imageStore(S112, ivec2(x, y), vec4(o112[0], o112[1], o112[2], 0.0));
imageStore(S120, ivec2(x, y), vec4(o120[0], o120[1], o120[2], 0.0));
if (((x | y) & 1) == 0) imageStore(S128, ivec2(x >> 1, y >> 1), vec4(cs[0], cs[1], cs[2], 0.0));
}
}
}
engine/rr/enc1_pack.py
+106
View File
@@ -0,0 +1,106 @@
#!/usr/bin/env python3
"""Pack encoder 1's extracted weights (engine/rr/enc1_ref.py --extract) for engine/rr/enc1_m.comp.
usage: enc1_pack.py WEIGHTS.npz OUT.bin OUT.h
The shader runs every product transposed (D^T = W^T X^T: weights are the A operand, activations the B operand), so a
result's rows are output channels and its columns tokens, and the result is the next product's B operand in place.
The rows come out in the order of W's columns, so every weight's columns are permuted here into the k order NVIDIA's
next product contracts them in (engine/rr/enc1_ref.py: q8(x)[:, inter(n)]), and the f16 vectors that act on those
rows likewise. Keys and queries are token slots: slot s is token inter(96)[s] (the prologue's key list), so the
attention bias is stored transposed in slot order.
A tiles W (K x N) -> W^T 16x16 tiles [mt][kt], row-major (m, k): 256 * (mt * K/16 + kt) bytes
B tiles the V projection runs X Wv (A = tokens): tiles [kt][nt], row-major (k, n)
f16 vectors as is; the attention bias as 96 x 64 [key slot][query slot]
"""
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from enc1_ref import inter # noqa: E402
I32, I64, I96 = inter(32), inter(64), inter(96)
def a_tiles(w):
k, n = w.shape
return np.ascontiguousarray(w.T).reshape(n // 16, 16, k // 16, 16).transpose(0, 2, 1, 3).reshape(-1)
def b_tiles(w):
k, n = w.shape
return w.reshape(k // 16, 16, n // 16, 16).transpose(0, 2, 1, 3).reshape(-1)
CUBIC_MAX = 0x3873 # binary16 bits of 0.55615234375, the clamp of the scaled score
def act_table():
"""the MLP's hidden activation and its FP8 conversion, q8(act(h)), for every binary16 h, indexed by its bits: one
byte load in place of the six binary16 operations and the conversion (the shaders' ACTQ)"""
from enc1_ref import act
from rrmath import q8
h = np.arange(65536, dtype=np.uint32).astype(np.uint16).view(np.float16)
with np.errstate(all='ignore'):
return np.asarray(q8(act(h).astype(np.float64)), np.uint8)
def cubic_table():
"""the cubic softmax numerator's bits (enc1_ref.softmax_cubic) for every binary16 score, indexed by its bits (the
scaling and the clamp included; NaN scores take NVIDIA's min / max, which return the other operand)"""
x = np.arange(65536, dtype=np.uint32).astype(np.uint16).view(np.float16)
with np.errstate(invalid='ignore', over='ignore'):
m = (x * np.float16(0.017181396484375)).astype(np.float16)
lim = np.float16(0.55615234375)
m = np.where(np.isnan(m), lim, np.clip(m, -lim, lim)).astype(np.float16)
def fma16(a, b, c):
return (a.astype(np.float64) * b.astype(np.float64) + c.astype(np.float64)).astype(np.float16)
t = fma16(m, -m, np.full_like(m, np.float16(0.92724609375)))
u = fma16(m, t, np.full_like(m, np.float16(1.375)))
return ((u.view(np.uint16).astype(np.uint32) << 5) & 0x7FE0).astype(np.uint16)
def main():
src, out_bin, out_h = sys.argv[1:4]
w = dict(np.load(src))
e = {}
for h in range(2):
e[f'wk{h}'] = ('A', np.vstack([w[f'wk{h}a'], w[f'wk{h}b']])[:, I32])
e[f'wv{h}'] = ('B', np.vstack([w[f'wv{h}a'], w[f'wv{h}b']])[:, I32])
e[f'wo{h}'] = ('A', w[f'wo{h}'][:, I64])
e[f'bias{h}'] = ('H', w[f'bias{h}'][np.ix_(I64, I96)].T)
for b in range(8):
e[f'w1_{b}'] = ('A', w[f'w1_{b}'][:, I32])
e[f'b1_{b}'] = ('H', w[f'b1_{b}'][I32])
e[f'w2_{b}'] = ('A', w[f'w2_{b}'][:, I64])
e['wm'], e['bm'] = ('A', w['wm'][:, I96]), ('H', w['bm'][I96])
e['wh'], e['bh'] = ('A', w['wh']), ('H', w['bh'])
e['gain0'] = ('H', w['gain0'])
for k in ('table7', 'gain14', 'table16'):
e[k] = ('H', w[k][I64])
e['cubic'] = ('U', cubic_table())
blob, defines = bytearray(), []
for name, (kind, a) in e.items():
while len(blob) % 256:
blob.append(0)
off = len(blob)
if kind == 'A':
blob += a_tiles(a).tobytes()
elif kind == 'B':
blob += b_tiles(a).tobytes()
elif kind == 'U':
blob += np.ascontiguousarray(a, np.uint16).tobytes()
else:
blob += np.ascontiguousarray(a, np.float16).tobytes()
defines.append(f'#define W_{name.upper()} {off}u // {kind} {a.shape}')
open(out_bin, 'wb').write(bytes(blob))
open(out_h, 'w').write('// generated by engine/rr/enc1_pack.py\n' + '\n'.join(defines) + '\n')
print(f'{len(blob)} bytes, {len(e)} entries')
if __name__ == '__main__':
main()
engine/rr/enc1_ref.py
+330
View File
@@ -0,0 +1,330 @@
#!/usr/bin/env python3
"""Exact numpy reference of the denoiser's encoder 1 (cuda_dldn_engine_swin_enc1_kernel), one 8x8 window at a time,
from its weights and real inputs only: the specification the native shader implements.
Derived phase by phase against ZLUDA replays (engine/rr/phase_dump.py), every step bit-exact:
prologue the 96 key tokens (the window plus a one-token border, corners excluded; mirrored at the image edge)
from the arena's two 32-channel planes, f16 RMS norm (NVIDIA's summation order) times the gain at W+0
attention two heads: QK and V projections (K = 64 in two k32 steps), scores = q8(QK[:64]) q8(QK)^T + bias,
the cubic-exponent softmax, P V, out-projection onto the residual (input tile + table at W+36992)
MLP RMS norm times the gain at W+37120, eight branches of 64 -> 32 (bias, ramp activation) -> 64,
accumulated onto the residual (+ table at W+41408)
outputs q8 tokens to param+48, the 2x2 merge (16 x 256 -> 96 + bias) to param+56, the 64 -> 16 head and its
3x3 kernel-prediction filter of texture 88 to surfaces 112, 120 and (half resolution) 128
usage: enc1_ref.py CAPTURE FRAME WEIGHTS.npz [--blocks X,Y ...] (verifies against the capture)
WEIGHTS.npz: engine/rr/enc1_ref.py --extract DUMP.npz CAPTURE FRAME OUT.npz builds it from a phase dump
"""
import argparse
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rrmath import e4m3, f16, q8, mma_k32 # noqa: E402
H = np.float16
F = np.float32
def inter(n):
perm = []
for b in range(n // 32):
for m in range(8):
perm += [32 * b + 2 * m, 32 * b + 2 * m + 1, 32 * b + 16 + 2 * m, 32 * b + 17 + 2 * m]
return perm
INTER64 = inter(64)
# arena channel order (NVIDIA's A-fragment order) -> natural: natural channel k is arena channel ARENA[k]
ARENA = [0, 1, 4, 5, 8, 9, 12, 13, 16, 17, 20, 21, 24, 25, 28, 29, 2, 3, 6, 7, 10, 11, 14, 15, 18, 19, 22, 23, 26, 27,
30, 31]
ARENA = ARENA + [32 + c for c in ARENA]
def add(a, b):
return (a + b).astype(H)
def mm(a_codes, b_codes, c):
"""a (M x K e4m3 codes, NVIDIA k order) . b (K x N codes) + c (M x N f16 values) with ZLUDA's k32 rounding"""
return f16(mma_k32(a_codes, b_codes, np.asarray(c, H).view(np.uint16))).astype(H)
def nv_rowsum_sq(x, eps=2.0 ** -13):
"""f16(rsqrt(sum of squares + eps)) per row of a 64-column C/D tile, NVIDIA's order"""
sq = (x * x).astype(H)
def tile(j, t):
return sq[:, 8 * j + 2 * t: 8 * j + 2 * t + 2]
h = []
for t in range(4):
a, b = add(tile(0, t), tile(2, t)), add(tile(1, t), tile(3, t))
c, d = add(tile(4, t), tile(6, t)), add(tile(5, t), tile(7, t))
h.append(add(add(a, c), add(b, d)))
h1 = [add(h[t], h[t ^ 1]) for t in range(4)]
h2 = add(h1[0], h1[2])
s = add(add(h2[:, 0], h2[:, 1]), np.full(x.shape[0], H(eps)))
return (F(1) / np.sqrt(s.astype(F))).astype(H)
def input_norm(tok_arena, gain):
"""the prologue's per-token norm: tok_arena (2 planes x 32, arena order) -> f16 normalised token (natural)"""
sq = [(v * v).astype(H) for v in tok_arena]
h = []
for t in range(4):
def pair(p, w):
s = sq[p][4 * w: 4 * w + 4]
return add(s[0:2], s[2:4])
h.append(add(add(pair(0, t), pair(1, t)), add(pair(0, t + 4), pair(1, t + 4))))
h01 = [add(h[t], h[t ^ 1]) for t in range(4)]
h2 = add(h01[0], h01[2])
s = H(H(h2[0] + h2[1]) + H(2.0 ** -13))
r = H(F(1) / np.sqrt(F(s)))
x = np.concatenate(tok_arena)[ARENA]
return (x * (r * gain).astype(H)).astype(H), x
def softmax_cubic(x):
"""E5: the cubic-exponent softmax of a 64 x 96 score block (natural columns)"""
def fma16(a, b, c):
return (a.astype(np.float64) * b.astype(np.float64) + c.astype(np.float64)).astype(H)
m = np.clip((x * H(0.017181396484375)).astype(H), H(-0.55615234375), H(0.55615234375))
t = fma16(m, -m, np.full_like(m, H(0.92724609375)))
u = fma16(m, t, np.full_like(m, H(1.375)))
e = ((((u.view(np.uint16).astype(np.uint32)) << 5) & 0x7FE0).astype(np.uint16)).view(H)
def tile(j, tt, half):
return e[:, 8 * j + 2 * tt + half]
s = {}
for tt in range(4):
for half in range(2):
ev = add(add(add(tile(0, tt, half), tile(2, tt, half)), add(tile(4, tt, half), tile(6, tt, half))),
add(tile(8, tt, half), tile(10, tt, half)))
od = add(add(add(tile(1, tt, half), tile(3, tt, half)), add(tile(5, tt, half), tile(7, tt, half))),
add(tile(9, tt, half), tile(11, tt, half)))
s[tt, half] = add(ev, od)
out = [add(add(s[0, h_], s[1, h_]), add(s[2, h_], s[3, h_])) for h_ in range(2)]
rs = add(out[0], out[1])
r = (F(1) / rs.astype(F)).astype(H)
return (e * r[:, None]).astype(H)
def act(x):
y = np.clip(x, H(-2), H(2)).astype(H)
g = add(H(0.5), (y * (H(0.412109375) - (H(0.0810546875) * np.abs(y)).astype(H)).astype(H)).astype(H))
return (x * g).astype(H)
MERGE_ORDER = [(0, 0), (0, 1), (1, 0), (1, 1), (2, 0), (2, 1), (3, 0), (3, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 2),
(2, 3), (3, 2), (3, 3)]
def mirror(v, n):
v = abs(v)
return min(v, 2 * n - 2 - v)
def window(w, planes, tex, x0, y0, ex, ey, dbg=None):
"""one window: returns (tokens out [64 x 64 codes, arena order], merged [16 x 96 codes, A order],
surf112 / surf120 [64 x 3 f16], surf128 contributions)"""
border = ([(x, y0 - 1) for x in range(x0, x0 + 8)] + [(x, y0 + 8) for x in range(x0, x0 + 8)]
+ [(x0 - 1, y) for y in range(y0, y0 + 8)] + [(x0 + 8, y) for y in range(y0, y0 + 8)])
coords = [(x, y) for y in range(y0, y0 + 8) for x in range(x0, x0 + 8)] + border
norm, raw = [], []
for (x, y) in coords:
xx, yy = mirror(x, ex), mirror(y, ey)
n_, r_ = input_norm([e4m3(planes[p, yy, xx]).astype(H) for p in (0, 1)], w['gain0'])
norm.append(n_)
raw.append(r_)
X = q8(np.array(norm))[:, INTER64] # 96 x 64 codes, NVIDIA k order
res = np.array(raw[:64])
dbg = {} if dbg is None else dbg
dbg['XQ'], dbg['RES0'] = X, add(res, w['table7'][None, :])
R = None
for h in range(2):
qk = mm(X[:, 32:], w[f'wk{h}b'], mm(X[:, :32], w[f'wk{h}a'], np.zeros((96, 32))))
v = mm(X[:, 32:], w[f'wv{h}b'], mm(X[:, :32], w[f'wv{h}a'], np.zeros((96, 32))))
Q = q8(qk)[:, inter(32)]
dbg[f'QK{h}'], dbg[f'V{h}'] = Q, q8(v)[inter(96)]
S = mm(Q[:64], Q.T, w[f'bias{h}'])
P = softmax_cubic(S)
O = mm(q8(P)[:, inter(96)], q8(v)[inter(96)], np.zeros((64, 32)))
dbg[f'S{h}'], dbg[f'P{h}'], dbg[f'O{h}'] = S, q8(P)[:, inter(96)], O
C = add(res, w['table7'][None, :]) if h == 0 else R
R = mm(q8(O)[:, inter(32)], w[f'wo{h}'], C)
dbg[f'R{h}'] = R
N = q8((R * (nv_rowsum_sq(R)[:, None] * w['gain14'][None, :]).astype(H)).astype(H))[:, INTER64]
acc = add(R, w['table16'][None, :])
for b in range(8):
hid = act(mm(N, w[f'w1_{b}'], np.repeat(w[f'b1_{b}'][None, :], 64, 0)))
acc = mm(q8(hid)[:, inter(32)], w[f'w2_{b}'], acc)
dbg['ACC'] = acc
tokens = q8(acc)[:, INTER64]
Mg = np.zeros((16, 256), np.uint8)
for r, (mx, my) in enumerate(MERGE_ORDER):
for blk, (dx, dy) in enumerate([(0, 0), (1, 0), (0, 1), (1, 1)]):
Mg[r, 64 * blk: 64 * blk + 64] = tokens[(2 * my + dy) * 8 + 2 * mx + dx]
merged = q8(mm(Mg, w['wm'], np.repeat(w['bm'][None, :], 16, 0)))[:, inter(96)]
head = mm(tokens, w['wh'], np.repeat(w['bh'][None, :], 64, 0))
s112, s120 = filter_head(head, tex, x0, y0, ex, ey)
return tokens, merged, s112, s120
def filter_head(D, tex, x0, y0, ex, ey, ungated=False):
"""surfaces 112 and 120 of the encoders (the gate's split of the filtered value); ungated (the bottleneck enc5):
the filtered value itself, for surface 112"""
def sig(v):
return H(F(1) / (np.exp2(F(v) * F(-1.4426950216293335)) + F(1)))
def ex2(v):
return np.exp2(F(v) * F(1.4426950216293335))
out112, out120 = np.zeros((64, 3), H), np.zeros((64, 3), H)
for tk in range(64):
x, y = x0 + tk % 8, y0 + tk // 8
l_ = D[tk].astype(H)
m01 = np.maximum(np.maximum(np.maximum(l_[0:2], l_[2:4]), l_[4:6]), l_[6:8])
m = max(max(m01[0], m01[1]), l_[8])
e = [ex2(H(l_[i] - m)) for i in range(9)]
pair = [F(e[2 * k] + e[2 * k + 1]) for k in range(4)]
s = F(F(pair[0]) + F(0))
for k in (1, 2, 3):
s = F(s + pair[k])
s = F(s + e[8])
r = H(F(1) / s)
wgt = [H(H(e[i]) * r) for i in range(9)]
def cl(v, n):
return 1 if v < 0 else (n - 2 if v >= n else v)
xm, xc, xp = cl(x - 1, ex), cl(x, ex), cl(x + 1, ex)
ym, yc, yp = cl(y - 1, ey), cl(y, ey), cl(y + 1, ey)
def T(xx, yy):
return tex[yy, xx].astype(H)
pts = [((xm, ym), (xc, ym)), ((xp, ym), (xm, yc)), ((xc, yc), (xp, yc)), ((xm, yp), (xc, yp))]
acc = np.zeros((3, 2), H)
for k, (pa, pb) in enumerate(pts):
for c in range(3):
for h_, pt in enumerate((pa, pb)):
acc[c, h_] = H(F(T(*pt)[c]) * F(wgt[2 * k + h_]) + F(acc[c, h_]))
f = [H(acc[c, 0] + acc[c, 1]) for c in range(3)]
f = [H(F(T(xp, yp)[c]) * F(wgt[8]) + F(f[c])) for c in range(3)]
if ungated:
out112[tk] = f
continue
g = [sig(l_[9]), sig(l_[10]), sig(l_[11])]
a = [H(g[c] * f[c]) for c in range(3)]
out120[tk] = a
out112[tk] = [H(f[c] - a[c]) for c in range(3)]
return out112, out120
def extract(dump, cap, frame, out):
"""the weights, from a phase dump of any block (they do not depend on the block) and the weight image"""
z = np.load(dump)
def B(p):
return z[f'p{p}_w0_B']
def C(p):
return f16(z[f'p{p}_w0_C']).astype(H)
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rr_capture import Capture
c = Capture(os.path.expanduser(cap))
W = np.frombuffer(c.find('swin_enc1', frame).weights, np.uint8)
def table(off, n=64):
return W[off:off + 2 * n].view(np.float16)
w = {'gain0': table(0), 'gain14': table(37120), 'table7': table(36992), 'table16': table(41408)}
for h, base in ((0, 1), (1, 8)):
w[f'wk{h}a'], w[f'wv{h}a'], w[f'wk{h}b'], w[f'wv{h}b'] = B(base), B(base + 1), B(base + 2), B(base + 3)
w[f'bias{h}'] = C(base + 4)
w[f'wo{h}'] = B(base + 6)
for b in range(8):
w[f'w1_{b}'], w[f'b1_{b}'] = B(15 + 2 * b), C(15 + 2 * b)[0]
w[f'w2_{b}'] = B(16 + 2 * b)
w['wm'], w['bm'] = B(31), C(31)[0]
w['wh'], w['bh'] = B(32), C(32)[0]
np.savez_compressed(out, **w)
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('frame', type=int)
ap.add_argument('weights')
ap.add_argument('--extract', metavar='DUMP')
ap.add_argument('--blocks', nargs='*', default=['5,5'])
o = ap.parse_args()
if o.extract:
extract(o.extract, o.capture, o.frame, o.weights)
return
from rr_capture import Capture
c = Capture(os.path.expanduser(o.capture))
l = c.find('swin_enc1', o.frame)
ex, ey = (int(v) for v in np.frombuffer(l.args[:8], np.uint32))
sx, sy = (int(v) for v in np.frombuffer(l.args[80:84], np.uint16))
pre = np.frombuffer(l.pre(l.arena_alloc), np.uint8)
post = np.frombuffer(l.post_buffer(c, l.arena_alloc), np.uint8)
ptr = {p.offset: p.delta for p in l.pointers}
planes = pre[ptr[8]:ptr[8] + 2 * ex * ey * 32].reshape(2, ey, ex, 32)
outp = post[ptr[48]:ptr[48] + 2 * ex * ey * 32].reshape(2, ey, ex, 32)
mrg = post[ptr[56]:ptr[56] + 3 * (ex // 2) * (ey // 2) * 32].reshape(3, ey // 2, ex // 2, 32)
img = {i.offset: i for i in l.images}
def surf(off):
i = img[off]
return np.frombuffer(i.path.read_bytes(), np.float16).reshape(i.height, i.width, 4)
tex, s112, s120 = surf(88), surf(112), surf(120)
w = dict(np.load(o.weights))
for blk in o.blocks:
bx, by = (int(v) for v in blk.split(','))
x0, y0 = 8 * bx - sx, 8 * by - sy
tokens, merged, o112, o120 = window(w, planes, tex, x0, y0, ex, ey)
tk_ok = sum(np.array_equal(tokens[t], np.concatenate([outp[0, y0 + t // 8, x0 + t % 8],
outp[1, y0 + t // 8, x0 + t % 8]]))
for t in range(64) if 0 <= x0 + t % 8 < ex and 0 <= y0 + t // 8 < ey)
mg_ok = sum(np.array_equal(merged[r], np.concatenate([mrg[p, y0 // 2 + my, x0 // 2 + mx] for p in range(3)]))
for r, (mx, my) in enumerate(MERGE_ORDER)
if 0 <= x0 // 2 + mx < ex // 2 and 0 <= y0 // 2 + my < ey // 2)
inside = [t for t in range(64) if 0 <= x0 + t % 8 < ex and 0 <= y0 + t // 8 < ey]
s_ok = sum(np.array_equal(o112[t], s112[y0 + t // 8, x0 + t % 8, :3]) and
np.array_equal(o120[t], s120[y0 + t // 8, x0 + t % 8, :3]) for t in inside)
print(f'block {bx},{by}: tokens {tk_ok}/{len(inside)}, merged {mg_ok}, surfaces {s_ok}/{len(inside)}')
if __name__ == '__main__':
main()
def pairs_tree(v):
"""NVIDIA's reduction over a list: (v0 + v1) + v2 for three, (v0 + v1) + (v2 + v3) for four, then the same
recursively on halves rounded up ((v0 + v1) + (v2 + v3)) + v4 for five"""
if len(v) == 1:
return v[0]
if len(v) == 2:
return add(v[0], v[1])
half = 1 << ((len(v) - 1).bit_length() - 1)
return add(pairs_tree(v[:half]), pairs_tree(v[half:]))
def rowsum_sq(x, eps=2.0 ** -13):
"""nv_rowsum_sq for any C = 8 n8-tiles: per lane word t the tiles of each parity (tile pairs 2 apart, tree),
the parities, then t ^ 1, t ^ 2 and the halves"""
sq = (x * x).astype(H)
nt = x.shape[1] // 8
def tile(j, t):
return sq[:, 8 * j + 2 * t: 8 * j + 2 * t + 2]
h = []
for t in range(4):
ev = pairs_tree([add(tile(j, t), tile(j + 2, t)) for j in range(0, nt, 4)] if nt % 4 == 0 else
[tile(j, t) for j in range(0, nt, 2)])
od = pairs_tree([add(tile(j, t), tile(j + 2, t)) for j in range(1, nt, 4)] if nt % 4 == 0 else
[tile(j, t) for j in range(1, nt, 2)])
h.append(add(ev, od))
h1 = [add(h[t], h[t ^ 1]) for t in range(4)]
h2 = add(h1[0], h1[2])
s = add(add(h2[:, 0], h2[:, 1]), np.full(x.shape[0], H(eps)))
return (F(1) / np.sqrt(s.astype(F))).astype(H)
engine/rr/f16_to_e4m3.npy
BIN
View File
Binary file not shown.
engine/rr/gen_enc0_m.py
+45
View File
@@ -0,0 +1,45 @@
#!/usr/bin/env python3
"""Generate engine/rr/enc0_m.comp's source: encoder 0 with ptx2rr.py's translation of NVIDIA's per-pixel feature gather
(run by wave 0, as NVIDIA runs it in one warp) and the native rest (engine/rr/enc0_m_decl.glsl, enc0_m_tail.glsl).
usage: gen_enc0_m.py ENC0.generated.comp OUT.comp
The translation's main is cut at its first workgroup barrier: by then the 10 x 10 pixels' features and the noise grid
are in its shared buffer sh and the state surfaces 304 / 1016 are written. Model data derived from the user's DLL is in
the translated part; do not commit the output.
"""
import re
import sys
def main():
src, out = sys.argv[1:3]
lines = open(src).read().split('\n')
m = next(i for i, l in enumerate(lines) if l.startswith('void main()'))
head = '\n'.join(lines[:m]).replace('layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in;',
'layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in;')
body = lines[m + 2:]
first = next(i for i, l in enumerate(body) if not re.match(r' (precise |bool |uint64_t |float |uint |int )', l))
bar = next(i for i, l in enumerate(body) if l.strip() == 'memoryBarrierShared(); barrier();')
decls, prologue = body[:first], body[first:bar]
# NVIDIA's warp gathers the 96 key pixels in a loop of three 32-pixel steps, then the four corners: the steps go to
# waves 0-2 (the setup before the loop runs in every wave, the corners in wave 0)
do_ = next(i for i, l in enumerate(prologue) if l.strip() == 'do')
end = next(i for i, l in enumerate(prologue) if l.strip().startswith('} while (again_'))
loop = '\n'.join(prologue[do_ - 1:end + 1])
inc = re.search(r'(r\d+) = uint\((r\d+) \+ 32u\);\n\s+(p\d+) = \(\1 < 96u\);', loop)
counter = inc.group(1)
loop = loop.replace(inc.group(0), inc.group(0).replace('+ 32u', '+ 96u'))
pre = [l.replace(f'{counter} = uint(0u);', f'{counter} = uint(32u * gl_SubgroupID);') for l in prologue[:do_ - 1]]
post = prologue[end + 1:]
head = head.replace('shared uint mms[384];', 'shared uint mms[384 * 4];')
text = (head + '\n#extension GL_GOOGLE_include_directive : require\n#include "enc0_m_decl.glsl"\n\nvoid main()\n{\n' +
'\n'.join(decls) + '\n // the translated feature gather (NVIDIA\'s one warp, its pixel loop split over the waves)\n' +
'\n'.join(pre) + '\n if (gl_SubgroupID < 3u) {\n' + loop + '\n }\n if (gl_SubgroupID == 0u) {\n' +
'\n'.join(post) + '\n }\n#include "enc0_m_tail.glsl"\n}\n')
open(out, 'w').write(text)
print(f'{out}: {len(decls)} declaration lines, {len(prologue)} prologue lines')
if __name__ == '__main__':
main()
engine/rr/hkpn.comp
+298
View File
@@ -0,0 +1,298 @@
#version 460
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
#extension GL_EXT_shader_16bit_storage : require
// Ray Reconstruction hkpn output stage, preset E (no diamond-wallaby branch).
// One workgroup per 16x16 output tile. Stage 1 stages the render-space neighbourhood of
// the tile into shared memory (colour log, per-texel weights, anisotropy); stage 2 gathers
// a 5x5 elliptical kernel per output pixel and writes the two output surfaces.
// Parameter block (binding 0) holds NVIDIA's kernel argument block; word index = byte / 4.
layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in;
layout(std430, binding = 0) readonly buffer Params { uint P[]; };
layout(std430, binding = 2) buffer Arena { uint A[]; };
layout(binding = 3) uniform sampler2D colourTex; // arg 176, 640x360 rgba16f
layout(binding = 5) uniform sampler2D exposureTex; // arg 200, 1x1 rgba16f
layout(binding = 8) uniform sampler2D albedoA; // arg 280, 640x360 rgba8unorm
layout(binding = 9) uniform sampler2D albedoB; // arg 288, 640x360 rgba8unorm
layout(binding = 11) uniform sampler2D decTex; // arg 448, 640x384 rgba16f
layout(rgba16f, binding = 12) uniform writeonly image2D outLog; // arg 224
layout(rgba16f, binding = 13) uniform writeonly image2D outLin; // arg 240
#define F_LOG2E uintBitsToFloat(0x3FB8AA3Bu)
#define F_2LOG2E uintBitsToFloat(0x4038AA3Bu)
#define F_LOG10_2 uintBitsToFloat(0x3E9A209Bu)
#define F_0_36 uintBitsToFloat(0x3EB851ECu)
#define F_100 uintBitsToFloat(0x42C80000u)
#define F_1E4 uintBitsToFloat(0x38D1B717u)
#define F_0_005 uintBitsToFloat(0x3BA3D70Au)
#define F_0_1 uintBitsToFloat(0x3DCCCCCDu)
#define F_1_7 uintBitsToFloat(0x3FD9999Au)
#define F_0_17 uintBitsToFloat(0x3E2E147Bu)
#define F_0_95 uintBitsToFloat(0x3F733333u)
#define F_0_6 uintBitsToFloat(0x3F19999Au)
#define F_9 uintBitsToFloat(0x41102CB4u)
// shared tile: 21 entries per row, up to 21x21 entries (render scale <= 1)
#define TILE_STRIDE 21
// binary16 numbers, stored as such: half the shared-memory traffic of the 5x5 gather
shared f16vec4 tLog[TILE_STRIDE * TILE_STRIDE]; // (L_R, L_G, L_B, alpha = 0) log-domain colour
shared f16vec4 tSig[TILE_STRIDE * TILE_STRIDE]; // (s0, s1, s2, s3) sigmoid of arena bytes 0..3
shared vec4 tAn[TILE_STRIDE * TILE_STRIDE]; // (p0, p1, p2, p3) anisotropy parameters
float pf(uint idx) { return uintBitsToFloat(P[idx]); }
// e4m3 byte -> float (exact); e4m3fn NaN for 0x7F/0xFF.
float e4m3(uint b) {
uint e = (b >> 3u) & 15u;
uint m = b & 7u;
if (e == 15u && m == 7u) return uintBitsToFloat(0x7FC00000u);
float v = (e == 0u) ? float(m) * uintBitsToFloat(0x3B000000u)
: (1.0 + float(m) * 0.125) * exp2(float(e) - 7.0);
return ((b & 128u) != 0u) ? -v : v;
}
// round to nearest even binary16, returned as f32: the hardware conversion (RNE), kept from being folded back to x by
// OR-ing an opaque zero into its bits (RADV folds a register-only f32 -> f16 -> f32)
// Beyond the binary16 range NVIDIA's arithmetic keeps a finite rounded value: that case takes the explicit rounding.
uint16_t zero16;
float f16rnSlow(float x) {
if (x == 0.0 || isnan(x) || isinf(x)) return x;
int e;
frexp(x, e);
int ee = max(e - 11, -24);
float q = uintBitsToFloat(uint(ee + 127) << 23);
return roundEven(x / q) * q;
}
float f16rn(float x) {
if (abs(x) >= 65504.0) return f16rnSlow(x);
return float(uint16BitsToFloat16(float16BitsToUint16(float16_t(x)) | zero16));
}
// Sigmoid as evaluated in binary16 on this lowering: r = f16(2^(-x*log2e)), den = f16(1+r),
// s = f16(1/den), with one Newton step for small s and the observed fix-ups of r.
float sig16(float e) {
float x = f16rn(e);
float r = f16rn(exp2(-x * F_LOG2E));
uint hb = uint(float16BitsToUint16(float16_t(r)));
if (hb == 0x1F79u || hb == 0x25CFu) {
r = f16rn(r - 0.0009765625);
} else if (hb == 0xC13Bu) {
r = f16rn(r + 0.00006103515625);
} else if (hb == 0xC1EFu) {
r = f16rn(r + 0.000030517578125);
}
float d32 = f16rn(1.0 + r);
float rc = 1.0 / d32;
float s16 = f16rn(rc);
uint sb = uint(float16BitsToUint16(float16_t(s16))) & 0x7FFFu;
if (sb != 0u && sb < 0x8Fu) {
float ee = fma(-d32, s16, 1.0);
s16 = f16rn(fma(rc, ee, s16));
}
return s16;
}
// tanh as evaluated: polynomial for small |x|, exponential form otherwise.
float tanhApprox(float x) {
float ax = abs(x);
float x2 = x * x;
float q = fma(uintBitsToFloat(0x3C80F082u), x2, uintBitsToFloat(0xBD563CAEu));
q = fma(q, x2, uintBitsToFloat(0x3E085941u));
q = fma(q, x2, uintBitsToFloat(0xBEAAA9EDu));
q = fma(q, x2, 0.0);
float small = fma(q, x, x);
float ex = exp2(ax * F_2LOG2E);
float big = fma(1.0 / (ex + 1.0), -2.0, 1.0);
if (ax >= F_9) big = 1.0;
float bigSigned = uintBitsToFloat((floatBitsToUint(x) & 0x80000000u) | floatBitsToUint(big));
return (ax >= F_0_6) ? bigSigned : small;
}
// round half away from zero
int roundAway(float v) {
return int(trunc(v + ((v < 0.0) ? -0.5 : 0.5)));
}
void main() {
zero16 = uint16_t(P[0] >> 24u); // the output width is below 2^24
// ---- parameters (byte offsets / 4) ----
uint outW = P[0];
uint outH = P[1];
float sx = pf(6u); // 24: render/out scale x
float sy = pf(7u); // 28
float ox = pf(10u); // 40: jitter offset x
float oy = pf(11u); // 44
float expScale = pf(37u); // 148
uint decMaxX = P[14] - 1u; // 56: dec0 width
uint decMaxY = P[15] - 1u; // 60: dec0 height
uint colMaxX = P[35]; // 140: colour/albedo max x
uint colMaxY = P[36]; // 144: colour/albedo max y
uint arenaBase = P[64]; // 256: arena offset
uint arenaPitch = P[4]; // 16: render width (640)
// ---- tile origin / size in render space ----
uvec2 wg = gl_WorkGroupID.xy;
float tileX = ox + fma(float(wg.x * 16u) + 0.5, sx, -0.5);
float tileY = oy + fma(float(wg.y * 16u) + 0.5, sy, -0.5);
int c0 = int(floor(tileX)) - 2;
int r0 = int(floor(tileY)) - 2;
int tw = int(ceil(16.0 * sx)) + 5;
int th = int(ceil(16.0 * sy)) + 5;
// ---- stage 1: tile load ----
uint total = uint(tw * th);
for (uint idx = gl_LocalInvocationIndex; idx < total; idx += 256u) {
int col = int(idx) % tw;
int row = int(idx) / tw;
int ix = clamp(c0 + col, 0, int(decMaxX));
int iy = clamp(r0 + row, 0, int(decMaxY));
vec3 d = texelFetch(decTex, ivec2(ix, iy), 0).rgb;
int ax = min(ix, int(colMaxX));
int ay = min(iy, int(colMaxY));
vec3 alb = texelFetch(albedoA, ivec2(ax, ay), 0).rgb
+ texelFetch(albedoB, ivec2(ax, ay), 0).rgb;
float albSum = (alb.r + alb.g) + alb.b;
if (albSum != 0.0) {
d.r *= max(alb.r, F_0_005);
d.g *= max(alb.g, F_0_005);
d.b *= max(alb.b, F_0_005);
}
vec3 lg;
lg.r = f16rn(F_0_36 * (log2(max(fma(d.r, F_100, 1.0), F_1E4)) * F_LOG10_2));
lg.g = f16rn(F_0_36 * (log2(max(fma(d.g, F_100, 1.0), F_1E4)) * F_LOG10_2));
lg.b = f16rn(F_0_36 * (log2(max(fma(d.b, F_100, 1.0), F_1E4)) * F_LOG10_2));
uint wordIdx = (arenaBase >> 2u) + uint(arenaPitch * uint(iy) + uint(ix)) * 2u;
uint lo = A[wordIdx];
uint hi = A[wordIdx + 1u];
vec4 sg = vec4(sig16(e4m3(lo & 0xFFu)),
sig16(e4m3((lo >> 8u) & 0xFFu)),
sig16(e4m3((lo >> 16u) & 0xFFu)),
sig16(e4m3(lo >> 24u)));
vec4 an = vec4(e4m3(hi & 0xFFu),
e4m3((hi >> 8u) & 0xFFu),
e4m3((hi >> 16u) & 0xFFu),
e4m3(hi >> 24u));
uint se = uint(row * TILE_STRIDE + col);
tLog[se] = f16vec4(lg, 0.0);
tSig[se] = f16vec4(sg);
tAn[se] = an;
}
barrier();
// ---- stage 2: per output pixel ----
uvec2 pix = gl_GlobalInvocationID.xy;
if (pix.x >= outW || pix.y >= outH) return;
float expo = texelFetch(exposureTex, ivec2(0, 0), 0).r;
float e41 = ((expo == 0.0) ? 1.0 : expo) * expScale;
float fx = (float(pix.x) + 0.5) * sx;
float fy = (float(pix.y) + 0.5) * sy;
float cx = fx + ox;
float cy = fy + oy;
// colour fetch: log of exposure-scaled colour at the nearest render texel
int ux = clamp(roundAway((fx - 0.5) + ox), 0, int(colMaxX));
int uy = clamp(roundAway((fy - 0.5) + oy), 0, int(colMaxY));
vec3 col = texelFetch(colourTex, ivec2(ux, uy), 0).rgb;
vec3 gcol;
gcol.r = F_0_36 * (log2(max(fma(col.r * e41, F_100, 1.0), F_1E4)) * F_LOG10_2);
gcol.g = F_0_36 * (log2(max(fma(col.g * e41, F_100, 1.0), F_1E4)) * F_LOG10_2);
gcol.b = F_0_36 * (log2(max(fma(col.b * e41, F_100, 1.0), F_1E4)) * F_LOG10_2);
// centre entry, anisotropy
float flx = floor(cx);
float fly = floor(cy);
int ix0 = int(flx) - c0;
int iy0 = int(fly) - r0;
float dx0 = (flx + 0.5) - cx;
float dy0 = (fly + 0.5) - cy;
float ks = (sy >= 0.5) ? 1.0 : (0.5 / sy);
vec4 an = tAn[iy0 * TILE_STRIDE + ix0];
float sa = fma(1.0 / (1.0 + exp2(-F_LOG2E * an.x)), F_1_7, F_0_17);
float sb = fma(1.0 / (1.0 + exp2(-F_LOG2E * an.y)), F_1_7, F_0_17);
float rho = F_0_95 * tanhApprox(an.z);
float sa2 = sa * sa;
float sb2 = sb * sb;
float denom = (sa2 * sb2) * (1.0 - rho * rho);
float invDenom = -(1.0 / denom);
float cross2 = rho * (2.0 * sqrt(sa2 * sb2));
// 5x5 gather
float sumw = 0.0;
vec3 accL = vec3(0.0);
vec4 accS = vec4(0.0);
for (int c = -2; c <= 2; c++) {
for (int r = -2; r <= 2; r++) {
float dx = ks * (dx0 + float(c));
float dy = ks * (dy0 + float(r));
uint se = uint((iy0 + r) * TILE_STRIDE + (ix0 + c));
precise float t_c = (cross2 * dx) * dy;
precise float t_b = dx * (sb2 * dx);
precise float t_a = dy * (sa2 * dy);
precise float num = (t_c - t_b) - t_a;
precise float Q = invDenom * num;
float w = exp2((Q * -0.5) * F_LOG2E);
sumw += w;
vec4 lv = vec4(tLog[se]);
vec4 sv = vec4(tSig[se]);
accL = fma(vec3(w), lv.rgb, accL);
accS = fma(vec4(w), sv, accS);
}
}
float inv = 1.0 / sumw;
vec3 Lm = accL * inv;
float S0 = accS.x * inv;
float S1 = accS.y * inv;
float S2 = accS.z * inv;
// blend weight
float q = F_0_1 / (sx * sy);
float t1 = sumw * q;
float t2 = S0 * t1;
float t3 = (sx * sy) * t2;
float t4 = 4.0 * t3;
float t5 = (sy < 0.5) ? t4 : t2;
float wgt = (t5 < 0.0) ? 0.0 : ((t5 > 1.0) ? 1.0 : t5);
float fA = fma(1.0 - S1, wgt, S1);
float fB = fma(1.0 - S2, wgt, S2);
// surface 224 (E): fA * D + (1 - fA) * gcol; D = Lm (f16 tile log channels),
// gcol = log of the exposure-scaled colour texel (PTX linear texel fetch path); D.a = 0 so aA = 1 - fA
vec3 cA = fma(Lm, vec3(fA), gcol * (1.0 - fA));
float aA = 1.0 - fA;
float sumA = cA.r + cA.g + cA.b + aA;
if (isnan(sumA) || isinf(sumA)) {
cA = Lm;
aA = 0.0;
}
imageStore(outLog, ivec2(pix), vec4(f16rn(cA.r), f16rn(cA.g), f16rn(cA.b), f16rn(aA)));
// surface 240 (E, PTX BB0_83): {R,G,B,1} of lin = (10^(clamp(c)/0.36)-1)/100/e41,
// c = fma(Lm, fB, gcol*(1-fB)); c = Lm when the sum is not finite; stored RNE f16
vec3 cB = fma(Lm, vec3(fB), gcol * (1.0 - fB));
float sumB = cB.r + cB.g + cB.b;
if (isnan(sumB) || isinf(sumB)) cB = Lm;
cB = clamp(cB, vec3(F_1E4), vec3(3.0));
// PTX divides by constants and takes lg2(10) with the hardware approximations at run time; keep the
// constants opaque so the compiler cannot fold them with different rounding.
uint zero = P[0] >> 24u;
float c036 = uintBitsToFloat(0x3EB851ECu | zero);
float c100 = uintBitsToFloat(0x42C80000u | zero);
float lg10 = log2(uintBitsToFloat(0x41200000u | zero));
float rc036 = 1.0 / c036;
float rc100 = 1.0 / c100;
float rce41 = 1.0 / e41;
precise vec3 lin = (exp2((cB * rc036) * lg10) - 1.0) * rc100 * rce41;
imageStore(outLin, ivec2(pix), vec4(float(float16_t(lin.r)), float(float16_t(lin.g)), float(float16_t(lin.b)), 1.0));
}
engine/rr/native_kernels.txt
+15
View File
@@ -0,0 +1,15 @@
# SHA-256 of the PTX modules of the network kernels engine/rr's native shaders (swin_m.comp, dec0_m.comp, enc0_m)
# were written against and verified with: nvngx_dlssd.dll 310.7. engine/rr/dll_weights.py DLL --kernels prints a DLL's.
cuda_dldn_engine_swin_dec0_diamond_wallaby_kernel f9420970b6e29c1cd100057ba9eb56a0465a95a4aaeaa21dd355605b75160e5b
cuda_dldn_engine_swin_dec0_kernel 5e533151fd1d469bb344548b113e67b5fc798db807cebaf627864ce877f01242
cuda_dldn_engine_swin_dec1_kernel ac293a9826ea7c8be44168698e79f122d0bb2cb23dbf55999a094876e04ce32f
cuda_dldn_engine_swin_dec2_kernel 282d8f093bf343c7a9080c7379ff933443943d503c7dea647535444711494528
cuda_dldn_engine_swin_dec3_kernel f2325cc78cac7c3f20d4b3b5d6604b8a350774dd83642da10ab525bd3fd1e5b0
cuda_dldn_engine_swin_dec4_kernel 9bb7cfeba4380b023f2a1b2db5fa311edd40d4e4fdce62af2fe9227a83c52213
cuda_dldn_engine_swin_enc0_diamond_wallaby_kernel 35a2b205de3feb52dacfd54ebdcb2eda4dc75fe78a181a28687876f9a1d2a17d
cuda_dldn_engine_swin_enc0_kernel e0927a196eeffc93996a0c53337fe6fecc8a869145d832a02c59d9e12741db9d
cuda_dldn_engine_swin_enc1_kernel c9df416b891e4dddbeb6311ce110b8362c1c7f1d9e9efb2595b9b4ccae155fce
cuda_dldn_engine_swin_enc2_kernel 58702ea68b478da919a31cb9052fdb82eba3d7d1ec0d94ef2bb2827686a3b36a
cuda_dldn_engine_swin_enc3_kernel 7f07b4cf84adcdc6bfb367c00767b1f124bbdb1956f704a5dbc72f9e1d3fb9ac
cuda_dldn_engine_swin_enc4_kernel d1466bfa00c702d1a67820d1ad12e351fc4185d90695b21286c1d9523e772d6e
cuda_dldn_engine_swin_enc5_kernel 720c531da5552ff9a0304c8628aea46d9eaeae3c60f9ee0cdb4f94bebfe1a6e0
engine/rr/native_pack.py
+46
View File
@@ -0,0 +1,46 @@
#!/usr/bin/env python3
"""Pack a denoiser stage's weights (engine/rr/enc1_ref.py --extract) for the native stage shaders.
usage: native_pack.py WEIGHTS.npz OUT.bin OUT.h
Every FP8 matrix W (K x N, rows in NVIDIA's k order) becomes 16x16 tiles, tile (kt, nt) at
256 * (kt * N/16 + nt) bytes, row-major inside the tile (stride 16): a B operand of the cooperative matrix loads.
Every f16 matrix is stored as is (row-major); an f16 bias vector (a 1-D entry whose name starts with 'b') as 16
equal rows, so that it loads as an accumulator tile. OUT.h lists the byte offset of every entry as a #define,
for the shader's include.
"""
import sys
import numpy as np
def tiles(w):
k, n = w.shape
assert k % 16 == 0 and n % 16 == 0, w.shape
return w.reshape(k // 16, 16, n // 16, 16).transpose(0, 2, 1, 3).reshape(-1)
def main():
src, out_bin, out_h = sys.argv[1:4]
w = dict(np.load(src))
blob = bytearray()
defines = []
for name in sorted(w):
a = w[name]
while len(blob) % 256:
blob.append(0)
off = len(blob)
if a.dtype == np.uint8:
blob += tiles(a).tobytes()
elif a.ndim == 1 and name.startswith('b'):
blob += np.tile(np.asarray(a, np.float16), (16, 1)).tobytes()
else:
blob += np.asarray(a, np.float16).tobytes()
defines.append(f'#define W_{name.upper()} {off}u // {a.dtype} {a.shape}')
open(out_bin, 'wb').write(bytes(blob))
open(out_h, 'w').write('// generated by engine/rr/native_pack.py\n' + '\n'.join(defines) + '\n')
print(f'{len(blob)} bytes, {len(w)} entries')
if __name__ == '__main__':
main()
engine/rr/phase_dump.py
+203
View File
@@ -0,0 +1,203 @@
#!/usr/bin/env python3
"""Dump every mma of a denoiser layer as matrices, from a ZLUDA replay of a captured launch: the oracle for porting
a layer to native shaders.
usage: phase_dump.py CAPTURE STAGE FRAME BX,BY OUT.npz [--phases A-B]
For one thread block, debug stores added to the PTX record every mma's A, B, C and D registers (each mma's last
execution) into padding behind the plane arena; the replay runs through ZLUDA (engine/rr/zluda_replay.py). The
fragments are then put back together with the m16n8k32 layouts (A row-major 16x32 e4m3, B column-major 32x8 e4m3,
C/D 16x8 f16) into one matrix per phase: phase P's mma are grouped into m-tiles by A fragment and into n-tiles by
B fragment, k-steps follow the C <- D chains. OUT.npz holds, per phase p and wave w: p{p}_w{w}_A (M x K, e4m3 codes),
_B (K x N codes), _C (M x N f16 of the first k-step), _D (M x N f16 of the last).
"""
import argparse
import os
import re
import subprocess
import sys
from pathlib import Path
import numpy as np
HERE = Path(__file__).resolve().parent
ROOT = HERE.parents[1]
sys.path.insert(0, str(HERE))
from rr_capture import Capture # noqa: E402
REG = r'%r\d+'
def kernel_module(kernel):
for p in sorted(Path('~/.cache/d4r-rr-corpus').expanduser().glob('*.ptx')):
if f'.entry {kernel}(' in p.read_text(errors='replace'):
return p
raise SystemExit(f'no corpus module for {kernel}')
def mma_statements(lines, start, end):
"""(first line, last line, D, A, B, C) of every mma in [start, end)"""
out, i = [], start
while i < end:
if 'mma.sync.aligned.m16n8k32' in lines[i]:
j, text = i, ''
while True:
text += lines[j]
if lines[j].rstrip().endswith(';'):
break
j += 1
groups = [re.findall(REG, g) for g in re.findall(r'\{([^}]*)\}', text)]
out.append((i + 1, j + 1, *groups[:4]))
i = j + 1
else:
i += 1
return out
def phases_of(mmas, lines):
"""maximal runs of mma separated only by blank or comment lines"""
phases, cur = [], []
for m in mmas:
if cur:
gap = [l.strip() for l in lines[cur[-1][1]:m[0] - 1]]
if any(g and not g.startswith('//') for g in gap):
phases.append(cur)
cur = []
cur.append(m)
if cur:
phases.append(cur)
return phases
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('stage')
ap.add_argument('frame', type=int)
ap.add_argument('block')
ap.add_argument('out')
ap.add_argument('--phases', help='first-last (1-based), default all')
o = ap.parse_args()
# the capture decides the kernel: a preset D capture runs cuda_dldn_engine_swin_{enc0,dec0}_diamond_wallaby_kernel
cap = Capture(os.path.expanduser(o.capture))
launch = [l for l in cap.launches if f'swin_{o.stage}' in l.kernel][o.frame - 1]
kernel = launch.kernel
module = kernel_module(kernel)
lines = module.read_text().replace('\r', '').split('\n')
k = next(i for i, l in enumerate(lines) if f'.entry {kernel}(' in l)
e = next((i for i in range(k + 1, len(lines)) if '.entry ' in lines[i]), len(lines))
phases = phases_of(mma_statements(lines, k, e), lines)
first, last = (int(x) for x in o.phases.split('-')) if o.phases else (1, len(phases))
T = launch.block[0] * launch.block[1] * launch.block[2]
arena_bytes = len(launch.pre(launch.arena_alloc))
# any argument pointing into the arena locates its base (enc0 has none at 8)
aptr = next(p for p in launch.pointers if p.alloc == launch.arena_alloc)
delta8, off8 = aptr.delta, aptr.offset
bx, by = (int(v) for v in o.block.split(','))
# one probe slot per register per mma: after the mma, all of A, B, C, D (C read after the mma is still C
# unless C and D share registers, which the k-step chains never do for the first step)
slots, inserts = [], {}
pname = f'{kernel}_param_0'
for p in range(first, last + 1):
for (l0, l1, D, A, B, C) in phases[p - 1]:
code = []
for role, regs in (('D', D), ('A', A), ('B', B), ('C', C)):
for r in regs:
s = len(slots)
slots.append((p, l0, role, r))
code.append(f'mov.b32 %dv, {r}; mad.wide.u32 %da, %dt, 4, %db; '
f'add.s64 %da, %da, {4 * T * s}; @%dq2 st.global.u32 [%da], %dv;')
inserts[l1] = ('{ .reg .b64 %da, %db; .reg .b32 %dv, %dt, %dx, %dy; .reg .pred %dq, %dq2;\n'
f'mov.u32 %dx, %ctaid.x; mov.u32 %dy, %ctaid.y; setp.eq.u32 %dq, %dx, {bx}; '
f'setp.eq.and.u32 %dq2, %dy, {by}, %dq;\n'
f'mov.u32 %dt, %tid.z; mov.u32 %dx, %tid.y; mov.u32 %dy, %tid.x; '
f'mad.lo.u32 %dt, %dt, {launch.block[1]}, %dx; mad.lo.u32 %dt, %dt, {launch.block[0]}, %dy;\n'
f'ld.param.u64 %db, [{pname}+{off8}]; cvta.to.global.u64 %db, %db; '
f'add.s64 %db, %db, {arena_bytes - delta8};\n' + '\n'.join(code) + '\n}')
out = list(lines)
for l1 in sorted(inserts, reverse=True):
out[l1:l1] = [inserts[l1]]
work = ROOT / 'build/rr' / f'phase-{o.stage}'
work.mkdir(parents=True, exist_ok=True)
(work / 'probe.ptx').write_text('\n'.join(out))
pad = 4 * T * len(slots)
env = dict(os.environ, D4R_ZLUDA_WMMA='1', D4R_ZLUDA_WMMA_FP8='1', D4R_ZLUDA_WMMA_FP8_NATIVE='1',
D4R_ZLUDA_IMPLICIT_MAX_BLOCK='256')
cmd = ['systemd-run', '--user', '--scope', '-q', '-p', 'MemoryMax=24G', '-p', 'MemorySwapMax=0']
cmd += [f'--setenv={k_}={v_}' for k_, v_ in env.items() if k_.startswith('D4R_ZLUDA')]
r = subprocess.run(cmd + [sys.executable, str(HERE / 'zluda_replay.py'), os.path.expanduser(o.capture),
f'swin_{o.stage}' + ('_diamond' if 'diamond' in kernel else ''), str(o.frame), str(work),
'--ptx', str(work / 'probe.ptx'),
'--pad', str(pad)], capture_output=True, text=True)
if 'failed' in r.stdout + r.stderr:
sys.exit(r.stdout[-2000:] + r.stderr[-3000:])
raw = np.frombuffer((work / 'arena.out.bin').read_bytes()[arena_bytes:], np.uint32).reshape(len(slots), T)
val = {(p, l0, role, r): raw[s] for s, (p, l0, role, r) in enumerate(slots)}
result = {}
waves = T // 32
for p in range(first, last + 1):
mm = phases[p - 1]
# k-step chains: C of a later mma = D of an earlier one
dset = {tuple(m[2]): n for n, m in enumerate(mm)}
prev = {n: dset.get(tuple(m[5])) for n, m in enumerate(mm)}
step = {}
for n in range(len(mm)):
s_, q = 0, n
while prev[q] is not None:
q = prev[q]
s_ += 1
step[n] = s_
ksteps = max(step.values()) + 1
a_keys = list(dict.fromkeys(tuple(m[3]) for m in mm))
b_keys = list(dict.fromkeys(tuple(m[4]) for m in mm))
# m-tile of an A fragment: its index among the A fragments of the same k-step
a_tile = {}
for s_ in range(ksteps):
for t_, key in enumerate(dict.fromkeys(tuple(m[3]) for n, m in enumerate(mm) if step[n] == s_)):
a_tile[(s_, key)] = t_
b_tile = {}
for s_ in range(ksteps):
for t_, key in enumerate(dict.fromkeys(tuple(m[4]) for n, m in enumerate(mm) if step[n] == s_)):
b_tile[(s_, key)] = t_
mt = max(t for (s_, _), t in a_tile.items()) + 1
nt = max(t for (s_, _), t in b_tile.items()) + 1
for w in range(waves):
lanes = np.arange(32) + 32 * w
A = np.zeros((16 * mt, 32 * ksteps), np.uint8)
B = np.zeros((32 * ksteps, 8 * nt), np.uint8)
C = np.zeros((16 * mt, 8 * nt), np.uint16)
D = np.zeros((16 * mt, 8 * nt), np.uint16)
for n, (l0, l1, Dr, Ar, Br, Cr) in enumerate(mm):
s_ = step[n]
i, j = a_tile[(s_, tuple(Ar))], b_tile[(s_, tuple(Br))]
for lane in range(32):
g, t = lane >> 2, lane & 3
av = [int(val[(p, l0, 'A', r)][lanes[lane]]) for r in Ar]
bv = [int(val[(p, l0, 'B', r)][lanes[lane]]) for r in Br]
for q in range(4): # A regs: (g, k 4t..), (g+8, 4t..), (g, 16+4t..), (g+8, 16+4t..)
row = 16 * i + g + 8 * (q & 1)
col = 32 * s_ + 4 * t + 16 * (q >> 1)
for byte in range(4):
A[row, col + byte] = (av[q] >> (8 * byte)) & 0xff
for q in range(2): # B regs: column g, k 4t.. and 16+4t..
for byte in range(4):
B[32 * s_ + 4 * t + 16 * q + byte, 8 * j + g] = (bv[q] >> (8 * byte)) & 0xff
for q in range(2): # C/D regs: row g (+8), columns 2t, 2t+1
row = 16 * i + g + 8 * q
if s_ == 0:
cv = int(val[(p, l0, 'C', Cr[q])][lanes[lane]])
C[row, 8 * j + 2 * t] = cv & 0xffff
C[row, 8 * j + 2 * t + 1] = cv >> 16
if s_ == ksteps - 1:
dv = int(val[(p, l0, 'D', Dr[q])][lanes[lane]])
D[row, 8 * j + 2 * t] = dv & 0xffff
D[row, 8 * j + 2 * t + 1] = dv >> 16
for name, mat in (('A', A), ('B', B), ('C', C), ('D', D)):
result[f'p{p}_w{w}_{name}'] = mat
print(f'phase {p}: {len(mm)} mma, {mt} m-tiles x {nt} n-tiles x {ksteps} k-steps -> '
f'{16 * mt}x{8 * nt}x{32 * ksteps}, lines {mm[0][0]}-{mm[-1][1]}')
np.savez_compressed(o.out, **result)
if __name__ == '__main__':
main()
engine/rr/ptx2rr.py
Executable
+729
View File
@@ -0,0 +1,729 @@
#!/usr/bin/env python3
"""Generate a GLSL compute shader for one Ray Reconstruction (DLSS-D) PTX kernel.
python3 engine/rr/ptx2rr.py MODULE.ptx KERNEL OUT.comp [--block X,Y,Z]
The translation core is vulkan/ptx2glsl.py. This module subclasses its translator and
implements the stage contract of engine/rr/README.md: binding 0 = parameter block,
binding 1 = weight buffer (argument 40), binding 2 = plane arena, then textures and
surfaces in ascending argument offset. One PTX block is one workgroup of X*Y*Z invocations.
The generated shader is derived from the user's own DLL and is model data: write it to the
model/build directory only; do not commit or redistribute it.
"""
import argparse
import importlib.util
import json
import re
import struct
import sys
from pathlib import Path
import numpy as np
ROOT = Path(__file__).resolve().parents[2]
_spec = importlib.util.spec_from_file_location('ptx2glsl', ROOT / 'vulkan' / 'ptx2glsl.py')
base = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(base)
# Workgroup shapes of the 15 Swin/HKPN kernels, from the captured launch manifests.
BLOCKS = {
'enc0': (32, 1, 1), 'enc1': (32, 1, 1), 'enc2': (32, 1, 4), 'enc3': (32, 1, 4), 'enc4': (32, 1, 8),
'enc5': (32, 1, 8), 'dec0': (32, 1, 1), 'dec1': (32, 1, 2), 'dec2': (32, 1, 4), 'dec3': (32, 1, 4),
'dec4': (32, 1, 8), 'hkpn': (16, 16, 1),
}
def default_block(kernel):
m = re.search(r'swin_(enc|dec)(\d)', kernel)
key = f'{m.group(1)}{m.group(2)}' if m else ('hkpn' if 'hkpn_output' in kernel else None)
if key not in BLOCKS:
raise SystemExit(f'no default block shape for {kernel}; pass --block X,Y,Z')
return BLOCKS[key]
def braces(tok):
s = tok.strip()
if not (s.startswith('{') and s.endswith('}')):
raise base.Err(f'expected a {{...}} vector operand, got {tok}')
return [x.strip() for x in s[1:-1].split(',')]
def mem_shape(o):
"""(bytes per element, element count) of ld/st.global modifiers o."""
count = 2 if 'v2' in o else 4 if 'v4' in o else 1
size = int(re.sub(r'\D', '', o[-1])) // 8
if size not in (2, 4):
raise base.Err(f'unsupported global access width in {".".join(o)}')
return size, count
RR_RUNTIME = r'''
float16_t H(uint u) { return uint16BitsToFloat16(uint16_t(u)); }
uint HU(float16_t h) { return uint(float16BitsToUint16(h)); }
f16vec2 H2(uint u) { return uint16BitsToFloat16(u16vec2(uint16_t(u & 0xffffu), uint16_t(u >> 16))); }
uint HU2_(u16vec2 u) { return uint(u.x) | (uint(u.y) << 16); }
uint H2U(f16vec2 h) { return HU2_(float16BitsToUint16(h)); }
int s16(uint u) { return (int(u) << 16) >> 16; }
float rcp_a(float x) { return 1.0 / x; }
// x + copysign(0.5, x) rounded toward zero and then truncated: round half away from zero, exactly
float round_half_away(float x) { float t = trunc(x); return abs(x - t) >= 0.5 ? t + sign(x) : t; }
uint PARAM32(uint o) { return param.w[o >> 2]; }
uint64_t PARAM64(uint o) { return uint64_t(param.w[o >> 2]) | (uint64_t(param.w[(o >> 2) + 1u]) << 32); }
uint PARAM16(uint o) { return (param.w[o >> 2] >> ((o & 2u) * 8u)) & 0xffffu; }
uint PARAM8(uint o) { return (param.w[o >> 2] >> ((o & 3u) * 8u)) & 0xffu; }
uint SH32(uint a) { return sh[a >> 2]; }
uint SH16(uint a) { return (sh[a >> 2] >> ((a & 2u) * 8u)) & 0xffffu; }
uint bfi32(uint ins, uint base, uint pos, uint len)
{
pos &= 0xffu; len &= 0xffu;
if (pos >= 32u) return base;
uint mask = len >= 32u ? (0xffffffffu << pos) : (((1u << len) - 1u) << pos);
return (~mask & base) | (mask & (ins << pos));
}
// Global memory: weight buffer (binding 1, read-only) and plane arena (binding 2). Addresses are byte offsets.
#ifdef OOB_REPORT
// debugging (ptx2rr.py --oob-report): per buffer (0 weights read, 1 arena read, 2 arena write) a count and the lowest
// and highest out-of-range byte address, in dbg.d[3k .. 3k+2]
void OOB(uint k, uint a, uint size) { if (a + 2u > size) { atomicAdd(dbg.d[3u * k], 1u); atomicMin(dbg.d[3u * k + 1u], a); atomicMax(dbg.d[3u * k + 2u], a); } }
#define CHECKW(a) OOB(0u, a, uint(wbuf.wh.length()) * 2u)
#define CHECKR(a) OOB(1u, a, uint(abuf.ah.length()) * 2u)
#define CHECKS(a) OOB(2u, a, uint(abuf.ah.length()) * 2u)
#else
#define CHECKW(a)
#define CHECKR(a)
#define CHECKS(a)
#endif
uint GW16(uint a) { CHECKW(a); return uint(wbuf.wh[a >> 1]); }
uint GW32(uint a) { CHECKW(a + 2u); return uint(wbuf.wh[a >> 1]) | (uint(wbuf.wh[(a >> 1) + 1u]) << 16); }
uint GA16(uint a) { CHECKR(a); return uint(abuf.ah[a >> 1]); }
uint GA32(uint a) { CHECKR(a + 2u); return uint(abuf.ah[a >> 1]) | (uint(abuf.ah[(a >> 1) + 1u]) << 16); }
void SA16(uint a, uint v) { CHECKS(a); abuf.ah[a >> 1] = uint16_t(v & 0xffffu); }
void SA32(uint a, uint v) { CHECKS(a + 2u); abuf.ah[a >> 1] = uint16_t(v & 0xffffu); abuf.ah[(a >> 1) + 1u] = uint16_t(v >> 16); }
// Surface stores with .zero out-of-bounds semantics.
void sust_p4(writeonly image2D s, int x, int y, vec4 v)
{
ivec2 n = imageSize(s);
if (uint(x) < uint(n.x) && uint(y) < uint(n.y)) imageStore(s, ivec2(x, y), v);
}
void sust_p1(writeonly image2D s, int x, int y, float v) { sust_p4(s, x, y, vec4(v, 0.0, 0.0, 0.0)); }
// Binary stores write the halves' bits. RADV folds float(float16_t(x)) back to x when the half came from a
// cvt.rn.f16.f32 in a register, and the image store then truncates x instead (the output stage's history surface
// was one binary16 ulp low in half its values, which the next frame's encoder 0 amplifies): OR-ing in a zero the
// compiler cannot see (the dispatch's z count is 1) keeps the conversion's rounding.
#define SUST_Z (gl_NumWorkGroups.z - 1u)
void sust_b2(writeonly image2D s, int x, int y, uint a, uint b)
{
sust_p4(s, x >> 2, y, vec4(float(H(a | SUST_Z)), float(H(b | SUST_Z)), 0.0, 0.0));
}
void sust_b4(writeonly image2D s, int x, int y, uint a, uint b, uint c, uint d)
{
sust_p4(s, x >> 3, y, vec4(float(H(a | SUST_Z)), float(H(b | SUST_Z)), float(H(c | SUST_Z)), float(H(d | SUST_Z))));
}
// shfl.sync.up/down within a 32-lane warp (ZLUDA's SHFL_SYNC_IMPL). in_bounds is the predicate output.
uint shfl_dir(uint v, uint delta, uint opts, bool down, out bool in_bounds)
{
int section_mask = int((opts >> 8) & 31u), warp_end = int(opts & 31u);
int self = int(gl_SubgroupInvocationID & 31u);
int d = int(delta & 31u);
int idx = down ? self + d : self - d;
int sub_end = (section_mask & self) | (~section_mask & (down ? warp_end : 0));
bool oob = down ? idx > sub_end : idx < sub_end;
if (oob) idx = self;
in_bounds = !oob;
return subgroupShuffle(v, uint(idx) & 31u);
}
// shfl.sync.idx (PTX ISA): j = minLane | (b & ~segmask), in bounds while j <= maxLane, otherwise the lane's own value
uint shfl_idx(uint v, uint b, uint opts, out bool in_bounds)
{
uint mask = (opts >> 8) & 31u, lane = gl_SubgroupInvocationID & 31u;
uint maxLane = (lane & mask) | (opts & 31u & ~mask), j = (lane & mask) | (b & 31u & ~mask);
in_bounds = j <= maxLane;
return subgroupShuffle(v, in_bounds ? j : lane);
}
// movmatrix.sync.trans.aligned.m8n8.b16 on one b32 fragment register per lane: lane (g = lane>>2, t = lane&3)
// holds row g, columns 2t and 2t+1. The transpose gathers (row 2t, col g) and (row 2t+1, col g).
uint movm_trans(uint src)
{
// through the subgroup's mma scratch rather than two shuffles: RADV miscompiled the shuffle form in the large
// stages (enc1's movmatrix results were wrong unless a store followed them; ZLUDA replay, 2026-10-11)
const uint base = gl_SubgroupID * 384u + 256u;
uint lane = gl_SubgroupInvocationID & 31u;
uint g = lane >> 2u, t = lane & 3u;
// order this operation's writes after every lane's reads of the previous one's results (the same scratch)
memoryBarrierShared(); subgroupBarrier();
mms[base + lane] = src;
memoryBarrierShared(); subgroupBarrier();
uint w1 = mms[base + 8u * t + (g >> 1u)];
uint w2 = mms[base + 8u * t + 4u + (g >> 1u)];
uint s = 16u * (g & 1u);
return ((w1 >> s) & 0xffffu) | (((w2 >> s) & 0xffffu) << 16);
}
// PTX min/max: a NaN operand yields the other one (GLSL leaves NaN undefined, so the result depended on codegen)
float16_t pmin(float16_t a, float16_t b) { return isnan(a) ? b : isnan(b) ? a : min(a, b); }
float16_t pmax(float16_t a, float16_t b) { return isnan(a) ? b : isnan(b) ? a : max(a, b); }
f16vec2 pmin2(f16vec2 a, f16vec2 b) { return f16vec2(pmin(a.x, b.x), pmin(a.y, b.y)); }
f16vec2 pmax2(f16vec2 a, f16vec2 b) { return f16vec2(pmax(a.x, b.x), pmax(a.y, b.y)); }
float pminf(float a, float b) { return isnan(a) ? b : isnan(b) ? a : min(a, b); }
float pmaxf(float a, float b) { return isnan(a) ? b : isnan(b) ? a : max(a, b); }
// e4m3 <-> f16 (PTX cvt semantics; RNE, satfinite; NaN -> 0x7f e4m3 / 0x7e00 f16).
uint e4m3_to_h(uint b)
{
uint e = (b >> 3) & 15u, m = b & 7u, s = (b & 0x80u) << 8;
if (e == 15u && m == 7u) return s | 0x7e00u;
float f = e == 0u ? ldexp(float(m), -9) : ldexp(float(8u + m), int(max(e, 1u)) - 10);
return s | HU(float16_t(f));
}
#ifdef SOFTWARE_E4M3
f16vec2 e4m3x2_to_f16x2(uint v) { return H2(e4m3_to_h(v & 0xffu) | (e4m3_to_h((v >> 8) & 0xffu) << 16)); }
#else
// the hardware conversion: exact for every non-NaN code (2026-10-12); NaN keeps e4m3_to_h's sign | 0x7e00
uint e4m3_to_h_hw(uint b)
{
if ((b & 0x7fu) == 0x7fu) return ((b & 0x80u) << 8) | 0x7e00u;
return uint(float16BitsToUint16(float16_t(uintBitsToFloate4m3EXT(uint8_t(b)))));
}
f16vec2 e4m3x2_to_f16x2(uint v) { return H2(e4m3_to_h_hw(v & 0xffu) | (e4m3_to_h_hw((v >> 8) & 0xffu) << 16)); }
#endif
uint h2e4m3(float f, bool relu)
{
if (isnan(f)) return 0x7fu;
if (relu && !(f > 0.0)) return 0u;
uint s = f < 0.0 ? 0x80u : 0u;
float a = abs(f);
if (a >= 448.0) return s | 0x7eu;
if (a == 0.0) return s;
int ex;
frexp(a, ex);
int qe = max(ex - 4, -9);
uint n = uint(roundEven(ldexp(a, -qe)));
uint bits;
if (n < 8u) bits = n;
else { int p = findMSB(n); bits = (uint(qe + p + 7) << 3) | ((n >> uint(p - 3)) & 7u); }
return s | min(bits, 0x7eu);
}
#ifdef SOFTWARE_E4M3
uint f16x2_to_e4m3x2(f16vec2 h, bool relu) { return h2e4m3(float(h.x), relu) | (h2e4m3(float(h.y), relu) << 8); }
#else
// cvt.rn.satfinite(.relu).e4m3x2.f16x2 on the hardware conversion: round to nearest even after clamping to +-448
// (satfinite), NaN to 0x7f. Equal to the software h2e4m3 for all 65536 f16 inputs (2026-10-12) and about two thirds
// of a stage's time faster (dec0: 2.85 -> 1.03 ms).
uint h2e4m3_hw(float16_t h, bool relu)
{
if (isnan(h)) return 0x7fu;
if (relu) h = h > float16_t(0.0) ? h : float16_t(0.0);
return uint(floate4m3BitsToUintEXT(floate4m3_t(clamp(h, float16_t(-448.0), float16_t(448.0)))));
}
uint f16x2_to_e4m3x2(f16vec2 h, bool relu) { return h2e4m3_hw(h.x, relu) | (h2e4m3_hw(h.y, relu) << 8); }
#endif
// mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 on the FP8 cooperative matrices, with ZLUDA's lowering of
// it (engine/layer_m.comp's TEXL): C enters the first of two 16x16x16 products, the sum stays f32 and is rounded to
// binary16 once. The operands pass through the subgroup's own scratch (MMS): NVIDIA's fragments are written in
// place (A row-major 16x32, B column-major 32x8 in a 16-column tile, C/D row-major f16 16x16), then loaded as
// whole matrices. Lane l = 4g + t holds A rows g and g+8 (k 4t.. and 16+4t..), B column g (same k), C/D row g
// and g+8, columns 2t and 2t+1.
#define MMA coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define MMB coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#define MMF coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define MMH coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
void mma_e4m3_f16(uint a0, uint a1, uint a2, uint a3, uint b0, uint b1, uint c0, uint c1, out uint d0, out uint d1)
{
const uint base = gl_SubgroupID * 384u, A = base, B = base + 128u, C = base + 256u;
const uint lane = gl_SubgroupInvocationID & 31u, g = lane >> 2u, t = lane & 3u;
// order this operation's writes after every lane's reads of the previous one's results (the same scratch)
memoryBarrierShared(); subgroupBarrier();
mms[A + g * 8u + t] = a0; mms[A + (g + 8u) * 8u + t] = a1;
mms[A + g * 8u + 4u + t] = a2; mms[A + (g + 8u) * 8u + 4u + t] = a3;
mms[B + g * 8u + t] = b0; mms[B + g * 8u + 4u + t] = b1;
mms[B + (g + 8u) * 8u + t] = 0u; mms[B + (g + 8u) * 8u + 4u + t] = 0u;
mms[C + g * 8u + t] = c0; mms[C + (g + 8u) * 8u + t] = c1;
mms[C + g * 8u + 4u + t] = 0u; mms[C + (g + 8u) * 8u + 4u + t] = 0u;
memoryBarrierShared(); subgroupBarrier();
MMA x0, x1; MMB y0, y1; MMH c;
coopMatLoad(x0, mms, A, 8u, gl_CooperativeMatrixLayoutRowMajor);
coopMatLoad(x1, mms, A + 4u, 8u, gl_CooperativeMatrixLayoutRowMajor);
coopMatLoad(y0, mms, B, 8u, gl_CooperativeMatrixLayoutColumnMajor);
coopMatLoad(y1, mms, B + 4u, 8u, gl_CooperativeMatrixLayoutColumnMajor);
coopMatLoad(c, mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor);
MMF acc = coopMatMulAdd(x1, y1, coopMatMulAdd(x0, y0, MMF(c)));
memoryBarrierShared(); subgroupBarrier();
coopMatStore(MMH(acc), mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor);
memoryBarrierShared(); subgroupBarrier();
d0 = mms[C + g * 8u + t]; d1 = mms[C + (g + 8u) * 8u + t];
}
// Two m16n8k32 products with the same A fragment (NVIDIA issues every one of them in such adjacent pairs): one 16x16
// product, the first's B, C and D in columns 0-7, the second's in columns 8-15.
void mma_e4m3_f16_pair(uint a0, uint a1, uint a2, uint a3, uint b0, uint b1, uint c0, uint c1,
uint e0, uint e1, uint f0, uint f1, out uint d0, out uint d1, out uint h0, out uint h1)
{
const uint base = gl_SubgroupID * 384u, A = base, B = base + 128u, C = base + 256u;
const uint lane = gl_SubgroupInvocationID & 31u, g = lane >> 2u, t = lane & 3u;
memoryBarrierShared(); subgroupBarrier();
mms[A + g * 8u + t] = a0; mms[A + (g + 8u) * 8u + t] = a1;
mms[A + g * 8u + 4u + t] = a2; mms[A + (g + 8u) * 8u + 4u + t] = a3;
mms[B + g * 8u + t] = b0; mms[B + g * 8u + 4u + t] = b1;
mms[B + (g + 8u) * 8u + t] = e0; mms[B + (g + 8u) * 8u + 4u + t] = e1;
mms[C + g * 8u + t] = c0; mms[C + (g + 8u) * 8u + t] = c1;
mms[C + g * 8u + 4u + t] = f0; mms[C + (g + 8u) * 8u + 4u + t] = f1;
memoryBarrierShared(); subgroupBarrier();
MMA x0, x1; MMB y0, y1; MMH c;
coopMatLoad(x0, mms, A, 8u, gl_CooperativeMatrixLayoutRowMajor);
coopMatLoad(x1, mms, A + 4u, 8u, gl_CooperativeMatrixLayoutRowMajor);
coopMatLoad(y0, mms, B, 8u, gl_CooperativeMatrixLayoutColumnMajor);
coopMatLoad(y1, mms, B + 4u, 8u, gl_CooperativeMatrixLayoutColumnMajor);
coopMatLoad(c, mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor);
MMF acc = coopMatMulAdd(x1, y1, coopMatMulAdd(x0, y0, MMF(c)));
memoryBarrierShared(); subgroupBarrier();
coopMatStore(MMH(acc), mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor);
memoryBarrierShared(); subgroupBarrier();
d0 = mms[C + g * 8u + t]; d1 = mms[C + (g + 8u) * 8u + t];
h0 = mms[C + g * 8u + 4u + t]; h1 = mms[C + (g + 8u) * 8u + 4u + t];
}
'''
def statement_lines(ptx, kernel):
"""the source line of every statement of KERNEL, in the order vulkan/ptx2glsl.py's tokenizer yields them"""
text = re.sub(r'//[^\n]*', '', ptx.replace('\r', ''))
m = re.search(r'\.entry\s+' + re.escape(kernel) + r'\s*\((.*?)\)', text, re.S)
start = text.index('{', m.end())
depth, k = 0, start
while True:
if text[k] == '{':
depth += 1
elif text[k] == '}':
depth -= 1
if depth == 0:
break
k += 1
lines, i, line = [], start + 1, text.count('\n', 0, start + 1) + 1
while i < k:
c = text[i]
if c.isspace():
line += c == '\n'
i += 1
continue
if c in '{}':
i += 1
continue
lm = re.match(r'(\$?[A-Za-z_][\w$]*):', text[i:i + 200])
if lm:
i += lm.end()
continue
j = text.index(';', i)
lines.append(line)
line += text.count('\n', i, j)
i = j + 1
return lines
def base_helper(name):
"""one helper function of vulkan/ptx2glsl.py's runtime, with its leading comment"""
lines = base.RUNTIME.split('\n')
start = next(i for i, l in enumerate(lines) if re.match(r'\w+ ' + name + r'\(', l))
while lines[start - 1].startswith('//'):
start -= 1
end = start
if not lines[start].rstrip().endswith('}'): # one-line helpers end on their own line
while lines[end] != '}':
end += 1
return '\n'.join(lines[start:end + 1]) + '\n'
RR_RUNTIME += base_helper('prmt') + base_helper('shfl_bfly').replace(
'return subgroupShuffle(v, (gl_SubgroupInvocationID & 32u) | uint(idx));', 'return subgroupShuffle(v, uint(idx) & 31u);')
class RRT(base.T):
def __init__(self, ptx, kernel, block):
bx, by, bz = block
if bx * by * bz % 32:
raise base.Err('block size must be a multiple of the 32-lane subgroup')
self.bx, self.by, self.bz = bx, by, bz
self.sreg = {
'%tid.x': f'(gl_LocalInvocationIndex % {bx}u)',
'%tid.y': f'((gl_LocalInvocationIndex / {bx}u) % {by}u)',
'%tid.z': f'(gl_LocalInvocationIndex / {bx * by}u)',
'%ntid.x': f'{bx}u', '%ntid.y': f'{by}u', '%ntid.z': f'{bz}u',
'%ctaid.x': 'gl_WorkGroupID.x', '%ctaid.y': 'gl_WorkGroupID.y', '%ctaid.z': 'gl_WorkGroupID.z',
'%nctaid.x': 'gl_NumWorkGroups.x', '%nctaid.y': 'gl_NumWorkGroups.y', '%nctaid.z': 'gl_NumWorkGroups.z',
'%laneid': 'gl_SubgroupInvocationID',
'WARP_SZ': '32u',
}
self.origin = {}
self.f64 = {} # %fd registers: the kernels only load f64 immediates into them, to round them to f16
self.probes = {} # debugging: PTX line -> registers stored to the debug buffer after it executes
self.probe_block = (0, 0)
self.probe_slots = 0
self.oob = False
# every register precise, the integer ones too: they carry f16x2 and f32 bit patterns between operations, and
# without it the compiler may contract a product and a sum across them depending on scheduling (results then
# changed with unrelated code and with the device's robustness features)
self.precise_all = True
self.stmt_lines = statement_lines(ptx, kernel)
self.stmt_index = 0
super().__init__(ptx, kernel)
self.pair_mma()
self.origins()
def stmt(self, s):
line = self.stmt_lines[self.stmt_index] if self.stmt_index < len(self.stmt_lines) else None
self.stmt_index += 1
super().stmt(s)
for reg in self.probes.get(line, ()):
# the debug buffer: 32-bit words [slot][local invocation], the last execution of the statement
n = self.reg(reg)
g = self.regs[n]
v = (f'floatBitsToUint({n})' if g == 'float' else f'uint({n})' if g in ('uint', 'uint64_t', 'bool') else
f'HU({n})' if g == 'float16_t' else f'H2U({n})')
bx, by = self.probe_block
self.E(f'if (gl_WorkGroupID.xy == uvec2({bx}u, {by}u)) dbg.d[{self.probe_slots}u * (gl_WorkGroupSize.x * gl_WorkGroupSize.y) + '
f'gl_LocalInvocationIndex] = {v};')
self.probe_slots += 1
def pair_mma(self):
"""merge adjacent mma statements with the same A fragment and independent operands into one paired product;
the second statement becomes a no-op so statement numbering (probes) is unchanged"""
stmts = [k for k, (kind, _) in enumerate(self.toks) if kind == 'stmt']
n = 0
while n + 1 < len(stmts):
i, j = stmts[n], stmts[n + 1]
x, y = self.toks[i][1], self.toks[j][1]
if x.startswith('mma.sync.aligned.m16n8k32') and y.startswith('mma.sync.aligned.m16n8k32'):
gx = [braces(t) for t in base.split_operands(x.split(None, 1)[1])]
gy = [braces(t) for t in base.split_operands(y.split(None, 1)[1])]
if gx[1] == gy[1] and not set(gx[0]) & set(gy[2] + gy[3] + gy[1]):
ops = ', '.join('{' + ', '.join(g) + '}' for g in (gx[0], gy[0], gx[1], gx[2], gx[3], gy[2], gy[3]))
self.toks[i] = ('stmt', 'mma.pair.e4m3 ' + ops)
self.toks[j] = ('stmt', 'mma.paired.noop')
n += 2
continue
n += 1
# ------------------------------------------------ SREG remap onto the flat workgroup
def is_reg(self, tok):
return tok.strip() not in self.sreg and super().is_reg(tok)
def val(self, tok, t):
tok = tok.strip()
if re.fullmatch(r'(0[xX][0-9a-fA-F]+|-?\d+)[uU]', tok):
tok = tok[:-1]
tok = re.sub(r'^0X', '0x', tok)
if tok in self.sreg:
e = self.sreg[tok]
return f'int({e})' if t == 's32' else f'uint64_t({e})' if t in ('u64', 'b64') else f'uint({e})'
return super().val(tok, t)
# ------------------------------------------------ global memory origin: weight buffer or arena
def origins(self):
stmts = []
for kind, s in self.toks:
if kind != 'stmt' or s.startswith('.'):
continue
s = re.sub(r'^@!?%p\d+\s+', '', s)
parts = s.split(None, 1)
stmts.append((parts[0], base.split_operands(parts[1]) if len(parts) > 1 else []))
org = self.origin
changed = True
while changed:
changed = False
for op, a in stmts:
if not a or not re.fullmatch(r'%rd\d+', a[0]) or op.startswith(('ld.', 'st.', 'bra')):
if op != 'ld.param.u64':
continue
if op == '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'}
else:
new = set()
for x in a[1:]:
if re.fullmatch(r'%rd\d+', x):
new |= org.get(x, set())
cur = org.setdefault(a[0], set())
if new - cur:
cur |= new
changed = True
def gorigin(self, reg):
s = self.origin.get(reg, set())
if len(s) != 1:
raise base.Err(f'global pointer {reg} has origin {sorted(s)} (need exactly one of w/a)')
return next(iter(s))
def gaddr(self, tok):
m = re.fullmatch(r'\[\s*(%rd\d+)\s*(?:\+\s*(\d+))?\s*\]', tok.strip())
if not m:
raise base.Err(f'unsupported global address {tok}')
return m.group(1), int(m.group(2) or 0)
def gbyte(self, reg, off):
return f'uint({self.val(reg, "u64")}) + {off}u'
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'
dsts = braces(a[0]) if count > 1 else [a[0]]
for i, d in enumerate(dsts):
self.A(d, f'{fn}{size * 8}({self.gbyte(reg, off + i * size)})', 'u16' if size == 2 else 'u32')
def st(self, o, a):
if 'global' not in o:
return super().st(o, a)
size, count = mem_shape(o)
reg, off = self.gaddr(a[0])
if self.gorigin(reg) != 'a':
raise base.Err('store through the read-only weight pointer')
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")});')
else:
self.E(f'SA32({self.gbyte(reg, off + i * size)}, {self.val(v, "u32")});')
# ------------------------------------------------ textures and surfaces
def handle_off(self, h):
off = self.handle.get(h)
if off is None:
raise base.Err(f'{h} is not a texture or surface handle loaded from the parameter block')
return off
def set_surface(self, off, fmt):
old = self.surfaces.get(off)
if old is not None and old != fmt:
raise base.Err(f'surface {off} used with formats {old} and {fmt}')
self.surfaces[off] = fmt
def tex(self, op, a):
o = op.split('.')[1:]
dst = braces(a[0])
h, c = self.texref(a[1])
off = self.handle_off(h)
self.textures.add(off)
if 'base' in o and o[-1] == 's32':
# integer coordinates still go through the texture's address mode in CUDA (clamp for every denoiser
# texture); texelFetch does not clamp, and what it returns out of range depends on the device's robustness
fetch = (f'texelFetch(T{off}, clamp(ivec2({self.val(c[0], "s32")}, {self.val(c[1], "s32")}), ivec2(0), '
f'textureSize(T{off}, 0) - 1), 0)')
elif 'base' in o and o[-1] == 'f32':
fetch = f'textureLod(T{off}, vec2({self.val(c[0], "f32")}, {self.val(c[1], "f32")}), 0.0)'
elif 'level' in o and o[-1] == 'f32':
fetch = (f'textureLod(T{off}, vec2({self.val(c[0], "f32")}, {self.val(c[1], "f32")}), '
f'{self.val(a[2], "f32")})')
else:
raise base.Err(f'unsupported texture op {op}')
f16 = 'f16' in o
self.E(f'{{ vec4 t_ = {fetch};')
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')
self.E('}')
def sust(self, op, a):
h, c = self.texref(a[0])
off = self.handle_off(h)
v = braces(a[1])
x, y = self.val(c[0], 's32'), self.val(c[1], 's32')
S = f'S{off}'
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)}));')
elif op == 'sust.p.2d.b32.zero':
self.set_surface(off, 'r16f')
self.E(f'sust_p1({S}, {x}, {y}, {self.val(v[0], "f32")});')
elif op == 'sust.b.2d.v2.b16.zero':
self.set_surface(off, 'rg16f')
self.E(f'sust_b2({S}, {x}, {y}, {", ".join(self.val(q, "u16") for q in v)});')
elif op == 'sust.b.2d.v4.b16.zero':
self.set_surface(off, 'rgba16f')
self.E(f'sust_b4({S}, {x}, {y}, {", ".join(self.val(q, "u16") for q in v)});')
else:
raise base.Err(f'unsupported surface op {op}')
# ------------------------------------------------ warp-level ops
def mma(self, a):
d, x, b, c = (braces(t) for t in a)
if not (len(d) == 2 and len(x) == 4 and len(b) == 2 and len(c) == 2):
raise base.Err('mma operand shape')
args = ', '.join(self.val(r, 'u32') for r in x + b + c)
self.E('{ uint md0_, md1_;')
self.E(f'mma_e4m3_f16({args}, md0_, md1_);')
self.A(d[0], 'md0_', 'u32')
self.A(d[1], 'md1_', 'u32')
self.E('}')
def shfl_dir(self, op, a):
down = op.split('.')[2] == 'down'
d = a[0].split('|')
self.E('{ bool ib_;')
self.A(d[0], f'shfl_dir({self.val(a[1], "u32")}, {self.val(a[2], "u32")}, {self.val(a[3], "u32")}, '
f'{"true" if down else "false"}, ib_)', 'u32')
if len(d) > 1:
self.A(d[1], 'ib_', 'bool')
self.E('}')
def cvt_to_e4m3(self, op, a):
relu = 'relu' in op.split('.')
self.A(a[0], f'f16x2_to_e4m3x2({self.val(a[1], "f16x2")}, {"true" if relu else "false"})', 'u16')
def cvt(self, op, o, a):
types = [t for t in o if re.match(r'^(f|s|u|b)\d+(x2)?$', t)]
if types[:2] in (['u32', 'u64'], ['s32', 's64'], ['u32', 's64'], ['s32', 'u64']):
return self.A(a[0], f'uint({self.val(a[1], "u64")})', 'u32')
return super().cvt(op, o, a)
def f64_op(self, op, a):
if op == 'mov.f64' and re.fullmatch(r'0[dD][0-9A-Fa-f]{16}', a[1].strip()):
self.f64[a[0].strip()] = struct.unpack('<d', int(a[1].strip()[2:], 16).to_bytes(8, 'little'))[0]
return True
if op == 'cvt.rn.f16.f64':
v = self.f64.get(a[1].strip())
if v is None:
raise base.Err(f'{a[1]} is not a known f64 constant')
bits = int(np.float16(v).view(np.uint16))
self.A(a[0], f'{bits}u', 'u16')
return True
return False
def op(self, op, a):
parts = op.split('.')
if '%fd' in ','.join(a) and self.f64_op(op, a):
return
if op.startswith('mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16'):
return self.mma(a)
if op == 'mma.pair.e4m3':
d, h, x, b, c, e, f = (braces(t) for t in a)
args = ', '.join(self.val(r, 'u32') for r in x + b + c + e + f)
self.E('{ uint md0_, md1_, mh0_, mh1_;')
self.E(f'mma_e4m3_f16_pair({args}, md0_, md1_, mh0_, mh1_);')
for reg, v in zip(d + h, ('md0_', 'md1_', 'mh0_', 'mh1_')):
self.A(reg, v, 'u32')
return self.E('}')
if op == 'mma.paired.noop':
return
if op == 'movmatrix.sync.trans.aligned.m8n8.b16':
return self.A(a[0], f'movm_trans({self.val(a[1], "u32")})', 'u32')
# PTX clamps shift amounts at the register width (shl/shr.u give 0, shr.s the sign); GLSL leaves them undefined
if parts[0] in ('shl', 'shr') and len(parts) == 2 and parts[1] in ('b32', 'u32', 's32'):
x, n = self.val(a[1], 's32' if parts[1] == 's32' else 'u32'), self.val(a[2], 'u32')
if parts[0] == 'shl':
return self.A(a[0], f'(({n}) >= 32u ? 0u : ({x}) << ({n}))', 'u32')
if parts[1] == 's32':
return self.A(a[0], f'uint(({x}) >> min({n}, 31u))', 'u32')
return self.A(a[0], f'(({n}) >= 32u ? 0u : ({x}) >> ({n}))', 'u32')
if op == 'shfl.sync.idx.b32':
d = a[0].split('|')
self.E('{ bool ib_;')
self.A(d[0], f'shfl_idx({self.val(a[1], "u32")}, {self.val(a[2], "u32")}, {self.val(a[3], "u32")}, ib_)', 'u32')
if len(d) > 1:
self.A(d[1], 'ib_', 'bool')
return self.E('}')
if op in ('shfl.sync.down.b32', 'shfl.sync.up.b32'):
return self.shfl_dir(op, a)
if parts[0] == 'sust':
return self.sust(op, a)
if parts[0] == 'tex':
return self.tex(op, a)
if parts[0] == 'cvt' and [t for t in parts[1:] if t in ('f16x2', 'e4m3x2')] == ['e4m3x2', 'f16x2']:
return self.cvt_to_e4m3(op, a)
if op == 'neg.f16':
return self.A(a[0], f'-{self.val(a[1], "f16")}', 'f16')
if parts[0] in ('min', 'max') and 'NaN' not in parts and any(t in parts for t in ('f16', 'f16x2', 'f32')):
t = 'f16x2' if 'f16x2' in parts else 'f16' if 'f16' in parts else 'f32'
fn = {'f16x2': 'p{}2', 'f16': 'p{}', 'f32': 'p{}f'}[t].format(parts[0])
x, y = self.val(a[1], t), self.val(a[2], t)
return self.A(a[0], f'{fn}({x}, {y})', 'float' if t == 'f32' else t)
if op == 'neg.f16x2':
return self.A(a[0], f'-{self.val(a[1], "f16x2")}', 'f16x2')
if op == 'abs.f16x2':
return self.A(a[0], f'abs({self.val(a[1], "f16x2")})', 'f16x2')
if parts[0] == 'set' and parts[2] == 'f16':
c = ('(isnan({0}) || isnan({1}))' if parts[1] == 'nan' else
'!(isnan({0}) || isnan({1}))' if parts[1] == 'num' else self.cmp(set(parts), True))
return self.A(a[0], f'(({c.format(self.val(a[1], "f16"), self.val(a[2], "f16"))}) ? 1.0hf : 0.0hf)', 'f16')
return super().op(op, a)
# ------------------------------------------------ stage contract
def preamble(self, w, tex_filter):
w('#version 460')
w('#extension GL_EXT_shader_explicit_arithmetic_types : require')
w('#extension GL_EXT_shader_16bit_storage : require')
w('#extension GL_KHR_shader_subgroup_basic : require')
w('#extension GL_KHR_shader_subgroup_shuffle : require')
w('#extension GL_KHR_cooperative_matrix : require')
w('#extension GL_KHR_memory_scope_semantics : require')
w('#extension GL_EXT_float_e4m3 : require')
w(f'// Generated by engine/rr/ptx2rr.py from {self.kernel}. Model data derived from the user\'s DLL; do not commit.')
w(f'layout(local_size_x = {self.bx * self.by * self.bz}, local_size_y = 1, local_size_z = 1) in;')
w('layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint w[]; } param;')
w('layout(set = 0, binding = 1, std430) readonly buffer WeightBuf { uint16_t wh[]; } wbuf;')
w('layout(set = 0, binding = 2, std430) buffer ArenaBuf { uint16_t ah[]; } abuf;')
n = 3
for off in sorted(self.textures):
w(f'layout(set = 0, binding = {n}) uniform sampler2D T{off};')
n += 1
for off in sorted(self.surfaces):
w(f'layout(set = 0, binding = {n}, {self.surfaces[off]}) uniform writeonly image2D S{off};')
n += 1
w(f'shared uint sh[{max(1, (self.shared_bytes + 3) // 4)}];')
if self.probe_slots or self.oob:
w(f'layout(set = 0, binding = {n}, std430) buffer Debug {{ uint d[]; }} dbg;')
if self.oob:
w('#define OOB_REPORT 1')
w(f'shared uint mms[{384 * (self.bx * self.by * self.bz // 32)}]; // mma_e4m3_f16 scratch, per subgroup')
def runtime(self):
return RR_RUNTIME
def main():
ap = argparse.ArgumentParser(description='PTX kernel -> GLSL compute shader (Ray Reconstruction stage contract)')
ap.add_argument('module')
ap.add_argument('kernel')
ap.add_argument('out')
ap.add_argument('--block', help='workgroup shape X,Y,Z (default: from the kernel table)')
ap.add_argument('--surface', action='append', default=[], metavar='OFFSET=FORMAT',
help="the GLSL image format of a surface whose stores are formatted (sust.p), when it is not the "
"store's own guess, e.g. 144=rgba8_snorm (repeatable)")
ap.add_argument('--probe', action='append', default=[], metavar='LINE:%REG,...',
help='debugging: store the registers after the statement at that module line executes (last '
'execution) for workgroup --probe-block, into a debug buffer bound after the surfaces')
ap.add_argument('--probe-block', default='0,0')
ap.add_argument('--oob-report', action='store_true', help='debugging: count out-of-range buffer accesses')
ns = ap.parse_args()
block = tuple(int(x) for x in ns.block.split(',')) if ns.block else default_block(ns.kernel)
formats = {int(k): v for k, v in (x.split('=') for x in ns.surface)}
src = open(ns.module).read()
probes = {}
for spec in ns.probe:
line, regs = spec.split(':', 1)
probes[int(line)] = [r.strip() for r in regs.split(',')]
def make():
t = RRT(src, ns.kernel, block)
t.probes = probes
t.probe_block = tuple(int(x) for x in ns.probe_block.split(','))
t.oob = ns.oob_report
return t
t = base.build(make, False, False, False)
for off, fmt in formats.items():
if off in t.surfaces:
t.surfaces[off] = fmt
open(ns.out, 'w').write(t.emit({}))
# 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,
'probes': [[line, reg] for line in sorted(probes) for reg in probes[line]], 'oob': ns.oob_report}) + '\n')
sys.stderr.write(f'block={block} textures={sorted(t.textures)} surfaces={t.surfaces} shared={t.shared_bytes}\n')
if __name__ == '__main__':
main()
engine/rr/ptx_expr.py
+108
View File
@@ -0,0 +1,108 @@
#!/usr/bin/env python3
"""Print the expression tree that defines a PTX register, followed back through arithmetic, pair packs/unpacks,
conversions and warp shuffles down to loads, constants and registers named as leaves.
usage: ptx_expr.py STAGE LINE %REG [--depth N] [--stop %r1,%r2]
LINE is a line at or after the register's definition (the search goes backwards from it). Brace-scoped temporaries
(`{ .reg .b16 hl, hu; ... }`) are resolved within their block. Shuffles print as `shfl.bfly^k(expr)`: the same
expression on the lane `laneid ^ k`.
"""
import argparse
import re
from pathlib import Path
def module_of(kernel):
for p in sorted(Path('~/.cache/d4r-rr-corpus').expanduser().glob('*.ptx')):
if f'.entry {kernel}(' in p.read_text(errors='replace'):
return p
class Tracer:
def __init__(self, lines, start):
self.lines = lines
self.start = start
self.consts = {}
self.cache = {}
def find_def(self, reg, before, block_lo):
"""the statement (line index) that last writes REG before line BEFORE"""
pat = re.compile(r'(^|[\s{,|])' + re.escape(reg) + r'([\s,|}]|$)')
lo = block_lo if not reg.startswith('%') else self.start
for i in range(before - 1, lo - 1, -1):
s = re.sub(r'//.*', '', self.lines[i]).strip().lstrip('{').strip()
s = re.sub(r'^@!?%p\d+\s+', '', s)
if not s or s.startswith(('.', '$', '}')):
continue
parts = s.split(None, 1)
if len(parts) < 2:
continue
rest = parts[1]
dsts = ([x.strip() for x in rest[1:rest.index('}')].split(',')] if rest.startswith('{')
else [x.strip() for x in rest.split(',')[0].split('|')])
if reg in dsts and not parts[0].startswith(('st.', 'sust', 'bra')):
return i, parts[0], rest.rstrip(';').strip(), dsts
return None
def block_start(self, line):
depth = 0
for i in range(line, self.start, -1):
depth += self.lines[i].count('}') - self.lines[i].count('{')
if depth < 0:
return i
return self.start
def expr(self, reg, before, depth, stop, half=None):
reg = reg.strip()
if re.fullmatch(r'-?\d+|0[fFdD][0-9A-Fa-f]+|0x[0-9a-fA-F]+', reg):
return reg
if reg in stop or depth <= 0:
return reg + (f'.{half}' if half is not None else '')
d = self.find_def(reg, before, self.block_start(before - 1))
if d is None:
return reg + '?'
i, op, rest, dsts = d
args = [a.strip() for a in re.split(r',(?![^{]*\})', rest)]
srcs = args[1:] if not rest.startswith('{') else args[1:]
e = lambda r, h=None: self.expr(r, i, depth - 1, stop, h)
if op.startswith('mov.b32') and rest.startswith('{'): # unpack: {lo, hi} <- r
which = 'lo' if dsts.index(reg) == 0 else 'hi'
return e(srcs[0], which)
if op.startswith('mov.b32') and srcs and srcs[0].startswith('{'): # pack r <- {lo, hi}
lo, hi = [x.strip() for x in srcs[0][1:-1].split(',')]
if half == 'lo':
return e(lo)
if half == 'hi':
return e(hi)
return f'pack({e(lo)}, {e(hi)})'
if op.startswith(('mov.', 'cvt.')) and len(srcs) == 1:
inner = e(srcs[0], half)
return inner if op.startswith('mov.') else f'{op.split(".")[0]}.{op.split(".")[-2]}({inner})'
if op.startswith('shfl.sync.bfly'):
return f'shfl^{srcs[1]}({e(srcs[0], half)})'
if op.startswith('ld.'):
return f'load@L{i + 1}[{srcs[0]}]' + (f'.{half}' if half else '')
name = op.split('.')[0]
if op.endswith('x2') and half:
return f'{name}(' + ', '.join(e(s, half) for s in srcs) + ')'
return f'{name}(' + ', '.join(e(s, half) for s in srcs) + ')'
def main():
ap = argparse.ArgumentParser()
ap.add_argument('stage')
ap.add_argument('line', type=int)
ap.add_argument('reg')
ap.add_argument('--depth', type=int, default=30)
ap.add_argument('--stop', default='')
o = ap.parse_args()
kernel = f'cuda_dldn_engine_swin_{o.stage}_kernel'
lines = module_of(kernel).read_text().replace('\r', '').split('\n')
k = next(i for i, l in enumerate(lines) if f'.entry {kernel}(' in l)
t = Tracer(lines, k)
print(t.expr(o.reg, o.line, o.depth, set(filter(None, o.stop.split(',')))))
if __name__ == '__main__':
main()
engine/rr/relations.py
+55
View File
@@ -0,0 +1,55 @@
#!/usr/bin/env python3
"""Which earlier result each phase operand of a dumped layer (engine/rr/phase_dump.py) comes from.
usage: relations.py DUMP.npz [WAVE]"""
import sys
import numpy as np
sys.path.insert(0, __import__('os').path.dirname(__file__))
from rrmath import f16, q8
z = np.load(sys.argv[1]); w = int(sys.argv[2]) if len(sys.argv) > 2 else 0
phases = sorted({int(k.split('_')[0][1:]) for k in z.files})
P = lambda p, n: z[f'p{p}_w{w}_{n}']
def inter(n):
"""the C/D -> A reinterpretation within each 32 block: column 4m+j <- 2m+j, 4m+2+j <- 16+2m+j"""
perm = []
for b in range(n // 32):
for m in range(8):
perm += [32 * b + 2 * m, 32 * b + 2 * m + 1, 32 * b + 16 + 2 * m, 32 * b + 16 + 2 * m + 1]
return perm
def rel(op, M, p):
found = []
if not M.any():
return ['zero']
for q in range(p - 1, 0, -1):
D = P(q, 'D')
if op == 'C':
for name, X in (('D', D),):
if X.shape == M.shape and np.array_equal(X, M):
found.append(f'{name}{q}')
continue
Q = q8(f16(D))
cands = []
if Q.shape[1] % 32 == 0:
cands += [('q8(D%d)[:, inter]' % q, Q[:, inter(Q.shape[1])])]
cands += [('q8(D%d)' % q, Q)]
for name, X in cands:
for tname, Y in ((name, X), (name + '.T', X.T)):
if Y.shape == M.shape and np.array_equal(Y, M):
found.append(tname)
elif Y.shape[1] == M.shape[1] and Y.shape[0] > M.shape[0] and np.array_equal(Y[:M.shape[0]], M):
found.append(tname + f'[:{M.shape[0]}]')
if op == 'B' and Y.shape[0] == M.shape[0] and Y.shape[1] > M.shape[1] and np.array_equal(Y[:, :M.shape[1]], M):
found.append(tname + f'[:, :{M.shape[1]}]')
if found:
break
return found or ['?']
for p in phases:
A, B, C = P(p, 'A'), P(p, 'B'), P(p, 'C')
print(f'phase {p:2d} {A.shape[0]}x{B.shape[1]}x{A.shape[1]}: A <- {rel("A", A, p)} B <- {rel("B", B, p)} '
f'C <- {rel("C", C, p)}')
engine/rr/rr_attn.glsl
+71
View File
@@ -0,0 +1,71 @@
// The attention core shared by the 32-channel-head stages (engine/rr/enc1_m.comp, dec0_m.comp), included inside the
// head loop: this wave's 16 query slots against the window's 96 key slots. Expects KS (K, token slot x 32 FP8
// channels in k order) and VS (V^T, 32 channels x 96 token slots) filled for head hd, the bias at W_BIAS0 + HEAD * hd
// (enc1_pack: 96 x 64, [key slot][query slot]) and the cubic table at W_CUBIC. Leaves Ob[2]: O^T (32 channels x
// this wave's queries) as the out-projection's B operand. RR_KS / RR_VS (+ word offsets RR_KS0 / RR_VS0) name
// the arrays when they are not KS and VS.
#ifndef RR_KS
#define RR_KS KS
#define RR_KS0 0u
#endif
#ifndef RR_VS
#define RR_VS VS
#define RR_VS0 0u
#endif
B8 Qb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) coopMatLoad(Qb[c], RR_KS, RR_KS0 + wv * 128u + 4u * c, 8u, CM);
// scores S^T (keys x this wave's queries) + bias, then the cubic softmax down each column
B8 Pb[6];
{
float16_t e[6][8];
[[unroll]] for (uint kt = 0u; kt < 6u; ++kt) {
HAC c;
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
// 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]))]));
}
// key slot 16kt + 8hl + i is NVIDIA's column 8j + 2tt + half with j = 4(kt >> 1) + (kt & 1) + 2((i >> 1) & 1),
// tt = 2hl + (i >> 2), half = i & 1 (enc1_ref.softmax_cubic's tree)
// packed: component h of a pair is column parity h (elements i = 4tq + h, i + 2), so the tree runs on f16x2
f16vec2 ep[6][4];
[[unroll]] for (uint kt = 0u; kt < 6u; ++kt)
[[unroll]] for (uint q = 0u; q < 4u; ++q) ep[kt][q] = f16vec2(e[kt][2u * q], e[kt][2u * q + 1u]);
precise f16vec2 st[2];
[[unroll]] for (uint tq = 0u; tq < 2u; ++tq) { // pairs q = 2tq (i0) and 2tq + 1 (i1 = i0 + 2)
const uint q0 = 2u * tq, q1 = q0 + 1u;
precise f16vec2 ev = ((ep[0][q0] + ep[0][q1]) + (ep[2][q0] + ep[2][q1])) + (ep[4][q0] + ep[4][q1]);
precise f16vec2 od = ((ep[1][q0] + ep[1][q1]) + (ep[3][q0] + ep[3][q1])) + (ep[5][q0] + ep[5][q1]);
st[tq] = ev + od;
}
precise f16vec2 loc = st[0] + st[1];
const uint other = subgroupShuffleXor(packFloat2x16(loc), 16u);
precise f16vec2 o = loc + unpackFloat2x16(other);
precise float16_t rs = o.x + o.y;
// f16(1 / rs) for a binary16 rs: x * m = 1 + d with |d| >= 2^-22 for every binary16 midpoint m, so the f32
// reciprocal (within an ulp) already rounds correctly
const float16_t r = float16_t(1.0 / float(rs));
const f16vec2 r2 = f16vec2(r, r);
[[unroll]] for (uint kt = 0u; kt < 6u; ++kt)
[[unroll]] for (uint q = 0u; q < 4u; ++q) {
precise f16vec2 pv = ep[kt][q] * r2;
// a probability is in [0, 1]: the plain conversion, no saturation clamp, both in one instruction
{ const fe4m3vec2 t_ = fe4m3vec2(pv); Pb[kt][2u * q] = t_.x; Pb[kt][2u * q + 1u] = t_.y; }
}
}
// 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);
[[unroll]] for (uint ks = 0u; ks < 3u; ++ks) {
A8 v0, v1;
coopMatLoad(v0, RR_VS, RR_VS0 + c * 384u + 8u * ks, 24u, RM);
coopMatLoad(v1, RR_VS, RR_VS0 + c * 384u + 8u * ks + 4u, 24u, RM);
K32(o, v0, Pb[2u * ks], v1, Pb[2u * ks + 1u])
}
Q8(Ob[c], o)
}
engine/rr/rr_capture.py
+156
View File
@@ -0,0 +1,156 @@
"""Ray Reconstruction captures: the launches of one denoiser frame as the CUDA bridge recorded them.
A capture directory holds what the bridge writes with
D4R_CUDA_VERBOSE=1 D4R_CUDA_REPLAY_DUMP_DIR=CAP/replay D4R_CUDA_REPLAY_DUMP_FILTER=cuda_dldn
D4R_CUDA_LAUNCH_STATS=1 D4R_CUDA_LAUNCH_STATS_FILTER=cuda_dldn D4R_CUDA_LAUNCH_DUMP_DIR=CAP/dump
plus the run's log as CAP/run/run.log (texture sampler settings and array shapes come from the verbose
cuTexObjectCreate / cuSurfObjectCreate / cuArrayCreate lines).
replay/replay-NNNNNN-KERNEL/ args.bin, manifest.txt and every buffer the launch points into, *before* it ran
dump/launch-NNN-KERNEL-argOOO-{buffer,texture,surface}-... every resource *after* the launch ran
Captures are reference data only: nothing at runtime reads them.
"""
from dataclasses import dataclass, field
from pathlib import Path
import re
import struct
# CUarray_format codes the denoiser uses, and how a Vulkan image of the same texels is described
# (kvk's BITS column: 8 = UNORM bytes, 9 = SNORM bytes, 16 = binary16, 32 = binary32)
ARRAY_FORMATS = {16: 16, 32: 32, 1: 8, 200: 9, 194: 8}
@dataclass
class Image:
offset: int
kind: str # 'texture' or 'surface'
width: int
height: int
channels: int
format: int # CUarray_format
path: Path # post-launch contents
linear: bool = False # textures: bilinear filtering
@property
def bits(self):
return ARRAY_FORMATS[self.format]
@property
def texel_bytes(self):
return self.channels * (1 if self.bits in (8, 9) else self.bits // 8)
@dataclass
class Pointer:
offset: int
alloc: int
delta: int
@dataclass
class Launch:
sequence: int
kernel: str
directory: Path
grid: tuple
block: tuple
args: bytes
allocs: dict = field(default_factory=dict) # index -> (base, size)
pointers: list = field(default_factory=list)
images: list = field(default_factory=list)
def pre(self, alloc):
"""an allocation's contents before the launch"""
return (self.directory / f'alloc-{alloc}.bin').read_bytes()
def pointer(self, offset):
return next(p for p in self.pointers if p.offset == offset)
def image(self, offset):
return next(i for i in self.images if i.offset == offset)
@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)
@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]
return max(set(candidates), key=lambda a: self.allocs[a][1]) if candidates else None
def post_buffer(self, capture, alloc):
base = self.allocs[alloc][0]
for path in (capture.root / 'dump').glob(f'launch-{self.sequence:03d}-{self.kernel}-arg*-buffer-*'):
m = re.search(r'-base0x([0-9a-f]+)\+', path.name)
if m and int(m.group(1), 16) == base:
return path.read_bytes()
raise FileNotFoundError(f'no post-launch dump of allocation {alloc} for launch {self.sequence}')
def u32(self, offset):
return struct.unpack_from('<I', self.args, offset)[0]
def f32(self, offset):
return struct.unpack_from('<f', self.args, offset)[0]
class Capture:
def __init__(self, root):
self.root = Path(root)
log_path = self.root / 'run' / 'run.log'
self.log = log_path.read_text(errors='replace') if log_path.exists() else ''
# (log position, texture handle, filter mode) of every texture object creation, in order
self._tex_events = [(m.start(), int(m.group(2), 16), int(m.group(1))) for m in re.finditer(
r'cuTexObjectCreate .*? filter=(\d+) .*? object=(0x[0-9a-f]+)', self.log)]
self.launches = []
for d in sorted((self.root / 'replay').iterdir()):
self.launches.append(self._launch(d))
def filter_at(self, handle, position):
value = None
for pos, h, f in self._tex_events:
if pos > position:
break
if h == handle:
value = f
return value
def _launch(self, directory):
manifest = (directory / 'manifest.txt').read_text().split('\n')
kernel = manifest[0].split()[1]
launch = [int(x) for x in manifest[1].split()[1:]]
sequence = int(re.search(r'replay-(\d+)-', directory.name).group(1))
args = (directory / 'args.bin').read_bytes()
result = Launch(sequence, kernel, directory, tuple(launch[0:3]), tuple(launch[3:6]), args)
for line in manifest[2:]:
f = line.split()
if not f:
continue
if f[0] == 'alloc':
result.allocs[int(f[1])] = (int(f[2], 16), int(f[3]))
elif f[0] == 'pointer':
result.pointers.append(Pointer(int(f[1]), int(f[2]), int(f[3])))
# the log position of this launch, so a reused texture handle resolves to its settings at the time
position = self.log.find(f'launch[{sequence}] {kernel}')
if position < 0:
position = len(self.log)
for path in sorted((self.root / 'dump').glob(f'launch-{sequence:03d}-{kernel}-arg*')):
m = re.search(r'-arg(\d+)-(texture|surface)-(\d+)x(\d+)x(\d+)-fmt(-?\d+)\.bin$', path.name)
if not m:
continue
offset = int(m.group(1))
image = Image(offset, m.group(2), int(m.group(3)), int(m.group(4)), int(m.group(5)), int(m.group(6)), path)
if image.kind == 'texture':
handle = struct.unpack_from('<Q', args, offset)[0]
image.linear = self.filter_at(handle, position) == 1
result.images.append(image)
result.images.sort(key=lambda i: (i.kind != 'texture', i.offset))
return result
def find(self, kernel_part, frame=None):
"""the launches whose kernel name contains KERNEL_PART, in order; FRAME (1-based) selects one"""
found = [l for l in self.launches if kernel_part in l.kernel]
return found if frame is None else found[frame - 1]
engine/rr/rr_decfilter.glsl
+65
View File
@@ -0,0 +1,65 @@
// The decoders' output filter (engine/rr/dec0_ref.outputs' surface 112), shared by dec1_m.comp and swin_m.comp: included
// after the head has been written to OUTF (16 f16 per window token, raster order) and a barrier. Texture 88 (one level
// down, bilinear) at the window's 10 x 10 tap positions into TAP; then two threads per token: the softmax of channels
// 0-8 over the 3x3 taps, one f16 fma chain per colour channel, blended with texture 96 by the gate (channel 9) plus
// texture 104, to surface 112.
// texture 88 (one level down) at the window's 10 x 10 full-resolution tap positions (clamped like the filter's
// coordinates), through the bilinear sampler at NVIDIA's coordinates: the reference's texture unit does not round
// like f32 arithmetic (computing the 1/4, 3/4 blend by hand left ~0.2% of surface 112 one ulp off, capDT 2026-10-12)
if (tid < 100u) {
const int xi = clampTap(x0 - 1 + int(tid % 10u), ex), yi = clampTap(y0 - 1 + int(tid / 10u), ey);
const float iw = uintBitsToFloat(param.p[34]), ih = uintBitsToFloat(param.p[35]);
precise float cu = fma(iw, float(xi), iw * 0.5), cv = fma(ih, float(yi), ih * 0.5);
const vec4 bl = textureLod(T88, vec2(cu, cv), 0.0);
[[unroll]] for (uint c = 0u; c < 3u; ++c) TAP[tid * 4u + c] = float16_t(bl[c]);
}
barrier();
// two threads per token: half 0 filters channels 0 and 1, half 1 channel 2
{
const uint t = tid >> 1u, hf = tid & 1u;
const int x = x0 + int(t & 7u), y = y0 + int(t >> 3u);
if (x < 0 || y < 0 || x >= ex || y >= ey) return;
float16_t a[10];
[[unroll]] for (uint c = 0u; c < 10u; ++c) a[c] = OUTF[t * 16u + c];
// surface 112: softmax of channels 0-8 over the 3x3 taps, one f16 fma chain per channel; the gate (channel 9)
// blends it with texture 96; + texture 104
const float16_t ma = max(max(max(a[0], a[2]), a[4]), a[6]), mb = max(max(max(a[1], a[3]), a[5]), a[7]);
const float16_t mx = max(max(ma, mb), a[8]);
float e[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) { precise float16_t d = a[i] - mx; e[i] = exp2(float(d) * 1.4426950216293335); }
precise float s = (e[0] + e[1]) + 0.0;
s = s + (e[2] + e[3]);
s = s + (e[4] + e[5]);
s = s + (e[6] + e[7]);
s = s + e[8];
const float16_t rr = rcp16(s);
float16_t wg[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) wg[i] = h16(r16(e[i]) * float(rr));
const uint tx = uint(x - x0), ty = uint(y - y0), c0 = 2u * hf, nc = 2u - hf;
float16_t f[2] = float16_t[2](0.0hf, 0.0hf);
bool mid = false;
[[unroll]] for (uint k = 0u; k < 9u; ++k) {
const uint at = ((ty + k / 3u) * 10u + tx + k % 3u) * 4u + c0;
[[unroll]] for (uint c = 0u; c < 2u; ++c) f[c] = fma16q(TAP[at + c], wg[k], f[c], mid);
}
if (mid) { // an f32 result on a binary16 midpoint: redo the chain with the exact fma (rare)
f[0] = 0.0hf; f[1] = 0.0hf;
for (uint k = 0u; k < 9u; ++k) {
const uint at = ((ty + k / 3u) * 10u + tx + k % 3u) * 4u + c0;
for (uint c = 0u; c < 2u; ++c) f[c] = fma16(TAP[at + c], wg[k], f[c]);
}
}
const float16_t g = sig(a[9]);
precise float16_t om = float16_t(1.0) - g;
const vec4 t96 = texelFetch(T96, ivec2(x, y), 0), t104 = texelFetch(T104, ivec2(x, y), 0);
float16_t o[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
precise float16_t fg = f[c] * g;
float16_t v = fma16(float16_t(t96[c0 + c]), om, fg);
const float16_t t4 = float16_t(t104[c0 + c]);
if (!isnan(t4)) { precise float16_t vv = v + t4; v = vv; }
o[c] = v;
}
const uint other = subgroupShuffleXor(bh(o[0]), 1u);
if (hf == 0u) imageStore(S112, ivec2(x, y), vec4(o[0], o[1], hb(other), 0.0));
}
engine/rr/rr_exact.glsl
+147
View File
@@ -0,0 +1,147 @@
// Exact NVIDIA arithmetic helpers for the native Ray Reconstruction stages (engine/rr/*_m.comp, enc1_native.comp).
// Requires GL_EXT_shader_explicit_arithmetic_types, GL_EXT_float_e4m3.
//
// RADV folds a register-only float(float16_t(x)) back to x, so every binary16 rounding whose result is used again goes
// through a conversion with an opaque zero (RR_ZERO16, set in main from something the compiler cannot evaluate, e.g.
// gl_NumWorkGroups.z - 1) OR-ed into its bits: one instruction more than the plain conversion.
uint16_t RR_ZERO16;
#if defined(RR_RELAXED)
// Relaxed rounding (build_native_model.sh with D4R_RR_RELAXED=1): the accumulators stay f32 between 32-deep steps and
// are converted to FP8 in one rounding instead of NVIDIA's two (binary16, then FP8). About 1% of the FP8 codes of a
// layer then differ by one step, so a stage no longer reproduces the reference bit for bit; it saves the conversions
// that are a quarter of a stage's instructions.
float rnd16(float x) { return x; }
#elif defined(RR_EXACT_CONVERT)
// the module's conversions are decorated NoContraction (engine/rr/spv_exact_convert.py): the driver keeps them
float rnd16(float x) { return float(float16_t(x)); }
#else
float rnd16(float x) { return float(uint16BitsToFloat16(float16BitsToUint16(float16_t(x)) | RR_ZERO16)); }
#endif
// x rounded to binary16 (nearest even), as f32 / as f16, through a conversion the compiler keeps (see rnd16)
float r16(float x) { return rnd16(x); }
#if (defined(RR_RELAXED) && RR_RELAXED >= 2) || defined(RR_EXACT_CONVERT)
float16_t h16(float x) { return float16_t(x); }
#else
float16_t h16(float x) { return uint16BitsToFloat16(float16BitsToUint16(float16_t(x)) | RR_ZERO16); }
#endif
// NVIDIA's k order within a 32-block: slot j holds natural channel inter(j); uninter is the inverse (and also the
// arena's channel order: arena slot uninter(k) holds natural channel k)
uint inter(uint j) { return (j & ~31u) | (((j & 31u) >> 2u) << 1u) | (j & 1u) | (((j >> 1u) & 1u) << 4u); }
uint uninter(uint c) { return (c & ~31u) | (((c & 15u) >> 1u) << 2u) | (c & 1u) | (((c >> 4u) & 1u) << 1u); }
int mirror(int v, int n) { v = abs(v); return min(v, 2 * n - 2 - v); }
// cvt.rn.satfinite.e4m3 (NaN -> 0x7f)
floate4m3_t q8(float16_t h)
{
return isnan(h) ? uintBitsToFloate4m3EXT(uint8_t(0x7fu)) : floate4m3_t(clamp(h, float16_t(-448.0), float16_t(448.0)));
}
// the same from f32 / f16 values in one saturating hardware conversion (satfinite; for an f32 value that is a binary16
// number the result equals the f16 conversion's)
floate4m3_t q8s(float x) { floate4m3_t r; saturatedConvertEXT(r, x); return r; }
floate4m3_t q8h(float16_t x) { floate4m3_t r; saturatedConvertEXT(r, x); return r; }
uint code(floate4m3_t v) { return uint(floate4m3BitsToUintEXT(v)); }
float16_t hb(uint b) { return uint16BitsToFloat16(uint16_t(b)); }
uint bh(float16_t h) { return uint(float16BitsToUint16(h)); }
// m * x compared with 1, exactly (the fma recovers the product's rounding error): -1, 0 or 1
int cmp1(float m, float x)
{
precise float t = m * x;
precise float e = fma(m, x, -t);
return t > 1.0 ? 1 : (t < 1.0 ? -1 : (e > 0.0 ? 1 : (e < 0.0 ? -1 : 0)));
}
// 1 / x of the filters' softmax and gate, rounded to binary16. The reference is the GPU's own reciprocal instruction
// (rcp.approx; ZLUDA lowers it to llvm.amdgcn.rcp.f32, which is what 1.0 / x compiles to here), not the correctly
// rounded quotient: correcting it left ~0.01% of the filtered surfaces one ulp off (capDT, 2026-10-12).
float16_t rcp16(float x) { return h16(1.0 / x); }
// binary16 rounding of 1 / sqrt(s) (s > 0)
float16_t rsq16(float16_t s16)
{
const float s = float(s16);
float16_t q = h16(inversesqrt(s));
#if defined(RR_RELAXED) && RR_RELAXED >= 2
return q;
#endif
for (int i = 0; i < 2; ++i) {
const uint b = bh(q);
const float lo = (float(q) + float(hb(b - 1u))) * 0.5, hi = (float(q) + float(hb(b + 1u))) * 0.5;
const int cl = cmp1(lo * lo, s), ch = cmp1(hi * hi, s);
if (cl > 0 || (cl == 0 && (b & 1u) != 0u)) q = hb(b - 1u);
else if (ch < 0 || (ch == 0 && (b & 1u) != 0u)) q = hb(b + 1u);
}
return q;
}
// fma of binary16 operands with one binary16 rounding, exactly. The product is exact in f32, so an f32 fma rounds the
// exact result once; converting that to binary16 rounds twice, which differs only when the f32 value lands exactly
// on a binary16 midpoint: then the sum's rounding error (TwoSum, the rare branch) decides the direction.
float16_t fma16(float16_t a, float16_t b, float16_t c)
{
const float fa = float(a), fb = float(b), fc = float(c);
precise float s = fma(fa, fb, fc);
#if defined(RR_RELAXED) && RR_RELAXED >= 2
return float16_t(s);
#endif
const uint u = floatBitsToUint(s);
if ((u & 0x1fffu) == 0x1000u && abs(s) >= 6.103515625e-05) {
precise float p = fa * fb;
precise float bb = s - p;
precise float err = (p - (s - bb)) + (fc - bb);
if (err != 0.0) return h16(uintBitsToFloat(err * s > 0.0 ? u + 1u : u - 1u));
}
return h16(s);
}
// the same in one f32 fma, rounded once more to binary16: mid is set when the f32 result lies on a binary16 midpoint,
// where that double rounding can differ (the caller then redoes the computation with fma16)
float16_t fma16q(float16_t a, float16_t b, float16_t c, inout bool mid)
{
precise float s = fma(float(a), float(b), float(c));
#if defined(RR_RELAXED) && RR_RELAXED >= 2
return float16_t(s);
#endif
mid = mid || (floatBitsToUint(s) & 0x1fffu) == 0x1000u;
#ifdef RR_EXACT_CONVERT
return float16_t(s);
#else
return hb(bh(float16_t(s)) | uint(RR_ZERO16));
#endif
}
// The saturating FP8 conversion of a whole accumulator, two elements per conversion instruction
#define Q8PAIRS(dst, src) [[unroll]] for (int i_ = 0; i_ < 8; i_ += 2) { fe4m3vec2 t_; saturatedConvertEXT(t_, vec2((src)[i_], (src)[i_ + 1])); (dst)[i_] = t_.x; (dst)[i_ + 1] = t_.y; }
// q8(act(h)) from the packed table (enc1_pack.act_table; the weight buffer's W8U view at W_ACT)
#define ACTQ(h) uintBitsToFloate4m3EXT(w8u.w[W_ACT + bh(h)])
float16_t act(float16_t x)
{
precise float16_t y = clamp(x, float16_t(-2.0), float16_t(2.0));
precise float16_t a = float16_t(0.0810546875) * abs(y);
precise float16_t b = float16_t(0.412109375) - a;
precise float16_t c = y * b;
precise float16_t g = float16_t(0.5) + c;
precise float16_t r = x * g;
return r;
}
float16_t sig(float16_t v) { return rcp16(exp2(float(v) * -1.4426950216293335) + 1.0); }
// the cubic-exponent softmax numerator (enc1_ref.softmax_cubic): a cubic in the clamped scaled score whose binary16
// bits, shifted into the exponent field, are the exponential
// the same in plain f32 fmas (double rounding); bad is set when an f32 result lies exactly on a binary16 midpoint,
// the only case where that differs: the caller then recomputes with cubic_e (subgroupAny, rare)
float16_t cubic_e_fast(float16_t x, inout bool bad)
{
precise float16_t m = clamp(x * float16_t(0.017181396484375), float16_t(-0.55615234375), float16_t(0.55615234375));
const float fm = float(m);
precise float t = fma(fm, -fm, 0.92724609375);
const float tr = rnd16(t);
precise float u = fma(fm, tr, 1.375);
// t - c and u - 1.375 are exact (Sterbenz), so one more fma gives the rounding error's sign: only an inexact
// midpoint is a double rounding
precise float et = fma(fm, fm, t - 0.92724609375), eu = fma(fm, -tr, u - 1.375);
bad = bad || ((floatBitsToUint(t) & 0x1fffu) == 0x1000u && et != 0.0) || ((floatBitsToUint(u) & 0x1fffu) == 0x1000u && eu != 0.0);
return hb((bh(float16_t(u)) << 5u) & 0x7FE0u);
}
float16_t cubic_e(float16_t x)
{
precise float16_t m = clamp(x * float16_t(0.017181396484375), float16_t(-0.55615234375), float16_t(0.55615234375));
const float16_t t = fma16(m, -m, float16_t(0.92724609375));
const float16_t u = fma16(m, t, float16_t(1.375));
return hb((bh(u) << 5u) & 0x7FE0u);
}
engine/rr/rr_inputs.py
+65
View File
@@ -0,0 +1,65 @@
#!/usr/bin/env python3
"""Write the game inputs of a capture's denoiser frames for engine/replay_rr.cpp (d4r-rr).
usage: rr_inputs.py CAPTURE OUT_DIR [--frames N] [--no-reset]
The inputs are enc0's textures (colour 160, motion 168, normals 184, roughness 192, the albedos 232 and 240, depth
512); the frame settings come from enc0's argument block: jitter (320), motion scale (344, in output pixels per
render pixel), depth inversion (528) and the camera, whose rotation-times-projection (596) and projection (660) give
back a world-to-view rotation and the view-to-clip matrix that reproduce the block. The harness passes reset on every
frame (scripts/run_rr_engine_compare.sh); --no-reset clears it after the first frame.
"""
import argparse
import os
import struct
import sys
from pathlib import Path
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rr_capture import Capture # noqa: E402
OPTIONAL = {'hitdistance': 248, 'disocclusion': 296}
INPUTS = {'color': 160, 'motion': 168, 'normals': 184, 'roughness': 192, 'diffuse': 232, 'specular': 240, 'depth': 512}
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('out', type=Path)
ap.add_argument('--frames', type=int, default=3)
ap.add_argument('--no-reset', action='store_true')
o = ap.parse_args()
cap = Capture(o.capture)
o.out.mkdir(parents=True, exist_ok=True)
lines = []
for f in range(1, o.frames + 1):
enc0 = cap.find('swin_enc0', f)
a = enc0.args
u32 = lambda off: struct.unpack_from('<I', a, off)[0]
f32 = lambda off: struct.unpack_from('<f', a, off)[0]
mat = lambda off: np.array(struct.unpack_from('<16f', a, off), np.float64).reshape(4, 4)
rw, rh, ow, oh = u32(144), u32(148), u32(472), u32(476)
if f == 1:
lines.append(f'size {rw} {rh} {ow} {oh}')
images = {i.offset: i for i in enc0.images}
for name, off in INPUTS.items():
(o.out / f'{f}-{name}.bin').write_bytes(images[off].path.read_bytes())
for name, off in OPTIONAL.items(): # DLSSD.SpecularHitDistance, DLSS.DisocclusionMask when the game passed them
if off in images:
(o.out / f'{f}-{name}.bin').write_bytes(images[off].path.read_bytes())
clip = mat(660)
view = mat(596) @ np.linalg.inv(clip)
view[3, :3] = 0.0
jx, jy = f32(320), f32(324)
mvx, mvy = f32(344) * rw / ow, f32(348) * rh / oh
reset = 1 if f == 1 or not o.no_reset else 0
values = ' '.join(f'{v:.9g}' for v in list(view.flatten()) + list(clip.flatten()))
lines.append(f'frame {f} {jx:.9g} {jy:.9g} {mvx:.9g} {mvy:.9g} 1 {reset} {u32(528)} {values}')
(o.out / 'frames.txt').write_text('\n'.join(lines) + '\n')
print(f'{o.frames} frames of {o.capture} in {o.out}')
if __name__ == '__main__':
main()
engine/rr/rrmath.py
+40
View File
@@ -0,0 +1,40 @@
"""Shared numeric helpers for the Ray Reconstruction port: e4m3 codes, f16, and ZLUDA's m16n8k32 accumulation."""
import numpy as np
_E4M3 = np.zeros(256)
for _c in range(256):
_s, _e, _m = _c >> 7, (_c >> 3) & 15, _c & 7
_v = np.nan if (_e == 15 and _m == 7) else (_m * 2.0 ** -9 if _e == 0 else (1 + _m / 8) * 2.0 ** (_e - 7))
_E4M3[_c] = -_v if _s else _v
def e4m3(codes):
return _E4M3[np.asarray(codes, np.uint8)]
def f16(bits):
return np.asarray(bits, np.uint16).view(np.float16).astype(np.float64)
def f16bits(x):
return np.asarray(x, np.float64).astype(np.float16).view(np.uint16)
def mma_k32(a_codes, b_codes, c_f16bits, ksteps=None):
"""D = A.B + C with ZLUDA's lowering: per 32-deep k-step, f32 sum of the exact products and C, one f16 rounding;
the next k-step takes that f16 value as C. (f64 stands in for f32: the products and short sums are exact.)"""
a, b = e4m3(a_codes), e4m3(b_codes)
acc = f16(c_f16bits)
for k in range(0, a.shape[1], 32):
acc = (a[:, k:k + 32] @ b[k:k + 32, :] + acc).astype(np.float32).astype(np.float16).astype(np.float64)
return f16bits(acc)
# cvt.rn.satfinite.e4m3 of every f16 bit pattern (the hardware conversion, exhaustively equal to PTX's semantics;
# NaN -> 0x7f), see engine/rr/ptx2rr.py h2e4m3_hw
import os as _os
_Q8 = np.load(_os.path.join(_os.path.dirname(__file__), 'f16_to_e4m3.npy'))
def q8(x_f16):
"""f16 values (float array) -> e4m3 codes"""
return _Q8[np.asarray(x_f16, np.float64).astype(np.float16).view(np.uint16)]
engine/rr/spv_exact_convert.py
+41
View File
@@ -0,0 +1,41 @@
#!/usr/bin/env python3
"""Decorate every float conversion of a SPIR-V module NoContraction, so the driver keeps it.
usage: spv_exact_convert.py IN.spv [OUT.spv]
RADV folds float(float16_t(x)) back to x unless the conversions are marked exact; glslang's `precise` does not reach
conversions. With the decoration a binary16 rounding is two instructions (engine/rr/rr_exact.glsl, RR_EXACT_CONVERT)
instead of three with an opaque zero OR-ed into the bits.
"""
import struct
import sys
OP_DECORATE, OP_FCONVERT, NO_CONTRACTION = 71, 115, 42
def main():
src = sys.argv[1]
dst = sys.argv[2] if len(sys.argv) > 2 else src
data = open(src, 'rb').read()
words = list(struct.unpack(f'<{len(data) // 4}I', data))
if words[0] != 0x07230203:
sys.exit(f'{src}: not SPIR-V')
insts, i = [], 5
while i < len(words):
count = words[i] >> 16
insts.append((words[i] & 0xffff, words[i:i + count]))
i += count
ids = [w[2] for op, w in insts if op == OP_FCONVERT]
first_type = next(k for k, (op, _) in enumerate(insts) if 19 <= op <= 39) # annotations end where types begin
out = list(words[:5])
for k, (op, w) in enumerate(insts):
if k == first_type:
for result in ids:
out += [3 << 16 | OP_DECORATE, result, NO_CONTRACTION]
out += w
open(dst, 'wb').write(struct.pack(f'<{len(out)}I', *out))
print(f'{dst}: {len(ids)} conversions marked exact')
if __name__ == '__main__':
main()
engine/rr/spv_float_controls.py
+71
View File
@@ -0,0 +1,71 @@
#!/usr/bin/env python3
"""Give a compute SPIR-V module PTX's floating-point environment: RoundingModeRTE for 16- and 32-bit results and
DenormPreserve for 16-bit values (SPV_KHR_float_controls, core in SPIR-V 1.4).
usage: spv_float_controls.py IN.spv [OUT.spv]
Without these execution modes the rounding of float conversions is the driver's choice, where every cvt.rn of
NVIDIA's kernels rounds to nearest even. glslang keeps only one execution mode per mode id, so the modes are inserted
into the compiled module instead.
"""
import struct
import sys
OP_EXTENSION, OP_CAPABILITY, OP_ENTRY_POINT, OP_EXECUTION_MODE = 10, 17, 15, 16
CAP_DENORM_PRESERVE, CAP_ROUNDING_RTE = 4464, 4467
MODES = ((4462, 16), (4462, 32), (4459, 16)) # RoundingModeRTE 16/32, DenormPreserve 16
EXTENSION = 'SPV_KHR_float_controls'
def instructions(words):
i = 5
while i < len(words):
count = words[i] >> 16
yield words[i] & 0xffff, words[i:i + count]
i += count
def string_words(text):
data = text.encode() + b'\0'
data += b'\0' * (-len(data) % 4)
return list(struct.unpack(f'<{len(data) // 4}I', data))
def patch(words):
insts = list(instructions(words))
caps = {w[1] for op, w in insts if op == OP_CAPABILITY}
has_ext = any(op == OP_EXTENSION and EXTENSION.encode() in struct.pack(f'<{len(w) - 1}I', *w[1:]) for op, w in insts)
entry = next(w[2] for op, w in insts if op == OP_ENTRY_POINT)
modes = {(w[2], w[3]) for op, w in insts if op == OP_EXECUTION_MODE and w[1] == entry and len(w) > 3}
last_cap = max(i for i, (op, _) in enumerate(insts) if op == OP_CAPABILITY)
first_mode = min(i for i, (op, _) in enumerate(insts) if op == OP_EXECUTION_MODE)
out = list(words[:5])
for i, (op, w) in enumerate(insts):
if i == first_mode:
for mode, width in MODES:
if (mode, width) not in modes:
out += [4 << 16 | OP_EXECUTION_MODE, entry, mode, width]
out += w
if i == last_cap:
for cap in (CAP_DENORM_PRESERVE, CAP_ROUNDING_RTE):
if cap not in caps:
out += [2 << 16 | OP_CAPABILITY, cap]
if not has_ext:
name = string_words(EXTENSION)
out += [(1 + len(name)) << 16 | OP_EXTENSION, *name]
return out
def main():
src = sys.argv[1]
dst = sys.argv[2] if len(sys.argv) > 2 else src
data = open(src, 'rb').read()
words = list(struct.unpack(f'<{len(data) // 4}I', data))
if words[0] != 0x07230203:
sys.exit(f'{src}: not SPIR-V')
out = patch(words)
open(dst, 'wb').write(struct.pack(f'<{len(out)}I', *out))
if __name__ == '__main__':
main()
engine/rr/swin_m.comp
+471
View File
@@ -0,0 +1,471 @@
#version 460
// Native Swin stage of the DLSS-D (preset E) denoiser, register-resident, for any stage of engine/rr/swin_ref.STAGES
// (enc2-enc5, dec2-dec4; enc1_m.comp / dec1_m.comp are the hand-written C = 64 instances it generalises). The stage's
// shape and weight offsets come from the include generated by engine/rr/swin_pack.py:
// KIND 0 encoder (tokens from the arena, 2x2 merge to CX channels, enc1's gated filter; CX = 0: the bottleneck,
// ungated filter, no merge), 1 decoder (expand of the CX channels one level down + skip, dec0's filter)
// C channels (KC = C / 32 planes, MT = C / 16 tiles), NH heads of 32 channels, C / 8 MLP branches
// The arithmetic is engine/rr/swin_ref.py's; enc1_m.comp's header explains the transposed products and slot orders.
// One workgroup per 8x8-token window, four waves, wave w owning query slots 16w..16w+15.
#extension GL_KHR_cooperative_matrix : require
#extension GL_KHR_memory_scope_semantics : require
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_8bit_storage : require
#extension GL_EXT_float_e4m3 : require
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_EXT_control_flow_attributes : require
#extension GL_GOOGLE_include_directive : require
layout(local_size_x = 128) in;
#include "swin_stage.h"
#include "rr_exact.glsl"
#define KC (C / 32)
#define MT (C / 16)
#define NB (C / 8)
layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param;
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;
layout(set = 0, binding = 3) uniform sampler2D T88;
#if KIND == 1
layout(set = 0, binding = 4) uniform sampler2D T96;
layout(set = 0, binding = 5) uniform sampler2D T104;
layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S112;
#else
layout(set = 0, binding = 4, rgba16f) uniform writeonly image2D S112;
#if CX
layout(set = 0, binding = 5, rgba16f) uniform writeonly image2D S120;
layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S128;
#endif
#endif
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define HAC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CM gl_CooperativeMatrixLayoutColumnMajor
#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM)
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
#define Q8(dst, src) Q8PAIRS(dst, src)
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) }
const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0;
// NVIDIA's reduction over a short list (swin_ref.pairs_tree)
#if KC == 1
#define TREE(v) (v[0])
#elif KC == 2
#define TREE(v) (v[0] + v[1])
#elif KC == 3
#define TREE(v) ((v[0] + v[1]) + v[2])
#elif KC == 4
#define TREE(v) ((v[0] + v[1]) + (v[2] + v[3]))
#else
#define TREE(v) (((v[0] + v[1]) + (v[2] + v[3])) + v[4])
#endif
#if KIND == 1
// one pool, by phase: [EX | LOW] for the expand; the prologue reads EX and writes XT over LOW (dead) and beyond;
// the attention's K and V^T reuse EX (dead after the residual)
#define EX_W (48u * C)
#define POOL_W (EX_W + (12u * CX > 24u * C ? 12u * CX : 24u * C))
shared uint POOL[POOL_W];
#define EX POOL // the expansion: 48 rows x 4 sub-positions x C FP8 channels in slot order
#define LOW_AT EX_W // the low-resolution tokens (6 x 6, padded): CX FP8 channels, arena order
#define XT_AT EX_W // the normalised key tokens: slot-major, C FP8 channels in slot order
#define XTS POOL
#define RR_KS POOL
#define RR_KS0 0u
#define RR_VS POOL
#define RR_VS0 768u
shared float16_t TAP[100 * 4]; // texture 88 bilinear at the window's 10 x 10 filter tap positions
shared float16_t OUTF[64 * 16]; // the head per window token (raster order)
#else
shared uint XT[96 * C / 4]; // the normalised key tokens: slot-major, C FP8 channels in slot order
#define XT_AT 0u
#define XTS XT
shared uint KVS[2 * 768]; // the current head's K, then V^T
#define RR_KS KVS
#define RR_KS0 0u
#define RR_VS KVS
#define RR_VS0 768u
#define MG XT // the merge rows (16 x 4C FP8) reuse XT, dead after the projections
#if CX
shared float16_t O112[64 * 4]; // surface 112 per token, for surface 128's 2x2 sums
#endif
#endif
ivec2 keyXY(uint t, int x0, int y0)
{
if (t < 64u) return ivec2(x0 + int(t & 7u), y0 + int(t >> 3u));
const int i = int(t & 7u), side = int((t - 64u) >> 3u);
return ivec2(side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8), side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i));
}
int clampTap(int v, int n) { return v < 0 ? 1 : (v >= n ? n - 2 : v); }
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)); }
void main()
{
const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID;
const uint hl = lane >> 4u, col = lane & 15u;
RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u);
const int ex = int(param.p[0]), ey = int(param.p[1]), hx = ex / 2, hy = ey / 2;
const uint inOff = param.p[2], skOff = param.p[6], outOff = param.p[12], mrgOff = param.p[14], shift = param.p[20];
const int x0 = 8 * int(gl_WorkGroupID.x) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u);
const uint plane = uint(ex * ey) * 32u, lplane = uint(hx * hy) * 32u;
#if KIND == 1
// ---------------------------------------------------------------- expand: wave w computes sub-position w
const int X0 = (x0 >> 1) - 1, Y0 = (y0 >> 1) - 1;
for (uint i = tid; i < 48u * (CX / 4u); i += 128u) {
const uint r = i / (CX / 4u), wd = i % (CX / 4u);
uint u = 0u;
if (r < 36u) {
const int X = mirror(X0 + int(r % 6u), hx), Y = mirror(Y0 + int(r / 6u), hy);
u = arena.w[(inOff + (wd >> 3u) * lplane + uint(Y * hx + X) * 32u) / 4u + (wd & 7u)];
}
POOL[LOW_AT + i] = u;
}
barrier();
[[unroll]] for (uint n = 0u; n < 3u; ++n) {
B8 lb[CX / 16];
[[unroll]] for (uint kt = 0u; kt < CX / 16; ++kt) 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)]);
[[unroll]] for (uint kt = 0u; kt < CX / 16; kt += 2u) {
A8 w0, w1;
LDA(w0, W_WE, CX, m, kt); LDA(w1, W_WE, CX, m, kt + 1u);
K32(a, w0, lb[kt], w1, lb[kt + 1u])
}
B8 q;
Q8(q, a)
coopMatStore(q, EX, n * 16u * C + 4u * m, C, CM);
}
}
barrier();
#endif
// ---------------------------------------------------------------- prologue: one key token per thread
if (tid < 96u) {
const ivec2 k = keyXY(inter(tid), x0, y0);
const int xx = mirror(k.x, ex), yy = mirror(k.y, ey);
const uint tb = ((KIND == 1 ? skOff : inOff) + uint(yy * ex + xx) * 32u) / 4u;
#if KIND == 1
const uint r = uint(((k.y >> 1) - Y0) * 6 + ((k.x >> 1) - X0)), s = uint(2 * (k.y & 1) + (k.x & 1));
#endif
float16_t v[C];
[[unroll]] for (uint wd = 0u; wd < C / 4; ++wd) {
const uint t = arena.w[tb + (wd >> 3u) * (plane >> 2u) + (wd & 7u)];
#if KIND == 1
const uint e = EX[r * C + s * (C / 4u) + wd];
[[unroll]] for (uint b = 0u; b < 4u; ++b) { precise float16_t x = e4(e, b) + e4(t, b); v[4u * wd + b] = x; }
#else
[[unroll]] for (uint b = 0u; b < 4u; ++b) v[4u * wd + b] = e4(t, b);
#endif
}
// swin_ref.input_norm: per lane word q (and q + 4) the pair sums of the planes (TREE), the halves, then
// q ^ 1, q ^ 2, the components
precise f16vec2 h[4];
[[unroll]] for (uint q = 0u; q < 4u; ++q) {
precise f16vec2 hh[2];
[[unroll]] for (uint half_ = 0u; half_ < 2u; ++half_) {
precise f16vec2 pr[KC];
[[unroll]] for (uint p = 0u; p < KC; ++p) {
const uint o = 32u * p + 4u * (q + 4u * half_);
precise f16vec2 a = f16vec2(v[o], v[o + 1u]) * f16vec2(v[o], v[o + 1u]);
precise f16vec2 c = f16vec2(v[o + 2u], v[o + 3u]) * f16vec2(v[o + 2u], v[o + 3u]);
pr[p] = a + c;
}
precise f16vec2 tr = TREE(pr);
hh[half_] = tr;
}
h[q] = hh[0] + hh[1];
}
precise f16vec2 h0 = h[0] + h[1], h2_ = h[2] + h[3];
precise f16vec2 hs = h0 + h2_;
precise float16_t sm = hs.x + hs.y;
precise float16_t se = sm + float16_t(0.0001220703125);
const float16_t rr = rsq16(se);
[[unroll]] for (uint wd = 0u; wd < C / 4; ++wd) {
floate4m3_t c[4];
[[unroll]] for (uint b = 0u; b < 4u; ++b) {
const uint j = 4u * wd + b;
precise float16_t g = rr * w16.w[W_GAIN0 / 2u + inter(j)];
precise float16_t n = v[j] * g;
c[b] = q8h(n);
}
XTS[XT_AT + tid * (C / 4u) + wd] = pack4(c[0], c[1], c[2], c[3]);
}
}
barrier();
// ---------------------------------------------------------------- the residual: the window token + table 7
ACC R[MT];
{
const uint tq = inter(16u * wv + col);
const int qx = x0 + int(tq & 7u), qy = y0 + int(tq >> 3u);
const uint tb = ((KIND == 1 ? skOff : inOff) + uint(mirror(qy, ey) * ex + mirror(qx, ex)) * 32u) / 4u;
#if KIND == 1
const uint r = uint(((qy >> 1) - Y0) * 6 + ((qx >> 1) - X0)), s = uint(2 * (qy & 1) + (qx & 1));
#endif
[[unroll]] for (uint m = 0u; m < MT; ++m) {
const uint ta = tb + (m >> 1u) * (plane >> 2u) + 4u * (m & 1u) + 2u * hl;
const uvec2 t = uvec2(arena.w[ta], arena.w[ta + 1u]);
#if KIND == 1
const uint ea = r * C + s * (C / 4u) + 4u * m + 2u * hl;
const uvec2 e = uvec2(EX[ea], EX[ea + 1u]);
#endif
[[unroll]] for (int i = 0; i < 8; ++i) {
#if KIND == 1
precise float16_t x = e4(e[i >> 2], uint(i) & 3u) + e4(t[i >> 2], uint(i) & 3u);
#else
const float16_t x = e4(t[i >> 2], uint(i) & 3u);
#endif
precise float16_t c = x + w16.w[W_TABLE7 / 2u + 16u * m + 8u * hl + uint(i)];
R[m][i] = float(c);
}
}
}
// ---------------------------------------------------------------- attention, NH heads
[[dont_unroll]] for (uint hd = 0u; hd < NH; ++hd) {
barrier(); // every wave is done with the previous head's K / V (and, for the first, the expand's LOW)
for (uint j = 0u; j < 3u; ++j) { // 12 jobs (token tile n, channel tile c)
const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u;
ACC a = ACC(0.0), v = ACC(0.0);
[[unroll]] for (uint kt = 0u; kt < 2u * KC; kt += 2u) {
B8 xb0, xb1, wv0, wv1;
A8 xa0, xa1, wk0, wk1;
coopMatLoad(xb0, XTS, XT_AT + n * 4u * C + 4u * kt, C / 4u, CM);
coopMatLoad(xb1, XTS, XT_AT + n * 4u * C + 4u * kt + 4u, C / 4u, CM);
coopMatLoad(xa0, XTS, XT_AT + n * 4u * C + 4u * kt, C / 4u, RM);
coopMatLoad(xa1, XTS, XT_AT + n * 4u * C + 4u * kt + 4u, C / 4u, RM);
LDA(wk0, W_WK0 + HEAD * hd, C, c, kt); LDA(wk1, W_WK0 + HEAD * hd, C, c, kt + 1u);
LDB(wv0, W_WV0 + HEAD * hd, 32u, kt, c); LDB(wv1, W_WV0 + HEAD * hd, 32u, kt + 1u, c);
K32(a, wk0, xb0, wk1, xb1)
K32(v, xa0, wv0, xa1, wv1)
}
A8 ka, va;
Q8(ka, a)
Q8(va, v)
coopMatStore(ka, RR_KS, RR_KS0 + n * 128u + 4u * c, 8u, RM);
coopMatStore(va, RR_VS, RR_VS0 + c * 384u + 4u * n, 24u, RM);
}
barrier();
#include "rr_attn.glsl"
[[unroll]] for (uint m = 0u; m < MT; ++m) {
A8 a0, a1;
LDA(a0, W_WO0 + HEAD * hd, 32u, m, 0u); LDA(a1, W_WO0 + HEAD * hd, 32u, m, 1u);
K32(R[m], a0, Ob[0], a1, Ob[1])
}
}
// ---------------------------------------------------------------- MLP
B8 Nb[MT];
{
float16_t rv[MT][8];
[[unroll]] for (uint m = 0u; m < MT; ++m) [[unroll]] for (int i = 0; i < 8; ++i) rv[m][i] = float16_t(R[m][i]);
// swin_ref.rowsum_sq: slot 16m + 8hl + i is NVIDIA's channel 8j + 2t + c with j = 4(m >> 1) + (m & 1) +
// 2((i >> 1) & 1), t = 2hl + (i >> 2), c = i & 1: per parity (m & 1) the tile pairs (i, i + 2) of every
// 32-channel block, TREE over the blocks; the parities; t ^ 1 (local), t ^ 2 (the other lane half)
float16_t h1[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
float16_t ht[2];
[[unroll]] for (uint tq = 0u; tq < 2u; ++tq) {
const uint i0 = 4u * tq + c, i1 = i0 + 2u;
precise float16_t par[2];
[[unroll]] for (uint pa = 0u; pa < 2u; ++pa) {
precise float16_t pr[KC];
[[unroll]] for (uint b = 0u; b < KC; ++b)
pr[b] = rv[2u * b + pa][i0] * rv[2u * b + pa][i0] + rv[2u * b + pa][i1] * rv[2u * b + pa][i1];
precise float16_t tr = TREE(pr);
par[pa] = tr;
}
precise float16_t hv = par[0] + par[1];
ht[tq] = hv;
}
precise float16_t hh = ht[0] + ht[1];
h1[c] = hh;
}
const uint other = subgroupShuffleXor(bh(h1[0]) | (bh(h1[1]) << 16u), 16u);
precise float16_t g0 = h1[0] + hb(other & 0xffffu), g1 = h1[1] + hb(other >> 16u);
precise float16_t s = g0 + g1;
precise float16_t se = s + float16_t(0.0001220703125);
const float16_t r = rsq16(se);
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) {
const uint k = 16u * m + 8u * hl + uint(i);
precise float16_t g = r * w16.w[W_GAIN14 / 2u + k];
precise float16_t n = rv[m][i] * g;
Nb[m][i] = q8h(n);
precise float16_t a = rv[m][i] + w16.w[W_TABLE16 / 2u + k];
R[m][i] = float(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)]);
[[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(R[m], a0, Hb[0], a1, Hb[1])
}
}
// ---------------------------------------------------------------- outputs
const uint t = inter(16u * wv + col);
const int tx = x0 + int(t & 7u), ty = y0 + int(t >> 3u);
const bool inside = tx >= 0 && ty >= 0 && tx < ex && ty < ey;
B8 Tb[MT];
{
const uint o = (outOff + uint(ty * ex + tx) * 32u) >> 2u;
#if KIND == 0 && CX
const uint ty2 = t >> 3u, tx2 = t & 7u;
const uint mr = (ty2 >> 1u) < 2u ? 2u * (tx2 >> 1u) + (ty2 >> 1u) : 8u + 2u * (tx2 >> 1u) + (ty2 >> 1u) - 2u;
const uint blk = 2u * (ty2 & 1u) + (tx2 & 1u);
barrier(); // every wave is done with XT (the merge rows reuse it)
#endif
[[unroll]] for (uint m = 0u; m < MT; ++m) {
Q8(Tb[m], R[m])
const uint w0 = pack4(Tb[m][0], Tb[m][1], Tb[m][2], Tb[m][3]), w1 = pack4(Tb[m][4], Tb[m][5], Tb[m][6], Tb[m][7]);
const uint at = 4u * (m & 1u) + 2u * hl;
if (inside) { arena.w[o + (m >> 1u) * (plane >> 2u) + at] = w0; arena.w[o + (m >> 1u) * (plane >> 2u) + at + 1u] = w1; }
#if KIND == 0 && CX
const uint g = mr * C + blk * (C / 4u) + 4u * m + 2u * hl;
MG[g] = w0; MG[g + 1u] = w1;
#endif
}
}
// the C -> 16 head
ACC hv;
[[unroll]] for (int i = 0; i < 8; ++i) hv[i] = float(w16.w[W_BH / 2u + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) {
A8 a0, a1;
LDA(a0, W_WH, C, 0u, kt); LDA(a1, W_WH, C, 0u, kt + 1u);
K32(hv, a0, Tb[kt], a1, Tb[kt + 1u])
}
#if KIND == 0
{ // enc1's kernel-prediction filter (lanes 0-15; 16-31 hold head channels 8-15)
float16_t l[12];
[[unroll]] for (int i = 0; i < 8; ++i) l[i] = float16_t(hv[i]);
const uint u0 = subgroupShuffleXor(bh(l[0]) | (bh(l[1]) << 16u), 16u), u1 = subgroupShuffleXor(bh(l[2]) | (bh(l[3]) << 16u), 16u);
l[8] = hb(u0 & 0xffffu); l[9] = hb(u0 >> 16u); l[10] = hb(u1 & 0xffffu); l[11] = hb(u1 >> 16u);
const float16_t ma = max(max(max(l[0], l[2]), l[4]), l[6]), mb = max(max(max(l[1], l[3]), l[5]), l[7]);
const float16_t mx = max(max(ma, mb), l[8]);
float e[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) { precise float16_t d = l[i] - mx; e[i] = exp2(float(d) * 1.4426950216293335); }
precise float s = (e[0] + e[1]) + 0.0;
s = s + (e[2] + e[3]);
s = s + (e[4] + e[5]);
s = s + (e[6] + e[7]);
s = s + e[8];
const float16_t rr = rcp16(s);
float wg[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) wg[i] = r16(r16(e[i]) * float(rr));
const int x = tx, y = ty;
const int xm = x - 1 < 0 ? 1 : x - 1, xc = x < 0 ? 1 : (x >= ex ? ex - 2 : x), xp = x + 1 >= ex ? ex - 2 : x + 1;
const int ym = y - 1 < 0 ? 1 : y - 1, yc = y < 0 ? 1 : (y >= ey ? ey - 2 : y), yp = y + 1 >= ey ? ey - 2 : y + 1;
const ivec2 pts[9] = ivec2[9](ivec2(xm, ym), ivec2(xc, ym), ivec2(xp, ym), ivec2(xm, yc), ivec2(xc, yc),
ivec2(xp, yc), ivec2(xm, yp), ivec2(xc, yp), ivec2(xp, yp));
float acc[3][2] = float[3][2](float[2](0.0, 0.0), float[2](0.0, 0.0), float[2](0.0, 0.0));
[[unroll]] for (uint k = 0u; k < 8u; ++k) {
const vec4 tv = texelFetch(T88, pts[k], 0);
[[unroll]] for (uint c = 0u; c < 3u; ++c) { precise float f = tv[c] * wg[k] + acc[c][k & 1u]; acc[c][k & 1u] = r16(f); }
}
const vec4 t8 = texelFetch(T88, pts[8], 0);
float o112[3], o120[3];
[[unroll]] for (uint c = 0u; c < 3u; ++c) {
precise float f0 = acc[c][0] + acc[c][1];
precise float f1 = t8[c] * wg[8] + r16(f0);
const float f = r16(f1);
#if CX
precise float a = float(sig(l[9u + c])) * f;
const float ar = r16(a);
precise float b = f - ar;
o120[c] = ar; o112[c] = r16(b);
#else
o112[c] = f; // the bottleneck: the filtered value itself
#endif
}
if (hl == 0u && inside) {
imageStore(S112, ivec2(x, y), vec4(o112[0], o112[1], o112[2], 0.0));
#if CX
imageStore(S120, ivec2(x, y), vec4(o120[0], o120[1], o120[2], 0.0));
#endif
}
#if CX
if (hl == 0u) [[unroll]] for (uint c = 0u; c < 3u; ++c) O112[t * 4u + c] = float16_t(o112[c]);
#endif
}
#if CX
barrier();
{ // the 2x2 merge: merged^T (CX x 16 rows), channel tiles m = wv, wv + 4, ...
const uvec2 mxy = mergeXY(col);
const int X = x0 / 2 + int(mxy.x), Y = y0 / 2 + int(mxy.y);
const uint o = (mrgOff + uint(Y * hx + X) * 32u) >> 2u;
for (uint m = wv; m < CX / 16u; m += 4u) {
ACC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_BM / 2u + 16u * m + 8u * hl + uint(i)]);
for (uint ks = 0u; ks < C / 8u; ++ks) {
B8 g0, g1; A8 w0, w1;
coopMatLoad(g0, MG, 8u * ks, C, CM);
coopMatLoad(g1, MG, 8u * ks + 4u, C, CM);
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 (wv == 3u && lane < 16u) { // surface 128: 2x2 sums of 0.25 * surface 112, rightwards first, then downwards
const uint mx = lane & 3u, my = lane >> 2u;
const int x = x0 + 2 * int(mx), y = y0 + 2 * int(my);
if (x >= 0 && y >= 0 && x < ex && y < ey) {
const uint t00 = (2u * my) * 8u + 2u * mx;
vec4 out_ = vec4(0.0);
[[unroll]] for (uint c = 0u; c < 3u; ++c) {
precise float16_t a00 = O112[t00 * 4u + c] * float16_t(0.25), a01 = O112[(t00 + 1u) * 4u + c] * float16_t(0.25);
precise float16_t a10 = O112[(t00 + 8u) * 4u + c] * float16_t(0.25), a11 = O112[(t00 + 9u) * 4u + c] * float16_t(0.25);
precise float16_t b0 = a00 + a01, b1 = a10 + a11;
precise float16_t cc = b0 + b1;
out_[c] = float(cc);
}
imageStore(S128, ivec2(x >> 1, y >> 1), out_);
}
}
}
#endif
#else
[[unroll]] for (int i = 0; i < 8; ++i) OUTF[t * 16u + 8u * hl + uint(i)] = float16_t(hv[i]);
barrier();
#define DEC_FILTER
#include "rr_decfilter.glsl"
#endif
}
engine/rr/swin_pack.py
+75
View File
@@ -0,0 +1,75 @@
#!/usr/bin/env python3
"""Pack a Swin stage's extracted weights (engine/rr/swin_ref.py --extract) for engine/rr/swin_m.comp.
usage: swin_pack.py STAGE WEIGHTS.npz OUT.bin OUT.h
enc1_pack.py's layout for any stage of engine/rr/swin_ref.STAGES: transposed products, every weight's columns in the
k order of the product that consumes its result (slot j holds natural channel inter(n)[j]); the expand's 4C columns
are four sub-positions of C channels, each permuted the same way. OUT.h also defines the stage's shape for the shader
(KIND, C, CX, NH).
"""
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from enc1_ref import inter # noqa: E402
from enc1_pack import a_tiles, b_tiles, cubic_table, act_table # noqa: E402
from swin_ref import STAGES # noqa: E402
def main():
stage, src, out_bin, out_h = sys.argv[1:5]
kind, C, CX, NH, _ = STAGES[stage]
w = dict(np.load(src))
I32, I96, IC = inter(32), inter(96), inter(C)
e = {}
if kind == 'dec':
I4C = [C * (j // C) + IC[j % C] for j in range(4 * C)]
e['we'], e['be'] = ('A', w['we'][:, I4C]), ('H', w['be'][I4C])
for h in range(NH):
e[f'wk{h}'] = ('A', w[f'wk{h}'][:, I32])
e[f'wv{h}'] = ('B', w[f'wv{h}'][:, I32])
e[f'wo{h}'] = ('A', w[f'wo{h}'][:, IC])
e[f'bias{h}'] = ('H', w[f'bias{h}'][np.ix_(inter(64), I96)].T)
for b in range(C // 8):
e[f'w1_{b}'] = ('A', w[f'w1_{b}'][:, I32])
e[f'b1_{b}'] = ('H', w[f'b1_{b}'][I32])
e[f'w2_{b}'] = ('A', w[f'w2_{b}'][:, IC])
if kind == 'enc' and CX:
e['wm'], e['bm'] = ('A', w['wm'][:, inter(CX)]), ('H', w['bm'][inter(CX)])
e['wh'], e['bh'] = ('A', w['wh']), ('H', w['bh'])
e['gain0'] = ('H', w['gain0'])
for k in ('table7', 'gain14', 'table16'):
e[k] = ('H', w[k][IC])
e['cubic'] = ('U', cubic_table())
e['act'] = ('Y', act_table())
blob, defines = bytearray(), []
for name, (kd, a) in e.items():
while len(blob) % 256:
blob.append(0)
off = len(blob)
if kd == 'A':
blob += a_tiles(a).tobytes()
elif kd == 'B':
blob += b_tiles(a).tobytes()
elif kd == 'Y':
blob += np.ascontiguousarray(a, np.uint8).tobytes()
elif kd == 'U':
blob += np.ascontiguousarray(a, np.uint16).tobytes()
else:
blob += np.ascontiguousarray(a, np.float16).tobytes()
defines.append(f'#define W_{name.upper()} {off}u // {kd} {a.shape}')
shape = [f'#define KIND {0 if kind == "enc" else 1}', f'#define C {C}', f'#define CX {CX}', f'#define NH {NH}']
if NH < 2:
defines.append('#define W_WK1 (W_WK0 + 1u)')
if C // 8 < 2:
defines.append('#define W_W1_1 (W_W1_0 + 1u)')
open(out_bin, 'wb').write(bytes(blob))
open(out_h, 'w').write(f'// generated by engine/rr/swin_pack.py for {stage}\n' + '\n'.join(shape + defines) + '\n')
print(f'{stage}: {len(blob)} bytes, {len(e)} entries')
if __name__ == '__main__':
main()
engine/rr/swin_ref.py
+334
View File
@@ -0,0 +1,334 @@
#!/usr/bin/env python3
"""Exact numpy reference of the denoiser's Swin stages enc1-enc5 and dec1-dec4 (cuda_dldn_engine_swin_*_kernel), one
8x8 window at a time, from weights and real inputs only. Generalises engine/rr/enc1_ref.py (encoders) and
engine/rr/dec1_ref.py (decoders); every stage was checked against its capture (main).
A stage has C channels (32 * KC), NH heads of 32 channels, NB = C / 8 MLP branches of C -> 32 -> C:
prologue encoders: the 96 key tokens from the arena (C / 32 planes); decoders: the expand of the 6 x 6 tokens
one level down (C_LO channels -> 4 sub-positions x C) + the skip token. RMS norm (input_norm: per lane word
the planes summed pairwise (pairs_tree), the two word halves, then lanes ^1, ^2 and the halves) times gain0; the
residual is the token + table7
heads Q/K and V projections over KC k32 steps, scores + bias, the cubic softmax, P V, out-projection onto the
residual, head after head
MLP RMS norm (rowsum_sq) times gain14, residual + table16, NB branches
outputs the tokens; encoders: the 2x2 merge (4C -> C_M) and the 64 -> 16 head with enc1's filter (surfaces 112,
120, 128); decoders: the head with dec0's filter (surface 112)
usage: swin_ref.py CAPTURE STAGE FRAME WEIGHTS.npz [--blocks X,Y ...] (verifies against the capture)
swin_ref.py CAPTURE STAGE FRAME OUT.npz --extract DUMP.npz (weights from a phase dump)
"""
import argparse
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rrmath import e4m3, f16, q8 # noqa: E402
from enc1_ref import (inter, mm, softmax_cubic, act, mirror, add, rowsum_sq, pairs_tree, filter_head, # noqa: E402
MERGE_ORDER, H, F)
from dec0_ref import keys # noqa: E402
from dec1_ref import filter112 # noqa: E402
# kind, C, C_LO (decoders: channels one level down) or C_M (encoders: merged channels), heads, warps
STAGES = {
'enc1': ('enc', 64, 96, 2, 1), 'enc2': ('enc', 96, 128, 4, 4), 'enc3': ('enc', 128, 160, 4, 4),
'enc4': ('enc', 160, 160, 8, 8), 'enc5': ('enc', 160, 0, 8, 8),
'dec4': ('dec', 160, 160, 8, 8), 'dec3': ('dec', 128, 160, 4, 4), 'dec2': ('dec', 96, 128, 4, 4),
'dec1': ('dec', 64, 96, 2, 2),
}
def input_norm(xn, gain):
"""xn: one token, natural channels (f16). NVIDIA's tree over the arena order; returns the normalised token
(natural, f16)"""
c = len(xn)
xa = xn[inter(c)]
sq = (xa * xa).astype(H)
def pair(p, w):
return add(sq[32 * p + 4 * w:32 * p + 4 * w + 2], sq[32 * p + 4 * w + 2:32 * p + 4 * w + 4])
fold = pairs_tree # (p0 + p1) + p2, (p0 + p1) + (p2 + p3), ((p0 + p1) + (p2 + p3)) + p4
h = [add(fold([pair(p, t) for p in range(c // 32)]), fold([pair(p, t + 4) for p in range(c // 32)]))
for t in range(4)]
h01 = [add(h[t], h[t ^ 1]) for t in range(4)]
h2 = add(h01[0], h01[2])
s = H(H(h2[0] + h2[1]) + H(2.0 ** -13))
r = H(F(1) / np.sqrt(F(s)))
return (xn * (r * gain).astype(H)).astype(H)
def block(stage, w, x0, y0, ex, ey, inp):
"""the window: returns (tokens 64 x C codes in arena order, merged 16 x C_M codes or None, head 64 x 16 f16).
inp: 'tok' [C/32][ey][ex][32] (encoders) or 'lo' [C_LO/32][ey/2][ex/2][32] and 'skip' [C/32][ey][ex][32]"""
kind, C, CX, NH, _ = STAGES[stage]
IC = inter(C)
UN = [IC.index(i) for i in range(C)]
xs = []
if kind == 'dec':
hx, hy = ex // 2, ey // 2
X0, Y0 = x0 // 2 - 1, y0 // 2 - 1
low = np.zeros((48, CX), np.uint8)
for r in range(36):
X, Y = mirror(X0 + r % 6, hx), mirror(Y0 + r // 6, hy)
low[r] = np.concatenate([inp['lo'][p, Y, X] for p in range(CX // 32)])
Ex = q8(mm(low, w['we'], np.repeat(w['be'][None, :], 48, 0)))
for (x, y) in keys(x0, y0):
xx, yy = mirror(x, ex), mirror(y, ey)
r, s = (y // 2 - Y0) * 6 + (x // 2 - X0), 2 * (y & 1) + (x & 1)
sk = np.concatenate([inp['skip'][p, yy, xx] for p in range(C // 32)])[UN]
xs.append((e4m3(Ex[r, C * s:C * s + C]).astype(H) + e4m3(sk).astype(H)).astype(H))
else:
for (x, y) in keys(x0, y0):
xx, yy = mirror(x, ex), mirror(y, ey)
xs.append(e4m3(np.concatenate([inp['tok'][p, yy, xx] for p in range(C // 32)])[UN]).astype(H))
X = np.array([q8(input_norm(x, w['gain0']))[IC] for x in xs])
R = add(np.array(xs[:64]), w['table7'][None, :])
I32, I96 = inter(32), inter(96)
for h in range(NH):
qk, v = np.zeros((96, 32)), np.zeros((96, 32))
for c in range(C // 32):
qk = mm(X[:, 32 * c:32 * c + 32], w[f'wk{h}'][32 * c:32 * c + 32], qk)
v = mm(X[:, 32 * c:32 * c + 32], w[f'wv{h}'][32 * c:32 * c + 32], v)
Q = q8(qk)[:, I32]
P = softmax_cubic(mm(Q[:64], Q.T, w[f'bias{h}']))
O = mm(q8(P)[:, I96], q8(v)[I96], np.zeros((64, 32)))
R = mm(q8(O)[:, I32], w[f'wo{h}'], R)
N = q8((R * (rowsum_sq(R)[:, None] * w['gain14'][None, :]).astype(H)).astype(H))[:, IC]
acc = add(R, w['table16'][None, :])
for b in range(C // 8):
hid = act(mm(N, w[f'w1_{b}'], np.repeat(w[f'b1_{b}'][None, :], 64, 0)))
acc = mm(q8(hid)[:, I32], w[f'w2_{b}'], acc)
tokens = q8(acc)[:, IC]
merged = None
if kind == 'enc' and CX:
Mg = np.zeros((16, 4 * C), np.uint8)
for r, (mx, my) in enumerate(MERGE_ORDER):
for blk, (dx, dy) in enumerate([(0, 0), (1, 0), (0, 1), (1, 1)]):
Mg[r, C * blk:C * blk + C] = tokens[(2 * my + dy) * 8 + 2 * mx + dx]
merged = q8(mm(Mg, w['wm'], np.repeat(w['bm'][None, :], 16, 0)))[:, inter(CX)]
head = mm(tokens, w['wh'], np.repeat(w['bh'][None, :], 64, 0))
return tokens, merged, head
def phase_map(stage):
"""dump phase numbers: projections [h][c] = (QK, V), bias[h] (scores phase, wave), out[h], mlp up / down, merge,
head, expand"""
kind, C, CX, NH, waves = STAGES[stage]
KC = C // 32
m = {'proj': {}, 'bias': {}, 'out': {}}
if waves == 1: # heads in sequence: per head 2 KC projections, scores, P V, out-projection
p = 1
for h in range(NH):
for c in range(KC):
m['proj'][h, c] = (p, p + 1, 0)
p += 2
m['bias'][h] = (p, 0)
m['out'][h] = p + 2
p += 3
else: # one head per warp
p = 2 if kind == 'dec' else 1
for h in range(NH):
for c in range(KC):
m['proj'][h, c] = (p + 2 * c, p + 2 * c + 1, h)
p += 2 * KC
for h in range(NH):
m['bias'][h] = (p, h)
m['out'][h] = p + 2 + h
p += 2 + NH
m['mlp'] = [(p + 2 * b, p + 2 * b + 1) for b in range(C // 8)]
p += 2 * (C // 8)
if kind == 'enc' and CX:
m['merge'] = p
p += 1
m['head'] = p
return m
def read_b(W, base, K):
"""an FP8 B matrix (K x 32) as NVIDIA stores it in the weight image: 4-byte k groups, offset 512 (kg >> 3) +
16 (kg & 3) + 4 ((kg >> 2) & 1) + 64 (n & 7) + 8 ((n >> 3) & 1) + 16 K (n >> 4)"""
b = np.zeros((K, 32), np.uint8)
for kg in range(K // 4):
for n in range(32):
o = base + 512 * (kg >> 3) + 16 * (kg & 3) + 4 * ((kg >> 2) & 1) + 64 * (n & 7) + 8 * ((n >> 3) & 1) + 16 * K * (n >> 4)
b[4 * kg:4 * kg + 4, n] = W[o:o + 4]
return b
def merge_weights(z, W, p, waves, C, CX):
"""the merge (4C -> CX) of a multi-warp encoder: warp w computes the 32-column chunks w, w + warps, ...; the dump
holds each warp's last chunk, so the chunks come from the weight image (located through warp 1's, chunk 1) and the
dumped ones check the layout"""
K, size = 4 * C, 4 * C * 32
W32 = W[:len(W) // 4 * 4].view(np.uint32)
first = z[f'p{p}_w1_B'][0:4, 0].copy().view(np.uint32)[0]
pos1 = int(next(i for i in np.nonzero(W32 == first)[0] * 4 if np.array_equal(read_b(W, int(i), K), z[f'p{p}_w1_B'])))
base0 = pos1 - size
chunks = [read_b(W, base0 + size * j, K) for j in range(CX // 32)]
for v in range(waves):
last = max(j for j in range(CX // 32) if j % waves == v) if any(j % waves == v for j in range(CX // 32)) else None
if last is not None:
assert np.array_equal(chunks[last], z[f'p{p}_w{v}_B']), (v, last)
b1 = f16(z[f'p{p}_w1_C']).astype(H)[0].tobytes()
bpos1 = W.tobytes().find(b1)
bias = np.concatenate([W[bpos1 + 64 * (j - 1):bpos1 + 64 * j].view(H) for j in range(CX // 32)])
for v in range(waves):
last = max((j for j in range(CX // 32) if j % waves == v), default=None)
if last is not None:
assert np.array_equal(bias[32 * last:32 * last + 32].view(np.uint16), z[f'p{p}_w{v}_C'][0]), ('bias', v)
return np.hstack(chunks), bias
def extract(stage, dump, weights, out, inp, x0, y0, ex, ey):
"""weights from the phase dump; gains and tables are found in the weight image by matching the dump"""
kind, C, CX, NH, waves = STAGES[stage]
z = np.load(dump)
W = np.frombuffer(weights, np.uint8)
def B(p, wv=0):
return z[f'p{p}_w{wv}_B']
def Cm(p, wv=0):
return f16(z[f'p{p}_w{wv}_C']).astype(H)
pm = phase_map(stage)
w = {}
if kind == 'dec':
w['we'] = np.hstack([B(1, v) for v in range(waves)])
w['be'] = np.concatenate([Cm(1, v)[0] for v in range(waves)])
for h in range(NH):
w[f'wk{h}'] = np.vstack([B(pm['proj'][h, c][0], pm['proj'][h, c][2]) for c in range(C // 32)])
w[f'wv{h}'] = np.vstack([B(pm['proj'][h, c][1], pm['proj'][h, c][2]) for c in range(C // 32)])
w[f'bias{h}'] = Cm(*pm['bias'][h])
w[f'wo{h}'] = B(pm['out'][h])
for b, (up, dn) in enumerate(pm['mlp']):
w[f'w1_{b}'], w[f'b1_{b}'], w[f'w2_{b}'] = B(up), Cm(up)[0], B(dn)
if 'merge' in pm:
mw = min(waves, 4) # the merge runs on (at most) four warps
if waves > 1 and CX // 32 > mw: # a warp computes several chunks: read them from the weight image
w['wm'], w['bm'] = merge_weights(z, W, pm['merge'], mw, C, CX)
elif waves > 1:
w['wm'] = np.hstack([B(pm['merge'], v) for v in range(CX // 32)])
w['bm'] = np.concatenate([Cm(pm['merge'], v)[0] for v in range(CX // 32)])
else:
w['wm'], w['bm'] = B(pm['merge']), Cm(pm['merge'])[0]
w['wh'], w['bh'] = B(pm['head']), Cm(pm['head'])[0]
# the four f16 vectors: candidates at every even offset, checked against the first token rows of the dump
tables = [W[o:o + 2 * C].view(H) for o in range(0, len(W) - 2 * C + 1, 2)]
w0 = {k: v for k, v in w.items()}
# gain0: the key tokens' norm against the first projection's A (window row 0..)
xs = key_tokens(stage, w0, x0, y0, ex, ey, inp)
IC = inter(C)
A1 = np.hstack([z[f'p{pm["proj"][0, c][0]}_w{pm["proj"][0, c][2]}_A'] for c in range(C // 32)])
w['gain0'] = next(t for t in tables if np.array_equal(q8(input_norm(xs[0], t))[IC], A1[0]) and
np.array_equal(q8(input_norm(xs[5], t))[IC], A1[5]))
p_out = pm['out'][0]
C0 = Cm(p_out, 0)
rows = np.array(xs[:64])
w['table7'] = next(t for t in tables if np.array_equal(add(rows[:len(C0)][:2], t[None]).view(np.uint16), C0[:2].view(np.uint16)))
up, dn = pm['mlp'][0]
Rt = f16(z[f'p{pm["out"][NH - 1]}_w0_D']).astype(H)
Aup, Cdn = z[f'p{up}_w0_A'], Cm(dn)
w['gain14'] = next(t for t in tables if np.array_equal(
q8((Rt[:2] * (rowsum_sq(Rt[:2])[:, None] * t[None]).astype(H)).astype(H))[:, IC], Aup[:2]))
w['table16'] = next(t for t in tables if np.array_equal(add(Rt[:2], t[None]).view(np.uint16), Cdn[:2].view(np.uint16)))
np.savez_compressed(out, **w)
def key_tokens(stage, w, x0, y0, ex, ey, inp):
kind, C, CX, NH, _ = STAGES[stage]
IC = inter(C)
UN = [IC.index(i) for i in range(C)]
xs = []
if kind == 'dec':
hx, hy = ex // 2, ey // 2
X0, Y0 = x0 // 2 - 1, y0 // 2 - 1
low = np.zeros((48, CX), np.uint8)
for r in range(36):
X, Y = mirror(X0 + r % 6, hx), mirror(Y0 + r // 6, hy)
low[r] = np.concatenate([inp['lo'][p, Y, X] for p in range(CX // 32)])
Ex = q8(mm(low, w['we'], np.repeat(w['be'][None, :], 48, 0)))
for (x, y) in keys(x0, y0):
xx, yy = mirror(x, ex), mirror(y, ey)
if kind == 'dec':
r, s = (y // 2 - Y0) * 6 + (x // 2 - X0), 2 * (y & 1) + (x & 1)
sk = np.concatenate([inp['skip'][p, yy, xx] for p in range(C // 32)])[UN]
xs.append((e4m3(Ex[r, C * s:C * s + C]).astype(H) + e4m3(sk).astype(H)).astype(H))
else:
xs.append(e4m3(np.concatenate([inp['tok'][p, yy, xx] for p in range(C // 32)])[UN]).astype(H))
return xs
def load(cap, stage, frame):
from rr_capture import Capture
c = Capture(os.path.expanduser(cap))
l = c.find(f'swin_{stage}', frame)
kind, C, CX, NH, _ = STAGES[stage]
ex, ey = (int(v) for v in np.frombuffer(l.args[:8], np.uint32))
sx, sy = (int(v) for v in np.frombuffer(l.args[80:84], np.uint16))
ptr = {p.offset: p.delta for p in l.pointers}
pre = np.frombuffer(l.pre(l.arena_alloc), np.uint8)
post = np.frombuffer(l.post_buffer(c, l.arena_alloc), np.uint8)
inp = {}
if kind == 'dec':
inp['lo'] = pre[ptr[8]:ptr[8] + CX * (ex // 2) * (ey // 2)].reshape(CX // 32, ey // 2, ex // 2, 32)
inp['skip'] = pre[ptr[24]:ptr[24] + C * ex * ey].reshape(C // 32, ey, ex, 32)
else:
inp['tok'] = pre[ptr[8]:ptr[8] + C * ex * ey].reshape(C // 32, ey, ex, 32)
outs = {'tok': post[ptr[48]:ptr[48] + C * ex * ey].reshape(C // 32, ey, ex, 32)}
if kind == 'enc' and CX:
outs['mrg'] = post[ptr[56]:ptr[56] + CX * (ex // 2) * (ey // 2)].reshape(CX // 32, ey // 2, ex // 2, 32)
img = {i.offset: i for i in l.images}
def surf(off):
i = img[off]
return np.frombuffer(i.path.read_bytes(), np.float16).reshape(i.height, i.width, 4)
tex = {k: surf(k) for k in img}
return l, ex, ey, sx, sy, inp, outs, tex
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('stage')
ap.add_argument('frame', type=int)
ap.add_argument('weights')
ap.add_argument('--extract', metavar='DUMP')
ap.add_argument('--dump-block', default='1,1', help='the block the dump was taken of (for --extract)')
ap.add_argument('--blocks', nargs='*', default=['1,1'])
o = ap.parse_args()
l, ex, ey, sx, sy, inp, outs, tex = load(o.capture, o.stage, o.frame)
kind, C, CX, NH, _ = STAGES[o.stage]
if o.extract:
bx, by = (int(v) for v in o.dump_block.split(','))
extract(o.stage, o.extract, l.weights, o.weights, inp, 8 * bx - sx, 8 * by - sy, ex, ey)
return
w = dict(np.load(o.weights))
for blk in o.blocks:
bx, by = (int(v) for v in blk.split(','))
x0, y0 = 8 * bx - sx, 8 * by - sy
tokens, merged, head = block(o.stage, w, x0, y0, ex, ey, inp)
inside = [t for t in range(64) if 0 <= x0 + t % 8 < ex and 0 <= y0 + t // 8 < ey]
tk = sum(np.array_equal(tokens[t], np.concatenate([outs['tok'][p, y0 + t // 8, x0 + t % 8] for p in range(C // 32)]))
for t in inside)
msg = f'block {bx},{by}: tokens {tk}/{len(inside)}'
if merged is not None:
mg = sum(np.array_equal(merged[r], np.concatenate([outs['mrg'][p, y0 // 2 + my, x0 // 2 + mx] for p in range(CX // 32)]))
for r, (mx, my) in enumerate(MERGE_ORDER) if 0 <= x0 // 2 + mx < ex // 2 and 0 <= y0 // 2 + my < ey // 2)
msg += f', merged {mg}'
if kind == 'enc' and not CX: # the bottleneck: the ungated filter to surface 112 only
o112, _ = filter_head(head, tex[88], x0, y0, ex, ey, ungated=True)
ok = sum(np.array_equal(o112[t], tex[112][y0 + t // 8, x0 + t % 8, :3]) for t in inside)
elif kind == 'enc':
o112, o120 = filter_head(head, tex[88], x0, y0, ex, ey)
ok = sum(np.array_equal(o112[t], tex[112][y0 + t // 8, x0 + t % 8, :3]) and
np.array_equal(o120[t], tex[120][y0 + t // 8, x0 + t % 8, :3]) for t in inside)
else:
s112 = filter112(head, tex[88], tex[96], tex[104], x0, y0, ex, ey)
ok = sum(np.array_equal(s112[t], tex[112][y0 + t // 8, x0 + t % 8, :3]) for t in inside)
print(msg + f', surfaces {ok}/{len(inside)}')
if __name__ == '__main__':
main()
engine/rr/test_layer.py
Executable
+207
View File
@@ -0,0 +1,207 @@
#!/usr/bin/env python3
"""Run one Ray Reconstruction stage shader on a captured launch and compare it with NVIDIA's result.
usage: test_layer.py CAPTURE KERNEL FRAME SHADER.spv WORK [--kvk build/rr/kvk] [--time N] [--subgroup 32|64]
KERNEL is a substring of the denoiser kernel's name (e.g. swin_enc1, hkpn_output_kernel_transformer), FRAME the
1-based occurrence in the capture. The shader follows the RR stage contract (engine/rr/README.md):
binding 0 the parameter block: NVIDIA's argument block with every pointer into the plane arena
replaced by its byte offset in the arena, the weight pointer and every texture or surface
handle by zero (stage_params below; the runtime builds the same block)
binding 1 the weight image exactly as NGX uploads it (argument 40), read-only
binding 2 the plane arena, read/write
bindings 3, 4, ... the textures in ascending argument offset, then the surfaces in ascending argument offset
Every surface starts zeroed and the arena starts as it was before NVIDIA's launch; afterwards each surface and the
whole arena are compared with what NVIDIA's launch left. Exit status 0 only when everything is byte-identical.
"""
import argparse
import json
import os
import struct
import subprocess
import sys
from pathlib import Path
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rr_capture import Capture # noqa: E402
# GLSL image formats of the stage shaders' surfaces -> (channels, kvk bits)
SURFACE_FORMATS = {'rgba16f': (4, 16), 'rg16f': (2, 16), 'r16f': (1, 16), 'rgba8_snorm': (4, 9), 'rgba8': (4, 8)}
def stage_params(launch):
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
struct.pack_into('<Q', args, p.offset, value)
# texture and surface handles: 1 where NGX bound one (the kernels test optional inputs' handles for null, e.g.
# enc0's specular hit distance at argument 248), 0 elsewhere
for image in launch.images:
struct.pack_into('<Q', args, image.offset, 1)
return bytes(args)
def e4m3_table():
t = np.zeros(256)
for c in range(256):
s, e, m = c >> 7, (c >> 3) & 15, c & 7
v = np.nan if (e == 15 and m == 7) else (m * 2.0 ** -9 if e == 0 else (1 + m / 8) * 2.0 ** (e - 7))
t[c] = -v if s else v
return t
def psnr(a, b):
finite = np.isfinite(a) & np.isfinite(b)
if not finite.any():
return float('nan')
err = np.mean((a[finite] - b[finite]) ** 2)
peak = max(np.max(np.abs(b[finite])), 1e-6)
return float('inf') if err == 0 else 10 * np.log10(peak * peak / err)
def compare_image(name, image, got, ref, ignore_channels=()):
"""IGNORE_CHANNELS: channels NVIDIA's kernel fills from an undefined register (left out of the comparison)"""
if image.bits == 16:
a, b = np.frombuffer(got, np.float16).astype(np.float64), np.frombuffer(ref, np.float16).astype(np.float64)
elif image.bits == 32:
a, b = np.frombuffer(got, np.float32).astype(np.float64), np.frombuffer(ref, np.float32).astype(np.float64)
else:
dt = np.int8 if image.bits == 9 else np.uint8
a, b = np.frombuffer(got, dt).astype(np.float64), np.frombuffer(ref, dt).astype(np.float64)
keep = np.ones(len(a), bool)
for c in ignore_channels:
keep[c::image.channels] = False
a, b = a[keep], b[keep]
element_bytes = len(got) // len(keep)
same = (np.frombuffer(got, np.uint8).reshape(-1, element_bytes) == np.frombuffer(ref, np.uint8).reshape(-1, element_bytes))[keep]
elements = len(a)
equal = np.sum(a == b) + np.sum(np.isnan(a) & np.isnan(b))
worst = np.nanmax(np.abs(a - b)) if elements else 0
print(f' {name}: {image.width}x{image.height}x{image.channels} fmt{image.format}'
f'{" (channels " + ",".join(map(str, ignore_channels)) + " ignored)" if ignore_channels else ""}: '
f'{100.0 * equal / elements:.4f}% identical values, max |diff| {worst:.6g}, PSNR {psnr(a, b):.2f} dB, '
f'nonfinite ours/ref {int(np.sum(~np.isfinite(a)))}/{int(np.sum(~np.isfinite(b)))}')
return bool(same.all())
def compare_arena(got, before, ref):
g, p, r = (np.frombuffer(x, np.uint8) for x in (got, before, ref))
written = (r != p) | (g != p)
differ = g != r
print(f' arena: {int(written.sum())} bytes written by either side, {int(differ.sum())} differ', end='')
if differ.any():
idx = np.nonzero(differ)[0]
t = e4m3_table()
dv = np.abs(t[g[idx]] - t[r[idx]])
print(f' (first at byte {idx[0]:#x}, last at {idx[-1]:#x}; as e4m3 values: max |diff| {np.nanmax(dv):.6g}, '
f'{int(np.sum(dv <= 0.0625 * np.abs(t[r[idx]]) + 1e-12))} within one e4m3 step)')
else:
print()
return not differ.any()
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('kernel')
ap.add_argument('frame', type=int)
ap.add_argument('shader')
ap.add_argument('work')
ap.add_argument('--kvk', default=str(Path(__file__).resolve().parents[2] / 'build/rr/kvk'))
ap.add_argument('--time', type=int, default=1)
ap.add_argument('--subgroup', type=int, default=32)
ap.add_argument('--weights', help='bind this file at binding 1 instead of NGX\'s weight image (native stages)')
ap.add_argument('--grid', help='X,Y workgroups (default: the launch grid)')
ap.add_argument('--ignore', action='append', default=[], metavar='SURFACE:CHANNEL',
help='leave a channel out of the comparison, e.g. 112:3 for a value NVIDIA stores from an '
'undefined register (repeatable)')
o = ap.parse_args()
cap = Capture(o.capture)
launch = [l for l in cap.launches if o.kernel in l.kernel and 'diamond' not in l.kernel or
('diamond' in o.kernel and o.kernel in l.kernel)][o.frame - 1]
work = Path(o.work)
work.mkdir(parents=True, exist_ok=True)
print(f'{launch.kernel} (launch {launch.sequence}, grid {launch.grid}, block {launch.block})')
(work / 'params.bin').write_bytes(stage_params(launch))
(work / 'weights.bin').write_bytes(Path(o.weights).read_bytes() if o.weights else launch.weights)
arena = launch.arena_alloc
arena_before = launch.pre(arena)
(work / 'arena.bin').write_bytes(arena_before)
gx, gy = (int(v) for v in o.grid.split(',')) if o.grid else launch.grid[:2]
lines = [f'param {work}/params.bin', f'grid {gx} {gy}']
# Bind in the shader's order (its ptx2rr.py sidecar, SHADER.json: textures then surfaces by argument offset).
# A handle the shader reads but NGX left null gets a zeroed 1x1 image.
sidecar = Path(o.shader).with_suffix('.json')
if sidecar.exists():
info = json.loads(sidecar.read_text())
order = [('texture', off, None) for off in info['textures']] + [('surface', off, f) for off, f in info['surfaces']]
else:
order = [(i.kind, i.offset, None) for i in launch.images]
by_key = {(i.kind, i.offset): i for i in launch.images}
(work / 'zero.bin').write_bytes(bytes(64))
for kind, offset, fmt in order:
image = by_key.get((kind, offset))
if image is None:
ch, bits = SURFACE_FORMATS.get(fmt, (4, 16))
lines.append(f'tex {offset} {work}/zero.bin 1 1 4 16 nearest' if kind == 'texture' else
f'surf {offset} 1 1 {ch} {bits} - {work}/dummy{offset}.bin')
elif kind == 'texture':
lines.append(f'tex {image.offset} {image.path} {image.width} {image.height} {image.channels} {image.bits} '
f'{"linear" if image.linear else "nearest"}')
else:
lines.append(f'surf {image.offset} {image.width} {image.height} {image.channels} {image.bits} - '
f'{work}/surf{image.offset}.bin')
n_images = len(order)
lines.append(f'lut {work}/weights.bin')
lines.append(f'inout {work}/arena.bin {work}/arena.out.bin')
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]
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 ''))
(work / 'manifest').write_text('\n'.join(lines) + '\n')
r = subprocess.run([o.kvk, o.shader, str(work / 'manifest'), str(o.time), str(o.subgroup)],
capture_output=True, text=True)
sys.stdout.write(r.stdout)
if r.returncode != 0:
sys.stderr.write(r.stderr)
sys.exit(f'kvk failed ({r.returncode})')
identical = True
bound = {(kind, offset) for kind, offset, _ in order}
for image in launch.images:
if image.kind != 'surface':
continue
if ('surface', image.offset) not in bound:
print(f' surface {image.offset}: not written by this kernel')
continue
got = (work / f'surf{image.offset}.bin').read_bytes()
waived = [int(s.split(':')[1]) for s in o.ignore if int(s.split(':')[0]) == image.offset]
identical &= compare_image(f'surface {image.offset}', image, got, image.path.read_bytes(), waived)
if arena is not None:
identical &= compare_arena((work / 'arena.out.bin').read_bytes(), arena_before,
launch.post_buffer(cap, arena))
if oob:
r = struct.unpack('<9I', (work / 'oob.bin').read_bytes()[:36])
for k, name in enumerate(('weight reads', 'arena reads', 'arena writes')):
print(f' out of range {name}: {r[3 * k]}' + (f' (bytes {r[3 * k + 1]:#x}..{r[3 * k + 2]:#x}, size '
f'{len(launch.weights) if k == 0 else len(arena_before):#x})' if r[3 * k] else ''))
print('IDENTICAL' if identical else 'DIFFERENT')
sys.exit(0 if identical else 1)
if __name__ == '__main__':
main()
engine/rr/zluda_replay.py
+84
View File
@@ -0,0 +1,84 @@
#!/usr/bin/env python3
"""Replay one captured denoiser launch through ZLUDA on Linux (tools/replay_rr_launch.cpp) and compare it with the
capture, like test_layer.py does for a stage shader.
usage: zluda_replay.py CAPTURE KERNEL FRAME WORK [--ptx MODULE] [--zluda DIR]
KERNEL and FRAME select the launch as in test_layer.py. MODULE defaults to the corpus module holding the kernel
(~/.cache/d4r-rr-corpus); pass an edited copy to add debug stores. ZLUDA's own settings (D4R_ZLUDA_WMMA,
D4R_ZLUDA_WMMA_FP8, D4R_ZLUDA_WMMA_FP8_NATIVE, ...) are taken from the environment. Outputs land in WORK
(surfOFFSET.bin, arena.out.bin) so the two replays can be compared with each other too.
"""
import argparse
import os
import subprocess
import sys
from pathlib import Path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rr_capture import Capture # noqa: E402
from test_layer import compare_arena, compare_image # noqa: E402
ROOT = Path(__file__).resolve().parents[2]
def main():
ap = argparse.ArgumentParser()
ap.add_argument('capture')
ap.add_argument('kernel')
ap.add_argument('frame', type=int)
ap.add_argument('work')
ap.add_argument('--ptx')
ap.add_argument('--zluda', default=os.path.expanduser('~/.cache/d4r-zluda-rr'))
ap.add_argument('--rocm', default='/opt/rocm/lib')
ap.add_argument('--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]
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'))
if f'.entry {launch.kernel}(' in p.read_text(errors='replace'))
(work / 'args.bin').write_bytes(launch.args)
arena = launch.arena_alloc
lines = [f'args {work}/args.bin', 'grid ' + ' '.join(map(str, launch.grid)),
'block ' + ' '.join(map(str, launch.block))]
if o.pad:
(work / 'arena-pad.bin').write_bytes(launch.pre(arena) + bytes(o.pad))
for p in launch.pointers:
out = f' {work}/arena.out.bin' if p.alloc == arena else ''
source = f'{work}/arena-pad.bin' if o.pad and p.alloc == arena else f'{launch.directory}/alloc-{p.alloc}.bin'
lines.append(f'buf {p.offset} {source} {p.delta}{out}')
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
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
lines.append(f'surf {image.offset} {image.width} {image.height} {image.channels} {fmt} '
f'{work}/surf{image.offset}.bin')
(work / 'manifest').write_text('\n'.join(lines) + '\n')
env = dict(os.environ, LD_LIBRARY_PATH=f'{o.zluda}:{o.rocm}:' + os.environ.get('LD_LIBRARY_PATH', ''))
print(f'{launch.kernel} (launch {launch.sequence}) through ZLUDA {o.zluda}')
r = subprocess.run([str(ROOT / 'build/replay_rr_launch'), module, launch.kernel, str(work / 'manifest')],
env=env, capture_output=True, text=True)
sys.stdout.write(r.stdout)
if r.returncode:
sys.stderr.write(r.stderr[-4000:])
sys.exit('replay failed')
same = True
for image in launch.images:
if image.kind == 'surface':
same &= compare_image(f'surface {image.offset}', image, (work / f'surf{image.offset}.bin').read_bytes(),
image.path.read_bytes())
if arena is not None and not o.pad:
same &= compare_arena((work / 'arena.out.bin').read_bytes(), launch.pre(arena), launch.post_buffer(cap, arena))
print('IDENTICAL' if same else 'DIFFERENT')
if __name__ == '__main__':
main()