ENH: update dpnp.fft; MKL integration (#820)

* ENH: update dpnp.fft; MKL integration (comlex to complex, 64/128)

Author: samir-nasibli

dpnp/backend/include/dpnp_iface_fft.hpp

@@ -70,6 +70,7 @@
  * @param[in]  axis            Axis ID to compute by.
  * @param[in]  input_boundarie Limit number of elements for @ref axis.
  * @param[in]  inverse         Using inverse algorithm.
+ * @param[in]  norm            Normalization mode. 0 - backward, 1 - forward.
  */
 template <typename _DataType>
 INP_DLLEXPORT void dpnp_fft_fft_c(const void* array_in,
@@ -79,6 +80,6 @@ INP_DLLEXPORT void dpnp_fft_fft_c(const void* array_in,
                                   size_t shape_size,
                                   long axis,
                                   long input_boundarie,
-                                  size_t inverse);
-
+                                  size_t inverse,
+                                  const size_t norm);
 #endif // BACKEND_IFACE_FFT_H

dpnp/backend/kernels/dpnp_krnl_fft.cpp

@@ -33,6 +33,11 @@
 
 namespace mkl_dft = oneapi::mkl::dft;
 
+typedef mkl_dft::descriptor<mkl_dft::precision::DOUBLE, mkl_dft::domain::COMPLEX> desc_dp_cmplx_t;
+typedef mkl_dft::descriptor<mkl_dft::precision::SINGLE, mkl_dft::domain::COMPLEX> desc_sp_cmplx_t;
+typedef mkl_dft::descriptor<mkl_dft::precision::DOUBLE, mkl_dft::domain::REAL> desc_dp_real_t;
+typedef mkl_dft::descriptor<mkl_dft::precision::SINGLE, mkl_dft::domain::REAL> desc_sp_real_t;
+
 #ifdef _WIN32
 #ifndef M_PI // Windows compatibility
 #define M_PI 3.14159265358979323846
@@ -43,23 +48,24 @@ template <typename _KernelNameSpecialization1, typename _KernelNameSpecializatio
 class dpnp_fft_fft_c_kernel;
 
 template <typename _DataType_input, typename _DataType_output>
-void dpnp_fft_fft_c(const void* array1_in,
-                    void* result1,
-                    const long* input_shape,
-                    const long* output_shape,
-                    size_t shape_size,
-                    long axis,
-                    long input_boundarie,
-                    size_t inverse)
+void dpnp_fft_fft_sycl_c(const void* array1_in,
+                         void* result1,
+                         const long* input_shape,
+                         const long* output_shape,
+                         size_t shape_size,
+                         const size_t result_size,
+                         const size_t input_size,
+                         long axis,
+                         long input_boundarie,
+                         size_t inverse)
 {
-    const size_t input_size = std::accumulate(input_shape, input_shape + shape_size, 1, std::multiplies<size_t>());
-    const size_t result_size = std::accumulate(output_shape, output_shape + shape_size, 1, std::multiplies<size_t>());
     if (!(input_size && result_size && shape_size))
     {
         return;
     }
 
     cl::sycl::event event;
+
     const double kernel_pi = inverse ? -M_PI : M_PI;
 
     DPNPC_ptr_adapter<_DataType_input> input1_ptr(array1_in, input_size);
@@ -148,21 +154,139 @@ void dpnp_fft_fft_c(const void* array1_in,
     };
 
     event = DPNP_QUEUE.submit(kernel_func);
+    event.wait();
+
+    dpnp_memory_free_c(input_shape_offsets);
+    dpnp_memory_free_c(output_shape_offsets);
+    dpnp_memory_free_c(axis_iterator);
+
+    return;
+}
 
-#if 0 // keep this code
-    oneapi::mkl::dft::descriptor<mkl_dft::precision::DOUBLE, mkl_dft::domain::COMPLEX> desc(result_size);
-    desc.set_value(mkl_dft::config_param::FORWARD_SCALE, static_cast<double>(result_size));
-    desc.set_value(mkl_dft::config_param::PLACEMENT, DFTI_NOT_INPLACE); // enum value from math library C interface
+template <typename _DataType_input, typename _DataType_output, typename _Descriptor_type>
+void dpnp_fft_fft_mathlib_compute_c(const void* array1_in,
+                                    void* result1,
+                                    const size_t shape_size,
+                                    const size_t result_size,
+                                    _Descriptor_type& desc,
+                                    const size_t norm)
+{
+    cl::sycl::event event;
+
+    DPNPC_ptr_adapter<_DataType_input> input1_ptr(array1_in, result_size);
+    DPNPC_ptr_adapter<_DataType_output> result_ptr(result1, result_size);
+    _DataType_input* array_1 = input1_ptr.get_ptr();
+    _DataType_output* result = result_ptr.get_ptr();
+
+    desc.set_value(mkl_dft::config_param::BACKWARD_SCALE, (1.0 / result_size));
+    // enum value from math library C interface
+    // instead of mkl_dft::config_value::NOT_INPLACE
+    desc.set_value(mkl_dft::config_param::PLACEMENT, DFTI_NOT_INPLACE);
     desc.commit(DPNP_QUEUE);
 
     event = mkl_dft::compute_forward(desc, array_1, result);
-#endif
 
     event.wait();
 
-    dpnp_memory_free_c(input_shape_offsets);
-    dpnp_memory_free_c(output_shape_offsets);
-    dpnp_memory_free_c(axis_iterator);
+    return;
+}
+
+// norm: backward - 0, forward is 1
+template <typename _DataType_input, typename _DataType_output>
+void dpnp_fft_fft_mathlib_c(const void* array1_in,
+                            void* result1,
+                            const long* input_shape,
+                            const size_t shape_size,
+                            const size_t result_size,
+                            const size_t norm)
+{
+    if (!shape_size || !result_size || !array1_in || !result1)
+    {
+        return;
+    }
+    std::vector<std::int64_t> dimensions(input_shape, input_shape + shape_size);
+
+    if constexpr (std::is_same<_DataType_input, std::complex<double>>::value &&
+                  std::is_same<_DataType_output, std::complex<double>>::value)
+    {
+        if (shape_size == 1)
+        {
+            desc_dp_cmplx_t desc(result_size);
+            dpnp_fft_fft_mathlib_compute_c<_DataType_input, _DataType_output, desc_dp_cmplx_t>(
+                array1_in, result1, shape_size, result_size, desc, norm);
+        }
+        else
+        {
+            desc_dp_cmplx_t desc(dimensions);
+            dpnp_fft_fft_mathlib_compute_c<_DataType_input, _DataType_output, desc_dp_cmplx_t>(
+                array1_in, result1, shape_size, result_size, desc, norm);
+        }
+    }
+    else if (std::is_same<_DataType_input, std::complex<float>>::value &&
+             std::is_same<_DataType_output, std::complex<float>>::value)
+    {
+        if (shape_size == 1)
+        {
+            desc_sp_cmplx_t desc(result_size);
+            dpnp_fft_fft_mathlib_compute_c<_DataType_input, _DataType_output, desc_sp_cmplx_t>(
+                array1_in, result1, shape_size, result_size, desc, norm);
+        }
+        else
+        {
+            desc_sp_cmplx_t desc(dimensions);
+            dpnp_fft_fft_mathlib_compute_c<_DataType_input, _DataType_output, desc_sp_cmplx_t>(
+                array1_in, result1, shape_size, result_size, desc, norm);
+        }
+    }
+    return;
+}
+
+template <typename _DataType_input, typename _DataType_output>
+void dpnp_fft_fft_c(const void* array1_in,
+                    void* result1,
+                    const long* input_shape,
+                    const long* output_shape,
+                    size_t shape_size,
+                    long axis,
+                    long input_boundarie,
+                    size_t inverse,
+                    const size_t norm)
+{
+    if (!shape_size)
+    {
+        return;
+    }
+
+    const size_t result_size = std::accumulate(output_shape, output_shape + shape_size, 1, std::multiplies<size_t>());
+    const size_t input_size = std::accumulate(input_shape, input_shape + shape_size, 1, std::multiplies<size_t>());
+
+    if (!input_size || !result_size || !array1_in || !result1)
+    {
+        return;
+    }
+
+    if (((std::is_same<_DataType_input, std::complex<double>>::value &&
+          std::is_same<_DataType_output, std::complex<double>>::value) ||
+         (std::is_same<_DataType_input, std::complex<float>>::value &&
+          std::is_same<_DataType_output, std::complex<float>>::value)) &&
+        (shape_size <= 3))
+    {
+        dpnp_fft_fft_mathlib_c<_DataType_input, _DataType_output>(
+            array1_in, result1, input_shape, shape_size, result_size, norm);
+    }
+    else
+    {
+        dpnp_fft_fft_sycl_c<_DataType_input, _DataType_output>(array1_in,
+                                                               result1,
+                                                               input_shape,
+                                                               output_shape,
+                                                               shape_size,
+                                                               result_size,
+                                                               input_size,
+                                                               axis,
+                                                               input_boundarie,
+                                                               inverse);
+    }
 
     return;
 }
@@ -173,12 +297,12 @@ void func_map_init_fft_func(func_map_t& fmap)
                                                              (void*)dpnp_fft_fft_c<int, std::complex<double>>};
     fmap[DPNPFuncName::DPNP_FN_FFT_FFT][eft_LNG][eft_LNG] = {eft_C128,
                                                              (void*)dpnp_fft_fft_c<long, std::complex<double>>};
