yzhliu commented on a change in pull request #4886: [WIP][POC]First pass a
defining at non-recursive Graph Vistor and Rewriter
URL: https://github.com/apache/incubator-tvm/pull/4886#discussion_r394726606
##########
File path: src/relay/ir/expr_functor.cc
##########
@@ -29,8 +29,135 @@
#include <tvm/relay/expr_functor.h>
#include <tvm/relay/pattern_functor.h>
+#include <stack>
+
namespace tvm {
namespace relay {
+/*!
+ * \brief A function to non-recursively traverse dataflow regions of a graph
+ *
+ * ExpandDatflow 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 pushed those
arguments to the stack
+ * and continues non-recursively 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.
+ */
+template <typename FCheckVisited, typename FVisitLeaf>
+void ExpandDataflow(Expr expr, FCheckVisited fcheck_visited, FVisitLeaf
fvisit_leaf) {
+ std::stack<std::pair<Expr, bool>> stack;
+ // The second state of the stack indicate whether the child has been
+ // expanded in the pre-order.
+ // NOTE: function will be inlined.
+ auto fpush_to_stack = [&fcheck_visited, &stack](const Expr& expr) {
+ if (!fcheck_visited(expr)) {
+ stack.push({expr, false});
+ }
+ };
+ fpush_to_stack(expr);
+ while (stack.size() > 0) {
+ auto node = stack.top().first;
+ // if this node was visited through another path
+ // after being added to the stack ignore it.
+ if (fcheck_visited(expr)) {
+ stack.pop();
+ } else if (stack.top().second) {
+ // all the children has 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
+ for (auto it = op->args.rbegin(); it != op->args.rend(); ++it) {
+ fpush_to_stack(*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
+ for (auto it = op->fields.rbegin(); it != op->fields.rend(); ++it) {
+ fpush_to_stack(*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.
+ // terminal leaf, directly use visited.
+ fvisit_leaf(node);
+ stack.pop();
+ }
+ }
+}
+
+DataflowVisitor::DataflowVisitor(int visit_limit) {
+ CHECK(visit_limit > 0) << "Dataflow visit limit must be greater than 0";
+ CHECK(visit_limit < 10) << "Dataflow visit limit must be less than 10";
+ visit_limit_ = visit_limit;
+}
+
+void DataflowVisitor::VisitLeaf(const Expr& expr) {
+ if (visit_counter_[expr.get()] == 0) {
+ ExprFunctor::VisitExpr(expr);
+ }
+ visit_counter_[expr.get()]++;
+}
+
+bool DataflowVisitor::CheckVisited(const Expr& expr) {
+ if (visit_counter_[expr.get()] < visit_limit_) {
+ return false;
+ } else {
+ visit_counter_[expr.get()]++;
+ return true;
+ }
+}
+
+void DataflowVisitor::VisitExpr(const Expr& expr) {
+ auto fcheck_visited = [this](const Expr& expr) { return
this->CheckVisited(expr); };
+ auto fvisit_leaf = [this](const Expr& expr) { return this->VisitLeaf(expr);
};
+ if (visit_counter_[expr.get()] < 1) {
Review comment:
not sure, do we need to respect `visit_limit_` here?
----------------------------------------------------------------
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
With regards,
Apache Git Services
