Cycles Denoising: Correctly handle target buffer in tile unmapping and move device swap logic to the device_memory

4 files changed, 41 insertions, 31 deletions

diff --git a/intern/cycles/device/device_memory.cpp b/intern/cycles/device/device_memory.cpp
index c6248fcf88b..482af685011 100644
--- a/intern/cycles/device/device_memory.cpp
+++ b/intern/cycles/device/device_memory.cpp
@@ -104,6 +104,26 @@ void device_memory::device_zero()
 	}
 }
 
+void device_memory::swap_device(Device *new_device,
+                                size_t new_device_size,
+                                device_ptr new_device_ptr)
+{
+	original_device = device;
+	original_device_size = device_size;
+	original_device_ptr = device_pointer;
+
+	device = new_device;
+	device_size = new_device_size;
+	device_pointer = new_device_ptr;
+}
+
+void device_memory::restore_device()
+{
+	device = original_device;
+	device_size = original_device_size;
+	device_pointer = original_device_ptr;
+}
+
 /* Device Sub Ptr */
 
 device_sub_ptr::device_sub_ptr(device_memory& mem, int offset, int size)
diff --git a/intern/cycles/device/device_memory.h b/intern/cycles/device/device_memory.h
index d8fe41e78bb..1138964f18c 100644
--- a/intern/cycles/device/device_memory.h
+++ b/intern/cycles/device/device_memory.h
@@ -200,6 +200,9 @@ public:
 
 	virtual ~device_memory();
 
+	void swap_device(Device *new_device, size_t new_device_size, device_ptr new_device_ptr);
+	void restore_device();
+
 protected:
 	friend class CUDADevice;
 
@@ -222,6 +225,10 @@ protected:
 	void device_copy_to();
 	void device_copy_from(int y, int w, int h, int elem);
 	void device_zero();
+
+	device_ptr original_device_ptr;
+	size_t original_device_size;
+	Device *original_device;
 };
 
 /* Device Only Memory
diff --git a/intern/cycles/device/device_multi.cpp b/intern/cycles/device/device_multi.cpp
index 3a4c08b6eb2..1353797edb4 100644
--- a/intern/cycles/device/device_multi.cpp
+++ b/intern/cycles/device/device_multi.cpp
@@ -282,26 +282,27 @@ public:
 					mem.copy_from_device(0, mem.data_size, 1);
 				}
 
-				Device *original_device = mem.device;
-				device_ptr original_ptr = mem.device_pointer;
-				size_t original_size = mem.device_size;
-
-				mem.device = sub_device;
-				mem.device_pointer = 0;
-				mem.device_size = 0;
+				mem.swap_device(sub_device, 0, 0);
 
 				mem.copy_to_device();
 				tiles[i].buffer = mem.device_pointer;
+				tiles[i].device_size = mem.device_size;
 
-				mem.device = original_device;
-				mem.device_pointer = original_ptr;
-				mem.device_size = original_size;
+				mem.restore_device();
 			}
 		}
 	}
 
 	void unmap_neighbor_tiles(Device * sub_device, RenderTile * tiles)
 	{
+		/* Copy denoised result back to the host. */
+		device_vector<float> &mem = tiles[9].buffers->buffer;
+		mem.swap_device(sub_device, tiles[9].device_size, tiles[9].buffer);
+		mem.copy_from_device(0, mem.data_size, 1);
+		mem.restore_device();
+		/* Copy denoised result to the original device. */
+		mem.copy_to_device();
+
 		for(int i = 0; i < 9; i++) {
 			if(!tiles[i].buffers) {
 				continue;
@@ -309,28 +310,9 @@ public:
 
 			device_vector<float> &mem = tiles[i].buffers->buffer;
 			if(mem.device != sub_device) {
-				Device *original_device = mem.device;
-				device_ptr original_ptr = mem.device_pointer;
-				size_t original_size = mem.device_size;
-
-				mem.device = sub_device;
-				mem.device_pointer = tiles[i].buffer;
-
-				/* Copy denoised tile to the host. */
-				if(i == 4) {
-					mem.copy_from_device(0, mem.data_size, 1);
-				}
-
+				mem.swap_device(sub_device, tiles[i].device_size, tiles[i].buffer);
 				sub_device->mem_free(mem);
-
-				mem.device = original_device;
-				mem.device_pointer = original_ptr;
-				mem.device_size = original_size;
-
-				/* Copy denoised tile to the original device. */
-				if(i == 4) {
-					mem.copy_to_device();
-				}
+				mem.restore_device();
 			}
 		}
 	}
diff --git a/intern/cycles/render/buffers.h b/intern/cycles/render/buffers.h
index dfc98fe2061..ee7c0490dfd 100644
--- a/intern/cycles/render/buffers.h
+++ b/intern/cycles/render/buffers.h
@@ -137,6 +137,7 @@ public:
 	int tile_index;
 
 	device_ptr buffer;
+	int device_size;
 
 	RenderBuffers *buffers;