05 · 算子怎么注册进 PyTorch

⏱ 80 分钟 ｜ ⭐ 一条报错里同时给出了算子名和分派键 —— 那正好是 dispatcher 那张表的两个坐标轴

🎯 一句话

PyTorch 的算子分派就是一张二维表：行是算子名（aten::gelu），列是分派键（CPU / SparseCPU / Autograd），格子里放一个函数指针。 所谓「注册一个算子」就是往表里填格子；TORCH_LIBRARY 和 TORCH_LIBRARY_IMPL 这两个宏，一个负责新开一行，一个负责填一个格子。

🧩 一、⭐ 先看一条查表失败的报错

本章 Python 侧的输出全是实跑的，环境 torch 2.13.0+cpu + CPython 3.13，CUDA 不可用。

# c05_miss.py
import torch

x = torch.ones(2, 2).to_sparse()      # ⭐ 转成稀疏张量，仅此而已
torch.nn.functional.gelu(x)

实跑报错（截前两句）：

NotImplementedError: Could not run 'aten::gelu' with arguments from the 'SparseCPU' backend.
This could be because the operator doesn't exist for this backend, or was omitted during ...
'aten::gelu' is only available for these backends: [CPU, CUDA, HIP, MPS, ... ]

⭐ 这一句里有三样东西，正好是本章的三条主线：

⚠️ 两个容易记错的细节：

• 异常类型是 NotImplementedError 而不是 RuntimeError。⭐ 但 Python 里前者是后者的子类（实跑 isinstance(e, RuntimeError) 为 True），所以 except RuntimeError 照样接得住 —— 这也是为什么到处都有人写成 RuntimeError。
• ⭐ 不需要 GPU。 很多人以为要看见 aten:: 名字得先把 tensor 扔到 CUDA 上，其实普通 CPU 机器上 to_sparse() 一下就能复现。

🚦 ⭐ 它不是「稀疏张量不支持激活函数」

同一个稀疏张量，换个算子就能跑：

# c05_contrast.py
import torch

s = torch.ones(2, 2).to_sparse()
print(s.mm(torch.ones(2, 2)))    # ⭐ 能跑，输出 [[2,2],[2,2]]
try:
    s.sigmoid()                  # ⚠️ 不能跑
except NotImplementedError as e:
    print(str(e).split("This could be")[0])

⭐ 同一列（SparseCPU）上，mm 那格填了，sigmoid 那格是空的。 这张表是逐格填的，不存在「某个后端整体支持 / 不支持」这回事。

不用猜，直接查格子在不在：

# c05_probe.py
import torch

for op in ["aten::gelu", "aten::sigmoid", "aten::mm"]:
    print(op, [(k, torch._C._dispatch_has_kernel_for_dispatch_key(op, k))
               for k in ("CPU", "SparseCPU", "QuantizedCPU", "CUDA", "Meta")])

实跑结果（✅ = 这一格有 kernel）：

⚠️⚠️ CUDA 那一列的 ❌ 是本机的假象，不是 PyTorch 的事实。 本机装的是 CPU 版轮子，torch_cuda 那批库根本不在包里，那些格子从没被填过。⭐ 这恰好演示了更重要的一件事：这张表不是编译期写死的常量表，是运行时一批一批被「加载进来」的 —— 第五节会看到填格子的动作发生在哪一刻。

⭐ 顺带记住 Meta 这一列：「只算形状不算数值」的假后端，形状推导和 torch.compile 都靠它，三个算子全有。

🧩 二、行坐标：算子名和它的 schema

aten::gelu 不是随手起的字符串，它背后挂着一条签名（schema）。Python 侧直接能问：

# c05_schema.py
import torch

print(torch.ops.aten.add.Tensor._schema)
print(torch.ops.aten.gelu.default._schema)
print(len(torch.ops.aten.add.overloads()), "个 add 重载")
print(len(torch._C._jit_get_all_schemas()), "条 schema")

实跑输出：

⚠️ 4362 随版本变（torch 2.13.0+cpu 上是这个数，你的一定不同），要看的是量级：四千多条。

把第一行拆开：

⭐ 重载名是字符串，这点反直觉但重要：aten::add 实跑有 16 个重载（Tensor / Scalar / out / default …），而每个重载在 dispatcher 眼里都是独立的一行、独立的一张表。所以 aten::add 不是一行，是十六行。

schema 里几个必须认得的记号：

⚠️ 别名标记不是文档，是给编译器看的合同。 torch.compile、functionalization、torch.jit 全靠它判断「这个算子会不会偷改我的输入」。⭐ 自己写算子时填错它的后果是别的通道上出现随机错误，而不是当场报错。

⭐ schema 是唯一的真相源：C++ 那边的函数签名、Python 这边的 torch.ops.aten.gelu、类型存根 .pyi、还有 06 章会翻的那堆 ATen/ops/*.h，全是从同一份 schema 生成的。这就是为什么改一个算子的签名在 PyTorch 里是件大事。

🧩 三、列坐标：DispatchKeySet

一个 tensor 身上带的是一组分派键，不是一个：

# c05_keyset.py
import torch

print("普通 CPU     :", torch._C._dispatch_key_set(torch.ones(2, 2)))
print("稀疏         :", torch._C._dispatch_key_set(torch.ones(2, 2).to_sparse()))
print("requires_grad:", torch._C._dispatch_key_set(torch.ones(2, 2, requires_grad=True)))

实跑：

⭐ ① 稀疏那行只改了第一个键（CPU → SparseCPU），后面三个一模一样。第一节那条报错的全部病因就是：换了列，而那一列的格子是空的。

⚠️⚠️ ② requires_grad=True 和普通 CPU 的 keyset 完全一样。 很多人以为「要不要求导」是靠分派键区分的 —— 实测不是。AutogradCPU 一直都在，走不走那条路取决于运行时的 GradMode 和 tensor 自己的 requires_grad 标志。

🚦 dispatcher 怎么用这一组键

⭐ 规则只有一句：所有键有一个全局优先级，从高到低找第一个「这一行填了格子」的键。 大致层次：

⭐ 上面几层「干完自己的事再往下发一次」的动作叫 redispatch。 ⚠️ 由此得到一条很实用的推论：「注册一个 CPU kernel」和「让它可导」是两件独立的事 —— 你填的是 CPU 那格，Autograd 那格还空着（第六节）。

🧩 四、⭐ 直接把那张表 dump 出来

前面都在推理，现在直接看。PyTorch 有个内部接口能打印某一行的全部格子：

# c05_dump.py
import torch

print(torch._C._dispatch_dump("aten::gelu"))

实跑输出（⭐ 每一行都在，只把长路径前缀省成 …，本机 site-packages 路径写成 <site-packages>）：

🚦 这段输出的五个读法

① ⭐ 从头到尾没有 SparseCPU 那一行。 第一节的报错到此闭环 —— 不是「不支持稀疏」这种模糊说法，是表上那个格子确实空着，你亲眼看见了。

② 每行都写着「注册在哪个文件的哪一行」。 ⭐ 这是 06 章那条定位路线的起点。⚠️ 原始输出里那些 C:\actions-runner\_work\pytorch\... 是 PyTorch 官方 CI 构建机上的路径，你本地绝对找不到这些文件 —— 但文件名和行号是真的，去 GitHub 上对应 tag 里搜就能找到。

③ 路径里带 build\ 的都是【生成】出来的。 RegisterCPU_3.cpp、RegisterQuantizedCPU_0.cpp、VariableType_1.cpp 在 PyTorch 仓库里根本不存在，是构建时由代码生成器按 schema 吐出来的。⭐ 这也解释了它们为什么动辄几千行、名字还带编号 —— 那是为了切开编译单元。

④ [alias] 是别名键，一次填一批格子。 Autograd[alias] 不是某个具体后端的 autograd，而是「所有后端都用这个」。⭐ PyTorch 靠它避免为 30 多个后端各写一遍同样的代码。

⑤ ⭐ Meta 那行注册在 _meta_registrations.py —— 一个 Python 文件。 说明这张表不是 C++ 独占的，Python 也能填格子。⚠️ 注意它下面还跟着 Meta (inactive)：C++ 那边本来也注册了一个，被 Python 这个覆盖掉了，dispatcher 老老实实把被盖住的那个标成 inactive 留在原地。

🛑 读到这里可以停 —— 前半章讲完了（约 35 分钟）。 后半章还有：⭐ C++ 侧：两个宏，各管一件事 · ⭐ 让它可导：C++ 侧的那一半 · CUDA 那一格，和不想编译的路子 回来的时候不用重读，直接从下一节接着看就行。

🧩 五、⭐ C++ 侧：两个宏，各管一件事

⚠️ 本节及以后的 C++ 代码全部标 🗓️ 未实跑 —— 本机无编译环境（MinGW g++ 配 MSVC 编译的 CPython + CPU 版 torch，编不出能导入的扩展）。写法取自本机 torch/include/torch/library.h 的官方注释和 PyTorch 文档，作者没有在本机编译运行过。

// myops.cpp   🗓️ 未实跑 —— 本机无编译环境
#include <torch/library.h>
#include <ATen/ATen.h>

at::Tensor scaled_add_cpu(const at::Tensor& a, const at::Tensor& b, double s) {
    return a + b * s;
}

TORCH_LIBRARY(myops, m) {                                   // ⭐ ① 新开一行
    m.def("scaled_add(Tensor a, Tensor b, float s) -> Tensor");
}

TORCH_LIBRARY_IMPL(myops, CPU, m) {                         // ⭐ ② 填 CPU 那一格
    m.impl("scaled_add", scaled_add_cpu);
}

⭐ 四个入口，对应表上四种动作：

⚠️ TORCH_LIBRARY 对同一个命名空间只能出现一次 —— library.h 的注释里明写着「There may only be one TORCH_LIBRARY() for any given namespace」。分好几个文件写就用 TORCH_LIBRARY_FRAGMENT。

🚦 ⭐ 宏展开成什么 —— 一个静态初始化对象

这是本章最该记住的机制。下面是 torch/include/torch/library.h 里的真实源码（本机头文件第 975 行起）：

#define TORCH_LIBRARY(ns, m)                                                   \
  static void TORCH_LIBRARY_init_##ns(torch::Library&);                        \
  static const torch::detail::TorchLibraryInit TORCH_LIBRARY_static_init_##ns( \
      torch::Library::DEF,                                                     \
      &TORCH_LIBRARY_init_##ns,                                                \
      C10_STRINGIZE(ns),                                                       \
      std::nullopt,                                                            \
      __FILE__,                                                                \
      __LINE__);                                                               \
  void TORCH_LIBRARY_init_##ns(torch::Library& m)

拆开看：

⭐ 重点是中间那行：全局对象的构造函数会在这个动态库被加载时自动执行。

⭐ 注册不是在你调用算子的时候发生的，是在那个 .so / .dll 被载入进程的那一刻发生的。 库没加载 = 表上根本没这一行 = 你得到的是「算子不存在」，而不是「算子出错」。

这一条把好几件事串起来了：

⚠️ 静态初始化的老毛病这里同样有：多个库注册同一个格子会互相覆盖（第四节那个 Meta (inactive) 就是活例子），而谁先谁后取决于加载顺序，不是你能控制的。

🚦 给已有的 aten 算子换实现

同一套机制，命名空间换成 aten：

// 🗓️ 未实跑
TORCH_LIBRARY_IMPL(aten, CPU, m) {
    m.impl("gelu", my_faster_gelu);     // ⚠️ 覆盖官方的 CPU gelu
}

⚠️ 能做，但基本不该做。 它是进程全局的：库一被加载，全进程所有 gelu 都换成你的了，包括别的库调的那些。⭐ 想加速某个算子，正路是注册到你自己的命名空间再显式调用，或者走 AI 基础设施 07 讲的编译器路线。

🧩 六、⭐ 让它可导：C++ 侧的那一半

⚠️ 边界先说清楚：autograd.Function 怎么写、forward 存什么 backward 还什么 —— 那是框架的 Python 接口，归 《PyTorch 这个框架本身》09。本节只讲在 C++ 这边，反向是怎么被挂进那张表的。

回到第三节的优先级：Autograd 在 backend 之上。所以「填了 CPU 格、没填 Autograd 格」的算子，dispatcher 会直接掉进 CPU kernel，没人记录反向图。三条路：

// 🗓️ 未实跑 —— 本机无编译环境
#include <torch/library.h>
#include <torch/autograd.h>

class ScaledAdd : public torch::autograd::Function<ScaledAdd> {
 public:
  static at::Tensor forward(torch::autograd::AutogradContext* ctx,
                            const at::Tensor& a, const at::Tensor& b, double s) {
    ctx->saved_data["s"] = s;                    // 反向要用的标量存这里
    at::AutoDispatchBelowADInplaceOrView guard;  // ⭐ 往下 redispatch，跳过 autograd 层
    return scaled_add_cpu(a, b, s);              //     不写这行就无限递归
  }
  static torch::autograd::variable_list backward(
      torch::autograd::AutogradContext* ctx, torch::autograd::variable_list grad) {
    double s = ctx->saved_data["s"].toDouble();
    return {grad[0], grad[0] * s, at::Tensor()};  // ⚠️ 对 double 参数返回空 Tensor 占位
  }
};

TORCH_LIBRARY_IMPL(myops, Autograd, m) {          // ⭐ 填 Autograd 那一格
  m.impl("scaled_add", [](const at::Tensor& a, const at::Tensor& b, double s) {
    return ScaledAdd::apply(a, b, s);
  });
}

⭐ AutoDispatchBelowADInplaceOrView guard; 是全段重点，而且形状你已经见过 —— 它是个 RAII 对象（和 01 章的 gil_scoped_release 同一个套路），作用域内把分派起点压到 autograd 之下。没有它，forward 里再调一次这个算子会又进 autograd kernel。

⚠️ 只填 CPU 不填 Autograd 的症状：前向完全正常，反向时要么报 element 0 of tensors does not require grad，要么梯度悄悄变成 0。💀 后者不报错，训练看起来在跑、loss 就是不降 —— 它长得不像 bug，所以调试成本极高。

🧩 七、CUDA 那一格，和不想编译的路子

⭐ 好消息：CUDA 没有新机制，同一套宏换个键：

// 🗓️ 未实跑 —— 本机无 nvcc，torch 也是 CPU 版
TORCH_LIBRARY_IMPL(myops, CUDA, m) {
    m.impl("scaled_add", scaled_add_cuda);   // ⭐ 唯一的区别就是这个键
}

写 CUDA 难在 kernel 本身（SM、warp、合并访存归 AI 基础设施 02），不难在注册 —— 注册就是上面这三行。

只想先试一下机制的话，有两条不用自己配编译工程的路：

⭐ 建议顺序：先用 Python 接口把注册跑通、确认分派真的进了你的函数，再决定要不要下到 C++。 这样出问题时你知道是注册错了还是 kernel 错了。具体写法归 《PyTorch 这个框架本身》09。

⚠️⚠️ 动手之前先确认你真的需要：AI 基础设施 07 给了「什么时候值得手写 kernel」的三个条件，不满足就别写。这和 01 章第三节那张反面表是同一件事的两个版本。

🔗 这一章连到哪里

✅ 检查点

1. torch.nn.functional.gelu(torch.ones(2,2).to_sparse()) 报什么错？这条报错里的两个片段各对应 dispatcher 的哪个坐标轴？异常类型到底是什么？
2. 为什么「稀疏张量不支持激活函数」这个说法是错的？本章用哪组实跑对照证明了它？
3. aten::add.Tensor(...) 里 aten / add / .Tensor 各是什么？aten::add 实跑有几个重载，这对「表的行」意味着什么？
4. schema 里的 Tensor(a!) 是什么意思？填错它为什么特别难查？
5. 本章实跑的三个 keyset 里，requires_grad=True 那个和普通 CPU 差在哪？说明了什么？
6. torch._C._dispatch_dump("aten::gelu") 的输出里哪一行没有出现？路径里带 build\ 的文件有什么特点？Meta 那行为什么值得注意？
7. TORCH_LIBRARY / TORCH_LIBRARY_IMPL / m.fallback 各对应表上的什么动作？
8. TORCH_LIBRARY 宏展开成一个什么东西？由此推出注册发生在哪一刻？这怎么解释「装了扩展但 torch.ops.myops 报 AttributeError」？
9. 只注册 CPU kernel 不注册 Autograd 的症状是什么？为什么这类 bug 特别贵？AutoDispatchBelowADInplaceOrView guard; 不写会怎样？

🛑 可以停在这里

⚡ 走神救援

🎯 dispatcher 是一张二维表：行是算子名、列是分派键、格子里是函数指针，注册就是填格子。

开场那条报错一次给出两个坐标：「Could not run aten::gelu with arguments from the SparseCPU backend」。⭐ CPU 机器就能复现，不需要 GPU。 它不是「稀疏不支持激活函数」——同一个稀疏张量有的算子能跑有的不能，表是逐格填的。

⚠️ 本机 CUDA 那一列全空只是因为装的是 CPU 版轮子——⭐ 这反过来说明表不是编译期常量，是运行时被一批批加载进来的。

📐 行坐标是 schema，而且每个重载都是独立的一行（aten::add 实测有十几行，不是一行）。⭐ schema 是唯一真相源：C++ 签名、torch.ops.*、类型存根、头文件全从它生成。记号里最要命的是 Tensor(a!) 的那个 !（会就地写）——它是给编译器的合同不是文档，⭐ 填错的后果是别的通道随机出错，而不是当场报错。

列坐标是 DispatchKeySet。⚠️⚠️ requires_grad=True 的 keyset 和普通 CPU 一模一样——求不求导不靠分派键区分。⭐ 所以「注册 CPU kernel」和「让它可导」是两件独立的事。

🔍 _dispatch_dump 能把整行 dump 出来，三条读法值得记：带 build\ 的文件在仓库里根本不存在（按 schema 生成的）；[alias] 是别名键，一次填一批格子；而 (inactive) 标记的是被覆盖掉的注册。

下一节 👉 06-读源码路线图.md