Repository: flatironinstitute/jax-finufft (language: python)

Issue:
Update vendored finufft and add GPU support
In recent months, cufinufft has been merged into the primary finufft codebase (thanks to @blackwer), which itself has matured a lot over the last few years. The vendored finufft hasn't been updated in a long time, so this PR does that, fixes some minor compatibility issues, and removes the deprecated vendored cufinufft. It also adds GPU support.

The GPU support comes from #4 but with an important update, which is that we can now pass the GPU stream that JAX gives us to the cufinufft library via flatironinstitute/finufft#330. This is probably required for multi-GPU support (although I haven't tested it) and may help with performance too. The only other changes were various updates to match the new API, build fixes, etc.

Some notes about the CMake build: cufinufft no longer requires building two static libraries with different macro definitions (one for float and one for double). The multiple precisions are supported via C++ templates, so the library can be built all at once. Also, now that finufft itself uses CMake, we might prefer to include it as a CMake sub-package rather than itemizing source files. But I left the itemization approach for now, since it required fewer changes.

For the moment, the vendored finufft is maintained in my own fork while we wait for some important PRs (including flatironinstitute/finufft#330 and flatironinstitute/finufft#354) to be merged. It should be easy to re-target the primary repo later.

The tests pass on the GPU and CPU, but more help testing the GPU in particular would be great! For anybody looking to run this on the Flatiron clusters, this is my build environment:
<details>
<summary>env.sh</summary>

```
ml modules/2.2
ml gcc
ml python/3.11
ml fftw
ml cmake
ml cuda/12
ml cudnn
```
</details>


Discussion / hints:
This is awesome! I wonder if it's worth getting the Flatiron Jenkins CI set up to test the GPU ops? I've been doing that with the JAX CUDA ops in my `exoplanet-core` package: https://github.com/exoplanet-dev/exoplanet-core/blob/main/ci/Jenkinsfile

I'd like to update all the custom call stuff because the best practices have changed a little in the meantime, but that should be a separate PR.

For now, I'm keen to merge this with the only question being about running a GPU-enabled CI.
Yeah, I think it definitely makes sense to set this up on Jenkins! I actually started looking into this, but @Matematija reported an non-deterministic crash that I'm looking into first, might be when using the latest JAX. I'll keep you posted.
The problems appear at JAX 0.4.9, which is also the version where JAX starts using a non-blocking CUDA stream, according to google/jax#16580. I've gone through and fixed some cufinufft stream race conditions for kernel launches, `cudaMemcpyAsync`, Thrust, and (maybe) cufft, which resolves the particular crash that @Matematija reported, but there are still problems, because the tests don't pass. I'll need to dig into why (I'll probably be delayed in working on this by jury duty).
I think this is fixed now! The relevant patch is in the finufft submodule, so don't forget to update submodules when you pull.  I'll work on Jenkins next.
@dfm Can you add the Jenkins webhook to this project, or give me permissions to do so? I think I would need to be a repo admin.

Fix patch:
diff --git a/.github/dependabot.yml b/.github/dependabot.yml
new file mode 100644
index 0000000..8ac6b8c
--- /dev/null
+++ b/.github/dependabot.yml
@@ -0,0 +1,6 @@
+version: 2
+updates:
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "monthly"
diff --git a/.github/workflows/gpu-build.yml b/.github/workflows/gpu-build.yml
new file mode 100644
index 0000000..2780280
--- /dev/null
+++ b/.github/workflows/gpu-build.yml
@@ -0,0 +1,39 @@
+name: GPU Build
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+
+jobs:
+  build:
+    runs-on: ubuntu-22.04
+    container: nvidia/cuda:12.2.2-devel-ubuntu22.04
+
+    steps:
+      - name: Install git and FFTW
+        run: apt-get update && apt-get install -y git libfftw3-dev
+        env:
+          DEBIAN_FRONTEND: "noninteractive"
+
+      - name: Get the version info for NVCC
+        run: nvcc --version
+
+      - uses: actions/checkout@v4
+        with:
+          submodules: true
+          fetch-depth: 0
+
+      - uses: actions/setup-python@v4
+        with:
+          python-version: 3.11
+
+      - name: Install Python dependencies
+        run: python -m pip install -U pip
+
+      - name: Compile the extension
+        run: python -m pip install -v .
+
+      - name: Check that the GPU extension was built
+        run: python -c "import jax_finufft.jax_finufft_gpu"
diff --git a/.gitignore b/.gitignore
index 8feb13b..d995aca 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,3 +6,4 @@ dist
 MANIFEST
 __pycache__/
 *.egg-info
+.*_cache
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 9e349b8..7a15776 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -1,124 +1,63 @@
-cmake_minimum_required(VERSION 3.15)
+cmake_minimum_required(VERSION 3.15...3.26)
 project(${SKBUILD_PROJECT_NAME} LANGUAGES C CXX)
 message(STATUS "Using CMake version: " ${CMAKE_VERSION})
 
-# Add the /cmake directory to the module path so that we can find FFTW
-set(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} ${CMAKE_CURRENT_LIST_DIR}/cmake)
+# Workaround for LTO applied incorrectly to CUDA fatbin
+# https://github.com/pybind/pybind11/issues/4825
+set(CMAKE_INTERPROCEDURAL_OPTIMIZATION OFF)
 
