caffe代码阅读9：SyncedMemory的实现细节-2016.3.28

一、SyncedMemory的作用简介

SyncedMemory类主要负责在GPU或者CPU上分配内存以及保持数据的同步作用。

SyncedMemory类主要应用在BLOB类中，我们可以在BLOB类中看出一些使用方法

比如：

// blob中的reshape 的具体实现
template <typename Dtype>
void Blob<Dtype>::Reshape(const vector<int>& shape) {
  CHECK_LE(shape.size(), kMaxBlobAxes); //是否小于规定的最大BLOB的维度(35维)
  count_ = 1;
  shape_.resize(shape.size());// 首先将大小设置为vector<int> shape_; 即新的形状数据的大小
  if (!shape_data_ || shape_data_->size() < shape.size() * sizeof(int)) {
    shape_data_.reset(new SyncedMemory(shape.size() * sizeof(int)));//  shared_ptr<SyncedMemory> shape_data_;
  }
  int* shape_data = static_cast<int*>(shape_data_->mutable_cpu_data());
  for (int i = 0; i < shape.size(); ++i) {
    // 检查形状数据是否合法
    CHECK_GE(shape[i], 0);
    CHECK_LE(shape[i], INT_MAX / count_) << "blob size exceeds INT_MAX";
    // 计算数据个数
    count_ *= shape[i];
    // 复制shape到新的和旧的形状数据
    shape_[i] = shape[i];
    shape_data[i] = shape[i];
  }
  // 判断是否大于存储的容量
  if (count_ > capacity_) {
    capacity_ = count_;
    // 重新分配内存
    data_.reset(new SyncedMemory(capacity_ * sizeof(Dtype)));
    diff_.reset(new SyncedMemory(capacity_ * sizeof(Dtype)));
  }
}

// 注意这里的 gpu_data和mutable_gpu_data的区别是
// mutable_gpu_data是设置了head_ 的而
// gpu_data是不设置head_
// 但是我并不明白为啥要有个mutable
// 难道是因为设置了head_就可以实现互斥了?
// 可能的解释是设置了head_可以保护gpu或者内存中的数据不被销毁

template <typename Dtype>
const int* Blob<Dtype>::gpu_shape() const {
  CHECK(shape_data_);
  // shared_ptr<SyncedMemory> shape_data_;
  // 因此也分gpu_data和cpu_data
  return (const int*)shape_data_->gpu_data();
}

另一个blob中使用的例子：

// Update是计算data=-1 * diff + data
template <typename Dtype>
void Blob<Dtype>::Update() {
  // We will perform update based on where the data is located.
  switch (data_->head()) {
  case SyncedMemory::HEAD_AT_CPU:
    // perform computation on CPU
    // axpby即alpha * x plus beta *y 这个含义,blas的函数命名真是见名知意
    // template <> void caffe_axpy<float>(const int N, const float alpha, const float* X, float* Y) { cblas_saxpy(N, alpha, X, 1, Y, 1); }
    // caffe_axpy计算的是Y=alpha * X + Y ，其中alpha=-1了这里
    // 存储的时候用到了mutable_cpu_data，防止其他线程访问
    caffe_axpy<Dtype>(count_, Dtype(-1),
        static_cast<const Dtype*>(diff_->cpu_data()),
        static_cast<Dtype*>(data_->mutable_cpu_data()));
    break;
  case SyncedMemory::HEAD_AT_GPU:
  case SyncedMemory::SYNCED:
#ifndef CPU_ONLY
    // perform computation on GPU
    // Y=alpha * X + Y ，其中alpha=-1了这里
    caffe_gpu_axpy<Dtype>(count_, Dtype(-1),
        static_cast<const Dtype*>(diff_->gpu_data()),
        static_cast<Dtype*>(data_->mutable_gpu_data()));
#else
    NO_GPU;
#endif
    break;
  default:
    LOG(FATAL) << "Syncedmem not initialized.";
  }
}

二、SyncedMemory类的详细介绍

1）构造函数

