Handle edge case in CumSum causing overflow (#13174)

### Description
<!-- Describe your changes. -->
Add special case handling for exclusive + reverse where axis has dim
value of 1.

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
#13165
This commit is contained in:
Scott McKay 2022-10-06 07:18:02 +10:00 committed by GitHub
parent 4e37464cc5
commit cf075fcbad
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 106 additions and 65 deletions

View file

@ -74,74 +74,90 @@ Status GetAxis(const Tensor* axis_tensor, int64_t input_rank, int64_t& axis_out)
} // namespace cumsum_op
ONNX_CPU_OPERATOR_VERSIONED_TYPED_KERNEL(CumSum,
11,
13,
float,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::GetTensorType<float>())
.TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(), DataTypeImpl::GetTensorType<int64_t>()}),
CumSum<float>);
ONNX_CPU_OPERATOR_VERSIONED_TYPED_KERNEL(
CumSum,
11,
13,
float,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::GetTensorType<float>())
.TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),
CumSum<float>);
ONNX_CPU_OPERATOR_VERSIONED_TYPED_KERNEL(CumSum,
11,
13,
double,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::GetTensorType<double>())
.TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(), DataTypeImpl::GetTensorType<int64_t>()}),
CumSum<double>);
ONNX_CPU_OPERATOR_VERSIONED_TYPED_KERNEL(
CumSum,
11,
13,
double,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::GetTensorType<double>())
.TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),
CumSum<double>);
ONNX_CPU_OPERATOR_VERSIONED_TYPED_KERNEL(CumSum,
11,
13,
int32_t,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::GetTensorType<int32_t>())
.TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(), DataTypeImpl::GetTensorType<int64_t>()}),
CumSum<int32_t>);
ONNX_CPU_OPERATOR_VERSIONED_TYPED_KERNEL(
CumSum,
11,
13,
int32_t,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::GetTensorType<int32_t>())
.TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),
CumSum<int32_t>);
ONNX_CPU_OPERATOR_VERSIONED_TYPED_KERNEL(CumSum,
11,
13,
int64_t,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::GetTensorType<int64_t>())
.TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(), DataTypeImpl::GetTensorType<int64_t>()}),
CumSum<int64_t>);
ONNX_CPU_OPERATOR_VERSIONED_TYPED_KERNEL(
CumSum,
11,
13,
int64_t,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::GetTensorType<int64_t>())
.TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),
CumSum<int64_t>);
// Opset 14 kernels
ONNX_CPU_OPERATOR_TYPED_KERNEL(CumSum,
14,
float,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::GetTensorType<float>())
.TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(), DataTypeImpl::GetTensorType<int64_t>()}),
CumSum<float>);
ONNX_CPU_OPERATOR_TYPED_KERNEL(
CumSum,
14,
float,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::GetTensorType<float>())
.TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),
CumSum<float>);
ONNX_CPU_OPERATOR_TYPED_KERNEL(CumSum,
14,
double,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::GetTensorType<double>())
.TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(), DataTypeImpl::GetTensorType<int64_t>()}),
CumSum<double>);
ONNX_CPU_OPERATOR_TYPED_KERNEL(
CumSum,
14,
double,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::GetTensorType<double>())
.TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),
CumSum<double>);
ONNX_CPU_OPERATOR_TYPED_KERNEL(CumSum,
14,
int32_t,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::GetTensorType<int32_t>())
.TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(), DataTypeImpl::GetTensorType<int64_t>()}),
CumSum<int32_t>);
ONNX_CPU_OPERATOR_TYPED_KERNEL(
CumSum,
14,
int32_t,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::GetTensorType<int32_t>())
.TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),
CumSum<int32_t>);
ONNX_CPU_OPERATOR_TYPED_KERNEL(CumSum,
14,
int64_t,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::GetTensorType<int64_t>())
.TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(), DataTypeImpl::GetTensorType<int64_t>()}),
CumSum<int64_t>);
ONNX_CPU_OPERATOR_TYPED_KERNEL(
CumSum,
14,
int64_t,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::GetTensorType<int64_t>())
.TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),
CumSum<int64_t>);
template <typename T>
CumSum<T>::CumSum(const OpKernelInfo& info) : OpKernel(info), exclusive_(), reverse_() {
@ -223,7 +239,8 @@ Status CumSum<T>::Compute(OpKernelContext* ctx) const {
::ZeroOutSliceAtIndex<T>(output_tensor, rank, axis, index, slice_dims, steps, slice_size);
--index;
}
{
if (index >= 0) {
// The next slice is a copy of the input (if exclusive == false then this is the first slice)
auto input_starts(::GetStarts(rank, axis, dim - 1));
auto output_starts(::GetStarts(rank, axis, index));

View file

@ -12,12 +12,12 @@ namespace test {
TEST(CumSumTest, _1DTest) {
OpTester test("CumSum", 11, onnxruntime::kOnnxDomain);
test.AddInput<float>("x", {5}, {1., 2., 3., 4., 5.});
// Pass in 0D Axis for all OpenVINO tests, and keep one 1D Axis test for coverage.
#ifdef USE_OPENVINO
// Pass in 0D Axis for all OpenVINO tests, and keep one 1D Axis test for coverage.
#ifdef USE_OPENVINO
test.AddInput<int32_t>("axis", {}, {0});
#else
#else
test.AddInput<int32_t>("axis", {1}, {0});
#endif
#endif
test.AddOutput<float>("y", {5}, {1., 3., 6., 10., 15.});
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});
}
@ -52,6 +52,30 @@ TEST(CumSumTest, _1DTestExclusive) {
test.AddOutput<float>("y", {5}, {0., 1., 3., 6., 10.});
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});
}
// GH13165.
TEST(CumSumTest, _1DTestExclusiveAxisHasSingleValue) {
{
// forward
OpTester test("CumSum", 11, onnxruntime::kOnnxDomain);
test.AddAttribute<int64_t>("exclusive", 1);
test.AddInput<float>("x", {1, 2}, {1., 2.});
test.AddInput<int32_t>("axis", {}, {0}); // dim value of axis is 1
test.AddOutput<float>("y", {1, 2}, {0., 0.});
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});
}
{
// reverse
OpTester test("CumSum", 11, onnxruntime::kOnnxDomain);
test.AddAttribute<int64_t>("exclusive", 1);
test.AddAttribute<int64_t>("reverse", 1);
test.AddInput<float>("x", {1, 2}, {1., 2.});
test.AddInput<int32_t>("axis", {}, {0}); // dim value of axis is 1
test.AddOutput<float>("y", {1, 2}, {0., 0.});
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});
}
}
TEST(CumSumTest, _2DTestAxis0) {
OpTester test("CumSum", 11, onnxruntime::kOnnxDomain);
test.AddInput<float>("x", {2, 3}, {1., 2., 3., 4., 5., 6.});