前言

昇腾CANN作为昇腾异构计算架构，昇腾CANN作为昇腾异构计算架构，ops-nn是其中最核心的神经网络算子库。很多人以为神经网络算子就是卷积、池化、激活函数这些名字听起来很熟悉的东西，但真正上手做算子开发的时候才发现，同样的卷积运算，在CPU上写和在NPU上写完全是两码事。CPU上你可以用MKL或者手写循环，NPU上你需要理解什么是Cube单元、什么是Vector单元、为什么需要把卷积转换成矩阵乘法、什么是Im2Col和Col2Im。

ops-nn这个仓库，本质上是把PyTorch torch.nn模块里那些常见的层，用昇腾NPU的硬件特性重新实现了一遍。卷积、全连接、BatchNorm、LayerNorm、Softmax、Pooling、激活函数，这些算子都在ops-nn里有对应的实现。但ops-nn不是简单的翻译，它针对昇腾NPU的达芬奇架构做了深度优化，包括算子融合、内存复用、流水线并行等技术。

一、ops-nn在CANN架构中的定位

1.1 从一个简单的问题开始

先回答一个问题：为什么需要ops-nn？PyTorch不是已经有torch.nn了吗？

答案在于硬件差异。PyTorch的nn模块主要针对GPU（CUDA）优化，而昇腾NPU的硬件架构跟GPU完全不同。GPU有大量的CUDA核心，每个核心可以做独立的标量运算。昇腾NPU有Cube单元（专门做矩阵乘法）和Vector单元（做向量运算），这两种单元的调度方式、内存访问模式、计算效率都跟GPU不一样。

同样的卷积运算，在GPU上可能用cuDNN的卷积实现，在NPU上用ops-nn的实现。两者的性能差异可能达到2-5倍，这不是因为NPU比GPU快，而是因为算子实现是否充分利用了硬件特性。

# 为什么ops-nn比直接用PyTorch快？

# 方式一：PyTorch + GPU后端
import torch
import torch.nn as nn

# 创建卷积层
conv = nn.Conv2d(3, 64, kernel_size=3, padding=1).cuda()
x = torch.randn(1, 3, 224, 224).cuda()

# 执行卷积
output = conv(x)  # 这个卷积在GPU上用cuDNN实现

# 方式二：昇腾NPU + ops-nn
import torch_npu  # 昇腾PyTorch扩展

conv_npu = nn.Conv2d(3, 64, kernel_size=3, padding=1).npu()
x_npu = torch.randn(1, 3, 224, 224).npu()

output_npu = conv_npu(x_npu)  # 这个卷积在NPU上用ops-nn实现

# 表面上看代码一模一样，但底层实现完全不同。
# GPU版本调用cuDNN的卷积库，NPU版本调用ops-nn的卷积算子。
# ops-nn的卷积算子针对昇腾的Cube单元做了优化，
# 把卷积转换成矩阵乘法（Im2Col），然后用Cube单元高效计算。
# 同时，ops-nn还做了内存复用和流水线并行优化。

1.2 CANN五层架构中的位置

CANN的架构分五层，ops-nn在第二层，属于算子库层。完整架构是这样的：

第一层是AscendCL，这是给应用开发者用的接口层。你用C语言的API调用推理、训练、算子执行等功能。

第二层是算子库层，包括ops-nn（神经网络算子）、ops-math（数学算子）、ops-blas（线性代数算子）、ops-cv（视觉算子）、HCCL（通信算子）等。ops-nn是其中最常用的，因为它包含了深度学习90%的算子。

第三层是编译层，把你的模型转换成NPU能执行的形式。GE（图引擎）在这一层做算子融合、内存规划等优化。

第四层是执行层，Runtime负责设备管理、任务调度，HCCL负责分布式通信。

第五层是基础层，驱动、固件这些底层组件。

# ops-nn与其他算子库的关系

# ops-nn：神经网络算子（卷积、全连接、归一化、激活函数）
# ops-math：数学算子（exp, log, sqrt, erf, gamma等超越函数）
# ops-blas：线性代数算子（矩阵乘法、矩阵分解）
# ops-cv：视觉算子（resize, normalize, color jitter）
# ops-transformer：Transformer专用算子（FlashAttention, MoE等）

# 为什么要分这么多算子库？
# 因为不同类型的算子有不同的计算特征和优化策略。
# ops-nn里的卷积需要Im2Col转换，用Cube单元加速。
# ops-math里的超越函数用Vector单元加速，每个像素独立计算。
# ops-blas里的矩阵乘法直接用Cube单元，不需要额外转换。
# 分库管理可以让每个算子库专注一类算子的优化。

二、ops-nn支持的核心算子

2.1 卷积类算子

卷积是神经网络的基础，ops-nn支持多种卷积变体。标准卷积（Conv2d）、深度可分离卷积（DepthwiseConv2d）、转置卷积（ConvTranspose2d）、分组卷积（GroupedConv2d）都在支持范围内。

卷积在NPU上的实现有个关键优化：Im2Col。原始的卷积运算是滑动窗口，每次取一小块输入，乘以卷积核，输出一个值。这种操作在GPU上可以做并行，但在NPU上效率不高，因为每次滑动窗口都要访问不连续的内存地址。

Im2Col把滑动窗口转换成矩阵乘法。先把输入数据重排成一个矩阵，每一行对应一个滑动窗口位置。然后把卷积核也展开成矩阵。这样卷积运算就变成了矩阵乘法，可以用Cube单元高效计算。

# Im2Col的原理

import numpy as np

def im2col_simple(input_data, kernel_h, kernel_w, stride=1, padding=0):
    """
    Im2Col：把卷积转换成矩阵乘法
    
    原始卷积：
    - 输入：(N, C, H, W)
    - 卷积核：(K, C, KH, KW)
    - 输出：(N, K, OH, OW)
    
    Im2Col后：
    - 输入矩阵：(OH*OW, C*KH*KW)  # 每行是一个滑动窗口
    - 卷积核矩阵：(K, C*KH*KW)     # 每行是一个输出通道的核
    - 输出矩阵：(OH*OW, K)         # 矩阵乘法结果
    """
    N, C, H, W = input_data.shape
    
    # 计算输出尺寸
    OH = (H + 2*padding - kernel_h) // stride + 1
    OW = (W + 2*padding - kernel_w) // stride + 1
    
    # 创建输出矩阵
    col = np.zeros((OH * OW, C * kernel_h * kernel_w))
    
    # 填充矩阵（这里是简化版，实际实现更复杂）
    idx = 0
    for i in range(OH):
        for j in range(OW):
            # 提取滑动窗口
            h_start = i * stride
            h_end = h_start + kernel_h
            w_start = j * stride
            w_end = w_start + kernel_w
            
            patch = input_data[:, :, h_start:h_end, w_start:w_end]
            col[idx, :] = patch.flatten()
            idx += 1
    
    return col

# Im2Col看起来很浪费，因为输入数据被复制了很多次。
# 原始输入大小是N*C*H*W，Im2Col后变成N*OH*OW*C*KH*KW，大了好几倍。
# 但这个额外开销是值得的，因为矩阵乘法可以用Cube单元加速，
# Cube单元一次可以算16x16的矩阵块，效率远高于滑动窗口逐点计算。
# 实际上，ops-nn会用更智能的内存布局减少数据复制开销。

2.2 归一化算子

BatchNorm和LayerNorm是深度学习里最常用的归一化层。ops-nn对这两种归一化都做了优化实现。

BatchNorm的问题在于训练和推理的行为不同。训练时用batch统计量，推理时用移动平均统计量。ops-nn实现了融合BatchNorm，把卷积+BatchNorm+ReLU融合成一个算子，减少内存访问次数。

LayerNorm在Transformer里用得很多。ops-nn的LayerNorm实现针对大hidden_dim做了优化，用Vector单元并行计算均值和方差。

# 融合BatchNorm的原理

import torch

class FusedConvBNReLU(torch.nn.Module):
    """
    融合卷积 + BatchNorm + ReLU
    
    传统做法：
    x -> Conv -> y -> BN -> z -> ReLU -> out
    
    融合做法：
    x -> FusedConvBNReLU -> out
    
    融合后的效果：
    1. 只需一次内存读写（原来要三次）
    2. 减少中间结果存储
    3. 可以用流水线并行
    """
    
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
        super().__init__()
        self.conv = torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)
        self.bn = torch.nn.BatchNorm2d(out_channels)
        self.relu = torch.nn.ReLU(inplace=True)
    
    def forward(self, x):
        # 方式一：逐层调用（未融合）
        # x = self.conv(x)
        # x = self.bn(x)
        # x = self.relu(x)
        # return x
        
        # 方式二：融合调用（ops-nn提供的接口）
        return torch.ops.npu.fused_conv_bn_relu(
            x, 
            self.conv.weight, 
            self.conv.bias,
            self.bn.weight,
            self.bn.bias,
            self.bn.running_mean,
            self.bn.running_var,
            stride=self.conv.stride,
            padding=self.conv.padding
        )

