d4r/kernels/support/fsr_rr_prepare.comp

#version 450
// FSR Ray Regeneration backend, stage 1: a game's DLSS Ray Reconstruction inputs -> the inputs of
// AMD's FidelityFX denoiser 1.2 (ffx_denoiser.h), at render resolution.
//
// The split of the game's single noisy colour into a diffuse and a specular signal follows the
// DLSS-RR -> FSR-RR adapter of Zach Hembree's OptiScaler branch ffx-denoise-experimental (GPL-3.0,
// 18ac686, FSRDInputConv.hlsl, "Mode 2") and burak113's continuation (5a6adb1): each lobe gets the
// colour in proportion to its albedo and is demodulated by it. That is an approximation - one
// colour cannot be split exactly - and what does not survive demodulation goes to the residual,
// which bypasses the denoiser. Unlike that reference this writes FSR-RR 1.2's contract: signed
// linear depth reconstructed from hardware or linear depth, motion as previous-minus-current UV,
// linear (non-gamma) albedo and the colour before particles as the denoised base.
layout(local_size_x = 8, local_size_y = 8) in;

layout(binding = 0) uniform sampler2D inColor;
layout(binding = 1) uniform sampler2D inDepth;
layout(binding = 2) uniform sampler2D inMotion;
layout(binding = 3) uniform sampler2D inNormals;
layout(binding = 4) uniform sampler2D inRoughness;       // the normals again when roughness is packed
layout(binding = 5) uniform sampler2D inDiffuseAlbedo;
layout(binding = 6) uniform sampler2D inSpecularAlbedo;
layout(binding = 7) uniform sampler2D inSpecularHit;     // bound to the normals when absent
layout(binding = 8) uniform sampler2D inDiffuseHit;      // bound to the normals when absent
layout(binding = 9) uniform sampler2D inBeforeParticles; // bound to the colour when absent

layout(binding = 10, r32f) uniform writeonly image2D outLinearDepth;
layout(binding = 11, rgba16f) uniform writeonly image2D outMotion;
layout(binding = 12, rgba16f) uniform writeonly image2D outNormals;
layout(binding = 13, rgba16f) uniform writeonly image2D outDiffuseAlbedo;
layout(binding = 14, rgba16f) uniform writeonly image2D outSpecularAlbedo;
layout(binding = 15, rgba16f) uniform writeonly image2D outDiffuse;
layout(binding = 16, rgba16f) uniform writeonly image2D outSpecular;
layout(binding = 17, rgba16f) uniform writeonly image2D outResidual;

// Matrices are uploaded as the row-major, row-vector (v' = vM) float[16] NGX and FFX use. Read as a
// GLSL (column-major) mat4 that is the transpose, so M * v here is the row-vector product v * M.
layout(push_constant) uniform Constants
{
    mat4 clipToView;          // inverse(ViewToClipMatrix)
    mat4 viewToPreviousView;  // inverse(WorldToView) * previous WorldToView
    vec4 renderSize;          // width, height, 1 / width, 1 / height
    vec4 motion;              // xy: MV texel -> UV delta (MV.Scale / MV resolution); zw: MV texel per render pixel
    uvec4 colorDepthBase;     // colour subrect origin xy, depth subrect origin xy
    uvec4 motionBaseFlags;    // motion subrect origin xy, flags, unused
    vec4 depthParams;         // x: sign of view z in front of the camera, y: sky depth (hardware), z: far marker
    vec4 detail;              // x: share of the colour that bypasses the denoiser (0: none); y: smallest albedo a
                              // signal is divided by
} pc;

const uint kHardwareDepth = 1u;
const uint kRoughnessPacked = 2u;
const uint kSpecularHit = 4u;
const uint kSpecularHitInW = 8u;   // DLSSD.SpecularRayDirectionHitDistance: the distance is .w
const uint kDiffuseHit = 16u;
const uint kDiffuseHitInW = 32u;
const uint kNormalsUnorm = 64u;    // [0,1]-encoded normals (UNORM formats)
const uint kBeforeParticles = 128u;
const uint kInvalidNormalPassthrough = 256u;

const float kHalfMax = 65000.0;

