Repository: pytorch
Issue:  torch.jit.script(model) and torch.jit.trace(model) performance degradation.
Labels: oncall: jit

### 🐛 Describe the bug

I used https://github.com/MichaelMonashev/bench_models/blob/main/bench_models.py to benchmark new torch version on vision neural networks. Here is results:

**Torch 2.0.1**:
```
CPU: 12th Gen Intel(R) Core(TM) i9-12900K cores: 24 

Motherboard vendor: Gigabyte Technology Co., Ltd.
Motherboard model name: Z690 UD  version: x.x 

GPU: NVIDIA RTX A5000   24564 MiB
GPU driver version: 525.89.02
PCI-E max: 4@16x

OS: Ubuntu 20.04 focal
Linux-5.4.0-144-generic-x86_64-with-glibc2.29 

Python version: 3.8.10
Torch version: 2.0.1+cu118
Torch GPU name: NVIDIA RTX A5000
Torch available memory: 24 Gb
Torch CUDA version: 11.8 

Sending data to GPU benchmark.
Warming up batches 3
Benchmark batches 5 

                                       Pinned                Not pinned
Batch size                       Blocked Not blocked     Blocked Not blocked
torch.Size([64, 3, 32, 32])         4378        3780        3502        3469 MB/s
torch.Size([64, 3, 64, 64])         9364        9437        7973        8228 MB/s
torch.Size([64, 3, 128, 128])      17568       17489       15235       15135 MB/s
torch.Size([64, 3, 256, 256])      17245       21214       13322       13275 MB/s
torch.Size([64, 3, 512, 512])      22743       22787       13678       13673 MB/s
torch.Size([64, 3, 1024, 1024])    23011       23024       13904       13931 MB/s


Getting data from GPU benchmark.
Warming up batches 3
Benchmark batches 5 

Batch size               Blocked   Not blocked
torch.Size([64, 32])         217      109 MB/s
torch.Size([64, 64])         353      242 MB/s
torch.Size([64, 128])        780      478 MB/s
torch.Size([64, 256])       1488      977 MB/s
torch.Size([64, 512])       2482     1729 MB/s
torch.Size([64, 1024])      3914     3771 MB/s
torch.Size([64, 2048])      5535     7679 MB/s
torch.Size([64, 4096])      7027    13738 MB/s
torch.Size([64, 8192])      9248    17599 MB/s
torch.Size([64, 16384])     7256    20670 MB/s
torch.Size([64, 32768])     4375    21368 MB/s
torch.Size([64, 65536])    13547    22286 MB/s


Downloading model weights...
GPU inference benchmark (forward pass speed).
Batch size torch.Size([64, 3, 332, 332])
Warming up batches 3
Benchmark batches 5 

Source   Model name                     Top1                                   |     torch.jit.script(model)      |      torch.jit.trace(model)      | with torch.cuda.graph()
                                              Float32 Float32 Float16 BFloat16 | Float32 Float32 Float16 BFloat16 | Float32 Float32 Float16 BFloat16 | Float32 Float32 Float16
                                                       + TF32   (AMP)          |          + TF32   (AMP)          |          + TF32   (AMP)          |          + TF32   (AMP)
timm     tf_efficientnetv2_s_in21ft1k   84.9%     492     596     968      855 |     519     638     801      772 |     649     655     814      781 |     496     597     960 img/s
timm       channels last                84.9%     417     689    1255     1298 |     448     791    1269     1385 |     821     823    1334     1427 |     420     695    1258 img/s
timm     tf_efficientnetv2_m_in21ft1k   86.2%     273     347     559      495 |     283     367     458      440 |     375     377     474      451 |     278     351     560 img/s
timm       channels last                86.2%     232     410     742      770 |     246     460     730      794 |     475     477     777      831 |     235     413     741 img/s
timm     tf_efficientnetv2_l_in21ft1k   86.8%     143     191     316      281 |     149     202     263      253 |     207     208     271      258 |     146     194     317 img/s
timm       channels last                86.8%     116     227     417      442 |     121     252     411      453 |     262     263     439      476 |     117     229     417 img/s
timm     tf_efficientnetv2_xl_in21ft1k  87.3%      99     134     224      199 |     103     142     186      176 |     145     145     191      180 |     100     136     225 img/s
timm       channels last                87.3%      85     162     297      316 |      88     178     288      314 |     185     185     310      333 |      86     164     297 img/s
timm     gernet_s                       75.7%    2063    2624    4322     4180 |    2485    3416    4987     5277 |    3490    3485    5140     5408 |    2096    2667    4331 img/s
timm       channels last                75.7%    1533    2877    5406     5536 |    1759    3877    6529     7536 |    3982    3986    7026     7829 |    1551    2906    5438 img/s
timm     tinynet_a                      77.7%     997    1086    1904     1453 |    1135    1251    1768     1404 |    1259    1266    1801     1422 |    1011    1099    1921 img/s
timm       channels last                77.7%    1011    1327    2411     2482 |    1148    1594    2662     2713 |    1613    1614    2769     2755 |    1030    1339    2421 img/s
timm     tinynet_b                      75.0%    1298    1370    2372     1852 |    1530    1631    2299     1852 |    1643    1646    2363     1869 |    1312    1383    2385 img/s
timm       channels last                75.0%    1311    1625    2922     3022 |    1545    1982    3313     3339 |    2005    2006    3441     3410 |    1341    1640    2913 img/s
timm     tinynet_c                      71.2%    1950    2030    3535     2762 |    2328    2449    3534     2797 |    2467    2459    3566     2844 |    1976    2053    3562 img/s
timm       channels last                71.2%    1980    2391    4228     4311 |    2366    2934    4798     4852 |    2952    2958    4969     4910 |    2053    2427    4244 img/s
timm     tinynet_d                      67.0%    2832    2943    4879     3987 |    3418    3588    4956     4246 |    3591    3564    5147     4241 |    2865    2968    4865 img/s
timm       channels last                67.0%    2864    3382    5932     5996 |    3367    4181    6863     7063 |    4170    4170    7085     7029 |    2879    3423    5914 img/s
timm     tinynet_e                      59.9%    3008    3127    5133     4286 |    3668    3863    5369     4540 |    3854    3869    5453     4537 |    3039    3151    5149 img/s
timm       channels last                59.9%    2967    3543    6255     6270 |    3580    4416    7252     7352 |    4411    4417    7536     7408 |    3051    3588    6192 img/s
```

**torch 2.2.0**:
```
CPU: 12th Gen Intel(R) Core(TM) i9-12900K cores: 24 

Motherboard vendor: Gigabyte Technology Co., Ltd.
Motherboard model name: Z690 UD  version: x.x 

GPU: NVIDIA RTX A5000   24564 MiB
GPU driver version: 545.29.06
PCI-E max: 4@16x

OS: Ubuntu 20.04 focal
Linux-5.4.0-144-generic-x86_64-with-glibc2.29 

Python version: 3.8.10
Torch version: 2.2.0+cu121
Torch GPU name: NVIDIA RTX A5000
Torch available memory: 24 Gb
Torch CUDA version: 12.1 

Sending data to GPU benchmark.
Warming up batches 3
Benchmark batches 5 

                                       Pinned                Not pinned
Batch size                       Blocked Not blocked     Blocked Not blocked
torch.Size([64, 3, 32, 32])         4386        3834        3560        3396 MB/s
torch.Size([64, 3, 64, 64])         9960        9906        8447        8326 MB/s
torch.Size([64, 3, 128, 128])      17892       16246       15136       15132 MB/s
torch.Size([64, 3, 256, 256])      21276       21376       13924       14034 MB/s
torch.Size([64, 3, 512, 512])      22685       22732       13989       13909 MB/s
torch.Size([64, 3, 1024, 1024])    23093       23024       14244       14250 MB/s


Getting data from GPU benchmark.
Warming up batches 3
Benchmark batches 5 

Batch size               Blocked   Not blocked
torch.Size([64, 32])         211      120 MB/s
torch.Size([64, 64])         409      208 MB/s
torch.Size([64, 128])        799      471 MB/s
torch.Size([64, 256])       1509      925 MB/s
torch.Size([64, 512])       2520     1853 MB/s
torch.Size([64, 1024])      4051     3823 MB/s
torch.Size([64, 2048])      5659     7377 MB/s
torch.Size([64, 4096])      7160    13868 MB/s
torch.Size([64, 8192])      9294    17597 MB/s
torch.Size([64, 16384])     5043    20815 MB/s
torch.Size([64, 32768])     5952    22316 MB/s
torch.Size([64, 65536])    13044    23810 MB/s


Downloading model weights...
GPU inference benchmark (forward pass speed).
Batch size torch.Size([64, 3, 332, 332])
Warming up batches 3
Benchmark batches 5 

Source   Model name                     Top1                                   |     torch.jit.script(model)      |      torch.jit.trace(model)      | with torch.cuda.graph()
                                              Float32 Float32 Float16 BFloat16 | Float32 Float32 Float16 BFloat16 | Float32 Float32 Float16 BFloat16 | Float32 Float32 Float16
                                                       + TF32   (AMP)          |          + TF32   (AMP)          |          + TF32   (AMP)          |          + TF32   (AMP)
timm     tf_efficientnetv2_s_in21ft1k   84.9%     477     609     968      850 |     476     607     880      828 |     615     616     904      837 |     476     608     956 img/s
timm       channels last                84.9%     361     685    1252     1296 |     359     684    1123     1218 |     701     702    1173     1286 |     362     692    1244 img/s
timm     tf_efficientnetv2_m_in21ft1k   86.2%     264     350     557      491 |     263     350     510      477 |     355     356     525      483 |     264     353     554 img/s
timm       channels last                86.2%     200     407     729      759 |     200     404     662      714 |     414     415     691      757 |     203     410     731 img/s
timm     tf_efficientnetv2_l_in21ft1k   86.8%     139     191     311      276 |     139     191     287      271 |     195     195     297      275 |     141     194     312 img/s
timm       channels last                86.8%     110     225     410      436 |     110     223     374      412 |     229     229     394      438 |     112     227     411 img/s
timm     tf_efficientnetv2_xl_in21ft1k  87.3%      95     134     218      194 |      95     134     203      190 |     137     137     210      193 |      96     136     219 img/s
timm       channels last                87.3%      77     160     293      313 |      77     159     268      294 |     164     164     281      313 |      78     162     294 img/s
timm     gernet_s                       75.7%    2041    2695    4297     4049 |    2064    2709    3575     4019 |    2956    2956    3805     4408 |    2088    2743    4289 img/s
timm       channels last                75.7%    1388    2839    5445     5433 |    1395    2842    4380     5391 |    3125    3127    4741     5969 |    1409    2862    5430 img/s
timm     tinynet_a                      77.7%     987    1086    1865     1436 |     988    1083    1769     1399 |    1090    1091    1799     1409 |     998    1100    1884 img/s
timm       channels last                77.7%    1000    1321    2374     2453 |     995    1308    2239     2269 |    1330    1330    2285     2407 |    1015    1335    2400 img/s
timm     tinynet_b                      75.0%    1281    1350    2329     1834 |    1287    1345    2202     1785 |    1353    1357    2250     1798 |    1301    1365    2332 img/s
timm       channels last                75.0%    1278    1614    2892     2978 |    1281    1601    2729     2761 |    1622    1623    2779     2920 |    1304    1629    2907 img/s
timm     tinynet_c                      71.2%    1939    2013    3494     2661 |    1935    2013    3244     2612 |    2020    2025    3340     2629 |    1970    2043    3472 img/s
timm       channels last                71.2%    1952    2379    4182     4230 |    1967    2366    3954     3951 |    2396    2400    3961     4154 |    2009    2417    4193 img/s
timm     tinynet_d                      67.0%    2813    2926    4811     3842 |    2808    2916    4614     3795 |    2925    2925    4583     3819 |    2853    2962    4765 img/s
timm       channels last                67.0%    2800    3367    5867     5878 |    2789    3358    5648     5513 |    3370    3373    5572     5757 |    2832    3413    5892 img/s
timm     tinynet_e                      59.9%    2993    3105    5123     4130 |    2986    3094    4823     4114 |    3100    3101    4849     4087 |    3027    3142    5092 img/s
timm       channels last                59.9%    2928    3535    6220     6190 |    2960    3519    5911     5803 |    3536    3539    5888     6062 |    3002    3579    6153 img/s
```

