Support linking code for device functions in declaration

Author: gmarkall

numba_cuda/numba/cuda/compiler.py

@@ -570,16 +570,16 @@ def compile_ptx_for_current_device(pyfunc, sig, debug=None, lineinfo=False,
                        abi=abi, abi_info=abi_info)
 
 
-def declare_device_function(name, restype, argtypes):
-    return declare_device_function_template(name, restype, argtypes).key
+def declare_device_function(name, restype, argtypes, link):
+    return declare_device_function_template(name, restype, argtypes, link).key
 
 
-def declare_device_function_template(name, restype, argtypes):
+def declare_device_function_template(name, restype, argtypes, link):
     from .descriptor import cuda_target
     typingctx = cuda_target.typing_context
     targetctx = cuda_target.target_context
     sig = typing.signature(restype, *argtypes)
-    extfn = ExternFunction(name, sig)
+    extfn = ExternFunction(name, sig, link)
 
     class device_function_template(ConcreteTemplate):
         key = extfn
@@ -593,7 +593,8 @@ class device_function_template(ConcreteTemplate):
     return device_function_template
 
 
-class ExternFunction(object):
-    def __init__(self, name, sig):
+class ExternFunction:
+    def __init__(self, name, sig, link):
         self.name = name
         self.sig = sig
+        self.link = link

numba_cuda/numba/cuda/cudadecl.py

@@ -403,16 +403,20 @@ def _genfp16_binary_operator(op):
 
 
 def _resolve_wrapped_unary(fname):
+    link = tuple()
     decl = declare_device_function_template(f'__numba_wrapper_{fname}',
                                             types.float16,
-                                            (types.float16,))
+                                            (types.float16,),
+                                            link)
     return types.Function(decl)
 
 
 def _resolve_wrapped_binary(fname):
+    link = tuple()
     decl = declare_device_function_template(f'__numba_wrapper_{fname}',
                                             types.float16,
-                                            (types.float16, types.float16,))
+                                            (types.float16, types.float16,),
+                                            link)
     return types.Function(decl)
 
 

numba_cuda/numba/cuda/decorators.py

@@ -173,7 +173,7 @@ def autojitwrapper(func):
                 return disp
 
 
-def declare_device(name, sig):
+def declare_device(name, sig, link=None):
     """
     Declare the signature of a foreign function. Returns a descriptor that can
     be used to call the function from a Python kernel.
@@ -182,9 +182,15 @@ def declare_device(name, sig):
     :type name: str
     :param sig: The Numba signature of the function.
     """
+    if link is None:
+        link = tuple()
+    else:
+        if not isinstance(link, (list, tuple, set)):
+            link = (link,)
+
     argtypes, restype = sigutils.normalize_signature(sig)
     if restype is None:
         msg = 'Return type must be provided for device declarations'
         raise TypeError(msg)
 
-    return declare_device_function(name, restype, argtypes)
+    return declare_device_function(name, restype, argtypes, link)

numba_cuda/numba/cuda/dispatcher.py

@@ -11,10 +11,11 @@
 from numba.core.dispatcher import Dispatcher
 from numba.core.errors import NumbaPerformanceWarning
 from numba.core.typing.typeof import Purpose, typeof
-
+from numba.core.types.functions import Function
 from numba.cuda.api import get_current_device
 from numba.cuda.args import wrap_arg
-from numba.cuda.compiler import compile_cuda, CUDACompiler, kernel_fixup
+from numba.cuda.compiler import (compile_cuda, CUDACompiler, kernel_fixup,
+                                 ExternFunction)
 from numba.cuda.cudadrv import driver
 from numba.cuda.cudadrv.devices import get_context
 from numba.cuda.descriptor import cuda_target
