Fix out of bounds error. Attempt to fix buck2 issues

Author: jgibson2

kernels/portable/cpu/op_grid_sampler_2d.cpp

@@ -115,22 +115,31 @@ void grid_sample_2d_bilinear_kernel_impl_nchw(
                   se_weight;
             }
           } else {
-            // For border/reflection padding, coordinates are already clipped
+            // For border/reflection padding, clip corner indices to valid range
+            // Even though source coordinates are clipped, adding 1 can push corners out of bounds
+            const int64_t ix_nw_safe = clip_coordinates(ix_nw, inp_W);
+            const int64_t iy_nw_safe = clip_coordinates(iy_nw, inp_H);
+            const int64_t ix_ne_safe = clip_coordinates(ix_ne, inp_W);
+            const int64_t iy_ne_safe = clip_coordinates(iy_ne, inp_H);
+            const int64_t ix_sw_safe = clip_coordinates(ix_sw, inp_W);
+            const int64_t iy_sw_safe = clip_coordinates(iy_sw, inp_H);
+            const int64_t ix_se_safe = clip_coordinates(ix_se, inp_W);
+            const int64_t iy_se_safe = clip_coordinates(iy_se, inp_H);
             out_val = in_data
-                          [in_channel_offset + iy_nw * in.strides()[2] +
-                           ix_nw * in.strides()[3]] *
+                          [in_channel_offset + iy_nw_safe * in.strides()[2] +
+                           ix_nw_safe * in.strides()[3]] *
                     nw_weight +
                 in_data
-                        [in_channel_offset + iy_ne * in.strides()[2] +
-                         ix_ne * in.strides()[3]] *
+                        [in_channel_offset + iy_ne_safe * in.strides()[2] +
+                         ix_ne_safe * in.strides()[3]] *
                     ne_weight +
                 in_data
-                        [in_channel_offset + iy_sw * in.strides()[2] +
-                         ix_sw * in.strides()[3]] *
+                        [in_channel_offset + iy_sw_safe * in.strides()[2] +
+                         ix_sw_safe * in.strides()[3]] *
                     sw_weight +
                 in_data
-                        [in_channel_offset + iy_se * in.strides()[2] +
-                         ix_se * in.strides()[3]] *
+                        [in_channel_offset + iy_se_safe * in.strides()[2] +
+                         ix_se_safe * in.strides()[3]] *
                     se_weight;
           }
 

kernels/portable/cpu/util/targets.bzl

@@ -36,6 +36,7 @@ def define_common_targets():
             "//executorch/kernels/portable/cpu/util:elementwise_util",
             "//executorch/kernels/portable/cpu/util:upsample_util",
             "//executorch/kernels/portable/cpu/util:vectorized_math",
+            "//executorch/kernels/portable/cpu/util:grid_sampler_2d_util",
         ],
         visibility = ["//executorch/...", "@EXECUTORCH_CLIENTS"],
     )
@@ -342,6 +343,16 @@ def define_common_targets():
         ],
     )
 
+    runtime.cxx_library(
+        name = "grid_sampler_2d_util",
+        srcs = ["grid_sampler_2d_util.cpp"],
+        exported_headers = ["grid_sampler_2d_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 get_aten_mode_options():
         suffix = "_aten" if aten_mode else ""

kernels/portable/test/op_grid_sampler_2d_test.py

@@ -7,7 +7,6 @@
 # pyre-unsafe
 
 import itertools
-import os
 import unittest
 
 import torch

kernels/portable/test/test_grid_sampler_2d_executorch.py

@@ -12,7 +12,6 @@
 import itertools
 import sys
 import unittest
-from typing import Tuple
 
 import torch
 import torch.nn as nn
@@ -87,7 +86,7 @@ def run_executorch_test(
                 "forward"
             )
             if fwd_method is None:
-                self.fail(f"Failed to load forward method")
+                self.fail("Failed to load forward method")
             executorch_output = fwd_method.execute((input_tensor, grid))[0]
 
             # Compare results

shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl

@@ -628,6 +628,13 @@ ATEN_OPS = (
             "//executorch/runtime/core/exec_aten/util:tensor_util",
         ],
     ),
+    op_target(
+            name = "op_grid_sampler_2d",
+            deps = [
+                "//executorch/kernels/portable/cpu/util:grid_sampler_2d_util",
+                "//executorch/runtime/core/exec_aten/util:tensor_util",
+            ],
+        ),
     op_target(
         name = "op_gt",
         deps = [