add set_cpu_data to Blob and SyncedMemory
authorJonathan L Long <jonlong@cs.berkeley.edu>
Fri, 4 Apr 2014 02:19:47 +0000 (19:19 -0700)
committerJonathan L Long <jonlong@cs.berkeley.edu>
Fri, 2 May 2014 20:24:08 +0000 (13:24 -0700)
This allows a blob to be updated without copy to use already existing
memory (and will support MemoryDataLayer).

include/caffe/blob.hpp
include/caffe/syncedmem.hpp
src/caffe/blob.cpp
src/caffe/syncedmem.cpp

index 2f6b8f8..7510146 100644 (file)
@@ -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;
index e265c45..600feb3 100644 (file)
@@ -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
index f1fe98d..444e9cf 100644 (file)
@@ -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();
index c33f3e6..fec37d6 100644 (file)
@@ -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_;