在 C++ 中为新的后端扩展 dispatcher¶

创建于：2021 年 2 月 1 日 | 最后更新：2024 年 9 月 23 日 | 最后验证：2024 年 11 月 05 日

在本教程中，我们将逐步介绍扩展 dispatcher 以添加一个新的设备的所有必要步骤，该设备位于 pytorch/pytorch 仓库之外，并对其进行维护以与原生 PyTorch 设备保持同步。在这里，我们假设您熟悉如何在 C++ 中注册 dispatched operator 以及如何编写自定义 autograd 函数。

注意

本教程涉及 PyTorch 内部的许多组件，这些组件正在积极改进，如果您决定遵循本教程，请预期 API 会发生更改。我们将使本教程与最新的 API 保持同步。

什么是新的后端？¶

向 PyTorch 添加新的后端需要后端扩展者进行大量的开发和维护。在添加新的后端之前，我们首先考虑一些常见的用例和推荐的解决方案

如果您有现有 PyTorch 运算符的新算法，请向 PyTorch 发送 PR。
如果您想提议一个新的运算符，请向 PyTorch 发送功能请求/PR。
如果您想添加对新设备/硬件（如 Google TPU 和自定义芯片）的支持，这通常需要使用特定于硬件的 API 来编写内核，请按照本教程操作，并将树外后端添加到 PyTorch。
如果您想添加对现有运算符的支持，但使用不同的张量布局/表示形式（如稀疏和量化），这会强制您的内核以更有效的方式编写，考虑到布局/表示形式的限制，请按照本教程操作，并将树外后端添加到 PyTorch。

在本教程中，我们将主要关注在下面添加新的树外设备。为不同的张量布局添加树外支持可能与设备共享许多共同步骤，但我们尚未看到此类集成的示例，因此可能需要 PyTorch 方面的额外工作来支持它。

获取后端的 dispatch key¶

PyTorch 运算符在 C++ 中实现，并通过 Python 绑定在 Python 前端中可用。PyTorch dispatcher 将运算符的实现划分为多个内核，每个内核都与特定的 dispatch key 相关联。在 PyTorch 中支持新的后端本质上意味着为每个 PyTorch 运算符在 C++ 中编写内核，然后将它们注册到表示您的自定义后端的 dispatcher 中的 dispatch key。

Dispatch key 是您在 dispatcher 系统中的标识符。Dispatcher 查看输入张量上携带的 dispatch key，并相应地调用正确的内核。PyTorch 提供了三个保留的 dispatch key（及其对应的 Autograd key），用于原型化树外后端扩展

PrivateUse1/AutogradPrivateUse1
PrivateUse2/AutogradPrivateUse2
PrivateUse3/AutogradPrivateUse3

您可以选择上面的任何 key 来原型化您的自定义后端。要在 PrivateUse1 后端上创建张量，您需要在 TensorImpl 构造函数中设置 dispatch key。

/* Example TensorImpl constructor */
TensorImpl(
    Storage&& storage,
    DispatchKeySet ks,
    const caffe2::TypeMeta data_type);

// To create a TensorImpl on PrivateUse1 backend, pass in the following ks to TensorImpl creation.
DispatchKeySet ks = c10::DispatchKeySet{c10::DispatchKey::PrivateUse1, c10::DispatchKey::AutogradPrivateUse1};

请注意，上面的 TensorImpl 类假定您的张量由 CPU/CUDA 等存储支持。我们还为没有存储的后端提供了 OpaqueTensorImpl。您可能需要调整/覆盖某些方法以适应您的自定义硬件。pytorch 仓库中的一个示例是Vulkan TensorImpl。

注意

原型完成后，如果您计划为您的后端扩展进行定期发布，请随时提交 PR 到 pytorch/pytorch 以保留您的后端的专用 dispatch key。

获取 PyTorch 运算符的完整列表¶

PyTorch 在生成的文件 build/aten/src/ATen/RegistrationDeclarations.h 中提供了可扩展 C++ 运算符的完整列表。此文件仅在从源代码构建 PyTorch 后可用。以下是该文件的代码片段

