Add portable adaptive avg pool2d kernel

Author: GregoryComer

kernels/portable/cpu/op__adaptive_avg_pool2d.cpp

@@ -0,0 +1,127 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <cmath>
+
+#include <executorch/kernels/portable/cpu/util/kernel_ops_util.h>
+#include <executorch/runtime/kernel/kernel_includes.h>
+
+namespace torch {
+namespace executor {
+namespace native {
+
+using Tensor = executorch::aten::Tensor;
+using ScalarType = executorch::aten::ScalarType;
+using IntArrayRef = executorch::aten::ArrayRef<int64_t>;
+
+namespace {
+
+inline int64_t
+adaptive_start_index(int64_t out_idx, int64_t out_size, int64_t in_size) {
+  return static_cast<int64_t>(
+      std::floor(static_cast<float>(out_idx * in_size) / out_size));
+}
+
+inline int64_t
+adaptive_end_index(int64_t out_idx, int64_t out_size, int64_t in_size) {
+  return static_cast<int64_t>(
+      std::ceil(static_cast<float>((out_idx + 1) * in_size) / out_size));
+}
+
+} // namespace
+
+Tensor& _adaptive_avg_pool2d_out(
+    KernelRuntimeContext& ctx,
+    const Tensor& in,
+    IntArrayRef output_size,
+    Tensor& out) {
+  ET_KERNEL_CHECK(
+      ctx,
+      check_adaptive_avg_pool2d_args(in, output_size, out),
+      InvalidArgument,
+      out);
+
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
+
+  ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out);
+
+  size_t output_ndim = 0;
+  executorch::aten::SizesType output_sizes[kTensorDimensionLimit];
+  get_adaptive_avg_pool2d_out_target_size(
+      in, output_size, output_sizes, &output_ndim);
+
+  ET_KERNEL_CHECK(
+      ctx,
+      output_size_is_valid({output_sizes, output_ndim}, 2),
+      InvalidArgument,
+      out);
+
+  ET_KERNEL_CHECK(
+      ctx,
+      resize_tensor(out, {output_sizes, output_ndim}) == Error::Ok,
+      InvalidArgument,
+      out);
+
+  ScalarType in_type = in.scalar_type();
+
+  // @lint-ignore CLANGTIDY facebook-hte-CArray
+  static constexpr const char op_name[] = "_adaptive_avg_pool2d.out";
+
+  ET_SWITCH_FLOATHBF16_TYPES_AND(Long, in_type, ctx, op_name, CTYPE, [&]() {
+    const CTYPE* const in_ptr = in.const_data_ptr<CTYPE>();
+    CTYPE* const out_ptr = out.mutable_data_ptr<CTYPE>();
+
+    const size_t ndim = in.dim();
+    const int64_t in_H = in.size(ndim - 2);
+    const int64_t in_W = in.size(ndim - 1);
+    const int64_t out_H = output_size[0];
+    const int64_t out_W = output_size[1];
+
+    const size_t channels = in.size(ndim - 3);
+    const size_t batch_size = ndim == 4 ? in.size(0) : 1;
+
+    const size_t in_plane_size = in_H * in_W;
+    const size_t out_plane_size = out_H * out_W;
+
+    for (size_t b = 0; b < batch_size; ++b) {
+      for (size_t c = 0; c < channels; ++c) {
+        const size_t plane_idx = b * channels + c;
+        const CTYPE* plane_in = in_ptr + plane_idx * in_plane_size;
+        CTYPE* plane_out = out_ptr + plane_idx * out_plane_size;
+
+        for (int64_t oh = 0; oh < out_H; ++oh) {
+          int64_t ih0 = adaptive_start_index(oh, out_H, in_H);
+          int64_t ih1 = adaptive_end_index(oh, out_H, in_H);
+
+          for (int64_t ow = 0; ow < out_W; ++ow) {
+            int64_t iw0 = adaptive_start_index(ow, out_W, in_W);
+            int64_t iw1 = adaptive_end_index(ow, out_W, in_W);
+
+            float sum = 0;
+            for (int64_t ih = ih0; ih < ih1; ++ih) {
+              for (int64_t iw = iw0; iw < iw1; ++iw) {
+                sum += plane_in[ih * in_W + iw];
+              }
+            }
+
+            int64_t count = (ih1 - ih0) * (iw1 - iw0);
+            plane_out[oh * out_W + ow] =
+                static_cast<CTYPE>(sum / static_cast<float>(count));
+          }
+        }
+      }
+    }
+  });
+
+  return out;
+}
+
+} // namespace native
+} // namespace executor
+} // namespace torch

