该开发流程以为基础,除了需要提供工程化算子开发中的算子实现文件外,还需要额外交付算子入图的代码文件。本节仅提供算子入图代码文件的开发指导。
假设下图是我们需要使用的网络模型,您可能会想直接逐个算子调用,根据输入tensor得到输出tensor就可以完成网络的运行,但在图模式场景下,实际的网络模型生成过程中,会先进行tensor shape以及datatype的推导。这样可以让我们在图执行之前,就知道各tensor的数据类型和形状,提前校验其正确性;同时提前推理出算子的输出张量描述,包括张量的形状、数据类型及数据排布格式等信息,算子构图准备阶段就可以为所有的张量静态分配内存,避免动态内存分配带来的开销。
下面的网络模型经过shape和datatype推导之后,可以得到灰色底纹框中的推导信息:
图 1 shape与datatype推导示意图[object Object][object Object]
除了tiling实现外,算子入图时需要额外提供的实现代码有以下几种:
- datatype推导:根据算子的输入datatype、算子逻辑及算子属性等信息,推理出算子的输出张量datatype。
- shape推导:根据算子的输入shape、算子逻辑及算子属性等信息,推理出算子的输出张量shape。
- ShapeRange推导:编译时无法推导输出shape,只能推导输出shape range,执行完才能得出输出shape。 在下发时需要按照输出shape range来申请最大输出内存,该类算子需要提供ShapeRange推导函数。
- 声明数据依赖:部分算子在InferShape时,需要依赖某个输入的具体值才可以进行,这类算子被称为“数据依赖算子”,对应的输入被称为“数据依赖输入”。该类算子在注册时,需要声明其数据依赖输入。
下表列出了不同类型的算子对上述实现代码的要求。
表 1 不同的类型的算子对入图实现代码的要求
[object Object][object Object]
[object Object]实际开发时通过固定的datatype和shape推导原型实现推导函数,然后再通过SetInferShape、SetInferDataType接口来关联对应的shape推导函数,样例如下。
以AddCustom算子为例,InferDataType的实现如下所示。该样例中输出tensor的数据类型与输入tensor的数据类型相同,所以直接将任意一个输入tensor的数据类型赋给输出tensor即可。
如下示例则给出了更灵活的datatype推导样例,当输入的数据类型为DT_INT4时,其输出的数据类型为DT_INT32。
简单的shape推导逻辑可以使用Follow接口来表达,比如输出shape和输入shape相同的情况。示例如下:输出“y1”Follow输入“x1”场景,指定Follow模式为SHAPE,此时“y1”的shape将会和“x1”保持一致。
无法在原型定义中通过Follow表达的情况需要开发者编写InferShape函数,InferShape函数的原型是固定的,如下示例,接受一个InferShapeContext作为输入,从此context上可以获取到输入、输出的shape指针等内容。输入shape为const类型,因此InferShape时,输入shape是只读、不允许修改的。InferShape成功后,返回ge::GRAPH_SUCCESS,其他返回值被认为推导失败。推导失败后,执行过程结束退出。
以ReShape算子为例,InferShape的实现如下所示。根据第1个输入(shape输入)的值,Reshape算子将第0个输入(x输入)的shape做变换,并输出到其第0个输出(y输出)上。Reshape的InferShape实现为:
InferShapeContext public继承自ExtendedKernelContext,因此ExtendedKernelContext中提供的方法如获取算子type、name、属性等接口均可以在InferShapeContext实例中调用。
[object Object]
某些算子的输出Shape在计算完成后才能确定。比如unique算子,其Shape的推导逻辑如下:
给定一维Tensor x,找到其中不重复的元素,返回去重后的Tensor y,输出idx与输入x大小相同,保存x每个元素在y中的索引。
由此可知,y的shape在编译时为[-1],unique执行后shape才确定。
在入图场景执行时,需要在执行前分配输出内存,而内存的大小依赖于输出Shape和数据类型。对于此类算子,由于输出Shape在执行后才能确定,因此需要根据输出Shape的范围,按照最大范围申请输出内存,以确保有足够的空间供计算函数写入输出Tensor。
这种场景下,开发者需要自行实现InferShapeRange函数,来推导输出Shape的范围。下面以unique算子为例子,介绍InferShapeRange函数的实现方法。
在InferShape、Tiling时,可以通过context实例获取算子IR属性值,所谓IR属性,是指在IR注册时定义的属性,以TransData算子为例:
其原型定义中声明了src_format、dst_format、group三个属性,可以通过如下方式获取算子属性:
通过index而不是字符串name来索引输入输出,对于带有OPTIONAL、DYNAMIC类型输入的算子,可能出现实例化后,单纯通过index无法索引到具体输入的问题,以DynamicRNNV3算子为例:
由于DynamicRNNV3算子有连续的多个optional输入,这导致init_h及其后面的输入的实例化后index都是不确定的,对于这种类型的算子,可以通过GetOptionalInputShape传入原型对应的index来获取对应的输入shape等数据,以InferShape为例:
对于dynamic类型的输入,实例化后的输入可能是一到多个,对于此类输入,获取方式为:
本节举例的获取optional、dynamic输入的方式,在InferShape、Tiling函数中均可以调用。
一般来说,具备输入shape后,算子可以通过InferShape推导出输出shape。然而部分算子在InferShape时,需要依赖某个输入的具体值才可以进行,这类算子被称为“数据依赖算子”,对应的输入被称为“数据依赖输入”。以Reshape算子为例,其依据shape输入的描述,对输入的shape做调整,因此Reshape算子依赖shape输入的值。这类算子需要在原型定义时通过ValueDepend接口声明对应的输入为数据依赖输入。
根据第1个输入(shape输入)的值,Reshape算子将第0个输入(x输入)的shape做变换,并输出到其第0个输出(y输出)上。Reshape的InferShape实现为:
[object Object]