advanced/dispatcher

在 Google Colab 中运行

Colab

下载 Notebook

Notebook

在 GitHub 上查看

GitHub

在 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。）

在本教程中，我们将描述如何构造自定义运算符注册以使用分派器来组织各种组件。我们将假设您熟悉如何 注册运算符 以及如何编写 自定义自动梯度函数。

定义模式和后端实现

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

让我们看看实现这一点的各种部分。首先，我们必须定义相关运算符的模式。与简单的 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 块吗？XLA 对 PyTorch 的支持就是这样实现的：torch_xla 库包含一个 TORCH_LIBRARY_IMPL，它在 XLA 分派键上为所有基本运算符提供实现。

对于不需要自动梯度的运算符

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

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

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

上面的行注册了一个 Autograd 内核，它在前向传播时附加一个虚拟的 NotImplemented 节点（保留输入的 require_grad 属性）。在后向传播时，NotImplemented 节点会引发错误。这对于在较大的模型中调试可能很有帮助，因为在以前很难精确地确定 requires_grad 属性在前向传播过程中丢失的位置。

原地或视图运算符

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

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 会将您送回此函数（因为优先级最高的_分派键_仍然是自动梯度）。使用 (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. 它是分布式的。您无需编写一个引用所有这些内容的集中式 if 语句，即可组合运算符的所有部分（CPU、CUDA、自动梯度）。重要的是，第三方可以注册额外的实现以处理其他方面，而无需修补运算符的原始定义。我们将在 在 C++ 中为新后端扩展分派器 中更详细地讨论扩展分派器。

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

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

以下是一些您可能需要为其定义运算符的特定_分派键_：

自动转换

Autocast _分派键_实现了对 自动混合精度 (AMP) 的支持。Autocast 包装器内核通常会在运行 op 之前将传入的 float16 或 float32 CUDA 张量转换为首选精度。例如，浮点 CUDA 张量上的 matmuls 和 convolutions 通常在 float16 中运行速度更快，内存占用更少，而不会损害收敛性。Autocast 包装器仅在 启用自动转换的上下文中 有效。

这是一个假设的自定义 matmul 的自动转换包装器及其注册：

// 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 保持不变（参见 原生支持自动转换的 op 的合格策略）。这确保了如果网络在任何 float16 和 float32 CUDA 张量的混合体上调用 mymatmul，mymatmul 将以 float16 运行。同时，对非 CUDA、整数类型或 float64 输入的 mymatmul 调用不受影响。建议使用 cached_cast 来遵循您自己的自动转换包装器中的原生合格策略，但这并非强制要求。例如，如果您想强制所有输入类型的 float16 执行，您可以 return mymatmul(self.half(), other.half()); 而不是使用 cached_cast。

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

默认情况下，如果未提供自动转换包装器，我们将直接回退到常规运算符实现（不进行自动转换）。（我们没有为此示例使用 myadd，因为逐点加法不需要自动转换，应该直接回退。）

何时应注册自动转换包装器？不幸的是，关于 op 的首选精度没有明确的规则。您可以通过查看 转换列表 来了解一些原生 op 的首选精度。一般指导：

• 执行归约的 op 应该在 float32 中执行，

• 任何在内部执行卷积或 gemm 的 op 应该在 float16 中执行，并且

• 具有多个浮点张量输入的其他 op 应将其标准化为通用精度（除非实现支持不同精度的输入）。

如果您的自定义 op 属于第三类，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));
}

如果您的自定义 op 已 启用自动梯度，您只需为主自动梯度包装器注册的相同名称编写并注册一个自动转换包装器。例如，如果您想要上面自动梯度部分所示的 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);
}

没有单独的技巧可以使后向方法兼容自动转换。但是，您自定义自动梯度函数中定义的后向方法将在与自动转换为前向方法设置的 dtype 相同的 dtype 中运行，因此您应该选择适合您的前向和后向方法的 <desired dtype>。

批处理

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

追踪器

Tracer _分派键_支持在运行 torch.jit.trace 时将运算符的调用记录到跟踪中。我们打算提供一个盒装回退，它将实现任意操作的跟踪，请参阅 issue #41478 以跟踪进度。