Skip to content

Commit

Permalink
Merge branch 'heffte' of github.com:Haavaan/Warpx into heffte
Browse files Browse the repository at this point in the history
  • Loading branch information
@Haavaan
Haavaan committed Aug 19, 2024
2 parents 985858e + 2d077f6 commit 7c16d44
Show file tree
Hide file tree
Showing 3 changed files with 60 additions and 45 deletions.
32 changes: 5 additions & 27 deletions GNUmakefile
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,8 @@ AMREX_HOME ?= ../amrex
PICSAR_HOME ?= ../picsar
OPENBC_HOME ?= ../openbc_poisson

#DEBUG = FALSE
DEBUG = TRUE
DEBUG = FALSE
#DEBUG = TRUE

WARN_ALL = TRUE
#WARN_ERROR=TRUE
Expand Down Expand Up @@ -33,37 +33,15 @@ USE_PYTHON_MAIN = FALSE

USE_SENSEI_INSITU = FALSE
USE_ASCENT_INSITU = FALSE
USE_OPENPMD = TRUE
USE_OPENPMD = FALSE

WarpxBinDir = Bin

USE_FFT =TRUE
USE_FFT = FALSE
USE_HEFFTE = FALSE
USE_RZ = FALSE

USE_EB = FALSE

WARPX_HOME := .
include $(WARPX_HOME)/Source/Make.WarpX





VPATH_LOCATIONS += $(HEFFTE_HOME)/include
INCLUDE_LOCATIONS += $(HEFFTE_HOME)/include
LIBRARY_LOCATIONS += $(HEFFTE_HOME)/lib

libraries += -lheffte

ifeq ($(USE_CUDA),TRUE)
libraries += -lcufft
else ifeq ($(USE_HIP),TRUE)
# Use rocFFT. ROC_PATH is defined in amrex
INCLUDE_LOCATIONS += $(ROC_PATH)/rocfft/include
LIBRARY_LOCATIONS += $(ROC_PATH)/rocfft/lib
LIBRARIES += -L$(ROC_PATH)/rocfft/lib -lrocfft
else
libraries += -lfftw3_mpi -lfftw3f -lfftw3
endif

include $(AMREX_HOME)/Tools/GNUMake/Make.rules
22 changes: 11 additions & 11 deletions Regression/Checksum/benchmarks_json/openbc_poisson_solver_heffte.json
View file
Original file line number Diff line number Diff line change
@@ -1,20 +1,20 @@
{
"lev=0": {
"Bx": 100915975.15349947,
"By": 157610677.3147944,
"Bz": 1.2351730578146108e-11,
"Ex": 4.725066923360429e+16,
"Ey": 3.025396149301539e+16,
"Ez": 3276584.4371608277,
"Bx": 100915975.15352048,
"By": 157610677.3147938,
"Bz": 9.053087335578475e-12,
"Ex": 4.72506692336041e+16,
"Ey": 3.025396149302092e+16,
"Ez": 3276584.4383433764,
"rho": 10994013582437.197
},
"electron": {
"particle_momentum_x": 5.701279599510115e-19,
"particle_momentum_y": 3.6504531724372975e-19,
"particle_momentum_x": 5.701279599509562e-19,
"particle_momentum_y": 3.65045317231828e-19,
"particle_momentum_z": 1.145432768297242e-10,
"particle_position_x": 17.314086912497846,
"particle_position_x": 17.31408691249785,
"particle_position_y": 0.2583691267187801,
"particle_position_z": 10066.329600000006,
"particle_weight": 19969036501.910973
"particle_position_z": 10066.329600000008,
"particle_weight": 19969036501.910976
}
}
51 changes: 44 additions & 7 deletions Source/ablastr/fields/IntegratedGreenFunctionSolver.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@
#include <AMReX_REAL.H>
#include <AMReX_PlotFileUtil.H>

#if defined(ABLASTR_FFT) && defined(ABLASTR_HEFFTE)
#if defined(ABLASTR_USE_FFT) && defined(ABLASTR_USE_HEFFTE)
#include <heffte.h>
#endif

