Refactor IR Pass

HalideIR @ 4becbde6

-Subproject commit 89b7939957d66a37dd6083ad6b09a5644e73fd8b
+Subproject commit 4becbde67c8aa565941b02648cea90f50211f8dc

include/tvm/expr.h

@@ -27,6 +27,7 @@ using Halide::abs;
 using Halide::select;
 
 using Halide::Expr;
+using Halide::IR::FunctionBaseNode;
 using Halide::Internal::Stmt;
 
 class Var : public Halide::VarExpr {

include/tvm/ir_mutator.h

@@ -29,7 +29,7 @@ class IRMutator {
    * \brief mutate expression
    * \return the mutated expr
    */
-  virtual Expr mutate(Expr expr) {
+  virtual Expr Mutate(Expr expr) {
     static const FMutateExpr& f = vtable_expr();
     return f(expr, expr, this);
   }
@@ -37,7 +37,7 @@ class IRMutator {
    * \brief mutate expression
    * \return the mutated stmt
    */
-  virtual Stmt mutate(Stmt stmt) {
+  virtual Stmt Mutate(Stmt stmt) {
     static const FMutateStmt& f = vtable_stmt();
     return f(stmt, stmt, this);
   }
@@ -58,28 +58,21 @@ class IRMutator {
  */
 class IRMutatorExample : public IRMutator {
  public:
-  Expr mutate(Expr expr) final {
+  Expr Mutate(Expr expr) final {
     static const FMutateExpr& f = IRMutatorExample::vtable_expr();
     return (f.can_dispatch(expr) ?
-            f(expr, expr, this) : IRMutator::mutate(expr));
+            f(expr, expr, this) : IRMutator::Mutate(expr));
   }
-  Stmt mutate(Stmt stmt) final {
+  Stmt Mutate(Stmt stmt) final {
     static const FMutateStmt& f = IRMutatorExample::vtable_stmt();
     return (f.can_dispatch(stmt) ?
-            f(stmt, stmt, this) : IRMutator::mutate(stmt));
+            f(stmt, stmt, this) : IRMutator::Mutate(stmt));
   }
   // to be implemented by child class
   static FMutateExpr& vtable_expr();  // NOLINT(*)
   static FMutateStmt& vtable_stmt();  // NOLINT(*)
 };
 
-/*!
- * \brief Substitute occurance of IRNode to be expr
- * \param replacements The replacement rule of substitution
- * \param expr The expression to be substituted.
- */
-Expr Substitute(const std::unordered_map<const IRNode*, Expr>& replacements, Expr expr);
-
 }  // namespace ir
 }  // namespace tvm
 #endif  // TVM_IR_MUTATOR_H_

include/tvm/ir_pass.h

+/*!
+ *  Copyright (c) 2016 by Contributors
+ * \file ir_pass.h
+ * \brief Collection of IR pass functions and visit functions
+ */
+#ifndef TVM_IR_PASS_H_
+#define TVM_IR_PASS_H_
+
+#include <tvm/ir_node.h>
+#include <unordered_map>
+#include "./expr.h"
+
+namespace tvm {
+namespace ir {
+
+/*!
+ * \brief Substitute occurance of IRNode in expr
+ * \param replacements The replacement rule of substitution
+ * \param expr The expression to be substituted.
+ */
+Expr Substitute(const std::unordered_map<const IRNode*, Expr>& replacements, Expr expr);
+
+}  // namespace ir
+}  // namespace tvm
+
+#endif  // TVM_IR_PASS_H_

include/tvm/ir_visitor.h

@@ -24,7 +24,7 @@ class IRVisitor {
   /*!
    * \brief recursively visit an IR node
    */
-  virtual void visit(const IRNodeRef& node) {
+  virtual void Visit(const IRNodeRef& node) {
     static const FVisit& f = vtable();
     if (node.defined()) f(node, this);
   }

include/tvm/tensor.h

@@ -101,7 +101,7 @@ class Tensor : public FunctionRef {
 };
 
 /*! \brief Node to represent a tensor */
-class TensorNode : public Node {
+class TensorNode : public FunctionBaseNode {
  public:
   /*! \brief The shape of the tensor */
   Array<Expr> shape;
@@ -125,6 +125,12 @@ class TensorNode : public Node {
     v->Visit("dim_var", &dim_var);
     v->Visit("source", &source);
   }
+  const std::string& func_name() const final {
+    return name;
+  }
+  int outputs() const final {
+    return 1;
+  }
   static Tensor make(Array<Expr> shape,
                      std::string name,
                      Type dtype,

src/pass/ir_mutator.cc

@@ -8,32 +8,6 @@
 namespace tvm {
 namespace ir {
 
-namespace {
-// visitor to implement apply
-class IRSubstitute : public IRMutator {
- public:
-  Expr mutate(Expr expr) final {
-    const IRNode* v = expr.get();
-    if (v != nullptr) {
-      auto it = replacements_.find(v);
-      if (it != replacements_.end()) {
-        return it->second;
-      }
-    }
-    return IRMutator::mutate(expr);
-  }
-  explicit IRSubstitute(const std::unordered_map<const IRNode*, Expr>& replacements)
-      : replacements_(replacements) {}
-
- private:
-  const std::unordered_map<const IRNode*, Expr>& replacements_;
-};
-}  // namespace
-
-Expr Substitute(const std::unordered_map<const IRNode*, Expr>& replacements, Expr expr) {
-  return IRSubstitute(replacements).mutate(expr);
-}
-
 IRMutator::FMutateExpr& IRMutator::vtable_expr() {  // NOLINT(*)
   static FMutateExpr inst; return inst;
 }
@@ -57,7 +31,7 @@ inline Array<Expr> MutateArray(Array<Expr> arr, IRMutator *m) {
   bool changed = false;
   for (size_t i = 0; i < arr.size(); i++) {
     Expr old_elem = arr[i];
-    Expr new_elem = m->mutate(old_elem);
+    Expr new_elem = m->Mutate(old_elem);
     if (!new_elem.same_as(old_elem)) changed = true;
     new_arr[i] = new_elem;
   }
@@ -73,8 +47,8 @@ inline RDomain MutateRDom(RDomain rdom, IRMutator *m) {
   bool changed = false;
   for (size_t i = 0; i < rdom->domain.size(); i++) {
     Range r = rdom->domain[i];
-    Expr new_min = m->mutate(r->min);
-    Expr new_extent = m->mutate(r->extent);
+    Expr new_min = m->Mutate(r->min);
+    Expr new_extent = m->Mutate(r->extent);
     if (!r->min.same_as(new_min)) changed = true;
     if (!r->extent.same_as(new_extent)) changed = true;
     new_dom[i] = Range::make_with_min_extent(new_min, new_extent);
@@ -89,7 +63,7 @@ inline RDomain MutateRDom(RDomain rdom, IRMutator *m) {
 TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_expr)
 .set_dispatch<Reduce>([](const Reduce* op, const Expr& e, IRMutator* m) {
     RDomain new_rdom  = MutateRDom(op->rdom, m);
-    Expr new_source  = m->mutate(op->source);
+    Expr new_source  = m->Mutate(op->source);
     if (op->rdom.same_as(new_rdom) &&
         op->source.same_as(new_source)) {
       return e;
@@ -107,7 +81,7 @@ TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_expr)
 
 TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_expr)
 .set_dispatch<Cast>([](const Cast* op, const Expr& e, IRMutator* m) {
-    Expr value = m->mutate(op->value);
+    Expr value = m->Mutate(op->value);
     if (value.same_as(op->value)) {
       return e;
     } else {
@@ -118,8 +92,8 @@ TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_expr)
 // 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);
+  Expr a = m->Mutate(op->a);
+  Expr b = m->Mutate(op->b);
   if (a.same_as(op->a) &&
       b.same_as(op->b)) {
     return e;
@@ -147,7 +121,7 @@ TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_expr)
 
 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 a = m->Mutate(op->a);
     if (a.same_as(op->a)) {
       return e;
     } else {
@@ -155,9 +129,9 @@ TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_expr)
     }
   })
 .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 cond = m->Mutate(op->condition);
+    Expr t = m->Mutate(op->true_value);
+    Expr f = m->Mutate(op->false_value);
     if (cond.same_as(op->condition) &&
         t.same_as(op->true_value) &&
         f.same_as(op->false_value)) {
@@ -167,7 +141,7 @@ TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_expr)
     }
   })
 .set_dispatch<Load>([](const Load *op, const Expr& e, IRMutator* m) {
-    Expr index = m->mutate(op->index);
+    Expr index = m->Mutate(op->index);
     if (index.same_as(op->index)) {
       return e;
     } else {
@@ -175,8 +149,8 @@ TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_expr)
     }
   })
 .set_dispatch<Ramp>([](const Ramp *op, const Expr& e, IRMutator* m) {
-    Expr base = m->mutate(op->base);
-    Expr stride = m->mutate(op->stride);
+    Expr base = m->Mutate(op->base);
+    Expr stride = m->Mutate(op->stride);
     if (base.same_as(op->base) &&
         stride.same_as(op->stride)) {
       return e;
@@ -185,7 +159,7 @@ TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_expr)
     }
   })
 .set_dispatch<Broadcast>([](const Broadcast *op, const Expr& e, IRMutator* m) {
-    Expr value = m->mutate(op->value);
+    Expr value = m->Mutate(op->value);
     if (value.same_as(op->value)) {
       return e;
     } else {
@@ -202,8 +176,8 @@ TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_expr)
     }
   })
 .set_dispatch<Let>([](const Let *op, const Expr& e, IRMutator* m) {
-    Expr value = m->mutate(op->value);
-    Expr body = m->mutate(op->body);
+    Expr value = m->Mutate(op->value);
+    Expr body = m->Mutate(op->body);
     if (value.same_as(op->value) &&
         body.same_as(op->body)) {
       return e;
@@ -214,8 +188,8 @@ TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_expr)
 
 TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_stmt)
 .set_dispatch<LetStmt>([](const LetStmt *op, const Stmt& s, IRMutator* m) {
-    Expr value = m->mutate(op->value);
-    Stmt body = m->mutate(op->body);
+    Expr value = m->Mutate(op->value);
+    Stmt body = m->Mutate(op->body);
     if (value.same_as(op->value) &&
         body.same_as(op->body)) {
       return s;
@@ -224,8 +198,8 @@ 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);
+    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;
@@ -234,7 +208,7 @@ TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_stmt)
     }
   })
 .set_dispatch<ProducerConsumer>([](const ProducerConsumer *op, const Stmt& s, IRMutator* m) {
-    Stmt body = m->mutate(op->body);
+    Stmt body = m->Mutate(op->body);
     if (body.same_as(op->body)) {
       return s;
     } else {
@@ -242,9 +216,9 @@ TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_stmt)
     }
   })
 .set_dispatch<For>([](const For *op, const Stmt& s, IRMutator* m) {
-    Expr min = m->mutate(op->min);
-    Expr extent = m->mutate(op->extent);
-    Stmt body = m->mutate(op->body);
+    Expr min = m->Mutate(op->min);
+    Expr extent = m->Mutate(op->extent);
+    Stmt body = m->Mutate(op->body);
     if (min.same_as(op->min) &&
         extent.same_as(op->extent) &&
         body.same_as(op->body)) {
@@ -255,8 +229,8 @@ TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_stmt)
     }
   })
 .set_dispatch<Store>([](const Store *op, const Stmt& s, IRMutator* m) {
-    Expr value = m->mutate(op->value);
-    Expr index = m->mutate(op->index);
+    Expr value = m->Mutate(op->value);
+    Expr index = m->Mutate(op->index);
     if (value.same_as(op->value) && index.same_as(op->index)) {
       return s;
     } else {
@@ -276,14 +250,14 @@ TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_stmt)
     std::vector<Expr> new_extents;
     bool all_extents_unmodified = true;
     for (size_t i = 0; i < op->extents.size(); i++) {
-        new_extents.push_back(m->mutate(op->extents[i]));
+        new_extents.push_back(m->Mutate(op->extents[i]));
         all_extents_unmodified &= new_extents[i].same_as(op->extents[i]);
     }
-    Stmt body = m->mutate(op->body);
-    Expr condition = m->mutate(op->condition);
+    Stmt body = m->Mutate(op->body);
+    Expr condition = m->Mutate(op->condition);
     Expr new_expr;
     if (op->new_expr.defined()) {
-      new_expr = m->mutate(op->new_expr);
+      new_expr = m->Mutate(op->new_expr);
     }
     if (all_extents_unmodified &&
         body.same_as(op->body) &&
@@ -308,16 +282,16 @@ TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_stmt)
     for (size_t i = 0; i < op->bounds.size(); i++) {
         Expr old_min = op->bounds[i]->min;
         Expr old_extent = op->bounds[i]->extent;
-        Expr new_min = m->mutate(old_min);
-        Expr new_extent = m->mutate(old_extent);
+        Expr new_min = m->Mutate(old_min);
+        Expr new_extent = m->Mutate(old_extent);
         if (!new_min.same_as(old_min))  bounds_changed = true;
         if (!new_extent.same_as(old_extent)) bounds_changed = true;
         new_bounds.push_back(
             Range::make_by_min_extent(new_min, new_extent));
     }
 
-    Stmt body = m->mutate(op->body);
-    Expr condition = m->mutate(op->condition);
+    Stmt body = m->Mutate(op->body);
+    Expr condition = m->Mutate(op->condition);
     if (!bounds_changed &&
         body.same_as(op->body) &&
         condition.same_as(op->condition)) {
@@ -328,8 +302,8 @@ TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_stmt)
     }
   })
 .set_dispatch<Block>([](const Block *op, const Stmt& s, IRMutator* m) {
-    Stmt first = m->mutate(op->first);
-    Stmt rest = m->mutate(op->rest);
+    Stmt first = m->Mutate(op->first);
+    Stmt rest = m->Mutate(op->rest);
     if (first.same_as(op->first) &&
         rest.same_as(op->rest)) {
       return s;
@@ -338,9 +312,9 @@ TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_stmt)
     }
   })
 .set_dispatch<IfThenElse>([](const IfThenElse *op, const Stmt& s, IRMutator* m) {
-    Expr condition = m->mutate(op->condition);
-    Stmt then_case = m->mutate(op->then_case);
-    Stmt else_case = m->mutate(op->else_case);
+    Expr condition = m->Mutate(op->condition);
+    Stmt then_case = m->Mutate(op->then_case);
+    Stmt else_case = m->Mutate(op->else_case);
     if (condition.same_as(op->condition) &&
         then_case.same_as(op->then_case) &&
         else_case.same_as(op->else_case)) {
@@ -350,7 +324,7 @@ TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_stmt)
     }
   })
 .set_dispatch<Evaluate>([](const Evaluate *op, const Stmt& s, IRMutator* m) {
-    Expr v = m->mutate(op->value);
+    Expr v = m->Mutate(op->value);
     if (v.same_as(op->value)) {
       return s;
     } else {

src/pass/ir_pass.cc

+/*!
+ *  Copyright (c) 2016 by Contributors
+ * \file ir_pass.cc
+ */
+#include <tvm/ir.h>
+#include <tvm/ir_visitor.h>
+#include <tvm/ir_mutator.h>
+#include <unordered_set>
+
+namespace tvm {
+namespace ir {
+namespace {
+// visitor to implement apply
+class IRSubstitute : public IRMutator {
+ public:
+  Expr Mutate(Expr expr) final {
+    const IRNode* v = expr.get();
+    if (v != nullptr) {
+      auto it = replacements_.find(v);
+      if (it != replacements_.end()) {
+        return it->second;
+      }
+    }
+    return IRMutator::Mutate(expr);
+  }
+  explicit IRSubstitute(const std::unordered_map<const IRNode*, Expr>& replacements)
+      : replacements_(replacements) {}
+
+ private:
+  const std::unordered_map<const IRNode*, Expr>& replacements_;
+};
+}  // namespace
+
+Expr Substitute(const std::unordered_map<const IRNode*, Expr>& replacements, Expr expr) {
+  return IRSubstitute(replacements).Mutate(expr);
+}
+
+}  // namespace ir
+}  // namespace tvm

src/pass/ir_visitor.cc

@@ -14,10 +14,10 @@ class IRApplyVisit : public IRVisitor {
  public:
   explicit IRApplyVisit(std::function<void(const IRNodeRef&)> f) : f_(f) {}
 
-  void visit(const IRNodeRef& node) final {
+  void Visit(const IRNodeRef& node) final {
     if (visited_.count(node.get()) != 0) return;
     visited_.insert(node.get());
-    IRVisitor::visit(node);
+    IRVisitor::Visit(node);
     f_(node);
   }
 
@@ -25,18 +25,18 @@ class IRApplyVisit : public IRVisitor {
   std::function<void(const IRNodeRef&)> f_;
   std::unordered_set<const Node*> visited_;
 };
+
 }  // namespace
 
+void PostOrderVisit(const IRNodeRef& node, std::function<void(const IRNodeRef&)> fvisit) {
+  IRApplyVisit(fvisit).Visit(node);
+}
+
 IRVisitor::FVisit& IRVisitor::vtable() {  // NOLINT(*)
   static FVisit inst; return inst;
 }
 
 
-void PostOrderVisit(const IRNodeRef& node, std::function<void(const IRNodeRef&)> fvisit) {
-  IRApplyVisit v(fvisit);
-  v.visit(node);
-}
-
 // namespace to register the functors.
 namespace {
 
@@ -47,22 +47,22 @@ void NoOp(const IRNodeRef& n, IRVisitor* v) {
 
 inline void VisitArray(Array<Expr> arr, IRVisitor* v) {
   for (size_t i = 0; i < arr.size(); i++) {
-    v->visit(arr[i]);
+    v->Visit(arr[i]);
   }
 }
 
 inline void VisitRDom(RDomain rdom, IRVisitor* v) {
   for (size_t i = 0; i < rdom->domain.size(); i++) {
     Range r = rdom->domain[i];
-    v->visit(r->min);
-    v->visit(r->extent);
+    v->Visit(r->min);
+    v->Visit(r->extent);
   }
 }
 
 TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable)
 .set_dispatch<Reduce>([](const Reduce* op, IRVisitor* v) {
     VisitRDom(op->rdom, v);
-    v->visit(op->source);
+    v->Visit(op->source);
   });
 
 TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable)
@@ -74,14 +74,14 @@ TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable)
 
 TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable)
 .set_dispatch<Cast>([](const Cast* op, IRVisitor* v) {
-    v->visit(op->value);
+    v->Visit(op->value);
   });
 
 // binary operator
 template<typename T>
 inline void Binary(const T* op, IRVisitor* v) {
-  v->visit(op->a);
-  v->visit(op->b);
+  v->Visit(op->a);
+  v->Visit(op->b);
 }
 
 TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable)
@@ -103,51 +103,51 @@ TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable)
 
 TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable)
 .set_dispatch<Not>([](const Not* op, IRVisitor* v) {
-    v->visit(op->a);
+    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);
+    v->Visit(op->condition);
+    v->Visit(op->true_value);
+    v->Visit(op->false_value);
   })
 .set_dispatch<Load>([](const Load *op, IRVisitor* v) {
-    v->visit(op->index);
+    v->Visit(op->index);
   })
 .set_dispatch<Ramp>([](const Ramp *op, IRVisitor* v) {
-    v->visit(op->base);
-    v->visit(op->stride);
+    v->Visit(op->base);
+    v->Visit(op->stride);
   })
 .set_dispatch<Broadcast>([](const Broadcast *op, IRVisitor* v) {
-    v->visit(op->value);
+    v->Visit(op->value);
   })
 .set_dispatch<Call>([](const Call *op, IRVisitor* v) {
     VisitArray(op->args, v);
   })
 .set_dispatch<Let>([](const Let *op, IRVisitor* v) {
-    v->visit(op->value);
-    v->visit(op->body);
+    v->Visit(op->value);
+    v->Visit(op->body);
   });
 
 TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable)
 .set_dispatch<LetStmt>([](const LetStmt *op, IRVisitor* v) {
-    v->visit(op->value);
-    v->visit(op->body);
+    v->Visit(op->value);
+    v->Visit(op->body);
   })
 .set_dispatch<AssertStmt>([](const AssertStmt *op, IRVisitor* v) {
-    v->visit(op->condition);
-    v->visit(op->message);
+    v->Visit(op->condition);
+    v->Visit(op->message);
   })
 .set_dispatch<ProducerConsumer>([](const ProducerConsumer *op, IRVisitor* v) {
-    v->visit(op->body);
+    v->Visit(op->body);
   })
 .set_dispatch<For>([](const For *op, IRVisitor* v) {
-    v->visit(op->min);
-    v->visit(op->extent);
-    v->visit(op->body);
+    v->Visit(op->min);
+    v->Visit(op->extent);
+    v->Visit(op->body);
   })
 .set_dispatch<Store>([](const Store *op, IRVisitor* v) {
-    v->visit(op->value);
-    v->visit(op->index);
+    v->Visit(op->value);
+    v->Visit(op->index);
   })
 .set_dispatch<Provide>([](const Provide *op, IRVisitor* v) {
     VisitArray(op->args, v);
@@ -155,36 +155,36 @@ TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable)
   })
 .set_dispatch<Allocate>([](const Allocate *op, IRVisitor* v) {
     for (size_t i = 0; i < op->extents.size(); i++) {
-      v->visit(op->extents[i]);
+      v->Visit(op->extents[i]);
     }
-    v->visit(op->body);
-    v->visit(op->condition);
+    v->Visit(op->body);
+    v->Visit(op->condition);
     if (op->new_expr.defined()) {
-      v->visit(op->new_expr);
+      v->Visit(op->new_expr);
     }
   })
 .set_dispatch<Free>(NoOp)
 .set_dispatch<Realize>([](const Realize *op, IRVisitor* v) {
     // 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);
+      v->Visit(op->bounds[i]->min);
+      v->Visit(op->bounds[i]->extent);
     }
 
-    v->visit(op->body);
-    v->visit(op->condition);
+    v->Visit(op->body);
+    v->Visit(op->condition);
   })
 .set_dispatch<Block>([](const Block *op, IRVisitor* v) {
-    v->visit(op->first);
-    v->visit(op->rest);
+    v->Visit(op->first);
+    v->Visit(op->rest);
   })
 .set_dispatch<IfThenElse>([](const IfThenElse *op, IRVisitor* v) {
-    v->visit(op->condition);
-    v->visit(op->then_case);
-    v->visit(op->else_case);
+    v->Visit(op->condition);
+    v->Visit(op->then_case);
+    v->Visit(op->else_case);
   })
 .set_dispatch<Evaluate>([](const Evaluate *op, IRVisitor* v) {
-    v->visit(op->value);
+    v->Visit(op->value);
   });
 
 }  // namespace

