Repository: numba/numba

Issue:
jit cache pickle exception on many local cuda arrays
Anyone facing the same ?
confirmed on py3.9/3.10 and nb 0.56.0

```python
import numba as nb
from numba import cuda


@cuda.jit(cache=True)
def test():
    aa = cuda.local.array((1, 1), nb.float64)
    ab = cuda.local.array((1, 1), nb.float64)
    ac = cuda.local.array((1, 1), nb.float64)
    ad = cuda.local.array((1, 1), nb.float64)
    ae = cuda.local.array((1, 1), nb.float64)
    af = cuda.local.array((1, 1), nb.float64)
    ag = cuda.local.array((1, 1), nb.float64)
    ah = cuda.local.array((1, 1), nb.float64)
    ai = cuda.local.array((1, 1), nb.float64)
    aj = cuda.local.array((1, 1), nb.float64)
    ak = cuda.local.array((1, 1), nb.float64)
    al = cuda.local.array((1, 1), nb.float64)
    am = cuda.local.array((1, 1), nb.float64)
    an = cuda.local.array((1, 1), nb.float64)
    ao = cuda.local.array((1, 1), nb.float64)
    ap = cuda.local.array((1, 1), nb.float64)
    ar = cuda.local.array((1, 1), nb.float64)
    at = cuda.local.array((1, 1), nb.float64)
    au = cuda.local.array((1, 1), nb.float64)
    av = cuda.local.array((1, 1), nb.float64)
    aw = cuda.local.array((1, 1), nb.float64)
    ax = cuda.local.array((1, 1), nb.float64)
    ay = cuda.local.array((1, 1), nb.float64)
    az = cuda.local.array((1, 1), nb.float64)


    aa[:] = 0
    ab[:] = 0
    ac[:] = 0
    ad[:] = 0
    ae[:] = 0
    af[:] = 0
    ag[:] = 0
    ah[:] = 0
    ai[:] = 0
    aj[:] = 0
    ak[:] = 0
    al[:] = 0
    am[:] = 0
    an[:] = 0
    ao[:] = 0
    ap[:] = 0
    ar[:] = 0
    at[:] = 0
    au[:] = 0
    av[:] = 0
    aw[:] = 0
    ax[:] = 0
    ay[:] = 0
    az[:] = 0



test[1, 32]()
```

```
/usr/local/lib/python3.9/site-packages/numba/cuda/dispatcher.py:488: NumbaPerformanceWarning: Grid size 1 will likely result in GPU under-utilization due to low occupancy.
  warn(NumbaPerformanceWarning(msg))
Traceback (most recent call last):
  File "/usr/local/lib/python3.9/site-packages/numba/cloudpickle/cloudpickle_fast.py", line 568, in dump
    return Pickler.dump(self, obj)
  File "/usr/local/lib/python3.9/site-packages/numba/core/serialize.py", line 186, in reducer_override
    return super().reducer_override(obj)
  File "/usr/local/lib/python3.9/site-packages/numba/cloudpickle/cloudpickle_fast.py", line 644, in reducer_override
    if sys.version_info[:2] < (3, 7) and _is_parametrized_type_hint(obj):  # noqa  # pragma: no branch
RecursionError: maximum recursion depth exceeded in comparison

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/tmp/pycharm_project_703/nb_err.py", line 60, in <module>
    test[1, 32]()
  File "/usr/local/lib/python3.9/site-packages/numba/cuda/dispatcher.py", line 491, in __call__
    return self.dispatcher.call(args, self.griddim, self.blockdim,
  File "/usr/local/lib/python3.9/site-packages/numba/cuda/dispatcher.py", line 625, in call
    kernel = _dispatcher.Dispatcher._cuda_call(self, *args)
  File "/usr/local/lib/python3.9/site-packages/numba/cuda/dispatcher.py", line 633, in _compile_for_args
    return self.compile(tuple(argtypes))
  File "/usr/local/lib/python3.9/site-packages/numba/cuda/dispatcher.py", line 797, in compile
    self._cache.save_overload(sig, kernel)
  File "/usr/local/lib/python3.9/site-packages/numba/core/caching.py", line 652, in save_overload
    self._save_overload(sig, data)
  File "/usr/local/lib/python3.9/site-packages/numba/core/caching.py", line 662, in _save_overload
    self._cache_file.save(key, data)
  File "/usr/local/lib/python3.9/site-packages/numba/core/caching.py", line 479, in save
    self._save_data(data_name, data)
  File "/usr/local/lib/python3.9/site-packages/numba/core/caching.py", line 537, in _save_data
    data = self._dump(data)
  File "/usr/local/lib/python3.9/site-packages/numba/core/caching.py", line 550, in _dump
    return dumps(obj)
  File "/usr/local/lib/python3.9/site-packages/numba/core/serialize.py", line 57, in dumps
    p.dump(obj)
  File "/usr/local/lib/python3.9/site-packages/numba/cloudpickle/cloudpickle_fast.py", line 575, in dump
    raise pickle.PicklingError(msg) from e
_pickle.PicklingError: Could not pickle object as excessively deep recursion required.
```



