第6章 ATen 算子库与代码生成

"ATen 'native' functions are the modern mechanism for adding operators and functions to ATen. Native functions are declared in native_functions.yaml and have implementations defined in one of the cpp files in this directory."

—— aten/src/ATen/native/README.md（v2.11.0）

6.1 问题：近 2700 个算子背后的人力黑洞

第 5 章我们看到 PyTorch 注册算子要做的事：

• 写 schema（输入输出类型）
• 写每个 backend 的实现（CPU / CUDA / MPS / XPU / ...）
• 写 autograd 反向规则
• 注册到 dispatcher
• 生成 Python 包装（让 torch.add 能调用）
• 生成 C++ 包装（让 at::add 能调用）
• 生成 Tensor::add 方法
• 生成 aten::add schema 字符串（给 JIT / TorchScript / FX）
• 生成 functionalization 规则
• 生成 vmap 的 plumbing 胶水
• ……

每个算子都要做以上 10+ 件事。v2.11 的 native_functions.yaml 里有 2674 条算子 schema（§6.4.4）。如果靠手写，工作量是几万次重复劳动 —— 每次新增一个 dispatch key（如某国产芯片）都要动近 2700 处。

更糟糕的是，每条衍生路径之间必须保持一致：

• Python 端 torch.add(a, b, alpha=2) 与 C++ 端 at::add(a, b, 2) 必须接收相同参数顺序
• loss.backward() 触发的反向必须与前向接收的张量类型/形状对得上
• torch.jit.trace 序列化的 IR 节点必须包含完整 schema 信息
• functionalize 规则必须知道哪些参数是 inplace 的

人工维护这种一致性几乎不可能。任何一处疏忽都会让用户看到"前向能跑反向报错""Python 能调 C++ 不能调"等怪异 bug。这就是为什么 PyTorch 必须用代码生成 —— 不只是省人力，更是用机器生成保证一致性。

PyTorch 团队的解法是 代码生成：把所有重复部分用一份 YAML 声明，让 Python 脚本在编译时生成 C++ 代码。这一章拆这个体系。

flowchart TB
    Y1[native_functions.yaml<br/>16172 行]
    Y2[derivatives.yaml<br/>3253 行]
    Y3[tags.yaml]
    Y1 --> TG[torchgen<br/>Python 生成器<br/>约 25000 行]
    Y2 --> TG
    Y3 --> TG
    TG --> O1["build/aten/src/ATen/Operators_*.cpp<br/>(5 个分片)"]
    TG --> O2[RegisterCPU_0..3.cpp / RegisterCUDA_0.cpp / ...]
    TG --> O3[Functions.h / NativeFunctions.h]
    TG --> O4[VariableType_*.cpp<br/>autograd 包装]
    TG --> O5[python_torch_functions_*.cpp<br/>Python 端绑定]
    TG --> O6[python_variable_methods.cpp<br/>Tensor 方法]
    TG --> O7[RegisterFunctionalization_*.cpp<br/>functionalization]
    style TG fill:#fef3c7,stroke:#f59e0b,stroke-width:3px
    style Y1 fill:#dbeafe,stroke:#3b82f6
    style Y2 fill:#dbeafe,stroke:#3b82f6

6.2 native_functions.yaml 的语法

打开 aten/src/ATen/native/native_functions.yaml，第一感觉是 庞大 —— 16172 行，全是 YAML 条目。每个条目对应一个 ATen 算子。让我们看一个真实的例子（mm —— 矩阵乘）：

- func: mm(Tensor self, Tensor mat2) -> Tensor
  structured_delegate: mm.out
  variants: function, method
  dispatch:
    SparseCPU, SparseCUDA, SparseMPS: _sparse_mm
    SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: _sparse_csr_mm
  tags: core

- func: mm.out(Tensor self, Tensor mat2, *, Tensor(a!) out) -> Tensor(a!)
  structured: True
  dispatch:
    CPU: mm_out_cpu
    CUDA: mm_out_cuda
    MTIA: mm_out_mtia
    MPS: mm_out_mps
    XPU: mm_out_xpu
    SparseCPU, SparseCUDA, SparseMPS: _sparse_mm_out
    SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: _sparse_csr_mm_out

——aten/src/ATen/native/native_functions.yaml 中 mm 与 mm.out 的真实条目

每个字段的语义：

• func：算子签名。仿 TorchScript IR 类型语法
• structured_delegate：把实现委托给另一个算子（典型是 out= 版本）
• structured: True：这是结构化算子，使用统一的 broadcast / dtype 提升 / output 分配框架
• variants: function, method：生成 at::mm(a, b)（function）和 Tensor::mm(b)（method）两种 C++ API
• dispatch：哪个 dispatch key 调用哪个 C++ 实现
• tags: core：分类标签，给算子打语义标记；core 表示它属于 Core ATen 算子集 —— 经 aten 分解与功能化后保留下来、供编译器后端对接的纯函数算子子集（定义见 aten/src/ATen/native/tags.yaml）

6.2.1 schema 类型语法

func 字段里的类型是仿 TorchScript IR 的语法，而不是 C++ 语法：

YAML	C++	含义
Tensor	const Tensor&	张量（不可空）
Tensor?	const std::optional<Tensor>&	可选张量
Tensor(a!)	Tensor&	可变引用，参数 a 可以被修改
Tensor[]	TensorList	张量列表
int	int64_t	整数
int[]	IntArrayRef	整数数组
Scalar	const Scalar&	标量（任意 dtype）
bool	bool	布尔
*	(kwargs 分隔符)	* 之后的参数必须用关键字传

Tensor(a!) 是 ATen 特有的"别名"标记，告诉系统这是 inplace 修改。它在 functionalization、autograd 等多处影响代码生成。

6.2.2 dispatch 表的写法

dispatch 字段的左边是 dispatch key 列表，右边是 C++ 函数名。多个 key 可以共用一个实现：

dispatch:
  CPU, CUDA, MPS: my_kernel                   # 三个 key 共用 my_kernel
  SparseCPU, SparseCUDA: my_sparse_kernel     # 稀疏后端另走一份实现

注意 dispatch 表不支持通配符，每个 backend key 都得逐个列出来 —— torchgen/model.py 解析这一栏时只按逗号切分并逐个查 DispatchKey 枚举。

如果不写 dispatch 字段，算子默认走 CompositeImplicitAutograd —— 也就是用其他算子组合实现的算子，autograd 自动正确（第 5 章 §5.2.2 提过）。这种"组合算子"通常不写自己的 kernel，而是转调其他算子，例如 absolute（abs 的别名）在 native_functions.yaml 里就只有 variants: function, method 而没有 dispatch 栏。aten/src/ATen/native/README.md 对此写得很明确："if you provide no dispatch table, we default to registering your kernel as CompositeImplicitAutograd"。

6.2.3 schema 中的"出参"与"别名规则"

PyTorch schema 有一套精细的"别名规则"用于功能化（functionalization）。看几个例子：

- func: foo(Tensor self) -> Tensor
  # foo 接收 self，返回新张量。无别名

- func: foo_(Tensor(a!) self) -> Tensor(a!)
  # foo_ 是 inplace 版本：self 被修改，返回的也是 self（同一个内存）

- func: foo.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)
  # foo.out 是 out= 版本：把结果写到 out，返回 out

- func: as_strided(Tensor(a) self, int[] size, int[] stride) -> Tensor(a)
  # 创建 view: 返回的张量与 self 共享内存（无 ! 表示只读共享）

(a) 是"别名但不写"标记（与输入共享 storage，但本算子不修改它），(a!) 是"可变别名"（共享并可能修改）。这套语法让 functionalization pass（第 13 章 AOTAutograd）能在编译期把所有 inplace 改写成纯函数式 —— 不需要运行期分析就能知道哪些张量参数是潜在的"修改源"。

