diff --git a/.gitignore b/.gitignore index 9c628ff..6174540 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,8 @@ *.cubin *.fatbin *.nvngx +# NVIDIA's NGX/DLSS SDK headers (the probes in tools/ build against the user's own copy) +nvsdk_ngx*.h # Local captures, dumps, traces, and build products /captures/ @@ -35,4 +37,11 @@ videos/ /kernels/out/ *.hsaco *.ptx -/kernels/extracted/ \ No newline at end of file +/kernels/extracted/ +# Native engine models: weights and tables read from NVIDIA's DLSS, and shaders compiled next to them +weights.bin +offsets.bin +lut.bin +hipnet.bin +model.json +*.spv diff --git a/NOTICE b/NOTICE index 88cf5d3..1837ac3 100644 --- a/NOTICE +++ b/NOTICE @@ -1,5 +1,5 @@ d4r -Copyright 2026 countervolts +Copyright 2026 countervolts and the d4r contributors This product is licensed under the GNU General Public License, version 3 only (GPL-3.0-only; see LICENSE), except for the files @@ -27,7 +27,8 @@ This repository contains no NVIDIA software. The NVIDIA DLSS feature library (nv the NGX core (_nvngx.dll) are NVIDIA's; you supply your own copies, and their use is governed by NVIDIA's license terms. At build time, kernels/tools/extract_dlss_ptx.py reads PTX from your copy of nvngx_dlss.dll into kernels/extracted/, which is not part of the repository and must not be -redistributed. +redistributed. The same holds for the model files the native engine's tools write (engine/README.md): +weights.bin, offsets.bin, lut.bin, hipnet.bin and model.json hold values read from NVIDIA's DLSS. NVIDIA and DLSS are trademarks of NVIDIA Corporation. AMD, Radeon and FSR are trademarks of Advanced Micro Devices, Inc. This project is not affiliated with or endorsed by either company. diff --git a/README.md b/README.md index 5c37431..3766e28 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,13 @@ See [supported games](SUPPORTED_GAMES.md) for the tested games and DLSS models. The development shim also has [experimental native Vulkan support](docs/vulkan.md). +A [standalone native inference engine](engine/README.md) runs presets K and M from weights extracted from your own +`nvngx_dlss.dll`, with d4r's own Vulkan input, network and reconstruction shaders (optionally the native HIP network +layers) and GPU-resident temporal history. The game shim uses the CUDA/ZLUDA path described below by default; +`[Engine] Enabled = true` (`D4R_ENGINE=1`) records supported K and M evaluations with the native engine into the +game's own command list instead and hands anything else to the CUDA path (see the engine's README). Presets E and L +remain CUDA-only. + > **Not affiliated with NVIDIA or AMD.** @@ -160,6 +167,7 @@ Gameplay recordings of each DLSS model at the Quality, Performance and Ultra Per |---|---| | `tools/` | the NGX shim, the Wine CUDA bridge, a D3D12 DLSS harness, probes and kernel replay tools | | `kernels/` | native RDNA3/RDNA4 kernels (`k/` DLSS 4, `m/` DLSS 4.5, `tex/` texture-kernel parts, `common/` WMMA layouts), their build script, validation tools and numpy reference models | +| `engine/` | standalone weight-based preset K and M inference runtime (Vulkan shaders, optional native HIP network), model packaging and full-frame validation | | `patches/` | ZLUDA and vkd3d-proton patches | | `scripts/` | build, install, launch and probe scripts | | `config/` | the default `d4r.ini` for the developer launcher | diff --git a/config/d4r.ini.default b/config/d4r.ini.default index e37e406..f1f74b4 100644 --- a/config/d4r.ini.default +++ b/config/d4r.ini.default @@ -97,6 +97,29 @@ ElideNgxSync = true EvalSync = false MarkerPollUs = 200 +[Engine] +; Experimental. Runs DLSS 4 and 4.5 (models K and M) on d4r's own Vulkan inference engine instead of NVIDIA's +; kernels through ZLUDA: the whole evaluation is recorded into the game's own command list, the game shows the +; result of the frame it just rendered, and FrameAge and the [Interop] hand-off settings do not apply. +; Needs VkD3DDir, a driver with VK_KHR_cooperative_matrix (Mesa RADV on RDNA3/RDNA4) and the model folders +; (RuntimeDir/engine/k, k-ldr, m, m-ldr; scripts/install_d4r_runtime.sh). A game or setting the engine does not +; cover uses the normal path, with the reason in d4r_nvngx.log: other models than K and M, display-resolution +; motion vectors with M, or a K model built before the current engine/compile_k.py for display-resolution +; motion, non-inverted depth or exposure supplied by the game. Details and measurements: engine/README.md. +; (D4R_ENGINE) +Enabled = false +; The model K folder for games with the HDR flag: the folder engine/compile_k.py writes (weights extracted +; from your nvngx_dlss.dll plus d4r's shaders). auto = RuntimeDir/engine/k, where scripts/install_d4r_runtime.sh +; puts D4R_ENGINE_MODEL. The other folders: D4R_ENGINE_MODEL_LDR_DIR, D4R_ENGINE_MODEL_M_DIR and +; D4R_ENGINE_MODEL_M_LDR_DIR. (D4R_ENGINE_MODEL_DIR) +ModelDir = auto +; Where the network between the engine's Vulkan stages runs (K and M). hip = d4r's native HIP layer kernels +; (M: about 1.7x faster than the shader, 3.8 instead of 5.5 ms per frame at 1080p -> 4K on an RX 9070 XT); +; the game's command list is split around the network like a split frame, so it needs VkD3DDir, the CUDA +; bridge and a model built with --hip. vulkan = the engine's cooperative-matrix shader, everything in one +; command list. auto and hip = hip when available, else vulkan (the log says why). (D4R_ENGINE_NETWORK) +Network = auto + [Paths] ; auto = defaults of scripts/d4r_proton_env.sh RuntimeDir = auto diff --git a/docs/architecture.md b/docs/architecture.md index 5b99040..3cc797f 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -72,7 +72,7 @@ The release zip (`scripts/package_release.sh`) is unpacked into the folder that | `d4r\zluda\libcuda.so` | ZLUDA | | `d4r\kernels\\` | native kernels and their manifest `d4r-kernels.txt` | | `d4r\d4r.ini` | this game's settings | -| `d4r\nvngx_dlss.dll`, `d4r\ngx\_nvngx.dll` | NVIDIA's files (the packager's `D4R_BUNDLE_NVIDIA=0` leaves them for the user to add) | +| `d4r\nvngx_dlss.dll`, `d4r\ngx\_nvngx.dll` | NVIDIA's files, downloaded by `install.sh` (`D4R_BUNDLE_NVIDIA=installer`, the default), bundled (`=1`) or added by the user (`=0`) | **NGX routing.** OptiScaler tries `_nvngx.dll` before `nvngx.dll`, with system fallback for each name. `OptiDllPath=d4r` alone can therefore select Proton's system `_nvngx.dll` before reaching the shim. A file-valued `NvngxPath` takes priority on the first probe, regardless of the probe's name. It must select `d4r\nvngx.dll`, never `d4r\ngx\_nvngx.dll`: the latter is NVIDIA's core, which the shim loads internally. See OptiScaler 0.9.4's [NGX loader](https://github.com/optiscaler/OptiScaler/blob/v0.9.4/OptiScaler/proxies/NVNGX_Proxy.h) and [override handling](https://github.com/optiscaler/OptiScaler/blob/v0.9.4/OptiScaler/Util.cpp). diff --git a/docs/building.md b/docs/building.md index 9f9cfb0..a48dad9 100644 --- a/docs/building.md +++ b/docs/building.md @@ -102,9 +102,7 @@ In the game, pick DLSS as the upscaler (OptiScaler intercepts it). Every setting scripts/fetch_rocm_runtime.sh # -> ~/.cache/d4r-rocm-runtime (AMD's ROCm 7.2.4 runtime, checksummed) D4R_OPTISCALER=/path/to/OptiScaler_0.9.4.7z \ D4R_DLSS_DLLS=/path/to/310.7/nvngx_dlss.dll:/path/to/310.9/nvngx_dlss.dll \ -D4R_BUNDLE_DLSS=/path/to/nvngx_dlss.dll D4R_BUNDLE_NGX=/path/to/_nvngx.dll D4R_BUNDLE_TEX=kernels/out/native \ D4R_ZLUDA_DIR=~/.cache/d4r-zluda-current D4R_VKD3D_DIR=~/.cache/d4r-vkd3d-d4r D4R_ROCM_DIR=/opt/rocm \ -D4R_ZLUDA_EMIT=/path/to/zluda/target/release/examples/d4r_emit \ scripts/package_release.sh # -> dist/d4r-.zip ``` @@ -113,6 +111,7 @@ The script: - writes the kernel manifest from the DLLs you list (it records hashes of their PTX, nothing else); - stages OptiScaler as `dxgi.dll` with the settings in `packaging/optiscaler.settings`, applied by `scripts/configure_optiscaler.py`; both full and clean ZIPs set `NvngxPath=d4r\nvngx.dll` explicitly, retaining `OptiDllPath=d4r` for other libraries; - includes `install.sh` instead of NVIDIA's two DLLs by default, retaining any supplied texture kernels (built from NVIDIA's PTX by `kernels/build.sh tex`); builds accuracy texture sets with `D4R_ZLUDA_EMIT` unless marked prebuilt sets are supplied in `D4R_BUNDLE_TEX/accuracy/`; +- with `D4R_BUNDLE_ENGINE`, adds the native engine's shaders as `d4r/engine/k`; the model's weights and tables are never packaged; - zips the result together with the ZLUDA and vkd3d-proton builds, the ROCm runtime (as `d4r/rocm`), `packaging/d4r.ini`, the licenses and the patches. `D4R_BUNDLE_NVIDIA=installer` is the default: the ZIP includes `install.sh`, and users run `bash install.sh` after extraction to fetch the pinned official NVIDIA DLLs. This mode requires `D4R_BUNDLE_DLSS` to match the DLSS 310.7.0 SHA-256 pinned in the installer, so bundled texture kernels match the downloaded library; `D4R_BUNDLE_NGX` is only needed with `D4R_BUNDLE_NVIDIA=1`, which retains the old DLL-bundling mode. diff --git a/engine/INPUT_K_NOTES.md b/engine/INPUT_K_NOTES.md new file mode 100644 index 0000000..65847a1 --- /dev/null +++ b/engine/INPUT_K_NOTES.md @@ -0,0 +1,44 @@ +# K input kernel (hiluma_engine_input_depthinv_mvlo_hdr_v2_rel) - reverse-engineering notes +P = parameter block (248 bytes). Textures: 176 history colour (Y,Co,Cg,a compressed; linear), 184 motion, 192 depth, 200/208/216 exposure (1x1). +Filter modes are NGX's own (cuTexObjectCreate): 168, 176 and 224 linear, the others point; all clamped, normalized coordinates. +200 = current exposure, 208 = previous exposure, 216 = current exposure. +## entry + a = tex208 (0 -> 1), b = tex200 (0 -> 1), c = tex216 (0 -> 1) + f152 = P144*a ; f158 = P144*b ; f1 = P156*f158 (0 -> 1) ; f2 = (P148*c)/f158 ; f3 = f158/f152 + One thread per 2x2-pixel sub-block: gx = 16*ctaid.x + tid.x, gy = 16*ctaid.y + tid.y ; + X = 4*(gx>>2) + 2*(tid.x&1), Y = 4*gy + (tid.x&2) -> token (gx>>2, gy), sub index = (tid.x&1) + 2*((tid.x>>1)&1) + f4 = X+.5, f5 = Y+.5 ; (ix,iy) = roundHalfAway((f4,f5)*P40 - .5 + P24) + depth (tex192) at (ix,iy) and the four diagonal neighbours, each clamped as v > 0 ? min(P88 + v, P88 + P96) : P88 + same dilation rule as the output kernel -> (dx,dy) ; mv = texelFetch(tex184, clampfn(P104 + (ix+dx,iy+dy))), max = P104 + P112 + q = (f4,f5) + P8*(mv + P16) ; hist = textureLod(tex176, q*P120) (f22..f25) ; P136 loaded (f201,f202) +## the four pixels j of the sub-block: (X,Y), (X+1,Y), (X,Y+1), (X+1,Y+1) + each: nearest render texel, dilated motion mv_j (as above), q_j = pixel + .5 + P8*(mv_j + P16), hist_j = textureLod(tex176, q_j*P120) + lin_j = hist_j.rgb * expandK(hist_j.r) * f3 -> shared W[j][tid.y*16+tid.x] (3 floats) + local[j] = ((token centre in px) + P8*(mv_j + P16)) * P128, token centre = 4*(X/4 + .5), 4*(Y/4 + .5) + invalid |= reset(P32 u8) || q_j*P136 outside [0,1] (accumulated over the four pixels: p3) + r32 = 1 - subY ; r44 = r32 + X ; r45 = subX + Y +## BB12: current colour statistics, 3x3 render texels around the sub-block centre + (cx,cy) = roundHalfAway((X+1, Y+1)*P40 - .5 + P24) ; colour tex168, clamp origin P72, max P72+P80 ; f481 = P160 + per tap (order (-1,-1),(0,-1),(1,-1),(-1,0),(0,0),(1,0),(-1,1),(0,1),(1,1)): rgb = f1*max(c,0) ; Y,Co,Cg as in the output kernel + wA = max(1/(Y+1), P160) ; S += wA*(Y,Co,Cg) + t = log2(12*max(Y,0)+1)*.6931472 / (max(Y,0)+1e-6) ; L = (Y,Co,Cg)*t ; running min and max of L over the taps + after the 9 taps: m = S/9 ; avg = m / max(1 - m.Y, P160) ; A = avg * t(avg.Y) (t as above: log compression) + halfr = .5*(max - min) ; den = 8*(max - min) + 1 (per channel) + for each j: Hl = lin_j * t(lin_j.Y) ; d = max(|A - Hl| - halfr, 0)/den ; e_j = min(2*d.Y + .5*d.Co + .5*d.Cg, 1) + mismatch = .25 * sum_j e_j + representative pixel of the sub-block: offset (1 - subY, subX), i.e. j' = 2*subX + (1 - subY) + Yc = exposed luma (f1 * (.25R + .5G + .25B) of max(c,0)) of the colour texel nearest to that pixel + feat = textureLod(tex224 (previous token feature, 4 ch), local[j']) + if invalid: mismatch = 0, Yh = Yc, feat = 0 else Yh = lin_j'.Y + featSum = sum of feat over the token's four threads (shfl bfly 3, then 1) + tc = (f2*Yc) * max(1/(f2*Yc+1), P160) ; th = (f2*Yh) * max(1/(f2*Yh+1), P160) + then a switch on the sub index +## output: 4 f16 per sub-block into buffer P240 at ((tokY*P64 + tokX)*16 + sub*4) (16 channels per token: the network input) + G(x) = x > .0031308 ? 1.055*x^(1/2.4) - .055 : 12.92*x (sRGB transfer) + out0 = G(sat(tc)) ; out1 = G(sat(th)) ; out2 = .25*featSum[sub] (non-finite -> 0) ; out3 = mismatch (non-finite -> 0) + No barrier: the shared array is only per-thread scratch. + +## display-resolution motion (frame flag 4, `D4RO0003`) + Vectors are output-sized and already in output pixels (scale 1 by default). Each output pixel j of a token uses its own + vector: mv_j = mvScale*(tex184(clamp(X+j)) + mvOffs), with no diagonal depth dilation and no render-texel lookup. + The rest (q_j, history fetch, local terms) is unchanged. The package must carry `D4RO0003`. diff --git a/engine/M_NOTES.md b/engine/M_NOTES.md new file mode 100644 index 0000000..63c16fa --- /dev/null +++ b/engine/M_NOTES.md @@ -0,0 +1,134 @@ +# Preset M (rrlite) - reverse-engineering notes + +Working notes for the engine's own preset M stages. Sources: captured launches (`build/engine-work/mq`), the +translated kernels (`vulkan/ptx2glsl.py` + `vulkan/condense.py` on the DLL's PTX; M uses the `_folded` entry +points: `rrlite_enc0_4x4_mvlo_hdr_folded`, `rrlite_dec0_4x4_folded`, `rrlite_post_3_1_mvlo_hdr_folded`), and +`kernels/tools/swin_model.py` for the network. + +## Frame (640x360 render -> 1280x720 output in the capture) +1. exposure: `cuda_dldn_engine_luma_convert_kernel`, 2x `reduce_sum`, `auto_exposure_copy` (every frame for M) +2. `rrlite_enc0_*_folded` grid 81x49 blocks of 32 lanes: features from textures -> 16 tokens x 128 FP8 features per + block -> GEMM 128 -> 64 (+ bias, clamp +-2pi) -> tokens 320x192 x 64 ch (FP8, planes of 32) ; also writes three + 640x360 2-channel f16 surfaces (args 280/288/296) +3. enc1 (C 64, 2 heads, shift 2, merge -> 96), enc2 (C 96, 4 heads, shift 2, merge -> 128), 6 tube blocks (C 128, + 4 heads, 80x48 tokens, shifts 2,0,2,0,2,0, own weights each), dec2 (C 96, expand from 128, shift 0), + dec1 (C 64, expand from 96, shift 0). Fixed shifts (no per-frame schedule seen). +4. `rrlite_dec0_4x4_folded` grid 80x48 blocks of 32 lanes, P16 = (640, 384) = the token grid x2: per block the 16 + low-resolution tokens (mirrored at the borders) x 64 ch -> GEMM to 4 sub-positions x 32 values (weights at + w + 16 lane + 512 i, f16 bias pairs at w + 8192) -> per render pixel outputs on four 640x360 surfaces: + 128 SNORM8x4, 136 UNORM8x4, 144 UNORM8x4, 152 RGBA16F. Textures 56 (2ch f16, point) and 64 (SNORM8x4, linear). +5. `rrlite_post_3_1_mvlo_hdr_folded` grid 120x34 blocks of 16x16 threads, two output rows per thread: the + reconstruction at 3x render resolution (1920x1080). Textures: 0 colour (linear), 16/24 exposure (1x1), + 32/40 UNORM8x4 from dec0 (point), 48 RGBA16F from dec0 (point), 56 previous 1920x1080 result (linear), + 64/72 2ch f16 (point), 80 2ch f16 (linear). Writes surface 96 (1920x1080 RGBA16F, next history), surface 112, + and a packed 11-11-10 copy to the buffer at 104. Shared tiles (16 rows x 10 x 8 bytes): colour + mvec.x, + mvec.yz + ajab.xy, ajab.z + aniso. +6. `rrlite_downsample_kernel_static_hdr`: packed 1920x1080 -> 1280x720 output. Own version: `down_m.comp`. + +NGX texture filters (cuTexObjectCreate): enc0: 80 linear, 96 point, 112/120 point (exposure), 128 linear, +136 point, 152 linear, 160 linear. dec0: 56 point, 64 linear. post: see above. downsample: 48 point, 56 linear. + +Array formats in the dumps: fmt194 = UNORM8x4, fmt200 = SNORM8x4, fmt16 = f16, fmt32 = f32. + +## dec0 (`rrlite_dec0_4x4_folded`), verified numerically against a capture (`build/engine-work/mre/dec0_check.py`) +Weights: FP8 at w, the Swin fragment layout `woff_table(0, 512, 1024, K 64, N 128)`; bias 128 f16 at w + 8192, +plain order. E = tokens(dec1 output, 64 ch) W + bias, per low-resolution token; render pixel (x, y) takes columns +32 q + j of token (x/2, y/2), q = 2 (y & 1) + (x & 1), each rounded to FP8: v[j]. Only j = 0..16 are used. +- S128 (SNORM8x4, the recurrent feature): out_i = (1 - g_i) hist_i + g_i tanh(v[2 + i]), + g_i = .5 tanh(.5 v[6 + i]) + .5, i = 0..3, f16 arithmetic. hist = previous S128 (tex 64, linear) at the + reprojected position, 0 when reset (P120 byte) or off screen: + f = (pixel + .5 - P72) * P88 + mv (tex 56, texelFetch at the pixel; P88 = 3: units of the 3x image), + off screen if f < 0 or f > P96 (1920, 1080); uv = P104 * ((f / P96 + P80 / P112) * P112), P112 = (640, 360), + P104 = 1 / P112. +- S136 (UNORM8x4): sigmoid(v0), sigmoid(v1), sigmoid(v10), 0 (constant in the capture). +- S144 (UNORM8x4): softmax(v11, v12, v13), alpha 0. +- S152 (RGBA16F): a = P160 f16(sigmoid(clamp(v14, +-20))), b likewise with v15, t = v16 * P164, + rho = tanh(t) (|t| >= 9.01: 1; >= .6: 1 - 2 / (exp2(2.88539 |t|) + 1) with the sign of t; else t + t * poly(t^2), + poly = ((.0157397 s - .052304) s + .133153) s - .3333277) s); out = (a^2, b^2, a b rho, 0). P160 = 8, P164 = .0786. + +## post (`rrlite_post_3_1_mvlo_hdr_folded`) - in progress +Block 16x16 threads; thread (tx, ty) handles hi-res pixels (x, 2y) and (x, 2y + 1) (pairs as f16x2). Tile of +render texels, 10 columns x 16 rows, origin (floor(P124 * 16 bx) - 2, floor(P128 * 32 by) - 2), clamped to +[0, P148 - 1] x [0, P152 - 1] (+ colour origin P8/P12), three 8-byte-per-texel arrays: + tile0: range-compressed colour c' = log10(max(100 * min(e1 * c, 32752) + 1, 1e-4)) * 0.36 (rgb, f16; e1 = tex16, 0 -> 1), T40.x + tile1: T40.y, T40.z, T32.x, T32.y (T32, T40: the UNORM8x4 outputs of dec0) + tile2: T32.z, -0.721325 * cov.x, -0.721325 * cov.y, -1.44265 * cov.z (cov = tex48 = dec0's covariance) +Render position of a pixel centre: f4 = (x + .5) * P124 + P132 (P124 = 1/3), f9/f11 likewise in y with P128, P136. +Main loop: the 3x3 render texels around floor(f - .5); per texel an anisotropic Gaussian +w = exp2(A dx^2 + B dy^2 + C dx dy) with (A, B, C) that texel's tile2 values and (dx, dy) = texel centre - position +(f16); every tile quantity is accumulated with w (kernel-predicted resampling). +After the loop (per pixel, f32): cur = sum(w * colour') / sum(w) (rgb, compressed), m = resampled T40.xyz, +g = resampled T32.xyz. +Motion: three candidates, mvA = texelFetch(tex64) and mvB = texelFetch(tex72) at the nearest render texel +(floor(f4), floor(f9)), mvC = tex80 (linear, uv = position * P160/P164); paired with the three weights m; the two +with the largest weights are kept (w1 >= w2). k = clamp(1.1 w - .05, 0, 1); b2 = k2 / max(k1 + k2, 1e-10), set to 0 +when below 1/255 (then b1 = 1), b1 = k1 / (k1 + k2). +History sample for a candidate: q = pixel + .5 + mv (hi-res pixels). Off screen (q.x < 0, q.y < 0, q.x > P140, +...): the current colour texel instead (nearest render texel, times e2 = tex24, compressed as above). Otherwise +tex56 (previous result) with the 5-tap bilinear Catmull-Rom cross: base = floor(q - .5) + .5, t = sat(q - base), +w0 = t^2 - (t + t^3)/2, w1 = 1.5 t^3 - 2.5 t^2 + 1, w3 = (t^3 - t^2)/2, w2 = 1 - w0 - w1 - w3, w12 = w1 + w2, +taps at (base.x - 1, base.y + w2y/w12y), (base.x + w2x/w12x, base.y - 1), centre, (.., base.y + 2), (base.x + 2, ..) +with weights w0x w12y, w12x w0y, w12x w12y, w12x w3y, w3x w12y, normalised by their sum. uv = q' / (P140, P144). +Details of the loop: base texel r6 = floor(f4 - .5), r8 = floor(f9 - .5); taps are texels r6 + i, r8 + j, i, j in +-1..1, dx = (r6 + .5 - f4) + i, dy likewise (f16). e = (dx^2 A + dy^2 B) + (dx dy) C in f16, w = exp2(e): the +products use f16(w), the weight sum W the unrounded f32 values. Accumulators are f16 fma chains (i outer, j inner). +tile0 -> cur'.rgb, m.x ; tile1 -> m.y, m.z, g.x, g.y ; tile2 -> g.z (then A, B, C). All divided by W. +T40 = dec0's mix (softmax) -> m ; T32 = dec0's gates -> g. +Candidate order: a = (mvA, m.x), c = (mvC, m.y), b = (mvB, m.z). big, small = (a.w < c.w) ? (c, a) : (a, c); +s = !(small.w < b.w) ? small : b; if (big.w < s.w) swap(big, s). Candidate 1 = big, candidate 2 = s. +Candidate 2's history is 0 (not the fallback) when q.y > P144 or the reset byte (P168) is set. +Blend (f32): hist' = b1 h1 + b2 h2 ; expand(v) = (10^(clamp(v, 1e-4, 3) / .36) - 1) / 100 ; +histLin = e1 * expand(hist') / e2 ; curLin = expand(cur') ; a0 = clamp(P120 * W * g.x, 0, 1) ; +a1 = a0 (1 - g.y) + g.y ; a2 = a0 (1 - g.z) + g.z ; outA = hist + (cur - hist) a1 ; outB = hist + (cur - hist) a2 ; +a non-finite out is replaced by cur. compress(v) = log10(max(100 v + 1, 1e-4)) * .36. +Outputs: surface 96 (x, y) = (compress(outA), 0) - the next frame's history ; +buffer 104 [y * P140 + x] = 11-11-10 bits of compress(outB) * .4081632 (trunc(v * 2047 + .5) clamped; blue 1023) ; +surface 112 (x / 2, y / 2) = compress(mean of outA over the 2x2 pixels) * .4081632. +e1 = tex16 (current exposure), e2 = tex24 (previous), 0 -> 1. + +## enc0 (`rrlite_enc0_4x4_mvlo_hdr_folded`) - in progress +Block = 8x8 render pixels (32 lanes x 2 pixels) = 4x4 tokens of 2x2 pixels; pixel = 8 * block - P8 (2, 2) + (lane & 7, +lane >> 3) (+ 4 rows for the lane's second pixel), mirrored into the image (|v|, and 2 n - v - 2 beyond P56 = render +size); mirrored pixels are flagged (no surface stores). +Per pixel: exposures tex112 (e1) / tex120 ; colour tex80 texel (origin P88, max P232) ; depth tex136 over the 3x3 +neighbourhood (origin P144), keeping the positions of the smallest and largest depth (the byte at P216 swaps their +roles) ; motion tex96 at those two positions and at the pixel -> the three candidates written to surfaces +280 / 288 / 296 (two f16 each; post's mvA / mvB / mvC). +Motion: mv = P200 * (P208 + tex96 texel) (P200 = 3: hi-res pixels). A = at the 3x3 neighbour with the smallest depth +when the P216 byte is set (largest otherwise), B = at the other extreme, C = at the pixel. Strict comparisons, the +centre first, then (-1,-1) (0,-1) (1,-1) (-1,0) (1,0) (-1,1) (0,1) (1,1). +32 feature slots per pixel, FP8, two planes of 16 bytes: + plane 0: R' G' | f.x f.y | B' h1.R | f.z f.w | h1.G h1.B | 0 0 | h2.R h2.G | 0 0 + plane 1: h2.B h3.R | 0 0 | h3.G h3.B | 0 0 | h4.R h4.G | 0 0 | h4.B mf | 0 0 + (R', G', B') = compress(e1 * colour) (with the 32752 clamp and non-finite -> 1 before the log, as in post) + h1..h4 = compress(e1 * s_k), s_k = history colour along a candidate: q = mv + (pixel + .5 - P168) * P192 for + k = 1 (A), 2 (C), 3 (B); k = 4: C again with q = mv + (pixel + .5) * P192 (no jitter term). + off screen (q < 0, q > P248) or reset byte P316: s = the pixel's own colour texel (linear); + otherwise tex152 (the half-resolution output of the previous post, linear filter, uv = q * P256) * 2.45, and when + e2 = tex120 > 0: expand(.) / e2. + f = tex160 (dec0's recurrent feature, SNORM8, linear) at uv = (q_A / P248 + P176 / P56) ; 0 when q_A is off + screen or reset. + mf = P272 * ln(|d| + 1) (0 when reset), d = P268 * 540 / P60 * P184 * (mvA - prevA), prevA = tex128 (the previous + frame's surface 280, linear) at uv = clamp((pixel + .5 - P168 + P176 + mvA * P184) * P64, 0, P56 * P64). + Each value: non-finite -> 0, min(v, 1024), then dithered truncation to FP8: n = u16 at w[2 * i] >> 9 with + i = (P312 << 8) | (|y| & 15) << 4 | (|x| & 15) (to be confirmed), v' = bits((v + f16(bits(v) & 0xFC00 | n)) - + f16(bits(v) & 0xFC00)) & 0xFF80 in f16 arithmetic. +Token = its 2x2 pixels: K index 32 q + 16 plane + byte, q = 2 (py & 1) + (px & 1). Weights at w + 4096: +`woff_table(4096, 512, 2048, K 128, N 64)`, bias 64 f16 at w + 12288; output clamp(E, +-2 pi) stored as FP8. +The N columns (and the bias) are in pair order: natural channel c is column gp(c) (`swin_model.gp`). Verified by +regression on the reset frame (`build/engine-work/mre/enc0_check.py`): slots 0,1,4,5,8,9,12,13,16,17,20,21,24,25,28. +The first 4096 bytes of the blob are the dither table (8 x 256 u16). + +## exposure (cuda_dldn_engine_* kernels), every frame +sum = mean over render pixels of min(max(.25 R + .5 G + .25 B, 0), 10000) ; exposure = 0.23547077 / max(sum, 0.005) +(f64 division, stored as f16). Capture: mean 0.3455 -> 0.6816. The 0.2195 in the parameter block is not used. + +## Sizes and parameters across modes (harness, 2026-10-08) +- Token pixel grid P16 = render size rounded up to multiples of 32; tokens = / 2; the tube works at / 8. +- The reconstruction image is 1.5 x the OUTPUT (2560x1440 -> 3840x2160, 3840x2160 -> 5760x3240), not 3 x render; the + downsample is always the static 3 -> 2 one. P124 (post) = render / reconstruction size. +- `rrlite_post_3_1` runs when reconstruction / render >= 3 (2x and 3x upscaling); `rrlite_post_3_2` when it is smaller + (1.5x: 2.25, balanced: 2.59). post_m.comp is post_3_1 only: the engine covers Performance and Ultra Performance. +- Constants in every mode: post gain 2.75301 ; enc0 change scale 9.04507 and gain 0.22357 ; dec0 sigma 8, rho 0.07863. +- jitter = NGX's Jitter.Offset as given ; enc0 P176 / dec0 P80 = the previous frame's jitter ; enc0/dec0 motion scale + and render->hi factors = reconstruction / render ; downsample P80 = render / output, P88 = jitter. diff --git a/engine/OUTPUT_K_NOTES.md b/engine/OUTPUT_K_NOTES.md new file mode 100644 index 0000000..b9cec7b --- /dev/null +++ b/engine/OUTPUT_K_NOTES.md @@ -0,0 +1,89 @@ +# K output kernel (hiluma_engine_output_depthinv_mvlo_hdr_max_v2_rel) - reverse-engineering notes +P = parameter block (328 bytes). rd1 = P+256, rd17 = P. Offsets below are absolute in P. +Textures: 200 color (linear), 224 history color (linear), 232 hi-res history 2x (linear), 240 motion (point), 248 depth (point), + 256 exposure prev? (1x1), 264 exposure cur? (1x1). Buffers: 272 (network head, 40 ch?), 280. Surfaces: 288 out/history color, + 296 hi-res history (R16, byte stores), 304 320x256 x4, 312 final output. +## entry + eA = tex264(.5,.5).r (0 -> 1), eB = tex256(.5,.5).r (0 -> 1) + f497 = P160.x*eA ; f503 = P160.x*eB ; f1 = P168*f503 (0 -> 1) ; f2 = f503 / f497 + block origin (r1,r2) = 16*ctaid. render-space tile: u = (px+0.5)*P40 + P16 ; r3,r4 = roundHalfAway(u(first)) - 4 ; r5,r6 = roundHalfAway(u(last)) + 4 + r7 = r5-r3+1 (tile width) ; r8 = r6-r4 ; path B (shared tile) if r5-r3 < 24 and r8 < 24 +## tile load (path B), thread t = tid.y*16+tid.x, for idx = t; idx < w*h; idx += 256 + x = idx % w + r3 ; y = idx / w + r4 ; cx = x > 0 ? min(P208.x + x, P216.x) : P208.x (same for y) [P208 v2u32 origin, P216 v2u32 max] + c = texelFetch(T200 color, cx, cy) ; rgb = max(c, 0) + Y = .25R + .5G + .25B ; Co = .5R - .5B ; Cg = .5G - .25R - .25B (YCoCg) + L = f1 * Y ; k = L < 1000 ? 1/(L+1) : (BB127 const) ... + k = L < 1000 ? 1/(L+1) : (log2(9L+1)*0.6931472/9 - 0.012676) / L (range compression of the exposed luma) + sharedColor[row*192 + col*8] = f16 { Y' = L*k, Co' = f1*Co*k, Cg' = f1*Cg*k, alpha = c.a } (order in memory: Y', Co', Cg', a) + depth d = texelFetch(T248, clamp like above with origin P96, max P104).r -> sharedDepth[row*96 + col*4] (f32) + bar.sync +## main (after barrier), per output pixel (px,py) = 16*ctaid + tid + token tx = int(px*P64.x)/2 (P64 = 0.5 -> px/4), sub = (int(py*.5)&1)*2 + (int(px*.5)&1) ; head H = buf P272, 40 f16 per token = 4 sub-blocks x 10 + a9 = H[(ty*P80 + tx)*40 + sub*10 + 9] ; x = .5*a9 ; f1800 = x<-3 ? 0 : x>3 ? 1 : .5 + .5*x(x^2+27)/(9x^2+27) (sigmoid(a9)) + T240 motion, T224 history color, T232 hi-res history (2x) + f292 = px+.5, f293 = py+.5 ; (ix,iy) = roundHalfAway(f292*P56 + P16) (nearest render pixel) ; tile pos (ix-r3, iy-r4) + depth dilation on the tile: d0 centre, diagonals Dmm(-1,-1) Dpm(row-1,col+1) Dmp(row+1,col-1) Dpp(+1,+1) + use = d0 < (Dmm+Dpp)/2 - .001*d0 || d0 < (Dpm+Dmp)/2 - .001*d0 + pick the largest of d0, Dmm, Dpm(row-1,col+1), Dmp(row+1,col-1), Dpp in that order (strict >) -> (dx,dy); (0,0) if !use + mv = texelFetch(T240, clampfn(P112 + (ix+dx, iy+dy))) [clampfn(v): v > 0 ? min(v + P112, P120) : P112] + q = (f292,f293) + P0 * (mv.xy + P8) (previous position in output pixels; P0 = mv scale, P8 = mv offset) + uvh = q * P144 ; uv2 = q * P152 ; p5 = P32(u8) != 0 || uv2.x < 0 || uv2.x > 1 (.. more below) + hist = textureLod(T224, uvh) -> f295..f298 + hi-res history (T232): h = (q-.5)*2+.5 ; base = floor(h-.5)+.5 ; t = sat(h-base) ; Catmull-Rom w0..w3(t) per axis + (w0 = t^2-.5(t+t^3), w1 = 1.5t^3-2.5t^2+1, w3 = .5(t^3-t^2), w2 = 1-w0-w1-w3) + gathers (textureGather .r) at uv = (P144*.5) * (base + {-.5, 1.5, 3.5}) in x and y ; each value v: v < .99 ? v : .99 + (v-.99)*.001 + f1858 = P172 + gather footprint: 5x5 hi-res texels T[c][r], c,r = 0..4, T[0][0] = texel (floor(h-.5)-1) i.e. gather(x0,y0).w ; all softclamped (.99 rule) + a = w1/(w1+w2), b = 1-a (per axis), s = P172 (sharpness) + row(r, c0) = a*T[c0+1][r] + b*T[c0+2][r] + s*( w0*sat(T[c0][r]) + (w1-a)*sat(T[c0+1][r]) + (w2-b)*sat(T[c0+2][r]) + w3*sat(T[c0+3][r]) ) (c0 = 0 or 1) + col(c0, r0) = ay*row(r0+1,c0) + by*row(r0+2,c0) + s*( wy0*sat(row(r0,c0)) + (wy1-ay)*sat(row(r0+1,c0)) + (wy2-by)*sat(row(r0+2,c0)) + wy3*sat(row(r0+3,c0)) ) + Hs[sx][sy] = clamp(col(sx, sy), min/max of T[sx+1..sx+2][sy+1..sy+2]) for sx,sy in {0,1}: f300=(0,0) f301=(1,0) f302=(0,1) f303=(1,1) + re-exposure of each: g(v) = v < .999 ? v/(1-v) : (exp2((v+.012676)*9*1.442695)-1)/9 ; Lh = f2*g(v) ; v' = Lh * k(Lh) (k as in tile load) + -> f314 (0,0), f322 (1,0), f330 (0,1), 4th from f331 in next block + p6 = p5 || uv2.y < 0 || uv2.y > 1 (history invalid) +## BB162-164 + f342 = Hs'(1,1). nonfinite = !( |f314*f322 + f330*f342 + hist.r*hist.g + sig*hist.b| < inf ) ; p7 = p6 || nonfinite + sig = sat(f1800) ; s1 = p7 ? 1 : sig ; ha = p7 ? 1 : hist.a ; f343 = s1*(1-ha) + ha ; f344 = ha*(1-s1^2) + s1^2 ; f345 = P176 + current colour, bilinear from the tile: u = (px+.5,py+.5)*P40 + P16 - (r3,r4) [to f16] ; i = floor(u) ; f = sat(u - i) (f16) + texel words at tile[i.y*192 + i.x*8 (+8, +192, +200)], weights (1-fx)(1-fy), fx(1-fy), (1-fx)fy, fx*fy, f16 fma chain + -> cur = (Yb, Cob, Cgb, ab) = (f2251, f2252, f2253, f2254) + (a0,a1) = H[tok*40 + sub*10 + 0,1] ; bright = Yb > .8 ; quad = OR of bright over the 2x2 pixel quad (px^1, py^1) [shfl bfly 16 then 1] = r255 + m0 = a0/(1+|a0|), m1 = a1/(1+|a1|) (f16 softsign; rs606, rs607) + Hq[4] = p7 ? 0 : Hs' ; hc = p7 ? 0 : hist.rgb + f367 = max(f344/(f344+1), P176) +## BB166: kernel-predicted hi-res luma filter (4 sub-pixels of the output pixel, 12 taps of the tile's Y') + i0 = clamp(roundEven(m0*P188)+8, 0, 15), i1 likewise with m1 (P188 f16 = 8.0) + A = (sig < .25 && f343 < .25) ; B = (P180 > f343) ; idx = ((quad<<2 | A<<1 | B) << 8) | (i0<<4 | i1) + LUT at buffer P280: 16-byte entries: f16 e0 (lo of word0), e1 (hi word0), e2 (lo word1), e3 (hi word1) = slope, u16 e4 (lo word2) = orientation, e5 (hi word2), e6 (lo word3) + sub-pixel positions in tile coords: xs[j] = (2px+j+.5)*P48.x + P16.x - r3 (j=0,1), ys[j] likewise (f16) + if e4 == 0: A-axis = y, B-axis = x else A = x, B = y. cA, cB = pixel-centre tile coords (rs329/rs330 = u from the bilinear step) + rows: A_r = floor(cA) - 1 + r, r = 0..3 ; centre line cB_r = cB + (A_0 - cA)*e3 + r*e3 ; B start = roundEven(cB_r) - 1 ; taps B, B+1, B+2 + for each sub-pixel s: dA = A_r - subA_s ; dB = B - subB_s ; E = e0*dA^2 + e2*dA*dB + e1*dB^2 (f16 fma chain) ; w = exp2(E) + sumY += w * Y'(tap) ; sumW += w ; track the tap with the largest E (strict >, first wins) and its Y' (Ymax) + F[s] = sumY/sumW (f16; f32 rcp with one refinement step for tiny quotients) -> rs608 (x0,y0), rs609 (x1,y0), rs610 (x0,y1), rs611 (x1,y1) + Emax pairs r260 (subs 0,1), r261 (subs 2,3); Ymax pairs r262, r263 +## BB174: blend with history, hi-res store + Fsel[s] = (Emax[s] < -12) ? Ymax[s] : F[s] + confidence per sub-pixel (f16): rA = roundEven(subA) ; dA = rA - subA ; line = subB + dA*e3 ; dB = |roundEven(line) - line| + KA = (e4==0) ? P192 : P190 ; KB = (e4==0) ? P190 : P192 (f16 params at 190 / 192) + cB = sat(dB*KB + 4.5) ; cA = sat(|dA|*KA + 4.5) ; conf[s] = (cA*e5 + e6) * cB ; avg = .25 * sum(conf) + v = f16(f343) ; v2 = .1 + .9*v ; lo = (v - P184)*P186 ; hi = (v2 - P184)*P186 (P184, P186 f16 params) + alphaC = sat(avg*(v2 - hi) + hi) (f2349) ; bl[s] = sat(conf[s]*(v - lo) + lo) (f2350..3) + nonfinite2 = !(|Yb*Fsel0*Fsel1*Fsel2*Fsel3| < inf) -> cur and Fsel all 0 + N[s] = Hq[s] + bl[s]*(Fsel[s] - Hq[s]) (new hi-res luma, f384..f387) + store S296 (R16F 2x): texel (2px+sx, 2py+sy) = enc(N[s]), enc(v) = v < .99 ? v : .99 + (v-.99)*1000 + colour: C = alphaC*cur + (1-alphaC)*hc (Y,Co,Cg) ; rCo = C.Co/max(C.Y,1e-5), rCg likewise ; Nm = .25*sum(N) ; f391 = rCo*Nm ; f392 = rCg*Nm + p9 = nonfinite || nonfinite2 ; P90(u8) == 0 -> BB203 (the path taken in the captures; P89, P90, P91 all 0). P90 != 0 path not analysed. +## tails (P90 == 0) + BB203: (Yn, Con, Cgn) = (f390 = Nm, f391, f392) [debug colours if P89 && p9] ; P91 != 0 -> BB208 else BB205..207 -> final output S312 + BB210 history colour S288(px,py): Y = max(Nm,0) ; lim = Y ; if (Cgn - 1e-5 < -Y) lim = 1e-5 - Cgn ; if (Con - Cgn - 1e-5 < -lim) lim = Cgn - Con + 1e-5 ; + if (-Con - Cgn - 1e-5 < -lim) lim = Con + Cgn + 1e-5 ; sc = (lim <= Y) ? 1 : Y/lim ; store (Y, Con*sc, Cgn*sc, sat(f367)) + BB213 S304 (token grid): only for px < P80 && py < P84 (px,py used as token coords): m_s = .2 * sum_{k=3..7} H[(py*P80+px)*40 + s*10 + k], s=0..3 ; + n = rsqrt(sum m_s^2 + 1) ; store (m0,m1,m2,m3)*n + final S312: if P91 == 0: (Y,Co,Cg) = (Yn,Con,Cgn) * ginv(Yn) / f1, ginv(v) = v < .999 ? 1/(1-v) : (exp2((v+.012676)*9*1.442695)-1)/(9v) + RGB = (Y + Co - Cg, Y + Cg, Y - Co - Cg) ; if px < P136.x && py < P136.y: store S312(P128 + (px,py)) = (R, G, B, ab) + +## display-resolution motion (frame flag 4, `D4RO0003`) + Motion is read at the output pixel itself: mv = tex240(clampOrigin(px, P112, P120)), with no depth tile load and no + nearest-depth selection (so depth is not read). P0 (scale) and P8 (offset) are used as before, with the motion + texel in output pixels. diff --git a/engine/README.md b/engine/README.md new file mode 100644 index 0000000..867c095 --- /dev/null +++ b/engine/README.md @@ -0,0 +1,874 @@ +# Native weight-based inference engine + +This directory implements standalone preset K and M inference engines. NGX is used **once to obtain weights +and reference data**. K loads extracted f16 matrices and a reconstruction table; M loads FP8 matrices and +f16 vectors. Both run d4r's own Vulkan shaders. M can run its network on d4r's native HIP kernels instead; +neither engine translates PTX or launches NVIDIA kernels. + +The reusable library is `libd4r_engine.so`; `d4r-k` and `d4r-m` provide full-frame replay and benchmarking. +The D3D12 game shim supports both as opt-in engine backends (see [In a game](#in-a-game-d3d12-shim)). +The CUDA/ZLUDA backend remains the default and fallback. Presets L/E and RTX/game-quality parity are not +implemented or established here. Exposure is supplied by the application or measured on the GPU. + +## Architecture + +``` +extracted model ──► shared Model (weights + LUT + 13 compiled pipelines) + │ +frame color/depth/motion ──► Engine input stage ──► 11 fused Swin layers ──► reconstruction ──► output + ▲ │ │ + │ encoder skip buffers │ + └──── previous color / luma / token features ◄─┘ +``` + +- `Model` owns weights in device-local memory and compiled pipelines. Multiple `Engine` instances can share it. + Shader compilation and weight uploads happen at initialization. +- Each `Engine` owns two reusable activation buffers, five encoder skip buffers, and two sets of temporal history. + The input stage writes directly into the first activation buffer; the head feeds reconstruction directly. + Intermediate activations and history stay on the GPU. +- A frame records thirteen dispatches into a caller-owned command buffer, with dependencies between stages and + before reusing buffers. It performs no allocations, queue submissions, CPU waits, or intermediate readbacks. +- Input images are borrowed Vulkan views. The caller provides the device, synchronization, and command buffer, + so a future game adapter can schedule inference with its own rendering work. +- Per-frame parameters include exposure, jitter, motion conventions, reset, and each layer's window shifts. + Window shifts can change without rebuilding shaders or buffers. Resolution changes require a new `Engine`. +- Frame parameters and window shifts are written by commands inside the frame's own command buffer, so a frame can + be recorded while earlier ones are still queued on the GPU. Each input/output view and layout combination gets + immutable descriptors, cached and reused; first use allocates descriptors, steady recording does not. Borrowed + images and views must outlive the `Engine`. Keep the `Model` alive until all engines using it have been destroyed. + Submit recorded stateful frames once, in order. + +The network shader uses 16x16 cooperative matrices with f32 accumulation, four wave32 subgroups per 8x8 token +window, fused attention/MLP/patch operations, padded LDS rows, and a residual stream held in registers. The output +shader uses wave64 and computes reconstruction and temporal accumulation in one dispatch. + +## Build and extract a model + +Requirements: C++17, Vulkan headers/loader, Python with NumPy, and `glslangValidator`. The GPU/driver must support +Vulkan 1.3, `VK_KHR_cooperative_matrix` with 16x16x16 f16 operands/f32 accumulation, wave32/wave64 subgroup control, +f16 storage/arithmetic, shader int16, and extended storage-image formats. The executable checks these capabilities. + +A model directory holds values read from NVIDIA's DLSS (`weights.bin`, `offsets.bin`, `lut.bin`, `hipnet.bin`, +`model.json`) next to d4r's compiled shaders. Build it from your own `nvngx_dlss.dll` and keep it on your machine: +those files are not d4r's to license, `.gitignore` excludes them, and `scripts/package_release.sh` never packages them. + +```sh +bash engine/build.sh # build/native-engine/{libd4r_engine.so,d4r-k,d4r-m} +# With ROCm installed (D4R_ROCM_DIR, default /opt/rocm) d4r-k and d4r-m also link the native HIP network. + +# CAP contains replay/ and dump/ for the K input, 11 layers, and output launches. +# All replay allocations must be present for weight extraction. +python3 engine/extract_k.py "$CAP/replay" build/model-k +python3 engine/prepare_k.py "$CAP" build/k-case --frames 6 +python3 engine/compile_k.py build/model-k build/k-case/lut.bin build/k-model + +build/native-engine/d4r-k build/k-model build/k-case/case.txt build/k-results 25 +python3 engine/test_frame.py build/k-case build/k-results +``` + +Capture the standard K HDR, inverse-depth, render-resolution-motion input/output kernels: +`hiluma_engine_input_depthinv_mvlo_hdr_v2_rel` and `hiluma_engine_output_depthinv_mvlo_hdr_max_v2_rel`. +The bridge's `D4R_CUDA_REPLAY_DUMP_DIR`, `D4R_CUDA_REPLAY_DUMP_LIMIT`, and `D4R_CUDA_LAUNCH_DUMP_DIR` capture the +arguments, uploaded weight allocations, textures, and reference outputs. Take references with the native K layers +and ZLUDA patch 0008; otherwise dec5's native phase output can be zero. Large captures require substantial disk space. + +The reconstruction table is extracted separately from the output launch's argument 280, at its allocation offset. +Model packaging validates tensor dimensions, alignment and complete spans, records file hashes in `model.json`, +and compiles all shaders. The runtime also checks the binary format and every used tensor's bounds before upload. +Weights/tables are proprietary inputs and are not committed in this repository. + +`test_frame.py` compares all four outputs across the sequence, using the engine's own history after the first frame. +It fails for missing/wrong-sized results, nonfinite values, or PSNR below its threshold (50 dB by default). +`test_k.py --chain` remains available for isolated-layer diagnosis; it now selects the first launch consistently and +fails on dispatch errors rather than continuing with a broken chain. + +## Ordinary frames, without a capture + +Create a JSON description with raw linear HDR color, inverse depth, and render-resolution motion vectors: + +```json +{ + "output": [1280, 720], + "frames": [ + {"color": "color-0.bin", "motion": "motion-0.bin", "depth": "depth-0.bin", "render": [640, 360]}, + {"color": "color-1.bin", "motion": "motion-1.bin", "depth": "depth-1.bin", "render": [640, 360], "jitter": [0.25, -0.25]} + ] +} +``` + +Color is interleaved RGBA16F, motion RG16F, and depth R32F, tightly packed. `bits: 32` selects f32 color/motion. +Paths are relative to the JSON. The first frame resets; subsequent frames feed back the engine's own history. +Optional exposure/motion/sharpness settings are described by `prepare_frames.py --help` and `FrameSettings`. + +```sh +python3 engine/prepare_frames.py frames.json build/frames.case +build/native-engine/d4r-k build/k-model build/frames.case build/frame-output +# GPU timing with no output readback or output files: +build/native-engine/d4r-k build/k-model build/frames.case - 25 +``` + +`makeKShape` uses a token grid padded to multiples of 32, with a minimum of 256x256 tokens, and a four-token +window shift. NGX changes the shifts of the five deepest layers every frame: the bottleneck layer follows an +eight-frame cycle and each shallower layer's shift is (2 * shift + 4) mod 8 of the one below. `makeKWindows` +reproduces that cycle (observed over 64 evaluations); `d4r-k` replays use the captured shifts or, for ordinary +frames, the fixed one. On a captured sequence the fixed shift costs 4-8 dB against NGX's. The helper supports output dimensions at least +256x256 and render dimensions no larger than output. Device image/buffer limits can impose a lower maximum than +the helper's 8192-pixel dimension limit. Only standard HDR output with zero output origin is supported. + +The eleven layers can also run on the native HIP kernels (`kernels/k`, the accuracy HSACOs) with the stages +around them still on Vulkan, through two storage buffers shared by file descriptor. `compile_k.py --hip +KERNEL_DIR --capture CAPTURE_DIR` adds `hipnet.bin` (the layers' weight allocations as NGX uploads them) and +`hip/dltss_pwin_*_layer.hsaco` copied from `KERNEL_DIR`: + +```sh +python3 engine/compile_k.py build/model-k build/k-case/lut.bin build/k-model --hip build/kernels-gfx1201 --capture "$CAP" +build/native-engine/d4r-k build/k-model build/k-case/case.txt build/k-results 25 hip +``` + +`Engine` takes the caller's `ExternalNetwork` (`tokens`, `Engine::networkTokenBytes`; `head`, +`Engine::networkHeadBytes`) and records the frame in two parts, `recordFront` (exposure and the input stage, which +writes the tokens) and `recordBack` (the final store, which reads the head). The network derives its token grid +from the output dimensions, runs the layers between the two parts on the shared buffers, and picks one of eight +immutable graphs from an evaluation counter that advances every evaluation and is never restarted by a reset. +`HipNetwork` (`hip_net.h`, native Linux) and the CUDA bridge's `d4rEngineNet*` (the game path) are the same +logic; both dispatch on `hipnet.bin`'s magic, so one bridge serves preset M and preset K. + +The optional HIP replay accepts the native eight-frame window sequence starting at evaluation zero. +Custom windows or a capture beginning mid-cycle are rejected rather than silently using a different graph; +use the Vulkan replay for those cases. + +### Shared M/K engine cutover and K accuracy limits + +Preset K packages now use `D4RK0002`; rebuild old K packages for the current runtime. Automatic exposure +samples the output/2 grid with the native coordinate and luma associations, truncates the measured exposure +to half, and measures on alternate evaluations starting with the first. Reconstruction confidence derives +its two geometry scales from output/render rather than assuming a fixed 2× ratio. The filtered current luma +is rounded to half before the f32 blend; the blend, bright-luma encoding and raw four-lane mean retain the +native operation order. These changes apply to both native-network and Vulkan-network K. + +M and K share `ExternalNetwork`, direct RGB10A2 stores, lazy canonical input allocation, direct-input +barrier elision, native package loading, GPU input readiness and callback-released split completion. +K's native backend includes all three bottleneck post kernels and eight immutable window graphs; M retains +two graphs. External-network engines do not allocate the unused Vulkan network activation buffers. +This does **not** mean zero network VRAM: the native K backend owns its own activations and weights. + +The old bundled ZLUDA library omitted K's bottleneck post phases; its captured decoder input and bottleneck +output were zero. It is not a valid K accuracy reference. The current comparison uses the existing +phase-capable bridge with the same accuracy HSACOs. Its three captured frames match the independent +eleven-frame reference byte-for-byte. + +With captured native tokens, all three tested native K network heads are byte-identical to the reference. +With captured native head/history, current K reconstruction is byte-identical for the three LDR colour, +luma and final-output surfaces; HDR final-output stage PSNR is 116.1–123.4 dB. The final scalar FMA +normalisation removes all 34 first-frame feature differences (983,040 values), but does not remove the +coupled drift. **These isolated checks do not prove lossless temporal feedback.** +The full, independently evolving eleven-frame sequence then had first-frame RGB PSNR 114.86 dB HDR and exact +LDR, but subsequent frames fell to 47.30–50.95 dB HDR and 47.36–52.02 dB LDR (fixed peak 1, unclipped RGB; +`build/engine-quality/k-accuracy-v2/native-engine-fixed-reference-quality.json`). + +#### Temporal feedback made exact (2026-10-10) + +That fall was not accumulating error. K is recurrent and its reconstruction takes discrete decisions, so any +one-code half-float difference entering the history is spread over most tokens by the next frame's network and +turned into sparse large errors: about a thousand differing texels do as much damage as thirty thousand. Three +causes were found and removed, none of which costs time: + +- **One reconstruction table for every ratio.** NGX holds sixteen 32 KiB tables on a uniform grid of the + upscaling ratio (0.5 + index × 3/28) and gives the output kernel the nearest even one for the geometric mean + of the two axes' ratios, at most index 14 (ratio 2). The engine packaged the single table of a 2× capture and + used it everywhere, so at Quality (1.5×, table 10) every frame with history filtered its luma with the wrong + table. `lut.bin` now holds all sixteen (`prepare_k.py` writes NGX's whole allocation), the frame parameters + carry the table in bits 8–11 of the output flags (`kReconstructionTable`), and a single-table package keeps + the old behaviour. The rule matches NGX at all 49 measured sizes, including unequal axis ratios; HDR and LDR + use the same tables. **Rebuild K models to get this.** +- **Input stage, mismatch channel.** NGX rounds the compressed history before subtracting it from the current + average; the compiler fused that product into the subtraction, leaving the channel one code off in about 3% of + the sub-blocks of every frame with history (30,805 of 983,040 values on the test frame). Now `precise`. +- **Output stage, sub-pixel positions.** NGX fuses the scale and the jitter offset and subtracts the tile origin + afterwards. At some size ratios a whole output column lands on a half-float rounding tie, where any other + order moves the luma filter's samples by one half-float step (252 history and 909 luma texels on the test frame). + +Measured on the RX 9070 XT with the native HIP network, 1706×960 → 2560×1440, create flags 0x4b, against the +phase-capable CUDA reference (`hdr-phase-accuracy-hip-11`): the per-stage tests show zero differing codes on +every surface of frames 0–2; the three-frame replay on the engine's own history is byte-identical on all four +surfaces of every frame; and all eleven frames through the shim's game path are byte-identical, where the same +run with the previous shaders and package reproduces 114.86 dB then 49.64–47.30 dB. Whole-frame time of the +three-frame replay is unchanged within run-to-run noise (about 1.6–1.9 ms). Not measured: LDR end to end, other +ratios end to end, and the Vulkan network, which is not bit-exact against NGX and therefore still settles near +50 dB on later frames. + +K gameplay HUD before the final feature-normalisation association fix: 1.52 ms current / 1.58 ms average, effective 1712×960 → 2560×1440 +Quality, preset 11, LDR, native HIP network and direct RGB10A2 output. This is total in-game upscaler time, +not a kernel sum. Evidence: `build/engine-quality/ghost-k-final-gameplay.png` and +`build/engine-quality/ghost-k-final/d4r_nvngx.log`. +The final feature-normalisation package subsequently ran 20,400 evaluations through the in-game native K +path and shut down cleanly; no HUD screenshot was captured from that run, so the earlier timing is not +relabelled as a new final-build measurement. + +Final shared-runtime verification preserved all 22 complete M HDR/LDR outputs byte-for-byte, passed all +114 six-frame direct/fallback adapter cases (57 per colour mode), and passed the nine model-corruption tests. +All 64 packaged SPIR-V modules were validated, including a subsequent validation of the four rebuilt K +output modules. The three-frame native K replay completed with repeated submissions, and incompatible +offset window schedules were explicitly rejected. + + + +## Embedding + +Use [runtime.h](runtime.h). With a suitable Vulkan device and a command buffer owned by the caller: + +```cpp +d4r::Model model(physicalDevice, device, modelDirectory, pipelineCache); +auto shape = d4r::makeKShape(1280, 720); +d4r::Engine engine(model, shape); +model.recordUpload(initializationCommand); +engine.recordInitialize(initializationCommand); +// Submit initialization and wait for its fence. + +d4r::FrameSettings settings{640, 360}; +settings.reset = true; +auto parameters = d4r::makeKFrameParams(shape, settings); +engine.setFrame({colorView, motionView, depthView}, parameters.input, parameters.output); +engine.recordFrame(frameCommand); +// engine.outputImage() is RGBA16F in GENERAL, ready for a transfer read. +// Submit with the caller's rendering dependencies. Later frames can be recorded while this one is queued. +``` + +Borrowed images must be sampleable, accessible to the recording queue family, and alive (including their views) +until Engine destruction after GPU completion. `FrameImages` supplies each sampled layout, defaulting to `GENERAL`. +The caller must enable the device features listed in `runtime.h`. Optional `FrameOutput{view, format}` passed to +`setFrame` receives reconstruction directly: RGBA16F or RGB10A2 storage, `GENERAL`, covering +`outOrigin + outSize`. In this mode `outputImage()` still names the owned, unwritten fallback image. + +[game.h](game.h) is the adapter for game textures: provide `GameImage::view` for compatible sampled inputs and +an RGBA16F or RGB10A2 storage output. It samples inputs and writes the output directly, respecting rectangle origins. +Null views retain per-plane copies. Other output formats and packages without the RGB10A2 variant retain the final conversion blit; +input/output image aliasing also uses the output blit. Views and images must outlive the adapter. Input formats +are sampled in their original precision, so f32 colour/motion avoid the former f16 intermediate rounding. + +`compile_k.py` and `compile_m.py` package both final-store formats; `hasRgb10Output()` reports the packed variant. +Both compilers accept `--ldr` for NGX evaluations without the HDR flag. RGB10A2 stores share `output_store.glsl`: +the mantissa mask reproduces RADV's truncating RGBA16F store before UNORM conversion, rather than changing rounding +when removing the intermediate image. A GPU corpus of 324,614 values (all half patterns, signed half ties and their +adjacent floats, RGB/alpha quantization boundaries, random f32 and low-payload NaNs) matched the old half-store + +blit on every pixel. This format-conversion proof applies to the tested RX 9070 XT/RADV driver, not all vendors. + +Paired output-path timing (2026-10-10, LDR 1712×960 → 2560×1440, 20 alternating samples per mode, +Vulkan network, warm GPU timestamps; not the in-game acceptance number): + +| Final output path | Downsample median | Tail / conversion median | Whole engine median | +|---|---:|---:|---:| +| RGBA16F + RGB10A2 blit | 160.96 µs | 36.12 µs | 3327.32 µs | +| Direct RGB10A2 | 134.68 µs | 3.80 µs | 3272.04 µs | + + +K and M share `game_internal.h`'s canonical input resources and copy barriers. Each conversion plane is allocated +and transitioned on its first fallback frame; an all-direct frame allocates none of these planes and records no +extra transfer barrier. Depth staging is allocated only for a copied depth-aspect input. Later direct/copy changes +remain supported, including first fallback use after earlier direct frames were recorded. This avoids unused +allocations; it is not a measured steady-state speedup. + +## In a game (D3D12 shim) + +`scripts/build_d4r_nvngx_shim.sh` compiles the engine into `d4r_nvngx.dll`. It is off unless requested: + +```sh +D4R_ENGINE=1 D4R_ENGINE_MODEL_DIR=/path/to/build/k-model # the output of compile_k.py +``` + +The shim loads each model on first use from a folder next to it, or from a variable: + +| Preset, colour | Folder | Variable | Built by | +|---|---|---|---| +| K, HDR flag set | `engine\k` | `D4R_ENGINE_MODEL_DIR` (`[Engine] ModelDir`) | `compile_k.py` | +| K, no HDR flag | `engine\k-ldr` | `D4R_ENGINE_MODEL_LDR_DIR` | `compile_k.py --ldr` | +| M, HDR flag set | `engine\m` | `D4R_ENGINE_MODEL_M_DIR` | `compile_m.py` | +| M, no HDR flag | `engine\m-ldr` | `D4R_ENGINE_MODEL_M_LDR_DIR` | `compile_m.py --ldr` | + +`scripts/install_d4r_runtime.sh` stages them from `D4R_ENGINE_MODEL`, `D4R_ENGINE_MODEL_LDR`, `D4R_ENGINE_MODEL_M` +and `D4R_ENGINE_MODEL_M_LDR`, copying each folder whole (with `hipnet.bin` and `hip/` from `--hip`). A model whose +`hipnet.bin` belongs to the other preset keeps the Vulkan network. `D4R_ENGINE_NETWORK=vulkan` (`[Engine] Network`) +forces the Vulkan network for both presets; `D4R_SHIM_ENGINE_GPU_WAIT=0` makes the HIP network wait for its inputs on +the CPU-visible marker instead of a GPU-side wait, for comparisons. + +With the engine enabled, an evaluation of preset K is recorded entirely into the game's own command list through vkd3d-proton's +`BeginVkCommandBufferInterop`: thirteen to fifteen dispatches (automatic exposure adds one or two), with direct +sampled inputs and a direct RGBA16F or RGB10A2 storage output when compatible. Unsupported usages/views keep their individual +input copies or output blit. There is no NGX +evaluation, CUDA launch, worker thread, frame marker, shared-memory import or split frame, and the game presents +the result of the frame it just rendered. The shim log says `engine backend: feature N runs preset K in the game's +command list`. + +The shim hands a feature to the CUDA backend, with a log line giving the reason, when: the preset is not K or M; +the feature is preset M with display-resolution motion vectors (MVLowRes clear), which M does not implement yet; the +motion vectors are display-resolution and the K package predates `D4RO0003` (K with a current package takes them, and +jittered vectors are handled by adding the jitter change, as NGX does); the depth is regular (not inverted) and the K +package predates `D4RO0002`; HDR K runs without AutoExposure and either has no float exposure texture the engine can +sample directly or its package predates `D4RO0003`; a non-positive or non-finite exposure scale or pre-exposure is +passed; or +a texture format is not handled (colour/output RGBA16F, RGBA32F, R11G11B10F, RGB10A2, +RGBA8/BGRA8 UNORM; motion RG16F/RG32F/RGBA16F/RGBA32F; depth D32F, D32F_S8, R32F, R16F, R16_UNORM). + +Limits of this first integration, all untested in a real game: + +- vkd3d-proton does not enable `VK_KHR_cooperative_matrix` on the game's device. RADV compiles and runs the + shaders anyway (measured: `D4R_ENGINE_TEST_NO_COOPMAT=1 d4r-k ...` creates such a device and gives byte-identical + results), but that is outside the Vulkan specification and may not hold for other drivers or Mesa versions. + The proper fix is a vkd3d-proton patch that enables the feature. +- Display-resolution motion (preset K, flag 0x49 without MVLowRes, `D4RO0003`) is byte-exact against the CUDA + backend under the project's parity recipe: eleven frames each at 1706x960 -> 2560x1440 with create flags 0x49 (HDR) + and at 1712x960 -> 2560x1440 with 0x68 (LDR), plus three frames at 1280x720 -> 2560x1440, are byte-identical on + every frame, with output-sized vectors in output pixels and no depth. The CUDA reference must use a phase-capable + ZLUDA (patch 0008, with the native K accuracy layers); the old bundled library omits K's bottleneck post + phases and makes its own output differ by about 33 dB. The harness's panning scene (`D4R_HARNESS_QUALITY_PAN=1,0`, + content moving one output pixel per frame) keeps frame 1 exact and then drifts to about 70-73 dB, with the same + figures for render-resolution motion on the same scene, so that residual is the moving-content history drift and + not the display-resolution path. M still uses CUDA for it. +- With MVJittered set (create flags 0x4f render-resolution or 0x4d display-resolution) the harness's jitter scene + (`D4R_HARNESS_JITTER_SCENE=1`) is 109.8 dB on frame 1 and then drifts to about 69-72 dB; the same scene with + MVJittered clear shows the same drift, so that is the scene's own content/rounding limit, not the offset. The + jitter change is added to the raw vectors as `delta jitter / MV.Scale` in both motion modes, before + `makeKFrameParams` applies the output-pixel conversion; applying an extra output/render factor there moves the + render-resolution case from about 70 dB to about 33 dB. +- NGX and the CUDA bridge are still loaded for initialization/capability APIs. With the engine enabled, K's CUDA + feature, parameters, 64 MiB scratch buffer and completion/profile events are created only when a frame actually + needs fallback. Supported K frames never create that feature or compile its inference shaders. Other presets and + engine-disabled features retain eager CUDA creation. The first fallback frame resets CUDA temporal history; + failed creation frees partial resources so a later evaluation can retry. +- Exposure: HDR K reproduces automatic exposure (AutoExposure) and a game's own exposure texture (AutoExposure off; a float + texture whose (0,0) texel the engine samples directly; `D4RO0003` packages). The texture replaces the measurement. The + game's exposure scale and pre-exposure apply as `source × scale / pre-exposure` in both paths. Two details were needed for + byte-exact results: the automatic measurement scales the luma it averages by `1 / pre-exposure` (as NGX's + `cuda_luma_convert_kernel` does), and the previous frame's game exposure is carried into the ratio as a half with NGX's + truncating store. Measured against the CUDA backend at 1706×960 → 2560×1440 with the native HIP network over eleven + frames: automatic exposure with scale and pre-exposure 1, 0.5, 0.7, 2, 3 and combinations is byte-identical; HDR game + exposure is byte-identical for texture values 0.4, 0.5, 0.6, 0.75, 0.8, 1, 1.2, 1.25, 1.5, 1.6, 2.5, 3, 5, 8, 16, 32, + and for texture 1 with scale 0.5/2/4 or pre-exposure 2/0.5. The one exception: an effective exposure of exactly 2 or 4 + leaves a one-code difference in a few hundred to a few thousand texels on the first frame (104.7 dB and 112.8 dB RGB PSNR, + first three differing frames 104.7/49.5 dB and 112.8/48.2 dB), which the recurrent network spreads over later frames + (≈49–63 dB). The CUDA launches carry the same exposure, ratio and scale parameters, so that gap is an unresolved + half-rounding tie in the reconstruction, not the exposure formula; the same values stay exact with automatic exposure. + A zero texel reads as 1; negative and +∞ pass through and are byte-identical; a NaN passes through but the backends keep + different NaN bit patterns (both outputs are NaN). LDR ignores the exposure entirely and is byte-identical for any value, + scale, pre-exposure and AutoExposure flag. Preset M ignores the AutoExposure flag, the texture, the pre-exposure and the + scale, so the shim does not refuse them (see the Preset M section). +- DLSS indicator axes (`DLSS.Indicator.Invert.X/Y.Axis`) only place NGX's debug indicator: with them set NGX's K and M + outputs are byte-identical to unset runs on the harness (`D4R_HARNESS_INDICATOR_INVERT=1`, CUDA backend), and so are + the engine's, so the shim no longer refuses them. +- Regular depth (K): 47–49 dB against the CUDA backend on the moving-box harness scene, the same as inverted depth; the wrong + convention gives about 41 dB. Output subrectangles (K and M) are passed to the adapters, whose regression test + covers nonzero output origins; they have not been run through the harness. +- NGX's eight-frame window cycle is counted from the feature's first evaluation; whether NGX restarts it on a + reset was not checked. + +Lazy-fallback verification: eleven moving/jittered Vulkan frames remain byte-identical to eager feature creation, +with no `NVSDK_NGX_CUDA_CreateFeature` call. Releasing a recorded-but-discarded evaluation and a newly created, +unevaluated feature both complete without CUDA creation. Changing exposure scale at frame 3 creates the fallback +exactly once; its first three outputs match a fresh CUDA feature's three outputs byte-for-byte, establishing the +history reset. A first-frame fallback also matches eager creation. `D4R_HARNESS_EXPOSURE_CHANGE_FRAME=N` changes +exposure scale to 0.5 at frame N for this transition probe. + +The KCD2 portable installation was checked with an empty CUDA cache at 1706x960 -> 2560x1440: three Vulkan frames, +no CUDA feature creation, direct inputs/output active, and a byte-identical final image to the eager installed shim. +The whole isolated process took 11.2 seconds including Proton startup and NGX initialization; that is not a measured +in-game loading time. NGX initialization may still compile its basic CUDA support code, but the unused fallback +inference feature and its shader compilation are skipped. + +The shim caches borrowed texture views and retains their D3D12 resources until the feature's recorded lists have +retired. Immutable descriptor bindings allow texture rotation and queued frames. New views are capped at 64 and +512 MiB of image memory requirements per feature; new textures beyond that budget use copies. Cached textures +stay allocated until feature retirement. The shim sends aliased input/output resources to the CUDA backend because +their D3D12 state conventions are ambiguous. Inputs denying shader-resource usage keep the copy path; direct outputs +need RGBA16F or RGB10A2 and `ALLOW_UNORDERED_ACCESS`. `D4R_ENGINE_DIRECT_INPUTS=0` and `D4R_ENGINE_DIRECT_OUTPUT=0` independently +force the old paths for comparisons. Older model packages keep copies for nonzero input origins. Recompile with +`compile_k.py` to obtain `direct_origins.bin` and shaders with corrected motion-origin handling. + +The native adapter regression test covers six temporal frames per case, rotating three texture sets, +nonzero K origins, odd K output sizes, K render size changes, resets, automatic and game exposure (K), display-resolution +motion (K), separate direct/copy planes, R32F/D32F depth, read-only input layouts, RGBA32F output fallback, and RGBA32F +colour/RG32F motion against the unquantized Engine. K and M switch each input to fallback for the first time after +startup, then back again. Both also rotate RGB10A2 outputs and mixed RGBA16F/RGB10A2 outputs, with queued and +individually submitted recordings, comparing every output byte against the old conversion path. The K checks need a +package from the current `compile_k.py` (`D4RO0003`). HDR and LDR model pairs each passed every case of six temporal +frames on the RX 9070 XT: + +```sh +bash engine/build.sh +# Reuse the replay runner's device/test helpers for the adapter checks. +g++ -std=c++17 -O3 -Wall -Wextra -Wno-missing-field-initializers engine/test_game.cpp \ + -Lbuild/native-engine -ld4r_engine -lvulkan '-Wl,-rpath,$ORIGIN' -o build/native-engine/test-game +build/native-engine/test-game build/k-model +# Include M's direct-input and late-fallback checks: +build/native-engine/test-game build/k-model build/m-model +``` + +### Direct-texture measurements (2026-10-08, RX 9070 XT) + +D3D12 harness, automatic exposure, three rotating-order rounds per mode/resolution, 140 frames per round with +frames 1–40 discarded and no pause between frames. Medians of the three run medians, evaluation call through the +completed submission fence (includes CPU recording/submission overhead, not game FPS): + +| Adapter mode | 640x360 -> 1280x720 | 1920x1080 -> 3840x2160 | +|---|---:|---:| +| Input copies + output blit | 0.857 ms | 3.240 ms | +| Direct inputs + output blit | 0.924 ms | 3.212 ms | +| Input copies + direct output | 0.752 ms | 3.017 ms | +| Direct inputs + direct output | 0.728 ms | 2.989 ms | + +The 4K runs are stable: about 0.028 ms saved by inputs, 0.223 ms by output, and 0.252 ms (7.8%) with both. +720p is bimodal/clock-sensitive: copying run medians range 0.837–1.017 ms and direct-input-only medians range +0.851–0.987 ms, so there is no established input-only speedup there. Both/direct-output medians range +0.717–0.745 / 0.743–0.761 ms, respectively. These are adapter measurements; do not add them to inference-only +benchmarks or treat the differences as isolated transfer timings. + +Eleven moving/jittered quality frames match the copy adapter byte-for-byte, including a separate run recreating +the output every frame and using RGBA16F motion textures. Seventeen native adapter comparisons pass across six +frames each (including multiple frames recorded before submission). Preserved RGBA32F colour/RG32F motion matches +the unquantized Engine; it is not expected to match the old f16 conversion path. Results establish correctness +on this device; RTX perceptual parity remains unverified. A 20-frame 4K harness run recreating its output each +frame exercised the retained-view budget and switched new outputs to blits with a byte-identical final result. +RGB10A2 output fallback likewise matched the copy adapter. + +### Harness results (2026-10-08, RX 9070 XT) + +`tools/d3d12_dlss_harness.cpp`, quality scene (blinds, thin poles and fine diagonals near the render pixel pitch) +panning 2,1 output pixels per frame, 640x360 -> 1280x720, eleven frames, engine against the CUDA backend: + +| Reference | Final RGB PSNR | +|---|---| +| CUDA backend with d4r's native K layer kernels | 53.2 dB on the reset frame, then 56.5-60.9 dB | +| CUDA backend with ZLUDA's translation of NVIDIA's K layers | 44.6-52.8 dB | +| the two CUDA variants against each other | 44.7-52.9 dB | + +These earlier eleven-frame comparisons do not establish general numerical agreement; the longer, matched +accuracy-HIP comparisons below found larger Vulkan differences. Which variant is closer to an NVIDIA GPU is +not established here. Against the earlier scene's analytic ground truth the engine scores 20.07 dB and the CUDA +backend 19.48-19.59 dB; those scores alone do not establish equivalent image quality. + +#### Matched accuracy-HIP comparison (2026-10-08) + +All three paths ran identical panning (2,1 pixels/frame), jittered quality-scene inputs with automatic exposure. +HIP used split frames, so outputs are aligned to the current evaluation, as with Vulkan. The reference used +`PreferAccuracy=1`, the marked conservative native K layer set, and translated remaining kernels with the +accuracy compiler/synchronization policy. Fast HIP used `k-shipping-deep2-fast` with `PreferAccuracy=0`; +Vulkan used the folded-gain `build/k-store/final-model`, including the enc0 store optimization. +Logs verified all eleven native HIP overrides and the respective policies. The 1440p accuracy reference +was repeated and all sixteen output frames matched byte-for-byte. + +RGB PSNR against accuracy HIP, fixed peak 1, without clipping; alpha excluded. Post-reset PSNR pools +squared error over all subsequent frames rather than averaging their dB values: + +| Case | Fast HIP reset | Vulkan reset | Fast HIP post-reset | Vulkan post-reset | +|---|---:|---:|---:|---:| +| 640x360 -> 1280x720, 32 frames | 52.41 dB | 40.89 dB | 60.89 dB | 46.43 dB | +| 1280x720 -> 2560x1440, 16 frames | 52.77 dB | 41.49 dB | 59.74 dB | 48.02 dB | + +Individual post-reset frames ranged 55.87–69.77 dB for fast HIP and 34.98–60.31 dB for Vulkan at 720p; +56.45–63.73 dB and 44.58–52.65 dB respectively at 1440p. Fast HIP is substantially closer to accuracy HIP +in these tests. These are synthetic temporal scenes on RX 9070 XT, not game or NVIDIA reference captures. +Per-frame metrics, exact kernel/model hashes and the comparison script are in `build/k-accuracy-compare/`; +raw outputs and run logs are `build/engine-work/hout/compare-{720,1440}-{accuracy,fast,vulkan}.*`. + +Wall time of one evaluation in the harness, from the call to the completed fence, with both backends presenting +the frame's own result (CUDA: `D4R_SHIM_VRAM_INTEROP=1 D4R_SHIM_SPLIT_FRAME=1`, native K kernels); medians of 200 +frames after 40 warm-up frames: + +| Resolution | Engine | CUDA backend | +|---|---:|---:| +| 640x360 -> 1280x720 | 0.98 ms (bimodal, 0.74-1.03) | 1.18 ms | +| 1706x960 -> 2560x1440 | 1.78 ms | 2.28 ms | +| 1920x1080 -> 3840x2160 | 3.44 ms | 4.25 ms | + +The harness submits one list per frame and waits for it, so these figures include submission and fence overhead +that a game shares with its own rendering. They are not game frame rates. + +## Numerics and measured results (2026-10-08) + +GPU: RX 9070 XT, RADV/Mesa 26.2.4. Timings use GPU timestamps around the **complete** input/network/reconstruction +graph. Nine measured batches follow three warmups; each batch repeats a fixed input as 25 successive temporal +frames, preserving dependencies and feedback. Uploads, compilation and validation readbacks are excluded. +These are sustained benchmarks, not game timings or matched-clock comparisons against another engine. + +| Configuration | 1280x720 output | 3840x2160 output | Activation buffers at 4K | +|---|---:|---:|---:| +| folded gains | about 0.60 ms | 2.88 ms | 109.6 MiB | +| explicit f16 gains/residual boundaries | about 0.65 ms | 2.95 ms | 109.6 MiB | + +Activation memory excludes weights, temporal images, output, and borrowed frame inputs. At 720p it is 17.2 MiB. +The 4K timing uses ordinary-frame parameters with the fixed window schedule. Full temporal capture comparison +was performed at 720p, where final RGB is 59.9–61.8 dB PSNR over six frames for folded gains, and 60.3–62.4 dB for +explicit gains. All outputs are finite. The first frame's network head is byte-identical between the persistent +runtime and the isolated shader chain (3,276,800 f16 values). + +The default extractor folds per-channel gains into weight rows, rounding once to f16. `extract_k.py --unfold-gains` +preserves explicit f16 gain multiplication. Both modes now preserve the f16 residual additions at the MLP boundary +and use f32 matrix accumulation. Neither is bit-identical to NVIDIA: small head differences +can cross discrete reconstruction-table boundaries and cause larger errors at isolated pixels. PSNR is a numerical +check on these captures, not evidence of perceptual parity or state-of-the-art quality. + +The input shader was additionally compared on synthetic motion/depth: current/history luma matched the translated +reference exactly, token feature and mismatch at 122 dB and 131 dB PSNR. It samples the history colour and the +previous token feature bilinearly, as NGX's texture objects are created; with point sampling (an earlier mistake +that zero-motion captures could not expose) a panning capture's token feature dropped to 23 dB. The output shader on synthetic motion/depth, +given the captured head, measured 78.4 dB for hi-res luma and 86.8 dB for final RGB. Translation is used only for +these reference comparisons, never by the engine. Both model modes' 26 SPIR-V modules passed `spirv-val`. + +```sh +python3 -m unittest discover -s engine -p test_model.py +python3 -m unittest discover -s tests +``` + +## Per-stage K profiling + +Set `D4R_ENGINE_PROFILE=1` when running `d4r-k` to record GPU timestamps at frame entry and after setup, +exposure, input, each of the eleven network layers, reconstruction, and final synchronization: + +```sh +D4R_ENGINE_PROFILE=1 build/native-engine/d4r-k build/k-model build/k-case/case.txt - 25 +``` + +With timing repeats, stdout reports each stage's median across the nine measured batches, after three warmups. +Each batch averages successive temporal frames. With readback enabled, `stages.csv` also records the initial +sequence's per-frame intervals in microseconds; these initial frames are not warmed-up benchmark samples. +Without repeats, the per-frame stages are printed. The default run records no extra timestamps. Embedders can +use the optional `FrameProfiler` callback in `runtime.h` to record their own markers. + +Intervals include the dependencies preceding each dispatch and timestamp overhead; they are not isolated +kernel timings. Timestamps can perturb scheduling, so use profiling to locate costs and ordinary, uninstrumented +runs to assess an optimization. Input uploads, game adapter copies/blits, and readbacks are outside these timings. +The exposure interval contains only marker overhead when automatic exposure is disabled. + +Checked on 2026-10-08, RX 9070 XT: profiling on and off both produced byte-identical head, final output and all +three histories to the pre-instrumentation executable over a five-frame 720p capture (25 binary files per mode). +All four reference checks passed the default 50 dB threshold; final RGB was 59.9–61.8 dB. + +Three alternating profiling-on/off runs per resolution, 25 frames per batch, folded gains: + +| Stage/group | 640x360 -> 1280x720 capture | 1920x1080 -> 3840x2160 synthetic | +|---|---:|---:| +| Input | 41.9 us | 316.1 us | +| Eleven network layers (sum of stage medians) | 409.9 us | 1456.0 us | +| Reconstruction | 93.8 us | 903.8 us | +| Whole frame, profiling enabled | 551.0 us | 2681.1 us | +| Whole frame, profiling disabled | 558.0 us | 2762.9 us | + +The 4K case uses constant HDR colour, inverse depth and zero motion with the fixed window schedule, not a game +or quality reference. Automatic exposure is disabled in both cases. Whole-frame figures are medians of the three +run medians; stages are medians of the three stage medians and need not sum exactly to the whole-frame median. +The instrumented totals were 1.3% and 3.0% lower; this is a measurement perturbation, not an inference speedup. +Reconstruction accounts for about 34% of the profiled 4K total. The largest individual network stages there +were dec0 (300.1 us), enc0 (284.9 us), and dec1 (216.1 us). + +### Reconstruction tuning (2026-10-08) + +The Gaussian filter calls `exp2` directly on its f16 exponent, avoiding an explicit promotion to f32 and +conversion back to f16. On RX 9070 XT/RADV this produced byte-identical final output, head and all three histories +on both the five-frame capture and a five-frame synthetic stress sequence (50 binary files checked). The latter +covers fractional motion, off-screen reprojection, jitter, HDR highlights, exposure changes, mid-sequence reset, +and odd 1283x723 output dimensions. The captured reference quality checks still pass. This establishes numerical +agreement on this device, not on other drivers. + +Five alternating baseline/candidate runs per resolution, profiling disabled, 50 temporal frames per batch: + +| Case | Baseline run medians | Candidate run medians | Median paired reduction | +|---|---:|---:|---:| +| 720p capture | 582.4–604.8 us | 578.0–602.0 us | 3.1 us (about 0.5%) | +| 4K synthetic, as above | 2738.6–2781.8 us | 2732.5–2759.7 us | 23.0 us (about 0.8%) | + +The 720p runs were bimodal, so comparing unpaired medians would overstate the gain. This is a small performance +change with no observed quality change. Explicit history/filter loop unrolling did not improve the 4K stage; +forcing loops to remain rolled increased it to about 1.14 ms. Reconstruction with wave32 was also slower, +about 1.06 ms versus the wave64 baseline's 0.90–0.91 ms. Those variants are not used. + +### enc0 matrix stores and accuracy investigation (2026-10-08) + +The 32-channel enc0 block now uses padded column-major shared token tiles. This lets RADV pack +cooperative-matrix stores into wider LDS writes without changing arithmetic or accumulation order. +On RX 9070 XT, static 16-bit LDS store instructions fell from 183 to 48, allocated VGPRs from +192 to 120, and shared memory from 10,240 to 9,216 bytes. Addressing uses cooperative-matrix +layouts, with no hardware fragment mapping. Other layers keep their existing shared layout. +The analogous dec0 variants did not show a useful improvement and were discarded. + +Five alternating baseline/candidate runs, profiling disabled, 50 temporal frames per batch: + +| Case | Baseline run medians | Candidate run medians | Median paired reduction | +|---|---:|---:|---:| +| 720p capture | 595.0–610.2 us | 584.2–600.3 us | 10.2 us (about 1.7%) | +| 1920x1080 -> 2560x1440 synthetic | 1418.8–1421.9 us | 1405.0–1409.0 us | 12.7 us (about 0.9%) | +| 1920x1080 -> 3840x2160 synthetic | 2770.8–2782.5 us | 2751.8–2756.3 us | 20.6 us (about 0.7%) | + +Earlier 720p trials showed no gain, so the low-resolution benefit remains uncertain. Synthetic +cases use constant inputs and the fixed window schedule; these are engine GPU timings, not game +frame timings. Results and diagnostic variants are in `build/k-store/`, including +`final-benchmark.json`; the validated folded package is `build/k-store/final-model`. + +Head, output and all three histories were byte-identical to the baseline across five captured, +two independent captured, and five synthetic stress frames (60 files). Explicit-gain output +also remained byte-identical across the five-frame capture (25 files). All 17 native game-adapter +comparisons passed, as did nine model tests and SPIR-V validation of both packages. Eleven moving, +jittered D3D harness frames and the final output matched the baseline byte-for-byte. The independent +capture's first-frame hi-res luma scores 49.2 dB against NGX, below the default 50 dB threshold in +both baseline and candidate; this change preserves that existing discrepancy. + +Accuracy experiments covered explicit gains and f16 normalization/residual boundaries, per layer +and in combinations. Explicit enc1+dec0 improved mean RGB PSNR by 0.44 dB on the first capture, +but reduced it by 0.83 dB on the independent capture. All-explicit gains improved the first by +0.70 dB but reduced the second by 0.16 dB. No tested boundary variant improved both captures +consistently, so no accuracy variant became the default. `--unfold-gains` remains available; +the retained enc0 change preserves numerical behavior on this device. + +## Preset M (runtime and shim path exist; experimental) + +Preset M (DLSS 4.5, `rrlite_*`) is a different pipeline: an input stage that embeds 2x2-pixel tokens, a network of +4x4-token-window Swin blocks with FP8 (e4m3) weights and activations (enc1, enc2, six "tube" blocks, dec2, dec1), +a stage that expands the tokens to per-pixel reconstruction inputs, a reconstruction at 1.5 times the output +resolution, and a 3 -> 2 downsample. `M_NOTES.md` has the algorithms. Every stage has its own shader and a test +against captured launches: + +| Stage | Shader | Against the capture (640x360 -> 1280x720, 2026-10-08) | +|---|---|---| +| input | `enc0_m.comp` | motion surfaces identical; tokens 99.2% identical on the reset frame, 94.3% with history | +| network | `layer_m.comp` | per layer on the captured inputs, all ten layers and all three captured frames: byte-identical to the capture (2026-10-09) | +| expansion | `dec0_m.comp` | 96-99.6% of the 8-bit values identical; covariance 70-73 dB | +| reconstruction | `post_m.comp` | 61.5 dB on the reset frame, 76.9 dB with history; packed output 94-98% identical | +| downsample | `down_m.comp` | 70-71 dB, 92% identical | + +`test_frame_m.py` chains all of them over captured frames with the engine's own intermediate data and history. +Final image, three frames of the panning stress scene: + +| Compared with | Final RGB PSNR | +|---|---| +| CUDA backend, d4r's native (fast) M kernels | 36.2, 35.9, 35.9 dB | +| CUDA backend, ZLUDA's translation of NVIDIA's kernels (exact arithmetic) | 33.1, 29.4 dB | +| the two CUDA variants against each other | 32.8, 29.2 dB | + +The loss is in the network: its activations are FP8, one rounding flip is a 6-12% change, and the flips compound +through ten layers (the network output is at 24-30 dB whichever two implementations are compared). The engine is as +far from NVIDIA's exact arithmetic as d4r's shipping fast kernels are, no further. That was the state in 2026-10-08. +`layer_m.comp` now reproduces those ten layers bit for bit (2026-10-09): NVIDIA's binary16 rounding after every 32 +products (through a per-invocation scratch the compiler cannot fold), the native norms' summation order +(`tok_rsq` on the register element order `tok_xnat` implies, which is not the shader's own accumulator order) and +the native attention tree. Every h1 and h2 FP8 operand, both residual streams and every output of all ten layers +are byte-identical to the capture, and the whole coupled engine (front, network, back) is byte-identical to the +same engine running the native HIP network on the three captured frames. + +Current matched full-frame comparison against the old accuracy-HIP path (RX 9070 XT, 2026-10-10), fixed peak 1, +RGB only, unclipped RGBA16F outputs, eleven quality-scene frames with independent temporal feedback: + +| Case | RGB PSNR range | Reset RGB PSNR | Last-frame RGB PSNR | Maximum absolute RGB error | +|---|---:|---:|---:|---:| +| HDR, 1706×960 → 2560×1440 | 68.65–70.36 dB | 69.12 dB | 69.66 dB | 0.008789 | +| LDR, 1712×960 → 2560×1440 | 79.19–80.76 dB | 79.19 dB | 80.45 dB | 0.001465 | + +The previous feedback divergence is fixed. Native RGB10 half-history uses RGBA32F backing: packed Vulkan UNORM +bilinear filtering differs by up to three float ulps. M now uses the same float backing, with correctly rounded +`code/1023` normalization (a two-FMA residual correction; plain constant division is lowered to an inexact +reciprocal multiply). All 1024 codes and 524,288 native/Vulkan sampling coordinates match bit-for-bit. + +Native paired reconstruction rows also round the odd row's own-cell offset to binary16 **before** subtracting +one cell in binary16. A single rounding relative to the even row's cell changed Gaussian weights, first appearing +at frame nine in this scene and causing the frame-ten feedback cliff. The corrected order reproduces all three +post outputs on the failing captured frames exactly. The packed and half-mean paths retain their explicit +`precise` rounding boundaries. + +These are near-lossless comparisons, not bit-identical final output or an in-game timing claim. Raw comparisons: +`build/engine-quality/m-{hdr,ldr}-paired-f32-quality-11.raw.frameN` against +`m-{hdr,ldr}-matched-accuracy-11.raw.frameN`; metrics: `m-paired-f32-quality-11.json`. +Archived `.frame.001` scene captures are not interchangeable with this current harness scene. + +Runtime: `MModel` / `MEngine` in `runtime.h` (implemented in `runtime_m.inc`), `GameUpscalerM` in `game_m.h`, the +model directory from `compile_m.py CAPTURE OUT` (weights of one captured frame plus the compiled shaders), and the +runner `d4r-m`. Exposure is measured on the GPU every frame (preset M's formula, see `M_NOTES.md`). NGX does that whatever +the AutoExposure flag, the game's exposure texture, the pre-exposure and the exposure scale say (CUDA output is identical for +each), and the engine output is byte-identical for each too. The shim records +preset M through `engine_evaluate_m` when `[Engine] Enabled` is set and the model is in `engine\m` next to the DLL +(or `D4R_ENGINE_MODEL_M_DIR`). + +M packages now use `D4RM0003`. Rebuild with the current `compile_m.py`: the blocked FP8 operands, +RGBA32F half-history bindings and 8×64 downsample dispatch must match the runtime. `D4RM0001` and +`D4RM0002` packages are rejected rather than silently loading incompatible shaders. Native `hipnet.bin` +remains `D4RMHIP1`; its tensor and weight layout is unchanged. + +Compile `--ldr` and stage the package in `engine\m-ldr` (or set `D4R_ENGINE_MODEL_M_LDR_DIR`) for evaluations without +NGX's HDR flag. The display's HDR setting does not determine this flag. M's adapter and native HIP network are +initialized from the first evaluation's effective render extent, not the nominal feature-creation dimensions: +Ghost of Tsushima creates a 1706×960 feature but evaluates 1712×960 inputs. Subsequent M extent changes still +require a new feature. K's adapter capacity also covers the first effective extent when it exceeds the nominal size. + +D3D12 harness, engine against the CUDA backend with native M kernels (2026-10-09): + +| | Engine | CUDA backend | +|---|---:|---:| +| final RGB PSNR between them, 11 panning frames at 640x360 -> 1280x720 | 32.2-37.3 dB | | +| against the scene's ground truth | 18.64 dB | 18.52 dB | +| one evaluation, call to fence, 640x360 -> 1280x720 | 1.00 ms | 1.33 ms | +| one evaluation, call to fence, 1920x1080 -> 3840x2160 | 6.21 ms | 5.96 ms | + +That early harness measured an untuned HDR-only path, with copied inputs and no game verification. Those +limitations and the then-unverified half-history reconstruction do not describe the current implementation; +the matched temporal accuracy and actual game observations above and below supersede them. + +### The network on the native HIP layers (2026-10-09) + +The Vulkan layer shader runs preset M's network about 1.3x slower than d4r's native HIP layer kernels +(`kernels/m`, `rrlite_*_4x4.hsaco`): a paired same-session engine benchmark on the three captured frames +(1706x960 -> 2560x1440, 30 timed frames each, Vulkan and HIP alternating) measured 2.14-2.18 ms against +1.65-1.68 ms before the tiled-weight cutover. Both produced byte-identical images. +The FP8 weight packer now stores logical W[K][N] as [K/16][N/16][N16][K16]. Each 256-byte tile supplies +the shader's A[N16][K16] operand through a fixed-stride-16 row-major load, rather than a stride-N +column-major load. This is a byte permutation: FP8 values, f16 vectors, offsets, activation layouts, +front-end weights and the native HIP network package are unchanged. Compiler output replaces the +transposed global weight loads with ordinary loads. + +Paired full-network measurements of that cutover on this GPU: + +| Weight operands | Network median, first session | Network median, second session | +|---|---:|---:| +| Plain, column-major load | 2174.4 µs | 2181.4 µs | +| Tiled, row-major load | 2117.3 µs | 2119.0 µs | + +All ten native captured layers across three frames, including merged outputs, remain byte-identical. +The integrated packages also preserve all three coupled full-frame outputs with either network backend. +These are network/component proofs, not an in-game total or a fix for the longer recurrent-quality drift. +`layer_m.comp`'s remaining FP8 conversion and register-scheduling costs are optimization targets. The Vulkan FP8 +cooperative-matrix API exposes K16 on this device; these measurements do not establish a hardware limit or prove +that better Vulkan kernels cannot match HIP. + +So the network can run on those HIP kernels while the stages around it, which are faster on Vulkan than their HIP +counterparts, stay where they are (`[Engine] Network`, default `auto` = HIP when available): + +- `compile_m.py CAPTURE OUT --hip KERNEL_DIR` adds `hipnet.bin` (the ten layers' weight allocations as NGX + uploads them, shifts and grid divisors), `hip/*.hsaco`, and the input and expansion stages compiled with + `-DPLANES=1` (token buffers in NGX's layout, two planes of 32 channels, which the kernels read and write). +- `MEngine` takes two caller-owned buffers (`ExternalNetwork`) and records a frame in two parts, `recordFront` + (exposure, input stage) and `recordBack` (expansion, reconstruction, downsample). +- `hip_net.h` (`HipNetwork`, native Linux, used by `d4r-m ... REPEATS hip`) and the CUDA bridge's `d4rEngineNet*` + (the same logic in C for the game path) load the code objects with HIP directly, upload the weights once, run each + layer's weight-preparation kernel once (the tube keeps one prepared image per weights pointer), and launch the ten + layers per frame. No NGX, no PTX and no ZLUDA translation are involved; the bridge's library only supplies the + HIP context. +- In the shim the game's command list is split at the network: front, frame marker, + `SplitCommandListForExternalWait`, back. The native HIP stream waits for the front's GPU-published frame number. + A native HIP completion callback signals the Vulkan timeline directly; Wine-worker polling retires resources + but is not needed to release the back half. The game still shows the frame it just rendered. This needs the + d4r vkd3d-proton and VRAM interop; otherwise the Vulkan network is used and the log says why. +- The native network uses a nonblocking stream at the device's greatest supported priority: it is on the frame's + critical path between the two Vulkan parts. This changes scheduling, not the layer arithmetic. The front's + invocation-private half-rounding scratch uses memory fences without subgroup execution barriers; all three + coupled HDR frame outputs remained byte-identical to the prior front after this change. + +The input GEMM keeps four accumulators live and shares each activation load across them. Captured tokens and +motion outputs remain exact in both token layouts. A further store-epilogue change hoists each lane's token +base address and uses one uniform bounds check for fully in-screen blocks; border blocks retain individual +row checks. Paired RX 9070 XT probes at 1706×960 measured 152.2 → 144.2 µs for interleaved tokens and +156.1 → 146.1 µs for plane-layout tokens. All eleven fresh captured frames preserved both output buffers +byte-for-byte in each layout. This is a stage result, not an in-game total. +The integrated `D4RM0002` runtime and both model packages also preserved all 22 complete independent-feedback +HDR/LDR frame outputs after the address hoist. Current HDR/LDR adapter checks passed all 114 six-frame cases, +including queued mixed-format output and late per-input fallback. + +The native accuracy network pairs adjacent K16 operand loads in `dec1` and `enc3_tube`. The tube's MLP uses +two-chunk unrolling without changing accumulation order. Six alternating, clean network-only runs measured +median 1630.05 → 1611.05 µs (19 µs, 1.17%). All 22 complete HDR/LDR independent-feedback outputs remained +byte-identical after this native scheduling change. This is not an in-game total. + +The expansion shares each activation tile across three accumulator tiles, then the remaining two, while +retaining the binary16 boundary after each K32 chunk. Six paired 32-lane replay rounds preserved all four +output buffers in both token layouts; native plane-layout medians were 100.55 → 97.95 µs. + +Reconstruction uses 128 invocations per 16×16-pixel block, with each invocation owning a vertical pixel pair. +It materializes the vertical column sum before adding the adjacent column, retaining the native mean's +association and odd-row residual rounding. Captured HDR/LDR outputs and 15 partial-edge geometries remained +byte-identical. An initial isolated timing gain was not reproducible, so no stage-speed claim is made for this +change; acceptance depends on the integrated in-game total. + +Downsampling uses an 8×64 output block, reducing the repeated horizontal filtering of border rows. All decoded +values, filter associations, tap-range clamps and final output rounding are unchanged. Bounds are checked +after both shared barriers against the feature output, not the borrowed image's potentially larger extent. +The combined `D4RM0003` shaders/runtime preserved all 22 complete HDR/LDR independent-feedback outputs, +byte-for-byte against the preceding native-tuned package. Weight payloads and offsets after the magic are +unchanged. The current runtime also passed 114 direct/fallback adapter cases, each with six temporal frames. +The standalone three-frame replay retained exact front tokens, network outputs and all expansion outputs; +its final output remained 68.8–71.5 dB against the captured accuracy HIP output. + +Actual Ghost of Tsushima HUD observations (2026-10-10), preset 13/M, Quality, effective +1712×960 → 2560×1440, LDR, direct RGB10A2 output: + +| Observation | Network | Current upscaler time | Average upscaler time | +|---|---|---:|---:| +| Verified gameplay before compact reconstruction tier | Native HIP | 2.61 ms | 2.63 ms | +| Later user HUD with compact tier; scene not shown in crop | Native HIP | 2.89 ms | 2.89 ms | +| User menu HUD with compact tier, no HIP handoff | Vulkan | 3.05 ms | 3.08 ms | +| Verified gameplay with corrected float half-history and paired-row rounding, before store-address hoisting | Native HIP | 2.59 ms | 2.53 ms | +| Verified gameplay with tiled Vulkan weights, corrected histories and store-address hoisting | Vulkan | 2.85 ms | 2.87 ms | +| Verified gameplay with corrected histories and store-address hoisting, before the final native load-scheduling change | Native HIP | 2.49 ms | 2.52 ms | +| Verified gameplay after the final native load-scheduling change, first observation | Native HIP | 2.52 ms | 2.50 ms | +| Same gameplay after settling, second observation | Native HIP | 2.51 ms | 2.51 ms | +| Verified gameplay with the combined `D4RM0003` reconstruction kernels | Native HIP | 2.44 ms | 2.47 ms | +| Same `D4RM0003` gameplay after settling, second observation | Native HIP | 2.49 ms | 2.49 ms | +| Final shared M/K runtime recheck, first gameplay observation | Native HIP | 2.47 ms | 2.54 ms | +| Final shared M/K runtime recheck, settled observation | Native HIP | 2.59 ms | 2.56 ms | + +The combined `D4RM0003` package reached the 2.5 ms total in-game target in the earlier gameplay samples. +The final shared-runtime recheck averaged 2.54–2.56 ms, so a sustained ≤2.5 ms claim is **not** established. +The final samples are `build/engine-quality/ghost-m-final-gameplay{,-settled}.png`; the log confirms preset +13/M, Quality, effective 1712×960 → 2560×1440, native HIP network and direct RGB10A2 output. +No scene-independent timing guarantee is implied, and these measurements do not isolate whether the small +change comes from shared-runtime overhead, GPU clock state or the evolving scene. +The compact 11×11 reconstruction tile and conditional fallback fetch save about 26 µs in paired LDR stage +probes at this geometry, but that isolated saving alone did **not** establish an in-game improvement. +Removing the HIP handoff did not improve the observed total either; the native-network engine remains +the default. These observations do not isolate each new kernel's contribution. +The menu is not an acceptance scene; stage sums and callback launch-to-completion spans are not substitutes +for the HUD number. The runtime retains 10×10 for render/reconstruction ≤ 1/3 and 15×15 for wider ratios; +11×11 requires both ratios ≤ 0.45 and both jitter components in [-0.5, 0.5]. +The compact tier and conditional fallback fetch preserved every output byte over two complete +eleven-frame feedback sequences (HDR and LDR), compared with the preceding engine package. The subsequent +float-history and paired-row rounding corrections fix the accuracy-HIP divergence independently of that tier. + +Earlier measurements, before the current accuracy and output changes, on an RX 9070 XT: +`d4r-m`, synthetic frames, wall clock from recording to completion (GPU stage sum): + +| | Vulkan network | HIP network | +|---|---:|---:| +| 1280x720 -> 3840x2160 | 3.22 ms (2.78) | 2.39 ms (2.19) | +| 1920x1080 -> 3840x2160 | 5.52 ms (5.16) | 3.81 ms (3.62) | +| network alone at 1080p | 3.57 ms | 2.07 ms | + +D3D12 harness, one evaluation from the call to the fence, every path showing the frame's own result: + +| | Engine, HIP network | Engine, Vulkan network | CUDA backend, native kernels | +|---|---:|---:|---:| +| 1280x720 -> 2560x1440 | 2.19 ms | 2.91 ms | 2.87 ms | +| 1280x720 -> 3840x2160 | 3.32 ms | 3.91 ms | 4.00 ms | +| 1920x1080 -> 3840x2160 | 4.54 ms | 6.23 ms | 5.38 ms | + +Those earlier harness images measured 44–47 dB against the CUDA backend (HDR) and 36–37 dB without the HDR +flag; they are not current accuracy results. Feature release and re-creation at alternating render sizes ran +clean. The current game timing and longer accuracy comparison are reported above. The kernels are built per +GPU architecture (gfx1201 with FP8 here); another GPU falls back to Vulkan when they do not load. + +## Files + +| File | Role | +|---|---| +| `runtime.h`, `runtime.cpp` | Embeddable model/engine, persistent resources, frame parameter helpers, temporal loop | +| `game.h`, `game.cpp` | Adapter for a game's textures: format conversion, frame, output blit, in one command buffer | +| `vk_dispatch.h`, `vk_dispatch.cpp` | Vulkan function table for builds without a Vulkan import library (the MinGW shim) | +| `replay.cpp`, `build.sh` | Full-frame execution, capability checks, correctness readbacks, GPU timing | +| `extract_k.py`, `compile_k.py` | Weight extraction, validation, model/shader packaging | +| `capture.py`, `prepare_k.py` | Capture decoding and complete-frame reference case preparation | +| `prepare_frames.py` | Ordinary raw HDR frame preparation without NGX parameter blocks | +| `input_k.comp`, `layer.comp`, `output_k.comp` | Own input, fused network, and reconstruction shaders | +| `m_model.py`, `layer_m.comp`, `test_m.py` | Preset M: logical weights and numpy model, layer shader, tiled FP8 packer and per-layer test | +| `hip_net.h`, `hip_net.cpp` | Shared M/K native HIP network loading and immutable graph execution, for native Linux programs | +| `enc0_m.comp`, `dec0_m.comp`, `post_m.comp`, `down_m.comp` | Preset M: input, expansion, reconstruction and downsample stages | +| `activation_m.glsl`, `gen_activation_m.py` | Preset M: sigmoid/tanh/gate table of the 256 E4M3 codes and the script that computes it (`--check` verifies the table) | +| `test_enc0_m.py`, `test_dec0_m.py`, `test_post_m.py`, `test_down_m.py`, `test_frame_m.py` | Preset M: per-stage and whole-frame tests | +| `M_NOTES.md` | Preset M: the algorithms as reverse-engineered | +| `exposure_k.comp` | Automatic exposure: log-luma block sums, then the exposure state and the stages' exposure fields | +| `test_frame.py`, `test_model.py` | Temporal result checks and model corruption tests | +| `test_game.cpp` | Queued direct-texture, rectangle, layout, conversion and precision regression checks | +| `test_k.py`, `test_in*.py`, `test_out*.py` | Isolated shader diagnosis and reference comparisons | + +`vulkan/ptx2glsl.py` and `vulkan/kvk.cpp` remain reference/research tools. They are not the engine backend. diff --git a/engine/activation_m.glsl b/engine/activation_m.glsl new file mode 100644 index 0000000..552c8a7 --- /dev/null +++ b/engine/activation_m.glsl @@ -0,0 +1,270 @@ +// Expansion activations are functions of the 256 E4M3 codes, not arbitrary floats. Each row packs binary16 +// sigmoid, tanh and the separate-half-op gate .5*tanh(.5*x)+.5. This avoids GLSL half transcendental lowering, +// which differs from the accuracy reference by one half ulp and then crosses normalized-store thresholds. +// Computed by gen_activation_m.py from those formulas alone (sigmoid's f32 exp2(-1.442695*x), correctly rounded +// binary16 tanh); `gen_activation_m.py --check` verifies this table. All four expansion surfaces replay +// bit-identically against the captured reset frame with it. +const uvec2 ACT[256] = uvec2[256]( + uvec2(0x00003800u, 0x00003800u), + uvec2(0x18003801u, 0x00003801u), + uvec2(0x1c003802u, 0x00003802u), + uvec2(0x1e003803u, 0x00003803u), + uvec2(0x20003804u, 0x00003804u), + uvec2(0x21003805u, 0x00003805u), + uvec2(0x22003806u, 0x00003806u), + uvec2(0x23003807u, 0x00003807u), + uvec2(0x24003808u, 0x00003808u), + uvec2(0x24803809u, 0x00003809u), + uvec2(0x2500380au, 0x0000380au), + uvec2(0x2580380bu, 0x0000380bu), + uvec2(0x2600380cu, 0x0000380cu), + uvec2(0x2680380du, 0x0000380du), + uvec2(0x2700380eu, 0x0000380eu), + uvec2(0x277f380fu, 0x0000380fu), + uvec2(0x27ff3810u, 0x00003810u), + uvec2(0x28803812u, 0x00003812u), + uvec2(0x28ff3814u, 0x00003814u), + uvec2(0x297f3816u, 0x00003816u), + uvec2(0x29ff3818u, 0x00003818u), + uvec2(0x2a7f381au, 0x0000381au), + uvec2(0x2afe381cu, 0x0000381cu), + uvec2(0x2b7e381eu, 0x0000381eu), + uvec2(0x2bfd3820u, 0x00003820u), + uvec2(0x2c7e3824u, 0x00003824u), + uvec2(0x2cfd3828u, 0x00003828u), + uvec2(0x2d7d382cu, 0x0000382cu), + uvec2(0x2dfc3830u, 0x00003830u), + uvec2(0x2e7a3834u, 0x00003834u), + uvec2(0x2ef93838u, 0x00003838u), + uvec2(0x2f77383cu, 0x0000383cu), + uvec2(0x2ff53840u, 0x00003840u), + uvec2(0x30783848u, 0x00003848u), + uvec2(0x30f63850u, 0x00003850u), + uvec2(0x31723858u, 0x00003858u), + uvec2(0x31ee3860u, 0x00003860u), + uvec2(0x32693868u, 0x00003868u), + uvec2(0x32e43870u, 0x00003870u), + uvec2(0x335e3877u, 0x00003877u), + uvec2(0x33d6387fu, 0x0000387fu), + uvec2(0x3463388fu, 0x0000388fu), + uvec2(0x34d8389fu, 0x0000389fu), + uvec2(0x354b38aeu, 0x000038aeu), + uvec2(0x35bc38beu, 0x000038beu), + uvec2(0x362a38cdu, 0x000038cdu), + uvec2(0x369638dcu, 0x000038dcu), + uvec2(0x36ff38ecu, 0x000038ecu), + uvec2(0x376538fbu, 0x000038fbu), + uvec2(0x38143919u, 0x00003919u), + uvec2(0x38703936u, 0x00003936u), + uvec2(0x38c53953u, 0x00003953u), + uvec2(0x3915396fu, 0x0000396fu), + uvec2(0x395e398bu, 0x0000398au), + uvec2(0x39a239a5u, 0x000039a6u), + uvec2(0x39df39c0u, 0x000039c0u), + uvec2(0x3a1839d9u, 0x000039d9u), + uvec2(0x3a793a0au, 0x00003a0au), + uvec2(0x3ac93a38u, 0x00003a38u), + uvec2(0x3b0a3a63u, 0x00003a62u), + uvec2(0x3b3e3a8au, 0x00003a8au), + uvec2(0x3b673aafu, 0x00003aafu), + uvec2(0x3b883ad1u, 0x00003ad1u), + uvec2(0x3ba23af0u, 0x00003af0u), + uvec2(0x3bb63b0cu, 0x00003b0cu), + uvec2(0x3bd33b3du, 0x00003b3cu), + uvec2(0x3be53b65u, 0x00003b64u), + uvec2(0x3bef3b85u, 0x00003b85u), + uvec2(0x3bf63b9fu, 0x00003b9fu), + uvec2(0x3bfa3bb4u, 0x00003bb4u), + uvec2(0x3bfc3bc4u, 0x00003bc4u), + uvec2(0x3bfe3bd1u, 0x00003bd1u), + uvec2(0x3bff3bdbu, 0x00003bdbu), + uvec2(0x3bff3be9u, 0x00003beau), + uvec2(0x3c003bf2u, 0x00003bf2u), + uvec2(0x3c003bf8u, 0x00003bf8u), + uvec2(0x3c003bfbu, 0x00003bfbu), + uvec2(0x3c003bfdu, 0x00003bfdu), + uvec2(0x3c003bfeu, 0x00003bfeu), + uvec2(0x3c003bffu, 0x00003bffu), + uvec2(0x3c003bffu, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x3c003c00u, 0x00003c00u), + uvec2(0x7e007e00u, 0x00007e00u), + uvec2(0x80003800u, 0x00003800u), + uvec2(0x980037feu, 0x000037feu), + uvec2(0x9c0037fcu, 0x000037fcu), + uvec2(0x9e0037fau, 0x000037fau), + uvec2(0xa00037f8u, 0x000037f8u), + uvec2(0xa10037f6u, 0x000037f6u), + uvec2(0xa20037f4u, 0x000037f4u), + uvec2(0xa30037f2u, 0x000037f2u), + uvec2(0xa40037f0u, 0x000037f0u), + uvec2(0xa48037eeu, 0x000037eeu), + uvec2(0xa50037ecu, 0x000037ecu), + uvec2(0xa58037eau, 0x000037eau), + uvec2(0xa60037e8u, 0x000037e8u), + uvec2(0xa68037e6u, 0x000037e6u), + uvec2(0xa70037e4u, 0x000037e4u), + uvec2(0xa77f37e2u, 0x000037e2u), + uvec2(0xa7ff37e0u, 0x000037e0u), + uvec2(0xa88037dcu, 0x000037dcu), + uvec2(0xa8ff37d8u, 0x000037d8u), + uvec2(0xa97f37d4u, 0x000037d4u), + uvec2(0xa9ff37d0u, 0x000037d0u), + uvec2(0xaa7f37ccu, 0x000037ccu), + uvec2(0xaafe37c8u, 0x000037c8u), + uvec2(0xab7e37c4u, 0x000037c4u), + uvec2(0xabfd37c0u, 0x000037c0u), + uvec2(0xac7e37b8u, 0x000037b8u), + uvec2(0xacfd37b0u, 0x000037b0u), + uvec2(0xad7d37a8u, 0x000037a8u), + uvec2(0xadfc37a0u, 0x000037a0u), + uvec2(0xae7a3798u, 0x00003798u), + uvec2(0xaef93790u, 0x00003790u), + uvec2(0xaf773788u, 0x00003788u), + uvec2(0xaff53780u, 0x00003780u), + uvec2(0xb0783770u, 0x00003770u), + uvec2(0xb0f63760u, 0x00003760u), + uvec2(0xb1723750u, 0x00003750u), + uvec2(0xb1ee3741u, 0x00003740u), + uvec2(0xb2693731u, 0x00003731u), + uvec2(0xb2e43721u, 0x00003721u), + uvec2(0xb35e3711u, 0x00003711u), + uvec2(0xb3d63701u, 0x00003701u), + uvec2(0xb46336e2u, 0x000036e2u), + uvec2(0xb4d836c3u, 0x000036c2u), + uvec2(0xb54b36a3u, 0x000036a4u), + uvec2(0xb5bc3684u, 0x00003684u), + uvec2(0xb62a3666u, 0x00003666u), + uvec2(0xb6963647u, 0x00003647u), + uvec2(0xb6ff3629u, 0x00003628u), + uvec2(0xb765360au, 0x0000360au), + uvec2(0xb81435cfu, 0x000035ceu), + uvec2(0xb8703594u, 0x00003594u), + uvec2(0xb8c5355au, 0x0000355au), + uvec2(0xb9153522u, 0x00003522u), + uvec2(0xb95e34ebu, 0x000034ebu), + uvec2(0xb9a234b5u, 0x000034b5u), + uvec2(0xb9df3481u, 0x00003480u), + uvec2(0xba18344eu, 0x0000344eu), + uvec2(0xba7933d8u, 0x000033d8u), + uvec2(0xbac93320u, 0x00003320u), + uvec2(0xbb0a3275u, 0x00003276u), + uvec2(0xbb3e31d6u, 0x000031d6u), + uvec2(0xbb673144u, 0x00003144u), + uvec2(0xbb8830bdu, 0x000030bcu), + uvec2(0xbba23041u, 0x00003042u), + uvec2(0xbbb62fa1u, 0x00002fa0u), + uvec2(0xbbd32e1au, 0x00002e1cu), + uvec2(0xbbe52cdbu, 0x00002cdcu), + uvec2(0xbbef2bb1u, 0x00002bb0u), + uvec2(0xbbf62a12u, 0x00002a10u), + uvec2(0xbbfa28c7u, 0x000028c8u), + uvec2(0xbbfc2781u, 0x00002780u), + uvec2(0xbbfe25e2u, 0x000025e0u), + uvec2(0xbbff249bu, 0x000024a0u), + uvec2(0xbbff21a0u, 0x000021a0u), + uvec2(0xbc001edau, 0x00001ec0u), + uvec2(0xbc001c2bu, 0x00001c40u), + uvec2(0xbc001910u, 0x00001900u), + uvec2(0xbc001626u, 0x00001600u), + uvec2(0xbc001377u, 0x00001400u), + uvec2(0xbc001087u, 0x00001000u), + uvec2(0xbc000d7fu, 0x00000c00u), + uvec2(0xbc00080bu, 0x00000c00u), + uvec2(0xbc0002fau, 0x00000000u), + uvec2(0xbc000118u, 0x00000000u), + uvec2(0xbc000067u, 0x00000000u), + uvec2(0xbc000026u, 0x00000000u), + uvec2(0xbc00000eu, 0x00000000u), + uvec2(0xbc000005u, 0x00000000u), + uvec2(0xbc000002u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0xbc000000u, 0x00000000u), + uvec2(0x7e00fe00u, 0x00007e00u) +); +uint actIndex(float16_t v) { uint h = uint(float16BitsToUint16(v)); uint a = h & 32767u; if (a > 0x7c00u) return ((h >> 8u) & 128u) | 127u; return ((h >> 8u) & 128u) | (a >= 0x2400u ? (a >> 7u) - 64u : uint(abs(float(v)) * 512.0)); } +float16_t sig16(float16_t v) { return uint16BitsToFloat16(uint16_t(ACT[actIndex(v)].x & 65535u)); } +float16_t tan16(float16_t v) { return uint16BitsToFloat16(uint16_t(ACT[actIndex(v)].x >> 16u)); } +float16_t gate16(float16_t v) { return uint16BitsToFloat16(uint16_t(ACT[actIndex(v)].y)); } +f16vec4 tan16(f16vec4 v) { return f16vec4(tan16(v.x),tan16(v.y),tan16(v.z),tan16(v.w)); } +f16vec4 gate16(f16vec4 v) { return f16vec4(gate16(v.x),gate16(v.y),gate16(v.z),gate16(v.w)); } diff --git a/engine/build.sh b/engine/build.sh new file mode 100755 index 0000000..02068a4 --- /dev/null +++ b/engine/build.sh @@ -0,0 +1,20 @@ +#!/usr/bin/env bash +set -euo pipefail +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +OUT="${1:-$ROOT/build/native-engine}" +CXX="${CXX:-g++}" +mkdir -p "$OUT" +# The native HIP network (hip_net.h) is shared by both runners when ROCm is installed; without it they fall back +# to the engine's own Vulkan network. The library itself never links HIP. +ROCM="${D4R_ROCM_DIR:-/opt/rocm}" +HIP=() +if [[ -f "$ROCM/include/hip/hip_runtime_api.h" ]]; then + HIP=(-DD4R_ENGINE_HIP "$ROOT/engine/hip_net.cpp" -I"$ROCM/include" -L"$ROCM/lib" -lamdhip64 "-Wl,-rpath,$ROCM/lib") +fi +"$CXX" -std=c++17 -O3 -Wall -Wextra -Wno-missing-field-initializers -fPIC -shared \ + "$ROOT/engine/runtime.cpp" "$ROOT/engine/game.cpp" "$ROOT/engine/game_m.cpp" -lvulkan -o "$OUT/libd4r_engine.so" +"$CXX" -std=c++17 -O3 -Wall -Wextra -Wno-missing-field-initializers \ + "$ROOT/engine/replay.cpp" "${HIP[@]}" -L"$OUT" -ld4r_engine -lvulkan '-Wl,-rpath,$ORIGIN' -o "$OUT/d4r-k" +"$CXX" -std=c++17 -O3 -Wall -Wextra -Wno-missing-field-initializers \ + "$ROOT/engine/replay_m.cpp" "${HIP[@]}" -L"$OUT" -ld4r_engine -lvulkan '-Wl,-rpath,$ORIGIN' -o "$OUT/d4r-m" +printf 'Built %s and %s\n' "$OUT/d4r-k" "$OUT/d4r-m" diff --git a/engine/capture.py b/engine/capture.py new file mode 100644 index 0000000..96bcf01 --- /dev/null +++ b/engine/capture.py @@ -0,0 +1,130 @@ +"""Capture decoding shared by full-frame validation and model packaging. + +Captures are reference data only. The runtime has no dependency on NGX, CUDA, +ZLUDA, PTX or the pointer addresses contained in a capture. +""" +from dataclasses import dataclass +from pathlib import Path +import re +import struct +import numpy as np + +INPUT = 'hiluma_engine_input_depthinv_mvlo_hdr_v2_rel' +OUTPUT = 'hiluma_engine_output_depthinv_mvlo_hdr_max_v2_rel' +ORDER = ['enc0', 'enc1', 'enc2', 'enc3', 'enc4', 'dec5', 'dec4', 'dec3', 'dec2', 'dec1', 'dec0'] + + +@dataclass +class Resource: + path: Path + width: int + height: int + channels: int + bits: int + + def validate(self): + if self.path.stat().st_size != self.width * self.height * self.channels * self.bits // 8: + raise ValueError(f'truncated texture/surface: {self.path}') + + +class Launch: + def __init__(self, capture, directory): + self.capture, self.directory = Path(capture), Path(directory) + self.sequence = int(re.search(r'replay-(\d+)-', self.directory.name).group(1)) + self.args = (self.directory / 'args.bin').read_bytes() + self.manifest = (self.directory / 'manifest.txt').read_text().splitlines() + self.kernel = self.manifest[0].split()[1] + self.resources = {} + for f in (self.capture / 'dump').glob(f'launch-{self.sequence:03d}-{self.kernel}-arg*'): + match = re.search(r'-arg(\d+)-(texture|surface)-(\d+)x(\d+)x(\d+)-fmt(\d+)\.bin$', f.name) + if match: + offset, kind, w, h, c, fmt = match.groups() + resource = Resource(f.resolve(), int(w), int(h), int(c), 32 if fmt == '32' else 16) + self.resources[int(offset)] = resource + + def f2(self, offset): return struct.unpack_from('<2f', self.args, offset) + def i2(self, offset): return struct.unpack_from('<2i', self.args, offset) + def f1(self, offset): return struct.unpack_from(' ptr] + end = min(following) if following else base + size + with open(path, 'rb') as f: + f.seek(ptr - base) + w = bytearray(f.read(end - ptr)) + pad = (-len(w)) % 64 + role = HIP_ROLE.get(name) + sx, sy = hip_shift(role, 0) if role is not None else (4, 4) + head += struct.pack('<32siiIIQQ', HIP_KERNEL.format(name).encode(), sx, sy, HIP_DIVISOR[name], 0, at, len(w)) + blobs += [w, bytes(pad)]; at += len(w) + pad + (output / 'hipnet.bin').write_bytes(head + b''.join(blobs)) + (output / 'hip').mkdir(exist_ok=True) + for name in HIP_ORDER: + src = kernels / f'{HIP_KERNEL.format(name)}.hsaco' + if not src.exists(): + raise ValueError(f'{src} is missing') + (output / 'hip' / src.name).write_bytes(src.read_bytes()) + + +def tensor_spans(e): + c, h, x = e['C'], e['H'], e['X'] + for key in ('g1', 'bo', 'g2', 'b2'): + yield e[key], 16 * c + for key, size in (('qkv', c * 32), ('table', 64 * 64), ('wo', 32 * c), + ('w1', c * 32), ('b1', 16 * 32), ('w2', 32 * c)): + expected = {'qkv': h * 3, 'table': h, 'wo': h, 'w1': c // 8, 'b1': c // 8, 'w2': c // 8}[key] + if len(e[key]) != expected: + raise ValueError(f'{key}: expected {expected} tensors') + for offset in e[key]: + yield offset, size + for key, size in {'embed_w': 16 * c, 'embed_b': 16 * c, + 'merge_w': 4 * c * e.get('merge_n', 0), 'merge_b': 16 * e.get('merge_n', 0), + 'expand_w': x * 4 * c, 'expand_b': 16 * 4 * c, + 'head_w': c * 48, 'head_b': 16 * 48}.items(): + if key in e: + yield e[key], size + + +def validate(index, weight_bytes): + if list(index) != list(LAYERS): + raise ValueError('expected the eleven preset K layers in network order') + if weight_bytes <= 0 or weight_bytes % 2: + raise ValueError('weights must be nonempty f16 data') + for name, (kind, heads, c, x, pos) in LAYERS.items(): + e = index[name] + if (e['kind'], e['H'], e['C'], e['X'], e['posattn']) != (kind, heads, c, x, pos): + raise ValueError(f'{name}: incompatible layer shape') + required = ('embed_w', 'embed_b') if kind == 'enc0' else () + required += ('merge_w', 'merge_b', 'merge_n') if kind in ('enc', 'enc0') else () + required += ('expand_w', 'expand_b') if kind in ('dec', 'dec0') else () + required += ('head_w', 'head_b') if kind == 'dec0' else () + for key in required: + if key not in e: + raise ValueError(f'{name}: missing {key}') + if kind in ('enc', 'enc0') and e['merge_n'] != (x + 15) // 16 * 16: + raise ValueError(f'{name}: invalid merge stride') + for offset, size in tensor_spans(e): + if not isinstance(offset, int) or offset < 0 or offset % 16 or (offset + size) * 2 > weight_bytes: + raise ValueError(f'{name}: tensor at {offset} with {size} elements is outside/alignment-invalid') + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument('model_prefix', type=Path) + p.add_argument('lut', type=Path) + p.add_argument('output', type=Path) + p.add_argument('--ldr', action='store_true', help='display-referred NGX LDR input and reconstruction') + p.add_argument('--hip', type=Path, metavar='KERNEL_DIR', help='also package the network for the native HIP layers: ' + 'hipnet.bin (the eleven layers and their weight allocations as NGX uploads them) and ' + 'hip/dltss_pwin_*_layer.hsaco from KERNEL_DIR (kernels/build.sh output)') + p.add_argument('--capture', type=Path, metavar='CAPTURE_DIR', help='with --hip: the dump directory holding one ' + 'captured frame of the eleven dltss_pwin_*_layer launches (D4R_CUDA_REPLAY_DUMP_DIR)') + a = p.parse_args() + if a.hip and a.capture is None: + p.error('--hip needs --capture (the weights come from a captured frame)') + weights = Path(str(a.model_prefix) + '.bin') + index = json.loads(Path(str(a.model_prefix) + '.json').read_text()) + validate(index, weights.stat().st_size) + # All sixteen tables (prepare_k.py's lut.bin; NGX picks one by upscaling ratio) or, from older captures, + # the single table of the captured ratio, which the engine then uses for every ratio. + if a.lut.stat().st_size not in (32768, 16 * 32768): + raise ValueError('the K reconstruction tables must be one or sixteen tables of 2048 sixteen-byte entries') + if a.lut.stat().st_size == 32768: + print('warning: one reconstruction table only; output at other upscaling ratios will not match NGX', file=sys.stderr) + a.output.mkdir(parents=True, exist_ok=True) + shutil.copyfile(weights, a.output / 'weights.bin') + shutil.copyfile(a.lut, a.output / 'lut.bin') + # Only offsets; dimensions and window shifts belong to the Engine, not the Model. + class ZeroShape: + def i32x2(self, offset): + return (0, 0) + offsets = b''.join(params(ZeroShape(), e)[0][16:] for e in index.values()) + (a.output / 'offsets.bin').write_bytes(MAGIC + struct.pack(' pixels of the 3x image (the motion vectors' unit) + vec2 hiSize; + vec2 histOffs; // added to the history uv, in render pixels + float sigma, rho; + vec2 renderInv; // host-rounded texture coordinates, distinct from GPU rcp.approx offsets +} p; +layout(set = 0, binding = 1) uniform sampler2D motionTex; // texelFetch +layout(set = 0, binding = 2) uniform sampler2D featTex; // the previous feature, bilinear +layout(set = 0, binding = 3, rgba8_snorm) uniform writeonly image2D featOut; +layout(set = 0, binding = 4, rgba8) uniform writeonly image2D gateOut; +layout(set = 0, binding = 5, rgba8) uniform writeonly image2D mixOut; +layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D covOut; +layout(set = 0, binding = 7, std430) readonly buffer Tokens { floate4m3_t a[]; } tok; // [tokH][tokW][64] +layout(set = 0, binding = 8, std430) readonly buffer W8 { floate4m3_t w[]; } w8; // [64][80]: 20 columns per sub-position +layout(set = 0, binding = 9, std430) readonly buffer W16 { f16vec4 w[]; } w16; // bias, 16 equal rows of 80 +#define ACC coopmat +#define HACC coopmat +#define A8 coopmat +#define B8 coopmat +#define RM gl_CooperativeMatrixLayoutRowMajor +shared floate4m3_t TK[16 * 72]; +shared f16vec4 EO[16 * 88 / 4]; +float16_t q8(float16_t v) { return float16_t(floate4m3_t(clamp(v, float16_t(-448.0), float16_t(448.0)))); } +f16vec4 q8(f16vec4 v) { return f16vec4(q8(v.x), q8(v.y), q8(v.z), q8(v.w)); } +// NVIDIA's expansion GEMM is chained m16n8k32 mma with an f16 accumulator: the bias is the initial C and every +// 32 products are summed with one f16 rounding (swin_model.py mma_chain; the same k32 step the native kernels' +// SWIN_EXACT build implements). K32_ROUND reproduces that: after each pair of k16 cooperative products the +// accumulator is rounded to f16 through the shared EO region (ROUND_HALF_STORE). +#ifndef K32_ROUND +#define K32_ROUND 1 +#endif +#if K32_ROUND +#define R32(m, column) ROUND_HALF_STORE(m, EO, (column) * 4u, 22u) +#else +#define R32(m, column) +#endif +#ifndef PLANES +#define PLANES 0 +#endif +void main() +{ + const uint lane = gl_LocalInvocationID.x, row = lane & 15u, hf = lane >> 4; + const ivec2 t0 = 4 * ivec2(gl_WorkGroupID.xy); +#if PLANES + // NGX's layout, two planes of 32 channels (the HIP network's output); the dispatch covers the token grid + const uint src = (hf * 16u * gl_NumWorkGroups.x * gl_NumWorkGroups.y + uint((t0.y + int(row >> 2)) * p.tokW + t0.x + int(row & 3u))) * 32u; +#else + const uint src = uint((t0.y + int(row >> 2)) * p.tokW + t0.x + int(row & 3u)) * 64u + hf * 32u; +#endif + for (uint c = 0u; c < 32u; ++c) TK[row * 72u + hf * 32u + c] = tok.a[src + c]; + memoryBarrierShared(); barrier(); + // Share each activation tile across three output tiles, then the remaining two. + // Five live accumulators cost more registers than the saved loads on gfx1201. + { + HACC b0; coopMatLoad(b0, w16.w, 0u, 20u, RM); + ACC a0 = ACC(b0); + HACC b1; coopMatLoad(b1, w16.w, 4u, 20u, RM); + ACC a1 = ACC(b1); + HACC b2; coopMatLoad(b2, w16.w, 8u, 20u, RM); + ACC a2 = ACC(b2); + for (uint kt = 0u; kt < 4u; ++kt) + { + A8 a; coopMatLoad(a, TK, kt * 16u, 72u, RM); + B8 w0; coopMatLoad(w0, w8.w, kt * 16u * 80u, 80u, RM); + a0 = coopMatMulAdd(a, w0, a0); + B8 w1; coopMatLoad(w1, w8.w, kt * 16u * 80u + 16u, 80u, RM); + a1 = coopMatMulAdd(a, w1, a1); + B8 w2; coopMatLoad(w2, w8.w, kt * 16u * 80u + 32u, 80u, RM); + a2 = coopMatMulAdd(a, w2, a2); + if ((kt & 1u) != 0u) { R32(a0, 0u); R32(a1, 1u); R32(a2, 2u); } + } + coopMatStore(HACC(a0), EO, 0u, 22u, RM); + coopMatStore(HACC(a1), EO, 4u, 22u, RM); + coopMatStore(HACC(a2), EO, 8u, 22u, RM); + } + { + HACC b0; coopMatLoad(b0, w16.w, 12u, 20u, RM); + ACC a0 = ACC(b0); + HACC b1; coopMatLoad(b1, w16.w, 16u, 20u, RM); + ACC a1 = ACC(b1); + for (uint kt = 0u; kt < 4u; ++kt) + { + A8 a; coopMatLoad(a, TK, kt * 16u, 72u, RM); + B8 w0; coopMatLoad(w0, w8.w, kt * 16u * 80u + 48u, 80u, RM); + a0 = coopMatMulAdd(a, w0, a0); + B8 w1; coopMatLoad(w1, w8.w, kt * 16u * 80u + 64u, 80u, RM); + a1 = coopMatMulAdd(a, w1, a1); + if ((kt & 1u) != 0u) { R32(a0, 3u); R32(a1, 4u); } + } + coopMatStore(HACC(a0), EO, 12u, 22u, RM); + coopMatStore(HACC(a1), EO, 16u, 22u, RM); + } + memoryBarrierShared(); barrier(); + for (uint pi = lane; pi < 64u; pi += 32u) + { + const uint px = pi & 7u, py = pi >> 3; + const ivec2 at = 2 * t0 + ivec2(px, py); + if (at.x >= p.size.x || at.y >= p.size.y) continue; + const uint e = (((py >> 1) * 4u + (px >> 1)) * 88u + ((py & 1u) * 2u + (px & 1u)) * 20u) / 4u; + const f16vec4 v0 = q8(EO[e]), v1 = q8(EO[e + 1u]), v2 = q8(EO[e + 2u]), v3 = q8(EO[e + 3u]), v4 = q8(EO[e + 4u]); + // values: 0, 1, 10 gates; 2..5 feature; 6..9 its blend; 11..13 mix; 14, 15, 16 covariance + imageStore(gateOut, at, unorm8Value(vec4(sig16(v0.x), sig16(v0.y), sig16(v2.z), 0.0))); + const vec3 m = vec3(v2.w, v3.x, v3.y), ex = exp2((m - max(max(m.x, m.y), m.z)) * 1.442695); + imageStore(mixOut, at, unorm8Value(vec4(ex * (1.0 / (ex.z + (ex.x + ex.y))), 0.0))); + const float a = p.sigma * float(sig16(v3.z)); + const float b = p.sigma * float(sig16(v3.w)); + const float t = float(v4.x) * p.rho, s = t * t, m_ = abs(t); + const float big = m_ >= 9.0109138 ? 1.0 : fma(1.0 / (exp2(m_ * 2.88539) + 1.0), -2.0, 1.0); + const float rho = 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); + imageStore(covOut, at, vec4(a * a, b * b, a * b * rho, 0.0)); + precise vec2 f = fma(vec2(at) + 0.5 - p.rndOffs, p.rnd2hi, texelFetch(motionTex, at, 0).xy); + f16vec4 hist = f16vec4(0.0); + if (p.reset == 0u && all(greaterThanEqual(f, vec2(0.0))) && all(lessThanEqual(f, p.hiSize))) { + precise vec2 normalized = fma(f, 1.0 / p.hiSize, p.histOffs * (1.0 / vec2(p.size))); + precise vec2 texels = normalized * vec2(p.size); + precise vec2 uv = texels * p.renderInv; + hist = f16vec4(roundHalfValue(textureLod(featTex, uv, 0.0))); + } + // the recurrent feature: gated mix of its reprojected history and the new value. The reference's PTX + // computes it with four separate f16 operations (sub, mul, mul, add: + // build/gfx1201-perf/ptx/dlss-0162-00.ptx:1416-1447), so `precise` keeps the compiler from fusing a + // product into an fma and dropping a rounding. + const f16vec4 g = gate16(f16vec4(v1.z, v1.w, v2.x, v2.y)); + precise f16vec4 blend = (f16vec4(1.0) - g) * hist + g * tan16(f16vec4(v0.z, v0.w, v1.x, v1.y)); + imageStore(featOut, at, snorm8Value(vec4(blend))); + } +} diff --git a/engine/down_m.comp b/engine/down_m.comp new file mode 100644 index 0000000..1f37bd7 --- /dev/null +++ b/engine/down_m.comp @@ -0,0 +1,76 @@ +#version 460 +// d4r engine, preset M: the last stage. The reconstruction works at three times the render resolution; this +// resamples its packed result (11-11-10 bits, range-compressed) to the output resolution, 3 -> 2, with a separable +// five-tap filter whose weights depend on the output pixel's parity, limits each pass to the range of its taps +// (no ringing), and undoes the range compression and the exposure. +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_GOOGLE_include_directive : require +#include "output_store.glsl" +// LDR (NGX's *_ldr_* kernel): the packed image is the display-referred colour itself. +#ifndef LDR +#define LDR 0 +#endif +// D4RM0003: the runtime dispatches 8x64 output pixels per workgroup. +#define BLOCK_X 8 +#define BLOCK_Y 64 +#define TILE_W (3 * BLOCK_X / 2 + 4) +#define TILE_H (3 * BLOCK_Y / 2 + 4) +layout(local_size_x = BLOCK_X, local_size_y = BLOCK_Y) in; +layout(set = 0, binding = 0, std430) readonly buffer Params +{ + ivec2 srcSize; // of the packed image + ivec2 outOrigin; + ivec2 colOrigin, colSize; + vec2 out2rnd, rndOffs; // render texel of an output pixel centre: floor(centre * out2rnd + rndOffs) + float expo; +} p; +layout(set = 0, binding = 1) uniform sampler2D colorTex; // render resolution; its alpha is passed through +layout(set = 0, binding = 2, FINAL_OUTPUT_FORMAT) uniform writeonly image2D outImg; +layout(set = 0, binding = 3, std430) readonly buffer Src { uint w[]; } src; +shared f16vec4 tile[TILE_W * TILE_H]; // source texels around the block +shared f16vec4 rows[TILE_H * BLOCK_X]; // after the horizontal pass +// weights of taps -2 .. +2 for an even output coordinate; an odd one uses them mirrored +const float16_t W0 = float16_t(-0.041347626596689224), W1 = float16_t(0.08728941529989243), W2 = float16_t(0.6202144026756287), + W3 = float16_t(0.3721286356449127), W4 = float16_t(-0.03828485310077667); +f16vec3 filter5(f16vec3 a, f16vec3 b, f16vec3 c, f16vec3 d, f16vec3 e, bool odd) +{ + const f16vec3 w0 = f16vec3(odd ? W4 : W0), w1 = f16vec3(odd ? W3 : W1), w3 = f16vec3(odd ? W1 : W3), w4 = f16vec3(odd ? W0 : W4); + f16vec3 s = w0 * a; + s = fma(w1, b, s); s = fma(f16vec3(W2), c, s); s = fma(w3, d, s); s = fma(w4, e, s); + return min(max(s, min(min(min(a, b), min(c, d)), e)), max(max(max(a, b), max(c, d)), e)); +} +void main() +{ + const ivec2 lid = ivec2(gl_LocalInvocationID.xy), pix = ivec2(gl_GlobalInvocationID.xy); + const ivec2 base = ivec2(3 * BLOCK_X / 2, 3 * BLOCK_Y / 2) * ivec2(gl_WorkGroupID.xy) - 2; + for (int i = lid.y * BLOCK_X + lid.x; i < TILE_W * TILE_H; i += BLOCK_X * BLOCK_Y) + { + const ivec2 at = clamp(base + ivec2(i % TILE_W, i / TILE_W), ivec2(0), p.srcSize - 1); + const uint w = src.w[at.y * p.srcSize.x + at.x]; +#if LDR + tile[i] = f16vec4(vec4(float(w & 2047u) * (1.0 / 2047.0), float((w >> 11) & 2047u) * (1.0 / 2047.0), float(w >> 22) * (1.0 / 1023.0), 0.0)); +#else + tile[i] = f16vec4(vec4(float(w & 2047u) * 0.0011968734906986356, float((w >> 11) & 2047u) * 0.0011968734906986356, float(w >> 22) * 0.002394916955381632, 0.0)); +#endif + } + memoryBarrierShared(); barrier(); + const ivec2 c = ivec2(floor((vec2(pix) + 0.5) * 1.5)) - base; + for (int y = lid.y; y < TILE_H; y += BLOCK_Y) + { + const int o = y * TILE_W + c.x; + rows[y * BLOCK_X + lid.x].rgb = filter5(tile[o - 2].rgb, tile[o - 1].rgb, tile[o].rgb, tile[o + 1].rgb, tile[o + 2].rgb, (pix.x & 1) != 0); + } + memoryBarrierShared(); barrier(); + // After both shared barriers: borrowed images may extend beyond the feature's true output rectangle. + // M accepts even output dimensions and srcSize = outputSize/2*3, so this bounds check is exact. + if (3 * pix.x >= 2 * p.srcSize.x || 3 * pix.y >= 2 * p.srcSize.y) return; + const int o = c.y * BLOCK_X + lid.x; +#if LDR + const vec3 lin = vec3(filter5(rows[o - 2 * BLOCK_X].rgb, rows[o - BLOCK_X].rgb, rows[o].rgb, rows[o + BLOCK_X].rgb, rows[o + 2 * BLOCK_X].rgb, (pix.y & 1) != 0)); +#else + const vec3 v = clamp(vec3(filter5(rows[o - 2 * BLOCK_X].rgb, rows[o - BLOCK_X].rgb, rows[o].rgb, rows[o + BLOCK_X].rgb, rows[o + 2 * BLOCK_X].rgb, (pix.y & 1) != 0)), 1e-4, 3.0); + const vec3 lin = (exp2(v * (1.0 / 0.36) * log2(10.0)) - 1.0) * (1.0 / 100.0) * (1.0 / (p.expo == 0.0 ? 1.0 : p.expo)); +#endif + const ivec2 r = clamp(ivec2(floor((vec2(pix) + 0.5) * p.out2rnd + p.rndOffs)), ivec2(0), p.colSize - 1) + p.colOrigin; + imageStore(outImg, p.outOrigin + pix, finalOutputValue(vec4(lin, texelFetch(colorTex, r, 0).a))); +} diff --git a/engine/enc0_m.comp b/engine/enc0_m.comp new file mode 100644 index 0000000..4a206f9 --- /dev/null +++ b/engine/enc0_m.comp @@ -0,0 +1,263 @@ +#version 460 +// d4r engine, preset M: the input stage, written from the algorithm in M_NOTES.md. Per render pixel it gathers +// 20 features (the colour, four colours of the previous result along three motion candidates, the recurrent +// feature, a motion-change measure), rounds them to FP8 with NGX's dither, and embeds each 2x2-pixel token +// (128 feature slots) into 64 channels with one matrix product per 4x4-token block. It also writes the three +// motion candidates the reconstruction uses. +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_float_e4m3 : require +#extension GL_GOOGLE_include_directive : require +#include "round_half.glsl" +layout(local_size_x = 32) in; // one wave: 16 tokens = 8x8 render pixels, two pixels per lane +layout(set = 0, binding = 0, std430) readonly buffer Params +{ + ivec2 size; // render size + ivec2 shift; // pixels the block grid is shifted by + ivec2 tok; // token grid + vec2 jitter, prevJitter; + vec2 mvScale, mvOffs; // motion texel -> pixels of the 3x image + vec2 rnd2hi, hi2rnd; + vec2 hiSize; + float changeScale, changeGain; + uint depthFlag, reset, frame; + float expo, expoPrev; + float halfInvX, halfInvY, renderInvX, renderInvY; // host-rounded CUDA texture coordinates +} p; +layout(set = 0, binding = 1) uniform sampler2D colorTex; +layout(set = 0, binding = 2) uniform sampler2D motionTex; +layout(set = 0, binding = 3) uniform sampler2D depthTex; +layout(set = 0, binding = 4) uniform sampler2D prevMvTex; // candidate A of the previous frame, bilinear +layout(set = 0, binding = 5) uniform sampler2D halfTex; // half-resolution previous result, bilinear +layout(set = 0, binding = 6) uniform sampler2D featTex; // recurrent feature, bilinear +layout(set = 0, binding = 7, rg16f) uniform writeonly image2D mvAOut; +layout(set = 0, binding = 8, rg16f) uniform writeonly image2D mvBOut; +layout(set = 0, binding = 9, rg16f) uniform writeonly image2D mvCOut; +layout(set = 0, binding = 10, std430) readonly buffer W8 { floate4m3_t w[]; } w8; // [128][64] +layout(set = 0, binding = 11, std430) readonly buffer W16 { f16vec4 w[]; } w16; // bias, 16 equal rows of 64 +layout(set = 0, binding = 12, std430) readonly buffer Noise { uint16_t n[]; } noise; // 8 x 16 x 16 +layout(set = 0, binding = 13, std430) writeonly buffer Out { floate4m3_t a[]; } outb; // [tok.y][tok.x][64] +#define ACC coopmat +#define HACC coopmat +#define A8 coopmat +#define B8 coopmat +#define RM gl_CooperativeMatrixLayoutRowMajor +// NVIDIA's embedding GEMM is chained m16n8k32 mma with an f16 accumulator: the bias is the initial C and every +// 32 products are summed with one f16 rounding (swin_model.py mma_chain; the same k32 step the native kernels' +// SWIN_EXACT build implements). K32_ROUND reproduces that: after each pair of k16 cooperative products the +// accumulator is rounded to f16 through a per-invocation f16 scratch slot (ROUND_HALF_SHARED). +#ifndef K32_ROUND +#define K32_ROUND 1 +#endif +#if K32_ROUND +#define R32(m) ROUND_HALF_SHARED(m, RS) +#else +#define R32(m) +#endif +shared floate4m3_t FE[16 * 136]; +shared float16_t RS[gl_WorkGroupSize.x * 8u]; +const float L10 = 3.3219281; +// LDR (NGX's *_ldr_* kernels, the game did not set the HDR flag): display-referred colour, no exposure and no +// range compression anywhere. +#ifndef LDR +#define LDR 0 +#endif +// PLANES: the tokens in NGX's layout, two planes of 32 channels ([2][tok.y][tok.x][32]), for the HIP network +#ifndef PLANES +#define PLANES 0 +#endif +#if LDR +float compress1(float v) { return v; } +#else +float compress1(float v) { + precise float mapped = fma(v, 100.0, 1.0); + precise float lg = log2(max(mapped, 1e-4)); + precise float base10 = lg * 0.30103; + precise float result = base10 * 0.36; + return result; +} +#endif +vec3 compress(vec3 v) { return vec3(compress1(v.x), compress1(v.y), compress1(v.z)); } +#if LDR +float colour1(float v) { return isinf(v) || isnan(v) ? 0.0 : min(v, 32752.0); } +#else +float colour1(float v) { return compress1(isinf(v) || isnan(v) ? 0.0 : min(v, 32752.0)); } +#endif +ivec2 texel(ivec2 t) { return ivec2(t.x > 0 ? min(t.x, p.size.x - 1) : 0, t.y > 0 ? min(t.y, p.size.y - 1) : 0); } +vec2 motion(ivec2 t) { return p.mvScale * (p.mvOffs + texelFetch(motionTex, texel(t), 0).xy); } +// NGX's dithered truncation of an f16 value to the FP8 grid +floate4m3_t dither(float v, uint n) +{ + const float16_t h = float16_t(isinf(v) || isnan(v) ? 0.0 : min(v, 1024.0)); + const uint e = uint(float16BitsToUint16(h)) & 0xFC00u; + precise float16_t added = h + uint16BitsToFloat16(uint16_t(e | n)); + precise float16_t t = added - uint16BitsToFloat16(uint16_t(e)); + const float16_t r = uint16BitsToFloat16(uint16_t(uint(float16BitsToUint16(t)) & 0xFF80u)); + return floate4m3_t(clamp(r, float16_t(-448.0), float16_t(448.0))); +} +void gatherPixel(uint pi, ivec2 block, float e1, float e2, bool fusedFourth) +{ + const uint px = pi & 7u, py = pi >> 3; + const ivec2 raw = block + ivec2(px, py), a = abs(raw); + const ivec2 m = ivec2(a.x < p.size.x ? a.x : 2 * p.size.x - a.x - 2, a.y < p.size.y ? a.y : 2 * p.size.y - a.y - 2); + const bool outside = raw.x < 0 || raw.y < 0 || a.x >= p.size.x || a.y >= p.size.y; + // the neighbours with the smallest and the largest depth + ivec2 lo = m, hi = m; + float dlo = texelFetch(depthTex, texel(m), 0).r, dhi = dlo; + const ivec2 offs[8] = ivec2[8](ivec2(-1, -1), ivec2(0, -1), ivec2(1, -1), ivec2(-1, 0), ivec2(1, 0), ivec2(-1, 1), ivec2(0, 1), ivec2(1, 1)); + for (int k = 0; k < 8; ++k) + { + const float d = texelFetch(depthTex, texel(m + offs[k]), 0).r; + if (d < dlo) { dlo = d; lo = m + offs[k]; } + if (d > dhi) { dhi = d; hi = m + offs[k]; } + } + const vec2 mvA = motion(p.depthFlag != 0u ? lo : hi), mvB = motion(p.depthFlag != 0u ? hi : lo), mvC = motion(m); + if (!outside) + { + imageStore(mvAOut, m, vec4(mvA, 0.0, 0.0)); imageStore(mvBOut, m, vec4(mvB, 0.0, 0.0)); imageStore(mvCOut, m, vec4(mvC, 0.0, 0.0)); + } + const vec3 own = texelFetch(colorTex, texel(m), 0).rgb; + const vec2 centre = vec2(m) + 0.5, f = centre - p.jitter, inv = 1.0 / vec2(p.size); + // how much the motion changed since the previous frame + const vec2 prev = textureLod(prevMvTex, clamp(((f + p.prevJitter) + mvA * p.hi2rnd) * vec2(p.renderInvX, p.renderInvY), vec2(0.0), vec2(p.size) * vec2(p.renderInvX, p.renderInvY)), 0.0).xy; + const vec2 dm = (p.changeScale * 540.0 / float(p.size.y)) * (mvA * p.hi2rnd - p.hi2rnd * prev); + const float d2 = fma(dm.x, dm.x, dm.y * dm.y); + const float change = p.changeGain * (log2(d2 > 0.0 ? sqrt(d2) + 1.0 : 1.0) * 0.6931472); + // the previous result along the candidates + precise vec2 q[4]; + q[0] = mvA + f * p.rnd2hi; q[1] = mvC + f * p.rnd2hi; q[2] = mvB + f * p.rnd2hi; + // The captured PTX fuses the first pixel's fourth coordinate, but reuses a rounded + // centre*scale product for the second pixel. A single FMA for both changes FP8 dither. + q[3] = fusedFourth ? fma(centre, p.rnd2hi, mvC) : centre * p.rnd2hi + mvC; + vec3 h[4]; + for (int k = 0; k < 4; ++k) + { + precise vec3 s = own; + if (!(q[k].x < 0.0 || q[k].y < 0.0 || q[k].x > p.hiSize.x || q[k].y > p.hiSize.y || p.reset != 0u)) +#if LDR + s = textureLod(halfTex, q[k] * vec2(p.halfInvX, p.halfInvY), 0.0).rgb; +#else + { + precise vec3 packedHistory = textureLod(halfTex, q[k] * vec2(p.halfInvX, p.halfInvY), 0.0).rgb * 2.45; + precise vec3 scaled = clamp(packedHistory, 1e-4, 3.0) * (1.0 / 0.36); + precise vec3 exponent = scaled * L10; + precise vec3 linear = (exp2(exponent) - 1.0) * (1.0 / 100.0); + s = linear * (1.0 / e2); + } +#endif + h[k] = compress(e1 * s); + } + vec4 feat = vec4(0.0); + float mf = 0.0; + if (p.reset == 0u) + { + mf = change; + if (!(q[0].x < 0.0 || q[0].y < 0.0 || q[0].x > p.hiSize.x || q[0].y > p.hiSize.y)) + { + precise vec2 normalized = fma(q[0], 1.0 / p.hiSize, p.prevJitter * inv); + precise vec2 texels = normalized * vec2(p.size); + precise vec2 uv = texels * vec2(p.renderInvX, p.renderInvY); + feat = textureLod(featTex, uv, 0.0); + } + } + const vec3 c = e1 * own; + const vec3 col = vec3(colour1(c.r), colour1(c.g), colour1(c.b)); + const uint n = uint(noise.n[(p.frame << 8) | uint((a.y & 15) << 4) | uint(a.x & 15)]) >> 9; + // 32 slots: 16 per plane + float v[32]; + for (int k = 0; k < 32; ++k) v[k] = 0.0; + v[0] = col.r; v[1] = col.g; v[2] = feat.x; v[3] = feat.y; v[4] = col.b; v[5] = h[0].r; v[6] = feat.z; v[7] = feat.w; + v[8] = h[0].g; v[9] = h[0].b; v[12] = h[1].r; v[13] = h[1].g; + v[16] = h[1].b; v[17] = h[2].r; v[20] = h[2].g; v[21] = h[2].b; v[24] = h[3].r; v[25] = h[3].g; v[28] = h[3].b; v[29] = mf; + const uint at = ((py >> 1) * 4u + (px >> 1)) * 136u + ((py & 1u) * 2u + (px & 1u)) * 32u; + for (uint k = 0u; k < 32u; ++k) FE[at + k] = dither(v[k], n); +} +void main() +{ + const uint lane = gl_LocalInvocationID.x; + const ivec2 block = 8 * ivec2(gl_WorkGroupID.xy) - p.shift; + const float e1 = p.expo == 0.0 ? 1.0 : p.expo, e2 = p.expoPrev == 0.0 ? 1.0 : p.expoPrev; + // Static variants preserve the first pixel's FMA; RADV merges a runtime FMA/mul-add selection. + gatherPixel(lane, block, e1, e2, true); + gatherPixel(lane + 32u, block, e1, e2, false); + memoryBarrierShared(); barrier(); + // Four output accumulators live at once: every FE k16 A operand is loaded once and applied to the four + // weight B operands. Each output column group keeps its own K32 chain: two k16 products, then the previous + // accumulator added and rounded to f16 (ROUND_HALF_SHARED), exactly as the sequential version did. + HACC b0; coopMatLoad(b0, w16.w, 0u, 16u, RM); + HACC b1; coopMatLoad(b1, w16.w, 4u, 16u, RM); + HACC b2; coopMatLoad(b2, w16.w, 8u, 16u, RM); + HACC b3; coopMatLoad(b3, w16.w, 12u, 16u, RM); + ACC a0 = ACC(b0), a1 = ACC(b1), a2 = ACC(b2), a3 = ACC(b3); + for (uint kt = 0u; kt < 8u; kt += 2u) + { + A8 x0, x1; + coopMatLoad(x0, FE, kt * 16u, 136u, RM); + coopMatLoad(x1, FE, (kt + 1u) * 16u, 136u, RM); + ACC p0 = ACC(0.0), p1 = ACC(0.0), p2 = ACC(0.0), p3 = ACC(0.0); + B8 w0; coopMatLoad(w0, w8.w, kt * 16u * 64u + 0u, 64u, RM); p0 = coopMatMulAdd(x0, w0, p0); + B8 w1; coopMatLoad(w1, w8.w, kt * 16u * 64u + 16u, 64u, RM); p1 = coopMatMulAdd(x0, w1, p1); + B8 w2; coopMatLoad(w2, w8.w, kt * 16u * 64u + 32u, 64u, RM); p2 = coopMatMulAdd(x0, w2, p2); + B8 w3; coopMatLoad(w3, w8.w, kt * 16u * 64u + 48u, 64u, RM); p3 = coopMatMulAdd(x0, w3, p3); + B8 y0; coopMatLoad(y0, w8.w, (kt + 1u) * 16u * 64u + 0u, 64u, RM); p0 = coopMatMulAdd(x1, y0, p0); + B8 y1; coopMatLoad(y1, w8.w, (kt + 1u) * 16u * 64u + 16u, 64u, RM); p1 = coopMatMulAdd(x1, y1, p1); + B8 y2; coopMatLoad(y2, w8.w, (kt + 1u) * 16u * 64u + 32u, 64u, RM); p2 = coopMatMulAdd(x1, y2, p2); + B8 y3; coopMatLoad(y3, w8.w, (kt + 1u) * 16u * 64u + 48u, 64u, RM); p3 = coopMatMulAdd(x1, y3, p3); + // The k32 product is independent of C: add the previous half accumulator only after both k16s. + a0 = p0 + a0; R32(a0); + a1 = p1 + a1; R32(a1); + a2 = p2 + a2; R32(a2); + a3 = p3 + a3; R32(a3); + } + // Each lane owns one accumulator column (lane & 15) and eight rows (lane >= 16 ? 8..15 : 0..7). The clamped + // accumulator is converted to FP8 straight into the packed token, so no staging region or second barrier. + // + // A lane's eight rows are the tokens (block >> 1) + (r & 3, r >> 2) in the low half and (r & 3, 2 + (r >> 2)) + // in the high half, so their byte offsets from the lane's first row are affine in r, (r >> 2) * tok.x + + // (r & 3) tokens. One base word and eight constants therefore address all eight rows, and a block whose + // whole 4x4 token square is on screen (every block but the image border) needs no per-row bounds test. +#if PLANES +#define TOK_STRIDE 32u +#else +#define TOK_STRIDE 64u +#endif + const uint c16 = lane & 15u, rbase = (lane >> 4u) << 3u; + const ivec2 bt = block >> 1; + const uint b8 = uint((int(bt.y) + int(rbase >> 2)) * p.tok.x + int(bt.x)) * TOK_STRIDE + c16; + const uint X8 = uint(p.tok.x) * TOK_STRIDE; +#if PLANES +#define DO_STORE(r, off) { \ + const uint tb_ = b8 + (off), pb_ = uint(p.tok.x * p.tok.y) * 32u; \ + outb.a[tb_] = floate4m3_t(clamp(a0[r], -6.2831855, 6.2831855)); \ + outb.a[tb_ + 16u] = floate4m3_t(clamp(a1[r], -6.2831855, 6.2831855)); \ + outb.a[tb_ + pb_] = floate4m3_t(clamp(a2[r], -6.2831855, 6.2831855)); \ + outb.a[tb_ + pb_ + 16u] = floate4m3_t(clamp(a3[r], -6.2831855, 6.2831855)); } +#else +#define DO_STORE(r, off) { \ + const uint tb_ = b8 + (off); \ + outb.a[tb_] = floate4m3_t(clamp(a0[r], -6.2831855, 6.2831855)); \ + outb.a[tb_ + 16u] = floate4m3_t(clamp(a1[r], -6.2831855, 6.2831855)); \ + outb.a[tb_ + 32u] = floate4m3_t(clamp(a2[r], -6.2831855, 6.2831855)); \ + outb.a[tb_ + 48u] = floate4m3_t(clamp(a3[r], -6.2831855, 6.2831855)); } +#endif +#define ROW_OFF(r) ((r) < 4 ? (r) * TOK_STRIDE : X8 + ((r) - 4) * TOK_STRIDE) + if (bt.x >= 0 && bt.y >= 0 && bt.x + 3 < p.tok.x && bt.y + 3 < p.tok.y) + { + DO_STORE(0, ROW_OFF(0)); DO_STORE(1, ROW_OFF(1)); DO_STORE(2, ROW_OFF(2)); DO_STORE(3, ROW_OFF(3)); + DO_STORE(4, ROW_OFF(4)); DO_STORE(5, ROW_OFF(5)); DO_STORE(6, ROW_OFF(6)); DO_STORE(7, ROW_OFF(7)); + } + else + { +#define STORE_ROW(r) { \ + const ivec2 t_ = bt + ivec2(int(rbase + uint(r)) & 3, int(rbase + uint(r)) >> 2); \ + if (t_.x >= 0 && t_.y >= 0 && t_.x < p.tok.x && t_.y < p.tok.y) DO_STORE(r, ROW_OFF(r)) } + STORE_ROW(0); STORE_ROW(1); STORE_ROW(2); STORE_ROW(3); + STORE_ROW(4); STORE_ROW(5); STORE_ROW(6); STORE_ROW(7); +#undef STORE_ROW + } +#undef DO_STORE +#undef ROW_OFF +#undef TOK_STRIDE +} diff --git a/engine/evk.cpp b/engine/evk.cpp new file mode 100644 index 0000000..2f8a2d3 --- /dev/null +++ b/engine/evk.cpp @@ -0,0 +1,130 @@ +// Minimal Vulkan compute runner for the engine's kernels. +// evk SHADER.spv GX GY REPEATS in:BINDING=FILE ... out:BINDING=FILE:BYTES[:INITFILE] ... +// Every binding is a storage buffer in VRAM. Outputs are zero-filled (or INITFILE), and written back after one dispatch. +#include +#include +#include +#include +#include +#include +#include +#include +#include +#define CK(x) do { VkResult r_ = (x); if (r_ != VK_SUCCESS) { fprintf(stderr, "%s failed: %d\n", #x, r_); exit(1); } } while (0) +static std::vector slurp(const std::string& p) +{ + std::ifstream f(p, std::ios::binary); + if (!f) { fprintf(stderr, "cannot read %s\n", p.c_str()); exit(1); } + return std::vector((std::istreambuf_iterator(f)), {}); +} +struct B { uint32_t binding; std::string file, out; VkDeviceSize size; VkBuffer b; void* map; }; +int main(int argc, char** argv) +{ + if (argc < 6) return 1; + const uint32_t gx = atoi(argv[2]), gy = atoi(argv[3]); const int repeats = atoi(argv[4]); + std::vector bs; + for (int i = 5; i < argc; ++i) + { + std::string a = argv[i]; B b{}; + const bool out = a.rfind("out:", 0) == 0; a = a.substr(a.find(':') + 1); + b.binding = atoi(a.substr(0, a.find('=')).c_str()); a = a.substr(a.find('=') + 1); + if (out) + { + b.out = a.substr(0, a.find(':')); a = a.substr(a.find(':') + 1); + b.size = strtoull(a.c_str(), nullptr, 10); + if (a.find(':') != std::string::npos) b.file = a.substr(a.find(':') + 1); + } + else b.file = a; + bs.push_back(b); + } + VkApplicationInfo app{VK_STRUCTURE_TYPE_APPLICATION_INFO}; app.apiVersion = VK_API_VERSION_1_3; + VkInstanceCreateInfo ici{VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO}; ici.pApplicationInfo = &app; + VkInstance inst; CK(vkCreateInstance(&ici, nullptr, &inst)); + uint32_t np = 8; VkPhysicalDevice pds[8]; vkEnumeratePhysicalDevices(inst, &np, pds); + VkPhysicalDevice pd = pds[0]; + for (uint32_t i = 0; i < np; ++i) { VkPhysicalDeviceProperties p; vkGetPhysicalDeviceProperties(pds[i], &p); if (p.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU) pd = pds[i]; } + VkPhysicalDeviceShaderFloat8FeaturesEXT f8{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_FLOAT8_FEATURES_EXT}; + VkPhysicalDeviceCooperativeMatrixFeaturesKHR cm{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR, &f8}; + VkPhysicalDeviceVulkan13Features v13{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_FEATURES, &cm}; + VkPhysicalDeviceVulkan12Features v12{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES, &v13}; + VkPhysicalDeviceVulkan11Features v11{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES, &v12}; + VkPhysicalDeviceFeatures2 f2{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2, &v11}; + vkGetPhysicalDeviceFeatures2(pd, &f2); + VkPhysicalDeviceMemoryProperties mp; vkGetPhysicalDeviceMemoryProperties(pd, &mp); + uint32_t nq = 8; VkQueueFamilyProperties qf[8]; vkGetPhysicalDeviceQueueFamilyProperties(pd, &nq, qf); + uint32_t qi = 0; for (uint32_t i = 0; i < nq; ++i) if (qf[i].queueFlags & VK_QUEUE_COMPUTE_BIT) { qi = i; break; } + float prio = 1; VkDeviceQueueCreateInfo qci{VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO}; qci.queueFamilyIndex = qi; qci.queueCount = 1; qci.pQueuePriorities = &prio; + const char* exts[] = {"VK_KHR_cooperative_matrix", "VK_EXT_shader_float8"}; + VkDeviceCreateInfo dci{VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO, &f2}; dci.queueCreateInfoCount = 1; dci.pQueueCreateInfos = &qci; dci.enabledExtensionCount = 2; dci.ppEnabledExtensionNames = exts; + VkDevice dev; CK(vkCreateDevice(pd, &dci, nullptr, &dev)); + VkQueue q; vkGetDeviceQueue(dev, qi, 0, &q); + std::vector lb; + for (B& b : bs) + { + std::vector data; + if (!b.file.empty()) data = slurp(b.file); + if (b.out.empty()) b.size = data.size(); + b.size = std::max((b.size + 3) & ~3ull, 16); + VkBufferCreateInfo bci{VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO}; bci.size = b.size; bci.usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT; + CK(vkCreateBuffer(dev, &bci, nullptr, &b.b)); + VkMemoryRequirements mr; vkGetBufferMemoryRequirements(dev, b.b, &mr); + const VkMemoryPropertyFlags want = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT; // VRAM, not system memory + uint32_t mt = ~0u; + for (uint32_t t = 0; t < mp.memoryTypeCount; ++t) if ((mr.memoryTypeBits >> t & 1) && (mp.memoryTypes[t].propertyFlags & want) == want) { mt = t; break; } + if (mt == ~0u) { fprintf(stderr, "no host-visible VRAM type\n"); return 1; } + VkMemoryAllocateInfo mai{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO}; mai.allocationSize = mr.size; mai.memoryTypeIndex = mt; + VkDeviceMemory mem; CK(vkAllocateMemory(dev, &mai, nullptr, &mem)); CK(vkBindBufferMemory(dev, b.b, mem, 0)); + CK(vkMapMemory(dev, mem, 0, b.size, 0, &b.map)); + memset(b.map, 0, b.size); memcpy(b.map, data.data(), std::min(data.size(), b.size)); + lb.push_back({b.binding, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, VK_SHADER_STAGE_COMPUTE_BIT, nullptr}); + } + auto code = slurp(argv[1]); + VkShaderModuleCreateInfo smi{VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO}; smi.codeSize = code.size(); smi.pCode = (const uint32_t*)code.data(); + VkShaderModule sm; CK(vkCreateShaderModule(dev, &smi, nullptr, &sm)); + VkDescriptorSetLayoutCreateInfo dli{VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO}; dli.bindingCount = lb.size(); dli.pBindings = lb.data(); + VkDescriptorSetLayout dl; CK(vkCreateDescriptorSetLayout(dev, &dli, nullptr, &dl)); + VkPipelineLayoutCreateInfo pli{VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO}; pli.setLayoutCount = 1; pli.pSetLayouts = &dl; + VkPipelineLayout pl; CK(vkCreatePipelineLayout(dev, &pli, nullptr, &pl)); + VkPipelineShaderStageRequiredSubgroupSizeCreateInfo rs{VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_REQUIRED_SUBGROUP_SIZE_CREATE_INFO}; rs.requiredSubgroupSize = 32; + VkComputePipelineCreateInfo cpi{VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO}; + cpi.stage = {VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO, &rs, VK_PIPELINE_SHADER_STAGE_CREATE_REQUIRE_FULL_SUBGROUPS_BIT, VK_SHADER_STAGE_COMPUTE_BIT, sm, "main", nullptr}; + cpi.layout = pl; + VkPipeline pipe; CK(vkCreateComputePipelines(dev, VK_NULL_HANDLE, 1, &cpi, nullptr, &pipe)); + VkDescriptorPoolSize dps{VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 16}; + VkDescriptorPoolCreateInfo dpi{VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO}; dpi.maxSets = 1; dpi.poolSizeCount = 1; dpi.pPoolSizes = &dps; + VkDescriptorPool dp; CK(vkCreateDescriptorPool(dev, &dpi, nullptr, &dp)); + VkDescriptorSetAllocateInfo dai{VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO}; dai.descriptorPool = dp; dai.descriptorSetCount = 1; dai.pSetLayouts = &dl; + VkDescriptorSet ds; CK(vkAllocateDescriptorSets(dev, &dai, &ds)); + std::vector bi(bs.size()); std::vector wr; + for (size_t i = 0; i < bs.size(); ++i) + { + bi[i] = {bs[i].b, 0, VK_WHOLE_SIZE}; + wr.push_back({VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET, nullptr, ds, bs[i].binding, 0, 1, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, nullptr, &bi[i], nullptr}); + } + vkUpdateDescriptorSets(dev, wr.size(), wr.data(), 0, nullptr); + VkCommandPoolCreateInfo cpci{VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO}; cpci.queueFamilyIndex = qi; cpci.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT; + VkCommandPool cp; CK(vkCreateCommandPool(dev, &cpci, nullptr, &cp)); + VkCommandBufferAllocateInfo cai{VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO}; cai.commandPool = cp; cai.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; cai.commandBufferCount = 1; + VkCommandBuffer cb; CK(vkAllocateCommandBuffers(dev, &cai, &cb)); + auto submit = [&](int n) { + VkCommandBufferBeginInfo cbi{VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO}; CK(vkBeginCommandBuffer(cb, &cbi)); + vkCmdBindPipeline(cb, VK_PIPELINE_BIND_POINT_COMPUTE, pipe); + vkCmdBindDescriptorSets(cb, VK_PIPELINE_BIND_POINT_COMPUTE, pl, 0, 1, &ds, 0, nullptr); + for (int r = 0; r < n; ++r) vkCmdDispatch(cb, gx, gy, 1); + CK(vkEndCommandBuffer(cb)); + VkSubmitInfo si{VK_STRUCTURE_TYPE_SUBMIT_INFO}; si.commandBufferCount = 1; si.pCommandBuffers = &cb; + auto t0 = std::chrono::steady_clock::now(); + CK(vkQueueSubmit(q, 1, &si, VK_NULL_HANDLE)); CK(vkQueueWaitIdle(q)); + return std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + }; + submit(1); + for (B& b : bs) if (!b.out.empty()) std::ofstream(b.out, std::ios::binary).write((const char*)b.map, b.size); + if (repeats > 1) + { + std::vector t; + for (int it = 0; it < 12; ++it) { double us = submit(repeats) / repeats; if (it >= 3) t.push_back(us); } + std::sort(t.begin(), t.end()); + printf("%s: %.1f us per dispatch (median, min %.1f)\n", argv[1], t[t.size() / 2], t[0]); + } + return 0; +} diff --git a/engine/exposure_k.comp b/engine/exposure_k.comp new file mode 100644 index 0000000..2c6f636 --- /dev/null +++ b/engine/exposure_k.comp @@ -0,0 +1,152 @@ +#version 460 +// d4r engine: automatic exposure, as NGX computes it (luma convert, two reduce-sum passes, auto-exposure copy). +// MODE 0 (preset K): exposure = key * exp(-mean of the log luma). +// MODE 1 (preset M, the cuda_dldn_engine_* kernels): exposure = key / max(mean of the luma clamped to +// [0, 10000], 0.005). +// Both modes measure the same grid: the colour rectangle sampled at half the output size with an +// endpoint-aligned nearest-texel lookup (the reference's cuda_luma_convert_kernel; the grid equals the +// rectangle itself only when the render size is half the output size). +// The result is kept as a half float, through the reference's truncating formatted surface store. +// PASS 0: one workgroup sums the mapped luma over a 32x32-pixel block of the measured grid. +// PASS 1: one workgroup adds the block sums, advances the exposure state and writes the exposure fields of the +// input and output stages' parameter blocks. It also runs on frames that do not measure. +// With flag 4 (MODE 0) the game's 1x1 exposure texture replaces the measurement; nothing is measured then. +#extension GL_EXT_shader_explicit_arithmetic_types : require +#ifndef MODE +#define MODE 0 +#endif +layout(local_size_x = 64) in; +layout(set = 0, binding = 0, std430) readonly buffer Params +{ + ivec2 origin, size; // rectangle of the colour texture; the grid maps onto it + float invCount; // 1 / (sampleSize.x * sampleSize.y) + float scale; // applied to the luma before the logarithm + float key; // exposure * geometric mean + uint flags; // 1: measure this frame; 2: reset (previous exposure = current); 4: game exposure (MODE 0) + uint blocks, blocksX; + // The measured grid: the reference's cuda_luma_convert_kernel samples the colour rectangle at half the + // output size with an endpoint-aligned nearest-texel lookup, not at every render texel. size is the + // source rectangle; invCount is 1 / (sampleSize.x * sampleSize.y) and the workgroups cover the grid. + ivec2 sampleSize; + // The game's exposure scale and 1 / pre-exposure (MODE 0). NGX keeps the measured exposure unscaled and + // applies both when it fills the stages' parameters, in this order: (1 / pre) * (scale * exposure). + float gameScale, invPre; +} p; +shared float part[64]; +// The half surface the reference's auto_exposure_copy_kernel writes truncates toward zero. An explicit +// float16_t cast rounds to nearest instead and lands on the next code whenever the exact value sits below +// the midpoint of its binary16 interval - on ordinary frames the measured exposure, so emulate the store. +float truncatedHalf(float v) +{ + uint bits = packHalf2x16(vec2(v, 0.0)) & 65535u; + if (unpackHalf2x16(bits).x > v) --bits; + return unpackHalf2x16(bits).x; +} +#if PASS == 0 +layout(set = 0, binding = 1) uniform sampler2D colorTex; +layout(set = 0, binding = 2, std430) writeonly buffer Sums { float v[]; } sums; +void main() +{ + const uint t = gl_LocalInvocationID.x; + const ivec2 base = ivec2(gl_WorkGroupID.xy) * 32 + ivec2(t & 7u, t >> 3) * 4; + const ivec2 grid = p.sampleSize; + float sum = 0.0; + for (int y = 0; y < 4; ++y) + for (int x = 0; x < 4; ++x) + { + const ivec2 at = base + ivec2(x, y); + if (at.x >= grid.x || at.y >= grid.y) continue; + // endpoint-aligned nearest-texel lookup from the grid to the colour rectangle + const ivec2 sampleAt = ivec2(floor(fma(vec2(at) / vec2(grid - 1), vec2(p.size - 1), vec2(0.5)))); + const vec3 c = texelFetch(colorTex, p.origin + sampleAt, 0).rgb; + // the reference accumulates the logarithm of the exposed luma (cuda_luma_convert_kernel). The + // constant is applied to Y = fma(b, 0.25, fma(r, 0.25, g * 0.5)); MODE 1 keeps its own expression. +#if MODE == 0 + precise const float luma = p.scale * fma(c.b, 0.25, fma(c.r, 0.25, 0.5 * c.g)); + // log(luma + 1e-6) as lg2(x) * ln2, and a non-finite or non-positive luma contributes nothing + if (luma > 0.0 && !isinf(luma)) + sum += (log2(luma + 1e-6) * uintBitsToFloat(0x3F317218u)) * p.invCount; +#else + const float luma = p.scale * (0.25 * c.r + 0.5 * c.g + 0.25 * c.b); + sum += p.invCount * min(max(luma, 0.0), 10000.0); +#endif + } + part[t] = sum; + barrier(); + if (t == 0u) + { + float total = 0.0; + for (uint i = 0u; i < 64u; ++i) total += part[i]; + sums.v[gl_WorkGroupID.y * p.blocksX + gl_WorkGroupID.x] = total; + } +} +#else +layout(set = 0, binding = 1, std430) readonly buffer Sums { float v[]; } sums; +layout(set = 0, binding = 2, std430) buffer State { float current, previous; } state; +#if MODE == 0 +layout(set = 0, binding = 3, std430) buffer InputParams { float w[]; } inp; // expo, expoIn, expoRatio at words 28..30 +layout(set = 0, binding = 4, std430) buffer OutputParams { float w[]; } outp; // expo, expoRatio at words 37, 38 +layout(set = 0, binding = 5) uniform sampler2D gameExposure; // the game's 1x1 exposure (flag 4) +#else +layout(set = 0, binding = 3, std430) buffer InputParams { float w[]; } inp; // enc0_m: expo, expoPrev at words 25, 26 +layout(set = 0, binding = 4, std430) buffer PostParams { float w[]; } post; // post_m: words 14, 15 +layout(set = 0, binding = 5, std430) buffer DownParams { float w[]; } down; // down_m: word 12 +#endif +void main() +{ + const uint t = gl_LocalInvocationID.x; + float sum = 0.0; + if ((p.flags & 1u) != 0u) + for (uint i = t; i < p.blocks; i += 64u) sum += sums.v[i]; + part[t] = sum; + barrier(); + if (t != 0u) return; + float current = state.current, previous = current; +#if MODE == 0 + if ((p.flags & 4u) != 0u) + { + float value = texelFetch(gameExposure, ivec2(0), 0).r; + if (value == 0.0) value = 1.0; // NGX's texture read maps a zero exposure to 1 + // NGX copies the exposure into a half surface with the same truncating store as the measured + // one, and the next frame's ratio uses that half value; a reset carries the live f32 value. + previous = (p.flags & 2u) != 0u ? value : truncatedHalf(current); + current = value; + } +#endif + if ((p.flags & 1u) != 0u) + { + float mean = 0.0; + for (uint i = 0u; i < 64u; ++i) mean += part[i]; +#if MODE == 0 + // auto_exposure_copy_kernel: key * exp2(-log2e * sum), written to a half surface. The formatted + // store truncates toward zero, as in MODE 1 (measured on the captured K exposure: the store lands + // one binary16 code below the correctly rounded value). + float measured = truncatedHalf(p.key * exp2(uintBitsToFloat(0xBFB8AA3Bu) * mean)); +#else + // The reference's formatted half surface store truncates positive exposure; an explicit f16 cast + // rounds to nearest instead and moves FP8 embedding codes across thresholds. + const float full = p.key / max(mean, 0.005); + float measured = truncatedHalf(full); +#endif + if (!(measured > 0.0) || isinf(measured)) measured = 1.0; +#if MODE == 0 + previous = (p.flags & 2u) != 0u ? measured : current; +#else + previous = (p.flags & 2u) != 0u ? 0.0 : current; // 0: no earlier exposure (the stages read it as 1) +#endif + current = measured; + } + state.current = current; state.previous = previous; +#if MODE == 0 + precise const float scaled = p.gameScale * current, scaledPrevious = p.gameScale * previous; + precise const float exposure = p.invPre * scaled; + precise const float ratio = scaled / scaledPrevious; + inp.w[28] = exposure; inp.w[29] = 1.0; inp.w[30] = ratio; + outp.w[37] = exposure; outp.w[38] = ratio; +#else + inp.w[25] = current; inp.w[26] = previous; + post.w[14] = current; post.w[15] = previous; + down.w[12] = current; +#endif +} +#endif diff --git a/engine/extract_k.py b/engine/extract_k.py new file mode 100644 index 0000000..85db11d --- /dev/null +++ b/engine/extract_k.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python3 +"""Extract the preset K network from one captured frame into d4r's own model file. + +usage: extract_k.py REPLAY_DIR OUT_PREFIX + REPLAY_DIR D4R_CUDA_REPLAY_DUMP_DIR of a run with D4R_CUDA_REPLAY_DUMP_FILTER=dltss_pwin (the 11 layer launches) +Writes OUT_PREFIX.bin (f16) and OUT_PREFIX.json (per layer: shape parameters and the offset of every tensor, in f16 +elements). Tensors are plain row-major matrices W[K][N] (y = x W), unlike the pre-swizzled fragments NGX keeps on +the GPU; per-channel vectors are stored 16 times (one 16-row tile) so a kernel can load them as a matrix operand. +No NVIDIA code is involved: only the weight values, read from the buffers NGX uploaded. +""" +import argparse, glob, json, os, sys +import numpy as np +sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'kernels', 'tools')) +import pwin_model as pm + +F16 = np.float16 +# layer -> (kind, heads, C, COUT or CL, position-only attention) +LAYERS = {'enc0': ('enc0', 2, 32, 64, True), 'enc1': ('enc', 2, 64, 64, True), 'enc2': ('enc', 2, 64, 96, False), + 'enc3': ('enc', 4, 96, 128, False), 'enc4': ('enc', 4, 128, 160, False), 'dec5': ('plain', 8, 160, 0, False), + 'dec4': ('dec', 4, 128, 160, False), 'dec3': ('dec', 4, 96, 128, False), 'dec2': ('dec', 2, 64, 96, False), + 'dec1': ('dec', 2, 64, 64, False), 'dec0': ('dec0', 2, 32, 64, True)} + + +class Blob: + def __init__(self): + self.parts, self.n = [], 0 + + def add(self, a): + a = np.ascontiguousarray(a, F16).ravel() + pad = (-self.n) % 16 # keep every tensor 32-byte aligned + if pad: + self.parts.append(np.zeros(pad, F16)); self.n += pad + off = self.n + self.parts.append(a); self.n += a.size + return off + + +def tile16(v): + return np.tile(np.asarray(v, F16)[None, :], (16, 1)) + + +def pad_cols(W, n): + out = np.zeros((W.shape[0], n), F16); out[:, :W.shape[1]] = W + return out + + +def core(P, L, blob, posattn, fold_gains=True): + C, H = L.C, L.H + # The two per-channel scales are folded into the rows of the weights they feed (x * g) W = x (diag(g) W): + # norm 1's gain into Q, K, V and the MLP input gain into fc1. Products in f64, rounded once to f16. + g1 = pm.vec(P, L.G1, C).astype(np.float64)[:, None] + g2 = pm.vec(P, L.G2, C).astype(np.float64)[:, None] + t = {'g1': blob.add(tile16(pm.vec(P, L.G1, C)))} + t['qkv'], t['table'], t['wo'] = [], [], [] + for h in range(H): + hb = L.QKV + L.HEAD * h + for j in range(3): + W = np.concatenate([pm.weights(P, hb + 2048 * (3 * ch + j), 32, 32, 1024) for ch in range(C // 32)], axis=0) # [C][32] + t['qkv'].append(blob.add((W.astype(np.float64) * g1).astype(F16) if fold_gains else W)) + t['table'].append(blob.add(pm.bias_table(P, L.BIAS + 8192 * h))) # [64][64] + t['wo'].append(blob.add(pm.weights(P, L.WO + 64 * C * h, 32, C, 1024))) # [32][C] + t['bo'] = blob.add(tile16(pm.vec(P, L.BO, C))) + t['g2'] = blob.add(tile16(pm.vec(P, L.G2, C))) + t['b2'] = blob.add(tile16(pm.vec(P, L.B2, C))) + t['w1'] = [] + for c in range(C // 8): + W = pm.weights(P, L.W1 + 64 * C * c, C, 32, 32 * C) + t['w1'].append(blob.add((W.astype(np.float64) * g2).astype(F16) if fold_gains else W)) + t['b1'] = [blob.add(tile16(pm.vec(P, L.B1 + 64 * c, 32))) for c in range(C // 8)] + t['w2'] = [blob.add(pm.weights(P, L.W2 + 64 * C * c, 32, C, 1024)) for c in range(C // 8)] # [32][C] + return t + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('replay_dir') + parser.add_argument('out_prefix') + parser.add_argument('--unfold-gains', action='store_true', help='preserve explicit f16 gain and residual rounding instead of folding gains into weights') + args = parser.parse_args() + rdir, prefix = args.replay_dir, args.out_prefix + blob, index = Blob(), {} + for name, (kind, H, C, X, posattn) in LAYERS.items(): + matches = sorted(glob.glob(f'{rdir}/replay-*-dltss_pwin_{name}_layer')) + if not matches: + raise ValueError(f'missing preset K layer {name} in {rdir}') + d = matches[0] + P = pm.Dump(d) + e = {'kind': kind, 'H': H, 'C': C, 'X': X, 'posattn': posattn, 'fold_gains': not args.unfold_gains} + if kind in ('enc0', 'enc'): + L = pm.Layout(H, C, X) + Pw = P + if kind == 'enc0': + CIN = 16 + e['embed_w'] = blob.add(pm.weights(P, 2 * C, CIN, C, 512 * (CIN // 16))) # [16][C] + e['embed_b'] = blob.add(tile16(pm.vec(P, 0, C))) + Pw = pm.Shifted(P, 2 * C + 2 * CIN * C) + e.update(core(Pw, L, blob, posattn, not args.unfold_gains)) + Wp = pm.weights(Pw, L.PM, 4 * C, L.NPA * H, 128 * C) # [4C][NPA*H] + bias = pm.vec(Pw, L.PMB, L.NPA * H) if kind == 'enc0' or name != 'enc1x' else None + cols = np.concatenate([Wp[:, L.NPA * w:L.NPA * w + L.NPW] for w in range(H)], axis=1) # [4C][COUT] + b = np.concatenate([bias[L.NPA * w:L.NPA * w + L.NPW] for w in range(H)]) + n16 = -(-X // 16) * 16 + e['merge_w'] = blob.add(pad_cols(cols, n16)); e['merge_b'] = blob.add(tile16(np.concatenate([b, np.zeros(n16 - X, F16)]))) + e['merge_n'] = n16 + elif kind == 'plain': + e.update(core(P, pm.Layout(H, C, C), blob, posattn, not args.unfold_gains)) + else: + CL = X + EXP = 2 * CL * 4 * C + e['expand_w'] = blob.add(pm.weights(P, 0, CL, 4 * C, 512 * (CL // 16))) # [CL][4C] + e['expand_b'] = blob.add(tile16(pm.vec(P, EXP, 4 * C))) + L = pm.Layout(H, C, C) + Pw = pm.Shifted(P, EXP + 8 * C) + e.update(core(Pw, L, blob, posattn, not args.unfold_gains)) + if kind == 'dec0': + NOUTA = 48 + e['head_w'] = blob.add(pm.weights(Pw, L.PM + 2 * NOUTA, C, NOUTA, 1024)) # [C][48] + e['head_b'] = blob.add(tile16(pm.vec(Pw, L.PM, NOUTA))) + index[name] = e + data = np.concatenate(blob.parts) + data.tofile(prefix + '.bin') + with open(prefix + '.json', 'w') as f: + json.dump(index, f, indent=1) + print(f'{prefix}.bin: {data.size * 2} bytes, {len(index)} layers') + + +if __name__ == '__main__': + main() diff --git a/engine/game.cpp b/engine/game.cpp new file mode 100644 index 0000000..e1a021b --- /dev/null +++ b/engine/game.cpp @@ -0,0 +1,126 @@ +#include "game.h" +#include "game_internal.h" + +namespace d4r { +namespace { +bool depthAspect(VkFormat format) { return format == VK_FORMAT_D32_SFLOAT || format == VK_FORMAT_D32_SFLOAT_S8_UINT; } +} // namespace + +bool gameColorFormat(VkFormat format) { + switch (format) { + case VK_FORMAT_R16G16B16A16_SFLOAT: case VK_FORMAT_R32G32B32A32_SFLOAT: case VK_FORMAT_B10G11R11_UFLOAT_PACK32: + case VK_FORMAT_A2B10G10R10_UNORM_PACK32: case VK_FORMAT_R8G8B8A8_UNORM: case VK_FORMAT_B8G8R8A8_UNORM: + return true; + default: return false; + } +} +bool gameMotionFormat(VkFormat format) { + return format == VK_FORMAT_R16G16_SFLOAT || format == VK_FORMAT_R32G32_SFLOAT || + format == VK_FORMAT_R16G16B16A16_SFLOAT || format == VK_FORMAT_R32G32B32A32_SFLOAT; +} +bool gameDepthFormat(VkFormat format, VkImageAspectFlags aspect) { + if (aspect & VK_IMAGE_ASPECT_DEPTH_BIT) return depthAspect(format); + return format == VK_FORMAT_R32_SFLOAT || format == VK_FORMAT_R16_SFLOAT || format == VK_FORMAT_R16_UNORM; +} +bool gameOutputFormat(VkFormat format) { return gameColorFormat(format); } + +struct GameUpscaler::Impl { + Model& model; + Shape shape; + Engine engine; + uint32_t maxWidth, maxHeight; + game_detail::CanonicalInputs fallback; + bool started = false; + uint32_t frames = 0; + Impl(Model& m, VkPhysicalDevice p, VkDevice d, uint32_t ow, uint32_t oh, uint32_t rw, uint32_t rh, const ExternalNetwork* external) + : model(m), shape(makeKShape(ow, oh)), engine(m, shape, external), maxWidth(rw), maxHeight(rh), + fallback(p, d, rw, rh, ow, oh) {} + void inputs(VkCommandBuffer cb, const GameFrame& frame); + void result(VkCommandBuffer cb, const GameFrame& frame); +}; +GameUpscaler::GameUpscaler(Model& model, VkPhysicalDevice physical, VkDevice device, uint32_t outputWidth, uint32_t outputHeight, + uint32_t maxRenderWidth, uint32_t maxRenderHeight, const ExternalNetwork* external) + : impl(std::make_unique(model, physical, device, outputWidth, outputHeight, maxRenderWidth, maxRenderHeight, external)) { + if (!maxRenderWidth || !maxRenderHeight || maxRenderWidth > outputWidth || maxRenderHeight > outputHeight) + throw std::runtime_error("render dimensions must be positive and no larger than the output"); +} +GameUpscaler::~GameUpscaler() = default; + +// Everything before the network: input conversion, exposure setup and the input stage (and a reset's clears). +void GameUpscaler::Impl::inputs(VkCommandBuffer cb, const GameFrame& frame) { + auto& g = *this; + const uint32_t w = frame.settings.renderWidth, h = frame.settings.renderHeight; + if (!w || !h || w > g.maxWidth || h > g.maxHeight) throw std::runtime_error("render size exceeds the adapter's inputs"); + if (!gameColorFormat(frame.color.format) || !gameMotionFormat(frame.motion.format) || + !gameDepthFormat(frame.depth.format, frame.depth.aspect) || !gameOutputFormat(frame.output.format)) + throw std::runtime_error("unsupported game texture format"); + const bool gameExposure = frame.exposure.view != VK_NULL_HANDLE; + if (gameExposure && frame.settings.autoExposure) throw std::runtime_error("game exposure and automatic exposure are exclusive"); + // Validates the settings before anything is recorded. + FrameParams parameters = makeKFrameParams(g.shape, frame.settings); + const bool directColor = frame.color.view != VK_NULL_HANDLE; + const bool directMotion = frame.motion.view != VK_NULL_HANDLE; + const bool directDepth = frame.depth.view != VK_NULL_HANDLE; + const bool directOutput = game_detail::directOutput(frame.output, frame.color, frame.motion, frame.depth, g.model.hasRgb10Output()); + auto origin = [&](const GameImage& image, bool direct, int32_t* inOrigin, int32_t* outOrigin, int32_t* outMax) { + if (!direct) return; + if ((image.x || image.y) && !g.model.hasDirectOrigins()) + throw std::runtime_error("recompile the model for direct texture rectangles"); + if (uint64_t(image.x) + w > INT32_MAX || uint64_t(image.y) + h > INT32_MAX) + throw std::runtime_error("texture rectangle exceeds shader coordinate range"); + inOrigin[0] = outOrigin[0] = int32_t(image.x); inOrigin[1] = outOrigin[1] = int32_t(image.y); + outMax[0] += int32_t(image.x); outMax[1] += int32_t(image.y); + }; + origin(frame.color, directColor, parameters.input.colOrigin, parameters.output.colOrigin, parameters.output.colMax); + origin(frame.motion, directMotion, parameters.input.mvOrigin, parameters.output.mvOrigin, parameters.output.mvMax); + origin(frame.depth, directDepth, parameters.input.depOrigin, parameters.output.depOrigin, parameters.output.depMax); + if (directOutput) { + if (uint64_t(frame.output.x) + g.shape.outputWidth > INT32_MAX || uint64_t(frame.output.y) + g.shape.outputHeight > INT32_MAX) + throw std::runtime_error("output rectangle exceeds shader coordinate range"); + parameters.output.outOrigin[0] = int32_t(frame.output.x); parameters.output.outOrigin[1] = int32_t(frame.output.y); + } + if (!g.started) { + if (!g.model.uploadRecorded()) g.model.recordUpload(cb); + g.engine.recordInitialize(cb); + g.started = true; + } + const uint32_t mw = frame.settings.displayMotion ? g.shape.outputWidth : w, mh = frame.settings.displayMotion ? g.shape.outputHeight : h; + g.fallback.record(cb, frame.color, frame.motion, frame.depth, w, h, mw, mh, directColor, directMotion, directDepth); + ExposureControl exposure; + exposure.size[0] = int32_t(w); exposure.size[1] = int32_t(h); + if (directColor) { exposure.origin[0] = int32_t(frame.color.x); exposure.origin[1] = int32_t(frame.color.y); } + exposure.game = gameExposure; + exposure.measure = !gameExposure && (g.frames % 2 == 0 || frame.settings.reset); + exposure.reset = frame.settings.reset; + exposure.exposureScale = frame.settings.exposureScale; exposure.preExposure = frame.settings.preExposure; + // The window schedule follows the evaluation counter, which advances every frame and is not restarted by a + // reset; the external network's phase counter must advance in lockstep, so this is the only place it moves. + auto windows = makeKWindows(g.shape, g.frames++); + FrameImages images{directColor ? frame.color.view : g.fallback.color.view, directMotion ? frame.motion.view : g.fallback.motion.view, + directDepth ? frame.depth.view : g.fallback.depth.view, directColor ? frame.color.layout : VK_IMAGE_LAYOUT_GENERAL, + directMotion ? frame.motion.layout : VK_IMAGE_LAYOUT_GENERAL, directDepth ? frame.depth.layout : VK_IMAGE_LAYOUT_GENERAL, + gameExposure ? frame.exposure.view : VK_NULL_HANDLE, frame.exposure.layout}; + g.engine.setFrame(images, parameters.input, parameters.output, &windows, + (frame.settings.autoExposure || gameExposure) ? &exposure : nullptr, + directOutput ? FrameOutput{frame.output.view, frame.output.format} : FrameOutput{}); +} +// Everything after the network: the copy into the game's output texture when the store was not direct. +void GameUpscaler::Impl::result(VkCommandBuffer cb, const GameFrame& frame) { + auto& g = *this; + if (game_detail::directOutput(frame.output, frame.color, frame.motion, frame.depth, g.model.hasRgb10Output())) return; + game_detail::copyOutput(cb, g.engine.outputImage(), frame.output, g.shape.outputWidth, g.shape.outputHeight); +} +void GameUpscaler::record(VkCommandBuffer cb, const GameFrame& frame) { + impl->inputs(cb, frame); + impl->engine.recordFrame(cb); // begins and ends with barriers against all other work + impl->result(cb, frame); +} +void GameUpscaler::recordFront(VkCommandBuffer cb, const GameFrame& frame) { + impl->inputs(cb, frame); + impl->engine.recordFront(cb); +} +void GameUpscaler::recordBack(VkCommandBuffer cb, const GameFrame& frame) { + impl->engine.recordBack(cb); + impl->result(cb, frame); +} +} // namespace d4r diff --git a/engine/game.h b/engine/game.h new file mode 100644 index 0000000..74fee74 --- /dev/null +++ b/engine/game.h @@ -0,0 +1,51 @@ +// Adapter between a game's own textures and the preset K engine. Everything is recorded into the +// game's command buffer: direct texture access (or format conversion), inference and output. +#pragma once +#include "runtime.h" + +namespace d4r { +struct GameImage { + VkImage image = VK_NULL_HANDLE; + VkFormat format = VK_FORMAT_UNDEFINED; + VkImageLayout layout = VK_IMAGE_LAYOUT_GENERAL; // at record time; must match the supplied sampled views or allow transfers + VkImageAspectFlags aspect = VK_IMAGE_ASPECT_COLOR_BIT; + uint32_t x = 0, y = 0; // origin of the used rectangle + // Optional compatible sampled view (inputs) or RGBA16F/RGB10A2 storage view (output, GENERAL). + // The caller guarantees usage/format compatibility and keeps image + view alive until adapter destruction. + // Null selects the existing copy/conversion path. Sampled views use layout above. + VkImageView view = VK_NULL_HANDLE; +}; +struct GameFrame { + GameImage color, depth, motion, output; + // Optional 1x1 float exposure (texel 0 is read): the game's value replaces measurement. Exclusive with settings.autoExposure. + GameImage exposure; + FrameSettings settings; // render size, jitter, motion scale, reset +}; +// Formats record() converts. Depth may be a depth-aspect D32 image or a one-channel colour image. +bool gameColorFormat(VkFormat format); +bool gameMotionFormat(VkFormat format); +bool gameDepthFormat(VkFormat format, VkImageAspectFlags aspect); +bool gameOutputFormat(VkFormat format); + +class GameUpscaler { +public: + // The model must outlive this object. Inputs up to maxRenderWidth x maxRenderHeight. With an external + // network (ExternalNetwork) the eleven layers run outside the engine and the game's command list is split + // around them (recordFront/recordBack); without one record() runs the whole network. + GameUpscaler(Model& model, VkPhysicalDevice physical, VkDevice device, uint32_t outputWidth, uint32_t outputHeight, + uint32_t maxRenderWidth, uint32_t maxRenderHeight, const ExternalNetwork* external = nullptr); + ~GameUpscaler(); // only once every recorded frame has completed on the GPU + GameUpscaler(const GameUpscaler&) = delete; + GameUpscaler& operator=(const GameUpscaler&) = delete; + // Frames may be recorded while earlier ones are still queued; they must execute in recording order. + void record(VkCommandBuffer command, const GameFrame& frame); + // With an external network: the part before it (input copies, exposure, input stage; writes the tokens) and + // the part after it (final store; reads the head). The network runs between the two on the GPU and the + // caller orders them (the shim splits the game's command list there). + void recordFront(VkCommandBuffer command, const GameFrame& frame); + void recordBack(VkCommandBuffer command, const GameFrame& frame); +private: + struct Impl; + std::unique_ptr impl; +}; +} // namespace d4r diff --git a/engine/game_internal.h b/engine/game_internal.h new file mode 100644 index 0000000..6f340b1 --- /dev/null +++ b/engine/game_internal.h @@ -0,0 +1,139 @@ +#pragma once +#include "game.h" +#include "vk_dispatch.h" +#include + +namespace d4r::game_detail { +inline void check(VkResult result, const char* operation) { + if (result != VK_SUCCESS) throw std::runtime_error(std::string(operation) + ": VkResult " + std::to_string(result)); +} +inline uint32_t memoryType(VkPhysicalDevice physical, uint32_t bits) { + VkPhysicalDeviceMemoryProperties memory; + vkGetPhysicalDeviceMemoryProperties(physical, &memory); + for (uint32_t i = 0; i < memory.memoryTypeCount; ++i) + if ((bits & (1u << i)) && (memory.memoryTypes[i].propertyFlags & VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT)) return i; + throw std::runtime_error("no device-local Vulkan memory type"); +} +// Canonical fallback plane; direct sampled inputs do not use it. +struct Plane { + VkDevice device = VK_NULL_HANDLE; + VkImage image = VK_NULL_HANDLE; + VkDeviceMemory memory = VK_NULL_HANDLE; + VkImageView view = VK_NULL_HANDLE; + ~Plane() { + if (view) vkDestroyImageView(device, view, nullptr); + if (image) vkDestroyImage(device, image, nullptr); + if (memory) vkFreeMemory(device, memory, nullptr); + } + void create(VkPhysicalDevice physical, VkDevice dev, VkFormat format, uint32_t width, uint32_t height) { + device = dev; + VkImageCreateInfo info{VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO}; + info.imageType = VK_IMAGE_TYPE_2D; info.format = format; info.extent = {width, height, 1}; + info.mipLevels = info.arrayLayers = 1; info.samples = VK_SAMPLE_COUNT_1_BIT; + info.usage = VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT; + check(vkCreateImage(dev, &info, nullptr, &image), "vkCreateImage"); + VkMemoryRequirements requirements; vkGetImageMemoryRequirements(dev, image, &requirements); + VkMemoryAllocateInfo allocation{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO}; + allocation.allocationSize = requirements.size; allocation.memoryTypeIndex = memoryType(physical, requirements.memoryTypeBits); + check(vkAllocateMemory(dev, &allocation, nullptr, &memory), "vkAllocateMemory"); + check(vkBindImageMemory(dev, image, memory, 0), "vkBindImageMemory"); + VkImageViewCreateInfo vi{VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO}; + vi.image = image; vi.viewType = VK_IMAGE_VIEW_TYPE_2D; vi.format = format; + vi.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + check(vkCreateImageView(dev, &vi, nullptr, &view), "vkCreateImageView"); + } +}; +inline void barrier(VkCommandBuffer cb, VkPipelineStageFlags source, VkPipelineStageFlags destination, VkAccessFlags src, VkAccessFlags dst) { + VkMemoryBarrier b{VK_STRUCTURE_TYPE_MEMORY_BARRIER}; + b.srcAccessMask = src; b.dstAccessMask = dst; + vkCmdPipelineBarrier(cb, source, destination, 0, 1, &b, 0, nullptr, 0, nullptr); +} +inline bool directOutput(const GameImage& output, const GameImage& color, const GameImage& motion, const GameImage& depth, bool rgb10) { + return output.view != VK_NULL_HANDLE && (output.format == VK_FORMAT_R16G16B16A16_SFLOAT || + (rgb10 && output.format == VK_FORMAT_A2B10G10R10_UNORM_PACK32)) && + output.layout == VK_IMAGE_LAYOUT_GENERAL && output.image != color.image && + output.image != motion.image && output.image != depth.image; +} +inline void copyInput(VkCommandBuffer cb, const GameImage& source, VkImage destination, uint32_t width, uint32_t height, + VkBuffer depthStaging = VK_NULL_HANDLE) { + if (source.aspect & VK_IMAGE_ASPECT_DEPTH_BIT) { + VkBufferImageCopy region{}; + region.imageSubresource = {VK_IMAGE_ASPECT_DEPTH_BIT, 0, 0, 1}; + region.imageOffset = {int32_t(source.x), int32_t(source.y), 0}; region.imageExtent = {width, height, 1}; + vkCmdCopyImageToBuffer(cb, source.image, source.layout, depthStaging, 1, ®ion); + barrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_TRANSFER_READ_BIT); + region.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT; region.imageOffset = {0, 0, 0}; + vkCmdCopyBufferToImage(cb, depthStaging, destination, VK_IMAGE_LAYOUT_GENERAL, 1, ®ion); + } else { + VkImageBlit region{}; + region.srcSubresource = region.dstSubresource = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1}; + region.srcOffsets[0] = {int32_t(source.x), int32_t(source.y), 0}; + region.srcOffsets[1] = {int32_t(source.x + width), int32_t(source.y + height), 1}; + region.dstOffsets[1] = {int32_t(width), int32_t(height), 1}; + vkCmdBlitImage(cb, source.image, source.layout, destination, VK_IMAGE_LAYOUT_GENERAL, 1, ®ion, VK_FILTER_NEAREST); + } +} +// Both adapters allocate conversion resources only when an input cannot be borrowed. +struct CanonicalInputs { + VkPhysicalDevice physical; + VkDevice device; + uint32_t width, height, motionWidth, motionHeight; + Plane color, motion, depth; + VkBuffer depthStaging = VK_NULL_HANDLE; + VkDeviceMemory depthMemory = VK_NULL_HANDLE; + CanonicalInputs(VkPhysicalDevice p, VkDevice d, uint32_t w, uint32_t h, uint32_t mw, uint32_t mh) + : physical(p), device(d), width(w), height(h), motionWidth(mw), motionHeight(mh) {} + ~CanonicalInputs() { + if (depthStaging) vkDestroyBuffer(device, depthStaging, nullptr); + if (depthMemory) vkFreeMemory(device, depthMemory, nullptr); + } + void ensure(VkCommandBuffer cb, Plane& plane, VkFormat format, uint32_t w, uint32_t h) { + if (plane.image) return; + plane.create(physical, device, format, w, h); + VkImageMemoryBarrier layout{VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER}; + layout.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; + layout.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED; layout.newLayout = VK_IMAGE_LAYOUT_GENERAL; + layout.srcQueueFamilyIndex = layout.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; + layout.image = plane.image; layout.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + vkCmdPipelineBarrier(cb, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, + 0, 0, nullptr, 0, nullptr, 1, &layout); + } + void record(VkCommandBuffer cb, const GameImage& col, const GameImage& mv, const GameImage& dep, + uint32_t w, uint32_t h, uint32_t mw, uint32_t mh, bool directColor, bool directMotion, bool directDepth) { + if (directColor && directMotion && directDepth) return; + if (!directColor) ensure(cb, color, VK_FORMAT_R16G16B16A16_SFLOAT, width, height); + if (!directMotion) ensure(cb, motion, VK_FORMAT_R16G16_SFLOAT, motionWidth, motionHeight); + if (!directDepth) { + ensure(cb, depth, VK_FORMAT_R32_SFLOAT, width, height); + if ((dep.aspect & VK_IMAGE_ASPECT_DEPTH_BIT) && !depthStaging) { + VkBufferCreateInfo info{VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO}; + info.size = VkDeviceSize(width) * height * 4; + info.usage = VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT; + check(vkCreateBuffer(device, &info, nullptr, &depthStaging), "vkCreateBuffer"); + VkMemoryRequirements requirements; vkGetBufferMemoryRequirements(device, depthStaging, &requirements); + VkMemoryAllocateInfo allocation{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO}; + allocation.allocationSize = requirements.size; + allocation.memoryTypeIndex = memoryType(physical, requirements.memoryTypeBits); + check(vkAllocateMemory(device, &allocation, nullptr, &depthMemory), "vkAllocateMemory"); + check(vkBindBufferMemory(device, depthStaging, depthMemory, 0), "vkBindBufferMemory"); + } + } + barrier(cb, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, + VK_ACCESS_MEMORY_WRITE_BIT | VK_ACCESS_MEMORY_READ_BIT, + VK_ACCESS_TRANSFER_READ_BIT | VK_ACCESS_TRANSFER_WRITE_BIT); + if (!directColor) copyInput(cb, col, color.image, w, h); + if (!directMotion) copyInput(cb, mv, motion.image, mw, mh); + if (!directDepth) copyInput(cb, dep, depth.image, w, h, depthStaging); + } +}; +inline void copyOutput(VkCommandBuffer cb, VkImage source, const GameImage& destination, uint32_t width, uint32_t height) { + VkImageBlit region{}; + region.srcSubresource = region.dstSubresource = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1}; + region.srcOffsets[1] = {int32_t(width), int32_t(height), 1}; + region.dstOffsets[0] = {int32_t(destination.x), int32_t(destination.y), 0}; + region.dstOffsets[1] = {int32_t(destination.x + width), int32_t(destination.y + height), 1}; + vkCmdBlitImage(cb, source, VK_IMAGE_LAYOUT_GENERAL, destination.image, destination.layout, 1, ®ion, VK_FILTER_NEAREST); + barrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, + VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT); +} +} // namespace d4r::game_detail diff --git a/engine/game_m.cpp b/engine/game_m.cpp new file mode 100644 index 0000000..a1c213d --- /dev/null +++ b/engine/game_m.cpp @@ -0,0 +1,74 @@ +#include "game_m.h" +#include "game_internal.h" + +namespace d4r { + +struct GameUpscalerM::Impl { + MModel& model; + MEngine engine; + uint32_t width, height, outputWidth, outputHeight; + game_detail::CanonicalInputs fallback; + bool started = false; + Impl(MModel& m, VkPhysicalDevice p, VkDevice d, uint32_t ow, uint32_t oh, uint32_t rw, uint32_t rh, const ExternalNetwork* external) + : model(m), engine(m, rw, rh, ow, oh, external), width(rw), height(rh), outputWidth(ow), outputHeight(oh), + fallback(p, d, rw, rh, rw, rh) {} + void inputs(VkCommandBuffer cb, const GameFrameM& frame); + void result(VkCommandBuffer cb, const GameFrameM& frame); +}; +GameUpscalerM::GameUpscalerM(MModel& model, VkPhysicalDevice physical, VkDevice device, uint32_t outputWidth, uint32_t outputHeight, + uint32_t renderWidth, uint32_t renderHeight, const ExternalNetwork* external) + : impl(std::make_unique(model, physical, device, outputWidth, outputHeight, renderWidth, renderHeight, external)) {} +GameUpscalerM::~GameUpscalerM() = default; + +// the game's inputs into the engine's images, and the frame's settings +void GameUpscalerM::Impl::inputs(VkCommandBuffer cb, const GameFrameM& frame) { + auto& g = *this; + const uint32_t w = g.width, h = g.height; + if (!gameColorFormat(frame.color.format) || !gameMotionFormat(frame.motion.format) || + !gameDepthFormat(frame.depth.format, frame.depth.aspect) || !gameOutputFormat(frame.output.format)) + throw std::runtime_error("unsupported game texture format"); + // Nonzero input origins retain conversion copies until M's shader parameter blocks carry origins. + const bool directColor = frame.color.view != VK_NULL_HANDLE && !frame.color.x && !frame.color.y; + const bool directMotion = frame.motion.view != VK_NULL_HANDLE && !frame.motion.x && !frame.motion.y; + const bool directDepth = frame.depth.view != VK_NULL_HANDLE && !frame.depth.x && !frame.depth.y; + const bool directOutput = game_detail::directOutput(frame.output, frame.color, frame.motion, frame.depth, g.model.hasRgb10Output()) && + !frame.output.x && !frame.output.y; + if (!g.started) { + if (!g.model.uploadRecorded()) g.model.recordUpload(cb); + g.engine.recordInitialize(cb); + g.started = true; + } + g.fallback.record(cb, frame.color, frame.motion, frame.depth, w, h, w, h, directColor, directMotion, directDepth); + MFrameSettings s; + s.jitter[0] = frame.jitter[0]; s.jitter[1] = frame.jitter[1]; + s.motionScale[0] = frame.motionToRender[0]; s.motionScale[1] = frame.motionToRender[1]; + s.motionOffset[0] = frame.motionOffset[0]; s.motionOffset[1] = frame.motionOffset[1]; + s.reset = frame.reset; s.depthInverted = frame.depthInverted; + FrameImages images{}; + images.color = directColor ? frame.color.view : g.fallback.color.view; + images.motion = directMotion ? frame.motion.view : g.fallback.motion.view; + images.depth = directDepth ? frame.depth.view : g.fallback.depth.view; + images.colorLayout = directColor ? frame.color.layout : VK_IMAGE_LAYOUT_GENERAL; + images.motionLayout = directMotion ? frame.motion.layout : VK_IMAGE_LAYOUT_GENERAL; + images.depthLayout = directDepth ? frame.depth.layout : VK_IMAGE_LAYOUT_GENERAL; + g.engine.setFrame(images, s, directOutput ? FrameOutput{frame.output.view, frame.output.format} : FrameOutput{}); +} +void GameUpscalerM::Impl::result(VkCommandBuffer cb, const GameFrameM& frame) { + auto& g = *this; + if (game_detail::directOutput(frame.output, frame.color, frame.motion, frame.depth, g.model.hasRgb10Output()) && !frame.output.x && !frame.output.y) return; + game_detail::copyOutput(cb, g.engine.outputImage(), frame.output, g.outputWidth, g.outputHeight); +} +void GameUpscalerM::record(VkCommandBuffer cb, const GameFrameM& frame) { + impl->inputs(cb, frame); + impl->engine.recordFrame(cb); // begins and ends with barriers against all other work + impl->result(cb, frame); +} +void GameUpscalerM::recordFront(VkCommandBuffer cb, const GameFrameM& frame) { + impl->inputs(cb, frame); + impl->engine.recordFront(cb); +} +void GameUpscalerM::recordBack(VkCommandBuffer cb, const GameFrameM& frame) { + impl->engine.recordBack(cb); + impl->result(cb, frame); +} +} // namespace d4r diff --git a/engine/game_m.h b/engine/game_m.h new file mode 100644 index 0000000..b6c2cc5 --- /dev/null +++ b/engine/game_m.h @@ -0,0 +1,33 @@ +// Adapter between a game's own textures and the preset M engine: input conversion, the frame and the copy into +// the game's output texture, all recorded into the game's command buffer. Separate from GameUpscaler (preset K). +#pragma once +#include "game.h" + +namespace d4r { +struct GameFrameM { + GameImage color, depth, motion, output; // as for GameUpscaler; views are not used (inputs are always copied) + float jitter[2] = {0, 0}; + float motionToRender[2] = {1, 1}; // (motion texel + motionOffset) * this = render pixels + float motionOffset[2] = {0, 0}; // jittered motion vectors: jitter - previous jitter, motion texel units + bool reset = false; + bool depthInverted = true; +}; +class GameUpscalerM { +public: + // The model must outlive this object. The render size is fixed and at most half the output size. + GameUpscalerM(MModel& model, VkPhysicalDevice physical, VkDevice device, uint32_t outputWidth, uint32_t outputHeight, + uint32_t renderWidth, uint32_t renderHeight, const ExternalNetwork* external = nullptr); + ~GameUpscalerM(); // only once every recorded frame has completed on the GPU + GameUpscalerM(const GameUpscalerM&) = delete; + GameUpscalerM& operator=(const GameUpscalerM&) = delete; + void record(VkCommandBuffer command, const GameFrameM& frame); + // With an external network (ExternalNetwork): the part before it (input copies, exposure, input stage) and the + // part after it (expansion, reconstruction, downsample, the copy into the game's output). The network must run + // between the two on the GPU; the caller orders them (the shim splits the game's command list there). + void recordFront(VkCommandBuffer command, const GameFrameM& frame); + void recordBack(VkCommandBuffer command, const GameFrameM& frame); +private: + struct Impl; + std::unique_ptr impl; +}; +} // namespace d4r diff --git a/engine/gen_activation_m.py b/engine/gen_activation_m.py new file mode 100755 index 0000000..c2e8119 --- /dev/null +++ b/engine/gen_activation_m.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python3 +"""Generate activation_m.glsl's ACT table: binary16 sigmoid, tanh and gate of every E4M3 code. + +usage: gen_activation_m.py print the ACT declaration + gen_activation_m.py --check exit 1 unless engine/activation_m.glsl holds exactly this table + +Row c is uvec2(tanh << 16 | sigmoid, gate) for the E4M3 value of code c: + sigmoid 1 / (1 + exp2(-1.442695 * x)) in f32, rounded to binary16 + tanh correctly rounded binary16 tanh(x) + gate .5 * tanh(.5 * x) + .5 with each operation rounded to binary16 +The NaN codes (127, 255) give quiet NaN; the sigmoid keeps the input's sign, as f32 arithmetic propagates it. +""" +import math +from pathlib import Path +import re +import sys + +import numpy as np + +GLSL = Path(__file__).resolve().parent / 'activation_m.glsl' +NAN_ROWS = {127: (0x7e007e00, 0x7e00), 255: (0x7e00fe00, 0x7e00)} + + +def e4m3(code): + sign = -1.0 if code & 128 else 1.0 + exponent, mantissa = (code >> 3) & 15, code & 7 + if exponent == 0: + return sign * mantissa / 8 * 2.0 ** -6 + return sign * (1 + mantissa / 8) * 2.0 ** (exponent - 7) + + +def bits(value): + return int(np.array(value, dtype=np.float16).view(np.uint16)) + + +def row(code): + if code in NAN_ROWS: + return NAN_ROWS[code] + x = e4m3(code) + with np.errstate(over='ignore'): + sigmoid = np.float32(1) / (np.float32(1) + np.exp2(np.float32(-1.442695) * np.float32(x), dtype=np.float32)) + # float64 tanh rounded once to binary16 is correctly rounded for every E4M3 input (no value lies near a tie). + tanh = np.float16(math.tanh(x)) + half = np.float16(0.5) + gate = np.float16(np.float16(half * np.float16(math.tanh(float(half * np.float16(x))))) + half) + return bits(tanh) << 16 | bits(sigmoid), bits(gate) + + +def declaration(): + rows = ',\n'.join(f' uvec2(0x{x:08x}u, 0x{y:08x}u)' for x, y in map(row, range(256))) + return f'const uvec2 ACT[256] = uvec2[256](\n{rows}\n);' + + +def main(argv): + if argv == ['--check']: + match = re.search(r'const uvec2 ACT\[256\] = uvec2\[256\]\(.*?\n\);', GLSL.read_text(), re.S) + if match is None or match.group(0) != declaration(): + print(f'{GLSL} does not hold the generated table', file=sys.stderr) + return 1 + return 0 + if argv: + sys.exit(__doc__) + print(declaration()) + return 0 + + +if __name__ == '__main__': + sys.exit(main(sys.argv[1:])) diff --git a/engine/hip_net.cpp b/engine/hip_net.cpp new file mode 100644 index 0000000..edade73 --- /dev/null +++ b/engine/hip_net.cpp @@ -0,0 +1,332 @@ +// d4r engine: the DLSS network on d4r's native HIP layer kernels. See hip_net.h. +// +// The backend is shared by the two presets and only the plan differs: preset M runs ten layers with the +// kernels/m parameter block (kernels/m/swin_common.h) in two window-alignment phases, preset K runs eleven +// layers with kernels/k/pwin_common.h's PwinParams (the bottleneck layer as its four phase kernels) in the +// eight-frame window cycle of makeKWindows. Everything else - the stream, the immutable captured graphs, the +// events and the per-evaluation frame counter - is the same code. +#define __HIP_PLATFORM_AMD__ +#include "hip_net.h" +#include +#include +#include +#include +#include +#include +namespace d4r { +namespace { +void check(hipError_t e, const char* what) { + if (e != hipSuccess) throw std::runtime_error(std::string("HIP: ") + what + " failed: " + hipGetErrorString(e)); +} +#define HIPCHECK(x) check((x), #x) +std::vector slurp(const std::string& path) { + std::ifstream f(path, std::ios::binary); + if (!f) throw std::runtime_error("cannot read " + path); + return std::vector((std::istreambuf_iterator(f)), {}); +} +// preset M's kernels' parameter block (kernels/m/swin_common.h: CommonParams, TubeParams). Both blocks are +// plain aggregates: every parameter set is value-initialized (all bytes zero) before its fields are filled. +struct LayerParamsM { + void* w; int32_t sx, sy, tw, th; void* in; void* skip; void* out; void* merged; + int32_t r69, r70, r71, r72; void* inFlags; void* outFlags; +}; +static_assert(sizeof(LayerParamsM) == 88 && offsetof(LayerParamsM, in) == 24, "M param block"); +// preset K's kernels' parameter block (kernels/k/pwin_layer.h: PwinParams) +struct LayerParamsK { + int32_t W, H; void* in; void* skip; void* out24; void* out32; uint64_t pad40, pad48; int32_t sx, sy; void* w; + uint8_t rest[104]; +}; +static_assert(sizeof(LayerParamsK) == 176 && offsetof(LayerParamsK, in) == 8 && offsetof(LayerParamsK, w) == 64, "K param block"); +union Params { + LayerParamsM m; + LayerParamsK k; + uint8_t raw[176]; +}; +// one code object of a layer: the main entry, the weight-image preparation, and the K bottleneck's post phases +struct Kernel { + hipModule_t module = nullptr; + hipFunction_t main = nullptr, prep = nullptr; + hipFunction_t post[8] = {}; + uint32_t prepBlocks = 0, blockZ = 4, gridX = 0, gridY = 0; + uint32_t postGridX[8] = {}, postGridY[8] = {}, postBlockX[8] = {}, postBlockY[8] = {}, postBlockZ[8] = {}; + int postCount = 0; +}; +// one immutable launch: entry point, grid, block and the (frozen) parameter block +struct Node { + hipFunction_t function = nullptr; + uint32_t gx = 1, gy = 1, bx = 32, by = 1, bz = 1; + Params params{}; +}; +uint32_t global32(hipModule_t module, const char* name, uint32_t fallback) { + hipDeviceptr_t at = nullptr; size_t bytes = 0; uint32_t v = fallback; + if (hipModuleGetGlobal(&at, &bytes, module, name) == hipSuccess && bytes == 4) { + if (hipMemcpyDtoH(&v, at, 4) != hipSuccess || v == 0) v = fallback; + } + return v; +} +void* import(int fd, uint64_t bytes, hipExternalMemory_t& memory) { + hipExternalMemoryHandleDesc desc{}; desc.type = hipExternalMemoryHandleTypeOpaqueFd; desc.handle.fd = fd; desc.size = bytes; + HIPCHECK(hipImportExternalMemory(&memory, &desc)); + hipExternalMemoryBufferDesc buffer{}; buffer.offset = 0; buffer.size = bytes; + void* at = nullptr; HIPCHECK(hipExternalMemoryGetMappedBuffer(&at, memory, &buffer)); + return at; +} +// makeKWindows (engine/runtime.cpp): dec5 moves once per frame around an eight-frame cycle and every +// shallower layer follows the one below it as (2 * shift + 4) mod 8. Role 0 = dec5, 1 = enc4/dec4, +// 2 = enc3/dec3; every other layer keeps shift 4. +constexpr int32_t KCycle[8][2] = {{0, 2}, {5, 6}, {4, 0}, {1, 4}, {7, 5}, {2, 1}, {3, 7}, {6, 3}}; +void kShift(int role, uint32_t frame, int32_t& sx, int32_t& sy) { + if (role < 0) { sx = sy = 4; return; } + int32_t x = KCycle[frame % 8][0], y = KCycle[frame % 8][1]; + for (int depth = 0; depth < role; ++depth) { x = (2 * x + 4) & 7; y = (2 * y + 4) & 7; } + sx = x; sy = y; +} +// preset K's eleven layers in network order: token-grid divisor, input channels, core channels, the channels +// written to out24 (the patch-merged skip for encoders, the block output for the decoders, the 40-channel head +// for dec0) and the window-schedule role. +struct KLayer { + const char* name; uint32_t divisor, cin, c, out24ch; int role; +}; +constexpr KLayer KPlan[11] = { + {"enc0", 1, 16, 32, 64, -1}, {"enc1", 2, 64, 64, 64, -1}, {"enc2", 4, 64, 64, 96, -1}, + {"enc3", 8, 96, 96, 128, 2}, {"enc4", 16, 128, 128, 160, 1}, {"dec5", 32, 160, 160, 160, 0}, + {"dec4", 16, 160, 128, 128, 1}, {"dec3", 8, 128, 96, 96, 2}, {"dec2", 4, 96, 64, 64, -1}, + {"dec1", 2, 64, 64, 64, -1}, {"dec0", 1, 64, 32, 40, -1}, +}; +constexpr uint32_t KLayers = 11, MLayers = 10, MaxGraphs = 8; +} // namespace + +struct HipNetwork::Impl { + std::vector kernels; + std::vector preps; + std::vector nodes[MaxGraphs]; + std::vector owned; + hipExternalMemory_t shared[2]{}; + hipStream_t stream = nullptr; + hipGraph_t graph[MaxGraphs]{}; + hipGraphExec_t executable[MaxGraphs]{}; + uint32_t graphs = 2, frame = 0; + hipEvent_t start = nullptr, stop = nullptr; + bool pending = false; + ~Impl() { + if (stream) (void)hipStreamSynchronize(stream); + for (auto e : executable) if (e) (void)hipGraphExecDestroy(e); + for (auto g : graph) if (g) (void)hipGraphDestroy(g); + for (void* p : owned) (void)hipFree(p); + for (auto& k : kernels) if (k.module) (void)hipModuleUnload(k.module); + for (auto m : shared) if (m) (void)hipDestroyExternalMemory(m); + if (start) (void)hipEventDestroy(start); + if (stop) (void)hipEventDestroy(stop); + if (stream) (void)hipStreamDestroy(stream); + } + void run(const Node& n) { + void* args[] = {const_cast(&n.params)}; + HIPCHECK(hipModuleLaunchKernel(n.function, n.gx, n.gy, 1, n.bx, n.by, n.bz, 0, stream, args, nullptr)); + } +}; + +HipNetwork::HipNetwork(const std::string& directory, uint32_t width, uint32_t height, int tokensFd, int outputFd, + uint64_t tokensBytes, uint64_t outputBytes) : impl(std::make_unique()) { + auto& n = *impl; + HIPCHECK(hipInit(0)); + int leastPriority = 0, greatestPriority = 0; + HIPCHECK(hipDeviceGetStreamPriorityRange(&leastPriority, &greatestPriority)); + HIPCHECK(hipStreamCreateWithPriority(&n.stream, hipStreamNonBlocking, greatestPriority)); + HIPCHECK(hipEventCreate(&n.start)); HIPCHECK(hipEventCreate(&n.stop)); + const auto file = slurp(directory + "/hipnet.bin"); + if (file.size() < 16) throw std::runtime_error("incompatible hipnet.bin"); + const bool kPreset = std::memcmp(file.data(), "D4RKHIP1", 8) == 0; + if (!kPreset && std::memcmp(file.data(), "D4RMHIP1", 8)) throw std::runtime_error("incompatible hipnet.bin"); + uint32_t count = 0; + std::memcpy(&count, file.data() + 8, 4); + const uint32_t layers = kPreset ? KLayers : MLayers; + if (count != layers || file.size() < 16 + 64ull * count) throw std::runtime_error("invalid hipnet.bin"); + void* tokens = import(tokensFd, tokensBytes, n.shared[0]); + void* output = import(outputFd, outputBytes, n.shared[1]); + // Both dimensions describe the layer-0 token grid: preset M's render dimensions, preset K's output + // dimensions (makeKShape: max 256, align 32 of the output quarter resolution). + uint32_t tw = 0, th = 0; + if (kPreset) { + if (width < 256 || height < 256 || width > 8192 || height > 8192) + throw std::runtime_error("standard K output dimensions must be 256..8192"); + tw = std::max(256u, ((width + 3) / 4 + 31) / 32 * 32); + th = std::max(256u, ((height + 3) / 4 + 31) / 32 * 32); + } else { + tw = (width + 31) / 32 * 16; th = (height + 31) / 32 * 16; + } + // Kernel lookup by name; the M tube shares one code object across six layers, the K layer names are unique. + std::vector names; + auto kernelOf = [&](const char* name) -> uint32_t { + size_t k = 0; + while (k < names.size() && names[k] != name) ++k; + if (k < names.size()) return uint32_t(k); + names.emplace_back(name); n.kernels.emplace_back(); + auto& kernel = n.kernels.back(); + const auto code = slurp(directory + "/hip/" + name + ".hsaco"); + HIPCHECK(hipModuleLoadData(&kernel.module, code.data())); + HIPCHECK(hipModuleGetFunction(&kernel.main, kernel.module, name)); + const std::string base(name); + if (hipModuleGetFunction(&kernel.prep, kernel.module, (base + "_prep").c_str()) != hipSuccess) kernel.prep = nullptr; + kernel.prepBlocks = global32(kernel.module, "d4r_prep_blocks", 0); + kernel.blockZ = global32(kernel.module, "d4r_block_z", 4); + kernel.gridX = global32(kernel.module, "d4r_grid_x", 0); + kernel.gridY = global32(kernel.module, "d4r_grid_y", 0); + if (kernel.prep && !kernel.prepBlocks) throw std::runtime_error(base + ": prep kernel without d4r_prep_blocks"); + // The deep layers' code objects also carry their post phases; their grids are compile-time. + for (int j = 1; j <= 8; ++j) { + const std::string tag = base + "_post" + std::to_string(j); + if (hipModuleGetFunction(&kernel.post[kernel.postCount], kernel.module, tag.c_str()) != hipSuccess) break; + const std::string suffix = "_post" + std::to_string(j); + kernel.postGridX[kernel.postCount] = global32(kernel.module, ("d4r" + suffix + "_grid_x").c_str(), 0); + kernel.postGridY[kernel.postCount] = global32(kernel.module, ("d4r" + suffix + "_grid_y").c_str(), 0); + kernel.postBlockX[kernel.postCount] = global32(kernel.module, ("d4r" + suffix + "_block_x").c_str(), 32); + kernel.postBlockY[kernel.postCount] = global32(kernel.module, ("d4r" + suffix + "_block_y").c_str(), 1); + kernel.postBlockZ[kernel.postCount] = global32(kernel.module, ("d4r" + suffix + "_block_z").c_str(), 1); + if (!kernel.postGridX[kernel.postCount]) kernel.postGridX[kernel.postCount] = kernel.gridX; + ++kernel.postCount; + } + // a module keeps one prepared weight image per weights pointer while its slots last; the M tube needs six + if (global32(kernel.module, "d4r_prep_key_slots", 0) < 6 && base.find("tube") != std::string::npos) + throw std::runtime_error(base + ": fewer than six weight image slots"); + return uint32_t(names.size() - 1); + }; + auto allocate = [&](uint64_t bytes) { + void* p = nullptr; HIPCHECK(hipMalloc(&p, size_t(bytes) + 4096)); n.owned.push_back(p); return p; + }; + std::vector layerKernel(layers), divisor(layers); + std::vector shiftX(layers), shiftY(layers); + std::vector weights(layers); + for (uint32_t i = 0; i < count; ++i) { + const char* e = file.data() + 16 + 64 * i; + char name[33]{}; std::memcpy(name, e, 32); + int32_t sx, sy; uint32_t scale; uint64_t at, bytes; + std::memcpy(&sx, e + 32, 4); std::memcpy(&sy, e + 36, 4); std::memcpy(&scale, e + 40, 4); + std::memcpy(&at, e + 48, 8); std::memcpy(&bytes, e + 56, 8); + if (!scale || at > file.size() || bytes > file.size() - at) throw std::runtime_error("invalid hipnet.bin layer"); + if (kPreset && std::strcmp(name, ("dltss_pwin_" + std::string(KPlan[i].name) + "_layer").c_str())) + throw std::runtime_error(std::string("unexpected K layer order: ") + name); + layerKernel[i] = kernelOf(name); + void* w = allocate(bytes); + HIPCHECK(hipMemcpyHtoD(w, file.data() + at, bytes)); + weights[i] = w; divisor[i] = scale; shiftX[i] = sx; shiftY[i] = sy; + } + // The layer plan: activation buffers and the frozen parameter block of every layer at frame/phase 0. + std::vector base(layers); + if (kPreset) { + // NGX packs the encoder skips, the merged outputs and the decoder outputs into one scratch allocation; + // separate buffers only have to respect the channel counts and the wiring, not the original aliasing. + std::vector act(layers), full(5, nullptr); + for (uint32_t i = 0; i < layers; ++i) { + const auto& p = KPlan[i]; + const uint32_t w = tw / p.divisor, h = th / p.divisor; + // an encoder's out24 is the patch-merged 2x2 output, half the resolution of its input + const uint32_t ow = i < 5 ? tw / (2 * p.divisor) : w, oh = i < 5 ? th / (2 * p.divisor) : h; + // dec0 writes the head straight into the caller's output buffer + if (i + 1 != layers) act[i] = allocate(uint64_t(ow) * oh * p.out24ch * 2); + if (i < 5) full[i] = allocate(uint64_t(w) * h * p.c * 2); + } + for (uint32_t i = 0; i < layers; ++i) { + const auto& p = KPlan[i]; + auto& q = base[i].k; + q.W = int32_t(tw / p.divisor); q.H = int32_t(th / p.divisor); + q.in = i == 0 ? tokens : act[i - 1]; + q.skip = i >= 6 ? full[10 - i] : nullptr; + q.out24 = i + 1 == layers ? output : act[i]; + q.out32 = i < 5 ? full[i] : nullptr; + kShift(p.role, 0, q.sx, q.sy); + q.w = weights[i]; + } + } else { + void* skip1 = allocate(uint64_t(tw) * th * 64); + void* merged1 = allocate(uint64_t(tw / 2) * (th / 2) * 96); + void* skip2 = allocate(uint64_t(tw / 2) * (th / 2) * 96); + void* merged2 = allocate(uint64_t(tw / 4) * (th / 4) * 128); + void* tubeA = allocate(uint64_t(tw / 4) * (th / 4) * 128); + void* tubeB = allocate(uint64_t(tw / 4) * (th / 4) * 128); + void* dec2 = allocate(uint64_t(tw / 2) * (th / 2) * 96); + // tokens -> enc1 -> enc2 -> six tube blocks (ping-pong) -> dec2 -> dec1 + void* in[MLayers] = {tokens, merged1, merged2, tubeA, tubeB, tubeA, tubeB, tubeA, tubeB, dec2}; + void* skip[MLayers] = {nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, skip2, skip1}; + void* out[MLayers] = {skip1, skip2, tubeA, tubeB, tubeA, tubeB, tubeA, tubeB, dec2, output}; + void* merged[MLayers] = {merged1, merged2, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr}; + for (uint32_t i = 0; i < layers; ++i) { + auto& q = base[i].m; + q.w = weights[i]; q.sx = shiftX[i]; q.sy = shiftY[i]; + q.tw = int32_t(tw / divisor[i]); q.th = int32_t(th / divisor[i]); + q.in = in[i]; q.skip = skip[i]; q.out = out[i]; q.merged = merged[i]; + } + } + // Immutable graphs: preset M has two horizontal window-alignment phases, preset K the eight-frame cycle. + n.graphs = kPreset ? MaxGraphs : 2; + for (uint32_t phase = 0; phase < n.graphs; ++phase) { + auto& nodes = n.nodes[phase]; + for (uint32_t i = 0; i < layers; ++i) { + const auto& kernel = n.kernels[layerKernel[i]]; + Params params = base[i]; + if (kPreset) kShift(KPlan[i].role, phase, params.k.sx, params.k.sy); + else params.m.sx = shiftX[i] ^ int32_t(phase * 2); // NGX alternates horizontal alignment every evaluation + if (kernel.postCount) { + Node node; node.function = kernel.main; node.params = params; + node.gx = kernel.gridX; node.gy = kernel.gridY ? kernel.gridY : 1; node.bz = kernel.blockZ; + nodes.push_back(node); + for (int j = 0; j < kernel.postCount; ++j) { + Node post; post.function = kernel.post[j]; post.params = params; + post.gx = kernel.postGridX[j]; post.gy = kernel.postGridY[j] ? kernel.postGridY[j] : 1; + post.bx = kernel.postBlockX[j]; post.by = kernel.postBlockY[j]; post.bz = kernel.postBlockZ[j]; + nodes.push_back(post); + } + } else { + Node node; node.function = kernel.main; node.params = params; + if (kPreset) { + node.gx = (params.k.W + params.k.sx + 7) / 8; node.gy = (params.k.H + params.k.sy + 7) / 8; + } else { + node.gx = (params.m.tw + params.m.sx + 7) / 8; node.gy = (params.m.th + params.m.sy + 7) / 8; + } + node.bz = kernel.blockZ; + nodes.push_back(node); + } + } + } + // The weights never change: every layer's weight image is prepared once, here, and found by its pointer. + // A shared code object (the M tube) gets one preparation per layer, i.e. per distinct weights pointer. + for (uint32_t i = 0; i < layers; ++i) { + const auto& kernel = n.kernels[layerKernel[i]]; + if (!kernel.prep) continue; + Node prep; prep.function = kernel.prep; prep.gx = kernel.prepBlocks; prep.bx = 128; prep.params = base[i]; + n.preps.push_back(prep); + } + for (const auto& prep : n.preps) n.run(prep); + HIPCHECK(hipStreamSynchronize(n.stream)); + for (uint32_t phase = 0; phase < n.graphs; ++phase) { + HIPCHECK(hipStreamBeginCapture(n.stream, hipStreamCaptureModeRelaxed)); + try { + for (const auto& node : n.nodes[phase]) n.run(node); + } catch (...) { + (void)hipStreamEndCapture(n.stream, &n.graph[phase]); + throw; + } + HIPCHECK(hipStreamEndCapture(n.stream, &n.graph[phase])); + HIPCHECK(hipGraphInstantiate(&n.executable[phase], n.graph[phase], nullptr, nullptr, 0)); + HIPCHECK(hipGraphUpload(n.executable[phase], n.stream)); + } + HIPCHECK(hipStreamSynchronize(n.stream)); +} +HipNetwork::~HipNetwork() = default; +void HipNetwork::launch() { + auto& n = *impl; + HIPCHECK(hipEventRecord(n.start, n.stream)); + HIPCHECK(hipGraphLaunch(n.executable[n.frame % n.graphs], n.stream)); + ++n.frame; + HIPCHECK(hipEventRecord(n.stop, n.stream)); + n.pending = true; +} +bool HipNetwork::done() { + auto& n = *impl; + if (!n.pending) return true; + const hipError_t e = hipEventQuery(n.stop); + if (e == hipErrorNotReady) return false; + check(e, "hipEventQuery"); n.pending = false; return true; +} +void HipNetwork::wait() { HIPCHECK(hipStreamSynchronize(impl->stream)); impl->pending = false; } +float HipNetwork::lastMilliseconds() { float ms = 0; HIPCHECK(hipEventElapsedTime(&ms, impl->start, impl->stop)); return ms; } +} // namespace d4r diff --git a/engine/hip_net.h b/engine/hip_net.h new file mode 100644 index 0000000..9962bee --- /dev/null +++ b/engine/hip_net.h @@ -0,0 +1,41 @@ +// d4r engine: the DLSS network on d4r's native HIP layer kernels instead of the Vulkan layer shader. HIP reaches +// the hardware's packed FP8 conversions and keeps a block's activations in registers, which a cooperative-matrix +// shader cannot, so the network runs about 1.7x faster here. The stages around it stay on Vulkan: the engine's +// front writes the tokens into memory shared with HIP, launch() runs the layers, the engine's back reads their +// output (ExternalNetwork in runtime.h). +// +// The model directory needs hipnet.bin and hip/*.hsaco (compile_m.py --hip for preset M, compile_k.py --hip for +// preset K) and the header + entry names select the preset: +// D4RMHIP1 (10 layers) - M on the kernels/m layer kernels. width/height are the render dimensions; two +// window-alignment phases, one per evaluation. +// D4RKHIP1 (11 layers) - K on the kernels/k DltssPaddedWinLayer kernels, the bottleneck layer running as its +// main kernel plus post1..3. width/height are the OUTPUT dimensions (the token grid is max(256, align 32 of +// their quarter resolution, makeKShape) and the input/output buffers hold 16 / 40 f16 channels per token, +// row-major, with no layout copy. Eight immutable window graphs, one per evaluation of makeKWindows' cycle. +// Either way every graph is uploaded once and one evaluation is a single hipGraphLaunch: no per-frame allocation, +// no per-frame parameter copy. The frame counter advances once per launch() and is never reset (a reset +// evaluation is still an evaluation), so the caller must pair one launch() with each network evaluation. +#pragma once +#include +#include +#include +namespace d4r { +class HipNetwork { +public: + // tokensFd, outputFd: opaque file descriptors of the two shared Vulkan allocations (ownership passes to HIP); + // allocationBytes: the size of each allocation; the buffers start at offset 0. width/height are preset + // dependent, as above. + HipNetwork(const std::string& modelDirectory, uint32_t width, uint32_t height, int tokensFd, int outputFd, + uint64_t tokensAllocationBytes, uint64_t outputAllocationBytes); + ~HipNetwork(); + HipNetwork(const HipNetwork&) = delete; + HipNetwork& operator=(const HipNetwork&) = delete; + void launch(); // queues the network; input tokens must already be complete + bool done(); // true once the queued layers have finished + void wait(); // blocks until they have + float lastMilliseconds(); // GPU time of the last finished launch +private: + struct Impl; + std::unique_ptr impl; +}; +} // namespace d4r diff --git a/engine/input_k.comp b/engine/input_k.comp new file mode 100644 index 0000000..b84ec7d --- /dev/null +++ b/engine/input_k.comp @@ -0,0 +1,180 @@ +#version 460 +// d4r engine, preset K: the input stage, written from the algorithm described in INPUT_K_NOTES.md. It turns the frame +// into the network's input: 16 values per token (4x4 output pixels), four per 2x2-pixel sub-block: +// the current luma and the reprojected history luma of one pixel (range-compressed, sRGB-encoded), +// one channel of the previous frame's token feature, reprojected and averaged over the token, +// and how far the history's colour lies outside the current 3x3 neighbourhood. +// One thread computes one token (the reference uses one thread per sub-block and warp shuffles for the average). +#extension GL_EXT_shader_explicit_arithmetic_types : require +layout(local_size_x = 8, local_size_y = 8) in; +layout(set = 0, binding = 0, std430) readonly buffer Params +{ + vec2 out2rnd, rndOffs; // render position = (pixel + .5) * out2rnd - .5 + rndOffs + vec2 mvScale, mvOffs; // previous position = pixel + .5 + mvScale * (mv + mvOffs) + vec2 invHist, invValid; // previous position -> history uv; -> [0, 1] on-screen test + vec2 tok2uv; // output pixels -> uv of the token feature texture + ivec2 colOrigin, colSize, depOrigin, depSize, mvOrigin, mvSize; + int tokW; + uint flags; // bit 0: reset; bit 1: regular depth (larger is farther), otherwise inverted; + // bit 2: display-resolution motion (one vector per output pixel, no depth) + float expo, expoIn, expoRatio; // exposure of the statistics; of the luma channels; previous / current + float minW; +} p; +layout(set = 0, binding = 1) uniform sampler2D colorTex; +layout(set = 0, binding = 2) uniform sampler2D histTex; // compressed Y, Co, Cg; bilinear, as NGX creates it +layout(set = 0, binding = 3) uniform sampler2D motionTex; +layout(set = 0, binding = 4) uniform sampler2D depthTex; +layout(set = 0, binding = 5) uniform sampler2D featTex; // previous frame's token feature; bilinear +layout(set = 0, binding = 6, std430) writeonly buffer Out { f16vec4 v[]; } outb; + +float roundHalfAway(float x) { const float t = trunc(x); return abs(x - t) >= 0.5 ? t + sign(x) : t; } +int clamp1(int v, int origin, int hi) { return v > 0 ? min(v + origin, hi) : origin; } +// The reference's expansion of the stored history luma: rcp.approx(1 - v) below 0.99900102 (its threshold, +// one ulp above 0.999), otherwise ex2(((v + 0x3C4FAE29) * 9) * log2e) - 1) / (9 v). +float expandK(float v) +{ + if (v < uintBitsToFloat(0x3F7FBE88u)) + return 1.0 / (1.0 - v); + precise const float y = (v + uintBitsToFloat(0x3C4FAE29u)) * 9.0; + return (exp2(y * uintBitsToFloat(0x3FB8AA3Bu)) - 1.0) / (9.0 * v); +} +vec3 ycocg(vec3 c) { return vec3(0.25 * c.r + 0.5 * c.g + 0.25 * c.b, 0.5 * c.r - 0.5 * c.b, 0.5 * c.g - 0.25 * c.r - 0.25 * c.b); } +// LDR (the game did not set the HDR flag; NGX's *_ldr_* kernels): the colour is display-referred, so there is no +// exposure and no range compression: plain averages, a linear contrast weight, and the luma channels go through +// an inverse and a forward Reinhard curve. +#ifndef LDR +#define LDR 0 +#endif +#if LDR +float logK(float Y) { return fma(Y, 3.0, 1.0); } +#else +float logK(float Y) { precise const float m = max(Y, 0.0); return (log2(fma(m, 12.0, 1.0)) * uintBitsToFloat(0x3F317218u)) / (m + 1e-6); } +#endif +float srgb(float v) { const float x = clamp(v, 0.0, 1.0); return x > 0.0031308 ? fma(exp2(log2(x) * (1.0 / 2.4)), 1.055, -0.055) : 12.92 * x; } +ivec2 nearest(vec2 pixCentre) { const vec2 r = pixCentre * p.out2rnd - 0.5 + p.rndOffs; return ivec2(roundHalfAway(r.x), roundHalfAway(r.y)); } +// Depth with larger = nearer: regular depth is negated, which turns its "nearest" comparisons into the inverted ones. +float depthAt(ivec2 t) +{ + const ivec2 hi = p.depOrigin + p.depSize; + const float d = texelFetch(depthTex, ivec2(clamp1(t.x, p.depOrigin.x, hi.x), clamp1(t.y, p.depOrigin.y, hi.y)), 0).r; + return (p.flags & 2u) != 0u ? -d : d; +} +// motion vector of a texel, clamped to the motion rectangle, in output pixels +vec2 motionTexel(ivec2 t) +{ + const ivec2 hi = p.mvOrigin + p.mvSize; + return p.mvScale * (texelFetch(motionTex, ivec2(clamp1(t.x, p.mvOrigin.x, hi.x), clamp1(t.y, p.mvOrigin.y, hi.y)), 0).xy + p.mvOffs); +} +// motion of the render texel nearest to a pixel, taken from the nearest of its diagonal depth neighbours; +// with display-resolution vectors (flag 4) t is the pixel itself, and its own vector is used +vec2 motionOf(ivec2 t) +{ + if ((p.flags & 4u) != 0u) return motionTexel(t); + const float d0 = depthAt(t), dmm = depthAt(t + ivec2(-1, -1)), dpm = depthAt(t + ivec2(1, -1)), dmp = depthAt(t + ivec2(-1, 1)), dpp = depthAt(t + ivec2(1, 1)); + ivec2 d = ivec2(0); + if (d0 < 0.5 * (dmm + dpp) - 0.001 * abs(d0) || d0 < 0.5 * (dpm + dmp) - 0.001 * abs(d0)) + { + float m = d0; + if (dmm > m) { m = dmm; d = ivec2(-1, -1); } + if (dpm > m) { m = dpm; d = ivec2(1, -1); } + if (dmp > m) { m = dmp; d = ivec2(-1, 1); } + if (dpp > m) { d = ivec2(1, 1); } + } + const ivec2 v = t + d, hi = p.mvOrigin + p.mvSize; + return p.mvScale * (texelFetch(motionTex, ivec2(clamp1(v.x, p.mvOrigin.x, hi.x), clamp1(v.y, p.mvOrigin.y, hi.y)), 0).xy + p.mvOffs); +} +vec3 colourAt(ivec2 t) // exposed, non-negative +{ + const ivec2 hi = p.colOrigin + p.colSize; + return p.expo * max(texelFetch(colorTex, ivec2(clamp1(t.x, p.colOrigin.x, hi.x), clamp1(t.y, p.colOrigin.y, hi.y)), 0).rgb, 0.0); +} +void main() +{ + const ivec2 tok = ivec2(gl_GlobalInvocationID.xy); + const vec2 tokCentre = (vec2(tok) + 0.5) * 4.0; + vec4 feat[4]; // the token feature seen from each sub-block + vec4 res[4]; + for (int s = 0; s < 4; ++s) + { + const ivec2 sub = ivec2(s & 1, s >> 1); + const ivec2 X = tok * 4 + sub * 2; // top-left pixel of the sub-block + // history colour of the four pixels, in linear exposed units of this frame + vec3 lin[4]; + vec2 dq[4]; + bool invalid = (p.flags & 1u) != 0u; + for (int j = 0; j < 4; ++j) + { + const vec2 pc = vec2(X + ivec2(j & 1, j >> 1)) + 0.5; + // display-resolution vectors belong to the pixel itself; render-resolution ones to its nearest texel + dq[j] = motionOf((p.flags & 4u) != 0u ? X + ivec2(j & 1, j >> 1) : nearest(pc)); + const vec2 q = pc + dq[j], qv = q * p.invValid; + invalid = invalid || qv.x < 0.0 || qv.x > 1.0 || qv.y < 0.0 || qv.y > 1.0; + const vec3 h = textureLod(histTex, q * p.invHist, 0.0).rgb; +#if LDR + lin[j] = h; +#else + lin[j] = (h * expandK(h.r)) * p.expoRatio; +#endif + } + // the current 3x3 neighbourhood around the sub-block's centre + const ivec2 c0 = nearest(vec2(X) + 1.0); + precise vec3 S = vec3(0.0); + vec3 lo = vec3(1e30), hi = vec3(-1e30); + for (int k = 0; k < 9; ++k) + { + const vec3 c = ycocg(colourAt(c0 + ivec2(k % 3 - 1, k / 3 - 1))); +#if LDR + S += c; +#else + S = fma(vec3(max(1.0 / (c.x + 1.0), p.minW)), c, S); // fused in the reference +#endif + const vec3 L = c * logK(c.x); + lo = min(lo, L); hi = max(hi, L); + } + const vec3 m = S * (1.0 / 9.0); +#if LDR + const vec3 avg = m; +#else + const vec3 avg = m / max(1.0 - m.x, p.minW); +#endif + const vec3 A = avg * logK(avg.x); + const vec3 halfr = 0.5 * (hi - lo), den = 8.0 * (hi - lo) + 1.0; + precise float mismatch = 0.0; + for (int j = 0; j < 4; ++j) + { + // The reference rounds the compressed history before subtracting it. Left to the compiler, the product + // is fused into the subtraction, and after the cancellation that leaves the mismatch one half-float + // code off in about 3% of the sub-blocks, which the recurrent network then spreads over the frame. + precise const vec3 Hl = lin[j] * logK(lin[j].x); + precise const vec3 diff = A - Hl; + const vec3 d = max(abs(diff) - halfr, 0.0) / den; + precise const float e = fma(d.z, 0.5, fma(d.x, 2.0, d.y * 0.5)); + mismatch += min(e, 1.0); + } + mismatch *= 0.25; + // the sub-block's representative pixel + const int jr = 2 * sub.x + (1 - sub.y); + const ivec2 pr = X + ivec2(1 - sub.y, sub.x); + const vec3 cr = colourAt(nearest(vec2(pr) + 0.5)); + const float Yc = 0.25 * cr.r + 0.5 * cr.g + 0.25 * cr.b; + feat[s] = invalid ? vec4(0.0) : textureLod(featTex, (tokCentre + dq[jr]) * p.tok2uv, 0.0); + const float Yh = invalid ? Yc : lin[jr].x; + if (invalid) mismatch = 0.0; + // luma channels: exposure ratio of the network's own scale (p.expoIn relative to p.expo) +#if LDR + const float tc = p.expoIn * (Yc / max(1.0 - Yc, 0.2)), th = p.expoIn * (Yh / max(1.0 - Yh, 0.2)); +#else + const float tc = p.expoIn * Yc, th = p.expoIn * Yh; +#endif + res[s] = vec4(srgb(tc * max(1.0 / (tc + 1.0), p.minW)), srgb(th * max(1.0 / (th + 1.0), p.minW)), 0.0, mismatch); + } + const vec4 fs = 0.25 * (feat[0] + feat[1] + feat[2] + feat[3]); + const uint o = uint(tok.y * p.tokW + tok.x) * 4u; + for (int s = 0; s < 4; ++s) + { + float f = fs[s], mm = res[s].w; + if (!(abs(f) < uintBitsToFloat(0x7F800000u))) f = 0.0; + if (!(abs(mm) < uintBitsToFloat(0x7F800000u))) mm = 0.0; + outb.v[o + uint(s)] = f16vec4(res[s].x, res[s].y, f, mm); + } +} diff --git a/engine/layer.comp b/engine/layer.comp new file mode 100644 index 0000000..6fd68c7 --- /dev/null +++ b/engine/layer.comp @@ -0,0 +1,426 @@ +#version 460 +// d4r engine, preset K: one Swin block (8x8-token window) written from the layer's mathematics, not from NVIDIA's +// kernel. Weights come from the engine's own model file (extract_k.py): plain row-major matrices, with the two +// per-channel gains folded into Q/K/V and fc1. +// Build-time shape: C, NH, KIND (0 plain, 1 encoder with patch merge, 2 decoder with patch expand + skip), +// POSATTN (position-only attention), EMBED (enc0: 16-channel input embedding + ReLU), HEAD (dec0: output head), +// XC (merged channels or low-resolution channels), MERGE_N (XC rounded up to 16). +// - Four waves per window, one per 16-token tile. The residual stream of a tile stays in registers (one f32 +// accumulator per 16 channels) from the input to the output. +// - Shared memory only stages what has to be read back as a matrix operand or by another wave: one buffer BF and +// the Q/K/V regions. It is kept small on purpose: the kernel waits on memory a lot, so its speed depends on how +// many windows fit on a compute unit (64 KB of shared memory each). +// - Rows are padded: with power-of-two row strides the 16 rows of an operand load land in the same LDS banks. +// - Global arrays use f16vec4 units. Shared token tiles are row-major except enc0, which uses padded +// column-major scalar-f16 arrays so cooperative-matrix stores can pack adjacent token values. +// - FOLD_GAINS=0 keeps the per-channel gains explicit instead of folding them into weights. +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_shuffle : require +#ifndef PADX +#define PADX 8 +#endif +#ifndef FOLD_GAINS +#define FOLD_GAINS 1 +#endif +#ifndef PADT +#define PADT 8 +#endif +#ifndef PADS +#define PADS 8 +#endif +#ifndef ABL_HEADS +#define ABL_HEADS uint(NH) +#endif +#ifndef ABL_CHUNKS +#define ABL_CHUNKS (UC / 8u) +#endif +layout(local_size_x = 128) in; +layout(set = 0, binding = 0, std430) readonly buffer Pb { int W; int Hh; int sx; int sy; uint o[108]; } p; +layout(set = 0, binding = 1, std430) readonly buffer Wb { f16vec4 w[]; } wb; +layout(set = 0, binding = 2, std430) readonly buffer Ib { f16vec4 a[]; } inb; +layout(set = 0, binding = 3, std430) readonly buffer Sb { f16vec4 a[]; } skb; +layout(set = 0, binding = 4, std430) buffer O0 { f16vec4 a[]; } out0; +layout(set = 0, binding = 5, std430) buffer O1 { f16vec4 a[]; } out1; +#define ACC coopmat +#define HACC coopmat +#define MA coopmat +#define MB coopmat +#define RM gl_CooperativeMatrixLayoutRowMajor +#define CMJ gl_CooperativeMatrixLayoutColumnMajor +#define LD(m, arr, off, stride) coopMatLoad(m, arr, (off) / 4u, (stride) / 4u, RM) +#define ST(m, arr, off, stride) coopMatStore(HACC(m), arr, (off) / 4u, (stride) / 4u, RM) +#define SYNC memoryBarrierShared(); barrier(); +const uint UC = uint(C), NT = uint(C) / 16u, HC = uint(C) / 2u; +const uint SX = UC + uint(PADX); // row stride of the token layout (64 rows of C) +const uint S4 = 4u * UC + 8u; // row stride of the 16 x 4C layout (patch expand / merge) +const uint SQ = 32u + uint(PADT); // row stride of Q, K, V, attention output and MLP hidden (64 rows each) +const uint RQ = 64u * SQ; // one such region +const uint SLO = uint(XC) + 8u, SEM = 24u, SHD = 56u, SMG = uint(MERGE_N) + 8u; // low-res input, embedding input, head, merged +#if POSATTN +#define NREG 1 +#define VOFF 0u +#else +#define NREG 2 +#define VOFF RQ +#endif +// TT must also hold the token layout (decoders: x0), the 16 x 4C layout (encoders: merge input) or the head's output +#if KIND == 2 +#define TT_ALT (16 * (C + PADX)) +#elif KIND == 1 +#define TT_ALT (4 * (4 * C + 8)) +#else +#define TT_ALT 0 +#endif +#if HEAD +#define TT_HEAD (16 * 56) +#else +#define TT_HEAD 0 +#endif +#define TT_REG (NREG * 16 * (32 + PADT)) +#define MAX2(a, b) ((a) > (b) ? (a) : (b)) +// BF: this tile's rows hold, in turn, the normalised input, Q, the scores / probabilities, the attention output, the +// MLP input and the block output; as a whole it is also the 16 x 4C scratch. TT: K | V shared by the four waves; its +// first region doubles as MLP hidden and scratch. +// enc0's 32-channel, position-only block benefits from column-major shared tiles. All addressing uses +// cooperative-matrix layouts rather than hardware fragment indices; other layers retain their existing layout. +#ifndef COLUMN_MAJOR_ENC0 +#define COLUMN_MAJOR_ENC0 (EMBED && C == 32) +#endif +#if COLUMN_MAJOR_ENC0 && (!EMBED || C != 32 || !POSATTN || KIND != 1 || HEAD) +#error Column-major shared tiles require the 32-channel positional enc0 block +#endif +#if COLUMN_MAJOR_ENC0 +const uint SP = 72u; // 64 token rows plus padding; patch scratch is transposed with 16 rows +shared float16_t BF[32 * 72]; +shared float16_t TT[32 * 72]; +#define LOAD_B_TOKEN(m,r,c) coopMatLoad(m,BF,(c)*SP+(r),SP,CMJ) +#define STORE_B_TOKEN(m,r,c) coopMatStore(HACC(m),BF,(c)*SP+(r),SP,CMJ) +#define LOAD_T_TOKEN(m,r,c,s) coopMatLoad(m,TT,(c)*SP+(r),SP,CMJ) +#define STORE_T_TOKEN(m,r,c,s) coopMatStore(HACC(m),TT,(c)*SP+(r),SP,CMJ) +#define STORE_B_PATCH(m,r,c,s) coopMatStore(HACC(m),BF,(c)*16u+(r),16u,CMJ) +#define LOAD_T_RAW(m,off,s) coopMatLoad(m,TT,off,s,RM) +f16vec4 readBToken(uint r,uint c) {uint i=c*SP+r;return f16vec4(BF[i],BF[i+SP],BF[i+2u*SP],BF[i+3u*SP]);} +f16vec4 readTToken(uint r,uint c,uint stride) {uint i=c*SP+r;return f16vec4(TT[i],TT[i+SP],TT[i+2u*SP],TT[i+3u*SP]);} +f16vec4 readBPatch(uint r,uint c,uint stride) {uint i=c*16u+r;return f16vec4(BF[i],BF[i+16u],BF[i+32u],BF[i+48u]);} +void writeBToken(uint r,uint c,f16vec4 v) {uint i=c*SP+r;BF[i]=v.x;BF[i+SP]=v.y;BF[i+2u*SP]=v.z;BF[i+3u*SP]=v.w;} +void writeTToken(uint r,uint c,uint stride,f16vec4 v) {uint i=c*SP+r;TT[i]=v.x;TT[i+SP]=v.y;TT[i+2u*SP]=v.z;TT[i+3u*SP]=v.w;} +void writeTRaw(uint off,f16vec4 v) {TT[off]=v.x;TT[off+1u]=v.y;TT[off+2u]=v.z;TT[off+3u]=v.w;} +#else +shared f16vec4 BF[16 * (C + PADX)]; +shared f16vec4 TT[MAX2(TT_REG, MAX2(TT_ALT, TT_HEAD))]; +#define LOAD_B_TOKEN(m,r,c) LD(m,BF,(r)*SX+(c),SX) +#define STORE_B_TOKEN(m,r,c) ST(m,BF,(r)*SX+(c),SX) +#define LOAD_T_TOKEN(m,r,c,s) LD(m,TT,(r)*(s)+(c),s) +#define STORE_T_TOKEN(m,r,c,s) ST(m,TT,(r)*(s)+(c),s) +#define STORE_B_PATCH(m,r,c,s) ST(m,BF,(r)*(s)+(c),s) +#define LOAD_T_RAW(m,off,s) LD(m,TT,off,s) +f16vec4 readBToken(uint r,uint c) {return BF[(r*SX+c)/4u];} +f16vec4 readTToken(uint r,uint c,uint stride) {return TT[(r*stride+c)/4u];} +f16vec4 readBPatch(uint r,uint c,uint stride) {return BF[(r*stride+c)/4u];} +void writeBToken(uint r,uint c,f16vec4 v) {BF[(r*SX+c)/4u]=v;} +void writeTToken(uint r,uint c,uint stride,f16vec4 v) {TT[(r*stride+c)/4u]=v;} +void writeTRaw(uint off,f16vec4 v) {TT[off/4u]=v;} +#endif +int mirror(int v, int n) { v = abs(v); return min(v, 2 * n - 2 - v); } +int tdiv2(int v) { return v < 0 ? -((-v) / 2) : v / 2; } +ACC bias_tile(uint off, uint stride) { HACC b; LD(b, wb.w, off, stride); return ACC(b); } // vectors are stored as 16 equal rows +f16vec4 softmax_exp(f16vec4 s) +{ + const f16vec4 t = clamp(s, f16vec4(uint16BitsToFloat16(uint16_t(0xCD69))), f16vec4(uint16BitsToFloat16(uint16_t(0x4D69)))); + const f16vec4 v = fma(t * t, f16vec4(uint16BitsToFloat16(uint16_t(0x91D5))), f16vec4(1.0)); + return f16vec4(exp2(vec4(fma(t, v, f16vec4(uint16BitsToFloat16(uint16_t(0xC50D))))))); +} +f16vec4 gelu(f16vec4 x) +{ + const f16vec4 c = clamp(x, f16vec4(-2.0), f16vec4(2.0)); + // the reference rounds after every binary16 operation (its compiler has fp contraction off): keep the + // same steps, or a contracted fma changes the last bit of most outputs + precise f16vec4 u = f16vec4(0.41216981) - f16vec4(0.08108133) * abs(c); + precise f16vec4 y = f16vec4(0.5) + c * u; + return x * y; +} +void main() +{ + const uint mt = gl_LocalInvocationID.x >> 5, lane = gl_LocalInvocationID.x & 31u; // wave = token tile + const uint row = lane & 15u, hf = lane >> 4, tok = 16u * mt + row; // two lanes per token + const int X0 = 8 * int(gl_WorkGroupID.x) - p.sx, Y0 = 8 * int(gl_WorkGroupID.y) - p.sy; + const int ty = Y0 + int(tok >> 3), tx = X0 + int(tok & 7u); + const uint ch0 = hf * HC; // this lane's first channel + ACC acc[C / 16]; // the residual stream of this tile + HACC residual[C / 16]; // x0, added back at the MLP boundary as one f16 add + // ------------------------------------------------------------ input: x0 into acc, x0 * rsqrt(sum x0^2) into BF +#if KIND == 2 + { + // low-resolution tokens (one per row index), channels split over the 8 lane groups of the window + const int W2 = p.W / 2, H2 = p.Hh / 2; + const uint part = uint(XC) / 8u, c0 = (mt * 2u + hf) * part; + const uint src = uint(mirror(tdiv2(Y0) + int(row >> 2), H2) * W2 + mirror(tdiv2(X0) + int(row & 3u), W2)) * uint(XC); + for (uint c = c0; c < c0 + part; c += 4u) writeTRaw(row*SLO+c,inb.a[(src+c)/4u]); + SYNC + // expand: this wave computes quarter mt of the 4C columns + for (uint nt = 0u; nt < NT; ++nt) acc[nt] = bias_tile(p.o[105] + mt * UC + nt * 16u, 4u * UC); + for (uint kt = 0u; kt < uint(XC) / 16u; ++kt) + { + MA a; LOAD_T_RAW(a,kt*16u,SLO); + MB b[C / 16]; + for (uint nt = 0u; nt < NT; ++nt) LD(b[nt], wb.w, p.o[104] + kt * 16u * 4u * UC + mt * UC + nt * 16u, 4u * UC); + for (uint nt = 0u; nt < NT; ++nt) acc[nt] = coopMatMulAdd(a, b[nt], acc[nt]); + } + SYNC // everyone has read the low-resolution tokens + for (uint nt = 0u; nt < NT; ++nt) STORE_B_PATCH(acc[nt],0u,mt*UC+nt*16u,S4); + SYNC + const uint ly = (tok >> 3) >> 1, lx = (tok & 7u) >> 1, q = ((tok >> 3) & 1u) * 2u + (tok & 1u); + const uint sk = uint(mirror(ty, p.Hh) * p.W + mirror(tx, p.W)) * UC + ch0; + for (uint c = 0u; c < HC; c += 4u) writeTToken(tok,ch0+c,SX,readBPatch(ly*4u+lx,q*UC+ch0+c,S4)+skb.a[(sk+c)/4u]); + SYNC // BF (the expanded tokens) is free + } +#define LOAD_X0(m,c) LOAD_T_TOKEN(m,mt*16u,c,SX) +#define READ_X0(r,c) readTToken(r,c,SX) +#elif EMBED + { + const uint src = uint(mirror(ty, p.Hh) * p.W + mirror(tx, p.W)) * 16u + hf * 8u; + for (uint c = 0u; c < 8u; c += 4u) writeTRaw(tok*SEM+hf*8u+c,inb.a[(src+c)/4u]); + MA a; LOAD_T_RAW(a,mt*16u*SEM,SEM); + for (uint nt = 0u; nt < NT; ++nt) + { + MB b; LD(b, wb.w, p.o[106] + nt * 16u, UC); + STORE_B_TOKEN(coopMatMulAdd(a,b,bias_tile(p.o[107]+nt*16u,UC)),mt*16u,nt*16u); + } + for (uint c = 0u; c < HC; c += 4u) writeBToken(tok,ch0+c,max(readBToken(tok,ch0+c),f16vec4(0.0))); + SYNC // TT is laid out differently from here on + } +#define LOAD_X0(m,c) LOAD_B_TOKEN(m,mt*16u,c) +#define READ_X0(r,c) readBToken(r,c) +#else + { + const uint src = uint(mirror(ty, p.Hh) * p.W + mirror(tx, p.W)) * UC + ch0; + for (uint c = 0u; c < HC; c += 4u) writeBToken(tok,ch0+c,inb.a[(src+c)/4u]); + } +#define LOAD_X0(m,c) LOAD_B_TOKEN(m,mt*16u,c) +#define READ_X0(r,c) readBToken(r,c) +#endif + { + // the residual stream starts as x0 + the output projection's bias + for (uint nt = 0u; nt < NT; ++nt) { + HACC x; LOAD_X0(x,nt*16u); + residual[nt] = x; + acc[nt] = bias_tile(p.o[1] + nt * 16u, UC); + } + vec4 s4 = vec4(0.0); + f16vec4 x[C / 8]; + for (uint c = 0u; c < HC; c += 4u) { + x[c / 4u] = READ_X0(tok,ch0+c); +#if FOLD_GAINS + const vec4 f = vec4(x[c / 4u]); s4 += f * f; +#else + s4 += vec4(x[c / 4u] * x[c / 4u]); +#endif + } + float s = s4.x + s4.y + s4.z + s4.w; + s += subgroupShuffleXor(s, 16u); + const float16_t r = float16_t(inversesqrt(float(float16_t(s)))); + for (uint c = 0u; c < HC; c += 4u) { +#if FOLD_GAINS + writeBToken(tok,ch0+c,x[c/4u]*r); +#else + writeBToken(tok,ch0+c,x[c/4u]*(r*wb.w[(p.o[0]+ch0+c)/4u])); +#endif + } +#if KIND == 2 + SYNC // x0 lived in TT; the heads overwrite it +#endif + } + MA xn[C / 16]; // the normalised input of this tile, for every head's Q, K, V + for (uint kt = 0u; kt < NT; ++kt) LOAD_B_TOKEN(xn[kt],mt*16u,kt*16u); + // ------------------------------------------------------------ attention, one head at a time + for (uint h = 0u; h < ABL_HEADS; ++h) + { + const uint otab = p.o[7u + 5u * h], owo = p.o[8u + 5u * h]; + { + // Q, K, V of this tile, two K steps at a time: their weight tiles are requested before the first product + ACC q0 = ACC(0.0), q1 = ACC(0.0), k0 = ACC(0.0), k1 = ACC(0.0), v0 = ACC(0.0), v1 = ACC(0.0); + for (uint kt = 0u; kt < NT; kt += 2u) + { + MB bv[4]; + for (uint i = 0u; i < 2u; ++i) + { + LD(bv[2u * i], wb.w, p.o[6u + 5u * h] + (kt + i) * 512u, 32u); + LD(bv[2u * i + 1u], wb.w, p.o[6u + 5u * h] + (kt + i) * 512u + 16u, 32u); + } +#if !POSATTN + MB bq[4], bk[4]; + for (uint i = 0u; i < 2u; ++i) + { + LD(bq[2u * i], wb.w, p.o[4u + 5u * h] + (kt + i) * 512u, 32u); + LD(bq[2u * i + 1u], wb.w, p.o[4u + 5u * h] + (kt + i) * 512u + 16u, 32u); + LD(bk[2u * i], wb.w, p.o[5u + 5u * h] + (kt + i) * 512u, 32u); + LD(bk[2u * i + 1u], wb.w, p.o[5u + 5u * h] + (kt + i) * 512u + 16u, 32u); + } + for (uint i = 0u; i < 2u; ++i) + { + q0 = coopMatMulAdd(xn[kt + i], bq[2u * i], q0); q1 = coopMatMulAdd(xn[kt + i], bq[2u * i + 1u], q1); + k0 = coopMatMulAdd(xn[kt + i], bk[2u * i], k0); k1 = coopMatMulAdd(xn[kt + i], bk[2u * i + 1u], k1); + } +#endif + for (uint i = 0u; i < 2u; ++i) { v0 = coopMatMulAdd(xn[kt + i], bv[2u * i], v0); v1 = coopMatMulAdd(xn[kt + i], bv[2u * i + 1u], v1); } + } +#if !POSATTN + STORE_B_TOKEN(q0,mt*16u,0u); STORE_B_TOKEN(q1,mt*16u,16u); // Q: this wave's own rows + STORE_T_TOKEN(k0,mt*16u,0u,SQ); STORE_T_TOKEN(k1,mt*16u,16u,SQ); +#endif + STORE_T_TOKEN(v0,VOFF/SQ+mt*16u,0u,SQ); STORE_T_TOKEN(v1,VOFF/SQ+mt*16u,16u,SQ); + } + SYNC +#if !POSATTN + { + ACC s[4]; + for (uint nt = 0u; nt < 4u; ++nt) { HACC tb; LD(tb, wb.w, otab + mt * 1024u + nt * 16u, 64u); s[nt] = ACC(tb); } + MA qa[2]; LOAD_B_TOKEN(qa[0],mt*16u,0u); LOAD_B_TOKEN(qa[1],mt*16u,16u); + for (uint kt = 0u; kt < 2u; ++kt) + { + MB b[4]; + for (uint nt = 0u; nt < 4u; ++nt) coopMatLoad(b[nt], TT, (nt * 16u * SQ + kt * 16u) / 4u, SQ / 4u, CMJ); // K transposed + for (uint nt = 0u; nt < 4u; ++nt) s[nt] = coopMatMulAdd(qa[kt], b[nt], s[nt]); + } + for (uint nt = 0u; nt < 4u; ++nt) STORE_B_TOKEN(s[nt],mt*16u,nt*16u); // over Q, which is in registers + // softmax over the token's 64 scores: 32 per lane + f16vec4 e[8]; + vec4 sum = vec4(0.0); + for (uint j = 0u; j < 8u; ++j) { e[j] = softmax_exp(readBToken(tok,hf*32u+j*4u)); sum += vec4(e[j]); } + float tot = sum.x + sum.y + sum.z + sum.w; + tot += subgroupShuffleXor(tot, 16u); + const float16_t rc = float16_t(1.0 / float(float16_t(tot))); + for (uint j = 0u; j < 8u; ++j) writeBToken(tok,hf*32u+j*4u,e[j]*rc); + } +#endif + { + ACC o0 = ACC(0.0), o1 = ACC(0.0); + // two K steps at a time: few live operands keep the register count, and with it the occupancy, in check + for (uint k2 = 0u; k2 < 4u; k2 += 2u) + { + MA a[2]; MB b[4]; + for (uint i = 0u; i < 2u; ++i) + { + const uint kt = k2 + i; +#if POSATTN + LD(a[i], wb.w, otab + mt * 1024u + kt * 16u, 64u); +#else + LOAD_B_TOKEN(a[i],mt*16u,kt*16u); +#endif + LOAD_T_TOKEN(b[2u*i],VOFF/SQ+kt*16u,0u,SQ); LOAD_T_TOKEN(b[2u*i+1u],VOFF/SQ+kt*16u,16u,SQ); + } + for (uint i = 0u; i < 2u; ++i) { o0 = coopMatMulAdd(a[i], b[2u * i], o0); o1 = coopMatMulAdd(a[i], b[2u * i + 1u], o1); } + } + // the two K tiles of this head for the output projection: through memory, to become A operands + STORE_B_TOKEN(o0,mt*16u,0u); STORE_B_TOKEN(o1,mt*16u,16u); + MA a0, a1; LOAD_B_TOKEN(a0,mt*16u,0u); LOAD_B_TOKEN(a1,mt*16u,16u); + for (uint n0 = 0u; n0 < NT; n0 += 2u) + { + MB b0[2], b1[2]; + for (uint i = 0u; i < 2u; ++i) { LD(b0[i], wb.w, owo + (n0 + i) * 16u, UC); LD(b1[i], wb.w, owo + 16u * UC + (n0 + i) * 16u, UC); } + for (uint i = 0u; i < 2u; ++i) { + acc[n0 + i] = coopMatMulAdd(a1, b1[i], coopMatMulAdd(a0, b0[i], acc[n0 + i])); + } + } + } + SYNC // every wave is done with K and V before the next head overwrites them + } + // ------------------------------------------------------------ MLP: its input is the stream itself (gain folded into fc1) + { + for (uint nt = 0u; nt < NT; ++nt) acc[nt] = ACC(HACC(acc[nt]) + residual[nt]); + for (uint nt = 0u; nt < NT; ++nt) STORE_B_TOKEN(acc[nt],mt*16u,nt*16u); +#if !FOLD_GAINS + for (uint c = 0u; c < HC; c += 4u) writeBToken(tok,ch0+c,readBToken(tok,ch0+c)*wb.w[(p.o[2]+ch0+c)/4u]); +#endif + MA m[C / 16]; + for (uint kt = 0u; kt < NT; ++kt) LOAD_B_TOKEN(m[kt],mt*16u,kt*16u); + for (uint nt = 0u; nt < NT; ++nt) { + acc[nt] = ACC(HACC(acc[nt]) + HACC(bias_tile(p.o[3] + nt * 16u, UC))); + } + for (uint ch = 0u; ch < ABL_CHUNKS; ++ch) + { + const uint ow1 = p.o[44u + 3u * ch], ob1 = p.o[45u + 3u * ch], ow2 = p.o[46u + 3u * ch]; + ACC h0 = bias_tile(ob1, 32u), h1 = bias_tile(ob1 + 16u, 32u); + for (uint kt = 0u; kt < NT; kt += 2u) + { + MB b[4]; + for (uint i = 0u; i < 2u; ++i) { LD(b[2u * i], wb.w, ow1 + (kt + i) * 512u, 32u); LD(b[2u * i + 1u], wb.w, ow1 + (kt + i) * 512u + 16u, 32u); } + for (uint i = 0u; i < 2u; ++i) { h0 = coopMatMulAdd(m[kt + i], b[2u * i], h0); h1 = coopMatMulAdd(m[kt + i], b[2u * i + 1u], h1); } + } + STORE_T_TOKEN(h0,mt*16u,0u,SQ); STORE_T_TOKEN(h1,mt*16u,16u,SQ); + // fc2's first weight tiles are requested before the activation pass, so they arrive while it runs + MB w2a[2], w2b[2]; + for (uint i = 0u; i < 2u; ++i) { LD(w2a[i], wb.w, ow2 + i * 16u, UC); LD(w2b[i], wb.w, ow2 + 16u * UC + i * 16u, UC); } + for (uint j = 0u; j < 4u; ++j) writeTToken(tok,hf*16u+j*4u,SQ,gelu(readTToken(tok,hf*16u+j*4u,SQ))); + MA a0, a1; LOAD_T_TOKEN(a0,mt*16u,0u,SQ); LOAD_T_TOKEN(a1,mt*16u,16u,SQ); + for (uint i = 0u; i < 2u; ++i) acc[i] = coopMatMulAdd(a1, w2b[i], coopMatMulAdd(a0, w2a[i], acc[i])); + for (uint n0 = 2u; n0 < NT; n0 += 2u) + { + for (uint i = 0u; i < 2u; ++i) { LD(w2a[i], wb.w, ow2 + (n0 + i) * 16u, UC); LD(w2b[i], wb.w, ow2 + 16u * UC + (n0 + i) * 16u, UC); } + for (uint i = 0u; i < 2u; ++i) acc[n0 + i] = coopMatMulAdd(a1, w2b[i], coopMatMulAdd(a0, w2a[i], acc[n0 + i])); + } + } + for (uint nt = 0u; nt < NT; ++nt) STORE_B_TOKEN(acc[nt],mt*16u,nt*16u); + } + // ------------------------------------------------------------ outputs (the block output is in BF, token layout) + const bool inside = ty >= 0 && ty < p.Hh && tx >= 0 && tx < p.W; + const uint dst = uint(ty * p.W + tx); +#if HEAD + SYNC // the head's rows of TT overlap other tiles' MLP rows + { + ACC ho[3]; + for (uint nt = 0u; nt < 3u; ++nt) ho[nt] = bias_tile(p.o[107] + nt * 16u, 48u); + for (uint kt = 0u; kt < NT; ++kt) + { + MA a; LOAD_B_TOKEN(a,mt*16u,kt*16u); + MB b[3]; + for (uint nt = 0u; nt < 3u; ++nt) LD(b[nt], wb.w, p.o[106] + kt * 16u * 48u + nt * 16u, 48u); + for (uint nt = 0u; nt < 3u; ++nt) ho[nt] = coopMatMulAdd(a, b[nt], ho[nt]); + } + for (uint nt = 0u; nt < 3u; ++nt) STORE_T_TOKEN(ho[nt],mt*16u,nt*16u,SHD); + if (inside) + for (uint c = hf * 20u; c < hf * 20u + 20u; c += 4u) out0.a[(dst * 40u + c) / 4u] = readTToken(tok,c,SHD); + } +#elif KIND == 1 + if (inside) + for (uint c = 0u; c < HC; c += 4u) out1.a[(dst * UC + ch0 + c) / 4u] = readBToken(tok,ch0+c); +#else + if (inside) + for (uint c = 0u; c < HC; c += 4u) out0.a[(dst * UC + ch0 + c) / 4u] = readBToken(tok,ch0+c); +#endif +#if KIND == 1 + SYNC + { + // patch merge: wave mt gathers sub-token mt of the 16 merged tokens into TT, then computes output tiles mt, mt + 4, ... + const uint sourceToken = 8u * (2u * (row >> 2) + (mt >> 1)) + 2u * (row & 3u) + (mt & 1u); + for (uint c = 0u; c < HC; c += 4u) + writeTRaw(row * S4 + mt * UC + ch0 + c, readBToken(sourceToken, ch0 + c)); + SYNC + const uint MN = uint(MERGE_N), NI = (MN / 16u + 3u) / 4u, ntn = (MN / 16u + 3u - mt) / 4u; + ACC mg[(MERGE_N / 16 + 3) / 4]; + for (uint i = 0u; i < NI; ++i) mg[i] = bias_tile(p.o[105] + min(mt + 4u * i, MN / 16u - 1u) * 16u, MN); + for (uint kt = 0u; kt < 4u * UC / 16u; kt += 2u) + { + MA a[2]; MB b[2 * ((MERGE_N / 16 + 3) / 4)]; + for (uint j = 0u; j < 2u; ++j) + { + LOAD_T_RAW(a[j],(kt+j)*16u,S4); + for (uint i = 0u; i < NI; ++i) LD(b[j * NI + i], wb.w, p.o[104] + (kt + j) * 16u * MN + min(mt + 4u * i, MN / 16u - 1u) * 16u, MN); + } + for (uint j = 0u; j < 2u; ++j) + for (uint i = 0u; i < NI; ++i) mg[i] = coopMatMulAdd(a[j], b[j * NI + i], mg[i]); + } + // BF is free (every wave gathered before the barrier above): the merged tokens go there + for (uint i = 0u; i < NI; ++i) + if (i < ntn) STORE_B_PATCH(mg[i],0u,(mt+4u*i)*16u,SMG); + SYNC + const int ym = tdiv2(Y0) + int(row >> 2), xm = tdiv2(X0) + int(row & 3u); + const uint part = uint(XC) / 8u, c0 = (mt * 2u + hf) * part; + if (ym >= 0 && ym < p.Hh / 2 && xm >= 0 && xm < p.W / 2) + for (uint c = c0; c < c0 + part; c += 4u) out0.a[(uint(ym * (p.W / 2) + xm) * uint(XC) + c) / 4u] = readBPatch(row,c,SMG); + } +#endif +} diff --git a/engine/layer_m.comp b/engine/layer_m.comp new file mode 100644 index 0000000..da4d608 --- /dev/null +++ b/engine/layer_m.comp @@ -0,0 +1,446 @@ +#version 460 +// d4r engine, preset M (rrlite): one Swin block over 4x4-token windows, register-resident. +// +// The cooperative matrix is opaque in GLSL, but RADV's gfx1201 lowering has a fixed lane/slot <-> element +// mapping, the same for every 16x16x16 use. It was proven on this GPU with the probes in +// build/engine-quality/rm-probe (each lane writes its own slot values into an operand and the resulting image is +// read back): +// +// RowMajor sources and destinations are indexed [row][col] +// ColumnMajor sources and destinations are indexed [col][row] +// accumulator and B operand: element i of lane l is (row = 8*(l>>4) + i, col = l & 15) +// A operand: element i of lane l is (row = l & 15, col = 8*(l>>4) + i) +// +// Every product here is therefore weights x activations (A = weights, B = activations): a product's rows are the +// output channels and its columns the tokens, and the per-lane elements of a result ARE the per-lane elements of +// the next product's B operand. Cooperative-matrix subscripts read and write those elements in place, so the +// attention and the MLP chain never leave the registers. Only the patch expand (decoders) and the patch merge +// (encoders) move data between waves, and only those use shared memory and a workgroup barrier. +// +// NVIDIA's arithmetic (the semantics of kernels/m/swin_common.h K32_MIX, which the captured launches agree with) +// rounds the accumulator to binary16 after every logical 32 products; the 16x16x16 cells have no K32 form, so a +// logical k32 is two K16 steps summed in f32 and rounded once. RADV folds a register-only f32 -> f16 -> f32 +// conversion away, so the rounding of P + the previous accumulator passes through a per-invocation binary16 +// scratch slot (no cross-lane traffic: a subgroup memory barrier is enough). The tile loops are written out one +// tile at a time because a cooperative matrix indexed by a loop variable is not a register any more. +// +// The two norms do NOT use the natural element order of this shader's accumulators: the native register element +// (tile T, half hf, element I) holds the natural channel tok_xnat(T, 8hf+I) = 16T + 8*((I>>2)&1) + 4hf + (I&3), +// and tok_rsq's summation tree groups the lanes' squares in that order. Norm 1 therefore reads the input at that +// channel (XVALN, the absolute one, not the lane's own half) and norm 2 rebuilds the native order from the +// accumulator registers with a pair-aligned half-wave exchange. Without this the ten layers reach only 41-46 dB +// against the capture even though every product, weight map and bias is already the native one. +// +// Build-time shape: C, NH, KIND (0 plain, 1 encoder with 2x2 patch merge, 2 decoder with patch expand + skip), +// XC (merged channels or low-resolution channels). One workgroup is an 8x8-token block: four waves, one per +// 4x4-token window. Bindings and the parameter block are unchanged; the FP8 weight matrices are the packed +// fixed-16 tiles of test_m.pack8 (the f16 vectors stay row-major). +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_float_e4m3 : require +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_shuffle : require +layout(local_size_x = 128) in; +// o: 0 g1, 1 bo, 2 g2, 3 b2 (f16); per head h at 4 + 4h: q, v (FP8), bias (f16), wo (FP8); +// per MLP chunk c at 20 + 3c: w1 (FP8), b1 (f16), w2 (FP8); 68, 69: merge or expand weights (FP8) and bias (f16) +layout(set = 0, binding = 0, std430) readonly buffer Pb { int W; int Hh; int sx; int sy; uint o[72]; } p; +layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8; +layout(set = 0, binding = 2, std430) readonly buffer W16 { f16vec4 w[]; } w16; +layout(set = 0, binding = 3, std430) readonly buffer Ib { floate4m3_t a[]; } inb; // [H][W][C]; decoders: [H/2][W/2][XC] +layout(set = 0, binding = 4, std430) readonly buffer Sb { floate4m3_t a[]; } skb; // decoders: skip [H][W][C] +layout(set = 0, binding = 5, std430) buffer O0 { floate4m3_t a[]; } out0; // [H][W][C] +layout(set = 0, binding = 6, std430) buffer O1 { floate4m3_t a[]; } out1; // encoders: merged [H/2][W/2][XC] +#define ACC coopmat +#define HAC coopmat +#define A8 coopmat +#define B8 coopmat +#define A16 coopmat +#define B16 coopmat +#define RM gl_CooperativeMatrixLayoutRowMajor +#define CM gl_CooperativeMatrixLayoutColumnMajor +const uint UC = uint(C), NT = uint(C) / 16u, UXC = uint(XC); +const uint SLO = UXC + 8u; // row stride of the low-resolution token image +const uint S4 = 4u * UC + 8u; // row stride of the 4C rows (patch expand / merge) +const uint VWD = 16u; // row width of the per-wave V transpose buffer, in f16 +int mirror(int v, int n) { v = abs(v); return min(v, 2 * n - 2 - v); } +#define Q16(v) float16_t(v) +#define CL8(v) clamp(v, float16_t(-448.0), float16_t(448.0)) +// NVIDIA's accumulator rounding: one binary16 rounding of the whole tile after every logical 32 products. +// RADV folds a register-only f32->f16->f32 conversion away (measured: half of all values and every tie +// midpoint unchanged), so the rounding has to pass through binary16 memory. Each invocation owns 8 halves of +// its own scratch: no cross-lane traffic, so a subgroup fence is enough (no workgroup barrier). +shared f16vec4 RS[128 * 2]; +#define R16(m) { \ + RS[gl_LocalInvocationID.x * 2u + 0u] = f16vec4(float16_t((m)[0]), float16_t((m)[1]), float16_t((m)[2]), float16_t((m)[3])); \ + RS[gl_LocalInvocationID.x * 2u + 1u] = f16vec4(float16_t((m)[4]), float16_t((m)[5]), float16_t((m)[6]), float16_t((m)[7])); \ + subgroupMemoryBarrierShared(); \ + const f16vec4 r0_ = RS[gl_LocalInvocationID.x * 2u + 0u], r1_ = RS[gl_LocalInvocationID.x * 2u + 1u]; \ + (m)[0] = float(r0_.x); (m)[1] = float(r0_.y); (m)[2] = float(r0_.z); (m)[3] = float(r0_.w); \ + (m)[4] = float(r1_.x); (m)[5] = float(r1_.y); (m)[6] = float(r1_.z); (m)[7] = float(r1_.w); } +// one logical k32, exactly as the native K32_MIX: P = A0*B0 + A1*B1 from a zero seed, then one binary16 +// rounding of P + the previous accumulator (the native v_fma_mix), through the scratch so it is not folded. +// Packed fixed-16 weight tiles (test_m.pack8): matrix W[K][N] as [K/16][N/16][N16][K16], so one 16x16 A tile +// is 256 contiguous bytes at base + 256*(kTile*(N/16) + nOffset/16), read RowMajor with stride 16. Same A +// operand values as the old plain W[K][N] ColumnMajor stride-N load, but a plain (non-transposed) load. +#define LDA8(dst, base, ks, stride, moff) coopMatLoad((dst), w8.w, (base) + 256u * (uint(ks) * (uint(stride) / 16u) + (uint(moff) / 16u)), 16u, RM) + +#define K32H(dst, a0, b0, a1, b1) { \ + ACC P_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(0.0))); \ + RS[gl_LocalInvocationID.x * 2u + 0u] = f16vec4( \ + float16_t(P_[0] + float((dst)[0])), float16_t(P_[1] + float((dst)[1])), \ + float16_t(P_[2] + float((dst)[2])), float16_t(P_[3] + float((dst)[3]))); \ + RS[gl_LocalInvocationID.x * 2u + 1u] = f16vec4( \ + float16_t(P_[4] + float((dst)[4])), float16_t(P_[5] + float((dst)[5])), \ + float16_t(P_[6] + float((dst)[6])), float16_t(P_[7] + float((dst)[7]))); \ + subgroupMemoryBarrierShared(); \ + const f16vec4 u0_ = RS[gl_LocalInvocationID.x * 2u + 0u], u1_ = RS[gl_LocalInvocationID.x * 2u + 1u]; \ + (dst)[0] = u0_.x; (dst)[1] = u0_.y; (dst)[2] = u0_.z; (dst)[3] = u0_.w; \ + (dst)[4] = u1_.x; (dst)[5] = u1_.y; (dst)[6] = u1_.z; (dst)[7] = u1_.w; } +// a tiled f16 vector (16 equal rows) as the initial value of an accumulator whose rows are its channels +HAC bias16(uint off, uint stride) { HAC b; coopMatLoad(b, w16.w, off / 4u, stride / 4u, CM); return b; } +#define MK8_F8(dst, a) { \ + const float16_t q0_ = CL8(Q16((a)[0])), q1_ = CL8(Q16((a)[1])), q2_ = CL8(Q16((a)[2])), q3_ = CL8(Q16((a)[3])); \ + const float16_t q4_ = CL8(Q16((a)[4])), q5_ = CL8(Q16((a)[5])), q6_ = CL8(Q16((a)[6])), q7_ = CL8(Q16((a)[7])); \ + (dst)[0] = floate4m3_t(q0_); (dst)[1] = floate4m3_t(q1_); (dst)[2] = floate4m3_t(q2_); (dst)[3] = floate4m3_t(q3_); \ + (dst)[4] = floate4m3_t(q4_); (dst)[5] = floate4m3_t(q5_); (dst)[6] = floate4m3_t(q6_); (dst)[7] = floate4m3_t(q7_); } +// NVIDIA's softmax numerator, two neighbouring columns in one 32-bit word (the even column's carry reaches the odd) +void softmax2(float16_t s0, float16_t s1, out float16_t e0, out float16_t e1) +{ + const float16_t lim = float16_t(0.55615234375); + const float16_t cs = float16_t(0.0972222164273262 / 5.656854152679443); + const float16_t t0 = clamp(s0 * cs, -lim, lim), t1 = clamp(s1 * cs, -lim, lim); + const float16_t p0 = fma(t0, fma(t0, -t0, float16_t(0.92730712890625)), float16_t(1.375)); + const float16_t p1 = fma(t1, fma(t1, -t1, float16_t(0.92730712890625)), float16_t(1.375)); + const uint w = ((uint(float16BitsToUint16(p0)) | (uint(float16BitsToUint16(p1)) << 16)) << 5) + 0x7FF88000u; + e0 = uint16BitsToFloat16(uint16_t(w & 0xFFFFu)); e1 = uint16BitsToFloat16(uint16_t(w >> 16)); +} +float16_t gelu1(float16_t x) +{ + const float16_t c = clamp(x, float16_t(-2.0), float16_t(2.0)); + const float16_t ac = abs(c); + precise float16_t inner = float16_t(0.41216981) - float16_t(0.08108133) * ac; + precise float16_t outer = float16_t(0.5) + c * inner; + return x * outer; +} +#if C == 64 +#define EACH_T(X) X(0) X(1) X(2) X(3) +#define EACH_T16(X) X(0, 1) X(2, 3) +#elif C == 96 +#define EACH_T(X) X(0) X(1) X(2) X(3) X(4) X(5) +#define EACH_T16(X) X(0, 1) X(2, 3) X(4, 5) +#else +#define EACH_T(X) X(0) X(1) X(2) X(3) X(4) X(5) X(6) X(7) +#define EACH_T16(X) X(0, 1) X(2, 3) X(4, 5) X(6, 7) +#endif + +// ------------------------------------------------------------------ native arithmetic order +// The accuracy build (kernels/build.sh D4R_PREFER_ACCURACY=1 -> -DD4R_ACCURACY -> SWIN_EXACT, and +// D4R_NATIVE_FP8=1 -> -DD4R_FP8_WMMA -> SWIN_A8/K32_MIX) runs the token-lane block of kernels/m/swin_tok.h. +// Its two non-obvious orderings are reproduced here in binary16: +// * tok_rsq: per tile pp = a*a + b*b over the pairs (0,1)-(2,3) and (4,5)-(6,7), then a tree that splits the +// tiles by parity, then a + other(a) across the half-waves, then u[0] + u[1], then the two components. +// * the softmax: v1[k] = wg[k] + other(wg[k]) elementwise first, then (v1[0]+v1[1]) + (v1[2]+v1[3]). +// f16 arithmetic is marked precise so nothing is contracted into an FMA (the native kernel builds with +// #pragma clang fp contract(off)). +#define SH16(v) subgroupShuffleXor(packHalf2x16(vec2(float((v).x), float((v).y))), 16u) +#define UN16(p) f16vec2(float16_t(unpackHalf2x16(p).x), float16_t(unpackHalf2x16(p).y)) +#if C == 64 +#define RSQ_TREE(OUT) { \ + f16vec2 P0_[NT], P1_[NT]; \ + EACH_T(RSQ_TILE) \ + precise f16vec2 W00_ = P0_[0] + P0_[2], W01_ = P1_[0] + P1_[2]; \ + precise f16vec2 W10_ = P0_[1] + P0_[3], W11_ = P1_[1] + P1_[3]; \ + precise f16vec2 A0_ = W00_ + W10_, A1_ = W01_ + W11_; \ + precise f16vec2 U0_ = A0_ + UN16(SH16(A0_)), U1_ = A1_ + UN16(SH16(A1_)); \ + precise f16vec2 S_ = U0_ + U1_; \ + OUT = float16_t(1.0 / sqrt(float(float16_t((S_.x + S_.y) + float16_t(0.0001220703125))))); } +#elif C == 96 +#define RSQ_TREE(OUT) { \ + f16vec2 P0_[NT], P1_[NT]; \ + EACH_T(RSQ_TILE) \ + precise f16vec2 W00_ = (P0_[0] + P0_[2]) + P0_[4], W01_ = (P1_[0] + P1_[2]) + P1_[4]; \ + precise f16vec2 W10_ = (P0_[1] + P0_[3]) + P0_[5], W11_ = (P1_[1] + P1_[3]) + P1_[5]; \ + precise f16vec2 A0_ = W00_ + W10_, A1_ = W01_ + W11_; \ + precise f16vec2 U0_ = A0_ + UN16(SH16(A0_)), U1_ = A1_ + UN16(SH16(A1_)); \ + precise f16vec2 S_ = U0_ + U1_; \ + OUT = float16_t(1.0 / sqrt(float(float16_t((S_.x + S_.y) + float16_t(0.0001220703125))))); } +#else +#define RSQ_TREE(OUT) { \ + f16vec2 P0_[NT], P1_[NT]; \ + EACH_T(RSQ_TILE) \ + precise f16vec2 W00_ = (P0_[0] + P0_[2]) + (P0_[4] + P0_[6]), W01_ = (P1_[0] + P1_[2]) + (P1_[4] + P1_[6]); \ + precise f16vec2 W10_ = (P0_[1] + P0_[3]) + (P0_[5] + P0_[7]), W11_ = (P1_[1] + P1_[3]) + (P1_[5] + P1_[7]); \ + precise f16vec2 A0_ = W00_ + W10_, A1_ = W01_ + W11_; \ + precise f16vec2 U0_ = A0_ + UN16(SH16(A0_)), U1_ = A1_ + UN16(SH16(A1_)); \ + precise f16vec2 S_ = U0_ + U1_; \ + OUT = float16_t(1.0 / sqrt(float(float16_t((S_.x + S_.y) + float16_t(0.0001220703125))))); } +#endif +#define RSQ_TILE(T) { \ + precise f16vec2 A_ = f16vec2(EV(T, 0), EV(T, 1)), B_ = f16vec2(EV(T, 2), EV(T, 3)); \ + precise f16vec2 C_ = f16vec2(EV(T, 4), EV(T, 5)), D_ = f16vec2(EV(T, 6), EV(T, 7)); \ + P0_[T] = (A_ * A_) + (B_ * B_); P1_[T] = (C_ * C_) + (D_ * D_); } + +shared f16vec4 VSH[4 * 2 * 16 * VWD / 4]; // per wave and head: the V tiles, transposed through shared memory +#if KIND == 1 +shared floate4m3_t NB[32 * S4]; // 16 merged rows, preceded by 16 rows of slack for the gather +#endif +#if KIND == 2 +shared floate4m3_t NLI[16 * SLO]; // the block's 16 low-resolution tokens +shared floate4m3_t NEX[16 * S4]; // their expanded 4C columns +#endif +#define SYNC memoryBarrierShared(); barrier(); +void main() +{ + const uint mt = gl_LocalInvocationID.x >> 5, lane = gl_LocalInvocationID.x & 31u; + const uint row = lane & 15u, hf = lane >> 4, ch8 = 8u * hf; // one token per lane, its channels 8h..8h+7 + const int X0 = 8 * int(gl_WorkGroupID.x) - p.sx, Y0 = 8 * int(gl_WorkGroupID.y) - p.sy; + const uint bxr = 4u * (mt & 1u) + (row & 3u), byr = 4u * (mt >> 1) + (row >> 2); + const int tx = X0 + int(bxr), ty = Y0 + int(byr); + const bool inside = ty >= 0 && ty < p.Hh && tx >= 0 && tx < p.W; + const uint iy = uint(mirror(ty, p.Hh)), ix = uint(mirror(tx, p.W)); + HAC res[NT]; // the residual stream: rows are channels, columns are tokens + B8 h1[NT]; // the normalised activation, the K operand of every product + // ------------------------------------------------------------ patch expand: x0 = q8(E) + skip +#if KIND == 2 + { + const uint W2 = uint(p.W) / 2u, H2 = uint(p.Hh) / 2u, part = UXC / 8u, c0 = (mt * 2u + hf) * part; + const uint lo = uint(mirror((Y0 >> 1) + int(row >> 2), int(H2)) * int(W2) + mirror((X0 >> 1) + int(row & 3u), p.W / 2)) * UXC; + for (uint c = c0; c < c0 + part; ++c) NLI[row * SLO + c] = inb.a[lo + c]; + SYNC + // this wave computes sub-position mt of every low-resolution token's 4C columns + HAC ef[NT]; +#define EF_INIT(N) ef[N] = bias16(p.o[69] + 16u * (mt * NT + (N)), 4u * UC); +#define EF_K32(N) { A8 xa0, xa1; \ + LDA8(xa0, p.o[68], kt, 4u * UC, 16u * (mt * NT + (N))); \ + LDA8(xa1, p.o[68], kt + 1u, 4u * UC, 16u * (mt * NT + (N))); \ + K32H(ef[N], xa0, b0, xa1, b1) } +#define EF_STORE(N) { A8 ey; MK8_F8(ey, ef[N]); coopMatStore(ey, NEX, mt * UC + 16u * (N), S4, RM); } + EACH_T(EF_INIT) + for (uint kt = 0u; kt < UXC / 16u; kt += 2u) + { + B8 b0, b1; + coopMatLoad(b0, NLI, kt * 16u, SLO, CM); + coopMatLoad(b1, NLI, (kt + 1u) * 16u, SLO, CM); + EACH_T(EF_K32) + } + EACH_T(EF_STORE) + SYNC + } +#endif + // ------------------------------------------------------------ norm 1: sum of squares, residual and h1 +#if KIND == 2 + const uint lb0 = ((byr >> 1) * 4u + (bxr >> 1)) * S4 + ((byr & 1u) * 2u + (bxr & 1u)) * UC + ch8; + const uint sb0 = (iy * uint(p.W) + ix) * UC + ch8; +#define XVAL(C) (float16_t(NEX[lb0 + (C)]) + float16_t(skb.a[sb0 + (C)])) +#else + const uint ib0 = (iy * uint(p.W) + ix) * UC + ch8; +#define XVAL(C) float16_t(inb.a[ib0 + (C)]) +#endif +// The native's register element (tile T, half hf, element I) holds the natural channel tok_xnat(T, 8hf+I) = +// 16T + 8*((I>>2)&1) + 4hf + (I&3) (kernels/m/swin_tok.h tok_rsq on the register layout the input plane implies). +// This shader's own accumulators and A rows keep natural channels, so h1 and the residual use XVAL (the lane's +// own half); the norm tree, however, groups per the native element order, so it reads the absolute channel. +#if KIND == 2 +#define XVALN(C) (float16_t(NEX[(lb0 - ch8) + (C)]) + float16_t(skb.a[(sb0 - ch8) + (C)])) +#else +#define XVALN(C) float16_t(inb.a[(iy * uint(p.W) + ix) * UC + (C)]) +#endif +#define EV(T, I) XVALN(16u * (T) + 8u * ((uint(I) >> 2) & 1u) + 4u * hf + (uint(I) & 3u)) + float16_t rn; RSQ_TREE(rn) +#undef EV +#define NM1(N) { \ + const uint wb = (p.o[0] + 16u * (N) + ch8) / 4u, bb = (p.o[1] + 16u * (N) + ch8) / 4u; \ + const f16vec4 g0 = w16.w[wb], g1 = w16.w[wb + 1u], d0 = w16.w[bb], d1 = w16.w[bb + 1u]; \ + const float16_t x0 = XVAL(16u * (N) + 0u), x1 = XVAL(16u * (N) + 1u), x2 = XVAL(16u * (N) + 2u), x3 = XVAL(16u * (N) + 3u); \ + const float16_t x4 = XVAL(16u * (N) + 4u), x5 = XVAL(16u * (N) + 5u), x6 = XVAL(16u * (N) + 6u), x7 = XVAL(16u * (N) + 7u); \ + res[N][0] = float16_t(x0 + d0.x); res[N][1] = float16_t(x1 + d0.y); \ + res[N][2] = float16_t(x2 + d0.z); res[N][3] = float16_t(x3 + d0.w); \ + res[N][4] = float16_t(x4 + d1.x); res[N][5] = float16_t(x5 + d1.y); \ + res[N][6] = float16_t(x6 + d1.z); res[N][7] = float16_t(x7 + d1.w); \ + h1[N][0] = floate4m3_t(CL8(x0 * float16_t(rn * g0.x))); h1[N][1] = floate4m3_t(CL8(x1 * float16_t(rn * g0.y))); \ + h1[N][2] = floate4m3_t(CL8(x2 * float16_t(rn * g0.z))); h1[N][3] = floate4m3_t(CL8(x3 * float16_t(rn * g0.w))); \ + h1[N][4] = floate4m3_t(CL8(x4 * float16_t(rn * g1.x))); h1[N][5] = floate4m3_t(CL8(x5 * float16_t(rn * g1.y))); \ + h1[N][6] = floate4m3_t(CL8(x6 * float16_t(rn * g1.z))); h1[N][7] = floate4m3_t(CL8(x7 * float16_t(rn * g1.w))); } + EACH_T(NM1) + // ------------------------------------------------------------ attention, one head at a time + for (uint hd = 0u; hd < uint(NH); ++hd) + { + const uint oq = p.o[4u + 4u * hd], ov = p.o[5u + 4u * hd], ob = p.o[6u + 4u * hd], ow = p.o[7u + 4u * hd]; + HAC q[2], v[2]; + q[0] = HAC(0.0); q[1] = HAC(0.0); v[0] = HAC(0.0); v[1] = HAC(0.0); +#define QV_G(A, B) { \ + const uint kt = A; \ + A8 aq0, aq1, av0, av1; \ + LDA8(aq0, oq, kt, 32u, 0u); \ + LDA8(aq1, oq, kt + 1u, 32u, 0u); \ + LDA8(av0, ov, kt, 32u, 0u); \ + LDA8(av1, ov, kt + 1u, 32u, 0u); \ + K32H(q[0], aq0, h1[A], aq1, h1[B]) \ + K32H(v[0], av0, h1[A], av1, h1[B]) \ + LDA8(aq0, oq, kt, 32u, 16u); \ + LDA8(aq1, oq, kt + 1u, 32u, 16u); \ + LDA8(av0, ov, kt, 32u, 16u); \ + LDA8(av1, ov, kt + 1u, 32u, 16u); \ + K32H(q[1], aq0, h1[A], aq1, h1[B]) \ + K32H(v[1], av0, h1[A], av1, h1[B]) } + EACH_T16(QV_G) + // scores: q8(Q) against itself, rows are the queries and columns the keys, plus the relative position bias + A8 qa[2]; B8 qb[2]; + for (uint m = 0u; m < 2u; ++m) + { + const float16_t q0 = CL8(q[m][0]), q1 = CL8(q[m][1]), q2 = CL8(q[m][2]), q3 = CL8(q[m][3]); + const float16_t q4 = CL8(q[m][4]), q5 = CL8(q[m][5]), q6 = CL8(q[m][6]), q7 = CL8(q[m][7]); + qa[m][0] = floate4m3_t(q0); qb[m][0] = qa[m][0]; qa[m][1] = floate4m3_t(q1); qb[m][1] = qa[m][1]; + qa[m][2] = floate4m3_t(q2); qb[m][2] = qa[m][2]; qa[m][3] = floate4m3_t(q3); qb[m][3] = qa[m][3]; + qa[m][4] = floate4m3_t(q4); qb[m][4] = qa[m][4]; qa[m][5] = floate4m3_t(q5); qb[m][5] = qa[m][5]; + qa[m][6] = floate4m3_t(q6); qb[m][6] = qa[m][6]; qa[m][7] = floate4m3_t(q7); qb[m][7] = qa[m][7]; + } + // the bias is read column-major, so the accumulator holds S^T: this lane holds the keys 8h..8h+7 of one query + HAC s = bias16(ob, 16u); + K32H(s, qa[0], qb[0], qa[1], qb[1]) + float16_t e[8]; + softmax2(float16_t(s[0]), float16_t(s[1]), e[0], e[1]); + softmax2(float16_t(s[2]), float16_t(s[3]), e[2], e[3]); + softmax2(float16_t(s[4]), float16_t(s[5]), e[4], e[5]); + softmax2(float16_t(s[6]), float16_t(s[7]), e[6], e[7]); + precise f16vec2 G0_ = f16vec2(e[0], e[1]), G1_ = f16vec2(e[2], e[3]); + precise f16vec2 G2_ = f16vec2(e[4], e[5]), G3_ = f16vec2(e[6], e[7]); + precise f16vec2 V0_ = G0_ + UN16(SH16(G0_)), V1_ = G1_ + UN16(SH16(G1_)); + precise f16vec2 V2_ = G2_ + UN16(SH16(G2_)), V3_ = G3_ + UN16(SH16(G3_)); + precise f16vec2 VV_ = (V0_ + V1_) + (V2_ + V3_); + const float16_t rc = float16_t(1.0 / float(VV_.x + VV_.y)); + B16 pb; + pb[0] = G0_.x * rc; pb[1] = G0_.y * rc; pb[2] = G1_.x * rc; pb[3] = G1_.y * rc; + pb[4] = G2_.x * rc; pb[5] = G2_.y * rc; pb[6] = G3_.x * rc; pb[7] = G3_.y * rc; + // V through this wave's shared buffer, then O^T = V^T x P^T: rows are the head channels, columns the tokens + ACC o[2]; + coopMatStore(v[0], VSH, (mt * 2u + 0u) * 64u, VWD / 4u, RM); + coopMatStore(v[1], VSH, (mt * 2u + 1u) * 64u, VWD / 4u, RM); + A16 vt0, vt1; + coopMatLoad(vt0, VSH, (mt * 2u + 0u) * 64u, VWD / 4u, RM); + coopMatLoad(vt1, VSH, (mt * 2u + 1u) * 64u, VWD / 4u, RM); + o[0] = coopMatMulAdd(vt0, pb, ACC(0.0)); + o[1] = coopMatMulAdd(vt1, pb, ACC(0.0)); + R16(o[0]); R16(o[1]); + B8 ob8[2]; + MK8_F8(ob8[0], o[0]); MK8_F8(ob8[1], o[1]); + // the output projection of this head's 32 attention channels, straight into the residual stream +#define WO_N(N) { \ + A8 a0, a1; \ + LDA8(a0, ow, 0u, UC, 16u * (N)); \ + LDA8(a1, ow, 1u, UC, 16u * (N)); \ + K32H(res[N], a0, ob8[0], a1, ob8[1]) } + EACH_T(WO_N) + } + // ------------------------------------------------------------ norm 2 + // the native element (T, 8h+i) holds the natural channel tok_xnat(T, 8h+i) = 16T + 4h + (i&3) + 8*((i>>2)&1), + // which lives in my natural-order accumulator at the half ((i>>2)&1) and the element 4h+(i&3). The tree needs + // the pairs in that order, so fetch the (pair 2h, pair 2h+1) of the half ((i>>2)&1) and select by this lane's half. + f16vec2 XA_[NT], XB_[NT], XC_[NT], XD_[NT]; + for (uint t_ = 0u; t_ < NT; ++t_) + { + const f16vec2 p0_ = f16vec2(float16_t(res[t_][0]), float16_t(res[t_][1])); + const f16vec2 p1_ = f16vec2(float16_t(res[t_][2]), float16_t(res[t_][3])); + const f16vec2 p2_ = f16vec2(float16_t(res[t_][4]), float16_t(res[t_][5])); + const f16vec2 p3_ = f16vec2(float16_t(res[t_][6]), float16_t(res[t_][7])); + const f16vec2 o0_ = UN16(SH16(p0_)), o1_ = UN16(SH16(p1_)), o2_ = UN16(SH16(p2_)), o3_ = UN16(SH16(p3_)); + if (hf == 0u) { XA_[t_] = p0_; XB_[t_] = p1_; XC_[t_] = o0_; XD_[t_] = o1_; } + else { XA_[t_] = o2_; XB_[t_] = o3_; XC_[t_] = p2_; XD_[t_] = p3_; } + } +#define EV(T, I) ((uint(I) < 2u) ? XA_[T][uint(I) & 1u] : (uint(I) < 4u) ? XB_[T][uint(I) & 1u] : \ + (uint(I) < 6u) ? XC_[T][uint(I) & 1u] : XD_[T][uint(I) & 1u]) + float16_t rn2; RSQ_TREE(rn2) +#undef EV +#define N2(N) { \ + const uint wb = (p.o[2] + 16u * (N) + ch8) / 4u, bb = (p.o[3] + 16u * (N) + ch8) / 4u; \ + const f16vec4 g0 = w16.w[wb], g1 = w16.w[wb + 1u], d0 = w16.w[bb], d1 = w16.w[bb + 1u]; \ + const float16_t x0 = res[N][0], x1 = res[N][1], x2 = res[N][2], x3 = res[N][3]; \ + const float16_t x4 = res[N][4], x5 = res[N][5], x6 = res[N][6], x7 = res[N][7]; \ + h1[N][0] = floate4m3_t(CL8(x0 * float16_t(rn2 * g0.x))); h1[N][1] = floate4m3_t(CL8(x1 * float16_t(rn2 * g0.y))); \ + h1[N][2] = floate4m3_t(CL8(x2 * float16_t(rn2 * g0.z))); h1[N][3] = floate4m3_t(CL8(x3 * float16_t(rn2 * g0.w))); \ + h1[N][4] = floate4m3_t(CL8(x4 * float16_t(rn2 * g1.x))); h1[N][5] = floate4m3_t(CL8(x5 * float16_t(rn2 * g1.y))); \ + h1[N][6] = floate4m3_t(CL8(x6 * float16_t(rn2 * g1.z))); h1[N][7] = floate4m3_t(CL8(x7 * float16_t(rn2 * g1.w))); \ + res[N][0] = float16_t(x0 + d0.x); res[N][1] = float16_t(x1 + d0.y); \ + res[N][2] = float16_t(x2 + d0.z); res[N][3] = float16_t(x3 + d0.w); \ + res[N][4] = float16_t(x4 + d1.x); res[N][5] = float16_t(x5 + d1.y); \ + res[N][6] = float16_t(x6 + d1.z); res[N][7] = float16_t(x7 + d1.w); } + EACH_T(N2) + // ------------------------------------------------------------ the MLP, 32 hidden channels at a time + for (uint c = 0u; c < UC / 8u; ++c) + { + const uint ow1 = p.o[20u + 3u * c], ob1 = p.o[21u + 3u * c], ow2 = p.o[22u + 3u * c]; + HAC hm[2]; + hm[0] = bias16(ob1, 32u); hm[1] = bias16(ob1 + 16u, 32u); +#define MLP_G(A, B) { \ + const uint kt = A; \ + A8 a0, a1; \ + LDA8(a0, ow1, kt, 32u, 0u); \ + LDA8(a1, ow1, kt + 1u, 32u, 0u); \ + K32H(hm[0], a0, h1[A], a1, h1[B]) \ + LDA8(a0, ow1, kt, 32u, 16u); \ + LDA8(a1, ow1, kt + 1u, 32u, 16u); \ + K32H(hm[1], a0, h1[A], a1, h1[B]) } + EACH_T16(MLP_G) + B8 g0, g1; + { const float16_t v0 = gelu1(hm[0][0]), v1 = gelu1(hm[0][1]), v2 = gelu1(hm[0][2]), v3 = gelu1(hm[0][3]); + const float16_t v4 = gelu1(hm[0][4]), v5 = gelu1(hm[0][5]), v6 = gelu1(hm[0][6]), v7 = gelu1(hm[0][7]); + g0[0] = floate4m3_t(CL8(v0)); g0[1] = floate4m3_t(CL8(v1)); g0[2] = floate4m3_t(CL8(v2)); g0[3] = floate4m3_t(CL8(v3)); + g0[4] = floate4m3_t(CL8(v4)); g0[5] = floate4m3_t(CL8(v5)); g0[6] = floate4m3_t(CL8(v6)); g0[7] = floate4m3_t(CL8(v7)); } + { const float16_t v0 = gelu1(hm[1][0]), v1 = gelu1(hm[1][1]), v2 = gelu1(hm[1][2]), v3 = gelu1(hm[1][3]); + const float16_t v4 = gelu1(hm[1][4]), v5 = gelu1(hm[1][5]), v6 = gelu1(hm[1][6]), v7 = gelu1(hm[1][7]); + g1[0] = floate4m3_t(CL8(v0)); g1[1] = floate4m3_t(CL8(v1)); g1[2] = floate4m3_t(CL8(v2)); g1[3] = floate4m3_t(CL8(v3)); + g1[4] = floate4m3_t(CL8(v4)); g1[5] = floate4m3_t(CL8(v5)); g1[6] = floate4m3_t(CL8(v6)); g1[7] = floate4m3_t(CL8(v7)); } +#define MLP2_N(N) { \ + A8 a0, a1; \ + LDA8(a0, ow2, 0u, UC, 16u * (N)); \ + LDA8(a1, ow2, 1u, UC, 16u * (N)); \ + K32H(res[N], a0, g0, a1, g1) } + EACH_T(MLP2_N) + } + // ------------------------------------------------------------ the block output, FP8 + const uint dst = uint(ty * p.W + tx) * UC + ch8; + if (inside) + for (uint nt = 0u; nt < NT; ++nt) + { +#define OUT8(I) out0.a[dst + 16u * nt + uint(I)] = floate4m3_t(CL8(res[nt][I])); + OUT8(0) OUT8(1) OUT8(2) OUT8(3) OUT8(4) OUT8(5) OUT8(6) OUT8(7) +#undef OUT8 + } + // ------------------------------------------------------------ patch merge: 2x2 tokens (4C) -> XC, half resolution +#if KIND == 1 + { + // the 16 merged rows are gathered from the four waves: the row is the merged token of the lane's pixel + const uint mrow = (byr >> 1) * 4u + (bxr >> 1), qq = (byr & 1u) * 2u + (bxr & 1u); +#define MG_W(N) { A8 y; MK8_F8(y, res[N]); coopMatStore(y, NB, 16u * S4 + mrow * S4 + qq * UC + 16u * (N) - row * S4, S4, RM); } + EACH_T(MG_W) + SYNC + const uint MT = UXC / 16u, NI = (MT + 3u) / 4u; + const int ym = (Y0 >> 1) + int(row >> 2), xm = (X0 >> 1) + int(row & 3u); + const bool mok = ym >= 0 && ym < p.Hh / 2 && xm >= 0 && xm < p.W / 2; + const uint md = uint(ym * (p.W / 2) + xm) * UXC + ch8; + for (uint i = 0u; i < NI; ++i) + { + const uint u = mt + 4u * i; + if (u < MT) + { + HAC y = bias16(p.o[69] + 16u * u, UXC); + for (uint kt = 0u; kt < 4u * UC / 16u; kt += 2u) + { + B8 b0, b1; coopMatLoad(b0, NB, 16u * S4 + kt * 16u, S4, CM); + coopMatLoad(b1, NB, 16u * S4 + (kt + 1u) * 16u, S4, CM); + A8 a0, a1; LDA8(a0, p.o[68], kt, UXC, 16u * u); + LDA8(a1, p.o[68], kt + 1u, UXC, 16u * u); + K32H(y, a0, b0, a1, b1) + } + if (mok) + { +#define MGO(I) out1.a[md + 16u * u + uint(I)] = floate4m3_t(CL8(y[I])); + MGO(0) MGO(1) MGO(2) MGO(3) MGO(4) MGO(5) MGO(6) MGO(7) +#undef MGO + } + } + } + } +#endif +} diff --git a/engine/m_model.py b/engine/m_model.py new file mode 100644 index 0000000..906062f --- /dev/null +++ b/engine/m_model.py @@ -0,0 +1,186 @@ +#!/usr/bin/env python3 +"""Preset M (rrlite): plain weights of the Swin layers and a plain numpy model of them. + +NGX keeps the weights as FP8 (e4m3) bytes in a layout arranged for its tensor-core fragments, and half of the +activations in a permuted ("pair") channel order. plain_layer() undoes both: every tensor becomes a row-major +matrix W[K][N] (y = x W) over natural channel order, f16 exact (every e4m3 value is an f16 value). +forward() is the layer's mathematics on those matrices. Every activation that feeds a matrix product as its left +operand is rounded to FP8 first, as in NVIDIA's arithmetic (its tensor-core operands are FP8) and in d4r's native +FP8 kernels; sums are accumulated without the intermediate f16 rounding of NVIDIA's k32 steps. + +usage: m_model.py CAPTURE_DIR [frame] compares forward() with the captured outputs of each layer launch +""" +import glob, os, struct, sys +import numpy as np +sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'kernels', 'tools')) +import model_enc3 as me +import swin_model as sm + +F16 = np.float16 +# kernel -> (C, heads, merged channels, low-resolution input channels) +LAYERS = {'enc1': (64, 2, 96, 0), 'enc2': (96, 4, 128, 0), 'enc3_tube': (128, 4, 0, 0), 'dec2': (96, 4, 0, 128), 'dec1': (64, 2, 0, 96)} +ORDER = ['enc1', 'enc2'] + ['enc3_tube'] * 6 + ['dec2', 'dec1'] + + +def gp(n): + return sm.gp(np.arange(n)) + + +def plain_layer(P, name): + """dict of plain f16 tensors of one layer launch (P: model_enc3.Params)""" + C, heads, NPM, CIN = LAYERS[name] + fp8 = lambda T: me.E4M3[P.wbuf[T]] + pre = 4 * C * CIN + 8 * C if CIN else 0 + L = sm.layout(C, heads, pre) + g = gp(C) + t = {'g1': P.f16vec(L['g1'], C)[g], 'bo': P.f16vec(L['bo'], C)[g], 'g2': P.f16vec(L['g2'], C)[g], 'b2': P.f16vec(L['b2'], C)[g], + 'q': [], 'v': [], 'bias': [], 'w1': [], 'b1': [], 'w2': []} + wo = np.zeros((32 * heads, C)) + g32 = gp(32) + for h in range(heads): + hb = L['head'][h] + t['q'].append(fp8(sm.woff_table(hb, 1024, 512, C, 32))) + t['v'].append(fp8(sm.woff_table(hb + 32 * C, 1024, 512, C, 32))) + t['bias'].append(sm.rel_bias(P, hb + 64 * C)) + raw = fp8(sm.woff_table(hb + 64 * C + 512, 0, 512, 32, C)) # rows: pair-ordered attention outputs + wo[32 * h + g32] = raw[:, g] # Wo[gp(a)][c] = raw[a][gp(c)] + t['wo'] = wo + for B1 in L['b1']: + t['w1'].append(fp8(sm.woff_table(B1, 512, 16 * C, C, 32))) + t['b1'].append(P.f16vec(B1 + 32 * C, 32)) + raw = fp8(sm.woff_table(B1 + 32 * C + 64, 0, 512, 32, C)) + w2 = np.zeros((32, C)); w2[g32] = raw[:, g] + t['w2'].append(w2) + if NPM: + PM0 = L['end'] + W = np.concatenate([fp8(sm.woff_table(PM0 + k * 128 * C, 512, 64 * C, 4 * C, 32)) for k in range(NPM // 32)], 1) + b = np.concatenate([P.f16vec(PM0 + 4 * C * NPM + 64 * k, 32) for k in range(NPM // 32)]) + gm = gp(NPM) + t['merge_w'], t['merge_b'] = W[:, gm], b[gm] + if CIN: + cols = 4 * C // heads + W = np.concatenate([fp8(sm.woff_table(w * CIN * cols, 512, 16 * CIN, CIN, cols)) for w in range(heads)], 1) # [CIN][4C] + b = np.concatenate([P.f16vec(4 * C * CIN + 2 * cols * w, cols) for w in range(heads)]) + idx = (np.arange(4)[:, None] * C + g[None, :]).ravel() # x0[.., q, c] = E[q C + gp(c)] + t['expand_w'], t['expand_b'] = W[:, idx], b[idx] + return {k: ([np.asarray(a, F16) for a in v] if isinstance(v, list) else np.asarray(v, F16)) for k, v in t.items()} + + +def h(a): + return np.asarray(a, np.float64).astype(F16).astype(np.float64) + + +def q8(a): + return me.q8(a) + + +def softmax_rows(S): + t = np.clip(h(S * me.C_SCALE), -float(me.CLAMP), float(me.CLAMP)) + poly = h(t * h(t * (-t) + me.P1) + float(me.P0)) + w = me.exp_trick(poly).astype(np.float64) + s = h(w.sum(1, keepdims=True)) + return h(w * h(1.0 / s)) + + +def block(t, x0, C, heads): + """x0: [windows][16][C] f16 values; returns the block output before the FP8 rounding""" + f = lambda a: np.asarray(a, np.float64) + r = h(1.0 / np.sqrt(h((x0 * x0).sum(-1, keepdims=True) + 2.0 ** -13))) + h1 = q8(h(x0 * h(r * f(t['g1'])))) + O = np.zeros(x0.shape[:2] + (32 * heads,)) + for k in range(heads): + Q = h(h1 @ f(t['q'][k])); V = h(h1 @ f(t['v'][k])) + for w in range(x0.shape[0]): + Qq = q8(Q[w]) + O[w, :, 32 * k:32 * k + 32] = q8(h(softmax_rows(h(Qq @ Qq.T + f(t['bias'][k]))) @ V[w])) + x1 = h(O @ f(t['wo']) + x0 + f(t['bo'])) + r = h(1.0 / np.sqrt(h((x1 * x1).sum(-1, keepdims=True) + 2.0 ** -13))) + h2 = q8(h(x1 * h(r * f(t['g2'])))) + x2 = x1 + f(t['b2']) + for w1, b1, w2 in zip(t['w1'], t['b1'], t['w2']): + x2 = x2 + q8(me.gelu_poly(h(h2 @ f(w1) + f(b1)))) @ f(w2) + return h(x2) + + +def tokens(buf, off, W, H, C): + """planes of 32 e4m3 channels at byte `off` -> [H][W][C] values""" + a = me.E4M3[buf[off:off + (C // 32) * H * W * 32]].reshape(C // 32, H, W, 32) + return a.transpose(1, 2, 0, 3).reshape(H, W, C) + + +def mirror(v, n): + v = np.abs(v); return np.minimum(v, 2 * n - 2 - v) + + +def forward(t, name, x, skip, sx, sy): + """x: [H][W][C] input tokens (decoders: [H/2][W/2][CIN]); returns (out [H][W][C], merged or None), FP8-rounded""" + C, heads, NPM, CIN = LAYERS[name] + f = lambda a: np.asarray(a, np.float64) + if CIN: + H, W = skip.shape[:2] + else: + H, W = x.shape[:2] + gx, gy = (W + sx + 7) // 8, (H + sy + 7) // 8 + Y, X = np.meshgrid(np.arange(8 * gy) - sy, np.arange(8 * gx) - sx, indexing='ij') + if CIN: + e = h(x @ f(t['expand_w']) + f(t['expand_b'])).reshape(x.shape[0], x.shape[1], 4, C) + # the low-resolution token of block row r is taken at (8 bx - sx) // 2 + r % 4 (floor), then mirrored + Yb, Xb = (Y + sy) // 8 * 8 - sy, (X + sx) // 8 * 8 - sx + ly, lx = mirror(Yb // 2 + (Y - Yb) // 2, H // 2), mirror(Xb // 2 + (X - Xb) // 2, W // 2) + q = ((Y - Yb) & 1) * 2 + ((X - Xb) & 1) + x0 = h(q8(e[ly, lx, q]) + skip[mirror(Y, H), mirror(X, W)]) + else: + x0 = x[mirror(Y, H), mirror(X, W)] + win = x0.reshape(2 * gy, 4, 2 * gx, 4, C).transpose(0, 2, 1, 3, 4).reshape(-1, 16, C) + y = block(t, win, C, heads).reshape(2 * gy, 2 * gx, 4, 4, C).transpose(0, 2, 1, 3, 4).reshape(8 * gy, 8 * gx, C) + out = q8(y) + merged = None + if NPM: + A = out.reshape(4 * gy, 2, 4 * gx, 2, C).transpose(0, 2, 1, 3, 4).reshape(4 * gy, 4 * gx, 4 * C) + m = q8(h(A @ f(t['merge_w']) + f(t['merge_b']))) + # merged token (my, mx) of a block sits at (8 by - sy) // 2 + my + my = ((np.arange(4 * gy) // 4) * 8 - sy) // 2 + np.arange(4 * gy) % 4 + mx = ((np.arange(4 * gx) // 4) * 8 - sx) // 2 + np.arange(4 * gx) % 4 + merged = np.zeros((H // 2, W // 2, NPM)) + oky, okx = (my >= 0) & (my < H // 2), (mx >= 0) & (mx < W // 2) + merged[np.ix_(my[oky], mx[okx])] = m[np.ix_(oky, okx)] + oy, ox = (Y[:, 0] >= 0) & (Y[:, 0] < H), (X[0] >= 0) & (X[0] < W) + return out[np.ix_(oy, ox)], merged + + +def launches(cap, frame=0): + """the ten layer launches of one frame, in network order, as model_enc3.Params""" + dirs = sorted(d for d in glob.glob(f'{cap}/replay/replay-*-rrlite_*_4x4') if any(d.endswith(f'rrlite_{k}_4x4') for k in LAYERS)) + if len(dirs) < 10 * (frame + 1): raise ValueError('capture lacks the layer launches of that frame') + return [me.Params(d) for d in dirs[10 * frame:10 * frame + 10]], dirs[10 * frame:10 * frame + 10] + + +def psnr(a, b): + d = (a - b) ** 2 + return 10 * np.log10(np.abs(b).max() ** 2 / max(d.mean(), 1e-30)) + + +def main(): + cap = sys.argv[1]; frame = int(sys.argv[2]) if len(sys.argv) > 2 else 0 + Ps, dirs = launches(cap, frame) + for i, (P, d, name) in enumerate(zip(Ps, dirs, ORDER)): + C, heads, NPM, CIN = LAYERS[name] + seq = int(os.path.basename(d).split('-')[1]) + after = np.fromfile(glob.glob(f'{cap}/dump/launch-{seq:03d}-*-arg024-*')[0], np.uint8) + t = plain_layer(P, name) + W, H = P.W, P.H + if CIN: + x = tokens(P.arena, P.inp - P.abase, W // 2, H // 2, CIN); skip = tokens(P.arena, P.p32 - P.abase, W, H, C) + else: + x = tokens(P.arena, P.inp - P.abase, W, H, C); skip = None + out, merged = forward(t, name, x, skip, P.sx, P.sy) + ref = tokens(after, P.out - P.abase, W, H, C) + line = f'{name:10s} launch {seq:3d} {W}x{H}x{C} shift {P.sx},{P.sy}: output {psnr(out, ref):5.1f} dB, identical {100 * (out == ref).mean():5.1f}%' + if NPM: + mref = tokens(after, P.p48 - P.abase, W // 2, H // 2, NPM) + line += f'; merged {psnr(merged, mref):5.1f} dB, identical {100 * (merged == mref).mean():5.1f}%' + print(line, flush=True) + + +if __name__ == '__main__': + main() diff --git a/engine/normalized_store.glsl b/engine/normalized_store.glsl new file mode 100644 index 0000000..762061c --- /dev/null +++ b/engine/normalized_store.glsl @@ -0,0 +1,22 @@ +// Match formatted CUDA stores before the Vulkan image conversion. The reference rounds halfway away from +// zero; feeding an exact normalized code to imageStore also avoids implementation-specific UNORM/SNORM rounding. +vec4 unorm8Value(vec4 v) { + v = clamp(mix(v, vec4(0.0), isnan(v)), 0.0, 1.0); + precise vec4 scaled = v * 255.0; + return floor(scaled + 0.5) * (1.0 / 255.0); +} +vec4 snorm8Value(vec4 v) { + v = clamp(mix(v, vec4(0.0), isnan(v)), -1.0, 1.0); + precise vec4 scaled = abs(v) * 127.0; + return sign(v) * floor(scaled + 0.5) * (1.0 / 127.0); +} +vec3 unorm10Value(vec3 v) { + v = clamp(mix(v, vec3(0.0), isnan(v)), 0.0, 1.0); + precise vec3 scaled = v * 1023.0; + precise vec3 code = floor(scaled + 0.5); + precise vec3 value = code * (1.0 / 1023.0); + // CUDA's float-backed formatted surface stores code/1023 with correctly rounded division. RADV lowers + // even precise constant division to reciprocal multiplication (24/1024 codes differ by one ulp). + // Correct the quotient with its fused residual; all 1024 normalized codes then match bit-for-bit. + return fma(fma(value, vec3(-1023.0), code), vec3(1.0 / 1023.0), value); +} diff --git a/engine/output_k.comp b/engine/output_k.comp new file mode 100644 index 0000000..1c835cd --- /dev/null +++ b/engine/output_k.comp @@ -0,0 +1,397 @@ +#version 460 +// d4r engine, preset K: the output stage (temporal accumulation and upsampling), written from the algorithm +// described in OUTPUT_K_NOTES.md. One 16x16 workgroup shades a 16x16 block of output pixels. +// 1. The block's footprint of the render-resolution colour (up to 24x24 texels) is range-compressed, converted to +// YCoCg and staged in shared memory (with depth, unless the motion is display-resolution). +// 2. Motion comes from the nearest of the diagonal depth neighbours (display-resolution: the pixel's own vector); it locates the history. +// 3. The 2x-resolution luma history is resampled with a sharpened Catmull-Rom filter at the four sub-pixels. +// 4. The current luma is filtered at the four sub-pixels with twelve taps of an anisotropic Gaussian whose shape +// comes from a table indexed by two network outputs. +// 5. History and current are blended per sub-pixel with a confidence from the tap geometry; chroma follows luma. +// Not handled (the reference kernel has them, the captured frames do not use them): blocks whose footprint exceeds +// 24x24 texels, and the alternative output modes selected by three flag bytes. +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_shuffle : require +#extension GL_GOOGLE_include_directive : require +#include "output_store.glsl" +#include "round_half.glsl" +layout(local_size_x = 16, local_size_y = 16) in; +layout(set = 0, binding = 0, std430) readonly buffer Params +{ + vec2 out2rnd, rndOffs; // render position = (pixel + .5) * out2rnd + rndOffs + vec2 near2rnd, sub2rnd; // the same scale for the nearest-texel lookup; for 2x sub-pixels + vec2 out2tok; // token column = int(pixel * out2tok) / 2 + vec2 mvScale, mvOffs; // previous position = pixel + .5 + mvScale * (mv + mvOffs) + vec2 invOut, invOutValid; // 1 / output size, for the history uv and for the on-screen test + ivec2 colOrigin, colMax, depOrigin, depMax, mvOrigin, mvMax, outOrigin, outSize; + int tokW, tokH; + uint flags; // bit 0: reset; bit 1: regular depth (larger is farther), otherwise inverted; + // bit 2: display-resolution motion (one vector per output pixel, no depth); + // bits 8..11: which of the model's reconstruction tables this frame uses + float expo, expoRatio; // exposure scale of this frame; previous / current + float sharp, minAlpha, alphaThr; + float lutScale, kA, kB, c184, c186; +} p; +layout(set = 0, binding = 1) uniform sampler2D colorTex; // render resolution, linear RGB + alpha +layout(set = 0, binding = 2) uniform sampler2D histTex; // output resolution: compressed Y, Co, Cg, alpha (bilinear) +layout(set = 0, binding = 3) uniform sampler2D hiresTex; // 2x output resolution: compressed luma +layout(set = 0, binding = 4) uniform sampler2D motionTex; +layout(set = 0, binding = 5) uniform sampler2D depthTex; +layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D histOut; +layout(set = 0, binding = 7, r16ui) uniform writeonly uimage2D hiresOut; +layout(set = 0, binding = 8, rgba16f) uniform writeonly image2D tokOut; +layout(set = 0, binding = 9, FINAL_OUTPUT_FORMAT) uniform writeonly image2D finalOut; +layout(set = 0, binding = 10, std430) readonly buffer Head { float16_t h[]; } head; // 40 values per token: 4 sub-blocks x 10 +layout(set = 0, binding = 11, std430) readonly buffer Lut { uvec4 e[]; } lut; +shared f16vec4 tileC[576]; // compressed Y, Co, Cg and alpha of the footprint, 24 texels per row +shared float tileD[576]; + +// LDR (NGX's *_ldr_* kernels, the game did not set the HDR flag): display-referred colour, no range compression +#ifndef LDR +#define LDR 0 +#endif +#if LDR +float compressK(float L) { return 1.0; } +float expandK(float v) { return 1.0; } +#else +// The reference's constants and association (output_color_pack): rcp.approx(y + 1), or above 1000 +// rcp.approx(y) * (lg2.approx(fma(y, 9, 1)) * ln2 / 9 + 0xBC4FAE29). Its offset differs by 1.7e-5 +// relative from 0.012676, and folding ln2/9 into one constant moves the result a further ulp. +float compressK(float L) // L * compressK(L) maps [0, inf) to [0, 1): Reinhard below 1000, logarithmic above +{ + if (L < 1000.0) + return 1.0 / (L + 1.0); + precise const float t = (log2(fma(L, 9.0, 1.0)) * uintBitsToFloat(0x3F317218u)) / 9.0 + + uintBitsToFloat(0xBC4FAE29u); + return t / L; +} +float expandK(float v) // v * expandK(v) is the inverse +{ + if (v < uintBitsToFloat(0x3F7FBE88u)) + return 1.0 / (1.0 - v); + precise const float y = (v + uintBitsToFloat(0x3C4FAE29u)) * 9.0; + return (exp2(y * uintBitsToFloat(0x3FB8AA3Bu)) - 1.0) / (9.0 * v); +} +#endif +float roundHalfAway(float x) { const float t = trunc(x); return abs(x - t) >= 0.5 ? t + sign(x) : t; } +ivec2 clampOrigin(ivec2 v, ivec2 origin, ivec2 hi) // the reference's clamp: origin for v <= 0, else min(v + origin, hi) +{ + return ivec2(v.x > 0 ? min(v.x + origin.x, hi.x) : origin.x, v.y > 0 ? min(v.y + origin.y, hi.y) : origin.y); +} +vec4 softIn(vec4 v) { return mix(0.99 + (v - 0.99) * 0.001, v, lessThan(v, vec4(0.99))); } // stored luma -> working luma +// Catmull-Rom weights of the four taps around a fraction t +void crWeights(float t, out float w0, out float w1, out float w2, out float w3) +{ + const float t2 = t * t, t3 = t2 * t; + w0 = t2 - 0.5 * (t + t3); + w1 = 1.5 * t3 - 2.5 * t2 + 1.0; + w3 = 0.5 * (t3 - t2); + w2 = 1.0 - w0 - w1 - w3; +} +// The reference's six coefficients for one filter axis, computed once per pixel: the bilinear split +// mid = w1 / (w1 + w2) and the four sharpened weights, kept separate from the bilinear pair (folding +// sharp*(w1-mid) into mid changes the row result by up to a few ulp, which the half-encoded highlight +// range then amplifies). +struct CrAxis { float mid, split, s0, s1, s2, s3; }; +CrAxis crAxis(float w0, float w1, float w2, float w3, float mid) +{ + return CrAxis(mid, 1.0 - mid, p.sharp * w0, p.sharp * (w1 - mid), p.sharp * (w2 - (1.0 - mid)), p.sharp * w3); +} +// One pass along one axis: the two middle taps are bilinear, the four sharpened taps are saturated, and +// every term is the reference's own rounded product/fma. +float crPass(float t0, float t1, float t2, float t3, CrAxis c) +{ + float t = c.split * t2; + t = fma(c.mid, t1, t); + t = fma(c.s0, clamp(t0, 0.0, 1.0), t); + t = fma(c.s1, clamp(t1, 0.0, 1.0), t); + t = fma(c.s2, clamp(t2, 0.0, 1.0), t); + t = fma(c.s3, clamp(t3, 0.0, 1.0), t); + return t; +} +void main() +{ + const ivec2 pix = ivec2(gl_GlobalInvocationID.xy); + const vec2 pc = vec2(pix) + 0.5; + const bool displayMv = (p.flags & 4u) != 0u; + // ---------------------------------------------------------------- 1. footprint of the block + const vec2 b0 = (vec2(gl_WorkGroupID.xy * 16u) + 0.5) * p.out2rnd + p.rndOffs; + const vec2 b1 = (vec2(gl_WorkGroupID.xy * 16u + 15u) + 0.5) * p.out2rnd + p.rndOffs; + const ivec2 t0 = ivec2(vec2(roundHalfAway(b0.x), roundHalfAway(b0.y)) - 4.0); + const ivec2 t1 = ivec2(vec2(roundHalfAway(b1.x), roundHalfAway(b1.y)) + 4.0); + const int tw = t1.x - t0.x + 1, th = t1.y - t0.y + 1; + for (int i = int(gl_LocalInvocationIndex); i < tw * th; i += 256) + { + const ivec2 t = ivec2(i % tw, i / tw), g = t + t0; + const vec4 c = texelFetch(colorTex, clampOrigin(g, p.colOrigin, p.colMax), 0); + const vec3 rgb = max(c.rgb, 0.0); + const float Y = 0.25 * rgb.r + 0.5 * rgb.g + 0.25 * rgb.b; + const float L = p.expo * Y; + const float k = compressK(L); + tileC[t.y * 24 + t.x] = f16vec4(L * k, p.expo * (0.5 * rgb.r - 0.5 * rgb.b) * k, p.expo * (0.5 * rgb.g - 0.25 * rgb.r - 0.25 * rgb.b) * k, c.a); + // larger = nearer: regular depth is negated, which turns its "nearest" comparisons into the inverted ones + if (!displayMv) + { + const float depth = texelFetch(depthTex, clampOrigin(g, p.depOrigin, p.depMax), 0).r; + tileD[t.y * 24 + t.x] = (p.flags & 2u) != 0u ? -depth : depth; + } + } + memoryBarrierShared(); barrier(); + // ---------------------------------------------------------------- network outputs of this pixel's 2x2 sub-block + const ivec2 tk = ivec2(vec2(pix) * p.out2tok); + const uint hbase = uint((tk.y / 2) * p.tokW + tk.x / 2) * 40u + uint((tk.y & 1) * 2 + (tk.x & 1)) * 10u; + // sigmoid of a head word (the reference's fma chain: y = x * (x^2 + 27) / fma(x^2, 9, 27), sat(0.5*y + 0.5)) + const float x = 0.5 * float(head.h[hbase + 9u]); + precise const float x2 = x * x; + const float sig = clamp(x < -3.0 ? 0.0 : x > 3.0 ? 1.0 + : fma((x * (x2 + 27.0)) / fma(x2, 9.0, 27.0), 0.5, 0.5), 0.0, 1.0); + // ---------------------------------------------------------------- 2. motion + // Render-resolution vectors come from the nearest diagonal depth neighbour; display-resolution ones belong to the pixel. + ivec2 mvAt = pix; + if (!displayMv) + { + const vec2 nr = pc * p.near2rnd + p.rndOffs; + const ivec2 ni = ivec2(roundHalfAway(nr.x), roundHalfAway(nr.y)); + const int di = (ni.y - t0.y) * 24 + (ni.x - t0.x); + const float d0 = tileD[di], dmm = tileD[di - 25], dpp = tileD[di + 25], dpm = tileD[di - 23], dmp = tileD[di + 23]; + ivec2 dsel = ivec2(0); + if (d0 < 0.5 * (dmm + dpp) - 0.001 * abs(d0) || d0 < 0.5 * (dpm + dmp) - 0.001 * abs(d0)) + { + float m = d0; + if (dmm > m) { m = dmm; dsel = ivec2(-1, -1); } + if (dpm > m) { m = dpm; dsel = ivec2(1, -1); } + if (dmp > m) { m = dmp; dsel = ivec2(-1, 1); } + if (dpp > m) { dsel = ivec2(1, 1); } + } + mvAt = ni + dsel; + } + const vec2 mv = texelFetch(motionTex, clampOrigin(mvAt, p.mvOrigin, p.mvMax), 0).xy; + const vec2 q = pc + p.mvScale * (mv + p.mvOffs); // where this pixel was, in output pixels + const vec2 qv = q * p.invOutValid; + bool noHist = (p.flags & 1u) != 0u || qv.x < 0.0 || qv.x > 1.0 || qv.y < 0.0 || qv.y > 1.0; + const vec4 hist = textureLod(histTex, q * p.invOut, 0.0); + // ---------------------------------------------------------------- 3. luma history at the four sub-pixels + vec4 Hq; + { + const vec2 h = (q - 0.5) * 2.0 + 0.5; + const vec2 base = floor(h - 0.5) + 0.5; + const vec2 fr = clamp(h - base, 0.0, 1.0); + float x0, x1, x2, x3, y0, y1, y2, y3; + crWeights(fr.x, x0, x1, x2, x3); crWeights(fr.y, y0, y1, y2, y3); + const float mx = x1 / (x1 + x2), my = y1 / (y1 + y2); + const vec2 ts = p.invOut * 0.5; + // 5x5 texels T[col][row] from nine gathers (gather: x = (0,1), y = (1,1), z = (1,0), w = (0,0)) + vec4 g00 = textureGather(hiresTex, ts * (base + vec2(-0.5, -0.5)), 0); + vec4 g10 = textureGather(hiresTex, ts * (base + vec2(1.5, -0.5)), 0); + vec4 g20 = textureGather(hiresTex, ts * (base + vec2(3.5, -0.5)), 0); + vec4 g01 = textureGather(hiresTex, ts * (base + vec2(-0.5, 1.5)), 0); + vec4 g11 = textureGather(hiresTex, ts * (base + vec2(1.5, 1.5)), 0); + vec4 g21 = textureGather(hiresTex, ts * (base + vec2(3.5, 1.5)), 0); + vec4 g02 = textureGather(hiresTex, ts * (base + vec2(-0.5, 3.5)), 0); + vec4 g12 = textureGather(hiresTex, ts * (base + vec2(1.5, 3.5)), 0); + vec4 g22 = textureGather(hiresTex, ts * (base + vec2(3.5, 3.5)), 0); + // stored luma at 0.99 and above is an encoding of [0.99, 1): rare enough to branch on + const vec4 gm = max(max(max(g00, g10), max(g20, g01)), max(max(g11, g21), max(max(g02, g12), g22))); + if (!(max(max(gm.x, gm.y), max(gm.z, gm.w)) < 0.99)) + { + g00 = softIn(g00); g10 = softIn(g10); g20 = softIn(g20); g01 = softIn(g01); g11 = softIn(g11); + g21 = softIn(g21); g02 = softIn(g02); g12 = softIn(g12); g22 = softIn(g22); + } + float T[5][5]; // [row][col] + T[0][0] = g00.w; T[0][1] = g00.z; T[1][0] = g00.x; T[1][1] = g00.y; + T[0][2] = g10.w; T[0][3] = g10.z; T[1][2] = g10.x; T[1][3] = g10.y; + T[0][4] = g20.w; T[1][4] = g20.x; + T[2][0] = g01.w; T[2][1] = g01.z; T[3][0] = g01.x; T[3][1] = g01.y; + T[2][2] = g11.w; T[2][3] = g11.z; T[3][2] = g11.x; T[3][3] = g11.y; + T[2][4] = g21.w; T[3][4] = g21.x; + T[4][0] = g02.w; T[4][1] = g02.z; T[4][2] = g12.w; T[4][3] = g12.z; T[4][4] = g22.w; + // Each row filtered for the two sub-pixel columns. Stored luma lies in [0, 1), so the clamps of the + // sharpened taps do nothing here; the column pass below keeps them (rows can overshoot). + const CrAxis cx = crAxis(x0, x1, x2, x3, mx), cy = crAxis(y0, y1, y2, y3, my); + float R[5][2]; + for (int r = 0; r < 5; ++r) + for (int c = 0; c < 2; ++c) + R[r][c] = crPass(T[r][c], T[r][c + 1], T[r][c + 2], T[r][c + 3], cx); + for (int sy = 0; sy < 2; ++sy) + for (int sx = 0; sx < 2; ++sx) + { + const float v = crPass(R[sy][sx], R[sy + 1][sx], R[sy + 2][sx], R[sy + 3][sx], cy); + const float lo = min(min(T[sy + 1][sx + 1], T[sy + 1][sx + 2]), min(T[sy + 2][sx + 1], T[sy + 2][sx + 2])); + const float hi = max(max(T[sy + 1][sx + 1], T[sy + 1][sx + 2]), max(T[sy + 2][sx + 1], T[sy + 2][sx + 2])); + Hq[sy * 2 + sx] = max(min(v, hi), lo); + } + // the history was compressed with last frame's exposure. The reference always re-exposes + // (v' = L * k(L) with L = ratio * (v * g(v))); skipping it when the ratio is exactly 1 leaves the + // pre-rounding value, which is not what k(g(v)) costs on a real frame. + for (int s = 0; s < 4; ++s) + { + const float L = p.expoRatio * (Hq[s] * expandK(Hq[s])); + Hq[s] = L * compressK(L); + } + } + { + precise const float chk = fma(sig, hist.b, fma(hist.r, hist.g, fma(Hq.x, Hq.y, Hq.z * Hq.w))); + noHist = noHist || !(abs(chk) < uintBitsToFloat(0x7F800000u)); + } + const float s1 = noHist ? 1.0 : sig, ha = noHist ? 1.0 : hist.a; + precise const float f343 = fma(s1, 1.0 - ha, ha); + precise const float f344 = fma(ha, 1.0 - s1 * s1, s1 * s1); + if (noHist) Hq = vec4(0.0); + const vec3 hc = noHist ? vec3(0.0) : hist.rgb; + // ---------------------------------------------------------------- current colour, bilinear from the tile + // Positions, weights and sums of this part are half floats, as in the reference: the filter below is sensitive + // to where the samples sit, and half precision moves them by up to 1/64 texel. + const f16vec2 u = f16vec2(pc * p.out2rnd + p.rndOffs - vec2(t0)); + const ivec2 ui = ivec2(floor(vec2(u))); + const f16vec2 uf = clamp(u - f16vec2(ui), float16_t(0.0), float16_t(1.0)), ug = float16_t(1.0) - uf; + const int ci = ui.y * 24 + ui.x; + f16vec4 curh = tileC[ci] * (ug.x * ug.y); + curh = fma(tileC[ci + 1], f16vec4(uf.x * ug.y), curh); + curh = fma(tileC[ci + 24], f16vec4(ug.x * uf.y), curh); + curh = fma(tileC[ci + 25], f16vec4(uf.x * uf.y), curh); + vec4 cur = vec4(curh); + // "bright" over the 2x2 pixel quad: the lanes of a subgroup are the pixels of two rows, x + 16 * (y & 1) + uint bright = cur.x > 0.8 ? 1u : 0u; + bright |= subgroupShuffleXor(bright, 16u); + bright |= subgroupShuffleXor(bright, 1u); + // ---------------------------------------------------------------- 4. kernel-predicted luma filter + const f16vec2 hm = f16vec2(head.h[hbase], head.h[hbase + 1u]); + const f16vec2 m = f16vec2(vec2(hm) / vec2(float16_t(1.0) + abs(hm))); + const ivec2 mi = clamp(ivec2(roundEven(vec2(m * float16_t(p.lutScale)))) + 8, 0, 15); + const uint fl = (bright << 2) | ((sig < 0.25 && f343 < 0.25) ? 2u : 0u) | (p.alphaThr > f343 ? 1u : 0u); + const uvec4 le = lut.e[(((p.flags >> 8) & 15u) << 11) | (fl << 8) | uint(mi.x << 4) | uint(mi.y)]; + const float16_t e0 = uint16BitsToFloat16(uint16_t(le.x & 0xffffu)), e1 = uint16BitsToFloat16(uint16_t(le.x >> 16)); + const float16_t e2 = uint16BitsToFloat16(uint16_t(le.y & 0xffffu)), slope = uint16BitsToFloat16(uint16_t(le.y >> 16)); + const bool alongX = (le.z & 0xffffu) == 0u; // rows of taps step in y, taps within a row step in x + const float16_t e5 = uint16BitsToFloat16(uint16_t(le.z >> 16)), e6 = uint16BitsToFloat16(uint16_t(le.w & 0xffffu)); + // sub-pixel positions in tile coordinates: s = sy * 2 + sx + // The reference fuses the scale and the offset and subtracts the tile origin afterwards. The order matters + // where a position falls on a half-float rounding tie (a whole output column at some size ratios): any other + // order moves those samples by one half-float step, which the recurrent history then spreads over the frame. + precise const vec2 sa0 = fma(vec2(pix * 2) + 0.5, p.sub2rnd, p.rndOffs), sa1 = fma(vec2(pix * 2 + 1) + 0.5, p.sub2rnd, p.rndOffs); + precise const vec2 sq0 = sa0 - vec2(t0), sq1 = sa1 - vec2(t0); + const f16vec2 sp0 = f16vec2(sq0), sp1 = f16vec2(sq1); + const f16vec4 subX = f16vec4(sp0.x, sp1.x, sp0.x, sp1.x), subY = f16vec4(sp0.y, sp0.y, sp1.y, sp1.y); + const f16vec4 subA = alongX ? subY : subX, subB = alongX ? subX : subY; + const float16_t cA = alongX ? u.y : u.x, cB = alongX ? u.x : u.y; + const int stepB = alongX ? 1 : 24, stepA = alongX ? 24 : 1; + const int a0 = int(floor(float(cA))) - 1; + // Every product and sum of the filter is rounded to half precision, as in the reference: with single precision + // the result differs visibly more (49 dB against 116 dB on the luma history of the test frames). + f16vec4 sumY = f16vec4(0.0), sumW = f16vec4(0.0); + const float16_t line0 = cB + (float16_t(a0) - cA) * slope; + float16_t line = line0; + for (int r = 0; r < 4; ++r) + { + const int bs = int(roundEven(float(line))) - 1; + const f16vec4 dA = float16_t(a0 + r) - subA; + f16vec4 db = float16_t(bs) - subB; + const f16vec4 cross = dA * e2; + f16vec4 lin = fma(db, cross, (dA * dA) * e0); + const int ti = (a0 + r) * stepA + bs * stepB; + for (int j = 0; j < 3; ++j) + { + // Keep the rounded exponent in f16 through exp2; avoids f16/f32 conversions. + const f16vec4 w = exp2(fma(db * db, f16vec4(e1), lin)); + sumY = fma(w, f16vec4(tileC[ti + j * stepB].x), sumY); sumW += w; + db += float16_t(1.0); lin += cross; + } + line += slope; + } + vec4 F = vec4(sumY) / vec4(sumW); + // Where every exponent is below -12 the result is the tap with the largest exponent instead. That needs all + // twelve weights below 2^-12, so the sum of weights rules it out for almost every pixel. + if (any(lessThan(sumW, f16vec4(12.5 / 4096.0)))) + { + f16vec4 eMax = f16vec4(-60000.0), yMax = f16vec4(0.0); + line = line0; + for (int r = 0; r < 4; ++r) + { + const int bs = int(roundEven(float(line))) - 1; + const f16vec4 dA = float16_t(a0 + r) - subA; + f16vec4 db = float16_t(bs) - subB; + const f16vec4 cross = dA * e2; + f16vec4 lin = fma(db, cross, (dA * dA) * e0); + const int ti = (a0 + r) * stepA + bs * stepB; + for (int j = 0; j < 3; ++j) + { + const f16vec4 E = fma(db * db, f16vec4(e1), lin); + const bvec4 gt = not(greaterThan(eMax, E)); // a later tap wins ties + eMax = mix(eMax, E, gt); yMax = mix(yMax, f16vec4(tileC[ti + j * stepB].x), gt); + db += float16_t(1.0); lin += cross; + } + line += slope; + } + F = mix(F, vec4(yMax), lessThan(eMax, f16vec4(-12.0))); + } + // The reference holds the filtered current luma in half precision: the blend at the end of the stage reads + // it as an f16 value widened to f32 (dlss-0046 converts every lane with cvt.rn.f16.f32 before the + // `fma bl, F - Hq, Hq` that forms N), and the normalisation and the low-exponent fallback also produce + // halves. Keeping it f32 pulls N a fraction of a half code towards the current luma wherever the history + // and the current frame differ, which the recurrent history then carries into the next frames. + F = roundHalfValue(F); + // ---------------------------------------------------------------- 5. blend + f16vec4 conf; + { + const f16vec4 dA = roundEven(subA) - subA; + const f16vec4 ln = fma(dA, f16vec4(slope), subB); + const f16vec4 dB = abs(roundEven(ln) - ln); + const float16_t kA = float16_t(alongX ? p.kB : p.kA), kB = float16_t(alongX ? p.kA : p.kB); + const f16vec4 sA = clamp(fma(abs(dA), f16vec4(kA), f16vec4(4.5)), float16_t(0.0), float16_t(1.0)); + const f16vec4 sB = clamp(fma(dB, f16vec4(kB), f16vec4(4.5)), float16_t(0.0), float16_t(1.0)); + conf = fma(sA, f16vec4(e5), f16vec4(e6)) * sB; + } + const float16_t hv = float16_t(f343), hv2 = float16_t(0.1) + float16_t(0.9) * hv; + const float16_t lo = (hv - float16_t(p.c184)) * float16_t(p.c186), hi = (hv2 - float16_t(p.c184)) * float16_t(p.c186); + const float16_t avg = float16_t(0.25) * ((conf.x + conf.y) + (conf.z + conf.w)); + const float alphaC = float(clamp(fma(avg, hv2 - hi, hi), float16_t(0.0), float16_t(1.0))); + const vec4 bl = vec4(clamp(fma(conf, f16vec4(hv - lo), f16vec4(lo)), float16_t(0.0), float16_t(1.0))); + if (!(abs(cur.x * F.x * F.y * F.z * F.w) < uintBitsToFloat(0x7F800000u))) { cur.xyz = vec3(0.0); F = vec4(0.0); } + // a single-rounded fma, as the reference: fma.rn.ftz.f32 N, bl, F - Hq, Hq + const vec4 N = fma(bl, F - Hq, Hq); + // Native averages the four f32 blend registers, before the encoded half-surface stores. + { + const vec4 enc = mix(fma(N - 0.99, vec4(1000.0), vec4(0.99)), N, lessThan(N, vec4(0.99))); + const u16vec4 b = float16BitsToUint16(f16vec4(enc)); + imageStore(hiresOut, pix * 2, uvec4(b.x)); imageStore(hiresOut, pix * 2 + ivec2(1, 0), uvec4(b.y)); + imageStore(hiresOut, pix * 2 + ivec2(0, 1), uvec4(b.z)); imageStore(hiresOut, pix * 2 + ivec2(1, 1), uvec4(b.w)); + } + // the reference fuses the history product: fma(hc, 1 - alphaC, alphaC * cur) + precise const vec3 C = fma(hc, vec3(1.0 - alphaC), alphaC * cur.xyz); + precise const float Nm = 0.25 * (((N.x + N.y) + N.z) + N.w); + const vec2 chroma = C.yz / max(C.x, 1e-5) * Nm; + { + // history colour: chroma limited so that the RGB it stands for stays non-negative + const float Y = max(Nm, 0.0); + float lim = Y; + if (chroma.y - 1e-5 < -Y) lim = 1e-5 - chroma.y; + if (chroma.x - chroma.y - 1e-5 < -lim) lim = chroma.y - chroma.x + 1e-5; + if (-chroma.x - chroma.y - 1e-5 < -lim) lim = chroma.x + chroma.y + 1e-5; + const float sc = lim <= Y ? 1.0 : Y / lim; + imageStore(histOut, pix, vec4(Y, chroma * sc, clamp(max(f344 / (f344 + 1.0), p.minAlpha), 0.0, 1.0))); + } + if (pix.x < p.outSize.x && pix.y < p.outSize.y) + { + // the reference multiplies by ginv(Nm) first and divides by the exposure afterwards; the other + // association (ginv / exposure first) rounds differently before the RGB reconstruction + const vec3 ycc = vec3(Nm, chroma) * expandK(Nm) / p.expo; + imageStore(finalOut, p.outOrigin + pix, finalOutputValue(vec4(ycc.x + ycc.y - ycc.z, ycc.x + ycc.z, ycc.x - ycc.y - ycc.z, cur.a))); + } + // ---------------------------------------------------------------- per-token feature for the next frame's input stage + if (pix.x < p.tokW && pix.y < p.tokH) + { + const uint tb = uint(pix.y * p.tokW + pix.x) * 40u; + precise vec4 ms; + for (uint s = 0u; s < 4u; ++s) + ms[s] = 0.2 * (float(head.h[tb + s * 10u + 3u]) + float(head.h[tb + s * 10u + 4u]) + float(head.h[tb + s * 10u + 5u]) + + float(head.h[tb + s * 10u + 6u]) + float(head.h[tb + s * 10u + 7u])); + // Match the native scalar FMA chain; dot() permits a different association at half-store boundaries. + precise float norm = fma(ms.x, ms.x, ms.y * ms.y); + norm = fma(ms.z, ms.z, norm); + norm = fma(ms.w, ms.w, norm); + norm += 1.0; + precise const vec4 feature = ms * inversesqrt(norm); + imageStore(tokOut, pix, feature); + } +} diff --git a/engine/output_store.glsl b/engine/output_store.glsl new file mode 100644 index 0000000..6f96680 --- /dev/null +++ b/engine/output_store.glsl @@ -0,0 +1,17 @@ +// RADV's rgba16f image stores truncate before the copy path's UNORM conversion. +// Preserve that half-store + blit result when writing RGB10A2 directly. Below +// binary16's normal range both 10-bit RGB and 2-bit alpha quantize to zero, and +// overflow saturates, so a normal-range mantissa mask suffices for this format. +#ifndef OUTPUT_RGB10 +#define OUTPUT_RGB10 0 +#endif +#if OUTPUT_RGB10 +#define FINAL_OUTPUT_FORMAT rgb10_a2 +vec4 finalOutputValue(vec4 value) { + vec4 rounded = uintBitsToFloat(floatBitsToUint(value) & 0xffffe000u); + return mix(rounded, value, isnan(value)); +} +#else +#define FINAL_OUTPUT_FORMAT rgba16f +#define finalOutputValue(value) (value) +#endif diff --git a/engine/post_m.comp b/engine/post_m.comp new file mode 100644 index 0000000..10ac505 --- /dev/null +++ b/engine/post_m.comp @@ -0,0 +1,236 @@ +#version 460 +// d4r engine, preset M: the reconstruction, written from the algorithm in M_NOTES.md. Every pixel resamples the +// frame's colour and the network's per-pixel outputs with the anisotropic Gaussians the network predicted for the +// 3x3 render texels around it, picks the two most likely of three motion candidates, fetches the previous result +// along them (bicubic) and blends it with the resampled colour. +// +// NGX ships two kernels with this exact body, differing only in the shared render-texel tile: rrlite_post_3_1 +// (10x16, reconstruction/render >= 3) and rrlite_post_3_2 (15x27, 1.5 <= reconstruction/render < 3, sized for the +// worst case of its range, 1.5x). TILE_W/TILE_H select the tile for the engine's own 16x16-pixel block: 10x10 +// covers reconstruction/render >= 3; 11x11 covers render/reconstruction <= 0.45 with jitter in [-0.5, 0.5]; +// 15x15 covers the whole post_3_2 range (render/reconstruction <= 2/3). +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_GOOGLE_include_directive : require +#include "normalized_store.glsl" +// One invocation owns a vertical pixel pair; the dispatch still covers 16x16 reconstruction pixels. +layout(local_size_x = 16, local_size_y = 8) in; +layout(set = 0, binding = 0, std430) readonly buffer Params +{ + ivec2 colOrigin; + ivec2 hiSize, rndSize; + vec2 hi2rnd, rndOffs; // render position of a pixel centre: centre * hi2rnd + rndOffs + vec2 invRnd; + float gain; // of the first gate + uint reset; + float expo, expoPrev; +} p; +layout(set = 0, binding = 1) uniform sampler2D colorTex; // render resolution, linear HDR +layout(set = 0, binding = 2) uniform sampler2D mixTex; // three candidate weights +layout(set = 0, binding = 3) uniform sampler2D gateTex; // three gates +layout(set = 0, binding = 4) uniform sampler2D covTex; // covariance (xx, yy, xy) +layout(set = 0, binding = 5) uniform sampler2D histTex; // previous result, range-compressed; bilinear +layout(set = 0, binding = 6) uniform sampler2D mvATex; +layout(set = 0, binding = 7) uniform sampler2D mvBTex; +layout(set = 0, binding = 8) uniform sampler2D mvCTex; // bilinear +layout(set = 0, binding = 9, rgba16f) uniform writeonly image2D histOut; +layout(set = 0, binding = 10, rgba32f) uniform writeonly image2D halfOut; +layout(set = 0, binding = 11, std430) writeonly buffer Packed { uint w[]; } packed; +// the render-texel tile around the block; the taps of every pixel stay inside it (see the file comment) +#ifndef TILE_W +#define TILE_W 10 +#endif +#ifndef TILE_H +#define TILE_H 10 +#endif +shared f16vec4 tile0[TILE_W * TILE_H], tile1[TILE_W * TILE_H], tile2[TILE_W * TILE_H]; +shared vec3 colSum[128]; // vertical pair first, then the neighbouring column +const float L10 = 3.3219281; // log2(10) +// LDR (NGX's *_ldr_* kernels, the game did not set the HDR flag): display-referred colour, no exposure and no +// range compression anywhere. +#ifndef LDR +#define LDR 0 +#endif +#if LDR +vec3 compress(vec3 v) { return v; } +vec3 expand(vec3 v) { return v; } +const float PACK = 1.0; +#else +vec3 compress(vec3 v) +{ + precise vec3 mapped = fma(v, vec3(100.0), vec3(1.0)); + precise vec3 log10 = log2(max(mapped, vec3(1e-4))) * 0.30103; + precise vec3 compressed = log10 * 0.36; + return compressed; +} +vec3 expand(vec3 v) +{ + precise vec3 scaled = clamp(v, 1e-4, 3.0) * (1.0 / 0.36); + precise vec3 exponent = scaled * L10; + return (exp2(exponent) - 1.0) * (1.0 / 100.0); +} +const float PACK = 0.4081632490158081; +#endif +ivec2 colourTexel(ivec2 t) +{ + return ivec2(t.x > 0 ? min(p.colOrigin.x + t.x, p.colOrigin.x + p.rndSize.x - 1) : p.colOrigin.x, + t.y > 0 ? min(p.colOrigin.y + t.y, p.colOrigin.y + p.rndSize.y - 1) : p.colOrigin.y); +} +// the previous result at q (pixels), five bilinear taps of a Catmull-Rom filter +vec3 history(vec2 q) +{ + const vec2 inv = 1.0 / vec2(p.hiSize); + const vec2 base = floor(q - 0.5) + 0.5, t = clamp(q - base, 0.0, 1.0), t2 = t * t, t3 = t * t2; + const vec2 w0 = t2 - (t + t3) * 0.5, w1 = (t3 * 1.5 - t2 * 2.5) + 1.0, w3 = (t3 - t2) * 0.5, w2 = ((1.0 - w0) - w1) - w3; + const vec2 w12 = w1 + w2, mid = inv * (base + w2 / w12), lo = inv * (base - 1.0), hi = inv * (base + 2.0); + const float a = w0.x * w12.y, b = w0.y * w12.x, c = w12.x * w12.y, d = w3.y * w12.x, e = w3.x * w12.y; + const vec3 s = fma(textureLod(histTex, vec2(hi.x, mid.y), 0.0).rgb, vec3(e), fma(textureLod(histTex, vec2(mid.x, hi.y), 0.0).rgb, vec3(d), + fma(textureLod(histTex, mid, 0.0).rgb, vec3(c), fma(textureLod(histTex, vec2(mid.x, lo.y), 0.0).rgb, vec3(b), + textureLod(histTex, vec2(lo.x, mid.y), 0.0).rgb * a)))); + return s * (1.0 / (e + (d + (c + (b + a))))); +} +void main() +{ + const ivec2 lid = ivec2(gl_LocalInvocationID.xy), gid = ivec2(gl_GlobalInvocationID.xy); + const ivec2 origin = ivec2(p.hi2rnd * vec2(16 * ivec2(gl_WorkGroupID.xy))) - 2; + const float e1 = p.expo == 0.0 ? 1.0 : p.expo, e2 = p.expoPrev == 0.0 ? 1.0 : p.expoPrev; + const int li = lid.y * 16 + lid.x; +#if TILE_W * TILE_H > 256 +#error "the tile fill assumes at most two tile entries per thread" +#endif + // 128 invocations cover the tile in at most two entries each (TILE_W*TILE_H <= 225 here) + for (int e = 0; e < 2; ++e) + { + const int te = e * 128 + li; + if (te < TILE_W * TILE_H) + { + const ivec2 rel = origin + ivec2(te % TILE_W, te / TILE_W); + const ivec2 t = ivec2(rel.x > 0 ? min(rel.x, p.rndSize.x - 1) : 0, rel.y > 0 ? min(rel.y, p.rndSize.y - 1) : 0); + const vec3 c = e1 * texelFetch(colorTex, colourTexel(rel), 0).rgb; + vec3 l; +#if LDR + for (int k = 0; k < 3; ++k) l[k] = isinf(c[k]) || isnan(c[k]) ? 0.0 : min(c[k], 32752.0); +#else + for (int k = 0; k < 3; ++k) l[k] = max(isinf(c[k]) || isnan(c[k]) ? 1.0 : fma(min(c[k], 32752.0), 100.0, 1.0), 1e-4); +#endif + const vec4 m = texelFetch(mixTex, t, 0), g = texelFetch(gateTex, t, 0), cov = texelFetch(covTex, t, 0); +#if LDR + tile0[te] = f16vec4(vec4(l, m.x)); +#else + precise vec3 log10 = log2(l) * 0.30103; + precise vec3 compressed = log10 * 0.36; + tile0[te] = f16vec4(vec4(compressed, m.x)); +#endif + tile1[te] = f16vec4(m.y, m.z, g.x, g.y); + tile2[te] = f16vec4(vec4(g.z, cov.x * -0.721325, cov.y * -0.721325, cov.z * -1.44265)); + } + } + memoryBarrierShared(); barrier(); + // the two rows of the original 16x16 workgroup's row pair; y0 is even, so the odd row keeps the even row's + // three render texel rows and its own-cell residual rounding exactly as in the 16x16 shader + const int y0 = 2 * gid.y; + const bool inside0 = gid.x < p.hiSize.x && y0 < p.hiSize.y; + vec3 outA0 = vec3(0.0), outA1 = vec3(0.0); + for (int r = 0; r < 2; ++r) + { + const ivec2 pix = ivec2(gid.x, y0 + r); + vec3 outA = vec3(0.0); + const bool inside = pix.x < p.hiSize.x && pix.y < p.hiSize.y; + if (inside) + { + const vec2 centre = vec2(pix) + 0.5; + const vec2 pos = fma(centre, p.hi2rnd, p.rndOffs); + // NVIDIA's kernel handles rows in pairs and gives the odd row the even row's three texel rows + const float evenY = fma(float(pix.y & ~1) + 0.5, p.hi2rnd.y, p.rndOffs.y); + const ivec2 near = ivec2(floor(pos)), base = ivec2(floor(pos.x - 0.5), floor(evenY - 0.5)); + f16vec2 d0 = f16vec2((vec2(base) + 0.5) - pos); + if ((pix.y & 1) != 0) + { + // Native paired rows round the odd row's own-cell residual first, then adjust it in f16. + // Casting the even-cell residual once changes Gaussian weights at half-rounding boundaries. + const int ownBaseY = int(floor(pos.y - 0.5)); + if (ownBaseY != base.y) + d0.y = float16_t((float(ownBaseY) + 0.5) - pos.y) - float16_t(1.0); + } + f16vec4 a0 = f16vec4(0.0), a1 = f16vec4(0.0); + float16_t a2 = float16_t(0.0); + float W = 0.0; + for (int i = -1; i <= 1; ++i) + { + const float16_t dx = d0.x + float16_t(i), sq = dx * dx; + for (int j = -1; j <= 1; ++j) + { + const int k = (base.y - origin.y + j) * TILE_W + (base.x - origin.x + i); + const f16vec4 t2 = tile2[k]; + const float16_t dy = d0.y + float16_t(j); + precise float16_t e = (sq * t2.y + (dy * dy) * t2.z) + (dx * dy) * t2.w; + const float w = exp2(float(e)); + const float16_t wh = float16_t(w); + W += w; + a0 = fma(f16vec4(wh), tile0[k], a0); a1 = fma(f16vec4(wh), tile1[k], a1); a2 = fma(wh, t2.x, a2); + } + } + const float inv = 1.0 / W; + const vec3 cur = inv * vec3(a0.rgb), m = inv * vec3(a0.w, a1.x, a1.y), g = inv * vec3(a1.z, a1.w, a2); + // the two most likely motion candidates + const vec2 mvA = texelFetch(mvATex, near, 0).xy, mvB = texelFetch(mvBTex, near, 0).xy, mvC = textureLod(mvCTex, pos * p.invRnd, 0.0).xy; + vec3 big = vec3(mvA, m.x), small = vec3(mvC, m.y); + if (big.z < small.z) { const vec3 t = big; big = small; small = t; } + vec3 second = !(small.z < m.z) ? small : vec3(mvB, m.z); + if (big.z < second.z) { const vec3 t = big; big = second; second = t; } + const float k1 = clamp(fma(big.z, 1.1, -0.05), 0.0, 1.0), k2 = clamp(fma(second.z, 1.1, -0.05), 0.0, 1.0); + const float kn = 1.0 / max(k1 + k2, 1e-10); + float b2 = k2 * kn, b1 = k1 * kn; + if (b2 < (1.0 / 255.0)) { b2 = 0.0; b1 = 1.0; } + const vec2 size = vec2(p.hiSize); + const vec2 q1 = centre + big.xy; + // the fallback colour is only fetched where the history sample leaves the image or the history was reset: + // the value is the same either way, so computing it in the branch saves the fetch on ordinary pixels + const bool out1 = q1.x < 0.0 || q1.y < 0.0 || q1.x > size.x || q1.y > size.y || p.reset != 0u; + const vec3 h1 = out1 ? compress(e2 * texelFetch(colorTex, colourTexel(near), 0).rgb) : history(q1); + vec3 h2 = vec3(0.0); + if (b2 > 0.0) + { + const vec2 q2 = centre + second.xy; + if (q2.x < 0.0 || q2.y < 0.0 || q2.x > size.x || q2.y > size.y || p.reset != 0u) + h2 = out1 ? h1 : compress(e2 * texelFetch(colorTex, colourTexel(near), 0).rgb); + else h2 = history(q2); + } + const vec3 hist = e1 * (expand(fma(vec3(b1), h1, b2 * h2)) * (1.0 / e2)); + const vec3 lin = expand(cur); + const float g0 = clamp(p.gain * W * g.x, 0.0, 1.0); + const float ga = fma(g0, 1.0 - g.y, g.y), gb = fma(g0, 1.0 - g.z, g.z); + outA = fma(lin - hist, vec3(ga), hist); + if (isinf(outA.x + outA.y + outA.z) || isnan(outA.x + outA.y + outA.z)) outA = lin; + vec3 outB = fma(lin - hist, vec3(gb), hist); + if (isinf(outB.x + outB.y + outB.z) || isnan(outB.x + outB.y + outB.z)) outB = lin; + imageStore(histOut, pix, vec4(compress(outA), 0.0)); + // precise: the reference rounds compress(outB) * PACK to f32 (mul.ftz.f32) and only then scales it with an + // fma. Without the qualifier RADV folds the constant PACK*2047 into one f32 constant and reassociates the + // chain, which biases the scaled value by up to one ulp and flips the 11-11-10 truncation at ~2e-5 of + // pixels (178 of 8.29M on the captured frame, all one step). + precise vec3 pk = compress(outB) * PACK; + const uint r = uint(min(max(fma(pk.r, 2047.0, 0.5), 0.0), 2047.0)), gg = uint(min(max(fma(pk.g, 2047.0, 0.5), 0.0), 2047.0)), + b = uint(min(max(fma(pk.b, 1023.0, 0.5), 0.0), 1023.0)); + packed.w[pix.y * p.hiSize.x + pix.x] = r | (gg << 11) | (b << 22); + } + if (r == 0) outA0 = outA; else outA1 = outA; + } + // the half-resolution mean of the new history. The reference adds a column's two rows first (the thread's own + // pair) and then the paired column (a shfl.sync.bfly of that sum); the row pair's vertical sum is materialised + // here and the even x lane adds the neighbouring column's sum, so the addition order is the same. + colSum[li] = outA0 + outA1; + memoryBarrierShared(); barrier(); + if (inside0 && (lid.x & 1) == 0) + { + const bool x1 = gid.x + 1 < p.hiSize.x, y1 = y0 + 1 < p.hiSize.y; + vec3 sum = colSum[li]; + float n = y1 ? 2.0 : 1.0; + if (x1) { sum += colSum[li + 1]; n += y1 ? 2.0 : 1.0; } + // precise: the reference keeps compress()'s first fma (mean * 100 + 1) fused. Passing the mean as the + // product chain sum * (1/n) lets RADV evaluate that multiply and add with separate roundings in this + // shader, which lands the mapped value one ulp high and flips the 10-bit code at knife-edge pixels + // (1 of 2.07M on the captured frame). Materialising the mean first keeps the fma fused. + precise vec3 halfMean = sum * (1.0 / n); + imageStore(halfOut, ivec2(gid.x, y0) >> 1, vec4(unorm10Value(compress(halfMean) * PACK), 0.0)); + } +} diff --git a/engine/prepare_frames.py b/engine/prepare_frames.py new file mode 100644 index 0000000..eb0bd2b --- /dev/null +++ b/engine/prepare_frames.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python3 +"""Run ordinary linear HDR frames through the native engine without an NGX capture. + +prepare_frames.py FRAMES.json CASE.txt +JSON has output: [width,height] and frames: [{color, motion, depth, render: [w,h], +bits: 16, color_exposure: 1, history_exposure_ratio: 1, network_exposure_scale: 1, +jitter: [0,0], motion_scale: [0,0], motion_offset: [0,0], sharpness: 1.37, reset: true}]. +Paths are relative to the JSON. Color is interleaved RGBA16F/32F, motion RG16F/32F, +depth R32F (inverse depth). Subsequent frames keep the engine's own history. +The first frame resets. Attention windows default to a fixed four-token shift; +captured NGX window scheduling is available separately through prepare_k.py. +""" +import argparse +import json +import math +from pathlib import Path + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument('frames', type=Path) + p.add_argument('output', type=Path) + a = p.parse_args() + config = json.loads(a.frames.read_text()) + ow, oh = config['output'] + if not all(isinstance(x, int) and 256 <= x <= 8192 for x in (ow,oh)): + raise ValueError('output dimensions must be integers in 256..8192') + frames = config['frames'] + if not frames: raise ValueError('no frames') + lines = [f'output {ow} {oh}'] + for i, frame in enumerate(frames): + rw, rh = frame['render']; bits = frame.get('bits',16) + if not isinstance(rw,int) or not isinstance(rh,int) or not 0 < rw <= ow or not 0 < rh <= oh or bits not in (16,32): + raise ValueError('invalid render dimensions or component size') + for j, (key, channels, b) in enumerate((('color',4,bits),('motion',2,bits),('depth',1,32))): + path = (a.frames.resolve().parent / frame[key]).resolve() + if path.stat().st_size != rw * rh * channels * b // 8: + raise ValueError(f'{path}: unexpected frame image size') + lines.append(f'image {i*3+j} {json.dumps(str(path))} {rw} {rh} {channels} {b}') + reset = frame.get('reset', i == 0) + if not isinstance(reset,bool): raise ValueError('reset must be a boolean') + if i == 0 and not reset: raise ValueError('the first ordinary frame must reset') + exposure = frame.get('color_exposure',1) + ratio = frame.get('history_exposure_ratio',1) + network = frame.get('network_exposure_scale',1) + if not all(math.isfinite(x) and x > 0 for x in (exposure,ratio,network)): + raise ValueError('exposure factors must be finite and positive') + vectors = [frame.get(k,[0,0]) for k in ('jitter','motion_scale','motion_offset')] + if any(len(v) != 2 or not all(math.isfinite(x) for x in v) for v in vectors): + raise ValueError('jitter and motion vectors must contain two finite components') + sharpness = frame.get('sharpness',1.37) + if not math.isfinite(sharpness): raise ValueError('sharpness must be finite') + values = [i*3,i*3+1,i*3+2,exposure,ratio,network,int(reset),*vectors[0],*vectors[1],*vectors[2],sharpness] + lines.append('frame_auto ' + ' '.join(map(str,values))) + a.output.parent.mkdir(parents=True,exist_ok=True) + a.output.write_text('\n'.join(lines)+'\n') + print(f'Prepared {len(frames)} ordinary frames for {ow}x{oh} native inference') + + +if __name__ == '__main__': main() diff --git a/engine/prepare_k.py b/engine/prepare_k.py new file mode 100644 index 0000000..a3e0354 --- /dev/null +++ b/engine/prepare_k.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +"""Prepare a full-frame native K replay from a capture. Does not translate kernels. + +prepare_k.py CAPTURE_DIR OUTPUT_DIR [--frames N] [--start FRAME] +The model's reconstruction tables can be packaged using OUTPUT_DIR/lut.bin. +""" +import argparse +import json +from pathlib import Path +import struct +from capture import INPUT, OUTPUT, ORDER, launches, input_params, output_params + + +def quote(path): + return json.dumps(str(Path(path).resolve())) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('capture', type=Path) + parser.add_argument('output', type=Path) + parser.add_argument('--frames', type=int, default=0, help='0 = all captured frames') + parser.add_argument('--start', type=int, default=0) + a = parser.parse_args() + inputs, outputs = launches(a.capture, INPUT), launches(a.capture, OUTPUT) + if not inputs or not outputs or len(outputs) > len(inputs): + raise ValueError('capture must contain complete K input/output frame pairs') + pairs = list(zip(inputs, outputs)) + for i, (ip, op) in enumerate(pairs): + if not ip.sequence < op.sequence or (i+1 < len(inputs) and op.sequence >= inputs[i+1].sequence): + raise ValueError('capture has a missing or out-of-order input/output launch') + if a.start < 0 or a.start >= len(pairs) or a.frames < 0: + raise ValueError('invalid frame selection') + end = len(pairs) if a.frames == 0 else min(len(pairs), a.start + a.frames) + selected = pairs[a.start:end] + # Unselected later frames may be incomplete (for example a capture stopped while writing them). + for ip, op in selected: + for resource in (*ip.resources.values(), *op.resources.values()): resource.validate() + layer_launches = {name: launches(a.capture, 'dltss_pwin_'+name+'_layer') for name in ORDER} + first_in, first_out = selected[0] + geometry = [] + for name in ORDER: + matches = [p for p in layer_launches[name] if first_in.sequence < p.sequence < first_out.sequence] + if len(matches) != 1: + raise ValueError(f'frame {a.start}: expected one {name} launch, found {len(matches)}') + geometry.append((*matches[0].i2(0), *matches[0].i2(56))) + final = first_out.resources[312] + a.output.mkdir(parents=True, exist_ok=True) + (a.output / 'shape.bin').write_bytes(struct.pack('<2I', final.width, final.height) + b''.join(struct.pack('<4i', *s) for s in geometry)) + # NGX's whole allocation: sixteen reconstruction tables, of which a launch uses the one for its upscaling ratio + (a.output / 'lut.bin').write_bytes(first_out.allocation(280, 16 * 32768)) + lines = [f'shape {quote(a.output / "shape.bin")}', + f'history {quote(first_in.resources[176].path)} {quote(first_out.resources[232].path)} {quote(first_in.resources[224].path)}'] + resource_ids = {} + references = [] + for number, (ip, op) in enumerate(selected): + if not ip.sequence < op.sequence or (number and ip.sequence <= selected[number-1][1].sequence): + raise ValueError('capture frames are not in network order') + if op.resources[312].width != final.width or op.resources[312].height != final.height: + raise ValueError('resolution changes require another Engine') + frame_geometry = [] + for name, expected in zip(ORDER, geometry): + matches = [p for p in layer_launches[name] if ip.sequence < p.sequence < op.sequence] + if len(matches) != 1 or matches[0].i2(0) != expected[:2]: + raise ValueError(f'frame {number}: missing or changed {name} geometry') + frame_geometry.append((*matches[0].i2(0), *matches[0].i2(56))) + ids = [] + for off in (168, 184, 192): + r = ip.resources[off] + if r.path not in resource_ids: + i = len(resource_ids); resource_ids[r.path] = i + lines.append(f'image {i} {quote(r.path)} {r.width} {r.height} {r.channels} {r.bits}') + ids.append(resource_ids[r.path]) + inp, outp = a.output / f'frame-{number}.input', a.output / f'frame-{number}.output' + inp.write_bytes(input_params(ip)); outp.write_bytes(output_params(op)) + windows = a.output / f'frame-{number}.windows' + windows.write_bytes(b''.join(struct.pack('<4i', *s) for s in frame_geometry)) + lines.append(f'frame {quote(inp)} {quote(outp)} {quote(windows)} ' + ' '.join(map(str, ids))) + references.append({str(off): str(op.resources[off].path) for off in (288,296,304,312)}) + (a.output / 'case.txt').write_text('\n'.join(lines) + '\n') + (a.output / 'references.json').write_text(json.dumps(references, indent=2) + '\n') + print(f'Prepared {len(selected)} frames, {final.width}x{final.height}, token grid {geometry[0][0]}x{geometry[0][1]}') + + +if __name__ == '__main__': main() diff --git a/engine/replay.cpp b/engine/replay.cpp new file mode 100644 index 0000000..cacc2e0 --- /dev/null +++ b/engine/replay.cpp @@ -0,0 +1,494 @@ +// Full-frame native K validation and GPU timing. No CUDA/NGX/PTX execution. +// With the optional trailing `hip` argument (and a model built by compile_k.py --hip) the eleven layers run on +// the native HIP kernels (hip_net.h) between the engine's Vulkan stages, whose command buffer is split around +// them through two buffers shared by file descriptor. The engine's own Vulkan network is the default. +#include "runtime.h" +#ifdef D4R_ENGINE_HIP +#include "hip_net.h" +#endif +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { +void check(VkResult r, const char* operation) { + if (r != VK_SUCCESS) throw std::runtime_error(std::string(operation) + ": VkResult " + std::to_string(r)); +} +#define CK(call) check((call), #call) +std::vector read(const std::string& file) { + std::ifstream stream(file, std::ios::binary | std::ios::ate); + if (!stream) throw std::runtime_error("cannot read " + file); + auto bytes = stream.tellg(); + if (bytes <= 0 || bytes > 1024ll * 1024 * 1024) throw std::runtime_error("invalid file size: " + file); + std::vector data(static_cast(bytes)); stream.seekg(0); + if (!stream.read(data.data(), bytes)) throw std::runtime_error("truncated file: " + file); + return data; +} +template T structure(const std::string& file) { + auto data = read(file); + if (data.size() != sizeof(T)) throw std::runtime_error("invalid structure size: " + file); + T result; std::memcpy(&result, data.data(), sizeof(result)); return result; +} +struct Context { + VkInstance instance = VK_NULL_HANDLE; + VkPhysicalDevice physical = VK_NULL_HANDLE; + VkDevice device = VK_NULL_HANDLE; + VkQueue queue = VK_NULL_HANDLE; + VkCommandPool pool = VK_NULL_HANDLE; + VkCommandBuffer command = VK_NULL_HANDLE; + VkFence fence = VK_NULL_HANDLE; + VkQueryPool timestamps = VK_NULL_HANDLE; + VkPipelineCache cache = VK_NULL_HANDLE; + VkPhysicalDeviceProperties properties{}; + VkPhysicalDeviceMemoryProperties memory{}; + uint32_t timestampBits = 64; + ~Context() { + if (device) vkDeviceWaitIdle(device); + if (cache) vkDestroyPipelineCache(device, cache, nullptr); + if (timestamps) vkDestroyQueryPool(device, timestamps, nullptr); + if (fence) vkDestroyFence(device, fence, nullptr); + if (pool) vkDestroyCommandPool(device, pool, nullptr); + if (device) vkDestroyDevice(device, nullptr); + if (instance) vkDestroyInstance(instance, nullptr); + } + void create(uint32_t queryCount = 2, bool fp8 = false, bool externalMemory = false) { + VkApplicationInfo application{VK_STRUCTURE_TYPE_APPLICATION_INFO}; + application.pApplicationName = "d4r native inference"; application.apiVersion = VK_API_VERSION_1_3; + VkInstanceCreateInfo info{VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO}; info.pApplicationInfo = &application; + CK(vkCreateInstance(&info, nullptr, &instance)); + uint32_t count = 0; CK(vkEnumeratePhysicalDevices(instance, &count, nullptr)); + std::vector devices(count); CK(vkEnumeratePhysicalDevices(instance, &count, devices.data())); + if (devices.empty()) throw std::runtime_error("no Vulkan GPU"); + for (auto gpu : devices) { + VkPhysicalDeviceProperties p; vkGetPhysicalDeviceProperties(gpu, &p); + if (!physical || p.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU) physical = gpu; + } + vkGetPhysicalDeviceProperties(physical, &properties); vkGetPhysicalDeviceMemoryProperties(physical, &memory); + if (properties.apiVersion < VK_API_VERSION_1_3) throw std::runtime_error("Vulkan 1.3 required"); + uint32_t extensionCount = 0; CK(vkEnumerateDeviceExtensionProperties(physical, nullptr, &extensionCount, nullptr)); + std::vector extensions(extensionCount); + CK(vkEnumerateDeviceExtensionProperties(physical, nullptr, &extensionCount, extensions.data())); + if (std::none_of(extensions.begin(), extensions.end(), [](const auto& e) { return !std::strcmp(e.extensionName, VK_KHR_COOPERATIVE_MATRIX_EXTENSION_NAME); })) + throw std::runtime_error("VK_KHR_cooperative_matrix required"); + if (fp8 && std::none_of(extensions.begin(), extensions.end(), [](const auto& e) { return !std::strcmp(e.extensionName, VK_EXT_SHADER_FLOAT8_EXTENSION_NAME); })) + throw std::runtime_error("VK_EXT_shader_float8 required for M"); + if (externalMemory && std::none_of(extensions.begin(), extensions.end(), [](const auto& e) { return !std::strcmp(e.extensionName, VK_KHR_EXTERNAL_MEMORY_FD_EXTENSION_NAME); })) + throw std::runtime_error("VK_KHR_external_memory_fd required for the native HIP network"); + VkPhysicalDeviceShaderFloat8FeaturesEXT float8{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_FLOAT8_FEATURES_EXT}; + VkPhysicalDeviceCooperativeMatrixFeaturesKHR matrix{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR, fp8 ? &float8 : nullptr}; + VkPhysicalDeviceVulkan13Features f13{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_FEATURES, &matrix}; + VkPhysicalDeviceVulkan12Features f12{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES, &f13}; + VkPhysicalDeviceVulkan11Features f11{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES, &f12}; + VkPhysicalDeviceFeatures2 features{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2, &f11}; + vkGetPhysicalDeviceFeatures2(physical, &features); + if (fp8 && (!float8.shaderFloat8 || !float8.shaderFloat8CooperativeMatrix)) + throw std::runtime_error("GPU lacks FP8 cooperative matrices"); + if (!matrix.cooperativeMatrix || !features.features.shaderInt16 || !features.features.shaderStorageImageExtendedFormats || !f11.storageBuffer16BitAccess || !f12.shaderFloat16 || !f13.subgroupSizeControl || !f13.computeFullSubgroups) + throw std::runtime_error("GPU lacks f16 storage/arithmetic, cooperative matrix or subgroup control"); + auto matrixProperties = reinterpret_cast(vkGetInstanceProcAddr(instance, "vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR")); + if (!matrixProperties) throw std::runtime_error("cooperative matrix query unavailable"); + uint32_t matrixCount = 0; CK(matrixProperties(physical, &matrixCount, nullptr)); + std::vector matrices(matrixCount, {VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_KHR}); + CK(matrixProperties(physical, &matrixCount, matrices.data())); + if (std::none_of(matrices.begin(), matrices.end(), [](const auto& p) { + return p.MSize == 16 && p.NSize == 16 && p.KSize == 16 && p.scope == VK_SCOPE_SUBGROUP_KHR && + p.AType == VK_COMPONENT_TYPE_FLOAT16_KHR && p.BType == VK_COMPONENT_TYPE_FLOAT16_KHR && + p.CType == VK_COMPONENT_TYPE_FLOAT32_KHR && p.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR; + })) throw std::runtime_error("16x16x16 f16 -> f32 cooperative matrices required"); + VkPhysicalDeviceSubgroupSizeControlProperties subgroup{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SUBGROUP_SIZE_CONTROL_PROPERTIES}; + VkPhysicalDeviceProperties2 p2{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2, &subgroup}; + vkGetPhysicalDeviceProperties2(physical, &p2); + if (subgroup.minSubgroupSize > 32 || subgroup.maxSubgroupSize < 64 || !(subgroup.requiredSubgroupSizeStages & VK_SHADER_STAGE_COMPUTE_BIT)) + throw std::runtime_error("32/64-lane compute subgroups required"); + // Enable only the features this engine needs, rather than every supported feature. + features.features = {}; features.features.shaderInt16 = VK_TRUE; + features.features.shaderStorageImageExtendedFormats = VK_TRUE; + f11 = {VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES, &f12}; f11.storageBuffer16BitAccess = VK_TRUE; + f12 = {VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES, &f13}; f12.shaderFloat16 = VK_TRUE; + f13 = {VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_FEATURES, &matrix}; f13.subgroupSizeControl = f13.computeFullSubgroups = VK_TRUE; + matrix = {VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR, fp8 ? &float8 : nullptr}; matrix.cooperativeMatrix = VK_TRUE; + if (fp8) { + float8 = {VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_FLOAT8_FEATURES_EXT}; + float8.shaderFloat8 = float8.shaderFloat8CooperativeMatrix = VK_TRUE; + } + uint32_t queueCount = 0; vkGetPhysicalDeviceQueueFamilyProperties(physical, &queueCount, nullptr); + std::vector families(queueCount); vkGetPhysicalDeviceQueueFamilyProperties(physical, &queueCount, families.data()); + uint32_t family = queueCount; + for (uint32_t i = 0; i < queueCount; ++i) if ((families[i].queueFlags & VK_QUEUE_COMPUTE_BIT) && families[i].timestampValidBits) { family = i; break; } + if (family == queueCount) throw std::runtime_error("no timestamp-capable compute queue"); + timestampBits = families[family].timestampValidBits; + float priority = 1; + VkDeviceQueueCreateInfo qi{VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO}; qi.queueFamilyIndex = family; qi.queueCount = 1; qi.pQueuePriorities = &priority; + const char* enabled[3] = {VK_KHR_COOPERATIVE_MATRIX_EXTENSION_NAME}; + uint32_t enabledCount = 1; + if (fp8) enabled[enabledCount++] = VK_EXT_SHADER_FLOAT8_EXTENSION_NAME; + if (externalMemory) enabled[enabledCount++] = VK_KHR_EXTERNAL_MEMORY_FD_EXTENSION_NAME; + VkDeviceCreateInfo di{VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO, &features}; + di.queueCreateInfoCount = 1; di.pQueueCreateInfos = &qi; di.enabledExtensionCount = enabledCount; + di.ppEnabledExtensionNames = enabled; + // Diagnostic: what a host device that never asked for cooperative matrices would give the shaders. + if (std::getenv("D4R_ENGINE_TEST_NO_COOPMAT")) { f13.pNext = nullptr; di.enabledExtensionCount = 0; } + CK(vkCreateDevice(physical, &di, nullptr, &device)); vkGetDeviceQueue(device, family, 0, &queue); + VkCommandPoolCreateInfo pi{VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO}; pi.queueFamilyIndex = family; pi.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT; + CK(vkCreateCommandPool(device, &pi, nullptr, &pool)); + VkCommandBufferAllocateInfo ci{VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO}; ci.commandPool = pool; ci.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; ci.commandBufferCount = 1; + CK(vkAllocateCommandBuffers(device, &ci, &command)); + VkFenceCreateInfo fi{VK_STRUCTURE_TYPE_FENCE_CREATE_INFO}; CK(vkCreateFence(device, &fi, nullptr, &fence)); + VkQueryPoolCreateInfo ti{VK_STRUCTURE_TYPE_QUERY_POOL_CREATE_INFO}; ti.queryType = VK_QUERY_TYPE_TIMESTAMP; ti.queryCount = queryCount; + CK(vkCreateQueryPool(device, &ti, nullptr, ×tamps)); + VkPipelineCacheCreateInfo pci{VK_STRUCTURE_TYPE_PIPELINE_CACHE_CREATE_INFO}; CK(vkCreatePipelineCache(device, &pci, nullptr, &cache)); + std::cerr << "GPU: " << properties.deviceName << '\n'; + } + VkDeviceMemory allocate(VkMemoryRequirements requirements, VkMemoryPropertyFlags flags) { + for (uint32_t i = 0; i < memory.memoryTypeCount; ++i) if ((requirements.memoryTypeBits & (1u << i)) && (memory.memoryTypes[i].propertyFlags & flags) == flags) { + VkMemoryAllocateInfo info{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO}; info.allocationSize = requirements.size; info.memoryTypeIndex = i; + VkDeviceMemory result; CK(vkAllocateMemory(device, &info, nullptr, &result)); return result; + } + throw std::runtime_error("no compatible memory type"); + } + void begin() { + CK(vkResetCommandBuffer(command, 0)); + VkCommandBufferBeginInfo info{VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO}; info.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT; + CK(vkBeginCommandBuffer(command, &info)); + } + void submit() { + CK(vkEndCommandBuffer(command)); CK(vkResetFences(device, 1, &fence)); + VkSubmitInfo info{VK_STRUCTURE_TYPE_SUBMIT_INFO}; info.commandBufferCount = 1; info.pCommandBuffers = &command; + CK(vkQueueSubmit(queue, 1, &info, fence)); CK(vkWaitForFences(device, 1, &fence, VK_TRUE, UINT64_MAX)); + } + double elapsed(uint64_t begin, uint64_t end) const { + const uint64_t mask = timestampBits == 64 ? UINT64_MAX : (uint64_t(1) << timestampBits) - 1; + return ((end - begin) & mask) * double(properties.limits.timestampPeriod) / 1000.0; + } + double gpuMicros() { + uint64_t times[2]; CK(vkGetQueryPoolResults(device, timestamps, 0, 2, sizeof(times), times, sizeof(uint64_t), VK_QUERY_RESULT_64_BIT | VK_QUERY_RESULT_WAIT_BIT)); + return elapsed(times[0], times[1]); + } +}; +struct Buffer { + Context& context; + VkBuffer handle = VK_NULL_HANDLE; + VkDeviceMemory memory = VK_NULL_HANDLE; + void* map = nullptr; + VkDeviceSize size; + Buffer(Context& c, VkDeviceSize bytes) : context(c), size(bytes) { + VkBufferCreateInfo info{VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO}; info.size = bytes; info.usage = VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT; + CK(vkCreateBuffer(c.device, &info, nullptr, &handle)); + VkMemoryRequirements requirements; vkGetBufferMemoryRequirements(c.device, handle, &requirements); + memory = c.allocate(requirements, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT); + CK(vkBindBufferMemory(c.device, handle, memory, 0)); CK(vkMapMemory(c.device, memory, 0, VK_WHOLE_SIZE, 0, &map)); + } + ~Buffer() { + if (map) vkUnmapMemory(context.device, memory); + if (handle) vkDestroyBuffer(context.device, handle, nullptr); + if (memory) vkFreeMemory(context.device, memory, nullptr); + } + void load(const std::string& file) { + auto bytes = read(file); if (bytes.size() != size) throw std::runtime_error("unexpected image bytes: " + file); + std::memcpy(map, bytes.data(), bytes.size()); + } + void save(const std::string& file) { + std::ofstream stream(file, std::ios::binary); if (!stream.write(static_cast(map), size)) throw std::runtime_error("cannot write " + file); + } +}; +#ifdef D4R_ENGINE_HIP +// A storage buffer whose memory another API (HIP) imports through a file descriptor, for the native network. +struct Shared { VkBuffer buffer{}; VkDeviceMemory memory{}; VkDeviceSize allocation = 0; int fd = -1; }; +static Shared shared(const Context& c, VkDeviceSize size) { + Shared s; + VkExternalMemoryBufferCreateInfo ei{VK_STRUCTURE_TYPE_EXTERNAL_MEMORY_BUFFER_CREATE_INFO}; + ei.handleTypes = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT; + VkBufferCreateInfo bi{VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO, &ei}; bi.size = size; bi.usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT; + CK(vkCreateBuffer(c.device, &bi, nullptr, &s.buffer)); + VkMemoryRequirements r; vkGetBufferMemoryRequirements(c.device, s.buffer, &r); + VkExportMemoryAllocateInfo xi{VK_STRUCTURE_TYPE_EXPORT_MEMORY_ALLOCATE_INFO}; + xi.handleTypes = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT; + VkMemoryAllocateInfo ai{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO, &xi}; ai.allocationSize = s.allocation = r.size; + uint32_t index = c.memory.memoryTypeCount; + for (uint32_t i = 0; i < c.memory.memoryTypeCount; ++i) + if ((r.memoryTypeBits & (1u << i)) && (c.memory.memoryTypes[i].propertyFlags & VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT)) { index = i; break; } + if (index == c.memory.memoryTypeCount) throw std::runtime_error("no device-local memory type"); + ai.memoryTypeIndex = index; + CK(vkAllocateMemory(c.device, &ai, nullptr, &s.memory)); CK(vkBindBufferMemory(c.device, s.buffer, s.memory, 0)); + VkMemoryGetFdInfoKHR gi{VK_STRUCTURE_TYPE_MEMORY_GET_FD_INFO_KHR}; + gi.memory = s.memory; gi.handleType = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT; + const auto getFd = reinterpret_cast(vkGetDeviceProcAddr(c.device, "vkGetMemoryFdKHR")); + if (!getFd) throw std::runtime_error("vkGetMemoryFdKHR is unavailable"); + CK(getFd(c.device, &gi, &s.fd)); + return s; +} +#endif +void barrier(VkCommandBuffer cb, VkPipelineStageFlags source, VkPipelineStageFlags dest, VkAccessFlags src, VkAccessFlags dst) { + VkMemoryBarrier b{VK_STRUCTURE_TYPE_MEMORY_BARRIER}; b.srcAccessMask = src; b.dstAccessMask = dst; + vkCmdPipelineBarrier(cb, source, dest, 0, 1, &b, 0, nullptr, 0, nullptr); +} +VkBufferImageCopy region(uint32_t w, uint32_t h) { + VkBufferImageCopy r{}; r.imageSubresource = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1}; r.imageExtent = {w, h, 1}; return r; +} +struct Image { + Context& context; + VkImage handle = VK_NULL_HANDLE; + VkImageView view = VK_NULL_HANDLE; + VkDeviceMemory memory = VK_NULL_HANDLE; + uint32_t w, h; + Buffer upload; + Image(Context& c, const std::string& file, uint32_t width, uint32_t height, int channels, int bits) + : context(c), w(width), h(height), upload(c, VkDeviceSize(width) * height * channels * bits / 8) { + if (!w || !h || w > c.properties.limits.maxImageDimension2D || h > c.properties.limits.maxImageDimension2D || + (channels != 1 && channels != 2 && channels != 4) || (bits != 16 && bits != 32)) throw std::runtime_error("unsupported frame image format/dimensions"); + VkFormat format = bits == 16 ? (channels == 1 ? VK_FORMAT_R16_SFLOAT : channels == 2 ? VK_FORMAT_R16G16_SFLOAT : VK_FORMAT_R16G16B16A16_SFLOAT) + : (channels == 1 ? VK_FORMAT_R32_SFLOAT : channels == 2 ? VK_FORMAT_R32G32_SFLOAT : VK_FORMAT_R32G32B32A32_SFLOAT); + VkImageCreateInfo info{VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO}; info.imageType = VK_IMAGE_TYPE_2D; info.format = format; + info.extent = {w,h,1}; info.mipLevels = info.arrayLayers = 1; info.samples = VK_SAMPLE_COUNT_1_BIT; + info.usage = VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT; + CK(vkCreateImage(c.device, &info, nullptr, &handle)); + VkMemoryRequirements requirements; vkGetImageMemoryRequirements(c.device, handle, &requirements); + memory = c.allocate(requirements, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT); CK(vkBindImageMemory(c.device, handle, memory, 0)); + VkImageViewCreateInfo vi{VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO}; vi.image = handle; vi.viewType = VK_IMAGE_VIEW_TYPE_2D; vi.format = format; + vi.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT,0,1,0,1}; CK(vkCreateImageView(c.device, &vi, nullptr, &view)); + upload.load(file); + } + ~Image() { + if (view) vkDestroyImageView(context.device, view, nullptr); + if (handle) vkDestroyImage(context.device, handle, nullptr); + if (memory) vkFreeMemory(context.device, memory, nullptr); + } + void recordUpload() { + VkImageMemoryBarrier b{VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER}; b.image = handle; + b.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED; b.newLayout = VK_IMAGE_LAYOUT_GENERAL; b.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; + b.srcQueueFamilyIndex = b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; b.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT,0,1,0,1}; + vkCmdPipelineBarrier(context.command, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, nullptr, 0, nullptr, 1, &b); + auto r = region(w,h); vkCmdCopyBufferToImage(context.command, upload.handle, handle, VK_IMAGE_LAYOUT_GENERAL, 1, &r); + } +}; +struct Frame { d4r::InputParams input; d4r::OutputParams output; std::array windows; unsigned color, motion, depth; }; + +int run(const std::string& modelDir, const std::string& caseFile, const std::string& outputDir, int repeats, bool hip) { +#ifndef D4R_ENGINE_HIP + if (hip) throw std::runtime_error("this build has no HIP support (rebuild with ROCm installed)"); +#endif + const bool readback = outputDir != "-"; + const bool profile = std::getenv("D4R_ENGINE_PROFILE") && std::string(std::getenv("D4R_ENGINE_PROFILE")) == "1"; + constexpr unsigned boundaries = d4r::KTimingStages + 1; + const uint32_t queryCount = profile ? 2 + boundaries * repeats : 2; + Context context; context.create(queryCount, false, hip); + const std::array stageNames = {"setup", "exposure", "input", + "enc0", "enc1", "enc2", "enc3", "enc4", "dec5", "dec4", "dec3", "dec2", "dec1", "dec0", "output", "finish"}; + struct Marks { VkQueryPool pool; uint32_t first; } marks{context.timestamps, 2}; + d4r::FrameProfiler profiler{[](VkCommandBuffer cb, uint32_t boundary, void* user) { + auto& m = *static_cast(user); + vkCmdWriteTimestamp(cb, VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT, m.pool, m.first + boundary); + }, &marks}; + + d4r::Shape shape{}; + std::array historyFiles; + std::vector> images; + std::vector frames; + std::ifstream manifest(caseFile); + if (!manifest) throw std::runtime_error("cannot open case manifest"); + std::string line; + while (std::getline(manifest,line)) { + std::istringstream s(line); std::string kind; s >> kind; + if (kind == "shape") { + std::string file; s >> std::quoted(file); shape = structure(file); + } else if (kind == "output") { + uint32_t w,h; s >> w >> h; + if (!s) throw std::runtime_error("malformed output dimensions"); + shape = d4r::makeKShape(w,h); + } else if (kind == "history") { + for (auto& file : historyFiles) s >> std::quoted(file); + } else if (kind == "image") { + unsigned id, w, h; int channels, bits; std::string file; + s >> id >> std::quoted(file) >> w >> h >> channels >> bits; + if (!s || id != images.size()) throw std::runtime_error("malformed image declaration"); + images.push_back(std::make_unique(context,file,w,h,channels,bits)); + } else if (kind == "frame") { + Frame f{}; std::string input,output,windows; s >> std::quoted(input) >> std::quoted(output) >> std::quoted(windows) >> f.color >> f.motion >> f.depth; + if (!s) throw std::runtime_error("malformed frame declaration"); + f.input = structure(input); f.output = structure(output); + f.windows = structure>(windows); frames.push_back(f); + } else if (kind == "frame_auto") { + Frame f{}; d4r::FrameSettings settings{}; unsigned reset; + s >> f.color >> f.motion >> f.depth >> settings.colorExposure >> settings.historyExposureRatio >> settings.networkExposureScale >> reset + >> settings.jitter[0] >> settings.jitter[1] >> settings.motionScale[0] >> settings.motionScale[1] + >> settings.motionOffset[0] >> settings.motionOffset[1] >> settings.sharpness; + if (!s || f.color >= images.size() || reset > 1) throw std::runtime_error("malformed standard frame declaration"); + settings.renderWidth = images[f.color]->w; settings.renderHeight = images[f.color]->h; settings.reset = reset != 0; + auto params = d4r::makeKFrameParams(shape,settings); f.input = params.input; f.output = params.output; f.windows = shape.layers; + frames.push_back(f); + } else if (!kind.empty()) throw std::runtime_error("unknown case declaration " + kind); + } + if (frames.empty() || !shape.outputWidth) throw std::runtime_error("incomplete case manifest"); + const bool seedHistory = !historyFiles[0].empty(); + if (!seedHistory && !(frames[0].input.flags & d4r::FrameFlagReset)) throw std::runtime_error("the first frame must reset or provide history"); + for (auto& f : frames) if (f.color >= images.size() || f.motion >= images.size() || f.depth >= images.size()) throw std::runtime_error("invalid frame image id"); + if (hip) { + for (size_t i = 0; i < frames.size(); ++i) { + const auto expected = d4r::makeKWindows(shape, uint32_t(i)); + for (unsigned layer = 0; layer < d4r::KLayers; ++layer) { + const auto& a = frames[i].windows[layer]; + const auto& b = expected[layer]; + if (a.w != b.w || a.h != b.h || a.sx != b.sx || a.sy != b.sy) + throw std::runtime_error("HIP replay requires the native window cycle starting at evaluation zero; use Vulkan for custom or offset windows"); + } + } + } + const auto start = std::chrono::steady_clock::now(); + d4r::Model model(context.physical,context.device,modelDir,context.cache); + d4r::ExternalNetwork external{}; +#ifdef D4R_ENGINE_HIP + std::unique_ptr network; + if (hip) { + // The native K network reads the tokens and writes the head at the output-derived token grid. + Shared tokens = shared(context, d4r::Engine::networkTokenBytes(shape.outputWidth, shape.outputHeight)); + Shared head = shared(context, d4r::Engine::networkHeadBytes(shape.outputWidth, shape.outputHeight)); + external = {tokens.buffer, head.buffer}; + network = std::make_unique(modelDir, shape.outputWidth, shape.outputHeight, tokens.fd, head.fd, + tokens.allocation, head.allocation); + } +#endif + d4r::Engine engine(model,shape,hip ? &external : nullptr); + auto recordFrame = [&](uint32_t frame) { + marks.first = 2 + frame * boundaries; + if (profile) engine.recordFrame(context.command, profiler); + else engine.recordFrame(context.command); + }; + auto stageMicros = [&](uint32_t count) { + std::vector times(boundaries * count); + CK(vkGetQueryPoolResults(context.device, context.timestamps, 2, times.size(), times.size() * sizeof(uint64_t), + times.data(), sizeof(uint64_t), VK_QUERY_RESULT_64_BIT | VK_QUERY_RESULT_WAIT_BIT)); + std::array result{}; + for (uint32_t frame = 0; frame < count; ++frame) + for (unsigned stage = 0; stage < result.size(); ++stage) + result[stage] += context.elapsed(times[frame * boundaries + stage], times[frame * boundaries + stage + 1]) / count; + return result; + }; + std::cerr << "Model/engine setup: " << std::chrono::duration(std::chrono::steady_clock::now()-start).count() + << " ms; activation buffers " << engine.activationBytes() / double(1024*1024) << " MiB\n"; + uint32_t w = shape.outputWidth, h = shape.outputHeight; + auto base = shape.layers[0]; + Buffer seedColor(context,seedHistory ? VkDeviceSize(w)*h*8 : 16), seedLuma(context,seedHistory ? VkDeviceSize(w)*h*8 : 16), seedFeature(context,seedHistory ? VkDeviceSize(base.w)*base.h*8 : 16); + if (seedHistory) { seedColor.load(historyFiles[0]); seedLuma.load(historyFiles[1]); seedFeature.load(historyFiles[2]); } + Buffer result(context,readback ? VkDeviceSize(w)*h*8 : 16), color(context,readback ? VkDeviceSize(w)*h*8 : 16), + luma(context,readback ? VkDeviceSize(w)*h*8 : 16), feature(context,readback ? VkDeviceSize(base.w)*base.h*8 : 16), + head(context,readback ? engine.headBytes() : 16); + context.begin(); model.recordUpload(context.command); engine.recordInitialize(context.command); + for (auto& image : images) image->recordUpload(); + // Seed copies overwrite the initialization clears, so order transfer writes explicitly. + barrier(context.command,VK_PIPELINE_STAGE_TRANSFER_BIT,VK_PIPELINE_STAGE_TRANSFER_BIT,VK_ACCESS_TRANSFER_WRITE_BIT,VK_ACCESS_TRANSFER_WRITE_BIT); + auto r = region(w,h); + if (seedHistory) { + auto seed = engine.initialHistoryImages(); + vkCmdCopyBufferToImage(context.command,seedColor.handle,seed.color,VK_IMAGE_LAYOUT_GENERAL,1,&r); + r = region(w*2,h*2); vkCmdCopyBufferToImage(context.command,seedLuma.handle,seed.luma,VK_IMAGE_LAYOUT_GENERAL,1,&r); + r = region(base.w,base.h); vkCmdCopyBufferToImage(context.command,seedFeature.handle,seed.feature,VK_IMAGE_LAYOUT_GENERAL,1,&r); + } + barrier(context.command,VK_PIPELINE_STAGE_TRANSFER_BIT,VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,VK_ACCESS_TRANSFER_WRITE_BIT,VK_ACCESS_SHADER_READ_BIT); + context.submit(); + std::ofstream timings, stages; + if (profile && readback) { + std::filesystem::create_directories(outputDir); stages.open(outputDir + "/stages.csv"); + stages << "frame"; for (auto name : stageNames) stages << "," << name << "_us"; stages << '\n'; + } + if (readback) { std::filesystem::create_directories(outputDir); timings.open(outputDir+"/timings.csv"); timings << "frame,gpu_us\n"; } + auto setFrame = [&](const Frame& f) { engine.setFrame({images[f.color]->view,images[f.motion]->view,images[f.depth]->view}, f.input, f.output, &f.windows); }; + // With the native network the frame is two command buffers: the front, the HIP layers on the shared buffers, + // then the back. The tokens and head are complete before each submission returns. + auto recordSplit = [&] { + engine.recordFront(context.command); + context.submit(); +#ifdef D4R_ENGINE_HIP + network->launch(); network->wait(); +#else + throw std::runtime_error("this build has no HIP support"); +#endif + context.begin(); + engine.recordBack(context.command); + }; + for (unsigned i = 0; i < frames.size(); ++i) { + setFrame(frames[i]); context.begin(); + vkCmdResetQueryPool(context.command,context.timestamps,0,queryCount); + vkCmdWriteTimestamp(context.command,VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,context.timestamps,0); + if (hip) recordSplit(); else recordFrame(0); + vkCmdWriteTimestamp(context.command,VK_PIPELINE_STAGE_ALL_COMMANDS_BIT,context.timestamps,1); + if (readback) { + auto history = engine.historyImages(); + r = region(w,h); vkCmdCopyImageToBuffer(context.command,engine.outputImage(),VK_IMAGE_LAYOUT_GENERAL,result.handle,1,&r); + vkCmdCopyImageToBuffer(context.command,history.color,VK_IMAGE_LAYOUT_GENERAL,color.handle,1,&r); + r = region(w*2,h*2); vkCmdCopyImageToBuffer(context.command,history.luma,VK_IMAGE_LAYOUT_GENERAL,luma.handle,1,&r); + r = region(base.w,base.h); vkCmdCopyImageToBuffer(context.command,history.feature,VK_IMAGE_LAYOUT_GENERAL,feature.handle,1,&r); + VkBufferCopy copy{0,0,head.size}; vkCmdCopyBuffer(context.command,engine.headBuffer(),head.handle,1,©); + barrier(context.command,VK_PIPELINE_STAGE_TRANSFER_BIT,VK_PIPELINE_STAGE_HOST_BIT,VK_ACCESS_TRANSFER_WRITE_BIT,VK_ACCESS_HOST_READ_BIT); + } + context.submit(); + double us = context.gpuMicros(); timings << i << ',' << us << '\n'; + if (profile && !hip) { + auto values = stageMicros(1); + if (readback) { stages << i; for (auto value : values) stages << ',' << value; stages << '\n'; } + if (repeats == 1) for (unsigned stage = 0; stage < values.size(); ++stage) + std::cout << " " << stageNames[stage] << ": " << std::fixed << std::setprecision(1) << values[stage] << " us\n"; + } + std::cout << "frame " << i << ": " << std::fixed << std::setprecision(1) << us << " us GPU" + << (hip ? " (Vulkan front + HIP network + Vulkan back)" : "") << '\n'; + if (readback) { + const std::string prefix = outputDir+"/frame-"+std::to_string(i); + result.save(prefix+".312.bin"); color.save(prefix+".288.bin"); luma.save(prefix+".296.bin"); feature.save(prefix+".304.bin"); head.save(prefix+".head.bin"); + } + } + std::array, d4r::KTimingStages> stageTrials; + if (repeats > 1) { + // Repeat the last input as successive frames: both history sets and all data dependencies remain active. + if (!hip) setFrame(frames.back()); + std::vector gpu, wall; + for (unsigned trial = 0; trial < 12; ++trial) { + context.begin(); vkCmdResetQueryPool(context.command,context.timestamps,0,queryCount); + vkCmdWriteTimestamp(context.command,VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,context.timestamps,0); + for (int i = 0; i < repeats; ++i) { + if (hip) { setFrame(frames.back()); recordSplit(); } + else recordFrame(i); + } + vkCmdWriteTimestamp(context.command,VK_PIPELINE_STAGE_ALL_COMMANDS_BIT,context.timestamps,1); + auto t0 = std::chrono::steady_clock::now(); context.submit(); + double cpu = std::chrono::duration(std::chrono::steady_clock::now()-t0).count()/repeats; + if (trial >= 3) { + gpu.push_back(context.gpuMicros()/repeats); wall.push_back(cpu); + if (profile && !hip) { auto values = stageMicros(repeats); + for (unsigned stage = 0; stage < values.size(); ++stage) stageTrials[stage].push_back(values[stage]); } + } + } + std::sort(gpu.begin(),gpu.end()); std::sort(wall.begin(),wall.end()); + std::cout << "full native frame: " << gpu[gpu.size()/2] << " us GPU median (min " << gpu.front() + << "); " << wall[wall.size()/2] << " us amortized submit+wait; " << repeats << " frames/submit\n"; + } + if (profile && !hip && repeats > 1) { + std::cout << "stage GPU medians (us/frame, includes preceding dependencies; timestamp instrumentation enabled):\n"; + for (unsigned stage = 0; stage < stageTrials.size(); ++stage) { + auto& values = stageTrials[stage]; std::sort(values.begin(), values.end()); + std::cout << " " << stageNames[stage] << ": " << values[values.size()/2] << '\n'; + } + } + return 0; +} +} // namespace +int main(int argc, char** argv) { + try { + if (argc != 4 && argc != 5 && argc != 6) { std::cerr << "usage: d4r-k MODEL_DIR CASE.txt OUTPUT_DIR|- [TIMING_REPEATS [hip]]\n" + "The trailing 'hip' runs the eleven layers on the native HIP kernels (compile_k.py --hip) through two\n" + "buffers shared with the engine; without it the engine's own Vulkan network is used.\n" + "Set D4R_ENGINE_PROFILE=1 for stage timings (stages.csv with readback; Vulkan network only).\n"; return 2; } + int repeats = argc >= 5 ? std::stoi(argv[4]) : 1; + if (repeats < 1 || repeats > 1000) throw std::runtime_error("timing repeats must be 1..1000"); + const bool hip = argc == 6 && std::string(argv[5]) == "hip"; + if (argc == 6 && !hip) throw std::runtime_error("the optional sixth argument must be 'hip'"); + return run(argv[1],argv[2],argv[3],repeats,hip); + } catch (const std::exception& e) { std::cerr << "d4r-k: " << e.what() << '\n'; return 1; } +} diff --git a/engine/replay_m.cpp b/engine/replay_m.cpp new file mode 100644 index 0000000..56ed381 --- /dev/null +++ b/engine/replay_m.cpp @@ -0,0 +1,202 @@ +// Preset M on the engine's runtime: runs raw frames and writes the output images. No CUDA/NGX. +// d4r-m MODEL_DIR RENDER_W RENDER_H OUT_W OUT_H OUTPUT_PREFIX [repeats [hip]] -- COLOR MOTION DEPTH JX JY ... +// "hip" runs the network on the native HIP layer kernels (hip_net.h; the model needs compile_m.py --hip) between the +// engine's Vulkan stages, through two buffers shared by file descriptor. +// COLOR: RGBA16F, MOTION: RG16F, DEPTH: R32F, all at render size, tightly packed. The first frame resets. +#include "runtime.h" +#ifdef D4R_ENGINE_HIP +#include "hip_net.h" +#endif +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#define CK(x) do { VkResult r_ = (x); if (r_ != VK_SUCCESS) throw std::runtime_error(std::string(#x) + " failed: " + std::to_string(r_)); } while (0) +static std::vector slurp(const std::string& p) { + std::ifstream f(p, std::ios::binary); + if (!f) throw std::runtime_error("cannot read " + p); + return std::vector((std::istreambuf_iterator(f)), {}); +} +struct Ctx { + VkInstance instance{}; VkPhysicalDevice physical{}; VkDevice device{}; VkQueue queue{}; VkCommandPool pool{}; VkPhysicalDeviceMemoryProperties memory{}; + uint32_t type(uint32_t bits, VkMemoryPropertyFlags flags) const { + for (uint32_t i = 0; i < memory.memoryTypeCount; ++i) if ((bits & (1u << i)) && (memory.memoryTypes[i].propertyFlags & flags) == flags) return i; + throw std::runtime_error("no memory type"); + } +}; +struct Staging { VkBuffer buffer{}; VkDeviceMemory memory{}; void* map{}; }; +static Staging staging(const Ctx& c, VkDeviceSize size) { + Staging s; VkBufferCreateInfo bi{VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO}; bi.size = size; bi.usage = VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT; + CK(vkCreateBuffer(c.device, &bi, nullptr, &s.buffer)); + VkMemoryRequirements r; vkGetBufferMemoryRequirements(c.device, s.buffer, &r); + VkMemoryAllocateInfo ai{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO}; ai.allocationSize = r.size; ai.memoryTypeIndex = c.type(r.memoryTypeBits, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT); + CK(vkAllocateMemory(c.device, &ai, nullptr, &s.memory)); CK(vkBindBufferMemory(c.device, s.buffer, s.memory, 0)); CK(vkMapMemory(c.device, s.memory, 0, VK_WHOLE_SIZE, 0, &s.map)); + return s; +} +// a storage buffer whose memory another API imports +struct Shared { VkBuffer buffer{}; VkDeviceMemory memory{}; VkDeviceSize allocation = 0; int fd = -1; }; +static Shared shared(const Ctx& c, VkDeviceSize size) { + Shared s; VkExternalMemoryBufferCreateInfo ei{VK_STRUCTURE_TYPE_EXTERNAL_MEMORY_BUFFER_CREATE_INFO}; ei.handleTypes = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT; + VkBufferCreateInfo bi{VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO, &ei}; bi.size = size; bi.usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT; + CK(vkCreateBuffer(c.device, &bi, nullptr, &s.buffer)); + VkMemoryRequirements r; vkGetBufferMemoryRequirements(c.device, s.buffer, &r); + VkExportMemoryAllocateInfo xi{VK_STRUCTURE_TYPE_EXPORT_MEMORY_ALLOCATE_INFO}; xi.handleTypes = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT; + VkMemoryAllocateInfo ai{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO, &xi}; ai.allocationSize = s.allocation = r.size; + ai.memoryTypeIndex = c.type(r.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT); + CK(vkAllocateMemory(c.device, &ai, nullptr, &s.memory)); CK(vkBindBufferMemory(c.device, s.buffer, s.memory, 0)); + VkMemoryGetFdInfoKHR gi{VK_STRUCTURE_TYPE_MEMORY_GET_FD_INFO_KHR}; gi.memory = s.memory; gi.handleType = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT; + const auto getFd = reinterpret_cast(vkGetDeviceProcAddr(c.device, "vkGetMemoryFdKHR")); + if (!getFd) throw std::runtime_error("vkGetMemoryFdKHR is unavailable"); + CK(getFd(c.device, &gi, &s.fd)); + return s; +} +struct Image { VkImage image{}; VkDeviceMemory memory{}; VkImageView view{}; }; +static Image image(const Ctx& c, uint32_t w, uint32_t h, VkFormat f) { + Image im; VkImageCreateInfo ci{VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO}; ci.imageType = VK_IMAGE_TYPE_2D; ci.format = f; ci.extent = {w, h, 1}; ci.mipLevels = ci.arrayLayers = 1; + ci.samples = VK_SAMPLE_COUNT_1_BIT; ci.usage = VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT; + CK(vkCreateImage(c.device, &ci, nullptr, &im.image)); + VkMemoryRequirements r; vkGetImageMemoryRequirements(c.device, im.image, &r); + VkMemoryAllocateInfo ai{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO}; ai.allocationSize = r.size; ai.memoryTypeIndex = c.type(r.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT); + CK(vkAllocateMemory(c.device, &ai, nullptr, &im.memory)); CK(vkBindImageMemory(c.device, im.image, im.memory, 0)); + VkImageViewCreateInfo vi{VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO}; vi.image = im.image; vi.viewType = VK_IMAGE_VIEW_TYPE_2D; vi.format = f; vi.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + CK(vkCreateImageView(c.device, &vi, nullptr, &im.view)); + return im; +} +int main(int argc, char** argv) try { + int sep = 0; + for (int i = 1; i < argc; ++i) if (!std::strcmp(argv[i], "--")) sep = i; + if (sep < 7 || (argc - sep - 1) % 5 || argc - sep - 1 == 0) { std::fprintf(stderr, "usage: d4r-m MODEL_DIR RW RH OW OH OUTPUT_PREFIX [repeats [hip]] -- COLOR MOTION DEPTH JX JY ...\n"); return 2; } + const bool hip = sep > 8 && !std::strcmp(argv[8], "hip"); +#ifndef D4R_ENGINE_HIP + if (hip) throw std::runtime_error("built without the HIP network (engine/build.sh needs ROCm)"); +#endif + const std::string model = argv[1], prefix = argv[6]; + const uint32_t rw = std::atoi(argv[2]), rh = std::atoi(argv[3]), ow = std::atoi(argv[4]), oh = std::atoi(argv[5]); + const int repeats = sep > 7 ? std::atoi(argv[7]) : 0; + const int frames = (argc - sep - 1) / 5; + Ctx c; + VkApplicationInfo app{VK_STRUCTURE_TYPE_APPLICATION_INFO}; app.apiVersion = VK_API_VERSION_1_3; + VkInstanceCreateInfo ici{VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO}; ici.pApplicationInfo = &app; + CK(vkCreateInstance(&ici, nullptr, &c.instance)); + uint32_t n = 8; VkPhysicalDevice pds[8]; vkEnumeratePhysicalDevices(c.instance, &n, pds); + if (!n) throw std::runtime_error("no Vulkan GPU"); + c.physical = pds[0]; + for (uint32_t i = 0; i < n; ++i) { VkPhysicalDeviceProperties p; vkGetPhysicalDeviceProperties(pds[i], &p); if (p.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU) c.physical = pds[i]; } + vkGetPhysicalDeviceMemoryProperties(c.physical, &c.memory); + VkPhysicalDeviceShaderFloat8FeaturesEXT f8{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_FLOAT8_FEATURES_EXT}; + VkPhysicalDeviceCooperativeMatrixFeaturesKHR cm{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR, &f8}; + VkPhysicalDeviceVulkan13Features v13{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_FEATURES, &cm}; + VkPhysicalDeviceVulkan12Features v12{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES, &v13}; + VkPhysicalDeviceVulkan11Features v11{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES, &v12}; + VkPhysicalDeviceFeatures2 f2{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2, &v11}; + vkGetPhysicalDeviceFeatures2(c.physical, &f2); + if (!cm.cooperativeMatrix || !f8.shaderFloat8 || !f8.shaderFloat8CooperativeMatrix) throw std::runtime_error("GPU lacks FP8 cooperative matrices (VK_EXT_shader_float8)"); + uint32_t nq = 8; VkQueueFamilyProperties qf[8]; vkGetPhysicalDeviceQueueFamilyProperties(c.physical, &nq, qf); + uint32_t qi = 0; for (uint32_t i = 0; i < nq; ++i) if ((qf[i].queueFlags & VK_QUEUE_COMPUTE_BIT) && qf[i].timestampValidBits) { qi = i; break; } + float prio = 1; VkDeviceQueueCreateInfo qci{VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO}; qci.queueFamilyIndex = qi; qci.queueCount = 1; qci.pQueuePriorities = &prio; + const char* exts[] = {"VK_KHR_cooperative_matrix", "VK_EXT_shader_float8", "VK_KHR_external_memory_fd"}; + VkDeviceCreateInfo dci{VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO, &f2}; dci.queueCreateInfoCount = 1; dci.pQueueCreateInfos = &qci; dci.enabledExtensionCount = hip ? 3 : 2; dci.ppEnabledExtensionNames = exts; + CK(vkCreateDevice(c.physical, &dci, nullptr, &c.device)); vkGetDeviceQueue(c.device, qi, 0, &c.queue); + VkCommandPoolCreateInfo pi{VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO}; pi.queueFamilyIndex = qi; pi.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT; + CK(vkCreateCommandPool(c.device, &pi, nullptr, &c.pool)); + VkCommandBufferAllocateInfo cai{VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO}; cai.commandPool = c.pool; cai.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; cai.commandBufferCount = 1; + VkCommandBuffer cb; CK(vkAllocateCommandBuffers(c.device, &cai, &cb)); + VkQueryPoolCreateInfo qpi{VK_STRUCTURE_TYPE_QUERY_POOL_CREATE_INFO}; qpi.queryType = VK_QUERY_TYPE_TIMESTAMP; qpi.queryCount = 7; + VkQueryPool qp; CK(vkCreateQueryPool(c.device, &qpi, nullptr, &qp)); + VkPhysicalDeviceProperties props; vkGetPhysicalDeviceProperties(c.physical, &props); + auto submit = [&] { + CK(vkEndCommandBuffer(cb)); + VkSubmitInfo si{VK_STRUCTURE_TYPE_SUBMIT_INFO}; si.commandBufferCount = 1; si.pCommandBuffers = &cb; + CK(vkQueueSubmit(c.queue, 1, &si, VK_NULL_HANDLE)); CK(vkQueueWaitIdle(c.queue)); + }; + auto begin = [&] { CK(vkResetCommandBuffer(cb, 0)); VkCommandBufferBeginInfo bi{VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO}; CK(vkBeginCommandBuffer(cb, &bi)); }; + { + d4r::MModel m(c.physical, c.device, model); + d4r::ExternalNetwork external; +#ifdef D4R_ENGINE_HIP + std::unique_ptr network; + if (hip) { + Shared tokens = shared(c, d4r::MEngine::networkBufferBytes(rw, rh)), result = shared(c, d4r::MEngine::networkBufferBytes(rw, rh)); + external = {tokens.buffer, result.buffer}; + network = std::make_unique(model, rw, rh, tokens.fd, result.fd, tokens.allocation, result.allocation); + } +#endif + d4r::MEngine engine(m, rw, rh, ow, oh, hip ? &external : nullptr); + Image color = image(c, rw, rh, VK_FORMAT_R16G16B16A16_SFLOAT), motion = image(c, rw, rh, VK_FORMAT_R16G16_SFLOAT), depth = image(c, rw, rh, VK_FORMAT_R32_SFLOAT); + const VkDeviceSize sizes[3] = {VkDeviceSize(rw) * rh * 8, VkDeviceSize(rw) * rh * 4, VkDeviceSize(rw) * rh * 4}; + Staging in[3] = {staging(c, sizes[0]), staging(c, sizes[1]), staging(c, sizes[2])}, out = staging(c, VkDeviceSize(ow) * oh * 8); + Image* images[3] = {&color, &motion, &depth}; + begin(); m.recordUpload(cb); engine.recordInitialize(cb); + for (Image* im : images) { + VkImageMemoryBarrier b{VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER}; b.newLayout = VK_IMAGE_LAYOUT_GENERAL; b.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; + b.srcQueueFamilyIndex = b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; b.image = im->image; b.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + vkCmdPipelineBarrier(cb, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, nullptr, 0, nullptr, 1, &b); + } + submit(); + std::vector> frameInputs; + frameInputs.reserve(size_t(frames) * 3); + for (int f = 0; f < frames; ++f) + for (int k = 0; k < 3; ++k) { + auto data = slurp(argv[sep + 1 + 5 * f + k]); + if (data.size() != sizes[k]) throw std::runtime_error("wrong input size"); + frameInputs.push_back(std::move(data)); + } + std::vector times, walls; std::vector> stage(6); + const auto now = [] { return std::chrono::duration(std::chrono::steady_clock::now().time_since_epoch()).count(); }; + for (int pass = 0; pass <= repeats; ++pass) + for (int f = 0; f < frames; ++f) { + char** a = argv + sep + 1 + 5 * f; + if (pass == 0 || frames > 1) + for (int k = 0; k < 3; ++k) + std::memcpy(in[k].map, frameInputs[3 * f + k].data(), sizes[k]); + begin(); + VkBufferImageCopy region{}; region.imageSubresource = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1}; region.imageExtent = {rw, rh, 1}; + if (pass == 0 || frames > 1) + for (int k = 0; k < 3; ++k) vkCmdCopyBufferToImage(cb, in[k].buffer, images[k]->image, VK_IMAGE_LAYOUT_GENERAL, 1, ®ion); + d4r::MFrameSettings s; s.jitter[0] = std::atof(a[3]); s.jitter[1] = std::atof(a[4]); s.reset = pass == 0 && f == 0; + const double t0 = now(); + engine.setFrame({color.view, motion.view, depth.view}, s); + vkCmdResetQueryPool(cb, qp, 0, 7); + double networkUs = 0; + if (hip) { +#ifdef D4R_ENGINE_HIP + engine.recordFront(cb, qp, 0); + submit(); + network->launch(); + network->wait(); + networkUs = network->lastMilliseconds() * 1000.0; + begin(); + engine.recordBack(cb, qp, 3); +#endif + } else engine.recordFrame(cb, qp, 0); + region.imageExtent = {ow, oh, 1}; + if (pass == 0) vkCmdCopyImageToBuffer(cb, engine.outputImage(), VK_IMAGE_LAYOUT_GENERAL, out.buffer, 1, ®ion); + submit(); + if (pass) walls.push_back(now() - t0); + uint64_t t[7]; CK(vkGetQueryPoolResults(c.device, qp, 0, 7, sizeof(t), t, 8, VK_QUERY_RESULT_64_BIT | VK_QUERY_RESULT_WAIT_BIT)); + const double tick = props.limits.timestampPeriod / 1000.0; + double d[6]; for (int k = 0; k < 6; ++k) d[k] = double(t[k + 1] - t[k]) * tick; + if (hip) d[2] = networkUs; // measured HIP network time; the complete span includes host handoffs + const double us = double(t[6] - t[0]) * tick; + if (pass) for (int k = 0; k < 6; ++k) stage[k].push_back(d[k]); + if (pass == 0) { + std::ofstream(prefix + "-" + std::to_string(f) + ".rgba16f", std::ios::binary).write(static_cast(out.map), std::streamsize(ow) * oh * 8); + std::printf("frame %d: %.1f us GPU\n", f, us); + } else times.push_back(us); + } + if (!times.empty()) { + std::sort(times.begin(), times.end()); std::printf("%zu timed frames: median %.1f us, min %.1f us GPU\n", times.size(), times[times.size() / 2], times[0]); + std::sort(walls.begin(), walls.end()); std::printf(" wall clock, record to completion: median %.1f us, min %.1f us\n", walls[walls.size() / 2], walls[0]); + const char* names[6] = {"exposure", "input", "network", "expansion", "reconstruction", "downsample"}; + for (int k = 0; k < 6; ++k) { std::sort(stage[k].begin(), stage[k].end()); std::printf(" %-14s median %.1f us\n", names[k], stage[k][stage[k].size() / 2]); } + } + CK(vkDeviceWaitIdle(c.device)); + } + return 0; +} catch (const std::exception& e) { std::fprintf(stderr, "d4r-m: %s\n", e.what()); return 1; } diff --git a/engine/round_half.glsl b/engine/round_half.glsl new file mode 100644 index 0000000..47ecc61 --- /dev/null +++ b/engine/round_half.glsl @@ -0,0 +1,42 @@ +// Round the components of a float32 cooperative accumulator to binary16 (round-to-nearest-even) and widen them +// back to f32. A register-only conversion, or a f32 -> f16 -> f32 component conversion, is folded away by RADV on +// gfx1201, and so is the packHalf2x16 round-trip (which is round-toward-zero anyway): the rounding only happens +// once the binary16 value lands in binary16 memory. Scalar f32 -> f16 -> f32 casts are folded the same way, +// notably for filtered recurrent features, hence the integer helper below. +// Cooperative matrices are composites of implementation-dependent components: every invocation owns and may +// subscript a fixed number of them (length(), 8 for a 16x16 subgroup-scope accumulator on a 32-lane wave). +float roundHalfValue(float v) { + uint bits = floatBitsToUint(v), magnitude = bits & 0x7fffffffu; + if (magnitude < 0x38800000u) return roundEven(v * 16777216.0) * (1.0 / 16777216.0); + if (magnitude >= 0x477ff000u) + return magnitude > 0x7f800000u ? v : uintBitsToFloat((bits & 0x80000000u) | 0x7f800000u); + return uintBitsToFloat((bits + 4095u + ((bits >> 13u) & 1u)) & 0xffffe000u); +} +vec4 roundHalfValue(vec4 v) { + return vec4(roundHalfValue(v.x), roundHalfValue(v.y), roundHalfValue(v.z), roundHalfValue(v.w)); +} + +// The cooperative-matrix store/read-back through a wave-owned row-major region: the cheapest form where the region +// already exists for the stage's own staging store (dec0 reuses EO). storage/offset/stride are in storage's +// element units, as in the engine's LD16/ST16. +#define ROUND_HALF_STORE(m, storage, offset, stride) { \ + coopMatStore(HACC(m), storage, offset, stride, RM); \ + subgroupMemoryBarrierShared(); subgroupBarrier(); \ + HACC rh_st_; coopMatLoad(rh_st_, storage, offset, stride, RM); \ + m = ACC(rh_st_); \ + subgroupMemoryBarrierShared(); subgroupBarrier(); \ +} + +// The hardware binary16 conversion through one private scratch slot per invocation, ~2x cheaper than a +// cooperative-matrix store/read-back to shared memory. storage must hold 8 binary16 values per invocation of the +// workgroup: shared float16_t RS[gl_WorkGroupSize.x * 8u]; (8 = components per invocation, 256 / subgroup size). +// The memory fences prevent compiler forwarding from removing the rounding. Slots are invocation-private, +// so no subgroup execution barrier is needed (unlike ROUND_HALF_STORE's cross-lane cooperative loads). +#define ROUND_HALF_SHARED(m, storage) { \ + for (int rh_i_ = 0, rh_n_ = (m).length(); rh_i_ < rh_n_; ++rh_i_) \ + storage[gl_LocalInvocationID.x * 8u + uint(rh_i_)] = float16_t((m)[rh_i_]); \ + subgroupMemoryBarrierShared(); \ + for (int rh_i_ = 0, rh_n_ = (m).length(); rh_i_ < rh_n_; ++rh_i_) \ + (m)[rh_i_] = float(storage[gl_LocalInvocationID.x * 8u + uint(rh_i_)]); \ + subgroupMemoryBarrierShared(); \ +} diff --git a/engine/runtime.cpp b/engine/runtime.cpp new file mode 100644 index 0000000..1377736 --- /dev/null +++ b/engine/runtime.cpp @@ -0,0 +1,782 @@ +#include "runtime.h" +#include "vk_dispatch.h" +#include +#include +#include +#include +#include +#include + +namespace d4r { +static_assert(sizeof(InputParams) == 128); +static_assert(sizeof(OutputParams) == 188); +static_assert(offsetof(OutputParams, expo) == 148); +static_assert(sizeof(LayerShape) == 16); +static_assert(sizeof(Shape) == 184); +namespace { +void check(VkResult result, const char* operation) { + if (result != VK_SUCCESS) throw std::runtime_error(std::string(operation) + ": VkResult " + std::to_string(result)); +} +#define VKCHECK(call) check((call), #call) +std::vector read(const std::string& file) { + std::ifstream stream(file, std::ios::binary | std::ios::ate); + if (!stream) throw std::runtime_error("cannot open " + file); + auto size = stream.tellg(); + if (size <= 0 || size > 256 * 1024 * 1024) throw std::runtime_error("invalid size: " + file); + std::vector data(static_cast(size)); + stream.seekg(0); + if (!stream.read(data.data(), size)) throw std::runtime_error("cannot read " + file); + return data; +} +struct Device { + VkDevice handle; + VkPhysicalDeviceMemoryProperties memory; + VkPhysicalDeviceProperties properties; + Device(VkPhysicalDevice physical, VkDevice device) : handle(device) { + vkGetPhysicalDeviceMemoryProperties(physical, &memory); + vkGetPhysicalDeviceProperties(physical, &properties); + } + VkDeviceMemory allocate(VkMemoryRequirements requirements, VkMemoryPropertyFlags flags) const { + for (uint32_t i = 0; i < memory.memoryTypeCount; ++i) { + if ((requirements.memoryTypeBits & (1u << i)) && (memory.memoryTypes[i].propertyFlags & flags) == flags) { + VkMemoryAllocateInfo info{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO}; + info.allocationSize = requirements.size; info.memoryTypeIndex = i; + VkDeviceMemory result; + VKCHECK(vkAllocateMemory(handle, &info, nullptr, &result)); + return result; + } + } + throw std::runtime_error("no compatible Vulkan memory type"); + } +}; +struct Buffer { + const Device* device = nullptr; + VkBuffer handle = VK_NULL_HANDLE; + VkDeviceMemory memory = VK_NULL_HANDLE; + VkDeviceSize size = 0; + void* mapped = nullptr; + Buffer() = default; + Buffer(const Buffer&) = delete; + Buffer& operator=(const Buffer&) = delete; + ~Buffer() { + if (!device) return; + if (mapped) vkUnmapMemory(device->handle, memory); + if (handle) vkDestroyBuffer(device->handle, handle, nullptr); + if (memory) vkFreeMemory(device->handle, memory, nullptr); + } + void create(const Device& d, VkDeviceSize bytes, VkBufferUsageFlags usage, bool host = false) { + device = &d; size = std::max(16, (bytes + 3) & ~VkDeviceSize(3)); + VkBufferCreateInfo info{VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO}; + info.size = size; info.usage = usage; + VKCHECK(vkCreateBuffer(d.handle, &info, nullptr, &handle)); + VkMemoryRequirements requirements; vkGetBufferMemoryRequirements(d.handle, handle, &requirements); + memory = d.allocate(requirements, host ? VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT + : VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT); + VKCHECK(vkBindBufferMemory(d.handle, handle, memory, 0)); + if (host) VKCHECK(vkMapMemory(d.handle, memory, 0, VK_WHOLE_SIZE, 0, &mapped)); + } +}; +struct Image { + const Device* device = nullptr; + VkImage handle = VK_NULL_HANDLE; + VkDeviceMemory memory = VK_NULL_HANDLE; + VkImageView sample = VK_NULL_HANDLE, storage = VK_NULL_HANDLE; + uint32_t w = 0, h = 0; + ~Image() { + if (!device) return; + if (storage && storage != sample) vkDestroyImageView(device->handle, storage, nullptr); + if (sample) vkDestroyImageView(device->handle, sample, nullptr); + if (handle) vkDestroyImage(device->handle, handle, nullptr); + if (memory) vkFreeMemory(device->handle, memory, nullptr); + } + void create(const Device& d, uint32_t width, uint32_t height, bool luma = false) { + device = &d; w = width; h = height; + if (!w || !h || w > d.properties.limits.maxImageDimension2D || h > d.properties.limits.maxImageDimension2D) + throw std::runtime_error("image dimensions exceed device limits"); + VkImageCreateInfo info{VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO}; + info.flags = luma ? VK_IMAGE_CREATE_MUTABLE_FORMAT_BIT : 0; + info.imageType = VK_IMAGE_TYPE_2D; info.format = luma ? VK_FORMAT_R16_SFLOAT : VK_FORMAT_R16G16B16A16_SFLOAT; + info.extent = {w, h, 1}; info.mipLevels = info.arrayLayers = 1; info.samples = VK_SAMPLE_COUNT_1_BIT; + info.usage = VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_STORAGE_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT; + VKCHECK(vkCreateImage(d.handle, &info, nullptr, &handle)); + VkMemoryRequirements requirements; vkGetImageMemoryRequirements(d.handle, handle, &requirements); + memory = d.allocate(requirements, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT); + VKCHECK(vkBindImageMemory(d.handle, handle, memory, 0)); + VkImageViewCreateInfo view{VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO}; + view.image = handle; view.viewType = VK_IMAGE_VIEW_TYPE_2D; view.format = info.format; + view.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + VKCHECK(vkCreateImageView(d.handle, &view, nullptr, &sample)); + storage = sample; + if (luma) { + view.format = VK_FORMAT_R16_UINT; + VKCHECK(vkCreateImageView(d.handle, &view, nullptr, &storage)); + } + } +}; +void imageBarrier(VkCommandBuffer cb, VkImage image, VkImageLayout before, VkImageLayout after, + VkPipelineStageFlags source, VkPipelineStageFlags destination, VkAccessFlags src, VkAccessFlags dst) { + VkImageMemoryBarrier barrier{VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER}; + barrier.image = image; barrier.oldLayout = before; barrier.newLayout = after; + barrier.srcAccessMask = src; barrier.dstAccessMask = dst; + barrier.srcQueueFamilyIndex = barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; + barrier.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + vkCmdPipelineBarrier(cb, source, destination, 0, 0, nullptr, 0, nullptr, 1, &barrier); +} +void memoryBarrier(VkCommandBuffer cb, VkPipelineStageFlags source, VkPipelineStageFlags destination, + VkAccessFlags src, VkAccessFlags dst) { + VkMemoryBarrier barrier{VK_STRUCTURE_TYPE_MEMORY_BARRIER}; + barrier.srcAccessMask = src; barrier.dstAccessMask = dst; + vkCmdPipelineBarrier(cb, source, destination, 0, 1, &barrier, 0, nullptr, 0, nullptr); +} +struct Pipeline { + VkDevice device = VK_NULL_HANDLE; + VkDescriptorSetLayout descriptors = VK_NULL_HANDLE; + VkPipelineLayout layout = VK_NULL_HANDLE; + VkPipeline handle = VK_NULL_HANDLE; + ~Pipeline() { + if (handle) vkDestroyPipeline(device, handle, nullptr); + if (layout) vkDestroyPipelineLayout(device, layout, nullptr); + if (descriptors) vkDestroyDescriptorSetLayout(device, descriptors, nullptr); + } + void create(VkDevice dev, const std::string& file, const std::vector& types, + uint32_t subgroup, VkPipelineCache cache) { + device = dev; + std::vector bindings; + for (uint32_t i = 0; i < types.size(); ++i) + bindings.push_back({i, types[i], 1, VK_SHADER_STAGE_COMPUTE_BIT, nullptr}); + VkDescriptorSetLayoutCreateInfo di{VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO}; + di.bindingCount = static_cast(bindings.size()); di.pBindings = bindings.data(); + VKCHECK(vkCreateDescriptorSetLayout(dev, &di, nullptr, &descriptors)); + VkPipelineLayoutCreateInfo li{VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO}; + li.setLayoutCount = 1; li.pSetLayouts = &descriptors; + VKCHECK(vkCreatePipelineLayout(dev, &li, nullptr, &layout)); + auto bytes = read(file); + if (bytes.size() % 4 || bytes.size() < 20) throw std::runtime_error("invalid SPIR-V: " + file); + std::vector code(bytes.size() / 4); std::memcpy(code.data(), bytes.data(), bytes.size()); + if (code[0] != 0x07230203) throw std::runtime_error("invalid SPIR-V magic: " + file); + VkShaderModuleCreateInfo si{VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO}; + si.codeSize = bytes.size(); si.pCode = code.data(); + VkShaderModule shader; + VKCHECK(vkCreateShaderModule(dev, &si, nullptr, &shader)); + VkPipelineShaderStageRequiredSubgroupSizeCreateInfo required{VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_REQUIRED_SUBGROUP_SIZE_CREATE_INFO}; + required.requiredSubgroupSize = subgroup; + VkComputePipelineCreateInfo ci{VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO}; + ci.layout = layout; + ci.stage = {VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO, &required, + VK_PIPELINE_SHADER_STAGE_CREATE_REQUIRE_FULL_SUBGROUPS_BIT, VK_SHADER_STAGE_COMPUTE_BIT, shader, "main", nullptr}; + VkResult result = vkCreateComputePipelines(dev, cache, 1, &ci, nullptr, &handle); + vkDestroyShaderModule(dev, shader, nullptr); + check(result, "vkCreateComputePipelines"); + } +}; +VkDescriptorSet descriptorSet(VkDevice device, VkDescriptorPool pool, VkDescriptorSetLayout layout) { + VkDescriptorSetAllocateInfo info{VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO}; + info.descriptorPool = pool; info.descriptorSetCount = 1; info.pSetLayouts = &layout; + VkDescriptorSet result; + VKCHECK(vkAllocateDescriptorSets(device, &info, &result)); + return result; +} +void bindBuffer(VkDevice device, VkDescriptorSet set, uint32_t binding, const Buffer& buffer, VkDeviceSize size = VK_WHOLE_SIZE) { + VkDescriptorBufferInfo info{buffer.handle, 0, size}; + VkWriteDescriptorSet write{VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET}; + write.dstSet = set; write.dstBinding = binding; write.descriptorCount = 1; + write.descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER; write.pBufferInfo = &info; + vkUpdateDescriptorSets(device, 1, &write, 0, nullptr); +} +void bindImage(VkDevice device, VkDescriptorSet set, uint32_t binding, VkImageView view, VkSampler sampler = VK_NULL_HANDLE, + VkImageLayout layout = VK_IMAGE_LAYOUT_GENERAL) { + VkDescriptorImageInfo info{sampler, view, layout}; + VkWriteDescriptorSet write{VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET}; + write.dstSet = set; write.dstBinding = binding; write.descriptorCount = 1; + write.descriptorType = sampler ? VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER : VK_DESCRIPTOR_TYPE_STORAGE_IMAGE; + write.pImageInfo = &info; + vkUpdateDescriptorSets(device, 1, &write, 0, nullptr); +} +void dispatch(VkCommandBuffer cb, const Pipeline& pipeline, VkDescriptorSet set, uint32_t x, uint32_t y) { + vkCmdBindPipeline(cb, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline.handle); + vkCmdBindDescriptorSets(cb, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline.layout, 0, 1, &set, 0, nullptr); + vkCmdDispatch(cb, x, y, 1); +} +constexpr const char* Names[KLayers] = {"enc0", "enc1", "enc2", "enc3", "enc4", "dec5", "dec4", "dec3", "dec2", "dec1", "dec0"}; +constexpr uint32_t Channels[KLayers] = {32, 64, 64, 96, 128, 160, 128, 96, 64, 64, 32}; +constexpr uint32_t Heads[KLayers] = {2, 2, 2, 4, 4, 8, 4, 4, 2, 2, 2}; +constexpr uint32_t Merged[5] = {64, 64, 96, 128, 160}; +struct LayerParams { LayerShape shape; std::array offsets; }; +struct ExposureParams { int32_t origin[2], size[2]; float invCount, scale, key; uint32_t flags, blocks, blocksX; int32_t sampleSize[2]; + float gameScale = 1, invPre = 1; }; // the last two: preset K, shader revision 2 +static_assert(sizeof(ExposureParams) == 56); +bool exists(const std::string& file) { return bool(std::ifstream(file, std::ios::binary)); } +static_assert(sizeof(LayerParams) == 448); +} // namespace + +Shape makeKShape(uint32_t width, uint32_t height) { + if (width < 256 || height < 256 || width > 8192 || height > 8192) + throw std::runtime_error("standard K output dimensions must be 256..8192"); + Shape shape{}; shape.outputWidth = width; shape.outputHeight = height; + uint32_t tw = std::max(256u, ((width + 3) / 4 + 31) / 32 * 32); + uint32_t th = std::max(256u, ((height + 3) / 4 + 31) / 32 * 32); + for (unsigned i = 0; i < KLayers; ++i) { + unsigned scale = i <= 5 ? i : 10 - i; + shape.layers[i] = {int32_t(tw >> scale), int32_t(th >> scale), 4, 4}; + } + return shape; +} +std::array makeKWindows(const Shape& shape, uint32_t frameIndex) { + auto windows = shape.layers; + // Observed over 64 consecutive NGX evaluations: an eight-frame cycle that visits every shift once per axis. + static constexpr int32_t Cycle[8][2] = {{0, 2}, {5, 6}, {4, 0}, {1, 4}, {7, 5}, {2, 1}, {3, 7}, {6, 3}}; + int32_t x = Cycle[frameIndex % 8][0], y = Cycle[frameIndex % 8][1]; + for (unsigned depth = 0; depth < 3; ++depth) { // dec5, then enc4/dec4, then enc3/dec3 + windows[5 - depth].sx = windows[5 + depth].sx = x; windows[5 - depth].sy = windows[5 + depth].sy = y; + x = (2 * x + 4) & 7; y = (2 * y + 4) & 7; + } + return windows; +} +uint32_t kReconstructionTable(uint32_t outputWidth, uint32_t outputHeight, uint32_t renderWidth, uint32_t renderHeight) { + const double ratio = std::sqrt(double(outputWidth) * outputHeight / (double(renderWidth) * renderHeight)); + const double position = (ratio - 0.5) * 28.0 / 3.0; // table index on NGX's ratio grid + const double even = 2.0 * std::floor(position * 0.5 + 0.5); // nearest even table; a tie takes the upper + return uint32_t(std::clamp(even, 0.0, 14.0)); +} +FrameParams makeKFrameParams(const Shape& shape, const FrameSettings& settings) { + if (!settings.renderWidth || !settings.renderHeight || settings.renderWidth > shape.outputWidth || settings.renderHeight > shape.outputHeight) + throw std::runtime_error("render dimensions must be positive and no larger than the output"); + FrameParams params{}; auto& in = params.input; auto& out = params.output; + auto base = shape.layers[0]; + in.tokW = out.tokW = base.w; out.tokH = base.h; + in.flags = out.flags = (settings.reset ? FrameFlagReset : 0) | (settings.depthInverted ? 0 : FrameFlagRegularDepth) | (settings.displayMotion ? FrameFlagDisplayMotion : 0); + out.flags |= kReconstructionTable(shape.outputWidth, shape.outputHeight, settings.renderWidth, settings.renderHeight) << FrameFlagTableShift; + in.expo = out.expo = settings.colorExposure; + in.expoRatio = out.expoRatio = settings.historyExposureRatio; + in.expoIn = settings.networkExposureScale; in.minW = 1e-4f; + out.sharp = settings.sharpness; out.minAlpha = .02f; out.alphaThr = .0200042724609375f; + out.lutScale = 8; out.c184 = .040008544921875f; out.c186 = 1.0419921875f; + for (unsigned axis = 0; axis < 2; ++axis) { + float render = axis ? settings.renderHeight : settings.renderWidth, output = axis ? shape.outputHeight : shape.outputWidth; + float tokens = axis ? base.h : base.w; + (axis ? out.kB : out.kA) = -16.f * output / render; + in.out2rnd[axis] = out.out2rnd[axis] = out.near2rnd[axis] = render / output; + out.sub2rnd[axis] = .5f * render / output; out.out2tok[axis] = .5f; + in.rndOffs[axis] = settings.jitter[axis]; out.rndOffs[axis] = settings.jitter[axis] - .5f; + in.mvScale[axis] = out.mvScale[axis] = settings.motionScale[0] == 0 && settings.motionScale[1] == 0 ? (settings.displayMotion ? 1.f : output / render) : settings.motionScale[axis]; + in.mvOffs[axis] = out.mvOffs[axis] = settings.motionOffset[axis]; + in.invHist[axis] = in.invValid[axis] = out.invOut[axis] = out.invOutValid[axis] = 1.f / output; + in.tok2uv[axis] = 1.f / (tokens * 4.f); + in.colSize[axis] = in.depSize[axis] = out.colMax[axis] = out.depMax[axis] = int32_t(render) - 1; + in.mvSize[axis] = out.mvMax[axis] = int32_t(settings.displayMotion ? output : render) - 1; + out.outSize[axis] = int32_t(output); + } + return params; +} + +struct Model::Impl { + Device device; + Buffer weights, lut, weightUpload, lutUpload; + std::array, KLayers> offsets{}; + std::array network; + Pipeline input, output, outputRgb10; + std::array exposure; + bool uploaded = false, hasExposure = false, hasDirectOrigins = false, hasRegularDepth = false, hasDisplayMotion = false, + hasGameExposure = false, hasRatioTables = false; + Impl(VkPhysicalDevice physical, VkDevice dev) : device(physical, dev) {} +}; +Model::Model(VkPhysicalDevice physical, VkDevice device, const std::string& directory, VkPipelineCache cache) + : impl(std::make_unique(physical, device)) { + auto& m = *impl; + if (exists(directory + "/direct_origins.bin")) { + auto marker = read(directory + "/direct_origins.bin"); + // Shader revision: 1 has the corrected motion origins; 2 also reads the frame flags (regular depth); + // 3 also reads the display-motion flag and the game's exposure texture. + if (marker.size() != 8 || std::memcmp(marker.data(), "D4RO000", 7) || marker[7] < '1' || marker[7] > '3') + throw std::runtime_error("incompatible direct-origin shader marker"); + m.hasDirectOrigins = true; + m.hasRegularDepth = marker[7] >= '2'; + m.hasDisplayMotion = m.hasGameExposure = marker[7] == '3'; + } + auto offsets = read(directory + "/offsets.bin"); + uint32_t count = 0, weightSize = 0; + if (offsets.size() != 16 + KLayers * 108 * 4 || std::memcmp(offsets.data(), "D4RK0002", 8)) + throw std::runtime_error("incompatible K model metadata"); + std::memcpy(&count, offsets.data() + 8, 4); std::memcpy(&weightSize, offsets.data() + 12, 4); + auto weights = read(directory + "/weights.bin"), lut = read(directory + "/lut.bin"); + // One reconstruction table (older packages: the ratio it was captured at serves every ratio) or all sixteen. + if (count != KLayers || weights.size() != weightSize || weights.size() % 2 || + (lut.size() != KTableBytes && lut.size() != size_t(KTableBytes) * KTableCount)) + throw std::runtime_error("invalid K model sizes"); + m.hasRatioTables = lut.size() != KTableBytes; + std::memcpy(m.offsets.data(), offsets.data() + 16, sizeof(m.offsets)); + for (unsigned i = 0; i < KLayers; ++i) { + const auto& layer = m.offsets[i]; uint32_t c = Channels[i]; + auto span = [&](unsigned slot, uint64_t elements) { + uint32_t offset = layer[slot]; + if (offset % 16 || (uint64_t(offset) + elements) * 2 > weights.size()) + throw std::runtime_error(std::string("invalid tensor span in ") + Names[i]); + }; + for (unsigned j = 0; j < 4; ++j) span(j, 16 * c); + for (unsigned h = 0; h < Heads[i]; ++h) { + for (unsigned j = 0; j < 3; ++j) span(4 + 5 * h + j, c * 32); + span(7 + 5 * h, 4096); span(8 + 5 * h, 32 * c); + } + for (unsigned ch = 0; ch < c / 8; ++ch) { + span(44 + 3 * ch, c * 32); span(45 + 3 * ch, 512); span(46 + 3 * ch, 32 * c); + } + if (i < 5) { span(104, 4 * c * Merged[i]); span(105, 16 * Merged[i]); } + if (i >= 6) { span(104, Channels[i - 1] * 4 * c); span(105, 16 * 4 * c); } + if (i == 0) { span(106, 16 * c); span(107, 16 * c); } + if (i == 10) { span(106, c * 48); span(107, 16 * 48); } + } + if (weights.size() > m.device.properties.limits.maxStorageBufferRange) + throw std::runtime_error("model weights exceed the device's storage buffer range"); + m.weights.create(m.device, weights.size(), VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT); + m.lut.create(m.device, lut.size(), VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT); + m.weightUpload.create(m.device, weights.size(), VK_BUFFER_USAGE_TRANSFER_SRC_BIT, true); + m.lutUpload.create(m.device, lut.size(), VK_BUFFER_USAGE_TRANSFER_SRC_BIT, true); + std::memcpy(m.weightUpload.mapped, weights.data(), weights.size()); std::memcpy(m.lutUpload.mapped, lut.data(), lut.size()); + const auto b = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, t = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, s = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE; + for (unsigned i = 0; i < KLayers; ++i) m.network[i].create(device, directory + "/" + Names[i] + ".spv", {b,b,b,b,b,b}, 32, cache); + m.input.create(device, directory + "/input_k.spv", {b,t,t,t,t,t,b}, 64, cache); + m.output.create(device, directory + "/output_k.spv", {b,t,t,t,t,t,s,s,s,s,b,b}, 64, cache); + if (exists(directory + "/output_k_rgb10.spv")) + m.outputRgb10.create(device, directory + "/output_k_rgb10.spv", {b,t,t,t,t,t,s,s,s,s,b,b}, 64, cache); + if (exists(directory + "/exposure_k0.spv") && exists(directory + "/exposure_k1.spv")) { + m.exposure[0].create(device, directory + "/exposure_k0.spv", {b,t,b}, 64, cache); + m.exposure[1].create(device, directory + "/exposure_k1.spv", {b,b,b,b,b,t}, 64, cache); + m.hasExposure = true; + } +} +bool Model::hasAutoExposure() const { return impl->hasExposure; } +bool Model::hasDirectOrigins() const { return impl->hasDirectOrigins; } +bool Model::hasRegularDepth() const { return impl->hasRegularDepth; } +bool Model::hasDisplayMotion() const { return impl->hasDisplayMotion; } +bool Model::hasGameExposure() const { return impl->hasGameExposure; } +bool Model::hasRatioTables() const { return impl->hasRatioTables; } +bool Model::hasRgb10Output() const { return impl->outputRgb10.handle != VK_NULL_HANDLE; } +Model::~Model() = default; +bool Model::uploadRecorded() const { return impl->uploaded; } +void Model::recordUpload(VkCommandBuffer cb) { + auto& m = *impl; + if (m.uploaded) throw std::runtime_error("model upload already recorded"); + VkBufferCopy copy{0, 0, m.weightUpload.size}; vkCmdCopyBuffer(cb, m.weightUpload.handle, m.weights.handle, 1, ©); + copy.size = m.lutUpload.size; vkCmdCopyBuffer(cb, m.lutUpload.handle, m.lut.handle, 1, ©); + memoryBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT); + m.uploaded = true; +} + +struct Engine::Impl { + Model::Impl& model; + Shape shape; + Buffer inParams, outParams, dummy; + Buffer exposureParams, exposureSums, exposureState; + std::array exposureSets{}; + ExposureParams exposure{}; + bool exposureOn = false; + std::array layerParams; + std::array main; + std::array skip; + struct History { Image color, luma, feature; }; + std::array history; + Image final; + VkDescriptorPool pool = VK_NULL_HANDLE; + std::array layerSets{}; + std::array inputSets{}, outputSets{}; + VkSampler nearest = VK_NULL_HANDLE, linear = VK_NULL_HANDLE; + unsigned previous = 0; + bool initialized = false, frameSet = false, reset = false, windowsChanged = false, rgb10Output = false; + bool external = false; // the caller owns the token and head buffers; the layers run outside + bool frontDone = false; // recordFront has run and recordBack is due + // Recorded into each frame's command buffer, so a frame still queued on the GPU keeps its own values. + InputParams input{}; OutputParams output{}; + struct Bindings { + FrameImages images{}; + FrameOutput finalOutput{}; + VkDescriptorPool pool = VK_NULL_HANDLE; // the initial binding uses Impl::pool + std::array input{}, output{}; + VkDescriptorSet exposure = VK_NULL_HANDLE; + VkDescriptorSet exposureUpdate = VK_NULL_HANDLE; // the exposure state update (pass 1), per binding + }; + std::vector bindings; + VkDeviceSize activations = 0; + Impl(Model::Impl& m, const Shape& s) : model(m), shape(s) {} + ~Impl() { + auto dev = model.device.handle; + for (auto& binding : bindings) if (binding.pool) vkDestroyDescriptorPool(dev, binding.pool, nullptr); + if (pool) vkDestroyDescriptorPool(dev, pool, nullptr); + if (nearest) vkDestroySampler(dev, nearest, nullptr); + if (linear) vkDestroySampler(dev, linear, nullptr); + } +}; +Engine::Engine(Model& model, const Shape& shape, const ExternalNetwork* external) : impl(std::make_unique(*model.impl, shape)) { + auto& e = *impl; const auto& d = e.model.device; + auto dev = d.handle; + if (!shape.outputWidth || !shape.outputHeight || shape.outputWidth > d.properties.limits.maxImageDimension2D / 2 || + shape.outputHeight > d.properties.limits.maxImageDimension2D / 2) throw std::runtime_error("unsupported output dimensions"); + auto base = shape.layers[0]; + if (base.w < 256 || base.h < 256 || base.w % 32 || base.h % 32 || uint64_t(base.w) * 4 < shape.outputWidth || uint64_t(base.h) * 4 < shape.outputHeight || + uint32_t(base.w) > shape.outputWidth || uint32_t(base.h) > shape.outputHeight) + throw std::runtime_error("K token grid must be a multiple of 32, at least 256x256, and cover the output"); + e.external = external != nullptr; + if (e.external && (!external->tokens || !external->head)) throw std::runtime_error("external network without its buffers"); + for (unsigned i = 0; i < KLayers; ++i) { + auto layer = shape.layers[i]; unsigned scale = i <= 5 ? i : 10 - i; + if (layer.w != base.w / (1 << scale) || layer.h != base.h / (1 << scale) || layer.sx < 0 || layer.sx > 7 || layer.sy < 0 || layer.sy > 7 || (i != 5 && (layer.sx % 2 || layer.sy % 2))) + throw std::runtime_error(std::string("invalid geometry for ") + Names[i]); + if (e.external) continue; // the layers run outside; their buffer geometry is not this engine's concern + uint32_t gx = (layer.w + layer.sx + 7) / 8, gy = (layer.h + layer.sy + 7) / 8; + if (gx > d.properties.limits.maxComputeWorkGroupCount[0] || gy > d.properties.limits.maxComputeWorkGroupCount[1]) + throw std::runtime_error("dispatch exceeds device limits"); + } + e.inParams.create(d, sizeof(InputParams), VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT); + e.outParams.create(d, sizeof(OutputParams), VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT); + if (!e.external) e.dummy.create(d, 16, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT); + std::array capacity{VkDeviceSize(base.w) * base.h * 16 * 2, 16}; + for (unsigned i = 0; i < KLayers && !e.external; ++i) { + const auto& ls = shape.layers[i]; + VkDeviceSize bytes = VkDeviceSize(ls.w) * ls.h * (i < 5 ? Merged[i] / 4 : i == 10 ? 40 : Channels[i]) * 2; + capacity[(i + 1) % 2] = std::max(capacity[(i + 1) % 2], bytes); + } + for (unsigned i = 0; i < 2; ++i) { + if (e.external) { // the caller's buffers: main[0] holds the tokens, main[1] the head the final store reads + e.main[i].handle = i ? external->head : external->tokens; + e.main[i].size = i ? networkHeadBytes(shape.outputWidth, shape.outputHeight) : networkTokenBytes(shape.outputWidth, shape.outputHeight); + continue; + } + if (capacity[i] > d.properties.limits.maxStorageBufferRange) throw std::runtime_error("activation buffer exceeds device range"); + e.main[i].create(d, capacity[i], VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT); + e.activations += e.main[i].size; + } + for (unsigned i = 0; i < 5 && !e.external; ++i) { + auto ls = shape.layers[i]; VkDeviceSize bytes = VkDeviceSize(ls.w) * ls.h * Channels[i] * 2; + if (bytes > d.properties.limits.maxStorageBufferRange) throw std::runtime_error("skip buffer exceeds device range"); + e.skip[i].create(d, bytes, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT); + e.activations += e.skip[i].size; + } + VkDescriptorPoolSize sizes[] = {{VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, e.external ? 16u : 96u}, {VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 24}, {VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 8}}; + VkDescriptorPoolCreateInfo pi{VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO}; + pi.maxSets = e.external ? 6u : 17u; pi.poolSizeCount = 3; pi.pPoolSizes = sizes; + VKCHECK(vkCreateDescriptorPool(dev, &pi, nullptr, &e.pool)); + for (unsigned i = 0; i < KLayers && !e.external; ++i) { + e.layerParams[i].create(d, sizeof(LayerParams), VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, true); + LayerParams params{shape.layers[i], e.model.offsets[i]}; + std::memcpy(e.layerParams[i].mapped, ¶ms, sizeof(params)); + auto set = e.layerSets[i] = descriptorSet(dev, e.pool, e.model.network[i].descriptors); + bindBuffer(dev, set, 0, e.layerParams[i]); bindBuffer(dev, set, 1, e.model.weights); + bindBuffer(dev, set, 2, e.main[i % 2]); + bindBuffer(dev, set, 3, i >= 6 ? e.skip[10 - i] : e.dummy); + bindBuffer(dev, set, 4, e.main[(i + 1) % 2]); + bindBuffer(dev, set, 5, i < 5 ? e.skip[i] : e.dummy); + } + VkSamplerCreateInfo si{VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO}; + si.addressModeU = si.addressModeV = si.addressModeW = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE; + VKCHECK(vkCreateSampler(dev, &si, nullptr, &e.nearest)); + si.magFilter = si.minFilter = VK_FILTER_LINEAR; + VKCHECK(vkCreateSampler(dev, &si, nullptr, &e.linear)); + e.final.create(d, shape.outputWidth, shape.outputHeight); + for (unsigned i = 0; i < 2; ++i) { + auto& h = e.history[i]; + h.color.create(d, shape.outputWidth, shape.outputHeight); + h.luma.create(d, shape.outputWidth * 2, shape.outputHeight * 2, true); + h.feature.create(d, base.w, base.h); + } + if (e.model.hasExposure) { + // One partial sum per 32x32 block of the native output/2 exposure sample grid. + VkDeviceSize blocks = VkDeviceSize((shape.outputWidth / 2 + 31) / 32) * ((shape.outputHeight / 2 + 31) / 32); + e.exposureParams.create(d, sizeof(ExposureParams), VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT); + e.exposureSums.create(d, blocks * 4, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT); + e.exposureState.create(d, 8, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT); + for (unsigned i = 0; i < 2; ++i) { + e.exposureSets[i] = descriptorSet(dev, e.pool, e.model.exposure[i].descriptors); + bindBuffer(dev, e.exposureSets[i], 0, e.exposureParams); + } + bindBuffer(dev, e.exposureSets[0], 2, e.exposureSums); + } + for (unsigned i = 0; i < 2; ++i) { + auto& h = e.history[i]; + e.inputSets[i] = descriptorSet(dev, e.pool, e.model.input.descriptors); + e.outputSets[i] = descriptorSet(dev, e.pool, e.model.output.descriptors); + bindBuffer(dev, e.inputSets[i], 0, e.inParams); bindBuffer(dev, e.inputSets[i], 6, e.main[0]); + // NGX creates both of these texture objects with linear filtering. + bindImage(dev, e.inputSets[i], 2, h.color.sample, e.linear); + bindImage(dev, e.inputSets[i], 5, h.feature.sample, e.linear); + auto out = e.outputSets[i]; auto& dest = e.history[1 - i]; + bindBuffer(dev, out, 0, e.outParams); bindBuffer(dev, out, 10, e.main[1]); bindBuffer(dev, out, 11, e.model.lut); + bindImage(dev, out, 2, h.color.sample, e.linear); bindImage(dev, out, 3, h.luma.sample, e.linear); + bindImage(dev, out, 6, dest.color.storage); bindImage(dev, out, 7, dest.luma.storage); + bindImage(dev, out, 8, dest.feature.storage); bindImage(dev, out, 9, e.final.storage); + } +} +Engine::~Engine() = default; +void Engine::recordInitialize(VkCommandBuffer cb) { + auto& e = *impl; + if (e.initialized) throw std::runtime_error("engine initialization already recorded"); + VkClearColorValue zero{}; + VkImageSubresourceRange range{VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + auto init = [&](Image& image) { + imageBarrier(cb, image.handle, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_GENERAL, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, + VK_PIPELINE_STAGE_TRANSFER_BIT, 0, VK_ACCESS_TRANSFER_WRITE_BIT); + vkCmdClearColorImage(cb, image.handle, VK_IMAGE_LAYOUT_GENERAL, &zero, 1, &range); + }; + for (auto& history : e.history) { init(history.color); init(history.luma); init(history.feature); } + init(e.final); + if (e.model.hasExposure) { + const float one[2] = {1, 1}; + vkCmdUpdateBuffer(cb, e.exposureState.handle, 0, sizeof(one), one); + } + memoryBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT); + e.initialized = true; +} +void Engine::setFrame(const FrameImages& images, const InputParams& input, const OutputParams& output, + const std::array* windows, const ExposureControl* exposure, + FrameOutput finalOutput) { + auto& e = *impl; + if (finalOutput.view && finalOutput.format != VK_FORMAT_R16G16B16A16_SFLOAT && + (finalOutput.format != VK_FORMAT_A2B10G10R10_UNORM_PACK32 || !e.model.outputRgb10.handle)) + throw std::runtime_error("unsupported borrowed output format or missing RGB10A2 shader"); + if (exposure) { + if (!e.model.hasExposure) throw std::runtime_error("the model has no exposure shaders (recompile it with compile_k.py)"); + if (exposure->game && (!e.model.hasGameExposure || !images.exposure)) throw std::runtime_error("game exposure needs recompiled K shaders and an exposure view"); + if (exposure->size[0] <= 0 || exposure->size[1] <= 0 || exposure->origin[0] < 0 || exposure->origin[1] < 0 || + uint32_t(exposure->size[0]) > e.shape.outputWidth || uint32_t(exposure->size[1]) > e.shape.outputHeight || + !(exposure->scale > 0) || !(exposure->key > 0) || !std::isfinite(exposure->scale) || !std::isfinite(exposure->key) || + !(exposure->exposureScale > 0) || !(exposure->preExposure > 0) || !std::isfinite(exposure->exposureScale) || + !std::isfinite(exposure->preExposure)) + throw std::runtime_error("unsupported exposure rectangle or scale"); + } + if (!images.color || !images.motion || !images.depth) throw std::runtime_error("frame has a null image view"); + if (input.tokW != e.shape.layers[0].w || output.tokW != input.tokW || output.tokH != e.shape.layers[0].h || + input.flags != (output.flags & ~FrameFlagTableMask) || (input.flags & ~(FrameFlagReset | FrameFlagRegularDepth | FrameFlagDisplayMotion)) || output.outSize[0] != int32_t(e.shape.outputWidth) || output.outSize[1] != int32_t(e.shape.outputHeight) || + (finalOutput.view ? output.outOrigin[0] < 0 || output.outOrigin[1] < 0 : output.outOrigin[0] || output.outOrigin[1]) || output.out2tok[0] != .5f || output.out2tok[1] != .5f) + throw std::runtime_error("frame parameters do not match engine dimensions"); + if ((input.flags & FrameFlagRegularDepth) && !e.model.hasRegularDepth) + throw std::runtime_error("regular depth needs a model compiled by the current compile_k.py"); + if ((input.flags & FrameFlagDisplayMotion) && !e.model.hasDisplayMotion) + throw std::runtime_error("display-resolution motion needs a model compiled by the current compile_k.py"); + if (!(output.out2rnd[0] > 0 && output.out2rnd[0] <= 1 && output.out2rnd[1] > 0 && output.out2rnd[1] <= 1) || + output.near2rnd[0] != output.out2rnd[0] || output.near2rnd[1] != output.out2rnd[1] || + !std::isfinite(output.sharp) || + !(std::isfinite(input.expo) && input.expo > 0 && std::isfinite(input.expoIn) && input.expoIn > 0 && + std::isfinite(output.expo) && output.expo > 0 && std::isfinite(input.expoRatio) && input.expoRatio > 0 && + std::isfinite(output.expoRatio) && output.expoRatio > 0)) throw std::runtime_error("unsupported scale or exposure"); + if (windows) { + for (unsigned i = 0; i < KLayers; ++i) { + auto old = e.shape.layers[i], ls = (*windows)[i]; + if (ls.w != old.w || ls.h != old.h || ls.sx < 0 || ls.sx > 7 || ls.sy < 0 || ls.sy > 7 || (i != 5 && (ls.sx % 2 || ls.sy % 2))) + throw std::runtime_error(std::string("unsupported frame window geometry: ") + Names[i]); + } + e.shape.layers = *windows; e.windowsChanged = true; + } + e.input = input; e.output = output; + // A single-table package has only table 0, and its shaders may predate the table bits. + if (!e.model.hasRatioTables) e.output.flags &= ~FrameFlagTableMask; + e.exposureOn = exposure != nullptr; + if (exposure) { + const uint32_t sw = e.shape.outputWidth / 2, sh = e.shape.outputHeight / 2; + const uint32_t bx = (sw + 31) / 32, by = (sh + 31) / 32; + e.exposure = {{exposure->origin[0], exposure->origin[1]}, {exposure->size[0], exposure->size[1]}, + 1.f / (float(sw) * float(sh)), exposure->scale / exposure->preExposure, exposure->key, + (exposure->measure && !exposure->game ? 1u : 0u) | (exposure->reset ? 2u : 0u) | (exposure->game ? 4u : 0u), bx * by, bx, + {int32_t(sw), int32_t(sh)}, exposure->exposureScale, 1.f / exposure->preExposure}; + } + auto dev = e.model.device.handle; + auto sameImages = [&](const FrameImages& b) { + return images.color == b.color && images.motion == b.motion && images.depth == b.depth && + images.colorLayout == b.colorLayout && images.motionLayout == b.motionLayout && images.depthLayout == b.depthLayout && + images.exposure == b.exposure && images.exposureLayout == b.exposureLayout; + }; + auto found = std::find_if(e.bindings.begin(), e.bindings.end(), [&](const auto& b) { + return sameImages(b.images) && finalOutput.view == b.finalOutput.view && finalOutput.format == b.finalOutput.format; + }); + if (found == e.bindings.end()) { + Impl::Bindings binding; + binding.images = images; binding.finalOutput = finalOutput; + if (e.bindings.empty()) { + binding.input = e.inputSets; binding.output = e.outputSets; binding.exposure = e.exposureSets[0]; + binding.exposureUpdate = e.exposureSets[1]; + } else { + VkDescriptorPoolSize sizes[] = {{VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 24}, + {VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 24}, {VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 8}}; + VkDescriptorPoolCreateInfo info{VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO}; + info.maxSets = 6; info.poolSizeCount = 3; info.pPoolSizes = sizes; + VKCHECK(vkCreateDescriptorPool(dev, &info, nullptr, &binding.pool)); + } + // Only new descriptor sets are written. Existing command buffers retain their original textures. + try { + if (e.model.hasExposure) { + if (binding.pool) { + binding.exposure = descriptorSet(dev, binding.pool, e.model.exposure[0].descriptors); + binding.exposureUpdate = descriptorSet(dev, binding.pool, e.model.exposure[1].descriptors); + } + bindBuffer(dev, binding.exposure, 0, e.exposureParams); + bindImage(dev, binding.exposure, 1, images.color, e.nearest, images.colorLayout); + bindBuffer(dev, binding.exposure, 2, e.exposureSums); + // The sampled exposure view is read only in game mode; the colour image stands in otherwise. + const bool game = images.exposure != VK_NULL_HANDLE; + bindBuffer(dev, binding.exposureUpdate, 0, e.exposureParams); + bindBuffer(dev, binding.exposureUpdate, 1, e.exposureSums); + bindBuffer(dev, binding.exposureUpdate, 2, e.exposureState); + bindBuffer(dev, binding.exposureUpdate, 3, e.inParams); + bindBuffer(dev, binding.exposureUpdate, 4, e.outParams); + bindImage(dev, binding.exposureUpdate, 5, game ? images.exposure : images.color, e.nearest, game ? images.exposureLayout : images.colorLayout); + } + for (unsigned i = 0; i < 2; ++i) { + if (binding.pool) { + binding.input[i] = descriptorSet(dev, binding.pool, e.model.input.descriptors); + binding.output[i] = descriptorSet(dev, binding.pool, e.model.output.descriptors); + } + auto in = binding.input[i], out = binding.output[i]; + auto& h = e.history[i]; auto& dest = e.history[1 - i]; + bindBuffer(dev, in, 0, e.inParams); bindBuffer(dev, in, 6, e.main[0]); + bindImage(dev, in, 1, images.color, e.nearest, images.colorLayout); + bindImage(dev, in, 2, h.color.sample, e.linear); + bindImage(dev, in, 3, images.motion, e.nearest, images.motionLayout); + bindImage(dev, in, 4, images.depth, e.nearest, images.depthLayout); + bindImage(dev, in, 5, h.feature.sample, e.linear); + bindBuffer(dev, out, 0, e.outParams); bindBuffer(dev, out, 10, e.main[1]); bindBuffer(dev, out, 11, e.model.lut); + bindImage(dev, out, 1, images.color, e.nearest, images.colorLayout); + bindImage(dev, out, 2, h.color.sample, e.linear); bindImage(dev, out, 3, h.luma.sample, e.linear); + bindImage(dev, out, 4, images.motion, e.nearest, images.motionLayout); + bindImage(dev, out, 5, images.depth, e.nearest, images.depthLayout); + bindImage(dev, out, 6, dest.color.storage); bindImage(dev, out, 7, dest.luma.storage); + bindImage(dev, out, 8, dest.feature.storage); + bindImage(dev, out, 9, finalOutput.view ? finalOutput.view : e.final.storage); + } + e.bindings.push_back(binding); + } catch (...) { + if (binding.pool) vkDestroyDescriptorPool(dev, binding.pool, nullptr); + throw; + } + found = e.bindings.end() - 1; + } + e.inputSets = found->input; e.outputSets = found->output; + e.exposureSets[0] = found->exposure; + e.exposureSets[1] = found->exposureUpdate; + e.rgb10Output = finalOutput.view && finalOutput.format == VK_FORMAT_A2B10G10R10_UNORM_PACK32; + e.reset = (input.flags & FrameFlagReset) != 0; e.frameSet = true; +} +void Engine::recordFrame(VkCommandBuffer cb) { + recordFrame(cb, {}); +} +void Engine::recordFrame(VkCommandBuffer cb, const FrameProfiler& profiler) { + auto& e = *impl; + if (e.external) throw std::runtime_error("this engine's network is external: record the front and the back"); + if (!e.initialized || !e.model.uploaded || !e.frameSet) throw std::runtime_error("initialize/upload/setFrame before recording inference"); + auto mark = [&](uint32_t boundary) { if (profiler.mark) profiler.mark(cb, boundary, profiler.user); }; + mark(0); + // Includes reads-to-writes when reusing scratch and histories, and copies made by the caller after the last frame. + memoryBarrier(cb, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, + VK_ACCESS_MEMORY_WRITE_BIT | VK_ACCESS_MEMORY_READ_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT | VK_ACCESS_TRANSFER_WRITE_BIT); + vkCmdUpdateBuffer(cb, e.inParams.handle, 0, sizeof(InputParams), &e.input); + vkCmdUpdateBuffer(cb, e.outParams.handle, 0, sizeof(OutputParams), &e.output); + if (e.exposureOn) vkCmdUpdateBuffer(cb, e.exposureParams.handle, 0, sizeof(ExposureParams), &e.exposure); + if (e.windowsChanged) { + for (unsigned i = 0; i < KLayers; ++i) vkCmdUpdateBuffer(cb, e.layerParams[i].handle, 0, sizeof(LayerShape), &e.shape.layers[i]); + e.windowsChanged = false; + } + if (e.reset) { + VkClearColorValue zero{}; VkImageSubresourceRange range{VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + for (auto& h : e.history) for (auto* im : {&h.color, &h.luma, &h.feature}) + vkCmdClearColorImage(cb, im->handle, VK_IMAGE_LAYOUT_GENERAL, &zero, 1, &range); + } + memoryBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT); + mark(1); + if (e.exposureOn) { + const auto compute = VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT; + if (e.exposure.flags & 1u) { + dispatch(cb, e.model.exposure[0], e.exposureSets[0], e.exposure.blocksX, e.exposure.blocks / e.exposure.blocksX); + memoryBarrier(cb, compute, compute, VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT); + } + dispatch(cb, e.model.exposure[1], e.exposureSets[1], 1, 1); + memoryBarrier(cb, compute, compute, VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT); + } + mark(2); + auto base = e.shape.layers[0]; + dispatch(cb, e.model.input, e.inputSets[e.previous], uint32_t(base.w) / 8, uint32_t(base.h) / 8); + mark(3); + for (unsigned i = 0; i < KLayers; ++i) { + memoryBarrier(cb, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, + VK_ACCESS_SHADER_WRITE_BIT | VK_ACCESS_SHADER_READ_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT); + auto ls = e.shape.layers[i]; + dispatch(cb, e.model.network[i], e.layerSets[i], (ls.w + ls.sx + 7) / 8, (ls.h + ls.sy + 7) / 8); + mark(4 + i); + } + memoryBarrier(cb, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT); + dispatch(cb, e.rgb10Output ? e.model.outputRgb10 : e.model.output, e.outputSets[e.previous], + (e.shape.outputWidth + 15) / 16, (e.shape.outputHeight + 15) / 16); + mark(4 + KLayers); + memoryBarrier(cb, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, + VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT); + mark(KTimingStages); + e.previous = 1 - e.previous; +} +// The front: setup, exposure and the input stage, which writes the tokens the external network reads. +void Engine::recordFront(VkCommandBuffer cb, VkQueryPool stamps, uint32_t first) { + auto& e = *impl; + if (!e.external) throw std::runtime_error("this engine's network is internal: record the whole frame"); + if (!e.initialized || !e.model.uploaded || !e.frameSet || e.frontDone) throw std::runtime_error("initialize/upload/setFrame before recording inference"); + const auto compute = VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT; + memoryBarrier(cb, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, compute | VK_PIPELINE_STAGE_TRANSFER_BIT, + VK_ACCESS_MEMORY_WRITE_BIT | VK_ACCESS_MEMORY_READ_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT | VK_ACCESS_TRANSFER_WRITE_BIT); + vkCmdUpdateBuffer(cb, e.inParams.handle, 0, sizeof(InputParams), &e.input); + vkCmdUpdateBuffer(cb, e.outParams.handle, 0, sizeof(OutputParams), &e.output); + if (e.exposureOn) vkCmdUpdateBuffer(cb, e.exposureParams.handle, 0, sizeof(ExposureParams), &e.exposure); + if (e.reset) { + VkClearColorValue zero{}; VkImageSubresourceRange range{VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + for (auto& h : e.history) for (auto* im : {&h.color, &h.luma, &h.feature}) + vkCmdClearColorImage(cb, im->handle, VK_IMAGE_LAYOUT_GENERAL, &zero, 1, &range); + } + memoryBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, compute, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT); + uint32_t query = first; + auto stamp = [&] { if (stamps) vkCmdWriteTimestamp(cb, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, stamps, query++); }; + stamp(); // setup + if (e.exposureOn) { + if (e.exposure.flags & 1u) { + dispatch(cb, e.model.exposure[0], e.exposureSets[0], e.exposure.blocksX, e.exposure.blocks / e.exposure.blocksX); + memoryBarrier(cb, compute, compute, VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT); + } + dispatch(cb, e.model.exposure[1], e.exposureSets[1], 1, 1); + memoryBarrier(cb, compute, compute, VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT); + } + stamp(); // exposure + auto base = e.shape.layers[0]; + dispatch(cb, e.model.input, e.inputSets[e.previous], uint32_t(base.w) / 8, uint32_t(base.h) / 8); + stamp(); // input + memoryBarrier(cb, compute, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT); + e.frontDone = true; +} +// The back: the final store, which reads the head the external network wrote. +void Engine::recordBack(VkCommandBuffer cb, VkQueryPool stamps, uint32_t first) { + auto& e = *impl; + if (!e.frontDone) throw std::runtime_error("recordFront before recordBack"); + const auto compute = VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT; + memoryBarrier(cb, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, compute, VK_ACCESS_MEMORY_WRITE_BIT | VK_ACCESS_MEMORY_READ_BIT, + VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT); + uint32_t query = first; + auto stamp = [&] { if (stamps) vkCmdWriteTimestamp(cb, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, stamps, query++); }; + stamp(); // before the final store + dispatch(cb, e.rgb10Output ? e.model.outputRgb10 : e.model.output, e.outputSets[e.previous], + (e.shape.outputWidth + 15) / 16, (e.shape.outputHeight + 15) / 16); + memoryBarrier(cb, compute, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT); + stamp(); // after the final store + e.previous = 1 - e.previous; e.frameSet = false; e.frontDone = false; +} +VkImage Engine::outputImage() const { return impl->final.handle; } +HistoryImages Engine::historyImages() const { + auto& h = impl->history[impl->previous]; return {h.color.handle, h.luma.handle, h.feature.handle}; +} +HistoryImages Engine::initialHistoryImages() const { + auto& h = impl->history[0]; return {h.color.handle, h.luma.handle, h.feature.handle}; +} +VkBuffer Engine::headBuffer() const { return impl->main[1].handle; } +VkDeviceSize Engine::headBytes() const { auto s = impl->shape.layers[0]; return VkDeviceSize(s.w) * s.h * 40 * 2; } +// The external network's buffers: the input stage's 16 f16 per token, the head's 40 f16 per token, both at the +// token grid makeKShape derives from the output. +VkDeviceSize Engine::networkTokenBytes(uint32_t outputWidth, uint32_t outputHeight) { + auto base = makeKShape(outputWidth, outputHeight).layers[0]; + return VkDeviceSize(base.w) * base.h * 16 * 2; +} +VkDeviceSize Engine::networkHeadBytes(uint32_t outputWidth, uint32_t outputHeight) { + auto base = makeKShape(outputWidth, outputHeight).layers[0]; + return VkDeviceSize(base.w) * base.h * 40 * 2; +} +VkDeviceSize Engine::activationBytes() const { return impl->activations; } +#include "runtime_m.inc" +} // namespace d4r diff --git a/engine/runtime.h b/engine/runtime.h new file mode 100644 index 0000000..9cac77f --- /dev/null +++ b/engine/runtime.h @@ -0,0 +1,237 @@ +// GPU-resident preset K inference. All Vulkan handles belong to the caller's device. +#pragma once +#include +#include +#include +#include +#include + +namespace d4r { +constexpr unsigned KLayers = 11; +// These are the shader ABI, not NGX/CUDA parameter blocks. Sizes are checked in runtime.cpp. +struct InputParams { + float out2rnd[2], rndOffs[2], mvScale[2], mvOffs[2], invHist[2], invValid[2], tok2uv[2]; + int32_t colOrigin[2], colSize[2], depOrigin[2], depSize[2], mvOrigin[2], mvSize[2]; + int32_t tokW; uint32_t flags; // FrameFlagReset, FrameFlagRegularDepth, FrameFlagDisplayMotion + float expo, expoIn, expoRatio, minW; +}; +struct OutputParams { + float out2rnd[2], rndOffs[2], near2rnd[2], sub2rnd[2], out2tok[2], mvScale[2], mvOffs[2], invOut[2], invOutValid[2]; + int32_t colOrigin[2], colMax[2], depOrigin[2], depMax[2], mvOrigin[2], mvMax[2], outOrigin[2], outSize[2]; + int32_t tokW, tokH; uint32_t flags; // as InputParams::flags (both must agree), plus the table bits + float expo, expoRatio, sharp, minAlpha, alphaThr, lutScale, kA, kB, c184, c186; +}; +constexpr uint32_t FrameFlagReset = 1, FrameFlagRegularDepth = 2, FrameFlagDisplayMotion = 4; +// OutputParams::flags only: the reconstruction table of this frame, 0..15 (kReconstructionTable). A model +// holding a single table ignores it. +constexpr uint32_t FrameFlagTableShift = 8, FrameFlagTableMask = 15u << FrameFlagTableShift; +constexpr uint32_t KTableBytes = 32768, KTableCount = 16; +struct LayerShape { int32_t w, h, sx, sy; }; +struct Shape { + uint32_t outputWidth, outputHeight; + std::array layers; +}; +struct FrameImages { + VkImageView color, motion, depth; + VkImageLayout colorLayout = VK_IMAGE_LAYOUT_GENERAL; + VkImageLayout motionLayout = VK_IMAGE_LAYOUT_GENERAL; + VkImageLayout depthLayout = VK_IMAGE_LAYOUT_GENERAL; + // The game's 1x1 exposure texture (a float view; texel 0 is read), used when ExposureControl::game is set. + VkImageView exposure = VK_NULL_HANDLE; + VkImageLayout exposureLayout = VK_IMAGE_LAYOUT_GENERAL; +}; +// A borrowed storage view in GENERAL. Its format selects the final store pipeline. +struct FrameOutput { + VkImageView view = VK_NULL_HANDLE; + VkFormat format = VK_FORMAT_R16G16B16A16_SFLOAT; +}; +struct HistoryImages { VkImage color, luma, feature; }; +// Standard linear HDR, render-resolution motion vectors. Coordinates start at (0,0). +// Motion is in render pixels by default; provide a scale to convert another convention to output pixels. +// Exposure is supplied by the application. Ratios refer to the compressed history already in this Engine. +struct FrameSettings { + uint32_t renderWidth, renderHeight; + float jitter[2] = {0, 0}; + float motionScale[2] = {0, 0}; // (0,0) chooses output/render, or 1 with displayMotion + float motionOffset[2] = {0, 0}; + float colorExposure = 1, historyExposureRatio = 1, networkExposureScale = 1; + float sharpness = 1.37f; + bool reset = false; + bool depthInverted = true; // false: regular depth (larger is farther); needs Model::hasRegularDepth + // Motion vectors at output resolution (one texel per output pixel, in output pixels); needs Model::hasDisplayMotion. + bool displayMotion = false; + bool autoExposure = false; // GameUpscaler only: measure the exposure instead of using the three fields above; exclusive with GameFrame::exposure + // GameUpscaler with autoExposure: the game's exposure scale and the factor it pre-multiplied the colour by. + // Values other than 1 need Model::hasRegularDepth (the same shader revision). + float exposureScale = 1, preExposure = 1; +}; +struct FrameParams { InputParams input; OutputParams output; }; +// The network run by something else (the native HIP layers, engine/hip_net.h and the CUDA bridge's +// d4rEngineNet*): the caller owns two storage buffers of the model's own size, shared with that network and +// written/read by the engine's stages as row-major f16. Preset M: `tokens` and `head` are each +// MEngine::networkBufferBytes; the input stage writes the tokens and the expansion reads the head. Preset K: +// `tokens` is Engine::networkTokenBytes and `head` is Engine::networkHeadBytes; the input stage writes the +// tokens and the final store reads the head. +struct ExternalNetwork { VkBuffer tokens = VK_NULL_HANDLE, head = VK_NULL_HANDLE; }; + +// Automatic K exposure: key / geometric mean of luma on an output/2 grid sampled across the colour rectangle, +// kept as a half float inside the Engine. When passed to setFrame, the frame's exposure fields (expo, expoIn, +// expoRatio) are written on the GPU and the values in InputParams/OutputParams are ignored. NGX measures on +// every second evaluation, starting with the first; a frame that does not measure keeps the last value. +struct ExposureControl { + int32_t origin[2] = {0, 0}, size[2] = {0, 0}; + float scale = 1; // applied to the luma before the logarithm, together with 1 / preExposure + float key = 0.18f / 0.82f; + // The game's. NGX scales the luma it measures by 1 / pre (cuda_luma_convert_kernel), so the kept measurement is + // pre times the unscaled one, and the frame's exposure is exposureScale * measured / pre. + float exposureScale = 1, preExposure = 1; + bool measure = true; + bool reset = false; // the history was produced with no earlier exposure: previous = current + // The game supplies the exposure in FrameImages::exposure instead of measuring it (measure is then unused). + bool game = false; +}; +Shape makeKShape(uint32_t outputWidth, uint32_t outputHeight); +// NGX keeps sixteen reconstruction tables on a uniform grid of the upscaling ratio (0.5 + index * 3/28) and +// uses the nearest even one for the geometric mean of the two axes' ratios, at most index 14 (ratio 2). +uint32_t kReconstructionTable(uint32_t outputWidth, uint32_t outputHeight, uint32_t renderWidth, uint32_t renderHeight); +FrameParams makeKFrameParams(const Shape& shape, const FrameSettings& settings); +// NGX's window schedule: the bottleneck layer's shift follows an eight-frame cycle, counted from the feature's +// first evaluation, and every shallower layer follows from the one below it as (2 * shift + 4) mod 8, which +// leaves the three outer layers on each side at (4,4). +std::array makeKWindows(const Shape& shape, uint32_t frameIndex); + +// Device must enable shaderInt16, shaderStorageImageExtendedFormats, +// storageBuffer16BitAccess, shaderFloat16, subgroupSizeControl, +// computeFullSubgroups and VK_KHR_cooperative_matrix (16x16x16 f16 inputs / f32 accumulation). +// Caller serializes recording and waits for the GPU before destruction. Frames execute in recording order. +class Model { +public: + Model(VkPhysicalDevice physical, VkDevice device, const std::string& directory, VkPipelineCache cache = VK_NULL_HANDLE); + ~Model(); + Model(const Model&) = delete; + Model& operator=(const Model&) = delete; + // Once, before any Engine using this model. The upload staging memory lives until Model destruction. + void recordUpload(VkCommandBuffer command); + bool uploadRecorded() const; + bool hasAutoExposure() const; // the model directory holds the exposure shaders + bool hasDirectOrigins() const; // packaged shaders have the corrected nonzero motion-origin handling + bool hasRegularDepth() const; // packaged shaders read the frame flags, so regular depth can be used + bool hasDisplayMotion() const; // packaged shaders read FrameFlagDisplayMotion (render-resolution motion otherwise) + bool hasGameExposure() const; // packaged shaders read FrameImages::exposure (exposure flag 4) + bool hasRatioTables() const; // lut.bin holds all sixteen reconstruction tables, not only one ratio's + bool hasRgb10Output() const; // a packaged RGB10A2 final-store variant is available +private: + struct Impl; + std::unique_ptr impl; + friend class Engine; +}; + +// Optional diagnostic markers: 0 at entry, then setup, exposure, input, eleven layers, +// reconstruction, and final synchronization (16 intervals). The callback records commands only. +constexpr unsigned KTimingStages = KLayers + 5; +struct FrameProfiler { + void (*mark)(VkCommandBuffer, uint32_t boundary, void* user) = nullptr; + void* user = nullptr; +}; + +class Engine { +public: + // With an external network the caller owns the token and head buffers (ExternalNetwork) and the eleven + // layers run outside the engine: recordFront writes the tokens, recordBack reads the head. Without one the + // engine allocates its own activation buffers and recordFrame runs the whole network internally. + Engine(Model& model, const Shape& shape, const ExternalNetwork* external = nullptr); + ~Engine(); + Engine(const Engine&) = delete; + Engine& operator=(const Engine&) = delete; + // The two buffers of an external network, in the row-major f16 layout the kernels read and write: + // 16 values per token for the input stage's tokens, 40 per token for the head the final store reads. + static VkDeviceSize networkTokenBytes(uint32_t outputWidth, uint32_t outputHeight); + static VkDeviceSize networkHeadBytes(uint32_t outputWidth, uint32_t outputHeight); + // Once, before recordFrame. Clears both sets of temporal history on the GPU. + void recordInitialize(VkCommandBuffer command); + // Parameters and window shifts are recorded into the command buffer by recordFrame, so frames may be recorded + // while earlier ones are still queued. Each image/view/layout combination gets immutable cached descriptors. + // Borrowed images and views must stay alive until Engine destruction and be on the recording queue family. + // Images must be sampleable and in the supplied descriptor layouts (GENERAL by default). + // Window shifts may change each frame; dimensions stay fixed for the lifetime of this Engine. + // Optional borrowed RGBA16F or RGB10A2 storage view in GENERAL, covering outOrigin + outSize. + // Views must outlive the Engine: immutable bindings allow rotation while earlier frames remain queued. + void setFrame(const FrameImages& images, const InputParams& input, const OutputParams& output, + const std::array* windows = nullptr, const ExposureControl* exposure = nullptr, + FrameOutput finalOutput = {}); + // Thirteen dispatches (up to two more with automatic exposure), with GPU dependencies. No submission, allocation, host wait or readback. + // A reset frame clears history first. Records are stateful: submit each recorded frame exactly once in order. + void recordFrame(VkCommandBuffer command); + void recordFrame(VkCommandBuffer command, const FrameProfiler& profiler); + // The same frame in two parts, for an external network that runs between them: the front (exposure and the + // input stage, which writes the tokens; three stamps) and the back (the final store, which reads the head; + // two stamps, the first before the store). The caller runs the network on the GPU between the two parts. + void recordFront(VkCommandBuffer command, VkQueryPool stamps = VK_NULL_HANDLE, uint32_t first = 0); + void recordBack(VkCommandBuffer command, VkQueryPool stamps = VK_NULL_HANDLE, uint32_t first = 0); + VkImage outputImage() const; // RGBA16F, GENERAL, ready for a transfer read after recordFrame + HistoryImages historyImages() const; // most recently recorded frame, GENERAL + VkBuffer headBuffer() const; + VkDeviceSize headBytes() const; + VkDeviceSize activationBytes() const; + // Seed initial history for replay/validation. Caller records copies after recordInitialize, before the first frame. + HistoryImages initialHistoryImages() const; +private: + struct Impl; + std::unique_ptr impl; +}; + +// ---------------------------------------------------------------- preset M (DLSS 4.5) +// The same conventions as the preset K classes above: the caller owns the device and the command buffer, frames +// are recorded with their parameters and execute in recording order. Reconstruction works at 1.5 x the output +// size, selecting NGX's post_3_1/post_3_2 tile class from the render ratio. Render dimensions must not exceed +// the output. Exposure is measured on the GPU every frame, as NGX does for this preset. +struct MFrameSettings { + float jitter[2] = {0, 0}; // render-pixel offsets in [-0.5, 0.5], as NGX's reconstruction tiles require + float motionScale[2] = {1, 1}; // (motion texel + offset) * scale = render pixels + float motionOffset[2] = {0, 0}; // jittered motion vectors: jitter - previous jitter, in motion texel units + bool reset = false; + bool depthInverted = true; +}; +class MModel { +public: + MModel(VkPhysicalDevice physical, VkDevice device, const std::string& directory, VkPipelineCache cache = VK_NULL_HANDLE); + ~MModel(); + MModel(const MModel&) = delete; + MModel& operator=(const MModel&) = delete; + void recordUpload(VkCommandBuffer command); + bool uploadRecorded() const; + bool hasExternalNetworkStages() const; // the model directory has the plane-layout input and expansion stages + bool hasRgb10Output() const; +private: + struct Impl; + std::unique_ptr impl; + friend class MEngine; +}; +// The network run by something else (the native HIP layers, engine/hip_net.h): the caller owns the two +// ExternalNetwork buffers, each of MEngine::networkBufferBytes shared with it. The input stage writes the +// tokens, the expansion reads the head, both in NGX's layout ([2][H][W][32] FP8, H x W = the render size +// rounded up to 32, halved). +class MEngine { +public: + MEngine(MModel& model, uint32_t renderWidth, uint32_t renderHeight, uint32_t outputWidth, uint32_t outputHeight, const ExternalNetwork* external = nullptr); + static VkDeviceSize networkBufferBytes(uint32_t renderWidth, uint32_t renderHeight); + ~MEngine(); + MEngine(const MEngine&) = delete; + MEngine& operator=(const MEngine&) = delete; + void recordInitialize(VkCommandBuffer command); + // Sampleable render-resolution views with the supplied layouts. A borrowed finalOutput writes + // RGBA16F or RGB10A2 directly in GENERAL at origin zero. Views must outlive this engine. + void setFrame(const FrameImages& images, const MFrameSettings& settings, FrameOutput finalOutput = {}); + // With a timestamp query pool, seven stamps are written from query `first`: start, after exposure, input, + // network, expansion, reconstruction and downsample. + void recordFrame(VkCommandBuffer command, VkQueryPool stamps = VK_NULL_HANDLE, uint32_t first = 0); + // The same frame in two parts, for an external network that runs between them: the front (exposure, input stage; + // three stamps) and the back (expansion, reconstruction, downsample; four stamps, the first before the expansion). + void recordFront(VkCommandBuffer command, VkQueryPool stamps = VK_NULL_HANDLE, uint32_t first = 0); + void recordBack(VkCommandBuffer command, VkQueryPool stamps = VK_NULL_HANDLE, uint32_t first = 0); + VkImage outputImage() const; // RGBA16F, GENERAL +private: + struct Impl; + std::unique_ptr impl; +}; +} // namespace d4r diff --git a/engine/runtime_m.inc b/engine/runtime_m.inc new file mode 100644 index 0000000..c93e702 --- /dev/null +++ b/engine/runtime_m.inc @@ -0,0 +1,419 @@ +// Preset M runtime; included by runtime.cpp inside namespace d4r (it shares that file's Vulkan helpers). +// GCC treats an included file as a header and warns that these types use the including file's anonymous namespace. +#pragma GCC diagnostic ignored "-Wsubobject-linkage" +namespace { +constexpr unsigned MLayers = 10; +constexpr const char* MShader[MLayers] = {"m_enc1", "m_enc2", "m_tube", "m_tube", "m_tube", "m_tube", "m_tube", "m_tube", "m_dec2", "m_dec1"}; +constexpr int32_t MShift[MLayers] = {2, 2, 2, 0, 2, 0, 2, 0, 0, 0}; +constexpr unsigned MScale[MLayers] = {1, 2, 4, 4, 4, 4, 4, 4, 2, 1}; // token grid divisor +constexpr unsigned MChannels[MLayers] = {64, 96, 128, 128, 128, 128, 128, 128, 96, 64}; +struct MEnc0Params { + int32_t size[2], shift[2], tok[2]; + float jitter[2], prevJitter[2], mvScale[2], mvOffs[2], rnd2hi[2], hi2rnd[2], hiSize[2]; + float changeScale, changeGain; + uint32_t depthFlag, reset, frame; + float expo, expoPrev; + float halfInv[2], renderInv[2]; // CUDA texture coordinates use host-rounded reciprocals +}; +struct MDec0Params { int32_t size[2], tokW; uint32_t reset; float rndOffs[2], rnd2hi[2], hiSize[2], histOffs[2], sigma, rho, renderInv[2]; }; +struct MPostParams { int32_t colOrigin[2], hiSize[2], rndSize[2]; float hi2rnd[2], rndOffs[2], invRnd[2], gain; uint32_t reset; float expo, expoPrev; }; +struct MDownParams { int32_t srcSize[2], outOrigin[2], colOrigin[2], colSize[2]; float out2rnd[2], rndOffs[2], expo; }; +struct MLayerParams { int32_t w, h, sx, sy; uint32_t o[72]; }; +static_assert(sizeof(MEnc0Params) == 124 && offsetof(MEnc0Params, expo) == 100 && offsetof(MEnc0Params, halfInv) == 108, "enc0_m.comp and exposure_k.comp (words 25, 26)"); +static_assert(sizeof(MDec0Params) == 64, "dec0_m.comp"); +static_assert(sizeof(MPostParams) == 64 && offsetof(MPostParams, expo) == 56, "post_m.comp and exposure_k.comp (words 14, 15)"); +static_assert(sizeof(MDownParams) == 52 && offsetof(MDownParams, expo) == 48, "down_m.comp and exposure_k.comp (word 12)"); +static_assert(sizeof(MLayerParams) == 304, "layer_m.comp"); +struct Tex { + const Device* device = nullptr; + VkImage handle = VK_NULL_HANDLE; + VkDeviceMemory memory = VK_NULL_HANDLE; + VkImageView view = VK_NULL_HANDLE; + ~Tex() { + if (!device) return; + if (view) vkDestroyImageView(device->handle, view, nullptr); + if (handle) vkDestroyImage(device->handle, handle, nullptr); + if (memory) vkFreeMemory(device->handle, memory, nullptr); + } + void create(const Device& d, uint32_t w, uint32_t h, VkFormat format) { + device = &d; + if (!w || !h || w > d.properties.limits.maxImageDimension2D || h > d.properties.limits.maxImageDimension2D) + throw std::runtime_error("image dimensions exceed device limits"); + VkImageCreateInfo info{VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO}; + info.imageType = VK_IMAGE_TYPE_2D; info.format = format; info.extent = {w, h, 1}; info.mipLevels = info.arrayLayers = 1; + info.samples = VK_SAMPLE_COUNT_1_BIT; + info.usage = VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_STORAGE_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT; + VKCHECK(vkCreateImage(d.handle, &info, nullptr, &handle)); + VkMemoryRequirements requirements; vkGetImageMemoryRequirements(d.handle, handle, &requirements); + memory = d.allocate(requirements, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT); + VKCHECK(vkBindImageMemory(d.handle, handle, memory, 0)); + VkImageViewCreateInfo vi{VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO}; + vi.image = handle; vi.viewType = VK_IMAGE_VIEW_TYPE_2D; vi.format = format; vi.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + VKCHECK(vkCreateImageView(d.handle, &vi, nullptr, &view)); + } +}; +void bindRange(VkDevice device, VkDescriptorSet set, uint32_t binding, const Buffer& buffer, VkDeviceSize offset) { + VkDescriptorBufferInfo info{buffer.handle, offset, VK_WHOLE_SIZE}; + VkWriteDescriptorSet write{VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET}; + write.dstSet = set; write.dstBinding = binding; write.descriptorCount = 1; + write.descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER; write.pBufferInfo = &info; + vkUpdateDescriptorSets(device, 1, &write, 0, nullptr); +} +} // namespace + +struct MModel::Impl { + Device device; + Buffer w8, w16, up8, up16; + uint32_t enc0[3]{}, dec0[2]{}; + std::array, MLayers> layers{}; + std::array network; // one per layer (five distinct shaders) + Pipeline input, expand, post, postQuality, postCompact, down, downRgb10; + Pipeline inputPlanes, expandPlanes; // token buffers in NGX's plane layout, for an external (HIP) network + bool planes = false; + std::array exposure; + bool uploaded = false, ldr = false; + Impl(VkPhysicalDevice physical, VkDevice dev) : device(physical, dev) {} +}; +MModel::MModel(VkPhysicalDevice physical, VkDevice device, const std::string& directory, VkPipelineCache cache) + : impl(std::make_unique(physical, device)) { + auto& m = *impl; + m.ldr = exists(directory + "/ldr"); // the stages were compiled for a game without the HDR flag: no exposure + auto head = read(directory + "/offsets.bin"); + const size_t expect = 16 + 5 * 4 + MLayers * 74 * 4; + if (head.size() != expect || std::memcmp(head.data(), "D4RM0003", 8)) throw std::runtime_error("incompatible M model metadata"); + uint32_t n8 = 0, n16 = 0; + std::memcpy(&n8, head.data() + 8, 4); std::memcpy(&n16, head.data() + 12, 4); + auto b8 = read(directory + "/weights8.bin"), b16 = read(directory + "/weights16.bin"); + if (b8.size() != n8 || b16.size() != n16 || n16 % 64) throw std::runtime_error("invalid M model sizes"); + std::memcpy(m.enc0, head.data() + 16, 12); std::memcpy(m.dec0, head.data() + 28, 8); + std::memcpy(m.layers.data(), head.data() + 36, sizeof(m.layers)); + auto inside = [&](uint32_t off8, uint32_t off16) { if (off8 >= n8 || off16 >= n16 || off8 % 64 || off16 % 64) throw std::runtime_error("invalid M model offsets"); }; + inside(m.enc0[0], m.enc0[1]); inside(m.dec0[0], m.dec0[1]); inside(0, m.enc0[2]); + for (auto& l : m.layers) inside(l[0], l[1]); + const VkBufferUsageFlags usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT; + m.w8.create(m.device, b8.size(), usage); m.w16.create(m.device, b16.size(), usage); + m.up8.create(m.device, b8.size(), VK_BUFFER_USAGE_TRANSFER_SRC_BIT, true); m.up16.create(m.device, b16.size(), VK_BUFFER_USAGE_TRANSFER_SRC_BIT, true); + std::memcpy(m.up8.mapped, b8.data(), b8.size()); std::memcpy(m.up16.mapped, b16.data(), b16.size()); + const auto b = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, t = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, s = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE; + for (unsigned i = 0; i < MLayers; ++i) m.network[i].create(device, directory + "/" + MShader[i] + ".spv", {b,b,b,b,b,b,b}, 32, cache); + m.input.create(device, directory + "/enc0_m.spv", {b,t,t,t,t,t,t,s,s,s,b,b,b,b}, 32, cache); + m.expand.create(device, directory + "/dec0_m.spv", {b,t,t,s,s,s,s,b,b,b}, 32, cache); + if ((m.planes = exists(directory + "/enc0_m_planes.spv") && exists(directory + "/dec0_m_planes.spv"))) { + m.inputPlanes.create(device, directory + "/enc0_m_planes.spv", {b,t,t,t,t,t,t,s,s,s,b,b,b,b}, 32, cache); + m.expandPlanes.create(device, directory + "/dec0_m_planes.spv", {b,t,t,s,s,s,s,b,b,b}, 32, cache); + } + m.post.create(device, directory + "/post_m.spv", {b,t,t,t,t,t,t,t,t,s,s,b}, 64, cache); + m.postQuality.create(device, directory + "/post_m_q.spv", {b,t,t,t,t,t,t,t,t,s,s,b}, 64, cache); + if (exists(directory + "/post_m_c.spv")) + m.postCompact.create(device, directory + "/post_m_c.spv", {b,t,t,t,t,t,t,t,t,s,s,b}, 64, cache); + m.down.create(device, directory + "/down_m.spv", {b,t,s,b}, 64, cache); + if (exists(directory + "/down_m_rgb10.spv")) + m.downRgb10.create(device, directory + "/down_m_rgb10.spv", {b,t,s,b}, 64, cache); + m.exposure[0].create(device, directory + "/exposure_m0.spv", {b,t,b}, 64, cache); + m.exposure[1].create(device, directory + "/exposure_m1.spv", {b,b,b,b,b,b}, 64, cache); +} +MModel::~MModel() = default; +bool MModel::uploadRecorded() const { return impl->uploaded; } +bool MModel::hasExternalNetworkStages() const { return impl->planes; } +bool MModel::hasRgb10Output() const { return impl->downRgb10.handle != VK_NULL_HANDLE; } +void MModel::recordUpload(VkCommandBuffer cb) { + auto& m = *impl; + if (m.uploaded) throw std::runtime_error("model upload already recorded"); + VkBufferCopy copy{0, 0, m.up8.size}; vkCmdCopyBuffer(cb, m.up8.handle, m.w8.handle, 1, ©); + copy.size = m.up16.size; vkCmdCopyBuffer(cb, m.up16.handle, m.w16.handle, 1, ©); + memoryBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT); + m.uploaded = true; +} + +struct MEngine::Impl { + MModel::Impl& model; + uint32_t rw, rh, ow, oh, hw, hh, gw, gh; // render, output, reconstruction, token pixel grid + Buffer enc0P, dec0P, postP, downP, expP, expSums, expState, dummy, packed; + std::array layerP; + Buffer tokens, skip1, merged1, skip2, merged2, tubeA, tubeB, dec2Out, dec1Out; + std::array mvA, feat, hist, half; + Tex mvB, mvC, gates, mix, cov, output; + VkDescriptorPool pool = VK_NULL_HANDLE; + std::array layerSets{}; + std::array enc0Sets{}, dec0Sets{}, postSets{}, expSets{}; + VkDescriptorSet downSet = VK_NULL_HANDLE; + VkSampler nearest = VK_NULL_HANDLE, linear = VK_NULL_HANDLE; + MEnc0Params enc0{}; MDec0Params dec0{}; MPostParams post{}; MDownParams down{}; ExposureParams exposure{}; + struct Bindings { + FrameImages images{}; + FrameOutput finalOutput{}; + VkDescriptorPool pool = VK_NULL_HANDLE; + std::array input{}, post{}; + VkDescriptorSet exposure = VK_NULL_HANDLE, down = VK_NULL_HANDLE; + }; + std::vector bindings; + float prevJitter[2] = {0, 0}; + unsigned current = 0, frame = 0; + bool initialized = false, frameSet = false, external = false, frontDone = false, quality = false, compact = false, rgb10Output = false; + Impl(MModel::Impl& m) : model(m) {} + ~Impl() { + auto dev = model.device.handle; + for (auto& binding : bindings) if (binding.pool) vkDestroyDescriptorPool(dev, binding.pool, nullptr); + if (pool) vkDestroyDescriptorPool(dev, pool, nullptr); + if (nearest) vkDestroySampler(dev, nearest, nullptr); + if (linear) vkDestroySampler(dev, linear, nullptr); + } +}; +VkDeviceSize MEngine::networkBufferBytes(uint32_t renderWidth, uint32_t renderHeight) { + return VkDeviceSize((renderWidth + 31) / 32 * 16) * ((renderHeight + 31) / 32 * 16) * 64 + 64; +} +MEngine::MEngine(MModel& model, uint32_t renderWidth, uint32_t renderHeight, uint32_t outputWidth, uint32_t outputHeight, const ExternalNetwork* external) + : impl(std::make_unique(*model.impl)) { + auto& e = *impl; const auto& d = e.model.device; auto dev = d.handle; + if (external && (!external->tokens || !external->head)) throw std::runtime_error("external network without its buffers"); + if (external && !e.model.planes) throw std::runtime_error("the M model has no stages for an external network"); + e.external = external != nullptr; + if (renderWidth < 64 || renderHeight < 64 || outputWidth % 2 || outputHeight % 2 || outputWidth > 16384 || outputHeight > 16384) + throw std::runtime_error("unsupported preset M dimensions"); + e.rw = renderWidth; e.rh = renderHeight; e.ow = outputWidth; e.oh = outputHeight; + e.hw = outputWidth / 2 * 3; e.hh = outputHeight / 2 * 3; + if (renderWidth > outputWidth || renderHeight > outputHeight) + throw std::runtime_error("preset M render dimensions must not exceed the output"); + e.quality = uint64_t(e.hw) < 3ull * renderWidth || uint64_t(e.hh) < 3ull * renderHeight; + // An 11x11 render tile covers a 16x16 reconstruction block when h <= 0.45; keep margin below 7/15.5. + e.compact = e.quality && e.model.postCompact.handle && + 20ull * renderWidth <= 9ull * e.hw && 20ull * renderHeight <= 9ull * e.hh; + e.gw = (renderWidth + 31) / 32 * 32; e.gh = (renderHeight + 31) / 32 * 32; + const uint32_t tw = e.gw / 2, th = e.gh / 2; + const VkBufferUsageFlags sb = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, sbt = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT; + e.enc0P.create(d, sizeof(MEnc0Params), sbt); e.dec0P.create(d, sizeof(MDec0Params), sbt); e.postP.create(d, sizeof(MPostParams), sbt); + e.downP.create(d, sizeof(MDownParams), sbt); e.expP.create(d, sizeof(ExposureParams), sbt); + e.expSums.create(d, VkDeviceSize((outputWidth / 2 + 31) / 32) * ((outputHeight / 2 + 31) / 32) * 4, sb); + e.expState.create(d, 8, sbt); e.dummy.create(d, 256, sb); + e.packed.create(d, VkDeviceSize(e.hw) * e.hh * 4, sb); + auto tok = [&](Buffer& b, unsigned scale, unsigned channels) { b.create(d, VkDeviceSize(tw / scale) * (th / scale) * channels + 64, sb); }; + if (external) { + // the caller's buffers (shared with its network): not owned, the layers' buffers are not needed + e.tokens.handle = external->tokens; e.dec1Out.handle = external->head; e.tokens.size = e.dec1Out.size = networkBufferBytes(renderWidth, renderHeight); + } else { + tok(e.tokens, 1, 64); tok(e.skip1, 1, 64); tok(e.merged1, 2, 96); tok(e.skip2, 2, 96); tok(e.merged2, 4, 128); + tok(e.tubeA, 4, 128); tok(e.tubeB, 4, 128); tok(e.dec2Out, 2, 96); tok(e.dec1Out, 1, 64); + } + for (auto& t : e.mvA) t.create(d, renderWidth, renderHeight, VK_FORMAT_R16G16_SFLOAT); + e.mvB.create(d, renderWidth, renderHeight, VK_FORMAT_R16G16_SFLOAT); e.mvC.create(d, renderWidth, renderHeight, VK_FORMAT_R16G16_SFLOAT); + for (auto& t : e.feat) t.create(d, renderWidth, renderHeight, VK_FORMAT_R8G8B8A8_SNORM); + e.gates.create(d, renderWidth, renderHeight, VK_FORMAT_R8G8B8A8_UNORM); e.mix.create(d, renderWidth, renderHeight, VK_FORMAT_R8G8B8A8_UNORM); + e.cov.create(d, renderWidth, renderHeight, VK_FORMAT_R16G16B16A16_SFLOAT); + for (auto& t : e.hist) t.create(d, e.hw, e.hh, VK_FORMAT_R16G16B16A16_SFLOAT); + // The native formatted RGB10 history has RGBA32F backing. UNORM filtering differs by up to three ulps. + for (auto& t : e.half) t.create(d, e.hw / 2, e.hh / 2, VK_FORMAT_R32G32B32A32_SFLOAT); + e.output.create(d, outputWidth, outputHeight, VK_FORMAT_R16G16B16A16_SFLOAT); + VkDescriptorPoolSize sizes[] = {{VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 128}, {VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 48}, {VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 32}}; + VkDescriptorPoolCreateInfo pi{VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO}; + pi.maxSets = 24; pi.poolSizeCount = 3; pi.pPoolSizes = sizes; + VKCHECK(vkCreateDescriptorPool(dev, &pi, nullptr, &e.pool)); + VkSamplerCreateInfo si{VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO}; + si.addressModeU = si.addressModeV = si.addressModeW = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE; + VKCHECK(vkCreateSampler(dev, &si, nullptr, &e.nearest)); + si.magFilter = si.minFilter = VK_FILTER_LINEAR; + VKCHECK(vkCreateSampler(dev, &si, nullptr, &e.linear)); + auto& m = e.model; + // the network: tokens -> enc1 -> enc2 -> six tube blocks (ping-pong) -> dec2 -> dec1 + const Buffer* in[MLayers] = {&e.tokens, &e.merged1, &e.merged2, &e.tubeA, &e.tubeB, &e.tubeA, &e.tubeB, &e.tubeA, &e.tubeB, &e.dec2Out}; + const Buffer* skip[MLayers] = {&e.dummy, &e.dummy, &e.dummy, &e.dummy, &e.dummy, &e.dummy, &e.dummy, &e.dummy, &e.skip2, &e.skip1}; + const Buffer* out[MLayers] = {&e.skip1, &e.skip2, &e.tubeA, &e.tubeB, &e.tubeA, &e.tubeB, &e.tubeA, &e.tubeB, &e.dec2Out, &e.dec1Out}; + const Buffer* merged[MLayers] = {&e.merged1, &e.merged2, &e.dummy, &e.dummy, &e.dummy, &e.dummy, &e.dummy, &e.dummy, &e.dummy, &e.dummy}; + for (unsigned i = 0; i < MLayers && !external; ++i) { + e.layerP[i].create(d, sizeof(MLayerParams), sbt, true); + MLayerParams lp{int32_t(tw / MScale[i]), int32_t(th / MScale[i]), MShift[i], MShift[i], {}}; + std::memcpy(lp.o, m.layers[i].data() + 2, sizeof(lp.o)); + std::memcpy(e.layerP[i].mapped, &lp, sizeof(lp)); + auto set = e.layerSets[i] = descriptorSet(dev, e.pool, m.network[i].descriptors); + bindBuffer(dev, set, 0, e.layerP[i]); bindRange(dev, set, 1, m.w8, m.layers[i][0]); bindRange(dev, set, 2, m.w16, m.layers[i][1]); + bindBuffer(dev, set, 3, *in[i]); bindBuffer(dev, set, 4, *skip[i]); bindBuffer(dev, set, 5, *out[i]); bindBuffer(dev, set, 6, *merged[i]); + } + for (unsigned i = 0; i < 2; ++i) { // i: the index of this frame's history outputs + const unsigned prev = 1 - i; + auto s0 = e.enc0Sets[i] = descriptorSet(dev, e.pool, (external ? m.inputPlanes : m.input).descriptors); + bindBuffer(dev, s0, 0, e.enc0P); + bindImage(dev, s0, 4, e.mvA[prev].view, e.linear); bindImage(dev, s0, 5, e.half[prev].view, e.linear); bindImage(dev, s0, 6, e.feat[prev].view, e.linear); + bindImage(dev, s0, 7, e.mvA[i].view); bindImage(dev, s0, 8, e.mvB.view); bindImage(dev, s0, 9, e.mvC.view); + bindRange(dev, s0, 10, m.w8, m.enc0[0]); bindRange(dev, s0, 11, m.w16, m.enc0[1]); bindRange(dev, s0, 12, m.w16, m.enc0[2]); bindBuffer(dev, s0, 13, e.tokens); + auto s1 = e.dec0Sets[i] = descriptorSet(dev, e.pool, (external ? m.expandPlanes : m.expand).descriptors); + bindBuffer(dev, s1, 0, e.dec0P); bindImage(dev, s1, 1, e.mvA[i].view, e.nearest); bindImage(dev, s1, 2, e.feat[prev].view, e.linear); + bindImage(dev, s1, 3, e.feat[i].view); bindImage(dev, s1, 4, e.gates.view); bindImage(dev, s1, 5, e.mix.view); bindImage(dev, s1, 6, e.cov.view); + bindBuffer(dev, s1, 7, e.dec1Out); bindRange(dev, s1, 8, m.w8, m.dec0[0]); bindRange(dev, s1, 9, m.w16, m.dec0[1]); + auto s2 = e.postSets[i] = descriptorSet(dev, e.pool, m.post.descriptors); + bindBuffer(dev, s2, 0, e.postP); + bindImage(dev, s2, 2, e.mix.view, e.nearest); bindImage(dev, s2, 3, e.gates.view, e.nearest); bindImage(dev, s2, 4, e.cov.view, e.nearest); + bindImage(dev, s2, 5, e.hist[prev].view, e.linear); bindImage(dev, s2, 6, e.mvA[i].view, e.nearest); bindImage(dev, s2, 7, e.mvB.view, e.nearest); + bindImage(dev, s2, 8, e.mvC.view, e.linear); bindImage(dev, s2, 9, e.hist[i].view); bindImage(dev, s2, 10, e.half[i].view); bindBuffer(dev, s2, 11, e.packed); + e.expSets[i] = descriptorSet(dev, e.pool, m.exposure[i].descriptors); + bindBuffer(dev, e.expSets[i], 0, e.expP); + } + bindBuffer(dev, e.expSets[0], 2, e.expSums); + bindBuffer(dev, e.expSets[1], 1, e.expSums); bindBuffer(dev, e.expSets[1], 2, e.expState); + bindBuffer(dev, e.expSets[1], 3, e.enc0P); bindBuffer(dev, e.expSets[1], 4, e.postP); bindBuffer(dev, e.expSets[1], 5, e.downP); + e.downSet = descriptorSet(dev, e.pool, m.down.descriptors); + bindBuffer(dev, e.downSet, 0, e.downP); bindImage(dev, e.downSet, 2, e.output.view); bindBuffer(dev, e.downSet, 3, e.packed); +} +MEngine::~MEngine() = default; +void MEngine::recordInitialize(VkCommandBuffer cb) { + auto& e = *impl; + if (e.initialized) throw std::runtime_error("engine initialization already recorded"); + VkClearColorValue zero{}; VkImageSubresourceRange range{VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + for (Tex* t : {&e.mvA[0], &e.mvA[1], &e.mvB, &e.mvC, &e.feat[0], &e.feat[1], &e.gates, &e.mix, &e.cov, &e.hist[0], &e.hist[1], &e.half[0], &e.half[1], &e.output}) { + imageBarrier(cb, t->handle, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_GENERAL, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, 0, VK_ACCESS_TRANSFER_WRITE_BIT); + vkCmdClearColorImage(cb, t->handle, VK_IMAGE_LAYOUT_GENERAL, &zero, 1, &range); + } + const float none[2] = {0, 0}; + vkCmdUpdateBuffer(cb, e.expState.handle, 0, sizeof(none), none); + memoryBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT); + e.initialized = true; +} +void MEngine::setFrame(const FrameImages& images, const MFrameSettings& s, FrameOutput finalOutput) { + auto& e = *impl; + if (!images.color || !images.motion || !images.depth) throw std::runtime_error("frame has a null image view"); + if (finalOutput.view && finalOutput.format != VK_FORMAT_R16G16B16A16_SFLOAT && + (finalOutput.format != VK_FORMAT_A2B10G10R10_UNORM_PACK32 || !e.model.downRgb10.handle)) + throw std::runtime_error("unsupported borrowed output format or missing RGB10A2 shader"); + const bool reset = s.reset || e.frame == 0; + const float fx = float(e.hw) / float(e.rw), fy = float(e.hh) / float(e.rh); + e.enc0 = {{int32_t(e.rw), int32_t(e.rh)}, {2, 2}, {int32_t(e.gw / 2), int32_t(e.gh / 2)}, {s.jitter[0], s.jitter[1]}, {e.prevJitter[0], e.prevJitter[1]}, + {s.motionScale[0] * fx, s.motionScale[1] * fy}, {s.motionOffset[0], s.motionOffset[1]}, {fx, fy}, {1.f / fx, 1.f / fy}, {float(e.hw), float(e.hh)}, 9.045065879821777f, 0.22356846928596497f, + s.depthInverted ? 1u : 0u, reset ? 1u : 0u, e.frame & 7u, 1.f, 0.f, {1.f / float(e.hw), 1.f / float(e.hh)}, {1.f / float(e.rw), 1.f / float(e.rh)}}; + e.dec0 = {{int32_t(e.rw), int32_t(e.rh)}, int32_t(e.gw / 2), reset ? 1u : 0u, {s.jitter[0], s.jitter[1]}, {fx, fy}, {float(e.hw), float(e.hh)}, + {e.prevJitter[0], e.prevJitter[1]}, 8.f, 0.07863261550664902f, {1.f / float(e.rw), 1.f / float(e.rh)}}; + e.post = {{0, 0}, {int32_t(e.hw), int32_t(e.hh)}, {int32_t(e.rw), int32_t(e.rh)}, {1.f / fx, 1.f / fy}, {s.jitter[0], s.jitter[1]}, + {1.f / float(e.rw), 1.f / float(e.rh)}, 2.7530112f, reset ? 1u : 0u, 1.f, 0.f}; + e.down = {{int32_t(e.hw), int32_t(e.hh)}, {0, 0}, {0, 0}, {int32_t(e.rw), int32_t(e.rh)}, {float(e.rw) / float(e.ow), float(e.rh) / float(e.oh)}, {s.jitter[0], s.jitter[1]}, 1.f}; + const uint32_t sw = e.ow / 2, sh = e.oh / 2; + const uint32_t bx = (sw + 31) / 32, by = (sh + 31) / 32; + e.exposure = {{0, 0}, {int32_t(e.rw), int32_t(e.rh)}, 1.f / (float(sw) * float(sh)), 1.f, 0.23547087609767914f, 1u | (reset ? 2u : 0u), bx * by, bx, {int32_t(sw), int32_t(sh)}}; + e.prevJitter[0] = s.jitter[0]; e.prevJitter[1] = s.jitter[1]; + auto dev = e.model.device.handle; + auto& m = e.model; + auto found = std::find_if(e.bindings.begin(), e.bindings.end(), [&](const auto& b) { + return b.images.color == images.color && b.images.motion == images.motion && b.images.depth == images.depth && + b.images.colorLayout == images.colorLayout && b.images.motionLayout == images.motionLayout && + b.images.depthLayout == images.depthLayout && b.finalOutput.view == finalOutput.view && + b.finalOutput.format == finalOutput.format; + }); + if (found == e.bindings.end()) { + Impl::Bindings b; + b.images = images; b.finalOutput = finalOutput; + if (e.bindings.empty()) { + b.input = e.enc0Sets; b.post = e.postSets; b.exposure = e.expSets[0]; b.down = e.downSet; + } else { + VkDescriptorPoolSize sizes[] = {{VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 32}, + {VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 32}, {VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 16}}; + VkDescriptorPoolCreateInfo info{VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO}; + info.maxSets = 6; info.poolSizeCount = 3; info.pPoolSizes = sizes; + VKCHECK(vkCreateDescriptorPool(dev, &info, nullptr, &b.pool)); + } + try { + if (b.pool) { + b.exposure = descriptorSet(dev, b.pool, m.exposure[0].descriptors); + b.down = descriptorSet(dev, b.pool, m.down.descriptors); + } + bindBuffer(dev, b.exposure, 0, e.expP); bindImage(dev, b.exposure, 1, images.color, e.nearest, images.colorLayout); + bindBuffer(dev, b.exposure, 2, e.expSums); + bindBuffer(dev, b.down, 0, e.downP); bindImage(dev, b.down, 1, images.color, e.nearest, images.colorLayout); + bindImage(dev, b.down, 2, finalOutput.view ? finalOutput.view : e.output.view); bindBuffer(dev, b.down, 3, e.packed); + for (unsigned i = 0; i < 2; ++i) { + if (b.pool) { + b.input[i] = descriptorSet(dev, b.pool, (e.external ? m.inputPlanes : m.input).descriptors); + b.post[i] = descriptorSet(dev, b.pool, m.post.descriptors); + } + auto in = b.input[i], post = b.post[i]; + bindBuffer(dev, in, 0, e.enc0P); + bindImage(dev, in, 1, images.color, e.nearest, images.colorLayout); + bindImage(dev, in, 2, images.motion, e.nearest, images.motionLayout); + bindImage(dev, in, 3, images.depth, e.nearest, images.depthLayout); + bindImage(dev, in, 4, e.mvA[1 - i].view, e.linear); bindImage(dev, in, 5, e.half[1 - i].view, e.linear); + bindImage(dev, in, 6, e.feat[1 - i].view, e.linear); + bindImage(dev, in, 7, e.mvA[i].view); bindImage(dev, in, 8, e.mvB.view); bindImage(dev, in, 9, e.mvC.view); + bindRange(dev, in, 10, m.w8, m.enc0[0]); bindRange(dev, in, 11, m.w16, m.enc0[1]); + bindRange(dev, in, 12, m.w16, m.enc0[2]); bindBuffer(dev, in, 13, e.tokens); + bindBuffer(dev, post, 0, e.postP); bindImage(dev, post, 1, images.color, e.nearest, images.colorLayout); + bindImage(dev, post, 2, e.mix.view, e.nearest); bindImage(dev, post, 3, e.gates.view, e.nearest); + bindImage(dev, post, 4, e.cov.view, e.nearest); bindImage(dev, post, 5, e.hist[1 - i].view, e.linear); + bindImage(dev, post, 6, e.mvA[i].view, e.nearest); bindImage(dev, post, 7, e.mvB.view, e.nearest); + bindImage(dev, post, 8, e.mvC.view, e.linear); bindImage(dev, post, 9, e.hist[i].view); + bindImage(dev, post, 10, e.half[i].view); bindBuffer(dev, post, 11, e.packed); + } + e.bindings.push_back(b); + } catch (...) { + if (b.pool) vkDestroyDescriptorPool(dev, b.pool, nullptr); + throw; + } + found = e.bindings.end() - 1; + } + e.enc0Sets = found->input; e.postSets = found->post; + e.expSets[0] = found->exposure; e.downSet = found->down; + e.rgb10Output = finalOutput.view && finalOutput.format == VK_FORMAT_A2B10G10R10_UNORM_PACK32; + e.frameSet = true; +} +namespace { +const auto compute = VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT; +void nextStage(VkCommandBuffer cb) { memoryBarrier(cb, compute, compute, VK_ACCESS_SHADER_WRITE_BIT | VK_ACCESS_SHADER_READ_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT); } +void openStages(VkCommandBuffer cb) { + memoryBarrier(cb, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, compute | VK_PIPELINE_STAGE_TRANSFER_BIT, + VK_ACCESS_MEMORY_WRITE_BIT | VK_ACCESS_MEMORY_READ_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT | VK_ACCESS_TRANSFER_WRITE_BIT); +} +void closeStages(VkCommandBuffer cb) { memoryBarrier(cb, compute, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT); } +} +// exposure and the input stage; writes three stamps +void MEngine::recordFront(VkCommandBuffer cb, VkQueryPool stamps, uint32_t first) { + auto& e = *impl; auto& m = e.model; + if (!e.initialized || !m.uploaded || !e.frameSet || e.frontDone) throw std::runtime_error("initialize/upload/setFrame before recording inference"); + openStages(cb); + vkCmdUpdateBuffer(cb, e.enc0P.handle, 0, sizeof(e.enc0), &e.enc0); vkCmdUpdateBuffer(cb, e.dec0P.handle, 0, sizeof(e.dec0), &e.dec0); + vkCmdUpdateBuffer(cb, e.postP.handle, 0, sizeof(e.post), &e.post); vkCmdUpdateBuffer(cb, e.downP.handle, 0, sizeof(e.down), &e.down); + vkCmdUpdateBuffer(cb, e.expP.handle, 0, sizeof(e.exposure), &e.exposure); + if (!e.external) for (unsigned l = 0; l < MLayers; ++l) { + const int32_t shifts[2] = {MShift[l] ^ int32_t((e.frame & 1u) * 2u), MShift[l]}; + vkCmdUpdateBuffer(cb, e.layerP[l].handle, offsetof(MLayerParams, sx), sizeof(shifts), shifts); + } + memoryBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, compute, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT); + uint32_t query = first; + auto stamp = [&] { if (stamps) vkCmdWriteTimestamp(cb, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, stamps, query++); }; + stamp(); + if (!m.ldr) { + dispatch(cb, m.exposure[0], e.expSets[0], e.exposure.blocksX, e.exposure.blocks / e.exposure.blocksX); nextStage(cb); + dispatch(cb, m.exposure[1], e.expSets[1], 1, 1); nextStage(cb); + } + stamp(); + dispatch(cb, e.external ? m.inputPlanes : m.input, e.enc0Sets[e.current], (e.gw + 2 + 7) / 8, (e.gh + 2 + 7) / 8); + stamp(); + if (e.external) closeStages(cb); else nextStage(cb); + e.frontDone = true; +} +// expansion, reconstruction and downsample; writes four stamps (the first before the expansion) +void MEngine::recordBack(VkCommandBuffer cb, VkQueryPool stamps, uint32_t first) { + auto& e = *impl; auto& m = e.model; + if (!e.frontDone) throw std::runtime_error("recordFront before recordBack"); + if (e.external) openStages(cb); + const unsigned i = e.current; + uint32_t query = first; + auto stamp = [&] { if (stamps) vkCmdWriteTimestamp(cb, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, stamps, query++); }; + stamp(); + dispatch(cb, e.external ? m.expandPlanes : m.expand, e.dec0Sets[i], e.gw / 8, e.gh / 8); nextStage(cb); + stamp(); + const bool compact = e.compact && e.post.rndOffs[0] >= -0.5f && e.post.rndOffs[0] <= 0.5f && + e.post.rndOffs[1] >= -0.5f && e.post.rndOffs[1] <= 0.5f; + dispatch(cb, compact ? m.postCompact : (e.quality ? m.postQuality : m.post), + e.postSets[i], (e.hw + 15) / 16, (e.hh + 15) / 16); nextStage(cb); + stamp(); + // D4RM0003's downsample workgroup covers 8x64 output pixels. + dispatch(cb, e.rgb10Output ? m.downRgb10 : m.down, e.downSet, (e.ow + 7) / 8, (e.oh + 63) / 64); + stamp(); + closeStages(cb); + e.current = 1 - e.current; ++e.frame; e.frameSet = false; e.frontDone = false; +} +void MEngine::recordFrame(VkCommandBuffer cb, VkQueryPool stamps, uint32_t first) { + auto& e = *impl; auto& m = e.model; + if (e.external) throw std::runtime_error("this engine's network is external: record the front and the back"); + recordFront(cb, stamps, first); + for (unsigned l = 0; l < MLayers; ++l) { + const uint32_t w = e.gw / 2 / MScale[l], h = e.gh / 2 / MScale[l]; + dispatch(cb, m.network[l], e.layerSets[l], (w + (MShift[l] ^ ((e.frame & 1u) * 2u)) + 7) / 8, (h + MShift[l] + 7) / 8); nextStage(cb); + } + recordBack(cb, stamps, stamps ? first + 3 : 0); +} +VkImage MEngine::outputImage() const { return impl->output.handle; } diff --git a/engine/test_dec0_m.py b/engine/test_dec0_m.py new file mode 100644 index 0000000..74da8e7 --- /dev/null +++ b/engine/test_dec0_m.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python3 +"""Run the engine's own preset M dec0 stage (dec0_m.comp) on a captured launch and compare its four surfaces. + +usage: test_dec0_m.py CAPTURE_DIR FRAME WORK_DIR KVK [repeats] +""" +import glob, os, re, struct, subprocess, sys +import numpy as np +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, HERE) +import m_model as mm +import model_enc3 as me +import swin_model as sm +from test_m import codes, token_codes +K = 'rrlite_dec0_4x4_folded' +cap, frame, work, kvk = sys.argv[1], int(sys.argv[2]), sys.argv[3], sys.argv[4] +reps = sys.argv[5] if len(sys.argv) > 5 else '1' +os.makedirs(work, exist_ok=True) +seq = sorted({int(re.match(r'launch-(\d+)-', os.path.basename(f)).group(1)) for f in glob.glob(f'{cap}/dump/launch-*-{K}-arg*')})[frame] +res = {int(re.search(r'-arg(\d+)-', f).group(1)): f for f in glob.glob(f'{cap}/dump/launch-{seq:03d}-{K}-arg*')} +d = f'{cap}/replay/replay-{seq:06d}-{K}' +P = open(d + '/args.bin', 'rb').read() +allocs = {} +for line in open(d + '/manifest.txt'): + q = line.split() + if q and q[0] == 'alloc': allocs[int(q[1])] = (int(q[2], 16), np.fromfile(f'{d}/alloc-{q[1]}.bin', np.uint8)) +find = lambda ptr: next((a, b) for a, b in allocs.values() if a <= ptr < a + b.size) +wptr, inp = struct.unpack_from(' 1 else 100 * (dd < 1e-3 * np.abs(a).max()).mean():5.1f}% max |diff| {dd.max():.4g}') diff --git a/engine/test_down_m.py b/engine/test_down_m.py new file mode 100644 index 0000000..76018a1 --- /dev/null +++ b/engine/test_down_m.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python3 +"""Run the engine's own preset M downsample stage (down_m.comp) on a captured launch and compare the output. + +usage: test_down_m.py CAPTURE_DIR FRAME WORK_DIR KVK [repeats] +""" +import glob, os, re, struct, subprocess, sys +import numpy as np +K = 'rrlite_downsample_kernel_static_hdr' +cap, frame, work, kvk = sys.argv[1], int(sys.argv[2]), sys.argv[3], sys.argv[4] +reps = sys.argv[5] if len(sys.argv) > 5 else '1' +HERE = os.path.dirname(os.path.abspath(__file__)) +os.makedirs(work, exist_ok=True) +seq = sorted({int(re.match(r'launch-(\d+)-', os.path.basename(f)).group(1)) for f in glob.glob(f'{cap}/dump/launch-*-{K}-arg*')})[frame] +res = {int(re.search(r'-arg(\d+)-', f).group(1)): f for f in glob.glob(f'{cap}/dump/launch-{seq:03d}-{K}-arg*')} +P = open(f'{cap}/replay/replay-{seq:06d}-{K}/args.bin', 'rb').read() +i2 = lambda o: struct.unpack_from('<2i', P, o); f2 = lambda o: struct.unpack_from('<2f', P, o) +expo = float(np.fromfile(res[48], np.float16)[0]) +open(f'{work}/params.bin', 'wb').write(struct.pack('<8i4ff', *i2(0), *i2(40), *i2(64), *i2(72), *f2(80), *f2(88), expo)) +dims = lambda f: [int(x) for x in re.search(r'-(\d+x\d+x\d+)-fmt', f).group(1).split('x')] +cw, ch, cc = dims(res[56]); ow, oh, oc = dims(res[32]) +open(f'{work}/manifest', 'w').write('\n'.join([ + f'param {work}/params.bin', f'grid {(ow + 7) // 8} {(oh + 63) // 64}', f'tex 1 {res[56]} {cw} {ch} {cc} 16 nearest', + f'surf 2 {ow} {oh} 4 16 - {work}/out.bin', f'lut {res[24]}']) + '\n') +spv = f'{work}/down_m.spv' +subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', f'{HERE}/down_m.comp', '-o', spv], check=True, stdout=subprocess.DEVNULL) +r = subprocess.run([kvk, spv, f'{work}/manifest', reps, '64'], capture_output=True, text=True) +if r.returncode: + print('FAILED', r.stderr[-400:]); sys.exit(1) +for l in r.stdout.split('\n'): + if 'us per' in l: print(l) +a = np.fromfile(res[32], np.float16).astype(np.float64); b = np.fromfile(f'{work}/out.bin', np.float16).astype(np.float64) +d = np.abs(a - b) +print(f'frame {frame} (launch {seq}) {ow}x{oh}: psnr {10 * np.log10(np.abs(a).max() ** 2 / max((d ** 2).mean(), 1e-30)):.1f} dB, ' + f'identical {100 * (np.fromfile(res[32], np.uint16) == np.fromfile(f"{work}/out.bin", np.uint16)).mean():.1f}%, max |diff| {d.max():.4g}') diff --git a/engine/test_enc0_m.py b/engine/test_enc0_m.py new file mode 100644 index 0000000..6babf25 --- /dev/null +++ b/engine/test_enc0_m.py @@ -0,0 +1,72 @@ +#!/usr/bin/env python3 +"""Run the engine's own preset M input stage (enc0_m.comp) on a captured launch and compare the tokens and the three +motion surfaces. + +usage: test_enc0_m.py CAPTURE_DIR FRAME WORK_DIR KVK [repeats] +The half-resolution previous result (texture 152) is not in the dumps; frames after the first use the engine's +own post stage output if WORK_DIR/half-prev.bin exists, zeros otherwise. +""" +import glob, os, re, struct, subprocess, sys +import numpy as np +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, HERE) +import m_model as mm +import model_enc3 as me +import swin_model as sm +from test_m import token_codes, report +K = 'rrlite_enc0_4x4_mvlo_hdr_folded' +cap, frame, work, kvk = sys.argv[1], int(sys.argv[2]), sys.argv[3], sys.argv[4] +reps = sys.argv[5] if len(sys.argv) > 5 else '1' +os.makedirs(work, exist_ok=True) +seq = sorted({int(re.match(r'launch-(\d+)-', os.path.basename(f)).group(1)) for f in glob.glob(f'{cap}/dump/launch-*-{K}-arg*')})[frame] +res = {int(re.search(r'-arg(\d+)-', f).group(1)): f for f in glob.glob(f'{cap}/dump/launch-{seq:03d}-{K}-arg*')} +d = f'{cap}/replay/replay-{seq:06d}-{K}' +P = open(d + '/args.bin', 'rb').read() +allocs = {} +for line in open(d + '/manifest.txt'): + q = line.split() + if q and q[0] == 'alloc': allocs[int(q[1])] = (int(q[2], 16), np.fromfile(f'{d}/alloc-{q[1]}.bin', np.uint8)) +find = lambda ptr: next((a, b) for a, b in allocs.values() if a <= ptr < a + b.size) +i2 = lambda o: struct.unpack_from('<2i', P, o); f2 = lambda o: struct.unpack_from('<2f', P, o); f1 = lambda o: struct.unpack_from('> shift) & 1023).astype(np.float32) / np.float32(1023) +values[:, 3] = (codes >> 30).astype(np.float32) / np.float32(3) +values.tofile(half) +lines = [f'param {work}/params.bin', f'grid {(gw + i2(8)[0] + 7) // 8} {(gh + i2(8)[1] + 7) // 8}', + f'tex 1 {res[80]} {rw} {rh} 4 16 nearest', f'tex 2 {res[96]} {rw} {rh} 2 16 nearest', f'tex 3 {res[136]} {rw} {rh} 1 32 nearest', + f'tex 4 {res[128]} {rw} {rh} 2 16 linear', f'tex 5 {half} {hw // 2} {hh // 2} 4 32 linear', f'tex 6 {res[160]} {rw} {rh} 4 9 linear'] +for i, off in enumerate((280, 288, 296)): lines.append(f'surf {7 + i} {rw} {rh} 2 16 - {work}/mv-{off}.bin') +lines += [f'lut {work}/w8.bin', f'lut {work}/w16.bin', f'lut {work}/noise.bin', f'outbuf {work}/tokens.bin {(gw // 2) * (gh // 2) * 64}'] +open(f'{work}/manifest', 'w').write('\n'.join(lines) + '\n') +spv = f'{work}/enc0_m.spv' +subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', f'{HERE}/enc0_m.comp', '-o', spv], check=True, stdout=subprocess.DEVNULL) +r = subprocess.run([kvk, spv, f'{work}/manifest', reps, '32'], capture_output=True, text=True) +if r.returncode: + print('FAILED', r.stderr[-600:]); sys.exit(1) +for l in r.stdout.split('\n'): + if 'us per' in l: print(l) +print(f'frame {frame} (launch {seq}) {rw}x{rh}, tokens {gw // 2}x{gh // 2}, reset {P[316]}, dither frame {i2(312)[0]}') +for off, name in ((280, 'A'), (288, 'B'), (296, 'C')): + a = np.fromfile(res[off], np.float16).astype(np.float64); b = np.fromfile(f'{work}/mv-{off}.bin', np.float16).astype(np.float64) + print(f' motion {name}: identical {100 * (a == b).mean():5.1f}% max |diff| {np.abs(a - b).max():.4g}') +after = np.fromfile(res[40], np.uint8); ab, _ = find(outp) +ref = token_codes(after, outp - ab, gw // 2, gh // 2, 64) +got = np.fromfile(f'{work}/tokens.bin', np.uint8).reshape(gh // 2, gw // 2, 64) +print(' ' + report('tokens', got, ref)) +vis = rh // 2 +print(' ' + report('tokens, rows inside the image', got[:vis], ref[:vis])) diff --git a/engine/test_frame.py b/engine/test_frame.py new file mode 100644 index 0000000..d11d26d --- /dev/null +++ b/engine/test_frame.py @@ -0,0 +1,44 @@ +#!/usr/bin/env python3 +"""Compare full native frame results, including the engine's own temporal feedback. + +test_frame.py CASE_DIR RESULT_DIR [--min-psnr 50] +Exits nonzero for missing data, nonfinite output, or a PSNR regression. +""" +import argparse +import json +from pathlib import Path +import numpy as np + + +def compare(reference, output): + a = np.fromfile(reference, np.float16).astype(np.float64) + b = np.fromfile(output, np.float16).astype(np.float64) + if a.shape != b.shape or not a.size: + raise ValueError(f'{output}: size differs from reference') + if not np.isfinite(a).all() or not np.isfinite(b).all(): + raise ValueError(f'{output}: nonfinite reference or result') + error = np.abs(a - b) + peak = max(np.abs(a).max(), 1e-9) + psnr = 10 * np.log10(peak**2 / max((error**2).mean(), 1e-30)) + return psnr, error.max(), error.mean() + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument('case', type=Path) + p.add_argument('results', type=Path) + p.add_argument('--min-psnr', type=float, default=50) + a = p.parse_args() + references = json.loads((a.case / 'references.json').read_text()) + if not references: raise ValueError('no reference frames') + failed = False + names = {'288': 'history color', '296': 'hi-res luma', '304': 'token feature', '312': 'final RGB'} + for i, refs in enumerate(references): + for off, reference in refs.items(): + psnr, maximum, mean = compare(reference, a.results / f'frame-{i}.{off}.bin') + print(f'frame {i} {names[off]:14s}: {psnr:6.1f} dB PSNR; max error {maximum:.5g}; mean {mean:.5g}') + failed |= psnr < a.min_psnr + if failed: raise SystemExit(f'frame validation failed: PSNR below {a.min_psnr:g} dB') + + +if __name__ == '__main__': main() diff --git a/engine/test_frame_m.py b/engine/test_frame_m.py new file mode 100644 index 0000000..72b5bcb --- /dev/null +++ b/engine/test_frame_m.py @@ -0,0 +1,150 @@ +#!/usr/bin/env python3 +"""Run whole preset M frames on the engine's own stages and compare the output with the capture. + +usage: test_frame_m.py CAPTURE_DIR WORK_DIR KVK EVK [frames] +Every stage runs as its own process (kvk for the texture stages, evk for the network layers). Only the frame's +colour, motion and depth, the exposure values and the kernels' parameter blocks come from the capture; all +intermediate data and all temporal history are the engine's own. +""" +import glob, os, re, struct, subprocess, sys +import numpy as np +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, HERE) +import m_model as mm +import swin_model as sm +from test_m import pack_layer, defines, token_codes, report +cap, work, kvk, evk = sys.argv[1:5] +frames = int(sys.argv[5]) if len(sys.argv) > 5 else 2 +os.makedirs(work, exist_ok=True) +LDR = os.environ.get('LDR') == '1' # a capture without the HDR flag: NGX's *_ldr_* kernels, the shaders' LDR paths +R = 'ldr' if LDR else 'hdr' +# below a reconstruction/render ratio of 3 the capture holds NGX's rrlite_post_3_2; the engine then needs the +# larger shared tile (post_m.comp's TILE_W/TILE_H) +POST = os.environ.get('POST_KERNEL') or next((k for k in (f'rrlite_post_3_1_mvlo_{R}_folded', f'rrlite_post_3_2_mvlo_{R}_folded') + if glob.glob(f'{cap}/dump/launch-*-{k}-arg*')), f'rrlite_post_3_1_mvlo_{R}_folded') +POST_TILE = ['-DTILE_W=15', '-DTILE_H=15'] if '3_2' in POST else [] +NAMES = {'enc0': f'rrlite_enc0_4x4_mvlo_{R}_folded', 'dec0': 'rrlite_dec0_4x4_folded', 'post': POST, 'down': f'rrlite_downsample_kernel_static_{R}'} +STAGE = ['-DLDR=1'] if LDR else [] + + +class Launch: + def __init__(self, kernel, frame): + seqs = sorted({int(re.match(r'launch-(\d+)-', os.path.basename(f)).group(1)) for f in glob.glob(f'{cap}/dump/launch-*-{kernel}-arg*')}) + self.seq = seqs[frame] + self.res = {int(re.search(r'-arg(\d+)-', f).group(1)): f for f in glob.glob(f'{cap}/dump/launch-{self.seq:03d}-{kernel}-arg*')} + self.dir = f'{cap}/replay/replay-{self.seq:06d}-{kernel}' + self.P = open(self.dir + '/args.bin', 'rb').read() + self.allocs = {} + for line in open(self.dir + '/manifest.txt'): + q = line.split() + if q and q[0] == 'alloc': self.allocs[int(q[1])] = (int(q[2], 16), np.fromfile(f'{self.dir}/alloc-{q[1]}.bin', np.uint8)) + i2 = lambda s, o: struct.unpack_from('<2i', s.P, o) + f2 = lambda s, o: struct.unpack_from('<2f', s.P, o) + f1 = lambda s, o: struct.unpack_from(' + +namespace { +constexpr unsigned rotatingMasks[] = {15u, 14u, 13u, 11u, 15u, 8u}; +struct TestImage { + Context& context; + VkImage image = VK_NULL_HANDLE; + VkImageView view = VK_NULL_HANDLE; + VkDeviceMemory memory = VK_NULL_HANDLE; + VkFormat format; + VkImageAspectFlags aspect; + uint32_t w, h, channels, bits; + Buffer upload; + TestImage(Context& c, VkFormat f, uint32_t width, uint32_t height, uint32_t nc, uint32_t nb, + VkImageAspectFlags a = VK_IMAGE_ASPECT_COLOR_BIT) + : context(c), format(f), aspect(a), w(width), h(height), channels(nc), bits(nb), + upload(c, VkDeviceSize(w) * h * nc * nb / 8) { + VkImageCreateInfo info{VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO}; + info.imageType = VK_IMAGE_TYPE_2D; info.format = f; info.extent = {w,h,1}; + info.mipLevels = info.arrayLayers = 1; info.samples = VK_SAMPLE_COUNT_1_BIT; + info.usage = VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT; + if (f == VK_FORMAT_R16G16B16A16_SFLOAT || f == VK_FORMAT_R32G32B32A32_SFLOAT || + f == VK_FORMAT_A2B10G10R10_UNORM_PACK32) info.usage |= VK_IMAGE_USAGE_STORAGE_BIT; + CK(vkCreateImage(c.device, &info, nullptr, &image)); + VkMemoryRequirements requirements; vkGetImageMemoryRequirements(c.device, image, &requirements); + memory = c.allocate(requirements, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT); + CK(vkBindImageMemory(c.device, image, memory, 0)); + VkImageViewCreateInfo vi{VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO}; + vi.image = image; vi.viewType = VK_IMAGE_VIEW_TYPE_2D; vi.format = f; + vi.subresourceRange = {a,0,1,0,1}; CK(vkCreateImageView(c.device, &vi, nullptr, &view)); + } + ~TestImage() { + if (view) vkDestroyImageView(context.device, view, nullptr); + if (image) vkDestroyImage(context.device, image, nullptr); + if (memory) vkFreeMemory(context.device, memory, nullptr); + } + void put(uint32_t x, uint32_t y, uint32_t channel, float value) { + const size_t i = (size_t(y)*w+x)*channels+channel; + if (bits == 32) static_cast(upload.map)[i] = value; + else static_cast<_Float16*>(upload.map)[i] = static_cast<_Float16>(value); + } + void initialize(VkImageLayout layout) { + VkImageMemoryBarrier b{VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER}; b.image = image; + b.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED; b.newLayout = VK_IMAGE_LAYOUT_GENERAL; + b.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; + b.srcQueueFamilyIndex = b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; b.subresourceRange = {aspect,0,1,0,1}; + vkCmdPipelineBarrier(context.command, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, + 0,0,nullptr,0,nullptr,1,&b); + auto r = region(w,h); r.imageSubresource.aspectMask = aspect; + vkCmdCopyBufferToImage(context.command, upload.handle, image, VK_IMAGE_LAYOUT_GENERAL, 1, &r); + if (layout != VK_IMAGE_LAYOUT_GENERAL) { + b.oldLayout = VK_IMAGE_LAYOUT_GENERAL; b.newLayout = layout; + b.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; b.dstAccessMask = VK_ACCESS_SHADER_READ_BIT; + vkCmdPipelineBarrier(context.command, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, + 0,0,nullptr,0,nullptr,1,&b); + } + } + d4r::GameImage game(uint32_t x, uint32_t y, bool direct, VkImageLayout layout) const { + return {image,format,layout,aspect,x,y,direct ? view : VK_NULL_HANDLE}; + } +}; +struct Case { + uint32_t ow = 1283, oh = 723; + bool depthAspect = false, wideColor = false, wideMotion = false, wideOutput = false, readonly = false; + bool rgb10Output = false, rotateOutputFormats = false; + bool regularDepth = false, swapDepth = false; // K: depth convention; near and far values exchanged + float exposureScale = 1, preExposure = 1; // K: the game's factors on the measured exposure + bool displayMotion = false; // K: motion vectors at output resolution, in output pixels + bool gameExposure = false; float exposureValue = 1; // K: the game's 1x1 exposure texture instead of measurement +}; +void validateFrames(const Buffer& results, VkDeviceSize stride, const std::array, 3>& output, unsigned count) { + for (unsigned frame = 0; frame < count; ++frame) { + const auto& image = *output[frame % output.size()]; + if (image.format == VK_FORMAT_A2B10G10R10_UNORM_PACK32) continue; + const auto* data = static_cast(results.map) + frame * stride; + for (VkDeviceSize i = 0; i < image.upload.size / (image.bits / 8); ++i) + if (!std::isfinite(image.bits == 32 ? reinterpret_cast(data)[i] : float(reinterpret_cast(data)[i]))) + throw std::runtime_error("nonfinite adapter output"); + } +} +std::vector evaluate(Context& c, d4r::Model& model, const Case& test, unsigned mask, + bool queued, bool bareEngine = false, bool rotateInputs = false) { + constexpr uint32_t rw = 640, rh = 360, count = 6, slots = 3; + std::array, slots> color, motion, depth, output; + TestImage exposureTexture(c, VK_FORMAT_R32_SFLOAT, 1, 1, 1, 32); + exposureTexture.put(0,0,0,test.exposureValue); + const auto inputLayout = test.readonly ? VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL : VK_IMAGE_LAYOUT_GENERAL; + for (unsigned slot = 0; slot < slots; ++slot) { + color[slot] = std::make_unique(c, test.wideColor ? VK_FORMAT_R32G32B32A32_SFLOAT : VK_FORMAT_R16G16B16A16_SFLOAT, rw+16,rh+16,4,test.wideColor ? 32 : 16); + const uint32_t mw = test.displayMotion ? test.ow+16 : rw+16, mh = test.displayMotion ? test.oh+16 : rh+16; + motion[slot] = std::make_unique(c, test.wideMotion ? VK_FORMAT_R32G32_SFLOAT : VK_FORMAT_R16G16_SFLOAT, mw,mh,2,test.wideMotion ? 32 : 16); + depth[slot] = std::make_unique(c, test.depthAspect ? VK_FORMAT_D32_SFLOAT : VK_FORMAT_R32_SFLOAT, rw+16,rh+16,1,32, + test.depthAspect ? VK_IMAGE_ASPECT_DEPTH_BIT : VK_IMAGE_ASPECT_COLOR_BIT); + const bool packed = test.rgb10Output && (!test.rotateOutputFormats || slot != 1); + output[slot] = std::make_unique(c, packed ? VK_FORMAT_A2B10G10R10_UNORM_PACK32 : + test.wideOutput ? VK_FORMAT_R32G32B32A32_SFLOAT : VK_FORMAT_R16G16B16A16_SFLOAT, + test.ow+16,test.oh+16,packed ? 1 : 4,packed || test.wideOutput ? 32 : 16); + for (uint32_t y = 0; y < rh+16; ++y) for (uint32_t x = 0; x < rw+16; ++x) { + const int cx = int(x)-7, cy = int(y)-9; + const bool inside = cx >= 0 && cy >= 0 && cx < int(rw) && cy < int(rh); + for (unsigned ch = 0; ch < 4; ++ch) color[slot]->put(x,y,ch, ch == 3 ? 1.f : + inside ? .1f + ((cx + 3*cy + int(slot)*13 + int(ch)*7)%19)*.08f + (test.wideColor ? .00012345f : 0) : 20.f); + depth[slot]->put(x,y,0, (((int(x)-11)/13)%2 != 0) != test.swapDepth ? .2f : .8f); + } + for (uint32_t y = 0; y < mh; ++y) for (uint32_t x = 0; x < mw; ++x) { + const int mx = int(x)-3, my = int(y)-5; + motion[slot]->put(x,y,0, .25f * ((mx+int(slot))%7-3) + (test.wideMotion ? .00012345f : 0)); + motion[slot]->put(x,y,1, .25f * ((my+int(slot))%9-4) + (test.wideMotion ? .00032109f : 0)); + } + std::memset(output[slot]->upload.map,0,output[slot]->upload.size); + } + // The borrowed views above outlive both the adapter/engine and their immutable descriptor cache. + d4r::GameUpscaler adapter(model,c.physical,c.device,test.ow,test.oh,rw,rh); + auto shape = d4r::makeKShape(test.ow,test.oh); + std::unique_ptr engine; + if (bareEngine) engine = std::make_unique(model,shape); + const auto bytes = std::max({output[0]->upload.size, output[1]->upload.size, output[2]->upload.size}); + Buffer results(c,bytes*count); + std::memset(results.map,0,results.size); + c.begin(); + for (unsigned slot = 0; slot < slots; ++slot) { + color[slot]->initialize(inputLayout); motion[slot]->initialize(inputLayout); + depth[slot]->initialize(inputLayout); output[slot]->initialize(VK_IMAGE_LAYOUT_GENERAL); + } + exposureTexture.initialize(inputLayout); + if (engine) engine->recordInitialize(c.command); + c.submit(); + if (queued) c.begin(); + for (unsigned frame = 0; frame < count; ++frame) { + unsigned slot = frame%slots; + if (!queued) c.begin(); + const unsigned frameMask = rotateInputs ? rotatingMasks[frame] : mask; + d4r::GameFrame f; + f.color = color[slot]->game(7,9,frameMask&1,inputLayout); + f.motion = motion[slot]->game(3,5,frameMask&2,inputLayout); + f.depth = depth[slot]->game(11,13,frameMask&4,inputLayout); + f.output = output[slot]->game(5,7,frameMask&8,VK_IMAGE_LAYOUT_GENERAL); + if (test.gameExposure) f.exposure = exposureTexture.game(0,0,true,inputLayout); + f.settings.renderWidth = rw-frame%2*8; f.settings.renderHeight = rh-frame%2*4; + f.settings.reset = frame == 0 || frame == 4; f.settings.autoExposure = !test.gameExposure; + f.settings.depthInverted = !test.regularDepth; f.settings.exposureScale = test.exposureScale; f.settings.preExposure = test.preExposure; + f.settings.displayMotion = test.displayMotion; + f.settings.jitter[0] = frame%2 ? .25f : -.125f; f.settings.jitter[1] = frame%3 ? .375f : -.25f; + if (engine) { + auto p = d4r::makeKFrameParams(shape,f.settings); + const int origins[3][2] = {{7,9},{3,5},{11,13}}; + int32_t* in[] = {p.input.colOrigin,p.input.mvOrigin,p.input.depOrigin}; + int32_t* out[] = {p.output.colOrigin,p.output.mvOrigin,p.output.depOrigin}; + int32_t* hi[] = {p.output.colMax,p.output.mvMax,p.output.depMax}; + for (unsigned k=0;k<3;++k) for(unsigned a=0;a<2;++a) {in[k][a]=out[k][a]=origins[k][a];hi[k][a]+=origins[k][a];} + p.output.outOrigin[0]=5;p.output.outOrigin[1]=7; + d4r::ExposureControl exposure; exposure.origin[0]=7;exposure.origin[1]=9; + exposure.size[0]=int32_t(f.settings.renderWidth);exposure.size[1]=int32_t(f.settings.renderHeight); + exposure.measure=frame%2==0||f.settings.reset;exposure.reset=f.settings.reset; + exposure.exposureScale=test.exposureScale;exposure.preExposure=test.preExposure; + exposure.game=test.gameExposure; + auto windows=d4r::makeKWindows(shape,frame); + engine->setFrame({f.color.view,f.motion.view,f.depth.view,inputLayout,inputLayout,inputLayout,f.exposure.view,inputLayout},p.input,p.output,&windows,&exposure,{f.output.view,f.output.format}); + engine->recordFrame(c.command); + } else adapter.record(c.command,f); + auto r=region(output[slot]->w,output[slot]->h);r.bufferOffset=frame*bytes; + vkCmdCopyImageToBuffer(c.command,output[slot]->image,VK_IMAGE_LAYOUT_GENERAL,results.handle,1,&r); + if (!queued) { + barrier(c.command,VK_PIPELINE_STAGE_TRANSFER_BIT,VK_PIPELINE_STAGE_HOST_BIT,VK_ACCESS_TRANSFER_WRITE_BIT,VK_ACCESS_HOST_READ_BIT); + c.submit(); + } + } + if (queued) { + barrier(c.command,VK_PIPELINE_STAGE_TRANSFER_BIT,VK_PIPELINE_STAGE_HOST_BIT,VK_ACCESS_TRANSFER_WRITE_BIT,VK_ACCESS_HOST_READ_BIT); + c.submit(); + } + validateFrames(results,bytes,output,count); + return {static_cast(results.map),static_cast(results.map)+results.size}; +} +std::vector evaluateM(Context& c, d4r::MModel& model, const Case& test, unsigned mask, bool queued, bool rotateInputs = false) { + constexpr uint32_t rw = 640, rh = 360, ow = 960, oh = 540, count = 6, slots = 3; + std::array, slots> color, motion, depth, output; + const auto layout = test.readonly ? VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL : VK_IMAGE_LAYOUT_GENERAL; + for (unsigned slot = 0; slot < slots; ++slot) { + color[slot] = std::make_unique(c, VK_FORMAT_R16G16B16A16_SFLOAT, rw,rh,4,16); + motion[slot] = std::make_unique(c, VK_FORMAT_R16G16_SFLOAT, rw,rh,2,16); + depth[slot] = std::make_unique(c, test.depthAspect ? VK_FORMAT_D32_SFLOAT : VK_FORMAT_R32_SFLOAT, rw,rh,1,32, + test.depthAspect ? VK_IMAGE_ASPECT_DEPTH_BIT : VK_IMAGE_ASPECT_COLOR_BIT); + const bool packed = test.rgb10Output && (!test.rotateOutputFormats || slot != 1); + output[slot] = std::make_unique(c, packed ? VK_FORMAT_A2B10G10R10_UNORM_PACK32 : + test.wideOutput ? VK_FORMAT_R32G32B32A32_SFLOAT : VK_FORMAT_R16G16B16A16_SFLOAT, + ow,oh,packed ? 1 : 4,packed || test.wideOutput ? 32 : 16); + for (uint32_t y = 0; y < rh; ++y) for (uint32_t x = 0; x < rw; ++x) { + for (unsigned ch = 0; ch < 4; ++ch) + color[slot]->put(x,y,ch,ch == 3 ? 1.f : .1f + ((x+3*y+slot*13+ch*7)%19)*.08f); + motion[slot]->put(x,y,0,.25f*(int((x+slot)%7)-3)); + motion[slot]->put(x,y,1,.25f*(int((y+slot)%9)-4)); + depth[slot]->put(x,y,0,(x/13)%2 ? .2f : .8f); + } + std::memset(output[slot]->upload.map,0,output[slot]->upload.size); + } + d4r::GameUpscalerM adapter(model,c.physical,c.device,ow,oh,rw,rh); + const auto bytes = std::max({output[0]->upload.size, output[1]->upload.size, output[2]->upload.size}); + Buffer results(c,bytes*count); + std::memset(results.map,0,results.size); + c.begin(); + for (unsigned slot = 0; slot < slots; ++slot) { + color[slot]->initialize(layout); motion[slot]->initialize(layout); + depth[slot]->initialize(layout); output[slot]->initialize(VK_IMAGE_LAYOUT_GENERAL); + } + c.submit(); + if (queued) c.begin(); + for (unsigned frame = 0; frame < count; ++frame) { + const unsigned slot = frame%slots; + if (!queued) c.begin(); + d4r::GameFrameM f; + const unsigned frameMask = rotateInputs ? rotatingMasks[frame] : mask; + f.color=color[slot]->game(0,0,frameMask&1,layout); f.motion=motion[slot]->game(0,0,frameMask&2,layout); + f.depth=depth[slot]->game(0,0,frameMask&4,layout); f.output=output[slot]->game(0,0,frameMask&8,VK_IMAGE_LAYOUT_GENERAL); + f.reset=frame==0||frame==4; + f.jitter[0]=frame%2 ? .25f : -.125f; f.jitter[1]=frame%3 ? .375f : -.25f; + adapter.record(c.command,f); + auto r=region(ow,oh); r.bufferOffset=frame*bytes; + vkCmdCopyImageToBuffer(c.command,output[slot]->image,VK_IMAGE_LAYOUT_GENERAL,results.handle,1,&r); + if (!queued) { + barrier(c.command,VK_PIPELINE_STAGE_TRANSFER_BIT,VK_PIPELINE_STAGE_HOST_BIT,VK_ACCESS_TRANSFER_WRITE_BIT,VK_ACCESS_HOST_READ_BIT); + c.submit(); + } + } + if (queued) { + barrier(c.command,VK_PIPELINE_STAGE_TRANSFER_BIT,VK_PIPELINE_STAGE_HOST_BIT,VK_ACCESS_TRANSFER_WRITE_BIT,VK_ACCESS_HOST_READ_BIT); + c.submit(); + } + validateFrames(results,bytes,output,count); + return {static_cast(results.map),static_cast(results.map)+results.size}; +} +void expectEqual(const std::vector& expected,const std::vector& actual,const std::string& name) { + if(expected!=actual) throw std::runtime_error(name+": output mismatch"); + std::cout< 5 else '1' +HERE = os.path.dirname(os.path.abspath(__file__)) +os.makedirs(work, exist_ok=True) +seqs = sorted({int(re.match(r'launch-(\d+)-', os.path.basename(f)).group(1)) for f in glob.glob(f'{cap}/dump/launch-*-{K}-arg*')}) +seq = seqs[frame] +res = {} +for f in glob.glob(f'{cap}/dump/launch-{seq:03d}-{K}-arg*'): + m = re.search(r'-arg(\d+)-(texture|surface|buffer)-(.*)-fmt(-?\d+)\.bin$', f) + res[int(m.group(1))] = (m.group(2), m.group(3), int(m.group(4)), f) +P = open(f'{cap}/replay/replay-{seq:06d}-{K}/args.bin', 'rb').read() +f2 = lambda o: struct.unpack_from('<2f', P, o) +i2 = lambda o: struct.unpack_from('<2i', P, o) +f1 = lambda o: struct.unpack_from(' 6 else '1' +HERE = os.path.dirname(os.path.abspath(__file__)) +os.makedirs(f'{work}/ref', exist_ok=True) +seqs = sorted({int(re.match(r'launch-(\d+)-', os.path.basename(f)).group(1)) for f in glob.glob(f'{cap}/dump/launch-*-{K}-arg*')}) +seq = seqs[int(frame)] +res = {} +for f in glob.glob(f'{cap}/dump/launch-{seq:03d}-{K}-arg*'): + m = re.search(r'-arg(\d+)-(texture|surface|buffer)-(.*)-fmt(-?\d+)\.bin$', f) + res[int(m.group(1))] = [m.group(2), m.group(3), int(m.group(4)), f] +override = '' +if seed != 'none': + rng = np.random.default_rng(int(seed)) + def smooth(h, w, cells, amp): + g = rng.uniform(-amp, amp, (cells + 1, cells + 1)) + y = np.linspace(0, cells, h)[:, None]; x = np.linspace(0, cells, w)[None, :] + y0 = np.minimum(y.astype(int), cells - 1); x0 = np.minimum(x.astype(int), cells - 1); fy = y - y0; fx = x - x0 + return g[y0, x0] * (1 - fy) * (1 - fx) + g[y0, x0 + 1] * (1 - fy) * fx + g[y0 + 1, x0] * fy * (1 - fx) + g[y0 + 1, x0 + 1] * fy * fx + w, h, c = (int(x) for x in res[184][1].split('x')) + mv = np.stack([smooth(h, w, 6, 3.0), smooth(h, w, 6, 3.0)], axis=-1) + mv[rng.uniform(size=(h, w)) < 0.02] += rng.uniform(-8, 8, 2) + mv.astype(np.float16).tofile(f'{work}/mv.bin'); res[184][3] = f'{work}/mv.bin' + w, h, c = (int(x) for x in res[192][1].split('x')) + depth = 0.3 + 0.1 * smooth(h, w, 5, 1.0); depth[smooth(h, w, 12, 1.0) > 0.35] += 0.4 + depth.astype(np.float32).tofile(f'{work}/depth.bin'); res[192][3] = f'{work}/depth.bin' + override = f'184={work}/mv.bin,192={work}/depth.bin' +rd = f'{cap}/replay/replay-{seq:06d}-{K}' +grid = next(l for l in open(f'{rd}/manifest.txt') if l.startswith('launch')).split() +n = os.path.getsize(res[240][3]) +open(f'{work}/ref/zero.bin', 'wb').write(bytes(n)) +lines = [f'param {rd}/args.bin', f'grid {grid[1]} {grid[2]}'] +for off in (168, 176, 184, 192, 200, 208, 216, 224): + kind, shape, fmt, f = res[off]; w, h, c = (int(x) for x in shape.split('x')) + # NGX's texture objects: history colour and token feature linear, the rest point + lines.append(f'tex {off} {f} {w} {h} {c} {32 if fmt == 32 else 16} {"linear" if off in (176, 224) else "nearest"}') +lines += [f'buf 240 {work}/ref/zero.bin 0', f'bufout 240 {work}/ref/out.bin'] +open(f'{work}/ref/manifest', 'w').write('\n'.join(lines) + '\n') +subprocess.run([kvk, ref_spv, f'{work}/ref/manifest', '1', '32'], check=True, capture_output=True) +a = np.fromfile(res[240][3], np.uint16); b = np.fromfile(f'{work}/ref/out.bin', np.uint16) +if seed == 'none': + print(f'translated reference vs captured buffer: {int((a != b).sum())} of {a.size} values differ') +env = dict(os.environ, IN_REF=f'{work}/ref/out.bin') +if override: env['IN_TEX_OVERRIDE'] = override +print(subprocess.run([sys.executable, f'{HERE}/test_in.py', cap, frame, work, kvk], env=env, capture_output=True, text=True).stdout.strip()) diff --git a/engine/test_k.py b/engine/test_k.py new file mode 100644 index 0000000..ad2642e --- /dev/null +++ b/engine/test_k.py @@ -0,0 +1,144 @@ +#!/usr/bin/env python3 +"""Run the engine's own preset K layers on a captured frame and compare with what NGX's pipeline produced. + +usage: test_k.py CAPTURE_DIR MODEL_PREFIX WORK_DIR [--chain] [--time N] [layer ...] + CAPTURE_DIR/replay: layer launches before they ran (inputs, parameters); CAPTURE_DIR/dump: buffers after each launch. + --chain feeds each layer the engine's own previous outputs instead of the captured inputs. +""" +import glob, json, os, re, struct, subprocess, sys +import numpy as np +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, os.path.join(HERE, '..', 'kernels', 'tools')) +import pwin_model as pm +F16 = np.float16 +ORDER = ['enc0', 'enc1', 'enc2', 'enc3', 'enc4', 'dec5', 'dec4', 'dec3', 'dec2', 'dec1', 'dec0'] + +def params(P, e): + W, H = P.i32x2(0); sx, sy = P.i32x2(56) + o = [0] * 108 + o[0], o[1], o[2], o[3] = e['g1'], e['bo'], e['g2'], e['b2'] + for h in range(e['H']): + o[4 + 5 * h:7 + 5 * h] = e['qkv'][3 * h:3 * h + 3]; o[7 + 5 * h] = e['table'][h]; o[8 + 5 * h] = e['wo'][h] + for c in range(e['C'] // 8): + o[44 + 3 * c], o[45 + 3 * c], o[46 + 3 * c] = e['w1'][c], e['b1'][c], e['w2'][c] + k = e['kind'] + if k in ('enc', 'enc0'): + o[104], o[105] = e['merge_w'], e['merge_b'] + if k in ('dec', 'dec0'): + o[104], o[105] = e['expand_w'], e['expand_b'] + if k == 'enc0': + o[106], o[107] = e['embed_w'], e['embed_b'] + if k == 'dec0': + o[106], o[107] = e['head_w'], e['head_b'] + return struct.pack('<4i108I', W, H, sx, sy, *o), (W, H) + +def defines(e, extra=()): + k = e['kind'] + kind = {'plain': 0, 'enc': 1, 'enc0': 1, 'dec': 2, 'dec0': 2}[k] + n16 = e.get('merge_n', 16) + return [f'-DC={e["C"]}', f'-DNH={e["H"]}', f'-DKIND={kind}', f'-DPOSATTN={int(e["posattn"])}', f'-DEMBED={int(k == "enc0")}', + f'-DHEAD={int(k == "dec0")}', f'-DXC={e["X"] or 16}', f'-DMERGE_N={n16}', f'-DFOLD_GAINS={int(e.get("fold_gains", True))}', + *extra] + +def ulps(m, g): + m, g = m.astype(np.float64), g.astype(np.float64) + scale = np.maximum(np.abs(m), 2.0 ** -10) + return np.abs(m - g) / 2.0 ** (np.floor(np.log2(scale)) - 10) + +def report(tag, ref, got): + u = ulps(ref, got); r, g = ref.astype(np.float64), got.astype(np.float64) + if not np.abs(r).max(): + print(f' {tag:12s} reference is all zero'); return 0.0 + psnr = 10 * np.log10(np.abs(r).max() ** 2 / max(((r - g) ** 2).mean(), 1e-30)) + print(f' {tag:12s} {ref.size:9d} values: exact {100 * (ref.view(np.uint16) == got.view(np.uint16)).mean():5.1f}% <=1ulp {100 * (u <= 1).mean():5.1f}% ' + f'<=4ulp {100 * (u <= 4).mean():6.2f}% max {u.max():8.1f} ulp psnr {psnr:6.1f} dB finite {100 * np.isfinite(got.astype(np.float32)).mean():.0f}%') + return psnr + +def main(): + a = sys.argv[1:] + chain = '--chain' in a + if chain: a.remove('--chain') + reps = 1 + if '--time' in a: + i = a.index('--time'); reps = int(a[i + 1]); del a[i:i + 2] + cap, prefix, work = a[:3] + layers = a[3:] or ORDER + os.makedirs(work, exist_ok=True) + index = json.load(open(prefix + '.json')) + open(f'{work}/dummy.bin', 'wb').write(b'\0' * 16) + prev = {} # (kind of buffer, layer) -> file of the engine's own output + total = 0.0 + for name in layers: + e = index[name] + d = sorted(glob.glob(f'{cap}/replay/replay-*-dltss_pwin_{name}_layer'))[0] # the main launch: params and inputs + seq = int(re.search(r'replay-(\d+)-', d).group(1)) + P = pm.Dump(d) + pb, (W, H) = params(P, e) + open(f'{work}/{name}.params', 'wb').write(pb) + spv = f'{work}/{name}.spv' + subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *defines(e), f'{HERE}/layer.comp', '-o', spv], check=True, stdout=subprocess.DEVNULL) + def slice_in(off, fn): + buf, o = P.ptr(off); buf[o:].tofile(fn); return fn + k, C, X = e['kind'], e['C'], e['X'] + fin = slice_in(8, f'{work}/{name}.in'); fskip = f'{work}/dummy.bin' + if k in ('dec', 'dec0'): + fskip = slice_in(16, f'{work}/{name}.skip') + if chain: + src = {'enc1': 'enc0', 'enc2': 'enc1', 'enc3': 'enc2', 'enc4': 'enc3', 'dec5': 'enc4', 'dec4': 'dec5', 'dec3': 'dec4', 'dec2': 'dec3', 'dec1': 'dec2', 'dec0': 'dec1'} + skip = {'dec4': 'enc4', 'dec3': 'enc3', 'dec2': 'enc2', 'dec1': 'enc1', 'dec0': 'enc0'} # skip of decN: the pre-merge output of encN + if name in src and ('main', src[name]) in prev: + fin = prev[('main', src[name])] + if name in skip and ('full', skip[name]) in prev: + fskip = prev[('full', skip[name])] + outs = [] # (binding, param offset, bytes, tag) + if k in ('enc', 'enc0'): + outs = [(4, 24, (H // 2) * (W // 2) * X * 2, 'merged'), (5, 32, H * W * C * 2, 'full')] + elif k == 'dec0': + outs = [(4, 24, H * W * 40 * 2, 'head')] + else: + outs = [(4, 24, H * W * C * 2, 'full')] + cmd = [f'{work}/evk', spv, str(P.grid[0]), str(P.grid[1]), str(reps), f'in:0={work}/{name}.params', f'in:1={prefix}.bin', f'in:2={fin}', f'in:3={fskip}'] + got = {} + for b, off, n, tag in outs: + cmd.append(f'out:{b}={work}/{name}.{tag}:{n}'); got[tag] = (off, n) + if len(outs) == 1: + cmd.append(f'out:5={work}/{name}.unused:16') + r = subprocess.run(cmd, capture_output=True, text=True) + if r.returncode: + raise RuntimeError(f'{name} FAILED: {r.stderr[-300:]}') + t = re.search(r'([0-9.]+) us per dispatch', r.stdout) + if t: total += float(t.group(1)) + print(f'{name}: {e["kind"]} C={C} H={e["H"]} grid {P.grid[0]}x{P.grid[1]} tokens {W}x{H}' + (f' {t.group(1)} us' if t else '')) + for tag, (off, n) in got.items(): + # The post-launch dump holds each allocation once, named by its base address. A phase-split layer + # writes its output in a later launch of the same layer (kernels/k/pwin_phase.h: NAME_post3 writes + # out24 for the bottleneck, NAME_post4 the merged output), which the bridge launches right after the + # main kernel, and the main launch's dump can only hold the output allocation's previous content. + # Take the last of this layer's launches in this frame that dumped the output allocation. + addr = P.u64(off) + base = next(b for b, buf in P.allocs.values() if b <= addr < b + buf.size) + f, launches = None, [] + for n2, k2 in [(seq, P.kernel)] + [(seq + j, f'{P.kernel}_post{j}') for j in range(1, 9)]: + if not os.path.isdir(f'{cap}/replay/replay-{n2:06d}-{k2}'): + continue + q = glob.glob(f'{cap}/dump/launch-{n2:03d}-{k2}-arg*-buffer-base0x{base:x}+*') + launches.append(f'{k2}#{n2}' + ('' if q else '(no dump)')) + if q: f = q[0] + if f is None: + raise ValueError(f'{name}.{tag}: no launch of {name} dumped the output allocation') + o = addr - base + ref = np.fromfile(f, dtype=np.uint8)[o:o + n].view(F16) + if not np.abs(ref.astype(np.float32)).max(): + print(f' {tag:12s} reference is all zero: {name} launches in this frame: {", ".join(launches)}') + result = np.fromfile(f'{work}/{name}.{tag}', dtype=F16)[:n // 2] + if result.shape != ref.shape or not np.isfinite(result).all(): + raise ValueError(f'{name}.{tag}: wrong size or nonfinite output') + report(tag, ref, result) + prev[('main' if tag in ('merged', 'head') or k in ('plain', 'dec') else 'full', name)] = f'{work}/{name}.{tag}' + if tag == 'full' and k in ('enc', 'enc0'): + prev[('full', name)] = f'{work}/{name}.{tag}' + if reps > 1: + print(f'sum of layers: {total / 1000:.3f} ms') + +if __name__ == '__main__': + main() diff --git a/engine/test_m.py b/engine/test_m.py new file mode 100644 index 0000000..84e6d7b --- /dev/null +++ b/engine/test_m.py @@ -0,0 +1,166 @@ +#!/usr/bin/env python3 +"""Run the engine's own preset M layers (layer_m.comp) on a captured frame and compare with the captured outputs and +with the plain numpy model (m_model.py). + +usage: test_m.py CAPTURE_DIR WORK_DIR EVK [--frame N] [--time N] [--model] [--chain] [layer index ...] + CAPTURE_DIR/replay: layer launches before they ran; CAPTURE_DIR/dump: buffers after each launch. + --model also runs the numpy model (slow) and reports the shader against it. + --chain feeds each layer the engine's own previous outputs instead of the captured inputs. + +The layers' FP8 matrices are exported in the shader's blocked fixed-16 tile layout (pack8 below); the front-end +embedding/expansion weights that compile_m.py packages for enc0_m/dec0_m stay in NGX's own plain layout. +""" +import glob, os, struct, subprocess, sys +import numpy as np +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, HERE) +import m_model as mm +import model_enc3 as me + +F16 = np.float16 + + +def codes(values): + """e4m3 codes of exact e4m3 values""" + v = np.asarray(values, np.float64); a = np.abs(v) + c = np.searchsorted(me._POS, a).astype(np.uint8) + assert np.array_equal(me._POS[c], a), 'not an FP8 value' + return (c | np.where(np.signbit(v) & (a > 0), 128, 0)).astype(np.uint8) + + +class Pack: + def __init__(self, dtype, align): + self.parts, self.n, self.dtype, self.align = [], 0, dtype, align + + def add(self, a): + a = np.ascontiguousarray(a, self.dtype).ravel() + pad = (-self.n) % self.align + if pad: self.parts.append(np.zeros(pad, self.dtype)); self.n += pad + off = self.n; self.parts.append(a); self.n += a.size + return off + + def data(self): + return np.concatenate(self.parts + [np.zeros((-self.n) % 64 + 64, self.dtype)]) + + +def pack8(values): + """blocked FP8 bytes of a logical weight matrix W[K][N] (K, N multiples of 16), the layout layer_m.comp loads: + [K/16][N/16][n16][k16], i.e. reshape(K//16, 16, N//16, 16).transpose(0, 2, 3, 1). contiguous (no arithmetic). + A fixed 16x16 tile is 256 bytes and starts at 256*(kblock*(N//16) + nblock); the shader's A operand is the tile + [N16][K16] read row-major with stride 16 (A[n][k] = W[kblock*16 + k][nblock*16 + n]), so a K16 row block of + every output channel tile stays contiguous and the per-head slice of wo is still base + 32*h*C.""" + W = codes(values) + K, N = W.shape + assert K % 16 == 0 and N % 16 == 0, (K, N) + return np.ascontiguousarray(W.reshape(K // 16, 16, N // 16, 16).transpose(0, 2, 3, 1)).ravel() + + +def pack_layer(t, name): + """(o[72], FP8 bytes, f16 words) of one layer's tensors; the FP8 matrices are blocked by pack8()""" + C, heads, NPM, CIN = mm.LAYERS[name] + b8, b16 = Pack(np.uint8, 16), Pack(F16, 16) + tile = lambda v: np.tile(np.asarray(v, F16)[None, :], (16, 1)) + o = [0] * 72 + o[0], o[1], o[2], o[3] = b16.add(t['g1']), b16.add(tile(t['bo'])), b16.add(t['g2']), b16.add(tile(t['b2'])) + assert t['wo'].shape == (32 * heads, C), t['wo'].shape + wo = b8.add(pack8(t['wo'])) + for h in range(heads): + assert t['q'][h].shape == (C, 32) and t['v'][h].shape == (C, 32), (t['q'][h].shape, t['v'][h].shape) + o[4 + 4 * h], o[5 + 4 * h], o[6 + 4 * h], o[7 + 4 * h] = b8.add(pack8(t['q'][h])), b8.add(pack8(t['v'][h])), b16.add(t['bias'][h]), wo + 32 * h * C + for c in range(C // 8): + assert t['w1'][c].shape == (C, 32) and t['w2'][c].shape == (32, C), (t['w1'][c].shape, t['w2'][c].shape) + o[20 + 3 * c], o[21 + 3 * c], o[22 + 3 * c] = b8.add(pack8(t['w1'][c])), b16.add(tile(t['b1'][c])), b8.add(pack8(t['w2'][c])) + if NPM: + assert t['merge_w'].shape == (4 * C, NPM), t['merge_w'].shape + o[68], o[69] = b8.add(pack8(t['merge_w'])), b16.add(tile(t['merge_b'])) + if CIN: + assert t['expand_w'].shape == (CIN, 4 * C), t['expand_w'].shape + o[68], o[69] = b8.add(pack8(t['expand_w'])), b16.add(tile(t['expand_b'])) + return o, b8.data(), b16.data() + + +def defines(name): + C, heads, NPM, CIN = mm.LAYERS[name] + return [f'-DC={C}', f'-DNH={heads}', f'-DKIND={1 if NPM else 2 if CIN else 0}', f'-DXC={NPM or CIN or 16}'] + os.environ.get('M_DEFS', '').split() + + +def token_codes(buf, off, W, H, C): + a = buf[off:off + (C // 32) * H * W * 32].reshape(C // 32, H, W, 32) + return np.ascontiguousarray(a.transpose(1, 2, 0, 3).reshape(H, W, C)) + + +def report(tag, got, ref): + rank = lambda v: np.searchsorted(me._POS, np.abs(v)) * np.sign(v) + g, r = me.E4M3[got].astype(np.float64), me.E4M3[ref].astype(np.float64) + steps = np.abs(rank(g) - rank(r)) + return (f'{tag} {mm.psnr(g, r):5.1f} dB, identical {100 * (got == ref).mean():5.1f}%, within one step {100 * (steps <= 1).mean():5.1f}%') + + +def main(): + a = sys.argv[1:] + def opt(flag, default): + if flag in a: + i = a.index(flag); v = a[i + 1]; del a[i:i + 2]; return v + return default + frame, reps = int(opt('--frame', 0)), opt('--time', '1') + model = '--model' in a + if model: a.remove('--model') + chain = '--chain' in a + if chain: a.remove('--chain') + own = {} # arena address -> the engine's own tokens there + cap, work, evk = a[:3] + only = [int(x) for x in a[3:]] + os.makedirs(work, exist_ok=True) + Ps, dirs = mm.launches(cap, frame) + total = 0.0 + for i, (P, d, name) in enumerate(zip(Ps, dirs, mm.ORDER)): + if only and i not in only: continue + C, heads, NPM, CIN = mm.LAYERS[name] + seq = int(os.path.basename(d).split('-')[1]) + after = np.fromfile(glob.glob(f'{cap}/dump/launch-{seq:03d}-*-arg024-*')[0], np.uint8) + t = mm.plain_layer(P, name) + o, b8, b16 = pack_layer(t, name) + W, H = P.W, P.H + pre = f'{work}/{i}-{name}' + open(pre + '.params', 'wb').write(struct.pack('<4i72I', W, H, P.sx, P.sy, *o)) + b8.tofile(pre + '.w8'); b16.tofile(pre + '.w16') + def source(ptr, w, h_, c): + return own[ptr] if chain and ptr in own else token_codes(P.arena, ptr - P.abase, w, h_, c) + if CIN: + source(P.inp, W // 2, H // 2, CIN).tofile(pre + '.in') + source(P.p32, W, H, C).tofile(pre + '.skip') + else: + source(P.inp, W, H, C).tofile(pre + '.in') + np.zeros(64, np.uint8).tofile(pre + '.skip') + spv = pre + '.spv' + subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *defines(name), f'{HERE}/layer_m.comp', '-o', spv], check=True, stdout=subprocess.DEVNULL) + gx, gy = (W + P.sx + 7) // 8, (H + P.sy + 7) // 8 + msize = (H // 2) * (W // 2) * NPM if NPM else 64 + r = subprocess.run([evk, spv, str(gx), str(gy), reps, f'in:0={pre}.params', f'in:1={pre}.w8', f'in:2={pre}.w16', f'in:3={pre}.in', + f'in:4={pre}.skip', f'out:5={pre}.out:{H * W * C}', f'out:6={pre}.merged:{msize}'], capture_output=True, text=True) + if r.returncode: + print(name, 'FAILED', r.stderr[-600:]); sys.exit(1) + us = [float(l.split(':')[1].split()[0]) for l in r.stdout.split('\n') if 'us per dispatch' in l] + total += us[0] if us else 0 + got = np.fromfile(pre + '.out', np.uint8).reshape(H, W, C) + own[P.out] = got + ref = token_codes(after, P.out - P.abase, W, H, C) + line = f'{i} {name:9s} {W}x{H}x{C} {us[0] if us else 0:7.1f} us: ' + report('vs capture', got, ref) + if NPM: + gm = np.fromfile(pre + '.merged', np.uint8).reshape(H // 2, W // 2, NPM) + own[P.p48] = gm + line += '; ' + report('merged', gm, token_codes(after, P.p48 - P.abase, W // 2, H // 2, NPM)) + print(line, flush=True) + if model: + if CIN: + out, merged = mm.forward(t, name, mm.tokens(P.arena, P.inp - P.abase, W // 2, H // 2, CIN), mm.tokens(P.arena, P.p32 - P.abase, W, H, C), P.sx, P.sy) + else: + out, merged = mm.forward(t, name, mm.tokens(P.arena, P.inp - P.abase, W, H, C), None, P.sx, P.sy) + line = ' ' + report('shader vs numpy model', got, codes(out)) + if NPM: line += '; ' + report('merged', gm, codes(merged)) + print(line, flush=True) + print(f'total {total:.1f} us') + + +if __name__ == '__main__': + main() diff --git a/engine/test_model.py b/engine/test_model.py new file mode 100644 index 0000000..3792431 --- /dev/null +++ b/engine/test_model.py @@ -0,0 +1,51 @@ +"""Structural model corruption must be rejected before compiling/dispatching shaders.""" +import unittest +from compile_k import validate +from extract_k import LAYERS + + +def minimal_index(): + index = {} + for name, (kind, h, c, x, pos) in LAYERS.items(): + e = {'kind': kind, 'H': h, 'C': c, 'X': x, 'posattn': pos, + 'g1': 0, 'bo': 0, 'g2': 0, 'b2': 0, + 'qkv': [0] * (h * 3), 'table': [0] * h, 'wo': [0] * h, + 'w1': [0] * (c // 8), 'b1': [0] * (c // 8), 'w2': [0] * (c // 8)} + if kind in ('enc', 'enc0'): e.update(merge_n=x, merge_w=0, merge_b=0) + if kind in ('dec', 'dec0'): e.update(expand_w=0, expand_b=0) + if kind == 'enc0': e.update(embed_w=0, embed_b=0) + if kind == 'dec0': e.update(head_w=0, head_b=0) + index[name] = e + return index + + +class ModelTests(unittest.TestCase): + def setUp(self): self.index = minimal_index() + def test_shared_zero_weights_are_structurally_valid(self): + validate(self.index, 4 * 1024 * 1024) + def test_matrix_end_past_blob_is_rejected(self): + self.index['dec4']['expand_w'] = 4 * 1024 * 1024 // 2 - 16 + with self.assertRaisesRegex(ValueError, 'outside'): validate(self.index, 4 * 1024 * 1024) + def test_weight_alignment_is_checked(self): + self.index['enc0']['embed_w'] = 1 + with self.assertRaisesRegex(ValueError, 'alignment'): validate(self.index, 4 * 1024 * 1024) + def test_wrong_network_architecture_is_rejected(self): + self.index['enc3']['H'] = 2 + with self.assertRaisesRegex(ValueError, 'shape'): validate(self.index, 4 * 1024 * 1024) + def test_missing_merge_is_rejected(self): + del self.index['enc2']['merge_w'] + with self.assertRaisesRegex(ValueError, 'missing'): validate(self.index, 4 * 1024 * 1024) + def test_incomplete_head_set_is_rejected(self): + self.index['dec5']['qkv'].pop() + with self.assertRaisesRegex(ValueError, 'expected'): validate(self.index, 4 * 1024 * 1024) + def test_wrong_layer_order_is_rejected(self): + index = dict(reversed(list(self.index.items()))) + with self.assertRaisesRegex(ValueError, 'order'): validate(index, 4 * 1024 * 1024) + def test_odd_blob_bytes_are_rejected(self): + with self.assertRaisesRegex(ValueError, 'f16'): validate(self.index, 4 * 1024 * 1024 - 1) + def test_invalid_merge_stride_is_rejected(self): + self.index['enc2']['merge_n'] = 128 + with self.assertRaisesRegex(ValueError, 'stride'): validate(self.index, 4 * 1024 * 1024) + + +if __name__ == '__main__': unittest.main() diff --git a/engine/test_out.py b/engine/test_out.py new file mode 100644 index 0000000..f365379 --- /dev/null +++ b/engine/test_out.py @@ -0,0 +1,80 @@ +#!/usr/bin/env python3 +"""Run the engine's own K output kernel on a captured launch of NVIDIA's and compare the four output surfaces. + +usage: test_out.py CAPTURE_DIR FRAME WORK_DIR KVK [repeats] + CAPTURE_DIR: dump/ and replay/ of a run with the output kernel captured (see vulkan/README.md); KVK: the vulkan/kvk binary. +""" +import glob, os, re, struct, subprocess, sys +import numpy as np +LDR = os.environ.get('LDR') == '1' # the *_ldr_* kernel variant and the shader's LDR path +K = 'hiluma_engine_output_depthinv_mvlo_ldr_max_v2_rel' if LDR else 'hiluma_engine_output_depthinv_mvlo_hdr_max_v2_rel' +cap, frame, work, kvk = sys.argv[1], int(sys.argv[2]), sys.argv[3], sys.argv[4] +reps = sys.argv[5] if len(sys.argv) > 5 else '1' +HERE = os.path.dirname(os.path.abspath(__file__)) +os.makedirs(work, exist_ok=True) +seqs = sorted({int(re.match(r'launch-(\d+)-', os.path.basename(f)).group(1)) for f in glob.glob(f'{cap}/dump/launch-*-{K}-arg*')}) +seq = seqs[frame] +res = {} +for f in glob.glob(f'{cap}/dump/launch-{seq:03d}-{K}-arg*'): + m = re.search(r'-arg(\d+)-(texture|surface|buffer)-(.*)-fmt(-?\d+)\.bin$', f) + res[int(m.group(1))] = (m.group(2), m.group(3), int(m.group(4)), f) +rd = f'{cap}/replay/replay-{seq:06d}-{K}' +P = open(f'{rd}/args.bin', 'rb').read() +man = open(f'{rd}/manifest.txt').read().split('\n') +grid = next(l for l in man if l.startswith('launch')).split() +ptr = {int(l.split()[1]): int(l.split()[3]) for l in man if l.startswith('pointer')} +f2 = lambda o: struct.unpack_from('<2f', P, o) +i2 = lambda o: struct.unpack_from('<2i', P, o) +h1 = lambda o: float(np.frombuffer(P, np.float16, 1, o)[0]) +def expo_tex(off): + if off not in res: return 1.0 # LDR: no exposure textures + v = float(np.fromfile(res[off][3], np.float16)[0]) + return 1.0 if v == 0 else v +eA, eB = expo_tex(264), expo_tex(256) +s160 = struct.unpack_from(' 6 else 1 +HERE = os.path.dirname(os.path.abspath(__file__)) +os.makedirs(f'{work}/ref', exist_ok=True) +rng = np.random.default_rng(seed) +seqs = sorted({int(re.match(r'launch-(\d+)-', os.path.basename(f)).group(1)) for f in glob.glob(f'{cap}/dump/launch-*-{K}-arg*')}) +seq = seqs[int(frame)] +def shape_of(off): + f = glob.glob(f'{cap}/dump/launch-{seq:03d}-{K}-arg{off}-*')[0] + w, h, c = (int(x) for x in re.search(r'-(\d+x\d+x\d+)-fmt', f).group(1).split('x')) + return f, w, h, c +def smooth(h, w, cells, amp): + # bilinear upsampling of a coarse random grid: a smooth field with fractional values + g = rng.uniform(-amp, amp, (cells + 1, cells + 1)) + y = np.linspace(0, cells, h)[:, None]; x = np.linspace(0, cells, w)[None, :] + y0 = np.minimum(y.astype(int), cells - 1); x0 = np.minimum(x.astype(int), cells - 1); fy = y - y0; fx = x - x0 + return g[y0, x0] * (1 - fy) * (1 - fx) + g[y0, x0 + 1] * (1 - fy) * fx + g[y0 + 1, x0] * fy * (1 - fx) + g[y0 + 1, x0 + 1] * fy * fx +f, w, h, c = shape_of(240) +mv = np.stack([smooth(h, w, 6, 3.0), smooth(h, w, 6, 3.0)], axis=-1) +mv[rng.uniform(size=(h, w)) < 0.02] += rng.uniform(-8, 8, 2) # outliers, also off-screen reprojections +mv.astype(np.float16).tofile(f'{work}/mv.bin') +f, w, h, c = shape_of(248) +depth = 0.3 + 0.1 * smooth(h, w, 5, 1.0) +blobs = smooth(h, w, 12, 1.0) > 0.35 # nearer objects with hard edges +depth[blobs] += 0.4 +depth.astype(np.float32).tofile(f'{work}/depth.bin') +# reference: the translated kernel on the modified inputs +out = subprocess.run([sys.executable, f'{HERE}/../vulkan/kcase.py', 'make', cap, K, frame, f'{work}/ref', '200:linear', '224:linear', '232:linear'], + capture_output=True, text=True, check=True).stdout +man = open(f'{work}/ref/manifest').read().split('\n') +for i, l in enumerate(man): + if l.startswith('tex 240 '): man[i] = re.sub(r'^tex 240 \S+', f'tex 240 {work}/mv.bin', l) + if l.startswith('tex 248 '): man[i] = re.sub(r'^tex 248 \S+', f'tex 248 {work}/depth.bin', l) + if l.startswith('surf '): man[i] = re.sub(r'^(surf \d+ \d+ \d+ \d+ \d+) \S+', r"\1 -", l) +open(f'{work}/ref/manifest', 'w').write('\n'.join(man)) +subprocess.run([kvk, ref_spv, f'{work}/ref/manifest', '1', '64'], check=True, capture_output=True) +env = dict(os.environ, OUT_TEX_OVERRIDE=f'240={work}/mv.bin,248={work}/depth.bin', OUT_REF_DIR=f'{work}/ref') +print(subprocess.run([sys.executable, f'{HERE}/test_out.py', cap, frame, work, kvk], env=env, capture_output=True, text=True).stdout.strip()) diff --git a/engine/test_post_m.py b/engine/test_post_m.py new file mode 100644 index 0000000..36ae00a --- /dev/null +++ b/engine/test_post_m.py @@ -0,0 +1,54 @@ +#!/usr/bin/env python3 +"""Run the engine's own preset M reconstruction (post_m.comp) on a captured launch and compare its outputs. + +usage: test_post_m.py CAPTURE_DIR FRAME WORK_DIR KVK [repeats] +""" +import glob, os, re, struct, subprocess, sys +import numpy as np +cap, frame, work, kvk = sys.argv[1], int(sys.argv[2]), sys.argv[3], sys.argv[4] +# the kernel name can be overridden; post_3_2 (reconstruction/render < 3) needs the larger shared tile +K = os.environ.get('POST_KERNEL') or next((k for k in ('rrlite_post_3_1_mvlo_hdr_folded', 'rrlite_post_3_2_mvlo_hdr_folded') + if glob.glob(f'{cap}/replay/replay-*-{k}')), 'rrlite_post_3_1_mvlo_hdr_folded') +TILE = ['-DTILE_W=15', '-DTILE_H=15'] if '3_2' in K else [] +reps = sys.argv[5] if len(sys.argv) > 5 else '1' +HERE = os.path.dirname(os.path.abspath(__file__)) +os.makedirs(work, exist_ok=True) +seq = sorted({int(re.match(r'launch-(\d+)-', os.path.basename(f)).group(1)) for f in glob.glob(f'{cap}/dump/launch-*-{K}-arg*')})[frame] +res = {int(re.search(r'-arg(\d+)-', f).group(1)): f for f in glob.glob(f'{cap}/dump/launch-{seq:03d}-{K}-arg*')} +P = open(f'{cap}/replay/replay-{seq:06d}-{K}/args.bin', 'rb').read() +i2 = lambda o: struct.unpack_from('<2i', P, o); f2 = lambda o: struct.unpack_from('<2f', P, o) +expo = lambda o: float(np.fromfile(res[o], np.float16)[0]) +hw, hh = i2(140); rw, rh = i2(148) +open(f'{work}/params.bin', 'wb').write(struct.pack('<2i2i2i2f2f2ffIff', *i2(8), hw, hh, rw, rh, *f2(124), *f2(132), *f2(160), + struct.unpack_from(' {hw}x{hh}, reset {P[168]}, exposure {expo(16):.4g} / {expo(24):.4g}') +a = np.fromfile(res[96], np.float16).astype(np.float64); b = np.fromfile(f'{work}/hist.bin', np.float16).astype(np.float64) +d = np.abs(a - b) +print(f' new history psnr {10 * np.log10(np.abs(a).max() ** 2 / max((d ** 2).mean(), 1e-30)):5.1f} dB identical {100 * (np.fromfile(res[96], np.uint16) == np.fromfile(f"{work}/hist.bin", np.uint16)).mean():5.1f}% max |diff| {d.max():.4g}') +def unpack(f): + w = np.fromfile(f, np.uint32) + return np.stack([w & 2047, (w >> 11) & 2047, w >> 22], -1).astype(np.float64) +a, b = unpack(res[104]), unpack(f'{work}/packed.bin') +d = np.abs(a - b) +print(f' packed output identical {100 * (d == 0).mean():5.1f}% within one step {100 * (d <= 1).mean():5.1f}% psnr {10 * np.log10(2047.0 ** 2 / max((d ** 2).mean(), 1e-30)):5.1f} dB max step {int(d.max())}') +if 112 in res: # the half-resolution output; the launch dumps do not always carry it + codes = np.fromfile(res[112], np.uint32) + a = np.stack([codes & 1023, (codes >> 10) & 1023, (codes >> 20) & 1023], -1).astype(np.float32) / np.float32(1023) + b = np.fromfile(f'{work}/half.bin', np.float32).reshape(-1, 4)[:, :3] + d = np.abs(a.astype(np.float64) - b) + exact = a.view(np.uint32) == b.view(np.uint32) + print(f' half output psnr {10 * np.log10(1.0 / max((d ** 2).mean(), 1e-30)):5.1f} dB identical {100 * exact.mean():5.1f}% max |diff| {d.max():.4g}') diff --git a/engine/vk_dispatch.cpp b/engine/vk_dispatch.cpp new file mode 100644 index 0000000..3fbd116 --- /dev/null +++ b/engine/vk_dispatch.cpp @@ -0,0 +1,55 @@ +#include "vk_dispatch.h" +#ifdef D4R_ENGINE_VK_TABLE +namespace d4r { +VkTable vkTable; +bool loadVulkan(PFN_vkGetInstanceProcAddr instanceProc, VkInstance instance, VkDevice device) { + if (!instanceProc) return false; + auto deviceProc = reinterpret_cast(instanceProc(instance, "vkGetDeviceProcAddr")); + if (!deviceProc) return false; + bool ok = true; + ok &= (vkTable.GetPhysicalDeviceFormatProperties = reinterpret_cast(instanceProc(instance, "vkGetPhysicalDeviceFormatProperties"))) != nullptr; + ok &= (vkTable.GetPhysicalDeviceMemoryProperties = reinterpret_cast(instanceProc(instance, "vkGetPhysicalDeviceMemoryProperties"))) != nullptr; + ok &= (vkTable.GetPhysicalDeviceProperties = reinterpret_cast(instanceProc(instance, "vkGetPhysicalDeviceProperties"))) != nullptr; + ok &= (vkTable.AllocateDescriptorSets = reinterpret_cast(deviceProc(device, "vkAllocateDescriptorSets"))) != nullptr; + ok &= (vkTable.AllocateMemory = reinterpret_cast(deviceProc(device, "vkAllocateMemory"))) != nullptr; + ok &= (vkTable.BindBufferMemory = reinterpret_cast(deviceProc(device, "vkBindBufferMemory"))) != nullptr; + ok &= (vkTable.BindImageMemory = reinterpret_cast(deviceProc(device, "vkBindImageMemory"))) != nullptr; + ok &= (vkTable.CmdBindDescriptorSets = reinterpret_cast(deviceProc(device, "vkCmdBindDescriptorSets"))) != nullptr; + ok &= (vkTable.CmdBindPipeline = reinterpret_cast(deviceProc(device, "vkCmdBindPipeline"))) != nullptr; + ok &= (vkTable.CmdBlitImage = reinterpret_cast(deviceProc(device, "vkCmdBlitImage"))) != nullptr; + ok &= (vkTable.CmdClearColorImage = reinterpret_cast(deviceProc(device, "vkCmdClearColorImage"))) != nullptr; + ok &= (vkTable.CmdCopyBuffer = reinterpret_cast(deviceProc(device, "vkCmdCopyBuffer"))) != nullptr; + ok &= (vkTable.CmdCopyBufferToImage = reinterpret_cast(deviceProc(device, "vkCmdCopyBufferToImage"))) != nullptr; + ok &= (vkTable.CmdCopyImageToBuffer = reinterpret_cast(deviceProc(device, "vkCmdCopyImageToBuffer"))) != nullptr; + ok &= (vkTable.CmdDispatch = reinterpret_cast(deviceProc(device, "vkCmdDispatch"))) != nullptr; + ok &= (vkTable.CmdPipelineBarrier = reinterpret_cast(deviceProc(device, "vkCmdPipelineBarrier"))) != nullptr; + ok &= (vkTable.CmdUpdateBuffer = reinterpret_cast(deviceProc(device, "vkCmdUpdateBuffer"))) != nullptr; + ok &= (vkTable.CmdWriteTimestamp = reinterpret_cast(deviceProc(device, "vkCmdWriteTimestamp"))) != nullptr; + ok &= (vkTable.CreateBuffer = reinterpret_cast(deviceProc(device, "vkCreateBuffer"))) != nullptr; + ok &= (vkTable.CreateComputePipelines = reinterpret_cast(deviceProc(device, "vkCreateComputePipelines"))) != nullptr; + ok &= (vkTable.CreateDescriptorPool = reinterpret_cast(deviceProc(device, "vkCreateDescriptorPool"))) != nullptr; + ok &= (vkTable.CreateDescriptorSetLayout = reinterpret_cast(deviceProc(device, "vkCreateDescriptorSetLayout"))) != nullptr; + ok &= (vkTable.CreateImage = reinterpret_cast(deviceProc(device, "vkCreateImage"))) != nullptr; + ok &= (vkTable.CreateImageView = reinterpret_cast(deviceProc(device, "vkCreateImageView"))) != nullptr; + ok &= (vkTable.CreatePipelineLayout = reinterpret_cast(deviceProc(device, "vkCreatePipelineLayout"))) != nullptr; + ok &= (vkTable.CreateSampler = reinterpret_cast(deviceProc(device, "vkCreateSampler"))) != nullptr; + ok &= (vkTable.CreateShaderModule = reinterpret_cast(deviceProc(device, "vkCreateShaderModule"))) != nullptr; + ok &= (vkTable.DestroyBuffer = reinterpret_cast(deviceProc(device, "vkDestroyBuffer"))) != nullptr; + ok &= (vkTable.DestroyDescriptorPool = reinterpret_cast(deviceProc(device, "vkDestroyDescriptorPool"))) != nullptr; + ok &= (vkTable.DestroyDescriptorSetLayout = reinterpret_cast(deviceProc(device, "vkDestroyDescriptorSetLayout"))) != nullptr; + ok &= (vkTable.DestroyImage = reinterpret_cast(deviceProc(device, "vkDestroyImage"))) != nullptr; + ok &= (vkTable.DestroyImageView = reinterpret_cast(deviceProc(device, "vkDestroyImageView"))) != nullptr; + ok &= (vkTable.DestroyPipeline = reinterpret_cast(deviceProc(device, "vkDestroyPipeline"))) != nullptr; + ok &= (vkTable.DestroyPipelineLayout = reinterpret_cast(deviceProc(device, "vkDestroyPipelineLayout"))) != nullptr; + ok &= (vkTable.DestroySampler = reinterpret_cast(deviceProc(device, "vkDestroySampler"))) != nullptr; + ok &= (vkTable.DestroyShaderModule = reinterpret_cast(deviceProc(device, "vkDestroyShaderModule"))) != nullptr; + ok &= (vkTable.FreeMemory = reinterpret_cast(deviceProc(device, "vkFreeMemory"))) != nullptr; + ok &= (vkTable.GetBufferMemoryRequirements = reinterpret_cast(deviceProc(device, "vkGetBufferMemoryRequirements"))) != nullptr; + ok &= (vkTable.GetImageMemoryRequirements = reinterpret_cast(deviceProc(device, "vkGetImageMemoryRequirements"))) != nullptr; + ok &= (vkTable.MapMemory = reinterpret_cast(deviceProc(device, "vkMapMemory"))) != nullptr; + ok &= (vkTable.UnmapMemory = reinterpret_cast(deviceProc(device, "vkUnmapMemory"))) != nullptr; + ok &= (vkTable.UpdateDescriptorSets = reinterpret_cast(deviceProc(device, "vkUpdateDescriptorSets"))) != nullptr; + return ok; +} +} // namespace d4r +#endif diff --git a/engine/vk_dispatch.h b/engine/vk_dispatch.h new file mode 100644 index 0000000..baef618 --- /dev/null +++ b/engine/vk_dispatch.h @@ -0,0 +1,97 @@ +// Vulkan entry points for the engine. A native build links the loader. A build without a Vulkan import +// library (the MinGW NGX shim) defines D4R_ENGINE_VK_TABLE and fills this table once with loadVulkan(). +#pragma once +#include +#ifdef D4R_ENGINE_VK_TABLE +namespace d4r { +struct VkTable { + PFN_vkGetPhysicalDeviceFormatProperties GetPhysicalDeviceFormatProperties = nullptr; + PFN_vkGetPhysicalDeviceMemoryProperties GetPhysicalDeviceMemoryProperties = nullptr; + PFN_vkGetPhysicalDeviceProperties GetPhysicalDeviceProperties = nullptr; + PFN_vkAllocateDescriptorSets AllocateDescriptorSets = nullptr; + PFN_vkAllocateMemory AllocateMemory = nullptr; + PFN_vkBindBufferMemory BindBufferMemory = nullptr; + PFN_vkBindImageMemory BindImageMemory = nullptr; + PFN_vkCmdBindDescriptorSets CmdBindDescriptorSets = nullptr; + PFN_vkCmdBindPipeline CmdBindPipeline = nullptr; + PFN_vkCmdBlitImage CmdBlitImage = nullptr; + PFN_vkCmdClearColorImage CmdClearColorImage = nullptr; + PFN_vkCmdCopyBuffer CmdCopyBuffer = nullptr; + PFN_vkCmdCopyBufferToImage CmdCopyBufferToImage = nullptr; + PFN_vkCmdCopyImageToBuffer CmdCopyImageToBuffer = nullptr; + PFN_vkCmdDispatch CmdDispatch = nullptr; + PFN_vkCmdPipelineBarrier CmdPipelineBarrier = nullptr; + PFN_vkCmdUpdateBuffer CmdUpdateBuffer = nullptr; + PFN_vkCmdWriteTimestamp CmdWriteTimestamp = nullptr; + PFN_vkCreateBuffer CreateBuffer = nullptr; + PFN_vkCreateComputePipelines CreateComputePipelines = nullptr; + PFN_vkCreateDescriptorPool CreateDescriptorPool = nullptr; + PFN_vkCreateDescriptorSetLayout CreateDescriptorSetLayout = nullptr; + PFN_vkCreateImage CreateImage = nullptr; + PFN_vkCreateImageView CreateImageView = nullptr; + PFN_vkCreatePipelineLayout CreatePipelineLayout = nullptr; + PFN_vkCreateSampler CreateSampler = nullptr; + PFN_vkCreateShaderModule CreateShaderModule = nullptr; + PFN_vkDestroyBuffer DestroyBuffer = nullptr; + PFN_vkDestroyDescriptorPool DestroyDescriptorPool = nullptr; + PFN_vkDestroyDescriptorSetLayout DestroyDescriptorSetLayout = nullptr; + PFN_vkDestroyImage DestroyImage = nullptr; + PFN_vkDestroyImageView DestroyImageView = nullptr; + PFN_vkDestroyPipeline DestroyPipeline = nullptr; + PFN_vkDestroyPipelineLayout DestroyPipelineLayout = nullptr; + PFN_vkDestroySampler DestroySampler = nullptr; + PFN_vkDestroyShaderModule DestroyShaderModule = nullptr; + PFN_vkFreeMemory FreeMemory = nullptr; + PFN_vkGetBufferMemoryRequirements GetBufferMemoryRequirements = nullptr; + PFN_vkGetImageMemoryRequirements GetImageMemoryRequirements = nullptr; + PFN_vkMapMemory MapMemory = nullptr; + PFN_vkUnmapMemory UnmapMemory = nullptr; + PFN_vkUpdateDescriptorSets UpdateDescriptorSets = nullptr; +}; +extern VkTable vkTable; +// False when an entry point is missing. Call before constructing a Model. +bool loadVulkan(PFN_vkGetInstanceProcAddr instanceProc, VkInstance instance, VkDevice device); +} // namespace d4r +#define vkGetPhysicalDeviceFormatProperties ::d4r::vkTable.GetPhysicalDeviceFormatProperties +#define vkGetPhysicalDeviceMemoryProperties ::d4r::vkTable.GetPhysicalDeviceMemoryProperties +#define vkGetPhysicalDeviceProperties ::d4r::vkTable.GetPhysicalDeviceProperties +#define vkAllocateDescriptorSets ::d4r::vkTable.AllocateDescriptorSets +#define vkAllocateMemory ::d4r::vkTable.AllocateMemory +#define vkBindBufferMemory ::d4r::vkTable.BindBufferMemory +#define vkBindImageMemory ::d4r::vkTable.BindImageMemory +#define vkCmdBindDescriptorSets ::d4r::vkTable.CmdBindDescriptorSets +#define vkCmdBindPipeline ::d4r::vkTable.CmdBindPipeline +#define vkCmdBlitImage ::d4r::vkTable.CmdBlitImage +#define vkCmdClearColorImage ::d4r::vkTable.CmdClearColorImage +#define vkCmdCopyBuffer ::d4r::vkTable.CmdCopyBuffer +#define vkCmdCopyBufferToImage ::d4r::vkTable.CmdCopyBufferToImage +#define vkCmdCopyImageToBuffer ::d4r::vkTable.CmdCopyImageToBuffer +#define vkCmdDispatch ::d4r::vkTable.CmdDispatch +#define vkCmdPipelineBarrier ::d4r::vkTable.CmdPipelineBarrier +#define vkCmdUpdateBuffer ::d4r::vkTable.CmdUpdateBuffer +#define vkCmdWriteTimestamp ::d4r::vkTable.CmdWriteTimestamp +#define vkCreateBuffer ::d4r::vkTable.CreateBuffer +#define vkCreateComputePipelines ::d4r::vkTable.CreateComputePipelines +#define vkCreateDescriptorPool ::d4r::vkTable.CreateDescriptorPool +#define vkCreateDescriptorSetLayout ::d4r::vkTable.CreateDescriptorSetLayout +#define vkCreateImage ::d4r::vkTable.CreateImage +#define vkCreateImageView ::d4r::vkTable.CreateImageView +#define vkCreatePipelineLayout ::d4r::vkTable.CreatePipelineLayout +#define vkCreateSampler ::d4r::vkTable.CreateSampler +#define vkCreateShaderModule ::d4r::vkTable.CreateShaderModule +#define vkDestroyBuffer ::d4r::vkTable.DestroyBuffer +#define vkDestroyDescriptorPool ::d4r::vkTable.DestroyDescriptorPool +#define vkDestroyDescriptorSetLayout ::d4r::vkTable.DestroyDescriptorSetLayout +#define vkDestroyImage ::d4r::vkTable.DestroyImage +#define vkDestroyImageView ::d4r::vkTable.DestroyImageView +#define vkDestroyPipeline ::d4r::vkTable.DestroyPipeline +#define vkDestroyPipelineLayout ::d4r::vkTable.DestroyPipelineLayout +#define vkDestroySampler ::d4r::vkTable.DestroySampler +#define vkDestroyShaderModule ::d4r::vkTable.DestroyShaderModule +#define vkFreeMemory ::d4r::vkTable.FreeMemory +#define vkGetBufferMemoryRequirements ::d4r::vkTable.GetBufferMemoryRequirements +#define vkGetImageMemoryRequirements ::d4r::vkTable.GetImageMemoryRequirements +#define vkMapMemory ::d4r::vkTable.MapMemory +#define vkUnmapMemory ::d4r::vkTable.UnmapMemory +#define vkUpdateDescriptorSets ::d4r::vkTable.UpdateDescriptorSets +#endif diff --git a/packaging/D4R_README.txt b/packaging/D4R_README.txt index 68d762a..a81782f 100644 --- a/packaging/D4R_README.txt +++ b/packaging/D4R_README.txt @@ -86,6 +86,15 @@ Settings PreferAccuracy = true. This uses original translated enc1/enc2 layers while retaining other native acceleration. It reduced trails in a captured Townfall sequence on RX 7700 XT and costs some GPU time. Restart after changing it; other scenes and RDNA4 hardware still need testing. +- [Engine] Enabled = true (experimental, off by default) runs DLSS 4 (K) and 4.5 (M) on d4r's own Vulkan + inference engine: the upscale is recorded into the game's own frame instead of going through CUDA + emulation. It needs Mesa RADV with VK_KHR_cooperative_matrix and the models in d4r/engine/k (K), + d4r/engine/k-ldr (K in games without HDR), d4r/engine/m and d4r/engine/m-ldr (M). Games and settings + it does not cover (other models, display-resolution motion vectors with M) keep using the normal path, + and d4r/d4r_nvngx.log says which one ran and why. Tested in one game so far. + This zip has only the engine's shaders in d4r/engine/k. The weights (weights.bin, offsets.bin, + lut.bin) are read from your own nvngx_dlss.dll with the tools in d4r's source repository + (engine/README.md); until they are there, d4r uses the normal path. - OptiScaler.ini: OptiScaler's settings, such as its fps overlay and the render resolution of each quality mode. Press Insert in game for OptiScaler's menu. [Libraries] NvngxPath=d4r\nvngx.dll explicitly selects d4r's shim; OptiDllPath=d4r is also @@ -152,7 +161,8 @@ Licenses d4r is GNU General Public License, version 3 only (GPL-3.0-only). The release also contains ZLUDA (Apache 2.0 or MIT), vkd3d-proton (LGPL 2.1, patched) and OptiScaler 0.9.4 (GPL 3.0, unmodified). d4r/source/SOURCES.txt says where each file comes from, and d4r/licenses has the license texts. d4r is not affiliated with NVIDIA, -AMD or the OptiScaler project. +AMD or the OptiScaler project, and is provided "as is", without warranty of any kind (GPL-3.0, +sections 15 and 16). @clean NVIDIA's files are yours to supply, under NVIDIA's terms. @full NVIDIA's files in this zip (d4r/nvngx_dlss.dll, d4r/ngx/_nvngx.dll) and the kernels built from @full NVIDIA's code (listed in d4r/source/SOURCES.txt) are NVIDIA's property, are not covered by diff --git a/packaging/RELEASE_NOTES.md b/packaging/RELEASE_NOTES.md index 9638211..c724b4b 100644 --- a/packaging/RELEASE_NOTES.md +++ b/packaging/RELEASE_NOTES.md @@ -1,5 +1,12 @@ # Unreleased +- `scripts/check_redistributable.py` refuses NVIDIA's binaries, NGX SDK headers, PTX and code objects, d4r's PTX translations and the engine's model data. `scripts/package_release.sh` runs it on every release before zipping, CI runs it on every push and pull request, and `git config core.hooksPath scripts/git-hooks` runs it before each commit. `LEGAL.md` describes what d4r publishes; `CONTRIBUTING.md` requires signed-off commits. +- Native engine fixes: the K adapter now passes the game's exposure scale to the exposure shader (it was silently ignored); the HIP network's shared buffers are reused across feature recreation instead of leaking their ROCm mapping (21 → 6 imports over seven K features in the harness); `hipnet.bin` must belong to the preset being created; an M-only game no longer loads the K model; `scripts/install_d4r_runtime.sh` stages `engine/k-ldr` and each model's `hip/` folder. +- Native engine, preset K: output is now byte-identical to the CUDA reference over an eleven-frame sequence at Quality (1706×960 → 2560×1440, native HIP network, RX 9070 XT), where later frames previously fell to 47–51 dB. The engine now selects NGX's reconstruction table by upscaling ratio instead of using one table for every ratio, and two rounding-order differences in the input and output stages are gone. K models must be rebuilt (`prepare_k.py`, `compile_k.py`); older packages still load and behave as before. +- Native engine, wider game coverage (K packages must be rebuilt with the current `compile_k.py`, marker `D4RO0003`; older packages keep refusing the new cases with a logged reason): + - Preset K accepts display-resolution motion vectors (MVLowRes clear; output-sized, in output pixels). Byte-identical to the phase-capable CUDA reference over eleven frames at 1706×960 → 2560×1440 (HDR, create flags 0x49) and 1712×960 → 2560×1440 (LDR, 0x68), and over three frames at 1280×720 → 2560×1440. The reference must use a ZLUDA with patch 0008; the old bundled library omits K's bottleneck post phases and is itself off by ~33 dB. Preset M still uses the CUDA backend for display-resolution vectors. + - Preset K accepts a game-supplied exposure texture (AutoExposure off) and any positive exposure scale and pre-exposure in both the automatic and the game path. With the native HIP network at 1706×960 → 2560×1440, HDR K is byte-identical to the CUDA backend over eleven frames for automatic exposure with scale and pre-exposure 1/0.5/0.7/2/3 (and combinations) and for game textures 0.4–32, except an effective exposure of exactly 2 or 4 (first differing frame 104.7 dB / 112.8 dB, then ≈49–63 dB); LDR ignores the exposure and is byte-identical for any value. Preset M accepts the AutoExposure flag, an exposure texture, pre-exposure and exposure scale, which NGX ignores for M (CUDA and engine output identical for each). + - Presets K and M accept the DLSS indicator's inverted-axis parameters, which only place NGX's debug indicator (NGX and engine output byte-identical with and without them), regular (non-inverted) depth (K), and output subrectangles (not yet run in the harness). - Merge the RDNA4 token-lane Swin and hardware-FP8 texture-tail optimizations, preserving MicroCUDA's allocation-identity tube weight cache. - Add opt-in packed R10G10B10A2 and in-place DLSS output, with vkd3d-proton patches `0001`–`0003`. In-place mappings retain their source textures until process exit; output recreation can increase retained VRAM. - Fix deferred input-marker ordering: wait for the game's GPU copies before CUDA uploads or diagnostic readbacks, not just before NGX evaluation. diff --git a/packaging/SOURCES.txt b/packaging/SOURCES.txt index e719ab5..3de2044 100644 --- a/packaging/SOURCES.txt +++ b/packaging/SOURCES.txt @@ -37,17 +37,19 @@ Third-party software, unmodified or patched by d4r: @clean Not included: NVIDIA's DLSS library (nvngx_dlss.dll) and NGX runtime (_nvngx.dll). You supply them @clean and are responsible for complying with NVIDIA's license terms. -@full NVIDIA's software, not covered by any license in d4r/licenses and included by whoever distributes -@full this zip, not by NVIDIA: +@full NVIDIA's software, not covered by any license in d4r/licenses and included by whoever built this +@full zip, not by NVIDIA. NVIDIA's license does not permit redistributing it; do not share this zip: @full d4r/nvngx_dlss.dll NVIDIA DLSS Super Resolution library @DLSS_VERSION@ @full d4r/ngx/_nvngx.dll NVIDIA NGX runtime from NVIDIA's Windows display driver, @NGX_VERSION@ @installer NVIDIA DLLs are not bundled. install.sh downloads DLSS 310.7.0 from NVIDIA/DLSS @installer commit a291cc7d2cc642a51566f3dfd5376f635cd1b284 and _nvngx.dll from NVIDIA driver 596.36. @installer URLs and SHA-256 hashes are pinned in install.sh. NVIDIA license terms apply. @installer Included NVIDIA PTX-derived kernels (not covered by d4r licenses): -@installer d4r/kernels/*/hiluma_engine_output_*, rrlite_enc0_*, rrlite_dec0_*, rrlite_post_*, rrlite_downsample_* -@installer compiled from DLSS PTX by ZLUDA (kernels/tex in this repository). -@full d4r/kernels/*/hiluma_engine_output_*, rrlite_enc0_*, rrlite_dec0_*, rrlite_post_*, rrlite_downsample_* -@full DLSS kernels: NVIDIA's PTX from nvngx_dlss.dll with parts replaced by d4r code -@full (kernels/tex in the d4r repository), compiled by ZLUDA +@installer d4r/kernels/*/hiluma_engine_*, rrlite_enc0_*, rrlite_dec0_*, rrlite_post_*, rrlite_downsample_* +@installer NVIDIA's PTX with parts replaced by d4r code (kernels/tex) or translated to +@installer HIP (kernels/native), compiled by ZLUDA or ROCm's clang. +@full d4r/kernels/*/hiluma_engine_*, rrlite_enc0_*, rrlite_dec0_*, rrlite_post_*, rrlite_downsample_* +@full DLSS kernels: NVIDIA's PTX from nvngx_dlss.dll, with parts replaced by d4r code +@full (kernels/tex in the d4r repository) or translated to HIP (kernels/native), +@full compiled by ZLUDA or ROCm's clang d4r is not affiliated with NVIDIA, AMD or the OptiScaler project. diff --git a/packaging/d4r.ini b/packaging/d4r.ini index dcc24f6..adc750d 100644 --- a/packaging/d4r.ini +++ b/packaging/d4r.ini @@ -60,6 +60,19 @@ ElideNgxSync = true ; effect only while d4r's own output kernel runs, which it does when the DLSS library matches. DirectOutput = auto +[Engine] +; Experimental. Runs DLSS 4 and 4.5 (models K and M) on d4r's own Vulkan inference engine instead of NVIDIA's +; kernels: the game shows the result of the frame it just rendered and [Latency] / [Interop] do not apply. +; Needs Mesa RADV with VK_KHR_cooperative_matrix and the models in d4r\engine\k, k-ldr, m and m-ldr +; (see D4R_README.txt). Games or settings it does not cover use the normal path; d4r_nvngx.log says why. +Enabled = false +; Another model K folder for games with the HDR flag (a path relative to the d4r folder, or a Windows path +; such as Z:\home\me\k) +ModelDir = auto +; auto / hip = the network on d4r's native HIP layer kernels when the model folder has them (faster), +; vulkan = on the engine's shader +Network = auto + [Paths] ; Linux paths. ; ROCm installation (default: the ROCm runtime in d4r\rocm, else /opt/rocm) diff --git a/scripts/build_d4r_nvngx_shim.sh b/scripts/build_d4r_nvngx_shim.sh index 7e8f843..034fbea 100755 --- a/scripts/build_d4r_nvngx_shim.sh +++ b/scripts/build_d4r_nvngx_shim.sh @@ -18,7 +18,9 @@ ln -sfn "$VULKAN_INCLUDE/vk_video" "$ROOT/build/vulkan-include/vk_video" PARAM_OBJECT="$ROOT/build/d4r_ngx_param_msvc.obj" "$CLANG_CL" --target=x86_64-pc-windows-msvc /nologo /std:c++20 /O2 /c /GS- /GR- /EHs-c- /Zl \ "/Fo$PARAM_OBJECT" -- "$ROOT/tools/d4r_ngx_param_msvc.cpp" -"$MINGW_CXX" -std=c++20 -O2 -Wall -Wextra -Wno-missing-field-initializers -I"$ROOT/build/vulkan-include" -I"$ROOT/build/vulkan-staging" -shared -static -static-libgcc -static-libstdc++ \ +# The engine (D4R_ENGINE) reaches Vulkan through a function table: MinGW has no Vulkan import library. +"$MINGW_CXX" -std=c++20 -O2 -Wall -Wextra -Wno-missing-field-initializers -DD4R_ENGINE_VK_TABLE -I"$ROOT/build/vulkan-include" -I"$ROOT/build/vulkan-staging" -shared -static -static-libgcc -static-libstdc++ \ "$ROOT/tools/d4r_nvngx_shim.cpp" "$ROOT/tools/d4r_ngx_param_host.cpp" "$PARAM_OBJECT" \ + "$ROOT/engine/runtime.cpp" "$ROOT/engine/game.cpp" "$ROOT/engine/game_m.cpp" "$ROOT/engine/vk_dispatch.cpp" \ -Wl,--entry,d4r_dll_entry -o "$ROOT/build/d4r_nvngx.dll" printf 'Built %s\n' "$ROOT/build/d4r_nvngx.dll" diff --git a/scripts/d4r_config.py b/scripts/d4r_config.py index a70b554..5bc166b 100755 --- a/scripts/d4r_config.py +++ b/scripts/d4r_config.py @@ -156,6 +156,16 @@ def main(argv): fail(f"[Interop] MarkerPollUs must be a number of microseconds, not {value!r}") env["D4R_SHIM_MARKER_POLL_US"] = value + value = flag("Engine", "Enabled") + env["D4R_ENGINE"] = "1" if value else "0" + if path_value("Engine", "ModelDir"): + env["D4R_ENGINE_MODEL_DIR"] = path_value("Engine", "ModelDir") + network = get("Engine", "Network") + if network is not None and network not in ("auto", "hip", "vulkan"): + fail("[Engine] Network must be auto, hip or vulkan") + if network in ("hip", "vulkan"): + env["D4R_ENGINE_NETWORK"] = network + engine = path_value("Game", "EngineIni") cvars = get("Game", "Cvars") if engine and cvars: diff --git a/scripts/install_d4r_runtime.sh b/scripts/install_d4r_runtime.sh index 0138b9f..297c421 100755 --- a/scripts/install_d4r_runtime.sh +++ b/scripts/install_d4r_runtime.sh @@ -7,6 +7,11 @@ set -euo pipefail # bin/d4r_nvngx.dll D3D12-to-CUDA NGX core shim (OptiScaler NvngxPath target) # ngx/_nvngx.dll official NGX core # dlss/nvngx_dlss.dll official DLSS SR feature DLL (also copied into bin/) +# engine/k/, k-ldr/ models of preset K for the native engine ([Engine] in d4r.ini), when D4R_ENGINE_MODEL / +# D4R_ENGINE_MODEL_LDR name the folders engine/compile_k.py wrote (--ldr for the second) +# engine/m/, m-ldr/ models of preset M, when D4R_ENGINE_MODEL_M / D4R_ENGINE_MODEL_M_LDR name the folders +# engine/compile_m.py wrote +# Each model folder is copied whole, with the native HIP network's hipnet.bin and hip/ when built with --hip. # usage: install_d4r_runtime.sh PATH_TO_NGX_CORE_DLL PATH_TO_NVNGX_DLSS_DLL if [[ $# -ne 2 ]]; then printf 'usage: %s PATH_TO_NGX_CORE_DLL PATH_TO_NVNGX_DLSS_DLL\n' "$0" >&2 @@ -25,5 +30,22 @@ copy_unless_same "$2" "$RUNTIME/dlss/nvngx_dlss.dll" # The NGX core's CUDA init ignores the feature search paths it is given and # scans the directory of the module that calls it, i.e. the shim's. cp -f "$RUNTIME/dlss/nvngx_dlss.dll" "$RUNTIME/bin/nvngx_dlss.dll" +# K and M models; the -ldr variants serve games without the HDR flag +for variant in k k-ldr m m-ldr; do + case "$variant" in + k) source_dir="${D4R_ENGINE_MODEL:-}" required="weights.bin offsets.bin lut.bin input_k.spv output_k.spv exposure_k0.spv exposure_k1.spv" compiler=compile_k.py ;; + k-ldr) source_dir="${D4R_ENGINE_MODEL_LDR:-}" required="weights.bin offsets.bin lut.bin input_k.spv output_k.spv" compiler="compile_k.py --ldr" ;; + m) source_dir="${D4R_ENGINE_MODEL_M:-}" required="offsets.bin enc0_m.spv" compiler=compile_m.py ;; + m-ldr) source_dir="${D4R_ENGINE_MODEL_M_LDR:-}" required="offsets.bin enc0_m.spv" compiler="compile_m.py --ldr" ;; + esac + [[ -n "$source_dir" ]] || continue + for f in $required; do + [[ -f "$source_dir/$f" ]] || { printf '%s lacks %s (run engine/%s)\n' "$source_dir" "$f" "$compiler" >&2; exit 2; } + done + rm -rf "$RUNTIME/engine/$variant" + mkdir -p "$RUNTIME/engine" + cp -r "$source_dir" "$RUNTIME/engine/$variant" + printf 'Installed the native engine model in %s\n' "$RUNTIME/engine/$variant" +done sha256sum "$RUNTIME"/ngx/_nvngx.dll "$RUNTIME"/dlss/nvngx_dlss.dll printf 'Installed d4r runtime in %s\n' "$RUNTIME" diff --git a/scripts/package_release.sh b/scripts/package_release.sh index 339f009..117238c 100755 --- a/scripts/package_release.sh +++ b/scripts/package_release.sh @@ -32,6 +32,10 @@ # Accuracy texture sets go in accuracy//, built with D4R_PREFER_ACCURACY=1. # D4R_ZLUDA_EMIT d4r_emit from the patched ZLUDA build; required in full builds when an accuracy # texture set or L's unfolded texture variants are not supplied in D4R_BUNDLE_TEX. +# D4R_BUNDLE_ENGINE optional: the native engine's preset K model (the folder engine/compile_k.py wrote), +# packaged as d4r/engine/k for d4r.ini [Engine]. Only d4r's shaders are taken from it: the +# weights and reconstruction table are values read from NVIDIA's DLSS and are never packaged, +# in any variant. Users produce them from their own nvngx_dlss.dll (engine/README.md). # D4R_BUNDLE_NVIDIA=installer (default) ships install.sh instead of NVIDIA DLLs, retaining texture kernels. # D4R_BUNDLE_NVIDIA=1 bundles both DLLs; =0 omits DLLs and texture kernels (nonvidia ZIP). set -euo pipefail @@ -119,6 +123,16 @@ if [[ "$VARIANT" != clean ]]; then else printf 'Put NVIDIA'"'"'s NGX runtime, _nvngx.dll, in this folder (see D4R_README.txt).\r\n' > "$STAGE/d4r/ngx/README.txt" fi +if [[ -n "${D4R_BUNDLE_ENGINE:-}" ]]; then + for f in input_k.spv output_k.spv exposure_k0.spv exposure_k1.spv; do + [[ -f "$D4R_BUNDLE_ENGINE/$f" ]] || { echo "D4R_BUNDLE_ENGINE lacks $f (run engine/compile_k.py)" >&2; exit 2; } + done + mkdir -p "$STAGE/d4r/engine/k" + cp "$D4R_BUNDLE_ENGINE"/*.spv "$STAGE/d4r/engine/k/" + # d4r's own shader-revision marker; every other .bin in the model holds values read from NVIDIA's DLSS + [[ ! -f "$D4R_BUNDLE_ENGINE/direct_origins.bin" ]] || cp "$D4R_BUNDLE_ENGINE/direct_origins.bin" "$STAGE/d4r/engine/k/" + printf 'The native engine needs weights.bin, offsets.bin and lut.bin here (see D4R_README.txt).\r\n' > "$STAGE/d4r/engine/k/README.txt" +fi IFS=: read -r -a DLLS <<< "$D4R_DLSS_DLLS" # one folder per target; RDNA4 targets also get -fp8 (native FP8 WMMA, d4r.ini NativeFp8), which the # bridge serves instead when that setting is on @@ -253,3 +267,6 @@ rm -f "$OUT/$NAME.zip" (cd "$STAGE" && find . -type f | LC_ALL=C sort | sed 's|^\./||' | zip -q -X -9 "$OUT/$NAME.zip" -@) (cd "$OUT" && sha256sum "$NAME.zip" > "$NAME.zip.sha256") printf 'Built %s (%s)\n' "$OUT/$NAME.zip" "$(du -h "$OUT/$NAME.zip" | cut -f1)" +if [[ "$VARIANT" == full ]]; then + echo "This zip contains NVIDIA's DLLs and kernels built from NVIDIA's code. Do not publish or share it." >&2 +fi diff --git a/scripts/run_game_d4r_optiscaler_proton.sh b/scripts/run_game_d4r_optiscaler_proton.sh index af0e166..e7a46c5 100755 --- a/scripts/run_game_d4r_optiscaler_proton.sh +++ b/scripts/run_game_d4r_optiscaler_proton.sh @@ -41,6 +41,8 @@ if [[ "${D4R_NO_CONFIG:-0}" != 1 ]]; then eval "$D4R_CONFIG_EXPORTS" fi source "$ROOT/scripts/d4r_proton_env.sh" +# [Engine] ModelDir = auto: the model scripts/install_d4r_runtime.sh staged +export D4R_ENGINE_MODEL_DIR="${D4R_ENGINE_MODEL_DIR:-$D4R_RUNTIME_DIR/engine/k}" PREFIX="$COMPAT_DIR/pfx" UMU="$PREFIX/drive_c/windows/system32/umu" diff --git a/tests/test_frame_completion.py b/tests/test_frame_completion.py index a0d1e69..a87e3d6 100644 --- a/tests/test_frame_completion.py +++ b/tests/test_frame_completion.py @@ -12,6 +12,9 @@ class FrameCompletionTests(unittest.TestCase): shim = (ROOT / 'tools/d4r_nvngx_shim.cpp').read_text() start = shim.index('static void release_split_frame(Feature* feature, uint32_t frame)\n{') end = shim.index('\nstatic void split_watchdog()', start) + # release_split_frame signals through this helper, defined after the engine network's exports + helper = shim.index('\nstatic VkResult signal_split_semaphore(Feature& feature, uint64_t value)\n{') + helper = shim[helper:shim.index('\n}\n', helper) + 3] source = r''' #include "d4r_frame_completion.h" #include @@ -25,7 +28,10 @@ struct Feature { D4rFrameCompletion splitCompletion; uint64_t splitSignalled = 0; int splitSemaphore = 1; + bool engineNetTimeline = false; // the Vulkan timeline of this test, not the native network's + int engineNet = 0; }; +static struct { int (*signalTimeline)(int, uint64_t) = nullptr; } g_engineNet; static uint64_t gpuTimeline = 0; static bool failSignal = false; static int signal(int, const VkSemaphoreSignalInfo* info) { @@ -36,7 +42,7 @@ static int signal(int, const VkSemaphoreSignalInfo* info) { } static struct { int device = 1; decltype(&signal) signalSemaphore = signal; } g_vk; static void logf(const char*, ...) {} -''' + shim[start:end] + r''' +''' + helper + shim[start:end] + r''' int main() { Feature feature; feature.splitCompletion.admit(1); // still running on the CUDA worker diff --git a/tests/test_optiscaler_packaging.py b/tests/test_optiscaler_packaging.py index 9edf722..e50162f 100644 --- a/tests/test_optiscaler_packaging.py +++ b/tests/test_optiscaler_packaging.py @@ -147,6 +147,11 @@ class OptiScalerConfigTests(unittest.TestCase): D4R_BUNDLE_DLSS=str(dlss), D4R_BUNDLE_NGX=str(core), D4R_BUNDLE_TEX=str(inputs / 'tex')) version = (repo / 'packaging/VERSION').read_text().strip() configs = [] + # A compiled engine model: only d4r's shaders and marker may reach a zip. + for name in ('input_k.spv', 'output_k.spv', 'exposure_k0.spv', 'exposure_k1.spv', 'enc0.spv', + 'direct_origins.bin', 'weights.bin', 'offsets.bin', 'lut.bin', 'hipnet.bin', 'model.json'): + fixture(inputs / 'engine' / name) + env['D4R_BUNDLE_ENGINE'] = str(inputs / 'engine') for bundled in ('1', '0'): with self.subTest(bundled=bundled): env['D4R_BUNDLE_NVIDIA'] = bundled @@ -165,6 +170,10 @@ class OptiScalerConfigTests(unittest.TestCase): self.assertNotIn('d4r/_nvngx.dll', archive.namelist()) self.assertEqual('d4r/ngx/_nvngx.dll' in archive.namelist(), bundled == '1') self.assertEqual('d4r/nvngx_dlss.dll' in archive.namelist(), bundled == '1') + engine = sorted(n for n in archive.namelist() if n.startswith('d4r/engine/')) + self.assertEqual(engine, ['d4r/engine/k/' + n for n in ( + 'README.txt', 'direct_origins.bin', 'enc0.spv', 'exposure_k0.spv', 'exposure_k1.spv', + 'input_k.spv', 'output_k.spv')]) self.assertEqual(archive.read('d4r/d4r-check.sh'), (ROOT / 'packaging/d4r-check.sh').read_bytes()) self.assertEqual(configs[0], configs[1]) diff --git a/tests/test_shim_config.py b/tests/test_shim_config.py index 59ac250..4204a19 100644 --- a/tests/test_shim_config.py +++ b/tests/test_shim_config.py @@ -25,6 +25,29 @@ class ShimConfigTests(unittest.TestCase): self.assertEqual(int(portable[1]), 200) self.assertEqual(int(runtime[1]), 200) + def test_engine_is_opt_in_and_configured_by_ini(self): + for name in ('packaging/d4r.ini', 'config/d4r.ini.default'): + ini = configparser.ConfigParser(inline_comment_prefixes=(';',)) + ini.read(ROOT / name) + self.assertFalse(ini.getboolean('Engine', 'Enabled'), name) + self.assertEqual(ini.get('Engine', 'ModelDir'), 'auto', name) + source = SHIM.read_text() + self.assertIn('portable_set("D4R_ENGINE", ini_flag(ini, "engine", "Enabled", 0) ? "1" : "0");', source) + self.assertIn('env_uint("D4R_ENGINE", 0)', source) + with tempfile.TemporaryDirectory() as temporary: + ini = Path(temporary) / 'd4r.ini' + env = {key: value for key, value in os.environ.items() if not key.startswith('D4R_')} + cmd = ['python3', str(ROOT / 'scripts/d4r_config.py'), '--config', str(ini)] + for text, enabled, model in (('', '0', None), ('[Engine]\nEnabled = true\n', '1', None), + ('[Engine]\nEnabled = true\nModelDir = /models/k\n', '1', '/models/k')): + ini.write_text(text) + result = subprocess.run(cmd, env=env, text=True, capture_output=True, check=True) + self.assertIn(f'export D4R_ENGINE={enabled}\n', result.stdout) + self.assertEqual('export D4R_ENGINE_MODEL_DIR=/models/k\n' in result.stdout, model is not None) + ini.write_text('[Engine]\nEnabled = maybe\n') + result = subprocess.run(cmd, env=env, text=True, capture_output=True) + self.assertEqual(result.returncode, 2) + def test_launcher_marker_poll_defaults_overrides_and_validation(self): with tempfile.TemporaryDirectory() as temporary: ini = Path(temporary) / 'd4r.ini' @@ -49,7 +72,8 @@ class ShimConfigTests(unittest.TestCase): # and a captured logger, so the test exercises the production message # and override behavior without requiring Windows, NGX or a GPU. source = SHIM.read_text() - start = source.index(' const char* presetNames[] =') + creation = source.index('static NgxResult create_feature(') + start = source.index(' const char* presetNames[] =', creation) end = source.index(' feature->preset = presets[1];', start) block = source[start:end] runner_source = r''' diff --git a/tools/d3d12_dlss_harness.cpp b/tools/d3d12_dlss_harness.cpp index 2a925d8..cafb0cc 100644 --- a/tools/d3d12_dlss_harness.cpp +++ b/tools/d3d12_dlss_harness.cpp @@ -594,6 +594,17 @@ int main(int argc, char** argv) std::vector motion(static_cast(motionWidth) * motionHeight * 2, 0); const float exposureValue = 1.0f; const std::array exposureRGBA = {exposureValue, 0.25f, 0.5f, 0.75f}; + // D4R_HARNESS_EXPOSURE_VALUES=v1,v2,...: exposure texture values, one per frame. + std::vector exposureSchedule; + if (const char* values = std::getenv("D4R_HARNESS_EXPOSURE_VALUES")) + for (const char* cursor = values;;) + { + char* end = nullptr; + exposureSchedule.push_back(std::strtof(cursor, &end)); + if (*end != ',') + break; + cursor = end + 1; + } std::vector outputInit(static_cast(outWidth) * outHeight * 4, 0); std::vector color8; std::vector outputInit8; @@ -734,6 +745,13 @@ int main(int argc, char** argv) // Default to the proven single-frame reference. Temporal probes reset only // the first frame and let the same official NGX feature accumulate history. d4r_ngx_set_int(parameters, "Reset", 1); + // D4R_HARNESS_INDICATOR_INVERT=1 sets NGX's DLSS.Indicator.Invert.X/Y.Axis, which games set to place the + // DLSS debug indicator; compare against an unset run to see whether they reach the image. + if (const char* invert = std::getenv("D4R_HARNESS_INDICATOR_INVERT"); invert != nullptr && std::atoi(invert) != 0) + { + d4r_ngx_set_int(parameters, "DLSS.Indicator.Invert.X.Axis", 1); + d4r_ngx_set_int(parameters, "DLSS.Indicator.Invert.Y.Axis", 1); + } // Quality scene ground truth (8x8 supersampled output pixels) and metrics. // D4R_HARNESS_QUALITY_PAN="vx,vy" scrolls the scene content by whole @@ -835,8 +853,13 @@ int main(int argc, char** argv) const char* multiplier = std::getenv("D4R_HARNESS_MV_MULT"); const float mvScale = multiplier != nullptr ? std::strtof(multiplier, nullptr) : 1.0f; const float unitX = motionHighRes ? 1.0f : 1.0f / scaleX, unitY = motionHighRes ? 1.0f : 1.0f / scaleY; - const uint16_t mvX = std::bit_cast(static_cast<_Float16>(-mvScale * panX * unitX)); - const uint16_t mvY = std::bit_cast(static_cast<_Float16>(-mvScale * panY * unitY)); + // D4R_HARNESS_MV_SCALE="sx,sy": the same motion expressed in other units, as games do + // (e.g. "-640,-360": negated, normalised vectors); MV.Scale carries the factor back. + float gameScaleX = 1.0f, gameScaleY = 1.0f; + if (const char* gameScale = std::getenv("D4R_HARNESS_MV_SCALE")) + std::sscanf(gameScale, "%f,%f", &gameScaleX, &gameScaleY); + const uint16_t mvX = std::bit_cast(static_cast<_Float16>(-mvScale * panX * unitX / gameScaleX)); + const uint16_t mvY = std::bit_cast(static_cast<_Float16>(-mvScale * panY * unitY / gameScaleY)); for (size_t pixel = 0; pixel < motion.size() / 2; ++pixel) { motion[pixel * 2 + 0] = mvX; @@ -846,6 +869,13 @@ int main(int argc, char** argv) } d4r_ngx_set_float(parameters, "Jitter.Offset.X", signX * jx); d4r_ngx_set_float(parameters, "Jitter.Offset.Y", signY * jy); + if (const char* gameScale = std::getenv("D4R_HARNESS_MV_SCALE")) + { + float gameScaleX = 1.0f, gameScaleY = 1.0f; + std::sscanf(gameScale, "%f,%f", &gameScaleX, &gameScaleY); + d4r_ngx_set_float(parameters, "MV.Scale.X", gameScaleX); + d4r_ngx_set_float(parameters, "MV.Scale.Y", gameScaleY); + } } if (jitterScene) { @@ -914,12 +944,51 @@ int main(int argc, char** argv) } else if (temporal) d4r_ngx_set_int(parameters, "Reset", forceReset || frame == 1 ? 1 : 0); + // D4R_HARNESS_PRE_EXPOSURE / D4R_HARNESS_EXPOSURE_SCALE: non-neutral values on every frame, as games + // with their own exposure pipeline pass them. + if (const char* value = std::getenv("D4R_HARNESS_PRE_EXPOSURE")) + d4r_ngx_set_float(parameters, "DLSS.Pre.Exposure", static_cast(std::atof(value))); + if (const char* value = std::getenv("D4R_HARNESS_EXPOSURE_SCALE")) + d4r_ngx_set_float(parameters, "DLSS.Exposure.Scale", static_cast(std::atof(value))); + // Change exposure after initial frames to exercise temporal/backend transitions. + // CUDA supports this scale; the native K adapter currently accepts only a neutral scale. + if (const char* changeFrame = std::getenv("D4R_HARNESS_EXPOSURE_CHANGE_FRAME")) + if (frame == std::atoi(changeFrame)) + d4r_ngx_set_float(parameters, "DLSS.Exposure.Scale", 0.5f); + // D4R_HARNESS_EXPOSURE_VALUES: the exposure texture's value on each frame (the last value repeats). + if (!exposureSchedule.empty()) + { + const float value = exposureSchedule[std::min(static_cast(frame - 1), exposureSchedule.size() - 1)]; + const std::array rgba = {value, 0.25f, 0.5f, 0.75f}; + update_texture(exposureTexture, exposureRGBA32 ? static_cast(rgba.data()) : static_cast(&value), + exposureRGBA32 ? 16 : 4, srv); + } const DWORD start = GetTickCount(); + LARGE_INTEGER counterStart, counterRecorded, counterEnd, counterFrequency; + QueryPerformanceFrequency(&counterFrequency); + QueryPerformanceCounter(&counterStart); result = evaluate(g_list, feature, parameters, nullptr); + QueryPerformanceCounter(&counterRecorded); if (const char* delay = std::getenv("D4R_HARNESS_DELAY_SUBMIT_MS")) Sleep(static_cast(std::clamp(std::atoi(delay), 0, 4000))); submit_and_wait(); - std::printf("frame %d: EvaluateFeature -> 0x%08x (%lu ms incl. submit)\n", frame, result, GetTickCount() - start); + QueryPerformanceCounter(&counterEnd); + // The list holds only the evaluation, so the second figure is its whole cost when it is recorded into + // the list (the engine backend): recording, submission, GPU execution and the fence. + std::printf("frame %d: EvaluateFeature -> 0x%08x (%lu ms incl. submit; call %.0f us, call + submit + wait %.0f us)\n", + frame, result, GetTickCount() - start, + 1e6 * static_cast(counterRecorded.QuadPart - counterStart.QuadPart) / counterFrequency.QuadPart, + 1e6 * static_cast(counterEnd.QuadPart - counterStart.QuadPart) / counterFrequency.QuadPart); + // D4R_HARNESS_TIMING=: one "frame call_us total_us start_us" line per frame. + if (const char* timingPath = std::getenv("D4R_HARNESS_TIMING")) + if (FILE* timingFile = std::fopen(timingPath, "a")) + { + std::fprintf(timingFile, "%d %.1f %.1f %.1f\n", frame, + 1e6 * static_cast(counterRecorded.QuadPart - counterStart.QuadPart) / counterFrequency.QuadPart, + 1e6 * static_cast(counterEnd.QuadPart - counterStart.QuadPart) / counterFrequency.QuadPart, + 1e6 * static_cast(counterStart.QuadPart) / counterFrequency.QuadPart); + std::fclose(timingFile); + } Sleep(frameWaitMs); // let the CUDA worker finish before the next presentation if (qualityScene && !poleScene && frame > frames - 17) { diff --git a/tools/d4r_nvngx_shim.cpp b/tools/d4r_nvngx_shim.cpp index 06ec568..4f4a8da 100644 --- a/tools/d4r_nvngx_shim.cpp +++ b/tools/d4r_nvngx_shim.cpp @@ -66,6 +66,9 @@ #include "d4r_frame_completion.h" #include "d4r_win32_wait.h" #include "d4r_vkd3d_interop.h" +#include "../engine/game.h" +#include "../engine/game_m.h" +#include "../engine/vk_dispatch.h" // --- NGX types (declared locally: the core's exports differ from the SDK) -- @@ -573,6 +576,17 @@ static void load_portable_config() // release cannot ship that kernel, which is built from NVIDIA's PTX. portable_set("D4R_SHIM_OUTPUT_DIRECT", ini_flag(ini, "interop", "DirectOutput", nativeOn) ? "1" : "0"); + // [Engine]: d4r's own Vulkan inference for DLSS 4 (K); the model is looked up in d4r\engine\k + portable_set("D4R_ENGINE", ini_flag(ini, "engine", "Enabled", 0) ? "1" : "0"); + if (const std::string engineModel = ini_value(ini, "engine", "ModelDir"); !engineModel.empty()) + portable_set(L"D4R_ENGINE_MODEL_DIR", portable_path(engineModel)); + // Network: preset M's network on the native HIP layers ("hip", the default when the model has them: faster) or + // on the engine's Vulkan shader ("vulkan") + if (const std::string network = ini_value(ini, "engine", "Network"); network == "hip" || network == "vulkan") + portable_set("D4R_ENGINE_NETWORK", network); + else if (!network.empty() && network != "auto") + g_portable.notes.push_back("[Engine] Network must be auto, hip or vulkan; using auto"); + portable_set("D4R_SHIM_WATERMARK", ini_flag(ini, "dlss", "ShowWatermark", 0) ? "1" : "0"); const std::string dilation = ini_value(ini, "interop", "MotionVectorDilation"); if (dilation.empty() || dilation == "0" || dilation == "1" || dilation == "2") @@ -1563,6 +1577,10 @@ struct VramBuffer CudaDevicePtr device = 0; void* external = nullptr; size_t bytes = 0; + void* mapped = nullptr; // engine readiness word: CPU-coherent, kept until HIP and Vulkan retire + // What a pooled buffer can be reused as: 0 the interop transfer buffers (create_vram_buffer), 1 and 2 the + // engine's device-local and host-coherent storage buffers (create_engine_shared_buffer). + uint8_t kind = 0; }; // A temporary image for colour or exposure format conversion. @@ -1696,6 +1714,22 @@ struct OutputSlot std::atomic producedFrame{0}; }; +// Views and their D3D12 image owners survive every command list recorded by this feature. +struct EngineBorrowedView +{ + ID3D12Resource* resource = nullptr; + VkDevice device = VK_NULL_HANDLE; + VkImageView view = VK_NULL_HANDLE; + VkFormat format = VK_FORMAT_UNDEFINED; + VkImageAspectFlags aspect = 0; + VkImageUsageFlags usage = 0; + ~EngineBorrowedView() + { + if (view) d4r::vkTable.DestroyImageView(device, view, nullptr); + if (resource) resource->Release(); + } +}; + struct Feature { NgxHandle handle = {}; // returned to the caller @@ -1708,7 +1742,9 @@ struct Feature unsigned int width = 0, height = 0, outWidth = 0, outHeight = 0; unsigned int cudaWidth = 0, cudaHeight = 0; // render size the NGX feature was created with (worker only) unsigned int refusedWidth = 0, refusedHeight = 0; // render size NGX would not create a feature for (worker only) - int quality = 0, flags = 0; + int quality = 0, flags = 0, subrects = 0; + std::array creationPresets{}; + bool cudaCreationDeferred = false; unsigned int preset = 0; CudaEvent profileStart = nullptr, profileEnd = nullptr; CudaEvent outputReadyEvent = nullptr; @@ -1774,6 +1810,29 @@ struct Feature std::atomic latestOutput{-1}; std::atomic completedFrames{0}; uint32_t evaluatedFrames = 0; // worker only + // D4R_ENGINE: d4r's own inference, recorded into the game's command lists (game thread only). Recorded + // lists retain `resources`, so this is destroyed only once the GPU has finished with them. + uint64_t engineViewBytes = 0; + bool engineViewBudgetReported = false; + std::vector> engineViews; // destroyed after engine descriptor pools + std::unique_ptr engine; + std::unique_ptr engineM; // preset M + unsigned int engineMRenderWidth = 0, engineMRenderHeight = 0; // first evaluation's effective extent + // Preset M and preset K with the network on the native HIP layers ([Engine] Network): the buffers shared + // with HIP and the bridge's network. The game's list is split around the network, as for the CUDA backend's + // split frames. `engineHead` is the network's result (M's expansion / K's final store reads it). + VramBuffer engineTokens, engineHead; + // Input readiness for the HIP network: engineSync holds the frame number the engine's front writes (a Vulkan + // fill) and the network's stream waits for it on the GPU (engineGpuWait), so the layers start without a + // CPU round trip. engineWaitFrame: the frame a queued stream wait is waiting for (0: none). + VramBuffer engineSync; + std::atomic engineGpuWait{false}; + bool engineNetTimeline = false; // native HIP callback and dropped-frame signals share one timeline lock + std::atomic engineWaitFrame{0}; + double engineNetWaitEmaUs = 0.0; // worker only: smoothed launch-to-completion wait, for adaptive polling + void* engineNet = nullptr; + bool engineRefused = false; + float engineJitter[2] = {0, 0}; // of the previous engine frame // Serializes picking the slot to present (game thread) against picking // the slot to overwrite (worker). std::mutex outputMutex; @@ -2183,6 +2242,8 @@ struct VulkanInterop PFN_vkBindImageMemory bindImage = nullptr; PFN_vkAllocateMemory allocate = nullptr; PFN_vkFreeMemory free = nullptr; + PFN_vkMapMemory map = nullptr; + PFN_vkUnmapMemory unmap = nullptr; PFN_vkCmdPipelineBarrier barrier = nullptr; PFN_vkCmdCopyImageToBuffer copyImageToBuffer = nullptr; PFN_vkCmdCopyBufferToImage copyBufferToImage = nullptr; @@ -2262,6 +2323,8 @@ static void init_vram_interop() load(g_vk.bindImage, "vkBindImageMemory"); load(g_vk.allocate, "vkAllocateMemory"); load(g_vk.free, "vkFreeMemory"); + load(g_vk.map, "vkMapMemory"); + load(g_vk.unmap, "vkUnmapMemory"); load(g_vk.barrier, "vkCmdPipelineBarrier"); load(g_vk.copyImageToBuffer, "vkCmdCopyImageToBuffer"); load(g_vk.copyBufferToImage, "vkCmdCopyBufferToImage"); @@ -2311,13 +2374,13 @@ static std::mutex g_vramPoolMutex; static std::vector g_vramPool; constexpr size_t kVramPoolLimit = 64; -// The smallest pooled buffer that holds `bytes`. -static bool take_pooled_vram_buffer(VramBuffer& target, size_t bytes) +// The smallest pooled buffer of `kind` that holds `bytes`. +static bool take_pooled_vram_buffer(VramBuffer& target, size_t bytes, uint8_t kind = 0) { std::lock_guard lock(g_vramPoolMutex); auto best = g_vramPool.end(); for (auto it = g_vramPool.begin(); it != g_vramPool.end(); ++it) - if (it->bytes >= bytes && (best == g_vramPool.end() || it->bytes < best->bytes)) + if (it->kind == kind && it->bytes >= bytes && (best == g_vramPool.end() || it->bytes < best->bytes)) best = it; if (best == g_vramPool.end()) return false; @@ -2387,6 +2450,8 @@ static void destroy_vram_buffer(VramBuffer& buffer) logf("VRAM interop: freeing %zu-byte buffer: cuMemFree %d, release %d", buffer.bytes, freed, released); return released; }); + if (buffer.mapped != nullptr) + g_vk.unmap(g_vk.device, buffer.memory); if (buffer.memory != VK_NULL_HANDLE) g_vk.free(g_vk.device, buffer.memory, nullptr); if (buffer.buffer != VK_NULL_HANDLE) @@ -2804,16 +2869,15 @@ static bool record_vram_output(ID3D12GraphicsCommandList* list, const VramBuffer // dropped frames advance the signal only past earlier retired frames. The // watchdog reports delays without releasing an unfinished producer. +static VkResult signal_split_semaphore(Feature& feature, uint64_t value); + static void release_split_frame(Feature* feature, uint32_t frame) { std::lock_guard lock(feature->splitMutex); const uint64_t value = feature->splitCompletion.retire(frame); if (value <= feature->splitSignalled) return; - VkSemaphoreSignalInfo info = {VK_STRUCTURE_TYPE_SEMAPHORE_SIGNAL_INFO}; - info.semaphore = feature->splitSemaphore; - info.value = value; - const VkResult result = g_vk.signalSemaphore(g_vk.device, &info); + const VkResult result = signal_split_semaphore(*feature, value); if (result == VK_SUCCESS) feature->splitSignalled = value; if (result != VK_SUCCESS) @@ -4782,6 +4846,89 @@ static float get_float_or(void* parameters, const char* name, float fallback) return d4r_ngx_get_float(parameters, name, &value) == NGX_SUCCESS ? value : fallback; } +static bool engine_requested(); + +// Called on the game thread with the feature API lock held. All CUDA resources are made on its worker. +// Vulkan-only K features retain just their creation settings until a frame actually needs CUDA. +static NgxResult create_cuda_feature(Feature& feature) +{ + const char* presetNames[] = {"DLSS.Hint.Render.Preset.DLAA", "DLSS.Hint.Render.Preset.Quality", + "DLSS.Hint.Render.Preset.Balanced", "DLSS.Hint.Render.Preset.Performance", + "DLSS.Hint.Render.Preset.UltraPerformance", "DLSS.Hint.Render.Preset.UltraQuality"}; + const NgxResult result = g.worker.call([&]() -> NgxResult { + NgxResult allocation = g.ngx.allocateParameters(&feature.cudaParams); + if (allocation != NGX_SUCCESS) + return allocation; + void* p = feature.cudaParams; + d4r_ngx_set_uint(p, "Width", feature.width); + d4r_ngx_set_uint(p, "Height", feature.height); + d4r_ngx_set_uint(p, "OutWidth", feature.outWidth); + d4r_ngx_set_uint(p, "OutHeight", feature.outHeight); + d4r_ngx_set_int(p, "PerfQualityValue", feature.quality); + d4r_ngx_set_int(p, "DLSS.Feature.Create.Flags", feature.flags); + d4r_ngx_set_int(p, "DLSS.Enable.Output.Subrects", feature.subrects); + d4r_ngx_set_uint(p, "CreationNodeMask", 1); + d4r_ngx_set_uint(p, "VisibilityNodeMask", 1); + for (int index = 0; index < 6; ++index) + d4r_ngx_set_uint(p, presetNames[index], feature.creationPresets[index]); + if (g.cu.memAlloc(&feature.scratch, 64ull * 1024 * 1024) == 0) + { + d4r_ngx_set_void(p, "Scratch", reinterpret_cast(static_cast(feature.scratch))); + d4r_ngx_set_ull(p, "Scratch.SizeInBytes", 64ull * 1024 * 1024); + } + else + feature.scratch = 0; + const NgxResult created = g.ngx.createFeature(NGX_FEATURE_SUPER_SAMPLING, p, &feature.cudaHandle); + if (created != NGX_SUCCESS) + { + if (feature.cudaHandle != nullptr) g.ngx.releaseFeature(feature.cudaHandle); + feature.cudaHandle = nullptr; + g.ngx.destroyParameters(feature.cudaParams); feature.cudaParams = nullptr; + if (feature.scratch != 0) g.cu.memFree(feature.scratch); + feature.scratch = 0; + return created; + } + feature.cudaWidth = feature.width; + feature.cudaHeight = feature.height; + if (created == NGX_SUCCESS && env_uint("D4R_SHIM_BLOCKING_SYNC", 1) != 0 && + g.cu.eventCreate != nullptr && g.cu.eventRecord != nullptr && + g.cu.eventSynchronize != nullptr && g.cu.eventQuery != nullptr && g.cu.eventDestroy != nullptr) + { + // Explicit query/sleep waits avoid HIP's spinning event-sync path. + // Timing is unnecessary for this completion event. + constexpr unsigned int kBlockingSyncNoTiming = 0x1u | 0x2u; + const int eventResult = g.cu.eventCreate(&feature.outputReadyEvent, kBlockingSyncNoTiming); + if (eventResult != 0) + { + logf("blocking output wait unavailable (cuEventCreate %d); using context sync", eventResult); + feature.outputReadyEvent = nullptr; + } + else + logf("output event query/sleep wait enabled for this feature"); + } + if (created == NGX_SUCCESS && profile_enabled() && g.cu.eventCreate != nullptr && + g.cu.eventRecord != nullptr && g.cu.eventElapsedTime != nullptr && g.cu.eventDestroy != nullptr) + { + const int startResult = g.cu.eventCreate(&feature.profileStart, 0); + const int endResult = startResult == 0 ? g.cu.eventCreate(&feature.profileEnd, 0) : -1; + feature.profileEventsReady = startResult == 0 && endResult == 0; + if (!feature.profileEventsReady) + { + logf("D4R_PROFILE: CUDA event creation failed (%d, %d)", startResult, endResult); + if (feature.profileStart != nullptr) + g.cu.eventDestroy(feature.profileStart); + if (feature.profileEnd != nullptr) + g.cu.eventDestroy(feature.profileEnd); + feature.profileStart = feature.profileEnd = nullptr; + } + } + return created; + }); + logf("NVSDK_NGX_CUDA_CreateFeature -> 0x%08x", result); + if (result == NGX_SUCCESS) feature.cudaCreationDeferred = false; + return result; +} + static NgxResult create_feature(unsigned int featureId, void* parameters, NgxHandle** handle, bool nativeVulkan) { std::lock_guard apiLock(g_featureApiMutex); @@ -4848,67 +4995,13 @@ static NgxResult create_feature(unsigned int featureId, void* parameters, NgxHan logf("NVSDK_NGX_D3D12_CreateFeature: %ux%u -> %ux%u quality=%d flags=0x%x subrects=%d preset=%u", feature->width, feature->height, feature->outWidth, feature->outHeight, quality, flags, subrects, presets[1]); - const NgxResult result = g.worker.call([&]() -> NgxResult { - NgxResult allocation = g.ngx.allocateParameters(&feature->cudaParams); - if (allocation != NGX_SUCCESS) - return allocation; - void* p = feature->cudaParams; - d4r_ngx_set_uint(p, "Width", feature->width); - d4r_ngx_set_uint(p, "Height", feature->height); - d4r_ngx_set_uint(p, "OutWidth", feature->outWidth); - d4r_ngx_set_uint(p, "OutHeight", feature->outHeight); - d4r_ngx_set_int(p, "PerfQualityValue", quality); - d4r_ngx_set_int(p, "DLSS.Feature.Create.Flags", flags); - d4r_ngx_set_int(p, "DLSS.Enable.Output.Subrects", subrects); - d4r_ngx_set_uint(p, "CreationNodeMask", 1); - d4r_ngx_set_uint(p, "VisibilityNodeMask", 1); - for (int index = 0; index < 6; ++index) - d4r_ngx_set_uint(p, presetNames[index], presets[index]); - if (g.cu.memAlloc(&feature->scratch, 64ull * 1024 * 1024) == 0) - { - d4r_ngx_set_void(p, "Scratch", reinterpret_cast(static_cast(feature->scratch))); - d4r_ngx_set_ull(p, "Scratch.SizeInBytes", 64ull * 1024 * 1024); - } - else - feature->scratch = 0; - const NgxResult created = g.ngx.createFeature(NGX_FEATURE_SUPER_SAMPLING, p, &feature->cudaHandle); - feature->cudaWidth = feature->width; - feature->cudaHeight = feature->height; - if (created == NGX_SUCCESS && env_uint("D4R_SHIM_BLOCKING_SYNC", 1) != 0 && - g.cu.eventCreate != nullptr && g.cu.eventRecord != nullptr && - g.cu.eventSynchronize != nullptr && g.cu.eventQuery != nullptr && g.cu.eventDestroy != nullptr) - { - // Explicit query/sleep waits avoid HIP's spinning event-sync path. - // Timing is unnecessary for this completion event. - constexpr unsigned int kBlockingSyncNoTiming = 0x1u | 0x2u; - const int eventResult = g.cu.eventCreate(&feature->outputReadyEvent, kBlockingSyncNoTiming); - if (eventResult != 0) - { - logf("blocking output wait unavailable (cuEventCreate %d); using context sync", eventResult); - feature->outputReadyEvent = nullptr; - } - else - logf("output event query/sleep wait enabled for this feature"); - } - if (created == NGX_SUCCESS && profile_enabled() && g.cu.eventCreate != nullptr && - g.cu.eventRecord != nullptr && g.cu.eventElapsedTime != nullptr && g.cu.eventDestroy != nullptr) - { - const int startResult = g.cu.eventCreate(&feature->profileStart, 0); - const int endResult = startResult == 0 ? g.cu.eventCreate(&feature->profileEnd, 0) : -1; - feature->profileEventsReady = startResult == 0 && endResult == 0; - if (!feature->profileEventsReady) - { - logf("D4R_PROFILE: CUDA event creation failed (%d, %d)", startResult, endResult); - if (feature->profileStart != nullptr) - g.cu.eventDestroy(feature->profileStart); - if (feature->profileEnd != nullptr) - g.cu.eventDestroy(feature->profileEnd); - feature->profileStart = feature->profileEnd = nullptr; - } - } - return created; - }); - logf("NVSDK_NGX_CUDA_CreateFeature -> 0x%08x", result); + feature->subrects = subrects; + std::copy(std::begin(presets), std::end(presets), feature->creationPresets.begin()); + // The engine records D3D12 evaluations only; the native Vulkan frontend keeps eager CUDA creation. + feature->cudaCreationDeferred = !nativeVulkan && engine_requested() && feature->preset == 11; + const NgxResult result = feature->cudaCreationDeferred ? NGX_SUCCESS : create_cuda_feature(*feature); + if (feature->cudaCreationDeferred) + logf("engine backend: CUDA feature creation deferred until fallback is needed"); if (result != NGX_SUCCESS) { release_feature(feature); @@ -4961,6 +5054,1107 @@ static ID3D12Resource* get_resource(void* parameters, const char* name) return nullptr; } +// --- Native engine backend (engine/): no NGX evaluation, CUDA, worker or frame split ------------- +// [Engine] Enabled in d4r.ini (D4R_ENGINE=1) runs presets K and M with d4r's own shaders inside the game's command +// list. The models (the output of engine/compile_k.py / compile_m.py) are the folders engine\k, engine\k-ldr, +// engine\m and engine\m-ldr next to this DLL, or [Engine] ModelDir (D4R_ENGINE_MODEL_DIR) and the +// D4R_ENGINE_MODEL_{LDR,M,M_LDR}_DIR variables. Anything the engine does not cover falls back to the CUDA backend. +struct EngineBackend +{ + std::once_flag vulkanOnce; // the game's device and the engine's Vulkan entry points, for every preset + std::once_flag once; + std::unique_ptr model; // shared by all features, kept for the life of the process + // The same network with NGX's LDR input and output stages, for games that do not set the HDR flag + // (engine\\k-ldr next to this DLL, or D4R_ENGINE_MODEL_LDR_DIR); loaded on first use. + std::once_flag ldrOnce; + std::unique_ptr modelLdr; + std::string directory, directoryLdr; // kept for the native HIP network's hipnet.bin and hip/*.hsaco +}; +static EngineBackend g_engine; + +static bool engine_requested() +{ + static const bool requested = env_uint("D4R_ENGINE", 0) != 0; + return requested; +} + +static void init_engine_vulkan() +{ + VkInstance instance = VK_NULL_HANDLE; + if (g_vk.interop == nullptr || g_vk.lifetime == nullptr || + FAILED(g_vk.interop->GetVulkanHandles(&instance, &g_vk.physical, &g_vk.device))) + { + logf("engine backend: the vkd3d-proton interop interfaces are unavailable"); + return; + } + HMODULE vulkan = GetModuleHandleA("vulkan-1.dll"); + PFN_vkGetInstanceProcAddr instanceProc = nullptr; + if (vulkan == nullptr || !load_export(vulkan, "vkGetInstanceProcAddr", instanceProc) || + !d4r::loadVulkan(instanceProc, instance, g_vk.device)) + { + logf("engine backend: Vulkan entry points are unavailable"); + return; + } +} + +static void init_engine_backend() +{ + // SetEnvironmentVariable (d4r.ini) does not reach the C runtime's getenv + std::string directory = env_string("D4R_ENGINE_MODEL_DIR"); + if (directory.empty()) + { + char folder[MAX_PATH * 3] = {}; + WideCharToMultiByte(CP_ACP, 0, (g_portable.dir + L"\\engine\\k").c_str(), -1, folder, sizeof(folder), nullptr, nullptr); + directory = folder; + } + if (g_vk.interop == nullptr || g_vk.lifetime == nullptr || g_vk.device == VK_NULL_HANDLE || d4r::vkTable.CreateBuffer == nullptr) + return; + // A Unix path names the same file through Wine's Z: drive. + const std::string path = directory[0] == '/' ? "Z:" + directory : directory; + const auto start = std::chrono::steady_clock::now(); + try + { + g_engine.model = std::make_unique(g_vk.physical, g_vk.device, path); + g_engine.directory = path; + } + catch (const std::exception& error) + { + logf("engine backend: cannot load the model from %s: %s (see D4R_README.txt / engine/README.md for how to " + "create it)", path.c_str(), error.what()); + return; + } + logf("engine backend: model %s loaded and pipelines compiled in %.0f ms", path.c_str(), + std::chrono::duration(std::chrono::steady_clock::now() - start).count()); +} + +static void init_engine_backend_ldr() +{ + std::string directory = env_string("D4R_ENGINE_MODEL_LDR_DIR"); + if (directory.empty()) + { + char folder[MAX_PATH * 3] = {}; + WideCharToMultiByte(CP_ACP, 0, (g_portable.dir + L"\\engine\\k-ldr").c_str(), -1, folder, sizeof(folder), nullptr, nullptr); + directory = folder; + } + if (g_vk.interop == nullptr || g_vk.lifetime == nullptr || g_vk.device == VK_NULL_HANDLE || d4r::vkTable.CreateBuffer == nullptr) + return; + const std::string path = directory[0] == '/' ? "Z:" + directory : directory; + try + { + g_engine.modelLdr = std::make_unique(g_vk.physical, g_vk.device, path); + g_engine.directoryLdr = path; + logf("engine backend: LDR model %s loaded", path.c_str()); + } + catch (const std::exception& error) + { + logf("engine backend: cannot load the LDR model from %s: %s", path.c_str(), error.what()); + } +} + +// K and M borrow compatible views through the same bounded, resource-owning cache. +static void prepare_engine_view(Feature& feature, ID3D12Resource* resource, const D3D12_RESOURCE_DESC& desc, + d4r::GameImage& image, bool isOutput, bool originSupported, + D3D12_RESOURCE_STATES& state) +{ + if (image.format == VK_FORMAT_D32_SFLOAT || image.format == VK_FORMAT_D32_SFLOAT_S8_UINT) + image.aspect = VK_IMAGE_ASPECT_DEPTH_BIT; + const VkImageUsageFlags usage = isOutput ? VK_IMAGE_USAGE_STORAGE_BIT : VK_IMAGE_USAGE_SAMPLED_BIT; + VkFormatProperties properties{}; + d4r::vkTable.GetPhysicalDeviceFormatProperties(g_vk.physical, image.format, &properties); + const bool enabled = env_uint(isOutput ? "D4R_ENGINE_DIRECT_OUTPUT" : "D4R_ENGINE_DIRECT_INPUTS", 1) != 0; + const bool compatible = desc.DepthOrArraySize == 1 && originSupported && (isOutput ? + (image.format == VK_FORMAT_R16G16B16A16_SFLOAT || image.format == VK_FORMAT_A2B10G10R10_UNORM_PACK32) && + (desc.Flags & D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS) && + (properties.optimalTilingFeatures & VK_FORMAT_FEATURE_STORAGE_IMAGE_BIT) : + !(desc.Flags & D3D12_RESOURCE_FLAG_DENY_SHADER_RESOURCE) && + (properties.optimalTilingFeatures & VK_FORMAT_FEATURE_SAMPLED_IMAGE_BIT)); + if (!enabled || !compatible) + return; + auto found = std::find_if(feature.engineViews.begin(), feature.engineViews.end(), [&](const auto& view) { + return view->resource == resource && view->format == image.format && view->aspect == image.aspect && + view->usage == usage; + }); + if (found == feature.engineViews.end()) + { + VkMemoryRequirements memory{}; + d4r::vkTable.GetImageMemoryRequirements(g_vk.device, image.image, &memory); + constexpr uint64_t cacheBudget = 512ull * 1024 * 1024; + if (feature.engineViews.size() < 64 && memory.size <= cacheBudget - feature.engineViewBytes) + { + auto view = std::make_unique(); + view->device = g_vk.device; view->format = image.format; + view->aspect = image.aspect; view->usage = usage; + VkImageViewUsageCreateInfo restricted{VK_STRUCTURE_TYPE_IMAGE_VIEW_USAGE_CREATE_INFO}; restricted.usage = usage; + VkImageViewCreateInfo info{VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO, &restricted}; + info.image = image.image; info.viewType = VK_IMAGE_VIEW_TYPE_2D; info.format = view->format; + info.subresourceRange = {view->aspect, 0, 1, 0, 1}; + if (d4r::vkTable.CreateImageView(g_vk.device, &info, nullptr, &view->view) == VK_SUCCESS) + { + resource->AddRef(); view->resource = resource; + feature.engineViews.push_back(std::move(view)); + feature.engineViewBytes += memory.size; + found = feature.engineViews.end() - 1; + } + } + else if (!feature.engineViewBudgetReported) + { + feature.engineViewBudgetReported = true; + logf("engine view cache budget reached; new textures use copies until feature retirement"); + } + } + if (found != feature.engineViews.end()) + { + image.view = (*found)->view; + state = isOutput ? D3D12_RESOURCE_STATE_UNORDERED_ACCESS : D3D12_RESOURCE_STATE_NON_PIXEL_SHADER_RESOURCE; + } +} + +// Preset M on the engine: its own model (engine\\m next to this DLL, or D4R_ENGINE_MODEL_M_DIR) and adapter. +struct EngineBackendM +{ + std::once_flag once, ldrOnce; + std::unique_ptr model, modelLdr; // LDR: engine\\m-ldr or D4R_ENGINE_MODEL_M_LDR_DIR + std::string directory, directoryLdr; +}; +static EngineBackendM g_engineM; + +static void init_engine_backend_m_variant(bool ldr) +{ + std::unique_ptr& slot = ldr ? g_engineM.modelLdr : g_engineM.model; + std::string directory = env_string(ldr ? "D4R_ENGINE_MODEL_M_LDR_DIR" : "D4R_ENGINE_MODEL_M_DIR"); + if (directory.empty()) + { + char folder[MAX_PATH * 3] = {}; + WideCharToMultiByte(CP_ACP, 0, (g_portable.dir + (ldr ? L"\\engine\\m-ldr" : L"\\engine\\m")).c_str(), -1, folder, sizeof(folder), nullptr, nullptr); + directory = folder; + } + if (g_vk.interop == nullptr || g_vk.lifetime == nullptr || g_vk.device == VK_NULL_HANDLE || d4r::vkTable.CreateBuffer == nullptr) + { + logf("engine backend (M): Vulkan interop is unavailable"); + return; + } + const std::string path = directory[0] == '/' ? "Z:" + directory : directory; + const auto start = std::chrono::steady_clock::now(); + try + { + slot = std::make_unique(g_vk.physical, g_vk.device, path); + (ldr ? g_engineM.directoryLdr : g_engineM.directory) = path; + } + catch (const std::exception& error) + { + logf("engine backend (M): cannot load the model from %s: %s", path.c_str(), error.what()); + return; + } + logf("engine backend (M): model %s loaded and pipelines compiled in %.0f ms", path.c_str(), + std::chrono::duration(std::chrono::steady_clock::now() - start).count()); +} + +// Preset M's network on d4r's native HIP layer kernels (bridge: d4rEngineNet*; engine/hip_net.h). HIP runs the +// network about 1.7x faster than the engine's cooperative-matrix shader; the stages around it stay on Vulkan in +// the game's command list, which is split at the network like a CUDA split frame. The ten layers wait on the GPU +// for the front's fill of the frame number (d4rEngineNetInputWaitReady/d4rEngineNetLaunchAt, D4R_SHIM_ENGINE_GPU_WAIT) +// when the bridge supports it; otherwise the worker waits for the CPU-visible marker before launching. +struct EngineNetCode +{ + const char* name; + const void* data; + uint64_t bytes; +}; +struct EngineNetExports +{ + std::once_flag once; + int(WINAPI* create)(const uint8_t*, uint64_t, const EngineNetCode*, uint32_t, uint32_t, uint32_t, CudaDevicePtr, + CudaDevicePtr, void**) = nullptr; + int(WINAPI* launch)(void*, uint32_t) = nullptr; + int(WINAPI* done)(void*, float*) = nullptr; + void(WINAPI* destroy)(void*) = nullptr; + // GPU-side readiness: the front fills a u32; the HIP stream waits on it instead of a CPU-visible marker. + int(WINAPI* inputWaitReady)(void*, CudaDevicePtr, uint32_t*) = nullptr; + int(WINAPI* launchAt)(void*, uint32_t) = nullptr; + int(WINAPI* releaseInput)(void*, uint32_t) = nullptr; + int(WINAPI* setTimeline)(void*, VkDevice, uint64_t) = nullptr; + int(WINAPI* signalTimeline)(void*, uint64_t) = nullptr; + bool ready = false; +}; +static EngineNetExports g_engineNet; + +static VkResult signal_split_semaphore(Feature& feature, uint64_t value) +{ + if (feature.engineNetTimeline) + return static_cast(g_engineNet.signalTimeline(feature.engineNet, value)); + VkSemaphoreSignalInfo info = {VK_STRUCTURE_TYPE_SEMAPHORE_SIGNAL_INFO}; + info.semaphore = feature.splitSemaphore; + info.value = value; + return g_vk.signalSemaphore(g_vk.device, &info); +} + +static bool read_file(const std::string& path, std::vector& bytes) +{ + FILE* file = fopen(path.c_str(), "rb"); + if (file == nullptr) + return false; + fseek(file, 0, SEEK_END); + const long size = ftell(file); + fseek(file, 0, SEEK_SET); + bytes.resize(size > 0 ? static_cast(size) : 0); + const bool ok = size > 0 && fread(bytes.data(), 1, bytes.size(), file) == bytes.size(); + fclose(file); + return ok; +} + +// A storage buffer the engine's shaders and HIP both use. Pooled like the interop buffers, for the same ROCm leak. +static bool create_engine_shared_buffer(VramBuffer& target, size_t bytes, bool host = false) +{ + const uint8_t kind = host ? 2 : 1; + if (take_pooled_vram_buffer(target, bytes, kind)) + { + if (target.mapped != nullptr) + { + std::memset(target.mapped, 0, target.bytes); + MemoryBarrier(); + } + return true; + } + VramBuffer buffer; + buffer.kind = kind; + VkExternalMemoryBufferCreateInfo external = {VK_STRUCTURE_TYPE_EXTERNAL_MEMORY_BUFFER_CREATE_INFO}; + external.handleTypes = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32_BIT; + VkBufferCreateInfo info = {VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO, &external}; + info.size = bytes; + info.usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT; + info.sharingMode = VK_SHARING_MODE_EXCLUSIVE; + VkResult result = g_vk.createBuffer(g_vk.device, &info, nullptr, &buffer.buffer); + VkMemoryRequirements requirements = {}; + if (result == VK_SUCCESS) + g_vk.bufferRequirements(g_vk.device, buffer.buffer, &requirements); + int type = device_local_type(requirements.memoryTypeBits); + if (host) + { + type = -1; + const auto coherent = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT; + for (int pass = 0; pass < 2 && type < 0; ++pass) + for (uint32_t index = 0; index < g_vk.memory.memoryTypeCount; ++index) + { + const auto flags = g_vk.memory.memoryTypes[index].propertyFlags; + if ((requirements.memoryTypeBits & (1u << index)) && (flags & coherent) == coherent && + (pass != 0 || !(flags & VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT))) + { + type = static_cast(index); + break; + } + } + } + VkMemoryDedicatedAllocateInfo dedicated = {VK_STRUCTURE_TYPE_MEMORY_DEDICATED_ALLOCATE_INFO}; + dedicated.buffer = buffer.buffer; + VkExportMemoryAllocateInfo exportInfo = {VK_STRUCTURE_TYPE_EXPORT_MEMORY_ALLOCATE_INFO, &dedicated}; + exportInfo.handleTypes = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32_BIT; + VkMemoryAllocateInfo allocation = {VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO, &exportInfo}; + allocation.allocationSize = requirements.size; + allocation.memoryTypeIndex = static_cast(type); + if (result == VK_SUCCESS && type < 0) + result = VK_ERROR_OUT_OF_DEVICE_MEMORY; + if (result == VK_SUCCESS) + result = g_vk.allocate(g_vk.device, &allocation, nullptr, &buffer.memory); + if (result == VK_SUCCESS) + result = g_vk.bindBuffer(g_vk.device, buffer.buffer, buffer.memory, 0); + if (result == VK_SUCCESS && host) + result = g_vk.map(g_vk.device, buffer.memory, 0, VK_WHOLE_SIZE, 0, &buffer.mapped); + int imported = -1; + if (result == VK_SUCCESS) + imported = g.worker.call([&] { + return g_vk.import(g_vk.device, reinterpret_cast(buffer.memory), requirements.size, &buffer.device, + &buffer.external); + }); + if (result != VK_SUCCESS || imported != 0) + { + logf("engine backend: shared buffer of %zu bytes failed (VkResult %d, import %d)", bytes, result, imported); + if (buffer.mapped != nullptr) + g_vk.unmap(g_vk.device, buffer.memory); + if (buffer.memory != VK_NULL_HANDLE) + g_vk.free(g_vk.device, buffer.memory, nullptr); + if (buffer.buffer != VK_NULL_HANDLE) + g_vk.destroyBuffer(g_vk.device, buffer.buffer, nullptr); + return false; + } + if (buffer.mapped != nullptr) + { + std::memset(buffer.mapped, 0, bytes); + MemoryBarrier(); + } + buffer.bytes = bytes; + target = buffer; + return true; +} + +// Only once no command list uses the engine and no network launch is pending. +static void release_engine_network(Feature& feature) +{ + if (feature.engineNet != nullptr && g_engineNet.destroy != nullptr) + g.worker.call([&feature] { + g_engineNet.destroy(feature.engineNet); + return 0; + }); + feature.engineNet = nullptr; + feature.engineNetTimeline = false; + for (VramBuffer* buffer : {&feature.engineTokens, &feature.engineHead, &feature.engineSync}) + recycle_vram_buffer(*buffer); +} + +// The HIP network of a feature: the buffers shared with HIP and the bridge's network over them. False: the +// engine's own Vulkan network is used instead (reason logged). `preset` is "M" or "K"; `width` +// and `height` are the dimensions the bridge derives the token grid from (M: render, K: output); `tokenBytes` +// and `headBytes` size the two shared buffers (M: both MEngine::networkBufferBytes, K: Engine::networkTokenBytes +// and Engine::networkHeadBytes). The layer kernel names come from hipnet.bin, so no preset list is repeated here. +static bool create_engine_network(Feature& feature, const std::string& directory, const char* preset, + bool stagesAvailable, unsigned int width, unsigned int height, size_t tokenBytes, + size_t headBytes) +{ + const std::string choice = env_string("D4R_ENGINE_NETWORK"); + if (choice == "vulkan") + return false; + auto unavailable = [&](const char* reason) { + logf("engine backend (%s): HIP network unavailable (%s); using the Vulkan network", preset, reason); + release_engine_network(feature); + return false; + }; + std::call_once(g_engineNet.once, [] { + g_engineNet.ready = load_export(g.cuda, "d4rEngineNetCreate", g_engineNet.create) && + load_export(g.cuda, "d4rEngineNetLaunch", g_engineNet.launch) && + load_export(g.cuda, "d4rEngineNetDone", g_engineNet.done) && + load_export(g.cuda, "d4rEngineNetDestroy", g_engineNet.destroy) && + load_export(g.cuda, "d4rEngineNetInputWaitReady", g_engineNet.inputWaitReady) && + load_export(g.cuda, "d4rEngineNetLaunchAt", g_engineNet.launchAt) && + load_export(g.cuda, "d4rEngineNetReleaseInput", g_engineNet.releaseInput) && + load_export(g.cuda, "d4rEngineNetSetTimeline", g_engineNet.setTimeline) && + load_export(g.cuda, "d4rEngineNetSignalTimeline", g_engineNet.signalTimeline); + }); + if (!g_engineNet.ready) + return unavailable("the CUDA bridge lacks d4rEngineNet*"); + if (!stagesAvailable) + return unavailable("the model has no stages for the native network; rebuild it with --hip"); + if (g_vk.split == nullptr || g_vk.import == nullptr || feature.marker == nullptr || feature.markerValue == nullptr) + return unavailable("split frames need the d4r vkd3d-proton and VRAM interop"); + std::vector net; + if (!read_file(directory + "\\hipnet.bin", net)) + return unavailable("no hipnet.bin in the model directory"); + // The bridge dispatches on the magic alone, and the shared buffers are sized for `preset`. + const char* const magic = std::strcmp(preset, "K") == 0 ? "D4RKHIP1" : "D4RMHIP1"; + if (net.size() < 16 || std::memcmp(net.data(), magic, 8) != 0) + return unavailable("hipnet.bin is not a native network of this preset"); + uint32_t count = 0; + std::memcpy(&count, net.data() + 8, sizeof(count)); + if (count == 0 || net.size() < 16 + 64ull * count) + return unavailable("invalid hipnet.bin"); + // The kernel names live in the layer table; load the code object of each distinct one (the M tube repeats). + std::vector names; + names.reserve(count); + for (uint32_t index = 0; index < count; ++index) + { + char name[33] = {}; + std::memcpy(name, net.data() + 16 + 64 * index, 32); + if (name[0] == '\0') + return unavailable("invalid kernel name in hipnet.bin"); + if (std::find(names.begin(), names.end(), name) == names.end()) + names.emplace_back(name); + } + std::vector> code(names.size()); + std::vector codes(names.size()); + for (size_t index = 0; index < names.size(); ++index) + { + if (!read_file(directory + "\\hip\\" + names[index] + ".hsaco", code[index])) + return unavailable("a layer kernel is missing in the model's hip folder"); + codes[index] = {names[index].c_str(), code[index].data(), code[index].size()}; + } + if (!create_engine_shared_buffer(feature.engineTokens, tokenBytes) || + !create_engine_shared_buffer(feature.engineHead, headBytes) || + !create_engine_shared_buffer(feature.engineSync, 256, true)) + return unavailable("shared buffers"); + const int created = g.worker.call([&] { + return g_engineNet.create(reinterpret_cast(net.data()), net.size(), codes.data(), + static_cast(codes.size()), width, height, feature.engineTokens.device, + feature.engineHead.device, &feature.engineNet); + }); + if (created != 0 || feature.engineNet == nullptr) + { + feature.engineNet = nullptr; + logf("engine backend (%s): d4rEngineNetCreate failed (%d)", preset, created); + return unavailable("the kernels could not be loaded for this GPU"); + } + // Input readiness on the GPU: the engine's front fills engineSync with the frame number (the shim records the + // fill right after recordFront) and the network's stream waits for it, so the layers start the moment the + // front is complete, without the CPU round trip through the readback marker. The bridge checks HIP's + // capability and a round trip on this buffer. Without wait-value support (or D4R_SHIM_ENGINE_GPU_WAIT=0) + // the worker keeps waiting for the CPU-visible marker before launching. + const bool gpuWaitEnabled = env_uint("D4R_SHIM_ENGINE_GPU_WAIT", 1) != 0; + const int waitProbe = gpuWaitEnabled ? g.worker.call([&] { + return g_engineNet.inputWaitReady(feature.engineNet, feature.engineSync.device, + static_cast(feature.engineSync.mapped)); + }) : 0; + if (waitProbe < 0) + return unavailable("the HIP stream cannot complete a wait-value test"); + feature.engineGpuWait = waitProbe == 1; + logf("engine backend (%s): HIP network input readiness on %s", preset, feature.engineGpuWait.load() + ? "the GPU (the front publishes the frame number; the layers wait on it)" + : "the CPU (marker polling)"); + if (feature.splitSemaphore == VK_NULL_HANDLE && !create_split_semaphore(feature)) + return unavailable("timeline semaphore"); + if (g.worker.call([&] { + return g_engineNet.setTimeline(feature.engineNet, g_vk.device, + reinterpret_cast(feature.splitSemaphore)); + }) != 0) + return unavailable("native HIP completion timeline"); + feature.engineNetTimeline = true; + logf("engine backend (%s): network completion releases the split timeline from the native HIP callback", preset); + return true; +} + +// Records the end of the front: a fill of this frame's number into the buffer the HIP stream waits on +// (d4rEngineNetLaunchAt). The engine's recordFront ends the input stage with its own ALL_COMMANDS barrier (the +// tokens are then available); the fill is ordered after it and, once it lands, the ten layers may read them. +static void record_engine_input_ready(Feature& feature, VkCommandBuffer command, uint32_t frame) +{ + const VkMemoryBarrier before = {VK_STRUCTURE_TYPE_MEMORY_BARRIER, nullptr, + VK_ACCESS_MEMORY_WRITE_BIT, VK_ACCESS_TRANSFER_WRITE_BIT}; + g_vk.barrier(command, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 1, &before, 0, nullptr, + 0, nullptr); + g_vk.fill(command, feature.engineSync.buffer, 0, sizeof(uint32_t), frame); + const VkMemoryBarrier after = {VK_STRUCTURE_TYPE_MEMORY_BARRIER, nullptr, VK_ACCESS_TRANSFER_WRITE_BIT, + VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT}; + g_vk.barrier(command, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, 0, 1, &after, 0, nullptr, + 0, nullptr); +} + +static bool wait_for_input_marker(Feature* feature, int slotIndex, uint32_t frame); +static void frame_retired(Feature* feature, uint32_t frame); + +// The network's inputs are written by the front of the game's list, so its stream waits for the frame's fill of +// the shared u32 (d4rEngineNetLaunchAt). A wait for a frame whose front never reached the GPU (the game recorded +// its list but did not execute it, or the game stalled) can only be satisfied by a later frame's fill: writing +// the value it waits for releases it, so the stream keeps moving. The frame's output is undefined either way. +// `disable`: the GPU-side wait itself looks broken for this feature, so later frames wait for the CPU marker. +static void release_engine_input_wait(Feature& feature, uint32_t frame, const char* reason, bool disable) +{ + if (feature.engineNet == nullptr || g_engineNet.releaseInput == nullptr) + return; + const int result = g_engineNet.releaseInput(feature.engineNet, frame); + if (disable) + feature.engineGpuWait.store(false); + logf("engine frame %u: %s; releasing the GPU-side input wait (%d)", frame, reason, result); +} + +// CUDA worker: runs the network of `frame` (behind the GPU-side input wait, or once the CPU marker shows the +// front has run), then lets the second part of the game's command list (the engine's reconstruction) run. +static void run_engine_network(Feature* feature, int slotIndex, uint32_t frame) +{ + const bool gpuWait = feature->engineGpuWait.load() && g_engineNet.launchAt != nullptr; + if (!gpuWait && !wait_for_input_marker(feature, slotIndex, frame)) + return; // dropped and retired there + const auto start = std::chrono::steady_clock::now(); + int launched; + if (gpuWait) + { + feature->engineWaitFrame.store(frame); + launched = g_engineNet.launchAt(feature->engineNet, frame); + if (launched == 801 /* CUDA_ERROR_NOT_SUPPORTED: no wait was queued */) + { + feature->engineWaitFrame.store(0); + // The bridge could not queue the wait (no hipStreamWaitValue32, or the stream rejected it): wait + // for the input marker on the CPU, as without the GPU-side wait. + feature->engineGpuWait.store(false); + logf("engine frame %u: GPU-side input wait unavailable (%d); waiting for the input marker on the CPU", frame, launched); + if (!wait_for_input_marker(feature, slotIndex, frame)) + return; + launched = g_engineNet.launch(feature->engineNet, frame); + } + } + else + launched = g_engineNet.launch(feature->engineNet, frame); + float gpuMs = 0.0f; + int done = launched == 0 ? 0 : -1; + bool released = false, discarded = false; + // Completion polling (as the CUDA path's output event): sleep while far from the predicted completion, + // yield around it (a Wine sleep rounds up to ~50 us), back off when the GPU runs late. The prediction + // learns from the earlier frames; with the GPU-side wait it also covers the front's remaining GPU time. + double lastNotReadyUs = -1.0; + while (done == 0) + { + done = g_engineNet.done(feature->engineNet, &gpuMs); + if (done != 0) + break; + const double elapsedUs = std::chrono::duration(std::chrono::steady_clock::now() - start).count(); + if (gpuWait && !released) + { + // This frame's own marker (written with the aggregate one, at the end of its front) tells whether + // its front ran at all. A later frame's front running while this one's never did means the game + // discarded this list: its wait can only be satisfied by that later front, which would make the + // network read another frame's tokens. Release it and drop the frame, as the CPU marker path does. + const uint32_t slotMarker = feature->markerValue[1 + slotIndex]; + if (slotMarker != frame && static_cast(*feature->markerValue - frame) > 0) + { + release_engine_input_wait(*feature, frame, "its command list was discarded", false); + discarded = true; + break; + } + // The front ran but the stream's wait did not start the network (a broken wait would leave every + // later frame queued behind it): release it and wait for the CPU marker from now on. + if (slotMarker == frame && elapsedUs > 1000000.0) + { + release_engine_input_wait(*feature, frame, "the GPU-side wait did not start the network", true); + released = true; + } + } + if (elapsedUs > 2000000.0) + { + done = -2; + break; + } + lastNotReadyUs = elapsedUs; + constexpr double kMarginUs = 120.0, kStepUs = 20.0, kChunkUs = 200.0, kSpinAfterUs = 150.0; + const double predictedUs = feature->engineNetWaitEmaUs; + const double remaining = predictedUs - kMarginUs - elapsedUs; + if (remaining > kStepUs) + d4r_sleep_us(static_cast(std::min(remaining, kChunkUs))); + else if (elapsedUs < predictedUs + kSpinAfterUs) + d4r_sleep_us(0); // polling retires resources; the native callback releases the GPU's split wait + else + d4r_sleep_us(static_cast(std::min( + kStepUs * (1.0 + (elapsedUs - predictedUs - kSpinAfterUs) / 160.0), kChunkUs))); + } + feature->engineWaitFrame.store(0); + const double waitedUs = std::chrono::duration(std::chrono::steady_clock::now() - start).count(); + if (done == 1) + { + // Learn the completion time, not our own wait: the work finished between the last NOT_READY poll and the + // successful one (at most one chunk apart), so the midpoint cannot ratchet upwards. + double sample = lastNotReadyUs >= 0.0 ? 0.5 * (lastNotReadyUs + waitedUs) : waitedUs; + sample = std::min(sample, 20000.0); + feature->engineNetWaitEmaUs = feature->engineNetWaitEmaUs == 0.0 + ? sample : 0.8 * feature->engineNetWaitEmaUs + 0.2 * sample; + } + if (discarded) + logf("engine frame %u: the game discarded its command list; the frame is dropped", frame); + else if (done != 1) + logf("engine frame %u: HIP network failed (launch %d, state %d); the frame's output is undefined", frame, launched, done); + else if (frame <= 3 || frame % 600 == 0) + logf("engine frame %u: HIP network %.2f ms GPU, %.2f ms %s", frame, gpuMs, waitedUs / 1000.0, + gpuWait ? "from launch to completion (GPU-side input wait)" : "from input readiness to completion"); + if (gpuWait && !discarded && done != 1) + release_engine_input_wait(*feature, frame, "the network failed before completing its input wait", true); + feature->inputs[slotIndex].busy = false; + frame_retired(feature, frame); +} + +// Preset K and preset M record a native-network frame the same way: the front half, this frame's marker, the +// split of the game's list, the back half, and the worker that runs the network between the two halves. +// `recordPart(front)` records one half into the game's list; `sources` is {colour, depth, motion, output}, the +// first three sampled through the front. `refuse` (nothing usable recorded; the CUDA backend takes over) and +// `abandon` (recorded but the frame is unusable) belong to the caller so its logging and flags stay preset-specific. +struct EngineFrameSource +{ + ID3D12Resource* resource; + d4r::GameImage* image; + unsigned int x, y, width, height; + D3D12_RESOURCE_STATES state, copyState; +}; + +template +static bool record_engine_network_frame(Feature& feature, ID3D12GraphicsCommandList* list, EngineFrameSource* sources, + uint32_t frame, const char* preset, int reset, float mvScaleX, float mvScaleY, + float jitterX, float jitterY, RecordPart&& recordPart, Refuse&& refuse, + Abandon&& abandon) +{ + // HIP waits on the GPU-published input value between the two Vulkan list parts. Its native completion + // callback releases the second part. The Wine worker retires resources, not frame timing. + const int slotIndex = static_cast(frame % kSlots); + InputSlot& slot = feature.inputs[slotIndex]; + const auto waitStart = std::chrono::steady_clock::now(); + while (slot.busy.load() || feature.inFlight.load() >= 2) + { + if (std::chrono::steady_clock::now() - waitStart > std::chrono::seconds(2)) + { + logf("engine frame %u: %d frames in flight; skipping evaluation", frame, feature.inFlight.load()); + return true; + } + d4r_sleep_us(200); + } + if (FAILED(g_vk.lifetime->RetainExternalResources(list, feature.resources))) + return refuse("the command list cannot retain the engine's resources"); + ID3D12GraphicsCommandList2* list2 = nullptr; + if (FAILED(list->QueryInterface(__uuidof(ID3D12GraphicsCommandList2), reinterpret_cast(&list2)))) + return refuse("ID3D12GraphicsCommandList2 is unavailable"); + for (int index = 0; index < 3; ++index) + transition(list, sources[index].resource, sources[index].state, sources[index].copyState); + const char* failure = recordPart(true); + if (failure != nullptr) + { + for (int index = 0; index < 3; ++index) + transition(list, sources[index].resource, sources[index].copyState, sources[index].state); + list2->Release(); + return refuse(failure); + } + const D3D12_GPU_VIRTUAL_ADDRESS markerAddress = feature.marker->GetGPUVirtualAddress(); + D3D12_WRITEBUFFERIMMEDIATE_PARAMETER markers[] = { + {markerAddress + sizeof(uint32_t) * (1 + slotIndex), frame}, + {markerAddress, frame}, + }; + D3D12_WRITEBUFFERIMMEDIATE_MODE modes[] = { + D3D12_WRITEBUFFERIMMEDIATE_MODE_MARKER_OUT, D3D12_WRITEBUFFERIMMEDIATE_MODE_MARKER_OUT, + }; + list2->WriteBufferImmediate(2, markers, modes); + list2->Release(); + // From here on this call owns the frame: a failure leaves this frame's output undefined and hands the next + // frame to the CUDA backend (which must not record its own marker and split for this frame). + if (FAILED(g_vk.split->SplitCommandListForExternalWait(list, reinterpret_cast(feature.splitSemaphore), frame))) + { + for (int index = 0; index < 3; ++index) + transition(list, sources[index].resource, sources[index].copyState, sources[index].state); + return abandon("the command list could not be split"); + } + transition(list, sources[3].resource, sources[3].state, sources[3].copyState); + failure = recordPart(false); + transition(list, sources[3].resource, sources[3].copyState, sources[3].state); + // The reconstruction still samples borrowed colour; keep its sampled layout through the back. + for (int index = 0; index < 3; ++index) + transition(list, sources[index].resource, sources[index].copyState, sources[index].state); + feature.split = true; + slot.busy = true; + feature.inFlight.fetch_add(1); + { + std::lock_guard lock(feature.splitMutex); + feature.splitCompletion.admit(frame); + } + Feature* const owner = &feature; + g.worker.post([owner, slotIndex, frame] { run_engine_network(owner, slotIndex, frame); }); + if (failure != nullptr) + return abandon(failure); + feature.engineJitter[0] = jitterX; + feature.engineJitter[1] = jitterY; + if (frame <= 3 || frame % 600 == 0) + logf("engine frame %u recorded, preset %s with the HIP network (jitter %.6f,%.6f, reset %d, mvScale %.3f,%.3f)", + frame, preset, jitterX, jitterY, reset, mvScaleX, mvScaleY); + return true; +} + +// Preset M: record the evaluation into `list`, borrowing compatible inputs and copying only fallbacks. +static bool engine_evaluate_m(Feature& feature, ID3D12GraphicsCommandList* list, ID3D12Resource* color, + ID3D12Resource* depth, ID3D12Resource* motion, ID3D12Resource* output, + const FrameParams& p, uint32_t frame) +{ + auto refuse = [&](const char* reason) { + feature.engineRefused = true; + logf("engine backend (M): %s; feature %u uses the CUDA backend from frame %u", reason, feature.handle.Id, frame); + return false; + }; + // NVSDK_NGX_DLSS_Feature_Flags: IsHDR 1, MVLowRes 2, MVJittered 4, DepthInverted 8, AutoExposure 0x40. + const bool ldr = (feature.flags & 0x1) == 0; // NGX's LDR stages: a model directory of its own, no exposure + if (ldr) + std::call_once(g_engineM.ldrOnce, [] { init_engine_backend_m_variant(true); }); + else + std::call_once(g_engineM.once, [] { init_engine_backend_m_variant(false); }); + d4r::MModel* const model = ldr ? g_engineM.modelLdr.get() : g_engineM.model.get(); + if (model == nullptr) + return refuse(ldr ? "no LDR model (engine\\m-ldr) for a game without the HDR flag" : "not available"); + if ((feature.flags & 0x2) == 0) + return refuse("needs render-resolution motion vectors"); + // NGX measures the exposure for this preset, so the AutoExposure flag, the game's exposure texture, the + // pre-exposure and the exposure scale are all ignored here (CUDA output is identical for each of them). + if (feature.engineM != nullptr && + (p.renderWidth != feature.engineMRenderWidth || p.renderHeight != feature.engineMRenderHeight)) + return refuse("changing the effective render resolution after initialization is not implemented"); + if (color == output || depth == output || motion == output) + return refuse("input/output texture aliasing is not supported"); + d4r::GameFrameM game; + using Source = EngineFrameSource; + const auto inputState = static_cast(env_uint("D4R_SHIM_INPUT_STATE", 0x40)); + const auto depthState = static_cast(env_uint("D4R_SHIM_DEPTH_STATE", 0x40)); + const auto outputState = static_cast(env_uint("D4R_SHIM_OUTPUT_STATE", 0x8)); + Source sources[4] = { + {color, &game.color, p.colorBaseX, p.colorBaseY, p.renderWidth, p.renderHeight, inputState, D3D12_RESOURCE_STATE_COPY_SOURCE}, + {depth, &game.depth, p.depthBaseX, p.depthBaseY, p.renderWidth, p.renderHeight, depthState, D3D12_RESOURCE_STATE_COPY_SOURCE}, + {motion, &game.motion, p.mvBaseX, p.mvBaseY, p.renderWidth, p.renderHeight, inputState, D3D12_RESOURCE_STATE_COPY_SOURCE}, + {output, &game.output, p.outputBaseX, p.outputBaseY, feature.outWidth, feature.outHeight, outputState, + D3D12_RESOURCE_STATE_COPY_DEST}, + }; + for (Source& source : sources) + { + UINT64 handle = 0, offset = 0; + D3D12_RESOURCE_DESC desc; + source.resource->GetDesc(&desc); + if (FAILED(g_vk.interop->GetVulkanResourceInfo1(source.resource, &handle, &offset, &source.image->format)) || + handle == 0 || desc.Dimension != D3D12_RESOURCE_DIMENSION_TEXTURE2D || desc.SampleDesc.Count != 1 || + source.x + source.width > desc.Width || source.y + source.height > desc.Height) + return refuse("a texture is not a single-sample 2D image covering its rectangle"); + source.image->image = reinterpret_cast(handle); + source.image->x = source.x; + source.image->y = source.y; + prepare_engine_view(feature, source.resource, desc, *source.image, source.image == &game.output, + source.x == 0 && source.y == 0, source.copyState); + g_vk.interop->GetVulkanImageLayout(source.resource, source.copyState, &source.image->layout); + } + if (game.depth.format == VK_FORMAT_D32_SFLOAT || game.depth.format == VK_FORMAT_D32_SFLOAT_S8_UINT) + game.depth.aspect = VK_IMAGE_ASPECT_DEPTH_BIT; + if (!d4r::gameColorFormat(game.color.format) || !d4r::gameMotionFormat(game.motion.format) || + !d4r::gameDepthFormat(game.depth.format, game.depth.aspect) || !d4r::gameOutputFormat(game.output.format)) + { + logf("engine backend (M): Vulkan formats colour %d, depth %d, motion %d, output %d", game.color.format, + game.depth.format, game.motion.format, game.output.format); + return refuse("a texture format is not implemented"); + } + game.jitter[0] = p.jitterX; + game.jitter[1] = p.jitterY; + game.motionToRender[0] = p.mvScaleX; + game.motionToRender[1] = p.mvScaleY; + if ((feature.flags & 0x4) != 0 && p.mvScaleX != 0.0f && p.mvScaleY != 0.0f) + { + game.motionOffset[0] = (p.jitterX - feature.engineJitter[0]) / p.mvScaleX; + game.motionOffset[1] = (p.jitterY - feature.engineJitter[1]) / p.mvScaleY; + } + game.reset = p.reset != 0 || feature.engineM == nullptr; + game.depthInverted = (feature.flags & 0x8) != 0; + try + { + if (feature.engineM == nullptr) + { + const size_t bytes = static_cast(d4r::MEngine::networkBufferBytes(p.renderWidth, p.renderHeight)); + const bool hip = create_engine_network(feature, ldr ? g_engineM.directoryLdr : g_engineM.directory, "M", + model->hasExternalNetworkStages(), p.renderWidth, p.renderHeight, + bytes, bytes); + const d4r::ExternalNetwork external = {feature.engineTokens.buffer, feature.engineHead.buffer}; + feature.engineM = std::make_unique(*model, g_vk.physical, g_vk.device, feature.outWidth, + feature.outHeight, p.renderWidth, p.renderHeight, + hip ? &external : nullptr); + feature.engineMRenderWidth = p.renderWidth; + feature.engineMRenderHeight = p.renderHeight; + logf("engine backend (M): feature %u runs preset M in the game's command list (%ux%u -> %ux%u), network on %s, RGB10A2 direct output %d", + feature.handle.Id, p.renderWidth, p.renderHeight, feature.outWidth, feature.outHeight, + hip ? "the native HIP layers (list split around it)" : "Vulkan", + model->hasRgb10Output() && game.output.view != VK_NULL_HANDLE && + game.output.format == VK_FORMAT_A2B10G10R10_UNORM_PACK32 && + game.output.layout == VK_IMAGE_LAYOUT_GENERAL && game.output.image != game.color.image && + game.output.image != game.motion.image && game.output.image != game.depth.image); + } + } + catch (const std::exception& error) + { + logf("engine backend (M): %s", error.what()); + release_engine_network(feature); + return refuse("the engine could not be created"); + } + if (feature.engineNet != nullptr) + { + // The native HIP network: the game's list is split around it. Its native completion callback releases + // the second part, which records expansion, reconstruction and the final store (or the conversion blit + // for incompatible outputs). The Wine worker retires resources, not frame timing. + auto recordPart = [&](bool front) -> const char* { + VkCommandBuffer command = VK_NULL_HANDLE; + if (FAILED(g_vk.interop->BeginVkCommandBufferInterop(list, &command))) + return "BeginVkCommandBufferInterop failed"; + const char* failure = nullptr; + try + { + if (front) + { + feature.engineM->recordFront(command, game); + // The HIP network's stream waits for this frame's number (record_engine_input_ready): the + // fill is part of the first half, so the network starts the moment the front's GPU work ends. + if (feature.engineGpuWait.load()) + record_engine_input_ready(feature, command, frame); + } + else + feature.engineM->recordBack(command, game); + } + catch (const std::exception& error) + { + logf("engine backend (M): %s", error.what()); + failure = "the frame could not be recorded"; + } + if (FAILED(g_vk.interop->EndVkCommandBufferInterop(list)) && failure == nullptr) + failure = "EndVkCommandBufferInterop failed"; + return failure; + }; + auto abandon = [&](const char* reason) { + feature.engineRefused = true; + logf("engine backend (M): %s; feature %u uses the CUDA backend after frame %u", reason, feature.handle.Id, frame); + return true; + }; + return record_engine_network_frame(feature, list, sources, frame, "M", game.reset, p.mvScaleX, p.mvScaleY, + p.jitterX, p.jitterY, recordPart, refuse, abandon); + } + if (FAILED(g_vk.lifetime->RetainExternalResources(list, feature.resources))) + return refuse("the command list cannot retain the engine's resources"); + for (const Source& source : sources) + transition(list, source.resource, source.state, source.copyState); + VkCommandBuffer command = VK_NULL_HANDLE; + bool recorded = SUCCEEDED(g_vk.interop->BeginVkCommandBufferInterop(list, &command)); + const char* failure = recorded ? nullptr : "BeginVkCommandBufferInterop failed"; + if (recorded) + { + try + { + feature.engineM->record(command, game); + } + catch (const std::exception& error) + { + logf("engine backend (M): %s", error.what()); + failure = "the frame could not be recorded"; + recorded = false; + } + if (FAILED(g_vk.interop->EndVkCommandBufferInterop(list)) && recorded) + { + failure = "EndVkCommandBufferInterop failed"; + recorded = false; + } + } + for (const Source& source : sources) + transition(list, source.resource, source.copyState, source.state); + if (!recorded) + return refuse(failure); + feature.engineJitter[0] = p.jitterX; + feature.engineJitter[1] = p.jitterY; + if (frame <= 3 || frame % 600 == 0) + logf("engine frame %u recorded, preset M (jitter %.6f,%.6f, reset %d, mvScale %.3f,%.3f)", frame, p.jitterX, p.jitterY, + game.reset, p.mvScaleX, p.mvScaleY); + return true; +} + +// Records the whole evaluation of `frame` into `list`. False: nothing usable was recorded and the +// CUDA backend takes this feature over. +static bool engine_evaluate(Feature& feature, ID3D12GraphicsCommandList* list, ID3D12Resource* color, + ID3D12Resource* depth, ID3D12Resource* motion, ID3D12Resource* output, + ID3D12Resource* exposure, const FrameParams& p, uint32_t frame) +{ + if (feature.engineRefused) + return false; + auto refuse = [&](const char* reason) { + feature.engineRefused = true; + logf("engine backend: %s; feature %u uses the CUDA backend from frame %u", reason, feature.handle.Id, frame); + if (feature.engine != nullptr) + logf("engine backend: the temporal history restarts in the CUDA backend"); + return false; + }; + std::call_once(g_engine.vulkanOnce, init_engine_vulkan); + if (feature.preset == 13) + return engine_evaluate_m(feature, list, color, depth, motion, output, p, frame); + // Without the HDR flag NGX runs its LDR input and output kernels: a model directory of its own. + const bool ldr = (feature.flags & 0x1) == 0; + if (ldr) + std::call_once(g_engine.ldrOnce, init_engine_backend_ldr); + else + std::call_once(g_engine.once, init_engine_backend); + d4r::Model* const model = ldr ? g_engine.modelLdr.get() : g_engine.model.get(); + if (model == nullptr) + return refuse(ldr ? "no LDR model (engine\\k-ldr) for a game without the HDR flag" : "not available"); + // NVSDK_NGX_DLSS_Feature_Flags: IsHDR 1, MVLowRes 2, MVJittered 4, DepthInverted 8. + if (feature.preset != 11) + return refuse("only preset K is implemented"); + // MVLowRes (0x2) clear: vectors are at output resolution (engine revision D4RO0003). + const bool displayMotion = (feature.flags & 0x2) == 0; + if (displayMotion && !model->hasDisplayMotion()) + return refuse("display-resolution motion needs a model compiled by the current compile_k.py"); + if ((feature.flags & 0x8) == 0 && !model->hasRegularDepth()) + return refuse("regular (non-inverted) depth needs a model compiled by the current compile_k.py"); + // DLSS.Indicator.Invert.X/Y.Axis (p.invertX/Y) only place NGX's debug indicator: NGX's K and M outputs are + // byte-identical with them set (harness D4R_HARNESS_INDICATOR_INVERT=1), so neither engine path reads them. + // Two exposure sources. AutoExposure (0x40) measures the frame's luma scaled by 1 / pre-exposure and multiplies + // the result by the exposure scale. Without it the game's ExposureTexture is the exposure: texel x scale / + // pre-exposure (read directly, not measured). + const bool autoExposure = (feature.flags & 0x40) != 0; + const bool gameExposure = !ldr && !autoExposure; + if (!ldr && autoExposure && !model->hasAutoExposure()) + return refuse("automatic exposure needs a model with exposure shaders"); + if (gameExposure && (!p.hasExposure || !model->hasGameExposure())) + return refuse("exposure without AutoExposure needs an exposure texture and a model compiled by the current compile_k.py"); + if (!ldr && !(std::isfinite(p.exposureScale) && p.exposureScale > 0.0f)) + return refuse("the exposure scale is not a positive number"); + if (!ldr && !(std::isfinite(p.preExposure) && p.preExposure > 0.0f)) + return refuse("the pre-exposure is not a positive number"); + + if (output == color || output == motion || output == depth) + return refuse("input/output texture aliasing is not supported"); + + d4r::GameFrame game; + using Source = EngineFrameSource; + const auto inputState = static_cast(env_uint("D4R_SHIM_INPUT_STATE", 0x40)); + const auto depthState = static_cast(env_uint("D4R_SHIM_DEPTH_STATE", 0x40)); + const auto outputState = static_cast(env_uint("D4R_SHIM_OUTPUT_STATE", 0x8)); + Source sources[5] = { + {color, &game.color, p.colorBaseX, p.colorBaseY, p.renderWidth, p.renderHeight, inputState, + D3D12_RESOURCE_STATE_COPY_SOURCE}, + {depth, &game.depth, p.depthBaseX, p.depthBaseY, p.renderWidth, p.renderHeight, depthState, + D3D12_RESOURCE_STATE_COPY_SOURCE}, + {motion, &game.motion, p.mvBaseX, p.mvBaseY, displayMotion ? feature.outWidth : p.renderWidth, + displayMotion ? feature.outHeight : p.renderHeight, inputState, D3D12_RESOURCE_STATE_COPY_SOURCE}, + {output, &game.output, p.outputBaseX, p.outputBaseY, feature.outWidth, feature.outHeight, outputState, + D3D12_RESOURCE_STATE_COPY_DEST}, + {exposure, &game.exposure, 0, 0, 1, 1, inputState, D3D12_RESOURCE_STATE_COPY_SOURCE}, + }; + const size_t sourceCount = gameExposure ? 5 : 4; + for (size_t index = 0; index < sourceCount; ++index) + { + Source& source = sources[index]; + UINT64 handle = 0, offset = 0; + D3D12_RESOURCE_DESC desc; + source.resource->GetDesc(&desc); + if (FAILED(g_vk.interop->GetVulkanResourceInfo1(source.resource, &handle, &offset, &source.image->format)) || + handle == 0 || desc.Dimension != D3D12_RESOURCE_DIMENSION_TEXTURE2D || desc.SampleDesc.Count != 1 || + source.x + source.width > desc.Width || source.y + source.height > desc.Height) + return refuse("a texture is not a single-sample 2D image covering its rectangle"); + source.image->image = reinterpret_cast(handle); + source.image->x = source.x; + source.image->y = source.y; + const bool isOutput = source.image == &game.output; + prepare_engine_view(feature, source.resource, desc, *source.image, isOutput, + (source.x == 0 && source.y == 0) || model->hasDirectOrigins(), source.copyState); + g_vk.interop->GetVulkanImageLayout(source.resource, source.copyState, &source.image->layout); + } + if (gameExposure && (game.exposure.view == VK_NULL_HANDLE || (game.exposure.format != VK_FORMAT_R32_SFLOAT && + game.exposure.format != VK_FORMAT_R16_SFLOAT && + game.exposure.format != VK_FORMAT_R32G32B32A32_SFLOAT && + game.exposure.format != VK_FORMAT_R16G16B16A16_SFLOAT))) + return refuse("the exposure texture must be a float texture the engine samples directly"); + if (game.depth.format == VK_FORMAT_D32_SFLOAT || game.depth.format == VK_FORMAT_D32_SFLOAT_S8_UINT) + game.depth.aspect = VK_IMAGE_ASPECT_DEPTH_BIT; + if (!d4r::gameColorFormat(game.color.format) || !d4r::gameMotionFormat(game.motion.format) || + !d4r::gameDepthFormat(game.depth.format, game.depth.aspect) || !d4r::gameOutputFormat(game.output.format)) + { + logf("engine backend: Vulkan formats colour %d, depth %d, motion %d, output %d", game.color.format, + game.depth.format, game.motion.format, game.output.format); + return refuse("a texture format is not implemented"); + } + if (frame <= 2) + logf("engine direct textures: color %d, motion %d, depth %d, output %d (format %d, RGB10A2 pipeline %d)", + game.color.view != VK_NULL_HANDLE, game.motion.view != VK_NULL_HANDLE, + game.depth.view != VK_NULL_HANDLE, game.output.view != VK_NULL_HANDLE, + game.output.format, model->hasRgb10Output()); + game.settings.renderWidth = p.renderWidth; + game.settings.renderHeight = p.renderHeight; + game.settings.jitter[0] = p.jitterX; + game.settings.jitter[1] = p.jitterY; + game.settings.displayMotion = displayMotion; + // The engine takes motion in output pixels; MV.Scale gives render pixels (display-resolution vectors are output pixels already). + const float toOutput = displayMotion ? 1.0f : static_cast(feature.outWidth) / static_cast(p.renderWidth); + const float toOutputY = displayMotion ? 1.0f : static_cast(feature.outHeight) / static_cast(p.renderHeight); + game.settings.motionScale[0] = p.mvScaleX * toOutput; + game.settings.motionScale[1] = p.mvScaleY * toOutputY; + // MVJittered: NGX adds the change of the jitter (render pixels) to the vectors before scaling them + // (observed with MV.Scale 1: offset = jitter - previous jitter, the first frame's previous being 0). + // The offset is added to the raw vector, before motionScale's own output-pixel conversion, so it carries + // no extra output/render factor in either motion mode. + if ((feature.flags & 0x4) != 0 && p.mvScaleX != 0.0f && p.mvScaleY != 0.0f) + { + game.settings.motionOffset[0] = (p.jitterX - feature.engineJitter[0]) / p.mvScaleX; + game.settings.motionOffset[1] = (p.jitterY - feature.engineJitter[1]) / p.mvScaleY; + } + game.settings.reset = p.reset != 0 || feature.engine == nullptr; + game.settings.depthInverted = (feature.flags & 0x8) != 0; + game.settings.autoExposure = !ldr && autoExposure; // LDR: no exposure at all + game.settings.exposureScale = ldr ? 1.0f : p.exposureScale; + // NGX's automatic K measurement scales the luma by 1 / pre-exposure before the logarithm + // (cuda_luma_convert_kernel), so the stored exposure is pre times the raw measurement and the stage's + // 1 / pre factor recovers it; the game texture path applies the same stage factor. + game.settings.preExposure = ldr ? 1.0f : p.preExposure; + try + { + if (feature.engine == nullptr) + { + const size_t tokenBytes = static_cast(d4r::Engine::networkTokenBytes(feature.outWidth, feature.outHeight)); + const size_t headBytes = static_cast(d4r::Engine::networkHeadBytes(feature.outWidth, feature.outHeight)); + const bool hip = create_engine_network(feature, ldr ? g_engine.directoryLdr : g_engine.directory, "K", true, + feature.outWidth, feature.outHeight, tokenBytes, headBytes); + const d4r::ExternalNetwork external = {feature.engineTokens.buffer, feature.engineHead.buffer}; + feature.engine = std::make_unique(*model, g_vk.physical, g_vk.device, + feature.outWidth, feature.outHeight, + std::max(feature.width, p.renderWidth), + std::max(feature.height, p.renderHeight), + hip ? &external : nullptr); + logf("engine backend: feature %u runs preset K in the game's command list (%ux%u -> %ux%u), network on %s", + feature.handle.Id, p.renderWidth, p.renderHeight, feature.outWidth, feature.outHeight, + hip ? "the native HIP layers (list split around it)" : "Vulkan"); + } + } + catch (const std::exception& error) + { + logf("engine backend: %s", error.what()); + release_engine_network(feature); + return refuse("the engine could not be created"); + } + if (feature.engineNet != nullptr) + { + // The native HIP network: the game's list is split around it. Its native completion callback releases + // the second part, which records the final store (or the conversion blit for incompatible outputs). + auto recordPart = [&](bool front) -> const char* { + VkCommandBuffer command = VK_NULL_HANDLE; + if (FAILED(g_vk.interop->BeginVkCommandBufferInterop(list, &command))) + return "BeginVkCommandBufferInterop failed"; + const char* failure = nullptr; + try + { + if (front) + { + feature.engine->recordFront(command, game); + // The HIP network's stream waits for this frame's number (record_engine_input_ready): the + // fill is part of the first half, so the network starts the moment the front's GPU work ends. + if (feature.engineGpuWait.load()) + record_engine_input_ready(feature, command, frame); + } + else + feature.engine->recordBack(command, game); + } + catch (const std::exception& error) + { + logf("engine backend (K): %s", error.what()); + failure = "the frame could not be recorded"; + } + if (FAILED(g_vk.interop->EndVkCommandBufferInterop(list)) && failure == nullptr) + failure = "EndVkCommandBufferInterop failed"; + return failure; + }; + auto abandon = [&](const char* reason) { + feature.engineRefused = true; + logf("engine backend (K): %s; feature %u uses the CUDA backend after frame %u", reason, feature.handle.Id, frame); + return true; + }; + return record_engine_network_frame(feature, list, sources, frame, "K", game.settings.reset, + game.settings.motionScale[0], game.settings.motionScale[1], p.jitterX, + p.jitterY, recordPart, refuse, abandon); + } + if (FAILED(g_vk.lifetime->RetainExternalResources(list, feature.resources))) + return refuse("the command list cannot retain the engine's resources"); + for (const Source& source : sources) + transition(list, source.resource, source.state, source.copyState); + VkCommandBuffer command = VK_NULL_HANDLE; + bool recorded = SUCCEEDED(g_vk.interop->BeginVkCommandBufferInterop(list, &command)); + const char* failure = recorded ? nullptr : "BeginVkCommandBufferInterop failed"; + if (recorded) + { + try + { + feature.engine->record(command, game); + } + catch (const std::exception& error) + { + logf("engine backend: %s", error.what()); + failure = "the frame could not be recorded"; + recorded = false; + } + if (FAILED(g_vk.interop->EndVkCommandBufferInterop(list)) && recorded) + { + failure = "EndVkCommandBufferInterop failed"; + recorded = false; + } + } + for (const Source& source : sources) + transition(list, source.resource, source.copyState, source.state); + if (!recorded) + return refuse(failure); + feature.engineJitter[0] = p.jitterX; + feature.engineJitter[1] = p.jitterY; + if (frame <= 3 || frame % 600 == 0) + logf("engine frame %u recorded (jitter %.6f,%.6f, reset %d, mvScale %.3f,%.3f, render %ux%u)", frame, + p.jitterX, p.jitterY, game.settings.reset, p.mvScaleX, p.mvScaleY, p.renderWidth, p.renderHeight); + return true; +} + static NgxResult evaluate(ID3D12GraphicsCommandList* list, const NgxHandle* handle, void* parameters) { std::lock_guard apiLock(g_featureApiMutex); @@ -5058,6 +6252,18 @@ static NgxResult evaluate(ID3D12GraphicsCommandList* list, const NgxHandle* hand p.outputBaseY = get_uint_or(parameters, "DLSS.Output.Subrect.Base.Y", 0); p.hasExposure = exposure != nullptr; + if (engine_requested() && engine_evaluate(*feature, list, color, depth, motion, output, exposure, p, frame)) + return NGX_SUCCESS; + + if (feature->cudaCreationDeferred) + { + logf("engine backend: creating CUDA fallback for feature %u at frame %u", feature->handle.Id, frame); + const NgxResult created = create_cuda_feature(*feature); + if (created != NGX_SUCCESS) return created; + p.reset = 1; // the freshly created CUDA feature has none of the Vulkan engine's history + logf("engine backend: CUDA fallback initialized; temporal history reset"); + } + const int slotIndex = static_cast(frame % kSlots); InputSlot& slot = feature->inputs[slotIndex]; // Throttle the game to DLSS speed when the pipeline still owns this slot @@ -5102,7 +6308,7 @@ static NgxResult evaluate(ID3D12GraphicsCommandList* list, const NgxHandle* hand feature->vramDecided = true; feature->vram = p.vram; feature->split = p.vram && env_uint("D4R_SHIM_SPLIT_FRAME", 0) != 0 && g_vk.split != nullptr && - create_split_semaphore(*feature); + (feature->splitSemaphore != VK_NULL_HANDLE || create_split_semaphore(*feature)); feature->linearInputs = feature->split && env_uint("D4R_SHIM_LINEAR_INPUTS", 0) != 0 && g.cu.registerLinearTexture != nullptr; if (feature->linearInputs) @@ -5372,6 +6578,14 @@ D4R_EXPORT NgxResult NVSDK_NGX_D3D12_EvaluateFeature_C(ID3D12GraphicsCommandList static void release_feature(Feature* feature) { + // A queued GPU-side input wait for a frame that never ran would stall the network's stream: release it + // before the drain below waits for the worker (the worker also releases waits it sees stuck). + if (feature->engineNet != nullptr) + { + const uint32_t waiting = feature->engineWaitFrame.load(); + if (waiting != 0) + release_engine_input_wait(*feature, waiting, "the feature is being released", false); + } // Wait for outstanding evaluations, then free CUDA resources on the worker. drain_pipeline(); g.worker.call([feature] { @@ -5429,7 +6643,10 @@ static void release_feature(Feature* feature) if (feature->resources != nullptr) feature->resources->Release(); else + { + release_engine_network(*feature); delete feature; + } } static void destroy_recorded_resources(Feature* feature) @@ -5439,8 +6656,12 @@ static void destroy_recorded_resources(Feature* feature) staging->release(); for (OutputSlot& slot : feature->outputs) slot.staging.release(); + release_engine_network(*feature); + const bool engineSemaphore = !feature->vramDecided && feature->splitSemaphore != VK_NULL_HANDLE; if (g_vk.ready && feature->vramDecided) release_vram(*feature); + else if (engineSemaphore) // created for the engine's HIP network; release_vram destroys it otherwise + g_vk.destroySemaphore(g_vk.device, feature->splitSemaphore, nullptr); if (feature->marker != nullptr) feature->marker->Release(); logf("feature %u recorded resources retired after allocator completion", feature->handle.Id); diff --git a/tools/wine_nvcuda_bridge.c b/tools/wine_nvcuda_bridge.c index 71dd9c0..18f1bae 100644 --- a/tools/wine_nvcuda_bridge.c +++ b/tools/wine_nvcuda_bridge.c @@ -58,7 +58,9 @@ enum CUDA_ERROR_OUT_OF_MEMORY = 2, CUDA_ERROR_NOT_FOUND = 500, CUDA_ERROR_NOT_READY = 600, + CUDA_ERROR_LAUNCH_FAILED = 719, CUDA_ERROR_NOT_SUPPORTED = 801, + CUDA_ERROR_UNKNOWN = 999, }; static void* cuda_library; @@ -1627,7 +1629,7 @@ static size_t array_element_bytes(int32_t format, uint32_t channels) case 0x01: case 0x08: return 1u * channels; case 0x02: case 0x09: case 0x10: return 2u * channels; case 0x03: case 0x0a: case 0x20: return 4u * channels; - case 80: return 16; /* packed 10:10:10:2 backed by RGBA32F */ + case 80: return 4; /* CUDA copies packed 10:10:10:2 even though ZLUDA's backing is RGBA32F */ case 192: case 198: return 1; case 193: case 199: return 2; case 194: case 200: return 4; @@ -3174,6 +3176,956 @@ CUresult WINAPI d4rImportVulkanMemory(void* client_device, uint64_t client_memor return CUDA_SUCCESS; } +/* d4r engine: the Swin network on d4r's native HIP layer kernels, launched directly through HIP (no NGX, no + PTX; engine/hip_net.cpp is the same logic for native Linux programs). The engine's Vulkan front writes the + tokens into `tokens`, the network's layers run here, and its back reads `output`; both are device pointers of + memory imported with d4rImportVulkanMemory. `net` is the model's hipnet.bin: + D4RMHIP1 (compile_m.py --hip): preset M, ten layers, width/height are the render dimensions; two immutable + window-alignment graphs, one per evaluation. + D4RKHIP1 (compile_k.py --hip): preset K, eleven layers (the bottleneck dec5 as its main kernel plus + post1..3), width/height are the output dimensions (the token grid is makeKShape), and `tokens`/`output` + hold 16 / 40 f16 channels per token row-major. Eight immutable graphs, one per makeKWindows evaluation. + Either way: {kernel name[32], shift x, y, grid divisor, pad, weight offset, weight bytes} per layer, then the + weight allocations as NGX uploads them. */ +typedef struct +{ + const char* name; /* kernel name, as in hipnet.bin */ + const void* data; /* its code object (.hsaco) */ + uint64_t bytes; +} D4rEngineNetCode; + +typedef struct +{ + void* w; + int32_t sx, sy, tw, th; + void* in; + void* skip; + void* out; + void* merged; + int32_t r69, r70, r71, r72; + void* in_flags; + void* out_flags; +} D4rEngineNetParams; /* kernels/m/swin_common.h: CommonParams, TubeParams */ + +typedef struct +{ + int32_t W, H; /* token grid of the layer, pixels */ + uint8_t* in; + uint8_t* skip; /* decoders: the full-resolution encoder output */ + uint8_t* out24; /* encoders: patch-merged output; decoders: output; dec0: the caller's head buffer */ + uint8_t* out32; /* encoders: full-resolution output (the skip of the matching decoder) */ + uint64_t pad40, pad48; + int32_t sx, sy; + uint8_t* w; + uint8_t rest[104]; +} D4rEngineNetPwinParams; /* kernels/k/pwin_layer.h: PwinParams */ +_Static_assert(sizeof(D4rEngineNetPwinParams) == 176, "K param block"); + +/* one immutable launch of one graph: entry point, grid, block and the frozen parameter block (the active + member follows the model's magic) */ +typedef union +{ + D4rEngineNetParams m; + D4rEngineNetPwinParams k; +} D4rEngineNetParamsU; + +typedef struct +{ + void* function; + uint32_t gx, gy, bx, by, bz; + D4rEngineNetParamsU params; +} D4rEngineNetNode; + +#define D4R_ENGINE_NET_M_LAYERS 10 +#define D4R_ENGINE_NET_K_LAYERS 11 +#define D4R_ENGINE_NET_LAYERS D4R_ENGINE_NET_K_LAYERS +#define D4R_ENGINE_NET_MODULES D4R_ENGINE_NET_K_LAYERS +#define D4R_ENGINE_NET_GRAPHS 8 +#define D4R_ENGINE_NET_NODES 16 /* 11 layers, the deepest of them with its three post phases */ +#define D4R_ENGINE_NET_OWNED 32 +#define D4R_ENGINE_NET_SIGNALS 16 +struct D4rEngineNet; +typedef struct +{ + struct D4rEngineNet* net; + uint64_t value; + int busy; +} D4rEngineNetSignal; + +typedef struct +{ + int sType; /* VK_STRUCTURE_TYPE_SEMAPHORE_SIGNAL_INFO = 1000207005 */ + const void* pNext; + uint64_t semaphore, value; +} D4rVkSemaphoreSignalInfo; +typedef int (*VK_SIGNAL_SEMAPHORE_FN)(void*, const D4rVkSemaphoreSignalInfo*); +typedef int (*VK_GET_SEMAPHORE_COUNTER_FN)(void*, uint64_t, uint64_t*); +typedef int (*HIP_LAUNCH_HOST_FN)(void*, void (*)(void*), void*); + +typedef struct D4rEngineNet +{ + int k_preset; /* D4RKHIP1: the token grid comes from the output dimensions and the cycle has eight graphs */ + uint32_t graphs; /* 2 for M, 8 for K */ + void* modules[D4R_ENGINE_NET_MODULES]; + void* owned[D4R_ENGINE_NET_OWNED]; + int module_count, owned_count; + uint32_t node_count[D4R_ENGINE_NET_GRAPHS]; + D4rEngineNetNode nodes[D4R_ENGINE_NET_GRAPHS][D4R_ENGINE_NET_NODES]; + /* weight-image preparations: one per layer, i.e. per distinct weights pointer, run once at create time */ + uint32_t prep_count; + D4rEngineNetNode preps[D4R_ENGINE_NET_LAYERS]; + void* stream; + void* graph[D4R_ENGINE_NET_GRAPHS]; + void* executable[D4R_ENGINE_NET_GRAPHS]; + uint32_t frame; + uint32_t* input_wait_host; + void* start; + void* stop; + /* Input readiness on the GPU (d4rEngineNetInputWaitReady / d4rEngineNetLaunchAt): the u32 the engine's + front fills with the frame number, and the value a queued stream wait is waiting for (0: none). */ + CUdeviceptr input_wait; + uint32_t input_wait_value; + void* timeline_device; + uint64_t timeline_semaphore, timeline_value; + pthread_mutex_t timeline_lock; + VK_SIGNAL_SEMAPHORE_FN signal_timeline; + HIP_LAUNCH_HOST_FN launch_host; + D4rEngineNetSignal signals[D4R_ENGINE_NET_SIGNALS]; + int timeline_error; +} D4rEngineNet; + +typedef int (*HIP_PP_FN)(void**); +typedef int (*HIP_P_FN)(void*); +typedef int (*HIP_MALLOC_FN)(void**, size_t); +typedef int (*HIP_MEMCPY_FN)(void*, const void*, size_t); +typedef int (*HIP_MODULE_LOAD_DATA_FN)(void**, const void*); +typedef int (*HIP_MODULE_GET_FUNCTION_FN)(void**, void*, const char*); +typedef int (*HIP_MODULE_GET_GLOBAL_FN)(void**, size_t*, void*, const char*); +typedef int (*HIP_MODULE_LAUNCH_FN)(void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, void*, + void**, void**); +typedef int (*HIP_EVENT_RECORD_FN)(void*, void*); +typedef int (*HIP_EVENT_ELAPSED_FN)(float*, void*, void*); +typedef int (*HIP_STREAM_WAIT_VALUE32_FN)(void*, void*, uint32_t, unsigned int, uint32_t); +typedef int (*HIP_STREAM_WRITE_VALUE32_FN)(void*, void*, uint32_t, unsigned int); +typedef int (*HIP_STREAM_CREATE_WITH_FLAGS_FN)(void**, unsigned int); +typedef int (*HIP_STREAM_SYNCHRONIZE_FN)(void*); +typedef int (*HIP_CAPTURE_BEGIN_FN)(void*, int); +typedef int (*HIP_CAPTURE_END_FN)(void*, void**); +typedef int (*HIP_GRAPH_INSTANTIATE_FN)(void**, void*, void**, char*, size_t); + + +CUresult WINAPI d4rEngineNetReleaseInput(void* handle, uint32_t value) +{ + D4rEngineNet* net = (D4rEngineNet*)handle; + if (net == NULL || net->input_wait_host == NULL) + return CUDA_ERROR_INVALID_VALUE; + /* Never rewind a newer front's publication when releasing a discarded older frame. */ + uint32_t current = __atomic_load_n(net->input_wait_host, __ATOMIC_ACQUIRE); + while (current < value && + !__atomic_compare_exchange_n(net->input_wait_host, ¤t, value, 0, __ATOMIC_RELEASE, __ATOMIC_RELAXED)) + {} + return CUDA_SUCCESS; +} + +/* preset K's eleven layers in network order (the same plan as engine/hip_net.cpp): token-grid divisor, input + channels, the core channels (the encoder's full-resolution output width), the channels written to out24 (the + patch-merged skip for encoders, the block output for the decoders, the 40-channel head for dec0) and the + window-schedule role. */ +typedef struct +{ + const char* name; + uint32_t divisor, cin, c, out24ch; + int role; +} D4rEngineNetKLayer; + +static const D4rEngineNetKLayer d4r_engine_net_k_plan[D4R_ENGINE_NET_K_LAYERS] = { + {"enc0", 1, 16, 32, 64, -1}, {"enc1", 2, 64, 64, 64, -1}, {"enc2", 4, 64, 64, 96, -1}, + {"enc3", 8, 96, 96, 128, 2}, {"enc4", 16, 128, 128, 160, 1}, {"dec5", 32, 160, 160, 160, 0}, + {"dec4", 16, 160, 128, 128, 1}, {"dec3", 8, 128, 96, 96, 2}, {"dec2", 4, 96, 64, 64, -1}, + {"dec1", 2, 64, 64, 64, -1}, {"dec0", 1, 64, 32, 40, -1}, +}; + +/* makeKWindows (engine/runtime.cpp): dec5 moves once per frame around an eight-frame cycle and every shallower + layer follows the one below it as (2 * shift + 4) mod 8; role 0 = dec5, 1 = enc4/dec4, 2 = enc3/dec3, + everything else keeps shift 4. */ +static const int32_t d4r_engine_net_k_cycle[8][2] = {{0, 2}, {5, 6}, {4, 0}, {1, 4}, {7, 5}, {2, 1}, {3, 7}, {6, 3}}; + +static void d4r_engine_net_k_shift(int role, uint32_t frame, int32_t* sx, int32_t* sy) +{ + if (role < 0) + { + *sx = 4; + *sy = 4; + return; + } + int32_t x = d4r_engine_net_k_cycle[frame % 8][0], y = d4r_engine_net_k_cycle[frame % 8][1]; + for (int depth = 0; depth < role; ++depth) + { + x = (2 * x + 4) & 7; + y = (2 * y + 4) & 7; + } + *sx = x; + *sy = y; +} + +/* one code object of a layer: the main entry, its weight-image preparation and the deep layers' post phases */ +typedef struct +{ + char name[33]; + void* module; + void* main; + void* prep; + void* post[8]; + uint32_t prep_blocks, block_z, grid_x, grid_y; + uint32_t post_gx[8], post_gy[8], post_bx[8], post_by[8], post_bz[8]; + int post_count; +} D4rEngineNetKernel; + +static uint32_t d4r_engine_net_global(void* module, const char* name, uint32_t fallback) +{ + HIP_MODULE_GET_GLOBAL_FN get = (HIP_MODULE_GET_GLOBAL_FN)hip_symbol("hipModuleGetGlobal"); + HIP_MEMCPY_FN to_host = (HIP_MEMCPY_FN)hip_symbol("hipMemcpyDtoH"); + void* at = NULL; + size_t bytes = 0; + uint32_t value = fallback; + if (get != NULL && to_host != NULL && get(&at, &bytes, module, name) == 0 && bytes == 4) + to_host(&value, at, 4); + return value; +} + +void WINAPI d4rEngineNetDestroy(void* handle) +{ + D4rEngineNet* net = (D4rEngineNet*)handle; + if (net == NULL) + return; + HIP_P_FN synchronize = (HIP_P_FN)hip_symbol("hipStreamSynchronize"); + HIP_P_FN free_device = (HIP_P_FN)hip_symbol("hipFree"); + HIP_P_FN unload = (HIP_P_FN)hip_symbol("hipModuleUnload"); + HIP_P_FN event_destroy = (HIP_P_FN)hip_symbol("hipEventDestroy"); + HIP_P_FN stream_destroy = (HIP_P_FN)hip_symbol("hipStreamDestroy"); + /* A queued input wait whose frame never reached the GPU would make the synchronize below wait forever: + satisfy it (the shim also releases waits it knows about, before it destroys the network). */ + if (net->input_wait != 0 && net->input_wait_value != 0) + d4rEngineNetReleaseInput(net, net->input_wait_value); + if (net->stream != NULL && synchronize != NULL) + synchronize(net->stream); + HIP_P_FN exec_destroy = (HIP_P_FN)hip_symbol("hipGraphExecDestroy"); + HIP_P_FN graph_destroy = (HIP_P_FN)hip_symbol("hipGraphDestroy"); + for (unsigned phase = 0; phase < D4R_ENGINE_NET_GRAPHS; ++phase) + { + if (net->executable[phase] != NULL && exec_destroy != NULL) + exec_destroy(net->executable[phase]); + if (net->graph[phase] != NULL && graph_destroy != NULL) + graph_destroy(net->graph[phase]); + } + for (int i = 0; i < net->owned_count && free_device != NULL; ++i) + free_device(net->owned[i]); + for (int i = 0; i < net->module_count && unload != NULL; ++i) + unload(net->modules[i]); + if (net->start != NULL && event_destroy != NULL) + event_destroy(net->start); + if (net->stop != NULL && event_destroy != NULL) + event_destroy(net->stop); + if (net->stream != NULL && stream_destroy != NULL) + stream_destroy(net->stream); + if (net->timeline_device != NULL) + pthread_mutex_destroy(&net->timeline_lock); + free(net); +} + +static CUresult d4r_engine_net_capture(D4rEngineNet* net, HIP_MODULE_LAUNCH_FN launch) +{ + HIP_CAPTURE_BEGIN_FN begin = (HIP_CAPTURE_BEGIN_FN)hip_symbol("hipStreamBeginCapture"); + HIP_CAPTURE_END_FN end = (HIP_CAPTURE_END_FN)hip_symbol("hipStreamEndCapture"); + HIP_GRAPH_INSTANTIATE_FN instantiate = (HIP_GRAPH_INSTANTIATE_FN)hip_symbol("hipGraphInstantiate"); + if (begin == NULL || end == NULL || instantiate == NULL) + return CUDA_ERROR_NOT_SUPPORTED; + /* Every frame of the alignment cycle has its own immutable graph (arguments and grids frozen, the K + bottleneck's post phases included); no per-frame node updates. */ + for (uint32_t phase = 0; phase < net->graphs; ++phase) + { + if (begin(net->stream, 2 /* hipStreamCaptureModeRelaxed */) != 0) + return CUDA_ERROR_UNKNOWN; + int queued = 0; + for (uint32_t i = 0; i < net->node_count[phase] && queued == 0; ++i) + { + D4rEngineNetNode* node = &net->nodes[phase][i]; + void* args[] = {&node->params}; + queued = launch(node->function, node->gx, node->gy, 1, node->bx, node->by, node->bz, 0, net->stream, args, + NULL); + } + const int captured = end(net->stream, &net->graph[phase]); + if (queued != 0 || captured != 0) + return CUDA_ERROR_LAUNCH_FAILED; + if (instantiate(&net->executable[phase], net->graph[phase], NULL, NULL, 0) != 0) + return CUDA_ERROR_UNKNOWN; + } + return CUDA_SUCCESS; +} + +CUresult WINAPI d4rEngineNetCreate(const uint8_t* file, uint64_t file_bytes, const D4rEngineNetCode* codes, + uint32_t code_count, uint32_t width, uint32_t height, CUdeviceptr tokens, + CUdeviceptr output, void** handle) +{ + ensure_context(); + if (file == NULL || codes == NULL || handle == NULL || tokens == 0 || output == 0) + return CUDA_ERROR_INVALID_VALUE; + if (context_setup_result != CUDA_SUCCESS) + return context_setup_result; + typedef int (*STREAM_PRIORITY_RANGE_FN)(int*, int*); + typedef int (*STREAM_CREATE_PRIORITY_FN)(void**, unsigned int, int); + STREAM_PRIORITY_RANGE_FN priority_range = (STREAM_PRIORITY_RANGE_FN)hip_symbol("hipDeviceGetStreamPriorityRange"); + STREAM_CREATE_PRIORITY_FN stream_create = (STREAM_CREATE_PRIORITY_FN)hip_symbol("hipStreamCreateWithPriority"); + HIP_PP_FN event_create = (HIP_PP_FN)hip_symbol("hipEventCreate"); + HIP_P_FN synchronize = (HIP_P_FN)hip_symbol("hipStreamSynchronize"); + HIP_MALLOC_FN allocate = (HIP_MALLOC_FN)hip_symbol("hipMalloc"); + HIP_MEMCPY_FN to_device = (HIP_MEMCPY_FN)hip_symbol("hipMemcpyHtoD"); + HIP_MODULE_LOAD_DATA_FN load = (HIP_MODULE_LOAD_DATA_FN)hip_symbol("hipModuleLoadData"); + HIP_MODULE_GET_FUNCTION_FN function = (HIP_MODULE_GET_FUNCTION_FN)hip_symbol("hipModuleGetFunction"); + HIP_MODULE_LAUNCH_FN launch = (HIP_MODULE_LAUNCH_FN)hip_symbol("hipModuleLaunchKernel"); + if (priority_range == NULL || stream_create == NULL || event_create == NULL || synchronize == NULL || allocate == NULL || to_device == NULL || + load == NULL || function == NULL || launch == NULL || hip_symbol("hipEventRecord") == NULL || + hip_symbol("hipEventQuery") == NULL) + return CUDA_ERROR_NOT_SUPPORTED; + if (file_bytes < 16) + return CUDA_ERROR_INVALID_VALUE; + const int k_preset = memcmp(file, "D4RKHIP1", 8) == 0; + if (!k_preset && memcmp(file, "D4RMHIP1", 8) != 0) + return CUDA_ERROR_INVALID_VALUE; + uint32_t count = 0; + memcpy(&count, file + 8, 4); + const uint32_t layers = k_preset ? D4R_ENGINE_NET_K_LAYERS : D4R_ENGINE_NET_M_LAYERS; + if (count != layers || file_bytes < 16 + 64ull * count) + return CUDA_ERROR_INVALID_VALUE; + int least_priority = 0, greatest_priority = 0; + if (priority_range(&least_priority, &greatest_priority) != 0) + return CUDA_ERROR_UNKNOWN; + D4rEngineNet* net = (D4rEngineNet*)calloc(1, sizeof(*net)); + if (net == NULL) + return CUDA_ERROR_OUT_OF_MEMORY; + net->k_preset = k_preset; + net->graphs = k_preset ? D4R_ENGINE_NET_GRAPHS : 2; + const char* failure = NULL; + /* The layer-0 token grid: M's render dimensions, K's output dimensions (makeKShape). */ + uint32_t tw, th; + if (k_preset) + { + if (width < 256 || height < 256 || width > 8192 || height > 8192) + failure = "output dimensions"; + tw = ((width + 3) / 4 + 31) / 32 * 32; + th = ((height + 3) / 4 + 31) / 32 * 32; + if (tw < 256) tw = 256; + if (th < 256) th = 256; + } + else + { + tw = (width + 31) / 32 * 16; + th = (height + 31) / 32 * 16; + } + /* This stream is on the frame's critical path, between two parts of the game's graphics list. */ + if (failure == NULL && (stream_create(&net->stream, 1 /* hipStreamNonBlocking */, greatest_priority) != 0 || + event_create(&net->start) != 0 || event_create(&net->stop) != 0)) + failure = "stream or events"; + void* const tok = (void*)(uintptr_t)tokens; + void* const result = (void*)(uintptr_t)output; + D4rEngineNetKernel kern[D4R_ENGINE_NET_MODULES]; + memset(kern, 0, sizeof(kern)); + int module_count = 0; + void* layer_in[D4R_ENGINE_NET_K_LAYERS] = {NULL}; + void* layer_skip[D4R_ENGINE_NET_K_LAYERS] = {NULL}; + void* layer_out[D4R_ENGINE_NET_K_LAYERS] = {NULL}; + void* layer_merged[D4R_ENGINE_NET_K_LAYERS] = {NULL}; + void* k_full[5] = {NULL}; + void* weights[D4R_ENGINE_NET_K_LAYERS] = {NULL}; + uint32_t divisor[D4R_ENGINE_NET_K_LAYERS] = {0}; + int32_t shift_sx[D4R_ENGINE_NET_K_LAYERS] = {0}, shift_sy[D4R_ENGINE_NET_K_LAYERS] = {0}; + uint32_t layer_kernel[D4R_ENGINE_NET_K_LAYERS] = {0}; + D4rEngineNetParamsU layer_params[D4R_ENGINE_NET_K_LAYERS]; + memset(layer_params, 0, sizeof(layer_params)); + if (k_preset && failure == NULL) + { + /* NGX packs the encoder skips, the merged outputs and the decoder outputs into one scratch allocation; + separate buffers only have to respect the channel counts and the wiring, not the original aliasing. */ + void* act[D4R_ENGINE_NET_K_LAYERS] = {NULL}; + for (uint32_t i = 0; i < layers && failure == NULL; ++i) + { + const D4rEngineNetKLayer* p = &d4r_engine_net_k_plan[i]; + const uint32_t w = tw / p->divisor, h = th / p->divisor; + /* an encoder's out24 is the patch-merged 2x2 output, half the resolution of its input */ + const uint32_t ow = i < 5 ? tw / (2 * p->divisor) : w, oh = i < 5 ? th / (2 * p->divisor) : h; + const int need = (i + 1 != layers ? 1 : 0) + (i < 5 ? 1 : 0); + if (net->owned_count + need > D4R_ENGINE_NET_OWNED) + { + failure = "activation buffers"; + break; + } + /* dec0 writes the head straight into the caller's output buffer */ + if (i + 1 != layers) + { + if (allocate(&act[i], (size_t)ow * oh * p->out24ch * 2 + 4096) != 0) + { + failure = "activation buffers"; + break; + } + net->owned[net->owned_count++] = act[i]; + } + if (i < 5) + { + if (allocate(&k_full[i], (size_t)w * h * p->c * 2 + 4096) != 0) + { + failure = "activation buffers"; + break; + } + net->owned[net->owned_count++] = k_full[i]; + } + layer_in[i] = i == 0 ? tok : act[i - 1]; + layer_skip[i] = i >= 6 ? k_full[10 - i] : NULL; + layer_out[i] = i + 1 == layers ? result : act[i]; + } + } + else if (failure == NULL) + { + /* skip1, merged1, skip2, merged2, tubeA, tubeB, dec2 */ + static const uint32_t scale[7] = {1, 2, 2, 4, 4, 4, 2}, channels[7] = {64, 96, 96, 128, 128, 128, 96}; + void* b[7] = {NULL}; + for (int i = 0; i < 7 && failure == NULL; ++i) + { + if (net->owned_count >= D4R_ENGINE_NET_OWNED || + allocate(&b[i], (size_t)(tw / scale[i]) * (th / scale[i]) * channels[i] + 4096) != 0) + failure = "activation buffers"; + else + net->owned[net->owned_count++] = b[i]; + } + /* tokens -> enc1 -> enc2 -> six tube blocks (ping-pong) -> dec2 -> dec1 */ + const int in_i[10] = {0, 1, 3, 4, 5, 4, 5, 4, 5, 6}; + const int skip_i[10] = {-1, -1, -1, -1, -1, -1, -1, -1, 2, 0}; + const int out_i[10] = {0, 2, 4, 5, 4, 5, 4, 5, 6, -1}; + const int merged_i[10] = {1, 3, -1, -1, -1, -1, -1, -1, -1, -1}; + for (int i = 0; i < 10; ++i) + { + layer_in[i] = in_i[i] < 0 ? NULL : b[in_i[i]]; + layer_skip[i] = skip_i[i] < 0 ? NULL : b[skip_i[i]]; + layer_out[i] = out_i[i] < 0 ? result : b[out_i[i]]; + layer_merged[i] = merged_i[i] < 0 ? NULL : b[merged_i[i]]; + } + layer_in[0] = tok; /* the engine's token buffer feeds the first layer, not a scratch buffer */ + } + for (uint32_t i = 0; i < count && failure == NULL; ++i) + { + const uint8_t* e = file + 16 + 64 * i; + char name[33] = {0}; + int32_t sx, sy; + uint32_t div; + uint64_t at, bytes; + memcpy(name, e, 32); + memcpy(&sx, e + 32, 4); + memcpy(&sy, e + 36, 4); + memcpy(&div, e + 40, 4); + memcpy(&at, e + 48, 8); + memcpy(&bytes, e + 56, 8); + if (div == 0 || at > file_bytes || bytes > file_bytes - at) + { + failure = "layer table"; + break; + } + if (k_preset) + { + char expect[48]; + snprintf(expect, sizeof(expect), "dltss_pwin_%s_layer", d4r_engine_net_k_plan[i].name); + if (strcmp(name, expect) != 0) + { + failure = "layer table"; + break; + } + } + int k = 0; + while (k < module_count && strcmp(kern[k].name, name) != 0) + ++k; + if (k == module_count) + { + const D4rEngineNetCode* code = NULL; + for (uint32_t c = 0; c < code_count; ++c) + if (codes[c].name != NULL && strcmp(codes[c].name, name) == 0) + code = &codes[c]; + if (code == NULL || k == D4R_ENGINE_NET_MODULES || load(&net->modules[k], code->data) != 0) + { + failure = "kernel code object"; + break; + } + ++net->module_count; + ++module_count; + strcpy(kern[k].name, name); + kern[k].module = net->modules[k]; + char prep_name[48]; + snprintf(prep_name, sizeof(prep_name), "%s_prep", name); + if (function(&kern[k].main, kern[k].module, name) != 0) + { + failure = "kernel entry point"; + break; + } + if (function(&kern[k].prep, kern[k].module, prep_name) != 0) + kern[k].prep = NULL; + kern[k].prep_blocks = d4r_engine_net_global(kern[k].module, "d4r_prep_blocks", 0); + kern[k].block_z = d4r_engine_net_global(kern[k].module, "d4r_block_z", 4); + kern[k].grid_x = d4r_engine_net_global(kern[k].module, "d4r_grid_x", 0); + kern[k].grid_y = d4r_engine_net_global(kern[k].module, "d4r_grid_y", 0); + if (kern[k].block_z == 0) + kern[k].block_z = 4; + /* one prepared weight image per weights pointer while a module's slots last: the tube needs six */ + if ((kern[k].prep != NULL && kern[k].prep_blocks == 0) || + (strstr(name, "tube") != NULL && d4r_engine_net_global(kern[k].module, "d4r_prep_key_slots", 0) < 6)) + { + failure = "weight image slots"; + break; + } + /* the deep layers' code objects also carry their post phases; their grids are compile-time */ + for (int j = 1; j <= 8; ++j) + { + char tag[48], sym[48]; + snprintf(tag, sizeof(tag), "%s_post%d", name, j); + if (function(&kern[k].post[kern[k].post_count], kern[k].module, tag) != 0) + break; + snprintf(sym, sizeof(sym), "d4r_post%d_grid_x", j); + kern[k].post_gx[kern[k].post_count] = d4r_engine_net_global(kern[k].module, sym, 0); + snprintf(sym, sizeof(sym), "d4r_post%d_grid_y", j); + kern[k].post_gy[kern[k].post_count] = d4r_engine_net_global(kern[k].module, sym, 0); + snprintf(sym, sizeof(sym), "d4r_post%d_block_x", j); + kern[k].post_bx[kern[k].post_count] = d4r_engine_net_global(kern[k].module, sym, 0); + snprintf(sym, sizeof(sym), "d4r_post%d_block_y", j); + kern[k].post_by[kern[k].post_count] = d4r_engine_net_global(kern[k].module, sym, 0); + snprintf(sym, sizeof(sym), "d4r_post%d_block_z", j); + kern[k].post_bz[kern[k].post_count] = d4r_engine_net_global(kern[k].module, sym, 0); + if (kern[k].post_gx[kern[k].post_count] == 0) + kern[k].post_gx[kern[k].post_count] = kern[k].grid_x; + if (kern[k].post_bx[kern[k].post_count] == 0) + kern[k].post_bx[kern[k].post_count] = 32; + if (kern[k].post_by[kern[k].post_count] == 0) + kern[k].post_by[kern[k].post_count] = 1; + if (kern[k].post_bz[kern[k].post_count] == 0) + kern[k].post_bz[kern[k].post_count] = 1; + ++kern[k].post_count; + } + } + void* w = NULL; + if (net->owned_count >= D4R_ENGINE_NET_OWNED || allocate(&w, bytes + 4096) != 0) + { + failure = "weights"; + break; + } + net->owned[net->owned_count++] = w; + if (to_device(w, file + at, bytes) != 0) + { + failure = "weights upload"; + break; + } + weights[i] = w; + divisor[i] = div; + shift_sx[i] = sx; + shift_sy[i] = sy; + layer_kernel[i] = (uint32_t)k; + } + /* The frozen parameter block of every layer at phase/frame 0; the graph loop copies and shifts it. */ + for (uint32_t i = 0; i < count && failure == NULL; ++i) + { + if (k_preset) + { + D4rEngineNetPwinParams* p = &layer_params[i].k; + p->W = (int32_t)(tw / divisor[i]); + p->H = (int32_t)(th / divisor[i]); + p->in = (uint8_t*)layer_in[i]; + p->skip = (uint8_t*)layer_skip[i]; + p->out24 = (uint8_t*)layer_out[i]; + p->out32 = i < 5 ? (uint8_t*)k_full[i] : NULL; + p->w = (uint8_t*)weights[i]; + d4r_engine_net_k_shift(d4r_engine_net_k_plan[i].role, 0, &p->sx, &p->sy); + } + else + { + D4rEngineNetParams* p = &layer_params[i].m; + p->w = weights[i]; + p->sx = shift_sx[i]; + p->sy = shift_sy[i]; + p->tw = (int32_t)(tw / divisor[i]); + p->th = (int32_t)(th / divisor[i]); + p->in = layer_in[i]; + p->skip = layer_skip[i]; + p->out = layer_out[i]; + p->merged = layer_merged[i]; + } + } + /* Every frame of the alignment cycle has its own immutable graph: frozen grids, blocks and arguments, + the K bottleneck's post phases included. No per-frame allocation or parameter copy. */ + for (uint32_t phase = 0; phase < net->graphs && failure == NULL; ++phase) + { + uint32_t n = 0; + for (uint32_t i = 0; i < count && failure == NULL; ++i) + { + const D4rEngineNetKernel* kk = &kern[layer_kernel[i]]; + if (n >= D4R_ENGINE_NET_NODES) + { + failure = "too many nodes"; + break; + } + D4rEngineNetNode* node = &net->nodes[phase][n++]; + node->function = kk->main; + node->bx = 32; + node->by = 1; + node->bz = kk->block_z; + node->params = layer_params[i]; + if (k_preset) + { + D4rEngineNetPwinParams* p = &node->params.k; + d4r_engine_net_k_shift(d4r_engine_net_k_plan[i].role, phase, &p->sx, &p->sy); + node->gx = kk->post_count != 0 ? (kk->grid_x != 0 ? kk->grid_x : 32) + : (uint32_t)(p->W + p->sx + 7) / 8; + node->gy = kk->post_count != 0 ? (kk->grid_y != 0 ? kk->grid_y : 1) + : (uint32_t)(p->H + p->sy + 7) / 8; + for (int j = 0; j < kk->post_count && failure == NULL; ++j) + { + if (n >= D4R_ENGINE_NET_NODES) + { + failure = "too many nodes"; + break; + } + D4rEngineNetNode* post = &net->nodes[phase][n++]; + post->function = kk->post[j]; + post->params = node->params; /* the same PwinParams block */ + post->gx = kk->post_gx[j]; + post->gy = kk->post_gy[j] != 0 ? kk->post_gy[j] : 1; + post->bx = kk->post_bx[j]; + post->by = kk->post_by[j]; + post->bz = kk->post_bz[j]; + } + } + else + { + D4rEngineNetParams* p = &node->params.m; + /* NGX alternates horizontal window alignment every evaluation, including reset evaluations. */ + p->sx = shift_sx[i] ^ (int32_t)(phase * 2); + node->gx = (uint32_t)(p->tw + p->sx + 7) / 8; + node->gy = (uint32_t)(p->th + shift_sy[i] + 7) / 8; + } + } + net->node_count[phase] = n; + } + /* The weights never change: every layer's weight image is prepared once, here, and found by its pointer. + A shared code object (the M tube) gets one preparation per layer, i.e. per distinct weights pointer. */ + for (uint32_t i = 0; i < count && failure == NULL; ++i) + { + const D4rEngineNetKernel* kk = &kern[layer_kernel[i]]; + if (kk->prep == NULL) + continue; + if (net->prep_count >= D4R_ENGINE_NET_LAYERS) + { + failure = "weight image slots"; + break; + } + D4rEngineNetNode* prep = &net->preps[net->prep_count++]; + prep->function = kk->prep; + prep->gx = kk->prep_blocks; + prep->bx = 128; + prep->by = 1; + prep->bz = 1; + prep->params = layer_params[i]; + } + for (uint32_t i = 0; i < net->prep_count && failure == NULL; ++i) + { + D4rEngineNetNode* prep = &net->preps[i]; + void* args[] = {&prep->params}; + if (launch(prep->function, prep->gx, 1, 1, prep->bx, prep->by, prep->bz, 0, net->stream, args, NULL) != 0) + failure = "weight preparation"; + } + if (failure == NULL && synchronize(net->stream) != 0) + failure = "weight preparation"; + if (failure == NULL && d4r_engine_net_capture(net, launch) != CUDA_SUCCESS) + failure = "network graph capture"; + if (failure == NULL) + { + HIP_EVENT_RECORD_FN upload = (HIP_EVENT_RECORD_FN)hip_symbol("hipGraphUpload"); + if (upload == NULL) + failure = "network graph upload"; + for (uint32_t phase = 0; phase < net->graphs && failure == NULL; ++phase) + if (upload(net->executable[phase], net->stream) != 0) + failure = "network graph upload"; + if (failure == NULL && synchronize(net->stream) != 0) + failure = "network graph upload"; + } + if (failure != NULL) + { + tracef("d4rEngineNetCreate: %s failed", failure); + d4rEngineNetDestroy(net); + return CUDA_ERROR_UNKNOWN; + } + *handle = net; + tracef("d4rEngineNetCreate: preset %s network on native HIP layers, %ux%u tokens, %u graphs", k_preset ? "K" : "M", + tw, th, net->graphs); + return CUDA_SUCCESS; +} + +/* All writers (including cancellation retirement in the shim) use this lock. The callback runs on a native + HIP host thread: no Wine/PE entry points and no HIP APIs are called from it. */ +int WINAPI d4rEngineNetSignalTimeline(void* handle, uint64_t value) +{ + D4rEngineNet* net = (D4rEngineNet*)handle; + if (net == NULL || net->timeline_device == NULL) + return -13; /* VK_ERROR_UNKNOWN */ + pthread_mutex_lock(&net->timeline_lock); + int result = 0; + if (value > net->timeline_value) + { + D4rVkSemaphoreSignalInfo info = {1000207005, NULL, net->timeline_semaphore, value}; + result = net->signal_timeline(net->timeline_device, &info); + if (result == 0) + net->timeline_value = value; + } + pthread_mutex_unlock(&net->timeline_lock); + return result; +} + +static void d4r_engine_net_completed(void* data) +{ + D4rEngineNetSignal* signal = (D4rEngineNetSignal*)data; + const int result = d4rEngineNetSignalTimeline(signal->net, signal->value); + if (result != 0) + __atomic_store_n(&signal->net->timeline_error, result, __ATOMIC_RELEASE); + __atomic_store_n(&signal->busy, 0, __ATOMIC_RELEASE); +} + +CUresult WINAPI d4rEngineNetSetTimeline(void* handle, void* client_device, uint64_t client_semaphore) +{ + D4rEngineNet* net = (D4rEngineNet*)handle; + if (net == NULL || net->timeline_device != NULL || client_device == NULL || client_semaphore == 0) + return CUDA_ERROR_INVALID_VALUE; + const uint64_t* device_object = (const uint64_t*)(uintptr_t)((const uint64_t*)client_device)[1]; + const uint64_t* semaphore_object = (const uint64_t*)(uintptr_t)client_semaphore; + if (device_object == NULL || device_object[1] != (uint64_t)(uintptr_t)client_device || + semaphore_object[1] != client_semaphore) + return CUDA_ERROR_NOT_SUPPORTED; + void* host_device = (void*)(uintptr_t)device_object[0]; + void* vulkan = dlopen("libvulkan.so.1", RTLD_NOW | RTLD_NOLOAD); + VK_GET_DEVICE_PROC_ADDR_FN get_proc = + vulkan != NULL ? (VK_GET_DEVICE_PROC_ADDR_FN)dlsym(vulkan, "vkGetDeviceProcAddr") : NULL; + HIP_LAUNCH_HOST_FN launch_host = (HIP_LAUNCH_HOST_FN)hip_symbol("hipLaunchHostFunc"); + if (get_proc == NULL || launch_host == NULL) + return CUDA_ERROR_NOT_SUPPORTED; + VK_SIGNAL_SEMAPHORE_FN signal = (VK_SIGNAL_SEMAPHORE_FN)get_proc(host_device, "vkSignalSemaphore"); + VK_GET_SEMAPHORE_COUNTER_FN counter = (VK_GET_SEMAPHORE_COUNTER_FN)get_proc(host_device, "vkGetSemaphoreCounterValue"); + uint64_t value = 0; + if (signal == NULL || counter == NULL || counter(host_device, semaphore_object[0], &value) != 0 || + pthread_mutex_init(&net->timeline_lock, NULL) != 0) + return CUDA_ERROR_NOT_SUPPORTED; + net->timeline_device = host_device; + net->timeline_semaphore = semaphore_object[0]; + net->timeline_value = value; + net->signal_timeline = signal; + net->launch_host = launch_host; + for (unsigned i = 0; i < D4R_ENGINE_NET_SIGNALS; ++i) + net->signals[i].net = net; + return CUDA_SUCCESS; +} + +/* Reserve before queuing this frame's input wait. A burst of discarded lists can leave callbacks in flight; + draining the previous work is safe here, but would deadlock on a newly queued wait for an unsubmitted list. */ +static D4rEngineNetSignal* d4r_engine_net_signal_slot(D4rEngineNet* net, uint32_t value) +{ + if (net->timeline_device == NULL || value == 0) + return NULL; + for (unsigned pass = 0; pass < 2; ++pass) + { + for (unsigned i = 0; i < D4R_ENGINE_NET_SIGNALS; ++i) + { + int free = 0; + if (__atomic_compare_exchange_n(&net->signals[i].busy, &free, 1, 0, __ATOMIC_ACQUIRE, __ATOMIC_RELAXED)) + { + net->signals[i].value = value; + return &net->signals[i]; + } + } + HIP_P_FN synchronize = (HIP_P_FN)hip_symbol("hipStreamSynchronize"); + if (pass != 0 || synchronize == NULL || synchronize(net->stream) != 0) + break; + } + return NULL; +} + +static CUresult d4r_engine_net_queue(D4rEngineNet* net, HIP_EVENT_RECORD_FN graph_launch, HIP_EVENT_RECORD_FN record, + D4rEngineNetSignal* signal) +{ + CUresult result = CUDA_ERROR_UNKNOWN; + if (record(net->start, net->stream) == 0) + { + /* One evaluation advances the alignment cycle by one frame; a reset does not restart it. */ + result = graph_launch(net->executable[net->frame % net->graphs], net->stream) == 0 + ? CUDA_SUCCESS : CUDA_ERROR_LAUNCH_FAILED; + if (result == CUDA_SUCCESS) + { + ++net->frame; + if (record(net->stop, net->stream) != 0 || + net->launch_host(net->stream, d4r_engine_net_completed, signal) != 0) + result = CUDA_ERROR_UNKNOWN; + } + } + if (result != CUDA_SUCCESS) + __atomic_store_n(&signal->busy, 0, __ATOMIC_RELEASE); + return result; +} + +CUresult WINAPI d4rEngineNetLaunch(void* handle, uint32_t value) +{ + D4rEngineNet* net = (D4rEngineNet*)handle; + HIP_EVENT_RECORD_FN graph_launch = (HIP_EVENT_RECORD_FN)hip_symbol("hipGraphLaunch"); + HIP_EVENT_RECORD_FN record = (HIP_EVENT_RECORD_FN)hip_symbol("hipEventRecord"); + if (net == NULL || graph_launch == NULL || record == NULL) + return CUDA_ERROR_INVALID_VALUE; + D4rEngineNetSignal* signal = d4r_engine_net_signal_slot(net, value); + if (signal == NULL) + return CUDA_ERROR_INVALID_VALUE; + return d4r_engine_net_queue(net, graph_launch, record, signal); +} + +/* Check HIP's wait-value capability and a round trip on this imported buffer before enabling LaunchAt. + Release the probe wait before querying its event: HIP may flush the blocked queue during a query. + The test leaves the word at 0; frame values then increase monotonically and use a Gte wait. */ +int WINAPI d4rEngineNetInputWaitReady(void* handle, CUdeviceptr sync, uint32_t* host) +{ + D4rEngineNet* net = (D4rEngineNet*)handle; + HIP_STREAM_WAIT_VALUE32_FN wait = (HIP_STREAM_WAIT_VALUE32_FN)hip_symbol("hipStreamWaitValue32"); + HIP_STREAM_WRITE_VALUE32_FN write = (HIP_STREAM_WRITE_VALUE32_FN)hip_symbol("hipStreamWriteValue32"); + HIP_EVENT_RECORD_FN record = (HIP_EVENT_RECORD_FN)hip_symbol("hipEventRecord"); + HIP_P_FN query = (HIP_P_FN)hip_symbol("hipEventQuery"); + HIP_PP_FN event_create = (HIP_PP_FN)hip_symbol("hipEventCreate"); + HIP_P_FN event_destroy = (HIP_P_FN)hip_symbol("hipEventDestroy"); + HIP_P_FN synchronize = (HIP_P_FN)hip_symbol("hipStreamSynchronize"); + HIP_MEMCPY_FN to_host = (HIP_MEMCPY_FN)hip_symbol("hipMemcpyDtoH"); + typedef int (*GET_ATTRIBUTE_FN)(int*, int, int); + typedef int (*GET_DEVICE_FN)(int*); + GET_ATTRIBUTE_FN attribute = (GET_ATTRIBUTE_FN)hip_symbol("hipDeviceGetAttribute"); + GET_DEVICE_FN device = (GET_DEVICE_FN)hip_symbol("hipGetDevice"); + int current = 0, supported = 0; + if (attribute == NULL || device == NULL || device(¤t) != 0 || + attribute(&supported, 10013 /* hipDeviceAttributeCanUseStreamWaitValue */, current) != 0 || !supported) + return 0; + if (net == NULL || sync == 0 || host == NULL || wait == NULL || write == NULL || record == NULL || query == NULL || + event_create == NULL || event_destroy == NULL || synchronize == NULL || to_host == NULL) + return 0; + /* Set for the probe's own releases; cleared again below unless the probe proves GPU-side waits work. */ + net->input_wait = sync; + net->input_wait_host = host; + void* probe = NULL; + if (event_create(&probe) != 0) + { + net->input_wait = 0; + net->input_wait_host = NULL; + return 0; + } + /* HIP lazily initializes queue/event storage. Do that before blocking a queue: allocating it after the + wait has been submitted can itself synchronize against the blocked queue. */ + if (d4rEngineNetReleaseInput(net, 0) != CUDA_SUCCESS || record(probe, net->stream) != 0 || + record(net->start, net->stream) != 0 || record(net->stop, net->stream) != 0 || + synchronize(net->stream) != 0) + { + event_destroy(probe); + net->input_wait = 0; + net->input_wait_host = NULL; + return 0; + } + int ready = 0, wedged = 0; + if (write(net->stream, (void*)(uintptr_t)sync, 0, 0) == 0 && synchronize(net->stream) == 0 && + wait(net->stream, (void*)(uintptr_t)sync, 1, 0 /* hipStreamWaitValueGte */, 0xffffffffu) == 0) + { + /* The wait is queued now: it must be satisfied before the stream is synchronized again. */ + const int recorded = record(probe, net->stream); + int state = 600; /* hipErrorNotReady */ + if (d4rEngineNetReleaseInput(net, 1) != CUDA_SUCCESS) + wedged = 1; + else if (recorded == 0) + { + for (int i = 0; i < 4000 && state == 600; ++i) + { + usleep(50); + state = query(probe); + } + if (state == 600) + wedged = 1; + else + ready = state == 0; + } + if (!wedged) + { + /* Back to 0 for the frames' own values, with no wait left queued. */ + if (write(net->stream, (void*)(uintptr_t)sync, 0, 0) != 0 || synchronize(net->stream) != 0) + ready = 0; + if (ready != 0) + { + uint32_t value = 0; + if (to_host(&value, (void*)(uintptr_t)sync, 4) != 0 || value != 0) + ready = 0; + } + } + } + event_destroy(probe); + if (wedged) + { + /* Destroy releases it (d4rEngineNetDestroy); the caller must not use this network. */ + net->input_wait = sync; + net->input_wait_value = 1; + } + else if (!ready) + { + /* d4rEngineNetLaunchAt refuses without a proven wait. */ + net->input_wait = 0; + net->input_wait_host = NULL; + } + tracef("d4rEngineNetInputWaitReady(%p): %s", (void*)(uintptr_t)sync, + wedged ? "the stream is stuck" : ready ? "GPU-side waits" : "not supported"); + return wedged ? -1 : ready; +} + +/* Queues a wait for `value` (the frame number the engine's front fills into the sync buffer) before this + evaluation's graph, so the network starts on the GPU the moment its inputs are complete. */ +CUresult WINAPI d4rEngineNetLaunchAt(void* handle, uint32_t value) +{ + D4rEngineNet* net = (D4rEngineNet*)handle; + HIP_EVENT_RECORD_FN graph_launch = (HIP_EVENT_RECORD_FN)hip_symbol("hipGraphLaunch"); + HIP_EVENT_RECORD_FN record = (HIP_EVENT_RECORD_FN)hip_symbol("hipEventRecord"); + HIP_STREAM_WAIT_VALUE32_FN wait = (HIP_STREAM_WAIT_VALUE32_FN)hip_symbol("hipStreamWaitValue32"); + if (net == NULL || graph_launch == NULL || record == NULL) + return CUDA_ERROR_INVALID_VALUE; + if (net->input_wait == 0 || wait == NULL) + return CUDA_ERROR_NOT_SUPPORTED; + D4rEngineNetSignal* signal = d4r_engine_net_signal_slot(net, value); + if (signal == NULL) + return CUDA_ERROR_INVALID_VALUE; + if (wait(net->stream, (void*)(uintptr_t)net->input_wait, value, 0 /* hipStreamWaitValueGte */, 0xffffffffu) != 0) + { + __atomic_store_n(&signal->busy, 0, __ATOMIC_RELEASE); + return CUDA_ERROR_NOT_SUPPORTED; + } + net->input_wait_value = value; + // Keep the queued wait tracked even if a later event/kernel enqueue fails; destroy must release it. + return d4r_engine_net_queue(net, graph_launch, record, signal); +} + +/* 1: the launched layers have finished; 0: still running; negative: error */ +int WINAPI d4rEngineNetDone(void* handle, float* milliseconds) +{ + D4rEngineNet* net = (D4rEngineNet*)handle; + HIP_P_FN query = (HIP_P_FN)hip_symbol("hipEventQuery"); + HIP_EVENT_ELAPSED_FN elapsed = (HIP_EVENT_ELAPSED_FN)hip_symbol("hipEventElapsedTime"); + if (net == NULL || query == NULL || __atomic_load_n(&net->timeline_error, __ATOMIC_ACQUIRE) != 0) + return -1; + const int result = query(net->stop); + if (result == 600 /* hipErrorNotReady */) + return 0; + if (result != 0) + return -1; + /* The layers ran, so the input wait (if any) is behind us. */ + net->input_wait_value = 0; + if (milliseconds != NULL && (elapsed == NULL || elapsed(milliseconds, net->start, net->stop) != 0)) + *milliseconds = 0.0f; + return 1; +} + /* d4r: asynchronous 2D copy between device memory and arrays on the null stream (ZLUDA forwards CUarray handles as hipArray_t and the null CUstream as HIP's null stream; ZLUDA itself has no cuMemcpy2DAsync). The CUDA and HIP descriptors share their layout; only the memory type codes @@ -3212,11 +4164,6 @@ CUresult WINAPI d4rMemcpy2DAsync(const D4rMemcpy2D* copy, CUstream stream) /* d4r: GPU-side handoff from the game's queue. The null stream (where NGX runs) waits until the u32 at `pointer` (device memory the game's queue writes with vkCmdFillBuffer) is >= value. d4rWriteValue32 writes such a value from a separate non-blocking stream (initialisation, or releasing a stuck wait). */ -typedef int (*HIP_STREAM_WAIT_VALUE32_FN)(void*, void*, uint32_t, unsigned int, uint32_t); -typedef int (*HIP_STREAM_WRITE_VALUE32_FN)(void*, void*, uint32_t, unsigned int); -typedef int (*HIP_STREAM_CREATE_WITH_FLAGS_FN)(void**, unsigned int); -typedef int (*HIP_STREAM_SYNCHRONIZE_FN)(void*); - CUresult WINAPI d4rStreamWaitValue32(CUdeviceptr pointer, uint32_t value) { static HIP_STREAM_WAIT_VALUE32_FN function; diff --git a/tools/wine_nvcuda_bridge.spec b/tools/wine_nvcuda_bridge.spec index 9c6dcd2..01da021 100644 --- a/tools/wine_nvcuda_bridge.spec +++ b/tools/wine_nvcuda_bridge.spec @@ -90,6 +90,15 @@ @ stdcall cuPointerGetAttribute(ptr ptr ptr) cuPointerGetAttribute @ stdcall d4rImportVulkanMemory(ptr int64 int64 ptr ptr) d4rImportVulkanMemory @ stdcall d4rReleaseVulkanMemory(ptr) d4rReleaseVulkanMemory +@ stdcall d4rEngineNetCreate(ptr int64 ptr long long long int64 int64 ptr) d4rEngineNetCreate +@ stdcall d4rEngineNetLaunch(ptr long) d4rEngineNetLaunch +@ stdcall d4rEngineNetDone(ptr ptr) d4rEngineNetDone +@ stdcall d4rEngineNetDestroy(ptr) d4rEngineNetDestroy +@ stdcall d4rEngineNetInputWaitReady(ptr int64 ptr) d4rEngineNetInputWaitReady +@ stdcall d4rEngineNetLaunchAt(ptr long) d4rEngineNetLaunchAt +@ stdcall d4rEngineNetReleaseInput(ptr long) d4rEngineNetReleaseInput +@ stdcall d4rEngineNetSetTimeline(ptr ptr int64) d4rEngineNetSetTimeline +@ stdcall d4rEngineNetSignalTimeline(ptr int64) d4rEngineNetSignalTimeline @ stdcall d4rMemcpy2DAsync(ptr ptr) d4rMemcpy2DAsync @ stdcall d4rCtxSynchronize() d4rCtxSynchronize @ stdcall d4rSetShutdownHook(ptr) d4rSetShutdownHook diff --git a/vulkan/README.md b/vulkan/README.md new file mode 100644 index 0000000..59eef77 --- /dev/null +++ b/vulkan/README.md @@ -0,0 +1,47 @@ +# Vulkan kernels (experimental) + +Tools for running DLSS texture kernels as Vulkan compute shaders instead of through CUDA/HIP. Nothing here is +wired into the shim yet; it is a kernel-level proof with a byte-exact check against captured frames. + +| File | Role | +|---|---| +| `ptx2glsl.py` | Translates one PTX kernel from your DLSS DLL into GLSL. Every instruction becomes the same operation `kernels/native/ptx2hip.py` emits, gotos become nested `if`/`else` and `do`/`while` (each branch rejoins at its immediate post-dominator), and registers that only hold half floats are typed `float16_t` / `f16vec2`. | +| `kvk.cpp` | Headless runner: binds captured textures, surfaces, buffers and the parameter block, dispatches, writes the surfaces back and times the dispatch. | +| `kcase.py` | Builds a `kvk` manifest from a capture and compares `kvk`'s output with the captured surfaces. | + +## Reproducing the K output kernel check + +```sh +# 1. capture (bridge options that already exist): every launch's resources after it ran, plus the arguments +D4R_CUDA_VERBOSE=1 D4R_CUDA_LAUNCH_STATS=1 D4R_CUDA_LAUNCH_DUMP_DIR=$CAP/dump \ +D4R_CUDA_REPLAY_DUMP_DIR=$CAP/replay D4R_CUDA_REPLAY_DUMP_FILTER=hiluma_engine_output D4R_CUDA_REPLAY_DUMP_LIMIT=8 \ + +# 2. translate and compile +K=hiluma_engine_output_depthinv_mvlo_hdr_max_v2_rel +python3 ptx2glsl.py MODULE.ptx $K > k.comp && glslangValidator -V --target-env vulkan1.3 k.comp -o k.spv +g++ -O2 -std=c++17 -o kvk kvk.cpp -lvulkan +# 3. run frame N and compare (the three linear-filtered textures are named by parameter offset) +python3 kcase.py make $CAP $K N case 200:linear 224:linear 232:linear +./kvk k.spv case/manifest 150 64 # 150 dispatches per submit, 64-wide subgroups +python3 kcase.py cmp case +``` + +Sampler settings come from the `cuTexObjectCreate` lines of the verbose log (all clamp, normalized coordinates; +filter per object). + +## Result (RX 9070 XT, RADV / Mesa 26.2.4, DLSS 310.7, 2026-10-08) + +Byte-identical to the native HIP kernel's four output surfaces on six consecutive frames at 640x360 -> 1280x720 and +one frame at 1920x1080 -> 3840x2160. At 4K: 1.18 ms per dispatch under sustained load (GPU clock 2.38 GHz), +1.05 ms with 0.3 ms gaps (2.7 GHz). The native kernel measures 1.16 ms in the live benchmark at about 3.1 GHz. The +two were not timed in the same run or clock state. + +## Pitfalls found + +- Shader-visible buffers must be allocated from a `DEVICE_LOCAL | HOST_VISIBLE` memory type. The first + host-visible type on RADV is system memory; reading the network output from it made every variant take 2.9 ms. +- A program-counter loop around the whole shader (`while (pc != END) { if (pc == k) ... }`) keeps every value + live across the back edge: 35 ms. Guards without the loop: 2.7 ms (VRAM buffers). Real nesting: 1.69 ms. +- Packing half floats into `uint` registers costs about 840 `v_mov_b16`; native half types: 1.35 -> 1.18 ms. +- A 64K-entry table for `ex2.approx.f16x2` is exact but slower than `v_exp_f32`; dropping `precise` gains under 1% + and changes about 0.005% of the output values. diff --git a/vulkan/condense.py b/vulkan/condense.py new file mode 100644 index 0000000..9bdf4bb --- /dev/null +++ b/vulkan/condense.py @@ -0,0 +1,50 @@ +#!/usr/bin/env python3 +"""Condense a ptx2glsl.py translation for reading: drops declarations, prints float constants as decimals and +substitutes registers that are assigned once and used once into their use. The result is not meant to compile. + +usage: condense.py KERNEL.comp > KERNEL.txt +""" +import re, struct, sys +src = open(sys.argv[1]).read().split('\n') +start = next(i for i, l in enumerate(src) if l.startswith('void main')) +body = [l for l in src[start:] if not re.match(r'\s*(precise )?(uint|float|bool|uint64_t|float16_t|f16vec2|int) [a-z0-9_]+ = [^;]*(, [a-z0-9_]+ = [^;,]*)*;\s*$', l) or '(' in l.split('=', 1)[1][:3]] +def fl(m): + v = struct.unpack(' 1e-5 or v == 0) else f'{v:.6g}' +body = [re.sub(r'uintBitsToFloat\(0x([0-9A-Fa-f]{8})u\)', fl, l) for l in body] +body = [re.sub(r'\b(\d+)u\b', r'\1', l) for l in body] +body = [re.sub(r' \((\d+) & 31\)', r' \1', l) for l in body] +tok = re.compile(r'\b(?:r|f|rs|rd|p)\d+\b') +for _ in range(12): + assigns, uses = {}, {} + for i, l in enumerate(body): + m = re.match(r'(\s*)((?:r|f|rs|rd|p)\d+) = (.*);\s*$', l) + rhs = l + if m: + assigns.setdefault(m.group(2), []).append(i); rhs = m.group(3) + for t in tok.findall(rhs): uses.setdefault(t, []).append(i) + changed = False + for name, at in assigns.items(): + u = uses.get(name, []) + if len(at) != 1 or len(u) != 1 or u[0] <= at[0]: continue + a, b = at[0], u[0] + if body[a] is None or body[b] is None: continue + ind = lambda s: len(s) - len(s.lstrip()) + # only within straight-line code at the same nesting, and short enough to stay readable + if any(body[k] is not None and ('{' in body[k] or '}' in body[k]) for k in range(a, b)): continue + m = re.match(r'(\s*)((?:r|f|rs|rd|p)\d+) = (.*);\s*$', body[a]) + expr = m.group(3) + if len(expr) > 90: continue + # the expression's inputs must not be reassigned in between + ins = set(tok.findall(expr)) + if any(body[k] is not None and re.match(r'\s*(%s) = ' % '|'.join(map(re.escape, ins)) if ins else r'$^', body[k]) for k in range(a + 1, b)): continue + simple = re.fullmatch(r'[\w.]+|\w+\([^()]*\)', expr) + body[b] = re.sub(r'\b%s\b' % name, expr if simple else f'({expr})', body[b], count=1) + body[a] = None; changed = True + body = [l for l in body if l is not None] + if not changed: break +for a, b in ((r'uint\(uint\(([^()]*)\)\)', r'\1'), (r'\buint\(((?:r|rs)\d+)\)', r'\1'), (r'uint64_t\((rd\d+)\)', r'\1')): + body = [re.sub(a, b, l) for l in body] +# scripts/check_redistributable.py refuses files carrying this line: the input was NVIDIA's code. +print(f'// Generated by vulkan/condense.py from NVIDIA PTX ({sys.argv[1]}); do not commit or redistribute.') +print('\n'.join(l.rstrip() for l in body if l.strip())) diff --git a/vulkan/kcase.py b/vulkan/kcase.py new file mode 100644 index 0000000..dde99ee --- /dev/null +++ b/vulkan/kcase.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python3 +"""Build a kvk manifest for one captured launch and compare kvk's outputs with the capture. + +usage: kcase.py make CAPTURE_DIR KERNEL FRAME OUTDIR [--tex OFF:linear ...] (FRAME counts that kernel's launches from 0) + kcase.py cmp OUTDIR +CAPTURE_DIR holds dump/ (D4R_CUDA_LAUNCH_DUMP_DIR) and replay/ (D4R_CUDA_REPLAY_DUMP_DIR) of the same run. +""" +import glob, os, re, sys +import numpy as np +PING = set() + +def launches(cap, kernel): + return sorted({int(re.match(r'launch-(\d+)-', os.path.basename(f)).group(1)) for f in glob.glob(f'{cap}/dump/launch-*-{kernel}-arg*')}) + +def res(cap, kernel, seq): + out = {} + for f in glob.glob(f'{cap}/dump/launch-{seq:03d}-{kernel}-arg*'): + m = re.search(r'-arg(\d+)-(texture|surface|buffer)-(.*)-fmt(-?\d+)\.bin$', f) + out[int(m.group(1))] = (m.group(2), m.group(3), int(m.group(4)), f) + return out + +if sys.argv[1] == 'make': + cap, kernel, frame, outdir = sys.argv[2], sys.argv[3], int(sys.argv[4]), sys.argv[5] + linear = {int(a.split(':')[0]) for a in sys.argv[6:] if a.endswith(':linear')} + os.makedirs(outdir, exist_ok=True) + seqs = launches(cap, kernel); seq = seqs[frame] + cur = res(cap, kernel, seq) + man = open(f'{cap}/replay/replay-{seq:06d}-{kernel}/manifest.txt').read().split('\n') + grid = next(l for l in man if l.startswith('launch')).split() + lines = [f'param {cap}/replay/replay-{seq:06d}-{kernel}/args.bin', f'grid {grid[1]} {grid[2]}'] + ptr = {int(l.split()[1]): int(l.split()[3]) for l in man if l.startswith('pointer')} + expect = [] + for off, (kind, shape, fmt, f) in sorted(cur.items()): + if kind == 'buffer': + lines.append(f'buf {off} {f} {ptr[off]}'); continue + w, h, c = (int(x) for x in shape.split('x')) + bits = {32: 32, 194: 8, 192: 8, 193: 8, 200: 9, 198: 9, 199: 9}.get(fmt, 16) # 19x: UNORM8, 198..200: SNORM8 + if kind == 'texture': + lines.append(f'tex {off} {f} {w} {h} {c} {bits} {"linear" if off in linear else "nearest"}') + else: + # the surface's contents before this launch: its dump from the latest earlier launch that had the same data layout + init = '-' + for s in reversed([s for s in seqs if s < seq]): + prev = res(cap, kernel, s) + # ping-pong surfaces alternate; take the one whose size matches and which is not this frame's input + if off in prev and prev[off][1] == shape and (s == seqs[frame - 2] if frame >= 2 and off in PING else True): + init = prev[off][3]; break + lines.append(f'surf {off} {w} {h} {c} {bits} {init} {outdir}/out-{off}.bin') + expect.append((off, f, c, bits)) + open(f'{outdir}/manifest', 'w').write('\n'.join(lines) + '\n') + open(f'{outdir}/expect', 'w').write('\n'.join(f'{o} {f} {c} {b}' for o, f, c, b in expect) + '\n') + print('\n'.join(lines)) +else: + outdir = sys.argv[2] + ok = True + for l in open(f'{outdir}/expect').read().split('\n'): + if not l: + continue + off, f, c, bits = l.split() + dt = np.uint16 if bits == '16' else np.uint32 + a = np.fromfile(f, dtype=dt); b = np.fromfile(f'{outdir}/out-{off}.bin', dtype=dt) + diff = a != b + fa = a.view(np.float16 if bits == '16' else np.float32).astype(np.float64); fb = b.view(np.float16 if bits == '16' else np.float32).astype(np.float64) + with np.errstate(invalid='ignore'): + d = np.abs(fa - fb); d = d[np.isfinite(d)] + ok &= not diff.any() + print(f'surface {off}: {a.size} values, {int(diff.sum())} differ ({diff.mean() * 100:.4f}%), max |diff| {d.max() if d.size else 0:.6g}, mean |ref| {np.nanmean(np.abs(fa[np.isfinite(fa)])):.4g}') + print('BYTE-IDENTICAL' if ok else 'DIFFERENT') diff --git a/vulkan/kvk.cpp b/vulkan/kvk.cpp new file mode 100644 index 0000000..be0f4c3 --- /dev/null +++ b/vulkan/kvk.cpp @@ -0,0 +1,269 @@ +// Runs one translated DLSS texture kernel as a Vulkan compute shader on captured inputs. +// kvk SHADER.spv MANIFEST [repeats] [subgroup] +// MANIFEST lines (bindings are assigned in the order ptx2glsl.py uses: parameter block, textures by +// offset, surfaces by offset): +// param FILE +// grid X Y +// tex OFF FILE W H CHANNELS BITS linear|nearest +// surf OFF W H CHANNELS BITS INITFILE|- OUTFILE +// buf OFF FILE DELTA (8-byte pointer at OFF = buffer address + DELTA) +// lut FILE (storage buffer bound after the surfaces) +// bufout OFF FILE (after the dispatch, write the pointer buffer of parameter offset OFF to FILE) +// outbuf FILE BYTES (storage buffer bound after those, written to FILE after the dispatch) +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#define CK(x) do { VkResult r_ = (x); if (r_ != VK_SUCCESS) { fprintf(stderr, "%s failed: %d\n", #x, r_); exit(1); } } while (0) +static std::vector slurp(const std::string& p) +{ + std::ifstream f(p, std::ios::binary); + if (!f) { fprintf(stderr, "cannot read %s\n", p.c_str()); exit(1); } + return std::vector((std::istreambuf_iterator(f)), {}); +} +static VkDevice dev; static VkPhysicalDeviceMemoryProperties mp; +static VkDeviceMemory alloc(VkMemoryRequirements mr, VkMemoryPropertyFlags want, bool bda) +{ + uint32_t mt = ~0u; + for (uint32_t t = 0; t < mp.memoryTypeCount; ++t) if ((mr.memoryTypeBits >> t & 1) && (mp.memoryTypes[t].propertyFlags & want) == want) { mt = t; break; } + VkMemoryAllocateFlagsInfo fi{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_FLAGS_INFO}; fi.flags = VK_MEMORY_ALLOCATE_DEVICE_ADDRESS_BIT; + VkMemoryAllocateInfo ai{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO, bda ? &fi : nullptr}; ai.allocationSize = mr.size; ai.memoryTypeIndex = mt; + VkDeviceMemory m; CK(vkAllocateMemory(dev, &ai, nullptr, &m)); return m; +} +struct Buf { VkBuffer b; VkDeviceMemory m; void* map; VkDeviceSize size; uint64_t addr; }; +static Buf make_buf(VkDeviceSize size, VkBufferUsageFlags usage) +{ + Buf r{}; r.size = size; + VkBufferCreateInfo ci{VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO}; ci.size = size; ci.usage = usage | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT; + CK(vkCreateBuffer(dev, &ci, nullptr, &r.b)); + VkMemoryRequirements mr; vkGetBufferMemoryRequirements(dev, r.b, &mr); + // shader-visible data must live in VRAM: the first host-visible type on RADV is system memory, which the GPU reads over PCIe + const VkMemoryPropertyFlags vram = VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT | VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT; + bool has_vram = false; + for (uint32_t t = 0; t < mp.memoryTypeCount; ++t) has_vram |= (mr.memoryTypeBits >> t & 1) && (mp.memoryTypes[t].propertyFlags & vram) == vram; + r.m = alloc(mr, has_vram ? vram : VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, true); + CK(vkBindBufferMemory(dev, r.b, r.m, 0)); CK(vkMapMemory(dev, r.m, 0, size, 0, &r.map)); + VkBufferDeviceAddressInfo ai{VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO}; ai.buffer = r.b; r.addr = vkGetBufferDeviceAddress(dev, &ai); + return r; +} +struct Img { VkImage i; VkImageView v; VkFormat f; uint32_t w, h, bpp; std::string init, out; bool surf, linear; int off; }; +static VkFormat fmt_of(int ch, int bits) +{ + // 8: UNORM, 9: SNORM, one byte per channel (CUDA's UNORM_INT8Xn / SNORM_INT8Xn arrays) + if (bits == 8) return ch == 1 ? VK_FORMAT_R8_UNORM : ch == 2 ? VK_FORMAT_R8G8_UNORM : VK_FORMAT_R8G8B8A8_UNORM; + if (bits == 9) return ch == 1 ? VK_FORMAT_R8_SNORM : ch == 2 ? VK_FORMAT_R8G8_SNORM : VK_FORMAT_R8G8B8A8_SNORM; + if (bits == 10 && ch == 4) return VK_FORMAT_A2B10G10R10_UNORM_PACK32; + if (bits == 16) return ch == 1 ? VK_FORMAT_R16_SFLOAT : ch == 2 ? VK_FORMAT_R16G16_SFLOAT : VK_FORMAT_R16G16B16A16_SFLOAT; + return ch == 1 ? VK_FORMAT_R32_SFLOAT : ch == 2 ? VK_FORMAT_R32G32_SFLOAT : VK_FORMAT_R32G32B32A32_SFLOAT; +} +int main(int argc, char** argv) +{ + if (argc < 3) return 1; + const int repeats = argc > 3 ? atoi(argv[3]) : 1; const uint32_t subgroup = argc > 4 ? atoi(argv[4]) : 32; + VkApplicationInfo app{VK_STRUCTURE_TYPE_APPLICATION_INFO}; app.apiVersion = VK_API_VERSION_1_3; + VkInstanceCreateInfo ici{VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO}; ici.pApplicationInfo = &app; + VkInstance inst; CK(vkCreateInstance(&ici, nullptr, &inst)); + uint32_t np = 8; VkPhysicalDevice pds[8]; vkEnumeratePhysicalDevices(inst, &np, pds); + VkPhysicalDevice pd = pds[0]; + for (uint32_t i = 0; i < np; ++i) { VkPhysicalDeviceProperties p; vkGetPhysicalDeviceProperties(pds[i], &p); if (p.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU) pd = pds[i]; } + VkPhysicalDeviceVulkan13Features v13{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_FEATURES}; + VkPhysicalDeviceVulkan12Features v12{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES, &v13}; + VkPhysicalDeviceVulkan11Features v11{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES, &v12}; + VkPhysicalDeviceFeatures2 f2{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2, &v11}; + vkGetPhysicalDeviceFeatures2(pd, &f2); + vkGetPhysicalDeviceMemoryProperties(pd, &mp); + uint32_t nq = 8; VkQueueFamilyProperties qf[8]; vkGetPhysicalDeviceQueueFamilyProperties(pd, &nq, qf); + uint32_t qi = 0; for (uint32_t i = 0; i < nq; ++i) if (qf[i].queueFlags & VK_QUEUE_COMPUTE_BIT) { qi = i; break; } + float prio = 1; VkDeviceQueueCreateInfo qci{VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO}; qci.queueFamilyIndex = qi; qci.queueCount = 1; qci.pQueuePriorities = &prio; + VkDeviceCreateInfo dci{VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO, &f2}; dci.queueCreateInfoCount = 1; dci.pQueueCreateInfos = &qci; + CK(vkCreateDevice(pd, &dci, nullptr, &dev)); + VkQueue q; vkGetDeviceQueue(dev, qi, 0, &q); + VkCommandPoolCreateInfo cpci{VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO}; cpci.queueFamilyIndex = qi; cpci.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT; + VkCommandPool cp; CK(vkCreateCommandPool(dev, &cpci, nullptr, &cp)); + VkCommandBufferAllocateInfo cai{VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO}; cai.commandPool = cp; cai.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; cai.commandBufferCount = 1; + VkCommandBuffer cb; CK(vkAllocateCommandBuffers(dev, &cai, &cb)); + auto run = [&](auto&& record) { + VkCommandBufferBeginInfo bi{VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO}; CK(vkBeginCommandBuffer(cb, &bi)); record(); CK(vkEndCommandBuffer(cb)); + VkSubmitInfo si{VK_STRUCTURE_TYPE_SUBMIT_INFO}; si.commandBufferCount = 1; si.pCommandBuffers = &cb; + CK(vkQueueSubmit(q, 1, &si, VK_NULL_HANDLE)); CK(vkQueueWaitIdle(q)); + }; + // ---------------------------------------------------------------- manifest + std::vector param; uint32_t gx = 1, gy = 1; + std::vector imgs; std::map> bufs; std::vector luts; std::vector lutOut; std::map bufOut; + std::ifstream mf(argv[2]); std::string line; + while (std::getline(mf, line)) + { + std::istringstream ls(line); std::string k; ls >> k; + if (k == "param") { std::string f; ls >> f; param = slurp(f); } + else if (k == "grid") ls >> gx >> gy; + else if (k == "tex" || k == "surf") + { + Img im{}; int ch, bits; std::string mode; im.surf = k == "surf"; + if (im.surf) ls >> im.off >> im.w >> im.h >> ch >> bits >> im.init >> im.out; + else { ls >> im.off >> im.init >> im.w >> im.h >> ch >> bits >> mode; im.linear = mode == "linear"; } + im.f = fmt_of(ch, bits); im.bpp = bits < 16 ? ch : ch * bits / 8; + imgs.push_back(im); + } + else if (k == "lut") + { + std::string f; ls >> f; auto data = slurp(f); + Buf b = make_buf(data.size(), VK_BUFFER_USAGE_STORAGE_BUFFER_BIT); memcpy(b.map, data.data(), data.size()); luts.push_back(b); lutOut.push_back(""); + } + else if (k == "bufout") { int off; std::string f; ls >> off >> f; bufOut[off] = f; } + else if (k == "outbuf") + { + std::string f; VkDeviceSize n; ls >> f >> n; + Buf b = make_buf(n, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT); memset(b.map, 0, n); luts.push_back(b); lutOut.push_back(f); + } + else if (k == "buf") + { + int off; std::string f; int64_t delta; ls >> off >> f >> delta; + auto data = slurp(f); Buf b = make_buf(data.size(), VK_BUFFER_USAGE_STORAGE_BUFFER_BIT); memcpy(b.map, data.data(), data.size()); + bufs[off] = {b, delta}; + } + } + for (auto& [off, b] : bufs) { uint64_t a = b.first.addr + b.second; memcpy(param.data() + off, &a, 8); } + std::sort(imgs.begin(), imgs.end(), [](const Img& a, const Img& b) { return a.surf != b.surf ? !a.surf : a.off < b.off; }); + Buf pbuf = make_buf((param.size() + 3) & ~3ull, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT); memcpy(pbuf.map, param.data(), param.size()); + // ---------------------------------------------------------------- images + VkSamplerCreateInfo sci{VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO}; + sci.addressModeU = sci.addressModeV = sci.addressModeW = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE; sci.maxLod = 0; + VkSampler nearest, linear; CK(vkCreateSampler(dev, &sci, nullptr, &nearest)); + sci.magFilter = sci.minFilter = VK_FILTER_LINEAR; CK(vkCreateSampler(dev, &sci, nullptr, &linear)); + std::vector staging; + for (Img& im : imgs) + { + const bool r16u = im.surf && im.f == VK_FORMAT_R16_SFLOAT; // byte-addressed surface: bound as R16_UINT + VkImageCreateInfo ci{VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO}; ci.imageType = VK_IMAGE_TYPE_2D; ci.format = im.f; ci.extent = {im.w, im.h, 1}; + ci.mipLevels = 1; ci.arrayLayers = 1; ci.samples = VK_SAMPLE_COUNT_1_BIT; + ci.usage = VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_STORAGE_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT; + if (r16u) ci.flags = VK_IMAGE_CREATE_MUTABLE_FORMAT_BIT; + CK(vkCreateImage(dev, &ci, nullptr, &im.i)); + VkMemoryRequirements mr; vkGetImageMemoryRequirements(dev, im.i, &mr); + CK(vkBindImageMemory(dev, im.i, alloc(mr, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, false), 0)); + VkImageViewCreateInfo vi{VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO}; vi.image = im.i; vi.viewType = VK_IMAGE_VIEW_TYPE_2D; + vi.format = r16u ? VK_FORMAT_R16_UINT : im.f; vi.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + CK(vkCreateImageView(dev, &vi, nullptr, &im.v)); + Buf st = make_buf((VkDeviceSize)im.w * im.h * im.bpp, VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT); + memset(st.map, 0, st.size); + if (im.init != "-") { auto d = slurp(im.init); memcpy(st.map, d.data(), std::min(d.size(), st.size)); } + staging.push_back(st); + } + auto barrier = [&](VkImage i, VkImageLayout a, VkImageLayout b) { + VkImageMemoryBarrier mb{VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER}; mb.oldLayout = a; mb.newLayout = b; mb.image = i; + mb.srcAccessMask = mb.dstAccessMask = VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT; + mb.srcQueueFamilyIndex = mb.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; mb.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + vkCmdPipelineBarrier(cb, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, 0, 0, nullptr, 0, nullptr, 1, &mb); + }; + auto upload = [&]() { + for (size_t i = 0; i < imgs.size(); ++i) + { + barrier(imgs[i].i, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL); + VkBufferImageCopy r{}; r.imageSubresource = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1}; r.imageExtent = {imgs[i].w, imgs[i].h, 1}; + vkCmdCopyBufferToImage(cb, staging[i].b, imgs[i].i, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &r); + barrier(imgs[i].i, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_GENERAL); + } + }; + run(upload); + // ---------------------------------------------------------------- pipeline + auto code = slurp(argv[1]); + VkShaderModuleCreateInfo smi{VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO}; smi.codeSize = code.size(); smi.pCode = (const uint32_t*)code.data(); + VkShaderModule sm; CK(vkCreateShaderModule(dev, &smi, nullptr, &sm)); + std::vector lb; + lb.push_back({0, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, VK_SHADER_STAGE_COMPUTE_BIT, nullptr}); + for (size_t i = 0; i < imgs.size(); ++i) + lb.push_back({(uint32_t)i + 1, imgs[i].surf ? VK_DESCRIPTOR_TYPE_STORAGE_IMAGE : VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 1, VK_SHADER_STAGE_COMPUTE_BIT, nullptr}); + for (size_t i = 0; i < luts.size(); ++i) + lb.push_back({(uint32_t)(imgs.size() + 1 + i), VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, VK_SHADER_STAGE_COMPUTE_BIT, nullptr}); + VkDescriptorSetLayoutCreateInfo dli{VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO}; dli.bindingCount = lb.size(); dli.pBindings = lb.data(); + VkDescriptorSetLayout dl; CK(vkCreateDescriptorSetLayout(dev, &dli, nullptr, &dl)); + VkPipelineLayoutCreateInfo pli{VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO}; pli.setLayoutCount = 1; pli.pSetLayouts = &dl; + VkPipelineLayout pl; CK(vkCreatePipelineLayout(dev, &pli, nullptr, &pl)); + VkPipelineShaderStageRequiredSubgroupSizeCreateInfo rs{VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_REQUIRED_SUBGROUP_SIZE_CREATE_INFO}; rs.requiredSubgroupSize = subgroup; + VkComputePipelineCreateInfo cpi{VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO}; + cpi.stage = {VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO, &rs, VK_PIPELINE_SHADER_STAGE_CREATE_REQUIRE_FULL_SUBGROUPS_BIT, VK_SHADER_STAGE_COMPUTE_BIT, sm, "main", nullptr}; + cpi.layout = pl; + VkPipeline pipe; + auto c0 = std::chrono::steady_clock::now(); + CK(vkCreateComputePipelines(dev, VK_NULL_HANDLE, 1, &cpi, nullptr, &pipe)); + fprintf(stderr, "pipeline compile: %.0f ms\n", std::chrono::duration(std::chrono::steady_clock::now() - c0).count()); + VkDescriptorPoolSize dps[3] = {{VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 8}, {VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 32}, {VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 32}}; + VkDescriptorPoolCreateInfo dpi{VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO}; dpi.maxSets = 1; dpi.poolSizeCount = 3; dpi.pPoolSizes = dps; + VkDescriptorPool dp; CK(vkCreateDescriptorPool(dev, &dpi, nullptr, &dp)); + VkDescriptorSetAllocateInfo dai{VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO}; dai.descriptorPool = dp; dai.descriptorSetCount = 1; dai.pSetLayouts = &dl; + VkDescriptorSet ds; CK(vkAllocateDescriptorSets(dev, &dai, &ds)); + VkDescriptorBufferInfo pbi{pbuf.b, 0, VK_WHOLE_SIZE}; + std::vector ii(imgs.size()); std::vector wr; + wr.push_back({VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET, nullptr, ds, 0, 0, 1, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, nullptr, &pbi, nullptr}); + for (size_t i = 0; i < imgs.size(); ++i) + { + ii[i] = {imgs[i].linear ? linear : nearest, imgs[i].v, VK_IMAGE_LAYOUT_GENERAL}; + wr.push_back({VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET, nullptr, ds, (uint32_t)i + 1, 0, 1, imgs[i].surf ? VK_DESCRIPTOR_TYPE_STORAGE_IMAGE : VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, &ii[i], nullptr, nullptr}); + } + std::vector li(luts.size()); + for (size_t i = 0; i < luts.size(); ++i) + { + li[i] = {luts[i].b, 0, VK_WHOLE_SIZE}; + wr.push_back({VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET, nullptr, ds, (uint32_t)(imgs.size() + 1 + i), 0, 1, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, nullptr, &li[i], nullptr}); + } + vkUpdateDescriptorSets(dev, wr.size(), wr.data(), 0, nullptr); + auto dispatch = [&](int n) { + vkCmdBindPipeline(cb, VK_PIPELINE_BIND_POINT_COMPUTE, pipe); + vkCmdBindDescriptorSets(cb, VK_PIPELINE_BIND_POINT_COMPUTE, pl, 0, 1, &ds, 0, nullptr); + for (int r = 0; r < n; ++r) vkCmdDispatch(cb, gx, gy, 1); + }; + run([&] { dispatch(1); }); + // ---------------------------------------------------------------- read back + run([&] { + for (size_t i = 0; i < imgs.size(); ++i) if (imgs[i].surf) + { + VkBufferImageCopy r{}; r.imageSubresource = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1}; r.imageExtent = {imgs[i].w, imgs[i].h, 1}; + vkCmdCopyImageToBuffer(cb, imgs[i].i, VK_IMAGE_LAYOUT_GENERAL, staging[i].b, 1, &r); + } + }); + for (size_t i = 0; i < imgs.size(); ++i) if (imgs[i].surf) + std::ofstream(imgs[i].out, std::ios::binary).write((const char*)staging[i].map, staging[i].size); + for (auto& [off, f] : bufOut) + std::ofstream(f, std::ios::binary).write((const char*)bufs[off].first.map, bufs[off].first.size); + for (size_t i = 0; i < luts.size(); ++i) if (!lutOut[i].empty()) + std::ofstream(lutOut[i], std::ios::binary).write((const char*)luts[i].map, luts[i].size); + if (const char* gap = getenv("KVK_GAP_US")) + { + // one dispatch per submit with an idle gap, like one kernel inside a frame + std::vector t; + for (int it = 0; it < 600; ++it) + { + VkCommandBufferBeginInfo bi{VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO}; CK(vkBeginCommandBuffer(cb, &bi)); dispatch(1); CK(vkEndCommandBuffer(cb)); + VkSubmitInfo si{VK_STRUCTURE_TYPE_SUBMIT_INFO}; si.commandBufferCount = 1; si.pCommandBuffers = &cb; + auto t0 = std::chrono::steady_clock::now(); + CK(vkQueueSubmit(q, 1, &si, VK_NULL_HANDLE)); CK(vkQueueWaitIdle(q)); + if (it >= 100) t.push_back(std::chrono::duration(std::chrono::steady_clock::now() - t0).count()); + auto until = std::chrono::steady_clock::now() + std::chrono::microseconds(atoi(gap)); + while (std::chrono::steady_clock::now() < until) {} + } + std::sort(t.begin(), t.end()); + printf("%s: %.1f us per submitted dispatch with %s us gaps (median of %zu, min %.1f)\n", argv[1], t[t.size() / 2], gap, t.size(), t[0]); + } + else if (repeats > 1) + { + std::vector t; + for (int it = 0; it < 12; ++it) + { + VkCommandBufferBeginInfo bi{VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO}; CK(vkBeginCommandBuffer(cb, &bi)); dispatch(repeats); CK(vkEndCommandBuffer(cb)); + VkSubmitInfo si{VK_STRUCTURE_TYPE_SUBMIT_INFO}; si.commandBufferCount = 1; si.pCommandBuffers = &cb; + auto t0 = std::chrono::steady_clock::now(); + CK(vkQueueSubmit(q, 1, &si, VK_NULL_HANDLE)); CK(vkQueueWaitIdle(q)); + if (it >= 3) t.push_back(std::chrono::duration(std::chrono::steady_clock::now() - t0).count() / repeats); + } + std::sort(t.begin(), t.end()); + printf("%s: %.1f us per dispatch (median of %zu submits of %d, min %.1f)\n", argv[1], t[t.size() / 2], t.size(), repeats, t[0]); + } + return 0; +} diff --git a/vulkan/ptx2glsl.py b/vulkan/ptx2glsl.py new file mode 100644 index 0000000..584bb7e --- /dev/null +++ b/vulkan/ptx2glsl.py @@ -0,0 +1,968 @@ +#!/usr/bin/env python3 +"""Translate one NVIDIA PTX texture kernel into a Vulkan GLSL compute shader. + +usage: ptx2glsl.py MODULE.ptx KERNEL [--tex OFF:linear|nearest ...] > KERNEL.comp + +Every PTX instruction becomes the GLSL expression of the same operation that kernels/native/ptx2hip.py +emits for HIP, so the shader computes what the translated kernel computes. Registers become locals and +basic blocks become guarded sections of a program-counter loop (GLSL has no goto). The 8-byte handles in +the parameter block become bindings: textures by their parameter offset (T), surfaces (S); +pointers stay 64-bit addresses and are read through buffer references. +""" +import re +import sys + +REG_RE = re.compile(r'%([a-zA-Z]+)(\d+)$') +GT = {'f': 'float', 'r': 'uint', 'rd': 'uint64_t', 'rs': 'uint', 'p': 'bool'} +SREG = {'%tid.x': 'gl_LocalInvocationID.x', '%tid.y': 'gl_LocalInvocationID.y', '%tid.z': 'gl_LocalInvocationID.z', + '%ctaid.x': 'gl_WorkGroupID.x', '%ctaid.y': 'gl_WorkGroupID.y', '%ctaid.z': 'gl_WorkGroupID.z', + '%ntid.x': 'gl_WorkGroupSize.x', '%ntid.y': 'gl_WorkGroupSize.y', + '%nctaid.x': 'gl_NumWorkGroups.x', '%nctaid.y': 'gl_NumWorkGroups.y', '%laneid': 'gl_SubgroupInvocationID'} + + +class Err(Exception): + pass + + +def tokenize(body): + out, i, n = [], 0, len(body) + while i < n: + c = body[i] + if c.isspace(): + i += 1; continue + if c == '{': + out.append(('open', None)); i += 1; continue + if c == '}': + out.append(('close', None)); i += 1; continue + m = re.match(r'(\$?[A-Za-z_][\w$]*):', body[i:]) + if m: + out.append(('label', m.group(1))); i += m.end(); continue + j = body.index(';', i) + out.append(('stmt', ' '.join(body[i:j].split()))) + i = j + 1 + return out + + +def split_operands(s): + ops, depth, cur = [], 0, '' + for ch in s: + if ch in '{[': + depth += 1 + elif ch in '}]': + depth -= 1 + if ch == ',' and depth == 0: + ops.append(cur.strip()); cur = '' + else: + cur += ch + if cur.strip(): + ops.append(cur.strip()) + return ops + + +class T: + def __init__(self, ptx, kernel): + text = re.sub(r'//[^\n]*', '', ptx) + m = re.search(r'\.entry\s+' + re.escape(kernel) + r'\s*\((.*?)\)', text, re.S) + if not m: + raise Err('kernel not found') + pm = re.search(r'\.param\s+\.align\s+(\d+)\s+\.b8\s+(\w+)\[(\d+)\]', m.group(1)) + self.kernel, self.pname, self.psize = kernel, pm.group(2), int(pm.group(3)) + rest = text[m.end():] + mt = re.search(r'\.maxntid\s+(\d+)\s*,\s*(\d+)\s*,\s*(\d+)', rest[:rest.index('{')]) + self.maxntid = [int(x) for x in mt.groups()] if mt else [256, 1, 1] + start = rest.index('{'); depth, k = 0, start + while True: + if rest[k] == '{': + depth += 1 + elif rest[k] == '}': + depth -= 1 + if depth == 0: + break + k += 1 + self.toks = tokenize(rest[start + 1:k]) + self.regs = {} + self.shared = {} # name -> byte base + self.shared_bytes = 0 + self.scopes = [{}] + self.scope_id = 0 + self.blocks = [[]] # lists of lines + self.labels = {} # label -> block index + self.textures, self.surfaces = set(), {} + self.dyn_tex = False + self.fast = False # --fast: no 'precise', the compiler may contract and reassociate + self.collect = False # first pass: gather how every 16/32-bit register is used + self.ctx, self.links, self.width, self.half = {}, [], {}, set() + self.exp_lut = False + self.prescan() + + # ---------------------------------------------------------------- static facts about registers + def prescan(self): + defs = {} + 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) + if len(parts) < 2: + continue + a = split_operands(parts[1]) + for d in re.findall(r'%rd\d+', a[0]) if not a[0].startswith('[') else []: + defs.setdefault(d, []).append((parts[0], a)) + self.pbase = {} # reg -> constant offset from the parameter block + self.handle = {} # reg -> parameter offset its 8-byte value was loaded from + changed = True + while changed: + changed = False + for r, ds in defs.items(): + if len(ds) != 1: + continue + op, a = ds[0] + if r not in self.pbase: + v = None + if op.startswith('mov') and a[1] == self.pname: + v = 0 + elif op in ('add.s64', 'add.u64') and a[1] in self.pbase and re.match(r'-?\d+$', a[2]): + v = self.pbase[a[1]] + int(a[2]) + if v is not None: + self.pbase[r] = v; changed = True + if r not in self.handle: + v = None + if op == 'ld.param.u64': + v = self.param_off(a[1]) + elif op.startswith('mov') and a[1] in self.handle: + v = self.handle[a[1]] + if v is not None: + self.handle[r] = v; changed = True + + def param_off(self, tok): + m = re.match(r'\[\s*([^\]+]+?)\s*(?:\+\s*(-?\d+))?\s*\]$', tok) + base, off = m.group(1), int(m.group(2) or 0) + if base == self.pname: + return off + if base in self.pbase: + return self.pbase[base] + off + return None + + # ---------------------------------------------------------------- operands + def reg(self, tok): + tok = tok.strip() + for sc in reversed(self.scopes): + if tok in sc: + return sc[tok] + m = REG_RE.match(tok) + if not m or m.group(1) not in GT: + raise Err('not a register: ' + tok) + name = m.group(1) + m.group(2) + if m.group(1) in ('r', 'rs'): + self.width[name] = 32 if m.group(1) == 'r' else 16 + self.regs[name] = self.htype(name) or GT[m.group(1)] + return name + + def htype(self, name): + return ('f16vec2' if self.width[name] == 32 else 'float16_t') if name in self.half else None + + def note(self, name, cls): + if self.collect: + self.ctx.setdefault(name, set()).add(cls) + + def hk(self, name): + return 'f16x2' if self.width.get(name) == 32 else 'f16' + + def is_reg(self, tok): + return tok.startswith('%') and tok not in SREG or any(tok in sc for sc in self.scopes) + + def val(self, tok, t): + tok = tok.strip() + if tok in SREG: + return f'uint({SREG[tok]})' if t != 's32' else f'int({SREG[tok]})' + if tok in ('%SP', '%SPL'): + return 'uint64_t(0ul)' # the local depot is its own array; addresses in it start at 0 + if self.is_reg(tok): + n = self.reg(tok); g = self.regs[n] + self.note(n, 'h' if t in ('f16', 'f16x2') else 'u') + if g in ('float16_t', 'f16vec2'): + if t in ('f16', 'f16x2'): + return n + bits = f'HU({n})' if g == 'float16_t' else f'H2U({n})' + if t == 's16': + return f's16({bits})' + if t == 's32': + return f'int({bits})' + if t in ('u64', 'b64'): + return f'uint64_t({bits})' + return bits + if t == 'f32': + return n if g == 'float' else f'uintBitsToFloat({n})' + if t in ('u32', 'b32'): + return f'floatBitsToUint({n})' if g == 'float' else (f'uint({n})' if g == 'uint64_t' else n) + if t == 's32': + return f'floatBitsToInt({n})' if g == 'float' else f'int({n})' + if t in ('u64', 'b64'): + return n if g == 'uint64_t' else f'uint64_t({n})' + if t == 's64': + return f'int64_t({n})' + if t in ('u16', 'b16'): + return n + if t == 's16': + return f's16({n})' + if t == 'f16': + return f'H({n})' + if t == 'f16x2': + return f'H2({n})' + if t == 'pred': + return n + raise Err('type ' + t) + if tok in self.shared: + return f'{self.shared[tok]}u' + if re.match(r'0[fF][0-9A-Fa-f]{8}$', tok): + bits = int(tok[2:], 16) + return f'uintBitsToFloat(0x{bits:08X}u)' if t == 'f32' else f'0x{bits:08X}u' + if re.match(r'-?\d+$', tok) or re.match(r'0x[0-9a-fA-F]+$', tok): + v = int(tok, 0) + if t == 'f32': + raise Err('int immediate as float') + if t in ('u64', 'b64'): + return f'uint64_t({v & 0xffffffffffffffff}ul)' + if t == 's64': + return f'int64_t({v}l)' + if t in ('u16', 'b16'): + return f'{v & 0xffff}u' + if t == 's16': + v &= 0xffff + return f'({v - 0x10000 if v >= 0x8000 else v})' + if t == 'f16': + return f'H({v & 0xffff}u)' + if t == 'f16x2': + return f'H2({v & 0xffffffff}u)' + if t == 's32': + v &= 0xffffffff + return f'int(0x{v:08X}u)' if v >= 1 << 31 else f'{v}' + return f'{v & 0xffffffff}u' + raise Err('operand ' + tok) + + def A(self, dst, expr, st): + n = self.reg(dst); g = self.regs[n] + self.note(n, 'h' if st in ('f16', 'f16x2') else 'u') + if g == 'bool' and st in ('u32', 'u16', 'u64'): + self.E(f'{n} = ({expr}) != 0u;'); return # mov.pred with an immediate + if g in ('float16_t', 'f16vec2'): + if st in ('f16', 'f16x2'): + self.E(f'{n} = {expr};') + else: + self.E(f'{n} = {"H" if g == "float16_t" else "H2"}(uint({expr}));') + return + if st == 'float': + e = expr if g == 'float' else f'floatBitsToUint({expr})' + elif st == 'u32': + e = f'uintBitsToFloat({expr})' if g == 'float' else (f'uint64_t({expr})' if g == 'uint64_t' else f'uint({expr})') + elif st == 'u64': + e = f'uint64_t({expr})' + elif st == 'u16': + e = f'(uint({expr}) & 0xffffu)' + elif st == 'f16': + e = f'HU({expr})' + elif st == 'f16x2': + e = f'H2U({expr})' + elif st == 'bool': + e = f'({expr})' + else: + raise Err('src type') + self.E(f'{n} = {e};') + + def E(self, s): + self.blocks[-1].append(s) + + # ---------------------------------------------------------------- driver + def translate(self): + for kind, text in self.toks: + if kind == 'open': + self.scopes.append({}); continue + if kind == 'close': + self.scopes.pop(); continue + if kind == 'label': + self.blocks.append([]) + self.labels[text] = len(self.blocks) - 1 + continue + self.stmt(text) + + def stmt(self, s): + if s.startswith('.reg'): + m = re.match(r'\.reg\s+\.(\w+)\s+(.*)$', s) + ty, names = m.groups() + if '<' in names: + return + self.scope_id += 1 + for nm in names.split(','): + nm = nm.strip() + c = f'l{self.scope_id}_' + re.sub(r'\W', '_', nm) + self.scopes[-1][nm] = c + if ty in ('f16', 'b16', 'f16x2', 'b32', 'u32', 's32'): + self.width[c] = 16 if ty in ('f16', 'b16') else 32 + self.regs[c] = self.htype(c) if c in self.half else {'f16': 'uint', 'b16': 'uint', 'f32': 'float', 'b32': 'uint', 'pred': 'bool', 'f16x2': 'uint', + 'u32': 'uint', 's32': 'uint', 'b64': 'uint64_t', 'u64': 'uint64_t'}[ty] + return + if s.startswith('.local'): + m = re.match(r'\.local\s+\.align\s+(\d+)\s+\.b8\s+(\w+)\[(\d+)\]', s) + self.local_bytes = max(getattr(self, 'local_bytes', 0), int(m.group(3))) + return + if re.match(r'mov\.u64 %SPL?, __local_depot', s): + return + if s.startswith('.shared'): + m = re.match(r'\.shared\s+\.align\s+(\d+)\s+\.b8\s+(\w+)\[(\d+)\]', s) + self.shared[m.group(2)] = self.shared_bytes + self.shared_bytes += (int(m.group(3)) + 3) & ~3 + return + if s.startswith('.'): + return + pred = None + m = re.match(r'@(!?)(%p\d+)\s+(.*)$', s) + if m: + pred = ('!' if m.group(1) else '') + self.reg(m.group(2)); s = m.group(3) + parts = s.split(None, 1) + op = parts[0] + a = split_operands(parts[1]) if len(parts) > 1 else [] + if op.startswith('bra'): + tgt = ('L', a[0]) + if pred is not None: + self.E(('BR', pred, tgt)) # conditional: falls through into a new block + else: + self.E(('BR', None, tgt)) + self.blocks.append([]) + return + if op == 'ret': + self.E(('BR', pred, ('END', None))) + self.blocks.append([]) + return + if pred is not None: + self.E(f'if ({pred}) {{') + self.op(op, a) + self.E('}') + return + self.op(op, a) + + def cmp(self, mods, isf): + table = {'lt': '{0} < {1}', 'gt': '{0} > {1}', 'le': '{0} <= {1}', 'ge': '{0} >= {1}', 'eq': '{0} == {1}', 'ne': '{0} != {1}', + 'ltu': '!({0} >= {1})', 'gtu': '!({0} <= {1})', 'leu': '!({0} > {1})', 'geu': '!({0} < {1})', + 'equ': '!({0} != {1})', 'neu': '!({0} == {1})'} + for k in table: + if k in mods: + if isf and k == 'ne': + return '({0} < {1} || {0} > {1})' + if isf and k == 'equ': + return '!({0} < {1} || {0} > {1})' + return table[k] + raise Err('compare') + + def texref(self, tok): + m = re.match(r'\[\s*(%rd\d+)\s*,\s*\{([^}]*)\}\s*\]', tok) + return m.group(1), [x.strip() for x in m.group(2).split(',')] + + def op(self, op, a): + o = op.split('.'); base = o[0]; mods = set(o[1:]) + E, A, V = self.E, self.A, self.val + ft = next((t for t in ('f32', 'f16x2', 'f16') if t in mods), None) + if base == 'bar': + E('memoryBarrierShared(); barrier();'); return + if base == 'mov': + dst, src = a + if dst.startswith('{'): + n = [x.strip() for x in dst[1:-1].split(',')] + sr = self.reg(src) if self.is_reg(src) else None + if self.collect and sr: + for x in [sr] + [self.reg(y) for y in n]: + self.ctx.setdefault(x, set()).add('h') + return + if sr and (sr in self.half or any(self.reg(y) in self.half for y in n)): + A(n[0], f'{V(src, "f16x2")}.x', 'f16'); A(n[1], f'{V(src, "f16x2")}.y', 'f16'); return + A(n[0], f'{V(src, "u32")} & 0xffffu', 'u16'); A(n[1], f'{V(src, "u32")} >> 16', 'u16'); return + if src.startswith('{'): + n = [x.strip() for x in src[1:-1].split(',')] + dr = self.reg(dst) + if self.collect: + for x in [dr] + [self.reg(y) for y in n if self.is_reg(y)]: + self.ctx.setdefault(x, set()).add('h') + return + if dr in self.half: + A(dst, f'f16vec2({V(n[0], "f16")}, {V(n[1], "f16")})', 'f16x2'); return + A(dst, f'({V(n[0], "u16")} & 0xffffu) | ({V(n[1], "u16")} << 16)', 'u32'); return + if self.is_reg(dst) and self.reg(dst) in self.width and not (mods & {'f32', 'u64', 'b64', 's64'}): + dr = self.reg(dst) + srcreg = self.reg(src) if self.is_reg(src) and src not in SREG else None + isimm = bool(re.match(r'-?\d+$', src) or re.match(r'0x[0-9a-fA-F]+$', src)) + if self.collect and srcreg and srcreg in self.width: + self.links.append((dr, srcreg)); return + if self.collect and isimm: + return + if dr in self.half and (srcreg or isimm): + A(dst, V(src, self.hk(dr)), self.hk(dr)); return + if src == self.pname: + A(dst, '0ul', 'u64'); return + if 'f32' in mods: + A(dst, V(src, 'f32'), 'float'); return + if mods & {'u64', 'b64', 's64'}: + A(dst, V(src, 'u64'), 'u64'); return + if mods & {'u16', 'b16', 's16'}: + A(dst, V(src, 'u16'), 'u16'); return + A(dst, V(src, 'u32'), 'u32'); return + # ------------------------------------------------ f32 + if ft == 'f32' and base in ('add', 'sub', 'mul'): + if 'rz' in mods: + A(a[0], f'round_half_away({V(a[1], "f32")})', 'float'); return # the roundf idiom: + copysign(0.5), then trunc + A(a[0], f'{V(a[1], "f32")} {"+-*"[["add", "sub", "mul"].index(base)]} {V(a[2], "f32")}', 'float'); return + if ft == 'f32' and base == 'fma': + A(a[0], f'fma({V(a[1], "f32")}, {V(a[2], "f32")}, {V(a[3], "f32")})', 'float'); return + if ft == 'f32' and base in ('max', 'min'): + A(a[0], f'{base}({V(a[1], "f32")}, {V(a[2], "f32")})', 'float'); return + if ft == 'f32' and base == 'neg': + A(a[0], f'-{V(a[1], "f32")}', 'float'); return + if ft == 'f32' and base == 'abs': + A(a[0], f'abs({V(a[1], "f32")})', 'float'); return + if ft == 'f32' and base in ('rcp', 'ex2', 'lg2', 'rsqrt', 'sqrt'): + fn = {'rcp': 'rcp_a', 'ex2': 'exp2', 'lg2': 'log2', 'rsqrt': 'inversesqrt', 'sqrt': 'sqrt'}[base] + A(a[0], f'{fn}({V(a[1], "f32")})', 'float'); return + if ft == 'f32' and base == 'div': + A(a[0], f'{V(a[1], "f32")} * rcp_a({V(a[2], "f32")})', 'float'); return + # ------------------------------------------------ f16 + if ft == 'f16x2' and base in ('add', 'sub', 'mul', 'fma'): + x = [V(t, 'f16x2') for t in a[1:]] + e = f'fma({x[0]}, {x[1]}, {x[2]})' if base == 'fma' else f'({x[0]} {"+-*"[["add", "sub", "mul"].index(base)]} {x[1]})' + if 'sat' in mods: + e = f'min(max({e}, f16vec2(0.0hf)), f16vec2(1.0hf))' + A(a[0], e, 'f16x2'); return + if ft == 'f16x2' and base == 'abs': + A(a[0], f'abs({V(a[1], "f16x2")})', 'f16x2'); return + if ft == 'f16x2' and base == 'ex2': + if self.exp_lut: + A(a[0], f'ex2_lut({V(a[1], "u32")})', 'u32'); return + A(a[0], f'f16vec2(exp2(vec2({V(a[1], "f16x2")})))', 'f16x2'); return + if ft in ('f16', 'f16x2') and base in ('min', 'max'): + A(a[0], f'{base}({V(a[1], ft)}, {V(a[2], ft)})', ft); return + if ft == 'f16' and base in ('add', 'sub', 'mul', 'fma'): + x = [V(t, 'f16') for t in a[1:]] + e = f'fma({x[0]}, {x[1]}, {x[2]})' if base == 'fma' else f'({x[0]} {"+-*"[["add", "sub", "mul"].index(base)]} {x[1]})' + if 'sat' in mods: + e = f'min(max({e}, 0.0hf), 1.0hf)' + A(a[0], e, 'f16'); return + if base == 'set' and 'f16x2' in mods and 'u32' in mods: + c = self.cmp(mods, True); x, y = V(a[1], 'f16x2'), V(a[2], 'f16x2') + A(a[0], f'(({c.format(x + ".x", y + ".x")}) ? 0xffffu : 0u) | (({c.format(x + ".y", y + ".y")}) ? 0xffff0000u : 0u)', 'u32'); return + # ------------------------------------------------ compares / selects + if base == 'setp': + isf = bool(mods & {'f32', 'f16'}) + t = next(t for t in ('f32', 'f16', 's32', 'u32', 'b32', 's16', 'u16', 'b16', 's64', 'u64') if t in mods) + if '|' in a[0]: + raise Err('setp form') + A(a[0], self.cmp(mods, isf).format(V(a[1], t), V(a[2], t)), 'bool'); return + if base == 'selp' and not (mods & {'f32', 'u64', 'b64', 's64'}) and self.reg(a[0]) in self.width: + dr = self.reg(a[0]) + srcs = [self.reg(x) for x in a[1:3] if self.is_reg(x) and x not in SREG] + if self.collect: + self.links += [(dr, x) for x in srcs if x in self.width] + V(a[3], 'pred') + if any(x not in self.width for x in srcs): + self.note(dr, 'u') + return + if dr in self.half: + k = self.hk(dr) + A(a[0], f'{V(a[3], "pred")} ? {V(a[1], k)} : {V(a[2], k)}', k); return + if base == 'selp': + t = 'f32' if 'f32' in mods else 'u16' if mods & {'u16', 'b16', 's16'} else 'u64' if mods & {'u64', 'b64', 's64'} else 'u32' + A(a[0], f'{V(a[3], "pred")} ? {V(a[1], t)} : {V(a[2], t)}', {'f32': 'float', 'u16': 'u16', 'u64': 'u64', 'u32': 'u32'}[t]); return + if 'pred' in mods and base in ('and', 'or', 'xor', 'not'): + if base == 'not': + A(a[0], f'!{V(a[1], "pred")}', 'bool'); return + A(a[0], f'{V(a[1], "pred")} {{"and": "&&", "or": "||", "xor": "!="}}'.replace('{"and": "&&", "or": "||", "xor": "!="}', {'and': '&&', 'or': '||', 'xor': '!='}[base]) + f' {V(a[2], "pred")}', 'bool'); return + # ------------------------------------------------ integer + if base in ('and', 'or', 'xor', 'not', 'shl', 'shr', 'add', 'sub', 'mul', 'mad', 'min', 'max', 'div', 'rem', 'neg', 'abs') and ft is None: + w = 64 if mods & {'u64', 's64', 'b64'} else 16 if mods & {'u16', 's16', 'b16'} else 32 + signed = bool(mods & {'s32', 's64', 's16'}) + ut, st = {64: ('u64', 's64'), 32: ('u32', 's32'), 16: ('u16', 's16')}[w] + if base == 'mul' and 'wide' in mods: + if 's32' in mods: + A(a[0], f'uint64_t(int64_t({V(a[1], "s32")}) * int64_t({V(a[2], "s32")}))', 'u64'); return + if 'u32' in mods: + A(a[0], f'uint64_t({V(a[1], "u32")}) * uint64_t({V(a[2], "u32")})', 'u64'); return + if 'u16' in mods: + A(a[0], f'{V(a[1], "u16")} * {V(a[2], "u16")}', 'u32'); return + raise Err(op) + if base == 'not': + A(a[0], f'~{V(a[1], ut)}', ut); return + if base == 'neg': + A(a[0], f'0u - {V(a[1], ut)}', ut); return + if base == 'shl': + A(a[0], f'{V(a[1], ut)} << ({V(a[2], "u32")} & {w - 1}u)', ut); return + if base == 'shr': + x = V(a[1], st if signed else ut) + A(a[0], f'{x} >> ({V(a[2], "u32")} & {w - 1}u)', ut); return + if base in ('and', 'or', 'xor'): + A(a[0], f'{V(a[1], ut)} {{"and": "&", "or": "|", "xor": "^"}}'.replace('{"and": "&", "or": "|", "xor": "^"}', {'and': '&', 'or': '|', 'xor': '^'}[base]) + f' {V(a[2], ut)}', ut); return + if base in ('add', 'sub'): + A(a[0], f'{V(a[1], ut)} {"+" if base == "add" else "-"} {V(a[2], ut)}', ut); return + if base == 'mul' and 'hi' in mods: + if signed: + A(a[0], f'uint(int((int64_t({V(a[1], "s32")}) * int64_t({V(a[2], "s32")})) >> 32))', 'u32'); return + A(a[0], f'uint((uint64_t({V(a[1], "u32")}) * uint64_t({V(a[2], "u32")})) >> 32)', 'u32'); return + if base == 'mul': + A(a[0], f'{V(a[1], ut)} * {V(a[2], ut)}', ut); return + if base == 'mad': + A(a[0], f'{V(a[1], ut)} * {V(a[2], ut)} + {V(a[3], ut)}', ut); return + if base == 'abs': + A(a[0], f'uint(abs({V(a[1], st)}))', ut); return + if base == 'neg' and len(a) == 2: + A(a[0], f'uint(-{V(a[1], st)})', ut); return + x, y = V(a[1], st if signed else ut), V(a[2], st if signed else ut) + if base in ('min', 'max'): + A(a[0], f'{base}({x}, {y})', ut); return + if base == 'div': + A(a[0], f'{x} / {y}', ut); return + if base == 'rem': + A(a[0], f'{x} % {y}', ut); return + if base == 'bfi': + A(a[0], f'bfi32({V(a[1], "u32")}, {V(a[2], "u32")}, {V(a[3], "u32")}, {V(a[4], "u32")})', 'u32'); return + if base == 'cvt': + return self.cvt(op, o[1:], a) + if base == 'cvta': + A(a[0], V(a[1], 'u64'), 'u64'); return + if base == 'ld': + return self.ld(o[1:], a) + if base == 'st': + return self.st(o[1:], a) + if base in ('tex', 'tld4'): + dst = [x.strip() for x in a[0][1:-1].split(',')] + h, c = self.texref(a[1]) + off = self.handle.get(h) + if off is None: + raise Err('texture handle of ' + h + ' is not a single parameter load') + self.textures.add(off) + if base == 'tld4': + E(f'{{ vec4 t_ = textureGather(T{off}, vec2({V(c[0], "f32")}, {V(c[1], "f32")}), 0);') + elif 'base' in mods: + E(f'{{ vec4 t_ = texelFetch(T{off}, ivec2({V(c[0], "s32")}, {V(c[1], "s32")}), 0);') + elif 'level' in mods: + E(f'{{ vec4 t_ = textureLod(T{off}, vec2({V(c[0], "f32")}, {V(c[1], "f32")}), {V(a[2], "f32")});') + else: + raise Err(op) + for i, d in enumerate(dst): + if 'f16' in o: + A(d, f'float16_t(t_[{i}])', 'f16') # the texel as a half float + else: + A(d, f't_[{i}]', 'float') + E('}') + return + if base == 'sust': + h, c = self.texref(a[0]) + off = self.handle.get(h) + if off is None: + raise Err('surface handle of ' + h) + v = [x.strip() for x in a[1][1:-1].split(',')] + if op == 'sust.p.2d.v4.b32.zero': + self.surfaces.setdefault(off, 'p'); assert self.surfaces[off] == 'p' + E(f'sust_p(S{off}, {V(c[0], "s32")}, {V(c[1], "s32")}, vec4({", ".join(V(x, "f32") for x in v)}));'); return + if op == 'sust.b.2d.b32.zero': + self.surfaces.setdefault(off, 'b'); assert self.surfaces[off] == 'b' + E(f'sust_b(S{off}, {V(c[0], "s32")}, {V(c[1], "s32")}, {V(v[0], "u32")});'); return + if op == 'sust.b.2d.v2.b16.zero': # reading only: two 16-bit values at a byte-addressed position + self.surfaces.setdefault(off, 'b') + E(f'sust_b2x16(S{off}, {V(c[0], "s32")}, {V(c[1], "s32")}, {V(v[0], "u16")}, {V(v[1], "u16")});'); return + raise Err(op) + if op == 'shfl.sync.bfly.b32': + d = a[0].split('|') + E('{ bool ib_;') + A(d[0], f'shfl_bfly({V(a[1], "u32")}, {V(a[2], "s32")}, {V(a[3], "u32")}, ib_)', 'u32') + if len(d) > 1: + A(d[1], 'ib_', 'bool') + E('}') + return + if op == 'shfl.sync.idx.b32': + d = a[0].split('|') + A(d[0], f'subgroupShuffle({V(a[1], "u32")}, {V(a[2], "u32")} & 31u)', 'u32') + if len(d) > 1: + A(d[1], 'true', 'bool') + return + if op == 'tanh.approx.f16': + A(a[0], f'tanh({V(a[1], "f16")})', 'f16'); return + if op == 'prmt.b32': + A(a[0], f'prmt({V(a[1], "u32")}, {V(a[2], "u32")}, {V(a[3], "u32")})', 'u32'); return + if op.startswith('mma.sync'): + # not executable here: the kernel is translated for reading only (tensor-core fragments are not emulated) + self.E(f'// {op} ' + ' '.join(a)) + for d in a[0][1:-1].split(','): + A(d.strip(), '0u', 'u32') + return + if op == 'vote.sync.ballot.b32': + # the warp is the subgroup (32 lanes in thread order), as for the shuffles + A(a[0], f'subgroupBallot({V(a[1], "pred")}).x', 'u32'); return + raise Err('unhandled ' + op) + + def cvt(self, op, o, a): + A, V = self.A, self.val + mods = set(o) + types = [t for t in o if re.match(r'^(f|s|u|b)\d+(x2)?$', t)] + types = [t for t in o if re.match(r'^(f|s|u|b)\d+(x2)?$|^e4m3x2$', t)] + dt, st = types[0], types[1] + # FP8 pairs (reading only: the helpers are not defined) + if dt == 'f16x2' and st == 'e4m3x2': + A(a[0], f'e4m3x2_to_f16x2({V(a[1], "u16")})', 'f16x2'); return + if dt == 'e4m3x2' and st == 'f16x2': + A(a[0], f'f16x2_to_e4m3x2({V(a[1], "f16x2")})', 'u16'); return + if dt == 'f32' and st == 'f32': + x = V(a[1], 'f32') + if 'sat' in mods: + A(a[0], f'min(max({x}, 0.0), 1.0)', 'float'); return + fn = 'trunc' if 'rzi' in mods else 'floor' if 'rmi' in mods else 'roundEven' if 'rni' in mods else 'ceil' if 'rpi' in mods else '' + A(a[0], f'{fn}({x})', 'float'); return + if dt == 's32' and st == 'f32': + x = V(a[1], 'f32') + fn = '' if 'rzi' in mods else 'floor' if 'rmi' in mods else 'roundEven' + A(a[0], f'uint(int({fn}({x})))', 'u32'); return + if dt == 'u32' and st == 'f32': + x = V(a[1], 'f32') + fn = 'trunc' if 'rzi' in mods else 'floor' if 'rmi' in mods else 'roundEven' + A(a[0], f'uint(clamp({fn}({x}), 0.0, 4294967040.0))', 'u32'); return # CUDA saturates + if dt == 'f32' and st == 's32': + A(a[0], f'float({V(a[1], "s32")})', 'float'); return + if dt == 'f32' and st == 'u32': + A(a[0], f'float({V(a[1], "u32")})', 'float'); return + if dt == 'f32' and st == 'f16': + A(a[0], f'float({V(a[1], "f16")})', 'float'); return + if dt == 'f16' and st == 'f32': + A(a[0], f'float16_t({V(a[1], "f32")})', 'f16'); return + if dt == 'f16x2' and st == 'f32': + A(a[0], f'f16vec2(float16_t({V(a[2], "f32")}), float16_t({V(a[1], "f32")}))', 'f16x2'); return + if dt == 'f16' and st == 'f16': + fn = 'roundEven' if 'rni' in mods else 'floor' + A(a[0], f'{fn}({V(a[1], "f16")})', 'f16'); return + if dt == 's32' and st == 'f16': + fn = 'roundEven' if 'rni' in mods else 'floor' + A(a[0], f'uint(int(float({fn}({V(a[1], "f16")}))))', 'u32'); return + if dt == 'f16' and st == 's32': + A(a[0], f'float16_t({V(a[1], "s32")})', 'f16'); return + if dt == 's64' and st == 's32': + A(a[0], f'uint64_t(int64_t({V(a[1], "s32")}))', 'u64'); return + if dt == 'u64' and st == 'u32': + A(a[0], f'uint64_t({V(a[1], "u32")})', 'u64'); return + if dt in ('u16', 's16') and st in ('u32', 's32'): + A(a[0], V(a[1], 'u32'), 'u16'); return + if dt in ('u32', 's32') and st in ('u16',): + A(a[0], V(a[1], 'u16'), 'u32'); return + if dt == 's32' and st == 's16': + A(a[0], f'uint({V(a[1], "s16")})', 'u32'); return + raise Err('cvt ' + op) + + def addr(self, tok): + m = re.match(r'\[\s*([^\]+]+?)\s*(?:\+\s*(-?\d+))?\s*\]$', tok) + return m.group(1), int(m.group(2) or 0) + + def ld(self, o, a): + A, V, E = self.A, self.val, self.E + mods = set(o) + ty = next(t for t in o if re.match(r'^(f|s|u|b)\d+$', t)) + vec = next((int(t[1]) for t in o if re.match(r'^v\d$', t)), 1) + size = int(ty[1:]) // 8 + dsts = [x.strip() for x in a[0][1:-1].split(',')] if a[0].startswith('{') else [a[0]] + base, off = self.addr(a[1]) + for i, d in enumerate(dsts): + if 'param' in mods: + po = self.param_off(a[1]) + if po is None: + raise Err('param address ' + a[1]) + e = f'PARAM{size * 8}({po + i * size}u)' + elif 'shared' in mods: + e = f'SH{size * 8}({V(base, "u32")} + {off + i * size}u)' + elif 'local' in mods: + e = f'loc[(uint({V(base, "u64")}) + {off + i * size}u) >> 2]' + else: + e = f'G{size * 8}({V(base, "u64")} + uint64_t({off + i * size}l))' + if size == 8: + A(d, e, 'u64') + elif size == 2: + A(d, e, 'u16') + elif size == 1: + A(d, e, 'u16') + else: + A(d, e, 'u32') + + def st(self, o, a): + V, E = self.val, self.E + mods = set(o) + ty = next(t for t in o if re.match(r'^(f|s|u|b)\d+$', t)) + size = int(ty[1:]) // 8 + base, off = self.addr(a[0]) + if 'shared' in mods: + # vector stores go word by word; 16-bit pairs share a word (the kernels keep them 4-byte aligned) + vals = [x.strip() for x in a[1][1:-1].split(',')] if a[1].startswith('{') else [a[1]] + if size == 4: + for i, v in enumerate(vals): + E(f'sh[({V(base, "u32")} + {off + 4 * i}u) >> 2] = {V(v, "u32")};') + return + if size == 2 and len(vals) % 2 == 0: + for i in range(0, len(vals), 2): + E(f'sh[({V(base, "u32")} + {off + 2 * i}u) >> 2] = (uint({V(vals[i], "u16")}) & 0xFFFFu) | (uint({V(vals[i + 1], "u16")}) << 16);') + return + if 'local' in mods and size == 4: + vals = [x.strip() for x in a[1][1:-1].split(',')] if a[1].startswith('{') else [a[1]] + for i, v in enumerate(vals): + E(f'loc[(uint({V(base, "u64")}) + {off + 4 * i}u) >> 2] = {V(v, "u32")};') + return + if 'global' in mods and size == 2 and not a[1].startswith('{'): + E(f'W16({V(base, "u64")} + uint64_t({off}l)).v[0] = uint16_t({V(a[1], "u16")});'); return + if 'global' in mods and size == 2 and a[1].startswith('{'): + vals = [x.strip() for x in a[1][1:-1].split(',')] + E(f'{{ W16 w_ = W16({V(base, "u64")} + uint64_t({off}l)); ' + ' '.join(f'w_.v[{i}] = uint16_t({V(v, "u16")});' for i, v in enumerate(vals)) + ' }'); return + if 'global' in mods and size == 4: + E(f'{{ W16 w_ = W16({V(base, "u64")} + uint64_t({off}l)); uint v_ = {V(a[1], "u32")}; w_.v[0] = uint16_t(v_ & 0xffffu); w_.v[1] = uint16_t(v_ >> 16); }}'); return + raise Err('store ' + '.'.join(o)) + + # ---------------------------------------------------------------- output + def emit(self, tex_filter): + nb = len(self.blocks) + out = [] + w = out.append + w('#version 460') + # scripts/check_redistributable.py refuses files carrying this line: they are NVIDIA's code. + w(f'// Generated by vulkan/ptx2glsl.py from NVIDIA PTX ({self.kernel}); do not commit or redistribute.') + for e in ('GL_EXT_shader_explicit_arithmetic_types', 'GL_EXT_buffer_reference', 'GL_EXT_buffer_reference2', 'GL_EXT_scalar_block_layout', + 'GL_KHR_shader_subgroup_basic', 'GL_KHR_shader_subgroup_shuffle', 'GL_KHR_shader_subgroup_ballot', 'GL_EXT_samplerless_texture_functions'): + w(f'#extension {e} : require') + bx = 16 if self.maxntid[0] == 256 else self.maxntid[0] + by = self.maxntid[0] // bx + w(f'layout(local_size_x = {bx}, local_size_y = {by}) in;') + w(f'layout(set = 0, binding = 0, std430) readonly buffer ParamBlock {{ uint w[{(self.psize + 3) // 4}]; }} param;') + b = 1 + for off in sorted(self.textures): + w(f'layout(set = 0, binding = {b}) uniform sampler2D T{off}; // {tex_filter.get(off, "?")}') + b += 1 + for off in sorted(self.surfaces): + fmt = 'r16ui' if self.surfaces[off] == 'b' else 'rgba16f' + typ = 'uimage2D' if self.surfaces[off] == 'b' else 'image2D' + w(f'layout(set = 0, binding = {b}, {fmt}) uniform writeonly {typ} S{off};') + b += 1 + if self.exp_lut: + w(f'layout(set = 0, binding = {b}, scalar) readonly buffer ExpLut {{ uint16_t e[65536]; }} explut;') + w('uint ex2_lut(uint u) { return uint(explut.e[u & 0xffffu]) | (uint(explut.e[u >> 16]) << 16); }') + w(f'shared uint sh[{max(self.shared_bytes // 4, 1)}];') + w(RUNTIME) + w('void main()\n{') + if getattr(self, 'local_bytes', 0): + w(f' uint loc[{(self.local_bytes + 3) // 4}];') + for g in ('float', 'float16_t', 'f16vec2', 'uint', 'uint64_t', 'bool'): + names = sorted(n for n, t in self.regs.items() if t == g) + for i in range(0, len(names), 24): + init = {'float': '0.0', 'float16_t': '0.0hf', 'f16vec2': 'f16vec2(0.0hf)', 'uint': '0u', 'uint64_t': '0ul', 'bool': 'false'}[g] + w(f' {"precise " if g in ("float", "float16_t", "f16vec2") and not self.fast else ""}{g} ' + ', '.join(f'{n} = {init}' for n in names[i:i + 24]) + ';') + # control flow graph over the blocks; END is block nb + succ = [] + for i, blk in enumerate(self.blocks): + br = blk[-1] if blk and isinstance(blk[-1], tuple) else None + if br is None: + succ.append([min(i + 1, nb)]) + else: + tgt = nb if br[2][0] == 'END' else self.labels[br[2][1]] + succ.append([tgt] if br[1] is None else [tgt, i + 1]) + # ---- structured emission: if/else arms that rejoin at the branch's immediate post-dominator -------------- + END = nb + # reachable nodes, back edges (to a node on the DFS stack) + color, back = {}, set() + st = [(0, iter(succ[0]))]; color[0] = 1 + while st: + n, it = st[-1] + for m in it: + if m == END: + continue + if color.get(m) == 1: + back.add((n, m)) + elif m not in color: + color[m] = 1; st.append((m, iter(succ[m]))); break + else: + color[n] = 2; st.pop() + nodes = sorted(color) + # each loop header h gets a synthetic "continue" node: back edges go there, and it flows to the loop's exit + headers = sorted({h for _, h in back}) + cont = {h: END + 1 + i for i, h in enumerate(headers)} + g = {n: [(cont[m] if (n, m) in back else m) for m in succ[n]] for n in nodes} + for h in headers: + latches = [n for n, m in back if m == h] + exits = {m for n in latches for m in succ[n] if (n, m) not in back} + if len(exits) != 1: + raise Err('loop without a single latch exit') + g[cont[h]] = [exits.pop()] + g[END] = [] + allnodes = list(g) + # immediate post-dominators (iterative, on the acyclic graph g) + rpo, seen = [], set() + def dfs(n): + stack = [(n, iter(g[n]))]; seen.add(n) + while stack: + x, it = stack[-1] + for m in it: + if m not in seen: + seen.add(m); stack.append((m, iter(g[m]))); break + else: + rpo.append(x); stack.pop() + dfs(0) + topo = rpo[::-1] # entry first + index = {n: i for i, n in enumerate(topo)} + ipdom = {END: END} + def inter(x, y): + while x != y: + while index[x] < index[y]: + x = ipdom[x] + while index[y] < index[x]: + y = ipdom[y] + return x + for n in reversed(topo): # END first + if n == END: + continue + d = None + for m in g[n]: + d = m if d is None else inter(d, m) + ipdom[n] = d + emitted = [0] + def body(n, ind): + if n > END: + w(ind + f'again_{n - END - 1} = true;'); return + for line in self.blocks[n]: + if not isinstance(line, tuple): + w(ind + line) + emitted[0] += 1 + def region(n, stop, ind): + while n != stop: + if n in cont and (n, 'open') not in region.active: + # loop: body once per iteration until no back edge is taken + e = g[cont[n]][0] + w(ind + f'bool again_{cont[n] - END - 1};'); w(ind + 'do\n' + ind + '{') + w(ind + f' again_{cont[n] - END - 1} = false;') + region.active.add((n, 'open')) + region(n, e, ind + ' ') + region.active.discard((n, 'open')) + w(ind + f'}} while (again_{cont[n] - END - 1});') + n = e + continue + body(n, ind) + sc = g[n] + if len(sc) == 1: + n = sc[0]; continue + br = self.blocks[n][-1] + m = ipdom[n] + t_empty, f_empty = sc[0] == m, sc[1] == m + if t_empty and f_empty: + pass + elif f_empty: + w(ind + f'if ({br[1]})\n' + ind + '{'); region(sc[0], m, ind + ' '); w(ind + '}') + elif t_empty: + w(ind + f'if (!({br[1]}))\n' + ind + '{'); region(sc[1], m, ind + ' '); w(ind + '}') + else: + w(ind + f'if ({br[1]})\n' + ind + '{'); region(sc[0], m, ind + ' ') + w(ind + '}\n' + ind + 'else\n' + ind + '{'); region(sc[1], m, ind + ' '); w(ind + '}') + n = m + region.active = set() + import sys as _s + _s.setrecursionlimit(20000) + region(0, END, ' ') + w('}') + sys.stderr.write(f'reachable blocks={len(nodes)} emitted block copies={emitted[0]} loops={len(headers)}\n') + return '\n'.join(out) + '\n' + + +RUNTIME = r''' +layout(buffer_reference, scalar, buffer_reference_align = 2) readonly buffer R16 { uint16_t v[4]; }; +layout(buffer_reference, scalar, buffer_reference_align = 2) writeonly buffer W16 { uint16_t v[4]; }; +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 G16(uint64_t a) { return uint(R16(a).v[0]); } +uint G32(uint64_t a) { R16 r = R16(a); return uint(r.v[0]) | (uint(r.v[1]) << 16); } +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)); +} +// prmt.b32 (default mode): byte i of the result is byte (selector >> 4i) & 7 of the eight bytes b:a +uint prmt(uint a, uint b, uint sel) +{ + uint r = 0u; + for (uint i = 0u; i < 4u; ++i) { const uint k = (sel >> (4u * i)) & 7u; r |= (((k < 4u ? a : b) >> (8u * (k & 3u))) & 0xffu) << (8u * i); } + return r; +} +// shfl.sync.bfly within a 32-thread warp (ZLUDA's SHFL_SYNC_IMPL) +uint shfl_bfly(uint v, int delta, uint opts, out bool in_bounds) +{ + int section_mask = int((opts >> 8) & 31u), warp_end = int(opts & 31u); + int self = int(gl_SubgroupInvocationID & 31u); + int sub_end = (section_mask & self) | (~section_mask & warp_end); + int idx = self ^ (delta & 31); + bool oob = idx > sub_end; + if (oob) idx = self; + in_bounds = !oob; + return subgroupShuffle(v, (gl_SubgroupInvocationID & 32u) | uint(idx)); +} +void sust_p(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); +} +// byte-addressed 32-bit store into a 16-bit single-channel surface: two texels +void sust_b(writeonly uimage2D s, int xb, int y, uint v) +{ + ivec2 n = imageSize(s); + int x = xb >> 1; + if (uint(y) < uint(n.y)) + { + if (uint(x) < uint(n.x)) imageStore(s, ivec2(x, y), uvec4(v & 0xffffu)); + if (uint(x + 1) < uint(n.x)) imageStore(s, ivec2(x + 1, y), uvec4(v >> 16)); + } +} +''' + +if __name__ == '__main__': + filt = {} + args = sys.argv[1:] + while '--tex' in args: + i = args.index('--tex'); k, v = args[i + 1].split(':'); filt[int(k)] = v; del args[i:i + 2] + lut = '--exp-lut' in args + if lut: + args.remove('--exp-lut') + fast = '--fast' in args + if fast: + args.remove('--fast') + src = open(args[0]).read() + t = T(src, args[1]) + t.exp_lut = lut + if '--no-half-types' not in args: + t.collect = True + t.translate() + parent = {} + def find(x): + while parent.setdefault(x, x) != x: + parent[x] = parent[parent[x]]; x = parent[x] + return x + for x, y in t.links: + if t.width.get(x) == t.width.get(y): + parent[find(x)] = find(y) + else: + t.ctx.setdefault(x, set()).add('u'); t.ctx.setdefault(y, set()).add('u') + groups = {} + for x in set(t.ctx) | set(parent): + groups.setdefault(find(x), set()).update(t.ctx.get(x, set())) + half = {x for x in set(t.ctx) | set(parent) if x in t.width and groups[find(x)] == {'h'}} + sys.stderr.write(f'half-typed registers: {len(half)} of {len(t.width)}\n') + lut2 = t.exp_lut + t = T(src, args[1]); t.exp_lut = lut2; t.half = half + t.fast = fast + t.translate() + sys.stdout.write(t.emit(filt)) + sys.stderr.write(f'blocks={len(t.blocks)} textures={sorted(t.textures)} surfaces={t.surfaces} shared={t.shared_bytes}\n')