You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2021/03/10 15:11:53 UTC

[GitHub] [tvm] tqchen commented on a change in pull request #7627: [TIR] Add PreOrderVisit and VisitPrimFuncs

tqchen commented on a change in pull request #7627:
URL: https://github.com/apache/tvm/pull/7627#discussion_r591603180



##########
File path: include/tvm/tir/stmt_functor.h
##########
@@ -386,6 +394,24 @@ inline T Substitute(T input, const std::unordered_map<const VarNode*, PrimExpr>&
   return Substitute(std::move(input), vmap);
 }
 
+/*!
+ * \brief Recursively visit the IR in pre DFS order node, apply fvisit.
+ * If fvisit returns false, it stops visit the children of the current node.
+ * \param node The IR to be visited, can be either Stmt or Expr.
+ * \param fvisit The visitor function to be applied. If fvisit returns false, it stops visit the

Review comment:
       it won't visit the children of the node

##########
File path: include/tvm/tir/stmt_functor.h
##########
@@ -386,6 +394,24 @@ inline T Substitute(T input, const std::unordered_map<const VarNode*, PrimExpr>&
   return Substitute(std::move(input), vmap);
 }
 
+/*!
+ * \brief Recursively visit the IR in pre DFS order node, apply fvisit.
+ * If fvisit returns false, it stops visit the children of the current node.
+ * \param node The IR to be visited, can be either Stmt or Expr.
+ * \param fvisit The visitor function to be applied. If fvisit returns false, it stops visit the
+ * children on that node.
+ */
+TVM_DLL void PreOrderVisit(const ObjectRef& node,
+                           const std::function<bool(const ObjectRef&)>& fvisit);
+
+/*!
+ * \brief Visit the PrimFuncs in the IRModule
+ * \param mod The IRModule to be visited
+ * \param fvisit The visitor to the PrimFuncs in the IRModule
+ */
+TVM_DLL void VisitPrimFuncs(const IRModule& mod,

Review comment:
       Consider make an inline function so we can templatize fvisit, and let compiler inline that lambda

##########
File path: include/tvm/tir/stmt_functor.h
##########
@@ -386,6 +394,24 @@ inline T Substitute(T input, const std::unordered_map<const VarNode*, PrimExpr>&
   return Substitute(std::move(input), vmap);
 }
 
+/*!
+ * \brief Recursively visit the IR in pre DFS order node, apply fvisit.
+ * If fvisit returns false, it stops visit the children of the current node.
+ * \param node The IR to be visited, can be either Stmt or Expr.
+ * \param fvisit The visitor function to be applied. If fvisit returns false, it stops visit the
+ * children on that node.
+ */
+TVM_DLL void PreOrderVisit(const ObjectRef& node,

Review comment:
       node=> stmt_or_expr

##########
File path: include/tvm/tir/stmt_functor.h
##########
@@ -386,6 +394,24 @@ inline T Substitute(T input, const std::unordered_map<const VarNode*, PrimExpr>&
   return Substitute(std::move(input), vmap);
 }
 
+/*!
+ * \brief Recursively visit the IR in pre DFS order node, apply fvisit.
+ * If fvisit returns false, it stops visit the children of the current node.
+ * \param node The IR to be visited, can be either Stmt or Expr.
+ * \param fvisit The visitor function to be applied. If fvisit returns false, it stops visit the
+ * children on that node.
+ */

Review comment:
       please add a testcase to https://github.com/apache/tvm/blob/main/tests/cpp/ir_functor_test.cc

##########
File path: include/tvm/tir/stmt_functor.h
##########
@@ -386,6 +394,24 @@ inline T Substitute(T input, const std::unordered_map<const VarNode*, PrimExpr>&
   return Substitute(std::move(input), vmap);
 }
 
+/*!
+ * \brief Recursively visit the IR in pre DFS order node, apply fvisit.
+ * If fvisit returns false, it stops visit the children of the current node.
+ * \param node The IR to be visited, can be either Stmt or Expr.
+ * \param fvisit The visitor function to be applied. If fvisit returns false, it stops visit the
+ * children on that node.
+ */
+TVM_DLL void PreOrderVisit(const ObjectRef& node,
+                           const std::function<bool(const ObjectRef&)>& fvisit);
+
+/*!
+ * \brief Visit the PrimFuncs in the IRModule
+ * \param mod The IRModule to be visited
+ * \param fvisit The visitor to the PrimFuncs in the IRModule
+ */
+TVM_DLL void VisitPrimFuncs(const IRModule& mod,

Review comment:
       Consider move to https://github.com/apache/tvm/blob/main/include/tvm/tir/analysis.h

##########
File path: include/tvm/tir/stmt_functor.h
##########
@@ -36,7 +36,15 @@
 #include <utility>
 
 namespace tvm {
+
+// Forward declaration
+class IRModule;

Review comment:
       simply include ir/module.h




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org