Repository: cupy/cupy

Issue:
Make `nvrtc.getCUBIN()` default on CUDA 11.1+?
In #4558 the new `nvrtcGetCUBIN()` API was exposed. The purpose was to avoid the driver from performing just-in-time compilation of PTX codes, which [was shown slow unless the jitted kernels are cached](https://github.com/cupy/cupy/pull/4558#issuecomment-765610834) (CUDA's default behavior). However, we decided to temporarily disable this feature in 8cb443f. This issue is for discussing the strategy to utilize this feature, which IMHO is a promising one 🙂


Discussion:
FWIW, when determining the performance of `getCUBIN()` it's best to set `CUDA_CACHE_DISABLE=1`. For fair comparison, it's important to always force JIT compiling PTX instead of fetching from cache. It is expected to benefit fresh kernels.
If `getCUBIN()` is enabled by default, this problem would be avoided: https://stackoverflow.com/a/66788739/2344149

tl;dr: To keep the kernel cache valid (without recompiling), users not only have to keep `~/.cupy/kernel_cache` but also `~/.nv/ComputeCache`, because CuPy only gets and caches the PTXs from `getPTX()`.

Resolving patch:
diff --git a/cupy/cuda/compiler.py b/cupy/cuda/compiler.py
--- a/cupy/cuda/compiler.py
+++ b/cupy/cuda/compiler.py
@@ -120,29 +120,49 @@ def _get_nvrtc_version():
 _tegra_archs = ('53', '62', '72')
 
 
-@_util.memoize(for_each_device=True)
-def _get_arch():
-    # See Supported Compile Options section of NVRTC User Guide for
-    # the maximum value allowed for `--gpu-architecture`.
+@_util.memoize()
+def _get_max_compute_capability():
     major, minor = _get_nvrtc_version()
     if major < 10 or (major == 10 and minor == 0):
         # CUDA 9.x / 10.0
-        _nvrtc_max_compute_capability = '70'
+        nvrtc_max_compute_capability = '70'
     elif major < 11:
         # CUDA 10.1 / 10.2
-        _nvrtc_max_compute_capability = '75'
+        nvrtc_max_compute_capability = '75'
     elif major == 11 and minor == 0:
         # CUDA 11.0
-        _nvrtc_max_compute_capability = '80'
+        nvrtc_max_compute_capability = '80'
     else:
         # CUDA 11.1 / 11.2 / 11.3 / 11.4
-        _nvrtc_max_compute_capability = '86'
+        nvrtc_max_compute_capability = '86'
+    return nvrtc_max_compute_capability
+
+
+@_util.memoize(for_each_device=True)
+def _get_arch():
+    # See Supported Compile Options section of NVRTC User Guide for
+    # the maximum value allowed for `--gpu-architecture`.
+    nvrtc_max_compute_capability = _get_max_compute_capability()
 
     arch = device.Device().compute_capability
     if arch in _tegra_archs:
         return arch
     else:
-        return min(arch, _nvrtc_max_compute_capability)
+        return min(arch, nvrtc_max_compute_capability)
+
+
+def _get_arch_for_options_for_nvrtc(arch=None):
+    # NVRTC in CUDA 11.3+ generates PTX that cannot be run an earlier driver
+    # version than the one included in the used CUDA version, as
+    # documented in:
+    # https://docs.nvidia.com/cuda/archive/11.3.0/nvrtc/index.html#versioning
+    # Here we use `-arch=sm_*` instead of `-arch=compute_*` to directly
+    # generate cubin (SASS) instead of PTX. See #5097 for details.
+    if arch is None:
+        arch = _get_arch()
+    if _cuda_hip_version >= 11010 and arch < _get_max_compute_capability():
+        return f'-arch=sm_{arch}', 'cubin'
+    return f'-arch=compute_{arch}', 'ptx'
 
 
 def _is_cudadevrt_needed(options):
@@ -236,12 +256,6 @@ def _jitify_prep(source, options, cu_path):
 def compile_using_nvrtc(source, options=(), arch=None, filename='kern.cu',
                         name_expressions=None, log_stream=None,
                         cache_in_memory=False, jitify=False):
-    # For hipRTC, arch is ignored
-    if not runtime.is_hip:
-        if not arch:
-            arch = _get_arch()
-        options += ('-arch=compute_{}'.format(arch),)
-
     def _compile(
             source, options, cu_path, name_expressions, log_stream, jitify):
 
