在章节已经介绍了host侧Tiling核心的实现方法,本章节侧重于介绍接入CANN框架时编程模式和API的使用。
大多数情况下,Local Memory的存储,无法完整的容纳算子的输入与输出,需要每次搬运一部分输入进行计算然后搬出,再搬运下一部分输入进行计算,直到得到完整的最终结果,这个数据切分、分块计算的过程称之为Tiling。根据算子的shape等信息来确定数据切分算法相关参数(比如每次搬运的块大小,以及总共循环多少次)的计算程序,称之为Tiling实现。
Tiling实现完成后,获取到的Tiling切分算法相关参数,会传递给kernel侧,用于指导并行数据的切分。由于Tiling实现中完成的均为标量计算,AI Core并不擅长,所以我们将其独立出来放在host CPU上执行。
图 1 Tiling实现的输入输出[object Object][object Object]
如上图所示,Tiling实现即为根据算子shape等信息来确定切分算法相关参数的过程,这里的算子shape等信息可以理解为是Tiling实现的输入,切分算法相关参数可以理解为是Tiling实现的输出。输入和输出都通过Tiling函数的参数(TilingContext* context上下文结构)来承载。也就是说,开发者可以从上下文结构中获取算子的输入、输出以及属性信息,也就是Tiling实现的输入,经过Tiling计算后,获取到TilingData数据结构(切分算法相关参数)、numBlocks变量、用于选择不同的kernel实现分支的TilingKey、算子workspace的大小,也就是Tiling实现的输出,并将这些输出设置到上下文结构中。
TilingData、numBlocks、TilingKey、workspace这些概念的具体解释如下:
TilingData:切分算法相关参数,比如每次搬运的块大小,以及总共循环多少次,通过结构体存储,由开发者自行设计。
TilingData结构定义支持单结构定义方法,也支持结构体嵌套:
单结构定义方法,以平铺的形式定义:
[object Object]Tiling实现函数中对tiling结构成员赋值的方式如下:
[object Object]支持结构体嵌套:
[object Object]Tiling实现函数中对tiling结构成员赋值的方式如下:
[object Object]
numBlocks:规定了核函数将会在几个核上执行。例如,需要计算8M的数据,每个核上计算1M的数据,numBlocks设置为8,但是为了充分利用硬件资源,一般将numBlocks设置为硬件平台的核数,根据核数进行数据切分。
[object Object]
TilingKey(可选):TilingKey是一个算子内为了区分不同的实现而将kernel代码进行区分的方法,该方法类似于C++的Template模板机制,可减少不必要的icache miss以及scalar耗时,有助于优化单次调用kernel的性能。不同的kernel实现分支可以通过TilingKey来标识,host侧设置TilingKey后,可以选择对应的分支。例如,一个算子在不同的shape下,有不同的算法逻辑,kernel侧可以通过TilingKey来选择不同的算法逻辑,在host侧Tiling算法也有差异,host/kernel侧通过相同的TilingKey进行关联。
假如有如下kernel代码:
[object Object]如果函数ProcessA、ProcessB两个函数是个非常大的函数,那么上述代码在编译后会变得更大,而每次kernel运行只会选择1个分支,条件的判断和跳转在代码大到一定程度(16-32K,不同芯片存在差异)后会出现icache miss。通过TilingKey可以对这种情况进行优化,给2个kernel的处理函数设置不同的TilingKey 1和2:
[object Object]这样device kernel编译时会自动识别到2个TilingKey并编译2个kernel入口函数,将条件判断进行常量折叠。同时需要和host tiling函数配合,判断走ProcessA的场景设置TilingKey为1,走ProcessB的场景设置TilingKey为2:
[object Object][object Object]
WorkspaceSize:workspace是设备侧Global Memory上的一块内存。在Tiling函数中可以设置workspace的大小。设置后:单算子API执行场景,可以通过单算子API调用第一段接口获取workspace的大小,然后由开发者申请对应大小的Global Memory;入图场景,框架会根据设置的大小自动申请对应大小的Global Memory。申请workspace后,在算子Kernel实现时,可以使用这块workspace内存。
workspace内存分为两部分:Ascend C API需要的workspace内存和算子实现使用到的workspace内存(按需)。
Ascend C API需要预留workspace内存
API在计算过程需要一些workspace内存作为缓存,因此算子Tiling函数需要为API预留workspace内存,预留内存大小通过接口获取。
算子实现使用到的workspace内存(按需)
算子内部需要通过额外的device内存进行数据交换或者缓存的时候才需要分配,根据算子计算的空间自行分配。
整体的workspace内存就是上述两部分之和,在Tiling函数中设置方法如下:
[object Object]
Tiling实现开发的流程图如下:
图 2 Tiling开发流程图[object Object][object Object]
下面将从一个简单的Add算子为例介绍Tiling的实现流程。本样例中待处理数据的Shape大小可以平均分配到每个核上,并且可以对齐到一个datablock(32B)的大小。
首先完成算子TilingData结构定义头文件的编写,该文件命名为“算子名称_tiling.h”,位于算子工程的op_host目录下。样例代码如下:
具体的编写步骤如下:
代码框架编写,需要增加#ifndef...的判断条件,防止头文件的重复包含;需要包含register/tilingdata_base.h头文件,tilingdata_base.h中定义了多个用于tilingdata注册的宏。样例代码如下:
[object Object]TilingData参数设计,TilingData参数本质上是和并行数据切分相关的参数,本示例算子使用了2个tiling参数:totalLength、tileNum。totalLength是指需要计算的数据量大小,tileNum是指每个核上总计算数据分块个数。比如,totalLength这个参数传递到kernel侧后,可以通过除以参与计算的核数,得到每个核上的计算量,这样就完成了多核数据的切分。
[object Object][object Object]TilingData结构定义,通过BEGIN_TILING_DATA_DEF接口定义一个TilingData的类,通过TILING_DATA_FIELD_DEF接口增加TilingData的两个字段totalLength、tileNum,通过END_TILING_DATA_DEF接口结束TilingData定义。相关接口的详细说明请参考TilingData结构定义。
[object Object]注册TilingData结构,通过REGISTER_TILING_DATA_CLASS接口,注册TilingData类,和自定义算子相关联。REGISTER_TILING_DATA_CLASS第一个参数为op_type(算子类型),本样例中传入AddCustom,第二个参数为TilingData的类名。REGISTER_TILING_DATA_CLASS接口介绍请参考TilingData结构注册。
[object Object]
然后完成算子host实现cpp文件中Tiling函数实现,该文件命名为“算子名称.cpp”,位于算子工程的op_host目录下。Tiling函数的原型是固定的,接受一个TilingContext作为输入,在此context上可以获取到输入、输出的Shape指针等内容。注册的Tiling函数由框架调用,调用时会传入TilingContext参数。样例代码如下:
具体步骤如下:
获取TilingContext的上下文,即Tiling函数的入参gert::TilingContext* context。
设置TilingData。在中定义了TilingData类后,可以创建该类的一个实例,并通过调用set_{field_name}方法来设置各个字段值(其中field_name是中定义的tiling字段名)。设置完tiling字段后,通过调用SaveToBuffer方法完成TilingData实例的序列化和保存。
通过上下文获取输入输出shape信息。本样例中通过TilingContext的GetInputShape接口获取输入的shape大小。
[object Object]设置TilingData。通过调用set_{field_name}方法来设置TilingData的字段值。
[object Object]调用TilingData类的SaveToBuffer接口完成序列化并保存至TilingContext上下文。SaveToBuffer的第一个参数为存储Buffer的首地址,第二个参数为Buffer的长度。通过调用GetRawTilingData获取无类型的TilingData的地址,再通过GetData获取数据指针,作为Buffer的首地址;通过调用GetRawTilingData获取无类型的TilingData的地址,再通过GetCapacity获取TilingData的长度,作为Buffer的长度。完成SaveToBuffer操作后需要通过SetDataSize设置TilingData的长度,该长度通过TilingData类的GetDataSize接口获取。
[object Object]
通过SetBlockDim接口设置numBlocks。
[object Object](可选)通过SetTilingKey设置TilingKey。
[object Object](可选)通过GetWorkspaceSizes获取workspace size指针,并设置size大小。此处仅作为举例,设置workspace的大小为0。
[object Object]