Pooling ceil mode (#615)

* support pooling ceil_mode

* clang format

* add unit test for pooling ceil mode

* clang format

* fix review comments

* clang format

* add more unit tests and fixed a bug in cpu pooling implementation

* clang format

* add one more unit test

* clang format

* fix cppcheck error

* fix cppcheck error

* fix cppcheck error

* fix review comments

* clang format

* remove the padding_mode attribute in pooling

* clang format

* clang format

* fix review comments

* clang format

* fix a cppcheck error

* fix review comments
Co-authored-by: mvermeulen <5479696+mvermeulen@users.noreply.github.com>

src/include/migraphx/op/pooling.hpp

@@ -23,16 +23,16 @@ struct pooling
     std::vector<std::size_t> padding = {0, 0};
     std::vector<std::size_t> stride  = {1, 1};
     std::vector<std::size_t> lengths = {1, 1};
-    padding_mode_t padding_mode      = default_;
+    bool ceil_mode                   = false;
 
     template <class Self, class F>
     static auto reflect(Self& self, F f)
     {
         return pack(f(self.mode, "mode"),
                     f(self.padding, "padding"),
-                    f(self.padding_mode, "padding_mode"),
                     f(self.stride, "stride"),
-                    f(self.lengths, "lengths"));
+                    f(self.lengths, "lengths"),
+                    f(self.ceil_mode, "ceil_mode"));
     }
 
     std::string name() const { return "pooling"; }
@@ -63,13 +63,12 @@ struct pooling
 
         for(size_t i = 0; i < kdims; i++)
         {
-            assert(lengths[i] <= input_lens[i + 2] + 2 * padding[i]);
+            std::ptrdiff_t dim_size = input_lens[i + 2] + 2 * padding[i] - lengths[i];
+            assert(dim_size >= 0);
+            std::size_t len = (ceil_mode) ? ceil_divide<std::ptrdiff_t>(dim_size, stride[i])
+                                          : floor_divide<std::ptrdiff_t>(dim_size, stride[i]);
 
-            output_lens.push_back(std::size_t(std::max<std::ptrdiff_t>(
-                1,
-                floor_divide<std::ptrdiff_t>(input_lens[i + 2] + 2 * padding[i] - lengths[i],
-                                             stride[i]) +
-                    1)));
+            output_lens.push_back(std::size_t(std::max<std::ptrdiff_t>(1, len + 1)));
         }
         return {t, output_lens};
     }

src/onnx/onnx.cpp

@@ -630,7 +630,6 @@ struct onnx_parser
         auto auto_pad = info.attributes["auto_pad"].s();
         if(auto_pad.find("SAME") != std::string::npos)
         {
-            v["padding_mode"]  = to_value(op::padding_mode_t::same);
             bool is_same_upper = (auto_pad.find("SAME_UPPER") != std::string::npos);
             paddings.resize(2 * kdims);
 
@@ -702,7 +701,6 @@ struct onnx_parser
         if(contains(info.attributes, "auto_pad"))
         {
             auto weight_lens = weights->get_shape().lens();
-
             std::vector<std::size_t> k_lens(weight_lens.begin() + 2, weight_lens.end());
             cal_auto_padding_size(info,
                                   values,
@@ -710,6 +708,11 @@ struct onnx_parser
                                   values["dilation"].to_vector<std::size_t>(),
                                   in_lens,
                                   padding);
+            auto auto_pad = info.attributes["auto_pad"].s();
+            if(auto_pad.find("SAME") != std::string::npos)
+            {
+                values["padding_mode"] = to_value(op::padding_mode_t::same);
+            }
         }
         check_asym_padding(l0, padding, values);
 
@@ -913,10 +916,7 @@ struct onnx_parser
         // does not support ceil_mode
         if(contains(info.attributes, "ceil_mode"))
         {
-            if(info.attributes.at("ceil_mode").i() == 1)
-            {
-                MIGRAPHX_THROW("PARSE_POOLING: pool does not support ceil_mode");
-            }
+            values["ceil_mode"] = static_cast<bool>(info.attributes.at("ceil_mode").i());
         }
 
         // count include padding, if count include pad is 1, we always use

src/targets/cpu/lowering.cpp

@@ -421,7 +421,11 @@ MIGRAPHX_REGISTER_OP(cpu_im2col)
 struct max_pool
 {
     static std::string name() { return "max"; }
-    static double start() { return std::numeric_limits<double>::lowest(); }
+    template <class T>
+    static T start()
+    {
+        return std::numeric_limits<T>::lowest();
+    }
 
     static double apply(double x, double y)
     {
@@ -429,17 +433,22 @@ struct max_pool
         return (m);
     }
 
-    static double final(double x, double) { return (x); }
+    static double final(double x, std::size_t) { return (x); }
 };
 
 struct avg_pool
 {
     static std::string name() { return "average"; }
-    static double start() { return 0.0; }
+
+    template <class T>
+    static double start()
+    {
+        return 0.0;
+    }
 
     static double apply(double x, double y) { return x + y; }
 
-    static double final(double x, double y) { return x / y; }
+    static double final(double x, std::size_t y) { return (y == 0) ? 0.0 : (x / y); }
 };
 
 template <class Op>
@@ -486,7 +495,7 @@ struct cpu_pooling : auto_register_op<cpu_pooling<Op>>
 
                 shape win_shape{output_shape.type(), win_size};
                 auto pool_size = win_shape.elements();
-                double acc     = Op::start();
+                double acc     = Op::template start<type>();
                 shape_for_each(win_shape, [&](auto idx_w) {
                     auto idx = idx_o;
                     std::transform(idx_w.begin(),

src/tf/tf.cpp

@@ -746,7 +746,6 @@ struct tf_parser
             const std::string& pad_mode = attributes.at("padding").s();
             if(pad_mode.find("SAME") != std::string::npos)
             {
-                op.padding_mode = op::padding_mode_t::same;
                 auto input_dims = l0->get_shape().lens();
                 std::vector<int64_t> pads(input_dims.size());
                 calculate_padding(0, pads, input_dims[2], op.stride[0], 1, op.lengths[0]);
@@ -764,10 +763,6 @@ struct tf_parser
                     op.padding[1] = pads[1];
                 }
             }
-            else if(pad_mode.find("VALID") != std::string::npos)
-            {
-                op.padding_mode = op::padding_mode_t::valid;
-            }
         }
         return prog.add_instruction(op, l0);
     }

test/cpu_ops_test.cpp

@@ -514,6 +514,41 @@ TEST_CASE(maxpool_test_1D_3D)
         EXPECT(migraphx::verify_range(results_vector, gold));
     }
 
+    // 1D case 2, input is 3D, ceil mode
+    {
+        migraphx::program p;
+        auto s       = migraphx::shape{migraphx::shape::float_type, {2, 2, 5}};
+        auto op      = migraphx::op::pooling{"max"};
+        op.lengths   = {2};
+        op.padding   = {0};
+        op.stride    = {2};
+        op.ceil_mode = true;
+
+        std::vector<float> data{0.4975, -0.1226, -0.0405, -0.2861, -0.1227, -0.6186, -0.9618,
+                                0.6022, -0.1912, 1.1925,  0.5493,  0.1692,  -0.8039, -1.0281,
+                                0.9907, 0.477,   1.5001,  -1.1603, -1.361,  1.2556};
+        auto l0 = p.add_literal(migraphx::literal{s, data});
+        p.add_instruction(op, l0);
+        p.compile(migraphx::cpu::target{});
+        auto result = p.eval({}).back();
+
+        std::vector<float> results_vector;
+        result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+        std::vector<float> gold{0.4975,
+                                -0.0405,
+                                -0.1227,
+                                -0.6186,
+                                0.6022,
+                                1.1925,
+                                0.5493,
+                                -0.8039,
+                                0.9907,
+                                1.5001,
+                                -1.1603,
+                                1.2556};
+        EXPECT(migraphx::verify_range(results_vector, gold));
+    }
+
     // 3D, input is 5D
     {
         migraphx::program p;

test/gpu/ops_test.cpp

@@ -1211,6 +1211,19 @@ struct test_global_max_pooling : verify_program<test_global_max_pooling>
     }
 };
 
