caffe SyncedMemory

用途

用于GPU和CPU之间的数据传递；

源码分析

#ifndef CAFFE_SYNCEDMEM_HPP_
#define CAFFE_SYNCEDMEM_HPP_

#include <cstdlib>

#ifdef USE_MKL
  #include "mkl.h"
#endif

#include "caffe/common.hpp"

namespace caffe {

inline void CaffeMallocHost(void** ptr, size_t size, bool* use_cuda) {
#ifndef CPU_ONLY
  if (Caffe::mode() == Caffe::GPU) {
    CUDA_CHECK(cudaMallocHost(ptr, size));//cudaMallocHost分配的内存是固定的，不用进行页交换；
    *use_cuda = true;
    return;
  }
#endif
#ifdef USE_MKL
  *ptr = mkl_malloc(size ? size:1, 64);
#else
  *ptr = malloc(size);
#endif
  *use_cuda = false;
  CHECK(*ptr) << "host allocation of size " << size << " failed";
}

inline void CaffeFreeHost(void* ptr, bool use_cuda) {
#ifndef CPU_ONLY
  if (use_cuda) {
    CUDA_CHECK(cudaFreeHost(ptr));
    return;
  }
#endif
#ifdef USE_MKL
  mkl_free(ptr);
#else
  free(ptr);
#endif
}

class SyncedMemory {
 public:
  SyncedMemory();
  explicit SyncedMemory(size_t size);
  ~SyncedMemory();
  const void* cpu_data();//先用to_cpu()更新cpu内的数据，然后返回指向cpu的指针；
  void set_cpu_data(void* data);//如果之前cpu里有数据，释放之前的数据，将cpu_ptr_指针指向data
  const void* gpu_data();
  void set_gpu_data(void* data);
  void* mutable_cpu_data();//先用to_cpu()更新cpu内的数据，设置head=  head_ = HEAD_AT_CPU;然后返回指向cpu的指针
  void* mutable_gpu_data();
  enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };//返回数据的状态分别为：未初始化；数据在CPU内是最新的；数据在GPU内是最新的；数据已同步
  SyncedHead head() { return head_; }
  size_t size() { return size_; }

#ifndef CPU_ONLY
  void async_gpu_push(const cudaStream_t& stream);//如果数据在CPU内是最新的，将CPU内的数据同步到GPU
#endif

 private:
  void check_device();

  void to_cpu();//更新cpu数据
  void to_gpu();//更新gpu数据
  void* cpu_ptr_;//指向cpu的内存指针
  void* gpu_ptr_;
  size_t size_;
  SyncedHead head_;
  bool own_cpu_data_;
  bool cpu_malloc_use_cuda_;
  bool own_gpu_data_;
  int device_;

  DISABLE_COPY_AND_ASSIGN(SyncedMemory);
};  // class SyncedMemory

}  // namespace caffe

#endif  // CAFFE_SYNCEDMEM_HPP_