numba-cuda 0.19.1__py3-none-any.whl → 0.20.1__py3-none-any.whl

numba_cuda/numba/cuda/core/event.py

@@ -0,0 +1,511 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+The ``numba.cuda.core.event`` module provides a simple event system for applications
+to register callbacks to listen to specific compiler events.
+
+The following events are built in:
+
+- ``"numba-cuda:compile"`` is broadcast when a dispatcher is compiling. Events of
+  this kind have ``data`` defined to be a ``dict`` with the following
+  key-values:
+
+  - ``"dispatcher"``: the dispatcher object that is compiling.
+  - ``"args"``: the argument types.
+  - ``"return_type"``: the return type.
+
+- ``"numba-cuda:compiler_lock"`` is broadcast when the internal compiler-lock is
+  acquired. This is mostly used internally to measure time spent with the lock
+  acquired.
+
+- ``"numba-cuda:run_pass"`` is broadcast when a compiler pass is running.
+
+    - ``"name"``: pass name.
+    - ``"qualname"``: qualified name of the function being compiled.
+    - ``"module"``: module name of the function being compiled.
+    - ``"flags"``: compilation flags.
+    - ``"args"``: argument types.
+    - ``"return_type"`` return type.
+
+Applications can register callbacks that are listening for specific events using
+``register(kind: str, listener: Listener)``, where ``listener`` is an instance
+of ``Listener`` that defines custom actions on occurrence of the specific event.
+"""
+
+import os
+import json
+import atexit
+import abc
+import enum
+import time
+import threading
+from timeit import default_timer as timer
+from contextlib import contextmanager, ExitStack
+from collections import defaultdict
+from numba.core import config
+from numba.cuda import utils
+
+
+class EventStatus(enum.Enum):
+    """Status of an event."""
+
+    START = enum.auto()
+    END = enum.auto()
+
+
+# Builtin event kinds.
+_builtin_kinds = frozenset(
+    [
+        "numba-cuda:compiler_lock",
+        "numba-cuda:compile",
+        "numba-cuda:run_pass",
+    ]
+)
+
+
+def _guard_kind(kind):
+    """Guard to ensure that an event kind is valid.
+
+    All event kinds with a "numba-cuda:" prefix must be defined in the pre-defined
+    ``numba.cuda.core.event._builtin_kinds``.
+    Custom event kinds are allowed by not using the above prefix.
+
+    Parameters
+    ----------
+    kind : str
+
+    Return
+    ------
+    res : str
+    """
+    if kind.startswith("numba-cuda:") and kind not in _builtin_kinds:
+        msg = (
+            f"{kind} is not a valid event kind, "
+            "it starts with the reserved prefix 'numba-cuda:'"
+        )
+        raise ValueError(msg)
+    if kind.startswith("numba:"):
+        msg = (
+            f"{kind} is using invalid reserved prefix 'numba:' "
+            "it should either start with the reserved prefix 'numba-cuda:'"
+            ", or be a custom event kind"
+        )
+        raise ValueError(msg)
+    return kind
+
+
+class Event:
+    """An event.
+
+    Parameters
+    ----------
+    kind : str
+    status : EventStatus
+    data : any; optional
+        Additional data for the event.
+    exc_details : 3-tuple; optional
+        Same 3-tuple for ``__exit__``.
+    """
+
+    def __init__(self, kind, status, data=None, exc_details=None):
+        self._kind = _guard_kind(kind)
+        self._status = status
+        self._data = data
+        self._exc_details = (
+            None
+            if exc_details is None or exc_details[0] is None
+            else exc_details
+        )
+
+    @property
+    def kind(self):
+        """Event kind
+
+        Returns
+        -------
+        res : str
+        """
+        return self._kind
+
+    @property
+    def status(self):
+        """Event status
+
+        Returns
+        -------
+        res : EventStatus
+        """
+        return self._status
+
+    @property
+    def data(self):
+        """Event data
+
+        Returns
+        -------
+        res : object
+        """
+        return self._data
+
+    @property
+    def is_start(self):
+        """Is it a *START* event?
+
+        Returns
+        -------
+        res : bool
+        """
+        return self._status == EventStatus.START
+
+    @property
+    def is_end(self):
+        """Is it an *END* event?
+
+        Returns
+        -------
+        res : bool
+        """
+        return self._status == EventStatus.END
+
+    @property
+    def is_failed(self):
+        """Is the event carrying an exception?
+
+        This is used for *END* event. This method will never return ``True``
+        in a *START* event.
+
+        Returns
+        -------
+        res : bool
+        """
+        return self._exc_details is None
+
+    def __str__(self):
+        data = (
+            f"{type(self.data).__qualname__}"
+            if self.data is not None
+            else "None"
+        )
+        return f"Event({self._kind}, {self._status}, data: {data})"
+
+    __repr__ = __str__
+
+
+_registered = defaultdict(list)
+
+
+def register(kind, listener):
+    """Register a listener for a given event kind.
+
+    Parameters
+    ----------
+    kind : str
+    listener : Listener
+    """
+    assert isinstance(listener, Listener)
+    kind = _guard_kind(kind)
+    _registered[kind].append(listener)
+
+
+def unregister(kind, listener):
+    """Unregister a listener for a given event kind.
+
+    Parameters
+    ----------
+    kind : str
+    listener : Listener
+    """
+    assert isinstance(listener, Listener)
+    kind = _guard_kind(kind)
+    lst = _registered[kind]
+    lst.remove(listener)
+
+
+def broadcast(event):
+    """Broadcast an event to all registered listeners.
+
+    Parameters
+    ----------
+    event : Event
+    """
+    for listener in _registered[event.kind]:
+        listener.notify(event)
+
+
+class Listener(abc.ABC):
+    """Base class for all event listeners."""
+
+    @abc.abstractmethod
+    def on_start(self, event):
+        """Called when there is a *START* event.
+
+        Parameters
+        ----------
+        event : Event
+        """
+        pass
+
+    @abc.abstractmethod
+    def on_end(self, event):
+        """Called when there is a *END* event.
+
+        Parameters
+        ----------
+        event : Event
+        """
+        pass
+
+    def notify(self, event):
+        """Notify this Listener with the given Event.
+
+        Parameters
+        ----------
+        event : Event
+        """
+        if event.is_start:
+            self.on_start(event)
+        elif event.is_end:
+            self.on_end(event)
+        else:
+            raise AssertionError("unreachable")
+
+
+class TimingListener(Listener):
+    """A listener that measures the total time spent between *START* and
+    *END* events during the time this listener is active.
+    """
+
+    def __init__(self):
+        self._depth = 0
+
+    def on_start(self, event):
+        if self._depth == 0:
+            self._ts = timer()
+        self._depth += 1
+
+    def on_end(self, event):
+        self._depth -= 1
+        if self._depth == 0:
+            last = getattr(self, "_duration", 0)
+            self._duration = (timer() - self._ts) + last
+
+    @property
+    def done(self):
+        """Returns a ``bool`` indicating whether a measurement has been made.
+
+        When this returns ``False``, the matching event has never fired.
+        If and only if this returns ``True``, ``.duration`` can be read without
+        error.
+        """
+        return hasattr(self, "_duration")
+
+    @property
+    def duration(self):
+        """Returns the measured duration.
+
+        This may raise ``AttributeError``. Users can use ``.done`` to check
+        that a measurement has been made.
+        """
+        return self._duration
+
+
+class RecordingListener(Listener):
+    """A listener that records all events and stores them in the ``.buffer``
+    attribute as a list of 2-tuple ``(float, Event)``, where the first element
+    is the time the event occurred as returned by ``time.time()`` and the second
+    element is the event.
+    """
+
+    def __init__(self):
+        self.buffer = []
+
+    def on_start(self, event):
+        self.buffer.append((time.time(), event))
+
+    def on_end(self, event):
+        self.buffer.append((time.time(), event))
+
+
+@contextmanager
+def install_listener(kind, listener):
+    """Install a listener for event "kind" temporarily within the duration of
+    the context.
+
+    Returns
+    -------
+    res : Listener
+        The *listener* provided.
+
+    Examples
+    --------
+
+    >>> with install_listener("numba-cuda:compile", listener):
+    >>>     some_code()  # listener will be active here.
+    >>> other_code()  # listener will be unregistered by this point.
+
+    """
+    register(kind, listener)
+    try:
+        yield listener
+    finally:
+        unregister(kind, listener)
+
+
+@contextmanager
+def install_timer(kind, callback):
+    """Install a TimingListener temporarily to measure the duration of
+    an event.
+
+    If the context completes successfully, the *callback* function is executed.
+    The *callback* function is expected to take a float argument for the
+    duration in seconds.
+
+    Returns
+    -------
+    res : TimingListener
+
+    Examples
+    --------
+
+    This is equivalent to:
+
+    >>> with install_listener(kind, TimingListener()) as res:
+    >>>    ...
+    """
+    tl = TimingListener()
+    with install_listener(kind, tl):
+        yield tl
+
+    if tl.done:
+        callback(tl.duration)
+
+
+@contextmanager
+def install_recorder(kind):
+    """Install a RecordingListener temporarily to record all events.
+
+    Once the context is closed, users can use ``RecordingListener.buffer``
+    to access the recorded events.
+
+    Returns
+    -------
+    res : RecordingListener
+
+    Examples
+    --------
+
+    This is equivalent to:
+
+    >>> with install_listener(kind, RecordingListener()) as res:
+    >>>    ...
+    """
+    rl = RecordingListener()
+    with install_listener(kind, rl):
+        yield rl
+
+
+def start_event(kind, data=None):
+    """Trigger the start of an event of *kind* with *data*.
+
+    Parameters
+    ----------
+    kind : str
+        Event kind.
+    data : any; optional
+        Extra event data.
+    """
+    evt = Event(kind=kind, status=EventStatus.START, data=data)
+    broadcast(evt)
+
+
+def end_event(kind, data=None, exc_details=None):
+    """Trigger the end of an event of *kind*, *exc_details*.
+
+    Parameters
+    ----------
+    kind : str
+        Event kind.
+    data : any; optional
+        Extra event data.
+    exc_details : 3-tuple; optional
+        Same 3-tuple for ``__exit__``. Or, ``None`` if no error.
+    """
+    evt = Event(
+        kind=kind,
+        status=EventStatus.END,
+        data=data,
+        exc_details=exc_details,
+    )
+    broadcast(evt)
+
+
+@contextmanager
+def trigger_event(kind, data=None):
+    """A context manager to trigger the start and end events of *kind* with
+    *data*. The start event is triggered when entering the context.
+    The end event is triggered when exiting the context.
+
+    Parameters
+    ----------
+    kind : str
+        Event kind.
+    data : any; optional
+        Extra event data.
+    """
+    with ExitStack() as scope:
+
+        @scope.push
+        def on_exit(*exc_details):
+            end_event(kind, data=data, exc_details=exc_details)
+
+        start_event(kind, data=data)
+        yield
+
+
+def _prepare_chrome_trace_data(listener: RecordingListener):
+    """Prepare events in `listener` for serializing as chrome trace data."""
+    # The spec for the trace event format can be found at:
+    # https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/edit   # noqa
+    # This code only uses the JSON Array Format for simplicity.
+    pid = os.getpid()
+    tid = threading.get_native_id()
+    evs = []
+    for ts, rec in listener.buffer:
+        data = rec.data
+        cat = str(rec.kind)
+        ts_scaled = ts * 1_000_000  # scale to microseconds
+        ph = "B" if rec.is_start else "E"
+        name = data["name"]
+        args = data
+        ev = dict(
+            cat=cat,
+            pid=pid,
+            tid=tid,
+            ts=ts_scaled,
+            ph=ph,
+            name=name,
+            args=args,
+        )
+        evs.append(ev)
+    return evs
+
+
+def _setup_chrome_trace_exit_handler():
+    """Setup a RecordingListener and an exit handler to write the captured
+    events to file.
+    """
+    listener = RecordingListener()
+    register("numba-cuda:run_pass", listener)
+    filename = config.CHROME_TRACE
+
+    @atexit.register
+    def _write_chrome_trace():
+        # The following output file is not multi-process safe.
+        evs = _prepare_chrome_trace_data(listener)
+        with open(filename, "w") as out:
+            json.dump(evs, out, cls=utils._LazyJSONEncoder)
+
+
+if config.CHROME_TRACE:
+    _setup_chrome_trace_exit_handler()