muxer.hpp

// FEAT3: Finite Element Analysis Toolbox, Version 3
// Copyright (C) 2010 by Stefan Turek & the FEAT group
// FEAT3 is released under the GNU General Public License version 3,
// see the file 'copyright.txt' in the top level directory for details.
 
#pragma once
 
#include <kernel/base_header.hpp>
#include <kernel/util/assertion.hpp>
#include <kernel/util/dist.hpp>
#include <kernel/lafem/dense_vector.hpp>
 
#include <vector>
 
namespace FEAT
{
  namespace Global
  {
    template<typename LocalVector_, typename Mirror_>
    class Muxer
    {
    public:
      typedef typename LocalVector_::DataType DataType;
      typedef typename LocalVector_::IndexType IndexType;
      typedef LocalVector_ LocalVectorType;
      typedef Mirror_ MirrorType;
      typedef LAFEM::DenseVector<DataType, IndexType> BufferType;
 
      template <typename LocalVector2_, typename Mirror2_>
      using MuxerType = Muxer<LocalVector2_, Mirror2_>;
 
      template <typename DataType2_, typename IndexType2_>
      using MuxerTypeByDI = Muxer<
        typename LocalVector_::template ContainerType<DataType2_, IndexType2_>,
        typename Mirror_::template MirrorType<DataType2_, IndexType2_> >;
 
    public:
      const Dist::Comm* _sibling_comm;
 
      int _parent_rank;
 
      Index _buffer_size;
 
      Mirror_ _parent_mirror;
 
      std::vector<Mirror_> _child_mirrors;
 
    public:
      Muxer() :
        _sibling_comm(nullptr),
        _parent_rank(-1),
        _buffer_size(0),
        _parent_mirror(),
        _child_mirrors()
      {
      }
 
      Muxer(Muxer&& other) :
        _sibling_comm(other._sibling_comm),
        _parent_rank(other._parent_rank),
        _buffer_size(other._buffer_size),
        _parent_mirror(std::forward<Mirror_>(other._parent_mirror)),
        _child_mirrors(std::forward<std::vector<Mirror_>>(other._child_mirrors))
      {
      }
 
      virtual ~Muxer()
      {
      }
 
      Muxer& operator=(Muxer&& other)
      {
        if(this == &other)
          return *this;
 
        _sibling_comm = other._sibling_comm;
        _parent_rank = other._parent_rank;
        _buffer_size = other._buffer_size;
        _parent_mirror = std::forward<Mirror_>(other._parent_mirror);
        _child_mirrors = std::forward<std::vector<Mirror_>>(other._child_mirrors);
 
        return *this;
      }
 
      template<typename LVT2_, typename MT2_>
      void convert(const Muxer<LVT2_, MT2_>& other)
      {
        if((void*)this == (void*)&other)
          return;
 
        this->_sibling_comm = other._sibling_comm;
        this->_parent_rank = other._parent_rank;
        this->_buffer_size = other._buffer_size;
        this->_parent_mirror.convert(other._parent_mirror);
 
        this->_child_mirrors.resize(other._child_mirrors.size());
        for(std::size_t i(0); i < other._child_mirrors.size(); ++i)
        {
          this->_child_mirrors.at(i).convert(other._child_mirrors.at(i));
        }
      }
 
      template<typename LVT2_>
      void convert(const Muxer<LVT2_, Mirror_>& other, LocalVector_&& vector, LAFEM::CloneMode mode = LAFEM::CloneMode::Shallow)
      {
        if((void*)this == (void*)&other)
          return;
 
        this->_sibling_comm = other._sibling_comm;
        this->_parent_rank = other._parent_rank;
        this->_buffer_size = Index(0); // will be rebuild by compile function
        this->_parent_mirror.clone(other._parent_mirror, mode);
 
        this->_child_mirrors.clear();
        this->_child_mirrors.resize(other._child_mirrors.size());
 
        for(std::size_t i(0); i < other._child_mirrors.size(); ++i)
        {
          this->_child_mirrors.at(i).clone(other._child_mirrors.at(i), mode);
        }
 
        // compile this muxer
        compile(std::forward<LocalVector_>(vector));
      }
 
      std::size_t bytes() const
      {
        std::size_t b = _parent_mirror.bytes();
        for(const auto& cm : _child_mirrors)
          b += cm.bytes();
        return b;
      }
 
      bool is_child() const
      {
        return (_sibling_comm != nullptr) && (_sibling_comm->size() > 0);
      }
 
      bool is_parent() const
      {
        return (_sibling_comm != nullptr) && (_sibling_comm->rank() == _parent_rank);
      }
 
      bool is_ghost() const
      {
        return is_child() && (!is_parent());
      }
 
      const Dist::Comm* get_sibling_comm() const
      {
        return _sibling_comm;
      }
 
      void set_parent(const Dist::Comm* sibling_comm_, int parent_rank, Mirror_&& parent_mirror)
      {
        _sibling_comm = sibling_comm_;
        XASSERT(_sibling_comm != nullptr);
        XASSERT(parent_rank >= 0);
        XASSERT(parent_rank < _sibling_comm->size());
        _parent_rank = parent_rank;
        _parent_mirror = std::forward<Mirror_>(parent_mirror);
      }
 
      int get_parent_rank() const
      {
        return _parent_rank;
      }
 
