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:
[email protected]