As you can see torch.jit.script(model) and torch.jit.trace(model) with torch 2.2.0 works **slowly** than Torch 2.0.1 .

### Versions

```
Collecting environment information...
PyTorch version: 2.0.1+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.25.0
Libc version: glibc-2.31

Python version: 3.8.10 (default, Nov 14 2022, 12:59:47)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.4.0-139-generic-x86_64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA RTX A5000
Nvidia driver version: 525.85.05
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
Address sizes:                   46 bits physical, 48 bits virtual
CPU(s):                          24
On-line CPU(s) list:             0-23
Thread(s) per core:              1
Core(s) per socket:              16
Socket(s):                       1
NUMA node(s):                    1
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           151
Model name:                      12th Gen Intel(R) Core(TM) i9-12900K
Stepping:                        2
CPU MHz:                         2568.291
CPU max MHz:                     6700.0000
CPU min MHz:                     800.0000
BogoMIPS:                        6374.40
Virtualization:                  VT-x
L1d cache:                       384 KiB
L1i cache:                       256 KiB
L2 cache:                        10 MiB
NUMA node0 CPU(s):               0-23
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Not affected
Vulnerability Retbleed:          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; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt clwb intel_pt sha_ni xsaveopt xsavec xgetbv1 xsaves dtherm ida arat pln pts hwp hwp_notify hwp_act_window hwp_epp hwp_pkg_req umip pku ospke waitpkg gfni vaes vpclmulqdq tme rdpid movdiri movdir64b md_clear pconfig flush_l1d arch_capabilities

Versions of relevant libraries:
[pip3] efficientnet-pytorch==0.7.1
[pip3] numpy==1.24.2
[pip3] pytorch-lightning==2.0.1
[pip3] segmentation-models-pytorch==0.3.3
[pip3] torch==2.0.1+cu118
[pip3] torchaudio==2.0.2+cu118
[pip3] torchmetrics==0.11.4
[pip3] torchvision==0.15.2+cu118
[pip3] triton==2.0.0
[conda] Could not collect
```

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel