mxnet原始碼閱讀筆記之include

摘要： 寫在前面 mxnet程式碼的規範性比Caffe2要好，看起來核心程式碼量也小很多，但由於對dmlc其它庫的依賴太強，程式碼的獨立性並不好。依賴的第三方庫包括： cub dlpack dmlc-core googletest mkldnn mshadow onnx-tensorrt op...

寫在前面

mxnet程式碼的規範性比Caffe2要好，看起來核心程式碼量也小很多，但由於對dmlc其它庫的依賴太強，程式碼的獨立性並不好。依賴的第三方庫包括：

cub
dlpack
dmlc-core
googletest
mkldnn
mshadow
onnx-tensorrt
openmp
ps-lite
tvm

如果對於這些第三方庫沒有足夠的理解，mxnet的核心程式碼看起來比較費勁。因此時間原因，本篇僅解析了mxnet對外的介面include目錄，並且對於嚴重依賴第三方庫的檔案沒有深入探究，只能算作一篇不完整的原始碼閱讀筆記了。後續有時間的話，再回來迭代。

目錄

• storage
• tensor_blob
• ndarray
• resource
• kvstore
• base
• operator
• engine
• executor
• rtc
• graph_attr_types
• op_attr_types
• imperative
• operator_util
• c_api

storage

Storage是一個跨裝置的記憶體管理類，它提供了記憶體分配和回收的功能，但並不儲存分配的記憶體，真正的記憶體指標分配在Storage類內部的Handle結構體中：

struct Handle {
void * dptr{nullptr}; //記憶體地址
size_t size{0};
Context ctx;
int shared_pid{-1};
int shared_id{-1};
};

class Storage {
public:
Handle Alloc(size_t size, Context ctx) {...};
virtual void Alloc(Handle* handle) = 0;
virtual void Free(Handle handle) = 0;
};

tensor_blob

TBlob類可以表示任意維度、在任意裝置上、任意資料型別的張量，它是NDArray的內部儲存，是mxnet中最底層的資料結構。但本質上它是對DLTensor的代理，DLTensor定義在第三方庫dlpack中的dlpack.h檔案中，以下是它們的關係：

graph LR NDArray-->|包含|TBlob TBlob-->|包含|DLTensor

ndarray

ndarray是mxnet中的核心資料結構，代表了多維資料，類似於Tensorflow中的Tensor。本質上它借鑑了numpy中關於ndarray的定義，一部分ndarray是包含實際資料的，另外一些ndarray並不包含實際資料，它們只是其他ndarray的檢視。舉例說明，ndarrayA是一個[1x12]的多維陣列，儲存了12個元素，ndarrayB是一個[3x4]的多維陣列，它底層的資料由ndarrayA提供，因此A和B共享了記憶體，B僅是A的一個檢視。

ndarray內部由chunk結構提供實際的資料儲存，先來看下chunk：

struct Chunk {
Storage::Handle shandle;
std::vector<Storage::Handle> aux_handles;
bool static_data; //如果為真，表示該資料是靜態的，並非來自Storage，不需要被釋放
bool delay_alloc; //資料分配是否需要延緩，注意對輔助資料aux data無效
NDArrayStorageType storage_type = kDefaultStorage;
std::vector<int> aux_types;
Context ctx;
TShape storage_shape;
std::vector<TShape> aux_shapes;
};

可見，Chunk結構仍然不是最終的資料儲存結構，本質上資料還是儲存在Storage結構中，如下所示：

graph LR NDArray-->|使用|Chunk Chunk-->|使用|Storage

在ndarray中，我們發現數據分為資料本身，以及輔助資料。輔助資料主要用於儲存稀疏資料的時候，資料本身放在data中，資料索引放在aux_data中。

最後看下NDArray的資料結構：

class NDArray {
std::shared_ptr<Chunk> ptr_{nullptr};
TShape shape_;
size_t byte_offset_ = 0;
int dtype_ = -1;
bool reuse_ = false;
nnvm::NodeEntry entry_;
mutable TBlob tblob_;
};

resource

在mxnet中，計算中用到的所有內容，除了ndarray之外，都可以被稱為資源。其中最常用的資源，就是隨機數生成器，分為CPU和GPU兩個版本，如下：

enum Type {
kRandom, //CPU版本隨機數生成器
kTempSpace, //動態隨機記憶體
kParallelRandom //可以在GPU中使用的並行隨機數生成器
};

另外，mxnet還為資源提供了一個管理器，ResourceManager，用於獲取資源。

kvstore

kv儲存的作用是儲存模型引數，以便在分散式的計算中，在多個裝置/機器之間進行資料同步。

