Browse Source

Convert IndirectDiffuseDenoise to HLSL

Panagiotis Christopoulos Charitos 3 years ago
parent
commit
d94dd5a707

+ 124 - 0
AnKi/Shaders/BilateralFilter.hlsl

@@ -0,0 +1,124 @@
+// Copyright (C) 2009-2022, Panagiotis Christopoulos Charitos and contributors.
+// All rights reserved.
+// Code licensed under the BSD License.
+// http://www.anki3d.org/LICENSE
+
+// Contains a bounch of edge stopping functions plus some other things
+
+#pragma once
+
+#include <AnKi/Shaders/Common.hlsl>
+
+// https://cs.dartmouth.edu/~wjarosz/publications/mara17towards.html
+F32 calculateBilateralWeightDepth(F32 depthCenter, F32 depthTap, F32 phi)
+{
+	const F32 diff = abs(depthTap - depthCenter);
+#if 0
+	return max(0.0, 1.0 - diff * phi);
+#else
+	return sqrt(1.0 / (kEpsilonF32 + diff)) * phi;
+#endif
+}
+
+// From the SVGF sample code. Depth is linear
+F32 calculateBilateralWeightDepth2(F32 depthCenter, F32 depthTap, F32 phi)
+{
+	return (phi == 0.0) ? 0.0 : abs(depthCenter - depthTap) / phi;
+}
+
+// https://cs.dartmouth.edu/~wjarosz/publications/mara17towards.html
+F32 calculateBilateralWeightNormal(Vec3 center, Vec3 tap, F32 phi)
+{
+	F32 normalCloseness = dot(tap, center);
+	normalCloseness = normalCloseness * normalCloseness;
+	normalCloseness = normalCloseness * normalCloseness;
+
+	const F32 normalError = (1.0 - normalCloseness);
+	return max((1.0 - normalError * phi), 0.0);
+}
+
+F32 calculateBilateralWeightNormalCos(Vec3 ref, Vec3 tap, F32 phi)
+{
+	return pow(saturate(dot(ref, tap)), phi);
+}
+
+// https://cs.dartmouth.edu/~wjarosz/publications/mara17towards.html
+F32 calculateBilateralWeightPlane(Vec3 positionCenter, Vec3 centerNormal, Vec3 positionTap, Vec3 normalTap, F32 phi)
+{
+	const F32 lowDistanceThreshold2 = 0.001;
+
+	// Change in position in camera space
+	const Vec3 dq = positionCenter - positionTap;
+
+	// How far away is this point from the original sample in camera space? (Max value is unbounded)
+	const F32 distance2 = dot(dq, dq);
+
+	// How far off the expected plane (on the perpendicular) is this point? Max value is unbounded.
+	const F32 planeError = max(abs(dot(dq, normalTap)), abs(dot(dq, centerNormal)));
+
+	return (distance2 < lowDistanceThreshold2) ? 1.0
+											   : pow(max(0.0, 1.0 - 2.0 * phi * planeError / sqrt(distance2)), 2.0);
+}
+
+// https://cs.dartmouth.edu/~wjarosz/publications/mara17towards.html
+F32 calculateBilateralWeightRoughness(F32 roughnessCenter, F32 roughnessTap, F32 phi)
+{
+	const F32 gDiff = abs(roughnessCenter - roughnessTap) * 10.0;
+	return max(0.0, 1.0 - (gDiff * phi));
+}
+
+// From SVGF sample code.
+F32 calculateBilateralWeightLinearDepthAndLuminance(F32 depthCenter, F32 luminanceCenter, F32 depthTap,
+													F32 luminanceTap, F32 phiDepth, F32 phiLuminance)
+{
+	const F32 wZ = calculateBilateralWeightDepth2(depthCenter, depthTap, phiDepth);
+	const F32 wL = abs(luminanceCenter - luminanceTap) / phiLuminance;
+	return exp(0.0 - max(wL, 0.0) - max(wZ, 0.0));
+}
+
+// Compute a weight given 2 viewspace positions.
+F32 calculateBilateralWeightViewspacePosition(Vec3 posCenter, Vec3 posTap, F32 sigma)
+{
+	const Vec3 t = posCenter - posTap;
+	const F32 dist2 = dot(t, t) + t.z * t.z;
+	return min(exp(-(dist2) / sigma), 1.0);
+}
+
+struct SpatialBilateralContext
+{
+	U32 sampleCount;
+	UVec2 referencePointUnormalizedTextureUv;
+	U32 maxRadiusPixels;
+	U32 spiralTurnCount;
+	F32 randomRotationAngle;
+};
+
+// Initialize the spatial bilateral context. Its purpose it to create samples that form a spiral around the reference
+// point. To experiment and play with the values use this: https://www.shadertoy.com/view/3s3BRs
+SpatialBilateralContext spatialBilateralInit(U32 sampleCount, UVec2 referencePointUnormalizedTextureUv,
+											 U32 maxRadiusPixels, U32 spiralTurnCount, F32 time)
+{
+	SpatialBilateralContext ctx;
+
+	ctx.sampleCount = sampleCount;
+	ctx.referencePointUnormalizedTextureUv = referencePointUnormalizedTextureUv;
+	ctx.maxRadiusPixels = maxRadiusPixels;
+	ctx.spiralTurnCount = spiralTurnCount;
+
+	const UVec2 ref = referencePointUnormalizedTextureUv;
+	ctx.randomRotationAngle = F32((3u * ref.x) ^ (ref.y + ref.x * ref.y)) + time;
+
+	return ctx;
+}
+
+// Iterate to get a new sample. See spatialBilateralInit()
+UVec2 spatialBilateralIterate(SpatialBilateralContext ctx, U32 iteration)
+{
+	const F32 alpha = F32(F32(iteration) + 0.5) * (1.0 / F32(ctx.sampleCount));
+	const F32 angle = alpha * F32(ctx.spiralTurnCount) * 6.28 + ctx.randomRotationAngle;
+
+	const Vec2 unitOffset = Vec2(cos(angle), sin(angle));
+
+	const IVec2 tapOffset = IVec2(alpha * F32(ctx.maxRadiusPixels) * unitOffset);
+	return UVec2(IVec2(ctx.referencePointUnormalizedTextureUv) + tapOffset);
+}

