Overview
The ATB defines three types of operators: single-operator (Operation), graph operator (GraphOperation), and custom plugin operator (PluginOperation).
- Single-operators (Operations) are a series of basic operators with functions such as matrix multiplication and transposition. For details, see atb/infer_op_params.h and atb/train_op_params.h.
- GraphOperation is a combination of multiple single-operators. It is used to simplify the call steps of multiple operators. The performance and memory of the graph operator are optimized in the ATB. You are advised to use the graph operator preferentially when multiple operators are combined.
You need to design and define the graph structure, that is, the combination and dependency of nodes in the graph structure, including the single-operator, input tensor, and output tensor of the nodes, these tensors are identified as the input tensor, output tensor, and intermediate tensor of the graph. The graph input tensors are all tensors that need to be input from the outside when the graph operator is used. The graph output tensors are all tensors that will not be used for the next operation when the graph operator is used. The graph intermediate tensor is a temporary tensor generated during graph operator operation. The tensor is neither the graph input nor the graph output.
As shown in Figure 1, the graph operator in this example consists of two nodes, which are both Elewise_Add operators. a, b, and c are the three input tensors of the graph operator, output is the output tensor of the graph operator, and a_add_b_output is the intermediate tensor of the graph operator. Inputs of the node 0 are a and b, and an output is a_add_b_output. Inputs of the node 1 are a_add_b_output and c, and an output is output.
- PluginOperation is a mechanism for users to implement specific functions. If some functions cannot be implemented using single-operators or graph operators, you can develop custom plugin operators to implement the corresponding functions.
The three types of operators are inherited from the same parent class Operation. The steps are basically the same except for the operator construction procedure.
