kernels/custom: add NEON sum.IntList_out for arm64

Drops native_call_sum.IntList_out cumulative time by ~65% on Pixel 9 for
the polycam depth model (5 calls, aggregate 3.0 ms -> 1.1 ms).

The portable kernel's fast path is a scalar inner loop that the compiler
may auto-vectorize but doesn't guarantee. This kernel uses explicit NEON
intrinsics for the two shapes the model actually hits — innermost-dim
reduction and middle-dim (strided) reduction — and falls back to the
portable-style parallel_for + MapReduce plan for everything else
(multi-dim reductions, dtype conversion, complex types).

Implementation details (kernels/custom/neon_sum.cpp):

 * sum_innermost: one reduction per outer row, NEON float32x4 / float16x8
   loads with horizontal reduction at the end. fp16 inputs accumulate in
   fp32 for precision (long reductions in pure fp16 lose meaningful
   trailing digits).
 * sum_strided: non-innermost reduction, vectorized across the contiguous
   inner axis rather than the reduction axis. 4 fp32 or 8 fp16 inner
   positions per iteration; same fp32 accumulator rule for fp16.
 * Scalar fallback mirrors both shapes with f64 accumulators for non-arm64
   hosts and for dtypes we don't special-case.

Integration follows the grid_sampler_2d pattern from the prior commit:
new source file added to custom_kernels, new optimized.yaml entry
pointing at custom::sum_dim_out.

Author: jgibson2

kernels/custom/neon_sum.cpp

