pool.hpp

#pragma once
#include <dlprim/tensor.hpp>
#include <dlprim/context.hpp>
namespace dlprim {
namespace core {

    inline int calc_pooling_output_size(int in_size,int kernel,int pad,int stride,bool ceil_mode)
    {
        int padded_size = in_size + pad*2;
        DLPRIM_CHECK(padded_size >= kernel);
        int offset = ceil_mode ? (stride-1) : 0;
        int size = (padded_size - kernel + offset) / stride + 1;
        if((size - 1) * stride >= in_size + pad) {
            size--;
        }
        return size;
    }


    class Pooling2DForward {
    public:
        virtual ~Pooling2DForward() {}
        // get workspace size
        virtual size_t workspace() = 0;

        virtual void enqueue(Tensor &X,Tensor &Y,ExecutionContext const &e) = 0;

        static std::unique_ptr<Pooling2DForward> create_max_pooling(
                            Context &ctx,
                            int kernel[2],int pad[2],int stride[2],
                            DataType dt=float_data);

        static std::unique_ptr<Pooling2DForward> create_avg_pooling(
                    Context &ctx,
                    int kernel[2],int pad[2],int stride[2],bool count_include_pad=false,
                    DataType dt=float_data);
        
        static std::unique_ptr<Pooling2DForward> create_global_max_pooling(
                    Context &ctx,Shape const &in_shape,DataType dt=float_data);

        static std::unique_ptr<Pooling2DForward> create_global_avg_pooling(
                    Context &ctx,Shape const &in_shape,DataType dt=float_data);
    };

    class Pooling2DBackwardBase {
    public:
        virtual size_t workspace() = 0;
        virtual ~Pooling2DBackwardBase() {}
        virtual void enqueue(Tensor &X,Tensor &dX,Tensor &dY,float factor,ExecutionContext const &e) = 0;
    };

    class MaxPooling2DBackward : public Pooling2DBackwardBase {
    public:
        virtual ~MaxPooling2DBackward() {}
        static std::unique_ptr<MaxPooling2DBackward> create(
                    Context &ctx,
                    int kernel[2],int pad[2],int stride[2],
                    DataType dt=float_data);

        static std::unique_ptr<MaxPooling2DBackward> create_global(
                    Context &ctx,Shape const &in_shape,DataType dt=float_data);
    };

    class AvgPooling2DBackward : public Pooling2DBackwardBase {
    public:
        virtual ~AvgPooling2DBackward() {}

        virtual void enqueue(Tensor &/*X*/,Tensor &dX,Tensor &dY,float factor,ExecutionContext const &e) 
        {
            enqueue(dX,dY,factor,e);
        }
        virtual void enqueue(Tensor &dX,Tensor &dY,float factor,ExecutionContext const &e)  = 0;
        
        static std::unique_ptr<AvgPooling2DBackward> create(
                    Context &ctx,
                    int kernel[2],int pad[2],int stride[2],bool count_include_pad=false,
                    DataType dt=float_data);

        static std::unique_ptr<AvgPooling2DBackward> create_global(
                    Context &ctx,Shape const &in_shape,DataType dt=float_data);
    };
} // core
} //dlprim