【Caffe代码解析】SyncedMemory

功能：

Caffe的底层数据的切换（cpu模式和gpu模式），需要用到内存同步模块。

源码：头文件

#ifndef CAFFE_SYNCEDMEM_HPP_#define CAFFE_SYNCEDMEM_HPP_#include <cstdlib>#include "caffe/common.hpp"#include "caffe/util/math_functions.hpp"namespace caffe {inline void CaffeMallocHost(void** ptr, size_t size) {  *ptr = malloc(size);  CHECK(*ptr) << "host allocation of size " << size << " failed";}inline void CaffeFreeHost(void* ptr) {  free(ptr);}/** * @brief Manages memory allocation and synchronization between the host (CPU) *        and device (GPU). * * TODO(dox): more thorough description. */class SyncedMemory { public:  SyncedMemory()      : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(0), head_(UNINITIALIZED),        own_cpu_data_(false) {}  explicit SyncedMemory(size_t size)      : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(size), head_(UNINITIALIZED),        own_cpu_data_(false) {}  ~SyncedMemory();  const void* cpu_data();//获取cpu数据，返回void * 指针  void set_cpu_data(void* data);//用一个void * 指针修改指针  const void* gpu_data();//获取gpu数据，返回void * 指针  void* mutable_cpu_data();//获取可以更改cpu数据，返回void * 指针  void* mutable_gpu_data();//获取可以更改gpu数据，返回void * 指针  enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };//enum枚举值  SyncedHead head() { return head_; }//获得枚举值  size_t size() { return size_; }//获得数据大小 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数据  DISABLE_COPY_AND_ASSIGN(SyncedMemory);};  // class SyncedMemory}  // namespace caffe#endif  // CAFFE_SYNCEDMEM_HPP_

实现文件：

#include <cstring>#include "caffe/common.hpp"#include "caffe/syncedmem.hpp"#include "caffe/util/math_functions.hpp"namespace caffe {//析构函数，调用caffe函数来释放空间SyncedMemory::~SyncedMemory() {  if (cpu_ptr_ && own_cpu_data_) {    CaffeFreeHost(cpu_ptr_);  }//如果有gpu数据，也进行释放。#ifndef CPU_ONLY  if (gpu_ptr_) {    CUDA_CHECK(cudaFree(gpu_ptr_));  }#endif  // CPU_ONLY}// 根据head信息，选择：1， 分类cpu空间，2.拷贝gpu数据值cpuinline void SyncedMemory::to_cpu() {  switch (head_) {  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    if (cpu_ptr_ == NULL) {      CaffeMallocHost(&cpu_ptr_, size_);      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;  }}// 根据head信息，选择：1， 分类gpu空间，2.拷贝cpu数据值gpuinline void SyncedMemory::to_gpu() {#ifndef CPU_ONLY  switch (head_) {  case UNINITIALIZED:    CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));    caffe_gpu_memset(size_, 0, gpu_ptr_);    head_ = HEAD_AT_GPU;    break;  case HEAD_AT_CPU:    if (gpu_ptr_ == NULL) {      CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));    }    caffe_gpu_memcpy(size_, cpu_ptr_, gpu_ptr_);    head_ = SYNCED;    break;  case HEAD_AT_GPU:  case SYNCED:    break;  }#else  NO_GPU;#endif}//返回cpu指针 void * 类型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_) {    CaffeFreeHost(cpu_ptr_);  }  cpu_ptr_ = data;//直接重置指针，  head_ = HEAD_AT_CPU;  own_cpu_data_ = false;//设false}//获得gpu指针const void* SyncedMemory::gpu_data() {#ifndef CPU_ONLY  to_gpu();  return (const void*)gpu_ptr_;#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;#endif}}  // namespace caffe