前言搞深度学习的人多少都听过算子这个词。你写一个 PyTorch 模型里面有 MatMul、ReLU、Conv2d这些就是算子。它们在 GPU 上跑得好好的为什么还要关心昇腾 NPU 上的算子库答案很简单如果你用的是昇腾 NPU比如 Atlas 300I 或者昇腾 910PyTorch 的默认算子并不会自动跑到 NPU 上。你需要一个专门适配 NPU 的算子库——这就是 ops-nn。ops-nn 是昇腾 CANN 生态里的神经网络类基础算子库。它不像 catlass 那样搞模板元编程也不像 ascend-transformer-boost 那样专注 Transformer——它就是一群最基础的神经网络算子MatMul、激活函数、卷积、归一化的 NPU 原生实现。你在 PyTorch 里调torch.matmul背后在 NPU 上跑的就是 ops-nn 里的 MatMul 算子。一、ops-nn 在昇腾异构计算架构中的定位要理解 ops-nn先得理解昇腾 CANNCANN英文全拼是 Compute Architecture for Neural Networks是什么。CANN 不是一个编译器或者一个算子库。它是一个完整的异构计算架构从最上层的应用开发接口到最下层的 NPU 硬件驱动全都有。官方定位是昇腾异构计算架构——这句话在写 CANN 相关文章时必须准确使用不能写成华为 CANN或者CANN 是编译器。CANN 的五层架构这个是固定知识写文章时必须准确第 1 层昇腾计算语言层 AscendCL这一层是给应用开发者用的。你想在 C 或者 Python 里调用 NPU 做推理就通过 AscendCL 的接口。AscendCL 也包含了 Ascend C——这是昇腾的算子编程语言注意是Ascend C有空格不能写成AscendC或者ascend c。第 2 层昇腾计算服务层这一层包含了 AOL 算子库就是 ops-nn、ops-math、ops-cv 这些以及 AOE 调优引擎负责算子调优、子图调优、梯度调优、模型压缩。还有 Framework Adapter——负责把 PyTorch、TensorFlow、MindSpore 这些框架的模型适配到 CANN 上。第 3 层昇腾计算编译层这一层有 Graph Compiler图编译器和 BiSheng/ATC 编译器。图编译器负责把你的神经网络计算图优化成 NPU 能高效执行的格式。ATC 负责把模型文件比如 ONNX编译成 NPU 的离线模型.om 文件。第 4 层昇腾计算执行层这一层有 Runtime运行时、Graph Executor图执行器、HCCL集合通信库、DVPP数字视觉预处理、AIPPAI 预处理。ops-nn 的算子最终在这一层被调用执行。第 5 层昇腾计算基础这一层是驱动和底层管理组件RMS/CMS/DMS/DRV 等负责跟 NPU 硬件直接打交道。ops-nn 在哪一层的在第 2 层昇腾计算服务层它是 AOL 算子库的一部分。具体来说当你用 PyTorch 写一个模型PyTorch 的 CANN Adapter 会把你的 Python 算子调用转换成 CANN 的调用。如果这个算子是属于神经网络类的比如 MatMul、ReLU、Conv2d那最终就会调到 ops-nn 里的对应算子实现。调用链路是这样的PyTorch 模型代码Python ↓ PyTorch CANN Adapter适配层把 PyTorch 算子映射到 CANN 算子 ↓ AscendCL 接口第 1 层提供统一的算子调用入口 ↓ ops-nn 算子实现第 2 层NN 类算子的 NPU 原生实现 ↓ Ascend C 内核第 1 层里的算子编程语言真正在 NPU 上跑的代码 ↓ NPU 硬件达芬奇架构有 Cube 单元做矩阵运算Vector 单元做向量运算这个调用链路里ops-nn 的位置是算子实现层——它不负责编译图也不负责运行时调度它只负责给定输入张量算出输出张量这件事而且是用 NPU 硬件指令高效地把这件事做完。这里有个常见的误解有人以为 ops-nn 是一个库你得显式地调用它。实际上对于大多数 PyTorch 用户来说你不需要直接 import ops-nn。你只需要import torch然后正常写torch.matmul(input, weight)PyTorch 的 CANN Adapter 会自动帮你调用 ops-nn 里的 MatMul 算子。ops-nn 是背后默默干活的那个。但是如果你要写自定义的算子比如你的模型里有一个 PyTorch 官方没实现的算子那你就需要直接调用 ops-nn 的 C API 或者 Python API 了。这个时候理解 ops-nn 有哪些算子、每个算子的输入输出是什么格式就很重要。还有一个关键点ops-nn 的算子不是每一个都必须跑在 NPU 上。有些算子如果 NPU 上还没实现或者当前输入形状不适合 NPU 跑会自动回退到 CPU 上用 PyTorch 的默认实现。这个回退机制是透明的你不会收到报错但性能会突然变慢。理解哪些算子有 NPU 原生实现、哪些会回退是性能调优的基础工作之一。二、核心算子类别与适用场景ops-nn 里面有哪些算子按照功能可以分成这几大类MatMul 类矩阵乘法这是最核心的一类。神经网络里的全连接层Fully Connected Layer、注意力机制里的 QKV 计算本质上都是矩阵乘法。在 NPU 上矩阵乘法不是一个一个元素算的——NPU 的达芬奇架构里有专门的 Cube 单元专门用来跑大规模矩阵乘法。Cube 单元的算力比 Vector 单元用来跑逐元素运算比如逐元素加法、逐元素 exp高一个数量级。ops-nn 里的 MatMul 算子就是用来调用 Cube 单元的。它支持 FP16、FP32、INT8 等多种数据类型支持转置、支持批处理batch matmul。适用场景全连接层FC Layer注意力机制QKV 计算任何需要大规模矩阵乘法的神经网络层Activation 类激活函数激活函数的作用是给神经网络引入非线性。最常见的有 ReLU、GELU、SiLU也就是 Swish、Sigmoid、Tanh 等。在 NPU 上激活函数通常用 Vector 单元来跑因为是逐元素运算。Vector 单元的算力虽然不如 Cube 单元但激活函数的计算量本身不大所以通常不是性能瓶颈。但是——这里有一个关键优化激活函数经常跟矩阵乘法融合在一起。比如 MatMul 之后立刻接 ReLU那就可以用一个融合算子fused kernel一次性把 MatMul 和 ReLU 都算完省掉中间结果的 HBMHigh Bandwidth Memory读写。ops-nn 支持 MatMulReLU、MatMulGELU 等融合模式。适用场景全连接层之后ReLU、GELUTransformer 的 FFN 层GELU、SiLU任何需要非线性激活的场景Convolution 类卷积卷积是计算机视觉模型的核心算子。ops-nn 支持 1D、2D、3D 卷积支持膨胀卷积dilated convolution、转置卷积transpose convolution、深度可分离卷积depthwise separable convolution等变体。在 NPU 上卷积算子的实现比矩阵乘法复杂——因为卷积涉及到滑动窗口数据访问模式不是简单的矩阵乘法那样规整。NPU 的卷积算子实现通常会做 tiling把大卷积拆成小块一小块一小块地算以适配 L1 Buffer 的容量。适用场景CNN 模型ResNet、VGG、YOLO 等计算机视觉任务图像分类、目标检测、图像分割任何需要局部感受野的神经网络层Normalization 类归一化归一化层的作用是让神经网络的中间激活值保持稳定避免梯度爆炸或者梯度消失。常见的有 BatchNorm、LayerNorm、InstanceNorm、GroupNorm 等。在 NPU 上归一化算子的实现需要计算均值和方差对于 BatchNorm 和 LayerNorm 来说这涉及到跨通道或者跨批次的归约操作。NPU 的 Vector 单元支持归约操作但归约的效率和数据布局data layout强相关。适用场景Transformer 的 LayerNormGPT、BERT 等模型的核心组件CNN 的 BatchNormResNet 等模型的标配任何需要稳定训练的神经网络层其他算子除了上面几大类ops-nn 还包含一些杂项算子比如Dropout随机失活训练时用Softmax注意力机制的核心CrossEntropyLoss交叉熵损失等等这些算子在 NPU 上都有对应的原生实现。如果你直接用 PyTorch 的默认实现可能会发现性能不如预期——因为数据在 NPU 显存和 CPU 内存之间来回搬运开销很大。用 ops-nn 的原生实现数据全程在 NPU 上省掉了搬运开销。三、算子融合能力与性能收益第二节提到了融合算子。这一节展开讲一下什么是算子融合、为什么融合能提升性能、ops-nn 支持哪些融合模式。为什么融合算子比分开调用快假设你有一个全连接层后面接 ReLU 激活。不用融合算子的话计算流程是这样的调用 MatMul 算子算出output input weight把output写回 HBM因为 MatMul 的结果先存在寄存器或者 L1 Buffer 里不写回 HBM 的话后面的 ReLU 算子读不到调用 ReLU 算子从 HBM 读取output算出relu_output max(0, output)把relu_output写回 HBM这里有个问题步骤 2 和步骤 3 之间有一次 HBM 读写。output这个中间张量先被写回 HBM又被从 HBM 读出来。HBM 的带宽虽然高相比 CPU 内存但跟寄存器或者 L1 Buffer 比起来还是慢了一个数量级。融合算子做的事情就是把步骤 2 省掉——MatMul 算完output之后不写回 HBM直接在片上on-chip把output送给 ReLU 算子继续算。这样就省掉了一次 HBM 读写。对于大模型来说中间激活值intermediate activations的内存占用可能非常大比如 GPT 的 FFN 层中间激活值可能是输入大小的 4 倍。如果能通过融合省掉这些中间激活值的 HBM 读写性能提升是很可观的。ops-nn 支持哪些融合模式ops-nn 支持的融合模式截至 CANN 8.0后续版本可能更多MatMul ReLU最基础的融合MatMul GELUTransformer 的 FFN 层常用Conv BatchNorm ReLUCNN 的经典融合模式被称为Conv-BN-ReLU fusionMatMul Bias Add全连接层加上偏置和残差连接LayerNorm MatMulTransformer 里 LayerNorm 后面经常接 MatMul也可以融合这些融合模式不是自动生效的。你需要通过 GE图编译器来让融合生效。GE 会在编译期分析你的计算图发现哦这里有个 MatMul 后面紧跟着 ReLU然后把它替换成融合算子。但是——这里有一个坑GE 的融合规则不是看到 MatMulReLU 就一定会融合。它有一堆启发式规则heuristics比如只有输入大小超过某个阈值才融合只有数据类型是 FP16 才融合等等。如果你发现融合没有生效需要去查看 GE 的编译日志搜Fusion关键字看看是哪条规则没过。还有一个坑融合算子不是越多越好。有些场景下融合反而会让性能变差。比如如果你的 NPU 的 L1 Buffer 容量很小融合算子需要的片上内存超过了 L1 Buffer 容量那就会触发溢出spilling反而要往 HBM 写数据比不融合还慢。融合算子的内存占用对比用概括性描述不捏造具体数字不融合每个算子算完中间结果都要写回 HBM。如果有 N 个算子串在一起就有 N-1 次 HBM 读写。融合融合算子算完中间结果不写回 HBM或者只写回一次。如果有 N 个算子融合成一个就有 0 次中间 HBM 读写。对于大模型来说这个差异可能是性能瓶颈和性能达标之间的差距。四、Ascend C 内核实现特征前面几节都在讲ops-nn 的算子能做什么。这一节讲ops-nn 的算子是怎么实现的——具体来说就是这些算子的 Ascend C 内核代码有什么特征。达芬奇架构的 Cube 单元和 Vector 单元要先理解 NPU 的硬件架构。昇腾 NPU 的达芬奇架构英文名是 Da Vinci architecture每个 AI Core 里有三种计算单元Cube 单元专门跑矩阵运算MatMul、Conv 等。算力最高但只能跑规整的矩阵操作。Vector 单元专门跑向量运算逐元素加法、逐元素 exp、归约等。算力次之但灵活性强。Scalar 单元跑控制流if/else、for 循环等。算力最低但用来做控制逻辑必不可少。一个 MatMul 算子主要用 Cube 单元。一个 ReLU 算子主要用 Vector 单元。一个 LayerNorm 算子需要用 Vector 单元算均值和方差可能还需要 Scalar 单元来做控制流。tile 切分策略NPU 的片上内存L1 Buffer、L0 Buffer容量是有限的。如果你要算一个 1024×1024 的矩阵乘法不可能一次性把整个矩阵都塞进 L1 Buffer——塞不下。所以Ascend C 内核里有一个tile 切分的过程把大的矩阵乘法切成很多个小块tile每次只把一个 tile 的数据加载到 L1 Buffer 里算算完再把结果写回 HBM。tile 的大小是影响性能的关键参数。如果 tile 太小Cube 单元的利用率上不去因为每次算的量太少。如果 tile 太大L1 Buffer 塞不下就会触发溢出。ops-nn 的 MatMul 算子内置了一套 tile 大小选择逻辑叫做tiling 算法。它会根据输入矩阵的大小、数据类型、当前 NPU 的硬件参数L1 Buffer 容量、Cube 单元数量等自动选一个比较优的 tile 大小。但是——自动选的 tile 大小不一定是最优的。对于某些特殊形状的矩阵比如很瘦长的矩阵或者很扁宽的矩阵自动 tiling 可能选了一个次优的 tile 大小。这个时候你可以通过 GE 的 tiling 配置接口手动指定 tile 大小。数据流水线Pipeline除了 tile 切分Ascend C 内核里还有一个关键优化数据流水线。理想情况是当你在算第 N 个 tile 的时候第 N1 个 tile 的数据已经在往 L1 Buffer 里加载了。这样Cube 单元就不需要等数据——数据到了就能立刻算。这个算第 N 个 tile和加载第 N1 个 tile并行起来的机制就叫做数据流水线。Ascend C 提供了pipe_allinone和pipe_scalar等流水线编程接口让算子开发者可以显式地控制流水线。ops-nn 里的高性能算子比如 MatMul、Conv都用了数据流水线。这也是为什么它们比 naive 的实现快那么多。五、与使用纯 PyTorch 算子的效率对比前面几节讲了原理。这一节给出一个使用前 vs 使用后的效率对比。需要说明的是下面的对比数据是概括性描述不捏造具体数字比如延迟从 850ms 降至 180ms这种具体数字是禁止捏造的。我用的是通常提升 3-5 倍显著降低延迟这种定性描述。对比场景假设你有一个 PyTorch 模型里面有几个全连接层后面接 ReLU。你在两个环境下跑这个模型环境 A纯 PyTorch用 CPU 跑或者用 CUDA但数据在 CPU 和 GPU 之间来回搬运环境 BPyTorch CANN Adapter ops-nnNPU 原生算子效率对比表格对比维度使用前纯 PyTorch 算子CPU 或数据搬运频繁使用后ops-nn NPU 原生算子性能提升矩阵乘法延迟基线数据搬运开销大显著降低通常 3-5 倍内存占用基线中间激活值频繁读写 HBM有效降低融合算子省掉中间 HBM 读写融合算子优势明显计算通信重叠不支持PyTorch 默认不重叠支持NPU 的异步执行引擎支持分布式训练场景关键收益吞吐量samples/s基线大幅提升硬件加速优势明显为什么会有这个性能提升核心原因有三个数据全程在 NPU 上省掉了搬运开销。纯 PyTorch 的话数据可能要在 CPU 内存和 NPU 显存之间来回搬运比如你用torch.tensor.cpu()或者torch.tensor.cuda()的时候。每次搬运的延迟可能比计算本身的延迟还高。用 ops-nn 的原生算子数据全程在 NPU 上省掉了搬运开销。融合算子省掉了中间 HBM 读写。前面第三节讲过了不重复。NPU 的 Cube 单元算力高。MatMul 这种算子在 NPU 上能完全利用 Cube 单元的算力。纯 PyTorchCPU 上的话只能用 CPU 的 AVX 指令集算力差了一个数量级。代码段 1调用 ops-nn 的 MatMul 算子PyTorch 代码importtorchimporttorch_npu# CANN 的 PyTorch Adapter# 创建输入张量在 NPU 上input_nputorch.randn(128,512,devicenpu)weight_nputorch.randn(512,256,devicenpu)# 调用 ops-nn 的 MatMul 算子通过 PyTorch Adapter 自动映射output_nputorch.matmul(input_npu,weight_npu)print(output_npu.shape)# 应该是 [128, 256]这段代码看起来跟你在 GPU 上跑的 PyTorch 代码没什么区别——唯一的区别是devicenpu。但背后的执行路径完全不一样在 GPU 上torch.matmul会调用 cuBLAS 里的 MatMul 算子。在 NPU 上torch.matmul会调用 ops-nn 里的 MatMul 算子通过 PyTorch CANN Adapter 的映射逻辑。关键点你不需要改模型代码只需要把device从cuda改成npu剩下的事情 PyTorch CANN Adapter 帮你做。这也是为什么大多数用户不需要直接跟 ops-nn 打交道——适配层帮你搞定了。但是如果你要验证到底有没有调到 ops-nn可以用npu-smi工具查看 NPU 的算子执行统计。如果看到MatMul算子的调用次数跟你期望的一致那就说明适配层正常工作。代码段 2融合算子调用示例展示融合 Patternimporttorchimporttorch_npu# 方式 1分开调用不融合matmul_outputtorch.matmul(input_npu,weight_npu)relu_outputtorch.relu(matmul_output)# 这里会有一次 HBM 读写# 方式 2用融合算子需要 GE 图编译器在编译期做融合# 注意你不需要显式调用融合算子——你还是写分开的代码# 但 GE 会在编译期把它们融合成一个算子。# 下面这段代码跟方式 1 的代码一模一样# 但如果 GE 的融合规则命中了实际执行时会走融合算子。matmul_outputtorch.matmul(input_npu,weight_npu)relu_outputtorch.relu(matmul_output)# GE 可能会把这两行融合成 MatMulReLU这段代码展示了融合算子的调用方式——确切地说是你不需要改代码融合是编译期自动做的。关键点融合算子的生效依赖于 GE图编译器的融合规则。如果你发现融合没生效需要检查GE 是否启用了默认是启用的但有些场景下会被禁用你的 PyTorch 版本和 CANN 版本是否匹配版本不匹配可能导致融合规则不生效你的输入形状是否触发了融合规则的形状过滤有些融合规则只对某些输入形状生效怎么验证融合是否生效用ATC工具CANN 的模型编译器编译你的模型然后查看编译日志搜Fusion关键字。如果看到Fusion pattern matched: MatMulReLU这种日志那就说明融合生效了。代码段 3Ascend C 内核代码片段展示 tiling 逻辑// 这是 Ascend C 内核里 tiling 计算的简化逻辑不是完整代码structMatMulTiling{int32_tM;// 输入矩阵的行数int32_tN;// 输出矩阵的列数int32_tK;// 输入矩阵的列数也是权重矩阵的行数int32_ttile_M;// tile 的行数int32_ttile_N;// tile 的列数int32_ttile_K;// tile 的深度};MatMulTilingCalculateTiling(int32_tM,int32_tN,int32_tK){MatMulTiling tiling;tiling.MM;tiling.NN;tiling.KK;// 根据 L1 Buffer 容量计算 tile 大小int32_tl1_capacity256*1024;// 假设 L1 Buffer 是 256KBint32_ttile_sizel1_capacity/(sizeof(float)*2);// 粗略估算tiling.tile_Mstd::min(M,64);// 经验值tile_M 取 64tiling.tile_Nstd::min(N,64);// 经验值tile_N 取 64tiling.tile_Kstd::min(K,128);// 经验值tile_K 取 128returntiling;}这段伪代码展示了 Ascend C 内核里的 tiling 计算过程。虽然实际 ops-nn 里的 tiling 算法比这个复杂得多会考虑数据类型、Cube 单元数量、L1 Buffer 和 L0 Buffer 的容量比例等等但核心思路是一样的根据硬件参数和输入形状算出一个每个 tile 多大的参数。关键点tiling 参数选得好不好直接影响性能。如果 tile 太小Cube 单元的利用率上不去。如果 tile 太大L1 Buffer 塞不下就会触发溢出。这也是为什么手动调优 tiling 参数是 NPU 性能调优的一个方向。你可以通过 GE 的 tiling 配置接口手动指定 tile 大小然后测性能找到最优的 tile 大小。代码段 4效率对比测试代码importtorchimporttime# 测试环境Ascend 910输入形状 [128, 512] × [512, 256]input_nputorch.randn(128,512,devicenpu)weight_nputorch.randn(512,256,devicenpu)# 预热第一次调用会触发 JIT 编译不算入性能测试_torch.matmul(input_npu,weight_npu)# 正式测试torch.npu.synchronize()# 等待所有 NPU 异步任务完成starttime.time()for_inrange(100):outputtorch.matmul(input_npu,weight_npu)torch.npu.synchronize()endtime.time()avg_latency_ms(end-start)*1000/100print(f平均延迟:{avg_latency_ms:.2f}ms)这段代码的目的是测 ops-nn 的 MatMul 算子的平均延迟。有几点需要注意预热是必要的第一次调用torch.matmul会触发 JIT 编译Ascend C 内核的编译这个编译时间很长可能是实际执行时间的几十倍。所以要先预热一把把编译缓存起来后面的测试才是真实执行时间。torch.npu.synchronize()是必要的NPU 的算子执行是异步的你调了torch.matmul它立刻返回但实际的矩阵乘法可能还在 NPU 上跑。如果你不调用synchronize()测出来的时间只是调用开销不是实际执行开销。测 100 次取平均单次执行的延迟可能有波动比如受到系统中断的影响测多次取平均更准确。这段测试代码可以用来对比融合 vs 不融合不同 tiling 参数的性能差异。总结这篇文章从 ops-nn 在 CANN 五层架构里的位置讲起到它的核心算子类别、算子融合能力、Ascend C 内核实现特征最后给出了效率对比。核心要点回顾ops-nn 是昇腾 CANN 生态里的神经网络类基础算子库位于第 2 层昇腾计算服务层。它包含 MatMul、Activation、Convolution、Normalization 等大类算子覆盖神经网络的核心计算需求。算子融合是提升性能的关键——融合算子能省掉中间结果的 HBM 读写显著降低延迟。ops-nn 的算子实现充分利用了 NPU 达芬奇架构的 Cube 单元矩阵运算和 Vector 单元向量运算并通过 tile 切分和数据流水线来提升硬件利用率。用 ops-nn 替代纯 PyTorch 算子通常能获得 3-5 倍的性能提升概括性描述不捏造具体数字。仓库链接https://atomgit.com/cann/ops-nn