Repository: pytorch
Issue: [PT2.0] [.Compile] [Dynamic] Pytorch FX/JIT graph's inputs/nodes ordering is changed when FX recompiles even though the graph operations are same
Labels: low priority,triaged,oncall: pt2,module: dynamic shapes,module: dynamo

### 🐛 Describe the bug

Consider a graph with 5 input tensors and its input shapes as below for iteration 1 and iternation 2:
Iteration 1: input shapes::{0: [], 1: [3], 2: [3], 3: [8, 3, 24, 24, 24], 4: [3], 5: [3]}
Iteration 2: input shapes::{0: [3], 1: [], 2: [3], 3: [8, 3, 24, 24, 71], 4: [3], 5: [3]}

With the above input shapes, Fx will perform recompilations even when dynamic=True. FX graph input for tensor 3 will be like:
iteration 1: [s0, 3, s2, s2, s2]
iteration 2: [s0, 3, s2, s2, s3]

Here the FX graph generated in iteration 1 and iterations 2 have a differences that, the graph input order and graph nodes order (randomly) got changed in recompilation. This makes different JIT graphs and not able to leverage the dynamic shape feature fully.

For example, the below images shows the FX graph of iteration 1 and iteration 2:
![image](https://github.com/pytorch/pytorch/assets/1856674/5b592959-7f94-4ef3-b9db-a3c63225695a)


The graph inputs are same but its order is changed, similarly the graph nodes and operations are same but its order is changed randomly. 

The JIT graph generated from the the FX graphs also reflects with the same differences.

Below testcase can reproduce the issue:
[test_batchnorm3d.zip](https://github.com/pytorch/pytorch/files/12546771/test_batchnorm3d.zip)

```
import torch
import torch.nn as nn
aten = torch.ops.aten
from torch._functorch.aot_autograd import aot_module_simplified
import numpy as np
import random

def my_aot_compiler(gm, example_inputs):
    def my_compiler(gm, example_inputs):
        # print(gm.code)
        # print(gm.graph)
        import copy
        from torch._functorch.compile_utils import strip_overloads
        from torch._functorch.compilers import _disable_jit_autocast

        module = copy.deepcopy(gm)

        with _disable_jit_autocast():
            strip_overloads(module)

            for node in module.graph.nodes:
                new_kwargs = {}
                for k, v in node.kwargs.items():
                    if isinstance(v, torch.device):
                        v = v.type
                    new_kwargs[k] = v
                node.kwargs = new_kwargs

            module.graph.lint()
            module.recompile()

            from collections import OrderedDict
            module._forward_hooks = OrderedDict()
            module._forward_pre_hooks = OrderedDict()
            print("Module is \n", module.print_readable())
            f = torch.jit.script(module)
            print("Converted Graph is \n", f.graph)
        return gm.forward


    # Invoke AOTAutograd
    return aot_module_simplified(
        gm,
        example_inputs,
        fw_compiler=my_compiler
    )

class OpWrapperModule(torch.nn.Module):
    def __init__(self, op):
        super().__init__()
        self.op = op

    def forward(self, inputs):
        result = self.op(inputs)
        return result

def test_dynamic_shape_topk_cpu():
    print("Starting the test.................")
    sizes = [
        (8, 3, 24, 24, 24),
        (8, 3, 24, 24, 71),
        (8, 3, 24, 24, 80),
    ]

    op = nn.BatchNorm3d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
    op = op.to('cpu')
    model_cpu = OpWrapperModule(op).to(dtype=torch.float32).to('cpu')

    compiled_function_training = torch.compile(model_cpu, backend=my_aot_compiler, dynamic=True)

    for s in sizes:
        t = torch.empty(size=s, dtype=torch.float32).uniform_(0, 1).to('cpu').requires_grad_()
        t = t.detach().requires_grad_()
        result_compile_train = compiled_function_training(t)
        grad_in = torch.empty(size=s, dtype=torch.float32).uniform_(0, 1).to('cpu')
        result_compile_train.backward(grad_in)

test_dynamic_shape_topk_cpu()

```

This issue even observed with dynamic=False case,
Please help to analyze this issue and suggest any graph normalization method is there in Pytorch?

### Versions

Collecting environment information...
PyTorch version: 2.0.1a0+git783db92
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: 14.0.5 (ssh://gerrit:29418/tpc_llvm10 02ad77b7f83f1afda4228414a6e3c917bb653668)
CMake version: version 3.27.2
Libc version: glibc-2.31

Python version: 3.8.16 (default, Jan 17 2023, 23:13:24)  [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.4.0-150-generic-x86_64-with-glibc2.17
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

cc @ezyang @msaroufim @bdhirsh @anijain2305 @zou3519 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng @wconstab @Xia-Weiwen