在 C++中注册调度操作符 §

创建于：2025 年 4 月 1 日 | 最后更新：2025 年 4 月 1 日 | 最后验证：2024 年 11 月 5 日

警告

本教程自 PyTorch 2.4 版本起已过时。请参阅 PyTorch 自定义算子以获取最新的扩展 PyTorch 自定义算子的指南。

分发器是 PyTorch 的一个内部组件，负责确定在调用类似 torch::add .这样的函数时应该实际运行什么代码。这可能相当复杂，因为 PyTorch 操作需要处理许多相互交织的跨领域问题。以下是它处理的一些示例：

• 根据输入张量的设备，在操作符的 CPU 和 CUDA 实现之间切换。

• 根据是否需要自动微分处理，在操作符的自动微分和后端实现之间切换。

• 在必要时应用自动类型转换以实现自动混合精度。

• 在操作符在 vmap 调用下运行时应用批处理规则。

• 如果您正在跟踪模型以进行导出，则跟踪操作的执行。

如果在您的自定义算子代码中，您发现自己需要手动编写 if 语句来处理这些情况，分发器 API 可以帮助组织您的代码。（相反，如果您的自定义算子非常简单，并且仅用于 CPU 推理，您可能不需要使用分发器，只需使用基本 API 即可。）

在本教程中，我们将介绍如何构建自定义算子注册以使用调度器来组织各种组件。我们将假设您熟悉如何注册算子以及如何编写自定义自动微分函数。

定义模式和后端实现

调度器背后的基本原理是将算子的实现划分为多个内核，每个内核实现特定调度键的功能，例如 CPU、CUDA。调度器确定在调用算子时最高优先级的调度键是什么（这是通过查看张量参数以及一些线程局部状态来完成的），并将控制权转移到该调度键的内核。最终效果是，当您调用算子时，我们首先执行 Autograd 内核，然后根据传入张量的设备类型重新调度到后端内核。

让我们看看实现这一功能所涉及的各个部分。首先，我们必须定义所涉及操作符的模式。与简单的 pybind11 风格操作符注册不同，我们在此处并没有提供操作符的实现；我们只是提供一个模式字符串，指定操作符的类型签名，所有其他内核都将遵循这个签名：

TORCH_LIBRARY(myops, m) {
  m.def("myadd(Tensor self, Tensor other) -> Tensor");
}

接下来，我们需要实际提供这个操作符的实现。为了具体化，这里是一个在 CPU 上实现加法的非常简单的示例：

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<float>();
  const float* other_ptr = other.data_ptr<float>();
  float* result_ptr = result.data_ptr<float>();
  for (int64_t i = 0; i < result.numel(); i++) {
    result_ptr[i] = self_ptr[i] + other_ptr[i];
  }
  return result;
}

我们希望将这个函数注册为 myops::myadd 的实现。然而，注册它的简单方法（ def("myadd", myadd_cpu) ）会将内核注册为在所有情况下运行，即使张量不是 CPU 张量！（内部，我们称这些为“捕获所有情况”的内核，因为它们捕获所有情况。）为了确保 myadd_cpu 只在 CPU 张量上运行，我们可以使用 TORCH_LIBRARY_IMPL 宏：

TORCH_LIBRARY_IMPL(myops, CPU, m) {
  m.impl("myadd", myadd_cpu);
}

TORCH_LIBRARY_IMPL 允许我们在特定的调度键（在这种情况下，CPU）上注册操作符的实现。每次调用 impl 都会将一个 CPU 内核与相应的操作符（我们之前在 TORCH_LIBRARY 块中定义的）关联起来。如果我们还有 CUDA 实现 myadd_cuda ，我们可以在单独的 TORCH_LIBRARY_IMPL 块中注册它：

TORCH_LIBRARY_IMPL(myops, CUDA, m) {
  m.impl("myadd", myadd_cuda);
}

这些注册可以分散到文件中，甚至跨越库边界；例如，您可以将这两个 TORCH_LIBRARY_IMPL 块编译成单独的 myops_cpu 和 myops_cuda 动态库。一般来说，您的注册结构将类似于以下：

1. 在一个集中位置列出您命名空间中每个自定义操作符的单个 TORCH_LIBRARY 。

2. 每个调度键一个 TORCH_LIBRARY_IMPL ，用于注册该键的实现（例如，CPU 或 CUDA）。如果您愿意，可以将 TORCH_LIBRARY_IMPL 块进一步细分为每个操作符一个块。如果您为每个操作符实现有一个单独的文件，但又不想在头文件中暴露操作符，这很方便；只需将注册放在定义您的操作符的 cpp 文件中即可。

备注

你知道吗，你还可以为现有的 PyTorch 核心操作符编写 TORCH_LIBRARY_IMPL 块。这就是 PyTorch 对 XLA 支持的实现方式： torch_xla 库包含一个 TORCH_LIBRARY_IMPL ，它为所有基本操作符在 XLA 调度键上提供了实现。

对于不需要 autograd 的操作符

注意：本节仅适用于 PyTorch >= 1.10 版本。

在下一节中，我们将讨论如何为操作符添加 autograd 支持。但对于不需要 autograd 支持的操作，以下内核的注册可以改善可用性，并使你的操作符表现得像 PyTorch 的内置操作符。

TORCH_LIBRARY_IMPL(myops, Autograd, m) {
  m.impl(op, autogradNotImplementedFallback());
}

上述行注册了一个 Autograd 内核，该内核在正向传播时附加一个虚拟的 NotImplemented 节点（保留输入的 require_grad -ness）。在反向传播时， NotImplemented 节点会引发错误。这对于调试较大的模型非常有帮助，因为在之前，很难确定在正向传播过程中哪里丢失了 requires_grad -ness。

原地或视图操作

为了确保正确性和最佳性能，如果您的操作原地修改输入或返回与输入之一别名化的张量，则应采取以下两个额外步骤：

1. 除了上面的 Autograd 内核外，还需要注册一个 ADInplaceOrView 内核。此内核处理必要的账目，以确保原地或视图操作的正确性。需要注意的是，此 ADInplaceOrView 内核仅应与 autogradNotImplementedFallback 一起使用。

TORCH_LIBRARY_IMPL(myops, Autograd, m) {
  m.impl(op, autogradNotImplementedFallback());
}
TORCH_LIBRARY_IMPL(myops, ADInplaceOrView, m) {
  m.impl(op, autogradNotImplementedInplaceOrViewFallback());
}

2. 上面的 Autograd 或 ADInplaceOrView 框选的核依赖于其逻辑中的操作符模式信息。如果你的操作符在原地修改输入或返回一个与输入之一有别名关系的张量，那么确保你的模式正确反映这一点是很重要的。有关如何注释模式的更多信息，请参阅此处。

添加自动微分支持

到目前为止，我们已经有一个同时具有 CPU 和 CUDA 实现的操作符。我们如何为它添加自动微分支持呢？正如你可能猜到的，我们将注册一个自动微分核（类似于自定义自动微分函数教程中描述的核）！然而，有一个转折点：与 CPU 和 CUDA 核不同，自动微分核需要重新调度：它需要回调到调度器以获取到推理核，例如 CPU 或 CUDA 实现。

因此，在我们编写自动微分核之前，让我们编写一个调度函数，该函数调用调度器以找到适合你的操作符的正确核。这个函数构成了你的操作符的公共 C++ API - 事实上，PyTorch 的所有 C++ API 中的张量函数都在幕后以相同的方式调用调度器。以下是调度函数的示例：

Tensor myadd(const Tensor& self, const Tensor& other) {
  static auto op = torch::Dispatcher::singleton()
    .findSchemaOrThrow("myops::myadd", "")
    .typed<decltype(myadd)>();
  return op.call(self, other);
}

让我们来分解一下：

• 在第一行，我们从调度器中查找一个与将要调度的操作员对应的已输入操作员句柄。 findSchemaOrThrow 接受两个参数：操作员的（命名空间限定）名称和操作符的重载名称（通常是空字符串）。 typed 将动态类型句柄转换为静态类型句柄（进行运行时测试以确保您已给出正确的 C++类型），这样我们就可以在它上面进行正常的 C++调用。我们传递 decltype(myadd) ，因为调度函数的类型与注册到调度器的底层内核的类型相同。

  为了性能，这个计算是在静态变量中完成的，这样我们只需要进行一次（慢速）查找。如果您输入了要调用的操作员名称，这次查找将在第一次调用此函数时出错。

• 在第二行，我们只需 call 操作员句柄和传递给调度函数的所有参数。这实际上将调用调度器，最终控制权将转移到适合此调用的任何内核。

拥有调度函数后，我们现在可以编写自动微分内核：

class MyAddFunction : public torch::autograd::Function<MyAddFunction> {
 public:
  static Tensor forward(
      AutogradContext *ctx, torch::Tensor self, torch::Tensor other) {
    at::AutoNonVariableTypeMode g;
    return myadd(self, other);
  }

  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];
}

自动微分函数的编写与平常一样使用 torch::autograd::Function ，只是我们不在 forward() 中直接编写实现：

1. 使用 at::AutoNonVariableTypeMode RAII 保护来关闭自动微分处理，然后

2. 调用调度函数 myadd 以回调到调度器。

没有使用（1），您的调用将无限循环（并发生栈溢出），因为 myadd 会将您送回此函数（因为最高优先级的调度键仍然是 autograd。）使用（1），autograd 被排除在考虑的调度键集合之外，我们将转到下一个处理器，这些处理器将是 CPU 和 CUDA。

我们现在可以以与注册 CPU/CUDA 函数相同的方式注册此函数：

TORCH_LIBRARY_IMPL(myops, Autograd, m) {
  m.impl("myadd", myadd_autograd);
}

备注

在此示例中，我们将内核注册到 Autograd ，将其安装为所有后端的 autograd 内核。您还可以通过使用相应的后端特定调度键来注册针对特定后端的优化内核，例如 AutogradCPU 或 AutogradCUDA 。要详细了解这些和其他调度键选项，请查看 torch/_python_dispatcher.py 中提供的 PythonDispatcher 工具。

超越 autograd

