Cycles Denoising: Split main function into logical steps

1 files changed, 145 insertions, 122 deletions

diff --git a/intern/cycles/device/device_denoising.cpp b/intern/cycles/device/device_denoising.cpp
index 644cf6cd10e..4d2ba508aec 100644
--- a/intern/cycles/device/device_denoising.cpp
+++ b/intern/cycles/device/device_denoising.cpp
@@ -20,12 +20,24 @@
 
 CCL_NAMESPACE_BEGIN
 
-DenoisingTask::DenoisingTask(Device *device)
+DenoisingTask::DenoisingTask(Device *device, const DeviceTask &task)
 : tiles_mem(device, "denoising tiles_mem", MEM_READ_WRITE),
   storage(device),
   buffer(device),
   device(device)
 {
+	radius = task.denoising_radius;
+	nlm_k_2 = powf(2.0f, lerp(-5.0f, 3.0f, task.denoising_strength));
+	if(task.denoising_relative_pca) {
+		pca_threshold = -powf(10.0f, lerp(-8.0f, 0.0f, task.denoising_feature_strength));
+	}
+	else {
+		pca_threshold = powf(10.0f, lerp(-5.0f, 3.0f, task.denoising_feature_strength));
+	}
+
+	render_buffer.pass_stride = task.pass_stride;
+	render_buffer.denoising_data_offset = task.pass_denoising_data;
+	render_buffer.denoising_clean_offset = task.pass_denoising_clean;
 }
 
 DenoisingTask::~DenoisingTask()
@@ -41,26 +53,6 @@ DenoisingTask::~DenoisingTask()
 	tiles_mem.free();
 }
 
-void DenoisingTask::init_from_devicetask(const DeviceTask &task)
-{
-	radius = task.denoising_radius;
-	nlm_k_2 = powf(2.0f, lerp(-5.0f, 3.0f, task.denoising_strength));
-	if(task.denoising_relative_pca) {
-		pca_threshold = -powf(10.0f, lerp(-8.0f, 0.0f, task.denoising_feature_strength));
-	}
-	else {
-		pca_threshold = powf(10.0f, lerp(-5.0f, 3.0f, task.denoising_feature_strength));
-	}
-
-	render_buffer.pass_stride = task.pass_stride;
-	render_buffer.denoising_data_offset  = task.pass_denoising_data;
-	render_buffer.denoising_clean_offset = task.pass_denoising_clean;
-
-	/* Expand filter_area by radius pixels and clamp the result to the extent of the neighboring tiles */
-	rect = rect_from_shape(filter_area.x, filter_area.y, filter_area.z, filter_area.w);
-	rect = rect_expand(rect, radius);
-	rect = rect_clip(rect, make_int4(tiles->x[0], tiles->y[0], tiles->x[3], tiles->y[3]));
-}
 
 void DenoisingTask::tiles_from_rendertiles(RenderTile *rtiles)
 {
@@ -88,120 +80,142 @@ void DenoisingTask::tiles_from_rendertiles(RenderTile *rtiles)
 	functions.set_tiles(buffers);
 }
 
-bool DenoisingTask::run_denoising()
+void DenoisingTask::setup_denoising_buffer()
 {
-	/* Allocate denoising buffer. */
+	/* Expand filter_area by radius pixels and clamp the result to the extent of the neighboring tiles */
+	rect = rect_from_shape(filter_area.x, filter_area.y, filter_area.z, filter_area.w);
+	rect = rect_expand(rect, radius);
+	rect = rect_clip(rect, make_int4(tiles->x[0], tiles->y[0], tiles->x[3], tiles->y[3]));
+
 	buffer.passes = 14;
 	buffer.width = rect.z - rect.x;
 	buffer.stride = align_up(buffer.width, 4);
 	buffer.h = rect.w - rect.y;
-	buffer.pass_stride = align_up(buffer.stride * buffer.h, divide_up(device->mem_sub_ptr_alignment(), sizeof(float)));
-	buffer.mem.alloc_to_device(buffer.pass_stride * buffer.passes, false);
+	int alignment_floats = divide_up(device->mem_sub_ptr_alignment(), sizeof(float));
+	buffer.pass_stride = align_up(buffer.stride * buffer.h, alignment_floats);
+	/* Pad the total size by four floats since the SIMD kernels might go a bit over the end. */
+	int mem_size = align_up(buffer.pass_stride * buffer.passes + 4, alignment_floats);
+	buffer.mem.alloc_to_device(mem_size, false);
+}
 
+void DenoisingTask::prefilter_shadowing()
+{
 	device_ptr null_ptr = (device_ptr) 0;
 
-	/* Prefilter shadow feature. */
-	{
-		device_sub_ptr unfiltered_a   (buffer.mem, 0,                    buffer.pass_stride);
-		device_sub_ptr unfiltered_b   (buffer.mem, 1*buffer.pass_stride, buffer.pass_stride);
-		device_sub_ptr sample_var     (buffer.mem, 2*buffer.pass_stride, buffer.pass_stride);
-		device_sub_ptr sample_var_var (buffer.mem, 3*buffer.pass_stride, buffer.pass_stride);
-		device_sub_ptr buffer_var     (buffer.mem, 5*buffer.pass_stride, buffer.pass_stride);
-		device_sub_ptr filtered_var   (buffer.mem, 6*buffer.pass_stride, buffer.pass_stride);
-		device_sub_ptr nlm_temporary_1(buffer.mem, 7*buffer.pass_stride, buffer.pass_stride);
-		device_sub_ptr nlm_temporary_2(buffer.mem, 8*buffer.pass_stride, buffer.pass_stride);
-		device_sub_ptr nlm_temporary_3(buffer.mem, 9*buffer.pass_stride, buffer.pass_stride);
-
-		nlm_state.temporary_1_ptr = *nlm_temporary_1;
-		nlm_state.temporary_2_ptr = *nlm_temporary_2;
-		nlm_state.temporary_3_ptr = *nlm_temporary_3;
-
-		/* Get the A/B unfiltered passes, the combined sample variance, the estimated variance of the sample variance and the buffer variance. */
-		functions.divide_shadow(*unfiltered_a, *unfiltered_b, *sample_var, *sample_var_var, *buffer_var);
-
-		/* Smooth the (generally pretty noisy) buffer variance using the spatial information from the sample variance. */
-		nlm_state.set_parameters(6, 3, 4.0f, 1.0f);
-		functions.non_local_means(*buffer_var, *sample_var, *sample_var_var, *filtered_var);
-
-		/* Reuse memory, the previous data isn't needed anymore. */
-		device_ptr filtered_a = *buffer_var,
-		           filtered_b = *sample_var;
-		/* Use the smoothed variance to filter the two shadow half images using each other for weight calculation. */
-		nlm_state.set_parameters(5, 3, 1.0f, 0.25f);
-		functions.non_local_means(*unfiltered_a, *unfiltered_b, *filtered_var, filtered_a);
-		functions.non_local_means(*unfiltered_b, *unfiltered_a, *filtered_var, filtered_b);
-
-		device_ptr residual_var = *sample_var_var;
-		/* Estimate the residual variance between the two filtered halves. */
-		functions.combine_halves(filtered_a, filtered_b, null_ptr, residual_var, 2, rect);
-
-		device_ptr final_a = *unfiltered_a,
-		           final_b = *unfiltered_b;
-		/* Use the residual variance for a second filter pass. */
-		nlm_state.set_parameters(4, 2, 1.0f, 0.5f);
-		functions.non_local_means(filtered_a, filtered_b, residual_var, final_a);
-		functions.non_local_means(filtered_b, filtered_a, residual_var, final_b);
-
-		/* Combine the two double-filtered halves to a final shadow feature. */
-		device_sub_ptr shadow_pass(buffer.mem, 4*buffer.pass_stride, buffer.pass_stride);
-		functions.combine_halves(final_a, final_b, *shadow_pass, null_ptr, 0, rect);
-	}
+	device_sub_ptr unfiltered_a   (buffer.mem, 0,                    buffer.pass_stride);
+	device_sub_ptr unfiltered_b   (buffer.mem, 1*buffer.pass_stride, buffer.pass_stride);
+	device_sub_ptr sample_var     (buffer.mem, 2*buffer.pass_stride, buffer.pass_stride);
+	device_sub_ptr sample_var_var (buffer.mem, 3*buffer.pass_stride, buffer.pass_stride);
+	device_sub_ptr buffer_var     (buffer.mem, 5*buffer.pass_stride, buffer.pass_stride);
+	device_sub_ptr filtered_var   (buffer.mem, 6*buffer.pass_stride, buffer.pass_stride);
+	device_sub_ptr nlm_temporary_1(buffer.mem, 7*buffer.pass_stride, buffer.pass_stride);
+	device_sub_ptr nlm_temporary_2(buffer.mem, 8*buffer.pass_stride, buffer.pass_stride);
+	device_sub_ptr nlm_temporary_3(buffer.mem, 9*buffer.pass_stride, buffer.pass_stride);
+
+	nlm_state.temporary_1_ptr = *nlm_temporary_1;
+	nlm_state.temporary_2_ptr = *nlm_temporary_2;
+	nlm_state.temporary_3_ptr = *nlm_temporary_3;
+
+	/* Get the A/B unfiltered passes, the combined sample variance, the estimated variance of the sample variance and the buffer variance. */
+	functions.divide_shadow(*unfiltered_a, *unfiltered_b, *sample_var, *sample_var_var, *buffer_var);
+
+	/* Smooth the (generally pretty noisy) buffer variance using the spatial information from the sample variance. */
+	nlm_state.set_parameters(6, 3, 4.0f, 1.0f);
+	functions.non_local_means(*buffer_var, *sample_var, *sample_var_var, *filtered_var);
+
+	/* Reuse memory, the previous data isn't needed anymore. */
+	device_ptr filtered_a = *buffer_var,
+	           filtered_b = *sample_var;
+	/* Use the smoothed variance to filter the two shadow half images using each other for weight calculation. */
+	nlm_state.set_parameters(5, 3, 1.0f, 0.25f);
+	functions.non_local_means(*unfiltered_a, *unfiltered_b, *filtered_var, filtered_a);
+	functions.non_local_means(*unfiltered_b, *unfiltered_a, *filtered_var, filtered_b);
+
+	device_ptr residual_var = *sample_var_var;
+	/* Estimate the residual variance between the two filtered halves. */
+	functions.combine_halves(filtered_a, filtered_b, null_ptr, residual_var, 2, rect);
+
+	device_ptr final_a = *unfiltered_a,
+	           final_b = *unfiltered_b;
+	/* Use the residual variance for a second filter pass. */
+	nlm_state.set_parameters(4, 2, 1.0f, 0.5f);
+	functions.non_local_means(filtered_a, filtered_b, residual_var, final_a);
+	functions.non_local_means(filtered_b, filtered_a, residual_var, final_b);
+
+	/* Combine the two double-filtered halves to a final shadow feature. */
+	device_sub_ptr shadow_pass(buffer.mem, 4*buffer.pass_stride, buffer.pass_stride);
+	functions.combine_halves(final_a, final_b, *shadow_pass, null_ptr, 0, rect);
+}
 
