目录
  • 教程 >
  • 在 C++中扩展用于新后端的调度器
快捷键
advanced/扩展调度器

在 C++中扩展用于新后端的调度器 ¶

创建于：2025 年 4 月 1 日 | 最后更新：2025 年 4 月 1 日 | 最后验证：2024 年 11 月 5 日

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

备注

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

什么是新的后端？

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

  • 如果您为现有的 PyTorch 运算符有新的算法，请向 PyTorch 提交一个 PR。

  • 如果您想提议一个新的操作符，请向 PyTorch 发送功能请求/PR。

  • 如果您想添加对新的设备/硬件（如 Google TPU 和定制芯片）的支持，这通常需要使用特定于硬件的 API 编写内核，请遵循本教程并为 PyTorch 添加一个树外后端。

  • 如果您想添加对现有操作符的支持，但使用不同的张量布局/表示（如稀疏和量化），这要求您的内核以更有效的方式编写，请遵循本教程并为 PyTorch 添加一个树外后端。

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

获取后端调度键

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

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

  • PrivateUse1/AutogradPrivateUse1

  • 私有用途 2/Autograd 私有用途 2

  • 私有用途 3/Autograd 私有用途 3

您可以选择上述任何键来原型化您定制的后端。要在 PrivateUse1 后端创建 Tensor，您需要在 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 类假设您的 Tensor 由 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!) 是表示操作符的唯一模式，它还包含与 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 算子可以分为两类：

  • 需要注册的 Ops：这些 Ops 的 PyTorch 原生实现是后端特定的，因此需要提供定制的后端内核。否则，在定制后端上调用此类 Ops 将导致错误。

    • 在 RegistrationDeclarations.h 中，这些算子的 dispatch 被设置为 True， default 被设置为 False，在它们的伴随注释中找到的元数据中。

  • 注册是可选的：后端扩展器可以不注册这些 Ops 而不会牺牲任何支持。然而，如果后端扩展器想要覆盖 PyTorch 提供的默认内核，他们仍然可以将他们的定制内核注册到他们的后端，并且调度器将仅为此后端使用它。例如，PyTorch 的当前实现 max_pool2d 将 indices 作为前向输出的部分返回，这为 torch_xla 创建了开销，因此 torch_xla 为其 max_pool2d 注册了自己的内核。

    • 在这些操作符的注释中找到的元数据中， RegistrationDeclarations.h 设置为 False 或 dispatch 设置为 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 反向内核实现也是设备特定的，以便可以从每个后端中挤出最大性能。对于这些算子，您将看到 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 到相应调度键（例如，如果您使用 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 仓库中的示例。

构建扩展程序

通过添加 C++扩展到 PyTorch 来支持树外后端。一旦你有内核和注册就绪，你可以通过编写一个使用 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 注册的内核支持以无包装和包装方式被调用。换句话说，您的自定义后端也可以与我们的 JIT 跟踪/脚本前端一起工作，就像树内后端（如 CPU 或 CUDA）一样。您还可以为您的后端在 JIT 图上编写专门的优化过程。但由于我们尚未最终确定 JIT 的集成点，所以当前的后端支持将专注于当前的急切前端。

测试您的后端与原生 PyTorch 后端

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

使用您的设备类型运行 PyTorch 的现有测试套件对于确保正确性很重要，但并非所有 PyTorch 功能都支持每种设备类型。通用设备类型测试框架允许进行相当程度的定制，以便设备类型可以选择要运行的测试、支持的 dtypes 以及在比较张量相等时使用的精度。

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

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

向后兼容性 ¶

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

因此，我们强烈建议树外后端扩展器仅与主要 PyTorch 版本同步，以最大限度地减少开发中断。PyTorch 的发布周期为每季度一次。后端扩展器应加入 pytorch.slack.com 的#announcement 频道，以获取有关发布的最新更新。

已知问题及附加说明 ¶

  • 并非所有测试套件都适用于所有设备。可以通过在 PyTorch 代码库中搜索 instantiate_device_type_tests 来找到可扩展的测试类，例如 TestTorchDeviceType, TestViewOps, TestTensorDeviceOps, TestTypePromotion 等。

  • 在 C++中，没有序列化 Python Tensor 对象到自定义后端的扩展点。目前，您只能通过修改 PyTorch Tensor 的__reduce_ex__方法或在树外仓库中进行猴子补丁来扩展它。

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

  • 如果您的后端与原生 PyTorch 优化器不兼容，例如需要将状态携带到反向传播中，如 torch-xla，则在 C++中没有优化器的扩展点。此类用例目前只能通过添加自定义 API 或在树外仓库中进行猴子补丁来实现。

未来工作

使 PyTorch 中的每个组件都支持树外后端无缝扩展，需要对 PyTorch 内部进行大量修改。以下是我们正在积极工作的几个可能改善未来体验的项目：

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

  • 提高内核覆盖率和更全面的测试，以确保内核行为与其他后端（如）匹配。

  • 尽可能重构 RegistrationDeclarations.h 以承载最小信息并尽可能多地重用 PyTorch 的 codegen。

  • 支持后端回退内核，自动将输入转换为 CPU，并将结果转换回自定义后端。即使没有为每个运算符编写内核，也能实现“全面”的运算符覆盖。

保持联系 ¶

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

如果您对上述任何未来工作项目感兴趣（例如，在 C++中为 PyTorch 运算符添加更多 Math 内核），请通过 GitHub 或 Slack 联系我们！

评分这个教程

© 版权所有 2024，PyTorch。

使用 Sphinx 构建，主题由 Read the Docs 提供。
//暂时添加调查链接

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取初学者和高级开发者的深入教程

查看教程

资源

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

查看资源

为了分析流量并优化您的体验，我们在本网站上提供 cookies。通过点击或导航，您同意允许我们使用 cookies。作为本站点的当前维护者，适用 Facebook 的 Cookies 政策。了解更多信息，包括可用的控制选项：Cookies 政策。