diff --git a/backends/cadence/fusion_g3/operators/op_slice_copy.cpp b/backends/cadence/fusion_g3/operators/op_slice_copy.cpp index a97f9beb0c7..504a00fcaee 100644 --- a/backends/cadence/fusion_g3/operators/op_slice_copy.cpp +++ b/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; diff --git a/backends/cadence/hifi/operators/op_slice_copy.cpp b/backends/cadence/hifi/operators/op_slice_copy.cpp index 014eaa6698b..ff447461d6e 100644 --- a/backends/cadence/hifi/operators/op_slice_copy.cpp +++ b/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; } diff --git a/kernels/portable/cpu/op_narrow_copy.cpp b/kernels/portable/cpu/op_narrow_copy.cpp index 960ea35efac..260ee9697db 100644 --- a/kernels/portable/cpu/op_narrow_copy.cpp +++ b/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; diff --git a/kernels/portable/cpu/op_slice_copy.cpp b/kernels/portable/cpu/op_slice_copy.cpp index 1d4e509e083..0baacb874e1 100644 --- a/kernels/portable/cpu/op_slice_copy.cpp +++ b/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; } diff --git a/kernels/portable/cpu/util/slice_util.cpp b/kernels/portable/cpu/util/slice_util.cpp index 909bd827d79..531218a3083 100644 --- a/kernels/portable/cpu/util/slice_util.cpp +++ b/kernels/portable/cpu/util/slice_util.cpp @@ -150,13 +150,31 @@ 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) { + // No slicing requested. + if (length <= 0) { + return; + } + + ET_KERNEL_CHECK_MSG( + ctx, + dim < in.dim(), + InvalidArgument, + /* void */, + "Requested dim is larger than input tensor dim"); size_t dim_length = in.size(dim); + ET_KERNEL_CHECK_MSG( + ctx, + (start + (length - 1) * step) < dim_length, + InvalidArgument, + /* void */, + "Requested slice is larger than the dim size"); size_t leading_dims = getLeadingDims(in, dim); size_t trailing_dims = getTrailingDims(in, dim); @@ -170,6 +188,12 @@ void compute_slice( const char* input_data = in.const_data_ptr(); char* dest = out.mutable_data_ptr(); + 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)) { diff --git a/kernels/portable/cpu/util/slice_util.h b/kernels/portable/cpu/util/slice_util.h index accfb387246..52a0b05c864 100644 --- a/kernels/portable/cpu/util/slice_util.h +++ b/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,