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,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
|
||||
}
|
||||
Reference in new issue
Block a user