model.cpp

#define USE_MEMPOOL

#include "model/model.h"
#include "aclnn/aclnn_gelu_operation.h"
#include "utils/utils.h"
#include "atb/atb_graph_op.h"
#include "memory/memory_utils.h"

void Model::InitResource(uint32_t deviceId)
{
    // Set deviceId.
    deviceId_ = deviceId;
    auto ret = aclrtSetDevice(deviceId_);
    CHECK_RET(ret, "aclrtSetDevice failed. ret: " + std::to_string(ret));

    // Create a context.
    ret = atb::CreateContext(&modeContext_);
    CHECK_RET(ret, "ATB CreateContext failed. ret: " + std::to_string(ret));

    // Create a stream.
    ret = aclrtCreateStream(&modelStream_);
    CHECK_RET(ret, "aclrtCreateStream failed. ret: " + std::to_string(ret));

    // Configure the stream.
    modeContext_->SetExecuteStream(modelStream_);
}

void Model::CreateModelGraph()
{
    LOG_INFO("CreateModelGraph start");
    // Assume that two nodes in the model are used as an example.
    nodes_.resize(2);
    for (size_t i = 0; i < nodes_.size(); i++) {
        auto node = Node();
        nodes_[i] = node;
    }

    modelInTensors_.resize(Mode_INPUT_SIZE);
    modelOutTensors_.resize(Mode_OUTPUT_SIZE);

    internalTensors_.resize(1);
    size_t nodeId = 0;
    CreateGraphOpLayer(nodeId++);

    // Step 2: Create a node for the aclnn operator.
    CreateAclnnOpLayer(nodeId);
    LOG_INFO("CreateModelGraph end");
}

void Model::CreateGraphOpLayer(size_t nodeId)
{
    // Create the operation of the graph operator.
    Node &graph_node = nodes_[nodeId];
    auto ret = CreateGraphOperation(&graph_node.operation_);
    CHECK_RET(ret, "CreateGraphOperation failed");
    graph_node.inTensors_.resize(graph_node.operation_->GetInputNum());

    // Set the input of the graph operator node.
    // The input of the graph operator is the input of the entire model. Therefore, the value is directly assigned from inTensors_ of the model.
    size_t layerInTensorId = 0;
    graph_node.inTensors_.at(layerInTensorId++) = &modelInTensors_.at(IN_TENSOR_A);
    graph_node.inTensors_.at(layerInTensorId++) = &modelInTensors_.at(IN_TENSOR_B);
    graph_node.inTensors_.at(layerInTensorId++) = &modelInTensors_.at(IN_TENSOR_C);
    graph_node.inTensors_.at(layerInTensorId++) = &modelInTensors_.at(IN_TENSOR_D);

    // Set the output of the graph operator node. There is only one intermediate node.
    graph_node.outTensors_ = {&internalTensors_.at(0)};
    graph_node.outTensorTypes_ = {TensorType::INTERNAL_TENSOR};
};

void Model::CreateAclnnOpLayer(size_t nodeId)
{
    // Create the operation of the aclnn operator.
    Node &aclnn_node = nodes_[nodeId];
    AclnnGeluParam AclnnGeluParam;
    AclnnGeluParam.geluApproximate = -1;
    aclnn_node.operation_ = new GeluOperation("Gelu", AclnnGeluParam);
    aclnn_node.inTensors_.resize(aclnn_node.operation_->GetInputNum());

    // Set the input of the aclnn operator node.
    // The output of the graph operator is the input of the aclnn operator.
    size_t layerInTensorId = 0;
    aclnn_node.inTensors_.at(layerInTensorId++) = &internalTensors_.at(0);

    // Set the output of the aclnn operator node, that is, the output of the model.
    aclnn_node.outTensors_ = {&modelOutTensors_.at(GLUE_OUT)};
    aclnn_node.outTensorTypes_ = {TensorType::NOT_INTERNAL_TENSOR};
}

void Model::CreateModelInput()
{
    LOG_INFO("CreateModelInput start");
    atb::SVector<atb::TensorDesc> intensorDescs;
    intensorDescs.resize(Mode_INPUT_SIZE);
    CreateInTensorDescs(intensorDescs);
    CreateInTensors(modelInTensors_, intensorDescs);
    LOG_INFO("CreateModelInput end");
}

void Model::CreateModelOutput()
{
    LOG_INFO("CreateModelOutput start");
    atb::SVector<atb::TensorDesc> outtensorDescs;
    outtensorDescs.resize(Mode_OUTPUT_SIZE);

    // Set the input description.
    atb::SVector<atb::TensorDesc> inTensorDescs;
    inTensorDescs.resize(Mode_INPUT_SIZE);
    for (size_t i = 0; i < modelInTensors_.size(); ++i) {
        inTensorDescs.at(i) = modelInTensors_.at(i).desc;
    }

    // Call infer shape to deduce the model output.
    InferShape(inTensorDescs, outtensorDescs);
    CreateOutTensors(modelOutTensors_, outtensorDescs);
    LOG_INFO("CreateModelOutput end");
}

atb::Status Model::InferShape(
    const atb::SVector<atb::TensorDesc> &inTensorDescs, atb::SVector<atb::TensorDesc> &outTensorDescs)
{
    // The output shape is the same as the input shape. You can use the first input.
    outTensorDescs.at(0) = modelInTensors_.at(0).desc;
    return atb::NO_ERROR;
}

void Model::Execute()
{
    LOG_INFO(modelName_ + " Execute start");
    for (size_t nodeId = 0; nodeId < nodes_.size(); ++nodeId) {
        BuildNodeVariantPack(nodeId);
        atb::Status status = ExecuteNode(nodeId);
        CHECK_RET(status, "ExecuteNode " + std::to_string(nodeId) + " failed. status: " + std::to_string(status));
    }

    WaitFinish();
    LOG_INFO(modelName_ + " Execute end");
}

