Table of Contents
Shortcuts
advanced/dispatcher

在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。调度器在你调用操作符时确定最高优先级的调度键（这是通过查看张量参数以及一些线程本地状态来完成的），并将控制权转移到该调度键的内核。最终的效果是，当你调用一个操作符时，我们首先执行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。

在下一节中，我们将讨论如何为操作符添加自动求导支持。 但对于不需要自动求导支持的操作，应注册以下内核以提高可用性，并使您的操作符表现得像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());
}

  1. 上面注册的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)，因为调度函数的类型与注册到调度器的基础内核的类型相同。

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

  • 在第二行，我们简单地调用操作符处理程序，并传入所有传递给调度函数的参数。这将实际调用调度器，最终控制权将转移到适合此调用的任何内核。

有了dispatch函数，我们现在可以编写autograd内核：

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. 调度器实现了对盒装回退函数的支持，这些函数可以一次性实现并适用于系统中的所有操作符。盒装回退可以用于为调度键提供默认行为；如果您使用调度器来实现您的操作符，您也将选择加入所有这些操作的回退。

以下是一些您可能需要为操作符定义的特定调度键。

自动转换

Autocast调度键实现了对自动混合精度（AMP）的支持。一个autocast包装器内核通常在运行操作之前将传入的float16或float32 CUDA张量转换为某种首选精度。例如，浮点CUDA张量上的矩阵乘法和卷积通常在float16中运行得更快，并且使用更少的内存，而不会影响收敛性。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) casts tensor to float16 if tensor is CUDA and float32, otherwise, it leaves tensor unchanged (c.f. the eligibility policy for natively autocasted ops). This ensures if the network calls mymatmul on any mixture of float16 and float32 CUDA tensors, mymatmul runs in float16. Meanwhile, calls to mymatmul with non-CUDA, integer-type, or float64 inputs are unaffected. Using cached_cast to follow the native eligibility policy in your own autocast wrapper is recommended, but not required. For example, if you wanted to force float16 execution for all input types, you could return mymatmul(self.half(), other.half()); instead of using cached_cast.

请注意，就像我们的自动梯度内核一样，我们在重新分发之前从分发中排除了Autocast键。

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

何时应注册自动类型转换包装器？不幸的是，对于操作的首选精度并没有一成不变的规则。你可以通过查看类型转换列表来了解一些原生操作的首选精度。一般指导原则如下：

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

没有单独的体操可以使反向方法与自动转换兼容。 但是，在自定义的自动梯度函数中定义的反向方法将以与自动转换为前向方法设置的相同数据类型运行，因此您应该选择一个 dtype>，该数据类型适用于您的前向和反向方法。

批量处理

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

追踪器

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

为了分析流量并优化您的体验，我们在本网站上使用cookies。通过点击或浏览，您同意我们使用cookies。作为本网站的当前维护者，Facebook的Cookies政策适用。了解更多信息，包括可用的控制选项：Cookies Policy。

优云智算