loadbalance.hh

// -*- tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 2 -*-
// vi: set et ts=4 sw=2 sts=2:
#ifndef DUNE_PDELAB_GRIDFUNCTIONSPACE_LOADBALANCE_HH
#define DUNE_PDELAB_GRIDFUNCTIONSPACE_LOADBALANCE_HH
 
#include <dune/geometry/dimension.hh>
#include <dune/grid/common/datahandleif.hh>
#include <dune/grid/common/partitionset.hh>
 
#include<dune/pdelab/common/polymorphicbufferwrapper.hh>
#include <dune/pdelab/gridfunctionspace/entityindexcache.hh>
 
namespace Dune {
  namespace PDELab {
 
#ifndef DOXYGEN
    namespace impl{
 
      template<typename... T>
      class LoadBalanceDataHandle
        : public Dune::CommDataHandleIF<LoadBalanceDataHandle<T...>,
                                        typename std::decay<typename std::tuple_element<0,std::tuple<T...>>::type>::type
                                        ::Map::mapped_type::value_type>
      {
      public:
 
        using R0 = typename std::decay<typename std::tuple_element<0,std::tuple<T...>>::type>::type
          ::Map::mapped_type::value_type;
 
        LoadBalanceDataHandle(std::tuple<T&...>&& mapTuple)
          : _mapTuple(mapTuple)
        {
        }
 
        bool contains (int dim, int codim) const
        {
          return true;
        }
 
        bool fixedSize (int dim, int codim) const
        {
          // We return false here, since gather and scatter is called
          // for all entities of all levels of the grid but we only
          // communicate data for a given gridView. This means there
          // will be gather and scatter calls for cells that don't exist
          // in our current grid view and we won't communicate anything
          // for these cells.
          return false;
        }
 
        // End of tmp recursion
        template <std::size_t I, typename EntityType>
        inline typename std::enable_if<I==sizeof...(T),void>::type
        sizeTMP(std::size_t& commSize, EntityType& e) const
        {
        }
 
        // Go through tuple and accumulate communication size
        template <std::size_t I=0, typename EntityType>
        inline typename std::enable_if<I<sizeof...(T),void>::type
        sizeTMP(std::size_t& commSize, EntityType& e) const
        {
          // Compare if data type for this entry equals data typo of first entry
          using R = typename std::decay<typename std::tuple_element<I,std::tuple<T...>>::type>::type
            ::Map::mapped_type::value_type;
          static_assert(std::is_same<R,R0>::value,"Different field type of vectors not supported by DH");
 
          // Current function space
          auto& gfs = std::get<I>(_mapTuple)._gfs;
 
          // Only communicate if there are dofs for this codimension.
          if (gfs.finiteElementMap().hasDOFs(e.codimension)){
 
            // We first have to communicate the size of data we send for this particular vector
            commSize += sizeof(R);
 
            // Only communicate data if entity lies in our entitySet
            if (gfs.entitySet().contains(e)){
              // Get some types
              using GFS = typename std::decay<decltype(gfs)>::type;
              using EntitySet = typename GFS::Traits::EntitySet;
              using IDSet = typename EntitySet::Traits::GridView::Grid::LocalIdSet;
 
              // Find element id in map
              const IDSet& idSet(gfs.entitySet().gridView().grid().localIdSet());
              auto& map = std::get<I>(_mapTuple)._map;
              auto find = map.find(idSet.id(e));
              assert (find!=map.end());
 
              // Add size of degrees of freedom vetor
              commSize += find->second.size()*sizeof(R);
            }
          }
 
          // Get size for next vector
          sizeTMP<I+1>(commSize,e);
        }
 
        template<class EntityType>
        std::size_t size (EntityType& e) const
        {
          // Use tmp to sum up sizes for all vectors
          std::size_t commSize(0.0);
          sizeTMP<0>(commSize,e);
          return commSize;
        }
 
 
        // End of tmp recursion
        template <std::size_t I, typename Buf, typename Entity>
        inline typename std::enable_if<I==sizeof...(T),void>::type
        gatherTMP(Buf& buf, const Entity& e) const
        {
        }
 
        // Go through tuple and call gather for all vectors
        template <std::size_t I=0, typename Buf, typename Entity>
        inline typename std::enable_if<I<sizeof...(T),void>::type
        gatherTMP(Buf& buf, const Entity& e) const
        {
          // Current gridFunctionSpace and dof map for the vector
          auto& gfs = std::get<I>(_mapTuple)._gfs;
          auto& map = std::get<I>(_mapTuple)._map;
 
          // Communication is called for every level and every entity of
          // every codim of all entities where load balance will result in
          // changes.   We only send dofs for current gridView/entitySet.
          if (gfs.finiteElementMap().hasDOFs(e.codimension)){
            if (gfs.entitySet().contains(e)){
              // Get important types
              using GFS = typename std::decay<decltype(gfs)>::type;
              using EntitySet = typename GFS::Traits::EntitySet;
              using IDSet = typename EntitySet::Traits::GridView::Grid::LocalIdSet;
 
              // Find element id in map
              const IDSet& idSet(gfs.entitySet().gridView().grid().localIdSet());
 
              // Find entity id in map.
              auto find = map.find(idSet.id(e));
              assert (find!=map.end());
 
              // Send size we need to communicate for this vector
              buf.write (static_cast<R0>(find->second.size()));
 
              // Send dofs
              for (size_t i=0; i<find->second.size(); ++i){
                buf.write(find->second[i]);
              }
            }
            else {
              // Only communicate that we don't communicate any DOFs
              R0 tmp(0);
              buf.write (tmp);
            }
          } // hasDOFs
 
          // Call gather for next vector
          gatherTMP<I+1> (buf,e);
        }
 
        template<class MessageBuffer, class EntityType>
        void gather (MessageBuffer& buff, const EntityType& e) const
        {
          // Communicate different things than char.
          using Buf = Dune::PDELab::PolymorphicBufferWrapper<MessageBuffer>;
          Buf bufWrapper(buff,Dune::PDELab::PolymorphicBufferWrapper<MessageBuffer>::Mode::send);
 
          // Call gather for all vectors using tmp
          gatherTMP<0> (bufWrapper,e);
        }
 
        // End of tmp reucrsion
        template <std::size_t I, typename Buf, typename Entity>
        inline typename std::enable_if<I==sizeof...(T),void>::type
        scatterTMP(Buf& buf, const Entity& e) const
        {
        }
 
        // Go through tuple receive DOFs
        template <std::size_t I=0, typename Buf, typename Entity>
        inline typename std::enable_if<I<sizeof...(T),void>::type
        scatterTMP(Buf& buf, const Entity& e) const
        {
          auto& gfs = std::get<I>(_mapTuple)._gfs;
          auto& map = std::get<I>(_mapTuple)._map;
          if (gfs.finiteElementMap().hasDOFs(e.codimension)){
 
            // Receive number of DOFs for this vector
            R0 tmp;
            buf.read(tmp);
            std::size_t numberOfEntries(0);
            numberOfEntries = (size_t) tmp;
 
            // Create vector of DOFs and receive DOFs
            std::vector<R0> dofs(numberOfEntries);
            for (size_t i=0; i<numberOfEntries; ++i){
              buf.read(dofs[i]);
            }
 
            // Store id and dofs in map
            const auto& id_set = gfs.entitySet().grid().localIdSet();
            map.insert({{id_set.id(e),dofs}});
          }
 
          // Call scatter for next vetcor using tmp
          scatterTMP<I+1> (buf,e);
        }
 
 
        template<class MessageBuffer, class EntityType>
        void scatter (MessageBuffer& buff, const EntityType& e, size_t n)
        {
          // Communicate different things than char.
          using Buf = Dune::PDELab::PolymorphicBufferWrapper<MessageBuffer>;
          Buf bufWrapper(buff,Dune::PDELab::PolymorphicBufferWrapper<MessageBuffer>::Mode::receive);
 
          // Call scatter for all vectors using tmp
          scatterTMP<0> (bufWrapper, e);
        }
 
      private:
        // Tuple storing GFSandMap structs for every vector that should get adapted
        std::tuple<T&...>& _mapTuple;
      };
 
 
      template <typename GFS, typename V, typename MAP, int codim>
      void loadBalanceMapFiller (const GFS& gfs, V& v, MAP& map)
      {
        using IndexCache = Dune::PDELab::EntityIndexCache<GFS>;
        using LocalView =  typename V::template LocalView<IndexCache>;
        IndexCache indexCache(gfs);
        LocalView localView(v);
        const auto& id_set = gfs.entitySet().grid().localIdSet();
 
        // Iterate over all interiorBorder entities of codim.
        for (const auto& e : entities(gfs.entitySet(),Dune::Codim<codim>(),Dune::Partitions::interiorBorder)){
          // Bind cache to entity.
          indexCache.update(e);
          localView.bind(indexCache);
 
          // Store dofs of entity in std::vector.
          std::vector<typename LocalView::ElementType> dofs;
          for (std::size_t i=0; i<localView.size(); ++i){
            dofs.push_back(localView[i]);
          }
 
          // Insert id and vector in map.
          map.insert ( {{id_set.id(e),dofs}});
 
          // Unbind cache.
          localView.unbind();
        }
      }
 
