课程目标
通过本节课的学习,能够掌握TBE DSL算子的实现流程,几个算子开发交付件的作用及代码实现方法,包括:
算子原型定义
算子原型定义规定了在异腾A处理器上可运行算子的约束,主要体现算子的数学含义,包含定义算子输入、输出和属性信息,基本参数的校验和shape的推导,原型定义的信息会被注册到GE的算子原型库中。网络模型生成时,GE会调用算子原型库的校验接口进行基本参数的校验,校验通过后,会根据原型库中的推导函数推导每个节点的输出shape与dtype,进行输出tensor的静态内存的分配
总结起来,原型主要做三件事:
(1)校验算子的输入张量元数据是否合法
(2)根据算子的输入张量元数据(输入的形状、数据类型、排布格式)推导输出张量的对应信息
(3)把算子注册进GE(图引擎)
算子原型在GE注册的流程
形状推导流程
形状推导(下文用InferShape表示)是算子原型中相对复杂的一个步骤,介绍InferShape流程前先了解如下背景知识。图(Graph)是承载编译优化阶段的基本结构,图由节点OP(Operator)连接而成,每个节点间的边代表执行时要传递的数据,即张量(Tensor),每个Tensor都包含三个主要属性:dtype(数据类型)、shape(形状)、format(数据排布格式).
图中每个OP都可能会有多个输入与多个输出,输入与输出的个数分别由输入TensorDesc与输出TensorDesc的个数来表示.
图中的两个OP通过某条边直连时,前导OP的输出TensorDesc与后继算子的输入TensorDesc应该完全一致。另外每个OP在获取了输入TensorDesc后,通常可以在编译阶段就推导出所有输出TensorDesc(例外场景:有些OP需要根据输入的实际值来推导输出TensorDesc,这类OP如果依赖的输入不是常量节点就无法推导出确切的输出TensorDesc)。
由上述背景介绍可知,只要全图所有首节点的TensorDesc确定了,就可以逐个向下传播,再由算子自身实现的Shape推导能力,就可以将全图所有OP的输入输出TensorDesc推导出来,这就是InferShape的推导流程。
InferShape有以下注意点:可进行InferShape的前提是图中所有首节点的TensorDesc都可确定,因此对于输入数据是动态shape的网络,静态编译的InferShape是无法生效的。
GE的InferShape流程负责推导TensorDesc中的dtype与shape,推导结束后,全图的dtype与shape的规格就完全连续了,如果生成网络模型时产生的GEDump图“ge_proto_000000xx_after_infershape.txt”中存在dtype与shape规格不连续的情况,说明InferShape处理有错误。
InferShape流程里,GE将前一个算子的输出Tesnor刷新到下一个算子的输入Tensor,算子根据输入推导并更新输出Tensor,若算子未自定义实现InferShape,则保留原图的输出Tesnor。
某些算子在进行输出TensorDesc的推导时依赖算子输入的实际值,此时就需要算子的输入必须是常量节点(Const)。如果算子的输入不为常量节点, 可能会造成推导输出Tensor的shape出错。
算子校验函数实现
算子Verify函数的实现使用如下接口:
IMPLEMT_VERIFIER(OpType,func name)
传入的OpType为基于Operator类派生出来的子类,会自动生成一个类型为此子类的对象op,可以使用子类的成员函数获取算子的相关属性
Verify函数主要校验算子内在关联关系,例如对于多输入算子,多个tensor的dtype需要保持一致,此时需要校验多个输入的dtype,其他情况dtype不需要校验。
IMPLEMT_VERIFER(PowPowverify){
DataType input type_x = op.GetnputDesc(“x”).GetDataType();
DataType input type_y = op.GetinputDesc(“y”).GetDataType();
if(input_type_x!= input_type_y){
return GRAPH FAILED:
}
return GRAPH SUCCESS
}
注册算子原型
GE提供REGOP宏,以“.”链接INPUT、OUTPUT、ATTR等接口注册算子的输入、输出和属性信息,最终以OPENDFACTORY_REG接口结束,完成算子的注册
中输入输出的描述信息顺序需要与算子实现中定义保持一致,ATTR的顺序可变
注册代码实现如下所示:
namespace get
REG_OP(OpType)//算子类型名称
.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)
}
算子实现
DSL 算子开发过程
TBE算子输入占位符
样例:data = tvm.placeholder(shape, name=”data”,dtype=dtype)
tvm.placeholder是tvm框架的API,用来为算子执行时接收的数据占位,通俗理解与C语言中%d%s一样,返回的是一个Tensor对象,上例中使用data表示;入参为shape,name,dtype,是为Tensor对象的属性。
这里的输入是指算子执行时的输入数据,与编译时期入参不同,编译时期入参(xy)是为了得到算子执行文件的入参。
TBE算子代码结构——构建
TBE算子中的构建操作
样例: dsl.build(sch,config)
TBE框架提供了build()API,传入schedule以及相关的配置项,即可完成编译,生成最终硬件可执行文件。
TBE算子编译过程
TBE算子编;译过程分为DSL->Schedule->pass->codegen四步
- Schedule:经过Schedule,计算过程描述逻辑发生了转变,针对硬件的存储上限做了切分、合并等操作
- Pass:Pass阶段会对计算过程描述进行指令替换,将数学方式表示的计算描述,,映射为硬件可以读懂的指令
且会对指令进行优化,以获得更高效的性能,经过pass后,计算过程变为IR表示。 - Codegen:Codegen是TBE框架执行的最后一步,将Pass产生的IR构建为cce代码,进而经过compile生成二进制的可执行文件。
DSL的功能调试
DSL的功能调试是在CPU上进行算子的功能验证:
支持算子在CPU上执行
支持打印中间Tensor
支持将结果数据与期望数据进行比对
DSL功能调试代码逻辑
算子信息库定义
算子信息库作为算子开发的交付件之一,主要体现算子在异腾AI处理器上的具体实现规格,包括算子支持输入输出dtype、format以及输入shape等信息。网络运行时,FE会根据算子信息库中的算子信息做基本校验,选择dtypeformat等信息,并根据算子信息库中信息找到对应的算子实现文件进行编译,用于生成算子二进制文件
关键字:FE(融合引擎)
一句话总结:通过一个配置文件配置一下自定义算子所支持的输入(和输出)的数据类型/数据排布格式组合。
[Add]
inputO.name=x1
inputo.dtype=float16,float16,float16,float16;float;float;float,float,int32,int32,int32,int32
inputo.shape=all
inputo.paramType=required
inputo.format=NCHW,NC1HWCO,NHWC,ND,NCHW,NC1HWCO,NHWCND,NCHW,NC1HWCO,NHWCN
input.name=x2
inputl.dtype=float16,float16,float16,float16,float,float,float,float,int32,int32,int32,int32
inputl.shape=all
input.paramType=required
inputl.format=NCHW,NCIHWCO,NHWC,ND,NCHWNCIHWCO,NHWC,ND,NCHW,NCIHWCO,NHWC,Noutputo.name=y
outputo.dtype=float16,float16float16,float16,floatfloat,float,float,int32,int32,int32,int32
outputo.shape=all
outputo.paramType=required
outputo.format=NCHW,NCHWCO,NHWC,ND,NCHW,NC1HWCO,NHWC,ND,NCHW,NCHWCO,NHWCopFile.value=add
opnterface.value=add
算子UT测试
基于Mindstudio进行算子开发的场景下,用户可基于MindStudio进行算子的UT测试,UT(UnitTest:单元测试)是开发人员进行算子代码验证的手段之一,主要目的是:
测试算子代码的正确性,验证输入输出结果与设计的一致性
UT侧重于保证算子程序能够跑通,选取的场景组合应能覆盖算子代码的所有分支(一般来说覆盖率要达到100%),从而降低不同场景下算子代码的编译失败率。
(测试类的详细定义可参见Ascend-cann-toolkit安装目录/ascend-toolkit/{version}/{arch}-linux/toolkit/python/site-packages/op_test_frame/ut/op_ut.py文件。)
测试流程:
选择算子
实现test_算子名称_impl.py文件
测试用例编写完毕后,在ut->sqrt目录上单击右键,Run TBE Operator’算子名称’ UT impl with Coverage