[Unity][Transform] Canonicalize and use CSE between pattern matches

python/tvm/relax/transform/remove_redundant_reshape.py

@@ -66,13 +66,18 @@ def transform_function(self, func: Expr, mod: IRModule, ctx: PassContext) -> IRM
                 continue
 
             def rewriter(expr, matches):
-                args = matches[self.pattern]
+                arg = matches[self.input1]
+
                 if self.repeated_reshape in matches:
-                    return relax.op.reshape(matches[self.input1], args.args[1])
+                    output_shape = matches[self.repeated_reshape].args[1]
+                    return relax.op.reshape(arg, output_shape)
+
                 elif self.no_op_reshape in matches:
-                    if args.args[0].struct_info.shape:
-                        if structural_equal(args.args[0].struct_info.shape, args.args[1]):
-                            return args.args[0]
+                    output_shape = matches[self.no_op_reshape].args[1]
+                    if arg.struct_info.shape and structural_equal(
+                        arg.struct_info.shape, output_shape
+                    ):
+                        return arg
                 return expr
 
             updated_func = rewrite_call(self.pattern, rewriter, funct)

src/relax/ir/dataflow_matcher.cc

@@ -42,6 +42,7 @@
 #include <utility>
 #include <vector>
 
+#include "../transform/utils.h"
 #include "dataflow_matcher_impl.h"
 
 namespace tvm {
@@ -937,57 +938,76 @@ class PatternRewriter : ExprMutator {
     return Downcast<Function>(RemoveAllUnused(rewriter.VisitExpr(f)));
   }
 
-  void VisitBinding_(const VarBindingNode* binding) final {
-    bindings_.Set(binding->var, binding->value);
-    ExprMutator::VisitBinding_(binding);
-    if (auto it = memo_.find(binding->value.get()); it != memo_.end()) {
-      // We need to update the binding to pass to ExtractMatchedExpr, so that the rewritten
-      // expression can be subject to further pattern matchings.
-      bindings_.Set(binding->var, it->second);
+  Expr VisitExpr_(const SeqExprNode* seq) override {
+    if (ctx_) {
+      return ExprMutator::VisitExpr_(seq);
     }
-  }
 
-  Expr VisitExpr(const Expr& expr) final {
-    auto node = ExprMutator::VisitExpr(expr);
-    if (pattern_) {
-      if (auto matches_opt = ExtractMatchedExpr(pattern_.value(), node, bindings_)) {
-        Expr rewritten_expr = rewriter_func_(node, matches_opt.value());
-        if (!rewritten_expr.same_as(node)) {
-          rewritten_expr = builder_->Normalize(rewritten_expr);
-
-          // If the rewriter returns a variable (e.g. when rewriting
-          // from `R.add(x, R.const(0.0))` to `x`), the variable
-          // should be dereferenced to avoid trivial `var_2 = var_1`
-          // bindings.  This lookup is done using the builder_ instead
-          // of the bindings_, as the previous `builder_->Normalize`
-          // call may have introduced variable bindings.
-          if (auto opt_var = rewritten_expr.as<Var>()) {
-            if (auto binding = builder_->LookupBinding(opt_var.value())) {
-              rewritten_expr = binding.value();
-            }
-          }
-          memo_[expr.get()] = rewritten_expr;
-          return rewritten_expr;
+    auto cache = bindings_;
+    SeqExpr prev = GetRef<SeqExpr>(seq);
+
+    StructuralEqual struct_equal;
+
+    while (true) {
+      SeqExpr next = Downcast<SeqExpr>(builder_->Normalize(ExprMutator::VisitExpr_(prev.get())));
+      if (struct_equal(prev, next)) {
+        return std::move(next);
+      }
+
+      // Canonicalization may result in two previously-different
+      // expressions being recognized as identical.  Elimination of
+      // common subexpressions may result in trival var-to-var
+      // bindings that can be canonicalized.  Therefore, iterate the
+      // simplification steps until converged.
+      while (true) {
+        auto start_of_loop = next;
+        next = Downcast<SeqExpr>(CanonicalizeBindings(next));
+        next = Downcast<SeqExpr>(EliminateCommonSubexpr(next));
+        next = Downcast<SeqExpr>(RemoveAllUnused(next));
+        if (struct_equal(start_of_loop, next)) {
+          break;
         }
       }
+
+      if (struct_equal(prev, next)) {
+        return std::move(next);
+      }
+
+      // Reset all knowledge of bindings that were collected from
+      // this DataflowBlock.  The collected bindings are only after
+      // the point where they were collected, and we are repeating
+      // the mutation of this DataflowBlock.
+      bindings_ = cache;
+      prev = next;
     }
-    return node;
   }
 
-  BindingBlock VisitBindingBlock_(const DataflowBlockNode* block_node) final {
+  BindingBlock VisitBindingBlock_(const DataflowBlockNode* block_node) override {
     if (ctx_) {
       return RewriteDataflowBlockFixedPoint(GetRef<DataflowBlock>(block_node));
+    } else {
+      return ExprMutator::VisitBindingBlock_(block_node);
     }
+  }
 
-    DataflowBlock prev = GetRef<DataflowBlock>(block_node);
-    while (true) {
-      DataflowBlock next = Downcast<DataflowBlock>(ExprMutator::VisitBindingBlock_(prev.get()));
-      if (StructuralEqual()(prev, next)) {
-        return std::move(next);
-      } else {
-        prev = next;
+  void VisitBinding_(const VarBindingNode* binding) override {
+    auto expr = VisitExpr(binding->value);
+    bindings_.Set(binding->var, expr);
+    ReEmitBinding(binding, expr);
+  }
+
+  Expr VisitExpr(const Expr& expr) override {
+    auto node = ExprMutator::VisitExpr(expr);
+
+    if (pattern_) {
+      if (auto matches_opt = ExtractMatchedExpr(pattern_.value(), node, bindings_)) {
+        Expr rewritten_expr = rewriter_func_(node, matches_opt.value());
+        if (!rewritten_expr.same_as(node)) {
+          return builder_->Normalize(rewritten_expr);
+        }
       }
     }
+    return node;
   }
 
  private:
@@ -1076,7 +1096,6 @@ class PatternRewriter : ExprMutator {
   PackedFunc rewriter_func_;
   std::unordered_set<const VarNode*> params_;
   Map<Var, Expr> bindings_;
-  std::unordered_map<const Object*, Expr> memo_;
 };
 
 Function RewriteBindings(const PatternContext& ctx, PackedFunc rewriter, Function f) {

src/relax/transform/eliminate_common_subexpr.cc

@@ -105,7 +105,27 @@ class SubexprCounter : public ExprVisitor {
         count_map_[e] = count + 1;
       }
     }
-    ExprVisitor::VisitExpr(e);
+
+    // Only visit the interior of objects that we might still keep
+    // around.  Otherwise, double-counting these would lead to extra
+    // variable bindings.
+    //
+    // Before:
+    //     y = f(a+b)
+    //     z = f(a+b)
+    //
+    // Expected:
+    //     y = f(a+b)  // De-duped from (y==z)
+    //     z = y
+    //
+    // Erroneous output:
+    //     c = a+b    // Incorrect, a+b only has a single usage.
+    //     y = f(c)   // De-duped from
+    //     z = y
+    //
+    if (auto it = count_map_.find(e); it == count_map_.end() || it->second < 2) {
+      ExprVisitor::VisitExpr(e);
+    }
   }
 
   // do not visit inner functions: we will do CSE within those