在 C++ 中注册 Dispatched 运算符¶

创建时间： 2020年7月22日 |上次更新时间： 2024 年 7 月 22 日 |上次验证： Nov 05， 2024

警告

本教程从 PyTorch 2.4 开始已弃用。有关使用自定义运算符扩展 PyTorch 的最新指南，请参阅 PyTorch 自定义运算符。

dispatcher 是 PyTorch 的一个内部组件，负责弄清楚在调用类似 .这可能非常重要，因为 PyTorch 操作需要处理许多 “分层” 在一个之上的横切关注点另一个。以下是它处理的一些事情的示例：torch::add

在运算符的 CPU 和 CUDA 实现之间切换，具体取决于在输入张量的设备上。
在 Operator 的 autograd 和 backend 实现之间切换，取决于是否需要 autograd 处理。
必要时应用自动转换以实现自动混合精度。
在调用下运行运算符时应用批处理规则。vmap
跟踪操作的执行（如果要跟踪要导出的模型）。

如果在自定义运算符代码中，您会发现自己手动编写 if 语句来处理这些情况，调度程序 API 可以帮助组织您的代码。（相反，如果您的自定义运算符非常简单并且仅用于 CPU 推理，则您可能不需要使用 Dispatcher，只需使用基本 API。

在本教程中，我们将介绍如何构建自定义运算符注册以使用 Dispatcher 组织各种组件。我们将假设您熟悉如何注册运算符以及如何编写自定义 autograd 函数。

定义 schema 和后端实现¶

调度程序背后的一般原则是，它将将 Operator 实现到多个内核中，每个内核都实现特定 dispatch key 的功能，例如 CPU、CUDA。调度程序确定当时优先级最高的 dispatch key 是什么您调用一个运算符（这是通过将两个 Tensor 参数视为以及一些线程本地状态），并将控制权转移给内核 dispatch 键。最终效果是，当你调用运算符时，我们首先执行 Autograd 内核，然后我们重新调度到后端内核取决于传入的张量的设备类型。

让我们来看看制作它所涉及的各个部分发生。首先，我们必须为相关运算符定义 schema。与简单的 pybind11 风格的运算符注册不同，我们实际上并没有此时提供我们 Operator 的实现;我们只是提供指定运算符类型签名的架构字符串我们所有其他内核都将遵守：

TORCH_LIBRARY(myops, m) {
  m.def("myadd(Tensor self, Tensor other) -> Tensor");
}

接下来，我们需要实际提供这个 Operator 的一些实现。具体来说，下面是一个非常简单的 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;
}

我们希望将此函数注册为 . 但是，注册它的简单方法（）会注册内核以在所有情况下运行，即使张量不是 CPU 张肌！（在内部，我们将这些称为 “catch-all” 内核，因为它们捕获所有 Cases。要确保仅对 CPU 张量，我们可以使用宏：myops::myadddef("myadd", myadd_cpu)myadd_cpuTORCH_LIBRARY_IMPL

TORCH_LIBRARY_IMPL(myops, CPU, m) {
  m.impl("myadd", myadd_cpu);
}

这让我们为特定的 dispatch key（在本例中为 CPU）。每次调用 to 都会将一个 CPU 内核与相应的运算符（我们之前定义）。如果我们还有一个 CUDA 实现，我们可以在一个单独的区块中注册它：TORCH_LIBRARY_IMPLimplTORCH_LIBRARYmyadd_cudaTORCH_LIBRARY_IMPL

TORCH_LIBRARY_IMPL(myops, CUDA, m) {
  m.impl("myadd", myadd_cuda);
}

这些注册可以跨文件拆分，甚至可以跨库边界拆分;所以例如，您可以编译这两个块转换为单独的动态库中。一般说，您的注册结构将如下所示：TORCH_LIBRARY_IMPLmyops_cpumyops_cuda

一个列出命名空间中每个自定义运算符的 single 在一个集中的地方。TORCH_LIBRARY
一个 per dispatch 键，用于注册该密钥（例如 CPU 或 CUDA）。如果您愿意，您可以进一步将块细分为每个运算符的一个块。这很方便如果你有一个单独的 File Per Operator 实现，但不想在标头中公开运算符;你可以把注册放在 cpp 文件。TORCH_LIBRARY_IMPLTORCH_LIBRARY_IMPL

注意

