Caffe源码SyncedMemory类

SyncedMemory类简介

最近在阅读caffe源码，代码来自BVLC/caffe，基本是参照网络上比较推荐的 Blob-->Layer-->Net-->Solver 的顺序来分析。其中SyncedMemory类是caffe中底层的结构，负责操作（申请、拷贝等）内存或显存中的数据。

syncedmem.cpp源码

SyncedMemory::SyncedMemory()    //构造函数,初始化内部的变量,size为0,指针为空等
  : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(0), head_(UNINITIALIZED),
    own_cpu_data_(false), cpu_malloc_use_cuda_(false), own_gpu_data_(false) {
#ifndef CPU_ONLY
#ifdef DEBUG
  CUDA_CHECK(cudaGetDevice(&device_));  //cudaGetDevice()函数会返回当前被使用的设备
#endif
#endif
}

SyncedMemory::SyncedMemory(size_t size)   //构造函数,设置size_的值(不会分配内存)
  : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(size), head_(UNINITIALIZED),
    own_cpu_data_(false), cpu_malloc_use_cuda_(false), own_gpu_data_(false) {
#ifndef CPU_ONLY
#ifdef DEBUG
  CUDA_CHECK(cudaGetDevice(&device_));
#endif
#endif
}

SyncedMemory::~SyncedMemory() {       //析构函数
  check_device();                     //检查gpu设备
  if (cpu_ptr_ && own_cpu_data_) {    //如果cpu数据的指针不为空并且数据为自身创建的
    CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);    //释放数据
  }

#ifndef CPU_ONLY
  if (gpu_ptr_ && own_gpu_data_) {    //同理,gpu数据指针不为空并且数据为自身创建的
    CUDA_CHECK(cudaFree(gpu_ptr_));   //释放
  }
#endif  // CPU_ONLY
}

//将数据转移到cpu中.如果还未创建内存则申请对应大小的内存,
//如果数据只在gpu中则将数据拷至cpu中,如果cpu中已存在则不处理
inline void SyncedMemory::to_cpu() {
  check_device();
  switch (head_) {          //当前数据的状态
  case UNINITIALIZED:       //未分配状态
    CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_);   //申请内存
    caffe_memset(size_, 0, cpu_ptr_);   //数据全部设置为0
    head_ = HEAD_AT_CPU;    //设置状态为数据位于内存中,由cpu处理
    own_cpu_data_ = true;   //数据由自身申请创建
    break;
  case HEAD_AT_GPU:         //当前数据位于gpu中
#ifndef CPU_ONLY
    if (cpu_ptr_ == NULL) {
      CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_);   //如果cpu数据指针为空,则申请内存
      own_cpu_data_ = true;
    }
    caffe_gpu_memcpy(size_, gpu_ptr_, cpu_ptr_);    //将gpu_ptr_中的数据复制到cpu_ptr_中,复制size_大小
    head_ = SYNCED;         //设置状态为已同步(cpu数据与gpu数据拥有相同的数据)
#else
    NO_GPU;     //数据在gpu中但是不支持gpu,错误
#endif
    break;
  case HEAD_AT_CPU:   //数据已经在cpu中,不进行处理
  case SYNCED:
    break;
  }
}

//同理,将数据转移到gpu中
inline void SyncedMemory::to_gpu() {
  check_device();
#ifndef CPU_ONLY
  switch (head_) {
  case UNINITIALIZED:   //未初始化
    CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));   //申请显存
    caffe_gpu_memset(size_, 0, gpu_ptr_);       //置为0
    head_ = HEAD_AT_GPU;    //设置状态为gpu
    own_gpu_data_ = true;
    break;
  case HEAD_AT_CPU:     //数据位于cpu中
    if (gpu_ptr_ == NULL) {
      CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));   //申请显存
      own_gpu_data_ = true;
    }
    caffe_gpu_memcpy(size_, cpu_ptr_, gpu_ptr_);  //将数据从cpu_ptr_拷至gpu_ptr_
    head_ = SYNCED;
    break;
  case HEAD_AT_GPU:
  case SYNCED:
    break;
  }
