[nnc] Add support for AvgPool2D and MaxPool2D (#7011)

These operations are the future replacement of the `Pool` operation.

Signed-off-by: Sergei Barannikov <s.barannikov@samsung.com>

17 files changed, 457 insertions, 53 deletions

diff --git a/compiler/nnc/include/passes/interpreter/Interpreter.h b/compiler/nnc/include/passes/interpreter/Interpreter.h
index 8441be764..a7881fd7a 100644
--- a/compiler/nnc/include/passes/interpreter/Interpreter.h
+++ b/compiler/nnc/include/passes/interpreter/Interpreter.h
@@ -39,6 +39,7 @@ public:
   mir::TensorVariant getResult(const mir::Operation::Output *tensor);
 
   void visit(mir::ops::AddOp &op) override;
+  void visit(mir::ops::AvgPool2DOp &op) override;
   void visit(mir::ops::CappedReluOp &op) override;
   void visit(mir::ops::ConcatOp &op) override;
   void visit(mir::ops::ConstantOp &op) override;
@@ -52,6 +53,7 @@ public:
   void visit(mir::ops::InputOp &op) override;
   void visit(mir::ops::LeakyReluOp &op) override;
   void visit(mir::ops::MaxOp &op) override;
+  void visit(mir::ops::MaxPool2DOp &op) override;
   void visit(mir::ops::MulOp &op) override;
   void visit(mir::ops::OutputOp &op) override;
   void visit(mir::ops::PadOp &op) override;
diff --git a/compiler/nnc/passes/acl_soft_backend/AclCppOpGenerator.cpp b/compiler/nnc/passes/acl_soft_backend/AclCppOpGenerator.cpp
index 9e1879eed..5fd4b082b 100644
--- a/compiler/nnc/passes/acl_soft_backend/AclCppOpGenerator.cpp
+++ b/compiler/nnc/passes/acl_soft_backend/AclCppOpGenerator.cpp
@@ -203,7 +203,7 @@ shared_ptr<ArtifactVariable> AclCppOpGenerator::genPadStrideInfo(const Op &op, c
 {
   using AF = ArtifactFactory;
 
-  const Shape &strides = transposeShape<1, 0>(op.getStrides());
+  const Shape strides(op.getStrides());
   assert(strides.rank() == 2);
   auto &padding_before = op.getPaddingBefore();
   auto &padding_after = op.getPaddingAfter();
@@ -213,8 +213,8 @@ shared_ptr<ArtifactVariable> AclCppOpGenerator::genPadStrideInfo(const Op &op, c
   string var_name = prefix + "_pad_stride_info";
 
   list<std::shared_ptr<ArtifactExpr>> var_init_params = {
-      AF::lit(to_string(op.getStrides().dim(1))),
-      AF::lit(to_string(op.getStrides().dim(0))),
+      AF::lit(to_string(strides.dim(1))),
+      AF::lit(to_string(strides.dim(0))),
       AF::lit(to_string(padding_before.at(1))),
       AF::lit(to_string(padding_after.at(1))),
       AF::lit(to_string(padding_before.at(0))),
@@ -317,6 +317,17 @@ void AclCppOpGenerator::visit(ops::PoolOp &op)
   genTensorDeallocation(_infBlock, transposed_output);
 }
 
+void AclCppOpGenerator::visit(ops::AvgPool2DOp &op)
+{
+  genPooling(op, "arm_compute::PoolingType::AVG", !op.getIncludePad());
+}
+
+void AclCppOpGenerator::visit(ops::MaxPool2DOp &op)
+{
+  // The value of 'exclude_padding' does not really matter for MAX pooling.
+  genPooling(op, "arm_compute::PoolingType::MAX", false);
+}
+
 void AclCppOpGenerator::visit(ops::FullyConnectedOp &op)
 {
   assert(op.getNumInputs() == 2);
@@ -493,6 +504,61 @@ void AclCppOpGenerator::visit(ops::PadOp &op)
 }
 
 template <typename Op>
+void AclCppOpGenerator::genPooling(Op &op, const std::string &pooling_type, bool exclude_padding)
+{
+  assert(op.getNumInputs() == 1);
+  const auto *ir_input = op.getInput(0)->getProducer();
+  const auto *ir_output = op.getOutput(0);
+
+  string in_name = tensorName(ir_input);
+  auto in_id = AF::id(in_name);
+
+  const string output_tensor_name = tensorName(ir_output);
+
+  // Transpose data from MIR format to format compatible with ACL
+  const string transposed_input_name = output_tensor_name + "transposed_input";
+  shared_ptr<ArtifactId> transposed_input =
+      genTransposeMIRtoACL(transposed_input_name, ir_input->getShape(), in_id);
+
+  const string layer_name = output_tensor_name + "_pooling_layer";
+
+  shared_ptr<ArtifactVariable> pad_stride_info_var = genPadStrideInfo(op, layer_name, _constrBlock);
+
+  shared_ptr<ArtifactId> pad_stride_info = pad_stride_info_var->use();
+
+  // Create kernel window info
+  shared_ptr<ArtifactVariable> kernel_window_var = _constrBlock->var(
+      "arm_compute::Size2D", layer_name + "_kernel_window", {},
+      {AF::lit(to_string(op.getWindowSize()[1])), AF::lit(to_string(op.getWindowSize()[0]))});
+  shared_ptr<ArtifactId> kernel_window = kernel_window_var->use();
+
+  // Create pooling info: pooling type, kernel info, strides, etc
+  shared_ptr<ArtifactVariable> pooling_info_var =
+      _constrBlock->var("arm_compute::PoolingLayerInfo", layer_name + "_pooling_info", {},
+                        {AF::lit(pooling_type), kernel_window, pad_stride_info,
+                         AF::lit(exclude_padding ? "true" : "false")});
+  shared_ptr<ArtifactId> pooling_info = pooling_info_var->use();
+
+  // Generate auxiliary tensor to hold transposed output of pool in NCHW format
+  Shape transposed_output_shape = transposeShape<0, 3, 1, 2>(ir_output->getShape());
+  shared_ptr<ArtifactId> transposed_output =
+      genTensor(layer_name + "_out_transpose", transposed_output_shape);
+
+  // Actual layer creation
+  shared_ptr<ArtifactId> layer =
+      genLayer("arm_compute::CLPoolingLayer", layer_name,
+               {AF::ref(transposed_input), AF::ref(transposed_output), pooling_info});
+  genTensorAllocation(_infBlock, transposed_output);
+  genLayerExecution(layer);
+
+  shared_ptr<ArtifactId> output =
+      genTransposeACLtoMIR(output_tensor_name, transposed_output_shape, transposed_output);
+
+  genTensorDeallocation(_infBlock, transposed_input);
+  genTensorDeallocation(_infBlock, transposed_output);
+}
+
+template <typename Op>
 void AclCppOpGenerator::genConvolution(Op &op, const string &acl_func_name, const string &suffix)
 {
   const auto *ir_input = op.getInput(0)->getProducer();
diff --git a/compiler/nnc/passes/acl_soft_backend/AclCppOpGenerator.h b/compiler/nnc/passes/acl_soft_backend/AclCppOpGenerator.h
index f6e004097..b67135b57 100644
--- a/compiler/nnc/passes/acl_soft_backend/AclCppOpGenerator.h
+++ b/compiler/nnc/passes/acl_soft_backend/AclCppOpGenerator.h
@@ -50,6 +50,7 @@ public:
    * @param op
    */
   void visit(mir::ops::AddOp &op) override;
+  void visit(mir::ops::AvgPool2DOp &op) override;
   void visit(mir::ops::CappedReluOp &op) override;
   void visit(mir::ops::ConcatOp &op) override;
   void visit(mir::ops::ConstantOp &op) override;
@@ -63,6 +64,7 @@ public:
   void visit(mir::ops::InputOp &op) override;
   void visit(mir::ops::LeakyReluOp &op) override;
   void visit(mir::ops::MaxOp &op) override;
+  void visit(mir::ops::MaxPool2DOp &op) override;
   void visit(mir::ops::MulOp &op) override;
   void visit(mir::ops::OutputOp &op) override;
   void visit(mir::ops::PadOp &op) override;
@@ -120,6 +122,9 @@ private:
   std::shared_ptr<ArtifactVariable> genPadStrideInfo(const Op &op, const std::string &prefix,
                                                      ArtifactBlock *block);
 
+  template <typename Op>
+  void genPooling(Op &op, const std::string &pooling_type, bool exclude_padding);
+
   /**
    * @brief The common part of the convolution and the depthwise convolution.
    */
diff --git a/compiler/nnc/passes/interpreter/Interpreter.cpp b/compiler/nnc/passes/interpreter/Interpreter.cpp
index f12f0e78b..48373a8d1 100644
--- a/compiler/nnc/passes/interpreter/Interpreter.cpp
+++ b/compiler/nnc/passes/interpreter/Interpreter.cpp
@@ -17,6 +17,7 @@
 #include "passes/interpreter/Interpreter.h"
 
 #include "ops/Add.h"
+#include "ops/AvgPool2D.h"
 #include "ops/Concat.h"
 #include "ops/Conv2D.h"
 #include "ops/DeConv2D.h"
@@ -25,6 +26,7 @@
 #include "ops/FullyConnected.h"
 #include "ops/Gather.h"
 #include "ops/Max.h"
+#include "ops/MaxPool2D.h"
 #include "ops/Mul.h"
 #include "ops/Pad.h"
 #include "ops/Pool.h"
@@ -84,6 +86,13 @@ void NNInterpreter::visit(ops::InputOp &op)
   setOutputTensors(op, {it->second});
 }
 
+void NNInterpreter::visit(ops::AvgPool2DOp &op)
+{
+  auto inputs = getInputTensors(op);
+  auto outputs = AvgPool2D(inputs[0], op)();
+  setOutputTensors(op, std::move(outputs));
+}
+
 void NNInterpreter::visit(ops::ConstantOp &op) { setOutputTensors(op, {op.getValue()}); }
 
 void NNInterpreter::visit(ops::ConcatOp &op)
@@ -100,6 +109,13 @@ void NNInterpreter::visit(ops::Conv2DOp &op)
   setOutputTensors(op, std::move(outputs));
 }
 
+void NNInterpreter::visit(ops::MaxPool2DOp &op)
+{
+  auto inputs = getInputTensors(op);
+  auto outputs = MaxPool2D(inputs[0], op)();
+  setOutputTensors(op, std::move(outputs));
+}
+
 void NNInterpreter::visit(ops::ReshapeOp &op)
 {
   auto inputs = getInputTensors(op);
diff --git a/compiler/nnc/passes/interpreter/ops/AvgPool2D.cpp b/compiler/nnc/passes/interpreter/ops/AvgPool2D.cpp
new file mode 100644
index 000000000..f96823cb5
--- /dev/null
+++ b/compiler/nnc/passes/interpreter/ops/AvgPool2D.cpp
@@ -0,0 +1,83 @@
+/*
+ * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "AvgPool2D.h"
+#include "common.h"
+
+#include "mir/ShapeRange.h"
+
+namespace nnc
+{
+
+using namespace mir;
+
+std::vector<TensorVariant> AvgPool2D::operator()()
+{
+  const auto &input_shape = _op.getInputShape(0);
+  const auto &output_shape = _op.getOutputShape(0);
+  const auto &window_size = _op.getWindowSize();
+  const auto &strides = _op.getStrides();
+  const auto &padding_before = _op.getPaddingBefore();
+  const auto &padding_after = _op.getPaddingAfter();
+
+  constexpr int num_spatial_dims = 2;
+  assert(_input.getShape().rank() == 4);
+  assert(window_size.size() == num_spatial_dims);
+  assert(strides.size() == num_spatial_dims);
+  assert(padding_before.size() == num_spatial_dims);
+  assert(padding_after.size() == num_spatial_dims);
+
+  auto res = allocate_tensor(output_shape);
+  Tensor<float> res_accessor(res);
+
+  ShapeRange in_range(input_shape);
+  Index in_index(input_shape.rank());
+
+  for (const auto &out_index : ShapeRange(output_shape))
+  {
+    float result = 0.0f;
+    size_t num_elements = 0;
+
+    // Assuming NHWC format.
+    in_index.at(0) = out_index.at(0);
+    in_index.at(3) = out_index.at(3);
+
+    for (const auto &window_index : ShapeRange(Shape(window_size)))
+    {
+      // Assuming NHWC format.
+      for (int i = 0; i < num_spatial_dims; ++i)
+        in_index.at(1 + i) =
+            out_index.at(1 + i) * strides[i] + window_index.at(i) - padding_before[i];
+
+      if (in_range.contains(in_index))
+      {
+        num_elements++;
+        result += _input.at(in_index);
+      }
+      else if (_op.getIncludePad())
+      {
+        num_elements++;
+      }
+    }
+
+    result /= num_elements;
+    res_accessor.at(out_index) = result;
+  }
+
+  return {res};
+}
+
+} // namespace nnc
diff --git a/compiler/nnc/passes/interpreter/ops/AvgPool2D.h b/compiler/nnc/passes/interpreter/ops/AvgPool2D.h
new file mode 100644
index 000000000..f6f175057
--- /dev/null
+++ b/compiler/nnc/passes/interpreter/ops/AvgPool2D.h
@@ -0,0 +1,44 @@
+/*
+ * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef _NNC_CORE_BACKEND_INTERPRETER_AVG_POOL_2D_
+#define _NNC_CORE_BACKEND_INTERPRETER_AVG_POOL_2D_
+
+#include "OperationImpl.h"
+#include "mir/ops/AvgPool2DOp.h"
+#include "mir/Tensor.h"
+
+namespace nnc
+{
+
+class AvgPool2D : public OperationImpl<float>
+{
+public:
+  std::vector<mir::TensorVariant> operator()() override;
+
+  AvgPool2D(const mir::TensorVariant &input, const mir::ops::AvgPool2DOp &op)
+      : _op(op), _input(input)
+  {
+  }
+
+private:
+  const mir::ops::AvgPool2DOp &_op;
+  const mir::Tensor<float> _input;
+};
+
+} // namespace nnc
+
+#endif //_NNC_CORE_BACKEND_INTERPRETER_AVG_POOL_2D_
diff --git a/compiler/nnc/passes/interpreter/ops/MaxPool2D.cpp b/compiler/nnc/passes/interpreter/ops/MaxPool2D.cpp
new file mode 100644
index 000000000..22eb369e2
--- /dev/null
+++ b/compiler/nnc/passes/interpreter/ops/MaxPool2D.cpp
@@ -0,0 +1,78 @@
+/*
+ * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "MaxPool2D.h"
+#include "common.h"
+
+#include "mir/ShapeRange.h"
+
+#include <limits>
+
+namespace nnc
+{
+
+using namespace mir;
+
+std::vector<TensorVariant> MaxPool2D::operator()()
+{
+  const auto &input_shape = _op.getInputShape(0);
+  const auto &output_shape = _op.getOutputShape(0);
+  const auto &window_size = _op.getWindowSize();
+  const auto &strides = _op.getStrides();
+  const auto &padding_before = _op.getPaddingBefore();
+  const auto &padding_after = _op.getPaddingAfter();
+
+  constexpr int num_spatial_dims = 2;
+  assert(_input.getShape().rank() == 4);
+  assert(window_size.size() == num_spatial_dims);
+  assert(strides.size() == num_spatial_dims);
+  assert(padding_before.size() == num_spatial_dims);
+  assert(padding_after.size() == num_spatial_dims);
+
+  auto res = allocate_tensor(output_shape);
+  Tensor<float> res_accessor(res);
+
+  ShapeRange in_range(input_shape);
+  Index in_index(input_shape.rank());
+
+  for (const auto &out_index : ShapeRange(output_shape))
+  {
+    float result = std::numeric_limits<float>::lowest();
+
+    // Assuming NHWC format.
+    in_index.at(0) = out_index.at(0);
+    in_index.at(3) = out_index.at(3);
+
+    for (const auto &window_index : ShapeRange(Shape(window_size)))
+    {
+      // Assuming NHWC format.
+      for (int i = 0; i < num_spatial_dims; ++i)
+        in_index.at(1 + i) =
+            out_index.at(1 + i) * strides[i] + window_index.at(i) - padding_before[i];
+
+      if (in_range.contains(in_index))
+      {
+        result = std::max(result, _input.at(in_index));
+      }
+    }
+
+    res_accessor.at(out_index) = result;
+  }
+
+  return {res};
+}
+
+} // namespace nnc
diff --git a/compiler/nnc/passes/interpreter/ops/MaxPool2D.h b/compiler/nnc/passes/interpreter/ops/MaxPool2D.h
new file mode 100644
index 000000000..9950af102
--- /dev/null
+++ b/compiler/nnc/passes/interpreter/ops/MaxPool2D.h
@@ -0,0 +1,44 @@
+/*
+ * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef _NNC_CORE_BACKEND_INTERPRETER_MAX_POOL_2D_
+#define _NNC_CORE_BACKEND_INTERPRETER_MAX_POOL_2D_
+
+#include "OperationImpl.h"
+#include "mir/ops/MaxPool2DOp.h"
+#include "mir/Tensor.h"
+
+namespace nnc
+{
+
+class MaxPool2D : public OperationImpl<float>
+{
+public:
+  std::vector<mir::TensorVariant> operator()() override;
+
+  MaxPool2D(const mir::TensorVariant &input, const mir::ops::MaxPool2DOp &op)
+      : _op(op), _input(input)
+  {
+  }
+
+private:
+  const mir::ops::MaxPool2DOp &_op;
+  const mir::Tensor<float> _input;
+};
+
+} // namespace nnc
+
+#endif //_NNC_CORE_BACKEND_INTERPRETER_MAX_POOL_2D_
diff --git a/compiler/nnc/passes/optimizations/SinkRelu.cpp b/compiler/nnc/passes/optimizations/SinkRelu.cpp
index f0d428061..0e2f47820 100644
--- a/compiler/nnc/passes/optimizations/SinkRelu.cpp
+++ b/compiler/nnc/passes/optimizations/SinkRelu.cpp
@@ -40,10 +40,7 @@ PassData SinkRelu::run(PassData data)
   auto is_relu = [](const Operation *op) { return op->getType() == Operation::Type::ReLU; };
   auto is_concat = [](const Operation *op) { return op->getType() == Operation::Type::concat; };
   auto is_max_pool = [](const Operation *op) {
-    auto *p_op = dynamic_cast<const ops::PoolOp *>(op);
-    if (!p_op)
-      return false;
-    return p_op->getPoolingType() == ops::PoolOp::PoolingType::MAX;
+    return op->getType() == Operation::Type::maxPool2D;
   };
   std::vector<std::pair<Operation *, Operation *>> matches;
 
diff --git a/compiler/nnc/passes/soft_backend/ModelAnalyzer.cpp b/compiler/nnc/passes/soft_backend/ModelAnalyzer.cpp
index 2518d4116..88d628232 100644
--- a/compiler/nnc/passes/soft_backend/ModelAnalyzer.cpp
+++ b/compiler/nnc/passes/soft_backend/ModelAnalyzer.cpp
@@ -337,6 +337,10 @@ void ModelAnalyzer::visit(ops::PoolOp &op)
   appendOperationToInference(&op, func_name);
 }
 
+void ModelAnalyzer::visit(ops::AvgPool2DOp &op) { appendOperationToInference(&op, "avgPool"); }
+
+void ModelAnalyzer::visit(ops::MaxPool2DOp &op) { appendOperationToInference(&op, "maxPool"); }
+
 void ModelAnalyzer::visit(ops::FullyConnectedOp &op)
 {
   appendOperationToInference(&op, "fullConnect");
diff --git a/compiler/nnc/passes/soft_backend/ModelAnalyzer.h b/compiler/nnc/passes/soft_backend/ModelAnalyzer.h
index 71eb44d71..d7f9d9c70 100644
--- a/compiler/nnc/passes/soft_backend/ModelAnalyzer.h
+++ b/compiler/nnc/passes/soft_backend/ModelAnalyzer.h
@@ -48,6 +48,7 @@ public:
   void analyze(const mir::Graph *g);
 
   void visit(mir::ops::AddOp &op) override;
+  void visit(mir::ops::AvgPool2DOp &op) override;
   void visit(mir::ops::CappedReluOp &op) override;
   void visit(mir::ops::ConcatOp &op) override;
   void visit(mir::ops::ConstantOp &op) override;
@@ -61,6 +62,7 @@ public:
   void visit(mir::ops::InputOp &op) override;
   void visit(mir::ops::LeakyReluOp &op) override;
   void visit(mir::ops::MaxOp &op) override;
+  void visit(mir::ops::MaxPool2DOp &op) override;
   void visit(mir::ops::MulOp &op) override;
   void visit(mir::ops::OutputOp &op) override;
   void visit(mir::ops::PadOp &op) override;
diff --git a/compiler/nnc/passes/soft_backend/SBSerializer.cpp b/compiler/nnc/passes/soft_backend/SBSerializer.cpp
index 054687316..c32e08854 100644
--- a/compiler/nnc/passes/soft_backend/SBSerializer.cpp
+++ b/compiler/nnc/passes/soft_backend/SBSerializer.cpp
@@ -189,6 +189,42 @@ void Serializer::visit(ops::PoolOp &op)
   serializeShape(op.getOutputShape(0));
 }
 
+void Serializer::visit(ops::AvgPool2DOp &op)
+{
+  _curOp->paramStartOffset = _buffer.size();
+  // serialize window shape
+  serializeShape(Shape(op.getWindowSize()));
+  // serialize strindes
+  serializeShape(Shape(op.getStrides()));
+  // serialize pads
+  int32_t number_of_pads = 2; // windowShape.rank();
+  serializePads(op, number_of_pads);
+  // serialize border type
+  PoolBorderType border_type =
+      op.getIncludePad() ? PoolBorderType::ZEROFILLED : PoolBorderType::EMPTY;
+
+  serializeT<int32_t>(etoi(border_type));
+  // serialize output shape
+  serializeShape(op.getOutputShape(0));
+}
+
+void Serializer::visit(ops::MaxPool2DOp &op)
+{
+  _curOp->paramStartOffset = _buffer.size();
+  // serialize window shape
+  serializeShape(Shape(op.getWindowSize()));
+  // serialize strindes
+  serializeShape(Shape(op.getStrides()));
+  // serialize pads
+  int32_t number_of_pads = 2; // windowShape.rank();
+  serializePads(op, number_of_pads);
+  // serialize border type
+  // TODO This parameter is not used for max pooling, remove it.
+  serializeT<int32_t>(etoi(PoolBorderType::EMPTY));
+  // serialize output shape
+  serializeShape(op.getOutputShape(0));
+}
+
 void Serializer::visit(ops::FullyConnectedOp &op)
 {
   _curOp->paramStartOffset = _buffer.size();
diff --git a/compiler/nnc/passes/soft_backend/SBSerializer.h b/compiler/nnc/passes/soft_backend/SBSerializer.h
index 5a64fb825..3c7fe709d 100644
--- a/compiler/nnc/passes/soft_backend/SBSerializer.h
+++ b/compiler/nnc/passes/soft_backend/SBSerializer.h
@@ -42,6 +42,7 @@ class Serializer : public mir::Visitor
 {
 public:
   void visit(mir::ops::AddOp &op) override;
+  void visit(mir::ops::AvgPool2DOp &op) override;
   void visit(mir::ops::CappedReluOp &op) override;
   void visit(mir::ops::ConcatOp &op) override;
   void visit(mir::ops::ConstantOp &op) override;
@@ -55,6 +56,7 @@ public:
   void visit(mir::ops::InputOp &op) override;
   void visit(mir::ops::LeakyReluOp &op) override;
   void visit(mir::ops::MaxOp &op) override;
+  void visit(mir::ops::MaxPool2DOp &op) override;
   void visit(mir::ops::MulOp &op) override;
   void visit(mir::ops::OutputOp &op) override;
   void visit(mir::ops::PadOp &op) override;
diff --git a/compiler/nnc/unittests/acl_backend/MIRToDOM.cpp b/compiler/nnc/unittests/acl_backend/MIRToDOM.cpp
index c1617f157..dcb88892d 100644
--- a/compiler/nnc/unittests/acl_backend/MIRToDOM.cpp
+++ b/compiler/nnc/unittests/acl_backend/MIRToDOM.cpp
@@ -37,6 +37,7 @@
 #include "mir/ops/EluOp.h"
 #include "mir/ops/FullyConnectedOp.h"
 #include "mir/ops/InputOp.h"
+#include "mir/ops/MaxPool2DOp.h"
 #include "mir/ops/OutputOp.h"
 #include "mir/ops/PadOp.h"
 #include "mir/ops/PoolOp.h"
@@ -365,16 +366,15 @@ TEST(acl_backend_mir_to_dom, fully_connected)
 
 TEST(acl_backend_mir_to_dom, maxpool)
 {
-  mir::Shape window_shape{3, 3}; // Height, Width
-  mir::Shape strides{1, 1};
+  vector<int32_t> window_size{3, 3}; // Height, Width
+  vector<int32_t> strides{1, 1};
 
   Graph g;
   OpConstructor op_generator =
-      [window_shape, strides](mir::Graph &g, const std::vector<mir::Operation::Output *> &inputs) {
+      [window_size, strides](mir::Graph &g, const std::vector<mir::Operation::Output *> &inputs) {
         std::vector<int32_t> padding{0, 0};
-        return g.create<mir::ops::PoolOp>(inputs[0], ops::PoolOp::PoolingType::MAX, window_shape,
-                                          strides, padding, padding,
-                                          mir::ops::PoolOp::BorderType::EMPTY);
+        return g.create<mir::ops::MaxPool2DOp>(inputs[0], window_size, strides, padding, padding,
+                                               mir::DataFormat::NHWC);
       };
 
   vector<Shape> input_shapes{{1, 10, 10, 3}};
diff --git a/compiler/nnc/unittests/optimizations/SinkTest.cpp b/compiler/nnc/unittests/optimizations/SinkTest.cpp
index 6344bf462..38f895a11 100644
--- a/compiler/nnc/unittests/optimizations/SinkTest.cpp
+++ b/compiler/nnc/unittests/optimizations/SinkTest.cpp
@@ -165,9 +165,9 @@ TEST(OptPass, sinkPoolReLU)
    */
   Operation *input = g.create<ops::InputOp>(Shape{1, 4, 4, 3});
   Operation *relu = g.create<ops::ReluOp>(input->getOutput(0));
-  Operation *mp = g.create<ops::PoolOp>(relu->getOutput(0), ops::PoolOp::PoolingType::MAX,
-                                        Shape{2, 2}, Shape{2, 2}, vector<int32_t>{0, 0},
-                                        vector<int32_t>{0, 0}, ops::PoolOp::BorderType::EMPTY);
+  Operation *mp =
+      g.create<ops::MaxPool2DOp>(relu->getOutput(0), vector<int32_t>{2, 2}, vector<int32_t>{2, 2},
+                                 vector<int32_t>{0, 0}, vector<int32_t>{0, 0}, DataFormat::NHWC);
   Operation *tanh = g.create<ops::TanhOp>(mp->getOutput(0));
   Operation *out = g.create<ops::OutputOp>(tanh->getOutput(0));
   (void)out;
@@ -179,7 +179,7 @@ TEST(OptPass, sinkPoolReLU)
   g.accept(&d);
 
   // tanh(relu(pool(input)))
-  ASSERT_EQ(getNext(g.getInputs()[0])->getType(), mir::Operation::Type::pool);
+  ASSERT_EQ(getNext(g.getInputs()[0])->getType(), mir::Operation::Type::maxPool2D);
   ASSERT_EQ(getPrev(g.getOutputs()[0])->getType(), mir::Operation::Type::tanh);
   ASSERT_EQ("i_0.p_5.r_6.th_3.", ss.str());
 }
diff --git a/compiler/nnc/unittests/optimizations/Util.h b/compiler/nnc/unittests/optimizations/Util.h
index de1a80c42..2b682ac71 100644
--- a/compiler/nnc/unittests/optimizations/Util.h
+++ b/compiler/nnc/unittests/optimizations/Util.h
@@ -16,16 +16,19 @@
 
 #ifndef NNCC_UTIL_H
 #define NNCC_UTIL_H
-#include "mir/ops/TransposeOp.h"
-#include "mir/ops/ReluOp.h"
-#include "mir/ops/ConstantOp.h"
-#include "mir/ops/TanhOp.h"
+
+#include "mir/ops/AddOp.h"
+#include "mir/ops/AvgPool2DOp.h"
 #include "mir/ops/ConcatOp.h"
+#include "mir/ops/ConstantOp.h"
+#include "mir/ops/Conv2DOp.h"
+#include "mir/ops/MaxPool2DOp.h"
+#include "mir/ops/MulOp.h"
 #include "mir/ops/OutputOp.h"
 #include "mir/ops/PoolOp.h"
-#include "mir/ops/AddOp.h"
-#include "mir/ops/MulOp.h"
-#include "mir/ops/Conv2DOp.h"
+#include "mir/ops/ReluOp.h"
+#include "mir/ops/TanhOp.h"
+#include "mir/ops/TransposeOp.h"
 #include "mir/Visitor.h"
 
 namespace nnc
@@ -48,6 +51,16 @@ public:
 
   void visit(mir::ops::PoolOp &op) override { _s << "p_" << std::to_string(op.getId()) << "."; }
 
+  void visit(mir::ops::AvgPool2DOp &op) override
+  {
+    _s << "p_" << std::to_string(op.getId()) << ".";
+  }
+
+  void visit(mir::ops::MaxPool2DOp &op) override
+  {
+    _s << "p_" << std::to_string(op.getId()) << ".";
+  }
+
   void visit(mir::ops::TransposeOp &op) override
   {
     _s << "t_" << std::to_string(op.getId()) << ".";
diff --git a/compiler/nnc/unittests/soft_backend/CPPOperations.cpp b/compiler/nnc/unittests/soft_backend/CPPOperations.cpp
index eb6c7c271..c97dd7405 100644
--- a/compiler/nnc/unittests/soft_backend/CPPOperations.cpp
+++ b/compiler/nnc/unittests/soft_backend/CPPOperations.cpp
@@ -57,6 +57,7 @@
 
 // operations part
 #include "mir/ops/AddOp.h"
+#include "mir/ops/AvgPool2DOp.h"
 #include "mir/ops/CappedReluOp.h"
 #include "mir/ops/ConcatOp.h"
 #include "mir/ops/Conv2DOp.h"
@@ -68,6 +69,7 @@
 #include "mir/ops/InputOp.h"
 #include "mir/ops/LeakyReluOp.h"
 #include "mir/ops/MaxOp.h"
+#include "mir/ops/MaxPool2DOp.h"
 #include "mir/ops/MulOp.h"
 #include "mir/ops/OutputOp.h"
 #include "mir/ops/PadOp.h"
@@ -691,18 +693,7 @@ TEST(cpp_operations_test, resize_NN_test_scales)
   }
 }
 
-template <irOps::PoolOp::PoolingType poolT>
-static mir::Operation *
-createPool(mir::Graph &g, const std::vector<mir::Operation::Output *> &inputs,
-           mir::Shape &window_shape, mir::Shape &strides, irOps::PoolOp::BorderType border)
-{
-  std::vector<int32_t> padding{0, 0};
-  return g.create<mir::ops::PoolOp>(inputs[0], poolT, window_shape, strides, padding, padding,
-                                    border);
-};
-
-template <irOps::PoolOp::PoolingType poolT, typename Func>
-static void genericPoolTest(Func test_func, const vector<irOps::PoolOp::BorderType> borders)
+TEST(cpp_operations_test, avgpool)
 {
   // Iterate over window width, window height
   // channels
@@ -717,35 +708,56 @@ static void genericPoolTest(Func test_func, const vector<irOps::PoolOp::BorderTy
           for (iT stride_w = 1; stride_w <= 3; ++stride_w)
           {
             vector<int> shape_data{3, 5, 7, static_cast<int>(channels)};
-            mir::Shape window_shape{windowH, windowW};
-            mir::Shape strides{stride_h, stride_w};
+            vector<int32_t> window_size{windowH, windowW};
+            vector<int32_t> strides{stride_h, stride_w};
             Tensor input_atensor;
             vector<unique_ptr<mir::TensorVariant>> input_ntensors(1);
             fillTensors(input_ntensors[0], input_atensor, shape_data, 1.0f);
 
-            for (auto border : borders)
+            for (const auto include_pad : {false, true})
             {
-              using namespace std::placeholders;
-
-              auto op_generator =
-                  std::bind(createPool<poolT>, _1, _2, window_shape, strides, border);
+              auto op_generator = [&window_size, &strides, include_pad](
+                  mir::Graph &g, const std::vector<mir::Operation::Output *> &inputs) {
+                std::vector<int32_t> padding{0, 0};
+                return g.create<mir::ops::AvgPool2DOp>(inputs[0], window_size, strides, padding,
+                                                       padding, include_pad, mir::DataFormat::NHWC);
+              };
 
-              createAndRunTestGraph(op_generator, test_func, input_ntensors, input_atensor);
+              createAndRunTestGraph(op_generator, avgPool, input_ntensors, input_atensor);
             }
           }
-};
+}
 
 TEST(cpp_operations_test, maxpool)
 {
-  vector<irOps::PoolOp::BorderType> border_types = {irOps::PoolOp::BorderType::EMPTY};
-  genericPoolTest<mir::ops::PoolOp::PoolingType::MAX>(maxPool, border_types);
-}
+  // Iterate over window width, window height
+  // channels
+  // stride width, stride height
+  // size 3 is chosen to cover all cases, where width bigger/smaller then height and equal/not equal
+  // to 1
+  using iT = int32_t;
+  for (iT windowH = 1; windowH <= 3; ++windowH)
+    for (iT windowW = 1; windowW <= 3; ++windowW)
+      for (iT channels = 1; channels <= 2; ++channels)
+        for (iT stride_h = 1; stride_h <= 3; ++stride_h)
+          for (iT stride_w = 1; stride_w <= 3; ++stride_w)
+          {
+            vector<int> shape_data{3, 5, 7, static_cast<int>(channels)};
+            vector<int32_t> window_size{windowH, windowW};
+            vector<int32_t> strides{stride_h, stride_w};
+            Tensor input_atensor;
+            vector<unique_ptr<mir::TensorVariant>> input_ntensors(1);
+            fillTensors(input_ntensors[0], input_atensor, shape_data, 1.0f);
 
-TEST(cpp_operations_test, avgpool)
-{
-  vector<irOps::PoolOp::BorderType> border_types = {irOps::PoolOp::BorderType::EMPTY,
-                                                    irOps::PoolOp::BorderType::ZEROFILLED};
-  genericPoolTest<mir::ops::PoolOp::PoolingType::AVG>(avgPool, border_types);
+            auto op_generator = [&window_size, &strides](
+                mir::Graph &g, const std::vector<mir::Operation::Output *> &inputs) {
+              std::vector<int32_t> padding{0, 0};
+              return g.create<mir::ops::MaxPool2DOp>(inputs[0], window_size, strides, padding,
+                                                     padding, mir::DataFormat::NHWC);
+            };
+
+            createAndRunTestGraph(op_generator, maxPool, input_ntensors, input_atensor);
+          }
 }
 
 TEST(cpp_operations_test, relu)