Do not copy beyond out tensor bounds

Differential Revision: D86436435

Pull Request resolved: #15648

Author: lucylq

kernels/portable/cpu/op_constant_pad_nd.cpp

@@ -30,6 +30,7 @@ void set_all_to_value(CTYPE* out_data, size_t step_len, CTYPE value) {
 
 template <typename CTYPE>
 void apply_padding_to_dim(
+    KernelRuntimeContext& ctx,
     size_t ndim,
     const CTYPE* self_data,
     IntArrayRef self_sizes,
@@ -57,7 +58,20 @@ void apply_padding_to_dim(
   size_t out_step_len = out_strides[dim];
   size_t in_step_len = self_strides[dim];
 
-  for ([[maybe_unused]] const auto i : c10::irange(pad_before)) {
+  // Do not copy padding beyond the out tensor bounds.
+  if (pad_before > 0) {
+    size_t numel = 1;
+    for (ET_UNUSED const auto i : c10::irange(out_sizes.size())) {
+      numel *= out_sizes[i];
+    }
+    ET_KERNEL_CHECK_MSG(
+        ctx,
+        numel >= pad_before * out_step_len,
+        InvalidArgument,
+        /* void */,
+        "Out tensor is too small for the requested padding.");
+  }
+  for (ET_UNUSED const auto i : c10::irange(pad_before)) {
     set_all_to_value(out_data, out_step_len, value);
     out_data += out_step_len;
   }
@@ -76,8 +90,9 @@ void apply_padding_to_dim(
   }
   // Otherwise, call this function recursively
   else {
-    for ([[maybe_unused]] const auto i : c10::irange(self_sizes[dim])) {
+    for (ET_UNUSED const auto i : c10::irange(self_sizes[dim])) {
       apply_padding_to_dim(
+          ctx,
           ndim,
           self_data,
           self_sizes,
@@ -95,14 +110,28 @@ void apply_padding_to_dim(
     }
   }
 
-  for ([[maybe_unused]] const auto i : c10::irange(pad_after)) {
+  // Do not copy padding beyond the out tensor bounds.
+  if (pad_after > 0) {
+    size_t numel = 1;
+    for (ET_UNUSED const auto i : c10::irange(out_sizes.size())) {
+      numel *= out_sizes[i];
+    }
+    ET_KERNEL_CHECK_MSG(
+        ctx,
+        numel >= pad_after * out_step_len,
+        InvalidArgument,
+        /* void */,
+        "Out tensor is too small for the requested padding.");
+  }
+  for (ET_UNUSED const auto i : c10::irange(pad_after)) {
     set_all_to_value(out_data, out_step_len, value);
     out_data += out_step_len;
   }
 }
 
 template <typename CTYPE>
 void constant_pad_nd_out_impl(
+    KernelRuntimeContext& ctx,
     const Tensor& self,
     IntArrayRef pad,
     CTYPE value_v,
@@ -145,6 +174,7 @@ void constant_pad_nd_out_impl(
   IntArrayRef out_strides_ref(out_strides, ndim);
 
   apply_padding_to_dim(
+      ctx,
       ndim,
       self_data,
       self_sizes_ref,
@@ -192,7 +222,7 @@ Tensor& constant_pad_nd_out(
         utils::internal::check_overflow_scalar_cast<CTYPE>(value);
     ET_KERNEL_CHECK(ctx, opt_value_casted.has_value(), InvalidArgument, );
     auto value_casted = opt_value_casted.value();
-    constant_pad_nd_out_impl<CTYPE>(in, pad, value_casted, out);
+    constant_pad_nd_out_impl<CTYPE>(ctx, in, pad, value_casted, out);
   });
 
   return out;