【2023 - CANN训练营第二季 】-- Ascend C （2)Tilling

CANN训练营第二季 – Ascend C （2)Tilling

​ 在之前的的基本概念介绍中，我们提到了Ascend C的并行计算需要将数据搬入Local Memory中计算，计算后再搬出。

​ 大多数情况下，Local Memory的存储，无法完整的容纳算子的输入与输出，需要每次搬运一部分输入进行计算然后搬出，再搬运下一部分输入进行计算，直到得到完整的最终结果，这个数据切分、分块计算的过程称之为Tiling。根据算子的shape等信息来确定数据切分算法相关参数（比如每次搬运的块大小，以及总共循环多少次）的计算程序，称之为Tiling实现。

​ Tiling实现完成后，获取到的Tiling切分算法相关参数，会传递给kernel侧，用于指导并行数据的切分。由于Tiling实现中完成的均为标量计算，AI Core并不擅长，所以我们将其独立出来放在host CPU上执行。

1.Tilling的基本流程

Tilling 的流程可以简化为输入输出。

如上图所示，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内存。

而细分之后，Tilling的基本流程如下所示：

1.1 Tilling的结构定义

下面将从一个简单的Add算子为例介绍Tiling的实现流程。本样例中待处理数据的Shape大小可以平均分配到每个核上，并且可以对齐到一个data block（32B）的大小。

1.2 示例代码

*首先完成算子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

然后完成算子host实现cpp文件中Tiling函数实现，该文件命名为“算子名称.cpp”，位于算子工程的op_host目录下。样例代码如下：

namespace optiling {
const uint32_t BLOCK_DIM = 8;
const uint32_t TILE_NUM = 8;
//获取TilingContext的上下文,即Tiling函数的入参gert::TilingContext* context
static ge::graphStatus TilingFunc(gert::TilingContext* context) 
{
    // 用TilingData定义一个具体的实例
    TilingData tiling;
    //获取输入shape信息
    uint32_t totalLength = context->GetInputTensor(0)->GetShapeSize();
    //通过SetBlockDim接口设置blockdim。
    context->SetBlockDim(BLOCK_DIM);
    // 设置TilingData
    tiling.set_totalLength(totalLength);
    tiling.set_tileNum(TILE_NUM);
    // 序列化并保存
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    //(可选)通过SetTilingKey设置TilingKey。通过GetWorkspaceSizes获取workspace size指针，并设置size大小。此处仅作为举例，设置workspace的大小为0。
    context->SetTilingKey(1);
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
    
    return ge::GRAPH_SUCCESS;
}
} // namespace optiling

注意：这是两个存放在op_host文件夹下的文件，op_host文件夹是Add算子的一个子文件夹，具体在之后add算子解析中会介绍。

1.3 固定和动态shape

Tilling本身是数据搬运的算法，那么自然有搬固定和动态shape的区分，在实现逻辑上，固定shape由于输入的大小都是已知的，所以可以把详细信息都在编译时直接计算处理，动态shape则是去掉控制形状的变量的输入，改为在核函数中传入一个tIlling,来控制这几个变量。

1.4 非对齐shape

另外，针对一些非对齐shape，比如算子的输入shape为（1，1999），支持的数据类型为half类型，既无法对齐到一个block的大小（32B），也无法平均分配到每个核上，需要一些特殊的tiling处理方法。

1.定义和注册算子需要使用的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

2. 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

size_t *currentWorkspace = context->GetWorkspaceSizes(1);
currentWorkspace[0] = 0;
return ge::GRAPH_SUCCESS;

}
} // namespace optiling