numba-cuda 0.8.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

numba_cuda/numba/cuda/cudadrv/dummyarray.py

@@ -20,7 +20,7 @@ attempt_nocopy_reshape = ctypes.CFUNCTYPE(
     np.ctypeslib.ndpointer(np.ctypeslib.c_intp, ndim=1),  # newstrides
     ctypes.c_long,  # itemsize
     ctypes.c_int,  # is_f_order
-)(_helperlib.c_helpers['attempt_nocopy_reshape'])
+)(_helperlib.c_helpers["attempt_nocopy_reshape"])
 
 
 class Dim(object):
@@ -37,7 +37,8 @@ class Dim(object):
     stride:
         item stride
     """
-    __slots__ = 'start', 'stop', 'size', 'stride', 'single'
+
+    __slots__ = "start", "stop", "size", "stride", "single"
 
     def __init__(self, start, stop, size, stride, single):
         self.start = start
@@ -58,15 +59,11 @@ class Dim(object):
             else:
                 size = _compute_size(start, stop, stride)
             ret = Dim(
-                start=start,
-                stop=stop,
-                size=size,
-                stride=stride,
-                single=False
+                start=start, stop=stop, size=size, stride=stride, single=False
             )
             return ret
         else:
-            sliced = self[item:item + 1] if item != -1 else self[-1:]
+            sliced = self[item : item + 1] if item != -1 else self[-1:]
             if sliced.size != 1:
                 raise IndexError
             return Dim(
@@ -85,8 +82,13 @@ class Dim(object):
         return strfmt % (self.start, self.stop, self.size, self.stride)
 
     def normalize(self, base):
-        return Dim(start=self.start - base, stop=self.stop - base,
-                   size=self.size, stride=self.stride, single=self.single)
+        return Dim(
+            start=self.start - base,
+            stop=self.stop - base,
+            size=self.size,
+            stride=self.stride,
+            single=self.single,
+        )
 
     def copy(self, start=None, stop=None, size=None, stride=None, single=None):
         if start is None:
@@ -143,14 +145,16 @@ class Array(object):
     extent: (start, end)
         start and end offset containing the memory region
     """
+
     is_array = True
 
     @classmethod
     def from_desc(cls, offset, shape, strides, itemsize):
         dims = []
         for ashape, astride in zip(shape, strides):
-            dim = Dim(offset, offset + ashape * astride, ashape, astride,
-                      single=False)
+            dim = Dim(
+                offset, offset + ashape * astride, ashape, astride, single=False
+            )
             dims.append(dim)
             offset = 0  # offset only applies to first dimension
         return cls(dims, itemsize)
@@ -173,23 +177,23 @@ class Array(object):
 
         # Records have no dims, and we can treat them as contiguous
         if not self.dims:
-            return {'C_CONTIGUOUS': True, 'F_CONTIGUOUS': True}
+            return {"C_CONTIGUOUS": True, "F_CONTIGUOUS": True}
 
         # If this is a broadcast array then it is not contiguous
         if any([dim.stride == 0 for dim in self.dims]):
-            return {'C_CONTIGUOUS': False, 'F_CONTIGUOUS': False}
+            return {"C_CONTIGUOUS": False, "F_CONTIGUOUS": False}
 
-        flags = {'C_CONTIGUOUS': True, 'F_CONTIGUOUS': True}
+        flags = {"C_CONTIGUOUS": True, "F_CONTIGUOUS": True}
 
         # Check C contiguity
         sd = self.itemsize
         for dim in reversed(self.dims):
             if dim.size == 0:
                 # Contiguous by definition
-                return {'C_CONTIGUOUS': True, 'F_CONTIGUOUS': True}
+                return {"C_CONTIGUOUS": True, "F_CONTIGUOUS": True}
             if dim.size != 1:
                 if dim.stride != sd:
-                    flags['C_CONTIGUOUS'] = False
+                    flags["C_CONTIGUOUS"] = False
                 sd *= dim.size
 
         # Check F contiguity