# 为什么要融合？
# 因为内存访问是性能瓶颈。每个算子执行完，结果要写回HBM，
# 下一个算子再从HBM读出来。HBM带宽有限，这个过程很慢。
# 融合后，中间结果不需要写回HBM，直接在片上缓存传递给下一个算子。
# 对于BatchNorm，还可以进一步优化：推理时把BN的参数融合进卷积权重，
# 这样推理阶段完全跳过BN计算。

2.3 激活函数

ReLU、GELU、SiLU、Swish、HardSwish，这些激活函数ops-nn都支持。激活函数的计算相对简单，但有一个优化点：融合。

单独调用激活函数需要读写一次内存。如果跟前面的卷积或矩阵乘法融合，激活函数的计算可以在输出结果写回HBM之前完成，省掉一次内存访问。

# 激活函数融合示例

import torch
import torch_npu

# 未融合：逐算子调用
def forward_unfused(x, weight):
    x = torch.matmul(x, weight)
    x = torch.nn.functional.gelu(x)  # 单独调用GELU
    return x

# 融合：一次调用
def forward_fused(x, weight):
    return torch.ops.npu.fused_matmul_gelu(x, weight)

# 融合激活函数的收益有两点。
# 第一，省掉一次HBM读写。GELU原本要读输入、写输出，
# 融合后输入在片上缓存，不需要从HBM读。
# 第二，可以利用Vector单元的并行能力。
# GELU的计算是逐元素的，Vector单元一次可以处理多个元素。
# 如果GELU单独调用，要等前面的算子执行完才能开始。
# 融合后可以在矩阵乘法的输出写回之前，并行计算激活函数。

三、算子融合技术

3.1 什么是算子融合

算子融合是提升NPU性能的核心技术。原理很简单：如果两个算子可以合并成一个更大的算子执行，就能省掉中间结果的内存读写。

举个具体例子。假设你要做卷积、BatchNorm、ReLU三个操作。传统做法是分别调用三个算子，每个算子执行完把结果写回内存，下一个算子再读出来。融合做法是把这三个算子合并成一个算子，输入数据进来，一次计算完，输出结果直接写回内存。

# 算子融合的收益分析

import time

def benchmark_fusion():
    """
    对比融合前后的性能
    """
    import torch_npu
    
    batch_size = 32
    channels = 64
    height, width = 224, 224
    
    x = torch.randn(batch_size, channels, height, width).npu()
    conv_weight = torch.randn(128, 64, 3, 3).npu()
    bn_weight = torch.randn(128).npu()
    bn_bias = torch.randn(128).npu()
    
    # 未融合版本
    torch.npu.synchronize()
    start = time.perf_counter()
    
    for _ in range(100):
        out = torch.nn.functional.conv2d(x, conv_weight, padding=1)
        out = torch.nn.functional.batch_norm(out, torch.zeros(128).npu(), torch.ones(128).npu(), bn_weight, bn_bias)
        out = torch.nn.functional.relu(out)
    
    torch.npu.synchronize()
    unfused_time = time.perf_counter() - start
    
    # 融合版本
    torch.npu.synchronize()
    start = time.perf_counter()
    
    for _ in range(100):
        out = torch.ops.npu.fused_conv_bn_relu(x, conv_weight, bn_weight, bn_bias)
    
    torch.npu.synchronize()
    fused_time = time.perf_counter() - start
    
    print(f"未融合: {unfused_time/100*1000:.2f} ms/次")
    print(f"融合后: {fused_time/100*1000:.2f} ms/次")
    print(f"加速比: {unfused_time/fused_time:.2f}x")

# 融合能带来2-3倍的性能提升，主要来自三点。
# 第一，减少HBM访问次数。原来要写3次中间结果，融合后只写1次。
# 第二，减少kernel launch开销。每次调用算子都要从CPU发起，
# 融合后只发起一次。
# 第三，更好的缓存利用。融合算子可以把中间结果留在片上缓存，
# 不需要写回HBM再读回来。

3.2 支持的融合模式

ops-nn支持多种融合模式。常见的有Conv+BN+ReLU、Conv+ReLU、Linear+ReLU、MatMul+GELU、Softmax+Dropout等。融合模式的选择通常由GE图引擎自动完成，但有些场景你也可以手动指定。

使用前 vs 使用后：ops-nn带来的效率对比

指标 使用前（PyTorch默认） 使用后（ops-nn优化） 提升效果
ResNet-50推理延迟 约15ms/张 约5ms/张 显著降低
卷积+BN+ReLU内存访问 3次读写 1次读写 大幅减少
BatchNorm训练稳定性 需手动调参 内置优化 更稳定
激活函数计算效率 CPU/GPU混合 全NPU加速 明显提升
算子融合覆盖率 需手动实现 自动融合 更高效率

