Prevent pattern matches across mutation ops in Inductor pre-grad FX passes

[williamwen42]

The inductor pre-grad FX passes (https://github.com/pytorch/pytorch/blob/main/torch/_inductor/fx_passes/pre_grad.py) are unsafe with respect to mutation. We need to discard matches that cross a mutation op.

Example repro:

    def test_match_with_mutation(self):
        from torch._inductor.pattern_matcher import (
            CallFunction,
            KeywordArg,
            PatternMatcherPass,
            register_graph_pattern,
        )

        counter = 0
        test_pass = PatternMatcherPass()

        @register_graph_pattern(
            CallFunction(
                torch.add, KeywordArg("x"), CallFunction(torch.sin, KeywordArg("x"))
            ),
            pass_dict=test_pass,
        )
        def _test(match, x):
            nonlocal counter
            counter += 1

        def fn(x, y):
            a = torch.sin(x)
            x.copy_(y)
            b = torch.add(x, a)
            return b

        args = [
            torch.randn(5, 5, device="cuda"),
            torch.randn(5, 5, device="cuda"),
        ]

        with unittest.mock.patch(
            "torch._inductor.fx_passes.pre_grad.pattern_matcher_passes", [test_pass]
        ):
            expected = fn(*args)
            actual = torch.compile(fn)(*args)
            # should not match
            self.assertEqual(counter, 0)
            torch.testing.assert_close(actual, expected)       

reveals that pre-grad FX passes currently match across mutation ops.

To run, add it to a test file (e.g. to test/inductor/test_pattern_matcher.py) and run as any other unittest.

cc @ezyang @soumith @msaroufim @wconstab @ngimel @bdhirsh @anijain2305

[ezyang]

oh this is literally the thing we were talking about today @bdhirsh