從零開始山寨Caffe·貳：主存模型

你左手是內存，右手是顯存，內存可以打死顯存，顯存也可以打死內存。

　　　　　　　　　　　　　　　　　　　　　　 　　　　　　—— 請協調好你的主存

從硬件說起

物理之觴

大部分Caffe源碼解讀都喜歡跳過這部分，我不知道他們是什么心態，因為這恰恰是最重要的一部分。

內存的管理不擅，不僅會導致程序的立即崩潰，還會導致內存的泄露，當然，這只針對傳統CPU程序而言。

由於GPU的引入，我們需要同時操縱倆種不同的存儲體：

一個受北橋控制，與CPU之間架起地址總線、控制總線、數據總線。

一個受南橋控制，與CPU之間僅僅是一條可憐的PCI總線。

一個傳統的C++程序，在操作系統中，會被裝載至內存空間上。

一個有趣的問題，你覺得CPU能夠訪問顯存空間嘛？你覺得你的默認C++代碼能訪問顯存空間嘛？

結果顯然是否定的，問題就在於CPU和GPU之間只存在一條數據總線。

沒有地址總線和控制總線，你除了讓CPU發送數據拷貝指令外，別無其它用處。

這不是NVIDIA解決不了，AMD就能解決的問題。除非計算機體系結構再一次迎來變革，

AMD和NVIDIA的工程師聯名要求在CPU和GPU之間追加復雜的通信總線用於異構程序設計。

當然，你基本是想多了。

環境之艱

可憐的數據總線，加大了異構程序設計的難度。

於是我們看到，GPU的很大一部分時鍾周期，用在了和CPU互相交換數據。

也就是所謂的“內存與顯存之間友好♂關系”。

你不得不接受一個事實：

GPU最慢的存儲體，也就是片外顯存，得益於鎂光的GDDR技術，目前家用游戲顯卡的訪存速度也有150GB/S。

而我們可憐的內存呢，你以為配上Skylake后，DDR4已經很了不起了，實際上它只有可憐的48GB/S。

那么問題來了，內存如何去彌補與顯存的之間帶寬的差距？

答案很簡單：分時、異步、多線程。

換言之，如果GPU需要在接下來1秒內，獲得CPU的150GB數據，那么CPU顯然不能提前一秒去復制。

它需要提前3秒、甚至4秒。如果它當前還有其它串行任務，你就不得不設個線程去完成它。

這就是新版Caffe增加的新功能之一：多重預緩沖。

設置於DataLayer的分支線程，在GPU計算，CPU空閑期間，為顯存預先緩沖3~4個Batch的數據量，

來解決內存顯存帶寬不一致，導致的GPU時鍾周期浪費問題，也增加了CPU的利用率。

最終，你還是需要牢記一點：

不要嘗試以默認的C++代碼去訪問顯存空間，除非你把它們復制回內存空間上。

否則，就是一個毫無提示的程序崩潰問題（准確來說，是被CPU硬件中斷了【微機原理或是計算機組成原理說法】)

編程之繁

在傳統的CUDA程序設計里，我們往往經歷這樣一個步驟：

->計算前

cudaMalloc(....)          【分配顯存空間】

cudaMemset(....)　　 【顯存空間置0】

cudaMemcpy(....)　   【將數據從內存復制到顯存】

->計算后

cudaMemcpy(....)       【將數據從顯存復制回內存】

這些步驟相當得繁瑣，你不僅需要反復敲打，而且如果忘記其中一步，就是毀滅性的災難。

這還僅僅是GPU程序設計，如果考慮到CPU/GPU異構設計，那么就更麻煩了。

於是，聰明的人類就發明了主存管理自動機，按照按照一定邏輯設計狀態轉移代碼。

這是Caffe非常重要的部分，稱之為SyncedMemory(同步存儲體)。

 

主存模型

狀態轉移自動機

自動機共有四種狀態，以枚舉類型定義於類SyncedMemory中：

enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };

這四種狀態基本會被四個應用函數觸發：cpu_data()、gpu_data()、mutable_cpu_data()、mutable_gpu_data()

在它們之上，有四個狀態轉移函數：to_cpu()、to_gpu()、mutable_cpu()、mutable_gpu()

前兩個狀態轉移函數用於未進入Synced狀態之前的狀態機維護，后兩個用於從Synced狀態中打破出來。

具體細節見后文，因為Synced狀態會忽略to_cpu和to_gpu的行為，打破Synced狀態只能靠人工賦值，切換狀態頭head。

