[REFACTOR] Remove old Low-level Visitor/Mutator (#4612)

include/tvm/ir_mutator.h

-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements.  See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership.  The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License.  You may obtain a copy of the License at
- *
- *   http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied.  See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file tvm/ir_mutator.h
- * \brief Defines general IRMutation pass
- */
-#ifndef TVM_IR_MUTATOR_H_
-#define TVM_IR_MUTATOR_H_
-
-#include <unordered_map>
-#include <utility>
-#include "expr.h"
-#include "ir.h"
-#include "tvm/node/functor.h"
-
-namespace tvm {
-namespace ir {
-/*!
- * \brief a base class for mutator to iterative mutate the IR
- *
- *  This IRMutator is implemented via Visitor Pattern.
- *  Also you can implement via NodeFunctor.
- *  This enables easy extensions of possible new Node.
- *  It also makes changing return types easier.
- *
- * \note If you want to return a different type other than Expr and Stmt,
- *       Simply following the same pattern as IRMutator and create a seperate class.
- * \sa NodeFunctor
- */
-class TVM_DLL IRMutator {
- public:
-  /*!
-   * \brief mutate expression
-   * \return the mutated expr
-   */
-  virtual Expr Mutate(Expr expr) {
-    static const FMutateExpr& f = vtable_expr();
-    return f(expr, expr, this);
-  }
-  /*!
-   * \brief mutate expression
-   * \return the mutated stmt
-   */
-  virtual Stmt Mutate(Stmt stmt) {
-    static const FMutateStmt& f = vtable_stmt();
-    return f(stmt, stmt, this);
-  }
-  /*! \brief destructor */
-  virtual ~IRMutator() {}
-  /*! \brief functor type of expr mutation */
-  using FMutateExpr = NodeFunctor<Expr(const ObjectRef&, const Expr&, IRMutator*)>;
-  /*! \brief functor type of stmt mutation */
-  using FMutateStmt = NodeFunctor<Stmt(const ObjectRef&, const Stmt&, IRMutator*)>;
-  /*! \return internal vtable of expr */
-  static FMutateExpr& vtable_expr();  // NOLINT(*)
-  /*! \return internal stmt of expr */
-  static FMutateStmt& vtable_stmt();  // NOLINT(*)
-  // Set of overloadable functions
-  // The underscore allows Mutate not to be shadowed by inheritance
-  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 Allocate* op, const Stmt& s);
-  virtual Stmt Mutate_(const Store* op, const Stmt& s);
-  virtual Stmt Mutate_(const Free* op, const Stmt& s);
-  virtual Stmt Mutate_(const AssertStmt* op, const Stmt& s);
-  virtual Stmt Mutate_(const ProducerConsumer* op, const Stmt& s);
-  virtual Stmt Mutate_(const Provide* op, const Stmt& s);
-  virtual Stmt Mutate_(const Realize* op, const Stmt& s);
-  virtual Stmt Mutate_(const Prefetch* op, const Stmt& s);
-  virtual Stmt Mutate_(const Block* op, const Stmt& s);
-  virtual Stmt Mutate_(const Evaluate* op, const Stmt& s);
-
-  virtual Expr Mutate_(const Variable* op, const Expr& e);
-  virtual Expr Mutate_(const Load* op, const Expr& e);
-  virtual Expr Mutate_(const Let* 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 FloorDiv* op, const Expr& e);
-  virtual Expr Mutate_(const FloorMod* 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 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);
-  virtual Expr Mutate_(const Shuffle* op, const Expr& e);
-};
-}  // namespace ir
-}  // namespace tvm
-#endif  // TVM_IR_MUTATOR_H_

include/tvm/ir_visitor.h

