d-smirnov commented on a change in pull request #7817:
URL: https://github.com/apache/tvm/pull/7817#discussion_r611923032
##########
File path: include/tvm/relay/expr_functor.h
##########
@@ -410,72 +414,82 @@ Expr PostOrderRewrite(const Expr& expr, ExprRewriter*
rewriter);
*/
void PostOrderVisit(const Expr& node, std::function<void(const Expr&)> fvisit);
+/*!
+ * \brief A struct to keep info of traversed expr in ExpandDataflow function
+ */
+struct v_info {
+ explicit v_info(Expr node_) : node{node_} {}
+ v_info(Expr node_, bool children_expanded_)
+ : node{node_}, children_expanded{children_expanded_} {};
+ Expr node{};
+ bool children_expanded{false};
+};
+
/*!
* \brief A function to iteratively traverse dataflow regions of a graph
*
* ExpandDataflow manually manages a stack and performs DFS to determine the
processing
* order of nodes in an input graph.
*
- * If it finds a dataflow node (Call, Tuple, TupleGetItem), it checks if the
arguments to that node
- * need to be processed via fcheck_visited. If so, the function pushes those
arguments to the stack
- * and continues iteratively to process the top of the stack. When it finds a
node that doesn't
- * match the dataflow types, or a node who's inputs have all been processed,
it visits the current
- * leaf via fvisit_leaf.
+ * By default fexpand_expr implemented in a way that if it finds a dataflow
node (Call, Tuple,
+ * TupleGetItem), it checks if the arguments to that node need to be processed
via fcheck_visited.
+ * If so, the function pushes those arguments to the stack and continues
iteratively to process
+ * the top of the stack. When it finds a node that doesn't match the dataflow
types, or a node who's
+ * inputs have all been processed, it visits the current leaf via fvisit_leaf.
*
* This function should be used internally to other classes to implement
mixed-mode traversals. The
* expectation is that fvisit_leaf will perform recursive analysis within
mixed-mode traversal if it
* hits a non-dataflow node.
*
- * fcheck_visited and fvisit_leaf are templated to encourage compiler inlining.
+ * fcheck_visited, fvisit_leaf and fexpand_expr are templated to encourage
reusing.
*/
-template <typename FCheckVisited, typename FVisitLeaf>
-void ExpandDataflow(Expr expr, FCheckVisited fcheck_visited, FVisitLeaf
fvisit_leaf) {
- std::stack<std::pair<Expr, bool>> stack;
+template <typename FCheckVisited, typename FVisitLeaf, typename FExpandExpr>
+void ExpandDataflow(Expr expr, FCheckVisited fcheck_visited, FVisitLeaf
fvisit_leaf,
+ FExpandExpr fexpand_expr) {
+ std::deque<v_info> stack;
auto fpush_to_stack = [&fcheck_visited, &stack](const Expr& expr) {
- // The second state of the stack indicate whether the child has been
- // expanded in the pre-order.
- // NOTE: function will be inlined.
if (!fcheck_visited(expr)) {
- stack.push({expr, false});
+ stack.push_front(std::move(v_info(expr)));
Review comment:
Elaborate, please. As well as why one should want these optimisations to
be used.
--
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:
us...@infra.apache.org