#else
  NO_GPU;
#endif
}

//返回cpu数据的指针,指向的数据不可修改
const void* SyncedMemory::cpu_data() {
  check_device();   //检查设备是否出错
  to_cpu();         //数据转移至cpu中
  return (const void*)cpu_ptr_;
}

//将cpu的数据指针设置为data
void SyncedMemory::set_cpu_data(void* data) {
  check_device();   //检查
  CHECK(data);      //非空检查
  if (own_cpu_data_) {    //自身已经创建了cpu数据,先释放
    CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);
  }
  cpu_ptr_ = data;        //指向data
  head_ = HEAD_AT_CPU;    //修改状态
  own_cpu_data_ = false;  //数据并非自身申请创建的,在调用析构函数时,并不会释放cpu_ptr_指向的内存
}

//返回gpu数据的指针,指向的数据不可修改
const void* SyncedMemory::gpu_data() {
  check_device();
#ifndef CPU_ONLY
  to_gpu();   //转移到gpu中
  return (const void*)gpu_ptr_;
#else
  NO_GPU;
  return NULL;
#endif
}

//设置gpu数据的指针
void SyncedMemory::set_gpu_data(void* data) {
  check_device();
#ifndef CPU_ONLY
  CHECK(data);
  if (own_gpu_data_) {    //自身创建的gpu数据,先释放
    CUDA_CHECK(cudaFree(gpu_ptr_));
  }
  gpu_ptr_ = data;
  head_ = HEAD_AT_GPU;
  own_gpu_data_ = false;    //同样设置为false
#else
  NO_GPU;
#endif
}

//返回cpu上的数据指针,指向的数据可修改
void* SyncedMemory::mutable_cpu_data() {
  check_device();
  to_cpu();
  head_ = HEAD_AT_CPU;
  return cpu_ptr_;
}

//返回gpu上的数据指针,指向的数据可修改
void* SyncedMemory::mutable_gpu_data() {
  check_device();
#ifndef CPU_ONLY
  to_gpu();
  head_ = HEAD_AT_GPU;
  return gpu_ptr_;
#else
  NO_GPU;
  return NULL;
#endif
}

//从cpu中来拷贝数据至gpu,异步拷贝
#ifndef CPU_ONLY
void SyncedMemory::async_gpu_push(const cudaStream_t& stream) {
  check_device();
  CHECK(head_ == HEAD_AT_CPU);    //当前数据应在cpu中
  if (gpu_ptr_ == NULL) {
    CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));   //申请gpu显存
    own_gpu_data_ = true;
  }
  const cudaMemcpyKind put = cudaMemcpyHostToDevice;    //设置拷贝方向,Host To Device
  //Copies data between host and device.异步操作,可能在数据拷贝完成之前函数便返回
  //cudaMemcpy()为同步的,数据拷贝完后函数才会返回
  CUDA_CHECK(cudaMemcpyAsync(gpu_ptr_, cpu_ptr_, size_, put, stream));  //将cpu_ptr_数据拷贝至gpu_ptr_中
  // Assume caller will synchronize on the stream before use
  head_ = SYNCED;   //共享
}
#endif

void SyncedMemory::check_device() {   //检查设备,判断是否出错
#ifndef CPU_ONLY
#ifdef DEBUG
  int device;
  cudaGetDevice(&device);   //返回当前被使用的设备
  CHECK(device == device_);
  if (gpu_ptr_ && own_gpu_data_) {
    cudaPointerAttributes attributes;
    CUDA_CHECK(cudaPointerGetAttributes(&attributes, gpu_ptr_));  //返回gpu_ptr_指针的属性到attributes中
    CHECK(attributes.device == device_);    //检查指针所在的设备与类中保存的设备device_是否一致
  }
#endif
#endif
}

syncedmem.hpp

// If CUDA is available and in GPU mode, host memory will be allocated pinned,
// using cudaMallocHost. It avoids dynamic pinning for transfers (DMA).
// The improvement in performance seems negligible in the single GPU case,
// but might be more significant for parallel training. Most importantly,
// it improved stability for large models on many GPUs.
//申请内存
inline void CaffeMallocHost(void** ptr, size_t size, bool* use_cuda) {
#ifndef CPU_ONLY
  if (Caffe::mode() == Caffe::GPU) {        //gpu模式下
    CUDA_CHECK(cudaMallocHost(ptr, size));  //分配锁页内存
    *use_cuda = true;   //使用了cuda
    return;
  }
#endif
#ifdef USE_MKL          //使用了Intel的Math Kernel Library库
  *ptr = mkl_malloc(size ? size:1, 64);
#else
  *ptr = malloc(size);  //朴实无华的内存创建(分页内存)
#endif
  *use_cuda = false;    //未使用cuda
  CHECK(*ptr) << "host allocation of size " << size << " failed";
}

//释放内存
inline void CaffeFreeHost(void* ptr, bool use_cuda) {
#ifndef CPU_ONLY
  if (use_cuda) {   //使用了cuda,则使用cuda函数释放对应的内存
    CUDA_CHECK(cudaFreeHost(ptr));
    return;
  }
#endif
#ifdef USE_MKL
  mkl_free(ptr);
#else
  free(ptr);
#endif
}

/**
 * @brief Manages memory allocation and synchronization between the host (CPU)
 *        and device (GPU).
 *
 * TODO(dox): more thorough description.
 */
class SyncedMemory {
 public:
  SyncedMemory();
  explicit SyncedMemory(size_t size);
  ~SyncedMemory();
  const void* cpu_data();
  void set_cpu_data(void* data);
  const void* gpu_data();
  void set_gpu_data(void* data);
  void* mutable_cpu_data();
  void* mutable_gpu_data();
  //数据的几种状态,UNINITIALIZED(未初始化,内存或显存还未申请), HEAD_AT_CPU(数据在cpu中),
  //HEAD_AT_GPU(数据在gpu中), SYNCED(数据在cpu和gpu中都存在,并且内容相同)
  enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };
  SyncedHead head() const { return head_; }
  size_t size() const { return size_; }

#ifndef CPU_ONLY
  void async_gpu_push(const cudaStream_t& stream);
#endif

 private:
  void check_device();

  void to_cpu();        //数据转移到cpu中
  void to_gpu();        //数据转移到gpu中
  void* cpu_ptr_;       //cpu中的数据指针
  void* gpu_ptr_;       //gpu中的数据指针
  size_t size_;         //数据的大小
  SyncedHead head_;     //数据的状态,共SyncedHead中指示的四种
  //cpu中的数据是否有自身创建,还是外部传入的指针?(自身创建自己负责释放,外部传的指针析构时不会释放,由外部决定)
  bool own_cpu_data_;
  bool cpu_malloc_use_cuda_;    //申请cpu数据时是否使用了cuda
  bool own_gpu_data_;   //同理,gpu中的数据是否由自身创建
  int device_;          //当前使用的gpu设备

  DISABLE_COPY_AND_ASSIGN(SyncedMemory);  //禁止类的拷贝或者赋值操作
};  // class SyncedMemory

小结

1. cpu处理的数据对应内存数据，gpu处理的数据对应显存数据
2. 单纯创建SyncedMemory类的实例时并不会分配内存或显存，只有在实际需要访问数据的时候(如cpu_data()/mutable_gpu_data()等)时，才会在内部的to_cpu()或to_gpu()函数中分配对应的内存或显存
3. CaffeMallocHost()函数中使用cudaMallocHost()分配的锁页内存，这种内存可被gpu设备直接访问，读写速度比普通的分页内存(malloc申请)要快。关于CUDA的各种函数可参考官方提供的手册。

Caffe的源码笔者是第一次阅读，一边阅读一边记录，对代码的理解和分析可能会存在错误或遗漏，希望各位读者批评指正，谢谢支持！

参考

https://docs.nvidia.com/pdf/CUDA_Runtime_API.pdf
https://www.zhihu.com/question/27982282