TVM: VisitExpr流程分析

TVM源码中涉及到表达式遍历的地方，一般是适用VisitExpr接口进行，这个接口设计TVM的visitor模式，具体分析可参考：TVM:visitor设计模式

基类tvm::relay::ExprFunctor

适用visitor遍历的起点是调用VisitExpr接口，看下基类tvm::relay::ExprFunctor中这个方法的代码：

template <typename R, typename... Args>
class ExprFunctor<R(const Expr& n, Args...)> {
 private:
  using TSelf = ExprFunctor<R(const Expr& n, Args...)>;
  using FType = tvm::NodeFunctor<R(const ObjectRef& n, TSelf* self, Args...)>;

 public:
  ......
  /*!
   * \brief The functor call.
   * \param n The expression node.
   * \param args Additional arguments.
   * \return The result of the call
   */
  virtual R VisitExpr(const Expr& n, Args... args) {
    ICHECK(n.defined()) << "Found null pointer node while traversing AST. The previous pass may "
                           "have generated invalid data.";
    static FType vtable = InitVTable();
    return vtable(n, this, std::forward<Args>(args)...);
  }
  // Functions that can be overriden by subclass
  virtual R VisitExpr_(const ConstantNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const TupleNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
.......
    throw;
  }

 private:
  // initialize the vtable.
  static FType InitVTable() {
    FType vtable;
    // Set dispatch
    RELAY_EXPR_FUNCTOR_DISPATCH(ConstantNode);
    RELAY_EXPR_FUNCTOR_DISPATCH(TupleNode);
   ......
    return vtable;
  }
};

VisitExpr中调用InitVTable:

// initialize the vtable.
  static FType InitVTable() {
    FType vtable;
    // Set dispatch
    RELAY_EXPR_FUNCTOR_DISPATCH(ConstantNode);
    RELAY_EXPR_FUNCTOR_DISPATCH(TupleNode);
    RELAY_EXPR_FUNCTOR_DISPATCH(VarNode);
    RELAY_EXPR_FUNCTOR_DISPATCH(GlobalVarNode);
.....
    return vtable;
  }
  
#define RELAY_EXPR_FUNCTOR_DISPATCH(OP)                                                    \
  vtable.template set_dispatch<OP>([](const ObjectRef& n, TSelf* self, Args... args) {     \
    return self->VisitExpr_(static_cast<const OP*>(n.get()), std::forward<Args>(args)...); \
  });

template <typename R, typename... Args>
class NodeFunctor<R(const ObjectRef& n, Args...)> {
 private:
  /*! \brief internal function pointer type */
  typedef R (*FPointer)(const ObjectRef& n, Args...);
  /*! \brief refer to itself. */
  using TSelf = NodeFunctor<R(const ObjectRef& n, Args...)>;
  /*! \brief internal function table */
  std::vector<FPointer> func_;

 public:
......
  /*!
   * \brief set the dispacher for type TNode
   * \param f The function to be set.
   * \tparam TNode the type of Node to be dispatched.
   * \return reference to self.
   */
  template <typename TNode>
  TSelf& set_dispatch(FPointer f) {  // NOLINT(*)
    uint32_t tindex = TNode::RuntimeTypeIndex();
    if (func_.size() <= tindex) {
      func_.resize(tindex + 1, nullptr);
    }
    ICHECK(func_[tindex] == nullptr) << "Dispatch for " << TNode::_type_key << " is already set";
    func_[tindex] = f;
    return *this;
  }
  /*!
   * \brief unset the dispacher for type TNode
   *
   * \tparam TNode the type of Node to be dispatched.
   * \return reference to self.
   */
  template <typename TNode>
  TSelf& clear_dispatch() {  // NOLINT(*)
    uint32_t tindex = TNode::RuntimeTypeIndex();
    ICHECK_LT(tindex, func_.size()) << "clear_dispatch: index out of range";
    func_[tindex] = nullptr;
    return *this;
  }
};

InitVTable中调用NodeFunctor::set_dispatch接口，类型参数为tvm relay ir的各种表达式类型，传入set_dispatch的函数参数是lamad函数，lamad函数体中执行self->VisitExpr_()。self时传入的参数this，当从派生类中发起VisitExpr的时候，这个this将是派生类实例，而不是基类。

NodeFunctor::set_dispatch是在函数指针表func_中添加传入的lamad函数，表项索引为类型参数的id。

InitVTable在为所有类型都调用set_dispatch注册对应的visit调用后，返回了注册的NodeFunctor实例。而VisitExpr在调用InitVTable后return vtable(n, this, std::forward<Args>(args)...)。NodeFunctor中对()进行了运算符重载：

R operator()(const ObjectRef& n, Args... args) const {
    ICHECK(can_dispatch(n)) << "NodeFunctor calls un-registered function on type "
                            << n->GetTypeKey();
    return (*func_[n->type_index()])(n, std::forward<Args>(args)...);
  }

这里以传入的参数的类型id为索引，从func_表中获取对应的lamad函数体，并调用执行。也就是执行了类实例的VisitExpr_。因为一般来说发起VisitExpr调用的是以tvm::relay::ExprFunctor为基类，并在VisitExpr_中完成业务操作的类，所以这里VisitExpr_是调用的业务类中重载后的VisitExpr_方法。业务类对自己关注的类型的VisitExpr_进行重载，在其中完成自己的操作。

如果派生类不对各种类型重载VisitExpr_，就会调用到tvm::relay::ExprFunctor定义的VisitExpr_，抛出异常：

virtual R VisitExpr_(const ConstantNode* op, Args... args) { 
  return VisitExprDefault_(op, std::forward<Args>(args)...); 
};
 
virtual R VisitExpr_(const TupleNode* op, Args... args) { 
  return VisitExprDefault_(op, std::forward<Args>(args)...); 
};
 
