shaders/computeshader/lsr_lock.comp

#version 460 core
#extension GL_ARB_separate_shader_objects : enable
#extension GL_ARB_shading_language_420pack : enable

// Lume Super Resolution
// Create new locks mask

// includes
#include "render/shaders/common/render_color_conversion_common.h"
#include "render/shaders/common/render_post_process_structs_common.h"

#include "common/bloom_common.h"

// camera data
struct DefaultCameraMatrixStruct {
    mat4 view;
    mat4 proj;
    mat4 viewProj;

    mat4 viewInv;
    mat4 projInv;
    mat4 viewProjInv;

    mat4 viewPrevFrame;
    mat4 projPrevFrame;
    mat4 viewProjPrevFrame;

    mat4 shadowViewProj;
    mat4 shadowViewProjInv;

    // .xy = jitter offset, .zw = jitter offset with baked screen size
    vec4 jitter;
    vec4 jitterPrevFrame;

    // .xy = unique id (64-bit), .zw = layer mask (64 bit)
    uvec4 indices;
    // .x multi-view camera additional layer count, .yzw 3 multi-view camera indices
    // yzw are packed, use unpack functions
    uvec4 multiViewIndices;

    vec4 frustumPlanes[2];

    // .x environment count
    uvec4 counts;
    // padding to 256
    uvec4 pad0;
    mat4 matPad0;
    mat4 matPad1;
};

struct LockPassPushConstant {
    vec4 renderSizeInvSize;
    vec4 displaySizeInvSize;
};

// sets

layout(set = 0, binding = 0) uniform texture2D lockInputLumaTex;
layout(set = 0, binding = 1, r8) uniform writeonly image2D newLocksTex;
layout(set = 0, binding = 3, r32ui) uniform writeonly uimage2D prevDepthTex;
layout(set = 0, binding = 2) uniform sampler uSampler;

layout(push_constant, std430) uniform uPushConstantBlock
{
    LockPassPushConstant uPc;
};

layout(set = 0, binding = 12, std140) uniform uCameraMatrices
{
    DefaultCameraMatrixStruct uCameras[16];
};

layout(constant_id = 0) const uint CORE_POST_PROCESS_FLAGS = 0;

// Constants
const float cSimilarThreshold = 1.05f;

//---------------- Utility Functions ------------------//

ivec2 RenderSize() {
    return ivec2(uPc.renderSizeInvSize.x, uPc.renderSizeInvSize.y);
}

ivec2 DisplaySize() {
    return ivec2(uPc.displaySizeInvSize.x, uPc.displaySizeInvSize.y);
}

vec2 RenderSizeVec2() { return uPc.renderSizeInvSize.xy; }
vec2 DisplaySizeVec2() { return uPc.displaySizeInvSize.xy; }

bool OnScreen(ivec2 pos, ivec2 size) {
    return (pos.x >= 0 && pos.y >= 0 && pos.x < size.x && pos.y < size.y);
}

vec2 JitterPixels() { return uCameras[0].jitter.xy; }
//vec2 JitterPixels() { return vec2(0); }

// Load Luma from the input texture calculated in the Dilate/Reconstruct pass
float LoadLockInputLuma(ivec2 px)
{
    ivec2 sz = RenderSize();
    px       = clamp(px, ivec2(0), sz - ivec2(1));
    return texelFetch(sampler2D(lockInputLumaTex, uSampler), px, 0).r;
}

ivec2 HrPosFromLrPos(vec2 iPxLrPos)
{
    vec2 fSrcJitteredPos = vec2(iPxLrPos) + 0.5f - JitterPixels();
    vec2 fLrPosInHr = (fSrcJitteredPos / RenderSizeVec2()) * DisplaySizeVec2();
    ivec2 iPxHrPos = ivec2(floor(fLrPosInHr));
    return iPxHrPos;
}

void StoreNewLocks(ivec2 coord, float lockValue) {
    imageStore(newLocksTex, coord, vec4(lockValue, 0, 0, 0));
}


void GatherLockInputLumaRQuad(vec2 baseUv,
                              out float col00,
                              out float col10,
                              out float col01,
                              out float col11)
{
    vec4 g = textureGather(sampler2D(lockInputLumaTex, uSampler), baseUv, 0);

    col01 = g.w;
    col11 = g.z;
    col10 = g.y;
    col00 = g.x;
}

bool ComputeThinFeatureConfidence(ivec2 pos)
{
    const float kSimilarThreshold = 1.05;
    const float kFp16Max          = 65504.0;
    const uint  kSetBit4          = 1u << 4;

    const uint rejectionMasks[4] = {
        (1u<<0)|(1u<<1)|(1u<<3)|(1u<<4),   // upper-left  2�2
        (1u<<1)|(1u<<2)|(1u<<4)|(1u<<5),   // upper-right 2�2
        (1u<<3)|(1u<<4)|(1u<<6)|(1u<<7),   // lower-left  2�2
        (1u<<4)|(1u<<5)|(1u<<7)|(1u<<8)    // lower-right 2�2
    };


    float nucleusLuma      = LoadLockInputLuma(pos);
    float dissimMin        = kFp16Max;
    float dissimMax        = 0.0;
    uint  similarityMask   = kSetBit4;

    //------------------------------------------------------------------
    //  sample neighbourhood    0 1 2
    //                          3 4 5
    //                          6 7 8
    //------------------------------------------------------------------
    float luma[9];
    float dummy;

    vec2 invInputSize = 1.0 / vec2(RenderSize());
    vec2 baseUv       = vec2(pos) * invInputSize;
    vec2 unitUv       = invInputSize;          // (1 px, 1 px) in UV

    GatherLockInputLumaRQuad(baseUv,
                             luma[0], luma[1],
                             luma[3], luma[4]);
    GatherLockInputLumaRQuad(baseUv + unitUv,
                             dummy,
                             luma[5],
                             luma[7], luma[8]);

    luma[2] = LoadLockInputLuma(pos + ivec2( 1,-1));
    luma[6] = LoadLockInputLuma(pos + ivec2(-1, 1));

    int idx = 0;
    for (int y =-1; y <= 1; ++y)
        for (int x =-1; x <= 1; ++x, ++idx)
        {
            if (x==0 && y==0) continue;

            float s = luma[idx];
            float ratio = max(s, nucleusLuma) / min(s, nucleusLuma);

            if (ratio < kSimilarThreshold)
            {
                similarityMask |= 1u << idx;
            }
            else
            {
                dissimMin = min(dissimMin, s);
                dissimMax = max(dissimMax, s);
            }
        }

    bool isRidge = (nucleusLuma > dissimMax) || (nucleusLuma < dissimMin);
    if (!isRidge) return false;


    for (int i = 0; i < 4; ++i)
        if ((similarityMask & rejectionMasks[i]) == rejectionMasks[i])
            return false;

    return true;
}

#define cTgs 8

layout(local_size_x = 32, local_size_y = 8, local_size_z = 1) in;
void main()
{
    const ivec2 gid = ivec2(gl_GlobalInvocationID.xy);
    const ivec2 size = RenderSize();

    if (gid.x >= size.x || gid.y >= size.y) {
        return;
    }

    if(ComputeThinFeatureConfidence(gid))
    {
        StoreNewLocks(HrPosFromLrPos(gid), 1U);
    }

    // Clear depth map for next frame
    uint farPlaneBits = floatBitsToUint(1.0);
    imageStore(prevDepthTex, gid, uvec4(farPlaneBits, 0, 0, 0));
}