Table of Contents
Shortcuts
advanced/dispatcher

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

Created On: Jul 22, 2020 | Last Updated: Jul 22, 2024 | Last Verified: Nov 05, 2024

警告

自 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 张量！（内部我们将这些称为 “全捕捉” 内核，因为它们捕捉了所有情况。）为了确保 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 节点（保留输入的 requires_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)”传递给它，因为分派函数的类型与注册到分派器的底层内核类型相同。

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

  • 在第二行中，我们只需使用传递到分派函数中的所有参数来“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)

自动类型转换 (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”为CUDA且“float32”时，将其转换为“float16”；否则保持“tensor”不变。这确保了如果网络调用了“mymatmul”且任何张量为“float16”和“float32”的CUDA张量的混合， “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>``。

批处理模式 (Batched)

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

跟踪 (Tracer)

跟踪分派键实现了支持在运行“torch.jit.trace”时将运算符调用记录到跟踪中。我们计划提供一个箱式后备以实现对任意操作的跟踪，请参阅Issue #41478以跟踪进展。

文档

访问 PyTorch 的详细开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源

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