Add Module Setup and Teardown Callback to Linkable Code Interface

docs/source/user/cuda_ffi.rst

@@ -160,6 +160,23 @@ CUDA C/C++ source code will be compiled with the `NVIDIA Runtime Compiler
 kernel as either PTX or LTOIR, depending on whether LTO is enabled. Other files
 will be passed directly to the CUDA Linker.
 
+A ``LinkableCode`` object may have setup and teardown callback functions that
+perform module-specific initialization and cleanup tasks.
+
+* Setup functions are invoked once for every new module loaded.
+* Teardown functions are invoked just prior to module unloading.
+
+Both setup and teardown callbacks are called with a handle to the relevant
+module. In practice, Numba creates a new module each time a kernel is compiled
+for a specific set of argument types.
+
+The callbacks are defined as follows:
+
+.. code::
+
+  def setup_callback(mod: cuda.cudadrv.drvapi.cu_module):...
+  def teardown_callback(mod: cuda.cudadrv.drvapi.cu_module):...
+
 :class:`LinkableCode <numba.cuda.LinkableCode>` objects are initialized using
 the parameters of their base class:
 

numba_cuda/numba/cuda/codegen.py

@@ -4,6 +4,7 @@
 from numba.core.codegen import Codegen, CodeLibrary
 from .cudadrv import devices, driver, nvvm, runtime
 from numba.cuda.cudadrv.libs import get_cudalib
+from numba.cuda.cudadrv.linkable_code import LinkableCode
 
 import os
 import subprocess
@@ -94,6 +95,12 @@ def __init__(
         # Files to link with the generated PTX. These are linked using the
         # Driver API at link time.
         self._linking_files = set()
+        # List of setup functions to the loaded module
+        # the order is determined by the order they are added to the codelib.
+        self._setup_functions = []
+        # List of teardown functions to the loaded module
+        # the order is determined by the order they are added to the codelib.
+        self._teardown_functions = []
         # Should we link libcudadevrt?
         self.needs_cudadevrt = False
 
@@ -246,7 +253,8 @@ def get_cufunc(self):
             return cufunc
 
         cubin = self.get_cubin(cc=device.compute_capability)
-        module = ctx.create_module_image(cubin)
+        module = ctx.create_module_image(
+            cubin, self._setup_functions, self._teardown_functions)
 
         # Load
         cufunc = module.get_function(self._entry_name)
@@ -284,8 +292,14 @@ def add_linking_library(self, library):
 
         self._linking_libraries.add(library)
 
-    def add_linking_file(self, filepath):
-        self._linking_files.add(filepath)
+    def add_linking_file(self, path_or_obj):
+        if isinstance(path_or_obj, LinkableCode):
+            if path_or_obj.setup_callback:
+                self._setup_functions.append(path_or_obj.setup_callback)
+            if path_or_obj.teardown_callback:
+                self._teardown_functions.append(path_or_obj.teardown_callback)
+
+        self._linking_files.add(path_or_obj)
 
     def get_function(self, name):
         for fn in self._module.functions:

numba_cuda/numba/cuda/cudadrv/driver.py

@@ -1433,8 +1433,10 @@ def create_module_ptx(self, ptx):
             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:
@@ -1528,17 +1530,21 @@ def __ne__(self, other):
         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)()
@@ -1566,10 +1572,12 @@ def load_module_image_ctypes(context, image):
     info_log = jitinfo.value
 
     return CtypesModule(weakref.proxy(context), handle, info_log,
-                        _module_finalizer(context, handle))
+                        _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
     """
@@ -1601,7 +1609,8 @@ def load_module_image_cuda_python(context, image):
     info_log = jitinfo.decode('utf-8')
 
     return CudaPythonModule(weakref.proxy(context), handle, info_log,
-                            _module_finalizer(context, handle))
+                            _module_finalizer(context, handle),
+                            setup_callbacks, teardown_callbacks)
 
 
 def _alloc_finalizer(memory_manager, ptr, alloc_key, size):
@@ -2387,13 +2396,20 @@ 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,6 +2422,35 @@ def get_function(self, name):
     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.")
+
+        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):
 

numba_cuda/numba/cuda/cudadrv/linkable_code.py

@@ -7,11 +7,27 @@ class LinkableCode:
     :param data: A buffer containing the data to link.
     :param name: The name of the file to be referenced in any compilation or
                  linking errors that may be produced.
+    :param setup_callback: A function called prior to the launch of a kernel
+                           contained within a module that has this code object
+                           linked into it.
+    :param teardown_callback: A function called just prior to the unloading of
+                              a module that has this code object linked into
+                              it.
     """
 
-    def __init__(self, data, name=None):
+    def __init__(self, data, name=None,
+                 setup_callback=None,
+                 teardown_callback=None):
+
+        if setup_callback and not callable(setup_callback):
+            raise TypeError("setup_callback must be callable")
+        if teardown_callback and not callable(teardown_callback):
+            raise TypeError("teardown_callback must be callable")
+
         self.data = data
         self._name = name
+        self.setup_callback = setup_callback
+        self.teardown_callback = teardown_callback
 
     @property
     def name(self):

numba_cuda/numba/cuda/dispatcher.py

@@ -310,6 +310,10 @@ def bind(self):
         """
         cufunc = self._codelibrary.get_cufunc()
 
+        mod = cufunc.module
+        if not mod.initialized:
+            mod.setup()
+
         if (
             hasattr(self, "target_context")
             and self.target_context.enable_nrt