目录
目录

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

创建时间： Feb 01， 2021 |上次更新时间：2024 年 9 月 23 日 |上次验证： Nov 05， 2024

在本教程中，我们将演练将 Dispatcher 扩展到 添加位于存储库之外的新设备并维护它以保留在 与本机 PyTorch 设备同步。在这里，我们假设您熟悉 在 C++ 中注册调度运算符以及如何编写自定义 autograd 函数。pytorch/pytorch

注意

本教程涉及 PyTorch 中的许多内部组件，这些组件正在积极改进中。 如果您决定遵循本教程，请期待 API 的变化。我们将保留本教程 使用最新的 API 进行更新。

什么是新后端？

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

  • 如果您有现有 PyTorch 运算符的新算法，请向 PyTorch 发送 PR。

  • 如果要提议新的运算符，请向 PyTorch 发送功能请求/PR。

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

  • 如果要添加对现有运算符的支持，但具有不同的 Tensor 布局/表示 就像 Sparse 和 Quantized 一样，它强制以更高效的方式编写 kernel 鉴于布局/表示限制，请按照本教程向 PyTorch 添加树外后端。

在本教程中，我们将主要关注在下面添加新的树外设备。添加树外支持 对于不同的 Tensor 布局，可能会与设备共享许多常见步骤，但我们还没有看到 因此可能需要 PyTorch 的额外工作来支持它。

获取后端的 dispatch key

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

Dispatch key 是您在 Dispatcher 系统中的标识符。调度程序查看执行的 dispatch key input 张量并相应地调用正确的内核。PyTorch 提供 3 个预留的 dispatch key （及其相应的 Autograd 密钥）来构建树外后端扩展的原型：

  • PrivateUse1/AutogradPrivateUse1

  • PrivateUse2/AutogradPrivateUse2

  • PrivateUse3/AutogradPrivateUse3

您可以选择上述任何键来构建自定义后端的原型。 要在后端创建 Tensor，你需要在 constructor 中设置 dispatch key。PrivateUse1TensorImpl

