From cf075fcbad4cef486110f0bae0b6c15f7e01dd55 Mon Sep 17 00:00:00 2001
From: Scott McKay <skottmckay@gmail.com>
Date: Thu, 6 Oct 2022 07:18:02 +1000
Subject: [PATCH] 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
---
 onnxruntime/core/providers/cpu/math/cumsum.cc | 139 ++++++++++--------
 .../test/providers/cpu/math/cumsum_test.cc    |  32 +++-
 2 files changed, 106 insertions(+), 65 deletions(-)

diff --git a/onnxruntime/core/providers/cpu/math/cumsum.cc b/onnxruntime/core/providers/cpu/math/cumsum.cc
index 53a784a031..ccb5546b51 100644
--- a/onnxruntime/core/providers/cpu/math/cumsum.cc
+++ b/onnxruntime/core/providers/cpu/math/cumsum.cc
@@ -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));
diff --git a/onnxruntime/test/providers/cpu/math/cumsum_test.cc b/onnxruntime/test/providers/cpu/math/cumsum_test.cc
index 2149e27472..cbf2d2fa2b 100644
--- a/onnxruntime/test/providers/cpu/math/cumsum_test.cc
+++ b/onnxruntime/test/providers/cpu/math/cumsum_test.cc
@@ -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.});