第9章 C++扩展PyTorch Dispatcher支持新后端指南
在 C++ 中扩展 dispatcher 以支持新后端#Created On: Feb 01, 2021 | Last Updated: Sep 23, 2024 | Last Verified: Nov 05, 2024
本教程将带你走完扩展 PyTorch dispatcher 的全部步骤，添加一个位于 pytorch/pytorch 仓库之外的新设备，并维护它与原生 PyTorch 设备保持同步。这里假设你已经熟悉如何在 C++ 中注册 dispatched operator，以及如何编写自定义 autograd 函数。
注意
本教程涉及大量 PyTorch 内部组件，这些组件仍在积极改进中，如果你决定跟随本教程，API 可能会发生变化。我们会持续更新本教程以匹配最新 API。
什么是新后端？#
为 PyTorch 添加新后端需要后端扩展者投入大量开发和维护工作。在动手之前，先看看几种常见场景及其推荐方案：
如果你为已有的 PyTorch 算子实现了新算法，向 PyTorch 提交 PR。 如果你想提出新算子，向 PyTorch 提交 feature request/PR。 如果你想支持新设备/硬件（如 Google TPU 或定制芯片），通常需要用硬件专用 API 编写 kernel，请按本教程为 PyTorch 添加一个 out-of-tree 后端。 如果你想以不同的 Tensor 布局/表示（如 sparse、quantized）支持现有算子，且 kernel 需要针对布局/表示限制做更高效的实现，请按本教程为 PyTorch 添加一个 out-of-tree 后端。
本教程主要聚焦于添加一个 out-of-tree 新设备。为不同 tensor 布局添加 out-of-tree 支持与设备的步骤有许多共通之处，但目前还没有这类集成的先例，可能需要 PyTorch 侧额外的支持工作。
为你的后端获取 dispatch key#
PyTorch 算子用 C++ 实现，通过 Python 绑定暴露给 Python 前端。PyTorch dispatcher 把一个算子的实现拆分为多个 kernel，每个 kernel 关联一个特定的 dispatch key。在 PyTorch 中支持新后端，本质上就是用 C++ 为每个 PyTorch 算子编写 kernel，然后把它们注册到 dispatcher 中代表你自定义后端的 dispatch key 上。 Dispatch key 是你在 dispatcher 系统中的标识符。dispatcher 根据输入 tensor 携带的 dispatch key 调用对应的 kernel。PyTorch 为 out-of-tree 后端扩展的原型开发预留了三个 dispatch key（及其对应的 Autograd key）：
PrivateUse1/AutogradPrivateUse1 PrivateUse2/AutogradPrivateUse2 PrivateUse3/AutogradPrivateUse3
你可以任选其中一个 key 来搭建自定义后端的原型。 要在 PrivateUse1 后端上创建 Tensor，需要在 TensorImpl 构造函数中设置 dispatch key。 /* 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 类假设你的 Tensor 由类似 CPU/CUDA 的 storage 支撑。我们也为没有 storage 的后端提供了 OpaqueTensorImpl，你可能需要调整/重写某些方法以适配自定义硬件。pytorch 仓库中的一个例子是 Vulkan TensorImpl。
注意
原型完成后，如果你计划为后端扩展做正式发布，欢迎向 pytorch/pytorch 提交 PR，为你的后端申请一个专属的 dispatch key。
获取 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++ kernel 必须与之完全一致。 aten::abs.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) 是算子的唯一 schema，相比 C++ 签名还包含别名（aliasing）和变更（mutation）注解。这是 dispatcher 用来查找算子的唯一标识符。 dispatch 和 default 是布尔字段，描述原生 PyTorch kernel 的情况，据此可以判断后端扩展者是否必须实现该 kernel。更多细节见为新后端注册 kernel 一节。
为新后端注册 kernel#
可以按照在 C++ 中注册 Dispatched Operator 中介绍的 TORCH_LIBRARY_IMPL API 把 kernel 注册到 PyTorch dispatcher：
TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
m.impl(
接下来详细看看哪些算子需要自定义后端提供 kernel，以及 kernel 里到底要写什么。 PyTorch 目前有超过 1600 个算子，而且还在持续增长，让后端扩展跟上这个速度并不现实。即使是 CPU、CUDA 这类原生后端，为每个新算子编写专属 kernel 也往往需要大量工作。好在一些原生 PyTorch kernel 的实现方式会分解为若干已知算子的组合，也就是说，你只需实现一组已知算子（即下文"需要注册的算子"），而不必覆盖所有 PyTorch 算子。 PyTorch 算子可以分为两类：
需要注册的算子：PyTorch 对这些算子的原生实现是后端相关的，因此必须为自定义后端提供 kernel，否则在自定义后端上调用会报错。
在 RegistrationDeclarations.h 中，这类算子注释元数据里 dispatch 为 True、default 为 False。
注册可选的算子：后端扩展者可以跳过这些算子的注册，不影响任何功能支持。不过，如果你想覆盖 PyTorch 提供的默认 kernel，也可以把自定义 kernel 注册到自己的后端，dispatcher 将只对你的后端使用它。例如，PyTorch 当前的 max_pool2d 会在 forward 输出中返回 indices，这给 torch_xla 带来了额外开销，所以 torch_xla 为 max_pool2d 注册了自己的 kernel。
在 RegistrationDeclarations.h 中，这类算子注释元数据里 dispatch 为 False 或 default 为 True。
为新后端提供 Autograd 支持#
梯度公式大多是纯数学推导，对所有后端通用。PyTorch 通常把 kernel 注册到别名 dispatch key Autograd 上，这意味着所有后端都能使用它。对这类算子你无需关心导数公式，只需为 RegistrationDeclarations.h 中的算子编写 forward 定义，PyTorch 会自动处理 backward。
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(
有些情况下，PyTorch 的 backward kernel 实现也是设备相关的，以便榨取每个后端的最大性能。对这些算子，你会在 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(
少数情况下，PyTorch 对某些算子的梯度公式可能带有并不适用于所有后端的假设。此时后端扩展者可以选择覆盖 PyTorch 的 Autograd 层，即把从 torch::autograd::Function 派生的 kernel 注册到对应的 dispatch key（例如使用 PrivateUse1 时对应 AutogradPrivateUse1）：
class MyAddFunction : public torch::autograd::Function
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(
// Register the inference kernel to PrivateUse1
TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
m.impl(
借助这个技巧，你可以完全掌控 my_add 算子在你的后端上训练和推理的行为。pytorch/xla 仓库中有一个示例。
构建扩展#
PyTorch 通过 C++ extension 的方式支持 out-of-tree 后端。kernel 和注册代码就绪后，就可以编写使用 setuptools 编译 C++ 代码的 setup.py 脚本来构建 C++ extension。以下是 pytorch/xla 仓库中的一个简化示例： 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++ extension 教程。
自定义算子支持#
只要自定义算子由现有 PyTorch 算子组合而成（这些算子你的后端已支持），新后端就应该能无缝支持通过 Python 扩展的自定义算子，无需编写任何新 kernel。 对于通过 C++ 扩展的自定义算子，通常会有一个后端相关的 C++ kernel 实现（如 torchvision 中的 nms kernel），以及一个自定义 Python API（如 torch.ops.torchvision.nms）。要支持这类算子，后端扩展者需要为自己的后端编写 C++ kernel，并像支持 PyTorch 原生算子那样把它注册到 dispatcher 中对应的 namespace。或者，你也可以在自己的扩展里为这类临时需求添加自定义 API，例如 torch_xla.core.functions.nms。
JIT 支持#
正如在"在 C++ 中注册 Dispatched Operator"中提到的，通过 m.impl() 注册的 kernel 既支持 unboxed 调用也支持 boxed 调用。换句话说，你的自定义后端可以和 CPU、CUDA 等 in-tree 后端一样，配合 JIT tracing/scripting 前端工作。理论上你还可以针对你的后端在 JIT graph 上编写专门的优化 pass，但这里不展开，因为 JIT 的集成点尚未定型，目前后端支持以 eager 前端为主。
用原生 PyTorch 后端测试你的后端#
PyTorch 提供通用的设备类型测试框架，让测试能在多种设备类型上运行。你可以在这里了解测试如何使用该框架，以及如何添加新设备类型的说明。添加之后，使用该框架的 PyTorch 测试也会在你的设备类型上运行。这个 Wiki 页面展示了测试如何被实例化的示例。 让你的设备类型跑一遍 PyTorch 现有测试套件对保证正确性很重要，但并非所有 PyTorch 功能在每个设备类型上都支持。通用设备类型测试框架支持大量定制，设备类型可以选择要运行哪些测试、支持哪些 dtype，甚至比较 tensor 相等性时使用什么精度。 XLA 是一个使用该框架但不在 PyTorch 中内置的设备类型示例。可以查看它对该框架的扩展，其中包含 block list 测试、block list dtype、覆盖测试精度等示例。 通用设备类型测试框架仍在积极开发中，如需新功能请在 PyTorch 的 Github 上提 issue。
向后兼容性#
目前 PyTorch 无法保证已注册算子的向后兼容性。算子及其 schema 可能按需增删或修改。已注册的 kernel 必须与 PyTorch 版本完全一致——如果 PyTorch 为某算子新增参数（即使带默认值），你的旧注册在更新到匹配新签名之前将无法工作。 因此，我们强烈建议 out-of-tree 后端扩展者只跟随 PyTorch 的正式版本同步，以尽量减少开发中断。PyTorch 按季度发版，后端扩展者应加入 pytorch.slack.com 的 #announcement 频道获取最新发布动态。
已知问题与补充说明#
并非所有测试套件都已设备通用化。可在 PyTorch 代码库中搜索 instantiate_device_type_tests 找到可扩展的测试类，例如 TestTorchDeviceType、TestViewOps、TestTensorDeviceOps、TestTypePromotion 等。 C++ 中没有为自定义后端的 python Tensor 对象序列化提供扩展点。目前只能通过修改 PyTorch Tensor 的 __reduce_ex__ 方法，或在 out-of-tree 仓库中 monkey patch 来实现。 如果你的后端不支持直接内存访问，需要特别关注 view 类算子，因为它们应当共享 storage。对 view tensor 的修改需要同步传播到其 base tensor，反之亦然。 如果你的后端无法使用 PyTorch 原生 Optimizer（例如像 torch-xla 那样需要在 backward 中携带待更新的状态），C++ 中没有为 Optimizer 提供扩展点。这类场景目前只能通过添加自定义 API 或在 out-of-tree 仓库中 monkey patch 来实现。
后续工作#
要让 PyTorch 的每个组件都能无缝支持 out-of-tree 后端扩展，需要对内部做大量改动。以下是我们正在积极推进的项目，未来有望改善体验：
提高通用测试框架的测试覆盖率。 提高 Math kernel 覆盖率并补充更全面的测试，确保 Math kernel 行为与 CPU/CUDA 等其他后端一致。 重构 RegistrationDeclarations.h，只保留最少信息，并尽可能复用 PyTorch 的 codegen。 支持后端 fallback kernel，自动把输入转成 CPU、计算后再把结果转回自定义后端。这样即使你没有为每个算子编写 kernel，也能获得"完整"的算子覆盖。
保持联系#
如有问题和讨论，请使用 PyTorch dev discussions。如有功能需求或 bug 报告，请在 github 上提 issue。 如果你有兴趣参与上述任何一项后续工作（例如用 C++ 为 PyTorch 算子添加更多 Math kernel），欢迎通过 Github 或 Slack 联系我们！
