Integrate the Cabana load balancer

.github/workflows/CI.yml

@@ -30,23 +30,32 @@ jobs:
           cmake -B build -DCMAKE_INSTALL_PREFIX=$HOME/kokkos -DKokkos_CXX_STANDARD=14 -DKokkos_ENABLE_${{ matrix.backend }}=ON
           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/[email protected]
         with:
-          repository: ECP-copa/Cabana
-          ref: master
+          repository: aetx/Cabana
+          ref: cajita_loadbalancer
           path: Cabana
       - name: Build Cabana
         working-directory: Cabana
         run: |
-          cmake -B build -DCMAKE_INSTALL_PREFIX=$HOME/Cabana -DCMAKE_PREFIX_PATH="$HOME/kokkos" -DCabana_REQUIRE_${{ matrix.backend }}=ON
+          cmake -B build -DCMAKE_INSTALL_PREFIX=$HOME/Cabana -DCMAKE_PREFIX_PATH="$HOME/kokkos" -DCabana_REQUIRE_${{ matrix.backend }}=ON -DCabana_REQUIRE_ALL=ON
           cmake --build build --parallel 2
           cmake --install build
       - name: Checkout code
         uses: actions/[email protected]
       - name: Build
         run: |
-          cmake -B build -DCMAKE_PREFIX_PATH="$HOME/kokkos;$HOME/Cabana"
+          cmake -B build -DCMAKE_PREFIX_PATH="$HOME/kokkos;$HOME/Cabana;$HOME/ALL"
           cmake --build build --parallel 2
           cd build && ctest --output-on-failure
       - name: Format

.gitignore

@@ -1 +1,2 @@
-*.mp4
+*.mp4
+*.swp

CMakeLists.txt

@@ -16,6 +16,22 @@ if( NOT Cabana_ENABLE_CAJITA )
 endif()
 find_package(Silo REQUIRED)
 
+# optional dependencies (taken from Cabana)
+macro(ExaMPM_add_dependency)
+  cmake_parse_arguments(EXAMPM_DEPENDENCY "" "PACKAGE;VERSION" "" ${ARGN})
+  find_package( ${EXAMPM_DEPENDENCY_PACKAGE} ${EXAMPM_DEPENDENCY_VERSION} QUIET )
+  string(TOUPPER "${EXAMPM_DEPENDENCY_PACKAGE}" EXAMPM_DEPENDENCY_OPTION )
+  option(
+    ExaMPM_REQUIRE_${EXAMPM_DEPENDENCY_OPTION}
+    "Require ExaMPM to build with ${EXAMPM_DEPENDENCY_PACKAGE} support" ${EXAMPM_DEPENDENCY_PACKAGE}_FOUND)
+  if(ExaMPM_REQUIRE_${EXAMPM_DEPENDENCY_OPTION})
+    find_package( ${EXAMPM_DEPENDENCY_PACKAGE} ${EXAMPM_DEPENDENCY_VERSION} REQUIRED )
+  endif()
+  set(ExaMPM_ENABLE_${EXAMPM_DEPENDENCY_OPTION} ${${EXAMPM_DEPENDENCY_PACKAGE}_FOUND})
+endmacro()
+
+ExaMPM_add_dependency( PACKAGE "ALL" VERSION 0.9.2 )
+
 # find Clang Format
 find_package( CLANG_FORMAT 10 )
 

examples/dam_break.cpp

