Enable torch.jit.load for custom device (#99535)

Fixes #ISSUE_NUMBER
1、torch.jit.load for custom device
```
# custom device named `foo`
ts_model = torch.jit.script(mode.to(device="foo"))
ts_model.save("./ts.pt") # it is a script model on device `foo`

# and then we want to load it and run it
torch.jit.load("./ts.pt")
```
2、 add some extra key for custom device with `privateuse1`
Pull Request resolved: #99535
Approved by: https://github.com/albanD

c10/core/Backend.h

@@ -114,7 +114,8 @@ static inline Backend dispatchKeyToBackend(DispatchKey t) {
     return Backend::HPU;
   } else if (t == DispatchKey::MTIA || t == DispatchKey::AutogradMTIA) {
     return Backend::MTIA;
-  } else if (t == DispatchKey::PrivateUse1) {
+  } else if (
+      t == DispatchKey::PrivateUse1 || t == DispatchKey::AutogradPrivateUse1) {
     return Backend::PrivateUse1;
   } else if (t == DispatchKey::Undefined) {
     return Backend::Undefined;

c10/core/Device.h

@@ -81,6 +81,11 @@ struct C10_API Device final {
     return type_ == DeviceType::CUDA;
   }
 
+  /// Return true if the device is of PrivateUse1 type.
+  bool is_privateuseone() const noexcept {
+    return type_ == DeviceType::PrivateUse1;
+  }
+
   /// Return true if the device is of MPS type.
   bool is_mps() const noexcept {
     return type_ == DeviceType::MPS;

torch/csrc/autograd/init.cpp

@@ -175,7 +175,8 @@ PyObject* THPAutograd_initExtension(PyObject* _unused, PyObject* unused) {
       .value("HPU", c10::DeviceType::HPU)
       .value("VE", c10::DeviceType::VE)
       .value("Lazy", c10::DeviceType::Lazy)
-      .value("IPU", c10::DeviceType::IPU);
+      .value("IPU", c10::DeviceType::IPU)
+      .value("PrivateUse1", c10::DeviceType::PrivateUse1);
 
   py::class_<KinetoEvent>(m, "_KinetoEvent")
       // name of the event

torch/csrc/jit/serialization/unpickler.cpp

@@ -579,11 +579,13 @@ PickleOpCode Unpickler::readInstruction() {
       }
 
       if (device.is_cuda() || device.is_xpu() || device.is_meta() ||
-          device.is_hpu()) {
+          device.is_hpu() || device.is_privateuseone()) {
         tensor = tensor.to(device, tensor.scalar_type());
       } else if (device.type() != DeviceType::CPU) {
         AT_ERROR(
-            "supported devices include CPU, CUDA and HPU, however got ",
+            "supported devices include CPU, CUDA, HPU and ",
+            c10::get_privateuse1_backend(),
+            " however got ",
             DeviceTypeName(device.type(), false));
       }
       stack_.emplace_back(std::move(tensor));

torch/csrc/utils/python_dispatch.cpp

@@ -547,6 +547,7 @@ void initDispatchBindings(PyObject* module) {
       DEF_ONE(AutocastXPU)
       DEF_ONE(AutocastHPU)
       DEF_ONE(AutocastCUDA)
+      DEF_ONE(AutocastPrivateUse1)
   // clang-format on
 
 #define DEF_SINGLE(n, prefix) .value(#prefix #n, c10::DispatchKey::prefix##n)