-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements.  See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership.  The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License.  You may obtain a copy of the License at
- *
- *   http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied.  See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file tvm/ir_visitor.h
- * \brief Visitor to quickly visit IR trees
- */
-#ifndef TVM_IR_VISITOR_H_
-#define TVM_IR_VISITOR_H_
-
-#include "ir.h"
-#include "tvm/node/functor.h"
-
-namespace tvm {
-namespace ir {
-
-/*!
- * \brief a base class for visitor to iterative traverse the IR
- *
- *  This IRVisitor is implemented via NodeFunctor
- *  This enables extensions of possible new Node.
- *
- * \sa ExprFunctor, StmtFunctor, PostOrderVisit
- *
- * \note If you need to return values during Visit:
- *  - If it is mutation of the IR, use IRMutator
- *  - If you want to return other things, consider use ExprFunctor/StmtFunctor
- *  - Watch out for possible bug pattern if you use IRVisitor to simulate returns.
- *
- * \code
- *
- * // This is an example code to show cases for traps in IRVisitor
- * // The use case is to count number of Variables in the ir tree.
- * class MyCounter : public IRVisitor {
- *  public:
- *   int Count(const ObjectRef& n) {
- *     ret_ = 0;
- *     this->Visit(n);
- *     return ret_;
- *   }
- *   void Visit_(const Variable* op) final {
- *     ret_ = 1;
- *   }
- *   void Visit_(const Add* op) final {
- *     ret_ = count(op->a) + count(op->b);
- *   }
-
- *  private:
- *   int ret_;
- * };
- * MyCounter counter;
- * Var x("x");
- * // this returns 2
- * CHECK_EQ(counter.Count(x + x), 2);
- * // Think what is the result of the following count
- * counter.count(Max::make(x, x));
- * // The result is actually 1
- * // This is because Visit is not overriden for Max
- * // so it simply calls Visit for the left and right children
- * // and because Count is not called, ret_ is not cleared.
- * // There can also be cases where ret_ is forgetten to be set.
- *
- * // These traps may not happen if we program carefully
- * // But it is recommended to use ExprFunctor, which allows direct
- * // return the value, this helps us to avoid such problems.
- *
- * \endcode
- */
-class TVM_DLL IRVisitor {
- public:
-  /*!
-   * \brief recursively visit an IR node
-   */
-  virtual void Visit(const ObjectRef& node) {
-    static const FVisit& f = vtable();
-    if (node.defined()) f(node, this);
-  }
-  /*! \brief destructor */
-  virtual ~IRVisitor() {}
-  /*! \brief functor type of visitor */
-  using FVisit = NodeFunctor<void(const ObjectRef&, IRVisitor*)>;
-  /*! \return internal vtable*/
-  static FVisit& vtable();
-  // overloadable visit function.
-  virtual void Visit_(const Variable* 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 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 FloorDiv* op);
-  virtual void Visit_(const FloorMod* 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 Shuffle* 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 Prefetch* 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);
-};
-}  // namespace ir
-}  // namespace tvm
-
-#endif  // TVM_IR_VISITOR_H_

src/autotvm/feature_visitor.h

