add kernel finalizer

Author: isVoid

numba_cuda/numba/cuda/cudadrv/managed_module.py

@@ -1,13 +1,10 @@
 import weakref
 
-from numba import config
 from . import devices
-from .driver import CtypesModule
+from .driver import CtypesModule, USE_NV_BINDING
 
-USE_NV_BINDING = config.CUDA_USE_NVIDIA_BINDING
 
-
-class CuFuncProxy:
+class _CuFuncProxy:
     def __init__(self, module, cufunc):
         self._module = module
         self._cufunc = cufunc
@@ -73,7 +70,7 @@ def lazy_callback(callbacks, module, stream):
 
     def get_function(self, name):
         ctypesfunc = self._module.get_function(name)
-        return CuFuncProxy(self, ctypesfunc)
+        return _CuFuncProxy(self, ctypesfunc)
 
     def __getattr__(self, name):
         if name == "get_function":

numba_cuda/numba/cuda/dispatcher.py

@@ -1,4 +1,5 @@
 import numpy as np
+import weakref
 import os
 import re
 import sys
@@ -247,6 +248,7 @@ def _rebuild(cls, cooperative, name, signature, codelibrary,
         instance.lineinfo = lineinfo
         instance.call_helper = call_helper
         instance.extensions = extensions
+        instance.initialized = False
         return instance
 
     def _reduce_states(self):
@@ -1023,6 +1025,7 @@ def compile(self, sig):
                 raise RuntimeError("Compilation disabled")
 
             kernel = _Kernel(self.py_func, argtypes, **self.targetoptions)
+            weakref.finalize(kernel, _kernel_finalize_callback, kernel)
             # We call bind to force codegen, so that there is a cubin to cache
             kernel.bind()
             self._cache.save_overload(sig, kernel)
@@ -1148,3 +1151,18 @@ def _reduce_states(self):
         """
         return dict(py_func=self.py_func,
                     targetoptions=self.targetoptions)
+
+
+def _kernel_finalize_callback(kernel):
+    module = kernel.library.get_cufunc().module
+    try:
+        if driver.USE_NV_BINDING:
+            key = module.handle
+        else:
+            key = module.handle.value
+    except ReferenceError:
+        return
+
+    ctx = cuda.current_context()
+    if key in ctx.modules:
+        del ctx.modules[key]

numba_cuda/numba/cuda/tests/cudadrv/test_module_callbacks.py

@@ -1,3 +1,5 @@
+import gc
+
 import numpy as np
 
 from numba import cuda
@@ -31,10 +33,11 @@ def kernel():
         self.assertEqual(counter, 1)
         kernel[1, 1]() # cached
         self.assertEqual(counter, 1)
-        # del kernel
-        # gc.collect()
-        # cuda.current_context().deallocations.clear()
-        # self.assertEqual(counter, 0)
+        breakpoint()
+        del kernel
+        gc.collect()
+        cuda.current_context().deallocations.clear()
+        self.assertEqual(counter, 0)
         # We don't have a way to explicitly evict kernel and its modules at
         # the moment.