☰
ops-nn:用软硬协同优化神经网络算子的性能
2026/10/20 2:48:24 网站建设 项目流程

我早期碰过一个很现实的问题:同一段网络里的卷积和归一化,框架自带的实现跑一次只要几百微秒,而我自己老老实实按算法流程写的算子,要多花好几倍时间。函数逻辑一模一样,数值结果也一模一样,性能却差出两个量级。后来花了不少时间才想明白,神经网络算子的性能从来不是“算”出来的,而是“排”出来的——怎么摆放数据、怎么切分循环、怎么喂给硬件,这些软硬件协同的细节,决定了最终能跑多快。这篇文章想聊的 ops-nn,就是围绕这一目标设计的一套高性能神经网络算子库,主打把卷积、归一化、激活、矩阵乘这类常见算子,按硬件感知的方式重新实现和封装。如果你正在写推理引擎、算子库,或者只是想把模型里那几个慢算子优化一下,这篇文章里的思路和步骤应该对你有用。

1. 性能差距的来源:先搞懂算子的“天花板”在哪

1.1 纸面算力 vs 实际吞吐

性能优化的第一步不是写代码,是先搞清楚这台机器极限能做到多少。现代 CPU 或 GPU 的峰值算力通常很漂亮,但那只是理论极限。实际能用到的,取决于数据能不能及时送到计算单元手里。

我用工厂流水线来类比:计算单元是灶台,内存是仓库,总线是搬运工。灶台再多、厨师手艺再好,食材搬不进来,所有人都在等菜。神经网络算子也一样,很多时候慢的根源不在“算”,而在“搬”。

举例来说,一台设备假设浮点算力是 2 TFLOPS,内存带宽 100 GB/s。如果一个算子每做 1 次浮点运算,至少要读 1 字节的数据,那么这个算子每字节搬运只对应极少次计算,性能就会卡在带宽上,算力再高也白搭。

判断阈值有个简单方式:峰值算力除以带宽。2 TFLOPS 除以 100 GB/s,大概是 20 FLOP/Byte。一个算子的算术强度(每字节搬运对应多少次浮点运算)如果低于这个值,就是访存受限;高于这个值,才有机会接近算力上限。所以拿到一个算子,第一件事不是优化循环,而是先判断它是算得慢还是搬得慢。

1.2 给算子的瓶颈分个类

常见神经网络算子大致分两类,优化策略完全不同:

算子类型典型例子主要瓶颈优先优化方向
计算密集型矩阵乘、卷积、全连接指令效率、计算密度分块、向量化、数据复用
访存密集型归一化、激活、池化、逐点运算内存带宽算子融合、减少搬运次数、缓存友好

判断方法有个土办法:把实现里的循环去掉,只看输入、输出数据量。如果整个算子的数据总量(读加写)远大于运算量,那就别浪费时间抠指令数,先想办法少搬数据。

以逐点 ReLU 为例,输入 N 个 FP32,输出 N 个 FP32,运算只有 N 次比较和选择,但数据搬运是 8N 字节。这种算子你就算把指令优化到极致,收益也很小,真正有效的方法是把它和相邻算子融合,或者让它在缓存里和其他计算共用数据。

计算密集型的 GEMM 则相反,一个 1024x1024 的矩阵乘有约 20 亿次乘加,而输入数据才 8 MB。只要分块策略合理,每个数据可以被上千次运算复用,算力和带宽的天平完全倒向计算端。这时优化重点就变成指令调度、寄存器利用和缓存命中。

1.3 理论峰值 vs 实测:别高估优化空间

做优化之前,我习惯先测两个“参考值”。

第一个参考值是有效内存带宽。写一段简单的连续拷贝或者三元运算循环,测一下设备在理想情况下的实际带宽。大多数机器的实测带宽只有理论带宽的 70% 到 90%,这个数才是你的访存瓶颈基准。

第二个参考值是实际浮点峰值。用一组寄存器饱和的乘加循环,反复跑,取最高值。注意要避免数据被编译器优化掉,直接把累加结果输出到 volatile 变量里。

有了这两个数再去做优化,你就知道空间有多大。我常干的事是:优化完的算子和理论值对比,如果带宽利用率超过 80%,就说明软硬件协同基本到位,不用再折腾了,再抠也只剩几个点,性价比很低。反过来,如果利用率只有 20% 到 30%,那一定是某些关键环节出了问题,值得继续挖。

2. 算子在硬件上的真实“跑姿”:数据布局、SIMD 与 cache 的相互作用

