FFT. rfft(), rfft2() rfftn() kernel (#345)

Author: shssf

dpnp/backend/backend_iface_fft.hpp

@@ -62,19 +62,21 @@
  *
  * Compute the one-dimensional discrete Fourier Transform.
  *
- * @param[in]  array1_in     Input array 1.
- * @param[out] result1       Output array.
- * @param[in]  input_shape   Array with shape information for input array.
- * @param[in]  output_shape  Array with shape information for output array.
- * @param[in]  shape_size    Number of elements in @ref input_shape or @ref output_shape arrays.
- * @param[in]  axis          Axis ID to compute by.
+ * @param[in]  array_in        Input array.
+ * @param[out] result          Output array.
+ * @param[in]  input_shape     Array with shape information for input array.
+ * @param[in]  output_shape    Array with shape information for output array.
+ * @param[in]  shape_size      Number of elements in @ref input_shape or @ref output_shape arrays.
+ * @param[in]  axis            Axis ID to compute by.
+ * @param[in]  input_boundarie Limit number of elements for @ref axis.
  */
 template <typename _DataType>
-INP_DLLEXPORT void dpnp_fft_fft_c(const void* array1_in,
-                                  void* result1,
+INP_DLLEXPORT void dpnp_fft_fft_c(const void* array_in,
+                                  void* result,
                                   const long* input_shape,
                                   const long* output_shape,
                                   size_t shape_size,
-                                  long axis);
+                                  long axis,
+                                  long input_boundarie);
 
 #endif // BACKEND_IFACE_FFT_H

dpnp/backend/dpnp_kernels_fft.cpp

@@ -41,7 +41,8 @@ void dpnp_fft_fft_c(const void* array1_in,
                     const long* input_shape,
                     const long* output_shape,
                     size_t shape_size,
-                    long axis)
+                    long axis,
+                    long input_boundarie)
 {
     const size_t result_size = std::accumulate(output_shape, output_shape + shape_size, 1, std::multiplies<size_t>());
     if (!(result_size && shape_size))
@@ -80,7 +81,7 @@ void dpnp_fft_fft_c(const void* array1_in,
             axis_iterator_thread[i] = xyz_thread[i];
         }
 
-        const long axis_length = output_shape[axis];
+        const long axis_length = input_boundarie;
         for (long it = 0; it < axis_length; ++it)
         {
             double in_real = 0.0;

dpnp/fft/dpnp_algo_fft.pyx

@@ -42,29 +42,30 @@ __all__ = [
     "dpnp_fft"
 ]
 
-ctypedef void(*fptr_dpnp_fft_fft_t)(void * , void * , long * , long * , size_t, long)
+ctypedef void(*fptr_dpnp_fft_fft_t)(void *, void * , long * , long * , size_t, long, long)
 
 
-cpdef dparray dpnp_fft(dparray input, size_t axis_boundarie, long axis):
+cpdef dparray dpnp_fft(dparray input, size_t input_boundarie, size_t output_boundarie, long axis):
 
     cdef dparray_shape_type input_shape = input.shape
     cdef dparray_shape_type output_shape = input_shape
 
     cdef long axis_norm = normalize_axis((axis,), input_shape.size())[0]
-    output_shape[axis_norm] = axis_boundarie
+    output_shape[axis_norm] = output_boundarie
 
     # convert string type names (dparray.dtype) to C enum DPNPFuncType
     cdef DPNPFuncType param1_type = dpnp_dtype_to_DPNPFuncType(input.dtype)
 
     # get the FPTR data structure
     cdef DPNPFuncData kernel_data = get_dpnp_function_ptr(DPNP_FN_FFT_FFT, param1_type, param1_type)
 
-    result_type = dpnp_DPNPFuncType_to_dtype(< size_t > kernel_data.return_type)
+    result_type = dpnp_DPNPFuncType_to_dtype( < size_t > kernel_data.return_type)
     # ceate result array with type given by FPTR data
     cdef dparray result = dparray(output_shape, dtype=result_type)
 
     cdef fptr_dpnp_fft_fft_t func = <fptr_dpnp_fft_fft_t > kernel_data.ptr
     # call FPTR function
-    func(input.get_data(), result.get_data(), input_shape.data(), output_shape.data(), input_shape.size(), axis_norm)
+    func(input.get_data(), result.get_data(), input_shape.data(),
+         output_shape.data(), input_shape.size(), axis_norm, input_boundarie)
 
     return result

dpnp/fft/dpnp_iface_fft.py

@@ -87,18 +87,20 @@ def fft(x1, n=None, axis=-1, norm=None):
             axis_param = axis
 
         if n is None:
-            boundarie = x1.shape[axis_param]
+            input_boundarie = x1.shape[axis_param]
         else:
-            boundarie = n
+            input_boundarie = n
 
         if x1.size < 1:
             pass                 # let fallback to handle exception
-        elif boundarie < 1:
+        elif input_boundarie < 1:
             pass                 # let fallback to handle exception
         elif norm is not None:
             pass
         else:
-            return dpnp_fft(x1, boundarie, axis_param)
+            output_boundarie = input_boundarie
+
+            return dpnp_fft(x1, input_boundarie, output_boundarie, axis_param)
 
     return call_origin(numpy.fft.fft, x1, n, axis, norm)
 
@@ -202,18 +204,20 @@ def ifft(x1, n=None, axis=-1, norm=None):
             axis_param = axis
 
         if n is None:
-            boundarie = x1.shape[axis_param]
+            input_boundarie = x1.shape[axis_param]
         else:
-            boundarie = n
+            input_boundarie = n
 
         if x1.size < 1:
             pass                 # let fallback to handle exception
-        elif boundarie < 1:
+        elif input_boundarie < 1:
             pass                 # let fallback to handle exception
         elif norm is not None:
             pass
         else:
-            return dpnp_fft(x1, boundarie, axis_param)
+            output_boundarie = input_boundarie
+
+            return dpnp_fft(x1, input_boundarie, output_boundarie, axis_param)
 
     return call_origin(numpy.fft.ifft, x1, n, axis, norm)
 
@@ -317,18 +321,20 @@ def irfft(x1, n=None, axis=-1, norm=None):
             axis_param = axis
 
         if n is None:
-            boundarie = x1.shape[axis_param]
+            input_boundarie = x1.shape[axis_param]
         else:
-            boundarie = n
+            input_boundarie = n
 
         if x1.size < 1:
             pass                 # let fallback to handle exception
-        elif boundarie < 1:
+        elif input_boundarie < 1:
             pass                 # let fallback to handle exception
         elif norm is not None:
             pass
         else:
-            return dpnp_fft(x1, boundarie, axis_param)
+            output_boundarie = input_boundarie
+
+            return dpnp_fft(x1, input_boundarie, output_boundarie, axis_param)
 
     return call_origin(numpy.fft.irfft, x1, n, axis, norm)
 
@@ -408,6 +414,7 @@ def irfftn(x1, s=None, axes=None, norm=None):
 
     return call_origin(numpy.fft.irfftn, x1, s, axes, norm)
 
+
 def rfft(x1, n=None, axis=-1, norm=None):
     """
     Compute the one-dimensional discrete Fourier Transform for real input.
@@ -424,25 +431,27 @@ def rfft(x1, n=None, axis=-1, norm=None):
 
     is_x1_dparray = isinstance(x1, dparray)
 
-    if (not use_origin_backend(x1) and is_x1_dparray and 0):
+    if (not use_origin_backend(x1) and is_x1_dparray):
         if axis is None:
-            axis_param = -1      # the most right dimension (default value)
+            axis_param = -1                             # the most right dimension (default value)
         else:
             axis_param = axis
 
         if n is None:
-            boundarie = x1.shape[axis_param]
+            input_boundarie = x1.shape[axis_param]
         else:
-            boundarie = n
+            input_boundarie = n
 
         if x1.size < 1:
-            pass                 # let fallback to handle exception
-        elif boundarie < 1:
-            pass                 # let fallback to handle exception
+            pass                                        # let fallback to handle exception
+        elif input_boundarie < 1:
+            pass                                        # let fallback to handle exception
         elif norm is not None:
             pass
         else:
-            return dpnp_fft(x1, boundarie, axis_param)
+            output_boundarie = input_boundarie // 2 + 1  # rfft specific requirenment
+
+            return dpnp_fft(x1, input_boundarie, output_boundarie, axis_param)
 
     return call_origin(numpy.fft.rfft, x1, n, axis, norm)
 
@@ -465,11 +474,11 @@ def rfft2(x1, s=None, axes=(-2, -1), norm=None):
 
     is_x1_dparray = isinstance(x1, dparray)
 
-    if (not use_origin_backend(x1) and is_x1_dparray and 0):
+    if (not use_origin_backend(x1) and is_x1_dparray):
         if norm is not None:
             pass
         else:
-            return fftn(x1, s, axes, norm)
+            return rfftn(x1, s, axes, norm)
 
     return call_origin(numpy.fft.rfft2, x1, s, axes, norm)
 
@@ -492,7 +501,7 @@ def rfftn(x1, s=None, axes=None, norm=None):
 
     is_x1_dparray = isinstance(x1, dparray)
 
-    if (not use_origin_backend(x1) and is_x1_dparray and 0):
+    if (not use_origin_backend(x1) and is_x1_dparray):
         if s is None:
             boundaries = tuple([x1.shape[i] for i in range(x1.ndim)])
         else:
@@ -505,6 +514,8 @@ def rfftn(x1, s=None, axes=None, norm=None):
 
         if norm is not None:
             pass
+        elif len(axes) < 1:
+            pass                      # let fallback to handle exception
         else:
             x1_iter = x1
             iteration_list = list(range(len(axes_param)))
@@ -516,7 +527,7 @@ def rfftn(x1, s=None, axes=None, norm=None):
                 except IndexError:
                     checker_throw_axis_error("fft.rfftn", "is out of bounds", param_axis, f"< {len(boundaries)}")
 
-                x1_iter = fft(x1_iter, n=param_n, axis=param_axis, norm=norm)
+                x1_iter = rfft(x1_iter, n=param_n, axis=param_axis, norm=norm)
 
             return x1_iter