mxnet 源碼閱讀 6

標籤：clu   memory   contex   gpu   相關   pup   size_t   ice   storage   

mxnet/src/storage/storage.cc

mxnet/include/mxnet/storage.h

mxnet/include/mxnet/base.h

 

以上三個檔案共同描述了Storage虛類以及其執行個體化的結果，其中storage.h中定義了Storage抽象介面(虛函數) Alloc Free DirectFree 以及靜態方法Get ，其中定義了以後用以操作和管理記憶體的Handle控制代碼，其內部有兩個部分組成一個是描述記憶體本身屬性的包括地址 void *dptr 和大小 size_t size 以及描述記憶體所在裝置的Context ctx; 其中Context結構在base.h中定義，其成員有兩個一個是 dev_type (1-CPU 2-GPU 3-CPU_PINNED ) 另一是 dev_id ，同時在base中還有一些構建Context的相關函數，如GPU(...) CPU(...) CPUPinned(,,,) FromString(...) 等可以直接構建相關裝置環境。

StorageImpl是 Stroage的具體繼承類介面有 返回Handle 的 Alloc 函數 Free 函數 和 DirectFree函數，同時多了靜態函數 ActivateDevice用以切換GPU裝置，CPU裝置其實不用切換，所以實現也只是直接break而已。

從Alloc的實現看，其申請的記憶體包括 CPU記憶體，CPUpinned 記憶體 和 GPU記憶體池，沒有單純的申請一塊GPU記憶體的介面。

同時總結一下mxnet/src/storage 目錄中檔案和類的情況：

cpu_device_storage.h gpu_device_storage.h pinned_memory_storage.h 分別構建了負責CPU GPU 和 CPUpinned 記憶體申請的類。

storage_manager.h 是純粹的抽象的介面類。naive_storage_manager.h 和  pooled_storage_manager.h 是抽象類別storage_manager是具體實現。實質使用是通過C++多態，分別管理了 CPU（包括CPUpinned）和GPU(記憶體池)。

 

整個storage目錄中的內容在mxnet架構中屬於底層模組python等介面涉及不到該介面。

總結起來講，其提供了以Handle為操作的控制代碼的記憶體申請和釋放介面，即 Handle Alloc(size_t size, Context ctx)，void Free(Handle handle)，void DirectFree(Handle handle) 函數，統一了CPU GPU 介面。CPU 和CPUpinned 是隨時申請隨時釋放模式，GPU記憶體採用記憶體池的模式。

猜測gpu_device_storage.h是為了以後臨時資源的申請之用。

mxnet 源碼閱讀 6