Implemented reweighting

Author: Przemog1

31_HLSLPathTracer/app_resources/hlsl/pathtracer.hlsl

@@ -314,10 +314,19 @@ struct Unidirectional
         uint32_t base;
     };
 
-    // tmp
-    float calculateLumaRec709(float32_t4 color)
+    /**
+    * @brief Resets all buffers in the cascade to 0 at the given pixel coordinates.
+    *
+    * This function writes zero values to every buffer in the cascade
+    * for the specified 2D pixel location.
+    *
+    * @param coords Integer 2D coordinates of the pixel to reset.
+    * @param cascadeSize number of buffers in the cascade to clear.
+    */
+    void resetCascade(NBL_CONST_REF_ARG(int32_t2) coords, uint32_t cascadeSize)
     {
-        return 0.2126 * color.r + 0.7152 * color.g + 0.0722 * color.b;
+        for (int i = 0; i < 6; ++i)
+            cascade[uint3(coords.x, coords.y, i)] = float32_t4(0.0f, 0.0f, 0.0f, 0.0f);
     }
 
     void generateCascade(int32_t2 coords, uint32_t numSamples, uint32_t depth, NBL_CONST_REF_ARG(RWMCCascadeSettings) cascadeSettings, NBL_CONST_REF_ARG(scene_type) scene)
@@ -331,7 +340,6 @@ struct Unidirectional
             measure_type accumulation = getSingleSampleMeasure(i, depth, scene);
 
             const float luma = getLuma(accumulation);
-            //const float luma = calculateLumaRec709(float32_t4(accumulation, 1.0f));
 
             uint32_t lowerCascadeIndex = 0u;
             while (!(luma < upperScale) && lowerCascadeIndex < cascadeSettings.size - 2)

31_HLSLPathTracer/app_resources/hlsl/render.comp.hlsl

@@ -231,6 +231,7 @@ void main(uint32_t3 threadID : SV_DispatchThreadID)
         cascadeSettings.start = 1u;
         cascadeSettings.base = 8u;
 
+        pathtracer.resetCascade(coords, 6u);
         pathtracer.generateCascade(coords, pc.sampleCount, pc.depth, cascadeSettings, scene);
     }
 

31_HLSLPathTracer/app_resources/hlsl/render_common.hlsl

@@ -6,6 +6,9 @@ struct SPushConstants
     float32_t4x4 invMVP;
     int sampleCount;
     int depth;
+    uint32_t rwmcCascadeSize;
+    uint32_t rwmcCascadeStart;
+    uint32_t rwmcCascadeBase;
 };
 
 [[vk::push_constant]] SPushConstants pc;

31_HLSLPathTracer/app_resources/hlsl/reweighting.hlsl

@@ -1,5 +1,16 @@
 #include "nbl/builtin/hlsl/cpp_compat.hlsl"
+#include <nbl/builtin/hlsl/colorspace/encodeCIEXYZ.hlsl>
 
+struct SPushConstants
+{
+	uint32_t cascadeCount;
+	float base;
+	uint32_t sampleCount;
+	float minReliableLuma;
+	float kappa;
+};
+
+[[vk::push_constant]] SPushConstants pc;
 [[vk::image_format("rgba16f")]] [[vk::binding(0, 0)]] RWTexture2D<float32_t4> outImage;
 [[vk::image_format("rgba16f")]] [[vk::binding(1, 0)]] RWTexture2DArray<float32_t4> cascade;
 
@@ -10,6 +21,145 @@ NBL_CONSTEXPR uint32_t WorkgroupSize = 512;
 NBL_CONSTEXPR uint32_t MAX_DEPTH_LOG2 = 4;
 NBL_CONSTEXPR uint32_t MAX_SAMPLES_LOG2 = 10;
 
+struct RWMCReweightingParameters
+{
+    uint32_t lastCascadeIndex;
+    float initialEmin; // a minimum image brightness that we always consider reliable
+    float reciprocalBase;
+    float reciprocalN;
+    float reciprocalKappa;
+    float colorReliabilityFactor;
+    float NOverKappa;
+};
+
+RWMCReweightingParameters computeReweightingParameters(uint32_t cascadeCount, float base, uint32_t sampleCount, float minReliableLuma, float kappa)
+{
+    RWMCReweightingParameters retval;
+    retval.lastCascadeIndex = cascadeCount - 1u;
+    retval.initialEmin = minReliableLuma;
+    retval.reciprocalBase = 1.f / base;
+    const float N = float(sampleCount);
+    retval.reciprocalN = 1.f / N;
+    retval.reciprocalKappa = 1.f / kappa;
+    // if not interested in exact expected value estimation (kappa!=1.f), can usually accept a bit more variance relative to the image brightness we already have
+    // allow up to ~<cascadeBase> more energy in one sample to lessen bias in some cases
+    retval.colorReliabilityFactor = base + (1.f - base) * retval.reciprocalKappa;
+    retval.NOverKappa = N * retval.reciprocalKappa;
+
+    return retval;
+}
+
+struct RWMCCascadeSample
+{
+	float32_t3 centerValue;
+	float normalizedCenterLuma;
+	float normalizedNeighbourhoodAverageLuma;
+};
+
+// TODO: figure out what values should pixels outside have, 0.0f is incorrect
+float32_t3 RWMCsampleCascadeTexel(int32_t2 currentCoord, int32_t2 offset, uint32_t cascadeIndex)
+{
+	const int32_t2 texelCoord = currentCoord + offset;
+	if (any(texelCoord < int32_t2(0, 0)))
+		return float32_t3(0.0f, 0.0f, 0.0f);
+
+    float32_t4 output = cascade.Load(int32_t3(texelCoord, int32_t(cascadeIndex)));
+    return float32_t3(output.r, output.g, output.b);
+}
+
+float32_t calcLuma(in float32_t3 col)
+{
+    return hlsl::dot<float32_t3>(hlsl::transpose(colorspace::scRGBtoXYZ)[1], col);
+}
+
+RWMCCascadeSample RWMCSampleCascade(in int32_t2 coord, in uint cascadeIndex, in float reciprocalBaseI)
+{
+	float32_t3 neighbourhood[9];
+	neighbourhood[0] = RWMCsampleCascadeTexel(coord, int32_t2(-1, -1), cascadeIndex);
+    neighbourhood[1] = RWMCsampleCascadeTexel(coord, int32_t2(0, -1), cascadeIndex);
+    neighbourhood[2] = RWMCsampleCascadeTexel(coord, int32_t2(1, -1), cascadeIndex);
+    neighbourhood[3] = RWMCsampleCascadeTexel(coord, int32_t2(-1, 0), cascadeIndex);
+    neighbourhood[4] = RWMCsampleCascadeTexel(coord, int32_t2(0, 0), cascadeIndex);
+    neighbourhood[5] = RWMCsampleCascadeTexel(coord, int32_t2(1, 0), cascadeIndex);
+    neighbourhood[6] = RWMCsampleCascadeTexel(coord, int32_t2(-1, 1), cascadeIndex);
+    neighbourhood[7] = RWMCsampleCascadeTexel(coord, int32_t2(0, 1), cascadeIndex);
+    neighbourhood[8] = RWMCsampleCascadeTexel(coord, int32_t2(1, 1), cascadeIndex);
+
+	// numerical robustness
+	float32_t3 excl_hood_sum = ((neighbourhood[0] + neighbourhood[1]) + (neighbourhood[2] + neighbourhood[3])) +
+		((neighbourhood[5] + neighbourhood[6]) + (neighbourhood[7] + neighbourhood[8]));
+
+	RWMCCascadeSample retval;
+	retval.centerValue = neighbourhood[4];
+	retval.normalizedNeighbourhoodAverageLuma = retval.normalizedCenterLuma = calcLuma(neighbourhood[4]) * reciprocalBaseI;
+	retval.normalizedNeighbourhoodAverageLuma = (calcLuma(excl_hood_sum) * reciprocalBaseI + retval.normalizedNeighbourhoodAverageLuma) / 9.f;
+	return retval;
+}
+
+float32_t3 RWMCReweight(in RWMCReweightingParameters params, in int32_t2 coord)
+{
+	float reciprocalBaseI = 1.f;
+	RWMCCascadeSample curr = RWMCSampleCascade(coord, 0u, reciprocalBaseI);
+
+	float32_t3 accumulation = float32_t3(0.0f, 0.0f, 0.0f);
+	float Emin = params.initialEmin;
+
+	float prevNormalizedCenterLuma, prevNormalizedNeighbourhoodAverageLuma;
+	for (uint i = 0u; i <= params.lastCascadeIndex; i++)
+	{
+		const bool notFirstCascade = i != 0u;
+		const bool notLastCascade = i != params.lastCascadeIndex;
+
+		RWMCCascadeSample next;
+		if (notLastCascade)
+		{
+			reciprocalBaseI *= params.reciprocalBase;
+			next = RWMCSampleCascade(coord, i + 1u, reciprocalBaseI);
+		}
+
+
+		float reliability = 1.f;
+		// sample counting-based reliability estimation
+		if (params.reciprocalKappa <= 1.f)
+		{
+			float localReliability = curr.normalizedCenterLuma;
+			// reliability in 3x3 pixel block (see robustness)
+			float globalReliability = curr.normalizedNeighbourhoodAverageLuma;
+			if (notFirstCascade)
+			{
+				localReliability += prevNormalizedCenterLuma;
+				globalReliability += prevNormalizedNeighbourhoodAverageLuma;
+			}
+			if (notLastCascade)
+			{
+				localReliability += next.normalizedCenterLuma;
+				globalReliability += next.normalizedNeighbourhoodAverageLuma;
+			}
+			// check if above minimum sampling threshold (avg 9 sample occurences in 3x3 neighbourhood), then use per-pixel reliability (NOTE: tertiary op is in reverse)
+			reliability = globalReliability < params.reciprocalN ? globalReliability : localReliability;
+			{
+				const float accumLuma = calcLuma(accumulation);
+				if (accumLuma > Emin)
+					Emin = accumLuma;
+
+				const float colorReliability = Emin * reciprocalBaseI * params.colorReliabilityFactor;
+
+				reliability += colorReliability;
+				reliability *= params.NOverKappa;
+				reliability -= params.reciprocalKappa;
+				reliability = clamp(reliability * 0.5f, 0.f, 1.f);
+			}
+		}
+		accumulation += curr.centerValue * reliability;
+
+		prevNormalizedCenterLuma = curr.normalizedCenterLuma;
+		prevNormalizedNeighbourhoodAverageLuma = curr.normalizedNeighbourhoodAverageLuma;
+		curr = next;
+	}
+
+	return accumulation;
+}
+
 int32_t2 getCoordinates()
 {
     uint32_t width, height;
@@ -19,7 +169,7 @@ int32_t2 getCoordinates()
 
 // this function is for testing purpose
 // simply adds every cascade buffer, output shoud be nearly the same as output of default accumulator (RWMC off)
-void sumCascade(in const int32_t2 coords)
+float32_t3 sumCascade(in const int32_t2 coords)
 {
     float32_t3 accumulation = float32_t3(0.0f, 0.0f, 0.0f);
 
@@ -31,9 +181,7 @@ void sumCascade(in const int32_t2 coords)
 
     accumulation /= 32.0f;
 
-    float32_t4 output = float32_t4(accumulation, 1.0f);
-
-    outImage[coords] = output;
+	return accumulation;
 }
 
 [numthreads(WorkgroupSize, 1, 1)]
@@ -45,9 +193,11 @@ void main(uint32_t3 threadID : SV_DispatchThreadID)
         return;
 
     const int32_t2 coords = getCoordinates();
-    sumCascade(coords);
+    //float32_t3 color = sumCascade(coords);
 
-    // zero out cascade
-    for (int i = 0; i < 6; ++i)
-        cascade[uint3(coords.x, coords.y, i)] = float32_t4(0.0f, 0.0f, 0.0f, 0.0f);
+    RWMCReweightingParameters reweightingParameters = computeReweightingParameters(pc.cascadeCount, pc.base, pc.sampleCount, pc.minReliableLuma, pc.kappa);
+	float32_t3 color = RWMCReweight(reweightingParameters, coords);
+	color /= pc.sampleCount;
+
+	outImage[coords] = float32_t4(color, 1.0f);
 }

