Tiling实现
大多数情况下,Local Memory的存储,无法完整的容纳算子的输入与输出,需要每次搬运一部分输入进行计算然后搬出,再搬运下一部分输入进行计算,直到得到完整的最终结果,这个数据切分、分块计算的过程称之为Tiling。根据算子的shape等信息来确定数据切分算法相关参数(比如每次搬运的块大小,以及总共循环多少次)的计算程序,称之为Tiling实现。
Tiling实现完成后,获取到的Tiling切分算法相关参数,会传递给kernel侧,用于指导并行数据的切分。由于Tiling实现中完成的均为标量计算,AI Core并不擅长,所以我们将其独立出来放在host CPU上执行。
如上图所示,Tiling实现即为根据算子shape等信息来确定切分算法相关参数的过程,这里的算子shape等信息可以理解为是Tiling实现的输入,切分算法相关参数可以理解为是Tiling实现的输出。输入和输出都通过Tiling函数的参数(TilingContext* context上下文结构)来承载。也就是说,开发者可以从上下文结构中获取算子的输入、输出以及属性信息,也就是Tiling实现的输入,经过Tiling计算后,获取到TilingData数据结构(切分算法相关参数)、block_dim变量、用于选择不同的kernel实现分支的TilingKey、算子workspace的大小,也就是Tiling实现的输出,并将这些输出设置到上下文结构中。
TilingData、block_dim、TilingKey、workspace这些概念的具体解释如下:
- TilingData:切分算法相关参数,比如每次搬运的块大小,以及总共循环多少次,通过结构体存储,由开发者自行设计。
- block_dim:算子数据切分的份数。例如,需要计算8M的数据,每个核上计算1M的数据,block_dim设置为8,但是为了充分利用硬件资源,一般将block_dim设置为硬件平台的核数,根据核数进行数据切分。
- TilingKey(可选):不同的kernel实现分支可以通过TilingKey来标识,host侧设置TilingKey后,可以选择对应的分支。例如,一个算子在不同的shape下,有不同的算法逻辑,kernel侧可以通过TilingKey来选择不同的算法逻辑,在host侧Tiling算法也有差异,host/kernel侧通过相同的TilingKey进行关联。
- workspace size(可选):workspace是设备侧Global Memory上的一块内存。在Tiling函数中可以设置workspace的大小,框架侧会为其在申请对应大小的设备侧Global Memory,在对应的算子kernel侧实现时可以使用这块workspace内存。
Tiling实现基本流程
Tiling实现开发的流程图如下:
下面将从一个简单的Add算子为例介绍Tiling的实现流程。本样例中待处理数据的Shape大小可以平均分配到每个核上,并且可以对齐到一个data block(32B)的大小。
首先完成算子TilingData结构定义头文件的编写,该文件命名为“算子名称_tiling.h”,位于算子工程的op_host目录下。样例代码如下:
#ifndef ADD_CUSTOM_TILING_H #define ADD_CUSTOM_TILING_H #include "register/tilingdata_base.h" namespace optiling { BEGIN_TILING_DATA_DEF(TilingData) // 注册一个tiling的类,以tiling的名字作为入参 TILING_DATA_FIELD_DEF(uint32_t, totalLength); // 添加tiling字段,总计算数据量 TILING_DATA_FIELD_DEF(uint32_t, tileNum); // 添加tiling字段,每个核上总计算数据分块个数 END_TILING_DATA_DEF; // 注册算子tilingdata类到对应的AddCustom算子 REGISTER_TILING_DATA_CLASS(AddCustom, TilingData) } #endif // ADD_CUSTOM_TILING_H
具体的编写步骤如下:
- 代码框架编写:需要增加#ifndef...的判断条件,防止头文件的重复包含;需要包含register/tilingdata_base.h头文件,tilingdata_base.h中定义了多个用于tilingdata注册的宏。样例代码如下:
#ifndef ADD_CUSTOM_TILING_H #define ADD_CUSTOM_TILING_H #include "register/tilingdata_base.h" namespace optiling { // tiling结构定义和注册代码 // ... } #endif // ADD_CUSTOM_TILING_H
- TilingData参数设计,TilingData参数本质上是和并行数据切分相关的参数,本示例算子使用了2个tiling参数:totalLength、tileNum。totalLength是指需要计算的数据量大小,tileNum是指每个核上总计算数据分块个数。比如,totalLength这个参数传递到kernel侧后,可以通过除以参与计算的核数,得到每个核上的计算量,这样就完成了多核数据的切分。
- TilingData结构定义,通过BEGIN_TILING_DATA_DEF接口定义一个TilingData的类,通过TILING_DATA_FIELD_DEF接口增加TilingData的两个字段totalLength、tileNum,通过END_TILING_DATA_DEF接口结束TilingData定义。相关接口的详细说明请参考TilingData结构定义。
BEGIN_TILING_DATA_DEF(TilingData) // 注册一个tiling的类,以tiling的名字作为入参 TILING_DATA_FIELD_DEF(uint32_t, totalLength); // 添加tiling字段,总计算数据量 TILING_DATA_FIELD_DEF(uint32_t, tileNum); // 添加tiling字段,每个核上总计算数据分块个数 END_TILING_DATA_DEF;
- 注册TilingData结构,通过REGISTER_TILING_DATA_CLASS接口,注册TilingData类,和自定义算子相关联。REGISTER_TILING_DATA_CLASS第一个参数为op_type(算子类型),本样例中传入AddCustom,第二个参数为TilingData的类名。REGISTER_TILING_DATA_CLASS接口介绍请参考TilingData结构注册。
// 注册算子tilingdata类到对应的AddCustom算子 REGISTER_TILING_DATA_CLASS(AddCustom, TilingData)
namespace optiling { const uint32_t BLOCK_DIM = 8; const uint32_t TILE_NUM = 8; static ge::graphStatus TilingFunc(gert::TilingContext* context) { TilingData tiling; uint32_t totalLength = context->GetInputTensor(0)->GetShapeSize(); context->SetBlockDim(BLOCK_DIM); tiling.set_totalLength(totalLength); tiling.set_tileNum(TILE_NUM); tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); context->GetRawTilingData()->SetDataSize(tiling.GetDataSize()); context->SetTilingKey(1); size_t *currentWorkspace = context->GetWorkspaceSizes(1); currentWorkspace[0] = 0; return ge::GRAPH_SUCCESS; } } // namespace optiling
具体步骤如下:
- 获取TilingContext的上下文,即Tiling函数的入参gert::TilingContext* context。
- 设置TilingData。上文3中,定义了TilingData的类,此时可以用TilingData定义一个具体的实例,通过调用TilingData类的set_+field_name接口来设置TilingData的字段值,通过调用TilingData类的SaveToBuffer接口完成TilingData的序列化和保存。
- 通过上下文获取输入输出shape信息。本样例中通过TilingContext的GetInputTensor接口获取输入Tensor,再通过Tensor类的GetShapeSize获取当前Tensor的shape大小。
// 获取输入shape信息 uint32_t totalLength = context->GetInputTensor(0)->GetShapeSize();
- 设置TilingData。通过调用set_+field_name方法来设置TilingData的字段值。
// 用TilingData定义一个具体的实例 TilingData tiling; // 设置TilingData tiling.set_totalLength(totalLength); tiling.set_tileNum(TILE_NUM);
- 调用TilingData类的SaveToBuffer接口完成序列化并保存至TilingContext上下文。SaveToBuffer的第一个参数为存储Buffer的首地址,第二个参数为Buffer的长度。通过调用GetRawTilingData获取无类型的TilingData的地址,再通过GetData获取数据指针,作为Buffer的首地址;通过调用GetRawTilingData获取无类型的TilingData的地址,再通过GetCapacity获取TilingData的长度,作为Buffer的长度。完成SaveToBuffer操作后需要通过SetDataSize设置TilingData的长度,该长度通过TilingData类的GetDataSize接口获取。
// 序列化并保存 tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
- 通过上下文获取输入输出shape信息。本样例中通过TilingContext的GetInputTensor接口获取输入Tensor,再通过Tensor类的GetShapeSize获取当前Tensor的shape大小。
- 通过SetBlockDim接口设置blockdim。
context->SetBlockDim(BLOCK_DIM);
- (可选)通过SetTilingKey设置TilingKey。
context->SetTilingKey(1);
- (可选)通过GetWorkspaceSizes获取workspace size指针,并设置size大小。此处仅作为举例,设置workspace的大小为0。
size_t *currentWorkspace = context->GetWorkspaceSizes(1); currentWorkspace[0] = 0;
Tiling参数设计更多样例-非对齐shape
针对一些非对齐shape,比如算子的输入shape为(1,1999),支持的数据类型为half类型,既无法对齐到一个block的大小(32B),也无法平均分配到每个核上,需要一些特殊的tiling处理方法。
- 定义和注册算子需要使用的tiling参数,本示例非对齐的AddCustom算子使用了5个tiling参数:formerNum, tailNum, formerLength, tailLength,alignNum。
#ifndef ADD_CUSTOM_UNALIGN_TILING_H #define ADD_CUSTOM_UNALIGN_TILING_H #include "register/tilingdata_base.h" namespace optiling { BEGIN_TILING_DATA_DEF(TilingDataUnalign) TILING_DATA_FIELD_DEF(uint32_t, formerNum); // 添加tiling字段,分配到较多数据量的核心数,即大块 TILING_DATA_FIELD_DEF(uint32_t, tailNum); // 添加tiling字段,分配到较少数据量的核心数,即小块 TILING_DATA_FIELD_DEF(uint32_t, formerLength); // 添加tiling字段,大块的长度 TILING_DATA_FIELD_DEF(uint32_t, tailLength); // 添加tiling字段,小块的长度 TILING_DATA_FIELD_DEF(uint32_t, alignNum); // 添加tiling字段,需要对齐到的最小数据量 END_TILING_DATA_DEF; // 注册算子tilingdata类到对应的AddCustom算子 REGISTER_TILING_DATA_CLASS(AddCustomUnalign, TilingDataUnalign) } #endif // ADD_CUSTOM_UNALIGN_TILING_H
- tiling实现,在“op_host/add_custom.cpp”中注册tiling实现接口。
namespace optiling { constexpr uint32_t BLOCK_DIM = 8; constexpr uint32_t SIZE_OF_HALF = 2; constexpr uint32_t BLOCK_SIZE = 32; // shape需要对齐到的最小单位 constexpr uint32_t ALIGN_NUM = BLOCK_SIZE / SIZE_OF_HALF; static ge::graphStatus TilingFunc(gert::TilingContext *context) { TilingDataUnalign tiling; uint32_t totalLength = context->GetInputTensor(0)->GetShapeSize(); context->SetBlockDim(BLOCK_DIM); // 如果是非对齐的shape,需要向上对齐到最小单位 uint32_t totalLengthAligned = ((totalLength + ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM; // 把所有的数据尽可能均匀地分配到每个核上,如果不能均分的话,那么会有部分核多算一个最小单位ALIGN_NUM // 通过模的计算,可以得到多算一个最小单位的核的数量,也可以得到少算一个最小单位的核的数量 // eg:1999 对齐后的总数据量为2000个数,核心数为8,数据块的最小单位是16,那么: // 1、最小单位数据块的总数:2000 / 16 = 125 // 2、有5个核会分到16个最小单位的数据块:125 % 8 =5,可以称之为大块 // 3、有3个核会分到15个最小单位的数据块:8 - 5 = 3,可以称之为小块 uint32_t formerNum = (totalLengthAligned / ALIGN_NUM) % BLOCK_DIM; uint32_t tailNum = BLOCK_DIM - formerNum; // 计算大块和小块的数据量 uint32_t formerLength = ((totalLengthAligned / BLOCK_DIM + ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM; uint32_t tailLength = (totalLengthAligned / BLOCK_DIM / ALIGN_NUM) * ALIGN_NUM; tiling.set_formerNum(formerNum); tiling.set_tailNum(tailNum); tiling.set_formerLength(formerLength); tiling.set_tailLength(tailLength); tiling.set_alignNum(ALIGN_NUM); tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); context->GetRawTilingData()->SetDataSize(tiling.GetDataSize()); context->SetTilingKey(1); size_t *currentWorkspace = context->GetWorkspaceSizes(1); currentWorkspace[0] = 0; return ge::GRAPH_SUCCESS; } } // namespace optiling
Tiling参数设计更多样例-属性信息通过TilingData传递
如果算子包含属性信息,该属性信息可以通过TilingData传递到kernel侧,参与kernel侧算子核函数的计算。以ReduceMax算子为例,该算子用于对输入数据按维度dim返回最大值,并且返回索引。ReduceMax算子有两个属性,reduceDim和isKeepDim,reduceDim表示按照哪一个维度进行reduce操作;isKeepDim表示是否需要保持输出的维度与输入一样。本样例仅支持对最后一维做reduce操作,输入数据类型为half。
- ReduceMaxCustom算子TilingData的定义如下:这里我们重点关注reduceAxisLen。参数reduceAxisLen表示获取reduceDim轴的长度,这里也就是最后一维的长度。该参数后续会通过TilingData传递到kernel侧参与计算。
#ifndef REDUCE_MAX_CUSTOM_TILING_H #define REDUCE_MAX_CUSTOM_TILING_H #include "register/tilingdata_base.h" namespace optiling { BEGIN_TILING_DATA_DEF(ReduceMaxTilingData) TILING_DATA_FIELD_DEF(uint32_t, reduceAxisLen); // 添加tiling字段,reduceDim轴的长度 //其他TilingData参数的定义 ... END_TILING_DATA_DEF; // 注册算子tilingdata类到对应的ReduceMaxCustom算子 REGISTER_TILING_DATA_CLASS(ReduceMaxCustom, ReduceMaxTilingData) } #endif // REDUCE_MAX_CUSTOM_TILING_H
- ReduceMaxCustom算子的Tiling实现如下。这里我们重点关注属性信息通过TilingData传递的过程:首先通过TilingContext上下文从attr获取reduceDim属性值;然后根据reduceDim属性值获取reduceDim轴的长度并设置到TilingData中。
namespace optiling { static ge::graphStatus TilingFunc(gert::TilingContext* context) { ReduceMaxTilingData tiling; // 从attr获取reduceDim属性值,因为reduceDim是第一个属性,所以GetAttrPointer传入的索引值为0 const gert::RuntimeAttrs* attrs = context->GetAttrs(); const uint32_t* reduceDim = attrs->GetAttrPointer<uint32_t>(0); // 获取reduceDim轴的长度 const gert::StorageShape* xShapePtr = context->GetInputShape(0); const gert::Shape& xShape = xShapePtr->GetStorageShape(); const uint32_t reduceAxisLen = xShape.GetDim(*reduceDim); // 计算TilingData中除了reduceAxisLen之外其他成员变量的值 ... // 将reduceAxisLen设置到tiling结构体中,传递到kernel函数使用 tiling.set_reduceAxisLen(reduceAxisLen); // 设置TilingData中除了reduceAxisLen之外其他成员变量的值 ... // TilingData序列化保存 tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); context->GetRawTilingData()->SetDataSize(tiling.GetDataSize()); ... return ge::GRAPH_SUCCESS; }} // namespace optiling