Discussion:
Thanks for the report and nice reproducer! I can reproduce the issue.
I'm not sure this is a CUDA-specific issue per se, it seems like the cache fails when the LLVM IR gets too big. I guess there's no optimization of the IR within Numba for the CUDA target though, so the cache might easily see much larger IR than with the CPU target.
Is there any quick fix ?
I can't think of a quick fix off the top of my head - I might be able to come up with some other workaround with a bit more thought though.
Not pickling the LLVM IR for CUDA-jitted functions works around it:

```diff
diff --git a/numba/cuda/codegen.py b/numba/cuda/codegen.py
index 47c6d36a7..42308a25e 100644
--- a/numba/cuda/codegen.py
+++ b/numba/cuda/codegen.py
@@ -305,7 +305,7 @@ class CUDACodeLibrary(serialize.ReduceMixin, CodeLibrary):
             codegen=None,
             name=self.name,
             entry_name=self._entry_name,
-            module=self._module,
+            module=None,
             linking_libraries=self._linking_libraries,
             ptx_cache=self._ptx_cache,
             cubin_cache=self._cubin_cache,
```

Perhaps this is the right thing to do - I need to think a little more about this. If possible, can you try applying the above patch locally and see if it fixes your use case and avoids introducing any unexpected side effects? (if you have multiple GPUs of different compute capabilities this might break that in its current naive form, btw).

