Merge pull request ECP-copa#113 from cz4rs/cabana_lb

Incorporate Cabana::Grid load balancer

.github/workflows/CI.yml

@@ -94,6 +94,15 @@ jobs:
             -DCMAKE_BUILD_TYPE=${{ matrix.cmake_build_type }}
           cmake --build build --parallel 2
           cmake --install build
+      - name: Checkout ALL
+        run: |
+          git clone --depth 1 --branch master https://gitlab.jsc.fz-juelich.de/SLMS/loadbalancing ALL
+      - name: Build ALL
+        working-directory: ALL
+        run: |
+          cmake -B build -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=$HOME/ALL
+          cmake --build build --parallel 2
+          cmake --install build
       - name: Checkout Cabana
         uses: actions/checkout@v3
         with:
@@ -106,9 +115,8 @@ jobs:
           cmake -B build \
             -DCMAKE_INSTALL_PREFIX=$HOME/Cabana \
             -DMPIEXEC_MAX_NUMPROCS=2 -DMPIEXEC_PREFLAGS="--oversubscribe" \
-            -DCMAKE_PREFIX_PATH="$HOME/kokkos;$HOME/arborx" \
+            -DCMAKE_PREFIX_PATH="$HOME/kokkos;$HOME/arborx;$HOME/ALL" \
             -DCMAKE_CXX_COMPILER=${{ matrix.cxx }} \
-            -DCabana_DISABLE_CAJITA_DEPRECATION_WARNINGS=ON \
             -DCMAKE_BUILD_TYPE=${{ matrix.cmake_build_type }}
           cmake --build build --parallel 2
           cmake --install build
@@ -140,6 +148,7 @@ jobs:
             -DCabanaMD_LAYOUT=${{ matrix.layout }} \
             -DCabanaMD_VECTORLENGTH=${{ matrix.vector }} \
             -DCabanaMD_ENABLE_NNP=${{ matrix.nnp }} \
+            -DCabanaMD_ENABLE_LB=ON \
             -DN2P2_DIR=$HOME/n2p2 \
             -DCabanaMD_LAYOUT_NNP=${{ matrix.layout_nnp }} \
             -DCabanaMD_VECTORLENGTH_NNP=${{ matrix.vector_nnp }}

CMakeLists.txt

@@ -1,4 +1,4 @@
-cmake_minimum_required(VERSION 3.11)
+cmake_minimum_required(VERSION 3.14)
 project(CabanaMD LANGUAGES CXX VERSION 0.1.0)
 
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
@@ -18,6 +18,8 @@ set(CMAKE_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/cmake)
 ##---------------------------------------------------------------------------##
 find_package(Cabana REQUIRED)
 
+option(CabanaMD_ENABLE_LB "Utilize Cabana load balancer" OFF)
+
 ##---------------------------------------------------------------------------##
 # Set up optional libraries
 ##---------------------------------------------------------------------------##

input/in.lj

@@ -17,7 +17,10 @@ pair_coeff	1 1 1.0 1.0 2.5
 
 neighbor	0.3 bin
 neigh_modify	every 20 one 50
+comm_modify	cutoff * 20
 fix		1 all nve
 thermo 10
 
+dump dmpvtk all vtk 10 dump%_*.vtu
+
 run		100

src/CabanaMD_config.hpp.cmakein

@@ -5,6 +5,7 @@
 #define CabanaMD_GIT_COMMIT_HASH "@CabanaMD_GIT_COMMIT_HASH@"
 
 #cmakedefine CabanaMD_ENABLE_NNP
+#cmakedefine CabanaMD_ENABLE_LB
 
 #cmakedefine CabanaMD_LAYOUT @CabanaMD_LAYOUT@
 #cmakedefine CabanaMD_VECTORLENGTH "@CabanaMD_VECTORLENGTH@"

src/cabanamd.h

@@ -57,6 +57,11 @@
 #include <integrator_nve.h>
 #include <types.h>
 
+#ifdef CabanaMD_ENABLE_LB
+#include <Cabana_Grid_LoadBalancer.hpp>
+#include <Cabana_Grid_Types.hpp>
+#endif
+
 class CabanaMD
 {
   public:
@@ -81,6 +86,10 @@ class CbnMD : public CabanaMD
     Comm<t_System> *comm;
     Binning<t_System> *binning;
     InputFile<t_System> *input;
+#ifdef CabanaMD_ENABLE_LB
+    Cabana::Grid::Experimental::LoadBalancer<Cabana::Grid::UniformMesh<double>>
+        *lb;
+#endif
 
     void init( InputCL cl ) override;
     void run() override;

src/cabanamd_impl.h

@@ -51,15 +51,19 @@
 #include <Cabana_Core.hpp>
 #include <Kokkos_Core.hpp>
 
+#include <mpi.h>
+
 #include <output.h>
 #include <property_kine.h>
 #include <property_pote.h>
 #include <property_temperature.h>
 #include <read_data.h>
+#include <vtk_writer.h>
 
 #include <fstream>
 #include <iomanip>
 #include <iostream>
+#include <string>
 
 #define MAXPATHLEN 1024
 
@@ -224,6 +228,12 @@ void CbnMD<t_System, t_Neighbor>::init( InputCL commandline )
         comm->update_force();
     }
 
