numba-cuda 0.3.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

numba_cuda/numba/cuda/runtime/nrt.py

@@ -0,0 +1,318 @@
+import ctypes
+import os
+from functools import wraps
+import numpy as np
+
+from numba import cuda, config
+from numba.core.runtime.nrt import _nrt_mstats
+from numba.cuda.cudadrv.driver import Linker, driver, launch_kernel
+from numba.cuda.cudadrv import devices
+from numba.cuda.api import get_current_device
+from numba.cuda.utils import _readenv
+
+
+# Check environment variable or config for NRT statistics enablement
+NRT_STATS = (
+    _readenv("NUMBA_CUDA_NRT_STATS", bool, False) or
+    getattr(config, "NUMBA_CUDA_NRT_STATS", False)
+)
+if not hasattr(config, "NUMBA_CUDA_NRT_STATS"):
+    config.CUDA_NRT_STATS = NRT_STATS
+
+
+# Check environment variable or config for NRT enablement
+ENABLE_NRT = (
+    _readenv("NUMBA_CUDA_ENABLE_NRT", bool, False) or
+    getattr(config, "NUMBA_CUDA_ENABLE_NRT", False)
+)
+if not hasattr(config, "NUMBA_CUDA_ENABLE_NRT"):
+    config.CUDA_ENABLE_NRT = ENABLE_NRT
+
+
+# Protect method to ensure NRT memory allocation and initialization
+def _alloc_init_guard(method):
+    """
+    Ensure NRT memory allocation and initialization before running the method
+    """
+    @wraps(method)
+    def wrapper(self, *args, **kwargs):
+        self.ensure_allocated()
+        self.ensure_initialized()
+        return method(self, *args, **kwargs)
+    return wrapper
+
+
+class _Runtime:
+    """Singleton class for Numba CUDA runtime"""
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super(_Runtime, cls).__new__(cls, *args, **kwargs)
+        return cls._instance
+
+    def __init__(self):
+        """Initialize memsys module and variable"""
+        self._memsys_module = None
+        self._memsys = None
+        self._initialized = False
+
+    def _compile_memsys_module(self):
+        """
+        Compile memsys.cu and create a module from it in the current context
+        """
+        # Define the path for memsys.cu
+        memsys_mod = os.path.join(
+            os.path.dirname(os.path.abspath(__file__)),
+            "memsys.cu"
+        )
+        cc = get_current_device().compute_capability
+
+        # Create a new linker instance and add the cu file
+        linker = Linker.new(cc=cc)
+        linker.add_cu_file(memsys_mod)
+
+        # Complete the linker and create a module from it
+        cubin = linker.complete()
+        ctx = devices.get_context()
+        module = ctx.create_module_image(cubin)
+
+        # Set the memsys module
+        self._memsys_module = module
+
+    def ensure_allocated(self, stream=None):
+        """
+        If memsys is not allocated, allocate it; otherwise, perform a no-op
+        """
+        if self._memsys is not None:
+            return
+
+        # Allocate the memsys
+        self.allocate(stream)
+
+    def allocate(self, stream=None):
+        """
+        Allocate memsys on global memory
+        """
+        from numba.cuda import device_array
+
+        # Check if memsys module is defined
+        if self._memsys_module is None:
+            # Compile the memsys module if not defined
+            self._compile_memsys_module()
+
+        # Allocate space for NRT_MemSys
+        ptr, nbytes = self._memsys_module.get_global_symbol("memsys_size")
+        memsys_size = ctypes.c_uint64()
+        driver.cuMemcpyDtoH(ctypes.addressof(memsys_size),
+                            ptr.device_ctypes_pointer, nbytes)
+        self._memsys = device_array(
+            (memsys_size.value,), dtype="i1", stream=stream)
+        self.set_memsys_to_module(self._memsys_module, stream=stream)
+
+    def _single_thread_launch(self, module, stream, name, params=()):
+        """
+        Launch the specified kernel with only 1 thread
+        """
+        if stream is None:
+            stream = cuda.default_stream()
+
+        func = module.get_function(name)
+        launch_kernel(
+            func.handle,
+            1, 1, 1,
+            1, 1, 1,
+            0,
+            stream.handle,
+            params,
+            cooperative=False
+        )
+
+    def ensure_initialized(self, stream=None):
+        """
+        If memsys is not initialized, initialize memsys
+        """
+        if self._initialized:
+            return
+
+        # Initialize the memsys
+        self.initialize(stream)
+
+    def initialize(self, stream=None):
+        """
+        Launch memsys initialization kernel
+        """
+        self.ensure_allocated()
+
+        self._single_thread_launch(
+            self._memsys_module, stream, "NRT_MemSys_init")
+        self._initialized = True
+
+        if config.CUDA_NRT_STATS:
+            self.memsys_enable_stats()
+
+    @_alloc_init_guard
+    def memsys_enable_stats(self, stream=None):
+        """
+        Enable memsys statistics
+        """
+        self._single_thread_launch(
+            self._memsys_module, stream, "NRT_MemSys_enable_stats")
+
+    @_alloc_init_guard
+    def memsys_disable_stats(self, stream=None):
+        """
+        Disable memsys statistics
+        """
+        self._single_thread_launch(
+            self._memsys_module, stream, "NRT_MemSys_disable_stats")
+
+    @_alloc_init_guard
+    def memsys_stats_enabled(self, stream=None):
+        """
+        Return a boolean indicating whether memsys is enabled. Synchronizes
+        context
+        """
+        enabled_ar = cuda.managed_array(1, np.uint8)
+
+        self._single_thread_launch(
+            self._memsys_module,
+            stream,
+            "NRT_MemSys_stats_enabled",
+            (enabled_ar.device_ctypes_pointer,)
+        )
+
+        cuda.synchronize()
+        return bool(enabled_ar[0])
+
+    @_alloc_init_guard
+    def _copy_memsys_to_host(self, stream):
+        """
+        Copy all statistics of memsys to the host
+        """
+        dt = np.dtype([
+            ('alloc', np.uint64),
+            ('free', np.uint64),
+            ('mi_alloc', np.uint64),
+            ('mi_free', np.uint64)
+        ])
+
+        stats_for_read = cuda.managed_array(1, dt)
+
+        self._single_thread_launch(
+            self._memsys_module,
+            stream,
+            "NRT_MemSys_read",
+            [stats_for_read.device_ctypes_pointer]
+        )
+        cuda.synchronize()
+
+        return stats_for_read[0]
+
+    @_alloc_init_guard
+    def get_allocation_stats(self, stream=None):
+        """
+        Get the allocation statistics
+        """
+        enabled = self.memsys_stats_enabled(stream)
+        if not enabled:
+            raise RuntimeError("NRT stats are disabled.")
+        memsys = self._copy_memsys_to_host(stream)
+        return _nrt_mstats(
+            alloc=memsys["alloc"],
+            free=memsys["free"],
+            mi_alloc=memsys["mi_alloc"],
+            mi_free=memsys["mi_free"]
+        )
+
+    @_alloc_init_guard
+    def _get_single_stat(self, stat, stream=None):
+        """
+        Get a single stat from the memsys
+        """
+        got = cuda.managed_array(1, np.uint64)
+        self._single_thread_launch(
+            self._memsys_module,
+            stream,
+            f"NRT_MemSys_read_{stat}",
+            [got.device_ctypes_pointer]
+        )
+
+        cuda.synchronize()
+        return got[0]
+
+    @_alloc_init_guard
+    def memsys_get_stats_alloc(self, stream=None):
+        """
+        Get the allocation statistic
+        """
+        enabled = self.memsys_stats_enabled(stream)
+        if not enabled:
+            raise RuntimeError("NRT stats are disabled.")
+
+        return self._get_single_stat("alloc")
+
+    @_alloc_init_guard
+    def memsys_get_stats_free(self, stream=None):
+        """
+        Get the free statistic
+        """
+        enabled = self.memsys_stats_enabled(stream)
+        if not enabled:
+            raise RuntimeError("NRT stats are disabled.")
+
+        return self._get_single_stat("free")
+
+    @_alloc_init_guard
+    def memsys_get_stats_mi_alloc(self, stream=None):
+        """
+        Get the mi alloc statistic
+        """
+        enabled = self.memsys_stats_enabled(stream)
+        if not enabled:
+            raise RuntimeError("NRT stats are disabled.")
+
+        return self._get_single_stat("mi_alloc")
+
+    @_alloc_init_guard
+    def memsys_get_stats_mi_free(self, stream=None):
+        """
+        Get the mi free statistic
+        """
+        enabled = self.memsys_stats_enabled(stream)
+        if not enabled:
+            raise RuntimeError("NRT stats are disabled.")
+
+        return self._get_single_stat("mi_free")
+
+    def set_memsys_to_module(self, module, stream=None):
+        """
+        Set the memsys module. The module must contain `NRT_MemSys_set` kernel,
+        and declare a pointer to NRT_MemSys structure.
+        """
+        if self._memsys is None:
+            raise RuntimeError(
+                "Please allocate NRT Memsys first before setting to module.")
+
+        self._single_thread_launch(
+            module,
+            stream,
+            "NRT_MemSys_set",
+            [self._memsys.device_ctypes_pointer,]
+        )
+
+    @_alloc_init_guard
+    def print_memsys(self, stream=None):
+        """
+        Print the current statistics of memsys, for debugging purposes
+        """
+        cuda.synchronize()
+        self._single_thread_launch(
+            self._memsys_module,
+            stream,
+            "NRT_MemSys_print"
+        )
+
+
+# Create an instance of the runtime
+rtsys = _Runtime()