@@ -87,7 +87,8 @@ void damBreak( const double cell_size, const int ppc, const int halo_size,
     int comm_size;
     MPI_Comm_size( MPI_COMM_WORLD, &comm_size );
     std::array<int, 3> ranks_per_dim = { 1, comm_size, 1 };
-    Cajita::ManualPartitioner partitioner( ranks_per_dim );
+    std::shared_ptr<Cajita::ManualPartitioner> partitioner =
+        std::make_shared<Cajita::ManualPartitioner>( ranks_per_dim );
 
     // Material properties.
     double bulk_modulus = 1.0e5;

examples/free_fall.cpp

@@ -86,7 +86,8 @@ void freeFall( const double cell_size, const int ppc, const int halo_size,
     int comm_size;
     MPI_Comm_size( MPI_COMM_WORLD, &comm_size );
     std::array<int, 3> ranks_per_dim = { 1, comm_size, 1 };
-    Cajita::ManualPartitioner partitioner( ranks_per_dim );
+    std::shared_ptr<Cajita::ManualPartitioner> partitioner =
+        std::make_shared<Cajita::ManualPartitioner>( ranks_per_dim );
 
     // Material properties.
     double bulk_modulus = 5.0e5;

src/CMakeLists.txt

@@ -8,9 +8,16 @@ set(HEADERS
   ExaMPM_Solver.hpp
   ExaMPM_TimeIntegrator.hpp
   ExaMPM_Types.hpp
+  ExaMPM_VTKDomainWriter.hpp
   ExaMPM_VelocityInterpolation.hpp
   )
 
+if(ExaMPM_ENABLE_ALL)
+  list(APPEND HEADERS
+    ExaMPM_LoadBalancer.hpp
+    )
+endif()
+
 set(SOURCES
   ExaMPM_Mesh.cpp
   )
@@ -22,6 +29,10 @@ target_link_libraries(exampm
   Cabana::Cajita
   Silo::silo )
 
+if(ExaMPM_ENABLE_ALL)
+  target_link_libraries(exampm ALL::ALL)
+endif()
+
 target_include_directories(exampm
   PUBLIC
     ${CMAKE_CURRENT_SOURCE_DIR}

src/ExaMPM_LoadBalancer.hpp

@@ -0,0 +1,73 @@
+/****************************************************************************
+ * Copyright (c) 2018-2020 by the ExaMPM authors                            *
+ * All rights reserved.                                                     *
+ *                                                                          *
+ * This file is part of the ExaMPM library. ExaMPM is distributed under a   *
+ * BSD 3-clause license. For the licensing terms see the LICENSE file in    *
+ * the top-level directory.                                                 *
+ *                                                                          *
+ * SPDX-License-Identifier: BSD-3-Clause                                    *
+ ****************************************************************************/
+
+#ifndef EXAMPM_LOADBALANCER_HPP
+#define EXAMPM_LOADBALANCER_HPP
+
+#include <ExaMPM_Mesh.hpp>
+#include <ExaMPM_ProblemManager.hpp>
+#include <ExaMPM_VTKDomainWriter.hpp>
+
+#include <Cajita_LoadBalancer.hpp>
+
+#include <string>
+
+#include <mpi.h>
+
+namespace ExaMPM
+{
+template <class MemorySpace>
+class LoadBalancer
+{
+  public:
+    using mesh_type = Mesh<MemorySpace>;
+    // todo(sschulz): Allow for arbitrary dimension
+    LoadBalancer( MPI_Comm comm, std::shared_ptr<mesh_type>& mesh )
+        : _comm( comm )
+        , _mesh( mesh )
+    {
+        MPI_Comm_rank( comm, &_rank );
+        _lb = std::make_shared<Cajita::LoadBalancer<UniformMesh<double>>>(
+            _mesh->localGrid()->globalGrid(), 3. * _mesh->cellSize() );
+    }
+
+    // This will update the domain decomposition and also update the mesh
+    void balance( std::shared_ptr<ProblemManager<MemorySpace>> pm )
+    {
+        double work = pm->numParticle();
+        // todo(sschulz): How to save the partitioner, without requiring a
+        // shared ptr everywhere?
+        auto partitioner = auto global_grid =
+            _lb->createBalancedGlobalGrid( _mesh->globalMesh(), partitioner,
+                                           work )
+                _mesh->newGlobalGrid( global_grid );
+        pm->updateMesh( _mesh );
+    }
+
+    // Output the actual and internal load balancing grid to VTK files
+    void output( const int t ) const
+    {
+        std::string vtk_actual_domain_basename( "domain_act" );
+        std::string vtk_lb_domain_basename( "domain_lb" );
+        VTKDomainWriter::writeDomain( _comm, t, _lb->getVertices(),
+                                      vtk_actual_domain_basename );
+        VTKDomainWriter::writeDomain( _comm, t, _lb->getInternalVertices(),
+                                      vtk_lb_domain_basename );
+    }
+
+  private:
+    MPI_Comm _comm;
+    std::shared_ptr<mesh_type> _mesh;
+    std::shared_ptr<ALL::ALL<double, double>> _liball;
+    int _rank;
+};
+} // end namespace ExaMPM
+#endif // EXAMPM_LOADBALANCER_HPP

src/ExaMPM_Mesh.hpp

@@ -91,16 +91,25 @@ class Mesh
         }
 
         // Create the global mesh.
-        auto global_mesh = Cajita::createUniformGlobalMesh(
+        _global_mesh = Cajita::createUniformGlobalMesh(
             global_low_corner, global_high_corner, num_cell );
 
         // Build the global grid.
-        auto global_grid = Cajita::createGlobalGrid( comm, global_mesh,
-                                                     periodic, partitioner );
+        _global_grid = Cajita::createGlobalGrid( comm, _global_mesh, periodic,
+                                                 partitioner );
 
         // Build the local grid.
-        int halo_width = std::max( minimum_halo_cell_width, halo_cell_width );
-        _local_grid = Cajita::createLocalGrid( global_grid, halo_width );
+        _halo_width = std::max( minimum_halo_cell_width, halo_cell_width );
+        _local_grid = Cajita::createLocalGrid( _global_grid, _halo_width );
+    }
+
+    // Global grid updated and update local values
+    void newGlobalGrid(
+        std::shared_ptr<Cajita::GlobalGrid<Cajita::UniformMesh<double>>>&
+            global_grid )
+    {
+        _global_grid = global_grid;
+        _local_grid = Cajita::createLocalGrid( _global_grid, _halo_width );
     }
 
     // Get the local grid.
@@ -110,6 +119,19 @@ class Mesh
         return _local_grid;
     }
 
+    // Get the mutable global grid.
+    const std::shared_ptr<Cajita::GlobalGrid<Cajita::UniformMesh<double>>>&
+    globalGrid()
+    {
+        return _global_grid;
+    }
+
+    const std::shared_ptr<Cajita::GlobalMesh<Cajita::UniformMesh<double>>>&
+    globalMesh() const
+    {
+        return _global_mesh;
+    }
+
     // Get the cell size.
     double cellSize() const
     {
@@ -130,6 +152,12 @@ class Mesh
 
   public:
     std::shared_ptr<Cajita::LocalGrid<Cajita::UniformMesh<double>>> _local_grid;
+    std::shared_ptr<Cajita::GlobalGrid<Cajita::UniformMesh<double>>>
+        _global_grid;
+    std::shared_ptr<Cajita::GlobalMesh<Cajita::UniformMesh<double>>>
+        _global_mesh;
+
+    int _halo_width;
 
     Kokkos::Array<int, 3> _min_domain_global_node_index;
     Kokkos::Array<int, 3> _max_domain_global_node_index;

src/ExaMPM_ProblemManager.hpp

@@ -148,6 +148,39 @@ class ProblemManager
             *cell_scalar_layout, Cajita::FullHaloPattern() );
     }
 
