Caffe源码解读:syncedmem类

内存同步(syncedmem)类的作用在于管理主机(CPU)和设备(GPU)之间的内存分配和数据同步,封装了二者之间的交互操作。

这个类没有对应的ProtoBuffer描述,所以直接看./include/caffe/syncedmem.cpp文件:

#ifndef CAFFE_SYNCEDMEM_HPP_
#define CAFFE_SYNCEDMEM_HPP_

#include <cstdlib>

#include "caffe/common.hpp"

namespace caffe {

// 以页锁定方式分配内存,在单GPU时提示不明显,多GPU提升很多
// malloc分配标准可分页的主机内存,操作系统可能会将这种内存分页或者交换到磁盘上,需要的时候调回内存,这样可能会增加运行时间
// cudaMallocHost分配页锁定的主机内存,操作系统不会对这块内存分页或者交换到磁盘上,可以节省时间
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); // CPU模式下分配内存
  *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) {
    //如果分配内存用的是cudaMallocHost分配,即use_cuda为真,cpu中的数据也可以用cudaMallocHost分配内存  
    CUDA_CHECK(cudaFreeHost(ptr));  
    return;
  }
#endif
  free(ptr);  // 用的malloc分配的内存
}

// 负责内存分配和设备同步
class SyncedMemory {
 public:
  SyncedMemory()     // 默认构造函数 
      : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(0), head_(UNINITIALIZED),
        own_cpu_data_(false), cpu_malloc_use_cuda_(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), cpu_malloc_use_cuda_(false), own_gpu_data_(false),
        gpu_device_(-1) {}
  ~SyncedMemory();

  // 对CPU,GPU数据的读写,不赘述。为什么没有set_cpu_diff() ???
  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();

  // 共享内存的4种状态:未初始化,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 to_cpu();  // 数据同步至cpu
  void to_gpu();  // 数据同步至gpu
  void* cpu_ptr_;  // cpu中数据的指针
  void* gpu_ptr_;   // gpu中数据的指针
  size_t size_;   // 存储空间的大小
  SyncedHead head_;  // 共享内存的状态
  bool own_cpu_data_;  // cpu拥有数据所有权
  bool cpu_malloc_use_cuda_;   // 分配cpu内存是否用cudaMallocHost()分配。
  bool own_gpu_data_;  // gpu拥有数据所有权
  int gpu_device_;  // gpu设备号

  // 禁用拷贝构造以及赋值运算符
  // 使用grep可以查到,该宏定义在common.hpp第35行
  // 该宏就是把拷贝构造和赋值运算符设置为private而已
  DISABLE_COPY_AND_ASSIGN(SyncedMemory); 
};  // class SyncedMemory

}  // namespace caffe
#endif  // CAFFE_SYNCEDMEM_HPP_

这个类比Blob简单多了,下面看对应的./src/caffe/syncedmem.cpp文件:

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

namespace caffe {

SyncedMemory::~SyncedMemory() {
 // 如果cpu拥有数据所有权
 if (cpu_ptr_ && own_cpu_data_) {
   CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);
 }
 // 如果数据在gpu上
#ifndef CPU_ONLY
 if (gpu_ptr_ && own_gpu_data_) {
   int initial_device;
   cudaGetDevice(&initial_device);  // 获取使用的gpu设备号
   if (gpu_device_ != -1) {
     CUDA_CHECK(cudaSetDevice(gpu_device_));
   }
   CUDA_CHECK(cudaFree(gpu_ptr_));
   cudaSetDevice(initial_device);
 }
#endif  // CPU_ONLY
}

// 数据同步到cpu上
inline void SyncedMemory::to_cpu() {
 switch (head_) {
 case UNINITIALIZED:  // 如果是未初始化状态,只需分配下内存就行了
   CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_);
   caffe_memset(size_, 0, cpu_ptr_);  // 定义在math_functions.hpp第42行,调用了memset
   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_);// 要内存
     own_cpu_data_ = true;
   }
   caffe_gpu_memcpy(size_, gpu_ptr_, cpu_ptr_); // 数据复制
   head_ = SYNCED;
#else
   NO_GPU;
#endif
   break;
 // 数据已经为cpu拥有所有权或者在内存共享状态,则什么都不管
 case HEAD_AT_CPU:
 case SYNCED:
   break;
 }
}

// 原理同上
inline void SyncedMemory::to_gpu() {
#ifndef CPU_ONLY
 switch (head_) {
 case UNINITIALIZED:
   CUDA_CHECK(cudaGetDevice(&gpu_device_));
   CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));
   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_); //数据拷贝,调用了cudaMemcpy函数 
   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中的数据
void SyncedMemory::set_cpu_data(void* data) {
 CHECK(data);
 if (own_cpu_data_) {
   // 调用这个函数的时候,如果cpu内有数据会被直接清空,要注意
   CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_); 
 }
 cpu_ptr_ = data;
 head_ = HEAD_AT_CPU;
 own_cpu_data_ = false;
}
// 获取gpu中的数据,只读
const void* SyncedMemory::gpu_data() {
#ifndef CPU_ONLY
 to_gpu();
 return (const void*)gpu_ptr_;
#else
 NO_GPU;
 return NULL;
#endif
}
// 设置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_));
   }
 // 调用这个函数的时候,如果gpu内有数据会被直接清空,要注意
   CUDA_CHECK(cudaFree(gpu_ptr_));
   cudaSetDevice(initial_device);
 }
 gpu_ptr_ = data;
 head_ = HEAD_AT_GPU;
 own_gpu_data_ = false;
#else
 NO_GPU;
#endif
}
// 读写获取cpu数据
void* SyncedMemory::mutable_cpu_data() {
 to_cpu();
 head_ = HEAD_AT_CPU;
 return cpu_ptr_;
}
// 读写获取gpu数据
void* SyncedMemory::mutable_gpu_data() {
#ifndef CPU_ONLY
 to_gpu();
 head_ = HEAD_AT_GPU;
 return gpu_ptr_;
#else
 NO_GPU;
 return NULL;
#endif
}

// cuda中的流同步,这里传入一个异步流,在计算的时候向GPU复制数据。
#ifndef CPU_ONLY
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

syncedmem类比较简单,主要是完成cpu和gpu之间的数据交互问题~

参考资料

  1. 《21天实战caffe》
  2. (介绍了流同步)http://blog.csdn.net/sinat_22336563/article/details/68496919

猜你喜欢

转载自blog.csdn.net/xs11222211/article/details/82457357