@@ -197,7 +201,7 @@ class Array(object):
         for dim in self.dims:
             if dim.size != 1:
                 if dim.stride != sd:
-                    flags['F_CONTIGUOUS'] = False
+                    flags["F_CONTIGUOUS"] = False
                     return flags
                 sd *= dim.size
 
@@ -208,11 +212,11 @@ class Array(object):
         lastidx = [s - 1 for s in self.shape]
         start = compute_index(firstidx, self.dims)
         stop = compute_index(lastidx, self.dims) + self.itemsize
-        stop = max(stop, start)   # ensure positive extent
+        stop = max(stop, start)  # ensure positive extent
         return Extent(start, stop)
 
     def __repr__(self):
-        return '<Array dims=%s itemsize=%s>' % (self.dims, self.itemsize)
+        return "<Array dims=%s itemsize=%s>" % (self.dims, self.itemsize)
 
     def __getitem__(self, item):
         if not isinstance(item, tuple):
@@ -240,15 +244,14 @@ class Array(object):
 
     @property
     def is_c_contig(self):
-        return self.flags['C_CONTIGUOUS']
+        return self.flags["C_CONTIGUOUS"]
 
     @property
     def is_f_contig(self):
-        return self.flags['F_CONTIGUOUS']
+        return self.flags["F_CONTIGUOUS"]
 
     def iter_contiguous_extent(self):
-        """ Generates extents
-        """
+        """Generates extents"""
         if self.is_c_contig or self.is_f_contig:
             yield self.extent
         else:
@@ -279,11 +282,11 @@ class Array(object):
         if newdims == self.shape:
             return self, None
 
-        order = kws.pop('order', 'C')
+        order = kws.pop("order", "C")
         if kws:
-            raise TypeError('unknown keyword arguments %s' % kws.keys())
-        if order not in 'CFA':
-            raise ValueError('order not C|F|A')
+            raise TypeError("unknown keyword arguments %s" % kws.keys())
+        if order not in "CFA":
+            raise ValueError("order not C|F|A")
 
         # check for exactly one instance of -1 in newdims
         # https://github.com/numpy/numpy/blob/623bc1fae1d47df24e7f1e29321d0c0ba2771ce0/numpy/core/src/multiarray/shape.c#L470-L515   # noqa: E501
@@ -301,25 +304,28 @@ class Array(object):
         # compute the missing dimension
         if unknownidx >= 0:
             if knownsize == 0 or self.size % knownsize != 0:
-                raise ValueError("cannot infer valid shape "
-                                 "for unknown dimension")
+                raise ValueError(
+                    "cannot infer valid shape for unknown dimension"
+                )
             else:
-                newdims = newdims[0:unknownidx] \
-                    + (self.size // knownsize,) \
-                    + newdims[unknownidx + 1:]
+                newdims = (
+                    newdims[0:unknownidx]
+                    + (self.size // knownsize,)
+                    + newdims[unknownidx + 1 :]
+                )
 
         newsize = functools.reduce(operator.mul, newdims, 1)
 
-        if order == 'A':
-            order = 'F' if self.is_f_contig else 'C'
+        if order == "A":
+            order = "F" if self.is_f_contig else "C"
 
         if newsize != self.size:
             raise ValueError("reshape changes the size of the array")
 
         if self.is_c_contig or self.is_f_contig:
-            if order == 'C':
+            if order == "C":
                 newstrides = list(iter_strides_c_contig(self, newdims))
-            elif order == 'F':
+            elif order == "F":
                 newstrides = list(iter_strides_f_contig(self, newdims))
             else:
                 raise AssertionError("unreachable")
@@ -340,12 +346,16 @@ class Array(object):
                 newdims,
                 newstrides,
                 self.itemsize,
-                order == 'F',
+                order == "F",
             ):
-                raise NotImplementedError('reshape would require copy')
+                raise NotImplementedError("reshape would require copy")
 
-        ret = self.from_desc(self.extent.begin, shape=newdims,
-                             strides=newstrides, itemsize=self.itemsize)
+        ret = self.from_desc(
+            self.extent.begin,
+            shape=newdims,
+            strides=newstrides,
+            itemsize=self.itemsize,
+        )
 
         return ret, list(self.iter_contiguous_extent())
 
@@ -377,16 +387,21 @@ class Array(object):
         )
         return newarr, list(self.iter_contiguous_extent())
 