@@ -251,17 +265,22 @@ def _compile(
         else:
             headers = include_names = ()
 
+        if not runtime.is_hip:
+            arch_opt, method = _get_arch_for_options_for_nvrtc(arch)
+            options += (arch_opt,)
+
         prog = _NVRTCProgram(source, cu_path, headers, include_names,
-                             name_expressions=name_expressions)
+                             name_expressions=name_expressions, method=method)
+
         try:
-            ptx, mapping = prog.compile(options, log_stream)
+            compiled_obj, mapping = prog.compile(options, log_stream)
         except CompileException as e:
             dump = _get_bool_env_variable(
                 'CUPY_DUMP_CUDA_SOURCE_ON_ERROR', False)
             if dump:
                 e.dump(sys.stderr)
             raise
-        return ptx, mapping
+        return compiled_obj, mapping
 
     if not cache_in_memory:
         with tempfile.TemporaryDirectory() as root_dir:
@@ -368,8 +387,9 @@ def compile_using_nvcc(source, options=(), arch=None,
 
 def _preprocess(source, options, arch, backend):
     if backend == 'nvrtc':
+        # For the preprocess it is enough to use PTX method
+        # we don't need to explicitly obtain a CUBIN file.
         options += ('-arch=compute_{}'.format(arch),)
-
         prog = _NVRTCProgram(source)
         try:
             result, _ = prog.compile(options)
@@ -469,7 +489,8 @@ def _compile_with_cache_cuda(
     if jitify and backend != 'nvrtc':
         raise ValueError('jitify only works with NVRTC')
 
-    env = (arch, options, _get_nvrtc_version(), backend)
+    env = ((arch, options, _get_nvrtc_version(), backend)
+           + _get_arch_for_options_for_nvrtc(arch))
     base = _empty_file_preprocess_cache.get(env, None)
     if base is None:
         # This is for checking NVRTC/NVCC compiler internal version
@@ -598,7 +619,7 @@ def dump(self, f):
 class _NVRTCProgram(object):
 
     def __init__(self, src, name='default_program', headers=(),
-                 include_names=(), name_expressions=None):
+                 include_names=(), name_expressions=None, method='ptx'):
         self.ptr = None
 
         if isinstance(src, bytes):
@@ -610,6 +631,7 @@ def __init__(self, src, name='default_program', headers=(),
         self.name = name
         self.ptr = nvrtc.createProgram(src, name, headers, include_names)
         self.name_expressions = name_expressions
+        self.method = method
 
     def __del__(self, is_shutting_down=_util.is_shutting_down):
         if is_shutting_down():
@@ -630,8 +652,13 @@ def compile(self, options=(), log_stream=None):
                     mapping[ker] = nvrtc.getLoweredName(self.ptr, ker)
             if log_stream is not None:
                 log_stream.write(nvrtc.getProgramLog(self.ptr))
-            # TODO(leofang): use getCUBIN() for _cuda_hip_version >= 11010?
-            return nvrtc.getPTX(self.ptr), mapping
+            # This is to ensure backwards compatibility with nvrtc
+            if self.method == 'cubin':
+                return nvrtc.getCUBIN(self.ptr), mapping
+            elif self.method == 'ptx':
+                return nvrtc.getPTX(self.ptr), mapping
+            else:
+                raise RuntimeError('Unknown NVRTC compile method')
         except nvrtc.NVRTCError:
             log = nvrtc.getProgramLog(self.ptr)
             raise CompileException(log, self.src, self.name, options,


Test patch:
diff --git a/tests/cupy_tests/core_tests/test_raw.py b/tests/cupy_tests/core_tests/test_raw.py
--- a/tests/cupy_tests/core_tests/test_raw.py
+++ b/tests/cupy_tests/core_tests/test_raw.py
@@ -16,6 +16,7 @@
 from cupy._core import _accelerator
 from cupy.cuda import compiler
 from cupy.cuda import memory
+from cupy_backends.cuda.api import driver
 
 
 _test_source1 = r'''
@@ -633,7 +634,8 @@ def test_dynamical_parallelism_compile_failure(self):
         N = 10
         inner_chunk = 2
         x = cupy.zeros((N,), dtype=cupy.float32)
-        if self.backend == 'nvrtc' and not cupy.cuda.runtime.is_hip:
+        if (self.backend == 'nvrtc' and not cupy.cuda.runtime.is_hip
+                and driver.get_build_version() < 11010):
             # raised when calling ls.complete()
             error = cupy.cuda.driver.CUDADriverError
         else:  # nvcc, hipcc, hiprtc