From c1c3025bde19881e7778d293c2b3d0df49f43b95 Mon Sep 17 00:00:00 2001 From: Geoffroy Lesur Date: Sat, 8 Aug 2026 14:54:31 +0200 Subject: [PATCH 1/5] first implementation --- .gitmodules | 3 + CMakeLists.txt | 13 + src/gravity/CMakeLists.txt | 4 + src/gravity/gravity.cpp | 11 +- src/gravity/gravity.hpp | 2 +- src/gravity/selfGravity.cpp | 453 +----------------- src/gravity/selfGravity.hpp | 49 +- src/gravity/selfGravityFFT.cpp | 237 +++++++++ src/gravity/selfGravityFFT.hpp | 51 ++ src/gravity/selfGravityIterative.cpp | 450 +++++++++++++++++ src/gravity/selfGravityIterative.hpp | 49 ++ src/kokkos-fft | 1 + src/real_types.hpp | 2 + src/timeIntegrator.cpp | 8 +- src/utils/fft/CMakeLists.txt | 5 + src/utils/fft/fft.cpp | 279 +++++++++++ src/utils/fft/fft.hpp | 107 +++++ src/utils/fft/transpose.hpp | 236 +++++++++ src/utils/iterativesolver/iterativesolver.hpp | 2 +- 19 files changed, 1481 insertions(+), 481 deletions(-) create mode 100644 src/gravity/selfGravityFFT.cpp create mode 100644 src/gravity/selfGravityFFT.hpp create mode 100644 src/gravity/selfGravityIterative.cpp create mode 100644 src/gravity/selfGravityIterative.hpp create mode 160000 src/kokkos-fft create mode 100644 src/utils/fft/CMakeLists.txt create mode 100644 src/utils/fft/fft.cpp create mode 100644 src/utils/fft/fft.hpp create mode 100644 src/utils/fft/transpose.hpp diff --git a/.gitmodules b/.gitmodules index 83fb2456f..00e8fa079 100644 --- a/.gitmodules +++ b/.gitmodules @@ -4,3 +4,6 @@ [submodule "reference"] path = reference url = https://github.com/idefix-code/reference +[submodule "src/kokkos-fft"] + path = src/kokkos-fft + url = https://github.com/kokkos/kokkos-fft.git diff --git a/CMakeLists.txt b/CMakeLists.txt index fad0058ba..15b5a6e1c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -16,6 +16,7 @@ option(Idefix_DEBUG "Enable Idefix debug features (makes the code very slow)" OF option(Idefix_RUNTIME_CHECKS "Enable runtime sanity checks" OFF) option(Idefix_WERROR "Treat compiler warnings as errors" OFF) option(Idefix_PYTHON "Enable python bindings (requires pybind11)" OFF) +option(Idefix_FFT "Enable FFT (requires dedicated FFT library)" OFF) set(Idefix_PROBLEM_DIR "${CMAKE_BINARY_DIR}" CACHE STRING "Problem directory to build for.") set(Idefix_CXX_FLAGS "" CACHE STRING "Additional compiler/linker flag") set(Idefix_DEFS "definitions.hpp" CACHE FILEPATH "Problem definition header file") @@ -54,6 +55,12 @@ endif() add_subdirectory(src/kokkos ${CMAKE_BINARY_DIR}/build/kokkos) include_directories(${Kokkos_INCLUDE_DIRS_RET}) +# Add FFT library if requested +if(Idefix_FFT) + add_compile_definitions("WITH_FFT") + add_subdirectory(src/kokkos-fft) +endif() + # Add Idefix CXX Flags add_compile_options(${Idefix_CXX_FLAGS}) @@ -248,12 +255,17 @@ target_include_directories(idefix PUBLIC src/rkl src/gravity src/utils + src/utils/fft src/utils/iterativesolver src/mpi src ) target_link_libraries(idefix Kokkos::kokkos) +if(Idefix_FFT) + target_link_libraries(idefix KokkosFFT::fft) + add_subdirectory(src/utils/fft) +endif() message(STATUS "Idefix final configuration") if(Idefix_EVOLVE_VECTOR_POTENTIAL) @@ -263,6 +275,7 @@ else() endif() message(STATUS " MPI: ${Idefix_MPI}") message(STATUS " HDF5: ${Idefix_HDF5}") +message(STATUS " FFT: ${Idefix_FFT}") message(STATUS " Python: ${Idefix_PYTHON}") message(STATUS " Reconstruction: ${Idefix_RECONSTRUCTION}") message(STATUS " Precision: ${Idefix_PRECISION}") diff --git a/src/gravity/CMakeLists.txt b/src/gravity/CMakeLists.txt index ad2585791..4ab0ab37a 100644 --- a/src/gravity/CMakeLists.txt +++ b/src/gravity/CMakeLists.txt @@ -5,4 +5,8 @@ target_sources(idefix PUBLIC ${CMAKE_CURRENT_LIST_DIR}/laplacian.hpp PUBLIC ${CMAKE_CURRENT_LIST_DIR}/selfGravity.hpp PUBLIC ${CMAKE_CURRENT_LIST_DIR}/selfGravity.cpp + PUBLIC ${CMAKE_CURRENT_LIST_DIR}/selfGravityIterative.hpp + PUBLIC ${CMAKE_CURRENT_LIST_DIR}/selfGravityIterative.cpp + PUBLIC ${CMAKE_CURRENT_LIST_DIR}/selfGravityFFT.hpp + PUBLIC ${CMAKE_CURRENT_LIST_DIR}/selfGravityFFT.cpp ) diff --git a/src/gravity/gravity.cpp b/src/gravity/gravity.cpp index a39b00063..80ed1521e 100644 --- a/src/gravity/gravity.cpp +++ b/src/gravity/gravity.cpp @@ -98,7 +98,10 @@ Gravity::Gravity(Input &input, DataBlock *datain) { // Check SelfGravity object if(haveSelfGravityPotential) { - selfGravity.Init(input, this->data); + if(!haveInitialisedSelfGravity) { + selfGravity = SelfGravity::Create(input, this->data); + } + selfGravity->Init(input, this->data); haveInitialisedSelfGravity = true; } @@ -125,7 +128,7 @@ void Gravity::ShowConfig() { } if(haveSelfGravityPotential) { idfx::cout << "Gravity: self-gravity ENABLED." << std::endl; - selfGravity.ShowConfig(); + selfGravity->ShowConfig(); } if(havePlanetsPotential) { idfx::cout << "Gravity: planet(s) potential ENABLED." << std::endl; @@ -165,10 +168,10 @@ void Gravity::ComputeGravity(int stepNumber) { } if(haveSelfGravityPotential) { // Solving Poisson for the current gas density distribution - if(stepNumber % selfGravity.skipSelfGravity == 0) selfGravity.SolvePoisson(); + if(stepNumber % selfGravity->skipSelfGravity == 0) selfGravity->SolvePoisson(); // Adding gas self-gravity contribution to global gravity potential - selfGravity.AddSelfGravityPotential(phiP); + selfGravity->AddSelfGravityPotential(phiP); } } if(haveBodyForce) { diff --git a/src/gravity/gravity.hpp b/src/gravity/gravity.hpp index aa0184aa6..71f29b096 100644 --- a/src/gravity/gravity.hpp +++ b/src/gravity/gravity.hpp @@ -52,7 +52,7 @@ class Gravity { IdefixArray4D bodyForceVector; // Self gravity - SelfGravity selfGravity; + std::unique_ptr selfGravity; // JM : moved in public class to handle changing centralMass during computation real centralMass{1.0}; ///< central mass parameter when central mass potential diff --git a/src/gravity/selfGravity.cpp b/src/gravity/selfGravity.cpp index 43f764713..292cc2f91 100644 --- a/src/gravity/selfGravity.cpp +++ b/src/gravity/selfGravity.cpp @@ -7,444 +7,27 @@ #include #include -#include #include "selfGravity.hpp" -#include "dataBlock.hpp" -#include "fluid.hpp" -#include "vector.hpp" -#include "bicgstab.hpp" -#include "cg.hpp" -#include "minres.hpp" -#include "jacobi.hpp" - - -void SelfGravity::Init(Input &input, DataBlock *datain) { - idfx::pushRegion("SelfGravity::Init"); - - // Save the parents data objects - this->data = datain; - - // Initialise (default) solver parameters - this->dt = 0.; - this->isPeriodic = true; - - // Update targetError when provided - real targetError = input.GetOrSet("SelfGravity","targetError",0,1e-2); - - // Get maxiter when provided - real maxiter = input.GetOrSet("SelfGravity","maxIter",0,1000); - - // Get the number of skipped cycles when provided and check consistency - this->skipSelfGravity = input.GetOrSet("SelfGravity","skip",0,1); - if(skipSelfGravity<1) { - IDEFIX_ERROR("[SelfGravity]:skip should be a strictly positive integer"); - } - - // Get the gravity-related boundary conditions - for (int dir = 0 ; dir < 3 ; dir++) { - this->lbound[dir] = Laplacian::LaplacianBoundaryType::undefined; - this->rbound[dir] = Laplacian::LaplacianBoundaryType::undefined; - } - for(int dir = 0 ; dir < DIMENSIONS ; dir++) { - std::string label = std::string("boundary-X")+std::to_string(dir+1)+std::string("-beg"); - std::string boundary = input.Get("SelfGravity",label,0); - - if(boundary.compare("nullpot") == 0) { - this->lbound[dir] = Laplacian::LaplacianBoundaryType::nullpot; - this->isPeriodic = false; - } else if(boundary.compare("periodic") == 0) { - this->lbound[dir] = Laplacian::LaplacianBoundaryType::periodic; - } else if(boundary.compare("nullgrad") == 0) { - this->lbound[dir] = Laplacian::LaplacianBoundaryType::nullgrad; - this->isPeriodic = false; - } else if(boundary.compare("internalgrav") == 0) { - this->lbound[dir] = Laplacian::LaplacianBoundaryType::internalgrav; - this->isPeriodic = false; - } else if(boundary.compare("userdef") == 0) { - this->lbound[dir] = Laplacian::LaplacianBoundaryType::userdef; - this->isPeriodic = false; - } else if(boundary.compare("axis") == 0) { - this->lbound[dir] = Laplacian::LaplacianBoundaryType::axis; - this->isPeriodic = false; - } else if(boundary.compare("origin") == 0) { - this->lbound[dir] = Laplacian::LaplacianBoundaryType::origin; - this->isPeriodic = false; - // origin only compatible with spherical & axis=IDIR - #if GEOMETRY != SPHERICAL - IDEFIX_ERROR("Origin boundary conditions are working in spherical coordinates"); - #endif - if(dir != IDIR) { - IDEFIX_ERROR("Origin boundary conditions are meaningful only on the X1 direction"); - } - } else { - std::stringstream msg; - msg << "SelfGravity:: Unknown boundary type " << boundary; - IDEFIX_ERROR(msg); - } - - label = std::string("boundary-X")+std::to_string(dir+1)+std::string("-end"); - boundary = input.Get("SelfGravity",label,0); - if(boundary.compare("nullpot") == 0) { - this->rbound[dir] = Laplacian::LaplacianBoundaryType::nullpot; - this->isPeriodic = false; - } else if(boundary.compare("periodic") == 0) { - this->rbound[dir] = Laplacian::LaplacianBoundaryType::periodic; - } else if(boundary.compare("nullgrad") == 0) { - this->rbound[dir] = Laplacian::LaplacianBoundaryType::nullgrad; - this->isPeriodic = false; - } else if(boundary.compare("internalgrav") == 0) { - this->rbound[dir] = Laplacian::LaplacianBoundaryType::internalgrav; - this->isPeriodic = false; - } else if(boundary.compare("userdef") == 0) { - this->rbound[dir] = Laplacian::LaplacianBoundaryType::userdef; - this->isPeriodic = false; - } else if(boundary.compare("axis") == 0) { - this->rbound[dir] = Laplacian::LaplacianBoundaryType::axis; - this->isPeriodic = false; - } else { - std::stringstream msg; - msg << "SelfGravity:: Unknown boundary type " << boundary; - IDEFIX_ERROR(msg); - } - } - - // Update solver when provided - if(input.CheckEntry("SelfGravity","solver") >= 0) { - std::string strSolver = input.Get("SelfGravity","solver",0); - if(strSolver.compare("Jacobi")==0) { - solver = JACOBI; - } else if(strSolver.compare("BICGSTAB")==0) { - solver = BICGSTAB; - } else if(strSolver.compare("PBICGSTAB")==0) { - solver = PBICGSTAB; - } else if(strSolver.compare("CG")==0) { - solver = CG; - } else if(strSolver.compare("PCG")==0) { - solver = PCG; - } else if(strSolver.compare("MINRES")==0) { - solver = MINRES; - } else if(strSolver.compare("PMINRES")==0) { - solver = PMINRES; - } else { - try { - // Try to use the old solver definition with integer (deprecated) - int s = std::stoi(strSolver); - if(s<0 || s > 2) throw std::runtime_error("Unknown solver number (should be 0,1 or 2)"); - this->solver = static_cast (s); - IDEFIX_DEPRECATED("The use of integer to define self-gravity solver is deprecated."); - } catch(const std::exception& e) { - std::stringstream msg; - msg << "SelfGravity: Unknown solver \"" << strSolver << "\"." - << "Use \"Jacobi\", \"BICGSTAB\" or \"PBICGSTAB\"." - << std::endl; - IDEFIX_ERROR(msg); - } - } +#include "selfGravityIterative.hpp" +#ifdef WITH_FFT + #include "selfGravityFFT.hpp" +#endif + +std::unique_ptr SelfGravity::Create(Input &input, DataBlock *data) { + std::string strSolver = input.GetOrSet("SelfGravity","solver",0,"BICGSTAB"); + + std::unique_ptr ptr; + if(strSolver == "FFT" || strSolver == "fft") { + #ifdef WITH_FFT + ptr = std::make_unique(); + #else + IDEFIX_ERROR("[SelfGravity]: FFT solver requested but Idefix was not compiled with FFT support."); + #endif } else { - this->solver = BICGSTAB; - } - - // Enable preconditionner - if(this->solver==PBICGSTAB || this->solver == PCG || this->solver == PMINRES) { - this->havePreconditioner = true; + ptr = std::make_unique(); } - // Make the Laplacian operator - laplacian = std::make_unique(data, lbound, rbound, this->havePreconditioner ); - - np_tot = laplacian->np_tot; - - // Instantiate the bicgstab solver - if(solver == BICGSTAB || solver == PBICGSTAB) { - iterativeSolver = new Bicgstab(*laplacian.get(), targetError, maxiter, - laplacian->np_tot, laplacian->beg, laplacian->end); - } else if(solver == CG || solver == PCG) { - iterativeSolver = new Cg(*laplacian.get(), targetError, maxiter, - laplacian->np_tot, laplacian->beg, laplacian->end); - } else if(solver == MINRES || solver == PMINRES) { - iterativeSolver = new Minres(*laplacian.get(), - targetError, maxiter, - laplacian->np_tot, laplacian->beg, laplacian->end); - } else { - real step = laplacian->ComputeCFL(); - iterativeSolver = new Jacobi(*laplacian.get(), targetError, maxiter, step, - laplacian->np_tot, laplacian->beg, laplacian->end); - } - - - // Arrays initialisation - this->density = IdefixArray3D ("Density", this->np_tot[KDIR], - this->np_tot[JDIR], - this->np_tot[IDIR]); - // Fill density array with 0 - { - auto d = this->density; - idefix_for("InitDensity",0,this->np_tot[KDIR],0,this->np_tot[JDIR],0,this->np_tot[IDIR], - KOKKOS_LAMBDA (int k, int j, int i) { - d(k,j,i) = 0.0; - }); - } - - this->potential = IdefixArray3D ("Potential", this->np_tot[KDIR], - this->np_tot[JDIR], - this->np_tot[IDIR]); - - - idfx::popRegion(); -} - - - -void SelfGravity::ShowConfig() { - idfx::cout << "SelfGravity: Using "; - switch(solver) { - case JACOBI: - idfx::cout << "Jacobi"; - break; - case BICGSTAB: - idfx::cout << "unpreconditionned BICGSTAB"; - break; - case PBICGSTAB: - idfx::cout << "preconditionned BICGSTAB"; - break; - case PCG: - idfx::cout << "preconditionned CG"; - break; - case CG: - idfx::cout << "unpreconditionned CG"; - break; - case MINRES: - idfx::cout << "unpreconditionned MinRes"; - break; - case PMINRES: - idfx::cout << "preconditionned MinRes"; - break; - default: - IDEFIX_ERROR("SelfGravity:: Unknown solver"); - } - idfx::cout << " solver." << std::endl; - // idfx::cout << "SelfGravity: target L2 norm error=" << targetError << "." << std::endl; - // idfx::cout << "SelfGravity: 4piG=" << gravCst << "." << std::endl; - - // The setup is periodic if it passes the previous boundary loading - if(this->isPeriodic == true) { - idfx::cout << "SelfGravity: Setup is periodic, using specific mass" - << " re-normalisation." << std::endl; - } - - if(this->lbound[IDIR] == Laplacian::LaplacianBoundaryType::origin) { - idfx::cout << "SelfGravity: using origin boundary with " << laplacian->loffset[IDIR] - << " additional radial points." << std::endl; - } - - if(this->skipSelfGravity>1) { - idfx::cout << "SelfGravity: self-gravity field will be updated every " << skipSelfGravity - << " cycles." << std::endl; - } - iterativeSolver->ShowConfig(); -} - - - -void SelfGravity::InitSolver() { - idfx::pushRegion("SelfGravity::InitSolver"); - - // Loading needed attributes - IdefixArray3D density = this->density; - IdefixArray4D Vc = data->hydro->Vc; - - // Initialise the density field - // todo: check bounds - int ioffset = laplacian->loffset[IDIR]; - int joffset = laplacian->loffset[JDIR]; - int koffset = laplacian->loffset[KDIR]; - - idefix_for("InitDensity", data->beg[KDIR], data->end[KDIR], - data->beg[JDIR], data->end[JDIR], - data->beg[IDIR], data->end[IDIR], - KOKKOS_LAMBDA (int k, int j, int i) { - density(k+koffset, j+joffset, i+ioffset) = Vc(RHO, k, j, i); - }); - - // Make sure that dust mass contributes to the self-gravitating field - if(data->haveDust) { - for(int i = 0 ; i < data->dust.size() ; i++) { - IdefixArray4D VcDust = data->dust[i]->Vc; - idefix_for("InitDustDensity", data->beg[KDIR], data->end[KDIR], - data->beg[JDIR], data->end[JDIR], - data->beg[IDIR], data->end[IDIR], - KOKKOS_LAMBDA (int k, int j, int i) { - density(k+koffset, j+joffset, i+ioffset) += VcDust(RHO, k, j, i); - }); - } - } - - // Deal with the mean issue for periodic density distribution - if(this->isPeriodic == true) { - SubstractMeanDensity(); // Remove density mean - } - - // divide density by preconditionner if we're doing the preconditionned version - if(havePreconditioner) { - int ibeg, iend, jbeg, jend, kbeg, kend; - ibeg = laplacian->beg[IDIR]; - iend = laplacian->end[IDIR]; - jbeg = laplacian->beg[JDIR]; - jend = laplacian->end[JDIR]; - kbeg = laplacian->beg[KDIR]; - kend = laplacian->end[KDIR]; - IdefixArray3D P = laplacian->precond; - idefix_for("Precond density", kbeg, kend, jbeg, jend, ibeg, iend, - KOKKOS_LAMBDA (int k, int j, int i) { - density(k, j, i) = density(k,j,i) / P(k,j,i); - }); - } - - // Look for Nans in the input field - int nanDensity = 0; - idefix_reduce("checkNanVc",0, this->np_tot[KDIR], 0, this->np_tot[JDIR], 0, this->np_tot[IDIR], - KOKKOS_LAMBDA (int k, int j, int i, int &nnan) { - if(std::isnan(density(k,j,i))) nnan++; - }, Kokkos::Sum(nanDensity) // reduction variable - ); - #ifdef WITH_MPI - MPI_Allreduce(MPI_IN_PLACE, &nanDensity,1,MPI_INT, MPI_SUM, MPI_COMM_WORLD); - #endif - - if(nanDensity>0) { - std::stringstream msg; - msg << "Input density in self-gravity contains "<< nanDensity << " NaNs" << std::endl; - throw std::runtime_error(msg.str()); - } - - idfx::popRegion(); -} - - -void SelfGravity::SubstractMeanDensity() { - idfx::pushRegion("SelfGravity::SubstractMeanDensity"); - - // Loading needed attributes - IdefixArray3D density = this->density; - IdefixArray3D dV = laplacian->dV; - - int ibeg, iend, jbeg, jend, kbeg, kend; - ibeg = laplacian->beg[IDIR]; - iend = laplacian->end[IDIR]; - jbeg = laplacian->beg[JDIR]; - jend = laplacian->end[JDIR]; - kbeg = laplacian->beg[KDIR]; - kend = laplacian->end[KDIR]; - - // Do the reduction on a vector - MyVector meanDensityVector; - - // Sum the density over the grid, weighted by cell volume - // and compute the total grid volume as a normalisation constant - // both stored in a 2D reduction vector - idefix_reduce("SumWeightedRho", - kbeg, kend, - jbeg, jend, - ibeg, iend, - KOKKOS_LAMBDA (int k, int j, int i, MyVector &localVector) { - localVector.v[0] += density(k,j,i) * dV(k,j,i); - localVector.v[1] += dV(k,j,i); - }, - Kokkos::Sum(meanDensityVector)); - - // Reduction on the whole grid - #ifdef WITH_MPI - MPI_Allreduce(MPI_IN_PLACE, &meanDensityVector.v, 2, realMPI, MPI_SUM, MPI_COMM_WORLD); - #endif - - real mean = meanDensityVector.v[0] / meanDensityVector.v[1]; - - // Remove the mean value of the density field - idefix_for("SubstractMeanDensity", - 0, this->np_tot[KDIR], - 0, this->np_tot[JDIR], - 0, this->np_tot[IDIR], - KOKKOS_LAMBDA (int k, int j, int i) { - density(k, j, i) -= mean; - }); - - idfx::popRegion(); -} - -void SelfGravity::EnrollUserDefBoundary(Laplacian::UserDefBoundaryFunc myFunc) { - laplacian->EnrollUserDefBoundary(myFunc); -} - - - -void SelfGravity::SolvePoisson() { - idfx::pushRegion("SelfGravity::SolvePoisson"); - - Kokkos::Timer timer; - - elapsedTime -= timer.seconds(); - - InitSolver(); // (Re)initialise the solver - - this->nsteps = iterativeSolver->Solve(potential, density); - if (this->nsteps<0) { - idfx::cout << "SelfGravity:: BICGSTAB failed, resetting potential" << std::endl; - - // Look for Nans to explain the repetitive failing - if(data->CheckNan()>0) { - std::stringstream msg; - msg << "Nan found after BICGSTAB failed at time " << data->t << std::endl; - throw std::runtime_error(msg.str()); - } - - // Re-initialise potential - IdefixArray3D potential = this->potential; - - idefix_for("ResetPotential", - 0, this->np_tot[KDIR], - 0, this->np_tot[JDIR], - 0, this->np_tot[IDIR], - KOKKOS_LAMBDA (int k, int j, int i) { - potential(k, j, i) = ZERO_F; - }); - - // Try again ! - this->nsteps = iterativeSolver->Solve(this->potential, density); - if (this->nsteps<0) { - IDEFIX_ERROR("SelfGravity:: BICGSTAB failed despite restart"); - } - } - - currentError = iterativeSolver->GetError(); - - - elapsedTime += timer.seconds(); - idfx::popRegion(); -} - -void SelfGravity::AddSelfGravityPotential(IdefixArray3D &phiP) { - idfx::pushRegion("SelfGravity::AddSelfGravityPotential"); - - // Loading needed data - IdefixArray3D localPot = phiP; - IdefixArray3D potential = this->potential; - real gravCst = this->data->gravity->gravCst; - - // Updating ghost cells before to return potential - laplacian->SetBoundaries(potential); - - // Adding self-gravity contribution - int ioffset = laplacian->loffset[IDIR]; - int joffset = laplacian->loffset[JDIR]; - int koffset = laplacian->loffset[KDIR]; - idefix_for("AddSelfGravityPotential", 0, data->np_tot[KDIR], - 0, data->np_tot[JDIR], - 0, data->np_tot[IDIR], - KOKKOS_LAMBDA (int k, int j, int i) { - // Takes into account the unit conversion, scaled by the choice of gravCst - localPot(k, j, i) += 4.*M_PI*gravCst * potential(k+koffset, j+joffset, i+ioffset); - }); - - idfx::popRegion(); + ptr->Init(input, data); + return ptr; } diff --git a/src/gravity/selfGravity.hpp b/src/gravity/selfGravity.hpp index badd729ce..ade6df533 100644 --- a/src/gravity/selfGravity.hpp +++ b/src/gravity/selfGravity.hpp @@ -14,59 +14,36 @@ #include "idefix.hpp" #include "input.hpp" #include "grid.hpp" -#include "fluid_defs.hpp" -#include "iterativesolver.hpp" #include "laplacian.hpp" #ifdef WITH_MPI #include "mpi.hpp" #endif -// Forward class hydro declaration class DataBlock; class SelfGravity { public: - enum GravitySolver {JACOBI, BICGSTAB, PBICGSTAB, PCG, CG, PMINRES, MINRES}; + enum GravitySolver {JACOBI, BICGSTAB, PBICGSTAB, PCG, CG, PMINRES, MINRES, FFTSolver}; - void Init(Input &, DataBlock *); // Initialisation of the class attributes - void ShowConfig(); // display current configuration - void InitSolver(); // (Re)initialisation of the solver for a given density distribution + virtual ~SelfGravity() = default; - void SubstractMeanDensity(); // Compute and substract the average input density + virtual void Init(Input &, DataBlock *) = 0; + virtual void ShowConfig() = 0; + virtual void SolvePoisson() = 0; + virtual void AddSelfGravityPotential(IdefixArray3D &) = 0; + virtual void EnrollUserDefBoundary(Laplacian::UserDefBoundaryFunc myFunc) = 0; - void SolvePoisson(); // Solve Poisson equation - void AddSelfGravityPotential(IdefixArray3D &); + static std::unique_ptr Create(Input &, DataBlock *); - void EnrollUserDefBoundary(Laplacian::UserDefBoundaryFunc myFunc); // User-defined boundary - - IterativeSolver *iterativeSolver; - - // The linear operator involved in Poisson equation - std::unique_ptr laplacian; - - real currentError{0}; // last error of the iterative solver - int nsteps{0}; // # of steps of the latest iteration - double elapsedTime; // time spent solving self gravity - - // Whether we should skip self-gravity computation every n steps + real currentError{0}; + int nsteps{0}; + double elapsedTime{0.0}; int skipSelfGravity{1}; - private: - DataBlock *data; // My parent data object - IdefixArray3D potential; // Gravitational potential - IdefixArray3D density; // Density - real dt; // CFL timestep - - // Local potential array size - std::array np_tot; - - std::array lbound; // Boundary condition to the left - std::array rbound; // Boundary condition to the right + protected: + DataBlock *data{nullptr}; - bool isPeriodic; - bool havePreconditioner{false}; - GravitySolver solver; // The solver used to solve Poisson }; #endif // GRAVITY_SELFGRAVITY_HPP_ diff --git a/src/gravity/selfGravityFFT.cpp b/src/gravity/selfGravityFFT.cpp new file mode 100644 index 000000000..931ec67d5 --- /dev/null +++ b/src/gravity/selfGravityFFT.cpp @@ -0,0 +1,237 @@ +// *********************************************************************************** +// Idefix MHD astrophysical code +// Copyright(C) Geoffroy R. J. Lesur +// and other code contributors +// Licensed under CeCILL 2.1 License, see COPYING for more information +// *********************************************************************************** + +#ifdef WITH_FFT + +#include +#include +#include +#include + +#include "idefix.hpp" +#include "selfGravityFFT.hpp" +#include "dataBlock.hpp" +#include "fluid.hpp" +#include "vector.hpp" +#include "grid.hpp" + +void SelfGravityFFT::Init(Input &input, DataBlock *datain) { + idfx::pushRegion("SelfGravityFFT::Init"); + this->data = datain; + Grid *grid = data->mygrid; + CheckCompatibility(); + + // FFT path currently requires periodic BC in all active dimensions + for (int dir = 0; dir < DIMENSIONS; dir++) { + if(grid->lbound[dir] != periodic || grid->rbound[dir] != periodic) { + IDEFIX_ERROR("[SelfGravityFFT]: FFT solver requires periodic boundary conditions in all active dimensions."); + } + } + + // storage with laplacian local array + rho = IdefixArray3D("Density", data->np_int[KDIR], data->np_tot[JDIR], data->np_tot[IDIR]); + phi = IdefixArray3D("Potential", data->np_tot[KDIR], data->np_tot[JDIR], data->np_tot[IDIR]); + + npr_glob = {grid->np_int[IDIR], grid->np_int[JDIR], grid->np_int[KDIR]}; + npr = {grid->np_int[IDIR], grid->np_int[JDIR], grid->np_int[KDIR]/idfx::psize}; + npf = {grid->np_int[IDIR]/2+1, grid->np_int[JDIR], grid->np_int[KDIR]/idfx::psize}; + npf_t = {grid->np_int[IDIR]/2+1, grid->np_int[KDIR], grid->np_int[JDIR]/idfx::psize}; + + for(int dir = 0; dir < 3; dir++) { + real d = (grid->xend[dir]-grid->xbeg[dir])/(2.0*M_PI*static_cast(npr_glob[dir])); + kx_glob[dir] = KokkosFFT::fftfreq(Device(), npr_glob[dir], d); + int p = idfx::prank; + if(dir != JDIR) kx[dir] = kx_glob[dir]; + else kx[dir] = Kokkos::subview(kx_glob[dir], std::pair(p * npf_t[KDIR], (p+1) * npf_t[KDIR])); + } + + if(idfx::psize > 1) { + rhoF = IdefixArray3D("rhoHatFFT", npf_t[KDIR], npf_t[JDIR], npf_t[IDIR]); + phiF = IdefixArray3D("phiHatFFT", npf_t[KDIR], npf_t[JDIR], npf_t[IDIR]); + } else { + rhoF = IdefixArray3D("rhoHatFFT", npf[KDIR], npf[JDIR], npf[IDIR]); + phiF = IdefixArray3D("phiHatFFT", npf[KDIR], npf[JDIR], npf[IDIR]); + } + + std::array nfft_real = {npr_glob[KDIR], npr_glob[JDIR], npr_glob[IDIR]}; + std::array nfft_complex = {npr_glob[KDIR], npr_glob[JDIR], npr_glob[IDIR]/2+1}; + + this->fft = std::make_unique(nfft_real, nfft_complex); + idfx::popRegion(); +} + +void SelfGravityFFT::CheckCompatibility() { + #if GEOMETRY != CARTESIAN + IDEFIX_ERROR("SelfGravityFFT supports CARTESIAN geometry only."); + #endif + Grid *grid = data->mygrid; + if(grid->np_int[IDIR] != data->np_int[IDIR] || grid->np_int[JDIR] != data->np_int[JDIR]) { + IDEFIX_ERROR("SelfGravityFFT requires no domain decomposition in the X1 and X2 directions."); + } + if(grid->np_int[KDIR] % idfx::psize != 0 || grid->np_int[JDIR] % idfx::psize != 0) { + IDEFIX_ERROR("SelfGravityFFT requires that the number of grid points in the X3 and X2 directions are" + " divisible by the number of MPI processes."); + } + if(grid->np_int[IDIR] % 2 != 0) { + IDEFIX_ERROR("SelfGravityFFT requires an even number of grid points in the X1 direction."); + } + +} + +void SelfGravityFFT::ShowConfig() { + idfx::cout << "SelfGravity: Using FFT Poisson solver (periodic Cartesian)." << std::endl; + if(skipSelfGravity>1) { + idfx::cout << "SelfGravity: self-gravity updated every " << skipSelfGravity + << " cycles." << std::endl; + } +} + +void SelfGravityFFT::EnforcePeriodic(int dir, BoundarySide side, IdefixArray3D &array) { + idfx::pushRegion("Laplacian::EnforceBoundary"); + + IdefixArray3D localVar = array; + + // Number of active cells + const int nxi = data->np_int[IDIR]; + const int nxj = data->np_int[JDIR]; + const int nxk = data->np_int[KDIR]; + + // Number of ghost cells + const int ighost = data->nghost[IDIR]; + const int jghost = data->nghost[JDIR]; + const int kghost = data->nghost[KDIR]; + + // Boundaries of the loop + const int ibeg = (dir == IDIR) ? side*(ighost+nxi) : 0; + const int iend = (dir == IDIR) ? ighost + side*(ighost+nxi) : data->np_tot[IDIR]; + const int jbeg = (dir == JDIR) ? side*(jghost+nxj) : 0; + const int jend = (dir == JDIR) ? jghost + side*(jghost+nxj) : data->np_tot[JDIR]; + const int kbeg = (dir == KDIR) ? side*(kghost+nxk) : 0; + const int kend = (dir == KDIR) ? kghost + side*(kghost+nxk) : data->np_tot[KDIR]; + + // Periodicity already enforced by MPI calls + if(data->mygrid->nproc[dir] == 1) { + idefix_for("BoundaryPeriodic", kbeg, kend, jbeg, jend, ibeg, iend, + KOKKOS_LAMBDA (int k, int j, int i) { + int iref, jref, kref; + // This hack takes care of cases where we have more ghost zones than active zones + if(dir==IDIR) + iref = ighost + (i+ighost*(nxi-1))%nxi; + else + iref = i; + if(dir==JDIR) + jref = jghost + (j+jghost*(nxj-1))%nxj; + else + jref = j; + if(dir==KDIR) + kref = kghost + (k+kghost*(nxk-1))%nxk; + else + kref = k; + + localVar(k,j,i) = localVar(kref,jref,iref); + }); + } + idfx::popRegion(); +} + +void SelfGravityFFT::SetBoundaries(IdefixArray3D &arr) { + idfx::pushRegion("SelfGravityFFT::SetBoundaries"); + + #ifdef WITH_MPI + this->arr4D = IdefixArray4D (arr.data(), 1, data->np_tot[KDIR], + data->np_tot[JDIR], + data->np_tot[IDIR]); + #endif + + for(int dir = 0 ; dir < DIMENSIONS ; dir++) { + // MPI Exchange data when needed + #ifdef WITH_MPI + if(data->mygrid->nproc[dir]>1) { + switch(dir) { + case 0: + this->mpi.ExchangeX1(this->arr4D); + break; + case 1: + this->mpi.ExchangeX2(this->arr4D); + break; + case 2: + this->mpi.ExchangeX3(this->arr4D); + break; + } + } + #endif + + EnforcePeriodic(dir, left, arr); + EnforcePeriodic(dir, right, arr); + } + + idfx::popRegion(); +} + +void SelfGravityFFT::SolvePoisson() { + idfx::pushRegion("SelfGravityFFT::SolvePoisson"); + Kokkos::Timer timer; + elapsedTime -= timer.seconds(); + + // Make a view omitting the ghost cells + auto rhoReal = Kokkos::subview(data->hydro->Vc, + RHO, + std::pair(data->beg[KDIR],data->end[KDIR]), + std::pair(data->beg[JDIR],data->end[JDIR]), + std::pair(data->beg[IDIR],data->end[IDIR])); + + auto rhoF = this->rhoF; + auto phiF = this->phiF; + // Forward transform of the density field + fft->R2C(rhoReal, rhoF, false); + + // Inverse poisson in Fourier space + auto kx1 = kx[IDIR]; + auto kx2 = kx[JDIR]; + auto kx3 = kx[KDIR]; + idefix_for("PoissonFFT", 0, npf[KDIR], 0, npf[JDIR], 0, npf[IDIR], + KOKKOS_LAMBDA(int k, int j, int i) { + const real k2 = kx1(i)*kx1(i) + kx2(j)*kx2(j) + kx3(k)*kx3(k); + real inv_k2 = (k2 > 0.0) ? -1.0/k2 : 0.0; + phiF(k,j,i) = rhoF(k,j,i) * inv_k2; + } + ); + + // Backward transform of the potential field + auto phiReal = Kokkos::subview(phi, + std::pair(data->beg[KDIR],data->end[KDIR]), + std::pair(data->beg[JDIR],data->end[JDIR]), + std::pair(data->beg[IDIR],data->end[IDIR])); + + fft->C2R(phiF, phiReal, false); + // Need to apply the boundary conditions to the potential field, since the FFT does not know about the ghost cells + SetBoundaries(phi); + elapsedTime += timer.seconds(); + + idfx::popRegion(); +} + +void SelfGravityFFT::AddSelfGravityPotential(IdefixArray3D &phiP) { + IdefixArray3D localPot = phiP; + IdefixArray3D pot = this->phi; + real gravCst = this->data->gravity->gravCst; + + + idefix_for("AddSelfGravityPotentialFFT", 0, data->np_tot[KDIR], + 0, data->np_tot[JDIR], + 0, data->np_tot[IDIR], + KOKKOS_LAMBDA (int k, int j, int i) { + localPot(k,j,i) += 4.*M_PI*gravCst * pot(k, j, i); + }); +} + +void SelfGravityFFT::EnrollUserDefBoundary(Laplacian::UserDefBoundaryFunc myFunc) { + (void) myFunc; + IDEFIX_ERROR("SelfGravityFFT only supports periodic boundaries; userdef boundary is unsupported."); +} + +#endif // WITH_FFT diff --git a/src/gravity/selfGravityFFT.hpp b/src/gravity/selfGravityFFT.hpp new file mode 100644 index 000000000..c5623ceea --- /dev/null +++ b/src/gravity/selfGravityFFT.hpp @@ -0,0 +1,51 @@ +// *********************************************************************************** +// Idefix MHD astrophysical code +// Copyright(C) Geoffroy R. J. Lesur +// and other code contributors +// Licensed under CeCILL 2.1 License, see COPYING for more information +// *********************************************************************************** + +#ifndef GRAVITY_SELFGRAVITYFFT_HPP_ +#define GRAVITY_SELFGRAVITYFFT_HPP_ + +#include +#include +#include + +#include "fft.hpp" +#include "selfGravity.hpp" + + +class SelfGravityFFT final : public SelfGravity { + public: + void Init(Input &, DataBlock *) override; + void ShowConfig() override; + void SolvePoisson() override; + void AddSelfGravityPotential(IdefixArray3D &) override; + void EnrollUserDefBoundary(Laplacian::UserDefBoundaryFunc myFunc) override; + void SetBoundaries(IdefixArray3D &); + void EnforcePeriodic(int dir, BoundarySide side, IdefixArray3D &); + private: + void SubstractMeanDensity(); + void CheckCompatibility(); + + std::unique_ptr fft; ///< FFT wrapper + + // FFT work arrays on active domain (no ghosts) + IdefixArray3D rho; + IdefixArray3D phi; + IdefixArray3D rhoF; + IdefixArray3D phiF; + std::array npf{1,1,1}; // [k,j,i] + std::array npf_t{1,1,1}; // [j,k,i] + std::array npr{1,1,1}; // [k,j,i] + std::array npr_glob{1,1,1}; // [k,j,i] + + std::array,3> kx_glob; + std::array,3> kx; + + std::array begin{0,0,0}; // [k,j,i] + std::array end{0,0,0}; // [k,j,i] +}; + +#endif // GRAVITY_SELFGRAVITYFFT_HPP_ diff --git a/src/gravity/selfGravityIterative.cpp b/src/gravity/selfGravityIterative.cpp new file mode 100644 index 000000000..e540ed2a5 --- /dev/null +++ b/src/gravity/selfGravityIterative.cpp @@ -0,0 +1,450 @@ +// *********************************************************************************** +// Idefix MHD astrophysical code +// Copyright(C) Geoffroy R. J. Lesur +// and other code contributors +// Licensed under CeCILL 2.1 License, see COPYING for more information +// *********************************************************************************** + +#include +#include +#include + +#include "selfGravityIterative.hpp" +#include "dataBlock.hpp" +#include "fluid.hpp" +#include "vector.hpp" +#include "bicgstab.hpp" +#include "cg.hpp" +#include "minres.hpp" +#include "jacobi.hpp" + + +void SelfGravityIterative::Init(Input &input, DataBlock *datain) { + idfx::pushRegion("SelfGravityIterative::Init"); + + // Save the parents data objects + this->data = datain; + + // Initialise (default) solver parameters + this->dt = 0.; + this->isPeriodic = true; + + // Update targetError when provided + real targetError = input.GetOrSet("SelfGravity","targetError",0,1e-2); + + // Get maxiter when provided + real maxiter = input.GetOrSet("SelfGravity","maxIter",0,1000); + + // Get the number of skipped cycles when provided and check consistency + this->skipSelfGravity = input.GetOrSet("SelfGravity","skip",0,1); + if(skipSelfGravity<1) { + IDEFIX_ERROR("[SelfGravity]:skip should be a strictly positive integer"); + } + + // Get the gravity-related boundary conditions + for (int dir = 0 ; dir < 3 ; dir++) { + this->lbound[dir] = Laplacian::LaplacianBoundaryType::undefined; + this->rbound[dir] = Laplacian::LaplacianBoundaryType::undefined; + } + for(int dir = 0 ; dir < DIMENSIONS ; dir++) { + std::string label = std::string("boundary-X")+std::to_string(dir+1)+std::string("-beg"); + std::string boundary = input.Get("SelfGravity",label,0); + + if(boundary.compare("nullpot") == 0) { + this->lbound[dir] = Laplacian::LaplacianBoundaryType::nullpot; + this->isPeriodic = false; + } else if(boundary.compare("periodic") == 0) { + this->lbound[dir] = Laplacian::LaplacianBoundaryType::periodic; + } else if(boundary.compare("nullgrad") == 0) { + this->lbound[dir] = Laplacian::LaplacianBoundaryType::nullgrad; + this->isPeriodic = false; + } else if(boundary.compare("internalgrav") == 0) { + this->lbound[dir] = Laplacian::LaplacianBoundaryType::internalgrav; + this->isPeriodic = false; + } else if(boundary.compare("userdef") == 0) { + this->lbound[dir] = Laplacian::LaplacianBoundaryType::userdef; + this->isPeriodic = false; + } else if(boundary.compare("axis") == 0) { + this->lbound[dir] = Laplacian::LaplacianBoundaryType::axis; + this->isPeriodic = false; + } else if(boundary.compare("origin") == 0) { + this->lbound[dir] = Laplacian::LaplacianBoundaryType::origin; + this->isPeriodic = false; + // origin only compatible with spherical & axis=IDIR + #if GEOMETRY != SPHERICAL + IDEFIX_ERROR("Origin boundary conditions are working in spherical coordinates"); + #endif + if(dir != IDIR) { + IDEFIX_ERROR("Origin boundary conditions are meaningful only on the X1 direction"); + } + } else { + std::stringstream msg; + msg << "SelfGravityIterative:: Unknown boundary type " << boundary; + IDEFIX_ERROR(msg); + } + + label = std::string("boundary-X")+std::to_string(dir+1)+std::string("-end"); + boundary = input.Get("SelfGravity",label,0); + if(boundary.compare("nullpot") == 0) { + this->rbound[dir] = Laplacian::LaplacianBoundaryType::nullpot; + this->isPeriodic = false; + } else if(boundary.compare("periodic") == 0) { + this->rbound[dir] = Laplacian::LaplacianBoundaryType::periodic; + } else if(boundary.compare("nullgrad") == 0) { + this->rbound[dir] = Laplacian::LaplacianBoundaryType::nullgrad; + this->isPeriodic = false; + } else if(boundary.compare("internalgrav") == 0) { + this->rbound[dir] = Laplacian::LaplacianBoundaryType::internalgrav; + this->isPeriodic = false; + } else if(boundary.compare("userdef") == 0) { + this->rbound[dir] = Laplacian::LaplacianBoundaryType::userdef; + this->isPeriodic = false; + } else if(boundary.compare("axis") == 0) { + this->rbound[dir] = Laplacian::LaplacianBoundaryType::axis; + this->isPeriodic = false; + } else { + std::stringstream msg; + msg << "SelfGravityIterative:: Unknown boundary type " << boundary; + IDEFIX_ERROR(msg); + } + } + + // Update solver when provided + if(input.CheckEntry("SelfGravity","solver") >= 0) { + std::string strSolver = input.Get("SelfGravity","solver",0); + if(strSolver.compare("Jacobi")==0) { + solver = JACOBI; + } else if(strSolver.compare("BICGSTAB")==0) { + solver = BICGSTAB; + } else if(strSolver.compare("PBICGSTAB")==0) { + solver = PBICGSTAB; + } else if(strSolver.compare("CG")==0) { + solver = CG; + } else if(strSolver.compare("PCG")==0) { + solver = PCG; + } else if(strSolver.compare("MINRES")==0) { + solver = MINRES; + } else if(strSolver.compare("PMINRES")==0) { + solver = PMINRES; + } else { + try { + // Try to use the old solver definition with integer (deprecated) + int s = std::stoi(strSolver); + if(s<0 || s > 2) throw std::runtime_error("Unknown solver number (should be 0,1 or 2)"); + this->solver = static_cast (s); + IDEFIX_DEPRECATED("The use of integer to define self-gravity solver is deprecated."); + } catch(const std::exception& e) { + std::stringstream msg; + msg << "SelfGravity: Unknown solver \"" << strSolver << "\"." + << "Use \"Jacobi\", \"BICGSTAB\" or \"PBICGSTAB\"." + << std::endl; + IDEFIX_ERROR(msg); + } + } + } else { + this->solver = BICGSTAB; + } + + // Enable preconditionner + if(this->solver==PBICGSTAB || this->solver == PCG || this->solver == PMINRES) { + this->havePreconditioner = true; + } + + // Make the Laplacian operator + laplacian = std::make_unique(data, lbound, rbound, this->havePreconditioner ); + + np_tot = laplacian->np_tot; + + // Instantiate the bicgstab solver + if(solver == BICGSTAB || solver == PBICGSTAB) { + iterativeSolver = std::make_unique>(*laplacian.get(), targetError, maxiter, + laplacian->np_tot, laplacian->beg, laplacian->end); + } else if(solver == CG || solver == PCG) { + iterativeSolver = std::make_unique>(*laplacian.get(), targetError, maxiter, + laplacian->np_tot, laplacian->beg, laplacian->end); + } else if(solver == MINRES || solver == PMINRES) { + iterativeSolver = std::make_unique>(*laplacian.get(), + targetError, maxiter, + laplacian->np_tot, laplacian->beg, laplacian->end); + } else { + real step = laplacian->ComputeCFL(); + iterativeSolver = std::make_unique>(*laplacian.get(), targetError, maxiter, step, + laplacian->np_tot, laplacian->beg, laplacian->end); + } + + + // Arrays initialisation + this->density = IdefixArray3D ("Density", this->np_tot[KDIR], + this->np_tot[JDIR], + this->np_tot[IDIR]); + // Fill density array with 0 + { + auto d = this->density; + idefix_for("InitDensity",0,this->np_tot[KDIR],0,this->np_tot[JDIR],0,this->np_tot[IDIR], + KOKKOS_LAMBDA (int k, int j, int i) { + d(k,j,i) = 0.0; + }); + } + + this->potential = IdefixArray3D ("Potential", this->np_tot[KDIR], + this->np_tot[JDIR], + this->np_tot[IDIR]); + + + idfx::popRegion(); +} + + + +void SelfGravityIterative::ShowConfig() { + idfx::cout << "SelfGravity: Using "; + switch(solver) { + case JACOBI: + idfx::cout << "Jacobi"; + break; + case BICGSTAB: + idfx::cout << "unpreconditionned BICGSTAB"; + break; + case PBICGSTAB: + idfx::cout << "preconditionned BICGSTAB"; + break; + case PCG: + idfx::cout << "preconditionned CG"; + break; + case CG: + idfx::cout << "unpreconditionned CG"; + break; + case MINRES: + idfx::cout << "unpreconditionned MinRes"; + break; + case PMINRES: + idfx::cout << "preconditionned MinRes"; + break; + default: + IDEFIX_ERROR("SelfGravityIterative:: Unknown solver"); + } + idfx::cout << " solver." << std::endl; + // idfx::cout << "SelfGravity: target L2 norm error=" << targetError << "." << std::endl; + // idfx::cout << "SelfGravity: 4piG=" << gravCst << "." << std::endl; + + // The setup is periodic if it passes the previous boundary loading + if(this->isPeriodic == true) { + idfx::cout << "SelfGravity: Setup is periodic, using specific mass" + << " re-normalisation." << std::endl; + } + + if(this->lbound[IDIR] == Laplacian::LaplacianBoundaryType::origin) { + idfx::cout << "SelfGravity: using origin boundary with " << laplacian->loffset[IDIR] + << " additional radial points." << std::endl; + } + + if(this->skipSelfGravity>1) { + idfx::cout << "SelfGravity: self-gravity field will be updated every " << skipSelfGravity + << " cycles." << std::endl; + } + iterativeSolver->ShowConfig(); +} + + + +void SelfGravityIterative::InitSolver() { + idfx::pushRegion("SelfGravityIterative::InitSolver"); + + // Loading needed attributes + IdefixArray3D density = this->density; + IdefixArray4D Vc = data->hydro->Vc; + + // Initialise the density field + // todo: check bounds + int ioffset = laplacian->loffset[IDIR]; + int joffset = laplacian->loffset[JDIR]; + int koffset = laplacian->loffset[KDIR]; + + idefix_for("InitDensity", data->beg[KDIR], data->end[KDIR], + data->beg[JDIR], data->end[JDIR], + data->beg[IDIR], data->end[IDIR], + KOKKOS_LAMBDA (int k, int j, int i) { + density(k+koffset, j+joffset, i+ioffset) = Vc(RHO, k, j, i); + }); + + // Make sure that dust mass contributes to the self-gravitating field + if(data->haveDust) { + for(int i = 0 ; i < data->dust.size() ; i++) { + IdefixArray4D VcDust = data->dust[i]->Vc; + idefix_for("InitDustDensity", data->beg[KDIR], data->end[KDIR], + data->beg[JDIR], data->end[JDIR], + data->beg[IDIR], data->end[IDIR], + KOKKOS_LAMBDA (int k, int j, int i) { + density(k+koffset, j+joffset, i+ioffset) += VcDust(RHO, k, j, i); + }); + } + } + + // Deal with the mean issue for periodic density distribution + if(this->isPeriodic == true) { + SubstractMeanDensity(); // Remove density mean + } + + // divide density by preconditionner if we're doing the preconditionned version + if(havePreconditioner) { + int ibeg, iend, jbeg, jend, kbeg, kend; + ibeg = laplacian->beg[IDIR]; + iend = laplacian->end[IDIR]; + jbeg = laplacian->beg[JDIR]; + jend = laplacian->end[JDIR]; + kbeg = laplacian->beg[KDIR]; + kend = laplacian->end[KDIR]; + IdefixArray3D P = laplacian->precond; + idefix_for("Precond density", kbeg, kend, jbeg, jend, ibeg, iend, + KOKKOS_LAMBDA (int k, int j, int i) { + density(k, j, i) = density(k,j,i) / P(k,j,i); + }); + } + + // Look for Nans in the input field + int nanDensity = 0; + idefix_reduce("checkNanVc",0, this->np_tot[KDIR], 0, this->np_tot[JDIR], 0, this->np_tot[IDIR], + KOKKOS_LAMBDA (int k, int j, int i, int &nnan) { + if(std::isnan(density(k,j,i))) nnan++; + }, Kokkos::Sum(nanDensity) // reduction variable + ); + #ifdef WITH_MPI + MPI_Allreduce(MPI_IN_PLACE, &nanDensity,1,MPI_INT, MPI_SUM, MPI_COMM_WORLD); + #endif + + if(nanDensity>0) { + std::stringstream msg; + msg << "Input density in self-gravity contains "<< nanDensity << " NaNs" << std::endl; + throw std::runtime_error(msg.str()); + } + + idfx::popRegion(); +} + + +void SelfGravityIterative::SubstractMeanDensity() { + idfx::pushRegion("SelfGravityIterative::SubstractMeanDensity"); + + // Loading needed attributes + IdefixArray3D density = this->density; + IdefixArray3D dV = laplacian->dV; + + int ibeg, iend, jbeg, jend, kbeg, kend; + ibeg = laplacian->beg[IDIR]; + iend = laplacian->end[IDIR]; + jbeg = laplacian->beg[JDIR]; + jend = laplacian->end[JDIR]; + kbeg = laplacian->beg[KDIR]; + kend = laplacian->end[KDIR]; + + // Do the reduction on a vector + MyVector meanDensityVector; + + // Sum the density over the grid, weighted by cell volume + // and compute the total grid volume as a normalisation constant + // both stored in a 2D reduction vector + idefix_reduce("SumWeightedRho", + kbeg, kend, + jbeg, jend, + ibeg, iend, + KOKKOS_LAMBDA (int k, int j, int i, MyVector &localVector) { + localVector.v[0] += density(k,j,i) * dV(k,j,i); + localVector.v[1] += dV(k,j,i); + }, + Kokkos::Sum(meanDensityVector)); + + // Reduction on the whole grid + #ifdef WITH_MPI + MPI_Allreduce(MPI_IN_PLACE, &meanDensityVector.v, 2, realMPI, MPI_SUM, MPI_COMM_WORLD); + #endif + + real mean = meanDensityVector.v[0] / meanDensityVector.v[1]; + + // Remove the mean value of the density field + idefix_for("SubstractMeanDensity", + 0, this->np_tot[KDIR], + 0, this->np_tot[JDIR], + 0, this->np_tot[IDIR], + KOKKOS_LAMBDA (int k, int j, int i) { + density(k, j, i) -= mean; + }); + + idfx::popRegion(); +} + +void SelfGravityIterative::EnrollUserDefBoundary(Laplacian::UserDefBoundaryFunc myFunc) { + laplacian->EnrollUserDefBoundary(myFunc); +} + + + +void SelfGravityIterative::SolvePoisson() { + idfx::pushRegion("SelfGravityIterative::SolvePoisson"); + + Kokkos::Timer timer; + + elapsedTime -= timer.seconds(); + + InitSolver(); // (Re)initialise the solver + + this->nsteps = iterativeSolver->Solve(potential, density); + if (this->nsteps<0) { + idfx::cout << "SelfGravityIterative:: BICGSTAB failed, resetting potential" << std::endl; + + // Look for Nans to explain the repetitive failing + if(data->CheckNan()>0) { + std::stringstream msg; + msg << "Nan found after BICGSTAB failed at time " << data->t << std::endl; + throw std::runtime_error(msg.str()); + } + + // Re-initialise potential + IdefixArray3D potential = this->potential; + + idefix_for("ResetPotential", + 0, this->np_tot[KDIR], + 0, this->np_tot[JDIR], + 0, this->np_tot[IDIR], + KOKKOS_LAMBDA (int k, int j, int i) { + potential(k, j, i) = ZERO_F; + }); + + // Try again ! + this->nsteps = iterativeSolver->Solve(this->potential, density); + if (this->nsteps<0) { + IDEFIX_ERROR("SelfGravityIterative:: BICGSTAB failed despite restart"); + } + } + + currentError = iterativeSolver->GetError(); + + + elapsedTime += timer.seconds(); + idfx::popRegion(); +} + +void SelfGravityIterative::AddSelfGravityPotential(IdefixArray3D &phiP) { + idfx::pushRegion("SelfGravityIterative::AddSelfGravityPotential"); + + // Loading needed data + IdefixArray3D localPot = phiP; + IdefixArray3D potential = this->potential; + real gravCst = this->data->gravity->gravCst; + + // Updating ghost cells before to return potential + laplacian->SetBoundaries(potential); + + // Adding self-gravity contribution + int ioffset = laplacian->loffset[IDIR]; + int joffset = laplacian->loffset[JDIR]; + int koffset = laplacian->loffset[KDIR]; + idefix_for("AddSelfGravityPotential", 0, data->np_tot[KDIR], + 0, data->np_tot[JDIR], + 0, data->np_tot[IDIR], + KOKKOS_LAMBDA (int k, int j, int i) { + // Takes into account the unit conversion, scaled by the choice of gravCst + localPot(k, j, i) += 4.*M_PI*gravCst * potential(k+koffset, j+joffset, i+ioffset); + }); + + idfx::popRegion(); +} diff --git a/src/gravity/selfGravityIterative.hpp b/src/gravity/selfGravityIterative.hpp new file mode 100644 index 000000000..8b25c54d2 --- /dev/null +++ b/src/gravity/selfGravityIterative.hpp @@ -0,0 +1,49 @@ +// *********************************************************************************** +// Idefix MHD astrophysical code +// Copyright(C) Geoffroy R. J. Lesur +// and other code contributors +// Licensed under CeCILL 2.1 License, see COPYING for more information +// *********************************************************************************** + +#ifndef GRAVITY_SELFGRAVITY_ITERATIVE_HPP_ +#define GRAVITY_SELFGRAVITY_ITERATIVE_HPP_ + +#include +#include +#include +#include "idefix.hpp" +#include "selfGravity.hpp" +#include "iterativesolver.hpp" +#include "laplacian.hpp" + +template class IterativeSolver; + +class SelfGravityIterative final : public SelfGravity { + public: + void Init(Input &, DataBlock *) override; + void ShowConfig() override; + void InitSolver(); + void SolvePoisson() override; + void AddSelfGravityPotential(IdefixArray3D &) override; + void EnrollUserDefBoundary(Laplacian::UserDefBoundaryFunc myFunc) override; + void SubstractMeanDensity(); + + std::unique_ptr laplacian; + std::unique_ptr> iterativeSolver; + + + private: + real dt; // CFL timestep + + std::array lbound; + std::array rbound; + bool havePreconditioner{false}; + GravitySolver solver{BICGSTAB}; + + std::array np_tot{0,0,0}; + IdefixArray3D potential; + IdefixArray3D density; + bool isPeriodic{true}; +}; + +#endif // GRAVITY_SELFGRAVITY_ITERATIVE_HPP_ diff --git a/src/kokkos-fft b/src/kokkos-fft new file mode 160000 index 000000000..ee53f8c68 --- /dev/null +++ b/src/kokkos-fft @@ -0,0 +1 @@ +Subproject commit ee53f8c68af9f4febe51c406312d4497368b0ba8 diff --git a/src/real_types.hpp b/src/real_types.hpp index 440c7aac9..f386c74c4 100644 --- a/src/real_types.hpp +++ b/src/real_types.hpp @@ -14,11 +14,13 @@ #ifdef SINGLE_PRECISION using real = float; + using complex = Kokkos::complex; #ifdef WITH_MPI #define realMPI MPI_FLOAT #endif #else using real = double; + using complex = Kokkos::complex; #ifdef WITH_MPI #define realMPI MPI_DOUBLE #endif diff --git a/src/timeIntegrator.cpp b/src/timeIntegrator.cpp index 7079700d9..a75327151 100644 --- a/src/timeIntegrator.cpp +++ b/src/timeIntegrator.cpp @@ -98,9 +98,9 @@ void TimeIntegrator::ShowLog(DataBlock &data) { #endif double sgOverhead; if(data.haveGravity && data.gravity->haveSelfGravityPotential) { - double sgCycleTime = data.gravity->selfGravity.elapsedTime - lastSGLog; + double sgCycleTime = data.gravity->selfGravity->elapsedTime - lastSGLog; sgOverhead = 100.0 * sgCycleTime / (timer.seconds() - lastLog); - lastSGLog = data.gravity->selfGravity.elapsedTime; + lastSGLog = data.gravity->selfGravity->elapsedTime; } lastLog = timer.seconds(); @@ -179,9 +179,9 @@ void TimeIntegrator::ShowLog(DataBlock &data) { } if(data.haveGravity && data.gravity->haveSelfGravityPotential) { if(ncycles>=cyclePeriod) { - idfx::cout << " | " << std::setw(col_width) << data.gravity->selfGravity.nsteps; + idfx::cout << " | " << std::setw(col_width) << data.gravity->selfGravity->nsteps; idfx::cout << std::scientific; - idfx::cout << " | " << std::setw(col_width) << data.gravity->selfGravity.currentError; + idfx::cout << " | " << std::setw(col_width) << data.gravity->selfGravity->currentError; idfx::cout << std::fixed; idfx::cout << " | " << std::setw(col_width) << sgOverhead; } else { diff --git a/src/utils/fft/CMakeLists.txt b/src/utils/fft/CMakeLists.txt new file mode 100644 index 000000000..0f0218143 --- /dev/null +++ b/src/utils/fft/CMakeLists.txt @@ -0,0 +1,5 @@ +target_sources(idefix + PUBLIC ${CMAKE_CURRENT_LIST_DIR}/fft.cpp + PUBLIC ${CMAKE_CURRENT_LIST_DIR}/fft.hpp + PUBLIC ${CMAKE_CURRENT_LIST_DIR}/transpose.hpp + ) diff --git a/src/utils/fft/fft.cpp b/src/utils/fft/fft.cpp new file mode 100644 index 000000000..a5dfb0a51 --- /dev/null +++ b/src/utils/fft/fft.cpp @@ -0,0 +1,279 @@ +// *********************************************************************************** +// Idefix MHD astrophysical code +// Copyright(C) Geoffroy R. J. Lesur +// and other code contributors +// Licensed under CeCILL 2.1 License, see COPYING for more information +// *********************************************************************************** + +#include +#include +#include + +#include "idefix.hpp" +#include "fft.hpp" +#include "global.hpp" +#include "loop.hpp" +#include "transpose.hpp" + +template void ShowExtent(T array) { + for(int i = 0 ; i < 3 ; i++) { + idfx::cout << array.extent(i) << " "; + } + idfx::cout << std::endl; +} + + // Empty constructor +FFT::FFT() {}; + +FFT::FFT(std::array npr_glob, std::array npf_glob) { + this->npf_glob = npf_glob; + this->npr_glob = npr_glob; + // Local dimensions + // We assume a decomposition along the first dimension only for now + this->npr = npr_glob; + this->npf = npf_glob; + this->npr_t = npr_glob; + + #ifdef WITH_MPI + this->npr[0] = npr_glob[0]/idfx::psize; + this->npf[0] = npf_glob[0]/idfx::psize; + this->npr_t[0] = npr_glob[1]/idfx::psize; + this->npr_t[1] = npr_glob[0]; + #endif + + tempReal = IdefixArray3D("FFT temp real", npr[0],npr[1],npr[2]); + tempComplex = IdefixArray3D("FFT temp complex", npf[0],npf[1],npf[2]); + + // Create the FFT plans + this->r2cPlan = std::make_unique(Kokkos::DefaultExecutionSpace(), + tempReal, tempComplex, KokkosFFT::Direction::forward, std::array{-3,-2,-1}); + this->c2rPlan = std::make_unique(Kokkos::DefaultExecutionSpace(), + tempComplex, tempReal, KokkosFFT::Direction::backward, std::array{-3,-2,-1}); + + #ifdef WITH_MPI + // Allocate temporary arrays for domain-splited FFTs and transposes + this->tempTransposedComplex = IdefixArray3D("FFT transpose temp", npf[1]/idfx::psize, npf[0]*idfx::psize, npf[2]); + this->tempTransposedComplex2 = IdefixArray3D("FFT transpose temp2", npf[1]/idfx::psize, npf[0]*idfx::psize, npf[2]); + this->tempTransposedReal = IdefixArray3D("FFT transpose temp2", npr[1]/idfx::psize, npr[0]*idfx::psize, npr[2]); + this->tempT2Complex = IdefixArray3D("FFT temp2 complex", npf[0],npf[2], npf[1]); + this->tempT2Complex2 = IdefixArray3D("FFT temp2 complex2", npf[0],npf[2], npf[1]); + IdefixArray3D tempComplex2 = idfx::makeArray>("FFT temp complex", npf); + + // MPI C2R Plans + // Axis 1 transposed is axis2 for the fft library + this->c2ciMPIPlan_axis2 = std::make_unique(Kokkos::DefaultExecutionSpace(), + tempT2Complex, tempT2Complex2, KokkosFFT::Direction::backward, std::array{-1}); + this->c2rMPIPlan_axis1t3 = std::make_unique(Kokkos::DefaultExecutionSpace(), + tempTransposedComplex, tempTransposedReal, KokkosFFT::Direction::backward, std::array{-2,-1}); + + // MPI R2C plans + this->r2cMPIPlan_axis1t3 = std::make_unique(Kokkos::DefaultExecutionSpace(), + tempTransposedReal, tempTransposedComplex, KokkosFFT::Direction::forward, std::array{-2,-1}); + this->c2cfMPIPlan_axis2 = std::make_unique(Kokkos::DefaultExecutionSpace(), + tempT2Complex, tempT2Complex2, KokkosFFT::Direction::forward, std::array{-1}); + + this->transposeComplex = std::make_unique>(npf); + this->transposeReal = std::make_unique>(npr); + + #endif + havePlan = true; + }; + +// Perform a real-to-complex FFT +void FFT::R2C(const IdefixArray3D in, IdefixArray3D out, bool transpose) { + idfx::pushRegion("FFT::R2C"); + + #ifdef WITH_MPI + R2C_MPI(in, out, transpose); + #else + // Ensure that in array is not erased + //Kokkos::deep_copy(tempReal, in); + if(havePlan) { + KokkosFFT::execute(*(r2cPlan.get()), in, out); + } else { + KokkosFFT::rfftn(Kokkos::DefaultExecutionSpace(), in, out); + } + #endif + idfx::popRegion(); +} + +void FFT::R2C_MPI(const IdefixArray3D in, IdefixArray3D out, bool transpose) { + idfx::pushRegion("FFT::R2C_MPI"); + + if(transpose) { + this->transposeReal->Apply(in,tempTransposedReal); + idfx::pushRegion("FFT::R2C_MPI axis1t3"); + KokkosFFT::execute(*(r2cMPIPlan_axis1t3.get()), tempTransposedReal, tempTransposedComplex); + idfx::popRegion(); + } else { + idfx::pushRegion("FFT::R2C_MPI axis1t3"); + KokkosFFT::execute(*(r2cMPIPlan_axis1t3.get()), in, tempTransposedComplex); + idfx::popRegion(); + } + this->transposeComplex->Apply(tempTransposedComplex,tempComplex); + TransposeLocal(tempComplex,tempT2Complex); + idfx::pushRegion("FFT::R2C_MPI axis2"); + KokkosFFT::execute(*(c2cfMPIPlan_axis2.get()), tempT2Complex, tempT2Complex2); + idfx::popRegion(); + TransposeLocal(tempT2Complex2,out); + idfx::popRegion(); +} + +// Perform a complex-to-real inverse FFT +void FFT::C2R(const IdefixArray3D in, IdefixArray3D out, bool transpose) { + idfx::pushRegion("FFT::C2R"); + #ifdef WITH_MPI + C2R_MPI(in,out,transpose); + #else + // Ensure that in array is not erased + Kokkos::deep_copy(tempComplex, in); + if(havePlan) { + KokkosFFT::execute(*(c2rPlan.get()), tempComplex, out); + } else { + KokkosFFT::irfftn(Kokkos::DefaultExecutionSpace(), tempComplex, out); + } + #endif + + idfx::popRegion(); +} + +void FFT::C2R_MPI(const IdefixArray3D in, IdefixArray3D out, bool transpose) { + idfx::pushRegion("FFT::C2R_MPI"); + idfx::pushRegion("FFT::C2R_MPI axis2"); + TransposeLocal(in,tempT2Complex); + KokkosFFT::execute(*(c2ciMPIPlan_axis2.get()), tempT2Complex, tempT2Complex2); + TransposeLocal(tempT2Complex2,tempComplex); + idfx::popRegion(); + this->transposeComplex->Apply(tempComplex,tempTransposedComplex); + if(transpose) { + idfx::pushRegion("FFT::C2R_MPI axis1t3"); + KokkosFFT::execute(*(c2rMPIPlan_axis1t3.get()), tempTransposedComplex, tempTransposedReal); + idfx::popRegion(); + this->transposeReal->Apply(tempTransposedReal,out); + } else { + idfx::pushRegion("FFT::C2R_MPI axis1t3"); + KokkosFFT::execute(*(c2rMPIPlan_axis1t3.get()), tempTransposedComplex, out); + idfx::popRegion(); + } + idfx::popRegion(); +} + +// FFT on Host, using the device. +void FFT::R2C_Host(const IdefixHostArray3D in, IdefixHostArray3D out) { + idfx::pushRegion("FFT::R2C_Host"); + IdefixArray3D inDev = Kokkos::create_mirror_view_and_copy(Kokkos::DefaultExecutionSpace(), in); + IdefixArray3D outDev = Kokkos::create_mirror_view(Kokkos::DefaultExecutionSpace(), out); + this->R2C(inDev, outDev,true); + Kokkos::deep_copy(out, outDev); + idfx::popRegion(); +} + +void FFT::C2R_Host(const IdefixHostArray3D in, IdefixHostArray3D out) { + idfx::pushRegion("FFT::C2R_Host"); + IdefixArray3D inDev = Kokkos::create_mirror_view_and_copy(Kokkos::DefaultExecutionSpace(), in); + IdefixArray3D outDev = Kokkos::create_mirror_view(Kokkos::DefaultExecutionSpace(), out); + this->C2R(inDev, outDev,true); + Kokkos::deep_copy(out, outDev); + idfx::popRegion(); +} + + +void FFT::TestMPI() { +#ifdef WITH_MPI + if(npr_glob[0] % idfx::psize != 0 || npr_glob[1] % idfx::psize != 0) { + throw std::runtime_error("Global problem size must be dividible by the number of MPI processes"); + } + + // Make an array with the local problem size + std::array npr = npr_glob; + npr[0] /= idfx::psize; + std::array npf = npr; + npf[2] = npr[2]/2+1; + std::array npf_glob = npr_glob; + npf_glob[2] = npr_glob[2]/2+1; + + Kokkos::View localReal_right("local real array", npr[0], npr[1], npr[2]); + IdefixArray3D localReal = idfx::makeArray>("local real array", npr); + + Kokkos::View globalReal_right("global real array", npr_glob[0], npr_glob[1], npr_glob[2]); + IdefixArray3D globalReal = idfx::makeArray>("global real array", npr_glob); + + // Compute a dummy real array + if(idfx::prank == 0) { + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(globalReal_right, random_pool, 1); + } + // Scatter to make local arrays + // And broadcast the global the array + MPI_Scatter(globalReal_right.data(), npr[0]*npr[1]*npr[2], + MPI_Astra_real, + localReal_right.data(), npr[0]*npr[1]*npr[2], + MPI_Astra_real, + 0, MPI_COMM_WORLD); + + MPI_Bcast(globalReal_right.data(), npr_glob[0]*npr_glob[1]*npr_glob[2], + MPI_Astra_real, 0, MPI_COMM_WORLD); + + // Change array Layout + idefix_for("Reshape global",0, npr_glob[0], + 0, npr_glob[1], + 0, npr_glob[2], + KOKKOS_LAMBDA(int i, int j, int k) { + globalReal(i,j,k) = globalReal_right(i,j,k); + }); + + // Change array Layout + idefix_for("Reshape local",0, npr[0], + 0, npr[1], + 0, npr[2], + KOKKOS_LAMBDA(int i, int j, int k) { + localReal(i,j,k) = localReal_right(i,j,k); + }); + + // Create the complex arrays + IdefixArray3D localComplex = idfx::makeArray>("local complex array", npf); + IdefixArray3D globalComplex = idfx::makeArray>("global complex array", npf_glob); + + // Compute the full serial fft + KokkosFFT::rfftn(Kokkos::DefaultExecutionSpace(), globalReal, globalComplex); + + // Compute the parallele fft + this->R2C_MPI(localReal, localComplex); + + // Check on a per-process that they all agree + int offset = npf[0]*idfx::prank; + + idefix_for("Reshape local",0, npf[0], + 0, npf[1], + 0, npf[2], + KOKKOS_LAMBDA(int i, int j, int k) { + int iglob = i+offset; + real norm = std::pow(localComplex(i,j,k).real()-globalComplex(iglob,j,k).real(),2) + +std::pow(localComplex(i,j,k).real()-globalComplex(iglob,j,k).real(),2); + + norm=std::sqrt(norm); + + if(norm > 1e-8) { + Kokkos::abort("incoherent values after MPI ifft"); + } + }); + + // Compute the parallele fft + this->C2R_MPI(localComplex, localReal); + + // Check on a per-process that they all agree + offset = npr[0]*idfx::prank; + + idefix_for("Reshape local",0, npr[0], + 0, npr[1], + 0, npr[2], + KOKKOS_LAMBDA(int i, int j, int k) { + int iglob = i+offset; + real norm = std::fabs(localReal(i,j,k)-globalReal(iglob,j,k)); + + if(norm > 1e-8) { + Kokkos::abort("incoherent values after MPI ifft"); + } + }); +#endif +} diff --git a/src/utils/fft/fft.hpp b/src/utils/fft/fft.hpp new file mode 100644 index 000000000..b49e8297e --- /dev/null +++ b/src/utils/fft/fft.hpp @@ -0,0 +1,107 @@ +// *********************************************************************************** +// Idefix MHD astrophysical code +// Copyright(C) Geoffroy R. J. Lesur +// and other code contributors +// Licensed under CeCILL 2.1 License, see COPYING for more information +// *********************************************************************************** + +// *********************************************************************************** +// ASTRA spectral code +// Accelerated Spectral code for TuRbulent plasmA +// Copyright(C) Geoffroy R. J. Lesur +// and other code contributors +// Licensed under CeCILL 2.1 License, see COPYING for more information +// *********************************************************************************** + +#ifndef FFT_HPP_ +#define FFT_HPP_ + +#include +#include +#include "arrays.hpp" +#include "transpose.hpp" + +// A class that wraps KokkosFFT functionality +using PlanR2CType = KokkosFFT::Plan, IdefixArray3D,3>; +using PlanC2RType = KokkosFFT::Plan, IdefixArray3D,3>; + +using PlanC2CType1D = KokkosFFT::Plan, IdefixArray3D,1>; +using PlanC2RType1D = KokkosFFT::Plan, IdefixArray3D,1>; +using PlanR2CType1D = KokkosFFT::Plan, IdefixArray3D,1>; + +using PlanR2CType2D = KokkosFFT::Plan, IdefixArray3D,2>; +using PlanC2RType2D = KokkosFFT::Plan, IdefixArray3D,2>; + +class FFT { + public: + // Empty constructor + FFT(); + + FFT(std::array npr, std::array npf); + + // Perform a real-to-complex FFT + void R2C(const IdefixArray3D in, IdefixArray3D out, bool transpose = true); + void R2C_MPI(const IdefixArray3D in, IdefixArray3D out, bool transpose = true); + + // Perform a complex-to-real inverse FFT + void C2R(const IdefixArray3D in, IdefixArray3D out, bool transpose = true); + void C2R_MPI(const IdefixArray3D in, IdefixArray3D out, bool transpose = true); + + // FFT on Host, using the device. + void R2C_Host(const IdefixHostArray3D in, IdefixHostArray3D out); + void C2R_Host(const IdefixHostArray3D in, IdefixHostArray3D out); + void TestMPI(); + + // Exchange last two dimensions of a 3D array + template + void TransposeLocal(const IdefixArray3D in, IdefixArray3D out); + + private: + bool havePlan{false}; + std::unique_ptr r2cPlan; + std::unique_ptr c2rPlan; + + // MPI plans + // Backard (C2R) + std::unique_ptr c2ciMPIPlan_axis2; + std::unique_ptr c2rMPIPlan_axis1t3; + + // Forward (R2C) + std::unique_ptr r2cMPIPlan_axis1t3; + std::unique_ptr c2cfMPIPlan_axis2; + + // Temporary arrays for MPI FFTs + IdefixArray3D tempComplex; + IdefixArray3D tempTransposedComplex; + IdefixArray3D tempTransposedComplex2; + IdefixArray3D tempT2Complex; + IdefixArray3D tempT2Complex2; + + IdefixArray3D tempTransposedReal; + IdefixArray3D tempReal; + + std::unique_ptr> transposeComplex; + std::unique_ptr> transposeReal; + + std::array npr; // Local real space dimensions + std::array npr_t; // Local real space dimensions after transpose + std::array npf; // Local fourier space dimensions + std::array npf_glob; // Global fourier space dimensions + std::array npr_glob; // Global real space dimensions +}; + +// exchange the last two dimensions of a 3D array +template +void FFT::TransposeLocal(const IdefixArray3D in, IdefixArray3D out) { + idfx::pushRegion("FFT::TransposeLocal"); + idefix_for("TransposeLocal",0, in.extent(0), + 0, in.extent(1), + 0, in.extent(2), + KOKKOS_LAMBDA(int i, int j, int k) { + out(i,k,j) = in(i,j,k); + }); + idfx::popRegion(); +} + + +#endif // FFT_HPP_ diff --git a/src/utils/fft/transpose.hpp b/src/utils/fft/transpose.hpp new file mode 100644 index 000000000..93a8fe8cd --- /dev/null +++ b/src/utils/fft/transpose.hpp @@ -0,0 +1,236 @@ +// *********************************************************************************** +// Idefix MHD astrophysical code +// Copyright(C) Geoffroy R. J. Lesur +// and other code contributors +// Licensed under CeCILL 2.1 License, see COPYING for more information +// *********************************************************************************** + +// *********************************************************************************** +// ASTRA spectral code +// Accelerated Spectral code for TuRbulent plasmA +// Copyright(C) Geoffroy R. J. Lesur +// and other code contributors +// Licensed under CeCILL 2.1 License, see COPYING for more information +// *********************************************************************************** + +#ifndef TRANSPOSE_HPP_ +#define TRANSPOSE_HPP_ + +#ifdef WITH_MPI +#include +#endif +#include +#include + + +template +class Transpose { + public: + explicit Transpose(std::array n) { + // Allocate temporary arrays for domain-splited FFTs and transposes + int64_t n1 = n[0]*n[1]; // Block size + int64_t nk = n[2]; + this->tempB = Kokkos::View("FFT transpose tempB", n1,nk); + this->tempC = Kokkos::View("FFT transpose tempC", n1,nk); + } + void Apply(const IdefixArray3D& input, IdefixArray3D& output); + void Test(); + + private: + Kokkos::View tempB, tempC; +}; + + + +/* MPI Transposition routines. +From complex to real, we first transform along x2 +Then we transpose x2 and x1 to have x1 contiguous in memory +We then transform along x1 and finally x3 (this last one being a real fft). This gives a final + +Transposition is done as follows: +Star with a matrix A_ijk +we require the first two indices to be dividible by n. so that (i,j) can be divided into n^2 +blocks. Hence the indices reads + +A_mi'pj'k +where i'=i'/n and j'=j/n while m = i%n and p = j%n +Initially the matrix is distributed accross MPI processes along m, + +The first step is to transpose locally pj'=j and i' +B_mpj'i'k = A_mi'pj'k + +Then we do an MPI_Alltoall to exchange the data between processors, this transposes m and p +C_pmj'i'k = B_mpj'i'k + +Finally, we do a local block transpose between j' and m +D_pj'mi'k = C_pmj'i'k + +The Matrix D is the final result with mi' contiguous in memory and pj' distributed accross MPI processes. + +Example: +Assuming a full domain of 6x4 elements. + +Initial layout contiguous along the second index (x2) and decomposed along the first (x1): +The dots represent the block division inside each MPI process whch is just represented to guide the eye +A11 A12 . A13 A14 +A21 A22 . A23 A24 +A31 A32 . A33 A34 +------------- MPI_Proc division ------------- +A41 A42 . A43 A44 +A51 A52 . A53 A54 +A61 A62 . A63 A64 + +1st step, local transposition (in each processor): +A11 A21 A31 +A12 A22 A32 + . . . +A13 A23 A33 +A14 A24 A34 +------------- MPI_Proc division ------------- +A41 A51 A61 +A42 A52 A62 + . . . +A43 A53 A63 +A44 A54 A64 + + +Then we do a MPI_AllToall to exchange the data between processors. +A11 A21 A31 +A12 A22 A32 + . . . +A41 A51 A61 +A42 A52 A62 +------------- MPI_Proc division ------------- +A13 A23 A33 +A14 A24 A34 + . . . +A43 A53 A63 +A44 A54 A64 + +Finally the block transpose +A11 A21 A31 . A41 A51 A61 +A12 A22 A32 . A42 A52 A62 +-------------- MPI_Proc division ------------- +A13 A23 A33 . A43 A53 A63 +A14 A24 A34 . A44 A54 A64 + + +After local transpose: + + +From real to complex, we first transform along x3 (real fft) then x1 (which is transposed, so this +is really the second coordinate). +We then transpose x1 and x2 to have x2 contiguous in memory and finally transform along x2. + + +Then along x3 (real fft) +we first transform along x3 and x2. +Then we need to transpose x1 and x2 to have x1 contiguous in memory. + + +*/ + +#include "idefix.hpp" + +template +void Transpose::Apply(const IdefixArray3D& in, IdefixArray3D& out) { + idfx::pushRegion("Transpose::Apply"); + #ifdef WITH_MPI + const int64_t n = idfx::psize; // number of MPI processes + const int64_t ni = in.extent(0); // Size of local block + const int64_t nj = in.extent(1)/n; + const int64_t nk = in.extent(2); + const int64_t n1 = ni*nj*n; // Full size of the first dimension + + + //Array2D tempB("FFT transpose temp", n1,nk); + auto tempB = this->tempB; + if(tempB.extent(0) != n1 || tempB.extent(1) != nk) { + idfx:: cout << "!! tempB size: " << tempB.extent(0) << "," << tempB.extent(1) << std::endl; + idfx:: cout << "!! in size: " << in.extent(0) << "," << in.extent(1) << "," << in.extent(2) << std::endl; + throw std::runtime_error("Transpose::Apply: temporary array size does not match input size"); + } + + idefix_for("FFT::TransposeLocal1", 0,n1, 0, nk, + KOKKOS_LAMBDA(const int64_t ij, const int64_t k) { + // for in array, ij = j'+nj*p +i'*nj*n= j+i'*nj*n (given that j=j'+nj*p) + // For tempB, ij=i'+ni*(j'+nj*p) + int64_t i = ij / (nj*n); + int64_t j = ij- i*nj*n; + int64_t ijprime = i + ni*j; + tempB(ijprime, k) = in(i,j,k); + }); + Kokkos::fence(); + + auto tempC = this->tempC; + //Kokkos::View tempC("FFT transpose temp", n1,nk); + int ret = MPI_Alltoall(tempB.data(), ni*nj*nk*sizeof(T), + MPI_BYTE, + tempC.data(), ni*nj*nk*sizeof(T), + MPI_BYTE, + MPI_COMM_WORLD); + + idefix_for("FFT::TransposeBlock", 0,n1, 0, nk, + KOKKOS_LAMBDA(const int64_t ij, const int64_t k) { + // For tempC, ij = i' + ni*(j'+m*nj) + // For out, ij = i' + m*ni + j'*ni*n (given that i = i'+m*ni) + + const int64_t j = ij / (ni*n); + const int64_t i = ij - j*ni*n; + const int64_t m = i / ni; + const int64_t iprime = i - m*ni; + const int64_t ijprime = iprime +ni*(j + m*nj); + + out(j,i,k) = tempC(ijprime,k); + }); + Kokkos::fence(); + + #endif + idfx::popRegion(); +} + +template +void Transpose::Test() { + idfx::cout << "Testing FFT Transpose" << std::endl; + const int n = idfx::psize; // number of MPI processes + const int rank = idfx::prank; // rank of the current process + const int ni = 3; // Size of local block + const int nj = 4; + const int nk = 9; + + + Transpose myTranspose({ni,nj*n,nk}); + IdefixArray3D in("FFT transpose test input", ni,nj*n,nk); + IdefixArray3D out("FFT transpose test output", nj,ni*n,nk); + + // Fill the input array with some test values + idefix_for("FFT::TestTransposeFill", 0,ni, 0,nj*n, 0,nk, + KOKKOS_LAMBDA(const int64_t i, const int64_t j, const int64_t k) { + in(i,j,k) = k+10*j+100*(i+rank*ni); + }); + myTranspose.Apply(in, out); + IdefixHostArray3D outHost = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(),out); + + // Check the output values + bool error = false; + for(int64_t j = 0 ; j < nj ; j++) { + for(int64_t i = 0 ; i < ni*n ; i++) { + for(int64_t k = 0 ; k < nk ; k++) { + if(outHost(j,i,k) != k+10*(j+rank*nj)+100*(i)) { + idfx::cout << "rank " << rank << " Transpose error at (" << i << "," << j << "," << k + << "): got " << outHost(j,i,k) << " instead of " << k+10*(j+rank*nj)+100*(i) << std::endl; + error = true; + } + } + } + } + + if(!error) { + idfx::cout << "Transpose test passed" << std::endl; + } else { + throw std::runtime_error("Transpose test failed @ rank " + std::to_string(rank)); + } +} + + +#endif // TRANSPOSE_HPP_ diff --git a/src/utils/iterativesolver/iterativesolver.hpp b/src/utils/iterativesolver/iterativesolver.hpp index e128f7fff..3b67dcba8 100644 --- a/src/utils/iterativesolver/iterativesolver.hpp +++ b/src/utils/iterativesolver/iterativesolver.hpp @@ -17,7 +17,7 @@ class IterativeSolver { public: IterativeSolver(T &op, real error, int maxIter, std::array ntot, std::array beg, std::array end); - + virtual ~IterativeSolver() = default; real GetError(); // return the current error of the solver virtual int Solve(IdefixArray3D &guess, IdefixArray3D &rhs) = 0; From 73aee17eaf0c6e11104150ef21df2c9f3c588253 Mon Sep 17 00:00:00 2001 From: Geoffroy Lesur Date: Sat, 8 Aug 2026 15:04:26 +0200 Subject: [PATCH 2/5] using multiple stride arrays --- src/utils/fft/fft.cpp | 54 +--------- src/utils/fft/fft.hpp | 204 ++++++++++++++++++++++++++---------- src/utils/fft/transpose.hpp | 8 -- 3 files changed, 149 insertions(+), 117 deletions(-) diff --git a/src/utils/fft/fft.cpp b/src/utils/fft/fft.cpp index a5dfb0a51..8cf3dd425 100644 --- a/src/utils/fft/fft.cpp +++ b/src/utils/fft/fft.cpp @@ -79,23 +79,7 @@ FFT::FFT(std::array npr_glob, std::array npf_glob) { havePlan = true; }; -// Perform a real-to-complex FFT -void FFT::R2C(const IdefixArray3D in, IdefixArray3D out, bool transpose) { - idfx::pushRegion("FFT::R2C"); - - #ifdef WITH_MPI - R2C_MPI(in, out, transpose); - #else - // Ensure that in array is not erased - //Kokkos::deep_copy(tempReal, in); - if(havePlan) { - KokkosFFT::execute(*(r2cPlan.get()), in, out); - } else { - KokkosFFT::rfftn(Kokkos::DefaultExecutionSpace(), in, out); - } - #endif - idfx::popRegion(); -} + void FFT::R2C_MPI(const IdefixArray3D in, IdefixArray3D out, bool transpose) { idfx::pushRegion("FFT::R2C_MPI"); @@ -119,23 +103,6 @@ void FFT::R2C_MPI(const IdefixArray3D in, IdefixArray3D out, bool idfx::popRegion(); } -// Perform a complex-to-real inverse FFT -void FFT::C2R(const IdefixArray3D in, IdefixArray3D out, bool transpose) { - idfx::pushRegion("FFT::C2R"); - #ifdef WITH_MPI - C2R_MPI(in,out,transpose); - #else - // Ensure that in array is not erased - Kokkos::deep_copy(tempComplex, in); - if(havePlan) { - KokkosFFT::execute(*(c2rPlan.get()), tempComplex, out); - } else { - KokkosFFT::irfftn(Kokkos::DefaultExecutionSpace(), tempComplex, out); - } - #endif - - idfx::popRegion(); -} void FFT::C2R_MPI(const IdefixArray3D in, IdefixArray3D out, bool transpose) { idfx::pushRegion("FFT::C2R_MPI"); @@ -158,25 +125,6 @@ void FFT::C2R_MPI(const IdefixArray3D in, IdefixArray3D out, bool idfx::popRegion(); } -// FFT on Host, using the device. -void FFT::R2C_Host(const IdefixHostArray3D in, IdefixHostArray3D out) { - idfx::pushRegion("FFT::R2C_Host"); - IdefixArray3D inDev = Kokkos::create_mirror_view_and_copy(Kokkos::DefaultExecutionSpace(), in); - IdefixArray3D outDev = Kokkos::create_mirror_view(Kokkos::DefaultExecutionSpace(), out); - this->R2C(inDev, outDev,true); - Kokkos::deep_copy(out, outDev); - idfx::popRegion(); -} - -void FFT::C2R_Host(const IdefixHostArray3D in, IdefixHostArray3D out) { - idfx::pushRegion("FFT::C2R_Host"); - IdefixArray3D inDev = Kokkos::create_mirror_view_and_copy(Kokkos::DefaultExecutionSpace(), in); - IdefixArray3D outDev = Kokkos::create_mirror_view(Kokkos::DefaultExecutionSpace(), out); - this->C2R(inDev, outDev,true); - Kokkos::deep_copy(out, outDev); - idfx::popRegion(); -} - void FFT::TestMPI() { #ifdef WITH_MPI diff --git a/src/utils/fft/fft.hpp b/src/utils/fft/fft.hpp index b49e8297e..9052c80f0 100644 --- a/src/utils/fft/fft.hpp +++ b/src/utils/fft/fft.hpp @@ -5,23 +5,15 @@ // Licensed under CeCILL 2.1 License, see COPYING for more information // *********************************************************************************** -// *********************************************************************************** -// ASTRA spectral code -// Accelerated Spectral code for TuRbulent plasmA -// Copyright(C) Geoffroy R. J. Lesur -// and other code contributors -// Licensed under CeCILL 2.1 License, see COPYING for more information -// *********************************************************************************** - #ifndef FFT_HPP_ #define FFT_HPP_ #include +#include #include #include "arrays.hpp" #include "transpose.hpp" -// A class that wraps KokkosFFT functionality using PlanR2CType = KokkosFFT::Plan, IdefixArray3D,3>; using PlanC2RType = KokkosFFT::Plan, IdefixArray3D,3>; @@ -34,74 +26,174 @@ using PlanC2RType2D = KokkosFFT::Plan npr_glob, std::array npf_glob); - FFT(std::array npr, std::array npf); + template + void R2C(const InView &in, const OutView &out, bool transpose=false); - // Perform a real-to-complex FFT - void R2C(const IdefixArray3D in, IdefixArray3D out, bool transpose = true); - void R2C_MPI(const IdefixArray3D in, IdefixArray3D out, bool transpose = true); + template + void C2R(const InView &in, const OutView &out, bool transpose=false); - // Perform a complex-to-real inverse FFT - void C2R(const IdefixArray3D in, IdefixArray3D out, bool transpose = true); - void C2R_MPI(const IdefixArray3D in, IdefixArray3D out, bool transpose = true); + template + void R2C_Host(const InView &in, const OutView &out); - // FFT on Host, using the device. - void R2C_Host(const IdefixHostArray3D in, IdefixHostArray3D out); - void C2R_Host(const IdefixHostArray3D in, IdefixHostArray3D out); - void TestMPI(); + template + void C2R_Host(const InView &in, const OutView &out); - // Exchange last two dimensions of a 3D array - template - void TransposeLocal(const IdefixArray3D in, IdefixArray3D out); + template + void TransposeLocal(const ViewIn &in, const ViewOut &out); - private: - bool havePlan{false}; - std::unique_ptr r2cPlan; - std::unique_ptr c2rPlan; + void R2C_MPI(const IdefixArray3D in, IdefixArray3D out, bool transpose=false); + void C2R_MPI(const IdefixArray3D in, IdefixArray3D out, bool transpose=false); + void TestMPI(); - // MPI plans - // Backard (C2R) - std::unique_ptr c2ciMPIPlan_axis2; - std::unique_ptr c2rMPIPlan_axis1t3; + private: + template + using view_t = std::decay_t; + + template + static constexpr bool is_view_v = Kokkos::is_view>::value; + + template + static constexpr bool is_rank3_v = is_view_v && (view_t::rank == 3); + + template + static constexpr bool has_scalar_v = + is_view_v && + std::is_same_v::non_const_value_type, Scalar>; + + template + static constexpr bool device_accessible_v = + is_view_v && + Kokkos::SpaceAccessibility< + Kokkos::DefaultExecutionSpace, + typename view_t::memory_space>::accessible; + + template + static constexpr bool host_accessible_v = + is_view_v && + Kokkos::SpaceAccessibility< + Kokkos::DefaultHostExecutionSpace, + typename view_t::memory_space>::accessible; - // Forward (R2C) - std::unique_ptr r2cMPIPlan_axis1t3; - std::unique_ptr c2cfMPIPlan_axis2; + public: + std::array npr_glob, npf_glob; + std::array npr, npf, npr_t; + bool havePlan{false}; - // Temporary arrays for MPI FFTs + IdefixArray3D tempReal; IdefixArray3D tempComplex; - IdefixArray3D tempTransposedComplex; - IdefixArray3D tempTransposedComplex2; - IdefixArray3D tempT2Complex; - IdefixArray3D tempT2Complex2; +#ifdef WITH_MPI + IdefixArray3D tempTransposedComplex, tempTransposedComplex2; IdefixArray3D tempTransposedReal; - IdefixArray3D tempReal; + IdefixArray3D tempT2Complex, tempT2Complex2; std::unique_ptr> transposeComplex; std::unique_ptr> transposeReal; - std::array npr; // Local real space dimensions - std::array npr_t; // Local real space dimensions after transpose - std::array npf; // Local fourier space dimensions - std::array npf_glob; // Global fourier space dimensions - std::array npr_glob; // Global real space dimensions + std::unique_ptr c2ciMPIPlan_axis2; + std::unique_ptr c2rMPIPlan_axis1t3; + std::unique_ptr r2cMPIPlan_axis1t3; + std::unique_ptr c2cfMPIPlan_axis2; +#endif + + std::unique_ptr r2cPlan; + std::unique_ptr c2rPlan; }; -// exchange the last two dimensions of a 3D array -template -void FFT::TransposeLocal(const IdefixArray3D in, IdefixArray3D out) { - idfx::pushRegion("FFT::TransposeLocal"); - idefix_for("TransposeLocal",0, in.extent(0), - 0, in.extent(1), - 0, in.extent(2), +template +void FFT::R2C(const InView &in, const OutView &out, bool transpose) { + static_assert(is_rank3_v, "FFT::R2C: input must be a rank-3 Kokkos::View"); + static_assert(is_rank3_v, "FFT::R2C: output must be a rank-3 Kokkos::View"); + static_assert(has_scalar_v, "FFT::R2C: input scalar type must be real"); + static_assert(has_scalar_v, "FFT::R2C: output scalar type must be complex"); + static_assert(device_accessible_v, "FFT::R2C: input view must be accessible from DefaultExecutionSpace"); + static_assert(device_accessible_v, "FFT::R2C: output view must be accessible from DefaultExecutionSpace"); + + idfx::pushRegion("FFT::R2C"); + +#ifdef WITH_MPI + Kokkos::deep_copy(tempReal, in); + R2C_MPI(tempReal, tempComplex, transpose); + Kokkos::deep_copy(out, tempComplex); +#else + Kokkos::deep_copy(tempReal, in); + KokkosFFT::execute(*(r2cPlan.get()), tempReal, tempComplex); + Kokkos::deep_copy(out, tempComplex); +#endif + + idfx::popRegion(); +} + +template +void FFT::C2R(const InView &in, const OutView &out, bool transpose) { + static_assert(is_rank3_v, "FFT::C2R: input must be a rank-3 Kokkos::View"); + static_assert(is_rank3_v, "FFT::C2R: output must be a rank-3 Kokkos::View"); + static_assert(has_scalar_v, "FFT::C2R: input scalar type must be complex"); + static_assert(has_scalar_v, "FFT::C2R: output scalar type must be real"); + static_assert(device_accessible_v, "FFT::C2R: input view must be accessible from DefaultExecutionSpace"); + static_assert(device_accessible_v, "FFT::C2R: output view must be accessible from DefaultExecutionSpace"); + + idfx::pushRegion("FFT::C2R"); + +#ifdef WITH_MPI + Kokkos::deep_copy(tempComplex, in); + C2R_MPI(tempComplex, tempReal, transpose); + Kokkos::deep_copy(out, tempReal); +#else + Kokkos::deep_copy(tempComplex, in); + KokkosFFT::execute(*(c2rPlan.get()), tempComplex, tempReal); + Kokkos::deep_copy(out, tempReal); +#endif + + idfx::popRegion(); +} + +template +void FFT::R2C_Host(const InView &in, const OutView &out) { + static_assert(is_rank3_v, "FFT::R2C_Host: input must be a rank-3 Kokkos::View"); + static_assert(is_rank3_v, "FFT::R2C_Host: output must be a rank-3 Kokkos::View"); + static_assert(has_scalar_v, "FFT::R2C_Host: input scalar type must be real"); + static_assert(has_scalar_v, "FFT::R2C_Host: output scalar type must be complex"); + static_assert(host_accessible_v, "FFT::R2C_Host: input view must be accessible from host"); + static_assert(host_accessible_v, "FFT::R2C_Host: output view must be accessible from host"); + + IdefixArray3D inDev = Kokkos::create_mirror_view_and_copy(Kokkos::DefaultExecutionSpace(), in); + IdefixArray3D outDev = Kokkos::create_mirror_view(Kokkos::DefaultExecutionSpace(), out); + R2C(inDev, outDev, true); + Kokkos::deep_copy(out, outDev); +} + +template +void FFT::C2R_Host(const InView &in, const OutView &out) { + static_assert(is_rank3_v, "FFT::C2R_Host: input must be a rank-3 Kokkos::View"); + static_assert(is_rank3_v, "FFT::C2R_Host: output must be a rank-3 Kokkos::View"); + static_assert(has_scalar_v, "FFT::C2R_Host: input scalar type must be complex"); + static_assert(has_scalar_v, "FFT::C2R_Host: output scalar type must be real"); + static_assert(host_accessible_v, "FFT::C2R_Host: input view must be accessible from host"); + static_assert(host_accessible_v, "FFT::C2R_Host: output view must be accessible from host"); + + IdefixArray3D inDev = Kokkos::create_mirror_view_and_copy(Kokkos::DefaultExecutionSpace(), in); + IdefixArray3D outDev = Kokkos::create_mirror_view(Kokkos::DefaultExecutionSpace(), out); + C2R(inDev, outDev, true); + Kokkos::deep_copy(out, outDev); +} + +template +void FFT::TransposeLocal(const ViewIn &in, const ViewOut &out) { + static_assert(is_rank3_v, "FFT::TransposeLocal: input must be a rank-3 Kokkos::View"); + static_assert(is_rank3_v, "FFT::TransposeLocal: output must be a rank-3 Kokkos::View"); + static_assert(device_accessible_v, "FFT::TransposeLocal: input view must be device accessible"); + static_assert(device_accessible_v, "FFT::TransposeLocal: output view must be device accessible"); + + idefix_for("TransposeLocal", 0, in.extent(0), + 0, in.extent(1), + 0, in.extent(2), KOKKOS_LAMBDA(int i, int j, int k) { out(i,k,j) = in(i,j,k); }); - idfx::popRegion(); } - -#endif // FFT_HPP_ +#endif // FFT_HPP_ \ No newline at end of file diff --git a/src/utils/fft/transpose.hpp b/src/utils/fft/transpose.hpp index 93a8fe8cd..fb00cc697 100644 --- a/src/utils/fft/transpose.hpp +++ b/src/utils/fft/transpose.hpp @@ -5,14 +5,6 @@ // Licensed under CeCILL 2.1 License, see COPYING for more information // *********************************************************************************** -// *********************************************************************************** -// ASTRA spectral code -// Accelerated Spectral code for TuRbulent plasmA -// Copyright(C) Geoffroy R. J. Lesur -// and other code contributors -// Licensed under CeCILL 2.1 License, see COPYING for more information -// *********************************************************************************** - #ifndef TRANSPOSE_HPP_ #define TRANSPOSE_HPP_ From 1064a52a6d1f3e118c720ed8a800e17741ab1700 Mon Sep 17 00:00:00 2001 From: Geoffroy Lesur Date: Sat, 8 Aug 2026 16:24:49 +0200 Subject: [PATCH 3/5] mpi version --- src/gravity/selfGravityFFT.cpp | 42 +++++++++++++++++++++++++--------- src/gravity/selfGravityFFT.hpp | 7 ++++++ src/utils/fft/fft.cpp | 16 ++++++------- src/utils/fft/fft.hpp | 1 + 4 files changed, 47 insertions(+), 19 deletions(-) diff --git a/src/gravity/selfGravityFFT.cpp b/src/gravity/selfGravityFFT.cpp index 931ec67d5..83c1bd680 100644 --- a/src/gravity/selfGravityFFT.cpp +++ b/src/gravity/selfGravityFFT.cpp @@ -61,6 +61,15 @@ void SelfGravityFFT::Init(Input &input, DataBlock *datain) { std::array nfft_complex = {npr_glob[KDIR], npr_glob[JDIR], npr_glob[IDIR]/2+1}; this->fft = std::make_unique(nfft_real, nfft_complex); + + #ifdef WITH_MPI + + int ntarget = 0; + std::vector mapVars; + mapVars.push_back(ntarget); + + this->mpi.Init(data->mygrid, mapVars, data->nghost, data->np_int, data->lbound, data->rbound, false); + #endif idfx::popRegion(); } @@ -142,7 +151,7 @@ void SelfGravityFFT::SetBoundaries(IdefixArray3D &arr) { idfx::pushRegion("SelfGravityFFT::SetBoundaries"); #ifdef WITH_MPI - this->arr4D = IdefixArray4D (arr.data(), 1, data->np_tot[KDIR], + IdefixArray4D arr4D = IdefixArray4D (arr.data(), 1, data->np_tot[KDIR], data->np_tot[JDIR], data->np_tot[IDIR]); #endif @@ -153,13 +162,13 @@ void SelfGravityFFT::SetBoundaries(IdefixArray3D &arr) { if(data->mygrid->nproc[dir]>1) { switch(dir) { case 0: - this->mpi.ExchangeX1(this->arr4D); + this->mpi.ExchangeX1(arr4D); break; case 1: - this->mpi.ExchangeX2(this->arr4D); + this->mpi.ExchangeX2(arr4D); break; case 2: - this->mpi.ExchangeX3(this->arr4D); + this->mpi.ExchangeX3(arr4D); break; } } @@ -193,13 +202,24 @@ void SelfGravityFFT::SolvePoisson() { auto kx1 = kx[IDIR]; auto kx2 = kx[JDIR]; auto kx3 = kx[KDIR]; - idefix_for("PoissonFFT", 0, npf[KDIR], 0, npf[JDIR], 0, npf[IDIR], - KOKKOS_LAMBDA(int k, int j, int i) { - const real k2 = kx1(i)*kx1(i) + kx2(j)*kx2(j) + kx3(k)*kx3(k); - real inv_k2 = (k2 > 0.0) ? -1.0/k2 : 0.0; - phiF(k,j,i) = rhoF(k,j,i) * inv_k2; - } - ); + #ifdef WITH_MPI + // Work with transposed arrays + idefix_for("PoissonFFT", 0, npf_t[KDIR], 0, npf_t[JDIR], 0, npf_t[IDIR], + KOKKOS_LAMBDA(int j, int k, int i) { + const real k2 = kx1(i)*kx1(i) + kx2(j)*kx2(j) + kx3(k)*kx3(k); + real inv_k2 = (k2 > 0.0) ? -1.0/k2 : 0.0; + phiF(j,k,i) = rhoF(j,k,i) * inv_k2; + } + ); + #else + idefix_for("PoissonFFT", 0, npf[KDIR], 0, npf[JDIR], 0, npf[IDIR], + KOKKOS_LAMBDA(int k, int j, int i) { + const real k2 = kx1(i)*kx1(i) + kx2(j)*kx2(j) + kx3(k)*kx3(k); + real inv_k2 = (k2 > 0.0) ? -1.0/k2 : 0.0; + phiF(k,j,i) = rhoF(k,j,i) * inv_k2; + } + ); + #endif // Backward transform of the potential field auto phiReal = Kokkos::subview(phi, diff --git a/src/gravity/selfGravityFFT.hpp b/src/gravity/selfGravityFFT.hpp index c5623ceea..113c07717 100644 --- a/src/gravity/selfGravityFFT.hpp +++ b/src/gravity/selfGravityFFT.hpp @@ -14,6 +14,9 @@ #include "fft.hpp" #include "selfGravity.hpp" +#ifdef WITH_MPI +#include "mpi.hpp" +#endif class SelfGravityFFT final : public SelfGravity { @@ -46,6 +49,10 @@ class SelfGravityFFT final : public SelfGravity { std::array begin{0,0,0}; // [k,j,i] std::array end{0,0,0}; // [k,j,i] + + #ifdef WITH_MPI + Mpi mpi; // Mpi object when WITH_MPI is set + #endif }; #endif // GRAVITY_SELFGRAVITYFFT_HPP_ diff --git a/src/utils/fft/fft.cpp b/src/utils/fft/fft.cpp index 8cf3dd425..48b4cc313 100644 --- a/src/utils/fft/fft.cpp +++ b/src/utils/fft/fft.cpp @@ -57,7 +57,7 @@ FFT::FFT(std::array npr_glob, std::array npf_glob) { this->tempTransposedReal = IdefixArray3D("FFT transpose temp2", npr[1]/idfx::psize, npr[0]*idfx::psize, npr[2]); this->tempT2Complex = IdefixArray3D("FFT temp2 complex", npf[0],npf[2], npf[1]); this->tempT2Complex2 = IdefixArray3D("FFT temp2 complex2", npf[0],npf[2], npf[1]); - IdefixArray3D tempComplex2 = idfx::makeArray>("FFT temp complex", npf); + IdefixArray3D tempComplex2 = IdefixArray3D("FFT temp complex", npf[0],npf[1],npf[2]); // MPI C2R Plans // Axis 1 transposed is axis2 for the fft library @@ -141,10 +141,10 @@ void FFT::TestMPI() { npf_glob[2] = npr_glob[2]/2+1; Kokkos::View localReal_right("local real array", npr[0], npr[1], npr[2]); - IdefixArray3D localReal = idfx::makeArray>("local real array", npr); + IdefixArray3D localReal = IdefixArray3D("local real array", npr[0], npr[1], npr[2]); Kokkos::View globalReal_right("global real array", npr_glob[0], npr_glob[1], npr_glob[2]); - IdefixArray3D globalReal = idfx::makeArray>("global real array", npr_glob); + IdefixArray3D globalReal = IdefixArray3D("global real array", npr_glob[0], npr_glob[1], npr_glob[2]); // Compute a dummy real array if(idfx::prank == 0) { @@ -154,13 +154,13 @@ void FFT::TestMPI() { // Scatter to make local arrays // And broadcast the global the array MPI_Scatter(globalReal_right.data(), npr[0]*npr[1]*npr[2], - MPI_Astra_real, + realMPI, localReal_right.data(), npr[0]*npr[1]*npr[2], - MPI_Astra_real, + realMPI, 0, MPI_COMM_WORLD); MPI_Bcast(globalReal_right.data(), npr_glob[0]*npr_glob[1]*npr_glob[2], - MPI_Astra_real, 0, MPI_COMM_WORLD); + realMPI, 0, MPI_COMM_WORLD); // Change array Layout idefix_for("Reshape global",0, npr_glob[0], @@ -179,8 +179,8 @@ void FFT::TestMPI() { }); // Create the complex arrays - IdefixArray3D localComplex = idfx::makeArray>("local complex array", npf); - IdefixArray3D globalComplex = idfx::makeArray>("global complex array", npf_glob); + IdefixArray3D localComplex = IdefixArray3D("local complex array", npf[0], npf[1], npf[2]); + IdefixArray3D globalComplex = IdefixArray3D("global complex array", npf_glob[0], npf_glob[1], npf_glob[2]); // Compute the full serial fft KokkosFFT::rfftn(Kokkos::DefaultExecutionSpace(), globalReal, globalComplex); diff --git a/src/utils/fft/fft.hpp b/src/utils/fft/fft.hpp index 9052c80f0..6ce3e898c 100644 --- a/src/utils/fft/fft.hpp +++ b/src/utils/fft/fft.hpp @@ -181,6 +181,7 @@ void FFT::C2R_Host(const InView &in, const OutView &out) { Kokkos::deep_copy(out, outDev); } + template void FFT::TransposeLocal(const ViewIn &in, const ViewOut &out) { static_assert(is_rank3_v, "FFT::TransposeLocal: input must be a rank-3 Kokkos::View"); From b99b936f2ce6f3fd1383cdad2543faba11b9a36a Mon Sep 17 00:00:00 2001 From: Geoffroy Lesur Date: Mon, 10 Aug 2026 10:07:04 +0200 Subject: [PATCH 4/5] fix linting issues in the gravity module --- CPPLINT.cfg | 2 +- src/gravity/gravity.hpp | 1 + src/gravity/selfGravity.cpp | 3 +- src/gravity/selfGravity.hpp | 1 - src/gravity/selfGravityFFT.cpp | 41 ++++++++++++++---------- src/gravity/selfGravityFFT.hpp | 1 + src/gravity/selfGravityIterative.cpp | 5 +-- src/gravity/selfGravityIterative.hpp | 6 ++-- src/utils/fft/fft.cpp | 3 ++ src/utils/fft/fft.hpp | 47 +++++++++++++++++++++------- src/utils/fft/transpose.hpp | 9 ++++-- 11 files changed, 80 insertions(+), 39 deletions(-) diff --git a/CPPLINT.cfg b/CPPLINT.cfg index b02d056c4..29437a5df 100644 --- a/CPPLINT.cfg +++ b/CPPLINT.cfg @@ -1,7 +1,7 @@ # Don't search for additional CPPLINT.cfg in parent directories. set noparent headers=hpp -linelength=100 +linelength=140 # Don't use 'SRC_' as the cpp header guard prefix root=./src/ extensions=hpp,cpp diff --git a/src/gravity/gravity.hpp b/src/gravity/gravity.hpp index 71f29b096..1c8af2985 100644 --- a/src/gravity/gravity.hpp +++ b/src/gravity/gravity.hpp @@ -8,6 +8,7 @@ #ifndef GRAVITY_GRAVITY_HPP_ #define GRAVITY_GRAVITY_HPP_ +#include #include "idefix.hpp" #include "input.hpp" #include "selfGravity.hpp" diff --git a/src/gravity/selfGravity.cpp b/src/gravity/selfGravity.cpp index 292cc2f91..0f1dfc881 100644 --- a/src/gravity/selfGravity.cpp +++ b/src/gravity/selfGravity.cpp @@ -22,7 +22,8 @@ std::unique_ptr SelfGravity::Create(Input &input, DataBlock *data) #ifdef WITH_FFT ptr = std::make_unique(); #else - IDEFIX_ERROR("[SelfGravity]: FFT solver requested but Idefix was not compiled with FFT support."); + IDEFIX_ERROR("[SelfGravity]: FFT solver requested but " + "Idefix was not compiled with FFT support."); #endif } else { ptr = std::make_unique(); diff --git a/src/gravity/selfGravity.hpp b/src/gravity/selfGravity.hpp index ade6df533..ed3ea4dee 100644 --- a/src/gravity/selfGravity.hpp +++ b/src/gravity/selfGravity.hpp @@ -43,7 +43,6 @@ class SelfGravity { protected: DataBlock *data{nullptr}; - }; #endif // GRAVITY_SELFGRAVITY_HPP_ diff --git a/src/gravity/selfGravityFFT.cpp b/src/gravity/selfGravityFFT.cpp index 83c1bd680..950258451 100644 --- a/src/gravity/selfGravityFFT.cpp +++ b/src/gravity/selfGravityFFT.cpp @@ -8,7 +8,9 @@ #ifdef WITH_FFT #include +#include #include +#include #include #include @@ -28,13 +30,18 @@ void SelfGravityFFT::Init(Input &input, DataBlock *datain) { // FFT path currently requires periodic BC in all active dimensions for (int dir = 0; dir < DIMENSIONS; dir++) { if(grid->lbound[dir] != periodic || grid->rbound[dir] != periodic) { - IDEFIX_ERROR("[SelfGravityFFT]: FFT solver requires periodic boundary conditions in all active dimensions."); + IDEFIX_ERROR("[SelfGravityFFT]: FFT solver requires periodic boundary " + "conditions in all active dimensions."); } } // storage with laplacian local array - rho = IdefixArray3D("Density", data->np_int[KDIR], data->np_tot[JDIR], data->np_tot[IDIR]); - phi = IdefixArray3D("Potential", data->np_tot[KDIR], data->np_tot[JDIR], data->np_tot[IDIR]); + rho = IdefixArray3D("Density", data->np_int[KDIR], + data->np_tot[JDIR], + data->np_tot[IDIR]); + phi = IdefixArray3D("Potential", data->np_tot[KDIR], + data->np_tot[JDIR], + data->np_tot[IDIR]); npr_glob = {grid->np_int[IDIR], grid->np_int[JDIR], grid->np_int[KDIR]}; npr = {grid->np_int[IDIR], grid->np_int[JDIR], grid->np_int[KDIR]/idfx::psize}; @@ -46,7 +53,8 @@ void SelfGravityFFT::Init(Input &input, DataBlock *datain) { kx_glob[dir] = KokkosFFT::fftfreq(Device(), npr_glob[dir], d); int p = idfx::prank; if(dir != JDIR) kx[dir] = kx_glob[dir]; - else kx[dir] = Kokkos::subview(kx_glob[dir], std::pair(p * npf_t[KDIR], (p+1) * npf_t[KDIR])); + else kx[dir] = Kokkos::subview(kx_glob[dir], + std::pair(p * npf_t[KDIR], (p+1) * npf_t[KDIR])); } if(idfx::psize > 1) { @@ -61,14 +69,15 @@ void SelfGravityFFT::Init(Input &input, DataBlock *datain) { std::array nfft_complex = {npr_glob[KDIR], npr_glob[JDIR], npr_glob[IDIR]/2+1}; this->fft = std::make_unique(nfft_real, nfft_complex); - + #ifdef WITH_MPI int ntarget = 0; std::vector mapVars; mapVars.push_back(ntarget); - this->mpi.Init(data->mygrid, mapVars, data->nghost, data->np_int, data->lbound, data->rbound, false); + this->mpi.Init(data->mygrid, mapVars, data->nghost, + data->np_int, data->lbound, data->rbound, false); #endif idfx::popRegion(); } @@ -82,13 +91,12 @@ void SelfGravityFFT::CheckCompatibility() { IDEFIX_ERROR("SelfGravityFFT requires no domain decomposition in the X1 and X2 directions."); } if(grid->np_int[KDIR] % idfx::psize != 0 || grid->np_int[JDIR] % idfx::psize != 0) { - IDEFIX_ERROR("SelfGravityFFT requires that the number of grid points in the X3 and X2 directions are" - " divisible by the number of MPI processes."); + IDEFIX_ERROR("SelfGravityFFT requires that the number of grid points in the X3 and " + "X2 directions are divisible by the number of MPI processes."); } if(grid->np_int[IDIR] % 2 != 0) { IDEFIX_ERROR("SelfGravityFFT requires an even number of grid points in the X1 direction."); } - } void SelfGravityFFT::ShowConfig() { @@ -188,10 +196,10 @@ void SelfGravityFFT::SolvePoisson() { // Make a view omitting the ghost cells auto rhoReal = Kokkos::subview(data->hydro->Vc, - RHO, - std::pair(data->beg[KDIR],data->end[KDIR]), - std::pair(data->beg[JDIR],data->end[JDIR]), - std::pair(data->beg[IDIR],data->end[IDIR])); + RHO, + std::pair(data->beg[KDIR],data->end[KDIR]), + std::pair(data->beg[JDIR],data->end[JDIR]), + std::pair(data->beg[IDIR],data->end[IDIR])); auto rhoF = this->rhoF; auto phiF = this->phiF; @@ -221,14 +229,14 @@ void SelfGravityFFT::SolvePoisson() { ); #endif - // Backward transform of the potential field + // make a ghost-cell free view of the potential field to apply the inverse transform auto phiReal = Kokkos::subview(phi, std::pair(data->beg[KDIR],data->end[KDIR]), std::pair(data->beg[JDIR],data->end[JDIR]), std::pair(data->beg[IDIR],data->end[IDIR])); fft->C2R(phiF, phiReal, false); - // Need to apply the boundary conditions to the potential field, since the FFT does not know about the ghost cells + // Need to apply the boundary conditions to the potential field to fill the ghost cells, SetBoundaries(phi); elapsedTime += timer.seconds(); @@ -251,7 +259,8 @@ void SelfGravityFFT::AddSelfGravityPotential(IdefixArray3D &phiP) { void SelfGravityFFT::EnrollUserDefBoundary(Laplacian::UserDefBoundaryFunc myFunc) { (void) myFunc; - IDEFIX_ERROR("SelfGravityFFT only supports periodic boundaries; userdef boundary is unsupported."); + IDEFIX_ERROR("SelfGravityFFT only supports periodic boundaries; " + "userdef boundaries are not supported."); } #endif // WITH_FFT diff --git a/src/gravity/selfGravityFFT.hpp b/src/gravity/selfGravityFFT.hpp index 113c07717..0bc5b63d2 100644 --- a/src/gravity/selfGravityFFT.hpp +++ b/src/gravity/selfGravityFFT.hpp @@ -28,6 +28,7 @@ class SelfGravityFFT final : public SelfGravity { void EnrollUserDefBoundary(Laplacian::UserDefBoundaryFunc myFunc) override; void SetBoundaries(IdefixArray3D &); void EnforcePeriodic(int dir, BoundarySide side, IdefixArray3D &); + private: void SubstractMeanDensity(); void CheckCompatibility(); diff --git a/src/gravity/selfGravityIterative.cpp b/src/gravity/selfGravityIterative.cpp index e540ed2a5..181bd4a68 100644 --- a/src/gravity/selfGravityIterative.cpp +++ b/src/gravity/selfGravityIterative.cpp @@ -168,8 +168,9 @@ void SelfGravityIterative::Init(Input &input, DataBlock *datain) { laplacian->np_tot, laplacian->beg, laplacian->end); } else { real step = laplacian->ComputeCFL(); - iterativeSolver = std::make_unique>(*laplacian.get(), targetError, maxiter, step, - laplacian->np_tot, laplacian->beg, laplacian->end); + iterativeSolver = std::make_unique>(*laplacian.get(), + targetError, maxiter, step, + laplacian->np_tot, laplacian->beg, laplacian->end); } diff --git a/src/gravity/selfGravityIterative.hpp b/src/gravity/selfGravityIterative.hpp index 8b25c54d2..0293a89a9 100644 --- a/src/gravity/selfGravityIterative.hpp +++ b/src/gravity/selfGravityIterative.hpp @@ -5,8 +5,8 @@ // Licensed under CeCILL 2.1 License, see COPYING for more information // *********************************************************************************** -#ifndef GRAVITY_SELFGRAVITY_ITERATIVE_HPP_ -#define GRAVITY_SELFGRAVITY_ITERATIVE_HPP_ +#ifndef GRAVITY_SELFGRAVITYITERATIVE_HPP_ +#define GRAVITY_SELFGRAVITYITERATIVE_HPP_ #include #include @@ -46,4 +46,4 @@ class SelfGravityIterative final : public SelfGravity { bool isPeriodic{true}; }; -#endif // GRAVITY_SELFGRAVITY_ITERATIVE_HPP_ +#endif // GRAVITY_SELFGRAVITYITERATIVE_HPP_ diff --git a/src/utils/fft/fft.cpp b/src/utils/fft/fft.cpp index 48b4cc313..87da59de6 100644 --- a/src/utils/fft/fft.cpp +++ b/src/utils/fft/fft.cpp @@ -5,6 +5,9 @@ // Licensed under CeCILL 2.1 License, see COPYING for more information // *********************************************************************************** +// This file is originally from the ASTRA code +// https://github.com/glesur/astra + #include #include #include diff --git a/src/utils/fft/fft.hpp b/src/utils/fft/fft.hpp index 6ce3e898c..520f2a7a6 100644 --- a/src/utils/fft/fft.hpp +++ b/src/utils/fft/fft.hpp @@ -5,8 +5,11 @@ // Licensed under CeCILL 2.1 License, see COPYING for more information // *********************************************************************************** -#ifndef FFT_HPP_ -#define FFT_HPP_ +// This file is originally from the ASTRA code +// https://github.com/glesur/astra + +#ifndef UTILS_FFT_FFT_HPP_ +#define UTILS_FFT_FFT_HPP_ #include #include @@ -14,15 +17,35 @@ #include "arrays.hpp" #include "transpose.hpp" -using PlanR2CType = KokkosFFT::Plan, IdefixArray3D,3>; -using PlanC2RType = KokkosFFT::Plan, IdefixArray3D,3>; - -using PlanC2CType1D = KokkosFFT::Plan, IdefixArray3D,1>; -using PlanC2RType1D = KokkosFFT::Plan, IdefixArray3D,1>; -using PlanR2CType1D = KokkosFFT::Plan, IdefixArray3D,1>; - -using PlanR2CType2D = KokkosFFT::Plan, IdefixArray3D,2>; -using PlanC2RType2D = KokkosFFT::Plan, IdefixArray3D,2>; +using PlanR2CType = KokkosFFT::Plan, + IdefixArray3D, + 3>; +using PlanC2RType = KokkosFFT::Plan, + IdefixArray3D, + 3>; +using PlanC2CType1D = KokkosFFT::Plan, + IdefixArray3D, + 1>; +using PlanC2RType1D = KokkosFFT::Plan, + IdefixArray3D, + 1>; +using PlanR2CType1D = KokkosFFT::Plan, + IdefixArray3D, + 1>; + +using PlanR2CType2D = KokkosFFT::Plan, + IdefixArray3D, + 2>; +using PlanC2RType2D = KokkosFFT::Plan, + IdefixArray3D, + 2>; class FFT { public: @@ -197,4 +220,4 @@ void FFT::TransposeLocal(const ViewIn &in, const ViewOut &out) { }); } -#endif // FFT_HPP_ \ No newline at end of file +#endif // UTILS_FFT_FFT_HPP_ diff --git a/src/utils/fft/transpose.hpp b/src/utils/fft/transpose.hpp index fb00cc697..a89bb887a 100644 --- a/src/utils/fft/transpose.hpp +++ b/src/utils/fft/transpose.hpp @@ -5,8 +5,11 @@ // Licensed under CeCILL 2.1 License, see COPYING for more information // *********************************************************************************** -#ifndef TRANSPOSE_HPP_ -#define TRANSPOSE_HPP_ +// This file is originally from the ASTRA code +// https://github.com/glesur/astra + +#ifndef UTILS_FFT_TRANSPOSE_HPP_ +#define UTILS_FFT_TRANSPOSE_HPP_ #ifdef WITH_MPI #include @@ -225,4 +228,4 @@ void Transpose::Test() { } -#endif // TRANSPOSE_HPP_ +#endif // UTILS_FFT_TRANSPOSE_HPP_ From 0e33f6e37070be891a2f355bbbe150e3aae0a620 Mon Sep 17 00:00:00 2001 From: Geoffroy Lesur Date: Mon, 10 Aug 2026 12:29:29 +0200 Subject: [PATCH 5/5] rearrange FFT mpi routines so that the complex array is transposed only --- src/gravity/selfGravityFFT.cpp | 6 +-- src/utils/fft/fft.cpp | 83 +++++++++++++++++++--------------- src/utils/fft/fft.hpp | 28 ++++++++---- src/utils/fft/transpose.hpp | 4 +- 4 files changed, 70 insertions(+), 51 deletions(-) diff --git a/src/gravity/selfGravityFFT.cpp b/src/gravity/selfGravityFFT.cpp index 950258451..d1e47efd9 100644 --- a/src/gravity/selfGravityFFT.cpp +++ b/src/gravity/selfGravityFFT.cpp @@ -57,13 +57,13 @@ void SelfGravityFFT::Init(Input &input, DataBlock *datain) { std::pair(p * npf_t[KDIR], (p+1) * npf_t[KDIR])); } - if(idfx::psize > 1) { + #ifdef WITH_MPI rhoF = IdefixArray3D("rhoHatFFT", npf_t[KDIR], npf_t[JDIR], npf_t[IDIR]); phiF = IdefixArray3D("phiHatFFT", npf_t[KDIR], npf_t[JDIR], npf_t[IDIR]); - } else { + #else rhoF = IdefixArray3D("rhoHatFFT", npf[KDIR], npf[JDIR], npf[IDIR]); phiF = IdefixArray3D("phiHatFFT", npf[KDIR], npf[JDIR], npf[IDIR]); - } + #endif std::array nfft_real = {npr_glob[KDIR], npr_glob[JDIR], npr_glob[IDIR]}; std::array nfft_complex = {npr_glob[KDIR], npr_glob[JDIR], npr_glob[IDIR]/2+1}; diff --git a/src/utils/fft/fft.cpp b/src/utils/fft/fft.cpp index 87da59de6..4f825ceb0 100644 --- a/src/utils/fft/fft.cpp +++ b/src/utils/fft/fft.cpp @@ -35,13 +35,10 @@ FFT::FFT(std::array npr_glob, std::array npf_glob) { // We assume a decomposition along the first dimension only for now this->npr = npr_glob; this->npf = npf_glob; - this->npr_t = npr_glob; #ifdef WITH_MPI this->npr[0] = npr_glob[0]/idfx::psize; this->npf[0] = npf_glob[0]/idfx::psize; - this->npr_t[0] = npr_glob[1]/idfx::psize; - this->npr_t[1] = npr_glob[0]; #endif tempReal = IdefixArray3D("FFT temp real", npr[0],npr[1],npr[2]); @@ -55,23 +52,21 @@ FFT::FFT(std::array npr_glob, std::array npf_glob) { #ifdef WITH_MPI // Allocate temporary arrays for domain-splited FFTs and transposes - this->tempTransposedComplex = IdefixArray3D("FFT transpose temp", npf[1]/idfx::psize, npf[0]*idfx::psize, npf[2]); - this->tempTransposedComplex2 = IdefixArray3D("FFT transpose temp2", npf[1]/idfx::psize, npf[0]*idfx::psize, npf[2]); - this->tempTransposedReal = IdefixArray3D("FFT transpose temp2", npr[1]/idfx::psize, npr[0]*idfx::psize, npr[2]); - this->tempT2Complex = IdefixArray3D("FFT temp2 complex", npf[0],npf[2], npf[1]); - this->tempT2Complex2 = IdefixArray3D("FFT temp2 complex2", npf[0],npf[2], npf[1]); - IdefixArray3D tempComplex2 = IdefixArray3D("FFT temp complex", npf[0],npf[1],npf[2]); + this->tempTransposedComplex = IdefixArray3D("FFT transpose temp", npf_glob[1]/idfx::psize, npf_glob[0], npf_glob[2]); + this->tempTransposedComplex2 = IdefixArray3D("FFT transpose temp2", npf_glob[1]/idfx::psize, npf_glob[0], npf_glob[2]); + this->tempT2Complex = IdefixArray3D("FFT temp2 complex", npf_glob[1]/idfx::psize, npf_glob[2],npf_glob[0]); + this->tempT2Complex2 = IdefixArray3D("FFT temp2 complex2", npf_glob[1]/idfx::psize, npf_glob[2], npf_glob[0]); // MPI C2R Plans // Axis 1 transposed is axis2 for the fft library this->c2ciMPIPlan_axis2 = std::make_unique(Kokkos::DefaultExecutionSpace(), tempT2Complex, tempT2Complex2, KokkosFFT::Direction::backward, std::array{-1}); this->c2rMPIPlan_axis1t3 = std::make_unique(Kokkos::DefaultExecutionSpace(), - tempTransposedComplex, tempTransposedReal, KokkosFFT::Direction::backward, std::array{-2,-1}); + tempComplex, tempReal, KokkosFFT::Direction::backward, std::array{-2,-1}); // MPI R2C plans this->r2cMPIPlan_axis1t3 = std::make_unique(Kokkos::DefaultExecutionSpace(), - tempTransposedReal, tempTransposedComplex, KokkosFFT::Direction::forward, std::array{-2,-1}); + tempReal, tempComplex, KokkosFFT::Direction::forward, std::array{-2,-1}); this->c2cfMPIPlan_axis2 = std::make_unique(Kokkos::DefaultExecutionSpace(), tempT2Complex, tempT2Complex2, KokkosFFT::Direction::forward, std::array{-1}); @@ -87,22 +82,32 @@ FFT::FFT(std::array npr_glob, std::array npf_glob) { void FFT::R2C_MPI(const IdefixArray3D in, IdefixArray3D out, bool transpose) { idfx::pushRegion("FFT::R2C_MPI"); - if(transpose) { - this->transposeReal->Apply(in,tempTransposedReal); - idfx::pushRegion("FFT::R2C_MPI axis1t3"); - KokkosFFT::execute(*(r2cMPIPlan_axis1t3.get()), tempTransposedReal, tempTransposedComplex); - idfx::popRegion(); - } else { - idfx::pushRegion("FFT::R2C_MPI axis1t3"); - KokkosFFT::execute(*(r2cMPIPlan_axis1t3.get()), in, tempTransposedComplex); - idfx::popRegion(); - } - this->transposeComplex->Apply(tempTransposedComplex,tempComplex); - TransposeLocal(tempComplex,tempT2Complex); + + idfx::pushRegion("FFT::R2C_MPI axis1t3"); + // input is [n0/p,n1,n2] + KokkosFFT::execute(*(r2cMPIPlan_axis1t3.get()), in, tempComplex); + idfx::popRegion(); + // tempComplex is [n0/p,n1,n2/2+1] + + this->transposeComplex->Apply(tempComplex,tempTransposedComplex); + + // tempTransposedComplex is [n1/p,n0,n2/2+1] + TransposeLocal(tempTransposedComplex,tempT2Complex); + + // tempT2Complex is [n1/p,n2/2+1,n0] idfx::pushRegion("FFT::R2C_MPI axis2"); KokkosFFT::execute(*(c2cfMPIPlan_axis2.get()), tempT2Complex, tempT2Complex2); idfx::popRegion(); - TransposeLocal(tempT2Complex2,out); + + // tempT2Complex is [n1/p,n2/2+1,n0] + if(transpose) { + // Expect output to be [n0/p,n1,n2/2+1] + TransposeLocal(tempT2Complex2,tempTransposedComplex); + this->transposeComplex->Apply(tempTransposedComplex,out); + } else { + // Expect output to be [n1/p,n0,n2/2+1] + TransposeLocal(tempT2Complex2,out); + } idfx::popRegion(); } @@ -110,21 +115,25 @@ void FFT::R2C_MPI(const IdefixArray3D in, IdefixArray3D out, bool void FFT::C2R_MPI(const IdefixArray3D in, IdefixArray3D out, bool transpose) { idfx::pushRegion("FFT::C2R_MPI"); idfx::pushRegion("FFT::C2R_MPI axis2"); - TransposeLocal(in,tempT2Complex); - KokkosFFT::execute(*(c2ciMPIPlan_axis2.get()), tempT2Complex, tempT2Complex2); - TransposeLocal(tempT2Complex2,tempComplex); - idfx::popRegion(); - this->transposeComplex->Apply(tempComplex,tempTransposedComplex); - if(transpose) { - idfx::pushRegion("FFT::C2R_MPI axis1t3"); - KokkosFFT::execute(*(c2rMPIPlan_axis1t3.get()), tempTransposedComplex, tempTransposedReal); - idfx::popRegion(); - this->transposeReal->Apply(tempTransposedReal,out); + if(!transpose) { + // Expect input to be [n1/p,n0,n2/2+1] + TransposeLocal(in,tempT2Complex); } else { - idfx::pushRegion("FFT::C2R_MPI axis1t3"); - KokkosFFT::execute(*(c2rMPIPlan_axis1t3.get()), tempTransposedComplex, out); - idfx::popRegion(); + // Expect input to be [n0/p,n1,n2/2+1] + this->transposeComplex->Apply(in,tempTransposedComplex); + // tempTransposedComplex is [n1/p,n0,n2/2+1] + TransposeLocal(tempTransposedComplex,tempT2Complex); } + // tempT2Complex is [n1/p,n2/2+1,n0] + KokkosFFT::execute(*(c2ciMPIPlan_axis2.get()), tempT2Complex, tempT2Complex2); + TransposeLocal(tempT2Complex2,tempTransposedComplex); + idfx::popRegion(); + // tempTransposedComplex is [n1/p,n0,n2/2+1] + this->transposeComplex->Apply(tempTransposedComplex,tempComplex); + + // tempComplex is [n0/p,n1,n2/2+1] + KokkosFFT::execute(*(c2rMPIPlan_axis1t3.get()), tempComplex, out); + idfx::popRegion(); } diff --git a/src/utils/fft/fft.hpp b/src/utils/fft/fft.hpp index 520f2a7a6..5ebaf3776 100644 --- a/src/utils/fft/fft.hpp +++ b/src/utils/fft/fft.hpp @@ -53,10 +53,10 @@ class FFT { FFT(std::array npr_glob, std::array npf_glob); template - void R2C(const InView &in, const OutView &out, bool transpose=false); + void R2C(const InView &in, const OutView &out, bool transpose=true); template - void C2R(const InView &in, const OutView &out, bool transpose=false); + void C2R(const InView &in, const OutView &out, bool transpose=true); template void R2C_Host(const InView &in, const OutView &out); @@ -67,8 +67,8 @@ class FFT { template void TransposeLocal(const ViewIn &in, const ViewOut &out); - void R2C_MPI(const IdefixArray3D in, IdefixArray3D out, bool transpose=false); - void C2R_MPI(const IdefixArray3D in, IdefixArray3D out, bool transpose=false); + void R2C_MPI(const IdefixArray3D in, IdefixArray3D out, bool transpose=true); + void C2R_MPI(const IdefixArray3D in, IdefixArray3D out, bool transpose=true); void TestMPI(); private: @@ -102,7 +102,7 @@ class FFT { public: std::array npr_glob, npf_glob; - std::array npr, npf, npr_t; + std::array npr, npf; bool havePlan{false}; IdefixArray3D tempReal; @@ -139,8 +139,13 @@ void FFT::R2C(const InView &in, const OutView &out, bool transpose) { #ifdef WITH_MPI Kokkos::deep_copy(tempReal, in); - R2C_MPI(tempReal, tempComplex, transpose); - Kokkos::deep_copy(out, tempComplex); + if(!transpose) { + R2C_MPI(tempReal, tempTransposedComplex, transpose); + Kokkos::deep_copy(out, tempTransposedComplex); + } else { + R2C_MPI(tempReal, tempComplex, transpose); + Kokkos::deep_copy(out, tempComplex); + } #else Kokkos::deep_copy(tempReal, in); KokkosFFT::execute(*(r2cPlan.get()), tempReal, tempComplex); @@ -162,8 +167,13 @@ void FFT::C2R(const InView &in, const OutView &out, bool transpose) { idfx::pushRegion("FFT::C2R"); #ifdef WITH_MPI - Kokkos::deep_copy(tempComplex, in); - C2R_MPI(tempComplex, tempReal, transpose); + if(!transpose) { + Kokkos::deep_copy(tempTransposedComplex, in); + C2R_MPI(tempTransposedComplex, tempReal,transpose); + } else { + Kokkos::deep_copy(tempComplex, in); + C2R_MPI(tempComplex, tempReal, transpose); + } Kokkos::deep_copy(out, tempReal); #else Kokkos::deep_copy(tempComplex, in); diff --git a/src/utils/fft/transpose.hpp b/src/utils/fft/transpose.hpp index a89bb887a..c8ad9e532 100644 --- a/src/utils/fft/transpose.hpp +++ b/src/utils/fft/transpose.hpp @@ -28,7 +28,7 @@ class Transpose { this->tempB = Kokkos::View("FFT transpose tempB", n1,nk); this->tempC = Kokkos::View("FFT transpose tempC", n1,nk); } - void Apply(const IdefixArray3D& input, IdefixArray3D& output); + void Apply(const IdefixArray3D& input, const IdefixArray3D& output); void Test(); private: @@ -128,7 +128,7 @@ Then we need to transpose x1 and x2 to have x1 contiguous in memory. #include "idefix.hpp" template -void Transpose::Apply(const IdefixArray3D& in, IdefixArray3D& out) { +void Transpose::Apply(const IdefixArray3D& in, const IdefixArray3D& out) { idfx::pushRegion("Transpose::Apply"); #ifdef WITH_MPI const int64_t n = idfx::psize; // number of MPI processes