Resolving patch:
diff --git a/numba/cuda/codegen.py b/numba/cuda/codegen.py
--- a/numba/cuda/codegen.py
+++ b/numba/cuda/codegen.py
@@ -80,6 +80,8 @@ def __init__(self, codegen, name, entry_name=None, max_registers=None,
         # Driver API at link time.
         self._linking_files = set()
 
+        # Cache the LLVM IR string
+        self._llvm_strs = None
         # Maps CC -> PTX string
         self._ptx_cache = {}
         # Maps CC -> cubin
@@ -95,8 +97,14 @@ def __init__(self, codegen, name, entry_name=None, max_registers=None,
         self._nvvm_options = nvvm_options
         self._entry_name = entry_name
 
+    @property
+    def llvm_strs(self):
+        if self._llvm_strs is None:
+            self._llvm_strs = [str(mod) for mod in self.modules]
+        return self._llvm_strs
+
     def get_llvm_str(self):
-        return str(self._module)
+        return self.llvm_strs[0]
 
     def get_asm_str(self, cc=None):
         return self._join_ptxes(self._get_ptxes(cc=cc))
@@ -126,7 +134,7 @@ def _get_ptxes(self, cc=None):
                     warn(NumbaInvalidConfigWarning(msg))
             options['debug'] = False
 
-        irs = [str(mod) for mod in self.modules]
+        irs = self.llvm_strs
 
         if options.get('debug', False):
             # If we're compiling with debug, we need to compile modules with
@@ -298,15 +306,15 @@ def _reduce_states(self):
         after deserialization.
         """
         if self._linking_files:
-            msg = ('cannot pickle CUDACodeLibrary function with additional '
-                   'libraries to link against')
+            msg = 'Cannot pickle CUDACodeLibrary with linking files'
             raise RuntimeError(msg)
+        if not self._finalized:
+            raise RuntimeError('Cannot pickle unfinalized CUDACodeLibrary')
         return dict(
             codegen=None,
             name=self.name,
             entry_name=self._entry_name,
-            module=self._module,
-            linking_libraries=self._linking_libraries,
+            llvm_strs=self.llvm_strs,
             ptx_cache=self._ptx_cache,
             cubin_cache=self._cubin_cache,
             linkerinfo_cache=self._linkerinfo_cache,
@@ -315,28 +323,23 @@ def _reduce_states(self):
         )
 
     @classmethod
-    def _rebuild(cls, codegen, name, entry_name, module, linking_libraries,
-                 ptx_cache, cubin_cache, linkerinfo_cache, max_registers,
-                 nvvm_options):
+    def _rebuild(cls, codegen, name, entry_name, llvm_strs, ptx_cache,
+                 cubin_cache, linkerinfo_cache, max_registers, nvvm_options):
         """
         Rebuild an instance.
         """
-        instance = cls.__new__(cls)
-        super(cls, instance).__init__(codegen, name)
-        instance._entry_name = entry_name
-
-        instance._module = module
-        instance._linking_libraries = linking_libraries
-        instance._linking_files = set()
+        instance = cls(codegen, name, entry_name=entry_name)
 
+        instance._llvm_strs = llvm_strs
         instance._ptx_cache = ptx_cache
         instance._cubin_cache = cubin_cache
         instance._linkerinfo_cache = linkerinfo_cache
-        instance._cufunc_cache = {}
 
         instance._max_registers = max_registers
         instance._nvvm_options = nvvm_options
 
+        instance._finalized = True
+
         return instance
 
 


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
@@ -135,6 +135,59 @@ def ambiguous_function(r, x):
 renamed_function2 = CUDAUseCase(ambiguous_function)
 
 
+@cuda.jit(cache=True)
+def many_locals():
+    aa = cuda.local.array((1, 1), np.float64)
+    ab = cuda.local.array((1, 1), np.float64)
+    ac = cuda.local.array((1, 1), np.float64)
+    ad = cuda.local.array((1, 1), np.float64)
+    ae = cuda.local.array((1, 1), np.float64)
+    af = cuda.local.array((1, 1), np.float64)
+    ag = cuda.local.array((1, 1), np.float64)
+    ah = cuda.local.array((1, 1), np.float64)
+    ai = cuda.local.array((1, 1), np.float64)
+    aj = cuda.local.array((1, 1), np.float64)
+    ak = cuda.local.array((1, 1), np.float64)
+    al = cuda.local.array((1, 1), np.float64)
+    am = cuda.local.array((1, 1), np.float64)
+    an = cuda.local.array((1, 1), np.float64)
+    ao = cuda.local.array((1, 1), np.float64)
+    ap = cuda.local.array((1, 1), np.float64)
+    ar = cuda.local.array((1, 1), np.float64)
+    at = cuda.local.array((1, 1), np.float64)
+    au = cuda.local.array((1, 1), np.float64)
+    av = cuda.local.array((1, 1), np.float64)
+    aw = cuda.local.array((1, 1), np.float64)
+    ax = cuda.local.array((1, 1), np.float64)
+    ay = cuda.local.array((1, 1), np.float64)
+    az = cuda.local.array((1, 1), np.float64)
+
+    aa[:] = 0
+    ab[:] = 0
+    ac[:] = 0
+    ad[:] = 0
+    ae[:] = 0
+    af[:] = 0
+    ag[:] = 0
+    ah[:] = 0
+    ai[:] = 0
+    aj[:] = 0
+    ak[:] = 0
+    al[:] = 0
+    am[:] = 0
+    an[:] = 0
+    ao[:] = 0
+    ap[:] = 0
+    ar[:] = 0
+    at[:] = 0
+    au[:] = 0
+    av[:] = 0
+    aw[:] = 0
+    ax[:] = 0
+    ay[:] = 0
+    az[:] = 0
+
+
 # Simple use case for multiprocessing test
 
 @cuda.jit(cache=True)
diff --git a/numba/cuda/tests/cudapy/jitlink.ptx b/numba/cuda/tests/cudapy/jitlink.ptx
new file mode 100644
--- /dev/null
+++ b/numba/cuda/tests/cudapy/jitlink.ptx
@@ -0,0 +1,30 @@
+//
+// Generated by NVIDIA NVVM Compiler
+// Compiler built on Tue Apr  1 03:34:02 2014 (1396341242)
+// Cuda compilation tools, release 6.0, V6.0.1
+//
+
+.version 4.0
+.target sm_20
+.address_size 64
+
+
+.visible .func  (.param .b32 func_retval0) bar(
+	.param .b64 bar_param_0,
+	.param .b32 bar_param_1
+)
+{
+	.reg .s32 	%r<4>;
+	.reg .s64 	%rd<2>;
+
+
+	ld.param.u64 	%rd1, [bar_param_0];
+	ld.param.u32 	%r1, [bar_param_1];
+	shl.b32 	%r2, %r1, 1;
+	st.u32 	[%rd1], %r2;
+	mov.u32 	%r3, 0;
+	st.param.b32	[func_retval0+0], %r3;
+	ret;
+}
+
+
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,10 +6,12 @@
 
 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
 from numba.tests.support import SerialMixin
 from numba.tests.test_caching import (DispatcherCacheUsecasesTest,
                                       skip_bad_access)
+from pathlib import Path
 
 
 @skip_on_cudasim('Simulator does not implement caching')
@@ -58,6 +60,18 @@ def test_no_caching(self):
         self.assertPreciseEqual(f(2, 3), 6)
         self.check_pycache(0)
 
+    def test_many_locals(self):
+        # Declaring many local arrays creates a very large LLVM IR, which
+        # cannot be pickled due to the level of recursion it requires to
+        # pickle. This test ensures that kernels with many locals (and
+        # therefore large IR) can be cached. See Issue #8373:
+        # https://github.com/numba/numba/issues/8373
+        self.check_pycache(0)
+        mod = self.import_module()
+        f = mod.many_locals
+        f[1, 1]()
+        self.check_pycache(2) # 1 index, 1 data
+
     def test_closure(self):
         mod = self.import_module()
 
@@ -190,6 +204,15 @@ def test_non_writable_pycache(self):
 
         self._test_pycache_fallback()
 
+    def test_cannot_cache_linking_libraries(self):
+        link = os.path.join(Path(__file__).parent.parent,
+                            'cudadrv', 'data', 'jitlink.ptx')
+        msg = 'Cannot pickle CUDACodeLibrary with linking files'
+        with self.assertRaisesRegex(RuntimeError, msg):
+            @cuda.jit('void()', cache=True, link=[link])
+            def f():
+                pass
+
 
 @skip_on_cudasim('Simulator does not implement caching')
 class CUDAAndCPUCachingTest(SerialMixin, DispatcherCacheUsecasesTest):
@@ -305,7 +328,7 @@ def test_cache(self):
         mod = self.import_module()
         self.check_pycache(0)
 
-        # Populate the cache with the first GPU
+        # Step 1. Populate the cache with the first GPU
         with gpus[0]:
             f = mod.add_usecase
             self.assertPreciseEqual(f(2, 3), 6)
@@ -324,8 +347,8 @@ def test_cache(self):
             self.check_pycache(6)  # 2 index, 4 data
             self.check_hits(f.func, 0, 2)
 
-        # Run with the second GPU - under present behaviour this doesn't
-        # further populate the cache.
+        # Step 2. Run with the second GPU - under present behaviour this
+        # doesn't further populate the cache.
         with gpus[1]:
             f = mod.add_usecase
             self.assertPreciseEqual(f(2, 3), 6)
@@ -344,8 +367,8 @@ def test_cache(self):
             self.check_pycache(6)  # cache unchanged
             self.check_hits(f.func, 0, 2)
 
-        # Run in a separate module with the second GPU - this populates the
-        # cache for the second CC.
+        # Step 3. Run in a separate module with the second GPU - this populates
+        # the cache for the second CC.
         mod2 = self.import_module()
         self.assertIsNot(mod, mod2)
 
@@ -367,6 +390,45 @@ def test_cache(self):
             self.check_pycache(10)  # 2 index, 8 data
             self.check_hits(f.func, 0, 2)
 
+        # The following steps check that we can use the NVVM IR loaded from the
+        # cache to generate PTX for a different compute capability to the
+        # cached cubin's CC. To check this, we create another module that loads
+        # the cached version containing a cubin for GPU 1. There will be no
+        # cubin for GPU 0, so when we try to use it the PTX must be generated.
+
+        mod3 = self.import_module()
+        self.assertIsNot(mod, mod3)
+
+        # Step 4. Run with GPU 1 and get a cache hit, loading the cache created
+        # during Step 3.
+        with gpus[1]:
+            f = mod3.add_usecase
+            self.assertPreciseEqual(f(2, 3), 6)
+            self.assertPreciseEqual(f(2.5, 3), 6.5)
+
+            f = mod3.record_return_aligned
+            rec = f(mod.aligned_arr, 1)
+            self.assertPreciseEqual(tuple(rec), (2, 43.5))
+
+            f = mod3.record_return_packed
+            rec = f(mod.packed_arr, 1)
+            self.assertPreciseEqual(tuple(rec), (2, 43.5))
+
+        # Step 5. Run with GPU 0 using the module from Step 4, to force PTX
+        # generation from cached NVVM IR.
+        with gpus[0]:
+            f = mod3.add_usecase
+            self.assertPreciseEqual(f(2, 3), 6)
+            self.assertPreciseEqual(f(2.5, 3), 6.5)
+
+            f = mod3.record_return_aligned
+            rec = f(mod.aligned_arr, 1)
+            self.assertPreciseEqual(tuple(rec), (2, 43.5))
+
+            f = mod3.record_return_packed
+            rec = f(mod.packed_arr, 1)
+            self.assertPreciseEqual(tuple(rec), (2, 43.5))
+
 
 def child_initializer():
     # Disable occupancy and implicit copy warnings in processes in a
@@ -414,3 +476,18 @@ def test_multiprocessing(self):
         finally:
             pool.close()
         self.assertEqual(res, n * (n - 1) // 2)
+
+
+@skip_on_cudasim('Simulator does not implement the CUDACodeLibrary')
+class TestCUDACodeLibrary(CUDATestCase):
+    # For tests of miscellaneous CUDACodeLibrary behaviour that we wish to
+    # explicitly check
+
+    def test_cannot_serialize_unfinalized(self):
+        # Usually a CodeLibrary requires a real CodeGen, but since we don't
+        # interact with it, anything will do
+        codegen = object()
+        name = 'library'
+        cl = CUDACodeLibrary(codegen, name)
+        with self.assertRaisesRegex(RuntimeError, 'Cannot pickle unfinalized'):
+            cl._reduce_states()