尾轮负载均衡特性介绍【免费下载链接】cann-samples算子领域高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples1. 原理介绍1.1 背景在进行算子多核运算时通常会遇到最后一轮迭代中剩余任务块block数量小于可用计算核数量的情况。此时仅部分计算核参与计算其余计算核处于闲置状态导致整体计算效率下降。若能将这些剩余任务进一步拆分使其均匀分配到所有计算核上并行执行则可有效提升计算资源利用率显著改善整体性能。这一数据拆分与重分配的过程即为尾轮负载均衡。1.2 原理将末轮未完全分配的基本块进行重切分重新计算所需的总核数并重新分配计算核心使其均匀分布至各核从而充分发挥计算核算力。原理图如下如图所示重新分配基本快到多核后端到端的计算耗时减少计算效率提升。2. 实践使用尾轮负载均衡策略来优化matmul计算性能2.1 代码下面以典型添加了SWAT特性的MatMul算子计算代码为例演示数据切分及均匀分发至各核心的具体实现过程2.1.1 计算尾轮基本块需要拆分的数量尾轮负载均衡首先需要确定尾块拆分的大小即明确在M维度和N维度上分别需要切分的块数。以下函数基于当前可用的AI核心数量计算最优的尾块切分方案// 计算尾块拆分维度 __aicore__ inline void CalcTailBasicBlock(uint64_t mTileNum, uint64_t nTileNum, uint64_t aicNum, uint64_t tailMCnt, uint64_t tailNCnt) { uint64_t mnCnt mTileNum * nTileNum; uint64_t tailCnt mnCnt - aicNum * (CeilDiv(mnCnt, aicNum) - 1); tailMCnt 1UL; tailNCnt 1UL; if (tailCnt ! 0UL) { while ((tailMCnt 1UL) * tailNCnt * tailCnt aicNum) { tailMCnt 1UL; if (tailMCnt * (tailNCnt 1UL) * tailCnt aicNum) { tailNCnt 1UL; } } } }2.1.2 多核任务重分配尾块拆分后需要重新计算总任务块数并调整多核分配策略确保各核心负载均衡。具体实现如下// 重新计算总块数原始块数 尾块拆分后新增的块数 tileNum tileNum (tailCnt - 1) * perTailCnt; // Multi-core tile processing loop - distribute tiles across available cores for (uint64_t tileIdx curBlockIdx; tileIdx tileNum; tileIdx blockNum) { // 多核任务分配循环 - 将任务块均匀分发至各核心 if (tileIdx / blockNum (perCoreBlockNum - 1) tailCnt 1) { tileIdx (perCoreBlockNum - 1) * blockNum curBlockIdx / tailCnt; } // 单核计算逻辑 ... }2.1.3 尾块坐标重映射当处理到末轮拆分的尾块时需要重新计算当前核心负责的子块在M维度和N维度上的起始坐标与尺寸确保数据切分的正确性。具体实现如下// 判断是否为尾块拆分场景最后一个核心且存在尾块拆分 if (tileIdx / blockNum (perCoreBlockNum - 1) tailCnt 1) { // 计算拆分后每个子块在M维度和N维度上的尺寸 int64_t splitBlkM tool::CeilDiv(curM, tailMCnt); int64_t splitBlkN tool::CeilDiv(curN, tailNCnt); // 计算当前核心在尾块拆分中的索引位置 int64_t mSplitIdx (curBlockIdx % tailCnt) % tailMCnt; int64_t nSplitIdx (curBlockIdx % tailCnt) / tailMCnt; // 计算子块在M维度和N维度上的起始偏移 int64_t mSplitOffset mSplitIdx * splitBlkM; int64_t nSplitOffset nSplitIdx * splitBlkN; //跳过那些超出原始维度范围的无效子块 if (mSplitOffset curM || nSplitOffset curN) { continue; } // 更新当前子块的实际尺寸边界处理 curM (curM - mSplitOffset) splitBlkM ? (curM - mSplitOffset) : splitBlkM; curN (curN - nSplitOffset) splitBlkN ? (curN - nSplitOffset) : splitBlkN; // 重新设置张量GM地址指向当前子块对应的数据区域 tensorAGmBlock tensorAgm(AscendC::Te::MakeCoord(mTileIdx * baseM mSplitOffset, 0L), AscendC::Te::MakeShape(curM, k)); tensorBGmBlock tensorBgm(AscendC::Te::MakeCoord(0L, nTileIdx * baseN nSplitOffset), AscendC::Te::MakeShape(k, curN)); tensorCGmBlock tensorCgm(AscendC::Te::MakeCoord(mTileIdx * baseM mSplitOffset, nTileIdx * baseN nSplitOffset), AscendC::Te::MakeShape(curM, curN)); // 根据实际输出尺寸重新设置L0C布局 layoutL0C AscendC::Te::MakeL0CLayout(curM, curN); // L0C layout for output tensorL0C AscendC::Te::MakeTensor(AscendC::Te::MakeL0CmemPtrfloat(l0cOffset), layoutL0C); }关键改动点:尾块拆分维度计算采用贪心策略在确保拆分后子块总数不超过可用核心数的前提下尽可能扩大M、N维度的拆分数量充分利用闲置核心。多核任务重分配重构多核任务分配循环通过动态更新总块数并增加尾块场景的特殊映射将末轮任务均匀分发至所有可用核心。尾块坐标重映射新增子块索引计算逻辑根据拆分后的网格位置重新确定当前核心负责的M、N维度起始偏移与数据尺寸确保张量地址映射的正确性。3. 性能结果对比3.1 case前后性能以基础MatMul算子为例在相同输入规模M2560, K1024, N2560下进行性能测试通过Profiling工具采集硬件流水线执行状态。使用尾轮负载均衡策略优化后可以看到经过尾轮负载均衡后尾轮的计算时间显著缩短整体计算效率得到提升。4. 结论适用场景存在尾轮分配不均匀矩阵维度不是基本块大小的整数倍导致末轮存在不完整分配。尾轮负载均衡通过将末轮剩余任务块进一步拆分并均匀分配至所有计算核心有效消除多核计算中的负载不均问题在尾块占比较高且核心数较多的场景下可显著提升计算资源利用率和算子执行性能。5. 编译 执行编译样例从项目根目录启动构建参考项目README.md在仓库根目录下完成编译和安装后进入当前样例目录cmake -S . -B build -DNPU_ARCHdav-3510 cmake --build build --parallel cmake --install build --prefix ./build_out cd ./build_out/1_Features/system_optimization/tail_rebalance/如需单独编译当前样例可使用以下指令cmake --build build --target tail_rebalance cp ./Samples/1_Features/system_optimization/tail_rebalance/scripts/profile_matmul.py ./build/Samples/1_Features/system_optimization/tail_rebalance/ cd ./build/Samples/1_Features/system_optimization/tail_rebalance/运行样例使用可执行文件直接执行算子用例需要指定矩阵乘维度并随机生成输入数据。./tail_rebalance 1024 2048 4096打印如下执行结果证明样例执行成功。matmul run successfully!如果存在精度问题则会打印错误数据并显示如下结果。matmul run failed!测试性能 运行性能测试脚本指定矩阵乘法的维度后执行。python3 profile_matmul.py 1024 2048 4096打印如下执行结果证明样例性能测试成功。[Profile Breakdowm] --------------------------------------------------------------------------------------------- | candidate | kernel(us) | mac(us) | scalar(us) | mte1(us) | mte2(us) | fixpipe(us) | icache_miss(%) | | matmul | 82.135 | 41.781 | 1.863 | 10.539 | 33.148 | 2.132 | 2.500 | ---------------------------------------------------------------------------------------------与相同输入规模下的基础 matmul 算子相比[Profile Breakdowm] --------------------------------------------------------------------------------------------- | candidate | kernel(us) | mac(us) | scalar(us) | mte1(us) | mte2(us) | fixpipe(us) | icache_miss(%) | | matmul | 86.870 | 43.804 | 1.850 | 12.997 | 51.857 | 2.970 | 2.200 | ---- 可以看到由于尾轮的计算效率提升整体计算时间缩短性能有所提升。6. 支持架构NPU ARCH 3510【免费下载链接】cann-samples算子领域高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考