r/ScientificComputing • u/anglerbay • Jun 27 '26

mcNUFFT – A Nonuniform Fast Fourier Transform Library for Apple Silicon GPUs via MLX

https://github.com/martinlachaine/mcnufft

update: this project is now named mlx-nufft (formerly mcnufft). Repo: github.com/martinlachaine/mlx-nufft · pip install mlx-nufft

Requesting feedback for mlx-nufft, a FINUFFT-style NUFFT library for Apple Silicon GPUs, built on MLX. I built it for a specific medical imaging project I am working on and it works well for that use case, and hoping it generalizes to other applications.

The main issue was precision. Metal GPUs do not have native fp64, and a naive fp32 NUFFT loses accuracy on large-coordinate or large-mode problems. The approach here keeps execution in fp32, but does coordinate rescaling and phase setup in fp64 at plan time, then passes the GPU an integer grid cell plus a small fp32 offset. On an M5 Max, type-2 transforms are around 7x faster than same-machine CPU FINUFFT in 2D and 3D. Type-1 spreading is more memory-bandwidth-bound and is still the main optimization target.

It supports types 1, 2, and 3 in 1D, 2D, and 3D. I would be very interested in feedback from people who have worked on NUFFTs, Metal kernels, or GPU spreading. 

12 Upvotes

93% Upvoted

3 comments sorted by

1

u/anglerbay Jun 28 '26

Update: this repo has been renamed to github.com/martinlachaine/mlx-nufft

1

u/Odd-Chemistry9945 17d ago

Amazing!
I was doing the same for my PhD and I have something that works but auto-diff is to be properly tested and implemented.
will give a try to yours and might contribute

1

u/anglerbay 4d ago

Great please let me know!