ATB Graph Offloading

source ${toolkit installation directory}/set_env.sh # For example, source /usr/local/Ascend/ascend-toolkit/set_env.sh
source ${nnal installation directory}/atb/set_env.sh # For example, source /usr/local/Ascend/nnal/atb/set_env.sh

#include <atb/operation.h>
#include "atb/infer_op_params.h"
#include <atb/types.h>
#include <acl/acl.h>
#include <iostream>
#include <vector>
#include <atb/utils.h>
#include "experiment/runtime/runtime/rt_model.h"
#include "experiment/runtime/runtime/stream.h"
#include <unistd.h>
#include <acl/acl_mdl.h>
#include <cstdlib>  // std::getenv included.
#include <cstring>  // std::strcmp included (used for string comparison)

// Set the properties of each intensor.
void CreateInTensorDescs(atb::SVector<atb::TensorDesc> &intensorDescs) 
{
    for (size_t i = 0; i < intensorDescs.size(); i++) {
        intensorDescs.at(i).dtype = ACL_FLOAT16;
        intensorDescs.at(i).format = ACL_FORMAT_ND;
        intensorDescs.at(i).shape.dimNum = 2;
        intensorDescs.at(i).shape.dims[0] = 2;
        intensorDescs.at(i).shape.dims[1] = 2;
    }
}

// Set each intensor and allocate memory space for each intensor. The intensor is manually set. In the project implementation, you can use torchTensor conversion or other simple data structure conversion.
void CreateInTensors(atb::SVector<atb::Tensor> &inTensors, atb::SVector<atb::TensorDesc> &intensorDescs, uint32_t value)
{
    for (size_t i = 0; i < inTensors.size(); i++) {
        inTensors.at(i).desc = intensorDescs.at(i);
        inTensors.at(i).dataSize = atb::Utils::GetTensorSize(inTensors.at(i));
        std::vector<uint16_t> hostData(atb::Utils::GetTensorNumel(inTensors.at(i)), value);   // A host buffer filled with all twos.
        int ret = aclrtMalloc(&inTensors.at(i).deviceData, inTensors.at(i).dataSize, ACL_MEM_MALLOC_HUGE_FIRST); // Allocate NPU memory.
        if (ret != 0) {
            std::cout << "alloc error!";
            exit(0);
        }
        ret = aclrtMemcpy(inTensors.at(i).deviceData, inTensors.at(i).dataSize, hostData.data(), hostData.size() * sizeof(uint16_t), ACL_MEMCPY_HOST_TO_DEVICE); //Copy the CPU memory to the NPU.
    }
}

// Set each outtensor and allocate memory space for each outtensor. The setting is the same as that of an intensor.
void CreateOutTensors(atb::SVector<atb::Tensor> &outTensors, atb::SVector<atb::TensorDesc> &outtensorDescs)
{
    for (size_t i = 0; i < outTensors.size(); i++) {
        outTensors.at(i).desc = outtensorDescs.at(i);
        outTensors.at(i).dataSize = atb::Utils::GetTensorSize(outTensors.at(i));
        int ret = aclrtMalloc(&outTensors.at(i).deviceData, outTensors.at(i).dataSize, ACL_MEM_MALLOC_HUGE_FIRST);
        if (ret != 0) {
            std::cout << "alloc error!";
            exit(0);
        }
    }
}

// Build subgraphs.
void CreateGraphOperation(atb::GraphParam &opGraph, atb::Operation **operation)
{
    opGraph.inTensorNum = 4;
    opGraph.outTensorNum = 1;
    opGraph.internalTensorNum = 2;
    opGraph.nodes.resize(3);

    enum InTensorId {
        IN_TENSOR_A = 0,
        IN_TENSOR_B,
        IN_TENSOR_C,
        IN_TENSOR_D,
        ADD3_OUT,
        ADD1_OUT,
        ADD2_OUT
    };

    size_t nodeId = 0;
    atb::Node &addNode = opGraph.nodes.at(nodeId++);
    atb::Node &addNode2 = opGraph.nodes.at(nodeId++);
    atb::Node &addNode3 = opGraph.nodes.at(nodeId++);

    atb::infer::ElewiseParam addParam;
    addParam.elewiseType = atb::infer::ElewiseParam::ElewiseType::ELEWISE_ADD;
    atb::Status status = atb::CreateOperation(addParam, &addNode.operation);
    addNode.inTensorIds = {IN_TENSOR_A, IN_TENSOR_B};
    addNode.outTensorIds = {ADD1_OUT};

    atb::infer::ElewiseParam addParam2;
    addParam2.elewiseType = atb::infer::ElewiseParam::ElewiseType::ELEWISE_ADD;
    status = atb::CreateOperation(addParam2, &addNode2.operation);
    addNode2.inTensorIds = {IN_TENSOR_C, IN_TENSOR_D};
    addNode2.outTensorIds = {ADD2_OUT};

    atb::infer::ElewiseParam addParam3;
    addParam3.elewiseType = atb::infer::ElewiseParam::ElewiseType::ELEWISE_ADD;
    status = CreateOperation(addParam3, &addNode3.operation);
    addNode3.inTensorIds = {ADD1_OUT, ADD2_OUT};
    addNode3.outTensorIds = {ADD3_OUT};

    status = atb::CreateOperation(opGraph, operation);
}

