numba-cuda 0.8.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

numba_cuda/numba/cuda/cudadrv/driver.py

@@ -10,6 +10,7 @@ subsequent deallocation could further corrupt the CUDA context and causes the
 system to freeze in some cases.
 
 """
+
 import sys
 import os
 import ctypes
@@ -25,8 +26,17 @@ import tempfile
 import re
 from itertools import product
 from abc import ABCMeta, abstractmethod
-from ctypes import (c_int, byref, c_size_t, c_char, c_char_p, addressof,
-                    c_void_p, c_float, c_uint)
+from ctypes import (
+    c_int,
+    byref,
+    c_size_t,
+    c_char,
+    c_char_p,
+    addressof,
+    c_void_p,
+    c_float,
+    c_uint,
+)
 import contextlib
 import importlib
 import numpy as np
@@ -51,13 +61,14 @@ USE_NV_BINDING = config.CUDA_USE_NVIDIA_BINDING
 
 if USE_NV_BINDING:
     from cuda import cuda as binding
+
     # There is no definition of the default stream in the Nvidia bindings (nor
     # is there at the C/C++ level), so we define it here so we don't need to
     # use a magic number 0 in places where we want the default stream.
     CU_STREAM_DEFAULT = 0
 
 MIN_REQUIRED_CC = (3, 5)
-SUPPORTS_IPC = sys.platform.startswith('linux')
+SUPPORTS_IPC = sys.platform.startswith("linux")
 
 
 _py_decref = ctypes.pythonapi.Py_DecRef
@@ -71,10 +82,9 @@ _MVC_ERROR_MESSAGE = (
     "to be available"
 )
 
-ENABLE_PYNVJITLINK = (
-    _readenv("NUMBA_CUDA_ENABLE_PYNVJITLINK", bool, False)
-    or getattr(config, "CUDA_ENABLE_PYNVJITLINK", False)
-)
+ENABLE_PYNVJITLINK = _readenv(
+    "NUMBA_CUDA_ENABLE_PYNVJITLINK", bool, False
+) or getattr(config, "CUDA_ENABLE_PYNVJITLINK", False)
 if not hasattr(config, "CUDA_ENABLE_PYNVJITLINK"):
     config.CUDA_ENABLE_PYNVJITLINK = ENABLE_PYNVJITLINK
 
@@ -94,7 +104,7 @@ def make_logger():
         if config.CUDA_LOG_LEVEL:
             # create a simple handler that prints to stderr
             handler = logging.StreamHandler(sys.stderr)
-            fmt = '== CUDA [%(relativeCreated)d] %(levelname)5s -- %(message)s'
+            fmt = "== CUDA [%(relativeCreated)d] %(levelname)5s -- %(message)s"
             handler.setFormatter(logging.Formatter(fmt=fmt))
             logger.addHandler(handler)
         else:
@@ -122,50 +132,52 @@ class CudaAPIError(CudaDriverError):
 
 
 def locate_driver_and_loader():
-
     envpath = config.CUDA_DRIVER
 
-    if envpath == '0':
+    if envpath == "0":
         # Force fail
         _raise_driver_not_found()
 
     # Determine DLL type
-    if sys.platform == 'win32':
+    if sys.platform == "win32":
         dlloader = ctypes.WinDLL
-        dldir = ['\\windows\\system32']
-        dlnames = ['nvcuda.dll']
-    elif sys.platform == 'darwin':
+        dldir = ["\\windows\\system32"]
+        dlnames = ["nvcuda.dll"]
+    elif sys.platform == "darwin":
         dlloader = ctypes.CDLL
-        dldir = ['/usr/local/cuda/lib']
-        dlnames = ['libcuda.dylib']
+        dldir = ["/usr/local/cuda/lib"]
+        dlnames = ["libcuda.dylib"]
     else:
         # Assume to be *nix like
         dlloader = ctypes.CDLL
-        dldir = ['/usr/lib', '/usr/lib64']
-        dlnames = ['libcuda.so', 'libcuda.so.1']
+        dldir = ["/usr/lib", "/usr/lib64"]
+        dlnames = ["libcuda.so", "libcuda.so.1"]
 
     if envpath:
         try:
             envpath = os.path.abspath(envpath)
         except ValueError:
-            raise ValueError("NUMBA_CUDA_DRIVER %s is not a valid path" %
-                             envpath)
+            raise ValueError(
+                "NUMBA_CUDA_DRIVER %s is not a valid path" % envpath
+            )
         if not os.path.isfile(envpath):
-            raise ValueError("NUMBA_CUDA_DRIVER %s is not a valid file "
-                             "path.  Note it must be a filepath of the .so/"
-                             ".dll/.dylib or the driver" % envpath)
+            raise ValueError(
+                "NUMBA_CUDA_DRIVER %s is not a valid file "
+                "path.  Note it must be a filepath of the .so/"
+                ".dll/.dylib or the driver" % envpath
+            )
         candidates = [envpath]
     else:
         # First search for the name in the default library path.
         # If that is not found, try the specific path.
-        candidates = dlnames + [os.path.join(x, y)
-                                for x, y in product(dldir, dlnames)]
+        candidates = dlnames + [
+            os.path.join(x, y) for x, y in product(dldir, dlnames)
+        ]
 
     return dlloader, candidates
 
 
 def load_driver(dlloader, candidates):
-
     # Load the driver; Collect driver error information
     path_not_exist = []
     driver_load_error = []
@@ -184,7 +196,7 @@ def load_driver(dlloader, candidates):
     if all(path_not_exist):
         _raise_driver_not_found()
     else:
-        errmsg = '\n'.join(str(e) for e in driver_load_error)
+        errmsg = "\n".join(str(e) for e in driver_load_error)
         _raise_driver_error(errmsg)
 
 
@@ -216,7 +228,7 @@ def _raise_driver_error(e):
 
 
 def _build_reverse_error_map():
-    prefix = 'CUDA_ERROR'
+    prefix = "CUDA_ERROR"
     map = utils.UniqueDict()
     for name in dir(enums):
         if name.startswith(prefix):
@@ -236,6 +248,7 @@ class Driver(object):
     """
     Driver API functions are lazily bound.
     """
+
     _singleton = None
 
     def __new__(cls):
@@ -254,9 +267,11 @@ class Driver(object):
         self.pid = None
         try:
             if config.DISABLE_CUDA:
-                msg = ("CUDA is disabled due to setting NUMBA_DISABLE_CUDA=1 "
-                       "in the environment, or because CUDA is unsupported on "
-                       "32-bit systems.")
+                msg = (
+                    "CUDA is disabled due to setting NUMBA_DISABLE_CUDA=1 "
+                    "in the environment, or because CUDA is unsupported on "
+                    "32-bit systems."
+                )
                 raise CudaSupportError(msg)
             self.lib = find_driver()
         except CudaSupportError as e:
@@ -273,7 +288,7 @@ class Driver(object):
 
         self.is_initialized = True
         try:
-            _logger.info('init')
+            _logger.info("init")
             self.cuInit(0)
         except CudaAPIError as e:
             description = f"{e.msg} ({e.code})"
@@ -292,8 +307,9 @@ class Driver(object):
         self.ensure_initialized()
 
         if self.initialization_error is not None:
-            raise CudaSupportError("Error at driver init: \n%s:" %
-                                   self.initialization_error)
+            raise CudaSupportError(
+                "Error at driver init: \n%s:" % self.initialization_error
+            )
 
         if USE_NV_BINDING:
             return self._cuda_python_wrap_fn(fname)
@@ -317,12 +333,12 @@ class Driver(object):
 
         def verbose_cuda_api_call(*args):
             argstr = ", ".join([str(arg) for arg in args])
-            _logger.debug('call driver api: %s(%s)', libfn.__name__, argstr)
+            _logger.debug("call driver api: %s(%s)", libfn.__name__, argstr)
             retcode = libfn(*args)
             self._check_ctypes_error(fname, retcode)
 
         def safe_cuda_api_call(*args):
-            _logger.debug('call driver api: %s', libfn.__name__)
+            _logger.debug("call driver api: %s", libfn.__name__)
             retcode = libfn(*args)
             self._check_ctypes_error(fname, retcode)
 
@@ -340,11 +356,11 @@ class Driver(object):
 
         def verbose_cuda_api_call(*args):
             argstr = ", ".join([str(arg) for arg in args])
-            _logger.debug('call driver api: %s(%s)', libfn.__name__, argstr)
+            _logger.debug("call driver api: %s(%s)", libfn.__name__, argstr)
             return self._check_cuda_python_error(fname, libfn(*args))
 
         def safe_cuda_api_call(*args):
-            _logger.debug('call driver api: %s', libfn.__name__)
+            _logger.debug("call driver api: %s", libfn.__name__)
             return self._check_cuda_python_error(fname, libfn(*args))
 
         if config.CUDA_LOG_API_ARGS:
@@ -361,27 +377,30 @@ class Driver(object):
         # binding. For the NVidia binding, it handles linking to the correct
         # variant.
         if config.CUDA_PER_THREAD_DEFAULT_STREAM and not USE_NV_BINDING:
-            variants = ('_v2_ptds', '_v2_ptsz', '_ptds', '_ptsz', '_v2', '')
+            variants = ("_v2_ptds", "_v2_ptsz", "_ptds", "_ptsz", "_v2", "")
         else:
-            variants = ('_v2', '')
+            variants = ("_v2", "")
+
+        if fname in ("cuCtxGetDevice", "cuCtxSynchronize"):
+            return getattr(self.lib, fname)
 
         for variant in variants:
             try:
-                return getattr(self.lib, f'{fname}{variant}')
+                return getattr(self.lib, f"{fname}{variant}")
             except AttributeError:
                 pass
 
         # Not found.
         # Delay missing function error to use
         def absent_function(*args, **kws):
-            raise CudaDriverError(f'Driver missing function: {fname}')
+            raise CudaDriverError(f"Driver missing function: {fname}")
 
         setattr(self, fname, absent_function)
         return absent_function
 
     def _detect_fork(self):
         if self.pid is not None and _getpid() != self.pid:
-            msg = 'pid %s forked from pid %s after CUDA driver init'
+            msg = "pid %s forked from pid %s after CUDA driver init"
             _logger.critical(msg, _getpid(), self.pid)
             raise CudaDriverError("CUDA initialized before forking")
 
@@ -425,13 +444,11 @@ class Driver(object):
         return count.value
 
     def list_devices(self):
-        """Returns a list of active devices
-        """
+        """Returns a list of active devices"""
         return list(self.devices.values())
 
     def reset(self):
-        """Reset all devices
-        """
+        """Reset all devices"""
         for dev in self.devices.values():
             dev.reset()
 
@@ -449,8 +466,7 @@ class Driver(object):
                     return popped
 
     def get_active_context(self):
-        """Returns an instance of ``_ActiveContext``.
-        """
+        """Returns an instance of ``_ActiveContext``."""
         return _ActiveContext()
 
     def get_version(self):
@@ -477,12 +493,13 @@ class _ActiveContext(object):
     Once entering the context, it is assumed that the active CUDA context is
     not changed until the context is exited.
     """