      template <int codim>
      struct FillLoadBalanceDOFMap
      {
        template <typename GFS, typename V, typename MAP>
        static void fillMap (const GFS& gfs, V& v, MAP& map)
        {
          if (gfs.finiteElementMap().hasDOFs(codim)){
            loadBalanceMapFiller<GFS,V,MAP,codim>(gfs,v,map);
          }
          FillLoadBalanceDOFMap<codim-1>::fillMap(gfs,v,map);
        }
      };
      template <>
      struct FillLoadBalanceDOFMap<0>
      {
        template <typename GFS, typename V, typename MAP>
        static void fillMap (const GFS& gfs, V& v, MAP& map)
        {
          if (gfs.finiteElementMap().hasDOFs(0)){
            loadBalanceMapFiller<GFS,V,MAP,0>(gfs,v,map);
          }
        }
      };
 
      template <typename GFS, typename V, typename MAP, int codim>
      void loadBalanceMapReader (const GFS& gfs, V& v, MAP& map)
      {
        using IndexCache = Dune::PDELab::EntityIndexCache<GFS>;
        using LocalView =  typename V::template LocalView<IndexCache>;
        IndexCache indexCache(gfs);
        LocalView localView(v);
        const auto& id_set = gfs.entitySet().grid().localIdSet();
 
        // Iterate over all interiorBorder entities of codim.
        for (const auto& e : entities(gfs.entitySet(),Dune::Codim<codim>(),Dune::Partitions::interiorBorder)){
          // Bind cache to entity.
          indexCache.update(e);
          localView.bind(indexCache);
 
          // Find key in map and get vector of dofs.
          auto find = map.find(id_set.id(e));
          auto& dofs(find->second);
 
          // Assert that we found element and that sizes of dof vector and local view match.
          assert(find!=map.end());
          assert(dofs.size()==localView.size());
 
          // Store dofs in local view.
          for (std::size_t i=0; i<dofs.size(); ++i){
            localView[i]=dofs[i];
          }
 
          // Write changes to underlying container.
          localView.commit();
 
          // Unbind cache.
          localView.unbind();
        }
      }
 