Tensor abs(const Tensor & self); // {"schema": "aten::abs(Tensor self) -> Tensor", "dispatch": "True", "default": "True"}
Tensor & abs_(Tensor & self); // {"schema": "aten::abs_(Tensor(a!) self) -> Tensor(a!)", "dispatch": "True", "default": "True"}
Tensor & abs_out(Tensor & out, const Tensor & self); // {"schema": "aten::abs.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)", "dispatch": "True", "default": "False"}
Tensor absolute(const Tensor & self); // {"schema": "aten::absolute(Tensor self) -> Tensor", "dispatch": "False", "default": "False"}
Tensor & absolute_(Tensor & self); // {"schema": "aten::absolute_(Tensor(a!) self) -> Tensor(a!)", "dispatch": "False", "default": "False"}
Tensor & absolute_out(Tensor & out, const Tensor & self); // {"schema": "aten::absolute.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)", "dispatch": "False", "default": "False"}
Tensor angle(const Tensor & self); // {"schema": "aten::angle(Tensor self) -> Tensor", "dispatch": "True", "default": "True"}
Tensor & angle_out(Tensor & out, const Tensor & self); // {"schema": "aten::angle.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)", "dispatch": "True", "default": "False"}
Tensor sgn(const Tensor & self); // {"schema": "aten::sgn(Tensor self) -> Tensor", "dispatch": "True", "default": "True"}

单个运算符关联了多个字段。让我们以 abs_out 为例进行分解

Tensor & abs_out(Tensor & out, const Tensor & self); 是运算符的 C++ 签名，您的 C++ 内核应与此签名完全匹配。
aten::abs.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) 是表示运算符的唯一模式，与 C++ 签名相比，它还包含别名和突变注释。这是 dispatcher 用于查找运算符的唯一标识符。
dispatch 和 default 是布尔字段，提供有关原生 PyTorch 内核可以执行的操作的信息，因此暗示后端扩展者是否需要实现内核。更多详细信息可以在为新后端注册内核中找到。

为新后端注册内核¶

要将您的内核注册到 PyTorch dispatcher，您可以使用在 C++ 中注册 Dispatched Operator中描述的 TORCH_LIBRARY_IMPL API

TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
  m.impl(<schema_my_op1>, &my_op1);
  m.impl(<schema_my_op2>, &my_op2);
  m.impl(<schema_my_op2_backward>, &my_op2_backward);
}

现在让我们放大，了解哪些运算符需要来自自定义后端的内核，以及内核内部到底是什么。

PyTorch 目前有 1600 多个运算符，并且仍在增长。后端扩展要跟上这个速度是不现实的。即使对于 CPU 或 CUDA 等原生后端，为每个新操作编写专用内核通常也需要大量工作。

幸运的是，一些原生 PyTorch 内核的编写方式使其可以分解为几个已知运算符的组合。换句话说，您只需要实现一组已知运算符（下面需要注册的操作），而不是所有 PyTorch 运算符。

PyTorch 运算符可以分为两类

需要注册的操作：这些操作的 PyTorch 原生实现是特定于后端的，因此需要为自定义后端提供内核。否则，在自定义后端上调用此类操作将出错。
- 在 RegistrationDeclarations.h 中，这些运算符的 dispatch 设置为 True 且 default 在其随附注释中找到的元数据中设置为 False。
注册是可选的：后端扩展者可以跳过注册到这些操作，而不会牺牲任何支持。但是，如果后端扩展者想要覆盖 PyTorch 提供的默认内核，他们仍然可以将他们的自定义内核注册到他们的后端，并且 dispatcher 将仅将它用于您的后端。例如，当前 PyTorch 的 max_pool2d 实现返回 indices 作为前向输出的一部分，这会在 torch_xla 中产生开销，因此 torch_xla 为 max_pool2d 注册了自己的内核。
- 在 RegistrationDeclarations.h 中，这些运算符的 dispatch 设置为 False 或 default 在其随附注释中找到的元数据中设置为 True。

新后端的 Autograd 支持¶

