Operator Prototype Definition

The operator prototype describes the input, output, and attributes of the operator as well as the operator implementation on the AI processor, and associates with functions such as tiling implementation. The operator prototype is carried by a custom operator class that is inherited from the OpDef class. After defining the operator prototype, call the OP_ADD API to pass in the operator type (class name of the custom operator class) and register the operator prototype. The following is an example of defining and registering the Add operator prototype.

 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
namespace ops {
class AddCustom : public OpDef {
public:
    AddCustom(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("z")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        // Determine whether registration is required based on the operator calling mode. Registration is not required in single-operator API calling mode.
        this->SetInferShape(ge::InferShape);
        this->SetInferDataType(ge::InferDataType);  
        this->AICore()
            .SetTiling(optiling::TilingFunc);
        // Replace it with the actual Ascend AI Processor model.
        this->AICore().AddConfig("ascendxxx");
    }
};
OP_ADD(AddCustom);
} // namespace ops

After the operator type is registered, the framework obtains the operator registration information based on the operator type and matches the operator implementation file name and kernel function name based on certain rules during compilation and running. To ensure correct matching, the operator type, operator implementation file name, and kernel function name must comply with the following rules. Generally, you only need to ensure that the value of the operator type in the JSON prototype definition file is in upper camel case. The code automatically generated after the project is created meets this rule. When manually writing the operator prototype definition and operator implementation file, comply with the following rules:

  • Name the operator type in upper camel case and separate words with a single capitalized letter.
  • The operator implementation file name and kernel function name must be the same. They are the values after the operator type is converted using underscores (_). The following describes the process of converting the operator implementation file name and kernel function name through the operator type.
    • Replace the first uppercase letter with a lowercase letter. Example: Abc -> abc
    • If a lowercase character or digit is used before an uppercase character, an underscore (_) is inserted before the uppercase character, and the uppercase character is converted to a lowercase character. Example: AbcDef -> abc_def
    • If the character before an uppercase character is an uppercase character and the character after the uppercase character is a lowercase character, underscores (_) are inserted before the uppercase characters and the uppercase characters are converted to lowercase characters. Example: AbcAAc -> abc_a_ac
    • Other uppercase characters are converted to lowercase characters, and lowercase characters remain unchanged.

Operator Prototype Definition

The operator prototype definition describes the input, output, and attributes of the operator. The number of data types and formats supported by the input and output must be the same.

The following code snippet shows the description of input x of the Add operator.

1
2
3
4
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
Table 1 Input and output parameter description

Prototype Definition

Parameter

Description

Input/Output

ParamType

Parameter type. The options are OPTIONAL, REQUIRED, and DYNAMIC.

  • Similar to the Add sample, the input and output are required.
  • The number of inputs or outputs of some operators is dynamic. For example, AddN accumulates N input tensors and outputs one tensor. SplitV splits a tensor into N tensors on a certain axis for output.
  • The inputs/outputs of some operators are optional. For example, the BatchNorm operator does not have the mean value and variance input during training, but has the mean value and variance input during inference.

DataType

Data type supported by the operator input and output. For details about the value of datatype, see DataType.

Format

Format supported by the operator input and output. For details about the value of format, see Format.

According to the preceding prototype definition, all combinations of data types and formats are listed and correspond to each other. The following APIs can be used to simplify the code logic.

  • If an input/output datatype can be used together with other input/output data types/formats, the datatype can be expressed using DataTypeList. If an input/output format can be used together with all other input/output data types/formats, the format can be expressed using FormatList. The following is an example. The meanings of the two types of code are the same.
     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
    40
    // List all corresponding combinations.
    class XxxCustom : public OpDef {
    public:
        XxxCustom(const char* name) : OpDef(name)
        {
            this->Input("x")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
            this->Input("y")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
            this->Output("z")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
            ...
        }
    };
    // Use DataTypeList and FormatList to express the data type and format.
    class XxxCustom : public OpDef {
    public:
        XxxCustom(const char* name) : OpDef(name)
        {
            this->Input("x")
                .ParamType(REQUIRED)
                .DataTypeList({ge::DT_FLOAT16})
                .FormatList({ge::FORMAT_ND});
            this->Input("y")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
            this->Output("z")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
            ...
        }
    };
    
  • Call the Follow API to specify that the datatype, format, and shape of the current input or output are the same as those of a previously defined input. For example, if the output y1 follows the input x1, the datatype, format, and shape of y1 are the same as those of x1. Using the Follow API to specify shape consistency is simpler than using the InferShape function logic. If the logic can be expressed by the Follow API, you are advised to use the Follow API. In this way, you do not need to compile and register the InferShape function.
     1
     2
     3
     4
     5
     6
     7
     8
     9
    10
    11
    12
    this->Input("x1")
        .ParamType(REQUIRED)
        .DataType({ge::DT_FLOAT, ge::DT_FLOAT})
        .Format({ge::FORMAT_ND, ge::FORMAT_ND});
    this->Input("x2")
        .ParamType(REQUIRED)
        .DataType({ge::DT_FLOAT, ge::DT_FLOAT})
        .Format({ge::FORMAT_ND, ge::FORMAT_ND});
    this->Output("y1")
        .ParamType(REQUIRED)
        .Follow("x1")
        .OutputShapeDependOnCompute();
    

