Updating kokkos lib
git-svn-id: svn://svn.icms.temple.edu/lammps-ro/trunk@14918 f3b2605a-c512-4ea7-a41b-209d697bcdaa
This commit is contained in:
@ -1,294 +0,0 @@
|
||||
/*
|
||||
//@HEADER
|
||||
// ************************************************************************
|
||||
//
|
||||
// Kokkos v. 2.0
|
||||
// Copyright (2014) Sandia Corporation
|
||||
//
|
||||
// Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
|
||||
// the U.S. Government retains certain rights in this software.
|
||||
//
|
||||
// Redistribution and use in source and binary forms, with or without
|
||||
// modification, are permitted provided that the following conditions are
|
||||
// met:
|
||||
//
|
||||
// 1. Redistributions of source code must retain the above copyright
|
||||
// notice, this list of conditions and the following disclaimer.
|
||||
//
|
||||
// 2. Redistributions in binary form must reproduce the above copyright
|
||||
// notice, this list of conditions and the following disclaimer in the
|
||||
// documentation and/or other materials provided with the distribution.
|
||||
//
|
||||
// 3. Neither the name of the Corporation nor the names of the
|
||||
// contributors may be used to endorse or promote products derived from
|
||||
// this software without specific prior written permission.
|
||||
//
|
||||
// THIS SOFTWARE IS PROVIDED BY SANDIA CORPORATION "AS IS" AND ANY
|
||||
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
|
||||
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL SANDIA CORPORATION OR THE
|
||||
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
|
||||
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
|
||||
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
|
||||
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
|
||||
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
|
||||
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
|
||||
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
//
|
||||
// Questions? Contact H. Carter Edwards (hcedwar@sandia.gov)
|
||||
//
|
||||
// ************************************************************************
|
||||
//@HEADER
|
||||
*/
|
||||
|
||||
#ifndef KOKKOS_VECTORIMPORT_HPP
|
||||
#define KOKKOS_VECTORIMPORT_HPP
|
||||
|
||||
#include <utility>
|
||||
#include <limits>
|
||||
#include <iostream>
|
||||
#include <sstream>
|
||||
#include <stdexcept>
|
||||
|
||||
#include <Kokkos_Core.hpp>
|
||||
|
||||
#include <WrapMPI.hpp>
|
||||
|
||||
namespace Kokkos {
|
||||
namespace Example {
|
||||
|
||||
template< class CommMessageType , class CommIdentType , class VectorType >
|
||||
struct VectorImport ;
|
||||
|
||||
} // namespace Example
|
||||
} // namespace Kokkos
|
||||
|
||||
#if ! defined( KOKKOS_HAVE_MPI )
|
||||
|
||||
//----------------------------------------------------------------------------
|
||||
//----------------------------------------------------------------------------
|
||||
|
||||
namespace Kokkos {
|
||||
namespace Example {
|
||||
|
||||
template< class CommMessageType , class CommIdentType , class VectorType >
|
||||
struct VectorImport {
|
||||
|
||||
const MPI_Comm comm ;
|
||||
const unsigned count_owned ;
|
||||
const unsigned count_receive ;
|
||||
|
||||
VectorImport( MPI_Comm arg_comm ,
|
||||
const CommMessageType & ,
|
||||
const CommMessageType & ,
|
||||
const CommIdentType & ,
|
||||
const unsigned arg_count_owned ,
|
||||
const unsigned arg_count_receive )
|
||||
: comm( arg_comm )
|
||||
, count_owned( arg_count_owned )
|
||||
, count_receive( arg_count_receive )
|
||||
{}
|
||||
|
||||
inline
|
||||
void operator()( const VectorType & ) const {}
|
||||
};
|
||||
|
||||
|
||||
} // namespace Example
|
||||
} // namespace Kokkos
|
||||
|
||||
//----------------------------------------------------------------------------
|
||||
//----------------------------------------------------------------------------
|
||||
|
||||
#else /* defined( KOKKOS_HAVE_MPI ) */
|
||||
|
||||
namespace Kokkos {
|
||||
namespace Example {
|
||||
|
||||
template< class CommMessageType , class CommIdentType , class VectorType >
|
||||
class VectorImport {
|
||||
private:
|
||||
|
||||
// rank == 1 or array_layout == LayoutRight
|
||||
enum { OK = Kokkos::Impl::StaticAssert<
|
||||
( VectorType::rank == 1 ) ||
|
||||
Kokkos::Impl::is_same< typename VectorType::array_layout , Kokkos::LayoutRight >::value
|
||||
>::value };
|
||||
|
||||
typedef typename VectorType::HostMirror HostVectorType ;
|
||||
|
||||
enum { ReceiveInPlace =
|
||||
Kokkos::Impl::is_same< typename VectorType::memory_space ,
|
||||
typename HostVectorType::memory_space >::value };
|
||||
|
||||
const CommMessageType recv_msg ;
|
||||
const CommMessageType send_msg ;
|
||||
const CommIdentType send_nodeid ;
|
||||
VectorType send_buffer ;
|
||||
HostVectorType host_send_buffer ;
|
||||
HostVectorType host_recv_buffer ;
|
||||
unsigned chunk ;
|
||||
|
||||
public:
|
||||
|
||||
const MPI_Comm comm ;
|
||||
const unsigned count_owned ;
|
||||
const unsigned count_receive ;
|
||||
|
||||
struct Pack {
|
||||
typedef typename VectorType::execution_space execution_space ;
|
||||
const CommIdentType index ;
|
||||
const VectorType source ;
|
||||
const VectorType buffer ;
|
||||
|
||||
KOKKOS_INLINE_FUNCTION
|
||||
void operator()( const unsigned i ) const
|
||||
{ buffer( i ) = source( index(i) ); }
|
||||
|
||||
Pack( const CommIdentType & arg_index ,
|
||||
const VectorType & arg_source ,
|
||||
const VectorType & arg_buffer )
|
||||
: index( arg_index )
|
||||
, source( arg_source )
|
||||
, buffer( arg_buffer )
|
||||
{
|
||||
Kokkos::parallel_for( index.dimension_0() , *this );
|
||||
execution_space::fence();
|
||||
}
|
||||
};
|
||||
|
||||
VectorImport( MPI_Comm arg_comm ,
|
||||
const CommMessageType & arg_recv_msg ,
|
||||
const CommMessageType & arg_send_msg ,
|
||||
const CommIdentType & arg_send_nodeid ,
|
||||
const unsigned arg_count_owned ,
|
||||
const unsigned arg_count_receive )
|
||||
: recv_msg( arg_recv_msg )
|
||||
, send_msg( arg_send_msg )
|
||||
, send_nodeid( arg_send_nodeid )
|
||||
, send_buffer()
|
||||
, host_send_buffer()
|
||||
, host_recv_buffer()
|
||||
, comm( arg_comm )
|
||||
, count_owned( arg_count_owned )
|
||||
, count_receive( arg_count_receive )
|
||||
{
|
||||
if ( ! ReceiveInPlace ) {
|
||||
host_recv_buffer = HostVectorType("recv_buffer",count_receive);
|
||||
}
|
||||
|
||||
unsigned send_count = 0 ;
|
||||
for ( unsigned i = 0 ; i < send_msg.dimension_0() ; ++i ) { send_count += send_msg(i,1); }
|
||||
send_buffer = VectorType("send_buffer",send_count);
|
||||
host_send_buffer = Kokkos::create_mirror_view( send_buffer );
|
||||
}
|
||||
|
||||
inline
|
||||
void operator()( const VectorType & v ) const
|
||||
{
|
||||
typedef typename VectorType::value_type scalar_type ;
|
||||
|
||||
const int mpi_tag = 42 ;
|
||||
const unsigned chunk = v.dimension_1();
|
||||
|
||||
// Subvector for receives
|
||||
const std::pair<unsigned,unsigned> recv_range( count_owned , count_owned + count_receive );
|
||||
const VectorType recv_vector = Kokkos::subview( v , recv_range );
|
||||
|
||||
std::vector< MPI_Request > recv_request( recv_msg.dimension_0() , MPI_REQUEST_NULL );
|
||||
|
||||
{ // Post receives
|
||||
scalar_type * ptr =
|
||||
ReceiveInPlace ? recv_vector.ptr_on_device() : host_recv_buffer.ptr_on_device();
|
||||
|
||||
for ( size_t i = 0 ; i < recv_msg.dimension_0() ; ++i ) {
|
||||
const int proc = recv_msg(i,0);
|
||||
const int count = recv_msg(i,1) * chunk ;
|
||||
|
||||
MPI_Irecv( ptr , count * sizeof(scalar_type) , MPI_BYTE ,
|
||||
proc , mpi_tag , comm , & recv_request[i] );
|
||||
|
||||
ptr += count ;
|
||||
}
|
||||
}
|
||||
|
||||
MPI_Barrier( comm );
|
||||
|
||||
{ // Pack and send
|
||||
const Pack pack( send_nodeid , v , send_buffer );
|
||||
|
||||
Kokkos::deep_copy( host_send_buffer , send_buffer );
|
||||
|
||||
scalar_type * ptr = host_send_buffer.ptr_on_device();
|
||||
|
||||
for ( size_t i = 0 ; i < send_msg.dimension_0() ; ++i ) {
|
||||
const int proc = send_msg(i,0);
|
||||
const int count = send_msg(i,1) * chunk ;
|
||||
|
||||
// MPI_Ssend blocks until
|
||||
// (1) a receive is matched for the message and
|
||||
// (2) the send buffer can be re-used.
|
||||
//
|
||||
// It is suggested that MPI_Ssend will have the best performance:
|
||||
// http://www.mcs.anl.gov/research/projects/mpi/sendmode.html .
|
||||
|
||||
MPI_Ssend( ptr ,
|
||||
count * sizeof(scalar_type) , MPI_BYTE ,
|
||||
proc , mpi_tag , comm );
|
||||
|
||||
ptr += count ;
|
||||
}
|
||||
}
|
||||
|
||||
// Wait for receives and verify:
|
||||
|
||||
for ( size_t i = 0 ; i < recv_msg.dimension_0() ; ++i ) {
|
||||
MPI_Status recv_status ;
|
||||
int recv_which = 0 ;
|
||||
int recv_size = 0 ;
|
||||
|
||||
MPI_Waitany( recv_msg.dimension_0() , & recv_request[0] , & recv_which , & recv_status );
|
||||
|
||||
const int recv_proc = recv_status.MPI_SOURCE ;
|
||||
|
||||
MPI_Get_count( & recv_status , MPI_BYTE , & recv_size );
|
||||
|
||||
// Verify message properly received:
|
||||
|
||||
const int expected_proc = recv_msg(recv_which,0);
|
||||
const int expected_size = recv_msg(recv_which,1) * chunk * sizeof(scalar_type);
|
||||
|
||||
if ( ( expected_proc != recv_proc ) ||
|
||||
( expected_size != recv_size ) ) {
|
||||
|
||||
int local_rank = 0 ;
|
||||
|
||||
MPI_Comm_rank( comm , & local_rank );
|
||||
|
||||
std::ostringstream msg ;
|
||||
msg << "VectorImport error:"
|
||||
<< " P" << local_rank
|
||||
<< " received from P" << recv_proc
|
||||
<< " size " << recv_size
|
||||
<< " expected " << expected_size
|
||||
<< " from P" << expected_proc ;
|
||||
throw std::runtime_error( msg.str() );
|
||||
}
|
||||
}
|
||||
|
||||
// Copy received data to device memory.
|
||||
|
||||
if ( ! ReceiveInPlace ) { Kokkos::deep_copy( recv_vector , host_recv_buffer ); }
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace Example
|
||||
} // namespace Kokkos
|
||||
|
||||
#endif
|
||||
|
||||
//----------------------------------------------------------------------------
|
||||
|
||||
#endif /* #ifndef KOKKOS_VECTORIMPORT_HPP */
|
||||
|
||||
|
||||
Reference in New Issue
Block a user