Repository: pytorch
Issue: Slow convolution with large kernels, should be using FFT
Labels: module: performance,module: cudnn,module: convolution,triaged,module: fft

## 🐛 Bug

When using `Conv1d` with a large kernel size (1024 for instance) on gpu, the cudnn implementation is very slow and gets slower as I increase the kernel size. I thought it was using FFT but apparently not. If it were using FFT, the computation time should be independent of the kernel size, because the kernel is anyway padded to the length of the input.

I have tried benchmarking with both `torch.backends.cudnn.benchmark` set to `False` and `True`. My implementation using the FFT is significantly faster especially when using a stride of 1. You will find hereafter the code both for the FFT based convolution implementation I use and the profiling. My implementation is within ~5e-5 of the reference implementation for random weights and input. For a kernel size of 1024, with 64 channels, a stride of 1 and an input of length 64000, the default implementation is about 20 times slower than the FFT based one. When using a kernel size of 2048, it is 40 times slower.


## To Reproduce

Steps to reproduce the behavior:

1. Copy the code below in `profile_conv.py`
2. Run `python3 profile_conv.py`

```python
import time
from functools import partial
from itertools import product

import torch
from torch import nn
from torch.nn import functional as F


def compl_mul(a, b):
    """
    Given a and b two tensors of dimension 4
    with the last dimension being the real and imaginary part, 
    returns a multiplied by the conjugate of b, the multiplication
    being with respect to the second dimension.
    """
    op = partial(torch.einsum, "bct,dct->bdt")
    return torch.stack([
        op(a[..., 0], b[..., 0]) + op(a[..., 1], b[..., 1]),
        op(a[..., 1], b[..., 0]) - op(a[..., 0], b[..., 1])
    ],
                       dim=-1)


class FastConv(nn.Module):
    """
    Convoluton based on FFT, faster for large kernels and small strides.
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size,
                 stride=1,
                 bias=True):
        super().__init__()

        if bias:
            self.bias = nn.Parameter(torch.zeros(out_channels, 1))
        else:
            self.bias = None
        self.weight = nn.Parameter(
            torch.zeros(out_channels, in_channels, kernel_size))

        self.kernel_size = kernel_size
        self.stride = stride

    def forward(self, signal):
        padded = F.pad(self.weight,
                       (0, signal.size(-1) - self.weight.size(-1)))
        signal_fr = torch.rfft(signal, 1)
        weight_fr = torch.rfft(padded, 1)
        output_fr = compl_mul(signal_fr, weight_fr)
        output = torch.irfft(output_fr, 1, signal_sizes=(signal.size(-1), ))
        output = output[..., ::self.stride]
        target_length = (signal.size(-1) - self.kernel_size) // self.stride + 1
        output = output[..., :target_length].contiguous()
        if self.bias is not None:
            output += self.bias
        return output


def profile(module, *args, repetitions=10, warmup=1):
    """
    Given a module and args, apply repeatedly the module to the args,
    calling `torch.cuda.synchronize()` in between. Return the time per
    repetition. Not perfect profiling but gives a rough idea.
    """
    module(*args)
    begin = time.time()
    for _ in range(repetitions):
        module(*args)
        torch.cuda.synchronize()
    return (time.time() - begin) / repetitions


def human_seconds(seconds, display='.2f'):
    """
    Human readable string from a number of seconds.
    """
    value = seconds * 1e6
    ratios = [1e3, 1e3, 60, 60, 24]
    names = ['us', 'ms', 's', 'min', 'hrs', 'days']
    last = names.pop(0)
    for name, ratio in zip(names, ratios):
        if value / ratio < 0.3:
            break
        value /= ratio
        last = name
    return f"{format(value, display)} {last}"


def test_one(kernel_size=1024,
             channels=8,
             batch_size=16,
             length=16000 * 5,
             stride=1):
    print(f"Benchmark for kernel_size={kernel_size} "
          f"stride={stride} channels={channels}")
    device = "cuda"
    signal = torch.randn(batch_size, channels, length, device=device)
    conv = nn.Conv1d(
        channels, channels, kernel_size, stride=stride, bias=False).to(device)
    fft_conv = FastConv(
        channels, channels, kernel_size, stride=stride, bias=False).to(device)
    fft_conv.weight = conv.weight

    conv_output = conv(signal)
    fft_output = fft_conv(signal)
    error = torch.abs(conv_output - fft_output)
    print("\tMean error={:.2g}, max error={:.2g}".format(
        error.mean(), error.max()))
    torch.backends.cudnn.benchmark = False
    print("\tCudnn benchmark = False: {}".format(
        human_seconds(profile(conv, signal))))
    torch.backends.cudnn.benchmark = True
    print("\tCudnn benchmark = True: {}".format(
        human_seconds(profile(conv, signal))))
    print("\tFFT Conv: {}".format(human_seconds(profile(fft_conv, signal))))


def test():
    print("torch.backends.cudnn.is_available(): ",
          torch.backends.cudnn.is_available(),
          "\ntorch.backends.cudnn.version(): ", torch.backends.cudnn.version())
    grid = product([256, 1024, 2048], [64], [1, 16])
    for kernel_size, channels, stride in grid:
        test_one(kernel_size=kernel_size, channels=channels, stride=stride)


if __name__ == "__main__":
    test()
```

Output from the above code in my environment:
```
torch.backends.cudnn.is_available():  True
torch.backends.cudnn.version():  7501
Benchmark for kernel_size=256 stride=1 channels=64
        Mean error=9.5e-07, max error=2.5e-05
        Cudnn benchmark = False: 0.33 s
        Cudnn benchmark = True: 0.33 s
        FFT Conv: 160.82 ms
Benchmark for kernel_size=256 stride=16 channels=64
        Mean error=9.5e-07, max error=1.8e-05
        Cudnn benchmark = False: 23.70 ms
        Cudnn benchmark = True: 23.12 ms
        FFT Conv: 160.16 ms
Benchmark for kernel_size=1024 stride=1 channels=64
        Mean error=1.9e-06, max error=4.1e-05
        Cudnn benchmark = False: 3.28 s
        Cudnn benchmark = True: 3.31 s
        FFT Conv: 160.77 ms
Benchmark for kernel_size=1024 stride=16 channels=64
        Mean error=1.9e-06, max error=3.5e-05
        Cudnn benchmark = False: 213.10 ms
        Cudnn benchmark = True: 213.29 ms
        FFT Conv: 158.74 ms
Benchmark for kernel_size=2048 stride=1 channels=64
        Mean error=2.6e-06, max error=6e-05
        Cudnn benchmark = False: 6.68 s
        Cudnn benchmark = True: 6.68 s
        FFT Conv: 160.56 ms
Benchmark for kernel_size=2048 stride=16 channels=64
        Mean error=2.6e-06, max error=4.8e-05
        Cudnn benchmark = False: 0.43 s
        Cudnn benchmark = True: 0.43 s
        FFT Conv: 160.41 ms
```

## Expected behavior

When using a stride of 1 and large kernel size, the FFT implementation is much faster than the default one. The FFT one takes 160ms whatever the size of the kernel, versus 3.3 seconds (resp 6.7) for the default one with a kernel size of 1024 (resp 2048).  For large strides, the cudnn implementation is competitive or faster as expected (the FFT only has an interest if we want the convolution for all positions).
I would expect cudnn to provide a fast implementation for large kernels with low stride, which can be especially useful in audio (filters implementation). When talking about this around me, most people were surprised as it has been announced that an FFT based implementation was added to cudnn.

## Environment

Collecting environment information...
PyTorch version: 1.1.0
Is debug build: No
CUDA used to build PyTorch: 10.0.130

OS: Ubuntu 18.04.1 LTS
GCC version: (Ubuntu 7.3.0-27ubuntu1~18.04) 7.3.0
CMake version: version 3.13.4

Python version: 3.7
Is CUDA available: Yes
CUDA runtime version: Could not collect
GPU models and configuration:
GPU 0: Quadro GP100
GPU 1: Quadro GP100

Nvidia driver version: 410.79
cuDNN version: Could not collect

Versions of relevant libraries:
[pip3] numpy==1.15.4
[pip3] torch==1.1.0
[pip3] torchvision==0.2.2
[conda] blas                      1.0                         mkl
[conda] mkl                       2018.0.3                      1
[conda] mkl_fft                   1.0.6            py37h7dd41cf_0
[conda] mkl_random                1.0.1            py37h4414c95_1
[conda] pytorch                   1.1.0           py3.7_cuda10.0.130_cudnn7.5.1_0    pytorch
[conda] torchvision               0.2.2                      py_3    pytorch


cc @csarofeen @ptrblck @mruberry @peterbell10 @VitalyFedyunin @ngimel