/* 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};

请注意，上面的类假定您的 Tensor 由 CPU/CUDA 等存储提供支持。我们还 为没有存储的后端提供。你可能需要调整/覆盖某些 方法以适应您的自定义硬件。 pytorch 存储库中的一个示例是 Vulkan TensorImpl。TensorImplOpaqueTensorImpl

注意

原型完成后，您计划为后端扩展进行定期发布，请随时 提交 PR to 为您的后端保留专用的 dispatch key。pytorch/pytorch

获取 PyTorch 运算符的完整列表

PyTorch 在 generated file 中提供了可扩展 C++ 运算符的完整列表。 此文件仅在从源代码构建 PyTorch 后可用。 以下是该文件的一个片段：build/aten/src/ATen/RegistrationDeclarations.h

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"}

有多个字段与单个运算符相关联。让我们以 As example 来分解它：abs_out

  • Tensor & abs_out(Tensor & out, const Tensor & self);是运算符的 C++ 签名，即您的 C++ kernel 应该与这个签名完全匹配。

  • aten::abs.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)是表示运算符的唯一架构， 其中还包含与 C++ 签名相比的别名和突变注释。这是唯一标识符 调度程序用于查找运算符。

  • dispatch和 是布尔字段，提供有关本机 PyTorch 内核的信息 可以执行，因此意味着后端扩展器是否需要实现内核。 更多详细信息可以在 register kernels for the new backend 中找到。default

为新后端注册内核

要将内核注册到 PyTorch 调度程序，可以使用在 C++ 中注册调度的运算符中描述的 API：TORCH_LIBRARY_IMPL

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);
}

现在让我们放大一下，什么 Operator 需要来自自定义后端的内核，以及 确切地在内核内部。

PyTorch 目前有 1600 多个算子，并且还在不断增长。这是不现实的 让后端扩展跟上这个速度。即使对于 CPU 等原生后端也是如此 或 CUDA 编写专用内核，通常需要大量工作才能为每个新操作编写专用内核。

幸运的是，一些原生 PyTorch 内核的编写方式可以分解为 多个已知运算符的组合。换句话说，您只需要实现 一组已知的运算符（需要在下面注册的运算），而不是所有 PyTorch 运算符。

PyTorch 算子可以分为两类：

  • 需要注册的操作：这些操作的 PyTorch 本机实现是特定于后端的 因此，需要为自定义后端提供内核。否则调用此类 op 在自定义后端会出错。

    • 在这些运算符中，已设置为 True 并设置为 False 在他们随附的评论中找到的元数据中。RegistrationDeclarations.hdispatchdefault

  • 注册是可选的：后端扩展器可以跳过注册到这些操作，而不会牺牲任何支持。 但是，如果后端扩展器想要覆盖 PyTorch 提供的默认内核，它们仍然可以 将他们的自定义内核注册到他们的后端，Dispatcher 将仅将其用于您的后端。 例如，PyTorch 的当前实现将 return 作为 forward output 的一部分， 在 torch_xla 中产生开销，因此 torch_xla 会注册自己的内核。max_pool2dindicesmax_pool2d

    • 在这些运算符中，已设置为 False 或 True 在他们随附的评论中找到的元数据中。RegistrationDeclarations.hdispatchdefault

Autograd 对新后端的支持

梯度公式大多是纯数学公式，因此适用于所有后端。 PyTorch 通常会注册一个内核来别名 dispatch key Autograd，这意味着它可以被所有后端使用。

对于这些运算符，您不必担心它们的导数公式， 您只需在 PyTorch 句柄中为运算符编写正向定义 自动为您倒退。RegistrationDeclarations.h

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 用于后端）：

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);
}

使用此技巧，您可以完全控制后端 operator 的训练和推理行为。 这是存储库中的一个示例。my_addpytorch/xla

构建扩展

通过向 PyTorch 添加 C++ 扩展来支持树外后端。 准备好内核和注册后，您可以通过以下方式构建 C++ 扩展 编写用于编译 C++ 代码的脚本。以下是 pytorch/xla 存储库中的简化示例：setup.pysetuptools

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++ 内核实现，例如 torchvsion 中的 nms 内核，以及自定义的 Python API，例如 torch.ops.torchvision.nms。 为了支持这些运算符，后端扩展器需要为您的后端编写一个 C++ 内核，并且 将其注册到 Dispatcher 中的相应命名空间，类似于支持 PyTorch 本机运算符。 或者，您也可以在扩展中添加自定义 API，例如 这些临时请求。torch_xla.core.functions.nms

JIT 支持

正如我们在 C++ 中注册 Dispatched Operator 中提到的，通过 m.impl（） API 注册的内核 支持以 unboxed 和 boxed 方式调用。换句话说，您的自定义后端也可以与我们的 JIT 跟踪/脚本前端就像 CPU 或 CUDA 等树内后端一样。您还可以编写专门的优化 传递给 JIT 图上的后端。但我们不会在这里讨论它，因为我们还没有最终确定集成点 在 JIT 中，因此当前的后端支持现在将集中在 Eager 前端。

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

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

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

使用通用设备类型测试框架且未随 PyTorch 是 XLA。请参阅其对通用设备类型测试框架的扩展， 其中包含块列表测试、块列表 DTYPES 和覆盖测试精度的示例。

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

向后兼容性

目前，PyTorch 无法保证已注册的 Operator 的向后兼容性。 可以根据需要添加/修改/删除运算符及其架构。注册 kernels 必须与 PyTorch 版本完全相同。如果 PyTorch 添加了更多参数（ 即使有默认值），您的旧注册在更新之前也不会生效 以匹配 PyTorch 的新签名。

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

已知问题和额外说明

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

  • C++ 中没有用于在自定义后端序列化 python Tensor 对象的扩展点。现在 您只能通过修改 PyTorch Tensor __reduce_ex__ 方法或在树外存储库中修补 monkey 来扩展它。

  • 如果您的后端不允许直接访问内存，则应额外注意支持 查看 OPS，因为它们应该共享存储空间。对视图张量的更改需要传播到其 base tensor 的 base 张量，反之亦然。

  • 如果您的后端不能与原生 PyTorch 一起使用，则 C++ 中没有 Optimizer 的扩展点 优化器，例如需要像 torch-xla 一样携带要向后更新的状态。此类用例 目前只能通过在树外存储库中添加自定义 API 或 monkey 补丁来完成。

未来的工作

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

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

  • 提高内核覆盖率和更全面的测试，以确保内核行为与其他后端（如 .MathMathCPU/CUDA

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

  • 支持后端回退内核自动将输入转换为 CPU，并将 result 返回自定义后端。这将允许“完全”运维覆盖 尽管您没有为每个 Operator 编写内核。

保持联系

请使用 PyTorch dev discussions 来提出问题和讨论。如果你有 任何功能请求或错误报告，请在 GitHub 上提交问题。

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

文档

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

查看文档

教程

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

查看教程

资源

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

查看资源