通过PrivateUse1促进新的后端集成
创建于：2023年10月03日 | 最后更新：2024年5月07日 | 最后验证：2024年11月05日
在本教程中，我们将逐步介绍一些必要的步骤，以通过PrivateUse1集成一个位于pytorch/pytorch仓库之外的新后端。请注意，本教程假设您已经对PyTorch有基本的了解，并且是PyTorch的高级用户。
注意
本教程仅涉及与促进新设备集成的PrivateUse1机制相关的部分，其他部分将不涉及。同时，本教程涉及的所有模块并非都是必需的，您可以根据实际需求选择对您有帮助的模块。
什么是PrivateUse1？
在Pytorch 2.0之前，PyTorch提供了三个保留的调度键（及其对应的Autograd键）用于原型化树外后端扩展，这三个调度键如下：
PrivateUse1/AutogradPrivateUse1PrivateUse2/AutogradPrivateUse2PrivateUse3/AutogradPrivateUse3
原型验证通过后，您可以申请新后端的私钥，例如CUDA、XLA、MPS等。
然而，随着PyTorch的快速发展，越来越多的硬件制造商试图将他们的后端集成到PyTorch中，这可能会导致以下问题：
每个新的后端集成都涉及大量的文件修改
目前对调度键的数量有一个硬性限制（
DispatchKeySet64位限制）
注意
通过PrivateUse1键将新后端集成到PyTorch中也存在问题，因为无法同时集成多个后端。幸运的是，这些树外后端很少同时使用。
鉴于上述原因，社区开始推荐通过PrivateUse1将新的后端集成到PyTorch中。
然而，之前的PrivateUse1机制无法完全与新的后端集成，因为它在某些模块中缺乏相关的支持，例如存储、AMP、分布式等。
随着Pytorch 2.1.0的到来，PrivateUse1在新后端集成方面进行了一系列优化和增强，现在可以快速高效地支持新设备的集成。
如何通过PrivateUse1集成新的后端
在本节中，我们将讨论通过PrivateUse1将新后端集成到Pytorch中的细节，主要包括以下几个部分：
为新后端注册内核。
为新后端注册生成器。
为新后端注册设备保护。
为新后端元数据注册序列化和反序列化函数。
其他模块。
为新后端注册内核
新的后端可能会有一些高性能的操作符实现，可以通过在C++中注册一个分派操作符中描述的TORCH_LIBRARY_IMPL API注册到调度器。这涉及几种情况：
将所有新后端支持的前向操作符注册到调度器，并同时注册回退机制，以便当新后端不支持某些操作符时，这些操作符可以回退到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的内核注册到调度器，如果新后端需要覆盖PyTorch Autograd layer，调度器和自动求导系统将自动调用这些算子的前向和后向实现。
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));
...
}
注册希望支持自动混合精度（AMP）并希望通过
AutocastPrivateUse1回退机制到调度器的内核，自动转换系统将在需要时自动调用这些内核。
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中。
为新后端注册生成器
需要支持与新设备对应的生成器。目前，PrivateUse1 可以动态注册自定义生成器，主要分为以下步骤。
继承
GeneratorImpl类以实现与新后端对应的生成器类，并实现各种通用方法。定义一个新的后端
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)
为新后端注册设备保护
PyTorch 提供了通过 DeviceGuard 进行设备、流和事件切换的功能。
此功能也适用于 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。调用
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集成的一个例子是Ascend 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 提供了一个新的 C++/Python API 用于名为 PrivateUse1 的后端，使用起来非常简单。
torch.rename_privateuse1_backend("npu")
c10::register_privateuse1_backend("npu")
