From 2f62cb1d67e14acede3106e96192a7c25a2e5bba Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Fri, 7 Mar 2025 08:38:32 -0800
Subject: [PATCH 01/33] initial

---
 numba_cuda/numba/cuda/cudadrv/nvrtc.py |  2 +-
 numba_cuda/numba/cuda/dispatcher.py    |  6 ++++--
 numba_cuda/numba/cuda/runtime/nrt.cu   | 21 +++++++--------------
 numba_cuda/numba/cuda/runtime/nrt.cuh  | 23 +++++++++++++++++++++++
 4 files changed, 35 insertions(+), 17 deletions(-)
 create mode 100644 numba_cuda/numba/cuda/runtime/nrt.cuh

diff --git a/numba_cuda/numba/cuda/cudadrv/nvrtc.py b/numba_cuda/numba/cuda/cudadrv/nvrtc.py
index 5ab970c02..32b58632a 100644
--- a/numba_cuda/numba/cuda/cudadrv/nvrtc.py
+++ b/numba_cuda/numba/cuda/cudadrv/nvrtc.py
@@ -292,7 +292,7 @@ def compile(src, name, cc, ltoir=False):
     # Otherwise, if there's any content in the log, present it as a warning
     if log:
         msg = (f"NVRTC log messages whilst compiling {name}:\n\n{log}")
-        warnings.warn(msg)
+#        warnings.warn(msg)
 
     if ltoir:
         ltoir = nvrtc.get_lto(program)
diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py
index 69eca864e..6ad58a125 100644
--- a/numba_cuda/numba/cuda/dispatcher.py
+++ b/numba_cuda/numba/cuda/dispatcher.py
@@ -129,6 +129,8 @@ def __init__(self, py_func, argtypes, link=None, debug=False,
         if self.cooperative:
             lib.needs_cudadevrt = True
 
+        self._linked_nrt = False
+
         def link_to_library_functions(library_functions, library_path,
                                       prefix=None):
             """
@@ -201,11 +203,11 @@ def maybe_link_nrt(self, link, tgt_ctx, asm):
         )
 
         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')
             link.append(nrt_path)
+            self._linked_nrt = True
 
     @property
     def library(self):
@@ -267,7 +269,7 @@ def bind(self):
         """
         cufunc = self._codelibrary.get_cufunc()
 
-        if hasattr(self, "target_context") and self.target_context.enable_nrt:
+        if hasattr(self, "target_context") and self.target_context.enable_nrt and self._linked_nrt:
             rtsys.ensure_initialized()
             rtsys.set_memsys_to_module(cufunc.module)
             # We don't know which stream the kernel will be launched on, so
diff --git a/numba_cuda/numba/cuda/runtime/nrt.cu b/numba_cuda/numba/cuda/runtime/nrt.cu
index 879bf8d2f..887c5f89a 100644
--- a/numba_cuda/numba/cuda/runtime/nrt.cu
+++ b/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,12 @@ extern "C" __device__ void NRT_decref(NRT_MemInfo* mi)
   }
 }
 
+extern "C" int __device__ extern_NRT_Decref(int &retval, void *ptr)
+{
+    NRT_decref(reinterpret_cast<NRT_MemInfo*>(ptr));
+    return 0;
+}
+
 #endif
 
 extern "C" __device__ void NRT_incref(NRT_MemInfo* mi)
diff --git a/numba_cuda/numba/cuda/runtime/nrt.cuh b/numba_cuda/numba/cuda/runtime/nrt.cuh
new file mode 100644
index 000000000..894f45afd
--- /dev/null
+++ b/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);
+

From ec6bf1e8c308c284c3ddcf8320c5171a0082176c Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Mon, 17 Mar 2025 10:01:47 -0700
Subject: [PATCH 02/33] pass compute sanitizer on the python side

---
 numba_cuda/numba/cuda/decorators.py   |  3 ++-
 numba_cuda/numba/cuda/dispatcher.py   | 32 ++++++++++++---------------
 numba_cuda/numba/cuda/runtime/nrt.cu  | 11 +++++----
 numba_cuda/numba/cuda/runtime/nrt.cuh |  2 +-
 4 files changed, 24 insertions(+), 24 deletions(-)

diff --git a/numba_cuda/numba/cuda/decorators.py b/numba_cuda/numba/cuda/decorators.py
index db62fb96a..dbea27599 100644
--- a/numba_cuda/numba/cuda/decorators.py
+++ b/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)
 
diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py
index 6ad58a125..7ba8b8897 100644
--- a/numba_cuda/numba/cuda/dispatcher.py
+++ b/numba_cuda/numba/cuda/dispatcher.py
@@ -66,7 +66,7 @@ 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):
+                 max_registers=None, lto=False, opt=True, device=False, nrt=None):
 
         if device:
             raise RuntimeError('Cannot compile a device function as a kernel')
@@ -129,7 +129,7 @@ def __init__(self, py_func, argtypes, link=None, debug=False,
         if self.cooperative:
             lib.needs_cudadevrt = True
 
-        self._linked_nrt = False
+        self.nrt = nrt
 
         def link_to_library_functions(library_functions, library_path,
                                       prefix=None):
@@ -193,21 +193,16 @@ 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)
-            self._linked_nrt = True
+            if self.nrt is None:
+                self.nrt = True
 
     @property
     def library(self):
@@ -230,7 +225,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.
         """
@@ -248,6 +243,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):
@@ -269,7 +265,7 @@ def bind(self):
         """
         cufunc = self._codelibrary.get_cufunc()
 
-        if hasattr(self, "target_context") and self.target_context.enable_nrt and self._linked_nrt:
+        if hasattr(self, "target_context") and self.target_context.enable_nrt and self.nrt:
             rtsys.ensure_initialized()
             rtsys.set_memsys_to_module(cufunc.module)
             # We don't know which stream the kernel will be launched on, so
@@ -676,12 +672,12 @@ 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
         # faster kernel launches.
-
+        #if targetoptions['nrt']:
+        #    breakpoint()
         # Is this a specialized dispatcher?
         self._specialized = False
 
diff --git a/numba_cuda/numba/cuda/runtime/nrt.cu b/numba_cuda/numba/cuda/runtime/nrt.cu
index 887c5f89a..e5ca33ca7 100644
--- a/numba_cuda/numba/cuda/runtime/nrt.cu
+++ b/numba_cuda/numba/cuda/runtime/nrt.cu
@@ -159,15 +159,18 @@ extern "C" __device__ void* NRT_Allocate_External(size_t size) {
 extern "C" __device__ void NRT_decref(NRT_MemInfo* mi)
 {
   if (mi != NULL) {
+    //printf("NRT_decref %p\n", mi);
+    //printf("NRT_decref refct %d\n", mi->refct);
     mi->refct--;
     if (mi->refct == 0) { NRT_MemInfo_call_dtor(mi); }
   }
 }
 
-extern "C" int __device__ extern_NRT_Decref(int &retval, void *ptr)
-{
-    NRT_decref(reinterpret_cast<NRT_MemInfo*>(ptr));
-    return 0;
+extern "C" __device__ void NRT_print_refct(NRT_MemInfo* mi){
+  if (mi != NULL) {
+    printf("NRT_print_refct %p\n", mi);
+    printf("NRT_print_refct refct %d\n", mi->refct);
+  }
 }
 
 #endif
diff --git a/numba_cuda/numba/cuda/runtime/nrt.cuh b/numba_cuda/numba/cuda/runtime/nrt.cuh
index 894f45afd..7dda5fbe9 100644
--- a/numba_cuda/numba/cuda/runtime/nrt.cuh
+++ b/numba_cuda/numba/cuda/runtime/nrt.cuh
@@ -20,4 +20,4 @@ extern "C" __device__ void NRT_MemInfo_init(NRT_MemInfo* mi,
                                             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);

From f52a15c7a84c727fd56a2e2888dd081921ff93e3 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Tue, 18 Mar 2025 18:52:36 -0700
Subject: [PATCH 03/33] clean

---
 numba_cuda/numba/cuda/dispatcher.py  | 20 ++++++++++++--------
 numba_cuda/numba/cuda/runtime/nrt.cu |  8 --------
 2 files changed, 12 insertions(+), 16 deletions(-)

diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py
index 7ba8b8897..6fcb2a9f6 100644
--- a/numba_cuda/numba/cuda/dispatcher.py
+++ b/numba_cuda/numba/cuda/dispatcher.py
@@ -1,6 +1,5 @@
 import numpy as np
 import os
-import re
 import sys
 import ctypes
 import functools
@@ -65,8 +64,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, nrt=None):
+                 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')
@@ -195,11 +195,13 @@ def link_to_library_functions(library_functions, library_path,
     def maybe_link_nrt(self, link, tgt_ctx, asm):
         if not tgt_ctx.enable_nrt or self.nrt is False:
             return
-        
+
         basedir = os.path.dirname(os.path.abspath(__file__))
         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)):
