[Unity] Avoid trivial var2 = var1 bindings in pattern matcher

src/relax/ir/dataflow_matcher.cc

@@ -818,16 +818,31 @@ class PatternRewriter : ExprMutator {
     }
   }
 
-  Expr VisitExpr_(const CallNode* call_node) final {
-    auto call = ExprMutator::VisitExpr_(call_node);
-    if (!pattern_) {
-      return call;
-    } else if (auto matches_opt = ExtractMatchedExpr(pattern_.value(), call, bindings_)) {
-      auto rewriten_expr = rewriter_func_(call, matches_opt.value());
-      memo_[call_node] = rewriten_expr;
-      return rewriten_expr;
-    }
-    return call;
+  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;
+        }
+      }
+    }
+    return node;
   }
 
   BindingBlock VisitBindingBlock_(const DataflowBlockNode* block_node) final {

tests/python/relax/test_dataflow_pattern.py

@@ -1365,5 +1365,40 @@ def rewriter(_expr, matches):
     tvm.ir.assert_structural_equal(after, expected)
 
 
+def test_rewrite_without_trivial_binding():
+    """rewrite_call should avoid producing trivial "y = x" bindings"""
+
+    @R.function(private=True)
+    def before(x: R.Tensor((1024,))):
+        with R.dataflow():
+            a = R.add(x, x)
+            b = R.reshape(a, (1024,))
+            R.output(b)
+        return b
+
+    @R.function(private=True)
+    def expected(x: R.Tensor((1024,))):
+        with R.dataflow():
+            a = R.add(x, x)
+            R.output(a)
+        return a
+
+    pattern_arg = wildcard()
+    pattern_shape_expr = wildcard()
+    pattern = is_op("relax.reshape")(pattern_arg, pattern_shape_expr)
+
+    def rewriter(expr, matches):
+        arg = matches[pattern_arg]
+        shape_expr = matches[pattern_shape_expr]
+
+        if tvm.ir.structural_equal(arg.struct_info.shape, shape_expr):
+            return arg
+        else:
+            return expr
+
+    after = rewrite_call(pattern, rewriter, before)
+    tvm.ir.assert_structural_equal(after, expected)
+
+
 if __name__ == "__main__":
     tvm.testing.main()