Repository: pytorch
Issue: Poor torch.cat performance in the quantized Unet
Labels: oncall: quantization,low priority,triaged

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->
`torch.cat` in the quantized model is a lot slower than it is in a float model. For float unet it takes 196.221ms (avg 98.110ms per 1 call) on my PC. For quantized unet, it takes 1.872s (avg 935.796ms per call).

## To Reproduce

Code to reproduce issue:
```
import torch
import torch.nn as nn
import torch.quantization as tq
import torch.autograd.profiler as tp


class ConvBlock(nn.Module):
    def __init__(self, in_c, out_c, k_size=3, stride=1, padding=1, padding_mode='zeros'):
        super().__init__()
        self.c = nn.Conv2d(in_c, out_c, kernel_size=k_size, stride=stride, padding=padding, padding_mode=padding_mode)
        self.bn = nn.BatchNorm2d(out_c)
        self.a = nn.ReLU()

    def forward(self, x):
        return self.a(self.bn(self.c(x)))

    def merge_quantize(self):
        tq.fuse_modules(self, ['c', 'bn', 'a'], inplace=True)


class Unet(nn.Module):
    def __init__(self):
        super().__init__()
        self.init_c = ConvBlock(1, 32)
        self.c1 = ConvBlock(32, 32)
        self.p1 = nn.MaxPool2d(2)
        self.c2 = ConvBlock(32, 64)
        self.c3 = ConvBlock(64, 64)
        self.p2 = nn.MaxPool2d(2)
        self.c4 = ConvBlock(64, 128)
        self.c5 = ConvBlock(128, 128)
        self.c6 = ConvBlock(128, 128)
        self.c7 = ConvBlock(128, 64)
        self.u1 = nn.Upsample(scale_factor=2)
        self.c8 = ConvBlock(128, 64)
        self.c9 = ConvBlock(64, 64)
        self.c10 = ConvBlock(64, 32)
        self.u2 = nn.Upsample(scale_factor=2)
        self.c11 = ConvBlock(64, 32)
        self.c12 = ConvBlock(32, 32)
        self.out_c = nn.Conv2d(32, 11, kernel_size=3)
        self.out_bn = nn.BatchNorm2d(11)
        self.out_a = nn.Sigmoid()
        self.quant = tq.QuantStub()
        self.dequant = tq.DeQuantStub()

    def forward(self, x):
        y = self.quant(x)
        # down path
        y = self.init_c(y)
        y1 = self.c1(y)
        y = self.p1(y1)
        y = self.c2(y)
        y2 = self.c3(y)
        y = self.p2(y2)
        # middle
        y = self.c4(y)
        y = self.c5(y)
        y = self.c6(y)
        # up path
        y = self.c7(y)
        y = self.u1(y)
        y = torch.cat([y2, y], dim=1)
        y = self.c8(y)
        y = self.c9(y)
        y = self.c10(y)
        y = self.u2(y)
        y = torch.cat([y1, y], dim=1)
        y = self.c11(y)
        y = self.c12(y)
        y = self.out_c(y)
        y = self.out_bn(y)
        y = self.dequant(y)
        y = self.out_a(y)
        return y

    def merge_quantize(self):
        for m in self.modules():
            if type(m) is ConvBlock:
                m.merge_quantize()

        tq.fuse_modules(self, ['out_c', 'out_bn'], inplace=True)


m = Unet()
m.eval()

# FP32 model
with tp.profile() as prof:
    _ = m(torch.rand([16, 1, 512, 512]))
print(prof.key_averages().table(sort_by="self_cpu_time_total"))

m.qconfig = tq.get_default_qconfig('fbgemm')
m.merge_quantize()
mp = tq.prepare(m)
_ = mp(torch.rand([4, 1, 256, 256]))
mq = tq.convert(mp)

# INT8 MODEL
with tp.profile() as prof:
    _ = mq(torch.rand([16, 1, 512, 512]))
print(prof.key_averages().table(sort_by="self_cpu_time_total"))
```

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->
`torch.cat` for quantized models should be as fast as `torch.cat` for float models (or faster).

## Environment

PyTorch version: 1.8.1+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 10 Education
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect

Python version: 3.9 (64-bit runtime)
CPU only

Versions of relevant libraries:
[pip3] numpy==1.20.2
[pip3] torch==1.8.1+cu102
[pip3] torchaudio==0.8.1
[pip3] torchvision==0.9.1+cu102
[conda] Could not collect

## Additional context

<!-- Add any other context about the problem here. -->