The prototype definition also contains operator attributes. The following code snippet shows the description of the reduceDim and isKeepDim attributes of the ReduceMax operator.

1
2
3
4
5
6
        this->Attr("reduceDim")
            .AttrType(REQUIRED)
            .Int();
        this->Attr("isKeepDim")
            .AttrType(OPTIONAL)
            .Int(1);

The following table describes the parameters.

Table 2 Attribute parameters

Prototype Definition

Registration Method

Description

Attr

AttrType

Sets the operator attribute type. The value can be OPTIONAL or REQUIRED.

Bool/Float/Int...

Sets the operator attribute data type to Bool, Float, or Int. For details, see OpAttrDef.

Implementation on the AI Processor

Register the AI processor models supported by the operator and related configuration by calling AddConfig. The prototype of the AddConfig API is as follows: The soc parameter indicates the AI processor model, and the aicore_config parameter indicates other configuration.

1
2
void AddConfig(const char *soc);
void AddConfig(const char *soc, OpAICoreConfig &aicore_config);

The following is an example of registering an AI processor model using this API. Replace ascendxxx with the actual AI processor model.

1
        this->AICore().AddConfig("ascendxxx");

For details about how to configure other AI Core configuration, see OpAICoreConfig.

Associating Functions Like SetTiling and InferShape

Call SetInferShape, SetInferDataType, and SetTiling to associate the shape inference function with the tiling function. An example is as follows:

1
2
3
4
        this->SetInferShape(ge::InferShape);
        this->SetInferDataType(ge::InferDataType);
        this->AICore()
            .SetTiling(optiling::TilingFunc);

Registering Differentiated Operator Prototypes on Multiple Hardware Platforms

The operator class inherits the base class OpDef and uses Input, Output, and Attr to register the operator prototype information. If the hardware platform supports the same operator prototype, use AICore().AddConfig to add the supported AI processor model. Different hardware forms have different operator prototype definitions. You can add OpAICoreConfig to register differentiated operator prototypes for different AI processor models.

Rules for differentiated operator prototypes to take effect:

  • For the input and output prototype information of the operator class, if OpAICoreConfig is not configured, the prototype defined by OpDef is inherited. For example, if output y is defined in the operator class but not in OpAICoreConfig, OpAICoreConfig inherits the prototype definition of y.
  • If the operator class is the same as the operator prototype defined in OpAICoreConfig, the operator prototype information defined in OpAICoreConfig overwrites the prototype information defined in OpDef. For example, the operator class defines that input x supports the DT_FLOAT16 data type, the input x is also defined in the new OpAICoreConfig. However, the DT_FLOAT16 and DT_BF16 data types are supported. The new definition in OpAICoreConfig prevails.

In the following example, ascendxxx1 and ascendxxx2 (AI processor models) use the same operator prototype. The operator class inherits the base class OpDef and uses Input, Output, and Attr to register the operator prototype information, and call AICore().AddConfig to add supported AI processor models. The operator prototype supported by ascendxxx3 needs to be customized. The DT_BF16 type is added and registered by adding OpAICoreConfig. The definitions of x, y, and z overwrite the prototype information defined in the operator class.

 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
40
41
42
namespace ops {
class MyAdd : public OpDef {
public:
    MyAdd(const char* name) : OpDef(name)
    {
        // ascendxxx1 ascendxxx2 AI processor model prototype definition
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND});
        this->Input("y")
            .ParamType(OPTIONAL)
            .DataType({ge::DT_INT64})
            .ValueDepend(REQUIRED)
            .Format({ge::FORMAT_ND});
        this->Output("z")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND});
        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascendxxx1");
        this->AICore().AddConfig("ascendxxx2");
        // The OpAICoreConfig variable is defined for the ascendxxx3 chip to customize the prototype.
        OpAICoreConfig config;
        config.Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND});
        config.Input("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND});
        config.Output("z")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND});
        this->AICore().AddConfig("ascendxxx3", config);
    }
};
OP_ADD(MyAdd);
}

In the following example, only several parameter prototype information is inconsistent on different hardware platforms. You can use OpAICoreConfig to customize some operator prototype information and reuse other operator prototype information defined by OpDef to customize the hardware platform of some prototype information.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
class AddCustom : public OpDef {
public:
    AddCustom(const char* name) : OpDef(name)
    {
        this->Input("x").DataType({ ge::DT_FLOAT16 }).ParamType(OPTIONAL);
        this->Output("y").DataType({ ge::DT_FLOAT16 });
        OpAICoreConfig aicConfig1;
        OpAICoreConfig aicConfig2;
        aicConfig1.Input("x")
            .ParamType(OPTIONAL)
            .DataType({ ge::DT_FLOAT })
            .Format({ ge::FORMAT_ND });
        aicConfig2.Input("x")
            .ParamType(REQUIRED)
            .DataType({ ge::DT_INT32 })
            .Format({ ge::FORMAT_ND });
        this->AICore().AddConfig("ascendxxx1", aicConfig1);
        this->AICore().AddConfig("ascendxxx2", aicConfig2);
    }
};