caffe源码剖析（二）--SyncedMemory解析

今天我们分析caffe中的底层数据格式，在caffe中，数据均是以一维数组的形式存放的。并且因为涉及到CPU和GPU之间的数据交换，因此在caffe中定义了一个叫做SyncedMemory的类来进行数据的封装和设备之间的数据交换。

SyncedMemory类的声明与实现在文件syncedmemory.hpp（目录：include/caffe/syncedmemory.hpp）和syncedmemory.cpp（目录：src/caffe/syncedmemory.cpp）中

1.SyncedMemory

SyncedMemory类的定义如下所示（重要部分已经注上注释，英文部分为caffe原始注释）

/**
 * @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();
//    获取CPU中的数据
  const void* cpu_data();
  void set_cpu_data(void* data);
//    获取GPU中的数据
  const void* gpu_data();
  void set_gpu_data(void* data);
//    获取CPU中的数据，并且该数据是可变的
  void* mutable_cpu_data();
//    获取GPU中的数据，并且该数据是可变的
  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() { return head_; }
  size_t size() { return size_; }

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

 private:
// 检查设备
 void check_device();
// 将数据同步到CPU
  void to_cpu();
// 将数据同步到GPU
  void to_gpu();
// 指向CPU数据的指针
 void* cpu_ptr_;
// 指向GPU数据的指针
 void* gpu_ptr_;
//  数据大小（unsigned long）
  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

下面我们着重分析to_cpu()和cpu_data()两个函数的实现

2.void to_cpu()

实现代码如下：

inline void SyncedMemory::to_cpu() {
  // 第一步检查GPU设备
  check_device();
  // 对当前数据状态进行判断
  /**
   * 1.UNINITIALIZED
   *   表明数据尚未初始化，因此进行内存分配
   * 2.HEAD_AT_GPU
   *   数据只在GPU端存在，那么
   *   (1) 在CPU上进行内存分配
   *   (2) 把GPU端的数据复制到CPU
   *   (3) 将head_设置为SYNCED,表示数据已经进行同步
   * 3.HEAD_AT_CPU和SYNCED
   *   当前数据已经在在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;
  }
}

3. const void* SyncedMemory::cpu_data()

实现代码如下：

const void* SyncedMemory::cpu_data() {
  // 检查GPU设备
  check_device();
  // 将数据同步到CPU
  to_cpu();
  // 返回数据的const指针，表明数据不可变(与mutable_cpu_data()相反)
  return (const void*)cpu_ptr_;
}

4.总结

GPU端的数据获取与CPU基本一致，这里不再赘述，关于SyncedMemory的分析就到此为止。下一次我们将要分析Blob