Repository: pytorch
Issue: compile_fx inplace modifly the input graph
Labels: triaged,oncall: pt2,module: inductor

### 🐛 Describe the bug

I'm compiling a graph multiple times using inductor. I find it inplace modify the graph, and one of the graph's output changes from tensor to list of tensor.

example code:

```python
import torch
from torch import nn

@torch.library.custom_op("silly::attention", mutates_args=["out"])
def silly_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, out: torch.Tensor) -> None:
    print("silly")
    out.copy_(q)
    out[0] += 1

@silly_attention.register_fake
def _(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, out: torch.Tensor) -> None:
    return

class SillyModel(nn.Module):
    def __init__(self) -> None:
        super().__init__()

    def forward(
        self,
        x: torch.Tensor
    ) -> torch.Tensor:
        x = x + 1
        x = x + 2
        out = torch.empty_like(x)
        torch.ops.silly.attention(x, x, x, out)
        x = out
        x = x - 2
        x = x - 1
        out = torch.empty_like(x)
        torch.ops.silly.attention(x, x, x, out)
        x = out
        x = x + 1
        return x

model = SillyModel()

def mybackend(graph, example_inputs):
    graph.print_readable()

    subgraph_id = 0
    node_to_subgraph_id = {}
    attention_graphs = []
    for node in graph.graph.nodes:
        if node.op in ("output", "placeholder"):
            continue
        if node.op == 'call_function' and str(node.target) in ["silly.attention"]:
            subgraph_id += 1
            node_to_subgraph_id[node] = subgraph_id
            attention_graphs.append(subgraph_id)
            subgraph_id += 1
        else:
            node_to_subgraph_id[node] = subgraph_id

    # `keep_original_order` is important!
    # otherwise pytorch might reorder the nodes and
    # the semantics of the graph will change when we
    # have mutations in the graph
    split_gm = torch.fx.passes.split_module.split_module(
        graph,
        None,
        lambda node: node_to_subgraph_id[node],
        keep_original_order=True)

    split_gm.print_readable()

    for (name, module) in list(split_gm.named_modules()):
        if name == "":
            # stitching module
            continue
        if "." in name:
            # recursive child module
            continue
        graph_id = int(name.replace("submod_", ""))
        if graph_id not in attention_graphs:
            # cannot setattr to a module, so we need to set it to the dict
            split_gm.__dict__[name] = piecewise_backend(
                module)

    print("first compilation")
    output = split_gm(*example_inputs)
    print(output.__class__)

    return split_gm.forward

def piecewise_backend(model):

    compiled_once = True
    compiled_twice = True
    compiled_graph_1 = None
    compiled_graph_2 = None
    def f(*args):
        model.print_readable()
        nonlocal compiled_once, compiled_twice, compiled_graph_1, compiled_graph_2
        if compiled_once:
            compiled_once = False
            from torch._inductor.compile_fx import compile_fx
            compiled_graph = compile_fx(model, args)
            return compiled_graph(*args)
        if compiled_twice:
            compiled_once = False
            from torch._inductor.compile_fx import compile_fx
            compiled_graph_1 = compile_fx(model, args)
        return compiled_graph_1(*args)
    return f

model = torch.compile(model, backend=mybackend)

input_buffer = torch.randn(100)
output = model(input_buffer)
print(output.__class__)

output = model(input_buffer)
print(output.__class__)
```

Some key output logging:

this is the last piece of graph:

```python
    class submod_4(torch.nn.Module):
        def forward(self, out_1: "f32[100]"):
             # File: /data/youkaichao/tmp/vllm/test_split.py:32 in forward, code: x = x + 1
            add: "f32[100]" = out_1 + 1;  out_1 = None
            return add
```

after calling `compile_fx`, it becomes:

```python
class GraphModule(torch.nn.Module):
    def forward(self, out_1: "f32[100]"):
         # File: /data/youkaichao/tmp/vllm/test_split.py:32 in forward, code: x = x + 1
        add: "f32[100]" = out_1 + 1;  out_1 = None
        return [add]
```

is it a desired behavior?

if I want to avoid it, how to clone a graph, without cloning the module and parameters, if I want to send a copy to `compile_fx`?

### Versions

pytorch 2.4 , 2.5, and 2.6.0.dev20241022+cu124

cc @ezyang @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @aakhundov