评价此页
★ ★ ★ ★ ★
advanced/extend_dispatcher

在 C++ 中为新后端扩展分发器#

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

在本教程中，我们将介绍扩展分发器 (Dispatcher) 以添加一个新的设备（位于 pytorch/pytorch 代码库之外）所需的所有步骤，并介绍如何对其进行维护以保持与原生 PyTorch 设备同步。在此，我们假设您已经熟悉如何 在 C++ 中注册分发算子 以及如何编写 自定义自动求导函数 (Autograd function)。

注意

本教程涉及 PyTorch 内部的许多组件，这些组件正在被积极改进，如果您决定遵循本教程，请预料到 API 可能会发生变化。我们将根据最新的 API 持续更新本教程。

什么是新后端？#

为 PyTorch 添加新后端需要后端扩展者进行大量的开发和维护工作。在添加新后端之前，让我们先考虑一些常见的使用场景及其推荐的解决方案：

  • 如果您有针对现有 PyTorch 算子的新算法，请向 PyTorch 提交 PR。

  • 如果您想提议一个新的算子，请向 PyTorch 提交功能请求/PR。

  • 如果您想添加对新设备/硬件（如 Google TPU 和定制芯片）的支持，这通常需要使用特定于硬件的 API 来编写内核，请按照本教程添加一个树外 (out-of-tree) 后端到 PyTorch。

  • 如果您想添加对现有算子的支持，但采用不同的张量布局/表示（如稀疏和量化），这要求您的内核必须以针对布局/表示限制更有效的方式编写，请按照本教程添加一个树外后端到 PyTorch。

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

为您的后端获取分发键 (Dispatch Key)#

PyTorch 算子使用 C++ 实现，并通过 Python 绑定在 Python 前端中提供。PyTorch 分发器将算子的实现划分为多个内核，每个内核与一个特定的分发键相关联。在 PyTorch 中支持新后端，本质上就是为每个 PyTorch 算子编写一个 C++ 内核，然后将它们注册到分发器中代表您自定义后端的 dispatch key 上。

分发键是您在分发器系统中的标识符。分发器会查看输入张量上携带的分发键，并相应地调用正确的内核。PyTorch 提供了三个预留的分发键（及其对应的 Autograd 键），用于构建树外后端扩展原型：

  • PrivateUse1/AutogradPrivateUse1

  • PrivateUse2/AutogradPrivateUse2

  • PrivateUse3/AutogradPrivateUse3

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

/* 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。

注意

原型完成后，如果您计划为您的后端扩展进行定期发布，请随时向 pytorch/pytorch 提交 PR，为您的后端预留一个专用的分发键。

获取 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!) 是表示该算子的唯一模式 (schema)，与 C++ 签名相比，它还包含别名和变异注释。这是分发器用于查找算子的唯一标识符。

  • dispatch 和 default 是布尔字段，提供了有关原生 PyTorch 内核能做什么的信息，从而暗示了后端扩展者是否需要实现该内核。更多详细信息可以在 为新后端注册内核 中找到。

为新后端注册内核#

要将您的内核注册到 PyTorch 分发器，您可以使用 在 C++ 中注册分发算子 中描述的 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 提供的默认内核，他们仍然可以将自定义内核注册到他们的后端，分发器将仅为您的后端使用它。例如，当前 PyTorch 的 max_pool2d 实现返回 indices 作为前向输出的一部分，这在 torch_xla 中会造成开销，因此 torch_xla 为 max_pool2d 注册了自己的内核。

    • 在 RegistrationDeclarations.h 中，这些算子在其伴随注释的元数据中，dispatch 设置为 False 或 default 设置为 True。

为新后端提供 Autograd 支持#

梯度公式大多纯粹是数学性的，因此对所有后端都是通用的。PyTorch 通常将一个内核注册为别名分发键 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 的反向传播内核实现也是特定于设备的，以便它们可以从每个后端挤出最大的性能。对于这些算子，您也会在 RegistrationDeclarations.h 中看到 op_backward 显示为 *必须注册*。

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 注册一个内核到相应的分发键（例如，如果您为后端使用 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++ 扩展教程。

自定义算子支持#

只要自定义算子是由现有的 PyTorch 算子（这些算子已由您的后端支持）组成的，您的新后端就应该能够无缝地处理 在 Python 中扩展的自定义算子，而无需编写任何新内核。

对于 在 C++ 中扩展的自定义算子，它们通常带有 特定于后端的 C++ 内核实现（例如 torchvision 中的 nms 内核） 以及 自定义 Python API（例如 torch.ops.torchvision.nms）。为了支持这些算子，后端扩展者需要为您的后端编写一个 C++ 内核，并将其适当地注册到分发器中相应的命名空间，类似于支持 PyTorch 原生算子的方式。或者，您也可以在您的扩展中添加一个自定义 API，例如 torch_xla.core.functions.nms 来处理这些临时需求。

JIT 支持#

正如我们在 在 C++ 中注册分发算子 中提到的，通过 m.impl() API 注册的内核支持以装箱 (boxed) 和拆箱 (unboxed) 方式调用。换句话说，您的自定义后端也可以像 CPU 或 CUDA 等树内后端一样，与我们的 JIT 跟踪/脚本前端协同工作。您可能还可以为您的后端在 JIT 图上编写专门的优化通道。但我们不会在这里讨论它，因为我们还没有最终确定 JIT 中的集成点，所以目前的后端支持暂时将集中在 eager 前端。

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

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

使用您的设备类型运行 PyTorch 现有的测试套件对于确保正确性非常重要，但并非所有设备类型都支持所有 PyTorch 功能。通用设备类型测试框架允许进行相当多的定制，以便设备类型可以选择运行哪些测试、它们支持哪些数据类型 (dtypes)，甚至可以在比较张量相等性时使用哪些精度。

一个使用通用设备类型测试框架且不随 PyTorch 一起分发的示例设备类型是 XLA。请参阅 它对通用设备类型测试框架的扩展，其中包含屏蔽测试、屏蔽数据类型和覆盖测试精度的示例。

通用设备类型测试框架正在积极开发中。如需请求功能，请在 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 Tensor __reduce_ex__ 方法 或在树外仓库中进行 monkey patching 来扩展它。

  • 如果您的后端不允许直接内存访问，您应该额外注意支持 view 算子，因为它们应该共享存储。对视图张量的更改需要传播到其基础张量，反之亦然。

  • 如果您的后端不能与原生 PyTorch 优化器一起工作（例如，需要在反向传播中携带需要更新的状态，如 torch-xla），则 C++ 中没有用于优化器的扩展点。此类用例目前只能通过添加自定义 API 或在树外仓库中进行 monkey patching 来完成。

未来工作#

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

  • 提高通用测试框架的测试覆盖率。

  • 提高 Math 内核覆盖率，并进行更全面的测试，以确保 Math 内核行为与 CPU/CUDA 等其他后端匹配。

  • 重构 RegistrationDeclarations.h 以携带最少信息，并尽可能多地重用 PyTorch 的代码生成。

  • 支持后端回退内核，将输入自动转换为 CPU，并将结果转换回自定义后端。这将允许“完全”算子覆盖，即使您没有为每个算子编写内核。

保持联系#

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

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

在此页面上
PyTorch 库