Allow Numba NVRTC Binding Search Additional Paths

docs/source/reference/envvars.rst

@@ -131,3 +131,9 @@ target.
    Use ``pynvjitlink`` for minor version compatibility. Requires the ``pynvjitlink``
    package to be installed. Provides minor version compatibility for driver versions
    greater than 12.0.
+
+.. envvar:: NUMBA_CUDA_RTC_EXTRA_SEARCH_PATHS
+
+   A colon separated list of paths that Numba's NVRTC should search for when compiling
+   external functions. These folders are searched after the system cudatoolkit search
+   paths and Numba-CUDA's internal search paths.

docs/source/user/cuda_ffi.rst

@@ -233,7 +233,58 @@ of NVRTC subject to the following considerations:
   on Linux and ``$env:CUDA_PATH\include`` on Windows. It can be modified using
   the environment variable :envvar:`NUMBA_CUDA_INCLUDE_PATH`.
 - The CUDA include directory will be made available to NVRTC on the include
-  path; additional includes are not supported.
+  path.
+- Additional search paths can be set to the environment variable
+   :envvar:`NUMBA_CUDA_RTC_EXTRA_SEARCH_PATHS`. Multiple paths should be colon
+   separated.
+
+Extra Search Paths Example
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+This example demonstrates calling a foreign function which includes additional
+headers that does not exist in default Numba-CUDA search paths.
+
+The definition of the C++ template APIs are defined in two different locations:
+
+.. literalinclude:: ../../../numba_cuda/numba/cuda/tests/doc_examples/ffi/include/mul.cuh
+   :language: C
+   :caption: ``numba/cuda/tests/doc_examples/ffi/include/mul.cuh``
+   :linenos:
+
+.. literalinclude:: ../../../numba_cuda/numba/cuda/tests/data/include/add.cuh
+   :language: C
+   :caption: ``numba/cuda/tests/data/include/add.cuh``
+   :linenos:
+
+Neither header exists in the default search paths of Numba-CUDA, however, the
+foreign device function ``saxpy`` depends on them:
+
+.. literalinclude:: ../../../numba_cuda/numba/cuda/tests/doc_examples/ffi/saxpy.cu
+   :language: C
+   :caption: ``numba/cuda/tests/data/doc_examples/ffi/saxpy.cu``
+   :linenos:
+
+In the Python Code, assume ``mul_dir`` and ``add_dir`` are set to the folder that contains
+``mul.cuh`` and ``add.cuh`` respectively. Use ``:`` to join them to a single string and
+set to ``CUDA_RTC_EXTRA_SEARCH_PATHS`` in config.
+
+.. literalinclude:: ../../../numba_cuda/numba/cuda/tests/doc_examples/test_ffi.py
+   :language: python
+   :caption: from ``test_ex_extra_includes`` in ``numba/cuda/tests/doc_examples/test_ffi.py``
+   :start-after: magictoken.ex_extra_search_paths.begin
+   :end-before: magictoken.ex_extra_search_paths.end
+   :dedent: 8
+   :linenos:
+
+Next, use ``saxpy`` as intended.
+
+.. literalinclude:: ../../../numba_cuda/numba/cuda/tests/doc_examples/test_ffi.py
+   :language: python
+   :caption: from ``test_ex_extra_includes`` in ``numba/cuda/tests/doc_examples/test_ffi.py``
+   :start-after: magictoken.ex_extra_search_paths_kernel.begin
+   :end-before: magictoken.ex_extra_search_paths_kernel.end
+   :dedent: 8
+   :linenos:
 
 
 Complete Example

numba_cuda/numba/cuda/cudadrv/nvrtc.py

@@ -6,13 +6,21 @@
     NvrtcCompilationError,
     NvrtcSupportError,
 )
+from numba import config
 from numba.cuda.cuda_paths import get_cuda_paths
+from numba.cuda.utils import _readenv
 
 import functools
 import os
 import threading
 import warnings
 
+RTC_EXTRA_SEARCH_PATHS = _readenv(
+    "NUMBA_CUDA_RTC_EXTRA_SEARCH_PATHS", str, ""
+) or getattr(config, "NUMBA_CUDA_RTC_EXTRA_SEARCH_PATHS", "")
+if not hasattr(config, "NUMBA_CUDA_RTC_EXTRA_SEARCH_PATHS"):
+    config.CUDA_RTC_EXTRA_SEARCH_PATHS = RTC_EXTRA_SEARCH_PATHS
+
 # Opaque handle for compilation unit
 nvrtc_program = c_void_p
 
@@ -383,10 +391,24 @@ def compile(src, name, cc, ltoir=False):
     else:
         numba_include = f"-I{os.path.join(numba_cuda_path, 'include', '12')}"
 
+    if config.CUDA_RTC_EXTRA_SEARCH_PATHS:
+        extra_search_paths = config.CUDA_RTC_EXTRA_SEARCH_PATHS.split(":")
+        extra_includes = [f"-I{p}" for p in extra_search_paths]
+    else:
+        extra_includes = []
+
     nrt_path = os.path.join(numba_cuda_path, "runtime")
     nrt_include = f"-I{nrt_path}"
 
-    options = [arch, numba_include, *cuda_include, nrt_include, "-rdc", "true"]
+    options = [
+        arch,
+        numba_include,
+        *cuda_include,
+        nrt_include,
+        *extra_includes,
+        "-rdc",
+        "true",
+    ]
 
     if ltoir:
         options.append("-dlto")

numba_cuda/numba/cuda/tests/data/include/add.cuh

@@ -0,0 +1,3 @@
+// Templated addition function: myadd
+template <typename T>
+__device__ T myadd(T a, T b) { return a + b; }

numba_cuda/numba/cuda/tests/doc_examples/ffi/include/mul.cuh

@@ -0,0 +1,3 @@
+// Templated multiplication function: mymul
+template <typename T>
+__device__ T mymul(T a, T b) { return a * b; }

numba_cuda/numba/cuda/tests/doc_examples/ffi/saxpy.cu

@@ -0,0 +1,9 @@
+#include <add.cuh> // In numba/cuda/tests/data/include
+#include <mul.cuh> // In numba/cuda/tests/doc_examples/ffi/include
+
+extern "C"
+__device__ int saxpy(float *ret, float a, float x, float y)
+{
+    *ret = myadd(mymul(a, x), y);
+    return 0;
+}

numba_cuda/numba/cuda/tests/doc_examples/test_ffi.py

@@ -85,6 +85,56 @@ def reduction_caller(result, array):
         actual = r[()]
         np.testing.assert_allclose(expected, actual)
 
+    def test_ex_extra_includes(self):
+        import numpy as np
+        from numba import cuda, config
+        import os
+
+        basedir = os.path.dirname(os.path.abspath(__file__))
+        mul_dir = os.path.join(basedir, "ffi", "include")
+        saxpy_cu = os.path.join(basedir, "ffi", "saxpy.cu")
+
+        testdir = os.path.dirname(basedir)
+        add_dir = os.path.join(testdir, "data", "include")
+
+        old_setting = config.CUDA_RTC_EXTRA_SEARCH_PATHS
+
+        # magictoken.ex_extra_search_paths.begin
+        from numba import config
+
+        includedir = ":".join([mul_dir, add_dir])
+        config.CUDA_RTC_EXTRA_SEARCH_PATHS = includedir
+        # magictoken.ex_extra_search_paths.end
+
+        # magictoken.ex_extra_search_paths_kernel.begin
+        sig = "float32(float32, float32, float32)"
+        saxpy = cuda.declare_device("saxpy", sig=sig, link=saxpy_cu)
+
+        @cuda.jit
+        def vector_saxpy(a, x, y, res):
+            i = cuda.grid(1)
+            if i < len(res):
+                res[i] = saxpy(a, x[i], y[i])
+
+        # magictoken.ex_extra_search_paths_kernel.end
+
+        size = 10_000
+        a = 3.0
+        X = np.ones((size,), dtype="float32")
+        Y = np.ones((size,), dtype="float32")
+        R = np.zeros((size,), dtype="float32")
+
+        block_size = 32
+        num_blocks = (size // block_size) + 1
+
+        vector_saxpy[num_blocks, block_size](a, X, Y, R)
+
+        expected = a * X + Y
+        np.testing.assert_equal(R, expected)
+
+        # Reset config setting
+        config.CUDA_RTC_EXTRA_SEARCH_PATHS = old_setting
+
 
 if __name__ == "__main__":
     unittest.main()