-# Handle Python settings passed from scikit-build
-if(SKBUILD)
-    set(Python_EXECUTABLE "${PYTHON_EXECUTABLE}")
-    set(Python_INCLUDE_DIR "${PYTHON_INCLUDE_DIR}")
-    set(Python_LIBRARY "${PYTHON_LIBRARY}")
-    execute_process(
-        COMMAND "${PYTHON_EXECUTABLE}" -c "import pybind11; print(pybind11.get_cmake_dir())"
-        OUTPUT_VARIABLE _tmp_dir
-        OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ECHO STDOUT)
-      list(APPEND CMAKE_PREFIX_PATH "${_tmp_dir}")
+include(CheckLanguage)
+check_language(CUDA)
+
+if(CMAKE_CUDA_COMPILER)
+    enable_language(CUDA)
+    set(FINUFFT_USE_CUDA ON)
 else()
-    find_package(Python COMPONENTS Interpreter Development REQUIRED)
-    execute_process(
-        COMMAND "${Python_EXECUTABLE}" -c "import pybind11; print(pybind11.get_cmake_dir())"
-        OUTPUT_VARIABLE _tmp_dir
-        OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ECHO STDOUT)
-      list(APPEND CMAKE_PREFIX_PATH "${_tmp_dir}")
+    message(STATUS "No CUDA compiler found; GPU support will be disabled")
+    set(FINUFFT_USE_CUDA OFF)
 endif()
 
