博客：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”，這樣做是為了提高效率、節省內存。
Blobs conceal the computational and mental overhead of mixed CPU/GPU operation by synchronizing from the CPU host to the GPU device as needed. Memory on the host and device is allocated on demand (lazily) for efficient memory usage.
具體怎么實現的？我們接着往下看。
內存同步管理
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.
A rule of thumb is, always use the const call if you do not want to change the values, and never store the pointers in your own object. Every time you work on a blob, call the functions to get the pointers, as the SyncedMem will need this to figure out when to copy data.
以上。