31_HLSLPathTracer/main.cpp

@@ -15,10 +15,24 @@ using namespace asset;
 using namespace ui;
 using namespace video;
 
-struct PTPushConstant {
+static constexpr uint32_t CascadeSize = 6u;
+struct PTPushConstant
+{
 	matrix4SIMD invMVP;
 	int sampleCount;
 	int depth;
+	const uint32_t rwmcCascadeSize = CascadeSize;
+	uint32_t rwmcCascadeStart;
+	uint32_t rwmcCascadeBase;
+};
+
+struct RWMCPushConstants
+{
+	const uint32_t cascadeSize = CascadeSize;
+	float base;
+	uint32_t sampleCount;
+	float minReliableLuma;
+	float kappa;
 };
 
 // TODO: Add a QueryPool for timestamping once its ready
@@ -509,8 +523,14 @@ class HLSLComputePathtracer final : public examples::SimpleWindowedApplication,
 
 				// Create reweighting pipeline
 				{
+					const nbl::asset::SPushConstantRange pcRange = {
+							.stageFlags = IShader::E_SHADER_STAGE::ESS_COMPUTE,
+							.offset = 0,
+							.size = sizeof(RWMCPushConstants)
+					};
+
 					auto pipelineLayout = m_device->createPipelineLayout(
-						{},
+						{ &pcRange, 1 },
 						core::smart_refctd_ptr(gpuDescriptorSetLayout0)
 					);
 
@@ -1098,6 +1118,15 @@ class HLSLComputePathtracer final : public examples::SimpleWindowedApplication,
 			m_oracle.reportBeginFrameRecord();
 			m_camera.mapKeysToWASD();
 
+			// set initial push constants contents
+			rwmcPushConstants.base = 8.0f;
+			rwmcPushConstants.sampleCount = spp;
+			rwmcPushConstants.minReliableLuma = 1.0f;
+			rwmcPushConstants.kappa = 5.0f;
+
+			pc.rwmcCascadeStart = 1.0;
+			pc.rwmcCascadeBase = 8.0f;
+
 			return true;
 		}
 
@@ -1162,11 +1191,12 @@ class HLSLComputePathtracer final : public examples::SimpleWindowedApplication,
 				cmdbuf->reset(IGPUCommandBuffer::RESET_FLAGS::NONE);
 				// disregard surface/swapchain transformation for now
 				const auto viewProjectionMatrix = m_camera.getConcatenatedMatrix();
-				PTPushConstant pc;
 				viewProjectionMatrix.getInverseTransform(pc.invMVP);
 				pc.sampleCount = spp;
 				pc.depth = depth;
 
+				rwmcPushConstants.sampleCount = spp;
+
 				// safe to proceed
 				// upload buffer data
 				cmdbuf->beginDebugMarker("ComputeShaderPathtracer IMGUI Frame");
@@ -1293,6 +1323,7 @@ class HLSLComputePathtracer final : public examples::SimpleWindowedApplication,
 
 					cmdbuf->bindComputePipeline(pipeline);
 					cmdbuf->bindDescriptorSets(EPBP_COMPUTE, pipeline->getLayout(), 0u, 1u, &m_descriptorSet0.get());
+					cmdbuf->pushConstants(pipeline->getLayout(), IShader::E_SHADER_STAGE::ESS_COMPUTE, 0, sizeof(RWMCPushConstants), &rwmcPushConstants);
 					cmdbuf->dispatch(1 + (WindowDimensions.x * WindowDimensions.y - 1) / DefaultWorkGroupSize, 1u, 1u);
 				}
 
@@ -1573,7 +1604,6 @@ class HLSLComputePathtracer final : public examples::SimpleWindowedApplication,
 		smart_refctd_ptr<IGPUImageView> m_envMapView, m_scrambleView;
 		smart_refctd_ptr<IGPUBufferView> m_sequenceBufferView;
 		smart_refctd_ptr<IGPUImageView> m_outImgView;
-		static constexpr uint32_t CascadeSize = 6u;
 		smart_refctd_ptr<IGPUImageView> m_cascadeView;
 
 		// sync
@@ -1610,6 +1640,8 @@ class HLSLComputePathtracer final : public examples::SimpleWindowedApplication,
 		int spp = 32;
 		int depth = 3;
 		bool usePersistentWorkGroups = false;
+		RWMCPushConstants rwmcPushConstants;
+		PTPushConstant pc;
 
 		bool m_firstFrame = true;
 		IGPUCommandBuffer::SClearColorValue clearColor = { .float32 = {0.f,0.f,0.f,1.f} };