      const Mirror_& get_parent_mirror() const
      {
        return _parent_mirror;
      }
 
      void push_child(Mirror_&& child_mirror)
      {
        _child_mirrors.push_back(std::forward<Mirror_>(child_mirror));
      }
 
      const std::vector<Mirror_>& get_child_mirrors() const
      {
        return _child_mirrors;
      }
 
      void compile(const LocalVector_& vec_tmp_)
      {
        if(!is_child())
          return; // nothing to
 
        XASSERT(_sibling_comm != nullptr);
        XASSERT(_parent_rank >= 0);
        XASSERT(_parent_rank < _sibling_comm->size());
 
        Index bufsize(0);
 
        // Is this the parent?
        if(_sibling_comm->rank() == _parent_rank)
        {
          // loop over all child buffers and compute the largest buffer size
          for(const auto& cm : _child_mirrors)
          {
            bufsize = Math::max(bufsize, cm.buffer_size(vec_tmp_));
          }
        }
 
        // broadcast buffer size to all siblings
        _sibling_comm->bcast(&bufsize, std::size_t(1), _parent_rank);
 
        // verify against length of parent buffer
        XASSERT(bufsize >= _parent_mirror.buffer_size(vec_tmp_));
 
        // store buffer size
        this->_buffer_size = bufsize;
      }
 
      void join_send(const LocalVector_& vec_src) const
      {
        XASSERT(_sibling_comm != nullptr);
        XASSERT(_sibling_comm->size() > 1);
        XASSERT(_sibling_comm->rank() != _parent_rank); // parent must call join() instead
 
        // create parent buffer in device memory
        BufferType parent_buffer(_buffer_size);
 
        // gather source to parent buffer
        _parent_mirror.gather(parent_buffer, vec_src);
 
        // gather to parent sibling
        DataType* dummy = nullptr;
        _sibling_comm->gather(parent_buffer.elements(), _buffer_size, dummy, std::size_t(0), _parent_rank);
      }
 
      void join(const LocalVector_& vec_src, LocalVector_& vec_trg) const
      {
        // if this muxer is not a child, then this operation is a simple copy
        if((_sibling_comm == nullptr) || (_sibling_comm->size() <= 1))
        {
          vec_trg.copy(vec_src);
          return;
        }
 
        XASSERT(_sibling_comm != nullptr);
        XASSERT(_sibling_comm->size() > 0);
        XASSERT(_sibling_comm->rank() == _parent_rank);
 
        const Index num_children = Index(_child_mirrors.size());
 
        // allocate parent buffer
        BufferType parent_buffer(_buffer_size);
 
        // gather source to parent buffer
        _parent_mirror.gather(parent_buffer, vec_src);
 
        // allocate child buffers
        BufferType child_buffers(_buffer_size * num_children);
 
        // gather from siblings
        _sibling_comm->gather(parent_buffer.elements(), _buffer_size, child_buffers.elements(), _buffer_size, _parent_rank);
 
        // format target vector
        vec_trg.format();
 
        // scatter child buffers into target vector
        for(Index i(0); i < num_children; ++i)
        {
          // scatter buffers
          _child_mirrors.at(i).scatter_axpy(vec_trg, child_buffers, DataType(1), i*_buffer_size);
        }
      }
 
      void split_recv(LocalVector_& vec_trg) const
      {
        XASSERT(_sibling_comm != nullptr);
        XASSERT(_sibling_comm->size() > 1);
        XASSERT(_sibling_comm->rank() != _parent_rank); // parent must call split() instead
 
        // allocate parent buffer
        BufferType parent_buffer(_buffer_size);
 
        // receive scatter from parent sibling
        DataType* dummy = nullptr;
        _sibling_comm->scatter(dummy, std::size_t(0), parent_buffer.elements(), _buffer_size, _parent_rank);
 
        // scatter into target vector
        vec_trg.format();
        _parent_mirror.scatter_axpy(vec_trg, parent_buffer);
      }
 
      void split(LocalVector_& vec_trg, const LocalVector_& vec_src) const
      {
        // if this muxer is not a child, then this operation is a simple copy
        if((_sibling_comm == nullptr) || (_sibling_comm->size() <= 1))
        {
          vec_trg.copy(vec_src);
          return;
        }
 
        XASSERT(_sibling_comm != nullptr);
        XASSERT(_sibling_comm->size() > 0);
        XASSERT(_sibling_comm->rank() == _parent_rank);
 
        TimeStamp ts_execute;
 
        const Index num_children = Index(_child_mirrors.size());
 
        // allocate child buffers
        BufferType child_buffers(_buffer_size * num_children);
 
        // gather child buffers from target vector
        for(Index i(0); i < num_children; ++i)
        {
          _child_mirrors.at(i).gather(child_buffers, vec_src, i*_buffer_size);
        }
 
        // create parent buffer
        BufferType parent_buffer(_buffer_size);
 
        // scatter to siblings
        _sibling_comm->scatter(child_buffers.elements(), _buffer_size, parent_buffer.elements(), _buffer_size, _parent_rank);
 
        // scatter into target vector
        vec_trg.format();
        _parent_mirror.scatter_axpy(vec_trg, parent_buffer);
      }
    }; // class Muxer<...>
  } // namespace Global
} // namespace FEAT