Implementation

Registration Code Implementation for the Operator IR Header File

1. Define the macro.

   Use the following statement to define the operator IR registration macro. The macro name is fixed to GE_OP_OPERATORTYPE_H. OPERATORTYPE is the uppercase of OpType in the REG_OP(OpType) statement.

   #ifndef GE_OP_OPERATORTYPE_H       // conditional build
   #define GE_OP_OPERATORTYPE_H       // macro definition

2. Include header files.

   In the header of the operator IR implementation file, run the #include command to include the header files registered by the operator in the operator IR implementation file.

   #include "graph/operator_reg.h"

   Find the operator_reg.h file in the /include/graph/ directory in the ATC installation path. Include this file to use functions, macros, and structures related to the operator type registration.

3. Register the prototype.

   GE provides the REG_OP macro and the INPUT, OUTPUT, and ATTR APIs (cascaded using dots) to register the input, output, and attribute information of the operator. Operator registration is ended with the OP_END_FACTORY_REG API.

   The sequence of the input/output arguments must be the same as that defined in the operator implementation. The sequence of ATTR arguments is variable.

   The registration code is as follows:

   namespace ge{
       REG_OP(OpType) // operator type
       .INPUT(x1, TensorType({ DT_FLOAT, DT_INT32 }))
       .INPUT(x2, TensorType({ DT_FLOAT, DT_INT32 })) 
       // .DYNAMIC_INPUT(x, TensorType{DT_FLOAT, DT_INT32})
       // .OPTIONAL_INPUT(b, TensorType{DT_FLOAT})        
       .OUTPUT(y, TensorType({ DT_FLOAT, DT_INT32 })) 
       // .DYNAMIC_OUTPUT(y, TensorType{DT_FLOAT, DT_INT32})
       .ATTR(x, Type, DefaultValue)
       // .REQUIRED_ATTR(x, Type)
       // .GRAPH(z1)
       // .DYNAMIC_GRAPH(z2)
       .OP_END_FACTORY_REG(OpType)
   }

   • REG_OP(OpType)

     OpType: operator type registered with the custom operator library of the Ascend AI Processor. The value must be the same as the operator type name defined by REGISTER_CUSTOM_OP("OpType") in Operator Plug-in Implementation.

   • INPUT(x1, TensorType({ DT_FLOAT,DT_UINT8,... }))
     Registers an operator input.

     • x: user-defined macro parameter that indicates the input name of the operator.
     • TensorType({ DT_FLOAT,DT_UINT8,... }): The data in { } indicates the list of supported input data types.

     If the operator has multiple inputs, describe each input with one statement INPUT(x, TensorType({ DT_FLOAT,DT_UINT8,...})).

   • DYNAMIC_INPUT(x, TensorType{DT_FLOAT, DT_INT32, ...})
     Registers input information in the dynamic multi-input scenario.

     • x: macro parameter, which indicates the input name of the operator. During graph execution, x0, x1, x2, and more are automatically generated based on the number of inputs. The inputs are numbered starting at 0 in ascending order.
     • TensorType({ DT_FLOAT,DT_UINT8,... }): The data in { } indicates the list of supported input data types.
   • OPTIONAL_INPUT(x, TensorType{DT_FLOAT, ...})
     Registers an optional operator input.

     • x: a macro parameter for the input name
     • TensorType{DT_FLOAT, ...}: a list of data types supported by the input.
   • OUTPUT(y, TensorType({ DT_FLOAT,DT_UINT8,... }))
     Registers an operator output.

     • y: user-defined macro parameter that indicates the output name of the operator.
     • TensorType({ DT_FLOAT,DT_UINT8,... }): The data in { } indicates the list of supported output data types.

     If the operator has multiple outputs, register each output with one statement OUTPUT(x, TensorType({ DT_FLOAT,DT_UINT8, ...})).

   • DYNAMIC_OUTPUT(y, TensorType{DT_FLOAT, DT_INT32})
     Registers output information in the dynamic multi-output scenario.

     • y: macro parameter, which indicates the output name of the operator. During graph execution, y0, y1, y2, and more are automatically generated based on the number of outputs. The outputs are numbered starting at 0 in ascending order.
     • TensorType({ DT_FLOAT,DT_UINT8,... }): The data in { } indicates the list of supported output data types.
   • ATTR(x, Type, DefaultValue)

     Registers an operator attribute, including attribute name, type, and default value. If you do not set the attribute value of an operator object, the default value is used. For details about the values of Type, see Prototype Definition API (REG_OP).

     For example, ATTR(mode, Int, 1) registers a mode attribute of type int64_t with default value 1.

     If the operator has multiple attributes, register each attribute separately in the format of ATTR(x, Type, DefaultValue) or REQUIRED_ATTR (x, Type).

   • REQUIRED_ATTR (x, Type)

     Registers an operator attribute, including the attribute name and type, without a default value, which means that the attribute must be specified with a value. For details about the values of Type, see Prototype Definition API (REG_OP).

     If the operator has multiple attributes, register each attribute separately in the format of ATTR(x, Type, DefaultValue) or REQUIRED_ATTR (x, Type).

   • GRAPH(z1)

     Registers the subgraph information contained in the operator. Replace z1 with the subgraph name, which is applicable to control operators such as branch operators and loop operators.

     After the registration is complete, the APIs (as described in GRAPH) for obtaining subgraph names and descriptions, and setting subgraph descriptions are automatically generated, which can be used to build IR models. For a single operator, the name of the registered operator subgraph must be unique.

   • DYNAMIC_GRAPH(z2)

     Registers the subgraph information in a dynamic operator. Replace z2 with the subgraph name, which is applicable to control operators such as branch operators and loop operators.

     After the registration is complete, the APIs (as described in DYNAMIC_GRAPH) for creating dynamic subgraphs and setting subgraph descriptions are automatically generated, which can be used to build IR models. For a single operator, the name of the registered operator subgraph must be unique.

   • OP_END_FACTORY_REG(OpType): ends operator registration. OpType must be consistent with OpType in REG_OP(OpType).

4. End conditional build.

   #endif

Implementation of the .cpp Registration Code Defined by the Operator IR

The .cpp file implemented by the IR offers the following functions:

• Verifies operator parameters, implementing program robustness and improving locating efficiency.
• Infers the output tensor description of the operator based on the input tensor description, operator logic, and operator attributes. The output tensor description includes the tensor shape, data type, and data layout format. In this way, all tensors can be statically allocated with memory during the preparation for graph construction, thereby avoiding overhead caused by dynamic memory allocation.

When implementing the Verify and InferShape methods in /op_proto/operator_name.cpp, you do not need to declare the methods.

1. Include header files.

   #include "operator name.h"     
   #include <vector>
   #include <string>

   Table 8-1 Description of header files

   Header File	Category	Function
   operator name.h	IR header file implemented in Registration Code Implementation for the Operator IR Header File	The object op of the class Operator registered in this file or the subclass op derived can be called once this head file is included.
   string	C++ standard library	Class string can be used to construct objects APIs of class string can be called once this header file is included.
   vector	C++ standard library	Vector templates can be used and APIs of class vector can be called once this header file is included.

2. Implement the InferShape method.

   The following APIs can be used to define InferShape in the operator IR:

   • IMPLEMT_COMMON_INFERFUNC(func_name): automatically generates the object op of the Operator class. You can directly call an API of class Class Operator to implement InferShape. If the InferShape method is universal and can be called by the prototypes of multiple operators, you can use this API.

     func_name: user-defined function name.

     The following is an example of assigning the input description to the output description:

     IMPLEMT_COMMON_INFERFUNC(SoftmaxInferShape)
     {
       TensorDesc tensordesc_output = op.GetOutputDesc("y");
     
       tensordesc_output.SetShape(op.GetInputDesc("x").GetShape());
       tensordesc_output.SetDataType(op.GetInputDesc("x").GetDataType());
       tensordesc_output.SetFormat(op.GetInputDesc("x").GetFormat());
     
       (void)op.UpdateOutputDesc("y", tensordesc_output);
       return GRAPH_SUCCESS;
     }

     The output description is computed based on the operator logic. The following is an example.

     IMPLEMT_COMMON_INFERFUNC(NotEqualInferShape) {
       Shape x_shape = op.GetInputDesc("x1").GetShape();
       Shape y_shape = op.GetInputDesc("x2").GetShape();
       TensorDesc td = op.GetOutputDesc("y");
       std::vector<int64_t> dims_x = x_shape.GetDims();
       std::vector<int64_t> dims_y = y_shape.GetDims();
       if (dims_x.size() < dims_y.size()) {
         std::vector<int64_t> dims_tmp = dims_x;
         dims_x = dims_y;
         dims_y = dims_tmp;
       }
       if (dims_x.size() != dims_y.size()) {
         int dec = dims_x.size() - dims_y.size();
         for (int i = 0; i < dec; i++) {
           dims_y.insert(dims_y.begin(), (int64_t)1);
         }
       }
       std::vector<int64_t> dim_vec;
       for (size_t i = 0; i < dims_x.size(); i++) {
         if ((dims_x[i] != dims_y[i]) && (dims_x[i] != 1) && (dims_y[i] != 1)) {
           printf ("The %s op dimensions does not match the broadcast rule(%lu %lu).", op.GetName().c_str(), dims_x[i], dims_y[i]);
         }
         int64_t dims = dims_x[i] > dims_y[i] ? dims_x[i] : dims_y[i];
         dim_vec.push_back(dims);
       }
       td.SetShape(ge::Shape(dim_vec));
       td.SetDataType(DT_BOOL);
       (void)op.UpdateOutputDesc("y", td);
       return GRAPH_SUCCESS;
     }

   • IMPLEMT_INFERFUNC(OpType, func_name): The input OpType is a subclass derived from class Operator. An op object of this subclass is automatically generated. You can use the member functions of the subclass to obtain the input and output description methods. IMPLEMT_INFERFUNC describes the available functions.

     OpType indicates the custom operator type, which must be consistent with OpType in REG_OP(OpType).

     func_name: user-defined function name.

     The member functions of the subclass op derived from OpType are as follows:

     • op.set_input_x(Operator &v, const string &srcName): sets the output srcName of operator v on the network to the input x of the current operator.
     • op.get_input_desc_x(): obtains the description of the input x of the operator and returns an object of the TensorDesc type.

       op.update_input_desc_x(const TensorDesc& tensorDesc): updates the description of input x, including the shape, data type, and format.

     • op.get_output_desc_y(): obtains the description of the output y of the operator and returns an object of the TensorDesc type.
     • op.update_output_desc_y(const TensorDesc& tensorDesc): updates the description of output y, including the shape, data type, and format.
     • op.get_attr_attr1(): obtains the value of the operator attribute attr1.

     The implementation example is as follows:

     IMPLEMT_INFERFUNC(ReduceSum, ReduceSumInfer) {
       Tensor data;
       op.GetInputConstData("axis", data);
       DataType dtype = op.GetInputDesc("axis").GetDataType();
       std::vector<int64_t> const_vec {};
       size_t size = data.GetSize();
       if (size != 0) {
         CalcReduceSum(data, dtype, const_vec);
       }
     
       auto tensordesc = op.get_input_desc_x();
       auto shape = tensordesc.GetShape();
       std::vector<int64_t> shapeVector = shape.GetDims();
       int64_t dimNum = shape.GetDimNum();
       std::vector<int64_t> axis(const_vec);
       bool keep_dims = op.get_attr_keep_dims();
     
       if (axis.size() == 0) {
         for (size_t i = 0; i < shapeVector.size(); ++i) {
           axis.push_back(i);
         }
       }
     
       for (size_t i = 0; i < axis.size(); ++i) {
         if (axis[i] < 0) {
           axis[i] = dimNum + axis[i];
         }
       }
     
       std::vector<int64_t> oShapeVector;
       std::vector<int64_t>::iterator tmp;
       for (int64_t item = 0; item < dimNum; ++item) {
         tmp = std::find(axis.begin(), axis.end(), item);
         if (tmp != axis.end()) {
           // item in axis
           if (keep_dims == true) {
               // If keepDims is true, current dimesion set to 1
               oShapeVector.push_back(1);
           }
         } else {
           // item is not in ConstValueAxis
           oShapeVector.push_back(shapeVector[item]);
         }
       }
     
       Shape oShape(oShapeVector);
       tensordesc.SetShape(oShape);
       op.update_output_desc_y(tensordesc);
       return GRAPH_SUCCESS;
     }
     INFER_FUNC_REG(ReduceSum, ReduceSumInfer);

