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

[GitBox <gi...@apache.org>]

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



##########
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:
       Sounds good. I am going to move it to analysis.h and change it from std::function to FLambda




----------------------------------------------------------------
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