44 #ifndef KOKKOS_VECTORIMPORT_HPP
45 #define KOKKOS_VECTORIMPORT_HPP
53 #include <Kokkos_Core.hpp>
60 template<
class CommMessageType ,
class CommIdentType ,
class VectorType >
66 #if ! defined( KOKKOS_ENABLE_MPI )
74 template<
class CommMessageType ,
class CommIdentType ,
class VectorType >
82 const CommMessageType & ,
83 const CommMessageType & ,
84 const CommIdentType & ,
85 const unsigned arg_count_owned ,
86 const unsigned arg_count_receive )
110 template<
class CommMessageType ,
class CommIdentType ,
class VectorType >
115 enum { OK = Kokkos::Impl::StaticAssert<
116 ( VectorType::rank == 1 ) ||
117 Kokkos::Impl::is_same< typename VectorType::array_layout , Kokkos::LayoutRight >::value
120 typedef typename VectorType::HostMirror HostVectorType ;
121 typedef typename CommMessageType::HostMirror HostCommMessageType;
123 enum { ReceiveInPlace =
124 Kokkos::Impl::is_same<
typename VectorType::memory_space ,
125 typename HostVectorType::memory_space >::value };
127 const CommMessageType recv_msg ;
128 const CommMessageType send_msg ;
129 const CommIdentType send_nodeid ;
130 HostCommMessageType host_recv_msg ;
131 HostCommMessageType host_send_msg ;
132 VectorType send_buffer ;
133 HostVectorType host_send_buffer ;
134 HostVectorType host_recv_buffer ;
145 const CommIdentType index ;
146 const VectorType source ;
147 const VectorType buffer ;
149 KOKKOS_INLINE_FUNCTION
151 { buffer( i ) = source( index(i) ); }
153 Pack(
const CommIdentType & arg_index ,
154 const VectorType & arg_source ,
155 const VectorType & arg_buffer )
157 , source( arg_source )
158 , buffer( arg_buffer )
160 Kokkos::parallel_for( index.extent(0) , *this );
161 execution_space::fence();
166 const CommMessageType & arg_recv_msg ,
167 const CommMessageType & arg_send_msg ,
168 const CommIdentType & arg_send_nodeid ,
169 const unsigned arg_count_owned ,
170 const unsigned arg_count_receive )
171 : recv_msg( arg_recv_msg )
172 , send_msg( arg_send_msg )
173 , send_nodeid( arg_send_nodeid )
180 , count_owned( arg_count_owned )
181 , count_receive( arg_count_receive )
187 if ( ! ReceiveInPlace ) {
188 host_recv_buffer = HostVectorType(
"recv_buffer",count_receive);
191 unsigned send_count = 0 ;
192 for (
unsigned i = 0 ; i < send_msg.extent(0) ; ++i ) { send_count += host_send_msg(i,1); }
193 send_buffer = VectorType(
"send_buffer",send_count);
203 const Teuchos::MpiComm<int> & teuchos_mpi_comm =
dynamic_cast< const Teuchos::MpiComm<int> &
>( *comm );
205 MPI_Comm mpi_comm = * teuchos_mpi_comm.getRawMpiComm();
207 const int mpi_tag = 42 ;
208 const unsigned chunk = v.extent(1);
211 const std::pair<unsigned,unsigned> recv_range( count_owned , count_owned + count_receive );
212 const VectorType recv_vector = Kokkos::subview( v , recv_range );
214 std::vector< MPI_Request > recv_request( recv_msg.extent(0) , MPI_REQUEST_NULL );
217 if (ReceiveInPlace) {
218 scalar_type * ptr = recv_vector.data();
220 for (
size_t i = 0 ; i < recv_msg.extent(0) ; ++i ) {
221 const int proc = host_recv_msg(i,0);
222 const int count = host_recv_msg(i,1) * chunk ;
224 MPI_Irecv( ptr , count *
sizeof(scalar_type) , MPI_BYTE ,
225 proc , mpi_tag , mpi_comm , & recv_request[i] );
231 host_scalar_type * ptr = host_recv_buffer.data();
233 for (
size_t i = 0 ; i < recv_msg.extent(0) ; ++i ) {
234 const int proc = host_recv_msg(i,0);
235 const int count = host_recv_msg(i,1) * chunk ;
237 MPI_Irecv( ptr , count *
sizeof(host_scalar_type) , MPI_BYTE ,
238 proc , mpi_tag , mpi_comm , & recv_request[i] );
245 MPI_Barrier( mpi_comm );
248 const Pack pack( send_nodeid , v , send_buffer );
252 host_scalar_type * ptr = host_send_buffer.data();
254 for (
size_t i = 0 ; i < send_msg.extent(0) ; ++i ) {
255 const int proc = host_send_msg(i,0);
256 const int count = host_send_msg(i,1) * chunk ;
266 count *
sizeof(host_scalar_type) , MPI_BYTE ,
267 proc , mpi_tag , mpi_comm );
275 for (
size_t i = 0 ; i < recv_msg.extent(0) ; ++i ) {
276 MPI_Status recv_status ;
280 MPI_Waitany( recv_msg.extent(0) , & recv_request[0] , & recv_which , & recv_status );
282 const int recv_proc = recv_status.MPI_SOURCE ;
284 MPI_Get_count( & recv_status , MPI_BYTE , & recv_size );
288 const int expected_proc = host_recv_msg(recv_which,0);
289 const int expected_size = host_recv_msg(recv_which,1) * chunk *
sizeof(
scalar_type);
291 if ( ( expected_proc != recv_proc ) ||
292 ( expected_size != recv_size ) ) {
296 MPI_Comm_rank( mpi_comm , & local_rank );
298 std::ostringstream msg ;
299 msg <<
"VectorImport error:"
300 <<
" P" << local_rank
301 <<
" received from P" << recv_proc
302 <<
" size " << recv_size
303 <<
" expected " << expected_size
304 <<
" from P" << expected_proc ;
305 throw std::runtime_error( msg.str() );
Kokkos::DefaultExecutionSpace execution_space
void operator()(const VectorType &) const
const Teuchos::RCP< const Teuchos::Comm< int > > comm
void deep_copy(const Stokhos::CrsMatrix< ValueType, DstDevice, Layout > &dst, const Stokhos::CrsMatrix< ValueType, SrcDevice, Layout > &src)
VectorImport(const Teuchos::RCP< const Teuchos::Comm< int > > arg_comm, const CommMessageType &, const CommMessageType &, const CommIdentType &, const unsigned arg_count_owned, const unsigned arg_count_receive)
Stokhos::CrsMatrix< ValueType, Device, Layout >::HostMirror create_mirror_view(const Stokhos::CrsMatrix< ValueType, Device, Layout > &A)
const unsigned count_receive
const unsigned count_owned