numba-cuda 0.0.1__py3-none-any.whl → 0.0.13__py3-none-any.whl

numba_cuda/numba/cuda/cudadrv/libs.py

@@ -0,0 +1,176 @@
+"""CUDA Toolkit libraries lookup utilities.
+
+CUDA Toolkit libraries can be available via either:
+
+- the `cuda-nvcc` and `cuda-nvrtc` conda packages for CUDA 12,
+- the `cudatoolkit` conda package for CUDA 11,
+- a user supplied location from CUDA_HOME,
+- a system wide location,
+- package-specific locations (e.g. the Debian NVIDIA packages),
+- or can be discovered by the system loader.
+"""
+
+import os
+import sys
+import ctypes
+
+from numba.misc.findlib import find_lib
+from numba.cuda.cuda_paths import get_cuda_paths
+from numba.cuda.cudadrv.driver import locate_driver_and_loader, load_driver
+from numba.cuda.cudadrv.error import CudaSupportError
+
+
+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'
+
+
+def get_libdevice():
+    d = get_cuda_paths()
+    paths = d['libdevice'].info
+    return paths
+
+
+def open_libdevice():
+    with open(get_libdevice(), 'rb') as bcfile:
+        return bcfile.read()
+
+
+def get_cudalib(lib, static=False):
+    """
+    Find the path of a CUDA library based on a search of known locations. If
+    the search fails, return a generic filename for the library (e.g.
+    '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'
+    else:
+        dir_type = 'static_cudalib_dir' if static else 'cudalib_dir'
+        libdir = get_cuda_paths()[dir_type].info
+
+    candidates = find_lib(lib, libdir, static=static)
+    namepattern = _staticnamepattern if static else _dllnamepattern
+    return max(candidates) if candidates else namepattern % lib
+
+
+def open_cudalib(lib):
+    path = get_cudalib(lib)
+    return ctypes.CDLL(path)
+
+
+def check_static_lib(path):
+    if not os.path.isfile(path):
+        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
+    else:
+        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.
+    """
+    failed = False
+
+    # Check for the driver
+    try:
+        dlloader, candidates = locate_driver_and_loader()
+        print('Finding driver from candidates:')
+        for location in candidates:
+            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}')
+    except CudaSupportError as e:
+        print(f'\tERROR: failed to open driver: {e}')
+        failed = True
+
+    # Find the absolute location of the driver on Linux. Various driver-related
+    # issues have been reported by WSL2 users, and it is almost always due to a
+    # Linux (i.e. not- WSL2) driver being installed in a WSL2 system.
+    # 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:
+        pid = os.getpid()
+        mapsfile = os.path.join(os.path.sep, 'proc', f'{pid}', 'maps')
+        try:
+            with open(mapsfile) as f:
+                maps = f.read()
+        # It's difficult to predict all that might go wrong reading the maps
+        # file - in case various error conditions ensue (the file is not found,
+        # not readable, etc.) we use OSError to hopefully catch any of them.
+        except OSError:
+            # 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')
+        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:')
+            for location in locations:
+                print(f'\t\t{location}')
+
+    # Checks for dynamic libraries
+    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)
+
+        try:
+            print('\tTrying to open library', end='...')
+            open_cudalib(lib)
+            print('\tok')
+        except OSError as e:
+            print('\tERROR: failed to open %s:\n%s' % (lib, e))
+            failed = True
+
+    # Check for cudadevrt (the only static library)
+    lib = 'cudadevrt'
+    path = get_cudalib(lib, static=True)
+    print('Finding {} from {}'.format(lib, _get_source_variable(lib,
+                                                                static=True)))
+    print('\tLocated at', path)
+
+    try:
+        print('\tChecking library', end='...')
+        check_static_lib(path)
+        print('\tok')
+    except FileNotFoundError as 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}')
+    path = get_libdevice()
+    print('\tLocated at', path)
+
+    try:
+        print('\tChecking library', end='...')
+        check_static_lib(path)
+        print('\tok')
+    except FileNotFoundError as e:
+        print('\tERROR: failed to find %s:\n%s' % (lib, e))
+        failed = True
+
+    return not failed