-    fmap[DPNPFuncName::DPNP_FN_FFT_FFT][eft_FLT][eft_FLT] = {eft_C128,
-                                                             (void*)dpnp_fft_fft_c<float, std::complex<double>>};
+    fmap[DPNPFuncName::DPNP_FN_FFT_FFT][eft_FLT][eft_FLT] = {eft_C64,
+                                                             (void*)dpnp_fft_fft_c<float, std::complex<float>>};
     fmap[DPNPFuncName::DPNP_FN_FFT_FFT][eft_DBL][eft_DBL] = {eft_C128,
                                                              (void*)dpnp_fft_fft_c<double, std::complex<double>>};
     fmap[DPNPFuncName::DPNP_FN_FFT_FFT][eft_C64][eft_C64] = {
-        eft_C128, (void*)dpnp_fft_fft_c<std::complex<float>, std::complex<double>>};
+        eft_C64, (void*)dpnp_fft_fft_c<std::complex<float>, std::complex<float>>};
     fmap[DPNPFuncName::DPNP_FN_FFT_FFT][eft_C128][eft_C128] = {
         eft_C128, (void*)dpnp_fft_fft_c<std::complex<double>, std::complex<double>>};
     return;

dpnp/fft/dpnp_algo_fft.pyx

@@ -41,14 +41,15 @@ __all__ = [
     "dpnp_fft"
 ]
 
-ctypedef void(*fptr_dpnp_fft_fft_t)(void *, void * , long * , long * , size_t, long, long, size_t)
+ctypedef void(*fptr_dpnp_fft_fft_t)(void *, void * , long * , long * , size_t, long, long, size_t, size_t)
 
 
 cpdef utils.dpnp_descriptor dpnp_fft(utils.dpnp_descriptor input,
                                      size_t input_boundarie,
                                      size_t output_boundarie,
                                      long axis,
-                                     size_t inverse):
+                                     size_t inverse,
+                                     size_t norm):
 
     cdef shape_type_c input_shape = input.shape
     cdef shape_type_c output_shape = input_shape
@@ -68,6 +69,6 @@ cpdef utils.dpnp_descriptor dpnp_fft(utils.dpnp_descriptor input,
     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, input_boundarie, inverse)
+         output_shape.data(), input_shape.size(), axis_norm, input_boundarie, inverse, norm)
 
     return result

dpnp/fft/dpnp_iface_fft.py

@@ -76,15 +76,19 @@ def fft(x1, n=None, axis=-1, norm=None):
     Limitations
     -----------
     Parameter ``norm`` is unsupported.
-    Parameter ``x1`` supports ``dpnp.int32``, ``dpnp.int64``, ``dpnp.float32``, ``dpnp.float64`` and
-    ``dpnp.complex128`` datatypes only.
+    Parameter ``x1`` supports ``dpnp.int32``, ``dpnp.int64``, ``dpnp.float32``, ``dpnp.float64``,
+    ``dpnp.complex64`` and ``dpnp.complex128`` datatypes only.
 
     For full documentation refer to :obj:`numpy.fft.fft`.
 
     """
 
     x1_desc = dpnp.get_dpnp_descriptor(x1)
     if x1_desc:
+        # if norm is None or norm is 'backward':
+        #     norm_val = 0
+        # else:
+        #     norm_val = 1
         if axis is None:
             axis_param = -1      # the most right dimension (default value)
         else:
@@ -104,7 +108,7 @@ def fft(x1, n=None, axis=-1, norm=None):
         else:
             output_boundarie = input_boundarie
 
-            return dpnp_fft(x1_desc, input_boundarie, output_boundarie, axis_param, False).get_pyobj()
+            return dpnp_fft(x1_desc, input_boundarie, output_boundarie, axis_param, False, 0).get_pyobj()
 
     return call_origin(numpy.fft.fft, x1, n, axis, norm)
 

tests/skipped_tests_gpu.tbl

@@ -399,8 +399,6 @@ tests/test_dparray.py::test_flatten[[]-int32]
 tests/test_dparray.py::test_flatten[[]-bool]
 tests/test_dparray.py::test_flatten[[]-bool_]
 tests/test_dparray.py::test_flatten[[]-complex]
-tests/test_fft.py::test_fft[complex128]
-tests/test_fft.py::test_fft[complex64]
 tests/test_fft.py::test_fft[float32]
 tests/test_fft.py::test_fft[float64]
 tests/test_fft.py::test_fft[int32]

tests/test_fft.py

@@ -15,7 +15,7 @@ def test_fft(type):
     dpnp_data = dpnp.array(data)
 
     np_res = numpy.fft.fft(data)
-    dpnp_res = dpnp.fft.fft(dpnp_data)
+    dpnp_res = dpnp.asnumpy(dpnp.fft.fft(dpnp_data))
 
     numpy.testing.assert_allclose(dpnp_res, np_res, rtol=1e-4, atol=1e-7)
     assert dpnp_res.dtype == np_res.dtype