anki-3d-engine/programs/Reflections.ankiprog

<!-- 
Copyright (C) 2009-2018, Panagiotis Christopoulos Charitos and contributors.
All rights reserved.
Code licensed under the BSD License.
http://www.anki3d.org/LICENSE
-->
<shaderProgram>
	<mutators>
		<mutator name="VARIANT" values="0 1"/>
	</mutators>

	<shaders>
		<shader type="comp">
			<inputs>
				<input name="FB_SIZE" type="uvec2" const="1"/>
				<input name="WORKGROUP_SIZE" type="uvec2" const="1"/>
				<input name="MAX_STEPS" type="uint" const="1"/>
				<input name="LIGHT_BUFFER_MIP_COUNT" type="uint" const="1"/>
				<input name="HIZ_MIP_COUNT" type="uint" const="1"/>
				<input name="CLUSTER_COUNT_X" type="uint" const="1"/>
				<input name="CLUSTER_COUNT_Y" type="uint" const="1"/>
				<input name="CLUSTER_COUNT_Z" type="uint" const="1"/>
				<input name="IR_MIPMAP_COUNT" type="uint" const="1"/>
			</inputs>

			<source><![CDATA[
// if VARIANT==0 then the checkerboard pattern is (render on 'v'):
// -----
// |v| |
// | |v|
// -----


#include "shaders/Functions.glsl"
#include "shaders/Pack.glsl"
#include "shaders/Clusterer.glsl"

#define LIGHT_SET 0
#define LIGHT_SS_BINDING 0
#define LIGHT_UBO_BINDING 0
#define LIGHT_TEX_BINDING 6
#define LIGHT_INDIRECT
#define LIGHT_COMMON_UNIS
#include "shaders/ClusterLightCommon.glsl"

const ivec2 HIZ_SIZE = ivec2(FB_SIZE) >> 1;

layout(local_size_x = WORKGROUP_SIZE.x, local_size_y = WORKGROUP_SIZE.y, local_size_z = 1) in;

layout(ANKI_TEX_BINDING(0, 0)) uniform sampler2D u_gbufferRt0;
layout(ANKI_TEX_BINDING(0, 1)) uniform sampler2D u_gbufferRt1;
layout(ANKI_TEX_BINDING(0, 2)) uniform sampler2D u_gbufferRt2;
layout(ANKI_TEX_BINDING(0, 3)) uniform sampler2D u_depthRt;
layout(ANKI_TEX_BINDING(0, 4)) uniform sampler2D u_hizRt;
layout(ANKI_TEX_BINDING(0, 5)) uniform sampler2D u_lightBufferRt;

layout(ANKI_IMAGE_BINDING(0, 0)) writeonly uniform image2D out_reflAndIndirect;

layout(ANKI_SS_BINDING(0, 2)) buffer b_
{
	vec4 u_irradianceSh[];
};

// Temp buffer to hold the indirect color
shared vec3 s_pixels[WORKGROUP_SIZE.y][WORKGROUP_SIZE.x];

#define u_normalMat mat3(u_viewMat)

vec4 returnSslrColor(vec3 raySample, float factor, float roughness)
{
	// Re-project previous frame
	vec4 v4 = u_prevViewProjMatMulInvViewProjMat * vec4(UV_TO_NDC(raySample.xy), raySample.z, 1.0);
	raySample.xy = NDC_TO_UV(v4.xy / v4.w);
	raySample.xy = saturate(raySample.xy);

	vec2 ndc = abs(UV_TO_NDC(raySample.xy));
	float contribution = max(ndc.x, ndc.y);
	contribution = 1.0 - contribution * contribution;
	contribution *= factor;

	float lod = float(LIGHT_BUFFER_MIP_COUNT - 1u) * roughness;
	vec3 color = textureLod(u_lightBufferRt, raySample.xy, lod).rgb;
	return vec4(color, contribution);
}

// Note: All calculations in view space
vec4 doSslr(vec3 r, vec3 n, vec3 viewPos, vec2 uv, float depth, float roughness)
{
	vec3 p0 = viewPos;

	// Check for view facing reflections [sakibsaikia]
	vec3 viewDir = normalize(viewPos);
	float cameraFacingReflectionAttenuation = 1.0 - smoothstep(0.25, 0.5, dot(-viewDir, r));
	if(cameraFacingReflectionAttenuation <= 0.0)
	{
		return vec4(0.0);
	}

	// Compute an end point p1. This point is supposed to fall in front of the near plane. Add a small padding to near
	// to avoid having p1 touching the near plane.
	vec3 p1 = p0 + r * (-p0.z - (u_near + 0.1));

	// Start point
	vec3 start = vec3(uv, depth);

	// Project end point
	vec4 end4 = u_projMat * vec4(p1, 1.0);
	vec3 end = end4.xyz / end4.w;
	end.xy = NDC_TO_UV(end.xy);

	// Compute the ray and step size
	vec3 ray = end - start;
	vec2 texelDims = abs(ray.xy) * vec2(HIZ_SIZE);
	float stepSize = length(ray.xy) / max(texelDims.x, texelDims.y);
	ray = normalize(ray);

	// Compute step
	const uint BIG_STEP_SKIP = 32u;
	uint stepSkip = BIG_STEP_SKIP;

	uint l = gl_GlobalInvocationID.x & 1u;
	uint j = gl_GlobalInvocationID.y & 1u;
	const uint STEPS_ARR[4] = uint[](6u, 25u, 13u, 18u);
	uint step = STEPS_ARR[l * 2u + j];
	
	// Iterate
	bool found = false;
	vec3 raySample;
	ANKI_LOOP for(uint iterations = 0u; iterations < MAX_STEPS; ++iterations)
	{
		raySample = start + ray * (float(step) * stepSize);

		// Check if it's out of the view
		if(raySample.x <= 0.0 || raySample.y <= 0.0 || raySample.x >= 1.0 || raySample.y >= 1.0)
		{
			break;
		}

		float depth = textureLod(u_hizRt, raySample.xy, 0.0).r;

		bool hit = raySample.z - depth >= 0.0;
		if(!hit)
		{
			step += stepSkip;
		}
		else if(stepSkip > 1)
		{
			step -= BIG_STEP_SKIP - 1u;
			stepSkip = 1u;
		}
		else
		{
			found = true;
			break;
		}
	}

	//return vec4(heatmap(float(iterations) / float(MAX_STEPS)), 1.0);

	if(found)
	{
		return returnSslrColor(raySample, cameraFacingReflectionAttenuation, roughness);
	}
	else
	{
		return vec4(0.0);
	}
}

// Note: All calculations in world space
void readReflectionsAndIrradianceFromProbes(
	uint idxOffset, vec3 worldPos, vec3 normal, float roughness, out vec3 specIndirect, out vec3 diffIndirect)
{
	specIndirect = vec3(0.0);
	diffIndirect = vec3(0.0);

	vec3 viewDir = normalize(worldPos - u_cameraPos);
	vec3 reflDir = reflect(viewDir, normal);

	float reflLod = float(IR_MIPMAP_COUNT - 1u) * roughness;

	// Check proxy
	uint count = u_lightIndices[idxOffset++];
	ANKI_LOOP while(count-- != 0)
	{
		uint probeIdx = u_lightIndices[idxOffset++];
		ReflectionProbe probe = u_reflectionProbes[probeIdx];

		float R2 = probe.positionRadiusSq.w;
		vec3 center = probe.positionRadiusSq.xyz;

		// Get distance from the center of the probe
		vec3 f = worldPos - center;

		// Cubemap UV in view space
		vec3 uv = computeCubemapVecAccurate(reflDir, R2, f);

		// Read!
		float cubemapIndex = probe.cubemapIndexPad3.x;
		vec3 c = textureLod(u_reflectionsTex, vec4(uv, cubemapIndex), reflLod).rgb;

		// Combine (lerp) with previous color
		float d = dot(f, f);
		float factor = d / R2;
		factor = min(factor, 1.0);
		specIndirect = mix(c, specIndirect, factor);

		// Do the same for diffuse
		if(gl_GlobalInvocationID.x > 2560. / 2. / 2.)
		{
			uv = computeCubemapVecAccurate(normal, R2, f);
			vec3 id = textureLod(u_irradianceTex, vec4(uv, cubemapIndex), 0.0).rgb;
			diffIndirect = mix(id, diffIndirect, factor);
		}
		else
		{
			vec3 nn = normal.yzx;
			vec3 id;
			for(uint i = 0u; i < 3u; ++i)
			{
				vec4 sh = u_irradianceSh[uint(cubemapIndex) * 3u + i];
				id[i] = dot(sh, vec4(1.0, nn));
				//id[i] = sh[i];
			}
			diffIndirect = mix(id, diffIndirect, factor);
		}
	}
}

void main()
{
	// Compute a global invocation ID that takes the checkerboard pattern into account
	ivec2 fixedInvocationId = ivec2(gl_GlobalInvocationID.xy);
	fixedInvocationId.x *= 2;
#if VARIANT == 0
	fixedInvocationId.x += ((fixedInvocationId.y + 1) & 1);
#else
	fixedInvocationId.x += ((fixedInvocationId.y + 0) & 1);
#endif

	if(fixedInvocationId.x >= int(FB_SIZE.x) || fixedInvocationId.y >= int(FB_SIZE.y))
	{
		// Skip threads outside the writable image
		return;
	}

	vec2 uv = (vec2(fixedInvocationId) + 0.5) / vec2(FB_SIZE);
	vec2 ndc = UV_TO_NDC(uv);

	// Read gbuffer
	GbufferInfo gbuffer;
	readGBuffer(u_gbufferRt0, u_gbufferRt1, u_gbufferRt2, uv, 0.0, gbuffer);

	// Get depth
	float depth = textureLod(u_depthRt, uv, 0.0).r;

	// Get world position
	vec4 worldPos4 = u_invViewProjMat * vec4(ndc, depth, 1.0);
	vec3 worldPos = worldPos4.xyz / worldPos4.w;

	// Compute env BRDF
	vec3 env;
	{
		vec3 viewDir = normalize(u_cameraPos - worldPos);
		float NoV = max(EPSILON, dot(gbuffer.normal, viewDir));
		env = envBRDF(gbuffer.specular, gbuffer.roughness, u_integrationLut, NoV);
	}

	// Try SSR
	float sslrFactor = 0.0;
	vec3 sslrCol = vec3(0.0);
	if(env.g > 0.05)
	{
		// Get view pos
		vec4 viewPos4 = u_invProjMat * vec4(UV_TO_NDC(uv), depth, 1.0);
		vec3 viewPos = viewPos4.xyz / viewPos4.w;

		// Do SSLR
		vec3 viewDir = normalize(viewPos);
		vec3 viewNormal = u_normalMat * gbuffer.normal;
		vec3 reflVec = reflect(viewDir, viewNormal);

		vec4 sslr = doSslr(reflVec, viewNormal, viewPos, uv, depth, gbuffer.roughness);
		sslrFactor = sslr.w;
		sslrCol = sslr.xyz;
		sslrCol = clamp(sslrCol, 0.0, FLT_MAX); // Fix the value just in case
	}

	// Read probes
	vec3 probeCol = vec3(0.0);
	vec3 indirectCol = vec3(0.0);
	{
		// Get first light index
		uint clusterIdx = computeClusterIndex(u_clustererMagic, uv, worldPos, CLUSTER_COUNT_X, CLUSTER_COUNT_Y);
		uint idxOffset = u_clusters[clusterIdx];

		// Skip decals
		uint count = u_lightIndices[idxOffset++];
		idxOffset += count;
		
		// Skip point lights
		count = u_lightIndices[idxOffset++];
		idxOffset += count;

		// Skip spot lights
		count = u_lightIndices[idxOffset++];
		idxOffset += count;

		// Do the probe read
		readReflectionsAndIrradianceFromProbes(
			idxOffset, worldPos, gbuffer.normal, gbuffer.roughness, probeCol, indirectCol);
	}

	// Combine the SSR and probe reflections and write the result
	vec3 finalRefl = mix(probeCol, sslrCol, sslrFactor);

	// Compute the final color
	vec3 outColor = indirectCol /* gbuffer.diffuse + finalRefl * env*/;

	// Store the color for the resolve
	s_pixels[gl_LocalInvocationID.y][gl_LocalInvocationID.x] = outColor;

	// Wait for all the threads to store their stuff
	memoryBarrierShared();
	barrier();

	// Compute the missing pixel by resolving with the right or left neighbour
	ivec2 readPixel, storePixel;
	readPixel.y = int(gl_LocalInvocationID.y);
	storePixel.y = fixedInvocationId.y;

#if VARIANT == 0
	bool pickRightNeighbour = (fixedInvocationId.y & 1) == 1;
#else
	bool pickRightNeighbour = (fixedInvocationId.y & 1) == 0;
#endif
	int xOffset = (pickRightNeighbour) ? 1 : -1;
	
	readPixel.x = int(gl_LocalInvocationID.x) + xOffset;
	readPixel.x = clamp(readPixel.x, 0, int(WORKGROUP_SIZE.x - 1));

	storePixel.x = fixedInvocationId.x + xOffset;

	vec3 missingColor = (outColor + s_pixels[readPixel.y][readPixel.x]) * 0.5; // average

	// Store both the pixels
	imageStore(out_reflAndIndirect, fixedInvocationId, vec4(outColor, 0.0));
	imageStore(out_reflAndIndirect, storePixel, vec4(missingColor, 0.0));
}
			]]></source>
		</shader>
	</shaders>
</shaderProgram>