+ 0 - 98
AnKi/Shaders/IndirectDiffuseDenoise.glsl

@@ -1,98 +0,0 @@
-// 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 BLUR_ORIENTATION 0 1 // 0: in X axis, 1: in Y axis
-
-#include <AnKi/Shaders/Include/MiscRendererTypes.h>
-#include <AnKi/Shaders/PackFunctions.glsl>
-#include <AnKi/Shaders/Functions.glsl>
-#include <AnKi/Shaders/BilateralFilter.glsl>
-
-layout(set = 0, binding = 0) uniform sampler u_linearAnyClampSampler;
-layout(set = 0, binding = 1) uniform ANKI_RP texture2D u_toDenoiseTex;
-layout(set = 0, binding = 2) uniform texture2D u_depthTex;
-
-#if defined(ANKI_COMPUTE_SHADER)
-const UVec2 kWorkgroupSize = UVec2(8u, 8u);
-layout(local_size_x = kWorkgroupSize.x, local_size_y = kWorkgroupSize.y) in;
-
-layout(set = 0, binding = 3) writeonly uniform ANKI_RP image2D u_outImg;
-#else
-layout(location = 0) in Vec2 in_uv;
-layout(location = 0) out Vec3 out_color;
-#endif
-
-layout(push_constant, std430, row_major) uniform b_pc
-{
-	IndirectDiffuseDenoiseUniforms u_unis;
-};
-
-Vec3 unproject(Vec2 ndc, F32 depth)
-{
-	const Vec4 worldPos4 = u_unis.m_invertedViewProjectionJitterMat * Vec4(ndc, depth, 1.0);
-	const Vec3 worldPos = worldPos4.xyz / worldPos4.w;
-	return worldPos;
-}
-
-void main()
-{
-#if defined(ANKI_COMPUTE_SHADER)
-	if(skipOutOfBoundsInvocations(kWorkgroupSize, u_unis.m_viewportSize))
-	{
-		return;
-	}
-
-	const Vec2 uv = (Vec2(gl_GlobalInvocationID.xy) + 0.5) / u_unis.m_viewportSizef;
-#else
-	const Vec2 uv = in_uv;
-#endif
-
-	// Reference
-	const F32 depthCenter = textureLod(u_depthTex, u_linearAnyClampSampler, uv, 0.0).r;
-	if(depthCenter == 1.0)
-	{
-#if defined(ANKI_COMPUTE_SHADER)
-		imageStore(u_outImg, IVec2(gl_GlobalInvocationID.xy), Vec4(0.0));
-#else
-		out_color = Vec3(0.0);
-#endif
-		return;
-	}
-
-	const Vec3 positionCenter = unproject(UV_TO_NDC(uv), depthCenter);
-
-	// Sample
-	ANKI_RP F32 weight = kEpsilonRp;
-	ANKI_RP Vec3 color = Vec3(0.0);
-
-	for(F32 i = -u_unis.m_sampleCountDiv2; i <= u_unis.m_sampleCountDiv2; i += 1.0)
-	{
-		const Vec2 texelSize = 1.0 / u_unis.m_viewportSizef;
-#if BLUR_ORIENTATION == 0
-		const Vec2 sampleUv = Vec2(uv.x + i * texelSize.x, uv.y);
-#else
-		const Vec2 sampleUv = Vec2(uv.x, uv.y + i * texelSize.y);
-#endif
-
-		const F32 depthTap = textureLod(u_depthTex, u_linearAnyClampSampler, sampleUv, 0.0).r;
-
-		ANKI_RP F32 w = calculateBilateralWeightDepth(depthCenter, depthTap, 1.0);
-		// w *= gaussianWeight(0.4, abs(F32(i)) / (u_unis.m_sampleCountDiv2 * 2.0 + 1.0));
-		weight += w;
-
-		color += textureLod(u_toDenoiseTex, u_linearAnyClampSampler, sampleUv, 0.0).xyz * w;
-	}
-
-	// Normalize and store
-	color /= weight;
-
-#if defined(ANKI_COMPUTE_SHADER)
-	imageStore(u_outImg, IVec2(gl_GlobalInvocationID.xy), Vec4(color, 0.0));
-	// imageStore(u_outImg, IVec2(gl_GlobalInvocationID.xy), textureLod(u_toDenoiseTex, u_linearAnyClampSampler, uv,
-	// 0.0));
-#else
-	out_color = color;
-#endif
-}

