[Device] Add support for PrivateUse1 device type in parse_type function (#157609)

This pull request refactors the `parse_type` function in `c10/core/Device.cpp` to improve the handling of the `PrivateUse1` device type. The main change involves reordering the logic to check for the `PrivateUse1` device type earlier in the function for better clarity and efficiency.

This help to migrate existed backend to PrivateUse1 smoothly.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/157609
Approved by: https://github.com/jgong5, https://github.com/albanD

c10/core/Device.cpp

@@ -41,6 +41,9 @@ DeviceType parse_type(const std::string& device_string) {
         "'mkldnn' is no longer used as device type. So torch.device('mkldnn') will be "
         "deprecated and removed in the future. Please use other valid device types instead.");
   }
+  if (device_string == get_privateuse1_backend()) {
+    return DeviceType::PrivateUse1;
+  }
   auto device = std::find_if(
       types.begin(),
       types.end(),
@@ -50,9 +53,6 @@ DeviceType parse_type(const std::string& device_string) {
   if (device != types.end()) {
     return device->second;
   }
-  if (device_string == get_privateuse1_backend()) {
-    return DeviceType::PrivateUse1;
-  }
   std::vector<const char*> device_names;
   for (const auto& it : types) {
     if (it.first) {

test/test_rename_privateuse1_to_existing_device.py

@@ -0,0 +1,59 @@
+# Owner(s): ["module: PrivateUse1"]
+
+import torch
+from torch.testing._internal.common_utils import run_tests, TestCase
+
+
+class DummyPrivateUse1Module:
+    @staticmethod
+    def is_available():
+        return True
+
+    @staticmethod
+    def is_autocast_enabled():
+        return True
+
+    @staticmethod
+    def get_autocast_dtype():
+        return torch.float16
+
+    @staticmethod
+    def set_autocast_enabled(enable):
+        pass
+
+    @staticmethod
+    def set_autocast_dtype(dtype):
+        pass
+
+    @staticmethod
+    def get_amp_supported_dtype():
+        return [torch.float16]
+
+
+class TestRenamePrivateuseoneToExistingBackend(TestCase):
+    def test_external_module_register_with_existing_backend(self):
+        torch.utils.rename_privateuse1_backend("maia")
+        with self.assertRaisesRegex(RuntimeError, "has already been set"):
+            torch.utils.rename_privateuse1_backend("dummmy")
+
+        custom_backend_name = torch._C._get_privateuse1_backend_name()
+        self.assertEqual(custom_backend_name, "maia")
+
+        with self.assertRaises(AttributeError):
+            torch.maia.is_available()
+
+        with self.assertRaisesRegex(AssertionError, "Tried to use AMP with the"):
+            with torch.autocast(device_type=custom_backend_name):
+                pass
+        torch._register_device_module("maia", DummyPrivateUse1Module)
+
+        torch.maia.is_available()  # type: ignore[attr-defined]
+        with torch.autocast(device_type=custom_backend_name):
+            pass
+
+        self.assertEqual(torch._utils._get_device_index("maia:1"), 1)
+        self.assertEqual(torch._utils._get_device_index(torch.device("maia:2")), 2)
+
+
+if __name__ == "__main__":
+    run_tests()