后兩個mutable函數會被整合在應用函數里，因為它們只需要簡單地為head賦個值，沒必要大費周章寫個函數封裝。

★UNINITIALIZED：

UNINITIALIZED狀態很有趣，它的生命周期是所有狀態里最短的，將隨着CPU或GPU其中的任一個申請內存而終結。

在整個內存周期里，我們並非一定要遵循着，數據一定要先申請內存，然后在申請顯存，最后拷貝過去。

實際上，在GPU工作的情況下，大部分主存儲體都是直接申請顯存的，如除去DataLayer的前向/反向傳播階段。

所以，UNINITIALIZED允許直接由to_gpu()申請顯存。

由此狀態轉移時，除了需要申請內存之外，通常還需要將內存置0。

 

★HEAD_AT_CPU：

該狀態表明最近一次數據的修改，是由CPU觸發的。

注意，它只表明最近一次是由誰修改，而不是誰訪問。

在GPU工作時，該狀態將成為所有狀態里生命周期第二短的，通常自動機都處於SYNCED和HEAD_AT_GPU狀態，

因為大部分數據的修改工作都是GPU觸發的。

該狀態只有三個來源：

I、由UNINITIALIZED轉移到：說白了，就是欽定你作為第一次內存的載體。

II、由mutable_cpu_data()強制修改得到：都要准備改數據了，顯然需要重置狀態。

cpu_data()及其子函數to_cpu()，只要不符合I條件，都不可能轉移到改狀態(因為訪問不會引起數據的修改)

 

★HEAD_AT_GPU：

該狀態表明最近一次數據的修改，是由GPU觸發的。

幾乎是與HEAD_AT_CPU對稱的。

 

★SYNCED：

最重要的狀態，也是唯一一個非必要的狀態。

單獨設立同步狀態的原因，是為了標記內存顯存的數據一致情況。

由於類SyncedMemory將同時管理兩種主存的指針，

如果遇到HEAD_AT_CPU，卻要訪問顯存。或是HEAD_AT_GPU，卻要訪問內存，那么理論上，得先進行主存復制。

這個復制操作是可以被優化的，因為如果內存和顯存的數據是一致的，就沒必要來回復制。

所以，使用SYNCED來標記數據一致的情況。

SYNCED只有兩種轉移來源：

I、由HEAD_AT_CPU+to_gpu()轉移到：

含義就是，CPU的數據比GPU新，且需要使用GPU，此時就必須同步主存。

II、由HEAD_AT_GPU+to_cpu()轉移到：

含義就是，GPU的數據比CPU新，且需要使用CPU，此時就必須同步主存。

在轉移至SYNCED期間，還需要做兩件准備工作：

I、檢查當前CPU/GPU態的指針是否分配主存，如果沒有，就重新分配。

II、復制主存至對應態。

處於SYNCED狀態后，to_cpu()和to_gpu()將會得到優化，跳過內部全部代碼。

自動機將不再運轉，因為，此時僅需要返回需要的主存指針就行了，不需要特別維護。

這種安寧期會被mutable前綴的函數打破，因為它們會強制修改至HEAD_AT_XXX，再次啟動自動機。

 

代碼實戰

主存操作函數

建立synced_memory.hpp，在操作主存之前，你需要封裝一些基礎函數。

CPU端的函數是C/C++標准的通用函數：

inline void dragonMalloc(void **ptr, size_t size){
    *ptr = malloc(size);
    CHECK(*ptr) << "host allocation of size " << size << " failed";
}
inline void dragonFree(void *ptr){
    free(ptr);
}
inline void dragonMemset(void *ptr,size_t size){
    memset(ptr, 0, size);
}
inline void dragonMemcpy(void* dest, void* src,size_t size){
    memcpy(dest, src, size);
}

CHECK宏由GLOG提供，條件為假時，會觸發assert，終結程序。

GPU端的函數由CUDA提供：

#ifndef CPU_ONLY
#include "cuda.h"
inline void cudaSetDevice(){
    int device;
    cudaGetDevice(&device);
    if (device != -1) return;
    CUDA_CHECK(cudaSetDevice(0));
}
inline void dragonGpuMalloc(void **ptr, size_t size){
    cudaSetDevice();
    CUDA_CHECK(cudaMalloc(ptr, size));
}
inline void dragonGpuFree(void *ptr){
    cudaSetDevice();
    CUDA_CHECK(cudaFree(ptr));
}
inline void dragonGpuMemset(void *ptr, size_t size){
    cudaSetDevice();
    CUDA_CHECK(cudaMemset(ptr, 0, size));
}
inline void dragonGpuMemcpy(void *dest, void* src, size_t size){
    cudaSetDevice();
    CUDA_CHECK(cudaMemcpy(dest, src, size, cudaMemcpyDefault));
}
#endif

