注册一个在C++中分派的操作符

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

警告

本教程自 PyTorch 2.4 起已弃用。请参阅 PyTorch 自定义操作，以获取使用自定义操作扩展 PyTorch 的最新指南。

调度器是PyTorch内部的一个组件，负责在调用类似torch::add的函数时确定实际运行的代码。这可能并非易事，因为PyTorch操作需要处理许多相互交叉的关注点，这些关注点“叠加”在一个又一个之上。以下是它处理的一些内容：

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

• 根据是否需要进行自动求导处理，在自动求导和后端实现之间切换操作符。

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

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

• 追踪操作的执行，如果您要为导出追踪模型。

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

在本教程中，我们将描述如何构建一个自定义操作符注册，以使用调度器来组织各种组件。我们将假设您已经熟悉如何 注册一个操作符 以及如何编写 自定义自动微分函数。

定义模式和后端实现

调度器的基本原理是将操作符的实现分割为多个内核，每个内核实现针对特定 dispatch key 的功能，例如 CPU、CUDA。调度器会确定在调用操作符时最高优先级的 dispatch key（这是通过查看张量参数以及一些线程本地状态来完成的），并将控制权转移到对应 dispatch key 的内核。最终效果是，当你调用一个操作符时，我们首先执行 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个代码块吗？这是XLA对PyTorch支持的实现方式：torch_xla库包含一个TORCH_LIBRARY_IMPL，它为所有基本操作符在XLA调度键上提供了实现。

对于不需要自动求导的操作符

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

在下一节，我们将讨论如何为一个操作符添加 autograd 支持。 但对于不需要 autograd 支持的操作符，应注册以下内核，以提高易用性，并使您的操作符行为如同 PyTorch 的内置操作符。

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

上面的代码注册了一个 Autograd 内核，该内核在前向传播时附加一个虚拟 NotImplemented 节点（保留输入的 require_grad 性质）。 在反向传播时，NotImplemented 节点会引发错误。这在调试大型模型时非常有用， 因为在之前的前向传播过程中很难准确地定位到 requires_grad 性质丢失的确切位置。

就地操作或视图操作

为确保正确性并获得最佳性能，如果你的操作会原地修改输入，或者返回一个与其中一个输入存在别名的张量，则应采取两个额外的步骤：

1. 注册一个 ADInplaceOrView 内核，除了上面的 Autograd 内核之外。这个内核负责必要的簿记工作，以确保就地操作或视图操作的正确性。需要注意的是，此 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实现的操作符。我们如何为它添加自动微分（autograd）支持呢？正如你可能猜到的那样，我们将注册一个自动微分内核（类似于自定义自动微分函数教程中所描述的内容）！然而，这里有一个关键点：与CPU和CUDA内核不同，自动微分内核需要重新调度：它需要回调调度器以调用推理内核，例如CPU或CUDA实现。

因此，在编写autograd内核之前，让我们先编写一个分发函数 来调用分发器以找到适用于您操作符的正确内核。 此函数构成了您操作符的公共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。这实际上会调用分派器，最终控制权将转移到适合此调用的内核。

有了 dispatch 函数，我们现在可以编写自动求导核函数了：

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，这会将其安装为所有后端的自动微分内核。你也可以通过使用对应的特定后端调度键来注册优化后的内核——例如，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、Autograd），而无需编写一个单一的、集中的 if 语句来引用所有这些组件。重要的是，第三方可以在不修改操作符原始定义的情况下注册其他方面的额外实现。我们将在 为新后端扩展调度器 中进一步讨论扩展调度器。

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

3. 调度器实现了对装箱后备函数的支持，这些函数可以一次实现并适用于系统中的所有运算符。装箱后备函数可用于为一个调度键提供默认行为；如果你使用调度器来实现你的运算符，你也将同时采用这些操作的所有后备函数。

以下是一些你可能需要为定义操作符而定义的特定分发键。

自动混合精度

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

以下是一个假设的自定义 matmul 的 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 转换为 float16，如果 tensor 是 CUDA 类型，并且 float32， 否则，它将保持 tensor 不变（参见 原生自动转换操作的兼容性策略）。 这确保了如果网络在任何混合的 float16 和 float32 CUDA 张量上调用 mymatmul， mymatmul 都会在 float16 中运行。同时，对 mymatmul 的调用，如果输入是非 CUDA 类型、整数类型或 float64， 则不受影响。建议在自己的自动转换包装器中使用 cached_cast 来遵循原生兼容性策略， 但这不是必须的。例如，如果你想强制所有输入类型都执行 float16， 你可以 return mymatmul(self.half(), other.half()); 而不是使用 cached_cast。

请注意，与我们的自动梯度内核一样，在重新分派之前，我们排除了 Autocast 键。

默认情况下，如果没有提供自动混合精度包装器， 我们将直接进入常规操作符实现（不会发生自动混合精度转换）。（在这个例子中我们没有使用myadd，因为逐点加法不需要自动混合精度转换，应该直接通过。）

何时应该注册自动混合精度（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 函数创建一个自动混合精度包装器，你所需要做的就是

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

在使反向方法与 autocast 兼容时，无需进行单独的处理。 然而，在自定义 autograd 函数中定义的反向方法将按照 autocast 为正向方法设置的数据类型运行，因此您应该选择一个 <desired dtype>，使其适合您的正向和反向方法。

批量

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

追踪器

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