Spaces:
Running
Running
File size: 6,954 Bytes
c61ccee |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 |
import inspect
from typing import Any, Callable, Dict, Iterable, Optional, Tuple, Type, Union
import torch
from torch._streambase import _EventBase, _StreamBase
get_cuda_stream: Optional[Callable[[int], int]]
if torch.cuda._is_compiled():
from torch._C import _cuda_getCurrentRawStream as get_cuda_stream
else:
get_cuda_stream = None
_device_t = Union[torch.device, str, int, None]
# Recording the device properties in the main process but used in worker process.
caching_worker_device_properties: Dict[str, Any] = {}
caching_worker_current_devices: Dict[str, int] = {}
class DeviceInterfaceMeta(type):
def __new__(metacls, *args, **kwargs):
class_member = args[2]
if "Event" in class_member:
assert inspect.isclass(class_member["Event"]) and issubclass(
class_member["Event"], _EventBase
), "DeviceInterface member Event should be inherit from _EventBase"
if "Stream" in class_member:
assert inspect.isclass(class_member["Stream"]) and issubclass(
class_member["Stream"], _StreamBase
), "DeviceInterface member Stream should be inherit from _StreamBase"
return super().__new__(metacls, *args, **kwargs)
class DeviceInterface(metaclass=DeviceInterfaceMeta):
"""
This is a simple device runtime interface for Inductor. It enables custom
backends to be integrated with Inductor in a device-agnostic semantic.
"""
class device:
def __new__(cls, device: _device_t):
raise NotImplementedError()
class Worker:
"""
Worker API to query device properties that will work in multi processing
workers that cannot use the GPU APIs (due to processing fork() and
initialization time issues). Properties are recorded in the main process
before we fork the workers.
"""
@staticmethod
def set_device(device: int):
raise NotImplementedError()
@staticmethod
def current_device() -> int:
raise NotImplementedError()
@staticmethod
def get_device_properties(device: _device_t = None):
raise NotImplementedError()
@staticmethod
def current_device():
raise NotImplementedError()
@staticmethod
def set_device(device: _device_t):
raise NotImplementedError()
@staticmethod
def device_count():
raise NotImplementedError()
@staticmethod
def is_available() -> bool:
raise NotImplementedError()
@staticmethod
def stream(stream: torch.Stream):
raise NotImplementedError()
@staticmethod
def current_stream():
raise NotImplementedError()
@staticmethod
def set_stream(stream: torch.Stream):
raise NotImplementedError()
@staticmethod
def _set_stream_by_id(stream_id: int, device_index: int, device_type: int):
raise NotImplementedError()
@staticmethod
def get_raw_stream():
raise NotImplementedError()
@staticmethod
def synchronize(device: _device_t = None):
raise NotImplementedError()
@staticmethod
def get_device_properties(device: _device_t = None):
raise NotImplementedError()
@staticmethod
def get_compute_capability(device: _device_t = None):
raise NotImplementedError()
class CudaInterface(DeviceInterface):
device = torch.cuda.device
# register Event and Stream class into the backend interface
# make sure Event and Stream are implemented and inherited from the _EventBase and _StreamBase
Event = torch.cuda.Event
Stream = torch.cuda.Stream
class Worker:
@staticmethod
def set_device(device: int):
caching_worker_current_devices["cuda"] = device
@staticmethod
def current_device() -> int:
if "cuda" in caching_worker_current_devices:
return caching_worker_current_devices["cuda"]
return torch.cuda.current_device()
@staticmethod
def get_device_properties(device: _device_t = None):
if device is not None:
if isinstance(device, str):
device = torch.device(device)
assert device.type == "cuda"
if isinstance(device, torch.device):
device = device.index
if device is None:
device = CudaInterface.Worker.current_device()
if "cuda" not in caching_worker_device_properties:
device_prop = [
torch.cuda.get_device_properties(i)
for i in range(torch.cuda.device_count())
]
caching_worker_device_properties["cuda"] = device_prop
return caching_worker_device_properties["cuda"][device]
current_device = staticmethod(torch.cuda.current_device)
set_device = staticmethod(torch.cuda.set_device)
device_count = staticmethod(torch.cuda.device_count)
stream = staticmethod(torch.cuda.stream) # type: ignore[assignment]
current_stream = staticmethod(torch.cuda.current_stream)
set_stream = staticmethod(torch.cuda.set_stream) # type: ignore[assignment]
_set_stream_by_id = staticmethod(torch.cuda._set_stream_by_id) # type: ignore[assignment]
synchronize = staticmethod(torch.cuda.synchronize)
get_device_properties = staticmethod(torch.cuda.get_device_properties) # type: ignore[assignment]
get_raw_stream = staticmethod(get_cuda_stream) # type: ignore[arg-type]
# Can be mock patched by @patch decorator.
@staticmethod
def is_available() -> bool:
return torch.cuda.is_available()
@staticmethod
def get_compute_capability(device: _device_t = None):
major, min = torch.cuda.get_device_capability(device)
return major * 10 + min
device_interfaces: Dict[str, Type[DeviceInterface]] = {}
def register_interface_for_device(
device: Union[str, torch.device], device_interface: Type[DeviceInterface]
):
if isinstance(device, torch.device):
device = str(device)
device_interfaces[device] = device_interface
def get_interface_for_device(device: Union[str, torch.device]) -> Type[DeviceInterface]:
if isinstance(device, torch.device):
device = str(device)
if device in device_interfaces:
return device_interfaces[device]
raise NotImplementedError(f"No interface for device {device}")
def get_registered_device_interfaces() -> Iterable[Tuple[str, Type[DeviceInterface]]]:
return device_interfaces.items()
register_interface_for_device("cuda", CudaInterface)
for i in range(torch.cuda.device_count()):
register_interface_for_device(f"cuda:{i}", CudaInterface)
|