理解这套语法是给 PyTorch 提交"自定义算子 PR"的基本功。新手最常见错误就是把 out 参数标成 Tensor 而不是 Tensor(a!)，导致 inplace 路径在 functionalize 时被错误优化。

6.3 derivatives.yaml：反向规则的声明

打开 tools/autograd/derivatives.yaml，这是一份纯数学声明，长这样：

- name: mm(Tensor self, Tensor mat2) -> Tensor
  self: mm_mat1_backward(grad, mat2, self.sym_sizes(), self.sym_strides(), self.layout(), 1)
  mat2: mm_mat2_backward(grad, self, mat2.sym_sizes(), mat2.sym_strides(), mat2.layout(), 1)
  result: at::mm(self_t, mat2_p) + at::mm(self_p, mat2_t)

——tools/autograd/derivatives.yaml 中 mm 的真实条目

字段含义：

• name：与 native_functions.yaml 中的 schema 完全一致
• self、mat2：每个输入参数的反向梯度表达式（被 torchgen 翻成 C++ 代码）
• result：前向 jacobian-vector product 表达式，用于 forward-mode AD（forward grad）

注意几个特殊符号：

符号	含义
grad	输出梯度（backward 输入）
self_t / mat2_t	输入参数的"切线"（forward AD）
self_p / mat2_p	输入参数的"基础值"（forward AD）
result	前向输出（公式里用到时由 codegen 保存以备反向用；多输出时写 result0、result1…）

代码生成器把这些表达式翻成 C++ 类（如 MmBackward0），第 7 章 autograd 章会展开。

注意 result 字段对 forward grad 的特殊语义：它返回的是输出的切线 (output tangent)，输入是 *_p（基础值）和 *_t（输入切线）。如果你看不懂这套数学，没关系 —— 新增算子时可以先只写 backward 字段（self、mat2 等）；逐元素或线性的算子还能直接把 forward 公式写成 auto_element_wise / auto_linear，由代码生成器自动推出。

6.3.1 反向规则与前向函数的对应关系

每个 derivatives.yaml 条目都必须对应一个 native_functions.yaml 里存在的算子。如果你写了一个新算子但没写反向规则，默认这个算子不能反向：前向照常执行，输出的 grad_fn 被设成 NotImplemented 节点，等到 backward() 时才报 "derivative for aten::xxx is not implemented"（tools/autograd/gen_variable_type.py 的 Note [ Registering AutogradNotImplemented boxed kernel ]）。

如果某算子是 composite（用其他可微算子组合实现），可以完全不写 derivatives 条目 —— autograd 通过子算子的反向自动推导。

如果某算子明确"不需要反向"，正确的做法不是加什么 tag —— aten/src/ATen/native/tags.yaml 里根本没有 "nondifferentiable" 这个标签（v2.11 全部合法 tag 只有 inplace_view、core、pointwise、nondeterministic_seeded 等 18 个）。真正的开关在 derivatives.yaml 一侧：把某个参数写成 self: non_differentiable，或者用 output_differentiability: [False] 声明输出不可导，代码生成器就不会为它生成 grad_fn。

6.3.2 derivatives.yaml 表达式的"语言"

derivatives.yaml 里的表达式 mm_mat1_backward(grad, mat2, ...) 不是任意 C++，而是一门特定子语言。它的特征：

• 所有变量都是 schema 里的参数名（self、mat2 等）
• 函数调用是 at::xxx(...) 或者一些 helper（maybe_multiply、mm_mat1_backward 等）
• 支持 ?: 三目、属性访问（self.sym_sizes()）、转 at:: 调用
• 不支持 if/for —— 控制流要写在 helper 函数里

如果反向规则非常简单（像 add 那样 grad），可以直接写表达式；如果复杂（如 softmax、layer_norm），就在 torch/csrc/autograd/FunctionsManual.cpp 里写 helper 函数，YAML 里只写一行调用。这种"声明式入口 + 实现式后端"分离让反向规则一眼读起来就是"数学公式"。

6.3.3 forward grad 的特殊性

注意 derivatives.yaml 里的 result 字段 —— 它不是反向 (backward) 而是 前向 AD (forward grad)。这是 PyTorch 较新的能力（torch/autograd/forward_ad.py 自 v1.8 起加入），允许用户做 jvp（Jacobian-vector product）而不是 vjp（vector-Jacobian product）。

forward grad 在大多数算子上是反向规则的"对偶"：mm 的反向给 self 的梯度是 grad @ mat2ᵀ，前向切线则是 self_t @ mat2_p + self_p @ mat2_t。derivatives.yaml 在同一个条目里把两套规则一起写，torchgen 一次生成两套 C++ 代码：反向公式变成 MmBackward0 节点类（Functions.cpp），前向公式则直接内联进 VariableType 里 mm 的包装函数，并没有单独的 "Jvp" 类。

对绝大多数 PyTorch 用户来说 forward grad 不常用（vjp 已经够用），但在二阶导数、Hessian-vector product 等场景很有用。第 7 章 autograd 那章会展开。

6.4 structured 算子：现代算子的统一框架

把目光回到 mm.out 那个条目：structured: True。这是 PyTorch 算子的现代框架，统一处理算子的"前奏"。

传统手写算子的 boilerplate（伪代码）：

Tensor mm_out_cpu(const Tensor& self, const Tensor& mat2, Tensor& out) {
    // 1. shape 检查与广播
    TORCH_CHECK(self.dim() == 2 && mat2.dim() == 2);
    TORCH_CHECK(self.size(1) == mat2.size(0));
    auto out_shape = {self.size(0), mat2.size(1)};

    // 2. dtype 检查与提升
    auto common_dtype = at::result_type(self, mat2);
    TORCH_CHECK(common_dtype == kFloat || common_dtype == kDouble || ...);

    // 3. 分配输出
    out.resize_(out_shape);
    if (out.dtype() != common_dtype) out = out.to(common_dtype);

    // 4. 真正的计算
    // ... CPU GEMM kernel ...
    return out;
}

每个算子都要写这一坨。structured 框架把它拆开：shape / dtype 推导集中写进一个 meta 函数，其余样板自动生成：

graph TB
    YAML["native_functions.yaml<br/>mm.out: structured True"]
    YAML --> Gen["torchgen 解析"]
    Gen --> Meta["生成 structured_mm 类骨架<br/>(functional / out / inplace 三种输出分配方式)"]
    Gen --> Impl["生成 wrapper<br/>先调 meta 再调 impl"]
    User["开发者写: TORCH_META_FUNC(mm)<br/>+ TORCH_IMPL_FUNC(mm_out_cpu)"] -.被 wrapper 调.-> Impl
    Meta --> Final["完整算子<br/>meta + impl 自动注册到 dispatcher"]
    Impl --> Final
    style YAML fill:#fef3c7
    style Gen fill:#dbeafe
    style User fill:#dcfce7
    style Final fill:#fce7f3

structured 框架的核心：把 shape 检查 / dtype 推导收拢进一个与后端无关的 meta 函数，把分配输出、校验并 resize 用户给的 out、包装与注册这些样板交给代码生成，让各后端的 impl 只写"真正的计算"——这与 §6.5 的 ufunc 思路一致，都是"代码生成抽象掉重复"。

// aten/src/ATen/native/LinearAlgebra.cpp:202 —— meta 函数由开发者写在 TORCH_META_FUNC 里，
// torchgen 负责把它包装成 structured_mm 基类并接到 dispatcher 上
TORCH_META_FUNC(mm)(const Tensor & self, const Tensor & mat2) {
  TORCH_CHECK(self.dim() == 2, "self must be a matrix");
  TORCH_CHECK(mat2.dim() == 2, "mat2 must be a matrix");
  TORCH_CHECK(
      self.sizes()[1] == mat2.sizes()[0], "mat1 and mat2 shapes cannot be multiplied (", ...);

  auto names = at::namedinference::compute_matmul_outnames(self, mat2);
  set_output_raw_strided(0, {self.sizes()[0], mat2.sizes()[1]}, {}, self.options(), names);
}

