init

torch_vacc/vacc/streams.py

@@ -0,0 +1,327 @@
+import ctypes
+from typing import Any, Optional
+
+import torch
+from packaging import version
+from torch._utils import _get_device_index
+
+try:
+    from torch._streambase import _StreamBase, _EventBase
+except ImportError:
+    # torch <= 2.1
+    _StreamBase = _EventBase = object
+
+import torch_vacc
+
+from torch_vacc._vacc_libs import _torch_vacc
+from ._device import device
+from .lazy_initialize import _lazy_init
+
+
+# remove torch version arch-suffix(i.e. +cpu)
+torch_version = torch.__version__.split('+')[0]
+
+class _StreamCommon:
+    """Wrapper around a VACC stream.
+
+    A VACC stream is a linear sequence of execution that belongs to a specific
+    device, independent from other streams.
+
+    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. Can be either
+            -1 (high priority) or 0 (low priority). By default, streams have
+            priority 0.
+    """
+
+    def __new__(cls, device=None, priority=0, **kwargs):
+        if device is None or ("stream_id" in kwargs and "device_index" in kwargs):
+            return super(Stream, cls).__new__(cls, priority=priority, **kwargs)
+        else:
+            with torch_vacc.vacc.device(device):
+                return super(Stream, cls).__new__(cls, priority=priority, **kwargs)
+
+    def wait_event(self, event):
+        event.wait(self)
+
+    def record_event(self, event=None):
+        """Records an event.
+
+        Args:
+            event (torch_vacc.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 wait_stream(self, stream):
+        """Synchronizes 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.
+        """
+        self.wait_event(stream.record_event())
+
+    def query(self):
+        return super().query()
+
+    def synchronize(self):
+        super().synchronize()
+
+    @property
+    def _as_parameter_(self):
+        return ctypes.c_void_p(self.vacc_stream)
+
+    def __eq__(self, o):
+        if isinstance(o, Stream):
+            return super().__eq__(o)
+        return False
+
+    def __hash__(self):
+        return hash((self.vacc_stream, self.device))
+
+    def __repr__(self):
+        return f"torch_vacc.vacc.Stream device={self.device} vacc_stream={self.vacc_stream:#x}"
+
+if version.parse(torch_version) <= version.parse("2.1"):
+    # torch <= 2.1
+    class Stream(_torch_vacc._VACCStreamBase, _StreamCommon):
+        pass
+elif version.parse(torch_version) < version.parse("2.6"):
+    # torch < 2.6
+    class Stream(_torch_vacc._VACCStreamBase, _StreamBase, _StreamCommon):
+        pass
+else:
+    # torch >= 2.6
+    class Stream(_torch_vacc._VACCStreamBase, _StreamCommon):
+        pass
+
+
+class _EventCommon:
+    """Wrapper around a VACC event.
+
+    VACC events are synchronization markers that can be used to monitor the
+    device's progress, to accurately measure timing, and to synchronize VACC
+    streams.
+
+    The underlying VACC 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:
+        calc_time (bool, optional): indicates if the event should measure time
+            (default: ``False``)
+        blocking (bool, optional): if ``True``, :meth:`wait` will be blocking (default: ``False``)
+    """
+
+    def __new__(cls, enable_timing=False, blocking=False):
+        return super(Event, cls).__new__(
+            cls,
+            calc_time=enable_timing,
+            blocking=blocking,
+        )
+
+    def record(self, stream=None):
+        """Records the event in a given stream.
+
+        Uses ``torch_vacc.vacc.current_stream()`` if no stream is specified. The
+        stream's device must match the event's device."""
+        if stream is None:
+            stream = torch_vacc.vacc.current_stream()
+        super().record(stream)
+
+    def wait(self, stream=None):
+        """Makes all future work submitted to the given stream wait for this
+        event.
+
+        Use ``torch_vacc.vacc.current_stream()`` if no stream is specified.
+
+        .. note:: This is a wrapper around ``vaccrtStreamWaitEvent()``
+        """
+        if stream is None:
+            stream = torch_vacc.vacc.current_stream()
+        super().wait(stream)
+
+    def query(self):
+        """Checks 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):
+        """Returns the time elapsed in milliseconds after the event was
+        recorded and before the end_event was recorded.
+        """
+        return super().elapsed_time(end_event)
+
+    def synchronize(self):
+        r"""Waits 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 ``vaccEventSynchronize()``.
+        """
+        super().synchronize()
+
+    @property
+    def _as_parameter_(self):
+        return ctypes.c_void_p(self.vacc_event)
+
+    def __repr__(self):
+        if self.vacc_event:
+            return f"<torch_vacc.vacc.Event {self._as_parameter_.value:#x}>"
+        else:
+            return "<torch_vacc.vacc.Event uninitialized>"
+
+if version.parse(torch_version) <= version.parse("2.1"):
+    # torch <= 2.1
+    class Event(_torch_vacc._VACCEventBase, _EventCommon):
+        pass
+elif version.parse(torch_version) < version.parse("2.6"):
+    # torch < 2.6
+    class Event(_torch_vacc._VACCEventBase, _EventBase, _EventCommon):
+        pass
+else:
+    # torch >= 2.6
+    class Event(_torch_vacc._VACCEventBase, _EventCommon):
+        pass
+
+class StreamContext:
+    r"""Context-manager that selects a given stream.
+
+    All VACC kernels queued within its context will be enqueued on a selected
+    stream.
+
+    Args:
+        stream (stream): selected stream. This manager is a no-op if it's
+            ``None``.
+    .. note:: Streams are per-device.
+    """
+    cur_stream: Optional["torch_vacc.vacc.Stream"]
+
+    def __init__(self, stream: Optional["torch_vacc.vacc.Stream"]):
+        self.stream = stream
+        self.idx = _get_device_index(None, True)
+        if not torch.jit.is_scripting():
+            if self.idx is None:
+                self.idx = -1
+
+        self.src_prev_stream = (
+            None
+            if not torch.jit.is_scripting()
+            else torch_vacc.vacc.default_stream(None)
+        )
+        self.dst_prev_stream = (
+            None
+            if not torch.jit.is_scripting()
+            else torch_vacc.vacc.default_stream(None)
+        )
+
+    def __enter__(self):
+        # Local cur_stream variable for type refinement
+        cur_stream = self.stream
+        # Return if stream is None or VACC device not available
+        if cur_stream is None or self.idx == -1:
+            return
+        self.src_prev_stream = torch_vacc.vacc.current_stream(None)
+
+        # If the stream is not on the current device, then
+        # set the current stream on the device
+        if self.src_prev_stream.device != cur_stream.device:
+            with device(cur_stream.device):
+                self.dst_prev_stream = torch_vacc.vacc.current_stream(cur_stream.device)
+        torch_vacc.vacc.set_stream(cur_stream)
+
+    def __exit__(self, type: Any, value: Any, traceback: Any):
+        # Local cur_stream variable for type refinement
+        cur_stream = self.stream
+        # If stream is None or no VACC device available, return
+        if cur_stream is None or self.idx == -1:
+            return
+
+        # Reset the stream on the original device
+        # and destination device
+        if self.src_prev_stream.device != cur_stream.device:  # type: ignore[union-attr]
+            torch_vacc.vacc.set_stream(self.dst_prev_stream)  # type: ignore[arg-type]
+        torch_vacc.vacc.set_stream(self.src_prev_stream)  # type: ignore[arg-type]
+
+
+def stream(stream: Optional["torch_vacc.vacc.Stream"]) -> StreamContext:
+    r"""Wrapper around the Context-manager StreamContext that
+    selects a given stream.
+
+    Arguments:
+        stream (Stream): selected stream. This manager is a no-op if it's
+            ``None``.
+    """
+    return StreamContext(stream)
+
+
+def set_stream(stream: Stream):
+    r"""Sets the current stream.This is a wrapper API to set the stream.
+        Usage of this function is discouraged in favor of the ``stream``
+        context manager.
+
+    Args:
+        stream (Stream): selected stream. This function is a no-op
+            if this argument is ``None``.
+    """
+    if stream is None:
+        return
+    _torch_vacc._vacc_setStream(
+        stream_id=stream.stream_id,
+        device_index=stream.device_index,
+        device_type=stream.device_type,
+    )
+
+
+def current_stream(device=None) -> Stream:
+    r"""Returns the currently selected :class:`Stream` for a given device.
+
+    Args:
+        device (torch.device or int, optional): selected device. Returns
+            the currently selected :class:`Stream` for the current device, given
+            by :func:`~torch_vacc.vacc.current_device`, if :attr:`device` is ``None``
+            (default).
+    """
+    _lazy_init()
+    streamdata = _torch_vacc._vacc_getCurrentStream(
+        _get_device_index(device, optional=True)
+    )
+    return Stream(
+        stream_id=streamdata[0], device_index=streamdata[1], device_type=streamdata[2]
+    )
+
+
+def default_stream(device=None) -> Stream:
+    r"""Returns the default :class:`Stream` for a given device.
+
+    Args:
+        device (torch.device or int, optional): selected device. Returns
+            the default :class:`Stream` for the current device, given by
+            :func:`_torch_vacc.current_device`, if :attr:`device` is ``None``
+            (default).
+    """
+    _lazy_init()
+    streamdata = _torch_vacc._vacc_getDefaultStream(
+        _get_device_index(device, optional=True)
+    )
+
+    return Stream(
+        stream_id=streamdata[0], device_index=streamdata[1], device_type=streamdata[2]
+    )