Repository: pytorch
Issue: Group convolution slower than manually running separate convolutions in CUDA streams
Labels: module: performance,module: cudnn,module: cuda,module: convolution,triaged

### 🐛 Describe the bug

Calling `nn.Conv2d(..., group=n)` is slower than manually executing `n` separate convolutions.



```python
import torch as ch
from torch import nn
from torch.profiler import profile, record_function, ProfilerActivity

num_convs = 200
in_channels = 512
out_channels = 512
device = 'cuda:0'
x = ch.randn(1, in_channels, 32, 32).to(device)

convs = [nn.Conv2d(in_channels, out_channels, kernel_size=3,
         bias=False).to(device) for _ in range(num_convs)]

fused_conv = nn.Conv2d(in_channels=convs[0].in_channels * num_convs,
                       out_channels=convs[0].out_channels * num_convs,
                       groups=convs[0].groups * num_convs,
                       kernel_size=convs[0].kernel_size,
                       bias=convs[0].bias).to(device)

def run_in_streams(convs, x):
    out = []
    for ind, stream in enumerate([ch.cuda.Stream(device) for _ in range(len(convs))]):
        with ch.cuda.stream(stream):
            out.append(convs[ind](x))
    ch.cuda.synchronize()
    return ch.stack(out)
```

Running both the separate convolutions and the fused convolution through the torch profiler as follows
```python
for _ in range(2):  # warm-up CUDA
    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
        with record_function("model_inference"):
            run_in_streams(convs, x)
print('Manually put each convolution in a CUDA stream:')
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

for _ in range(2):  # warm-up CUDA
    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
        with record_function("model_inference"):
            out_fused = fused_conv(ch.cat([x] * num_convs, dim=1))
print('Use a single convolution with groups')
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
```

shows that the single grouped convolution is slower:
```
Manually put each convolution in a CUDA stream:
...
Self CPU time total: 42.293ms
Self CUDA time total: 94.768ms

Use a single convolution with groups
...
Self CPU time total: 42.354ms
Self CUDA time total: 117.627ms
```
Additionally, in both cases there are `num_convs` calls to the same CUDA kernel (`sm80_xmma_fprop_implicit_gemm`).

As an additional sanity check, proof that the methods above compute the same function:
```python
from torch.nn.utils import parameters_to_vector as ptv
from torch.nn.utils import vector_to_parameters as vtp

params = ch.cat([ptv(conv.parameters()) for conv in convs])
vtp(params, fused_conv.parameters())

out_stream = []
for ind, stream in enumerate(streams):
    with ch.cuda.stream(stream):
        out_stream.append(convs[ind](x))
ch.cuda.synchronize()
out_stream = ch.stack(out_stream)

out_fused = fused_conv(ch.cat([x] * num_convs, dim=1))
sh = list(out_fused.shape)
out_fused = out_fused.view([sh[0], num_convs,
                            sh[1] // num_convs] + sh[2:]).permute([1, 0, 2, 3, 4])
err = ((out_fused - out_stream).abs() / (out_fused.abs() + 1e-8)).mean().cpu().item()
print(f'mean relative difference: {err}')
```
```
mean relative difference: 0.0
```

### Versions

PyTorch version: 1.10.2
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: (GCC) 9.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.27

Python version: 3.9.10 | packaged by conda-forge | (main, Feb  1 2022, 21:24:11)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-4.15.0-169-generic-x86_64-with-glibc2.27
Is CUDA available: True
CUDA runtime version: 11.3.58
GPU models and configuration:
GPU 0: NVIDIA A100-PCIE-40GB
GPU 1: NVIDIA A100-PCIE-40GB
GPU 2: NVIDIA A100-PCIE-40GB
GPU 3: NVIDIA A100-PCIE-40GB
GPU 4: NVIDIA A100-PCIE-40GB
GPU 5: NVIDIA A100-PCIE-40GB
GPU 6: NVIDIA A100-PCIE-40GB
GPU 7: NVIDIA A100-PCIE-40GB
GPU 8: NVIDIA A100-PCIE-40GB

Nvidia driver version: 465.19.01
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.5
[pip3] pytorch-pfn-extras==0.5.6
[pip3] torch==1.10.2
[pip3] torchmetrics==0.7.2
[pip3] torchvision==0.11.3
[conda] blas                      2.113                       mkl    conda-forge
[conda] blas-devel                3.9.0            13_linux64_mkl    conda-forge
[conda] cudatoolkit               11.3.1              ha36c431_10    conda-forge
[conda] libblas                   3.9.0            13_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            13_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            13_linux64_mkl    conda-forge
[conda] liblapacke                3.9.0            13_linux64_mkl    conda-forge
[conda] mkl                       2022.0.1           h8d4b97c_803    conda-forge
[conda] mkl-devel                 2022.0.1           ha770c72_804    conda-forge
[conda] mkl-include               2022.0.1           h8d4b97c_803    conda-forge
[conda] mypy_extensions           0.4.3            py39hf3d152e_4    conda-forge
[conda] numpy                     1.21.5           py39haac66dc_0    conda-forge
[conda] pytorch                   1.10.2          py3.9_cuda11.3_cudnn8.2.0_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] pytorch-pfn-extras        0.5.6                    pypi_0    pypi
[conda] torchmetrics              0.7.2                    pypi_0    pypi
[conda] torchvision               0.11.3               py39_cu113    pytorch

cc @VitalyFedyunin @ngimel @csarofeen @ptrblck @xwang233