prefilter

prefilter_cubemap.glsl

#version 450

#extension GL_EXT_nonuniform_qualifier : enable

layout (local_size_x = 32, local_size_y = 32, local_size_z = 1) in;

layout (binding = 0) uniform samplerCube cube_map;
layout (binding = 1, rgba32f) uniform imageCube specular_map;

vec3 view_pos = vec3(0, 0, 0);

vec3 view_target[6] = {
    vec3(1, 0, 0),
    vec3(-1, 0, 0),
    vec3(0, -1, 0),
    vec3(0, 1, 0),
    vec3(0, 0, 1),
    vec3(0, 0, -1)
};

vec3 view_up[6] = {
    vec3(0, 1, 0),
    vec3(0, 1, 0),
    vec3(0, 0, 1), // +Y
    vec3(0, 0, -1), // -Y
    vec3(0, 1, 0),
    vec3(0, 1, 0)
};

mat4 look_at(vec3 pos, vec3 target, vec3 world_up) {
    vec3 fwd = normalize(target - pos);
    vec3 right = cross(world_up, fwd);
    vec3 up = cross(fwd, right);
    return mat4(vec4(right, 0), vec4(up, 0), vec4(fwd, 0), vec4(pos, 1));
}

// Remaps uv coordinates to a position in range [-1, 1]
vec2 uv_to_quad_pos(vec2 uv) {
    return uv * 2.0f - 1.0f;
}

vec3 to_local_direction(vec2 uv, int face) {
     mat4 view = look_at(view_pos, view_target[face], view_up[face]);
     return (view * vec4(uv_to_quad_pos(uv), 1, 1)).xyz;
}

float radicalinverse_vdc(uint bits) {
    bits = (bits << 16u) | (bits >> 16u);
    bits = ((bits & 0x55555555u) << 1u) | ((bits & 0xAAAAAAAAu) >> 1u);
    bits = ((bits & 0x33333333u) << 2u) | ((bits & 0xCCCCCCCCu) >> 2u);
    bits = ((bits & 0x0F0F0F0Fu) << 4u) | ((bits & 0xF0F0F0F0u) >> 4u);
    bits = ((bits & 0x00FF00FFu) << 8u) | ((bits & 0xFF00FF00u) >> 8u);
    return float(bits) * 2.3283064365386963e-10; // / 0x100000000
}

vec2 hammersley(uint i, uint N) {
    return vec2(float(i)/float(N), radicalinverse_vdc(i));
}  

#define PI 3.14159265359

vec3 importance_sample_ggx(vec2 Xi, vec3 N, float roughness) {
    float a = roughness * roughness;
	
    float phi = 2.0 * PI * Xi.x;
    float cos_theta = sqrt((1.0 - Xi.y) / (1.0 + (a * a - 1.0) * Xi.y));
    float sin_theta = sqrt(1.0 - cos_theta * cos_theta);
	
    // from spherical coordinates to cartesian coordinates
    vec3 H;
    H.x = cos(phi) * sin_theta;
    H.y = sin(phi) * sin_theta;
    H.z = cos_theta;
	
    // from tangent-space vector to world-space sample vector
    vec3 up = abs(N.z) < 0.999 ? vec3(0.0, 0.0, 1.0) : vec3(1.0, 0.0, 0.0);
    vec3 tangent = normalize(cross(up, N));
    vec3 bitangent = cross(N, tangent);
	
    vec3 sample_vec = tangent * H.x + bitangent * H.y + N * H.z;
    return normalize(sample_vec);
}  

// Credit goes to https://github.com/nem0/LumixEngine/blob/master/data/pipelines/ibl_filter.shd#L120
float D_GGX(float ndoth, float roughness) {
	float a = roughness * roughness;
	float a2 = a * a;
	float f = max(1e-5, (ndoth * ndoth) * (a2 - 1) + 1);
	return a2 / (f * f * PI);
}

// Credit goes to https://github.com/nem0/LumixEngine/blob/master/data/pipelines/ibl_filter.shd#L120
float prefiltered_importance_sampling(float ipdf) {
	const float numSamples = 1024.0f;
	const float invNumSamples = 1.0 / 1024.0f;
	const float dim = textureSize(cube_map, 0).x;
	const float omegaP = (4.0 * PI) / (6.0 * dim * dim);
	const float invOmegaP = 1.0 / omegaP;
	const float K = 4.0;
	const float iblRoughnessOneLevel = 7;
	float omegaS = invNumSamples * ipdf;
	float mipLevel = clamp(log2(K * omegaS * invOmegaP) * 0.5, 0.0, iblRoughnessOneLevel);
	return mipLevel;
}

vec3 apply_specular_filter(vec3 local_direction, float roughness) {
    vec3 N = normalize(local_direction);    
    vec3 R = N;
    vec3 V = R;

    const uint SAMPLE_COUNT = 1024u;
    float total_weight = 0.0;   
    vec3 prefiltered_color = vec3(0.0);     
    for(uint i = 0u; i < SAMPLE_COUNT; ++i) {
        vec2 Xi = hammersley(i, SAMPLE_COUNT);
        vec3 H = importance_sample_ggx(Xi, N, roughness);
        vec3 L = normalize(2.0 * dot(V, H) * H - V);

        float n_dot_l = max(dot(N, L), 0.0);
        if(n_dot_l > 0.0) {
            // Credit for this approach goes to https://github.com/nem0/LumixEngine/blob/master/data/pipelines/ibl_filter.shd#L120
            float ipdf = (4.0 * dot(L, H)) / (D_GGX(dot(N, H), roughness) * dot(N, H));
            float mip_level = prefiltered_importance_sampling(ipdf);
            prefiltered_color += textureLod(cube_map, L, mip_level).rgb * n_dot_l;
            total_weight += n_dot_l;
        }
    }
    prefiltered_color = prefiltered_color / total_weight;
    return prefiltered_color;
}

layout(push_constant) uniform PC {
    float roughness;
} pc;

void main() {
    for (int face = 0; face < 6; ++face) {
        int face_size = imageSize(specular_map).x;
        vec2 uv = clamp(vec2(gl_GlobalInvocationID.xy) / float(face_size), vec2(0), vec2(1));
        vec3 local_direction = normalize(to_local_direction(uv, face));
        local_direction.y = -local_direction.y;

        vec4 color = vec4(apply_specular_filter(local_direction, pc.roughness), 1);

        ivec3 cube_texels = ivec3(uv * face_size, face);
        imageStore(specular_map, cube_texels, color);
    }
}