Repository: pytorch
Issue: CPU execution/dispatch time dominates and slows down small TorchScript GPU models 
Labels: oncall: jit

### 🐛 Describe the bug

In small models that run on the GPU, the CPU time to dispatch/launch the kernels dominates - and is significantly slower than I would expect, even when using TorchScript (`script` or `trace`). Eg it takes ~.07ms for the CPU to execute a single `aten:addmm` operation (just dispatching to the GPU), while the GPU kernels themselves take only like .02ms. Here's a screenshot of a profiling trace of one `forward` pass from the below example: 

![image](https://user-images.githubusercontent.com/5598968/153672594-10cbdeac-cd16-4e58-a7ae-68d0d8d2ead4.png)


The whole forward pass takes .2ms (and about .1ms when run without the profiler), and is completely CPU-bound - you can increase `BATCH_SIZE` from 100 to 10000 with barely a change in execution time. I guess I just don't understand what takes that long - .2ms is like 600,000 CPU clock cycles! Is the CPU actually doing real work here that ought to take so long? My understanding is that TorchScript stuff is happening in C++ so there's not python slowness to blame it on.

This might sound plenty fast, but I'm trying to reproduce a popular RL model with a recurrent component, and the original [tensorflow implementation](https://github.com/danijar/dreamer) is able to run the inner loop of their algorithm (forward and backward pass, significantly more complicated than this example) in ~1ms. This is necessary to get 10 steps/sec each with a recurrence length (a loop where each computation depends on the last) of 100 timesteps. So a single `addmm` operation taking .07 ms is not really workable. `addmm` is just an example, though, all the ops take significantly more CPU time than GPU time.

Let me know if there's anything I can do to optimize this further!

Here's the python example: 

```python
from tqdm import tqdm
import torch
from torch import nn
from torch.nn import functional as F

device = torch.device("cuda")
BATCH_SIZE = 100
SIZE = 100

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(SIZE, SIZE)
        self.linear2 = nn.Linear(SIZE, SIZE)

    def forward(self, x):
        x = self.linear1(x)
        x = F.relu(x)
        return self.linear2(x)

mlp = MLP().to(device)
data = torch.randn(BATCH_SIZE, SIZE).to(device)
mlp = torch.jit.trace_module(mlp, {"forward": data})


for s in tqdm(range(50000)):
    mlp(data)
```


### Versions

```
% python collect_env.py
Collecting environment information...
PyTorch version: 1.10.0+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.10.2
Libc version: glibc-2.27

Python version: 3.9.2 (default, Jun 15 2021, 23:11:16)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-1064-gcp-x86_64-with-glibc2.27
Is CUDA available: True
CUDA runtime version: 11.3.109
GPU models and configuration:
GPU 0: NVIDIA A100-SXM4-40GB
GPU 1: NVIDIA A100-SXM4-40GB

Nvidia driver version: 465.19.01
cuDNN version: Probably one of the following:
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn.so.8
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_adv_train.so.8
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_ops_train.so.8
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy==0.902
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.19.5
[pip3] numpy-quaternion==2020.11.2.17.0.49
[pip3] nvidia-dlprof-pytorch-nvtx==1.8.0
[pip3] pytorch-lightning==1.3.7.post0
[pip3] torch==1.10.0+cu113
[pip3] torch-tb-profiler==0.3.1
[pip3] torchaudio==0.10.0+cu113
[pip3] torchmetrics==0.3.2
[pip3] torchtext==0.10.0
[pip3] torchvision==0.11.1+cu113
[conda] Could not collect
```