Maxpool-8 shell (openvinotoolkit#6332)

* Initial version of v8::MaxPool op class

* Type instead of Type_t to indicate element type

* Attribute visitor test

* Common MaxPoolBase base class

* More refactoring

* v8::MaxPool cleanup

* Pooling ops inference helper extension - window dilation

* New MaxPool 3D type prop tests

* Common part of MaxPool validation part extracted to the base class

* MaxPool-8 shape inference with base class utils

* infer_batched_pooling_forward arguments reorder to avoid compilation errors

* Align the rounding type attribute name for both MaxPool version

* MaxPool-8 axis attribute

* Missing attributes

* Code formatting

* PR feedback

* MaxPool-1 RTTI definition adjustment

ngraph/core/include/ngraph/op/max_pool.hpp

@@ -4,8 +4,9 @@
 
 #pragma once
 
-#include "ngraph/op/op.hpp"
-#include "ngraph/op/util/attr_types.hpp"
+#include <limits>
+
+#include "ngraph/op/util/max_pool_base.hpp"
 
 namespace ngraph
 {
@@ -14,7 +15,7 @@ namespace ngraph
         namespace v1
         {
             /// \brief Batched max pooling operation.
-            class NGRAPH_API MaxPool : public Op
+            class NGRAPH_API MaxPool : public op::util::MaxPoolBase
             {
             public:
                 NGRAPH_RTTI_DECLARATION;
@@ -29,44 +30,23 @@ namespace ngraph
                 /// \param pads_begin The beginning of padding shape.
                 /// \param pads_end The end of padding shape.
                 /// \param kernel The kernel shape.
-                /// \param rounding_mode Whether to use ceiling or floor rounding type while
+                /// \param rounding_type Whether to use ceiling or floor rounding type while
                 /// computing output shape.
                 /// \param auto_pad The pad type for automatically computing padding sizes.
                 MaxPool(const Output<Node>& arg,
                         const Strides& strides,
                         const Shape& pads_begin,
                         const Shape& pads_end,
                         const Shape& kernel,
-                        op::RoundingType rounding_mode = op::RoundingType::FLOOR,
-                        const PadType& auto_pad = op::PadType::EXPLICIT);
+                        const op::RoundingType rounding_type = op::RoundingType::FLOOR,
+                        const PadType auto_pad = op::PadType::EXPLICIT);
 
                 bool visit_attributes(AttributeVisitor& visitor) override;
                 void validate_and_infer_types() override;
 
                 virtual std::shared_ptr<Node>
                     clone_with_new_inputs(const OutputVector& new_args) const override;
 
-                /// \return The kernel shape.
-                const Shape& get_kernel() const { return m_kernel; }
-                void set_kernel(const Shape& kernel) { m_kernel = kernel; }
-                /// \return The strides.
-                const Strides& get_strides() const { return m_strides; }
-                void set_strides(const Strides& strides) { m_strides = strides; }
-                /// \return The beginning of padding shape.
-                const Shape& get_pads_begin() const { return m_pads_begin; }
-                void set_pads_begin(const Shape& pads_begin) { m_pads_begin = pads_begin; }
-                /// \return The end of padding shape.
-                const Shape& get_pads_end() const { return m_pads_end; }
-                void set_adding_above(const Shape& pads_end) { m_pads_end = pads_end; }
-                /// \return The pad type for pooling.
-                const PadType& get_auto_pad() const { return m_auto_pad; }
-                void set_auto_pad(const PadType& auto_pad) { m_auto_pad = auto_pad; }
-                /// \return The ceiling mode being used for output shape computations
-                op::RoundingType get_rounding_type() const { return m_rounding_type; }
-                void set_rounding_type(op::RoundingType rounding_mode)
-                {
-                    m_rounding_type = rounding_mode;
-                }
                 /// \return The default value for MaxPool.
                 NGRAPH_SUPPRESS_DEPRECATED_START
                 virtual std::shared_ptr<Node> get_default_value() const override;
@@ -76,21 +56,85 @@ namespace ngraph
                               const HostTensorVector& inputs) const override;
                 bool has_evaluate() const override;
 
-            protected:
-                Shape m_kernel;
-                Strides m_strides;
-                Shape m_pads_begin;
-                Shape m_pads_end;
-                PadType m_auto_pad;
-                op::RoundingType m_rounding_type;
-
             private:
-                bool update_auto_padding(const PartialShape& in_shape,
-                                         Shape& new_pads_end,
-                                         Shape& new_pads_begin) const;
                 bool evaluate_maxpool(const HostTensorVector& outputs,
                                       const HostTensorVector& inputs) const;
             };
         } // namespace v1
