Date: 2026-07-30T21:14:13+00:00

Actual source code: aijkok.hpp

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

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

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

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

     27: using KokkosCsrGraph     = KokkosCsrGraphType<DefaultMemorySpace>;
     28: using KokkosCsrGraphHost = KokkosCsrGraphType<HostMirrorMemorySpace>;

     30: using KokkosCsrMatrix     = KokkosCsrMatrixType<DefaultMemorySpace>;
     31: using KokkosCsrMatrixHost = KokkosCsrMatrixType<HostMirrorMemorySpace>;

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

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

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

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

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

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

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

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

     67:   MatRowMapKokkosViewHost iL_h, iU_h, iLt_h, iUt_h; // temp. buffers when we do factorization with PETSc on host. We copy L, U to these buffers and then copy to device
     68:   MatColIdxKokkosViewHost jL_h, jU_h, jLt_h, jUt_h;
     69:   MatScalarKokkosViewHost aL_h, aU_h, aLt_h, aUt_h, D_h; // D is for LDLT factorization
     70:   MatScalarKokkosView     D_d;
     71:   Mat                     L, U, Lt, Ut; // MATSEQAIJ on host if needed. Their arrays are alias to (iL_h, jL_h, aL_h), (iU_h, jU_h, aU_h) and their transpose.
     72:                                         // MatTranspose() on host might be faster than KK's csr transpose on device.

     74:   PetscIntKokkosView rowperm, colperm; // row permutation and column permutation

     76:   Mat_SeqAIJKokkosTriFactors(PetscInt n) : workVector("workVector", n)
     77:   {
     78:     L = U = Lt = Ut   = nullptr;
     79:     transpose_updated = sptrsv_symbolic_completed = PETSC_FALSE;
     80:   }

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

     84:   void Destroy()
     85:   {
     86:     PetscFunctionBeginUser;
     87:     kh.destroy_spiluk_handle();
     88:     khL.destroy_sptrsv_handle();
     89:     khU.destroy_sptrsv_handle();
     90:     khLt.destroy_sptrsv_handle();
     91:     khUt.destroy_sptrsv_handle();
     92:     PetscCallVoid(MatDestroy(&L));
     93:     PetscCallVoid(MatDestroy(&U));
     94:     PetscCallVoid(MatDestroy(&Lt));
     95:     PetscCallVoid(MatDestroy(&Ut));
     96:     L = U = Lt = Ut   = nullptr;
     97:     transpose_updated = sptrsv_symbolic_completed = PETSC_FALSE;
     98:     PetscFunctionReturnVoid();
     99:   }
    100: };

    102: /* For mat->spptr of a regular matrix */
    103: struct Mat_SeqAIJKokkos {
    104:   MatRowMapKokkosDualView i_dual;
    105:   MatColIdxKokkosDualView j_dual;
    106:   MatScalarKokkosDualView a_dual;
    107:   PetscBool               host_aij_allocated_by_kokkos = PETSC_FALSE; /* Are host views of a, i, j in the duals allocated by Kokkos? */

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

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

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

    118:   /* Construct a nrows by ncols matrix with given aseq on host. Caller also specifies a nonzero state */
    119:   Mat_SeqAIJKokkos(Mat A, PetscInt nrows, PetscInt ncols, Mat_SeqAIJ *aseq, PetscObjectState nzstate, PetscBool copyValues = PETSC_TRUE)
    120:   {
    121:     auto exec = PetscGetKokkosExecutionSpace();

    123:     // a->diag must exist before new Mat_SeqAIJKokkos()
    124:     const PetscInt *adiag;
    125:     PetscCallVoid(MatGetDiagonalMarkers_SeqAIJ(A, &adiag, NULL));

    127:     MatScalarKokkosViewHost a_h(aseq->a, aseq->nz);
    128:     MatRowMapKokkosViewHost i_h(const_cast<MatRowMapType *>(aseq->i), nrows + 1);
    129:     MatColIdxKokkosViewHost j_h(aseq->j, aseq->nz);
    130:     MatRowMapKokkosViewHost diag_h(aseq->diag, nrows);

    132:     auto a_d    = Kokkos::create_mirror_view(Kokkos::WithoutInitializing, exec, a_h);
    133:     auto i_d    = Kokkos::create_mirror_view_and_copy(exec, i_h);
    134:     auto j_d    = Kokkos::create_mirror_view_and_copy(exec, j_h);
    135:     auto diag_d = Kokkos::create_mirror_view_and_copy(exec, diag_h);
    136:     a_dual      = MatScalarKokkosDualView(a_d, a_h);
    137:     i_dual      = MatRowMapKokkosDualView(i_d, i_h);
    138:     j_dual      = MatColIdxKokkosDualView(j_d, j_h);
    139:     diag_dual   = MatColIdxKokkosDualView(diag_d, diag_h);

    141:     a_dual.modify_host(); /* Since caller provided values on host */
    142:     if (copyValues) (void)KokkosDualViewSyncDevice(a_dual, exec);

    144:     csrmat = KokkosCsrMatrix("csrmat", ncols, a_d, KokkosCsrGraph(j_d, i_d));
    145:     Init(nzstate);
    146:   }

    148:   /* Construct with a KokkosCsrMatrix. For performance, only i, j are copied to host, but not the matrix values. */
    149:   Mat_SeqAIJKokkos(const KokkosCsrMatrix &csr) : csrmat(csr) /* Shallow-copy csr's views to csrmat */
    150:   {
    151:     auto a_d = csr.values;
    152:     /* Get a non-const version since I don't want to deal with DualView<const T*>, which is not well defined */
    153:     MatRowMapKokkosView i_d(const_cast<MatRowMapType *>(csr.graph.row_map.data()), csr.graph.row_map.extent(0));
    154:     auto                j_d = csr.graph.entries;
    155:     auto                a_h = Kokkos::create_mirror_view(Kokkos::WithoutInitializing, HostMirrorMemorySpace(), a_d);
    156:     auto                i_h = Kokkos::create_mirror_view_and_copy(HostMirrorMemorySpace(), i_d);
    157:     auto                j_h = Kokkos::create_mirror_view_and_copy(HostMirrorMemorySpace(), j_d);

    159:     // diag_dual is set until MatAssemblyEnd() where we copy diag from host to device
    160:     a_dual = MatScalarKokkosDualView(a_d, a_h);
    161:     a_dual.modify_device(); /* since we did not copy a_d to a_h, we mark device has the latest data */
    162:     i_dual                       = MatRowMapKokkosDualView(i_d, i_h);
    163:     j_dual                       = MatColIdxKokkosDualView(j_d, j_h);
    164:     host_aij_allocated_by_kokkos = PETSC_TRUE; /* That means after deleting aijkok, one shouldn't access aijseq->{a,i,j} anymore! */
    165:     Init();
    166:   }

    168:   // Don't use DualView argument types as we want to be sure that a,i,j on host are allocated by Mat_SeqAIJKokkos itself (vs. by users)
    169:   Mat_SeqAIJKokkos(PetscInt nrows, PetscInt ncols, PetscInt nnz, const MatRowMapKokkosView &i_d, const MatColIdxKokkosView &j_d, const MatScalarKokkosView &a_d) : Mat_SeqAIJKokkos(KokkosCsrMatrix("csrmat", nrows, ncols, nnz, a_d, i_d, j_d)) { }

    171:   MatScalarType *a_host_data() { return a_dual.view_host().data(); }
    172:   MatRowMapType *i_host_data() { return i_dual.view_host().data(); }
    173:   MatColIdxType *j_host_data() { return j_dual.view_host().data(); }

    175:   MatScalarType *a_device_data() { return a_dual.view_device().data(); }
    176:   MatRowMapType *i_device_data() { return i_dual.view_device().data(); }
    177:   MatColIdxType *j_device_data() { return j_dual.view_device().data(); }

    179:   MatColIdxType nrows() { return csrmat.numRows(); }
    180:   MatColIdxType ncols() { return csrmat.numCols(); }
    181:   MatRowMapType nnz() { return csrmat.nnz(); }

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

    186:   void SetDiagonal(const MatRowMapType *diag)
    187:   {
    188:     MatRowMapKokkosViewHost diag_h(const_cast<MatRowMapType *>(diag), nrows());
    189:     auto                    diag_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), diag_h);
    190:     diag_dual                      = MatRowMapKokkosDualView(diag_d, diag_h);
    191:   }

    193:   /* Shared init stuff */
    194:   void Init(PetscObjectState nzstate = 0)
    195:   {
    196:     nonzerostate      = nzstate;
    197:     transpose_updated = PETSC_FALSE;
    198:     hermitian_updated = PETSC_FALSE;
    199:   }

    201:   PetscErrorCode DestroyMatTranspose(void)
    202:   {
    203:     PetscFunctionBegin;
    204:     csrmatT = KokkosCsrMatrix(); /* Overwrite with empty matrices */
    205:     csrmatH = KokkosCsrMatrix();
    206:     PetscFunctionReturn(PETSC_SUCCESS);
    207:   }
    208: };

    210: struct MatProductCtx_SeqAIJKokkos {
    211:   KernelHandle kh;
    212:   PetscBool    reusesym;
    213:   MatProductCtx_SeqAIJKokkos() : reusesym(PETSC_FALSE) { }
    214: };

    216: PETSC_INTERN PetscErrorCode MatSetSeqAIJKokkosWithCSRMatrix(Mat, Mat_SeqAIJKokkos *);
    217: PETSC_INTERN PetscErrorCode MatCreateSeqAIJKokkosWithCSRMatrix(MPI_Comm, Mat_SeqAIJKokkos *, Mat *);
    218: PETSC_INTERN PetscErrorCode MatSeqAIJKokkosMergeMats(Mat, Mat, MatReuse, Mat *);
    219: PETSC_INTERN PetscErrorCode MatSeqAIJKokkosSyncDevice(Mat);
    220: PETSC_INTERN PetscErrorCode MatSeqAIJKokkosGetKokkosCsrMatrix(Mat, KokkosCsrMatrix *);
    221: PETSC_INTERN PetscErrorCode MatCreateSeqAIJKokkosWithKokkosCsrMatrix(MPI_Comm, KokkosCsrMatrix, Mat *);
    222: PETSC_INTERN PetscErrorCode PrintCsrMatrix(const KokkosCsrMatrix &csrmat);
    223: PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqAIJKokkos(Mat, MatType, MatReuse, Mat *);
    224: PETSC_INTERN PetscErrorCode MatSeqAIJKokkosModifyDevice(Mat);
    225: PETSC_INTERN PetscErrorCode MatSeqAIJKokkosGenerateTranspose_Private(Mat, KokkosCsrMatrix *);

    227: PETSC_INTERN PetscErrorCode MatSeqAIJGetKokkosView(Mat, MatScalarKokkosView *);
    228: PETSC_INTERN PetscErrorCode MatSeqAIJRestoreKokkosView(Mat, MatScalarKokkosView *);
    229: PETSC_INTERN PetscErrorCode MatSeqAIJGetKokkosViewWrite(Mat, MatScalarKokkosView *);
    230: PETSC_INTERN PetscErrorCode MatSeqAIJRestoreKokkosViewWrite(Mat, MatScalarKokkosView *);
    231: PETSC_INTERN PetscErrorCode MatInvertVariableBlockDiagonal_SeqAIJKokkos(Mat, const PetscIntKokkosView &, const PetscIntKokkosView &, const PetscIntKokkosView &, PetscScalarKokkosView &, PetscScalarKokkosView &);
    232: PETSC_INTERN PetscErrorCode MatBindToCPU_SeqAIJKokkos(Mat, PetscBool);