尾核Tiling【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit对于不同shape的输入进行数据切分时可能会发生数据无法平均分配到多个核的情况。例如当算子的输入shape为[1, 1999]使用核数为8数据类型为half时需要计算的数据总量为1 * 1999 * sizeof(half) 3998字节3998字节既不满足32字节对齐也无法平均分配到8个核上。因此该场景下对数据进行多核切分后每个核的计算数据量不同。此种情况下应该尽可能均匀的分配数据所有核上的计算数据量有两种情况将计算量较多的核称为整核计算量较少的核称为尾核。图1数据对齐示意图Tiling实现因为AI处理器在进行数据搬运和Vector计算时对于搬运的数据长度和Unified Buffer首地址都有必须32字节对齐的要求首先待处理数据需要先保证向上对齐到32字节的大小。该场景下后续搬运和计算的处理细节请参考非对齐场景。如下代码片段展示了将数据对齐到datablock大小的示例constexpr uint32_t SIZE_OF_HALF 2; constexpr uint32_t BLOCK_SIZE 32; constexpr uint32_t NUM_BLOCKS 8; constexpr uint32_t ALIGN_NUM BLOCK_SIZE / SIZE_OF_HALF; // shape需要对齐到的32字节假设原totalLength为1999向上满足32字节对齐后为2000 uint32_t totalLengthAligned (totalLength % ALIGN_NUM 0U) ? static_castuint32_t(totalLength) : ((static_castuint32_t(totalLength) ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM;满足32字节对齐后的数据应尽可能的均分到每个核上。如果无法均分那么先将可以均分的部分平均分配剩余的部分分配给部分核会有部分核多算一个datablock。为了保证切分后的数据仍是满足32字节对齐的以ALIGN_NUMALIGN_NUM个数据为32字节为粒度将数据分配到所有核上。在本样例中数据类型为halfALIGN_NUM BLOCK_SIZE / sizeof(half) 16。将对齐后的数据总量按ALIGN_NUM为粒度分成x个数据块x 2000 / 16 125。AI处理器的核数NUM_BLOCKS为8无法将125个数据块均分到8个核上。按照以下步骤将数据块尽可能的均分到每个核上计算x / NUM_BLOCKS 15计算x % NUM_BLOCKS 5。根据上述步骤得出如果每个核上分配15个数据块那么将有5个数据块剩余。将这5个剩余的数据块分配到5个核上这样可以得到5个计算16个数据块的整核和3个计算15个数据块的尾核。下图展示了数据无法均分时多核切分的示例。图2无法均分到每个核上的示例基于上文设计如下的算子Tiling结构体成员formerNum分配到数据量较多的核数即整核的核数。tailNum分配到数据量较少的核数即尾核的核数。formerLength整核计算的数据长度。tailLength尾核计算的数据长度。Tiling参数的计算代码如下constexpr uint32_t NUM_BLOCKS 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; ... void GenerateTilingData(uint8_t* tilingBuf, uint32_t numBlocks) { // shape需要对齐到的datablock,假设原totalLength为1999向上满足32字节对齐后为2000 uint32_t totalLengthAligned (totalLength % ALIGN_NUM 0U) ? static_castuint32_t(totalLength) : ((static_castuint32_t(totalLength) ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM; // 核心数为8一个datablock包含16个数那么datablock的总数2000 / 16 125 // 有5个核会分到16个datablock125 % 8 5可以称之为整核 // 有3个核会分到15个datablock8 - 5 3可以称之为尾核 uint32_t formerNum (totalLengthAligned / ALIGN_NUM) % numBlocks; uint32_t tailNum numBlocks - formerNum; // 整核计算的数据长度totalLengthAligned / NUM_BLOCKS为每个核上计算的元素个数formerLength为上述元素个数向上32字节对齐的结果 uint32_t formerLength static_castuint32_t(((totalLengthAligned numBlocks - 1) / numBlocks ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM; // 尾核计算的数据长度totalLengthAligned / NUM_BLOCKS为每个核上计算的元素个数tailLength为上述元素个数向下32字节对齐的结果 uint32_t tailLength (totalLengthAligned / numBlocks / ALIGN_NUM) * ALIGN_NUM; ... }算子类实现在Kernel侧的Init函数中计算输入在Global Memory上的内存偏移地址时应对整核和尾核加以区分。整核上输入的内存偏移地址计算代码如下xGm.SetGlobalBuffer((__gm__ T*)x formerLength * AscendC::GetBlockIdx(), formerLength);尾核上计算输入的内存偏移地址时需在全部整核的数据长度基础上加上尾核的偏移量代码如下xGm.SetGlobalBuffer((__gm__ T*)x formerLength * formerNum tailLength * (AscendC::GetBlockIdx() - formerNum), tailLength);完整的Init函数实现代码如下__aicore__ inline void Init(__gm__ uint8_t* x, __gm__ uint8_t* y, __gm__ uint8_t* z, AddCustomTilingData tiling, AscendC::TPipe* pipeIn) { pipe pipeIn; if (AscendC::GetBlockIdx() tiling.formerNum) { this-tileLength tiling.formerLength; uint64_t offset tiling.formerLength * AscendC::GetBlockIdx(); xGm.SetGlobalBuffer((__gm__ half *)x offset, tiling.formerLength); yGm.SetGlobalBuffer((__gm__ half *)y offset, tiling.formerLength); zGm.SetGlobalBuffer((__gm__ half *)z offset, tiling.formerLength); } else { this-tileLength tiling.tailLength; uint64_t offset tiling.formerLength * tiling.formerNum tiling.tailLength * (AscendC::GetBlockIdx() - tiling.formerNum); xGm.SetGlobalBuffer((__gm__ half *)x offset, tiling.tailLength); yGm.SetGlobalBuffer((__gm__ half *)y offset, tiling.tailLength); zGm.SetGlobalBuffer((__gm__ half *)z offset, tiling.tailLength); } pipe-InitBuffer(inQueueX, 1, this-tileLength * sizeof(half)); pipe-InitBuffer(inQueueY, 1, this-tileLength * sizeof(half)); pipe-InitBuffer(outQueueZ, 1, this-tileLength * sizeof(half)); }其余实现与多核Tiling中的实现一致这里不重复进行说明。【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考