Cycles: Improve denoising speed on GPUs with small tile sizes

Previously, the NLM kernels would be launched once per offset with one thread per pixel.
However, with the smaller tile sizes that are now feasible, there wasn't enough work to fully occupy GPUs which results in a significant slowdown.

Therefore, the kernels are now launched in a single call that handles all offsets at once.
This has two downsides: Memory accesses to accumulating buffers are now atomic, and more importantly, the temporary memory now has to be allocated for every shift at once, increasing the required memory.
On the other hand, of course, the smaller tiles significantly reduce the size of the memory.

The main bottleneck right now is the construction of the transformation - there is nothing to be parallelized there, one thread per pixel is the maximum.
I tried to parallelize the SVD implementation by storing the matrix in shared memory and launching one block per pixel, but that wasn't really going anywhere.

To make the new code somewhat readable, the handling of rectangular regions was cleaned up a bit and commented, it should be easier to understand what's going on now.
Also, some variables have been renamed to make the difference between buffer width and stride more apparent, in addition to some general style cleanup.

8 files changed, 312 insertions, 183 deletions

diff --git a/intern/cycles/kernel/CMakeLists.txt b/intern/cycles/kernel/CMakeLists.txt
index de056ce97f0..5f10bdf2041 100644
--- a/intern/cycles/kernel/CMakeLists.txt
+++ b/intern/cycles/kernel/CMakeLists.txt
@@ -254,6 +254,7 @@ set(SRC_UTIL_HEADERS
 	../util/util_math_int3.h
 	../util/util_math_int4.h
 	../util/util_math_matrix.h
+	../util/util_rect.h
 	../util/util_static_assert.h
 	../util/util_transform.h
 	../util/util_texture.h
diff --git a/intern/cycles/kernel/filter/filter_nlm_cpu.h b/intern/cycles/kernel/filter/filter_nlm_cpu.h
index 5e989331bc2..e2da0fd872b 100644
--- a/intern/cycles/kernel/filter/filter_nlm_cpu.h
+++ b/intern/cycles/kernel/filter/filter_nlm_cpu.h
@@ -21,7 +21,7 @@ ccl_device_inline void kernel_filter_nlm_calc_difference(int dx, int dy,
                                                          const float *ccl_restrict variance_image,
                                                          float *difference_image,
                                                          int4 rect,
-                                                         int w,
+                                                         int stride,
                                                          int channel_offset,
                                                          float a,
                                                          float k_2)
@@ -31,15 +31,15 @@ ccl_device_inline void kernel_filter_nlm_calc_difference(int dx, int dy,
 			float diff = 0.0f;
 			int numChannels = channel_offset? 3 : 1;
 			for(int c = 0; c < numChannels; c++) {
-				float cdiff = weight_image[c*channel_offset + y*w+x] - weight_image[c*channel_offset + (y+dy)*w+(x+dx)];
-				float pvar = variance_image[c*channel_offset + y*w+x];
-				float qvar = variance_image[c*channel_offset + (y+dy)*w+(x+dx)];
+				float cdiff = weight_image[c*channel_offset + y*stride + x] - weight_image[c*channel_offset + (y+dy)*stride + (x+dx)];
+				float pvar = variance_image[c*channel_offset + y*stride + x];
+				float qvar = variance_image[c*channel_offset + (y+dy)*stride + (x+dx)];
 				diff += (cdiff*cdiff - a*(pvar + min(pvar, qvar))) / (1e-8f + k_2*(pvar+qvar));
 			}
 			if(numChannels > 1) {
 				diff *= 1.0f/numChannels;
 			}
-			difference_image[y*w+x] = diff;
+			difference_image[y*stride + x] = diff;
 		}
 	}
 }
@@ -47,7 +47,7 @@ ccl_device_inline void kernel_filter_nlm_calc_difference(int dx, int dy,
 ccl_device_inline void kernel_filter_nlm_blur(const float *ccl_restrict difference_image,
                                               float *out_image,
                                               int4 rect,
-                                              int w,
+                                              int stride,
                                               int f)
 {
 	int aligned_lowx = rect.x / 4;
@@ -56,17 +56,17 @@ ccl_device_inline void kernel_filter_nlm_blur(const float *ccl_restrict differen
 		const int low = max(rect.y, y-f);
 		const int high = min(rect.w, y+f+1);
 		for(int x = rect.x; x < rect.z; x++) {
-			out_image[y*w+x] = 0.0f;
+			out_image[y*stride + x] = 0.0f;
 		}
 		for(int y1 = low; y1 < high; y1++) {
-			float4* out_image4 = (float4*)(out_image + y*w);
-			float4* difference_image4 = (float4*)(difference_image + y1*w);
+			float4* out_image4 = (float4*)(out_image + y*stride);
+			float4* difference_image4 = (float4*)(difference_image + y1*stride);
 			for(int x = aligned_lowx; x < aligned_highx; x++) {
 				out_image4[x] += difference_image4[x];
 			}
 		}
 		for(int x = rect.x; x < rect.z; x++) {
-			out_image[y*w+x] *= 1.0f/(high - low);
+			out_image[y*stride + x] *= 1.0f/(high - low);
 		}
 	}
 }
@@ -74,12 +74,12 @@ ccl_device_inline void kernel_filter_nlm_blur(const float *ccl_restrict differen
 ccl_device_inline void kernel_filter_nlm_calc_weight(const float *ccl_restrict difference_image,
                                                      float *out_image,
                                                      int4 rect,
-                                                     int w,
+                                                     int stride,
                                                      int f)
 {
 	for(int y = rect.y; y < rect.w; y++) {
 		for(int x = rect.x; x < rect.z; x++) {
-			out_image[y*w+x] = 0.0f;
+			out_image[y*stride + x] = 0.0f;
 		}
 	}
 	for(int dx = -f; dx <= f; dx++) {
@@ -87,7 +87,7 @@ ccl_device_inline void kernel_filter_nlm_calc_weight(const float *ccl_restrict d
 		int neg_dx = min(0, dx);
 		for(int y = rect.y; y < rect.w; y++) {
 			for(int x = rect.x-neg_dx; x < rect.z-pos_dx; x++) {
-				out_image[y*w+x] += difference_image[y*w+dx+x];
+				out_image[y*stride + x] += difference_image[y*stride + x+dx];
 			}
 		}
 	}
@@ -95,7 +95,7 @@ ccl_device_inline void kernel_filter_nlm_calc_weight(const float *ccl_restrict d
 		for(int x = rect.x; x < rect.z; x++) {
 			const int low = max(rect.x, x-f);
 			const int high = min(rect.z, x+f+1);
-			out_image[y*w+x] = fast_expf(-max(out_image[y*w+x] * (1.0f/(high - low)), 0.0f));
+			out_image[y*stride + x] = fast_expf(-max(out_image[y*stride + x] * (1.0f/(high - low)), 0.0f));
 		}
 	}
 }
@@ -106,7 +106,7 @@ ccl_device_inline void kernel_filter_nlm_update_output(int dx, int dy,
                                                        float *out_image,
                                                        float *accum_image,
                                                        int4 rect,
-                                                       int w,
+                                                       int stride,
                                                        int f)
 {
 	for(int y = rect.y; y < rect.w; y++) {
@@ -115,11 +115,11 @@ ccl_device_inline void kernel_filter_nlm_update_output(int dx, int dy,
 			const int high = min(rect.z, x+f+1);
 			float sum = 0.0f;
 			for(int x1 = low; x1 < high; x1++) {
-				sum += difference_image[y*w+x1];
+				sum += difference_image[y*stride + x1];
 			}
 			float weight = sum * (1.0f/(high - low));
-			accum_image[y*w+x] += weight;
-			out_image[y*w+x] += weight*image[(y+dy)*w+(x+dx)];
+			accum_image[y*stride + x] += weight;
+			out_image[y*stride + x] += weight*image[(y+dy)*stride + (x+dx)];
 		}
 	}
 }
@@ -132,31 +132,31 @@ ccl_device_inline void kernel_filter_nlm_construct_gramian(int dx, int dy,
                                                            float *XtWX,
                                                            float3 *XtWY,
                                                            int4 rect,
-                                                           int4 filter_rect,
-                                                           int w, int h, int f,
+                                                           int4 filter_window,
+                                                           int stride, int f,
                                                            int pass_stride)
 {
+	int4 clip_area = rect_clip(rect, filter_window);
 	/* fy and fy are in filter-window-relative coordinates, while x and y are in feature-window-relative coordinates. */
-	for(int fy = max(0, rect.y-filter_rect.y); fy < min(filter_rect.w, rect.w-filter_rect.y); fy++) {
-		int y = fy + filter_rect.y;
-		for(int fx = max(0, rect.x-filter_rect.x); fx < min(filter_rect.z, rect.z-filter_rect.x); fx++) {
-			int x = fx + filter_rect.x;
+	for(int y = clip_area.y; y < clip_area.w; y++) {
+		for(int x = clip_area.x; x < clip_area.z; x++) {
 			const int low = max(rect.x, x-f);
 			const int high = min(rect.z, x+f+1);
 			float sum = 0.0f;
 			for(int x1 = low; x1 < high; x1++) {
-				sum += difference_image[y*w+x1];
+				sum += difference_image[y*stride + x1];
 			}
 			float weight = sum * (1.0f/(high - low));
 
-			int storage_ofs = fy*filter_rect.z + fx;
+			int storage_ofs = coord_to_local_index(filter_window, x, y);
 			float  *l_transform = transform + storage_ofs*TRANSFORM_SIZE;
 			float  *l_XtWX = XtWX + storage_ofs*XTWX_SIZE;
 			float3 *l_XtWY = XtWY + storage_ofs*XTWY_SIZE;
 			int    *l_rank = rank + storage_ofs;
 
 			kernel_filter_construct_gramian(x, y, 1,
-			                                dx, dy, w, h,
+			                                dx, dy,
+			                                stride,
 			                                pass_stride,
 			                                buffer,
 			                                l_transform, l_rank,
diff --git a/intern/cycles/kernel/filter/filter_nlm_gpu.h b/intern/cycles/kernel/filter/filter_nlm_gpu.h
index 2c5ac807051..4ca49ea6733 100644
--- a/intern/cycles/kernel/filter/filter_nlm_gpu.h
+++ b/intern/cycles/kernel/filter/filter_nlm_gpu.h
@@ -16,57 +16,114 @@
 
 CCL_NAMESPACE_BEGIN
 
+/* Determines pixel coordinates and offset for the current thread.
+ * Returns whether the thread should do any work.
+ *
+ * All coordinates are relative to the denoising buffer!
+ *
+ * Window is the rect that should be processed.
+ * co is filled with (x, y, dx, dy).
+ */
+ccl_device_inline bool get_nlm_coords_window(int w, int h, int r, int stride,
+                                             int4 *rect, int4 *co, int *ofs,
+                                             int4 window)
+{
+	/* Determine the pixel offset that this thread should apply. */
+	int s = 2*r+1;
+	int si = ccl_global_id(1);
+	int sx = si % s;
+	int sy = si / s;
+	if(sy >= s) {
+		return false;
+	}
+	co->z = sx-r;
+	co->w = sy-r;
+
+	/* Pixels still need to lie inside the denoising buffer after applying the offset,
+	 * so determine the area for which this is the case. */
+	*rect = make_int4(max(0, -co->z),     max(0, -co->w),
+	              w - max(0,  co->z), h - max(0,  co->w));
+
+	/* Find the intersection of the area that we want to process (window) and the area
+	 * that can be processed (rect) to get the final area for this offset. */
+	int4 clip_area = rect_clip(window, *rect);
+
+	/* If the radius is larger than one of the sides of the window,
+	 * there will be shifts for which there is no usable pixel at all. */
+	if(!rect_is_valid(clip_area)) {
+		return false;
+	}
+
+	/* Map the linear thread index to pixels inside the clip area. */
+	int x, y;
+	if(!local_index_to_coord(clip_area, ccl_global_id(0), &x, &y)) {
+		return false;
+	}
+	co->x = x;
+	co->y = y;
+
+	*ofs = (sy*s + sx) * stride;
+
+	return true;
+}
+
+ccl_device_inline bool get_nlm_coords(int w, int h, int r, int stride,
+                                      int4 *rect, int4 *co, int *ofs)
+{
+	return get_nlm_coords_window(w, h, r, stride, rect, co, ofs, make_int4(0, 0, w, h));
+}
+
 ccl_device_inline void kernel_filter_nlm_calc_difference(int x, int y,
                                                          int dx, int dy,
                                                          const ccl_global float *ccl_restrict weight_image,
                                                          const ccl_global float *ccl_restrict variance_image,
                                                          ccl_global float *difference_image,
-                                                         int4 rect, int w,
+                                                         int4 rect, int stride,
                                                          int channel_offset,
                                                          float a, float k_2)
 {
 	float diff = 0.0f;
 	int numChannels = channel_offset? 3 : 1;
 	for(int c = 0; c < numChannels; c++) {
-		float cdiff = weight_image[c*channel_offset + y*w+x] - weight_image[c*channel_offset + (y+dy)*w+(x+dx)];
-		float pvar = variance_image[c*channel_offset + y*w+x];
-		float qvar = variance_image[c*channel_offset + (y+dy)*w+(x+dx)];
+		float cdiff = weight_image[c*channel_offset + y*stride + x] - weight_image[c*channel_offset + (y+dy)*stride + (x+dx)];
+		float pvar = variance_image[c*channel_offset + y*stride + x];
+		float qvar = variance_image[c*channel_offset + (y+dy)*stride + (x+dx)];
 		diff += (cdiff*cdiff - a*(pvar + min(pvar, qvar))) / (1e-8f + k_2*(pvar+qvar));
 	}
 	if(numChannels > 1) {
 		diff *= 1.0f/numChannels;
 	}
-	difference_image[y*w+x] = diff;
+	difference_image[y*stride + x] = diff;
 }
 
 ccl_device_inline void kernel_filter_nlm_blur(int x, int y,
                                               const ccl_global float *ccl_restrict difference_image,
                                               ccl_global float *out_image,
-                                              int4 rect, int w, int f)
+                                              int4 rect, int stride, int f)
 {
 	float sum = 0.0f;
 	const int low = max(rect.y, y-f);
 	const int high = min(rect.w, y+f+1);
 	for(int y1 = low; y1 < high; y1++) {
-		sum += difference_image[y1*w+x];
+		sum += difference_image[y1*stride + x];
 	}
 	sum *= 1.0f/(high-low);
-	out_image[y*w+x] = sum;
+	out_image[y*stride + x] = sum;
 }
 
 ccl_device_inline void kernel_filter_nlm_calc_weight(int x, int y,
                                                      const ccl_global float *ccl_restrict difference_image,
                                                      ccl_global float *out_image,
-                                                     int4 rect, int w, int f)
+                                                     int4 rect, int stride, int f)
 {
 	float sum = 0.0f;
 	const int low = max(rect.x, x-f);
 	const int high = min(rect.z, x+f+1);
 	for(int x1 = low; x1 < high; x1++) {
-		sum += difference_image[y*w+x1];
+		sum += difference_image[y*stride + x1];
 	}
 	sum *= 1.0f/(high-low);
-	out_image[y*w+x] = fast_expf(-max(sum, 0.0f));
+	out_image[y*stride + x] = fast_expf(-max(sum, 0.0f));
 }
 
 ccl_device_inline void kernel_filter_nlm_update_output(int x, int y,
@@ -75,25 +132,25 @@ ccl_device_inline void kernel_filter_nlm_update_output(int x, int y,
                                                        const ccl_global float *ccl_restrict image,
                                                        ccl_global float *out_image,
                                                        ccl_global float *accum_image,
-                                                       int4 rect, int w, int f)
+                                                       int4 rect, int stride, int f)
 {
 	float sum = 0.0f;
 	const int low = max(rect.x, x-f);
 	const int high = min(rect.z, x+f+1);
 	for(int x1 = low; x1 < high; x1++) {
-		sum += difference_image[y*w+x1];
+		sum += difference_image[y*stride + x1];
 	}
 	sum *= 1.0f/(high-low);
 	if(out_image) {
-		accum_image[y*w+x] += sum;
-		out_image[y*w+x] += sum*image[(y+dy)*w+(x+dx)];
+		atomic_add_and_fetch_float(accum_image + y*stride + x, sum);
+		atomic_add_and_fetch_float(out_image + y*stride + x, sum*image[(y+dy)*stride + (x+dx)]);
 	}
 	else {
-		accum_image[y*w+x] = sum;
+		accum_image[y*stride + x] = sum;
 	}
 }
 
-ccl_device_inline void kernel_filter_nlm_construct_gramian(int fx, int fy,
+ccl_device_inline void kernel_filter_nlm_construct_gramian(int x, int y,
                                                            int dx, int dy,
                                                            const ccl_global float *ccl_restrict difference_image,
                                                            const ccl_global float *ccl_restrict buffer,
@@ -102,30 +159,31 @@ ccl_device_inline void kernel_filter_nlm_construct_gramian(int fx, int fy,
                                                            ccl_global float *XtWX,
                                                            ccl_global float3 *XtWY,
                                                            int4 rect,
-                                                           int4 filter_rect,
-                                                           int w, int h, int f,
+                                                           int4 filter_window,
+                                                           int stride, int f,
                                                            int pass_stride,
                                                            int localIdx)
 {
-	int y = fy + filter_rect.y;
-	int x = fx + filter_rect.x;
 	const int low = max(rect.x, x-f);
 	const int high = min(rect.z, x+f+1);
 	float sum = 0.0f;
 	for(int x1 = low; x1 < high; x1++) {
-		sum += difference_image[y*w+x1];
+		sum += difference_image[y*stride + x1];
 	}
 	float weight = sum * (1.0f/(high - low));
 
-	int storage_ofs = fy*filter_rect.z + fx;
+	/* Reconstruction data is only stored for pixels inside the filter window,
+	 * so compute the pixels's index in there. */
+	int storage_ofs = coord_to_local_index(filter_window, x, y);
 	transform += storage_ofs;
 	rank += storage_ofs;
 	XtWX += storage_ofs;
 	XtWY += storage_ofs;
 
 	kernel_filter_construct_gramian(x, y,
-	                                filter_rect.z*filter_rect.w,
-	                                dx, dy, w, h,
+	                                rect_size(filter_window),
+	                                dx, dy,
+	                                stride,
 	                                pass_stride,
 	                                buffer,
 	                                transform, rank,
@@ -136,9 +194,9 @@ ccl_device_inline void kernel_filter_nlm_construct_gramian(int fx, int fy,
 ccl_device_inline void kernel_filter_nlm_normalize(int x, int y,
                                                    ccl_global float *out_image,
                                                    const ccl_global float *ccl_restrict accum_image,
-                                                   int4 rect, int w)
+                                                   int stride)
 {
-	out_image[y*w+x] /= accum_image[y*w+x];
+	out_image[y*stride + x] /= accum_image[y*stride + x];
 }
 
 CCL_NAMESPACE_END
diff --git a/intern/cycles/kernel/filter/filter_reconstruction.h b/intern/cycles/kernel/filter/filter_reconstruction.h
index 25a3025056c..b7bf322f9ce 100644
--- a/intern/cycles/kernel/filter/filter_reconstruction.h
+++ b/intern/cycles/kernel/filter/filter_reconstruction.h
@@ -19,7 +19,7 @@ CCL_NAMESPACE_BEGIN
 ccl_device_inline void kernel_filter_construct_gramian(int x, int y,
                                                        int storage_stride,
                                                        int dx, int dy,
-                                                       int w, int h,
+                                                       int buffer_stride,
                                                        int pass_stride,
                                                        const ccl_global float *ccl_restrict buffer,
                                                        const ccl_global float *ccl_restrict transform,
@@ -33,8 +33,8 @@ ccl_device_inline void kernel_filter_construct_gramian(int x, int y,
 		return;
 	}
 
-	int p_offset =  y    *w +  x;
-	int q_offset = (y+dy)*w + (x+dx);
+	int p_offset =  y     * buffer_stride +  x;
+	int q_offset = (y+dy) * buffer_stride + (x+dx);
 
 #ifdef __KERNEL_GPU__
 	const int stride = storage_stride;
@@ -65,7 +65,7 @@ ccl_device_inline void kernel_filter_construct_gramian(int x, int y,
 	math_vec3_add_strided(XtWY, (*rank)+1, design_row, weight * q_color, stride);
 }
 
-ccl_device_inline void kernel_filter_finalize(int x, int y, int w, int h,
+ccl_device_inline void kernel_filter_finalize(int x, int y,
                                               ccl_global float *buffer,
                                               ccl_global int *rank,
                                               int storage_stride,
diff --git a/intern/cycles/kernel/kernels/cpu/filter_cpu.h b/intern/cycles/kernel/kernels/cpu/filter_cpu.h
index bf13ba62806..4231aba88d7 100644
--- a/intern/cycles/kernel/kernels/cpu/filter_cpu.h
+++ b/intern/cycles/kernel/kernels/cpu/filter_cpu.h
@@ -74,7 +74,7 @@ void KERNEL_FUNCTION_FULL_NAME(filter_nlm_calc_difference)(int dx,
                                                            float *variance,
                                                            float *difference_image,
                                                            int* rect,
-                                                           int w,
+                                                           int stride,
                                                            int channel_offset,
                                                            float a,
                                                            float k_2);
@@ -82,13 +82,13 @@ void KERNEL_FUNCTION_FULL_NAME(filter_nlm_calc_difference)(int dx,
 void KERNEL_FUNCTION_FULL_NAME(filter_nlm_blur)(float *difference_image,
                                                 float *out_image,
                                                 int* rect,
-                                                int w,
+                                                int stride,
                                                 int f);
 
 void KERNEL_FUNCTION_FULL_NAME(filter_nlm_calc_weight)(float *difference_image,
                                                        float *out_image,
                                                        int* rect,
-                                                       int w,
+                                                       int stride,
                                                        int f);
 
 void KERNEL_FUNCTION_FULL_NAME(filter_nlm_update_output)(int dx,
@@ -98,7 +98,7 @@ void KERNEL_FUNCTION_FULL_NAME(filter_nlm_update_output)(int dx,
                                                          float *out_image,
                                                          float *accum_image,
                                                          int* rect,
-                                                         int w,
+                                                         int stride,
                                                          int f);
 
 void KERNEL_FUNCTION_FULL_NAME(filter_nlm_construct_gramian)(int dx,
@@ -110,22 +110,19 @@ void KERNEL_FUNCTION_FULL_NAME(filter_nlm_construct_gramian)(int dx,
                                                              float *XtWX,
                                                              float3 *XtWY,
                                                              int *rect,
-                                                             int *filter_rect,
-                                                             int w,
-                                                             int h,
+                                                             int *filter_window,
+                                                             int stride,
                                                              int f,
                                                              int pass_stride);
 
 void KERNEL_FUNCTION_FULL_NAME(filter_nlm_normalize)(float *out_image,
                                                      float *accum_image,
                                                      int* rect,
-                                                     int w);
+                                                     int stride);
 
 void KERNEL_FUNCTION_FULL_NAME(filter_finalize)(int x,
                                                 int y,
                                                 int storage_ofs,
-                                                int w,
-                                                int h,
                                                 float *buffer,
                                                 int *rank,
                                                 float *XtWX,
diff --git a/intern/cycles/kernel/kernels/cpu/filter_cpu_impl.h b/intern/cycles/kernel/kernels/cpu/filter_cpu_impl.h
index 2fbb0ea2bdb..ab39260784b 100644
--- a/intern/cycles/kernel/kernels/cpu/filter_cpu_impl.h
+++ b/intern/cycles/kernel/kernels/cpu/filter_cpu_impl.h
@@ -150,7 +150,7 @@ void KERNEL_FUNCTION_FULL_NAME(filter_nlm_calc_difference)(int dx,
                                                            float *variance,
                                                            float *difference_image,
                                                            int *rect,
-                                                           int w,
+                                                           int stride,
                                                            int channel_offset,
                                                            float a,
                                                            float k_2)
@@ -158,33 +158,33 @@ void KERNEL_FUNCTION_FULL_NAME(filter_nlm_calc_difference)(int dx,
 #ifdef KERNEL_STUB
 	STUB_ASSERT(KERNEL_ARCH, filter_nlm_calc_difference);
 #else
-	kernel_filter_nlm_calc_difference(dx, dy, weight_image, variance, difference_image, load_int4(rect), w, channel_offset, a, k_2);
+	kernel_filter_nlm_calc_difference(dx, dy, weight_image, variance, difference_image, load_int4(rect), stride, channel_offset, a, k_2);
 #endif
 }
 
 void KERNEL_FUNCTION_FULL_NAME(filter_nlm_blur)(float *difference_image,
                                                 float *out_image,
                                                 int *rect,
-                                                int w,
+                                                int stride,
                                                 int f)
 {
 #ifdef KERNEL_STUB
 	STUB_ASSERT(KERNEL_ARCH, filter_nlm_blur);
 #else
-	kernel_filter_nlm_blur(difference_image, out_image, load_int4(rect), w, f);
+	kernel_filter_nlm_blur(difference_image, out_image, load_int4(rect), stride, f);
 #endif
 }
 
 void KERNEL_FUNCTION_FULL_NAME(filter_nlm_calc_weight)(float *difference_image,
                                                        float *out_image,
                                                        int *rect,
-                                                       int w,
+                                                       int stride,
                                                        int f)
 {
 #ifdef KERNEL_STUB
 	STUB_ASSERT(KERNEL_ARCH, filter_nlm_calc_weight);
 #else
-	kernel_filter_nlm_calc_weight(difference_image, out_image, load_int4(rect), w, f);
+	kernel_filter_nlm_calc_weight(difference_image, out_image, load_int4(rect), stride, f);
 #endif
 }
 
@@ -195,13 +195,13 @@ void KERNEL_FUNCTION_FULL_NAME(filter_nlm_update_output)(int dx,
                                                          float *out_image,
                                                          float *accum_image,
                                                          int *rect,
-                                                         int w,
+                                                         int stride,
                                                          int f)
 {
 #ifdef KERNEL_STUB
 	STUB_ASSERT(KERNEL_ARCH, filter_nlm_update_output);
 #else
-	kernel_filter_nlm_update_output(dx, dy, difference_image, image, out_image, accum_image, load_int4(rect), w, f);
+	kernel_filter_nlm_update_output(dx, dy, difference_image, image, out_image, accum_image, load_int4(rect), stride, f);
 #endif
 }
 
@@ -214,36 +214,33 @@ void KERNEL_FUNCTION_FULL_NAME(filter_nlm_construct_gramian)(int dx,
                                                              float *XtWX,
                                                              float3 *XtWY,
                                                              int *rect,
-                                                             int *filter_rect,
-                                                             int w,
-                                                             int h,
+                                                             int *filter_window,
+                                                             int stride,
                                                              int f,
                                                              int pass_stride)
 {
 #ifdef KERNEL_STUB
 	STUB_ASSERT(KERNEL_ARCH, filter_nlm_construct_gramian);
 #else
-    kernel_filter_nlm_construct_gramian(dx, dy, difference_image, buffer, transform, rank, XtWX, XtWY, load_int4(rect), load_int4(filter_rect), w, h, f, pass_stride);
+	kernel_filter_nlm_construct_gramian(dx, dy, difference_image, buffer, transform, rank, XtWX, XtWY, load_int4(rect), load_int4(filter_window), stride, f, pass_stride);
 #endif
 }
 
 void KERNEL_FUNCTION_FULL_NAME(filter_nlm_normalize)(float *out_image,
                                                      float *accum_image,
                                                      int *rect,
-                                                     int w)
+                                                     int stride)
 {
 #ifdef KERNEL_STUB
 	STUB_ASSERT(KERNEL_ARCH, filter_nlm_normalize);
 #else
-	kernel_filter_nlm_normalize(out_image, accum_image, load_int4(rect), w);
+	kernel_filter_nlm_normalize(out_image, accum_image, load_int4(rect), stride);
 #endif
 }
 
 void KERNEL_FUNCTION_FULL_NAME(filter_finalize)(int x,
                                                 int y,
                                                 int storage_ofs,
-                                                int w,
-                                                int h,
                                                 float *buffer,
                                                 int *rank,
                                                 float *XtWX,
@@ -257,7 +254,7 @@ void KERNEL_FUNCTION_FULL_NAME(filter_finalize)(int x,
 	XtWX += storage_ofs*XTWX_SIZE;
 	XtWY += storage_ofs*XTWY_SIZE;
 	rank += storage_ofs;
-	kernel_filter_finalize(x, y, w, h, buffer, rank, 1, XtWX, XtWY, load_int4(buffer_params), sample);
+	kernel_filter_finalize(x, y, buffer, rank, 1, XtWX, XtWY, load_int4(buffer_params), sample);
 #endif
 }
 
diff --git a/intern/cycles/kernel/kernels/cuda/filter.cu b/intern/cycles/kernel/kernels/cuda/filter.cu
index c8172355a7f..035f0484488 100644
--- a/intern/cycles/kernel/kernels/cuda/filter.cu
+++ b/intern/cycles/kernel/kernels/cuda/filter.cu
@@ -134,95 +134,140 @@ kernel_cuda_filter_construct_transform(float const* __restrict__ buffer,
 
 extern "C" __global__ void
 CUDA_LAUNCH_BOUNDS(CUDA_THREADS_BLOCK_WIDTH, CUDA_KERNEL_MAX_REGISTERS)
-kernel_cuda_filter_nlm_calc_difference(int dx, int dy,
-                                       const float *ccl_restrict weight_image,
+kernel_cuda_filter_nlm_calc_difference(const float *ccl_restrict weight_image,
                                        const float *ccl_restrict variance_image,
                                        float *difference_image,
-                                       int4 rect, int w,
+                                       int w,
+                                       int h,
+                                       int stride,
+                                       int shift_stride,
+                                       int r,
                                        int channel_offset,
-                                       float a, float k_2)
+                                       float a,
+                                       float k_2)
 {
-	int x = blockDim.x*blockIdx.x + threadIdx.x + rect.x;
-	int y = blockDim.y*blockIdx.y + threadIdx.y + rect.y;
-	if(x < rect.z && y < rect.w) {
-		kernel_filter_nlm_calc_difference(x, y, dx, dy, weight_image, variance_image, difference_image, rect, w, channel_offset, a, k_2);
+	int4 co, rect;
+	int ofs;
+	if(get_nlm_coords(w, h, r, shift_stride, &rect, &co, &ofs)) {
+		kernel_filter_nlm_calc_difference(co.x, co.y, co.z, co.w,
+		                                  weight_image,
+		                                  variance_image,
+		                                  difference_image + ofs,
+		                                  rect, stride,
+		                                  channel_offset, a, k_2);
 	}
 }
 
 extern "C" __global__ void
 CUDA_LAUNCH_BOUNDS(CUDA_THREADS_BLOCK_WIDTH, CUDA_KERNEL_MAX_REGISTERS)
-kernel_cuda_filter_nlm_blur(const float *ccl_restrict difference_image, float *out_image, int4 rect, int w, int f)
+kernel_cuda_filter_nlm_blur(const float *ccl_restrict difference_image,
+                            float *out_image,
+                            int w,
+                            int h,
+                            int stride,
+                            int shift_stride,
+                            int r,
+                            int f)
 {
-	int x = blockDim.x*blockIdx.x + threadIdx.x + rect.x;
-	int y = blockDim.y*blockIdx.y + threadIdx.y + rect.y;
-	if(x < rect.z && y < rect.w) {
-		kernel_filter_nlm_blur(x, y, difference_image, out_image, rect, w, f);
+	int4 co, rect;
+	int ofs;
+	if(get_nlm_coords(w, h, r, shift_stride, &rect, &co, &ofs)) {
+		kernel_filter_nlm_blur(co.x, co.y,
+		                       difference_image + ofs,
+		                       out_image + ofs,
+		                       rect, stride, f);
 	}
 }
 
 extern "C" __global__ void
 CUDA_LAUNCH_BOUNDS(CUDA_THREADS_BLOCK_WIDTH, CUDA_KERNEL_MAX_REGISTERS)
-kernel_cuda_filter_nlm_calc_weight(const float *ccl_restrict difference_image, float *out_image, int4 rect, int w, int f)
+kernel_cuda_filter_nlm_calc_weight(const float *ccl_restrict difference_image,
+                                   float *out_image,
+                                   int w,
+                                   int h,
+                                   int stride,
+                                   int shift_stride,
+                                   int r,
+                                   int f)
 {
-	int x = blockDim.x*blockIdx.x + threadIdx.x + rect.x;
-	int y = blockDim.y*blockIdx.y + threadIdx.y + rect.y;
-	if(x < rect.z && y < rect.w) {
-		kernel_filter_nlm_calc_weight(x, y, difference_image, out_image, rect, w, f);
+	int4 co, rect;
+	int ofs;
+	if(get_nlm_coords(w, h, r, shift_stride, &rect, &co, &ofs)) {
+		kernel_filter_nlm_calc_weight(co.x, co.y,
+		                              difference_image + ofs,
+		                              out_image + ofs,
+		                              rect, stride, f);
 	}
 }
 
 extern "C" __global__ void
 CUDA_LAUNCH_BOUNDS(CUDA_THREADS_BLOCK_WIDTH, CUDA_KERNEL_MAX_REGISTERS)
-kernel_cuda_filter_nlm_update_output(int dx, int dy,
-                                     const float *ccl_restrict difference_image,
+kernel_cuda_filter_nlm_update_output(const float *ccl_restrict difference_image,
                                      const float *ccl_restrict image,
-                                     float *out_image, float *accum_image,
-                                     int4 rect, int w,
+                                     float *out_image,
+                                     float *accum_image,
+                                     int w,
+                                     int h,
+                                     int stride,
+                                     int shift_stride,
+                                     int r,
                                      int f)
 {
-	int x = blockDim.x*blockIdx.x + threadIdx.x + rect.x;
-	int y = blockDim.y*blockIdx.y + threadIdx.y + rect.y;
-	if(x < rect.z && y < rect.w) {
-		kernel_filter_nlm_update_output(x, y, dx, dy, difference_image, image, out_image, accum_image, rect, w, f);
+	int4 co, rect;
+	int ofs;
+	if(get_nlm_coords(w, h, r, shift_stride, &rect, &co, &ofs)) {
+		kernel_filter_nlm_update_output(co.x, co.y, co.z, co.w,
+		                                difference_image + ofs,
+		                                image,
+		                                out_image,
+		                                accum_image,
+		                                rect, stride, f);
 	}
 }
 
 extern "C" __global__ void
 CUDA_LAUNCH_BOUNDS(CUDA_THREADS_BLOCK_WIDTH, CUDA_KERNEL_MAX_REGISTERS)
-kernel_cuda_filter_nlm_normalize(float *out_image, const float *ccl_restrict accum_image, int4 rect, int w)
+kernel_cuda_filter_nlm_normalize(float *out_image,
+                                 const float *ccl_restrict accum_image,
+                                 int w,
+                                 int h,
+                                 int stride)
 {
-	int x = blockDim.x*blockIdx.x + threadIdx.x + rect.x;
-	int y = blockDim.y*blockIdx.y + threadIdx.y + rect.y;
-	if(x < rect.z && y < rect.w) {
-		kernel_filter_nlm_normalize(x, y, out_image, accum_image, rect, w);
+	int x = blockDim.x*blockIdx.x + threadIdx.x;
+	int y = blockDim.y*blockIdx.y + threadIdx.y;
+	if(x < w && y < h) {
+		kernel_filter_nlm_normalize(x, y, out_image, accum_image, stride);
 	}
 }
 
 extern "C" __global__ void
 CUDA_LAUNCH_BOUNDS(CUDA_THREADS_BLOCK_WIDTH, CUDA_KERNEL_MAX_REGISTERS)
-kernel_cuda_filter_nlm_construct_gramian(int dx, int dy,
-                                         const float *ccl_restrict difference_image,
+kernel_cuda_filter_nlm_construct_gramian(const float *ccl_restrict difference_image,
                                          const float *ccl_restrict buffer,
                                          float const* __restrict__ transform,
                                          int *rank,
                                          float *XtWX,
                                          float3 *XtWY,
-                                         int4 rect,
-                                         int4 filter_rect,
-                                         int w, int h, int f,
+                                         int4 filter_window,
+                                         int w,
+                                         int h,
+                                         int stride,
+                                         int shift_stride,
+                                         int r,
+                                         int f,
                                          int pass_stride)
 {
-	int x = blockDim.x*blockIdx.x + threadIdx.x + max(0, rect.x-filter_rect.x);
-	int y = blockDim.y*blockIdx.y + threadIdx.y + max(0, rect.y-filter_rect.y);
-	if(x < min(filter_rect.z, rect.z-filter_rect.x) && y < min(filter_rect.w, rect.w-filter_rect.y)) {
-		kernel_filter_nlm_construct_gramian(x, y,
-		                                    dx, dy,
-		                                    difference_image,
+	int4 co, rect;
+	int ofs;
+	if(get_nlm_coords_window(w, h, r, shift_stride, &rect, &co, &ofs, filter_window)) {
+		kernel_filter_nlm_construct_gramian(co.x, co.y,
+		                                    co.z, co.w,
+		                                    difference_image + ofs,
 		                                    buffer,
 		                                    transform, rank,
 		                                    XtWX, XtWY,
-		                                    rect, filter_rect,
-		                                    w, h, f,
+		                                    rect, filter_window,
+		                                    stride, f,
 		                                    pass_stride,
 		                                    threadIdx.y*blockDim.x + threadIdx.x);
 	}
@@ -230,10 +275,12 @@ kernel_cuda_filter_nlm_construct_gramian(int dx, int dy,
 
 extern "C" __global__ void
 CUDA_LAUNCH_BOUNDS(CUDA_THREADS_BLOCK_WIDTH, CUDA_KERNEL_MAX_REGISTERS)
-kernel_cuda_filter_finalize(int w, int h,
-                            float *buffer, int *rank,
-                            float *XtWX, float3 *XtWY,
-                            int4 filter_area, int4 buffer_params,
+kernel_cuda_filter_finalize(float *buffer,
+                            int *rank,
+                            float *XtWX,
+                            float3 *XtWY,
+                            int4 filter_area,
+                            int4 buffer_params,
                             int sample)
 {
 	int x = blockDim.x*blockIdx.x + threadIdx.x;
@@ -243,7 +290,10 @@ kernel_cuda_filter_finalize(int w, int h,
 		rank += storage_ofs;
 		XtWX += storage_ofs;
 		XtWY += storage_ofs;
-		kernel_filter_finalize(x, y, w, h, buffer, rank, filter_area.z*filter_area.w, XtWX, XtWY, buffer_params, sample);
+		kernel_filter_finalize(x, y, buffer, rank,
+		                       filter_area.z*filter_area.w,
+		                       XtWX, XtWY,
+		                       buffer_params, sample);
 	}
 }
 
diff --git a/intern/cycles/kernel/kernels/opencl/filter.cl b/intern/cycles/kernel/kernels/opencl/filter.cl
index 7a7b596a350..2b77807c38b 100644
--- a/intern/cycles/kernel/kernels/opencl/filter.cl
+++ b/intern/cycles/kernel/kernels/opencl/filter.cl
@@ -126,113 +126,136 @@ __kernel void kernel_ocl_filter_construct_transform(const ccl_global float *ccl_
 	}
 }
 
-__kernel void kernel_ocl_filter_nlm_calc_difference(int dx,
-                                                    int dy,
-                                                    const ccl_global float *ccl_restrict weight_image,
+__kernel void kernel_ocl_filter_nlm_calc_difference(const ccl_global float *ccl_restrict weight_image,
                                                     const ccl_global float *ccl_restrict variance_image,
                                                     ccl_global float *difference_image,
-                                                    int4 rect,
                                                     int w,
+                                                    int h,
+                                                    int stride,
+                                                    int shift_stride,
+                                                    int r,
                                                     int channel_offset,
                                                     float a,
                                                     float k_2)
 {
-	int x = get_global_id(0) + rect.x;
-	int y = get_global_id(1) + rect.y;
-	if(x < rect.z && y < rect.w) {
-		kernel_filter_nlm_calc_difference(x, y, dx, dy, weight_image, variance_image, difference_image, rect, w, channel_offset, a, k_2);
+	int4 co, rect;
+	int ofs;
+	if(get_nlm_coords(w, h, r, shift_stride, &rect, &co, &ofs)) {
+		kernel_filter_nlm_calc_difference(co.x, co.y, co.z, co.w,
+		                                  weight_image,
+		                                  variance_image,
+		                                  difference_image + ofs,
+		                                  rect, stride,
+		                                  channel_offset, a, k_2);
 	}
 }
 
 __kernel void kernel_ocl_filter_nlm_blur(const ccl_global float *ccl_restrict difference_image,
                                          ccl_global float *out_image,
-                                         int4 rect,
                                          int w,
+                                         int h,
+                                         int stride,
+                                         int shift_stride,
+                                         int r,
                                          int f)
 {
-	int x = get_global_id(0) + rect.x;
-	int y = get_global_id(1) + rect.y;
-	if(x < rect.z && y < rect.w) {
-		kernel_filter_nlm_blur(x, y, difference_image, out_image, rect, w, f);
+	int4 co, rect;
+	int ofs;
+	if(get_nlm_coords(w, h, r, shift_stride, &rect, &co, &ofs)) {
+		kernel_filter_nlm_blur(co.x, co.y,
+		                       difference_image + ofs,
+		                       out_image + ofs,
+		                       rect, stride, f);
 	}
 }
 
 __kernel void kernel_ocl_filter_nlm_calc_weight(const ccl_global float *ccl_restrict difference_image,
                                                 ccl_global float *out_image,
-                                                int4 rect,
                                                 int w,
+                                                int h,
+                                                int stride,
+                                                int shift_stride,
+                                                int r,
                                                 int f)
 {
-	int x = get_global_id(0) + rect.x;
-	int y = get_global_id(1) + rect.y;
-	if(x < rect.z && y < rect.w) {
-		kernel_filter_nlm_calc_weight(x, y, difference_image, out_image, rect, w, f);
+	int4 co, rect;
+	int ofs;
+	if(get_nlm_coords(w, h, r, shift_stride, &rect, &co, &ofs)) {
+		kernel_filter_nlm_calc_weight(co.x, co.y,
+		                              difference_image + ofs,
+		                              out_image + ofs,
+		                              rect, stride, f);
 	}
 }
 
-__kernel void kernel_ocl_filter_nlm_update_output(int dx,
-                                                  int dy,
-                                                  const ccl_global float *ccl_restrict difference_image,
+__kernel void kernel_ocl_filter_nlm_update_output(const ccl_global float *ccl_restrict difference_image,
                                                   const ccl_global float *ccl_restrict image,
                                                   ccl_global float *out_image,
                                                   ccl_global float *accum_image,
-                                                  int4 rect,
                                                   int w,
+                                                  int h,
+                                                  int stride,
+                                                  int shift_stride,
+                                                  int r,
                                                   int f)
 {
-	int x = get_global_id(0) + rect.x;
-	int y = get_global_id(1) + rect.y;
-	if(x < rect.z && y < rect.w) {
-		kernel_filter_nlm_update_output(x, y, dx, dy, difference_image, image, out_image, accum_image, rect, w, f);
+	int4 co, rect;
+	int ofs;
+	if(get_nlm_coords(w, h, r, shift_stride, &rect, &co, &ofs)) {
+		kernel_filter_nlm_update_output(co.x, co.y, co.z, co.w,
+		                                difference_image + ofs,
+		                                image,
+		                                out_image,
+		                                accum_image,
+		                                rect, stride, f);
 	}
 }
 
 __kernel void kernel_ocl_filter_nlm_normalize(ccl_global float *out_image,
                                               const ccl_global float *ccl_restrict accum_image,
-                                              int4 rect,
-                                              int w)
+                                              int w,
+                                              int h,
+                                              int stride)
 {
-	int x = get_global_id(0) + rect.x;
-	int y = get_global_id(1) + rect.y;
-	if(x < rect.z && y < rect.w) {
-		kernel_filter_nlm_normalize(x, y, out_image, accum_image, rect, w);
+	int x = get_global_id(0);
+	int y = get_global_id(1);
+	if(x < w && y < h) {
+		kernel_filter_nlm_normalize(x, y, out_image, accum_image, stride);
 	}
 }
 
-__kernel void kernel_ocl_filter_nlm_construct_gramian(int dx,
-                                                      int dy,
-                                                      const ccl_global float *ccl_restrict difference_image,
+__kernel void kernel_ocl_filter_nlm_construct_gramian(const ccl_global float *ccl_restrict difference_image,
                                                       const ccl_global float *ccl_restrict buffer,
                                                       const ccl_global float *ccl_restrict transform,
                                                       ccl_global int *rank,
                                                       ccl_global float *XtWX,
                                                       ccl_global float3 *XtWY,
-                                                      int4 rect,
-                                                      int4 filter_rect,
+                                                      int4 filter_window,
                                                       int w,
                                                       int h,
+                                                      int stride,
+                                                      int shift_stride,
+                                                      int r,
                                                       int f,
                                                       int pass_stride)
 {
-	int x = get_global_id(0) + max(0, rect.x-filter_rect.x);
-	int y = get_global_id(1) + max(0, rect.y-filter_rect.y);
-	if(x < min(filter_rect.z, rect.z-filter_rect.x) && y < min(filter_rect.w, rect.w-filter_rect.y)) {
-		kernel_filter_nlm_construct_gramian(x, y,
-		                                    dx, dy,
-		                                    difference_image,
+	int4 co, rect;
+	int ofs;
+	if(get_nlm_coords_window(w, h, r, shift_stride, &rect, &co, &ofs, filter_window)) {
+		kernel_filter_nlm_construct_gramian(co.x, co.y,
+		                                    co.z, co.w,
+		                                    difference_image + ofs,
 		                                    buffer,
 		                                    transform, rank,
 		                                    XtWX, XtWY,
-		                                    rect, filter_rect,
-		                                    w, h, f,
+		                                    rect, filter_window,
+		                                    stride, f,
 		                                    pass_stride,
 		                                    get_local_id(1)*get_local_size(0) + get_local_id(0));
 	}
 }
 
-__kernel void kernel_ocl_filter_finalize(int w,
-                                         int h,
-                                         ccl_global float *buffer,
+__kernel void kernel_ocl_filter_finalize(ccl_global float *buffer,
                                          ccl_global int *rank,
                                          ccl_global float *XtWX,
                                          ccl_global float3 *XtWY,
@@ -247,7 +270,10 @@ __kernel void kernel_ocl_filter_finalize(int w,
 		rank += storage_ofs;
 		XtWX += storage_ofs;
 		XtWY += storage_ofs;
-		kernel_filter_finalize(x, y, w, h, buffer, rank, filter_area.z*filter_area.w, XtWX, XtWY, buffer_params, sample);
+		kernel_filter_finalize(x, y, buffer, rank,
+		                       filter_area.z*filter_area.w,
+		                       XtWX, XtWY,
+		                       buffer_params, sample);
 	}
 }