-set(PYBIND11_NEWPYTHON ON)
-find_package(pybind11 CONFIG REQUIRED)
-find_package(FFTW REQUIRED COMPONENTS FLOAT_LIB DOUBLE_LIB)
-link_libraries(${FFTW_FLOAT_LIB} ${FFTW_DOUBLE_LIB})
-
-# Work out compiler flags
-set(CMAKE_POSITION_INDEPENDENT_CODE ON)
-add_compile_options(-Wall -Wno-unknown-pragmas -O3 -funroll-loops -fdiagnostics-color)
-set(FINUFFT_INCLUDE_DIRS
-    ${CMAKE_CURRENT_LIST_DIR}/lib
-    ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/include
-    ${FFTW_INCLUDE_DIRS})
-
-# Build single and double point versions of the FINUFFT library
-add_library(finufft STATIC
-    ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/spreadinterp.cpp
-    ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/utils.cpp
-    ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/finufft.cpp
-    ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/simpleinterfaces.cpp)
-target_include_directories(finufft PRIVATE ${FINUFFT_INCLUDE_DIRS})
+if (${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
+    # TODO(dfm): OpenMP segfaults on my system - can we enable this somehow?
+    set(FINUFFT_USE_OPENMP OFF)
+else()
+    set(FINUFFT_USE_OPENMP ON)
+endif()
 
-add_library(finufft_32 STATIC
-    ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/spreadinterp.cpp
-    ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/utils.cpp
-    ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/finufft.cpp
-    ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/simpleinterfaces.cpp)
-target_compile_definitions(finufft_32 PUBLIC SINGLE)
-target_include_directories(finufft_32 PRIVATE ${FINUFFT_INCLUDE_DIRS})
+# Add the FINUFFT project using the vendored version
+add_subdirectory("${CMAKE_CURRENT_LIST_DIR}/vendor/finufft")
 
-# Build the XLA bindings to those libraries
-pybind11_add_module(jax_finufft_cpu
-    ${CMAKE_CURRENT_LIST_DIR}/lib/jax_finufft_cpu.cc
-    ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/utils_precindep.cpp
-    ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/contrib/legendre_rule_fast.cpp)
+# Find pybind11
+set(PYBIND11_NEWPYTHON ON)
+find_package(pybind11 CONFIG REQUIRED)
 
-target_link_libraries(jax_finufft_cpu PRIVATE finufft finufft_32)
-target_include_directories(jax_finufft_cpu PRIVATE ${FINUFFT_INCLUDE_DIRS})
-install(TARGETS jax_finufft_cpu DESTINATION .)
+# Build the CPU XLA bindings
+pybind11_add_module(jax_finufft_cpu ${CMAKE_CURRENT_LIST_DIR}/lib/jax_finufft_cpu.cc)
+target_link_libraries(jax_finufft_cpu PRIVATE finufft_static)
+install(TARGETS jax_finufft_cpu LIBRARY DESTINATION .)
 
+# Include the CUDA extensions if possible
 include(CheckLanguage)
 check_language(CUDA)
-if (CMAKE_CUDA_COMPILER)
+
+if(CMAKE_CUDA_COMPILER)
     enable_language(CUDA)
     set(CMAKE_CUDA_SEPARABLE_COMPILATION ON)
-    if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
-        set(CMAKE_CUDA_ARCHITECTURES "60;61;70;75;80;90")
-    endif()
 
-    # Find cufft
-    find_package(CUDAToolkit)
-
-    set(CUFINUFFT_INCLUDE_DIRS
-        ${CMAKE_CURRENT_LIST_DIR}/lib
-        # ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/include/cufinufft/contrib/
-        ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/include/cufinufft/contrib/cuda_samples
+    # TODO(dfm): The ${CUFINUFFT_INCLUDE_DIRS} variable doesn't seem to get set
+    # properly when FINUFFT is included as a submodule (maybe because of the use
+    # of ${PROJECT_SOURCE_DIR}). This is just copied from there, linking to the
+    # appropriate vendored directories.
+    set(CUFINUFFT_VENDORED_INCLUDE_DIRS
         ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/include
         ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/contrib
-        ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
-
-    set(CUFINUFFT_SOURCES
-        # TODO: 1D not supported via JAX, but needed for compilation
-        ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/cuda/1d/spread1d_wrapper.cu
-        ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/cuda/1d/interp1d_wrapper.cu
-        ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/cuda/1d/cufinufft1d.cu
-
-        ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/cuda/2d/spread2d_wrapper.cu
-        ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/cuda/2d/interp2d_wrapper.cu
-        ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/cuda/2d/cufinufft2d.cu
-
-        ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/cuda/3d/spread3d_wrapper.cu
-        ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/cuda/3d/interp3d_wrapper.cu
-        ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/cuda/3d/cufinufft3d.cu
-
-        ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/cuda/deconvolve_wrapper.cu
-        ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/cuda/memtransfer_wrapper.cu
-
-        ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/cuda/spreadinterp.cpp
-
-        ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/cuda/precision_independent.cu
-        ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/contrib/legendre_rule_fast.cpp
-        ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/cuda/utils.cpp
-        ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/cuda/cufinufft.cu
-        ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/src/cuda/common.cu)
-
-    add_library(cufinufft STATIC ${CUFINUFFT_SOURCES})
-    target_include_directories(cufinufft PRIVATE ${CUFINUFFT_INCLUDE_DIRS})
-
+        ${CMAKE_CURRENT_LIST_DIR}/vendor/finufft/include/cufinufft/contrib/cuda_samples
+    )
     pybind11_add_module(jax_finufft_gpu
         ${CMAKE_CURRENT_LIST_DIR}/lib/jax_finufft_gpu.cc
         ${CMAKE_CURRENT_LIST_DIR}/lib/kernels.cc.cu)
-    target_link_libraries(jax_finufft_gpu PRIVATE cufinufft)
-    target_link_libraries(jax_finufft_gpu PRIVATE ${CUDA_cufft_LIBRARY} ${CUDA_nvToolsExt_LIBRARY})
-    target_include_directories(jax_finufft_gpu PRIVATE ${CUFINUFFT_INCLUDE_DIRS})
-    install(TARGETS jax_finufft_gpu DESTINATION .)
-
-else()
-    message(STATUS "No CUDA compiler found; GPU support will be disabled")
+    target_include_directories(jax_finufft_gpu PUBLIC ${CUFINUFFT_INCLUDE_DIRS})
+    target_include_directories(jax_finufft_gpu PUBLIC ${CUFINUFFT_VENDORED_INCLUDE_DIRS})
+    target_link_libraries(jax_finufft_gpu PRIVATE cufinufft_static)
+    install(TARGETS jax_finufft_gpu LIBRARY DESTINATION .)
 endif()
diff --git a/README.md b/README.md
index dc6f6c0..4837751 100644
--- a/README.md
+++ b/README.md
@@ -11,8 +11,9 @@ rules for the transforms.
 
 ## Included features
 
-_This library is currently CPU-only, but GPU support is in the works using the
-[cuFINUFFT library](https://github.com/flatironinstitute/cufinufft)._
+This library includes CPU and GPU (CUDA) support. GPU support is implemented
+through the [cuFINUFFT interface](https://finufft.readthedocs.io/en/latest/c_gpu.html)
+of the FINUFFT library.
 
 [Type 1 and 2](https://finufft.readthedocs.io/en/latest/math.html) transforms
 are supported in 1-, 2-, and 3-dimensions. All of these functions support
@@ -24,24 +25,58 @@ forward, reverse, and higher-order differentiation, as well as batching using
 _For now, only a source build is supported._
 
 For building, you should only need a recent version of Python (>3.6) and
-[FFTW](https://www.fftw.org/). At runtime, you'll need `numpy`, `scipy`, and
-`jax`. To set up such an environment, you can use `conda` (but you're welcome to
-use whatever workflow works for you!):
+[FFTW](https://www.fftw.org/). GPU-enabled builds also require a working CUDA
+compiler (i.e. the CUDA Toolkit), CUDA >= 11.8, and a compatible cuDNN (older versions of CUDA may work but
+are untested). At runtime, you'll need `numpy`, `scipy`, and `jax`.
 
+First, clone the repo and `cd` into the repo root (don't forget the `--recursive` flag because FINUFFT is included as a submodule):
 ```bash
-conda create -n jax-finufft -c conda-forge python=3.9 numpy scipy fftw
-python -m pip install "jax[cpu]"
+git clone --recursive https://github.com/dfm/jax-finufft
+cd jax-finufft
 ```
 
-Then you can install from source using (don't forget the `--recursive` flag
-because FINUFFT is included as a submodule):
+Then, you can use `conda` to set up a build environment (but you're welcome to
+use whatever workflow works for you!). For example, for a CPU build, you can use:
 
 ```bash
-git clone --recursive https://github.com/dfm/jax-finufft
-cd jax-finufft
+conda create -n jax-finufft -c conda-forge python=3.10 numpy scipy fftw cxx-compiler
+conda activate jax-finufft
+python -m pip install "jax[cpu]"
 python -m pip install .
 ```
 
+For a GPU build, while the CUDA libraries and compiler are nominally available through conda,
+our experience trying to install them through conda suggests that the "traditional"
+way of obtaining the [CUDA Toolkit](https://developer.nvidia.com/cuda-downloads) directly
+from NVIDIA may work best (see [related advice for Horovod](https://horovod.readthedocs.io/en/stable/conda_include.html)). After installing the CUDA Toolkit, one can set up the rest of the dependencies with:
+
+```bash
+conda create -n gpu-jax-finufft -c conda-forge python=3.10 numpy scipy fftw 'gxx<12'
+conda activate gpu-jax-finufft
+python -m pip install "jax[cuda11_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+CMAKE_ARGS=-DCMAKE_CUDA_ARCHITECTURES=70 python -m pip install .
+```
+
+In the last line, you'll need to select the CUDA architecture(s) you wish to compile for. See the [FINUFFT docs](https://finufft.readthedocs.io/en/latest/install_gpu.html#cmake-installation).
+
+For Flatiron users, the following environment setup script can be used instead of conda:
+<details>
+<summary>Environment script</summary>
+
+```bash
+ml modules/2.2
+ml gcc
+ml python/3.11
+ml fftw
+ml cuda/11
+ml cudnn
+ml nccl
+
+export LD_LIBRARY_PATH=$CUDA_HOME/extras/CUPTI/lib64:$LD_LIBRARY_PATH
+export CMAKE_ARGS="-DCMAKE_CUDA_ARCHITECTURES=60;70;80;90"
+```
+</details>
+
 ## Usage
 
 This library provides two high-level functions (and these should be all that you
diff --git a/cmake/FindFFTW.cmake b/cmake/FindFFTW.cmake
deleted file mode 100644
index ec1981e..0000000
--- a/cmake/FindFFTW.cmake
+++ /dev/null
@@ -1,419 +0,0 @@
-# - Find the FFTW library
-#
-# Original version of this file:
-#   Copyright (c) 2015, Wenzel Jakob
-#   https://github.com/wjakob/layerlab/blob/master/cmake/FindFFTW.cmake, commit 4d58bfdc28891b4f9373dfe46239dda5a0b561c6
-# Modifications:
-#   Copyright (c) 2017, Patrick Bos
-#
-# Usage:
-#   find_package(FFTW [REQUIRED] [QUIET] [COMPONENTS component1 ... componentX] )
-#
-# It sets the following variables:
-#   FFTW_FOUND                  ... true if fftw is found on the system
-#   FFTW_[component]_LIB_FOUND  ... true if the component is found on the system (see components below)
-#   FFTW_LIBRARIES              ... full paths to all found fftw libraries
-#   FFTW_[component]_LIB        ... full path to one of the components (see below)
-#   FFTW_INCLUDE_DIRS           ... fftw include directory paths
-#
-# The following variables will be checked by the function
-#   FFTW_USE_STATIC_LIBS        ... if true, only static libraries are found, otherwise both static and shared.
-#   FFTW_ROOT                   ... if set, the libraries are exclusively searched
-#                                   under this path
-#
-# This package supports the following components:
-#   FLOAT_LIB
-#   DOUBLE_LIB
-#   LONGDOUBLE_LIB
-#   FLOAT_THREADS_LIB
-#   DOUBLE_THREADS_LIB
-#   LONGDOUBLE_THREADS_LIB
-#   FLOAT_OPENMP_LIB
-#   DOUBLE_OPENMP_LIB
-#   LONGDOUBLE_OPENMP_LIB
-#
-
-# TODO (maybe): extend with ExternalProject download + build option
-# TODO: put on conda-forge
-
-
-if( NOT FFTW_ROOT AND DEFINED ENV{FFTWDIR} )
-    set( FFTW_ROOT $ENV{FFTWDIR} )
-endif()
-
-# Check if we can use PkgConfig
-find_package(PkgConfig)
-
-#Determine from PKG
-if( PKG_CONFIG_FOUND AND NOT FFTW_ROOT )
-    pkg_check_modules( PKG_FFTW QUIET "fftw3" )
-endif()
-
-#Check whether to search static or dynamic libs
-set( CMAKE_FIND_LIBRARY_SUFFIXES_SAV ${CMAKE_FIND_LIBRARY_SUFFIXES} )
-
-if( ${FFTW_USE_STATIC_LIBS} )
-    set( CMAKE_FIND_LIBRARY_SUFFIXES ${CMAKE_STATIC_LIBRARY_SUFFIX} )
-else()
-    set( CMAKE_FIND_LIBRARY_SUFFIXES ${CMAKE_FIND_LIBRARY_SUFFIXES_SAV} )
-endif()
-
-if( FFTW_ROOT )
-    # find libs
-
-    find_library(
-        FFTW_DOUBLE_LIB
-        NAMES "fftw3" libfftw3-3
-        PATHS ${FFTW_ROOT}
-        PATH_SUFFIXES "lib" "lib64"
-        NO_DEFAULT_PATH
-        )
-
-    find_library(
-        FFTW_DOUBLE_THREADS_LIB
-        NAMES "fftw3_threads"
-        PATHS ${FFTW_ROOT}
-        PATH_SUFFIXES "lib" "lib64"
-        NO_DEFAULT_PATH
-        )
-
-    find_library(
-        FFTW_DOUBLE_OPENMP_LIB
-        NAMES "fftw3_omp"
-        PATHS ${FFTW_ROOT}
-        PATH_SUFFIXES "lib" "lib64"
-        NO_DEFAULT_PATH
-        )
-
-    find_library(
-        FFTW_DOUBLE_MPI_LIB
-        NAMES "fftw3_mpi"
-        PATHS ${FFTW_ROOT}
-        PATH_SUFFIXES "lib" "lib64"
-        NO_DEFAULT_PATH
-        )
-
-    find_library(
-        FFTW_FLOAT_LIB
-        NAMES "fftw3f" libfftw3f-3
-        PATHS ${FFTW_ROOT}
-        PATH_SUFFIXES "lib" "lib64"
-        NO_DEFAULT_PATH
-        )
-
-    find_library(
-        FFTW_FLOAT_THREADS_LIB
-        NAMES "fftw3f_threads"
-        PATHS ${FFTW_ROOT}
-        PATH_SUFFIXES "lib" "lib64"
-        NO_DEFAULT_PATH
-        )
-
-    find_library(
-        FFTW_FLOAT_OPENMP_LIB
-        NAMES "fftw3f_omp"
-        PATHS ${FFTW_ROOT}
-        PATH_SUFFIXES "lib" "lib64"
-        NO_DEFAULT_PATH
-        )
-
-    find_library(
-        FFTW_FLOAT_MPI_LIB
-        NAMES "fftw3f_mpi"
-        PATHS ${FFTW_ROOT}
-        PATH_SUFFIXES "lib" "lib64"
-        NO_DEFAULT_PATH
-        )
-
-    find_library(
-        FFTW_LONGDOUBLE_LIB
-        NAMES "fftw3l" libfftw3l-3
-        PATHS ${FFTW_ROOT}
-        PATH_SUFFIXES "lib" "lib64"
-        NO_DEFAULT_PATH
-        )
-
-    find_library(
-        FFTW_LONGDOUBLE_THREADS_LIB
-        NAMES "fftw3l_threads"
-        PATHS ${FFTW_ROOT}
-        PATH_SUFFIXES "lib" "lib64"
-        NO_DEFAULT_PATH
-        )
-
-    find_library(
-        FFTW_LONGDOUBLE_OPENMP_LIB
-        NAMES "fftw3l_omp"
-        PATHS ${FFTW_ROOT}
-        PATH_SUFFIXES "lib" "lib64"
-        NO_DEFAULT_PATH
-        )
-
-    find_library(
-        FFTW_LONGDOUBLE_MPI_LIB
-        NAMES "fftw3l_mpi"
-        PATHS ${FFTW_ROOT}
-        PATH_SUFFIXES "lib" "lib64"
-        NO_DEFAULT_PATH
-        )
-
-    #find includes
-    find_path(FFTW_INCLUDE_DIRS
-        NAMES "fftw3.h"
-        PATHS ${FFTW_ROOT}
-        PATH_SUFFIXES "include"
-        NO_DEFAULT_PATH
-        )
-
-else()
-
-    find_library(
-        FFTW_DOUBLE_LIB
-        NAMES "fftw3"
-        PATHS ${PKG_FFTW_LIBRARY_DIRS} ${LIB_INSTALL_DIR}
-        )
-
-    find_library(
-        FFTW_DOUBLE_THREADS_LIB
-        NAMES "fftw3_threads"
-        PATHS ${PKG_FFTW_LIBRARY_DIRS} ${LIB_INSTALL_DIR}
-        )
-
-    find_library(
-        FFTW_DOUBLE_OPENMP_LIB
-        NAMES "fftw3_omp"
-        PATHS ${PKG_FFTW_LIBRARY_DIRS} ${LIB_INSTALL_DIR}
-        )
-
-    find_library(
-        FFTW_DOUBLE_MPI_LIB
-        NAMES "fftw3_mpi"
-        PATHS ${PKG_FFTW_LIBRARY_DIRS} ${LIB_INSTALL_DIR}
-        )
-
-    find_library(
-        FFTW_FLOAT_LIB
-        NAMES "fftw3f"
-        PATHS ${PKG_FFTW_LIBRARY_DIRS} ${LIB_INSTALL_DIR}
-        )
-
-    find_library(
-        FFTW_FLOAT_THREADS_LIB
-        NAMES "fftw3f_threads"
-        PATHS ${PKG_FFTW_LIBRARY_DIRS} ${LIB_INSTALL_DIR}
-        )
-
-    find_library(
-        FFTW_FLOAT_OPENMP_LIB
-        NAMES "fftw3f_omp"
-        PATHS ${PKG_FFTW_LIBRARY_DIRS} ${LIB_INSTALL_DIR}
-        )
-
-    find_library(
-        FFTW_FLOAT_MPI_LIB
-        NAMES "fftw3f_mpi"
-        PATHS ${PKG_FFTW_LIBRARY_DIRS} ${LIB_INSTALL_DIR}
-        )
-
-    find_library(
-        FFTW_LONGDOUBLE_LIB
-        NAMES "fftw3l"
-        PATHS ${PKG_FFTW_LIBRARY_DIRS} ${LIB_INSTALL_DIR}
-        )
-
-    find_library(
-        FFTW_LONGDOUBLE_THREADS_LIB
-        NAMES "fftw3l_threads"
-        PATHS ${PKG_FFTW_LIBRARY_DIRS} ${LIB_INSTALL_DIR}
-        )
-
-    find_library(FFTW_LONGDOUBLE_OPENMP_LIB
-        NAMES "fftw3l_omp"
-        PATHS ${PKG_FFTW_LIBRARY_DIRS} ${LIB_INSTALL_DIR}
-        )
-
-    find_library(FFTW_LONGDOUBLE_MPI_LIB
-        NAMES "fftw3l_mpi"
-        PATHS ${PKG_FFTW_LIBRARY_DIRS} ${LIB_INSTALL_DIR}
-        )
-
-    find_path(FFTW_INCLUDE_DIRS
-        NAMES "fftw3.h"
-        PATHS ${PKG_FFTW_INCLUDE_DIRS} ${INCLUDE_INSTALL_DIR}
-        )
-
-endif( FFTW_ROOT )
-
-#--------------------------------------- components
-
-if (FFTW_DOUBLE_LIB)
-    set(FFTW_DOUBLE_LIB_FOUND TRUE)
-    set(FFTW_LIBRARIES ${FFTW_LIBRARIES} ${FFTW_DOUBLE_LIB})
-    add_library(FFTW::Double INTERFACE IMPORTED)
-    set_target_properties(FFTW::Double
-        PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${FFTW_INCLUDE_DIRS}"
-        INTERFACE_LINK_LIBRARIES "${FFTW_DOUBLE_LIB}"
-        )
-else()
-    set(FFTW_DOUBLE_LIB_FOUND FALSE)
-endif()
-
-if (FFTW_FLOAT_LIB)
-    set(FFTW_FLOAT_LIB_FOUND TRUE)
-    set(FFTW_LIBRARIES ${FFTW_LIBRARIES} ${FFTW_FLOAT_LIB})
-    add_library(FFTW::Float INTERFACE IMPORTED)
-    set_target_properties(FFTW::Float
-        PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${FFTW_INCLUDE_DIRS}"
-        INTERFACE_LINK_LIBRARIES "${FFTW_FLOAT_LIB}"
-        )
-else()
-    set(FFTW_FLOAT_LIB_FOUND FALSE)
-endif()
-
-if (FFTW_LONGDOUBLE_LIB)
-    set(FFTW_LONGDOUBLE_LIB_FOUND TRUE)
-    set(FFTW_LIBRARIES ${FFTW_LIBRARIES} ${FFTW_LONGDOUBLE_LIB})
-    add_library(FFTW::LongDouble INTERFACE IMPORTED)
-    set_target_properties(FFTW::LongDouble
-        PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${FFTW_INCLUDE_DIRS}"
-        INTERFACE_LINK_LIBRARIES "${FFTW_LONGDOUBLE_LIB}"
-        )
-else()
-    set(FFTW_LONGDOUBLE_LIB_FOUND FALSE)
-endif()
-
-if (FFTW_DOUBLE_THREADS_LIB)
-    set(FFTW_DOUBLE_THREADS_LIB_FOUND TRUE)
-    set(FFTW_LIBRARIES ${FFTW_LIBRARIES} ${FFTW_DOUBLE_THREADS_LIB})
-    add_library(FFTW::DoubleThreads INTERFACE IMPORTED)
-    set_target_properties(FFTW::DoubleThreads
-        PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${FFTW_INCLUDE_DIRS}"
-        INTERFACE_LINK_LIBRARIES "${FFTW_DOUBLE_THREADS_LIB}"
-        )
-else()
-    set(FFTW_DOUBLE_THREADS_LIB_FOUND FALSE)
-endif()
-
-if (FFTW_FLOAT_THREADS_LIB)
-    set(FFTW_FLOAT_THREADS_LIB_FOUND TRUE)
-    set(FFTW_LIBRARIES ${FFTW_LIBRARIES} ${FFTW_FLOAT_THREADS_LIB})
-    add_library(FFTW::FloatThreads INTERFACE IMPORTED)
-    set_target_properties(FFTW::FloatThreads
-        PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${FFTW_INCLUDE_DIRS}"
-        INTERFACE_LINK_LIBRARIES "${FFTW_FLOAT_THREADS_LIB}"
-        )
-else()
-    set(FFTW_FLOAT_THREADS_LIB_FOUND FALSE)
-endif()
-
-if (FFTW_LONGDOUBLE_THREADS_LIB)
-    set(FFTW_LONGDOUBLE_THREADS_LIB_FOUND TRUE)
-    set(FFTW_LIBRARIES ${FFTW_LIBRARIES} ${FFTW_LONGDOUBLE_THREADS_LIB})
-    add_library(FFTW::LongDoubleThreads INTERFACE IMPORTED)
-    set_target_properties(FFTW::LongDoubleThreads
-        PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${FFTW_INCLUDE_DIRS}"
-        INTERFACE_LINK_LIBRARIES "${FFTW_LONGDOUBLE_THREADS_LIB}"
-        )
-else()
-    set(FFTW_LONGDOUBLE_THREADS_LIB_FOUND FALSE)
-endif()
-
-if (FFTW_DOUBLE_OPENMP_LIB)
-    set(FFTW_DOUBLE_OPENMP_LIB_FOUND TRUE)
-    set(FFTW_LIBRARIES ${FFTW_LIBRARIES} ${FFTW_DOUBLE_OPENMP_LIB})
-    add_library(FFTW::DoubleOpenMP INTERFACE IMPORTED)
-    set_target_properties(FFTW::DoubleOpenMP
-        PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${FFTW_INCLUDE_DIRS}"
-        INTERFACE_LINK_LIBRARIES "${FFTW_DOUBLE_OPENMP_LIB}"
-        )
-else()
-    set(FFTW_DOUBLE_OPENMP_LIB_FOUND FALSE)
-endif()
-
-if (FFTW_FLOAT_OPENMP_LIB)
-    set(FFTW_FLOAT_OPENMP_LIB_FOUND TRUE)
-    set(FFTW_LIBRARIES ${FFTW_LIBRARIES} ${FFTW_FLOAT_OPENMP_LIB})
-    add_library(FFTW::FloatOpenMP INTERFACE IMPORTED)
-    set_target_properties(FFTW::FloatOpenMP
-        PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${FFTW_INCLUDE_DIRS}"
-        INTERFACE_LINK_LIBRARIES "${FFTW_FLOAT_OPENMP_LIB}"
-        )
-else()
-    set(FFTW_FLOAT_OPENMP_LIB_FOUND FALSE)
-endif()
-
-if (FFTW_LONGDOUBLE_OPENMP_LIB)
-    set(FFTW_LONGDOUBLE_OPENMP_LIB_FOUND TRUE)
-    set(FFTW_LIBRARIES ${FFTW_LIBRARIES} ${FFTW_LONGDOUBLE_OPENMP_LIB})
-    add_library(FFTW::LongDoubleOpenMP INTERFACE IMPORTED)
-    set_target_properties(FFTW::LongDoubleOpenMP
-        PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${FFTW_INCLUDE_DIRS}"
-        INTERFACE_LINK_LIBRARIES "${FFTW_LONGDOUBLE_OPENMP_LIB}"
-        )
-else()
-    set(FFTW_LONGDOUBLE_OPENMP_LIB_FOUND FALSE)
-endif()
-
-if (FFTW_DOUBLE_MPI_LIB)
-    set(FFTW_DOUBLE_MPI_LIB_FOUND TRUE)
-    set(FFTW_LIBRARIES ${FFTW_LIBRARIES} ${FFTW_DOUBLE_MPI_LIB})
-    add_library(FFTW::DoubleMPI INTERFACE IMPORTED)
-    set_target_properties(FFTW::DoubleMPI
-        PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${FFTW_INCLUDE_DIRS}"
-        INTERFACE_LINK_LIBRARIES "${FFTW_DOUBLE_MPI_LIB}"
-        )
-else()
-    set(FFTW_DOUBLE_MPI_LIB_FOUND FALSE)
-endif()
-
-if (FFTW_FLOAT_MPI_LIB)
-    set(FFTW_FLOAT_MPI_LIB_FOUND TRUE)
-    set(FFTW_LIBRARIES ${FFTW_LIBRARIES} ${FFTW_FLOAT_MPI_LIB})
-    add_library(FFTW::FloatMPI INTERFACE IMPORTED)
-    set_target_properties(FFTW::FloatMPI
-        PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${FFTW_INCLUDE_DIRS}"
-        INTERFACE_LINK_LIBRARIES "${FFTW_FLOAT_MPI_LIB}"
-        )
-else()
-    set(FFTW_FLOAT_MPI_LIB_FOUND FALSE)
-endif()
-
-if (FFTW_LONGDOUBLE_MPI_LIB)
-    set(FFTW_LONGDOUBLE_MPI_LIB_FOUND TRUE)
-    set(FFTW_LIBRARIES ${FFTW_LIBRARIES} ${FFTW_LONGDOUBLE_MPI_LIB})
-    add_library(FFTW::LongDoubleMPI INTERFACE IMPORTED)
-    set_target_properties(FFTW::LongDoubleMPI
-        PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${FFTW_INCLUDE_DIRS}"
-        INTERFACE_LINK_LIBRARIES "${FFTW_LONGDOUBLE_MPI_LIB}"
-        )
-else()
-    set(FFTW_LONGDOUBLE_MPI_LIB_FOUND FALSE)
-endif()
-
-#--------------------------------------- end components
-
-set( CMAKE_FIND_LIBRARY_SUFFIXES ${CMAKE_FIND_LIBRARY_SUFFIXES_SAV} )
-
-include(FindPackageHandleStandardArgs)
-
-find_package_handle_standard_args(FFTW
-    REQUIRED_VARS FFTW_INCLUDE_DIRS
-    HANDLE_COMPONENTS
-    )
-
-mark_as_advanced(
-    FFTW_INCLUDE_DIRS
-    FFTW_LIBRARIES
-    FFTW_FLOAT_LIB
-    FFTW_DOUBLE_LIB
-    FFTW_LONGDOUBLE_LIB
-    FFTW_FLOAT_THREADS_LIB
-    FFTW_DOUBLE_THREADS_LIB
-    FFTW_LONGDOUBLE_THREADS_LIB
-    FFTW_FLOAT_OPENMP_LIB
-    FFTW_DOUBLE_OPENMP_LIB
-    FFTW_LONGDOUBLE_OPENMP_LIB
-    FFTW_FLOAT_MPI_LIB
-    FFTW_DOUBLE_MPI_LIB
-    FFTW_LONGDOUBLE_MPI_LIB
-    )


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 480a99a..4d3f996 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -16,19 +16,19 @@ jobs:
         os: [ubuntu-latest, macos-latest]
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           submodules: true
           fetch-depth: 0
 
       - name: Set up Python
-        uses: actions/setup-python@v1
+        uses: actions/setup-python@v4
         with:
           python-version: "3.9"
 
       - name: Install fftw on ubuntu
         if: ${{ matrix.os == 'ubuntu-latest' }}
-        run: sudo apt-get install libfftw3-dev
+        run: sudo apt-get install -y libfftw3-dev
 
       - name: Install fftw on macos
         if: ${{ matrix.os == 'macos-latest' }}
@@ -37,7 +37,7 @@ jobs:
       - name: Build
         run: |
           python -m pip install -U pip
-          python -m pip install .[test]
+          python -m pip install -v .[test]
 
       - name: Run tests
         run: python -m pytest -v tests
diff --git a/tests/ops_test.py b/tests/ops_test.py
index c60c62e..8e7d3f1 100644
--- a/tests/ops_test.py
+++ b/tests/ops_test.py
@@ -277,8 +277,8 @@ def test_multi_transform():
     calc1 = nufft1(n_k, c, x)
     calc2 = nufft2(f, x)
     for n in range(n_tr):
-        np.testing.assert_allclose(calc1[:, n], nufft1(n_k, c[:, n], x))
-        np.testing.assert_allclose(calc2[:, n], nufft2(f[:, n], x))
+        np.testing.assert_allclose(calc1[:, n], nufft1(n_k, c[:, n], x), rtol=1e-4)
+        np.testing.assert_allclose(calc2[:, n], nufft2(f[:, n], x), rtol=1e-4)
 
 
 def test_issue14():