From c8e6e940e0f0f2a0937ba8d62726e1e178ff3cd9 Mon Sep 17 00:00:00 2001
From: Jonathan L Long <jonlong@cs.berkeley.edu>
Date: Thu, 3 Apr 2014 19:19:47 -0700
Subject: [PATCH] add set_cpu_data to Blob and SyncedMemory

This allows a blob to be updated without copy to use already existing
memory (and will support MemoryDataLayer).
---
 include/caffe/blob.hpp      |  1 +
 include/caffe/syncedmem.hpp |  8 ++++++--
 src/caffe/blob.cpp          |  6 ++++++
 src/caffe/syncedmem.cpp     | 14 +++++++++++++-
 4 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/include/caffe/blob.hpp b/include/caffe/blob.hpp
index 2f6b8f8..7510146 100644
--- a/include/caffe/blob.hpp
+++ b/include/caffe/blob.hpp
@@ -63,6 +63,7 @@ class Blob {
   }
 
   const Dtype* cpu_data() const;
+  void set_cpu_data(Dtype* data);
   const Dtype* gpu_data() const;
   const Dtype* cpu_diff() const;
   const Dtype* gpu_diff() const;
diff --git a/include/caffe/syncedmem.hpp b/include/caffe/syncedmem.hpp
index e265c45..600feb3 100644
--- a/include/caffe/syncedmem.hpp
+++ b/include/caffe/syncedmem.hpp
@@ -35,11 +35,14 @@ inline void CaffeFreeHost(void* ptr) {
 class SyncedMemory {
  public:
   SyncedMemory()
-      : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(0), head_(UNINITIALIZED) {}
+      : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(0), head_(UNINITIALIZED),
+        own_cpu_data_(false) {}
   explicit SyncedMemory(size_t size)
-      : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(size), head_(UNINITIALIZED) {}
+      : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(size), head_(UNINITIALIZED),
+        own_cpu_data_(false) {}
   ~SyncedMemory();
   const void* cpu_data();
+  void set_cpu_data(void* data);
   const void* gpu_data();
   void* mutable_cpu_data();
   void* mutable_gpu_data();
@@ -53,6 +56,7 @@ class SyncedMemory {
   void* gpu_ptr_;
   size_t size_;
   SyncedHead head_;
+  bool own_cpu_data_;
 
   DISABLE_COPY_AND_ASSIGN(SyncedMemory);
 };  // class SyncedMemory
diff --git a/src/caffe/blob.cpp b/src/caffe/blob.cpp
index f1fe98d..444e9cf 100644
--- a/src/caffe/blob.cpp
+++ b/src/caffe/blob.cpp
@@ -49,6 +49,12 @@ const Dtype* Blob<Dtype>::cpu_data() const {
 }
 
 template <typename Dtype>
+void Blob<Dtype>::set_cpu_data(Dtype* data) {
+  CHECK(data);
+  data_->set_cpu_data(data);
+}
+
+template <typename Dtype>
 const Dtype* Blob<Dtype>::gpu_data() const {
   CHECK(data_);
   return (const Dtype*)data_->gpu_data();
diff --git a/src/caffe/syncedmem.cpp b/src/caffe/syncedmem.cpp
index c33f3e6..fec37d6 100644
--- a/src/caffe/syncedmem.cpp
+++ b/src/caffe/syncedmem.cpp
@@ -10,7 +10,7 @@
 namespace caffe {
 
 SyncedMemory::~SyncedMemory() {
-  if (cpu_ptr_) {
+  if (cpu_ptr_ && own_cpu_data_) {
     CaffeFreeHost(cpu_ptr_);
   }
 
@@ -25,10 +25,12 @@ inline void SyncedMemory::to_cpu() {
     CaffeMallocHost(&cpu_ptr_, size_);
     memset(cpu_ptr_, 0, size_);
     head_ = HEAD_AT_CPU;
+    own_cpu_data_ = true;
     break;
   case HEAD_AT_GPU:
     if (cpu_ptr_ == NULL) {
       CaffeMallocHost(&cpu_ptr_, size_);
+      own_cpu_data_ = true;
     }
     CUDA_CHECK(cudaMemcpy(cpu_ptr_, gpu_ptr_, size_, cudaMemcpyDeviceToHost));
     head_ = SYNCED;
@@ -64,6 +66,16 @@ const void* SyncedMemory::cpu_data() {
   return (const void*)cpu_ptr_;
 }
 
+void SyncedMemory::set_cpu_data(void* data) {
+  CHECK(data);
+  if (own_cpu_data_) {
+    CaffeFreeHost(cpu_ptr_);
+  }
+  cpu_ptr_ = data;
+  head_ = HEAD_AT_CPU;
+  own_cpu_data_ = false;
+}
+
 const void* SyncedMemory::gpu_data() {
   to_gpu();
   return (const void*)gpu_ptr_;
-- 
2.7.4