@@ -158,6 +159,16 @@ def link_to_library_functions(library_functions, library_path,
 
         self.maybe_link_nrt(link, tgt_ctx, asm)
 
+        for k, v in cres.fndesc.typemap.items():
+            if not isinstance(v, Function):
+                continue
+
+            if not isinstance(v.typing_key, ExternFunction):
+                continue
+
+            for obj in v.typing_key.link:
+                lib.add_linking_file(obj)
+
         for filepath in link:
             lib.add_linking_file(filepath)
 

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

@@ -1,11 +1,13 @@
 import re
-import types
+import cffi
 
 import numpy as np
 
-from numba.cuda.testing import unittest, skip_on_cudasim, CUDATestCase
-from numba import cuda, jit, float32, int32
+from numba.cuda.testing import (skip_on_cudasim, test_data_dir, unittest,
+                                CUDATestCase)
+from numba import cuda, jit, float32, int32, types
 from numba.core.errors import TypingError
+from types import ModuleType
 
 
 class TestDeviceFunc(CUDATestCase):
@@ -92,7 +94,7 @@ def test_cpu_dispatcher_other_module(self):
         def add(a, b):
             return a + b
 
-        mymod = types.ModuleType(name='mymod')
+        mymod = ModuleType(name='mymod')
         mymod.add = add
         del add
 
@@ -192,31 +194,128 @@ def rgba_caller(x, channels):
 
         self.assertEqual(0x04010203, x[0])
 
-    def _test_declare_device(self, decl):
+
+times2_cu = cuda.CUSource("""
+extern "C" __device__
+int times2(int *out, int a)
+{
+  *out = a * 2;
+  return 0;
+}
+""")
+
+
+times4_cu = cuda.CUSource("""
+extern "C" __device__
+int times2(int *out, int a);
+
+extern "C" __device__
+int times4(int *out, int a)
+{
+  int tmp;
+  times2(&tmp, a);
+  *out = tmp * 2;
+  return 0;
+}
+""")
+
+jitlink_user_cu = cuda.CUSource("""
+extern "C" __device__
+int array_mutator(void *out, int *a);
+
+extern "C" __device__
+int use_array_mutator(void *out, int *a) {
+  array_mutator(out, a);
+  return 0;
+}
+""")
+
+
+@skip_on_cudasim('External functions unsupported in the simulator')
+class TestDeclareDevice(CUDATestCase):
+
+    def check_api(self, decl):
         self.assertEqual(decl.name, 'f1')
         self.assertEqual(decl.sig.args, (float32[:],))
         self.assertEqual(decl.sig.return_type, int32)
 
-    @skip_on_cudasim('cudasim does not check signatures')
     def test_declare_device_signature(self):
         f1 = cuda.declare_device('f1', int32(float32[:]))
-        self._test_declare_device(f1)
+        self.check_api(f1)
 
-    @skip_on_cudasim('cudasim does not check signatures')
     def test_declare_device_string(self):
         f1 = cuda.declare_device('f1', 'int32(float32[:])')
-        self._test_declare_device(f1)
+        self.check_api(f1)
 
-    @skip_on_cudasim('cudasim does not check signatures')
     def test_bad_declare_device_tuple(self):
         with self.assertRaisesRegex(TypeError, 'Return type'):
             cuda.declare_device('f1', (float32[:],))
 
-    @skip_on_cudasim('cudasim does not check signatures')
     def test_bad_declare_device_string(self):
         with self.assertRaisesRegex(TypeError, 'Return type'):
             cuda.declare_device('f1', '(float32[:],)')
 
+    def test_link_cu_source(self):
+        times2 = cuda.declare_device('times2', 'int32(int32)', link=times2_cu)
+
+        @cuda.jit
+        def kernel(r, x):
+            i = cuda.grid(1)
+            if i < len(r):
+                r[i] = times2(x[i])
+
+        x = np.arange(10, dtype=np.int32)
+        r = np.empty_like(x)
+
+        kernel[1, 32](r, x)
+
+        np.testing.assert_equal(r, x * 2)
+
+    def _test_link_multiple_sources(self, link_type):
+        link = link_type([times2_cu, times4_cu])
+        times4 = cuda.declare_device('times4', 'int32(int32)', link=link)
+
+        @cuda.jit
+        def kernel(r, x):
+            i = cuda.grid(1)
+            if i < len(r):
+                r[i] = times4(x[i])
+
+        x = np.arange(10, dtype=np.int32)
+        r = np.empty_like(x)
+
+        kernel[1, 32](r, x)
+
+        np.testing.assert_equal(r, x * 4)
+
+    def test_link_multiple_sources_set(self):
+        self._test_link_multiple_sources(set)
+
+    def test_link_multiple_sources_tuple(self):
+        self._test_link_multiple_sources(tuple)
+
+    def test_link_multiple_sources_list(self):
+        self._test_link_multiple_sources(list)
+
+    def test_link_sources_in_memory_and_on_disk(self):
+        jitlink_cu = str(test_data_dir / "jitlink.cu")
+        link = [jitlink_cu, jitlink_user_cu]
+        sig = types.void(types.CPointer(types.int32))
+        ext_fn = cuda.declare_device("use_array_mutator", sig, link=link)
+
+        ffi = cffi.FFI()
+
+        @cuda.jit
+        def kernel(x):
+            ptr = ffi.from_buffer(x)
+            ext_fn(ptr)
+
+        x = np.arange(2, dtype=np.int32)
+        kernel[1, 1](x)
+
+        expected = np.ones(2, dtype=np.int32)
+        np.testing.assert_equal(x, expected)
+
 
 if __name__ == '__main__':
     unittest.main()