      template <int codim>
      struct ReadLoadBalanceDOFMap
      {
        template <typename GFS, typename V, typename MAP>
        static void readMap (const GFS& gfs, V& v, MAP& map)
        {
          if (gfs.finiteElementMap().hasDOFs(codim)){
            loadBalanceMapReader<GFS,V,MAP,codim>(gfs,v,map);
          }
          ReadLoadBalanceDOFMap<codim-1>::readMap(gfs,v,map);
        }
      };
      template <>
      struct ReadLoadBalanceDOFMap<0>
      {
        template <typename GFS, typename V, typename MAP>
        static void readMap (const GFS& gfs, V& v, MAP& map)
        {
          if (gfs.finiteElementMap().hasDOFs(0)){
            loadBalanceMapReader<GFS,V,MAP,0>(gfs,v,map);
          }
        }
      };
 
 
      // Store reference to function space and map
      template <typename G, typename M>
      struct GFSAndMap
      {
        // Export types
        using GFS = G;
        using Map = M;
 
        GFSAndMap (GFS& gfs, Map& m) : _gfs(gfs), _map(m)
        {
        }
 
        GFS& _gfs;
        Map& _map;
      };
 
      // Create a GFSAndMap struct
      template <typename GFS, typename M>
      GFSAndMap<GFS,M> packGFSAndMap(GFS& gfs, M& m)
      {
        GFSAndMap<GFS,M> pack(gfs,m);
        return pack;
      }
 
 
      // Forward declarations needed for the tmp recursion
      template <typename Grid, typename... T>
      void iteratePacks(Grid& grid, std::tuple<T&...>&& mapTuple);
      template <typename Grid, typename... T, typename X, typename... XS>
      void iteratePacks(Grid& grid, std::tuple<T&...>&& mapTuple, X& x, XS&... xs);
 
