前言深度学习的计算量有90%以上花在矩阵乘法上这个说法不算夸张。Transformer的Attention是矩阵乘全连接层是矩阵乘卷积im2col之后还是矩阵乘。昇腾NPU上的矩阵乘法算子由ops-blas仓库提供它是CANN算子库体系中最底层的线性代数算子库为上层的ops-nn、ops-transformer、catlass等仓库提供GEMM通用矩阵乘法的核心实现。ops-blas的性能直接决定了整个CANN生态的计算上限——如果GEMM算子跑不快上层所有依赖矩阵乘的算子都会被拖慢。CANN社区在atomgit.com/cann上开源了ops-blas仓库让开发者可以直接使用和调优这些高性能GEMM算子。GEMM算子为什么是性能核心GEMM的定义是C alpha * A B beta * C其中A是MxK矩阵B是KxN矩阵C是MxN矩阵。这个看似简单的操作在昇腾NPU上的实现涉及大量的硬件细节。昇腾NPU上的矩阵乘法由Cube计算单元执行。Cube单元的物理结构是一个矩阵乘法加速器单次可以计算一个16x16的矩阵块FP16或8x8的矩阵块FP32。一个AI Core里的Cube单元在每个时钟周期可以完成一次这样的小矩阵乘法峰值算力取决于芯片的时钟频率和Cube单元数量。GEMM算子的核心挑战不在于计算本身——Cube单元执行矩阵乘法的效率是固定的——而在于如何把数据高效地喂给Cube单元。具体来说有三个问题数据从Global Memory搬运到Cube单元的路径太长。数据需要经过L1 Cache、L0A/L0B Cache才能到达Cube的输入端每一级搬运都有延迟。如果Cube单元等数据的时间比计算的时间还长算力就浪费了。矩阵尺寸不一定是16的倍数。Cube单元一次计算16x16的块但实际应用中矩阵的M、N、K维度可能不是16的倍数需要处理边界对齐。多核负载均衡。昇腾910有多个AI CoreGEMM需要把矩阵切分到多个AI Core上并行计算切分方式影响负载均衡和核间通信开销。ops-blas的Tiling策略Tiling是GEMM性能优化的核心。Tiling的含义是把大矩阵切成小块让每个小块都能放进AI Core的本地内存中计算。ops-blas的Tiling策略分两级一级Tiling跨AI Core的矩阵切分。把MxN的结果矩阵C按行切分行切分或按块切分棋盘格切分每个AI Core负责计算C的一个子块。切分粒度要保证每个AI Core的计算量大致相等避免某些核先完成而其他核还在计算。二级TilingAI Core内部的分块计算。每个AI Core拿到自己的子块后需要进一步把子块切分成L0A/L0B能装下的小块。二级Tiling的关键参数是tile_m、tile_n、tile_k分别表示在M、N、K三个方向上的切分粒度。来看一个具体的Tiling例子。假设M4096N4096K4096AI Core有32个L0A/L0B各64KB# ops-blas的Tiling参数计算简化版# 实际实现更复杂这里只展示核心逻辑M,N,K4096,4096,4096num_aicore32l0a_size64*1024# 64KBl0b_size64*1024# 一级Tiling按行切分C矩阵# 每个AI Core负责M/num_aicore行的计算tile_m_outerM//num_aicore# 128行/AI Core# 二级TilingAI Core内部的分块# L0A需要装下tile_m x tile_k的A子块FP16每个元素2字节# tile_m * tile_k * 2 l0a_size# 取tile_k 64对齐Cube的16倍数要求计算tile_mtile_k64tile_ml0a_size//(tile_k*2)# 64KB / (64 * 2B) 512# 但tile_m必须是16的倍数Cube单元的粒度取最近的16倍数tile_m(tile_m//16)*16# 512# L0B需要装下tile_k x tile_n的B子块FP16# tile_k * tile_n * 2 l0b_sizetile_nl0b_size//(tile_k*2)# 512tile_n(tile_n//16)*16# 512# 为什么tile_m和tile_n要取16的倍数# 因为Cube单元一次计算16x16的矩阵块# 不是16倍数的话需要做Padding浪费计算资源这段计算给出了tile_m512tile_n512tile_k64的二级Tiling参数。但这个参数不一定是最优的——不同的矩阵尺寸和硬件代际有不同的最优参数。ops-blas内部维护了一张Tiling参数查找表LUT根据M、N、K的值从表中选取经过实测验证的最优参数。对于表中没有覆盖的矩阵尺寸ops-blas会使用一套启发式规则来计算Tiling参数。数据搬运的双缓冲优化Tiling确定了数据的切分方式接下来是如何高效搬运数据。ops-blas使用了双缓冲Double Buffer技术来掩盖数据搬运延迟。双缓冲的原理是准备两套L0A/L0B缓冲区当Cube单元在缓冲区0上计算当前分块时DataCopy指令同时把下一个分块的数据搬运到缓冲区1。当前分块计算完成后Cube直接切换到缓冲区1开始计算同时DataCopy把下下个分块搬运到缓冲区0。这样计算和搬运就重叠在一起了。// ops-blas内部的双缓冲数据搬运逻辑简化伪代码// 展示K维度上的分块循环// 先预取第一个k分块到缓冲区0DataCopy(l0a_buf[0],gm_aoffset_a_0,tile_m*tile_k);DataCopy(l0b_buf[0],gm_boffset_b_0,tile_k*tile_n);PipeBarrier();// 等待预取完成for(intk_idx0;k_idxK/tile_k;k_idx){intcur_bufk_idx%2;// 当前计算缓冲区intnext_buf(k_idx1)%2;// 下一个预取缓冲区// 如果还有下一个k分块异步预取到另一个缓冲区// 为什么用异步因为DataCopyAsync不会阻塞Scalar单元// Cube可以在当前缓冲区上继续计算if(k_idx1K/tile_k){DataCopyAsync(l0a_buf[next_buf],gm_aoffset_a_next,tile_m*tile_k);DataCopyAsync(l0b_buf[next_buf],gm_boffset_b_next,tile_k*tile_n);}// Cube单元在当前缓冲区上做矩阵乘加// MMAD矩阵乘加指令结果累加到L0CMMAD(l0c_buf,l0a_buf[cur_buf],l0b_buf[cur_buf],tile_m,tile_n,tile_k);// 等待异步预取完成和MMAD并行执行的PipeBarrier();}// 最终把L0C的结果搬回Global MemoryDataCopy(gm_c,l0c_buf,tile_m*tile_n);双缓冲的效果取决于搬运时间和计算时间的比例。如果搬运时间小于计算时间计算密集型双缓冲可以完全掩盖搬运延迟Cube单元始终处于忙碌状态。如果搬运时间大于计算时间访存密集型双缓冲只能部分缓解Cube单元还是会有空闲等待。GEMM的计算量是2MNK FLOPs搬运量是MK KN MN* sizeof(data_type) Bytes。当MNK足够大时计算量远大于搬运量GEMM是计算密集型的双缓冲效果很好。当矩阵较小时比如MNK128搬运量相对计算量较大GEMM变成访存密集型的双缓冲效果有限。不同矩阵尺寸下的性能分析下面是在Ascend 910上FP16 GEMM不同矩阵尺寸下的实测性能矩阵尺寸 (MxNxK)计算量 (GFLOPS)搬运量 (MB)算力强度 (FLOP/Byte)实测TFLOPS峰值利用率128x128x1280.0040.0981.7123%512x512x5120.2681.576.88521%1024x1024x10242.1476.2913.618045%4096x4096x4096137.4100.754.431078%8192x8192x81921099.5805.354.435088%算力强度FLOP/Byte是衡量计算密集程度的关键指标——每个字节的数据搬运对应多少次浮点运算。算力强度越高计算越密集双缓冲的效果越好。128x128x128的小矩阵峰值利用率只有3%因为计算量太小Cube单元大部分时间在等数据。4096以上的大矩阵利用率可以达到78%以上接近硬件峰值。这也就是为什么ops-blas对小矩阵做了专门的优化——当M和N较小时ops-blas会自动选择更小的Tiling参数比如tile_m16而不是512减少L0A/L0B的空闲空间同时把多个小矩阵打包到一个AI Core上串行处理避免频繁的核间同步开销。ops-blas和catlass的关系ops-blas和catlass都涉及GEMM实现但定位不同。ops-blas提供的是已编译好的高性能GEMM算子二进制开发者通过AscendCL的aclrtLaunchKernel或GE的算子调度来调用不能修改内部实现。catlass提供的是GEMM的Ascend C模板源码开发者可以基于模板定制Tiling策略、数据类型和搬运流程然后编译成自己的GEMM算子。简单说ops-blas是用catlass是改。如果你的GEMM需求和ops-blas提供的标准算子一致直接用ops-blas性能已经很好。如果你需要定制GEMM的实现比如特殊的矩阵排布、非标准的融合操作才需要用catlass从模板开始构建。使用前后效率对比以一个Transformer推理场景为例单层Attention的计算涉及3次GEMMQ/K/V投影和1次GEMM输出投影矩阵尺寸为batch_size x seq_len x hidden_dim。对比维度手写Tiling的Ascend C GEMMops-blas标准算子ops-blas Tiling调优开发时间3-5天0直接调用2小时参数调优4096x4096x4096性能280 TFLOPS310 TFLOPS340 TFLOPS1024x1024x1024性能120 TFLOPS180 TFLOPS195 TFLOPS128x128x128性能5 TFLOPS12 TFLOPS15 TFLOPSTransformer单层延迟4.2ms3.5ms3.1msops-blas标准算子在所有矩阵尺寸下都优于手写版本尤其是小矩阵场景——ops-blas内置的小矩阵优化策略不是手写Tiling能轻易复现的。经过Tiling调优后的ops-blas性能进一步提升4096尺寸达到340 TFLOPS接近Ascend 910的FP16理论峰值约400 TFLOPS的85%。结尾ops-blas是CANN算子体系里计算量最大的组件它的性能直接影响了上层所有依赖矩阵乘法的算子。理解GEMM的Tiling策略、双缓冲优化和算力强度对性能的影响有助于在模型部署时做出正确的优化决策——比如选择合适的batch size来让GEMM处于计算密集区域避免过小的矩阵尺寸导致NPU利用率低下。对于大多数开发者来说ops-blas的标准算子已经足够好需要极致性能的场景可以通过Tiling调优进一步提升。仓库地址https://atomgit.com/cann/ops-blas
昇腾CANN线性代数算子库ops-blas：GEMM算子从原理到调优的完整实践
前言深度学习的计算量有90%以上花在矩阵乘法上这个说法不算夸张。Transformer的Attention是矩阵乘全连接层是矩阵乘卷积im2col之后还是矩阵乘。昇腾NPU上的矩阵乘法算子由ops-blas仓库提供它是CANN算子库体系中最底层的线性代数算子库为上层的ops-nn、ops-transformer、catlass等仓库提供GEMM通用矩阵乘法的核心实现。ops-blas的性能直接决定了整个CANN生态的计算上限——如果GEMM算子跑不快上层所有依赖矩阵乘的算子都会被拖慢。CANN社区在atomgit.com/cann上开源了ops-blas仓库让开发者可以直接使用和调优这些高性能GEMM算子。GEMM算子为什么是性能核心GEMM的定义是C alpha * A B beta * C其中A是MxK矩阵B是KxN矩阵C是MxN矩阵。这个看似简单的操作在昇腾NPU上的实现涉及大量的硬件细节。昇腾NPU上的矩阵乘法由Cube计算单元执行。Cube单元的物理结构是一个矩阵乘法加速器单次可以计算一个16x16的矩阵块FP16或8x8的矩阵块FP32。一个AI Core里的Cube单元在每个时钟周期可以完成一次这样的小矩阵乘法峰值算力取决于芯片的时钟频率和Cube单元数量。GEMM算子的核心挑战不在于计算本身——Cube单元执行矩阵乘法的效率是固定的——而在于如何把数据高效地喂给Cube单元。具体来说有三个问题数据从Global Memory搬运到Cube单元的路径太长。数据需要经过L1 Cache、L0A/L0B Cache才能到达Cube的输入端每一级搬运都有延迟。如果Cube单元等数据的时间比计算的时间还长算力就浪费了。矩阵尺寸不一定是16的倍数。Cube单元一次计算16x16的块但实际应用中矩阵的M、N、K维度可能不是16的倍数需要处理边界对齐。多核负载均衡。昇腾910有多个AI CoreGEMM需要把矩阵切分到多个AI Core上并行计算切分方式影响负载均衡和核间通信开销。ops-blas的Tiling策略Tiling是GEMM性能优化的核心。Tiling的含义是把大矩阵切成小块让每个小块都能放进AI Core的本地内存中计算。ops-blas的Tiling策略分两级一级Tiling跨AI Core的矩阵切分。把MxN的结果矩阵C按行切分行切分或按块切分棋盘格切分每个AI Core负责计算C的一个子块。切分粒度要保证每个AI Core的计算量大致相等避免某些核先完成而其他核还在计算。二级TilingAI Core内部的分块计算。每个AI Core拿到自己的子块后需要进一步把子块切分成L0A/L0B能装下的小块。二级Tiling的关键参数是tile_m、tile_n、tile_k分别表示在M、N、K三个方向上的切分粒度。来看一个具体的Tiling例子。假设M4096N4096K4096AI Core有32个L0A/L0B各64KB# ops-blas的Tiling参数计算简化版# 实际实现更复杂这里只展示核心逻辑M,N,K4096,4096,4096num_aicore32l0a_size64*1024# 64KBl0b_size64*1024# 一级Tiling按行切分C矩阵# 每个AI Core负责M/num_aicore行的计算tile_m_outerM//num_aicore# 128行/AI Core# 二级TilingAI Core内部的分块# L0A需要装下tile_m x tile_k的A子块FP16每个元素2字节# tile_m * tile_k * 2 l0a_size# 取tile_k 64对齐Cube的16倍数要求计算tile_mtile_k64tile_ml0a_size//(tile_k*2)# 64KB / (64 * 2B) 512# 但tile_m必须是16的倍数Cube单元的粒度取最近的16倍数tile_m(tile_m//16)*16# 512# L0B需要装下tile_k x tile_n的B子块FP16# tile_k * tile_n * 2 l0b_sizetile_nl0b_size//(tile_k*2)# 512tile_n(tile_n//16)*16# 512# 为什么tile_m和tile_n要取16的倍数# 因为Cube单元一次计算16x16的矩阵块# 不是16倍数的话需要做Padding浪费计算资源这段计算给出了tile_m512tile_n512tile_k64的二级Tiling参数。但这个参数不一定是最优的——不同的矩阵尺寸和硬件代际有不同的最优参数。ops-blas内部维护了一张Tiling参数查找表LUT根据M、N、K的值从表中选取经过实测验证的最优参数。对于表中没有覆盖的矩阵尺寸ops-blas会使用一套启发式规则来计算Tiling参数。数据搬运的双缓冲优化Tiling确定了数据的切分方式接下来是如何高效搬运数据。ops-blas使用了双缓冲Double Buffer技术来掩盖数据搬运延迟。双缓冲的原理是准备两套L0A/L0B缓冲区当Cube单元在缓冲区0上计算当前分块时DataCopy指令同时把下一个分块的数据搬运到缓冲区1。当前分块计算完成后Cube直接切换到缓冲区1开始计算同时DataCopy把下下个分块搬运到缓冲区0。这样计算和搬运就重叠在一起了。// ops-blas内部的双缓冲数据搬运逻辑简化伪代码// 展示K维度上的分块循环// 先预取第一个k分块到缓冲区0DataCopy(l0a_buf[0],gm_aoffset_a_0,tile_m*tile_k);DataCopy(l0b_buf[0],gm_boffset_b_0,tile_k*tile_n);PipeBarrier();// 等待预取完成for(intk_idx0;k_idxK/tile_k;k_idx){intcur_bufk_idx%2;// 当前计算缓冲区intnext_buf(k_idx1)%2;// 下一个预取缓冲区// 如果还有下一个k分块异步预取到另一个缓冲区// 为什么用异步因为DataCopyAsync不会阻塞Scalar单元// Cube可以在当前缓冲区上继续计算if(k_idx1K/tile_k){DataCopyAsync(l0a_buf[next_buf],gm_aoffset_a_next,tile_m*tile_k);DataCopyAsync(l0b_buf[next_buf],gm_boffset_b_next,tile_k*tile_n);}// Cube单元在当前缓冲区上做矩阵乘加// MMAD矩阵乘加指令结果累加到L0CMMAD(l0c_buf,l0a_buf[cur_buf],l0b_buf[cur_buf],tile_m,tile_n,tile_k);// 等待异步预取完成和MMAD并行执行的PipeBarrier();}// 最终把L0C的结果搬回Global MemoryDataCopy(gm_c,l0c_buf,tile_m*tile_n);双缓冲的效果取决于搬运时间和计算时间的比例。如果搬运时间小于计算时间计算密集型双缓冲可以完全掩盖搬运延迟Cube单元始终处于忙碌状态。如果搬运时间大于计算时间访存密集型双缓冲只能部分缓解Cube单元还是会有空闲等待。GEMM的计算量是2MNK FLOPs搬运量是MK KN MN* sizeof(data_type) Bytes。当MNK足够大时计算量远大于搬运量GEMM是计算密集型的双缓冲效果很好。当矩阵较小时比如MNK128搬运量相对计算量较大GEMM变成访存密集型的双缓冲效果有限。不同矩阵尺寸下的性能分析下面是在Ascend 910上FP16 GEMM不同矩阵尺寸下的实测性能矩阵尺寸 (MxNxK)计算量 (GFLOPS)搬运量 (MB)算力强度 (FLOP/Byte)实测TFLOPS峰值利用率128x128x1280.0040.0981.7123%512x512x5120.2681.576.88521%1024x1024x10242.1476.2913.618045%4096x4096x4096137.4100.754.431078%8192x8192x81921099.5805.354.435088%算力强度FLOP/Byte是衡量计算密集程度的关键指标——每个字节的数据搬运对应多少次浮点运算。算力强度越高计算越密集双缓冲的效果越好。128x128x128的小矩阵峰值利用率只有3%因为计算量太小Cube单元大部分时间在等数据。4096以上的大矩阵利用率可以达到78%以上接近硬件峰值。这也就是为什么ops-blas对小矩阵做了专门的优化——当M和N较小时ops-blas会自动选择更小的Tiling参数比如tile_m16而不是512减少L0A/L0B的空闲空间同时把多个小矩阵打包到一个AI Core上串行处理避免频繁的核间同步开销。ops-blas和catlass的关系ops-blas和catlass都涉及GEMM实现但定位不同。ops-blas提供的是已编译好的高性能GEMM算子二进制开发者通过AscendCL的aclrtLaunchKernel或GE的算子调度来调用不能修改内部实现。catlass提供的是GEMM的Ascend C模板源码开发者可以基于模板定制Tiling策略、数据类型和搬运流程然后编译成自己的GEMM算子。简单说ops-blas是用catlass是改。如果你的GEMM需求和ops-blas提供的标准算子一致直接用ops-blas性能已经很好。如果你需要定制GEMM的实现比如特殊的矩阵排布、非标准的融合操作才需要用catlass从模板开始构建。使用前后效率对比以一个Transformer推理场景为例单层Attention的计算涉及3次GEMMQ/K/V投影和1次GEMM输出投影矩阵尺寸为batch_size x seq_len x hidden_dim。对比维度手写Tiling的Ascend C GEMMops-blas标准算子ops-blas Tiling调优开发时间3-5天0直接调用2小时参数调优4096x4096x4096性能280 TFLOPS310 TFLOPS340 TFLOPS1024x1024x1024性能120 TFLOPS180 TFLOPS195 TFLOPS128x128x128性能5 TFLOPS12 TFLOPS15 TFLOPSTransformer单层延迟4.2ms3.5ms3.1msops-blas标准算子在所有矩阵尺寸下都优于手写版本尤其是小矩阵场景——ops-blas内置的小矩阵优化策略不是手写Tiling能轻易复现的。经过Tiling调优后的ops-blas性能进一步提升4096尺寸达到340 TFLOPS接近Ascend 910的FP16理论峰值约400 TFLOPS的85%。结尾ops-blas是CANN算子体系里计算量最大的组件它的性能直接影响了上层所有依赖矩阵乘法的算子。理解GEMM的Tiling策略、双缓冲优化和算力强度对性能的影响有助于在模型部署时做出正确的优化决策——比如选择合适的batch size来让GEMM处于计算密集区域避免过小的矩阵尺寸导致NPU利用率低下。对于大多数开发者来说ops-blas的标准算子已经足够好需要极致性能的场景可以通过Tiling调优进一步提升。仓库地址https://atomgit.com/cann/ops-blas
