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

numba_cuda/numba/cuda/cudadrv/nvrtc.py

@@ -1,8 +1,13 @@
 from ctypes import byref, c_char, c_char_p, c_int, c_size_t, c_void_p, POINTER
 from enum import IntEnum
-from numba.cuda.cudadrv.error import (NvrtcError, NvrtcCompilationError,
-                                      NvrtcSupportError)
+from numba.cuda.cudadrv.error import (
+    NvrtcError,
+    NvrtcBuiltinOperationFailure,
+    NvrtcCompilationError,
+    NvrtcSupportError,
+)
 from numba.cuda.cuda_paths import get_cuda_paths
+
 import functools
 import os
 import threading
@@ -39,6 +44,7 @@ class NvrtcProgram:
     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
@@ -62,46 +68,67 @@ class NVRTC:
     (for Numba) open_cudalib function to load the NVRTC library.
     """
 
+    _CU11_2ONLY_PROTOTYPES = {
+        # nvrtcResult nvrtcGetNumSupportedArchs(int *numArchs);
+        "nvrtcGetNumSupportedArchs": (nvrtc_result, POINTER(c_int)),
+        # nvrtcResult nvrtcGetSupportedArchs(int *supportedArchs);
+        "nvrtcGetSupportedArchs": (nvrtc_result, POINTER(c_int)),
+    }
+
     _CU12ONLY_PROTOTYPES = {
         # nvrtcResult nvrtcGetLTOIRSize(nvrtcProgram prog, size_t *ltoSizeRet);
         "nvrtcGetLTOIRSize": (nvrtc_result, nvrtc_program, POINTER(c_size_t)),
         # nvrtcResult nvrtcGetLTOIR(nvrtcProgram prog, char *lto);
-        "nvrtcGetLTOIR": (nvrtc_result, nvrtc_program, c_char_p)
+        "nvrtcGetLTOIR": (nvrtc_result, nvrtc_program, c_char_p),
     }
 
     _PROTOTYPES = {
         # nvrtcResult nvrtcVersion(int *major, int *minor)
-        'nvrtcVersion': (nvrtc_result, POINTER(c_int), POINTER(c_int)),
+        "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)),
+        "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)),
+        "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)),
+        "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)),
+        "nvrtcGetPTXSize": (nvrtc_result, nvrtc_program, POINTER(c_size_t)),
         # nvrtcResult nvrtcGetPTX(nvrtcProgram prog, char *ptx);
-        'nvrtcGetPTX': (nvrtc_result, nvrtc_program, c_char_p),
+        "nvrtcGetPTX": (nvrtc_result, nvrtc_program, c_char_p),
         # nvrtcResult nvrtcGetCUBINSize(nvrtcProgram prog,
         #                               size_t *cubinSizeRet);
-        'nvrtcGetCUBINSize': (nvrtc_result, nvrtc_program, POINTER(c_size_t)),
+        "nvrtcGetCUBINSize": (nvrtc_result, nvrtc_program, POINTER(c_size_t)),
         # nvrtcResult nvrtcGetCUBIN(nvrtcProgram prog, char *cubin);
-        'nvrtcGetCUBIN': (nvrtc_result, nvrtc_program, c_char_p),
+        "nvrtcGetCUBIN": (nvrtc_result, nvrtc_program, c_char_p),
         # nvrtcResult nvrtcGetProgramLogSize(nvrtcProgram prog,
         #                                    size_t *logSizeRet);
-        'nvrtcGetProgramLogSize': (nvrtc_result, nvrtc_program,
-                                   POINTER(c_size_t)),
+        "nvrtcGetProgramLogSize": (
+            nvrtc_result,
+            nvrtc_program,
+            POINTER(c_size_t),
+        ),
         # nvrtcResult nvrtcGetProgramLog(nvrtcProgram prog, char *log);
-        'nvrtcGetProgramLog': (nvrtc_result, nvrtc_program, c_char_p),
+        "nvrtcGetProgramLog": (nvrtc_result, nvrtc_program, c_char_p),
     }
 
     # Singleton reference
@@ -111,14 +138,18 @@ class NVRTC:
         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')
+                    lib = open_cudalib("nvrtc")
                 except OSError as e:
                     cls.__INSTANCE = None
                     raise NvrtcSupportError("NVRTC cannot be loaded") from e
 
                 from numba.cuda.cudadrv.runtime import get_version
+
+                if get_version() >= (11, 2):
+                    inst._PROTOTYPES |= inst._CU11_2ONLY_PROTOTYPES
                 if get_version() >= (12, 0):
                     inst._PROTOTYPES |= inst._CU12ONLY_PROTOTYPES
 
@@ -133,19 +164,73 @@ class NVRTC:
                         error = func(*args)
                         if error == NvrtcResult.NVRTC_ERROR_COMPILATION:
                             raise NvrtcCompilationError()
+                        elif (
+                            error
+                            == NvrtcResult.NVRTC_ERROR_BUILTIN_OPERATION_FAILURE
+                        ):
+                            raise NvrtcBuiltinOperationFailure()
                         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}'
+                                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_supported_archs(self):
+        """
+        Get Supported Architectures by NVRTC as list of arch tuples.
+        """
+        ver = self.get_version()
+        if ver < (11, 0):
+            raise RuntimeError(
+                "Unsupported CUDA version. CUDA 11.0 or higher is required."
+            )
+        elif ver == (11, 0):
+            return [
+                (3, 0),
+                (3, 2),
+                (3, 5),
+                (3, 7),
+                (5, 0),
+                (5, 2),
+                (5, 3),
+                (6, 0),
+                (6, 1),
+                (6, 2),
+                (7, 0),
+                (7, 2),
+                (7, 5),
+            ]
+        elif ver == (11, 1):
+            return [
+                (3, 5),
+                (3, 7),
+                (5, 0),
+                (5, 2),
+                (5, 3),
+                (6, 0),
+                (6, 1),
+                (6, 2),
+                (7, 0),
+                (7, 2),
+                (7, 5),
+                (8, 0),
+            ]
+        else:
+            num = c_int()
+            self.nvrtcGetNumSupportedArchs(byref(num))
+            archs = (c_int * num.value)()
+            self.nvrtcGetSupportedArchs(archs)
+            return [(archs[i] // 10, archs[i] % 10) for i in range(num.value)]
+
     def get_version(self):
         """
         Get the NVRTC version as a tuple (major, minor).
