通过 PrivateUse1 简化新后端集成¶
创建日期：2023 年 10 月 03 日 | 最后更新日期：2024 年 05 月 07 日 | 最后验证日期：2024 年 11 月 05 日
在本教程中，我们将逐步讲解通过 PrivateUse1 将位于 pytorch/pytorch 仓库外部的新后端集成的必要步骤。请注意，本教程假设你已对 PyTorch 有基本了解，并且是 PyTorch 的高级用户。
注意
本教程仅涉及与 PrivateUse1 机制相关的部分，该机制用于简化新设备的集成，其他部分将不予介绍。同时，本教程中涉及的所有模块并非都是必需的，你可以根据实际需求选择对你有帮助的模块。
什么是 PrivateUse1？¶
在 PyTorch 2.0 之前，PyTorch 提供了三个保留的分发键（及其对应的 Autograd 键），用于原型化 out-of-tree 后端扩展，这三个分发键如下所示
PrivateUse1/AutogradPrivateUse1PrivateUse2/AutogradPrivateUse2PrivateUse3/AutogradPrivateUse3
原型验证通过后，你可以为新后端申请一个私有键，例如 CUDA、XLA、MPS 等等。
然而，随着 PyTorch 的快速发展，越来越多的硬件制造商试图将他们的后端集成到 PyTorch 中，这可能会导致以下问题
每一次新的后端集成都会涉及大量的文件修改
目前对分发键的数量存在硬性限制（
DispatchKeySet64 位限制）
注意
通过 PrivateUse1 键将新后端集成到 PyTorch 中也存在问题，因为无法同时集成许多后端。幸运的是，这些 out-of-tree 后端很少同时使用。
鉴于上述原因，社区开始推荐通过 PrivateUse1 将新后端集成到 PyTorch 中。
然而，以前的 PrivateUse1 机制尚不足以完全与新后端集成，因为它在某些模块中缺乏相关的支持，例如 Storage、AMP、Distributed 等等。
随着 PyTorch 2.1.0 的发布，在新的后端集成方面对 PrivateUse1 进行了一系列优化和增强，现在可以快速高效地支持新设备的集成。
如何通过 PrivateUse1 集成新后端¶
在本节中，我们将讨论通过 PrivateUse1 将新后端集成到 PyTorch 中的详细信息，主要包括以下几个部分
为新后端注册 kernel。
为新后端注册 generator。
为新后端注册设备 guard。
为新后端元数据注册序列化和反序列化函数。
其他模块。
为新后端注册 kernel¶
新后端可能有一些高性能的算子实现，可以通过 TORCH_LIBRARY_IMPL API 将其注册到 dispatcher，这在 在 C++ 中注册分发算子 中有所描述。这涉及几种情况
将新后端支持的所有前向算子注册到 dispatcher，同时注册 fallback，这样当新后端不支持某些算子时，这些算子可以回退到 CPU 执行，以确保功能的可用性。
at::Tensor wrapper_Custom_Tensor_add(const at::Tensor & self, const at::Tensor & other, const at::Scalar & alpha) {
// Implementation of add kernel in new backend
...
}
TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
...
m.impl("add.Tensor", TORCH_FN(wrapper_Custom_Tensor_add));
...
}
void custom_cpu_fallback(const c10::OperatorHandle& op, torch::jit::Stack* stack) {
// Add some hints about new devices that do not support and need to fall back to cpu
at::native::cpu_fallback(op, stack);
}
TORCH_LIBRARY_IMPL(_, PrivateUse1, m) {
m.fallback(torch::CppFunction::makeFromBoxedFunction<&custom_cpu_fallback>());
}
通过
AutogradPrivateUse1将来自torch::autograd::Function的 kernel 注册到 dispatcher，如果新后端需要覆盖PyTorch Autograd layer，dispatcher 和 autograd 系统将自动调用这些算子的前向和反向实现。
class CumtomSeluFunction : public torch::autograd::Function<CumtomSeluFunction> {
// Implementation of selu kernel in new backend
}
at::Tensor wrapper_AutogradCumstom__selu(const at::Tensor & self) {
return CumtomSeluFunction::apply(self);
}
TORCH_LIBRARY_IMPL(aten, AutogradPrivateUse1, m) {
...
m.impl("selu", TORCH_FN(wrapper_AutogradCustom__selu));
...
}
通过
AutocastPrivateUse1将支持 自动混合精度 (AMP) 和 fallback 机制的 kernel 注册到 dispatcher，autocast 系统将在需要时自动调用这些 kernel。
TORCH_LIBRARY_IMPL(aten, AutocastPrivateUse1, m) {
...
KERNEL_PRIVATEUSEONE(<operator>, <policy>)
...
}
TORCH_LIBRARY_IMPL(_, AutocastPrivateUse1, m) {
m.fallback(torch::CppFunction::makeFallthrough());
}
需要补充的是，如果新后端想支持 AMP，需要通过 torch._register_device_module("backend_name", BackendModule) 注册一个新的 BackendModule，并且该 BackendModule 需要具有以下 API
get_amp_supported_dtype() -> List[torch.dtype]获取新后端在 AMP 中支持的 dtypes，它可能支持一个或多个
dtype。
is_autocast_enabled() -> bool检查新后端上的 AMP 是否已启用。
get_autocast_dtype() -> torch.dtype获取新后端在 AMP 中支持的
dtype，它由set_autocast_dtype设置或使用默认的dtype，默认dtype是torch.float16。
set_autocast_enabled(bool) -> None在新后端上启用或禁用 AMP。
set_autocast_dtype(dtype) -> None设置新后端在 AMP 中支持的
dtype，并且该dtype必须包含在从get_amp_supported_dtype获取的dtypes中。
为新后端注册 generator¶
有必要支持对应新设备的 generator。目前，PrivateUse1 可以动态注册自定义 generator，主要分为以下几个步骤。
继承
GeneratorImpl类以实现对应新后端的 generator 类，并实现各种通用方法。定义一个新的后端
builder，带有一个参数：device index。调用
REGISTER_GENERATOR_PRIVATEUSE1宏以完成动态注册。
struct CustomGeneratorImpl : public c10::GeneratorImpl {
// Implementation of generator in new backend
}
at::Generator make_custom_generator(c10::DeviceIndex device_index) {
return at::make_generator<CustomGeneratorImpl>(device_index);
}
REGISTER_GENERATOR_PRIVATEUSE1(make_cumstom_generator)
为新后端注册设备 guard¶
PyTorch 通过 DeviceGuard 提供与设备、stream 和事件切换相关的功能。此功能也适用于 PrivateUse1 键。
继承
DeviceGuardImplInterface类以实现对应新后端的各种通用方法。调用
C10_REGISTER_GUARD_IMPL宏以完成动态注册。
struct CustomGuardImpl final : public c10::impl::DeviceGuardImplInterface {
// Implementation of guard in new backend
}
C10_REGISTER_GUARD_IMPL(PrivateUse1, CustomGuardImpl);
为新后端元数据注册序列化和反序列化函数¶
PyTorch 目前能够动态注册序列化/反序列化函数，以支持类 TensorImpl.ExtraMeta 中名为 backend_meta_ 的新后端附加元数据的序列化和反序列化。你可以参考以下步骤
继承
BackendMeta类以实现对应新后端的CustomBackendMetadata，并且可以在该类中自定义新后端的各种字段。实现新后端的序列化和反序列化函数，函数签名如下：
void(const at::Tensor&, std::unordered_map<std::string, bool>&)。调用
TensorBackendMetaRegistry宏以完成动态注册。
struct CustomBackendMetadata : public c10::BackendMeta {
// Implementation of backend metadata in new backend
}
void for_serialization(const at::Tensor& t, std::unordered_map<std::string, bool>& m) {
// Implementation of serialization
}
void for_deserialization(const at::Tensor& t, std::unordered_map<std::string, bool>& m) {
// Implementation of deserialization
}
TensorBackendMetaRegistry(c10::DeviceType::PrivateUse1, &for_serialization, &for_deserialization);
其他模块¶
除了上述部分之外，还有一些其他模块可以通过 PrivateUse1 进行扩展，例如 distributed collective communication、benchmark timer 等，这些将在未来添加。一个关于 PrivateUse1 集成的示例是 昇腾 NPU。
如何使用 Privateuse1 改进用户体验¶
通过 PrivateUse1 集成新设备的主要目标是满足基本的功能需求，接下来要做的是提高可用性，主要涉及以下几个方面。
将新后端模块注册到 PyTorch。
将 PrivateUse1 重命名为新后端的自定义名称。
生成与新后端相关的方法和属性。
将新后端模块注册到 PyTorch¶
PyTorch 中一些 CUDA 相关的接口可以通过以下形式调用：torch.cuda.xxx。因此，为了符合用户习惯，通过 PrivateUse1 机制实现的新后端也应提供类似的接口。
例如，使用 Ascend NPU
torch._register_device_module('npu', torch_npu.npu)
完成上述操作后，用户可以通过 torch.npu.xxx 调用 Ascend NPU 的一些独有 API
将 PrivateUse1 重命名为新后端的自定义名称¶
PrivateUse1 键是将新后端集成到 PyTorch 中的内部机制。对于用户来说，相比 PrivateUse1，与新后端紧密相关的自定义名称应该更友好。
以 Ascend NPU 为例，第一种用法将更用户友好。
torch.rand((2,2),device='npu:0')
torch.rand((2,2),device='privateuse1:0')
现在，PyTorch 为自命名 PrivateUse1 后端提供了一个新的 C++/Python API，使用起来非常简单。
torch.rename_privateuse1_backend("npu")
c10::register_privateuse1_backend("npu")
未来工作¶
PrivateUse1 机制的改进仍在进行中，因此将陆续添加新模块的 PrivateUse1 集成方法。以下是我们正在积极开展的一些事项
添加
distributed collective communication的集成方法。添加
benchmark timer的集成方法。
结论¶
本教程向你介绍了通过 PrivateUse1 将新后端集成到 PyTorch 中的过程，包括但不限于算子注册、generator 注册、设备 guard 注册等。同时，还介绍了一些改进用户体验的方法。
