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

numba_cuda/numba/cuda/dispatcher.py

@@ -1,27 +1,25 @@
 import numpy as np
 import os
-import re
 import sys
 import ctypes
 import functools
-from collections import defaultdict
 
-from numba.core import config, ir, serialize, sigutils, types, typing, utils
+from numba.core import config, serialize, sigutils, types, typing, utils
 from numba.core.caching import Cache, CacheImpl
 from numba.core.compiler_lock import global_compiler_lock
 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,
-    ExternFunction,
 )
+import re
 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 (
@@ -29,7 +27,7 @@ from numba.cuda.errors import (
     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.locks import module_init_lock
 
 from numba import cuda
@@ -59,54 +57,6 @@ cuda_fp16_math_funcs = [
 reshape_funcs = ["nocopy_empty_reshape", "numba_attempt_nocopy_reshape"]
 
 
-def get_cres_link_objects(cres):
-    """Given a compile result, return a set of all linkable code objects that
-    are required for it to be fully linked."""
-
-    link_objects = set()
-
-    # List of calls into declared device functions
-    device_func_calls = [
-        (name, v)
-        for name, v in cres.fndesc.typemap.items()
-        if (isinstance(v, cuda_types.CUDADispatcher))
-    ]
-
-    # List of tuples with SSA name of calls and corresponding signature
-    call_signatures = [
-        (call.func.name, sig)
-        for call, sig in cres.fndesc.calltypes.items()
-        if (isinstance(call, ir.Expr) and call.op == "call")
-    ]
-
-    # Map SSA names to all invoked signatures
-    call_signature_d = defaultdict(list)
-    for name, sig in call_signatures:
-        call_signature_d[name].append(sig)
-
-    # Add the link objects from the current function's callees
-    for name, v in device_func_calls:
-        for sig in call_signature_d.get(name, []):
-            called_cres = v.dispatcher.overloads[sig.args]
-            called_link_objects = get_cres_link_objects(called_cres)
-            link_objects.update(called_link_objects)
-
-    # From this point onwards, we are only interested in ExternFunction
-    # declarations - these are the calls made directly in this function to
-    # them.
-    for name, 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:
-            link_objects.add(obj)
-
-    return link_objects
-
-
 class _Kernel(serialize.ReduceMixin):
     """
     CUDA Kernel specialized for a given set of argument types. When called, this
@@ -238,9 +188,6 @@ class _Kernel(serialize.ReduceMixin):
 
         self.maybe_link_nrt(link, tgt_ctx, asm)
 
-        for obj in get_cres_link_objects(cres):
-            lib.add_linking_file(obj)
-
         for filepath in link:
             lib.add_linking_file(filepath)
 
@@ -263,6 +210,13 @@ class _Kernel(serialize.ReduceMixin):
         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.
+        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
 
@@ -272,13 +226,19 @@ class _Kernel(serialize.ReduceMixin):
             + all_nrt
             + r")\s*\([^)]*\)\s*;"
         )
-
+        link_nrt = False
         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)
+        if len(nrt_in_asm) > 0:
+            link_nrt = True
+        if not link_nrt:
+            for file in link:
+                if isinstance(file, LinkableCode):
+                    if file.nrt:
+                        link_nrt = True
+                        break
+
+        if link_nrt:
+            link.append(NRT_LIBRARY)
 
     @property
     def library(self):

numba_cuda/numba/cuda/runtime/nrt.cu

@@ -4,30 +4,14 @@
 #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)
 {
   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)
 {
@@ -177,6 +161,7 @@ extern "C" __device__ void NRT_decref(NRT_MemInfo* mi)
   }
 }
 
+
 #endif
 
 extern "C" __device__ void NRT_incref(NRT_MemInfo* mi)

numba_cuda/numba/cuda/runtime/nrt.cuh

@@ -0,0 +1,41 @@
+#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);
+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_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);

numba_cuda/numba/cuda/runtime/nrt.py

@@ -13,7 +13,8 @@ from numba.cuda.cudadrv.driver import (
 )
 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
@@ -32,6 +33,11 @@ if not hasattr(config, "NUMBA_CUDA_ENABLE_NRT"):
     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):
     """
@@ -340,3 +346,9 @@ class _Runtime:
 
 # 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)

numba_cuda/numba/cuda/stubs.py

@@ -129,12 +129,16 @@ class shared(Stub):
     _description_ = "<shared>"
 
     @stub_function
-    def array(shape, dtype):
+    def array(shape, dtype, alignment=None):
         """
-        Allocate a shared array of the given *shape* and *type*. *shape* is
-        either an integer or a tuple of integers representing the array's
-        dimensions.  *type* is a :ref:`Numba type <numba-types>` of the
-        elements needing to be stored in the array.
+        Allocate a shared array of the given *shape*, *type*, and, optionally,
+        *alignment*.  *shape* is either an integer or a tuple of integers
+        representing the array's dimensions.  *type* is a :ref:`Numba type
+        <numba-types>` of the elements needing to be stored in the array.
+        *alignment* is an optional integer specifying the byte alignment of
+        the array.  When specified, it must be a power of two, and a multiple
+        of the size of a pointer (8 bytes).  When not specified, the array is
+        allocated with an alignment appropriate for the supplied *dtype*.
 
         The returned array-like object can be read and written to like any
         normal device array (e.g. through indexing).
@@ -149,12 +153,20 @@ class local(Stub):
     _description_ = "<local>"
 
     @stub_function
-    def array(shape, dtype):
-        """
-        Allocate a local array of the given *shape* and *type*. The array is
-        private to the current thread, and resides in global memory. An
-        array-like object is returned which can be read and written to like any
-        standard array (e.g.  through indexing).
+    def array(shape, dtype, alignment=None):
+        """
+        Allocate a local array of the given *shape*, *type*, and, optionally,
+        *alignment*.  *shape* is either an integer or a tuple of integers
+        representing the array's dimensions.  *type* is a :ref:`Numba type
+        <numba-types>` of the elements needing to be stored in the array.
+        *alignment* is an optional integer specifying the byte alignment of
+        the array.  When specified, it must be a power of two, and a multiple
+        of the size of a pointer (8 bytes).  When not specified, the array is
+        allocated with an alignment appropriate for the supplied *dtype*.
+
+        The array is private to the current thread, and resides in global
+        memory.  An array-like object is returned which can be read and
+        written to like any standard array (e.g. through indexing).
         """
 
 

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()