Repository: pytorch
Issue: [dynamo] Detect functionalized functions with HOPs with mutations at tracing time
Labels: triaged,actionable,oncall: pt2,module: functorch,module: dynamo,module: pt2-dispatcher

Discovered with https://github.com/pytorch/pytorch/pull/126341.

Original repro command: `PYTORCH_TEST_WITH_DYNAMO=1 pytest test/functorch/test_control_flow.py::TestControlFlowTraced::test_cond_functionalized_input_mutation_on_true_branch`

These tests also fail with dynamo likely for a similar reason:
```
FAILED [3.5028s] test/functorch/test_control_flow.py::TestControlFlowTraced::test_cond_functionalized_input_aliasing_with_aot_func - torch._dynamo.exc.TorchDynamoException: Unexpected exception when running generated GraphModule
FAILED [1.0649s] test/functorch/test_control_flow.py::TestControlFlowTraced::test_cond_functionalized_input_mutation_on_false_branch - torch._dynamo.exc.TorchDynamoException: Unexpected exception when running generated GraphModule
FAILED [0.7992s] test/functorch/test_control_flow.py::TestControlFlowTraced::test_cond_functionalized_input_mutation_on_true_branch - torch._dynamo.exc.TorchDynamoException: Unexpected exception when running generated GraphModule
FAILED [7.0281s] test/functorch/test_control_flow.py::TestControlFlowTraced::test_cond_functionalized_nested_input_mutation - torch._dynamo.exc.TorchDynamoException: Unexpected exception when running generated GraphModule
FAILED [0.7040s] test/functorch/test_control_flow.py::TestControlFlowTraced::test_cond_functionalized_nested_input_mutation_with_aot_func - torch._dynamo.exc.TorchDynamoException: Unexpected exception when running generated GraphModule
FAILED [0.4711s] test/functorch/test_control_flow.py::TestControlFlowTraced::test_cond_functionalized_output_alias_input - torch._dynamo.exc.TorchDynamoException: Unexpected exception when running generated GraphModule
FAILED [3.3982s] test/functorch/test_control_flow.py::TestControlFlowTraced::test_map_functionalized_arg_mutation - torch._dynamo.exc.TorchDynamoException: Unexpected exception when running generated GraphModule
FAILED [0.8395s] test/functorch/test_control_flow.py::TestControlFlowTraced::test_map_functionalized_elem_alias - torch._dynamo.exc.TorchDynamoException: Unexpected exception when running generated GraphModule
FAILED [0.8148s] test/functorch/test_control_flow.py::TestControlFlowTraced::test_map_functionalized_elem_mutation - torch._dynamo.exc.TorchDynamoException: Unexpected exception when running generated GraphModule
```

Repro:
```python
import torch
from functorch.experimental.control_flow import cond

def true_fn(x):
    view_x = x.view(x.shape)
    view_x.add_(1)
    return view_x.sin().sum()

def false_fn(x):
    return x.cos().sum()

def f(x):
    pred = x.shape[0] == 4
    try:
        y = cond(pred, true_fn, false_fn, [x])
    except:
        y = y + 1
    return y

functional_f = torch.func.functionalize(f)
opt_f = torch.compile(functional_f, backend="eager")
opt_f(torch.randn(4, 5))
```

Output:
```
$  TORCH_LOGS="graph_code" python playground3.py 
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code] TRACED GRAPH
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]  ===== __compiled_fn_3 =====
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]  /data/users/williamwen/pytorch2/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]     def forward(self, L_x_ : torch.Tensor):
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]         l_x_ = L_x_
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]         
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]         # File: /data/users/williamwen/pytorch2/torch/_higher_order_ops/cond.py:111 in cond, code: return cond_op(pred, true_fn, false_fn, operands)
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]         cond_true_0 = self.cond_true_0
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]         cond_false_0 = self.cond_false_0
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]         cond = torch.ops.higher_order.cond(True, cond_true_0, cond_false_0, [l_x_]);  cond_true_0 = cond_false_0 = l_x_ = None
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]         y = cond[0];  cond = None
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]         return (y,)
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]         
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]     class GraphModule(torch.nn.Module):
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]         def forward(self, l_x_):
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]             l_x__1 = l_x_
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]             
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]             # File: /data/users/williamwen/pytorch2/playground3.py:5 in true_fn, code: view_x = x.view(x.shape)
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]             view = l_x__1.view((4, 5));  l_x__1 = None
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]             
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]             # File: /data/users/williamwen/pytorch2/playground3.py:6 in true_fn, code: view_x.add_(1)
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]             add_ = view.add_(1)
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]             
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]             # File: /data/users/williamwen/pytorch2/playground3.py:7 in true_fn, code: return view_x.sin().sum()
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]             sin = view.sin();  view = None
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]             sum_1 = sin.sum();  sin = None
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]             return (sum_1,)
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]             
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]     class GraphModule(torch.nn.Module):
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]         def forward(self, l_x_):
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]             l_x__1 = l_x_
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]             
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]             # File: /data/users/williamwen/pytorch2/playground3.py:10 in false_fn, code: return x.cos().sum()
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]             cos = l_x__1.cos();  l_x__1 = None
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]             sum_1 = cos.sum();  cos = None
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]             return (sum_1,)
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code]             
V0523 13:15:10.230000 139986969671488 torch/_dynamo/output_graph.py:1277] [1/0] [__graph_code] 
Traceback (most recent call last):
  File "/data/users/williamwen/pytorch2/playground3.py", line 22, in <module>
    opt_f(torch.randn(4, 5))
  File "/data/users/williamwen/pytorch2/torch/_dynamo/eval_frame.py", line 414, in _fn
    return fn(*args, **kwargs)
  File "/data/users/williamwen/pytorch2/torch/_dynamo/external_utils.py", line 37, in inner
    return fn(*args, **kwargs)
  File "/data/users/williamwen/pytorch2/torch/_functorch/vmap.py", line 48, in fn
    return f(*args, **kwargs)
  File "/data/users/williamwen/pytorch2/torch/_functorch/eager_transforms.py", line 1691, in wrapped
    func_outputs = func(*func_args, **func_kwargs)
  File "/data/users/williamwen/pytorch2/playground3.py", line 12, in f
    def f(x):
  File "/data/users/williamwen/pytorch2/torch/_dynamo/eval_frame.py", line 548, in _fn
    return fn(*args, **kwargs)
  File "<eval_with_key>.3", line 8, in forward
  File "/data/users/williamwen/pytorch2/torch/_ops.py", line 379, in __call__
    return wrapper()
  File "/data/users/williamwen/pytorch2/torch/_dynamo/eval_frame.py", line 548, in _fn
    return fn(*args, **kwargs)
  File "/data/users/williamwen/pytorch2/torch/_ops.py", line 375, in wrapper
    return self.dispatch(
  File "/data/users/williamwen/pytorch2/torch/_ops.py", line 299, in dispatch
    return dispatch_functorch(self, args, kwargs)
  File "/data/users/williamwen/pytorch2/torch/_functorch/pyfunctorch.py", line 292, in dispatch_functorch
    return interpreter.process(op, args, kwargs)
  File "/data/users/williamwen/pytorch2/torch/_functorch/pyfunctorch.py", line 231, in process
    return kernel(self, *args, **kwargs)
  File "/data/users/williamwen/pytorch2/torch/_ops.py", line 155, in functionalize_functorch_fn
    return fn(_FunctorchFunctionalizeAPI(interpreter), *args, **kwargs)
  File "/data/users/williamwen/pytorch2/torch/_higher_order_ops/cond.py", line 300, in cond_func
    raise UnsupportedAliasMutationException(
torch._higher_order_ops.utils.UnsupportedAliasMutationException: One of torch.cond branch might be modifying the input!
```

cc @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang @zou3519 @Chillee @samdow @kshitij12345 @janeyx99 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng