numba-cuda 0.10.0__py3-none-any.whl → 0.11.0__py3-none-any.whl

numba_cuda/numba/cuda/tests/cudapy/test_array_alignment.py

@@ -0,0 +1,236 @@
+import re
+import itertools
+import numpy as np
+from numba import cuda
+from numba.core.errors import TypingError
+from numba.cuda.testing import CUDATestCase
+import unittest
+
+
+# Set to true if you want to see dots printed for each subtest.
+NOISY = False
+
+
+# In order to verify the alignment of the local and shared memory arrays, we
+# inspect the LLVM IR of the generated kernel using the following regexes.
+
+# Shared memory example:
+# @"_cudapy_smem_38" = addrspace(3) global [1 x i8] undef, align 16
+SMEM_PATTERN = re.compile(
+    r'^@"_cudapy_smem_\d+".*?align (\d+)',
+    re.MULTILINE,
+)
+
+# Local memory example:
+# %"_cudapy_lmem" = alloca [1 x i8], align 64
+LMEM_PATTERN = re.compile(
+    r'^\s*%"_cudapy_lmem".*?align (\d+)',
+    re.MULTILINE,
+)
+
+
+DTYPES = [np.uint8, np.uint32, np.uint64]
+
+# Add in some record dtypes with and without alignment.
+for align in (True, False):
+    DTYPES += [
+        np.dtype(
+            [
+                ("a", np.uint8),
+                ("b", np.int32),
+                ("c", np.float64),
+            ],
+            align=align,
+        ),
+        np.dtype(
+            [
+                ("a", np.uint32),
+                ("b", np.uint8),
+            ],
+            align=align,
+        ),
+        np.dtype(
+            [
+                ("a", np.uint8),
+                ("b", np.int32),
+                ("c", np.float64),
+                ("d", np.complex64),
+                ("e", (np.uint8, 5)),
+            ],
+            align=align,
+        ),
+    ]
+
+# N.B. We name the test class TestArrayAddressAlignment to avoid name conflict
+#      with the test_alignment.TestArrayAlignment class.
+
+
+class TestArrayAddressAlignment(CUDATestCase):
+    """
+    Test cuda.local.array and cuda.shared.array support for an alignment
+    keyword argument.
+    """
+
+    def test_array_alignment_1d(self):
+        shapes = (1, 8, 50)
+        alignments = (None, 16, 64, 256)
+        array_types = [(0, "local"), (1, "shared")]
+        self._do_test(array_types, shapes, DTYPES, alignments)
+
+    def test_array_alignment_2d(self):
+        shapes = ((2, 3),)
+        alignments = (None, 16, 64, 256)
+        array_types = [(0, "local"), (1, "shared")]
+        self._do_test(array_types, shapes, DTYPES, alignments)
+
+    def test_array_alignment_3d(self):
+        shapes = ((2, 3, 4), (1, 4, 5))
+        alignments = (None, 16, 64, 256)
+        array_types = [(0, "local"), (1, "shared")]
+        self._do_test(array_types, shapes, DTYPES, alignments)
+
+    def _do_test(self, array_types, shapes, dtypes, alignments):
+        items = itertools.product(array_types, shapes, dtypes, alignments)
+
+        for (which, array_type), shape, dtype, alignment in items:
+            with self.subTest(
+                array_type=array_type,
+                shape=shape,
+                dtype=dtype,
+                alignment=alignment,
+            ):
+
+                @cuda.jit
+                def f(loc, shrd, which):
+                    i = cuda.grid(1)
+                    if which == 0:
+                        local_array = cuda.local.array(
+                            shape=shape,
+                            dtype=dtype,
+                            alignment=alignment,
+                        )
+                        if i == 0:
+                            loc[0] = local_array.ctypes.data
+                    else:
+                        shared_array = cuda.shared.array(
+                            shape=shape,
+                            dtype=dtype,
+                            alignment=alignment,
+                        )
+                        if i == 0:
+                            shrd[0] = shared_array.ctypes.data
+
+                loc = np.zeros(1, dtype=np.uint64)
+                shrd = np.zeros(1, dtype=np.uint64)
+                f[1, 1](loc, shrd, which)
+
+                kernel = f.overloads[f.signatures[0]]
+                llvm_ir = kernel.inspect_llvm()
+
+                if alignment is None:
+                    if which == 0:
+                        # Local memory shouldn't have any alignment information
+                        # when no alignment is specified.
+                        match = LMEM_PATTERN.findall(llvm_ir)
+                        self.assertEqual(len(match), 0)
+                    else:
+                        # Shared memory should at least have a power-of-two
+                        # alignment when no alignment is specified.
+                        match = SMEM_PATTERN.findall(llvm_ir)
+                        self.assertEqual(len(match), 1)
+
+                        alignment = int(match[0])
+                        # Verify alignment is a power of two.
+                        self.assertTrue(alignment & (alignment - 1) == 0)
+                else:
+                    # Verify alignment is in the LLVM IR.
+                    if which == 0:
+                        match = LMEM_PATTERN.findall(llvm_ir)
+                        self.assertEqual(len(match), 1)
+                        actual_alignment = int(match[0])
+                        self.assertEqual(alignment, actual_alignment)
+                    else:
+                        match = SMEM_PATTERN.findall(llvm_ir)
+                        self.assertEqual(len(match), 1)
+                        actual_alignment = int(match[0])
+                        self.assertEqual(alignment, actual_alignment)
+
+                    # Also verify that the address of the array is aligned.
+                    # If this fails, there problem is likely with NVVM.
+                    address = loc[0] if which == 0 else shrd[0]
+                    alignment_mod = int(address % alignment)
+                    self.assertEqual(alignment_mod, 0)
+
+                if NOISY:
+                    print(".", end="", flush=True)
+
+    def test_invalid_aligments(self):
+        shapes = (1, 50)
+        dtypes = (np.uint8, np.uint64)
+        invalid_alignment_values = (-1, 0, 3, 17, 33)
+        invalid_alignment_types = ("1.0", "1", "foo", 1.0, 1.5, 3.2)
+        alignments = invalid_alignment_values + invalid_alignment_types
+        array_types = [(0, "local"), (1, "shared")]
+
+        # Use regex pattern to match error message, handling potential ANSI
+        # color codes which appear on CI.
+        expected_invalid_type_error_regex = (
+            r"RequireLiteralValue:.*alignment must be a constant integer"
+        )
+
+        items = itertools.product(array_types, shapes, dtypes, alignments)
+
+        for (which, array_type), shape, dtype, alignment in items:
+            with self.subTest(
+                array_type=array_type,
+                shape=shape,
+                dtype=dtype,
+                alignment=alignment,
+            ):
+                if which == 0:
+
+                    @cuda.jit
+                    def f(dest_array):
+                        i = cuda.grid(1)
+                        local_array = cuda.local.array(
+                            shape=shape,
+                            dtype=dtype,
+                            alignment=alignment,
+                        )
+                        if i == 0:
+                            dest_array[0] = local_array.ctypes.data
+                else:
+
+                    @cuda.jit
+                    def f(dest_array):
+                        i = cuda.grid(1)
+                        shared_array = cuda.shared.array(
+                            shape=shape,
+                            dtype=dtype,
+                            alignment=alignment,
+                        )
+                        if i == 0:
+                            dest_array[0] = shared_array.ctypes.data
+
+                array = np.zeros(1, dtype=np.uint64)
+
+                # The type of error we expect differs between an invalid value
+                # that is still an int, and an invalid type.
+                if isinstance(alignment, int):
+                    self.assertRaisesRegex(
+                        ValueError, r"Alignment must be.*", f[1, 1], array
+                    )
+                else:
+                    self.assertRaisesRegex(
+                        TypingError,
+                        expected_invalid_type_error_regex,
+                        f[1, 1],
+                        array,
+                    )
+
+                if NOISY:
+                    print(".", end="", flush=True)
+
+
+if __name__ == "__main__":
+    unittest.main()

numba_cuda/numba/cuda/tests/cudapy/test_extending.py

@@ -1,7 +1,10 @@
 from numba.cuda.testing import skip_on_cudasim, unittest, CUDATestCase
+from llvmlite import ir
 
 import numpy as np
+import os
 from numba import config, cuda, njit, types
+from numba.extending import overload
 
 
 class Interval:
@@ -160,5 +163,142 @@ class TestExtending(CUDATestCase):
         np.testing.assert_allclose(r, expected)
 
 
+TEST_BIN_DIR = os.getenv("NUMBA_CUDA_TEST_BIN_DIR")
+if TEST_BIN_DIR:
+    test_device_functions_a = os.path.join(
+        TEST_BIN_DIR, "test_device_functions.a"
+    )
+    test_device_functions_cubin = os.path.join(
+        TEST_BIN_DIR, "test_device_functions.cubin"
+    )
+    test_device_functions_cu = os.path.join(
+        TEST_BIN_DIR, "test_device_functions.cu"
+    )
+    test_device_functions_fatbin = os.path.join(
+        TEST_BIN_DIR, "test_device_functions.fatbin"
+    )
+    test_device_functions_fatbin_multi = os.path.join(
+        TEST_BIN_DIR, "test_device_functions_multi.fatbin"
+    )
+    test_device_functions_o = os.path.join(
+        TEST_BIN_DIR, "test_device_functions.o"
+    )
+    test_device_functions_ptx = os.path.join(
+        TEST_BIN_DIR, "test_device_functions.ptx"
+    )
+    test_device_functions_ltoir = os.path.join(
+        TEST_BIN_DIR, "test_device_functions.ltoir"
+    )
+
+
+class TestExtendingLinkage(CUDATestCase):
+    def test_extension_adds_linkable_code(self):
+        cuda_major_version = cuda.runtime.get_version()[0]
+
+        if cuda_major_version < 12:
+            self.skipTest("CUDA 12 required for linking in-memory data")
+
+        files = (
+            (test_device_functions_a, cuda.Archive),
+            (test_device_functions_cubin, cuda.Cubin),
+            (test_device_functions_cu, cuda.CUSource),
+            (test_device_functions_fatbin, cuda.Fatbin),
+            (test_device_functions_o, cuda.Object),
+            (test_device_functions_ptx, cuda.PTXSource),
+            (test_device_functions_ltoir, cuda.LTOIR),
+        )
+
+        lto = config.CUDA_ENABLE_PYNVJITLINK
+
+        for path, ctor in files:
+            if ctor == cuda.LTOIR and not lto:
+                # Don't try to test with LTOIR if LTO is not enabled
+                continue
+
+            with open(path, "rb") as f:
+                code_object = ctor(f.read())
+
+            def external_add(x, y):
+                return x + y
+
+            @type_callable(external_add)
+            def type_external_add(context):
+                def typer(x, y):
+                    if x == types.uint32 and y == types.uint32:
+                        return types.uint32
+
+                return typer
+
+            @lower_builtin(external_add, types.uint32, types.uint32)
+            def lower_external_add(context, builder, sig, args):
+                context.active_code_library.add_linking_file(code_object)
+                i32 = ir.IntType(32)
+                fnty = ir.FunctionType(i32, [i32, i32])
+                fn = cgutils.get_or_insert_function(
+                    builder.module, fnty, "add_cabi"
+                )
+                return builder.call(fn, args)
+
+            @cuda.jit(lto=lto)
+            def use_external_add(r, x, y):
+                r[0] = external_add(x[0], y[0])
+
+            r = np.zeros(1, dtype=np.uint32)
+            x = np.ones(1, dtype=np.uint32)
+            y = np.ones(1, dtype=np.uint32) * 2
+
+            use_external_add[1, 1](r, x, y)
+
+            np.testing.assert_equal(r[0], 3)
+
+            @cuda.jit(lto=lto)
+            def use_external_add_device(x, y):
+                return external_add(x, y)
+
+            @cuda.jit(lto=lto)
+            def use_external_add_kernel(r, x, y):
+                r[0] = use_external_add_device(x[0], y[0])
+
+            r = np.zeros(1, dtype=np.uint32)
+            x = np.ones(1, dtype=np.uint32)
+            y = np.ones(1, dtype=np.uint32) * 2
+
+            use_external_add_kernel[1, 1](r, x, y)
+
+            np.testing.assert_equal(r[0], 3)
+
+    def test_linked_called_through_overload(self):
+        cu_code = cuda.CUSource("""
+            extern "C" __device__
+            int bar(int *out, int a)
+            {
+              *out = a * 2;
+              return 0;
+            }
+        """)
+
+        bar = cuda.declare_device("bar", "int32(int32)", link=cu_code)
+
+        def bar_call(val):
+            pass
+
+        @overload(bar_call, target="cuda")
+        def ol_bar_call(a):
+            return lambda a: bar(a)
+
+        @cuda.jit("void(int32[::1], int32[::1])")
+        def foo(r, x):
+            i = cuda.grid(1)
+            if i < len(r):
+                r[i] = bar_call(x[i])
+
+        x = np.arange(10, dtype=np.int32)
+        r = np.empty_like(x)
+
+        foo[1, 32](r, x)
+
+        np.testing.assert_equal(r, x * 2)
+
+
 if __name__ == "__main__":
     unittest.main()