+        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
@@ -265,7 +267,10 @@ def bind(self):
         """
         cufunc = self._codelibrary.get_cufunc()
 
-        if hasattr(self, "target_context") and self.target_context.enable_nrt and self.nrt:
+        if (
+            hasattr(self, "target_context")
+            and self.target_context.enable_nrt and self.nrt
+        ):
             rtsys.ensure_initialized()
             rtsys.set_memsys_to_module(cufunc.module)
             # We don't know which stream the kernel will be launched on, so
@@ -676,8 +681,7 @@ def __init__(self, py_func, targetoptions, pipeline_class=CUDACompiler):
         # specialized CUDADispatcher is one that is compiled for exactly one
         # set of argument types, and bypasses some argument type checking for
         # faster kernel launches.
-        #if targetoptions['nrt']:
-        #    breakpoint()
+
         # Is this a specialized dispatcher?
         self._specialized = False
 
diff --git a/numba_cuda/numba/cuda/runtime/nrt.cu b/numba_cuda/numba/cuda/runtime/nrt.cu
index e5ca33ca7..abf86a7f8 100644
--- a/numba_cuda/numba/cuda/runtime/nrt.cu
+++ b/numba_cuda/numba/cuda/runtime/nrt.cu
@@ -159,19 +159,11 @@ extern "C" __device__ void* NRT_Allocate_External(size_t size) {
 extern "C" __device__ void NRT_decref(NRT_MemInfo* mi)
 {
   if (mi != NULL) {
-    //printf("NRT_decref %p\n", mi);
-    //printf("NRT_decref refct %d\n", mi->refct);
     mi->refct--;
     if (mi->refct == 0) { NRT_MemInfo_call_dtor(mi); }
   }
 }
 
-extern "C" __device__ void NRT_print_refct(NRT_MemInfo* mi){
-  if (mi != NULL) {
-    printf("NRT_print_refct %p\n", mi);
-    printf("NRT_print_refct refct %d\n", mi->refct);
-  }
-}
 
 #endif
 

From 01be5c2811d79debb6bcd9ef614ece8388e95d6a Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Tue, 18 Mar 2025 18:55:43 -0700
Subject: [PATCH 04/33] reset files

---
 numba_cuda/numba/cuda/cudadrv/nvrtc.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/numba_cuda/numba/cuda/cudadrv/nvrtc.py b/numba_cuda/numba/cuda/cudadrv/nvrtc.py
index 32b58632a..5ab970c02 100644
--- a/numba_cuda/numba/cuda/cudadrv/nvrtc.py
+++ b/numba_cuda/numba/cuda/cudadrv/nvrtc.py
@@ -292,7 +292,7 @@ def compile(src, name, cc, ltoir=False):
     # Otherwise, if there's any content in the log, present it as a warning
     if log:
         msg = (f"NVRTC log messages whilst compiling {name}:\n\n{log}")
-#        warnings.warn(msg)
+        warnings.warn(msg)
 
     if ltoir:
         ltoir = nvrtc.get_lto(program)

From 0b23a097fb391424ea830975ab6bbc06e6721d6e Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Wed, 19 Mar 2025 12:57:26 -0700
Subject: [PATCH 05/33] check all ptx files in maybe_link_nrt

---
 numba_cuda/numba/cuda/dispatcher.py | 24 +++++++++++++++---------
 1 file changed, 15 insertions(+), 9 deletions(-)

diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py
index e302750f5..886d6b2ce 100644
--- a/numba_cuda/numba/cuda/dispatcher.py
+++ b/numba_cuda/numba/cuda/dispatcher.py
@@ -208,7 +208,6 @@ def link_to_library_functions(library_functions, library_path,
         link_to_library_functions(cuda_fp16_math_funcs,
                                   'cpp_function_wrappers.cu',
                                   '__numba_wrapper_')
-
         self.maybe_link_nrt(link, tgt_ctx, asm)
 
         for obj in get_cres_link_objects(cres):
@@ -239,15 +238,22 @@ def maybe_link_nrt(self, link, tgt_ctx, asm):
         if not tgt_ctx.enable_nrt or self.nrt is False:
             return
 
-        basedir = os.path.dirname(os.path.abspath(__file__))
-        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)
-        ):
+        nrt_in_asm = lambda asm: any(fn in asm for fn in self.NRT_functions)
+        link_nrt = nrt_in_asm(asm)
+        if not link_nrt:
+            for file in link:
+                if file.endswith('ptx'):
+                    with open(file) as f:
+                        asm = f.read()
+                        # TODO: double read
+                    if nrt_in_asm(asm):
+                        link_nrt = True
+                        break
+
+        if link_nrt:
+            basedir = os.path.dirname(os.path.abspath(__file__))
+            nrt_path = os.path.join(basedir, 'runtime', 'nrt.cu')
             link.append(nrt_path)
-            if self.nrt is None:
-                self.nrt = True
 
     @property
     def library(self):

From dd568e74ded42adf99aed5cdc7f809872d895ce6 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Wed, 19 Mar 2025 13:01:41 -0700
Subject: [PATCH 06/33] detect nrt in all linked ptx files

---
 numba_cuda/numba/cuda/decorators.py | 3 +--
 numba_cuda/numba/cuda/dispatcher.py | 8 +++-----
 2 files changed, 4 insertions(+), 7 deletions(-)

diff --git a/numba_cuda/numba/cuda/decorators.py b/numba_cuda/numba/cuda/decorators.py
index dbea27599..db62fb96a 100644
--- a/numba_cuda/numba/cuda/decorators.py
+++ b/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, nrt=None, **kws):
+        opt=None, lineinfo=False, cache=False, **kws):
     """
     JIT compile a Python function for CUDA GPUs.
 
@@ -114,7 +114,6 @@ def _jit(func):
             targetoptions['fastmath'] = fastmath
             targetoptions['device'] = device
             targetoptions['extensions'] = extensions
-            targetoptions['nrt'] = nrt   
 
             disp = CUDADispatcher(func, targetoptions=targetoptions)
 
diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py
index 886d6b2ce..76feaa2c7 100644
--- a/numba_cuda/numba/cuda/dispatcher.py
+++ b/numba_cuda/numba/cuda/dispatcher.py
@@ -116,7 +116,7 @@ class _Kernel(serialize.ReduceMixin):
     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, nrt=None):
+                 device=False):
 
         if device:
             raise RuntimeError('Cannot compile a device function as a kernel')
@@ -179,8 +179,6 @@ 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):
             """
@@ -235,7 +233,7 @@ 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 or self.nrt is False:
+        if not tgt_ctx.enable_nrt:
             return
 
         nrt_in_asm = lambda asm: any(fn in asm for fn in self.NRT_functions)
@@ -318,7 +316,7 @@ def bind(self):
 
         if (
             hasattr(self, "target_context")
-            and self.target_context.enable_nrt and self.nrt
+            and self.target_context.enable_nrt
             and config.CUDA_NRT_STATS
         ):
             rtsys.ensure_initialized()

From af9045dfccf8423c83e1d74e0923ace5416ee53e Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Mon, 24 Mar 2025 05:58:54 -0700
Subject: [PATCH 07/33] read_file

---
 numba_cuda/numba/cuda/cudadrv/driver.py | 15 +++++++++------
 numba_cuda/numba/cuda/dispatcher.py     |  7 ++-----
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/numba_cuda/numba/cuda/cudadrv/driver.py b/numba_cuda/numba/cuda/cudadrv/driver.py
index 1641bf779..99acaa10b 100644
--- a/numba_cuda/numba/cuda/cudadrv/driver.py
+++ b/numba_cuda/numba/cuda/cudadrv/driver.py
@@ -103,6 +103,12 @@ def make_logger():
     return logger
 
 
+@functools.lru_cache(maxsize=None)
+def read_file(filepath, how='r'):
+    with open(filepath, how) as f:
+        return f.read()
+
+
 class DeadMemoryError(RuntimeError):
     pass
 
@@ -2654,8 +2660,7 @@ def add_file(self, path, kind):
         """Add code from a file to the link"""
 
     def add_cu_file(self, path):
-        with open(path, 'rb') as f:
-            cu = f.read()
+        cu = read_file(path, how='rb')
         self.add_cu(cu, os.path.basename(path))
 
     def add_file_guess_ext(self, path_or_code, ignore_nonlto=False):
@@ -2804,8 +2809,7 @@ def add_file(self, path, kind):
             raise ImportError(_MVC_ERROR_MESSAGE) from err
 
         try:
-            with open(path, 'rb') as f:
-                data = f.read()
+            data = read_file(path, how='rb')
         except FileNotFoundError:
             raise LinkerError(f'{path} not found')
 
@@ -3088,8 +3092,7 @@ def add_object(self, obj, name="<external-object>"):
 
     def add_file(self, path, kind):
         try:
-            with open(path, "rb") as f:
-                data = f.read()
+            data = read_file(path, 'rb')
         except FileNotFoundError:
             raise LinkerError(f"{path} not found")
 
diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py
index 76feaa2c7..1cba785d9 100644
--- a/numba_cuda/numba/cuda/dispatcher.py
+++ b/numba_cuda/numba/cuda/dispatcher.py
@@ -241,9 +241,7 @@ def maybe_link_nrt(self, link, tgt_ctx, asm):
         if not link_nrt:
             for file in link:
                 if file.endswith('ptx'):
-                    with open(file) as f:
-                        asm = f.read()
-                        # TODO: double read
+                    asm = driver.read_file(file)
                     if nrt_in_asm(asm):
                         link_nrt = True
                         break
@@ -274,7 +272,7 @@ def argument_types(self):
 
     @classmethod
     def _rebuild(cls, cooperative, name, signature, codelibrary,
-                 debug, lineinfo, call_helper, extensions, nrt):
+                 debug, lineinfo, call_helper, extensions):
         """
         Rebuild an instance.
         """
@@ -292,7 +290,6 @@ 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):

From f2e9b2972880159289c6a2968f8747befd6abf94 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Mon, 31 Mar 2025 08:04:30 -0700
Subject: [PATCH 08/33] tests, fill out headerA

---
 numba_cuda/numba/cuda/runtime/nrt.cu        |  3 -
 numba_cuda/numba/cuda/runtime/nrt.cuh       | 20 ++++++-
 numba_cuda/numba/cuda/runtime/nrt.py        |  5 ++
 numba_cuda/numba/cuda/tests/nrt/test_nrt.py | 64 ++++++++++++++++++++-
 4 files changed, 86 insertions(+), 6 deletions(-)

diff --git a/numba_cuda/numba/cuda/runtime/nrt.cu b/numba_cuda/numba/cuda/runtime/nrt.cu
index abf86a7f8..e9f0aefa7 100644
--- a/numba_cuda/numba/cuda/runtime/nrt.cu
+++ b/numba_cuda/numba/cuda/runtime/nrt.cu
@@ -12,9 +12,6 @@ extern "C" __global__ void NRT_MemSys_set(NRT_MemSys *memsys_ptr)
   TheMSys = memsys_ptr;
 }
 
-static __device__ void *nrt_allocate_meminfo_and_data_align(size_t size, unsigned align, NRT_MemInfo **mi);
-static __device__ void *nrt_allocate_meminfo_and_data(size_t size, NRT_MemInfo **mi_out);
-extern "C" __device__ void* NRT_Allocate_External(size_t size);
 
 extern "C" __device__ void* NRT_Allocate(size_t size)
 {
diff --git a/numba_cuda/numba/cuda/runtime/nrt.cuh b/numba_cuda/numba/cuda/runtime/nrt.cuh
index 7dda5fbe9..3f0070b39 100644
--- a/numba_cuda/numba/cuda/runtime/nrt.cuh
+++ b/numba_cuda/numba/cuda/runtime/nrt.cuh
@@ -19,5 +19,23 @@ extern "C" __device__ void NRT_MemInfo_init(NRT_MemInfo* mi,
                                             size_t size,
                                             NRT_dtor_function dtor,
                                             void* dtor_info);
+static __device__ void *nrt_allocate_meminfo_and_data_align(size_t size, unsigned align, NRT_MemInfo **mi);
+static __device__ void *nrt_allocate_meminfo_and_data(size_t size, NRT_MemInfo **mi_out);
+extern "C" __device__ void* NRT_Allocate_External(size_t size);
 extern "C" __device__ void NRT_decref(NRT_MemInfo* mi);
-extern "C" __device__ void NRT_print_refct(NRT_MemInfo* mi);
+extern "C" __device__ void NRT_incref(NRT_MemInfo* mi);
+extern "C" __device__ void* NRT_Allocate_External(size_t size);
+static __device__ void *nrt_allocate_meminfo_and_data(size_t size, NRT_MemInfo **mi_out);
+static __device__ void *nrt_allocate_meminfo_and_data_align(size_t size, unsigned align, NRT_MemInfo **mi);
+extern "C" __device__ NRT_MemInfo *NRT_MemInfo_alloc_aligned(size_t size, unsigned align);
+extern "C" __device__ void* NRT_MemInfo_data_fast(NRT_MemInfo *mi);
+extern "C" __device__ void NRT_MemInfo_call_dtor(NRT_MemInfo* mi);
+extern "C" __device__ void NRT_MemInfo_destroy(NRT_MemInfo* mi);
+extern "C" __device__ void NRT_dealloc(NRT_MemInfo* mi);
+extern "C" __device__ void NRT_Free(void* ptr);
+extern "C" __device__ NRT_MemInfo* NRT_MemInfo_new(void* data, size_t size, NRT_dtor_function dtor, void* dtor_info);
+extern "C" __device__ void NRT_MemInfo_init(NRT_MemInfo* mi,
+                                            void* data,
+                                            size_t size,
+                                            NRT_dtor_function dtor,
+                                            void* dtor_info);
diff --git a/numba_cuda/numba/cuda/runtime/nrt.py b/numba_cuda/numba/cuda/runtime/nrt.py
index 0e981baac..1c487fd90 100644
--- a/numba_cuda/numba/cuda/runtime/nrt.py
+++ b/numba_cuda/numba/cuda/runtime/nrt.py
@@ -29,6 +29,11 @@
     config.CUDA_ENABLE_NRT = ENABLE_NRT
 
 
+def get_include():
+    """Return the include path for the NRT header"""
+    return os.path.dirname(os.path.abspath(__file__))
+
+
 # Protect method to ensure NRT memory allocation and initialization
 def _alloc_init_guard(method):
     """
diff --git a/numba_cuda/numba/cuda/tests/nrt/test_nrt.py b/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
index acdaa0c8d..e84a235d8 100644
--- a/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
+++ b/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
@@ -4,12 +4,53 @@
 import numpy as np
 import unittest
 from numba.cuda.testing import CUDATestCase
-
 from numba.tests.support import run_in_subprocess, override_config
-
+from numba.cuda import get_current_device
+from numba.cuda.cudadrv.nvrtc import compile
+from numba import types
+from numba.cuda.cudadecl import registry as cuda_decl_registry
+from numba.core.typing import signature
+from numba.cuda.cudaimpl import lower as cuda_lower
 from numba import cuda
 from numba.cuda.runtime.nrt import rtsys
 
+from numba.core.typing.templates import AbstractTemplate
+
+
+def allocate_shim():
+    pass
+
+
+@cuda_decl_registry.register_global(allocate_shim)
+class AllocateShimImpl(AbstractTemplate):
+    def generic(self, args, kws):
+        return signature(types.void)
+
+
+# let there be a __device__ function expecting
+# a string_view* and returning size_type
+_allocate_shim = cuda.declare_device("extern_func", types.int32())
+
+
+# wrapper to turn the above into a python callable
+def call_allocate_shim():
+    return _allocate_shim()
+
+
+@cuda_lower(allocate_shim)
+def allocate_shim_impl(context, builder, sig, args):
+    sig_ = types.int32()
+    # call the external function, passing the pointer
+    result = context.compile_internal(
+        builder,
+        call_allocate_shim,
+        sig_,
+        (),
+    )
+
+    # return len(&input)
+    return result
+
 
 class TestNrtBasic(CUDATestCase):
     def run(self, result=None):
@@ -229,6 +270,25 @@ def test_nrt_explicit_stats_query_raises_exception_when_disabled(self):
                     stats_func()
                 self.assertIn("NRT stats are disabled.", str(raises.exception))
 
+    def test_nrt_detect_linked_ptx_file(self):
+        src = """;
+                 extern "C" __device__ void* NRT_Allocate(size_t size);
+                 extern "C" __device__ int extern_func(int* nb_retval){
+                     NRT_Allocate(1);
+                     return 0;
+                 }
+        """
+        cc = get_current_device().compute_capability
+        ptx, _ = compile(src, 'external_nrt.cu', cc)
+        with open('external_nrt.ptx', 'w') as f:
+            f.write(ptx)
+
+        @cuda.jit(link=['external_nrt.ptx'])
+        def kernel():
+            allocate_shim()
+
+        kernel[1,1]()
+
 
 if __name__ == '__main__':
     unittest.main()

From e65088fe5afb171bb533f68a2a4d1e8c9f66cf81 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Mon, 31 Mar 2025 08:48:59 -0700
Subject: [PATCH 09/33] temp file

---
 numba_cuda/numba/cuda/tests/nrt/test_nrt.py | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/numba_cuda/numba/cuda/tests/nrt/test_nrt.py b/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
index c53bbcfff..15a94f7d7 100644
--- a/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
+++ b/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
@@ -13,7 +13,7 @@
 from numba.cuda.cudaimpl import lower as cuda_lower
 from numba import cuda
 from numba.cuda.runtime.nrt import rtsys
-
+import tempfile
 from numba.core.typing.templates import AbstractTemplate
 
 
