[inductor] Fix pattern replacements with multiple users (pytorch#129689)

Fixes pytorch#129685

After matching a pattern, we currently try to remove all the nodes of that
pattern, which doesn't work if any intermediate node has users outside of the
pattern. In which case we can't delete those particular nodes.

Pull Request resolved: pytorch#129689
Approved by: https://github.com/shunting314

test/inductor/test_torchinductor.py

@@ -10607,6 +10607,20 @@ def fn(primals_5):
         actual = compiled_fn(torch.ones(s0, s1))
         self.assertTrue((actual == 1).all())
 
+    def test_pattern_matcher_multi_user(self):
+        # Reproducer for https://github.com/pytorch/pytorch/issues/129685
+
+        def forward(float_1, view_1):
+            logits = float_1 / 64.0
+            loss = torch.nn.functional.cross_entropy(logits, view_1, ignore_index=5)
+            logsumexp = logits.logsumexp(dim=-1)
+            return [loss, logsumexp]
+
+        a = torch.randn(512, 4096, requires_grad=True)
+        b = torch.randint(size=(512,), low=0, high=4095)
+
+        self.common(forward, (a, b))
+
 
 @dataclasses.dataclass
 class TestFailure:

torch/_inductor/pattern_matcher.py

@@ -202,7 +202,7 @@ def __repr__(self) -> str:
 
     def erase_nodes(self, graph: torch.fx.Graph) -> None:
         for n in reversed(self.nodes):
-            if not n._erased:
+            if not n._erased and not n.users:
                 graph.erase_node(n)
 
     def output_nodes(self) -> List[Optional[torch.fx.Node]]: