woutermans
50fa65d87b
- 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
299 lines
12 KiB
Plaintext
299 lines
12 KiB
Plaintext
#version 460
|
|
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
|
|
#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
|
|
#extension GL_EXT_shader_16bit_storage : require
|
|
|
|
// Ray Reconstruction hkpn output stage, preset E (no diamond-wallaby branch).
|
|
// One workgroup per 16x16 output tile. Stage 1 stages the render-space neighbourhood of
|
|
// the tile into shared memory (colour log, per-texel weights, anisotropy); stage 2 gathers
|
|
// a 5x5 elliptical kernel per output pixel and writes the two output surfaces.
|
|
// Parameter block (binding 0) holds NVIDIA's kernel argument block; word index = byte / 4.
|
|
|
|
layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in;
|
|
|
|
layout(std430, binding = 0) readonly buffer Params { uint P[]; };
|
|
layout(std430, binding = 2) buffer Arena { uint A[]; };
|
|
layout(binding = 3) uniform sampler2D colourTex; // arg 176, 640x360 rgba16f
|
|
layout(binding = 5) uniform sampler2D exposureTex; // arg 200, 1x1 rgba16f
|
|
layout(binding = 8) uniform sampler2D albedoA; // arg 280, 640x360 rgba8unorm
|
|
layout(binding = 9) uniform sampler2D albedoB; // arg 288, 640x360 rgba8unorm
|
|
layout(binding = 11) uniform sampler2D decTex; // arg 448, 640x384 rgba16f
|
|
layout(rgba16f, binding = 12) uniform writeonly image2D outLog; // arg 224
|
|
layout(rgba16f, binding = 13) uniform writeonly image2D outLin; // arg 240
|
|
|
|
#define F_LOG2E uintBitsToFloat(0x3FB8AA3Bu)
|
|
#define F_2LOG2E uintBitsToFloat(0x4038AA3Bu)
|
|
#define F_LOG10_2 uintBitsToFloat(0x3E9A209Bu)
|
|
#define F_0_36 uintBitsToFloat(0x3EB851ECu)
|
|
#define F_100 uintBitsToFloat(0x42C80000u)
|
|
#define F_1E4 uintBitsToFloat(0x38D1B717u)
|
|
#define F_0_005 uintBitsToFloat(0x3BA3D70Au)
|
|
#define F_0_1 uintBitsToFloat(0x3DCCCCCDu)
|
|
#define F_1_7 uintBitsToFloat(0x3FD9999Au)
|
|
#define F_0_17 uintBitsToFloat(0x3E2E147Bu)
|
|
#define F_0_95 uintBitsToFloat(0x3F733333u)
|
|
#define F_0_6 uintBitsToFloat(0x3F19999Au)
|
|
#define F_9 uintBitsToFloat(0x41102CB4u)
|
|
|
|
// shared tile: 21 entries per row, up to 21x21 entries (render scale <= 1)
|
|
#define TILE_STRIDE 21
|
|
// binary16 numbers, stored as such: half the shared-memory traffic of the 5x5 gather
|
|
shared f16vec4 tLog[TILE_STRIDE * TILE_STRIDE]; // (L_R, L_G, L_B, alpha = 0) log-domain colour
|
|
shared f16vec4 tSig[TILE_STRIDE * TILE_STRIDE]; // (s0, s1, s2, s3) sigmoid of arena bytes 0..3
|
|
shared vec4 tAn[TILE_STRIDE * TILE_STRIDE]; // (p0, p1, p2, p3) anisotropy parameters
|
|
|
|
float pf(uint idx) { return uintBitsToFloat(P[idx]); }
|
|
|
|
// e4m3 byte -> float (exact); e4m3fn NaN for 0x7F/0xFF.
|
|
float e4m3(uint b) {
|
|
uint e = (b >> 3u) & 15u;
|
|
uint m = b & 7u;
|
|
if (e == 15u && m == 7u) return uintBitsToFloat(0x7FC00000u);
|
|
float v = (e == 0u) ? float(m) * uintBitsToFloat(0x3B000000u)
|
|
: (1.0 + float(m) * 0.125) * exp2(float(e) - 7.0);
|
|
return ((b & 128u) != 0u) ? -v : v;
|
|
}
|
|
|
|
// round to nearest even binary16, returned as f32: the hardware conversion (RNE), kept from being folded back to x by
|
|
// OR-ing an opaque zero into its bits (RADV folds a register-only f32 -> f16 -> f32)
|
|
// Beyond the binary16 range NVIDIA's arithmetic keeps a finite rounded value: that case takes the explicit rounding.
|
|
uint16_t zero16;
|
|
float f16rnSlow(float x) {
|
|
if (x == 0.0 || isnan(x) || isinf(x)) return x;
|
|
int e;
|
|
frexp(x, e);
|
|
int ee = max(e - 11, -24);
|
|
float q = uintBitsToFloat(uint(ee + 127) << 23);
|
|
return roundEven(x / q) * q;
|
|
}
|
|
float f16rn(float x) {
|
|
if (abs(x) >= 65504.0) return f16rnSlow(x);
|
|
return float(uint16BitsToFloat16(float16BitsToUint16(float16_t(x)) | zero16));
|
|
}
|
|
|
|
// Sigmoid as evaluated in binary16 on this lowering: r = f16(2^(-x*log2e)), den = f16(1+r),
|
|
// s = f16(1/den), with one Newton step for small s and the observed fix-ups of r.
|
|
float sig16(float e) {
|
|
float x = f16rn(e);
|
|
float r = f16rn(exp2(-x * F_LOG2E));
|
|
uint hb = uint(float16BitsToUint16(float16_t(r)));
|
|
if (hb == 0x1F79u || hb == 0x25CFu) {
|
|
r = f16rn(r - 0.0009765625);
|
|
} else if (hb == 0xC13Bu) {
|
|
r = f16rn(r + 0.00006103515625);
|
|
} else if (hb == 0xC1EFu) {
|
|
r = f16rn(r + 0.000030517578125);
|
|
}
|
|
float d32 = f16rn(1.0 + r);
|
|
float rc = 1.0 / d32;
|
|
float s16 = f16rn(rc);
|
|
uint sb = uint(float16BitsToUint16(float16_t(s16))) & 0x7FFFu;
|
|
if (sb != 0u && sb < 0x8Fu) {
|
|
float ee = fma(-d32, s16, 1.0);
|
|
s16 = f16rn(fma(rc, ee, s16));
|
|
}
|
|
return s16;
|
|
}
|
|
|
|
// tanh as evaluated: polynomial for small |x|, exponential form otherwise.
|
|
float tanhApprox(float x) {
|
|
float ax = abs(x);
|
|
float x2 = x * x;
|
|
float q = fma(uintBitsToFloat(0x3C80F082u), x2, uintBitsToFloat(0xBD563CAEu));
|
|
q = fma(q, x2, uintBitsToFloat(0x3E085941u));
|
|
q = fma(q, x2, uintBitsToFloat(0xBEAAA9EDu));
|
|
q = fma(q, x2, 0.0);
|
|
float small = fma(q, x, x);
|
|
float ex = exp2(ax * F_2LOG2E);
|
|
float big = fma(1.0 / (ex + 1.0), -2.0, 1.0);
|
|
if (ax >= F_9) big = 1.0;
|
|
float bigSigned = uintBitsToFloat((floatBitsToUint(x) & 0x80000000u) | floatBitsToUint(big));
|
|
return (ax >= F_0_6) ? bigSigned : small;
|
|
}
|
|
|
|
// round half away from zero
|
|
int roundAway(float v) {
|
|
return int(trunc(v + ((v < 0.0) ? -0.5 : 0.5)));
|
|
}
|
|
|
|
void main() {
|
|
zero16 = uint16_t(P[0] >> 24u); // the output width is below 2^24
|
|
// ---- parameters (byte offsets / 4) ----
|
|
uint outW = P[0];
|
|
uint outH = P[1];
|
|
float sx = pf(6u); // 24: render/out scale x
|
|
float sy = pf(7u); // 28
|
|
float ox = pf(10u); // 40: jitter offset x
|
|
float oy = pf(11u); // 44
|
|
float expScale = pf(37u); // 148
|
|
uint decMaxX = P[14] - 1u; // 56: dec0 width
|
|
uint decMaxY = P[15] - 1u; // 60: dec0 height
|
|
uint colMaxX = P[35]; // 140: colour/albedo max x
|
|
uint colMaxY = P[36]; // 144: colour/albedo max y
|
|
uint arenaBase = P[64]; // 256: arena offset
|
|
uint arenaPitch = P[4]; // 16: render width (640)
|
|
|
|
// ---- tile origin / size in render space ----
|
|
uvec2 wg = gl_WorkGroupID.xy;
|
|
float tileX = ox + fma(float(wg.x * 16u) + 0.5, sx, -0.5);
|
|
float tileY = oy + fma(float(wg.y * 16u) + 0.5, sy, -0.5);
|
|
int c0 = int(floor(tileX)) - 2;
|
|
int r0 = int(floor(tileY)) - 2;
|
|
int tw = int(ceil(16.0 * sx)) + 5;
|
|
int th = int(ceil(16.0 * sy)) + 5;
|
|
|
|
// ---- stage 1: tile load ----
|
|
uint total = uint(tw * th);
|
|
for (uint idx = gl_LocalInvocationIndex; idx < total; idx += 256u) {
|
|
int col = int(idx) % tw;
|
|
int row = int(idx) / tw;
|
|
int ix = clamp(c0 + col, 0, int(decMaxX));
|
|
int iy = clamp(r0 + row, 0, int(decMaxY));
|
|
|
|
vec3 d = texelFetch(decTex, ivec2(ix, iy), 0).rgb;
|
|
int ax = min(ix, int(colMaxX));
|
|
int ay = min(iy, int(colMaxY));
|
|
vec3 alb = texelFetch(albedoA, ivec2(ax, ay), 0).rgb
|
|
+ texelFetch(albedoB, ivec2(ax, ay), 0).rgb;
|
|
float albSum = (alb.r + alb.g) + alb.b;
|
|
if (albSum != 0.0) {
|
|
d.r *= max(alb.r, F_0_005);
|
|
d.g *= max(alb.g, F_0_005);
|
|
d.b *= max(alb.b, F_0_005);
|
|
}
|
|
vec3 lg;
|
|
lg.r = f16rn(F_0_36 * (log2(max(fma(d.r, F_100, 1.0), F_1E4)) * F_LOG10_2));
|
|
lg.g = f16rn(F_0_36 * (log2(max(fma(d.g, F_100, 1.0), F_1E4)) * F_LOG10_2));
|
|
lg.b = f16rn(F_0_36 * (log2(max(fma(d.b, F_100, 1.0), F_1E4)) * F_LOG10_2));
|
|
|
|
uint wordIdx = (arenaBase >> 2u) + uint(arenaPitch * uint(iy) + uint(ix)) * 2u;
|
|
uint lo = A[wordIdx];
|
|
uint hi = A[wordIdx + 1u];
|
|
vec4 sg = vec4(sig16(e4m3(lo & 0xFFu)),
|
|
sig16(e4m3((lo >> 8u) & 0xFFu)),
|
|
sig16(e4m3((lo >> 16u) & 0xFFu)),
|
|
sig16(e4m3(lo >> 24u)));
|
|
vec4 an = vec4(e4m3(hi & 0xFFu),
|
|
e4m3((hi >> 8u) & 0xFFu),
|
|
e4m3((hi >> 16u) & 0xFFu),
|
|
e4m3(hi >> 24u));
|
|
|
|
uint se = uint(row * TILE_STRIDE + col);
|
|
tLog[se] = f16vec4(lg, 0.0);
|
|
tSig[se] = f16vec4(sg);
|
|
tAn[se] = an;
|
|
}
|
|
barrier();
|
|
|
|
// ---- stage 2: per output pixel ----
|
|
uvec2 pix = gl_GlobalInvocationID.xy;
|
|
if (pix.x >= outW || pix.y >= outH) return;
|
|
|
|
float expo = texelFetch(exposureTex, ivec2(0, 0), 0).r;
|
|
float e41 = ((expo == 0.0) ? 1.0 : expo) * expScale;
|
|
|
|
float fx = (float(pix.x) + 0.5) * sx;
|
|
float fy = (float(pix.y) + 0.5) * sy;
|
|
float cx = fx + ox;
|
|
float cy = fy + oy;
|
|
|
|
// colour fetch: log of exposure-scaled colour at the nearest render texel
|
|
int ux = clamp(roundAway((fx - 0.5) + ox), 0, int(colMaxX));
|
|
int uy = clamp(roundAway((fy - 0.5) + oy), 0, int(colMaxY));
|
|
vec3 col = texelFetch(colourTex, ivec2(ux, uy), 0).rgb;
|
|
vec3 gcol;
|
|
gcol.r = F_0_36 * (log2(max(fma(col.r * e41, F_100, 1.0), F_1E4)) * F_LOG10_2);
|
|
gcol.g = F_0_36 * (log2(max(fma(col.g * e41, F_100, 1.0), F_1E4)) * F_LOG10_2);
|
|
gcol.b = F_0_36 * (log2(max(fma(col.b * e41, F_100, 1.0), F_1E4)) * F_LOG10_2);
|
|
|
|
// centre entry, anisotropy
|
|
float flx = floor(cx);
|
|
float fly = floor(cy);
|
|
int ix0 = int(flx) - c0;
|
|
int iy0 = int(fly) - r0;
|
|
float dx0 = (flx + 0.5) - cx;
|
|
float dy0 = (fly + 0.5) - cy;
|
|
float ks = (sy >= 0.5) ? 1.0 : (0.5 / sy);
|
|
|
|
vec4 an = tAn[iy0 * TILE_STRIDE + ix0];
|
|
float sa = fma(1.0 / (1.0 + exp2(-F_LOG2E * an.x)), F_1_7, F_0_17);
|
|
float sb = fma(1.0 / (1.0 + exp2(-F_LOG2E * an.y)), F_1_7, F_0_17);
|
|
float rho = F_0_95 * tanhApprox(an.z);
|
|
|
|
float sa2 = sa * sa;
|
|
float sb2 = sb * sb;
|
|
float denom = (sa2 * sb2) * (1.0 - rho * rho);
|
|
float invDenom = -(1.0 / denom);
|
|
float cross2 = rho * (2.0 * sqrt(sa2 * sb2));
|
|
|
|
// 5x5 gather
|
|
float sumw = 0.0;
|
|
vec3 accL = vec3(0.0);
|
|
vec4 accS = vec4(0.0);
|
|
for (int c = -2; c <= 2; c++) {
|
|
for (int r = -2; r <= 2; r++) {
|
|
float dx = ks * (dx0 + float(c));
|
|
float dy = ks * (dy0 + float(r));
|
|
uint se = uint((iy0 + r) * TILE_STRIDE + (ix0 + c));
|
|
precise float t_c = (cross2 * dx) * dy;
|
|
precise float t_b = dx * (sb2 * dx);
|
|
precise float t_a = dy * (sa2 * dy);
|
|
precise float num = (t_c - t_b) - t_a;
|
|
precise float Q = invDenom * num;
|
|
float w = exp2((Q * -0.5) * F_LOG2E);
|
|
sumw += w;
|
|
vec4 lv = vec4(tLog[se]);
|
|
vec4 sv = vec4(tSig[se]);
|
|
accL = fma(vec3(w), lv.rgb, accL);
|
|
accS = fma(vec4(w), sv, accS);
|
|
}
|
|
}
|
|
|
|
float inv = 1.0 / sumw;
|
|
vec3 Lm = accL * inv;
|
|
float S0 = accS.x * inv;
|
|
float S1 = accS.y * inv;
|
|
float S2 = accS.z * inv;
|
|
|
|
// blend weight
|
|
float q = F_0_1 / (sx * sy);
|
|
float t1 = sumw * q;
|
|
float t2 = S0 * t1;
|
|
float t3 = (sx * sy) * t2;
|
|
float t4 = 4.0 * t3;
|
|
float t5 = (sy < 0.5) ? t4 : t2;
|
|
float wgt = (t5 < 0.0) ? 0.0 : ((t5 > 1.0) ? 1.0 : t5);
|
|
|
|
float fA = fma(1.0 - S1, wgt, S1);
|
|
float fB = fma(1.0 - S2, wgt, S2);
|
|
|
|
// surface 224 (E): fA * D + (1 - fA) * gcol; D = Lm (f16 tile log channels),
|
|
// gcol = log of the exposure-scaled colour texel (PTX linear texel fetch path); D.a = 0 so aA = 1 - fA
|
|
vec3 cA = fma(Lm, vec3(fA), gcol * (1.0 - fA));
|
|
float aA = 1.0 - fA;
|
|
float sumA = cA.r + cA.g + cA.b + aA;
|
|
if (isnan(sumA) || isinf(sumA)) {
|
|
cA = Lm;
|
|
aA = 0.0;
|
|
}
|
|
imageStore(outLog, ivec2(pix), vec4(f16rn(cA.r), f16rn(cA.g), f16rn(cA.b), f16rn(aA)));
|
|
|
|
// surface 240 (E, PTX BB0_83): {R,G,B,1} of lin = (10^(clamp(c)/0.36)-1)/100/e41,
|
|
// c = fma(Lm, fB, gcol*(1-fB)); c = Lm when the sum is not finite; stored RNE f16
|
|
vec3 cB = fma(Lm, vec3(fB), gcol * (1.0 - fB));
|
|
float sumB = cB.r + cB.g + cB.b;
|
|
if (isnan(sumB) || isinf(sumB)) cB = Lm;
|
|
cB = clamp(cB, vec3(F_1E4), vec3(3.0));
|
|
// PTX divides by constants and takes lg2(10) with the hardware approximations at run time; keep the
|
|
// constants opaque so the compiler cannot fold them with different rounding.
|
|
uint zero = P[0] >> 24u;
|
|
float c036 = uintBitsToFloat(0x3EB851ECu | zero);
|
|
float c100 = uintBitsToFloat(0x42C80000u | zero);
|
|
float lg10 = log2(uintBitsToFloat(0x41200000u | zero));
|
|
float rc036 = 1.0 / c036;
|
|
float rc100 = 1.0 / c100;
|
|
float rce41 = 1.0 / e41;
|
|
precise vec3 lin = (exp2((cB * rc036) * lg10) - 1.0) * rc100 * rce41;
|
|
imageStore(outLin, ivec2(pix), vec4(float(float16_t(lin.r)), float(float16_t(lin.g)), float(float16_t(lin.b)), 1.0));
|
|
}
|