CAFFE源码学习笔记之二-Syncmemory

一、前言

在CPU与GPU合作的异构体系中，CPU负责逻辑性强的事务处理，GPU只负责高度线程化的数据处理。在这个体系中，CPU和GPU之间将会存在大量的数据交换，其中涉及到的内存操作就有

Malloc(),

Free(),

memset(),

cudaMalloc(),

cudaMemset(),

cudaMencpy(),

cudaFree()……..

在C++中仅仅是new和delete都会因为在不注意的情况下犯错，这么多繁琐的操作必然是个大隐患。

所以我觉得Syncmemory类就是将这些操作设计成状态机，和verilog中的状态机很像，可以帮助用户更好的管理内存。

另外还有GPU与CPU之间的同步，毕竟两者带宽存在着倍数的差异。CPU带宽在80GB/s,GPU带宽在150GB/s。就是因为CPU带宽差距大，所以有可能GPU已经扑哧扑哧地算完，结果发现没有数据了。。。所以在后面caffe有个basePrefetchingLayer在GPU计算的时候预先取若干batch的数据。

二、源码分析

1、头文件：

mkl是和cblas相似的数学计算库，比cblas性能强，全称英特尔数学核心函数库，是一套经过高度优化和广泛线程化的数学例程，专为需要极致性能的科学、工程及金融等领域的应用而设计。当然平时用cblas就够了。

#ifdef USE_MKL
#include "mkl.h"
#endif

cudaMalloc()是在GPU上使用的Malloc函数，而cudaMallocHost()同样也是Malloc函数，不同的是后者表示分配的显存将会为锁页内存。锁页内存可以使主机和设备之间的传输带宽增加，在某些设备上还可以通过zero-copy功能映射到设备地址空间，CPU可以直接访问，省去了CPU和GPU之间的数据传输。

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
#ifdef USE_MKL
*ptr = mkl_malloc(size ? size:1, 64);
#else
*ptr = malloc(size);
#endif
*use_cuda = false;
CHECK(*ptr) << "host allocation of size " << size << " failed";
}

有Malloc()就有Free():

inline void CaffeFreeHost(void* ptr, bool use_cuda) {
#ifndef CPU_ONLY
if (use_cuda) {
CUDA_CHECK(cudaFreeHost(ptr));
return;
}
#endif
#ifdef USE_MKL
mkl_free(ptr);
#else
free(ptr);
#endif
}

那么Syncedmemory类的头文件如下：

class SyncedMemory {
public:
SyncedMemory();
explicit SyncedMemory(size_t size);
~SyncedMemory();
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();
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();

void to_cpu();
void to_gpu();
void* cpu_ptr_;
void* gpu_ptr_;
size_t size_;
SyncedHead head_;//状态
bool own_cpu_data_;//cpu是否保有数据
bool cpu_malloc_use_cuda_;
bool own_gpu_data_;//gpu是否保有数据
int device_;

DISABLE_COPY_AND_ASSIGN(SyncedMemory);
};  // class SyncedMemory

cpu_data()就是指向CPU内存中的数据，gpu_ptr()则是指向显存中的数据。有const标志显示两者都是只读的。

而mutable_cpu_data()/mutable_gpu_data()则是可以改写的，这样的标记很清楚。

2、实现

构造函数初始化数据默认为空，注意的是head_的初始化态是UNINITIALIZED。

SyncedMemory::SyncedMemory()
: cpu_ptr_(NULL), gpu_ptr_(NULL), size_(0), head_(UNINITIALIZED),
own_cpu_data_(false), cpu_malloc_use_cuda_(false), own_gpu_data_(false) {
#ifndef CPU_ONLY
#ifdef DEBUG
CUDA_CHECK(cudaGetDevice(&device_));
#endif
#endif
}

SyncedMemory::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) {
#ifndef CPU_ONLY
#ifdef DEBUG
CUDA_CHECK(cudaGetDevice(&device_));
#endif
#endif
}

SyncedMemory::~SyncedMemory() {
check_device();
if (cpu_ptr_ && own_cpu_data_) {
CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);
}//注意条件，cpu_ptr()为NULL时，数据可能在GPU内，故own_cpu_data_必须同时为true

#ifndef CPU_ONLY
if (gpu_ptr_ && own_gpu_data_) {
CUDA_CHECK(cudaFree(gpu_ptr_));
}
#endif  // CPU_ONLY
}

接下来就是该类的核心-内存管理状态机了。

四种状态

UNINITIALIZED，初始化状态

HEAD_AT_GPU，表明在上一次的数据是由GPU完成的，状态内将GPU内的数据复制到CPU内。

HEAD_AT_CPU，与上相反。

SYNCED，表明GPU与CPU数据一致的情况，就不需要来回复制了。

inline void SyncedMemory::to_cpu() {
check_device();
switch (head_) {
case UNINITIALIZED:
CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_);//在cpu上开辟内存
caffe_memset(size_, 0, cpu_ptr_);//初始化内存
head_ = HEAD_AT_CPU;//目前在cpu上操作
own_cpu_data_ = true;//cpu保有数据
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_);//将cpu内存中的数据复制到gpu显存中
head_ = SYNCED;//cpu刚刚将数据复制到gpu上，两者的数据相同，避免重复再复制
#else
NO_GPU;
#endif
break;
case HEAD_AT_CPU:
case SYNCED:
break;
}
}

忘了最基本的了，gpu_data()/cpu_data(),根据此可以推测内存都是线性增长。注意该函数返回的数据不可更改。

const void* SyncedMemory::cpu_data() {
check_device();
to_cpu();
return (const void*)cpu_ptr_;//表明只读
}

而相应的可以更改版本：

void* SyncedMemory::mutable_cpu_data() {
check_device();
to_cpu();
head_ = HEAD_AT_CPU;
return cpu_ptr_;//返回的指针指向的对象可以更改
}

同时，还有set，可以用于偏移进行内存共享，传入指针前需要先释放旧数据：

void SyncedMemory::set_cpu_data(void* data) {
check_device();
CHECK(data);
if (own_cpu_data_) {
CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);
}
cpu_ptr_ = data;
head_ = HEAD_AT_CPU;
own_cpu_data_ = false;
}

流同步：CUDA中的程序是通过流来管理的，每个流顺序执行，不同流乱序执行，所以不同流的计算与传输的时候可以同时进行，从而提高GPU利用率。这里传入一个异步流，在计算的时候像GPU复制数据。

void SyncedMemory::async_gpu_push(const cudaStream_t& stream) {
check_device();
CHECK(head_ == HEAD_AT_CPU);
if (gpu_ptr_ == NULL) {
CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));
own_gpu_data_ = true;
}
const cudaMemcpyKind put = cudaMemcpyHostToDevice;
CUDA_CHECK(cudaMemcpyAsync(gpu_ptr_, cpu_ptr_, size_, put, stream));//cuda中的流同步
// Assume caller will synchronize on the stream before use
head_ = SYNCED;
}

最后，CUDA_CHECK()相当于ASSERT断言。

三、总结

CPU内存管理，GPU显存管理，CPU与GPU数据交换，同步。