Skip to content
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

[BugFix] Robustize cuFINUFFT Python guru interface. #426

Merged
merged 1 commit into from
Apr 23, 2024

Conversation

SepandKashani
Copy link
Contributor

As per the docs:

  • finufft.Plan(n_modes_or_dim) accepts (int, tuple[int])
  • cufinufft.Plan(n_modes) accepts (int, tuple[int])

In practice n_modes_or_dim is transformed internally to a NumPy array, hence providing any iterable to finufft.Plan() works.
However n_modes in cufinufft.Plan() is more restrictive: it works with tuple[int], fails planning if list[int], and wrongfully passes planning if ndarray[int], but then seg-faults at execute().

The examples below showcase the problem:

### mve_finufft.py
import finufft
import numpy as np

xp, fi = np, finufft

M, D = 1_000, 2
N = (20, 21, 22)[:D]
cdtype = np.cdouble

plan = fi.Plan(
    nufft_type=1,
    n_modes_or_dim=tuple(N),  # ok
    # n_modes_or_dim=list(N),  # ok
    # n_modes_or_dim=np.array(N),  # ok
    dtype=cdtype,
    isign=1,
)

rng = xp.random.default_rng(0)
x = rng.uniform(-np.pi, np.pi, size=(D, M))
plan.setpts(**dict(zip("xyz"[:D], x[:D])))

w = rng.standard_normal(2 * M).view(cdtype)
out = plan.execute(w)
### mve_cufinufft.py
import cufinufft
import cupy as cp
import numpy as np

xp, fi = cp, cufinufft

M, D = 1_000, 2
N = (20, 21, 22)[:D]
cdtype = np.cdouble

plan = fi.Plan(
    nufft_type=1,
    n_modes=tuple(N),  # ok
    # n_modes=list(N),  # planning error
    # n_modes=np.array(N),  # execute error + seg-fault (sometimes)
    dtype=cdtype,
    isign=1,
)

rng = xp.random.default_rng(0)
x = rng.uniform(-np.pi, np.pi, size=(D, M))
plan.setpts(**dict(zip("xyz"[:D], x[:D])))

w = rng.standard_normal(2 * M).view(cdtype)
out = plan.execute(w)

This PR increases the robustness of cufinufft.Plan() to match the flexibility of the CPU interface.

As per the docs:
* finufft.Plan(n_modes_or_dim) accepts (int, tuple[int])
* cufinufft.Plan(n_modes) accepts (int, tuple[int])

In practice `n_modes_or_dim` is transformed internally to a NumPy array, hence providing any iterable to finufft.Plan() works.
However `n_modes` in cufinufft.Plan() is more restrictive: it works with tuple[int], fails planning if list[int], and wrongfully passes planning if ndarray[int], but then seg-faults at execute().

This commit increases the robustness of cufinufft.Plan() to match the flexibility of the CPU interface.
@ahbarnett ahbarnett requested a review from janden April 22, 2024 16:44
Copy link
Collaborator

@janden janden left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks great. Thanks!

@ahbarnett ahbarnett merged commit aff74f2 into flatironinstitute:master Apr 23, 2024
9 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants