1 files changed, 13 insertions, 31 deletions

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();
 			}
 		}
 	}