utils.h

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
#ifndef UTILS_H
#define UTILS_H

#include <string>
#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"

#define CHECK_RET(cond, str) \
    do                       \
    {                        \
        if (cond)            \
        {                    \
            LOG_ERROR(str);  \
            exit(0);         \
        }                    \
    } while (0)

// Set the attributes of each input tensor.
void CreateInTensorDescs(atb::SVector<atb::TensorDesc> &intensorDescs);

// Set each input tensor and allocate memory space for each input tensor. Here, you need to manually set the input tensors. For implementation, you can use torchTensor conversion or other simple data structure conversion methods.
void CreateInTensors(atb::SVector<atb::Tensor> &inTensors, atb::SVector<atb::TensorDesc> &intensorDescs);

// Set each output tensor and allocate memory space for each output tensor. The setting is the same as that of an input tensor.
void CreateOutTensors(atb::SVector<atb::Tensor> &outTensors, atb::SVector<atb::TensorDesc> &outtensorDescs);

void CreateTensorFromDesc(atb::Tensor &tensor, atb::TensorDesc &tensorDescs);

// Print the output.
void PrintOutTensorValue(atb::Tensor &outTensor);

// Create a graph operator.
atb::Status CreateGraphOperation(atb::Operation **operation);

#endif