Full profiler output for the float model:
```
---------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
                             Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls  
---------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
         aten::mkldnn_convolution        66.74%        7.632s        66.74%        7.633s     545.192ms            14  
                  aten::threshold        13.52%        1.546s        13.52%        1.547s     118.970ms            13  
          aten::native_batch_norm        11.89%        1.359s        11.89%        1.360s      97.156ms            14  
    aten::max_pool2d_with_indices         2.05%     234.629ms         2.05%     234.748ms     117.374ms             2  
         aten::upsample_nearest2d         1.91%     218.017ms         1.91%     218.017ms     109.008ms             2  
                    aten::sigmoid         1.79%     205.106ms         1.79%     205.142ms     205.142ms             1  
                       aten::_cat         1.71%     195.907ms         1.71%     196.056ms      98.028ms             2  
                   aten::uniform_         0.20%      22.728ms         0.20%      22.728ms      22.728ms             1  
               aten::_convolution         0.09%      10.504ms        66.83%        7.643s     545.942ms            14  
                 aten::max_pool2d         0.07%       7.872ms         2.12%     242.620ms     121.310ms             2  
                      aten::empty         0.01%       1.684ms         0.01%       1.684ms      18.507us            91  
                       aten::relu         0.00%     547.500us        13.53%        1.547s     119.012ms            13  
                 aten::empty_like         0.00%     299.500us         0.01%     788.600us      18.776us            42  
     aten::_batch_norm_impl_index         0.00%     297.700us        11.90%        1.361s      97.182ms            14  
                 aten::batch_norm         0.00%     174.800us        11.90%        1.361s      97.195ms            14  
                    aten::resize_         0.00%     174.800us         0.00%     174.800us      29.133us             6  
                        aten::cat         0.00%     165.000us         1.72%     196.221ms      98.110ms             2  
                     aten::conv2d         0.00%     152.000us        66.84%        7.643s     545.962ms            14  
                aten::convolution         0.00%     129.400us        66.83%        7.643s     545.951ms            14  
                aten::as_strided_         0.00%     115.700us         0.00%     115.700us       8.264us            14  
                     aten::narrow         0.00%      35.700us         0.00%      69.100us      34.550us             2  
                       aten::rand         0.00%      31.100us         0.20%      22.780ms      22.780ms             1  
                      aten::slice         0.00%      25.200us         0.00%      33.400us      16.700us             2  
                 aten::as_strided         0.00%       8.200us         0.00%       8.200us       4.100us             2  
---------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
Self CPU time total: 11.436s
```

Full profiler output for the quantized model:

```
---------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
                             Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls  
---------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
           quantized::conv2d_relu        66.93%        5.265s        66.94%        5.266s     405.075ms            13  
                       aten::_cat        23.79%        1.871s        23.79%        1.872s     935.776ms             2  
                quantized::conv2d         3.76%     296.157ms         3.77%     296.213ms     296.213ms             1  
       aten::quantized_max_pool2d         2.22%     174.662ms         2.22%     174.730ms      87.365ms             2  
                    aten::sigmoid         1.22%      95.644ms         1.22%      95.677ms      95.677ms             1  
         aten::upsample_nearest2d         0.88%      69.058ms         0.88%      69.121ms      34.560ms             2  
                 aten::dequantize         0.46%      36.519ms         0.46%      36.546ms      36.546ms             1  
                   aten::uniform_         0.29%      23.069ms         0.29%      23.069ms      23.069ms             1  
        aten::quantize_per_tensor         0.29%      22.483ms         0.29%      22.483ms      22.483ms             1  
                 aten::max_pool2d         0.15%      11.444ms         2.37%     186.174ms      93.087ms             2  
    aten::_empty_affine_quantized         0.01%     608.800us         0.01%     608.800us      30.440us            20  
                      aten::empty         0.00%     325.100us         0.00%     325.100us      19.124us            17  
                       aten::item         0.00%     120.500us         0.00%     131.700us      65.850us             2  
                    aten::q_scale         0.00%     108.300us         0.00%     108.300us       4.512us            24  
               aten::q_zero_point         0.00%      63.700us         0.00%      63.700us       2.654us            24  
                        aten::cat         0.00%      38.500us        23.79%        1.872s     935.796ms             2  
                       aten::rand         0.00%      26.600us         0.29%      23.120ms      23.120ms             1  
        aten::_local_scalar_dense         0.00%      11.200us         0.00%      11.200us       5.600us             2  
                    aten::qscheme         0.00%       4.300us         0.00%       4.300us       2.150us             2  
---------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
Self CPU time total: 7.867s
```

cc @jerryzh168 @jianyuh @raghuramank100 @jamesr66a @vkuzo @jgong5 @Xia-Weiwen @leslie-fang-intel