您知道吗，您也可以为现有的 PyTorch 中的核心操作符？这就是 PyTorch 的 XLA 支持方式已实现：该库包含一个为 XLA 分派上的所有基本运算符提供实现钥匙。TORCH_LIBRARY_IMPLtorch_xlaTORCH_LIBRARY_IMPL

对于不需要 autograd 的 Operator¶

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

在下一节中，我们将讨论如何为 Operator 添加 autograd 支持。但是对于不需要 autograd 支持的 ops，应该使用以下内核 registered 提高可用性，并使您的 op 行为类似于 PyTorch 的内置运维。

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

以上几行注册了一个内核，该内核在 forward 时附加了一个虚拟节点（保留 inputs 的 -ness）。向后时，节点会引发错误。这可能会有所帮助用于在以前难以确定的大型模型中进行调试正是在前向传递期间 -ness 丢失的位置。AutogradNotImplementedrequire_gradNotImplementedrequires_grad

就地或查看操作¶

为了确保正确性和最佳性能，如果您的运算会更改 input 就地或返回一个与其中一个输入别名的张量，两个额外的应采取的步骤：

除了内核之外，还注册一个内核以上。这个内核处理必要的记账以确保正确性的就地或视图操作。请务必注意，此 ADInplaceOrView kernel 只能与一起使用。ADInplaceOrViewAutogradautogradNotImplementedFallback

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

上面注册的 or 盒装内核依赖其 Logi 中的 Operator 架构信息。如果你的操作改变一个输入就地或返回一个张量，该张量与它很重要的输入之一进行别名确保您的架构正确反映了这一点。有关如何注释架构的更多信息，请参阅此处。AutogradADInplaceOrView

添加 autograd 支持¶

此时，我们有一个同时具有 CPU 和 CUDA 实现的运算符。如何我们可以为其添加 Autograd 支持吗？正如您可能猜到的那样，我们将注册一个 autograd 内核（类似于自定义 autograd 函数教程中描述的内容）！但是，有一个转折点：与 CPU 和 CUDA 内核不同，autograd 内核需要重新调度：它需要回调调度程序以到达推理内核，例如 CPU 或 CUDA 实现。

因此，在我们编写 autograd 内核之前，让我们编写一个调度函数，该函数调用 Dispatcher 以为您的 Operator 找到合适的内核。此函数构成了运算符的公共 C++ API，实际上是所有 PyTorch 的 C++ API 中的 tensor 函数都以相同的方式调用 Dispatcher 在引擎盖下。调度函数如下所示：

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

让我们来分析一下：

在第一行中，我们从调度程序中查找类型化运算符句柄对应于我们要 dispatch 到的 operator。采用两个参数：（namespace qualified） name 的运算符，以及运算符的重载名称（通常只是空字符串）。将动态类型化的句柄转换为静态类型的句柄（执行运行时测试以确保您已给出正确的 C++ 类型），这样我们就可以对它进行普通的 C++ 调用。我们传递它，因为调度函数的类型是与注册到 Dispatcher 的基础内核的类型相同。findSchemaOrThrowtypeddecltype(myadd)

为了提高性能，此计算在 static 变量中完成，因此我们只需要做一次（慢速）查找。如果您输入了运算符，则此查找将在您第一次调用此功能。
在第二行中，我们简单地使用所有传递给调度函数的参数。这实际上会调用 Dispatcher 和 In the End Control 将转移到任何内核适用于此调用。call

有了 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];
}

autograd 函数是正常使用，不同之处在于，不是直接在中编写实现，我们：torch::autograd::Functionforward()

使用 RAII 关闭 autograd 处理 guard，然后at::AutoNonVariableTypeMode
调用 dispatch 函数以回调到 Dispatcher 中。myadd

如果没有（1），你的调用将无限循环（和堆栈溢出），因为会将你送回这个函数（作为最高优先级的 dispatch key 仍将为 autograd。其中（1）、 autograd 被排除在正在考虑的 dispatch 键集中，并且我们将转到下一个处理程序，这将是 CPU 和 CUDA。myadd

我们现在可以像注册 CPU/CUDA 一样注册这个函数功能：

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

注意

在此示例中，我们将内核注册到，后者将其安装为 autograd 内核。您还可以为特定的 backends 使用相应的特定于后端的 dispatch key - 例如或 .浏览这些和其他调度键选项更详细地查看 torch/_python_dispatcher.py 中提供的工具。AutogradAutogradCPUAutogradCUDAPythonDispatcher

超越 autograd¶

从某种意义上说，调度程序并没有做那么多事情：它所做的只是实现一个美化的 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());
    }
  }
  ...
}