kv儲存可以有多種型別，比如：

• 'local'或者'local_update_cpu‘或者'local_allreduce_cpu'，表明這是一個單機的kv儲存，並且僅使用cpu做kv的allreduce；
• 'device'或者'local_allreduce_device'，也是單機的kv儲存，只不過使用gpu做kv的allreduce；
• 'dist_*'，分散式的kv儲存；

每個kv儲存中都有一個更新器，它定義了，針對指定的key，當新value來到時，如何與舊value進行融合。這一點非常重要，因為在深度學習模型的訓練中，需要迭代式的對模型引數進行更新，而更新的方式就是通過更新器來定義。

kv儲存中，key通常是整型或者字串，而value是NDArray，因此，有兩種更新器的定義：

typedef std::function<void(int, const NDArray&, NDArray*)> Updater;
typedef std::function<void(const std::string&, const NDArray&, NDArray*)> StrUpdater;

最後，kv儲存在底層用到了ps-lite來作資料同步。

class KVStore {
public:
static KVStore *Create(const char *type = "local");

virtual void Init(const std::vector<int>& keys, const std::vector<NDArray>& values) = 0;
virtual void Init(const std::vector<std::string>& str_keys, const std::vector<NDArray>& values) = 0;

virtual void Push(...) = 0;
virtual void Pull(...) = 0;
virtual void PullRowSparse(...) = 0;

virtual void set_updater(...);
};

base

引入了兩個類，執行環境的上下文資訊類Context，實際執行時的上下文類RunContext，後者包含前者。首先看下Context類的定義：

struct Context {
DeviceType dev_type;
int32_t dev_id;
inline void Save(dmlc::Stream *strm) const {...}; //將Context資訊記入二進位制流
inline bool Load(dmlc::Stream *strm) {...}; //從二進位制流中載入Context資訊
inline static Context Create(DeviceType dev_type, int32_t dev_id = -1); //構造一個新的Context
inline static Context CPU(int32_t dev_id = 0);
inline static Context GPU(int32_t dev_id=-1);
inline static int32_t GetGPUCount(); //獲取GPU的數量
inline static void GetGPUMemoryInformation(int dev, int *free, int *total);
inline static Context CPUPinned(int32_t dev_id = -1);
inline static Context CPUShared(int32_t dev_id = 0);
inline static Context FromString(const std::string& str);
};

而RunContext就相對簡單了，它包含了一個Context和一個流指標：

struct RunContext {
Context ctx;
void *stream;
//...
};

operator

Operator定義了mxnet計算圖中基礎的操作單位。相當於Tensorflow中的kernel，和Caffe2中的Operator。但它與Tensorflow和Caffe2中的操作有本質區別，在Tensorflow中，操作本身和它對應的求導操作是分開的，而在mxnet中，這兩者是結合在一起的，分別使用Forward和Backward兩個函式實現，因此，mxnet在操作的實現上更加緊湊，與Tensorflow相比減少了一些對計算圖進行裁剪的額外開銷，效能上有優勢，但也同時限制了自己的計算邊界，靈活性不足。

class Operator {
public:
//進行前向計算，將計算結果儲存在TBlob中
virtual void Forward(const OpContext &ctx, const std::vector<TBlob> &in_data, const std::vector<OpReqType> &req, const std::vector<TBlob> &out_data, const std::vector<TBlob> &aux_states) = 0;

//進行後向計算，將梯度寫入in_grad
virtual void Backward(const OpContext &ctx, const std::vector<TBlob> &out_grad, const std::vector<TBlob> &in_data, const std::vector<TBlob> &out_data, const std::vector<OpReqType> &req, const std::vector<TBlob> &in_grad, const std::vector<TBlob> &aux_states);
};

Operator中僅包含了操作計算的介面，對於操作的描述儲存在OperatorProperty類中，它負責儲存所有與Operator有關的資訊，且能夠產生裝置相關的Operator。同時，它還為計算引擎提供了一些可以優化操作計算的函式。

class OperatorProperty {
public:
//初始化Operator時需要用到的引數
virtual void Init(const std::vector<std::pair<std::string, std::string>>& kwargs) = 0;
//獲取為Operator準備的引數
virtual std::map<std::string, std::string> GetParams() const = 0;

virtual int NumOutputs() const {...}
//進行Operator的形狀推斷，類似於Tensorflow的ShapeInference
virtual bool InferShape(std::vector<TShape> *in_shape, std::vector<TShape> *out_shape, std::vector<TShape> *aux_shape) const = 0;
//進行Operator的型別推斷
virtual bool InferType(...);

//構建Operator
virtual Operator* CreateOperator(Context ctx) const = 0;
};

