上个月帮一个做推荐算法的朋友排查线上推理变慢的问题。他给我看模型代码，前向算下来也就几十个算子调用，怎么看都不该慢成那样。结果问题不出在模型结构，而是落在某个自定义算子没有适配推理引擎的高效执行路径上，框架兜底走了一条通用但很慢的 kernel。这件事让我想写一篇关于“AI算子”的入门文章：很多算法工程师每天都在调用算子，但算子本身是什么、为什么同一个数学公式会有完全不同的执行速度、以及当框架没有现成算子时该怎么自己动手写，这些其实很少有人系统性讲过。
这篇内容适合几类人：刚开始接触深度学习框架、好奇底层实现的在校学生；工作中频繁踩到性能问题但不太清楚从哪入手的算法工程师；以及做过传统开发、想了解 AI 计算核心逻辑的老朋友。我会从概念边界讲起，给出手写算子的完整示例，也会分享我在实际项目里踩过的坑和排查思路。看完之后，你至少能从“会用算子”迈入“看得懂算子、改得动算子”这个区间。
1. 算子不是“算法”——先把概念边界划清楚
1.1 一句话认识算子：输入Tensor，输出Tensor
“算子”这个词源自动态数学里的 operator 翻译，在深度学习领域，它被用于描述 Tensor 到 Tensor 的一个计算单元。你不用把它想得多神秘：一个 python 函数接收若干 Tensor，经过数学变换，再输出 Tensor，这个函数对应的底层实现，就是算子。
举几个常见的例子你立刻就能对上号：
- 元素级算子：ReLU、Sigmoid、Add、Mul，输入和输出 shape 相同，逐个元素做计算；
- 规约算子：Sum、Mean、Softmax、LayerNorm，往往在某个维度上聚合信息；
- 矩阵算子：MatMul、Conv2D、ConvTranspose，整个推荐系统、CV 模型的核心计算都靠它；
- 数据搬运算子：Reshape、Transpose、Permute，不改变数值本身，但改变张量在内存中的排列方式。
你平时写torch.relu(x)、F.linear(x, w, b)，其实就是在向框架请求执行某个算子。框架做的主要事情，就是把你 Python 端的“算子请求”翻译成具体硬件上的执行代码，并像工厂流水线一样把一个个算子的输入输出接好。
1.2 Kernel、OP、算子：同一个词的三副面孔
这在社区里非常容易绕晕，因为不同资料用词不一样。实际上：
- OP（Operation），偏上层概念，表达“我要做什么”；比如 PyTorch 里的
torch.nn.functional.relu是一个 op。 - Kernel，偏执行层概念，指“在硬件上具体怎么算”。同一款 op，有不同的 kernel：CPU kernel、GPU CUDA kernel，甚至 CPU 上还有针对不同指令集（AVX、NEON）优化的 kernel。
- “算子”这个词在很多时候是 OP 和 Kernel 的总称，看语境偏哪一层。
一个入门的诀窍：当资料说“添加一个新算子”时，通常要同时完成上层 OP 定义和底层 Kernel 实现，以及把它们注册到框架的调度器里。研究框架源码时，你会发现aten/src/ATen/native目录下存放的就是各类 CPU 算子实现，而MACE/ONEAPI/TVM这类项目里讨论的 kernel 则偏向下层硬件代码。
1.3 算子与计算图：自动微分依赖的“砖块”
另一个常见误解，是把算子等同于神经网络的“层”。其实层是逻辑单元，算子才是真实计算单元。例如一个nn.Linear层，内部至少包含matmul与add（bias）两个算子，反向还对应matmul的另一个隐式算子。模型训练中所谓反向传播，就是沿计算图把算子之间的梯度求导规则逐一应用。框架的 autograd 引擎维护完整的计算图，图上每个节点就是一个算子执行点，边是 Tensor 依赖。
因此，算子不只是在做数学运算，它还承接了两个额外职责：
- 作为可微函数，提供 forward 与 backward 两个方向的实现；
- 接受框架调度器的参数选择，不同 dtype、device、layout 下自动匹配不同 kernel。
这就是为什么从业人员常说“算子是全栈能力的电子侦查集合点”：往下看是硬件指令，往上看是模型语义。
2. 为什么“三行代码”的性能差几十倍：算子是性能的第一现场
2.1 同样一个卷积，三重循环和库实现差距在哪
我先说一个反直觉的事实：一个数学定义上完全相同的卷积，手写三重循环与高性能库实现做对比，性能差距可以到几十倍或上百倍。比如一个 3x3 卷积，朴素实现按 5 重循环展开每个输出像素，对每个输入通道取值、乘加、写回内存；而 cudnn 或 oneDNN 的实现会考虑内存布局、数据复用、SIMD 向量化、分块（tiling）等，把同一个公式以完全不同的方式映射到硬件。
用做饭做个类比：数学公式是菜谱，硬件是灶台，而算子就是厨师。同一个酸辣土豆丝配方，新手灶台前手忙脚乱是一桌菜，经验丰富的师傅掂勺翻锅、火候顺着锅气走，出锅完全两回事。菜谱没变，问题在于“执行策略”。
为什么要强调这一点？因为很多入门者调试性能问题时，第一反应是“模型结构是不是有问题”，很少去怀疑算子本身的执行策略。但在深度学习中，算子执行的效率往往比算子“数量”更影响最终速度。
2.2 框架调用的实际路径：Python 端的“水面之下”
PyTorch 哲学的入门通常只说了 Python 用法，但实际一次torch.relu(x)调用背后，会经过这么几步：
- 解析 Python 参数，绑定到
torch.relu的 C++ 入口； - 进入 ATen dispatcher，根据 Tensor 的 dtype、device、layout 找到对应的 kernel 注册项；
- 如果是 CUDA tensor，则从算子库 dispatch 到 CUDA kernel；
- kernel 在 GPU 上按线程网格并行执行，结果写进输出 Tensor 内存；
- 返回 Python 对象。
这个调度过程里很多策略与性能强相关。比如算子融合就是“减少调度次数”的典型思路：把add + relu融合为add_relu一个 kernel，减少一次 Tensor 读取、一次写回以及一次 kernel launch 的开销。框架层面有手动融合，也有torch.compile这样的编译器自动融合。理解了调度机制，你会更容易看懂框架为什么某次升级后模型快了不少——不一定是模型层面变了，很可能是某个算子的 kernel 被换成了更高效版本。
2.3 什么情况下你必须研究算子
并不是每个项目都需要深入算子，算法工程师可以先跑通再说。但下面这些场景出现时，你绕不开：
- 模型里有一种标准算子无法表达的操作，需要自定义实现；
- 已有算子拼出来的功能性能不理想，需要把它融合成更高画面效率的单个算子；
- 从训练框架切到推理引擎（TensorRT/ONNX Runtime）后，自定义算子不被原生支持，导致整图优化被打断；
- 遇到精度问题，需要对算子内部的累加顺序、归一化位置做精细化控制。
提前知道这些事情，至少能让你在遇到模型变慢或者部署失败时，少走很多弯路。
3. 手写一个 ReLU：从数学到可运行实现
3.1 先写一个“慢但绝对正确”的版本
网上讲算子往往一上来就上 CUDA，容易劝退初学者。我自己带人入门算子开发时，会先要求他在 Python 侧把算子的语义写对，再做性能版本。这个“慢但正确”的版本，是所有后续工作的校验基准。
挑选 ReLU 来实操，是因为它足够简单又包含关键要素：前向要对原始输入做max(0, x)，反向要把上游梯度原样传给输入中大于 0 的位置，等于 0 的位置梯度可以约定为 0（不同框架可能微调，但 PyTorch 沿用的是 0）。用 NumPy 实现：
import numpy as np def relu_forward(x: np.ndarray) -> np.ndarray: return np.maximum(x, 0.0) def relu_backward(dout: np.ndarray, x: np.ndarray) -> np.ndarray: return dout * (x > 0).astype(dout.dtype)别小看这几行代码。它明确了一个重要原则：算子的反向不是独立定义的，而是对应前向的“分段导数规则”。写算子前，先把 forward 与 backward 的数学逻辑写死在注释里，这比一上来就写 CUDA 值得多。
3.2 用 autograd.Function 接入反向传播
PyTorch 给自定义算子提供了标准入口，就是继承torch.autograd.Function。我们在这里实现前向与反向，而不是简单调用 Python 函数，这样框架的 autograd 引擎就能识别算子的求导路径。
import torch class ReLUFunc(torch.autograd.Function): @staticmethod def forward(ctx, x): # ctx 用于在 forward 中保存反向需要用到的信息 ctx.save_for_backward(x) return x.clamp_min(0) @staticmethod def backward(ctx, grad_output): (x,) = ctx.saved_tensors # 这里一定要记得把 mask 转成与梯度一致的类型 return grad_output * (x > 0).to(grad_output.dtype)之后在模型里就能这样调用：
y = ReLUFunc.apply(x)这样，算子逻辑虽然写在 Python 端，却已经通过autograd.Function具备“可微分性”了。需要特别注意的是：forward接收几个 Tensor 输入，backward就应返回对应数量的梯度；非 Tensor 的上下文性输入（比如一个浮点参数）通常返回None，这一点很多入门者在这里被刺。
3.3 小实验验证前向与反向
代码写出来不代表对，必须有验证实验。最直接的方式是和框架自带实现对照：
torch.manual_seed(42) x = torch.randn(4, 4, requires_grad=True) out_custom = ReLUFunc.apply(x) x2 = x.detach().clone().requires_grad_(True) out_ref = torch.relu(x2) # 前向一致性 print("forward match:", torch.allclose(out_custom, out_ref)) # 反向一致性，给一个任意上游梯度 grad_seed = torch.randn(4, 4) out_custom.backward(grad_seed) out_ref.backward(grad_seed) print("backward match:", torch.allclose(x.grad, x2.grad))“反向一致性”是算子开源社区非常常见的 CI 测试方式。凡是写新算子，我都会把这种对照测试固化到测试文件里，避免后续优化时把一个看似等价的新 kernel 换进来，结果把梯度悄悄“优化”错了。
3.4 算子注册到框架的关键点
如果你想让做出来的算子成为一个真正可被框架调度的算子，而非仅在一个模型文件中可用，就需要完成注册。PyTorch 的实践路径大致是：
- 用 C++ 扩展或
torch.ops注册命名空间和算子函数； - 注册一个或者多个 kernel，每个 kernel 对应一种设备/数据类型；
- 对需要在 TorchScript 或 ONNX 里用的算子，还要补充 schema（输入输出类型约束）与 shape 推断逻辑。
C++ 自定义算子典型的骨架是这样（这里只展示 CPU 版本）：
#include <torch/extension.h> torch::Tensor relu_forward(torch::Tensor x) { return at::clamp_min(x, 0); } TORCH_LIBRARY(my_ops, m) { m.def("relu_forward", &relu_forward); }写 Python 也好、写 C++ 也好，你的目标其实相同：让框架知道“有这么个算子，这个算子的计算语义是什么，在什么设备上该执行哪个实现”。真正难的东西不在注册代码本身，而在于写出的实现在不同 dtype、不同 shape、不同内存格式下都能保持正确。
4. 让算子变快之前，先理解计算与存储的“瓶颈”
4.1 算子的速度瓶颈往往在内存带宽，而不是算力
我见过很多刚接触性能优化的人，总是下意识认为优化就是把“计算次数”变少，于是先数 FLOPs。但实际上，对大部分 elementwise 算子和不少规约算子来说，瓶颈根本不在计算，而在内存访问。
算一下就明白了：假设要处理 1GB 的 float 数据，计算只是让每个元素和 0 比大小。如果内存带宽是 50GB/s，读一遍数据就要 20ms；你即便把“和 0 比大小”这个操作优化到零成本，总时间也不会低于 20ms。这时候性能天花板叫“内存墙”，有些资料也直接叫带宽瓶颈。
所以对 elementwise 类算子，真正有效的手段是减少内存访问次数。这就是为什么add+relu融合有意义：原本读两个数、写两个数，现在读两个数、写一个数，节省了一次写回和一次 kernel launch。而在显存带宽受限的场景下，这种融合往往比减少 FLOPs 更直接。
4.2 三个最常用的优化方向：融合、分块、向量化
- 融合（Fusion）：把多个算子的数据流合并到一个 kernel，减少中间 Tensor 的读写。框架层面
torch.compile已经在做，手写自定义算子时也可以主动设计。 - 分块（Tiling）：主要针对 MatMul 这类计算量大、单个矩阵放不进缓存的问题。把输出矩阵切块，让计算过程尽量复用 cache 里的数据块，从而降低对内存带宽的依赖。
- 向量化（Vectorization）：一次处理多个元素的加载和计算。CPU 上依赖 SSE/AVX 指令，GPU 上则是让线程做连续地址访问，保证内存 coalescing（合并访存）。
对一个入门项目而言，我的建议流程是：先确认内存访问模式是否连续 → 再考虑能否和相邻算子融合 → 然后才考虑向量化等微优化。顺序错了，经常白忙活。
4.3 学会测量算子的真实开销
优化之前必须会测量。入门阶段做两件事就够了：
- 用
torch.profiler或框架自带的 profile 工具，记录各算子的耗时占比； - 单独测一个算子的纯执行时间，用多次重复加
torch.cuda.synchronize()消除异步执行的影响。
import torch import time x = torch.randn(1024, 1024, device="cuda") # 预热 for _ in range(10): y = torch.relu(x) torch.cuda.synchronize() start = time.perf_counter() for _ in range(1000): y = torch.relu(x) torch.cuda.synchronize() avg_ms = (time.perf_counter() - start) / 1000 * 1000 print(f"avg: {avg_ms:.3f} ms")此时算一下数据的有效读写量，再除一下显存带宽，就能大概估算“理论下限”。如果实际耗时是理论上限的几倍，说明还有很大优化空间；如果已经贴近带宽上限，那就说明这个算子本身能挖的不多了。建立这种“先估上限再对比实测”的思维，比盲目套优化手段要有用得多。
5. 算子开发四大坑：我的踩坑记录
5.1 dtype不一致：输出正常，梯度全乱
第一次写自定义算子时，我在反向里直接写了grad_output * (x > 0)，前向和反向都通过了我的对照测试，因为当时输入也是 float32。后来换成混合精度训练，输入 x 变成 float16，但x > 0得到的是 bool 张量，和 float16 的梯度相乘时类型提升了一阵，看似没报错，实际上梯度值全部悄悄变成了 float32，再传给下游算子时就出现了精度衰退问题。
真正的教训是：反向 mask 一定要显式转成与 grad_output 相同的 dtype。这也是为什么我在上文代码里特地写了.to(grad_output.dtype)。算子实现里，dtype 是所有后续逻辑的地基，这个地基出了偏差，问题往往要到几个算子之后才爆发。
5.2 backward 返回值数量必须“斤斤计较”
另一个高频报错来自autograd.Function的 backward 返回值约定。写得复杂一点的算子，forward 输入往往是多个：
@staticmethod def forward(ctx, x, bias, scalar_param): ctx.save_for_backward(x, bias) return x + bias + scalar_param如果 backward 里写了三个 return，但 scalar_param 本身又不是 Tensor 参数（比如普通浮点），框架会直接报错。正确写法是：
@staticmethod def backward(ctx, grad_output): x, bias = ctx.saved_tensors return grad_output, grad_output, None这条规则想通之后就很简单，却真的很容易在实际算子代码里搞漏。每当我看到一串 “return grad_x, grad_w, None, None” 的代码，都会想起来“有几个输入就还几个梯度”这句话，真的写着写着就容易手滑。
5.3 动态 shape 与 JIT 编译平台的冲突
我曾经把一个小算子封装到 TorchScript 里用于生产推理。本地跑没问题，一到线上容器里新传入的 batch 大小和之前不同，就报找不到对应 kernel 的错误。
原因在于 TorchScript 在推理时会针对特定 shape 做静态化编译，而我的自定义算子没有提供足够的 shape 提示或 fallback 路径。框架默认只有在 shape 匹配的情况下才用编译缓存。解决办法：对自定义算子的 JIT 使用场景补充 schema 标注，并且明确告诉框架这个运行时的动态范围。更稳妥的是在部署前把动态维度固定或做 padding 对齐，减少跨 shape 的编译消耗。
这块暴露出的深层问题是：算子开发者不能只保证“Python 里能跑”，还得考虑它在其他执行引擎里的可编译性。
5.4 ONNX/TorchScript 导出时 shape 推导失效
把包含自定义算子的模型导出成 ONNX 时，框架并不知道你这个算子的输出 shape 应该是什么。于是导出的计算图里要么多了一个未知节点导致推理引擎无法继续做图优化，要么完全导出失败。
正确的工程做法是：为自定义算子注册 shape 推断函数。现实中很多团队图省事，选择在导出前把自定义算子替换成标准算子组合，比如把融合算子“展开”成 add + relu。这虽牺牲推理性能，但换来跨平台兼容性，属于部署阶段的合理折中。当你学习算子开发从“跑通”走向“上线”，这类问题就会成为主要的工作内容。
6. 如果我现在重新入门算子开发，会怎样安排学习路径
6.1 先选一个熟悉框架，吃透一个简单算子的完整链路
如果让我重新入门，我不会同时追 PyTorch、TensorFlow、JAX 三家生态，而是挑自己最常用的一个框架，用 ReLU 或 Softmax 做研究对象。具体做法：
- 在 Python 里调用这个算子，用 profile 工具看耗时时长；
- 找到它对应的 C++/CUDA 实现源码路径；
- 用上文讲的方法自己重写一份等价实现并做正确性对比；
- 再尝试和相邻算子做融合或向量化，比较前后性能变化。
这套循环下来，你对“框架如何调算子、算子在硬件上如何执行”的理解，会比读十篇博客更扎实。
6.2 从“调用算子”到“阅读算子实现”
阅读实现源码没有想象中困难。拿 PyTorch 举例，aten/src/ATen/native/cpu/SoftmaxKernel.cpp这类文件中，你能看到实际 kernel 的循环写法、内存布局处理以及针对最后维度优化的特殊情况。刚开始看不懂没关系，按“找到入口 → 找到核心循环 → 理解数据索引 → 理解并行方式”这个顺序去读即可。
值得提醒的是，源码里大量存在“特例”。例如某些算子在最后一个维度上能走快速路径，在中间维度上则要转置或复制数据。这些特例行代码往往才是性能的秘密所在，也最容易让初学者困惑。我的建议是：先忽略特例，抓住通用路径，理解后再回头研究为什么需要特例。
6.3 站在生态上做选择
最后说点我的个人感受：算子开发并不是“必须手写 CUDA 才算入门”。如今很多工作可以在高层完成，比如用torch.ops注册现有 C++ 实现、用 Triton 编写 GPU 计算、依赖编译器来自动优化。但有一点始终不会变：你必须对你实现的算子在内存和计算上的行为有准确预期，否则框架再智能也无法帮你兜底。
把基础概念、执行流程、调试方法掌握好，之后无论是转向推理引擎算子适配、编译器算子实现，还是继续做模型侧算法开发，这一层认知都会时刻发挥作用。这是我认为这部分知识对算法工程师最有价值的长期回报。
