[REFACTOR] Migrate Low-level IR Passes into the New Stmt/Expr Mutator

include/tvm/ir_functor_ext.h

@@ -546,6 +546,35 @@ class StmtExprMutator :
   }
 };
 
+/*!
+ * \brief recursively visit the ir in post DFS order node, and transform it
+ *
+ * \param node The ir to be transformed.
+ * \param preorder The function called in before recursive mutation
+ *          If preorder returns None, then the transform will proceed to recursive call.
+ *          If preorder returns a not None Stmt/Expr, the transformer will simply return it and
+ *          won't do further recursion.
+ * \param postorder The function called after recursive mutation.
+ *          The recursive mutation result is passed to postorder for further mutation.
+ * \param only_enable List of StringImm.
+ *          If it is empty, all IRNode will call preorder/postorder
+ *          If it is not empty, preorder/postorder will only be called
+ *          when the IRNode's type key is in the list.
+ */
+TVM_DLL Stmt IRTransform(Stmt node,
+                         const runtime::PackedFunc& preorder,
+                         const runtime::PackedFunc& postorder,
+                         const Array<Expr>& only_enable = {});
+
+/*!
+ * \brief recursively visit the ir in post DFS order node, apply fvisit
+ * Each node is guaranteed to be visited only once.
+ * \param node The ir to be visited.
+ * \param fvisit The visitor function to be applied.
+ */
+TVM_DLL void PostOrderVisit(const ObjectRef& node, std::function<void(const ObjectRef&)> fvisit);
+
+
 }  // namespace ir
 }  // namespace tvm
 #endif  // TVM_IR_FUNCTOR_EXT_H_

include/tvm/ir_mutator.h

@@ -122,27 +122,6 @@ class TVM_DLL IRMutator {
   virtual Expr Mutate_(const StringImm* op, const Expr& e);
   virtual Expr Mutate_(const Shuffle* op, const Expr& e);
 };
-
-
-/*!
- * \brief recursively visit the ir in post DFS order node, and transform it
- *
- * \param node The ir to be transformed.
- * \param preorder The function called in before recursive mutation
- *          If preorder returns None, then the transform will proceed to recursive call.
- *          If preorder returns a not None Stmt/Expr, the transformer will simply return it and
- *          won't do further recursion.
- * \param postorder The function called after recursive mutation.
- *          The recursive mutation result is passed to postorder for further mutation.
- * \param only_enable List of StringImm.
- *          If it is empty, all IRNode will call preorder/postorder
- *          If it is not empty, preorder/postorder will only be called
- *          when the IRNode's type key is in the list.
- */
-Stmt IRTransform(Stmt node,
-                 const runtime::PackedFunc& preorder,
-                 const runtime::PackedFunc& postorder,
-                 const Array<Expr>& only_enable = {});
 }  // namespace ir
 }  // namespace tvm
 #endif  // TVM_IR_MUTATOR_H_

include/tvm/ir_visitor.h

@@ -145,15 +145,6 @@ class TVM_DLL IRVisitor {
   virtual void Visit_(const FloatImm* op);
   virtual void Visit_(const StringImm* op);
 };
-
-/*!
- * \brief recursively visit the ir in post DFS order node, apply fvisit
- * Each node is guaranteed to be visited only once.
- * \param node The ir to be visited.
- * \param fvisit The visitor function to be applied.
- */
-TVM_DLL void PostOrderVisit(const ObjectRef& node, std::function<void(const ObjectRef&)> fvisit);
-
 }  // namespace ir
 }  // namespace tvm
 

src/api/api_pass.cc

@@ -25,8 +25,7 @@
 #include <tvm/ir.h>
 #include <tvm/attrs.h>
 #include <tvm/ir_pass.h>
-#include <tvm/ir_visitor.h>
-#include <tvm/ir_mutator.h>
+#include <tvm/ir_functor_ext.h>
 #include <tvm/api_registry.h>
 
 namespace tvm {

src/arithmetic/bound_deducer.cc

@@ -23,7 +23,7 @@
  */
 #include <tvm/expr.h>
 #include <tvm/ir_pass.h>
-#include <tvm/ir_visitor.h>
+#include <tvm/ir_functor_ext.h>
 #include <tvm/arithmetic.h>
 #include <tvm/api_registry.h>
 
@@ -38,17 +38,17 @@ using namespace ir;
 
 // a visitor to find the path to the target variable
 // from a expression.
-class VariablePathFinder: public IRVisitor {
+class VariablePathFinder: public ExprVisitor {
  public:
   explicit VariablePathFinder(Expr target) : target_(target) {}
 
-  void Visit(const ObjectRef& node) final {
+  void VisitExpr(const Expr& node) final {
     if (visited_.count(node.get()) != 0) return;
     visited_.insert(node.get());
 
     if (!found_) path_.push_back(node.get());
     if (node.same_as(target_)) found_ = true;
-    IRVisitor::Visit(node);
+    ExprVisitor::VisitExpr(node);
     if (!found_) path_.pop_back();
   }
 
@@ -64,14 +64,14 @@ class VariablePathFinder: public IRVisitor {
 // return empty vector to represent failure
 std::vector<const Object*> GetPath(Expr target, Expr expr) {
   VariablePathFinder v(target);
-  v.Visit(expr);
+  v(expr);
   return v.path_;
 }
 
 enum CompareOp {kGreater, kLess, kEqual};
 
 // a visitor to deduce the bound of a variable from a expression
-class BoundDeducer: public IRVisitor {
+class BoundDeducer: public ExprVisitor {
  public:
   friend class BoundDeduceInputChecker;
   friend class Converter;
@@ -82,39 +82,39 @@ class BoundDeducer: public IRVisitor {
 
   void Deduce();
 
-  void Visit(const ObjectRef& e) final {
+  void VisitExpr(const Expr& e) final {
     if (!success_) return;
     if (e.get() == path_[iter_++]) {
-      IRVisitor::Visit(e);
+      ExprVisitor::VisitExpr(e);
     } else {
       success_ = false;
       return;
     }
   }
 
-  void Visit_(const LT* op) final {
+  void VisitExpr_(const LT* op) final {
     LOG(FATAL) << "unable to deduce due to multiple comparison operator";
   }
 
-  void Visit_(const LE* op) final {
+  void VisitExpr_(const LE* op) final {
     LOG(FATAL) << "unable to deduce due to multiple comparison operator";
   }
 
-  void Visit_(const GT* op) final {
+  void VisitExpr_(const GT* op) final {
     LOG(FATAL) << "unable to deduce due to multiple comparison operator";
   }
 
-  void Visit_(const GE* op) final {
+  void VisitExpr_(const GE* op) final {
     LOG(FATAL) << "unable to deduce due to multiple comparison operator";
   }
 
-  void Visit_(const Add* op) final {
+  void VisitExpr_(const Add* op) final {
     bool left = op->a.get() == path_[iter_];
     result_ -= left ? op->b : op->a;
-    Visit(left ? op->a : op->b);
+    this->VisitExpr(left ? op->a : op->b);
   }
 
-  void Visit_(const Sub* op) final {
+  void VisitExpr_(const Sub* op) final {
     bool left = op->a.get() == path_[iter_];
     if (left) {
       result_ += op->b;
@@ -123,10 +123,10 @@ class BoundDeducer: public IRVisitor {
       result_ = - result_;
       comp_op = ReverseOp(comp_op);
     }
-    Visit(left ? op->a : op->b);
+    this->VisitExpr(left ? op->a : op->b);
   }
 
-  void Visit_(const Mul* op) final {
+  void VisitExpr_(const Mul* op) final {
     bool left = op->a.get() == path_[iter_];
     Expr operand = left ? op->b : op->a;
     Expr target_var = left ? op->a : op->b;
@@ -171,7 +171,7 @@ class BoundDeducer: public IRVisitor {
         // ( x <= -3/-2 --> x <= 1)
       }
     }
-    Visit(left ? op->a : op->b);
+    this->VisitExpr(left ? op->a : op->b);
   }
 
   Expr result_;
@@ -194,17 +194,17 @@ class BoundDeducer: public IRVisitor {
   Analyzer analyzer_;
 };
 
-class BoundDeduceInputChecker: public IRVisitor {
+class BoundDeduceInputChecker: public ExprVisitor {
  public:
   bool Check(BoundDeducer* deducer) {
     deducer_ = deducer;
-    Visit(deducer_->expr_);
+    this->VisitExpr(deducer_->expr_);
     return target_count == 1;
   }
 
-  void Visit(const ObjectRef& e) final {
+  void VisitExpr(const Expr& e) final {
     if (e.same_as(deducer_->target_)) ++target_count;
-    IRVisitor::Visit(e);
+    ExprVisitor::VisitExpr(e);
   }
 
  private:
@@ -305,7 +305,7 @@ void BoundDeducer::Deduce() {
   }
   expr_map_ = EvalSetForEachSubExpr(expr_, hint_map_);
 
-  Visit(expr_);
+  this->VisitExpr(expr_);
 }
 
 void BoundDeducer::Relax() {