ops-nn的效率提升主要来自三个方面。第一，针对昇腾NPU硬件特性的深度优化，包括Im2Col、内存复用、流水线并行等技术。第二，算子融合减少了中间结果的内存访问，这是性能瓶颈所在。第三，自动化的融合策略选择，让用户无需手动调优就能获得接近最优的性能。

四、开发自定义算子

4.1 什么时候需要自定义算子

ops-nn已经覆盖了大部分常用算子，但有些场景你可能需要自己写算子。比如新的激活函数、特殊的归一化方法、自定义的注意力机制等。

开发自定义算子需要用Ascend C语言，这是昇腾提供的算子编程语言。Ascend C对C++做了扩展，可以直接访问NPU的硬件单元。

# 自定义算子的开发流程

custom_op_workflow = """
自定义算子开发步骤：

第一步：定义算子接口
- 在metadef中定义算子的输入、输出、属性
- 描述算子的数学语义
- 指定支持的数据类型

第二步：用Ascend C实现算子
- 编写计算逻辑
- 调用Cube或Vector单元
- 处理内存布局

第三步：注册算子到框架
- 用torch_npu注册PyTorch扩展
- 定义forward和backward

第四步：测试与调优
- 与参考实现对比精度
- 用profiler分析性能瓶颈
- 优化内存访问模式
"""

# 为什么要自定义算子而不是组合现有算子？
# 因为组合现有算子会有额外的内存访问开销。
# 比如你实现了一个新的激活函数，用现有的逐元素运算组合，
# 中间结果要写回HBM再读出来。
# 如果用自定义算子，可以在片上缓存中完成计算，
# 不需要中间结果的内存访问。
# 性能差异可能达到2-3倍。

五、性能调优建议

5.1 选择合适的数据类型

ops-nn支持float32、float16、bfloat16等数据类型。float16是最常用的，因为昇腾NPU对float16有硬件加速。但float16精度有限，有些场景需要用float32。

# 数据类型选择建议

dtype_guide = """
数据类型选择：

float16（推荐大部分场景）：
- 推理：几乎都可以用
- 训练：大部分模型可以用
- 优点：速度快、内存占用少
- 缺点：精度有限，可能出现数值溢出

float32：
- 训练：大模型、高精度要求场景
- 推理：精度敏感的场景（如医疗影像）
- 优点：精度高
- 缺点：速度慢、内存占用大

bfloat16：
- 训练：大模型训练
- 优点：比float16精度高，不易溢出
- 缺点：部分老硬件不支持

建议：
1. 先用float16试试，如果精度不够再换float32
2. 训练大模型用bfloat16
3. 推理场景优先用float16
"""

# 数据类型的选择影响性能和精度。
# float16的计算速度比float32快2倍左右，内存占用也减半。
# 但float16只有10位尾数，精度有限。
# 昇腾NPU的Cube单元对float16矩阵乘法有特殊优化，
# 一次可以算更大的矩阵块。
# 所以如果精度允许，优先用float16。

5.2 利用自动融合

ops-nn的算子融合通常由GE图引擎自动完成。你不需要手动指定哪些算子要融合，GE会分析计算图，自动选择最优的融合策略。

但有些情况自动融合效果不好，需要手动干预。比如你有自定义算子，GE不知道怎么融合。这时候可以用torch_npu的接口手动指定融合模式。

六、总结

ops-nn是昇腾CANN生态里最核心的算子库之一，它把深度学习常用的神经网络算子用昇腾NPU的硬件特性重新实现了一遍。卷积、归一化、激活函数这些算子在ops-nn里都有优化版本，性能通常比PyTorch默认实现快2-3倍。

算子融合是ops-nn的核心优化技术，把多个算子合并成一个更大的算子执行，能省掉中间结果的内存访问。Conv+BN+ReLU、MatMul+GELU这些融合模式ops-nn都支持，大部分场景下GE会自动选择最优的融合策略。

仓库链接：https://atomgit.com/cann/ops-nn

Logo
CANN开发者社区

CANN开发者社区旨在汇聚广大开发者，围绕CANN架构重构、算子开发、部署应用优化等核心方向，展开深度交流与思想碰撞，携手共同促进CANN开放生态突破！

更多推荐

cover

考研英语作文模板万能句型|pdf|电子版|资料已整理

avatar CANN开发者社区
cover

昇腾CANN集合通信库HCCL架构解析：分布式训练通信原语与拓扑感知优化技术

avatar CANN开发者社区
cover

昇腾CANN数学函数库ops-math深度解析：超越函数在NPU上的高效实现技术

avatar CANN开发者社区