+
     _tls_cache = threading.local()
 
     def __enter__(self):
         is_top = False
         # check TLS cache
-        if hasattr(self._tls_cache, 'ctx_devnum'):
+        if hasattr(self._tls_cache, "ctx_devnum"):
             hctx, devnum = self._tls_cache.ctx_devnum
         # Not cached. Query the driver API.
         else:
@@ -515,11 +532,10 @@ class _ActiveContext(object):
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         if self._is_top:
-            delattr(self._tls_cache, 'ctx_devnum')
+            delattr(self._tls_cache, "ctx_devnum")
 
     def __bool__(self):
-        """Returns True is there's a valid and active CUDA context.
-        """
+        """Returns True is there's a valid and active CUDA context."""
         return self.context_handle is not None
 
     __nonzero__ = __bool__
@@ -533,7 +549,7 @@ def _build_reverse_device_attrs():
     map = utils.UniqueDict()
     for name in dir(enums):
         if name.startswith(prefix):
-            map[name[len(prefix):]] = getattr(enums, name)
+            map[name[len(prefix) :]] = getattr(enums, name)
     return map
 
 
@@ -545,6 +561,7 @@ class Device(object):
     The device object owns the CUDA contexts.  This is owned by the driver
     object.  User should not construct devices directly.
     """
+
     @classmethod
     def from_identity(self, identity):
         """Create Device object from device identity created by
@@ -579,15 +596,17 @@ class Device(object):
         self.attributes = {}
 
         # Read compute capability
-        self.compute_capability = (self.COMPUTE_CAPABILITY_MAJOR,
-                                   self.COMPUTE_CAPABILITY_MINOR)
+        self.compute_capability = (
+            self.COMPUTE_CAPABILITY_MAJOR,
+            self.COMPUTE_CAPABILITY_MINOR,
+        )
 
         # Read name
         bufsz = 128
 
         if USE_NV_BINDING:
             buf = driver.cuDeviceGetName(bufsz, self.id)
-            name = buf.decode('utf-8').rstrip('\0')
+            name = buf.decode("utf-8").rstrip("\0")
         else:
             buf = (c_char * bufsz)()
             driver.cuDeviceGetName(buf, bufsz, self.id)
@@ -604,31 +623,31 @@ class Device(object):
             driver.cuDeviceGetUuid(byref(uuid), self.id)
             uuid_vals = tuple(bytes(uuid))
 
-        b = '%02x'
+        b = "%02x"
         b2 = b * 2
         b4 = b * 4
         b6 = b * 6
-        fmt = f'GPU-{b4}-{b2}-{b2}-{b2}-{b6}'
+        fmt = f"GPU-{b4}-{b2}-{b2}-{b2}-{b6}"
         self.uuid = fmt % uuid_vals
 
         self.primary_context = None
 
     def get_device_identity(self):
         return {
-            'pci_domain_id': self.PCI_DOMAIN_ID,
-            'pci_bus_id': self.PCI_BUS_ID,
-            'pci_device_id': self.PCI_DEVICE_ID,
+            "pci_domain_id": self.PCI_DOMAIN_ID,
+            "pci_bus_id": self.PCI_BUS_ID,
+            "pci_device_id": self.PCI_DEVICE_ID,
         }
 
     def __repr__(self):
         return "<CUDA device %d '%s'>" % (self.id, self.name)
 
     def __getattr__(self, attr):
-        """Read attributes lazily
-        """
+        """Read attributes lazily"""
         if USE_NV_BINDING:
-            code = getattr(binding.CUdevice_attribute,
-                           f'CU_DEVICE_ATTRIBUTE_{attr}')
+            code = getattr(
+                binding.CUdevice_attribute, f"CU_DEVICE_ATTRIBUTE_{attr}"
+            )
             value = driver.cuDeviceGetAttribute(code, self.id)
         else:
             try:
@@ -698,17 +717,18 @@ class Device(object):
 
 def met_requirement_for_device(device):
     if device.compute_capability < MIN_REQUIRED_CC:
-        raise CudaSupportError("%s has compute capability < %s" %
-                               (device, MIN_REQUIRED_CC))
+        raise CudaSupportError(
+            "%s has compute capability < %s" % (device, MIN_REQUIRED_CC)
+        )
 
 
 class BaseCUDAMemoryManager(object, metaclass=ABCMeta):
     """Abstract base class for External Memory Management (EMM) Plugins."""
 
     def __init__(self, *args, **kwargs):
-        if 'context' not in kwargs:
+        if "context" not in kwargs:
             raise RuntimeError("Memory manager requires a context")
-        self.context = kwargs.pop('context')
+        self.context = kwargs.pop("context")
 
     @abstractmethod
     def memalloc(self, size):
@@ -864,8 +884,7 @@ class HostOnlyCUDAMemoryManager(BaseCUDAMemoryManager):
             else:
                 raise
 
