Skip to content

Update dpnp.fft to run on CUDA #2332

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Feb 24, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions dpnp/fft/dpnp_utils_fft.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
9 changes: 4 additions & 5 deletions dpnp/tests/skipped_tests_cuda.tbl
Original file line number Diff line number Diff line change
Expand Up @@ -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]
17 changes: 0 additions & 17 deletions dpnp/tests/test_fft.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@
get_all_dtypes,
get_complex_dtypes,
get_float_dtypes,
is_cuda_device,
)


Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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)

Expand All @@ -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)

Expand Down Expand Up @@ -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)
Expand Down
10 changes: 1 addition & 9 deletions dpnp/tests/third_party/cupy/fft_tests/test_fft.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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":
Expand Down Expand Up @@ -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":
Expand Down Expand Up @@ -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":
Expand Down Expand Up @@ -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":
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading