ParallelComm.hpp

/*------------------------------------------------------------------------*/
/*                 Copyright 2010 Sandia Corporation.                     */
/*  Under terms of Contract DE-AC04-94AL85000, there is a non-exclusive   */
/*  license for use of this work by or on behalf of the U.S. Government.  */
/*  Export of this program may require a license from the                 */
/*  United States Government.                                             */
/*------------------------------------------------------------------------*/

#ifndef stk_util_parallel_ParallelComm_hpp
#define stk_util_parallel_ParallelComm_hpp

#include <cstddef>
#include <iosfwd>
#include <stk_util/parallel/Parallel.hpp>

//------------------------------------------------------------------------

namespace stk_classic {

class CommAll ;

class CommBuffer ;

bool comm_sizes( ParallelMachine ,
                 const unsigned   num_msg_bound ,
                       unsigned & num_msg_maximum ,
                 const unsigned * const send_size ,
                       unsigned * const recv_size ,
                 bool local_flag = false );

bool comm_dense_sizes( ParallelMachine ,
                       const unsigned * const send_size ,
                             unsigned * const recv_size ,
                       bool local_flag = false );

//------------------------------------------------------------------------

class CommBuffer {
public:

  template<typename T> CommBuffer &pack( const T & value );

  template<typename T> CommBuffer &pack( const T * value , size_t number );

  template<typename T> CommBuffer &unpack( T & value );

  template<typename T> CommBuffer &unpack( T * value , size_t number );

  template<typename T> CommBuffer &peek( T & value );

  template<typename T> CommBuffer &peek( T * value , size_t number );

  template<typename T> CommBuffer &skip( size_t number );

  void reset();

  size_t capacity() const ;

  size_t size() const ;

  ptrdiff_t remaining() const ;

  void * buffer() const ;

  ~CommBuffer();
  CommBuffer();

private:
  friend class CommAll ;
  friend class CommGather ;
  friend class CommBroadcast ;

  static CommBuffer * allocate( const unsigned, const unsigned * const );
  static void deallocate( const unsigned , CommBuffer * );

  void pack_overflow() const ;
  void unpack_overflow() const ;

  CommBuffer( const CommBuffer & );
  CommBuffer & operator = ( const CommBuffer & );

  typedef unsigned char * ucharp ;

  ucharp m_beg ;
  ucharp m_ptr ;
  ucharp m_end ;
};

//------------------------------------------------------------------------

class CommAll {
public:

  ParallelMachine parallel()      const { return m_comm ; }
  unsigned        parallel_size() const { return m_size ; }
  unsigned        parallel_rank() const { return m_rank ; }

  CommBuffer & send_buffer( unsigned ) const ;

  CommBuffer & recv_buffer( unsigned ) const ;

  //----------------------------------------
  CommAll();

  bool allocate_buffers( ParallelMachine ,
                         const unsigned num_msg_bounds ,
                         const unsigned * const send_size ,
                         const unsigned * const recv_size ,
                         const bool local_flag = false );

  //----------------------------------------
  explicit CommAll( ParallelMachine );

  bool allocate_buffers( const unsigned num_msg_bounds ,
                         const bool symmetric = false ,
                         const bool local_flag = false );

  //----------------------------------------
  void communicate();

  //----------------------------------------
  void swap_send_recv();

  void reset_buffers();

  ~CommAll();

private:

  CommAll( const CommAll & );
  CommAll & operator = ( const CommAll & );

  void rank_error( const char * , unsigned ) const ;

  bool allocate_buffers( const unsigned * const send_size ,
                         const unsigned * const recv_size ,
                         bool local_flag );

  ParallelMachine m_comm ;
  unsigned        m_size ;
  unsigned        m_rank ;
  unsigned        m_bound ;
  unsigned        m_max ;
  CommBuffer    * m_send ;
  CommBuffer    * m_recv ;
};

//------------------------------------------------------------------------

class CommBroadcast {
public:

  ParallelMachine parallel()      const { return m_comm ; }
  unsigned        parallel_size() const { return m_size ; }
  unsigned        parallel_rank() const { return m_rank ; }

  CommBuffer & send_buffer();

  CommBuffer & recv_buffer();

  //----------------------------------------

  CommBroadcast( ParallelMachine , unsigned root_rank );

  void communicate();

  bool allocate_buffer( const bool local_flag = false );

  ~CommBroadcast();

private:

  CommBroadcast();
  CommBroadcast( const CommBroadcast & );
  CommBroadcast & operator = ( const CommBroadcast & );

  ParallelMachine m_comm ;
  unsigned        m_size ;
  unsigned        m_rank ;
  unsigned        m_root_rank ;
  CommBuffer      m_buffer ;
};

//----------------------------------------------------------------------

class CommGather {
public:

  ParallelMachine parallel()      const { return m_comm ; }
  unsigned        parallel_size() const { return m_size ; }
  unsigned        parallel_rank() const { return m_rank ; }

  ~CommGather();

  CommGather( ParallelMachine , unsigned root_rank , unsigned send_size );

  CommBuffer & send_buffer() { return m_send ; }

  void communicate();

  CommBuffer & recv_buffer( unsigned );

  void reset(); 

private:

  CommGather();
  CommGather( const CommBroadcast & );
  CommGather & operator = ( const CommBroadcast & );

