flower/Install/Shader/Source/TAAMain.comp

#version 460

#include "Common.glsl"
#include "Sampler.glsl"

// if disable, will sample cloest velocity 3x3.
// use longest velocity get better motion picture, but tiny blur.
#ifndef USE_LONGEST_VELOCITY_3x3
#define USE_LONGEST_VELOCITY_3x3 0
#endif

// Temporal Anti-Alias With Upsample.
#define LOCAL_SIZE_XY 16
layout (local_size_x = LOCAL_SIZE_XY, local_size_y = LOCAL_SIZE_XY, local_size_z = 1) in;

#define VIEW_DATA_SET 0
#include "ViewData.glsl"

#define FRAME_DATA_SET 2
#include "FrameData.glsl"

layout (set = 1, binding = 0, rgba16f) uniform image2D outHdrColor;
layout (set = 1, binding = 1) uniform sampler2D inDepth;  
layout (set = 1, binding = 2) uniform sampler2D inHistory; 
layout (set = 1, binding = 3) uniform sampler2D inVelocity;
layout (set = 1, binding = 4) uniform sampler2D inHdrColor;

struct TAAPushConstant
{
    uint firstRender;
    uint camMove;
};

layout(push_constant) uniform block
{
	TAAPushConstant pushConstant;
};

bool IsFirstFrame()
{
    return pushConstant.firstRender != 0;
}

bool IsCamMove()
{
    return pushConstant.camMove != 0;
}

// TAA random offset within one pixel range.
// so use 3x3 tap to get safe value.
const ivec2 kPattern3x3[9] = {
    ivec2(-1,-1),
    ivec2(-1, 0),
    ivec2(-1, 1),
    ivec2( 0, 1),
    ivec2( 0, 0),
    ivec2( 0,-1),
    ivec2( 1, 1),
    ivec2( 1, 0),
    ivec2( 1,-1)
};
const float kRpc9 = 1.0f / 9.0f;

// keep 1 pixel border for lds, to keep edge tap safe.
const int kBorderSize = 1;
const int kGroupSize  = LOCAL_SIZE_XY;
const int kLdsLength  = kGroupSize + kBorderSize * 2;
const int kLdsArea    = kLdsLength * kLdsLength;

const float kTinyFloat   = 1e-8f;
const float kMaxFloat16  = 32767.0f;
const float kMaxFloat16u = 65535.0f;

shared vec3  sharedColor[kLdsLength][kLdsLength];
shared float sharedDepth[kLdsLength][kLdsLength];
#if USE_LONGEST_VELOCITY_3x3
    shared vec2  sharedVelocity[kLdsLength][kLdsLength];
#endif

vec3 reinhard(vec3 hdr)
{
    return hdr / (hdr + 1.0f);
}

vec3 ldsLoadColor(ivec2 GId)
{
    GId += ivec2(kBorderSize);
    return sharedColor[GId.x][GId.y];
}

// store color with reinhard tonemmaper.
void ldsStoreColor(ivec2 GId, vec3 color)
{
    sharedColor[GId.x][GId.y] = reinhard(color);
}

float ldsLoadDepth(ivec2 GId)
{
    GId += ivec2(kBorderSize);
    return sharedDepth[GId.x][GId.y];
}

void ldsStoreDepth(ivec2 GId, float depth)
{
    sharedDepth[GId.x][GId.y] = depth;
}

#if USE_LONGEST_VELOCITY_3x3
    vec2 ldsLoadVelocity(ivec2 GId)
    {
        GId += ivec2(kBorderSize);
        return sharedVelocity[GId.x][GId.y];
    }

    void ldsStoreVelocity(ivec2 GId, vec2 v)
    {
        sharedVelocity[GId.x][GId.y] = v;
    }
#endif
    

void storeLdsInfos(ivec2 GId, ivec2 TId, ivec2 size)
{
    TId = clamp(TId, ivec2(0,0), size - ivec2(1,1));

    // store color.
    ldsStoreColor(GId, texelFetch(inHdrColor, TId, 0).rgb);

    // store linear z.
    float linearZ = linearizeDepth(texelFetch(inDepth, TId, 0).r, viewData.camInfo.z, viewData.camInfo.w);
    ldsStoreDepth(GId, linearZ);

    #if USE_LONGEST_VELOCITY_3x3
        vec2 velocity = texelFetch(inVelocity, TId, 0).rg;
        ldsStoreVelocity(GId, velocity);
    #endif
}

void prepareLds(ivec2 topLeft, ivec2 workSize, int groupIndex)
{
    // 4 sample per pixel.
    if(groupIndex < (kLdsArea >> 2)) // 1 / 4 area pixel work.
    {
        // sample [0, 0]
        ivec2 id0 = ivec2(
            groupIndex                     % kLdsLength,                      
            groupIndex                     / kLdsLength);
        storeLdsInfos(id0, topLeft + id0, workSize);

        // sample [0.25,  0.25]
        ivec2 id1 = ivec2(
            (groupIndex + (kLdsArea >> 2))  % kLdsLength, 
            (groupIndex + (kLdsArea >> 2))  / kLdsLength);
        storeLdsInfos(id1, topLeft + id1, workSize);

        // sample [0.5, 0.5]
        ivec2 id2 = ivec2(
            (groupIndex + (kLdsArea >> 1))  % kLdsLength, 
            (groupIndex + (kLdsArea >> 1))  / kLdsLength);
        storeLdsInfos(id2, topLeft + id2, workSize);

        // sample [0.75, 0.75]
        ivec2 id3 = ivec2(
            (groupIndex + kLdsArea * 3 / 4) % kLdsLength, 
            (groupIndex + kLdsArea * 3 / 4) / kLdsLength);
        storeLdsInfos(id3, topLeft + id3, workSize);
    }
}

#if USE_LONGEST_VELOCITY_3x3
    void velocityGetLongest(ivec2 pos, inout float longestVelocity, inout ivec2 longestPos, inout vec2 longestV)
    {
        vec2 v = ldsLoadVelocity(pos);

        float v2 = dot(v.xy, v.xy);
        if(v2 > longestVelocity)
        {
            longestVelocity = v2;
            longestPos = pos;
            longestV = v;
        }
    }

    // 3x3 tap get longest velocity.
    float velocitySample3x3Longest(ivec2 groupPos, ivec2 topLeft, inout vec2 velocity)
    {
        float longestV2  = -1.0f;
        vec2  longestV   = vec2(0);
        ivec2 longestPos = groupPos;

        velocityGetLongest(groupPos + kPattern3x3[0], longestV2, longestPos, longestV);
        velocityGetLongest(groupPos + kPattern3x3[1], longestV2, longestPos, longestV);
        velocityGetLongest(groupPos + kPattern3x3[2], longestV2, longestPos, longestV);
        velocityGetLongest(groupPos + kPattern3x3[3], longestV2, longestPos, longestV);
        velocityGetLongest(groupPos + kPattern3x3[4], longestV2, longestPos, longestV);
        velocityGetLongest(groupPos + kPattern3x3[5], longestV2, longestPos, longestV);
        velocityGetLongest(groupPos + kPattern3x3[6], longestV2, longestPos, longestV);
        velocityGetLongest(groupPos + kPattern3x3[7], longestV2, longestPos, longestV);
        velocityGetLongest(groupPos + kPattern3x3[8], longestV2, longestPos, longestV);

        velocity = longestV;

        // when use longest sample3x3, use center depth.
        return ldsLoadDepth(longestPos);
    }
#endif

// linear z cloest test.
void depthGetClosest(ivec2 pos, inout float cloestDepth, inout ivec2 cloestPos)
{
    float d = ldsLoadDepth(pos);

    if(d < cloestDepth)
    {
        cloestDepth = d;
        cloestPos = pos;
    }
}

// 3x3 tap get closet pos depth.
float velocitySample3x3Closest(ivec2 groupPos, ivec2 topLeft, inout vec2 velocity)
{
    float minDepth = 1.0f;
    ivec2 minPos   = groupPos;

    depthGetClosest(groupPos + kPattern3x3[0], minDepth, minPos);
    depthGetClosest(groupPos + kPattern3x3[1], minDepth, minPos);
    depthGetClosest(groupPos + kPattern3x3[2], minDepth, minPos);
    depthGetClosest(groupPos + kPattern3x3[3], minDepth, minPos);
    depthGetClosest(groupPos + kPattern3x3[4], minDepth, minPos);
    depthGetClosest(groupPos + kPattern3x3[5], minDepth, minPos);
    depthGetClosest(groupPos + kPattern3x3[6], minDepth, minPos);
    depthGetClosest(groupPos + kPattern3x3[7], minDepth, minPos);
    depthGetClosest(groupPos + kPattern3x3[8], minDepth, minPos);

    velocity = texelFetch(inVelocity, topLeft + minPos, 0).xy;

    return minDepth;
}