@@ -127,14 +127,20 @@ def test_nrt_detect_linked_ptx_file(self):
         """
         cc = get_current_device().compute_capability
         ptx, _ = compile(src, 'external_nrt.cu', cc)
-        with open('external_nrt.ptx', 'w') as f:
+        with tempfile.NamedTemporaryFile(
+            delete=True,
+            mode="w+",
+            suffix='.ptx',
+            encoding='utf-8'
+        ) as f:
             f.write(ptx)
+            f.flush()
 
-        @cuda.jit(link=['external_nrt.ptx'])
-        def kernel():
-            allocate_shim()
+            @cuda.jit(link=[f.name])
+            def kernel():
+                allocate_shim()
 
-        kernel[1,1]()
+            kernel[1,1]()
 
 
 class TestNrtStatistics(CUDATestCase):

From 5bd9b7a4dd8ef01d600b4c00d43a3acd55d4869a Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Wed, 2 Apr 2025 07:55:10 -0700
Subject: [PATCH 10/33] move and rename read_file

---
 numba_cuda/numba/cuda/cudadrv/driver.py | 14 ++++----------
 numba_cuda/numba/cuda/utils.py          |  7 +++++++
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/numba_cuda/numba/cuda/cudadrv/driver.py b/numba_cuda/numba/cuda/cudadrv/driver.py
index 19c20d070..2712ed4a9 100644
--- a/numba_cuda/numba/cuda/cudadrv/driver.py
+++ b/numba_cuda/numba/cuda/cudadrv/driver.py
@@ -39,7 +39,7 @@
 from .drvapi import cu_occupancy_b2d_size, cu_stream_callback_pyobj, cu_uuid
 from .mappings import FILE_EXTENSION_MAP
 from .linkable_code import LinkableCode, LTOIR, Fatbin, Object
-from numba.cuda.utils import _readenv
+from numba.cuda.utils import _readenv, cached_file_read
 from numba.cuda.cudadrv import enums, drvapi, nvrtc
 
 try:
@@ -103,12 +103,6 @@ def make_logger():
     return logger
 
 
-@functools.lru_cache(maxsize=None)
-def read_file(filepath, how='r'):
-    with open(filepath, how) as f:
-        return f.read()
-
-
 class DeadMemoryError(RuntimeError):
     pass
 
@@ -2663,7 +2657,7 @@ def add_file(self, path, kind):
         """Add code from a file to the link"""
 
     def add_cu_file(self, path):
-        cu = read_file(path, how='rb')
+        cu = cached_file_read(path, how='rb')
         self.add_cu(cu, os.path.basename(path))
 
     def add_file_guess_ext(self, path_or_code, ignore_nonlto=False):
@@ -2812,7 +2806,7 @@ def add_file(self, path, kind):
             raise ImportError(_MVC_ERROR_MESSAGE) from err
 
         try:
-            data = read_file(path, how='rb')
+            data = cached_file_read(path, how='rb')
         except FileNotFoundError:
             raise LinkerError(f'{path} not found')
 
@@ -3095,7 +3089,7 @@ def add_object(self, obj, name="<external-object>"):
 
     def add_file(self, path, kind):
         try:
-            data = read_file(path, 'rb')
+            data = cached_file_read(path, 'rb')
         except FileNotFoundError:
             raise LinkerError(f"{path} not found")
 
diff --git a/numba_cuda/numba/cuda/utils.py b/numba_cuda/numba/cuda/utils.py
index 48ce2b011..ebf0c9d6b 100644
--- a/numba_cuda/numba/cuda/utils.py
+++ b/numba_cuda/numba/cuda/utils.py
@@ -1,6 +1,7 @@
 import os
 import warnings
 import traceback
+import functools
 
 
 def _readenv(name, ctor, default):
@@ -20,3 +21,9 @@ def _readenv(name, ctor, default):
             RuntimeWarning
         )
         return default
+
+
+@functools.lru_cache(maxsize=None)
+def cached_file_read(filepath, how='r'):
+    with open(filepath, how) as f:
+        return f.read()

From e3042a643c33bdef800a3e78b8db8f0d3daf4ca9 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <53796099+brandon-b-miller@users.noreply.github.com>
Date: Wed, 2 Apr 2025 09:55:47 -0500
Subject: [PATCH 11/33] Update numba_cuda/numba/cuda/tests/nrt/test_nrt.py

Co-authored-by: Michael Wang <13521008+isVoid@users.noreply.github.com>
---
 numba_cuda/numba/cuda/tests/nrt/test_nrt.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/numba_cuda/numba/cuda/tests/nrt/test_nrt.py b/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
index 15a94f7d7..7a6e30154 100644
--- a/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
+++ b/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
@@ -48,7 +48,6 @@ def allocate_shim_impl(context, builder, sig, args):
         (),
     )
 
-    # return len(&input)
     return result
 
 

From 501b4da47ef9db0f3b3262a5429417a40a0b9988 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Wed, 2 Apr 2025 08:15:51 -0700
Subject: [PATCH 12/33] get_include() in tests

---
 numba_cuda/numba/cuda/dispatcher.py         | 3 ++-
 numba_cuda/numba/cuda/tests/nrt/test_nrt.py | 6 +++---
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py
index 1cba785d9..c9bcf531b 100644
--- a/numba_cuda/numba/cuda/dispatcher.py
+++ b/numba_cuda/numba/cuda/dispatcher.py
@@ -23,6 +23,7 @@
                                normalize_kernel_dimensions)
 from numba.cuda import types as cuda_types
 from numba.cuda.runtime.nrt import rtsys
+from numba.cuda.utils import cached_file_read
 
 from numba import cuda
 from numba import _dispatcher
@@ -241,7 +242,7 @@ def maybe_link_nrt(self, link, tgt_ctx, asm):
         if not link_nrt:
             for file in link:
                 if file.endswith('ptx'):
-                    asm = driver.read_file(file)
+                    asm = cached_file_read(file)
                     if nrt_in_asm(asm):
                         link_nrt = True
                         break
diff --git a/numba_cuda/numba/cuda/tests/nrt/test_nrt.py b/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
index 7a6e30154..7361e205f 100644
--- a/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
+++ b/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
@@ -12,7 +12,7 @@
 from numba.core.typing import signature
 from numba.cuda.cudaimpl import lower as cuda_lower
 from numba import cuda
-from numba.cuda.runtime.nrt import rtsys
+from numba.cuda.runtime.nrt import rtsys, get_include
 import tempfile
 from numba.core.typing.templates import AbstractTemplate
 
@@ -117,8 +117,8 @@ def g(out_ary):
         self.assertEqual(out_ary[0], 1)
 
     def test_nrt_detect_linked_ptx_file(self):
-        src = """;
-                 extern "C" __device__ void* NRT_Allocate(size_t size);
+        src = f"#include <{get_include()}/nrt.cuh>"
+        src += """;
                  extern "C" __device__ int extern_func(int* nb_retval){
                      NRT_Allocate(1);
                      return 0;

From ff478f462ebf186f59a0104db9762ad396cc276f Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Wed, 2 Apr 2025 08:57:29 -0700
Subject: [PATCH 13/33] handle LinkableCode

---
 numba_cuda/numba/cuda/dispatcher.py | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py
index c9bcf531b..f8727d571 100644
--- a/numba_cuda/numba/cuda/dispatcher.py
+++ b/numba_cuda/numba/cuda/dispatcher.py
@@ -17,6 +17,7 @@
 from numba.cuda.compiler import (compile_cuda, CUDACompiler, kernel_fixup,
                                  ExternFunction)
 from numba.cuda.cudadrv import driver
+from numba.cuda.cudadrv.linkable_code import LinkableCode
 from numba.cuda.cudadrv.devices import get_context
 from numba.cuda.descriptor import cuda_target
 from numba.cuda.errors import (missing_launch_config_msg,
@@ -241,11 +242,17 @@ def maybe_link_nrt(self, link, tgt_ctx, asm):
         link_nrt = nrt_in_asm(asm)
         if not link_nrt:
             for file in link:
-                if file.endswith('ptx'):
-                    asm = cached_file_read(file)
-                    if nrt_in_asm(asm):
-                        link_nrt = True
-                        break
+                if isinstance(file, LinkableCode):
+                    asm = file.data.decode('utf-8')
+                else:
+                    if file.endswith('ptx') or file.endswith('cu'):
+                        asm = cached_file_read(file)
+                    else:
+                        # Not a PTX or CUDA source, skip it
+                        continue
+                if nrt_in_asm(asm):
+                    link_nrt = True
+                    break
 
         if link_nrt:
             basedir = os.path.dirname(os.path.abspath(__file__))

From 43b647b77b2cb1e2da2b0f5e6884d2e46da59820 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Mon, 7 Apr 2025 08:47:35 -0700
Subject: [PATCH 14/33] convert to LinkableCode

---
 numba_cuda/numba/cuda/dispatcher.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py
index f8727d571..5d56483ee 100644
--- a/numba_cuda/numba/cuda/dispatcher.py
+++ b/numba_cuda/numba/cuda/dispatcher.py
@@ -17,7 +17,7 @@
 from numba.cuda.compiler import (compile_cuda, CUDACompiler, kernel_fixup,
                                  ExternFunction)
 from numba.cuda.cudadrv import driver
-from numba.cuda.cudadrv.linkable_code import LinkableCode
+from numba.cuda.cudadrv.linkable_code import LinkableCode, PTXSource, CUSource
 from numba.cuda.cudadrv.devices import get_context
 from numba.cuda.descriptor import cuda_target
 from numba.cuda.errors import (missing_launch_config_msg,
@@ -241,12 +241,18 @@ def maybe_link_nrt(self, link, tgt_ctx, asm):
         nrt_in_asm = lambda asm: any(fn in asm for fn in self.NRT_functions)
         link_nrt = nrt_in_asm(asm)
         if not link_nrt:
-            for file in link:
+            for i, file in enumerate(link):
                 if isinstance(file, LinkableCode):
                     asm = file.data.decode('utf-8')
                 else:
                     if file.endswith('ptx') or file.endswith('cu'):
                         asm = cached_file_read(file)
+                        lc = (
+                            PTXSource(asm.encode())
+                            if file.endswith('ptx')
+                            else CUSource(asm.encode())
+                        )
+                        link[i] = lc
                     else:
                         # Not a PTX or CUDA source, skip it
                         continue

From eab1d9e8cce13fde9cfb6e74c8a6d737550e37ef Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Mon, 7 Apr 2025 09:26:43 -0700
Subject: [PATCH 15/33] Revert "convert to LinkableCode"

This reverts commit 43b647b77b2cb1e2da2b0f5e6884d2e46da59820.
---
 numba_cuda/numba/cuda/dispatcher.py | 10 ++--------
 1 file changed, 2 insertions(+), 8 deletions(-)

diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py
index 5d56483ee..f8727d571 100644
--- a/numba_cuda/numba/cuda/dispatcher.py
+++ b/numba_cuda/numba/cuda/dispatcher.py
@@ -17,7 +17,7 @@
 from numba.cuda.compiler import (compile_cuda, CUDACompiler, kernel_fixup,
                                  ExternFunction)
 from numba.cuda.cudadrv import driver
-from numba.cuda.cudadrv.linkable_code import LinkableCode, PTXSource, CUSource
+from numba.cuda.cudadrv.linkable_code import LinkableCode
 from numba.cuda.cudadrv.devices import get_context
 from numba.cuda.descriptor import cuda_target
 from numba.cuda.errors import (missing_launch_config_msg,
@@ -241,18 +241,12 @@ def maybe_link_nrt(self, link, tgt_ctx, asm):
         nrt_in_asm = lambda asm: any(fn in asm for fn in self.NRT_functions)
         link_nrt = nrt_in_asm(asm)
         if not link_nrt:
-            for i, file in enumerate(link):
+            for file in link:
                 if isinstance(file, LinkableCode):
                     asm = file.data.decode('utf-8')
                 else:
                     if file.endswith('ptx') or file.endswith('cu'):
                         asm = cached_file_read(file)
-                        lc = (
-                            PTXSource(asm.encode())
-                            if file.endswith('ptx')
-                            else CUSource(asm.encode())
-                        )
-                        link[i] = lc
                     else:
                         # Not a PTX or CUDA source, skip it
                         continue

From 0dff0b6705b4250d981dbbc3f6e35b91be21ea41 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Tue, 8 Apr 2025 06:23:00 -0700
Subject: [PATCH 16/33] style

---
 numba_cuda/numba/cuda/runtime/nrt.cuh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/numba_cuda/numba/cuda/runtime/nrt.cuh b/numba_cuda/numba/cuda/runtime/nrt.cuh
index 3f0070b39..8944733f1 100644
--- a/numba_cuda/numba/cuda/runtime/nrt.cuh
+++ b/numba_cuda/numba/cuda/runtime/nrt.cuh
@@ -3,7 +3,7 @@
 typedef void (*NRT_dtor_function)(void* ptr, size_t size, void* info);
 typedef void (*NRT_dealloc_func)(void* ptr, void* dealloc_info);
 
-extern "C" 
+extern "C"
 struct MemInfo {
   cuda::atomic<size_t, cuda::thread_scope_device> refct;
   NRT_dtor_function dtor;

From 11a014aa23445320846af6550c8a8c93a6a5b3b7 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Wed, 9 Apr 2025 15:16:26 -0700
Subject: [PATCH 17/33] LinkableCode accepts nrt parameter

---
 .../numba/cuda/cudadrv/linkable_code.py       |  3 +-
 numba_cuda/numba/cuda/dispatcher.py           |  8 +-
 numba_cuda/numba/cuda/tests/nrt/test_nrt.py   | 81 +++++++++++++++++++
 .../tests/test_binary_generation/Makefile     |  9 +++
 .../test_binary_generation/nrt_extern.cu      |  6 ++
 5 files changed, 104 insertions(+), 3 deletions(-)
 create mode 100644 numba_cuda/numba/cuda/tests/test_binary_generation/nrt_extern.cu

diff --git a/numba_cuda/numba/cuda/cudadrv/linkable_code.py b/numba_cuda/numba/cuda/cudadrv/linkable_code.py
index d1d715930..a015b4a72 100644
--- a/numba_cuda/numba/cuda/cudadrv/linkable_code.py
+++ b/numba_cuda/numba/cuda/cudadrv/linkable_code.py
@@ -9,8 +9,9 @@ class LinkableCode:
                  linking errors that may be produced.
     """
 
-    def __init__(self, data, name=None):
+    def __init__(self, data, name=None, nrt=False):
         self.data = data
+        self.nrt = nrt
         self._name = name
 
     @property
diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py
index 3c3c76662..6d6a9e2ca 100644
--- a/numba_cuda/numba/cuda/dispatcher.py
+++ b/numba_cuda/numba/cuda/dispatcher.py
@@ -21,7 +21,7 @@
     ExternFunction,
 )
 from numba.cuda.cudadrv import driver