目前看來，mxnet中Operator與OperatorProperty的關係，與Tensorflow中OpKernel與Op的關係不太一樣，後者與Caffe2中的Operator和OpSchema的關係更加相似，有機會我們會詳細比較下，這三種框架關於操作定義於使用的區別。

engine

引擎是執行核心之一，它負責對計算圖中的操作進行排程。引擎中的兩大關鍵元素是操作和變數，操作定義了計算圖每一個節點需要實際執行的動作，變數定義了動作之間的依賴關係。

首先，mxnet定義了一個，被非同步函式在執行結束時呼叫的回撥函式類，通過對()的過載，用類對回撥函式進行了一層封裝：

class CallbackOnComplete {
public:
inline void operator()() const {
(*callback_)(engine_, param_);
}
private:
friend class ::mxnet::Engine;
void (*callback_)(Engine *, void *);
Engine* engine_;
void* param_;
};

列舉類FnProperty介紹了常用的函式型別：

enum class FnProperty {
kNormal, //一般操作
kCopyFromGPU, //從GPU上拷貝內容到其它裝置的操作
kCopyToGPU, //從其它裝置向GPU拷貝內容的操作
kCPUPrioritized, //CPU上優先選擇的同步操作
kAsync, //非同步函式呼叫
kDeleteVar, //用來刪除變數的函式
kGPUPrioritized, //GPU上優先選擇的同步操作
};

engine的含義是，對操作進行排程執行的引擎。回想一下，在Tensorflow中，為了正確執行使用者設計好的計算圖，我們需要對原始計算圖進行一些迭代修改，在Engine類中提供了這樣的介面：

class Engine {
public:
//定義執行結束時的回撥類
typedef engine::CallbackOnComplete CallbackOnComplete;
//定義傳遞給引擎的同步操作函式
typedef std::function<void(RunContext)> SyncFn;
//定義傳遞給引擎的非同步操作函式
typedef std::function<void(RunContext, CallbackOnComplete)> AsyncFn;
//定義變數指標
typedef engine::VarHandle VarHandle;
//定義操作指標
typedef engine::OprHandle OprHandle;

//停止引擎中的所有worker
virtual void Stop() {}
//啟動引擎中的所有worker
virtual void Start() {}

//分配一個新的變數，該變數可以被用來根據依賴關係，輔助對引擎中的操作進行排程
virtual VarHandle NewVariable() = 0;
//構建一個操作，該操作定義在外部，從而我們可以在排程中重複使用
virtual OprHandle NewOperator(...) = 0;
//刪除一個操作，它不會立刻進行，而是直到所有使用該操作的動作執行結束之後再進行
virtual void DeleteOperator(OpHandle op) = 0;
//將一個操作加入引擎
virtual void Push(...);
//將一個非同步操作加入引擎
virtual void PushAsync(...);
//將一個同步操作加入引擎
virtual void PushSync(...);
//刪除一個變數，它不會立刻進行，而是直到所有依賴該變數的操作完成之後再進行
virtual void DeleteVariable(...) = 0;
//等待一個變數準備完成
virtual void WaitForVar(...) = 0;
//等待引擎中所有的活動都結束時再返回
virtual void WaitForAll() = 0;

//返回引擎的單例物件
static Engine* Get();
//用來生成OnComplete回撥的工廠函式
inline CallbackOnComplete CreateCallback(...);
};

executor

mxnet的執行器介面，用於對計算圖進行執行。執行的機制與Operator的設計相合，同樣提供了前向和後向兩種介面，如下：

class Executor {
public:
virtual void Forward(bool is_train) = 0;
virtual void PartialForward(bool is_train, int step, int *step_left) = 0;
virtual void Backward(const std::vector<NDArray> &head_grads, bool is_train = true) = 0;
};

rtc

包含了Cuda執行時的編譯模組CudaModule。

graph_attr_types

獲取圖相關屬性的輔助結構。對於一張計算圖中的節點，通常會關注兩種資訊，一種是計算圖中節點的儲存型別，一種是節點的排程模式，分別將結果儲存在StorageTypeVector和DispatchModeVector中，這兩種結構的定義如下：

using StorageTypeVector = std::vector<int>;
using DispatchModeVector = std::vector<DispatchMode>;

op_attr_types

有關操作的額外屬性，與nvvm有關，目前看不懂。

imperative

與NDArray有關的執行時函式，目前看不懂。

operator_util

輔助快速構建operator的功能和註冊器。

c_api

定義了mxnet後端"C++"程式碼的介面。