File size: 1,730 Bytes
9c6594c |
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 |
import torch
class Event:
r"""Wrapper around an MPS event.
MPS events are synchronization markers that can be used to monitor the
device's progress, to accurately measure timing, and to synchronize MPS streams.
Args:
enable_timing (bool, optional): indicates if the event should measure time
(default: ``False``)
"""
def __init__(self, enable_timing: bool = False) -> None:
self.__eventId = torch._C._mps_acquireEvent(enable_timing)
def __del__(self) -> None:
# checks if torch._C is already destroyed
if hasattr(torch._C, "_mps_releaseEvent") and self.__eventId > 0:
torch._C._mps_releaseEvent(self.__eventId)
def record(self) -> None:
r"""Records the event in the default stream."""
torch._C._mps_recordEvent(self.__eventId)
def wait(self) -> None:
r"""Makes all future work submitted to the default stream wait for this event."""
torch._C._mps_waitForEvent(self.__eventId)
def query(self) -> bool:
r"""Returns True if all work currently captured by event has completed."""
return torch._C._mps_queryEvent(self.__eventId)
def synchronize(self) -> None:
r"""Waits until the completion of all work currently captured in this event.
This prevents the CPU thread from proceeding until the event completes.
"""
torch._C._mps_synchronizeEvent(self.__eventId)
def elapsed_time(self, end_event: "Event") -> float:
r"""Returns the time elapsed in milliseconds after the event was
recorded and before the end_event was recorded.
"""
return torch._C._mps_elapsedTimeOfEvents(self.__eventId, end_event.__eventId)
|