From 874cd289bc5e93936f4c1fc9be4ebd0a25541220 Mon Sep 17 00:00:00 2001
From: cmoineau <cyril.moineau@cea.fr>
Date: Thu, 19 Oct 2023 15:04:10 +0000
Subject: [PATCH] [MaxPooling] Add support for ceil_mode parameter.

---
 include/aidge/operator/MaxPooling.hpp         | 32 ++++++++++++-------
 include/aidge/operator/MetaOperatorDefs.hpp   | 10 +++---
 python_binding/operator/pybind_MaxPooling.cpp | 17 +++++-----
 .../operator/pybind_MetaOperatorDefs.cpp      | 28 ++++++++--------
 4 files changed, 49 insertions(+), 38 deletions(-)

diff --git a/include/aidge/operator/MaxPooling.hpp b/include/aidge/operator/MaxPooling.hpp
index 874ea8177..d224ee0cd 100644
--- a/include/aidge/operator/MaxPooling.hpp
+++ b/include/aidge/operator/MaxPooling.hpp
@@ -26,14 +26,15 @@
 #include "aidge/utils/Types.h"
 
 namespace Aidge {
-enum class MaxPoolingAttr { StrideDims, KernelDims };
+enum class MaxPoolingAttr { StrideDims, KernelDims, CeilMode };
 
 template <DimIdx_t DIM>
 class MaxPooling_Op : public Operator,
                 public Registrable<MaxPooling_Op<DIM>, std::string, std::unique_ptr<OperatorImpl>(const MaxPooling_Op<DIM> &)>,
                 public StaticAttributes<MaxPoolingAttr,
                                        std::array<DimSize_t, DIM>,
-                                       std::array<DimSize_t, DIM>> {
+                                       std::array<DimSize_t, DIM>,
+                                       bool> {
 private:
     // FIXME: change accessibility
     std::shared_ptr<Tensor> mInput = std::make_shared<Tensor>();
@@ -46,15 +47,18 @@ public:
 
     using Attributes_ = StaticAttributes<MaxPoolingAttr,
                                              std::array<DimSize_t, DIM>,
-                                             std::array<DimSize_t, DIM>>;
+                                             std::array<DimSize_t, DIM>,
+                                             bool>;
     template <MaxPoolingAttr e>
     using attr = typename Attributes_::template attr<e>;
 
     constexpr MaxPooling_Op(const std::array<DimSize_t, DIM> &kernel_dims,
-                            const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t,DIM>(1))
+                            const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t,DIM>(1),
+                            bool ceil_mode = false)
         : Operator(Type),
           Attributes_(attr<MaxPoolingAttr::StrideDims>(stride_dims),
-                      attr<MaxPoolingAttr::KernelDims>(kernel_dims)),
+                      attr<MaxPoolingAttr::KernelDims>(kernel_dims),
+                      attr<MaxPoolingAttr::CeilMode>(ceil_mode)),
           mOutput(std::make_shared<Tensor>()) {
         setDatatype(DataType::Float32);
     }
@@ -92,10 +96,12 @@ public:
     void computeOutputDims() override final {
         if (!mInput->empty()) {
             std::array<DimSize_t, DIM + 2> outputDims = {};
-
+            auto roundingFunction = (this->template getAttr<MaxPoolingAttr::CeilMode>()) ?
+                [](float x){return std::ceil(x);} :
+                [](float x){return std::floor(x);};
             for (std::size_t dim = 0; dim < this->template getAttr<MaxPoolingAttr::KernelDims>().size() ; ++dim) {
                 outputDims[dim+2] = 1 + static_cast<DimSize_t>(
-                                            std::floor(static_cast<float>(mInput->dims()[dim+2] -
+                                            roundingFunction(static_cast<float>(mInput->dims()[dim+2] -
                                                                     this->template getAttr<MaxPoolingAttr::KernelDims>()[dim]) /
                                             static_cast<float>(this->template getAttr<MaxPoolingAttr::StrideDims>()[dim])));
             }
@@ -169,9 +175,10 @@ public:
 template <std::array<DimSize_t, 1>::size_type DIM>
 inline std::shared_ptr<Node> MaxPooling(const std::array<DimSize_t, DIM> &kernel_dims,
                                            const std::string& name = "",
-                                           const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t,DIM>(1)) {
+                                           const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t,DIM>(1),
+                                           bool ceil_mode=false) {
     static_assert(DIM<=MaxDim,"Too many kernel dimensions required by MaxPooling, not supported");
-    return std::make_shared<Node>(std::make_shared<MaxPooling_Op<static_cast<DimIdx_t>(DIM)>>(kernel_dims, stride_dims), name);
+    return std::make_shared<Node>(std::make_shared<MaxPooling_Op<static_cast<DimIdx_t>(DIM)>>(kernel_dims, stride_dims, ceil_mode), name);
 }
 
 // helper with C-style array instead of std::array for kernel_dims to allow automatic template DIM deduction
@@ -179,15 +186,16 @@ template <DimSize_t DIM>
 inline std::shared_ptr<Node> MaxPooling(
     DimSize_t const (&kernel_dims)[DIM],
     const std::string& name = "",
-    const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t,DIM>(1)) {
+    const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t,DIM>(1),
+    bool ceil_mode = false) {
     static_assert(DIM<=MaxDim,"Too many kernel dimensions required by MaxPooling, not supported");
-    return MaxPooling(to_array(kernel_dims), name, stride_dims);
+    return MaxPooling(to_array(kernel_dims), name, stride_dims, ceil_mode);
 }
 }  // namespace Aidge
 
 namespace {
 template <>
-const char *const EnumStrings<Aidge::MaxPoolingAttr>::data[] = {"StrideDims", "KernelDims"};
+const char *const EnumStrings<Aidge::MaxPoolingAttr>::data[] = {"StrideDims", "KernelDims", "CeilMode"};
 }
 
 #endif /* AIDGE_CORE_OPERATOR_MAXPOOLING_H_ */
diff --git a/include/aidge/operator/MetaOperatorDefs.hpp b/include/aidge/operator/MetaOperatorDefs.hpp
index 6da76c930..73feb1348 100644
--- a/include/aidge/operator/MetaOperatorDefs.hpp
+++ b/include/aidge/operator/MetaOperatorDefs.hpp
@@ -115,11 +115,12 @@ template <std::array<DimSize_t, 1>::size_type DIM>
 inline std::shared_ptr<Node> PaddedMaxPooling(const std::array<DimSize_t, DIM> &kernel_dims,
                                   const std::string& name = "",
                                   const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t,DIM>(1),
-                                  const std::array<DimSize_t, 2*DIM> &padding_dims = create_array<DimSize_t,2*DIM>(0))
+                                  const std::array<DimSize_t, 2*DIM> &padding_dims = create_array<DimSize_t,2*DIM>(0),
+                                  bool ceil_mode = false)
 {
     auto graph = Sequential({
         Pad<DIM>(padding_dims, (!name.empty()) ? name + "_pad" : ""),
-        MaxPooling(kernel_dims, (!name.empty()) ? name + "_maxpooling" : "", stride_dims)
+        MaxPooling(kernel_dims, (!name.empty()) ? name + "_maxpooling" : "", stride_dims, ceil_mode)
     });
 
     return MetaOperator("PaddedMaxPooling", graph, name);
@@ -131,9 +132,10 @@ inline std::shared_ptr<Node> PaddedMaxPooling(
     DimSize_t const (&kernel_dims)[DIM],
     const std::string& name = "",
     const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t,DIM>(1),
-    const std::array<DimSize_t, 2*DIM> &padding_dims = create_array<DimSize_t,2*DIM>(0))
+    const std::array<DimSize_t, 2*DIM> &padding_dims = create_array<DimSize_t,2*DIM>(0),
+    bool ceil_mode= false)
 {
-    return PaddedMaxPooling(to_array(kernel_dims), name, stride_dims, padding_dims);
+    return PaddedMaxPooling(to_array(kernel_dims), name, stride_dims, padding_dims, ceil_mode);
 }
 }  // namespace Aidge
 
diff --git a/python_binding/operator/pybind_MaxPooling.cpp b/python_binding/operator/pybind_MaxPooling.cpp
index c83dfaa36..907e8cfaa 100644
--- a/python_binding/operator/pybind_MaxPooling.cpp
+++ b/python_binding/operator/pybind_MaxPooling.cpp
@@ -30,22 +30,26 @@ template <DimIdx_t DIM> void declare_MaxPoolingOp(py::module &m) {
     m, ("MaxPoolingOp" + std::to_string(DIM) + "D").c_str(),
     py::multiple_inheritance())
   .def(py::init<const std::array<DimSize_t, DIM> &,
-                const std::array<DimSize_t, DIM> &>(),
+                const std::array<DimSize_t, DIM> &,
+                bool>(),
         py::arg("kernel_dims"),
-        py::arg("stride_dims"))
+        py::arg("stride_dims"),
+        py::arg("ceil_mode"))
   .def("get_inputs_name", &MaxPooling_Op<DIM>::getInputsName)
   .def("get_outputs_name", &MaxPooling_Op<DIM>::getOutputsName);
 
   m.def(("MaxPooling" + std::to_string(DIM) + "D").c_str(), [](const std::vector<DimSize_t>& kernel_dims,
                                                                   const std::string& name,
-                                                                  const std::vector<DimSize_t> &stride_dims) {
+                                                                  const std::vector<DimSize_t> &stride_dims,
+                                                                  bool ceil_mode) {
         AIDGE_ASSERT(kernel_dims.size() == DIM, "kernel_dims size [%ld] does not match DIM [%d]", kernel_dims.size(), DIM);
         AIDGE_ASSERT(stride_dims.size() == DIM, "stride_dims size [%ld] does not match DIM [%d]", stride_dims.size(), DIM);
 
-        return MaxPooling<DIM>(to_array<DIM>(kernel_dims.begin()), name, to_array<DIM>(stride_dims.begin()));
+        return MaxPooling<DIM>(to_array<DIM>(kernel_dims.begin()), name, to_array<DIM>(stride_dims.begin()), ceil_mode);
     }, py::arg("kernel_dims"),
        py::arg("name") = "",
-       py::arg("stride_dims") = std::vector<DimSize_t>(DIM,1));
+       py::arg("stride_dims") = std::vector<DimSize_t>(DIM,1),
+       py::arg("ceil_mode") = false);
 
 }
 
