ToyEngine/Exp2/Engine/Shaders/pbr.comp

#version 450 core

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

layout(rgba8, binding = 0) uniform image2D gBaseColor;
uniform sampler2D gNormal;
uniform sampler2D gPosition;
uniform sampler2D gMetallicRoughness;
uniform sampler2DArray gShadowMap;

uniform samplerCube irradianceMap;
uniform samplerCube prefilterMap;
uniform sampler2D brdfLUT;

layout (std140, binding = 0) uniform LightSpaceMatrices
{
    mat4 lightSpaceMatrices[16];
};
uniform mat4 view;
uniform float farPlane;
uniform float shadowCascadePlaneDistances[16];
uniform float shadowBiases[16];
uniform int shadowCascadeCount;
uniform float shadowBlendRatio;

uniform vec3 mainLightDirection;
uniform vec3 mainLightRadiance;

uniform vec3 camPos;

const float PI = 3.14159265359;

float DistributionGGX(vec3 N, vec3 H, float roughness)
{
	float a = roughness*roughness;
	float a2 = a*a;
	float NdotH = max(dot(N, H), 0.0);
	float NdotH2 = NdotH*NdotH;

	float nom   = a2;
	float denom = (NdotH2 * (a2 - 1.0) + 1.0);
	denom = PI * denom * denom;

	return nom / denom;
}
// ----------------------------------------------------------------------------
float GeometrySchlickGGX(float NdotV, float roughness)
{
	float r = (roughness + 1.0);
	float k = (r*r) / 8.0;

	float nom   = NdotV;
	float denom = NdotV * (1.0 - k) + k;

	return nom / denom;
}
// ----------------------------------------------------------------------------
float GeometrySmith(vec3 N, vec3 V, vec3 L, float roughness)
{
	float NdotV = max(dot(N, V), 0.0);
	float NdotL = max(dot(N, L), 0.0);
	float ggx2 = GeometrySchlickGGX(NdotV, roughness);
	float ggx1 = GeometrySchlickGGX(NdotL, roughness);

	return ggx1 * ggx2;
}
// ----------------------------------------------------------------------------
vec3 fresnelSchlick(float cosTheta, vec3 F0)
{
	return F0 + (1.0 - F0) * pow(1.0 - cosTheta, 5.0);
}
vec3 fresnelSchlickRoughness(float cosTheta, vec3 F0, float roughness)
{
    return F0 + (max(vec3(1.0 - roughness), F0) - F0) * pow(1.0 - cosTheta, 5.0);
}   

vec3 ambientLighting(vec3 N, vec3 V, vec3 F0, vec3 albedo, float metallic, float roughness, float ao)
{
	// ambient lighting (we now use IBL as the ambient term)
    vec3 F = fresnelSchlickRoughness(clamp(dot(N, V),0.,1.), F0, roughness);
    
    vec3 kS = F;
    vec3 kD = 1.0 - kS;
    kD *= 1.0 - metallic;	  
    
    vec3 irradiance = texture(irradianceMap, N).rgb;
    vec3 diffuse    = irradiance * albedo;
    
    // sample both the pre-filter map and the BRDF lut and combine them together as per the Split-Sum approximation to get the IBL specular part.
    const float MAX_REFLECTION_LOD = 4.0;
	vec3 R = reflect(-V, N); 
    vec3 prefilteredColor = textureLod(prefilterMap, R,  roughness * MAX_REFLECTION_LOD).rgb;    
    vec2 brdf  = texture(brdfLUT, vec2(clamp(dot(N, V),0.,1.), roughness)).rg;
    vec3 specular = prefilteredColor * (F * brdf.x + brdf.y);

	return (kD * diffuse + specular) * ao;
}


const int pcfRadius = 3;
float getShadowFromLayer(vec3 fragPosWorldSpace, vec3 normal,int layer)
{
    float normalBias = 1.2*4. /** (1+pcfRadius)*/ * shadowBiases[layer]*max((1.0 - dot(normal, mainLightDirection)), 0.1)/textureSize(gShadowMap, 0).x;
	normalBias = max(normalBias, 0.05);
    vec4 fragPosLightSpace = lightSpaceMatrices[layer] * vec4(fragPosWorldSpace+normal*normalBias, 1.0);
    // perform perspective divide
    vec3 projCoords = fragPosLightSpace.xyz / fragPosLightSpace.w;
    // transform to [0,1] range
    projCoords = projCoords * 0.5 + 0.5;

    // get depth of current fragment from light's perspective
    float currentDepth = projCoords.z;

    // keep the shadow at 0.0 when outside the far_plane region of the light's frustum.
    if (currentDepth > 1.0)
    {
        return 0.0;
    }

    // PCF
    float shadow = 0.0;
    vec2 texelSize = 1.0 / vec2(textureSize(gShadowMap, 0));
    for(int x = -pcfRadius; x <= pcfRadius; ++x)
    {
        for(int y = -pcfRadius; y <= pcfRadius; ++y)
        {
            float pcfDepth = texture(gShadowMap, vec3(projCoords.xy + vec2(x, y) * texelSize, layer)).r; 
			shadow += currentDepth > pcfDepth ? 1.0 : 0.0;    
        }    
    }
    shadow /= (2*pcfRadius+1)*(2*pcfRadius+1);
	return shadow;
}

