第8章 在 C++ 中注册 PyTorch 派发算子
第8章：在 C++ 中注册派发的算子
创建时间：2020年7月22日 | 最后更新：2024年7月22日 | 最后验证：2024年11月5日
警告
本教程在 PyTorch 2.4 中已被弃用。请参阅 PyTorch 自定义算子以获取关于扩展 PyTorch 最新指南。
Dispatcher（派发器）是 PyTorch 的内部组件，负责在调用如 `torch::add` 这样的函数时确定实际应运行哪段代码。由于 PyTorch 操作需要处理许多叠加在基础之上的“横切关注点”（cross-cutting concerns），这一过程往往并非简单直观。以下是它处理的一些事项示例：
- 根据输入张量的设备类型，在算子的 CPU 和 CUDA 实现之间切换。 - 根据是否需要自动求导处理，在算子的 autograd 和后端实现之间切换。 - 在自动混合精度（AMP）所需时应用自动类型转换（autocasting）。 - 当算子在 `vmap` 调用下运行时，应用批处理规则。 - 如果正在追踪模型以用于导出，则追踪操作的执行过程。
如果在编写自定义算子代码时，发现自己在手动编写 if 语句来处理这些情况，Dispatcher API 可以帮助你组织代码。（反之，如果你的自定义算子非常简单且仅用于 CPU 推理，可能不需要使用 Dispatcher，直接使用基础 API 即可。）
在本教程中，我们将描述如何结构化自定义算子的注册，以便利用 Dispatcher 组织各个组件。我们假设你熟悉如何注册算子以及如何编写自定义 autograd 函数。
定义 Schema 和后端实现
Dispatcher 背后的通用原则是将算子的实现划分为多个 kernel（内核），每个 kernel 实现特定 dispatch key（派发键，如 CPU、CUDA）的功能。Dispatcher 在你调用算子时确定优先级最高的 dispatch key（通过检查张量参数以及某些线程局部状态来完成），并将控制权转移给该 dispatch key 对应的 kernel。最终效果是，当你调用一个算子时，首先执行 Autograd kernel，然后根据传入张量的设备类型重新派发（redispatch）到相应的后端 kernel。
让我们看看实现这一机制涉及的各个部分。首先，我们必须定义相关算子的 schema。与简单的 pybind11 风格的算子注册不同，我们在此时并不提供算子的实际实现；我们只是提供一个 schema 字符串，指定所有其他 kernel 都将遵守的算子类型签名：
```cpp TORCH_LIBRARY(myops, m) { m.def("myadd(Tensor self, Tensor other) -> Tensor"); } ```
接下来，我们需要实际提供这个算子的一些实现。为了具体化，这里是一个简单的 CPU 加法实现：
```cpp
Tensor myadd_cpu(const Tensor& self_, const Tensor& other_) {
TORCH_CHECK(self_.sizes() == other_.sizes());
TORCH_INTERNAL_ASSERT(self_.device().type() == DeviceType::CPU);
TORCH_INTERNAL_ASSERT(other_.device().type() == DeviceType::CPU);
Tensor self = self_.contiguous();
Tensor other = other_.contiguous();
Tensor result = torch::empty(self.sizes(), self.options());
const float* self_ptr = self.data_ptr
我们希望将此函数注册为 `myops::myadd` 的一个实现。然而，简单的注册方式（`def("myadd", myadd_cpu)`）会将该 kernel 注册为在所有情况下运行，即使张量不是 CPU 张量！（内部，我们将这些称为“兜底”（catch-all）kernel，因为它们捕获所有情况。）为了确保 `myadd_cpu` 仅对 CPU 张量运行，我们可以使用 `TORCH_LIBRARY_IMPL` 宏：
```cpp TORCH_LIBRARY_IMPL(myops, CPU, m) { m.impl("myadd", myadd_cpu); } ```
`TORCH_LIBRARY_IMPL` 允许我们为特定 dispatch key（在此情况下为 CPU）的算子注册实现。每次调用 `impl` 都会将一个 CPU kernel 与相应的算子关联起来（该算子之前已在 `TORCH_LIBRARY` 块中定义）。如果我们还有一个 CUDA 实现 `myadd_cuda`，可以在单独的 `TORCH_LIBRARY_IMPL` 块中注册它：
```cpp TORCH_LIBRARY_IMPL(myops, CUDA, m) { m.impl("myadd", myadd_cuda); } ```
这些注册可以分布在不同的文件甚至库边界之间；例如，你可以将这两个 `TORCH_LIBRARY_IMPL` 块编译到单独的 `myops_cpu` 和 `myops_cuda` 动态库中。通常来说，你的注册结构如下：
- 一个集中的 `TORCH_LIBRARY`，在命名空间中列出的每个自定义算子。 - 每个 dispatch key（如 CPU 或 CUDA）对应一个 `TORCH_LIBRARY_IMPL`，用于注册该 key 的实现。如果你愿意，可以进一步将 `TORCH_LIBRARY_IMPL` 块细分为每个算子一个块。这对于每个算子实现位于单独文件但不想在头文件中暴露算子的情况非常方便；你只需将注册代码放在定义算子的 cpp 文件中即可。
注意
你知道可以为 PyTorch 中现有的核心算子编写 `TORCH_LIBRARY_IMPL` 块吗？这就是 PyTorch 中 XLA 支持实现的方式：`torch_xla` 库包含一个提供 XLA dispatch key 上所有基本算子实现的 `TORCH_LIBRARY_IMPL`。
对于不需要 autograd 的算子
注意： 本节仅适用于 PyTorch >= 1.10 版本。
在下一节中，我们将讨论如何为算子添加 autograd 支持。但对于不需要 autograd 支持的操作，注册以下 kernel 可以提高可用性，并使你的操作表现得像 PyTorch 的内置算子。
```cpp TORCH_LIBRARY_IMPL(myops, Autograd, m) { m.impl(op, autogradNotImplementedFallback()); } ```
上述代码注册了一个 Autograd kernel，它在前向传播时添加一个虚拟的 `NotImplemented` 节点（保留输入的 `require_grad` 属性）。在反向传播时，`NotImplemented` 节点会引发错误。这在调试较大模型时很有帮助，因为在之前很难准确定位 `requires_grad` 属性在前向传播过程中是在哪里丢失的。
原地操作或视图操作
为了确保正确性和最佳性能，如果你的算子原地修改输入或返回与其中一个输入存在别名关系的张量，应采取两个额外步骤：
- 除了上面的 Autograd kernel，还应注册一个 `ADInplaceOrView` kernel。此 kernel 处理必要的簿记工作，以确保原地或视图操作的正确性。需要注意的是，此 `ADInplaceOrView` kernel 仅应与 `autogradNotImplementedFallback` 一起使用。
```cpp TORCH_LIBRARY_IMPL(myops, Autograd, m) { m.impl(op, autogradNotImplementedFallback()); } TORCH_LIBRARY_IMPL(myops, ADInplaceOrView, m) { m.impl(op, autogradNotImplementedInplaceOrViewFallback()); } ```
上面注册的 Autograd 或 `ADInplaceOrView` boxed kernel 依赖于其逻辑中的算子 schema 信息。如果你的算子原地修改输入或返回与输入之一存在别名关系的张量，确保 schema 正确反映这一点非常重要。有关如何标注 schema 的更多信息，请参见此处。
添加 autograd 支持
到目前为止，我们拥有一个具有 CPU 和 CUDA 实现的算子。如何为它添加 autograd 支持？正如你可能猜测的那样，我们将注册一个 autograd kernel（类似于自定义 autograd 函数教程中描述的内容）！然而，有一个转折：与 CPU 和 CUDA kernel 不同，autograd kernel 需要重新派发：它需要回调到 Dispatcher 以获取推理 kernel，例如 CPU 或 CUDA 实现。
因此，在编写 autograd kernel 之前，让我们编写一个派发函数，它调用 Dispatcher 以找到算子的正确 kernel。这个函数构成了你算子的公开 C++ API——事实上，PyTorch C++ API 中所有的张量函数在底层都以相同的方式调用 Dispatcher。派发函数的样子如下：
```cpp
Tensor myadd(const Tensor& self, const Tensor& other) {
static auto op = torch::Dispatcher::singleton()
.findSchemaOrThrow("myops::myadd", "")
.typed
让我们分解一下：
- 在第一行中，我们从 Dispatcher 中查找对应于我们要派发的算子的类型化算子句柄。`findSchemaOrThrow` 接受两个参数：算子（带命名空间的）名称和算子的重载名称（通常只是空字符串）。`typed` 将动态类型句柄转换为静态类型句柄（进行运行时测试以确保你提供了正确的 C++ 类型），以便我们可以对其进行普通的 C++ 调用。我们传入 `decltype(myadd)`，因为派发函数的类型与注册到 Dispatcher 的底层 kernel 的类型相同。出于性能考虑，此计算在静态变量中进行，因此我们只需要执行一次（缓慢的）查找。如果你拼错了想要调用的算子名称，在第一次调用此函数时，此查找会引发错误。 - 在第二行中，我们简单地使用传入派发函数的所有参数调用算子句柄。这将实际调用 Dispatcher，最终控制权将转移到此调用适当的 kernel。
有了派发函数，我们现在可以编写 autograd kernel：
```cpp
class MyAddFunction : public torch::autograd::Function
static tensor_list backward(AutogradContext *ctx, tensor_list grad_outputs) { auto grad_output = grad_outputs[0]; return {grad_output, grad_output}; } };
Tensor myadd_autograd(const Tensor& self, const Tensor& other) { return MyAddFunction::apply(self, other)[0]; } ```
autograd 函数使用 `torch::autograd::Function` 正常编写，只是在 `forward()` 中不直接编写实现，而是：
1. 使用 `at::AutoNonVariableTypeMode` RAII 守卫关闭 autograd 处理，然后 2. 调用派发函数 `myadd` 以回调到 Dispatcher。
如果没有第 1 步，你的调用将无限循环（并导致堆栈溢出），因为 `myadd` 会将你送回此函数（因为最高的优先级 dispatch key 仍然是 autograd）。有了第 1 步，autograd 被排除在考虑的 dispatch key 集合之外，我们将进入下一个处理器，即 CPU 或 CUDA。
现在我们可以以注册 CPU/CUDA 函数的方式注册此函数：
```cpp TORCH_LIBRARY_IMPL(myops, Autograd, m) { m.impl("myadd", myadd_autograd); } ```
注意
在这个例子中，我们将 kernel 注册到 Autograd，这将其安装为所有后端的 autograd kernel。你也可以使用相应的特定后端 dispatch key（例如，`AutogradCPU` 或 `AutogradCUDA`）为特定后端注册优化 kernel。要更详细地探索这些以及其他 dispatch key 选项，请查看 `torch/_python_dispatcher.py` 中提供的 PythonDispatcher 工具。
超越 autograd
从某种意义上说，Dispatcher 做的并不多：它只是实现了一个华丽的 if 语句，大致如下：
```cpp class MyAddFunction : ... { public: static Tensor forward( AutogradContext *ctx, torch::Tensor self, torch::Tensor other) {
if (self.device().type() == DeviceType::CPU) { return add_cpu(self, other); } else if (self.device().type() == DeviceType::CUDA) { return add_cuda(self, other); } else { TORCH_CHECK(0, "Unsupported device ", self.device().type()); } } ... } ```
那么为什么使用 Dispatcher？有以下几个原因：
- 它是去中心化的。你可以组装算子的所有部分（CPU、CUDA、Autograd），而不必编写一个引用所有这些的集中式 if 语句。重要的是，第三方可以为其他方面注册额外的实现，而无需修补算子的原始定义。我们将在“为新的后端扩展 Dispatcher”中更多地讨论扩展 Dispatcher。 - 它支持比 CPU、CUDA 和 Autograd 更多的 dispatch key。你可以在 `c10/core/DispatchKey.h` 中查看 PyTorch 中当前实现的所有 dispatch key 的完整列表。这些 dispatch key 为算子实现了各种可选功能，如果你决定让你的自定义算子支持此功能，你只需为相应的 key 注册一个 kernel。 - Dispatcher 实现了 boxed fallback 函数的支持，这些函数可以一次性实现并应用于系统中的所有算子。Boxed fallback 可用于提供 dispatch key 的默认行为；如果使用 Dispatcher 实现算子，你也选择了所有这些操作的 fallback。
以下是一些你可能需要定义算子的特定 dispatch key。
Autocast
Autocast dispatch key 实现了对自动混合精度（AMP）的支持。 autocast 包装 kernel 通常在运行算子之前将传入的 float16 或 float32 CUDA 张量转换为某种首选精度。 例如，在浮点 CUDA 张量上进行的矩阵乘法（matmuls）和卷积（convolutions）通常在 float16 下运行更快且占用内存更少，同时不影响收敛。 Autocast 包装器仅在启用 autocast 的上下文中生效。 以下是假设的自定义 matmul 的 autocast 包装器及其注册：
```cpp
// Autocast-specific helper functions
#include
Tensor mymatmul_autocast(const Tensor& self, const Tensor& other) { c10::impl::ExcludeDispatchKeyGuard no_autocast(c10::DispatchKey::Autocast); return mymatmul(at::autocast::cached_cast(at::kHalf, self), at::autocast::cached_cast(at::kHalf, other)); }
TORCH_LIBRARY_IMPL(myops, Autocast, m) { m.impl("mymatmul", mymatmul_autocast); } ```
`cached_cast(kHalf, tensor)` 将 tensor 转换为 float16，如果 tensor 是 CUDA 且为 float32；否则，它保持 tensor 不变（参见原生 autocast 算子的适用性策略）。 这确保如果网络在任何 float16 和 float32 CUDA 张量的混合上调用 `mymatmul`，`mymatmul` 将以 float16 运行。同时，使用非 CUDA、整数类型或 float64 输入调用 `mymatmul` 不受影响。 建议在自定义 autocast 包装器中使用 `cached_cast` 遵循原生适用性策略，但不是强制要求的。例如，如果你希望强制所有输入类型执行 float16，可以返回 `mymatmul(self.half(), other.half());` 而不是使用 `cached_cast`。 注意，与我们的 autograd kernel 一样，我们在重新派发之前排除了 Autocast key。 默认情况下，如果没有提供 autocast 包装器，我们直接回落到常规算子实现（不发生自动类型转换）。（在这个例子中我们没有使用 `myadd`，因为逐元素加法不需要自动类型转换，应该直接回落。）
何时应该注册 autocast 包装器？不幸的是，关于算子首选精度没有明确的规则。通过查看转换列表（cast lists），你可以了解一些原生算子的首选精度。 总体指导原则：
- 执行归约（reductions）的算子可能应该在 float32 中执行。 - 底层执行卷积或 gemm 的任何算子可能应该在 float16 中执行。 - 具有多个浮点张量输入的其他算子应将其标准化为一种通用精度（除非实现支持不同精度的输入）。
如果你的自定义算子属于第三类，`promote_type` 模板有助于确定输入张量中存在的最大浮点类型，这是执行类型的最安全选择：
```cpp
#include
Tensor my_multiple_input_op_autocast(const Tensor& t0, const Tensor& t1) { c10::impl::ExcludeDispatchKeyGuard no_autocast(c10::DispatchKey::Autocast); // 所需的 at::kHalf 参数是一个乐观的初始猜测。 auto exec_type = at::autocast::promote_type(at::kHalf, t0, t1); return my_multiple_input_op(at::autocast::cached_cast(exec_type, t0), at::autocast::cached_cast(exec_type, t1)); } ```
如果你的自定义算子启用了 autograd，你只需要为 autograd 包装器注册的同名操作编写并注册一个 autocast 包装器。例如，如果你希望为 autograd 部分中显示的 `myadd` 函数创建一个 autocast 包装器，你只需要：
```cpp
Tensor myadd_autocast(const Tensor& self, const Tensor& other) {
c10::impl::ExcludeDispatchKeyGuard no_autocast(c10::DispatchKey::Autocast);
return myadd(at::autocast::cached_cast(
TORCH_LIBRARY_IMPL(myops, Autocast, m) { m.impl("myadd", myadd_autocast); } ```
使反向方法兼容 autocast 不需要额外的复杂操作。然而，自定义 autograd 函数中定义的反向方法将以 autocast 为前向方法设置的相同 dtype 运行，因此你应该选择适合前向和反向方法的 `
Batched
批处理张量允许你按示例（per-example）方式编写代码，然后在 `vmap` 调用下自动进行批处理。编写批处理规则的 API 目前处于开发阶段，但一旦稳定，你就可以通过在 Batched dispatch key 注册 kernel 为你的算子添加对 `vmap` 的支持。
Tracer
Tracer dispatch key 实现了在运行 `torch.jit.trace` 时将算子调用记录到 trace 中的支持。我们打算提供 boxed fallback 来实现任意操作的追踪，请参见 issue #41478 以跟踪进度。
