[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/op/_transform.py

@@ -52,7 +52,7 @@
 _reg.register_injective_schedule("take")
 _reg.register_injective_schedule("transpose")
 _reg.register_injective_schedule("stack")
-_reg.register_injective_schedule("_contrib_reverse_reshape")
+_reg.register_injective_schedule("contrib_reverse_reshape")
 _reg.register_injective_schedule("gather")
 _reg.register_injective_schedule("gather_nd")
 _reg.register_injective_schedule("sequence_mask")

python/tvm/relay/op/transform.py

@@ -908,7 +908,7 @@ def reverse_reshape(data, newshape):
     """
     if isinstance(newshape, int):
         newshape = [newshape]
-    return _make._contrib_reverse_reshape(data, list(newshape))
+    return _make.contrib_reverse_reshape(data, list(newshape))
 
 
 def gather(data, axis, indices):

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,14 +739,17 @@ 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 = this->VisitExpr(post);
         count++;
       }
-    } while (last != post || count >= 100);
+    } while (!StructuralEqual()(last, post) || count >= 100);
     if (count >= 100) {
       throw("Observed 100 rewrite passes, possible conflicting passes?");
     }
@@ -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.

src/relay/op/tensor/transform.cc

@@ -2573,13 +2573,13 @@ Expr MakeReverseReshape(Expr data, Array<Integer> newshape) {
   auto attrs = make_object<ReshapeAttrs>();
   attrs->newshape = std::move(newshape);
   attrs->reverse = true;
-  static const Op& op = Op::Get("_contrib_reverse_reshape");
+  static const Op& op = Op::Get("contrib_reverse_reshape");
   return Call(op, {data}, Attrs(attrs), {});
 }
 
-TVM_REGISTER_GLOBAL("relay.op._make._contrib_reverse_reshape").set_body_typed(MakeReverseReshape);
+TVM_REGISTER_GLOBAL("relay.op._make.contrib_reverse_reshape").set_body_typed(MakeReverseReshape);
 
-RELAY_REGISTER_OP("_contrib_reverse_reshape")
+RELAY_REGISTER_OP("contrib_reverse_reshape")
     .describe(R"code(Reshapes the input array where the special values are inferred from
 right to left.
 

src/relay/transforms/simplify_expr.cc

@@ -0,0 +1,124 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file src/relay/transforms/simplify_expr.cc
+ * \brief A pass for simplifying the Relay expression.
+ */
+
+#include <tvm/relay/expr.h>
+#include <tvm/relay/expr_functor.h>
+#include <tvm/relay/transform.h>
+#include <tvm/relay/dataflow_matcher.h>
+#include <tvm/support/logging.h>
+#include "../op/tensor/transform.h"
+
+namespace tvm {
+namespace relay {
+
+static Op reshape_op = Op::Get("reshape");
+static Op reverse_reshape_op = Op::Get("contrib_reverse_reshape");
+
+/*!
+ * \brief SimplifyReshape matches the pattern of consecutive reshape or reverse_reshape ops,
+ *   and merges into one reshape op.
+ */
+class SimplifyReshape {
+ public:
+  SimplifyReshape() {
+    x_ = WildcardPattern(make_object<WildcardPatternNode>());
+    auto reshape1 = AltPattern(ExprPattern(reshape_op), ExprPattern(reverse_reshape_op));
+    auto reshape2 = AltPattern(ExprPattern(reshape_op), ExprPattern(reverse_reshape_op));
+    pattern_ = CallPattern(reshape1, {CallPattern(reshape2, {x_}, Attrs{}, {})}, Attrs{}, {});
+  }
+
+  Expr callback(const Expr& pre, const Expr& post, const Map<DFPattern, Array<Expr>>& node_map) {
+    auto x = node_map[x_][0];
+    bool const_shape = true;
+    Array<Integer> newshape;
+    for (auto dim : Downcast<TensorType>(pre->checked_type())->shape) {
+      if (dim.as<IntImmNode>() == nullptr) {
+        const_shape = false;
+        break;
+      }
+      newshape.push_back(Downcast<Integer>(dim));
+    }
+    if (const_shape) {
+      return MakeReshape(x, newshape);
+    }
+    return post;
+  }
+
+  DFPattern pattern() const { return pattern_; }
+
+ private:
+  /*! \brief Pattern input */
+  DFPattern x_;
+  /*! \brief Pattern for consecutive reshape or reverse_reshape ops */
+  DFPattern pattern_;
+
+};
+
+/*!
+ * \brief ExprSimplifier simplifies the Relay expression.
+ */
+class ExprSimplifier {
+ public:
+  ExprSimplifier() {
+    auto reshape_func = [this](TVMArgs args, TVMRetValue* rv) {
+      Expr pre = args[0];
+      Expr post = args[1];
+      Map<DFPattern, Array<Expr>> node_map = args[2];
+      *rv = simplify_reshape_.callback(pre, post, node_map);
+    };
+    callbacks_.push_back(DFPatternCallback(simplify_reshape_.pattern(), PackedFunc(reshape_func),
+                                           true));
+  }
+
+  Expr Simplify(const Expr& expr) {
+    return RewritePatterns(callbacks_, expr);
+  }
+
+ private:
+  /*! \brief Simplify reshape pattern */
+  SimplifyReshape simplify_reshape_;
+  /*! \brief Callbacks for expr simplification */
+  Array<DFPatternCallback> callbacks_;
+};
+
+Expr SimplifyExpr(const Expr& expr, const IRModule& module) {
+  return ExprSimplifier().Simplify(expr);
+}
+
+namespace transform {
+
+Pass SimplifyExpr() {
+  runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func =
+      [=](Function f, IRModule m, PassContext pc) {
+        return Downcast<Function>(SimplifyExpr(f, m));
+      };
+  return CreateFunctionPass(pass_func, 0, "SimplifyExpr", {"InferType"});
+}
+
+TVM_REGISTER_GLOBAL("relay._transform.SimplifyExpr").set_body_typed(SimplifyExpr);
+
+}  // namespace transform
+
+}  // namespace relay
+}  // namespace tvm