+#ifdef CabanaMD_ENABLE_LB
+    lb = new Cabana::Grid::Experimental::LoadBalancer<
+        Cabana::Grid::UniformMesh<double>>( MPI_COMM_WORLD,
+                                            system->global_grid );
+#endif
+
     // Initial output
     int step = 0;
     if ( input->thermo_rate > 0 )
@@ -236,10 +246,19 @@ void CbnMD<t_System, t_Neighbor>::init( InputCL commandline )
         auto KE = kine.compute( system ) / system->N;
         if ( !_print_lammps )
         {
+#ifdef CabanaMD_ENABLE_LB
+            log( out, "\n", std::fixed, std::setprecision( 6 ),
+                 "#Timestep Temperature PotE ETot Time Atomsteps/s "
+                 "LBImbalance\n",
+                 step, " ", T, " ", PE, " ", PE + KE, " ",
+                 std::setprecision( 2 ), 0.0, " ", std::scientific, 0.0, " ",
+                 std::setprecision( 2 ), 0.0 );
+#else
             log( out, "\n", std::fixed, std::setprecision( 6 ),
                  "#Timestep Temperature PotE ETot Time Atomsteps/s\n", step,
                  " ", T, " ", PE, " ", PE + KE, " ", std::setprecision( 2 ),
                  0.0, " ", std::scientific, 0.0 );
+#endif
         }
         else
         {
@@ -264,6 +283,7 @@ template <class t_System, class t_Neighbor>
 void CbnMD<t_System, t_Neighbor>::run()
 {
     std::ofstream out( input->output_file, std::ofstream::app );
+    std::ofstream err( input->error_file, std::ofstream::app );
 
     auto neigh_cutoff = input->force_cutoff + input->neighbor_skin;
     bool half_neigh = input->force_iteration_type == FORCE_ITER_NEIGH_HALF;
@@ -272,15 +292,19 @@ void CbnMD<t_System, t_Neighbor>::run()
     PotE<t_System, t_Neighbor> pote( comm );
     KinE<t_System> kine( comm );
 
+    std::string vtk_actual_domain_basename( "domain_act" );
+    std::string vtk_lb_domain_basename( "domain_lb" );
+
     double force_time = 0;
     double comm_time = 0;
     double neigh_time = 0;
     double integrate_time = 0;
+    double lb_time = 0;
     double other_time = 0;
 
     double last_time = 0;
     Kokkos::Timer timer, force_timer, comm_timer, neigh_timer, integrate_timer,
-        other_timer;
+        lb_timer, other_timer;
 
     // Main timestep loop
     for ( int step = 1; step <= nsteps; step++ )
@@ -292,6 +316,16 @@ void CbnMD<t_System, t_Neighbor>::run()
 
         if ( step % input->comm_exchange_rate == 0 && step > 0 )
         {
+            // Update domain decomposition
+            lb_timer.reset();
+#ifdef CabanaMD_ENABLE_LB
+            double work = system->N_local + system->N_ghost;
+            auto new_global_grid = lb->createBalancedGlobalGrid(
+                system->global_mesh, *system->partitioner, work );
+            system->update_global_grid( new_global_grid );
+#endif
+            lb_time += lb_timer.seconds();
+
             // Exchange atoms across MPI ranks
             comm_timer.reset();
             comm->exchange();
@@ -361,9 +395,16 @@ void CbnMD<t_System, t_Neighbor>::run()
                 double time = timer.seconds();
                 double rate =
                     1.0 * system->N * input->thermo_rate / ( time - last_time );
+#ifdef CabanaMD_ENABLE_LB
+                log( out, std::fixed, std::setprecision( 6 ), step, " ", T, " ",
+                     PE, " ", PE + KE, " ", std::setprecision( 2 ), time, " ",
+                     std::scientific, rate, " ", std::setprecision( 2 ),
+                     lb->getImbalance() );
+#else
                 log( out, std::fixed, std::setprecision( 6 ), step, " ", T, " ",
                      PE, " ", PE + KE, " ", std::setprecision( 2 ), time, " ",
                      std::scientific, rate );
+#endif
                 last_time = time;
             }
             else
@@ -373,8 +414,22 @@ void CbnMD<t_System, t_Neighbor>::run()
                      " ", T, " ", PE, " ", PE + KE, " ", time );
                 last_time = time;
             }
+#ifdef CabanaMD_ENABLE_LB
+            double work = system->N_local + system->N_ghost;
+            std::array<double, 6> vertices;
+            vertices = lb->getVertices();
+            VTKWriter::writeDomain( MPI_COMM_WORLD, step, vertices, work,
+                                    vtk_actual_domain_basename );
+            vertices = lb->getInternalVertices();
+            VTKWriter::writeDomain( MPI_COMM_WORLD, step, vertices, work,
+                                    vtk_lb_domain_basename );
+#endif
         }
 
