"""Functions and classes for getting and setting computing devices.
"""
import numpy as np
from dosma.utils import env
if env.cupy_available():
import cupy as cp
if env.sigpy_available():
import sigpy as sp
if env.torch_available():
import torch
__all__ = ["Device", "get_device", "to_device"]
[docs]class Device(object):
"""Device class.
This class extends ``cupy.Device`` and can also be used to interface with
``torch.Device`` and ``sigpy.Device``. This class contains a device type
('cpu' or 'cuda') and optional device ordinal (i.e. the id) for the device type.
This class can also be constructed using only the ordinal id, where id >= 0
representing the id-th GPU, and id = -1 representing CPU. cupy must be installed to use GPUs.
The array module for the corresponding device can be obtained via ``device.xp``.
Similar to cupy.Device, the Device object can be used as a context:
>>> device = Device(1) # gpu 1
>>> xp = device.xp # xp is cupy.
>>> with device:
>>> x = xp.array([1, 2, 3])
>>> x += 1
Args:
id_or_device (int or Device or cupy.cuda.Device): id > 0 represents
the corresponding GPUs, and id = -1 represents CPU.
Attributes:
type (str): Device type. Either ``'cpu'`` or ``'cuda'``.
index (int): index = -1 represents CPU, and others represents the id_th ordinances.
Note:
This class is heavily based on
`sigpy.Device <https://sigpy.readthedocs.io/en/latest/generated/sigpy.Device.html>`_.
"""
[docs] def __init__(self, id_or_device):
_type, id = None, None
if isinstance(id_or_device, int):
id = id_or_device
elif id_or_device == "cpu":
_type, id = id_or_device, -1
elif isinstance(id_or_device, Device):
_type, id = id_or_device.type, id_or_device.id
elif env.cupy_available() and isinstance(id_or_device, cp.cuda.Device):
_type, id = "cuda", id_or_device.id
elif env.sigpy_available() and isinstance(id_or_device, sp.Device):
id = id_or_device.id
elif env.torch_available() and isinstance(id_or_device, torch.device):
_type, id = id_or_device.type, id_or_device.index
if id is None:
if _type == "cuda":
id = torch.cuda.current_device()
elif _type == "cpu":
id = -1
else:
raise ValueError(f"Unsupported device type: {_type}")
else:
raise ValueError(
f"Accepts int, Device, cupy.cuda.Device, or torch.device" f"got {id_or_device}"
)
assert id >= -1
if _type is None:
_type = "cpu" if id == -1 else "cuda"
cpdevice = None
if id != -1:
if env.cupy_available():
cpdevice = cp.cuda.Device(id)
else:
raise ValueError("cupy not installed, but set device {}.".format(id))
self._type = _type
self._id = id
self._cpdevice = cpdevice
@property
def id(self):
"""int: The device ordinal."""
return self._id
@property
def type(self):
"""str: Type of device. Either ``"cpu"`` or ``"cuda"``."""
return self._type
@property
def index(self):
"""int: Alias for ``self.id``."""
return self.id
@property
def cpdevice(self):
"""cupy.Device: The equivalent ```cupy.Device```."""
return self._cpdevice
@property
def ptdevice(self):
"""torch.device: The equivalent ```torch.device```."""
if not env.torch_available():
raise RuntimeError("`torch` not installed.")
if self.id == -1:
return torch.device("cpu")
return torch.device(f"{self.type}:{self.id}")
@property
def spdevice(self):
"""sigpy.Device: The equivalent ```sigpy.Device```."""
if not env.sigpy_available():
raise RuntimeError("`sigpy` not installed.")
if self.id >= 0 and self.type != "cuda":
raise RuntimeError(f"sigpy.Device does not support type {self.type}")
return sp.Device(self.id)
@property
def xp(self):
"""module: numpy or cupy module for the device."""
if self.id == -1:
return np
return cp
def use(self):
"""Use computing device.
All operations after use() will use the device when using ``cupy``.
"""
if self.id > 0:
self.cpdevice.use()
def __int__(self):
return self.id
def __eq__(self, other):
if other == -1:
# other integers are not compared as self.type may be subject to change
return self.id == other
elif isinstance(other, Device):
return self.type == other.type and self.id == other.id
elif env.cupy_available() and isinstance(other, cp.cuda.Device):
return self.type == "cuda" and self.id == other.id
elif env.sigpy_available() and isinstance(other, sp.Device):
try:
return self.spdevice == other
except RuntimeError:
return False
elif env.torch_available() and isinstance(other, torch.device):
return self.ptdevice == other
else:
return False
def __ne__(self, other):
return not self == other
def __enter__(self):
if self.id == -1:
return None
return self.cpdevice.__enter__()
def __exit__(self, *args):
if self.id == -1:
pass
else:
self.cpdevice.__exit__()
def __repr__(self):
if self.id == -1:
return "Device(type='cpu')"
return f"Device(type='{self.type}', index={self.id})"
cpu_device = Device(-1)
def get_array_module(array):
"""Gets an appropriate module from :mod:`numpy` or :mod:`cupy`.
This is almost equivalent to :func:`cupy.get_array_module`. The differences
are that this function can be used even if cupy is not available.
Adapted from :mod:`sigpy`.
Args:
array: Input array.
Returns:
module: :mod:`cupy` or :mod:`numpy` is returned based on input.
"""
if env.cupy_available():
return cp.get_array_module(array)
else:
return np
[docs]def get_device(array):
"""Get Device from input array.
Adapted from :mod:`sigpy`.
Args:
array (array): Array.
Returns:
Device.
"""
if hasattr(array, "device"):
return Device(array.device)
else:
return cpu_device
[docs]def to_device(input, device=cpu_device):
"""Move input to device. Does not copy if same device.
Adapted from :mod:`sigpy`.
Args:
input (array): Input.
device (int or Device or cupy.Device): Output device.
Returns:
array: Output array placed in device.
"""
idevice = get_device(input)
odevice = Device(device)
if idevice == odevice:
return input
if odevice == cpu_device:
with idevice:
return input.get()
else:
with odevice:
return cp.asarray(input)