@@ -0,0 +1,346 @@
+// NEON-optimized sum.IntList_out for ExecuTorch.
+//
+// The portable kernel hands off to a generic parallel_for + MapReduce plan,
+// which the compiler may auto-vectorize but doesn't guarantee. This kernel
+// picks the common single-dim cases and does them with explicit NEON
+// intrinsics, falling back to a scalar loop for anything else (multi-dim
+// reductions, dtype mismatch, complex types).
+//
+// Signature matches kernels/portable/cpu/op_sum.cpp::sum_dim_out.
+
+#include <executorch/runtime/kernel/kernel_includes.h>
+#include <executorch/kernels/portable/cpu/util/reduce_util.h>
+
+#ifdef __aarch64__
+#include <arm_neon.h>
+#endif
+
+#include <cmath>
+#include <cstring>
+#include <optional>
+
+namespace custom {
+namespace native {
+
+using exec_aten::ArrayRef;
+using exec_aten::ScalarType;
+using exec_aten::Tensor;
+using executorch::aten::optional;
+using torch::executor::KernelRuntimeContext;
+
+namespace {
+
+// Normalize a possibly-negative dim into [0, ndim).
+inline int64_t normalize_dim(int64_t dim, int64_t ndim) {
+  return dim < 0 ? dim + ndim : dim;
+}
+
+// Compute (outer_size, reduce_size, inner_size) around a single reduction dim.
+inline void compute_partition_sizes(
+    const Tensor& in,
+    int64_t dim,
+    int64_t& outer_size,
+    int64_t& reduce_size,
+    int64_t& inner_size) {
+  outer_size = 1;
+  reduce_size = in.size(dim);
+  inner_size = 1;
+  for (int64_t d = 0; d < dim; ++d) {
+    outer_size *= in.size(d);
+  }
+  for (int64_t d = dim + 1; d < in.dim(); ++d) {
+    inner_size *= in.size(d);
+  }
+}
+
+#ifdef __aarch64__
+
+// Sum along the innermost (contiguous) dim: one reduction per outer row.
+//   in   : shape [outer_size, reduce_size]
+//   out  : shape [outer_size]
+template <typename T>
+inline void sum_innermost(
+    const T* in, T* out, int64_t outer_size, int64_t reduce_size);
+
+template <>
+inline void sum_innermost<float>(
+    const float* in, float* out, int64_t outer_size, int64_t reduce_size) {
+  for (int64_t i = 0; i < outer_size; ++i) {
+    const float* row = in + i * reduce_size;
+    float32x4_t acc = vdupq_n_f32(0.0f);
+    int64_t j = 0;
+    for (; j + 15 < reduce_size; j += 16) {
+      acc = vaddq_f32(acc, vld1q_f32(row + j));
+      acc = vaddq_f32(acc, vld1q_f32(row + j + 4));
+      acc = vaddq_f32(acc, vld1q_f32(row + j + 8));
+      acc = vaddq_f32(acc, vld1q_f32(row + j + 12));
+    }
+    for (; j + 3 < reduce_size; j += 4) {
+      acc = vaddq_f32(acc, vld1q_f32(row + j));
+    }
+    float sum = vaddvq_f32(acc);
+    for (; j < reduce_size; ++j) {
+      sum += row[j];
+    }
+    out[i] = sum;
+  }
+}
+
+template <>
+inline void sum_innermost<__fp16>(
+    const __fp16* in, __fp16* out, int64_t outer_size, int64_t reduce_size) {
+  // Accumulate in fp32 for precision; long reductions in pure fp16 lose
+  // significant trailing digits.
+  for (int64_t i = 0; i < outer_size; ++i) {
+    const __fp16* row = in + i * reduce_size;
+    float32x4_t acc = vdupq_n_f32(0.0f);
+    int64_t j = 0;
+    for (; j + 7 < reduce_size; j += 8) {
+      float16x8_t v = vld1q_f16(row + j);
+      acc = vaddq_f32(acc, vcvt_f32_f16(vget_low_f16(v)));
+      acc = vaddq_f32(acc, vcvt_f32_f16(vget_high_f16(v)));
+    }
+    float sum = vaddvq_f32(acc);
+    for (; j < reduce_size; ++j) {
+      sum += static_cast<float>(row[j]);
+    }
+    out[i] = static_cast<__fp16>(sum);
+  }
+}
+
+// Sum along a non-innermost dim: reduce over `reduce_size` elements spaced
+// `inner_size` apart, for each (outer, inner) pair.
+//   in   : shape [outer_size, reduce_size, inner_size]
+//   out  : shape [outer_size, inner_size]
+// Vectorized over the contiguous `inner` axis.
+template <typename T>
+inline void sum_strided(
+    const T* in,
+    T* out,
+    int64_t outer_size,
+    int64_t reduce_size,
+    int64_t inner_size);
+
+template <>
+inline void sum_strided<float>(
+    const float* in,
+    float* out,
+    int64_t outer_size,
+    int64_t reduce_size,
+    int64_t inner_size) {
+  const int64_t reduce_stride = inner_size;
+  const int64_t outer_stride = reduce_size * inner_size;
+  for (int64_t o = 0; o < outer_size; ++o) {
+    const float* in_o = in + o * outer_stride;
+    float* out_o = out + o * inner_size;
+    int64_t j = 0;
+    for (; j + 3 < inner_size; j += 4) {
+      float32x4_t acc = vdupq_n_f32(0.0f);
+      for (int64_t k = 0; k < reduce_size; ++k) {
+        acc = vaddq_f32(acc, vld1q_f32(in_o + k * reduce_stride + j));
+      }
+      vst1q_f32(out_o + j, acc);
+    }
+    for (; j < inner_size; ++j) {
+      float sum = 0.0f;
+      for (int64_t k = 0; k < reduce_size; ++k) {
+        sum += in_o[k * reduce_stride + j];
+      }
+      out_o[j] = sum;
+    }
+  }
+}
+
+template <>
+inline void sum_strided<__fp16>(
+    const __fp16* in,
+    __fp16* out,
+    int64_t outer_size,
+    int64_t reduce_size,
+    int64_t inner_size) {
+  const int64_t reduce_stride = inner_size;
+  const int64_t outer_stride = reduce_size * inner_size;
+  for (int64_t o = 0; o < outer_size; ++o) {
+    const __fp16* in_o = in + o * outer_stride;
+    __fp16* out_o = out + o * inner_size;
+    int64_t j = 0;
+    // Accumulate in fp32 for precision, 8 contiguous inner positions per iter.
+    for (; j + 7 < inner_size; j += 8) {
+      float32x4_t acc_lo = vdupq_n_f32(0.0f);
+      float32x4_t acc_hi = vdupq_n_f32(0.0f);
+      for (int64_t k = 0; k < reduce_size; ++k) {
+        float16x8_t v = vld1q_f16(in_o + k * reduce_stride + j);
+        acc_lo = vaddq_f32(acc_lo, vcvt_f32_f16(vget_low_f16(v)));
+        acc_hi = vaddq_f32(acc_hi, vcvt_f32_f16(vget_high_f16(v)));
+      }
+      vst1q_f16(
+          out_o + j,
+          vcombine_f16(vcvt_f16_f32(acc_lo), vcvt_f16_f32(acc_hi)));
+    }
+    for (; j < inner_size; ++j) {
+      float sum = 0.0f;
+      for (int64_t k = 0; k < reduce_size; ++k) {
+        sum += static_cast<float>(in_o[k * reduce_stride + j]);
+      }
+      out_o[j] = static_cast<__fp16>(sum);
+    }
+  }
+}
+
+#endif // __aarch64__
+
+// Scalar fallback: equivalent semantics to the NEON paths, used on non-arm64
+// hosts and for dtypes we don't special-case.
+template <typename T_IN, typename T_OUT>
+inline void sum_scalar(
+    const T_IN* in,
+    T_OUT* out,
+    int64_t outer_size,
+    int64_t reduce_size,
+    int64_t inner_size) {
+  const int64_t reduce_stride = inner_size;
+  const int64_t outer_stride = reduce_size * inner_size;
+  for (int64_t o = 0; o < outer_size; ++o) {
+    const T_IN* in_o = in + o * outer_stride;
+    T_OUT* out_o = out + o * inner_size;
+    for (int64_t j = 0; j < inner_size; ++j) {
+      double sum = 0.0;
+      for (int64_t k = 0; k < reduce_size; ++k) {
+        sum += static_cast<double>(in_o[k * reduce_stride + j]);
+      }
+      out_o[j] = static_cast<T_OUT>(sum);
+    }
+  }
+}
+
+} // namespace
+
+Tensor& sum_dim_out(
+    KernelRuntimeContext& ctx,
+    const Tensor& in,
+    optional<ArrayRef<int64_t>> dim_list,
+    bool keepdim,
+    optional<ScalarType> dtype,
+    Tensor& out) {
+  ET_KERNEL_CHECK(
+      ctx,
+      torch::executor::check_reduction_args(in, dim_list, keepdim, dtype, out),
+      InvalidArgument,
+      out);
+
+  ET_KERNEL_CHECK(
+      ctx,
+      torch::executor::resize_reduction_out(in, dim_list, keepdim, out) ==
+          torch::executor::Error::Ok,
+      InvalidArgument,
+      out);
+
+  ET_KERNEL_CHECK(
+      ctx,
+      torch::executor::tensors_have_same_dim_order(in, out),
+      InvalidArgument,
+      out);
+
+  ET_KERNEL_CHECK(
+      ctx,
+      torch::executor::tensor_is_default_dim_order(in),
+      InvalidArgument,
+      out);
+
+  if (in.numel() == 0) {
+    if (out.numel() > 0) {
+      std::memset(out.mutable_data_ptr(), 0, out.nbytes());
+    }
+    return out;
+  }
+
+  // We only fast-path the common case: single reduction dim, matching dtype,
+  // non-complex type, contiguous input. Everything else falls through to the
+  // scalar path which has identical semantics but no vectorization.
+  const bool fast_eligible = dim_list.has_value() &&
+      dim_list.value().size() == 1 &&
+      in.scalar_type() == out.scalar_type() &&
+      !executorch::runtime::isComplexType(in.scalar_type()) &&
+      torch::executor::tensor_is_contiguous(in);
+
+  if (fast_eligible) {
+    const int64_t dim = normalize_dim(dim_list.value()[0], in.dim());
+    int64_t outer_size, reduce_size, inner_size;
+    compute_partition_sizes(in, dim, outer_size, reduce_size, inner_size);
+
+#ifdef __aarch64__
+    if (in.scalar_type() == ScalarType::Float) {
+      const float* ip = in.const_data_ptr<float>();
+      float* op = out.mutable_data_ptr<float>();
+      if (inner_size == 1) {
+        sum_innermost<float>(ip, op, outer_size, reduce_size);
+      } else {
+        sum_strided<float>(ip, op, outer_size, reduce_size, inner_size);
+      }
+      return out;
+    } else if (in.scalar_type() == ScalarType::Half) {
+      static_assert(sizeof(__fp16) == 2, "expected __fp16 == 2 bytes");
+      const __fp16* ip =
+          reinterpret_cast<const __fp16*>(in.const_data_ptr<uint16_t>());
+      __fp16* op =
+          reinterpret_cast<__fp16*>(out.mutable_data_ptr<uint16_t>());
+      if (inner_size == 1) {
+        sum_innermost<__fp16>(ip, op, outer_size, reduce_size);
+      } else {
+        sum_strided<__fp16>(ip, op, outer_size, reduce_size, inner_size);
+      }
+      return out;
+    }
+#endif
+
+    // Scalar fast path for dtypes we don't have NEON for, or non-arm64 hosts.
+    // @lint-ignore CLANGTIDY facebook-hte-CArray
+    static constexpr const char op_name[] = "sum.IntList_out";
+    ET_SWITCH_REALHBBF16_TYPES(in.scalar_type(), ctx, op_name, CTYPE, [&] {
+      const CTYPE* ip = in.const_data_ptr<CTYPE>();
+      CTYPE* op = out.mutable_data_ptr<CTYPE>();
+      sum_scalar<CTYPE, CTYPE>(ip, op, outer_size, reduce_size, inner_size);
+    });
+    return out;
+  }
+
+  // Generic fallback: multi-dim reduction, dtype conversion, or any other
+  // case the fast path rejects. Semantically matches the portable kernel.
+  // @lint-ignore CLANGTIDY facebook-hte-CArray
+  static constexpr const char op_name[] = "sum.IntList_out";
+  std::optional<torch::executor::MapReduceOverDimListPlan> plan;
+  plan.emplace(in, dim_list);
+  ET_SWITCH_REALHBBF16_TYPES(in.scalar_type(), ctx, op_name, CTYPE_IN, [&] {
+    ET_SWITCH_REALHBBF16_TYPES(
+        out.scalar_type(), ctx, op_name, CTYPE_OUT, [&] {
+          CTYPE_OUT* out_data = out.mutable_data_ptr<CTYPE_OUT>();
+          const bool success =
+              torch::executor::
+                  parallel_for_each_reduce_over_dim_list_output_index(
+                      in,
+                      dim_list,
+                      out,
+                      [&](const auto begin, const auto end) {
+                        for (const auto out_ix :
+                             c10::irange(begin, end)) {
+                          CTYPE_OUT sum = 0;
+                          sum = plan->execute<CTYPE_IN, CTYPE_OUT>(
+                              [](CTYPE_IN v) {
+                                return static_cast<CTYPE_OUT>(v);
+                              },
+                              [](CTYPE_OUT outv, CTYPE_OUT acc) {
+                                return acc + outv;
+                              },
+                              out_ix);
+                          out_data[out_ix] = sum;
+                        }
+                      });
+          ET_KERNEL_CHECK_MSG(
+              ctx, success, Internal, , "parallel_for failed");
+        });
+  });
+  return out;
+}
+
+} // namespace native
+} // namespace custom

kernels/optimized/CMakeLists.txt

@@ -81,6 +81,7 @@ target_compile_options(optimized_kernels PUBLIC ${_common_compile_options})
 add_library(
   custom_kernels STATIC
   ${EXECUTORCH_ROOT}/kernels/custom/neon_grid_sampler_2d.cpp
+  ${EXECUTORCH_ROOT}/kernels/custom/neon_sum.cpp
 )
 target_link_libraries(custom_kernels PUBLIC executorch_core)
 target_compile_options(custom_kernels PUBLIC ${_common_compile_options})

kernels/optimized/optimized.yaml

@@ -102,6 +102,11 @@
     - arg_meta: null
       kernel_name: torch::executor::opt_sub_out
 
+- op: sum.IntList_out
+  kernels:
+    - arg_meta: null
+      kernel_name: custom::sum_dim_out
+
 - op: sub.Scalar_out
   kernels:
     - arg_meta: null