|
|
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: |
|
|
|
|
|
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) |
|
|
|
