[MXNet Code Analysis] Template Techniques in MShadow

Expression Template

This kind of technique, overall, follows the idea of Lazy Evaluation. It can be seen as using template type deduction at the C++ syntax level to statically build an expression tree for a matrix-computation expression to be evaluated, and then performing lazy evaluation of the expression tree only after the full expression has been constructed and needs to be assigned to a Tensor.

Basic Principles

The official MShadow guide provides an excellent tutorial. Here I will analyze it directly using my own line of reasoning.

:::C++
template<typename SubType>
struct Exp {
  // returns const reference of the actual type of this expression
  inline const SubType& self(void) const {
    return *static_cast<const SubType*>(this);
  }
};

First, we construct a base class Exp for expressions. It is a template base class whose only purpose is to provide a method that returns a reference to the actual derived type—essentially, automatically casting a base-class pointer to a derived-class pointer.

:::C++
template<typename OP, typename TLhs, typename TRhs>
struct BinaryMapExp: public Exp<BinaryMapExp<OP, TLhs, TRhs>> {
  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));
  }
};

Next we define the type for binary operators. Still following the usual template-generic pattern, this template takes three type parameters: the operation OP of the current binary computation (add/subtract/multiply/divide/max/min) and the left and right operand. Its implementation actually calls OP::Map to perform the binary operation.

For example, a multiplication operation type can be defined like this:

:::C++
struct mul {
  inline static float Map(float a, float b) {
    return a * b;
  }
};

At this point we already have enough tools to define abstract binary operations. However, with the template classes above, we can only define an abstract expression tree, but we still cannot evaluate it, because Exp does not implement an Eval method. It is only an abstract concept of an expression and cannot be used for evaluation. So now we introduce an evaluatable entity Vec that defines this method:

:::C++
struct Vec: public Exp<Vec> {
  int len;
  float* dptr;
  Vec(void) {}
  Vec(float *dptr, int len)
      : len(len), dptr(dptr) {}
  // here is where evaluation happens
  template<typename EType>
  inline Vec& operator=(const Exp<EType>& src_) {
    const EType &src = 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];
  }
};

Add helper template functions and operator overloads as syntactic sugar, so expressions read more naturally:

:::C++
template<typename OP, typename TLhs, typename TRhs>
inline BinaryMapExp<OP, TLhs, TRhs>
F(const Exp<TLhs>& lhs, const Exp<TRhs>& rhs) {
  return BinaryMapExp<OP, TLhs, TRhs>(lhs.self(), rhs.self());
}

template<typename TLhs, typename TRhs>
inline BinaryMapExp<mul, TLhs, TRhs>
operator*(const Exp<TLhs>& lhs, const Exp<TRhs>& rhs) {
  return F<mul>(lhs, rhs);
}

Finally, we can write vector computation expressions using the following concise syntax, and we do not need to allocate any extra space at all, because parsing the expression and fusing operations are both completed during static deduction.

:::C++
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<maximum>(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;
}

With optimizations enabled, the compiled binary does not include things like the object constructions above. It is more or less just a for loop computing the result, and the loop may even be unrolled directly. If we bring out a disassembler, we will find that the code above is actually compiled into the following pseudocode:

:::C++
int _main() {
    xmm0 = intrinsic_movsd(xmm0, *0x100000f60);
    xmm2 = intrinsic_movsd(xmm2, *0x100000f68);
    xmm1 = intrinsic_movsd(xmm1, *0x100000f70);
    xmm3 = intrinsic_movapd(xmm3, xmm1);
    printf("%d:%f == %f * max(%f, %f)\n", 0x0, rdx, rcx, r8, r9);
    xmm0 = intrinsic_movsd(xmm0, *0x100000f78);
    xmm2 = intrinsic_movsd(xmm2, *0x100000f80);
    xmm1 = intrinsic_movsd(xmm1, *0x100000f68);
    xmm3 = intrinsic_movapd(xmm3, xmm1);
    printf("%d:%f == %f * max(%f, %f)\n", 0x1, rdx, rcx, r8, r9);
    intrinsic_movsd(xmm0, *0x100000f88);
    intrinsic_movsd(xmm2, *0x100000f90);
    intrinsic_movapd(xmm3, intrinsic_movsd(xmm1, *0x100000f80));
    printf("%d:%f == %f * max(%f, %f)\n", 0x2, rdx, rcx, r8, r9);
    return 0x0;
}

MShadow Implementation

Basic Logic

Shared CPU/GPU Code

Limitations of Template Expr

comments powered by Disqus
Published:
2017-07-10
Category:
Tag: