Repository: pytorch
Issue: torch.jit.trace has incorrect execution for += operation during compilation
Labels: oncall: jit

### 🐛 Describe the bug

When there is ```+=``` in model code, torch.compile and torch.jit.trace may have incorrect execution and result in some runtime errors. The following is the sample code.

```python
import torch

class Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.emb = torch.nn.Embedding(16, 16)

    def forward(self, input):
        idx = input
        # print(idx) # print-1
        idx += input
        # print(idx) # print-2
        ret = self.emb(idx)
        return ret
    
model = Model().cuda()
data = torch.tensor([6], dtype=torch.long).cuda()

# correct result
# out = model(data)
# print(out)
# out = model(data)
# print(out)

# correct execution during compilation
# model_opt = torch.compile(model)
# out = model_opt(data)
# print(out)

# incorrect execution during compilation
# out = model(data)
# model_opt = torch.compile(model)
# out = model_opt(data)
# print(out)

# incorrect execution during compilation
# model_opt = torch.jit.trace(model, data)
# out = model_opt(data)
# print(out)

# incorrect execution during jit compilation and onnx export
# torch.onnx.export(model, data, "test.onnx")
```

As shown in the above code, when execute the code in PyTorch eager mode, the execution is correct:
```
# print-1
tensor([6], device='cuda:0')
# print-2
tensor([12], device='cuda:0')
# print(out)
tensor([[ 1.5988, -0.4076,  1.0858,  0.3278,  2.6424, -0.6067, -0.6254,  0.3936,
         -1.2936,  0.1778,  0.1374, -0.3317, -0.6924, -1.6892, -0.3555,  1.4910]],
       device='cuda:0', grad_fn=<EmbeddingBackward0>)
```

When directly executing torch.compile, the execution is also correct:
```
# print-1
tensor([6], device='cuda:0')
# print-2
tensor([12], device='cuda:0')
# print(out)
tensor([[-0.2688,  0.5553,  1.6430, -0.2665,  0.4982,  1.1613, -1.2119,  1.4576,
          1.1647,  2.7468, -1.2716, -1.2882,  1.4957,  1.0340, -1.5124,  0.4420]],
       device='cuda:0', grad_fn=<CompiledFunctionBackward>)
```

However, when executing eager mode and then executing torch.compile, the ```idx``` value will be incorrect during compilation and result in cuda kernel failure.
```
# eagar mode
tensor([6], device='cuda:0')
tensor([12], device='cuda:0')
tensor([[ 1.1823, -0.1317,  1.6275,  0.6011,  1.6439, -1.4418,  0.6417,  0.7029,
          0.3030, -1.5715,  1.1036,  1.4204,  0.0075, -0.6817, -2.5617, -0.4169]],
       device='cuda:0', grad_fn=<EmbeddingBackward0>)
# incorrect results during compilation of torch.compile
tensor([12], device='cuda:0')
tensor([24], device='cuda:0')
<frozen importlib._bootstrap_external>:850: _call_with_frames_removed: block: [0,0,0], thread: [0,0,0] Assertion `index out of bounds: 0 <= tmp4 < 16` failed.
```

torch.jit.trace and torch.onnx.export (onnx.export leverages torch.jit.trace) will have incorrect ```idx``` value and lead to runtime error.
```
tensor([6], device='cuda:0')
tensor([12], device='cuda:0')
tensor([12], device='cuda:0')
tensor([24], device='cuda:0')
../aten/src/ATen/native/cuda/Indexing.cu:1237: indexSelectSmallIndex: block: [0,0,0], thread: [0,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
```


However, when using ```idx = idx + input``` instead of ```idx += input```, everything works well:
```python
class Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.emb = torch.nn.Embedding(16, 16)

    def forward(self, input):
        idx = input
        # print(idx) # print-1
        idx = idx + input
        # print(idx) # print-2
        ret = self.emb(idx)
        return ret
```

Both PyTorch-2.1.0 release and PyTorch nightly build 2.2.0.dev20231117 has been tested.

### Versions

PyTorch version: 2.2.0.dev20231117+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.25.0
Libc version: glibc-2.27

Python version: 3.9.12 (main, Apr  5 2022, 06:56:58)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-150-generic-x86_64-with-glibc2.27
Is CUDA available: True
CUDA runtime version: 12.2.128
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA RTX A6000
GPU 1: NVIDIA RTX A6000
GPU 2: NVIDIA RTX A6000
GPU 3: NVIDIA RTX A6000

Nvidia driver version: 535.54.03
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):              64
On-line CPU(s) list: 0-63
Thread(s) per core:  2
Core(s) per socket:  32
Socket(s):           1
NUMA node(s):        1
Vendor ID:           AuthenticAMD
CPU family:          25
Model:               1
Model name:          AMD EPYC 7513 32-Core Processor
Stepping:            1
CPU MHz:             1507.004
CPU max MHz:         2600.0000
CPU min MHz:         1500.0000
BogoMIPS:            5200.26
Virtualization:      AMD-V
L1d cache:           32K
L1i cache:           32K
L2 cache:            512K
L3 cache:            32768K
NUMA node0 CPU(s):   0-63
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 rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 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_llc mwaitx cpb cat_l3 cdp_l3 invpcid_single hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr wbnoinvd arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca

Versions of relevant libraries:
[pip3] numpy==1.24.1
[pip3] onnx==1.14.0
[pip3] onnxruntime-gpu==1.15.0
[pip3] onnxsim==0.4.28
[pip3] pytorch-triton==2.1.0+6e4932cda8
[pip3] torch==2.2.0.dev20231117+cu121
[pip3] triton==2.1.0
[conda] numpy                     1.24.1                   pypi_0    pypi
[conda] pytorch-triton            2.1.0+6e4932cda8          pypi_0    pypi
[conda] torch                     2.2.0.dev20231117+cu121          pypi_0    pypi
[conda] triton                    2.1.0                    pypi_0    pypi


cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @zou3519