一、算子优化概述

什么是算子

先说说算子是什么。说白了，算子就是深度学习网络中的一个基本计算单元。

你想想看，一个神经网络模型，不管多复杂，拆开来看，无非就是卷积、池化、全连接、激活函数这些基本操作。每个这样的操作，我们就叫它一个算子。

我记得刚入行那会儿，有个同事问我：「算子跟函数有什么区别？」嗯，这个问题问得好。函数是编程概念，算子是计算概念。一个算子可能对应一个函数，也可能对应一组函数。比如卷积算子，底层可能拆成 im2col、矩阵乘法、权重重排等多个步骤。

常见的算子有这些：

• 计算类算子：卷积、矩阵乘法、BatchNorm、LayerNorm
• 激活类算子：ReLU、Sigmoid、GELU、Swish
• 池化类算子：MaxPool、AvgPool、GlobalAvgPool
• 形状变换类：Reshape、Transpose、Squeeze、Expand
• 逐元素操作：Add、Mul、Div、Softmax

算子优化的目标

算子优化到底在优化什么？我把它总结成三个维度：

维度	目标	典型手段
计算速度	减少延迟，提高吞吐	算子融合、循环展开、向量化
显存占用	降低峰值显存，避免OOM	显存复用、梯度检查点、重计算
带宽利用	提高数据搬运效率	访存优化、数据布局调整、预取

这三个目标有时候是矛盾的。比如你用重计算来省显存，那就得多算一遍，速度就慢了。我在项目中遇到过这种情况：一个 7B 的模型，显存死活放不下，最后用梯度检查点省了 40% 显存，代价是训练时间多了 15%。这个 trade-off 值不值？得看你的瓶颈在哪。

我个人习惯是：先看显存够不够，不够就省显存；显存够了，再全力优化速度。别一上来就追求极致速度，结果 OOM 了，白忙活。

算子优化的意义

为什么要花时间优化算子？我给你三个理由：

1. 性能差距巨大：一个没优化的算子，跟一个精心优化的算子，性能可能差 10 倍甚至更多。我见过有人用 PyTorch 写了个自定义算子，跑一次要 200ms，我帮他改成 CUDA 实现后，直接降到 8ms。嗯，25 倍的差距。
2. 框架的瓶颈往往在算子：深度学习框架的调度开销其实很小，真正吃时间的是算子本身的执行。你优化了框架调度，可能只省了 5%；但优化一个热点算子，可能省 50%。
3. 硬件潜力要靠算子来释放：GPU 的算力越来越强，但如果你写的算子不能充分利用 Tensor Core、不能有效利用共享内存，那再强的硬件也白搭。

算子优化在框架中的位置

为了让你更直观地理解，我画了一张图：

从这张图你能看到，算子优化处在框架层和硬件层之间。它向上承接框架的调度指令，向下对接硬件的计算能力。

具体来说，算子优化在框架中扮演这几个角色：

• 算子注册：把优化后的算子注册到框架中，让框架能调用它
• 算子选择：根据输入形状、数据类型、硬件型号，自动选择最优的算子实现
• 算子融合：把多个连续的小算子合并成一个，减少访存和 kernel launch 开销
• 自动调优：通过搜索或机器学习，找到当前硬件上最优的算子参数

算子优化的常见手段

我简单列一下常用的优化手段，后面章节会逐个展开：

手段	说明	典型收益
算子融合	将多个连续算子合并为一个 kernel	减少访存 30%-70%
显存复用	同一块显存被多个算子交替使用	降低峰值显存 20%-50%
梯度检查点	前向时不保存中间结果，反向时重算	显存节省 30%-60%
混合精度	FP16/BF16 代替 FP32	速度提升 1.5-3 倍
循环展开	减少循环控制开销	速度提升 10%-30%
向量化	利用 SIMD 指令一次处理多个数据	速度提升 2-4 倍

我个人习惯是：先做算子融合，收益最直接。然后看显存，不够就上梯度检查点。最后再考虑向量化和循环展开这些细节优化。

好了，这一章就到这里。算子优化是个系统工程，后面我们会一步步深入。记住核心：先 profiling，再动手；先显存，再速度；先融合，再细节。