CommonMatmul【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass1 模板说明在泛化Matmul中存在两个CommonMatmul模板一个是纯CUBE类型名为CommonMatmulKernel一个是MIX类型名为PaddingCommonMatmulKernel如果需要用到AIV进行数据格式转换参考2.3节则会进入PaddingCommonMatmulKernel之所以区分这两个模板是因为编译一个Kernel的时候要么指定这个Kernel是MIX类型的要么指定这个Kernel是纯CUBE或纯AIV的而MIX算子由于需要同时启动AIC和AIV启动开销会比纯CUBE或纯AIV大尤其在小shape场景下性能影响显著如果不需要用到AIV就不需要启动AIV进入纯Cube类型的CommonMatmulKernel会有更好的性能。CommonMatmul的典型特征是使用L1上的基本块直接切分矩阵A、B、C先将C矩阵划分为若干个大小为L1TileM x L1TileN的基本任务块然后以L1TileM x L1TileK为粒度读取A矩阵到L1以L1TileK x L1TileN为粒度读取B矩阵到L1。每个任务块的完整计算结果在一个AI Core上产生。为了避免不必要的L1-L0的重复数据搬运在CommonMatmul中L0TileML1TileML0TileNL1TileNL0TileK在L0的空间约束下取最大值。2 优化点此处介绍的优化点是00_basic_matmul里面没有但是CommonMatmul中有的优化至于00_basic_matmul中已有的基础优化在此不做介绍。00_basic_matmul用到的优化点为CommonMatmul的子集。2.1 PreloadCommonMatmul中的Preload实现伪代码如下... uint32_t kTileCount CeilDiv(actualShape.k(), l1TileShape.k()); ... // block层的k循环 for (uint32_t kLoopIdx 0; kLoopIdx kTileCount; kLoopIdx) { // 如果是当前核心计算的第一个C Block块而且kLoopIdx 0的时候需要加载当前计算的Tile块的数据。 // 如果当前计算的不是第一个C Block块或者kLoopIdx 0的时候就不需要加载当前计算的Tile块数据因为数据已经在上一轮循环加载好了。 if (isFirstBlock kLoopIdx 0) { // 根据kLoopIdx计算需要载入GM块的地址 copyGmToL1A; copyGmToL1B; } // 块内的Preload预先加载当前计算的C Block的下一个A Tile块和B Tile块。 if ((kLoopIdx ! kTileCount - 1)) { // 根据kLoopIdx 1计算需要载入GM块的地址因为载入的是下一个需要计算的块的数据。 copyGmToL1A; copyGmToL1B; } // 块间的Preload预先加载下一个需要计算的C Block对应的第一个A Tile块和B Tile块。 if ((kLoopIdx ! kTileCount - 1) hasNextBlock) { // 加载下一个需要计算的C Block的对应第一个A Tile块和B Tile块 copyGmToL1A; copyGmToL1B; } ... tileMmad; ... } ... copyL0CToGm; ...实际代码实现见block_mmad_dynamic_common.h。isFirstBlock是否为当前核心计算的第一个C Block。如果是则为true。hasNextBlock当前核心是否存在下一个需要计算的C Block如果存在则为true。isFirstBlock和hasNextBlock参数由kernel层传到Block层。Preload的核心思想为在计算当前Tile的Matmul前预先加载下一个需要计算的Tile块的数据这样就做到了搬运指令的提前发射可减少MTE2流水的空泡。关于Preload更细节的描述可以参考矩阵乘模板总结-流水优化(Preload)。2.2 ShuffleK多核心同时读取同一片GM数据的时候存在数据读取冲突导致读取带宽降低。通常所有核心都从K方向第一个分块开始计算为了缓解上面的问题让不同的核心从不同的分块开始计算这样各个核心不再同时读取同一片GM数据。图中C矩阵中的数字为核心序号AB矩阵中的数字为任务块序号左图为通常的计算方式图中2号核心和3号核心同时按0 1 2 3的顺序访问A矩阵的同一片GM数据。右图为采用ShuffleK优化后的计算方式右图中2号核心按2 3 0 1的顺序访问A矩阵基本块而3号核心按3 0 1 2的顺序访问A矩阵基本块在时间上错开从而避免同地址访问冲突。图中的分块计算顺序采用GemmIdentityBlockSwizzle2,1请参考swizzle_explanation2.3 Padding在A2或A3上当A或B矩阵为NDRowMajor或ColumnMajor格式时如果矩阵的Stride为非512B对齐则ND2NZ搬运接口的带宽会显著下降。为了规避这个问题采用AIV提前对A或B矩阵进行数据格式转换或数据填充目的是让GM2L1搬运时候避免用非512B对齐的Stride访问GM数据。2.3.1 矩阵A或B的Padding模式说明当前泛化Matmul支持三种Padding方式在Padding_matmul.hpp中定义了枚举值enum class PaddingTag { NO_PADDING, PADDING_ND, PADDING_BLOCK_ND, PADDING_NZ };以上的Padding操作都是全局的操作即针对整个A矩阵或B矩阵需要开辟A矩阵和B矩阵对齐后同等大小的workspace。由AIV进行Padding处理AIC等待处理完成后开始Matmul计算。PADDING_NDPADDING_ND只是简单把Stride方向填充对齐到512B对齐例如图中的512 x 511的shapeStride为511经过PADDING_ND后Stride变为512对齐到了512B对于half类型也就是要求Stride为256倍数。PADDING_BLOCK_NDPADDING_BLOCK_ND会重排矩阵数据以L1Tile块的粒度对数据进行重组如果原矩阵为RowMajor如图所示重组后每个Tile块内RowMajor排布然后Tile块间也是RowMajor排布是一个嵌套的数据格式。这样做的好处是不管原矩阵的Stride为多少通过这个变换后每个Tile块内的Stride都变为L1TileK一般是512B对齐的规避了Stride过大超过65536或者非对齐造成的带宽劣化。PADDING_NZPADDING_NZ直接将RowMajor的矩阵转换为zN的格式如果是ColumnMajor的话就是nZ这样GM2L1读取的时候不再需要ND2NZ的随路转换也就规避了ND2NZ存在的所有劣化问题。PADDING_NZ有两种实现分别用于不同的场景矩阵的Stride方向矩阵的内轴较大的时候AIV一次读取16行数据repeat读取虽然repeat的stride非对齐但是每次repeat读取的数据量大非对齐对带宽影响不大在UB上通过Copy指令将数据重排为zN然后写出到workspace。矩阵的Stride方向矩阵的内轴较小的时候一般是小于96如果按第一种方案读取数据AIV的读取效率也很低。这时候将非对齐的矩阵数据连续读取到UB采用TransDataTo5HD指令经过两次转置操作完成数据的Padding同时顺带完成ND数据转NZ。由于PADDING_NZ有较优的总体性能所以泛化Matmul中实际采用的是PADDING_NZ至于其他两种Padding方式在某些特定的Shape上可能有比PADDING_NZ更好的性能可以多尝试进行调优。2.3.2 C矩阵的Padding模式说明由于基本块的计算结果在L0C上是zN排布的L0C到UB也需要经过NZ2ND的数据转换如果C矩阵的Stride非512B对齐同样也存在带宽明显劣化的问题。为了规避上面问题当前CommonMatmul的做法是申请与C矩阵同样大小对齐后的workspace当L0C上结果计算完成后Fixpipe按对齐的Stride先写出到workspace然后用AIV进行去Padding操作将最终结果写回到GM C。当前在CommonMatmul中是否对C矩阵进行Padding的判断条件如下PaddingTag paddingTagC PaddingTag::PADDING_NONE; // 要求C矩阵不能太小且要求非256B对齐 if (static_castsize_t(m) * n 2048 * 2048 n 256 (n % 128 ! 0)) { size_t totalDataSize static_castsize_t(m) * k * CeilDiv(n, n1) * 2 static_castsize_t(k) * n * CeilDiv(m, m1) * 2 static_castsize_t(m) * n * 2; // 要求读写总数据量小于L2 Cache大小 if (totalDataSize 192 * 1024 * 1024) { // L2 cache size paddingTagC PaddingTag::PADDING_ND; } }如果总数据量大小超过了L2 Cache大小那么AIV进行去Padding的时候就可能从GM读取数据这时候去Padding开销就很大以至于开销大于收益。2.3.3 Padding建模决定A矩阵或B矩阵是否需要Padding是否Padding影响Matmul的性能Padding有额外开销如果Padding后带来的带宽收益无法抵消Padding的额外开销那么Padding就会是负收益。Padding的开销主要来自下面两个方面AIV的启动开销只要启动了mix kernel就会有1~5us左右的额外开销。Padding操作的开销即AIV的读写计算耗时。在实际的测试中有下面的规律一般在非对齐的时候需要进行Padding优化非对齐矩阵的读取带宽但是如果该非对齐矩阵在Matmul计算过程中没有重复的数据读取就不应该进行Padding。如果矩阵的Stride方向长度小于64nd2nz的随路转换读取带宽会非常差此时应该进行Padding即使该矩阵没有重复的数据读取。如果矩阵过于小即矩阵的rows和cols都小于某个数此时也不需要进行Padding一方面小矩阵可以常驻L1不用重复读取一方面Padding的收益无法抵消额外开销。在矩阵重复读取量不大的情况下如果矩阵的Stride方向为256B对齐不Padding会更好。但是很难通过上面的规律总结出一个确定的Padding策略有很多的阈值难以确定。所以需要对Padding过程进行建模。设置以下量B_aivAIV的单核读取带宽。B_aic512AIC nd2nz stride 512B对齐情况下的单核读取带宽。B_aicunalign AIC stride非512B对齐情况下的单核读取带宽。T_headcost启动mix kernel带来的额外开销。M,N,K,m1,n1,k1矩阵shape和L1 Tile块大小。CCUBE核心数量VVECTOR核心数量。单核最大的计算轮次为R_max RoundUp(CeilDiv(M, m1) * CeilDiv(N, n1) / C)单核AIC A和B矩阵最大的数据读取量分别为CRead_a R_max * (m1 * K)CRead_b R_max * (n1 * K)假设A矩阵和B矩阵都进行Padding那么AIV的单核读取数据量为VRead_a M * K / VVRead_b K * N / V那么AB都Padding的AIC和AIV总读取耗时为T_11 (CRead_a CRead_b) / B_aic512 (VRead_a VRead_b) / B_aiv T_headcost如果不进行Padding那么总读取耗时为T_00 (CRead_a CRead_b) / B_aicunalign如果Padding A矩阵但是不Padding B矩阵总读取耗时为T_10 CRead_a / B_aic512 VRead_a / B_aiv CRead_b / B_aicunalign T_headcost如果Padding B矩阵但是不Padding A矩阵总读取耗时为T_01 CRead_b / B_aic512 VRead_b / B_aiv CRead_a / B_aicunalign T_headcost比较T00、T01、T10、T11这四个时间哪个时间最短就采用哪种Padding策略。B_aiv、B_aic512、T_headcost这几个值在硬件相同时可以简化认为是固定值可通过实测获得。而B_aicunalign的值和nd2nz的参数有关nValue、dValue和srcDValue需要通过实测数据进行曲线拟合得到计算公式。以上逻辑为Padding建模的简化过程具体实现有更细节的考虑参考select_kernel_bf16.h其中的多项式拟合公式仅适用于A2、A3如果是其他型号需要自测数据拟合曲线。2.4 特殊场景的读取优化2.4.1 场景一ND2NZ可以使用普通间隔搬运的DataCopy取代即for循环每次搬运一行每一行调用DataCopy Repeat多次进行搬运代码如下// 使用普通间隔搬运的DataCopy取代ND2NZ以half为例 for (int i 0; i nValue; i) { AscendC::DataCopyParams dataCopyParams { (dValue 15) / 16, // repeat times 1, // 一次搬运的数据量以32B为单位 0, // Gm上的Stride以32B为单位 (256 - 16) / 16 // L1上的Stride以32B为单位 }; AscendC::DataCopy(buffer[i * 16], gmSrc[i * srcDValue], dataCopyParams); }在M很小的时候通常是M 8ND2NZ的随路转换的搬运效率没有上面的逐行拷贝效率高所以CommonMatmul在这种场景下采用上面的间隔拷贝的DataCopy替换ND2NZ以获取更高的搬运效率。2.4.2 场景二如图所示当K16时候矩阵数据在GM上的数据排布与在L1上的数据排布相同此时不需要使用随路ND2NZ接口进行搬运直接进行连续的数据拷贝即可。直接的连续拷贝肯定比ND2NZ的随路转换带宽更高。2.4.3 场景三当矩阵A为ColumnMajor且M1的时候此时GM2L1读取的时候每行只有一个元素读取效率会非常低此时将A矩阵当作一个1 x K的RowMajor矩阵进行计算因为A矩阵相当于是一个向量既可以是行优先也可以是列优先两者是等价的这样读取A矩阵的时候A矩阵就是一个行向量读取效率就会高很多。同理当矩阵B为RowMajor且N1的时候此时将B矩阵当作一个K x 1的ColumnMajor矩阵进行计算。此处的逻辑参考select_kernel_bf16.h。至于K1的时候需要交给特殊的模板采用AIV计算处理。【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考