// Build graphs.
void CreateMultiGraphOperation(atb::GraphParam &opGraph, atb::Operation **operation)
{
    opGraph.inTensorNum = 4;
    opGraph.outTensorNum = 1;
    opGraph.internalTensorNum = 4;
    opGraph.nodes.resize(5);

    enum InTensorId {
        IN_TENSOR_A = 0,
        IN_TENSOR_B,
        IN_TENSOR_C,
        IN_TENSOR_D,
        ADD5_OUT,
        ADD1_INTER,
        ADD2_INTER,
        ADD3_INTER,
        ADD4_INTER
    };

    size_t nodeId = 0;
    atb::Node &addGraphNode1 = opGraph.nodes.at(nodeId++);
    atb::Node &addGraphNode2 = opGraph.nodes.at(nodeId++);
    atb::Node &addGraphNode3 = opGraph.nodes.at(nodeId++);
    atb::Node &addGraphNode4 = opGraph.nodes.at(nodeId++);
    atb::Node &addGraphNode5 = opGraph.nodes.at(nodeId++);

    atb::GraphParam addGraphParam1;
    CreateGraphOperation(addGraphParam1, &addGraphNode1.operation);
    addGraphNode1.inTensorIds = {IN_TENSOR_A, IN_TENSOR_B, IN_TENSOR_C, IN_TENSOR_D};
    addGraphNode1.outTensorIds = {ADD1_INTER};

    atb::GraphParam addGraphParam2;
    CreateGraphOperation(addGraphParam2, &addGraphNode2.operation);
    addGraphNode2.inTensorIds = {IN_TENSOR_A, IN_TENSOR_B, IN_TENSOR_C, IN_TENSOR_D};
    addGraphNode2.outTensorIds = {ADD2_INTER};

    atb::GraphParam addGraphParam3;
    CreateGraphOperation(addGraphParam3, &addGraphNode3.operation);
    addGraphNode3.inTensorIds = {IN_TENSOR_A, IN_TENSOR_B, IN_TENSOR_C, IN_TENSOR_D};
    addGraphNode3.outTensorIds = {ADD3_INTER};

    atb::GraphParam addGraphParam4;
    CreateGraphOperation(addGraphParam4, &addGraphNode4.operation);
    addGraphNode4.inTensorIds = {IN_TENSOR_A, IN_TENSOR_B, IN_TENSOR_C, IN_TENSOR_D};
    addGraphNode4.outTensorIds = {ADD4_INTER};

    atb::GraphParam addGraphParam5;
    CreateGraphOperation(addGraphParam5, &addGraphNode5.operation);
    addGraphNode5.inTensorIds = {ADD1_INTER, ADD2_INTER, ADD3_INTER, ADD4_INTER};
    addGraphNode5.outTensorIds = {ADD5_OUT};

    atb::Status status = atb::CreateOperation(opGraph, operation);
}

// Print the result.
void PrintOutTensorValue(atb::Tensor &outTensor)
{
    std::vector<uint16_t> outBuffer(atb::Utils::GetTensorNumel(outTensor));
    int ret = aclrtMemcpy(outBuffer.data(), outBuffer.size() * sizeof(uint16_t), outTensor.deviceData, outTensor.dataSize, ACL_MEMCPY_DEVICE_TO_HOST);
    if (ret != 0) {
        std::cout << "copy error!";
        exit(0);
    }
    for (size_t i = 0; i < outBuffer.size(); i = i + 1) {
        std::cout << "out[" << i << "] = " << (uint32_t)outBuffer.at(i) << std::endl;
    }
}

