mxnet源码阅读笔记之include

写在前面

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