开发者
资源
[object Object]

算子原型主要描述了算子的输入输出、属性等信息以及算子在AI处理器上相关实现信息,并关联等函数。算子原型通过自定义的算子类来承载,该算子类继承自OpDef类。完成算子的原型定义等操作后,需要调用OP_ADD接口,传入算子类型(自定义算子类的类名),进行算子原型注册。下面是一个简单的Add算子原型定义和注册的例子。

[object Object]
[object Object]
[object Object]

算子原型定义描述了算子的输入输出、属性等信息。输入输出支持的datatype、format格式的数量需要一致,并保持一一对应的关系。

如下的代码片段呈现了Add算子输入x的描述信息。

[object Object]

表 1 输入输出参数说明

[object Object][object Object]

[object Object]

从上文的原型定义中可以看出,列出了输入输出所有datatype和format的组合,保持一一对应。使用如下接口,可以达到简化这种代码逻辑的目的。

  • 在指定输入/输出的datatype信息时,如果某个输入/输出的datatype支持和其他所有输入/输出的datatype/format组合使用,其datatype可以通过DataTypeList来表达;在指定输入/输出的format信息时,如果某个输入/输出的format支持和其他所有输入/输出的datatype/format组合使用,其format可以通过FormatList来表达。示例如下,以下两种代码表达含义相同。

    [object Object]
  • 通过Follow接口指定当前输入/输出的datatype/format/shape信息与之前定义过的某个输入一致。示例如下:输出“y1”Follow输入“x1”场景,此时“y1”的datatype、format以及shape都将会和“x1”保持一致。使用Follow接口指定shape一致时通常比函数逻辑更加简单,能用Follow表达的逻辑,建议使用Follow接口,则无需再编写注册InferShape函数。

    [object Object]

原型定义中还包括算子属性信息,如下的代码片段呈现了ReduceMax算子的属性reduceDim和isKeepDim的描述信息。

[object Object]

具体参数说明如下:

表 2 属性参数说明

[object Object][object Object]

[object Object][object Object]

通过AddConfig注册算子支持的AI处理器型号以及相关的配置信息。AddConfig接口原型如下:soc参数表示AI处理器型号,aicore_config表示其他配置信息。

[object Object]

通过该接口注册AI处理器型号的样例如下,ascendxxx填写规则请参考算子工程目录下编译配置项文件CMakePresets.json中的ASCEND_COMPUTE_UNIT字段,该字段取值在时自动生成。

[object Object]

其他AI Core配置信息的配置方式请参考OpAICoreConfig。

[object Object]

通过SetInferShape、SetInferDataType、SetTiling接口来注册对应的Tiling实现和Shape推导等函数,样例如下。注册的Tiling实现等函数由框架侧进行调用,并在调用时传入对应的Context上下文,供开发者使用。Tiling函数的实现方法请参考,入图相关的Shape推导等函数实现请参考

[object Object]
[object Object]

算子类继承基类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的定义会覆盖算子类中对应定义的原型信息。

[object Object]

如下的样例中,只有几个参数原型信息在不同硬件平台不一致,开发者也可以通过OpAICoreConfig定制部分算子原型信息,复用OpDef定义的其他算子原型信息,达到部分原型信息硬件平台定制化的目的。

[object Object]