#ifndef CPU ONLY  ..... #endif 確保本段代碼不會被非CUDA模式所編譯

cudaSetDevice()是一個通用函數，在后期，你應該移至common.hpp中。

該函數不是必要的，目的只是對當前執行GPU的一個慣性檢查，檢查失敗則終結程序。

需要注意的是cudaMemcpy的最后一個參數，Flag：cudaMemcpyDefault，在CUDA 6.0之后才被使用。

在6.0版本之前，cudaMemcpy需要指明dest和src的來源，是host向device，還是device向host，還是device向device？

所以，早期的CUDA代碼可能需要三個if來指明Flag的值，而cudaMemcpyDefault會自動檢測，相當智能。

數據結構

class SyncedMemory
{
public:
    SyncedMemory():cpu_ptr(NULL), gpu_ptr(NULL), size_(0), head_(UNINITIALIZED) {}
    SyncedMemory(size_t size) :cpu_ptr(NULL), gpu_ptr(NULL), size_(size), head_(UNINITIALIZED) {}
    void to_gpu();
    void to_cpu();
    const void* cpu_data();
    const void* gpu_data();
    void set_cpu_data(void *data);
    void set_gpu_data(void *data);
    void* mutable_cpu_data();
    void* mutable_gpu_data();
#ifndef CPU_ONLY
    void async_gpu_data(const cudaStream_t& stream);
#endif
    enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };
    void *cpu_ptr, *gpu_ptr;
    size_t size_;
    bool own_cpu_data, own_gpu_data;
    SyncedHead head_;
    SyncedHead head() { return head_; }
    size_t size() { return size_; }
    ~SyncedMemory();
};

成員變量包括：

★兩個主存指針：cpu_ptr、gpu_ptr

★主存大小size以及狀態標記head

★共享標記：own_cpu_data、 own_gpu_data

成員函數包括：

★狀態轉移函數：void to_gpu()、void to_cpu()

★常訪問函數：const void* cpu_data()、 const void* gpu_data()

★修改函數：void* mutable_cpu_data()、void* mutable_gpu_data()

★共享函數：void set_cpu_data(void *data)、void set_gpu_data(void *data)

★封裝訪問函數：SyncedHead head()、size_t size()

★異步流同步函數與析構函數：void async_gpu_data(const cudaStream_t& stream)、~SyncedMemory()

值得注意的是，兩個共享函數以及共享標記不屬於自動機范圍。

共享函數的唯一用處是用於局部主存的共享，只用於DataLayer的Transformer中。

在Blob級別的共享中，存在兩種共享：

I、共享另一個Blob的全部數據：只需要讓SyncedMemory指針重新指向另一個Blob的SyncedMemory指針

II、共享另一個Blob的部分數據：

這利用了C/C++內存指針的一個Trick，內存首指針可以做代數加減運算，做一定偏移。

set_xxx_data(void *data)提供了最底層的指針修改，可以直接指向偏移之后的內存，而共享部分數據。

實現

建立synced_memory.cpp。

void SyncedMemory::to_cpu()
{
    switch (head_){
    case UNINITIALIZED:
        dragonMalloc(&cpu_ptr, size_);
        dragonMemset(cpu_ptr, size_);
        head_ = HEAD_AT_CPU;
        own_cpu_data = true;
        break;
    case HEAD_AT_GPU:
#ifndef CPU_ONLY
        if (cpu_ptr == NULL){
            dragonMalloc(&cpu_ptr, size_);
            own_cpu_data = true;
        }
        dragonGpuMemcpy(cpu_ptr,gpu_ptr,size_);
        head_ = SYNCED;
#endif
        break;
    case HEAD_AT_CPU:
    case SYNCED:
        break;
    }
}

需要注意的是共享標記own_xxx_data，只要申請了內存，就必須做標記。

共享標記在析構時是必要的，因為你不能將宿主數據一並釋放掉。

