model.h

#ifndef MODEL_H
#define MODEL_H

#include <map>
#include <acl/acl.h>
#include <atb/atb_infer.h>
#include <atb/types.h>
#include <atb/utils.h>
#include "atb/infer_op_params.h"
#include "utils/log.h"

enum class TensorType
{
    INTERNAL_TENSOR = 0,
    NOT_INTERNAL_TENSOR,
};

// Graph node. Each node indicates an operation or graph operation.
struct Node
{
    // Operation or graph operation corresponding to a node.
    atb::Operation *operation_ = nullptr;

    // Input tensors of the node.
    atb::SVector<atb::Tensor *> inTensors_{};

    // Output tensors of the node.
    atb::SVector<atb::Tensor *> outTensors_{};

    // The output of the node is of the intermediate tensor type.
    atb::SVector<TensorType> outTensorTypes_{};

    atb::VariantPack variantPack_{};

    uint64_t workspaceSize_ = 0;
    int workspaceBlockId_ = -1;
    void *workspace_ = nullptr;
};

// All nodes form a complete graph.
class Model
{
public:
    // Describe the input of the model.
    enum InTensorId : int
    { // Define tensor IDs.
        IN_TENSOR_A = 0,
        IN_TENSOR_B,
        IN_TENSOR_C,
        IN_TENSOR_D,
        Mode_INPUT_SIZE,
    };

    enum OutTensorId : int
    {
        GLUE_OUT = 0,
        Mode_OUTPUT_SIZE,
    };

    explicit Model(std::string &&modelName = "") : modelName_(std::move(modelName))
    {
        LOG_INFO("Create model: " + modelName_);
    }

    // Initialize the model and set the device ID of the model.
    void InitResource(uint32_t deviceId);

    // Create a model graph.
    void CreateModelGraph();

    // Create the input tensors of the model.
    void CreateModelInput();

    // Create the output tensors of the model.
    void CreateModelOutput();

    // Execute the model.
    void Execute();

    // Synchronize streams.
    void WaitFinish();

    // Release the resources.
    void FreeResource();

    // Input tensors of the model.
    atb::SVector<atb::Tensor> modelInTensors_;

    // Output tensors of the model.
    atb::SVector<atb::Tensor> modelOutTensors_;

private:
    // Create the operation of the graph operator.
    void CreateGraphOpLayer(size_t nodeId);

    // Create the operation of the aclnn operator.
    void CreateAclnnOpLayer(size_t nodeId);

    // Construct the VariantPack of the node corresponding to nodeId.
    void BuildNodeVariantPack(int nodeId);

    // Deliver the operation corresponding to nodeId.
    atb::Status ExecuteNode(int nodeId);

    // Workspace creation function.
    void CreateWorkspaceBuffer(int nodeId, int workspaceSizeNeeded);

    // Shape inference function of the model graph.
    atb::Status InferShape(
        const atb::SVector<atb::TensorDesc> &inTensorDescs, atb::SVector<atb::TensorDesc> &outTensorDescs);

    std::string modelName_;
    uint32_t deviceId_ = 1;
    atb::Context *modeContext_ = nullptr;
    aclrtStream modelStream_ = nullptr;
    std::vector<Node> nodes_;

    // Intermediate tensors of the model. Layers are connected by internalTensors. Pay attention to the sequence.
    std::vector<atb::Tensor> internalTensors_;
};

#endif