float ShadowCalculation(vec3 fragPosWorldSpace, vec3 normal, out int layer)
{
    // select cascade layer
    vec4 fragPosViewSpace = view * vec4(fragPosWorldSpace, 1.0);
    float depthValue = abs(fragPosViewSpace.z);

    layer = -1;
    for (int i = 0; i < shadowCascadeCount; ++i)
    {
        if (depthValue < shadowCascadePlaneDistances[i])
        {
            layer = i;
            break;
        }
    }
//    if (layer == -1)
//    {
//        layer = shadowCascadeCount;
//    } 
	 
    float shadow = getShadowFromLayer(fragPosWorldSpace,normal,layer);
        

	float nextLayerBeginDepth = layer==0? (1-shadowBlendRatio)*shadowCascadePlaneDistances[layer]
										:(1-shadowBlendRatio)*shadowCascadePlaneDistances[layer]+shadowBlendRatio*shadowCascadePlaneDistances[layer-1];
	if(depthValue > nextLayerBeginDepth)
	{
		float shadowNext = getShadowFromLayer(fragPosWorldSpace,normal,layer+1);
		shadow = mix(shadow,shadowNext, (depthValue-nextLayerBeginDepth)/(shadowCascadePlaneDistances[layer]-nextLayerBeginDepth));
	}
    return shadow;
}



vec4 encodeRGBM(vec3 color)
{
	if(dot(color,color)==0)
		return vec4(0,0,0,1);
	vec4 rgbm;
	float range = 8;
	color /= range;
    rgbm.a = clamp(max(max(color.r, color.g), max(color.b, 1e-6)), 0.0, 1.0);
    rgbm.a = ceil(rgbm.a * 255.0) / 255.0;
    rgbm.rgb = color / rgbm.a;
	return rgbm;
}

void main()
{
    ivec2 pixelLocation = ivec2(gl_GlobalInvocationID.xy);

	
    //vec3  albedo    = pow(texelFetch(gBaseColor, pixelLocation, 0).rgb, vec3(2.2));
	vec3  albedo    = pow(imageLoad(gBaseColor, pixelLocation).rgb, vec3(2.2));
	float metallic  = texelFetch(gMetallicRoughness, pixelLocation, 0).r;
	float roughness = texelFetch(gMetallicRoughness, pixelLocation, 0).g;

    vec3 worldPos = texelFetch(gPosition, pixelLocation,0).xyz;
    vec3 normal = texelFetch(gNormal, pixelLocation,0).xyz;

	if(normal==vec3(0))
	{
		//vec3 color = mainLightRadiance;
		//imageStore(gBaseColor, pixelLocation, vec4(color, 1.0));
		imageStore(gBaseColor, pixelLocation, vec4(0));
		return;
    }
	normal = normalize(normal);

	vec3 V = normalize(camPos - worldPos);

	vec3 F0 = vec3(0.04); 
	F0 = mix(F0, albedo, metallic);

	// calculate light radiance
	vec3 L = normalize(mainLightDirection);
	vec3 H = normalize(V + L);
	//float distance    = length(lightPositions[i] - WorldPos);
	//float attenuation = 1.0 / (distance * distance);
	vec3  radiance    = mainLightRadiance ;//* attenuation;        

	// cook-torrance brdf
	float NDF = DistributionGGX(normal, H, roughness);        
	float G   = GeometrySmith(normal, V, L, roughness);      
	vec3  F   = fresnelSchlick(clamp(dot(H, V),0.,1.), F0);
	//F = clamp(F,vec3(0),vec3(1));

	vec3 kS = F;
	vec3 kD = vec3(1.0) - kS;
	kD *= 1.0 - metallic;     

	vec3  nominator   = NDF * G * F;
	float denominator = 4.0 * max(dot(normal, V), 0.0) * max(dot(normal, L), 0.0) + 0.001; 
	vec3  specular    = nominator / denominator;

	float NdotL = max(dot(normal, L), 0.0);                
	vec3 Lo = (kD * albedo / PI + specular) * radiance * NdotL; 

	vec3 ambient = ambientLighting(normal, V, F0, albedo, metallic, roughness, 1); 

    int debugLayer;
	float shadow = ShadowCalculation(worldPos, normal, debugLayer);

	vec3 color = (1-shadow)*Lo + ambient;
    imageStore(gBaseColor, pixelLocation, encodeRGBM(color));
}