那么为什么要使用 Dispatcher 呢？有几个原因：

它是去中心化的。您可以组装运算符的所有部分（CPU、CUDA、Autograd），而无需编写单个集中式 if 语句。重要的是，第三方可以为其他方面注册额外的实现，而不必修补运算符的原始定义。我们将详细讨论如何扩展 Dispatcher 中的扩展 Dispatcher 以用于新后端。
它支持的调度密钥比 CPU、CUDA 和 Autograd 多。您可以查看当前已实现的 Dispatch Key 的完整列表在 PyTorch 中。这些 dispatch key 为操作员实现各种可选功能，并且如果你决定您希望自定义 Operator 支持此功能，您只需为适当的密钥注册一个内核即可。c10/core/DispatchKey.h
调度程序实现了对装箱回退函数的支持，该函数是可以实现一次并应用于所有运算符的函数在系统中。盒装回退可用于提供默认行为用于调度密钥;如果您使用 Dispatcher 实现 Operator，则您还可以选择所有这些操作的回退。

以下是一些特定的 dispatch 键，您可能需要定义运算符为。

自动投射¶

Autocast 调度键实现对自动混合精度（AMP）的支持。自动转换包装器内核通常强制转换传入张量或 CUDA 张量复制到某个首选精度。例如，浮点 CUDA 张量上的 matmuls 和卷积通常运行得更快并在不影响收敛的情况下使用更少的内存。 Autocast wrappers 仅在启用了 autocast 的上下文中有效。float16float32float16

下面是一个假设的自定义 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)转换为 if 为 CUDA 和，否则，它将保持不变（参见 Natively Autocasted Ops 的资格策略）。这确保了如果网络调用任何 CUDA 张量的混合，则在中运行。同时，使用非 CUDA、整数类型或输入的调用不受影响。用于在您自己的 autocast 包装器中遵循本机资格策略是推荐的，但不是必需的。例如，如果要强制执行所有输入类型，您可以代替使用 .tensorfloat16tensorfloat32tensormymatmulfloat16float32mymatmulfloat16mymatmulfloat64cached_castfloat16return mymatmul(self.half(), other.half());cached_cast

请注意，就像我们的 autograd 内核一样，我们从 dispatch 之前。Autocast

默认情况下，如果未提供 autocast wrapper，则我们直接跳转到常规的 Operator 实现（否 autocast 发生）。（我们没有在这个例子中使用，因为 pointwise add 不需要自动转换，应该会失败。myadd

什么时候应该注册 autocast wrapper？不幸的是，没有 OP 首选精度的 cut-and-dry 规则。您可以通过查看演员表来了解一些 Native Ops 的首选精度。一般指导：

执行 reduction 的操作可能应该在中执行，float32
任何在后台执行卷积或 gemm 的操作都应该可能在中执行，而float16
其他具有多个浮点张量输入的运算应标准化它们设置为通用精度（除非 implementation 支持具有不同精度的 Importing）。

如果您的自定义操作属于第三类，则模板有助于找出输入张量中存在的最宽浮点类型，即执行类型的最安全选择：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));
}

如果您的自定义运算已启用 autograd，则只需写入和注册 autograd 包装器注册到的同名 autocast 包装器。例如，如果您希望为所示的函数提供自动转换包装器在 autograd 部分，您只需要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);
}

没有单独的体操来使向后方法自动转换兼容。但是，自定义 autograd 函数中定义的 backward 方法将在同一 dtype 设置为 autocast 设置，因此您应该选择适合 forward 和 backward 方法的 AUTOCAST 方法。<desired dtype>

批¶

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

示踪¶

Tracer dispatch key 实现了对记录运算符调用的支持到跟踪中。我们打算提供将实现对任意操作的跟踪的 boxed fallback，请参阅 issue #41478 进行跟踪进展。torch.jit.trace

在 C++ 中注册 Dispatched 运算符¶

定义 schema 和后端实现¶

对于不需要 autograd 的 Operator¶

就地或查看操作¶

添加 autograd 支持¶

超越 autograd¶

自动投射¶

批¶

示踪¶

文档

教程

资源

在 C++ 中注册 Dispatched 运算符¶

定义 schema 和后端实现¶

对于不需要 autograd 的 Operator¶

就地或查看操作¶

添加 autograd 支持¶

超越 autograd¶

自动投射¶

批¶

示 踪¶

文档

教程

资源

示踪¶