From 0dc38eac2e383c8a089010ced49146f30d93b010 Mon Sep 17 00:00:00 2001
From: chxin66 <57057788+chxin66@users.noreply.github.com>
Date: Wed, 13 Apr 2022 22:16:47 +0800
Subject: [PATCH] Added unit test for maxpool (#361)

https://github.com/VeriSilicon/TIM-VX/issues/318

Signed-off-by: Chen Xin <jack.chen@verisilicon.com>
---
 include/tim/vx/ops/pool2d.h     | 30 +++++++++----------
 src/tim/vx/ops/max_pool_test.cc | 52 +++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 15 deletions(-)
 create mode 100644 src/tim/vx/ops/max_pool_test.cc

diff --git a/include/tim/vx/ops/pool2d.h b/include/tim/vx/ops/pool2d.h
index b5ba042..9e3bd76 100644
--- a/include/tim/vx/ops/pool2d.h
+++ b/include/tim/vx/ops/pool2d.h
@@ -37,61 +37,61 @@ namespace ops {
  * ## Pool2d
  *
  * ### Classic Pool2d
- * 
+ *
  * Performs an 2-D pooling operation.
  *
  * - type : MAX, AVG, L2 or AVG_ANDROID.
  * - padding : AUTO, VALID or SAME.
+ * - pad : Specify the number of pad values for left, right, top, and bottom.
  * - ksize : filter size.
  * - stride : stride along each spatial axis.
  * - round_type : CEILING or FLOOR.
- * 
+ *
  * ### Global Pool2d
- * 
+ *
  * - type : MAX, AVG, L2 or AVG_ANDROID.
  * - input_size : input size(only [W， H])
  * - round_type : CEILING or FLOOR.
- * 
+ *
  * ### Adaptive Pool2d
- * 
+ *
  * Same as torch.nn.AdaptiveXXXPool2d.
- * 
+ *
  * - type : MAX, AVG, L2 or AVG_ANDROID.
  * - input_size : input size(only [W， H])
  * - output_size : output size(only [W， H])
  * - round_type : CEILING or FLOOR.
- * 
+ *
  */
 
 class Pool2d : public DirectMapOp {
  public:
-  // for Classic Pool2d
+  /* for Classic Pool2d, pool does not support auto-completion of pad value,
+  you need to specify pad size explicitly, it is recommended to use the second api.*/
   Pool2d(Graph* graph, PoolType type, PadType padding,
          const std::array<uint32_t, 2>& ksize,
          const std::array<uint32_t, 2>& stride,
          RoundType round_type = RoundType::FLOOR,
          DataLayout layout = DataLayout::WHCN);
-  Pool2d(Graph* graph, PoolType type,
-         const std::array<uint32_t, 4>& pad,
+  Pool2d(Graph* graph, PoolType type, const std::array<uint32_t, 4>& pad,
          const std::array<uint32_t, 2>& ksize,
          const std::array<uint32_t, 2>& stride,
          RoundType round_type = RoundType::FLOOR,
          DataLayout layout = DataLayout::WHCN);
 
   // for Global Pool2d
-  Pool2d(Graph* graph, PoolType type,
-         const std::array<uint32_t, 2>& input_size,
+  Pool2d(Graph* graph, PoolType type, const std::array<uint32_t, 2>& input_size,
          RoundType round_type = RoundType::FLOOR,
          DataLayout layout = DataLayout::WHCN);
 
   // for Adaptive Pool2d
-  Pool2d(Graph* graph, PoolType type,
-         const std::array<uint32_t, 2>& input_size,
+  Pool2d(Graph* graph, PoolType type, const std::array<uint32_t, 2>& input_size,
          const std::array<uint32_t, 2>& output_size,
          RoundType round_type = RoundType::FLOOR,
          DataLayout layout = DataLayout::WHCN);
 
-  std::shared_ptr<Operation> Clone(std::shared_ptr<Graph>& graph) const override;
+  std::shared_ptr<Operation> Clone(
+      std::shared_ptr<Graph>& graph) const override;
   void Init();
 
  protected:
diff --git a/src/tim/vx/ops/max_pool_test.cc b/src/tim/vx/ops/max_pool_test.cc
new file mode 100644
index 0000000..c245414
--- /dev/null
+++ b/src/tim/vx/ops/max_pool_test.cc
@@ -0,0 +1,52 @@
+#include "tim/vx/context.h"
+#include "tim/vx/graph.h"
+#include "tim/vx/ops/pool2d.h"
+#include <iostream>
+#include "gtest/gtest.h"
+#include "test_utils.h"
+
+TEST(MAX, shape_6_6_1_1_fp32_kernel_3_stride_2) {
+  auto ctx = tim::vx::Context::Create();
+  auto graph = ctx->CreateGraph();
+
+  tim::vx::ShapeType in_shape({6, 6, 1, 1});
+  tim::vx::ShapeType out_shape({3, 3, 1, 1});
+  tim::vx::TensorSpec input_spec(tim::vx::DataType::FLOAT32, in_shape,
+                                 tim::vx::TensorAttribute::INPUT);
+  tim::vx::TensorSpec output_spec(tim::vx::DataType::FLOAT32, out_shape,
+                                  tim::vx::TensorAttribute::OUTPUT);
+
+  auto input_tensor = graph->CreateTensor(input_spec);
+  auto output_tensor = graph->CreateTensor(output_spec);
+
+  std::vector<float> in_data = {
+      -1,  -2,  -3,  -4,  -5,  -5,
+      -6,  -7,  -8,  -9,  -10, -10,
+      -11, -12, -13, -14, -15, -15,
+      -16, -17, -18, -19, -20, -20,
+      -21, -22, -23, -24, -25, -20,
+      -21, -22, -23, -24, -25, -20,
+  };
+  std::vector<float> golden = {
+      -1, -3, -5,
+      -11, -13, -15,
+      -21, -23, -20,
+  };
+
+  EXPECT_TRUE(
+      input_tensor->CopyDataToTensor(in_data.data(), in_data.size() * 4));
+  std::array<uint32_t, 4> pad = {0, 1, 0, 1};
+  std::array<uint32_t, 2> ksize = {3, 3};
+  std::array<uint32_t, 2> stride = {2, 2};
+  auto round_type = tim::vx::RoundType::CEILING;
+  auto op = graph->CreateOperation<tim::vx::ops::Pool2d>(
+      tim::vx::PoolType::MAX, pad, ksize, stride, round_type);
+  (*op).BindInputs({input_tensor}).BindOutputs({output_tensor});
+
+  EXPECT_TRUE(graph->Compile());
+  EXPECT_TRUE(graph->Run());
+
+  std::vector<float> output(golden.size());
+  EXPECT_TRUE(output_tensor->CopyDataFromTensor(output.data()));
+  EXPECT_EQ(golden, output);
+}
\ No newline at end of file