MaxPooling.cpp

/********************************************************************************
 * Copyright (c) 2023 CEA-List
 *
 * This program and the accompanying materials are made available under the
 * terms of the Eclipse Public License 2.0 which is available at
 * http://www.eclipse.org/legal/epl-2.0.
 *
 * SPDX-License-Identifier: EPL-2.0
 *
 ********************************************************************************/

#include "aidge/operator/MaxPooling.hpp"

#include <algorithm>
#include <string>
#include <vector>

#include "aidge/data/Tensor.hpp"
#include "aidge/utils/Types.h"
#include "aidge/utils/ErrorHandling.hpp"

template <Aidge::DimIdx_t DIM>
const std::string Aidge::MaxPooling_Op<DIM>::Type = "MaxPooling" + std::to_string(DIM) + "D";

template <Aidge::DimIdx_t DIM>
Aidge::MaxPooling_Op<DIM>::MaxPooling_Op(const std::array<Aidge::DimSize_t, DIM> &kernel_dims,
                            const std::array<Aidge::DimSize_t, DIM> &stride_dims,
                            bool ceil_mode)
    : OperatorTensor(Type, {InputCategory::Data}, 1),
    mAttributes(std::make_shared<Attributes_>(
    attr<MaxPoolingAttr::StrideDims>(stride_dims),
    attr<MaxPoolingAttr::KernelDims>(kernel_dims),
    attr<MaxPoolingAttr::CeilMode>(ceil_mode)))
{}

template <Aidge::DimIdx_t DIM>
Aidge::MaxPooling_Op<DIM>::MaxPooling_Op(const Aidge::MaxPooling_Op<DIM>& op)
    : OperatorTensor(op),
    mAttributes(op.mAttributes)
{
    if (op.mImpl) {
        SET_IMPL_MACRO(MaxPooling_Op<DIM>, *this, op.backend());
    } else {
        mImpl = nullptr;
    }
}

template <Aidge::DimIdx_t DIM>
std::shared_ptr<Aidge::Operator> Aidge::MaxPooling_Op<DIM>::clone() const {
    return std::make_shared<MaxPooling_Op<DIM>>(*this);
}

template <Aidge::DimIdx_t DIM>
bool Aidge::MaxPooling_Op<DIM>::forwardDims(bool /*allowDataDependency*/) {
    if (inputsAssociated()) {
        std::array<DimSize_t, DIM + 2> outputDims{};
        const std::array<DimSize_t, DIM + 2> inputDims(getInput(0)->template dims<DIM+2>());

        std::function<float(float)> roundingFunction;
        if (mAttributes->template getAttr<MaxPoolingAttr::CeilMode>()) {
            roundingFunction = [](float x) { return std::ceil(x); };
        } else {
            roundingFunction = [](float x) { return std::floor(x); };
        }

        for (std::size_t dim = 0; dim < mAttributes->template getAttr<MaxPoolingAttr::KernelDims>().size() ; ++dim) {
            outputDims[dim+2] = 1 + static_cast<DimSize_t>(
                                        roundingFunction(static_cast<float>(inputDims[dim+2] -
                                                                mAttributes->template getAttr<MaxPoolingAttr::KernelDims>()[dim]) /
                                        static_cast<float>(mAttributes->template getAttr<MaxPoolingAttr::StrideDims>()[dim])));
        }
        outputDims[1] = inputDims[1];
        outputDims[0] = inputDims[0];
        mOutputs[0]->resize(outputDims);
        return true;
    }
    return false;
}

template <Aidge::DimIdx_t DIM>
void Aidge::MaxPooling_Op<DIM>::setBackend(const std::string &name, Aidge::DeviceIdx_t device) {
    SET_IMPL_MACRO(MaxPooling_Op<DIM>, *this, name);
    mOutputs[0]->setBackend(name, device);
}

template <Aidge::DimIdx_t DIM>
std::set<std::string> Aidge::MaxPooling_Op<DIM>::getAvailableBackends() const {
    return Registrar<MaxPooling_Op<DIM>>::getKeys();
}

template class Aidge::MaxPooling_Op<1>;
template class Aidge::MaxPooling_Op<2>;
template class Aidge::MaxPooling_Op<3>;

///////////////////////////////////////////

template <std::array<Aidge::DimSize_t, 1>::size_type DIM>
std::shared_ptr<Aidge::Node> Aidge::MaxPooling(const std::array<Aidge::DimSize_t, DIM> &kernel_dims,
                                           const std::string& name,
                                           const std::array<Aidge::DimSize_t, DIM> &stride_dims,
                                           bool ceil_mode)
{
    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, ceil_mode), name);
}

template std::shared_ptr<Aidge::Node> Aidge::MaxPooling<1>(const std::array<Aidge::DimSize_t, 1>&, const std::string&, const std::array<Aidge::DimSize_t, 1>&, bool);
template std::shared_ptr<Aidge::Node> Aidge::MaxPooling<2>(const std::array<Aidge::DimSize_t, 2>&, const std::string&, const std::array<Aidge::DimSize_t, 2>&, bool);
template std::shared_ptr<Aidge::Node> Aidge::MaxPooling<3>(const std::array<Aidge::DimSize_t, 3>&, const std::string&, const std::array<Aidge::DimSize_t, 3>&, bool);