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

创建于：2021年2月1日 | 最后更新：2024年9月23日 | 最后验证：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 操作符编写一个 C++ 内核，然后将它们注册到调度器中代表你自定义后端的调度键。

调度键是您在调度系统中的标识符。调度器会查看输入张量上携带的调度键，并相应地调用正确的内核。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。

注意

一旦原型完成并且您计划为您的后端扩展进行定期发布，请随时提交一个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 内核可以做什么的信息，因此暗示了后端扩展者是否需要实现该内核。更多详细信息可以在为新后端注册内核中找到。

为新后端注册内核¶

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

JIT支持¶

正如我们在Registering a Dispatched Operator in C++中提到的，通过m.impl() API注册的内核支持以未装箱和装箱的方式调用。换句话说，您的自定义后端也可以像CPU或CUDA等内置后端一样与我们的JIT跟踪/脚本前端一起工作。您还可以为JIT图上的后端编写专门的优化过程。但我们不会在这里讨论它，因为我们还没有在JIT中确定集成点，所以目前的后端支持将主要关注急切前端。

测试你的后端与原生PyTorch后端的对比¶

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

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

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

通用设备类型测试框架正在积极开发中。要请求功能，请在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内部进行大量更改。以下是我们正在积极改进的几个项目，可能会在未来提升体验：

提高通用测试框架的测试覆盖率。
提高Math内核的覆盖率和更全面的测试，以确保Math内核行为与其他后端如CPU/CUDA相匹配。
重构RegistrationDeclarations.h以携带最少的信息，并尽可能重用PyTorch的代码生成。
支持后端回退内核，自动将输入转换为CPU并将结果转换回自定义后端。这将允许“完整”的操作符覆盖，即使您没有为每个操作符编写内核。

保持联系¶

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

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