评价此页
★ ★ ★ ★ ★
advanced/dispatcher

在 C++ 中注册分发算子#

创建日期：2020年7月22日 | 最后更新：2024年7月22日 | 最后验证：2024年11月5日

警告

本教程自 PyTorch 2.4 起已弃用。请参阅 PyTorch 自定义算子 (PyTorch Custom Operators) 获取有关使用自定义算子扩展 PyTorch 的最新指南。

分发器（dispatcher）是 PyTorch 的一个内部组件，负责确定当你调用诸如 torch::add 之类的函数时，实际应该运行什么代码。这并非易事，因为 PyTorch 操作需要处理许多“层叠”在一起的跨领域问题。以下是它所处理事项的一些示例：

  • 根据输入张量的设备类型，在算子的 CPU 和 CUDA 实现之间进行切换。

  • 根据是否需要自动求导（autograd）处理，在算子的自动求导实现和后端实现之间进行切换。

  • 在需要自动混合精度（AMP）时应用自动转换（autocasting）。

  • 当算子在 vmap 调用下运行时，应用批处理（batching）规则。

  • 如果你正在追踪（tracing）模型以进行导出，则追踪操作的执行过程。

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

在本教程中，我们将描述如何构建自定义算子注册以使用分发器来组织各个组件。我们假设你熟悉如何 注册算子 以及如何编写 自定义自动求导函数。

定义模式和后端实现#

分发器背后的总体原则是，它将算子的实现划分为多个内核（kernel），每个内核实现特定 分发键（dispatch key） 的功能，例如 CPU 或 CUDA。分发器会在你调用算子时确定优先级最高的分发键（这是通过查看张量参数以及某些线程局部状态来完成的），并将控制权转移给该分发键对应的内核。最终效果是，当你调用一个算子时，我们首先执行自动求导内核，然后根据传入张量的设备类型重定向（redispatch）到后端内核。

让我们看看实现这一目标所涉及的各个部分。首先，我们必须为相关的算子定义模式（schema）。与简单的 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 张量！（在内部，我们将这些称为“全捕获”（catch-all）内核，因为它们捕获所有情况。）为了确保 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 分发键上的所有基本算子提供了实现。

对于不需要自动求导的算子#

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

在下一节中，我们将讨论如何为算子添加自动求导支持。但对于不需要自动求导支持的算子，应注册以下内核，以提高可用性并使你的算子表现得像 PyTorch 的内置算子。

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

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

原地（In-place）或视图（View）操作#

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

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

  1. 上述注册的 Autograd 或 ADInplaceOrView 盒装（boxed）内核依赖于其逻辑中的算子模式信息。如果你的操作原地修改输入或返回与输入之一别名的张量，确保你的模式正确反映这一点非常重要。有关如何注释模式的更多信息，请参见此处。

添加自动求导支持#

此时，我们已经拥有了 CPU 和 CUDA 实现的算子。如何为其添加自动求导支持？正如你可能猜到的，我们将注册一个自动求导内核（类似于 自定义自动求导函数 教程中所述）！然而，有一个转折：与 CPU 和 CUDA 内核不同，自动求导内核需要 重定向（redispatch）：它需要回调分发器以找到推理内核（例如 CPU 或 CUDA 实现）。

因此，在编写自动求导内核之前，让我们编写一个 分发函数（dispatching function），它调用分发器来为你的算子找到正确的内核。此函数构成了你的算子的公共 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 会将你送回此函数（因为最高优先级的分发键仍然是自动求导）。有了 (1)，自动求导被排除在考虑的分发键集合之外，我们将进入下一个处理程序，即 CPU 或 CUDA。

我们现在可以像注册 CPU/CUDA 函数一样注册此函数：

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

注意

在此示例中，我们将内核注册到 Autograd，它将其安装为所有后端的自动求导内核。你也可以通过使用相应的特定后端分发键（例如 AutogradCPU 或 AutogradCUDA）来为特定后端注册优化的内核。要更详细地探索这些和其他分发键选项，请查看 torch/_python_dispatcher.py 中提供的 PythonDispatcher 工具。

超越自动求导#

从某种意义上说，分发器并没有做太多的事情：它所做的只是实现了一个美化的 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、自动求导），而无需编写一个引用所有部分的单一中心化 if 语句。重要的是，第三方可以注册其他方面的额外实现，而无需修补算子的原始定义。我们将在 为新后端扩展分发器 中更多地讨论扩展分发器。

  2. 它支持比 CPU、CUDA 和自动求导更多的分发键。你可以在 c10/core/DispatchKey.h 中查看 PyTorch 当前实现的所有分发键的完整列表。这些分发键为算子实现了各种可选功能，如果你决定让你的自定义算子支持此功能，只需为相应的键注册一个内核即可。

  3. 分发器实现了对盒装（boxed）回退（fallback）函数的支持，这些函数可以实现一次并应用于系统中的所有算子。盒装回退可用于为分发键提供默认行为；如果你使用分发器来实现你的算子，你也选择了所有这些操作的回退。

以下是你可能需要为其定义算子的某些特定分发键。

Autocast#

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

这是一个假设的自定义矩阵乘法的 autocast 包装器及其注册：

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

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

请注意，与我们的自动求导内核一样，我们在重定向之前从分发中排除了 Autocast 键。

默认情况下，如果未提供 autocast 包装器，我们会直接回退到常规算子实现（不发生 autocasting）。（我们在本例中没有使用 myadd，因为点积加法不需要 autocasting，应该直接回退。）

什么时候应该注册 autocast 包装器？不幸的是，对于算子的首选精度没有硬性规定。你可以通过查看 类型转换列表 来了解某些原生算子的首选精度。一般建议：

  • 执行归约的操作可能应该在 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));
}

如果你的自定义算子是 启用自动求导的，你只需要为注册了自动求导包装器的名称编写并注册一个 autocast 包装器。例如，如果你想为自动求导部分中显示的 myadd 函数提供 autocast 包装器，你只需要：

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

没有单独的机制来使 backward 方法兼容 autocast。但是，你的自定义自动求导函数中定义的 backward 方法将以与 autocast 为 forward 方法设置的相同 dtype 运行，因此你应该选择一个适用于你的 forward 和 backward 方法的 <desired dtype>。

Batched#

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

Tracer#

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

在此页面上
PyTorch 库