在某种意义上，调度器并没有做太多：它所做的只是实现了一个被美化的 if 语句，类似于以下这种：

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());
    }
  }
  ...
}

那么为什么还要使用调度器呢？有几个原因：

1. 它是去中心化的。你可以组装操作符（CPU、CUDA、Autograd）的所有部件，而不需要编写一个单独的、集中的 if 语句来引用它们。重要的是，第三方可以为操作符的其他方面注册额外的实现，而无需修补操作符的原始定义。我们将在扩展调度器以支持新的后端中进一步讨论扩展调度器。

2. 它支持的调度键比 CPU、CUDA 和 Autograd 更多。你可以在 PyTorch 中查看当前已实现的完整调度键列表，见 c10/core/DispatchKey.h 。这些调度键为操作符实现了各种可选功能，如果你决定你想你的自定义操作符支持这些功能，你只需为适当的关键注册一个内核即可。

3. 分发器实现了对封装回退函数的支持，这些函数只需实现一次即可应用于系统中的所有运算符。封装回退函数可以用于为分发键提供默认行为；如果您使用分发器来实现您的运算符，您也将选择所有这些操作的回退。

这里有一些您可能需要为它们定义运算符的特定分发键。

自动转换

自动转换分发键实现了对自动混合精度（AMP）的支持。自动转换包装内核通常在运行运算符之前将传入的 float16 或 float32 CUDA 张量转换为某种首选精度。例如，在浮点 CUDA 张量上进行的矩阵乘法和卷积通常在 float16 中运行更快，使用更少的内存，而不会损害收敛性。自动转换包装器仅在自动转换启用的上下文中才有作用。

这里是一个假设的定制矩阵乘法的自动转换包装器及其注册：

// Autocast-specific helper functions
#include <ATen/autocast_mode.h>

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);
}

如果 tensor 是 CUDA 和 float32 ，则将 tensor 转换为 float16 ，否则保持 tensor 不变（参照原生自动转换操作的资格政策）。这确保了如果网络在 float16 和 float32 的任何混合 CUDA 张量上调用 mymatmul ， mymatmul 将在 float16 中运行。同时，对非 CUDA、整型或 float64 输入的 mymatmul 调用不受影响。在您的自动转换包装器中使用 cached_cast 来遵循原生资格政策是推荐的，但不是必需的。例如，如果您想强制所有输入类型执行 float16 ，则可以使用 return mymatmul(self.half(), other.half()); 而不是使用 cached_cast 。

注意，就像我们的自动微分内核一样，我们在重新调度之前排除了 Autocast 键。

默认情况下，如果没有提供自动转换包装器，我们将直接跳转到常规操作符实现（不进行自动转换）。（我们没有使用 myadd ，因为点积加法不需要自动转换，应该直接跳过。）

当应该注册自动类型转换包装器？遗憾的是，对于操作符首选精度的规则并不明确。您可以通过查看类型列表来了解一些原生操作符的首选精度。一般指导原则：

• 执行归约操作的运算符可能应该在 float32 中执行，

• 任何在底层执行卷积或 gemm 操作的运算符可能应该在 float16 中执行，

• 其他具有多个浮点张量输入的运算符应将它们标准化为共同精度（除非实现支持不同精度的输入）。

如果您的自定义操作属于第三类， promote_type 模板可以帮助确定输入张量中存在的最宽浮点类型，这是执行类型的最安全选择：

#include <ATen/autocast_mode.h>

Tensor my_multiple_input_op_autocast(const Tensor& t0, const Tensor& t1) {
  c10::impl::ExcludeDispatchKeyGuard no_autocast(c10::DispatchKey::Autocast);
  // The required at::kHalf argument is an optimistic initial guess.
  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));
}

如果您的自定义操作启用了自动微分，您只需要为注册了自动微分包装器的同名函数编写并注册一个自动类型转换包装器。例如，如果您想要自动类型转换包装器，就像自动微分部分中显示的 myadd 函数，您只需要

Tensor myadd_autocast(const Tensor& self, const Tensor& other) {
  c10::impl::ExcludeDispatchKeyGuard no_autocast(c10::DispatchKey::Autocast);
  return myadd(at::autocast::cached_cast(<desired dtype>, self),
               at::autocast::cached_cast(<desired dtype>, other));
}

TORCH_LIBRARY_IMPL(myops, Autocast, m) {
  m.impl("myadd", myadd_autocast);
}

没有单独的技巧可以使反向方法与自动类型转换兼容。然而，在您的自定义自动微分函数中定义的反向方法将以与自动类型转换为正向方法设置的相同数据类型运行，因此您应该选择一个既适合您的正向方法也适合您的反向方法的 <desired dtype> 。

批处理

批处理张量允许您以每个示例的方式编写代码，然后在运行 vmap 时自动批处理。当前正在开发编写批处理规则的 API，一旦稳定，您可以通过在批处理调度键中注册内核来为您的算子添加对 vmap 的支持。

跟踪器

跟踪器调度键实现了在运行 torch.jit.trace 时将算子的调用记录到跟踪中的支持。我们打算提供一个包装的回退实现，以实现对任意操作的跟踪，请参阅问题 #41478 跟踪进度。