[Compile] NameError: name 'buf0' is not defined (raised in ddp-training)
Labels: triaged, oncall: pt2

### 🐛 Describe the bug

The demo code:

```python
from mmengine.dist import all_gather, broadcast, get_rank, init_dist

import torch


def batch_shuffle_ddp(x: torch.Tensor):
    """Batch shuffle, for making use of BatchNorm.
    Args:
        x (torch.Tensor): Data in each GPU.
    Returns:
        Tuple[torch.Tensor, torch.Tensor]: Output of shuffle operation.
            - x_gather[idx_this]: Shuffled data.
            - idx_unshuffle: Index for restoring.
    """
    # gather from all gpus
    batch_size_this = x.shape[0]
    x_gather = torch.cat(all_gather(x), dim=0)
    batch_size_all = x_gather.shape[0]

    num_gpus = batch_size_all // batch_size_this

    # random shuffle index
    idx_shuffle = torch.randperm(batch_size_all)

    # broadcast to all gpus
    broadcast(idx_shuffle, src=0)

    # index for restoring
    idx_unshuffle = torch.argsort(idx_shuffle)

    # shuffled index for this gpu
    gpu_idx = get_rank()
    idx_this = idx_shuffle.view(num_gpus, -1)[gpu_idx]

    return x_gather[idx_this], idx_unshuffle


if __name__ == '__main__':
    init_dist(launcher='pytorch')
    func = torch.compile(batch_shuffle_ddp)
    func(torch.ones(1, 1, 1, 1))
    
```

Training with `torchrun`
```python
torchrun --master_port 29501 --nproc_per_node 2  demo.py
```

### Versions

```bash
Collecting environment information...
PyTorch version: 2.1.0a0+git92eb9d3
Is debug build: True
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: CentOS Linux 7 (Core) (x86_64)
GCC version: (GCC) 8.5.0
Clang version: Could not collect
CMake version: version 3.25.0
Libc version: glibc-2.17

Python version: 3.10.9 (main, Jan 11 2023, 15:21:40) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-3.10.0-957.el7.x86_64-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.6.124
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-80GB
GPU 1: NVIDIA A100-SXM4-80GB
GPU 2: NVIDIA A100-SXM4-80GB
GPU 3: NVIDIA A100-SXM4-80GB
GPU 4: NVIDIA A100-SXM4-80GB
GPU 5: NVIDIA A100-SXM4-80GB
GPU 6: NVIDIA A100-SXM4-80GB
GPU 7: NVIDIA A100-SXM4-80GB

Nvidia driver version: 515.65.01
cuDNN version: Could not collect
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
Byte Order:            Little Endian
CPU(s):                256
On-line CPU(s) list:   0-255
Thread(s) per core:    2
Core(s) per socket:    64
Socket(s):             2
NUMA node(s):          2
Vendor ID:             AuthenticAMD
CPU family:            23
Model:                 49
Model name:            AMD EPYC 7742 64-Core Processor
Stepping:              0
CPU MHz:               2250.000
CPU max MHz:           2250.0000
CPU min MHz:           1500.0000
BogoMIPS:              4491.59
Virtualization:        AMD-V
L1d cache:             32K
L1i cache:             32K
L2 cache:              512K
L3 cache:              16384K
NUMA node0 CPU(s):     0-63,128-191
NUMA node1 CPU(s):     64-127,192-255
Flags:                 fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc art rep_good nopl xtopology nonstop_tsc extd_apicid aperfmperf eagerfpu pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_l2 cpb cat_l3 cdp_l3 hw_pstate sme retpoline_amd ssbd ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif umip overflow_recov succor smca

Versions of relevant libraries:
[pip3] numpy==1.23.5
[pip3] torch==2.1.0a0+git92eb9d3
[pip3] torchvision==0.14.1
[pip3] triton==2.0.0
[conda] blas                      1.0                         mkl    defaults
[conda] magma-cuda117             2.6.1                         1    pytorch
[conda] mkl                       2021.4.0           h06a4308_640    defaults
[conda] mkl-include               2022.1.0           h06a4308_224    defaults
[conda] mkl-service               2.4.0           py310h7f8727e_0    defaults
[conda] mkl_fft                   1.3.1           py310hd6ae3a3_0    defaults
[conda] mkl_random                1.2.2           py310h00e6091_0    defaults
[conda] numpy                     1.23.5          py310hd5efca6_0    defaults
[conda] numpy-base                1.23.5          py310h8e6c178_0    defaults
[conda] torch                     2.1.0a0+git92eb9d3           dev_0    <develop>
[conda] torchvision               0.14.1                   pypi_0    pypi
[conda] triton                    2.0.0                    pypi_0    pypi
```

cc @ezyang @soumith @msaroufim @wconstab @ngimel @bdhirsh