Cycles Denoising: Refactor denoiser tile handling
[blender.git] / intern / cycles / device / device_denoising.cpp
1 /*
2  * Copyright 2011-2017 Blender Foundation
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16
17 #include "device/device_denoising.h"
18
19 #include "kernel/filter/filter_defines.h"
20
21 CCL_NAMESPACE_BEGIN
22
23 DenoisingTask::DenoisingTask(Device *device, const DeviceTask &task)
24 : tiles_mem(device, "denoising tiles_mem", MEM_READ_WRITE),
25   storage(device),
26   buffer(device),
27   device(device)
28 {
29         radius = task.denoising_radius;
30         nlm_k_2 = powf(2.0f, lerp(-5.0f, 3.0f, task.denoising_strength));
31         if(task.denoising_relative_pca) {
32                 pca_threshold = -powf(10.0f, lerp(-8.0f, 0.0f, task.denoising_feature_strength));
33         }
34         else {
35                 pca_threshold = powf(10.0f, lerp(-5.0f, 3.0f, task.denoising_feature_strength));
36         }
37
38         render_buffer.pass_stride = task.pass_stride;
39         render_buffer.offset = task.pass_denoising_data;
40
41         target_buffer.pass_stride = task.pass_stride;
42         target_buffer.denoising_clean_offset = task.pass_denoising_clean;
43
44         functions.map_neighbor_tiles = function_bind(task.map_neighbor_tiles, _1, device);
45         functions.unmap_neighbor_tiles = function_bind(task.unmap_neighbor_tiles, _1, device);
46 }
47
48 DenoisingTask::~DenoisingTask()
49 {
50         storage.XtWX.free();
51         storage.XtWY.free();
52         storage.transform.free();
53         storage.rank.free();
54         storage.temporary_1.free();
55         storage.temporary_2.free();
56         storage.temporary_color.free();
57         buffer.mem.free();
58         tiles_mem.free();
59 }
60
61 void DenoisingTask::set_render_buffer(RenderTile *rtiles)
62 {
63         tiles = (TilesInfo*) tiles_mem.alloc(sizeof(TilesInfo)/sizeof(int));
64
65         device_ptr buffers[9];
66         for(int i = 0; i < 9; i++) {
67                 buffers[i] = rtiles[i].buffer;
68                 tiles->offsets[i] = rtiles[i].offset;
69                 tiles->strides[i] = rtiles[i].stride;
70         }
71         tiles->x[0] = rtiles[3].x;
72         tiles->x[1] = rtiles[4].x;
73         tiles->x[2] = rtiles[5].x;
74         tiles->x[3] = rtiles[5].x + rtiles[5].w;
75         tiles->y[0] = rtiles[1].y;
76         tiles->y[1] = rtiles[4].y;
77         tiles->y[2] = rtiles[7].y;
78         tiles->y[3] = rtiles[7].y + rtiles[7].h;
79
80         target_buffer.offset = rtiles[9].offset;
81         target_buffer.stride = rtiles[9].stride;
82         target_buffer.ptr    = rtiles[9].buffer;
83
84         functions.set_tiles(buffers);
85 }
86
87 void DenoisingTask::setup_denoising_buffer()
88 {
89         /* Expand filter_area by radius pixels and clamp the result to the extent of the neighboring tiles */
90         rect = rect_from_shape(filter_area.x, filter_area.y, filter_area.z, filter_area.w);
91         rect = rect_expand(rect, radius);
92         rect = rect_clip(rect, make_int4(tiles->x[0], tiles->y[0], tiles->x[3], tiles->y[3]));
93
94         buffer.passes = 14;
95         buffer.width = rect.z - rect.x;
96         buffer.stride = align_up(buffer.width, 4);
97         buffer.h = rect.w - rect.y;
98         int alignment_floats = divide_up(device->mem_sub_ptr_alignment(), sizeof(float));
99         buffer.pass_stride = align_up(buffer.stride * buffer.h, alignment_floats);
100         /* Pad the total size by four floats since the SIMD kernels might go a bit over the end. */
101         int mem_size = align_up(buffer.pass_stride * buffer.passes + 4, alignment_floats);
102         buffer.mem.alloc_to_device(mem_size, false);
103 }
104
105 void DenoisingTask::prefilter_shadowing()
106 {
107         device_ptr null_ptr = (device_ptr) 0;
108
109         device_sub_ptr unfiltered_a   (buffer.mem, 0,                    buffer.pass_stride);
110         device_sub_ptr unfiltered_b   (buffer.mem, 1*buffer.pass_stride, buffer.pass_stride);
111         device_sub_ptr sample_var     (buffer.mem, 2*buffer.pass_stride, buffer.pass_stride);
112         device_sub_ptr sample_var_var (buffer.mem, 3*buffer.pass_stride, buffer.pass_stride);
113         device_sub_ptr buffer_var     (buffer.mem, 5*buffer.pass_stride, buffer.pass_stride);
114         device_sub_ptr filtered_var   (buffer.mem, 6*buffer.pass_stride, buffer.pass_stride);
115         device_sub_ptr nlm_temporary_1(buffer.mem, 7*buffer.pass_stride, buffer.pass_stride);
116         device_sub_ptr nlm_temporary_2(buffer.mem, 8*buffer.pass_stride, buffer.pass_stride);
117         device_sub_ptr nlm_temporary_3(buffer.mem, 9*buffer.pass_stride, buffer.pass_stride);
118
119         nlm_state.temporary_1_ptr = *nlm_temporary_1;
120         nlm_state.temporary_2_ptr = *nlm_temporary_2;
121         nlm_state.temporary_3_ptr = *nlm_temporary_3;
122
123         /* Get the A/B unfiltered passes, the combined sample variance, the estimated variance of the sample variance and the buffer variance. */
124         functions.divide_shadow(*unfiltered_a, *unfiltered_b, *sample_var, *sample_var_var, *buffer_var);
125
126         /* Smooth the (generally pretty noisy) buffer variance using the spatial information from the sample variance. */
127         nlm_state.set_parameters(6, 3, 4.0f, 1.0f);
128         functions.non_local_means(*buffer_var, *sample_var, *sample_var_var, *filtered_var);
129
130         /* Reuse memory, the previous data isn't needed anymore. */
131         device_ptr filtered_a = *buffer_var,
132                    filtered_b = *sample_var;
133         /* Use the smoothed variance to filter the two shadow half images using each other for weight calculation. */
134         nlm_state.set_parameters(5, 3, 1.0f, 0.25f);
135         functions.non_local_means(*unfiltered_a, *unfiltered_b, *filtered_var, filtered_a);
136         functions.non_local_means(*unfiltered_b, *unfiltered_a, *filtered_var, filtered_b);
137
138         device_ptr residual_var = *sample_var_var;
139         /* Estimate the residual variance between the two filtered halves. */
140         functions.combine_halves(filtered_a, filtered_b, null_ptr, residual_var, 2, rect);
141
142         device_ptr final_a = *unfiltered_a,
143                    final_b = *unfiltered_b;
144         /* Use the residual variance for a second filter pass. */
145         nlm_state.set_parameters(4, 2, 1.0f, 0.5f);
146         functions.non_local_means(filtered_a, filtered_b, residual_var, final_a);
147         functions.non_local_means(filtered_b, filtered_a, residual_var, final_b);
148
149         /* Combine the two double-filtered halves to a final shadow feature. */
150         device_sub_ptr shadow_pass(buffer.mem, 4*buffer.pass_stride, buffer.pass_stride);
151         functions.combine_halves(final_a, final_b, *shadow_pass, null_ptr, 0, rect);
152 }
153
154 void DenoisingTask::prefilter_features()
155 {
156         device_sub_ptr unfiltered     (buffer.mem,  8*buffer.pass_stride, buffer.pass_stride);
157         device_sub_ptr variance       (buffer.mem,  9*buffer.pass_stride, buffer.pass_stride);
158         device_sub_ptr nlm_temporary_1(buffer.mem, 10*buffer.pass_stride, buffer.pass_stride);
159         device_sub_ptr nlm_temporary_2(buffer.mem, 11*buffer.pass_stride, buffer.pass_stride);
160         device_sub_ptr nlm_temporary_3(buffer.mem, 12*buffer.pass_stride, buffer.pass_stride);
161
162         nlm_state.temporary_1_ptr = *nlm_temporary_1;
163         nlm_state.temporary_2_ptr = *nlm_temporary_2;
164         nlm_state.temporary_3_ptr = *nlm_temporary_3;
165
166         int mean_from[]     = { 0, 1, 2, 12, 6,  7, 8 };
167         int variance_from[] = { 3, 4, 5, 13, 9, 10, 11};
168         int pass_to[]       = { 1, 2, 3, 0,  5,  6,  7};
169         for(int pass = 0; pass < 7; pass++) {
170                 device_sub_ptr feature_pass(buffer.mem, pass_to[pass]*buffer.pass_stride, buffer.pass_stride);
171                 /* Get the unfiltered pass and its variance from the RenderBuffers. */
172                 functions.get_feature(mean_from[pass], variance_from[pass], *unfiltered, *variance);
173                 /* Smooth the pass and store the result in the denoising buffers. */
174                 nlm_state.set_parameters(2, 2, 1.0f, 0.25f);
175                 functions.non_local_means(*unfiltered, *unfiltered, *variance, *feature_pass);
176         }
177 }
178
179 void DenoisingTask::prefilter_color()
180 {
181         int mean_from[]     = {20, 21, 22};
182         int variance_from[] = {23, 24, 25};
183         int mean_to[]       = { 8,  9, 10};
184         int variance_to[]   = {11, 12, 13};
185         int num_color_passes = 3;
186
187         storage.temporary_color.alloc_to_device(3*buffer.pass_stride, false);
188         device_sub_ptr nlm_temporary_1(storage.temporary_color, 0*buffer.pass_stride, buffer.pass_stride);
189         device_sub_ptr nlm_temporary_2(storage.temporary_color, 1*buffer.pass_stride, buffer.pass_stride);
190         device_sub_ptr nlm_temporary_3(storage.temporary_color, 2*buffer.pass_stride, buffer.pass_stride);
191
192         nlm_state.temporary_1_ptr = *nlm_temporary_1;
193         nlm_state.temporary_2_ptr = *nlm_temporary_2;
194         nlm_state.temporary_3_ptr = *nlm_temporary_3;
195
196         for(int pass = 0; pass < num_color_passes; pass++) {
197                 device_sub_ptr color_pass(storage.temporary_color, pass*buffer.pass_stride, buffer.pass_stride);
198                 device_sub_ptr color_var_pass(buffer.mem, variance_to[pass]*buffer.pass_stride, buffer.pass_stride);
199                 functions.get_feature(mean_from[pass], variance_from[pass], *color_pass, *color_var_pass);
200         }
201
202         device_sub_ptr depth_pass    (buffer.mem,                                 0,   buffer.pass_stride);
203         device_sub_ptr color_var_pass(buffer.mem, variance_to[0]*buffer.pass_stride, 3*buffer.pass_stride);
204         device_sub_ptr output_pass   (buffer.mem,     mean_to[0]*buffer.pass_stride, 3*buffer.pass_stride);
205         functions.detect_outliers(storage.temporary_color.device_pointer, *color_var_pass, *depth_pass, *output_pass);
206
207         storage.temporary_color.free();
208 }
209
210 void DenoisingTask::construct_transform()
211 {
212         storage.w = filter_area.z;
213         storage.h = filter_area.w;
214
215         storage.transform.alloc_to_device(storage.w*storage.h*TRANSFORM_SIZE, false);
216         storage.rank.alloc_to_device(storage.w*storage.h, false);
217
218         functions.construct_transform();
219 }
220
221 void DenoisingTask::reconstruct()
222 {
223
224         device_only_memory<float> temporary_1(device, "Denoising NLM temporary 1");
225         device_only_memory<float> temporary_2(device, "Denoising NLM temporary 2");
226         temporary_1.alloc_to_device(buffer.pass_stride, false);
227         temporary_2.alloc_to_device(buffer.pass_stride, false);
228         reconstruction_state.temporary_1_ptr = temporary_1.device_pointer;
229         reconstruction_state.temporary_2_ptr = temporary_2.device_pointer;
230
231         storage.XtWX.alloc_to_device(storage.w*storage.h*XTWX_SIZE, false);
232         storage.XtWY.alloc_to_device(storage.w*storage.h*XTWY_SIZE, false);
233
234         reconstruction_state.filter_window = rect_from_shape(filter_area.x-rect.x, filter_area.y-rect.y, storage.w, storage.h);
235         int tile_coordinate_offset = filter_area.y*target_buffer.stride + filter_area.x;
236         reconstruction_state.buffer_params = make_int4(target_buffer.offset + tile_coordinate_offset,
237                                                        target_buffer.stride,
238                                                        target_buffer.pass_stride,
239                                                        target_buffer.denoising_clean_offset);
240         reconstruction_state.source_w = rect.z-rect.x;
241         reconstruction_state.source_h = rect.w-rect.y;
242
243         device_sub_ptr color_ptr    (buffer.mem,  8*buffer.pass_stride, 3*buffer.pass_stride);
244         device_sub_ptr color_var_ptr(buffer.mem, 11*buffer.pass_stride, 3*buffer.pass_stride);
245         functions.reconstruct(*color_ptr, *color_var_ptr, target_buffer.ptr);
246 }
247
248 void DenoisingTask::run_denoising(RenderTile *tile)
249 {
250         RenderTile rtiles[10];
251         rtiles[4] = *tile;
252         functions.map_neighbor_tiles(rtiles);
253         set_render_buffer(rtiles);
254
255         setup_denoising_buffer();
256
257         prefilter_shadowing();
258         prefilter_features();
259         prefilter_color();
260
261         construct_transform();
262         reconstruct();
263
264         functions.unmap_neighbor_tiles(rtiles);
265 }
266
267 CCL_NAMESPACE_END