[Relay][Pass] Merge two consecutive reshape ops

include/tvm/relay/dataflow_matcher.h

@@ -42,11 +42,16 @@ class DFPatternCallback;
 class DFPatternCallbackNode : public Object {
  public:
   /*! \brief Pattern this callback matches */
-  DFPattern pattern_;
+  DFPattern pattern;
   /*! \brief Function to call when finding a matched expression */
-  PackedFunc function_;
+  PackedFunc function;
+  /*! \brief Require InferType to be run before the callback */
+  bool require_type;
 
-  void VisitAttrs(tvm::AttrVisitor* v) {}
+  void VisitAttrs(tvm::AttrVisitor* v) {
+    v->Visit("pattern", &pattern);
+    v->Visit("require_type", &require_type);
+  }
 
   static constexpr const char* _type_key = "DFPatternCallbackNode";
   TVM_DECLARE_BASE_OBJECT_INFO(DFPatternCallbackNode, Object);
@@ -58,7 +63,7 @@ class DFPatternCallbackNode : public Object {
  */
 class DFPatternCallback : public ObjectRef {
  public:
-  TVM_DLL DFPatternCallback(DFPattern pattern, PackedFunc callback);
+  TVM_DLL DFPatternCallback(DFPattern pattern, PackedFunc callback, bool require_type);
   TVM_DEFINE_OBJECT_REF_METHODS(DFPatternCallback, ObjectRef, DFPatternCallbackNode);
 };
 

python/tvm/relay/dataflow_pattern/__init__.py

@@ -687,7 +687,15 @@ class DFPatternCallback:
     the callback returns.
 
     Users are expect to inherit from this class and provide a "self.pattern" to match
+
+    Parameters
+    ----------
+    require_type: bool
+        Whether InferType is required to be run before the callback.
     """
+    def __init__(self, require_type=False):
+        self.pattern = None
+        self.require_type = require_type
 
     def rewrite(self, expr: Expr) -> Expr:
         """
@@ -727,8 +735,8 @@ def callback(self, pre: Expr, post: Expr, node_map: tvm.ir.container.Map) -> Exp
 
 class _DFPatternCallback(Object):
     """C++ implemenation"""
-    def __init__(self, pattern, callback):
-        self.__init_handle_by_constructor__(ffi.DFPatternCallback, pattern, callback)
+    def __init__(self, pattern, callback, require_type):
+        self.__init_handle_by_constructor__(ffi.DFPatternCallback, pattern, callback, require_type)
 
 
 def rewrite(callbacks, expr: Expr) -> Expr:
@@ -748,11 +756,14 @@ def rewrite(callbacks, expr: Expr) -> Expr:
         The Expression with matched subgraphs rewritten by the callbacks.
     """
     if isinstance(callbacks, DFPatternCallback):
-        tmp = [_DFPatternCallback(callbacks.pattern, callbacks.callback)]
+        assert callbacks.pattern is not None
+        tmp = [_DFPatternCallback(callbacks.pattern, callbacks.callback, callbacks.require_type)]
     else:
         tmp = []
         for callback in callbacks:
-            tmp.append(_DFPatternCallback(callback.pattern, callback.callback))
+            assert callback.pattern is not None
+            tmp.append(_DFPatternCallback(callback.pattern, callback.callback,
+                                          callback.require_type))
 
     return ffi.rewrite(tmp, expr)
 

python/tvm/relay/transform/__init__.py

@@ -19,4 +19,3 @@
 # transformation passes
 from .transform import *
 from . import memory_alloc
-from .simplify_expr import SimplifyExpr

python/tvm/relay/transform/transform.py

@@ -909,6 +909,7 @@ def DenseToSparse(weight_name, weight_shape):
     """
     return _ffi_api.DenseToSparse(weight_name, weight_shape)
 
+
 def SimplifyFCTranspose(target_weight_name):
     """
     Rewrite ```y = nn.dense(x, transpose(w, [1, 0]))``` to ```y = nn.dense(x, wt)```
@@ -926,3 +927,15 @@ def SimplifyFCTranspose(target_weight_name):
         The registered SimplifyFCTranspose pass.
     """
     return _ffi_api.SimplifyFCTranspose(target_weight_name)
+
+
+def SimplifyExpr():
+    """
+    Simplify the Relay expression, including merging consecutive reshapes.
+
+    Returns
+    -------
+    ret : tvm.transform.Pass
+        The registered SimplifyExpr pass.
+    """
+    return _ffi_api.SimplifyExpr()

src/relay/ir/dataflow_matcher.cc

@@ -436,7 +436,8 @@ bool MatchPattern(DFPattern pattern, Expr expr) {
 
 TVM_REGISTER_GLOBAL("relay.dataflow_pattern.match").set_body_typed(MatchPattern);
 
-/* \brief PatternGrouper does pre-rewriting pattern matching and analysis
+/*!
+ * \brief PatternGrouper does pre-rewriting pattern matching and analysis
  *
  * This class creates a number of groups of matched expressions, ensures they don't overlap, and
  * returns them to the caller for post-analysis rewriting.
@@ -446,7 +447,7 @@ TVM_REGISTER_GLOBAL("relay.dataflow_pattern.match").set_body_typed(MatchPattern)
  */
 class PatternGrouper {
  public:
-  /* \brief Internal Group class for storing analysis */
+  /*! \brief Internal Group class for storing analysis */
   struct Group {
     Expr root_node;
     int gid;
@@ -456,11 +457,11 @@ class PatternGrouper {
     Array<Expr> args;
   };
 
-  /* \brief Return the group assignments of expressions */
+  /*! \brief Return the group assignments of expressions */
   const std::unordered_map<Expr, int, ObjectPtrHash, ObjectPtrEqual>& GetGIDAssignments() {
     return gid_assignments_;
   }
-  /* \brief Group expressions that match the pattern */
+  /*! \brief Group expressions that match the pattern */
   const std::unordered_map<int, Group>& GroupMatches(const DFPattern& pattern, const Expr& pre) {
     groups_.clear();
     gid_assignments_.clear();
@@ -474,7 +475,7 @@ class PatternGrouper {
   }
 
  protected:
-  /* \brief Iteratively traverse the Expression in pre-order to find subgraphs
+  /*! \brief Iteratively traverse the Expression in pre-order to find subgraphs
    *
    * If we traverse the graph in post-order, we can run into situtations where a small subgraph will
    * match the pattern. Due to options like AltPattern, a larger subgraph with more nodes later in
@@ -501,7 +502,7 @@ class PatternGrouper {
       }
     }
   }
-  /* \brief Creates a new set of nodes based on Group inputs, used to create functions and perform
+  /*! \brief Creates a new set of nodes based on Group inputs, used to create functions and perform
    * group overlap analysis */
   class MatchExtractor : public ExprMutator {
    public:
@@ -563,7 +564,7 @@ class PatternGrouper {
     const std::unordered_map<Expr, Var, ObjectPtrHash, ObjectPtrEqual> inputs_;
   };
 
-  /* \brief Create a group based on a matched expression */
+  /*! \brief Create a group based on a matched expression */
   void CreateGroup(const Expr& expr) {
     int var_number = 0;
 
@@ -661,7 +662,7 @@ class PatternGrouper {
     groups_[group.gid] = std::move(group);
   }
 
-  /* \brief EmbedConst implements rules for embedding constants into partitioned functions or
+  /*! \brief EmbedConst implements rules for embedding constants into partitioned functions or
    * lifting them into the function arguments.
    *
    * The rules depend on what pattern the ConstantNode matched.
@@ -703,21 +704,23 @@ class PatternGrouper {
 
 // Rewrite
 
-DFPatternCallback::DFPatternCallback(DFPattern pattern, PackedFunc function) {
+DFPatternCallback::DFPatternCallback(DFPattern pattern, PackedFunc function, bool require_type) {
   ObjectPtr<DFPatternCallbackNode> n = make_object<DFPatternCallbackNode>();
-  n->pattern_ = std::move(pattern);
-  n->function_ = std::move(function);
+  n->pattern = std::move(pattern);
+  n->function = std::move(function);
+  n->require_type = require_type;
   data_ = std::move(n);
 }
 
 TVM_REGISTER_NODE_TYPE(DFPatternCallbackNode);
 
 TVM_REGISTER_GLOBAL("relay.dataflow_pattern.DFPatternCallback")
-    .set_body_typed([](DFPattern pattern, PackedFunc function) {
-      return DFPatternCallback(pattern, function);
+    .set_body_typed([](DFPattern pattern, PackedFunc function, bool require_type) {
+      return DFPatternCallback(pattern, function, require_type);
     });
 
-/* \brief PatternRewriter rewrites the expression by finding matches and allowing user callback
+/*!
+ * \brief PatternRewriter rewrites the expression by finding matches and allowing user callback
  * function to rewrite those matches
  *
  * The class uses PatternGrouper to support the dominator pattern.
@@ -736,11 +739,14 @@ class PatternRewriter : protected MixedModeMutator {
       last = post;
       for (auto callback : callbacks) {
         callback_ = callback;
+        if (callback_->require_type) {
+          post = InferType(post);
+        }
         auto grouper = PatternGrouper();
-        groups_ = grouper.GroupMatches(callback_->pattern_, post);
+        groups_ = grouper.GroupMatches(callback_->pattern, post);
         gid_assignments_ = grouper.GetGIDAssignments();
         memo_.clear();
-        post = InferType(this->VisitExpr(post));
+        post = this->VisitExpr(post);
         count++;
       }
     } while (!StructuralEqual()(last, post) || count >= 100);
@@ -765,7 +771,7 @@ class PatternRewriter : protected MixedModeMutator {
         node_map.insert({kv.first, tmp});
       }
       // run the user callback function
-      return callback_->function_(pre, post, Map<DFPattern, Array<Expr>>(node_map));
+      return callback_->function(pre, post, Map<DFPattern, Array<Expr>>(node_map));
     }
     return post;
   }
@@ -781,7 +787,8 @@ Expr RewritePatterns(Array<DFPatternCallback> callbacks, Expr expr) {
 
 TVM_REGISTER_GLOBAL("relay.dataflow_pattern.rewrite").set_body_typed(RewritePatterns);
 
-/* \brief PatternPartitioner replaces expressions that match a pattern with function call that
+/*!
+ * \brief PatternPartitioner replaces expressions that match a pattern with function call that
  * perform the same computation but allow for further analysis and lowering.
  *
  * The class uses PatternGrouper to support the dominator pattern.