Repository: cupy/cupy
Pull request: Support 0-size ndarray and fix possible error in __del__ at fft
Fix #3241


Problem statement:
`Invalid number of FFT data points` is not always raised
`cupy.fft.fft(cupy.array([]))` should raise `ValueError` even if `n` is default.

<details>

```
>>> np.fft.fft(np.array([]))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<__array_function__ internals>", line 5, in fft
  File "/path/to/python3.8/site-packages/numpy-1.18.2-py3.8-linux-x86_64.egg/numpy/fft/_pocketfft.py", line 188, in fft
    output = _raw_fft(a, n, axis, False, True, inv_norm)
  File "/path/to/python3.8/site-packages/numpy-1.18.2-py3.8-linux-x86_64.egg/numpy/fft/_pocketfft.py", line 57, in _raw_fft
    raise ValueError("Invalid number of FFT data points (%d) specified."
ValueError: Invalid number of FFT data points (0) specified.
>>> cupy.fft.fft(cupy.array([]))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/path/to/cupy/cupy/fft/fft.py", line 496, in fft
    return _fft(a, (n,), (axis,), norm, cupy.cuda.cufft.CUFFT_FORWARD)
  File "/path/to/cupy/cupy/fft/fft.py", line 175, in _fft
    a = _fft_c2c(a, direction, norm, axes, overwrite_x, plan=plan)
  File "/path/to/cupy/cupy/fft/fft.py", line 142, in _fft_c2c
    a = _exec_fft(a, direction, 'C2C', norm, axis, overwrite_x, plan=plan)
  File "/path/to/cupy/cupy/fft/fft.py", line 91, in _exec_fft
    batch = a.size // a.shape[-1]
ZeroDivisionError: integer division or modulo by zero
>>> cupy.fft.fft(cupy.array([]), n=0)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/path/to/cupy/cupy/fft/fft.py", line 496, in fft
    return _fft(a, (n,), (axis,), norm, cupy.cuda.cufft.CUFFT_FORWARD)
  File "/path/to/cupy/cupy/fft/fft.py", line 155, in _fft
    raise ValueError(
ValueError: Invalid number of FFT data points (0) specified.
```

Hints from issue discussion:

Commit message: Support 0-size ndarray at fft

Patch to review:
diff --git a/cupy/cuda/cufft.pyx b/cupy/cuda/cufft.pyx
index db12bc52b85..83d7474473d 100644
--- a/cupy/cuda/cufft.pyx
+++ b/cupy/cuda/cufft.pyx
@@ -252,17 +252,32 @@ class Plan1d(object):
         cdef Handle plan
         cdef bint use_multi_gpus = 0 if devices is None else 1
 
+        self.plan = 0
+        self.xtArr = <intptr_t>0  # pointer to metadata for multi-GPU buffer
+        self.xtArr_buffer = None  # actual multi-GPU intermediate buffer
+
         with nogil:
             result = cufftCreate(&plan)
             if result == 0:
                 result = cufftSetAutoAllocation(plan, 0)
         check_result(result)
 
-        # set plan, work_area, gpus, streams, and events
-        if not use_multi_gpus:
-            self._single_gpu_get_plan(plan, nx, fft_type, batch)
-        else:
-            self._multi_gpu_get_plan(plan, nx, fft_type, batch, devices, out)
+        self.plan = plan
+        self.work_area = None
+        self.gpus = None
+
+        self.gather_streams = None
+        self.gather_events = None
+        self.scatter_streams = None
+        self.scatter_events = None
+
+        if batch != 0:
+            # set plan, work_area, gpus, streams, and events
+            if not use_multi_gpus:
+                self._single_gpu_get_plan(plan, nx, fft_type, batch)
+            else:
+                self._multi_gpu_get_plan(
+                    plan, nx, fft_type, batch, devices, out)
 
         self.nx = nx
         self.fft_type = fft_type
@@ -270,8 +285,6 @@ class Plan1d(object):
 
         self._use_multi_gpus = use_multi_gpus
         self.batch_share = None
-        self.xtArr = <intptr_t>0  # pointer to metadata for multi-GPU buffer
-        self.xtArr_buffer = None  # actual multi-GPU intermediate buffer
 
     def _single_gpu_get_plan(self, Handle plan, int nx, int fft_type,
                              int batch):
@@ -298,14 +311,7 @@ class Plan1d(object):
             result = cufftSetWorkArea(plan, <void*>(ptr))
         check_result(result)
 
-        self.plan = plan
         self.work_area = work_area  # this is for cuFFT plan
-        self.gpus = None
-
-        self.gather_streams = None
-        self.gather_events = None
-        self.scatter_streams = None
-        self.scatter_events = None
 
     def _multi_gpu_get_plan(self, Handle plan, int nx, int fft_type, int batch,
                             devices, out):
@@ -378,7 +384,6 @@ class Plan1d(object):
             result = cufftXtSetWorkArea(plan, work_area_ptr.data())
         check_result(result)
 
-        self.plan = plan
         self.work_area = work_area  # this is for cuFFT plan
         self.gpus = list(gpus)
 
@@ -413,9 +418,11 @@ class Plan1d(object):
         cdef int dev = runtime.getDevice()
         cdef int result
 
-        with nogil:
-            result = cufftDestroy(plan)
-        check_result(result)
+        if plan != 0:
+            with nogil:
+                result = cufftDestroy(plan)
+            check_result(result)
+            self.plan = 0
 
         # cuFFT bug: after cufftDestroy(), the current device is mistakenly
         # set to the last device in self.gpus, so we must correct it. See
@@ -695,6 +702,8 @@ class PlanNd(object):
         cdef int* shape_ptr = shape_arr.data()
         cdef int* inembed_ptr
         cdef int* onembed_ptr
+        self.plan = 0
+
         ndim = len(shape)
 
         if inembed is None:
@@ -713,23 +722,28 @@ class PlanNd(object):
             result = cufftCreate(&plan)
             if result == 0:
                 result = cufftSetAutoAllocation(plan, 0)
-            if result == 0:
-                result = cufftMakePlanMany(plan, ndim, shape_ptr,
-                                           inembed_ptr, istride, idist,
-                                           onembed_ptr, ostride, odist,
-                                           <Type>fft_type, batch,
-                                           &work_size)
+        check_result(result)
 
-        # cufftMakePlanMany could use a large amount of memory
-        if result == 2:
-            cupy.get_default_memory_pool().free_all_blocks()
+        if batch == 0:
+            work_size = 0
+        else:
             with nogil:
                 result = cufftMakePlanMany(plan, ndim, shape_ptr,
                                            inembed_ptr, istride, idist,
                                            onembed_ptr, ostride, odist,
                                            <Type>fft_type, batch,
                                            &work_size)
-        check_result(result)
+
+            # cufftMakePlanMany could use a large amount of memory
+            if result == 2:
+                cupy.get_default_memory_pool().free_all_blocks()
+                with nogil:
+                    result = cufftMakePlanMany(plan, ndim, shape_ptr,
+                                               inembed_ptr, istride, idist,
+                                               onembed_ptr, ostride, odist,
+                                               <Type>fft_type, batch,
+                                               &work_size)
+            check_result(result)
 
         # TODO: for CUDA>=9.2 could also allow setting a work area policy
         # result = cufftXtSetWorkAreaPolicy(plan, policy, &work_size)
@@ -749,9 +763,11 @@ class PlanNd(object):
 
     def __del__(self):
         cdef Handle plan = self.plan
-        with nogil:
-            result = cufftDestroy(plan)
-        check_result(result)
+        if plan != 0:
+            with nogil:
+                result = cufftDestroy(plan)
+            check_result(result)
+            self.plan = 0
 
     def __enter__(self):
         _thread_local._current_plan = self
