Add portable upsample_bilinear2d kernel

Summary:
Add a upsample_bilinear2d kernel to the portable kernel library. This implementation re-uses some of the inner logic from the ATen implementation (see Upsample.h and UpsampleKernel.cpp), however I have not ported the outer kernel structure as it relies on TensorIterator and runtime allocation.

It may be worth re-visiting this in the future, either by looking at pulling in more of the ATen implementation or adding an optimized variant.

Differential Revision: D65756150

Author: GregoryComer

kernels/portable/cpu/op_upsample_bilinear2d.cpp

@@ -0,0 +1,139 @@
+/*
+ * 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 <executorch/kernels/portable/cpu/util/upsample_util.h>
+#include <executorch/runtime/kernel/kernel_includes.h>
+#include <algorithm>
+#include <array>
+
+namespace torch {
+namespace executor {
+namespace native {
+
+using exec_aten::ArrayRef;
+using exec_aten::optional;
+using exec_aten::SizesType;
+
+namespace {
+template <typename CTYPE>
+void upsample_bilinear2d_kernel_impl(
+    const Tensor& in,
+    bool align_corners,
+    const float scale_h,
+    const float scale_w,
+    Tensor& out) {
+  const auto in_data = in.const_data_ptr<CTYPE>();
+  auto out_data = out.mutable_data_ptr<CTYPE>();
+
+  auto in_plane = in_data;
+  for (auto n = 0; n < out.size(0); n++) {
+    for (auto c = 0; c < out.size(1); c++) {
+      for (auto h = 0; h < out.size(2); h++) {
+        for (auto w = 0; w < out.size(3); w++) {
+          // Compute source index.
+          // See area_pixel_compute_source_index in
+          // pytorch/aten/src/ATen/native/UpSample.h
+          float in_h =
+              area_pixel_compute_source_index(scale_h, h, align_corners, false);
+          float in_w =
+              area_pixel_compute_source_index(scale_w, w, align_corners, false);
+
+          int64_t in_h1, in_h2, in_w1, in_w2;
+          float weight_h, inv_weight_h, weight_w, inv_weight_w;
+
+          compute_source_index_and_lambda(
+              in_h1,
+              in_h2,
+              weight_h,
+              inv_weight_h,
+              scale_h,
+              h,
+              in.sizes()[2],
+              out.sizes()[2],
+              align_corners);
+
+          compute_source_index_and_lambda(
+              in_w1,
+              in_w2,
+              weight_w,
+              inv_weight_w,
+              scale_w,
+              w,
+              in.sizes()[3],
+              out.sizes()[3],
+              align_corners);
+
+          const auto top_left =
+              in_plane[in_h1 * in.strides()[2] + in_w1 * in.strides()[3]];
+          const auto top_right =
+              in_plane[in_h1 * in.strides()[2] + in_w2 * in.strides()[3]];
+          const auto bottom_left =
+              in_plane[in_h2 * in.strides()[2] + in_w1 * in.strides()[3]];
+          const auto bottom_right =
+              in_plane[in_h2 * in.strides()[2] + in_w2 * in.strides()[3]];
+
+          const auto top = top_left * weight_w + top_right * inv_weight_w;
+          const auto bottom =
+              bottom_left * weight_w + bottom_right * inv_weight_w;
+          const auto val = top * weight_h + bottom * inv_weight_h;
+
+          *out_data = val;
+          out_data++;
+        }
+      }
+
+      in_plane += in.strides()[1];
+    }
+  }
+}
+} // namespace
+
+Tensor& upsample_bilinear2d_out(
+    KernelRuntimeContext& ctx,
+    const Tensor& in,
+    const exec_aten::ArrayRef<int64_t> output_size,
+    bool align_corners,
+    const optional<double> scale_h,
+    const optional<double> scale_w,
+    Tensor& out) {
+  // Preconditions (checked in check_..._args):
+  //  In and out tensors have same dtype.
+  //  In and out tensors are rank 4 and have same dim[0] and dim[1].
+  //  In and out tensors are default dim order (NCHW).
+  ET_KERNEL_CHECK(
+      ctx,
+      check_upsample_bilinear2d_args(
+          in, output_size, align_corners, scale_h, scale_w, out),
+      InvalidArgument,
+      out);
+
+  ET_KERNEL_CHECK_MSG(
+      ctx,
+      resize_upsample_2d(in, output_size, scale_h, scale_w, out) == Error::Ok,
+      InvalidArgument,
+      out,
+      "Failed to resize output tensor");
+
+  const auto dim = in.dim();
+  const auto kernel_scale_h = area_pixel_compute_scale<double>(
+      in.sizes()[2], out.sizes()[2], align_corners, scale_h);
+  const auto kernel_scale_w = area_pixel_compute_scale<double>(
+      in.sizes()[3], out.sizes()[3], align_corners, scale_w);
+
+  ET_SWITCH_REAL_TYPES(
+      in.scalar_type(), ctx, "upsample_bilinear2d.out", CTYPE, [&]() {
+        upsample_bilinear2d_kernel_impl<CTYPE>(
+            in, align_corners, kernel_scale_h, kernel_scale_w, out);
+      });
+
+  return out;
+}
+
+} // namespace native
+} // namespace executor
+} // namespace torch

kernels/portable/cpu/util/targets.bzl

@@ -31,6 +31,7 @@ def define_common_targets():
             "//executorch/kernels/portable/cpu/util:advanced_index_util",
             "//executorch/kernels/portable/cpu/util:slice_util",
             "//executorch/kernels/portable/cpu/util:elementwise_util",
+            "//executorch/kernels/portable/cpu/util:upsample_util",
         ],
         visibility = ["//executorch/...", "@EXECUTORCH_CLIENTS"],
     )
@@ -266,6 +267,16 @@ def define_common_targets():
         visibility = ["//executorch/kernels/portable/cpu/..."],
     )
 
+    runtime.cxx_library(
+        name = "upsample_util",
+        srcs = ["upsample_util.cpp"],
+        exported_headers = ["upsample_util.h"],
+        deps = [
+            "//executorch/runtime/kernel:kernel_includes",
+        ],
+        visibility = ["//executorch/kernels/portable/cpu/..."],
+    )
+
     # Utility functions that can be used by operators that perform reduction
     for aten_mode in [True, False]:
         suffix = "_aten" if aten_mode else ""

kernels/portable/cpu/util/upsample_util.cpp

@@ -0,0 +1,95 @@
+/*
+ * 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 <executorch/kernels/portable/cpu/util/upsample_util.h>
+#include <executorch/runtime/core/exec_aten/util/tensor_util.h>
+
+namespace torch {
+namespace executor {
+
+bool check_upsample_2d_common_args(
+    const Tensor& in,
+    const exec_aten::ArrayRef<int64_t> output_size,
+    const exec_aten::optional<double> scale_h,
+    const exec_aten::optional<double> scale_w,
+    Tensor& out) {
+  ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, out));
+  ET_LOG_AND_RETURN_IF_FALSE(in.dim() == 4);
+  ET_LOG_AND_RETURN_IF_FALSE(out.dim() == 4);
+  ET_LOG_AND_RETURN_IF_FALSE(tensor_is_default_dim_order(in));
+  ET_LOG_AND_RETURN_IF_FALSE(tensor_is_default_dim_order(out));
+  ET_LOG_AND_RETURN_IF_FALSE(
+      (output_size.size() == 2 && !scale_h.has_value() &&
+       !scale_w.has_value()) ||
+      (output_size.size() == 0 && scale_h.has_value() && scale_w.has_value()));
+  ET_LOG_AND_RETURN_IF_FALSE(!scale_h.has_value() || scale_h.value() > 0);
+  ET_LOG_AND_RETURN_IF_FALSE(!scale_w.has_value() || scale_w.value() > 0);
+  ET_LOG_AND_RETURN_IF_FALSE(output_size.size() < 1 || output_size[0] > 0);
+  ET_LOG_AND_RETURN_IF_FALSE(output_size.size() < 2 || output_size[1] > 0);
+
+  return true;
+}
+
+bool check_upsample_bilinear2d_args(
+    const Tensor& in,
+    const exec_aten::ArrayRef<int64_t> output_size,
+    ET_UNUSED const bool align_corners,
+    const exec_aten::optional<double> scale_h,
+    const exec_aten::optional<double> scale_w,
+    Tensor& out) {
+  return check_upsample_2d_common_args(in, output_size, scale_h, scale_w, out);
+}
+
+bool check_upsample_nearest2d_args(
+    const Tensor& in,
+    const exec_aten::ArrayRef<int64_t> output_size,
+    const exec_aten::optional<double> scale_h,
+    const exec_aten::optional<double> scale_w,
+    Tensor& out) {
+  return check_upsample_2d_common_args(in, output_size, scale_h, scale_w, out);
+}
+
+Error resize_upsample_2d(
+    const Tensor& in,
+    const exec_aten::ArrayRef<int64_t> output_size,
+    const exec_aten::optional<double> scale_h,
+    const exec_aten::optional<double> scale_w,
+    Tensor& out) {
+  // Either output_size or scale_factors are provided, not both. This
+  // is checked in check_..._args.
+  // Scales are transformed according to align_corners.
+  std::array<Tensor::SizesType, kTensorDimensionLimit> target_size;
+
+  const auto dim = in.dim();
+  std::copy(in.sizes().cbegin(), in.sizes().cend(), target_size.begin());
+
+  if (scale_h.has_value() && scale_w.has_value()) {
+    target_size[dim - 2] =
+        static_cast<Tensor::SizesType>(in.sizes()[dim - 2] * scale_h.value());
+    target_size[dim - 1] =
+        static_cast<Tensor::SizesType>(in.sizes()[dim - 1] * scale_w.value());
+  } else if (output_size.size() == 2) {
+    target_size[dim - 2] = output_size[0];
+    target_size[dim - 1] = output_size[1];
+  } else {
+    ET_LOG(Error, "Invalid output_size or scale_factors");
+    return Error::InvalidArgument;
+  }
+
+  ET_CHECK_OR_RETURN_ERROR(
+      target_size[dim - 2] > 0 && target_size[dim - 1] > 0,
+      InvalidArgument,
+      "Upsampled output size must be non-empty, but was %ld x %ld.",
+      static_cast<long>(target_size[dim - 2]),
+      static_cast<long>(target_size[dim - 1]));
+
+  return resize_tensor(out, {target_size.data(), static_cast<size_t>(dim)});
+}
+
+} // namespace executor
+} // namespace torch