目录
  • 教程 >
  • 在 C++ 中注册分发操作符
快捷方式
advanced/dispatcher

在 C++ 中注册分发操作符

警告

本教程自 PyTorch 2.4 起已弃用。请参阅 PyTorch 自定义操作符，了解有关使用自定义操作符扩展 PyTorch 的最新指南。

调度程序是 PyTorch 的内部组件，负责确定在您调用 torch::add 之类的函数时实际运行的代码。这可能非同小可，因为 PyTorch 操作需要处理许多相互“叠加”的跨领域问题。以下是一些它处理的事项的示例

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

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

  • 在需要时为自动混合精度应用自动转换。

  • 在操作符在 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 张量！（在内部，我们将这些称为“通用”内核，因为它们捕获所有情况。）为了确保 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 调度键上的所有基本运算符提供实现。

对于不需要自动微分的运算符

注意：本节仅适用于 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 盒式内核依赖于其 logi 中的运算符模式信息。如果您的运算符就地修改输入或返回一个与其中一个输入别名的张量，则务必确保您的模式正确地反映了这一点。有关如何注释模式的更多信息，请参阅 这里。

添加自动微分支持

此时，我们有一个运算符，它同时具有 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. 它是分散的。您可以组装运算符的所有部分（CPU、CUDA、自动微分），而无需编写引用所有部分的单个集中式 if 语句。重要的是，第三方可以为其他方面注册额外的实现，而无需修补运算符的原始定义。我们将在 为新后端扩展调度器 中更多地讨论扩展调度器。

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

  3. 调度器实现了对打包的回退函数的支持，这些函数可以实现一次，并适用于系统中的所有算子。打包的回退可用于为调度键提供默认行为；如果您使用调度器来实现您的算子，您也选择加入这些操作的所有回退。

以下是您可能需要为其定义算子的特定调度键。

自动转换

Autocast 调度键实现了对 自动混合精度 (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);
}

cached_cast(kHalf, tensor) 将 tensor 转换为 float16，如果 tensor 是 CUDA 并且 float32，否则它保持 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);
}

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

批量

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

跟踪器

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

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获得面向初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得解答

查看资源

为了分析流量并优化您的体验，我们在本网站上提供 Cookie。点击或导航，即表示您同意允许我们使用 Cookie。作为本网站的当前维护者，Facebook 的 Cookie 政策适用。了解更多信息，包括有关可用控制的信息：Cookie 政策。