prefilteredmap.comp

#version 450
#extension GL_GOOGLE_include_directive : require
#if AMD_SHADER_IMAGE_LOAD_STORE_LOD == 1
#extension GL_AMD_shader_image_load_store_lod : enable
#endif

#include "cubemap.glsl"
#include "pbr.glsl"

layout (constant_id = 0) const uint ROUGHNESS_LEVELS = 8;
layout (constant_id = 1) const uint SAMPLES = 1024;

layout (set = 0, binding = 0) uniform samplerCube cubemapSampler;
#if AMD_SHADER_IMAGE_LOAD_STORE_LOD == 1
layout (set = 0, binding = 1) writeonly uniform imageCube prefilteredmapImage;
#else
layout (set = 0, binding = 1) writeonly uniform imageCube prefilteredmapMipImages[ROUGHNESS_LEVELS];
#endif

layout (push_constant) uniform PushConstant {
    int mipLevel;
} pc;

layout (local_size_x = 16, local_size_y = 16) in;

void main(){
    const float PI = 3.14159265359;

    uint prefilteredmapImageSize
#if AMD_SHADER_IMAGE_LOAD_STORE_LOD == 1
        = imageSize(prefilteredmapImage).x >> pc.mipLevel;
#else
        = imageSize(prefilteredmapMipImages[0]).x >> pc.mipLevel;
#endif
    if (gl_GlobalInvocationID.x >= prefilteredmapImageSize || gl_GlobalInvocationID.y >= prefilteredmapImageSize){
        return;
    }

    float roughness = float(pc.mipLevel) / float(ROUGHNESS_LEVELS - 1U);
    uint cubemapSize = textureSize(cubemapSampler, 0).x;
    float saTexel  = 4.0 * PI / (6 * cubemapSize * cubemapSize);

    // tagent space from origin point
    vec3 N = getWorldDirection(ivec3(gl_GlobalInvocationID), prefilteredmapImageSize);
    // assume view direction always equal to outgoing direction
    vec3 V = N;

    float totalWeight = 0.0;
    vec3 prefilteredColor = vec3(0.0);
    for (uint i = 0; i < SAMPLES; ++i){
        // generate sample vector towards the alignment of the specular lobe
        vec2 Xi = hammersleySequence(i, SAMPLES);
        vec3 H = importanceSampleGGX(Xi, N, roughness);
        vec3 L = reflect(-V, H);

        float dotNL = dot(N, L);
        if (dotNL > 0.0){
            float dotNH = max(dot(N, H), 0.0);
            // sample from the environment's mip level based on roughness/pdf
            float D = dGgx(dotNH, roughness);
            float pdf = 0.25 * D + 0.0001;

            float saSample = 1.0 / (SAMPLES * pdf + 0.0001);
            float mipLevel = roughness == 0.0 ? 0.0 : 0.5 * log2(saSample / saTexel);

            // TODO: don't know why, but use original L flips the prefilteredmap in y-axis.
            prefilteredColor += textureLod(cubemapSampler, vec3(L.x, -L.y, L.z), mipLevel).rgb * dotNL;
            totalWeight += dotNL;
        }
    }
    prefilteredColor /= totalWeight;

#if AMD_SHADER_IMAGE_LOAD_STORE_LOD == 1
    imageStoreLodAMD(prefilteredmapImage, ivec3(gl_GlobalInvocationID), pc.mipLevel, vec4(prefilteredColor, 1.0));
#else
    imageStore(prefilteredmapMipImages[pc.mipLevel], ivec3(gl_GlobalInvocationID), vec4(prefilteredColor, 1.0));
#endif
}