Date: 2025-03-28T21:10:59+00:00

Actual source code: aijkok.hpp

      1: #pragma once
      2: #include <petsc_kokkos.hpp>
      3: #include <petsc/private/kokkosimpl.hpp>
      4: #include <../src/mat/impls/aij/seq/aij.h>
      5: #include <KokkosSparse_CrsMatrix.hpp>
      6: #include <KokkosSparse_spiluk.hpp>
      7: #include <string>

      9: namespace
     10: {
     11: PETSC_NODISCARD inline decltype(auto) NoInit(std::string label)
     12: {
     13:   return Kokkos::view_alloc(Kokkos::WithoutInitializing, std::move(label));
     14: }
     15: } // namespace

     17: using MatRowMapType = PetscInt;
     18: using MatColIdxType = PetscInt;
     19: using MatScalarType = PetscScalar;

     21: template <class MemorySpace>
     22: using KokkosCsrMatrixType = typename KokkosSparse::CrsMatrix<MatScalarType, MatColIdxType, MemorySpace, void /* MemoryTraits */, MatRowMapType>;
     23: template <class MemorySpace>
     24: using KokkosCsrGraphType = typename KokkosCsrMatrixType<MemorySpace>::staticcrsgraph_type;

     26: using KokkosCsrGraph     = KokkosCsrGraphType<DefaultMemorySpace>;
     27: using KokkosCsrGraphHost = KokkosCsrGraphType<Kokkos::HostSpace>;

     29: using KokkosCsrMatrix     = KokkosCsrMatrixType<DefaultMemorySpace>;
     30: using KokkosCsrMatrixHost = KokkosCsrMatrixType<Kokkos::HostSpace>;

     32: using MatRowMapKokkosView = KokkosCsrGraph::row_map_type::non_const_type;
     33: using MatColIdxKokkosView = KokkosCsrGraph::entries_type::non_const_type;
     34: using MatScalarKokkosView = KokkosCsrMatrix::values_type::non_const_type;

     36: using MatRowMapKokkosViewHost = KokkosCsrGraphHost::row_map_type::non_const_type;
     37: using MatColIdxKokkosViewHost = KokkosCsrGraphHost::entries_type::non_const_type;
     38: using MatScalarKokkosViewHost = KokkosCsrMatrixHost::values_type::non_const_type;

     40: using ConstMatRowMapKokkosView = KokkosCsrGraph::row_map_type::const_type;
     41: using ConstMatColIdxKokkosView = KokkosCsrGraph::entries_type::const_type;
     42: using ConstMatScalarKokkosView = KokkosCsrMatrix::values_type::const_type;

     44: using ConstMatRowMapKokkosViewHost = KokkosCsrGraphHost::row_map_type::const_type;
     45: using ConstMatColIdxKokkosViewHost = KokkosCsrGraphHost::entries_type::const_type;
     46: using ConstMatScalarKokkosViewHost = KokkosCsrMatrixHost::values_type::const_type;

     48: using MatRowMapKokkosDualView = Kokkos::DualView<MatRowMapType *>;
     49: using MatColIdxKokkosDualView = Kokkos::DualView<MatColIdxType *>;
     50: using MatScalarKokkosDualView = Kokkos::DualView<MatScalarType *>;

     52: using KernelHandle = KokkosKernels::Experimental::KokkosKernelsHandle<MatRowMapType, MatColIdxType, MatScalarType, DefaultExecutionSpace, DefaultMemorySpace, DefaultMemorySpace>;

     54: using KokkosTeamMemberType = Kokkos::TeamPolicy<DefaultExecutionSpace>::member_type;

     56: /* For mat->spptr of a factorized matrix */
     57: struct Mat_SeqAIJKokkosTriFactors {
     58:   MatRowMapKokkosView   iL_d, iU_d, iLt_d, iUt_d;  /* rowmap for L, U, L^t, U^t of A=LU */
     59:   MatColIdxKokkosView   jL_d, jU_d, jLt_d, jUt_d;  /* column ids */
     60:   MatScalarKokkosView   aL_d, aU_d, aLt_d, aUt_d;  /* matrix values */
     61:   KernelHandle          kh, khL, khU, khLt, khUt;  /* Kernel handles for A, L, U, L^t, U^t */
     62:   PetscBool             transpose_updated;         /* Are L^T, U^T updated wrt L, U*/
     63:   PetscBool             sptrsv_symbolic_completed; /* Have we completed the symbolic solve for L and U */
     64:   PetscScalarKokkosView workVector;

     66:   Mat_SeqAIJKokkosTriFactors(PetscInt n) : transpose_updated(PETSC_FALSE), sptrsv_symbolic_completed(PETSC_FALSE), workVector("workVector", n) { }

     68:   ~Mat_SeqAIJKokkosTriFactors() { Destroy(); }

     70:   void Destroy()
     71:   {
     72:     kh.destroy_spiluk_handle();
     73:     khL.destroy_sptrsv_handle();
     74:     khU.destroy_sptrsv_handle();
     75:     khLt.destroy_sptrsv_handle();
     76:     khUt.destroy_sptrsv_handle();
     77:     transpose_updated = sptrsv_symbolic_completed = PETSC_FALSE;
     78:   }
     79: };

     81: /* For mat->spptr of a regular matrix */
     82: struct Mat_SeqAIJKokkos {
     83:   MatRowMapKokkosDualView i_dual;
     84:   MatColIdxKokkosDualView j_dual;
     85:   MatScalarKokkosDualView a_dual;

     87:   MatRowMapKokkosDualView diag_dual; /* Diagonal pointer, built on demand */

     89:   KokkosCsrMatrix  csrmat;       /* The CSR matrix, used to call KK functions */
     90:   PetscObjectState nonzerostate; /* State of the nonzero pattern (graph) on device */

     92:   KokkosCsrMatrix     csrmatT, csrmatH;                     /* Transpose and Hermitian of the matrix (built on demand) */
     93:   PetscBool           transpose_updated, hermitian_updated; /* Are At, Ah updated wrt the matrix? */
     94:   MatRowMapKokkosView transpose_perm;                       // A permutation array making Ta(i) = Aa(perm(i)), where T = A^t

     96:   /* Construct a nrows by ncols matrix with given aseq on host. Caller also specifies a nonzero state */
     97:   Mat_SeqAIJKokkos(PetscInt nrows, PetscInt ncols, Mat_SeqAIJ *aseq, PetscObjectState nzstate, PetscBool copyValues = PETSC_TRUE)
     98:   {
     99:     auto &exec = PetscGetKokkosExecutionSpace();

    101:     MatScalarKokkosViewHost a_h(aseq->a, aseq->nz);
    102:     MatRowMapKokkosViewHost i_h(const_cast<MatRowMapType *>(aseq->i), nrows + 1);
    103:     MatColIdxKokkosViewHost j_h(aseq->j, aseq->nz);
    104:     MatRowMapKokkosViewHost diag_h(aseq->diag, nrows);

    106:     auto a_d    = Kokkos::create_mirror_view(Kokkos::WithoutInitializing, exec, a_h);
    107:     auto i_d    = Kokkos::create_mirror_view_and_copy(exec, i_h);
    108:     auto j_d    = Kokkos::create_mirror_view_and_copy(exec, j_h);
    109:     auto diag_d = Kokkos::create_mirror_view_and_copy(exec, diag_h);
    110:     a_dual      = MatScalarKokkosDualView(a_d, a_h);
    111:     i_dual      = MatRowMapKokkosDualView(i_d, i_h);
    112:     j_dual      = MatColIdxKokkosDualView(j_d, j_h);
    113:     diag_dual   = MatColIdxKokkosDualView(diag_d, diag_h);

    115:     a_dual.modify_host(); /* Since caller provided values on host */
    116:     if (copyValues) a_dual.sync_device(exec);

    118:     csrmat            = KokkosCsrMatrix("csrmat", ncols, a_d, KokkosCsrGraph(j_d, i_d));
    119:     nonzerostate      = nzstate;
    120:     transpose_updated = hermitian_updated = PETSC_FALSE;
    121:   }

    123:   /* Construct with a KokkosCsrMatrix. For performance, only i, j are copied to host, but not the matrix values. */
    124:   Mat_SeqAIJKokkos(const KokkosCsrMatrix &csr) : csrmat(csr) /* Shallow-copy csr's views to csrmat */
    125:   {
    126:     auto a_d = csr.values;
    127:     /* Get a non-const version since I don't want to deal with DualView<const T*>, which is not well defined */
    128:     MatRowMapKokkosView i_d(const_cast<MatRowMapType *>(csr.graph.row_map.data()), csr.graph.row_map.extent(0));
    129:     auto                j_d = csr.graph.entries;
    130:     auto                a_h = Kokkos::create_mirror_view(Kokkos::WithoutInitializing, Kokkos::HostSpace(), a_d);
    131:     auto                i_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), i_d);
    132:     auto                j_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), j_d);

    134:     // diag_dual is set until MatAssemblyEnd() where we copy diag from host to device
    135:     a_dual = MatScalarKokkosDualView(a_d, a_h);
    136:     a_dual.modify_device(); /* since we did not copy a_d to a_h, we mark device has the latest data */
    137:     i_dual = MatRowMapKokkosDualView(i_d, i_h);
    138:     j_dual = MatColIdxKokkosDualView(j_d, j_h);
    139:     Init();
    140:   }

    142:   Mat_SeqAIJKokkos(PetscInt nrows, PetscInt ncols, PetscInt nnz, MatRowMapKokkosDualView &i, MatColIdxKokkosDualView &j, MatScalarKokkosDualView a) : i_dual(i), j_dual(j), a_dual(a)
    143:   {
    144:     csrmat = KokkosCsrMatrix("csrmat", nrows, ncols, nnz, a.view_device(), i.view_device(), j.view_device());
    145:     Init();
    146:   }

    148:   MatScalarType *a_host_data() { return a_dual.view_host().data(); }
    149:   MatRowMapType *i_host_data() { return i_dual.view_host().data(); }
    150:   MatColIdxType *j_host_data() { return j_dual.view_host().data(); }

    152:   MatScalarType *a_device_data() { return a_dual.view_device().data(); }
    153:   MatRowMapType *i_device_data() { return i_dual.view_device().data(); }
    154:   MatColIdxType *j_device_data() { return j_dual.view_device().data(); }

    156:   MatColIdxType nrows() { return csrmat.numRows(); }
    157:   MatColIdxType ncols() { return csrmat.numCols(); }
    158:   MatRowMapType nnz() { return csrmat.nnz(); }

    160:   /* Change the csrmat size to n */
    161:   void SetColSize(MatColIdxType n) { csrmat = KokkosCsrMatrix("csrmat", n, a_dual.view_device(), csrmat.graph); }

    163:   void SetDiagonal(const MatRowMapType *diag)
    164:   {
    165:     MatRowMapKokkosViewHost diag_h(const_cast<MatRowMapType *>(diag), nrows());
    166:     auto                    diag_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), diag_h);
    167:     diag_dual                      = MatRowMapKokkosDualView(diag_d, diag_h);
    168:   }

    170:   /* Shared init stuff */
    171:   void Init(void)
    172:   {
    173:     transpose_updated = hermitian_updated = PETSC_FALSE;
    174:     nonzerostate                          = 0;
    175:   }

    177:   PetscErrorCode DestroyMatTranspose(void)
    178:   {
    179:     PetscFunctionBegin;
    180:     csrmatT = KokkosCsrMatrix(); /* Overwrite with empty matrices */
    181:     csrmatH = KokkosCsrMatrix();
    182:     PetscFunctionReturn(PETSC_SUCCESS);
    183:   }
    184: };

    186: struct MatProductData_SeqAIJKokkos {
    187:   KernelHandle kh;
    188:   PetscBool    reusesym;
    189:   MatProductData_SeqAIJKokkos() : reusesym(PETSC_FALSE) { }
    190: };

    192: PETSC_INTERN PetscErrorCode MatSetSeqAIJKokkosWithCSRMatrix(Mat, Mat_SeqAIJKokkos *);
    193: PETSC_INTERN PetscErrorCode MatCreateSeqAIJKokkosWithCSRMatrix(MPI_Comm, Mat_SeqAIJKokkos *, Mat *);
    194: PETSC_INTERN PetscErrorCode MatSeqAIJKokkosMergeMats(Mat, Mat, MatReuse, Mat *);
    195: PETSC_INTERN PetscErrorCode MatSeqAIJKokkosSyncDevice(Mat);
    196: PETSC_INTERN PetscErrorCode MatSeqAIJKokkosGetKokkosCsrMatrix(Mat, KokkosCsrMatrix *);
    197: PETSC_INTERN PetscErrorCode MatCreateSeqAIJKokkosWithKokkosCsrMatrix(MPI_Comm, KokkosCsrMatrix, Mat *);
    198: PETSC_INTERN PetscErrorCode PrintCsrMatrix(const KokkosCsrMatrix &csrmat);
    199: PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqAIJKokkos(Mat, MatType, MatReuse, Mat *);
    200: PETSC_INTERN PetscErrorCode MatSeqAIJKokkosModifyDevice(Mat);
    201: PETSC_INTERN PetscErrorCode MatSeqAIJKokkosGenerateTranspose_Private(Mat, KokkosCsrMatrix *);

    203: PETSC_INTERN PetscErrorCode MatSeqAIJGetKokkosView(Mat, MatScalarKokkosView *);
    204: PETSC_INTERN PetscErrorCode MatSeqAIJRestoreKokkosView(Mat, MatScalarKokkosView *);
    205: PETSC_INTERN PetscErrorCode MatSeqAIJGetKokkosViewWrite(Mat, MatScalarKokkosView *);
    206: PETSC_INTERN PetscErrorCode MatSeqAIJRestoreKokkosViewWrite(Mat, MatScalarKokkosView *);
    207: PETSC_INTERN PetscErrorCode MatInvertVariableBlockDiagonal_SeqAIJKokkos(Mat, const PetscIntKokkosView &, const PetscIntKokkosView &, const PetscIntKokkosView &, PetscScalarKokkosView &, PetscScalarKokkosView &);