目录
目录

在C++中扩展调度器以支持新后端

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

在本教程中，我们将逐步介绍将调度程序扩展以添加一个位于 pytorch/pytorch 仓库之外的新设备并保持其与原生 PyTorch 设备同步所需的所有必要步骤。这里假设您已经熟悉如何在 C++ 中 注册分派操作符 以及如何编写 自定义自动求导函数。

注意

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

什么是新后端?

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

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

  • 如果你想提出一个新操作符，请向 PyTorch 提交功能请求或提交拉取请求。

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

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

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

获取您后端的调度密钥

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

Dispatch key 是你在调度系统中的标识符。调度器会查看输入张量所携带的 dispatch keys，并据此调用相应的内核。PyTorch 为原型化树外后端扩展提供了三个保留的 dispatch keys（以及对应的 Autograd keys）：

  • PrivateUse1/AutogradPrivateUse1

  • PrivateUse2/AutogradPrivateUse2

  • PrivateUse3/AutogradPrivateUse3

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

/* 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以保留一个专用的调度键用于你的后端。

获取完整的PyTorch算子列表

PyTorch 提供了一个完整的可扩展 C++ 操作符列表，生成的文件 build/aten/src/ATen/RegistrationDeclarations.h。 此文件仅在从源代码构建 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调度器，您可以使用 TORCH_LIBRARY_IMPL 在 C++中注册调度操作符 中描述的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 原生实现是后端特定的，因此需要为自定义后端提供内核。否则，在自定义后端上调用此类操作将导致错误。

    • In RegistrationDeclarations.h these operators have dispatch set to True and default set to False in the metadata found in their accompanying comments.

  • 注册是可选的：后端扩展者可以跳过对这些操作的注册而不会丧失任何支持。 然而，如果后端扩展者希望覆盖PyTorch提供的默认内核，他们仍然可以 在自己的后端注册自定义内核，并且调度器将仅为此后端使用它。 例如，PyTorch当前的max_pool2d实现返回indices作为前向输出的一部分， 这在torch_xla中会产生开销，因此torch_xla注册了自己的max_pool2d内核。

    • In RegistrationDeclarations.h these operators have dispatch set to False or default set to True in the metadata found in their accompanying comments.

PyTorch新后端的自动求导支持

梯度公式大多是纯粹数学的，因此适用于所有后端。 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 自动微分层：

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来使用外部后端。 一旦你准备好内核和注册，就可以通过编写一个setup.py脚本来构建C++扩展，该脚本使用setuptools来编译C++代码。以下是一个简化示例，来自 pytorch/xla repo：

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++内核实现，例如torchvision中的nms内核 以及一个自定义的Python API，例如torch.ops.torchvision.nms。 为了支持这些操作符，后端扩展者需要为你的后端编写一个C++内核，并正确地 将其注册到调度程序中的相应命名空间，类似于支持PyTorch原生操作符。 或者，你也可以在扩展中添加一个自定义API，例如torch_xla.core.functions.nms， 以应对这些临时请求。

JIT 支持

正如我们在在C++中注册分发操作符中提到的，通过m.impl() API注册的内核支持以未装箱和装箱两种方式调用。换句话说，你的自定义后端也可以像树中的后端（如CPU或CUDA）一样与我们的JIT追踪/脚本前端一起工作。你还可以为你的后端在JIT图上编写专门的优化过程。但由于我们尚未确定JIT中的集成点，因此当前的后端支持将暂时集中在eager前端。

对原生PyTorch后端进行后端测试

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

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

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

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

向后兼容性

目前 PyTorch 无法保证已注册操作符的向后兼容性。 操作符及其 schema 可能会根据需要被添加、修改或删除。已注册的内核必须与 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 Optimizers，并且需要像 torch-xla 那样在反向传播中携带需要更新的状态，C++ 中没有为 Optimizer 提供扩展点。此类用例目前只能通过在外部仓库中添加自定义 API 或 monkey patching 来实现。

未来工作

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

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

  • 提升 Math 核心覆盖率和更全面的测试，以确保 Math 核心行为与其他后端如 CPU/CUDA 一致。

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

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

保持联系

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

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

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

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

查看教程

资源

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

查看资源