atb::Tensor ToNpuTensor(const Data &data) {
    atb::Tensor tensor;
    tensor.hostData = data.cpuData;
    tensor.dataSize = data.GetBytes();
    tensor.desc.dtype = ACL_FLOAT16;
    AssertInFastLLM(data.dims.size() <= atb::MAX_DIM,
                    "Dims " + std::to_string(data.dims.size()) + " exceed MAX_DIM(8)");
    tensor.desc.shape.dimNum = data.dims.size();
    std::cout << "dims: ";
    for (int i = 0; i < data.dims.size(); ++i) {
        tensor.desc.shape.dims[i] = data.dims[i];
        std::cout << tensor.desc.shape.dims[i] << ", ";
    }
    std::cout << std::endl;
    tensor.desc.format = ACL_FORMAT_ND;
    return tensor;
}

void RunNpuRMSNormOp(Data& input, Data& weight, Data& output) {
      atb::infer::RmsNormParam op;
    op.layerType = atb::infer::RmsNormParam::RmsNormType::RMS_NORM_NORM;
    op.normParam.quantType = atb::infer::QuantType::QUANT_UNDEINFED;
    op.normParam.epsilon = eps;
    atb::VariantPack data;
    atb::SVector<atb::Tensor> &in = data.inTensors;
    input.Reshape({1, 1, 1, 5120});
    auto x_tensor = ToNpuTensor(input);
    in.push_back(x_tensor);  // x
    weight.Reshape({1, weight.Shape()[0]});
    auto gamma_tensor = ToNpuTensor(weight);
    in.push_back(gamma_tensor);  // gamma(weight)

        // 返回时清理相关资源
    struct ScopedCleanUp {
        ~ScopedCleanUp() {
            if (operation) atb::DestroyOperation(operation);
            if (workspace) aclrtFree(workspace);
            if (context) atb::DestroyContext(context);
            if (stream) aclrtDestroyStream(stream);
        }

        atb::Operation *operation = nullptr;
        void *workspace = nullptr;
        atb::Context *context = nullptr;
        void *stream = nullptr;
    } res;
    // 构造算子对应的Operation
    RETURN_IF_ATB_ERROR(atb::CreateOperation(op, &res.operation), "operation create");
    std::cout << res.operation->GetName() << std::endl;
    std::cout << "input tensor num: " << data.inTensors.size() << std::endl;
    std::cout << "op input num: " << res.operation->GetInputNum() << std::endl;
    std::cout << "op output num: " << res.operation->GetOutputNum() << std::endl;
    // 分配Npu内存
    uint64_t workspaceSize = 0;
    RETURN_IF_ATB_ERROR(res.operation->Setup(data, workspaceSize), "operation setup");
    RETURN_IF_ACL_ERROR(aclrtMalloc(&res.workspace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST), "npu memory malloc");
    // 执行算子
    RETURN_IF_ATB_ERROR(atb::CreateContext(&res.context), "context create");
    RETURN_IF_ACL_ERROR(aclrtCreateStream(&res.stream), "stream create");
    RETURN_IF_ATB_ERROR(res.context->SetExecuteStream(res.stream), "stream set");
    RETURN_IF_ATB_ERROR(res.operation->Execute(data, static_cast<uint8_t *>(res.workspace), workspaceSize, res.context),
                        "operation execute");
    output.cpuData = static_cast<uint8_t *>(data.outTensors[0].hostData);
}