Caffe 代码学习 3: SyncedMemory

SyncedMemory类定义在syncedmem.hpp/cpp里, 负责caffe底层的内存管理.

内存分配与释放

内存分配与释放由两个(不属于SyncedMemory类)的内联函数完成. 代码简单直观: 如果是CPU模式, 那么调用malloc和free来申请/释放内存, 否则调用CUDA的cudaMallocHost和cudaFreeHost来申请/释放显存.

// ------ 分配内存 ------ 
inline void CaffeMallocHost(void** ptr, size_t size, bool* use_cuda) {
#ifndef CPU_ONLY
  if (Caffe::mode() == Caffe::GPU) {
    CUDA_CHECK(cudaMallocHost(ptr, size));
    *use_cuda = true;
    return;
  }
#endif
  *ptr = malloc(size);
  *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
  free(ptr);
}

类成员变量

  void* cpu_ptr_;  // cpu 内存地址
  void* gpu_ptr_;  // gpu 内存地址
  size_t size_;  // 数据大小
  SyncedHead head_;  // 当前数据同步状态
  bool own_cpu_data_;  //  是否是自己的cpu data? (例如set_cpu_data就是false)
  bool cpu_malloc_use_cuda_;
  bool own_gpu_data_;  // 是否已经申请gpu内存空间
  int gpu_device_;  // 

值得稍加注意的是SyncedHead head_. 该变量的作用会在数据同步部分说明.

get and set 方法

cpu_data, gpu_data或者mutable_cpu_data, mutable_gpu_data方法返回cpu或者gpu内存指针, 前者是const void*, 不可对返回内存进行修改; 后者为void*, 可以修改.

set方法比较特别, 方法参数是指向另一段内存空间的地址:

void SyncedMemory::set_cpu_data(void* data) {
  CHECK(data);
  if (own_cpu_data_) {
    CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);
  }
  cpu_ptr_ = data;
  head_ = HEAD_AT_CPU;
  own_cpu_data_ = false;
}

该函数首先释放自己申请的内存空间, 然后直接指向参数传入的内存空间 (并不是重新申请空间, 并copy数据). 最后将 own_cpu_data_设置为false, 表示外来数据(?).

保持数据同步

在调用cpu_data或者gpu_data方法时, 需要确保cpu, gpu数据内容是一致的. 这里用到了前面提到的枚举类型来记录当前同步状态

  enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };

以to_cpu()方法为例: 检查head_所处状态, 若UNINITIALIZED, 则分配内存空间(置0); 若HEAD_AT_GPU, 则需要从GPU内存同步数据到CPU; HEAD_AT_CPU, 则说明目前最新的数据是在CPU的, 无须进行任何操作 (虽然并不知道GPU的数据是否和CPU一致, 因为当前我们并不关心GPU数据); 若SYNCED, 则CPU/GPU数据一致, 无须进行任何操作.

inline void SyncedMemory::to_cpu() {
  switch (head_) {
  case UNINITIALIZED:
    CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_);
    caffe_memset(size_, 0, cpu_ptr_);
    head_ = HEAD_AT_CPU;
    own_cpu_data_ = true;
    break;
  case HEAD_AT_GPU:
#ifndef CPU_ONLY
    if (cpu_ptr_ == NULL) {
      CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_);
      own_cpu_data_ = true;
    }
    caffe_gpu_memcpy(size_, gpu_ptr_, cpu_ptr_);
    head_ = SYNCED;
#else
    NO_GPU;
#endif
    break;
  case HEAD_AT_CPU:
  case SYNCED:
    break;
  }
}

References

1. [Caffe源码解析2：SycedMem