diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 28e02bbee7..f7e9b314bd 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -971,33 +971,53 @@ if(PKG_KOKKOS) endif() endif() if(PKG_KSPACE) - message(STATUS "<<< FFT settings >>> --- Primary FFT lib: ${FFT}") - if(FFT_SINGLE) - message(STATUS "Using single precision FFTs") - else() - message(STATUS "Using double precision FFTs") - endif() - if(FFT_FFTW_THREADS OR FFT_MKL_THREADS) - message(STATUS "Using threaded FFTs") - else() - message(STATUS "Using non-threaded FFTs") - endif() - if(PKG_KOKKOS) - if(Kokkos_ENABLE_CUDA) - if(FFT STREQUAL "KISS") - message(STATUS "Kokkos FFT: KISS") - else() - message(STATUS "Kokkos FFT: cuFFT") - endif() - elseif(Kokkos_ENABLE_HIP) - if(FFT STREQUAL "KISS") - message(STATUS "Kokkos FFT: KISS") - else() - message(STATUS "Kokkos FFT: hipFFT") - endif() + if (LMP_HEFFTE) + message(STATUS "<<< FFT settings >>> +-- Primary FFT lib: heFFTe") + if (HEFFTE_BACKEND) + message(STATUS "heFFTe backend: ${HEFFTE_BACKEND}") else() - message(STATUS "Kokkos FFT: ${FFT}") + message(STATUS "heFFTe backend: stock (builtin FFT implementation, tested for corrected but not optimized for production)") + endif() + if(FFT_SINGLE) + message(STATUS "Using single precision FFTs") + else() + message(STATUS "Using double precision FFTs") + endif() + else() + message(STATUS "<<< FFT settings >>> +-- Primary FFT lib: ${FFT}") + if(FFT_SINGLE) + message(STATUS "Using single precision FFTs") + else() + message(STATUS "Using double precision FFTs") + endif() + if(FFT_FFTW_THREADS OR FFT_MKL_THREADS) + message(STATUS "Using threaded FFTs") + else() + message(STATUS "Using non-threaded FFTs") + endif() + if (FFT_HEFFTE) + message(STATUS "Using distributed algorithms from heFTTe") + else() + message(STATUS "Using builtin distributed algorithms") + endif() + if(PKG_KOKKOS) + if(Kokkos_ENABLE_CUDA) + if(FFT STREQUAL "KISS") + message(STATUS "Kokkos FFT: KISS") + else() + message(STATUS "Kokkos FFT: cuFFT") + endif() + elseif(Kokkos_ENABLE_HIP) + if(FFT STREQUAL "KISS") + message(STATUS "Kokkos FFT: KISS") + else() + message(STATUS "Kokkos FFT: hipFFT") + endif() + else() + message(STATUS "Kokkos FFT: ${FFT}") + endif() endif() endif() endif() diff --git a/cmake/Modules/Packages/KSPACE.cmake b/cmake/Modules/Packages/KSPACE.cmake index de7e7e5b20..9c9c879cd4 100644 --- a/cmake/Modules/Packages/KSPACE.cmake +++ b/cmake/Modules/Packages/KSPACE.cmake @@ -46,6 +46,42 @@ else() target_compile_definitions(lammps PRIVATE -DFFT_KISS) endif() +option(FFT_USE_HEFFTE "Use heFFTe as the distributed FFT engine, overrides the FFT option." OFF) +if(FFT_USE_HEFFTE) + # if FFT_HEFFTE is enabled, switch the builtin FFT engine with Heffte + set(FFT_HEFFTE_BACKEND_VALUES FFTW MKL) + set(FFT_HEFFTE_BACKEND "" CACHE STRING "Select heFFTe backend, e.g., FFTW or MKL") + set_property(CACHE FFT_HEFFTE_BACKEND PROPERTY STRINGS ${FFT_HEFFTE_BACKEND_VALUES}) + + if(FFT_HEFFTE_BACKEND STREQUAL "FFTW") # respect the backend choice, FFTW or MKL + set(HEFFTE_COMPONENTS "FFTW") + set(Heffte_ENABLE_FFTW "ON" CACHE BOOL "Enables FFTW backend for heFFTe") + elseif(FFT_HEFFTE_BACKEND STREQUAL "MKL") + set(HEFFTE_COMPONENTS "MKL") + set(Heffte_ENABLE_MKL "ON" CACHE BOOL "Enables MKL backend for heFFTe") + else() + message(WARNING "FFT_HEFFTE_BACKEND not selected, defaulting to the builtin 'stock' backend, which is intended for testing and is not optimized for production runs") + endif() + + find_package(Heffte 2.4.0 QUIET COMPONENTS ${HEFFTE_COMPONENTS}) + if (NOT Heffte_FOUND) # download and build + include(FetchContent) + FetchContent_Declare(HEFFTE_PROJECT # using v2.4.0 + URL "https://github.com/icl-utk-edu/heffte/archive/refs/tags/v2.4.0.tar.gz" + URL_HASH SHA256=02310fb4f9688df02f7181667e61c3adb7e38baf79611d80919d47452ff7881d + ) + FetchContent_Populate(HEFFTE_PROJECT) + add_subdirectory(${heffte_project_SOURCE_DIR} ${heffte_project_BINARY_DIR}) + set_target_properties(lmp PROPERTIES INSTALL_RPATH "${CMAKE_INSTALL_PREFIX}/lib") + set_target_properties(lammps PROPERTIES INSTALL_RPATH "${CMAKE_INSTALL_PREFIX}/lib") + add_library(Heffte::Heffte INTERFACE IMPORTED GLOBAL) + target_link_libraries(Heffte::Heffte INTERFACE Heffte) + endif() + + target_compile_definitions(lammps PRIVATE -DFFT_HEFFTE "-DFFT_HEFFTE_${FFT_HEFFTE_BACKEND}") + target_link_libraries(lammps PRIVATE Heffte::Heffte) +endif() + set(FFT_PACK "array" CACHE STRING "Optimization for FFT") set(FFT_PACK_VALUES array pointer memcpy) set_property(CACHE FFT_PACK PROPERTY STRINGS ${FFT_PACK_VALUES}) diff --git a/doc/src/Build_settings.rst b/doc/src/Build_settings.rst index 7576cae3eb..e2b096c6e0 100644 --- a/doc/src/Build_settings.rst +++ b/doc/src/Build_settings.rst @@ -43,6 +43,12 @@ When the KSPACE package is included in a LAMMPS build, the require use of an FFT library to compute 1d FFTs. The KISS FFT library is included with LAMMPS, but other libraries can be faster. LAMMPS can use them if they are available on your system. +Alternatively, LAMMPS can use the +`heFFTe `_ +library for the MPI communication algorithms, +which comes with many optimizations for special cases, +e.g., leveraging 2D and 3D backend transforms and +better pipelining for packing and communication. .. tabs:: @@ -53,6 +59,7 @@ LAMMPS can use them if they are available on your system. -D FFT=value # FFTW3 or MKL or KISS, default is FFTW3 if found, else KISS -D FFT_SINGLE=value # yes or no (default), no = double precision -D FFT_PACK=value # array (default) or pointer or memcpy + -D FFT_USE_HEFFTE=value # yes or no (default), yes links to heFFTe .. note:: @@ -76,6 +83,15 @@ LAMMPS can use them if they are available on your system. -D MKL_INCLUDE_DIR=path # ditto for Intel MKL library -D FFT_MKL_THREADS=on # enable using threaded FFTs with MKL libraries -D MKL_LIBRARY=path # path to MKL libraries + -D FFT_HEFFTE_BACKEND=value # FFTW or MKL or empty/undefined for the stock backend + -D Heffte_ROOT=path # path to an existing heFFTe installation + + .. note:: + + heFFTe comes with a builtin stock backend for FFTs; however, the backend + is intended for testing purposes and is not performance optimized + for large scale production runs. + .. tab:: Traditional make @@ -111,6 +127,24 @@ LAMMPS can use them if they are available on your system. files in its default search path. You must specify ``FFT_LIB`` with the appropriate FFT libraries to include in the link. + Traditional make can also link to heFFTe using an existing installation + + .. code-block:: make + + include /share/heffte/HeffteMakefile.in + FFT_INC = -DFFT_HEFFTE -DFFT_HEFFTE_FFTW $(heffte_include) + FFT_PATH = + FFT_LIB = $(heffte_link) $(heffte_libs) + + The heFFTe install path will contain `HeffteMakefile.in`. + which will define the `heffte_` include variables needed to link to heFFTe from + an external project using traditional make. + The `-DFFT_HEFFTE` is required to switch to using heFFTe, while the optional `-DFFT_HEFFTE_FFTW` + selects the desired heFFTe backend, e.g., `-DFFT_HEFFTE_FFTW` or `-DFFT_HEFFTE_MKL`, + omitting the variable will default to the `stock` backend. + The heFFTe `stock` backend is intended to be used for testing and debugging, + but is not performance optimized for large scale production runs. + The `KISS FFT library `_ is included in the LAMMPS distribution. It is portable across all platforms. Depending on the size of the FFTs and the number of @@ -170,6 +204,16 @@ Depending on the machine, the size of the FFT grid, the number of processors used, one option may be slightly faster. The default is ARRAY mode. +When using ``-DFFT_HEFFTE`` CMake will first look for an existing install +with hints provided by ``-DHeffte_ROOT``, as recommended by the CMake +standard and note that the name is case sensitive. If CMake cannot find +a heFFTe installation with the correct backend (e.g., FFTW or MKL), +it will attempt to download and build the library automatically. +In this case, LAMMPS CMake will also accept all heFFTe specific variables +listed in the +`heFFTe documentation `_ +and those variables will be passed into the heFFTe build. + ---------- .. _size: diff --git a/src/KSPACE/fft3d_wrap.cpp b/src/KSPACE/fft3d_wrap.cpp index 478cf6fc9d..7b00543eea 100644 --- a/src/KSPACE/fft3d_wrap.cpp +++ b/src/KSPACE/fft3d_wrap.cpp @@ -27,30 +27,66 @@ FFT3d::FFT3d(LAMMPS *lmp, MPI_Comm comm, int nfast, int nmid, int nslow, int out_klo, int out_khi, int scaled, int permute, int *nbuf, int usecollective) : Pointers(lmp) { + #ifndef FFT_HEFFTE plan = fft_3d_create_plan(comm,nfast,nmid,nslow, in_ilo,in_ihi,in_jlo,in_jhi,in_klo,in_khi, out_ilo,out_ihi,out_jlo,out_jhi,out_klo,out_khi, scaled,permute,nbuf,usecollective); if (plan == nullptr) error->one(FLERR,"Could not create 3d FFT plan"); + #else + heffte::plan_options options = heffte::default_options(); + options.algorithm = (usecollective == 0) ? + heffte::reshape_algorithm::p2p_plined + : heffte::reshape_algorithm::alltoallv; + options.use_reorder = (permute != 0); + hscale = (scaled == 0) ? heffte::scale::none : heffte::scale::full; + + heffte_plan = std::unique_ptr>( + new heffte::fft3d( + heffte::box3d<>({in_ilo,in_jlo,in_klo}, {in_ihi, in_jhi, in_khi}), + heffte::box3d<>({out_ilo,out_jlo,out_klo}, {out_ihi, out_jhi, out_khi}), + comm, options) + ); + *nbuf = heffte_plan->size_workspace(); + heffte_workspace.resize(heffte_plan->size_workspace()); + #endif } /* ---------------------------------------------------------------------- */ FFT3d::~FFT3d() { + #ifndef FFT_HEFFTE fft_3d_destroy_plan(plan); + #endif } /* ---------------------------------------------------------------------- */ void FFT3d::compute(FFT_SCALAR *in, FFT_SCALAR *out, int flag) { + #ifndef FFT_HEFFTE fft_3d((FFT_DATA *) in,(FFT_DATA *) out,flag,plan); + #else + if (flag == 1) + heffte_plan->forward(reinterpret_cast*>(in), + reinterpret_cast*>(out), + reinterpret_cast*>(heffte_workspace.data()) + ); + else + heffte_plan->backward(reinterpret_cast*>(in), + reinterpret_cast*>(out), + reinterpret_cast*>(heffte_workspace.data()), + hscale + ); + #endif } /* ---------------------------------------------------------------------- */ void FFT3d::timing1d(FFT_SCALAR *in, int nsize, int flag) { + #ifndef FFT_HEFFTE fft_1d_only((FFT_DATA *) in,nsize,flag,plan); + #endif } diff --git a/src/KSPACE/fft3d_wrap.h b/src/KSPACE/fft3d_wrap.h index f72cfd4622..04b828b7de 100644 --- a/src/KSPACE/fft3d_wrap.h +++ b/src/KSPACE/fft3d_wrap.h @@ -17,6 +17,19 @@ #include "fft3d.h" // IWYU pragma: export #include "pointers.h" +#ifdef FFT_HEFFTE +#include "heffte.h" +// select the backend +#if defined(FFT_HEFFTE_FFTW) +using heffte_backend = heffte::backend::fftw; +#elif defined(FFT_HEFFTE_MKL) +using heffte_backend = heffte::backend::mkl; +#else +using heffte_backend = heffte::backend::stock; +#endif + +#endif // FFT_HEFFTE + namespace LAMMPS_NS { class FFT3d : protected Pointers { @@ -30,7 +43,14 @@ class FFT3d : protected Pointers { void timing1d(FFT_SCALAR *, int, int); private: + #ifdef FFT_HEFFTE + // the heFFTe plan supersedes the internal fft_plan_3d + std::unique_ptr> heffte_plan; + std::vector> heffte_workspace; + heffte::scale hscale; + #else struct fft_plan_3d *plan; + #endif }; } // namespace LAMMPS_NS