aclnn_operation_base.h

#ifndef ACLNN_OPERATION_BASE_H
#define ACLNN_OPERATION_BASE_H

#include <acl/acl.h>
#include <aclnn/acl_meta.h>
#include <atb/atb_infer.h>
#include <atb/types.h>
#include <atb/utils.h>
#include "atb/infer_op_params.h"

// Wrapper around atb::tensor.
struct AclnnTensor
{
public:
    atb::Tensor atbTensor; 
    aclTensor *tensor = nullptr;
    int tensorIdx = -1; // Index of the aclTensor in the aclExecutor.
    bool needUpdateTensorDataPtr = false;
    atb::SVector<int64_t> strides = {};
};

// Keep the unified API call with the atb operator.
class AclnnBaseOperation : public atb::Operation
{
public:
    explicit AclnnBaseOperation(const std::string &opName);
    ~AclnnBaseOperation() override;
    std::string GetName() const override;

    // Simulate the atb API to obtain the workspace size.
    atb::Status Setup(const atb::VariantPack &variantPack, uint64_t &workspaceSize, atb::Context *context) override;

    // Simulate the atb API to execute the operator.
    atb::Status Execute(const atb::VariantPack &variantPack, uint8_t *workspace, uint64_t workspaceSize,
                        atb::Context *context) override;

    // Create an input aclnnTensor.
    virtual atb::Status CreateAclnnVariantPack(const atb::VariantPack &variantPack) = 0;

    // Calculate the workspace size.
    virtual atb::Status SetAclnnWorkspaceExecutor() = 0;

    // Execute aclnn op.
    virtual atb::Status ExecuteAclnnOp(uint8_t *workspace, aclrtStream &stream) = 0;

    // Update the addresses of the aclnn input and output tensors.
    atb::Status UpdateAclnnVariantPack(const atb::VariantPack &variantPack);

    std::string opName_;
    aclOpExecutor *aclExecutor_ = nullptr;
    atb::SVector<std::shared_ptr<AclnnTensor>> aclInTensors_;
    atb::SVector<std::shared_ptr<AclnnTensor>> aclOutTensors_;
    uint64_t workspaceSize_;
    int workspaceBlockId_ = -1;
};

#endif