算子原型定义
算子原型主要描述了算子的输入输出、属性等信息以及算子在AI处理器上相关实现信息,并关联tiling实现等函数。算子原型通过自定义的算子类来承载,该算子类继承自OpDef类。完成算子的原型定义等操作后,需要调用OP_ADD接口,传入算子类型(自定义算子类的类名),进行算子原型注册。下面是一个简单的Add算子原型定义和注册的例子。
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}); // 根据用户的算子调用方式决定需不需要注册 单算子API调用方式下不需要 this->SetInferShape(ge::InferShape); this->SetInferDataType(ge::InferDataType); this->AICore() .SetTiling(optiling::TilingFunc); // 请替换为实际的昇腾AI处理器型号 this->AICore().AddConfig("ascendxxx"); } }; OP_ADD(AddCustom); } // namespace ops
注册算子类型后,框架会根据算子类型获取算子注册信息,同时在编译和运行时按照一定的规则匹配算子实现文件名称和kernel侧核函数名称。为了保证正确匹配,算子类型、算子实现文件名称和核函数名称需要遵循如下定义规则。通常情况下,开发者只需要保证创建算子工程时原型定义json文件中算子类型op的参数值为大驼峰命名方式即可,工程创建后自动生成的代码即满足该规则。在手动编写算子原型定义和算子实现文件时需要按照如下规则定义。
- 算子类型需要采用大驼峰的命名方式,即采用大写字符区分不同的语义。
- 算子实现文件名称、核函数名称需相同,均为算子类型转换为下划线命名方式后的值。下文描述了通过算子类型转换成算子实现文件名称和核函数名称的过程:
- 首字符的大写字符转换为小写字符。例如:Abc -> abc。
- 大写字符的前一个字符为小写字符或数字,则在大写字符前插一个下划线“_”,并将该字符转换为小写字符。例如:AbcDef -> abc_def。
- 大写字符前一个字符为大写字符且后一个字符是小写字符,则在大写字符前插一个下划线“_”,并将该字符转换为小写字符。例如:AbcAAc -> abc_a_ac。
- 其他大写字符转换为小写字符,小写字符保持不变。
算子原型定义
算子原型定义描述了算子的输入输出,属性等信息。输入输出支持的数据类型、格式的数量需要保持一致,保持一一对应的关系。
如下的代码片段呈现了Add算子输入x的描述信息。
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});
如下的代码片段呈现了ReduceMax算子的属性reduceDim和isKeepDim的描述信息。
this->Attr("reduceDim") .AttrType(REQUIRED) .Int(); this->Attr("isKeepDim") .AttrType(OPTIONAL) .Int(1);
具体参数说明如下:
原型定义 |
注册方式 |
具体描述 |
---|---|---|
Input/Output |
ParamType |
参数类型,Option取值为:OPTIONAL(可选)、REQUIRED(必选)、DYNAMIC(动态输入)。
|
DataType |
||
Format |
||
Attr |
AttrType |
设置算子属性类型,取值为:OPTIONAL(可选)、REQUIRED(必选)。 |
Bool/Float/Int... |
AI处理器上相关实现信息
通过AddConfig注册算子支持的AI处理器型号以及相关的配置信息。AddConfig接口原型如下:soc参数表示AI处理器型号,aicore_config表示其他配置信息。
void AddConfig(const char *soc); void AddConfig(const char *soc, OpAICoreConfig &aicore_config);
通过该接口注册AI处理器型号的样例如下,ascendxxx请替换为实际的AI处理器型号。
this->AICore().AddConfig("ascendxxx");
关联Tiling实现、Shape推导等函数
通过SetInferShape、SetInferDataType、SetTiling接口来关联对应的shape推导函数和Tiling函数,样例如下。
this->SetInferShape(ge::InferShape); this->SetInferDataType(ge::InferDataType); this->AICore() .SetTiling(optiling::TilingFunc);
多硬件平台注册差异化的算子原型
算子类继承基类OpDef,使用Input、Output、Attr等注册算子原型信息,硬件平台支持相同的算子原型的情况下,直接通过AICore().AddConfig添加支持的AI处理器型号即可;不同的硬件形态算子原型定义不同的情况,可以通过新增OpAICoreConfig的方式,针对不同的AI处理器型号注册差异化的算子原型。
差异化的算子原型生效规则如下:
- 对于算子类的输入输出原型信息,OpAICoreConfig未配置的会继承OpDef定义的原型,比如算子类中定义了输出y,OpAICoreConfig中没有定义输出y,OpAICoreConfig会继承y的原型定义;
- 对于算子类和新增OpAICoreConfig中定义的算子原型相同的情况,新增OpAICoreConfig中定义的算子原型信息会覆盖OpDef定义的原型信息,比如算子类中定义了输入x支持DT_FLOAT16数据类型,新增OpAICoreConfig中也定义了输入x,但是支持DT_FLOAT16、DT_BF16数据类型,则以OpAICoreConfig新增定义为准。
如下样例中ascendxxx1、ascendxxx2(AI处理器型号)使用相同的算子原型,算子类通过继承基类OpDef,使用Input、Output、Attr等注册算子原型信息,再通过AICore().AddConfig添加支持的AI处理器型号;对于ascendxxx3支持的算子原型需要定制化处理,新增了DT_BF16的类型,通过新增OpAICoreConfig的方式进行注册,x,y,z的定义会覆盖算子类中对应定义的原型信息。
namespace ops { class MyAdd : public OpDef { public: MyAdd(const char* name) : OpDef(name) { // ascendxxx1 ascendxxx2 AI处理器型号原型定义 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"); // ascendxxx3芯片定义OpAICoreConfig变量,定制化原型 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); }
如下的样例中,只有几个参数原型信息在不同硬件平台不一致,开发者也可以通过OpAICoreConfig定制部分算子原型信息,复用OpDef定义的其他算子原型信息,达到部分原型信息硬件平台定制化的目的。
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); } };