-    def memhostalloc(self, size, mapped=False, portable=False,
-                     wc=False):
+    def memhostalloc(self, size, mapped=False, portable=False, wc=False):
         """Implements the allocation of pinned host memory.
 
         It is recommended that this method is not overridden by EMM Plugin
@@ -880,6 +899,7 @@ class HostOnlyCUDAMemoryManager(BaseCUDAMemoryManager):
             flags |= enums.CU_MEMHOSTALLOC_WRITECOMBINED
 
         if USE_NV_BINDING:
+
             def allocator():
                 return driver.cuMemHostAlloc(size, flags)
 
@@ -946,16 +966,19 @@ class HostOnlyCUDAMemoryManager(BaseCUDAMemoryManager):
         ctx = weakref.proxy(self.context)
 
         if mapped:
-            mem = MappedMemory(ctx, pointer, size, owner=owner,
-                               finalizer=finalizer)
+            mem = MappedMemory(
+                ctx, pointer, size, owner=owner, finalizer=finalizer
+            )
             self.allocations[alloc_key] = mem
             return mem.own()
         else:
-            return PinnedMemory(ctx, pointer, size, owner=owner,
-                                finalizer=finalizer)
+            return PinnedMemory(
+                ctx, pointer, size, owner=owner, finalizer=finalizer
+            )
 
     def memallocmanaged(self, size, attach_global):
         if USE_NV_BINDING:
+
             def allocator():
                 ma_flags = binding.CUmemAttach_flags
 
@@ -1014,8 +1037,7 @@ class HostOnlyCUDAMemoryManager(BaseCUDAMemoryManager):
 
 
 class GetIpcHandleMixin:
-    """A class that provides a default implementation of ``get_ipc_handle()``.
-    """
+    """A class that provides a default implementation of ``get_ipc_handle()``."""
 
     def get_ipc_handle(self, memory):
         """Open an IPC memory handle by using ``cuMemGetAddressRange`` to
@@ -1034,8 +1056,9 @@ class GetIpcHandleMixin:
             offset = memory.handle.value - base
         source_info = self.context.device.get_device_identity()
 
-        return IpcHandle(memory, ipchandle, memory.size, source_info,
-                         offset=offset)
+        return IpcHandle(
+            memory, ipchandle, memory.size, source_info, offset=offset
+        )
 
 
 class NumbaCUDAMemoryManager(GetIpcHandleMixin, HostOnlyCUDAMemoryManager):
@@ -1050,6 +1073,7 @@ class NumbaCUDAMemoryManager(GetIpcHandleMixin, HostOnlyCUDAMemoryManager):
 
     def memalloc(self, size):
         if USE_NV_BINDING:
+
             def allocator():
                 return driver.cuMemAlloc(size)
 
@@ -1098,7 +1122,7 @@ def _ensure_memory_manager():
     if _memory_manager:
         return
 
-    if config.CUDA_MEMORY_MANAGER == 'default':
+    if config.CUDA_MEMORY_MANAGER == "default":
         _memory_manager = NumbaCUDAMemoryManager
         return
 
@@ -1106,8 +1130,9 @@ def _ensure_memory_manager():
         mgr_module = importlib.import_module(config.CUDA_MEMORY_MANAGER)
         set_memory_manager(mgr_module._numba_memory_manager)
     except Exception:
-        raise RuntimeError("Failed to use memory manager from %s" %
-                           config.CUDA_MEMORY_MANAGER)
+        raise RuntimeError(
+            "Failed to use memory manager from %s" % config.CUDA_MEMORY_MANAGER
+        )
 
 
 def set_memory_manager(mm_plugin):
@@ -1124,8 +1149,10 @@ def set_memory_manager(mm_plugin):
     dummy = mm_plugin(context=None)
     iv = dummy.interface_version
     if iv != _SUPPORTED_EMM_INTERFACE_VERSION:
-        err = "EMM Plugin interface has version %d - version %d required" \
-              % (iv, _SUPPORTED_EMM_INTERFACE_VERSION)
+        err = "EMM Plugin interface has version %d - version %d required" % (
+            iv,
+            _SUPPORTED_EMM_INTERFACE_VERSION,
+        )
         raise RuntimeError(err)
 
     _memory_manager = mm_plugin
@@ -1140,7 +1167,7 @@ class _SizeNotSet(int):
         return super().__new__(cls, 0)
 
     def __str__(self):
-        return '?'
+        return "?"
 
 
 _SizeNotSet = _SizeNotSet()
@@ -1153,6 +1180,7 @@ class _PendingDeallocs(object):
     modified later once the driver is initialized and the total memory capacity
     known.
     """
+
     def __init__(self, capacity=_SizeNotSet):
         self._cons = deque()
         self._disable_count = 0
@@ -1172,11 +1200,13 @@ class _PendingDeallocs(object):
         byte size of the resource added.  It is an optional argument.  Some
         resources (e.g. CUModule) has an unknown memory footprint on the device.
         """
-        _logger.info('add pending dealloc: %s %s bytes', dtor.__name__, size)
+        _logger.info("add pending dealloc: %s %s bytes", dtor.__name__, size)
         self._cons.append((dtor, handle, size))
         self._size += int(size)
-        if (len(self._cons) > config.CUDA_DEALLOCS_COUNT or
-                self._size > self._max_pending_bytes):
+        if (
+            len(self._cons) > config.CUDA_DEALLOCS_COUNT
+            or self._size > self._max_pending_bytes
+        ):
             self.clear()
 
     def clear(self):
@@ -1187,7 +1217,7 @@ class _PendingDeallocs(object):
         if not self.is_disabled:
             while self._cons:
                 [dtor, handle, size] = self._cons.popleft()
-                _logger.info('dealloc: %s %s bytes', dtor.__name__, size)
+                _logger.info("dealloc: %s %s bytes", dtor.__name__, size)
                 dtor(handle)
             self._size = 0
 
@@ -1251,19 +1281,19 @@ class Context(object):
         Clean up all owned resources in this context.
         """
         # Free owned resources
-        _logger.info('reset context of device %s', self.device.id)
+        _logger.info("reset context of device %s", self.device.id)
         self.memory_manager.reset()
         self.modules.clear()
         # Clear trash
         self.deallocations.clear()
 
     def get_memory_info(self):
-        """Returns (free, total) memory in bytes in the context.
-        """
+        """Returns (free, total) memory in bytes in the context."""
         return self.memory_manager.get_memory_info()
 
-    def get_active_blocks_per_multiprocessor(self, func, blocksize, memsize,
-                                             flags=None):
+    def get_active_blocks_per_multiprocessor(
+        self, func, blocksize, memsize, flags=None
+    ):
         """Return occupancy of a function.
         :param func: kernel for which occupancy is calculated
         :param blocksize: block size the kernel is intended to be launched with
@@ -1275,8 +1305,9 @@ class Context(object):
         else:
             return self._ctypes_active_blocks_per_multiprocessor(*args)
 
-    def _cuda_python_active_blocks_per_multiprocessor(self, func, blocksize,
-                                                      memsize, flags):
+    def _cuda_python_active_blocks_per_multiprocessor(
+        self, func, blocksize, memsize, flags
+    ):
         ps = [func.handle, blocksize, memsize]
 
         if not flags:
@@ -1285,8 +1316,9 @@ class Context(object):
         ps.append(flags)
         return driver.cuOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(*ps)
 
-    def _ctypes_active_blocks_per_multiprocessor(self, func, blocksize,
-                                                 memsize, flags):
+    def _ctypes_active_blocks_per_multiprocessor(
+        self, func, blocksize, memsize, flags
+    ):
         retval = c_int()
         args = (byref(retval), func.handle, blocksize, memsize)
 
@@ -1297,8 +1329,9 @@ class Context(object):
 
         return retval.value
 
