Do not copy beyond out tensor bounds (#15648)

Summary:

`set_all_to_value(out_data, out_step_len, value)` loops over the padding and copies a row at a time.

Add some checks to make sure we do not exceed the out tensor (out_data) size when copying.

Differential Revision: D86436435

Author: lucylq

kernels/portable/cpu/op_constant_pad_nd.cpp

@@ -57,7 +57,17 @@ 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) {
+    auto numel = 1;
+    for (ET_UNUSED const auto i : c10::irange(out_sizes.size())) {
+      numel *= out_sizes[i];
+    }
+    ET_CHECK_MSG(
+        numel * sizeof(CTYPE) >= pad_before * out_step_len,
+        "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,7 +86,7 @@ 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(
           ndim,
           self_data,
@@ -95,7 +105,17 @@ 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) {
+    auto numel = 1;
+    for (ET_UNUSED const auto i : c10::irange(out_sizes.size())) {
+      numel *= out_sizes[i];
+    }
+    ET_CHECK_MSG(
+        numel * sizeof(CTYPE) >= pad_after * out_step_len,
+        "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;
   }