+
+        namespace v8
+        {
+            /// \brief MaxPooling operation with values and indices calculated as individual outputs
+            class NGRAPH_API MaxPool : public op::util::MaxPoolBase
+            {
+            public:
+                NGRAPH_RTTI_DECLARATION;
+
+                /// \brief Constructs an empty MaxPool operation.
+                MaxPool() = default;
+
+                /// \brief Constructs a parametrized MaxPool operation.
+                ///
+                /// \param arg Output of a node producing the feature tensor to be pooled.
+                /// \param strides The strides of the pooling filter.
+                /// \param dilations The dilations of the pooling filter.
+                /// \param pads_begin Paddings at the beginning of each spatial axis.
+                /// \param pads_end Paddings at the end of each spatial axis.
+                /// \param kernel The kernel shape.
+                /// \param rounding_type Whether to use ceiling or floor rounding type while
+                ///                      computing the output shape.
+                /// \param auto_pad The pad type for automatic calculation of the padding sizes.
+                /// \param index_element_type The data type used by the second output tensor
+                ///                           containing the selected indices.
+                /// \param axis Indicates a dimension in the input data shape which should be used
+                ///             as a starting point for calculation of the upper bound of allowed
+                ///             values of the indices output.
+                MaxPool(const Output<Node>& arg,
+                        const Strides& strides,
+                        const Strides& dilations,
+                        const Shape& pads_begin,
+                        const Shape& pads_end,
+                        const Shape& kernel,
+                        const op::RoundingType rounding_type = op::RoundingType::FLOOR,
+                        const PadType auto_pad = op::PadType::EXPLICIT,
+                        const element::Type index_element_type = element::i64,
+                        const int64_t axis = 0,
+                        const float pads_value = -std::numeric_limits<float>::infinity());
+
+                bool visit_attributes(AttributeVisitor& visitor) override;
+                void validate_and_infer_types() override;
+
+                virtual std::shared_ptr<Node>
+                    clone_with_new_inputs(const OutputVector& new_args) const override;
+
+                /// \return The pooling filter's dilations.
+                const Strides& get_dilations() const noexcept { return m_dilations; }
+                void set_dilations(const Strides& dilations) { m_dilations = dilations; }
+
+                /// \return The data type of the second output tensor (indices).
+                element::Type get_index_element_type() const noexcept
+                {
+                    return m_index_element_type;
+                }
+                void set_index_element_type(const element::Type index_element_type)
+                {
+                    m_index_element_type = index_element_type;
+                }
+
+                // \return The 'axis' attribute value.
+                int64_t get_axis() const { return m_axis; }
+                void set_axis(const int64_t axis) { m_axis = axis; }
+
+                // \return The value stored in the padding cells.
+                float get_pads_value() const { return m_pads_value; }
+                void set_pads_value(const float pads_value) { m_pads_value = pads_value; }
+
+            private:
+                Strides m_dilations;
+                element::Type m_index_element_type{element::i32};
+                int64_t m_axis{0};
+                float m_pads_value{-std::numeric_limits<float>::infinity()};
+            };
+        } // namespace v8
     }     // namespace op
 } // namespace ngraph

ngraph/core/include/ngraph/op/util/max_pool_base.hpp

