寫在前面

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文件中，以下是它們的關系：

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結構中，如下所示：

在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++"代碼的接口。

mxnet源碼閱讀筆記之include