src/caffe/syncedmem.hpp

   1 // Copyright 2013 Yangqing Jia
   2
   3 #ifndef CAFFE_SYNCEDMEM_HPP_
   4 #define CAFFE_SYNCEDMEM_HPP_
   5
   6 namespace caffe {
   7
   8 class SyncedMemory {
   9  public:
  10   SyncedMemory()
  11       : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(0), head_(UNINITIALIZED) {}
  12   explicit SyncedMemory(size_t size)
  13       : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(size), head_(UNINITIALIZED) {}
  14   ~SyncedMemory();
  15   const void* cpu_data();
  16   const void* gpu_data();
  17   void* mutable_cpu_data();
  18   void* mutable_gpu_data();
  19   enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };
  20   SyncedHead head() { return head_; }
  21   size_t size() { return size_; }
  22  private:
  23   void to_cpu();
  24   void to_gpu();
  25   void* cpu_ptr_;
  26   void* gpu_ptr_;
  27   size_t size_;
  28   SyncedHead head_;
  29
  30   DISABLE_COPY_AND_ASSIGN(SyncedMemory);
  31 };  // class SyncedMemory
  32
  33 }  // namespace caffe
  34
  35 #endif  // CAFFE_SYNCEDMEM_HPP_