[quant] Support integer implementations for adaptive_avg_pool2d (pytorch#104226)

Summary:
This is needed for representing quantized model in pt2 export quantization flow

Test Plan:
tested by opinfo, python test/test_ops.py

Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

Pull Request resolved: pytorch#104226
Approved by: https://github.com/jgong5, https://github.com/andrewor14

Author: jerryzh168

aten/src/ATen/native/cpu/AdaptiveAvgPoolKernel.cpp

@@ -367,7 +367,7 @@ void adaptive_avg_pool2d_kernel_impl(
     IntArrayRef output_size) {
   switch (input.suggest_memory_format()) {
     case at::MemoryFormat::Contiguous: {
-      AT_DISPATCH_FLOATING_TYPES_AND(ScalarType::BFloat16, input.scalar_type(), "adaptive_avg_pool2d", [&] {
+      AT_DISPATCH_ALL_TYPES_AND(ScalarType::BFloat16, input.scalar_type(), "adaptive_avg_pool2d", [&] {
         if (input.scalar_type() == ScalarType::BFloat16) {
           cpu_adaptive_avg_pool<BFloat16, /*accscalar_t*/float>(output, input, output_size);
         } else {
@@ -377,7 +377,7 @@ void adaptive_avg_pool2d_kernel_impl(
       break;
     }
     case at::MemoryFormat::ChannelsLast: {
-      AT_DISPATCH_FLOATING_TYPES_AND(ScalarType::BFloat16, input.scalar_type(), "adaptive_avg_pool2d_channels_last", [&]{
+      AT_DISPATCH_ALL_TYPES_AND(ScalarType::BFloat16, input.scalar_type(), "adaptive_avg_pool2d_channels_last", [&]{
         cpu_adaptive_avg_pool_channels_last<scalar_t>(output, input, output_size);
       });
       break;

test/nn/test_pooling.py

@@ -754,6 +754,9 @@ def test_adaptive_avg_pool3d_output_size_one(self, device):
     def test_adaptive_pooling_no_suppot_input(self, device, dtype):
         for numel in (2, 3):
             for pool_type in ('Max', 'Avg'):
+                # adapative_avg_pool2d for int is implemented
+                if numel == 2 and pool_type == 'Avg':
+                    continue
                 cls_name = 'Adaptive{}Pool{}d'.format(pool_type, numel)
                 module_cls = getattr(nn, cls_name)
                 output_size = (2,) * numel

torch/_decomp/decompositions.py

@@ -1842,6 +1842,10 @@ def adaptive_avg_pool2d(input: Tensor, output_size: Tuple[int, int]):
             f"non-batch dimensions, but input has shape {tuple(shape)}.",
         )
 
+    # TODO: decompose integer path
+    if input.dtype in [torch.int8, torch.uint8, torch.int16, torch.int32, torch.int64]:
+        return torch.nn.functional.adaptive_avg_pool2d(input, output_size)
+
     # Optimisation (we should also do this in the kernel implementation)
     if shape[-2] % output_size[-2] == 0 and shape[-1] % output_size[-1] == 0:
         stride = tuple(i // o for i, o in zip(shape[-2:], output_size))

torch/testing/_internal/common_methods_invocations.py

@@ -12009,7 +12009,7 @@ def reference_flatten(input, start_dim=0, end_dim=-1):
            error_inputs_func=error_inputs_adaptive_avg_pool1d,
            sample_inputs_func=sample_inputs_adaptive_avg_pool1d),
     OpInfo('nn.functional.adaptive_avg_pool2d',
-           dtypes=floating_types_and(torch.bfloat16),
+           dtypes=all_types_and(torch.bfloat16),
            dtypesIfCUDA=floating_types_and(torch.half, torch.bfloat16),
            decorators=(
                # RuntimeError: