Engine: native presets K and M in the D3D12 shim, wider game coverage
A standalone Vulkan inference engine (engine/) for DLSS presets K and M that runs from weights extracted from the user's own nvngx_dlss.dll, with d4r's own input, network and reconstruction shaders or the native HIP network layers, recorded into the game's command list by the shim. Opt-in ([Engine] Enabled / D4R_ENGINE=1); anything uncovered falls back to the CUDA/ZLUDA path with the reason logged. Byte-identical to NGX's CUDA backend under the phase-capable reference (1706x960 -> 2560x1440, native HIP network, eleven frames, RX 9070 XT): - K render- and display-resolution motion vectors, HDR and LDR - K game-supplied exposure and automatic exposure with any pre-exposure and exposure scale (one residual: effective exposure exactly 2 or 4) - K and M with the DLSS indicator's inverted-axis parameters - M with any exposure settings (NGX ignores them for M) K packages need the current compile_k.py (marker D4RO0003). Also: HIP shared buffers are pooled across feature recreation (ROCm leaks destroyed imports), hipnet.bin must match the preset, M-only games do not load the K model, install_d4r_runtime.sh stages k, k-ldr, m, m-ldr with hip/, the activation table has a generator (gen_activation_m.py), and package_release.sh can bundle the engine's shaders (never its weights). Signed-off-by: woutermans <[email protected]>
This commit is contained in:
woutermans committed
1 parent
16cfc1f063
commit
e0d1549a21
82 files changed
+13335
-96
No files matched your search
@@ -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`.
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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)); }
|
||||
Executable
+20
@@ -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"
|
||||
@@ -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('<f', self.args, offset)[0]
|
||||
def h1(self, offset): return float(np.frombuffer(self.args, np.float16, 1, offset)[0])
|
||||
def exposure(self, offset):
|
||||
self.resources[offset].validate()
|
||||
value = float(np.fromfile(self.resources[offset].path, np.float16, 1)[0])
|
||||
return value if value != 0 else 1.0
|
||||
|
||||
def buffer(self, offset, size):
|
||||
pointer = struct.unpack_from('<Q', self.args, offset)[0]
|
||||
# Post-launch allocation captures include each entire allocation, not just this argument's slice.
|
||||
for path in (self.capture / 'dump').glob(f'launch-{self.sequence:03d}-{self.kernel}-arg*-buffer-*'):
|
||||
base = int(re.search(r'-base0x([0-9a-f]+)', path.name).group(1), 16)
|
||||
delta = pointer - base
|
||||
if 0 <= delta and delta + size <= path.stat().st_size:
|
||||
with path.open('rb') as stream:
|
||||
stream.seek(delta)
|
||||
return stream.read(size)
|
||||
for line in self.manifest:
|
||||
fields = line.split()
|
||||
if fields[0] == 'pointer' and int(fields[1]) == offset:
|
||||
path = self.directory / f'alloc-{fields[2]}.bin'
|
||||
if path.exists():
|
||||
with path.open('rb') as stream:
|
||||
stream.seek(int(fields[3])); data = stream.read(size)
|
||||
if len(data) == size: return data
|
||||
raise ValueError(f'{self.directory}: missing {size}-byte buffer at argument {offset}')
|
||||
|
||||
def buffer_location(self, offset):
|
||||
"""(file holding the whole allocation, byte offset of the argument's pointer within it)"""
|
||||
pointer = struct.unpack_from('<Q', self.args, offset)[0]
|
||||
for path in (self.capture / 'dump').glob(f'launch-{self.sequence:03d}-{self.kernel}-arg*-buffer-*'):
|
||||
base = int(re.search(r'-base0x([0-9a-f]+)', path.name).group(1), 16)
|
||||
if 0 <= pointer - base < path.stat().st_size:
|
||||
return path, pointer - base
|
||||
for line in self.manifest:
|
||||
fields = line.split()
|
||||
if fields[0] == 'pointer' and int(fields[1]) == offset and (self.directory / f'alloc-{fields[2]}.bin').exists():
|
||||
return self.directory / f'alloc-{fields[2]}.bin', int(fields[3])
|
||||
raise ValueError(f'{self.directory}: missing buffer at argument {offset}')
|
||||
|
||||
def buffer_offset(self, offset):
|
||||
return self.buffer_location(offset)[1]
|
||||
|
||||
def allocation(self, offset, size):
|
||||
data = self.buffer_location(offset)[0].read_bytes()
|
||||
if len(data) != size:
|
||||
raise ValueError(f'{self.directory}: the allocation at argument {offset} has {len(data)} bytes, expected {size}')
|
||||
return data
|
||||
|
||||
|
||||
def launches(capture, kernel):
|
||||
return [Launch(capture, d) for d in sorted((Path(capture) / 'replay').glob('replay-*-'+kernel))]
|
||||
|
||||
|
||||
def input_params(p):
|
||||
f = np.float32
|
||||
a, b, c = (f(p.exposure(o)) for o in (208, 200, 216))
|
||||
f152, f158 = f(p.f1(144)) * a, f(p.f1(144)) * b
|
||||
expo = f(p.f1(156)) * f158
|
||||
if expo == 0: expo = f(1)
|
||||
expo_in, ratio = f(p.f1(148)) * c / f158, f158 / f152
|
||||
blob = struct.pack('<14f', *p.f2(40), *p.f2(24), *p.f2(8), *p.f2(16), *p.f2(120), *p.f2(136), *p.f2(128))
|
||||
blob += struct.pack('<12i', *p.i2(72), *p.i2(80), *p.i2(88), *p.i2(96), *p.i2(104), *p.i2(112))
|
||||
blob += struct.pack('<iI4f', struct.unpack_from('<i', p.args, 64)[0], p.args[32], float(expo), float(expo_in), float(ratio), p.f1(160))
|
||||
assert len(blob) == 128
|
||||
return blob
|
||||
|
||||
|
||||
def output_params(p):
|
||||
if any(p.args[89:92]):
|
||||
raise ValueError('native K supports standard HDR output only (flag bytes 89..91 must be zero)')
|
||||
f = np.float32
|
||||
a, b = f(p.exposure(264)), f(p.exposure(256))
|
||||
s160, s168 = f(p.f1(160)), f(p.f1(168))
|
||||
expo = s168 * (s160 * b)
|
||||
if expo == 0: expo = f(1)
|
||||
ratio = (s160 * b) / (s160 * a)
|
||||
blob = struct.pack('<18f', *p.f2(40), *p.f2(16), *p.f2(56), *p.f2(48), *p.f2(64), *p.f2(0), *p.f2(8), *p.f2(144), *p.f2(152))
|
||||
blob += struct.pack('<16i', *p.i2(208), *p.i2(216), *p.i2(96), *p.i2(104), *p.i2(112), *p.i2(120), *p.i2(128), *p.i2(136))
|
||||
# flags: reset, and in bits 8..11 the reconstruction table the launch pointed at within NGX's sixteen
|
||||
blob += struct.pack('<2iI', *p.i2(80), p.args[32] | (p.buffer_offset(280) // 32768) << 8)
|
||||
blob += struct.pack('<10f', float(expo), float(ratio), *struct.unpack_from('<3f', p.args, 172), p.h1(188), p.h1(190), p.h1(192), p.h1(184), p.h1(186))
|
||||
assert len(blob) == 188
|
||||
return blob
|
||||
@@ -0,0 +1,219 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Package extracted K weights and compile d4r's shaders. No PTX translator is used.
|
||||
|
||||
compile_k.py MODEL_PREFIX LUT.bin OUTPUT_DIRECTORY
|
||||
LUT.bin holds NGX's sixteen anisotropic reconstruction tables (2048 sixteen-byte entries each): the whole
|
||||
allocation behind the output launch's argument 280, as prepare_k.py writes it. NGX picks one by upscaling
|
||||
ratio. A single 32768-byte table (older captures) is accepted and then serves every ratio.
|
||||
|
||||
--hip KERNEL_DIR --capture CAPTURE_DIR also packages the network for the native
|
||||
HIP layer kernels (engine/hip_net.h): hipnet.bin ("D4RKHIP1", the eleven layers in
|
||||
network order with the weight allocations exactly as NGX uploads them, which the
|
||||
native kernels read) and hip/dltss_pwin_*_layer.hsaco copied from KERNEL_DIR.
|
||||
CAPTURE_DIR/replay must hold one captured launch of each dltss_pwin_*_layer with
|
||||
its weight allocation (D4R_CUDA_REPLAY_DUMP_DIR, filter dltss_pwin). HDR and LDR
|
||||
share the weights, so one hipnet.bin serves both variants.
|
||||
"""
|
||||
# D4RK0002 couples the package to the exposure stage's sample grid: the engine passes automatic exposure the
|
||||
# native output/2 source grid (mandatory dispatch size and uniform sampleSize), and the K1 package's exposure
|
||||
# shader and dispatch assume the render grid, so a K1 offsets blob is not interchangeable with this runtime.
|
||||
# Nothing else in the package changed meaning: the tensor offsets and the layer/input/output stages are the same.
|
||||
import argparse
|
||||
import glob
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
import struct
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
from extract_k import LAYERS
|
||||
from test_k import defines, params
|
||||
|
||||
HERE = Path(__file__).resolve().parent
|
||||
MAGIC = b'D4RK0002'
|
||||
MANIFEST_VERSION = 2
|
||||
|
||||
# Native HIP packaging (hipnet.bin). The layer order is the network order of extract_k.LAYERS; NGX uploads every
|
||||
# layer's weights back to back into one allocation in the layout the native kernels read
|
||||
# (kernels/tools/pwin_model.Layout), so each blob is the launch's PwinParams.w up to the next layer's weights
|
||||
# pointer. `divisor` is the token-grid divisor of makeKShape (2**shift); the sx/sy written per entry are the
|
||||
# frame-0 window shifts (the backend recomputes the whole eight-frame cycle itself).
|
||||
HIP_ORDER = ['enc0', 'enc1', 'enc2', 'enc3', 'enc4', 'dec5', 'dec4', 'dec3', 'dec2', 'dec1', 'dec0']
|
||||
HIP_KERNEL = 'dltss_pwin_{}_layer'
|
||||
HIP_DIVISOR = {'enc0': 1, 'enc1': 2, 'enc2': 4, 'enc3': 8, 'enc4': 16, 'dec5': 32,
|
||||
'dec4': 16, 'dec3': 8, 'dec2': 4, 'dec1': 2, 'dec0': 1}
|
||||
# makeKWindows (engine/runtime.cpp): an eight-frame cycle of per-axis shifts; the five deepest layers move and
|
||||
# every shallower one follows from the one below as (2 * shift + 4) mod 8. depth 0 = dec5, 1 = enc4/dec4,
|
||||
# 2 = enc3/dec3; the rest keep shift 4.
|
||||
HIP_CYCLE = ((0, 2), (5, 6), (4, 0), (1, 4), (7, 5), (2, 1), (3, 7), (6, 3))
|
||||
HIP_ROLE = {'dec5': 0, 'enc4': 1, 'dec4': 1, 'enc3': 2, 'dec3': 2}
|
||||
|
||||
|
||||
def hip_shift(role, frame):
|
||||
x, y = HIP_CYCLE[frame % 8]
|
||||
for _ in range(role):
|
||||
x, y = (2 * x + 4) % 8, (2 * y + 4) % 8
|
||||
return x, y
|
||||
|
||||
|
||||
def hip_weight_place(launch, offset=64):
|
||||
"""The weight allocation a launch's PwinParams.w (argument at `offset`) points into: (base, size, path, ptr)."""
|
||||
ptr = struct.unpack_from('<Q', open(f'{launch}/args.bin', 'rb').read(), offset)[0]
|
||||
if ptr == 0:
|
||||
raise ValueError(f'{launch}: no weight pointer')
|
||||
for line in open(f'{launch}/manifest.txt'):
|
||||
q = line.split()
|
||||
if q and q[0] == 'alloc':
|
||||
path = f'{launch}/alloc-{q[1]}.bin'; base = int(q[2], 16); size = os.path.getsize(path)
|
||||
if base <= ptr < base + size:
|
||||
return (base, size, path, ptr)
|
||||
raise ValueError(f'{launch}: the weights allocation is missing')
|
||||
|
||||
|
||||
def package_hip(output, kernels, capture):
|
||||
"""Writes hipnet.bin + hip/*.hsaco. The kernels' PwinParams.w is the launch's argument at byte 64.
|
||||
NGX packs the eleven layers back to back in one allocation, so a layer's blob is its region up to the
|
||||
next layer's weights pointer (the last one to the end of the allocation)."""
|
||||
place = {}
|
||||
for name in HIP_ORDER:
|
||||
found = sorted(glob.glob(f'{capture}/replay/replay-*-{HIP_KERNEL.format(name)}'))
|
||||
if not found:
|
||||
raise ValueError(f'{capture}: no captured launch of {HIP_KERNEL.format(name)}')
|
||||
place[name] = hip_weight_place(found[0])
|
||||
# The per-layer region rule below needs the layers to be packed in network order inside one allocation;
|
||||
# NGX does that (and packs it identically for HDR and LDR), so a different order must be noticed here.
|
||||
offsets = [(place[name][3] - place[name][0]) for name in HIP_ORDER]
|
||||
if offsets != sorted(offsets) or len(set(offsets)) != len(offsets) or offsets[0] != 0:
|
||||
raise ValueError(f'{capture}: the layers are not packed in network order: {offsets}')
|
||||
blobs, at = [], 16 + 64 * len(HIP_ORDER)
|
||||
head = struct.pack('<8sI4x', b'D4RKHIP1', len(HIP_ORDER))
|
||||
for name in HIP_ORDER:
|
||||
base, size, path, ptr = place[name]
|
||||
following = [p for b, _, _, p in place.values() if b == base and p > 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('<II', 11, weights.stat().st_size) + offsets)
|
||||
for name, e in index.items():
|
||||
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *defines(e),
|
||||
str(HERE / 'layer.comp'), '-o', str(a.output / (name + '.spv'))], check=True)
|
||||
for name in ('input_k', 'output_k'):
|
||||
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *(['-DLDR=1'] if a.ldr else []),
|
||||
str(HERE / (name + '.comp')), '-o', str(a.output / (name + '.spv'))], check=True)
|
||||
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', '-DOUTPUT_RGB10=1',
|
||||
*(['-DLDR=1'] if a.ldr else []), str(HERE / 'output_k.comp'),
|
||||
'-o', str(a.output / 'output_k_rgb10.spv')], check=True)
|
||||
for stage in '01':
|
||||
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', '-DPASS=' + stage, str(HERE / 'exposure_k.comp'),
|
||||
'-o', str(a.output / f'exposure_k{stage}.spv')], check=True)
|
||||
# Shader revision marker. Packages without it double-added the motion origin (the adapter uses copies for
|
||||
# their nonzero rectangles); revision 1 shaders ignore the frame flags, so they cannot take regular depth.
|
||||
# Revision 2 reads the frame flags; revision 3 also reads the display-motion flag and the game's exposure
|
||||
# texture (exposure flag 4).
|
||||
(a.output / 'direct_origins.bin').write_bytes(b'D4RO0003')
|
||||
if a.hip:
|
||||
package_hip(a.output, a.hip, a.capture)
|
||||
files = sorted(a.output.glob('*.bin')) + sorted(a.output.glob('*.spv'))
|
||||
if a.hip:
|
||||
files += sorted(a.output.glob('hip/*.hsaco'))
|
||||
manifest = {'format': 'd4r-k', 'version': MANIFEST_VERSION, 'layers': list(index),
|
||||
'variant': 'ldr' if a.ldr else 'hdr',
|
||||
'numerics': 'folded-gains' if all(e.get('fold_gains', True) for e in index.values()) else 'explicit-f16-gains',
|
||||
'sha256': {f.name: hashlib.sha256(f.read_bytes()).hexdigest() for f in files}}
|
||||
(a.output / 'model.json').write_text(json.dumps(manifest, indent=2) + '\n')
|
||||
print(f'Compiled native preset K model in {a.output}' + (f' with native HIP kernels from {a.hip}' if a.hip else ''))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,117 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Build the engine's preset M model directory from a capture of one frame.
|
||||
|
||||
usage: compile_m.py CAPTURE_DIR OUT_DIR [--frame N]
|
||||
CAPTURE_DIR/replay must hold the frame's rrlite_enc0 and rrlite_dec0 launches (the `_folded` kernels) and its
|
||||
ten Swin layer launches, with their weight allocations (D4R_CUDA_REPLAY_DUMP_DIR, filter `rrlite_`).
|
||||
Writes weights8.bin (the layers' FP8 matrices in the blocked fixed-16 tile layout of test_m.pack8, with the front-end
|
||||
enc0/dec0 matrices plain and NGX-ordered as their own stages read them), weights16.bin (f16 vectors and the dither
|
||||
table), offsets.bin and the compiled shaders. The weights are values read from NVIDIA's DLSS and are not redistributable
|
||||
with d4r; the shaders are d4r's own.
|
||||
"""
|
||||
# D4RM0003 couples the blocked FP8 operands, RGBA32F half-history backing and 8x64 downsample dispatch;
|
||||
# the previous package's shader workgroup dimensions are not interchangeable with this runtime.
|
||||
import argparse, glob, hashlib, json, os, struct, subprocess, sys
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
HERE = Path(__file__).resolve().parent
|
||||
sys.path.insert(0, str(HERE))
|
||||
import m_model as mm
|
||||
import model_enc3 as me
|
||||
import swin_model as sm
|
||||
from test_m import pack_layer, defines
|
||||
|
||||
SHADERS = {'enc1': 'm_enc1', 'enc2': 'm_enc2', 'enc3_tube': 'm_tube', 'dec2': 'm_dec2', 'dec1': 'm_dec1'}
|
||||
|
||||
|
||||
def blob_of(directory, ptr_offset=0):
|
||||
args = open(f'{directory}/args.bin', 'rb').read()
|
||||
ptr = struct.unpack_from('<Q', args, ptr_offset)[0]
|
||||
for line in open(f'{directory}/manifest.txt'):
|
||||
q = line.split()
|
||||
if q and q[0] == 'alloc':
|
||||
base = int(q[2], 16); size = os.path.getsize(f'{directory}/alloc-{q[1]}.bin')
|
||||
if base <= ptr < base + size:
|
||||
return np.fromfile(f'{directory}/alloc-{q[1]}.bin', np.uint8)[ptr - base:]
|
||||
raise ValueError(f'{directory}: the weights allocation is missing')
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(description=__doc__)
|
||||
ap.add_argument('capture'); ap.add_argument('output', type=Path); ap.add_argument('--frame', type=int, default=0)
|
||||
ap.add_argument('--ldr', action='store_true', help="NGX's *_ldr_* stages, for games without the HDR flag (capture such a frame)")
|
||||
ap.add_argument('--hip', type=Path, metavar='KERNEL_DIR', help='also package the network for the native HIP layers: the layer '
|
||||
'weights as NGX uploads them (hipnet.bin) and the rrlite_*_4x4.hsaco kernels of KERNEL_DIR (kernels/build.sh output)')
|
||||
a = ap.parse_args()
|
||||
a.output.mkdir(parents=True, exist_ok=True)
|
||||
def one(kernel):
|
||||
d = sorted(glob.glob(f'{a.capture}/replay/replay-*-{kernel}'))
|
||||
if len(d) <= a.frame: raise ValueError(f'no launch of {kernel} for frame {a.frame}')
|
||||
return d[a.frame]
|
||||
R = 'ldr' if a.ldr else 'hdr'
|
||||
w8, w16 = [], []
|
||||
n8 = n16 = 0
|
||||
def add8(x):
|
||||
nonlocal n8
|
||||
x = np.ascontiguousarray(x, np.uint8).ravel(); pad = (-x.size) % 64
|
||||
off = n8; w8.append(x); w8.append(np.zeros(pad + 64, np.uint8)); n8 += x.size + pad + 64
|
||||
return off
|
||||
def add16(x):
|
||||
nonlocal n16
|
||||
x = np.ascontiguousarray(x).view(np.uint8).ravel(); pad = (-x.size) % 64
|
||||
off = n16; w16.append(x); w16.append(np.zeros(pad + 64, np.uint8)); n16 += x.size + pad + 64
|
||||
return off
|
||||
table = [] # (w8 offset, w16 offset) per stage, then the layers' offset blocks
|
||||
# input stage: embedding [128][64] in natural output order, bias tile, dither table
|
||||
wb = blob_of(one(f'rrlite_enc0_4x4_mvlo_{R}_folded'))
|
||||
g = sm.gp(np.arange(64))
|
||||
enc0 = (add8(wb[sm.woff_table(4096, 512, 2048, 128, 64)][:, g]), add16(np.tile(wb[12288:12288 + 128].view(np.float16)[g][None, :], (16, 1))), add16(wb[:4096]))
|
||||
# expansion: the 20 columns per sub-position the stage uses
|
||||
wb = blob_of(one('rrlite_dec0_4x4_folded'))
|
||||
cols = (np.arange(4)[:, None] * 32 + np.arange(20)[None, :]).ravel()
|
||||
dec0 = (add8(wb[sm.woff_table(0, 512, 1024, 64, 128)][:, cols]), add16(np.tile(wb[8192:8192 + 256].view(np.float16)[cols][None, :], (16, 1))))
|
||||
Ps, dirs = mm.launches(a.capture, a.frame)
|
||||
layers = []
|
||||
for P, name in zip(Ps, mm.ORDER):
|
||||
o, b8, b16 = pack_layer(mm.plain_layer(P, name), name)
|
||||
layers.append((add8(b8), add16(b16), o))
|
||||
blob8, blob16 = np.concatenate(w8), np.concatenate(w16)
|
||||
blob8.tofile(a.output / 'weights8.bin'); blob16.tofile(a.output / 'weights16.bin')
|
||||
head = struct.pack('<8sII', b'D4RM0003', blob8.size, blob16.size) + struct.pack('<5I', *enc0, *dec0)
|
||||
for o8, o16, o in layers: head += struct.pack('<2I72I', o8, o16, *o)
|
||||
(a.output / 'offsets.bin').write_bytes(head)
|
||||
def glsl(src, out, flags=()):
|
||||
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *flags, str(HERE / src), '-o', str(a.output / out)], check=True, stdout=subprocess.DEVNULL)
|
||||
for name, out in SHADERS.items(): glsl('layer_m.comp', out + '.spv', defines(name))
|
||||
for s in ('enc0_m', 'dec0_m', 'post_m', 'down_m'): glsl(s + '.comp', s + '.spv', ['-DLDR=1'] if a.ldr and s != 'dec0_m' else [])
|
||||
glsl('down_m.comp', 'down_m_rgb10.spv', ['-DOUTPUT_RGB10=1'] + (['-DLDR=1'] if a.ldr else []))
|
||||
glsl('post_m.comp', 'post_m_q.spv', ['-DTILE_W=15', '-DTILE_H=15'] + (['-DLDR=1'] if a.ldr else []))
|
||||
glsl('post_m.comp', 'post_m_c.spv', ['-DTILE_W=11', '-DTILE_H=11'] + (['-DLDR=1'] if a.ldr else []))
|
||||
# the same stages with the token buffers in NGX's plane layout, for a network that runs outside the engine
|
||||
for s in ('enc0_m', 'dec0_m'): glsl(s + '.comp', s + '_planes.spv', ['-DPLANES=1'] + (['-DLDR=1'] if a.ldr and s != 'dec0_m' else []))
|
||||
if a.hip:
|
||||
# hipnet.bin: "D4RMHIP1", layer count, then per layer {kernel name[32], shift x, y, token grid divisor, weight
|
||||
# offset, weight bytes} and the weight allocations as NGX uploads them (the native kernels read that layout)
|
||||
scale = {'enc1': 1, 'enc2': 2, 'enc3_tube': 4, 'dec2': 2, 'dec1': 1}
|
||||
blobs, head, at = [], struct.pack('<8sI4x', b'D4RMHIP1', len(Ps)), 16 + 64 * len(Ps)
|
||||
for P, name in zip(Ps, mm.ORDER):
|
||||
w = np.ascontiguousarray(P.wbuf); pad = (-w.size) % 64
|
||||
head += struct.pack('<32siiIIQQ', f'rrlite_{name}_4x4'.encode(), P.sx, P.sy, scale[name], 0, at, w.size)
|
||||
blobs += [w, np.zeros(pad, np.uint8)]; at += w.size + pad
|
||||
(a.output / 'hipnet.bin').write_bytes(head + np.concatenate(blobs).tobytes())
|
||||
(a.output / 'hip').mkdir(exist_ok=True)
|
||||
for name in sorted(set(mm.ORDER)):
|
||||
src = a.hip / f'rrlite_{name}_4x4.hsaco'
|
||||
if not src.exists(): raise ValueError(f'{src} is missing')
|
||||
(a.output / 'hip' / src.name).write_bytes(src.read_bytes())
|
||||
marker = a.output / 'ldr' # tells the runtime to skip exposure
|
||||
if a.ldr: marker.write_text('1\n')
|
||||
elif marker.exists(): marker.unlink()
|
||||
for p in '01': glsl('exposure_k.comp', f'exposure_m{p}.spv', ['-DMODE=1', '-DPASS=' + p])
|
||||
files = sorted(a.output.glob('*.bin')) + sorted(a.output.glob('*.spv')) + sorted(a.output.glob('hip/*.hsaco'))
|
||||
(a.output / 'model.json').write_text(json.dumps({'format': 'D4RM0003', 'preset': 'M', 'files': {str(f.relative_to(a.output)): hashlib.sha256(f.read_bytes()).hexdigest() for f in files}}, indent=1))
|
||||
print(f'Compiled native preset M model in {a.output} ({blob8.size} FP8 bytes, {blob16.size} bytes of f16 data)')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,147 @@
|
||||
#version 460
|
||||
// d4r engine, preset M: the stage after the network. Expands the last decoder's tokens (64 channels, FP8) to 17
|
||||
// values per render pixel (one matrix product per 4x4-token block) and turns them into the reconstruction's
|
||||
// per-pixel inputs: a recurrent four-channel feature blended with its reprojected history, three gate values,
|
||||
// three softmax weights and a 2x2 covariance. See M_NOTES.md.
|
||||
#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"
|
||||
#include "normalized_store.glsl"
|
||||
#include "activation_m.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
|
||||
int tokW; uint reset;
|
||||
vec2 rndOffs; // jitter
|
||||
vec2 rnd2hi; // render pixels -> 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<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
|
||||
#define HACC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
|
||||
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
|
||||
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
|
||||
#define RM gl_CooperativeMatrixLayoutRowMajor
|
||||
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)));
|
||||
}
|
||||
}
|
||||
@@ -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)));
|
||||
}
|
||||
@@ -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<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
|
||||
#define HACC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
|
||||
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
|
||||
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
|
||||
#define RM gl_CooperativeMatrixLayoutRowMajor
|
||||
// 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
|
||||
}
|
||||
+130
@@ -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 <vulkan/vulkan.h>
|
||||
#include <algorithm>
|
||||
#include <chrono>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <fstream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#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<char> 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<char>((std::istreambuf_iterator<char>(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<B> 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<VkDescriptorSetLayoutBinding> lb;
|
||||
for (B& b : bs)
|
||||
{
|
||||
std::vector<char> data;
|
||||
if (!b.file.empty()) data = slurp(b.file);
|
||||
if (b.out.empty()) b.size = data.size();
|
||||
b.size = std::max<VkDeviceSize>((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<size_t>(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<VkDescriptorBufferInfo> bi(bs.size()); std::vector<VkWriteDescriptorSet> 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<double, std::micro>(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<double> 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;
|
||||
}
|
||||
@@ -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
|
||||
@@ -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()
|
||||
+126
@@ -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<Impl>(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
|
||||
@@ -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> impl;
|
||||
};
|
||||
} // namespace d4r
|
||||
@@ -0,0 +1,139 @@
|
||||
#pragma once
|
||||
#include "game.h"
|
||||
#include "vk_dispatch.h"
|
||||
#include <stdexcept>
|
||||
|
||||
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
|
||||
@@ -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<Impl>(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
|
||||
@@ -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> impl;
|
||||
};
|
||||
} // namespace d4r
|
||||
Executable
+68
@@ -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:]))
|
||||
@@ -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 <hip/hip_runtime_api.h>
|
||||
#include <algorithm>
|
||||
#include <cstring>
|
||||
#include <fstream>
|
||||
#include <stdexcept>
|
||||
#include <vector>
|
||||
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<char> slurp(const std::string& path) {
|
||||
std::ifstream f(path, std::ios::binary);
|
||||
if (!f) throw std::runtime_error("cannot read " + path);
|
||||
return std::vector<char>((std::istreambuf_iterator<char>(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<Kernel> kernels;
|
||||
std::vector<Node> preps;
|
||||
std::vector<Node> nodes[MaxGraphs];
|
||||
std::vector<void*> 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<Params*>(&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<Impl>()) {
|
||||
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<std::string> 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<uint32_t> layerKernel(layers), divisor(layers);
|
||||
std::vector<int32_t> shiftX(layers), shiftY(layers);
|
||||
std::vector<void*> 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<Params> 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<void*> 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
|
||||
@@ -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 <cstdint>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
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> impl;
|
||||
};
|
||||
} // namespace d4r
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
|
||||
#define HACC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
|
||||
#define MA coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
|
||||
#define MB coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
|
||||
#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
|
||||
}
|
||||
@@ -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<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
|
||||
#define HAC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
|
||||
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
|
||||
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
|
||||
#define A16 coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
|
||||
#define B16 coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
|
||||
#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
|
||||
}
|
||||
@@ -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()
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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 <algorithm>
|
||||
#include <chrono>
|
||||
#include <cstring>
|
||||
#include <cstdlib>
|
||||
#include <filesystem>
|
||||
#include <fstream>
|
||||
#include <iomanip>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <sstream>
|
||||
#include <stdexcept>
|
||||
#include <vector>
|
||||
|
||||
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<char> 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<char> data(static_cast<size_t>(bytes)); stream.seekg(0);
|
||||
if (!stream.read(data.data(), bytes)) throw std::runtime_error("truncated file: " + file);
|
||||
return data;
|
||||
}
|
||||
template<class T> 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<VkPhysicalDevice> 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<VkExtensionProperties> 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<PFN_vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR>(vkGetInstanceProcAddr(instance, "vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR"));
|
||||
if (!matrixProperties) throw std::runtime_error("cooperative matrix query unavailable");
|
||||
uint32_t matrixCount = 0; CK(matrixProperties(physical, &matrixCount, nullptr));
|
||||
std::vector<VkCooperativeMatrixPropertiesKHR> 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<VkQueueFamilyProperties> 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<const char*>(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<PFN_vkGetMemoryFdKHR>(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<d4r::LayerShape,d4r::KLayers> 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<const char*, d4r::KTimingStages> 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<Marks*>(user);
|
||||
vkCmdWriteTimestamp(cb, VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT, m.pool, m.first + boundary);
|
||||
}, &marks};
|
||||
|
||||
d4r::Shape shape{};
|
||||
std::array<std::string,3> historyFiles;
|
||||
std::vector<std::unique_ptr<Image>> images;
|
||||
std::vector<Frame> 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<d4r::Shape>(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<Image>(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<d4r::InputParams>(input); f.output = structure<d4r::OutputParams>(output);
|
||||
f.windows = structure<std::array<d4r::LayerShape,d4r::KLayers>>(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<d4r::HipNetwork> 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<d4r::HipNetwork>(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<uint64_t> 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<double, d4r::KTimingStages> 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<double,std::milli>(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<std::vector<double>, 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<double> 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<double,std::micro>(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; }
|
||||
}
|
||||
@@ -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 <chrono>
|
||||
#include <unistd.h>
|
||||
#include <algorithm>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <fstream>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#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<char> slurp(const std::string& p) {
|
||||
std::ifstream f(p, std::ios::binary);
|
||||
if (!f) throw std::runtime_error("cannot read " + p);
|
||||
return std::vector<char>((std::istreambuf_iterator<char>(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<PFN_vkGetMemoryFdKHR>(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<d4r::HipNetwork> 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<d4r::HipNetwork>(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<std::vector<char>> 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<double> times, walls; std::vector<std::vector<double>> stage(6);
|
||||
const auto now = [] { return std::chrono::duration<double, std::micro>(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<const char*>(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; }
|
||||
@@ -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(); \
|
||||
}
|
||||
@@ -0,0 +1,782 @@
|
||||
#include "runtime.h"
|
||||
#include "vk_dispatch.h"
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
#include <fstream>
|
||||
#include <stdexcept>
|
||||
#include <vector>
|
||||
|
||||
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<char> 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<char> data(static_cast<size_t>(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<VkDeviceSize>(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<VkDescriptorType>& types,
|
||||
uint32_t subgroup, VkPipelineCache cache) {
|
||||
device = dev;
|
||||
std::vector<VkDescriptorSetLayoutBinding> 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<uint32_t>(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<uint32_t> 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<uint32_t, 108> 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<LayerShape, KLayers> 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<std::array<uint32_t, 108>, KLayers> offsets{};
|
||||
std::array<Pipeline, KLayers> network;
|
||||
Pipeline input, output, outputRgb10;
|
||||
std::array<Pipeline, 2> 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<Impl>(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<VkDescriptorSet, 2> exposureSets{};
|
||||
ExposureParams exposure{};
|
||||
bool exposureOn = false;
|
||||
std::array<Buffer, KLayers> layerParams;
|
||||
std::array<Buffer, 2> main;
|
||||
std::array<Buffer, 5> skip;
|
||||
struct History { Image color, luma, feature; };
|
||||
std::array<History, 2> history;
|
||||
Image final;
|
||||
VkDescriptorPool pool = VK_NULL_HANDLE;
|
||||
std::array<VkDescriptorSet, KLayers> layerSets{};
|
||||
std::array<VkDescriptorSet, 2> 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<VkDescriptorSet, 2> input{}, output{};
|
||||
VkDescriptorSet exposure = VK_NULL_HANDLE;
|
||||
VkDescriptorSet exposureUpdate = VK_NULL_HANDLE; // the exposure state update (pass 1), per binding
|
||||
};
|
||||
std::vector<Bindings> 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<Impl>(*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<VkDeviceSize, 2> 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<LayerShape, KLayers>* 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
|
||||
@@ -0,0 +1,237 @@
|
||||
// GPU-resident preset K inference. All Vulkan handles belong to the caller's device.
|
||||
#pragma once
|
||||
#include <vulkan/vulkan.h>
|
||||
#include <array>
|
||||
#include <cstdint>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
|
||||
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<LayerShape, KLayers> 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<LayerShape, KLayers> 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> 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<LayerShape, KLayers>* 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> 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> 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> impl;
|
||||
};
|
||||
} // namespace d4r
|
||||
@@ -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<std::array<uint32_t, 74>, MLayers> layers{};
|
||||
std::array<Pipeline, MLayers> 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<Pipeline, 2> 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<Impl>(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<Buffer, MLayers> layerP;
|
||||
Buffer tokens, skip1, merged1, skip2, merged2, tubeA, tubeB, dec2Out, dec1Out;
|
||||
std::array<Tex, 2> mvA, feat, hist, half;
|
||||
Tex mvB, mvC, gates, mix, cov, output;
|
||||
VkDescriptorPool pool = VK_NULL_HANDLE;
|
||||
std::array<VkDescriptorSet, MLayers> layerSets{};
|
||||
std::array<VkDescriptorSet, 2> 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<VkDescriptorSet, 2> input{}, post{};
|
||||
VkDescriptorSet exposure = VK_NULL_HANDLE, down = VK_NULL_HANDLE;
|
||||
};
|
||||
std::vector<Bindings> 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<Impl>(*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; }
|
||||
@@ -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('<Q', P, 0)[0], struct.unpack_from('<Q', P, 24)[0]
|
||||
i2 = lambda o: struct.unpack_from('<2i', P, o); f2 = lambda o: struct.unpack_from('<2f', P, o)
|
||||
gw, gh = i2(16); rw, rh = i2(112)
|
||||
wa, wbuf = find(wptr); wbuf = wbuf[wptr - wa:]
|
||||
ab, arena = find(inp)
|
||||
token_codes(arena, inp - ab, gw // 2, gh // 2, 64).tofile(f'{work}/tokens.bin')
|
||||
# plain weights [64][128] and bias; the stage uses columns 32 q + j, j < 20
|
||||
W = wbuf[sm.woff_table(0, 512, 1024, 64, 128)]
|
||||
bias = wbuf[8192:8192 + 256].view(np.float16)
|
||||
cols = (np.arange(4)[:, None] * 32 + np.arange(20)[None, :]).ravel()
|
||||
np.concatenate([W[:, cols].ravel(), np.zeros(64, np.uint8)]).tofile(f'{work}/w8.bin')
|
||||
np.concatenate([np.tile(bias[cols][None, :], (16, 1)).ravel(), np.zeros(32, np.float16)]).tofile(f'{work}/w16.bin')
|
||||
open(f'{work}/params.bin', 'wb').write(struct.pack('<2iiI2f2f2f2f2f2f', rw, rh, gw // 2, P[120], *f2(72), *f2(88), float(i2(96)[0]), float(i2(96)[1]), *f2(80), *f2(160), *f2(104)))
|
||||
lines = [f'param {work}/params.bin', f'grid {gw // 8} {gh // 8}',
|
||||
f'tex 1 {res[56]} {rw} {rh} 2 16 nearest', f'tex 2 {res[64]} {rw} {rh} 4 9 linear']
|
||||
outs = ((128, 9, np.int8, 127.0, 'feature (SNORM8)'), (136, 8, np.uint8, 255.0, 'gates (UNORM8)'), (144, 8, np.uint8, 255.0, 'mix (UNORM8)'), (152, 16, np.float16, 1.0, 'covariance (f16)'))
|
||||
for i, (off, bits, dt, scale, name) in enumerate(outs):
|
||||
lines.append(f'surf {3 + i} {rw} {rh} 4 {bits} - {work}/out-{off}.bin')
|
||||
lines += [f'lut {work}/tokens.bin', f'lut {work}/w8.bin', f'lut {work}/w16.bin']
|
||||
open(f'{work}/manifest', 'w').write('\n'.join(lines) + '\n')
|
||||
spv = f'{work}/dec0_m.spv'
|
||||
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', f'{HERE}/dec0_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}), reset {P[120]}')
|
||||
for off, bits, dt, scale, name in outs:
|
||||
a = np.fromfile(res[off], dt).astype(np.float64) / scale; b = np.fromfile(f'{work}/out-{off}.bin', dt).astype(np.float64) / scale
|
||||
dd = np.abs(a - b)
|
||||
print(f' {name:18s} psnr {10 * np.log10(max(np.abs(a).max(), 1e-9) ** 2 / max((dd ** 2).mean(), 1e-30)):5.1f} dB identical {100 * (dd == 0).mean():5.1f}% '
|
||||
f'within one step {100 * (dd <= 1.0 / scale + 1e-9).mean() if scale > 1 else 100 * (dd < 1e-3 * np.abs(a).max()).mean():5.1f}% max |diff| {dd.max():.4g}')
|
||||
@@ -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}')
|
||||
@@ -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('<f', P, o)[0]
|
||||
wptr, outp = struct.unpack_from('<Q', P, 0)[0], struct.unpack_from('<Q', P, 48)[0]
|
||||
wa, wbuf = find(wptr); wbuf = wbuf[wptr - wa:]
|
||||
rw, rh = i2(56); gw, gh = i2(16)
|
||||
expo = lambda o: float(np.fromfile(res[o], np.float16)[0])
|
||||
g = sm.gp(np.arange(64)) # the embedding's columns are in the network's pair order
|
||||
np.ascontiguousarray(wbuf[sm.woff_table(4096, 512, 2048, 128, 64)][:, g]).tofile(f'{work}/w8.bin')
|
||||
np.tile(wbuf[12288:12288 + 128].view(np.float16)[g][None, :], (16, 1)).tofile(f'{work}/w16.bin')
|
||||
wbuf[:4096].tofile(f'{work}/noise.bin')
|
||||
open(f'{work}/params.bin', 'wb').write(struct.pack('<2i2i2i2f2f2f2f2f2f2fff3Iffffff', rw, rh, *i2(8), gw // 2, gh // 2, *f2(168), *f2(176), *f2(200), *f2(208),
|
||||
*f2(192), *f2(184), float(i2(248)[0]), float(i2(248)[1]), f1(268), f1(272), P[216], P[316], i2(312)[0], expo(112), expo(120), *f2(256), *f2(64)))
|
||||
hw, hh = i2(248)
|
||||
# The native RGB10 formatted history has RGBA32F backing, not a Vulkan UNORM sampler.
|
||||
codes = np.fromfile(res[152], np.uint32) if 152 in res else np.zeros((hh // 2) * (hw // 2), np.uint32)
|
||||
half = f'{work}/half-prev.bin'
|
||||
values = np.empty((codes.size, 4), np.float32)
|
||||
for channel, shift in enumerate((0, 10, 20)):
|
||||
values[:, channel] = ((codes >> 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]))
|
||||
@@ -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()
|
||||
@@ -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('<f', s.P, o)[0]
|
||||
u8 = lambda s, o: struct.unpack_from('<Q', s.P, o)[0]
|
||||
expo = lambda s, o: float(np.fromfile(s.res[o], np.float16)[0]) if o in s.res else 1.0
|
||||
def blob(self, ptr):
|
||||
a, b = next((a, b) for a, b in self.allocs.values() if a <= ptr < a + b.size)
|
||||
return b[ptr - a:]
|
||||
|
||||
|
||||
def shader(name, flags=()):
|
||||
spv = f'{work}/{name}{"".join(flags).replace("-D", "_").replace("=", "")}.spv'
|
||||
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *flags, f'{HERE}/{name}.comp', '-o', spv], check=True, stdout=subprocess.DEVNULL)
|
||||
return spv
|
||||
|
||||
|
||||
def run(cmd):
|
||||
r = subprocess.run(cmd, capture_output=True, text=True)
|
||||
if r.returncode:
|
||||
print('FAILED', ' '.join(cmd[:3]), r.stderr[-500:]); sys.exit(1)
|
||||
|
||||
|
||||
def zeros(path, n, dtype=np.float16):
|
||||
np.zeros(n, dtype).tofile(path); return path
|
||||
|
||||
|
||||
prev = None # the previous frame's files
|
||||
for f in range(frames):
|
||||
D = f'{work}/f{f}'; os.makedirs(D, exist_ok=True)
|
||||
# ---------------------------------------------------------------- input stage
|
||||
L = Launch(NAMES['enc0'], f)
|
||||
rw, rh = L.i2(56); gw, gh = L.i2(16); tw, th = gw // 2, gh // 2
|
||||
hw, hh = L.i2(248)
|
||||
wb = L.blob(L.u8(0))
|
||||
g = sm.gp(np.arange(64))
|
||||
np.ascontiguousarray(wb[sm.woff_table(4096, 512, 2048, 128, 64)][:, g]).tofile(f'{D}/e0.w8')
|
||||
np.tile(wb[12288:12288 + 128].view(np.float16)[g][None, :], (16, 1)).tofile(f'{D}/e0.w16')
|
||||
wb[:4096].tofile(f'{D}/e0.noise')
|
||||
open(f'{D}/e0.params', 'wb').write(struct.pack('<2i2i2i2f2f2f2f2f2f2fff3Iffffff', rw, rh, *L.i2(8), tw, th, *L.f2(168), *L.f2(176), *L.f2(200), *L.f2(208),
|
||||
*L.f2(192), *L.f2(184), float(hw), float(hh), L.f1(268), L.f1(272), L.P[216], L.P[316], L.i2(312)[0], L.expo(112), L.expo(120), *L.f2(256), *L.f2(64)))
|
||||
colour, motion, depth = L.res[80], L.res[96], L.res[136]
|
||||
e1, e2 = L.expo(112), L.expo(120)
|
||||
pmv = prev['mvA'] if prev else zeros(f'{D}/zero.mv', rw * rh * 2)
|
||||
phalf = prev['half'] if prev else zeros(f'{D}/zero.half', (hw // 2) * (hh // 2) * 4, np.float32)
|
||||
pfeat = prev['feat'] if prev else zeros(f'{D}/zero.feat', rw * rh * 4, np.int8)
|
||||
phist = prev['hist'] if prev else zeros(f'{D}/zero.hist', hw * hh * 4)
|
||||
open(f'{D}/e0.manifest', 'w').write('\n'.join([
|
||||
f'param {D}/e0.params', f'grid {(gw + L.i2(8)[0] + 7) // 8} {(gh + L.i2(8)[1] + 7) // 8}',
|
||||
f'tex 1 {colour} {rw} {rh} 4 16 nearest', f'tex 2 {motion} {rw} {rh} 2 16 nearest', f'tex 3 {depth} {rw} {rh} 1 32 nearest',
|
||||
f'tex 4 {pmv} {rw} {rh} 2 16 linear', f'tex 5 {phalf} {hw // 2} {hh // 2} 4 32 linear', f'tex 6 {pfeat} {rw} {rh} 4 9 linear',
|
||||
f'surf 7 {rw} {rh} 2 16 - {D}/mvA', f'surf 8 {rw} {rh} 2 16 - {D}/mvB', f'surf 9 {rw} {rh} 2 16 - {D}/mvC',
|
||||
f'lut {D}/e0.w8', f'lut {D}/e0.w16', f'lut {D}/e0.noise', f'outbuf {D}/tokens {tw * th * 64}']) + '\n')
|
||||
run([kvk, shader('enc0_m', STAGE), f'{D}/e0.manifest', '1', '32'])
|
||||
ref_tokens = token_codes(np.fromfile(L.res[40], np.uint8), L.u8(48) - next(a for a, b in L.allocs.values() if a <= L.u8(48) < a + b.size), tw, th, 64)
|
||||
print(f'frame {f}: ' + report('input tokens', np.fromfile(f'{D}/tokens', np.uint8).reshape(th, tw, 64), ref_tokens), flush=True)
|
||||
# ---------------------------------------------------------------- network
|
||||
Ps, dirs = mm.launches(cap, f)
|
||||
own = {Ps[0].inp: f'{D}/tokens'}
|
||||
for i, (P, d, name) in enumerate(zip(Ps, dirs, mm.ORDER)):
|
||||
C, heads, NPM, CIN = mm.LAYERS[name]
|
||||
o, b8, b16 = pack_layer(mm.plain_layer(P, name), name)
|
||||
W, H = P.W, P.H
|
||||
pre = f'{D}/l{i}'
|
||||
open(pre + '.params', 'wb').write(struct.pack('<4i72I', W, H, P.sx, P.sy, *o)); b8.tofile(pre + '.w8'); b16.tofile(pre + '.w16')
|
||||
skip = own[P.p32] if CIN else zeros(pre + '.noskip', 64, np.uint8)
|
||||
msize = (H // 2) * (W // 2) * NPM if NPM else 64
|
||||
run([evk, shader('layer_m', defines(name)), str((W + P.sx + 7) // 8), str((H + P.sy + 7) // 8), '1', f'in:0={pre}.params', f'in:1={pre}.w8', f'in:2={pre}.w16',
|
||||
f'in:3={own[P.inp]}', f'in:4={skip}', f'out:5={pre}.out:{H * W * C}', f'out:6={pre}.merged:{msize}'])
|
||||
own[P.out] = pre + '.out'
|
||||
if NPM: own[P.p48] = pre + '.merged'
|
||||
seq = int(os.path.basename(dirs[-1]).split('-')[1])
|
||||
after = np.fromfile(glob.glob(f'{cap}/dump/launch-{seq:03d}-*-arg024-*')[0], np.uint8)
|
||||
P = Ps[-1]
|
||||
print(f'frame {f}: ' + report('network output', np.fromfile(own[P.out], np.uint8).reshape(P.H, P.W, 64), token_codes(after, P.out - P.abase, P.W, P.H, 64)), flush=True)
|
||||
# ---------------------------------------------------------------- dec0
|
||||
L = Launch(NAMES['dec0'], f)
|
||||
wb = L.blob(L.u8(0))
|
||||
cols = (np.arange(4)[:, None] * 32 + np.arange(20)[None, :]).ravel()
|
||||
np.concatenate([wb[sm.woff_table(0, 512, 1024, 64, 128)][:, cols].ravel(), np.zeros(64, np.uint8)]).tofile(f'{D}/d0.w8')
|
||||
np.concatenate([np.tile(wb[8192:8192 + 256].view(np.float16)[cols][None, :], (16, 1)).ravel(), np.zeros(32, np.float16)]).tofile(f'{D}/d0.w16')
|
||||
open(f'{D}/d0.params', 'wb').write(struct.pack('<2iiI2f2f2f2f2f2f', rw, rh, tw, L.P[120], *L.f2(72), *L.f2(88), float(L.i2(96)[0]), float(L.i2(96)[1]), *L.f2(80), *L.f2(160), *L.f2(104)))
|
||||
open(f'{D}/d0.manifest', 'w').write('\n'.join([
|
||||
f'param {D}/d0.params', f'grid {gw // 8} {gh // 8}', f'tex 1 {D}/mvA {rw} {rh} 2 16 nearest', f'tex 2 {pfeat} {rw} {rh} 4 9 linear',
|
||||
f'surf 3 {rw} {rh} 4 9 - {D}/feat', f'surf 4 {rw} {rh} 4 8 - {D}/gates', f'surf 5 {rw} {rh} 4 8 - {D}/mix', f'surf 6 {rw} {rh} 4 16 - {D}/cov',
|
||||
f'lut {own[Ps[-1].out]}', f'lut {D}/d0.w8', f'lut {D}/d0.w16']) + '\n')
|
||||
run([kvk, shader('dec0_m'), f'{D}/d0.manifest', '1', '32'])
|
||||
for off, name, dt in ((128, 'feat', np.int8), (136, 'gates', np.uint8), (144, 'mix', np.uint8)):
|
||||
a = np.fromfile(L.res[off], dt).astype(np.int32); b = np.fromfile(f'{D}/{name}', dt).astype(np.int32)
|
||||
print(f'frame {f}: dec0 {name:5s} identical {100 * (a == b).mean():5.1f}%, within one step {100 * (np.abs(a - b) <= 1).mean():5.1f}%, mean |diff| {np.abs(a - b).mean():.3f} steps', flush=True)
|
||||
# ---------------------------------------------------------------- reconstruction and output
|
||||
L = Launch(NAMES['post'], f)
|
||||
open(f'{D}/post.params', 'wb').write(struct.pack('<2i2i2i2f2f2ffIff', *L.i2(8), hw, hh, rw, rh, *L.f2(124), *L.f2(132), *L.f2(160), L.f1(120), L.P[168], L.expo(16), L.expo(24)))
|
||||
open(f'{D}/post.manifest', 'w').write('\n'.join([
|
||||
f'param {D}/post.params', f'grid {(hw + 15) // 16} {(hh + 15) // 16}',
|
||||
f'tex 1 {colour} {rw} {rh} 4 16 nearest', f'tex 2 {D}/mix {rw} {rh} 4 8 nearest', f'tex 3 {D}/gates {rw} {rh} 4 8 nearest', f'tex 4 {D}/cov {rw} {rh} 4 16 nearest',
|
||||
f'tex 5 {phist} {hw} {hh} 4 16 linear', f'tex 6 {D}/mvA {rw} {rh} 2 16 nearest', f'tex 7 {D}/mvB {rw} {rh} 2 16 nearest', f'tex 8 {D}/mvC {rw} {rh} 2 16 linear',
|
||||
f'surf 9 {hw} {hh} 4 16 - {D}/hist', f'surf 10 {hw // 2} {hh // 2} 4 32 - {D}/half', f'outbuf {D}/packed {hw * hh * 4}']) + '\n')
|
||||
run([kvk, shader('post_m', STAGE + POST_TILE), f'{D}/post.manifest', '1', '64'])
|
||||
L = Launch(NAMES['down'], f)
|
||||
ow, oh = L.i2(8)
|
||||
open(f'{D}/down.params', 'wb').write(struct.pack('<8i4ff', *L.i2(0), *L.i2(40), *L.i2(64), *L.i2(72), *L.f2(80), *L.f2(88), L.expo(48)))
|
||||
open(f'{D}/down.manifest', 'w').write('\n'.join([
|
||||
f'param {D}/down.params', f'grid {(ow + 7) // 8} {(oh + 63) // 64}', f'tex 1 {colour} {rw} {rh} 4 16 nearest',
|
||||
f'surf 2 {ow} {oh} 4 16 - {D}/output', f'lut {D}/packed']) + '\n')
|
||||
run([kvk, shader('down_m', STAGE), f'{D}/down.manifest', '1', '64'])
|
||||
a = np.fromfile(L.res[32], np.float16).astype(np.float64).reshape(oh, ow, 4)[..., :3]
|
||||
b = np.fromfile(f'{D}/output', np.float16).astype(np.float64).reshape(oh, ow, 4)[..., :3]
|
||||
d = a - b
|
||||
print(f'frame {f}: FINAL OUTPUT {ow}x{oh}: {10 * np.log10(1.0 / np.mean(d * d)):.1f} dB PSNR (peak 1.0), max |diff| {np.abs(d).max():.3f}, non-finite {int((~np.isfinite(b)).sum())}', flush=True)
|
||||
prev = {'mvA': f'{D}/mvA', 'half': f'{D}/half', 'feat': f'{D}/feat', 'hist': f'{D}/hist'}
|
||||
@@ -0,0 +1,360 @@
|
||||
// Native regression test for direct game textures, rectangle origins and queued descriptor rotation.
|
||||
// Reuse the standalone replay runner's Vulkan context/readback helpers, without another device harness.
|
||||
#define main d4r_replay_main
|
||||
#include "replay.cpp"
|
||||
#undef main
|
||||
#include "game.h"
|
||||
#include "game_m.h"
|
||||
#include <cmath>
|
||||
|
||||
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<float*>(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<std::unique_ptr<TestImage>, 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<const char*>(results.map) + frame * stride;
|
||||
for (VkDeviceSize i = 0; i < image.upload.size / (image.bits / 8); ++i)
|
||||
if (!std::isfinite(image.bits == 32 ? reinterpret_cast<const float*>(data)[i] : float(reinterpret_cast<const _Float16*>(data)[i])))
|
||||
throw std::runtime_error("nonfinite adapter output");
|
||||
}
|
||||
}
|
||||
std::vector<char> 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<std::unique_ptr<TestImage>, 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<TestImage>(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<TestImage>(c, test.wideMotion ? VK_FORMAT_R32G32_SFLOAT : VK_FORMAT_R16G16_SFLOAT, mw,mh,2,test.wideMotion ? 32 : 16);
|
||||
depth[slot] = std::make_unique<TestImage>(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<TestImage>(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<d4r::Engine> engine;
|
||||
if (bareEngine) engine = std::make_unique<d4r::Engine>(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<char*>(results.map),static_cast<char*>(results.map)+results.size};
|
||||
}
|
||||
std::vector<char> 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<std::unique_ptr<TestImage>, 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<TestImage>(c, VK_FORMAT_R16G16B16A16_SFLOAT, rw,rh,4,16);
|
||||
motion[slot] = std::make_unique<TestImage>(c, VK_FORMAT_R16G16_SFLOAT, rw,rh,2,16);
|
||||
depth[slot] = std::make_unique<TestImage>(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<TestImage>(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<char*>(results.map),static_cast<char*>(results.map)+results.size};
|
||||
}
|
||||
void expectEqual(const std::vector<char>& expected,const std::vector<char>& actual,const std::string& name) {
|
||||
if(expected!=actual) throw std::runtime_error(name+": output mismatch");
|
||||
std::cout<<name<<": byte-identical (six temporal frames)\n";
|
||||
}
|
||||
} // namespace
|
||||
int main(int argc,char** argv) {
|
||||
try {
|
||||
if(argc!=2 && argc!=3) throw std::runtime_error("usage: test-game K_MODEL_DIR [M_MODEL_DIR]");
|
||||
// The reconstruction table NGX itself selected at these sizes (measured through its launch arguments),
|
||||
// on both sides of every threshold and for unequal axis ratios: {render w, h, output w, h, table}.
|
||||
{
|
||||
const uint32_t probes[][5] = {{1280,720,1280,720,4},{1238,696,1280,720,4},{1232,693,1280,720,6},{1030,580,1280,720,6},
|
||||
{1024,576,1280,720,8},{880,495,1280,720,8},{872,491,1280,720,10},{764,430,1280,720,10},{760,428,1280,720,12},
|
||||
{680,383,1280,720,12},{675,380,1280,720,14},{426,240,1280,720,14},{640,720,1280,720,8},{640,648,1280,720,10},
|
||||
{640,432,1280,720,12},{1706,960,2560,1440,10},{1280,720,3840,2160,14}};
|
||||
for (const auto& q : probes)
|
||||
if (d4r::kReconstructionTable(q[2],q[3],q[0],q[1]) != q[4])
|
||||
throw std::runtime_error("reconstruction table selection differs from NGX at render "+std::to_string(q[0])+"x"+std::to_string(q[1]));
|
||||
std::cout<<"reconstruction table selection matches NGX at "<<std::size(probes)<<" measured sizes\n";
|
||||
}
|
||||
Context c;c.create(2,argc==3);d4r::Model model(c.physical,c.device,argv[1],c.cache);
|
||||
if (!model.hasDirectOrigins()) throw std::runtime_error("recompile the model with compile_k.py before testing direct rectangles");
|
||||
c.begin();model.recordUpload(c.command);c.submit();
|
||||
for(bool depthAspect:{false,true}) {
|
||||
Case test;test.depthAspect=depthAspect;
|
||||
auto reference=evaluate(c,model,test,0,false);
|
||||
for(unsigned mask:{1u,2u,4u,8u,7u,15u}) expectEqual(reference,evaluate(c,model,test,mask,true),"direct mask "+std::to_string(mask)+(depthAspect?" D32":" R32"));
|
||||
for (bool queued : {false, true})
|
||||
expectEqual(reference,evaluate(c,model,test,15,queued,false,true),"late fallback allocation "+std::string(queued ? "queued" : "submitted")+(depthAspect ? " D32" : " R32"));
|
||||
if (!model.hasRgb10Output()) throw std::runtime_error("recompile the K model for RGB10A2 checks");
|
||||
for (bool mixed : {false, true}) {
|
||||
test.rgb10Output=true;test.rotateOutputFormats=mixed;
|
||||
const auto packedReference=evaluate(c,model,test,0,false);
|
||||
for (bool queued : {false,true})
|
||||
expectEqual(packedReference,evaluate(c,model,test,15,queued,false,true),
|
||||
std::string(mixed ? "mixed RGBA16F/RGB10A2" : "RGB10A2")+ " late fallback "+(queued ? "queued" : "submitted")+(depthAspect ? " D32" : " R32"));
|
||||
}
|
||||
test.rgb10Output=false;test.rotateOutputFormats=false;
|
||||
test.readonly=true;
|
||||
expectEqual(reference,evaluate(c,model,test,15,true),"read-only input layouts");
|
||||
test.readonly=false;test.wideOutput=true;
|
||||
expectEqual(evaluate(c,model,test,0,false),evaluate(c,model,test,15,true),"RGBA32F output conversion fallback");
|
||||
}
|
||||
// Regular depth picks the nearest neighbour as the smallest depth. With near and far exchanged it must
|
||||
// choose exactly the texels inverted depth chose; inverted depth on the exchanged values must not.
|
||||
if (!model.hasRegularDepth()) throw std::runtime_error("recompile the K model for regular-depth checks");
|
||||
{
|
||||
Case test;
|
||||
const auto reference=evaluate(c,model,test,0,false);
|
||||
test.swapDepth=true;
|
||||
if (evaluate(c,model,test,0,false)==reference) throw std::runtime_error("depth does not affect the output; the regular-depth check is vacuous");
|
||||
test.regularDepth=true;
|
||||
for (unsigned mask:{0u,15u}) expectEqual(reference,evaluate(c,model,test,mask,mask!=0),"regular depth mirrors inverted depth, direct mask "+std::to_string(mask));
|
||||
}
|
||||
// The game's exposure scale multiplies the measured exposure: the adapter must equal an Engine given
|
||||
// the scaled key, and differ from the neutral result.
|
||||
{
|
||||
Case test;test.wideColor=true;test.wideMotion=true;test.readonly=true;
|
||||
const auto neutral=evaluate(c,model,test,15,true);
|
||||
test.exposureScale=.7f;test.preExposure=3.f;
|
||||
const auto scaled=evaluate(c,model,test,15,true);
|
||||
if (scaled==neutral) throw std::runtime_error("the exposure scale and pre-exposure do not affect the output");
|
||||
expectEqual(evaluate(c,model,test,15,false,true),scaled,"exposure scale 0.7 and pre-exposure 3 match the Engine");
|
||||
}
|
||||
// A game exposure texture replaces the measurement (AutoExposure off): the adapter must equal an Engine given
|
||||
// the same value, and the value must reach the output. A non-positive texel is read as 1.
|
||||
if (!model.hasGameExposure()) throw std::runtime_error("recompile the K model for game exposure checks");
|
||||
{
|
||||
Case test;test.wideColor=true;test.wideMotion=true;test.readonly=true;test.gameExposure=true;test.exposureValue=.6f;
|
||||
const auto game=evaluate(c,model,test,15,true);
|
||||
expectEqual(evaluate(c,model,test,15,false,true),game,"game exposure 0.6 matches the Engine");
|
||||
test.exposureValue=1.f;
|
||||
const auto unit=evaluate(c,model,test,15,true);
|
||||
if (game==unit) throw std::runtime_error("the game exposure does not affect the output");
|
||||
test.exposureValue=0.f;
|
||||
expectEqual(evaluate(c,model,test,15,true),unit,"non-positive game exposure reads as 1");
|
||||
}
|
||||
Case precise;precise.wideColor=true;precise.wideMotion=true;precise.readonly=true;
|
||||
expectEqual(evaluate(c,model,precise,15,false,true),evaluate(c,model,precise,15,true),"RGBA32F color/RG32F motion match unquantized Engine");
|
||||
// Display-resolution motion (flag 4): output-sized vectors in output pixels. The render-resolution
|
||||
// result must differ, or the check is vacuous. The direct/copy comparisons use the same input formats on
|
||||
// both sides (a wide source may only be borrowed, not copied: the fallback copy quantizes to RGBA16F/RG16F).
|
||||
if (!model.hasDisplayMotion()) throw std::runtime_error("recompile the K model with compile_k.py for display-resolution motion checks");
|
||||
{
|
||||
Case display;display.displayMotion=true;
|
||||
const auto reference=evaluate(c,model,display,0,false);
|
||||
Case render;
|
||||
if (evaluate(c,model,render,0,false)==reference) throw std::runtime_error("display-resolution motion does not affect the output; the check is vacuous");
|
||||
for (unsigned mask:{2u,15u}) expectEqual(reference,evaluate(c,model,display,mask,true),"display-resolution motion direct mask "+std::to_string(mask));
|
||||
expectEqual(reference,evaluate(c,model,display,15,false,true),"display-resolution motion matches the Engine");
|
||||
Case wide;wide.displayMotion=true;wide.wideColor=true;wide.wideMotion=true;wide.readonly=true;
|
||||
expectEqual(evaluate(c,model,wide,15,false,true),evaluate(c,model,wide,15,true),"display-resolution motion RGBA32F/RG32F match unquantized Engine");
|
||||
}
|
||||
if (argc==3) {
|
||||
d4r::MModel m(c.physical,c.device,argv[2],c.cache);
|
||||
c.begin();m.recordUpload(c.command);c.submit();
|
||||
for (bool depthAspect:{false,true}) {
|
||||
Case test;test.depthAspect=depthAspect;
|
||||
auto reference=evaluateM(c,m,test,0,false);
|
||||
for (unsigned mask:{1u,2u,4u,8u,7u,15u})
|
||||
expectEqual(reference,evaluateM(c,m,test,mask,true),"M Quality direct mask "+std::to_string(mask)+(depthAspect?" D32":" R32"));
|
||||
for (bool queued : {false, true})
|
||||
expectEqual(reference,evaluateM(c,m,test,15,queued,true),"M late fallback allocation "+std::string(queued ? "queued" : "submitted")+(depthAspect ? " D32" : " R32"));
|
||||
if (!m.hasRgb10Output()) throw std::runtime_error("recompile the M model for RGB10A2 checks");
|
||||
for (bool mixed : {false,true}) {
|
||||
test.rgb10Output=true;test.rotateOutputFormats=mixed;
|
||||
const auto packedReference=evaluateM(c,m,test,0,false);
|
||||
for (bool queued : {false,true})
|
||||
expectEqual(packedReference,evaluateM(c,m,test,15,queued,true),
|
||||
std::string(mixed ? "M mixed RGBA16F/RGB10A2" : "M RGB10A2")+" late fallback "+(queued ? "queued" : "submitted")+(depthAspect ? " D32" : " R32"));
|
||||
}
|
||||
test.rgb10Output=false;test.rotateOutputFormats=false;
|
||||
test.readonly=true;
|
||||
expectEqual(reference,evaluateM(c,m,test,15,true),"M read-only input layouts");
|
||||
test.readonly=false;test.wideOutput=true;
|
||||
expectEqual(evaluateM(c,m,test,0,false),evaluateM(c,m,test,15,true),"M RGBA32F output conversion fallback");
|
||||
}
|
||||
}
|
||||
std::cout<<"All direct-texture checks passed\n";
|
||||
return 0;
|
||||
}catch(const std::exception& e){std::cerr<<"test-game: "<<e.what()<<'\n';return 1;}
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Run the engine's own K input kernel on a captured launch of NVIDIA's and compare the network input it writes.
|
||||
|
||||
usage: test_in.py CAPTURE_DIR FRAME WORK_DIR KVK [repeats]
|
||||
"""
|
||||
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_input_depthinv_mvlo_ldr_v2_rel' if LDR else 'hiluma_engine_input_depthinv_mvlo_hdr_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)
|
||||
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('<f', P, 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
|
||||
F = np.float32
|
||||
a, b, c = expo_tex(208), expo_tex(200), expo_tex(216)
|
||||
f152, f158 = F(f1(144)) * F(a), F(f1(144)) * F(b)
|
||||
expo = F(f1(156)) * f158; expo = F(1) if expo == 0 else expo
|
||||
expo_in = F(f1(148)) * F(c) / f158
|
||||
ratio = f158 / f152
|
||||
blob = struct.pack('<14f', *f2(40), *f2(24), *f2(8), *f2(16), *f2(120), *f2(136), *f2(128))
|
||||
blob += struct.pack('<12i', *i2(72), *i2(80), *i2(88), *i2(96), *i2(104), *i2(112))
|
||||
blob += struct.pack('<iI4f', struct.unpack_from('<i', P, 64)[0], P[32], float(expo), float(expo_in), float(ratio), f1(160))
|
||||
open(f'{work}/params.bin', 'wb').write(blob)
|
||||
tokW = struct.unpack_from('<i', P, 64)[0]
|
||||
kind, shape, fmt, fref = res[240]
|
||||
nbytes = os.path.getsize(fref)
|
||||
tokH = 256 if tokW == 320 else None
|
||||
# IN_TEX_OVERRIDE="184=FILE,192=FILE" replaces captured textures; IN_REF compares against that file (the translated
|
||||
# reference kernel's output on the same inputs) instead of the captured buffer.
|
||||
for kv in filter(None, os.environ.get('IN_TEX_OVERRIDE', '').split(',')):
|
||||
o, fn = kv.split('='); res[int(o)] = res[int(o)][:3] + (fn,)
|
||||
fref = os.environ.get('IN_REF', fref)
|
||||
lines = [f'param {work}/params.bin']
|
||||
dims = {}
|
||||
for i, (off, filt) in enumerate(((168, 'nearest'), (176, 'linear'), (184, 'nearest'), (192, 'nearest'), (224, 'linear'))):
|
||||
kind, shape, fmt, f = res[off]; w, h, ch = (int(x) for x in shape.split('x'))
|
||||
dims[off] = (w, h)
|
||||
lines.append(f'tex {i + 1} {f} {w} {h} {ch} {32 if fmt == 32 else 16} {filt}')
|
||||
tokH = dims[224][1]
|
||||
lines.insert(1, f'grid {tokW // 8} {tokH // 8}')
|
||||
lines.append(f'outbuf {work}/out.bin {tokW * tokH * 32}')
|
||||
open(f'{work}/manifest', 'w').write('\n'.join(lines) + '\n')
|
||||
spv = f'{work}/input_k.spv'
|
||||
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *(['-DLDR=1'] if LDR else []), f'{HERE}/input_k.comp', '-o', spv], check=True, stdout=subprocess.DEVNULL)
|
||||
# K requires a 32-lane subgroup (the reconstruction's 2x2 "bright" quad is a subgroup shuffle and the
|
||||
# network layers use 16x16x16 cooperative matrices); wave64 changes both.
|
||||
r = subprocess.run([kvk, spv, f'{work}/manifest', reps, '32'], 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)
|
||||
ref = np.fromfile(fref, np.float16)[:tokW * tokH * 16].astype(np.float64).reshape(tokH, tokW, 4, 4)
|
||||
got = np.fromfile(f'{work}/out.bin', np.float16)[:tokW * tokH * 16].astype(np.float64).reshape(tokH, tokW, 4, 4)
|
||||
print(f'frame {frame} (launch {seq}), tokens {tokW}x{tokH}, reset {P[32]}, exposure {float(expo):.4g} / {float(expo_in):.4g}, ratio {float(ratio):.6g}')
|
||||
for ch, name in enumerate(('current luma', 'history luma', 'token feature', 'mismatch')):
|
||||
a, b = ref[..., ch], got[..., ch]
|
||||
d = np.abs(a - b); peak = max(np.abs(a).max(), 1e-9)
|
||||
print(f' {name:14s} psnr {10 * np.log10(peak ** 2 / max((d ** 2).mean(), 1e-30)):6.1f} dB max |diff| {d.max():.4g} mean |ref| {np.abs(a).mean():.4g} non-finite {int((~np.isfinite(b)).sum())}')
|
||||
@@ -0,0 +1,53 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Compare the own K input kernel with the translated reference kernel on a captured frame, optionally with the motion
|
||||
and depth replaced by synthetic fields (the captured scenes barely move).
|
||||
|
||||
usage: test_in_synth.py CAPTURE_DIR FRAME WORK_DIR KVK REF_SPV [seed | none]
|
||||
"""
|
||||
import glob, os, re, subprocess, sys
|
||||
import numpy as np
|
||||
K = 'hiluma_engine_input_depthinv_mvlo_hdr_v2_rel'
|
||||
cap, frame, work, kvk, ref_spv = sys.argv[1:6]
|
||||
seed = sys.argv[6] if len(sys.argv) > 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())
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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('<f', P, 160)[0]; s168 = struct.unpack_from('<f', P, 168)[0]
|
||||
expo = np.float32(s168) * np.float32(s160 * eB); expo = 1.0 if expo == 0 else float(expo)
|
||||
ratio = float(np.float32(s160 * eB) / np.float32(s160 * eA))
|
||||
blob = struct.pack('<18f', *f2(40), *f2(16), *f2(56), *f2(48), *f2(64), *f2(0), *f2(8), *f2(144), *f2(152))
|
||||
blob += struct.pack('<16i', *i2(208), *i2(216), *i2(96), *i2(104), *i2(112), *i2(120), *i2(128), *i2(136))
|
||||
blob += struct.pack('<2iI', *i2(80), P[32])
|
||||
blob += struct.pack('<10f', expo, ratio, *struct.unpack_from('<3f', P, 172), h1(188), h1(190), h1(192), h1(184), h1(186))
|
||||
open(f'{work}/params.bin', 'wb').write(blob)
|
||||
flags = (P[89], P[90], P[91])
|
||||
if any(flags):
|
||||
print('note: flag bytes 89..91 =', flags, '- the own kernel does not implement these modes')
|
||||
# OUT_TEX_OVERRIDE="240=FILE,248=FILE" replaces captured textures; OUT_REF_DIR compares against DIR/out-OFF.bin
|
||||
# (outputs of the translated reference kernel on the same inputs) instead of the captured surfaces.
|
||||
for kv in filter(None, os.environ.get('OUT_TEX_OVERRIDE', '').split(',')):
|
||||
o, fn = kv.split('='); res[int(o)] = res[int(o)][:3] + (fn,)
|
||||
lines = [f'param {work}/params.bin', f'grid {grid[1]} {grid[2]}']
|
||||
for i, (off, filt) in enumerate(((200, 'nearest'), (224, 'linear'), (232, 'linear'), (240, 'nearest'), (248, 'nearest'))):
|
||||
kind, shape, fmt, f = res[off]; w, h, c = (int(x) for x in shape.split('x'))
|
||||
lines.append(f'tex {i + 1} {f} {w} {h} {c} {32 if fmt == 32 else 16} {filt}')
|
||||
expect = []
|
||||
for i, off in enumerate((288, 296, 304, 312)):
|
||||
kind, shape, fmt, f = res[off]; w, h, c = (int(x) for x in shape.split('x'))
|
||||
lines.append(f'surf {i + 6} {w} {h} {c} 16 - {work}/out-{off}.bin')
|
||||
expect.append((off, f'{os.environ["OUT_REF_DIR"]}/out-{off}.bin' if 'OUT_REF_DIR' in os.environ else f, c))
|
||||
for off, name in ((272, 'head'), (280, 'lut')):
|
||||
data = np.fromfile(res[off][3], np.uint8)[ptr[off]:]
|
||||
data.tofile(f'{work}/{name}.bin'); lines.append(f'lut {work}/{name}.bin')
|
||||
open(f'{work}/manifest', 'w').write('\n'.join(lines) + '\n')
|
||||
spv = f'{work}/output_k.spv'
|
||||
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *os.environ.get('OUT_DEFS', '').split(), *(['-DLDR=1'] if LDR else []), f'{HERE}/output_k.comp', '-o', spv], check=True, stdout=subprocess.DEVNULL)
|
||||
# The reconstruction's 2x2 "bright" quad is a subgroup shuffle: 32 lanes put the pixel rows in the
|
||||
# subgroup's two halves (x + 16 * (y & 1)), so the test must request the engine's subgroup size.
|
||||
r = subprocess.run([kvk, spv, f'{work}/manifest', reps, '32'], 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)
|
||||
names = {288: 'history colour', 296: 'hi-res luma', 304: 'token feature', 312: 'final output'}
|
||||
print(f'frame {frame} (launch {seq}), exposure scale {expo:.4g}, ratio {ratio:.6g}')
|
||||
for off, f, c in expect:
|
||||
a = np.fromfile(f, np.float16).astype(np.float64); b = np.fromfile(f'{work}/out-{off}.bin', np.float16).astype(np.float64)
|
||||
ok = np.isfinite(a) & np.isfinite(b)
|
||||
d = np.abs(a - b)[ok]
|
||||
peak = np.abs(a[ok]).max()
|
||||
psnr = 10 * np.log10(peak ** 2 / max((d ** 2).mean(), 1e-30))
|
||||
exact = (np.fromfile(f, np.uint16) == np.fromfile(f'{work}/out-{off}.bin', np.uint16)).mean()
|
||||
print(f' {names[off]:15s} psnr {psnr:6.1f} dB exact {100 * exact:5.1f}% max |diff| {d.max():.4g} mean |diff| {d.mean():.3g} peak {peak:.3g} non-finite in own {int((~np.isfinite(b)).sum())}')
|
||||
@@ -0,0 +1,48 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Compare the own K output kernel with the translated reference kernel on a captured frame whose motion and depth are
|
||||
replaced by synthetic fields (the captured scenes barely move, so they do not exercise the reprojection).
|
||||
|
||||
usage: test_out_synth.py CAPTURE_DIR FRAME WORK_DIR KVK REF_SPV [seed]
|
||||
REF_SPV: vulkan/ptx2glsl.py's translation of the reference kernel, compiled (see vulkan/README.md).
|
||||
"""
|
||||
import glob, os, re, subprocess, sys
|
||||
import numpy as np
|
||||
K = 'hiluma_engine_output_depthinv_mvlo_hdr_max_v2_rel'
|
||||
cap, frame, work, kvk, ref_spv = sys.argv[1:6]
|
||||
seed = int(sys.argv[6]) if len(sys.argv) > 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())
|
||||
@@ -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('<f', P, 120)[0], P[168], expo(16), expo(24)))
|
||||
def tex(binding, off, ch, bits, mode, w=rw, h=rh): return f'tex {binding} {res[off]} {w} {h} {ch} {bits} {mode}'
|
||||
lines = [f'param {work}/params.bin', f'grid {(hw + 15) // 16} {(hh + 15) // 16}',
|
||||
tex(1, 0, 4, 16, 'nearest'), tex(2, 40, 4, 8, 'nearest'), tex(3, 32, 4, 8, 'nearest'), tex(4, 48, 4, 16, 'nearest'),
|
||||
tex(5, 56, 4, 16, 'linear', hw, hh), tex(6, 64, 2, 16, 'nearest'), tex(7, 72, 2, 16, 'nearest'), tex(8, 80, 2, 16, 'linear'),
|
||||
f'surf 9 {hw} {hh} 4 16 - {work}/hist.bin', f'surf 10 {(hw + 1) // 2} {(hh + 1) // 2} 4 32 - {work}/half.bin',
|
||||
f'outbuf {work}/packed.bin {hw * hh * 4}']
|
||||
open(f'{work}/manifest', 'w').write('\n'.join(lines) + '\n')
|
||||
spv = f'{work}/post_m.spv'
|
||||
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *TILE, f'{HERE}/post_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[-600:]); sys.exit(1)
|
||||
for l in r.stdout.split('\n'):
|
||||
if 'us per' in l: print(l)
|
||||
print(f'frame {frame} (launch {seq}) {K} {rw}x{rh} -> {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}')
|
||||
@@ -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<PFN_vkGetDeviceProcAddr>(instanceProc(instance, "vkGetDeviceProcAddr"));
|
||||
if (!deviceProc) return false;
|
||||
bool ok = true;
|
||||
ok &= (vkTable.GetPhysicalDeviceFormatProperties = reinterpret_cast<PFN_vkGetPhysicalDeviceFormatProperties>(instanceProc(instance, "vkGetPhysicalDeviceFormatProperties"))) != nullptr;
|
||||
ok &= (vkTable.GetPhysicalDeviceMemoryProperties = reinterpret_cast<PFN_vkGetPhysicalDeviceMemoryProperties>(instanceProc(instance, "vkGetPhysicalDeviceMemoryProperties"))) != nullptr;
|
||||
ok &= (vkTable.GetPhysicalDeviceProperties = reinterpret_cast<PFN_vkGetPhysicalDeviceProperties>(instanceProc(instance, "vkGetPhysicalDeviceProperties"))) != nullptr;
|
||||
ok &= (vkTable.AllocateDescriptorSets = reinterpret_cast<PFN_vkAllocateDescriptorSets>(deviceProc(device, "vkAllocateDescriptorSets"))) != nullptr;
|
||||
ok &= (vkTable.AllocateMemory = reinterpret_cast<PFN_vkAllocateMemory>(deviceProc(device, "vkAllocateMemory"))) != nullptr;
|
||||
ok &= (vkTable.BindBufferMemory = reinterpret_cast<PFN_vkBindBufferMemory>(deviceProc(device, "vkBindBufferMemory"))) != nullptr;
|
||||
ok &= (vkTable.BindImageMemory = reinterpret_cast<PFN_vkBindImageMemory>(deviceProc(device, "vkBindImageMemory"))) != nullptr;
|
||||
ok &= (vkTable.CmdBindDescriptorSets = reinterpret_cast<PFN_vkCmdBindDescriptorSets>(deviceProc(device, "vkCmdBindDescriptorSets"))) != nullptr;
|
||||
ok &= (vkTable.CmdBindPipeline = reinterpret_cast<PFN_vkCmdBindPipeline>(deviceProc(device, "vkCmdBindPipeline"))) != nullptr;
|
||||
ok &= (vkTable.CmdBlitImage = reinterpret_cast<PFN_vkCmdBlitImage>(deviceProc(device, "vkCmdBlitImage"))) != nullptr;
|
||||
ok &= (vkTable.CmdClearColorImage = reinterpret_cast<PFN_vkCmdClearColorImage>(deviceProc(device, "vkCmdClearColorImage"))) != nullptr;
|
||||
ok &= (vkTable.CmdCopyBuffer = reinterpret_cast<PFN_vkCmdCopyBuffer>(deviceProc(device, "vkCmdCopyBuffer"))) != nullptr;
|
||||
ok &= (vkTable.CmdCopyBufferToImage = reinterpret_cast<PFN_vkCmdCopyBufferToImage>(deviceProc(device, "vkCmdCopyBufferToImage"))) != nullptr;
|
||||
ok &= (vkTable.CmdCopyImageToBuffer = reinterpret_cast<PFN_vkCmdCopyImageToBuffer>(deviceProc(device, "vkCmdCopyImageToBuffer"))) != nullptr;
|
||||
ok &= (vkTable.CmdDispatch = reinterpret_cast<PFN_vkCmdDispatch>(deviceProc(device, "vkCmdDispatch"))) != nullptr;
|
||||
ok &= (vkTable.CmdPipelineBarrier = reinterpret_cast<PFN_vkCmdPipelineBarrier>(deviceProc(device, "vkCmdPipelineBarrier"))) != nullptr;
|
||||
ok &= (vkTable.CmdUpdateBuffer = reinterpret_cast<PFN_vkCmdUpdateBuffer>(deviceProc(device, "vkCmdUpdateBuffer"))) != nullptr;
|
||||
ok &= (vkTable.CmdWriteTimestamp = reinterpret_cast<PFN_vkCmdWriteTimestamp>(deviceProc(device, "vkCmdWriteTimestamp"))) != nullptr;
|
||||
ok &= (vkTable.CreateBuffer = reinterpret_cast<PFN_vkCreateBuffer>(deviceProc(device, "vkCreateBuffer"))) != nullptr;
|
||||
ok &= (vkTable.CreateComputePipelines = reinterpret_cast<PFN_vkCreateComputePipelines>(deviceProc(device, "vkCreateComputePipelines"))) != nullptr;
|
||||
ok &= (vkTable.CreateDescriptorPool = reinterpret_cast<PFN_vkCreateDescriptorPool>(deviceProc(device, "vkCreateDescriptorPool"))) != nullptr;
|
||||
ok &= (vkTable.CreateDescriptorSetLayout = reinterpret_cast<PFN_vkCreateDescriptorSetLayout>(deviceProc(device, "vkCreateDescriptorSetLayout"))) != nullptr;
|
||||
ok &= (vkTable.CreateImage = reinterpret_cast<PFN_vkCreateImage>(deviceProc(device, "vkCreateImage"))) != nullptr;
|
||||
ok &= (vkTable.CreateImageView = reinterpret_cast<PFN_vkCreateImageView>(deviceProc(device, "vkCreateImageView"))) != nullptr;
|
||||
ok &= (vkTable.CreatePipelineLayout = reinterpret_cast<PFN_vkCreatePipelineLayout>(deviceProc(device, "vkCreatePipelineLayout"))) != nullptr;
|
||||
ok &= (vkTable.CreateSampler = reinterpret_cast<PFN_vkCreateSampler>(deviceProc(device, "vkCreateSampler"))) != nullptr;
|
||||
ok &= (vkTable.CreateShaderModule = reinterpret_cast<PFN_vkCreateShaderModule>(deviceProc(device, "vkCreateShaderModule"))) != nullptr;
|
||||
ok &= (vkTable.DestroyBuffer = reinterpret_cast<PFN_vkDestroyBuffer>(deviceProc(device, "vkDestroyBuffer"))) != nullptr;
|
||||
ok &= (vkTable.DestroyDescriptorPool = reinterpret_cast<PFN_vkDestroyDescriptorPool>(deviceProc(device, "vkDestroyDescriptorPool"))) != nullptr;
|
||||
ok &= (vkTable.DestroyDescriptorSetLayout = reinterpret_cast<PFN_vkDestroyDescriptorSetLayout>(deviceProc(device, "vkDestroyDescriptorSetLayout"))) != nullptr;
|
||||
ok &= (vkTable.DestroyImage = reinterpret_cast<PFN_vkDestroyImage>(deviceProc(device, "vkDestroyImage"))) != nullptr;
|
||||
ok &= (vkTable.DestroyImageView = reinterpret_cast<PFN_vkDestroyImageView>(deviceProc(device, "vkDestroyImageView"))) != nullptr;
|
||||
ok &= (vkTable.DestroyPipeline = reinterpret_cast<PFN_vkDestroyPipeline>(deviceProc(device, "vkDestroyPipeline"))) != nullptr;
|
||||
ok &= (vkTable.DestroyPipelineLayout = reinterpret_cast<PFN_vkDestroyPipelineLayout>(deviceProc(device, "vkDestroyPipelineLayout"))) != nullptr;
|
||||
ok &= (vkTable.DestroySampler = reinterpret_cast<PFN_vkDestroySampler>(deviceProc(device, "vkDestroySampler"))) != nullptr;
|
||||
ok &= (vkTable.DestroyShaderModule = reinterpret_cast<PFN_vkDestroyShaderModule>(deviceProc(device, "vkDestroyShaderModule"))) != nullptr;
|
||||
ok &= (vkTable.FreeMemory = reinterpret_cast<PFN_vkFreeMemory>(deviceProc(device, "vkFreeMemory"))) != nullptr;
|
||||
ok &= (vkTable.GetBufferMemoryRequirements = reinterpret_cast<PFN_vkGetBufferMemoryRequirements>(deviceProc(device, "vkGetBufferMemoryRequirements"))) != nullptr;
|
||||
ok &= (vkTable.GetImageMemoryRequirements = reinterpret_cast<PFN_vkGetImageMemoryRequirements>(deviceProc(device, "vkGetImageMemoryRequirements"))) != nullptr;
|
||||
ok &= (vkTable.MapMemory = reinterpret_cast<PFN_vkMapMemory>(deviceProc(device, "vkMapMemory"))) != nullptr;
|
||||
ok &= (vkTable.UnmapMemory = reinterpret_cast<PFN_vkUnmapMemory>(deviceProc(device, "vkUnmapMemory"))) != nullptr;
|
||||
ok &= (vkTable.UpdateDescriptorSets = reinterpret_cast<PFN_vkUpdateDescriptorSets>(deviceProc(device, "vkUpdateDescriptorSets"))) != nullptr;
|
||||
return ok;
|
||||
}
|
||||
} // namespace d4r
|
||||
#endif
|
||||
@@ -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 <vulkan/vulkan.h>
|
||||
#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
|
||||
Reference in new issue
Block a user