+struct test_max_pooling_ceil_3d : verify_program<test_max_pooling_ceil_3d>
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+        auto input =
+            p.add_parameter("x", migraphx::shape{migraphx::shape::float_type, {1, 3, 5, 5, 5}});
+        auto op = migraphx::op::pooling{"max", {1, 1, 1}, {3, 3, 3}, {3, 3, 3}, true};
+        p.add_instruction(op, input);
+        return p;
+    }
+};
+
 struct test_avg_pooling_1d : verify_program<test_avg_pooling_1d>
 {
     migraphx::program create_program() const
@@ -1230,7 +1243,20 @@ struct test_avg_pooling_3d : verify_program<test_avg_pooling_3d>
         migraphx::program p;
         auto input =
             p.add_parameter("x", migraphx::shape{migraphx::shape::float_type, {1, 3, 5, 5, 5}});
-        auto op = migraphx::op::pooling{"average", {0, 0, 0}, {1, 1, 1}, {3, 3, 3}};
+        auto op = migraphx::op::pooling{"average", {1, 1, 1}, {3, 3, 3}, {3, 3, 3}};
+        p.add_instruction(op, input);
+        return p;
+    }
+};
+
+struct test_avg_pooling_ceil_3d : verify_program<test_avg_pooling_ceil_3d>
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+        auto input =
+            p.add_parameter("x", migraphx::shape{migraphx::shape::float_type, {1, 3, 5, 5, 5}});
+        auto op = migraphx::op::pooling{"average", {1, 1, 1}, {3, 3, 3}, {3, 3, 3}, true};
         p.add_instruction(op, input);
         return p;
     }
@@ -2438,9 +2464,8 @@ struct test_pooling_autopad : verify_program<test_pooling_autopad>
         migraphx::shape s0{migraphx::shape::float_type, {1, 3, 63, 63}};
         auto l0 = p.add_parameter("x", s0);
         migraphx::op::pooling op{"max"};
-        op.padding_mode = migraphx::op::padding_mode_t::same;
-        op.lengths      = {2, 2};
-        op.stride       = {2, 2};
+        op.lengths = {2, 2};
+        op.stride  = {2, 2};
         p.add_instruction(op, l0);
         return p;
     }

test/onnx/onnx_test.cpp

@@ -205,11 +205,8 @@ TEST_CASE(averagepool_same_lower_test)
 {
     migraphx::program p;
     auto input = p.add_parameter("x", migraphx::shape{migraphx::shape::float_type, {1, 1, 5, 5}});
-    auto ins   = p.add_instruction(
-        migraphx::op::pooling{
-            "average", {1, 1}, {1, 1}, {2, 2}, migraphx::op::padding_mode_t::same},
-        input);
-    auto ret = p.add_instruction(migraphx::op::slice{{2, 3}, {0, 0}, {5, 5}}, ins);
+    auto ins   = p.add_instruction(migraphx::op::pooling{"average", {1, 1}, {1, 1}, {2, 2}}, input);
+    auto ret   = p.add_instruction(migraphx::op::slice{{2, 3}, {0, 0}, {5, 5}}, ins);
     p.add_return({ret});
     auto prog = migraphx::parse_onnx("averagepool_same_lower_test.onnx");
 
@@ -222,10 +219,7 @@ TEST_CASE(averagepool_sl_cip_test)
     auto input = p.add_parameter("x", migraphx::shape{migraphx::shape::float_type, {1, 1, 5, 5}});
     std::vector<int64_t> pads = {0, 0, 1, 1, 0, 0, 0, 0};
     auto ins_pad              = p.add_instruction(migraphx::op::pad{pads}, input);
-    auto ret                  = p.add_instruction(
-        migraphx::op::pooling{
-            "average", {0, 0}, {1, 1}, {2, 2}, migraphx::op::padding_mode_t::same},
-        ins_pad);
+    auto ret = p.add_instruction(migraphx::op::pooling{"average", {0, 0}, {1, 1}, {2, 2}}, ins_pad);
     p.add_return({ret});
     auto prog = migraphx::parse_onnx("averagepool_sl_cip_test.onnx");
 
@@ -236,11 +230,8 @@ TEST_CASE(averagepool_same_upper_test)
 {
     migraphx::program p;
     auto input = p.add_parameter("x", migraphx::shape{migraphx::shape::float_type, {1, 1, 5, 5}});
-    auto ins   = p.add_instruction(
-        migraphx::op::pooling{
-            "average", {1, 1}, {1, 1}, {2, 2}, migraphx::op::padding_mode_t::same},
-        input);
-    auto ret = p.add_instruction(migraphx::op::slice{{2, 3}, {1, 1}, {6, 6}}, ins);
+    auto ins   = p.add_instruction(migraphx::op::pooling{"average", {1, 1}, {1, 1}, {2, 2}}, input);
+    auto ret   = p.add_instruction(migraphx::op::slice{{2, 3}, {1, 1}, {6, 6}}, ins);
     p.add_return({ret});
     auto prog = migraphx::parse_onnx("averagepool_same_upper_test.onnx");
 
@@ -1338,9 +1329,7 @@ TEST_CASE(maxpool_same_upper_test)
     std::vector<int64_t> pads = {0, 0, 0, 0, 0, 0, 1, 1};
     float val                 = std::numeric_limits<float>::lowest();
     auto ins_pad              = p.add_instruction(migraphx::op::pad{pads, val}, input);
-    p.add_instruction(
-        migraphx::op::pooling{"max", {0, 0}, {1, 1}, {2, 2}, migraphx::op::padding_mode_t::same},
-        ins_pad);
+    p.add_instruction(migraphx::op::pooling{"max", {0, 0}, {1, 1}, {2, 2}}, ins_pad);
 
     auto prog = optimize_onnx("maxpool_same_upper_test.onnx");
 

test/op_shape_test.cpp

@@ -142,7 +142,10 @@ TEST_CASE(pooling_shape)
     migraphx::shape output{migraphx::shape::float_type, {4, 3, 1, 1}};
     migraphx::shape input{migraphx::shape::float_type, {4, 3, 3, 3}};
     throws_shape(migraphx::op::pooling{"max", {1}, {0}, {1}}, input);
-    expect_shape(output, migraphx::op::pooling{"max", {0, 0}, {1, 1}, {3, 3}}, input);
+    expect_shape(output, migraphx::op::pooling{"max", {0, 0}, {3, 3}, {1, 1}}, input);
+
+    migraphx::shape output1{migraphx::shape::float_type, {4, 3, 2, 2}};
+    expect_shape(output1, migraphx::op::pooling{"max", {0, 0}, {3, 3}, {1, 1}, true}, input);
 }
 
 TEST_CASE(inconsistent_attr_shape)

test/py/onnx_backend_test.py

@@ -204,7 +204,6 @@ def create_backend_test(testname=None, target_device=None):
         backend_test.exclude(r'test_softmax_default_axis_cpu')
 
         # error cases
-        backend_test.exclude(r'test_averagepool_2d_ceil_cpu')
         backend_test.exclude(r'test_clip_default_inbounds_cpu')
         backend_test.exclude(r'test_clip_default_int8_inbounds_cpu')
         backend_test.exclude(r'test_clip_default_int8_max_cpu')
@@ -245,7 +244,6 @@ def create_backend_test(testname=None, target_device=None):
         backend_test.exclude(r'test_max_float16_cpu')
         backend_test.exclude(r'test_max_int64_cpu')
         backend_test.exclude(r'test_max_uint64_cpu')
-        backend_test.exclude(r'test_maxpool_2d_ceil_cpu')
         backend_test.exclude(r'test_maxpool_2d_uint8_cpu')
         backend_test.exclude(r'test_mean_example_cpu')
         backend_test.exclude(r'test_mean_one_input_cpu')

test/tf/tf_test.cpp

@@ -407,12 +407,10 @@ TEST_CASE(pooling_test)
     auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 3, 16, 16}});
     migraphx::op::pooling avg_pool_op{"average"};
     migraphx::op::pooling max_pool_op{"max"};
-    avg_pool_op.padding_mode = migraphx::op::padding_mode_t::valid;
-    max_pool_op.padding_mode = migraphx::op::padding_mode_t::valid;
-    avg_pool_op.stride       = {2, 2};
-    max_pool_op.stride       = {2, 2};
-    avg_pool_op.lengths      = {2, 2};
-    max_pool_op.lengths      = {2, 2};
+    avg_pool_op.stride  = {2, 2};
+    max_pool_op.stride  = {2, 2};
+    avg_pool_op.lengths = {2, 2};
+    max_pool_op.lengths = {2, 2};
     p.add_instruction(max_pool_op, l0);
     auto prog = optimize_tf("pooling_test.pb", true);