@@ -51,12 +51,12 @@ enum AnnotationType {
 class FeatureVisitor : public StmtExprVisitor {
  public:
   // for loop
-  void VisitStmt_(const For *op);
-  void VisitStmt_(const AttrStmt *op);
+  void VisitStmt_(const For* op) final;
+  void VisitStmt_(const AttrStmt* op) final;
 
   // memory access
-  void VisitExpr_(const Load *op);
-  void VisitStmt_(const Store *op);
+  void VisitExpr_(const Load* op) final;
+  void VisitStmt_(const Store* op) final;
 
   using StmtExprVisitor::VisitStmt_;
   using StmtExprVisitor::VisitExpr_;

src/autotvm/touch_extractor.cc

@@ -51,7 +51,7 @@ class IndexParser: public ExprVisitor {
     this->VisitExpr(expr);
   }
 
-  void VisitExpr_(const Variable *op) {
+  void VisitExpr_(const Variable* op) final {
     // TODO(lmzheng): handle more index types (multiple occurrence)
     if (pattern_map.count(op) == 0) {
       pattern_map[op] = TouchPattern();
@@ -60,7 +60,7 @@ class IndexParser: public ExprVisitor {
     }
   }
 
-  void VisitExpr_(const Mul *op) {
+  void VisitExpr_(const Mul* op) final {
     if (op->a.as<Variable>()) {
       if (const auto stride = op->b.as<IntImm>()) {
         next_stride_ = stride->value;

src/autotvm/touch_extractor.h

@@ -90,31 +90,31 @@ class TouchExtractor : public FeatureVisitor {
   }
 
   // arithmetic stats
-  void VisitExpr_(const Add *op) {
+  void VisitExpr_(const Add* op) final {
     if (op->dtype.is_float())
       itervar_map[itervar_stack_.back()].add_ct++;
     FeatureVisitor::VisitExpr_(op);
   }
 
-  void VisitExpr_(const Sub *op) {
+  void VisitExpr_(const Sub* op) final {
     if (op->dtype.is_float())
       itervar_map[itervar_stack_.back()].add_ct++;
     FeatureVisitor::VisitExpr_(op);
   }
 
-  void VisitExpr_(const Mul *op) {
+  void VisitExpr_(const Mul* op) final {
     if (op->dtype.is_float())
       itervar_map[itervar_stack_.back()].mul_ct++;
     FeatureVisitor::VisitExpr_(op);
   }
 
-  void VisitExpr_(const Div *op) {
+  void VisitExpr_(const Div* op) final {
     if (op->dtype.is_float())
       itervar_map[itervar_stack_.back()].div_ct++;
     FeatureVisitor::VisitExpr_(op);
   }
 
-  void VisitExpr_(const Mod *op) {
+  void VisitExpr_(const Mod* op) final {
     if (op->dtype.is_float())
       itervar_map[itervar_stack_.back()].div_ct++;
     FeatureVisitor::VisitExpr_(op);

src/pass/ir_visitor.cc

-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements.  See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership.  The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License.  You may obtain a copy of the License at
- *
- *   http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied.  See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file ir_visitor.cc
- */
-#include <tvm/ir.h>
-#include <tvm/ir_visitor.h>
-#include <unordered_set>
-
-namespace tvm {
-namespace ir {
-
-IRVisitor::FVisit& IRVisitor::vtable() {  // NOLINT(*)
-  static FVisit inst; return inst;
-}
-
-inline void VisitArray(const Array<Expr>& arr, IRVisitor* v) {
-  for (size_t i = 0; i < arr.size(); i++) {
-    v->Visit(arr[i]);
-  }
-}
-
-inline void VisitRDom(const Array<IterVar>& rdom, IRVisitor* v) {
-  for (size_t i = 0; i < rdom.size(); i++) {
-    Range r = rdom[i]->dom;
-    v->Visit(r->min);
-    v->Visit(r->extent);
-  }
-}
-
-void IRVisitor::Visit_(const Variable* op) {}
-
-void IRVisitor::Visit_(const LetStmt* op) {
-  this->Visit(op->value);
-  this->Visit(op->body);
-}
-
-void IRVisitor::Visit_(const AttrStmt* op) {
-  this->Visit(op->value);
-  this->Visit(op->body);
-}
-
-void IRVisitor::Visit_(const For* op) {
-  IRVisitor* v = this;
-  v->Visit(op->min);
-  v->Visit(op->extent);
-  v->Visit(op->body);
-}
-
-void IRVisitor::Visit_(const Allocate* op) {
-  IRVisitor* v = this;
-  for (size_t i = 0; i < op->extents.size(); i++) {
-    v->Visit(op->extents[i]);
-  }
-  v->Visit(op->body);
-  v->Visit(op->condition);
-  if (op->new_expr.defined()) {
-    v->Visit(op->new_expr);
-  }
-}
-
-void IRVisitor::Visit_(const Load* op) {
-  this->Visit(op->index);
-  this->Visit(op->predicate);
-}
-
-void IRVisitor::Visit_(const Store* op) {
-  this->Visit(op->value);
-  this->Visit(op->index);
-  this->Visit(op->predicate);
-}
-
-void IRVisitor::Visit_(const IfThenElse* op) {
-  this->Visit(op->condition);
-  this->Visit(op->then_case);
-  if (op->else_case.defined()) {
-    this->Visit(op->else_case);
-  }
-}
-
-void IRVisitor::Visit_(const Let* op) {
-  this->Visit(op->value);
-  this->Visit(op->body);
-}
-
-void IRVisitor::Visit_(const Free* op) {}
-
-void IRVisitor::Visit_(const Call* op) {
-  VisitArray(op->args, this);
-}
-
-#define DEFINE_BINOP_VISIT_(OP)                     \
-  void IRVisitor::Visit_(const OP* op) {            \
-    this->Visit(op->a);                             \
-    this->Visit(op->b);                             \
-  }
-
-DEFINE_BINOP_VISIT_(Add)
-DEFINE_BINOP_VISIT_(Sub)
-DEFINE_BINOP_VISIT_(Mul)
-DEFINE_BINOP_VISIT_(Div)
-DEFINE_BINOP_VISIT_(Mod)
-DEFINE_BINOP_VISIT_(FloorDiv)
-DEFINE_BINOP_VISIT_(FloorMod)
-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);
-  VisitArray(op->source, this);
-  this->Visit(op->condition);
-}
-
-void IRVisitor::Visit_(const Cast* op) {
-  this->Visit(op->value);
-}
-
-void IRVisitor::Visit_(const Not* op) {
-  this->Visit(op->a);
-}
-
-void IRVisitor::Visit_(const Select* op) {
-  this->Visit(op->condition);
-  this->Visit(op->true_value);
-  this->Visit(op->false_value);
-}
-
-void IRVisitor::Visit_(const Ramp* op) {
-  this->Visit(op->base);
-  this->Visit(op->stride);
-}
-
-void IRVisitor::Visit_(const Shuffle* op) {
-  for (const auto& elem : op->indices)
-    this->Visit(elem);
-  for (const auto& elem : op->vectors)
-    this->Visit(elem);
-}
-
-void IRVisitor::Visit_(const Broadcast* op) {
-  this->Visit(op->value);
-}
-
-void IRVisitor::Visit_(const AssertStmt* op) {
-  this->Visit(op->condition);
-  this->Visit(op->message);
-  this->Visit(op->body);
-}
-
-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) {
-  for (size_t i = 0; i < op->bounds.size(); i++) {
-    this->Visit(op->bounds[i]->min);
-    this->Visit(op->bounds[i]->extent);
-  }
-
-  this->Visit(op->body);
-  this->Visit(op->condition);
-}
-
-void IRVisitor::Visit_(const Prefetch* op) {
-  for (size_t i = 0; i < op->bounds.size(); i++) {
-    this->Visit(op->bounds[i]->min);
-    this->Visit(op->bounds[i]->extent);
-  }
-}
-
-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 ObjectRef& node, IRVisitor* v) { \
-      v->Visit_(static_cast<const OP*>(node.get()));         \
-    })
-
-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(FloorDiv)
-.DISPATCH_TO_VISIT(FloorMod)
-.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(Shuffle)
-.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)
-.DISPATCH_TO_VISIT(Prefetch);
-
-}  // namespace ir
-}  // namespace tvm

tests/cpp/ir_functor_test.cc

@@ -41,6 +41,19 @@ TEST(IRF, Basic) {
   CHECK_EQ(f(z, 2),  4);
 }
 
+TEST(IRF, CountVar) {
+  using namespace tvm;
+  int n_var = 0;
+  Var x("x"), y;
+
+  auto z = x + 1 + y + y;
+  ir::PostOrderVisit(z, [&n_var](const ObjectRef& n) {
+      if (n.as<Variable>()) ++n_var;
+    });
+  CHECK_EQ(n_var, 2);
+}
+
+
 TEST(IRF, ExprTransform) {
   using namespace tvm;
   using namespace tvm::ir;

tests/cpp/ir_mutator_test.cc

-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements.  See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership.  The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License.  You may obtain a copy of the License at
- *
- *   http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied.  See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-#include <dmlc/logging.h>
-#include <gtest/gtest.h>
-#include <tvm/ir_mutator.h>
-#include <tvm/expr_operator.h>
-
-namespace {
-using namespace tvm;
-using namespace tvm::ir;
-
-// replace variable to constant
-class IRVar2Const : public IRMutator {
- public:
-  Var var;
-  int int_val;
-  Expr Mutate(Expr expr) final {
-    static const FMutateExpr& f = IRVar2Const::vtable_expr();
-    return (f.can_dispatch(expr) ?
-            f(expr, expr, this) : IRMutator::Mutate(expr));
-  }
-  static FMutateExpr &vtable_expr();
-};
-
-// implement vtable
-IRMutator::FMutateExpr &IRVar2Const::vtable_expr() {  // NOLINT(*)
-  static FMutateExpr inst; return inst;
-}
-
-TVM_STATIC_IR_FUNCTOR(IRVar2Const, vtable_expr)
-.set_dispatch<Variable>([](const ObjectRef& ref, const Expr &e, IRMutator* m) {
-    IRVar2Const* vm = static_cast<IRVar2Const*>(m);
-    if (e.same_as(vm->var)) {
-      return Expr(IntImm::make(DataType::Int(32), vm->int_val));
-    } else {
-      return e;
-    }
-  });
-
-}  // namespace
-
-TEST(IRMutator, Basic) {
-  using namespace tvm::ir;
-  using namespace tvm;
-  Var x("x"), y;
-  auto z = x + y;
-  IRVar2Const mu;
-  mu.var = y;
-  mu.int_val = 10;
-  auto zz = mu.Mutate(z);
-  std::ostringstream os;
-  os << zz;
-  CHECK(os.str() == "(x + 10)");
-}
-
-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

-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements.  See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership.  The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License.  You may obtain a copy of the License at
- *
- *   http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied.  See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-#include <dmlc/logging.h>
-#include <gtest/gtest.h>
-#include <tvm/ir_visitor.h>
-#include <tvm/ir_functor_ext.h>
-#include <tvm/ir_pass.h>
-
-TEST(IRVisitor, CountVar) {
-  using namespace tvm;
-  int n_var = 0;
-  Var x("x"), y;
-
-  auto z = x + 1 + y + y;
-  ir::PostOrderVisit(z, [&n_var](const ObjectRef& n) {
-      if (n.as<Variable>()) ++n_var;
-    });
-  CHECK_EQ(n_var, 2);
-}
-
-int main(int argc, char ** argv) {
-  testing::InitGoogleTest(&argc, argv);
-  testing::FLAGS_gtest_death_test_style = "threadsafe";
-  return RUN_ALL_TESTS();
-}