Merge pull request #799 from rodrigoacb/master

Add support for all heFFTe MPI communication backends
ECP-copa · Dec 9, 2024 · 8cf2599 · 8cf2599
2 parents 8d396b5 + 4689269
commit 8cf2599
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 16 deletions.
diff --git a/example/grid_tutorial/10_fft_heffte/heffte_fast_fourier_transform_example.cpp b/example/grid_tutorial/10_fft_heffte/heffte_fast_fourier_transform_example.cpp
@@ -104,9 +104,15 @@ void fastFourierTransformHeffteExample()
     */
     Cabana::Grid::Experimental::FastFourierTransformParams params;
 
-    // Set communication to use all-to-all MPI communication.
-    // Set this option to false for point-to-point communication
-    params.setAllToAll( true );
+    // Set this option to Cabana::Grid::Experimental::FFTCommPattern::alltoallv
+    // or true use all-to-all communication with padding
+    // Set this option to Cabana::Grid::Experimental::FFTCommPattern::p2p or
+    // false for point-to-point communication
+    // Set this option to Cabana::Grid::Experimental::FFTCommPattern::alltoall
+    // for all-to-all communication without padding
+    // Set this option to Cabana::Grid::Experimental::FFTCommPattern::p2p_plined
+    // for pipelined point-to-point communication
+    params.setAlltoAll( Cabana::Grid::Experimental::FFTCommPattern::alltoallv );
 
     // Set data exchange type to use pencil decomposition
     // Set this option to false to use slab decomposition

diff --git a/grid/src/Cabana_Grid_FastFourierTransform.hpp b/grid/src/Cabana_Grid_FastFourierTransform.hpp
@@ -19,8 +19,6 @@
 #include <Cabana_Grid_Array.hpp>
 #include <Cabana_Grid_Types.hpp>
 
-#include <Cabana_Utils.hpp> // FIXME: remove after next release.
-
 #include <Kokkos_Core.hpp>
 #include <Kokkos_Profiling_ScopedRegion.hpp>
 
@@ -98,13 +96,24 @@ struct is_matching_array<
 {
 };
 
+/*!
+  \brief Choices of heFFTe MPI communication patterns
+*/
+enum class FFTCommPattern : unsigned int
+{
+    alltoallv = 0u,
+    p2p = 1u,
+    alltoall = 2u,
+    p2p_plined = 3u
+};
+
 //---------------------------------------------------------------------------//
 /*!
   \brief Parameters controlling details for fast Fourier transforms.
 */
 class FastFourierTransformParams
 {
-    bool alltoall = true;
+    FFTCommPattern FFTcomm = FFTCommPattern::alltoallv;
     bool pencils = true;
     bool reorder = true;
 
@@ -113,7 +122,18 @@ class FastFourierTransformParams
       \brief Set MPI communication strategy.
       \param value Use all to all MPI communication.
     */
-    void setAllToAll( bool value ) { alltoall = value; }
+    void setAlltoAll( FFTCommPattern value ) { FFTcomm = value; }
+    /*!
+      \brief Set MPI communication strategy.
+      \param value Use all to all MPI communication.
+    */
+    void setAlltoAll( bool value )
+    {
+        if ( value )
+            FFTcomm = FFTCommPattern::alltoallv;
+        else
+            FFTcomm = FFTCommPattern::p2p;
+    }
     /*!
       \brief Set data exchange type (pencil or slab).
       \param value Use pencil (true) or slab (false) decomposition.
@@ -129,7 +149,7 @@ class FastFourierTransformParams
       \brief Get MPI communication strategy.
       \return Using AllToAll or not.
     */
-    bool getAllToAll() const { return alltoall; }
+    FFTCommPattern getAlltoAll() const { return FFTcomm; }
     /*!
       \brief Get data exchange type (pencil or slab).
       \return Using pencil (true) or slab (false) decomposition.
@@ -538,12 +558,26 @@ class HeffteFastFourierTransform
 
         heffte::plan_options heffte_params =
             heffte::default_options<heffte_backend_type>();
-        // TODO: use all three heffte options for algorithm
-        bool alltoall = params.getAllToAll();
-        if ( alltoall )
-            heffte_params.algorithm = heffte::reshape_algorithm::alltoallv;
-        else
+        auto FFTcomm = params.getAlltoAll();
+        switch ( FFTcomm )
+        {
+        case Cabana::Grid::Experimental::FFTCommPattern::p2p:
             heffte_params.algorithm = heffte::reshape_algorithm::p2p;
+            break;
+        case Cabana::Grid::Experimental::FFTCommPattern::alltoallv:
+            heffte_params.algorithm = heffte::reshape_algorithm::alltoallv;
+            break;
+        case Cabana::Grid::Experimental::FFTCommPattern::alltoall:
+            heffte_params.algorithm = heffte::reshape_algorithm::alltoall;
+            break;
+        case Cabana::Grid::Experimental::FFTCommPattern::p2p_plined:
+            heffte_params.algorithm = heffte::reshape_algorithm::p2p_plined;
+            break;
+        default:
+            heffte_params.algorithm = heffte::reshape_algorithm::alltoallv;
+            break;
+        }
+
         heffte_params.use_pencils = params.getPencils();
         heffte_params.use_reorder = params.getReorder();
 
@@ -698,8 +732,7 @@ auto createHeffteFastFourierTransform(
     const heffte::plan_options heffte_params =
         heffte::default_options<heffte_backend_type>();
     FastFourierTransformParams params;
-    // TODO: set appropriate default for AllToAll
-    params.setAllToAll( true );
+    params.setAlltoAll( FFTCommPattern::alltoallv );
     params.setPencils( heffte_params.use_pencils );
     params.setReorder( heffte_params.use_reorder );
 

diff --git a/grid/unit_test/tstFastFourierTransform.hpp b/grid/unit_test/tstFastFourierTransform.hpp
@@ -39,7 +39,7 @@ void calculateFFT( bool use_default, bool use_params,
     Experimental::FastFourierTransformParams params;
 
     // Set MPI communication
-    params.setAllToAll( true );
+    params.setAlltoAll( Cabana::Grid::Experimental::FFTCommPattern::alltoallv );
 
     // Set data exchange type (false uses slab decomposition)
     params.setPencils( true );