算子与内核 · ATen 与 native_functions.yaml · PyTorch 源码漫游

dispatcher 把请求路由到内核，那内核本身写在哪？看 PyTorch 的张量库 ATen、那本记着两千多个算子的「花名册」native_functions.yaml，以及 TensorIterator 怎么让一份内核搞定广播、类型提升与 CPU/GPU 并行——外加 PyTorch 大量代码是「生成」出来的这件事。

算子与内核 · ATen 与 native_functions.yaml

上一章的 dispatcher 负责把请求送到正确的内核。这一章往下走一层：那个真正执行加法的内核，写在哪、长什么样、怎么同时供 CPU 和 CUDA？顺带解开一个谜：为什么你 grep 半天，找不到 add 到底在哪实现。

ATen：算子们住的地方

PyTorch 的 C++ 张量库叫 ATen(A TENsor library)。你在 Python 里调的每个张量操作——add、matmul、softmax、conv2d——底下都对应 ATen 里的一个（或多个）C++ 算子(operator)。它是引擎真正「干活」的车间。

但 ATen 有两千多个算子，每个还可能有 CPU、CUDA、autograd 等多套实现。这么多东西，靠人手一个个写注册、写绑定，根本维护不动。 于是 PyTorch 用了一招：用一份「花名册」声明所有算子，其余样板代码全部自动生成。

native_functions.yaml：两千多个算子的花名册

这本花名册，就是 aten/src/ATen/native/native_functions.yaml。每个算子在里面登记一条，声明它的签名(schema)和哪些分发键有内核。比如 add 的条目，大意长这样：

- func: add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
  variants: function, method
  dispatch:
    CPU, CUDA: add_kernel        # CPU / CUDA 各自的内核
    SparseCPU, SparseCUDA: ...   # 稀疏张量另有实现

第一行是算子的类型签名：名字、参数、返回值——这是 dispatcher 认识它的「身份证」。
dispatch: 块把上一章的分发键映射到具体内核函数：CPU 和 CUDA 走同一个 add_kernel（内部再分流），稀疏张量走别的。

这解释了「为什么 grep 找不到 add」。 PyTorch 在编译时跑一套代码生成(codegen)：读 native_functions.yaml，自动生成海量样板——Python 绑定、dispatcher 注册、参数解析、autograd 包装……所以一个 add，是由「YAML 声明 + 生成的胶水 + 手写内核」拼起来的，不是一个你能一眼找到的函数。看 PyTorch 源码，得先接受**「大量代码是长出来的，不是写出来的」**。

内核长什么样：一份逻辑，多种硬件

真正的计算内核，分 CPU 和 CUDA 两类，住在 aten/src/ATen/native/(CPU)和 aten/src/ATen/native/cuda/(CUDA)下：

CPU 内核：一段在 CPU 上遍历元素、做加法的 C++ 代码，通常还做了多线程和 SIMD 向量化。
CUDA 内核：一段在 GPU 上并行执行的 kernel，成千上万个线程一人算几个元素。

同样的数学（逐元素相加），两种硬件、两套写法。 dispatcher 根据张量在哪，替你挑对应的那套——这就是为什么你写 a + b 从不用关心 a 在 CPU 还是 GPU。

TensorIterator：不必为每个算子手写循环

问题又来了：逐元素算子（add、mul、sin、比较……）有几百个，每个都要处理广播、类型提升（int + float 该升成 float）、非连续 stride、多线程、向量化——难道每个都手写一遍这些琐事？

不。PyTorch 有个主力工具 TensorIterator(aten/src/ATen/TensorIterator.h)：你只需告诉它「输入输出是哪些张量、每个元素怎么算」，它替你统一处理广播对齐、类型提升、stride 遍历、并行切分、向量化。绝大多数逐元素和归约内核，都是架在 TensorIterator 之上写的——这样「逐元素相加」这段核心逻辑只写一次，广播和并行的脏活交给它。这正是 Tensor 那章的 stride 机制在内核层的兑现：TensorIterator 就是靠 sizes/strides 来正确遍历各种「说明书」。

[1] PyTorch 源码 README：aten/src/ATen/native/README.md（native_functions.yaml 的写法说明）、aten/src/ATen/native/native_functions.yaml(pytorch/pytorch)。
[2] Edward Z. Yang. “PyTorch internals.” 2019. http://blog.ezyang.com/2019/05/pytorch-internals/ —— ATen、codegen、TensorIterator 的概览。
[3] PyTorch 文档：“Tensor Iterator internals.” https://github.com/pytorch/pytorch/wiki/How-to-use-TensorIterator —— TensorIterator 用法。
源码：aten/src/ATen/TensorIterator.h、tools/autograd/derivatives.yaml(pytorch/pytorch)。

算子与内核 · ATen 与 native_functions.yaml