1 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS 2 #include <ATen/core/Tensor.h> 3 #include <ATen/native/AdaptivePooling.h> 4 #include <ATen/native/xnnpack/Engine.h> 5 #include <c10/util/irange.h> 6 7 #ifndef AT_PER_OPERATOR_HEADERS 8 #include <ATen/Functions.h> 9 #include <ATen/NativeFunctions.h> 10 #else 11 #include <ATen/ops/_adaptive_avg_pool2d.h> 12 #include <ATen/ops/_adaptive_avg_pool2d_backward_native.h> 13 #include <ATen/ops/_adaptive_avg_pool2d_native.h> 14 #include <ATen/ops/adaptive_avg_pool2d_native.h> 15 #include <ATen/ops/empty.h> 16 #include <ATen/ops/mkldnn_adaptive_avg_pool2d.h> 17 #endif 18 19 20 namespace at::native { 21 22 namespace { 23 adaptive_avg_pool2d_out_cpu_template(at::Tensor & output,at::Tensor const & input,IntArrayRef output_size)24 void adaptive_avg_pool2d_out_cpu_template( 25 at::Tensor& output, 26 at::Tensor const& input, 27 IntArrayRef output_size) 28 { 29 TORCH_CHECK(output_size.size() == 2, "adaptive_avg_pool2d: output_size must be 2"); 30 int64_t ndim = input.dim(); 31 TORCH_CHECK((ndim == 3 || ndim == 4), 32 "adaptive_avg_pool2d(): Expected 3D or 4D tensor, but got ", input.sizes()); 33 for (const auto i : {-2, -1}) { 34 TORCH_CHECK(input.size(i) > 0, 35 "adaptive_avg_pool2d(): Expected input to have non-zero size for non-batch dimensions, " 36 "but input has sizes ", input.sizes(), " with dimension ", i + ndim, " being " 37 "empty"); 38 } 39 40 TORCH_CHECK(input.dtype() == output.dtype(), 41 "expected dtype ", input.dtype(), " for `output` but got dtype ", output.dtype()); 42 43 int64_t channels = input.size(-3); 44 int64_t output_height = output_size[0]; 45 int64_t output_width = output_size[1]; 46 47 if (ndim == 3) { 48 output.resize_({channels, output_height, output_width}); 49 } else { 50 int64_t nbatch = input.size(0); 51 output.resize_({nbatch, channels, output_height, output_width}, input.suggest_memory_format()); 52 } 53 54 if (output.numel() == 0) { 55 return; 56 } 57 58 adaptive_avg_pool2d_kernel(kCPU, output, input, output_size); 59 } 60 adaptive_avg_pool2d_backward_out_cpu_template(Tensor & grad_input,const Tensor & grad_output,const Tensor & input)61 Tensor& adaptive_avg_pool2d_backward_out_cpu_template( 62 Tensor& grad_input, 63 const Tensor& grad_output, 64 const Tensor& input) 65 { 66 int64_t ndim = grad_output.ndimension(); 67 for (const auto i : c10::irange(1, ndim)) { 68 TORCH_CHECK(grad_output.size(i) > 0, 69 "adaptive_avg_pool2d_backward(): Expected grad_output to have non-zero size for non-batch dimensions, " 70 "but grad_output has sizes ", grad_output.sizes(), " with dimension ", i, " being " 71 "empty"); 72 } 73 74 TORCH_CHECK((ndim == 3 || ndim == 4), 75 "adaptive_avg_pool2d_backward(): Expected 3D or 4D tensor, but got ", input.sizes()); 76 TORCH_CHECK(input.dtype() == grad_output.dtype(), 77 "expected dtype ", input.dtype(), " for `grad_output` but got dtype ", grad_output.dtype()); 78 TORCH_CHECK(input.dtype() == grad_input.dtype(), 79 "expected dtype ", input.dtype(), " for `grad_input` but got dtype ", grad_input.dtype()); 80 81 grad_input.resize_(input.sizes(), input.suggest_memory_format()); 82 grad_input.zero_(); 83 84 adaptive_avg_pool2d_backward_kernel(kCPU, grad_input, grad_output); 85 return grad_input; 86 } 87 88 } // namespace 89 adaptive_avg_pool2d_out_cpu(const Tensor & input,IntArrayRef output_size,Tensor & output)90 Tensor& adaptive_avg_pool2d_out_cpu(const Tensor& input, 91 IntArrayRef output_size, 92 Tensor& output) 93 { 94 adaptive_avg_pool2d_out_cpu_template( 95 output, input, output_size); 96 return output; 97 } 98 adaptive_avg_pool2d_cpu(at::Tensor const & input,IntArrayRef output_size)99 Tensor adaptive_avg_pool2d_cpu( 100 at::Tensor const& input, 101 IntArrayRef output_size) 102 { 103 auto output = at::empty({0}, input.options()); 104 adaptive_avg_pool2d_out_cpu_template( 105 output, input, output_size); 106 return output; 107 } 108 adaptive_avg_pool2d_symint(at::Tensor const & input,SymIntArrayRef output_size)109 Tensor adaptive_avg_pool2d_symint(at::Tensor const& input, SymIntArrayRef output_size) { 110 TORCH_CHECK(output_size.size() == 2, "adaptive_avg_pool2d: output_size must be 2"); 111 TORCH_CHECK( 112 (output_size[0] >= 0 && output_size[1] >= 0), 113 "adaptive_avg_pool2d: elements of output_size must be greater than or equal to 0 ", 114 "but received {", output_size[0], ", ", output_size[1], "}"); 115 116 if (input.is_mkldnn()) { 117 return at::mkldnn_adaptive_avg_pool2d(input, C10_AS_INTARRAYREF_SLOW(output_size)); 118 } 119 120 if (!input.is_quantized() && output_size[0] == 1 && output_size[1] == 1) { 121 // in this case, adaptive pooling is just computing mean over hw 122 // dimensions, which can be done more efficiently 123 #if defined(C10_MOBILE) && defined(USE_XNNPACK) 124 if (xnnpack::use_global_average_pool(input)) { 125 return xnnpack::global_average_pool(input); 126 } 127 #endif 128 129 Tensor out = input.mean({-1, -2}, /* keepdim = */ true); 130 if (input.suggest_memory_format() == at::MemoryFormat::ChannelsLast) { 131 // assert ndim == 4, since ndim = 3 doesn't give channels_last 132 const auto n = input.sym_size(0); 133 const auto c = input.sym_size(1); 134 out.as_strided__symint({n, c, 1, 1}, {c, 1, c, c}); 135 } 136 return out; 137 } else { 138 return _adaptive_avg_pool2d_symint(input, output_size); 139 } 140 } 141 adaptive_avg_pool2d_backward_cpu(const Tensor & grad_output,const Tensor & input)142 Tensor adaptive_avg_pool2d_backward_cpu( 143 const Tensor& grad_output, 144 const Tensor& input) 145 { 146 auto grad_input = at::empty({0}, input.options()); 147 adaptive_avg_pool2d_backward_out_cpu_template( 148 grad_input, grad_output, input); 149 return grad_input; 150 } 151 152 DEFINE_DISPATCH(adaptive_avg_pool2d_kernel); 153 DEFINE_DISPATCH(adaptive_avg_pool2d_backward_kernel); 154 155 } // namespace at::native 156