engine/rr: native Ray Reconstruction stages with capture/replay and shim integration
- engine/rr: d4r's own shaders for the denoiser frame (exposure, six encoders,
five decoders, hkpn output stage) at NGX layouts, plus PTX->GLSL/HIP packing
scripts, reference implementations and weight extraction
- runtime.h/runtime_rr.inc, game_rr.{h,cpp}: GameDenoiser wired into the engine
build (d4r-rr replay binary) and the D3D12 shim's feature-13 path
- shim: answer GetFeatureRequirements from d4r's own state only (the old
cross-binary CUDA handoff crashed OptiScaler's scan); drop the core's
CUDA GetFeatureRequirements import
- shim/harness: exact guide capture/replay of a frame's RR inputs
- kernels/rr: expanded enc0/engine prototypes; kernels/tools: ptxsim rewrite
plus rrsim and an engine test; vulkan: ptx2glsl FP8/cooperative-matrix
support, kvk inout buffers and KVK_DISABLE_OPTIMIZATION
- bridge: report the surface array in cuSurfObjectCreate traces
- packaging: RR settings and the ICU launch overrides in D4R_README
This commit is contained in:
woutermans committed
1 parent
30472ae195
commit
50fa65d87b
67 files changed
+11926
-504
No files matched your search
@@ -0,0 +1,295 @@
|
||||
#version 460
|
||||
// Native decoder 1 of the DLSS-D (preset E) denoiser, register-resident: cuda_dldn_engine_swin_dec1_kernel.
|
||||
// engine/rr/dec1_ref.py's arithmetic (bit-exact to NVIDIA except surface 112's bilinear taps, see dec0_m.comp);
|
||||
// binding 1 holds engine/rr/dec1_pack.py's packing. The expand and prologue are dec0_m.comp's (64 channels from dec2's
|
||||
// 96), the block is enc1_m.comp's (whose header explains the layouts), the outputs are the tokens (argument 48) and
|
||||
// dec0_m.comp's filter on the 64 -> 16 head.
|
||||
#extension GL_KHR_cooperative_matrix : require
|
||||
#extension GL_KHR_memory_scope_semantics : require
|
||||
#extension GL_EXT_shader_explicit_arithmetic_types : require
|
||||
#extension GL_EXT_shader_16bit_storage : require
|
||||
#extension GL_EXT_shader_8bit_storage : require
|
||||
#extension GL_EXT_float_e4m3 : require
|
||||
#extension GL_KHR_shader_subgroup_basic : require
|
||||
#extension GL_KHR_shader_subgroup_shuffle : require
|
||||
#extension GL_EXT_control_flow_attributes : require
|
||||
#extension GL_GOOGLE_include_directive : require
|
||||
layout(local_size_x = 128) in;
|
||||
#include "dec1_m_w.h"
|
||||
#include "rr_exact.glsl"
|
||||
|
||||
layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param;
|
||||
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
|
||||
layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16;
|
||||
layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u;
|
||||
layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
|
||||
layout(set = 0, binding = 3) uniform sampler2D T88;
|
||||
layout(set = 0, binding = 4) uniform sampler2D T96;
|
||||
layout(set = 0, binding = 5) uniform sampler2D T104;
|
||||
layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S112;
|
||||
|
||||
#define SKIP 0
|
||||
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
|
||||
#define HAC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
|
||||
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
|
||||
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
|
||||
#define RM gl_CooperativeMatrixLayoutRowMajor
|
||||
#define CM gl_CooperativeMatrixLayoutColumnMajor
|
||||
#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM)
|
||||
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
|
||||
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
|
||||
#define Q8(dst, src) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (dst)[i_] = q8s((src)[i_]);
|
||||
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) }
|
||||
const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0;
|
||||
|
||||
// shared memory, aliased by phase: expand [LOW | EX], attention [KS | VS], outputs [TAP | OUTF]
|
||||
shared uint POOL[48 * 24 + 48 * 64];
|
||||
#define LOW POOL // 48 low-resolution tokens (6 x 6, padded): 96 FP8 channels, arena order (24 words)
|
||||
#define EX_AT (48u * 24u) // their expansion: 4 sub-positions x 64 FP8 channels in slot order (64 words)
|
||||
#define KS POOL // the current head's K: token slot x 32 FP8 channels
|
||||
#define VS_AT 768u // the current head's V^T: 32 channels x 96 token slots
|
||||
#define RR_KS POOL
|
||||
#define RR_KS0 0u
|
||||
#define RR_VS POOL
|
||||
#define RR_VS0 VS_AT
|
||||
shared uint XT[96 * 16]; // the normalised key tokens: slot-major, 64 FP8 channels in slot order
|
||||
shared float16_t TAP[100 * 4]; // texture 88 bilinear at the window's 10 x 10 filter tap positions
|
||||
shared float16_t OUTF[64 * 16]; // the head per window token (raster order)
|
||||
|
||||
ivec2 keyXY(uint t, int x0, int y0)
|
||||
{
|
||||
if (t < 64u) return ivec2(x0 + int(t & 7u), y0 + int(t >> 3u));
|
||||
const int i = int(t & 7u), side = int((t - 64u) >> 3u);
|
||||
return ivec2(side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8), side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i));
|
||||
}
|
||||
int clampTap(int v, int n) { return v < 0 ? 1 : (v >= n ? n - 2 : v); }
|
||||
uint pack4(floate4m3_t a, floate4m3_t b, floate4m3_t c, floate4m3_t d)
|
||||
{
|
||||
return code(a) | (code(b) << 8u) | (code(c) << 16u) | (code(d) << 24u);
|
||||
}
|
||||
float16_t e4(uint w, uint b) { return float16_t(uintBitsToFloate4m3EXT(uint8_t((w >> (8u * b)) & 0xffu))); }
|
||||
|
||||
void main()
|
||||
{
|
||||
const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID;
|
||||
const uint hl = lane >> 4u, col = lane & 15u;
|
||||
RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u);
|
||||
const int ex = int(param.p[0]), ey = int(param.p[1]), hx = ex / 2, hy = ey / 2;
|
||||
const uint loOff = param.p[2], skOff = param.p[6], outOff = param.p[12], shift = param.p[20];
|
||||
const int x0 = 8 * int(gl_WorkGroupID.x) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u);
|
||||
const int X0 = (x0 >> 1) - 1, Y0 = (y0 >> 1) - 1;
|
||||
const uint lplane = uint(hx * hy) * 32u, plane = uint(ex * ey) * 32u;
|
||||
|
||||
// ---------------------------------------------------------------- expand: wave w computes sub-position w
|
||||
for (uint i = tid; i < 48u * 24u; i += 128u) {
|
||||
const uint r = i / 24u, wd = i % 24u;
|
||||
uint u = 0u;
|
||||
if (r < 36u) {
|
||||
const int X = mirror(X0 + int(r % 6u), hx), Y = mirror(Y0 + int(r / 6u), hy);
|
||||
u = arena.w[(loOff + (wd >> 3u) * lplane + uint(Y * hx + X) * 32u) / 4u + (wd & 7u)];
|
||||
}
|
||||
LOW[i] = u;
|
||||
}
|
||||
barrier();
|
||||
[[unroll]] for (uint n = 0u; n < 3u; ++n) {
|
||||
B8 lb[6];
|
||||
[[unroll]] for (uint kt = 0u; kt < 6u; ++kt) coopMatLoad(lb[kt], POOL, n * 384u + 4u * kt, 24u, CM);
|
||||
[[unroll]] for (uint mm = 0u; mm < 4u; ++mm) {
|
||||
const uint m = 4u * wv + mm;
|
||||
ACC a;
|
||||
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_BE / 2u + 16u * m + 8u * hl + uint(i)]);
|
||||
[[unroll]] for (uint kt = 0u; kt < 6u; kt += 2u) {
|
||||
A8 w0, w1;
|
||||
LDA(w0, W_WE, 96u, m, kt); LDA(w1, W_WE, 96u, m, kt + 1u);
|
||||
K32(a, w0, lb[kt], w1, lb[kt + 1u])
|
||||
}
|
||||
B8 q;
|
||||
Q8(q, a)
|
||||
coopMatStore(q, POOL, EX_AT + n * 1024u + 4u * m, 64u, CM);
|
||||
}
|
||||
}
|
||||
barrier();
|
||||
|
||||
// ---------------------------------------------------------------- prologue: one key token per thread
|
||||
if (tid < 96u) {
|
||||
const ivec2 k = keyXY(inter(tid), x0, y0);
|
||||
const uint r = uint(((k.y >> 1) - Y0) * 6 + ((k.x >> 1) - X0)), s = uint(2 * (k.y & 1) + (k.x & 1));
|
||||
const int xx = mirror(k.x, ex), yy = mirror(k.y, ey);
|
||||
const uint sb = (skOff + uint(yy * ex + xx) * 32u) / 4u;
|
||||
float16_t v[64];
|
||||
[[unroll]] for (uint wd = 0u; wd < 16u; ++wd) {
|
||||
const uint e = POOL[EX_AT + r * 64u + 16u * s + wd], sk = arena.w[sb + (wd >> 3u) * (plane >> 2u) + (wd & 7u)];
|
||||
[[unroll]] for (uint b = 0u; b < 4u; ++b) { precise float16_t x = e4(e, b) + e4(sk, b); v[4u * wd + b] = x; }
|
||||
}
|
||||
// enc1's input norm: per word t (and t + 4) the f16x2 pair sums of both planes, then lanes ^1, ^2, the halves
|
||||
precise f16vec2 h[4];
|
||||
[[unroll]] for (uint q = 0u; q < 4u; ++q) {
|
||||
precise f16vec2 p[4];
|
||||
[[unroll]] for (uint kk = 0u; kk < 4u; ++kk) {
|
||||
const uint o = 32u * (kk & 1u) + 4u * (q + 4u * (kk >> 1u));
|
||||
precise f16vec2 a = f16vec2(v[o], v[o + 1u]) * f16vec2(v[o], v[o + 1u]);
|
||||
precise f16vec2 c = f16vec2(v[o + 2u], v[o + 3u]) * f16vec2(v[o + 2u], v[o + 3u]);
|
||||
p[kk] = a + c;
|
||||
}
|
||||
precise f16vec2 s0 = p[0] + p[1], s1 = p[2] + p[3];
|
||||
h[q] = s0 + s1;
|
||||
}
|
||||
precise f16vec2 h0 = h[0] + h[1], h2_ = h[2] + h[3];
|
||||
precise f16vec2 hs = h0 + h2_;
|
||||
precise float16_t sm = hs.x + hs.y;
|
||||
precise float16_t se = sm + float16_t(0.0001220703125);
|
||||
const float16_t rr = rsq16(se);
|
||||
[[unroll]] for (uint wd = 0u; wd < 16u; ++wd) {
|
||||
floate4m3_t c[4];
|
||||
[[unroll]] for (uint b = 0u; b < 4u; ++b) {
|
||||
const uint j = 4u * wd + b;
|
||||
precise float16_t g = rr * w16.w[W_GAIN0 / 2u + inter(j)];
|
||||
precise float16_t n = v[j] * g;
|
||||
c[b] = q8h(n);
|
||||
}
|
||||
XT[tid * 16u + wd] = pack4(c[0], c[1], c[2], c[3]);
|
||||
}
|
||||
}
|
||||
barrier();
|
||||
|
||||
// ---------------------------------------------------------------- the residual: expanded + skip + table 7
|
||||
ACC R[4];
|
||||
{
|
||||
const uint tq = inter(16u * wv + col);
|
||||
const int qx = x0 + int(tq & 7u), qy = y0 + int(tq >> 3u);
|
||||
const uint r = uint(((qy >> 1) - Y0) * 6 + ((qx >> 1) - X0)), s = uint(2 * (qy & 1) + (qx & 1));
|
||||
const uint sb = (skOff + uint(mirror(qy, ey) * ex + mirror(qx, ex)) * 32u) / 4u;
|
||||
[[unroll]] for (uint m = 0u; m < 4u; ++m) {
|
||||
const uint ea = EX_AT + r * 64u + 16u * s + 4u * m + 2u * hl;
|
||||
const uint sa = sb + (m >> 1u) * (plane >> 2u) + 4u * (m & 1u) + 2u * hl;
|
||||
const uvec2 e = uvec2(POOL[ea], POOL[ea + 1u]), sk = uvec2(arena.w[sa], arena.w[sa + 1u]);
|
||||
[[unroll]] for (int i = 0; i < 8; ++i) {
|
||||
precise float16_t x = e4(e[i >> 2], uint(i) & 3u) + e4(sk[i >> 2], uint(i) & 3u);
|
||||
precise float16_t c = x + w16.w[W_TABLE7 / 2u + 16u * m + 8u * hl + uint(i)];
|
||||
R[m][i] = float(c);
|
||||
}
|
||||
}
|
||||
}
|
||||
barrier(); // EX is read; the attention reuses its memory
|
||||
|
||||
// ---------------------------------------------------------------- attention, two heads
|
||||
[[unroll]] for (uint hd = 0u; hd < ((SKIP & 64) != 0 ? 0u : 2u); ++hd) {
|
||||
// this head's Q/K and V, shared: 12 jobs (token tile n, channel tile c), three per wave. K = (QK^T)^T rows
|
||||
// to KS (token slot x 32 channels, also Q^T read column-major), V^T to VS (32 channels x 96 token slots)
|
||||
if (hd == 1u) barrier(); // every wave is done with head 0's KS / VS
|
||||
for (uint j = 0u; j < ((SKIP & 128) != 0 ? 0u : 3u); ++j) {
|
||||
const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u;
|
||||
A8 wk[4];
|
||||
B8 wvb[4], xb[4];
|
||||
A8 xa[4];
|
||||
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) {
|
||||
coopMatLoad(xb[kt], XT, n * 256u + 4u * kt, 16u, CM);
|
||||
coopMatLoad(xa[kt], XT, n * 256u + 4u * kt, 16u, RM);
|
||||
LDA(wk[kt], W_WK0 + HEAD * hd, 64u, c, kt);
|
||||
LDB(wvb[kt], W_WV0 + HEAD * hd, 32u, kt, c);
|
||||
}
|
||||
ACC a = ACC(0.0);
|
||||
K32(a, wk[0], xb[0], wk[1], xb[1])
|
||||
K32(a, wk[2], xb[2], wk[3], xb[3])
|
||||
A8 ka;
|
||||
Q8(ka, a)
|
||||
coopMatStore(ka, KS, n * 128u + 4u * c, 8u, RM);
|
||||
ACC v = ACC(0.0);
|
||||
K32(v, xa[0], wvb[0], xa[1], wvb[1])
|
||||
K32(v, xa[2], wvb[2], xa[3], wvb[3])
|
||||
A8 va;
|
||||
Q8(va, v)
|
||||
coopMatStore(va, POOL, VS_AT + c * 384u + 4u * n, 24u, RM);
|
||||
}
|
||||
barrier();
|
||||
#include "rr_attn.glsl"
|
||||
[[unroll]] for (uint m = 0u; m < 4u; ++m) {
|
||||
A8 a0, a1;
|
||||
LDA(a0, W_WO0 + HEAD * hd, 32u, m, 0u); LDA(a1, W_WO0 + HEAD * hd, 32u, m, 1u);
|
||||
K32(R[m], a0, Ob[0], a1, Ob[1])
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------- MLP
|
||||
B8 Nb[4];
|
||||
{
|
||||
float16_t rv[4][8];
|
||||
[[unroll]] for (uint m = 0u; m < 4u; ++m) [[unroll]] for (int i = 0; i < 8; ++i) rv[m][i] = float16_t(R[m][i]);
|
||||
// slot 16m + 8hl + i is NVIDIA's channel 8j + 2t + c with j = 4(m >> 1) + (m & 1) + 2((i >> 1) & 1),
|
||||
// t = 2hl + (i >> 2), c = i & 1 (enc1_ref.nv_rowsum_sq's tree)
|
||||
float16_t h1[2];
|
||||
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
|
||||
float16_t ht[2];
|
||||
[[unroll]] for (uint tq = 0u; tq < 2u; ++tq) {
|
||||
const uint i0 = 4u * tq + c, i1 = i0 + 2u;
|
||||
precise float16_t a = rv[0][i0] * rv[0][i0] + rv[0][i1] * rv[0][i1];
|
||||
precise float16_t b = rv[1][i0] * rv[1][i0] + rv[1][i1] * rv[1][i1];
|
||||
precise float16_t cc = rv[2][i0] * rv[2][i0] + rv[2][i1] * rv[2][i1];
|
||||
precise float16_t d = rv[3][i0] * rv[3][i0] + rv[3][i1] * rv[3][i1];
|
||||
precise float16_t hv = (a + cc) + (b + d);
|
||||
ht[tq] = hv;
|
||||
}
|
||||
precise float16_t hh = ht[0] + ht[1];
|
||||
h1[c] = hh;
|
||||
}
|
||||
const uint other = subgroupShuffleXor(bh(h1[0]) | (bh(h1[1]) << 16u), 16u);
|
||||
precise float16_t g0 = h1[0] + hb(other & 0xffffu), g1 = h1[1] + hb(other >> 16u);
|
||||
precise float16_t s = g0 + g1;
|
||||
precise float16_t se = s + float16_t(0.0001220703125);
|
||||
const float16_t r = rsq16(se);
|
||||
[[unroll]] for (uint m = 0u; m < 4u; ++m)
|
||||
[[unroll]] for (int i = 0; i < 8; ++i) {
|
||||
const uint k = 16u * m + 8u * hl + uint(i);
|
||||
precise float16_t g = r * w16.w[W_GAIN14 / 2u + k];
|
||||
precise float16_t n = rv[m][i] * g;
|
||||
Nb[m][i] = q8h(n);
|
||||
precise float16_t a = rv[m][i] + w16.w[W_TABLE16 / 2u + k];
|
||||
R[m][i] = float(a);
|
||||
}
|
||||
}
|
||||
for (uint br = 0u; br < ((SKIP & 4) != 0 ? 1u : 8u); ++br) {
|
||||
const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br;
|
||||
B8 Hb[2];
|
||||
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
|
||||
ACC h;
|
||||
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = float(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]);
|
||||
A8 a0, a1, a2, a3;
|
||||
LDA(a0, w1, 64u, c, 0u); LDA(a1, w1, 64u, c, 1u); LDA(a2, w1, 64u, c, 2u); LDA(a3, w1, 64u, c, 3u);
|
||||
K32(h, a0, Nb[0], a1, Nb[1])
|
||||
K32(h, a2, Nb[2], a3, Nb[3])
|
||||
[[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = q8h(act(float16_t(h[i])));
|
||||
}
|
||||
[[unroll]] for (uint m = 0u; m < 4u; ++m) {
|
||||
A8 a0, a1;
|
||||
LDA(a0, w2, 32u, m, 0u); LDA(a1, w2, 32u, m, 1u);
|
||||
K32(R[m], a0, Hb[0], a1, Hb[1])
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------- outputs
|
||||
{
|
||||
const uint t = inter(16u * wv + col);
|
||||
const int tx = x0 + int(t & 7u), ty = y0 + int(t >> 3u);
|
||||
const bool inside = tx >= 0 && ty >= 0 && tx < ex && ty < ey;
|
||||
B8 Tb[4];
|
||||
const uint o = (outOff + uint(ty * ex + tx) * 32u) >> 2u;
|
||||
[[unroll]] for (uint m = 0u; m < 4u; ++m) {
|
||||
Q8(Tb[m], R[m])
|
||||
const uint at = 4u * (m & 1u) + 2u * hl;
|
||||
if (inside) {
|
||||
arena.w[o + (m >> 1u) * (plane >> 2u) + at] = pack4(Tb[m][0], Tb[m][1], Tb[m][2], Tb[m][3]);
|
||||
arena.w[o + (m >> 1u) * (plane >> 2u) + at + 1u] = pack4(Tb[m][4], Tb[m][5], Tb[m][6], Tb[m][7]);
|
||||
}
|
||||
}
|
||||
// the 64 -> 16 head, per token to OUTF
|
||||
ACC hv;
|
||||
[[unroll]] for (int i = 0; i < 8; ++i) hv[i] = float(w16.w[W_BH / 2u + 8u * hl + uint(i)]);
|
||||
A8 a0, a1, a2, a3;
|
||||
LDA(a0, W_WH, 64u, 0u, 0u); LDA(a1, W_WH, 64u, 0u, 1u); LDA(a2, W_WH, 64u, 0u, 2u); LDA(a3, W_WH, 64u, 0u, 3u);
|
||||
K32(hv, a0, Tb[0], a1, Tb[1])
|
||||
K32(hv, a2, Tb[2], a3, Tb[3])
|
||||
[[unroll]] for (int i = 0; i < 8; ++i) OUTF[t * 16u + 8u * hl + uint(i)] = float16_t(hv[i]);
|
||||
}
|
||||
barrier();
|
||||
#include "rr_decfilter.glsl"
|
||||
}
|
||||
Reference in new issue
Block a user