// impl 函数：每个 backend 一份纯计算，out 已经由框架分配好了
TORCH_IMPL_FUNC(mm_out_cpu)(const Tensor & self, const Tensor & mat2, const Tensor & result) {
  {
    at::NoNamesGuard guard;
    addmm_impl_cpu_(const_cast<Tensor&>(result), result, self, mat2, 0, 1);
  }
}

——aten/src/ATen/native/LinearAlgebra.cpp 里 mm 的 meta（TORCH_META_FUNC(mm)）与 CPU impl（TORCH_IMPL_FUNC(mm_out_cpu)）真实代码（省略了报错字符串拼接）。注意 mm 不是 TensorIterator 系算子，它的 native_functions.yaml 条目里没有 structured_inherits；带 structured_inherits: TensorIteratorBase 的是 add.out 那类逐元素算子。

——开发者写一份 meta（shape 检查、输出元信息）和每个后端一份 impl；输出分配、out 校验与 resize、wrapper 与 dispatcher 注册全部代码生成，meta 也只写一次、被所有后端和 functional / out / inplace 三种变体共用。这在大量算子上节省了海量重复代码。

6.4.1 structured_delegate 的转发机制

回到 mm（无 out）：

- func: mm(Tensor self, Tensor mat2) -> Tensor
  structured_delegate: mm.out

这告诉 torchgen：mm 自己不写 kernel，复用 mm.out 那一组的 meta 与 impl。生成在 RegisterCPU_*.cpp 里的代码大致是（简化，出自 torchgen/dest/register_dispatch_key.py 的 structured 分支）：

// 生成的 mm 在 CPU 上的 wrapper（简化）
struct structured_mm_out_cpu_functional final : public at::native::structured_mm_out_cpu {
    // set_output_raw_strided 被覆写为"新分配输出张量"
    std::array<Tensor, 1> outputs_;
};
at::Tensor wrapper_CPU_mm(const at::Tensor& self, const at::Tensor& mat2) {
    structured_mm_out_cpu_functional op;
    op.meta(self, mat2);                   // TORCH_META_FUNC(mm)：检查形状并分配输出
    op.impl(self, mat2, op.outputs_[0]);   // TORCH_IMPL_FUNC(mm_out_cpu)：真正计算
    return std::move(op.outputs_[0]);
}

注意这里不会再经 dispatcher 调一次 mm.out，而是直接调同一组 meta / impl。

这种"无 out 版本委托给 out 版本"模式让所有算子的"无 out / 有 out"两个变体共享同一份核心实现，消除一半重复。

6.4.2 structured 框架与 inplace 三态

一个 structured 算子组最多有三个变体（下面借 mm 的名字示意）：

• mm(a, b) -> result —— 函数式：返回新张量
• mm.out(a, b, *, out) —— out= 版本：写入用户给的 out
• mm_(a, b) —— inplace：修改 a 自身（如果可行）

这三个变体共享同一个 meta（计算 shape）和 impl（真正的 kernel）。代码生成器为三个变体各生成一个 wrapper，区别只在输出从哪来（新分配 / 用户给的 out / self 本身），meta 与 impl 都复用 mm.out 那一份。用户只需写一份 meta + 每后端一份 impl，三个 API 自动具备。

不是所有算子都有这三个变体。mm 没有 mm_ 因为矩阵乘的输出形状与输入不同（不能 inplace）。但加减乘除这种"形状不变"的算子三态都有。native_functions.yaml 通过分别声明三个 schema 让代码生成器知道每个算子有哪些变体。

6.4.3 meta 函数与 FakeTensor 的关系

structured 框架里 meta 函数的存在不只是工程整洁，它是 FakeTensor / torch.compile 的关键依赖。

FakeTensor 的语义是"假执行" —— 只算 shape / dtype / stride，不真做数值计算。它怎么知道每个算子的输出 shape？答案就是 调 meta 函数。torch.compile 在 trace 阶段把每个 ATen 调用 redirect 到对应的 meta，meta 函数算出输出 shape 后返回一个 FakeTensor，trace 继续往下走。

这就是为什么 §6.13 那条"meta 里不能有计算"的红线如此关键 —— 一旦 meta 错误地做了真计算，FakeTensor 的"零成本 shape 推导"就垮了。第 12-14 章 torch.compile 编译器栈会反复看到 meta 函数的影子。

6.4.4 结构化算子的演进

structured 框架不是 PyTorch 一开始就有的。早期所有算子都是手写完整流水线，重复严重。2020-2021 年间社区提出并落地了 "structured kernels" 方案，把 ATen 算子逐步向 structured 迁移；今天它在源码里的入口是 torchgen/api/structured.py、torchgen/api/meta.py 与 torchgen/dest/register_dispatch_key.py 这几个生成模块，以及 TORCH_META_FUNC / TORCH_IMPL_FUNC 这对宏。到 v2.11，native_functions.yaml 里有 273 条 structured: True 和 383 条 structured_delegate:，合计约 656 条 —— 相对 2674 条 - func: 声明，大约四分之一的条目走 structured 路线；余下大量算子（尤其是 RNG、量化、稀疏、各类 composite）仍是手写风格。

理解这个迁移历史，你看源码时能区分"老式手写算子"和"现代 structured 算子"两种风格。两者常常同处 aten/src/ATen/native/ 下的同一个文件（如 aten/src/ATen/native/LinearAlgebra.cpp），认 TORCH_META_FUNC / TORCH_IMPL_FUNC 这对宏即可区分。

6.5 ufunc 内层循环：一个公式生成几十个 kernel

第 1 章 §1.5 介绍过 add 的真正数学藏在 aten/src/ATen/native/ufunc/add.h：

// aten/src/ATen/native/ufunc/add.h:14（省略了 __ubsan_ignore_undefined__ 标注）
template <typename T>
C10_HOST_DEVICE C10_ALWAYS_INLINE T add(T self, T other, T alpha) {
    return self + alpha * other;
}

native_functions.yaml 里 add.out 条目带这一段：

- func: add.out(Tensor self, Tensor other, *, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!)
  structured: True
  structured_inherits: TensorIteratorBase
  ufunc_inner_loop:
    Generic: add (AllAndComplex, BFloat16, Half, ComplexHalf)
    ScalarOnly: add (Bool)

ufunc_inner_loop 字段告诉代码生成器：用 ufunc::add 这个内层函数，给它生成 SIMD 向量化的 CPU kernel + CUDA kernel。torchgen/dest/ufunc.py 会把它展开成构建目录里的三个文件（下面是简化后的样子）：

// build/aten/src/ATen/UfuncCPU_add.cpp：CPU 的 structured impl 只是转给 add_stub
TORCH_IMPL_FUNC(ufunc_add_CPU)(const Tensor& self, const Tensor& other, const Scalar& alpha, const Tensor& out) {
    add_stub(device_type(), *this, alpha);
}

