[GitHub] [tvm] mbrookhart commented on a change in pull request #7817: RelayTextPrinter is now non-recursive. ExpandDataflow refactored

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

mbrookhart commented on a change in pull request #7817:
URL: https://github.com/apache/tvm/pull/7817#discussion_r611863397



##########
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};
+};

Review comment:
       I don't understand why one would add the complexity of this vs a pair.

##########
File path: include/tvm/relay/expr_functor.h
##########
@@ -32,11 +32,12 @@
 #include <tvm/relay/function.h>
 #include <tvm/relay/op.h>
 
+#include <deque>
 #include <stack>

Review comment:
       This stack is no longer used

##########
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)));
     }
   };
+
   fpush_to_stack(expr);
   while (stack.size() > 0) {
-    auto node = stack.top().first;
-    if (fcheck_visited(node)) {
-      // if this node was visited through another path
-      // after being added to the stack ignore it.
-      stack.pop();
-    } else if (stack.top().second) {
-      // all the children have already been expanded.
-      // we can just run post order visit on it.
-      fvisit_leaf(node);
-      stack.pop();
-    } else if (const CallNode* op = node.as<CallNode>()) {
-      // mark expanded = true
-      stack.top().second = true;
-      // push the children to the stack in reverse order
-      // to match recursive processing order
+    v_info* front = &stack.front();
+    if (fcheck_visited(front->node)) {
+      stack.pop_front();
+    } else if (front->children_expanded) {
+      fvisit_leaf(front->node);
+      // TODO(d-smirnov): this is for compatibility with current implementation of MixedModeVisitor
+      stack.pop_front();
+    } else {
+      front->children_expanded = true;
+      for (auto e : fexpand_expr(front->node)) {
+        fpush_to_stack(e);
+      }
+    }
+  }
+}
+
+template <typename FCheckVisited, typename FVisitLeaf>
+void ExpandDataflow(Expr expr, FCheckVisited fcheck_visited, FVisitLeaf fvisit_leaf) {
+  auto fexpand_expr = [](const Expr& expr) {
+    std::vector<Expr> result;
+    if (const CallNode* op = expr.as<CallNode>()) {
       for (auto it = op->args.rbegin(); it != op->args.rend(); ++it) {
-        fpush_to_stack(*it);
+        result.push_back(*it);
       }
-      fpush_to_stack(op->op);
-    } else if (const TupleNode* op = node.as<TupleNode>()) {
-      stack.top().second = true;
-      // push the children to the stack in reverse order
-      // to match recursive processing order
+      result.push_back(op->op);
+    } else if (const TupleNode* op = expr.as<TupleNode>()) {
       for (auto it = op->fields.rbegin(); it != op->fields.rend(); ++it) {
-        fpush_to_stack(*it);
+        result.push_back(*it);
       }
-    } else if (const TupleGetItemNode* op = node.as<TupleGetItemNode>()) {
-      stack.top().second = true;
-      fpush_to_stack(op->tuple);
-    } else {
-      // No need to expand the children directly run visit.
-      fvisit_leaf(node);
-      stack.pop();
+    } else if (const TupleGetItemNode* op = expr.as<TupleGetItemNode>()) {
+      result.push_back(op->tuple);
     }
-  }
+    return std::move(result);

Review comment:
       Don't move on return, it prevents the compiler from optimizing

##########
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:
       Don't move on construction, it prevents certain compiler operations




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