Merge pull request #104 from ajnonaka/heFFTe

MPI-Distributed heFFTe tutorials

Author: ajnonaka

Docs/source/ML_Tutorial.rst

@@ -10,7 +10,7 @@ Located in the directory ``amrex-tutorials/ExampleCodes/ML/PYTORCH``, this examp
 Here we use a 1-input, 2-output model to illustrate the interface between the PyTorch model and a MultiFab.
 
 **Beta Decay Reaction**
---------------
+-----------------------
 
 In this example, the machine learning model is a regression model pre-trained to solve a two-component ODE system describing beta decay.
 

Docs/source/heFFTe_Tutorial.rst

@@ -0,0 +1,48 @@
+.. role:: cpp(code)
+   :language: c++
+
+.. role:: fortran(code)
+   :language: fortran
+
+.. _tutorials_heffte:
+
+heFFTe
+==========================
+
+These tutorials demonstrate how to heFFTe to solve for and manipulate Fourier transform data using multiple MPI ranks/GPUs.
+
+There are two heFFTe tutorials, ``Basic`` and ``Poisson``:
+
+- ``Basic`` tutorial: The tutorial found in ``amrex-tutorials/ExampleCodes/heFFTe/Basic`` is
+  useful if the objective is to simply take a forward FFT of data, and the DFT's ordering in k-space
+  matters to the user.  This tutorial initializes a 3D or 2D :cpp:`MultiFab`, takes a forward FFT,
+  and then redistributes the data in k-space where the center cell in the domain corresponds to the k=0 mode.
+  The results are written to a plot file.
+
+- ``Poisson`` tutorial: This tutorial: ``amrex-tutorials/ExampleCodes/heFFTe/Poisson``
+  solves a Poisson equation with periodic boundary conditions.  In it, both a forward FFT and reverse FFT
+  are called to solve the equation, however, no reordering of the DFT data in k-space is performed.
+
+We note that both backends, fftw and cufft, assume a row-major ordering of data; since a :cpp:`MultiFab` is column major,
+the output to the spectral array is spatially-transposed.
+
+.. toctree::
+   :maxdepth: 1
+
+.. _section:heffte_tutorial:heffte_basic:
+
+Basic
+--------------------------
+
+This tutorial initializes a 3D or 2D :cpp:`MultiFab`, takes a forward FFT,
+and then redistributes the data in k-space where the center cell in the domain corresponds to the k=0 mode.
+The results are written to a plot file.
+
+.. _section:heffte_tutorial:heffte_pois:
+
+Poisson
+--------------------------
+
+In this test case we set up a right hand side (rhs), call the forward transform,
+modify the coefficients, then call the backward solver and output the solution
+to the discrete Poisson equation.

Docs/source/index.rst

@@ -45,8 +45,10 @@ sorted by the following categories:
   Conduit Mesh Blueprint for use with the ALPINE Ascent in situ visualization
   and analysis tool.
 - :ref:`EB<tutorials_eb>`  -- Examples of embedded boundaries.
+- :ref:`FFTW<tutorials_fftw>`  -- FFTW and cuFFT single-rank tutorials.
 - :ref:`ForkJoin<tutorials_forkjoin>` -- Parallel execution and subgrouping of MPI ranks.
 - :ref:`GPU<tutorials_gpu>`  -- Offload work to the GPUs using AMReX tools.
+- :ref:`heFFTe<tutorials_heffte>`  -- heFFTe distributed tutorials.
 - :ref:`Linear Solvers<tutorials_linearsolvers>`  -- Examples of several linear solvers.
 - :ref:`ML/PYTORCH<tutorials_ml>`  -- Use of pytorch models to replace point-wise computational kernels.
 - :ref:`MUI<tutorials_mui>`  -- Incorporates the MxUI/MUI (Multiscale Universal interface) frame into AMReX.
@@ -67,8 +69,10 @@ sorted by the following categories:
    Basic_Tutorial
    Blueprint_Tutorial
    EB_Tutorial
+   FFTW_Tutorial
    ForkJoin_Tutorial
    GPU_Tutorial
+   heFFTe_Tutorial
    LinearSolvers_Tutorial
    ML_Tutorial
    MUI_Tutorial
@@ -88,10 +92,14 @@ sorted by the following categories:
 
 .. _`EB`:  EB_Tutorial.html
 
+.. _`FFTW`:  FFTW_Tutorial.html
+
 .. _`ForkJoin`:  ForkJoin_Tutorial.html
 
 .. _`GPU`:  GPU_Tutorial.html
 
+.. _`heFFTe`:  heFFTe_Tutorial.html
+
 .. _`Linear Solvers`:  LinearSolvers_Tutorial.html
 
 .. _`MUI`: MUI_Tutorial.html

ExampleCodes/FFTW/Basic/main.cpp

@@ -108,7 +108,7 @@ int main (int argc, char* argv[])
     // INITIALIZE PHI
     // **********************************
 