diff --git a/cupy/fft/fft.py b/cupy/fft/fft.py
index f8c3df493e5..02c3f89d974 100644
--- a/cupy/fft/fft.py
+++ b/cupy/fft/fft.py
@@ -87,10 +87,16 @@ def _exec_fft(a, direction, value_type, norm, axis, overwrite_x,
     if a.base is not None or not a.flags.c_contiguous:
         a = a.copy()
 
+    n = a.shape[-1]
+    if n < 1:
+        raise ValueError(
+            'Invalid number of FFT data points (%d) specified.' % n)
+
     if out_size is None:
-        out_size = a.shape[-1]
+        out_size = n
+
+    batch = a.size // n
 
-    batch = a.size // a.shape[-1]
     curr_plan = cufft.get_current_plan()
     if curr_plan is not None:
         if plan is None:
@@ -123,11 +129,12 @@ def _exec_fft(a, direction, value_type, norm, axis, overwrite_x,
     else:
         out = plan.get_output_array(a)
 
-    plan.fft(a, out, direction)
+    if batch != 0:
+        plan.fft(a, out, direction)
 
     sz = out.shape[-1]
     if fft_type == cufft.CUFFT_R2C or fft_type == cufft.CUFFT_D2Z:
-        sz = a.shape[-1]
+        sz = n
     if norm is None:
         if direction == cufft.CUFFT_INVERSE:
             out /= sz
@@ -152,12 +159,6 @@ def _fft(a, s, axes, norm, direction, value_type='C2C', overwrite_x=False,
         raise ValueError('Invalid norm value %s, should be None or "ortho".'
                          % norm)
 
-    if s is not None:
-        for n in s:
-            if (n is not None) and (n < 1):
-                raise ValueError(
-                    'Invalid number of FFT data points (%d) specified.' % n)
-
     if (s is not None) and (axes is not None) and len(s) != len(axes):
         raise ValueError('Shape and axes have different lengths.')
 
@@ -367,6 +368,11 @@ def _get_cufft_plan_nd(shape, fft_type, axes=None, order='C', out_size=None):
                 'GPU case (Can only batch FFT over the first or last '
                 'spatial axes).')
 
+    for n in plan_dimensions:
+        if n < 1:
+            raise ValueError(
+                'Invalid number of FFT data points specified.')
+
     plan = cufft.PlanNd(shape=plan_dimensions,
                         inembed=inembed,
                         istride=istride,
@@ -451,7 +457,9 @@ def _exec_fftn(a, direction, value_type, norm, axes, overwrite_x,
         out = plan.get_output_array(a, order=order)
     else:
         plan.check_output_array(a, out)
-    plan.fft(a, out, direction)
+
+    if out.size != 0:
+        plan.fft(a, out, direction)
 
     # normalize by the product of the shape along the transformed axes
     arr = a if fft_type in (cufft.CUFFT_R2C, cufft.CUFFT_D2Z) else out
@@ -493,6 +501,12 @@ def _fftn(a, s, axes, norm, direction, value_type='C2C', order='A', plan=None,
     # Note: need to call _cook_shape prior to sorting the axes
     a = _cook_shape(a, s, axes, value_type, order=order)
 
+    print(a.shape, s, axes_sorted)
+    for n in a.shape:
+        if n < 1:
+            raise ValueError(
+                'Invalid number of FFT data points (%d) specified.' % n)
+
     if order == 'C' and not a.flags.c_contiguous:
         a = cupy.ascontiguousarray(a)
     elif order == 'F' and not a.flags.f_contiguous:
diff --git a/tests/cupy_tests/fft_tests/test_fft.py b/tests/cupy_tests/fft_tests/test_fft.py
index 04bf29c9561..8dbfbcd692c 100644
--- a/tests/cupy_tests/fft_tests/test_fft.py
+++ b/tests/cupy_tests/fft_tests/test_fft.py
@@ -86,7 +86,7 @@ def test_func(self, *args, **kw):
 
 @testing.parameterize(*testing.product({
     'n': [None, 0, 5, 10, 15],
-    'shape': [(10,), (10, 10)],
+    'shape': [(0,), (10, 0), (10,), (10, 10)],
     'norm': [None, 'ortho', ''],
 }))
 @testing.gpu
@@ -122,7 +122,7 @@ def test_ifft(self, xp, dtype):
 
 
 @testing.parameterize(*testing.product({
-    'shape': [(10, 10), (10, 5, 10)],
+    'shape': [(0, 10), (10, 0, 10), (10, 10), (10, 5, 10)],
     'data_order': ['F', 'C'],
     'axis': [0, 1, -1],
 }))
@@ -164,7 +164,7 @@ def test_ifft(self, xp, dtype):
 # 2. the tested parameter combinations are adjusted to meet the requirements
 @testing.parameterize(*testing.product({
     'n': [None, 0, 64],
-    'shape': [(64,), (4, 64)],
+    'shape': [(0,), (0, 10), (64,), (4, 64)],
     'norm': [None, 'ortho', ''],
 }))
 @testing.multi_gpu(2)
@@ -337,6 +337,11 @@ def test_fft_allocate(self):
     {'shape': (2, 3, 4), 's': None, 'axes': (), 'norm': None},
     {'shape': (2, 3, 4), 's': (2, 3), 'axes': (0, 1, 2), 'norm': 'ortho'},
     {'shape': (2, 3, 4, 5), 's': None, 'axes': None, 'norm': None},
+    {'shape': (0, 5), 's': None, 'axes': None, 'norm': None},
+    {'shape': (2, 0, 5), 's': None, 'axes': None, 'norm': None},
+    {'shape': (0, 0, 5), 's': None, 'axes': None, 'norm': None},
+    {'shape': (3, 4), 's': (0, 5), 'axes': None, 'norm': None},
+    {'shape': (3, 4), 's': (1, 0), 'axes': None, 'norm': None},
 )
 @testing.gpu
 class TestFft2(unittest.TestCase):
@@ -406,6 +411,9 @@ def test_ifft2(self, xp, dtype, order, enable_nd):
     {'shape': (2, 3, 4), 's': (2, 3), 'axes': (0, 1, 2), 'norm': 'ortho'},
     {'shape': (2, 3, 4), 's': (4, 3, 2), 'axes': (2, 0, 1), 'norm': 'ortho'},
     {'shape': (2, 3, 4, 5), 's': None, 'axes': None, 'norm': None},
+    {'shape': (0, 5), 's': None, 'axes': None, 'norm': None},
+    {'shape': (2, 0, 5), 's': None, 'axes': None, 'norm': None},
+    {'shape': (0, 0, 5), 's': None, 'axes': None, 'norm': None},
 )
 @testing.gpu
 class TestFftn(unittest.TestCase):
@@ -468,6 +476,9 @@ def test_ifftn(self, xp, dtype, order, enable_nd):
     {'shape': (2, 3, 4), 's': None, 'axes': None, 'norm': 'ortho'},
     {'shape': (2, 3, 4), 's': (2, 3), 'axes': None, 'norm': 'ortho'},
     {'shape': (2, 3, 4), 's': (2, 3), 'axes': (0, 1, 2), 'norm': 'ortho'},
+    {'shape': (0, 5), 's': None, 'axes': None, 'norm': None},
+    {'shape': (2, 0, 5), 's': None, 'axes': None, 'norm': None},
+    {'shape': (0, 0, 5), 's': None, 'axes': None, 'norm': None},
 )
 @testing.gpu
 class TestPlanCtxManagerFftn(unittest.TestCase):
@@ -515,6 +526,8 @@ def test_ifftn(self, xp, dtype, enable_nd):
     @nd_planning_states()
     @testing.for_complex_dtypes()
     def test_fftn_error_on_wrong_plan(self, dtype, enable_nd):
+        if 0 in self.shape:
+            raise unittest.SkipTest('0 in shape')
         # This test ensures the context manager plan is picked up
 
         from cupyx.scipy.fftpack import get_fft_plan


Code review comments:
diff_hunk: @@ -493,6 +501,12 @@ def _fftn(a, s, axes, norm, direction, value_type='C2C', order='A', plan=None,
     # Note: need to call _cook_shape prior to sorting the axes
     a = _cook_shape(a, s, axes, value_type, order=order)
 
+    print(a.shape, s, axes_sorted)
path: cupy/fft/fft.py
text: @user1:
```suggestion
```