目录
  • 教程 >
  • 在 C++ 中为新后端扩展调度器
快捷方式
advanced/extend_dispatcher

在 C++ 中为新后端扩展调度器

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

在本教程中，我们将逐步讲解扩展调度器以添加一个存在于 pytorch/pytorch 仓库之外的新设备并维护它以使其与原生 PyTorch 设备同步的所有必要步骤。我们假设你熟悉如何在 C++ 中注册调度算子以及如何编写自定义 autograd 函数。

注意

本教程涉及许多 PyTorch 内部组件，这些组件正在积极改进中。如果你决定按照本教程操作，请预期 API 会有所变化。我们将及时更新本教程，使其与最新 API 保持同步。

什么是新后端？

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

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

  • 如果你想提出新的算子，请向 PyTorch 提交功能请求/PR。

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

  • 如果你想为现有算子添加支持，但使用不同的张量布局/表示形式，例如稀疏和量化，这要求你的核函数在给定布局/表示限制的情况下以更高效的方式编写，请按照本教程为 PyTorch 添加一个树外 (out-of-tree) 后端。

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

为你的后端获取调度键

PyTorch 算子是用 C++ 实现的，并通过 Python 绑定在 Python 前端中提供。PyTorch 调度器将一个算子的实现划分为多个核函数，每个核函数都与一个特定的调度键关联。在 PyTorch 中支持新后端本质上意味着用 C++ 为每个 PyTorch 算子编写一个核函数，然后将它们注册到调度器中代表你定制后端的调度键。

调度键是你在调度器系统中的标识符。调度器查看输入张量携带的调度键，并据此调用正确的核函数。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 中，这些算子在其 accompanying comments 中的元数据里，dispatch 设置为 True 且 default 设置为 False。

  • 注册是可选的：后端扩展者可以跳过注册这些算子而不牺牲任何支持。然而，如果后端扩展者想要覆盖 PyTorch 提供的默认核函数，他们仍然可以将自己定制的核函数注册到其后端，调度器将仅为你的后端使用它。例如，PyTorch 当前的 max_pool2d 实现将 indices 作为前向输出的一部分返回，这在 torch_xla 中会产生开销，因此 torch_xla 转而为 max_pool2d 注册了自己的核函数。

    • 在 RegistrationDeclarations.h 中，这些算子在其 accompanying comments 中的元数据里，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 显示为 required registration。

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 注册的核函数支持以 unboxed 和 boxed 两种方式调用。换句话说，你的定制后端也可以像 CPU 或 CUDA 等内置 (in-tree) 后端一样与我们的 JIT 追踪/脚本化前端一起工作。你甚至可以在 JIT 图上为你的后端编写专门的优化 pass。但我们在此不讨论这一点，因为我们尚未最终确定 JIT 中的集成点，因此当前的后端支持将暂时侧重于 eager 前端。

针对原生 PyTorch 后端测试你的后端

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

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

使用通用设备类型测试框架但不随 PyTorch 一起发布的设备类型示例是 XLA。请参阅其对通用设备类型测试框架的扩展，其中包含阻止列表测试、阻止列表 dtypes 和覆盖测试精度的示例。

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

向后兼容性

目前 PyTorch 无法保证注册算子的向后兼容性。算子及其 schema 可能会根据需要添加/修改/删除。注册的核函数必须与 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 来扩展它。

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

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

未来工作

要使 PyTorch 的每个组件都能无缝地为树外后端提供扩展点，需要对 PyTorch 内部进行大量更改。以下是我们正在积极开展的一些可能在未来改善体验的工作项：

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

  • 改进 Math 核函数的覆盖率和更全面的测试，以确保 Math 核函数的行为与其他后端（如 CPU/CUDA）匹配。

  • 重构 RegistrationDeclarations.h，使其携带最少信息，并尽可能重用 PyTorch 的 codegen。

  • 支持后端回退 (fallback) 核函数，自动将输入转换为 CPU，并将结果转换回定制后端。即使你没有为每个算子编写核函数，这也将实现“完全”的算子覆盖。

保持联系

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

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

文档

获取 PyTorch 全面的开发者文档

查看文档

教程

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

查看教程

资源

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

查看资源

为了分析流量并优化您的体验，我们在本网站提供 cookie。点击或导航即表示您同意我们使用 cookie。作为本网站当前的维护者，Facebook 的 Cookies Policy 适用。了解更多，包括可用控件：Cookies Policy。