mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-12 17:57:38 +00:00
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:
parent
4e37464cc5
commit
cf075fcbad
2 changed files with 106 additions and 65 deletions
|
|
@ -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));
|
||||
|
|
|
|||
|
|
@ -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.});
|
||||
|
|
|
|||
Loading…
Reference in a new issue