Spaces:
Sleeping
Sleeping
import torch | |
from torch._C import _rename_privateuse1_backend, _get_privateuse1_backend_name | |
from typing import List, Optional, Union | |
__all__ = ["rename_privateuse1_backend", "generate_methods_for_privateuse1_backend"] | |
# TODO: Should use `torch._C._get_privateuse1_backend_name()` to get | |
# renamed-backend name for `privateuse1`, but the func will cause an | |
# error with torch.jit.script, so we use the global variable named | |
# `_privateuse1_backend_name`. | |
_privateuse1_backend_name = "privateuseone" | |
def rename_privateuse1_backend(backend_name: str) -> None: | |
r""" | |
Rename the privateuse1 backend device to make it more convenient to use as a device name within PyTorch APIs. | |
The steps are: | |
(1) (In C++) implement kernels for various torch operations, and register them | |
to the PrivateUse1 dispatch key. | |
(2) (In python) call torch.utils.rename_privateuse1_backend("foo") | |
You can now use "foo" as an ordinary device string in python. | |
Note: this API can only be called once per process. Attempting to change | |
the external backend after it's already been set will result in an error. | |
Note(AMP): If you want to support AMP on your device, you can register a custom backend module. | |
The backend must register a custom backend module with ``torch._register_device_module("foo", BackendModule)``. | |
BackendModule needs to have the following API's: | |
(1) ``get_amp_supported_dtype() -> List[torch.dtype]`` | |
get the supported dtypes on your "foo" device in AMP, maybe the "foo" device supports one more dtype. | |
(2) ``is_autocast_enabled() -> bool`` | |
check the AMP is enabled or not on your "foo" device. | |
(3) ``get_autocast_dtype() -> torch.dtype`` | |
get the supported dtype on your "foo" device in AMP, which is set by ``set_autocast_dtype`` or the | |
default dtype, and the default dtype is ``torch.float16``. | |
(4) ``set_autocast_enabled(bool) -> None`` | |
enable the AMP or not on your "foo" device. | |
(5) ``set_autocast_dtype(dtype) -> None`` | |
set the supported dtype on your "foo" device in AMP, and the dtype be contained in the dtypes got | |
from ``get_amp_supported_dtype``. | |
Note(random): If you want to support to set seed for your device, BackendModule needs to have the following API's: | |
(1) ``_is_in_bad_fork() -> bool`` | |
Return ``True`` if now it is in bad_fork, else return ``False``. | |
(2) ``manual_seed_all(seed int) -> None`` | |
Sets the seed for generating random numbers for your devices. | |
(3) ``device_count() -> int`` | |
Returns the number of "foo"s available. | |
(4) ``get_rng_state(device: Union[int, str, torch.device] = 'foo') -> Tensor`` | |
Returns a list of ByteTensor representing the random number states of all devices. | |
(5) ``set_rng_state(new_state: Tensor, device: Union[int, str, torch.device] = 'foo') -> None`` | |
Sets the random number generator state of the specified "foo" device. | |
And there are some common funcs: | |
(1) ``is_available() -> bool`` | |
Returns a bool indicating if "foo" is currently available. | |
(2) ``current_device() -> int`` | |
Returns the index of a currently selected device. | |
For more details, see https://pytorch.org/tutorials/advanced/extend_dispatcher.html#get-a-dispatch-key-for-your-backend | |
For an existing example, see https://github.com/bdhirsh/pytorch_open_registration_example | |
Example:: | |
>>> # xdoctest: +SKIP("failing") | |
>>> torch.utils.rename_privateuse1_backend("foo") | |
# This will work, assuming that you've implemented the right C++ kernels | |
# to implement torch.ones. | |
>>> a = torch.ones(2, device="foo") | |
""" | |
_rename_privateuse1_backend(backend_name) | |
global _privateuse1_backend_name | |
_privateuse1_backend_name = backend_name | |
def _check_register_once(module, attr): | |
if hasattr(module, attr): | |
raise RuntimeError(f"The custom device module of {module} has already been registered with {attr}") | |
def _normalization_device(custom_backend_name: str, device: Optional[Union[int, str, torch.device]] = None) -> int: | |
def _get_current_device_index(): | |
_get_device_index = "current_device" | |
if hasattr(torch, custom_backend_name) and \ | |
hasattr(getattr(torch, custom_backend_name), _get_device_index): | |
return getattr(getattr(torch, custom_backend_name), _get_device_index)() | |
else: | |
# The default device index is 0. | |
return 0 | |
if device is None: | |
return _get_current_device_index() | |
# if isinstance(device, str), this means that the parameter passed in is in the string format "foo:0" | |
# convert str object to torch.device object, and then process it uniformly | |
elif isinstance(device, str): | |
device = torch.device(device) | |
# variable devcie can only be torch.device type or int type | |
if isinstance(device, torch.device): | |
if device.type != custom_backend_name: | |
raise RuntimeError(f"Invalid device, must be {custom_backend_name} device") | |
elif device.index is None: | |
device_idx = _get_current_device_index() | |
else: | |
device_idx = device.index | |
# if isinstance(device, int), we can take the index number directly | |
else: | |
device_idx = device | |
return device_idx | |
def _generate_tensor_methods_for_privateuse1_backend(custom_backend_name: str) -> None: | |
# type: ignore[misc] | |
def wrap_tensor_backend(self: torch.Tensor) -> bool: | |
return self.device.type == custom_backend_name | |
_check_register_once(torch.Tensor, f'is_{custom_backend_name}') | |
setattr(torch.Tensor, f'is_{custom_backend_name}', wrap_tensor_backend) | |
def wrap_tensor_to(self: torch.Tensor, device: Optional[Union[int, torch.device]] = None, non_blocking=False, | |
**kwargs) -> torch.Tensor: | |
r"""Perform Tensor device conversion. Call the to operator implementation. | |
.. note:: | |
If the ``self`` Tensor already | |
has the correct :class:`torch.device`, then ``self`` is returned. | |
Otherwise, the returned tensor is a copy of ``self`` with the desired :class:`torch.device`. | |
Args: | |
device (int, optional): if specified, all parameters will be copied to that device | |
non_blocking (bool): If ``True`` and the source is in pinned memory, | |
the copy will be asynchronous with respect to the host. Otherwise, | |
the argument has no effect. | |
**kwargs (dict): For compatibility, may contain the key ``memory_format`` argument. | |
""" | |
device_idx = _normalization_device(custom_backend_name, device) | |
return self.to(device=torch.device(f'{custom_backend_name}:{device_idx}'), non_blocking=non_blocking, **kwargs) | |
_check_register_once(torch.Tensor, custom_backend_name) | |
setattr(torch.Tensor, custom_backend_name, wrap_tensor_to) | |
def _generate_module_methods_for_privateuse1_backend(custom_backend_name: str) -> None: | |
# Generate Module attributes and methods depends on Tensor methods, | |
# so we need to check whether Tensor methods is already registered. | |
if not hasattr(torch.Tensor, custom_backend_name): | |
raise RuntimeError( | |
f"Can not automatically generate {custom_backend_name}() method for torch.nn.Module." | |
f"Because torch.Tensor doesn't has the method {custom_backend_name}()." | |
f"For this error, you can try setting for_tensor=True.") | |
def wrap_module_to(self: torch.nn.modules.module.T, | |
device: Optional[Union[int, torch.device]] = None) -> torch.nn.modules.module.T: | |
r"""Move all model parameters and buffers to the custom device. | |
This also makes associated parameters and buffers different objects. So | |
it should be called before constructing optimizer if the module will | |
live on device while being optimized. | |
.. note:: | |
This method modifies the module in-place. | |
Args: | |
device (int, optional): if specified, all parameters will be copied to that device | |
""" | |
return self._apply(lambda t: getattr(t, custom_backend_name)(device)) | |
_check_register_once(torch.nn.Module, custom_backend_name) | |
setattr(torch.nn.Module, custom_backend_name, wrap_module_to) | |
def _generate_storage_methods_for_privateuse1_backend(custom_backend_name: str, | |
unsupported_dtype: Optional[List[torch.dtype]] = None) -> None: | |
# Attribute is registered in the _StorageBase class | |
# and UntypedStorage obtains through inheritance. | |
# type: ignore[misc] | |
def wrap_storage_backend(self: torch.storage._StorageBase) -> bool: | |
r"""Return the internal :class:`torch.UntypedStorage`.""" | |
return self.device.type == custom_backend_name | |
_check_register_once(torch.storage._StorageBase, f'is_{custom_backend_name}') | |
setattr(torch.storage._StorageBase, f'is_{custom_backend_name}', wrap_storage_backend) | |
def wrap_storage_to(self, device=None, non_blocking=False): | |
r"""Return a copy of this object in custom device memory. | |
If this object is already in device memory and on the correct device, then | |
no copy is performed and the original object is returned. | |
Args: | |
device (int): The destination device id. Defaults to the current device. | |
non_blocking (bool): If ``True`` and the source is in pinned memory, | |
the copy will be asynchronous with respect to the host. Otherwise, | |
the argument has no effect. | |
""" | |
# There should be a judgment related to storage device and a judgment related to storage type, | |
# but it depends on the extended function, so this part is temporarily omitted in the automatic generation. | |
device_idx = _normalization_device(custom_backend_name, device) | |
if getattr(self, f'is_{custom_backend_name}'): | |
# storage has already on expected device. | |
if self.get_device() == device_idx: | |
return self | |
# For sparse storage, custom need to extend the implementation by themselves. | |
if self.is_sparse: | |
raise RuntimeError(f"Can not support a sparse storage move to {custom_backend_name} backend") | |
# create untyped_storage and copy data | |
untyped_storage = torch.UntypedStorage( | |
self.size(), device=torch.device(f'{custom_backend_name}:{device_idx}') | |
) | |
untyped_storage.copy_(self, non_blocking) | |
return untyped_storage | |
_check_register_once(torch.storage._StorageBase, custom_backend_name) | |
setattr(torch.storage._StorageBase, custom_backend_name, wrap_storage_to) | |
# Register the corresponding attribute for the TypedStorage class. | |
# When the TypedStorage class is removed, the registration is also removed. | |
# type: ignore[misc] | |
def wrap_typed_storage_backend(self: torch.storage.TypedStorage) -> bool: | |
torch.storage._warn_typed_storage_removal() | |
return self._untyped_storage.device.type == custom_backend_name | |
_check_register_once(torch.TypedStorage, f'is_{custom_backend_name}') | |
setattr(torch.storage.TypedStorage, f'is_{custom_backend_name}', wrap_typed_storage_backend) | |
def wrap_typed_storage_to(self: torch.storage.TypedStorage, | |
device=None, non_blocking=False, **kwargs) -> torch.storage.TypedStorage: | |
torch.storage._warn_typed_storage_removal() | |
if unsupported_dtype and self.dtype in unsupported_dtype: | |
raise RuntimeError(f"Cannot create {custom_backend_name} storage " | |
f"as {self.dtype} dtype is not supported by this backend") | |
custom_backend_storage: torch.UntypedStorage = getattr( | |
self._untyped_storage, custom_backend_name)(device, non_blocking, **kwargs) | |
return self._new_wrapped_storage(custom_backend_storage) | |
_check_register_once(torch.TypedStorage, custom_backend_name) | |
setattr(torch.TypedStorage, custom_backend_name, wrap_typed_storage_to) | |
def generate_methods_for_privateuse1_backend(for_tensor: bool = True, for_module: bool = True, | |
for_storage: bool = False, | |
unsupported_dtype: Optional[List[torch.dtype]] = None) -> None: | |
r""" | |
Automatically generate attributes and methods for the custom backend after rename privateuse1 backend. | |
In the default scenario, storage-related methods will not be generated automatically. | |
When you implement kernels for various torch operations, and register them to the PrivateUse1 dispatch key. | |
And call the function torch.rename_privateuse1_backend("foo") to rename your backend name. | |
At this point, you can easily register specific methods and attributes by calling this function. | |
Just like torch.Tensor.foo(), torch.Tensor.is_foo, torch.Storage.foo(), torch.Storage.is_foo. | |
Note: We recommend you use generic functions (check devices are equal or to(device=)). | |
We provide these methods for convenience only and they will be "monkey patched" onto the objects | |
and so will not be properly typed. For Storage methods generate, if you need to support sparse data storage, | |
you need to extend the implementation yourself. | |
Args: | |
for_tensor (bool): whether register related methods for torch.Tensor class. | |
for_module (bool): whether register related methods for torch.nn.Module class. | |
for_storage (bool): whether register related methods for torch.Storage class. | |
unsupported_dtype (List[torch.dtype]): takes effect only when the storage method needs to be generated, | |
indicating that the storage does not support the torch.dtype type. | |
Example:: | |
>>> # xdoctest: +SKIP("failing") | |
>>> torch.utils.rename_privateuse1_backend("foo") | |
>>> torch.utils.generate_methods_for_privateuse1_backend() | |
# Then automatically generate backend-related attributes and methods. | |
>>> a = torch.tensor(2).foo() | |
>>> a.is_foo | |
>>> hasattr(torch.nn.Module, 'foo') | |
""" | |
custom_backend_name = _get_privateuse1_backend_name() | |
if for_tensor: | |
_generate_tensor_methods_for_privateuse1_backend(custom_backend_name) | |
if for_module: | |
_generate_module_methods_for_privateuse1_backend(custom_backend_name) | |
if for_storage: | |
_generate_storage_methods_for_privateuse1_backend(custom_backend_name, unsupported_dtype) | |
def _get_custom_mod_func(func_name: str): | |
r""" | |
Return the func named `func_name` defined in custom device module. If not defined, | |
return `None`. And the func is registered with `torch.utils.rename_privateuse1_backend('foo')` | |
and `torch._register_device_module('foo', BackendModule)`. | |
If the custom device module or the func is not defined, it will give warning or error message. | |
Args: | |
func_name (str): return the callable func named func_name defined in custom device module. | |
Example:: | |
class DummyfooModule: | |
@staticmethod | |
def is_available(): | |
return True | |
@staticmethod | |
def func_name(*args, **kwargs): | |
.... | |
torch.utils.rename_privateuse1_backend("foo") | |
torch._register_device_module("foo", DummyfooModule) | |
foo_is_available_func = torch.utils.backend_registration._get_custom_mod_func("is_available") | |
if foo_is_available_func: | |
foo_is_available = foo_is_available_func() | |
func_ = torch.utils.backend_registration._get_custom_mod_func("func_name") | |
if func_: | |
result = func_(*args, **kwargs) | |
Attention: This function is not meant to be used directly by users, which is why | |
it is marked as private. It is a convenience function for backend implementers to | |
more easily call the hooks into their backend extensions. | |
""" | |
assert isinstance(func_name, str), f"func_name must be `str`, but got `{type(func_name)}`." | |
backend_name = _get_privateuse1_backend_name() | |
custom_device_mod = getattr(torch, backend_name, None) # type: ignore[arg-type] | |
function = getattr(custom_device_mod, func_name, None) # type: ignore[arg-type] | |
if custom_device_mod is None or function is None: | |
message = f'Try to call torch.{backend_name}.{func_name}. The backend must register a custom backend ' | |
message += f"module with `torch._register_device_module('{backend_name}', BackendModule)`. And " | |
message += f"BackendModule needs to have the following API's:\n `{func_name}(*args, **kwargs)`. \n" | |
raise RuntimeError(message) | |
return function | |