WIP: rework

pytorch · bohnstingl · Feb 7, 2025 · Feb 7, 2025 · Feb 7, 2025 · Feb 7, 2025
commit 9e8651f4087d98d9a361ed85f7a22dfcb700c91e
diff --git a/test/functorch/test_control_flow.py b/test/functorch/test_control_flow.py
@@ -4710,7 +4710,7 @@ def forward(self, x_1):
         # torch.cond triggers the check of the branches because the predicate
         # is a SymBool.
         with self.assertRaisesRegex(
-            UnsupportedAliasMutationException, "One of torch.cond branch"
+            UnsupportedAliasMutationException, "A branch or combine_fn might be modifying the input!.*"
         ):
             make_fx(torch.func.functionalize(f), tracing_mode="symbolic")(
                 *example_inputs
@@ -4751,7 +4751,7 @@ def forward(self, x_1):
         # torch.cond triggers the check of the branches because the predicate
         # is a SymBool.
         with self.assertRaisesRegex(
-            UnsupportedAliasMutationException, "One of torch.cond branch"
+            UnsupportedAliasMutationException, "A branch or combine_fn might be modifying the input!.*"
         ):
             make_fx(torch.func.functionalize(f), tracing_mode="symbolic")(
                 *example_inputs
@@ -4785,7 +4785,7 @@ def forward(self, x_1):
         # torch.cond triggers the check of the branches because the predicate
         # is a SymBool.
         with self.assertRaisesRegex(
-            UnsupportedAliasMutationException, "One of torch.cond branch"
+            UnsupportedAliasMutationException, "Aliasing within branch or combine_fn might be occuring!.*"
         ):
             make_fx(torch.func.functionalize(f), tracing_mode="symbolic")(
                 *example_inputs
@@ -4813,7 +4813,7 @@ def f(x):
 
         example_inputs = (torch.ones(4, 5),)
         with self.assertRaisesRegex(
-            UnsupportedAliasMutationException, "One of torch.cond branch"
+            UnsupportedAliasMutationException, "A branch or combine_fn might be modifying the input!.*"
         ):
             make_fx(torch.func.functionalize(f), tracing_mode="symbolic")(
                 *example_inputs
@@ -4846,7 +4846,7 @@ def f(x):
             f(example_input_func)
 
             with self.assertRaisesRegex(
-                UnsupportedAliasMutationException, "One of torch.cond branch"
+                UnsupportedAliasMutationException, "A branch or combine_fn might be modifying the input!.*"
             ):
                 make_fx(f, tracing_mode="symbolic")(example_input_func)
         finally:
@@ -4864,7 +4864,7 @@ def wrapper(*args, **kwargs):
             return wrapper
 
         with self.assertRaisesRegex(
-            UnsupportedAliasMutationException, "One of torch.cond branch"
+            UnsupportedAliasMutationException, "A branch or combine_fn might be modifying the input!.*"
         ):
             make_fx(f_wrapper(f), tracing_mode="symbolic")(example_input_func)
 
@@ -4888,7 +4888,7 @@ def f(x):
             torch._enable_functionalization(reapply_views=False)
             with self.assertRaisesRegex(
                 UnsupportedAliasMutationException,
-                "One of torch.cond branch might be aliasing",
+                "Aliasing within branch or combine_fn might be occuring!.*",
             ):
                 f(example_input_func)
         finally:
@@ -4919,7 +4919,7 @@ def wrapper(*args, **kwargs):
 
         with self.assertRaisesRegex(
             UnsupportedAliasMutationException,
-            "One of torch.cond branch might be aliasing",
+            "Aliasing within branch or combine_fn might be occuring!.*",
         ):
             make_fx(f_wrapper(f), tracing_mode="symbolic")(example_input)
 
@@ -5513,7 +5513,7 @@ def f(xs, y):
         example_inputs = (torch.ones(3, 2, 4), torch.ones(4))
         functional_f = torch.func.functionalize(f)
         with self.assertRaisesRegex(
-            UnsupportedAliasMutationException, "torch.map is mutating the input!"
+            UnsupportedAliasMutationException, "A branch or combine_fn might be modifying the input!.*"
         ):
             functional_f(*example_inputs)
 
@@ -5530,7 +5530,7 @@ def f(xs, y):
         example_inputs = (torch.ones(3, 2, 4), torch.ones(4))
         functional_f = torch.func.functionalize(f)
         with self.assertRaisesRegex(
-            UnsupportedAliasMutationException, "torch.map is mutating the input!"
+            UnsupportedAliasMutationException, "A branch or combine_fn might be modifying the input!.*"
         ):
             functional_f(*example_inputs)
 
@@ -5570,7 +5570,7 @@ def f(xs):
         example_inputs = (torch.ones(3, 2, 4),)
         functional_f = torch.func.functionalize(f)
         with self.assertRaisesRegex(
-            UnsupportedAliasMutationException, "torch.map is aliasing the input!"
+            UnsupportedAliasMutationException, "Aliasing within branch or combine_fn might be occuring!.*"
         ):
             functional_f(*example_inputs)
 
@@ -6462,7 +6462,7 @@ def f(init, xs):
         functional_f = torch.func.functionalize(f)
         with self.assertRaisesRegex(
             UnsupportedAliasMutationException,
-            "Combine_fn might be modifying the input!",
+            "A branch or combine_fn might be modifying the input!.*",
         ):
             functional_f(example_init, example_inputs)
 
@@ -6476,7 +6476,7 @@ def f(init, xs):
         functional_f = torch.func.functionalize(f)
         with self.assertRaisesRegex(
             UnsupportedAliasMutationException,
-            "Combine_fn might be modifying the input!",
+            "A branch or combine_fn might be modifying the input!.*",
         ):
             functional_f(example_init, example_inputs)
 
@@ -6493,7 +6493,7 @@ def f(init, xs):
         example_init = torch.ones(5, 4)
         functional_f = torch.func.functionalize(f)
         with self.assertRaisesRegex(
-            UnsupportedAliasMutationException, "Combine_fn might be aliasing the input!"
+            UnsupportedAliasMutationException, "Aliasing within branch or combine_fn might be occuring!.*"
         ):
             functional_f(example_init, example_inputs)
 
@@ -7015,7 +7015,7 @@ def fn(f, *args):
         x = torch.randn(2, 2)
         for f in ALIAS_FN:
             with self.assertRaisesRegex(
-                torch._dynamo.exc.BackendCompilerFailed, "might be aliasing the input"
+                torch._dynamo.exc.BackendCompilerFailed, ".*Aliasing within branch or combine_fn might be occuring!.*"
             ):
                 torch.compile(fn)(f, x)
 
@@ -7031,7 +7031,7 @@ def f(arg1, arg2):
         # as a result of auto lifting.
         for view_f in ALIAS_FN[1:]:
             with self.assertRaisesRegex(
-                torch._dynamo.exc.BackendCompilerFailed, "might be aliasing the input"
+                torch._dynamo.exc.BackendCompilerFailed, ".*Aliasing within branch or combine_fn might be occuring!.*"
             ):
                 torch.compile(fn)(view_f, x)
 

diff --git a/torch/_dynamo/graph_deduplication.py b/torch/_dynamo/graph_deduplication.py
@@ -120,6 +120,7 @@ def _replace_region_with_subgraph(
     invoke_args = (get_subgraph_node, subgraph_name, tuple(sub_args))
     fake_inputs = [node.meta["example_value"] for node in sub_args]
 
+    # TODO: We don't care here about any output-output aliasing?
     if has_potential_input_alias_or_mutation(sub_gm, fake_inputs):
         log.debug(
             "NYI: Failed to substitute region %s due to input alias or mutation",

@@ -2874,6 +2874,7 @@ def call_function(
             for node in body_gmod.graph.nodes
             if node.op == "placeholder"
         ]
+        # TODO: We don't care here about any output-output aliasing?
         if has_potential_input_alias_or_mutation(body_gmod, fake_inputs):
             raise RuntimeError(
                 f"{self.value._name} where the inputs are mutated or the "

@@ -9,6 +9,7 @@
 import torch.utils._pytree as pytree
 from torch._C import DispatchKey
 from torch._higher_order_ops.utils import (
+    check_input_mutation_and_alias,
     _maybe_run_with_interpreter,
     _set_compilation_env,
     autograd_not_implemented,
@@ -415,6 +416,16 @@ def associative_scan_functionalize(ctx, combine_fn, xs):
         functional_combine_fn = ctx.functionalize(
             _maybe_run_with_interpreter(combine_fn)
         )
+        pre_dispatch = hasattr(ctx, "mode") and ctx.mode.pre_dispatch
+        sample_inputs = list(
+            itertools.chain(
+                [inp.clone() for inp in unwrapped_xs],
+                [inp.clone() for inp in unwrapped_xs],
+            )
+        )
+
+        check_input_mutation_and_alias(combine_fn, sample_inputs, pre_dispatch=pre_dispatch)
+
         ret = associative_scan_op(functional_combine_fn, unwrapped_xs)
     return ctx.wrap_tensors(ret)
 

@@ -19,8 +19,7 @@
 from torch._functorch.utils import exposed_in
 from torch._guards import detect_fake_mode
 from torch._higher_order_ops.utils import (
-    _has_potential_branch_input_alias,
-    _has_potential_branch_input_mutation,
+    check_input_mutation_and_alias,
     _maybe_run_with_interpreter,
     _set_compilation_env,
     reenter_make_fx,
@@ -483,22 +482,7 @@ def cond_func(ctx, pred, true_fn, false_fn, inputs):
         functional_false = ctx.functionalize(_maybe_run_with_interpreter(false_fn))
         pre_dispatch = hasattr(ctx, "mode") and ctx.mode.pre_dispatch
         for branch in [true_fn, false_fn]:
-            if _has_potential_branch_input_mutation(
-                branch, unwrapped_inputs, pre_dispatch=pre_dispatch
-            ):
-                raise UnsupportedAliasMutationException(
-                    "One of torch.cond branch might be modifying the input! "
-                    "Consider cloning the input before modifying it. "
-                )
-        for branch in [true_fn, false_fn]:
-            if _has_potential_branch_input_alias(
-                branch, unwrapped_inputs, pre_dispatch=pre_dispatch
-            ):
-                raise UnsupportedAliasMutationException(
-                    "One of torch.cond branch might be aliasing the input! "
-                    "If you are returning a view of the input, please make sure "
-                    "to clone it. "
-                )
+            check_input_mutation_and_alias(branch, unwrapped_inputs, pre_dispatch=pre_dispatch)
 
         cond_return = cond_op(
             unwrapped_pred, functional_true, functional_false, unwrapped_inputs

@@ -7,7 +7,7 @@
 from torch import Tensor
 from torch._C import DispatchKey
 from torch._higher_order_ops.utils import (
-    _has_potential_branch_input_mutation,
+    check_input_mutation_and_alias,
     _maybe_reenter_make_fx,
     autograd_not_implemented,
     reenter_make_fx,
@@ -420,13 +420,15 @@ def flex_attention_functionalize(
         functional_score_mod = ctx.functionalize(score_mod)
         pre_dispatch = hasattr(ctx, "mode") and ctx.mode.pre_dispatch
         with TransformGetItemToIndex():
-            mutates = _has_potential_branch_input_mutation(
-                score_mod, example_vals, pre_dispatch
-            )
-        # The only care about mutations of existing buffers since we can't replay these.
-        # However, we can just error if anything 
EED3
is detected
-        if mutates:
-            raise UnsupportedAliasMutationException("Mutations detected in score_mod")
+            check_input_mutation_and_alias(score_mod, example_vals, pre_dispatch)
+        #     mutates = _has_potential_branch_input_mutation(
+        #         score_mod, example_vals, pre_dispatch
+        #     )
+        # # The only care about mutations of existing buffers since we can't replay these.
+        # # However, we can just error if anything is detected
+        # if mutates:
+        #     raise UnsupportedAliasMutationException("Mutations detected in score_mod")
+
 
         out = flex_attention(
             query_unwrapped,

@@ -3,8 +3,7 @@
 import torch.utils._pytree as pytree
 from torch._C import DispatchKey
 from torch._higher_order_ops.utils import (
-    _has_potential_branch_input_alias,
-    _has_potential_branch_input_mutation,
+    check_input_mutation_and_alias,
     autograd_not_implemented,
     reenter_make_fx,
     unique_graph_id,
@@ -97,18 +96,19 @@ def hints_wrapper_functionalize(ctx, body_fn, args, kwargs, hints):
     with ctx.redispatch_to_next():
         functional_body_fn = ctx.functionalize(body_fn)
         pre_dispatch = hasattr(ctx, "mode") and ctx.mode.pre_dispatch
-        if _has_potential_branch_input_mutation(
-            body_fn, unwrapped_args, pre_dispatch=pre_dispatch
-        ):
-            raise UnsupportedAliasMutationException(
-                "body_fn of hints_wrapper might be modifying the input!"
-            )
-        if _has_potential_branch_input_alias(
-            body_fn, unwrapped_args, pre_dispatch=pre_dispatch
-        ):
-            raise UnsupportedAliasMutationException(
-                "body_fn of hints_wrapper might be aliasing the input!"
-            )
+        check_input_mutation_and_alias(body_fn, unwrapped_args, pre_dispatch=pre_dispatch)
+        # if _has_potential_branch_input_mutation(
+        #     body_fn, unwrapped_args, pre_dispatch=pre_dispatch
+        # ):
+        #     raise UnsupportedAliasMutationException(
+        #         "body_fn of hints_wrapper might be modifying the input!"
+        #     )
+        # if _has_potential_branch_input_alias(
+        #     body_fn, unwrapped_args, pre_dispatch=pre_dispatch
+        # ):
+        #     raise UnsupportedAliasMutationException(
+        #         "body_fn of hints_wrapper might be aliasing the input!"
+        #     )
         outputs = hints_wrapper(
             functional_body_fn,
             unwrapped_args,

@@ -5,8 +5,7 @@
 from torch._dispatch.python import suspend_functionalization
 from torch._functorch.aot_autograd import AOTConfig, create_joint
 from torch._higher_order_ops.utils import (
-    _has_potential_branch_input_alias,
-    _has_potential_branch_input_mutation,
+    check_input_mutation_and_alias,
     _maybe_run_with_interpreter,
     reenter_make_fx,
     UnsupportedAliasMutationException,
@@ -250,15 +249,16 @@ def map_functionalize(ctx, f, xs, pos_args):
         with disable_proxy_modes_tracing():
             example_inputs = (*_unstack_pytree(unwrapped_xs)[0], *unwrapped_args)
         pre_dispatch = hasattr(ctx, "mode") and ctx.mode.pre_dispatch
-        if _has_potential_branch_input_mutation(
-            f, example_inputs, pre_dispatch=pre_dispatch
-        ):
-            raise UnsupportedAliasMutationException("torch.map is mutating the input!")
-
-        if _has_potential_branch_input_alias(
-            f, example_inputs, pre_dispatch=pre_dispatch
-        ):
-            raise UnsupportedAliasMutationException("torch.map is aliasing the input!")
+        check_input_mutation_and_alias(f, example_inputs, pre_dispatch=pre_dispatch)
+        # if _has_potential_branch_input_mutation(
+        #     f, example_inputs, pre_dispatch=pre_dispatch
+        # ):
+        #     raise UnsupportedAliasMutationException("torch.map is mutating the input!")
+
+        # if _has_potential_branch_input_alias(
+        #     f, example_inputs, pre_dispatch=pre_dispatch
+        # ):
+        #     raise UnsupportedAliasMutationException("torch.map is aliasing the input!")
 
         map_return = map_impl(wrapped_fn, unwrapped_xs, unwrapped_args)
         return ctx.wrap_tensors(map_return)
@@ -9,8 +9,7 @@
 import torch.utils._pytree as pytree
 from torch._C import DispatchKey
 from torch._higher_order_ops.utils import (
-    _has_potential_branch_input_alias,
-    _has_potential_branch_input_mutation,
+    check_input_mutation_and_alias,
     _set_compilation_env,
     autograd_not_implemented,
     reenter_make_fx,
@@ -454,18 +453,9 @@ def scan_functionalize(ctx, combine_fn, init, xs, reverse, additional_inputs):
                 unwrapped_additional_inputs,
             )
         )
-        if _has_potential_branch_input_mutation(
-            combine_fn, sample_inputs, pre_dispatch=pre_dispatch
-        ):
-            raise UnsupportedAliasMutationException(
-                "Combine_fn might be modifying the input!"
-            )
-        if _has_potential_branch_input_alias(
-            combine_fn, sample_inputs, pre_dispatch=pre_dispatch
-        ):
-            raise UnsupportedAliasMutationException(
-                "Combine_fn might be aliasing the input!"
-            )
+
+        check_input_mutation_and_alias(combine_fn, sample_inputs, pre_dispatch=pre_dispatch)
+
         ret = scan_op(
             functional_combine_fn,
             unwrapped_init,