Actual source code: mpiaijkok.hpp

  1: #if !defined(MPIAIJKOEKOS_HPP_)
  2:   #define MPIAIJKOEKOS_HPP_

  4: #include <petscsf.h>
  5: #include <../src/mat/impls/aij/mpi/mpiaij.h>
  6: #include <../src/mat/impls/aij/seq/kokkos/aijkok.hpp>

  8: struct Mat_MPIAIJKokkos {
  9:   /* MatSetValuesCOO() related stuff on device */
 10:   PetscCountKokkosView Ajmap1_d, Aperm1_d;           /* Local entries to diag */
 11:   PetscCountKokkosView Bjmap1_d, Bperm1_d;           /* Local entries to offdiag */
 12:   PetscCountKokkosView Aimap2_d, Ajmap2_d, Aperm2_d; /* Remote entries to diag */
 13:   PetscCountKokkosView Bimap2_d, Bjmap2_d, Bperm2_d; /* Remote entries to offdiag */
 14:   PetscCountKokkosView Cperm1_d;                     /* Permutation to fill send buffer. 'C' for communication */
 15:   MatScalarKokkosView  sendbuf_d, recvbuf_d;         /* Buffers for remote values in MatSetValuesCOO() */

 17:   Mat_MPIAIJKokkos(const Mat_MPIAIJ *mpiaij) :
 18:     Ajmap1_d(Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(mpiaij->Ajmap1, mpiaij->Annz + 1))),
 19:     Aperm1_d(Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(mpiaij->Aperm1, mpiaij->Atot1))),

 21:     Bjmap1_d(Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(mpiaij->Bjmap1, mpiaij->Bnnz + 1))),
 22:     Bperm1_d(Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(mpiaij->Bperm1, mpiaij->Btot1))),

 24:     Aimap2_d(Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(mpiaij->Aimap2, mpiaij->Annz2))),
 25:     Ajmap2_d(Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(mpiaij->Ajmap2, mpiaij->Annz2 + 1))),
 26:     Aperm2_d(Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(mpiaij->Aperm2, mpiaij->Atot2))),

 28:     Bimap2_d(Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(mpiaij->Bimap2, mpiaij->Bnnz2))),
 29:     Bjmap2_d(Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(mpiaij->Bjmap2, mpiaij->Bnnz2 + 1))),
 30:     Bperm2_d(Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(mpiaij->Bperm2, mpiaij->Btot2))),

 32:     Cperm1_d(Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(mpiaij->Cperm1, mpiaij->sendlen))),

 34:     sendbuf_d(Kokkos::create_mirror_view(DefaultMemorySpace(), MatScalarKokkosViewHost(mpiaij->sendbuf, mpiaij->sendlen))),
 35:     recvbuf_d(Kokkos::create_mirror_view(DefaultMemorySpace(), MatScalarKokkosViewHost(mpiaij->recvbuf, mpiaij->recvlen)))
 36:   {
 37:   }
 38: };

 40: #endif