ayo
7c5d3c9157
Add native K/M/RR inference engines, HIP kernels, replay and validation harnesses, and Windows DLSS setup and runtime support. Include GPU profiling tools, build and packaging scripts, platform integration patches, and implementation documentation.
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));
|
|
}
|