Repository: pytorch
Issue: Performance bugs of transpose2d on A100
Labels: module: performance,module: cudnn,triaged

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->
The performance of running `transpose2d` with different datatypes on A100 is unexpected.

## To Reproduce
```
import torch
import numpy as np

def conv_transpose2d(H, W, N, C, K, kernel_size, stride, padding, dtype):
  A_np = np.random.uniform(-10, 10, [N, K, H, W]).astype("float32")
  B_np = np.random.uniform(-10, 10, [K, C, kernel_size, kernel_size]).astype("float32")

  torch.backends.cuda.matmul.allow_tf32 = False
  torch.backends.cudnn.allow_tf32 = False

  if dtype == "FP16":
    A_torch = torch.tensor(A_np).type(torch.float16).cuda()
    B_torch = torch.tensor(B_np).type(torch.float16).cuda()
  elif dtype == "BF16":
    A_torch = torch.tensor(A_np).type(torch.bfloat16).cuda()
    B_torch = torch.tensor(B_np).type(torch.bfloat16).cuda()
  elif dtype == "FP32":
    A_torch = torch.tensor(A_np).type(torch.float32).cuda()
    B_torch = torch.tensor(B_np).type(torch.float32).cuda()
  elif dtype == "TF32":
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    A_torch = torch.tensor(A_np).type(torch.float32).cuda()
    B_torch = torch.tensor(B_np).type(torch.float32).cuda()
  elif dtype == "FP64":
    A_torch = torch.tensor(A_np).type(torch.float64).cuda()
    B_torch = torch.tensor(B_np).type(torch.float64).cuda()
  else:
    assert False, "wrong type: " + dtype

  number = 10
  repeats = 10

  for i in range(repeats):
      time_record = []
      for j in range(number):
        torch.cuda.synchronize()
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()

        C_torch = torch.nn.functional.conv_transpose2d(A_torch, B_torch, bias=None, stride=stride, padding=padding)

        end.record()
        torch.cuda.synchronize()
        total = start.elapsed_time(end)
        time_record.append(total)
      if i == repeats - 1:
        mean_cost = np.mean(time_record)
  return mean_cost

_ = None
#  (  N,   C,   H,   W,   K, kernel_size, stride, padding, dilation)
transpose2d_config = [
    ( _,   3, 112, 112,  64,           3,      7,       3,        1), # stem

    ( _,  64,  56,  56,  64,           3,      1,       1,        1), # layer1 x 4

    ( _,  64,  56,  56, 128,           1,      2,       0,        1), # layer2 downsample
    
    ( _,  64,  56,  56, 128,           3,      2,       1,        1), # layer2
    ( _, 128,  28,  28, 128,           3,      1,       1,        1), # layer2 x 3

    ( _, 128,  28,  28, 256,           1,      2,       0,        1), # layer3 downsample
    ( _, 128,  28,  28, 256,           3,      2,       1,        1), # layer3
    ( _, 256,  14,  14, 256,           3,      1,       1,        1), # layer3 x 3

    ( _, 256,  14,  14, 512,           1,      2,       0,        1), # layer4 downsample
    ( _, 256,  14,  14, 512,           3,      2,       1,        1), # layer4
    ( _, 256,   7,   7, 512,           3,      1,       1,        1), # layer4 x 3
]

if __name__ == "__main__":
    assert torch.backends.cudnn.is_available()
    torch.backends.cudnn.enabled = True
    for dtype in ["FP16", "FP32", "TF32", "FP64", "BF16"]:
      costs = []
      for i, shape in enumerate(transpose2d_config):
          (  N,   C,   H,   W,   K, kernel_size, stride, padding, dilation) = shape
          iH = (H + 2 * padding - dilation * (kernel_size - 1) - 1) // stride + 1
          iW = (W + 2 * padding - dilation * (kernel_size - 1) - 1) // stride + 1
          N = 1
          cost = conv_transpose2d(iH, iW, N, C, K, kernel_size, stride, padding, dtype)
          costs.append(cost)
      print("dtype =", dtype, "cost =", np.mean(costs))
    print("cudnn: %s" % ("enabled" if torch.backends.cudnn.enabled else "disabled"))

```
The above code turns on CuDNN, and evaluates the performance with a reasonable number of warmup runs on GPUs.

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->
The following is the output on A100
```
dtype = FP16 cost = 0.37737396318804134
dtype = FP32 cost = 0.10127941830591722
dtype = TF32 cost = 0.10782109092582357
dtype = FP64 cost = 0.25560843660072846
dtype = BF16 cost = 0.18091752644289624
cudnn: enabled
```

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->
It's expected that `float16` will deliver the best performance with the minimum runtime cost. However, it turns out that the cost of `float16` is the largest. 
It's expected that `tf32` should be a bit faster than `fp32`, however, `tf32` turns out to be slightly slower.

## Environment
```
PyTorch version: 1.7.1+cu110
Is debug build: False
CUDA used to build PyTorch: 11.0
ROCM used to build PyTorch: N/A

OS: Red Hat Enterprise Linux (x86_64)
GCC version: (GCC) 7.2.0
Clang version: 9.0.0 (tags/RELEASE_900/final)
CMake version: version 3.12.2

Python version: 3.6 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: A100-PCIE-40GB
Nvidia driver version: 455.45.01
HIP runtime version: N/A
MIOpen runtime version: N/A
Versions of relevant libraries:
[pip3] numpy==1.19.5
[pip3] torch==1.7.1+cu110
[pip3] torchaudio==0.7.2
[pip3] torchvision==0.8.2+cu110

 - How you installed PyTorch (`conda`, `pip`, source): pip
 - CUDA/cuDNN version: 11.1

```
cc @csarofeen @ptrblck @xwang233 @VitalyFedyunin @ngimel