ICODE-ADC/lammps
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #3963 from mkstoyanov/heffte_integration

Browse Source 
Heffte integration
This commit is contained in:
Axel Kohlmeyer
2023-12-12 17:27:02 -05:00
committed by GitHub
parent 75c3959cea e41d9cb746
commit 8104f68bc1
5 changed files with 182 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

72
cmake/CMakeLists.txt
View File

@ -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()

36
cmake/Modules/Packages/KSPACE.cmake
View File

@ -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})

44
doc/src/Build_settings.rst
View File

@ -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 <https://icl-utk-edu.github.io/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 <path-to-heffte-installation>/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 <https://github.com/mborgerding/kissfft>`_ 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 <https://mkstoyanov.bitbucket.io/heffte/md_doxygen_installation.html>`_
and those variables will be passed into the heFFTe build.
----------
.. _size:

36
src/KSPACE/fft3d_wrap.cpp
View File

@ -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<heffte_backend>();
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<heffte::fft3d<heffte_backend>>(
new heffte::fft3d<heffte_backend>(
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<std::complex<FFT_SCALAR>*>(in),
reinterpret_cast<std::complex<FFT_SCALAR>*>(out),
reinterpret_cast<std::complex<FFT_SCALAR>*>(heffte_workspace.data())
);
else
heffte_plan->backward(reinterpret_cast<std::complex<FFT_SCALAR>*>(in),
reinterpret_cast<std::complex<FFT_SCALAR>*>(out),
reinterpret_cast<std::complex<FFT_SCALAR>*>(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
}

20
src/KSPACE/fft3d_wrap.h
View File

@ -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::fft3d<heffte_backend>> heffte_plan;
std::vector<std::complex<FFT_SCALAR>> heffte_workspace;
heffte::scale hscale;
#else
struct fft_plan_3d *plan;
#endif
};
} // namespace LAMMPS_NS