Repository: pytorch
Issue: DNNL's backward pass much slower when using nn.grad.conv2d_input and nn.grad.conv2d_weight
Labels: module: autograd,module: nn,triaged,module: mkldnn

## 🐛 Bug

Hello, DNNL's performance is much worse in the backward pass of a custom convolution operation. When stride > 1 performance further decreases. Number of calls to the dilated_conv2d operation increases heavily and the backward pass is about 100x slow. You can find the profiling data for the forward and backward pass below.


## To Reproduce
```
import torch
import torch.nn as nn
import torch.nn.functional as F

class op_Conv2D(torch.autograd.Function):
  
  @staticmethod
  def forward(ctx, input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
    ctx.save_for_backward(input, weight, bias)
    ctx.stride = stride
    ctx.padding = padding
    ctx.groups = groups
    ctx.dilation = dilation
    out = F.conv2d(input, weight, bias, stride, padding, dilation, groups)
    return out
    
  @staticmethod
  def backward(ctx, grad_output):
    input, weight, bias = ctx.saved_tensors
    grad_input = grad_weight= grad_bias = None

    if ctx.needs_input_grad[0]:
      grad_input = torch.nn.grad.conv2d_input(input.shape, weight, grad_output, ctx.stride, ctx.padding, ctx.dilation, ctx.groups)
        
    if ctx.needs_input_grad[1]:
      grad_weight = torch.nn.grad.conv2d_weight(input, weight.shape, grad_output, ctx.stride, ctx.padding, ctx.dilation, ctx.groups)
            
    if bias is not None and ctx.needs_input_grad[2]:
      grad_bias = grad_output.sum((0,2,3))
        
    if bias is not None:
      return grad_input, grad_weight, grad_bias, None, None, None, None
    else:
      return grad_input, grad_weight, None, None, None, None, None

w = nn.Parameter(torch.randn(32, 32, 3, 3, device=torch.device("cpu")))
b = nn.Parameter(torch.zeros(32, device=torch.device("cpu"))) 

data = torch.rand(32, 32, 3, 3)

with torch.autograd.profiler.profile(use_cuda=False) as prof:
  x_act = op_Conv2D.apply(data,w, b, 2, 0, 1, 1)
print("forward pass F.conv2d --- \n", prof.key_averages().table(sort_by="self_cpu_time_total"))

with torch.autograd.profiler.profile(use_cuda=False) as prof:
  x_act.sum().backward()
print("backward pass F.conv2d --- \n", prof.key_averages().table(sort_by="self_cpu_time_total"))


conv_block = nn.Conv2d(32, 32, 3, stride=2).cpu()

with torch.autograd.profiler.profile(use_cuda=False) as prof:
  x_act = conv_block(data)
print("forward pass nn.Conv2d --- \n", prof.key_averages().table(sort_by="self_cpu_time_total"))

with torch.autograd.profiler.profile(use_cuda=False) as prof:
  x_act.sum().backward()
print("backward pass nn.Conv2d --- \n", prof.key_averages().table(sort_by="self_cpu_time_total"))
```


## Expected behavior

Number of calls to dilated_conv2d in backward pass should be much less making backward pass much faster
## Environment
PyTorch version: 1.6.0a0+47fec01
Is debug build: No
CUDA used to build PyTorch: Could not collect

OS: Ubuntu 18.04.2 LTS
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
CMake version: version 3.17.0

Python version: 3.7
Is CUDA available: No
CUDA runtime version: 10.1.243
GPU models and configuration: GPU 0: GeForce GTX 850M
Nvidia driver version: 435.21
cuDNN version: /usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn.so.7.6.4

Versions of relevant libraries:
[pip3] numpy==1.13.3
[conda] mkl                       2020.1                      217    conda-forge
[conda] mkl-include               2020.1                      217    conda-forge
[conda] torch                     1.6.0a0+47fec01          pypi_0    pypi
[conda] torchvision               0.7.0a0+c4f77c2          pypi_0    pypi

## Additional Context
Forward Pass
```
 --------------------------  ---------------  ---------------  ---------------  ---------------  ---------------  ---------------  
Name                        Self CPU total %  Self CPU total   CPU total %      CPU total        CPU time avg     Number of Calls  
--------------------------  ---------------  ---------------  ---------------  ---------------  ---------------  ---------------  
mkldnn_convolution          36.62%           120.430ms        73.66%           242.227ms        5.767ms          42               
C_Connect                   28.26%           92.930ms         31.67%           104.140ms        4.734ms          22               
S_Connect                   18.63%           61.264ms         20.19%           66.377ms         8.297ms          8                
native_batch_norm           4.63%            15.231ms         10.89%           35.819ms         895.465us        40               
_cat                        3.58%            11.760ms         4.21%            13.851ms         364.492us        38               
add                         1.08%            3.536ms          2.51%            8.261ms          56.585us         146              
empty                       0.78%            2.562ms          0.79%            2.603ms          7.912us          329              
mul                         0.57%            1.861ms          1.16%            3.822ms          238.906us        16               
sub                         0.55%            1.796ms          1.12%            3.692ms          230.766us        16               
op_Conv2D                   0.47%            1.532ms          37.70%           123.974ms        5.904ms          21               
is_leaf                     0.41%            1.341ms          0.50%            1.648ms          0.659us          2500             
leaky_relu                  0.39%            1.277ms          0.84%            2.759ms          153.299us        18   

```
Backward Pass
```
-----------------------------------  ---------------  ---------------  ---------------  ---------------  ---------------  ---------------  
Name                                 Self CPU total %  Self CPU total   CPU total %      CPU total        CPU time avg     Number of Calls  
-----------------------------------  ---------------  ---------------  ---------------  ---------------  ---------------  ---------------  
slow_conv_dilated2d                  20.68%           1.287s           69.93%           4.352s           37.949us         114688           
slow_conv_transpose2d                13.49%           839.662ms        27.88%           1.735s           43.386ms         40               
copy_                                11.12%           691.946ms        11.37%           707.678ms        2.842ms          249              
mkldnn_convolution                   10.04%           624.944ms        20.11%           1.252s           41.719ms         30               
size                                 6.53%            406.650ms        6.53%            406.650ms        0.243us          1674695          
op_Conv2DBackward                    5.99%            372.616ms        95.47%           5.942s           282.958ms        21               
slice                                4.58%            285.351ms        8.29%            515.909ms        2.246us          229714           
_cat                                 4.14%            257.976ms        5.55%            345.569ms        12.342ms         28               
_convolution                         2.87%            178.902ms        77.67%           4.834s           117.910ms        41               
empty                                2.75%            171.139ms        2.78%            172.734ms        1.499us          115224           
fill_                                2.39%            149.045ms        2.39%            149.047ms        2.592us          57500            
as_strided                           2.36%            147.065ms        2.36%            147.065ms        0.631us          233126           
select                               2.33%            144.808ms        4.15%            258.508ms        2.193us          117878           
narrow                               1.96%            121.740ms        8.09%            503.789ms        4.384us          114916           
C_ConnectBackward                    1.95%            121.317ms        2.45%            152.702ms        6.941ms          22               
contiguous                           1.40%            86.923ms         1.56%            97.384ms         0.424us          229539
```


cc @ezyang @SsnL @albanD @zou3519 @gqchen @gujinghui @PenghuiCheng @XiaobingSuper @jianyuh @VitalyFedyunin @mruberry