void Model::BuildNodeVariantPack(int nodeId)
{
    LOG_INFO("buildNodeVariantPack nodes[" + std::to_string(nodeId) + "] start");

    auto &node = nodes_.at(nodeId);
    atb::SVector<atb::TensorDesc> inTensorDescs;
    node.variantPack_.inTensors.resize(node.operation_->GetInputNum());
    inTensorDescs.resize(node.operation_->GetInputNum());

    // Obtain the input tensor description of operation_ in the node.
    for (size_t i = 0; i < node.inTensors_.size(); ++i) {
        node.variantPack_.inTensors.at(i) = *node.inTensors_.at(i);
        inTensorDescs.at(i) = node.inTensors_.at(i)->desc;
    }

    atb::SVector<atb::TensorDesc> outTensorDescs;
    outTensorDescs.resize(node.operation_->GetOutputNum());

    // Call InferShape of operation_ to deduce the description of the output tensor.
    atb::Status st = node.operation_->InferShape(inTensorDescs, outTensorDescs);

    node.variantPack_.outTensors.resize(node.operation_->GetOutputNum());
    for (size_t i = 0; i < node.outTensors_.size(); ++i) {
        node.variantPack_.outTensors.at(i) = *node.outTensors_.at(i);
        if (node.outTensorTypes_.at(i) == TensorType::INTERNAL_TENSOR) {
            // Create the output tensor space.
            CreateTensorFromDesc(node.variantPack_.outTensors.at(i), outTensorDescs.at(i));
            *node.outTensors_.at(i) = node.variantPack_.outTensors.at(i);
        }
    }
    LOG_INFO("buildNodeVariantPack nodes[" + std::to_string(nodeId) + "] end");
}

atb::Status Model::ExecuteNode(int nodeId)
{
    auto &node = nodes_.at(nodeId);

    // Call the Setup API.
    uint64_t workspaceSize = 0;
    atb::Status status = node.operation_->Setup(node.variantPack_, workspaceSize, modeContext_);
    CHECK_RET(status, "Setup node " + std::to_string(nodeId) + " failed. status: " + std::to_string(status));

    LOG_INFO("Get node[" + std::to_string(nodeId) + "] workspace size:" + std::to_string(workspaceSize));

    // Allocate workspace.
#ifdef USE_MEMPOOL
    CreateWorkspaceBuffer(nodeId, workspaceSize);
#else
    if (workspaceSize != 0) {
        status = aclrtMalloc(&node.workspace_, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
        CHECK_RET(status, "alloc error!");
    }
#endif

    // Call the Execute API.
    LOG_INFO("Execute node[" + std::to_string(nodeId) + "] start");
    status = node.operation_->Execute(node.variantPack_, (uint8_t *)(node.workspace_), workspaceSize, modeContext_);
    CHECK_RET(status, "Execute node " + std::to_string(nodeId) + " failed. status: " + std::to_string(status));
    LOG_INFO("Execute node[" + std::to_string(nodeId) + "] end");
    return atb::NO_ERROR;
}

void Model::CreateWorkspaceBuffer(int nodeId, int workspaceSizeNeeded)
{
    auto &node = nodes_.at(nodeId);
    if (workspaceSizeNeeded == 0) {
        LOG_INFO("skip the workspacebuffer for size 0");
        return;
    }
    if (node.workspaceBlockId_ == -1 || node.workspaceSize_ == 0) {
        node.workspaceSize_ = workspaceSizeNeeded;
        GetMemoryManager().AllocateBlock(node.workspaceSize_, node.workspaceBlockId_);
    }
    if (node.workspaceSize_ < workspaceSizeNeeded) {
        GetMemoryManager().FreeBlock(node.workspaceBlockId_);
        GetMemoryManager().AllocateBlock(workspaceSizeNeeded, node.workspaceBlockId_);
        node.workspaceSize_ = workspaceSizeNeeded;
    }

    GetMemoryManager().GetBlockPtr(node.workspaceBlockId_, node.workspace_);
}

void Model::FreeResource()
{
    LOG_INFO("FreeResource start");
    auto status = aclrtDestroyStream(modelStream_);  // Destroy the stream.
    CHECK_RET(status, "aclrtDestroyStream failed");

    // Release the operation.
    for (auto &node : nodes_) {
        atb::DestroyOperation(node.operation_);
#ifdef USE_MEMPOOL
        GetMemoryManager().FreeBlock(node.workspaceBlockId_);
#endif
    }

    status = atb::DestroyContext (modeContext_); // Destroy the context.
    CHECK_RET(status, "aclrtDestroyStream failed");

    // Destroy the input tensor.
    for (size_t i = 0; i < modelInTensors_.size(); i++) {
        aclrtFree(modelInTensors_.at(i).deviceData);
    }

    // Destroy the output tensor.
    for (size_t i = 0; i < modelOutTensors_.size(); i++) {
        aclrtFree(modelOutTensors_.at(i).deviceData);
    }

    // Release the intermediate tensor.
    for (size_t i = 0; i < internalTensors_.size(); i++) {
        aclrtFree(internalTensors_.at(i).deviceData);
    }

    aclrtResetDevice (deviceId_); // Reset the device ID.
    LOG_INFO("FreeResource end");
}

void Model::WaitFinish()
{
    // Step 9: Destroy the created object and release the memory.
    // Synchronize the stream and wait until the computation on the device is complete.
    auto ret = aclrtSynchronizeStream(modelStream_);
    CHECK_RET(ret, "sync error!");
}