Repository: numba/numba

Issue:
CUDA: Recent commit causes infinite loop (or a lot of recompiling)
Looking into an issue using 0.51.0rc1 with cuDF, it seems that this commit: https://github.com/numba/numba/commit/e203c9d8c00590882fb50a9f0e70d09f255bb4ab causes an infinite loop, or a lot of recompilation. I'm investigating this now and will post further details here as I progress.
CUDA: Recent commit causes infinite loop (or a lot of recompiling)
Looking into an issue using 0.51.0rc1 with cuDF, it seems that this commit: https://github.com/numba/numba/commit/e203c9d8c00590882fb50a9f0e70d09f255bb4ab causes an infinite loop, or a lot of recompilation. I'm investigating this now and will post further details here as I progress.


Discussion:
The problem is with this part of the changes: https://github.com/numba/numba/pull/5709/files#diff-8ec5e1e22375684281bc2eb2e58cddd3L847-R811

```diff
     def specialize(self, *args):
         '''
         Compile and bind to the current context a version of this kernel
         specialized for the given *args*.
         Create a new instance of this dispatcher specialized for the given
         *args*.
         '''
         argtypes = tuple(
             [self.typingctx.resolve_argument_type(a) for a in args])
-        kernel = self.compile(argtypes)
-        return kernel
+        targetoptions = self.targetoptions
+        targetoptions['link'] = self.link
+        return Dispatcher(self.py_func, [types.void(*argtypes)], self._bind,
+                          targetoptions)
```

Calling a `ForAll` through `Dispatcher.forall` creates a new `Dispatcher` object, which always compiles the requested specialization. Previously, the `compile` method was called, which had access to a cache of previously-compiled specializations.
I don't have a simple reproducer or a proper fix yet, but a workaround that works for the case I have is:

```diff
diff --git a/numba/cuda/compiler.py b/numba/cuda/compiler.py
index dd57817e4..24f441bbf 100644
--- a/numba/cuda/compiler.py
+++ b/numba/cuda/compiler.py
@@ -787,6 +787,7 @@ class Dispatcher(serialize.ReduceMixin):
 
         # keyed by a `(compute capability, args)` tuple
         self.definitions = {}
+        self.specializations = {}
 
         self.targetoptions = targetoptions
 
@@ -862,12 +863,19 @@ class Dispatcher(serialize.ReduceMixin):
         Create a new instance of this dispatcher specialized for the given
         *args*.
         '''
+        cc = get_current_device().compute_capability
         argtypes = tuple(
             [self.typingctx.resolve_argument_type(a) for a in args])
+        specialization = self.specializations.get((cc, argtypes))
+        if specialization:
+            return specialization
+
         targetoptions = self.targetoptions
         targetoptions['link'] = self.link
-        return Dispatcher(self.py_func, [types.void(*argtypes)], self._bind,
-                          targetoptions)
+        specialization = Dispatcher(self.py_func, [types.void(*argtypes)],
+                                    self._bind, targetoptions)
+        self.specializations[cc, argtypes] = specialization
+        return specialization
 
     def disable_compile(self, val=True):
         self._can_compile = not val
```
The problem is with this part of the changes: https://github.com/numba/numba/pull/5709/files#diff-8ec5e1e22375684281bc2eb2e58cddd3L847-R811

```diff
     def specialize(self, *args):
         '''
         Compile and bind to the current context a version of this kernel
         specialized for the given *args*.
         Create a new instance of this dispatcher specialized for the given
         *args*.
         '''
         argtypes = tuple(
             [self.typingctx.resolve_argument_type(a) for a in args])
-        kernel = self.compile(argtypes)
-        return kernel
+        targetoptions = self.targetoptions
+        targetoptions['link'] = self.link
+        return Dispatcher(self.py_func, [types.void(*argtypes)], self._bind,
+                          targetoptions)
```

Calling a `ForAll` through `Dispatcher.forall` creates a new `Dispatcher` object, which always compiles the requested specialization. Previously, the `compile` method was called, which had access to a cache of previously-compiled specializations.
I don't have a simple reproducer or a proper fix yet, but a workaround that works for the case I have is:

```diff
diff --git a/numba/cuda/compiler.py b/numba/cuda/compiler.py
index dd57817e4..24f441bbf 100644
--- a/numba/cuda/compiler.py
+++ b/numba/cuda/compiler.py
@@ -787,6 +787,7 @@ class Dispatcher(serialize.ReduceMixin):
 
         # keyed by a `(compute capability, args)` tuple
         self.definitions = {}
+        self.specializations = {}
 
         self.targetoptions = targetoptions
 
@@ -862,12 +863,19 @@ class Dispatcher(serialize.ReduceMixin):
         Create a new instance of this dispatcher specialized for the given
         *args*.
         '''
+        cc = get_current_device().compute_capability
         argtypes = tuple(
             [self.typingctx.resolve_argument_type(a) for a in args])
+        specialization = self.specializations.get((cc, argtypes))
+        if specialization:
+            return specialization
+
         targetoptions = self.targetoptions
         targetoptions['link'] = self.link
-        return Dispatcher(self.py_func, [types.void(*argtypes)], self._bind,
-                          targetoptions)
+        specialization = Dispatcher(self.py_func, [types.void(*argtypes)],
+                                    self._bind, targetoptions)
+        self.specializations[cc, argtypes] = specialization
+        return specialization
 
     def disable_compile(self, val=True):
         self._can_compile = not val
```

Resolving patch:
diff --git a/numba/cuda/compiler.py b/numba/cuda/compiler.py
--- a/numba/cuda/compiler.py
+++ b/numba/cuda/compiler.py
@@ -366,7 +366,10 @@ def __call__(self, *args):
         if self.ntasks == 0:
             return
 
-        kernel = self.kernel.specialize(*args)
+        if self.kernel.specialized:
+            kernel = self.kernel
+        else:
+            kernel = self.kernel.specialize(*args)
         blockdim = self._compute_thread_per_block(kernel)
         griddim = (self.ntasks + blockdim - 1) // blockdim
 
@@ -787,6 +790,7 @@ def __init__(self, func, sigs, bind, targetoptions):
 
         # keyed by a `(compute capability, args)` tuple
         self.definitions = {}
+        self.specializations = {}
 
         self.targetoptions = targetoptions
 
@@ -850,7 +854,7 @@ def __call__(self, *args, **kwargs):
 
     def call(self, args, griddim, blockdim, stream, sharedmem):
         '''
-        Specialize and invoke this kernel with *args*.
+        Compile if necessary and invoke this kernel with *args*.
         '''
         argtypes = tuple(
             [self.typingctx.resolve_argument_type(a) for a in args])
@@ -862,12 +866,22 @@ def specialize(self, *args):
         Create a new instance of this dispatcher specialized for the given
         *args*.
         '''
+        cc = get_current_device().compute_capability
         argtypes = tuple(
             [self.typingctx.resolve_argument_type(a) for a in args])
+        if self.specialized:
+            raise RuntimeError('Dispatcher already specialized')
+
+        specialization = self.specializations.get((cc, argtypes))
+        if specialization:
+            return specialization
+
         targetoptions = self.targetoptions
         targetoptions['link'] = self.link
-        return Dispatcher(self.py_func, [types.void(*argtypes)], self._bind,
-                          targetoptions)
+        specialization = Dispatcher(self.py_func, [types.void(*argtypes)],
+                                    self._bind, targetoptions)
+        self.specializations[cc, argtypes] = specialization
+        return specialization
 
     def disable_compile(self, val=True):
         self._can_compile = not val


Test patch:
diff --git a/numba/cuda/tests/cudapy/test_dispatcher.py b/numba/cuda/tests/cudapy/test_dispatcher.py
new file mode 100644
--- /dev/null
+++ b/numba/cuda/tests/cudapy/test_dispatcher.py
@@ -0,0 +1,98 @@
+from numba import cuda, float32, int32
+from numba.cuda.testing import skip_on_cudasim, unittest, CUDATestCase
+
+
+@skip_on_cudasim('Dispatcher objects not used in the simulator')
+class TestDispatcher(CUDATestCase):
+    def _test_no_double_specialize(self, dispatcher, ty):
+
+        with self.assertRaises(RuntimeError) as e:
+            dispatcher.specialize(ty)
+
+        self.assertIn('Dispatcher already specialized', str(e.exception))
+
+    def test_no_double_specialize_sig_same_types(self):
+        # Attempting to specialize a kernel jitted with a signature is illegal,
+        # even for the same types the kernel is already specialized for.
+        @cuda.jit('void(float32[::1])')
+        def f(x):
+            pass
+
+        self._test_no_double_specialize(f, float32[::1])
+
+    def test_no_double_specialize_no_sig_same_types(self):
+        # Attempting to specialize an already-specialized kernel is illegal,
+        # even for the same types the kernel is already specialized for.
+        @cuda.jit
+        def f(x):
+            pass
+
+        f_specialized = f.specialize(float32[::1])
+        self._test_no_double_specialize(f_specialized, float32[::1])
+
+    def test_no_double_specialize_sig_diff_types(self):
+        # Attempting to specialize a kernel jitted with a signature is illegal.
+        @cuda.jit('void(int32[::1])')
+        def f(x):
+            pass
+
+        self._test_no_double_specialize(f, float32[::1])
+
+    def test_no_double_specialize_no_sig_diff_types(self):
+        # Attempting to specialize an already-specialized kernel is illegal.
+        @cuda.jit
+        def f(x):
+            pass
+
+        f_specialized = f.specialize(int32[::1])
+        self._test_no_double_specialize(f_specialized, float32[::1])
+
+    def test_specialize_cache_same(self):
+        # Ensure that the same dispatcher is returned for the same argument
+        # types, and that different dispatchers are returned for different
+        # argument types.
+        @cuda.jit
+        def f(x):
+            pass
+
+        self.assertEqual(len(f.specializations), 0)
+
+        f_float32 = f.specialize(float32[::1])
+        self.assertEqual(len(f.specializations), 1)
+
+        f_float32_2 = f.specialize(float32[::1])
+        self.assertEqual(len(f.specializations), 1)
+        self.assertIs(f_float32, f_float32_2)
+
+        f_int32 = f.specialize(int32[::1])
+        self.assertEqual(len(f.specializations), 2)
+        self.assertIsNot(f_int32, f_float32)
+
+    def test_specialize_cache_same_with_ordering(self):
+        # Ensure that the same dispatcher is returned for the same argument
+        # types, and that different dispatchers are returned for different
+        # argument types, taking into account array ordering and multiple
+        # arguments.
+        @cuda.jit
+        def f(x, y):
+            pass
+
+        self.assertEqual(len(f.specializations), 0)
+
+        # 'A' order specialization
+        f_f32a_f32a = f.specialize(float32[:], float32[:])
+        self.assertEqual(len(f.specializations), 1)
+
+        # 'C' order specialization
+        f_f32c_f32c = f.specialize(float32[::1], float32[::1])
+        self.assertEqual(len(f.specializations), 2)
+        self.assertIsNot(f_f32a_f32a, f_f32c_f32c)
+
+        # Reuse 'C' order specialization
+        f_f32c_f32c_2 = f.specialize(float32[::1], float32[::1])
+        self.assertEqual(len(f.specializations), 2)
+        self.assertIs(f_f32c_f32c, f_f32c_f32c_2)
+
+
+if __name__ == '__main__':
+    unittest.main()