目录
博客：blog.shinelee.me | 博客园 | CSDN
写在前面
在Caffe源码理解1中介绍了Blob类，其中的数据成员有
shared_ptr<SyncedMemory> data_;
shared_ptr<SyncedMemory> diff_;std::shared_ptr 是共享对象所有权的智能指针，当最后一个占有对象的shared_ptr被销毁或再赋值时，对象会被自动销毁并释放内存，见cppreference.com。而shared_ptr所指向的SyncedMemory即是本文要讲述的重点。
在Caffe中，SyncedMemory有如下两个特点：
- 屏蔽了CPU和GPU上的内存管理以及数据同步细节
- 通过惰性内存分配与同步，提高效率以及节省内存
背后是怎么实现的？希望通过这篇文章可以将以上两点讲清楚。
成员变量的含义及作用
SyncedMemory的数据成员如下：
enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };
void* cpu_ptr_; // CPU侧数据指针
void* gpu_ptr_; // GPU侧数据指针
size_t size_; // 数据所占用的内存大小
SyncedHead head_; // 指示再近一次数据更新发生在哪一侧，在调用另一侧数据时需要将该侧数据同步过去
bool own_cpu_data_; // 指示cpu_ptr_是否为对象内部调用CaffeMallocHost分配的CPU内存
bool cpu_malloc_use_cuda_; // 指示是否使用cudaMallocHost分配页锁定内存，系统malloc分配的是可分页内存，前者更快
bool own_gpu_data_; // 指示gpu_ptr_是否为对象内部调用cudaMalloc分配的GPU内存
int device_; // GPU设备号cpu_ptr_和gpu_ptr_所指向的数据空间有两种来源，一种是对象内部自己分配的，一种是外部指定的，为了区分这两种情况，于是有了own_cpu_data_和own_gpu_data_，当为true时表示是对象内部自己分配的，因此需要对象自己负责释放（析构函数），如果是外部指定的，则由外部负责释放，即谁分配谁负责释放。
外部指定数据时需调用set_cpu_data和set_gpu_data，代码如下：
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; // 指示CPU侧更新了数据
own_cpu_data_ = false; // 指示数据来源于外部
}
void SyncedMemory::set_gpu_data(void* data) {
check_device();
#ifndef CPU_ONLY
CHECK(data);
if (own_gpu_data_) { // 如果自己分配过内存，先释放，换外部指定数据
CUDA_CHECK(cudaFree(gpu_ptr_));
}
gpu_ptr_ = data; // 直接指向外部数据
head_ = HEAD_AT_GPU; // 指示GPU侧更新了数据
own_gpu_data_ = false; // 指示数据来源于外部
#else
NO_GPU;
#endif
}构造与析构
在SyncedMemory构造函数中，获取GPU设备（如果使用了GPU的话），注意构造时head_ = UNINITIALIZED，初始化成员变量，但并没有真正的分配内存。
// 构造
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(); // 校验当前GPU设备以及gpu_ptr_所指向的设备是不是构造时获取的GPU设备
if (cpu_ptr_ && own_cpu_data_) { // 自己分配的空间自己负责释放
CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);
}
#ifndef CPU_ONLY
if (gpu_ptr_ && own_gpu_data_) { // 自己分配的空间自己负责释放
CUDA_CHECK(cudaFree(gpu_ptr_));
}
#endif // CPU_ONLY
}
// 释放CPU内存
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
}
但是，在析构函数中，却释放了CPU和GPU的数据指针，那么是什么时候分配的内存呢？这就要提到，Caffe官网中说的“在需要时分配内存” ，以及“在需要时同步CPU和GPU”，这样做是为了提高效率、节省内存。
具体怎么实现的？我们接着往下看。
内存同步管理
SyncedMemory成员函数如下：
const void* cpu_data(); // to_cpu(); return (const void*)cpu_ptr_; 返回CPU const指针
void set_cpu_data(void* data);
const void* gpu_data(); // to_gpu(); return (const void*)gpu_ptr_; 返回GPU const指针
void set_gpu_data(void* data);
void* mutable_cpu_data(); // to_cpu(); head_ = HEAD_AT_CPU; return cpu_ptr_;
void* mutable_gpu_data(); // to_gpu(); head_ = HEAD_AT_GPU; return gpu_ptr_;
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其中，cpu_data()和gpu_data()返回const指针只读不写，mutable_cpu_data()和mutable_gpu_data()返回可写指针，它们4个在获取数据指针时均调用了to_cpu()或to_gpu()，两者内部逻辑一样，内存分配发生在第一次访问某一侧数据时分配该侧内存，如果不曾访问过则不分配内存，以此按需分配来节省内存。同时，用head_来指示最近一次数据更新发生在哪一侧，仅在调用另一侧数据时才将该侧数据同步过去，如果访问的仍是该侧，则不会发生同步，当两侧已同步都是最新时，即head_=SYNCED，访问任何一侧都不会发生数据同步。下面以to_cpu()为例，
inline void SyncedMemory::to_cpu() {
check_device();
switch (head_) {
case UNINITIALIZED: // 如果未分配过内存（构造函数后就是这个状态）
CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_); // to_CPU时为CPU分配内存
caffe_memset(size_, 0, cpu_ptr_); // 数据清零
head_ = HEAD_AT_CPU; // 指示CPU更新了数据
own_cpu_data_ = true;
break;
case HEAD_AT_GPU: // 如果GPU侧更新过数据，则同步到CPU
#ifndef CPU_ONLY
if (cpu_ptr_ == NULL) { // 如果CPU侧没分配过内存，分配内存
CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_);
own_cpu_data_ = true;
}
caffe_gpu_memcpy(size_, gpu_ptr_, cpu_ptr_); // 数据同步
head_ = SYNCED; // 指示CPU和GPU数据已同步一致
#else
NO_GPU;
#endif
break;
case HEAD_AT_CPU: // 如果CPU数据是最新的，不操作
case SYNCED: // 如果CPU和GPU数据都是最新的，不操作
break;
}
}
// 分配CPU内存
inline void CaffeMallocHost(void** ptr, size_t size, bool* use_cuda) {
#ifndef CPU_ONLY
if (Caffe::mode() == Caffe::GPU) {
CUDA_CHECK(cudaMallocHost(ptr, size)); // cuda malloc
*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";
}下面看一下head_状态是如何转换的，如下图所示：
若以X指代CPU或GPU，Y指代GPU或CPU，需要注意的是，如果HEAD_AT_X表明X侧为最新数据，调用mutable_Y_data()时，在to_Y()内部会将X侧数据同步至Y，会暂时将状态置为SYNCED，但退出to_Y()后最终仍会将状态置为HEAD_AT_Y，如mutable_cpu_data()代码所示，
void* SyncedMemory::mutable_cpu_data() {
check_device();
to_cpu();
head_ = HEAD_AT_CPU;
return cpu_ptr_;
}
不管之前是何种状态，只要调用了mutable_Y_data()，则head_就为HEAD_AT_Y。背后的思想是，无论当前是HEAD_AT_X还是SYNCED，只要调用了mutable_Y_data()，就意味着调用者可能会修改Y侧数据，所以认为接下来Y侧数据是最新的，因此将其置为HEAD_AT_Y。
至此，就可以理解Caffe官网上提供的何时发生内存同步的例子，以及为什么建议不修改数据时要调用const函数，不要调用mutable函数了。
// Assuming that data are on the CPU initially, and we have a blob.
const Dtype* foo;
Dtype* bar;
foo = blob.gpu_data(); // data copied cpu->gpu.
foo = blob.cpu_data(); // no data copied since both have up-to-date contents.
bar = blob.mutable_gpu_data(); // no data copied.
// ... some operations ...
bar = blob.mutable_gpu_data(); // no data copied when we are still on GPU.
foo = blob.cpu_data(); // data copied gpu->cpu, since the gpu side has modified the data
foo = blob.gpu_data(); // no data copied since both have up-to-date contents
bar = blob.mutable_cpu_data(); // still no data copied.
bar = blob.mutable_gpu_data(); // data copied cpu->gpu.
bar = blob.mutable_cpu_data(); // data copied gpu->cpu.以上。