[Kernels] Add validations to catch misformed edge cases (#23113)

## Summary

Harden four kernels against malformed shapes that previously caused
out-of-bounds writes or silent truncation:

- `cat.out`: reject concatenations whose summed dim overflows 32-bit
`SizesType` instead of truncating the output size.
- `unfold_copy.out`: reject rank-16 input, which needs 17 shape entries
for a 16-slot buffer (same guard the sibling `view_as_real` kernel
uses).
- `choose_qparams_per_token_asymmetric`: reject scalar input, which
indexed `output_sizes[dim - 1]` at `-1`.
- Optimized `grid_sampler_2d`: validate and resize `out` via the
portable util before the NEON fast path, which sized its writes from
`input`/`grid` alone.

Each fix carries a regression test, and the grid_sampler suite now also
runs against the optimized kernels. Commits are stacked one per fix in
the order above.

This PR was authored with Muse Code.

## Test Plan

```
cmake --build build --target portable_kernels_test quantized_kernels_test optimized_kernels_test
./build/kernels/test/portable_kernels_test
./build/kernels/test/quantized_kernels_test
./build/kernels/test/optimized_kernels_test
```

1554 / 75 / 1564 tests pass with no failures (arm64, exercising the NEON
fast path). Each new test was verified to fail with its fix reverted.
`lintrunner` is clean on all touched files.

Author: JakeStevens

backends/cadence/fusion_g3/operators/op_cat.cpp

@@ -48,8 +48,12 @@ Tensor& cat_out(
 
   Tensor::SizesType expected_out_size[kTensorDimensionLimit];
   size_t expected_out_dim = 0;
-  torch::executor::get_cat_out_target_size(
-      tensors, dim, expected_out_size, &expected_out_dim);
+  ET_KERNEL_CHECK(
+      ctx,
+      torch::executor::get_cat_out_target_size(
+          tensors, dim, expected_out_size, &expected_out_dim),
+      InvalidArgument,
+      out);
 
   ET_KERNEL_CHECK(
       ctx,

backends/cadence/hifi/operators/op_cat.cpp

@@ -40,7 +40,12 @@ Tensor& cat_out(
   Tensor::SizesType
       expected_out_size[executorch::runtime::kTensorDimensionLimit];
   size_t expected_out_dim = 0;
-  get_cat_out_target_size(tensors, dim, expected_out_size, &expected_out_dim);
+  ET_KERNEL_CHECK(
+      ctx,
+      get_cat_out_target_size(
+          tensors, dim, expected_out_size, &expected_out_dim),
+      InvalidArgument,
+      out);
 
   ET_KERNEL_CHECK(
       ctx,

kernels/optimized/cpu/op_grid_sampler_2d.cpp

@@ -28,6 +28,7 @@
 // targets, and any unsupported interpolation/padding/layout combination,
 // delegate to the portable kernel.
 
+#include <executorch/kernels/portable/cpu/util/grid_sampler_2d_util.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
 
 #ifdef __aarch64__
@@ -335,6 +336,15 @@ Tensor& opt_grid_sampler_2d_out(
     int64_t padding_mode,
     bool align_corners,
     Tensor& out) {
+  // Validate args and resize out exactly like the portable kernel. The NEON
+  // paths below size their writes from input/grid alone, so out must already
+  // match [N, C, H_out, W_out] before the fast path is reached.
+  ET_KERNEL_CHECK_MSG(
+      ctx,
+      check_grid_sampler_2d_args_and_resize_out(input, grid, out) == Error::Ok,
+      InvalidArgument,
+      out,
+      "Failed to validate arguments and resize output tensor");
   // The NEON paths index input/grid/out directly assuming a contiguous NCHW
   // default-dim-order layout — no use of .strides() or .dim_order(). Fall
   // back to portable for anything else.

kernels/portable/cpu/op_cat.cpp

@@ -34,7 +34,12 @@ Tensor& cat_out(
 
   Tensor::SizesType expected_out_size[kTensorDimensionLimit];
   size_t expected_out_dim = 0;
-  get_cat_out_target_size(tensors, dim, expected_out_size, &expected_out_dim);
+  ET_KERNEL_CHECK(
+      ctx,
+      get_cat_out_target_size(
+          tensors, dim, expected_out_size, &expected_out_dim),
+      InvalidArgument,
+      out);
 
   ET_KERNEL_CHECK(
       ctx,

kernels/portable/cpu/op_unfold_copy.cpp

@@ -25,6 +25,16 @@ Tensor& unfold_copy_out(
   if (dim < 0) {
     dim += nonzero_dim(self);
   }
+  // Unfold adds one output dim; the target-size buffer holds
+  // kTensorDimensionLimit entries.
+  ET_KERNEL_CHECK_MSG(
+      ctx,
+      static_cast<size_t>(self.dim()) < kTensorDimensionLimit,
+      InvalidArgument,
+      out,
+      "Output size buffer is too small. Expected at least %zd, got %zu",
+      self.dim() + 1,
+      kTensorDimensionLimit);
   // Calculate output size
   // @lint-ignore CLANGTIDY facebook-hte-CArray
   Tensor::SizesType expected_output_size[kTensorDimensionLimit];

kernels/portable/cpu/util/copy_ops_util.cpp

@@ -9,6 +9,7 @@
 
 #include <c10/util/irange.h>
 #include <cstring>
+#include <limits>
 
 #include <executorch/kernels/portable/cpu/util/copy_ops_util.h>
 #include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
@@ -127,7 +128,7 @@ bool check_cat_args(
   return true;
 }
 
-void get_cat_out_target_size(
+bool get_cat_out_target_size(
     executorch::aten::ArrayRef<Tensor> tensors,
     int64_t dim,
     executorch::aten::SizesType* out_sizes,
@@ -139,6 +140,13 @@ void get_cat_out_target_size(
   size_t cat_dim_size = 0;
   for (const auto i : c10::irange(tensors.size())) {
     if (tensors[i].numel() > 0) {
+      ET_CHECK_OR_RETURN_FALSE(
+          static_cast<size_t>(tensors[i].size(dim)) <=
+              static_cast<size_t>(
+                  std::numeric_limits<executorch::aten::SizesType>::max()) -
+                  cat_dim_size,
+          "Concatenated size along dim %" PRId64 " overflows SizesType",
+          dim);
       cat_dim_size += tensors[i].size(dim);
     }
     if (tensors[i].dim() != 1 || tensors[i].numel() != 0) {
@@ -152,9 +160,10 @@ void get_cat_out_target_size(
     if (static_cast<int64_t>(d) != dim) {
       out_sizes[d] = tensors[ref_i].size(d);
     } else {
-      out_sizes[d] = cat_dim_size;
+      out_sizes[d] = static_cast<executorch::aten::SizesType>(cat_dim_size);
     }
   }
+  return true;
 }
 bool check_expand_copy_args(
     const Tensor& input,

kernels/portable/cpu/util/copy_ops_util.h

@@ -106,7 +106,7 @@ bool check_cat_args(
     int64_t dim,
     Tensor& out);
 
-void get_cat_out_target_size(
+bool get_cat_out_target_size(
     executorch::aten::ArrayRef<Tensor> tensors,
     int64_t dim,
     executorch::aten::SizesType* out_sizes,

kernels/portable/cpu/util/targets.bzl

@@ -412,7 +412,7 @@ def define_common_targets():
         deps = [
             "//executorch/runtime/kernel:kernel_includes",
         ],
-        visibility = ["//executorch/kernels/portable/cpu/..."],
+        visibility = ["//executorch/kernels/portable/cpu/...", "//executorch/kernels/optimized/cpu/..."],
     )
 
     # Utility functions that can be used by operators that perform reduction

kernels/quantized/cpu/op_choose_qparams.cpp

@@ -281,6 +281,10 @@ std::tuple<Tensor&, Tensor&> choose_qparams_per_token_asymmetric_out(
     ScalarType dtype,
     Tensor& scale_out,
     Tensor& zero_point_out) {
+  ET_CHECK_MSG(
+      input.dim() >= 1,
+      "choose_qparams_per_token: input must have rank >= 1, got %zd",
+      input.dim());
   int64_t quant_min = -128;
   int64_t quant_max = 127;
   executorch::aten::SizesType output_sizes[kTensorDimensionLimit];

kernels/quantized/test/op_choose_qparams_test.cpp

@@ -260,3 +260,21 @@ TEST(
     EXPECT_EQ(zero_point_out.const_data_ptr<int64_t>()[i], nudged_zero_point);
   }
 }
+
+TEST(OpChooseQparamsPerTokenAsymmetricTensorOutTest, ScalarInputDies) {
+  et_pal_init();
+  TensorFactory<ScalarType::Float> tf_float;
+  TensorFactory<ScalarType::Double> tf_double;
+  TensorFactory<ScalarType::Long> tf_long;
+
+  // Per-token scales index input.dim() - 1; a scalar input must be rejected
+  // instead of indexing before the output-shape buffer.
+  Tensor input = tf_float.make({}, {1.0});
+  Tensor scale_out = tf_double.zeros({1});
+  Tensor zero_point_out = tf_long.zeros({1});
+
+  ET_EXPECT_DEATH(
+      choose_qparams_per_token_asymmetric_out(
+          input, ScalarType::Float, scale_out, zero_point_out),
+      "must have rank");
+}