Poisson: 2D FFT in x-y plus tridiagonal solve in z-direction

ExampleCodes/heFFTe/Poisson2/GNUmakefile

@@ -0,0 +1,35 @@
+AMREX_HOME ?= ../../../../amrex
+HEFFTE_HOME ?= ../../../../heffte/build
+
+DEBUG = FALSE
+DIM = 3
+COMP = gcc
+TINY_PROFILE = FALSE
+USE_MPI = TRUE
+USE_CUDA = FALSE
+USE_HIP = FALSE
+
+BL_NO_FORT = TRUE
+
+include $(AMREX_HOME)/Tools/GNUMake/Make.defs
+include $(AMREX_HOME)/Src/Base/Make.package
+include Make.package
+
+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

ExampleCodes/heFFTe/Poisson2/Make.package

@@ -0,0 +1 @@
+CEXE_sources += main.cpp

ExampleCodes/heFFTe/Poisson2/inputs

@@ -0,0 +1,11 @@
+n_cell_x = 64
+n_cell_y = 64
+n_cell_z = 64
+
+prob_lo_x = 0.
+prob_lo_y = 0.
+prob_lo_z = 0.
+
+prob_hi_x = 1.
+prob_hi_y = 1.
+prob_hi_z = 1.

ExampleCodes/heFFTe/Poisson2/main.cpp

