Update on "add Half support for AdaptiveAvgPool2d and AdaptiveMaxPool2d on CPU"

CaoE · CaoE · commit c5859865cf01 · 2023-10-23T01:25:18.000-07:00
### Testing

Single core:

AdaptiveMaxPool2d:
shape | fp32 forward / ms | fp16 forward / ms | bf16 forward / ms | fp32 backward / ms | fp16 backward / ms | bf16 backward / ms
-- | -- | -- | -- | -- | -- | --
input size: (2, 56, 264, 264), output size: (100, 100)  | 71.5826 | 78.7460 | 85.7195 | 7.3925 | 6.0618 | 6.2596
input size: (2, 56, 264, 264), output size: (50, 50)  | 28.122 | 30.8572 | 36.6366 | 6.2645 | 3.4781 | 3.6628
input size: (32, 32, 100, 100), output size: (50, 50)  | 109.2978 | 115.0330 | 121.9500 | 13.4329 | 10.2769 | 12.1975
input size: (16, 4, 300, 300), output size: (100, 100) | 34.1849 | 36.5876 | 40.9862 | 4.7719 | 4.3362 | 4.1417


28 cores:

AdaptiveMaxPool2d:

shape | fp32 forward / ms | fp16 forward / ms | bf16 forward / ms | fp32 backward / ms | fp16 backward / ms | bf16 backward / ms
-- | -- | -- | -- | -- | -- | --
input size: (2, 56, 264, 264), output size: (100, 100)  | 3.1809 | 3.5057 | 3.6728 | 0.6657 | 0.3138 | 0.2934
input size: (2, 56, 264, 264), output size: (50, 50)  | 1.2779 | 1.3869 | 1.5238 | 0.4223 | 0.1775 | 0.1825
input size: (32, 32, 100, 100), output size: (50, 50)  | 4.7942 | 4.9670 | 5.2330 | 1.7146 | 0.6477 | 0.7001
input size: (16, 4, 300, 300), output size: (100, 100) | 1.9522 | 2.0879 | 2.3155 | 0.4370 | 0.3175 | 0.2828


cc jgong5 mingfeima XiaobingSuper sanchitintel ashokei jingxu10

[ghstack-poisoned]
diff --git a/aten/src/ATen/native/cpu/AdaptiveAvgPoolKernel.cpp b/aten/src/ATen/native/cpu/AdaptiveAvgPoolKernel.cpp
@@ -373,7 +373,7 @@ void adaptive_avg_pool2d_kernel_impl(
     case at::MemoryFormat::Contiguous: {
       AT_DISPATCH_FLOATING_TYPES_AND2(ScalarType::BFloat16, ScalarType::Half, input.scalar_type(), "adaptive_avg_pool2d", [&] {
         using param_t = at::opmath_type<scalar_t>;
-        cpu_adaptive_avg_pool<scalar_t, /*opmath_t*/param_t>(output, input, output_size);
+        cpu_adaptive_avg_pool<scalar_t, /*accscalar_t*/param_t>(output, input, output_size);
       });
       break;
     }
diff --git a/aten/src/ATen/native/cpu/AdaptiveMaxPoolKernel.cpp b/aten/src/ATen/native/cpu/AdaptiveMaxPoolKernel.cpp
@@ -443,7 +443,7 @@ void adaptive_max_pool2d_kernel_impl(
     case at::MemoryFormat::Contiguous: {
       AT_DISPATCH_FLOATING_TYPES_AND2(ScalarType::BFloat16, ScalarType::Half, input.scalar_type(), "adaptive_max_pool2d", [&] {
         using param_t = at::opmath_type<scalar_t>;
-        cpu_adaptive_max_pool<scalar_t, /*opmath_t*/param_t>(output, indices, input, output_size);
+        cpu_adaptive_max_pool<scalar_t, /*accscalar_t*/param_t>(output, indices, input, output_size);
       });
       break;
     }
