Add pooling operator

src/targets/cpu/cpu_target.cpp

@@ -60,7 +60,11 @@ struct max_pool
     static std::string name() { return "max"; }
     static double start() { return std::numeric_limits<double>::lowest(); }
 
-    static double apply(double x, double y) { return x + y; }
+    static double apply(double x, double y)
+    {
+        double m = std::max(x, y);
+        return (m);
+    }
 
     static double final(double x, double) { return (x); }
 };
@@ -70,11 +74,7 @@ struct avg_pool
     static std::string name() { return "average"; }
     static double start() { return 0.0; }
 
-    static double apply(double x, double y)
-    {
-        double m = std::max(x, y);
-        return (m);
-    }
+    static double apply(double x, double y) { return x + y; }
 
     static double final(double x, double y) { return x / y; }
 };

src/targets/miopen/include/rtg/miopen/miopen.hpp

@@ -12,6 +12,8 @@ using miopen_handle     = RTG_MANAGE_PTR(miopenHandle_t, miopenDestroy);
 using tensor_descriptor = RTG_MANAGE_PTR(miopenTensorDescriptor_t, miopenDestroyTensorDescriptor);
 using convolution_descriptor = RTG_MANAGE_PTR(miopenConvolutionDescriptor_t,
                                               miopenDestroyConvolutionDescriptor);
+using pooling_descriptor = RTG_MANAGE_PTR(miopenPoolingDescriptor_t,
+                                              miopenDestroyPoolingDescriptor);
 using activation_descriptor  = RTG_MANAGE_PTR(miopenActivationDescriptor_t,
                                              miopenDestroyActivationDescriptor);
 
@@ -55,6 +57,23 @@ inline convolution_descriptor make_conv(const rtg::convolution& op)
     return c;
 }
 
+inline pooling_descriptor make_pooling(const rtg::pooling& op)
+{
+    miopenPoolingMode_t mode;
+    if(op.mode == "max") mode = miopenPoolingMax;
+    else mode = miopenPoolingAverage;
+    auto p = make_obj<pooling_descriptor>(&miopenCreatePoolingDescriptor);
+    miopenSet2dPoolingDescriptor(p.get(),
+                                                          mode,
+                                                          op.lengths[0],
+                                    op.lengths[1],
+                                                          op.padding[0],
+                                    op.padding[1],
+                                    op.stride[0],
+                                    op.stride[1]);
+    return p;
+}
+
 inline activation_descriptor make_relu()
 {
     auto ad = make_obj<activation_descriptor>(&miopenCreateActivationDescriptor);

src/targets/miopen/miopen_target.cpp

@@ -59,6 +59,40 @@ struct miopen_convolution
     }
 };
 
+struct miopen_pooling
+{
+    pooling op;
+    shared<pooling_descriptor> pd;
+
+    std::string name() const { return "miopen::pooling"; }
+    shape compute_shape(std::vector<shape> inputs) const
+    {
+        check_shapes{inputs}.has(3);
+        return op.compute_shape({inputs.at(1)});
+    }
+    argument compute(shape output_shape, std::vector<argument> args) const
+    {
+        auto x_desc = make_tensor(args[1].get_shape());
+        auto y_desc = make_tensor(output_shape);
+
+        float alpha = 1, beta = 0;
+
+        miopenPoolingForward(args[0].implicit(),
+                                                  pd.get(),
+                                                  &alpha,
+                                                  x_desc.get(),
+                                              args[1].implicit(),
+                                                  &beta,
+                                                  y_desc.get(),
+                                              args[2].implicit(),
+                                                  false,
+                                                  nullptr,
+                                                  0);
+
+        return args[2];
+    }
+};
+
 struct miopen_relu
 {
     shared<activation_descriptor> ad;
@@ -105,6 +139,10 @@ struct miopen_apply
             {
                 apply_activation(it);
             }
+            else if(it->op.name() == "pooling")
+            {
+                apply_pooling(it);
+            }
         }
     }
 
@@ -136,6 +174,19 @@ struct miopen_apply
                                   output);
     }
 
+    void apply_pooling(instruction_ref ins)
+    {
+        auto&& op   = any_cast<pooling>(ins->op);
+        auto pd     = make_pooling(op);
+        auto output = insert_allocation(ins, ins->result);
+
+        prog->replace_instruction(ins,
+                                  miopen_pooling{op, std::move(pd)},
+                                  handle,
+                                  ins->arguments.at(0),
+                                  output);
+    }
+
     void apply_activation(instruction_ref ins)
     {
         auto&& op = any_cast<activation>(ins->op);

test/miopen/miopen.cpp

@@ -50,7 +50,7 @@ void verify_program()
     visit_all(cpu_arg, gpu_arg)([](auto cpu, auto gpu) { EXPECT(test::verify_range(cpu, gpu)); });
 }
 
-struct test1
+struct test_conv_relu
 {
     rtg::program create_program() const
     {
@@ -71,4 +71,29 @@ struct test1
     }
 };
 
-int main() { verify_program<test1>(); }
+struct test_conv_pooling
+{
+    rtg::program create_program() const
+    {
+        rtg::program p;
+        auto input   = p.add_parameter("x", rtg::shape{rtg::shape::float_type, {4, 3, 32, 32}});
+        auto weights = p.add_parameter("w", rtg::shape{rtg::shape::float_type, {4, 3, 3, 3}});
+        auto conv    = p.add_instruction(rtg::convolution{}, input, weights);
+        auto pooling = p.add_instruction(rtg::pooling{"max"}, conv);
+        p.add_instruction(rtg::activation{"relu"}, pooling);
+        return p;
+    }
+
+    rtg::program::parameter_map create_params() const
+    {
+        rtg::program::parameter_map m;
+        m["x"] = rtg::generate_argument({rtg::shape::float_type, {4, 3, 32, 32}});
+        m["w"] = rtg::generate_argument({rtg::shape::float_type, {4, 3, 3, 3}});
+        return m;
+    }
+};
+
+int main() { 
+    verify_program<test_conv_relu>(); 
+    verify_program<test_conv_pooling>(); 
+}