Repository: pytorch
Issue: Model parameter and gradient memory formats are inconsistent with compiled autograd
Labels: triaged,oncall: pt2,module: aotdispatch,module: compiled autograd,module: pt2-dispatcher

### 🐛 Describe the bug

When compiled autograd is enabled and model contains operations preferring inputs in channels last format, gradient and parameters memory layouts are inconsistent:
- parameter: contiguous (default)
- gradient: channels last

Issue doesn't appear when layout optimization is disabled (```TORCHINDUCTOR_LAYOUT_OPTIMIZATION=0```).

Reproducer:
``` python

import torch
from torch._dynamo.utils import maybe_enable_compiled_autograd


class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = torch.nn.Conv2d(
            3,
            10,
            kernel_size=7,
            stride=[2, 2],
            padding=[3, 3],
            dilation=[1, 1],
            bias=False,
            groups=1,
        )

    def forward(self, x):
        x = self.conv1(x)
        return x


for compiled_autograd_enabled in [False, True]:
    input = torch.rand(16, 3, 224, 224)

    model_compiled = torch.compile(Model())
    with maybe_enable_compiled_autograd(compiled_autograd_enabled):
        out = model_compiled(input)
        loss = out.sum()
        loss.backward()

    param = next(model_compiled.parameters())
    grad = param.grad
    print("> compiled_autograd_enabled", compiled_autograd_enabled)
    print("parameter contiguous={}, channels_last={}".format(param.is_contiguous(memory_format=torch.contiguous_format),
          param.is_contiguous(memory_format=torch.channels_last)))
    print("grad contiguous={}, channels_last={}".format(grad.is_contiguous(memory_format=torch.contiguous_format),
          grad.is_contiguous(memory_format=torch.channels_last)))
```

Output:
```
> compiled_autograd_enabled False
parameter contiguous=True, channels_last=False
grad contiguous=True, channels_last=False

> compiled_autograd_enabled True
parameter contiguous=True, channels_last=False
grad contiguous=False, channels_last=True
```

Output with ```TORCHINDUCTOR_LAYOUT_OPTIMIZATION=0```:
```
> compiled_autograd_enabled False
parameter contiguous=True, channels_last=False
grad contiguous=True, channels_last=False

> compiled_autograd_enabled True
parameter contiguous=True, channels_last=False
grad contiguous=True, channels_last=False
```


### Versions

Collecting environment information...
PyTorch version: 2.3.0+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: 14.0.0-1ubuntu1.1
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-107-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

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      43 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             12
On-line CPU(s) list:                0-11
Vendor ID:                          GenuineIntel
Model name:                         Intel(R) Xeon(R) Gold 5220R CPU @ 2.20GHz
CPU family:                         6
Model:                              85
Thread(s) per core:                 1
Core(s) per socket:                 6
Socket(s):                          2
Stepping:                           0
BogoMIPS:                           4389.68
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon nopl xtopology tsc_reliable nonstop_tsc cpuid tsc_known_freq pni pclmulqdq vmx ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single pti ssbd ibrs ibpb stibp tpr_shadow vnmi ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 invpcid avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xsaves arat pku ospke md_clear flush_l1d arch_capabilities                                                                                                                                                                Virtualization:                     VT-x                                                                                                                                                                                                                                                Hypervisor vendor:                  VMware                                                                                                                                                                                                                                              Virtualization type:                full                                                                                                                                                                                                                                                L1d cache:                          384 KiB (12 instances)                                                                                                                                                                                                                              L1i cache:                          384 KiB (12 instances)                                                                                                                                                                                                                              L2 cache:                           12 MiB (12 instances)                                                                                                                                                                                                                               L3 cache:                           71.5 MiB (2 instances)                                                                                                                                                                                                                              NUMA node(s):                       1                                                                                                                                                                                                                                                   NUMA node0 CPU(s):                  0-11                                                                                                                                                                                                                                                Vulnerability Gather data sampling: Unknown: Dependent on hypervisor status                                                                                                                                                                                                             Vulnerability Itlb multihit:        KVM: Mitigation: VMX disabled                                                                                                                                                                                                                       Vulnerability L1tf:                 Mitigation; PTE Inversion; VMX flush not necessary, SMT disabled                                                                                                                                                                                    Vulnerability Mds:                  Mitigation; Clear CPU buffers; SMT Host state unknown                                                                                                                                                                                               Vulnerability Meltdown:             Mitigation; PTI                                                                                                                                                                                                                                     Vulnerability Mmio stale data:      Mitigation; Clear CPU buffers; SMT Host state unknown                                                                                                                                                                                               Vulnerability Retbleed:             Mitigation; IBRS                                                                                                                                                                                                                                    Vulnerability Spec rstack overflow: Not affected                                                                                                                                                                                                                                        Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp                                                                                                                                                                                 Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization                                                                                                                                                                                Vulnerability Spectre v2:           Mitigation; IBRS; IBPB conditional; STIBP disabled; RSB filling; PBRSB-eIBRS Not affected; BHI Syscall hardening, KVM SW loop                                                                                                                       Vulnerability Srbds:                Not affected                                                                                                                                                                                                                                        Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:                                                                                                                                                                                                                                                         [pip3] torch==2.3.0                                                                                                                                                                                                                                                                     [pip3] triton==2.3.0                                                                                                                                                                                                                                                                    [conda] Could not collect 

cc @ezyang @chauhang @penguinwu @xmfan @yf225 @zou3519 @bdhirsh @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames @rec @msaroufim @anijain2305 @peterbell10 @ipiszy @muchulee8 @ColinPeppler @desertfire