+ 92 - 0
AnKi/Shaders/IndirectDiffuseDenoise.hlsl

@@ -0,0 +1,92 @@
+// 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 hlsl
+
+#pragma anki mutator BLUR_ORIENTATION 0 1 // 0: in X axis, 1: in Y axis
+
+#include <AnKi/Shaders/Include/MiscRendererTypes.h>
+#include <AnKi/Shaders/PackFunctions.hlsl>
+#include <AnKi/Shaders/Functions.hlsl>
+#include <AnKi/Shaders/BilateralFilter.hlsl>
+
+[[vk::binding(0)]] SamplerState g_linearAnyClampSampler;
+[[vk::binding(1)]] Texture2D<RVec4> g_toDenoiseTex;
+[[vk::binding(2)]] Texture2D g_depthTex;
+
+#if defined(ANKI_COMPUTE_SHADER)
+#	define THREAD_GROUP_SIZE_SQRT 8
+[[vk::binding(3)]] RWTexture2D<RVec3> g_outUav;
+#endif
+
+[[vk::push_constant]] ConstantBuffer<IndirectDiffuseDenoiseUniforms> g_uniforms;
+
+Vec3 unproject(Vec2 ndc, F32 depth)
+{
+	const Vec4 worldPos4 = mul(g_uniforms.m_invertedViewProjectionJitterMat, Vec4(ndc, depth, 1.0));
+	const Vec3 worldPos = worldPos4.xyz / worldPos4.w;
+	return worldPos;
+}
+
+#if defined(ANKI_COMPUTE_SHADER)
+ANKI_NUMTHREADS(THREAD_GROUP_SIZE_SQRT, THREAD_GROUP_SIZE_SQRT, 1)
+void main(UVec3 svDispatchThreadId : SV_DISPATCHTHREADID)
+#else
+RVec3 main([[vk::location(0)]] Vec2 uv : TEXCOORD): SV_TARGET0
+#endif
+{
+#if defined(ANKI_COMPUTE_SHADER)
+	if(skipOutOfBoundsInvocations(UVec2(THREAD_GROUP_SIZE_SQRT, THREAD_GROUP_SIZE_SQRT), g_uniforms.m_viewportSize,
+								  svDispatchThreadId))
+	{
+		return;
+	}
+
+	const Vec2 uv = (Vec2(svDispatchThreadId.xy) + 0.5) / g_uniforms.m_viewportSizef;
+#endif
+
+	// Reference
+	const F32 depthCenter = g_depthTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0).r;
+	if(depthCenter == 1.0)
+	{
+#if defined(ANKI_COMPUTE_SHADER)
+		g_outUav[svDispatchThreadId.xy] = RVec3(0.0, 0.0, 0.0);
+		return;
+#else
+		return RVec3(0.0, 0.0, 0.0);
+#endif
+	}
+
+	// Sample
+	RF32 weight = kEpsilonRF32;
+	RVec3 color = Vec3(0.0, 0.0, 0.0);
+
+	for(F32 i = -g_uniforms.m_sampleCountDiv2; i <= g_uniforms.m_sampleCountDiv2; i += 1.0)
+	{
+		const Vec2 texelSize = 1.0 / g_uniforms.m_viewportSizef;
+#if BLUR_ORIENTATION == 0
+		const Vec2 sampleUv = Vec2(uv.x + i * texelSize.x, uv.y);
+#else
+		const Vec2 sampleUv = Vec2(uv.x, uv.y + i * texelSize.y);
+#endif
+
+		const F32 depthTap = g_depthTex.SampleLevel(g_linearAnyClampSampler, sampleUv, 0.0).r;
+
+		RF32 w = calculateBilateralWeightDepth(depthCenter, depthTap, 1.0);
+		// w *= gaussianWeight(0.4, abs(F32(i)) / (g_uniforms.m_sampleCountDiv2 * 2.0 + 1.0));
+		weight += w;
+
+		color += g_toDenoiseTex.SampleLevel(g_linearAnyClampSampler, sampleUv, 0.0).xyz * w;
+	}
+
+	// Normalize and store
+	color /= weight;
+
+#if defined(ANKI_COMPUTE_SHADER)
+	g_outUav[svDispatchThreadId.xy] = color;
+#else
+	return color;
+#endif
+}

+ 1 - 1
AnKi/Shaders/IndirectDiffuseDenoiseCompute.ankiprog

@@ -4,5 +4,5 @@
 // http://www.anki3d.org/LICENSE
 
 #pragma anki start comp
-#include <AnKi/Shaders/IndirectDiffuseDenoise.glsl>
+#include <AnKi/Shaders/IndirectDiffuseDenoise.hlsl>
 #pragma anki end

+ 2 - 2
AnKi/Shaders/IndirectDiffuseDenoiseRaster.ankiprog

@@ -4,9 +4,9 @@
 // http://www.anki3d.org/LICENSE
 
 #pragma anki start vert
-#include <AnKi/Shaders/QuadVert.glsl>
+#include <AnKi/Shaders/QuadVert.hlsl>
 #pragma anki end
 
 #pragma anki start frag
-#include <AnKi/Shaders/IndirectDiffuseDenoise.glsl>
+#include <AnKi/Shaders/IndirectDiffuseDenoise.hlsl>
 #pragma anki end