Repository: numba/numba

Issue:
Feature: support compilation caching for CUDA target
<!--

Thanks for opening an issue! To help the Numba team handle your information
efficiently, please first ensure that there is no other issue present that
already describes the issue you have
(search at https://github.com/numba/numba/issues?&q=is%3Aissue).

-->

## Feature request

As mentioned here: https://numba.discourse.group/t/cuda-jit-is-there-cache-true/163/4

there is no caching for CUDA jit compilation. To confirm in most recent CUDA, when I try, I get:
```sh
    def _reduce_states(self):
        """
        Reduce the instance for serialization. We retain the PTX and cubins,
        but loaded functions are discarded. They are recreated when needed
        after deserialization.
        """
        if self._linking_files:
            msg = 'Cannot pickle CUDACodeLibrary with linking files'
>           raise RuntimeError(msg)
E           RuntimeError: Cannot pickle CUDACodeLibrary with linking files
```

Our test suite now has >60 kernels for unit testing device functions. Many are taking ~5-10 sec to compile, so this is becoming a big issue for us.


Discussion:
Compilation caching is supported on CUDA: https://numba.readthedocs.io/en/latest/cuda/caching.html

However, when linking in external files the caching is not supported because Numba doesn't have a way to detect if these files changed between runs. Are you using the `link` kwarg of the `@cuda.jit` decorator, or `float16` in the kernels?

The docs should be updated to mention this limitation.
I'm using neither.

The kernel call tree is rather large (for me to post as an example) - is there a way to get more info about what is causing external files to be included?

**Update**: seems to be anything that uses cooperative groups :( - e.g. the following crashes with that error:
```py
@cuda.jit(cache=True)  # type: ignore
def _init_datastructures_kernel(
    host_context: host.HostContext,
    i_level: np.int16,
) -> None:
    g: cuda.cg.GridGroup = cuda.cg.this_grid()  # type: ignore
```
Perhaps we could have an option (default w/ warning?) to exclude linked files? The doc already says that cache might not detect change in libraries.

Resolving patch:
diff --git a/numba/cuda/codegen.py b/numba/cuda/codegen.py
--- a/numba/cuda/codegen.py
+++ b/numba/cuda/codegen.py
@@ -3,6 +3,7 @@
 from numba.core import config, serialize
 from numba.core.codegen import Codegen, CodeLibrary
 from .cudadrv import devices, driver, nvvm, runtime
+from numba.cuda.cudadrv.libs import get_cudalib
 
 import os
 import subprocess
@@ -74,6 +75,8 @@ def __init__(self, codegen, name, entry_name=None, max_registers=None,
         # Files to link with the generated PTX. These are linked using the
         # Driver API at link time.
         self._linking_files = set()
+        # Should we link libcudadevrt?
+        self.needs_cudadevrt = False
 
         # Cache the LLVM IR string
         self._llvm_strs = None
@@ -156,6 +159,8 @@ def get_cubin(self, cc=None):
             linker.add_ptx(ptx.encode())
         for path in self._linking_files:
             linker.add_file_guess_ext(path)
+        if self.needs_cudadevrt:
+            linker.add_file_guess_ext(get_cudalib('cudadevrt', static=True))
 
         cubin = linker.complete()
         self._cubin_cache[cc] = cubin
@@ -283,12 +288,14 @@ def _reduce_states(self):
             cubin_cache=self._cubin_cache,
             linkerinfo_cache=self._linkerinfo_cache,
             max_registers=self._max_registers,
-            nvvm_options=self._nvvm_options
+            nvvm_options=self._nvvm_options,
+            needs_cudadevrt=self.needs_cudadevrt
         )
 
     @classmethod
     def _rebuild(cls, codegen, name, entry_name, llvm_strs, ptx_cache,
-                 cubin_cache, linkerinfo_cache, max_registers, nvvm_options):
+                 cubin_cache, linkerinfo_cache, max_registers, nvvm_options,
+                 needs_cudadevrt):
         """
         Rebuild an instance.
         """
@@ -301,6 +308,7 @@ def _rebuild(cls, codegen, name, entry_name, llvm_strs, ptx_cache,
 
         instance._max_registers = max_registers
         instance._nvvm_options = nvvm_options
+        instance.needs_cudadevrt = needs_cudadevrt
 
         instance._finalized = True
 
diff --git a/numba/cuda/dispatcher.py b/numba/cuda/dispatcher.py
--- a/numba/cuda/dispatcher.py
+++ b/numba/cuda/dispatcher.py
@@ -16,7 +16,6 @@
 from numba.cuda.compiler import compile_cuda, CUDACompiler
 from numba.cuda.cudadrv import driver
 from numba.cuda.cudadrv.devices import get_context
-from numba.cuda.cudadrv.libs import get_cudalib
 from numba.cuda.descriptor import cuda_target
 from numba.cuda.errors import (missing_launch_config_msg,
                                normalize_kernel_dimensions)
@@ -104,7 +103,7 @@ def __init__(self, py_func, argtypes, link=None, debug=False,
         self.cooperative = 'cudaCGGetIntrinsicHandle' in lib.get_asm_str()
         # We need to link against cudadevrt if grid sync is being used.
         if self.cooperative:
-            link.append(get_cudalib('cudadevrt', static=True))
+            lib.needs_cudadevrt = True
 
         res = [fn for fn in cuda_fp16_math_funcs
                if (f'__numba_wrapper_{fn}' in lib.get_asm_str())]


Test patch:
diff --git a/numba/cuda/tests/cudapy/cache_usecases.py b/numba/cuda/tests/cudapy/cache_usecases.py
--- a/numba/cuda/tests/cudapy/cache_usecases.py
+++ b/numba/cuda/tests/cudapy/cache_usecases.py
@@ -198,6 +198,17 @@ def simple_usecase_kernel(r, x):
 simple_usecase_caller = CUDAUseCase(simple_usecase_kernel)
 
 
+# Usecase with cooperative groups
+
+@cuda.jit(cache=True)
+def cg_usecase_kernel(r, x):
+    grid = cuda.cg.this_grid()
+    grid.sync()
+
+
+cg_usecase = CUDAUseCase(cg_usecase_kernel)
+
+
 class _TestModule(CUDATestCase):
     """
     Tests for functionality of this module's functions.
diff --git a/numba/cuda/tests/cudapy/test_caching.py b/numba/cuda/tests/cudapy/test_caching.py
--- a/numba/cuda/tests/cudapy/test_caching.py
+++ b/numba/cuda/tests/cudapy/test_caching.py
@@ -6,8 +6,9 @@
 
 from numba import cuda
 from numba.core.errors import NumbaWarning
-from numba.cuda.codegen import CUDACodeLibrary
-from numba.cuda.testing import CUDATestCase, skip_on_cudasim, test_data_dir
+from numba.cuda.testing import (CUDATestCase, skip_on_cudasim,
+                                skip_unless_cc_60, skip_if_cudadevrt_missing,
+                                skip_if_mvc_enabled, test_data_dir)
 from numba.tests.support import SerialMixin
 from numba.tests.test_caching import (DispatcherCacheUsecasesTest,
                                       skip_bad_access)
@@ -153,6 +154,22 @@ def test_same_names(self):
         f = mod.renamed_function2
         self.assertPreciseEqual(f(2), 8)
 
+    @skip_unless_cc_60
+    @skip_if_cudadevrt_missing
+    @skip_if_mvc_enabled('CG not supported with MVC')
+    def test_cache_cg(self):
+        # Functions using cooperative groups should be cacheable. See Issue
+        # #8888: https://github.com/numba/numba/issues/8888
+        self.check_pycache(0)
+        mod = self.import_module()
+        self.check_pycache(0)
+
+        mod.cg_usecase(0)
+        self.check_pycache(2)  # 1 index, 1 data
+
+        # Check the code runs ok from another process
+        self.run_in_separate_process()
+
     def _test_pycache_fallback(self):
         """
         With a disabled __pycache__, test there is a working fallback
@@ -482,6 +499,10 @@ class TestCUDACodeLibrary(CUDATestCase):
     # explicitly check
 
     def test_cannot_serialize_unfinalized(self):
+        # The CUDA codegen failes to import under the simulator, so we cannot
+        # import it at the top level
+        from numba.cuda.codegen import CUDACodeLibrary
+
         # Usually a CodeLibrary requires a real CodeGen, but since we don't
         # interact with it, anything will do
         codegen = object()