-    def ravel(self, order='C'):
-        if order not in 'CFA':
-            raise ValueError('order not C|F|A')
+    def ravel(self, order="C"):
+        if order not in "CFA":
+            raise ValueError("order not C|F|A")
 
-        if (order in 'CA' and self.is_c_contig
-                or order in 'FA' and self.is_f_contig):
+        if (
+            order in "CA"
+            and self.is_c_contig
+            or order in "FA"
+            and self.is_f_contig
+        ):
             newshape = (self.size,)
             newstrides = (self.itemsize,)
-            arr = self.from_desc(self.extent.begin, newshape, newstrides,
-                                 self.itemsize)
+            arr = self.from_desc(
+                self.extent.begin, newshape, newstrides, self.itemsize
+            )
             return arr, list(self.iter_contiguous_extent())
 
         else:
@@ -394,8 +409,7 @@ class Array(object):
 
 
 def iter_strides_f_contig(arr, shape=None):
-    """yields the f-contiguous strides
-    """
+    """yields the f-contiguous strides"""
     shape = arr.shape if shape is None else shape
     itemsize = arr.itemsize
     yield itemsize
@@ -406,8 +420,7 @@ def iter_strides_f_contig(arr, shape=None):
 
 
 def iter_strides_c_contig(arr, shape=None):
-    """yields the c-contiguous strides
-    """
+    """yields the c-contiguous strides"""
     shape = arr.shape if shape is None else shape
     itemsize = arr.itemsize
 
@@ -438,8 +451,7 @@ def is_element_indexing(item, ndim):
 
 
 def _compute_size(start, stop, step):
-    """Algorithm adapted from cpython rangeobject.c
-    """
+    """Algorithm adapted from cpython rangeobject.c"""
     if step > 0:
         lo = start
         hi = stop

numba_cuda/numba/cuda/cudadrv/enums.py

@@ -140,7 +140,7 @@ CU_CTX_USER_COREDUMP_ENABLE = 0x40
 # Force synchronous blocking on cudaMemcpy/cudaMemset
 CU_CTX_SYNC_MEMOPS = 0x80
 
-CU_CTX_FLAGS_MASK = 0xff
+CU_CTX_FLAGS_MASK = 0xFF
 
 
 # DEFINES

numba_cuda/numba/cuda/cudadrv/error.py

@@ -12,7 +12,7 @@ class CudaSupportError(ImportError):
 
 class NvvmError(Exception):
     def __str__(self):
-        return '\n'.join(map(str, self.args))
+        return "\n".join(map(str, self.args))
 
 
 class NvvmSupportError(ImportError):
@@ -25,12 +25,16 @@ class NvvmWarning(Warning):
 
 class NvrtcError(Exception):
     def __str__(self):
-        return '\n'.join(map(str, self.args))
+        return "\n".join(map(str, self.args))
 
 
 class NvrtcCompilationError(NvrtcError):
     pass
 
 
+class NvrtcBuiltinOperationFailure(NvrtcError):
+    pass
+
+
 class NvrtcSupportError(ImportError):
     pass

numba_cuda/numba/cuda/cudadrv/libs.py

@@ -21,25 +21,25 @@ from numba.cuda.cudadrv.error import CudaSupportError
 from numba.core import config
 
 
-if sys.platform == 'win32':
-    _dllnamepattern = '%s.dll'
-    _staticnamepattern = '%s.lib'
-elif sys.platform == 'darwin':
-    _dllnamepattern = 'lib%s.dylib'
-    _staticnamepattern = 'lib%s.a'
+if sys.platform == "win32":
+    _dllnamepattern = "%s.dll"
+    _staticnamepattern = "%s.lib"
+elif sys.platform == "darwin":
+    _dllnamepattern = "lib%s.dylib"
+    _staticnamepattern = "lib%s.a"
 else:
-    _dllnamepattern = 'lib%s.so'
-    _staticnamepattern = 'lib%s.a'
+    _dllnamepattern = "lib%s.so"
+    _staticnamepattern = "lib%s.a"
 
 
 def get_libdevice():
     d = get_cuda_paths()
-    paths = d['libdevice'].info
+    paths = d["libdevice"].info
     return paths
 
 
 def open_libdevice():
-    with open(get_libdevice(), 'rb') as bcfile:
+    with open(get_libdevice(), "rb") as bcfile:
         return bcfile.read()
 
 
@@ -50,10 +50,10 @@ def get_cudalib(lib, static=False):
     'libnvvm.so' for 'nvvm') so that we may attempt to load it using the system
     loader's search mechanism.
     """
-    if lib == 'nvvm':
-        return get_cuda_paths()['nvvm'].info or _dllnamepattern % 'nvvm'
+    if lib in {"nvrtc", "nvvm"}:
+        return get_cuda_paths()[lib].info or _dllnamepattern % lib
     else:
-        dir_type = 'static_cudalib_dir' if static else 'cudalib_dir'
+        dir_type = "static_cudalib_dir" if static else "cudalib_dir"
         libdir = get_cuda_paths()[dir_type].info
 
     candidates = find_lib(lib, libdir, static=static)
@@ -68,7 +68,7 @@ def get_cuda_include_dir():
     configuration.
     """
 
-    return get_cuda_paths()['include_dir'].info
+    return get_cuda_paths()["include_dir"].info
 
 
 def check_cuda_include_dir(path):
@@ -86,39 +86,40 @@ def open_cudalib(lib):
 
 def check_static_lib(path):
     if not os.path.isfile(path):
-        raise FileNotFoundError(f'{path} not found')
+        raise FileNotFoundError(f"{path} not found")
 
 
 def _get_source_variable(lib, static=False):
-    if lib == 'nvvm':
-        return get_cuda_paths()['nvvm'].by
-    elif lib == 'libdevice':
-        return get_cuda_paths()['libdevice'].by
-    elif lib == 'include_dir':
-        return get_cuda_paths()['include_dir'].by
+    if lib == "nvvm":
+        return get_cuda_paths()["nvvm"].by
+    elif lib == "nvrtc":
+        return get_cuda_paths()["nvrtc"].by
+    elif lib == "libdevice":
+        return get_cuda_paths()["libdevice"].by
+    elif lib == "include_dir":
+        return get_cuda_paths()["include_dir"].by
     else:
-        dir_type = 'static_cudalib_dir' if static else 'cudalib_dir'
+        dir_type = "static_cudalib_dir" if static else "cudalib_dir"
         return get_cuda_paths()[dir_type].by
 
 
 def test():
-    """Test library lookup.  Path info is printed to stdout.
-    """
+    """Test library lookup.  Path info is printed to stdout."""
     failed = False
 
     # Check for the driver
     try:
         dlloader, candidates = locate_driver_and_loader()
-        print('Finding driver from candidates:')
+        print("Finding driver from candidates:")
         for location in candidates:
-            print(f'\t{location}')
-        print(f'Using loader {dlloader}')
-        print('\tTrying to load driver', end='...')
+            print(f"\t{location}")
+        print(f"Using loader {dlloader}")
+        print("\tTrying to load driver", end="...")
         dll, path = load_driver(dlloader, candidates)
-        print('\tok')
-        print(f'\t\tLoaded from {path}')
+        print("\tok")
+        print(f"\t\tLoaded from {path}")
     except CudaSupportError as e:
-        print(f'\tERROR: failed to open driver: {e}')
+        print(f"\tERROR: failed to open driver: {e}")
         failed = True
 
     # Find the absolute location of the driver on Linux. Various driver-related
@@ -127,9 +128,9 @@ def test():
     # Providing the absolute location of the driver indicates its version
     # number in the soname (e.g. "libcuda.so.530.30.02"), which can be used to
     # look up whether the driver was intended for "native" Linux.
