Skip to content

mshadow调研

Jump to bottom Edit New page
Tao Luo edited this page Dec 9, 2019 · 1 revision

mshadow除了Tensor的表示以外，还包含了计算的实现。为了保证计算高效，mshadow的整体设计参考了惰性计算（lazy operation）的原则。关于lazy operation的详细信息可以参考这里。

Exp

Exp是mshaow中最重要的部分，用于表示一个表达式，例如a+b。任何可以放在=两边的内容都可以被认为是Exp。显然，单一的一个tensor也是一个Exp。Exp的基类定义如下：

template<typename SubType, typename DType, int exp_type>
struct Exp {
 public:
  /*! \return  subtype instance of current class */
  inline const SubType& self(void) const {
    return *static_cast<const SubType*>(this);
  }
  /*! \return reference of subtype instance of current class */
  inline SubType* ptrself(void) {
    return static_cast<SubType*>(this);
  }
};

Exp本身并不会被直接使用，使用的是其经过CRTP派生后的基类。经过特化和派生，Exp的各种子类被用于表示各种各样的表达式。模板参数的意义如下：

  • SubType表示经过派生后的具体类型，例如Tensor
  • Dtype表示数据的类型
  • exp_type表示表达式自身的的类型，分为kRValue、kMapper、kChainer和kComplex四种。
    • kRValue：已经存在、可直接使用的数据。
    • kMapper：两个tensor之间element-wise的操作。
    • kChainer：由上面两种表达式组成的复杂表达式。
    • kComplex：其他表达式。例如点积操作。

Exp的主要派生类有：

  • TernaryMapExp：三元操作
  • BinaryMapExp：两元操作
  • UnaryMapExp：一元操作
  • RValueExp已经存在、可以直接使用的数据，其exp_type一定为kRValue。

另外mshadow也还提供了一些专门处理特定操作的Exp，例如：

  • ScalarExp表示一个标量
  • TypecastExp数据类型转换表达式
  • TransposeExp转置表达式
  • DotExp点积操作

其中RValueExp是Tensor、Tensor1D等具体数据类型的基类，同时它也是唯一一个可以放在等号左边的Exp。mshadow通过重载RValueExp的+=、-=等和赋值有关的operator实现了惰性计算（lazy operation）。

Shape

Shape本质上就是一个整数数组，用来表示Tensor的维度。

Tensor

mshadow中，RValueExp派生出TRValue，TRvalue又进一步派生出Tensor，它主要有如下模板参数：

  • typename Device：GPU / CPU
  • int dimension：维度
  • typename Dtype：数据类型

Tensor的主要成员变量有：

  • Dtype *dptr_指向数据的指针
  • Shape<dimension> shape_维度
  • int stride_：Tensor的最底层维度上相邻元素在内存上的距离。如果元素紧密排列，则stride_ = 1

*dptr_指向的内存的管理不由Tensor负责。因此Tensor的新建和销毁一般通过全局函数NewTensor和FreeSpace进行，这两个函数在构造或者析构Tensor之前会先进行内存的申请和销毁。

除了必要的构造函数等，Tensor本身还提供了其他一些成员函数，主要有：

  • MemSize返回某个维度上相邻元素在内存上的距离
  • FlatTo1D/FlatTo2D将Tensor在逻辑上转化为一维 / 二维
  • operator[]返回某个idx上维度为d-1的Tensor（浅拷贝）
  • Slice类似于Python的切片操作

Tensor本身没有包含任何计算的定义。

Plan

一个完整的计算表达式一般包含两个部分：被计算对象和计算操作，其中计算对象在Exp及其派生类中定义，计算操作则在Plan中定义。

Plan的定义为：

template<typename ExpType, typename DType>
class Plan {
 public:
  /*!
   * \brief evaluate the expression at index [y][x]
   *  to be implemented by SubType, for RValue, the return type will be DType &
   */
  MSHADOW_XINLINE DType Eval(index_t y, index_t x) const;
};

通过用ExpType作为模板参数进行特化，Plan与具体的Exp一一绑定，每一个Exp的派生类都需要定义自己的Plan。从逻辑上来说，任何表达式执行完成后，一定会得到一个tensor，对表达式计算操作的定义，其实就等价于定义结果tensor中每一个元素生成的方法。这个定义在Plan中由Eval函数完成。

例如，当表达式是一个简单的tensor时候的特化：

template <typename Device, int dim, typename DType>
class Plan<Tensor<Device, dim, DType>, DType> {
 public:
  explicit Plan(const Tensor<Device, dim, DType> &t)
      : dptr_(t.dptr_), stride_(t.stride_) {}
  // for RValue, the return type should be reference
  MSHADOW_XINLINE DType &REval(index_t y, index_t x) {
    return dptr_[y * stride_ + x];
  }
  // const evaluation
  MSHADOW_XINLINE const DType &Eval(index_t y, index_t x) const {
    return dptr_[y * stride_ + x];
  }

 private:
  DType  *dptr_;
  index_t stride_;
};

Eval函数所做的就是直接取出对应位置上的值。

需要特别指出，Eval函数只有两个参数输入，因此只能处理结果为二维tensor的表达式。事实上，为了更简单地利用多线程加速，mshadow会在使用CPU时将任何表达式的结果tensor在逻辑上转化为二维，再交由Plan处理。

ExpEngine

mshadow中定义了名为ExpEngine的struct，用来执行一个等式，这里的等式不一定要有等号，也可以是+=、-=、*=、/=这些赋值操作。上文已经提到，等式的左边一定是一个RValueExp，mshadow重载了RValueExp的赋值operator，在内部调用ExpEngine来完成等式右边表达式的计算和对左边表达式的赋值。

ExpEngine的核心是Eval成员函数。根据赋值操作的不同（=，+=，-=，*=，/=）和右侧表达式类型的不同（kMapper，kChainer，kRValue，kComplex），Eval函数有多种类型的特化：

template<typename SV, typename RV, typename DType>
struct ExpEngine {
  template<typename E>
  inline static void Eval(RV *dst,
                          const Exp<E, DType, type::kMapper> &exp) {
    MapExp<SV>(dst, exp);
  }
  template<typename E>
  inline static void Eval(RV *dst,
                          const Exp<E, DType, type::kChainer> &exp) {
    MapExp<SV>(dst, exp);
  }
  template<typename E>
  inline static void Eval(RV *dst,
                          const Exp<E, DType, type::kRValue> &exp) {
    MapExp<SV>(dst, exp);
  }
  template<typename E>
  inline static void Eval(RV *dst,
                          const Exp<E, DType, type::kComplex> &exp) {
    ExpComplexEngine<SV, RV, E, DType>::Eval(dst->ptrself(), exp.self());
  }
};

其中的模板参数SV表示具体的赋值操作，可以看到Eval函数会根据具体情况进一步调用MapExp或者将需要执行的等式转交给ExpComplexEngine执行。在CPU环境中，MapExp函数在经过一系列的检查后最终调用全局函数MapPlan：

template<typename Saver, typename R, int dim,
         typename DType, typename E>
inline void MapPlan(TRValue<R, cpu, dim, DType> *dst,
                    const expr::Plan<E, DType> &plan) {
  Shape<2> shape = expr::ShapeCheck<dim, R>::Check(dst->self()).FlatTo2D();
  expr::Plan<R, DType> dplan = expr::MakePlan(dst->self());
#if (MSHADOW_USE_CUDA == 0)
  #pragma omp parallel for
#endif
  // temp remove openmp, as default setting throttles CPU
  for (openmp_index_t y = 0; y < shape[0]; ++y) {
    for (index_t x = 0; x < shape[1]; ++x) {
      // trust your compiler! -_- they will optimize it
      Saver::template Save<DType>(dplan.REval(y, x), plan.Eval(y, x));
    }
  }
}

MapPlan函数将左边的Tensor和右边表达式的计算结果在逻辑上平摊至二维以便于多线程并行。

MapPlan的本质在于建立等式两边表达式的计算结果在元素上的映射关系。例如在上面的例子中：

Saver::template Save<DType>(dplan.REval(y, x), plan.Eval(y, x));

其中，plan是等式右边表达式的Plan，通过Eval获得了其计算结果在[y,x]位置上的元素；dplan是等式左边表达式（一般就是一个tensor）的Plan，REval获得了该tensor在[y,x]位置上元素的引用。Save函数进行两个元素之间的赋值操作，这个操作具体是+=还是-=，或者是别的什么操作，则由Save的特化来确定。

各个概念相互之间的关系

可以通过一个简单的例子来说明上面出现过的各个概念相互之间的关系：

A += B + C

其中A``B``C都是tensor，+为逐元素相加。

在这个例子中，A作为等式左边的接受计算结果的对象，是一个RValueExp，B + C则是一个一般类型的Exp。

A和B + C都有用自己的类型特化产生的Plan，Plan中有Eval函数，它定义了表达式执行后的结果tensor中每一个元素的产生规则。例如，A + B这个Exp特化产生的Plan<Exp>中的Eval，就会定义res[i,j] = B[i,j] + C[i,j]这样的规则。

MapPlan函数则定义了等式左右两边元素之间的赋值关系和具体的赋值操作。例如本例子中为对应元素之间的累加赋值。

A作为一个RValueExp，它的operator+=已经被重载，会在内部调用ExpEngine中的Eval函数，执行上述的计算操作。

TensorContainer

对Tesnor的进一步派生，自带内存的申请和释放，并以STL的风格管理内存，因此可以进行resize。

TBlob

用于暂存任意维度、设备和数据类型的Tensor。TBlob本身不提供任何算法，仅仅用于暂时持有一个Tensor。Tensor本身所具有的模板参数，在TBlob中成为变量保存。

Add a custom footer
Clone this wiki locally