void SyncedMemory::to_gpu()
{
#ifndef CPU_ONLY
    switch (head_){
    case UNINITIALIZED:
        dragonGpuMalloc(&gpu_ptr,size_);
        dragonGpuMemset(gpu_ptr, size_);
        head_ = HEAD_AT_GPU;
        own_gpu_data = true;
        break;
    case HEAD_AT_CPU:
        if (gpu_ptr == NULL){
            dragonGpuMalloc(&gpu_ptr,size_);
            own_gpu_data = true;
        }
        dragonGpuMemcpy(gpu_ptr, cpu_ptr, size_);
        head_ = SYNCED;
        break;
    case HEAD_AT_GPU:
    case SYNCED:
        break;
    }
#endif
}

GPU的轉移函數是與CPU版本對稱的。

const void* SyncedMemory::cpu_data(){
    to_cpu();
    return (const void*)cpu_ptr;
}

const void* SyncedMemory::gpu_data(){
    to_gpu();
    return (const void*)gpu_ptr;
}

常訪問函數，注意const指針的強制轉換，訪問之前需要運行一次自動機。

void* SyncedMemory::mutable_cpu_data(){
    to_cpu();
    head_ = HEAD_AT_CPU;
    return cpu_ptr;
}

void* SyncedMemory::mutable_gpu_data(){
#ifndef CPU_ONLY
    to_gpu();
    head_ = HEAD_AT_GPU;
    return gpu_ptr;
#endif
}

修改函數，運行自動機、強制修改自動機狀態，最后返回指針，用於從外部修改。

void SyncedMemory::set_cpu_data(void *data){
    if (own_cpu_data) dragonFree(cpu_ptr);
    cpu_ptr = data;
    head_ = HEAD_AT_CPU;
    own_cpu_data = false;
}

void SyncedMemory::set_gpu_data(void *data){
#ifndef CPU_ONLY
    if (own_gpu_data) dragonGpuFree(gpu_ptr);
    gpu_ptr = data;
    head_ = HEAD_AT_GPU;
    own_gpu_data = false;
#endif
}

共享函數，共享之前，先釋放舊主存，修改共享標記，強制修改自動機狀態。

SyncedMemory::~SyncedMemory(){
    if (cpu_ptr && own_cpu_data) dragonFree(cpu_ptr);
#ifndef CPU_ONLY
    if (gpu_ptr && own_gpu_data) dragonGpuFree(gpu_ptr);
#endif
}

析構函數，注意檢查共享標記，不能釋放宿主內存。

異步流同步

異步流概念，是CUDA 5.0中引入的。

與Intel CPU的流水線架構一樣，NVIDIA的GPU也采用了I/O和計算分離的流水線做法。

cudaMemcpy使用的是默認流cudaStreamDefault，編號為0。

異步流編程API開放之后，允許程序員在CPU端多線程編程中，向GPU提交異步的同步復制流，

以此增加GPU端的I/O利用率。

簡單來說，默認流只允許主進程與顯存復制數據，而我們實際上不可能這么干，原因有二：

I、效率低，主進程就是單線程啊。

II、很多情況下，數據復制完之前，需要阻塞。阻塞主進程不是一個好主意。

Caffe中只有一處是這么做的，那就是DataLayer正向傳播一個Batch的時候，這個阻塞是必然的。

但是，在構成Batch之前，只要采取多線程設計，那么異步流復制只會阻塞旁支線程，而不會影響主進程。

這是為什么NVIDIA開放異步流API的原因，它鼓勵了CPU用於多線程I/O，讓GPU計算如虎添翼。

#ifndef CPU_ONLY
void SyncedMemory::async_gpu_data(const cudaStream_t& stream){
    CHECK(head_ == HEAD_AT_CPU);
    //    first allocating memory
    if (gpu_ptr == NULL){
        dragonGpuMalloc(&gpu_ptr, size_);
        own_gpu_data = true;
    }
    const cudaMemcpyKind kind = cudaMemcpyHostToDevice;
    CUDA_CHECK(cudaMemcpyAsync(gpu_ptr, cpu_ptr, size_, kind, stream));
    head_ = SYNCED;
}
#endif

異步流的底層代碼接受一個異步流作為參數，使用cudaMemcpyAsync()向GPU提交復制任務。

它等效於HEAD_AT_CPU+to_gpu()，所以需要更新同步標記。

完整代碼將在DataLayer中完成，該函數將由多線程調用。

完整代碼

synced_mem.hpp:

https://github.com/neopenx/Dragon/blob/master/Dragon/include/synced_mem.hpp

synced_mem.cpp:

https://github.com/neopenx/Dragon/blob/master/Dragon/src/synced_mem.cpp