kernels/portable/cpu/util/kernel_ops_util.cpp

@@ -262,6 +262,53 @@ bool check_arange_args(double start, double end, double step, Tensor& out) {
   return true;
 }
 
+bool check_adaptive_avg_pool2d_args(
+    const Tensor& in,
+    const IntArrayRef output_size,
+    const Tensor& out) {
+  ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, out));
+
+  ET_LOG_AND_RETURN_IF_FALSE(tensor_is_default_or_channels_last_dim_order(in));
+  ET_LOG_AND_RETURN_IF_FALSE(tensor_is_default_or_channels_last_dim_order(out));
+
+  ET_CHECK_OR_RETURN_FALSE(
+      (in.dim() == 3 && in.size(0) > 0 && in.size(1) > 0 && in.size(2) > 0) ||
+          (in.dim() == 4 && in.size(1) > 0 && in.size(2) > 0 && in.size(3) > 0),
+      "Expected 3D or 4D (batch mode) tensor with optional 0 dim batch size for input; in.dim() = %" ET_PRI_TENSOR_DIM,
+      in.dim());
+
+  ET_CHECK_OR_RETURN_FALSE(
+      output_size.size() == 2,
+      "output_size must have exactly 2 elements, but got %zu",
+      output_size.size());
+
+  ET_CHECK_OR_RETURN_FALSE(
+      output_size[0] > 0 && output_size[1] > 0,
+      "output_size must be positive, but got (%" PRId64 ", %" PRId64 ")",
+      output_size[0],
+      output_size[1]);
+
+  return true;
+}
+
+void get_adaptive_avg_pool2d_out_target_size(
+    const Tensor& in,
+    const IntArrayRef output_size,
+    executorch::aten::SizesType* const out_sizes,
+    size_t* const out_ndim) {
+  *out_ndim = in.dim();
+
+  if (in.dim() == 4) {
+    out_sizes[0] = in.size(0);
+    out_sizes[1] = in.size(1);
+  } else {
+    out_sizes[0] = in.size(0);
+  }
+
+  out_sizes[*out_ndim - 2] = output_size[0];
+  out_sizes[*out_ndim - 1] = output_size[1];
+}
+
 bool check_avg_pool2d_args(
     const Tensor& in,
     const IntArrayRef kernel_size,

kernels/portable/cpu/util/kernel_ops_util.h

@@ -384,6 +384,17 @@ void apply_kernel_2d_reduce_then_map_fn(
 // Operator specific utility functions
 //
 
+bool check_adaptive_avg_pool2d_args(
+    const Tensor& in,
+    const IntArrayRef output_size,
+    const Tensor& out);
+
+void get_adaptive_avg_pool2d_out_target_size(
+    const Tensor& in,
+    const IntArrayRef output_size,
+    executorch::aten::SizesType* const out_sizes,
+    size_t* const out_ndim);
+
 bool check_arange_args(double start, double end, double step, Tensor& out);
 
 bool check_avg_pool2d_args(

kernels/portable/functions.yaml

@@ -17,6 +17,11 @@
 # See the README.md file in this directory for a description of the syntax used
 # by this file.
 
+- op: _adaptive_avg_pool2d.out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::_adaptive_avg_pool2d_out
+
 - op: _cdist_forward.out
   kernels:
     - arg_meta: null

kernels/test/CMakeLists.txt

@@ -157,6 +157,7 @@ set(all_test_sources
     "BinaryLogicalOpTest.cpp"
     "op__to_dim_order_copy_test.cpp"
     "op__clone_dim_order_test.cpp"
+    "op__adaptive_avg_pool2d_test.cpp"
     "op_abs_test.cpp"
     "op_acos_test.cpp"
     "op_acosh_test.cpp"