// build/aten/src/ATen/UfuncCPUKernel_add.cpp：真正的 CPU kernel，注册到 add_stub
namespace {
void add_kernel(TensorIteratorBase& iter, const at::Scalar& alpha) {
    AT_DISPATCH_SWITCH(iter.common_dtype(), "add_stub",
        AT_DISPATCH_CASE(at::ScalarType::Float, [&]() {
            auto _s_alpha = alpha.to<scalar_t>();
            auto _v_alpha = at::vec::Vectorized<scalar_t>(_s_alpha);
            cpu_kernel_vec(iter,
                [=](scalar_t self, scalar_t other) { return ufunc::add(self, other, _s_alpha); },
                [=](at::vec::Vectorized<scalar_t> self, at::vec::Vectorized<scalar_t> other) {
                    return ufunc::add(self, other, _v_alpha);
                });
        })
        /* ... 其余每个 dtype 一个 AT_DISPATCH_CASE ... */
    );
}
} // anonymous namespace
REGISTER_DISPATCH(add_stub, &add_kernel)

// build/aten/src/ATen/UfuncCUDA_add.cu：CUDA 侧生成 functor，impl 直接调 add_kernel，不经 add_stub
template <typename scalar_t>
struct CUDAFunctor_add {
    using opmath_t = at::opmath_type<scalar_t>;
    opmath_t alpha_;
    CUDAFunctor_add(opmath_t alpha) : alpha_(alpha) {}
    __device__ scalar_t operator()(scalar_t self, scalar_t other) {
        return ufunc::add(static_cast<opmath_t>(self), static_cast<opmath_t>(other), alpha_);
    }
};
// （另有 CUDAFunctorOnSelf_add / CUDAFunctorOnOther_add，处理某个输入是 CPU 标量的情形）
void add_kernel(TensorIteratorBase& iter, const at::Scalar& alpha) {
    AT_DISPATCH_SWITCH(iter.common_dtype(), "ufunc_add_CUDA",
        AT_DISPATCH_CASE(at::ScalarType::Float, [&]() {
            using opmath_t = at::opmath_type<scalar_t>;
            /* if (iter.is_cpu_scalar(1)) {...} else if (iter.is_cpu_scalar(2)) {...} else */
            gpu_kernel(iter, CUDAFunctor_add<scalar_t>((alpha).to<opmath_t>()));
        })
        /* ... */
    );
}
TORCH_IMPL_FUNC(ufunc_add_CUDA)(const Tensor& self, const Tensor& other, const Scalar& alpha, const Tensor& out) {
    add_kernel(*this, alpha);
}

一个几行的内层公式 → 13 种 dtype × CPU / CUDA 两个后端的 kernel 特化（AllAndComplex 9 种加 BFloat16、Half、ComplexHalf，再加 ScalarOnly 的 Bool）。这是 PyTorch 为逐元素算子设计的 ufunc 代码生成路线 —— 数学公式与 SIMD 化、CUDA 化、dtype 展开完全解耦。不过要注意，截至 v2.11，native_functions.yaml 里带 ufunc_inner_loop 的只有 add.out 这一条（aten/src/ATen/native/ufunc/ 下也只有 add.h）；mul、sub 等其余逐元素算子仍是手写的 TensorIterator kernel，经 *_stub 分发（如 aten/src/ATen/native/cpu/BinaryOpsKernel.cpp 里的 mul_kernel）。

6.5.1 ufunc 与 NumPy 的渊源

ufunc 这个名字直接借自 NumPy —— "universal function"。NumPy 的 ufunc 也是"对每个元素应用同一函数"的逐元素抽象，由 NumPy 自家代码生成器生成各 dtype 的 SIMD 实现。

PyTorch 的 ufunc 与 NumPy 同名同源，但实现细节有差异：

• NumPy 用 C 模板生成；PyTorch 用 C++ 模板 + Python codegen
• PyTorch 的 ufunc 同时支持 CPU SIMD 与 CUDA，NumPy 只有 CPU
• PyTorch 把 ufunc 与 TensorIterator 紧密集成，NumPy 用自家的 nditer

但思想完全一致：让用户写一行数学公式，机器把它扩展成几十个特化版本。这是数值计算库通用的工程模式。

6.5.2 ufunc 与 vec::Vectorized

在 ufunc 内层循环代码里你会反复看到 Vectorized<T> 这个类型 —— PyTorch 自家的 SIMD 抽象（aten/src/ATen/cpu/vec/）。它对每个 dtype 提供 Vectorized<float>、Vectorized<bfloat16> 等模板特化，每个特化在 x86 上用 AVX2/AVX512、ARM（含 Apple Silicon）上用 NEON / SVE 实现（分别在 vec256/、vec512/、vec128/、sve/ 子目录）。

ufunc 函数同时支持标量版本和 Vectorized 版本：

template <typename T>
T add(T self, T other, T alpha) { return self + alpha * other; }   // 标量

template <typename T>
Vectorized<T> add(Vectorized<T> self, Vectorized<T> other, Vectorized<T> alpha) {
    return vec::fmadd(other, alpha, self);   // 向量化版，用 FMA 指令
}

cpu_kernel_vec 在内层维度连续时按整块（每次两个 Vectorized 宽度）用向量版，对剩余的尾部元素用标量版；步长不连续时整段退回标量版（aten/src/ATen/native/cpu/Loops.h 的 vectorized_loop）。这种"vector + scalar tail"是 SIMD 编程的标准模式。

Vectorized<T> 还提供完整的数学 API（abs、min、max、log、exp、sin、cos、…），让用户写 ufunc 不需要自己处理 SIMD 细节。这是 PyTorch CPU 性能在数十种 dtype 上保持一致的根本工具。

6.5.3 哪些算子能 ufunc 化

不是所有算子都能写成 ufunc。要求是：

• 逐元素：每个输出元素只依赖对应位置的输入元素（最多再加广播）
• 无内部状态：没有累加、归约、依赖前一个输出
• 数学公式可向量化：能在 SIMD 寄存器里并行算

加减乘除、比较运算、激活函数（ReLU / Sigmoid / Tanh）、绝大多数 element-wise 数学函数（exp / log / sin / cos）都符合（只是如前所述，v2.11 里真正迁到 ufunc 的只有 add）。但 matmul、conv、softmax、layer_norm、reduce 都不行 —— 它们涉及跨元素的依赖。这类算子要手写 kernel + AT_DISPATCH_* 宏来处理 dtype 展开。

6.6 AT_DISPATCH_* 宏家族：手写算子的"开关"

不是所有算子都能 ufunc。复杂的算子（如卷积、attention、归约）需要手写。但手写算子也需要"对所有 dtype 展开"—— 这就是 AT_DISPATCH_* 宏的用武之地。

AT_DISPATCH_ALL_TYPES_AND2(kBFloat16, kHalf, iter.common_dtype(), "my_op", [&]() {
    // 在这个 lambda 里，宏已经把 scalar_t 定义为当前分支的具体 dtype
    // 写算子实现
});

宏展开后是一个巨大的 switch：

switch (iter.common_dtype()) {
    case kFloat:    { using scalar_t = float;    /* lambda 体 */ break; }
    case kDouble:   { using scalar_t = double;   /* lambda 体 */ break; }
    case kBFloat16: { using scalar_t = at::BFloat16; /* lambda 体 */ break; }
    case kHalf:     { using scalar_t = at::Half; /* lambda 体 */ break; }
    case kInt:      { using scalar_t = int32_t;  /* lambda 体 */ break; }
    case kLong:     { using scalar_t = int64_t;  /* lambda 体 */ break; }
    /* ... */
}

每个 case 分支编译时实例化一份 lambda 副本，对应那个 dtype 的特化代码。最终二进制里有 N 份针对不同 dtype 编译过的 kernel —— 运行时一次 switch 跳到正确版本。

AT_DISPATCH_* 有几十个变体：AT_DISPATCH_ALL_TYPES、AT_DISPATCH_FLOATING_TYPES、AT_DISPATCH_INTEGRAL_TYPES 等等。它们的区别就是 switch case 覆盖的 dtype 集合不同。这套宏让"运行期 dtype 选择"只剩一次 switch 的开销 —— 编译期把所有可能性展开成静态分支。

6.6.1 dispatch 宏的二进制膨胀代价

AT_DISPATCH_* 的代价是 二进制膨胀：每个用了这套宏的算子实现，编译产出 N 份代码（每个 dtype 一份）。AT_DISPATCH_CASE_ALL_TYPES 本身等于 INTEGRAL_TYPES（uint8/int8/int16/int32/int64，5 种）加 FLOATING_TYPES（float/double，2 种），共 7 份；再用 AT_DISPATCH_ALL_TYPES_AND3(...) 追加 3 个 dtype，就是 10 份特化。乘上几百个算子，整个 libtorch 的体积有相当一部分来自这种"特化爆炸"。

PyTorch 对此有一套"选择性构建"机制：aten/src/ATen/Dispatch.h 开头判断 TEMPLATE_SELECTIVE_BUILD 宏，打开时 include 由构建流程生成的 ATen/selected_mobile_ops.h，其中的 at::should_include_kernel_dtype(kernel_tag, dtype) 会在编译期把没被模型用到的 dtype 分支整段裁掉（AT_PRIVATE_CHECK_SELECTIVE_BUILD 用 if constexpr 判断，未选中的 dtype 直接变成一句 TORCH_CHECK(false, ...)）。不打开该宏时这个函数恒返回 true，全部 dtype 都编进去。这一层裁剪是移动端产物瘦身的手段之一（另一层是 torchgen 生成注册代码时按 selector 只保留选中的算子）—— 不是少了功能，而是少了用不到的 dtype 特化。

6.6.2 SymInt：符号 shape 的代码生成支持

SymInt 是一个能表示"具体整数"或"符号变量"的类型（2022 年由 PR #74861 引入，首个包含它的正式版是 v1.12）。它由 torch.compile / Dynamo 在 trace 阶段使用，让算子能接受 dynamic shapes 而不立即具体化。

代码生成器对 schema 里带 SymInt 的算子自动生成 SymInt 版本。在 native_functions.yaml 里你会看到很多算子有 SymInt[] 参数（如 view(Tensor(a) self, SymInt[] size)），torchgen 据此生成两套 C++ API 入口：at::reshape / Tensor::view 接收普通的 IntArrayRef，at::reshape_symint / Tensor::view_symint 接收 c10::SymIntArrayRef。前者只是把整数包成 SymInt（c10::fromIntArrayRefSlow）后调同一个 at::_ops::xxx::call —— dispatcher 这一层只有 SymInt 签名一套，符号 shape 因此能在 trace 时一路传到 Inductor。

第 14 章 TorchInductor 章会展开 SymInt 怎么让 dynamic shape 编译成为可能。

6.7 torchgen 工作流

把视角拉高，看 torchgen 是怎么把 YAML 变成 C++ 的。打开 torchgen/gen.py（全文 3067 行，主入口）和 torchgen/api/ 目录的多个文件，整体流程：

flowchart LR
    Y1[native_functions.yaml]
    Y2[derivatives.yaml]
    Y3[tags.yaml]
    Y1 --> Parse[gen.py: parse_native_yaml<br/>解析为 NativeFunction 对象]
    Y3 --> Parse
    Y2 --> LD[tools/autograd/load_derivatives.py<br/>解析反向规则]
    Parse --> Model[torchgen/model.py<br/>Python 中间表示<br/>NativeFunction / FunctionSchema]
    Model --> CG1["torchgen/api/cpp.py<br/>Tensor 方法 / at::xxx 签名"]
    Model --> CG2["torchgen/gen.py<br/>Operators_0..4.cpp / TensorBody.h"]
    Model --> CG3["torchgen/dest/register_dispatch_key.py<br/>RegisterCPU/CUDA/..."]
    Model --> CG4["tools/autograd/gen_variable_type.py<br/>VariableType_*.cpp"]
    Model --> CG5["tools/autograd/gen_python_functions.py<br/>Python 包装"]
    Model --> CG6["torchgen/gen_functionalization_type.py<br/>functionalize 规则"]
    Model --> CG7["torchgen/gen_vmap_plumbing.py<br/>vmap plumbing"]
    LD --> CG4
    style Model fill:#fef3c7,stroke:#f59e0b,stroke-width:2px

torchgen 使用 Jinja-like 模板 生成代码（自家轻量级 code_template.py）。整个生成过程在 PyTorch 编译时由 CMake/setup.py 触发，产物落到 build/aten/src/ATen/ 与 torch/csrc/autograd/generated/。

6.7.1 几个 gen_*.py 的分工

torchgen 里的几个核心生成模块（注意 autograd 与 Python 绑定这两块的生成器不在 torchgen/ 而在 tools/autograd/）：

文件	职责
torchgen/gen.py	主入口，编排所有生成步骤
torchgen/api/cpp.py	C++ API 签名生成（at::add、Tensor::add）
torchgen/api/dispatcher.py	dispatcher 调用约定（boxed / unboxed signatures）
torchgen/api/native.py	native 层签名（实现侧）
torchgen/api/structured.py	structured 算子的 meta + impl 框架
torchgen/api/python.py (1550 行)	Python 端绑定生成 (THPVariable_xxx)
torchgen/api/ufunc.py	ufunc 的签名与参数绑定（展开代码在 torchgen/dest/ufunc.py）
torchgen/api/translate.py	不同 binding 集合之间的翻译
torchgen/api/functionalization.py	functionalize 规则生成
torchgen/api/lazy.py	Lazy Tensor backend 支持
torchgen/dest/register_dispatch_key.py	各 backend 的 RegisterXxx_N.cpp 输出（含 structured wrapper）
torchgen/gen_functionalization_type.py	RegisterFunctionalization_0..3.cpp
torchgen/gen_backend_stubs.py	外部 backend（PrivateUse1）的 stub 生成
tools/autograd/gen_variable_type.py	VariableType_*.cpp（autograd 包装）
tools/autograd/gen_autograd_functions.py	Functions.cpp（MmBackward0 这类节点类）
tools/autograd/gen_python_functions.py	python_torch_functions_*.cpp / python_variable_methods.cpp

每个文件都有清晰职责。理解这套划分让你看 PyTorch 编译产物时能快速定位"这段生成代码出自哪里"。

6.7.2 一个具体跟踪：mm 在 torchgen 里的生命周期

跟踪 mm 算子从 YAML 到二进制：

1. 解析：parse_native_yaml 读 native_functions.yaml，把 mm 与 mm.out 解析成两个 NativeFunction 对象
2. 按 dispatch key 建索引：把 dispatch 字段按 key 拆进各自的 BackendIndex（CompositeExplicitAutograd 这类 alias key 也只是其中一个索引；把 alias key 展开到具体 backend，是运行期 dispatcher 在 aten/src/ATen/core/dispatch/OperatorEntry.cpp 里计算 dispatch 表时做的）
3. Operators_3.cpp：生成 at::_ops::mm::call 入口，里面用 Dispatcher::singleton().findSchemaOrThrow(mm::name, mm::overload_name).typed<mm::schema>() 取得算子句柄（缓存在函数内 static 变量里）再 op.call(...)；这里两个参数是 ("aten::mm", "")，第 1 章 §1.3 见过的 add 则是 ("aten::add", "Tensor")
4. RegisterCPU_0.cpp：torchgen/dest/register_dispatch_key.py 的 structured 分支生成 wrapper 函数 wrapper_CPU_mm（functional）与 wrapper_CPU_mm_out_out（out 版，命名是 wrapper_{DispatchKey}_ 加 native 名 mm_out 再加重载名 _out；非 structured 算子才用 wrapper_{DispatchKey}_{overload_name}_{基名} 的模板），再在 TORCH_LIBRARY_IMPL(aten, CPU, m) 里 m.impl("mm.out", TORCH_FN(wrapper_CPU_mm_out_out))
5. TensorBody.h：生成 Tensor::mm(const at::Tensor & mat2) const { return at::_ops::mm::call(const_cast<Tensor&>(*this), mat2); }
6. VariableType_3.cpp：生成 autograd 包装 VariableType::mm，使用 derivatives.yaml 翻译出的 MmBackward0 节点
7. python_torch_functions_1.cpp：生成 THPVariable_mm Python 函数，能从 Python 端调用

（分片号由 torchgen/utils.py 的 string_stable_hash(键名) % 分片数 决定；RegisterCPU_0.cpp、VariableType_3.cpp 可用已装 wheel 的 torch._C._dispatch_dump("aten::mm") 输出核对。）

每一步都有专门的 gen_*.py 文件负责。整个 torchgen 大约 25000 行 Python，但它的产出是数十万行 C++ —— 代码膨胀超过一个数量级。

6.7.3 torchgen 模板系统

torchgen 不依赖 Jinja2 等外部库 —— 它有自己的轻量级模板（torchgen/code_template.py，全文 109 行）。模板语法形如：

TEMPLATE = CodeTemplate("""
Tensor ${name}(${arguments}) {
    ${body}
}
""")

result = TEMPLATE.substitute(
    name="my_op",
    arguments="const Tensor& self",
    body="return self.add(1);"
)

风格与 Python 的 string.Template 类似，但值为列表时能按所在行的缩进展开成多行，或用 ${,foo} / ${foo,} 拼成逗号分隔的参数表。源码里没有交代为什么不用 Jinja2；值得一提的是，jinja2 如今已在 torch 的 install_requires 里（Inductor 的 .jinja 模板要用），torchgen 仍沿用这套百余行的自家模板。

6.7.4 一个有趣的设计：translate

torchgen/api/translate.py（437 行）是 torchgen 里一个特别巧妙的模块。它的工作是把"一种 binding 集合"翻成"另一种"。

举例：C++ API 允许把 dtype / layout / device / pin_memory 打包成一个 TensorOptions options 传入，dispatcher 调用约定却要求它们分开传；structured 的 wrapper 调 meta / impl 时，又要把 wrapper 手里的参数和 op.outputs_[0] 这类表达式对到 meta / impl 的形参上。生成这些调用代码时，需要把"手头有的表达式"拼成"对方要的参数" —— translate.py 就是这道题的引擎。

它以 NamedCType（带语义名的 C++ 类型）为单位，自动推断怎么用上下文里已有的表达式构造出目标参数：有的只是同名直取，有的要打包 / 拆包（源码注释里的 Gather、Scatter、Join，如从 options 里取出 dtype），还有 std::optional<Tensor> 到 OptionalTensorRef、SymInt 与 int64_t 之间的转换。这套机制让 PyTorch 能在 dispatcher / native / Python / functionalize 等多种约定之间自如转换，开发者不用手写每一对的转换代码。

6.7.5 torchgen 的 Python 中间表示

torchgen/model.py 定义了一组 dataclass 表示算子的"AST"：

@dataclass(frozen=True)
class NativeFunction:
    namespace: str                    # "aten"
    func: FunctionSchema              # 解析后的 schema
    variants: set[Variant]
    autogen: list[OperatorName]
    ufunc_inner_loop: dict[UfuncKey, UfuncInnerLoop]
    structured: bool
    structured_delegate: OperatorName | None
    structured_inherits: str | None
    tags: set[str]
    ...

——torchgen/model.py 中 NativeFunction 的部分字段。注意 dispatch 表并不挂在 NativeFunction 上：YAML 里的 dispatch 栏被抽出来放进独立的 BackendIndex / BackendMetadata，每个 dispatch key 一份索引，NativeFunction 上只留 has_composite_implicit_autograd_kernel 这类布尔摘要。这样"算子定义"与"后端实现映射"是两张可以分别演化的表。

YAML 解析完后，每个算子是一个 NativeFunction 实例。后续所有代码生成函数都接收 NativeFunction 列表，遍历产出代码字符串。这种"先解析成模型对象，再多 pass 生成"的设计与编译器后端的 IR 思想完全一致。第 14 章 TorchInductor 章会再看到这种 IR-Lowering 模式。

6.8 一个真实例子：从 YAML 到二进制的完整链

以 add 为例，把整条链路画出来：

flowchart TB
    subgraph Source["源码 (静态)"]
        Y1["native_functions.yaml<br/>add.Tensor / add.out"]
        Y2["derivatives.yaml<br/>add 反向规则"]
        Y3["aten/src/ATen/native/ufunc/add.h<br/>self + alpha*other"]
        Y4["aten/src/ATen/native/BinaryOps.cpp<br/>structured meta"]
    end
    subgraph Codegen["编译期 torchgen"]
        TG[torchgen/gen.py 跑]
    end
    subgraph Generated["生成的 C++ (build/)"]
        G1[Operators_2.cpp<br/>at::_ops::add_Tensor::call]
        G2[RegisterCPU_0.cpp<br/>m.impl 注册]
        G3[RegisterCUDA_0.cpp]
        G4[VariableType_2.cpp<br/>autograd 包装]
        G5[python_torch_functions_2.cpp<br/>Python 绑定]
        G6[UfuncCPUKernel_add.cpp<br/>SIMD 展开]
        G7[UfuncCUDA_add.cu<br/>CUDA 展开]
    end
    subgraph Binary["编译后 (libtorch_cpu.so 等)"]
        B1[at::add 符号]
        B2[Tensor::add 符号]
        B3[VariableType::add 符号]
        B4[THPVariable_add 符号]
        B5[各 dtype × 后端的 kernel]
    end
    Y1 --> TG
    Y2 --> TG
    TG --> G1
    TG --> G2
    TG --> G3
    TG --> G4
    TG --> G5
    TG --> G6
    TG --> G7
    Y3 -.被 include.-> G6
    Y3 -.被 include.-> G7
    Y4 -.手写,直接编译.-> Binary
    G1 --> B1
    G1 --> B2
    G4 --> B3
    G5 --> B4
    G6 --> B5
    G7 --> B5
    style TG fill:#fef3c7,stroke:#f59e0b,stroke-width:2px

这条链上的关键观察：

• 静态 YAML → 自动膨胀到二进制：一条 YAML 条目 → 分散在多个生成文件里的 C++ → 各 dtype × 后端的机器码
• 手写部分极少：开发者只需写 YAML 条目、meta 与内层公式（或各后端 impl）、反向公式这几段，其余全是模板/生成
• 跨语言一致性自动保证：Python torch.add、C++ at::add、Tensor::add 永远拥有完全一致的签名 —— 因为它们都从同一个 schema 派生
• 二进制可观察：调试时可以 nm -C libtorch_cpu.so | grep wrapper_CPU_add 看到生成的 wrapper 符号（wheel 里的 libtorch.so 只是个小壳，ATen 代码在 libtorch_cpu.so / libtorch_cuda.so）

6.8.1 编译时间的代价

代码生成虽然减少手写，但有一个常被忽略的代价：编译时间。从源码全量编译 PyTorch 在普通开发机上是以小时计的事，其中相当一部分花在编译这些"机器生成的几十万行 C++"上。具体痛点：

• 生成的 Operators_*.cpp 单文件动辄上万行，且每个都要展开大量模板
• 模板展开与重载决议是编译器内存占用的大头，并行编译时容易把机器内存吃满
• 一次性 link 所有生成的 .o 时，linker 是另一个瓶颈

PyTorch 团队为了缓解，做了几个工程优化：

