saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix integer overflow for large step for Slice OP (#9376)

Browse source
This commit is contained in:
Guoyu Wang 2021-10-15 09:42:53 -07:00 committed by GitHub
parent 901c7de918
commit 59dfab59dc
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 47 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

43
onnxruntime/core/providers/cpu/tensor/slice_helper.h
View file

@ -8,14 +8,6 @@
namespace onnxruntime {
// std::clamp doesn't exist until C++17 so create a local version
template <typename T>
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<int64_t>& raw_starts,
std::unordered_set<int64_t> 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<int64_t>(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<int64_t>& raw_starts,
std::unordered_set<int64_t> 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<int64_t>(dimension_count) : axes[axis_index];
const auto axis = axes[axis_index] < 0 ? axes[axis_index] + static_cast<int64_t>(dimension_count) : axes[axis_index];
if (axis >= static_cast<int64_t>(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<int64_t>& 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<int64_t>& raw_starts,
// it represent slicing to the end of the dimension
if (end == std::numeric_limits<int32_t>::max() ||
end == std::numeric_limits<int64_t>::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<int64_t>(ceil(1.0 * (compute_metadata.ends_[axis] - compute_metadata.starts_[axis]) / step));
const auto temp = static_cast<int64_t>(ceil(1.0 * (compute_metadata.ends_[axis] - compute_metadata.starts_[axis]) / step));
if (temp < 0)
compute_metadata.output_dims_[axis] = 0;
else

28
onnxruntime/test/providers/cpu/tensor/slice_op.test.cc
View file

@ -18,7 +18,8 @@ void RunSliceTest(const std::vector<int64_t>& input_dims,
const std::vector<int64_t>& steps,
const std::vector<int64_t>& output_dims,
const std::vector<T>& output_vals,
bool v10_only = false) {
bool v10_only = false,
const std::unordered_set<std::string>& excluded_providers_input = {}) {
std::unordered_set<std::string> excluded_providers;
if (!v10_only)
@ -26,6 +27,9 @@ void RunSliceTest(const std::vector<int64_t>& 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<float>({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<int64_t>::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