tests/cpp/ir_mutator_test.cc

@@ -13,10 +13,10 @@ class IRVar2Const : public IRMutator {
  public:
   VarExpr var;
   int int_val;
-  Expr mutate(Expr expr) final {
+  Expr Mutate(Expr expr) final {
     static const FMutateExpr& f = IRVar2Const::vtable_expr();
     return (f.can_dispatch(expr) ?
-            f(expr, expr, this) : IRMutator::mutate(expr));
+            f(expr, expr, this) : IRMutator::Mutate(expr));
   }
   static FMutateExpr &vtable_expr();
 };
@@ -46,31 +46,12 @@ TEST(IRMutator, Basic) {
   IRVar2Const mu;
   mu.var = y;
   mu.int_val = 10;
-  auto zz = mu.mutate(z);
+  auto zz = mu.Mutate(z);
   std::ostringstream os;
   os << zz;
   CHECK(os.str() == "(x + 10)");
 }
 
-TEST(IRMutator, Substitute) {
-  using namespace Halide::Internal;
-  using namespace tvm;
-  Var x("x"), y;
-  auto z = x + y;
-  {
-    auto zz = Substitute({{y.get(), 11}}, z);
-    std::ostringstream os;
-    os << zz;
-    CHECK(os.str() == "(x + 11)");
-  }
-  {
-    auto zz = Substitute({{z.get(), 11}}, z);
-    std::ostringstream os;
-    os << zz;
-    CHECK(os.str() == "11");
-  }
-}
-
 int main(int argc, char ** argv) {
   testing::InitGoogleTest(&argc, argv);
   testing::FLAGS_gtest_death_test_style = "threadsafe";

tests/cpp/ir_pass_test.cc

+#include <dmlc/logging.h>
+#include <gtest/gtest.h>
+#include <tvm/tvm.h>
+#include <tvm/ir_pass.h>
+
+TEST(IRPass, Substitute) {
+  using namespace Halide::Internal;
+  using namespace tvm;
+  Var x("x"), y;
+  auto z = x + y;
+  {
+    auto zz = ir::Substitute({{y.get(), 11}}, z);
+    std::ostringstream os;
+    os << zz;
+    CHECK(os.str() == "(x + 11)");
+  }
+  {
+    auto zz = ir::Substitute({{z.get(), 11}}, z);
+    std::ostringstream os;
+    os << zz;
+    CHECK(os.str() == "11");
+  }
+}
+
+int main(int argc, char ** argv) {
+  testing::InitGoogleTest(&argc, argv);
+  testing::FLAGS_gtest_death_test_style = "threadsafe";
+  return RUN_ALL_TESTS();
+}

tests/cpp/ir_visitor_test.cc

@@ -2,6 +2,7 @@
 #include <gtest/gtest.h>
 #include <tvm/tvm.h>
 #include <tvm/ir_visitor.h>
+#include <tvm/ir_pass.h>
 
 TEST(IRVisitor, CountVar) {
   using namespace Halide::Internal;