Repository: pytorch
Issue: [quant] Quantized AdaptivePool3d is much slower for ChannelsLast3d.
Labels: module: performance,oncall: quantization,low priority,triaged,module: memory format

cc @jerryzh168 @jianyuh @raghuramank100 @jamesr66a @vkuzo @jgong5 @Xia-Weiwen @leslie-fang-intel @dzhulgakov @VitalyFedyunin @ppwwyyxx @ngimel @z-a-f, @vkuzo.

## 🐛 Bug

As part of the refactoring effort in #40316, I ran some benchmarks on `adaptive_avg_pool(2/3)d` to see what is the difference in execution time for `MemoryFormat::Contiguous` vs `MemoryFormat::ChannelsLast3d` for quantized tensors. Results (for 3d) are shown below:

### Python
![before_py](https://user-images.githubusercontent.com/37529096/89712024-40219380-d9e2-11ea-873d-574ac19ec6b7.png)

### C++
![before_cpp](https://user-images.githubusercontent.com/37529096/89712030-49aafb80-d9e2-11ea-86b0-e3b3211ecc62.png)

## Expected Behaviour
`adaptive_avg_pool3d` for `channelslast3d` format should either be as fast as the `contiguous` format or faster.

## To Reproduce

Steps to reproduce the behavior:

1. Build ([73642d9425a358b51a683cf6f95852d06cba1096](https://github.com/pytorch/pytorch/tree/73642d9425a358b51a683cf6f95852d06cba1096)).

### Python
1. `pip install termcolor` and run Python program below
```
import time
import numpy as np
import torch
from termcolor import colored

def time_adaptive_avg_pool3d(X, output_size_d, output_size_h, output_size_w, iterations):
    X, (scale, zero_point, torch_type) = X
    qX_ncdhw = torch.quantize_per_tensor(torch.from_numpy(X), scale=scale,
                                    zero_point=zero_point, dtype=torch_type)
    qX_ndhwc = qX_ncdhw.contiguous(memory_format=torch.channels_last_3d)
    output_size = (output_size_d, output_size_h, output_size_w)
    assert(qX_ndhwc.stride() != sorted(qX_ndhwc.stride()))
    assert(qX_ncdhw.is_contiguous(memory_format=torch.contiguous_format))
    assert(qX_ndhwc.is_contiguous(memory_format=torch.channels_last_3d))
    start = time.time()
    for _ in range(iterations):
        X_hat = torch.nn.quantized.functional.adaptive_avg_pool3d(qX_ncdhw, output_size=output_size)
    qncdhw_end = time.time() - start
    start = time.time()
    for _ in range(iterations):
        X_hat = torch.nn.quantized.functional.adaptive_avg_pool3d(qX_ndhwc, output_size=output_size)
    qndhwc_end = time.time() - start
    return qncdhw_end*1000/iterations, qndhwc_end*1000/iterations

iterations = 10000
print("iterations = {}".format(iterations))
print("Benchmark", "Time(ms)", sep="\t\t\t\t\t")
for torch_type in (torch.qint8, torch.quint8, torch.qint32):
    for channel in (4,8,64,256):
        X = np.random.rand(1, channel, 56, 56, 4).astype(np.float32), (0.5, 1, torch_type)
        ts = time_adaptive_avg_pool3d(X, 5, 5, 5, iterations)
        print(colored("adaptive_avg_pool3d({}, {}, {})".format(str(torch_type), channel, "ncdhw"), 'green'), colored(ts[0], 'yellow'), sep="\t")
        print(colored("adaptive_avg_pool3d({}, {}, {})".format(str(torch_type), channel, "ndhwc"), 'green'), colored(ts[1], 'yellow'), sep="\t")
```

### C++
1. `git clone https://github.com/google/benchmark.git`
1. `git clone https://github.com/google/googletest.git benchmark/googletest`
1. Copy `CMakeLists.txt` and `main.cpp` in the same top level directory as `benchmark`.
1. `mkdir build && cd build`
1. ```cmake -DCMAKE_PREFIX_PATH=`python -c 'import torch;print(torch.utils.cmake_prefix_path)'` -DCUDA_TOOLKIT_ROOT_DIR=<where your CUDA directory is> ..``` (`CUDA_TOOLKIT_ROOT_DIR` can be omitted if not torch wasn't build with CUDA.)
1. `cmake --build . --config Release`
1. `./main`
```
# CMakeLists.txt
cmake_minimum_required(VERSION 3.10 FATAL_ERROR)
project(time_adaptive_average_pool VERSION 0.1.0)

find_package(Torch REQUIRED)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
add_subdirectory(benchmark)

add_executable(main main.cpp)
target_link_libraries(main ${TORCH_LIBRARIES})
set_property(TARGET main PROPERTY CXX_STANDARD 14)
target_link_libraries(main benchmark::benchmark)
```
```
// main.cpp
#include <benchmark/benchmark.h>
#include <torch/torch.h>

torch::Device device(torch::kCPU);

static void BM_TORCH_QADAPTIVE_AVG_POOL3D_NCDHW(benchmark::State& state) {
  auto x_ncdhw = torch::rand({1, state.range(0), 56, 56, 4}, device);
  auto qx_ncdhw = torch::quantize_per_tensor(x_ncdhw, 0.5, 1, torch::kQUInt8);
  torch::Tensor X_hat;
  for (auto _ : state)
    X_hat = torch::nn::functional::adaptive_avg_pool3d(
        qx_ncdhw, torch::nn::AdaptiveAvgPool3dOptions({5, 5, 5}));
}

static void BM_TORCH_QADAPTIVE_AVG_POOL3D_NDHWC(benchmark::State& state) {
  auto x_ncdhw = torch::rand({1, state.range(0), 56, 56, 4}, device);
  auto qx_ncdhw = torch::quantize_per_tensor(x_ncdhw, 0.5, 1, torch::kQUInt8);
  auto qx_ndhwc = qx_ncdhw.contiguous(torch::MemoryFormat::ChannelsLast3d);
  torch::Tensor X_hat;
  for (auto _ : state)
    X_hat = torch::nn::functional::adaptive_avg_pool3d(
        qx_ndhwc, torch::nn::AdaptiveAvgPool3dOptions({5, 5, 5}));
}

BENCHMARK(BM_TORCH_QADAPTIVE_AVG_POOL3D_NCDHW)
    ->RangeMultiplier(8)
    ->Range(4, 1024);
BENCHMARK(BM_TORCH_QADAPTIVE_AVG_POOL3D_NDHWC)
    ->RangeMultiplier(8)
    ->Range(4, 1024);
BENCHMARK_MAIN();
```
## Environment

```
PyTorch version: 1.7.0a0+73642d9
Is debug build: No
CUDA used to build PyTorch: 10.2

OS: Ubuntu 18.04.4 LTS
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
CMake version: version 3.14.0

Python version: 3.8
Is CUDA available: Yes
CUDA runtime version: 10.2.89
GPU models and configuration: GPU 0: GeForce RTX 2070

Nvidia driver version: 440.33.01
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5

Versions of relevant libraries:
[pip3] numpy==1.18.5
[pip3] torch==1.7.0a0+73642d9
[conda] blas                      1.0                         mkl  
[conda] magma-cuda102             2.5.2                         1    pytorch
[conda] mkl                       2020.1                      217  
[conda] mkl-include               2020.1                      217  
[conda] mkl-service               2.3.0            py38he904b0f_0  
[conda] mkl_fft                   1.1.0            py38h23d657b_0  
[conda] mkl_random                1.1.1            py38h0573a6f_0  
[conda] numpy                     1.18.5           py38ha1c710e_0  
[conda] numpy-base                1.18.5           py38hde5b4d6_0  
[conda] torch                     1.7.0a0+73642d9           dev_0    <develop>
```

## Notes
1. The python program was derived from the [test_quantized_op.py](https://github.com/pytorch/pytorch/blob/73642d9425a358b51a683cf6f95852d06cba1096/test/quantization/test_quantized_op.py#L1300)