Register kernels for the new backend#

The new backend may have some high-performance implementations of operator, which can be registered to the dispatcherbyTORCH_LIBRARY_IMPL API described inRegistering a Dispatched Operator in C++. This involvesseveral situations:

1. Register all the forward operators supported by the new backend to the dispatcher, and register the fallbackat the same time, so that when the new backend does not support some operators, these operators can fall backto the CPU for execution to ensure the availability of functions.

at::Tensorwrapper_Custom_Tensor_add(constat::Tensor&self,constat::Tensor&other,constat::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));...}voidcustom_cpu_fallback(constc10::OperatorHandle&op,torch::jit::Stack*stack){// Add some hints about new devices that do not support and need to fall back to cpuat::native::cpu_fallback(op,stack);}TORCH_LIBRARY_IMPL(_,PrivateUse1,m){m.fallback(torch::CppFunction::makeFromBoxedFunction<&custom_cpu_fallback>());}

2. Register kernels fromtorch::autograd::Function to the dispatcher byAutogradPrivateUse1, if it is necessary fornew backend to overridePyTorchAutogradlayer, the dispatcher and autograd system will automatically call the forward andbackward implementations of these operators.

classCumtomSeluFunction:publictorch::autograd::Function<CumtomSeluFunction>{// Implementation of selu kernel in new backend}at::Tensorwrapper_AutogradCumstom__selu(constat::Tensor&self){returnCumtomSeluFunction::apply(self);}TORCH_LIBRARY_IMPL(aten,AutogradPrivateUse1,m){...m.impl("selu",TORCH_FN(wrapper_AutogradCustom__selu));...}

3. Register kernels which want to supportautomatic mixed precision (AMP) andfallback mechanism to the dispatcher byAutocastPrivateUse1, the autocast system will automatically call these kernels when needed.

TORCH_LIBRARY_IMPL(aten,AutocastPrivateUse1,m){...KERNEL_PRIVATEUSEONE(<operator>,<policy>)...}TORCH_LIBRARY_IMPL(_,AutocastPrivateUse1,m){m.fallback(torch::CppFunction::makeFallthrough());}

What needs to be added is that if you want to support AMP in a new backend, you need to register a newBackendModule bytorch._register_device_module("backend_name",BackendModule), and theBackendModule needs to have the following APIs:

• get_amp_supported_dtype()->List[torch.dtype]

  get the supported dtypes on the new backend in AMP, which might support one moredtype.

• is_autocast_enabled()->bool

  check the AMP is enabled or not on the new backend.

• get_autocast_dtype()->torch.dtype

  get the supporteddtype on the new backend in AMP, which is set byset_autocast_dtype or thedefaultdtype, and the defaultdtype istorch.float16.

• set_autocast_enabled(bool)->None

  enable or disable AMP on the new backend.

• set_autocast_dtype(dtype)->None

  set the supporteddtype on the new backend in AMP, and thedtype be contained in thedtypes gotfromget_amp_supported_dtype.

Register generator for the new backend#

It is necessary to support generators corresponding to new devices. Currently,PrivateUse1 can dynamicallyregister custom generators, which are mainly divided into the following steps.

1. Inherit theGeneratorImpl class to implement the generator class corresponding to the new backend,and implement various general methods.

2. Define a new backendbuilder with a single parameter:deviceindex.

3. CallREGISTER_GENERATOR_PRIVATEUSE1 macro to complete dynamic registration.

structCustomGeneratorImpl:publicc10::GeneratorImpl{// Implementation of generator in new backend}at::Generatormake_custom_generator(c10::DeviceIndexdevice_index){returnat::make_generator<CustomGeneratorImpl>(device_index);}REGISTER_GENERATOR_PRIVATEUSE1(make_cumstom_generator)

Register device guard for the new backend#

PyTorch provides functionalities related to device, stream, and event switching viaDeviceGuard.This function is also applicable toPrivateUse1 Key.

1. Inherit theDeviceGuardImplInterface class to implement the various general methods corresponding to the new backend.

2. CallC10_REGISTER_GUARD_IMPL macro to complete dynamic registration.

structCustomGuardImplfinal:publicc10::impl::DeviceGuardImplInterface{// Implementation of guard in new backend}C10_REGISTER_GUARD_IMPL(PrivateUse1,CustomGuardImpl);

Register serialization and deserialization functions for new backend metadata#

PyTorch is currently able to dynamically register serialization/deserialization functions to support the serialization and deserializationof new backend additional metadata namedbackend_meta_ in classTensorImpl.ExtraMeta. You can refer to the following steps:

1. Inherit theBackendMeta class to implementCustomBackendMetadata corresponding to the new backend andvarious fields of the new backend can be customized in the class.

2. Implement the serialization and deserialization functions of the new backend, the function signatures arevoid(constat::Tensor&,std::unordered_map<std::string,bool>&).

3. Call theTensorBackendMetaRegistry macro to complete dynamic registration.

structCustomBackendMetadata:publicc10::BackendMeta{// Implementation of backend metadata in new backend}voidfor_serialization(constat::Tensor&t,std::unordered_map<std::string,bool>&m){// Implementation of serialization}voidfor_deserialization(constat::Tensor&t,std::unordered_map<std::string,bool>&m){// Implementation of deserialization}TensorBackendMetaRegistry(c10::DeviceType::PrivateUse1,&for_serialization,&for_deserialization);

Other Modules#

In addition to the above-mentioned parts, there are some other modules that can be expanded throughPrivateUse1,such asdistributedcollectivecommunication,benchmarktimer, and others, which will be added in the future.One example aboutPrivateUse1 integration isAscend NPU.

Register new backend module to Pytorch#

Some CUDA-related interfaces in PyTorch can be called through the following form:torch.cuda.xxx. Therefore, in order tocomply with user habits, the new backend implemented through thePrivateUse1 mechanism should also provide similar interfaces.

For example, usingAscendNPU:

torch._register_device_module('npu',torch_npu.npu)

After doing the above operations, users can call some exclusive APIs ofAscendNPU throughtorch.npu.xxx

Rename PrivateUse1 to a custom name for the new backend#

PrivateUse1 Key is the internal mechanism of the new backend integrated into PyTorch. For users, compared withPrivateUse1,the custom name strongly related to the new backend should be more friendly.

Taking theAscendNPU as an example, the first usage will be more user-friendly.

torch.rand((2,2),device='npu:0')torch.rand((2,2),device='privateuse1:0')

Now, PyTorch provides a new C++/Python API for the self-namedPrivateUse1 backend, which is very simple to use.

torch.rename_privateuse1_backend("npu")

c10::register_privateuse1_backend("npu")

Generate methods and properties related to the new backend#

After renamingPrivateUse1 to a custome name, automatically generate properties and methods related to the new backend namein theTensor,nn,Storage modules for the new backend.

Here is an example forAscendNPU:

torch.rename_privateuse1_backend("npu")unsupported_dtype=[torch.quint8]torch.utils.generate_methods_for_privateuse1_backend(for_tensor=True,for_module=True,for_storage=True,unsupported_dtype=unsupported_dtype)

Then, you can use the following methods and properties:

torch.Tensor.npu()torch.Tensor.is_nputorch.Storage.npu()torch.Storage.is_npu...