Optimize scale-factor-2 nearest-neighbor upsampling portable op (#22999)

Add portable NCHW and AArch64 FP16 NEON fast paths. Preserve the
portable kernel as the fallback for unsupported scales, dtypes, and
layouts.

Across five 20-iteration repeats, improved YOLO26 FP16 640x640 median
latency by 2.9%. Delegation and binary size unchanged.


cc @digantdesai @freddan80 @per @zingo @oscarandersson8218
@Sebastian-Larsson @robell @rascani

Signed-off-by: Måns Nilsson <mans.nilsson@arm.com>
Co-authored-by: Jacob Stevens <stevens.jacob1492@gmail.com>

Author: mansnils

kernels/optimized/CMakeLists.txt

@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -107,6 +108,25 @@ if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64" OR ANDROID_ABI STREQUAL
   )
 endif()
 
+# The nearest-2x fast path uses only baseline AArch64 Advanced SIMD. Keep
+# arm_neon.h and its implementation out of all other architecture builds.
+if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64" OR ANDROID_ABI STREQUAL
+                                                     "arm64-v8a"
+)
+  add_library(
+    upsample_nearest2d_neon_impl OBJECT
+    ${EXECUTORCH_ROOT}/kernels/optimized/cpu/op_upsample_nearest2d_neon.cpp
+  )
+  target_compile_options(
+    upsample_nearest2d_neon_impl PRIVATE ${_common_compile_options}
+  )
+  target_link_libraries(upsample_nearest2d_neon_impl PRIVATE executorch_core)
+  target_link_libraries(
+    optimized_kernels
+    PRIVATE $<BUILD_LOCAL_INTERFACE:upsample_nearest2d_neon_impl>
+  )
+endif()
+
 # Build a library for _optimized_kernels_srcs
 #
 # optimized_ops_lib: Register optimized ops kernels into Executorch runtime

kernels/optimized/cpu/op_upsample_nearest2d.cpp

@@ -0,0 +1,89 @@
+/*
+ * Copyright 2026 Arm Limited and/or its affiliates.
+ *
+ * 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 <executorch/runtime/kernel/thread_parallel_interface.h>
+
+#include <algorithm>
+
+#if defined(__aarch64__)
+#include <executorch/kernels/optimized/cpu/op_upsample_nearest2d_neon.h>
+#endif
+
+namespace torch {
+namespace executor {
+namespace native {
+
+using executorch::aten::ScalarType;
+using executorch::aten::Tensor;
+
+Tensor& upsample_nearest2d_vec_out(
+    KernelRuntimeContext& ctx,
+    const Tensor& in,
+    executorch::aten::OptionalArrayRef<int64_t> output_size,
+    executorch::aten::OptionalArrayRef<double> scale_factors,
+    Tensor& out);
+
+Tensor& opt_upsample_nearest2d_vec_out(
+    KernelRuntimeContext& ctx,
+    const Tensor& in,
+    const executorch::aten::OptionalArrayRef<int64_t> output_size,
+    const executorch::aten::OptionalArrayRef<double> scale_factors,
+    Tensor& out) {
+#if !defined(__aarch64__)
+  return upsample_nearest2d_vec_out(ctx, in, output_size, scale_factors, out);
+#else
+  if (in.scalar_type() != ScalarType::Half ||
+      !tensor_is_default_dim_order(in) || !tensor_is_default_dim_order(out) ||
+      !tensor_is_contiguous(in) || !tensor_is_contiguous(out) ||
+      !check_upsample_nearest2d_args(in, output_size, scale_factors, out)) {
+    return upsample_nearest2d_vec_out(ctx, in, output_size, scale_factors, out);
+  }
+
+  double scale_h;
+  double scale_w;
+  ET_KERNEL_CHECK_MSG(
+      ctx,
+      resize_upsample_2d(
+          in, output_size, scale_factors, scale_h, scale_w, out) == Error::Ok,
+      InvalidArgument,
+      out,
+      "Failed to resize output tensor");
+
+  const auto kernel_scale_h =
+      area_pixel_compute_scale<double>(in.size(2), out.size(2), false, scale_h);
+  const auto kernel_scale_w =
+      area_pixel_compute_scale<double>(in.size(3), out.size(3), false, scale_w);
+  if (!tensor_is_contiguous(out) || kernel_scale_h != 0.5 ||
+      kernel_scale_w != 0.5 || out.size(2) != in.size(2) * 2 ||
+      out.size(3) != in.size(3) * 2) {
+    return upsample_nearest2d_vec_out(ctx, in, output_size, scale_factors, out);
+  }
+
+  const int64_t output_plane_size = out.size(2) * out.size(3);
+  const int64_t grain_size = std::max<int64_t>(
+      1, ::executorch::extension::internal::GRAIN_SIZE / output_plane_size);
+  const int64_t planes = in.size(0) * in.size(1);
+  const auto* input = reinterpret_cast<const uint16_t*>(in.const_data_ptr());
+  auto* output = reinterpret_cast<uint16_t*>(out.mutable_data_ptr());
+
+  const bool success = ::executorch::extension::parallel_for(
+      0, planes, grain_size, [&](const auto begin, const auto end) {
+        opt_upsample_nearest2d_internal::upsample_nearest2d_2x_nchw_u16(
+            input, output, begin, end, in.size(2), in.size(3));
+      });
+  if (!success) {
+    return upsample_nearest2d_vec_out(ctx, in, output_size, scale_factors, out);
+  }
+  return out;
+#endif
+}
+
+} // namespace native
+} // namespace executor
+} // namespace torch

kernels/optimized/cpu/op_upsample_nearest2d_neon.cpp

@@ -0,0 +1,66 @@
+/*
+ * Copyright 2026 Arm Limited and/or its affiliates.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#if defined(__aarch64__)
+
+#include <executorch/kernels/optimized/cpu/op_upsample_nearest2d_neon.h>
+
+#include <arm_neon.h>
+
+#include <cstring>
+
+namespace torch {
+namespace executor {
+namespace native {
+namespace opt_upsample_nearest2d_internal {
+
+void upsample_nearest2d_2x_nchw_u16(
+    const uint16_t* input,
+    uint16_t* output,
+    const int64_t begin,
+    const int64_t end,
+    const int64_t input_height,
+    const int64_t input_width) {
+  const int64_t output_height = input_height * 2;
+  const int64_t output_width = input_width * 2;
+  const int64_t input_plane_size = input_height * input_width;
+  const int64_t output_plane_size = output_height * output_width;
+
+  for (int64_t plane = begin; plane < end; ++plane) {
+    const uint16_t* input_plane = input + plane * input_plane_size;
+    uint16_t* output_plane = output + plane * output_plane_size;
+
+    for (int64_t h = 0; h < input_height; ++h) {
+      const uint16_t* input_row = input_plane + h * input_width;
+      uint16_t* output_row = output_plane + (2 * h) * output_width;
+      int64_t w = 0;
+      for (; w + 8 <= input_width; w += 8) {
+        const uint16x8_t values = vld1q_u16(input_row + w);
+        uint16x8x2_t duplicated;
+        duplicated.val[0] = values;
+        duplicated.val[1] = values;
+        vst2q_u16(output_row + 2 * w, duplicated);
+      }
+      for (; w < input_width; ++w) {
+        const uint16_t value = input_row[w];
+        output_row[2 * w] = value;
+        output_row[2 * w + 1] = value;
+      }
+      std::memcpy(
+          output_row + output_width,
+          output_row,
+          output_width * sizeof(uint16_t));
+    }
+  }
+}
+
+} // namespace opt_upsample_nearest2d_internal
+} // namespace native
+} // namespace executor
+} // namespace torch
+
+#endif

kernels/optimized/cpu/op_upsample_nearest2d_neon.h

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2026 Arm Limited and/or its affiliates.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+#if defined(__aarch64__)
+
+namespace torch {
+namespace executor {
+namespace native {
+namespace opt_upsample_nearest2d_internal {
+
+void upsample_nearest2d_2x_nchw_u16(
+    const uint16_t* input,
+    uint16_t* output,
+    int64_t begin,
+    int64_t end,
+    int64_t input_height,
+    int64_t input_width);
+
+} // namespace opt_upsample_nearest2d_internal
+} // namespace native
+} // namespace executor
+} // namespace torch
+
+#endif

kernels/optimized/cpu/targets.bzl

@@ -109,6 +109,19 @@ def define_common_targets():
         ],
     )
 
+    runtime.cxx_library(
+        name = "upsample_nearest2d_neon_impl",
+        srcs = select({
+            "DEFAULT": [],
+            "ovr_config//cpu:arm64": ["op_upsample_nearest2d_neon.cpp"],
+        }),
+        exported_headers = ["op_upsample_nearest2d_neon.h"],
+        visibility = ["PUBLIC"],
+        exported_deps = [
+            "//executorch/runtime/kernel:kernel_includes",
+        ],
+    )
+
     # Used for dtype selective build. Collect source and header files.
     runtime.filegroup(
         name = "optimized_source_files",

kernels/optimized/optimized.yaml

@@ -1,4 +1,5 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This yaml file contains operators that have optimized kernels available.
 
@@ -121,3 +122,8 @@
   kernels:
     - arg_meta: null
       kernel_name: torch::executor::opt_where_out
+
+- op: upsample_nearest2d.vec_out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::opt_upsample_nearest2d_vec_out

kernels/portable/cpu/op_upsample_nearest2d.cpp

@@ -1,25 +1,95 @@
 /*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
  * All rights reserved.
+ * Copyright 2026 Arm Limited and/or its affiliates.
  *
  * 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 <executorch/runtime/kernel/thread_parallel_interface.h>
+
+#include <cstring>
 
 namespace torch {
 namespace executor {
 namespace native {
 
 namespace {
+
+template <typename CTYPE>
+bool upsample_nearest2d_2x_nchw_fast_path(
+    const Tensor& in,
+    const float scale_h,
+    const float scale_w,
+    Tensor& out) {
+  if (scale_h != 0.5f || scale_w != 0.5f || out.size(2) != in.size(2) * 2 ||
+      out.size(3) != in.size(3) * 2) {
+    return false;
+  }
+
+  const int64_t in_h = in.size(2);
+  const int64_t in_w = in.size(3);
+  const int64_t out_h = out.size(2);
+  const int64_t out_w = out.size(3);
+
+  const bool in_contiguous_nchw = in.strides()[3] == 1 &&
+      in.strides()[2] == in_w && in.strides()[1] == in_h * in_w &&
+      in.strides()[0] == in.size(1) * in_h * in_w;
+  const bool out_contiguous_nchw = out.strides()[3] == 1 &&
+      out.strides()[2] == out_w && out.strides()[1] == out_h * out_w &&
+      out.strides()[0] == out.size(1) * out_h * out_w;
+  if (!in_contiguous_nchw || !out_contiguous_nchw) {
+    return false;
+  }
+
+  const auto in_data = in.const_data_ptr<CTYPE>();
+  auto out_data = out.mutable_data_ptr<CTYPE>();
+  const int64_t planes = in.size(0) * in.size(1);
+  const int64_t in_plane_size = in_h * in_w;
+  const int64_t out_plane_size = out_h * out_w;
+  int64_t grain_size =
+      ::executorch::extension::internal::GRAIN_SIZE / out_plane_size;
+  if (grain_size < 1) {
+    grain_size = 1;
+  }
+
+  const bool success = ::executorch::extension::parallel_for(
+      0, planes, grain_size, [&](const auto begin, const auto end) {
+        for (const auto plane : c10::irange(begin, end)) {
+          const CTYPE* in_plane = in_data + plane * in_plane_size;
+          CTYPE* out_plane = out_data + plane * out_plane_size;
+
+          for (int64_t h = 0; h < in_h; ++h) {
+            const CTYPE* in_row = in_plane + h * in_w;
+            CTYPE* out_row0 = out_plane + (2 * h) * out_w;
+            CTYPE* out_row1 = out_row0 + out_w;
+
+            for (int64_t w = 0; w < in_w; ++w) {
+              const CTYPE value = in_row[w];
+              out_row0[2 * w] = value;
+              out_row0[2 * w + 1] = value;
+            }
+            std::memcpy(out_row1, out_row0, out_w * sizeof(CTYPE));
+          }
+        }
+      });
+
+  return success;
+}
+
 template <typename CTYPE>
 void upsample_nearest2d_kernel_impl_nchw(
     const Tensor& in,
     const float scale_h,
     const float scale_w,
     Tensor& out) {
+  if (upsample_nearest2d_2x_nchw_fast_path<CTYPE>(in, scale_h, scale_w, out)) {
+    return;
+  }
+
   const auto in_data = in.const_data_ptr<CTYPE>();
   auto out_data = out.mutable_data_ptr<CTYPE>();
 

kernels/test/CMakeLists.txt

@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -363,6 +364,7 @@ if(EXECUTORCH_BUILD_KERNELS_OPTIMIZED)
       "op_native_layer_norm_test.cpp"
       "op_neg_test.cpp"
       "op_sub_test.cpp"
+      "op_upsample_nearest2d_test.cpp"
       "op_where_test.cpp"
       "UnaryUfuncRealHBBF16ToFloatHBF16Test.cpp"
       ${CMAKE_CURRENT_BINARY_DIR}/include/optimized/executorch/kernels/test/supported_features.cpp