Repository: cupy/cupy

Issue:
ROCm: Support rocFFT
Currently we don't support the FFT library on ROCm, and attempting to do FFT would return a weird OOM error (presumably due to incorrect return code overflowing):
```python
>>> import cupy
>>>
>>> a = cupy.random.random(10)
>>> cupy.fft.fft(a)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/leo/dev/cupy_rocm35/cupy/fft/fft.py", line 578, in fft
    return _fft(a, (n,), (axis,), norm, cupy.cuda.cufft.CUFFT_FORWARD)
  File "/home/leo/dev/cupy_rocm35/cupy/fft/fft.py", line 187, in _fft
    a = _fft_c2c(a, direction, norm, axes, overwrite_x, plan=plan)
  File "/home/leo/dev/cupy_rocm35/cupy/fft/fft.py", line 157, in _fft_c2c
    a = _exec_fft(a, direction, 'C2C', norm, axis, overwrite_x, plan=plan)
  File "/home/leo/dev/cupy_rocm35/cupy/fft/fft.py", line 114, in _exec_fft
    plan = cufft.Plan1d(out_size, fft_type, batch, devices=devices)
  File "cupy/cuda/cufft.pyx", line 277, in cupy.cuda.cufft.Plan1d.__init__
    self._single_gpu_get_plan(plan, nx, fft_type, batch)
  File "cupy/cuda/cufft.pyx", line 308, in cupy.cuda.cufft.Plan1d._single_gpu_get_plan
    work_area = memory.alloc(work_size)
  File "cupy/cuda/memory.pyx", line 569, in cupy.cuda.memory.alloc
    return get_allocator()(size)
  File "cupy/cuda/memory.pyx", line 1241, in cupy.cuda.memory.MemoryPool.malloc
    cpdef MemoryPointer malloc(self, size_t size):
  File "cupy/cuda/memory.pyx", line 1262, in cupy.cuda.memory.MemoryPool.malloc
    return mp.malloc(size)
  File "cupy/cuda/memory.pyx", line 930, in cupy.cuda.memory.SingleDeviceMemoryPool.malloc
    return self._malloc(rounded_size)
  File "cupy/cuda/memory.pyx", line 950, in cupy.cuda.memory.SingleDeviceMemoryPool._malloc
    mem = self._try_malloc(size)
  File "cupy/cuda/memory.pyx", line 1201, in cupy.cuda.memory.SingleDeviceMemoryPool._try_malloc
    raise OutOfMemoryError(
cupy.cuda.memory.OutOfMemoryError: Out of memory allocating 140,289,776,869,376 bytes (allocated so far: 1,024 bytes).
```
Note that while there's no package named hipFFT, it is actually there, bundled with rocFFT:
https://github.com/ROCmSoftwarePlatform/rocFFT/blob/master/library/include/hipfft.h



Discussion:

