[RELAY] Non-recursive Graph Vistor and Rewriter

include/tvm/relay/analysis.h

@@ -30,6 +30,7 @@
 #include <tvm/ir/module.h>
 #include <tvm/relay/type.h>
 #include <string>
+#include <unordered_map>
 
 namespace tvm {
 namespace relay {
@@ -225,6 +226,16 @@ TVM_DLL Map<Expr, Integer> CollectDeviceAnnotationOps(const Expr& expr);
  */
 TVM_DLL Array<Pattern> UnmatchedCases(const Match& match, const IRModule& mod);
 
+/*!
+ * \brief Get reference counter of each internal ExprNode in body.
+ *
+ * \param body The body expression.
+ *
+ * \return The reference count mapping.
+ */
+TVM_DLL std::unordered_map<const Object*, size_t>
+GetExprRefCount(const Expr& body);
+
 }  // namespace relay
 }  // namespace tvm
 

include/tvm/relay/expr_functor.h

@@ -232,6 +232,188 @@ class ExprMutator
   std::unordered_map<Expr, Expr, ObjectHash, ObjectEqual> memo_;
 };
 
+/*!
+ * \brief A wrapper around ExprVisitor which traverses the Dataflow Normal AST.
+ *
+ * MixedModeVisitor treats Expr as dataflow graph, and visits in post-DFS order
+ *
+ * MixedModeVisitor provides the same recursive API as ExprVisitor, and uses
+ * recursion to traverse most forms of the IR, but under the hood it expands nested dataflow regions
+ * of the graph and processes them iteratatively to prevent stack overflows
+ */
+class MixedModeVisitor : public ::tvm::relay::ExprVisitor {
+ public:
+  /*! \brief The constructor of MixedModeVisitor
+   *  \param visit_limit The number of times to allow visitation to a note. Usually 1, ocassionally
+   * higher (i.e., 2 for dead code elimiation), limited to 10 as a sanity check.
+   */
+  MixedModeVisitor(int visit_limit = 1);
+
+  /*!
+   * \brief VisitExpr is finalized to preserve call expansion of dataflow regions
+   */
+  void VisitExpr(const Expr& expr) final;
+  void VisitExpr_(const CallNode* op) override;
+  void VisitExpr_(const TupleNode* op) override;
+  void VisitExpr_(const TupleGetItemNode* op) override;
+
+ protected:
+  /*!
+   * \brief A function to apply when reaching a leaf of the graph non-recursively
+   */
+  virtual void VisitLeaf(const Expr& expr);
+  /*!
+   * \brief A function to determine if an expression has already been visited or needs to be
+   * re-visited
+   */
+  virtual bool CheckVisited(const Expr& expr);
+  /*!
+   * \brief The max number of times to visit a node
+   */
+  size_t visit_limit_;
+};
+
+/*! \brief Non-recursive DFS Graph Traversal for Custom Rewriting Passes
+ *
+ * MixedModeMutator treats Expr as dataflow graph, and only Rewrites each Expr once.
+ * The mutated results are memoized in a map and reused so that
+ * local transformation on the dataflow preserves the graph structure.
+ *
+ * MixedModeMutator provides the same recursive API as ExprMutator, and uses
+ * recursion to traverse most forms of the IR, but under the hood it expands nested dataflow regions
+ * of the graph and processes them iteratatively to prevent stack overflows
+ *
+ * Uses Rewrite_ API of ExprRewriter for a cleaner split between recrusive and non-recursive behavior.
+ */
+class MixedModeMutator : public ::tvm::relay::ExprMutator {
+ public:
+  Expr VisitExpr(const Expr& expr) final;
+  virtual Expr DispatchVisitExpr(const Expr& expr);
+  Expr VisitExpr_(const TupleNode* op) final { return Rewrite(op); };
+  Expr VisitExpr_(const CallNode* call_node) final { return Rewrite(call_node); };
+  Expr VisitExpr_(const TupleGetItemNode* op) final { return Rewrite(op); };
+  /*!
+   *  Users should override Rewrite_ methods to implement their pass. Rewrite_ functions will be
+   * able to rewrite the op only with data about the original node `pre` and the same node with
+   * modified inputs `post` and should not recurse.
+   *
+   * \param pre The expression node before rewriting.
+   * \param post The expression with rewritten inputs.
+   */
+  virtual Expr Rewrite_(const TupleNode* pre, const Expr& post) { return post;}
+  virtual Expr Rewrite_(const CallNode* pre, const Expr& post) { return post; }
+  virtual Expr Rewrite_(const TupleGetItemNode* pre, const Expr& post) { return post; }
+
+ protected:
+  /*! \brief Implement Rewrite API by calling ExprMutator's VisitExpr_(op) to get a `post` node with
+   * changed inputs.
+   */
+  template <typename T>
+  Expr Rewrite(const T* op) {
+    Expr post = ExprMutator::VisitExpr_(op);
+    return Rewrite_(op, post);
+  }
+
+  virtual void VisitLeaf(const Expr& expr);
+  virtual bool CheckVisited(const Expr& expr);
+};
+
+/*! \brief A non-iterating Expression Rewriter
+ *
+ *  ExprRewriter provides a Rewrite interface for modifying graphs in Post-DFS order.
+ *  The expectation is that ExprRewriter objects will be passed to PostOrderRewrite, which will
+ * non-recursively unroll the graph and call Rewriting on inputs. It will then pass the original
+ * node, called `pre`, and a node recreated with any alterned inputs, called `post`, to the
+ * ExprRewriter. The ExprRewriter can then use the information in those two nodes to do more complex
+ * graph rewriting.
+ */
+#define RELAY_EXPR_REWRITER_DISPATCH(OP)                                                   \
+  vtable.template set_dispatch<OP>([](const ObjectRef& n, TSelf* self, const Expr& post) { \
+    return self->Rewrite_(static_cast<const OP*>(n.get()), post);                          \
+  });
+
+#define EXPR_REWRITER_REWRITE_DEFAULT \
+  { return post; }
+
+class ExprRewriter {
+ private:
+  using TSelf = ExprRewriter;
+  using FType = tvm::NodeFunctor<Expr(const ObjectRef& n, TSelf* self, const Expr& post)>;
+
+ public:
+  /*! \brief virtual destructor */
+  virtual ~ExprRewriter() {}
+  /*!
+   * \brief Same as call.
+   * \param pre The expression node before rewriting.
+   * \param post The expression node with rewritten inputs.
+   * \return The result of the call
+   */
+  Expr operator()(const Expr& pre, const Expr& post) {
+    return Rewrite(pre, post);
+  }
+  /*!
+   * \brief The functor call.
+   * \param pre The expression node before rewriting.
+   * \param post The expression node with rewritten inputs.
+   * \return The result of the call
+   */
+  virtual Expr Rewrite(const Expr& pre, const Expr& post) {
+    CHECK(pre.defined());
+    static FType vtable = InitVTable();
+    return vtable(pre, this, post);
+  }
+  // Functions that can be overriden by subclass, should not recurse
+  virtual Expr Rewrite_(const VarNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
+  virtual Expr Rewrite_(const GlobalVarNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
+  virtual Expr Rewrite_(const ConstantNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
+  virtual Expr Rewrite_(const TupleNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
+  virtual Expr Rewrite_(const FunctionNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
+  virtual Expr Rewrite_(const CallNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
+  virtual Expr Rewrite_(const LetNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
+  virtual Expr Rewrite_(const IfNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
+  virtual Expr Rewrite_(const OpNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
+  virtual Expr Rewrite_(const TupleGetItemNode* pre,
+                        const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
+  virtual Expr Rewrite_(const RefCreateNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
+  virtual Expr Rewrite_(const RefReadNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
+  virtual Expr Rewrite_(const RefWriteNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
+  virtual Expr Rewrite_(const ConstructorNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
+  virtual Expr Rewrite_(const MatchNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
+
+ private:
+  // initialize the vtable.
+  static FType InitVTable() {
+    FType vtable;
+    // Set dispatch
+    RELAY_EXPR_REWRITER_DISPATCH(ConstantNode);
+    RELAY_EXPR_REWRITER_DISPATCH(TupleNode);
+    RELAY_EXPR_REWRITER_DISPATCH(VarNode);
+    RELAY_EXPR_REWRITER_DISPATCH(GlobalVarNode);
+    RELAY_EXPR_REWRITER_DISPATCH(FunctionNode);
+    RELAY_EXPR_REWRITER_DISPATCH(CallNode);
+    RELAY_EXPR_REWRITER_DISPATCH(LetNode);
+    RELAY_EXPR_REWRITER_DISPATCH(IfNode);
+    RELAY_EXPR_REWRITER_DISPATCH(OpNode);
+    RELAY_EXPR_REWRITER_DISPATCH(TupleGetItemNode);
+    RELAY_EXPR_REWRITER_DISPATCH(RefCreateNode);
+    RELAY_EXPR_REWRITER_DISPATCH(RefReadNode);
+    RELAY_EXPR_REWRITER_DISPATCH(RefWriteNode);
+    RELAY_EXPR_REWRITER_DISPATCH(ConstructorNode);
+    RELAY_EXPR_REWRITER_DISPATCH(MatchNode);
+    return vtable;
+  }
+};
+
+/*! \brief Non-recursive DFS Graph Traversal for Custom Rewriting Passes
+ *
+ *  PostOrderRewrite does a non-recursive traversal of the graph in Post-DFS order and calls the
+ * ExprRewriter's Rewrite functions on nodes once their inputs are rewritten. At each rewrite call,
+ * PostOrderRewrite provides the original node and the node with altered inputs for use by the
+ * ExprRewriter.
+ */
+Expr PostOrderRewrite(const Expr& expr, ExprRewriter* rewriter);
+
 /*!
  * \brief recursively visit the ir in post DFS order node, apply fvisit
  * Each node is guaranteed to be visited only once.

src/relay/analysis/util.cc

@@ -330,7 +330,7 @@ TVM_REGISTER_GLOBAL("relay.analysis.all_type_vars")
  */
 std::unordered_map<const Object*, size_t>
 GetExprRefCount(const Expr& body) {
-  class ExprRefCounter : private ExprVisitor {
+  class ExprRefCounter : private MixedModeVisitor {
    public:
     std::unordered_map<const Object*, size_t>
     Get(const Expr& body) {