optimized: add grid_sampler_2d.out (NEON) and sum.IntList_out (Vectorized<float>)

[jgibson2]

Summary

Two new optimized CPU kernels registered alongside the existing optimized_kernels library. Both replace the portable reference kernel (still available as fallback for unsupported inputs) with vectorized implementations that accumulate in fp32, which also sidesteps the fp16 precision issue noted in #19117 for grid_sampler_2d bilinear.

Measured end-to-end on a real depth model (Pixel 9 / arm64-v8a, fp16 inputs, shapes representative of the model's hot path):

Op	Portable	This PR	Speedup
grid_sampler_2d.out	17.3 ms	3.4 ms	5.1×
sum.IntList_out (5 calls, aggregate)	3.0 ms	0.56 ms	5.4×

grid_sampler_2d.out

aarch64 NEON, bilinear + zeros padding only (the dominant mode for depth / MVS / spatial transformer networks). Processes 4 channels per iteration with a vectorized FMA chain. fp16 inputs are promoted to fp32 for weight computation and accumulation, cast back on store — the portable kernel's fp16 weight subtractions like (ix_se - ix) otherwise suffer catastrophic cancellation (same concern as #19117). Unsupported modes and non-aarch64 targets delegate to the portable kernel.

sum.IntList_out

at::vec::Vectorized<float>-based implementation of the single-dim reduction fast path (both innermost-contiguous and strided cases). Cross-architecture SIMD via PyTorch's existing vector abstraction; always accumulates in fp32 regardless of input dtype. Multi-dim reductions, dtype-converting reductions, and complex types delegate to portable.

Integration

• Sources added to OPTIMIZED_KERNELS_SRCS in build_variables.bzl and to OPTIMIZED_ATEN_OPS in op_registration_util.bzl. Single source of truth for both Buck and CMake builds.
• optimized.yaml registers the ops with the standard opt_* naming convention used by sibling kernels.
• kernels/optimized/CMakeLists.txt scopes the -march=armv8.2-a+fp16 flag to just op_grid_sampler_2d.cpp via set_source_files_properties, so x86_64 builds are unaffected. The kernel has #ifdef __aarch64__ guards and falls through to portable on non-arm64 targets.

Test plan

• Builds cleanly for Android arm64-v8a, Android x86_64 (via scripts/build_android_library.sh), and host (macOS / Apple Clang 21).
• Existing kernels/test/op_grid_sampler_2d_test.cpp and op_sum_test.cpp unit tests continue to pass — both target the aten::sum_outf / aten::grid_sampler_2d_outf codegen-dispatched entry points, so they automatically exercise the optimized kernels when linked.
• Numerical verification against an fp32 reference (run portable in fp32, cast to fp16) on the shapes the polycam depth model uses — all cases pass within fp16 ULP.
• End-to-end Pixel 9 latency on a representative trained model matches the handwritten-NEON reference implementation to within run-to-run noise while producing more accurate fp16 outputs (fp32 accumulation).

Candidate successor to #19117 for the grid_sampler half — applies the same precision fix but at the optimized-kernel layer, so callers who link optimized_ops_lib get both the correctness fix and the speedup.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/19119

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Rolling out OSDC (ARC) runners on pull & trunk workflows in PyTorch main

❌ 3 New Failures, 2 Unrelated Failures

As of commit 4e98f5a with merge base de8ce55 ():

NEW FAILURES - The following jobs have failed:

• pull / unittest / linux / linux-job (gh)
  backends/xnnpack/test/ops/test_conv2d.py::TestConv2d::test_qs8_conv2d_relu_multi_users
• pull / unittest / macos / macos-job (gh)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 8
• pull / unittest-editable / macos / macos-job (gh)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 8

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / unittest / windows / windows-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / unittest-editable / windows / windows-job (gh) (trunk failure)
  ##[error]The operation was canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[jgibson2]

@pytorchbot label "release notes: none"

[GregoryComer]

@JacobSzwejbka @manuelcandales @digantdesai Do you have any concerns with conditionally linking an op in optimized only on some architectures?

[GregoryComer]

Thanks for the PR! I left one comment about runtime gating the armv8.2+fp16 code. Other than that, it looks good.

We are currently looking at adding better support for in-tree CPU operator implementations with arch-specific dispatch, so this type of thing should become easier soon.

[GregoryComer]

Would it be possible to split out the native f16 path? Right now, it will potentially SIGILL on ARM hardware without f16 support. If possible, I'd recommend something like this:

• Move the native f16 impl into a separate source file. Scope the march +fp16 to just this file.
• Add a variant that does the f16<->f32 conversion in software.
• In the top-level kernel, check hardware support using cpuinfo_has_arm_neon_fp16 and route to the implementation.

[jgibson2]

Addressed the SIGILL concern — new commit 53697e9 splits the fp16 path into HW and SW variants with runtime dispatch:

• op_grid_sampler_2d_fp16_hw.cpp (new) — hardware fp16 fast path. Uses vld1_f16 / vcvt_f32_f16 / vcvt_f16_f32 / vst1_f16. Compiled with -march=armv8.2-a+fp16 (scoped via set_source_files_properties in CMake; compiler_flags gated on ovr_config//cpu:arm64 in the Buck target).
• op_grid_sampler_2d.cpp — now hosts the fp32 path, a new fp16 software-convert path, and the runtime dispatcher. Plain ARMv8 only. The SW path converts fp16↔fp32 via c10::Half's portable conversions (no hardware fp16 instructions) and does all compute on NEON fp32 lanes — slower per conversion but safe on any ARMv8 chip.

Dispatcher is at op_grid_sampler_2d.cpp:394-425. Key lines:

if (input.scalar_type() == ScalarType::Half) {
  if (cpuinfo_initialize() && cpuinfo_has_arm_neon_fp16()) {
    // HW path (in _fp16_hw.cpp)
  }
  // fall through to SW path (in this file)
}

cpuinfo is already transitively linked via extension_threadpool, so the dependency was a one-line Buck addition. No change to behavior on fp16-capable chips (Pixel 9, S24 FE, etc.); chips without the +fp16 extension now run the SW variant instead of raising SIGILL.

One thing I'd appreciate guidance on: the fp16 SW and HW paths share ~60 lines of loop body verbatim — same FMA chain, same indexing, same weight math. I kept them copy-pasted for clarity rather than macro/template gymnastics. Happy to DRY that up if you'd prefer.

[meta-codesync]

@GregoryComer has imported this pull request. If you are a Meta employee, you can view this in D102420839.