// 构造函数
  SyncedMemory()
        // 一开始的时候就是未初始化
      : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(0), head_(UNINITIALIZED),
        own_cpu_data_(false), own_gpu_data_(false), gpu_device_(-1) {}
  explicit SyncedMemory(size_t size)
      : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(size), head_(UNINITIALIZED),
        own_cpu_data_(false), own_gpu_data_(false), gpu_device_(-1) {}
   // 析构函数
  ~SyncedMemory();

2）成员变量

  

// 指向内存的指针
  void* cpu_ptr_;
  void* gpu_ptr_;
  // 数据大小
  size_t size_;
  // 同步状态
  SyncedHead head_;
  // 是否拥有cpu还是gpu数据
  bool own_cpu_data_;
  bool own_gpu_data_;
  // 设备编号
  int gpu_device_;

3）成员函数

  

// 获取CPUDATA
  const void* cpu_data();
  // 设置CPUDATA
  void set_cpu_data(void* data);
  // 获取GPUDATA
  const void* gpu_data();
  // 设置GPUDATA
  void set_gpu_data(void* data);
  // 获取互斥的CPU或者GPUDATA
  void* mutable_cpu_data();
  void* mutable_gpu_data();
  // 枚举类型，未初始化，在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:
 // 内部使用的到cpu还是gpu
  void to_cpu();
  void to_gpu();

具体的实现如下：

三、SyncedMemory类的具体实现

#include <cstring>

#include "caffe/common.hpp"
#include "caffe/syncedmem.hpp"
#include "caffe/util/math_functions.hpp"

namespace caffe {
// 析构函数就是释放内存
SyncedMemory::~SyncedMemory() {
  if (cpu_ptr_ && own_cpu_data_) {
    CaffeFreeHost(cpu_ptr_);
  }

#ifndef CPU_ONLY// 只要不是定义的CPU_ONLY的编译模式
  if (gpu_ptr_ && own_gpu_data_) {
    int initial_device;
    // 获取可用设备
    cudaGetDevice(&initial_device);
    if (gpu_device_ != -1) {
        // 当前所使用的设备
      CUDA_CHECK(cudaSetDevice(gpu_device_));
    }
    // 释放当前
    CUDA_CHECK(cudaFree(gpu_ptr_));
    cudaSetDevice(initial_device);
  }
#endif  // CPU_ONLY
}

// 内部使用的
// 如果当前未初始化，直接在内存分配空间
// 如果在GPU上则复制到内存
// 如果已经在内存则啥都不动
inline void SyncedMemory::to_cpu() {
  switch (head_) {
  // 如果当前是未初始化，直接分配CPU上的内存
  case UNINITIALIZED:
    CaffeMallocHost(&cpu_ptr_, size_);
    caffe_memset(size_, 0, cpu_ptr_);
    head_ = HEAD_AT_CPU;
    own_cpu_data_ = true;
    break;
  case HEAD_AT_GPU:
#ifndef CPU_ONLY
    // 如果当前数据在GPU，然后cpu_ptr为空
    if (cpu_ptr_ == NULL) {
        // 分配内存
      CaffeMallocHost(&cpu_ptr_, size_);
      own_cpu_data_ = true;
    }
    // 复制数据
    caffe_gpu_memcpy(size_, gpu_ptr_, cpu_ptr_);
    head_ = SYNCED;
#else// CPU_ONLY模式当然只能报错了
    NO_GPU;
#endif
    break;
  case HEAD_AT_CPU:
  case SYNCED:
    break;
  }
}

// 内部使用的
// 如果当前未初始化直接在GPU分配内存
// 如果当前在CPU，则在GPU上分配内存并且复制到GPU
// 如果数据已经在GPU则啥也不做
inline void SyncedMemory::to_gpu() {
#ifndef CPU_ONLY
  switch (head_) {
  case UNINITIALIZED:
    // 获取设备
    CUDA_CHECK(cudaGetDevice(&gpu_device_));
    // 在设备上分配内存
    CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));
    // 初始化为0
    caffe_gpu_memset(size_, 0, gpu_ptr_);
    head_ = HEAD_AT_GPU;
    own_gpu_data_ = true;
    break;
  case HEAD_AT_CPU:
    if (gpu_ptr_ == NULL) {
      CUDA_CHECK(cudaGetDevice(&gpu_device_));
      CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));
      own_gpu_data_ = true;
    }
    caffe_gpu_memcpy(size_, cpu_ptr_, gpu_ptr_);
    head_ = SYNCED;
    break;
  case HEAD_AT_GPU:
  case SYNCED:
    break;
  }
#else
  NO_GPU;
#endif
}

// 首先不管三七二十一将数据搞到内存上去
// 然后获取cpu上的数据
const void* SyncedMemory::cpu_data() {
  to_cpu();
  return (const void*)cpu_ptr_;
}


// 如果当前cpu_ptr_有内存上的数据则先释放
// 然后再将地址给内部变量cpu_ptr_
// 设置cpu上的数据
void SyncedMemory::set_cpu_data(void* data) {
  CHECK(data);
  if (own_cpu_data_) {
    CaffeFreeHost(cpu_ptr_);
  }
  cpu_ptr_ = data;
  head_ = HEAD_AT_CPU;
  own_cpu_data_ = false;
}

// 首先不管三七二十一将数据搞到GPU上去
// 然后在获取gpu上的数据
// 但是并没有改变head_的值(head_表明数据究竟在哪儿)
const void* SyncedMemory::gpu_data() {
#ifndef CPU_ONLY
  to_gpu();
  return (const void*)gpu_ptr_;
#else
  NO_GPU;
#endif
}

// 如果当前gpu_ptr_有内存上的数据则先释放
// 然后再将地址给内部变量gpu_ptr_
// 设置gpu上的数据
void SyncedMemory::set_gpu_data(void* data) {
#ifndef CPU_ONLY
  CHECK(data);
  if (own_gpu_data_) {
    int initial_device;
    cudaGetDevice(&initial_device);
    if (gpu_device_ != -1) {
      CUDA_CHECK(cudaSetDevice(gpu_device_));
    }
    CUDA_CHECK(cudaFree(gpu_ptr_));
    cudaSetDevice(initial_device);
  }
  gpu_ptr_ = data;
  head_ = HEAD_AT_GPU;
  own_gpu_data_ = false;
#else
  NO_GPU;
#endif
}

// 首先不管三七二十一先数据搞到CPU上去
// 然后返回互斥的cpu_ptr_指针
// mutable_cpu_data与cpu_data的区别就是是否设置head
void* SyncedMemory::mutable_cpu_data() {
  to_cpu();
  head_ = HEAD_AT_CPU;
  return cpu_ptr_;
}


// 首先不管三七二十一先数据搞到GPU上去
// 然后返回互斥的gpu_ptr_指针
// mutable_gpu_data与gpu_data的区别就是是否设置head
void* SyncedMemory::mutable_gpu_data() {
#ifndef CPU_ONLY
  to_gpu();
  head_ = HEAD_AT_GPU;
  return gpu_ptr_;
#else
  NO_GPU;
#endif
}

#ifndef CPU_ONLY
// 异步推送数据到gpu
void SyncedMemory::async_gpu_push(const cudaStream_t& stream) {
  CHECK(head_ == HEAD_AT_CPU);
  if (gpu_ptr_ == NULL) {
    CUDA_CHECK(cudaGetDevice(&gpu_device_));
    CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));
    own_gpu_data_ = true;
  }
  const cudaMemcpyKind put = cudaMemcpyHostToDevice;
  CUDA_CHECK(cudaMemcpyAsync(gpu_ptr_, cpu_ptr_, size_, put, stream));
  // Assume caller will synchronize on the stream before use
  head_ = SYNCED;
}
#endif

}  // namespace caffe

四、总结

该类主要就是在内存分配空间以及在GPU上分配空间，并且负责同步数据，此外我看mutable_cpu_data和cpu_data 这两个函数的主要区别就是head_是否改变，至于这两个函数的命名上的mutable是有着互斥的含义的。究竟体现在哪儿，我的感受是，这里的mutable的体现主要是在调用了mutable_cpu_data之后强制设置了head_为HEAD_AT_CPU，从而保护了cpu上的数据？

参考：

暂无