[mir_tflite] Switch to AvgPool2D and MaxPool2D ops (#7024)

Switch from `PoolOp` to `AvgPool2DOp` and `MaxPool2DOp`. The former ones are deprecated.

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

diff --git a/compiler/mir-tflite-importer/tflite_op_creator.cpp b/compiler/mir-tflite-importer/tflite_op_creator.cpp
index f74696d..3f4232b 100644 (file)
--- a/compiler/mir-tflite-importer/tflite_op_creator.cpp
+++ b/compiler/mir-tflite-importer/tflite_op_creator.cpp
@@ -18,6 +18,7 @@
 #include "schema_generated.h"
 
 #include "mir/ops/AddOp.h"
+#include "mir/ops/AvgPool2DOp.h"
 #include "mir/ops/CappedReluOp.h"
 #include "mir/ops/ConcatOp.h"
 #include "mir/ops/ConstantOp.h"
@@ -28,9 +29,9 @@
 #include "mir/ops/FullyConnectedOp.h"
 #include "mir/ops/LeakyReluOp.h"
 #include "mir/ops/MaxOp.h"
+#include "mir/ops/MaxPool2DOp.h"
 #include "mir/ops/MulOp.h"
 #include "mir/ops/PadOp.h"
-#include "mir/ops/PoolOp.h"
 #include "mir/ops/ReduceMeanOp.h"
 #include "mir/ops/ReluOp.h"
 #include "mir/ops/ReshapeOp.h"
@@ -54,25 +55,33 @@ namespace mir_tflite
 {
 
 static void calculatePadding(tflite::Padding padding, const Shape &input_shape,
-                             const Shape &window_shape, const Shape &strides,
-                             std::vector<int32_t> &padding_before,
-                             std::vector<int32_t> &padding_after)
+                             const std::vector<std::int32_t> &window_size,
+                             const std::vector<std::int32_t> &strides,
+                             std::vector<std::int32_t> &padding_before,
+                             std::vector<std::int32_t> &padding_after)
 {
+  constexpr int num_spatial_dims = 2;
+  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);
+
   switch (padding)
   {
     case tflite::Padding_SAME:
-      for (int i = 0; i < 2; ++i)
+      for (int i = 0; i < num_spatial_dims; ++i)
       {
-        int32_t padding;
-        padding = (input_shape.dim(1 + i) % strides.dim(i) == 0)
-                      ? std::max(0, window_shape.dim(i) - strides.dim(i))
-                      : std::max(0, window_shape.dim(i) - input_shape.dim(1 + i) % strides.dim(i));
-        padding_before[i] = padding / 2;
-        padding_after[i] = padding - padding_before[i];
+        // Assuming NHWC format.
+        const std::int32_t total_padding =
+            (input_shape.dim(1 + i) % strides[i] == 0)
+                ? std::max(0, window_size[i] - strides[i])
+                : std::max(0, window_size[i] - input_shape.dim(1 + i) % strides[i]);
+        padding_before[i] = total_padding / 2;
+        padding_after[i] = total_padding - padding_before[i];
       }
       break;
     case tflite::Padding_VALID:
-      for (int i = 0; i < 2; ++i)
+      for (int i = 0; i < num_spatial_dims; ++i)
       {
         padding_before[i] = 0;
         padding_after[i] = 0;
@@ -110,18 +119,19 @@ TFLiteOpCreator::convertConv2D(const Conv2DOptions *opts,
 
   kernel = createOp<ops::ConstantOp>(extractTensor(kernel))->getOutput(0);
 
-  Shape strides{opts->stride_h(), opts->stride_w()};
-  std::vector<int32_t> padding_before(2);
-  std::vector<int32_t> padding_after(2);
+  const std::vector<std::int32_t> strides{opts->stride_h(), opts->stride_w()};
+  std::vector<std::int32_t> padding_before(2);
+  std::vector<std::int32_t> padding_after(2);
 
   const auto &input_shape = input->getShape();
   const auto &kernel_shape = kernel->getShape();
-  Shape window_shape{kernel_shape.dim(1), kernel_shape.dim(2)};
-  calculatePadding(opts->padding(), input_shape, window_shape, strides, padding_before,
+  std::vector<std::int32_t> kernel_size{kernel_shape.dim(1), kernel_shape.dim(2)};
+  calculatePadding(opts->padding(), input_shape, kernel_size, strides, padding_before,
                    padding_after);
 
   auto result =
-      createOp<ops::Conv2DOp>(input, kernel, strides, padding_before, padding_after)->getOutput(0);
+      createOp<ops::Conv2DOp>(input, kernel, Shape(strides), padding_before, padding_after)
+          ->getOutput(0);
   result = createOp<ops::AddOp>(result, bias)->getOutput(0);
   return {addFusedActivation(result, opts->fused_activation_function())};
 }
@@ -138,18 +148,18 @@ TFLiteOpCreator::convertDepthwiseConv2D(const DepthwiseConv2DOptions *opts,
   const std::vector<std::size_t> axis_order{1, 2, 3, 0};
   kernel = createOp<ops::TransposeOp>(kernel, axis_order)->getOutput(0);
 
-  Shape strides{opts->stride_h(), opts->stride_w()};
-  std::vector<int32_t> padding_before(2);
-  std::vector<int32_t> padding_after(2);
+  const std::vector<std::int32_t> strides{opts->stride_h(), opts->stride_w()};
+  std::vector<std::int32_t> padding_before(2);
+  std::vector<std::int32_t> padding_after(2);
 
   const auto &input_shape = input->getShape();
   const auto &kernel_shape = kernel->getShape();
-  Shape window_shape{kernel_shape.dim(0), kernel_shape.dim(1)};
-  calculatePadding(opts->padding(), input_shape, window_shape, strides, padding_before,
+  std::vector<std::int32_t> kernel_size{kernel_shape.dim(0), kernel_shape.dim(1)};
+  calculatePadding(opts->padding(), input_shape, kernel_size, strides, padding_before,
                    padding_after);
 
   auto result =
-      createOp<ops::DepthwiseConv2DOp>(input, kernel, strides, padding_before, padding_after)
+      createOp<ops::DepthwiseConv2DOp>(input, kernel, Shape(strides), padding_before, padding_after)
           ->getOutput(0);
   result = createOp<ops::AddOp>(result, bias)->getOutput(0);
   return {addFusedActivation(result, opts->fused_activation_function())};
@@ -170,17 +180,16 @@ TFLiteOpCreator::convertMaxPool2D(const ::tflite::Pool2DOptions *opts,
   auto input = inputs.at(0);
 
   const auto &input_shape = input->getShape();
-  Shape window_shape{opts->filter_height(), opts->filter_width()};
-  Shape strides{opts->stride_h(), opts->stride_w()};
-  std::vector<int32_t> padding_before(2);
-  std::vector<int32_t> padding_after(2);
+  const std::vector<std::int32_t> window_size{opts->filter_height(), opts->filter_width()};
+  const std::vector<std::int32_t> strides{opts->stride_h(), opts->stride_w()};
+  std::vector<std::int32_t> padding_before(2);
+  std::vector<std::int32_t> padding_after(2);
 
-  calculatePadding(opts->padding(), input_shape, window_shape, strides, padding_before,
+  calculatePadding(opts->padding(), input_shape, window_size, strides, padding_before,
                    padding_after);
 
-  auto result =
-      createOp<ops::PoolOp>(input, ops::PoolOp::PoolingType::MAX, window_shape, strides,
-                            padding_before, padding_after, ops::PoolOp::BorderType::EMPTY);
+  auto result = createOp<ops::MaxPool2DOp>(input, window_size, strides, padding_before,
+                                           padding_after, mir::DataFormat::NHWC);
   return {addFusedActivation(result->getOutput(0), opts->fused_activation_function())};
 }
 
@@ -191,17 +200,16 @@ TFLiteOpCreator::convertAveragePool2D(const ::tflite::Pool2DOptions *opts,
   auto input = inputs.at(0);
 
   const auto &input_shape = input->getShape();
-  Shape window_shape{opts->filter_height(), opts->filter_width()};
-  Shape strides{opts->stride_h(), opts->stride_w()};
-  std::vector<int32_t> padding_before(2);
-  std::vector<int32_t> padding_after(2);
+  const std::vector<std::int32_t> window_size{opts->filter_height(), opts->filter_width()};
+  const std::vector<std::int32_t> strides{opts->stride_h(), opts->stride_w()};
+  std::vector<std::int32_t> padding_before(2);
+  std::vector<std::int32_t> padding_after(2);
 
-  calculatePadding(opts->padding(), input_shape, window_shape, strides, padding_before,
+  calculatePadding(opts->padding(), input_shape, window_size, strides, padding_before,
                    padding_after);
 
-  auto result =
-      createOp<ops::PoolOp>(input, ops::PoolOp::PoolingType::AVG, window_shape, strides,
-                            padding_before, padding_after, ops::PoolOp::BorderType::EMPTY);
+  auto result = createOp<ops::AvgPool2DOp>(input, window_size, strides, padding_before,
+                                           padding_after, false, mir::DataFormat::NHWC);
   return {addFusedActivation(result->getOutput(0), opts->fused_activation_function())};
 }