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

什么是新后端？

在 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 后端上创建 Tensor，您需要在 TensorImpl 构造函数中设置调度键。

/* Example TensorImpl constructor */
TensorImpl(
Storage&&storage,
DispatchKeySetks,
constcaffe2::TypeMetadata_type);

// To create a TensorImpl on PrivateUse1 backend, pass in the following ks to TensorImpl creation.
DispatchKeySetks=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 后可用。以下是该文件的片段：

Tensorabs(constTensor&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,constTensor&self);// {"schema": "aten::abs.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)", "dispatch": "True", "default": "False"}
Tensorabsolute(constTensor&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,constTensor&self);// {"schema": "aten::absolute.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)", "dispatch": "False", "default": "False"}
Tensorangle(constTensor&self);// {"schema": "aten::angle(Tensor self) -> Tensor", "dispatch": "True", "default": "True"}
Tensor&angle_out(Tensor&out,constTensor&self);// {"schema": "aten::angle.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)", "dispatch": "True", "default": "False"}
Tensorsgn(constTensor&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 会自动为您处理反向传播。

Tensormy_op1(constTensor&self,constTensor&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 中作为必需的注册出现。

Tensormy_op2_backward(constTensor&self,constTensor&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 层：

classMyAddFunction:publictorch::autograd::Function<MyAddFunction>{
public:
staticTensorforward(AutogradContext*ctx,torch::Tensorself,torch::Tensorother){
at::AutoNonVariableTypeModeg;
returnmyadd(self,other);
}

statictensor_listbackward(AutogradContext*ctx,tensor_listgrad_outputs){
autograd_output=grad_outputs[0];
return{grad_output,grad_output};
}
};

Tensormyadd_autograd(constTensor&self,constTensor&other){
returnMyAddFunction::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 的简化示例：

fromsetuptoolsimport setup
fromtorch.utils.cpp_extensionimport 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 支持

正如我们在在 C++ 中注册调度运算符中提到的，通过 m.impl() API 注册的内核支持以未装箱（unboxed）和装箱（boxed）的方式调用。换句话说，您的自定义后端也可以像内置的后端（如 CPU 或 CUDA）一样与我们的 JIT 追踪/脚本化前端协同工作。您甚至可以为您的后端在 JIT 图上编写专门的优化过程。但由于我们尚未在 JIT 中确定集成点，因此这里不会讨论此内容，当前的后端支持将主要关注于急切模式（eager mode）前端。

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

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

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

一个使用通用设备类型测试框架且不随 PyTorch 一起发布的设备类型示例是 XLA。请参阅它对通用设备类型测试框架的扩展，其中包含屏蔽测试、屏蔽 dtype 以及覆盖测试精度的示例。

通用设备类型测试框架正在积极开发中。如需请求功能，请在 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++ 中没有为 Optimizer 提供扩展点。目前，此类用例只能通过添加自定义 API 或在外部代码库中进行猴子补丁来实现。

未来工作

为了使 PyTorch 中的每个组件都能无缝扩展以支持树外后端，需要对 PyTorch 内部进行大量更改。以下是我们正在积极改进的一些项目，未来可能会带来更好的体验：

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

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

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

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

保持联系

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

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