CANN/ge Auto Tiling自动切分

Auto Tiling【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge在Schedule阶段系统会根据前端输入的AscIR图生成多种优化后的切分图Auto Tiling需要评估这些图的最优切分方案及其对应的Kernel性能从而选出最优模板。Tiling策略决定了Kernel执行时需要使用的核的数量每次从GM中搬运多少数据到UB内在UB内需要循环多少次每次在UB内搬运多少次到GM中等这些因素共同影响着Kernel的执行性能。Auto Tiling的目标是找到能够实现Kernel最佳执行性能的Tiling策略。核心流程图 1Auto Tiling核心流程 Kernel的执行逻辑会在AscIR中表达Auto Tiling会根据AscIR图的表达提取关键信息包括LocalBuffer占用建模Auto Tiling求解需确保每一级LocalBuffer的占用在硬件允许的范围内比如Kernel申请的TQue/TBuf及临时Buf申请大小之和不能超过硬件的UB大小限制AscIR会表达出每个Tensor的Location是在GM/UB上Tensor间的复用关系自动Tiling根据这些信息将各级LocalBuffer的约束进行符号化表达。关于TQue/TBuf详细介绍请参见《Ascend C算子开发指南》。耗时公式建模Auto Tiling会对各个API进行性能建模通过符号化的形式表达这些API在各个流水上的性能根据AscIR表达的循环轴来确定表达API的调用次数从而推导出该AscIR的所有API在各个流水上的总耗时此处假设瓶颈流水线的执行可以较好地掩盖非瓶颈流水线的执行因此任务执行的总耗时主要体现在瓶颈流水线的时间上如下图下图表达了三个pipe流水瓶颈流水为AIV_MTE2AIV_MTE2在执行时会掩盖AIV_MTE3和AIV_VEC的执行任务执行的总耗时体现在AIV_MTE2的执行耗时上。图 2pipe流水示意图 求解器求解TilingAuto Tiling以算子实现过程中的存储占用不超过NPU各级物理存储大小为约束条件以最小化瓶颈执行单元的耗时为优化目标通过建模数据优化模型进行求解。例如通过启发式求解方法从初始解开始在满足约束条件的可行域内根据性能公式建模的梯度下降方向搜索满足内存要求的解空间直到公式建模的值达到最小从而退出寻优流程返回最优Tiling。关键技术性能公式建模目标性能公式的仿真精度决定了模板选择的准确性及启发式求解的准确性因此需要对API进行较为精确的性能建模。实现针对AscendC稳定的基础API自动Tiling模块会根据不同的输入采集性能数据得到性能与输入的关系并建立性能模型通过符号化的技术表达出来。针对调用AscendC易变的API自动Tiling模块会基于API调用的逻辑计算其调用入参和调用次数生成该API的完整性能模型从而得到相对准确的性能建模。轴排序求解器目标基于性能公式梯度下降的求解方式高度依赖性能公式的精确建模并且Tiling求解时间存在不确定性容易出现host bound问题而API本身对轴的切分有特定的偏好因此设计了一种轴排序求解器作为迭代求解的替代方案。实现首先需要基于API来确定切分轴的优先级顺序如下父轴优先级高于子轴。规约化类轴高于非规约化类轴。广播轴高于非广播轴。非最内轴高于最内轴。其次再分成两个部分切分包括核内Tiling按照轴排序的逆序依次遍历优先将变量调整至最大值判断是否符合硬件约束条件若不满足则通过二分法调整该变量直到符合硬件约束条件为止随后调整下一个核内Tiling变量直至所有的变量均满足硬件约束条件比如s1tt2、s1tt、s1t、s2t是Tiling相关轴其切分流程如下优先遍历s2t调至最大值1024符合硬件约束条件然后遍历s1t调至最大值256符合硬件约束条件然后依次调整下个变量s1tts1tt2直到所有的变量均满足硬件约束条件。图 3核内Tiling示意图 多核Tiling识别与多核相关的变量按从大到小的顺序遍历这些变量找到更大核数占用的记录若超出物理核数以Atlas A2 训练系列产品/Atlas A2 推理系列产品为例NPU核数为48则返回。如下图所示bngs1T是多核切分轴其切分流程如下。选择策略为核数占用不同时优先选择占用核数更大的记录根据上述策略最终选定的占用核数为47。图 4多核Tiling示意图 【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考