Unfold.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/Unfold.hpp"

#include <cmath>      // std::floor
#include <cstddef>    // std::size_t
#include <stdexcept>  // std::runtime_error
#include <string>
#include <utility>    // std::pair
#include <vector>

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

template <Aidge::DimIdx_t DIM>
void Aidge::Unfold_OpImpl<DIM>::forward() {
    const Unfold_Op<DIM>& op = dynamic_cast<const Unfold_Op<DIM>&>(mOp);
    const auto kernelDims = op.kernelDims();
    const auto dilationDims = op.dilationDims();
    const auto strideDims = op.strideDims();
    const DimSize_t inHeight = op.getInput(0)->dims()[2];
    const DimSize_t inWidth = op.getInput(0)->dims()[3];
    const DimSize_t inChannels = op.getInput(0)->dims()[1];

    const DimSize_t kernelExtentHeight = op.dilationDims()[0] *
                                            (op.kernelDims()[0] - 1) + 1;
    const DimSize_t outHeight = 1 + static_cast<DimSize_t>(
                    floor(static_cast<float>(inHeight - kernelExtentHeight) /
                            static_cast<float>(op.strideDims()[0])));
    const DimSize_t kernelExtentWidth = op.dilationDims()[1] *
                                            (op.kernelDims()[1] - 1) + 1;
    const DimSize_t outWidth = 1 + static_cast<DimSize_t>(
                    floor(static_cast<float>(inWidth - kernelExtentWidth) /
                            static_cast<float>(op.strideDims()[1])));
    const DimSize_t outChannels = op.getOutput(0)->dims()[1];

    for (DimSize_t n = 0; n < op.getOutput(0)->dims()[0]; ++n) {
        for (DimSize_t outC = 0; outC < outChannels; ++outC) {
            const auto inOffsetW = outC % kernelDims[1];
            const auto inOffsetH = (outC / kernelDims[1]) % kernelDims[0];
            const auto inC = outC / kernelDims[0] / kernelDims[1];

            for (DimSize_t outH = 0; outH < outHeight; ++outH) {
                const auto inH = outH * strideDims[0] + inOffsetH * dilationDims[0];

                for (DimSize_t outW = 0; outW < outWidth; ++outW) {
                    const auto inW = outW * strideDims[1] + inOffsetW * dilationDims[1];

                    op.getOutput(0)->getImpl()->copy(op.getInput(0)->getImpl()->rawPtr(((n * inChannels + inC) * inHeight + inH) * inWidth + inW), 1,
                        ((n * outChannels + outC) * outHeight + outH) * outWidth + outW);
                }
            }
        }
    }
}

template class Aidge::Unfold_OpImpl<2>;

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

template <Aidge::DimIdx_t DIM>
const std::string Aidge::Unfold_Op<DIM>::Type = "Unfold";

template <Aidge::DimIdx_t DIM>
Aidge::Unfold_Op<DIM>::Unfold_Op(const std::array<Aidge::DimSize_t, DIM> &kernelDims,
                    const std::array<Aidge::DimSize_t, DIM> &strideDims,
                    const std::array<Aidge::DimSize_t, DIM> &dilationDims)
    : OperatorTensor(Type, {InputCategory::Data}, 1),
        mAttributes(std::make_shared<Attributes_>(
        attr<UnfoldAttr::StrideDims>(strideDims),
        attr<UnfoldAttr::DilationDims>(dilationDims),
        attr<UnfoldAttr::KernelDims>(kernelDims)))
{
    mImpl = std::make_shared<Unfold_OpImpl<DIM>>(*this);
}

template <Aidge::DimIdx_t DIM>
Aidge::Unfold_Op<DIM>::Unfold_Op(const Aidge::Unfold_Op<DIM> &op)
    : OperatorTensor(op),
        mAttributes(op.mAttributes)
{
    if (!op.backend().empty()) {
        SET_IMPL_MACRO(Unfold_Op<DIM>, *this, op.backend());
    }
    else {
        mImpl = std::make_shared<Unfold_OpImpl<DIM>>(*this);
    }
}

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

template <Aidge::DimIdx_t DIM>
bool Aidge::Unfold_Op<DIM>::forwardDims(bool /*allowDataDependency*/) {
    if (inputsAssociated()) {
        const std::array<DimSize_t, DIM + 2> inputDims(getInput(0)->template dims<DIM+2>());
        DimSize_t k = 1;
        DimSize_t l = 1;

        for (std::size_t dim = 0; dim < this->kernelDims().size() ; ++dim) {
            const DimSize_t kernelExtent = this->dilationDims()[dim] *
                                                    (this->kernelDims()[dim] - 1) + 1;

            k *= this->kernelDims()[dim];
            l *= 1 + static_cast<DimSize_t>(
                    floor(static_cast<float>(inputDims[dim+2] - kernelExtent) /
                            static_cast<float>(this->strideDims()[dim])));
        }

        mOutputs[0]->resize({inputDims[0], inputDims[1] * k, l});
        return true;
    }

    return false;
}

template <Aidge::DimIdx_t DIM>
void Aidge::Unfold_Op<DIM>::setBackend(const std::string &name, Aidge::DeviceIdx_t device) {
    if (Registrar<Unfold_Op<DIM>>::exists({name})){
        SET_IMPL_MACRO(Unfold_Op<DIM>, *this, name);
    }
    else {
        mImpl = std::make_shared<Unfold_OpImpl<DIM>>(*this);
    }
    mOutputs[0]->setBackend(name, device);
}

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

template class Aidge::Unfold_Op<2>;

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

template <std::array<Aidge::DimSize_t, 1>::size_type DIM>
std::shared_ptr<Aidge::Node> Aidge::Unfold(const std::array<Aidge::DimSize_t, DIM> &kernelDims,
                                  const std::string& name,
                                  const std::array<Aidge::DimSize_t, DIM> &strideDims,
                                  const std::array<Aidge::DimSize_t, DIM> &dilationDims) {
    static_assert(DIM<=MaxDim,"Too many kernel dimensions required by Unfold, not supported");
    return std::make_shared<Node>(std::make_shared<Unfold_Op<static_cast<DimIdx_t>(DIM)>>(kernelDims, strideDims, dilationDims), name);
}

template std::shared_ptr<Aidge::Node> Aidge::Unfold<2>(const std::array<Aidge::DimSize_t, 2>&,
                                  const std::string&,
                                  const std::array<Aidge::DimSize_t, 2>&,
                                  const std::array<Aidge::DimSize_t, 2>&);