通过 PrivateUse1 便捷整合新后端¶
Created On: Oct 03, 2023 | Last Updated: May 07, 2024 | Last Verified: Nov 05, 2024
在本教程中，我们将展示一些必要的步骤，以通过 PrivateUse1 将 pytorch/pytorch 仓库之外的新后端集成进来。请注意，本教程假定您已经对 PyTorch 有基本了解，并且是 PyTorch 的高级用户。
备注
本教程仅涉及与 PrivateUse1 机制相关的部分，该机制便于新设备的集成，其他部分不会涉及。同时，本教程所涉及的模块并非全部必需，您可以根据实际需要选择对您有帮助的模块。
什么是 PrivateUse1？¶
在 PyTorch 2.0 之前，PyTorch 提供了三个保留的分派键（及其对应的 Autograd 键）用于原型外部后端扩展，这三个分派键如下：
PrivateUse1/AutogradPrivateUse1PrivateUse2/AutogradPrivateUse2PrivateUse3/AutogradPrivateUse3
原型验证通过后，您可以为新后端申请私有键，例如 CUDA、XLA、MPS 等。
但是，随着 PyTorch 的快速发展，越来越多的硬件制造商试图将后端集成到 PyTorch 中，这可能导致以下问题：
每次新后端集成都涉及大量文件修改
目前，分派键的数量有硬性限制（
DispatchKeySet为 64 位限制）
备注
通过 PrivateUse1 键将新后端集成到 PyTorch 中也存在问题，因为无法同时集成多个后端。幸运的是，这些树外后端很少同时使用。
基于上述原因，社区开始推荐通过 PrivateUse1 将新后端集成到 PyTorch 中。
然而，之前的 PrivateUse1 机制并不能完全支持新后端的集成，因为它在某些模块中缺乏支持，比如 Storage、AMP 和分布式等。
随着 PyTorch 2.1.0 的到来，PrivateUse1 在新后端集成方面进行了一系列优化和增强，现在可以快速高效地支持新设备的集成。
如何通过 PrivateUse1 集成新后端¶
在本节中，我们将讨论通过 PrivateUse1 将新后端集成到 PyTorch 中的细节，这主要由以下几个部分组成：
为新后端注册内核。
为新后端注册生成器。
为新后端注册设备管理器。
为新后端的元数据注册序列化和反序列化功能。
其他模块。
为新后端注册内核¶
新后端可能有一些高性能的运算符实现，这些可以通过使用 TORCH_LIBRARY_IMPL API（描述于 Registering a Dispatched Operator in C++）注册到分派器中。这涉及几种情况：
将新后端支持的所有前向运算符注册到分派器，同时注册后备机制，这样在新后端不支持某些运算符时，这些运算符可以回退到 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 层，分派器和自动求导系统将自动调用这些运算符的前向和后向实现。
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) 和后备机制的内核注册到分派器中，当需要时，自动转换系统将自动调用这些内核。
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 中的数据类型，可能支持额外的
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类以实现与新后端对应的生成器类，并实现各种通用方法。定义一个参数为
device index的新后端builder。调用
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<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 扩展，例如 分布式集体通信、基准计时器 等，这些模块将在未来添加。关于 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")
未来工作¶
PrivateUse1 机制的改进仍在进行中，因此新模块的``PrivateUse1`` 集成方法将依次添加。以下是我们正在积极开展的工作项目：
添加
分布式集体通信的集成方法。添加
基准计时器的集成方法。
总结¶
本教程向您介绍了通过 PrivateUse1 集成新后端到 PyTorch 的过程，包括但不限于算子注册、生成器注册、设备保护注册等。同时介绍了一些方法以提升用户体验。