-    def get_max_potential_block_size(self, func, b2d_func, memsize,
-                                     blocksizelimit, flags=None):
+    def get_max_potential_block_size(
+        self, func, b2d_func, memsize, blocksizelimit, flags=None
+    ):
         """Suggest a launch configuration with reasonable occupancy.
         :param func: kernel for which occupancy is calculated
         :param b2d_func: function that calculates how much per-block dynamic
@@ -1315,13 +1348,20 @@ class Context(object):
         else:
             return self._ctypes_max_potential_block_size(*args)
 
-    def _ctypes_max_potential_block_size(self, func, b2d_func, memsize,
-                                         blocksizelimit, flags):
+    def _ctypes_max_potential_block_size(
+        self, func, b2d_func, memsize, blocksizelimit, flags
+    ):
         gridsize = c_int()
         blocksize = c_int()
         b2d_cb = cu_occupancy_b2d_size(b2d_func)
-        args = [byref(gridsize), byref(blocksize), func.handle, b2d_cb,
-                memsize, blocksizelimit]
+        args = [
+            byref(gridsize),
+            byref(blocksize),
+            func.handle,
+            b2d_cb,
+            memsize,
+            blocksizelimit,
+        ]
 
         if not flags:
             driver.cuOccupancyMaxPotentialBlockSize(*args)
@@ -1331,10 +1371,11 @@ class Context(object):
 
         return (gridsize.value, blocksize.value)
 
-    def _cuda_python_max_potential_block_size(self, func, b2d_func, memsize,
-                                              blocksizelimit, flags):
+    def _cuda_python_max_potential_block_size(
+        self, func, b2d_func, memsize, blocksizelimit, flags
+    ):
         b2d_cb = ctypes.CFUNCTYPE(c_size_t, c_int)(b2d_func)
-        ptr = int.from_bytes(b2d_cb, byteorder='little')
+        ptr = int.from_bytes(b2d_cb, byteorder="little")
         driver_b2d_cb = binding.CUoccupancyB2DSize(ptr)
         args = [func.handle, driver_b2d_cb, memsize, blocksizelimit]
 
@@ -1387,7 +1428,7 @@ class Context(object):
         Returns an *IpcHandle* from a GPU allocation.
         """
         if not SUPPORTS_IPC:
-            raise OSError('OS does not support CUDA IPC')
+            raise OSError("OS does not support CUDA IPC")
         return self.memory_manager.get_ipc_handle(memory)
 
     def open_ipc_handle(self, handle, size):
@@ -1400,13 +1441,13 @@ class Context(object):
             driver.cuIpcOpenMemHandle(byref(dptr), handle, flags)
 
         # wrap it
-        return MemoryPointer(context=weakref.proxy(self), pointer=dptr,
-                             size=size)
+        return MemoryPointer(
+            context=weakref.proxy(self), pointer=dptr, size=size
+        )
 
     def enable_peer_access(self, peer_context, flags=0):
-        """Enable peer access between the current context and the peer context
-        """
-        assert flags == 0, '*flags* is reserved and MUST be zero'
+        """Enable peer access between the current context and the peer context"""
+        assert flags == 0, "*flags* is reserved and MUST be zero"
         driver.cuCtxEnablePeerAccess(peer_context, flags)
 
     def can_access_peer(self, peer_device):
@@ -1415,26 +1456,34 @@ class Context(object):
         """
         if USE_NV_BINDING:
             peer_device = binding.CUdevice(peer_device)
-            can_access_peer = driver.cuDeviceCanAccessPeer(self.device.id,
-                                                           peer_device)
+            can_access_peer = driver.cuDeviceCanAccessPeer(
+                self.device.id, peer_device
+            )
         else:
             can_access_peer = c_int()
-            driver.cuDeviceCanAccessPeer(byref(can_access_peer),
-                                         self.device.id, peer_device,)
+            driver.cuDeviceCanAccessPeer(
+                byref(can_access_peer),
+                self.device.id,
+                peer_device,
+            )
 
         return bool(can_access_peer)
 
     def create_module_ptx(self, ptx):
         if isinstance(ptx, str):
-            ptx = ptx.encode('utf8')
+            ptx = ptx.encode("utf8")
         if USE_NV_BINDING:
             image = ptx
         else:
             image = c_char_p(ptx)
         return self.create_module_image(image)
 
-    def create_module_image(self, image):
-        module = load_module_image(self, image)
+    def create_module_image(
+        self, image, setup_callbacks=None, teardown_callbacks=None
+    ):
+        module = load_module_image(
+            self, image, setup_callbacks, teardown_callbacks
+        )
         if USE_NV_BINDING:
             key = module.handle
         else:
@@ -1481,8 +1530,11 @@ class Context(object):
         else:
             handle = drvapi.cu_stream()
             driver.cuStreamCreate(byref(handle), 0)
-        return Stream(weakref.proxy(self), handle,
-                      _stream_finalizer(self.deallocations, handle))
+        return Stream(
+            weakref.proxy(self),
+            handle,
+            _stream_finalizer(self.deallocations, handle),
+        )
 
     def create_external_stream(self, ptr):
         if not isinstance(ptr, int):
@@ -1491,8 +1543,7 @@ class Context(object):
             handle = binding.CUstream(ptr)
         else:
             handle = drvapi.cu_stream(ptr)
-        return Stream(weakref.proxy(self), handle, None,
-                      external=True)
+        return Stream(weakref.proxy(self), handle, None, external=True)
 
     def create_event(self, timing=True):
         flags = 0
@@ -1503,8 +1554,11 @@ class Context(object):
         else:
             handle = drvapi.cu_event()
             driver.cuEventCreate(byref(handle), flags)
-        return Event(weakref.proxy(self), handle,
-                     finalizer=_event_finalizer(self.deallocations, handle))
+        return Event(
+            weakref.proxy(self),
+            handle,
+            finalizer=_event_finalizer(self.deallocations, handle),
+        )
 
     def synchronize(self):
         driver.cuCtxSynchronize()
@@ -1528,17 +1582,25 @@ class Context(object):
         return not self.__eq__(other)
 
 
-def load_module_image(context, image):
+def load_module_image(
+    context, image, setup_callbacks=None, teardown_callbacks=None
+):
     """
     image must be a pointer
     """
     if USE_NV_BINDING:
-        return load_module_image_cuda_python(context, image)
+        return load_module_image_cuda_python(
+            context, image, setup_callbacks, teardown_callbacks
+        )
     else:
-        return load_module_image_ctypes(context, image)
+        return load_module_image_ctypes(
+            context, image, setup_callbacks, teardown_callbacks
+        )
 
 
-def load_module_image_ctypes(context, image):
+def load_module_image_ctypes(
+    context, image, setup_callbacks, teardown_callbacks
+):
     logsz = config.CUDA_LOG_SIZE
 
     jitinfo = (c_char * logsz)()
@@ -1557,19 +1619,28 @@ def load_module_image_ctypes(context, image):
 
     handle = drvapi.cu_module()
     try:
-        driver.cuModuleLoadDataEx(byref(handle), image, len(options),
-                                  option_keys, option_vals)
+        driver.cuModuleLoadDataEx(
+            byref(handle), image, len(options), option_keys, option_vals
+        )
     except CudaAPIError as e:
         msg = "cuModuleLoadDataEx error:\n%s" % jiterrors.value.decode("utf8")
         raise CudaAPIError(e.code, msg)
 
     info_log = jitinfo.value
 
-    return CtypesModule(weakref.proxy(context), handle, info_log,
-                        _module_finalizer(context, handle))
+    return CtypesModule(
+        weakref.proxy(context),
+        handle,
+        info_log,
+        _module_finalizer(context, handle),
+        setup_callbacks,
+        teardown_callbacks,
+    )
 
 
-def load_module_image_cuda_python(context, image):
+def load_module_image_cuda_python(
+    context, image, setup_callbacks, teardown_callbacks
+):
     """
     image must be a pointer
     """
@@ -1591,17 +1662,24 @@ def load_module_image_cuda_python(context, image):
     option_vals = [v for v in options.values()]
 
     try:
-        handle = driver.cuModuleLoadDataEx(image, len(options), option_keys,
-                                           option_vals)
+        handle = driver.cuModuleLoadDataEx(
+            image, len(options), option_keys, option_vals
+        )
     except CudaAPIError as e:
-        err_string = jiterrors.decode('utf-8')
+        err_string = jiterrors.decode("utf-8")
         msg = "cuModuleLoadDataEx error:\n%s" % err_string
         raise CudaAPIError(e.code, msg)
 
-    info_log = jitinfo.decode('utf-8')
+    info_log = jitinfo.decode("utf-8")
 
-    return CudaPythonModule(weakref.proxy(context), handle, info_log,
-                            _module_finalizer(context, handle))
+    return CudaPythonModule(
+        weakref.proxy(context),
+        handle,
+        info_log,
+        _module_finalizer(context, handle),
+        setup_callbacks,
+        teardown_callbacks,
+    )
 
 
 def _alloc_finalizer(memory_manager, ptr, alloc_key, size):
@@ -1704,6 +1782,7 @@ class _CudaIpcImpl(object):
     """Implementation of GPU IPC using CUDA driver API.
     This requires the devices to be peer accessible.
     """
+
     def __init__(self, parent):
         self.base = parent.base
         self.handle = parent.handle
@@ -1717,10 +1796,10 @@ class _CudaIpcImpl(object):
         Import the IPC memory and returns a raw CUDA memory pointer object
         """
         if self.base is not None:
