subgraph rewriter supports matched pattern with no users (#4)

cherry-pick from upstream: pytorch#143842

Author: internal developer

test/fx/test_subgraph_rewriter.py

@@ -678,6 +678,51 @@ def comparison(x1, x2, x3):
         test_outs = traced.forward(x1, x2, x3)
         self.assertEqual(ref_outs, test_outs)
 
+    def test_subgraph_rewriter_with_unused_results(self):
+        class M(torch.nn.Module):
+            def forward(self, x, y, cache):
+                m = torch.mul(x, y)
+                n = cache.index_copy(0, torch.tensor([0]), m)
+                p = torch.ops.aten.copy.default(cache, n)
+                q = torch.ops.aten.copy_.default(cache, p)  # noqa: F841
+                u = torch.relu(cache)
+                # check the result to ensure cache is updated before relu op
+                return u
+
+        def pattern(self_tensor, src_tensor):
+            p = torch.ops.aten.copy.default(self_tensor, src_tensor)
+            q = torch.ops.aten.copy_.default(self_tensor, p)
+            return q
+
+        def replacement(self_tensor, src_tensor):
+            q = torch.ops.aten.copy_.default(self_tensor, src_tensor)
+            return q
+
+        def comparison(x, y, cache):
+            m = torch.mul(x, y)
+            n = cache.index_copy(0, torch.tensor([0]), m)
+            q = torch.ops.aten.copy_.default(cache, n)  # noqa: F841
+            u = torch.relu(cache)
+            return u
+
+        traced = symbolic_trace(M())
+        comparison_fn = symbolic_trace(comparison)
+
+        subgraph_rewriter.replace_pattern(traced, pattern, replacement)
+
+        traced.graph.lint()
+
+        x = torch.randn(1, 8)
+        y = torch.randn(1, 8)
+        cache = torch.randn(2, 8)
+        x_clone = x.clone()
+        y_clone = y.clone()
+        cache_clone = cache.clone()
+
+        ref_outs = comparison_fn(x, y, cache)
+        test_outs = traced.forward(x_clone, y_clone, cache_clone)
+        self.assertEqual(ref_outs, test_outs)
+
     def test_subgraph_rewriter_call_method(self):
         class M(torch.nn.Module):
             def forward(self, x):

torch/fx/subgraph_rewriter.py

@@ -364,9 +364,11 @@ def _replace_pattern(
         user_nodes: Set[Node] = set()
         for n in match.returning_nodes:
             user_nodes.update(n.users)
-        assert user_nodes, "The returning_nodes should have at least one user node"
 
-        if len(user_nodes) == 1:
+        first_user_node = None
+        if len(user_nodes) == 0:
+            first_user_node = None
+        elif len(user_nodes) == 1:
             first_user_node = next(iter(user_nodes))
         else:
             # If there are multiple user nodes, we need to find the first user node
@@ -376,7 +378,22 @@ def _replace_pattern(
                     first_user_node = n
                     break
 
-        with original_graph.inserting_before(first_user_node):  # type: ignore[possibly-undefined]
+        first_next_node = None
+        if first_user_node is None:
+            # no users, so we insert the replacement graph before the first next
+            # node of returning nodes
+            next_node = None
+            for n in reversed(original_graph.nodes):
+                if n in match.returning_nodes:
+                    first_next_node = next_node
+                    break
+                else:
+                    next_node = n
+        insert_point = (
+            first_user_node if first_user_node is not None else first_next_node
+        )
+        assert insert_point is not None, "The insert point can't be None"
+        with original_graph.inserting_before(insert_point):
             copied_returning_nodes = original_graph.graph_copy(
                 replacement_graph, val_map
             )