Repository: pytorch
Issue: AOTAutograd breaks diffusers VAE decoder
Labels: triaged,oncall: pt2,module: pt2-dispatcher

### 🐛 Describe the bug

AOTAutograd backend to torch.compile creates an non-functional graph resulting is a tensor violating an assertion of a particular op that it is contiguous

### Error logs

Traceback (most recent call last):
  File "/data/bvhworkspace/pt2/contiguous_bug.py", line 14, in <module>
    y = f_aot(x)  # not contiguous :(
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 489, in _fn
    return fn(*args, **kwargs)
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/diffusers/models/autoencoders/vae.py", line 285, in forward
    def forward(
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 489, in _fn
    return fn(*args, **kwargs)
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/_dynamo/external_utils.py", line 17, in inner
    return fn(*args, **kwargs)
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 901, in forward
    return compiled_fn(full_args)
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/utils.py", line 81, in g
    return f(*args)
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 83, in runtime_wrapper
    all_outs = call_func_at_runtime_with_args(
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/utils.py", line 105, in call_func_at_runtime_with_args
    out = normalize_as_list(f(args))
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/utils.py", line 81, in g
    return f(*args)
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/autograd/function.py", line 553, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 408, in forward
    fw_outs = call_func_at_runtime_with_args(
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/utils.py", line 105, in call_func_at_runtime_with_args
    out = normalize_as_list(f(args))
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/_dynamo/backends/debugging.py", line 59, in run
    return torch.fx.Interpreter(fx_g).boxed_run(args)
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/fx/interpreter.py", line 170, in boxed_run
    return self.run(initial_env=env)
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/fx/interpreter.py", line 138, in run
    self.env[node] = self.run_node(node)
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/fx/interpreter.py", line 195, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/fx/interpreter.py", line 267, in call_function
    return target(*args, **kwargs)
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/_ops.py", line 513, in __call__
    return self._op(*args, **(kwargs or {}))
RuntimeError: Expected X.is_contiguous(memory_format) to be true, but got false.  (Could this error message be improved?  If so, please report an enhancement request to PyTorch.)

While executing %native_group_norm_3 : [num_users=3] = call_function[target=torch.ops.aten.native_group_norm.default](args = (%div_1, %primals_21, %primals_22, 1, 64, 4096, 32, 1e-06), kwargs = {})
Original traceback:
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/diffusers/models/autoencoders/vae.py", line 333, in forward
    sample = self.mid_block(sample, latent_embeds)
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/diffusers/models/unets/unet_2d_blocks.py", line 739, in forward
    hidden_states = resnet(hidden_states, temb)
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/bvhworkspace/.venv/lib/python3.10/site-packages/diffusers/models/resnet.py", line 332, in forward
    hidden_states = self.norm1(hidden_states)


### Minified repro

### [minifier did not produce anything, this is my own code]
```
import diffusers
import torch


if __name__ == "__main__":
    m = diffusers.models.AutoencoderKL().decoder

    f_vanilla = torch.compile(m, backend="eager")
    f_aot = torch.compile(m, backend="aot_eager")

    x = torch.randn(1, 4, 64, 64)

    y = f_vanilla(x)  # no problem
    y = f_aot(x)  # not contiguous :(
```

### Versions

Collecting environment information...
PyTorch version: 2.2.2+cpu
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.4 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: 9.0.0 (https://github.com/llvm-mirror/llvm c62b24f070c9a4bb1a76409e623042a740cac4cd)
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-91-generic-x86_64-with-glibc2.35
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

Versions of relevant libraries:
[pip3] numpy==1.23.1
[pip3] torch==2.2.2+cpu
[pip3] torchvision==0.17.2
[pip3] diffusers==0.27.2
[conda] Could not collect

cc @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang @zou3519