-    if sys.platform == 'linux' and not failed:
+    if sys.platform == "linux" and not failed:
         pid = os.getpid()
-        mapsfile = os.path.join(os.path.sep, 'proc', f'{pid}', 'maps')
+        mapsfile = os.path.join(os.path.sep, "proc", f"{pid}", "maps")
         try:
             with open(mapsfile) as f:
                 maps = f.read()
@@ -140,58 +141,61 @@ def test():
             # It's helpful to report that this went wrong to the user, but we
             # don't set failed to True because this doesn't have any connection
             # to actual CUDA functionality.
-            print(f'\tERROR: Could not open {mapsfile} to determine absolute '
-                  'path to libcuda.so')
+            print(
+                f"\tERROR: Could not open {mapsfile} to determine absolute "
+                "path to libcuda.so"
+            )
         else:
             # In this case we could read the maps, so we can report the
             # relevant ones to the user
-            locations = set(s for s in maps.split() if 'libcuda.so' in s)
-            print('\tMapped libcuda.so paths:')
+            locations = set(s for s in maps.split() if "libcuda.so" in s)
+            print("\tMapped libcuda.so paths:")
             for location in locations:
-                print(f'\t\t{location}')
+                print(f"\t\t{location}")
 
     # Checks for dynamic libraries
-    libs = 'nvvm nvrtc cudart'.split()
+    libs = "nvvm nvrtc cudart".split()
     for lib in libs:
         path = get_cudalib(lib)
-        print('Finding {} from {}'.format(lib, _get_source_variable(lib)))
-        print('\tLocated at', path)
+        print("Finding {} from {}".format(lib, _get_source_variable(lib)))
+        print("\tLocated at", path)
 
         try:
-            print('\tTrying to open library', end='...')
+            print("\tTrying to open library", end="...")
             open_cudalib(lib)
-            print('\tok')
+            print("\tok")
         except OSError as e:
-            print('\tERROR: failed to open %s:\n%s' % (lib, e))
+            print("\tERROR: failed to open %s:\n%s" % (lib, e))
             failed = True
 
     # Check for cudadevrt (the only static library)
-    lib = 'cudadevrt'
+    lib = "cudadevrt"
     path = get_cudalib(lib, static=True)
-    print('Finding {} from {}'.format(lib, _get_source_variable(lib,
-                                                                static=True)))
-    print('\tLocated at', path)
+    print(
+        "Finding {} from {}".format(lib, _get_source_variable(lib, static=True))
+    )
+    print("\tLocated at", path)
 
     try:
-        print('\tChecking library', end='...')
+        print("\tChecking library", end="...")
         check_static_lib(path)
-        print('\tok')
+        print("\tok")
     except FileNotFoundError as e:
-        print('\tERROR: failed to find %s:\n%s' % (lib, e))
+        print("\tERROR: failed to find %s:\n%s" % (lib, e))
         failed = True
 
     # Check for libdevice
-    where = _get_source_variable('libdevice')
-    print(f'Finding libdevice from {where}')
+    where = _get_source_variable("libdevice")
+    print(f"Finding libdevice from {where}")
     path = get_libdevice()
-    print('\tLocated at', path)
+    print("\tLocated at", path)
 
     try:
-        print('\tChecking library', end='...')
+        print("\tChecking library", end="...")
         check_static_lib(path)
-        print('\tok')
+        print("\tok")
     except FileNotFoundError as e:
-        print('\tERROR: failed to find %s:\n%s' % (lib, e))
+        print("\tERROR: failed to find %s:\n%s" % (lib, e))
         failed = True
 
     # Check cuda include paths
@@ -199,16 +203,16 @@ def test():
     print("Include directory configuration variable:")
     print(f"\tCUDA_INCLUDE_PATH={config.CUDA_INCLUDE_PATH}")
 
-    where = _get_source_variable('include_dir')
-    print(f'Finding include directory from {where}')
+    where = _get_source_variable("include_dir")
+    print(f"Finding include directory from {where}")
     include = get_cuda_include_dir()
