Migrate numba-cuda driver to use cuda.core.launch API

numba_cuda/numba/cuda/cudadrv/nvrtc.py

@@ -109,6 +109,22 @@ def compile(src, name, cc, ltoir=False, lineinfo=False, debug=False):
 
     includes = [numba_include, *cuda_includes, nrt_include, *extra_includes]
 
+    # TODO: move all this into Program/ProgramOptions
+    # logsz = config.CUDA_LOG_SIZE
+    #
+    # jitinfo = bytearray(logsz)
+    # jiterrors = bytearray(logsz)
+    #
+    # jit_option = binding.CUjit_option
+    # options = {
+    #     jit_option.CU_JIT_INFO_LOG_BUFFER: jitinfo,
+    #     jit_option.CU_JIT_INFO_LOG_BUFFER_SIZE_BYTES: logsz,
+    #     jit_option.CU_JIT_ERROR_LOG_BUFFER: jiterrors,
+    #     jit_option.CU_JIT_ERROR_LOG_BUFFER_SIZE_BYTES: logsz,
+    #     jit_option.CU_JIT_LOG_VERBOSE: config.CUDA_VERBOSE_JIT_LOG,
+    # }
+    # info_log = jitinfo.decode("utf-8")
+
     options = ProgramOptions(
         arch=arch,
         include_path=includes,

numba_cuda/numba/cuda/dispatcher.py

@@ -15,6 +15,8 @@
 import re
 from warnings import warn
 
+from cuda.core.experimental import launch
+
 from numba.cuda.core import errors
 from numba.cuda import serialize, utils
 from numba import cuda
@@ -39,6 +41,7 @@
 from numba.cuda.core import sigutils, config, entrypoints
 from numba.cuda.flags import Flags
 from numba.cuda.cudadrv import driver, nvvm
+from cuda.core.experimental import LaunchConfig
 from numba.cuda.locks import module_init_lock
 from numba.cuda.core.caching import Cache, CacheImpl, NullCache
 from numba.cuda.descriptor import cuda_target
@@ -475,18 +478,15 @@ def launch(self, args, griddim, blockdim, stream=0, sharedmem=0):
         for t, v in zip(self.argument_types, args):
             self._prepare_args(t, v, stream, retr, kernelargs)
 
-        stream_handle = driver._stream_handle(stream)
-
         # Invoke kernel
-        driver.launch_kernel(
-            cufunc.handle,
-            *griddim,
-            *blockdim,
-            sharedmem,
-            stream_handle,
-            kernelargs,
-            cooperative=self.cooperative,
+        config = LaunchConfig(
+            grid=griddim,
+            block=blockdim,
+            shmem_size=sharedmem,
+            cooperative_launch=self.cooperative,
         )
+        kernel = cufunc.kernel
+        launch(stream, config, kernel, *kernelargs)
 
         if self.debug:
             driver.device_to_host(ctypes.addressof(excval), excmem, excsz)
@@ -540,30 +540,26 @@ def _prepare_args(self, ty, val, stream, retr, kernelargs):
 
         if isinstance(ty, types.Array):
             devary = wrap_arg(val).to_device(retr, stream)
-            c_intp = ctypes.c_ssize_t
-
-            meminfo = ctypes.c_void_p(0)
-            parent = ctypes.c_void_p(0)
-            nitems = c_intp(devary.size)
-            itemsize = c_intp(devary.dtype.itemsize)
-
-            ptr = driver.device_pointer(devary)
-
-            ptr = int(ptr)
 
-            data = ctypes.c_void_p(ptr)
+            meminfo = 0
+            parent = 0
 
             kernelargs.append(meminfo)
             kernelargs.append(parent)
-            kernelargs.append(nitems)
-            kernelargs.append(itemsize)
-            kernelargs.append(data)
-            kernelargs.extend(map(c_intp, devary.shape))
-            kernelargs.extend(map(c_intp, devary.strides))
+
+            # non-pointer-arguments-without-ctypes might be dicey, since we're
+            # assuming shape, strides, size, and itemsize fit into intptr_t
+            # however, this saves a noticeable amount of overhead in kernel
+            # invocation
+            kernelargs.append(devary.size)
+            kernelargs.append(devary.dtype.itemsize)
+            kernelargs.append(devary.device_ctypes_pointer.value)
+            kernelargs.extend(devary.shape)
+            kernelargs.extend(devary.strides)
 
         elif isinstance(ty, types.CPointer):
             # Pointer arguments should be a pointer-sized integer
-            kernelargs.append(ctypes.c_uint64(val))
+            kernelargs.append(val)
 
         elif isinstance(ty, types.Integer):
             cval = getattr(ctypes, "c_%s" % ty)(val)
@@ -582,8 +578,7 @@ def _prepare_args(self, ty, val, stream, retr, kernelargs):
             kernelargs.append(cval)
 
         elif ty == types.boolean:
-            cval = ctypes.c_uint8(int(val))
-            kernelargs.append(cval)
+            kernelargs.append(val)
 
         elif ty == types.complex64:
             kernelargs.append(ctypes.c_float(val.real))
@@ -598,8 +593,7 @@ def _prepare_args(self, ty, val, stream, retr, kernelargs):
 
         elif isinstance(ty, types.Record):
             devrec = wrap_arg(val).to_device(retr, stream)
-            ptr = devrec.device_ctypes_pointer
-            kernelargs.append(ptr)
+            kernelargs.append(devrec.device_ctypes_pointer.value)
 
         elif isinstance(ty, types.BaseTuple):
             assert len(ty) == len(val)
@@ -671,7 +665,7 @@ def __init__(self, dispatcher, griddim, blockdim, stream, sharedmem):
         self.dispatcher = dispatcher
         self.griddim = griddim
         self.blockdim = blockdim
-        self.stream = stream
+        self.stream = driver._to_core_stream(stream)
         self.sharedmem = sharedmem
 
         if (
@@ -700,6 +694,16 @@ def __call__(self, *args):
             args, self.griddim, self.blockdim, self.stream, self.sharedmem
         )
 
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state["stream"] = int(state["stream"].handle)
+        return state
+
+    def __setstate__(self, state):
+        handle = state.pop("stream")
+        self.__dict__.update(state)
+        self.stream = driver._to_core_stream(handle)
+
 
 class CUDACacheImpl(CacheImpl):
     def reduce(self, kernel):

numba_cuda/numba/cuda/memory_management/nrt.py

@@ -13,9 +13,10 @@
 from numba.cuda.cudadrv.driver import (
     _Linker,
     driver,
-    launch_kernel,
+    _to_core_stream,
     _have_nvjitlink,
 )
+from cuda.core.experimental import LaunchConfig, launch
 from numba.cuda.cudadrv import devices
 from numba.cuda.api import get_current_device
 from numba.cuda.utils import _readenv, cached_file_read
@@ -179,20 +180,15 @@ def _single_thread_launch(self, module, stream, name, params=()):
             stream = cuda.default_stream()
 
         func = module.get_function(name)
-        launch_kernel(
-            func.handle,
-            1,
-            1,
-            1,
-            1,
-            1,
-            1,
-            0,
-            stream.handle.value,
-            params,
-            cooperative=False,
+        config = LaunchConfig(
+            grid=(1, 1, 1),
+            block=(1, 1, 1),
+            shmem_size=0,
+            cooperative_launch=False,
         )
 
+        launch(_to_core_stream(stream), config, func.kernel, *params)
+
     def ensure_initialized(self, stream=None):
         """
         If memsys is not initialized, initialize memsys

numba_cuda/numba/cuda/tests/cudadrv/test_cuda_driver.py

@@ -3,15 +3,15 @@
 
 from ctypes import c_int, sizeof
 
-from numba.cuda.cudadrv.driver import (
-    host_to_device,
-    device_to_host,
-    driver,
-    launch_kernel,
+from numba.cuda.cudadrv.driver import host_to_device, device_to_host, driver
+from cuda.core.experimental import (
+    LaunchConfig,
+    Stream as ExperimentalStream,
+    launch,
 )
 
 from numba import cuda
-from numba.cuda.cudadrv import devices, driver as _driver
+from numba.cuda.cudadrv import devices
 from numba.cuda.testing import unittest, CUDATestCase
 from numba.cuda.testing import skip_on_cudasim
 import contextlib
@@ -98,22 +98,15 @@ def test_cuda_driver_basic(self):
         host_to_device(memory, array, sizeof(array))
 
         ptr = memory.device_ctypes_pointer
-        stream = 0
-
-        stream = _driver.binding.CUstream(stream)
-
-        launch_kernel(
-            function.handle,  # Kernel
-            1,
-            1,
-            1,  # gx, gy, gz
-            100,
-            1,
-            1,  # bx, by, bz
-            0,  # dynamic shared mem
-            stream,  # stream
-            [ptr],
-        )  # arguments
+
+        config = LaunchConfig(
+            grid=(1, 1, 1),
+            block=(100, 1, 1),
+            shmem_size=0,
+            cooperative_launch=False,
+        )
+        exp_stream = ExperimentalStream.from_handle(0)
+        launch(exp_stream, config, function.kernel, ptr)
 
         device_to_host(array, memory, sizeof(array))
         for i, v in enumerate(array):
@@ -122,6 +115,8 @@ def test_cuda_driver_basic(self):
         module.unload()
 
     def test_cuda_driver_stream_operations(self):
+        from numba.cuda.cudadrv.driver import _to_core_stream
+
         module = self.context.create_module_ptx(self.ptx)
         function = module.get_function("_Z10helloworldPi")
 
@@ -135,21 +130,14 @@ def test_cuda_driver_stream_operations(self):
 
             ptr = memory.device_ctypes_pointer
 
-            stream_handle = stream.handle
-            stream_handle = stream_handle.value
-
-            launch_kernel(
-                function.handle,  # Kernel
-                1,
-                1,
-                1,  # gx, gy, gz
-                100,
-                1,
-                1,  # bx, by, bz
-                0,  # dynamic shared mem
-                stream_handle,  # stream
-                [ptr],
-            )  # arguments
+            config = LaunchConfig(
+                grid=(1, 1, 1),
+                block=(100, 1, 1),
+                shmem_size=0,
+                cooperative_launch=False,
+            )
+            # Convert numba Stream to ExperimentalStream
+            launch(_to_core_stream(stream), config, function.kernel, ptr)
 
         device_to_host(array, memory, sizeof(array), stream=stream)
 
@@ -177,18 +165,13 @@ def auto_synchronize(stream):
 
             ptr = memory.device_ctypes_pointer
 
-            launch_kernel(
-                function.handle,  # Kernel
-                1,
-                1,
-                1,  # gx, gy, gz
-                100,
-                1,
-                1,  # bx, by, bz
-                0,  # dynamic shared mem
-                stream.handle,  # stream
-                [ptr],
+            config = LaunchConfig(
+                grid=(1, 1, 1),
+                block=(100, 1, 1),
+                shmem_size=0,
+                cooperative_launch=False,
             )
+            launch(stream, config, function.kernel, ptr)
 
             device_to_host(array, memory, sizeof(array), stream=stream)
         for i, v in enumerate(array):

numba_cuda/numba/cuda/tests/cudadrv/test_linker.py

@@ -10,11 +10,12 @@
     skip_if_nvjitlink_missing,
 )
 from numba.cuda.testing import CUDATestCase, test_data_dir
-from numba.cuda.cudadrv.driver import CudaAPIError, _Linker, LinkerError
+from numba.cuda.cudadrv.driver import _Linker, LinkerError
 from numba.cuda import require_context
 from numba import cuda
 from numba.cuda import void, float64, int64, int32, float32
 from numba.cuda.typing.typeof import typeof
+from cuda.core.experimental._utils.cuda_utils import CUDAError
 
 CONST1D = np.arange(10, dtype=np.float64)
 
@@ -308,10 +309,8 @@ def test_max_threads_exceeded(self):
         max_threads = compiled.get_max_threads_per_block()
         nelem = max_threads + 1
         ary = np.empty(nelem, dtype=np.int32)
-        try:
+        with self.assertRaisesRegex(CUDAError, "CUDA_ERROR_INVALID_VALUE"):
             compiled[1, nelem](ary)
-        except CudaAPIError as e:
-            self.assertIn("cuLaunchKernel", e.msg)
 
     def test_get_local_mem_per_thread(self):
         sig = void(int32[::1], int32[::1], typeof(np.int32))

numba_cuda/numba/cuda/tests/cudadrv/test_module_callbacks.py

@@ -13,11 +13,10 @@
     CUDATestCase,
     skip_on_cudasim,
 )
+from cuda.core.experimental import ObjectCode
 
 if not config.ENABLE_CUDASIM:
-    from cuda.bindings.driver import cuModuleGetGlobal, cuMemcpyHtoD
-
-    from cuda.bindings.driver import CUmodule as cu_module_type
+    from cuda.bindings.driver import cuLibraryGetGlobal, cuMemcpyHtoD
 
 
 def wipe_all_modules_in_context():
@@ -31,22 +30,22 @@ def wipe_all_modules_in_context():
     ctx.reset()
 
 
-def get_hashable_handle_value(handle):
-    return handle
+def get_hashable_handle_value(object_code):
+    return object_code.handle
 
 
 @skip_on_cudasim("Module loading not implemented in the simulator")
 class TestModuleCallbacksBasic(CUDATestCase):
     def test_basic(self):
         counter = 0
 
-        def setup(handle):
-            self.assertTrue(isinstance(handle, cu_module_type))
+        def setup(object_code):
+            self.assertIsInstance(object_code, ObjectCode)
             nonlocal counter
             counter += 1
 
-        def teardown(handle):
-            self.assertTrue(isinstance(handle, cu_module_type))
+        def teardown(object_code):
+            self.assertIsInstance(object_code, ObjectCode)
             nonlocal counter
             counter -= 1
 
@@ -183,10 +182,10 @@ def setUp(self):
 }
 """
 
-        def set_forty_two(handle):
+        def set_forty_two(object_code):
             # Initialize 42 to global variable `num`
-            res, dptr, size = cuModuleGetGlobal(
-                get_hashable_handle_value(handle), "num".encode()
+            res, dptr, size = cuLibraryGetGlobal(
+                get_hashable_handle_value(object_code), b"num"
             )
 
             arr = np.array([42], np.int32)