cann/catlass多核切K矩阵乘法

cann/catlass多核切K矩阵乘法 MultiCoreSplitkMatmul【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass1 模板说明由于硬件指令限制m1和n1需为16的倍数m1和n1为L1Tile参数当C矩阵较小的时候例如C矩阵大小为16x16那么此时只能划分出一个基本任务块如果此时K很大那么计算效率会很低因为计算量大但是却只有一个核心参与计算造成了严重的算力和带宽的浪费。且m1和n1太小的时候可能会导致搬运指令的效率较低使用更大的m1和n1可能可以得到更好的效果但是更大的m1和n1会导致切分的任务块减少利用的核心数少从而存在更多的资源浪费。多核切K可以解决上面问题不仅沿M、N方向分核同时K方向也参与分核将K方向的计算任务分布到多个核心上当M和N较小且K较大时相比普通的模板会有很大优势。例如M16N16K10240的Matmulm116n116k1256此时如果只对M和N切分那么只能有一个任务块但是K方向有40个任务块假设核心数量为3那么可以将K方向分为3份每个核心计算其中一份最后再用Vector Core对部分和进行累加从而利用上所有的算力资源。具体的步骤如图所示下面对此模板的要点进行详细的说明对于K方向的切分以基本块为粒度进行切分例如M16N16K10240的Matmulm116n116k1256K方向切分出40个任务块假设核心数量为3那么将40个任务块分为3份分别为141313均匀划分如果有余数则将余数均分到前面的核心。这时0号核心计算0~14*256的K段1号核心计算14*256~27*256的K段2号核心计算27*256~40*256的K段。需要申请workspace将K方向分了多少段就要申请多少份workspace每个workspace大小为C_SIZE*sizeof(ElementAccumulator)需要的workspace较大所以此模板不适合C矩阵大的场景。Cube完成部分和的计算后需要进行全核心同步保证所有的部分和都计算完成此时Vector可以开始进行累加操作为保证计算结果的确定性Vector将按顺序进行累加即一次读取完所有的workspace的部分和结果到UB中由于UB大小限制一次只能读取每个workspace一小部分的部分和然后按顺序循环将后面的部分和累加到第一个部分和上得到部分最终结果如果C矩阵的输出类型是float16那么需要将float类型的结果在UB上Cast为float16最后将数据写回到GM C上。为了更高的数据写出和数据读取效率Vector累加部分和的时候按元素进行任务划分例如16x16的C矩阵假设有6个Vector那么每个Vector分到16*16/642个元素进一步考虑到Vector的指令效率设置一个最小划分块元素数量为256/sizeof(ElementAccumulator)如果ElementAccumulator为32位类型则最少一个核心划分得到64个元素16x16的C矩阵被分成4份由4个Vector进行处理另外两个Vector空闲。这样的划分将C矩阵视为一个连续的一维数组每个Vector处理其中连续的一段这样Vector读取和写出的数据就是连续的一段指令效率更高。这样也导致了此模板不适用于C矩阵非连续的场景例如C矩阵的Shape为16x16每行的Stride为17。MultiCoreSplitkMatmul模板中用到了Preload、ShuffleK、Padding以及特殊场景的读取优化等CommonMatmul中已有的优化点。2 适用场景C矩阵较小且K较大核心没用满的时候。要求C矩阵在GM上连续存储。【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考