optimized/grid_sampler_2d: address review feedback

Three changes consolidated for review:

1. Move the forward declaration of grid_sampler_2d_bilinear_fp16_hw out
   of op_grid_sampler_2d.cpp into a new header
   kernels/optimized/cpu/op_grid_sampler_2d_fp16_hw.h. The function has
   external linkage (the dispatcher in op_grid_sampler_2d.cpp calls into
   it across translation units), and prior to this its definition site
   had no prior prototype visible — which trips -Wmissing-prototypes on
   build configurations that enable it. Both .cpp files now include the
   shared header. The function body stays in op_grid_sampler_2d_fp16_hw.cpp
   because that TU is the only one compiled with -march=armv8.2-a+fp16,
   so it cannot be inlined into a header. The header itself uses void* for
   input/output buffers and is fp16-free, so callers don't need the
   +fp16 march flag just to declare or call it.

2. Split the fp16 HW path into its own CMake target. Previously the
   -march=armv8.2-a+fp16 flag was scoped per-source-file via
   set_source_files_properties on the sole TU inside the optimized_kernels
   library. That works for a clean non-LTO build, but with ThinLTO or
   cross-TU optimizations the flag boundary becomes fuzzy and the fallback
   path in op_grid_sampler_2d.cpp could in principle be auto-vectorized
   into fp16 NEON instructions — exactly the SIGILL hazard the runtime
   dispatch is meant to prevent. Build the file as an OBJECT library
   (grid_sampler_2d_fp16_hw_impl) with target-scoped -march flag and link
   it into optimized_kernels via $<BUILD_LOCAL_INTERFACE:...> so the
   object code is baked into liboptimized_kernels.a at archive time and
   the OBJECT target is kept out of the install EXPORT set. Mirrors the
   existing buck `grid_sampler_2d_fp16_hw_impl` cxx_library.

3. Gate the optimized fast paths on input/grid/out dtype match. Each fast
   path assumes a single dtype across all three tensors:

     fp32 NEON path: data_ptr<float>() on all three
     fp16 HW path:   void* pointers reinterpret_cast<__fp16*> on all three
     fp16 SW NEON:   data_ptr<c10::Half>() on all three

   Until now the dispatcher gated only on input.scalar_type(). The
   reinterpret_casts in the fp16 HW kernel are particularly load-bearing
   because their behavior on a mismatched dtype would be silent
   corruption (reading int64/double bytes as __fp16 stride). The
   data_ptr<T>() runtime check exists but is not guaranteed in release
   builds. Add a dtypes_match clause at the top of the fast-path
   eligibility check that requires all three scalar types equal; fall
   back to the portable kernel otherwise. The portable kernel handles
   arbitrary dtype combinations correctly.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: jgibson2

kernels/optimized/CMakeLists.txt

@@ -77,17 +77,34 @@ target_link_libraries(
 target_compile_options(optimized_kernels PUBLIC ${_common_compile_options})
 
 # op_grid_sampler_2d_fp16_hw.cpp uses hardware fp16 NEON intrinsics
-# (vcvt_f32_f16 / vld1_f16). Those are part of the ARMv8.2-a+fp16 extension and
-# raise SIGILL on chips without it. Scope the `-march` flag to just that
-# translation unit. The main op_grid_sampler_2d.cpp (which hosts the runtime
-# dispatcher via cpuinfo_has_arm_neon_fp16) and the fp16 software-convert path
-# stay on plain ARMv8 so they can run on any chip.
+# (vcvt_f32_f16 / vld1_f16). Those are part of the ARMv8.2-a+fp16 extension
+# and raise SIGILL on chips without it. Build it as a separate OBJECT
+# library so the `-march=armv8.2-a+fp16` flag stays strictly scoped to that
+# translation unit and never reaches the dispatcher / fallback code in
+# op_grid_sampler_2d.cpp (which would otherwise risk auto-vectorizing into
+# fp16 NEON instructions). The dispatcher chooses between this entry point
+# and the fp16 software-convert path at runtime via
+# cpuinfo_has_arm_neon_fp16(). Mirrors the buck `grid_sampler_2d_fp16_hw_impl`
+# library.
 if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64" OR ANDROID_ABI STREQUAL
                                                      "arm64-v8a"
 )
