目錄
部落格：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.
以上。