mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
c09eba877f
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
60 lines
1.7 KiB
Python
60 lines
1.7 KiB
Python
# 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()
|