Actual source code: aijkokkosimpl.hpp
petsc-3.14.6 2021-03-30
4: #include "Kokkos_Core.hpp"
5: #include <Kokkos_DualView.hpp>
6: #include "KokkosSparse_CrsMatrix.hpp"
7: #include <petscvec.hpp>
9: using MatRowMap_t = PetscInt; /* RowMap, not RowOffset, to distinguish from Kokkos OffsetView */
10: using MatColumnIndex_t = PetscInt;
11: using MatValue_t = PetscScalar;
12: using MatDevice_t = typename Kokkos::Device<DeviceExecutionSpace,DeviceMemorySpace>;
14: using KokkosCsrMatrix_t = typename KokkosSparse::CrsMatrix<MatValue_t,MatColumnIndex_t,MatDevice_t,void/* MemoryTraits */,MatRowMap_t>;
15: using KokkosCsrGraph_t = typename KokkosCsrMatrix_t::staticcrsgraph_type;
17: using MatColumnIndexViewDevice_t = typename KokkosCsrGraph_t::entries_type;
18: using MatRowMapViewDevice_t = typename KokkosCsrGraph_t::row_map_type;
19: using MatValueViewDevice_t = typename KokkosCsrMatrix_t::values_type;
21: using MatColumnIndexViewHost_t = MatColumnIndexViewDevice_t::HostMirror;
22: using MatRowMapViewHost_t = MatRowMapViewDevice_t::HostMirror;
23: using MatValueViewHost_t = MatValueViewDevice_t::HostMirror;
25: using MatValueDualView_t = Kokkos::DualView<MatValue_t*>;
28: struct Mat_SeqAIJKokkos {
29: MatRowMapViewHost_t i_h;
30: MatColumnIndexViewHost_t j_h;
31: MatValueViewHost_t a_h;
33: MatRowMapViewDevice_t i_d;
34: MatColumnIndexViewDevice_t j_d;
35: MatValueViewDevice_t a_d;
37: MatValueDualView_t a_dual;
39: KokkosCsrMatrix_t csr;
40: PetscObjectState nonzerostate; /* State of the nonzero pattern (graph) on device */
42: Mat_SeqAIJKokkos(MatColumnIndex_t nrows,MatColumnIndex_t ncols,MatRowMap_t nnz,MatRowMap_t *i,MatColumnIndex_t *j,MatValue_t *a)
43: : i_h(i,nrows+1),
44: j_h(j,nnz),
45: a_h(a,nnz),
46: i_d(Kokkos::create_mirror_view_and_copy(DeviceMemorySpace(),i_h)),
47: j_d(Kokkos::create_mirror_view_and_copy(DeviceMemorySpace(),j_h)),
48: a_d(Kokkos::create_mirror_view_and_copy(DeviceMemorySpace(),a_h)),
49: a_dual(a_d,a_h),
50: csr("AIJKokkos",nrows,ncols,nnz,a_d,i_d,j_d)
51: {};
52: };
54: #endif