Resolve "Add cuFINUFFT interface"

CMakeLists.txt

@@ -62,6 +62,13 @@ if (ENABLE_FFT)
     message (STATUS "Found Heffte_DIR: ${Heffte_DIR}")
 endif ()
 
+option (ENABLE_NUFFT "Enable NUFFT transform" OFF)
+if (ENABLE_NUFFT)
+    add_definitions (-DENABLE_NUFFT)
+    find_package(CUFINUFFT REQUIRED)
+    message (STATUS "Found CUFINUFFT_DIR: ${CUFINUFFT_DIR}")
+endif ()
+
 option (ENABLE_SOLVERS "Enable IPPL solvers" OFF)
 
 add_subdirectory (src)

CMakeModules/FindCUFINUFFT.cmake

@@ -0,0 +1,32 @@
+#
+# Find CUFINUFFT includes and library
+#
+# CUFINUFFT_INCLUDE_DIR - where to find cufinufft.h
+# CUFINUFFT_LIBRARY     - libcufinufft.so path
+# CUFINUFFT_FOUND       - do not attempt to use if "no" or undefined.
+
+FIND_PATH(CUFINUFFT_INCLUDE_DIR cufinufft.h
+    HINTS $ENV{CUFINUFFT_INCLUDE_PATH} $ENV{CUFINUFFT_INCLUDE_DIR} $ENV{CUFINUFFT_PREFIX}/include $ENV{CUFINUFFT_DIR}/include ${PROJECT_SOURCE_DIR}/include
+    PATHS ENV CPP_INCLUDE_PATH
+)
+#Static library has some issues and gives a cuda error at the end of compilation
+FIND_LIBRARY(CUFINUFFT_LIBRARY_DIR libcufinufft.so
+    HINTS $ENV{CUFINUFFT_LIBRARY_PATH} $ENV{CUFINUFFT_LIBRARY_DIR} $ENV{CUFINUFFT_PREFIX}/lib $ENV{CUFINUFFT_DIR}/lib $ENV{CUFINUFFT}/lib ${PROJECT_SOURCE_DIR}/lib
+    PATHS ENV LIBRARY_PATH
+)
+
+IF(CUFINUFFT_INCLUDE_DIR AND CUFINUFFT_LIBRARY_DIR)
+    SET( CUFINUFFT_FOUND "YES" )
+    SET( CUFINUFFT_DIR $ENV{CUFINUFFT_DIR} )
+ENDIF()
+
+IF (CUFINUFFT_FOUND)
+   IF (NOT CUFINUFFT_FIND_QUIETLY)
+       MESSAGE(STATUS "Found cufinufft library dir: ${CUFINUFFT_LIBRARY_DIR}")
+       MESSAGE(STATUS "Found cufinufft include dir: ${CUFINUFFT_INCLUDE_DIR}")
+   ENDIF (NOT CUFINUFFT_FIND_QUIETLY)
+ELSE (CUFINUFFT_FOUND)
+    IF (CUFINUFFT_FIND_REQUIRED)
+      MESSAGE(FATAL_ERROR "Could not find CUFINUFFT!")
+  ENDIF (CUFINUFFT_FIND_REQUIRED)
+ENDIF (CUFINUFFT_FOUND)

src/CMakeLists.txt

@@ -94,7 +94,14 @@ include_directories (
 
 add_library ( ippl ${IPPL_SRCS} ${IPPL_SRCS_FORT} )
 
-target_link_libraries(ippl PUBLIC Kokkos::kokkos ${HEFFTE_LIBRARY})
+
+if (ENABLE_NUFFT)
+    target_include_directories(ippl PUBLIC ${CUFINUFFT_INCLUDE_DIR})
+    target_link_libraries(ippl PUBLIC Kokkos::kokkos ${HEFFTE_LIBRARY} ${CUFINUFFT_LIBRARY_DIR})
+else()
+    target_link_libraries(ippl PUBLIC Kokkos::kokkos ${HEFFTE_LIBRARY})
+endif()
+
 
 install (TARGETS ippl DESTINATION lib)
 install (FILES ${IPPL_BASEDIR_HDRS} DESTINATION include)

src/FFT/FFT.h

@@ -30,12 +30,15 @@
 
 #include <heffte_fft3d.h>
 #include <heffte_fft3d_r2c.h>
+#include <cufinufft.h>
 #include <array>
 #include <memory>
+#include <functional>
 #include <type_traits>
 
 #include "FieldLayout/FieldLayout.h"
 #include "Field/Field.h"
+#include "Particle/ParticleAttrib.h"
 #include "Utility/ParameterList.h"
 #include "Utility/IpplException.h"
 
@@ -64,6 +67,12 @@ namespace ippl {
        Tag classes for Cosine transforms
     */
     class CosTransform {};
+#ifdef KOKKOS_ENABLE_CUDA
+    /**
+       Tag classes for Non-uniform type of Fourier transforms
+    */
+    class NUFFTransform {};
+#endif
 
     enum FFTComm {
         a2av = 0,
@@ -110,6 +119,37 @@ namespace ippl {
             using backendCos = heffte::backend::stock_cos;
         };
 #endif
+#endif
+
+#ifdef KOKKOS_ENABLE_CUDA
+        template <class T>
+        struct CufinufftType {};
+
+        template <>
+        struct CufinufftType<float> {
+            std::function<int(int, int, int*, int, int, 
+                              float, int, cufinufftf_plan*, cufinufft_opts*)> makeplan = cufinufftf_makeplan; 
+            std::function<int(int, float*, float*, float*, 
+                              int, float*, float*, float*, cufinufftf_plan)> setpts = cufinufftf_setpts; 
+            std::function<int(cuFloatComplex*, cuFloatComplex*, cufinufftf_plan)> execute = cufinufftf_execute; 
+            std::function<int(cufinufftf_plan)> destroy = cufinufftf_destroy;
+
+            using complexType = cuFloatComplex;
+            using plan_t      = cufinufftf_plan;
+        };
+
+        template <>
+        struct CufinufftType<double> {
+            std::function<int(int, int, int*, int, int, 
+                              double, int, cufinufft_plan*, cufinufft_opts*)> makeplan = cufinufft_makeplan; 
+            std::function<int(int, double*, double*, double*, 
+                              int, double*, double*, double*, cufinufft_plan)> setpts = cufinufft_setpts; 
+            std::function<int(cuDoubleComplex*, cuDoubleComplex*, cufinufft_plan)> execute = cufinufft_execute; 
+            std::function<int(cufinufft_plan)> destroy = cufinufft_destroy; 
+
+            using complexType = cuDoubleComplex;
+            using plan_t      = cufinufft_plan;
+        };
 #endif
     }
 
@@ -296,7 +336,56 @@ namespace ippl {
     };
 
 
+#ifdef KOKKOS_ENABLE_CUDA
+    /**
+       Non-uniform FFT class
+    */
+    template <size_t Dim, class T>
+    class FFT<NUFFTransform,Dim,T> {
+
+    public:
+
+        typedef FieldLayout<Dim> Layout_t;
+        typedef Kokkos::complex<T> KokkosComplex_t;
+        typedef Field<KokkosComplex_t,Dim> ComplexField_t;
+
+        using complexType = typename detail::CufinufftType<T>::complexType;
+        using plan_t = typename detail::CufinufftType<T>::plan_t;
+
+        /** Create a new FFT object with the layout for the input Field, type 
+         * (1 or 2) for the NUFFT and parameters for cuFINUFFT.
+        */
+        FFT(const Layout_t& layout, int type, const ParameterList& params);
+
+        // Destructor
+        ~FFT();
+
+        /** Do the NUFFT.
+        */
+        template<class PT1, class PT2, class... Properties>
+        void transform(const ParticleAttrib< Vector<PT1, Dim>, Properties... >& R, 
+                       ParticleAttrib<PT2, Properties... >& Q, ComplexField_t& f);
+
+
+    private:
+
+        /**
+           setup performs the initialization necessary.
+        */
+        void setup(std::array<int, 3>& nmodes,
+                   const ParameterList& params);
+
+        detail::CufinufftType<T> nufft_m;
+        plan_t plan_m;
+        int ier_m;
+        T tol_m;
+        int type_m;
+
+    };
+
+
 }
+#endif
 #include "FFT/FFT.hpp"
 #endif // IPPL_FFT_FFT_H