Add checks for compute_slice (#15647)

Summary:

Add safety checks to compute_slice, to ensure that we:

1. Do not read outside of the src tensor bounds
2. Do not write outside of the output tensor bounds

Also pass in KernelRuntimeContext to use ET_KERNEL_CHECK_MSG and make errors non-fatal.

Differential Revision: D86433966

Author: lucylq

backends/cadence/fusion_g3/operators/op_slice_copy.cpp

@@ -123,7 +123,7 @@ Tensor& slice_copy_Tensor_out(
         InvalidArgument,
         out);
 
-    torch::executor::compute_slice(in, dim, start, length, step, out);
+    torch::executor::compute_slice(ctx, in, dim, start, length, step, out);
   }
 
   return out;

backends/cadence/hifi/operators/op_slice_copy.cpp

@@ -64,7 +64,7 @@ Tensor& slice_copy_Tensor_out(
       InvalidArgument,
       out);
 
-  compute_slice(in, dim, start, length, step, out);
+  compute_slice(ctx, in, dim, start, length, step, out);
 
   return out;
 }

kernels/portable/cpu/op_narrow_copy.cpp

@@ -46,7 +46,7 @@ Tensor& narrow_copy_out(
       out);
 
   if (length != 0) {
-    compute_slice(in, dim, start, length, 1, out);
+    compute_slice(ctx, in, dim, start, length, 1, out);
   }
 
   return out;

kernels/portable/cpu/op_slice_copy.cpp

@@ -55,7 +55,7 @@ Tensor& slice_copy_Tensor_out(
       InvalidArgument,
       out);
 
-  compute_slice(in, dim, start, length, step, out);
+  compute_slice(ctx, in, dim, start, length, step, out);
 
   return out;
 }

kernels/portable/cpu/util/slice_util.cpp

@@ -150,14 +150,33 @@ int64_t adjust_slice_indices(
 }
 
 void compute_slice(
+    KernelRuntimeContext& ctx,
     const Tensor& in,
     int64_t dim,
     int64_t start,
     int64_t length,
     int64_t step,
     Tensor& out) {
+  ET_KERNEL_CHECK_MSG(
+      ctx,
+      dim < in.dim(),
+      InvalidArgument,
+      /* void */,
+      "Requested dim is larger than input tensor dim; dim = %" PRId64,
+      dim);
   size_t dim_length = in.size(dim);
-
+  ET_KERNEL_CHECK_MSG(
+      ctx,
+      start < dim_length,
+      InvalidArgument,
+      /* void */,
+      "Requested start is larger than the dim length.");
+  ET_KERNEL_CHECK_MSG(
+      ctx,
+      length * step < dim_length,
+      InvalidArgument,
+      /* void */,
+      "Requested length * step is larger than the dim size.");
   size_t leading_dims = getLeadingDims(in, dim);
   size_t trailing_dims = getTrailingDims(in, dim);
 
@@ -170,6 +189,12 @@ void compute_slice(
   const char* input_data = in.const_data_ptr<char>();
   char* dest = out.mutable_data_ptr<char>();
 
+  ET_KERNEL_CHECK_MSG(
+      ctx,
+      out.nbytes() >= (length * leading_dims * length_per_step),
+      InvalidArgument,
+      /* void */,
+      "out.nbytes() is smaller than the expected slice size.");
   for (const auto i : c10::irange(leading_dims)) {
     const char* src = input_data + (i * dim_length + start) * length_per_step;
     for ([[maybe_unused]] const auto j : c10::irange(length)) {

kernels/portable/cpu/util/slice_util.h

@@ -55,6 +55,7 @@ int64_t adjust_slice_indices(
     int64_t step);
 
 void compute_slice(
+    KernelRuntimeContext& ctx,
     const Tensor& in,
     int64_t dim,
     int64_t start,