-            raise ValueError('opening IpcHandle from original process')
+            raise ValueError("opening IpcHandle from original process")
 
         if self._opened_mem is not None:
-            raise ValueError('IpcHandle is already opened')
+            raise ValueError("IpcHandle is already opened")
 
         mem = context.open_ipc_handle(self.handle, self.offset + self.size)
         # this object owns the opened allocation
@@ -1731,7 +1810,7 @@ class _CudaIpcImpl(object):
 
     def close(self):
         if self._opened_mem is None:
-            raise ValueError('IpcHandle not opened')
+            raise ValueError("IpcHandle not opened")
         driver.cuIpcCloseMemHandle(self._opened_mem.handle)
         self._opened_mem = None
 
@@ -1740,6 +1819,7 @@ class _StagedIpcImpl(object):
     """Implementation of GPU IPC using custom staging logic to workaround
     CUDA IPC limitation on peer accessibility between devices.
     """
+
     def __init__(self, parent, source_info):
         self.parent = parent
         self.base = parent.base
@@ -1795,6 +1875,7 @@ class IpcHandle(object):
                    referred to by this IPC handle.
     :type offset: int
     """
+
     def __init__(self, base, handle, size, source_info=None, offset=0):
         self.base = base
         self.handle = handle
@@ -1818,12 +1899,11 @@ class IpcHandle(object):
         return context.can_access_peer(source_device.id)
 
     def open_staged(self, context):
-        """Open the IPC by allowing staging on the host memory first.
-        """
+        """Open the IPC by allowing staging on the host memory first."""
         self._sentry_source_info()
 
         if self._impl is not None:
-            raise ValueError('IpcHandle is already opened')
+            raise ValueError("IpcHandle is already opened")
 
         self._impl = _StagedIpcImpl(self, self.source_info)
         return self._impl.open(context)
@@ -1833,7 +1913,7 @@ class IpcHandle(object):
         Import the IPC memory and returns a raw CUDA memory pointer object
         """
         if self._impl is not None:
-            raise ValueError('IpcHandle is already opened')
+            raise ValueError("IpcHandle is already opened")
 
         self._impl = _CudaIpcImpl(self)
         return self._impl.open(context)
@@ -1864,12 +1944,13 @@ class IpcHandle(object):
             strides = dtype.itemsize
         dptr = self.open(context)
         # read the device pointer as an array
-        return devicearray.DeviceNDArray(shape=shape, strides=strides,
-                                         dtype=dtype, gpu_data=dptr)
+        return devicearray.DeviceNDArray(
+            shape=shape, strides=strides, dtype=dtype, gpu_data=dptr
+        )
 
     def close(self):
         if self._impl is None:
-            raise ValueError('IpcHandle not opened')
+            raise ValueError("IpcHandle not opened")
         self._impl.close()
         self._impl = None
 
@@ -1895,8 +1976,13 @@ class IpcHandle(object):
         else:
             handle = drvapi.cu_ipc_mem_handle()
         handle.reserved = handle_ary
-        return cls(base=None, handle=handle, size=size,
-                   source_info=source_info, offset=offset)
+        return cls(
+            base=None,
+            handle=handle,
+            size=size,
+            source_info=source_info,
+            offset=offset,
+        )
 
 
 class MemoryPointer(object):
@@ -1930,6 +2016,7 @@ class MemoryPointer(object):
     :param finalizer: A function that is called when the buffer is to be freed.
     :type finalizer: function
     """
+
     __cuda_memory__ = True
 
     def __init__(self, context, pointer, size, owner=None, finalizer=None):
@@ -1965,8 +2052,9 @@ class MemoryPointer(object):
     def memset(self, byte, count=None, stream=0):
         count = self.size if count is None else count
         if stream:
-            driver.cuMemsetD8Async(self.device_pointer, byte, count,
-                                   stream.handle)
+            driver.cuMemsetD8Async(
+                self.device_pointer, byte, count, stream.handle
+            )
         else:
             driver.cuMemsetD8(self.device_pointer, byte, count)
 
@@ -1980,12 +2068,12 @@ class MemoryPointer(object):
         if not self.device_pointer_value:
             if size != 0:
                 raise RuntimeError("non-empty slice into empty slice")
-            view = self      # new view is just a reference to self
+            view = self  # new view is just a reference to self
         # Handle normal case
         else:
             base = self.device_pointer_value + start
             if size < 0:
-                raise RuntimeError('size cannot be negative')
+                raise RuntimeError("size cannot be negative")
             if USE_NV_BINDING:
                 pointer = binding.CUdeviceptr()
                 ctypes_ptr = drvapi.cu_device_ptr.from_address(pointer.getPtr())
@@ -2021,6 +2109,7 @@ class AutoFreePointer(MemoryPointer):
 
     Constructor arguments are the same as for :class:`MemoryPointer`.
     """
+
     def __init__(self, *args, **kwargs):
         super(AutoFreePointer, self).__init__(*args, **kwargs)
         # Releease the self reference to the buffer, so that the finalizer
@@ -2063,8 +2152,9 @@ class MappedMemory(AutoFreePointer):
             self._bufptr_ = self.host_pointer.value
 
         self.device_pointer = devptr
-        super(MappedMemory, self).__init__(context, devptr, size,
-                                           finalizer=finalizer)
+        super(MappedMemory, self).__init__(
+            context, devptr, size, finalizer=finalizer
+        )
         self.handle = self.host_pointer
 
         # For buffer interface
@@ -2179,8 +2269,7 @@ class OwnedPointer(object):
         weakref.finalize(self, deref)
 
     def __getattr__(self, fname):
-        """Proxy MemoryPointer methods
-        """
+        """Proxy MemoryPointer methods"""
         return getattr(self._view, fname)
 
 
@@ -2211,18 +2300,15 @@ class Stream(object):
         if USE_NV_BINDING:
             default_streams = {
                 CU_STREAM_DEFAULT: "<Default CUDA stream on %s>",
-                binding.CU_STREAM_LEGACY:
-                    "<Legacy default CUDA stream on %s>",
-                binding.CU_STREAM_PER_THREAD:
-                    "<Per-thread default CUDA stream on %s>",
+                binding.CU_STREAM_LEGACY: "<Legacy default CUDA stream on %s>",
+                binding.CU_STREAM_PER_THREAD: "<Per-thread default CUDA stream on %s>",
             }
             ptr = int(self.handle) or 0
         else:
             default_streams = {
                 drvapi.CU_STREAM_DEFAULT: "<Default CUDA stream on %s>",
                 drvapi.CU_STREAM_LEGACY: "<Legacy default CUDA stream on %s>",
-                drvapi.CU_STREAM_PER_THREAD:
-                    "<Per-thread default CUDA stream on %s>",
+                drvapi.CU_STREAM_PER_THREAD: "<Per-thread default CUDA stream on %s>",
             }
             ptr = self.handle.value or drvapi.CU_STREAM_DEFAULT
 
@@ -2234,18 +2320,18 @@ class Stream(object):
             return "<CUDA stream %d on %s>" % (ptr, self.context)
 
     def synchronize(self):
