PaddedConvImpl.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/backend/cpu/operator/PaddedConvImpl.hpp"
#include "aidge/backend/cpu/operator/PaddedConvImpl_kernels.hpp"

#include <memory>
#include <vector>

#include "aidge/backend/cpu/data/GetCPUPtr.h"
#include "aidge/backend/OperatorImpl.hpp"
#include "aidge/data/Tensor.hpp"
#include "aidge/operator/MetaOperator.hpp"
#include "aidge/operator/Conv.hpp"
#include "aidge/operator/Pad.hpp"
#include "aidge/utils/ErrorHandling.hpp"
#include "aidge/utils/Registrar.hpp"
#include "aidge/utils/Types.h"
#include "aidge/backend/cpu/data/GetCPUPtr.h"

template <>
void Aidge::PaddedConvImpl1D_cpu::forward() {
    const auto& op_ = static_cast<const MetaOperator_Op&>(mOp);

    // FIXME: uncomment the following code once memory handling will work
    AIDGE_ASSERT(op_.getInput(0), "missing input #0 in Conv Operator.");
    AIDGE_ASSERT(op_.getInput(1), "missing input #1 in Conv Operator.");

    // Find the correct kernel type
    const auto impl = Registrar<PaddedConvImpl1D_cpu>::create(getBestMatch(getRequiredSpec()));

    // Convert input data (no overhead if not needed!)
    // TODO: right now, if needed, memory will be allocated/deallocated at each
    // call to forward(). We might put the following shared_ptr as members of
    // this class to avoid that.
    std::shared_ptr<Tensor> input0Fallback, input1Fallback, input2Fallback;
    const auto& input0 = op_.getInput(0)->refCastFrom(input0Fallback, *op_.getOutput(0));
    const auto& input1 = op_.getInput(1)->refCastFrom(input1Fallback, *op_.getOutput(0));
    const auto& input2 = (op_.getInput(2)) ? op_.getInput(2)->refCastFrom(input2Fallback, *op_.getOutput(0)) : Tensor();

    std::shared_ptr<Conv_Op<1>> conv_op;
    std::shared_ptr<Pad_Op<1>> pad_op;
    for (const auto& n : op_.getMicroGraph()->getNodes()) {
        if (n->getOperator()->type() == Conv_Op<1>::Type) {
            conv_op = std::static_pointer_cast<Conv_Op<1>>(n->getOperator());
        } else {
            pad_op =  std::static_pointer_cast<Pad_Op<1>>(n->getOperator());
        }
    }

    // Call kernel
    impl.forward(
            pad_op->beginEndBorders(),
            conv_op->strideDims(),
            conv_op->dilationDims(),
            conv_op->kernelDims(),
            op_.getInput(0)->template dims<3>(), // input dimensions
            conv_op->outChannels(), // outChannels
            input0.getImpl()->rawPtr(), // input
            input1.getImpl()->rawPtr(), // weight
            op_.getInput(2) ? input2.getImpl()->rawPtr() : nullptr, // bias
            getCPUPtr(mOp.getRawOutput(0)) // output
            );
}

template <>
void Aidge::PaddedConvImpl1D_cpu::backward() {
    AIDGE_THROW_OR_ABORT(std::runtime_error, "Backward not yet implemented for Conv_Op<1> on backend cpu");
}

template <>
void Aidge::PaddedConvImpl2D_cpu::forward() {
    const auto& op_ = dynamic_cast<const MetaOperator_Op&>(mOp);

    // FIXME: uncomment the following code once memory handling will work
    AIDGE_ASSERT(op_.getInput(0), "missing input #0 in Conv Operator.");
    AIDGE_ASSERT(op_.getInput(1), "missing input #1 in Conv Operator.");

    // Find the correct kernel type
    const auto impl = Registrar<PaddedConvImpl2D_cpu>::create(getBestMatch(getRequiredSpec()));

    // Convert input data (no overhead if not needed!)
    // TODO: right now, if needed, memory will be allocated/deallocated at each
    // call to forward(). We might put the following shared_ptr as members of
    // this class to avoid that.
    std::shared_ptr<Tensor> input0Fallback, input1Fallback, input2Fallback;
    const auto& input0 = op_.getInput(0)->refCastFrom(input0Fallback, *op_.getOutput(0));
    const auto& input1 = op_.getInput(1)->refCastFrom(input1Fallback, *op_.getOutput(0));
    const auto& input2 = (op_.getInput(2)) ? op_.getInput(2)->refCastFrom(input2Fallback, *op_.getOutput(0)) : Tensor();

    std::shared_ptr<Conv_Op<2>> conv_op;
    std::shared_ptr<Pad_Op<2>> pad_op;

    for (const auto& n : op_.getMicroGraph()->getNodes()) {
        if (n->getOperator()->type() == Conv_Op<2>::Type) {
            conv_op = std::static_pointer_cast<Conv_Op<2>>(n->getOperator());
        } else {
            pad_op =  std::static_pointer_cast<Pad_Op<2>>(n->getOperator());
        }
    }

    // Call kernel
    impl.forward(
            pad_op->beginEndBorders(),
            conv_op->strideDims(),
            conv_op->dilationDims(),
            conv_op->kernelDims(),
            op_.getInput(0)->template dims<4>(), // input dimensions
            conv_op->outChannels(), // outChannels
            input0.getImpl()->rawPtr(), // input
            input1.getImpl()->rawPtr(), // weight
            op_.getInput(2) ? input2.getImpl()->rawPtr() : nullptr, // bias
            getCPUPtr(mOp.getRawOutput(0)) // output
            );
}

template <>
void Aidge::PaddedConvImpl2D_cpu::backward() {
    AIDGE_THROW_OR_ABORT(std::runtime_error, "Backward not yet implemented for Conv_Op<2> on backend cpu");
}