vec3 safe(vec3 v)
{
    // NaN fails both comparisons and becomes 0
    return vec3(v.x >= 0.0 ? min(v.x, kHalfMax) : 0.0, v.y >= 0.0 ? min(v.y, kHalfMax) : 0.0,
                v.z >= 0.0 ? min(v.z, kHalfMax) : 0.0);
}

// AMD's NormalToOctahedronUv (FidelityFX denoising.md 5.3.1)
vec2 octahedron(vec3 n)
{
    n /= abs(n.x) + abs(n.y) + abs(n.z);
    vec2 k = vec2(n.x >= 0.0 ? 1.0 : -1.0, n.y >= 0.0 ? 1.0 : -1.0);
    vec2 xy = n.z < 0.0 ? (1.0 - abs(n.yx)) * k : n.xy;
    return xy * 0.5 + 0.5;
}

void main()
{
    const ivec2 px = ivec2(gl_GlobalInvocationID.xy);
    if (px.x >= int(pc.renderSize.x) || px.y >= int(pc.renderSize.y))
        return;
    const uint flags = pc.motionBaseFlags.z;
    const vec2 uv = (vec2(px) + 0.5) * pc.renderSize.zw;
    const vec2 ndc = vec2(uv.x * 2.0 - 1.0, 1.0 - uv.y * 2.0);

    // Signed view-space position. Hardware depth goes through the inverse projection, which carries
    // the game's handedness and (reversed) depth range; linear depth is the distance along the view
    // axis, given its sign by the projection.
    const float depth = texelFetch(inDepth, ivec2(pc.colorDepthBase.zw) + px, 0).r;
    vec3 viewPos;
    bool sky;
    if ((flags & kHardwareDepth) != 0u)
    {
        const vec4 p = pc.clipToView * vec4(ndc, depth, 1.0);
        viewPos = p.xyz / p.w;
        sky = depth == pc.depthParams.y || !(abs(p.w) > 0.0);
    }
    else
    {
        const vec4 p = pc.clipToView * vec4(ndc, 0.5, 1.0);
        const vec3 ray = p.xyz / p.w;
        const float z = depth * pc.depthParams.x;
        viewPos = ray * (z / ray.z);
        sky = !(depth > 0.0) || depth >= pc.depthParams.z;
    }

    const vec3 color = safe(texelFetch(inColor, ivec2(pc.colorDepthBase.xy) + px, 0).rgb);
    const vec3 base = (flags & kBeforeParticles) != 0u ? safe(texelFetch(inBeforeParticles, px, 0).rgb) : color;
    // Particles (colour after minus before) stay out of the denoiser; signed, so smoke can darken.
    const vec3 particles = color - base;

    vec3 diffuseAlbedo = clamp(texelFetch(inDiffuseAlbedo, px, 0).rgb, 0.0, 1.0);
    vec3 specularAlbedo = clamp(texelFetch(inSpecularAlbedo, px, 0).rgb, 0.0, 1.0);
    // Energy above 1 is taken from the diffuse lobe first (reference: albedo overshoot clamp).
    diffuseAlbedo = max(diffuseAlbedo - max(diffuseAlbedo + specularAlbedo - 1.0, 0.0), 0.0);
    const float albedoSum = dot(diffuseAlbedo + specularAlbedo, vec3(1.0));
    sky = sky || !(albedoSum > 3e-3) || any(isnan(viewPos)) || any(isinf(viewPos));
    const vec4 normalTexel = texelFetch(inNormals, px, 0);
    vec3 normal = (flags & kNormalsUnorm) != 0u ? normalTexel.xyz * 2.0 - 1.0 : normalTexel.xyz;
    const float normalLength = dot(normal, normal);
    // A zero/invalid normal is no surface, not an invented world-Z-facing mirror. Alan Wake 2
    // supplies zero normals for its sky with finite FP16-max depth, so the far bound alone misses it.
    if ((flags & kInvalidNormalPassthrough) != 0u)
        sky = sky || !(normalLength > 1e-8) || isnan(normalLength) || isinf(normalLength);

    if (sky)
    {
        // Outside linearDepthBounds: the denoiser passes these texels through, and the residual carries
        // the whole colour.
        imageStore(outLinearDepth, px, vec4(pc.depthParams.z * 10.0 * pc.depthParams.x));
        imageStore(outMotion, px, vec4(0.0));
        imageStore(outNormals, px, vec4(0.5, 0.5, 1.0, 0.0));
        imageStore(outDiffuseAlbedo, px, vec4(0.0));
        imageStore(outSpecularAlbedo, px, vec4(0.0));
        imageStore(outDiffuse, px, vec4(0.0));
        imageStore(outSpecular, px, vec4(0.0));
        imageStore(outResidual, px, vec4(color, 0.0));
        return;
    }
    imageStore(outLinearDepth, px, vec4(viewPos.z));

    // 2.5D motion: DLSS motion vectors point from this pixel to where it was (previous - current), in
    // the MV texture's pixels. The depth delta is the camera's: the previous view's z of this point
    // (moving geometry is not reprojected; there is no previous depth for it).
    const ivec2 mvTexel = ivec2(pc.motionBaseFlags.xy) + ivec2(vec2(px) * pc.motion.zw);
    const vec2 mv = texelFetch(inMotion, mvTexel, 0).xy * pc.motion.xy;
    const vec4 previous = pc.viewToPreviousView * vec4(viewPos, 1.0);
    imageStore(outMotion, px, vec4(mv, previous.z / previous.w - viewPos.z, 0.0));

    normal = dot(normal, normal) > 1e-8 ? normalize(normal) : vec3(0.0, 0.0, 1.0);
    const float roughness = clamp((flags & kRoughnessPacked) != 0u ? normalTexel.w : texelFetch(inRoughness, px, 0).r,
                                  0.0, 1.0);
    // material type 0 (the default material) everywhere
    imageStore(outNormals, px, vec4(octahedron(normal), roughness, 0.0));
    imageStore(outDiffuseAlbedo, px, vec4(diffuseAlbedo, 0.0));
    imageStore(outSpecularAlbedo, px, vec4(specularAlbedo, 0.0));

    // Split by albedo and demodulate (reference "Mode 2"). Per channel, so a tinted metal keeps its tint.
    const vec3 total = diffuseAlbedo + specularAlbedo;
    const vec3 specularShare = specularAlbedo / max(total, vec3(1e-4));
    const vec3 specularColor = base * specularShare;
    const vec3 diffuseColor = base - specularColor;
    // A dark albedo would turn light it did not shape (fog, sky light, the background at a thin object's edge) into
    // a huge signal that the denoiser clips (fog lost) or spreads (bright fringes): no albedo below the floor
    // divides a signal, and what the division cannot carry stays in the residual.
    vec3 diffuse = safe(diffuseColor / max(diffuseAlbedo, vec3(pc.detail.y)));
    vec3 specular = safe(specularColor / max(specularAlbedo, vec3(pc.detail.y)));
    if (pc.detail.x > 0.0)
    {
        // Fine detail the denoiser treats as noise (hair, water ripples): part of the colour goes around it,
        // through the residual.
        const float denoised = 1.0 - clamp(pc.detail.x, 0.0, 1.0);
        diffuse *= denoised;
        specular *= denoised;
    }
    const vec3 remodulated = diffuse * diffuseAlbedo + specular * specularAlbedo;

    // A: the ray hit distance for an indirect signal (ignored, but non-negative, for a direct one)
    float specularHit = 0.0, diffuseHit = 0.0;
    if ((flags & kSpecularHit) != 0u)
    {
        const vec4 hit = texelFetch(inSpecularHit, px, 0);
        specularHit = clamp((flags & kSpecularHitInW) != 0u ? hit.w : hit.x, 0.0, kHalfMax);
    }
    if ((flags & kDiffuseHit) != 0u)
    {
        const vec4 hit = texelFetch(inDiffuseHit, px, 0);
        diffuseHit = clamp((flags & kDiffuseHitInW) != 0u ? hit.w : hit.x, 0.0, kHalfMax);
    }
    imageStore(outDiffuse, px, vec4(diffuse, diffuseHit));
    imageStore(outSpecular, px, vec4(specular, specularHit));
    imageStore(outResidual, px, vec4(base - remodulated + particles, 0.0));
}