3. Implement the Verify function.

   The following APIs are used to implement the Verify function of an operator:

   IMPLEMT_VERIFIER (OpType, func_name)

   The input OpType is a subclass derived from the class Operator. An op object of this subclass is automatically generated. You can use the member functions of the subclass to obtain the operator attributes. For details about the member functions of the op object, see 2.

   • OpType: type of the custom operator
   • func_name: user-defined name of the Verify function.

   The verify function is used to verify the internal association relationship of operators. For example, for a multi-input operator, the data types of multiple tensors must be the same. In this case, the data types of multiple inputs need to be verified. In other cases, the data type does not need to be verified.

   The implementation example is as follows:

   IMPLEMT_VERIFIER(Pow, PowVerify) {
     DataType input_type_x = op.GetInputDesc("x").GetDataType();
     DataType input_type_y = op.GetInputDesc("y").GetDataType();
     if (input_type_x != input_type_y) {
       return GRAPH_FAILED;
     }
     return GRAPH_SUCCESS;
   }

4. Register the Infershape and Verify functions.

   Call the Infershape registration macro and Verify registration macro to register the Infershape and Verify functions, as shown in the following:

   COMMON_INFER_FUNC_REG(OpType, func_name);  
   // INFER_FUNC_REG(OpType, func_name);  
   VERIFY_FUNC_REG (OpType, func_name);

   COMMON_INFER_FUNC_REG(OpType, func_name) or INFER_FUNC_REG(OpType, func_name) are mutually exclusive. The func_name is consistent with that in 2 and 3.