[BugFix] Support rewrite_once when the number of callbacks > 1 (#14344)

* [BugFix] Support rewrite_once when the number of callbacks > 1

* callbacks_map -> done, swapping false and true

---------

Co-authored-by: Bin Li <[email protected]>

Author: sisleyli

src/relay/ir/dataflow_matcher.cc

@@ -796,24 +796,35 @@ Expr PatternRewriter::Rewrite(const Array<DFPatternCallback>& callbacks, const E
   bool equal = true;
   static auto* structural_equal = runtime::Registry::Get("node.StructuralEqual");
   ICHECK(structural_equal) << "node.StructuralEqual is not registered.";
+  // Keep track of callbacks that have finished rewriting
+  std::unordered_map<DFPatternCallback, bool, ObjectPtrHash, ObjectPtrEqual> done;
   do {
     last = post;
     for (auto callback : callbacks) {
-      callback_ = callback;
-      if (callback_->require_type) {
-        post = InferTypeWithModule(post, mod_);
-      }
-      auto grouper = PatternGrouper();
-      groups_ = grouper.GroupMatches(callback_->pattern, post);
-      gid_assignments_ = grouper.GetGIDAssignments();
-      memo_.clear();
-      VLOG(1) << "pre rewritten:" << std::endl << PrettyPrint(pre);
-      post = this->VisitExpr(post);
-      VLOG(1) << "post rewritten:" << std::endl << PrettyPrint(post);
-      count++;
+      if (!done[callback]) {
+        auto before = post;
+        callback_ = callback;
+        if (callback_->require_type) {
+          post = InferTypeWithModule(post, mod_);
+        }
+        auto grouper = PatternGrouper();
+        groups_ = grouper.GroupMatches(callback_->pattern, post);
+        gid_assignments_ = grouper.GetGIDAssignments();
+        memo_.clear();
+        VLOG(1) << "pre rewritten:" << std::endl << PrettyPrint(pre);
+        post = this->VisitExpr(post);
+        VLOG(1) << "post rewritten:" << std::endl << PrettyPrint(post);
+        count++;
+        if (callback_->rewrite_once) {
+          bool current_equal = (*structural_equal)(before, post, false, true);
+          if (!current_equal) {
+            done[callback] = true;
+          }
+        }
+      }
     }
     equal = (*structural_equal)(last, post, false, true);
-  } while (!equal && count < 100 && !callback_->rewrite_once);
+  } while (!equal && count < 100);
   if (count >= 100) {
     LOG(FATAL) << "Observed 100 rewrite passes, possible conflicting passes?";
   }

tests/python/relay/test_dataflow_pattern.py

@@ -1804,22 +1804,83 @@ def callback(self, pre, post, node_map):
             if new_args:
                 return relay.op.concatenate(relay.expr.Tuple(new_args), axis=0)
             else:
-                return concat_args
+                return concat_args[0]
 
     x = relay.var("x")
     y = relay.var("y")
     z = relay.var("z")
     concat = relay.op.concatenate(relay.expr.Tuple([x, y, z]), axis=0)
 
-    # Let the rewriter run recursively
-    out = rewrite(ConcatRewriter(False), concat)
-    expected = relay.expr.Tuple([x])
-    assert tvm.ir.structural_equal(out, expected)
+    def test_one_callback():
+        # Let the rewriter run recursively
+        out = rewrite(ConcatRewriter(False), concat)
+        expected = x
+        assert tvm.ir.structural_equal(out, expected)
+
+        # Run the rewriter once
+        out = rewrite(ConcatRewriter(True), concat)
+        expected = relay.op.concatenate(relay.expr.Tuple([x, y]), axis=0)
+        assert tvm.ir.structural_equal(out, expected)
+
+    def test_multi_callbacks():
+        # This class recursively add a nn.relu operator after nn.softmax
+        class OneMoreReluRewriter(DFPatternCallback):
+            def __init__(self, rewrite_once):
+                super().__init__(rewrite_once=rewrite_once)
+                self.pattern = is_op("nn.softmax")(None)
+
+            def callback(self, pre, post, node_map):
+                return relay.nn.relu(post)
+
+        def before():
+            # Before:
+            #    x    y    z
+            #    |    |    |
+            #       concat
+            #         |
+            #      softmax
+            return relay.nn.softmax(concat)
+
+        def once_concat():
+            # ConcatRewrite once, OneMoreReluRewrite once
+            # Expected:
+            #   x    y
+            #   |    |
+            #   concat
+            #      |
+            #   softmax
+            #      |
+            #    relu
+            return relay.nn.relu(
+                relay.nn.softmax(relay.op.concatenate(relay.expr.Tuple([x, y]), axis=0))
+            )
+
+        def recursive_concat():
+            # ConcatRewrite recursively, OneMoreReluRewrite once
+            # Expected:
+            #      x
+            #      |
+            #   softmax
+            #      |
+            #    relu
+            return relay.nn.relu(relay.nn.softmax(x))
+
+        # Run ConcatRewriter once, OneMoreReluRewriter once
+        out = rewrite(
+            [OneMoreReluRewriter(True), ConcatRewriter(True)],
+            before(),
+        )
+        assert tvm.ir.structural_equal(out, once_concat())
+
+        # Run ConcatRewriter recursively, OneMoreReluRewriter once
+        out = rewrite(
+            [OneMoreReluRewriter(True), ConcatRewriter(False)],
+            before(),
+        )
+        assert tvm.ir.structural_equal(out, recursive_concat())
 
-    # Run the rewriter once
-    out = rewrite(ConcatRewriter(True), concat)
-    expected = relay.op.concatenate(relay.expr.Tuple([x, y]), axis=0)
-    assert tvm.ir.structural_equal(out, expected)
+    test_one_callback()
+    test_multi_callbacks()
 
 
 def test_matched_outside_but_dominated():