Prevent signed integer overflow in pixel_shuffle size calculation (#16138)

Summary:

The size calculation in `get_pixel_shuffle_out_target_size` can trigger signed integer overflow (UB) with very large upscale_factors. This manifested as an integer division by zero fault during fuzzing.

Specifically, the calculation of `upscale_factor * upscale_factor` ([source](https://github.com/pytorch/executorch/blob/04f1e4d22383ffcbc770acf5002348e3f95082a2/kernels/portable/cpu/util/copy_ops_util.cpp#L364)) can overflow. In the motivating case, SizesType is a signed 32-bit integer and upscale_factor = 2^17.

Since this is an impractically large upscale factor, I'm just adding a constraint that upscale_factor < 32768 (2^15). In theory, SizesType could be defined as less than 32 bits, but this seems unlikely in practice.

Differential Revision: D88693324

Author: GregoryComer

kernels/portable/cpu/op_pixel_shuffle.cpp

@@ -84,8 +84,11 @@ Tensor& pixel_shuffle_out(
   ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out);
   Tensor::SizesType expected_out_size[kTensorDimensionLimit];
   size_t expected_out_dim = 0;
-  get_pixel_shuffle_out_target_size(
-      in, upscale_factor, expected_out_size, &expected_out_dim);
+  ET_KERNEL_CHECK(ctx,
+      get_pixel_shuffle_out_target_size(
+        in, upscale_factor, expected_out_size, &expected_out_dim),
+      InvalidArgument,
+      out);
 
   // Make sure the output tensor is the right size.
   ET_KERNEL_CHECK(

kernels/portable/cpu/util/copy_ops_util.cpp

@@ -346,11 +346,14 @@ bool check_pixel_unshuffle_args(
   return true;
 }
 
-void get_pixel_shuffle_out_target_size(
+bool get_pixel_shuffle_out_target_size(
     const Tensor& in,
     int64_t upscale_factor,
     executorch::aten::SizesType* out_sizes,
     size_t* out_ndim) {
+  // Prevent signed integer overflow when computing upscale_factor ^ 2.
+  ET_CHECK_OR_RETURN_FALSE(upscale_factor < 32768, "Upscale factor must be less than 32768.");
+
   *out_ndim = in.dim();
   const executorch::aten::SizesType casted_upscale_factor = upscale_factor;
 
@@ -366,6 +369,8 @@ void get_pixel_shuffle_out_target_size(
   out_sizes[i] = in.size(i) * casted_upscale_factor;
   i++;
   out_sizes[i] = in.size(i) * casted_upscale_factor;
+
+  return true;
 }
 
 void get_pixel_unshuffle_out_target_size(

kernels/portable/cpu/util/copy_ops_util.h

@@ -139,7 +139,7 @@ bool check_pixel_shuffle_args(
     int64_t upscale_factor,
     Tensor& out);
 
-void get_pixel_shuffle_out_target_size(
+bool get_pixel_shuffle_out_target_size(
     const Tensor& in,
     int64_t upscale_factor,
     executorch::aten::SizesType* out_sizes,