2.1 内存布局决定了一切

神经网络张量有多种主流内存布局。以 4 维数据为例,NCHW 布局按通道连续存储,NHWC 布局把通道维度放到最内层。逐点激活、归一化这类算子天然喜欢 NHWC,因为通道方向连续,处理一个像素时,需要的几个通道数据正好聚在一起,向量化非常容易。而卷积的某些实现更喜欢 NCHW,因为卷积在空间维度卷积时,连续读取同一通道的数据更顺手。

那到底差多少?在某些 CPU 上,同样的逐元素操作,NHWC 比 NCHW 快 2 到 3 倍是常见现象,原因就是 cache 命中率。如果你固定通道去遍历空间坐标,每个通道都是一整幅图像平面,跨平面访问会让缓存里的数据反复被冲掉,每次都要去主存搬,自然慢。

ops-nn 这类库的做法不是把一种布局锁死,而是在算子层面按需选择,配合算子融合减少布局转换的额外开销。布局转换本身也是访存操作,频繁转换得不偿失,所以一般是在网络层面统一规划,而不是每个算子自己转一遍。

2.2 向量化:让编译器“看得出来”

CPU 和 GPU 都有向量处理单元。以 x86 CPU 为例,一条 SIMD 指令一次能处理 4 个 FP32(SSE)或 8 个 FP32(AVX),ARM 上的 NEON 通常一次处理 4 个。想让循环自动向量化,关键不是手写内在函数,而是让编译器有把握做向量化。

编译器不是所有循环都会向量化,它保守得很。要满足的条件包括:数据连续并且按向量宽度对齐、循环计数据编译器能静态推断、指针没有别名问题、循环里没有复杂的条件分支。

我举一个最简单的 ReLU 实现:

void relu(float* out, const float* in, int n) { for (int i = 0; i < n; ++i) { out[i] = in[i] > 0.f ? in[i] : 0.f; } }

这段代码在很多编译器上能自动向量化。但如果你把in和out指向同一块内存,编译器会担心重叠写,直接放弃向量化。你如果用restrict声明out不与任何其他指针重叠,编译器才敢放手干。

还有一个容易被忽略的点:输入指针是否按 32 字节或 64 字节对齐。现代平台用aligned_alloc很容易做到 64 字节对齐,对齐之后 AVX 的加载指令可以一次搬一整行,不需要处理未对齐加载的额外开销。尾部的非对齐部分,单独用标量循环处理,不要混在主循环里打断向量化。

2.3 cache 层级与数据复用

现代 CPU 每个核的 L1 可能只有 32 KB,L2 有几百 KB 到 1 MB,L3 有几 MB 到几十 MB。cache miss 的代价是几十个甚至上百个周期,而一次 L1 命中只需要几个周期。所以算子优化的本质,是尽量让数据在同一层 cache 里多待一阵、被多用几次,而不是反复去主存搬。

经典的 im2col 做法,本质上就是为了把卷积变成 GEMM 的计算密集形态,把数据重新排列到连续内存,让矩阵乘可以吃满 cache。但 im2col 有个明显缺点:内存膨胀得厉害,一个大卷积的中间矩阵可能比原图大十几倍。所以更精细的实现是隐式 im2col,不真正展开数据,而是在计算每个 tile 时动态生成需要的矩阵块,让数据在 L2 里生成并消费掉,不写回主存。

这一节给出三条实战规则,我踩坑总结出来的:

  1. 优先保证连续访存,尽量让相邻线程或向量通道访问相邻地址。
  2. 注意 cache line 的大小,通常是 64 字节,也就是 16 个 FP32。让每个线程在一次迭代里至少消费完整的一行,避免多个线程反复争抢同一行数据。
  3. 别迷信理论热点,用性能计数器看真正的 cache miss 和带宽,再决定优化方向。

3. 软硬协同的四板斧:融合、分块、调度、缓冲

3.1 算子融合:减少搬运次数

模型推理里的算子链往往是“卷积后紧跟着归一化,再跟激活,再跟池化”。如果每个算子都各自读一遍、算一遍、写一遍,同一份数据要进出内存好几次。融合的本意,就是把这些算子合并成一个 kernel,数据在寄存器或 L1、L2 里流动,只写回主存一次。

举一个最简单的例子:Conv + BN + ReLU 在推理时的融合。BN 推理阶段其实是对卷积输出做一个仿射变换,即乘一个缩放因子再加一个偏移。这个缩放和偏移完全可以在卷积权重里提前算好,变成卷积本身的一部分。之后 ReLU 只需要判断结果是否小于 0,可以直接在卷积核的最后一条指令里带过。