-	/* Prefilter general features. */
-	{
-		device_sub_ptr unfiltered     (buffer.mem,  8*buffer.pass_stride, buffer.pass_stride);
-		device_sub_ptr variance       (buffer.mem,  9*buffer.pass_stride, buffer.pass_stride);
-		device_sub_ptr nlm_temporary_1(buffer.mem, 10*buffer.pass_stride, buffer.pass_stride);
-		device_sub_ptr nlm_temporary_2(buffer.mem, 11*buffer.pass_stride, buffer.pass_stride);
-		device_sub_ptr nlm_temporary_3(buffer.mem, 12*buffer.pass_stride, buffer.pass_stride);
-
-		nlm_state.temporary_1_ptr = *nlm_temporary_1;
-		nlm_state.temporary_2_ptr = *nlm_temporary_2;
-		nlm_state.temporary_3_ptr = *nlm_temporary_3;
-
-		int mean_from[]     = { 0, 1, 2, 12, 6,  7, 8 };
-		int variance_from[] = { 3, 4, 5, 13, 9, 10, 11};
-		int pass_to[]       = { 1, 2, 3, 0,  5,  6,  7};
-		for(int pass = 0; pass < 7; pass++) {
-			device_sub_ptr feature_pass(buffer.mem, pass_to[pass]*buffer.pass_stride, buffer.pass_stride);
-			/* Get the unfiltered pass and its variance from the RenderBuffers. */
-			functions.get_feature(mean_from[pass], variance_from[pass], *unfiltered, *variance);
-			/* Smooth the pass and store the result in the denoising buffers. */
-			nlm_state.set_parameters(2, 2, 1.0f, 0.25f);
-			functions.non_local_means(*unfiltered, *unfiltered, *variance, *feature_pass);
-		}
+void DenoisingTask::prefilter_features()
+{
+	device_sub_ptr unfiltered     (buffer.mem,  8*buffer.pass_stride, buffer.pass_stride);
+	device_sub_ptr variance       (buffer.mem,  9*buffer.pass_stride, buffer.pass_stride);
+	device_sub_ptr nlm_temporary_1(buffer.mem, 10*buffer.pass_stride, buffer.pass_stride);
+	device_sub_ptr nlm_temporary_2(buffer.mem, 11*buffer.pass_stride, buffer.pass_stride);
+	device_sub_ptr nlm_temporary_3(buffer.mem, 12*buffer.pass_stride, buffer.pass_stride);
+
+	nlm_state.temporary_1_ptr = *nlm_temporary_1;
+	nlm_state.temporary_2_ptr = *nlm_temporary_2;
+	nlm_state.temporary_3_ptr = *nlm_temporary_3;
+
+	int mean_from[]     = { 0, 1, 2, 12, 6,  7, 8 };
+	int variance_from[] = { 3, 4, 5, 13, 9, 10, 11};
+	int pass_to[]       = { 1, 2, 3, 0,  5,  6,  7};
+	for(int pass = 0; pass < 7; pass++) {
+		device_sub_ptr feature_pass(buffer.mem, pass_to[pass]*buffer.pass_stride, buffer.pass_stride);
+		/* Get the unfiltered pass and its variance from the RenderBuffers. */
+		functions.get_feature(mean_from[pass], variance_from[pass], *unfiltered, *variance);
+		/* Smooth the pass and store the result in the denoising buffers. */
+		nlm_state.set_parameters(2, 2, 1.0f, 0.25f);
+		functions.non_local_means(*unfiltered, *unfiltered, *variance, *feature_pass);
 	}
+}
 
-	/* Copy color passes. */
-	{
-		int mean_from[]     = {20, 21, 22};
-		int variance_from[] = {23, 24, 25};
-		int mean_to[]       = { 8,  9, 10};
-		int variance_to[]   = {11, 12, 13};
-		int num_color_passes = 3;
-
-		storage.temporary_color.alloc_to_device(3*buffer.pass_stride, false);
-
-		for(int pass = 0; pass < num_color_passes; pass++) {
-			device_sub_ptr color_pass(storage.temporary_color, pass*buffer.pass_stride, buffer.pass_stride);
-			device_sub_ptr color_var_pass(buffer.mem, variance_to[pass]*buffer.pass_stride, buffer.pass_stride);
-			functions.get_feature(mean_from[pass], variance_from[pass], *color_pass, *color_var_pass);
-		}
-
-		{
-			device_sub_ptr depth_pass    (buffer.mem,                                 0,   buffer.pass_stride);
-			device_sub_ptr color_var_pass(buffer.mem, variance_to[0]*buffer.pass_stride, 3*buffer.pass_stride);
-			device_sub_ptr output_pass   (buffer.mem,     mean_to[0]*buffer.pass_stride, 3*buffer.pass_stride);
-			functions.detect_outliers(storage.temporary_color.device_pointer, *color_var_pass, *depth_pass, *output_pass);
-		}
+void DenoisingTask::prefilter_color()
+{
+	int mean_from[]     = {20, 21, 22};
+	int variance_from[] = {23, 24, 25};
+	int mean_to[]       = { 8,  9, 10};
+	int variance_to[]   = {11, 12, 13};
+	int num_color_passes = 3;
+
+	storage.temporary_color.alloc_to_device(3*buffer.pass_stride, false);
+	device_sub_ptr nlm_temporary_1(storage.temporary_color, 0*buffer.pass_stride, buffer.pass_stride);
+	device_sub_ptr nlm_temporary_2(storage.temporary_color, 1*buffer.pass_stride, buffer.pass_stride);
+	device_sub_ptr nlm_temporary_3(storage.temporary_color, 2*buffer.pass_stride, buffer.pass_stride);
+
+	nlm_state.temporary_1_ptr = *nlm_temporary_1;
+	nlm_state.temporary_2_ptr = *nlm_temporary_2;
+	nlm_state.temporary_3_ptr = *nlm_temporary_3;
+
+	for(int pass = 0; pass < num_color_passes; pass++) {
+		device_sub_ptr color_pass(storage.temporary_color, pass*buffer.pass_stride, buffer.pass_stride);
+		device_sub_ptr color_var_pass(buffer.mem, variance_to[pass]*buffer.pass_stride, buffer.pass_stride);
+		functions.get_feature(mean_from[pass], variance_from[pass], *color_pass, *color_var_pass);
 	}
 
+	device_sub_ptr depth_pass    (buffer.mem,                                 0,   buffer.pass_stride);
+	device_sub_ptr color_var_pass(buffer.mem, variance_to[0]*buffer.pass_stride, 3*buffer.pass_stride);
+	device_sub_ptr output_pass   (buffer.mem,     mean_to[0]*buffer.pass_stride, 3*buffer.pass_stride);
+	functions.detect_outliers(storage.temporary_color.device_pointer, *color_var_pass, *depth_pass, *output_pass);
+
+	storage.temporary_color.free();
+}
+
+void DenoisingTask::construct_transform()
+{
 	storage.w = filter_area.z;
 	storage.h = filter_area.w;
+
 	storage.transform.alloc_to_device(storage.w*storage.h*TRANSFORM_SIZE, false);
 	storage.rank.alloc_to_device(storage.w*storage.h, false);
 
 	functions.construct_transform();
+}
+
+void DenoisingTask::reconstruct()
+{
 
 	device_only_memory<float> temporary_1(device, "Denoising NLM temporary 1");
 	device_only_memory<float> temporary_2(device, "Denoising NLM temporary 2");
@@ -222,13 +236,22 @@ bool DenoisingTask::run_denoising()
 	reconstruction_state.source_w = rect.z-rect.x;
 	reconstruction_state.source_h = rect.w-rect.y;
 
-	{
-		device_sub_ptr color_ptr    (buffer.mem,  8*buffer.pass_stride, 3*buffer.pass_stride);
-		device_sub_ptr color_var_ptr(buffer.mem, 11*buffer.pass_stride, 3*buffer.pass_stride);
-		functions.reconstruct(*color_ptr, *color_var_ptr, render_buffer.ptr);
-	}
+	device_sub_ptr color_ptr    (buffer.mem,  8*buffer.pass_stride, 3*buffer.pass_stride);
+	device_sub_ptr color_var_ptr(buffer.mem, 11*buffer.pass_stride, 3*buffer.pass_stride);
+	functions.reconstruct(*color_ptr, *color_var_ptr, render_buffer.ptr);
+}
+
+void DenoisingTask::run_denoising()
+{
+	setup_denoising_buffer();
+
+	prefilter_shadowing();
+	prefilter_features();
+	prefilter_color();
+
+	construct_transform();
+	reconstruct();
 
-	return true;
 }
 
 CCL_NAMESPACE_END