[PASS] Change IRVisitor interfaces to function override (#42)

* [PASS] Change IRVisitor interfaces to function override

* [PASS] Change IRMutator interfaces to overloadable function

include/tvm/ir_mutator.h

@@ -16,7 +16,8 @@ namespace ir {
 /*!
  * \brief a base class for mutator to iterative mutate the IR
  *
- *  This IRMutator is implemented via IRFunctor instead of Visitor Pattern.
+ *  This IRMutator is implemented via Visitor Pattern.
+ *  Also you can implement via IRFunctor.
  *  This enables easy extensions of possible new Node.
  *  It also makes changing return types easier.
  *
@@ -54,20 +55,91 @@ class IRMutator {
   static FMutateStmt& vtable_stmt();  // NOLINT(*)
   // Set of overloadable functions
   // The underscore allows Mutate not to be shadowed by inheritance
+  virtual Stmt Mutate_(const Variable* op, const Stmt& s);
   virtual Stmt Mutate_(const LetStmt* op, const Stmt& s);
   virtual Stmt Mutate_(const AttrStmt* op, const Stmt& s);
+  virtual Stmt Mutate_(const IfThenElse* op, const Stmt& s);
   virtual Stmt Mutate_(const For* op, const Stmt& s);
-  virtual Stmt Mutate_(const Provide* op, const Stmt& s);
   virtual Stmt Mutate_(const Allocate* op, const Stmt& s);
-  virtual Stmt Mutate_(const Realize* op, const Stmt& s);
+  virtual Stmt Mutate_(const Load* op, const Stmt& s);
   virtual Stmt Mutate_(const Store* op, const Stmt& s);
+  virtual Stmt Mutate_(const Let* op, const Stmt& s);
   virtual Stmt Mutate_(const Free* op, const Stmt& s);
-  virtual Stmt Mutate_(const IfThenElse* op, const Stmt& s);
+  virtual Stmt Mutate_(const Call* op, const Stmt& s);
+  virtual Stmt Mutate_(const Add* op, const Stmt& e);
+  virtual Stmt Mutate_(const Sub* op, const Stmt& e);
+  virtual Stmt Mutate_(const Mul* op, const Stmt& e);
+  virtual Stmt Mutate_(const Div* op, const Stmt& e);
+  virtual Stmt Mutate_(const Mod* op, const Stmt& e);
+  virtual Stmt Mutate_(const Min* op, const Stmt& e);
+  virtual Stmt Mutate_(const Max* op, const Stmt& e);
+  virtual Stmt Mutate_(const EQ* op, const Stmt& e);
+  virtual Stmt Mutate_(const NE* op, const Stmt& e);
+  virtual Stmt Mutate_(const LT* op, const Stmt& e);
+  virtual Stmt Mutate_(const LE* op, const Stmt& e);
+  virtual Stmt Mutate_(const GT* op, const Stmt& e);
+  virtual Stmt Mutate_(const GE* op, const Stmt& e);
+  virtual Stmt Mutate_(const And* op, const Stmt& e);
+  virtual Stmt Mutate_(const Or* op, const Stmt& e);
+  virtual Stmt Mutate_(const Reduce* op, const Stmt& s);
+  virtual Stmt Mutate_(const Cast* op, const Stmt& s);
+  virtual Stmt Mutate_(const Not* op, const Stmt& s);
+  virtual Stmt Mutate_(const Select* op, const Stmt& s);
+  virtual Stmt Mutate_(const Ramp* op, const Stmt& s);
+  virtual Stmt Mutate_(const Broadcast* op, const Stmt& e);
+  virtual Stmt Mutate_(const AssertStmt* op, const Stmt& e);
+  virtual Stmt Mutate_(const ProducerConsumer* op, const Stmt& e);
+  virtual Stmt Mutate_(const Provide* op, const Stmt& e);
+  virtual Stmt Mutate_(const Realize* op, const Stmt& s);
   virtual Stmt Mutate_(const Block* op, const Stmt& s);
-  virtual Expr Mutate_(const Call* op, const Expr& e);
-  virtual Expr Mutate_(const Load* op, const Expr& s);
+  virtual Stmt Mutate_(const Evaluate* op, const Stmt& e);
+  virtual Stmt Mutate_(const IntImm* op, const Stmt& e);
+  virtual Stmt Mutate_(const UIntImm* op, const Stmt& e);
+  virtual Stmt Mutate_(const FloatImm* op, const Stmt& e);
+  virtual Stmt Mutate_(const StringImm* op, const Stmt& e);
+
   virtual Expr Mutate_(const Variable* op, const Expr& e);
+  virtual Expr Mutate_(const LetStmt* op, const Expr& e);
+  virtual Expr Mutate_(const AttrStmt* op, const Expr& e);
+  virtual Expr Mutate_(const IfThenElse* op, const Expr& e);
+  virtual Expr Mutate_(const For* op, const Expr& e);
+  virtual Expr Mutate_(const Allocate* op, const Expr& e);
+  virtual Expr Mutate_(const Load* op, const Expr& e);
+  virtual Expr Mutate_(const Store* op, const Expr& e);
   virtual Expr Mutate_(const Let* op, const Expr& e);
+  virtual Expr Mutate_(const Free* op, const Expr& e);
+  virtual Expr Mutate_(const Call* op, const Expr& e);
+  virtual Expr Mutate_(const Add* op, const Expr& e);
+  virtual Expr Mutate_(const Sub* op, const Expr& e);
+  virtual Expr Mutate_(const Mul* op, const Expr& e);
+  virtual Expr Mutate_(const Div* op, const Expr& e);
+  virtual Expr Mutate_(const Mod* op, const Expr& e);
+  virtual Expr Mutate_(const Min* op, const Expr& e);
+  virtual Expr Mutate_(const Max* op, const Expr& e);
+  virtual Expr Mutate_(const EQ* op, const Expr& e);
+  virtual Expr Mutate_(const NE* op, const Expr& e);
+  virtual Expr Mutate_(const LT* op, const Expr& e);
+  virtual Expr Mutate_(const LE* op, const Expr& e);
+  virtual Expr Mutate_(const GT* op, const Expr& e);
+  virtual Expr Mutate_(const GE* op, const Expr& e);
+  virtual Expr Mutate_(const And* op, const Expr& e);
+  virtual Expr Mutate_(const Or* op, const Expr& e);
+  virtual Expr Mutate_(const Reduce* op, const Expr& e);
+  virtual Expr Mutate_(const Cast* op, const Expr& e);
+  virtual Expr Mutate_(const Not* op, const Expr& e);
+  virtual Expr Mutate_(const Select* op, const Expr& e);
+  virtual Expr Mutate_(const Ramp* op, const Expr& e);
+  virtual Expr Mutate_(const Broadcast* op, const Expr& e);
+  virtual Expr Mutate_(const AssertStmt* op, const Expr& e);
+  virtual Expr Mutate_(const ProducerConsumer* op, const Expr& e);
+  virtual Expr Mutate_(const Provide* op, const Expr& e);
+  virtual Expr Mutate_(const Realize* op, const Expr& e);
+  virtual Expr Mutate_(const Block* op, const Expr& e);
+  virtual Expr Mutate_(const Evaluate* op, const Expr& e);
+  virtual Expr Mutate_(const IntImm* op, const Expr& e);
+  virtual Expr Mutate_(const UIntImm* op, const Expr& e);
+  virtual Expr Mutate_(const FloatImm* op, const Expr& e);
+  virtual Expr Mutate_(const StringImm* op, const Expr& e);
 };
 
 /*!

include/tvm/ir_visitor.h

@@ -36,16 +36,47 @@ class IRVisitor {
   static FVisit& vtable();
   // overloadable visit function.
   virtual void Visit_(const Variable* op);
-  virtual void Visit_(const AttrStmt* op);
   virtual void Visit_(const LetStmt* op);
+  virtual void Visit_(const AttrStmt* op);
+  virtual void Visit_(const IfThenElse* op);
   virtual void Visit_(const For* op);
   virtual void Visit_(const Allocate* op);
-  virtual void Visit_(const IfThenElse* op);
   virtual void Visit_(const Load* op);
   virtual void Visit_(const Store* op);
   virtual void Visit_(const Let* op);
   virtual void Visit_(const Free* op);
   virtual void Visit_(const Call* op);
+  virtual void Visit_(const Add* op);
+  virtual void Visit_(const Sub* op);
+  virtual void Visit_(const Mul* op);
+  virtual void Visit_(const Div* op);
+  virtual void Visit_(const Mod* op);
+  virtual void Visit_(const Min* op);
+  virtual void Visit_(const Max* op);
+  virtual void Visit_(const EQ* op);
+  virtual void Visit_(const NE* op);
+  virtual void Visit_(const LT* op);
+  virtual void Visit_(const LE* op);
+  virtual void Visit_(const GT* op);
+  virtual void Visit_(const GE* op);
+  virtual void Visit_(const And* op);
+  virtual void Visit_(const Or* op);
+  virtual void Visit_(const Reduce* op);
+  virtual void Visit_(const Cast* op);
+  virtual void Visit_(const Not* op);
+  virtual void Visit_(const Select* op);
+  virtual void Visit_(const Ramp* op);
+  virtual void Visit_(const Broadcast* op);
+  virtual void Visit_(const AssertStmt* op);
+  virtual void Visit_(const ProducerConsumer* op);
+  virtual void Visit_(const Provide* op);
+  virtual void Visit_(const Realize* op);
+  virtual void Visit_(const Block* op);
+  virtual void Visit_(const Evaluate* op);
+  virtual void Visit_(const IntImm* op);
+  virtual void Visit_(const UIntImm* op);
+  virtual void Visit_(const FloatImm* op);
+  virtual void Visit_(const StringImm* op);
 };
 
 /*!

src/pass/ir_mutator.cc

@@ -16,11 +16,6 @@ IRMutator::FMutateStmt& IRMutator::vtable_stmt() {  // NOLINT(*)
   static FMutateStmt inst; return inst;
 }
 
-// const expr
-inline Expr ReturnSelfExpr(const NodeRef&, const Expr& e, IRMutator*) {
-  return e;
-}
-
 inline Array<Expr> MutateArray(Array<Expr> arr, IRMutator *m) {
   std::vector<Expr> new_arr(arr.size());
   bool changed = false;
@@ -58,47 +53,33 @@ inline Array<IterVar> MutateIterVarArr(Array<IterVar> rdom, IRMutator *m) {
   }
 }
 
+
+// Mutate Stmt
+
 #define DISPATCH_TO_MUTATE_STMT(OP)                                 \
   set_dispatch<OP>([](const OP* op, const Stmt& s, IRMutator* m) {  \
       return m->Mutate_(op, s);                                     \
     })
 
-#define DISPATCH_TO_MUTATE_EXPR(OP)                                 \
-  set_dispatch<OP>([](const OP* op, const Expr& e, IRMutator* m) {  \
-      return m->Mutate_(op, e);                                     \
-    })
-
-TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_stmt)
-.DISPATCH_TO_MUTATE_STMT(LetStmt)
-.DISPATCH_TO_MUTATE_STMT(AttrStmt)
-.DISPATCH_TO_MUTATE_STMT(Provide)
-.DISPATCH_TO_MUTATE_STMT(Realize)
-.DISPATCH_TO_MUTATE_STMT(Store)
-.DISPATCH_TO_MUTATE_STMT(IfThenElse)
-.DISPATCH_TO_MUTATE_STMT(For)
-.DISPATCH_TO_MUTATE_STMT(Allocate)
-.DISPATCH_TO_MUTATE_STMT(Block)
-.DISPATCH_TO_MUTATE_STMT(Free);
-
-Stmt IRMutator::Mutate_(const LetStmt *op, const Stmt& s) {
+Stmt IRMutator::Mutate_(const AttrStmt* op, const Stmt& s) {
   Expr value = this->Mutate(op->value);
   Stmt body = this->Mutate(op->body);
   if (value.same_as(op->value) &&
       body.same_as(op->body)) {
     return s;
   } else {
-    return LetStmt::make(op->var, value, body);
+    return AttrStmt::make(op->node, op->type_key, value, body);
   }
 }
 
-Stmt IRMutator::Mutate_(const AttrStmt* op, const Stmt& s) {
+Stmt IRMutator::Mutate_(const LetStmt *op, const Stmt& s) {
   Expr value = this->Mutate(op->value);
   Stmt body = this->Mutate(op->body);
   if (value.same_as(op->value) &&
       body.same_as(op->body)) {
     return s;
   } else {
-    return AttrStmt::make(op->node, op->type_key, value, body);
+    return LetStmt::make(op->var, value, body);
   }
 }
 
@@ -143,6 +124,36 @@ Stmt IRMutator::Mutate_(const Allocate* op, const Stmt& s) {
   }
 }
 
+Stmt IRMutator::Mutate_(const IfThenElse *op, const Stmt& s) {
+  Expr condition = this->Mutate(op->condition);
+  Stmt then_case = this->Mutate(op->then_case);
+  Stmt else_case;
+  if (else_case.defined()) {
+    else_case = this->Mutate(op->else_case);
+  }
+  if (condition.same_as(op->condition) &&
+      then_case.same_as(op->then_case) &&
+      else_case.same_as(op->else_case)) {
+    return s;
+  } else {
+    return IfThenElse::make(condition, then_case, else_case);
+  }
+}
+
+Stmt IRMutator::Mutate_(const Load *op, const Stmt& s) {
+  return s;
+}
+
+Stmt IRMutator::Mutate_(const Store *op, const Stmt& s) {
+  Expr value = this->Mutate(op->value);
+  Expr index = this->Mutate(op->index);
+  if (value.same_as(op->value) && index.same_as(op->index)) {
+    return s;
+  } else {
+    return Store::make(op->buffer_var, value, index);
+  }
+}
+
 Stmt IRMutator::Mutate_(const Provide* op, const Stmt& s) {
   auto new_args = MutateArray(op->args, this);
   auto new_value = this->Mutate(op->value);
@@ -183,61 +194,135 @@ Stmt IRMutator::Mutate_(const Realize* op, const Stmt& s) {
   }
 }
 
-Stmt IRMutator::Mutate_(const Store *op, const Stmt& s) {
-  Expr value = this->Mutate(op->value);
-  Expr index = this->Mutate(op->index);
-  if (value.same_as(op->value) && index.same_as(op->index)) {
+Stmt IRMutator::Mutate_(const Block* op, const Stmt& s) {
+  Stmt first = this->Mutate(op->first);
+  Stmt rest = this->Mutate(op->rest);
+  if (first.same_as(op->first) &&
+      rest.same_as(op->rest)) {
     return s;
   } else {
-    return Store::make(op->buffer_var, value, index);
+    return Block::make(first, rest);
   }
 }
 
-Stmt IRMutator::Mutate_(const Free *op, const Stmt& s) {
+Stmt IRMutator::Mutate_(const AssertStmt *op, const Stmt& s) {
+  Expr condition = this->Mutate(op->condition);
+  Expr message = this->Mutate(op->message);
+
+  if (condition.same_as(op->condition) && message.same_as(op->message)) {
     return s;
+  } else {
+    return AssertStmt::make(condition, message);
+  }
 }
 
-Stmt IRMutator::Mutate_(const IfThenElse *op, const Stmt& s) {
-  Expr condition = this->Mutate(op->condition);
-  Stmt then_case = this->Mutate(op->then_case);
-  Stmt else_case;
-  if (else_case.defined()) {
-    else_case = this->Mutate(op->else_case);
-  }
-  if (condition.same_as(op->condition) &&
-      then_case.same_as(op->then_case) &&
-      else_case.same_as(op->else_case)) {
+Stmt IRMutator::Mutate_(const ProducerConsumer *op, const Stmt& s) {
+  Stmt body = this->Mutate(op->body);
+  if (body.same_as(op->body)) {
     return s;
   } else {
-    return IfThenElse::make(condition, then_case, else_case);
+    return ProducerConsumer::make(op->func, op->is_producer, body);
   }
 }
 
-Stmt IRMutator::Mutate_(const Block* op, const Stmt& s) {
-  Stmt first = this->Mutate(op->first);
-  Stmt rest = this->Mutate(op->rest);
-  if (first.same_as(op->first) &&
-      rest.same_as(op->rest)) {
+Stmt IRMutator::Mutate_(const Evaluate *op, const Stmt& s) {
+  Expr v = this->Mutate(op->value);
+  if (v.same_as(op->value)) {
     return s;
   } else {
-    return Block::make(first, rest);
+    return Evaluate::make(v);
   }
 }
 
-TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_expr)
-.DISPATCH_TO_MUTATE_EXPR(Call)
-.DISPATCH_TO_MUTATE_EXPR(Let)
-.DISPATCH_TO_MUTATE_EXPR(Load)
-.DISPATCH_TO_MUTATE_EXPR(Variable);
+#define DEFINE_OP_RETURN_SELF_STMT_MUTATE_(OP)              \
+  Stmt IRMutator::Mutate_(const OP *op, const Stmt& s) {    \
+    return s;                                               \
+  }
 
-Expr IRMutator::Mutate_(const Call* op, const Expr& e) {
-  auto new_args = MutateArray(op->args, this);
-  if (op->args.same_as(new_args)) {
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(Variable)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(Let)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(Free)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(Call)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(Add)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(Sub)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(Mul)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(Div)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(Mod)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(Min)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(Max)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(EQ)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(NE)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(LT)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(LE)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(GT)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(GE)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(And)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(Or)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(Reduce)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(Cast)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(Not)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(Select)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(Ramp)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(Broadcast)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(IntImm)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(UIntImm)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(FloatImm)
+DEFINE_OP_RETURN_SELF_STMT_MUTATE_(StringImm)
+
+TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_stmt)
+.DISPATCH_TO_MUTATE_STMT(Variable)
+.DISPATCH_TO_MUTATE_STMT(LetStmt)
+.DISPATCH_TO_MUTATE_STMT(AttrStmt)
+.DISPATCH_TO_MUTATE_STMT(IfThenElse)
+.DISPATCH_TO_MUTATE_STMT(For)
+.DISPATCH_TO_MUTATE_STMT(Allocate)
+.DISPATCH_TO_MUTATE_STMT(Load)
+.DISPATCH_TO_MUTATE_STMT(Store)
+.DISPATCH_TO_MUTATE_STMT(Let)
+.DISPATCH_TO_MUTATE_STMT(Free)
+.DISPATCH_TO_MUTATE_STMT(Call)
+.DISPATCH_TO_MUTATE_STMT(Add)
+.DISPATCH_TO_MUTATE_STMT(Sub)
+.DISPATCH_TO_MUTATE_STMT(Mul)
+.DISPATCH_TO_MUTATE_STMT(Div)
+.DISPATCH_TO_MUTATE_STMT(Mod)
+.DISPATCH_TO_MUTATE_STMT(Min)
+.DISPATCH_TO_MUTATE_STMT(Max)
+.DISPATCH_TO_MUTATE_STMT(EQ)
+.DISPATCH_TO_MUTATE_STMT(NE)
+.DISPATCH_TO_MUTATE_STMT(LT)
+.DISPATCH_TO_MUTATE_STMT(LE)
+.DISPATCH_TO_MUTATE_STMT(GT)
+.DISPATCH_TO_MUTATE_STMT(GE)
+.DISPATCH_TO_MUTATE_STMT(And)
+.DISPATCH_TO_MUTATE_STMT(Or)
+.DISPATCH_TO_MUTATE_STMT(Reduce)
+.DISPATCH_TO_MUTATE_STMT(Cast)
+.DISPATCH_TO_MUTATE_STMT(Not)
+.DISPATCH_TO_MUTATE_STMT(Select)
+.DISPATCH_TO_MUTATE_STMT(Ramp)
+.DISPATCH_TO_MUTATE_STMT(Broadcast)
+.DISPATCH_TO_MUTATE_STMT(AssertStmt)
+.DISPATCH_TO_MUTATE_STMT(ProducerConsumer)
+.DISPATCH_TO_MUTATE_STMT(Provide)
+.DISPATCH_TO_MUTATE_STMT(Realize)
+.DISPATCH_TO_MUTATE_STMT(Block)
+.DISPATCH_TO_MUTATE_STMT(Evaluate)
+.DISPATCH_TO_MUTATE_STMT(IntImm)
+.DISPATCH_TO_MUTATE_STMT(UIntImm)
+.DISPATCH_TO_MUTATE_STMT(FloatImm)
+.DISPATCH_TO_MUTATE_STMT(StringImm);
+
+
+// Mutate Expr
+
+#define DISPATCH_TO_MUTATE_EXPR(OP)                                 \
+  set_dispatch<OP>([](const OP* op, const Expr& e, IRMutator* m) {  \
+      return m->Mutate_(op, e);                                     \
+    })
+
+Expr IRMutator::Mutate_(const Variable *op, const Expr& e) {
   return e;
-  } else {
-    return Call::make(op->type, op->name, new_args, op->call_type,
-                      op->func, op->value_index);
-  }
 }
 
 Expr IRMutator::Mutate_(const Load *op, const Expr& e) {
@@ -249,11 +334,6 @@ Expr IRMutator::Mutate_(const Load *op, const Expr& e) {
   }
 }
 
-
-Expr IRMutator::Mutate_(const Variable *op, const Expr& e) {
-  return e;
-}
-
 Expr IRMutator::Mutate_(const Let *op, const Expr& e) {
   Expr value = this->Mutate(op->value);
   Expr body = this->Mutate(op->body);
@@ -265,77 +345,77 @@ Expr IRMutator::Mutate_(const Let *op, const Expr& e) {
   }
 }
 
-TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_expr)
-.set_dispatch<Reduce>([](const Reduce* op, const Expr& e, IRMutator* m) {
-    Array<IterVar> new_axis  = MutateIterVarArr(op->axis, m);
-    Expr new_source  = m->Mutate(op->source);
-    if (op->axis.same_as(new_axis) &&
-        op->source.same_as(new_source)) {
+Expr IRMutator::Mutate_(const Call* op, const Expr& e) {
+  auto new_args = MutateArray(op->args, this);
+  if (op->args.same_as(new_args)) {
     return e;
   } else {
-      return Reduce::make(op->op, new_source, new_axis);
+    return Call::make(op->type, op->name, new_args, op->call_type,
+                      op->func, op->value_index);
   }
-  });
+}
 
-TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_expr)
-.set_dispatch<IntImm>(ReturnSelfExpr)
-.set_dispatch<UIntImm>(ReturnSelfExpr)
-.set_dispatch<FloatImm>(ReturnSelfExpr)
-.set_dispatch<StringImm>(ReturnSelfExpr);
+#define DEFINE_BIOP_EXPR_MUTATE_(OP)                        \
+  Expr IRMutator::Mutate_(const OP* op, const Expr& e) {    \
+    Expr a = this->Mutate(op->a);                           \
+    Expr b = this->Mutate(op->b);                           \
+    if (a.same_as(op->a) &&                                 \
+        b.same_as(op->b)) {                                 \
+      return e;                                             \
+    } else {                                                \
+      return OP::make(a, b);                                 \
+    }                                                       \
+  }
 
-TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_expr)
-.set_dispatch<Cast>([](const Cast* op, const Expr& e, IRMutator* m) {
-    Expr value = m->Mutate(op->value);
-    if (value.same_as(op->value)) {
+DEFINE_BIOP_EXPR_MUTATE_(Add)
+DEFINE_BIOP_EXPR_MUTATE_(Sub)
+DEFINE_BIOP_EXPR_MUTATE_(Mul)
+DEFINE_BIOP_EXPR_MUTATE_(Div)
+DEFINE_BIOP_EXPR_MUTATE_(Mod)
+DEFINE_BIOP_EXPR_MUTATE_(Min)
+DEFINE_BIOP_EXPR_MUTATE_(Max)
+DEFINE_BIOP_EXPR_MUTATE_(EQ)
+DEFINE_BIOP_EXPR_MUTATE_(NE)
+DEFINE_BIOP_EXPR_MUTATE_(LT)
+DEFINE_BIOP_EXPR_MUTATE_(LE)
+DEFINE_BIOP_EXPR_MUTATE_(GT)
+DEFINE_BIOP_EXPR_MUTATE_(GE)
+DEFINE_BIOP_EXPR_MUTATE_(And)
+DEFINE_BIOP_EXPR_MUTATE_(Or)
+
+Expr IRMutator::Mutate_(const Reduce *op, const Expr& e) {
+  Array<IterVar> new_axis  = MutateIterVarArr(op->axis, this);
+  Expr new_source = this->Mutate(op->source);
+  if (op->axis.same_as(new_axis) &&
+      op->source.same_as(new_source)) {
     return e;
   } else {
-      return Cast::make(op->type, value);
+    return Reduce::make(op->op, new_source, new_axis);
   }
-  });
+}
 
-// binary operator
-template<typename T>
-inline Expr Binary(const T* op, const Expr& e, IRMutator* m) {
-  Expr a = m->Mutate(op->a);
-  Expr b = m->Mutate(op->b);
-  if (a.same_as(op->a) &&
-      b.same_as(op->b)) {
+Expr IRMutator::Mutate_(const Cast *op, const Expr& e) {
+  Expr value = this->Mutate(op->value);
+  if (value.same_as(op->value)) {
     return e;
   } else {
-    return T::make(a, b);
+    return Cast::make(op->type, value);
   }
 }
 
-TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_expr)
-.set_dispatch<Add>(Binary<Add>)
-.set_dispatch<Sub>(Binary<Sub>)
-.set_dispatch<Mul>(Binary<Mul>)
-.set_dispatch<Div>(Binary<Div>)
-.set_dispatch<Mod>(Binary<Mod>)
-.set_dispatch<Min>(Binary<Min>)
-.set_dispatch<Max>(Binary<Max>)
-.set_dispatch<EQ>(Binary<EQ>)
-.set_dispatch<NE>(Binary<NE>)
-.set_dispatch<LT>(Binary<LT>)
-.set_dispatch<LE>(Binary<LE>)
-.set_dispatch<GT>(Binary<GT>)
-.set_dispatch<GE>(Binary<GE>)
-.set_dispatch<And>(Binary<And>)
-.set_dispatch<Or>(Binary<Or>);
-
-TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_expr)
-.set_dispatch<Not>([](const Not* op, const Expr& e, IRMutator* m) {
-    Expr a = m->Mutate(op->a);
+Expr IRMutator::Mutate_(const Not *op, const Expr& e) {
+  Expr a = this->Mutate(op->a);
   if (a.same_as(op->a)) {
     return e;
   } else {
     return Not::make(a);
   }
-  })
-.set_dispatch<Select>([](const Select *op, const Expr& e, IRMutator* m) {
-    Expr cond = m->Mutate(op->condition);
-    Expr t = m->Mutate(op->true_value);
-    Expr f = m->Mutate(op->false_value);
+}
+
+Expr IRMutator::Mutate_(const Select *op, const Expr& e) {
+  Expr cond = this->Mutate(op->condition);
+  Expr t = this->Mutate(op->true_value);
+  Expr f = this->Mutate(op->false_value);
   if (cond.same_as(op->condition) &&
       t.same_as(op->true_value) &&
       f.same_as(op->false_value)) {
@@ -343,52 +423,94 @@ TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_expr)
   } else {
     return Select::make(cond, t, f);
   }
-  })
-.set_dispatch<Ramp>([](const Ramp *op, const Expr& e, IRMutator* m) {
-    Expr base = m->Mutate(op->base);
-    Expr stride = m->Mutate(op->stride);
+}
+
+Expr IRMutator::Mutate_(const Ramp *op, const Expr& e) {
+  Expr base = this->Mutate(op->base);
+  Expr stride = this->Mutate(op->stride);
   if (base.same_as(op->base) &&
       stride.same_as(op->stride)) {
     return e;
   } else {
     return Ramp::make(base, stride, op->lanes);
   }
-  })
-.set_dispatch<Broadcast>([](const Broadcast *op, const Expr& e, IRMutator* m) {
-    Expr value = m->Mutate(op->value);
+}
+
+Expr IRMutator::Mutate_(const Broadcast *op, const Expr& e) {
+  Expr value = this->Mutate(op->value);
   if (value.same_as(op->value)) {
     return e;
   } else {
     return Broadcast::make(value, op->lanes);
   }
-  });
-
-TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_stmt)
-.set_dispatch<AssertStmt>([](const AssertStmt *op, const Stmt& s, IRMutator* m) {
-    Expr condition = m->Mutate(op->condition);
-    Expr message = m->Mutate(op->message);
+}
 
-    if (condition.same_as(op->condition) && message.same_as(op->message)) {
-      return s;
-    } else {
-      return AssertStmt::make(condition, message);
-    }
-  })
-.set_dispatch<ProducerConsumer>([](const ProducerConsumer *op, const Stmt& s, IRMutator* m) {
-    Stmt body = m->Mutate(op->body);
-    if (body.same_as(op->body)) {
-      return s;
-    } else {
-      return ProducerConsumer::make(op->func, op->is_producer, body);
-    }
-  })
-.set_dispatch<Evaluate>([](const Evaluate *op, const Stmt& s, IRMutator* m) {
-    Expr v = m->Mutate(op->value);
-    if (v.same_as(op->value)) {
-      return s;
-    } else {
-      return Evaluate::make(v);
+#define DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(OP)              \
+  Expr IRMutator::Mutate_(const OP *op, const Expr& e) {    \
+    return e;                                               \
   }
-  });
+
+DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(LetStmt)
+DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(AttrStmt)
+DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(For)
+DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(IfThenElse)
+DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(Allocate)
+DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(Store)
+DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(Free)
+DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(AssertStmt)
+DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(ProducerConsumer)
+DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(Provide)
+DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(Realize)
+DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(Block)
+DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(Evaluate)
+DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(IntImm)
+DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(UIntImm)
+DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(FloatImm)
+DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(StringImm)
+
+TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_expr)
+.DISPATCH_TO_MUTATE_EXPR(Variable)
+.DISPATCH_TO_MUTATE_EXPR(LetStmt)
+.DISPATCH_TO_MUTATE_EXPR(AttrStmt)
+.DISPATCH_TO_MUTATE_EXPR(IfThenElse)
+.DISPATCH_TO_MUTATE_EXPR(For)
+.DISPATCH_TO_MUTATE_EXPR(Allocate)
+.DISPATCH_TO_MUTATE_EXPR(Load)
+.DISPATCH_TO_MUTATE_EXPR(Store)
+.DISPATCH_TO_MUTATE_EXPR(Let)
+.DISPATCH_TO_MUTATE_EXPR(Free)
+.DISPATCH_TO_MUTATE_EXPR(Call)
+.DISPATCH_TO_MUTATE_EXPR(Add)
+.DISPATCH_TO_MUTATE_EXPR(Sub)
+.DISPATCH_TO_MUTATE_EXPR(Mul)
+.DISPATCH_TO_MUTATE_EXPR(Div)
+.DISPATCH_TO_MUTATE_EXPR(Mod)
+.DISPATCH_TO_MUTATE_EXPR(Min)
+.DISPATCH_TO_MUTATE_EXPR(Max)
+.DISPATCH_TO_MUTATE_EXPR(EQ)
+.DISPATCH_TO_MUTATE_EXPR(NE)
+.DISPATCH_TO_MUTATE_EXPR(LT)
+.DISPATCH_TO_MUTATE_EXPR(LE)
+.DISPATCH_TO_MUTATE_EXPR(GT)
+.DISPATCH_TO_MUTATE_EXPR(GE)
+.DISPATCH_TO_MUTATE_EXPR(And)
+.DISPATCH_TO_MUTATE_EXPR(Or)
+.DISPATCH_TO_MUTATE_EXPR(Reduce)
+.DISPATCH_TO_MUTATE_EXPR(Cast)
+.DISPATCH_TO_MUTATE_EXPR(Not)
+.DISPATCH_TO_MUTATE_EXPR(Select)
+.DISPATCH_TO_MUTATE_EXPR(Ramp)
+.DISPATCH_TO_MUTATE_EXPR(Broadcast)
+.DISPATCH_TO_MUTATE_EXPR(AssertStmt)
+.DISPATCH_TO_MUTATE_EXPR(ProducerConsumer)
+.DISPATCH_TO_MUTATE_EXPR(Provide)
+.DISPATCH_TO_MUTATE_EXPR(Realize)
+.DISPATCH_TO_MUTATE_EXPR(Block)
+.DISPATCH_TO_MUTATE_EXPR(Evaluate)
+.DISPATCH_TO_MUTATE_EXPR(IntImm)
+.DISPATCH_TO_MUTATE_EXPR(UIntImm)
+.DISPATCH_TO_MUTATE_EXPR(FloatImm)
+.DISPATCH_TO_MUTATE_EXPR(StringImm);
+
 }  // namespace ir
 }  // namespace tvm

src/pass/ir_visitor.cc

@@ -34,9 +34,6 @@ IRVisitor::FVisit& IRVisitor::vtable() {  // NOLINT(*)
   static FVisit inst; return inst;
 }
 
-void NoOp(const NodeRef& n, IRVisitor* v) {
-}
-
 inline void VisitArray(const Array<Expr>& arr, IRVisitor* v) {
   for (size_t i = 0; i < arr.size(); i++) {
     v->Visit(arr[i]);
@@ -51,24 +48,6 @@ inline void VisitRDom(const Array<IterVar>& rdom, IRVisitor* v) {
   }
 }
 
-#define DISPATCH_TO_VISIT(OP)                       \
-  set_dispatch<OP>([](const OP* op, IRVisitor* v) { \
-      v->Visit_(op);                                \
-    })
-
-TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable)
-.DISPATCH_TO_VISIT(Variable)
-.DISPATCH_TO_VISIT(LetStmt)
-.DISPATCH_TO_VISIT(AttrStmt)
-.DISPATCH_TO_VISIT(IfThenElse)
-.DISPATCH_TO_VISIT(For)
-.DISPATCH_TO_VISIT(Allocate)
-.DISPATCH_TO_VISIT(Load)
-.DISPATCH_TO_VISIT(Store)
-.DISPATCH_TO_VISIT(Let)
-.DISPATCH_TO_VISIT(Call)
-.DISPATCH_TO_VISIT(Free);
-
 void IRVisitor::Visit_(const Variable* op) {}
 
 void IRVisitor::Visit_(const LetStmt *op) {
@@ -128,91 +107,146 @@ void IRVisitor::Visit_(const Call *op) {
   VisitArray(op->args, this);
 }
 
-TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable)
-.set_dispatch<Reduce>([](const Reduce* op, IRVisitor* v) {
-    VisitRDom(op->axis, v);
-    v->Visit(op->source);
-  })
-.set_dispatch<IntImm>(NoOp)
-.set_dispatch<UIntImm>(NoOp)
-.set_dispatch<FloatImm>(NoOp)
-.set_dispatch<StringImm>(NoOp);
+#define DEFINE_BINOP_VISIT_(OP)                     \
+  void IRVisitor::Visit_(const OP* op) {            \
+    this->Visit(op->a);                             \
+    this->Visit(op->b);                             \
+  }
 
-TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable)
-.set_dispatch<Cast>([](const Cast* op, IRVisitor* v) {
-    v->Visit(op->value);
-  });
+DEFINE_BINOP_VISIT_(Add)
+DEFINE_BINOP_VISIT_(Sub)
+DEFINE_BINOP_VISIT_(Mul)
+DEFINE_BINOP_VISIT_(Div)
+DEFINE_BINOP_VISIT_(Mod)
+DEFINE_BINOP_VISIT_(Min)
+DEFINE_BINOP_VISIT_(Max)
+DEFINE_BINOP_VISIT_(EQ)
+DEFINE_BINOP_VISIT_(NE)
+DEFINE_BINOP_VISIT_(LT)
+DEFINE_BINOP_VISIT_(LE)
+DEFINE_BINOP_VISIT_(GT)
+DEFINE_BINOP_VISIT_(GE)
+DEFINE_BINOP_VISIT_(And)
+DEFINE_BINOP_VISIT_(Or)
+
+void IRVisitor::Visit_(const Reduce* op) {
+  VisitRDom(op->axis, this);
+  this->Visit(op->source);
+}
 
-// binary operator
-template<typename T>
-inline void Binary(const T* op, IRVisitor* v) {
-  v->Visit(op->a);
-  v->Visit(op->b);
+void IRVisitor::Visit_(const Cast* op) {
+  this->Visit(op->value);
 }
 
-TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable)
-.set_dispatch<Add>(Binary<Add>)
-.set_dispatch<Sub>(Binary<Sub>)
-.set_dispatch<Mul>(Binary<Mul>)
-.set_dispatch<Div>(Binary<Div>)
-.set_dispatch<Mod>(Binary<Mod>)
-.set_dispatch<Min>(Binary<Min>)
-.set_dispatch<Max>(Binary<Max>)
-.set_dispatch<EQ>(Binary<EQ>)
-.set_dispatch<NE>(Binary<NE>)
-.set_dispatch<LT>(Binary<LT>)
-.set_dispatch<LE>(Binary<LE>)
-.set_dispatch<GT>(Binary<GT>)
-.set_dispatch<GE>(Binary<GE>)
-.set_dispatch<And>(Binary<And>)
-.set_dispatch<Or>(Binary<Or>);
+void IRVisitor::Visit_(const Not* op) {
+  this->Visit(op->a);
+}
 
-TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable)
-.set_dispatch<Not>([](const Not* op, IRVisitor* v) {
-    v->Visit(op->a);
-  })
-.set_dispatch<Select>([](const Select *op, IRVisitor* v) {
-    v->Visit(op->condition);
-    v->Visit(op->true_value);
-    v->Visit(op->false_value);
-  })
-.set_dispatch<Ramp>([](const Ramp *op, IRVisitor* v) {
-    v->Visit(op->base);
-    v->Visit(op->stride);
-  })
-.set_dispatch<Broadcast>([](const Broadcast *op, IRVisitor* v) {
-    v->Visit(op->value);
-  });
+void IRVisitor::Visit_(const Select* op) {
+  this->Visit(op->condition);
+  this->Visit(op->true_value);
+  this->Visit(op->false_value);
+}
 
-TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable)
-.set_dispatch<AssertStmt>([](const AssertStmt *op, IRVisitor* v) {
-    v->Visit(op->condition);
-    v->Visit(op->message);
-  })
-.set_dispatch<ProducerConsumer>([](const ProducerConsumer *op, IRVisitor* v) {
-    v->Visit(op->body);
-  })
-.set_dispatch<Provide>([](const Provide *op, IRVisitor* v) {
-    VisitArray(op->args, v);
-    v->Visit(op->value);
-  })
-.set_dispatch<Realize>([](const Realize *op, IRVisitor* v) {
+void IRVisitor::Visit_(const Ramp *op) {
+  this->Visit(op->base);
+  this->Visit(op->stride);
+}
+
+void IRVisitor::Visit_(const Broadcast *op) {
+  this->Visit(op->value);
+}
+
+void IRVisitor::Visit_(const AssertStmt *op) {
+  this->Visit(op->condition);
+  this->Visit(op->message);
+}
+
+void IRVisitor::Visit_(const ProducerConsumer *op) {
+  this->Visit(op->body);
+}
+
+void IRVisitor::Visit_(const Provide *op) {
+  VisitArray(op->args, this);
+  this->Visit(op->value);
+}
+
+void IRVisitor::Visit_(const Realize *op) {
     // Mutate the bounds
   for (size_t i = 0; i < op->bounds.size(); i++) {
-      v->Visit(op->bounds[i]->min);
-      v->Visit(op->bounds[i]->extent);
+    this->Visit(op->bounds[i]->min);
+    this->Visit(op->bounds[i]->extent);
   }
 
-    v->Visit(op->body);
-    v->Visit(op->condition);
-  })
-.set_dispatch<Block>([](const Block *op, IRVisitor* v) {
-    v->Visit(op->first);
-    v->Visit(op->rest);
+  this->Visit(op->body);
+  this->Visit(op->condition);
+}
+
+void IRVisitor::Visit_(const Block *op) {
+  this->Visit(op->first);
+  this->Visit(op->rest);
+}
+
+void IRVisitor::Visit_(const Evaluate *op) {
+  this->Visit(op->value);
+}
+
+#define DEFINE_OP_NO_VISIT_(OP)                     \
+  void IRVisitor::Visit_(const OP* op) {}
+
+DEFINE_OP_NO_VISIT_(IntImm)
+DEFINE_OP_NO_VISIT_(UIntImm)
+DEFINE_OP_NO_VISIT_(FloatImm)
+DEFINE_OP_NO_VISIT_(StringImm)
+
+#define DISPATCH_TO_VISIT(OP)                       \
+  set_dispatch<OP>([](const OP* op, IRVisitor* v) { \
+      v->Visit_(op);                                \
     })
-.set_dispatch<Evaluate>([](const Evaluate *op, IRVisitor* v) {
-    v->Visit(op->value);
-  });
+
+TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable)
+.DISPATCH_TO_VISIT(Variable)
+.DISPATCH_TO_VISIT(LetStmt)
+.DISPATCH_TO_VISIT(AttrStmt)
+.DISPATCH_TO_VISIT(IfThenElse)
+.DISPATCH_TO_VISIT(For)
+.DISPATCH_TO_VISIT(Allocate)
+.DISPATCH_TO_VISIT(Load)
+.DISPATCH_TO_VISIT(Store)
+.DISPATCH_TO_VISIT(Let)
+.DISPATCH_TO_VISIT(Free)
+.DISPATCH_TO_VISIT(Call)
+.DISPATCH_TO_VISIT(Add)
+.DISPATCH_TO_VISIT(Sub)
+.DISPATCH_TO_VISIT(Mul)
+.DISPATCH_TO_VISIT(Div)
+.DISPATCH_TO_VISIT(Mod)
+.DISPATCH_TO_VISIT(Min)
+.DISPATCH_TO_VISIT(Max)
+.DISPATCH_TO_VISIT(EQ)
+.DISPATCH_TO_VISIT(NE)
+.DISPATCH_TO_VISIT(LT)
+.DISPATCH_TO_VISIT(LE)
+.DISPATCH_TO_VISIT(GT)
+.DISPATCH_TO_VISIT(GE)
+.DISPATCH_TO_VISIT(And)
+.DISPATCH_TO_VISIT(Or)
+.DISPATCH_TO_VISIT(Reduce)
+.DISPATCH_TO_VISIT(Cast)
+.DISPATCH_TO_VISIT(Not)
+.DISPATCH_TO_VISIT(Select)
+.DISPATCH_TO_VISIT(Ramp)
+.DISPATCH_TO_VISIT(Broadcast)
+.DISPATCH_TO_VISIT(AssertStmt)
+.DISPATCH_TO_VISIT(ProducerConsumer)
+.DISPATCH_TO_VISIT(Provide)
+.DISPATCH_TO_VISIT(Realize)
+.DISPATCH_TO_VISIT(Block)
+.DISPATCH_TO_VISIT(Evaluate)
+.DISPATCH_TO_VISIT(IntImm)
+.DISPATCH_TO_VISIT(UIntImm)
+.DISPATCH_TO_VISIT(FloatImm)
+.DISPATCH_TO_VISIT(StringImm);
 
 }  // namespace ir
 }  // namespace tvm