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文件中，以下是它们的关系：
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++"代码的接口。
