# 算子与内核 · ATen 与 native_functions.yaml

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

- 作者：David（道雾轩）
- 专栏：PyTorch 源码漫游（https://daiw.org/manual/pytorch-internals.md）
- 最后更新：2026-07-13
- 原文：https://daiw.org/manual/pytorch-internals/operators-and-kernels
- 转载与引用：请注明出处并附原文链接（https://daiw.org/about/copyright）

# 算子与内核 · ATen 与 native_functions.yaml

[上一章](https://daiw.org/manual/pytorch-internals/dispatch-journey)的 dispatcher 负责**把请求送到正确的内核**。这一章往下走一层：那个真正执行加法的**内核**，写在哪、长什么样、怎么同时供 CPU 和 CUDA？顺带解开一个谜：**为什么你 grep 半天，找不到 `add` 到底在哪实现。**

## ATen：算子们住的地方

PyTorch 的 C++ 张量库叫 **ATen**(**A** **TEN**sor 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` 的条目，大意长这样：

```yaml
- 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:` 块把[上一章的分发键](https://daiw.org/manual/pytorch-internals/dispatch-journey)映射到具体内核函数：`CPU` 和 `CUDA` 走同一个 `add_kernel`（内部再分流），稀疏张量走别的。

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

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

真正的计算内核，分 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 那章](https://daiw.org/manual/pytorch-internals/tensor-and-storage)的 stride 机制在内核层的兑现：TensorIterator 就是靠 sizes/strides 来正确遍历各种「说明书」。

## 顺带：反向公式也登记在册

还有一本关键的册子：**`tools/autograd/derivatives.yaml`**，登记每个算子的**反向（求导）公式**。比如乘法 `mul` 的条目会写明：对 `self` 的梯度是 `grad * other`，对 `other` 的梯度是 `grad * self`。

PyTorch 的 codegen 会**根据它，自动生成[上一章说的 Autograd 键那一层](https://daiw.org/manual/pytorch-internals/dispatch-journey)的内核**——也就是「前向时该记什么账、反向时该怎么算梯度」。**这本 `derivatives.yaml`，就是下一站 autograd 引擎的燃料。** 你自定义算子时想让它可导，往往就是在这类地方（或用 `torch.autograd.Function`）补上反向公式。

## 小结与过渡

到这里，「一次调用」的下半程清楚了：**dispatcher 路由 → ATen 内核落地计算，而海量算子靠 `native_functions.yaml` 声明 + codegen 生成 + TensorIterator 复用**。

我们还欠一笔账没细讲：前向时，那个 **Autograd 键**到底「记了什么账」？`loss.backward()` 又是怎么顺着这些账，把梯度算出来的？这是整台引擎最神奇的部分，也是接下来两章的主角。👉 [autograd（上）：前向如何悄悄建图](https://daiw.org/manual/pytorch-internals/autograd-graph)

## 参考文献

- [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)。