numba_cuda/numba/cuda/cudadrv/ndarray.py

@@ -0,0 +1,20 @@
+from numba.cuda.cudadrv import devices, driver
+from numba.core.registry import cpu_target
+
+
+def _calc_array_sizeof(ndim):
+    """
+    Use the ABI size in the CPU target
+    """
+    ctx = cpu_target.target_context
+    return ctx.calc_array_sizeof(ndim)
+
+
+def ndarray_device_allocate_data(ary):
+    """
+    Allocate gpu data buffer
+    """
+    datasize = driver.host_memory_size(ary)
+    # allocate
+    gpu_data = devices.get_context().memalloc(datasize)
+    return gpu_data

numba_cuda/numba/cuda/cudadrv/nvrtc.py

@@ -0,0 +1,260 @@
+from ctypes import byref, c_char, c_char_p, c_int, c_size_t, c_void_p, POINTER
+from enum import IntEnum
+from numba.core import config
+from numba.cuda.cudadrv.error import (NvrtcError, NvrtcCompilationError,
+                                      NvrtcSupportError)
+
+import functools
+import os
+import threading
+import warnings
+
+# Opaque handle for compilation unit
+nvrtc_program = c_void_p
+
+# Result code
+nvrtc_result = c_int
+
+
+class NvrtcResult(IntEnum):
+    NVRTC_SUCCESS = 0
+    NVRTC_ERROR_OUT_OF_MEMORY = 1
+    NVRTC_ERROR_PROGRAM_CREATION_FAILURE = 2
+    NVRTC_ERROR_INVALID_INPUT = 3
+    NVRTC_ERROR_INVALID_PROGRAM = 4
+    NVRTC_ERROR_INVALID_OPTION = 5
+    NVRTC_ERROR_COMPILATION = 6
+    NVRTC_ERROR_BUILTIN_OPERATION_FAILURE = 7
+    NVRTC_ERROR_NO_NAME_EXPRESSIONS_AFTER_COMPILATION = 8
+    NVRTC_ERROR_NO_LOWERED_NAMES_BEFORE_COMPILATION = 9
+    NVRTC_ERROR_NAME_EXPRESSION_NOT_VALID = 10
+    NVRTC_ERROR_INTERNAL_ERROR = 11
+
+
+_nvrtc_lock = threading.Lock()
+
+
+class NvrtcProgram:
+    """
+    A class for managing the lifetime of nvrtcProgram instances. Instances of
+    the class own an nvrtcProgram; when an instance is deleted, the underlying
+    nvrtcProgram is destroyed using the appropriate NVRTC API.
+    """
+    def __init__(self, nvrtc, handle):
+        self._nvrtc = nvrtc
+        self._handle = handle
+
+    @property
+    def handle(self):
+        return self._handle
+
+    def __del__(self):
+        if self._handle:
+            self._nvrtc.destroy_program(self)
+
+
+class NVRTC:
+    """
+    Provides a Pythonic interface to the NVRTC APIs, abstracting away the C API
+    calls.
+
+    The sole instance of this class is a process-wide singleton, similar to the
+    NVVM interface. Initialization is protected by a lock and uses the standard
+    (for Numba) open_cudalib function to load the NVRTC library.
+    """
+    _PROTOTYPES = {
+        # nvrtcResult nvrtcVersion(int *major, int *minor)
+        'nvrtcVersion': (nvrtc_result, POINTER(c_int), POINTER(c_int)),
+        # nvrtcResult nvrtcCreateProgram(nvrtcProgram *prog,
+        #                                const char *src,
+        #                                const char *name,
+        #                                int numHeaders,
+        #                                const char * const *headers,
+        #                                const char * const *includeNames)
+        'nvrtcCreateProgram': (nvrtc_result, nvrtc_program, c_char_p, c_char_p,
+                               c_int, POINTER(c_char_p), POINTER(c_char_p)),
+        # nvrtcResult nvrtcDestroyProgram(nvrtcProgram *prog);
+        'nvrtcDestroyProgram': (nvrtc_result, POINTER(nvrtc_program)),
+        # nvrtcResult nvrtcCompileProgram(nvrtcProgram prog,
+        #                                 int numOptions,
+        #                                 const char * const *options)
+        'nvrtcCompileProgram': (nvrtc_result, nvrtc_program, c_int,
+                                POINTER(c_char_p)),
+        # nvrtcResult nvrtcGetPTXSize(nvrtcProgram prog, size_t *ptxSizeRet);
+        'nvrtcGetPTXSize': (nvrtc_result, nvrtc_program, POINTER(c_size_t)),
+        # nvrtcResult nvrtcGetPTX(nvrtcProgram prog, char *ptx);
+        'nvrtcGetPTX': (nvrtc_result, nvrtc_program, c_char_p),
+        # nvrtcResult nvrtcGetCUBINSize(nvrtcProgram prog,
+        #                               size_t *cubinSizeRet);
+        'nvrtcGetCUBINSize': (nvrtc_result, nvrtc_program, POINTER(c_size_t)),
+        # nvrtcResult nvrtcGetCUBIN(nvrtcProgram prog, char *cubin);
+        'nvrtcGetCUBIN': (nvrtc_result, nvrtc_program, c_char_p),
+        # nvrtcResult nvrtcGetProgramLogSize(nvrtcProgram prog,
+        #                                    size_t *logSizeRet);
+        'nvrtcGetProgramLogSize': (nvrtc_result, nvrtc_program,
+                                   POINTER(c_size_t)),
+        # nvrtcResult nvrtcGetProgramLog(nvrtcProgram prog, char *log);
+        'nvrtcGetProgramLog': (nvrtc_result, nvrtc_program, c_char_p),
+    }
+
+    # Singleton reference
+    __INSTANCE = None
+
+    def __new__(cls):
+        with _nvrtc_lock:
+            if cls.__INSTANCE is None:
+                from numba.cuda.cudadrv.libs import open_cudalib
+                cls.__INSTANCE = inst = object.__new__(cls)
+                try:
+                    lib = open_cudalib('nvrtc')
+                except OSError as e:
+                    cls.__INSTANCE = None
+                    raise NvrtcSupportError("NVRTC cannot be loaded") from e
+
+                # Find & populate functions
+                for name, proto in inst._PROTOTYPES.items():
+                    func = getattr(lib, name)
+                    func.restype = proto[0]
+                    func.argtypes = proto[1:]
+
+                    @functools.wraps(func)
+                    def checked_call(*args, func=func, name=name):
+                        error = func(*args)
+                        if error == NvrtcResult.NVRTC_ERROR_COMPILATION:
+                            raise NvrtcCompilationError()
+                        elif error != NvrtcResult.NVRTC_SUCCESS:
+                            try:
+                                error_name = NvrtcResult(error).name
+                            except ValueError:
+                                error_name = ('Unknown nvrtc_result '
+                                              f'(error code: {error})')
+                            msg = f'Failed to call {name}: {error_name}'
+                            raise NvrtcError(msg)
+
+                    setattr(inst, name, checked_call)
+
+        return cls.__INSTANCE
+
+    def get_version(self):
+        """
+        Get the NVRTC version as a tuple (major, minor).
+        """
+        major = c_int()
+        minor = c_int()
+        self.nvrtcVersion(byref(major), byref(minor))
+        return major.value, minor.value
+
+    def create_program(self, src, name):
+        """
+        Create an NVRTC program with managed lifetime.
+        """
+        if isinstance(src, str):
+            src = src.encode()
+        if isinstance(name, str):
+            name = name.encode()
+
+        handle = nvrtc_program()
+
+        # The final three arguments are for passing the contents of headers -
+        # this is not supported, so there are 0 headers and the header names
+        # and contents are null.
+        self.nvrtcCreateProgram(byref(handle), src, name, 0, None, None)
+        return NvrtcProgram(self, handle)
+
+    def compile_program(self, program, options):
+        """
+        Compile an NVRTC program. Compilation may fail due to a user error in
+        the source; this function returns ``True`` if there is a compilation
+        error and ``False`` on success.
+        """
+        # We hold a list of encoded options to ensure they can't be collected
+        # prior to the call to nvrtcCompileProgram
+        encoded_options = [opt.encode() for opt in options]
+        option_pointers = [c_char_p(opt) for opt in encoded_options]
+        c_options_type = (c_char_p * len(options))
+        c_options = c_options_type(*option_pointers)
+        try:
+            self.nvrtcCompileProgram(program.handle, len(options), c_options)
+            return False
+        except NvrtcCompilationError:
+            return True
+
+    def destroy_program(self, program):
+        """
+        Destroy an NVRTC program.
+        """
+        self.nvrtcDestroyProgram(byref(program.handle))
+
+    def get_compile_log(self, program):
+        """
+        Get the compile log as a Python string.
+        """
+        log_size = c_size_t()
+        self.nvrtcGetProgramLogSize(program.handle, byref(log_size))
+
+        log = (c_char * log_size.value)()
+        self.nvrtcGetProgramLog(program.handle, log)
+
+        return log.value.decode()
+
+    def get_ptx(self, program):
+        """
+        Get the compiled PTX as a Python string.
+        """
+        ptx_size = c_size_t()
+        self.nvrtcGetPTXSize(program.handle, byref(ptx_size))
+
+        ptx = (c_char * ptx_size.value)()
+        self.nvrtcGetPTX(program.handle, ptx)
+
+        return ptx.value.decode()
+
+
+def compile(src, name, cc):
+    """
+    Compile a CUDA C/C++ source to PTX for a given compute capability.
+
+    :param src: The source code to compile
+    :type src: str
+    :param name: The filename of the source (for information only)
+    :type name: str
+    :param cc: A tuple ``(major, minor)`` of the compute capability
+    :type cc: tuple
+    :return: The compiled PTX and compilation log
+    :rtype: tuple
+    """
+    nvrtc = NVRTC()
+    program = nvrtc.create_program(src, name)
+
+    # Compilation options:
+    # - Compile for the current device's compute capability.
+    # - The CUDA include path is added.
+    # - Relocatable Device Code (rdc) is needed to prevent device functions
+    #   being optimized away.
+    major, minor = cc
+    arch = f'--gpu-architecture=compute_{major}{minor}'
+    include = f'-I{config.CUDA_INCLUDE_PATH}'
+
+    cudadrv_path = os.path.dirname(os.path.abspath(__file__))
+    numba_cuda_path = os.path.dirname(cudadrv_path)
+    numba_include = f'-I{numba_cuda_path}'
+    options = [arch, include, numba_include, '-rdc', 'true']
+
+    # Compile the program
+    compile_error = nvrtc.compile_program(program, options)
+
+    # Get log from compilation
+    log = nvrtc.get_compile_log(program)
+
+    # If the compile failed, provide the log in an exception
+    if compile_error:
+        msg = (f'NVRTC Compilation failure whilst compiling {name}:\n\n{log}')
+        raise NvrtcError(msg)
+
+    # Otherwise, if there's any content in the log, present it as a warning
+    if log:
+        msg = (f"NVRTC log messages whilst compiling {name}:\n\n{log}")
+        warnings.warn(msg)
+
+    ptx = nvrtc.get_ptx(program)
+    return ptx, log