-        '''
+        """
         Wait for all commands in this stream to execute. This will commit any
         pending memory transfers.
-        '''
+        """
         driver.cuStreamSynchronize(self.handle)
 
     @contextlib.contextmanager
     def auto_synchronize(self):
-        '''
+        """
         A context manager that waits for all commands in this stream to execute
         and commits any pending memory transfers upon exiting the context.
-        '''
+        """
         yield self
         self.synchronize()
 
@@ -2272,7 +2358,7 @@ class Stream(object):
         data = (self, callback, arg)
         _py_incref(data)
         if USE_NV_BINDING:
-            ptr = int.from_bytes(self._stream_callback, byteorder='little')
+            ptr = int.from_bytes(self._stream_callback, byteorder="little")
             stream_callback = binding.CUstreamCallback(ptr)
             # The callback needs to receive a pointer to the data PyObject
             data = id(data)
@@ -2373,9 +2459,9 @@ class Event(object):
 
 
 def event_elapsed_time(evtstart, evtend):
-    '''
+    """
     Compute the elapsed time between two events in milliseconds.
-    '''
+    """
     if USE_NV_BINDING:
         return driver.cuEventElapsedTime(evtstart.handle, evtend.handle)
     else:
@@ -2387,13 +2473,27 @@ def event_elapsed_time(evtstart, evtend):
 class Module(metaclass=ABCMeta):
     """Abstract base class for modules"""
 
-    def __init__(self, context, handle, info_log, finalizer=None):
+    def __init__(
+        self,
+        context,
+        handle,
+        info_log,
+        finalizer=None,
+        setup_callbacks=None,
+        teardown_callbacks=None,
+    ):
         self.context = context
         self.handle = handle
         self.info_log = info_log
         if finalizer is not None:
             self._finalizer = weakref.finalize(self, finalizer)
 
+        self.initialized = False
+        self.setup_functions = setup_callbacks
+        self.teardown_functions = teardown_callbacks
+
+        self._set_finalizers()
+
     def unload(self):
         """Unload this module from the context"""
         self.context.unload_module(self)
@@ -2406,36 +2506,66 @@ class Module(metaclass=ABCMeta):
     def get_global_symbol(self, name):
         """Return a MemoryPointer referring to the named symbol"""
 
+    def setup(self):
+        """Call the setup functions for the module"""
+        if self.initialized:
+            raise RuntimeError("The module has already been initialized.")
 
-class CtypesModule(Module):
+        if self.setup_functions is None:
+            return
+
+        for f in self.setup_functions:
+            f(self.handle)
 
+        self.initialized = True
+
+    def _set_finalizers(self):
+        """Create finalizers that tear down the module."""
+        if self.teardown_functions is None:
+            return
+
+        def _teardown(teardowns, handle):
+            for f in teardowns:
+                f(handle)
+
+        weakref.finalize(
+            self,
+            _teardown,
+            self.teardown_functions,
+            self.handle,
+        )
+
+
+class CtypesModule(Module):
     def get_function(self, name):
         handle = drvapi.cu_function()
-        driver.cuModuleGetFunction(byref(handle), self.handle,
-                                   name.encode('utf8'))
+        driver.cuModuleGetFunction(
+            byref(handle), self.handle, name.encode("utf8")
+        )
         return CtypesFunction(weakref.proxy(self), handle, name)
 
     def get_global_symbol(self, name):
         ptr = drvapi.cu_device_ptr()
         size = drvapi.c_size_t()
-        driver.cuModuleGetGlobal(byref(ptr), byref(size), self.handle,
-                                 name.encode('utf8'))
+        driver.cuModuleGetGlobal(
+            byref(ptr), byref(size), self.handle, name.encode("utf8")
+        )
         return MemoryPointer(self.context, ptr, size), size.value
 
 
 class CudaPythonModule(Module):
-
     def get_function(self, name):
-        handle = driver.cuModuleGetFunction(self.handle, name.encode('utf8'))
+        handle = driver.cuModuleGetFunction(self.handle, name.encode("utf8"))
         return CudaPythonFunction(weakref.proxy(self), handle, name)
 
     def get_global_symbol(self, name):
-        ptr, size = driver.cuModuleGetGlobal(self.handle, name.encode('utf8'))
+        ptr, size = driver.cuModuleGetGlobal(self.handle, name.encode("utf8"))
         return MemoryPointer(self.context, ptr, size), size
 
 
-FuncAttr = namedtuple("FuncAttr", ["regs", "shared", "local", "const",
-                                   "maxthreads"])
+FuncAttr = namedtuple(
+    "FuncAttr", ["regs", "shared", "local", "const", "maxthreads"]
+)
 
 
 class Function(metaclass=ABCMeta):
@@ -2458,8 +2588,9 @@ class Function(metaclass=ABCMeta):
         return self.module.context.device
 
     @abstractmethod
-    def cache_config(self, prefer_equal=False, prefer_cache=False,
-                     prefer_shared=False):
+    def cache_config(
+        self, prefer_equal=False, prefer_cache=False, prefer_shared=False
+    ):
         """Set the cache configuration for this function."""
 
     @abstractmethod
@@ -2473,9 +2604,9 @@ class Function(metaclass=ABCMeta):
 
 
 class CtypesFunction(Function):
-
-    def cache_config(self, prefer_equal=False, prefer_cache=False,
-                     prefer_shared=False):
+    def cache_config(
+        self, prefer_equal=False, prefer_cache=False, prefer_shared=False
+    ):
         prefer_equal = prefer_equal or (prefer_cache and prefer_shared)
         if prefer_equal:
             flag = enums.CU_FUNC_CACHE_PREFER_EQUAL
@@ -2498,15 +2629,17 @@ class CtypesFunction(Function):
         lmem = self.read_func_attr(enums.CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES)
         smem = self.read_func_attr(enums.CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES)
         maxtpb = self.read_func_attr(
-            enums.CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK)
-        return FuncAttr(regs=nregs, const=cmem, local=lmem, shared=smem,
-                        maxthreads=maxtpb)
+            enums.CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK
+        )
+        return FuncAttr(
+            regs=nregs, const=cmem, local=lmem, shared=smem, maxthreads=maxtpb
+        )
 
 
 class CudaPythonFunction(Function):
-
-    def cache_config(self, prefer_equal=False, prefer_cache=False,
-                     prefer_shared=False):
+    def cache_config(
+        self, prefer_equal=False, prefer_cache=False, prefer_shared=False
+    ):
         prefer_equal = prefer_equal or (prefer_cache and prefer_shared)
         attr = binding.CUfunction_attribute
         if prefer_equal:
@@ -2529,19 +2662,26 @@ class CudaPythonFunction(Function):
         lmem = self.read_func_attr(attr.CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES)
         smem = self.read_func_attr(attr.CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES)
         maxtpb = self.read_func_attr(
-            attr.CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK)
-        return FuncAttr(regs=nregs, const=cmem, local=lmem, shared=smem,
-                        maxthreads=maxtpb)
-
+            attr.CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK
+        )
+        return FuncAttr(
+            regs=nregs, const=cmem, local=lmem, shared=smem, maxthreads=maxtpb
+        )
 
-def launch_kernel(cufunc_handle,
-                  gx, gy, gz,
-                  bx, by, bz,
-                  sharedmem,
-                  hstream,
-                  args,
-                  cooperative=False):
 
+def launch_kernel(
+    cufunc_handle,
+    gx,
+    gy,
+    gz,
+    bx,
+    by,
+    bz,
+    sharedmem,
+    hstream,
+    args,
+    cooperative=False,
+):
     param_ptrs = [addressof(arg) for arg in args]
     params = (c_void_p * len(param_ptrs))(*param_ptrs)
 
@@ -2553,46 +2693,54 @@ def launch_kernel(cufunc_handle,
         extra = None
 
     if cooperative:
-        driver.cuLaunchCooperativeKernel(cufunc_handle,
-                                         gx, gy, gz,
-                                         bx, by, bz,
-                                         sharedmem,
-                                         hstream,
-                                         params_for_launch)
+        driver.cuLaunchCooperativeKernel(
+            cufunc_handle,
+            gx,
+            gy,
+            gz,
+            bx,
+            by,
+            bz,
+            sharedmem,
+            hstream,
+            params_for_launch,
+        )
     else:
-        driver.cuLaunchKernel(cufunc_handle,
-                              gx, gy, gz,
-                              bx, by, bz,
-                              sharedmem,
-                              hstream,
-                              params_for_launch,
-                              extra)
+        driver.cuLaunchKernel(
+            cufunc_handle,
+            gx,
+            gy,
+            gz,
+            bx,
+            by,
+            bz,
+            sharedmem,
+            hstream,
+            params_for_launch,
+            extra,
+        )
 
 
 class Linker(metaclass=ABCMeta):
     """Abstract base class for linkers"""
 
     @classmethod
-    def new(cls,
-            max_registers=0,
-            lineinfo=False,
-            cc=None,
-            lto=None,
-            additional_flags=None
-            ):
-
+    def new(
+        cls,
+        max_registers=0,
+        lineinfo=False,
+        cc=None,
+        lto=None,
+        additional_flags=None,
+    ):
         driver_ver = driver.get_version()
-        if (
-            config.CUDA_ENABLE_MINOR_VERSION_COMPATIBILITY
-            and driver_ver >= (12, 0)
+        if config.CUDA_ENABLE_MINOR_VERSION_COMPATIBILITY and driver_ver >= (
+            12,
+            0,
         ):
-            raise ValueError(
-                "Use CUDA_ENABLE_PYNVJITLINK for CUDA >= 12.0 MVC"
-            )
+            raise ValueError("Use CUDA_ENABLE_PYNVJITLINK for CUDA >= 12.0 MVC")
         if config.CUDA_ENABLE_PYNVJITLINK and driver_ver < (12, 0):
-            raise ValueError(
-                "Enabling pynvjitlink requires CUDA 12."
-            )
+            raise ValueError("Enabling pynvjitlink requires CUDA 12.")
         if config.CUDA_ENABLE_PYNVJITLINK:
             linker = PyNvJitLinker
 
@@ -2641,9 +2789,9 @@ class Linker(metaclass=ABCMeta):
         ptx, log = nvrtc.compile(cu, name, cc)
 
         if config.DUMP_ASSEMBLY:
-            print(("ASSEMBLY %s" % name).center(80, '-'))
+            print(("ASSEMBLY %s" % name).center(80, "-"))
             print(ptx)
-            print('=' * 80)
+            print("=" * 80)
 
         # Link the program's PTX using the normal linker mechanism
         ptx_name = os.path.splitext(name)[0] + ".ptx"
@@ -2654,7 +2802,7 @@ class Linker(metaclass=ABCMeta):
         """Add code from a file to the link"""
 
     def add_cu_file(self, path):
-        with open(path, 'rb') as f:
+        with open(path, "rb") as f:
             cu = f.read()
         self.add_cu(cu, os.path.basename(path))
 
@@ -2672,24 +2820,24 @@ class Linker(metaclass=ABCMeta):
 
         if isinstance(path_or_code, str):
             ext = pathlib.Path(path_or_code).suffix
-            if ext == '':
+            if ext == "":
                 raise RuntimeError(
                     "Don't know how to link file with no extension"
                 )
-            elif ext == '.cu':
+            elif ext == ".cu":
                 self.add_cu_file(path_or_code)
             else:
-                kind = FILE_EXTENSION_MAP.get(ext.lstrip('.'), None)
+                kind = FILE_EXTENSION_MAP.get(ext.lstrip("."), None)
                 if kind is None:
                     raise RuntimeError(
-                        "Don't know how to link file with extension "
-                        f"{ext}"
+                        f"Don't know how to link file with extension {ext}"
                     )
 
                 if ignore_nonlto:
                     warn_and_return = False
                     if kind in (
-                        FILE_EXTENSION_MAP["fatbin"], FILE_EXTENSION_MAP["o"]
+                        FILE_EXTENSION_MAP["fatbin"],
+                        FILE_EXTENSION_MAP["o"],
                     ):
                         entry_types = inspect_obj_content(path_or_code)
                         if "nvvm" not in entry_types:
@@ -2754,6 +2902,7 @@ class MVCLinker(Linker):
     Linker supporting Minor Version Compatibility, backed by the cubinlinker
     package.
     """
+
     def __init__(self, max_registers=None, lineinfo=False, cc=None):
         try:
             from cubinlinker import CubinLinker
@@ -2761,18 +2910,20 @@ class MVCLinker(Linker):
             raise ImportError(_MVC_ERROR_MESSAGE) from err
 
         if cc is None:
-            raise RuntimeError("MVCLinker requires Compute Capability to be "
-                               "specified, but cc is None")
+            raise RuntimeError(
+                "MVCLinker requires Compute Capability to be "
+                "specified, but cc is None"
+            )
 
         super().__init__(max_registers, lineinfo, cc)
 
         arch = f"sm_{cc[0] * 10 + cc[1]}"
-        ptx_compile_opts = ['--gpu-name', arch, '-c']
+        ptx_compile_opts = ["--gpu-name", arch, "-c"]
         if max_registers:
             arg = f"--maxrregcount={max_registers}"
             ptx_compile_opts.append(arg)
         if lineinfo:
-            ptx_compile_opts.append('--generate-line-info')
+            ptx_compile_opts.append("--generate-line-info")
         self.ptx_compile_options = tuple(ptx_compile_opts)
 
         self._linker = CubinLinker(f"--arch={arch}")
@@ -2785,7 +2936,7 @@ class MVCLinker(Linker):
     def error_log(self):
         return self._linker.error_log
 
-    def add_ptx(self, ptx, name='<cudapy-ptx>'):
+    def add_ptx(self, ptx, name="<cudapy-ptx>"):
         try:
             from ptxcompiler import compile_ptx
             from cubinlinker import CubinLinkerError
@@ -2804,19 +2955,19 @@ class MVCLinker(Linker):
             raise ImportError(_MVC_ERROR_MESSAGE) from err
 
         try:
-            with open(path, 'rb') as f:
+            with open(path, "rb") as f:
                 data = f.read()
         except FileNotFoundError:
-            raise LinkerError(f'{path} not found')
+            raise LinkerError(f"{path} not found")
 
         name = pathlib.Path(path).name
-        if kind == FILE_EXTENSION_MAP['cubin']:
+        if kind == FILE_EXTENSION_MAP["cubin"]:
             fn = self._linker.add_cubin
-        elif kind == FILE_EXTENSION_MAP['fatbin']:
+        elif kind == FILE_EXTENSION_MAP["fatbin"]:
             fn = self._linker.add_fatbin
-        elif kind == FILE_EXTENSION_MAP['a']:
+        elif kind == FILE_EXTENSION_MAP["a"]:
             raise LinkerError(f"Don't know how to link {kind}")
-        elif kind == FILE_EXTENSION_MAP['ptx']:
+        elif kind == FILE_EXTENSION_MAP["ptx"]:
             return self.add_ptx(data, name)
         else:
             raise LinkerError(f"Don't know how to link {kind}")
@@ -2842,6 +2993,7 @@ class CtypesLinker(Linker):
     """
     Links for current device if no CC given
     """
+
     def __init__(self, max_registers=0, lineinfo=False, cc=None):
         super().__init__(max_registers, lineinfo, cc)
 
@@ -2875,8 +3027,9 @@ class CtypesLinker(Linker):
         option_vals = (c_void_p * len(raw_values))(*raw_values)
 
         self.handle = handle = drvapi.cu_link_state()
-        driver.cuLinkCreate(len(raw_keys), option_keys, option_vals,
-                            byref(self.handle))
+        driver.cuLinkCreate(
+            len(raw_keys), option_keys, option_vals, byref(self.handle)
+        )
 
         weakref.finalize(self, driver.cuLinkDestroy, handle)
 
@@ -2887,19 +3040,27 @@ class CtypesLinker(Linker):
 
     @property
     def info_log(self):
-        return self.linker_info_buf.value.decode('utf8')
+        return self.linker_info_buf.value.decode("utf8")
 
     @property
     def error_log(self):
-        return self.linker_errors_buf.value.decode('utf8')
+        return self.linker_errors_buf.value.decode("utf8")
 
-    def add_ptx(self, ptx, name='<cudapy-ptx>'):
+    def add_ptx(self, ptx, name="<cudapy-ptx>"):
         ptxbuf = c_char_p(ptx)
-        namebuf = c_char_p(name.encode('utf8'))
+        namebuf = c_char_p(name.encode("utf8"))
         self._keep_alive += [ptxbuf, namebuf]
         try:
-            driver.cuLinkAddData(self.handle, enums.CU_JIT_INPUT_PTX,
-                                 ptxbuf, len(ptx), namebuf, 0, None, None)
+            driver.cuLinkAddData(
+                self.handle,
+                enums.CU_JIT_INPUT_PTX,
+                ptxbuf,
+                len(ptx),
+                namebuf,
+                0,
+                None,
+                None,
+            )
         except CudaAPIError as e:
             raise LinkerError("%s\n%s" % (e, self.error_log))
 
@@ -2911,7 +3072,7 @@ class CtypesLinker(Linker):
             driver.cuLinkAddFile(self.handle, kind, pathbuf, 0, None, None)
         except CudaAPIError as e:
             if e.code == enums.CUDA_ERROR_FILE_NOT_FOUND:
-                msg = f'{path} not found'
+                msg = f"{path} not found"
             else:
                 msg = "%s\n%s" % (e, self.error_log)
             raise LinkerError(msg)
@@ -2926,7 +3087,7 @@ class CtypesLinker(Linker):
             raise LinkerError("%s\n%s" % (e, self.error_log))
 
         size = size.value
-        assert size > 0, 'linker returned a zero sized cubin'
+        assert size > 0, "linker returned a zero sized cubin"
         del self._keep_alive[:]
 
         # We return a copy of the cubin because it's owned by the linker
@@ -2938,6 +3099,7 @@ class CudaPythonLinker(Linker):
     """
     Links for current device if no CC given
     """
+
     def __init__(self, max_registers=0, lineinfo=False, cc=None):
         super().__init__(max_registers, lineinfo, cc)
 
@@ -2964,8 +3126,9 @@ class CudaPythonLinker(Linker):
             options[jit_option.CU_JIT_TARGET_FROM_CUCONTEXT] = 1
         else:
             cc_val = cc[0] * 10 + cc[1]
-            cc_enum = getattr(binding.CUjit_target,
-                              f'CU_TARGET_COMPUTE_{cc_val}')
+            cc_enum = getattr(
+                binding.CUjit_target, f"CU_TARGET_COMPUTE_{cc_val}"
+            )
             options[jit_option.CU_JIT_TARGET] = cc_enum
 
         raw_keys = list(options.keys())
@@ -2982,19 +3145,20 @@ class CudaPythonLinker(Linker):
 
     @property
     def info_log(self):
-        return self.linker_info_buf.decode('utf8')
+        return self.linker_info_buf.decode("utf8")
 
     @property
     def error_log(self):
-        return self.linker_errors_buf.decode('utf8')
+        return self.linker_errors_buf.decode("utf8")
 
-    def add_ptx(self, ptx, name='<cudapy-ptx>'):
-        namebuf = name.encode('utf8')
+    def add_ptx(self, ptx, name="<cudapy-ptx>"):
+        namebuf = name.encode("utf8")
         self._keep_alive += [ptx, namebuf]
         try:
             input_ptx = binding.CUjitInputType.CU_JIT_INPUT_PTX
-            driver.cuLinkAddData(self.handle, input_ptx, ptx, len(ptx),
-                                 namebuf, 0, [], [])
+            driver.cuLinkAddData(
+                self.handle, input_ptx, ptx, len(ptx), namebuf, 0, [], []
+            )
         except CudaAPIError as e:
             raise LinkerError("%s\n%s" % (e, self.error_log))
 
@@ -3006,7 +3170,7 @@ class CudaPythonLinker(Linker):
             driver.cuLinkAddFile(self.handle, kind, pathbuf, 0, [], [])
         except CudaAPIError as e:
             if e.code == binding.CUresult.CUDA_ERROR_FILE_NOT_FOUND:
-                msg = f'{path} not found'
+                msg = f"{path} not found"
             else:
                 msg = "%s\n%s" % (e, self.error_log)
             raise LinkerError(msg)
@@ -3017,7 +3181,7 @@ class CudaPythonLinker(Linker):
         except CudaAPIError as e:
             raise LinkerError("%s\n%s" % (e, self.error_log))
 
-        assert size > 0, 'linker returned a zero sized cubin'
+        assert size > 0, "linker returned a zero sized cubin"
         del self._keep_alive[:]
         # We return a copy of the cubin because it's owned by the linker
         cubin_ptr = ctypes.cast(cubin_buf, ctypes.POINTER(ctypes.c_char))
@@ -3151,6 +3315,7 @@ class PyNvJitLinker(Linker):
         except NvJitLinkError as e:
             raise LinkerError from e
 
+
 # -----------------------------------------------------------------------------
 
 
@@ -3200,7 +3365,7 @@ def device_memory_size(devmem):
     The result is cached in the device memory object.
     It may query the driver for the memory size of the device memory allocation.
     """
-    sz = getattr(devmem, '_cuda_memsize_', None)
+    sz = getattr(devmem, "_cuda_memsize_", None)
     if sz is None:
         s, e = device_extents(devmem)
         if USE_NV_BINDING:
@@ -3213,10 +3378,9 @@ def device_memory_size(devmem):
 
 
 def _is_datetime_dtype(obj):
-    """Returns True if the obj.dtype is datetime64 or timedelta64
-    """
-    dtype = getattr(obj, 'dtype', None)
-    return dtype is not None and dtype.char in 'Mm'
+    """Returns True if the obj.dtype is datetime64 or timedelta64"""
+    dtype = getattr(obj, "dtype", None)
+    return dtype is not None and dtype.char in "Mm"
 
 
 def _workaround_for_datetime(obj):
@@ -3295,12 +3459,11 @@ def is_device_memory(obj):
     "device_pointer" which value is an int object carrying the pointer
     value of the device memory address.  This is not tested in this method.
     """
-    return getattr(obj, '__cuda_memory__', False)
+    return getattr(obj, "__cuda_memory__", False)
 
 
 def require_device_memory(obj):
-    """A sentry for methods that accept CUDA memory object.
-    """
+    """A sentry for methods that accept CUDA memory object."""
     if not is_device_memory(obj):
         raise Exception("Not a CUDA memory object.")
 
@@ -3391,16 +3554,16 @@ def device_memset(dst, val, size, stream=0):
 
 
 def profile_start():
-    '''
+    """
     Enable profile collection in the current context.
-    '''
+    """
     driver.cuProfilerStart()
 
 
 def profile_stop():
-    '''
+    """
     Disable profile collection in the current context.
-    '''
+    """
     driver.cuProfilerStop()
 
 
@@ -3427,18 +3590,21 @@ def inspect_obj_content(objpath: str):
     Given path to a fatbin or object, use `cuobjdump` to examine its content
     Return the set of entries in the object.
     """
-    code_types :set[str] = set()
+    code_types: set[str] = set()
 
     try:
-        out = subprocess.run(["cuobjdump", objpath], check=True,
-                             capture_output=True)
+        out = subprocess.run(
+            ["cuobjdump", objpath], check=True, capture_output=True
+        )
     except FileNotFoundError as e:
-        msg = ("cuobjdump has not been found. You may need "
-               "to install the CUDA toolkit and ensure that "
-               "it is available on your PATH.\n")
+        msg = (
+            "cuobjdump has not been found. You may need "
+            "to install the CUDA toolkit and ensure that "
+            "it is available on your PATH.\n"
+        )
         raise RuntimeError(msg) from e
 
-    objtable = out.stdout.decode('utf-8')
+    objtable = out.stdout.decode("utf-8")
     entry_pattern = r"Fatbin (.*) code"
     for line in objtable.split("\n"):
         if match := re.match(entry_pattern, line):