Repository: cupy/cupy
Pull request: Fix CUDA 11 multi-GPU FFT bug
Close #3757. See https://github.com/cupy/cupy/issues/3757#issuecomment-673184478 for the trigger condition. 

cc: @kmaehashi @pentschev @anaruse @maxpkatz

Problem statement:
CUDA 11 Test: `cupy_tests.fft_tests.test_fft.TestMultiGpu`
```
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuFft_param_2_{n=None, norm=None, shape=(64,)}::test_fft - AssertionError: Parameterized test failed.
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuFft_param_2_{n=None, norm=None, shape=(64,)}::test_ifft - AssertionError: Parameterized test failed.
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuFft_param_6_{n=None, norm='ortho', shape=(64,)}::test_fft - AssertionError: Parameterized test failed.
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuFft_param_6_{n=None, norm='ortho', shape=(64,)}::test_ifft - AssertionError: Parameterized test failed.
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuFft_param_24_{n=64, norm=None, shape=(0,)}::test_fft - AssertionError: Parameterized test failed.
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuFft_param_24_{n=64, norm=None, shape=(0,)}::test_ifft - AssertionError: Parameterized test failed.
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuFft_param_26_{n=64, norm=None, shape=(64,)}::test_fft - AssertionError: Parameterized test failed.
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuFft_param_26_{n=64, norm=None, shape=(64,)}::test_ifft - AssertionError: Parameterized test failed.
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuFft_param_28_{n=64, norm='ortho', shape=(0,)}::test_fft - AssertionError: Parameterized test failed.
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuFft_param_28_{n=64, norm='ortho', shape=(0,)}::test_ifft - AssertionError: Parameterized test failed.
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuFft_param_30_{n=64, norm='ortho', shape=(64,)}::test_fft - AssertionError: Parameterized test failed.
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuFft_param_30_{n=64, norm='ortho', shape=(64,)}::test_ifft - AssertionError: Parameterized test failed.
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuPlanCtxManagerFft_param_0_{n=None, norm=None, shape=(64,)}::test_fft - AssertionError: Parameterized test f...
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuPlanCtxManagerFft_param_0_{n=None, norm=None, shape=(64,)}::test_ifft - AssertionError: Parameterized test ...
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuPlanCtxManagerFft_param_1_{n=None, norm=None, shape=(128,)}::test_fft - AssertionError: Parameterized test ...
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuPlanCtxManagerFft_param_1_{n=None, norm=None, shape=(128,)}::test_ifft - AssertionError: Parameterized test...
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuPlanCtxManagerFft_param_2_{n=None, norm='ortho', shape=(64,)}::test_fft - AssertionError: Parameterized tes...
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuPlanCtxManagerFft_param_2_{n=None, norm='ortho', shape=(64,)}::test_ifft - AssertionError: Parameterized te...
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuPlanCtxManagerFft_param_3_{n=None, norm='ortho', shape=(128,)}::test_fft - AssertionError: Parameterized te...
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuPlanCtxManagerFft_param_3_{n=None, norm='ortho', shape=(128,)}::test_ifft - AssertionError: Parameterized t...
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuPlanCtxManagerFft_param_4_{n=64, norm=None, shape=(64,)}::test_fft - AssertionError: Parameterized test fai...
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuPlanCtxManagerFft_param_4_{n=64, norm=None, shape=(64,)}::test_ifft - AssertionError: Parameterized test fa...
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuPlanCtxManagerFft_param_5_{n=64, norm=None, shape=(128,)}::test_fft - AssertionError: Parameterized test fa...
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuPlanCtxManagerFft_param_5_{n=64, norm=None, shape=(128,)}::test_ifft - AssertionError: Parameterized test f...
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuPlanCtxManagerFft_param_6_{n=64, norm='ortho', shape=(64,)}::test_fft - AssertionError: Parameterized test ...
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuPlanCtxManagerFft_param_6_{n=64, norm='ortho', shape=(64,)}::test_ifft - AssertionError: Parameterized test...
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuPlanCtxManagerFft_param_7_{n=64, norm='ortho', shape=(128,)}::test_fft - AssertionError: Parameterized test...
FAILED tests/cupy_tests/fft_tests/test_fft.py::TestMultiGpuPlanCtxManagerFft_param_7_{n=64, norm='ortho', shape=(128,)}::test_ifft - AssertionError: Parameterized tes...
```

Hints from issue discussion:
Multi-gpu cuFFT failing with `CUFFT_EXEC_FAILED` for complex64.

cc/ @pentschev 
We should suggest the cuFFT team to include CuPy's extensive FFT tests as part of their CI...It seems to me they keep breaking the behaviors/performance/precision/etc in every release...A better documentation is the minimal requirement.
I don't have access to any machine with CUDA 11 yet, but tend to think CUDA 11 messed up multi-GPU FFT. In one of the failing cases, the parameter is clearly supported in all previous versions:
- batch = 1
- nx = 128 
- fft type = CUFFT_C2C

See https://docs.nvidia.com/cuda/cufft/index.html#multiple-GPU-cufft-supported-functionality

