Support for multidim arrays in dpnp.fft.fft (#1112)

* support for multidim arrays in dpnp.fft.fft

* fix comments above

* add check of shape

Co-authored-by: Alexander-Makaryev <alexander.makaryev@gmail.com>

Author: KsanaKozlova

dpnp/backend/kernels/dpnp_krnl_fft.cpp

@@ -169,12 +169,16 @@ void dpnp_fft_fft_sycl_c(const void* array1_in,
 template <typename _DataType_input, typename _DataType_output, typename _Descriptor_type>
 void dpnp_fft_fft_mathlib_compute_c(const void* array1_in,
                                     void* result1,
+                                    const shape_elem_type* input_shape,
                                     const size_t shape_size,
                                     const size_t result_size,
                                     _Descriptor_type& desc,
                                     const size_t norm)
 {
-    sycl::event event;
+    if (!shape_size)
+    {
+        return;
+    }
 
     DPNPC_ptr_adapter<_DataType_input> input1_ptr(array1_in, result_size);
     DPNPC_ptr_adapter<_DataType_output> result_ptr(result1, result_size);
@@ -187,9 +191,19 @@ void dpnp_fft_fft_mathlib_compute_c(const void* array1_in,
     desc.set_value(mkl_dft::config_param::PLACEMENT, DFTI_NOT_INPLACE);
     desc.commit(DPNP_QUEUE);
 
-    event = mkl_dft::compute_forward(desc, array_1, result);
+    const size_t n_iter =
+        std::accumulate(input_shape, input_shape + shape_size - 1, 1, std::multiplies<shape_elem_type>());
 
-    event.wait();
+    const size_t shift = input_shape[shape_size - 1];
+
+    std::vector<sycl::event> fft_events;
+    fft_events.reserve(n_iter);
+
+    for (size_t i = 0; i < n_iter; ++i) {
+        fft_events.push_back(mkl_dft::compute_forward(desc, array_1 + i * shift, result + i * shift));
+    }
+
+    sycl::event::wait(fft_events);
 
     return;
 }
@@ -207,39 +221,24 @@ void dpnp_fft_fft_mathlib_c(const void* array1_in,
     {
         return;
     }
-    std::vector<std::int64_t> dimensions(input_shape, input_shape + shape_size);
+    //will be used with strides
+    //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);
-        }
+        desc_dp_cmplx_t desc(input_shape[shape_size - 1]);
+
+        dpnp_fft_fft_mathlib_compute_c<_DataType_input, _DataType_output, desc_dp_cmplx_t>(
+            array1_in, result1, input_shape, 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);
-        }
+        desc_sp_cmplx_t desc(input_shape[shape_size - 1]);
+
+        dpnp_fft_fft_mathlib_compute_c<_DataType_input, _DataType_output, desc_sp_cmplx_t>(
+            array1_in, result1, input_shape, shape_size, result_size, desc, norm);
     }
     return;
 }
@@ -270,11 +269,10 @@ void dpnp_fft_fft_c(const void* array1_in,
         return;
     }
 
-    if (((std::is_same<_DataType_input, std::complex<double>>::value &&
+    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))
+          std::is_same<_DataType_output, std::complex<float>>::value))
     {
         dpnp_fft_fft_mathlib_c<_DataType_input, _DataType_output>(
             array1_in, result1, input_shape, shape_size, result_size, norm);

dpnp/fft/dpnp_iface_fft.py

@@ -76,6 +76,7 @@ def fft(x1, n=None, axis=-1, norm=None):
     Limitations
     -----------
     Parameter ``norm`` is unsupported.
+    Parameter ``axis`` is supported with its default value.
     Parameter ``x1`` supports ``dpnp.int32``, ``dpnp.int64``, ``dpnp.float32``, ``dpnp.float64``,
     ``dpnp.complex64`` and ``dpnp.complex128`` datatypes only.
 
@@ -105,11 +106,11 @@ def fft(x1, n=None, axis=-1, norm=None):
             pass                 # let fallback to handle exception
         elif norm is not None:
             pass
+        elif axis != -1:
+            pass
         else:
             output_boundarie = input_boundarie
-
             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

@@ -217,6 +217,22 @@ 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::test_fft[int64]
+tests/test_fft.py::test_fft_ndim[shape0-float32]
+tests/test_fft.py::test_fft_ndim[shape0-float64]
+tests/test_fft.py::test_fft_ndim[shape0-int32]
+tests/test_fft.py::test_fft_ndim[shape0-int64]
+tests/test_fft.py::test_fft_ndim[shape1-float32]
+tests/test_fft.py::test_fft_ndim[shape1-float64]
+tests/test_fft.py::test_fft_ndim[shape1-int32]
+tests/test_fft.py::test_fft_ndim[shape1-int64]
+tests/test_fft.py::test_fft_ndim[shape2-float32]
+tests/test_fft.py::test_fft_ndim[shape2-float64]
+tests/test_fft.py::test_fft_ndim[shape2-int32]
+tests/test_fft.py::test_fft_ndim[shape2-int64]
+tests/test_fft.py::test_fft_ndim[shape3-float32]
+tests/test_fft.py::test_fft_ndim[shape3-float64]
+tests/test_fft.py::test_fft_ndim[shape3-int32]
+tests/test_fft.py::test_fft_ndim[shape3-int64]
 tests/test_linalg.py::test_cond[-1-[[1, 2, 3], [4, 5, 6], [7, 8, 9]]]
 tests/test_linalg.py::test_cond[1-[[1, 2, 3], [4, 5, 6], [7, 8, 9]]]
 tests/test_linalg.py::test_cond[-2-[[1, 0, -1], [0, 1, 0], [1, 0, 1]]]

tests/test_fft.py

@@ -19,3 +19,16 @@ def test_fft(type):
 
     numpy.testing.assert_allclose(dpnp_res, np_res, rtol=1e-4, atol=1e-7)
     assert dpnp_res.dtype == np_res.dtype
+
+
+@pytest.mark.parametrize("type", ['complex128', 'complex64', 'float32', 'float64', 'int32', 'int64'])
+@pytest.mark.parametrize("shape", [(8, 8), (4, 16), (4, 4, 4), (2, 4, 4, 2)])
+def test_fft_ndim(type, shape):
+    np_data = numpy.arange(64, dtype=numpy.dtype(type)).reshape(shape)
+    dpnp_data = dpnp.arange(64, dtype=numpy.dtype(type)).reshape(shape)
+
+    np_res = numpy.fft.fft(np_data)
+    dpnp_res = 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