5 files changed, 30 insertions, 15 deletions

diff --git a/intern/cycles/device/device_cpu.cpp b/intern/cycles/device/device_cpu.cpp
index fd5ae1d7828..4623764d210 100644
--- a/intern/cycles/device/device_cpu.cpp
+++ b/intern/cycles/device/device_cpu.cpp
@@ -73,8 +73,8 @@ public:
 	void mem_alloc(device_memory& mem, MemoryType type)
 	{
 		mem.device_pointer = mem.data_pointer;
-
-		stats.mem_alloc(mem.memory_size());
+		mem.device_size = mem.memory_size();
+		stats.mem_alloc(mem.device_size);
 	}
 
 	void mem_copy_to(device_memory& mem)
@@ -94,9 +94,11 @@ public:
 
 	void mem_free(device_memory& mem)
 	{
-		mem.device_pointer = 0;
-
-		stats.mem_free(mem.memory_size());
+		if(mem.device_pointer) {
+			mem.device_pointer = 0;
+			stats.mem_free(mem.device_size);
+			mem.device_size = 0;
+		}
 	}
 
 	void const_copy_to(const char *name, void *host, size_t size)
@@ -108,15 +110,17 @@ public:
 	{
 		kernel_tex_copy(&kernel_globals, name, mem.data_pointer, mem.data_width, mem.data_height, mem.data_depth, interpolation);
 		mem.device_pointer = mem.data_pointer;
-
-		stats.mem_alloc(mem.memory_size());
+		mem.device_size = mem.memory_size();
+		stats.mem_alloc(mem.device_size);
 	}
 
 	void tex_free(device_memory& mem)
 	{
-		mem.device_pointer = 0;
-
-		stats.mem_free(mem.memory_size());
+		if(mem.device_pointer) {
+			mem.device_pointer = 0;
+			stats.mem_free(mem.device_size);
+			mem.device_size = 0;
+		}
 	}
 
 	void *osl_memory()
diff --git a/intern/cycles/device/device_cuda.cpp b/intern/cycles/device/device_cuda.cpp
index 1ed26717f4b..5de2efab8be 100644
--- a/intern/cycles/device/device_cuda.cpp
+++ b/intern/cycles/device/device_cuda.cpp
@@ -334,6 +334,7 @@ public:
 		size_t size = mem.memory_size();
 		cuda_assert(cuMemAlloc(&device_pointer, size));
 		mem.device_pointer = (device_ptr)device_pointer;
+		mem.device_size = size;
 		stats.mem_alloc(size);
 		cuda_pop_context();
 	}
@@ -381,7 +382,8 @@ public:
 
 			mem.device_pointer = 0;
 
-			stats.mem_free(mem.memory_size());
+			stats.mem_free(mem.device_size);
+			mem.device_size = 0;
 		}
 	}
 
@@ -473,6 +475,7 @@ public:
 				cuda_assert(cuTexRefSetFlags(texref, CU_TRSF_NORMALIZED_COORDINATES));
 
 				mem.device_pointer = (device_ptr)handle;
+				mem.device_size = size;
 
 				stats.mem_alloc(size);
 			}
@@ -540,7 +543,8 @@ public:
 				tex_interp_map.erase(tex_interp_map.find(mem.device_pointer));
 				mem.device_pointer = 0;
 
-				stats.mem_free(mem.memory_size());
+				stats.mem_free(mem.device_size);
+				mem.device_size = 0;
 			}
 			else {
 				tex_interp_map.erase(tex_interp_map.find(mem.device_pointer));
@@ -790,7 +794,8 @@ public:
 				mem.device_pointer = pmem.cuTexId;
 				pixel_mem_map[mem.device_pointer] = pmem;
 
-				stats.mem_alloc(mem.memory_size());
+				mem.device_size = mem.memory_size();
+				stats.mem_alloc(mem.device_size);
 
 				return;
 			}
@@ -847,7 +852,8 @@ public:
 				pixel_mem_map.erase(pixel_mem_map.find(mem.device_pointer));
 				mem.device_pointer = 0;
 
-				stats.mem_free(mem.memory_size());
+				stats.mem_free(mem.device_size);
+				mem.device_size = 0;
 
 				return;
 			}
diff --git a/intern/cycles/device/device_memory.h b/intern/cycles/device/device_memory.h
index 8d6f4a49a9c..8eee6a2c79e 100644
--- a/intern/cycles/device/device_memory.h
+++ b/intern/cycles/device/device_memory.h
@@ -167,6 +167,7 @@ public:
 	int data_elements;
 	device_ptr data_pointer;
 	size_t data_size;
+	size_t device_size;
 	size_t data_width;
 	size_t data_height;
 	size_t data_depth;
@@ -194,6 +195,7 @@ public:
 		data_elements = device_type_traits<T>::num_elements;
 		data_pointer = 0;
 		data_size = 0;
+		device_size = 0;
 		data_width = 0;
 		data_height = 0;
 		data_depth = 0;
diff --git a/intern/cycles/device/device_opencl.cpp b/intern/cycles/device/device_opencl.cpp
index 82419cd62b1..d950d084cd4 100644
--- a/intern/cycles/device/device_opencl.cpp
+++ b/intern/cycles/device/device_opencl.cpp
@@ -794,6 +794,7 @@ public:
 		opencl_assert_err(ciErr, "clCreateBuffer");
 
 		stats.mem_alloc(size);
+		mem.device_size = size;
 	}
 
 	void mem_copy_to(device_memory& mem)
@@ -825,7 +826,8 @@ public:
 			opencl_assert(clReleaseMemObject(CL_MEM_PTR(mem.device_pointer)));
 			mem.device_pointer = 0;
 
-			stats.mem_free(mem.memory_size());
+			stats.mem_free(mem.device_size);
+			mem.device_size = 0;
 		}
 	}
 
diff --git a/intern/cycles/util/util_stats.h b/intern/cycles/util/util_stats.h
index 62b1f1760d7..8758b823084 100644
--- a/intern/cycles/util/util_stats.h
+++ b/intern/cycles/util/util_stats.h
@@ -30,6 +30,7 @@ public:
 	}
 
 	void mem_free(size_t size) {
+		assert(mem_used >= size);
 		mem_used -= size;
 	}