C++ VectorType::ptr_on_device方法代码示例

  inline
  void operator()( const VectorType & v ) const
  {
    typedef typename VectorType::value_type  scalar_type ;

    const Teuchos::MpiComm<int> & teuchos_mpi_comm = dynamic_cast< const Teuchos::MpiComm<int> & >( *comm );

    MPI_Comm mpi_comm = * teuchos_mpi_comm.getRawMpiComm();

    const int mpi_tag = 42 ;
    const unsigned vchunk = 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< VectorType >( 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) * vchunk ;

        MPI_Irecv( ptr , count * sizeof(scalar_type) , MPI_BYTE ,
                   proc , mpi_tag , mpi_comm , & recv_request[i] );

        ptr += count ;
      }
    }

    MPI_Barrier( mpi_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) * vchunk ;

        // 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 , mpi_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) * vchunk * sizeof(scalar_type);

      if ( ( expected_proc != recv_proc ) ||
           ( expected_size != recv_size ) ) {

        int local_rank  = 0 ;

        MPI_Comm_rank( mpi_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 ); }
//.........这里部分代码省略.........