矩阵乘模板总结【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass当前库上examples内包含多种矩阵乘的样例模板其来源是不同的matmul理论模板与工程实践中发现的工程优化点的组合。在充分理解了各个理论模板和工程优化后开发者可以基于问题场景选择适合的样例模板、甚至进一步自行组合出库上没有的新的样例模板来达成矩阵乘的高性能优化。注意本文档仅总结矩阵乘方案相关的样例其他涉及量化、groupMatmul、后处理等的矩阵乘不在此处总结。样例模板清单00_basic_matmul理论模板Common模板工程优化流水优化Multi Buffer关键交付件host00_basic_matmulkernelbasic_matmul.hppblockMmadblock_mmad_pingpong.hppdispatchPolicyMmadAtlasA2Pingpong04_padding_matmul理论模板Common模板工程优化流水优化Multi Buffer读取带宽优化padding- PaddingMatrixND关键交付件host04_padding_matmulkernelpadding_matmul.hppblockMmadblock_mmad_pingpong.hppdispatchPolicyMmadAtlasA2Pingpong06_optimized_matmul理论模板Common模板工程优化流水优化Multi Buffer流水优化Preload读取带宽优化Padding- PaddingMatrixNZ读取带宽优化ShuffleK读取带宽优化小M下指令替换需要修改样例使能关键交付件host06_optimized_matmulkerneloptimized_matmul.hppPadding前处理组件padding_matmul.hppblockMmadblock_mmad_preload.hppdispatchPolicyMmadAtlasA2Preload⚠️ 注意即使没有使用PaddingMatrixNZ前处理依然会产生MIX算子编译和CV1:0启动的开销比仅AIC启动的开销大09_splitk_matmul理论模板多核切K模板 MultiCoreSplitK工程优化流水优化Multi Buffer关键交付件host09_splitk_matmulkernelsplitk_matmul.hppblockMmadblock_mmad_pingpong.hppdispatchPolicyMmadAtlasA2Pingpong21_basic_matmul_preload_zN此样例主要承载NZ排布输入的适配方法也可换成ND排布输入无MIX算子编译启动的开销理论模板Common模板工程优化流水优化Multi Buffer流水优化Preload读取带宽优化ShuffleK关键交付件host21_basic_matmul_preload_zNkernelbasic_matmul_preload.hppblockMmadblock_mmad_preload.hppdispatchPolicyMmadAtlasA2Preload22_padding_splitk_matmul理论模板多核切K模板 MultiCoreSplitK工程优化流水优化Multi Buffer读取带宽优化padding- PaddingMatrixND关键交付件host22_padding_splitk_matmulkernelpadding_splitk_matmul.hppPadding前处理组件padding_matmul.hppSplitkReduceAdd后处理组件splitk_matmul.hppblockMmadblock_mmad_pingpong.hppdispatchPolicyMmadAtlasA2Pingpong25_matmul_full_loadA此样例及相关组件仅适配了A矩阵全载实现需要实现B矩阵全载可参考关键交付件自行开发理论模板Common模板工程优化流水优化Multi Buffer(全载的A矩阵在L1上不使用多buffer)读取带宽优化L1常驻关键交付件host25_matmul_full_loadAkernelmatmul_full_loadA.hppblockMmadblock_mmad_pingpong_full_loadA.hppdispatchPolicyMmadAtlasA2FullLoadABlockSchedulerGemmIdentityBlockSwizzleL1FullLoad31_small_matmul理论模板Common模板工程优化流水优化Multi BufferScalar开销消减关键交付件host31_small_matmulkernelsmall_matmul.hppblockMmadblock_mmad_small.hppdispatchPolicyMmadAtlasA2SmallBlockSchedulerkernel内实际不使用34_single_core_splitk_matmul理论模板单核切K模板 SingleCoreSplitK工程优化流水优化Multi Buffer读取带宽优化Padding- PaddingMatrixNZ写出带宽优化关键交付件host34_single_core_splitk_matmulkernelsingle_core_slicek_matmul.hppPadding前处理组件和RemovePaddingNDAndCast后处理组件padding_matmul.hppblockMmadblock_mmad_single_core_splitk.hppdispatchPolicyMmadAtlasA2SingleCoreSplitkBlockSchedulerSingleCoreSplitkGemmIdentityBlockSwizzle理论模板清单Common模板Tiling建模如图展示一个常规的fp16的矩阵运算L0C上按照fp32累加定义相关参数问题shape$M$$N$$K$搬入L1Cache时的TileShape$m_1$$n_1$$k_1$搬入L0A/L0B、搬出L0C时的TileShape$m_0$$n_0$$k_0$采用 $M$、$N$ 方向分核按照$m_1$、$n_1$切分产生$\frac{MN}{m_1n_1}$个基本任务块、分配给AIC核完成搬运和计算每个基本任务块需要搬运$m_1KKn_1$的数据、计算得到$m_1n_1$的结果并搬出。由此产生约束$m_1k_1L1Stage_A n_1k_1L1Stage_B L1Size / 2Byte$$m_0k_0*L0AStage L0ASize / 2Byte$$n_0k_0*L0BStage L0BSize / 2Byte$$m_0n_0*L0CStage L0CSize / 4Byte$$m_0 m_1$$n_0 n_1$读取数据量每个基本任务块需要搬运$m_1KKn_1$的数据总读取量为$2Byte * [m_1KKn_1] * \frac{MN}{m_1n_1} 2Byte * MNK * [\frac{1}{m_1}\frac{1}{n_1}]$写出数据量每个基本任务块计算得到$m_1n_1$的结果并搬出总写出量为$2Byte * MN$计算量输出矩阵C的每个数据点需要$K$次乘加总计算量固定为$2MNK$计算耗时多数情况下为刚性时间仅与参与计算的AIC核数相关在各理论模板中都一样后续不再赘述。多核切K模板 MultiCoreSplitKTiling建模如图展示一个常规的fp16的矩阵运算L0C上按照fp32累加$MN$方向共切分12个基本任务块假设有24个AIC物理核此时负载不均衡故引入$K$轴切分成2个$k$、产生24个基本任务块在AIC上负载均衡。定义相关参数问题shape$M$$N$$K$搬入L1Cache时的TileShape$m_1$$n_1$$k_1$搬入L0A/L0B、搬出L0C时的TileShape$m_0$$n_0$$k_0$相较Common模板新增$K$方向切分长度$k$在$m_1$、$n_1$较大时可能存在负载不均问题即M,N方向切分的任务块远少于AIC核数导致读取带宽不高核数不够因此可以在K方向上分核。产生$\frac{MNK}{m_1n_1k}$个基本任务块、分配给AIC核完成搬运和计算每个基本任务块需要搬运$m_1kkn_1$的数据、计算得到$m_1n_1$的结果并搬出。硬件上约束与Common相同。读取数据量每个基本任务块需要搬运$m_1KKn_1$的数据总读取量与Common模板一致$2Byte * [m_1kkn_1] * \frac{MNK}{m_1n_1k} 2Byte * MNK * [\frac{1}{m_1}\frac{1}{n_1}]$写出数据量每个基本任务块计算得到$m_1n_1$的结果并搬出需要$\frac{K}{k}$个基本块累加来得到输出矩阵C的$m_1n_1$块的最终输出总写出量为$2Byte * MNK / k$定性分析相较Common模板搬入数据量不变写出数据量增加并产生后处理ReduceAdd的开销包含MIX算子编译启动的开销但切分基本块更多、更易负载均衡。单核切K模板 SingleCoreSplitKTiling建模如图展示一个常规的fp16的矩阵运算L0C上按照fp32累加定义相关参数问题shape$M$$N$$K$搬入L1Cache时的TileShape$m_1$$n_1$$k_1$搬入L0A/L0B、搬出L0C时的TileShape$m_0$$n_0$$k_0$相比Common模板为了减少读取数据量进一步增大抽象上的$m_1$、$n_1$考虑将$m_1k_1$的tile块直接与对应的所有$k_1n_1$的tile块完成计算等同于将$n_1$放大到$N$此时输出$m_0n_0$的tile块没法在$L0C$常驻累加需要及时搬出通过atomicAdd在GM上累加。硬件上约束如下$m_1k_1L1Stage_A n_1k_1L1Stage_B L1Size / 2Byte$$m_0k_0*L0AStage L0ASize / 2Byte$$n_0k_0*L0BStage L0BSize / 2Byte$$m_0n_0*L0CStage L0CSize / 4Byte$$m_0 m_1$$n_0 n_1$读取数据量在Common模板的读取数据量公式上将$n_1$放大到$N$或者从A矩阵分基本任务块来理解切分$\frac{MK}{m_1k_1}$个基本块每个基本块完成搬入此块A矩阵tile块以及对应全部的B矩阵tile块搬入的数据量为$m_1k_1k_1N$$2Byte * [m_1k_1k_1N] * \frac{MK}{m_1k_1} 2Byte * MNK * [\frac{1}{m_1}\frac{1}{N}]$写出数据量切分A矩阵分基本任务块共$\frac{MK}{m_1k_1}$个基本块每个基本任务块计算得到$m_1N$的结果并搬出总写出量为$2Byte * MNK / k_1$定性分析相较Common模板搬入数据量减少写出数据量增加与AIV无关。工程优化清单流水优化Multi Buffer现象分析如下图构造一个Common模板下的简单场景对于单个AIC处理一个基本任务块C需要的A/B矩阵Tile块较小可以直接全部放入L1且A/B矩阵从L1搬入L0时需要切4次搬入。各PIPE的指令流水图示例如下如果在AIC的L1/L0A/L0B/L0C上每次载入数据tile块时都尽量填满所有空间会导致各PIPE的流水串行整体效率低下优化方案使用常规优化手段Multi Buffer即在L1/L0A/L0B/L0C上启用多buffer使得流水尽可能并行提升效率这一优化策略如下图所示各PIPE的指令流水图示例如下MTE1上指令的0、1表示pingpong流水⚠️ 需要注意的是与L1常驻优化结合时常驻的A/B矩阵的tile块不启用多buffer。特性承载代码由于是常规优化手段所有blockMmad组件均使能。流水优化Preload现象分析通过仿真流水发现pingpong策略的blockMmad存在问题MTE2流水上“当前C矩阵基本块计算的最后一个A矩阵B矩阵的tile块”和“下一个C矩阵基本块计算的第一个A矩阵B矩阵的tile块”之间加载的空泡。优化方案针对GM-L1过程读取当前轮次的$m_1k_1$$k_1n_1$时计算上一轮读取的数据假设Preload一轮PRELOAD_STAGES 1步骤伪代码如下for ... { // 搬入当前轮次的数据 copyGM2L1A copyGM2L1B preload_count for (preload_count PRELOAD_STAGES) { // 计算前PRELOAD_STAGES轮次的数据 copyL12L0A copyL12L0B Mmad } }如下图构造一个Common模板下的简单场景对于单个AIC处理两个基本任务块C1和C2需要的A/B矩阵Tile块放入L1需要切分2次且A/B矩阵L1Tile块从L1搬入L0时需要切分4次可以先学习前文流水优化Multi Buffer来增强理解。以下给出MmadAtlasA2Pingpong和MmadAtlasA2Preload、MmadAtlasA2PreloadAsync的指令对比MmadAtlasA2Pingpong中调用两次blockMmad分别完成C1和C2的计算MmadAtlasA2Preload中两次blockMmad也是分别完成C1和C2的计算但A3/B3的GmToL1搬运提前到了第一次blockMmad中执行MmadAtlasA2PreloadAsync中调用两次blockMmad和一次收尾的SynchronizeBlock。A2/B2的L1ToL0搬运、tileMmad以及C1的搬出从第一次blockMmad中推迟到第二次blockMmad中A4/B4的L1ToL0搬运、tileMmad以及C2的搬出从第二次blockMmad中推迟到SynchronizeBlock中各PIPE的指令流水图示例如下最终达成了A3、B3块的GmToL1搬运提前减缓了MTE2上的搬运空泡特性承载代码block_mmad_preload.hpp对应dispatchPolicyMmadAtlasA2Preload需要在kernel内手动计算传入下一块预载数据的信息。block_mmad_preload_async.hpp对应dispatchPolicyMmadAtlasA2PreloadAsync通过异步控制无需手动计算下一块预载数据信息并支持mmad计算完成后的Callback传入。block_mmad_preload_async_with_callback.hpp对应dispatchPolicyMmadAtlasA2PreloadAsyncWithCallback通过异步控制无需手动计算下一块预载数据信息并支持blockMmad计算前后的Callback传入。读取带宽优化Padding现象分析当数据读取为主流水时优化读取带宽能带来性能收益以fp16的A矩阵为例目前有以下几种低带宽场景Stride非512B对齐导致的低带宽。搬运参数srcDValue详见昇腾文档DataCopy-随路转换ND2NZ搬运非512B对齐时带宽会有明显下降。搬运指令限制导致的低带宽。ND2NZ的搬运指令参数srcDValue是uint16类型最大值65535。当K65535时只能通过取ndNum 1在m方向循环调用搬运指令降低了读取带宽。相比ND2ND不转换排布ND2NZ随路转换有带宽损失。优化方案针对上述情况可以使用AIV对数据格式进行重排预处理动作当重排开销低于带宽损失时会有性能收益。从复杂度和能应对的场景出发有下列三种不同的重排方式。PaddingMatrixND将Stride方向按照512B对齐实现复杂度最低可以处理Stride非对齐导致的带宽下降。PaddingMatrixBlockND按$m_1*k_1$作为“block”粒度重排block内行优先、block间行优先且$k_1$为512B对齐实现复杂度适中可以处理Stride非对齐和Stride超过65535导致的带宽下降。PaddingMatrixNZ重排为zN格式实现复杂度最高在介绍的几种Padding策略内因为和L1上数据排布一致搬运带宽也最高可以处理Stride非对齐、Stride超过65535、ND2NZ随路转换导致的带宽下降。实际应用中不同case适合的padding方式不同暂无全局最优Padding选择。特性承载代码padding_matmul.hpp中包含Padding前处理组件。实际适配可参考06_optimized_matmul通过PaddingTag、PaddingBuilder组装出A矩阵/B矩阵的Padding前处理。读取带宽优化ShuffleK现象分析通常所有AIC核心都从$K$方向的第一个分块开始搬运计算会存在多个核心同时读取同一片GM上数据的情况产生数据读取冲突导致读取带宽降低。优化方案以Common模板为例如图$matC$中的$CoreX$表示该基本块分配给第$X$个AIC进行计算$Aj$表示A矩阵该$m_1$下沿$K$轴切分的第$j$个L1Tile基本块$Bij$表示B矩阵该$n_1$下沿$K$轴切分的第$j$个L1Tile基本块图中matC基本块分核采用Swizzle2, 1详见swizzle_explanation如图左原始方案$Core2$和$Core3$搬运A矩阵时均按照“$A0$-$A1$-$A2$-$A3$”的顺序产生了数据读取冲突。ShuffleK方案如图所示根据$CoreIdx$来偏移起始搬运序号$j$$Core2$搬运A矩阵时按照“$A2$-$A3$-$A0$-$A1$”的顺序对应B矩阵按照“$B02$-$B03$-$B00$-$B01$”的顺序$Core3$搬运A矩阵均按照“$A3$-$A0$-$A1$-$A2$”的顺序对应B矩阵按照“$B13$-$B10$-$B11$-$B12$”的顺序。从时间上错开避免同地址访问冲突。特性承载代码block_mmad_preload.hppblock_mmad_preload_async.hppblock_mmad_preload_async_with_callback.hpp上述Block层实现代码通过设置起始L1序号为CoreIdx/kTileCount来实现错位kTileCount CeilDivL1TileShape::K(actualShape.k()); startTileIdx AscendC::GetBlockIdx(); firstTileIdx startTileIdx % kTileCount;读取带宽优化小M下指令替换现象分析矩阵计算中当$M$很小时例如$M$ 8采用随路ND2NZ的DataCopy详见昇腾文档DataCopy-随路转换ND2NZ搬运效率不高。优化方案通过for循环每次搬运一行并在每一行调用DataCopy进行多次搬运采用普通间隔搬运。特性承载代码CopyGmToL1IntervalDataCopy实际适配可参考06_optimized_matmul在struct TileCopyOpt中手动替换using CopyGmToL1A Gemm::Tile::CopyGmToL1IntervalDataCopyArchTag, AType;而不是默认的using CopyGmToL1A typename Base::CopyGmToL1A;struct TileCopyOpt : public Catlass::Gemm::Tile::TileCopyArchTag, AType, BType, CType, BiasType { ... - // using CopyGmToL1A Gemm::Tile::CopyGmToL1IntervalDataCopyArchTag, AType; using CopyGmToL1A Gemm::Tile::CopyGmToL1IntervalDataCopyArchTag, AType; - using CopyGmToL1A typename Base::CopyGmToL1A; // using CopyGmToL1A typename Base::CopyGmToL1A; ... };读取带宽优化L1常驻优化方案实际开发时可采用tile块常驻L1的方式减少tile块数据的重复读取等效提升了读取带宽该特性需要结合不同理论模板来考虑实现方法。特性承载代码Common模板可参考25_matmul_full_loadA及相关交付件该样例通过$M$轴上的单核全载或多核全载来在特定场景下优化性能并配合专门的swizzle策略来提高L1上全载的A矩阵块的复用频率。kernelmatmul_full_loadA.hppblockMmadblock_mmad_pingpong_full_loadA.hppdispatchPolicyMmadAtlasA2FullLoadABlockSchedulerGemmIdentityBlockSwizzleL1FullLoad单核切K模板34_single_core_splitk_matmul在理论设计上就考虑了L1Tile块常驻的优化点。kernelsingle_core_slicek_matmul.hppblockMmadblock_mmad_single_core_splitk.hppdispatchPolicyMmadAtlasA2SingleCoreSplitkBlockSchedulerSingleCoreSplitkGemmIdentityBlockSwizzleScalar开销消减现象分析对于小Shape场景如Common模板中$M$、$N$方向分核数小于实际物理核数每个AIC物理核最多仅处理一个基本任务块$k_1$ $K$$K$方向无需切分后从GM搬入L1此时kernel总耗时较小scalar开销对性能影响显著。优化方案消减冗余的scalar计算kernel内不使用 BlockScheduler 来将任务块分配给物理核手动计算每个物理核对应的任务块kernel内消除基本块循环每个AIC仅处理1个任务块kernel内简化offset相关计算blockMmad内消除L1A/L1B的$m$、$n$的循环blockMmad内简化offset相关计算特性承载代码参考31_small_matmul可以和00_basic_matmul的相关交付件对比来加深理解kernelsmall_matmul.hppblockMmadblock_mmad_small.hpp写出带宽优化现象分析当数据写出为主流水时优化写出带宽能够带来性能收益。当写出时dstStride未进行512B对齐带宽有明显下降写出时用NZ2ND随路格式转换产生带宽损失优化方案针对上述情况可使用AIV对数据格式进行重排在重排开销低于带宽损失时会有性能收益。以下提供四种重排方式。↑方式一使用局部workSpaceND写出到GM时512B对齐随后按block块粒度在UB上重排再写出到GM上。↑方式二使用全量workSpaceND写出到GM时512B对齐等全量结果写完后再启动UB上数据重排写入到GM上。↑方式三使用局部workSpaceNZ写出到GM时512B对齐再按block块粒度在UB上重排ND写出到GM上。↑方式四使用全量workSpaceNZ写出到GM时512B对齐等全量结果写完后再启动UB上数据重排ND写出到GM上。特性承载代码padding_matmul.hpp中实现了包含方式二的RemovePaddingNDAndCast后处理组件实际适配可参考34_single_core_splitk_matmul模板应用浅述参考102_dynamic_optimized_matmul的select_kernel策略模板选择先尝试基于00_basic_matmul进行TileShape调优并获得性能基线可参考模板库优化指引依次识别是否满足各模板适合场景并和性能基线作比较31_small_matmul计算当前基本任务块 taskBlockstaskBlocks CeilDiv(M, m1) * CeilDiv(N, n1);基本任务块小于AIC数 $taskBlocks aicCoreNum$$K$轴较小$K k_1$09_splitk_matmul或22_padding_splitk_matmul带Padding前处理选择$m_1$、$n_1$、$k_1$先设置$m_1 128$、$n_1 256$、$k_1 256$满足下列场景其一修改$m_1 256$、$n_1 128$A矩阵和B矩阵均为ColumnMajorA矩阵为ColumnMajorB矩阵为RowMajor且$M N$计算当前基本任务块 taskBlockstaskBlocks CeilDiv(M, m1) * CeilDiv(N, n1);满足下列两种场景基本任务块小于一半AIC数且$K$轴够大$taskBlocks aicCoreNum / 2, K 5120$基本任务块小于3块且$K$轴不会小$taskBlocks 2, K 1024$06_optimized_matmul带Padding前处理和21_basic_matmul_preload_zN手动改为ND输入泛化性更强适用于剩余场景不需要使用Padding时为了节约MIX算子编译启动的开销建议使用21_basic_matmul_preload_zN模板⚠️ 全载特性的使用25_matmul_full_loadA 和 单核切K方案34_single_core_splitk_matmul的适用场景待完善Padding选择当Stride非512B对齐时可以考虑使用Padding前处理但需要考虑Padding带来的开销以及MIX算子编译启动的开销小shape31_small_matmul方法不推荐额外适配PaddingPaddingMatrixND、PaddingMatrixBlockND和PaddingMatrixNZ各自的适用场景待完善泛化上PaddingMatrixNZ更具有优势。【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考