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

#include <memory>
#include <string>

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

namespace Aidge {

Aidge::Elts_t Aidge::StackProdConso::getRequiredMemory(
    const Aidge::IOIndex_t inputIdx,
    const std::vector<DimSize_t> &inputsSize) const {
    assert(mOp.getRawInput(inputIdx) && "requires valid input");

    const StackOp &op = dynamic_cast<const StackOp &>(mOp);
    // The produced data after one forward pass is simply the input size,
    // we do not produced the whole output tensor everytime.
    // The output tensor it set to its max dimensions just to some dimensions
    // to forward.
    return Elts_t::DataElts(op.getInput(inputIdx)->size());
}

const std::string StackOp::s_type = "Stack";

void StackOpImpl::forward() {
    const StackOp &op = dynamic_cast<const StackOp &>(mOp);
    AIDGE_ASSERT(op.getInput(0), "missing input #0");
    AIDGE_ASSERT((op.forwardStep() < op.maxElements()),
                 "cannot forward anymore, number of cycles exceeded");

    op.getOutput(0)->getImpl()->copy(
        op.getInput(0)->getImpl()->rawPtr(),
        op.getInput(0)->size(),
        op.forwardStep() * op.getInput(0)->size());
}

StackOp::StackOp(std::uint32_t maxElements)
    : OperatorTensor(s_type, {InputCategory::Data}, 1),
      mAttributes(std::make_shared<Attributes_>(
          attr<StackAttr::MaxElements>(maxElements),
          attr<StackAttr::ForwardStep>(0))) {
    if (maxElements == 0) {
        AIDGE_THROW_OR_ABORT(std::invalid_argument, "StackOp creation failed: maxElements must be greater than 0.");
    }
    mImpl = std::make_shared<StackOpImpl>(*this);
}

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

std::shared_ptr<Aidge::Operator> Aidge::StackOp::clone() const {
    return std::make_shared<StackOp>(*this);
}

bool Aidge::StackOp::forwardDims(bool /*allowDataDependency*/) {
    if (inputsAssociated()) {
        auto inputDims = getInput(0)->dims();
        inputDims.insert(inputDims.begin(), maxElements());
        getOutput(0)->resize(inputDims);
        return true;
    }

    return false;
}

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

std::set<std::string> StackOp::getAvailableBackends() const {
    return Registrar<StackOp>::getKeys();
}

void StackOp::forward() {
    Log::info("fw step {}", forwardStep());
    Operator::forward();
    ++forwardStep();
}

std::shared_ptr<Node> stack(std::uint32_t maxElements,
                            const std::string &name) {
    return std::make_shared<Node>(std::make_shared<StackOp>(maxElements),
                                  name);
}
} // namespace Aidge