int main()
{
    // Use the environment variable ATB_CAPTURE_TYPE to determine whether the internal or external capture is used.
    // ATB_CAPTURE_TYPE=in: internal capture
    const char* captureType = std::getenv("ATB_CAPTURE_TYPE");
    aclInit(nullptr);
    // Set a device ID, create a stream, create a context, and set the stream.
    uint32_t deviceId = 1;
    aclrtSetDevice(deviceId);
    aclrtStream exeStream = nullptr;
    aclrtCreateStream(&exeStream);

    atb::Context *context = nullptr;
    atb::CreateContext(&context);
    context->SetExecuteStream(exeStream);
    context->SetLaunchMode(atb::GRAPH_LAUNCH_MODE);

    // Create a graph operator.
    atb::Operation *operation = nullptr;
    atb::GraphParam opGraph;
    CreateMultiGraphOperation(opGraph, &operation);

    // Prepare the input and output tensors.
    atb::VariantPack pack;
    atb::SVector<atb::TensorDesc> intensorDescs;
    atb::SVector<atb::TensorDesc> outtensorDescs;

    uint32_t inTensorNum = opGraph.inTensorNum;
    uint32_t outTensorNum = opGraph.outTensorNum;
    inTensorNum = operation->GetInputNum();
    outTensorNum = operation->GetOutputNum();

    pack.inTensors.resize(inTensorNum);
    intensorDescs.resize(inTensorNum);

    CreateInTensorDescs(intensorDescs);

    outtensorDescs.resize(outTensorNum);
    pack.outTensors.resize(outTensorNum);
    operation->InferShape(intensorDescs, outtensorDescs);
    CreateOutTensors(pack.outTensors, outtensorDescs);

    // Initialize the workspace.
    uint64_t workspaceSize = 0;
    void *workSpace = nullptr;
    void *workSpace1 = nullptr;
    void *workSpace2 = nullptr;
    int ret = 0;
    // Graph execution
    // Internal capture
    if (!captureType || std::strcmp(captureType, "in") == 0) {
        std::cout << "start test capture graph inside" << std::endl;
        for (size_t i = 0; i < 10; i++) {
            CreateInTensors(pack.inTensors, intensorDescs, i + 1);
            if (i == 0) {
                // Capture device-side tasks.
                operation->Setup(pack, workspaceSize, context);
                if (workspaceSize != 0 && workSpace == nullptr) {
                    ret = aclrtMalloc(&workSpace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
                    if (ret != 0) {
                        std::cout << "alloc error!";
                        exit(0);
                    }
                }
                std::cout << "workspace:" << workSpace << ", workspaceSize:" << workspaceSize << std::endl;
                context->SetExecuteType(atb::EXECUTE_PRELAUNCH);
                operation->Execute(pack, (uint8_t*)workSpace, workspaceSize, context);
                context->SetExecuteType(atb::EXECUTE_LAUNCH);
                operation->Execute(pack, (uint8_t*)workSpace, workspaceSize, context);
            } else if (i > 3 && i < 6) {
                // Subsequent replay execution. A new workspace can be passed.
                operation->Setup(pack, workspaceSize, context);
                if (workspaceSize != 0 && workSpace1 == nullptr) {
                    ret = aclrtMalloc(&workSpace1, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
                    if (ret != 0) {
                        std::cout << "alloc error!";
                        exit(0);
                    }
                }
                std::cout << "workspace1:" << workSpace1 << ", workspaceSize:" << workspaceSize << std::endl;
                context->SetExecuteType(atb::EXECUTE_PRELAUNCH);
                operation->Execute(pack, (uint8_t*)workSpace1, workspaceSize, context);
                context->SetExecuteType(atb::EXECUTE_LAUNCH);
                operation->Execute(pack, (uint8_t*)workSpace1, workspaceSize, context);
            } else if (i >= 6 && i < 8) {
                // Subsequent replay execution. A new workspace can be passed.
                operation->Setup(pack, workspaceSize, context);
                if (workspaceSize != 0 && workSpace2 == nullptr) {
                    ret = aclrtMalloc(&workSpace2, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
                    if (ret != 0) {
                        std::cout << "alloc error!";
                        exit(0);
                    }
                }
                std::cout << "workSpace2:" << workSpace2 << ", workspaceSize:" << workspaceSize << std::endl;
                context->SetExecuteType(atb::EXECUTE_PRELAUNCH);
                operation->Execute(pack, (uint8_t*)workSpace2, workspaceSize, context);
                context->SetExecuteType(atb::EXECUTE_LAUNCH);
                operation->Execute(pack, (uint8_t*)workSpace2, workspaceSize, context);
            } else {
                // Subsequent replay execution.
                std::cout << "workspace:" << workSpace << ", workspaceSize:" << workspaceSize << std::endl;
                operation->Setup(pack, workspaceSize, context);
                context->SetExecuteType(atb::EXECUTE_PRELAUNCH);
                operation->Execute(pack, (uint8_t*)workSpace, workspaceSize, context);
                context->SetExecuteType(atb::EXECUTE_LAUNCH);
                operation->Execute(pack, (uint8_t*)workSpace, workspaceSize, context);
            }
            // Synchronize the stream and wait until the computation on the device is complete.
            ret = aclrtSynchronizeStream(exeStream);
            if (ret != 0) {
                std::cout << "sync error!" << std::endl;
                exit(0);
            }
            std::cout << "aclrtSynchronizeStream success!" << std::endl;

            // Print the values of the output tensors.
            PrintOutTensorValue(pack.outTensors.at(0));
        }
    } else {
        // External capture
        std::cout << "start test capture graph outside" << std::endl;
        // Create a model running instance.
        aclmdlRI model = nullptr;
        for (size_t i = 0; i < 10; i++) {
            CreateInTensors(pack.inTensors, intensorDescs, i + 1);
            if (i == 0) {
                // Capture device-side tasks.
                // Start the capture.
                aclmdlRICaptureBegin(exeStream, ACL_MODEL_RI_CAPTURE_MODE_RELAXED);
                operation->Setup(pack, workspaceSize, context);
                if (workspaceSize != 0 && workSpace == nullptr) {
                    ret = aclrtMalloc(&workSpace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
                    if (ret != 0) {
                        std::cout << "alloc error!";
                        exit(0);
                    }
                }
                std::cout << "workspace:" << workSpace << ", workspaceSize:" << workspaceSize << std::endl;
                context->SetExecuteType(atb::EXECUTE_PRELAUNCH);
                operation->Execute(pack, (uint8_t*)workSpace, workspaceSize, context);
                context->SetExecuteType(atb::EXECUTE_LAUNCH);
                operation->Execute(pack, (uint8_t*)workSpace, workspaceSize, context);
                // Finish the capture.
                aclmdlRICaptureEnd(exeStream, &model);
                // Execute the model running instance.
                aclmdlRIExecuteAsync(model, exeStream);
            } else if (i > 3 && i < 6) {
                // Subsequent replay execution. A new workspace can be passed.
                // In the Setup phase, only the variantPack is verified.
                operation->Setup(pack, workspaceSize, context);
                if (workspaceSize != 0 && workSpace1 == nullptr) {
                    ret = aclrtMalloc(&workSpace1, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
                    if (ret != 0) {
                        std::cout << "alloc error!";
                        exit(0);
                    }
                }
                std::cout << "workspace1:" << workSpace1 << ", workspaceSize:" << workspaceSize << std::endl;
                // Call PreLaunch to update the parameters of the model running instance.
                context->SetExecuteType(atb::EXECUTE_PRELAUNCH);
                operation->Execute(pack, (uint8_t*)workSpace1, workspaceSize, context);
                // Replay the model running instance.
                aclmdlRIExecuteAsync(model, exeStream);
            } else if (i >= 6 && i < 8) {
                // Subsequent replay execution. A new workspace can be passed.
                operation->Setup(pack, workspaceSize, context);
                if (workspaceSize != 0 && workSpace2 == nullptr) {
                    ret = aclrtMalloc(&workSpace2, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
                    if (ret != 0) {
                        std::cout << "alloc error!";
                        exit(0);
                    }
                }
                std::cout << "workSpace2:" << workSpace2 << ", workspaceSize:" << workspaceSize << std::endl;
                context->SetExecuteType(atb::EXECUTE_PRELAUNCH);
                operation->Execute(pack, (uint8_t*)workSpace2, workspaceSize, context);
                aclmdlRIExecuteAsync(model, exeStream);
            } else {
                // Subsequent replay execution.
                std::cout << "workspace:" << workSpace << ", workspaceSize:" << workspaceSize << std::endl;
                operation->Setup(pack, workspaceSize, context);
                context->SetExecuteType(atb::EXECUTE_PRELAUNCH);
                operation->Execute(pack, (uint8_t*)workSpace, workspaceSize, context);
                aclmdlRIExecuteAsync(model, exeStream);
            }
            // Synchronize the stream and wait until the computation on the device is complete.
            ret = aclrtSynchronizeStream(exeStream);
            if (ret != 0) {
                std::cout << "sync error!" << std::endl;
                exit(0);
            }
            std::cout << "aclrtSynchronizeStream success!" << std::endl;
            // Print the values of the output tensors.
            PrintOutTensorValue(pack.outTensors.at(0));
        }
        // Destroy the model running instance.
        aclmdlRIDestroy(model);
    }

    // Release the resources.
    atb::DestroyOperation(operation);
    atb::DestroyContext(context);
    for (size_t i = 0; i < pack.inTensors.size(); i++) {
        aclrtFree(pack.inTensors.at(i).deviceData);
    }
    for (size_t i = 0; i < pack.outTensors.size(); i++) {
        aclrtFree(pack.outTensors.at(i).deviceData);
    }
    aclrtFree(workSpace);
    aclrtFree(workSpace1);
    aclrtFree(workSpace2);
    aclrtDestroyStream(exeStream);
    aclrtResetDevice(deviceId);
    aclFinalize();
}