virtual R VisitExpr_(const VarNode* op, Args... args) { 
  return VisitExprDefault_(op, std::forward<Args>(args)...); 
};
 
...
 
virtual R VisitExprDefault_(const Object* op, Args...) {  
 ::tvm::runtime::detail::LogFatal("/home/tvm/tvmsource/tvm/include/tvm/relay/expr_functor.h", 114).stream() << "Do not have a default for " << op->GetTypeKey();
    throw;
  }

派生类tvm::relay::ExprVisitor

ExprVisitor继承了ExprFunctor，并对VisitExpr和VisitExpr_进行了重载：

class ExprVisitor : public ::tvm::relay::ExprFunctor<void(const Expr& n)> {
 public:
  void VisitExpr(const Expr& expr) override;
  void VisitExpr_(const VarNode* op) override;
  void VisitExpr_(const GlobalVarNode* op) override;
  void VisitExpr_(const ConstantNode* op) override;
  void VisitExpr_(const TupleNode* op) override;
  void VisitExpr_(const FunctionNode* op) override;
  void VisitExpr_(const CallNode* op) override;
  void VisitExpr_(const LetNode* op) override;
  void VisitExpr_(const IfNode* op) override;
  void VisitExpr_(const OpNode* op) override;
  void VisitExpr_(const TupleGetItemNode* op) override;
  void VisitExpr_(const RefCreateNode* op) override;
  void VisitExpr_(const RefReadNode* op) override;
  void VisitExpr_(const RefWriteNode* op) override;
  void VisitExpr_(const ConstructorNode* op) override;
  void VisitExpr_(const MatchNode* op) override;
  virtual void VisitType(const Type& t);
  virtual void VisitClause(const Clause& c);
  virtual void VisitPattern(const Pattern& c);
  virtual void VisitSpan(const Span& span);
 
 protected:
  // Internal visiting counter
  std::unordered_map<const Object*, size_t> visit_counter_;
};
 
void ExprVisitor::VisitExpr(const Expr& expr) {
  auto it = visit_counter_.find(expr.get());
  if (it != visit_counter_.end()) {
    ++it->second;
  } else {
    using TParent = ExprFunctor<void(const Expr&)>;
    TParent::VisitExpr(expr);
    visit_counter_.insert({expr.get(), 1});
  }
}

visit_counter_表记录了每个表达式（注意不是每种）的访问历史。在VisitExpr中，如果发现该表达式已经访问过，则只是递增该表达式的访问计数，而不做实质的访问操作。如果发现表达式没有遍历过，则调用基类ExprFunctor的VisitExpr，进而调用到发起VisitExpr的某个派生类的VisitExpr_。

派生类tvm::relay::ExprMutator

派生类ExprMutator的定义跟ExprFunctor差不多：

class ExprMutator : public ::tvm::relay::ExprFunctor<Expr(const Expr&)> {
 public:
  /*!
   * \brief Mutate is alias for VisitExpr
   * \return expr.
   */
  Expr Mutate(const Expr& expr) { return this->VisitExpr(expr); }
  Expr VisitExpr(const Expr& expr) override;
  Expr VisitExpr_(const VarNode* op) override;
  Expr VisitExpr_(const ConstantNode* op) override;
  Expr VisitExpr_(const GlobalVarNode* op) override;
  Expr VisitExpr_(const OpNode* op) override;
  Expr VisitExpr_(const TupleNode* op) override;
  Expr VisitExpr_(const FunctionNode* op) override;
  Expr VisitExpr_(const CallNode* call_node) override;
  Expr VisitExpr_(const LetNode* op) override;
  Expr VisitExpr_(const IfNode* op) override;
  Expr VisitExpr_(const TupleGetItemNode* op) override;
  Expr VisitExpr_(const RefCreateNode* op) override;
  Expr VisitExpr_(const RefReadNode* op) override;
  Expr VisitExpr_(const RefWriteNode* op) override;
  Expr VisitExpr_(const ConstructorNode* op) override;
  Expr VisitExpr_(const MatchNode* op) override;
 
  /*!
   * \brief Used to visit the types inside of expressions.
   *
   * Can be overloaded to transform the types in arbitrary
   * ways, one way would be to define a sub-class of type
   * visitor for types which transform them appropriately.
   */
  virtual Type VisitType(const Type& t);
  virtual Clause VisitClause(const Clause& c);
  virtual Pattern VisitPattern(const Pattern& c);
 
 protected:
  /*! \brief Internal map used for memoization. */
  std::unordered_map<Expr, Expr, ObjectPtrHash, ObjectPtrEqual> memo_;
};
 
Expr ExprMutator::VisitExpr(const Expr& expr) {
  auto it = this->memo_.find(expr);
  if (it != this->memo_.end()) {
    return it->second;
  } else {
    Expr new_expr = ExprFunctor::VisitExpr(expr);
    memo_[expr] = new_expr;
    return new_expr;
  }
}

这里需要注意的是，ExprMutator的VisitExpr和VisitExpr_都是有返回值的，调用将返回遍历到的表达式，这样可以在VisitExpr_外对表达式做操作，比如说修改。

Codegen内存申请时的visitor模式使用

GraphPlanMemory分配流程中涉及的类关系图如下所示：

在该流程中分别从StorageAllocInit和StorageAllocator里面调用Run接口，Run接口调用VisitExpr，这个时候调用的是ExprVisitor::VisitExpr。而VisitExpr_则是调用的StorageAllocaBaseVisitor和DeviceAwareExprVisitor中重载的。

从这里也可以看到，ExprFunctor和ExprVisitor是纯粹作为visitor模式的实现而设计，具体的业务在各业务实现类中。

参考：
VisitExpr流程分析