MShadow中的表达式模板

表达式模板是Eigen、GSL和boost.uBLAS等高性能C++矩阵库的核心技术。本文基于MXNet给出的教程文档来阐述MXNet所依赖的高性能矩阵库MShadow背后的原理。
编写高效的机器学习代码
我们先来思考一个问题:如何才能编写出高效的机器学习代码?假设DNN模型按照下面的代码进行权重更新,其中weightgrad都是长度为n的vector:

weight = -eta * (grad + lambda * weight)

既然我们选择C++来实现矩阵计算,那么性能肯定是最先要考虑的因素。在C++编程中,一个重要的原则就是——预先分配好所需的内存,不要在运行时申请分配临时内存。因此,我们可能会实现一个名为UpdateWeight的函数,其形参gradweight均为指针:
void UpdateWeight(const float *grad, float eta, float lambda, int n, float *weight) { for (int i = 0; i < n; ++i) { weight[i] =-eta * (grad[i] + lambda * weight[i]); } }

指针gradweight所指向的内存空间都是预先分配好的,函数在运行期只需要执行计算。显然,上面的代码非常简单直观,但是,当我们反复编写它们时可能会很烦(指编写循环处理每个元素)。因此问题是,我们能否像下面这样编写代码,同时又能获得上述代码的性能呢?答案是肯定的,但是方法可能并没有那么直观。
void UpdateWeight(const Vec& grad, float eta, float lambda, Vec& weight) { weight = -eta * (grad + lambda * weight); }

运算符重载
运算符重载是一种非常容易想到的解决方案。通过重载相应的运算符,我们可以将元素处理的细节隐藏在运算符中,简单地调用运算符就可以实现相应的操作。
// Naive solution for vector operation overloading struct Vec { int len; float* dptr; Vec(int len) : len(len) { dptr = new float[len]; } Vec(const Vec& src) : len(src.len) { dptr = new float[len]; memcpy(dptr, src.dptr, sizeof(float)*len ); } ~Vec(void) { delete [] dptr; } }; inline Vec operator+(const Vec &lhs, const Vec &rhs) { Vec res(lhs.len); for (int i = 0; i < lhs.len; ++i) { res.dptr[i] = lhs.dptr[i] + rhs.dptr[i]; } return res; }

然而,这种方法并不高效,原因是每次调用运算符时都会有内存空间的申请和释放。另一种更高效的方法是仅重载运算符+=和-=,他们无需临时内存分配即可实现, 但这又限制了我们可以调用的运算符的数量,得不偿失。下一小节,我们将介绍如何利用表达式模板实现延迟计算。
延迟计算
在调用operator+时,因为我们不知道运算符的结果要赋值给哪个变量,所以需要申请一块临时内存空间把结果保存下来。否则,如果我们能提前直到运算结果要存放在哪个变量中,那么就可以直接将结果存储到相应的内存空间。下面的代码说明了这一情况:
// Example Lazy evaluation code // for simplicity, we use struct and make all members public #include struct Vec; // expression structure holds the expression struct BinaryAddExp { const Vec &lhs; const Vec &rhs; BinaryAddExp(const Vec &lhs, const Vec &rhs) : lhs(lhs), rhs(rhs) {} }; // no constructor and destructor to allocate and de-allocate memory, //allocation done by user struct Vec { int len; float* dptr; Vec(void) {} Vec(float *dptr, int len) : len(len), dptr(dptr) {} // here is where evaluation happens inline Vec &operator=(const BinaryAddExp &src) { for (int i = 0; i < len; ++i) { dptr[i] = src.lhs.dptr[i] + src.rhs.dptr[i]; } return *this; } }; // no evaluation happens here inline BinaryAddExp operator+(const Vec &lhs, const Vec &rhs) { return BinaryAddExp(lhs, rhs); }const int n = 3; int main(void) { float sa[n] = {1, 2, 3}; float sb[n] = {2, 3, 4}; float sc[n] = {3, 4, 5}; Vec A(sa, n), B(sb, n), C(sc, n); // run expression A = B + C; for (int i = 0; i < n; ++i) { printf("%d:%f==%f+%f\\n", i, A.dptr[i], B.dptr[i], C.dptr[i]); } return 0; }

