diff --git a/onnxruntime/core/providers/cpu/tensor/slice_helper.h b/onnxruntime/core/providers/cpu/tensor/slice_helper.h index a8bd118c55..c62212ec48 100644 --- a/onnxruntime/core/providers/cpu/tensor/slice_helper.h +++ b/onnxruntime/core/providers/cpu/tensor/slice_helper.h @@ -8,14 +8,6 @@ namespace onnxruntime { -// std::clamp doesn't exist until C++17 so create a local version -template -const T& clamp(const T& v, const T& lo, const T& hi) { - if (v < lo) return lo; - if (v > hi) return hi; - return v; -} - namespace SliceOp { // compute output_dims without steps (Slice V1-9 & DynamicSlice) // Please note this will not Flatten the output shape @@ -35,27 +27,28 @@ inline Status PrepareForComputeHelper(const std::vector& raw_starts, std::unordered_set unique_axes; const auto& dimension_count = compute_metadata.input_dimensions_.size(); for (size_t axis_index = 0, axes_count = axes.size(); axis_index < axes_count; ++axis_index) { - auto axis = HandleNegativeAxis(axes[axis_index], dimension_count); // handle negative and enforce axis is valid + const auto axis = HandleNegativeAxis(axes[axis_index], dimension_count); // handle negative and enforce axis is valid if (axis >= static_cast(dimension_count) || axis < 0) return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "'axes' has an axis outside of the tensor dimension count"); if (unique_axes.find(axis) != unique_axes.end()) return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "'axes' has duplicates"); unique_axes.insert(axis); + const auto dim_value = compute_metadata.input_dimensions_[axis]; // process start auto start = raw_starts[axis_index]; if (start < 0) - start += compute_metadata.input_dimensions_[axis]; - compute_metadata.starts_[axis] = clamp(start, int64_t{0}, compute_metadata.input_dimensions_[axis]); + start += dim_value; + compute_metadata.starts_[axis] = std::clamp(start, int64_t{0}, dim_value); // process end auto end = raw_ends[axis_index]; if (end < 0) - end += compute_metadata.input_dimensions_[axis]; - compute_metadata.ends_[axis] = clamp(end, int64_t{0}, compute_metadata.input_dimensions_[axis]); + end += dim_value; + compute_metadata.ends_[axis] = std::clamp(end, int64_t{0}, dim_value); // find output dim value for this axis - auto temp = compute_metadata.ends_[axis] - compute_metadata.starts_[axis]; + const auto temp = compute_metadata.ends_[axis] - compute_metadata.starts_[axis]; if (temp < 0) compute_metadata.output_dims_[axis] = 0; else @@ -85,15 +78,19 @@ inline Status PrepareForComputeHelper(const std::vector& raw_starts, std::unordered_set unique_axes; const auto& dimension_count = compute_metadata.input_dimensions_.size(); for (size_t axis_index = 0, axes_count = axes.size(); axis_index < axes_count; ++axis_index) { - auto axis = axes[axis_index] < 0 ? axes[axis_index] + static_cast(dimension_count) : axes[axis_index]; + const auto axis = axes[axis_index] < 0 ? axes[axis_index] + static_cast(dimension_count) : axes[axis_index]; if (axis >= static_cast(dimension_count) || axis < 0) return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "'axes' has an axis outside of the tensor dimension count"); if (unique_axes.find(axis) != unique_axes.end()) return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "'axes' has duplicates"); unique_axes.insert(axis); + const auto dim_value = compute_metadata.input_dimensions_[axis]; // process step auto step = axis_index < raw_steps.size() ? raw_steps[axis_index] : 1; + // clamp step to avoid overflow if there's a stupidly large value (which will be multiplied in SliceImpl) + // as long as the clamped value is >= the size of the dimension a single step will push us past the end + step = std::clamp(step, -dim_value, dim_value); if (step == 0) return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "'step' value cannot be 0"); compute_metadata.steps_[axis] = step; @@ -101,11 +98,11 @@ inline Status PrepareForComputeHelper(const std::vector& raw_starts, // process start auto start = raw_starts[axis_index]; if (start < 0) - start += compute_metadata.input_dimensions_[axis]; + start += dim_value; if (step < 0) - compute_metadata.starts_[axis] = clamp(start, int64_t{0}, compute_metadata.input_dimensions_[axis] - 1); + compute_metadata.starts_[axis] = std::clamp(start, int64_t{0}, dim_value - 1); else - compute_metadata.starts_[axis] = clamp(start, int64_t{0}, compute_metadata.input_dimensions_[axis]); + compute_metadata.starts_[axis] = std::clamp(start, int64_t{0}, dim_value); // process end auto end = raw_ends[axis_index]; @@ -114,20 +111,20 @@ inline Status PrepareForComputeHelper(const std::vector& raw_starts, // it represent slicing to the end of the dimension if (end == std::numeric_limits::max() || end == std::numeric_limits::max()) { - end = step < 0 ? -1 : compute_metadata.input_dimensions_[axis]; + end = step < 0 ? -1 : dim_value; } else { if (end < 0) - end += compute_metadata.input_dimensions_[axis]; + end += dim_value; if (step < 0) - end = clamp(end, int64_t{-1}, compute_metadata.input_dimensions_[axis]); + end = std::clamp(end, int64_t{-1}, dim_value); else - end = clamp(end, int64_t{0}, compute_metadata.input_dimensions_[axis]); + end = std::clamp(end, int64_t{0}, dim_value); } compute_metadata.ends_[axis] = end; // find output dim value for this axis - auto temp = static_cast(ceil(1.0 * (compute_metadata.ends_[axis] - compute_metadata.starts_[axis]) / step)); + const auto temp = static_cast(ceil(1.0 * (compute_metadata.ends_[axis] - compute_metadata.starts_[axis]) / step)); if (temp < 0) compute_metadata.output_dims_[axis] = 0; else diff --git a/onnxruntime/test/providers/cpu/tensor/slice_op.test.cc b/onnxruntime/test/providers/cpu/tensor/slice_op.test.cc index 2326876e82..6dfa5c3911 100644 --- a/onnxruntime/test/providers/cpu/tensor/slice_op.test.cc +++ b/onnxruntime/test/providers/cpu/tensor/slice_op.test.cc @@ -18,7 +18,8 @@ void RunSliceTest(const std::vector& input_dims, const std::vector& steps, const std::vector& output_dims, const std::vector& output_vals, - bool v10_only = false) { + bool v10_only = false, + const std::unordered_set& excluded_providers_input = {}) { std::unordered_set excluded_providers; if (!v10_only) @@ -26,6 +27,9 @@ void RunSliceTest(const std::vector& input_dims, else excluded_providers = {kTensorrtExecutionProvider}; + // merge with the excluded ep from input + excluded_providers.insert(excluded_providers_input.cbegin(), excluded_providers_input.cend()); + // NNAPI EP does not support empty output if (std::any_of(output_dims.cbegin(), output_dims.cend(), [](int64_t i) { return i == 0; })) { excluded_providers.insert(kNnapiExecutionProvider); @@ -604,5 +608,27 @@ TEST(SliceTest, Slice5D_SubsetOfAxes_Flatten2Dims_OffsetInput) { {-5.f, -6.f, -7.f, -8.f}, true); } + +// test where we use a ridiculously large step, using a large enough step has caused +// ORT crash due to integer overflow on 32bit system +// See, https://github.com/microsoft/onnxruntime/issues/9368 +TEST(SliceTest, Slice5D_LargeStep) { + RunSliceTest({1, 2, 2, 2, 2}, + {1.f, 2.f, 3.f, 4.f, + 5.f, 6.f, 7.f, 8.f, + -1.f, -2.f, -3.f, -4.f, + -5.f, -6.f, -7.f, -8.f}, + {0}, + {1}, + {1}, + {std::numeric_limits::max()}, + {1, 1, 2, 2, 2}, + {1.f, 2.f, 3.f, 4.f, + 5.f, 6.f, 7.f, 8.f}, + true, + // Nuphar EP cannot handle large steps, TODO, add step clamp to Nuphar EP + {kNupharExecutionProvider}); +} + } // namespace test } // namespace onnxruntime