+    void updateMesh( const std::shared_ptr<mesh_type>& mesh )
+    {
+        auto node_vector_layout =
+            Cajita::createArrayLayout( _mesh->localGrid(), 3, Cajita::Node() );
+        auto node_scalar_layout =
+            Cajita::createArrayLayout( _mesh->localGrid(), 1, Cajita::Node() );
+        auto cell_scalar_layout =
+            Cajita::createArrayLayout( _mesh->localGrid(), 1, Cajita::Cell() );
+
+        _momentum = Cajita::createArray<double, MemorySpace>(
+            "momentum", node_vector_layout );
+        _mass = Cajita::createArray<double, MemorySpace>( "mass",
+                                                          node_scalar_layout );
+        _force = Cajita::createArray<double, MemorySpace>( "force",
+                                                           node_vector_layout );
+        _velocity = Cajita::createArray<double, MemorySpace>(
+            "velocity", node_vector_layout );
+        _position_correction = Cajita::createArray<double, MemorySpace>(
+            "position_correction", node_vector_layout );
+        _density = Cajita::createArray<double, MemorySpace>(
+            "density", cell_scalar_layout );
+
+        _mark = Cajita::createArray<double, MemorySpace>( "mark",
+                                                          cell_scalar_layout );
+
+        _node_vector_halo = Cajita::createHalo<double, MemorySpace>(
+            *node_vector_layout, Cajita::FullHaloPattern() );
+        _node_scalar_halo = Cajita::createHalo<double, MemorySpace>(
+            *node_scalar_layout, Cajita::FullHaloPattern() );
+        _cell_scalar_halo = Cajita::createHalo<double, MemorySpace>(
+            *cell_scalar_layout, Cajita::FullHaloPattern() );
+    }
+
     std::size_t numParticle() const { return _particles.size(); }
 
     const std::shared_ptr<mesh_type>& mesh() const { return _mesh; }