@@ -0,0 +1,79 @@
+// Copyright (C) 2018-2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "ngraph/op/op.hpp"
+#include "ngraph/op/util/attr_types.hpp"
+
+namespace ngraph
+{
+    namespace op
+    {
+        namespace util
+        {
+            class NGRAPH_API MaxPoolBase : public Op
+            {
+            public:
+                NGRAPH_RTTI_DECLARATION;
+                MaxPoolBase() = default;
+
+                /// \param arg The node producing the input data batch tensor.
+                /// \param strides The strides.
+                /// \param pads_begin The beginning of padding shape.
+                /// \param pads_end The end of padding shape.
+                /// \param kernel The kernel shape.
+                /// \param rounding_mode Whether to use ceiling or floor rounding type while
+                /// computing output shape.
+                /// \param auto_pad The pad type for automatically computing padding sizes.
+                MaxPoolBase(const Output<Node>& arg,
+                            const Strides& strides,
+                            const Shape& pads_begin,
+                            const Shape& pads_end,
+                            const Shape& kernel,
+                            const op::RoundingType rounding_mode = op::RoundingType::FLOOR,
+                            const PadType auto_pad = op::PadType::EXPLICIT);
+
+                void validate_and_infer_types() override;
+
+                /// \return The kernel shape.
+                const Shape& get_kernel() const { return m_kernel; }
+                void set_kernel(const Shape& kernel) { m_kernel = kernel; }
+                /// \return The strides.
+                const Strides& get_strides() const { return m_strides; }
+                void set_strides(const Strides& strides) { m_strides = strides; }
+                /// \return The beginning of padding shape.
+                const Shape& get_pads_begin() const { return m_pads_begin; }
+                void set_pads_begin(const Shape& pads_begin) { m_pads_begin = pads_begin; }
+                /// \return The end of padding shape.
+                const Shape& get_pads_end() const { return m_pads_end; }
+                void set_adding_above(const Shape& pads_end) { m_pads_end = pads_end; }
+                /// \return The pad type for pooling.
+                PadType get_auto_pad() const { return m_auto_pad; }
+                void set_auto_pad(const PadType auto_pad) { m_auto_pad = auto_pad; }
+                /// \return The ceiling mode being used for output shape computations
+                op::RoundingType get_rounding_type() const { return m_rounding_type; }
+                void set_rounding_type(op::RoundingType rounding_type)
+                {
+                    m_rounding_type = rounding_type;
+                }
+
+            protected:
+                bool update_auto_padding(const PartialShape& in_shape,
+                                         const Strides& filter_dilations,
+                                         Shape& new_pads_end,
+                                         Shape& new_pads_begin) const;
+
+                PartialShape infer_output_shape(const Strides& dilations);
+
+                Shape m_kernel;
+                Strides m_strides;
+                Shape m_pads_begin;
+                Shape m_pads_end;
+                PadType m_auto_pad;
+                op::RoundingType m_rounding_type;
+            };
+        } // namespace util
+    }     // namespace op
+} // namespace ngraph

ngraph/core/include/ngraph/opsets/opset8_tbl.hpp

@@ -58,7 +58,6 @@ NGRAPH_OP(LogicalXor, ngraph::op::v1)
 NGRAPH_OP(LRN, ngraph::op::v0)
 NGRAPH_OP(LSTMCell, ngraph::op::v4)
 NGRAPH_OP(MatMul, ngraph::op::v0)
-NGRAPH_OP(MaxPool, ngraph::op::v1)
 NGRAPH_OP(Maximum, ngraph::op::v1)
 NGRAPH_OP(Minimum, ngraph::op::v1)
 NGRAPH_OP(Mod, ngraph::op::v1)
@@ -180,5 +179,6 @@ NGRAPH_OP(AdaptiveAvgPool, ngraph::op::v8)
 NGRAPH_OP(AdaptiveMaxPool, ngraph::op::v8)
 NGRAPH_OP(DeformableConvolution, ngraph::op::v8)
 NGRAPH_OP(MatrixNms, ngraph::op::v8)
+NGRAPH_OP(MaxPool, ngraph::op::v8)
 NGRAPH_OP(MulticlassNms, ngraph::op::v8)
 NGRAPH_OP(RandomUniform, ngraph::op::v8)

ngraph/core/include/ngraph/validation_util.hpp

@@ -85,7 +85,8 @@ namespace ngraph
                                                const PartialShape& window_shape,
                                                const Strides& window_strides,
                                                bool is_window_all_in_padding_allowed,
-                                               bool ceil_mode = false);
+                                               bool ceil_mode = false,
+                                               const Strides& window_dilation = Strides{});
 
     NGRAPH_API
     std::tuple<element::Type, PartialShape, PartialShape>

ngraph/core/src/op/avg_pool.cpp

@@ -147,7 +147,8 @@ void op::v1::AvgPool::validate_and_infer_types()
                                                         m_kernel,
                                                         m_strides,
                                                         !m_exclude_pad,
-                                                        m_rounding_type == op::RoundingType::CEIL)
+                                                        m_rounding_type == op::RoundingType::CEIL,
+                                                        Strides{}) // no dilation of the window
                         : output_shape);
 }