From 34175826dfd3043ec1ef28bab28d6441e9ab77dd Mon Sep 17 00:00:00 2001
From: Ryan Hill <38674843+RyanUnderhill@users.noreply.github.com>
Date: Thu, 13 Dec 2018 10:44:01 -0800
Subject: [PATCH] Ryanunderhill/pad fix (#151)

* Fix pad bug when there is negative padding

* Add in dimension count check
---
 onnxruntime/core/providers/cpu/tensor/pad.cc  |  3 +-
 onnxruntime/core/providers/cpu/tensor/utils.h | 36 +++++++++++++++
 .../test/providers/cpu/tensor/pad_test.cc     | 46 +++++++++++++++++++
 3 files changed, 84 insertions(+), 1 deletion(-)

diff --git a/onnxruntime/core/providers/cpu/tensor/pad.cc b/onnxruntime/core/providers/cpu/tensor/pad.cc
index 5d8b19c501..3dbd4872a1 100644
--- a/onnxruntime/core/providers/cpu/tensor/pad.cc
+++ b/onnxruntime/core/providers/cpu/tensor/pad.cc
@@ -53,6 +53,7 @@ Status Pad<float>::Compute(OpKernelContext* ctx) const {
   std::vector<int64_t> output_dims(input_tensor.Shape().GetDims());
   size_t dimension_count = output_dims.size();
 
+  ONNXRUNTIME_ENFORCE(dimension_count > 0, "Input tensor has no dimensions");
   ONNXRUNTIME_ENFORCE(dimension_count * 2 == pads_.size(), "'pads' attribute has wrong number of values");
 
   std::vector<int64_t> input_starts;
@@ -78,7 +79,7 @@ Status Pad<float>::Compute(OpKernelContext* ctx) const {
     alignSkip += pads_[i] * output_pitches[i];
 
   size_t inner_axis = dimension_count - 1;
-  TensorAxisCounters input_counters(input_tensor);
+  ExtentAxisCounters input_counters(input_extents);
 
   switch (mode_) {
     case Mode::Constant:
diff --git a/onnxruntime/core/providers/cpu/tensor/utils.h b/onnxruntime/core/providers/cpu/tensor/utils.h
index a59514d999..ec2976914a 100644
--- a/onnxruntime/core/providers/cpu/tensor/utils.h
+++ b/onnxruntime/core/providers/cpu/tensor/utils.h
@@ -83,6 +83,42 @@ struct TensorAxisCounters {
   std::vector<int64_t> indices_;  // There is no index for innermost axis since it's a special case
 };
 
+struct ExtentAxisCounters {
+  ExtentAxisCounters(gsl::span<const int64_t> extents) : extents_(extents) {
+    indices_.resize(extents_.size() - 1, 0);
+    axis_ = indices_.size();
+
+    // If a tensor has a shape, but one of the axes is 0 in size, there are no elements, so nothing to iterate
+    if (std::find(extents.cbegin(), extents.cend(), 0) != extents.cend())
+      running_ = false;
+  }
+
+  // Returns true if there was a carry to the next axis
+  bool Increment() {
+    if (axis_-- == 0) {
+      running_ = false;
+      return false;
+    }
+
+    if (++indices_[axis_] != extents_[axis_]) {
+      axis_ = indices_.size();
+      return false;
+    }
+
+    indices_[axis_] = 0;  // Reset the counter for this axis
+    return true;          // There was a carry
+  }
+
+  size_t Axis() const { return axis_; }
+  operator bool() const { return running_; }
+
+ private:
+  bool running_{true};
+  size_t axis_;
+  std::vector<int64_t> indices_;      // There is no index for innermost axis since it's a special case
+  gsl::span<const int64_t> extents_;  // The extents of each axis
+};
+
 // A std::vector that holds the number of entries to skip to go to the next axis start given an extent in each axis
 // This is used by the SliceIterator to iterate over a slice of a tensor
 struct SliceSkips : std::vector<int64_t> {
diff --git a/onnxruntime/test/providers/cpu/tensor/pad_test.cc b/onnxruntime/test/providers/cpu/tensor/pad_test.cc
index 8a625436ba..ed780d7097 100644
--- a/onnxruntime/test/providers/cpu/tensor/pad_test.cc
+++ b/onnxruntime/test/providers/cpu/tensor/pad_test.cc
@@ -7,6 +7,16 @@
 namespace onnxruntime {
 namespace test {
 
+TEST(TensorOpTest, Pad_Spec_Example) {
+  OpTester test("Pad");
+
+  test.AddAttribute("pads", std::vector<int64_t>{0, 2, 0, 0});
+  test.AddAttribute("value", 0.0f);
+  test.AddInput<float>("data", {3, 2}, {1.0f, 1.2f, 2.3f, 3.4f, 4.5f, 5.7f});
+  test.AddOutput<float>("output", {3, 4}, {0.0f, 0.0f, 1.0f, 1.2f, 0.0f, 0.0f, 2.3f, 3.4f, 0.0f, 0.0f, 4.5f, 5.7f});
+  test.Run();
+}
+
 TEST(TensorOpTest, Pad_Constant_1D) {
   OpTester test("Pad");
 
@@ -43,6 +53,42 @@ TEST(TensorOpTest, Pad_Constant_2D) {
   test.Run();
 }
 
+TEST(TensorOpTest, Pad_Constant_2D_negative) {
+  OpTester test("Pad");
+
+  test.AddAttribute("pads", std::vector<int64_t>{1, 2, 1, -1});
+  test.AddAttribute("value", 1234.0f);
+  test.AddInput<float>("data", {2, 3},
+                       {11.0f, 21.0f, 31.0f,
+                        12.0f, 22.0f, 32.0f});
+  test.AddOutput<float>("output", {4, 4},
+                        {1234.0f, 1234.0f, 1234.0f, 1234.0f,
+                         1234.0f, 1234.0f, 11.0f, 21.0f,
+                         1234.0f, 1234.0f, 12.0f, 22.0f,
+                         1234.0f, 1234.0f, 1234.0f, 1234.0f});
+  test.Run();
+}
+
+TEST(TensorOpTest, Pad_3D_complex) {
+  OpTester test("Pad");
+
+  test.AddAttribute("pads", std::vector<int64_t>{1, 0, 0, -1, 0, 0});
+  test.AddAttribute("value", 0.0f);
+  test.AddInput<float>("data", {2, 2, 2},
+                       {111.0f, 112.0f,
+                        121.0f, 122.0f,
+
+                        211.0f, 212.0f,
+                        221.0f, 222.0f});
+  test.AddOutput<float>("output", {2, 2, 2},
+                        {0.0f, 0.0f,
+                         0.0f, 0.0f,
+
+                         111.0f, 112.0f,
+                         121.0f, 122.0f});
+  test.Run();
+}
+
 TEST(TensorOpTest, Pad_Edge_2D) {
   OpTester test("Pad");