-    double omega = M_PI/2.0;
+    Real omega = M_PI/2.0;
 
     // loop over boxes
     for (MFIter mfi(phi); mfi.isValid(); ++mfi)
@@ -128,10 +128,10 @@ int main (int argc, char* argv[])
             Real z = (AMREX_SPACEDIM==3) ? (k+0.5) * dx[2] : 0.;
             phi_ptr(i,j,k) = std::sin(4*M_PI*x/prob_hi_x + omega);
             if (AMREX_SPACEDIM >= 2) {
-                phi_ptr(i,j,k) *= std::sin(124*M_PI*y/prob_hi_y + omega);
+                phi_ptr(i,j,k) *= std::sin(6*M_PI*y/prob_hi_y + omega);
             }
             if (AMREX_SPACEDIM == 3) {
-                phi_ptr(i,j,k) *= std::sin(2*M_PI*z/prob_hi_z + omega);
+                phi_ptr(i,j,k) *= std::sin(8*M_PI*z/prob_hi_z + omega);
             }
         });
     }

ExampleCodes/FFTW/Basic/myfunc.cpp

@@ -17,21 +17,21 @@ void ShiftFFT(MultiFab& dft_onegrid, const Geometry& geom, const int& zero_avg)
 
   MultiFab::Copy(dft_onegrid_temp,dft_onegrid,0,0,1,0);
 
-  // Shift DFT by N/2+1 (pi)
-  for (MFIter mfi(dft_onegrid); mfi.isValid(); ++mfi) {
+  // zero k=0 mode
+  if (zero_avg == 1) {
+      for (MFIter mfi(dft_onegrid); mfi.isValid(); ++mfi) {
 
-    const Box& bx = mfi.tilebox();
+          const Box& bx = mfi.tilebox();
 
-    const Array4<Real>& dft_temp = dft_onegrid_temp.array(mfi);
+          const Array4<Real>& dft_temp = dft_onegrid_temp.array(mfi);
 
-    if (zero_avg == 1) {
-      amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
-      {
-        if (i == 0 && j == 0 && k == 0) {
-        dft_temp(i,j,k) = 0.;
-        }
-      });
-    }
+          amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
+          {
+              if (i == 0 && j == 0 && k == 0) {
+                  dft_temp(i,j,k) = 0.;
+              }
+          });
+      }
   }
 
   // Shift DFT by N/2+1 (pi)

ExampleCodes/heFFTe/Basic/GNUmakefile

@@ -0,0 +1,35 @@
+AMREX_HOME ?= ../../../../amrex
+HEFFTE_HOME ?= ../../../../heffte/build
+
+DEBUG = FALSE
+DIM = 3
+COMP = gcc
+TINY_PROFILE = TRUE
+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/Basic/Make.package

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

ExampleCodes/heFFTe/Basic/README

@@ -0,0 +1,45 @@
+heFFTe instaallation guide:
+https://icl-utk-edu.github.io/heffte/md_doxygen_installation.html
+
+Installation
+
+>> git clone [email protected]:icl-utk-edu/heffte.git
+OR
+>> git clone https://github.com/icl-utk-edu/heffte.git
+
+>> cd heffte
+>> mkdir build
+>> cd build
+
+HOST BUILD
+
+>> cmake -DCMAKE_BUILD_TYPE=Release \
+   	 -DCMAKE_CXX_STANDARD=17 \
+   	 -DBUILD_SHARED_LIBS=OFF \
+  	 -DCMAKE_INSTALL_PREFIX=. \  # can choose any installation /path/to/heffte
+  	 -DHeffte_ENABLE_FFTW=ON \
+  	 -DHeffte_ENABLE_CUDA=OFF \
+  	 ..
+
+>> make -j4
+>> make install
+
+>> cd /path/to/amrex-tutorials/ExampleCodes/heFFTe/Basic
+>> make -j4 HEFFTE_HOME=/path/to/heffte
+
+NVIDIA/CUDA BUILD
+
+>> cmake -DCMAKE_BUILD_TYPE=Release \
+   	 -DCMAKE_CXX_STANDARD=17 \
+  	 -DBUILD_SHARED_LIBS=OFF \
+  	 -DCMAKE_INSTALL_PREFIX=. \  # can choose any installation /path/to/heffte
+  	 -DHeffte_ENABLE_FFTW=ON \
+  	 -DHeffte_ENABLE_CUDA=ON \
+  	 -DHeffte_DISABLE_GPU_AWARE_MPI=ON \  # OFF is not working on perlmutter
+  	 ..
+
+>> make -j4
+>> make install
+
+>> cd /path/to/amrex-tutorials/ExampleCodes/heFFTe/Basic/
+>> make -j4 USE_CUDA=TRUE HEFFTE_HOME=/path/to/heffte

ExampleCodes/heFFTe/Basic/inputs

@@ -0,0 +1,15 @@
+n_cell_x = 64
+n_cell_y = 64
+n_cell_z = 64
+
+max_grid_size_x = 32
+max_grid_size_y = 32
+max_grid_size_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.