CATLASS Single_core_splitK_Matmul 样例介绍【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass样例实现CATLASS34_single_core_splitk_matmul样例算子是基于CATLASS Gemm API实现的昇腾亲和Matmul算子针对大尺寸矩阵计算场景优化设计关键算子组件包括以下几部分:Example组装single_core_splitk.cppKernel实现主Kernel文件single_core_slicek_matmul.hpp复用Padding组件padding_matmul.hppBlock组件 block_mmad_single_core_splitk.hppExample组装与通用的模板库开发风格一致本样例的实现执行过程如上图所示简要说明如下构造输入生成要计算的左/右矩阵计算各输入矩阵的尺寸在Host侧生成各输入矩阵值构造Device侧输入组装Padding对象为便于数据搬运的对齐组装Padding对象定义不同的PaddingTag通过PaddingBuilder组装出对于矩阵A/B的Padding对象预定义的PaddingC对象决定是否启用PaddingC对象组装blockMmad为便于数据搬运的对齐先组装相关的Padding对象定义各输入的Layout特征声明不同矩阵所对应的类型情况设置Dispatch策略用于选取BlockMmad组件设置L1和L0的Tile尺寸用于优化从GM到L1的搬运过程使用优化后的TileCopy组件使用上述模板入参组装BlockMmad组装并执行Kernel组装Kernel并实例化该对象完成算子计算使用新的S型Swizzle策略使用前述模板组装Kernel将Kernel传入适配器并实例化构造入参参数向Kernel侧传入参数调用Kernel层Workspace计算在Device侧申请Workspace(如有必要)适配器初始化算子执行算子精度校验和空间释放完成最后结果的验证与回收处理将算子输出结果搬回host侧计算golden标杆精度比对释放输入输出和workspaceKernel实现以下介绍Kernel层级的结构体与关键函数以及AIC/AIV部分的简明计算流程并说明了所使用的优化策略Kernel层的主要结构体与函数以下是在Kernel层所实现的结构体与关键函数struct Params运行时算子执行所需的参数struct Arguments封装Host侧传入的参数static size_t GetWorkspaceSize预先计算对齐所需的空间static Params ToUnderlyingArguments将Host侧入参解析为算子侧的Params结构体void operator()AscendC::AIVAIV(Vector)部分执行代码void operator()AscendC::AIC: AIC(Cube)部分执行代码AIV/AIC部分计算流程以下是Kernel层所进行的AIC/AIV操作AIV上所执行的操作如果A或B的对齐处理启用初始化GlobalTensorgmA gmWA(或gmB, gmWB)实例化PaddingA对象完成资源申请以便在UB(Unified Buffer)上存储对齐前后的数据调用PaddingA对象的operator()方法即执行数据对齐操作进行核间同步设置标志位通知AIC对齐操作已完成(CrossCoreSetFlag)结果矩阵搬出等待AIC完成计算(CrossCoreWaitFlag)初始化GlobalTensorgmC(目的地址),gmWC(临时累加地址)将数据搬出到目的地址进行升精度映射(RemovePaddingNDAndCastC)AIC上所执行的操作如果对左/右矩阵有对齐处理核间同步等待AIV完成标志位(CrossCoreWaitFlag)初始化GlobalTensorgmWA,gmWB初始化GlobalTensor:gmA,gmB,gmC初始化BlockScheduler和BlockMmad对象获取当前AIC序号coreIdx、AIC总数coreNum在Swizzle策略内以及所需的coreLoops进入主循环循环次coreLoops计算当前AB矩阵读入的偏移量gmOffsetA,gmOffsetB以及下一块的偏移量gmOffsetNextA,gmOffsetNextB(如果开DoubleBuffer)注意优化算法下左矩阵重载因此gmOffsetNextA实际不会“启用”计算needLoadNextA、needLoadNextB用以标识是否预加载调用blockMmad进行一次AIC计算完成L1A与L1B上相对应的分形矩阵计算注意blockMmad会依据K轴方向上切K情况决定是否启用GM上的原子加设置标志位通知AIV 计算已完成关闭原子加单核切K算法如上图所示相较于经典的矩阵乘过程单核切K的Matmul模板采取了复用左矩阵的办法即对于某个AI Core其L1A上加载的分形矩阵是驻留的进而降低对GM的访问量。举例而言对于经典计算过程单核所进行的数据搬运过程包括1. 从GM搬运A0到L1AB0到L1B2. 从GM搬运A1到L1AB1到L1B..., N. 从GM搬运A(n-1)到L1A B(n-1)到L1B。在本单核切K策略下单个AI Core上的流程图如下所示显然该办法能有效减轻MTE2的搬运负担但是会加重L0C计算结果搬回到GM的负担需要在GM上开启原子加。// Atomic set once in mmad level if (atomicAdd) { AscendC::SetAtomicAddElementAccumulator(); } else { AscendC::SetAtomicNone(); }增大L1利用空间 简单建模可知见下述补充从L0C搬运回GM的数据量正比于$MNK/k_{\text{L1}}$其中$M$, $N$, $K$分别为输入的矩阵尺寸$k_{\text{L1}}$是K轴上L1Tile的尺寸大小。在符合L1A, L1B物理尺寸限制的条件下$k_{\text{L1}}$越大可减少写出次数推荐L1TileShape按下表进行配置。数据搬运理论建模我们首先分析下基础Matmul详见本仓00_basic_matmul样例的内存访问量。设矩阵乘的尺寸为$(M, N, K)$L1上的Tile尺寸相应为$m$, $n$, $k$。假设二者对齐。则单个分形结果矩阵从GMGlobal Memory至L1所需的搬运量是$K(mkkn)/k$加之分形矩阵个数为$MN/mn$因此总的访问量为$MNK(1/m1/n)$。另一方面由于在L0C上进行累加因此搬出的单位数据量即为$MN$。 再考虑使用本优化策略后的内存访问设复用左矩阵则单核总的数据搬运量为$mK KN$总的访问量是$(mK KN)M/m$也即$MNK(1/m1/N)$小于基础Matmul的情形。另外由于计算结果并不能在L0C上进行累加需“随算随搬”因此搬出的数据量为$MNK/k$ 较基础有一定升高。类别内存访问操作总量基础矩阵乘$MNK(1/m1/n) MN$单核切K$MNK(1/m1/N) MNK/k$数据类型L1TileShape::ML1TileShape::NL1TileShape::KFP16/BF16256128512FP32256128256数据对齐 如需对数据进行类型转换需采取数据对齐以提高处理流程中的数据搬运带宽。// RemovePaddingNDAndCast的核心处理逻辑 uint32_t loopsPerTile RoundUp(tileLen, COMPUTE_LENGTH); uint32_t coreLoops tilesPerAiv * loopsPerTile; for (uint32_t loopIdx 0; loopIdx coreLoops; loopIdx) { // 计算gmWC搬运的offset // 将GmWC(src)的数据搬运至inputBuffer上 copyGm2Ub(inputBuffer, src, ubLayout, srcLayout); // ... // Cast到half AscendC::Cast(outputBuffer, inputBuffer,AscendC::RoundMode::CAST_RINT, actualDataNum); // ... // 将outputBuffer搬出到gmC(dst) copyUb2Gm(dst, outputBuffer[bufferIndex], dstLayout, ubLayout); // ... }Swizzle排布方案在矩阵C上的swizzle策略采用S型特征相较于Z型的swizzle排布在换行处可节约一次搬运。如下图所示按该S型Swizzle策略前一组需要加载到L1A, L1B上的分形矩阵为A02, B20, B21, B22, B23后一次需要加载的是A12, B20, B21, B22, B23(以SwizzleOffset为4为例)。“换行”过程下可以固定L1B上的数据仅从GM上重读入左矩阵可减轻载入负担。说明该Swizzle策略在大尺寸下有效排布策略会优先将任务分配至不同的AI Core上。性能收益经过实测应用上述单核切K算法在大尺寸场景下较基础Matmul有正向收益且随着K轴的尺寸增大GM至L1节省的正向收益高过重复写入GM的负向收益可参考下表。不过需要说明的是在MN过小或者K偏低的情况下会出现分核负载不均衡的情况。| M | N | K | 耗时(us) | 耗时[标杆](us) | 加速比 | | -- | -- | ---- | ----- | ----- | ----- | | 2048 | 4096 | 4000 | 261 | 445 | 1.7049 | | 4096 | 4096 | 8000 | 917 | 1231 | 1.3424 | | 4096 | 4096 | 40000 | 4669 | 7775 | 1.6652 | | 2048 | 4096 | 80000 | 16850 | 57144 | 1.7499 |说明标杆为BasicMatmul算子统计耗时均为核函数总耗时使用msprof工具得到上述测试例中AB及C矩阵均为layout::RowMajor排布方式测试环境说明NPU型号为910B2CANN包版本为8.2.RC1【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
CANN/catlass单核切K矩阵乘样例
发布时间：2026/7/7 14:16:07
CATLASS Single_core_splitK_Matmul 样例介绍【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass样例实现CATLASS34_single_core_splitk_matmul样例算子是基于CATLASS Gemm API实现的昇腾亲和Matmul算子针对大尺寸矩阵计算场景优化设计关键算子组件包括以下几部分:Example组装single_core_splitk.cppKernel实现主Kernel文件single_core_slicek_matmul.hpp复用Padding组件padding_matmul.hppBlock组件 block_mmad_single_core_splitk.hppExample组装与通用的模板库开发风格一致本样例的实现执行过程如上图所示简要说明如下构造输入生成要计算的左/右矩阵计算各输入矩阵的尺寸在Host侧生成各输入矩阵值构造Device侧输入组装Padding对象为便于数据搬运的对齐组装Padding对象定义不同的PaddingTag通过PaddingBuilder组装出对于矩阵A/B的Padding对象预定义的PaddingC对象决定是否启用PaddingC对象组装blockMmad为便于数据搬运的对齐先组装相关的Padding对象定义各输入的Layout特征声明不同矩阵所对应的类型情况设置Dispatch策略用于选取BlockMmad组件设置L1和L0的Tile尺寸用于优化从GM到L1的搬运过程使用优化后的TileCopy组件使用上述模板入参组装BlockMmad组装并执行Kernel组装Kernel并实例化该对象完成算子计算使用新的S型Swizzle策略使用前述模板组装Kernel将Kernel传入适配器并实例化构造入参参数向Kernel侧传入参数调用Kernel层Workspace计算在Device侧申请Workspace(如有必要)适配器初始化算子执行算子精度校验和空间释放完成最后结果的验证与回收处理将算子输出结果搬回host侧计算golden标杆精度比对释放输入输出和workspaceKernel实现以下介绍Kernel层级的结构体与关键函数以及AIC/AIV部分的简明计算流程并说明了所使用的优化策略Kernel层的主要结构体与函数以下是在Kernel层所实现的结构体与关键函数struct Params运行时算子执行所需的参数struct Arguments封装Host侧传入的参数static size_t GetWorkspaceSize预先计算对齐所需的空间static Params ToUnderlyingArguments将Host侧入参解析为算子侧的Params结构体void operator()AscendC::AIVAIV(Vector)部分执行代码void operator()AscendC::AIC: AIC(Cube)部分执行代码AIV/AIC部分计算流程以下是Kernel层所进行的AIC/AIV操作AIV上所执行的操作如果A或B的对齐处理启用初始化GlobalTensorgmA gmWA(或gmB, gmWB)实例化PaddingA对象完成资源申请以便在UB(Unified Buffer)上存储对齐前后的数据调用PaddingA对象的operator()方法即执行数据对齐操作进行核间同步设置标志位通知AIC对齐操作已完成(CrossCoreSetFlag)结果矩阵搬出等待AIC完成计算(CrossCoreWaitFlag)初始化GlobalTensorgmC(目的地址),gmWC(临时累加地址)将数据搬出到目的地址进行升精度映射(RemovePaddingNDAndCastC)AIC上所执行的操作如果对左/右矩阵有对齐处理核间同步等待AIV完成标志位(CrossCoreWaitFlag)初始化GlobalTensorgmWA,gmWB初始化GlobalTensor:gmA,gmB,gmC初始化BlockScheduler和BlockMmad对象获取当前AIC序号coreIdx、AIC总数coreNum在Swizzle策略内以及所需的coreLoops进入主循环循环次coreLoops计算当前AB矩阵读入的偏移量gmOffsetA,gmOffsetB以及下一块的偏移量gmOffsetNextA,gmOffsetNextB(如果开DoubleBuffer)注意优化算法下左矩阵重载因此gmOffsetNextA实际不会“启用”计算needLoadNextA、needLoadNextB用以标识是否预加载调用blockMmad进行一次AIC计算完成L1A与L1B上相对应的分形矩阵计算注意blockMmad会依据K轴方向上切K情况决定是否启用GM上的原子加设置标志位通知AIV 计算已完成关闭原子加单核切K算法如上图所示相较于经典的矩阵乘过程单核切K的Matmul模板采取了复用左矩阵的办法即对于某个AI Core其L1A上加载的分形矩阵是驻留的进而降低对GM的访问量。举例而言对于经典计算过程单核所进行的数据搬运过程包括1. 从GM搬运A0到L1AB0到L1B2. 从GM搬运A1到L1AB1到L1B..., N. 从GM搬运A(n-1)到L1A B(n-1)到L1B。在本单核切K策略下单个AI Core上的流程图如下所示显然该办法能有效减轻MTE2的搬运负担但是会加重L0C计算结果搬回到GM的负担需要在GM上开启原子加。// Atomic set once in mmad level if (atomicAdd) { AscendC::SetAtomicAddElementAccumulator(); } else { AscendC::SetAtomicNone(); }增大L1利用空间 简单建模可知见下述补充从L0C搬运回GM的数据量正比于$MNK/k_{\text{L1}}$其中$M$, $N$, $K$分别为输入的矩阵尺寸$k_{\text{L1}}$是K轴上L1Tile的尺寸大小。在符合L1A, L1B物理尺寸限制的条件下$k_{\text{L1}}$越大可减少写出次数推荐L1TileShape按下表进行配置。数据搬运理论建模我们首先分析下基础Matmul详见本仓00_basic_matmul样例的内存访问量。设矩阵乘的尺寸为$(M, N, K)$L1上的Tile尺寸相应为$m$, $n$, $k$。假设二者对齐。则单个分形结果矩阵从GMGlobal Memory至L1所需的搬运量是$K(mkkn)/k$加之分形矩阵个数为$MN/mn$因此总的访问量为$MNK(1/m1/n)$。另一方面由于在L0C上进行累加因此搬出的单位数据量即为$MN$。 再考虑使用本优化策略后的内存访问设复用左矩阵则单核总的数据搬运量为$mK KN$总的访问量是$(mK KN)M/m$也即$MNK(1/m1/N)$小于基础Matmul的情形。另外由于计算结果并不能在L0C上进行累加需“随算随搬”因此搬出的数据量为$MNK/k$ 较基础有一定升高。类别内存访问操作总量基础矩阵乘$MNK(1/m1/n) MN$单核切K$MNK(1/m1/N) MNK/k$数据类型L1TileShape::ML1TileShape::NL1TileShape::KFP16/BF16256128512FP32256128256数据对齐 如需对数据进行类型转换需采取数据对齐以提高处理流程中的数据搬运带宽。// RemovePaddingNDAndCast的核心处理逻辑 uint32_t loopsPerTile RoundUp(tileLen, COMPUTE_LENGTH); uint32_t coreLoops tilesPerAiv * loopsPerTile; for (uint32_t loopIdx 0; loopIdx coreLoops; loopIdx) { // 计算gmWC搬运的offset // 将GmWC(src)的数据搬运至inputBuffer上 copyGm2Ub(inputBuffer, src, ubLayout, srcLayout); // ... // Cast到half AscendC::Cast(outputBuffer, inputBuffer,AscendC::RoundMode::CAST_RINT, actualDataNum); // ... // 将outputBuffer搬出到gmC(dst) copyUb2Gm(dst, outputBuffer[bufferIndex], dstLayout, ubLayout); // ... }Swizzle排布方案在矩阵C上的swizzle策略采用S型特征相较于Z型的swizzle排布在换行处可节约一次搬运。如下图所示按该S型Swizzle策略前一组需要加载到L1A, L1B上的分形矩阵为A02, B20, B21, B22, B23后一次需要加载的是A12, B20, B21, B22, B23(以SwizzleOffset为4为例)。“换行”过程下可以固定L1B上的数据仅从GM上重读入左矩阵可减轻载入负担。说明该Swizzle策略在大尺寸下有效排布策略会优先将任务分配至不同的AI Core上。性能收益经过实测应用上述单核切K算法在大尺寸场景下较基础Matmul有正向收益且随着K轴的尺寸增大GM至L1节省的正向收益高过重复写入GM的负向收益可参考下表。不过需要说明的是在MN过小或者K偏低的情况下会出现分核负载不均衡的情况。| M | N | K | 耗时(us) | 耗时[标杆](us) | 加速比 | | -- | -- | ---- | ----- | ----- | ----- | | 2048 | 4096 | 4000 | 261 | 445 | 1.7049 | | 4096 | 4096 | 8000 | 917 | 1231 | 1.3424 | | 4096 | 4096 | 40000 | 4669 | 7775 | 1.6652 | | 2048 | 4096 | 80000 | 16850 | 57144 | 1.7499 |说明标杆为BasicMatmul算子统计耗时均为核函数总耗时使用msprof工具得到上述测试例中AB及C矩阵均为layout::RowMajor排布方式测试环境说明NPU型号为910B2CANN包版本为8.2.RC1【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
 获取与实战)