在这段代码中,运算符operator+不进行实际的计算,只返回一个表达式结构BinaryAddExp,它里面保存了进行向量加法的两个操作数。当重载operator=时,我们就可以知道向量加法的目标变量以及对应的两个操作数,因此,在这种情况下,不需要任何(运行时)内存分配就可以执行计算操作!类似的,我们可以定义一个DotExp并在operator=处进行惰性求值,然后在内部调用BLAS实现矩阵乘法。
更复杂的表达式以及表达式模板
使用延迟计算能够避免运行期的临时内存分配。但是,上一小节的代码仍然面临以下两个问题:
  • 只能编写形如A=B+C的表达式,不能编写类似A=B+C+D等更加复杂的表达式
  • 如果想添加更多的表达式,就需要编写更多的operator=来执行相应的计算
上述问题的解决方法就是使用模板编程。我们将BinaryAddExp实现成一个模板类,它保存的两个操作数都是模板,这样就能够实现任意长度的加法表达式,具体代码如下。
// Example code, expression template, and more length equations // for simplicity, we use struct and make all members public #include // this is expression, all expressions must inheritate it, //and put their type in subtype template struct Exp { // returns const reference of the actual type of this expression inline const SubType& self(void) const { return *static_cast(this); } }; // binary add expression // note how it is inheritates from Exp // and put its own type into the template argument template struct BinaryAddExp: public Exp > { const TLhs &lhs; const TRhs &rhs; BinaryAddExp(const TLhs& lhs, const TRhs& rhs) : lhs(lhs), rhs(rhs) {} // evaluation function, evaluate this expression at position i inline float Eval(int i) const { return lhs.Eval(i) + rhs.Eval(i); } }; // no constructor and destructor to allocate // and de-allocate memory, allocation done by user struct Vec: public Exp { int len; float* dptr; Vec(void) {} Vec(float *dptr, int len) :len(len), dptr(dptr) {} // here is where evaluation happens template inline Vec& operator= (const Exp& src_) { const EType &src = https://www.it610.com/article/src_.self(); for (int i = 0; i < len; ++i) { dptr[i] = src.Eval(i); } return *this; } // evaluation function, evaluate this expression at position i inline float Eval(int i) const { return dptr[i]; } }; // template add, works for any expressions template inline BinaryAddExp operator+(const Exp &lhs, const Exp &rhs) { return BinaryAddExp(lhs.self(), rhs.self()); }const int n = 3; int main(void) { float sa[n] = {1, 2, 3}; float sb[n] = {2, 3, 4}; float sc[n] = {3, 4, 5}; Vec A(sa, n), B(sb, n), C(sc, n); // run expression, this expression is longer:) A = B + C + C; for (int i = 0; i < n; ++i) { printf("%d:%f == %f + %f + %f\\n", i, A.dptr[i], B.dptr[i], C.dptr[i], C.dptr[i]); } return 0; }

代码的关键思想是模板Exp将其派生类SubType的类型作为模板参数,因此它可以通过self()将其自身转换为SubType,这种模式被称为奇异递归模板模式,简称CRTP。BinaryAddExp现在是一个模板类,可以将多个表达式复合在一起。真正的计算操作是通过函数Eval完成的,该函数在BinaryAddExp中以递归方式实现,operator=中的函数调用src.Eval(i)会被编译成B.dptr[i] + C.dptr[i] + C.dptr[i]
灵活性
前面的示例让我们领略到模板编程的强大功能,而这最后一个示例则与MShadow的实现更为接近,它允许用户自定义二元操作符。
// Example code, expression template // with binary operator definition and extension // for simplicity, we use struct and make all members public #include // this is expression, all expressions must inheritate it, // and put their type in subtype template struct Exp{ // returns const reference of the actual type of this expression inline const SubType& self(void) const { return *static_cast(this); } }; // binary operators struct mul{ inline static float Map(float a, float b) { return a * b; } }; // binary add expression // note how it is inheritates from Exp // and put its own type into the template argument template struct BinaryMapExp: public Exp >{ const TLhs& lhs; const TRhs& rhs; BinaryMapExp(const TLhs& lhs, const TRhs& rhs) :lhs(lhs), rhs(rhs) {} // evaluation function, evaluate this expression at position i inline float Eval(int i) const { return OP::Map(lhs.Eval(i), rhs.Eval(i)); } }; // no constructor and destructor to allocate and de-allocate memory // allocation done by user struct Vec: public Exp{ int len; float* dptr; Vec(void) {} Vec(float *dptr, int len) : len(len), dptr(dptr) {} // here is where evaluation happens template inline Vec& operator=(const Exp& src_) { const EType &src = https://www.it610.com/article/src_.self(); for (int i = 0; i < len; ++i) { dptr[i] = src.Eval(i); } return *this; } // evaluation function, evaluate this expression at position i inline float Eval(int i) const { return dptr[i]; } }; // template binary operation, works for any expressions template inline BinaryMapExp F(const Exp& lhs, const Exp& rhs) { return BinaryMapExp(lhs.self(), rhs.self()); }template inline BinaryMapExp operator*(const Exp& lhs, const Exp& rhs) { return F(lhs, rhs); }// user defined operation struct maximum{ inline static float Map(float a, float b) { return a > b ? a : b; } }; const int n = 3; int main(void) { float sa[n] = {1, 2, 3}; float sb[n] = {2, 3, 4}; float sc[n] = {3, 4, 5}; Vec A(sa, n), B(sb, n), C(sc, n); // run expression, this expression is longer:) A = B * F(C, B); for (int i = 0; i < n; ++i) { printf("%d:%f == %f * max(%f, %f)\\n", i, A.dptr[i], B.dptr[i], C.dptr[i], B.dptr[i]); } return 0; }

这段代码与上一小节代码的主要区别是模板类BinaryMapExp可以接受任意类型的二元操作符,要求是该操作符必须要实现一个Map函数。个人理解,第62行实现的那个F函数主要是为了编写代码方便,如果没有它,那么第69行就要写成BinaryMapExp(lhs.self(), rhs.self()); ,写起来就比较麻烦。其他的地方基本上与前一小节的代码差不多,稍微一看就能明白。
小结
综上所述,表达式模板基本工作原理包括以下几点:
  • 延迟计算,允许我们提前知道运算符和目标变量
  • 组合模板与递归计算,允许我们执行任意element-wise操作的复合表达式
  • 由于模板和内联的存在,表达式模板能够像编写for循环那样高效的实现element-wise的计算
MShadow中的表达式模板
MShadow中的表达式模板的原理与文中介绍的基本一致,但实现上还是有一些微小的差别:
  • 将计算代码和表达式构造相分离
    • 没有把Eval函数实现在Exp类中,而是根据表达式创建一个Plan类,并用它计算结果
    • 这样做的一个目的是减少Plan类中的私有变量数量,比如不需要知道数组的长度就可以计算结果
    • 另一个原因是CUDA kernel不能处理包含const reference的类
    • 这种设计值得商榷,但是目前很有用
  • 延迟计算支持复杂的表达式,例如矩阵点乘
    • 除了element-wise的表达式,MShadow还计算实现形如A = dot(B.T(), C)的语法糖。
  • 支持类型检查和数组长度检查
后记
【MShadow中的表达式模板】C++11中引入了移动构造函数,可用于保存重复分配的内存,从而消除了一些需要用到表达式模板的情况。然而,内存空间仍然至少需要被分配一次。
  • This only removes the need of expression template then expression generate space, say dst = A + B + C, dst does not contain space allocated before assignment. (这句话没有理解它的意思,先把原文放这里吧)
  • 如果想要保留一切都是预先分配的这种syntax,并且表达式无需内存分配即可执行(这就是MShadow所做的事情),那么仍然需要表达式模板。

    推荐阅读