1. 拆分生成文件：torchgen/gen.py 用 cpu_fm.write_sharded("Operators.cpp", ..., num_shards=5) 把全部算子切成 Operators_0.cpp 到 Operators_4.cpp 五块（RegisterFunctionalization.cpp 同理切 4 块），让 make 能并行编
2. AT_PER_OPERATOR_HEADERS 宏：每个算子单独一个 .h 文件，用户代码可以"只 include 自己用到的"，减少 amplification
3. 预编译头：常用 c10 / aten 头文件做成 PCH

这些优化把全量编译时间压下来不少，但对大多数 PyTorch 改造工程来说，"编译一次去喝杯咖啡"仍是常态。

6.9 添加一个新算子：实战流程

把上面所有抽象落到工程实战。要在 PyTorch 主仓加一个新算子（比如 my_clamp），完整流程：

1. 在 native_functions.yaml 加 schema：

- func: my_clamp(Tensor self, Scalar min, Scalar max) -> Tensor
  variants: function, method
  dispatch:
    CPU: my_clamp_cpu
    CUDA: my_clamp_cuda
  tags: pointwise

2. 在 aten/src/ATen/native/MyClamp.cpp 写 CPU impl：

Tensor my_clamp_cpu(const Tensor& self, const Scalar& min, const Scalar& max) {
    // 真正的 CPU 实现
}

3. 在 aten/src/ATen/native/cuda/MyClamp.cu 写 CUDA impl：

Tensor my_clamp_cuda(const Tensor& self, const Scalar& min, const Scalar& max) {
    // 真正的 CUDA 实现
}

4. 在 derivatives.yaml 加反向规则（如果可微）：

- name: my_clamp(Tensor self, Scalar min, Scalar max) -> Tensor
  self: clamp_backward(grad, self, min, max)

5. 重新跑构建：python -m pip install --no-build-isolation -v -e .（v2.11 里 python setup.py develop 也会被重定向到这条命令）。torchgen 自动检测 YAML 变化，重新生成所有派生文件

6. 在 Python 端使用：torch.my_clamp(t, -1, 1) 就能调

注意 6 步里除了第 1 步的 YAML 声明，只有 2、3、4 是手写（impl + 反向公式），其余全部由代码生成器 / 构建系统接管。绝大多数代码由生成器代劳，开发者只写"算子语义"。这是 PyTorch 工程化的最重要范本之一。

6.9.1 第三方库的 torch.library API

如果你不想改 PyTorch 主仓（大部分场景），可以用 torch.library Python API 在自己的项目里注册算子：

import torch

# 在 Python 端定义算子
@torch.library.custom_op("mylib::my_clamp", mutates_args=())
def my_clamp(x: torch.Tensor, min: float, max: float) -> torch.Tensor:
    return torch.clamp(x, min, max)

@my_clamp.register_fake
def _(x, min, max):
    # FakeTensor / shape 推导版本（用于 torch.compile）
    return torch.empty_like(x)

这套 API 由 PR #122344 引入，首次随 v2.4.0 发布，让用户不需要 C++、不需要修改 PyTorch 源码就能注册新算子，同时还能与 torch.compile 完美兼容。它的底层实现仍然是调用 dispatcher 的 registerImpl 注册接口，但把繁琐的 schema 解析、注册流程封装成 Python 函数。

如果你的需求是"在不改 PyTorch 主仓的前提下加一个算子"，这是最佳路径。第 22 章自定义算子章会详细演示。

6.9.2 国产芯片厂商怎么用这套机制

国产 AI 芯片厂商（华为昇腾、寒武纪、壁仞等）以树外插件接入 PyTorch 时，常走同一条路：

1. 声明自己的 dispatch key：通过 PyTorch 的 PrivateUse1 机制注册一个新的 backend key
2. 跑 codegen：torchgen 提供 gen_backend_stubs.py，按厂商列出的算子清单生成 C++ 声明头与注册代码
3. 逐个实现：把 stub 列表里的算子一个个用自家硬件 SDK 实现
4. 作为独立 wheel 发布：用户 pip install torch_npu / pip install torch_mlu 等就能在自家芯片上跑 PyTorch

这套流程让国产芯片不需要 fork PyTorch 主仓就能完整接入。gen_backend_stubs.py 是这条路上的关键 —— 它读一份"我要支持的算子列表" YAML，吐出对应的 C++ stub 模板，厂商只填实现细节。

理解这条机制，你就能解释为什么国产芯片接入 PyTorch 的速度比想象的快 —— 厂商真正要写的只有 kernel 本身，签名、注册、autograd 包装这些胶水全由代码生成器代劳。

gen_backend_stubs.py 的工作流大致是：读厂商提供的"目标 backend YAML"（声明它支持哪些算子、是否支持 autograd 等），自动生成对应的 Register{Backend}.cpp 注册表与 stub 头文件。厂商团队拿到这些 stub 后只需要逐个填实现 —— 类型签名、dispatcher 注册、autograd 包装全自动。这条"以 codegen 降低生态准入门槛"的设计，让新硬件接入 PyTorch 的成本大为降低。

6.10 横向对比：其他框架怎么做

框架	算子注册方式	优劣
PyTorch	YAML + Python codegen	自动化高、跨语言一致；调试需要看生成代码
TensorFlow	C++ REGISTER_OP 宏 + Bazel codegen	编译时检查严格；YAML 缺失，schema 不集中
JAX	XLA HLO + Python primitives	算子直接是 HLO，很多算子是 Python 组合实现
MindSpore	Python OpDef + 自动生成	与 PyTorch 类似但更晚成熟
OneFlow	自有 IR + 手写 op_def	编译期优化更激进，但生态小

面对"数千算子的工程一致性"这个问题，声明式 schema + 代码生成是一条行之有效的路，国内框架里也能看到类似的做法。

值得多说一句的是 JAX 的对比。JAX 选择了不同的路线：算子是 Python 一等公民，所谓 "primitive" 加上一组规则（abstract eval、impl、jvp、transpose）。这种纯 Python 设计让自定义算子门槛极低（写几个 Python 函数就行），但代价是性能完全依赖 XLA 编译器 —— 没编译就没法跑。PyTorch 的 codegen 路线让 eager 模式能立即跑，编译只是 可选优化。两条路线没有绝对优劣，体现了 "动态优先" 与 "静态编译优先" 两种 ML 框架哲学的差异。

6.11 历史回顾：从 C++ 模板到 YAML codegen

PyTorch 的算子代码生成并不是后来才有的：早在 2017 年，ATen 就用 YAML 类文件声明算子 —— 一边是承接 TH 系 C 库的 aten/src/ATen/Declarations.cwrap，一边是 2017 年 11 月已在仓库里的 aten/src/ATen/native/native_functions.yaml（"native functions"），由 Python 脚本生成 C++。

真正的转折是 2020 年 8 月 Edward Yang 的 PR #42629 "Rewrite of ATen code generator"：它把生成器重写成以类型化 dataclass（即今天 torchgen/model.py 里的 NativeFunction / FunctionSchema）为中心的 tools/codegen 包，并删掉了 Declarations.cwrap；2022 年 4 月的 PR #76275 再把 tools/codegen 改名为顶层的 torchgen 包（首个包含它的正式版是 v1.12）。structured kernels（§6.4.4）也是在这套新生成器上落地的。这次重写是 PyTorch 工程化历史上的重要转折之一。

理解这条历史，你在旧资料里看到 tools/codegen/、Declarations.cwrap 这类路径时，就知道它们对应今天的 torchgen/ 与 native_functions.yaml。

6.12 跨专栏关联