If we don't get prompt replies from the cuFFT team, how about just skipping the tests and/or raising a warning if CUDA 11 + multi-GPU is in use?
I've been looking at this today and I noticed the failing tests are only those running on devices `1,0`, those running on `0,1` all pass for me. I tried to hack the `7_CUDALibraries/simpleCUFFT_MGPU` sample a bit to try multiple combinations of GPUs and I can't get it to fail based solely on that. Does anyone have an idea why it would only fail for devices `1,0`? Is there something in CuPy that may make the device ordering relevant for this case here?
@pentschev That narrows down the search quite a bit! What you found reminds me a bug I discovered when working on the multi-GPU support:
https://github.com/cupy/cupy/blob/99a0214073f372263b28b49acabcaa55a290b6fc/cupy/cuda/cufft.pyx#L427-L431
I have a wild theory that they fixed the bug in CUDA 11 by requiring that when a multi-GPU plan is launched, the present device context must be the same as the first one on the list. So, we fails because when we launch a plan for `1, 0`, we're still on device `0`. Peter, could you check if it's the case? 
My wild guess is wrong (I got a CUDA 11 env set up). Will try to look into that today or tomorrow. Could be that they updated how `cudaXtDesc`/`cudaLibXtDesc` is handled, which we manage ourselves.
With the tests added in #3766 I am able to further narrow it down. The condition to trigger the error is 
- batch = 1
- ~~C2C (complex64 to complex64; Z2Z is fine)~~ **UPDATE:** C2C/Z2Z
- GPU = 1, 0
- Any legitimate nx (tried 64 & 128)

Will report back if I find out more.
@pentschev You need to try my test case above. It's easy to tweak the example code `simpleCUFFT_MGPU.cu` (change `SIGNAL_SIZE` to `122`, which will be padded to `128`, and then `whichGPUs` to  `1, 0`). It's a regression on the cuFFT side:
```
$ ./simpleCUFFT_MGPU 

[simpleCUFFT_MGPU] is starting...

GPU Device 0: "GeForce RTX 2080 Ti" with compute capability 7.5
GPU Device 1: "GeForce RTX 2080 Ti" with compute capability 7.5
new size: 128

Running on GPUs
GPU Device 1: "GeForce RTX 2080 Ti" with compute capability 7.5
GPU Device 0: "GeForce RTX 2080 Ti" with compute capability 7.5
CUDA error at simpleCUFFT_MGPU.cu:194 code=6(CUFFT_EXEC_FAILED) "cufftXtExecDescriptorC2C(plan_input, d_signal, d_signal, CUFFT_FORWARD)"
```
cc: @anaruse 
> * C2C (complex64 to complex64; Z2Z is fine)

Sorry, Z2Z is not fine. It's also buggy.

Commit message: add ref and fix flake8

Patch to review:
diff --git a/cupy/testing/parameterized.py b/cupy/testing/parameterized.py
index a09123e6f8f..5efe58333c4 100644
--- a/cupy/testing/parameterized.py
+++ b/cupy/testing/parameterized.py
@@ -6,6 +6,7 @@
 import unittest
 
 from cupy.testing import _bundle
+from cupy.cuda import cufft
 from cupy.cuda import driver
 
 
@@ -91,7 +92,7 @@ def new_method(self, *args, **kwargs):
                 for k, v in sorted(param.items()):
                     s.write('  {}: {}\n'.format(k, v))
                 err_class = e.__class__
-                if isinstance(e, driver.CUDADriverError):
+                if isinstance(e, (driver.CUDADriverError, cufft.CuFFTError)):
                     err_class, = err_class.__bases__
                 raise err_class(s.getvalue()).with_traceback(e.__traceback__)
         return new_method
