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
最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 204,590评论 6 478
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 86,808评论 2 381
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 151,151评论 0 337
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 54,779评论 1 277
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 63,773评论 5 367
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 48,656评论 1 281
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 38,022评论 3 398
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 36,678评论 0 258
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 41,038评论 1 299
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 35,659评论 2 321
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 37,756评论 1 330
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 33,411评论 4 321
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 39,005评论 3 307
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 29,973评论 0 19
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 31,203评论 1 260
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 45,053评论 2 350
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 42,495评论 2 343

推荐阅读更多精彩内容