梯度公式大多是纯数学的，因此对于所有后端都是通用的。PyTorch 通常将内核注册到别名 dispatch key Autograd，这意味着它可以被所有后端使用。

对于这些运算符，您不必担心它们的导数公式，您只需在 RegistrationDeclarations.h 中编写运算符的前向定义，PyTorch 会自动为您处理反向传播。

Tensor my_op1(const Tensor& self, const Tensor& other) {
  // call your backend-specific APIs to implement my_op so that
  // it matches PyTorch's native behavior
}
TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
  m.impl(<schema_my_op1>, &my_op);
}

在某些情况下，PyTorch 反向传播内核实现也是特定于设备的，以便它们可以从每个后端榨取最大性能。对于这些运算符，您会看到 op_backward 在 RegistrationDeclarations.h 中显示为必需注册。

Tensor my_op2_backward(const Tensor& self, const Tensor& other) {
  // call your backend-specific APIs to implement my_op2_backward so that
  // it matches PyTorch's native behavior
}

// Note backward kernel is still registered to PrivateUse1 instead of AutogradPrivateUse1.
// PyTorch will wrap your backward kernel with proper autograd setup and then link to it in
// my_op2's AutogradPrivateUse1 kernel.
TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
  m.impl(<schema_my_op2>, &my_op2);
  m.impl(<schema_my_op2_backward>, &my_op2_backward);
}

在少数罕见情况下，PyTorch 针对某些运算符的梯度公式可能具有不适用于所有后端的假设。在这些情况下，后端扩展者可以选择通过从 torch::autograd::Function 注册内核到相应的 dispatch key（例如，如果您为您的后端使用 PrivateUse1，则为 AutogradPrivateUse1）来覆盖 PyTorch 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];
}

// Register the autograd kernel to AutogradPrivateUse1
TORCH_LIBRARY_IMPL(aten, AutogradPrivateUse1, m) {
  m.impl(<myadd_schema>, &myadd_autograd);
}

// Register the inference kernel to PrivateUse1
TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
  m.impl(<myadd_schema>, &myadd);
}

使用此技巧，您可以完全控制您的后端中 my_add 运算符的训练和推理行为。这是一个示例，位于 pytorch/xla 仓库中。

构建扩展¶

树外后端通过向 PyTorch 添加 C++ 扩展来支持。一旦您准备好内核和注册，您可以通过编写使用 setuptools 编译 C++ 代码的 setup.py 脚本来构建 C++ 扩展。以下是来自 pytorch/xla 仓库的简化示例

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension

setup(
    name='torch_xla',
    ext_modules=[
        CppExtension(
            '_XLAC',
            torch_xla_sources,
            include_dirs=include_dirs,
            extra_compile_args=extra_compile_args,
            library_dirs=library_dirs,
            extra_link_args=extra_link_args + \
                [make_relative_rpath('torch_xla/lib')],
        ),
    ],
    cmdclass={
        'build_ext': Build,  # Build is a derived class of BuildExtension
    }
    # more configs...
)

有关更多详细信息，请参阅我们的 C++ 扩展教程。

自定义运算符支持¶

您的新后端应与在 python 中扩展的自定义运算符无缝协作，而无需编写任何新的内核，只要自定义运算符由现有的 PyTorch 运算符（您的后端已支持）组成。

对于在 C++ 中扩展的自定义运算符，它们通常带有后端特定的 C++ 内核实现，例如 torchvsion 中的 nms 内核，以及自定义的 Python API，例如 torch.ops.torchvision.nms。为了支持这些运算符，后端扩展者将需要为您的后端编写 C++ 内核，并将其正确注册到 dispatcher 中相应的命名空间，类似于支持 PyTorch 原生运算符。或者，您也可以在您的扩展中添加自定义 API，例如 torch_xla.core.functions.nms 用于这些临时请求。

JIT 支持¶

