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