+        if ( step % input->vtk_rate == 0 )
+            VTKWriter::writeParticles( MPI_COMM_WORLD, step, system,
+                                       input->vtk_file, err );
+
         if ( input->dumpbinaryflag )
             dump_binary( step );
 
@@ -391,16 +446,18 @@ void CbnMD<t_System, t_Neighbor>::run()
     {
         double steps_per_sec = 1.0 * nsteps / time;
         double atom_steps_per_sec = system->N * steps_per_sec;
+        // todo(sschulz): Properly remove lb timing if not enabled.
         log( out, std::fixed, std::setprecision( 2 ),
-             "\n#Procs Atoms | Time T_Force T_Neigh T_Comm T_Int ",
+             "\n#Procs Atoms | Time T_Force T_Neigh T_Comm T_Int T_lb ",
              "T_Other |\n", comm->num_processes(), " ", system->N, " | ", time,
              " ", force_time, " ", neigh_time, " ", comm_time, " ",
-             integrate_time, " ", other_time, " | PERFORMANCE\n", std::fixed,
-             comm->num_processes(), " ", system->N, " | ", 1.0, " ",
+             integrate_time, " ", lb_time, " ", other_time, " | PERFORMANCE\n",
+             std::fixed, comm->num_processes(), " ", system->N, " | ", 1.0, " ",
              force_time / time, " ", neigh_time / time, " ", comm_time / time,
-             " ", integrate_time / time, " ", other_time / time,
-             " | FRACTION\n\n", "#Steps/s Atomsteps/s Atomsteps/(proc*s)\n",
-             std::scientific, steps_per_sec, " ", atom_steps_per_sec, " ",
+             " ", integrate_time / time, " ", lb_time / time, " ",
+             other_time / time, " | FRACTION\n\n",
+             "#Steps/s Atomsteps/s Atomsteps/(proc*s)\n", std::scientific,
+             steps_per_sec, " ", atom_steps_per_sec, " ",
              atom_steps_per_sec / comm->num_processes() );
     }
     else
@@ -409,6 +466,7 @@ void CbnMD<t_System, t_Neighbor>::run()
              " procs for ", nsteps, " steps with ", system->N, " atoms" );
     }
     out.close();
+    err.close();
 
     if ( input->write_data_flag )
         write_data( system, input->output_data_file );

src/comm_mpi_impl.h

@@ -230,7 +230,7 @@ void Comm<t_System>::exchange()
             // resized as well
             if ( pack_ranks_migrate_all.extent( 0 ) < x.size() )
             {
-                max_local *= 1.1;
+                max_local = x.size() * 1.1;
                 Kokkos::realloc( pack_ranks_migrate_all, max_local );
             }
             pack_ranks_migrate =

src/inputCL.cpp

@@ -48,6 +48,7 @@
 
 #include <Cabana_Core.hpp>
 
+#include <cstdlib>
 #include <inputCL.h>
 #include <output.h>
 #include <types.h>
@@ -116,6 +117,12 @@ void InputCL::read_args( int argc, char *argv[] )
                  "                                (N = positive integer)\n",
                  "                                (PATH = location of ",
                  "directory)\n" );
+            log( std::cout,
+                 "  --vacuum [N]:             Create a vacuum for "
+                 "an unbalanced system, enlarging the simulation "
+                 "box N times\n"
+                 "                                (N = floating-point "
+                 "multiplier, must be bigger than 1.0)\n" );
         }
 
         // Read Lammps input deck
@@ -233,6 +240,16 @@ void InputCL::read_args( int argc, char *argv[] )
             i += 3;
         }
 
+        else if ( ( strcmp( argv[i], "--vacuum" ) == 0 ) )
+        {
+            vacuum = true;
+            vacuum_rate = std::atof( argv[i + 1] );
+            if ( vacuum_rate <= 1.0 )
+                log_err( std::cout,
+                         "Vacuum multiplier must be bigger than 1.0" );
+            i += 2;
+        }
+
         else if ( ( strstr( argv[i], "--kokkos-" ) == NULL ) )
         {
             log_err( std::cout, "Unknown command line argument: ", argv[i] );

src/inputCL.h

@@ -63,6 +63,8 @@ class InputCL
     int layout_type;
     int nnp_layout_type;
     int device_type;
+    bool vacuum = false;
+    double vacuum_rate = 1.0;
 
     int dumpbinary_rate, correctness_rate;
     bool dumpbinaryflag, correctnessflag;

src/inputFile.h

@@ -181,6 +181,7 @@ class InputFile
 
     int comm_type = COMM_MPI;
     int comm_exchange_rate = 20;
+    double comm_ghost_cutoff;
 
     int force_type = FORCE_LJ;
     int force_iteration_type;
@@ -203,6 +204,9 @@ class InputFile
     std::string output_data_file;
     bool read_data_flag = false;
     bool write_data_flag = false;
+    bool write_vtk_flag = false;
+    int vtk_rate;
+    std::string vtk_file;
 
     InputFile( InputCL cl, t_System *s );
     void read_file( const char *filename = NULL );