正如我们在在 C++ 中注册 Dispatched Operator中提到的，通过 m.impl() API 注册的内核支持以非装箱和装箱两种方式调用。换句话说，您的自定义后端也可以像 CPU 或 CUDA 等树内后端一样，与我们的 JIT 跟踪/脚本前端一起工作。您还可以为 JIT 图上的后端编写专门的优化 pass。但我们在这里不会讨论它，因为我们尚未最终确定 JIT 中的集成点，因此当前的后端支持将侧重于 eager 前端。

针对原生 PyTorch 后端测试您的后端¶

PyTorch 允许使用其通用设备类型测试框架在多种设备类型上运行测试。您可以找到有关测试如何使用它的详细信息，以及有关如何添加新设备类型的信息。添加后，使用通用设备类型测试框架的 PyTorch 测试也将使用您的设备类型运行。有关如何实例化测试的示例，请参阅此 Wiki 页面。

使用您的设备类型运行 PyTorch 现有的测试套件对于确保正确性非常重要，但并非每个设备类型都支持所有 PyTorch 功能。通用设备类型测试框架允许进行相当大的自定义，以便设备类型可以选择要运行的测试、它们支持的 dtype，甚至是在比较张量是否相等时要使用的精度。

XLA 是一个使用通用设备类型测试框架且未随 PyTorch 提供的设备类型示例。请参阅其通用设备类型测试框架的扩展，其中包含阻止列出测试、阻止列出 dtype 和覆盖测试精度的示例。

通用设备类型测试框架正在积极开发中。要请求功能，请在 PyTorch 的 Github 上提交 issue。

向后兼容性¶

目前，PyTorch 无法保证已注册运算符的向后兼容性。运算符及其模式可能会根据需要添加/修改/删除。注册的内核必须与 PyTorch 版本完全相同。如果 PyTorch 为运算符添加更多参数（甚至带有默认值），您的旧注册将无法工作，直到更新它以匹配 PyTorch 的新签名。

因此，我们强烈建议树外后端扩展者仅与主要的 PyTorch 版本同步，以最大程度地减少开发中的中断。PyTorch 采用季度发布节奏。后端扩展者应加入 pytorch.slack.com 上的 #announcement 频道，以获取有关发布的最新更新。

已知问题和附加说明¶

并非所有测试套件都是设备通用的。可以通过在 PyTorch 代码库中搜索 instantiate_device_type_tests 来找到可扩展的测试类，例如 TestTorchDeviceType、 TestViewOps、 TestTensorDeviceOps、 TestTypePromotion 等。
C++ 中没有扩展点用于序列化自定义后端上的 python 张量对象。目前，您只能通过修改 PyTorch 张量 __reduce_ex__ 方法或在树外仓库中进行 monkey patching 来扩展它。
如果您的后端不允许直接内存访问，您应该格外注意支持 view 操作，因为它们应该共享存储。对 view 张量的更改需要传播到其基础张量，反之亦然。
如果您的后端不适用于原生 PyTorch 优化器，则 C++ 中没有优化器的扩展点，例如，需要像 torch-xla 一样携带要在反向传播中更新的状态。此类用例目前只能通过添加自定义 API 或在树外仓库中进行 monkey patching 来完成。

未来工作¶

使 PyTorch 中的每个组件对于树外后端无缝可扩展，需要对 PyTorch 内部结构进行大量更改。以下是我们正在积极研究的几个项目，这些项目可能会在未来改善体验

提高通用测试框架的测试覆盖率。
提高 Math 内核覆盖率和更全面的测试，以确保 Math 内核行为与其他后端（如 CPU/CUDA）匹配。
重构 RegistrationDeclarations.h 以携带最少的信息，并尽可能重用 PyTorch 的代码生成。
支持后端回退内核以自动将输入转换为 CPU，并将结果转换回自定义后端。即使您没有为每个运算符编写内核，这也将允许“完整”运算符覆盖。

保持联系¶

请使用PyTorch 开发者讨论来提出问题和进行讨论。如果您有任何功能请求或错误报告，请在 github 上提交 issue。

如果您有兴趣帮助完成上述任何未来工作项目（例如，为 C++ 中的 PyTorch 运算符添加更多 Math 内核），请通过 Github 或 Slack 与我们联系！