动态shape场景

在算子实现章节，已经介绍了简单的固定shape矢量算子的kernel侧实现，算子的shape、数据类型都是固定的；在实际的算子开发场景中，这些信息是支持动态变化的，场景会更加灵活和复杂。下文重点进行动态shape与固定shape差异点的介绍。

最主要的区别是：动态Shape场景下，输入的Shape是未知的。一些与输入Shape相关的变量（比如每次搬运的块大小等），需要通过Tiling计算出来，然后传递到kernel侧，kernel侧使用该参数进行后续的计算。

算子实现章节中固定shape的算子样例中，TILE_NUM（每个核上总计算数据分块个数）、BLOCK_LENGTH（每个核上总计算数据大小）、TILE_LENGTH（每个分块大小）等是固定的数值。

constexpr int32_t TOTAL_LENGTH = 8 * 2048;                            // total length of data
constexpr int32_t USE_CORE_NUM = 8;                                   // num of core used
constexpr int32_t BLOCK_LENGTH = TOTAL_LENGTH / USE_CORE_NUM;         // length computed of each core
constexpr int32_t TILE_NUM = 8;                                       // split data into 8 tiles for each core
constexpr int32_t BUFFER_NUM = 2;                                     // tensor num for each queue
constexpr int32_t TILE_LENGTH = BLOCK_LENGTH / TILE_NUM / BUFFER_NUM; // each tile length is seperated to 2 part, due to double buffer

如果需要将上述代码转换为动态shape，需要在核函数定义中增加Tiling参数，在host侧计算Tiling参数并传入，然后基于Tiling参数计算得到singleCoreSize（每个核上总计算数据大小）、tileNum（每个核上总计算数据分块个数）、tileLength（每个分块大小）等变量。

__aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t totalLength, uint32_t tileNum)
{
    ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
    this->blockLength = totalLength / GetBlockNum();
    this->tileNum = tileNum;
    ASSERT(tileNum != 0 && "tile num can not be zero!");
    this->tileLength = this->blockLength / tileNum / BUFFER_NUM;
    // ...

}

extern "C" __global__ __aicore__ void add_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z, GM_ADDR workspace, AddCustomTilingData tiling)
{
    KernelAdd op;
    op.Init(x, y, z, tiling.totalLength, tiling.tileNum);
    op.Process();
}

shape非对齐场景

上文描述的都是shape对齐的场景：如下图中的示例，算子的输入shape为（1，2048），支持的数据类型为half类型，可以对齐到一个datablock的大小（32B），也可以平均分配到每个核上（假设使用8个核），每个核上处理256个数，16个datablock。

图1 shape对齐场景
点击放大

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

因为昇腾AI处理器在进行数据搬运和Vector计算时，对于搬运的数据长度和UB首地址都有必须32B对齐的要求，首先待处理数据需要先保证对齐到datablock的大小。该场景下后续搬运和计算的处理细节请参考非对齐处理。下图和代码片段展示了将数据对齐到datablock大小的示例：
图2 对齐到datablock大小
```
constexpr uint32_t SIZE_OF_HALF = 2;
constexpr uint32_t BLOCK_SIZE = 32;
constexpr uint32_t ALIGN_NUM = BLOCK_SIZE / SIZE_OF_HALF;
// shape需要对齐到的datablock,假设原totalLength为1999，向上满足32字节对齐后为2000
uint32_t totalLengthAligned = ((totalLength + ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM;
```

满足datablock对齐后的数据，应尽可能的均分到每个核上。如果无法均分，那么先将可以均分的部分平均分配，剩余的部分分配给部分核，会有部分核多算一个datablock。下图展示了无法均分时将数据进行多核切分的示例。对齐到datablock后为2000个half类型的数据，共125个datablock。125%8结果为15，余数为5，说明：可以均分的部分平均分配，每个核分配到15个datablock；还剩余5个datablock，分配给5个核，所以会有5个核分配到16个datablock，剩余3个核分配到15个datablock。

图3 无法均分到每个核上的示例

基于上文的描述，可以设计如下的Tiling参数：

formerNum：分配到大块的核数
tailNum：分配到小块的核数
formerLength：大块计算的数据量
tailLength：小块计算的数据量
alignNum：一个datablock包含的元素个数

这些Tiling参数的计算方法如下：

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;
...
uint8_t *GenerateTiling()
{
    // shape需要对齐到的datablock,假设原totalLength为1999，向上满足32字节对齐后为2000
    uint32_t totalLengthAligned = ((totalLength + ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM;
    // 把所有的数据尽可能均匀地分配到每个核上
    // 如果不能均分,先将可以均分的部分平均分配，剩余的部分分配给部分核，会有部分核多算一个datablock 
    // 通过模的计算，可以得到多算一个datablock的核的数量，也可以得到剩余核的数量
    // eg：1999 对齐后的总数据量为2000个数，核心数为8，一个datablock包含16个数，那么：
    // datablock的总数：2000 / 16 = 125
    // 有5个核会分到16个datablock：125 % 8 =5，可以称之为大块
    // 有3个核会分到15个datablock：8 - 5 = 3，可以称之为小块
    uint32_t formerNum = (totalLengthAligned / ALIGN_NUM) % BLOCK_DIM;
    uint32_t tailNum = BLOCK_DIM - formerNum;
    // 大块计算的数据量：totalLengthAligned / BLOCK_DIM为每个核上计算的元素个数，formerLength为上述元素个数向上32字节对齐的结果
    uint32_t formerLength = ((totalLengthAligned / BLOCK_DIM + ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM;
    // 小块计算的数据量：totalLengthAligned / BLOCK_DIM为每个核上计算的元素个数，tailLength 为上述元素个数向下32字节对齐的结果
    uint32_t tailLength = (totalLengthAligned / BLOCK_DIM / ALIGN_NUM) * ALIGN_NUM;

...
}

相对应的，在Kernel侧，使用获取到的信息计算得到每个核上的偏移量、每个分块大小的样例如下。

__aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t formerNum, uint32_t tailNum, uint32_t formerLength, uint32_t tailLength, uint32_t alignNum)
{
    if (GetBlockIdx() < formerNum) {
        this->tileLength = formerLength;
        xGm.SetGlobalBuffer((__gm__ half *)x + formerLength * GetBlockIdx(), formerLength);
        yGm.SetGlobalBuffer((__gm__ half *)y + formerLength * GetBlockIdx(), formerLength);
        zGm.SetGlobalBuffer((__gm__ half *)z + formerLength * GetBlockIdx(), formerLength);
    } else {
        this->tileLength = tailLength;
        xGm.SetGlobalBuffer((__gm__ half *)x + formerLength * formerNum + tailLength * (GetBlockIdx() - formerNum), tailLength);
        yGm.SetGlobalBuffer((__gm__ half *)y + formerLength * formerNum + tailLength * (GetBlockIdx() - formerNum), tailLength);
        zGm.SetGlobalBuffer((__gm__ half *)z + formerLength * formerNum + tailLength * (GetBlockIdx() - formerNum), tailLength);
    }
    pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(half));
    pipe.InitBuffer(inQueueY, BUFFER_NUM, this->tileLength * sizeof(half));
    pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->tileLength * sizeof(half));
}

更多场景

动态shape场景

shape非对齐场景