cann/catlass多核切K矩阵乘法

MultiCoreSplitkMatmul

【免费下载链接】catlass本项目是CANN的算子模板库，提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass

1 模板说明

由于硬件指令限制，m1和n1需为16的倍数（m1和n1为L1Tile参数），当C矩阵较小的时候，例如C矩阵大小为16x16，那么此时只能划分出一个基本任务块，如果此时K很大，那么计算效率会很低，因为计算量大，但是却只有一个核心参与计算，造成了严重的算力和带宽的浪费。且m1和n1太小的时候可能会导致搬运指令的效率较低，使用更大的m1和n1可能可以得到更好的效果，但是更大的m1和n1会导致切分的任务块减少，利用的核心数少，从而存在更多的资源浪费。

多核切K可以解决上面问题，不仅沿M、N方向分核，同时K方向也参与分核，将K方向的计算任务分布到多个核心上，当M和N较小且K较大时，相比普通的模板会有很大优势。例如M=16，N=16，K=10240的Matmul，m1=16，n1=16，k1=256，此时如果只对M和N切分，那么只能有一个任务块，但是K方向有40个任务块，假设核心数量为3，那么可以将K方向分为3份，每个核心计算其中一份，最后再用Vector Core对部分和进行累加，从而利用上所有的算力资源。

具体的步骤如图所示：

下面对此模板的要点进行详细的说明：

1. 对于K方向的切分，以基本块为粒度进行切分，例如M=16，N=16，K=10240的Matmul，m1=16，n1=16，k1=256，K方向切分出40个任务块，假设核心数量为3，那么将40个任务块分为3份，分别为14，13，13（均匀划分，如果有余数，则将余数均分到前面的核心）。这时0号核心计算0~14*256的K段，1号核心计算14*256~27*256的K段，2号核心计算27*256~40*256的K段。
2. 需要申请workspace，将K方向分了多少段，就要申请多少份workspace，每个workspace大小为C_SIZE*sizeof(ElementAccumulator)，需要的workspace较大，所以此模板不适合C矩阵大的场景。
3. Cube完成部分和的计算后，需要进行全核心同步，保证所有的部分和都计算完成，此时Vector可以开始进行累加操作，为保证计算结果的确定性，Vector将按顺序进行累加，即一次读取完所有的workspace的部分和结果到UB中（由于UB大小限制，一次只能读取每个workspace一小部分的部分和），然后按顺序循环将后面的部分和累加到第一个部分和上，得到部分最终结果，如果C矩阵的输出类型是float16，那么需要将float类型的结果在UB上Cast为float16，最后将数据写回到GM C上。
4. 为了更高的数据写出和数据读取效率，Vector累加部分和的时候，按元素进行任务划分，例如16x16的C矩阵，假设有6个Vector，那么每个Vector分到16*16/6=42个元素，进一步考虑到Vector的指令效率，设置一个最小划分块，元素数量为256/sizeof(ElementAccumulator)，如果ElementAccumulator为32位类型，则最少一个核心划分得到64个元素，16x16的C矩阵被分成4份，由4个Vector进行处理，另外两个Vector空闲。这样的划分将C矩阵视为一个连续的一维数组，每个Vector处理其中连续的一段，这样Vector读取和写出的数据就是连续的一段，指令效率更高。这样也导致了此模板不适用于C矩阵非连续的场景，例如C矩阵的Shape为16x16，每行的Stride为17。
5. MultiCoreSplitkMatmul模板中用到了Preload、ShuffleK、Padding以及特殊场景的读取优化等CommonMatmul中已有的优化点。

2 适用场景

1. C矩阵较小且K较大，核心没用满的时候。
2. 要求C矩阵在GM上连续存储。

【免费下载链接】catlass本项目是CANN的算子模板库，提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass