目录

Operator Prototype Definition and Implementation

The operator prototype describes the input, output, and attributes of an operator as well as the implementation information of an operator on the AI processor, and associates with functions such as Tiling Implementation on the Host. The operator prototype is carried by a custom operator class that is inherited from the OpDef. After defining the operator prototype, call the Prototype Registration (OP_ADD) API to transfer 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.

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 Kirin AI processor model.
        this->AICore().AddConfig("kirin9020"); // Corresponding Kirin AI processor
    }
};
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, attributes, and other information 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.

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 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 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.
Input/Output DataType Data type supported by the operator input and output. For details about the value of datatype, see DataType.
Input/Output 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.

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. An example is as follows: If the output y1 follows the input x1, the data type, 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 Shape Inference 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.

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.

this->Attr("reduceDim")
    .AttrType(REQUIRED)
    .Int();
this->Attr("isKeepDim")
    .AttrType(OPTIONAL)
    .Int(1);

The specific parameter descriptions are as follows.

Table 2 Attribute parameters

Prototype Registration Mode Description
Attr AttrType Sets the operator attribute type. The value can be OPTIONAL or REQUIRED.
Attr 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 aicore_config indicates other configuration information.

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.

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

Associating Functions Like Tiling Implementation and Shape Inference

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

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:

In the following example, kirinxxx1 and kirinxxx2 (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.

namespace ops {
class MyAdd : public OpDef {
public:
    MyAdd(const char* name) : OpDef(name)
    {
        // ascendxxx1 and 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("kirin9020", config); // Corresponding Kirin AI Processor
    }
};
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.

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("kirinxxxx1", aicConfig1);
        this->AICore().AddConfig("kirinxxxx2", aicConfig2);
    }
};