      // This function is called after the last vector of the tuple.  Here
      // the next pack is called.  On the way back we update the current
      // function space.
      template<std::size_t I = 0, typename Grid, typename... T, typename X, typename... XS>
      inline typename std::enable_if<I == std::tuple_size<typename X::Tuple>::value, void>::type
      iterateTuple(Grid& grid, std::tuple<T&...>&& mapTuple, X& x, XS&... xs)
      {
        // Iterate next pack
        iteratePacks(grid,std::move(mapTuple),xs...);
 
        // On our way back we need to update the current function space
        x._gfs.update(true);
      }
 
      /* In this function we store the data of the current vector (indicated
       * by template parameter I) of the current pack. After recursively
       * iterating through the other packs and vectors we replay the data.
       *
       * @tparam I      std:size_t used for tmp
       * @tparam Grid   Grid type
       * @tparam T...   Types of tuple elements (for storing data transfer maps)
       * @tparam X      Current  pack
       * @tparam ...XS  Remaining packs
       */
      template<std::size_t I = 0, typename Grid, typename... T, typename X, typename... XS>
      inline typename std::enable_if<I < std::tuple_size<typename X::Tuple>::value, void>::type
      iterateTuple(Grid& grid, std::tuple<T&...>&& mapTuple, X& x, XS&... xs)
      {
        // Get some basic types
        using GFS = typename X::GFS;
        using Tuple = typename X::Tuple;
        using V  = typename std::decay<typename std::tuple_element<I,Tuple>::type>::type;
        using IDSet = typename Grid::LocalIdSet;
        using ID = typename IDSet::IdType;
        using R = typename V::field_type;
 
        // Store vector of dofs for every entitiy of every codim.
        using MAP = std::unordered_map<ID,std::vector<R>>;
        MAP map;
        FillLoadBalanceDOFMap<GFS::Traits::GridView::dimension>::fillMap(x._gfs,std::get<I>(x._tuple),map);
 
        // Pack gfs and tuple
        auto mapPack = packGFSAndMap(x._gfs,map);
        auto newMapTuple = std::tuple_cat(mapTuple,std::tie(mapPack));
 
        // Recursively iterate through remaining vectors (and packs). Load
        // balancing will be done at the end of recursion.
        iterateTuple<I+1>(grid,std::move(newMapTuple),x,xs...);
 
        // Restore solution from dof map.
        std::get<I>(x._tuple) = V(x._gfs,0.0);
        ReadLoadBalanceDOFMap<GFS::Traits::GridView::dimension>::readMap(x._gfs,std::get<I>(x._tuple),map);
      }
 
      template <typename... T>
      LoadBalanceDataHandle<T...> createLoadBalanceDataHandle (std::tuple<T&...>&& mapTuple)
      {
        LoadBalanceDataHandle<T...> dh(std::move(mapTuple));
        return dh;
      }
 
      // This gets called after the last pack.  After this function call we
      // have visited every vector of every pack and we will go back through
      // the recursive function calls.
      template <typename Grid, typename... T>
      void iteratePacks(Grid& grid, std::tuple<T&...>&& mapTuple)
      {
        // Create data handle and use load balancing with communication.
        auto dh = createLoadBalanceDataHandle(std::move(mapTuple));
 
        std::cout << "Calling load balance with data communication" << std::endl;
        grid.loadBalance(dh);
      }
 
      /* Use template meta programming to iterate over packs at compile time
       *
       * In order to adapt our grid and all vectors of all packs we need to
       * do the following:
       * - Iterate over all vectors of all packs.
       * - Store the data from the vectors where things could change.
       * - Load Balance our grid.
       * - Update function spaces and restore data.
       *
       * The key point is that we need the object that stores the data to
       * replay it.  Because of that we can not just iterate over the packs
       * and within each pack iterate over the vectors but we have to make
       * one big recursion.  Therefore we iterate over the vectors of the
       * current pack.
       */
      template <typename Grid, typename... T, typename X, typename... XS>
      void iteratePacks(Grid& grid, std::tuple<T&...>&& mapTuple, X& x, XS&... xs)
      {
        iterateTuple(grid,std::move(mapTuple),x,xs...);
      }
 
    } // namespace impl
#endif // DOXYGEN
 
 
    template <typename Grid, typename... X>
    void loadBalanceGrid(Grid& grid, X&... x)
    {
      // Create tuple where all data transfer maps get stored
      std::tuple<> mapTuple;
 
      // Iterate over packs
      impl::iteratePacks(grid,std::move(mapTuple),x...);
    }
 
  } // namespace PDELab
} // namespace Dune
 
#endif // DUNE_PDELAB_GRIDFUNCTIONSPACE_LOADBALANCE_HH