Update dpnp.fft to run on CUDA (#2332)

This PR suggests fixing the current issues with `dpnp.fft.fftn()` and
`dpnp.fft.rfftn()` on CUDA and removing the skip tests for them

The `incorrect result` issue for `dpnp.fft.fftn()` on cuda was because
preparing the input array when `batch_fft=True` could change it to `F
contiguous` array. cuFFT for correct execution requires `C contiguous`
array as input.

The issue with raising `Invalid strides` error is a bug in oneMath
([631](uxlfoundation/oneMath#631)).
As a workaround until this issue is solved it is suggested to use swap
of the last two axes if the last dimension is 1 and there are multiple
axes.
In this case the strides inside cuFFT are calculated correctly.

Additionally updated the arguments for `test_erf` in
`skipped_tests_cuda.tbl` to skip it on cuda.

Author: vlad-perevezentsev

dpnp/fft/dpnp_utils_fft.py

@@ -395,13 +395,33 @@ def _fft(a, norm, out, forward, in_place, c2c, axes, batch_fft=True):
         a = dpnp.reshape(a, local_shape)
         index = 1
 
+        # cuFFT requires input arrays to be C-contiguous (row-major)
+        # for correct execution
+        if (
+            dpnp.is_cuda_backend(a) and not a.flags.c_contiguous
+        ):  # pragma: no cover
+            a = dpnp.ascontiguousarray(a)
+
+    # w/a for cuFFT to avoid "Invalid strides" error when
+    # the last dimension is 1 and there are multiple axes
+    # by swapping the last two axes to correct the input.
+    # TODO: Remove this ones the OneMath issue is resolved
+    # https://github.com/uxlfoundation/oneMath/issues/631
+    cufft_wa = dpnp.is_cuda_backend(a) and a.shape[-1] == 1 and len(axes) > 1
+    if cufft_wa:  # pragma: no cover
+        a = dpnp.moveaxis(a, -1, -2)
+
     a_strides = _standardize_strides_to_nonzero(a.strides, a.shape)
     dsc, out_strides = _commit_descriptor(
         a, forward, in_place, c2c, a_strides, index, batch_fft
     )
     res = _compute_result(dsc, a, out, forward, c2c, out_strides)
     res = _scale_result(res, a.shape, norm, forward, index)
 
+    # Revert swapped axes
+    if cufft_wa:  # pragma: no cover
+        res = dpnp.moveaxis(res, -1, -2)
+
     if batch_fft:
         tmp_shape = a_shape_orig[:-1] + (res.shape[-1],)
         res = dpnp.reshape(res, tmp_shape)

dpnp/tests/skipped_tests_cuda.tbl

@@ -813,8 +813,7 @@ tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_3_{extern
 # erf
 tests/test_special.py::test_erf
 tests/test_special.py::test_erf_fallback
-tests/test_strides.py::test_strides_erf[(10,)-int32]
-tests/test_strides.py::test_strides_erf[(10,)-int64]
-tests/test_strides.py::test_strides_erf[(10,)-float32]
-tests/test_strides.py::test_strides_erf[(10,)-float64]
-tests/test_strides.py::test_strides_erf[(10,)-None]
+tests/test_strides.py::test_erf[int32]
+tests/test_strides.py::test_erf[int64]
+tests/test_strides.py::test_erf[float32]
+tests/test_strides.py::test_erf[float64]

dpnp/tests/test_fft.py

@@ -14,7 +14,6 @@
     get_all_dtypes,
     get_complex_dtypes,
     get_float_dtypes,
-    is_cuda_device,
 )
 
 
@@ -443,11 +442,6 @@ def setup_method(self):
     @pytest.mark.parametrize("norm", [None, "backward", "forward", "ortho"])
     @pytest.mark.parametrize("order", ["C", "F"])
     def test_fftn(self, dtype, axes, norm, order):
-        if is_cuda_device():
-            if order == "C" and axes == (0, 1, 2):
-                pass
-            else:
-                pytest.skip("SAT-7587")
         a_np = generate_random_numpy_array((2, 3, 4, 5), dtype, order)
         a = dpnp.array(a_np)
 
@@ -482,9 +476,6 @@ def test_fftn_repeated_axes(self, axes):
     @pytest.mark.parametrize("axes", [(2, 3, 3, 2), (0, 0, 3, 3)])
     @pytest.mark.parametrize("s", [(5, 4, 3, 3), (7, 8, 10, 9)])
     def test_fftn_repeated_axes_with_s(self, axes, s):
-        if is_cuda_device():
-            if axes == (0, 0, 3, 3) and s == (7, 8, 10, 9):
-                pytest.skip("SAT-7587")
         a_np = generate_random_numpy_array((2, 3, 4, 5), dtype=numpy.complex64)
         a = dpnp.array(a_np)
 
@@ -504,11 +495,6 @@ def test_fftn_repeated_axes_with_s(self, axes, s):
     @pytest.mark.parametrize("axes", [(0, 1, 2, 3), (1, 2, 1, 2), (2, 2, 2, 3)])
     @pytest.mark.parametrize("s", [(2, 3, 4, 5), (5, 4, 7, 8), (2, 5, 1, 2)])
     def test_fftn_out(self, axes, s):
-        if is_cuda_device():
-            if axes == (0, 1, 2, 3):
-                pytest.skip("SAT-7587")
-            elif s == (2, 5, 1, 2) and axes in [(1, 2, 1, 2), (2, 2, 2, 3)]:
-                pytest.skip("SAT-7587")
         a_np = generate_random_numpy_array((2, 3, 4, 5), dtype=numpy.complex64)
         a = dpnp.array(a_np)
 
@@ -1082,9 +1068,6 @@ def test_rfftn_repeated_axes_with_s(self, axes, s):
     @pytest.mark.parametrize("axes", [(0, 1, 2, 3), (1, 2, 1, 2), (2, 2, 2, 3)])
     @pytest.mark.parametrize("s", [(2, 3, 4, 5), (5, 6, 7, 9), (2, 5, 1, 2)])
     def test_rfftn_out(self, axes, s):
-        if is_cuda_device():
-            if axes == (0, 1, 2, 3) and s == (2, 5, 1, 2):
-                pytest.skip("SAT-7587")
         x = numpy.random.uniform(-10, 10, 120)
         a_np = numpy.array(x, dtype=numpy.float32).reshape(2, 3, 4, 5)
         a = dpnp.asarray(a_np)

dpnp/tests/third_party/cupy/fft_tests/test_fft.py

@@ -5,7 +5,7 @@
 import pytest
 
 import dpnp as cupy
-from dpnp.tests.helper import has_support_aspect64, is_cuda_device
+from dpnp.tests.helper import has_support_aspect64
 from dpnp.tests.third_party.cupy import testing
 from dpnp.tests.third_party.cupy.testing._loops import _wraps_partial
 
@@ -413,8 +413,6 @@ class TestFft2:
         type_check=has_support_aspect64(),
     )
     def test_fft2(self, xp, dtype, order, enable_nd):
-        if is_cuda_device() and self.shape == (2, 3, 4, 5):
-            pytest.skip("SAT-7587")
         # assert config.enable_nd_planning == enable_nd
         a = testing.shaped_random(self.shape, xp, dtype)
         if order == "F":
@@ -442,8 +440,6 @@ def test_fft2(self, xp, dtype, order, enable_nd):
         type_check=has_support_aspect64(),
     )
     def test_ifft2(self, xp, dtype, order, enable_nd):
-        if is_cuda_device() and self.shape == (2, 3, 4, 5):
-            pytest.skip("SAT-7587")
         # assert config.enable_nd_planning == enable_nd
         a = testing.shaped_random(self.shape, xp, dtype)
         if order == "F":
@@ -507,8 +503,6 @@ class TestFftn:
         type_check=has_support_aspect64(),
     )
     def test_fftn(self, xp, dtype, order, enable_nd):
-        if is_cuda_device() and self.shape == (2, 3, 4, 5):
-            pytest.skip("SAT-7587")
         # assert config.enable_nd_planning == enable_nd
         a = testing.shaped_random(self.shape, xp, dtype)
         if order == "F":
@@ -536,8 +530,6 @@ def test_fftn(self, xp, dtype, order, enable_nd):
         type_check=has_support_aspect64(),
     )
     def test_ifftn(self, xp, dtype, order, enable_nd):
-        if is_cuda_device() and self.shape == (2, 3, 4, 5):
-            pytest.skip("SAT-7587")
         # assert config.enable_nd_planning == enable_nd
         a = testing.shaped_random(self.shape, xp, dtype)
         if order == "F":

dpnp/tests/third_party/cupy/linalg_tests/test_decomposition.py

@@ -7,7 +7,6 @@
 from dpnp.tests.helper import (
     has_support_aspect64,
     is_cpu_device,
-    is_cuda_device,
 )
 from dpnp.tests.third_party.cupy import testing
 from dpnp.tests.third_party.cupy.testing import _condition