-    print('\tLocated at', include)
+    print("\tLocated at", include)
     try:
-        print('\tChecking include directory', end='...')
+        print("\tChecking include directory", end="...")
         check_cuda_include_dir(include)
-        print('\tok')
+        print("\tok")
     except FileNotFoundError as e:
-        print('\tERROR: failed to find cuda include directory:\n%s' % e)
+        print("\tERROR: failed to find cuda include directory:\n%s" % e)
         failed = True
 
     return not failed

numba_cuda/numba/cuda/cudadrv/linkable_code.py

@@ -7,11 +7,26 @@ class LinkableCode:
     :param data: A buffer containing the data to link.
     :param name: The name of the file to be referenced in any compilation or
                  linking errors that may be produced.
+    :param setup_callback: A function called prior to the launch of a kernel
+                           contained within a module that has this code object
+                           linked into it.
+    :param teardown_callback: A function called just prior to the unloading of
+                              a module that has this code object linked into
+                              it.
     """
 
-    def __init__(self, data, name=None):
+    def __init__(
+        self, data, name=None, setup_callback=None, teardown_callback=None
+    ):
+        if setup_callback and not callable(setup_callback):
+            raise TypeError("setup_callback must be callable")
+        if teardown_callback and not callable(teardown_callback):
+            raise TypeError("teardown_callback must be callable")
+
         self.data = data
         self._name = name
+        self.setup_callback = setup_callback
+        self.teardown_callback = teardown_callback
 
     @property
     def name(self):

numba_cuda/numba/cuda/cudadrv/mappings.py

@@ -1,24 +1,26 @@
 from numba import config
 from . import enums
+
 if config.CUDA_USE_NVIDIA_BINDING:
     from cuda import cuda
+
     jitty = cuda.CUjitInputType
     FILE_EXTENSION_MAP = {
-        'o': jitty.CU_JIT_INPUT_OBJECT,
-        'ptx': jitty.CU_JIT_INPUT_PTX,
-        'a': jitty.CU_JIT_INPUT_LIBRARY,
-        'lib': jitty.CU_JIT_INPUT_LIBRARY,
-        'cubin': jitty.CU_JIT_INPUT_CUBIN,
-        'fatbin': jitty.CU_JIT_INPUT_FATBINARY,
-        'ltoir': jitty.CU_JIT_INPUT_NVVM,
+        "o": jitty.CU_JIT_INPUT_OBJECT,
+        "ptx": jitty.CU_JIT_INPUT_PTX,
+        "a": jitty.CU_JIT_INPUT_LIBRARY,
+        "lib": jitty.CU_JIT_INPUT_LIBRARY,
+        "cubin": jitty.CU_JIT_INPUT_CUBIN,
+        "fatbin": jitty.CU_JIT_INPUT_FATBINARY,
+        "ltoir": jitty.CU_JIT_INPUT_NVVM,
     }
 else:
     FILE_EXTENSION_MAP = {
-        'o': enums.CU_JIT_INPUT_OBJECT,
-        'ptx': enums.CU_JIT_INPUT_PTX,
-        'a': enums.CU_JIT_INPUT_LIBRARY,
-        'lib': enums.CU_JIT_INPUT_LIBRARY,
-        'cubin': enums.CU_JIT_INPUT_CUBIN,
-        'fatbin': enums.CU_JIT_INPUT_FATBINARY,
-        'ltoir': enums.CU_JIT_INPUT_NVVM,
+        "o": enums.CU_JIT_INPUT_OBJECT,
+        "ptx": enums.CU_JIT_INPUT_PTX,
+        "a": enums.CU_JIT_INPUT_LIBRARY,
+        "lib": enums.CU_JIT_INPUT_LIBRARY,
+        "cubin": enums.CU_JIT_INPUT_CUBIN,
+        "fatbin": enums.CU_JIT_INPUT_FATBINARY,
+        "ltoir": enums.CU_JIT_INPUT_NVVM,
     }