Expand All @@ -46,7 +46,11 @@ computePhiIGF ( amrex::MultiFab const & rho,
{
using namespace amrex::literals;

#if defined(ABLASTR_FFT) && defined(ABLASTR_HEFFTE)
BL_PROFILE_VAR_NS("IGF FFTs", timer_ffts);
BL_PROFILE_VAR_NS("IGF FFT plans", timer_plans);
BL_PROFILE_VAR_NS("IGF parallel copies", timer_pcopies);

#if defined(ABLASTR_USE_FFT) && defined(ABLASTR_USE_HEFFTE)
{
BL_PROFILE("Integrated Green Function Solver");

Expand Down Expand Up @@ -118,10 +122,12 @@ computePhiIGF ( amrex::MultiFab const & rho,
amrex::MultiFab tmp_G = amrex::MultiFab(realspace_ba, realspace_dm, 1, 0);
tmp_G.setVal(0);

BL_PROFILE_VAR_START(timer_pcopies);
// Copy from rho including its ghost cells to tmp_rho
// w.z.: please check. I think we need to use rho.nGrowVect() otherwise
// the data outside the ba.minimalBox() will not be copied over.
tmp_rho.ParallelCopy( rho, 0, 0, 1, rho.nGrowVect(), amrex::IntVect::TheZeroVector() );
BL_PROFILE_VAR_STOP(timer_pcopies);


// Compute the integrated Green function
Expand Down Expand Up @@ -187,9 +193,10 @@ computePhiIGF ( amrex::MultiFab const & rho,
tmp_rho_fft.shift(realspace_box.smallEnd());
tmp_G_fft.shift(realspace_box.smallEnd());

#ifdef AMREX_USE_CUDA
BL_PROFILE_VAR_START(timer_plans);
#if defined(AMREX_USE_CUDA)
heffte::fft3d_r2c<heffte::backend::cufft> fft
#elif AMREX_USE_HIP
#elif defined(AMREX_USE_HIP)
heffte::fft3d_r2c<heffte::backend::rocfft> fft
#else
heffte::fft3d_r2c<heffte::backend::fftw> fft
Expand All @@ -199,14 +206,16 @@ computePhiIGF ( amrex::MultiFab const & rho,
{{c_local_box.smallEnd(0),c_local_box.smallEnd(1), c_local_box.smallEnd(2)},
{c_local_box.bigEnd(0) ,c_local_box.bigEnd(1) ,c_local_box.bigEnd(2)}},
0, amrex::ParallelDescriptor::Communicator());
BL_PROFILE_VAR_STOP(timer_plans);

using heffte_complex = typename heffte::fft_output<amrex::Real>::type;
heffte_complex* rho_fft_data = (heffte_complex*) tmp_rho_fft.dataPtr();
heffte_complex* G_fft_data = (heffte_complex*) tmp_G_fft.dataPtr();


BL_PROFILE_VAR_START(timer_ffts);
fft.forward(tmp_rho[local_boxid].dataPtr(), rho_fft_data);
fft.forward(tmp_G[local_boxid].dataPtr(), G_fft_data);
BL_PROFILE_VAR_STOP(timer_ffts);

// Multiply tmp_G_fft and tmp_rho_fft in spectral space
// Store the result in-place in Gtmp_G_fft, to save memory
Expand All @@ -217,17 +226,21 @@ computePhiIGF ( amrex::MultiFab const & rho,

// PRINT / SAVE G TIMES RHO

BL_PROFILE_VAR_START(timer_ffts);
fft.backward(G_fft_data, tmp_G[local_boxid].dataPtr());
BL_PROFILE_VAR_STOP(timer_ffts);

// Normalize, since (FFT + inverse FFT) results in a factor N
const amrex::Real normalization = 1._rt / realspace_box.numPts();
tmp_G.mult( normalization );

BL_PROFILE_VAR_START(timer_pcopies);
// Copy from tmp_G to phi
phi.ParallelCopy( tmp_G, 0, 0, 1, amrex::IntVect::TheZeroVector(), phi.nGrowVect());
BL_PROFILE_VAR_STOP(timer_pcopies);

}
#elif defined(ABLASTR_FFT)
#elif defined(ABLASTR_USE_FFT) && !defined(ABLASTR_USE_HEFFTE)
{
BL_PROFILE("Integrated Green Function Solver");

Expand Down Expand Up @@ -265,8 +278,10 @@ computePhiIGF ( amrex::MultiFab const & rho,
SpectralField tmp_rho_fft = SpectralField( spectralspace_ba, dm_global_fft, 1, 0 );
SpectralField tmp_G_fft = SpectralField( spectralspace_ba, dm_global_fft, 1, 0 );

BL_PROFILE_VAR_START(timer_pcopies);
// Copy from rho to tmp_rho
tmp_rho.ParallelCopy( rho, 0, 0, 1, amrex::IntVect::TheZeroVector(), amrex::IntVect::TheZeroVector() );
BL_PROFILE_VAR_STOP(timer_pcopies);

// Compute the integrated Green function
{
Expand Down Expand Up @@ -323,9 +338,12 @@ computePhiIGF ( amrex::MultiFab const & rho,


}
BL_PROFILE_VAR_START(timer_plans);
// Perform forward FFTs
auto forward_plan_rho = ablastr::math::anyfft::FFTplans(spectralspace_ba, dm_global_fft);
auto forward_plan_G = ablastr::math::anyfft::FFTplans(spectralspace_ba, dm_global_fft);
BL_PROFILE_VAR_STOP(timer_plans);

// Loop over boxes perform FFTs
for ( amrex::MFIter mfi(realspace_ba, dm_global_fft); mfi.isValid(); ++mfi ){

Expand All @@ -335,28 +353,40 @@ computePhiIGF ( amrex::MultiFab const & rho,
const amrex::IntVect fft_size = realspace_ba[mfi].length();

// FFT of rho
BL_PROFILE_VAR_START(timer_plans);
forward_plan_rho[mfi] = ablastr::math::anyfft::CreatePlan(
fft_size, tmp_rho[mfi].dataPtr(),
reinterpret_cast<ablastr::math::anyfft::Complex*>(tmp_rho_fft[mfi].dataPtr()),
ablastr::math::anyfft::direction::R2C, AMREX_SPACEDIM);
BL_PROFILE_VAR_STOP(timer_plans);

BL_PROFILE_VAR_START(timer_ffts);
ablastr::math::anyfft::Execute(forward_plan_rho[mfi]);
BL_PROFILE_VAR_STOP(timer_ffts);

// FFT of G
BL_PROFILE_VAR_START(timer_plans);
forward_plan_G[mfi] = ablastr::math::anyfft::CreatePlan(
fft_size, tmp_G[mfi].dataPtr(),
reinterpret_cast<ablastr::math::anyfft::Complex*>(tmp_G_fft[mfi].dataPtr()),
ablastr::math::anyfft::direction::R2C, AMREX_SPACEDIM);
BL_PROFILE_VAR_STOP(timer_plans);

BL_PROFILE_VAR_START(timer_ffts);
ablastr::math::anyfft::Execute(forward_plan_G[mfi]);
BL_PROFILE_VAR_STOP(timer_ffts);

}

// Multiply tmp_G_fft and tmp_rho_fft in spectral space
// Store the result in-place in Gtmp_G_fft, to save memory
amrex::Multiply( tmp_G_fft, tmp_rho_fft, 0, 0, 1, 0);


BL_PROFILE_VAR_START(timer_plans);
// Perform inverse FFT
auto backward_plan = ablastr::math::anyfft::FFTplans(spectralspace_ba, dm_global_fft);
BL_PROFILE_VAR_STOP(timer_plans);

// Loop over boxes perform FFTs
for ( amrex::MFIter mfi(spectralspace_ba, dm_global_fft); mfi.isValid(); ++mfi ){

Expand All @@ -366,18 +396,25 @@ computePhiIGF ( amrex::MultiFab const & rho,
const amrex::IntVect fft_size = realspace_ba[mfi].length();

// Inverse FFT: is done in-place, in the array of G
BL_PROFILE_VAR_START(timer_plans);
backward_plan[mfi] = ablastr::math::anyfft::CreatePlan(
fft_size, tmp_G[mfi].dataPtr(),
reinterpret_cast<ablastr::math::anyfft::Complex*>( tmp_G_fft[mfi].dataPtr()),
ablastr::math::anyfft::direction::C2R, AMREX_SPACEDIM);
BL_PROFILE_VAR_STOP(timer_plans);

BL_PROFILE_VAR_START(timer_ffts);
ablastr::math::anyfft::Execute(backward_plan[mfi]);
BL_PROFILE_VAR_STOP(timer_ffts);
}
// Normalize, since (FFT + inverse FFT) results in a factor N
const amrex::Real normalization = 1._rt / realspace_box.numPts();
tmp_G.mult( normalization );

BL_PROFILE_VAR_START(timer_pcopies);
// Copy from tmp_G to phi
phi.ParallelCopy( tmp_G, 0, 0, 1, amrex::IntVect::TheZeroVector(), phi.nGrowVect() );
BL_PROFILE_VAR_STOP(timer_pcopies);

// Loop to destroy FFT plans
for ( amrex::MFIter mfi(spectralspace_ba, dm_global_fft); mfi.isValid(); ++mfi ){
Expand Down

0 comments on commit 7c16d44

Please sign in to comment.