@@ -182,12 +267,12 @@ class NVRTC:
         # 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_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:
+        except (NvrtcCompilationError, NvrtcBuiltinOperationFailure):
             return True
 
     def destroy_program(self, program):
@@ -251,13 +336,37 @@ def compile(src, name, cc, ltoir=False):
     nvrtc = NVRTC()
     program = nvrtc.create_program(src, name)
 
+    version = nvrtc.get_version()
+    ver_str = lambda v: ".".join(v)
+    if version < (11, 0):
+        raise RuntimeError(
+            "Unsupported CUDA version. CUDA 11.0 or higher is required."
+        )
+    else:
+        supported_arch = nvrtc.get_supported_archs()
+        try:
+            found = max(filter(lambda v: v <= cc, [v for v in supported_arch]))
+        except ValueError:
+            raise RuntimeError(
+                f"Device compute capability {ver_str(cc)} is less than the "
+                f"minimum supported by NVRTC {ver_str(version)}. Supported "
+                "compute capabilities are "
+                f"{', '.join([ver_str(v) for v in supported_arch])}."
+            )
+
+        if found != cc:
+            warnings.warn(
+                f"Device compute capability {ver_str(cc)} is not supported by "
+                f"NVRTC {ver_str(version)}. Using {ver_str(found)} instead."
+            )
+
     # 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}'
+    major, minor = found
+    arch = f"--gpu-architecture=compute_{major}{minor}"
 
     cuda_include = [
         f"-I{get_cuda_paths()['include_dir'].info}",
@@ -265,12 +374,12 @@ def compile(src, name, cc, ltoir=False):
 
     cudadrv_path = os.path.dirname(os.path.abspath(__file__))
     numba_cuda_path = os.path.dirname(cudadrv_path)
-    numba_include = f'-I{numba_cuda_path}'
+    numba_include = f"-I{numba_cuda_path}"
 
     nrt_path = os.path.join(numba_cuda_path, "runtime")
-    nrt_include = f'-I{nrt_path}'
+    nrt_include = f"-I{nrt_path}"
 
-    options = [arch, *cuda_include, numba_include, nrt_include, '-rdc', 'true']
+    options = [arch, *cuda_include, numba_include, nrt_include, "-rdc", "true"]
 
     if ltoir:
         options.append("-dlto")
@@ -286,12 +395,12 @@ def compile(src, name, cc, ltoir=False):
 
     # If the compile failed, provide the log in an exception
     if compile_error:
-        msg = (f'NVRTC Compilation failure whilst compiling {name}:\n\n{log}')
+        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}")
+        msg = f"NVRTC log messages whilst compiling {name}:\n\n{log}"
         warnings.warn(msg)
 
     if ltoir: