Spaces:
Running
Running
import ctypes | |
import torch | |
from torch._streambase import _EventBase, _StreamBase | |
from .._utils import _dummy_type | |
if not hasattr(torch._C, "_CudaStreamBase"): | |
# Define dummy base classes | |
torch._C.__dict__["_CudaStreamBase"] = _dummy_type("_CudaStreamBase") | |
torch._C.__dict__["_CudaEventBase"] = _dummy_type("_CudaEventBase") | |
class Stream(torch._C._CudaStreamBase, _StreamBase): | |
r"""Wrapper around a CUDA stream. | |
A CUDA stream is a linear sequence of execution that belongs to a specific | |
device, independent from other streams. See :ref:`cuda-semantics` for | |
details. | |
Args: | |
device(torch.device or int, optional): a device on which to allocate | |
the stream. If :attr:`device` is ``None`` (default) or a negative | |
integer, this will use the current device. | |
priority(int, optional): priority of the stream, should be 0 or | |
negative, where negative numbers indicate higher priority. By default, | |
streams have priority 0. | |
""" | |
def __new__(cls, device=None, priority=0, **kwargs): | |
# setting device manager is expensive, so we avoid it unless necessary | |
if device is None or ("stream_id" in kwargs and "device_index" in kwargs): | |
return super().__new__(cls, priority=priority, **kwargs) | |
else: | |
with torch.cuda.device(device): | |
return super().__new__(cls, priority=priority, **kwargs) | |
def wait_event(self, event): | |
r"""Make all future work submitted to the stream wait for an event. | |
Args: | |
event (torch.cuda.Event): an event to wait for. | |
.. note:: This is a wrapper around ``cudaStreamWaitEvent()``: see | |
`CUDA Stream documentation`_ for more info. | |
This function returns without waiting for :attr:`event`: only future | |
operations are affected. | |
.. _CUDA Stream documentation: | |
https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__STREAM.html | |
""" | |
event.wait(self) | |
def wait_stream(self, stream): | |
r"""Synchronize with another stream. | |
All future work submitted to this stream will wait until all kernels | |
submitted to a given stream at the time of call complete. | |
Args: | |
stream (Stream): a stream to synchronize. | |
.. note:: This function returns without waiting for currently enqueued | |
kernels in :attr:`stream`: only future operations are affected. | |
""" | |
self.wait_event(stream.record_event()) | |
def record_event(self, event=None): | |
r"""Record an event. | |
Args: | |
event (torch.cuda.Event, optional): event to record. If not given, a new one | |
will be allocated. | |
Returns: | |
Recorded event. | |
""" | |
if event is None: | |
event = Event() | |
event.record(self) | |
return event | |
def query(self): | |
r"""Check if all the work submitted has been completed. | |
Returns: | |
A boolean indicating if all kernels in this stream are completed. | |
""" | |
return super().query() | |
def synchronize(self): | |
r"""Wait for all the kernels in this stream to complete. | |
.. note:: This is a wrapper around ``cudaStreamSynchronize()``: see | |
`CUDA Stream documentation`_ for more info. | |
""" | |
super().synchronize() | |
def _as_parameter_(self): | |
return ctypes.c_void_p(self.cuda_stream) | |
def __eq__(self, o): | |
if isinstance(o, Stream): | |
return super().__eq__(o) | |
return False | |
def __hash__(self): | |
return hash((self.cuda_stream, self.device)) | |
def __repr__(self): | |
return f"<torch.cuda.Stream device={self.device} cuda_stream={self.cuda_stream:#x}>" | |
class ExternalStream(Stream): | |
r"""Wrapper around an externally allocated CUDA stream. | |
This class is used to wrap streams allocated in other libraries in order | |
to facilitate data exchange and multi-library interactions. | |
.. note:: This class doesn't manage the stream life-cycle, it is the user | |
responsibility to keep the referenced stream alive while this class is | |
being used. | |
Args: | |
stream_ptr(int): Integer representation of the `cudaStream_t` value. | |
allocated externally. | |
device(torch.device or int, optional): the device where the stream | |
was originally allocated. if device is specified incorrectly, | |
subsequent launches using this stream may fail. | |
""" | |
def __new__(cls, stream_ptr, device=None, **kwargs): | |
with torch.cuda.device(device): | |
return super().__new__(cls, stream_ptr=stream_ptr, **kwargs) | |
class Event(torch._C._CudaEventBase, _EventBase): | |
r"""Wrapper around a CUDA event. | |
CUDA events are synchronization markers that can be used to monitor the | |
device's progress, to accurately measure timing, and to synchronize CUDA | |
streams. | |
The underlying CUDA events are lazily initialized when the event is first | |
recorded or exported to another process. After creation, only streams on the | |
same device may record the event. However, streams on any device can wait on | |
the event. | |
Args: | |
enable_timing (bool, optional): indicates if the event should measure time | |
(default: ``False``) | |
blocking (bool, optional): if ``True``, :meth:`wait` will be blocking (default: ``False``) | |
interprocess (bool): if ``True``, the event can be shared between processes | |
(default: ``False``) | |
.. _CUDA Event Documentation: | |
https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__EVENT.html | |
""" | |
def __new__(cls, enable_timing=False, blocking=False, interprocess=False): | |
return super().__new__( | |
cls, | |
enable_timing=enable_timing, | |
blocking=blocking, | |
interprocess=interprocess, | |
) | |
def from_ipc_handle(cls, device, handle): | |
r"""Reconstruct an event from an IPC handle on the given device.""" | |
return super().from_ipc_handle(device, handle) | |
def record(self, stream=None): | |
r"""Record the event in a given stream. | |
Uses ``torch.cuda.current_stream()`` if no stream is specified. The | |
stream's device must match the event's device. | |
""" | |
if stream is None: | |
stream = torch.cuda.current_stream() | |
super().record(stream) | |
def wait(self, stream=None): | |
r"""Make all future work submitted to the given stream wait for this event. | |
Use ``torch.cuda.current_stream()`` if no stream is specified. | |
.. note:: This is a wrapper around ``cudaStreamWaitEvent()``: see | |
`CUDA Event documentation`_ for more info. | |
""" | |
if stream is None: | |
stream = torch.cuda.current_stream() | |
super().wait(stream) | |
def query(self): | |
r"""Check if all work currently captured by event has completed. | |
Returns: | |
A boolean indicating if all work currently captured by event has | |
completed. | |
""" | |
return super().query() | |
def elapsed_time(self, end_event): | |
r"""Return the time elapsed. | |
Time reported in milliseconds after the event was recorded and | |
before the end_event was recorded. | |
""" | |
return super().elapsed_time(end_event) | |
def synchronize(self): | |
r"""Wait for the event to complete. | |
Waits until the completion of all work currently captured in this event. | |
This prevents the CPU thread from proceeding until the event completes. | |
.. note:: This is a wrapper around ``cudaEventSynchronize()``: see | |
`CUDA Event documentation`_ for more info. | |
""" | |
super().synchronize() | |
def ipc_handle(self): | |
r"""Return an IPC handle of this event. | |
If not recorded yet, the event will use the current device. | |
""" | |
return super().ipc_handle() | |
def _as_parameter_(self): | |
return ctypes.c_void_p(self.cuda_event) | |
def __repr__(self): | |
if self.cuda_event: | |
return f"<torch.cuda.Event {self._as_parameter_.value:#x}>" | |
else: | |
return "<torch.cuda.Event uninitialized>" | |