@@ -55,8 +59,5 @@ void init_MaxPooling(py::module &m) {
   declare_MaxPoolingOp<2>(m);
   declare_MaxPoolingOp<3>(m);
 
-  // FIXME:
-  // m.def("MaxPooling1D", static_cast<NodeAPI(*)(const char*, int, int, int const
-  // (&)[1])>(&MaxPooling));
 }
 } // namespace Aidge
diff --git a/python_binding/operator/pybind_MetaOperatorDefs.cpp b/python_binding/operator/pybind_MetaOperatorDefs.cpp
index 3372d50e1..aa9f3c50e 100644
--- a/python_binding/operator/pybind_MetaOperatorDefs.cpp
+++ b/python_binding/operator/pybind_MetaOperatorDefs.cpp
@@ -28,7 +28,7 @@ template <DimIdx_t DIM> void declare_PaddedConvOp(py::module &m) {
   m.def(("PaddedConv" + std::to_string(DIM) + "D").c_str(), [](DimSize_t in_channels,
                                                          DimSize_t out_channels,
                                                          const std::vector<DimSize_t>& kernel_dims,
-                                                         const std::string& name, 
+                                                         const std::string& name,
                                                          const std::vector<DimSize_t> &stride_dims,
                                                          const std::vector<DimSize_t> &padding_dims,
                                                          const std::vector<DimSize_t> &dilation_dims)
@@ -50,7 +50,7 @@ template <DimIdx_t DIM> void declare_PaddedConvOp(py::module &m) {
 
 template <DimIdx_t DIM> void declare_PaddedConvDepthWiseOp(py::module &m) {
   m.def(("PaddedConvDepthWise" + std::to_string(DIM) + "D").c_str(), [](const std::vector<DimSize_t>& kernel_dims,
-                                                         const std::string& name, 
+                                                         const std::string& name,
                                                          const std::vector<DimSize_t> &stride_dims,
                                                          const std::vector<DimSize_t> &padding_dims,
                                                          const std::vector<DimSize_t> &dilation_dims)
@@ -66,12 +66,12 @@ template <DimIdx_t DIM> void declare_PaddedConvDepthWiseOp(py::module &m) {
        py::arg("stride_dims") = std::vector<DimSize_t>(DIM,1),
        py::arg("padding_dims") = std::vector<DimSize_t>(2*DIM,0),
        py::arg("dilation_dims") = std::vector<DimSize_t>(DIM,1));
-  
+
 }
 
 template <DimIdx_t DIM> void declare_PaddedAvgPoolingOp(py::module &m) {
   m.def(("PaddedAvgPooling" + std::to_string(DIM) + "D").c_str(), [](const std::vector<DimSize_t>& kernel_dims,
-                                                         const std::string& name, 
+                                                         const std::string& name,
                                                          const std::vector<DimSize_t> &stride_dims,
                                                          const std::vector<DimSize_t> &padding_dims)
     {
@@ -84,25 +84,27 @@ template <DimIdx_t DIM> void declare_PaddedAvgPoolingOp(py::module &m) {
        py::arg("name") = "",
        py::arg("stride_dims") = std::vector<DimSize_t>(DIM,1),
        py::arg("padding_dims") = std::vector<DimSize_t>(2*DIM,0));
-  
+
 }
 
 template <DimIdx_t DIM> void declare_PaddedMaxPoolingOp(py::module &m) {
   m.def(("PaddedMaxPooling" + std::to_string(DIM) + "D").c_str(), [](const std::vector<DimSize_t>& kernel_dims,
-                                                         const std::string& name, 
+                                                         const std::string& name,
                                                          const std::vector<DimSize_t> &stride_dims,
-                                                         const std::vector<DimSize_t> &padding_dims)
+                                                         const std::vector<DimSize_t> &padding_dims,
+                                                         bool ceil_mode)
     {
         AIDGE_ASSERT(kernel_dims.size() == DIM, "kernel_dims size [%ld] does not match DIM [%d]", kernel_dims.size(), DIM);
         AIDGE_ASSERT(stride_dims.size() == DIM, "stride_dims size [%ld] does not match DIM [%d]", stride_dims.size(), DIM);
         AIDGE_ASSERT(padding_dims.size() == 2*DIM, "padding_dims size [%ld] does not match DIM [%d]", padding_dims.size(), 2*DIM);
 
-        return PaddedMaxPooling<DIM>(to_array<DIM>(kernel_dims.begin()), name, to_array<DIM>(stride_dims.begin()), to_array<2*DIM>(padding_dims.begin()));
+        return PaddedMaxPooling<DIM>(to_array<DIM>(kernel_dims.begin()), name, to_array<DIM>(stride_dims.begin()), to_array<2*DIM>(padding_dims.begin()), ceil_mode);
     }, py::arg("kernel_dims"),
        py::arg("name") = "",
        py::arg("stride_dims") = std::vector<DimSize_t>(DIM,1),
-       py::arg("padding_dims") = std::vector<DimSize_t>(2*DIM,0));
-  
+       py::arg("padding_dims") = std::vector<DimSize_t>(2*DIM,0),
+       py::arg("ceil_mode") = false);
+
 }
 
 void init_MetaOperatorDefs(py::module &m) {
@@ -118,9 +120,7 @@ void init_MetaOperatorDefs(py::module &m) {
   declare_PaddedMaxPoolingOp<1>(m);
   declare_PaddedMaxPoolingOp<2>(m);
   declare_PaddedMaxPoolingOp<3>(m);
- 
-  // FIXME:
-  // m.def("Conv1D", static_cast<NodeAPI(*)(const char*, int, int, int const
-  // (&)[1])>(&Conv));
+
+
 }
 } // namespace Aidge
-- 
GitLab