-from numba.cuda.cudadrv.linkable_code import LinkableCode
+from numba.cuda.cudadrv.linkable_code import LinkableCode, CUSource, PTXSource
 from numba.cuda.cudadrv.devices import get_context
 from numba.cuda.descriptor import cuda_target
 from numba.cuda.errors import (
@@ -270,7 +270,11 @@ def maybe_link_nrt(self, link, tgt_ctx, asm):
         if not link_nrt:
             for file in link:
                 if isinstance(file, LinkableCode):
-                    asm = file.data.decode("utf-8")
+                    if file.nrt:
+                        link_nrt = True
+                        break
+                    elif isinstance(file, (CUSource, PTXSource)):
+                        asm = file.data.decode("utf-8")
                 else:
                     if file.endswith("ptx") or file.endswith("cu"):
                         asm = cached_file_read(file)
diff --git a/numba_cuda/numba/cuda/tests/nrt/test_nrt.py b/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
index 737db5a0f..bd5162041 100644
--- a/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
+++ b/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
@@ -15,6 +15,39 @@
 from numba.cuda.runtime.nrt import rtsys, get_include
 import tempfile
 from numba.core.typing.templates import AbstractTemplate
+from numba.cuda.cudadrv.linkable_code import (
+    CUSource,
+    PTXSource,
+    Fatbin,
+    Cubin,
+    Archive,
+    Object,
+)
+
+
+TEST_BIN_DIR = os.getenv("NUMBA_CUDA_TEST_BIN_DIR")
+
+if TEST_BIN_DIR:
+
+    def make_linkable_code(name, kind, mode):
+        path = os.path.join(TEST_BIN_DIR, name)
+        with open(path, mode) as f:
+            contents = f.read()
+        return kind(contents, nrt=True)
+
+    nrt_extern_a = make_linkable_code("nrt_extern.a", Archive, "rb")
+    nrt_extern_cubin = make_linkable_code("nrt_extern.cubin", Cubin, "rb")
+    nrt_extern_cu = make_linkable_code(
+        "nrt_extern.cu",
+        CUSource,
+        "rb",
+    )
+    nrt_extern_fatbin = make_linkable_code("nrt_extern.fatbin", Fatbin, "rb")
+    nrt_extern_fatbin_multi = make_linkable_code(
+        "nrt_extern_multi.fatbin", Fatbin, "rb"
+    )
+    nrt_extern_o = make_linkable_code("nrt_extern.o", Object, "rb")
+    nrt_extern_ptx = make_linkable_code("nrt_extern.ptx", PTXSource, "rb")
 
 
 def allocate_shim():
@@ -139,6 +172,54 @@ def kernel():
             kernel[1, 1]()
 
 
+class TestNrtLinking(CUDATestCase):
+    def run(self, result=None):
+        with override_config("CUDA_ENABLE_NRT", True):
+            super(TestNrtLinking, self).run(result)
+
+    def test_nrt_detect_linked_ptx_file(self):
+        src = f"#include <{get_include()}/nrt.cuh>"
+        src += """;
+                 extern "C" __device__ int extern_func(int* nb_retval){
+                     NRT_Allocate(1);
+                     return 0;
+                 }
+        """
+        cc = get_current_device().compute_capability
+        ptx, _ = compile(src, "external_nrt.cu", cc)
+        with tempfile.NamedTemporaryFile(
+            delete=True, mode="w+", suffix=".ptx", encoding="utf-8"
+        ) as f:
+            f.write(ptx)
+            f.flush()
+
+            @cuda.jit(link=[f.name])
+            def kernel():
+                allocate_shim()
+
+            kernel[1, 1]()
+
+    @unittest.skipIf(not TEST_BIN_DIR, "necessary binaries not generated.")
+    def test_nrt_detect_linkable_code(self):
+        codes = (
+            nrt_extern_a,
+            nrt_extern_cubin,
+            nrt_extern_cu,
+            nrt_extern_fatbin,
+            nrt_extern_fatbin_multi,
+            nrt_extern_o,
+            nrt_extern_ptx,
+        )
+        for code in codes:
+            with self.subTest(code=code):
+
+                @cuda.jit(link=[code])
+                def kernel():
+                    allocate_shim()
+
+                kernel[1, 1]()
+
+
 class TestNrtStatistics(CUDATestCase):
     def setUp(self):
         self._stream = cuda.default_stream()
diff --git a/numba_cuda/numba/cuda/tests/test_binary_generation/Makefile b/numba_cuda/numba/cuda/tests/test_binary_generation/Makefile
index 145f8b024..6d5ad76f8 100644
--- a/numba_cuda/numba/cuda/tests/test_binary_generation/Makefile
+++ b/numba_cuda/numba/cuda/tests/test_binary_generation/Makefile
@@ -40,6 +40,8 @@ LTOIR_FLAGS := $(LTOIR_GENCODE) -dc
 
 OUTPUT_DIR := ./
 
+NRT_INCLUDE_DIR := $(shell python -c "from numba.cuda.runtime.nrt import get_include; print(get_include())")
+
 all:
 	@echo "GPU CC: $(GPU_CC)"
 	@echo "Alternative CC: $(ALT_CC)"
@@ -52,6 +54,13 @@ all:
 	nvcc $(NVCC_FLAGS) $(OBJECT_FLAGS) -o $(OUTPUT_DIR)/test_device_functions.o test_device_functions.cu
 	nvcc $(NVCC_FLAGS) $(LIBRARY_FLAGS) -o $(OUTPUT_DIR)/test_device_functions.a test_device_functions.cu
 
+	nvcc $(NVCC_FLAGS) $(CUBIN_FLAGS) -o $(OUTPUT_DIR)/nrt_extern.cubin nrt_extern.cu -I$(NRT_INCLUDE_DIR)
+	nvcc $(NVCC_FLAGS) $(FATBIN_FLAGS) -o $(OUTPUT_DIR)/nrt_extern.fatbin nrt_extern.cu -I$(NRT_INCLUDE_DIR)
+	nvcc $(NVCC_FLAGS) $(MULTI_FATBIN_FLAGS) -o $(OUTPUT_DIR)/nrt_extern_multi.fatbin nrt_extern.cu -I$(NRT_INCLUDE_DIR)
+	nvcc $(NVCC_FLAGS) $(PTX_FLAGS) -o $(OUTPUT_DIR)/nrt_extern.ptx nrt_extern.cu -I$(NRT_INCLUDE_DIR)
+	nvcc $(NVCC_FLAGS) $(OBJECT_FLAGS) -o $(OUTPUT_DIR)/nrt_extern.o nrt_extern.cu -I$(NRT_INCLUDE_DIR)
+	nvcc $(NVCC_FLAGS) $(LIBRARY_FLAGS) -o $(OUTPUT_DIR)/nrt_extern.a nrt_extern.cu -I$(NRT_INCLUDE_DIR)
+
 	# Generate LTO-IR wrapped in a fatbin
 	nvcc $(NVCC_FLAGS) $(LTOIR_FLAGS) -o $(OUTPUT_DIR)/test_device_functions.ltoir.o test_device_functions.cu
 	# Generate LTO-IR in a "raw" LTO-IR container
diff --git a/numba_cuda/numba/cuda/tests/test_binary_generation/nrt_extern.cu b/numba_cuda/numba/cuda/tests/test_binary_generation/nrt_extern.cu
new file mode 100644
index 000000000..c6ffa5704
--- /dev/null
+++ b/numba_cuda/numba/cuda/tests/test_binary_generation/nrt_extern.cu
@@ -0,0 +1,6 @@
+#include <nrt.cuh>
+
+extern "C" __device__ int extern_func(int* nb_retval){
+    NRT_Allocate(1);
+    return 0;
+}

From 03cfd924e61ba10809be743bc6866c6f4358a31e Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Mon, 14 Apr 2025 08:03:02 -0700
Subject: [PATCH 18/33] reorder steps in pynvjitlink wheel test job

---
 ci/test_wheel_pynvjitlink.sh | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/ci/test_wheel_pynvjitlink.sh b/ci/test_wheel_pynvjitlink.sh
index 55110234c..746d8a3a6 100755
--- a/ci/test_wheel_pynvjitlink.sh
+++ b/ci/test_wheel_pynvjitlink.sh
@@ -15,6 +15,12 @@ python -m pip install \
 rapids-logger "Install pynvjitlink"
 python -m pip install pynvjitlink-cu12
 
+
+rapids-logger "Install wheel"
+package=$(realpath wheel/numba_cuda*.whl)
+echo "Package path: $package"
+python -m pip install $package
+
 rapids-logger "Build tests"
 PY_SCRIPT="
 import numba_cuda
@@ -28,10 +34,6 @@ pushd $NUMBA_CUDA_TEST_BIN_DIR
 make
 popd
 
-rapids-logger "Install wheel"
-package=$(realpath wheel/numba_cuda*.whl)
-echo "Package path: $package"
-python -m pip install $package
 
 rapids-logger "Check GPU usage"
 nvidia-smi

From 613db4a84c3f01620ce914e075f2c4c6f4695922 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <53796099+brandon-b-miller@users.noreply.github.com>
Date: Mon, 14 Apr 2025 15:11:00 -0500
Subject: [PATCH 19/33] Update numba_cuda/numba/cuda/dispatcher.py

Co-authored-by: Michael Wang <13521008+isVoid@users.noreply.github.com>
---
 numba_cuda/numba/cuda/dispatcher.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py
index fcaad9e3b..fd7cc855d 100644
--- a/numba_cuda/numba/cuda/dispatcher.py
+++ b/numba_cuda/numba/cuda/dispatcher.py
@@ -276,6 +276,8 @@ def maybe_link_nrt(self, link, tgt_ctx, asm):
                         break
                     elif isinstance(file, (CUSource, PTXSource)):
                         asm = file.data.decode("utf-8")
+                    else:
+                        asm = ""
                 else:
                     if file.endswith("ptx") or file.endswith("cu"):
                         asm = cached_file_read(file)

From b05fbe8d9913d044aa9221245c9ff5cb6759a173 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Tue, 15 Apr 2025 07:52:15 -0700
Subject: [PATCH 20/33] dont leak memory, cleanup

---
 numba_cuda/numba/cuda/tests/nrt/test_nrt.py   | 52 +++++++------------
 .../test_binary_generation/nrt_extern.cu      |  3 +-
 2 files changed, 20 insertions(+), 35 deletions(-)

diff --git a/numba_cuda/numba/cuda/tests/nrt/test_nrt.py b/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
index c306ca792..b14b28eea 100644
--- a/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
+++ b/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
@@ -50,11 +50,14 @@ def make_linkable_code(name, kind, mode):
     nrt_extern_ptx = make_linkable_code("nrt_extern.ptx", PTXSource, "rb")
 
 
-def allocate_shim():
+def allocate_deallocate_handle():
+    """
+    Handle to call NRT_Allocate and NRT_Free
+    """
     pass
 
 
-@cuda_decl_registry.register_global(allocate_shim)
+@cuda_decl_registry.register_global(allocate_deallocate_handle)
 class AllocateShimImpl(AbstractTemplate):
     def generic(self, args, kws):
         return signature(types.void)
@@ -62,21 +65,23 @@ def generic(self, args, kws):
 
 # let there be a __device__ function expecting
 # a string_view* and returning size_type
-_allocate_shim = cuda.declare_device("extern_func", types.int32())
+device_fun_shim = cuda.declare_device(
+    "device_allocate_deallocate", types.int32()
+)
 
 
 # wrapper to turn the above into a python callable
-def call_allocate_shim():
-    return _allocate_shim()
+def call_device_fun_shim():
+    return device_fun_shim()
 
 
-@cuda_lower(allocate_shim)
-def allocate_shim_impl(context, builder, sig, args):
+@cuda_lower(allocate_deallocate_handle)
+def allocate_deallocate_impl(context, builder, sig, args):
     sig_ = types.int32()
     # call the external function, passing the pointer
     result = context.compile_internal(
         builder,
-        call_allocate_shim,
+        call_device_fun_shim,
         sig_,
         (),
     )
@@ -149,28 +154,6 @@ def g(out_ary):
 
         self.assertEqual(out_ary[0], 1)
 
-    def test_nrt_detect_linked_ptx_file(self):
-        src = f"#include <{get_include()}/nrt.cuh>"
-        src += """;
-                 extern "C" __device__ int extern_func(int* nb_retval){
-                     NRT_Allocate(1);
-                     return 0;
-                 }
-        """
-        cc = get_current_device().compute_capability
-        ptx, _ = compile(src, "external_nrt.cu", cc)
-        with tempfile.NamedTemporaryFile(
-            delete=True, mode="w+", suffix=".ptx", encoding="utf-8"
-        ) as f:
-            f.write(ptx)
-            f.flush()
-
-            @cuda.jit(link=[f.name])
-            def kernel():
-                allocate_shim()
-
-            kernel[1, 1]()
-
 
 class TestNrtLinking(CUDATestCase):
     def run(self, result=None):
@@ -180,8 +163,9 @@ def run(self, result=None):
     def test_nrt_detect_linked_ptx_file(self):
         src = f"#include <{get_include()}/nrt.cuh>"
         src += """;
-                 extern "C" __device__ int extern_func(int* nb_retval){
-                     NRT_Allocate(1);
+                 extern "C" __device__ int device_allocate_deallocate(int* nb_retval){
+                     auto ptr = NRT_Allocate(1);
+                     NRT_Free(ptr);
                      return 0;
                  }
         """
@@ -195,7 +179,7 @@ def test_nrt_detect_linked_ptx_file(self):
 
             @cuda.jit(link=[f.name])
             def kernel():
-                allocate_shim()
+                allocate_deallocate_handle()
 
             kernel[1, 1]()
 
@@ -215,7 +199,7 @@ def test_nrt_detect_linkable_code(self):
 
                 @cuda.jit(link=[code])
                 def kernel():
-                    allocate_shim()
+                    allocate_deallocate_handle()
 
                 kernel[1, 1]()
 
diff --git a/numba_cuda/numba/cuda/tests/test_binary_generation/nrt_extern.cu b/numba_cuda/numba/cuda/tests/test_binary_generation/nrt_extern.cu
index c6ffa5704..da0ad8d8c 100644
--- a/numba_cuda/numba/cuda/tests/test_binary_generation/nrt_extern.cu
+++ b/numba_cuda/numba/cuda/tests/test_binary_generation/nrt_extern.cu
@@ -1,6 +1,7 @@
 #include <nrt.cuh>
 
 extern "C" __device__ int extern_func(int* nb_retval){
-    NRT_Allocate(1);
+    auto ptr = NRT_Allocate(1);
+    NRT_Free(ptr);
     return 0;
 }

From 5134e6c389e57edd687daf5b9e9a7e8d6fd9fe56 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Tue, 15 Apr 2025 08:16:25 -0700
Subject: [PATCH 21/33] cache nrt source

---
 numba_cuda/numba/cuda/dispatcher.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py
index fd7cc855d..440f78599 100644
--- a/numba_cuda/numba/cuda/dispatcher.py
+++ b/numba_cuda/numba/cuda/dispatcher.py
@@ -291,7 +291,8 @@ def maybe_link_nrt(self, link, tgt_ctx, asm):
         if link_nrt:
             basedir = os.path.dirname(os.path.abspath(__file__))
             nrt_path = os.path.join(basedir, "runtime", "nrt.cu")
-            link.append(nrt_path)
+            nrt_src = cached_file_read(nrt_path)
+            link.append(CUSource(nrt_src, name="nrt.cu", nrt=True))
 
     @property
     def library(self):

From a2aa31e0198a2de80d726d2f99f170dd799fe033 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Mon, 21 Apr 2025 14:10:50 -0700
Subject: [PATCH 22/33] small fix

---
 .../numba/cuda/tests/test_binary_generation/nrt_extern.cu       | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/numba_cuda/numba/cuda/tests/test_binary_generation/nrt_extern.cu b/numba_cuda/numba/cuda/tests/test_binary_generation/nrt_extern.cu
index da0ad8d8c..c9f73908e 100644
--- a/numba_cuda/numba/cuda/tests/test_binary_generation/nrt_extern.cu
+++ b/numba_cuda/numba/cuda/tests/test_binary_generation/nrt_extern.cu
@@ -1,6 +1,6 @@
 #include <nrt.cuh>
 
-extern "C" __device__ int extern_func(int* nb_retval){
+extern "C" __device__ int device_allocate_deallocate(int* nb_retval){
     auto ptr = NRT_Allocate(1);
     NRT_Free(ptr);
     return 0;

From 2edbe6a5c2a591e5b1f63b1242d09c3825a93e49 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Mon, 21 Apr 2025 14:39:31 -0700
Subject: [PATCH 23/33] reorder test build steps

---
 ci/test_wheel_deps_wheels.sh | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/ci/test_wheel_deps_wheels.sh b/ci/test_wheel_deps_wheels.sh
index c8c135b49..ea7dcffa6 100755
--- a/ci/test_wheel_deps_wheels.sh
+++ b/ci/test_wheel_deps_wheels.sh
@@ -25,16 +25,18 @@ test_dir = root + \"numba/cuda/tests/test_binary_generation/\"
 print(test_dir)
 "
 
-NUMBA_CUDA_TEST_BIN_DIR=$(python -c "$PY_SCRIPT")
-pushd $NUMBA_CUDA_TEST_BIN_DIR
-make
-popd
 
 rapids-logger "Install wheel"
 package=$(realpath wheel/numba_cuda*.whl)
 echo "Package path: $package"
 python -m pip install $package
 
+NUMBA_CUDA_TEST_BIN_DIR=$(python -c "$PY_SCRIPT")
+pushd $NUMBA_CUDA_TEST_BIN_DIR
+make
+popd
+
+
 rapids-logger "Check GPU usage"
 nvidia-smi
 

From 34b710d8c77fc828386b889cb4da07d54bd75b1e Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Wed, 30 Apr 2025 07:08:07 -0700
Subject: [PATCH 24/33] partially address reviews

---
 ci/test_wheel_deps_wheels.sh                        | 13 ++++++-------
 numba_cuda/numba/cuda/cudadrv/linkable_code.py      |  2 ++
 numba_cuda/numba/cuda/tests/nrt/test_nrt.py         |  2 --
 .../cuda/tests/test_binary_generation/Makefile      |  2 ++
 .../test_binary_generation/generate_raw_ltoir.py    |  9 +++++++--
 5 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/ci/test_wheel_deps_wheels.sh b/ci/test_wheel_deps_wheels.sh
index ea7dcffa6..9cbce45af 100755
--- a/ci/test_wheel_deps_wheels.sh
+++ b/ci/test_wheel_deps_wheels.sh
@@ -17,6 +17,12 @@ python -m pip install \
     pytest
 
 
+rapids-logger "Install wheel"
+package=$(realpath wheel/numba_cuda*.whl)
+echo "Package path: $package"
+python -m pip install $package
+
+
 rapids-logger "Build tests"
 PY_SCRIPT="
 import numba_cuda
@@ -24,13 +30,6 @@ root = numba_cuda.__file__.rstrip('__init__.py')
 test_dir = root + \"numba/cuda/tests/test_binary_generation/\"
 print(test_dir)
 "
-
-
-rapids-logger "Install wheel"
-package=$(realpath wheel/numba_cuda*.whl)
-echo "Package path: $package"
-python -m pip install $package
-
 NUMBA_CUDA_TEST_BIN_DIR=$(python -c "$PY_SCRIPT")
 pushd $NUMBA_CUDA_TEST_BIN_DIR
 make
diff --git a/numba_cuda/numba/cuda/cudadrv/linkable_code.py b/numba_cuda/numba/cuda/cudadrv/linkable_code.py
index ad77c0c92..a183ec14d 100644
--- a/numba_cuda/numba/cuda/cudadrv/linkable_code.py
+++ b/numba_cuda/numba/cuda/cudadrv/linkable_code.py
@@ -13,6 +13,8 @@ class LinkableCode:
     :param teardown_callback: A function called just prior to the unloading of
                               a module that has this code object linked into
                               it.
+    :param nrt: If True, assume this object contains NRT function calls and
+                add NRT source code to the final link.
     """
 
     def __init__(
diff --git a/numba_cuda/numba/cuda/tests/nrt/test_nrt.py b/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
index 35089e838..09db65c9a 100644
--- a/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
+++ b/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
@@ -63,8 +63,6 @@ def generic(self, args, kws):
         return signature(types.void)
 
 
-# let there be a __device__ function expecting
-# a string_view* and returning size_type
 device_fun_shim = cuda.declare_device(
     "device_allocate_deallocate", types.int32()
 )
diff --git a/numba_cuda/numba/cuda/tests/test_binary_generation/Makefile b/numba_cuda/numba/cuda/tests/test_binary_generation/Makefile
index 6d5ad76f8..0d6b5f589 100644
--- a/numba_cuda/numba/cuda/tests/test_binary_generation/Makefile
+++ b/numba_cuda/numba/cuda/tests/test_binary_generation/Makefile
@@ -63,5 +63,7 @@ all:
 
 	# Generate LTO-IR wrapped in a fatbin
 	nvcc $(NVCC_FLAGS) $(LTOIR_FLAGS) -o $(OUTPUT_DIR)/test_device_functions.ltoir.o test_device_functions.cu
+	nvcc $(NVCC_FLAGS) $(LTOIR_FLAGS) -o $(OUTPUT_DIR)/nrt_extern.ltoir.o nrt_extern.cu -I$(NRT_INCLUDE_DIR)
 	# Generate LTO-IR in a "raw" LTO-IR container
 	python generate_raw_ltoir.py --arch sm_$(GPU_CC) -o $(OUTPUT_DIR)/test_device_functions.ltoir test_device_functions.cu
+	python generate_raw_ltoir.py --arch sm_$(GPU_CC) -o $(OUTPUT_DIR)/nrt_extern.ltoir nrt_extern.cu --nrt=True
diff --git a/numba_cuda/numba/cuda/tests/test_binary_generation/generate_raw_ltoir.py b/numba_cuda/numba/cuda/tests/test_binary_generation/generate_raw_ltoir.py
index b4d32a34c..41d394050 100644
--- a/numba_cuda/numba/cuda/tests/test_binary_generation/generate_raw_ltoir.py
+++ b/numba_cuda/numba/cuda/tests/test_binary_generation/generate_raw_ltoir.py
@@ -7,6 +7,7 @@
 import sys
 
 from cuda import nvrtc
+from numba.cuda.runtime.nrt import get_include
 
 # Magic number found at the start of an LTO-IR file
 LTOIR_MAGIC = 0x7F4E43ED
@@ -88,7 +89,9 @@ def get_ltoir(source, name, arch):
         nvrtc.nvrtcCreateProgram(source.encode(), name.encode(), 0, [], [])
     )
 
-    cuda_include_flags = determine_include_flags()
+    cuda_include_flags = determine_include_flags() + (
+        [f"-I{get_include()}"] if args.nrt else []
+    )
     if cuda_include_flags is None:
         print("Error determining CUDA include flags. Exiting.", file=sys.stderr)
         sys.exit(1)
@@ -160,7 +163,9 @@ def main(sourcepath, outputpath, arch):
         help="compute arch to target (e.g. sm_87). Defaults to sm_50.",
         default="sm_50",
     )
-
+    parser.add_argument(
+        "--nrt", type=lambda x: x.lower() == "true", default=False
+    )
     args = parser.parse_args()
     outputpath = args.output
 

From bd3da220a53f08b7c91eaeed899e317281de3ece Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Wed, 30 Apr 2025 07:27:39 -0700
Subject: [PATCH 25/33] go back to regex check

---
 numba_cuda/numba/cuda/dispatcher.py         | 21 ++++++++++++++++++++-
 numba_cuda/numba/cuda/tests/nrt/test_nrt.py |  2 +-
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py
index 440f78599..fc1711ec1 100644
--- a/numba_cuda/numba/cuda/dispatcher.py
+++ b/numba_cuda/numba/cuda/dispatcher.py
@@ -20,6 +20,7 @@
     kernel_fixup,
     ExternFunction,
 )
+import re
 from numba.cuda.cudadrv import driver
 from numba.cuda.cudadrv.linkable_code import LinkableCode, CUSource, PTXSource
 from numba.cuda.cudadrv.devices import get_context
@@ -263,10 +264,28 @@ def link_to_library_functions(
         self.reload_init = []
 
     def maybe_link_nrt(self, link, tgt_ctx, asm):
+        """
+        Add the NRT source code to the link if the neccesary conditions are met:
+        If NRT is enabled for the CUDATargetContext, return True if we either
+        detect .extern .func decls of NRT functions in any file in the link or
+        any of the passed LinkableCode objects have NRT enabled. In the special
+        case of PTXSource or CUSource objects, the source code will be inspected
+        even if the NRT flag is not set.
+        """
         if not tgt_ctx.enable_nrt:
             return
 
-        nrt_in_asm = lambda asm: any(fn in asm for fn in self.NRT_functions)
+        def nrt_in_asm(asm):
+            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)
+            return len(nrt_in_asm) > 0
+
         link_nrt = nrt_in_asm(asm)
         if not link_nrt:
             for file in link:
diff --git a/numba_cuda/numba/cuda/tests/nrt/test_nrt.py b/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
index 09db65c9a..104fe0d60 100644
--- a/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
+++ b/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
@@ -160,7 +160,7 @@ def run(self, result=None):
 
     def test_nrt_detect_linked_ptx_file(self):
         src = f"#include <{get_include()}/nrt.cuh>"
-        src += """;
+        src += """
                  extern "C" __device__ int device_allocate_deallocate(int* nb_retval){
                      auto ptr = NRT_Allocate(1);
                      NRT_Free(ptr);

From dc9680bf9bf2235baca693c6bfd91c0a58c5437a Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Wed, 30 Apr 2025 14:48:55 -0700
Subject: [PATCH 26/33] small fix

---
 numba_cuda/numba/cuda/dispatcher.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py
index fc1711ec1..9351680d2 100644
--- a/numba_cuda/numba/cuda/dispatcher.py
+++ b/numba_cuda/numba/cuda/dispatcher.py
@@ -294,7 +294,7 @@ def nrt_in_asm(asm):
                         link_nrt = True
                         break
                     elif isinstance(file, (CUSource, PTXSource)):
-                        asm = file.data.decode("utf-8")
+                        asm = file.data
                     else:
                         asm = ""
                 else:

From 71c280b50edbb2b50505a7c9e2824baae692405c Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Fri, 2 May 2025 06:26:14 -0700
Subject: [PATCH 27/33] preserve nrt library through serialization

---
 numba_cuda/numba/cuda/codegen.py     | 17 +++++++++++++++--
 numba_cuda/numba/cuda/dispatcher.py  |  7 ++-----
 numba_cuda/numba/cuda/runtime/nrt.py |  9 ++++++++-
 3 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/numba_cuda/numba/cuda/codegen.py b/numba_cuda/numba/cuda/codegen.py
index 768d87c77..04b835c5e 100644
--- a/numba_cuda/numba/cuda/codegen.py
+++ b/numba_cuda/numba/cuda/codegen.py
@@ -5,6 +5,7 @@
 from .cudadrv import devices, driver, nvvm, runtime
 from numba.cuda.cudadrv.libs import get_cudalib
 from numba.cuda.cudadrv.linkable_code import LinkableCode
+from numba.cuda.runtime.nrt import NRT_LIBRARY
 
 import os
 import subprocess
@@ -362,9 +363,17 @@ def _reduce_states(self):
         but loaded functions are discarded. They are recreated when needed
         after deserialization.
         """
+        nrt = False
         if self._linking_files:
-            msg = "Cannot pickle CUDACodeLibrary with linking files"
-            raise RuntimeError(msg)
+            if (
+                len(self._linking_files) != 1
+                and self._linking_files[0] != NRT_LIBRARY
+            ):
+                msg = "Cannot pickle CUDACodeLibrary with linking files"
+                raise RuntimeError(msg)
+            else:
+                nrt = True
+
         if not self._finalized:
             raise RuntimeError("Cannot pickle unfinalized CUDACodeLibrary")
         return dict(
@@ -378,6 +387,7 @@ def _reduce_states(self):
             max_registers=self._max_registers,
             nvvm_options=self._nvvm_options,
             needs_cudadevrt=self.needs_cudadevrt,
+            nrt=nrt,
         )
 
     @classmethod
@@ -393,6 +403,7 @@ def _rebuild(
         max_registers,
         nvvm_options,
         needs_cudadevrt,
+        nrt,
     ):
         """
         Rebuild an instance.
@@ -409,6 +420,8 @@ def _rebuild(
         instance.needs_cudadevrt = needs_cudadevrt
 
         instance._finalized = True
+        if nrt:
+            instance._linking_files = {NRT_LIBRARY}
 
         return instance
 
diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py
index 9351680d2..619a67e13 100644
--- a/numba_cuda/numba/cuda/dispatcher.py
+++ b/numba_cuda/numba/cuda/dispatcher.py
@@ -30,7 +30,7 @@
     normalize_kernel_dimensions,
 )
 from numba.cuda import types as cuda_types
-from numba.cuda.runtime.nrt import rtsys
+from numba.cuda.runtime.nrt import rtsys, NRT_LIBRARY
 from numba.cuda.utils import cached_file_read
 from numba.cuda.locks import module_init_lock
 
@@ -308,10 +308,7 @@ def nrt_in_asm(asm):
                     break
 
         if link_nrt:
-            basedir = os.path.dirname(os.path.abspath(__file__))
-            nrt_path = os.path.join(basedir, "runtime", "nrt.cu")
-            nrt_src = cached_file_read(nrt_path)
-            link.append(CUSource(nrt_src, name="nrt.cu", nrt=True))
+            link.append(NRT_LIBRARY)
 
     @property
     def library(self):
diff --git a/numba_cuda/numba/cuda/runtime/nrt.py b/numba_cuda/numba/cuda/runtime/nrt.py
index 514f6dfdb..c978ce023 100644
--- a/numba_cuda/numba/cuda/runtime/nrt.py
+++ b/numba_cuda/numba/cuda/runtime/nrt.py
@@ -13,7 +13,8 @@
 )
 from numba.cuda.cudadrv import devices
 from numba.cuda.api import get_current_device
-from numba.cuda.utils import _readenv
+from numba.cuda.utils import _readenv, cached_file_read
+from numba.cuda.cudadrv.linkable_code import CUSource
 
 
 # Check environment variable or config for NRT statistics enablement
@@ -345,3 +346,9 @@ def print_memsys(self, stream=None):
 
 # Create an instance of the runtime
 rtsys = _Runtime()
+
+
+basedir = os.path.dirname(os.path.abspath(__file__))
+nrt_path = os.path.join(basedir, "nrt.cu")
+nrt_src = cached_file_read(nrt_path)
+NRT_LIBRARY = CUSource(nrt_src, name="nrt.cu", nrt=True)

From eaebb9ae274101b8389bf5160d489cfd82095423 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Fri, 2 May 2025 08:35:37 -0700
Subject: [PATCH 28/33] pass full test suite with nrt and nvidia binding

---
 numba_cuda/numba/cuda/codegen.py    | 8 ++++----
 numba_cuda/numba/cuda/dispatcher.py | 4 +++-
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/numba_cuda/numba/cuda/codegen.py b/numba_cuda/numba/cuda/codegen.py
index 04b835c5e..4030c6452 100644
--- a/numba_cuda/numba/cuda/codegen.py
+++ b/numba_cuda/numba/cuda/codegen.py
@@ -366,13 +366,13 @@ def _reduce_states(self):
         nrt = False
         if self._linking_files:
             if (
-                len(self._linking_files) != 1
-                and self._linking_files[0] != NRT_LIBRARY
+                len(self._linking_files) == 1
+                and NRT_LIBRARY in self._linking_files
             ):
+                nrt = True
+            else:
                 msg = "Cannot pickle CUDACodeLibrary with linking files"
                 raise RuntimeError(msg)
-            else:
-                nrt = True
 
         if not self._finalized:
             raise RuntimeError("Cannot pickle unfinalized CUDACodeLibrary")
diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py
index 619a67e13..2db7880d6 100644
--- a/numba_cuda/numba/cuda/dispatcher.py
+++ b/numba_cuda/numba/cuda/dispatcher.py
@@ -298,7 +298,9 @@ def nrt_in_asm(asm):
                     else:
                         asm = ""
                 else:
-                    if file.endswith("ptx") or file.endswith("cu"):
+                    if isinstance(file, bytes):
+                        continue
+                    elif file.endswith("ptx") or file.endswith("cu"):
                         asm = cached_file_read(file)
                     else:
                         # Not a PTX or CUDA source, skip it

From ff12e16cc0ede0d35e5f71b3afbcd2fc3db52ee4 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Tue, 6 May 2025 07:16:28 -0700
Subject: [PATCH 29/33] drop autodetection of NRT in user supplied raw file
 paths

---
 numba_cuda/numba/cuda/dispatcher.py         | 18 +-----------------
 numba_cuda/numba/cuda/tests/nrt/test_nrt.py | 14 ++++----------
 2 files changed, 5 insertions(+), 27 deletions(-)

diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py
index 2db7880d6..5df32b0a4 100644
--- a/numba_cuda/numba/cuda/dispatcher.py
+++ b/numba_cuda/numba/cuda/dispatcher.py
@@ -22,7 +22,7 @@
 )
 import re
 from numba.cuda.cudadrv import driver
-from numba.cuda.cudadrv.linkable_code import LinkableCode, CUSource, PTXSource
+from numba.cuda.cudadrv.linkable_code import LinkableCode
 from numba.cuda.cudadrv.devices import get_context
 from numba.cuda.descriptor import cuda_target
 from numba.cuda.errors import (
@@ -31,7 +31,6 @@
 )
 from numba.cuda import types as cuda_types
 from numba.cuda.runtime.nrt import rtsys, NRT_LIBRARY
-from numba.cuda.utils import cached_file_read
 from numba.cuda.locks import module_init_lock
 
 from numba import cuda
@@ -293,21 +292,6 @@ def nrt_in_asm(asm):
                     if file.nrt:
                         link_nrt = True
                         break
-                    elif isinstance(file, (CUSource, PTXSource)):
-                        asm = file.data
-                    else:
-                        asm = ""
-                else:
-                    if isinstance(file, bytes):
-                        continue
-                    elif file.endswith("ptx") or file.endswith("cu"):
-                        asm = cached_file_read(file)
-                    else:
-                        # Not a PTX or CUDA source, skip it
-                        continue
-                if nrt_in_asm(asm):
-                    link_nrt = True
-                    break
 
         if link_nrt:
             link.append(NRT_LIBRARY)
diff --git a/numba_cuda/numba/cuda/tests/nrt/test_nrt.py b/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
index 104fe0d60..9bd01b37d 100644
--- a/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
+++ b/numba_cuda/numba/cuda/tests/nrt/test_nrt.py
@@ -13,7 +13,6 @@
 from numba.cuda.cudaimpl import lower as cuda_lower
 from numba import cuda
 from numba.cuda.runtime.nrt import rtsys, get_include
-import tempfile
 from numba.core.typing.templates import AbstractTemplate
 from numba.cuda.cudadrv.linkable_code import (
     CUSource,
@@ -169,17 +168,12 @@ def test_nrt_detect_linked_ptx_file(self):
         """
         cc = get_current_device().compute_capability
         ptx, _ = compile(src, "external_nrt.cu", cc)
-        with tempfile.NamedTemporaryFile(
-            delete=True, mode="w+", suffix=".ptx", encoding="utf-8"
-        ) as f:
-            f.write(ptx)
-            f.flush()
 
-            @cuda.jit(link=[f.name])
-            def kernel():
-                allocate_deallocate_handle()
+        @cuda.jit(link=[PTXSource(ptx.encode(), nrt=True)])
+        def kernel():
+            allocate_deallocate_handle()
 
-            kernel[1, 1]()
+        kernel[1, 1]()
 
     @unittest.skipIf(not TEST_BIN_DIR, "necessary binaries not generated.")
     def test_nrt_detect_linkable_code(self):

From d893d2bd77956716daf58bdb939982c063659a43 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <53796099+brandon-b-miller@users.noreply.github.com>
Date: Tue, 6 May 2025 09:22:30 -0500
Subject: [PATCH 30/33] Apply suggestions from code review

Co-authored-by: Graham Markall <535640+gmarkall@users.noreply.github.com>
---
 numba_cuda/numba/cuda/tests/test_binary_generation/Makefile   | 2 +-
 .../cuda/tests/test_binary_generation/generate_raw_ltoir.py   | 4 +---
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/numba_cuda/numba/cuda/tests/test_binary_generation/Makefile b/numba_cuda/numba/cuda/tests/test_binary_generation/Makefile
index 0d6b5f589..7d36e5cba 100644
--- a/numba_cuda/numba/cuda/tests/test_binary_generation/Makefile
+++ b/numba_cuda/numba/cuda/tests/test_binary_generation/Makefile
@@ -66,4 +66,4 @@ all:
 	nvcc $(NVCC_FLAGS) $(LTOIR_FLAGS) -o $(OUTPUT_DIR)/nrt_extern.ltoir.o nrt_extern.cu -I$(NRT_INCLUDE_DIR)
 	# Generate LTO-IR in a "raw" LTO-IR container
 	python generate_raw_ltoir.py --arch sm_$(GPU_CC) -o $(OUTPUT_DIR)/test_device_functions.ltoir test_device_functions.cu
-	python generate_raw_ltoir.py --arch sm_$(GPU_CC) -o $(OUTPUT_DIR)/nrt_extern.ltoir nrt_extern.cu --nrt=True
+	python generate_raw_ltoir.py --arch sm_$(GPU_CC) -o $(OUTPUT_DIR)/nrt_extern.ltoir nrt_extern.cu --nrt
diff --git a/numba_cuda/numba/cuda/tests/test_binary_generation/generate_raw_ltoir.py b/numba_cuda/numba/cuda/tests/test_binary_generation/generate_raw_ltoir.py
index 41d394050..d5e50e5c8 100644
--- a/numba_cuda/numba/cuda/tests/test_binary_generation/generate_raw_ltoir.py
+++ b/numba_cuda/numba/cuda/tests/test_binary_generation/generate_raw_ltoir.py
@@ -163,9 +163,7 @@ def main(sourcepath, outputpath, arch):
         help="compute arch to target (e.g. sm_87). Defaults to sm_50.",
         default="sm_50",
     )
-    parser.add_argument(
-        "--nrt", type=lambda x: x.lower() == "true", default=False
-    )
+    parser.add_argument("--nrt", action="store_true")
     args = parser.parse_args()
     outputpath = args.output
 

From ad0a0a7bf0860f988e6f94548cdb111244f1e0bc Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Wed, 7 May 2025 10:16:18 -0700
Subject: [PATCH 31/33] simplify and update docstring

---
 numba_cuda/numba/cuda/dispatcher.py | 31 +++++++++++++----------------
 1 file changed, 14 insertions(+), 17 deletions(-)

diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py
index 5df32b0a4..afb47ffc8 100644
--- a/numba_cuda/numba/cuda/dispatcher.py
+++ b/numba_cuda/numba/cuda/dispatcher.py
@@ -264,28 +264,25 @@ def link_to_library_functions(
 
     def maybe_link_nrt(self, link, tgt_ctx, asm):
         """
-        Add the NRT source code to the link if the neccesary conditions are met:
-        If NRT is enabled for the CUDATargetContext, return True if we either
-        detect .extern .func decls of NRT functions in any file in the link or
-        any of the passed LinkableCode objects have NRT enabled. In the special
-        case of PTXSource or CUSource objects, the source code will be inspected
-        even if the NRT flag is not set.
+        Add the NRT source code to the link if the neccesary conditions are met.
+        NRT must be enabled for the CUDATargetContext, and either NRT functions
+        must be detected in the kernel asm or an NRT enabled LinkableCode object
+        must be passed.
         """
+
         if not tgt_ctx.enable_nrt:
             return
 
-        def nrt_in_asm(asm):
-            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)
-            return len(nrt_in_asm) > 0
+        all_nrt = "|".join(self.NRT_functions)
+        pattern = (
+            r"\.extern\s+\.func\s+(?:\s*\(.+\)\s*)?("
+            + all_nrt
+            + r")\s*\([^)]*\)\s*;"
+        )
 
-        link_nrt = nrt_in_asm(asm)
+        nrt_in_asm = re.findall(pattern, asm)
+        if len(nrt_in_asm) > 0:
+            link_nrt = True
         if not link_nrt:
             for file in link:
                 if isinstance(file, LinkableCode):

From 19c6a5601a9b1b2d5d4c753e6cc54a8410756813 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Wed, 7 May 2025 10:36:24 -0700
Subject: [PATCH 32/33] small bugfix

---
 numba_cuda/numba/cuda/dispatcher.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py
index afb47ffc8..6e40ea5eb 100644
--- a/numba_cuda/numba/cuda/dispatcher.py
+++ b/numba_cuda/numba/cuda/dispatcher.py
@@ -279,7 +279,7 @@ def maybe_link_nrt(self, link, tgt_ctx, asm):
             + all_nrt
             + r")\s*\([^)]*\)\s*;"
         )
-
+        link_nrt = False
         nrt_in_asm = re.findall(pattern, asm)
         if len(nrt_in_asm) > 0:
             link_nrt = True

From 50b61acc6f87450ce4360db0598dae6f4f43eb95 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Wed, 7 May 2025 12:00:20 -0700
Subject: [PATCH 33/33] reoder steps in final CI script

---
 ci/test_wheel.sh | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/ci/test_wheel.sh b/ci/test_wheel.sh
index 113810665..9c3e9bfe8 100755
--- a/ci/test_wheel.sh
+++ b/ci/test_wheel.sh
@@ -13,6 +13,11 @@ python -m pip install \
     "cuda-python==${CUDA_VER_MAJOR_MINOR%.*}.*" \
     pytest
 
+rapids-logger "Install wheel"
+package=$(realpath wheel/numba_cuda*.whl)
+echo "Wheel path: $package"
+python -m pip install $package
+
 
 GET_TEST_BINARY_DIR="
 import numba_cuda
@@ -34,11 +39,6 @@ else
 fi
 
 
-rapids-logger "Install wheel"
-package=$(realpath wheel/numba_cuda*.whl)
-echo "Wheel path: $package"
-python -m pip install $package
-
 rapids-logger "Check GPU usage"
 nvidia-smi