numba-cuda 0.21.1__cp313-cp313-win_amd64.whl

numba_cuda/numba/cuda/cudadrv/devices.py

@@ -0,0 +1,249 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Expose each GPU devices directly.
+
+This module implements a API that is like the "CUDA runtime" context manager
+for managing CUDA context stack and clean up.  It relies on thread-local globals
+to separate the context stack management of each thread. Contexts are also
+shareable among threads.  Only the main thread can destroy Contexts.
+
+Note:
+- This module must be imported by the main-thread.
+
+"""
+
+import functools
+import threading
+from contextlib import contextmanager
+
+from .driver import driver
+
+
+class _DeviceList:
+    @property
+    @functools.cache
+    def lst(self):
+        return [
+            _DeviceContextManager(driver.get_device(devid))
+            for devid in range(driver.get_device_count())
+        ]
+
+    def __getitem__(self, devnum):
+        """
+        Returns the context manager for device *devnum*.
+        """
+        if not isinstance(devnum, (int, slice)):
+            devnum = int(devnum)
+        return self.lst[devnum]
+
+    def __str__(self):
+        return ", ".join([str(d) for d in self.lst])
+
+    def __iter__(self):
+        return iter(self.lst)
+
+    def __len__(self):
+        return len(self.lst)
+
+    @property
+    def current(self):
+        """Returns the active device or None if there's no active device"""
+        with driver.get_active_context() as ac:
+            devnum = ac.devnum
+            if devnum is not None:
+                return self[devnum]
+
+
+class _DeviceContextManager(object):
+    """
+    Provides a context manager for executing in the context of the chosen
+    device. The normal use of instances of this type is from
+    ``numba.cuda.gpus``. For example, to execute on device 2::
+
+       with numba.cuda.gpus[2]:
+           d_a = numba.cuda.to_device(a)
+
+    to copy the array *a* onto device 2, referred to by *d_a*.
+    """
+
+    def __init__(self, device):
+        self._device = device
+        # Forwarded directly, to avoid the performance overhead of
+        # `__getattr__` and method lookup for a commonly accessed method
+        self.get_primary_context = self._device.get_primary_context
+
+    def __getattr__(self, item):
+        return getattr(self._device, item)
+
+    def __enter__(self):
+        _runtime.get_or_create_context(self._device.id)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # this will verify that we are popping the right device context.
+        self.get_primary_context().pop()
+
+    def __str__(self):
+        return f"<Managed Device {self.id}>"
+
+
+class _Runtime(object):
+    """Emulate the CUDA runtime context management.
+
+    It owns all Devices and Contexts.
+    Keeps at most one Context per Device
+    """
+
+    def __init__(self):
+        self.gpus = _DeviceList()
+
+        # For caching the attached CUDA Context
+        self._tls = threading.local()
+
+        # Remember the main thread
+        # Only the main thread can *actually* destroy
+        self._mainthread = threading.current_thread()
+
+        # Avoid mutation of runtime state in multithreaded programs
+        self._lock = threading.RLock()
+
+    @contextmanager
+    def ensure_context(self):
+        """Ensure a CUDA context is available inside the context.
+
+        On entrance, queries the CUDA driver for an active CUDA context and
+        attaches it in TLS for subsequent calls so they do not need to query
+        the CUDA driver again.  On exit, detach the CUDA context from the TLS.
+
+        This will allow us to pickup thirdparty activated CUDA context in
+        any top-level Numba CUDA API.
+        """
+        with driver.get_active_context():
+            oldctx = self._get_attached_context()
+            newctx = self.get_or_create_context(None)
+            self._set_attached_context(newctx)
+            try:
+                yield
+            finally:
+                self._set_attached_context(oldctx)
+
+    def get_or_create_context(self, devnum):
+        """Returns the primary context and push+create it if needed
+        for *devnum*.  If *devnum* is None, use the active CUDA context (must
+        be primary) or create a new one with ``devnum=0``.
+        """
+        if devnum is None:
+            attached_ctx = self._get_attached_context()
+            if attached_ctx is None:
+                return self._get_or_create_context_uncached(devnum)
+            else:
+                return attached_ctx
+        else:
+            devnum = int(devnum)
+            with self._lock:
+                return self._activate_context_for(devnum)
+
+    def _get_or_create_context_uncached(self, devnum):
+        """See also ``get_or_create_context(devnum)``.
+        This version does not read the cache.
+        """
+        with self._lock:
+            # Try to get the active context in the CUDA stack or
+            # activate GPU-0 with the primary context
+            with driver.get_active_context() as ac:
+                if not ac:
+                    return self._activate_context_for(0)
+                else:
+                    # Get primary context for the active device
+                    ctx = self.gpus[ac.devnum].get_primary_context()
+                    # Is active context the primary context?
+                    ctx_handle = ctx.handle.value
+                    ac_ctx_handle = ac.context_handle.value
+                    if ctx_handle != ac_ctx_handle:
+                        raise RuntimeError(
+                            "Numba cannot operate on non-primary"
+                            f" CUDA context {ac_ctx_handle:x}"
+                        )
+                    # Ensure the context is ready
+                    ctx.prepare_for_use()
+                return ctx
+
+    def _activate_context_for(self, devnum):
+        gpu = self.gpus[devnum]
+        newctx = gpu.get_primary_context()
+        # Detect unexpected context switch
+        cached_ctx = self._get_attached_context()
+        if cached_ctx is not None and cached_ctx is not newctx:
+            raise RuntimeError("Cannot switch CUDA-context.")
+        newctx.push()
+        return newctx
+
+    def _get_attached_context(self):
+        try:
+            return self._tls.attached_context
+        except AttributeError:
+            return None
+
+    def _set_attached_context(self, ctx):
+        self._tls.attached_context = ctx
+
+    def reset(self):
+        """Clear all contexts in the thread.  Destroy the context if and only
+        if we are in the main thread.
+        """
+        # Pop all active context.
+        while driver.pop_active_context() is not None:
+            pass
+
+        # If it is the main thread
+        if threading.current_thread() == self._mainthread:
+            self._destroy_all_contexts()
+
+    def _destroy_all_contexts(self):
+        # Reset all devices
+        for gpu in self.gpus:
+            gpu.reset()
+
+
+_runtime = _Runtime()
+
+# ================================ PUBLIC API ================================
+
+gpus = _runtime.gpus
+
+
+def get_context(devnum=None):
+    """Get the current device or use a device by device number, and
+    return the CUDA context.
+    """
+    return _runtime.get_or_create_context(devnum)
+
+
+def require_context(fn):
+    """
+    A decorator that ensures a CUDA context is available when *fn* is executed.
+
+    Note: The function *fn* cannot switch CUDA-context.
+    """
+
+    @functools.wraps(fn)
+    def _require_cuda_context(*args, **kws):
+        with _runtime.ensure_context():
+            return fn(*args, **kws)
+
+    return _require_cuda_context
+
+
+def reset():
+    """Reset the CUDA subsystem for the current thread.
+
+    In the main thread:
+    This removes all CUDA contexts.  Only use this at shutdown or for
+    cleaning up between tests.
+
+    In non-main threads:
+    This clear the CUDA context stack only.
+
+    """
+    _runtime.reset()