Resolving patch:
diff --git a/cupy/fft/fft.py b/cupy/fft/fft.py
--- a/cupy/fft/fft.py
+++ b/cupy/fft/fft.py
@@ -92,12 +92,28 @@ def _exec_fft(a, direction, value_type, norm, axis, overwrite_x,
         # The input array may be modified in CUDA 10.1 and above.
         # See #3763 for the discussion.
         a = a.copy()
+    elif cupy.cuda.runtime.is_hip and value_type != 'C2C':
+        # hipFFT's R2C would overwrite input
+        # hipFFT's C2R needs a workaround (see below)
+        a = a.copy()
 
     n = a.shape[-1]
     if n < 1:
         raise ValueError(
             'Invalid number of FFT data points (%d) specified.' % n)
 
+    # Workaround for hipFFT/rocFFT:
+    # Both cuFFT and hipFFT/rocFFT have this requirement that 0-th and
+    # N/2-th element must be real, but cuFFT internally simply ignores it
+    # while hipFFT handles it badly in both Plan1d and PlanNd, so we must
+    # do the correction ourselves to ensure the condition is met.
+    if cupy.cuda.runtime.is_hip and value_type == 'C2R':
+        a[..., 0] = a[..., 0].real + 0j
+        if out_size is None:
+            a[..., -1] = a[..., -1].real + 0j
+        elif out_size % 2 == 0:
+            a[..., out_size // 2] = a[...,  out_size // 2].real + 0j
+
     if out_size is None:
         out_size = n
 
@@ -138,7 +154,7 @@ def _exec_fft(a, direction, value_type, norm, axis, overwrite_x,
                              out_size, plan.nx)
         if batch != plan.batch:
             raise ValueError('Batch size does not match the plan.')
-        if config.use_multi_gpus != plan._use_multi_gpus:
+        if config.use_multi_gpus != (plan.gpus is not None):
             raise ValueError('Unclear if multiple GPUs are to be used or not.')
 
     if overwrite_x and value_type == 'C2C':
@@ -442,6 +458,10 @@ def _exec_fftn(a, direction, value_type, norm, axes, overwrite_x,
         # The input array may be modified in CUDA 10.1 and above.
         # See #3763 for the discussion.
         a = a.copy()
+    elif cupy.cuda.runtime.is_hip and value_type != 'C2C':
+        # hipFFT's R2C would overwrite input
+        # hipFFT's C2R PlanNd is actually not in use so it's fine here
+        a = a.copy()
 
     # plan search precedence:
     # 1. plan passed in as an argument
@@ -578,6 +598,17 @@ def _default_fft_func(a, s=None, axes=None, plan=None, value_type='C2C'):
 
     _, axes_sorted = _prep_fftn_axes(a.ndim, s, axes, value_type)
     if len(axes_sorted) > 1 and _nd_plan_is_possible(axes_sorted, a.ndim):
+        # circumvent two potential hipFFT/rocFFT bugs as of ROCm 3.5.0
+        # TODO(leofang): understand hipFFT better and test newer ROCm versions
+        if cupy.cuda.runtime.is_hip:
+            if (0 == axes_sorted[0] and len(axes_sorted) != a.ndim
+                    and a.flags.c_contiguous):
+                return _fft
+
+            # For C2R, we don't use PlanNd; see the workaround in _exec_fft()
+            if value_type == 'C2R':
+                return _fft
+
         # prefer Plan1D in the 1D case
         return _fftn
     return _fft
diff --git a/cupy_setup_build.py b/cupy_setup_build.py
--- a/cupy_setup_build.py
+++ b/cupy_setup_build.py
@@ -104,6 +104,7 @@
             'hip/hiprtc.h',
             'hipblas.h',
             'hiprand/hiprand.h',
+            'hipfft.h',
             'roctx.h',
             'rocsolver.h',
         ],
@@ -112,6 +113,7 @@
             'hip_hcc',
             'hipblas',
             'hiprand',
+            'rocfft',
             'roctx64',
             'rocblas',
             'rocsolver',
@@ -729,6 +731,13 @@ def cythonize(extensions, arg_options):
     cythonize_options = {key: arg_options[key]
                          for key in cythonize_option_keys}
 
+    # pass use_hip to Cython code
+    compile_time_env = cythonize_options.get('compile_time_env')
+    if compile_time_env is None:
+        compile_time_env = {}
+        cythonize_options['compile_time_env'] = compile_time_env
+    compile_time_env['use_hip'] = arg_options['use_hip']
+
     return Cython.Build.cythonize(
         extensions, verbose=True, language_level=3,
         compiler_directives=directives, **cythonize_options)
diff --git a/cupyx/scipy/fftpack/_fft.py b/cupyx/scipy/fftpack/_fft.py
--- a/cupyx/scipy/fftpack/_fft.py
+++ b/cupyx/scipy/fftpack/_fft.py
@@ -115,6 +115,8 @@ def get_fft_plan(a, shape=None, axes=None, value_type='C2C'):
     # generate plan
     # (load from cache if it exists, otherwise create one but don't cache it)
     if n > 1:  # ND transform
+        if cupy.cuda.runtime.is_hip and value_type == 'C2R':
+            raise RuntimeError("hipFFT's C2R PlanNd is buggy and unsupported")
         out_size = _get_fftn_out_size(
             shape, transformed_shape, axes[-1], value_type)
         # _get_cufft_plan_nd handles the interaction with plan cache


Test patch:
diff --git a/tests/cupy_tests/fft_tests/test_cache.py b/tests/cupy_tests/fft_tests/test_cache.py
--- a/tests/cupy_tests/fft_tests/test_cache.py
+++ b/tests/cupy_tests/fft_tests/test_cache.py
@@ -439,6 +439,7 @@ def test_LRU_cache12(self):
         assert cache0.get_curr_size() == 0 <= cache0.get_size()
         assert cache1.get_curr_size() == 0 <= cache1.get_size()
 
+    @unittest.skipIf(runtime.is_hip, "rocFFT has different plan sizes")
     def test_LRU_cache13(self):
         # test if plan insertion respect the memory size limit
         cache = config.get_plan_cache()
diff --git a/tests/cupy_tests/fft_tests/test_fft.py b/tests/cupy_tests/fft_tests/test_fft.py
--- a/tests/cupy_tests/fft_tests/test_fft.py
+++ b/tests/cupy_tests/fft_tests/test_fft.py
@@ -175,6 +175,8 @@ def _skip_multi_gpu_bug(shape, gpus):
 # 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
+@unittest.skipIf(cupy.cuda.runtime.is_hip,
+                 'hipFFT does not support multi-GPU FFT')
 @testing.parameterize(*testing.product({
     'n': [None, 0, 64],
     'shape': [(0,), (0, 10), (64,), (4, 64)],
@@ -222,6 +224,8 @@ def test_ifft(self, xp, dtype):
 # Almost identical to the TestFftOrder class, except that
 # 1. multi-GPU cuFFT is used
 # 2. the tested parameter combinations are adjusted to meet the requirements
+@unittest.skipIf(cupy.cuda.runtime.is_hip,
+                 'hipFFT does not support multi-GPU FFT')
 @testing.parameterize(*testing.product({
     'shape': [(10, 10), (10, 5, 10)],
     'data_order': ['F', 'C'],
@@ -276,7 +280,11 @@ def test_default_fft_func(self, enable_nd):
         for axes in [(0, 1), (1, 2), None, (0, 1, 2)]:
             fft_func = _default_fft_func(ca, axes=axes)
             if enable_nd:
-                assert fft_func is _fftn
+                # TODO(leofang): test newer ROCm versions
+                if axes == (0, 1) and cupy.cuda.runtime.is_hip:
+                    assert fft_func is _fft
+                else:
+                    assert fft_func is _fftn
             else:
                 assert fft_func is _fft
 
@@ -300,7 +308,11 @@ def test_default_fft_func(self, enable_nd):
                            [(-2, -1), (0, 1), None]):
             fft_func = _default_fft_func(ca, s=s, axes=axes, value_type='R2C')
             if enable_nd:
-                assert fft_func is _fftn
+                # TODO(leofang): test newer ROCm versions
+                if axes == (0, 1) and cupy.cuda.runtime.is_hip:
+                    assert fft_func is _fft
+                else:
+                    assert fft_func is _fftn
             else:
                 assert fft_func is _fft
 
@@ -313,7 +325,12 @@ def test_default_fft_func(self, enable_nd):
                            [(-2, -1), (0, 1), None]):
             fft_func = _default_fft_func(ca, s=s, axes=axes, value_type='C2R')
             if enable_nd:
-                assert fft_func is _fftn
+                # To get around hipFFT's bug, we don't use PlanNd for C2R
+                # TODO(leofang): test newer ROCm versions
+                if cupy.cuda.runtime.is_hip:
+                    assert fft_func is _fft
+                else:
+                    assert fft_func is _fftn
             else:
                 assert fft_func is _fft
 
@@ -505,6 +522,14 @@ def test_ifftn(self, xp, dtype, order, enable_nd):
 @testing.gpu
 class TestPlanCtxManagerFftn(unittest.TestCase):
 
+    def setUp(self):
+        if cupy.cuda.runtime.is_hip:
+            # TODO(leofang): test newer ROCm versions
+            if (self.axes == (0, 1) and self.shape == (2, 3, 4)):
+                raise unittest.SkipTest("hipFFT's PlanNd for this case "
+                                        "is buggy, so Plan1d is generated "
+                                        "instead")
+
     @nd_planning_states()
     @testing.for_complex_dtypes()
     @testing.numpy_cupy_allclose(rtol=1e-4, atol=1e-7, accept_error=ValueError,
@@ -650,6 +675,8 @@ def test_fft_error_on_wrong_plan(self, dtype):
 # Almost identical to the TestPlanCtxManagerFft class, except that
 # 1. multi-GPU cuFFT is used
 # 2. the tested parameter combinations are adjusted to meet the requirements
+@unittest.skipIf(cupy.cuda.runtime.is_hip,
+                 'hipFFT does not support multi-GPU FFT')
 @testing.parameterize(*testing.product({
     'n': [None, 64],
     'shape': [(64,), (128,)],
@@ -843,6 +870,7 @@ def test_rfft(self, xp, dtype):
                                  contiguous_check=False)
     def test_irfft(self, xp, dtype):
         a = testing.shaped_random(self.shape, xp, dtype)
+
         if xp is cupy:
             from cupyx.scipy.fftpack import get_fft_plan
             shape = (self.n,) if self.n is not None else None
@@ -1040,6 +1068,14 @@ def test_irfftn(self, xp, dtype, order, enable_nd):
 @testing.gpu
 class TestPlanCtxManagerRfftn(unittest.TestCase):
 
+    def setUp(self):
+        if cupy.cuda.runtime.is_hip:
+            # TODO(leofang): test newer ROCm versions
+            if (self.axes == (0, 1) and self.shape == (2, 3, 4)):
+                raise unittest.SkipTest("hipFFT's PlanNd for this case "
+                                        "is buggy, so Plan1d is generated "
+                                        "instead")
+
     @nd_planning_states()
     @testing.for_all_dtypes(no_complex=True)
     @testing.numpy_cupy_allclose(rtol=1e-4, atol=1e-7, accept_error=ValueError,
@@ -1060,6 +1096,8 @@ def test_rfftn(self, xp, dtype, enable_nd):
 
         return out
 
+    @unittest.skipIf(cupy.cuda.runtime.is_hip,
+                     "hipFFT's PlanNd for C2R is buggy")
     @nd_planning_states()
     @testing.for_all_dtypes()
     @testing.numpy_cupy_allclose(rtol=1e-4, atol=1e-7, accept_error=ValueError,
diff --git a/tests/cupyx_tests/scipy_tests/fft_tests/test_fft.py b/tests/cupyx_tests/scipy_tests/fft_tests/test_fft.py
--- a/tests/cupyx_tests/scipy_tests/fft_tests/test_fft.py
+++ b/tests/cupyx_tests/scipy_tests/fft_tests/test_fft.py
@@ -546,7 +546,7 @@ def test_ifft2_backend_plan(self, xp, dtype):
     + testing.product({
         'shape': [(2, 3, 4)],
         's': [None, (1, 5), (1, 4, 10)],
-        'axes': [None, (-2, -1), (-1, -2, -3)],
+        'axes': [None, (0, 1), (-2, -1), (-1, -2, -3)],
         'norm': [None, 'ortho']
     })
     + testing.product({
@@ -970,6 +970,14 @@ def test_irfft_backend(self, xp, dtype):
         return _correct_np_dtype(xp, dtype, out)
 
 
+def _skip_hipFFT_PlanNd_bug(axes, shape):
+    if cp.cuda.runtime.is_hip:
+        # TODO(leofang): test newer ROCm versions
+        if (axes == (0, 1) and shape == (2, 3, 4)):
+            raise unittest.SkipTest("hipFFT's PlanNd for this case is buggy, "
+                                    "so Plan1d is generated instead")
+
+
 @testing.parameterize(
     {'shape': (3, 4), 's': None, 'axes': None, 'norm': None},
     {'shape': (3, 4), 's': (1, 5), 'axes': None, 'norm': None},
@@ -1014,6 +1022,7 @@ def test_rfft2_overwrite(self, xp, dtype):
     @testing.numpy_cupy_allclose(rtol=1e-4, atol=1e-7, accept_error=ValueError,
                                  contiguous_check=False)
     def test_rfft2_plan(self, xp, dtype):
+        _skip_hipFFT_PlanNd_bug(self.axes, self.shape)
         x = testing.shaped_random(self.shape, xp, dtype)
         x_orig = x.copy()
 
@@ -1037,6 +1046,7 @@ def test_rfft2_plan(self, xp, dtype):
     @testing.numpy_cupy_allclose(rtol=1e-4, atol=1e-7, accept_error=ValueError,
                                  contiguous_check=False)
     def test_rfft2_overwrite_plan(self, xp, dtype):
+        _skip_hipFFT_PlanNd_bug(self.axes, self.shape)
         x = testing.shaped_random(self.shape, xp, dtype)
 
         # hack: skip testing if getting a cuFFT plan is impossible
@@ -1058,6 +1068,7 @@ def test_rfft2_overwrite_plan(self, xp, dtype):
     @testing.numpy_cupy_allclose(rtol=1e-4, atol=1e-7, accept_error=ValueError,
                                  contiguous_check=False)
     def test_rfft2_plan_manager(self, xp, dtype):
+        _skip_hipFFT_PlanNd_bug(self.axes, self.shape)
         x = testing.shaped_random(self.shape, xp, dtype)
         x_orig = x.copy()
 
@@ -1122,6 +1133,8 @@ def test_irfft2_overwrite(self, xp, dtype):
 
     @pytest.mark.skipif(int(cp.cuda.device.get_compute_capability()) < 70,
                         reason="Known to fail with Pascal or older")
+    @unittest.skipIf(cp.cuda.runtime.is_hip,
+                     "hipFFT's PlanNd for C2R is buggy")
     @testing.for_all_dtypes()
     @testing.numpy_cupy_allclose(rtol=1e-4, atol=1e-7, accept_error=ValueError,
                                  contiguous_check=False)
@@ -1148,6 +1161,8 @@ def test_irfft2_plan(self, xp, dtype):
 
     @pytest.mark.skipif(int(cp.cuda.device.get_compute_capability()) < 70,
                         reason="Known to fail with Pascal or older")
+    @unittest.skipIf(cp.cuda.runtime.is_hip,
+                     "hipFFT's PlanNd for C2R is buggy")
     @testing.for_all_dtypes()
     @testing.numpy_cupy_allclose(rtol=1e-4, atol=1e-7, accept_error=ValueError,
                                  contiguous_check=False)
@@ -1171,6 +1186,8 @@ def test_irfft2_overwrite_plan(self, xp, dtype):
 
     @pytest.mark.skipif(int(cp.cuda.device.get_compute_capability()) < 70,
                         reason="Known to fail with Pascal or older")
+    @unittest.skipIf(cp.cuda.runtime.is_hip,
+                     "hipFFT's PlanNd for C2R is buggy")
     @testing.for_all_dtypes()
     @testing.numpy_cupy_allclose(rtol=1e-4, atol=1e-7, accept_error=ValueError,
                                  contiguous_check=False)
@@ -1258,6 +1275,7 @@ def test_rfftn_overwrite(self, xp, dtype):
     @testing.numpy_cupy_allclose(rtol=1e-4, atol=1e-7, accept_error=ValueError,
                                  contiguous_check=False)
     def test_rfftn_plan(self, xp, dtype):
+        _skip_hipFFT_PlanNd_bug(self.axes, self.shape)
         x = testing.shaped_random(self.shape, xp, dtype)
         x_orig = x.copy()
 
@@ -1281,6 +1299,7 @@ def test_rfftn_plan(self, xp, dtype):
     @testing.numpy_cupy_allclose(rtol=1e-4, atol=1e-7, accept_error=ValueError,
                                  contiguous_check=False)
     def test_rfftn_overwrite_plan(self, xp, dtype):
+        _skip_hipFFT_PlanNd_bug(self.axes, self.shape)
         x = testing.shaped_random(self.shape, xp, dtype)
 
         # hack: skip testing if getting a cuFFT plan is impossible
@@ -1302,6 +1321,7 @@ def test_rfftn_overwrite_plan(self, xp, dtype):
     @testing.numpy_cupy_allclose(rtol=1e-4, atol=1e-7, accept_error=ValueError,
                                  contiguous_check=False)
     def test_rfftn_plan_manager(self, xp, dtype):
+        _skip_hipFFT_PlanNd_bug(self.axes, self.shape)
         x = testing.shaped_random(self.shape, xp, dtype)
         x_orig = x.copy()
 
@@ -1366,6 +1386,8 @@ def test_irfftn_overwrite(self, xp, dtype):
 
     @pytest.mark.skipif(int(cp.cuda.device.get_compute_capability()) < 70,
                         reason="Known to fail with Pascal or older")
+    @unittest.skipIf(cp.cuda.runtime.is_hip,
+                     "hipFFT's PlanNd for C2R is buggy")
     @testing.for_all_dtypes()
     @testing.numpy_cupy_allclose(rtol=1e-4, atol=1e-7, accept_error=ValueError,
                                  contiguous_check=False)
@@ -1392,6 +1414,8 @@ def test_irfftn_plan(self, xp, dtype):
 
     @pytest.mark.skipif(int(cp.cuda.device.get_compute_capability()) < 70,
                         reason="Known to fail with Pascal or older")
+    @unittest.skipIf(cp.cuda.runtime.is_hip,
+                     "hipFFT's PlanNd for C2R is buggy")
     @testing.for_all_dtypes()
     @testing.numpy_cupy_allclose(rtol=1e-4, atol=1e-7, accept_error=ValueError,
                                  contiguous_check=False)
@@ -1415,6 +1439,8 @@ def test_irfftn_overwrite_plan(self, xp, dtype):
 
     @pytest.mark.skipif(int(cp.cuda.device.get_compute_capability()) < 70,
                         reason="Known to fail with Pascal or older")
+    @unittest.skipIf(cp.cuda.runtime.is_hip,
+                     "hipFFT's PlanNd for C2R is buggy")
     @testing.for_all_dtypes()
     @testing.numpy_cupy_allclose(rtol=1e-4, atol=1e-7, accept_error=ValueError,
                                  contiguous_check=False)