-  set_source_files_properties(
+  add_library(
+    grid_sampler_2d_fp16_hw_impl OBJECT
     ${EXECUTORCH_ROOT}/kernels/optimized/cpu/op_grid_sampler_2d_fp16_hw.cpp
-    PROPERTIES COMPILE_OPTIONS "-march=armv8.2-a+fp16"
+  )
+  target_compile_options(
+    grid_sampler_2d_fp16_hw_impl
+    PRIVATE -march=armv8.2-a+fp16 ${_common_compile_options}
+  )
+  target_link_libraries(grid_sampler_2d_fp16_hw_impl PRIVATE executorch_core)
+  # BUILD_LOCAL_INTERFACE: object files are baked into optimized_kernels.a
+  # at archive time, so this OBJECT target stays out of the install EXPORT
+  # set and downstream consumers of the installed optimized_kernels need
+  # no separate link entry.
+  target_link_libraries(
+    optimized_kernels
+    PRIVATE $<BUILD_LOCAL_INTERFACE:grid_sampler_2d_fp16_hw_impl>
   )
 endif()
 

kernels/optimized/cpu/op_grid_sampler_2d.cpp

@@ -33,6 +33,7 @@
 #ifdef __aarch64__
 #include <arm_neon.h>
 #include <cpuinfo.h>
+#include <executorch/kernels/optimized/cpu/op_grid_sampler_2d_fp16_hw.h>
 #endif
 
 #include <c10/util/Half.h>
@@ -56,25 +57,6 @@ Tensor& grid_sampler_2d_out(
     bool align_corners,
     Tensor& out);
 
-#ifdef __aarch64__
-namespace opt_grid_sampler_2d_internal {
-// Declared in op_grid_sampler_2d_fp16_hw.cpp, compiled separately with
-// `-march=armv8.2-a+fp16`. Only safe to call when
-// cpuinfo_has_arm_neon_fp16() is true.
-void grid_sampler_2d_bilinear_fp16_hw(
-    const void* input,
-    const void* grid,
-    void* output,
-    int N,
-    int C,
-    int H_in,
-    int W_in,
-    int H_out,
-    int W_out,
-    bool align_corners);
-} // namespace opt_grid_sampler_2d_internal
-#endif
-
 #ifdef __aarch64__
 namespace {
 
@@ -361,7 +343,18 @@ Tensor& opt_grid_sampler_2d_out(
       tensor_is_contiguous(input) && tensor_is_contiguous(grid) &&
       tensor_is_contiguous(out);
 
-  if (interpolation_mode != 0 || padding_mode != 0 || !fast_eligible) {
+  // The fast paths read input/grid and write out as a single dtype: float for
+  // the fp32 NEON path, fp16 for both the fp16 HW path (which raw-casts the
+  // void* pointers to __fp16*) and the SW fp16 NEON path (which uses
+  // data_ptr<c10::Half>(), whose runtime dtype check is not guaranteed in
+  // release builds). Reject any mixed-dtype call up front so none of those
+  // unchecked casts can be reached with a mismatched buffer.
+  const bool dtypes_match =
+      input.scalar_type() == grid.scalar_type() &&
+      input.scalar_type() == out.scalar_type();
+
+  if (interpolation_mode != 0 || padding_mode != 0 || !fast_eligible ||
+      !dtypes_match) {
     return grid_sampler_2d_out(
         ctx, input, grid, interpolation_mode, padding_mode, align_corners, out);
   }

kernels/optimized/cpu/op_grid_sampler_2d_fp16_hw.cpp

@@ -22,6 +22,8 @@
 
 #ifdef __aarch64__
 
+#include <executorch/kernels/optimized/cpu/op_grid_sampler_2d_fp16_hw.h>
+
 #include <arm_neon.h>
 #include <cmath>
 

kernels/optimized/cpu/op_grid_sampler_2d_fp16_hw.h

@@ -0,0 +1,44 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#ifdef __aarch64__
+
+namespace torch {
+namespace executor {
+namespace native {
+namespace opt_grid_sampler_2d_internal {
+
+// Hardware-fp16 NEON bilinear + zeros-padding fast path. Defined in
+// op_grid_sampler_2d_fp16_hw.cpp, which is the only translation unit
+// compiled with `-march=armv8.2-a+fp16`. Only safe to call when
+// cpuinfo_has_arm_neon_fp16() reports true — see the runtime dispatcher
+// in op_grid_sampler_2d.cpp.
+//
+// Input/output buffers are passed as void* (raw uint16_t storage
+// interpreted as __fp16) so this header doesn't need <arm_neon.h> and
+// callers don't need the +fp16 march flag just to declare it.
+void grid_sampler_2d_bilinear_fp16_hw(
+    const void* input,
+    const void* grid,
+    void* output,
+    int N,
+    int C,
+    int H_in,
+    int W_in,
+    int H_out,
+    int W_out,
+    bool align_corners);
+
+} // namespace opt_grid_sampler_2d_internal
+} // namespace native
+} // namespace executor
+} // namespace torch
+
+#endif // __aarch64__

shim_et/xplat/executorch/build/build_variables.bzl

@@ -268,7 +268,6 @@ OPTIMIZED_KERNELS_SRCS = [
     "kernels/optimized/cpu/op_fft_r2c.cpp",
     "kernels/optimized/cpu/op_gelu.cpp",
     "kernels/optimized/cpu/op_grid_sampler_2d.cpp",
-    "kernels/optimized/cpu/op_grid_sampler_2d_fp16_hw.cpp",
     "kernels/optimized/cpu/op_le.cpp",
     "kernels/optimized/cpu/op_linear.cpp",
     "kernels/optimized/cpu/op_log_softmax.cpp",