Actual source code: aijkokkosimpl.hpp

petsc-3.14.6 2021-03-30
Report Typos and Errors

  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