Optimize SSR

shaders/Ssr.ankiprog

@@ -11,132 +11,29 @@
 
 #pragma anki mutator VARIANT 0 1
 
-ANKI_SPECIALIZATION_CONSTANT_UVEC2(FB_SIZE, 0, UVec2(1));
-ANKI_SPECIALIZATION_CONSTANT_U32(MAX_STEPS, 2, 1);
-ANKI_SPECIALIZATION_CONSTANT_U32(LIGHT_BUFFER_MIP_COUNT, 3, 1);
-ANKI_SPECIALIZATION_CONSTANT_F32(HISTORY_COLOR_BLEND_FACTOR, 4, 0.5);
-
 #pragma anki start comp
 #include <shaders/Functions.glsl>
 #include <shaders/Pack.glsl>
 #include <shaders/glsl_cpp_common/Ssr.h>
 #include <shaders/Tonemapping.glsl>
+#include <shaders/SsRaymarching.glsl>
 
 const UVec2 WORKGROUP_SIZE = UVec2(16, 16);
 layout(local_size_x = WORKGROUP_SIZE.x, local_size_y = WORKGROUP_SIZE.y, local_size_z = 1) in;
 
 layout(set = 0, binding = 0) uniform sampler u_trilinearClampSampler;
-layout(set = 0, binding = 1) uniform texture2D u_gbufferRt1;
-layout(set = 0, binding = 2) uniform texture2D u_gbufferRt2;
-layout(set = 0, binding = 3) uniform texture2D u_depthRt;
-layout(set = 0, binding = 4) uniform texture2D u_lightBufferRt;
+layout(set = 0, binding = 1) uniform sampler u_nearestNearestClampSampler;
+layout(set = 0, binding = 2) uniform texture2D u_gbufferRt1;
+layout(set = 0, binding = 3) uniform texture2D u_gbufferRt2;
+layout(set = 0, binding = 4) uniform texture2D u_depthRt;
+layout(set = 0, binding = 5) uniform texture2D u_lightBufferRt;
 
-layout(set = 0, binding = 5, rgba16f) uniform image2D out_img;
+layout(set = 0, binding = 6, rgba16f) uniform image2D out_img;
 
-layout(set = 0, binding = 6, row_major) uniform u_
+layout(set = 0, binding = 7, row_major) uniform u_
 {
 	SsrUniforms u_unis;
 };
-#define u_prevViewProjMatMulInvViewProjMat u_unis.m_prevViewProjMatMulInvViewProjMat
-#define u_near u_unis.m_nearPad3.x
-#define u_projMat u_unis.m_projMat
-#define u_invProjMat u_unis.m_invProjMat
-#define u_normalMat u_unis.m_normalMat
-
-// Temp buffer to hold the indirect color
-shared Vec4 s_pixels[WORKGROUP_SIZE.y][WORKGROUP_SIZE.x];
-
-// Note: All calculations in view space
-// It returns the UV coordinates of the reflection (xy) and the contrubution factor (z)
-Vec3 raymarch(Vec3 r, Vec3 n, Vec3 viewPos, Vec2 uv, F32 depth)
-{
-	const Vec3 p0 = viewPos;
-
-	// Check for view facing reflections [sakibsaikia]
-	const Vec3 viewDir = normalize(viewPos);
-	const F32 cameraContribution = 1.0 - smoothstep(0.25, 0.5, dot(-viewDir, r));
-	if(cameraContribution <= 0.0)
-	{
-		return Vec3(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.
-	const F32 p0DistFromNear = -p0.z - (u_near + 0.1); // aka testPlane(near, p0)
-	const Vec3 p1 = p0 + r * p0DistFromNear; // Whatever the direction of r, p1 won't touch the near plane
-
-	// Start point
-	const Vec3 start = Vec3(uv, depth);
-
-	// Project end point
-	const 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;
-	const Vec2 texelDims = abs(ray.xy) * Vec2(FB_SIZE);
-	const F32 stepSize = length(ray.xy) / max(texelDims.x, texelDims.y);
-	ray = normalize(ray);
-
-	// Compute step
-	const U32 BIG_STEP_SKIP = 16u;
-	U32 stepSkip = BIG_STEP_SKIP;
-
-	const U32 STEP_FRACTION = BIG_STEP_SKIP / (4u + 1u);
-	const U32 STEPS_ARR[4] = U32[](STEP_FRACTION, 4 * STEP_FRACTION, 2 * STEP_FRACTION, 3 * STEP_FRACTION);
-	const U32 l = gl_GlobalInvocationID.x & 1u;
-	const U32 j = gl_GlobalInvocationID.y & 1u;
-	U32 step = STEPS_ARR[l * 2u + j];
-
-	// Iterate
-	F32 finalContribution = 0.0;
-	Vec3 raySample;
-	ANKI_LOOP for(U32 iterations = 0u; iterations < MAX_STEPS; ++iterations)
-	{
-		raySample = start + ray * (F32(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;
-		}
-
-		const F32 depth = textureLod(u_depthRt, u_trilinearClampSampler, raySample.xy, 0.0).r;
-
-		const Bool hit = raySample.z - depth >= 0.0;
-		if(!hit)
-		{
-			step += stepSkip;
-		}
-		else if(stepSkip > 1)
-		{
-			step -= BIG_STEP_SKIP - 1u;
-			stepSkip = 1u;
-		}
-		else
-		{
-			// Found it
-
-			// Re-project previous frame
-			const Vec4 v4 = u_prevViewProjMatMulInvViewProjMat * Vec4(UV_TO_NDC(raySample.xy), raySample.z, 1.0);
-			Vec2 ndc = v4.xy / v4.w;
-			raySample.xy = NDC_TO_UV(ndc);
-
-			// Compute the edge contribution
-			ndc = abs(ndc);
-			F32 screedEdgeContributionFactor = max(ndc.x, ndc.y);
-			screedEdgeContributionFactor = 1.0 - screedEdgeContributionFactor * screedEdgeContributionFactor;
-			screedEdgeContributionFactor = saturate(screedEdgeContributionFactor);
-			finalContribution = cameraContribution * screedEdgeContributionFactor;
-
-			break;
-		}
-	}
-
-	// Return the traced UV and the contribution factor
-	return Vec3(raySample.xy, finalContribution);
-}
 
 void main()
 {
@@ -149,13 +46,13 @@ void main()
 	fixedInvocationId.x += ((fixedInvocationId.y + 0) & 1);
 #endif
 
-	if(fixedInvocationId.x >= I32(FB_SIZE.x) || fixedInvocationId.y >= I32(FB_SIZE.y))
+	if(fixedInvocationId.x >= I32(u_unis.m_framebufferSize.x) || fixedInvocationId.y >= I32(u_unis.m_framebufferSize.y))
 	{
 		// Skip threads outside the writable image
 		return;
 	}
 
-	const Vec2 uv = (Vec2(fixedInvocationId.xy) + 0.5) / Vec2(FB_SIZE);
+	const Vec2 uv = (Vec2(fixedInvocationId.xy) + 0.5) / Vec2(u_unis.m_framebufferSize);
 
 	// Read part of the G-buffer
 	const F32 roughness = readRoughnessFromGBuffer(u_gbufferRt1, u_trilinearClampSampler, uv);
@@ -165,75 +62,62 @@ void main()
 	const F32 depth = textureLod(u_depthRt, u_trilinearClampSampler, uv, 0.0).r;
 
 	// Get view pos
-	const Vec4 viewPos4 = u_invProjMat * Vec4(UV_TO_NDC(uv), depth, 1.0);
+	const Vec4 viewPos4 = u_unis.m_invProjMat * Vec4(UV_TO_NDC(uv), depth, 1.0);
 	const Vec3 viewPos = viewPos4.xyz / viewPos4.w;
 
 	// Compute refl vector
 	const Vec3 viewDir = normalize(viewPos);
-	const Vec3 viewNormal = u_normalMat * worldNormal;
+	const Vec3 viewNormal = u_unis.m_normalMat * worldNormal;
 	const Vec3 reflVec = reflect(viewDir, viewNormal);
 
-	// Raymatch
-	const Vec3 ssr = raymarch(reflVec, viewNormal, viewPos, uv, depth);
-	const Vec2 reflUv = ssr.xy;
-	const F32 factor = ssr.z;
+	// Rand idx
+	U32 randIdx = (fixedInvocationId.x & 1) * 2 + (fixedInvocationId.y & 1);
+	randIdx = (randIdx + u_unis.m_frameCount & 3) & 3;
+
+	// Do the heavy work
+	Vec3 hitPoint;
+	F32 hitAttenuation;
+	const U32 lod = 1;
+	raymarchGroundTruth(viewPos,
+		reflVec,
+		uv,
+		depth,
+		u_unis.m_projMat,
+		randIdx,
+		u_unis.m_maxSteps,
+		u_depthRt,
+		u_nearestNearestClampSampler,
+		F32(lod),
+		u_unis.m_depthBufferSize >> lod,
+		16u,
+		hitPoint,
+		hitAttenuation);
 
 	// Read the reflection
 	Vec4 outColor;
-	ANKI_BRANCH if(factor > 0.0)
+	ANKI_BRANCH if(hitAttenuation > 0.0)
 	{
 		// Read the refl
-		const F32 lod = F32(LIGHT_BUFFER_MIP_COUNT - 1u) * roughness;
-		outColor.rgb = textureLod(u_lightBufferRt, u_trilinearClampSampler, reflUv, lod).rgb;
+
+		// Reproject the UV because you are reading the previous frame
+		const Vec4 v4 = u_unis.m_prevViewProjMatMulInvViewProjMat * Vec4(UV_TO_NDC(hitPoint.xy), hitPoint.z, 1.0);
+		hitPoint.xy = NDC_TO_UV(v4.xy / v4.w);
+
+		// Compute the LOD based on the roughness
+		const F32 lod = F32(u_unis.m_lightBufferMipCount - 1u) * roughness;
+
+		// Read the light buffer
+		outColor.rgb = textureLod(u_lightBufferRt, u_trilinearClampSampler, hitPoint.xy, lod).rgb;
 		outColor.rgb = clamp(outColor.rgb, 0.0, FLT_MAX); // Fix the value just in case
-		outColor.rgb *= factor;
-		outColor.a = 1.0 - factor;
+		outColor.rgb *= hitAttenuation;
+		outColor.a = 1.0 - hitAttenuation;
 	}
 	else
 	{
 		outColor = Vec4(0.0, 0.0, 0.0, 1.0);
 	}
 
-	// Use the prev color
-	if(HISTORY_COLOR_BLEND_FACTOR > 0.0)
-	{
-		const Vec4 v4 = u_prevViewProjMatMulInvViewProjMat * Vec4(UV_TO_NDC(uv), depth, 1.0);
-		const Vec2 oldNdc = v4.xy / v4.w;
-		const Vec2 oldUv = saturate(NDC_TO_UV(oldNdc));
-
-		const Vec4 prevColor = imageLoad(out_img, ivec2(oldUv * vec2(FB_SIZE)));
-
-		outColor = mix(prevColor, outColor, HISTORY_COLOR_BLEND_FACTOR);
-	}
-
-	// 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 = I32(gl_LocalInvocationID.y);
-	storePixel.y = fixedInvocationId.y;
-
-#if VARIANT == 0
-	const Bool pickRightNeighbour = (fixedInvocationId.y & 1) == 1;
-#else
-	const Bool pickRightNeighbour = (fixedInvocationId.y & 1) == 0;
-#endif
-	const I32 xOffset = (pickRightNeighbour) ? 1 : -1;
-
-	readPixel.x = I32(gl_LocalInvocationID.x) + xOffset;
-	readPixel.x = clamp(readPixel.x, 0, I32(WORKGROUP_SIZE.x - 1));
-
-	storePixel.x = fixedInvocationId.x + xOffset;
-
-	const Vec4 missingColor = (outColor + s_pixels[readPixel.y][readPixel.x]) * 0.5; // average
-
-	// Store both the pixels
+	// Store
 	imageStore(out_img, fixedInvocationId, outColor);
-	imageStore(out_img, storePixel, missingColor);
 }
 #pragma anki end

shaders/glsl_cpp_common/Ssr.h

@@ -12,7 +12,12 @@ ANKI_BEGIN_NAMESPACE
 // Screen space reflections uniforms
 struct SsrUniforms
 {
-	Vec4 m_nearPad3;
+	UVec2 m_depthBufferSize;
+	UVec2 m_framebufferSize;
+	U32 m_frameCount;
+	U32 m_depthMipCount;
+	U32 m_maxSteps;
+	U32 m_lightBufferMipCount;
 	Mat4 m_prevViewProjMatMulInvViewProjMat;
 	Mat4 m_projMat;
 	Mat4 m_invProjMat;

src/anki/gr/RenderGraph.cpp

@@ -1405,6 +1405,8 @@ StringAuto RenderGraph::bufferUsageToStr(StackAllocator<U8>& alloc, BufferUsageB
 Error RenderGraph::dumpDependencyDotFile(
 	const RenderGraphDescription& descr, const BakeContext& ctx, CString path) const
 {
+	ANKI_GR_LOGW("Running with debug code");
+
 	static const Array<const char*, 5> COLORS = {{"red", "green", "blue", "magenta", "cyan"}};
 	auto alloc = ctx.m_alloc;
 	StringListAuto slist(alloc);
@@ -1412,7 +1414,7 @@ Error RenderGraph::dumpDependencyDotFile(
 	slist.pushBackSprintf("digraph {\n");
 	slist.pushBackSprintf("\t//splines = ortho;\nconcentrate = true;\n");
 
-	for(U batchIdx = 0; batchIdx < ctx.m_batches.getSize(); ++batchIdx)
+	for(U32 batchIdx = 0; batchIdx < ctx.m_batches.getSize(); ++batchIdx)
 	{
 		// Set same rank
 		slist.pushBackSprintf("\t{rank=\"same\";");
@@ -1459,14 +1461,14 @@ Error RenderGraph::dumpDependencyDotFile(
 	// slist.pushBackSprintf("subgraph cluster_1 {\n");
 	StringAuto prevBubble(ctx.m_alloc);
 	prevBubble.create("START");
-	for(U batchIdx = 0; batchIdx < ctx.m_batches.getSize(); ++batchIdx)
+	for(U32 batchIdx = 0; batchIdx < ctx.m_batches.getSize(); ++batchIdx)
 	{
 		const Batch& batch = ctx.m_batches[batchIdx];
 
 		StringAuto batchName(ctx.m_alloc);
 		batchName.sprintf("batch%u", batchIdx);
 
-		for(U barrierIdx = 0; barrierIdx < batch.m_barriersBefore.getSize(); ++barrierIdx)
+		for(U32 barrierIdx = 0; barrierIdx < batch.m_barriersBefore.getSize(); ++barrierIdx)
 		{
 			const Barrier& barrier = batch.m_barriersBefore[barrierIdx];
 			StringAuto barrierName(ctx.m_alloc);
@@ -1503,7 +1505,7 @@ Error RenderGraph::dumpDependencyDotFile(
 			prevBubble = barrierName;
 		}
 
-		for(U passIdx : batch.m_passIndices)
+		for(U32 passIdx : batch.m_passIndices)
 		{
 			const RenderPassDescriptionBase& pass = *descr.m_passes[passIdx];
 			StringAuto passName(alloc);

src/anki/renderer/Common.h

@@ -74,9 +74,6 @@ constexpr U32 BLOOM_FRACTION = 4;
 /// Volumetric size is rendererSize/VOLUMETRIC_FRACTION.
 constexpr U32 VOLUMETRIC_FRACTION = 4;
 
-/// SSR size is rendererSize/SSR_FRACTION.
-constexpr U32 SSR_FRACTION = 2;
-
 /// Used to calculate the mipmap count of the HiZ map.
 constexpr U32 HIERARCHICAL_Z_MIN_HEIGHT = 80;
 

src/anki/renderer/Ssr.cpp

@@ -21,19 +21,20 @@ Ssr::~Ssr()
 
 Error Ssr::init(const ConfigSet& cfg)
 {
-	Error err = initInternal(cfg);
+	const Error err = initInternal(cfg);
 	if(err)
 	{
-		ANKI_R_LOGE("Failed to initialize reflection pass");
+		ANKI_R_LOGE("Failed to initialize SSR pass");
 	}
 	return err;
 }
 
 Error Ssr::initInternal(const ConfigSet& cfg)
 {
-	U32 width = m_r->getWidth() / SSR_FRACTION;
-	U32 height = m_r->getHeight() / SSR_FRACTION;
+	const U32 width = m_r->getWidth();
+	const U32 height = m_r->getHeight();
 	ANKI_R_LOGI("Initializing SSR pass (%ux%u)", width, height);
+	m_maxSteps = cfg.getNumberU32("r_ssrMaxSteps");
 
 	// Create RTs
 	TextureInitInfo texinit = m_r->create2DRenderTargetInitInfo(width,
@@ -48,10 +49,6 @@ Error Ssr::initInternal(const ConfigSet& cfg)
 	ANKI_CHECK(getResourceManager().loadResource("shaders/Ssr.ankiprog", m_prog));
 
 	ShaderProgramResourceVariantInitInfo variantInitInfo(m_prog);
-	variantInitInfo.addConstant("FB_SIZE", UVec2(width, height));
-	variantInitInfo.addConstant("MAX_STEPS", cfg.getNumberU32("r_ssrMaxSteps"));
-	variantInitInfo.addConstant("LIGHT_BUFFER_MIP_COUNT", m_r->getDownscaleBlur().getMipmapCount());
-	variantInitInfo.addConstant("HISTORY_COLOR_BLEND_FACTOR", cfg.getNumberF32("r_ssrHistoryBlendFactor"));
 	variantInitInfo.addMutation("VARIANT", 0);
 
 	const ShaderProgramResourceVariant* variant;
@@ -84,7 +81,8 @@ void Ssr::populateRenderGraph(RenderingContext& ctx)
 	rpass.newDependency({m_r->getGBuffer().getColorRt(1), TextureUsageBit::SAMPLED_COMPUTE});
 	rpass.newDependency({m_r->getGBuffer().getColorRt(2), TextureUsageBit::SAMPLED_COMPUTE});
 
-	TextureSubresourceInfo hizSubresource; // Only first mip
+	TextureSubresourceInfo hizSubresource;
+	hizSubresource.m_mipmapCount = m_r->getDepthDownscale().getMipmapCount();
 	rpass.newDependency({m_r->getDepthDownscale().getHiZRt(), TextureUsageBit::SAMPLED_COMPUTE, hizSubresource});
 
 	rpass.newDependency({m_r->getDownscaleBlur().getRt(), TextureUsageBit::SAMPLED_COMPUTE});
@@ -98,20 +96,27 @@ void Ssr::run(RenderPassWorkContext& rgraphCtx)
 
 	// Bind all
 	cmdb->bindSampler(0, 0, m_r->getSamplers().m_trilinearClamp);
+	cmdb->bindSampler(0, 1, m_r->getSamplers().m_nearestNearestClamp);
 
-	rgraphCtx.bindColorTexture(0, 1, m_r->getGBuffer().getColorRt(1));
-	rgraphCtx.bindColorTexture(0, 2, m_r->getGBuffer().getColorRt(2));
+	rgraphCtx.bindColorTexture(0, 2, m_r->getGBuffer().getColorRt(1));
+	rgraphCtx.bindColorTexture(0, 3, m_r->getGBuffer().getColorRt(2));
 
-	TextureSubresourceInfo hizSubresource; // Only first mip
-	rgraphCtx.bindTexture(0, 3, m_r->getDepthDownscale().getHiZRt(), hizSubresource);
+	TextureSubresourceInfo hizSubresource;
+	hizSubresource.m_mipmapCount = m_r->getDepthDownscale().getMipmapCount();
+	rgraphCtx.bindTexture(0, 4, m_r->getDepthDownscale().getHiZRt(), hizSubresource);
 
-	rgraphCtx.bindColorTexture(0, 4, m_r->getDownscaleBlur().getRt());
+	rgraphCtx.bindColorTexture(0, 5, m_r->getDownscaleBlur().getRt());
 
-	rgraphCtx.bindImage(0, 5, m_runCtx.m_rt, TextureSubresourceInfo());
+	rgraphCtx.bindImage(0, 6, m_runCtx.m_rt, TextureSubresourceInfo());
 
 	// Bind uniforms
-	SsrUniforms* unis = allocateAndBindUniforms<SsrUniforms*>(sizeof(SsrUniforms), cmdb, 0, 6);
-	unis->m_nearPad3 = Vec4(ctx.m_renderQueue->m_cameraNear);
+	SsrUniforms* unis = allocateAndBindUniforms<SsrUniforms*>(sizeof(SsrUniforms), cmdb, 0, 7);
+	unis->m_depthBufferSize = UVec2(m_r->getWidth(), m_r->getHeight()) >> 2u;
+	unis->m_framebufferSize = UVec2(m_r->getWidth(), m_r->getHeight());
+	unis->m_frameCount = m_r->getFrameCount() & MAX_U32;
+	unis->m_depthMipCount = m_r->getDepthDownscale().getMipmapCount();
+	unis->m_maxSteps = m_maxSteps;
+	unis->m_lightBufferMipCount = m_r->getDownscaleBlur().getMipmapCount();
 	unis->m_prevViewProjMatMulInvViewProjMat =
 		ctx.m_prevMatrices.m_viewProjection * ctx.m_matrices.m_viewProjectionJitter.getInverse();
 	unis->m_projMat = ctx.m_matrices.m_projectionJitter;
@@ -119,8 +124,7 @@ void Ssr::run(RenderPassWorkContext& rgraphCtx)
 	unis->m_normalMat = Mat3x4(ctx.m_matrices.m_view.getRotationPart());
 
 	// Dispatch
-	dispatchPPCompute(
-		cmdb, m_workgroupSize[0], m_workgroupSize[1], m_r->getWidth() / SSR_FRACTION, m_r->getHeight() / SSR_FRACTION);
+	dispatchPPCompute(cmdb, m_workgroupSize[0], m_workgroupSize[1], m_r->getWidth() / 2, m_r->getHeight());
 }
 
 } // end namespace anki

src/anki/renderer/Ssr.h

@@ -41,6 +41,7 @@ private:
 	TexturePtr m_rt;
 
 	Array<U32, 2> m_workgroupSize = {};
+	U32 m_maxSteps = 32;
 
 	class
 	{

src/anki/resource/ShaderProgramResource.h

@@ -268,12 +268,18 @@ inline ShaderProgramResourceVariantInitInfo& ShaderProgramResourceVariantInitInf
 	CString name, const T& value)
 {
 	const ShaderProgramResourceConstant* in = m_ptr->tryFindConstant(name);
-	ANKI_ASSERT(in);
-	ANKI_ASSERT(in->m_dataType == getShaderVariableTypeFromTypename<T>());
-	const U32 constIdx = U32(in - m_ptr->getConstants().getBegin());
-	m_constantValues[constIdx].m_constantIndex = constIdx;
-	memcpy(&m_constantValues[constIdx].m_int, &value, sizeof(T));
-	m_setConstants.set(constIdx);
+	if(in != nullptr)
+	{
+		ANKI_ASSERT(in->m_dataType == getShaderVariableTypeFromTypename<T>());
+		const U32 constIdx = U32(in - m_ptr->getConstants().getBegin());
+		m_constantValues[constIdx].m_constantIndex = constIdx;
+		memcpy(&m_constantValues[constIdx].m_int, &value, sizeof(T));
+		m_setConstants.set(constIdx);
+	}
+	else
+	{
+		ANKI_RESOURCE_LOGW("Constant not found: %s", name.cstr());
+	}
 	return *this;
 }