通过 PrivateUse1 简化新后端的集成

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 cpu
at::native::cpu_fallback(op,stack);
}

TORCH_LIBRARY_IMPL(_,PrivateUse1,m){
m.fallback(torch::CppFunction::makeFromBoxedFunction<&custom_cpu_fallback>());
}

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));
...
}

TORCH_LIBRARY_IMPL(aten,AutocastPrivateUse1,m){
...
KERNEL_PRIVATEUSEONE(<operator>,<policy>)
...
}

TORCH_LIBRARY_IMPL(_,AutocastPrivateUse1,m){
m.fallback(torch::CppFunction::makeFallthrough());
}

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)

structCustomGuardImplfinal:publicc10::impl::DeviceGuardImplInterface{
// Implementation of guard in new backend
}

C10_REGISTER_GUARD_IMPL(PrivateUse1,CustomGuardImpl);

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);

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

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

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)

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