void main() 
{
    ivec2 workSize = textureSize(inHdrColor, 0).xy;
    ivec2 topLeft  = ivec2(gl_WorkGroupID.xy) * kGroupSize - kBorderSize;
    ivec2 groupPos = ivec2(gl_LocalInvocationID.xy);
    int groupIndex = int(gl_LocalInvocationIndex);

    prepareLds(topLeft, workSize, groupIndex);

    groupMemoryBarrier();
    barrier();

    ivec2 pixelPos = ivec2(gl_GlobalInvocationID.xy);
    if (pixelPos.x >= workSize.x || pixelPos.y >= workSize.y)
    {
        return; 
    }

    vec2 texelSize = 1.0f / vec2(workSize);
    vec2 uv        = (vec2(pixelPos) + vec2(0.5)) * texelSize; 

    // get cloest velocity.
    vec2 velocity;
    
    #if USE_LONGEST_VELOCITY_3x3
        const float cloestDepth = velocitySample3x3Longest(groupPos, topLeft, velocity);
    #else
        const float cloestDepth = velocitySample3x3Closest(groupPos, topLeft, velocity);
    #endif
    
    const bool  bSky = cloestDepth <= BG_DEPTH;

    // reproject uv.
    vec2 reprojectedUV = uv - velocity;

    float velocityLerp = 0.0f;
    if(!IsFirstFrame())
    {
        velocityLerp = texture(inHistory, reprojectedUV).w;
    }

    const float ideaStaticBoxSize = 2.5f;
    float staticBoxSize = IsCamMove() ? ideaStaticBoxSize : mix(0.5f, ideaStaticBoxSize, velocityLerp);
    float boxSize = mix(0.5f, staticBoxSize, bSky ? 0.0f : smoothstep(0.02f, 0.0f, length(velocity)));

    // in center color.
    vec3 colorIn = ldsLoadColor(groupPos);

    // sample history color.
    vec3 colorHistory = catmullRom9Sample(inHistory, reprojectedUV, vec2(workSize));
    colorHistory = reinhard(colorHistory);

    // variance clamp.
    vec3 clampHistory;
    {
        float wsum = 0.0f;
        vec3 vsum  = vec3(0.0f, 0.0f, 0.0f);
        vec3 vsum2 = vec3(0.0f, 0.0f, 0.0f);

        for (int y = -1; y <= 1; ++y)
        {
            for (int x = -1; x <= 1; ++x)
            {
                const vec3 neigh = ldsLoadColor(groupPos + ivec2(x, y));
                const float w = exp(-0.75f * (x * x + y * y));

                vsum2 += neigh * neigh * w;
                vsum  += neigh * w;
                wsum  += w;
            }
        }

        const vec3 ex  = vsum / wsum;
        const vec3 ex2 = vsum2 / wsum;
        const vec3 dev = sqrt(max(ex2 - ex * ex, 0.0f));

        vec3 nmin = ex - dev * boxSize;
        vec3 nmax = ex + dev * boxSize;
        clampHistory = clamp(colorHistory, nmin, nmax);
    }

    // when camera move, use this.
    // when camera don't move, use more bigger blend factor if motion factor check.
    const float ideaLerpFactor = 0.01f;
    float blendFactor = ideaLerpFactor;
    {   
        const float threshold   = 0.5f;
        const float base        = 0.5f;
        const float gather      = 0.1666f;

        // subpixel flicker reduce
        float depth = linearizeDepth(cloestDepth, viewData.camInfo.z, viewData.camInfo.w);
        float texelVelMag    = length(velocity * vec2(workSize)) * depth;
        float subpixelMotion = clamp(threshold / (texelVelMag + kTinyFloat), 0.0f, 1.0f);

        // something moveing
        float dynamicBlendFactor = texelVelMag * base + subpixelMotion * gather;

        // lumiance bias correct.
        float luminanceHistory = luminance(clampHistory);
        float luminanceCurrent = luminance(colorIn);
        float unbiasedDifference = abs(luminanceCurrent - luminanceHistory) / ((max(luminanceCurrent, luminanceHistory) + 0.3));
        dynamicBlendFactor *= 1.0 - unbiasedDifference;

        // clamp
        dynamicBlendFactor = clamp(dynamicBlendFactor, 0.0f, 0.4f);

        // dynamic move.
        float moveFactor = max(abs(velocity.x * workSize.x), abs(velocity.y * workSize.y)) * 100.0f;

        // tiny move, so reset lerp factor.
        velocityLerp = moveFactor > 0.01f ? 0 : velocityLerp;

        // mix lerp factor by frames to get a good clip value.
        velocityLerp = mix(velocityLerp, 1.0f, ideaLerpFactor);

        float lerpFactor = clamp(1.0f - velocityLerp, 0.0f, 1.0f);
        lerpFactor *= lerpFactor;
        lerpFactor *= lerpFactor;
        lerpFactor *= lerpFactor;
        lerpFactor *= lerpFactor;
        lerpFactor = smoothstep(0.0f, 1.0f, lerpFactor);

        blendFactor = bSky ? blendFactor : mix(blendFactor, dynamicBlendFactor, lerpFactor);
    }

    vec3 colorResolve = mix(clampHistory, colorIn, blendFactor);

    // half16 safe clamp.
    colorResolve = min(vec3(65504.0f), colorResolve);

    imageStore(outHdrColor, ivec2(gl_GlobalInvocationID.xy), vec4(colorResolve, velocityLerp));
}