Table of Contents
Shortcuts
advanced/extend_dispatcher

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

Created On: Feb 01, 2021 | Last Updated: Sep 23, 2024 | Last Verified: Nov 05, 2024

在本教程中，我们将逐步介绍扩展调度器的所有必要步骤，以添加一个新的设备（生活在``pytorch/pytorch``库之外），并维护它以与原生 PyTorch 设备保持同步。在这里我们假设你已经熟悉如何`在 C++ 中注册调度操作符 <dispatcher>`_以及如何编写`自定义自动求导函数 <cpp_autograd>`_。

备注

本教程涵盖了 PyTorch 内部许多正在积极改进的组件，如果您决定遵循本教程，请预期 API 会有所变化。我们将根据最新的 API 更新本教程。

什么是新后端？

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

  • 如果您对现有 PyTorch 操作符有新的算法，请向 PyTorch 提交 PR。

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

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

  • 如果您想为现有操作符添加支持，但具有不同的张量布局/表示形式，例如稀疏和量化，这会要求您以一种给定布局/表示形式限制下更高效的方式编写内核，请遵循此教程并为 PyTorch 添加一个树外后端。

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

为您的后端获取调度键

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

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

  • 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 <https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/vulkan/VulkanOpaqueTensorImpl.h>`_。

备注

一旦原型完成并计划为您的后端扩展进行定期发布，请随时提交 PR 至``pytorch/pytorch``以为您的后端保留一个专用的调度键。

获取 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 内核能力的信息，因此暗示是否需要后端扩展者实现该内核。更多细节请参见:ref:为新后端注册内核<register-kernel>。

为新后端注册内核

要将您的内核注册到 PyTorch 调度器，您可以使用``TORCH_LIBRARY_IMPL`` API，该 API在`注册一个调度操作符 <dispatcher>`_中有所描述：

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。

为新后端提供自动求导支持

梯度公式主要是纯数学的，因此对所有后端都是通用的。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 的自动求导层：

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``仓库中的`一个例子 <https://github.com/pytorch/xla/blob/r1.7/torch_xla/csrc/aten_autograd_ops.h>`_。

构建扩展

通过向 PyTorch 添加 C++ 扩展支持树外后端。一旦您准备好内核和注册，可以通过编写使用``setuptools``编译C++代码的``setup.py``脚本构建一个C++扩展。以下是`pytorch/xla仓库 <https://github.com/pytorch/xla/blob/master/setup.py>`_中的一个简化示例：

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++ 扩展教程 <https://pytorch.org/tutorials/advanced/cpp_extension.html#building-with-setuptools>`_。

自定义操作符支持

您的新后端应能与`Python中扩展的自定义操作符 <https://pytorch.org/docs/stable/notes/extending.html>`_无缝工作，而无需编写任何新的内核，只要自定义操作符由现有 PyTorch 操作符组成（这些已经被您的后端支持）。

对于在 C++ 中扩展的`自定义操作符 <cpp_autograd>`_，通常会附带`具体后端的 C++ 内核实现，例如 torchvision 中的 nms 内核 <https://github.com/pytorch/vision/blob/master/torchvision/csrc/ops/cuda/nms_kernel.cu>`_，以及`自定义的 Python API，例如 torch.ops.torchvision.nms <https://github.com/pytorch/vision/blob/master/torchvision/csrc/ops/nms.cpp#L18>`_。为了支持这些操作符，后端扩展者需要为其后端编写一个 C++ 内核，并将其正确注册到调度器中的对应命名空间，与支持 PyTorch 原生操作符类似。或者，您也可以在扩展中添加自定义的 API，例如 torch_xla.core.functions.nms 以满足这些临时请求。

JIT 支持

正如我们在`在 C++ 中注册调度操作符 <dispatcher>`_中提到的，通过 m.impl() API 注册的内核支持以未封装方式和封装方式调用。换句话说，您的自定义后端也可以像内置的后端（如 CPU 或 CUDA）一样，与我们的 JIT 跟踪/脚本前端协同工作。您甚至可以为您的后端在 JIT 图上编写专门的优化流程。不过，我们不会在这里讨论，因为我们尚未在 JIT 中确定集成点，因此目前的后端支持将主要聚焦于即时前端。

根据 PyTorch 原生后端测试您的后端

PyTorch 允许使用其`通用设备类型测试框架 <https://github.com/pytorch/pytorch/blob/master/torch/testing/_internal/common_device_type.py>`_在多个设备类型上运行测试。您可以在`测试如何使用它 <https://github.com/pytorch/pytorch/blob/5a8198eb3c594aa18352930fd21f3c25bd7b7100/torch/testing/_internal/common_device_type.py#L23>`_和`如何添加新的设备类型的信息 <https://github.com/pytorch/pytorch/blob/5a8198eb3c594aa18352930fd21f3c25bd7b7100/torch/testing/_internal/common_device_type.py#L369>`_中找到详细信息。一旦添加了新的设备类型，使用通用设备类型测试框架的 PyTorch 测试也将在您的设备类型上运行。有关如何实例化测试的示例，请参阅`这篇 Wiki 页面 <https://github.com/pytorch/pytorch/wiki/Writing-tests-that-run-on-all-available-device-types>`_。

使用您的设备类型运行 PyTorch 的现有测试套件对于确保正确性至关重要，但并不是所有 PyTorch 功能都支持每个设备类型。通用设备类型测试框架允许相当程度的自定义，因此设备类型可以选择运行哪些测试、支持哪些数据类型甚至在比较张量是否相等时使用哪些精度。

一个不随 PyTorch 一起发布但使用通用设备类型测试框架的示例设备类型是 XLA。请参阅其`对通用设备类型测试框架的扩展 <https://github.com/pytorch/xla/blob/master/test/pytorch_test_base.py>`_，其中包含屏蔽测试、屏蔽数据类型以及覆盖测试精度的示例。

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

向后兼容性

目前，PyTorch 无法保证注册操作符的向后兼容性。操作符及其架构可能会根据需要添加、修改或删除。注册的内核必须与 PyTorch 的版本*完全一致*。如果 PyTorch 为某个操作符添加了更多参数（即使是带默认值的参数），您的旧注册将无法工作，直到它更新为符合 PyTorch 的新签名。

因此，我们*强烈建议*扩展的外部后端仅与主要 PyTorch 版本同步，以最大限度地减少开发中的中断。PyTorch 是一个季度发布周期。后端扩展者应加入`pytorch.slack.com <http://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 内部进行大量更改。以下是我们正在积极研究的一些可能改善未来体验的项目：

  • 改进通用测试框架的测试覆盖率。

  • 改进 Math 内核覆盖率以及更全面的测试，以确保 Math 内核行为与其他后端（如 CPU/CUDA）一致。

  • 重构 RegistrationDeclarations.h 以包含最少信息并尽可能重复使用 PyTorch 的代码生成。

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

保持联系

请使用`PyTorch 开发讨论 <https://dev-discuss.pytorch.org/>`_来提问和讨论。如果您有任何功能请求或错误报告，请在 GitHub 上`提交问题 <https://github.com/pytorch/pytorch/issues>`_。

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

文档

访问 PyTorch 的详细开发者文档

查看文档

教程

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

查看教程

资源

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

查看资源

为了分析流量并优化您的体验，我们在此网站上提供Cookie。点击或浏览即表示您同意我们使用Cookie。 作为本网站的当前维护者，Facebook的Cookie政策适用。了解更多信息，包括可用的控件： Cookies政策.