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

“Code generation is what lets PyTorch ship 3000 operators with consistent quality without 3000 humans writing repetitive boilerplate.”

—— Edward Yang, “PyTorch dev podcast: code generation”

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

第 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 规则
• 生成 PyTorch RPC 序列化
• ……

每个算子都要做以上 10+ 件事。PyTorch 注册了 3000+ 个算子。如果靠手写，工作量是 3 万次重复劳动 —— 每次新增一个 dispatch key（如某国产芯片）都要改 3000 个文件。

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

• 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/>(数十个文件，几十万行)"]
    TG --> O2[RegisterCPU.cpp / RegisterCUDA.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[FunctionalInverses.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

——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：分类标签，用于子集筛选（如 mobile build 只保留 core 标签）

6.2.1 schema 类型语法

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

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

6.2.2 dispatch 表的写法

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

dispatch:
  CPU, CUDA, MPS: my_kernel       # 三个 key 共用 my_kernel
  Sparse*: my_sparse_kernel       # 通配符

如果不写 dispatch 字段，算子默认走 CompositeImplicitAutograd —— 也就是用其他算子组合实现的算子，autograd 自动正确（第 5 章 §5.2.2 提过）。这种”组合算子”通常代码非常短，例如 torch.where 内部就是 cond * a + (1 - cond) * b。

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 但不能 inplace 改），(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）

注意几个特殊符号：

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

注意 result 字段对 forward grad 的特殊语义：它返回的是输出的切线 (output tangent)，输入是 *_p（基础值）和 *_t（输入切线）。如果你看不懂这套数学，没关系 —— 99% 的 PyTorch 算子作者只写 backward 字段（self、mat2 等），forward grad 由 PyTorch 团队的几个核心维护者集中处理。

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

每个 derivatives.yaml 条目都必须对应一个 native_functions.yaml 里存在的算子。如果你写了一个新算子但没写反向规则，默认这个算子不能反向（dispatcher 命中 Autograd key 时报错”…does not have a derivative”）。

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

如果某算子明确”不需要反向”（如 RNG、随机采样），在 native_functions.yaml 里加 tags: nondifferentiable 标记，避免代码生成器抱怨缺反向。

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），就在 tools/autograd/FunctionsManual.cpp 里写 helper 函数，YAML 里只写一行调用。这种”声明式入口 + 实现式后端”分离让反向规则一眼读起来就是”数学公式”。

6.3.3 forward grad 的特殊性

注意 derivatives.yaml 里的 result 字段 —— 它不是反向 (backward) 而是 前向 AD (forward grad)。这是 PyTorch 较新的能力（v1.10+），允许用户做 jvp（Jacobian-vector product）而不是 vjp（vector-Jacobian product）。

forward grad 在大多数算子上是反向规则的”对偶”：如果反向是 mat1.grad += grad @ mat2.T，前向就是 result.grad = self.grad @ mat2。derivatives.yaml 在同一个条目里把两套规则一起写，让 torchgen 一次生成两套 C++ 代码（MmBackward0、MmJvp 等）。

对绝大多数 PyTorch 用户来说 forward grad 不常用（vjp 已经够用），但对二阶导数、Hessian 计算等场景必不可少。第 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 框架把它自动生成：

graph TB
    YAML["native_functions.yaml<br/>mm.out: structured True"]
    YAML --> Gen["torchgen 解析"]
    Gen --> Meta["生成 meta 函数<br/>shape 推导 + dtype 检查 + 分配 out"]
    Gen --> Impl["生成 impl wrapper<br/>调用用户写的 mm_kernel"]

    User["用户只写: mm_kernel<br/>纯计算, 无 boilerplate"] -.被 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 框架的核心：把 boilerplate（shape 检查 / dtype 推导 / 分配输出）从用户代码中剥离，让用户只写”真正的计算”——这与 §6.5 的 ufunc 思路一致，都是”代码生成抽象掉重复”。

// 代码生成器吐出的简化版 (基于 native_functions.yaml 的 structured: True)
struct structured_mm_out final : public TensorIteratorBase {
    void meta(const Tensor& self, const Tensor& mat2);
    void impl(const Tensor& self, const Tensor& mat2, const Tensor& out);
};

// meta 函数：算 shape、dtype、分配 out（torchgen 生成）
void structured_mm_out::meta(const Tensor& self, const Tensor& mat2) {
    set_output_raw_strided(0, {self.size(0), mat2.size(1)},
                            {}, self.options());
}

// impl 函数：用户只需写这一段（CPU 在 mm_out_cpu, CUDA 在 mm_out_cuda）
TORCH_IMPL_FUNC(mm_out_cpu)(const Tensor& self, const Tensor& mat2, const Tensor& out) {
    at::native::mm_kernel(self, mat2, out);
}

——开发者只写 impl，meta、shape 检查、dtype 提升、out 分配全部代码生成。这在大量算子上节省了海量重复代码。

6.4.1 structured_delegate 的转发机制

回到 mm（无 out）：

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

这告诉 torchgen：mm 自己不实现，分配一个 out 张量后调 mm.out。生成的代码大致是：

// 生成的 mm 实现
Tensor mm(const Tensor& self, const Tensor& mat2) {
    Tensor result;
    at::mm_outf(self, mat2, result);  // 实际调用 mm.out
    return result;
}

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

6.4.2 structured 框架与 inplace 三态

每个 ATen 算子在 structured 框架下有三个变体：

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

这三个变体共享同一个 meta（计算 shape）和 impl（真正的 kernel）。代码生成器自动把 mm 转发到 mm.out、把 mm_ 转发到 mm.out。用户只需写一份 impl，三个 API 自动具备。

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

6.4.2.5 meta 函数与 FakeTensor 的关系

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

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

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

6.4.3 结构化算子的演进

structured 框架不是 PyTorch 一开始就有的。早期所有算子都是手写完整流水线，重复严重。2020-2021 年 Edward Yang 主导了 “structured kernels” 提案（PR #51277 系列），把所有 ATen 算子向 structured 迁移。今天 v2.11 大约 70%+ 的算子已经迁移完成，剩下的一些”特殊算子”（如 RNG、量化、复杂的 reduce）暂未迁移。

理解这个迁移历史，你看源码时能区分”老式手写算子”和”现代 structured 算子”两种风格。前者多在 aten/src/ATen/native/ 的老文件里，后者在新文件或 *_native.cpp 后缀里。

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

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

// aten/src/ATen/native/ufunc/add.h:14
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。代码生成器会自动展开为：

// 生成的 (简化) CPU SIMD kernel
void add_kernel_cpu(TensorIteratorBase& iter, const Scalar& alpha) {
    AT_DISPATCH_ALL_TYPES_AND2(kBFloat16, kHalf, iter.common_dtype(), "add", [&]() {
        auto alpha_val = alpha.to<scalar_t>();
        cpu_kernel_vec(iter,
            [=](scalar_t a, scalar_t b) -> scalar_t {
                return ufunc::add(a, b, alpha_val);
            },
            [=](Vectorized<scalar_t> a, Vectorized<scalar_t> b) {
                return ufunc::add(a, b, Vectorized<scalar_t>(alpha_val));
            });
    });
}

// 生成的 CUDA kernel
void add_kernel_cuda(TensorIteratorBase& iter, const Scalar& alpha) {
    AT_DISPATCH_ALL_TYPES_AND2(kBFloat16, kHalf, iter.common_dtype(), "add", [&]() {
        gpu_kernel_with_scalars(iter,
            [=]GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t {
                return ufunc::add(a, b, alpha.to<scalar_t>());
            });
    });
}

一个 10 行的内层公式 → 几十个 dtype × 后端的 kernel。这是 PyTorch 处理”逐元素操作”的统一机制 —— 数学公式与 SIMD 化、CUDA 化、dtype 展开完全解耦。

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.1.5 ufunc 与 vec::Vectorized

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

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 在循环内部对对齐+整 chunk 部分用向量版，对剩余的尾部元素用标量版。这种”vector + scalar tail”是 SIMD 编程的标准模式。

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

6.5.2 哪些算子能 ufunc 化

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

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

加减乘除、比较运算、激活函数（ReLU / Sigmoid / Tanh）、绝大多数 element-wise 数学函数（exp / log / sin / cos）都符合。但 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", [&]() {
    using scalar_t = ...;        // 在这个 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 选择”零开销 —— 编译期把所有可能性展开成静态分支。

6.6.1 dispatch 宏的二进制膨胀代价

AT_DISPATCH_* 的代价是 二进制膨胀：每个用了这套宏的算子实现，编译产出 N 份代码（每个 dtype 一份）。如果一个算子用 AT_DISPATCH_ALL_TYPES_AND4(...)，可能展开 12-15 份特化。乘上几百个算子，整个 libtorch 的体积有相当一部分来自这种”特化爆炸”。

PyTorch 在移动端有 C10_MOBILE 编译宏，会让 dispatch 宏只展开少数几个 dtype（如 fp32），减小体积。这是为什么 PyTorch Lite（移动端）的 .so 比桌面版小几倍 —— 不是少了功能，而是少了 dtype 特化。

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

PyTorch 2.0+ 引入 SymInt —— 一个能表示”具体整数”或”符号变量”的类型。它由 torch.compile / Dynamo 在 trace 阶段使用，让算子能接受 dynamic shapes 而不立即具体化。

代码生成器对每个算子自动生成 SymInt 版本。在 native_functions.yaml 里你会看到很多算子有 SymInt[] 参数（如 view(Tensor self, SymInt[] size)），torchgen 据此生成两套 C++ 实现：一套接收普通 int64_t、一套接收 c10::SymInt。运行期普通调用走前者；trace 阶段走后者，让符号 shape 一路传到 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 对象]
    Y2 --> Parse
    Y3 --> Parse

    Parse --> Model[torchgen/model.py<br/>Python 中间表示<br/>NativeFunction / FunctionSchema]

    Model --> CG1[gen_cpp.py<br/>生成 Tensor 方法 / at::xxx]
    Model --> CG2[gen_dispatch.py<br/>生成 Operators_*.cpp]
    Model --> CG3[gen_register_*.py<br/>生成 RegisterCPU/CUDA/...]
    Model --> CG4[gen_autograd_*.py<br/>生成 VariableType / Functions]
    Model --> CG5[gen_python_*.py<br/>生成 Python 包装]
    Model --> CG6[gen_functionalization_type.py<br/>functionalize 规则]
    Model --> CG7[gen_vmap_plumbing.py<br/>vmap 规则]

    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.0 几个 gen_*.py 的分工

torchgen 里的几个核心生成模块 ：

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

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

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

1. 解析：parse_native_yaml 读 native_functions.yaml，把 mm 与 mm.out 解析成两个 NativeFunction 对象
2. dispatch 表合并：把 dispatch 字段展开到 dispatchTable，包括 alias key（CompositeExplicitAutograd → 所有 backend）
3. Operators_*.cpp：生成 at::_ops::mm::call 入口，里面调 Dispatcher::singleton().findSchemaOrThrow(...).typed<...>().call(...)（第 1 章 §1.2 见过）
4. RegisterCPU.cpp：生成 TORCH_LIBRARY_IMPL(aten, CPU, m) { m.impl("mm.out", &at::native::structured_mm_out_cpu_functional::wrapper); }
5. TensorBody.h：生成 Tensor::mm(const Tensor& mat2) const { return at::_ops::mm::call(*this, mat2); }
6. VariableType_4.cpp：生成 VariableType::mm 的反向包装，使用 derivatives.yaml 翻译出的 MmBackward0 节点
7. python_torch_functions_2.cpp：生成 THPVariable_mm Python 函数，能从 Python 端调用

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

6.7.2 torchgen 模板系统

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

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

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

风格与 Python 的 string.Template 类似，但支持嵌套替换和列表展开。PyTorch 团队选择不依赖 Jinja2 是为了 减少 build-time 依赖——torchgen 是 PyTorch 编译过程的”临界路径”，每多一个 Python 包就多一个潜在的安装问题。

6.7.2.5 一个有趣的设计：translate

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

举例：dispatcher 调用约定的参数列表（(const Tensor&, const Tensor&, const Scalar&)）和 native 调用约定（(const Tensor& self, const Tensor& mat2, const Scalar& alpha)）类型一样但语义不同（前者无名、后者有名）。生成 RegisterCPU.cpp 里的 wrapper 函数时，需要把 dispatcher 收到的位置参数转成 native 函数的命名参数 —— translate.py 就是这道题的引擎。

它内部维护一组”binding 转换规则”（如 Tensor → const Tensor&、int[] → IntArrayRef），自动推断怎么把一组类型签名”翻”成另一组。这套机制让 PyTorch 能在 dispatcher / native / Python / functionalize 等多种约定之间自如转换，开发者不用手写每一对的转换代码。

6.7.3 torchgen 的 Python 中间表示

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

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

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 impl"]
    end

    subgraph Codegen["编译期 torchgen"]
        TG[torchgen/gen.py 跑]
    end

    subgraph Generated["生成的 C++ (build/)"]
        G1[Operators_4.cpp<br/>at::_ops::add::call]
        G2[RegisterCPU.cpp<br/>m.impl 注册]
        G3[RegisterCUDA.cpp]
        G4[VariableType_4.cpp<br/>autograd 包装]
        G5[python_torch_functions_2.cpp<br/>Python 绑定]
        G6[UfuncCPUKernel_add.cpp<br/>SIMD 展开]
        G7[UfuncCUDAKernel_add.cu<br/>CUDA 展开]
    end

    subgraph Binary["编译后 (libtorch.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
    Y3 --> G6
    Y3 --> 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 → 自动膨胀到二进制：3 行 YAML → 数千行 C++ → 数万行汇编
• 手写部分极少：开发者只需写”算子的 impl + 反向公式”两段，其余全是模板/生成
• 跨语言一致性自动保证：Python torch.add、C++ at::add、Tensor::add 永远拥有完全一致的签名 —— 因为它们都从同一个 schema 派生
• 二进制可观察：调试时可以 nm libtorch.so | grep add 看到生成代码产生的所有符号

6.8.1 编译时间的代价

代码生成虽然减少手写，但有一个常被忽略的代价：编译时间。PyTorch 全量编译要 1-2 小时，其中很大一部分是编译这些”机器生成的几十万行 C++“。具体痛点：

• Operators_*.cpp 一个文件展开后可能 5000+ 行，每个文件编译 5-10 秒
• 模板展开导致的 SFINAE 计算让编译器单文件吃 1-2 GB 内存
• 一次性 link 所有生成的 .o 时，linker 是另一个瓶颈

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

1. 拆分生成文件：Operators_0.cpp 到 Operators_4.cpp 把 3000 算子分成 5 块，让 make 能并行编
2. AT_PER_OPERATOR_HEADERS 宏：每个算子单独一个 .h 文件，用户代码可以”只 include 自己用到的”，减少 amplification
3. 预编译头：常用 c10 / aten 头文件做成 PCH

这些优化让全量编译从 v1.0 时代的 4-5 小时压到现在的 1-2 小时。但对国内大多数 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 setup.py develop。torchgen 自动检测 YAML 变化，重新生成所有派生文件

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

注意 6 步里只有 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（v2.4+ 稳定）让用户不需要 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，生成”我应该实现哪些算子”的 stub 列表
3. 逐个实现：把 stub 列表里的算子一个个用自家硬件 SDK 实现
4. 作为独立 wheel 发布：用户 pip install torch_npu / pip install torch_mlu 等就能在自家芯片上跑 PyTorch

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

理解这条机制，你就能解释为什么国产芯片接入 PyTorch 速度比想象的快 —— 代码生成器消除了 90% 的胶水代码工作。

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

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

PyTorch 的 YAML 路线被 MindSpore、PaddlePaddle 等国内框架借鉴，因为它在”数千算子的工程一致性”问题上确实领先。

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

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

PyTorch 早期（pre-1.0 / Torch7 时代）算子注册全部是手写 C++ 模板。每加一个算子要改 5-10 个文件，每加一个 dtype 要改几十处宏。这种状态在 2017-2018 年算子数量爆炸增长时变得不可维持。

2018-2019 Edward Yang、Will Feng、Sebastian Messmer 等人主导了从手写模板向 YAML codegen 的渐进迁移。早期的 codegen 用的是 Python 2 风格的脚本，慢慢演进成今天的 torchgen 包（Python 3.10+ 类型化 dataclass + 模块化）。这次迁移是 PyTorch 工程化历史上最重要的转折之一 —— 没有它，PyTorch 不可能维持每月一个 minor 版本、每个版本几百个 PR 的开发节奏。

理解这条历史，你看 PyTorch 旧代码（如 torch/csrc/autograd/VariableTypeUtils.cpp、各种 *_legacy.cpp）就能识别”这是迁移前的手写代码”。它们大部分还在，但都标了 legacy、deprecated，逐步在被新生成代码替换。

6.11 跨书关联

• 《Serde 元编程》全书：Serde 用 Rust 派生宏（procedural macro）做”序列化代码生成”，与 ATen 用 YAML + Python 做”算子代码生成”是同一思想的不同实现。Rust 的派生宏在编译期一次过；PyTorch 的 torchgen 是 Python 脚本独立先跑、产物再编译。各有取舍：派生宏类型安全、torchgen 调试友好
• 《Rust 编译器之路》第 X 章 内置 trait 与代码生成：Rust 编译器内部对一些 trait（如 Clone、Default）的派生也是代码生成，思想与 ATen 一致
• 《MCP 协议剖析》第 X 章 Server 注册：MCP 也用 schema-driven 注册（JSON schema → method handler），但量级远小于 PyTorch（每个 server 几十个 method vs 3000+ 算子）
• 《vLLM 内核探秘》第 17 章 API 服务器：vLLM OpenAI 兼容 API 的 schema 也是 schema-driven 自动生成绑定 —— 与 ATen 思想一致，规模一个量级以下

6.12 几条踩坑提示

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

1. schema 写错了类型：Tensor 与 Tensor(a!) 一字之差，functionalization 规则完全不同。如果你想写”修改输入”的算子，必须用 (a!) 标记 2. 忘了 derivatives.yaml：新算子默认不可反向。loss.backward() 时报 “does not have a derivative implemented” 是这个错的标志 3. dispatch 表漏了 backend：只写了 CPU 没写 CUDA，CUDA 张量上调用会报 “no kernel found”。最少要为 CPU + CUDA 两个 backend 写 impl，或注册一个 fallback 4. 改了 YAML 但没重编：torchgen 在 build 时跑，如果你只改 .yaml 不重新构建，生成的代码不会变。python setup.py develop 强制重跑 5. struct kernel 的 meta 与 impl 签名不一致：structured_delegate 模式下 meta 函数和 impl 函数的签名必须严格匹配 schema，否则编译报错

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

7. derivatives.yaml 引用了未保存的张量：反向规则只能引用前向保留的张量。如果你的反向需要前向中间值（如 softmax 的输出），要在 native_functions.yaml 里通过 output_differentiability 或者特殊保存机制让前向把它存下来 —— 否则反向会拿到无效数据

这些坑在 PyTorch 官方的 aten/src/ATen/native/README.md 与 tools/codegen/README.md 里有详尽指南，新增算子前建议通读。

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

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

调试时的标准操作：

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

# 2. 在生成的文件里找具体注册语句
grep -A5 "at::_ops::add" build/aten/src/ATen/Operators_4.cpp

# 3. 找你的算子最终落到哪个 RegisterXxx.cpp
grep -rn "TORCH_LIBRARY_IMPL.*aten.*CPU" build/aten/src/ATen/RegisterCPU.cpp | head

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

6.13 几条通用启示

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

第一：在工程化框架里，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 的 docstring 在某种程度上也由 schema 推导。当 schema 是单一真理源，连文档都能自动一致。这是文档准确性的强保证。

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

6.14 一个建议练习：从 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.15 一个反思：为什么 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 怎么变成真正的反向计算引擎。

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