numba_cuda/numba/cuda/tests/cudapy/test_inline.py

@@ -8,8 +8,8 @@ from numba.cuda.testing import (
 )
 
 
+@skip_on_cudasim("Cudasim does not support inline and forceinline")
 class TestCudaInline(CUDATestCase):
-    @skip_on_cudasim("Cudasim does not support inline")
     def _test_call_inline(self, inline):
         """Test @cuda.jit(inline=...)"""
         a = np.ones(2, dtype=np.int32)
@@ -42,6 +42,9 @@ class TestCudaInline(CUDATestCase):
             # check that call was not inlined
             self.assertIsNotNone(match, msg=llvm_ir)
 
+        # alwaysinline should not be in the IR when the inline kwarg is used
+        self.assertNotIn("alwaysinline", llvm_ir)
+
     def test_call_inline_always(self):
         self._test_call_inline("always")
 
@@ -54,6 +57,100 @@ class TestCudaInline(CUDATestCase):
     def test_call_inline_false(self):
         self._test_call_inline(False)
 
+    def _test_call_forceinline(self, forceinline):
+        """Test @cuda.jit(forceinline=...)"""
+        a = np.ones(2, dtype=np.int32)
+
+        sig = (types.int32[::1],)
+
+        @cuda.jit(forceinline=forceinline)
+        def set_zero(a):
+            a[0] = 0
+
+        @cuda.jit(sig)
+        def call_set_zero(a):
+            set_zero(a)
+
+        call_set_zero[1, 2](a)
+
+        expected = np.arange(2, dtype=np.int32)
+        self.assertTrue(np.all(a == expected))
+
+        llvm_ir = call_set_zero.inspect_llvm(sig)
+        pat = r"call [a-zA-Z0-9]* @"
+        match = re.compile(pat).search(llvm_ir)
+
+        # Check that call was not inlined at the Numba IR level - the call
+        # should still be present in the IR
+        self.assertIsNotNone(match)
+
+        # Check the definition of set_zero - it is a definition where the
+        # name does not include an underscore just before "set_zero", because
+        # that would match the "call_set_zero" definition
+        pat = r"define.*[^_]set_zero.*"
+        match = re.compile(pat).search(llvm_ir)
+        self.assertIsNotNone(match)
+        if forceinline:
+            self.assertIn("alwaysinline", match.group())
+        else:
+            self.assertNotIn("alwaysinline", match.group())
+
+        # The kernel, "call_set_zero", should never have "alwaysinline" set
+        pat = r"define.*call_set_zero.*"
+        match = re.compile(pat).search(llvm_ir)
+        self.assertIsNotNone(match)
+        self.assertNotIn("alwaysinline", match.group())
+
+    def test_call_forceinline_true(self):
+        self._test_call_forceinline(True)
+
+    def test_call_forceinline_false(self):
+        self._test_call_forceinline(False)
+
+    def test_compile_forceinline_ltoir_only(self):
+        def set_zero(a):
+            a[0] = 0
+
+        args = (types.float32[::1],)
+        msg = r"Can only designate forced inlining in LTO-IR"
+        with self.assertRaisesRegex(ValueError, msg):
+            cuda.compile(
+                set_zero,
+                args,
+                device=True,
+                forceinline=True,
+            )
+
+    def _compile_set_zero(self, forceinline):
+        def set_zero(a):
+            a[0] = 0
+
+        args = (types.float32[::1],)
+        ltoir, resty = cuda.compile(
+            set_zero,
+            args,
+            device=True,
+            output="ltoir",
+            forceinline=forceinline,
+        )
+
+        # Sanity check
+        self.assertEqual(resty, types.none)
+
+        return ltoir
+
+    def test_compile_forceinline(self):
+        ltoir_noinline = self._compile_set_zero(False)
+        ltoir_forceinline = self._compile_set_zero(True)
+
+        # As LTO-IR is opaque, the best we can do is check that changing the
+        # flag resulted in a change in the generated LTO-IR in some way.
+        self.assertNotEqual(
+            ltoir_noinline,
+            ltoir_forceinline,
+            "forceinline flag appeared to have no effect on LTO-IR",
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

numba_cuda/numba/cuda/tests/nrt/test_nrt.py

@@ -4,11 +4,86 @@ import os
 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.cuda.runtime.nrt import rtsys, get_include
+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_deallocate_handle():
+    """
+    Handle to call NRT_Allocate and NRT_Free
+    """
+    pass
+
+
+@cuda_decl_registry.register_global(allocate_deallocate_handle)
+class AllocateShimImpl(AbstractTemplate):
+    def generic(self, args, kws):
+        return signature(types.void)
+
+
+device_fun_shim = cuda.declare_device(
+    "device_allocate_deallocate", types.int32()
+)
+
+
+# wrapper to turn the above into a python callable
+def call_device_fun_shim():
+    return device_fun_shim()
+
+
+@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_device_fun_shim,
+        sig_,
+        (),
+    )
+
+    return result
 
 
 class TestNrtBasic(CUDATestCase):
@@ -77,6 +152,50 @@ class TestNrtBasic(CUDATestCase):
         self.assertEqual(out_ary[0], 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 device_allocate_deallocate(int* nb_retval){
+                     auto ptr = NRT_Allocate(1);
+                     NRT_Free(ptr);
+                     return 0;
+                 }
+        """
+        cc = get_current_device().compute_capability
+        ptx, _ = compile(src, "external_nrt.cu", cc)
+
+        @cuda.jit(link=[PTXSource(ptx.encode(), nrt=True)])
+        def kernel():
+            allocate_deallocate_handle()
+
+        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_deallocate_handle()
+
+                kernel[1, 1]()
+
+
 class TestNrtStatistics(CUDATestCase):
     def setUp(self):
         self._stream = cuda.default_stream()

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,7 +54,16 @@ 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
+	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

numba_cuda/numba/cuda/tests/test_binary_generation/generate_raw_ltoir.py

@@ -7,6 +7,7 @@ import subprocess
 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,7 @@ if __name__ == "__main__":
         help="compute arch to target (e.g. sm_87). Defaults to sm_50.",
         default="sm_50",
     )
-
+    parser.add_argument("--nrt", action="store_true")
     args = parser.parse_args()
     outputpath = args.output