融合之后,原本三段算子链至少三次全局内存写入,变成了一次。对带宽利用率的影响是决定性的,尤其是访存密集的逐点算子和轻量算子。注意,融合不是盲目把什么算子都合在一起,要确保融合后的单 kernel 不会把计算密度压低,也不会因为寄存器压力太大导致频繁溢出。如果融合后寄存器不够用,数据被挤到栈上,反而会变慢。

3.2 循环分块:让数据在 cache 里转圈

分块(tiling)是为了让读进来的数据块在 cache 里装得下,并且能重复使用。以 GEMM 为例,一个大的矩阵乘如果不分块,内层循环每次访问矩阵 A 的一行和矩阵 B 的一列,反复横跳,cache 命中率非常差。分了块以后,以 64x64 或 128x128 的 tile 为单位,内层计算在一个小块内反复读数据,算术强度被大幅拉高。

伪代码结构大致是这样:

for (int by = 0; by < M; by += BM) { for (int bx = 0; bx < N; bx += BN) { // 把 A 和 B 的一个 tile 搬到 cache 里 load_tile(A_block, B_block); for (int kk = 0; kk < K; kk += BK) { // 微型 kernel,计算一个小块 compute_micro_kernel(); } } }

tile 大小的选择必须匹配 cache 容量。比如 L2 是 2 MB,你一个 block 的数据加上输出累加器就占了 1.8 MB,那很快就会被其他数据冲掉,分块反而成了形式主义。一般建议 tile 总大小控制在 L2 容量的一半左右,留出余量。另外还需要考虑寄存器数量,微型 kernel 内部往往会有多个累加器同时驻留在寄存器里,寄存器不够就得减少累加器的数量,或者减小 tile 维度。

3.3 线程调度:核心多并不一定快

多线程不是万能。判断依据是可并行工作量和线程开销的比值。小 shape 的逐元素算子,比如只有几千个元素,开线程的同步和唤醒开销比计算本身还大,反而更慢。而大 GEMM 或大卷积,线程分割必须考虑数据局部性和处理器拓扑。

我在项目里一般直接用 OpenMP 跑一个扫描测试。对同一算子分别用static、dynamic、guided三种调度方式跑一遍,然后在运行时根据 shape 大小从配置表里选最优调度。static调度是最均衡的,适合计算量均匀的大算子;dynamic适合负载不均匀的情况;guided介于两者之间。对于 GPU,则是设置好网格大小和线程块大小,让相邻线程访问相邻内存,也就是全局内存合并访问。同一个 warp 里线程访问的内存地址如果很离散,内存事务会被拆成很多个小事务,带宽利用率直接腰斩。

还有一个多线程常见问题:伪共享。多个线程各自写数据,但地址落在同一条 cache line 上,每次写都会把整行锁定,其他线程全部等待。解决办法是让每个线程的输出缓冲区按 cache line 大小对齐,或者在线程之间留出间隔,避免共享一条 64 字节的行。

3.4 缓冲与内存分配:少碰分配器

很多人优化算子时只盯着计算循环,忽略了内存分配的开销。一个具备临时中间结果的算子,如果在每次执行时都调用malloc或new,分配器开销很可能占整体时间的 20% 以上,特别是在小算子上。

算子库的常规做法是:按线程预分配工作区,用内存池复用临时 buffer,分配时按 64 字节甚至更长对齐。如果中间 buffer 是固定尺寸的,直接在算子对象里预留一块内存作为成员变量,避免每次调用都去系统拿。如果是在 GPU 上,尽量避免 host 端和设备端之间的频繁交互,设备端的临时内存用预分配的缓存池。另外注意,临时 buffer 的读写如果只在单线程内使用,尽量让 buffer 局部在线程栈上,栈上的数据更容易被 CPU 缓存。

4. 一个可复现的调优案例:从朴素实现到贴近硬件上限

4.1 基线版本与瓶颈定位

空谈原则没有用,我拿一个 LayerNorm 加 ReLU 的融合算子作为实例,展示完整的优化链路。这个算子把输入归一到零均值单位方差后,再过一遍 ReLU。场景是 x86 CPU、24 线程、支持 AVX2。

第一版朴素实现完全按公式来,先循环一遍算均值和方差,再循环一遍做归一化,最后再循环一遍做 ReLU。仔细一看,输入数据被读了两次,输出写了一次,而整个算子本来可以只读一次输入、写一次输出的。在数据量特别大、远超 L2 缓存的情况下,多读一遍就是把主存带宽翻倍消耗。

用性能计数器一看,这个版本的内存带宽利用率很低,原因不仅是读了两遍,还有第二遍计算时数据早已被挤出缓存,每个元素都要重新从主存拉一次。基线时间我实测约 1.2 毫秒,操作规模是 512x256 的矩阵。

4.2 第一步:单趟遍历少读一遍

既然瓶颈是访存,那第一个优化方向就是能不能只读一次输入。实现上用 Welford 算法或者简单的两趟合并到单趟:先在一个循环里同时累加 sum 和 sum of squares,再回过神用同一批数据做归一化。但这严格来说还是需要数据在缓存里,理想是一趟把数据块放进 L1,洗完再用。

实际操作中是把循环按 block 切分,每个 block 大约 4 KB,读进 L1 后立即完成均值方差计算,然后马上对这块数据做归一化和 ReLU,写回输出,再处理下一块。这样每个数据从主存只读一次,写回一次,时间直接降到约 0.6 毫秒。

这一步最关键的点是“数据局部性”。同样的数学运算量,单纯改变循环顺序和分块方式,就能省掉接近一半的访存开销。这个收益不是微调,是数量级的改变。

4.3 第二步:向量化和尾部处理

下一步解决归一化循环本身的标量问题。LayerNorm 的归一化公式里有一个除法和一个乘加,除以标准差变成了乘以标准差倒数。把向量加载、乘加、比较选择都用 AVX2 内在函数替换。

这段代码最终结构看起来就像:

__m256 v = _mm256_loadu_ps(src + i); __m256 scaled = _mm256_mul_ps(v, _mm256_broadcast_ss(&inv_std)); __m256 shifted = _mm256_add_ps(scaled, _mm256_broadcast_ss(&bias)); __m256 out = _mm256_max_ps(shifted, _mm256_setzero_ps()); _mm256_storeu_ps(dst + i, out);

AVX2 一次处理 8 个 FP32,这一步让计算部分的吞吐直接翻了 8 倍,虽然访存仍然是瓶颈,但时间从 0.6 毫秒降到了约 0.32 毫秒。需要注意尾部处理:数据长度不能整除 8 时,主循环跑到n - 7为止,最后几个元素用标量处理,不要在向量循环里塞一堆条件判断,否则编译器会把你整个循环打断成标量。

4.4 第三步:多线程与缓存友好切块

第三步把数据按线程数切成连续区间,每个线程处理一块独立的 region,内部再做 block 切分。这里的关键是让一个线程处理的区间尽量连续,避免线程之间互相抢 cache,也避免输出缓冲落在同一 cache line。

另外,输出缓冲区需要按 64 字节对齐,而且每个线程不要在同一 cache line 里写相邻位置。实测 24 线程下的最佳切块因子和 AVX2 对齐场景匹配,时间从 0.32 毫秒降到了 0.15 毫秒。注意这里没有线性加速,因为内存带宽本身就是共享的,并行反而让带宽利用了更充分,所以才有明显收益。

4.5 第四步:与前后算子融合

最后一步是把前面的逐点线性变换或后面的相邻激活算子直接融合进这个 kernel。在这个例子里,核心里已经把 ReLU 融进去了。假如前面还有一个逐点缩放算子,那么那个缩放系数可以直接并进来,不需要单独的 kernel 读一遍再写一遍。融合后,整条链路的访存量再减半,整体耗时降到约 0.09 毫秒。

完整的对比表:

版本主要优化手段耗时相对基线提升
基线实现朴素三循环1.2 ms1x
优化 1单趟遍历 + block 切块0.6 ms2x
优化 2AVX2 向量化 + 尾部标量处理0.32 ms3.75x
优化 3多线程 + 缓存友好的线程切块0.15 ms8x
优化 4与相邻算子融合0.09 ms13x

整个过程中,运算量几乎没变,但吞吐提升了一个数量级,靠的全部是软硬协同的搬运策略。这也验证了开头的判断:算子优化的核心不是加速计算,而是减少搬运,让数据在正确的位置被消费。

5. 算子库的工程化落地:接口、测试与可维护性

5.1 接口设计:把调度作为配置而不是写死

单个算子优化只能算“点”,把这套方法沉淀成库才是“面”。ops-nn 这类算子库的接口设计,我认为最重要的是把“算法”和“实现”分离。

调用方不应该关心你内部用了几种 SIMD、分块多大、线程怎么调度。它只需要提供张量描述、输入输出和算子类型,剩下的由库的调度层决定。调度层根据设备能力、数据形状、布局信息,从内部配置表里选择最优的 kernel 变体。

我在项目里通常这样做:

  • 先用朴素但正确的实现作为默认 fallback。
  • 再为热点 shape 注册高性能变体。
  • 高性能变体的分块大小、SIMD 宽度、线程策略都作为参数,由外部配置驱动。
  • 运行时根据 shape 动态选择,同时允许用户覆盖策略,比如强制只用单线程模式做对比测试。

5.2 正确性验证与性能回归

算子库最容易翻车的是正确性验证和性能回归。正确性方面,凡是涉及浮点优化的实现,输出一定会和参考实现有微小差异,不能拿“完全一致”当标准,否则向量化和 FMA 全都不用做了。用相对误差阈值判断更现实,比如每个元素误差允许 1e-5 级别(FP32 环境)。

测例设计要覆盖几个容易出事的角落:

  • 数据量刚好超过 L2 cache 的 shape,验证分块逻辑是否有 bug。
  • 通道数不是向量宽度整数倍的 shape,验证尾部处理是否正确。
  • 归一化里方差为 0 的极端输入,验证行为是否和参考一致。
  • 包含 NaN、Inf 的输入,和参考实现比较语义是否一致。

性能回归同样重要。每次代码合并前跑既定算子的基准测试,如果性能降幅超过设定阈值就直接失败。基准测试不能只优化一次,要放在 CI 里固定机器上跑,避免多跑几轮让 CPU 自动超频干扰结果。

5.3 多硬件后端的抽象

一个现代化算子库通常要同时支持不同 CPU 架构,可能还要兼顾 GPU。推荐的抽象方式是用“同一算法、多后端 kernel”的编译期分发。x86 下区分 SSE、AVX、AVX-512,ARM 下选 NEON 或 SVE,GPU 端走统一的 kernel 接口。内核函数本身用模板参数表达 tile 大小和向量宽度,而不是为每个指令集复制一份完整代码,这样可以大幅减少重复维护的成本。

还要考虑能效。在端侧或移动场景,追求跑满理论峰值可能没有意义,更高的能效比才是目标。软硬协同的优化里,功耗和频率是调度层需要纳入考量的变量,同样是做一个算子,在服务器上要算得快,在手机上可能要算得稳。

6. 我在实践中踩过的坑

老实说,优化的经验都是踩坑踩出来的,最后分享几个我反复掉进去的坑。

第一个坑:对一个小算子疯狂手写汇编,结果收益只有 2%。原因很简单,那个算子的瓶颈是启动开销和内存分配,而不是指令效率。现在我的习惯是先做 profile,用硬件计数器看带宽利用率和 cache miss,再决定往哪里使劲。

第二个坑:为了处理尾部数据,在循环里加了if判断,结果编译器把整个循环的向量化都放弃了。应该把主循环写干净,单独处理尾部的标量部分。这个教训特别值钱,一个if能让几十倍的性能提升瞬间归零。

第三个坑:多线程并行后反而变慢。排查下来是伪共享,多个线程虽然写的是不同数据,但地址落在同一条 cache line 上,互相踩踏。解决办法很简单,输出缓冲区按 64 字节对齐,并且每个线程的起始位置间隔一条 cache line。这个坑不踩一次真的很难想到。

第四个坑:只测了一组 shape 就宣布优化成功。结果换一个更大的 shape 之后,性能直接暴跌。原因是我优化的分块大小只对那个特定尺寸友好。后来我建立了一个 shape 分布统计表,覆盖各种实际出现的尺寸,每次优化完都要过一遍这个表。

第五个坑:把 GPU 的 kernel 也像 CPU 一样疯狂开线程,结果反而更慢。GPU 端的线程不是越多越好,关键是合并访问和占用率之间的平衡,需要看设备的 occupancy,不是拍脑袋定的。

我自己在写完这套流程之后再回头看那些慢算子,第一反应已经不再是“调调参”,而是先问它到底是算得慢还是搬得慢。这个思维转变才是做高性能神经网络算子最有价值的部分。ops-nn 这类库对很多人来说可能只是一堆 kernel 的集合,但它背后真正值得学的,是你如何把一个算子从“能用”做到“贴近硬件上限”,并且让这套能力可以复制到所有算子上。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询