anki-3d-engine/AnKi/Shaders/IndirectDiffuseVrsSriGeneration.ankiprog

// Copyright (C) 2009-2022, Panagiotis Christopoulos Charitos and contributors.
// All rights reserved.
// Code licensed under the BSD License.
// http://www.anki3d.org/LICENSE

#pragma anki mutator SRI_TEXEL_DIMENSION 8 16
#pragma anki mutator SHARED_MEMORY 0 1
#pragma anki mutator LIMIT_RATE_TO_2X2 0 1

#pragma anki start comp

#include <AnKi/Shaders/Functions.glsl>

layout(set = 0, binding = 0) uniform texture2D u_inputTex;
layout(set = 0, binding = 1) uniform sampler u_nearestClampSampler;

#if SRI_TEXEL_DIMENSION == 8
const UVec2 REGION_SIZE = UVec2(2u, 2u);
#else
const UVec2 REGION_SIZE = UVec2(2u, 4u);
#endif

const UVec2 WORKGROUP_SIZE = UVec2(SRI_TEXEL_DIMENSION) / REGION_SIZE;
layout(local_size_x = WORKGROUP_SIZE.x, local_size_y = WORKGROUP_SIZE.y, local_size_z = 1) in;

layout(set = 0, binding = 2) uniform writeonly uimage2D u_sriImg;

layout(push_constant, std140, row_major) uniform b_pc
{
	Vec2 u_oneOverViewportSize;
	F32 u_thresholdMeters;
	F32 u_padding0;
	Mat4 u_invertedProjectionJitter;
};

#if SHARED_MEMORY
// Ideally, we'd be able to calculate the min/max/average using subgroup operations, but there's no guarantee
// subgroupSize is large enough so we need shared memory as a fallback. We need gl_NumSubgroups entries, but it is not a
// constant, so estimate it assuming a subgroupSize of at least 8.
const U32 SHARED_MEMORY_ENTRIES = WORKGROUP_SIZE.x * WORKGROUP_SIZE.y / 8u;
shared Vec2 s_maxDerivative[SHARED_MEMORY_ENTRIES];
#endif

F32 sampleViewPositionZ(Vec2 uv, I32 offsetX, I32 offsetY)
{
	uv += Vec2(offsetX, offsetY) * u_oneOverViewportSize;
	const Vec2 ndc = UV_TO_NDC(uv);
	const F32 depth = textureLod(sampler2D(u_inputTex, u_nearestClampSampler), uv, 0.0).x;

	const Vec4 v4 = u_invertedProjectionJitter * Vec4(ndc, depth, 1.0);
	return v4.z / v4.w;
}

void main()
{
	const Vec2 uv = (Vec2(gl_GlobalInvocationID.xy) * Vec2(REGION_SIZE) + 0.5) * u_oneOverViewportSize;
	const Vec2 ndc = UV_TO_NDC(uv);

#if SRI_TEXEL_DIMENSION == 8
	// Get positions
	// l0.z  l0.w
	// l0.x  l0.y
	Vec4 l0;
	l0.x = sampleViewPositionZ(uv, 0, 0);
	l0.y = sampleViewPositionZ(uv, 1, 0);
	l0.z = sampleViewPositionZ(uv, 0, 1);
	l0.w = sampleViewPositionZ(uv, 1, 1);

	// Calculate derivatives.
	Vec2 a = Vec2(l0.y, l0.z);
	Vec2 b = Vec2(l0.x, l0.w);
	const Vec2 dx = abs(a - b);

	a = Vec2(l0.z, l0.w);
	b = Vec2(l0.x, l0.y);
	const Vec2 dy = abs(a - b);

	F32 maxDerivativeX = max(dx.x, dx.y);
	F32 maxDerivativeY = max(dy.x, dy.y);
#else
	// Get positions
	// l1.z  l1.w
	// l1.x  l1.y
	// l0.z  l0.w
	// l0.x  l0.y
	Vec4 l0;
	l0.x = sampleViewPositionZ(uv, 0, 0);
	l0.y = sampleViewPositionZ(uv, 1, 0);
	l0.z = sampleViewPositionZ(uv, 0, 1);
	l0.w = sampleViewPositionZ(uv, 1, 1);

	Vec4 l1;
	l1.x = sampleViewPositionZ(uv, 0, 2);
	l1.y = sampleViewPositionZ(uv, 1, 2);
	l1.z = sampleViewPositionZ(uv, 0, 3);
	l1.w = sampleViewPositionZ(uv, 1, 3);

	// Calculate derivatives.
	Vec4 a = Vec4(l0.y, l0.z, l1.y, l1.z);
	Vec4 b = Vec4(l0.x, l0.w, l1.x, l1.w);
	const Vec4 dx = abs(a - b);

	a = Vec4(l0.z, l0.w, l1.z, l1.y);
	b = Vec4(l0.x, l0.y, l1.x, l1.w);
	const Vec4 dy = abs(a - b);

	F32 maxDerivativeX = max(max(dx.x, dx.y), max(dx.z, dx.w));
	F32 maxDerivativeY = max(max(dy.x, dy.y), max(dy.z, dy.w));
#endif

	maxDerivativeX = subgroupMax(maxDerivativeX);
	maxDerivativeY = subgroupMax(maxDerivativeY);

#if SHARED_MEMORY
	// Store results in shared memory.
	ANKI_BRANCH if(subgroupElect())
	{
		s_maxDerivative[gl_SubgroupID] = Vec2(maxDerivativeX, maxDerivativeY);
	}

	memoryBarrierShared();
	barrier();
#endif

	// Write the result
	ANKI_BRANCH if(gl_LocalInvocationIndex == 0u)
	{
		// Get max across all subgroups.
#if SHARED_MEMORY
		Vec2 maxDerivative = s_maxDerivative[0];

		for(U32 i = 1u; i < gl_NumSubgroups; ++i)
		{
			maxDerivative = max(maxDerivative, s_maxDerivative[i]);
		}
#else
		const Vec2 maxDerivative = Vec2(maxDerivativeX, maxDerivativeY);
#endif

		// Determine shading rate.
		const F32 threshold1 = u_thresholdMeters;
		const F32 threshold2 = threshold1 * 0.4;

		UVec2 rate;
		rate.x = (maxDerivative.x > threshold1) ? 1u : ((maxDerivative.x > threshold2) ? 2u : 4u);
		rate.y = (maxDerivative.y > threshold1) ? 1u : ((maxDerivative.y > threshold2) ? 2u : 4u);

#if LIMIT_RATE_TO_2X2
		rate = min(rate, UVec2(2u));
#endif

		const UVec2 outTexelCoord = gl_WorkGroupID.xy;
		imageStore(u_sriImg, IVec2(outTexelCoord), UVec4(encodeVrsRate(rate)));
	}
}

#pragma anki end