numba_cuda/numba/cuda/testing.py

@@ -115,12 +115,22 @@ def skip_on_arm(reason):
 def skip_if_cuda_includes_missing(fn):
     # Skip when cuda.h is not available - generally this should indicate
     # whether the CUDA includes are available or not
-    cuda_h = os.path.join(config.CUDA_INCLUDE_PATH, 'cuda.h')
+    cuda_include_path = libs.get_cuda_include_dir()
+    cuda_h = os.path.join(cuda_include_path, 'cuda.h')
     cuda_h_file = (os.path.exists(cuda_h) and os.path.isfile(cuda_h))
     reason = 'CUDA include dir not available on this system'
     return unittest.skipUnless(cuda_h_file, reason)(fn)
 
 
+def skip_if_curand_kernel_missing(fn):
+    cuda_include_path = libs.get_cuda_include_dir()
+    curand_kernel_h = os.path.join(cuda_include_path, 'curand_kernel.h')
+    curand_kernel_h_file = (os.path.exists(curand_kernel_h) and
+                            os.path.isfile(curand_kernel_h))
+    reason = 'curand_kernel.h not available on this system'
+    return unittest.skipUnless(curand_kernel_h_file, reason)(fn)
+
+
 def skip_if_mvc_enabled(reason):
     """Skip a test if Minor Version Compatibility is enabled"""
     return unittest.skipIf(config.CUDA_ENABLE_MINOR_VERSION_COMPATIBILITY,

numba_cuda/numba/cuda/tests/__init__.py

@@ -49,6 +49,7 @@ def load_tests(loader, tests, pattern):
         if gpus and gpus[0].compute_capability >= (2, 0):
             suite.addTests(load_testsuite(loader, join(this_dir, 'cudadrv')))
             suite.addTests(load_testsuite(loader, join(this_dir, 'cudapy')))
+            suite.addTests(load_testsuite(loader, join(this_dir, 'nrt')))
             suite.addTests(load_testsuite(loader, join(this_dir,
                                                        'doc_examples')))
         else:

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

@@ -4,6 +4,7 @@ from numba.cuda.cudadrv import devicearray
 from numba import cuda
 from numba.cuda.testing import unittest, CUDATestCase
 from numba.cuda.testing import skip_on_cudasim
+from numba.tests.support import IS_NUMPY_2
 
 
 class TestCudaNDArray(CUDATestCase):
@@ -456,6 +457,36 @@ class TestCudaNDArray(CUDATestCase):
         dev_array_from_host.copy_to_device(dev_array)
 
 
+class TestArrayMethod(CUDATestCase):
+    """Tests of the __array__() method via np.array"""
+
+    def test_np_array(self):
+        dev_array = cuda.to_device(np.asarray([1.0, 2.0, 3.0]))
+        host_array = np.array(dev_array)
+        np.testing.assert_equal(dev_array.copy_to_host(), host_array)
+
+    def test_np_array_dtype(self):
+        dtype = np.int32
+        dev_array = cuda.to_device(np.asarray([1.0, 2.0, 3.0]))
+        host_array = np.array(dev_array, dtype=dtype)
+        np.testing.assert_equal(
+            host_array,
+            dev_array.copy_to_host().astype(dtype)
+        )
+
+    @unittest.skipUnless(IS_NUMPY_2, "NumPy 1.x does not pass copy kwarg")
+    def test_np_array_copy_false(self):
+        dev_array = cuda.to_device(np.asarray([1.0, 2.0, 3.0]))
+        with self.assertRaisesRegex(ValueError, "`copy=False` is not"):
+            np.array(dev_array, copy=False)
+
+    @unittest.skipUnless(IS_NUMPY_2, "NumPy 1.x does not pass copy kwarg")
+    def test_np_array_copy_true(self):
+        dev_array = cuda.to_device(np.asarray([1.0, 2.0, 3.0]))
+        host_array = np.array(dev_array)
+        np.testing.assert_equal(dev_array.copy_to_host(), host_array)
+
+
 class TestRecarray(CUDATestCase):
     def test_recarray(self):
         # From issue #4111

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

@@ -1,11 +1,14 @@
 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_if_curand_kernel_missing, skip_on_cudasim,
+                                test_data_dir, unittest, CUDATestCase)
+from numba import cuda, jit, float32, int32, types
 from numba.core.errors import TypingError
+from numba.tests.support import skip_unless_cffi
+from types import ModuleType
 
 
 class TestDeviceFunc(CUDATestCase):
@@ -92,7 +95,7 @@ class TestDeviceFunc(CUDATestCase):
         def add(a, b):
             return a + b
 
-        mymod = types.ModuleType(name='mymod')
+        mymod = ModuleType(name='mymod')
         mymod.add = add
         del add
 
@@ -192,31 +195,162 @@ class TestDeviceFunc(CUDATestCase):
 
         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;
+}
+""")
+
+rng_cu = cuda.CUSource("""
+#include <curand_kernel.h>
+
+extern "C" __device__
+int random_number(unsigned int *out, unsigned long long seed)
+{
+  // Initialize state
+  curandStateXORWOW_t state;
+  unsigned long long sequence = 1;
+  unsigned long long offset = 0;
+  curand_init(seed, sequence, offset, &state);
+
+  // Generate one random number
+  *out = curand(&state);
+
+  // Report no exception
+  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)
+
+    @skip_unless_cffi
+    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)
+
+    @skip_if_curand_kernel_missing
+    def test_include_cuda_header(self):
+        sig = types.int32(types.uint64)
+        link = [rng_cu]
+        random_number = cuda.declare_device("random_number", sig, link=link)
+
+        @cuda.jit
+        def kernel(x, seed):
+            x[0] = random_number(seed)
+
+        x = np.zeros(1, dtype=np.uint32)
+        kernel[1, 1](x, 1)
+        np.testing.assert_equal(x[0], 323845807)
+
 
 if __name__ == '__main__':
     unittest.main()

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

@@ -15,16 +15,18 @@ class TestFFI(CUDATestCase):
         import numpy as np
         import os
 
-        # Declaration of the foreign function
-        mul = cuda.declare_device('mul_f32_f32', 'float32(float32, float32)')
-
         # Path to the source containing the foreign function
         # (here assumed to be in a subdirectory called "ffi")
         basedir = os.path.dirname(os.path.abspath(__file__))
         functions_cu = os.path.join(basedir, 'ffi', 'functions.cu')
 
-        # Kernel that links in functions.cu and calls mul
-        @cuda.jit(link=[functions_cu])
+        # Declaration of the foreign function
+        mul = cuda.declare_device('mul_f32_f32', 'float32(float32, float32)',
+                                  link=functions_cu)
+
+        # A kernel that calls mul; functions.cu is linked automatically due to
+        # the call to mul.
+        @cuda.jit
         def multiply_vectors(r, x, y):
             i = cuda.grid(1)
 
@@ -54,14 +56,15 @@ class TestFFI(CUDATestCase):
 
         # magictoken.ex_from_buffer_decl.begin
         signature = 'float32(CPointer(float32), int32)'
-        sum_reduce = cuda.declare_device('sum_reduce', signature)
+        sum_reduce = cuda.declare_device('sum_reduce', signature,
+                                         link=functions_cu)
         # magictoken.ex_from_buffer_decl.end
 
         # magictoken.ex_from_buffer_kernel.begin
         import cffi
         ffi = cffi.FFI()
 
-        @cuda.jit(link=[functions_cu])
+        @cuda.jit
         def reduction_caller(result, array):
             array_ptr = ffi.from_buffer(array)
             result[()] = sum_reduce(array_ptr, len(array))