• 《Serde 元编程》专栏第 12-13 章：Serde 用 Rust 派生宏（procedural macro）做"序列化代码生成"，与 ATen 用 YAML + Python 做"算子代码生成"是同一思想的不同实现。Rust 的派生宏在编译期一次过；PyTorch 的 torchgen 是 Python 脚本独立先跑、产物再编译。各有取舍：派生宏类型安全、torchgen 调试友好
• 《Rust 编译器与运行时揭秘》专栏第 14 章（宏系统）：Rust 编译器内部对一些 trait（如 Clone、Default）的派生也是代码生成，思想与 ATen 一致
• 《MCP 协议设计与实现》专栏第 5 章（Tool）：MCP 也用 schema-driven 注册（JSON schema → method handler），但量级远小于 PyTorch（每个 server 几十个 method vs 近 2700 条算子 schema）
• 《vLLM 推理内核深度解析》专栏第 17 章（API 服务器）：vLLM OpenAI 兼容 API 的 schema 也是 schema-driven 自动生成绑定 —— 与 ATen 思想一致，规模一个量级以下

6.13 几条踩坑提示

实战添加新算子时常见的坑：

1. schema 写错了类型：Tensor 与 Tensor(a!) 一字之差，functionalization 规则完全不同。如果你想写"修改输入"的算子，必须用 (a!) 标记 2. 忘了 derivatives.yaml：新算子默认不可反向。loss.backward() 时报 "derivative for aten::xxx is not implemented" 是这个错的标志 3. dispatch 表漏了 backend：只写了 CPU 没写 CUDA，CUDA 张量上调用会报 "Could not run 'aten::xxx' with arguments from the 'CUDA' backend"。最少要为 CPU + CUDA 两个 backend 写 impl，或注册一个 fallback 4. 改了 YAML 但没重编：torchgen 在 build 时跑，如果你只改 .yaml 不重新构建，生成的代码不会变。重新执行 python -m pip install --no-build-isolation -v -e . 才会重跑 5. structured kernel 的 meta 与 impl 签名不一致：structured: True 算子的 TORCH_META_FUNC 和 TORCH_IMPL_FUNC 签名必须严格匹配 torchgen 按 schema 生成的声明，否则编译报错

6. structured impl 误把计算放到 meta：meta 函数只能算 shape / dtype / 分配 out，绝不能做实际数值计算。FakeTensor / torch.compile 在 trace 阶段只跑 meta、不跑 impl —— 如果 meta 里有数值计算，trace 出来的图会错

7. derivatives.yaml 引用了作用域外的张量：反向公式能引用的是前向的输入参数和输出（result，多输出时 result0、result1…），公式里用到哪些，代码生成器就自动把它们存进 grad_fn（如 _softmax 的反向公式直接用 result）。如果反向还需要别的中间值，常见做法是让前向算子把它作为额外输出返回 —— 例如 native_layer_norm 返回 (Tensor, Tensor, Tensor)，反向公式用 result1、result2 取 mean 与 rstd

这些坑在 PyTorch 官方的 aten/src/ATen/native/README.md 与 tools/autograd/README.md 里有详尽指南，新增算子前建议通读（注意老资料里的 tools/codegen/ 目录早已改名为顶层的 torchgen/，那个路径现在是不存在的）。

6.14 调试技巧：怎么找生成的代码

新手最常迷惑的一件事：源码里 grep 不到 at::add 的实现。原因是它最终调用的 at::_ops::add_Tensor::call 定义在 build/aten/src/ATen/Operators_*.cpp 里 —— 生成的代码不在 git 仓库里。

调试时的标准操作：

# 1. 定位生成代码所在文件
find build/ -name "Operators_*.cpp" -exec grep -l "add_Tensor::call" {} \;

# 2. 在生成的文件里看 add.Tensor 的调用入口
grep -n -A5 "add_Tensor::call" build/aten/src/ATen/Operators_2.cpp

# 3. 找你的算子在哪个 RegisterCPU_N.cpp 里注册
grep -n 'm.impl("add.Tensor"' build/aten/src/ATen/RegisterCPU_*.cpp

如果你想跟踪"我的 YAML 改动产出了什么 C++ 代码"，最快的办法是改一行 YAML、跑构建、对比 build/ 目录的 diff。这是给 PyTorch 提交 codegen 相关 PR 的标配技能。

6.15 几条通用启示

读完本章，把代码生成思想抽象成可迁移的几条规则：

第一：在工程化框架里，schema 是真理。一旦你设计的系统涉及多语言绑定、序列化、自动化文档，把 schema 用一个权威源（YAML / JSON / proto）声明，所有派生路径都从它生成。这避免"手维护 N 个一致性"的灾难。

第二：生成代码而不是写库。当你发现自己在写"看似相同但参数不同"的 N 段代码，停下来 —— 把它声明化、用脚本生成。一个 1000 行的 codegen 比 10000 行的"近重复"库更易维护。

第三：保留逃生口（escape hatch）。PyTorch 允许某些算子用手写而不是 structured，是因为现实总有"你的框架没考虑到"的边角情况。codegen 框架要能优雅退路。

第四：模板小、生成器逻辑大。torchgen 的模板字符串很短，复杂度都在 Python 端的"分析 + 决策"逻辑里。这是因为模板系统本身是脆弱的（缩进、换行、字符串拼接），而 Python 逻辑可调试可重构。

第五：生成代码不入 git，但保留可观察性。PyTorch 把 build/aten/src/ATen/Operators_*.cpp 这类生成产物放到 .gitignore，避免每次改 YAML 都引发巨大 diff。但生成产物在编译产物里完整保留，开发者能 grep / 调试。这种"瞬态产物可观察"是 codegen 工程化的关键 trade-off。

第六：让 schema 覆盖尽量多的衍生物。Python 端 torch.add 的类型签名（torch/_C/_VariableFunctions.pyi，由 tools/pyi/gen_pyi.py 生成）就从 schema 推导；docstring 则仍在 torch/_torch_docs.py 里手写。schema 覆盖到的部分自动一致，手写的部分就得靠人维护。

把这些启示用到你自己的项目，能省下大量看似"必要"的重复劳动。

6.16 一个建议练习：从 schema 反推生成代码

要真正吃透本章，可以试一道反向题：从一段 PyTorch 文档里的算子接口反推它的 native_functions.yaml schema。

例如 torch.nn.functional.linear(input, weight, bias=None)：

• 三个 Tensor 参数，第三个可选（→ Tensor?）
• 返回单个 Tensor
• 推测的 schema：

- func: linear(Tensor input, Tensor weight, Tensor? bias=None) -> Tensor
  python_module: nn
  variants: function
  dispatch:
    CompositeImplicitAutograd: linear

实际打开 aten/src/ATen/native/native_functions.yaml 搜 linear，对照差异。这个练习能把"YAML schema 语法"内化成肌肉记忆。

完成几次这种反推，你就能给 PyTorch 提交自定义算子 PR 了 —— 源码里有 16172 行真实例子可以学。

6.17 一个反思：为什么 PyTorch 不用 protobuf

读到这里有人会问 —— 既然要用 schema 驱动，为什么不用更标准的 protobuf / FlatBuffers / Cap'n Proto？这些工具有更成熟的代码生成器。

我们没有找到 PyTorch 官方对此的正式说明，但从工程上可以看出几点取舍：

• YAML 比 protobuf 可读得多：算子作者每天打开 schema 文件，YAML 的零样板比 proto 的 message + field number 更友好
• PyTorch 的 schema 不需要跨语言序列化：proto 的核心价值是网络传输，PyTorch schema 只在编译期被消费
• 自家 codegen 比 protoc 灵活：torchgen 能生成 Python 端的 docstring、autograd 反向、functionalize 规则等"非协议"产出，proto 做不到

这是一条"专用工具胜过通用工具"的工程决策。如果你的项目 schema 主要给本进程用、不需要网络传输，自家轻量级 YAML + Python codegen 通常比 protobuf 更合适。

下一章拆 Autograd —— 这一章我们看到的 derivatives.yaml 怎么变成真正的反向计算引擎。