  ParallelMachine m_comm ;
  unsigned        m_size ;
  unsigned        m_rank ;
  unsigned        m_root_rank ;
  CommBuffer      m_send ;
  CommBuffer    * m_recv ;
  int           * m_recv_count ;
  int           * m_recv_displ ;
};

}

//----------------------------------------------------------------------
//----------------------------------------------------------------------
// Inlined template implementations for the CommBuffer

namespace stk_classic {

template<unsigned N> struct CommBufferAlign ;

template<>
struct CommBufferAlign<1> {
  static size_t align( size_t ) { return 0 ; }
};

template<unsigned N>
struct CommBufferAlign {
  static size_t align( size_t i ) { i %= N ; return i ? ( N - i ) : 0 ; }
};

template<typename T>
inline
CommBuffer &CommBuffer::pack( const T & value )
{
  enum { Size = sizeof(T) };
  size_t nalign = CommBufferAlign<Size>::align( m_ptr - m_beg );
  if ( m_beg ) {
    if ( m_end < m_ptr + nalign + Size ) { pack_overflow(); }
    while ( nalign ) { --nalign ; *m_ptr = 0 ; ++m_ptr ; }
    T * tmp = reinterpret_cast<T*>(m_ptr);
    *tmp = value ;
    m_ptr = reinterpret_cast<ucharp>( ++tmp );
  }
  else {
    m_ptr += nalign + Size ;
  }
  return *this;
}

template<typename T>
inline
CommBuffer &CommBuffer::pack( const T * value , size_t number )
{
  enum { Size = sizeof(T) };
  size_t nalign = CommBufferAlign<Size>::align( m_ptr - m_beg );
  if ( m_beg ) {
    if ( m_end < m_ptr + nalign + number * Size ) { pack_overflow(); }
    while ( nalign ) { --nalign ; *m_ptr = 0 ; ++m_ptr ; }
    T * tmp = reinterpret_cast<T*>(m_ptr);
    while ( number ) { --number ; *tmp = *value ; ++tmp ; ++value ; }
    m_ptr = reinterpret_cast<ucharp>( tmp );
  }
  else {
    m_ptr += nalign + number * Size ;
  }
  return *this;
}

template<typename T>
inline
CommBuffer &CommBuffer::skip( size_t number )
{
  enum { Size = sizeof(T) };
  m_ptr += CommBufferAlign<Size>::align( m_ptr - m_beg ) + Size * number ;
  if ( m_beg && m_end < m_ptr ) { unpack_overflow(); }
  return *this;
}

template<typename T>
inline
CommBuffer &CommBuffer::unpack( T & value )
{
  enum { Size = sizeof(T) };
  const size_t nalign = CommBufferAlign<Size>::align( m_ptr - m_beg );
  T * tmp = reinterpret_cast<T*>( m_ptr + nalign );
  value = *tmp ;
  m_ptr = reinterpret_cast<ucharp>( ++tmp );
  if ( m_end < m_ptr ) { unpack_overflow(); }
  return *this;
}

template<typename T>
inline
CommBuffer &CommBuffer::unpack( T * value , size_t number )
{
  enum { Size = sizeof(T) };
  const size_t nalign = CommBufferAlign<Size>::align( m_ptr - m_beg );
  T * tmp = reinterpret_cast<T*>( m_ptr + nalign );
  while ( number ) { --number ; *value = *tmp ; ++tmp ; ++value ; }
  m_ptr = reinterpret_cast<ucharp>( tmp );
  if ( m_end < m_ptr ) { unpack_overflow(); }
  return *this;
}

template<typename T>
inline
CommBuffer &CommBuffer::peek( T & value )
{
  enum { Size = sizeof(T) };
  const size_t nalign = CommBufferAlign<Size>::align( m_ptr - m_beg );
  T * tmp = reinterpret_cast<T*>( m_ptr + nalign );
  value = *tmp ;
  if ( m_end < reinterpret_cast<ucharp>(++tmp) ) { unpack_overflow(); }
  return *this;
}

template<typename T>
inline
CommBuffer &CommBuffer::peek( T * value , size_t number )
{
  enum { Size = sizeof(T) };
  const size_t nalign = CommBufferAlign<Size>::align( m_ptr - m_beg );
  T * tmp = reinterpret_cast<T*>( m_ptr + nalign );
  while ( number ) { --number ; *value = *tmp ; ++tmp ; ++value ; }
  if ( m_end < reinterpret_cast<ucharp>(tmp) ) { unpack_overflow(); }
  return *this;
}

inline
void CommBuffer::reset()
{ m_ptr = m_beg ; }

inline
size_t CommBuffer::capacity() const
{ return m_end - m_beg ; }

inline
size_t CommBuffer::size() const
{ return m_ptr - m_beg ; }

inline
ptrdiff_t CommBuffer::remaining() const
{ return m_end - m_ptr ; }

inline
void * CommBuffer::buffer() const
{ return static_cast<void*>( m_beg ); }

//----------------------------------------------------------------------
// Inline implementations for the CommAll

inline
CommBuffer & CommAll::send_buffer( unsigned p ) const
{
  if ( m_size <= p ) { rank_error("send_buffer",p); }
  return m_send[p] ;
}

inline
CommBuffer & CommAll::recv_buffer( unsigned p ) const
{
  if ( m_size <= p ) { rank_error("recv_buffer",p); }
  return m_recv[p] ;
}

}

//----------------------------------------------------------------------
//----------------------------------------------------------------------

#endif