diff --git a/tests/cupy_tests/fft_tests/test_fft.py b/tests/cupy_tests/fft_tests/test_fft.py
index 8dbfbcd692c..c5a02fa3d9a 100644
--- a/tests/cupy_tests/fft_tests/test_fft.py
+++ b/tests/cupy_tests/fft_tests/test_fft.py
@@ -53,7 +53,7 @@ def multi_gpu_config(gpu_configs=None):
     """Decorator for parameterized tests with different GPU configurations.
 
     Args:
-        gpu_configs(list of list): The GPUs to test.
+        gpu_configs (list of list): The GPUs to test.
 
     .. notes:
         The decorated tests are skipped if no or only one GPU is available.
@@ -71,6 +71,7 @@ def test_func(self, *args, **kw):
                         assert nGPUs >= 2, 'Must use at least two gpus'
                         config.use_multi_gpus = True
                         config.set_cufft_gpus(gpus)
+                        self.gpus = gpus
 
                         impl(self, *args, **kw)
                     except Exception:
@@ -79,6 +80,7 @@ def test_func(self, *args, **kw):
             finally:
                 config.use_multi_gpus = use_multi_gpus
                 config._devices = _devices
+                del self.gpus
 
         return test_func
     return decorator
@@ -159,6 +161,17 @@ def test_ifft(self, xp, dtype):
         return out
 
 
+# See #3757 and NVIDIA internal ticket 3093094
+def _skip_multi_gpu_bug(shape, gpus):
+    # avoid CUDA 11 bug triggered by
+    # - batch = 1
+    # - gpus = [1, 0]
+    if (cupy.cuda.runtime.runtimeGetVersion() == 11000
+            and len(shape) == 1
+            and gpus == [1, 0]):
+        raise unittest.SkipTest('avoid CUDA 11 bug')
+
+
 # Almost identical to the TestFft class, except that
 # 1. multi-GPU cuFFT is used
 # 2. the tested parameter combinations are adjusted to meet the requirements
@@ -169,11 +182,14 @@ def test_ifft(self, xp, dtype):
 }))
 @testing.multi_gpu(2)
 class TestMultiGpuFft(unittest.TestCase):
+
     @multi_gpu_config(gpu_configs=[[0, 1], [1, 0]])
     @testing.for_complex_dtypes()
     @testing.numpy_cupy_allclose(rtol=1e-4, atol=1e-7, accept_error=ValueError,
                                  contiguous_check=False)
     def test_fft(self, xp, dtype):
+        _skip_multi_gpu_bug(self.shape, self.gpus)
+
         a = testing.shaped_random(self.shape, xp, dtype)
         out = xp.fft.fft(a, n=self.n, norm=self.norm)
 
@@ -191,6 +207,8 @@ def test_fft(self, xp, dtype):
     @testing.with_requires('numpy!=1.17.0')
     @testing.with_requires('numpy!=1.17.1')
     def test_ifft(self, xp, dtype):
+        _skip_multi_gpu_bug(self.shape, self.gpus)
+
         a = testing.shaped_random(self.shape, xp, dtype)
         out = xp.fft.ifft(a, n=self.n, norm=self.norm)
 
@@ -216,6 +234,8 @@ class TestMultiGpuFftOrder(unittest.TestCase):
     @testing.numpy_cupy_allclose(rtol=1e-4, atol=1e-7, accept_error=ValueError,
                                  contiguous_check=False)
     def test_fft(self, xp, dtype):
+        _skip_multi_gpu_bug(self.shape, self.gpus)
+
         a = testing.shaped_random(self.shape, xp, dtype)
         if self.data_order == 'F':
             a = xp.asfortranarray(a)
@@ -232,6 +252,8 @@ def test_fft(self, xp, dtype):
     @testing.numpy_cupy_allclose(rtol=1e-4, atol=1e-7, accept_error=ValueError,
                                  contiguous_check=False)
     def test_ifft(self, xp, dtype):
+        _skip_multi_gpu_bug(self.shape, self.gpus)
+
         a = testing.shaped_random(self.shape, xp, dtype)
         if self.data_order == 'F':
             a = xp.asfortranarray(a)
@@ -640,6 +662,8 @@ class TestMultiGpuPlanCtxManagerFft(unittest.TestCase):
     @testing.numpy_cupy_allclose(rtol=1e-4, atol=1e-7, accept_error=ValueError,
                                  contiguous_check=False)
     def test_fft(self, xp, dtype):
+        _skip_multi_gpu_bug(self.shape, self.gpus)
+
         a = testing.shaped_random(self.shape, xp, dtype)
         if xp is cupy:
             from cupyx.scipy.fftpack import get_fft_plan
@@ -662,6 +686,8 @@ def test_fft(self, xp, dtype):
     @testing.numpy_cupy_allclose(rtol=1e-4, atol=1e-7, accept_error=ValueError,
                                  contiguous_check=False)
     def test_ifft(self, xp, dtype):
+        _skip_multi_gpu_bug(self.shape, self.gpus)
+
         a = testing.shaped_random(self.shape, xp, dtype)
         if xp is cupy:
             from cupyx.scipy.fftpack import get_fft_plan


Code review comments:
diff_hunk: @@ -159,6 +161,17 @@ def test_ifft(self, xp, dtype):
         return out
 
 
+# See #3757 and NVIDIA internal ticket 3093094
+def _skip_multi_gpu_bug(shape, gpus):
+    # avoid CUDA 11 bug triggered by
path: tests/cupy_tests/fft_tests/test_fft.py
text: @user1:
```suggestion
    # avoid CUDA 11.0 (will be fixed by CUDA 11.2) bug triggered by
```

diff_hunk: @@ -159,6 +161,17 @@ def test_ifft(self, xp, dtype):
         return out
 
 
+# See #3757 and NVIDIA internal ticket 3093094
+def _skip_multi_gpu_bug(shape, gpus):
+    # avoid CUDA 11 bug triggered by
+    # - batch = 1
+    # - gpus = [1, 0]
+    if (cupy.cuda.runtime.runtimeGetVersion() == 11000
path: tests/cupy_tests/fft_tests/test_fft.py
text: @user1:
```suggestion
    if (cupy.cuda.runtime.runtimeGetVersion() < 11200
```

@user1:
Originally I forgot to change the comparison operator, which I fixed now.

@author:
I think you meant this?
```suggestion
    if (11000 <= cupy.cuda.runtime.runtimeGetVersion() < 11200
```

@user1:
Ah yes, totally missed that, thanks for fixing!