@@ -0,0 +1,314 @@
+#include <heffte.h>
+
+#include <AMReX.H>
+#include <AMReX_MultiFab.H>
+#include <AMReX_ParmParse.H>
+#include <AMReX_GpuComplex.H>
+#include <AMReX_PlotFileUtil.H>
+
+using namespace amrex;
+
+static_assert(AMREX_SPACEDIM == 3);
+
+int main (int argc, char* argv[])
+{
+ amrex::Initialize(argc, argv); {
+
+ BL_PROFILE("main");
+
+ // **********************************
+ // DECLARE SIMULATION PARAMETERS
+ // **********************************
+
+ // number of cells on each side of the domain
+ int n_cell_x;
+ int n_cell_y;
+ int n_cell_z;
+
+ // physical dimensions of the domain
+ Real prob_lo_x = 0.;
+ Real prob_lo_y = 0.;
+ Real prob_lo_z = 0.;
+ Real prob_hi_x = 1.;
+ Real prob_hi_y = 1.;
+ Real prob_hi_z = 1.;
+
+ // **********************************
+ // READ PARAMETER VALUES FROM INPUTS FILE
+ // **********************************
+ {
+ // ParmParse is way of reading inputs from the inputs file
+ // pp.get means we require the inputs file to have it
+ // pp.query means we optionally need the inputs file to have it - but you should supply a default value above
+
+ ParmParse pp;
+
+ pp.get("n_cell_x",n_cell_x);
+ pp.get("n_cell_y",n_cell_y);
+ pp.get("n_cell_z",n_cell_z);
+
+ pp.query("prob_lo_x",prob_lo_x);
+ pp.query("prob_lo_y",prob_lo_y);
+ pp.query("prob_lo_z",prob_lo_z);
+
+ pp.query("prob_hi_x",prob_hi_x);
+ pp.query("prob_hi_y",prob_hi_y);
+ pp.query("prob_hi_z",prob_hi_z);
+ }
+
+ // Determine the domain length in each direction
+ Real L_x = std::abs(prob_hi_x - prob_lo_x);
+ Real L_y = std::abs(prob_hi_y - prob_lo_y);
+
+ // define lower and upper indices of domain
+ IntVect dom_lo(AMREX_D_DECL( 0, 0, 0));
+ IntVect dom_hi(AMREX_D_DECL(n_cell_x-1, n_cell_y-1, n_cell_z-1));
+
+ // Make a single box that is the entire domain
+ Box domain(dom_lo, dom_hi);
+
+ // Initialize the boxarray "ba" from the single box "domain" There are
+ // exactly nprocs boxes. The domain decomposition is done in the x- and
+ // y-directions, but not the z-direction.
+ BoxArray ba = amrex::decompose(domain, ParallelDescriptor::NProcs(),
+ {AMREX_D_DECL(true,true,false)});
+
+ // How Boxes are distrubuted among MPI processes
+ DistributionMapping dm(ba);
+
+ // This defines the physical box size in each direction
+ RealBox real_box({ AMREX_D_DECL(prob_lo_x, prob_lo_y, prob_lo_z)},
+ { AMREX_D_DECL(prob_hi_x, prob_hi_y, prob_hi_z)} );
+
+ // periodic in all direction
+ Array<int,AMREX_SPACEDIM> is_periodic{AMREX_D_DECL(1,1,1)};
+
+ // geometry object for real data
+ Geometry geom(domain, real_box, CoordSys::cartesian, is_periodic);
+
+ // extract dx from the geometry object
+ GpuArray<Real,AMREX_SPACEDIM> dx = geom.CellSizeArray();
+
+ MultiFab rhs(ba,dm,1,0);
+ MultiFab soln(ba,dm,1,0);
+
+ // check to make sure each MPI rank has exactly 1 box
+ AMREX_ALWAYS_ASSERT_WITH_MESSAGE(rhs.local_size() == 1, "Must have one Box per MPI process");
+
+ for (MFIter mfi(rhs); mfi.isValid(); ++mfi) {
+
+ Array4<Real> const& rhs_ptr = rhs.array(mfi);
+
+ const Box& bx = mfi.fabbox();
+
+ amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
+ {
+
+ // **********************************
+ // SET VALUES FOR EACH CELL
+ // **********************************
+
+ Real x = (i+0.5) * dx[0];
+ Real y = (AMREX_SPACEDIM>=2) ? (j+0.5) * dx[1] : 0.;
+ Real z = (AMREX_SPACEDIM==3) ? (k+0.5) * dx[2] : 0.;
+
+ rhs_ptr(i,j,k) = std::exp(-10.*((x-0.5)*(x-0.5)+(y-0.5)*(y-0.5)+(z-0.5)*(z-0.5)));
+
+ });
+ }
+
+ // Shift rhs so that its sum is zero.
+ auto rhosum = rhs.sum(0);
+ rhs.plus(-rhosum/geom.Domain().d_numPts(), 0, 1);
+
+ // since there is 1 MPI rank per box, here each MPI rank obtains its local box and the associated boxid
+ Box local_box;
+ int local_boxid;
+ {
+ for (int i = 0; i < ba.size(); ++i) {
+ Box b = ba[i];
+ // each MPI rank has its own local_box Box and local_boxid ID
+ if (ParallelDescriptor::MyProc() == dm[i]) {
+ local_box = b;
+ local_boxid = i;
+ }
+ }
+ }
+
+ // now each MPI rank works on its own box
+ // for real->complex fft's, the fft is stored in an (nx/2+1) x ny x nz dataset
+
+ // start by coarsening each box by 2 in the x-direction
+ Box c_local_box = amrex::coarsen(local_box, IntVect(AMREX_D_DECL(2,1,1)));
+
+ // if the coarsened box's high-x index is even, we shrink the size in 1 in x
+ // this avoids overlap between coarsened boxes
+ if (c_local_box.bigEnd(0) * 2 == local_box.bigEnd(0)) {
+ c_local_box.setBig(0,c_local_box.bigEnd(0)-1);
+ }
+ // for any boxes that touch the hi-x domain we
+ // increase the size of boxes by 1 in x
+ // this makes the overall fft dataset have size (Nx/2+1 x Ny x Nz)
+ if (local_box.bigEnd(0) == geom.Domain().bigEnd(0)) {
+ c_local_box.growHi(0,1);
+ }
+
+ // each MPI rank gets storage for its piece of the fft
+ BaseFab<GpuComplex<Real> > spectral_field(c_local_box, 1, The_Device_Arena());
+
+ // create real->complex fft objects with the appropriate backend and data about
+ // the domain size and its local box size
+ using fft_r2c_t =
+#ifdef AMREX_USE_CUDA
+ heffte::fft2d_r2c<heffte::backend::cufft>;
+#elif AMREX_USE_HIP
+ heffte::fft2d_r2c<heffte::backend::rocfft>;
+#else
+ heffte::fft2d_r2c<heffte::backend::fftw>;
+#endif
+
+ auto lo = amrex::lbound(local_box);
+ auto hi = amrex::ubound(local_box);
+ auto len = amrex::length(local_box);
+ auto clo = amrex::lbound(c_local_box);
+ auto chi = amrex::ubound(c_local_box);
+ auto clen = amrex::length(c_local_box);
+
+ auto fft = std::make_unique<fft_r2c_t>
+ (heffte::box3d({ lo.x, lo.y, 0},
+ { hi.x, hi.y, 0}),
+ heffte::box3d({clo.x, clo.y, 0},
+ {chi.x, chi.y, 0}),
+ 0, ParallelDescriptor::Communicator());
+
+ Real start_step = static_cast<Real>(ParallelDescriptor::second());
+ using heffte_complex = typename heffte::fft_output<Real>::type;
+ heffte_complex* spectral_data = (heffte_complex*) spectral_field.dataPtr();
+
+ int batch_size = n_cell_z;
+ Gpu::DeviceVector<heffte_complex> workspace(fft->size_workspace()*batch_size);
+
+ { BL_PROFILE("HEFFTE-total");
+ {
+ BL_PROFILE("ForwardTransform");
+ fft->forward(batch_size, rhs[local_boxid].dataPtr(), spectral_data,
+ workspace.data());
+ }
+
+ // Now we take the standard FFT and scale it by 1/k^2
+ Array4< GpuComplex<Real> > spectral = spectral_field.array();
+
+ FArrayBox tridiag_workspace(c_local_box,4);
+ auto const& ald = tridiag_workspace.array(0);
+ auto const& bd = tridiag_workspace.array(1);
+ auto const& cud = tridiag_workspace.array(2);
+ auto const& scratch = tridiag_workspace.array(3);
+
+ Gpu::DeviceVector<Real> delzv(n_cell_z, dx[2]);
+ auto const* delz = delzv.data();
+
+ auto xybox = amrex::makeSlab(c_local_box, 2, 0);
+ ParallelFor(xybox, [=] AMREX_GPU_DEVICE(int i, int j, int)
+ {
+ Real a = 2.*M_PI*i / L_x;
+ Real b = 2.*M_PI*j / L_y;
+
+ // the values in the upper-half of the spectral array in y and z are here interpreted as negative wavenumbers
+ if (j >= n_cell_y/2) b = 2.*M_PI*(n_cell_y-j) / L_y;
+
+ Real k2 = 2*(std::cos(a*dx[0])-1.)/(dx[0]*dx[0]) + 2*(std::cos(b*dx[1])-1.)/(dx[1]*dx[1]);
+
+ // Tridiagonal solve with homogeneous Neumann
+ for( int k=0; k<n_cell_z; k++) {
+ if(k==0) {
+ ald(i,j,k) = 0.;
+ cud(i,j,k) = 2.0 /(delz[k]*(delz[k]+delz[k+1]));
+ bd(i,j,k) = k2 -ald(i,j,k)-cud(i,j,k);
+ } else if (k == n_cell_z-1) {
+ ald(i,j,k) = 2.0 /(delz[k]*(delz[k]+delz[k-1]));
+ cud(i,j,k) = 0.;
+ bd(i,j,k) = k2 -ald(i,j,k)-cud(i,j,k);
+ if (i == 0 && j == 0) {
+ bd(i,j,k) *= 2.0;
+ }
+ } else {
+ ald(i,j,k) = 2.0 /(delz[k]*(delz[k]+delz[k-1]));
+ cud(i,j,k) = 2.0 /(delz[k]*(delz[k]+delz[k+1]));
+ bd(i,j,k) = k2 -ald(i,j,k)-cud(i,j,k);
+ }
+ }
+
+ scratch(i,j,0) = cud(i,j,0)/bd(i,j,0);
+ spectral(i,j,0) = spectral(i,j,0)/bd(i,j,0);
+
+ for (int k = 1; k < n_cell_z; k++) {
+ if (k < n_cell_z-1){
+ scratch(i,j,k) = cud(i,j,k) / (bd(i,j,k) - ald(i,j,k) * scratch(i,j,k-1));
+ }
+ spectral(i,j,k) = (spectral(i,j,k) - ald(i,j,k) * spectral(i,j,k - 1)) / (bd(i,j,k) - ald(i,j,k) * scratch(i,j,k-1));
+ }
+
+ for (int k = n_cell_z - 2; k >= 0; k--) {
+ spectral(i,j,k) -= scratch(i,j,k) * spectral(i,j,k + 1);
+ }
+ });
+ Gpu::streamSynchronize();
+
+ {
+ BL_PROFILE("BackwardTransform");
+ fft->backward(batch_size, spectral_data, soln[local_boxid].dataPtr(),
+ heffte::scale::full);
+ }
+ }
+
+ Real end_step = static_cast<Real>(ParallelDescriptor::second());
+ // amrex::Print() << "TIME IN SOLVE " << end_step - start_step << std::endl;
+
+ // storage for variables to write to plotfile
+ MultiFab plotfile(ba, dm, 2, 0);
+
+ // copy rhs and soln into plotfile
+ MultiFab::Copy(plotfile, rhs , 0, 0, 1, 0);
+ MultiFab::Copy(plotfile, soln, 0, 1, 1, 0);
+
+ // time and step are dummy variables required to WriteSingleLevelPlotfile
+ Real time = 0.;
+ int step = 0;
+
+ // arguments
+ // 1: name of plotfile
+ // 2: MultiFab containing data to plot
+ // 3: variables names
+ // 4: geometry object
+ // 5: "time" of plotfile; not relevant in this example
+ // 6: "time step" of plotfile; not relevant in this example
+ WriteSingleLevelPlotfile("plt", plotfile, {"rhs", "soln"}, geom, time, step);
+
+ {
+ MultiFab phi(soln.boxArray(), soln.DistributionMap(), 1, 1);
+ MultiFab res(soln.boxArray(), soln.DistributionMap(), 1, 0);
+ MultiFab::Copy(phi, soln, 0, 0, 1, 0);
+ phi.FillBoundary(geom.periodicity());
+ auto const& res_ma = res.arrays();
+ auto const& phi_ma = phi.const_arrays();
+ auto const& rhs_ma = rhs.const_arrays();
+ ParallelFor(res, [=] AMREX_GPU_DEVICE (int b, int i, int j, int k)
+ {
+ auto const& phia = phi_ma[b];
+ auto lap = (phia(i-1,j,k)-2.*phia(i,j,k)+phia(i+1,j,k)) / (dx[0]*dx[0])
+ + (phia(i,j-1,k)-2.*phia(i,j,k)+phia(i,j+1,k)) / (dx[1]*dx[1]);
+ if (k == 0) {
+ lap += (-phia(i,j,k)+phia(i,j,k+1)) / (dx[2]*dx[2]);
+ } else if (k == n_cell_z-1) {
+ lap += (phia(i,j,k-1)-phia(i,j,k)) / (dx[2]*dx[2]);
+ } else {
+ lap += (phia(i,j,k-1)-2.*phia(i,j,k)+phia(i,j,k+1)) / (dx[2]*dx[2]);
+ }
+ res_ma[b](i,j,k) = rhs_ma[b](i,j,k) - lap;
+ });
+ amrex::Print() << " rhs.min & max: " << rhs.min(0) << " " << rhs.max(0) << "\n"
+ << " res.min & max: " << res.min(0) << " " << res.max(0) << "\n";
+ }
+
+ } amrex::Finalize();
+}