Repository: pytorch
Issue: scripted fft Convolutions are faster than nn.Conv1d with large kernels
Labels: module: performance,module: cuda,module: convolution,triaged

### 🐛 Describe the bug

On running torchaudio code I noticed that some resampling operations are slower than they should be on the forward pass of the Resample transform.  I tracked the slowness to the use of functional.conv1d(). Since an fft convolution is as slow as a mult per group and a cat, its time is approximately constant regarding conv size.  Should pytorch not handle the choice ( serialized conv, or fft conv) internally?  Or should this issue be a torchaudio issue? 

```python
from typing import Optional
import time
import torch
from torch import Tensor
import torch.nn.functional as F

# pylint: disable=no-member
# pylint: disable=suppressed-message
def fftconv1d(x: Tensor, weight: Tensor,
              bias: Optional[Tensor] = None,
              padding: int = 0,
              groups: int = 1) -> Tensor:
    """
    Args
        x:       Tensor   (batch_size, in_channels, size)
        weight:  Tensor   (out_channels, in_channels//groups, kernel_size)
        bias:    Tensor   [None] out_channels
        padding  int      [0]
        groups   int      [1] in_channels, out _channels must be divisible by groups
            # stride and dilation = 1
        adapted from https://towardsdatascience.com/fourier-convolutions-in-pytorch-4cbd23c70005
        faster for large ones
    """
    assert x.ndim == 3, "x expedted shape: (N, C, L)"
    assert weight.ndim == 3, "weight expected (in_channels, out_channels, kernel)"
    _out, _in, _ = weight.shape

    if bias is not None:
        assert bias.ndim==1 and len(bias) == _out, "bias vector sized as out_channels reqd"
    assert not x.shape[1]%groups, f"in_channels must be mod groups {x.shape[1], groups}"
    assert not _out%groups, f"out_channels must be mod groups {_out, groups}"
    assert x.shape[1]  == groups*_in, f"Given groups={groups} and weight {tuple(weight.shape)}, \
        expected input {tuple(x.shape)} to have {groups*_in} channels"

    out = F.pad(x, [padding, padding])
    _pad = out.shape[-1] - weight.shape[-1]

    x_rfft = torch.fft.rfftn(out, dim=-1)
    w_rfft = torch.fft.rfftn(F.pad(weight, (0, _pad)), dim=-1)
    w_rfft.imag *= -1
    if groups == 1:
        x_rfft = torch.einsum("ab..., cb... -> ac...", x_rfft, w_rfft)
    else:
        _o = _out//groups
        x_rfft = torch.cat([torch.einsum("ab..., cb... -> ac...",
                                         x_rfft[:, _in*g:_in*(g+1)],
                                         w_rfft[_o*g:_o*(g+1)])
                             for g in range(groups)], dim=1)

    out = torch.fft.irfftn(x_rfft, dim=-1)[..., :_pad + 1].contiguous()

    if bias is not None:
        out = out + bias.view(1, -1, 1)

    return out

def _testconv(cuda=True, grad=True, pad=None, out_channels=4, in_channels=2,
              batch_size= 20, size = 4096, ksize = 1000, groups=1):
    if pad is None:
        pad = ksize//2

    signal = torch.randn(batch_size, in_channels, size)
    if grad:
        signal.requires_grad = True
    kernel = torch.randn(out_channels, in_channels//groups, ksize)
    bias = torch.randn(out_channels)

    print(f"\n signal: {tuple(signal.shape)}, kernel: {tuple(kernel.shape)}")
    if cuda:
        signal = signal.to(device="cuda")
        kernel = kernel.to(device="cuda")
        bias = bias.to(device="cuda")

    _start = time.time()
    y0 = F.conv1d(signal, kernel, bias=bias, padding=pad, groups=groups)
    if cuda:
        torch.cuda.synchronize()
    _fconv = time.time()

    y2 = fftconv1d(signal, kernel, bias=bias, padding=pad, groups=groups)
    if cuda:
        torch.cuda.synchronize()
    _fftconv = time.time()

    _test = f'test: cuda:{cuda}, grad:{grad}, pad{pad}, out:{out_channels}, in{in_channels}, groups{groups}'
    print(_test)
    _nntime = 1000*(_fconv - _start)
    _fftime = 1000*(_fftconv - _fconv)
    if _nntime < _fftime:
        _nn="\t\t\tnn.Conv1d is faster"
        _ff =""
    elif _fftime < _nntime:
        _nn = ""
        _ff = "\t\t\tFFT faster"

    print(f"  nn.Conv1d() time {1000*(_nntime):.1f} ms {_nn}")
    print(f"  fftconv1d time   {1000*(_fftime):.1f} ms {_ff}")
    assert torch.allclose(y0, y2, rtol=1e-3, atol=1e-3), _test

def test_conv_opt():
    cuda = [True, False]
    grad = [True, False]
    padding = [0, None, 100]
    groups = [1,2]
    out_channels = [4,2]
    in_channels = [2,8]

    batch_size = 20
    size = [4096, 14400]
    ksize = [9, 1000]
    for p in padding:
        for r in grad:
            for c in cuda:
                for g in groups:
                    for i in in_channels:
                        for o in out_channels:
                            for k in ksize:
                                for s in size:
                                    _testconv(cuda=c, grad=r, pad=p, out_channels=o, groups=g,
                                            in_channels=i, batch_size=batch_size, size=s, ksize=k)

```


### Versions

Collecting environment information...
PyTorch version: 1.11.0
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.22.4
Libc version: glibc-2.27

Python version: 3.9.5 (default, May 18 2021, 19:34:48)  [GCC 7.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-113-generic-x86_64-with-glibc2.27
Is CUDA available: True
CUDA runtime version: 11.6.55
GPU models and configuration: GPU 0: NVIDIA TITAN RTX
Nvidia driver version: 510.39.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.0.5
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] denoising-diffusion-pytorch==0.7.1.1
[pip3] numpy==1.20.3
[pip3] pytorch-fid==0.2.0
[pip3] pytorch-lightning==1.4.0
[pip3] pytorch3d==0.2.0
[pip3] torch==1.11.0
[pip3] torch-ema==0.2
[pip3] torch-fidelity==0.3.0
[pip3] torch-tb-profiler==0.4.0
[pip3] torchaudio==0.11.0
[pip3] torchmetrics==0.4.1
[pip3] torchnmf==0.3.5.dev0
[pip3] torchvision==0.12.0
[conda] _pytorch_select           0.1                       cpu_0  
[conda] blas                      2.114                       mkl    conda-forge
[conda] blas-devel                3.9.0            14_linux64_mkl    conda-forge
[conda] cudatoolkit               11.3.1               h2bc3f7f_2  
[conda] denoising-diffusion-pytorch 0.7.1.1                  pypi_0    pypi
[conda] libblas                   3.9.0            14_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            14_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            14_linux64_mkl    conda-forge
[conda] liblapacke                3.9.0            14_linux64_mkl    conda-forge
[conda] mkl                       2022.0.1           h06a4308_117  
[conda] mkl-devel                 2022.0.1           h66538d2_117  
[conda] mkl-include               2022.0.1           h06a4308_117  
[conda] numpy                     1.19.3                   pypi_0    pypi
[conda] pytorch                   1.11.0          py3.9_cuda11.3_cudnn8.2.0_0    pytorch
[conda] pytorch-fid               0.2.0                    pypi_0    pypi
[conda] pytorch-lightning         1.4.0                    pypi_0    pypi
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] pytorch3d                 0.2.0                    pypi_0    pypi
[conda] torch-ema                 0.2                      pypi_0    pypi
[conda] torch-fidelity            0.3.0                    pypi_0    pypi
[conda] torch-tb-profiler         0.4.0                    pypi_0    pypi
[conda] torchaudio                0.11.0               py39_cu113    pytorch
[conda] torchmetrics              0.4.1                    pypi_0    pypi
[conda] torchnmf                  0.3.5.dev0               pypi_0    pypi
[conda] torchvision               0.12.0               py39_cu113    pytorch


cc @VitalyFedyunin @ngimel