Support consuming NRT functions from external libraries

numba_cuda/numba/cuda/decorators.py

@@ -12,7 +12,7 @@
 
 
 def jit(func_or_sig=None, device=False, inline=False, link=[], debug=None,
-        opt=None, lineinfo=False, cache=False, **kws):
+        opt=None, lineinfo=False, cache=False, nrt=None, **kws):
     """
     JIT compile a Python function for CUDA GPUs.
 
@@ -114,6 +114,7 @@ def _jit(func):
             targetoptions['fastmath'] = fastmath
             targetoptions['device'] = device
             targetoptions['extensions'] = extensions
+            targetoptions['nrt'] = nrt   
 
             disp = CUDADispatcher(func, targetoptions=targetoptions)
 

numba_cuda/numba/cuda/dispatcher.py

@@ -1,6 +1,5 @@
 import numpy as np
 import os
-import re
 import sys
 import ctypes
 import functools
@@ -115,8 +114,9 @@ class _Kernel(serialize.ReduceMixin):
 
     @global_compiler_lock
     def __init__(self, py_func, argtypes, link=None, debug=False,
-                 lineinfo=False, inline=False, fastmath=False, extensions=None,
-                 max_registers=None, lto=False, opt=True, device=False):
+                 lineinfo=False, inline=False, fastmath=False,
+                 extensions=None, max_registers=None, lto=False, opt=True,
+                 device=False, nrt=None):
 
         if device:
             raise RuntimeError('Cannot compile a device function as a kernel')
@@ -179,6 +179,8 @@ def __init__(self, py_func, argtypes, link=None, debug=False,
         if self.cooperative:
             lib.needs_cudadevrt = True
 
+        self.nrt = nrt
+
         def link_to_library_functions(library_functions, library_path,
                                       prefix=None):
             """
@@ -234,21 +236,18 @@ def link_to_library_functions(library_functions, library_path,
         self.reload_init = []
 
     def maybe_link_nrt(self, link, tgt_ctx, asm):
-        if not tgt_ctx.enable_nrt:
+        if not tgt_ctx.enable_nrt or self.nrt is False:
             return
 
-        all_nrt = "|".join(self.NRT_functions)
-        pattern = (
-            r'\.extern\s+\.func\s+(?:\s*\(.+\)\s*)?('
-            + all_nrt + r')\s*\([^)]*\)\s*;'
-        )
-
-        nrt_in_asm = re.findall(pattern, asm)
-
         basedir = os.path.dirname(os.path.abspath(__file__))
-        if nrt_in_asm:
-            nrt_path = os.path.join(basedir, 'runtime', 'nrt.cu')
+        nrt_path = os.path.join(basedir, 'runtime', 'nrt.cu')
+
+        if self.nrt is True or (
+            self.nrt is None and any(fn in asm for fn in self.NRT_functions)
+        ):
             link.append(nrt_path)
+            if self.nrt is None:
+                self.nrt = True
 
     @property
     def library(self):
@@ -271,7 +270,7 @@ def argument_types(self):
 
     @classmethod
     def _rebuild(cls, cooperative, name, signature, codelibrary,
-                 debug, lineinfo, call_helper, extensions):
+                 debug, lineinfo, call_helper, extensions, nrt):
         """
         Rebuild an instance.
         """
@@ -289,6 +288,7 @@ def _rebuild(cls, cooperative, name, signature, codelibrary,
         instance.lineinfo = lineinfo
         instance.call_helper = call_helper
         instance.extensions = extensions
+        instance.nrt = nrt
         return instance
 
     def _reduce_states(self):
@@ -312,7 +312,7 @@ def bind(self):
 
         if (
             hasattr(self, "target_context")
-            and self.target_context.enable_nrt
+            and self.target_context.enable_nrt and self.nrt
             and config.CUDA_NRT_STATS
         ):
             rtsys.ensure_initialized()
@@ -721,7 +721,6 @@ class CUDADispatcher(Dispatcher, serialize.ReduceMixin):
     def __init__(self, py_func, targetoptions, pipeline_class=CUDACompiler):
         super().__init__(py_func, targetoptions=targetoptions,
                          pipeline_class=pipeline_class)
-
         # The following properties are for specialization of CUDADispatchers. A
         # specialized CUDADispatcher is one that is compiled for exactly one
         # set of argument types, and bypasses some argument type checking for

numba_cuda/numba/cuda/runtime/nrt.cu

@@ -4,21 +4,8 @@
 #include <cuda/atomic>
 
 #include "memsys.cuh"
+#include "nrt.cuh"
 
-typedef void (*NRT_dtor_function)(void* ptr, size_t size, void* info);
-typedef void (*NRT_dealloc_func)(void* ptr, void* dealloc_info);
-
-typedef struct MemInfo NRT_MemInfo;
-
-extern "C" {
-struct MemInfo {
-  cuda::atomic<size_t, cuda::thread_scope_device> refct;
-  NRT_dtor_function dtor;
-  void* dtor_info;
-  void* data;
-  size_t size;
-};
-}
 
 extern "C" __global__ void NRT_MemSys_set(NRT_MemSys *memsys_ptr)
 {
@@ -177,6 +164,7 @@ extern "C" __device__ void NRT_decref(NRT_MemInfo* mi)
   }
 }
 
+
 #endif
 
 extern "C" __device__ void NRT_incref(NRT_MemInfo* mi)

numba_cuda/numba/cuda/runtime/nrt.cuh

@@ -0,0 +1,23 @@
+#include <cuda/atomic>
+
+typedef void (*NRT_dtor_function)(void* ptr, size_t size, void* info);
+typedef void (*NRT_dealloc_func)(void* ptr, void* dealloc_info);
+
+extern "C" 
+struct MemInfo {
+  cuda::atomic<size_t, cuda::thread_scope_device> refct;
+  NRT_dtor_function dtor;
+  void* dtor_info;
+  void* data;
+  size_t size;
+};
+typedef struct MemInfo NRT_MemInfo;
+
+extern "C" __device__ void* NRT_Allocate(size_t size);
+extern "C" __device__ void NRT_MemInfo_init(NRT_MemInfo* mi,
+                                            void* data,
+                                            size_t size,
+                                            NRT_dtor_function dtor,
+                                            void* dtor_info);
+extern "C" __device__ void NRT_decref(NRT_MemInfo* mi);
+extern "C" __device__ void NRT_print_refct(NRT_MemInfo* mi);