Lines Matching refs:mm
316 MatMPIAIJKokkosReduce - Reduce rows of a MPIAIJKOKKOS matrix (E, in split form) to produce another matrix (F, also in split form, stored in mm)
332 - mm - to stash intermediate data structures for reuse
336 - mm - contains various info, such as garray2[], F (Fd, Fo) etc.
342 static PetscErrorCode MatMPIAIJKokkosReduceBegin(MPI_Comm comm, KokkosCsrMatrix A, KokkosCsrMatrix B, PetscInt cstart, PetscInt cend, const PetscInt *garray1, PetscSF ownerSF, MatReuse reuse, PetscInt *map, MatMatStruct_AtB *mm)
598 mm->sf = reduceSF;
599 mm->leafBuf = MatScalarKokkosView(NoInit("leafBuf"), nleaves);
600 mm->rootBuf = MatScalarKokkosView(NoInit("rootBuf"), nroots);
601 mm->garray = garray2; // give ownership, so no free
602 mm->n = n2;
603 mm->E_NzLeft = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), E_NzLeft_h);
604 mm->Fdjmap = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), Fdjmap_h);
605 mm->Fdjperm = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), Fdjperm_h);
606 mm->Fojmap = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), Fojmap_h);
607 mm->Fojperm = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), Fojperm_h);
617 PetscCallCXX(mm->Fd = KokkosCsrMatrix("Fd", Fm, cend - cstart, Fdnz, Fda_d, Fdi_d, Fdj_d));
618 PetscCallCXX(mm->Fo = KokkosCsrMatrix("Fo", Fm, n2, Fonz, Foa_d, Foi_d, Foj_d)); // Fo's column size is n2, length of garray2[]
625 mm->E_TeamSize = teamSize;
626 mm->E_VectorLength = vectorLength;
627 mm->E_RowsPerTeam = rowsPerTeam;
635 const auto &E_NzLeft = mm->E_NzLeft;
636 auto &leafBuf = mm->leafBuf;
637 auto &rootBuf = mm->rootBuf;
638 PetscSF reduceSF = mm->sf;
640 PetscInt teamSize = mm->E_TeamSize;
641 PetscInt vectorLength = mm->E_VectorLength;
642 PetscInt rowsPerTeam = mm->E_RowsPerTeam;
674 static PetscErrorCode MatMPIAIJKokkosReduceEnd(MPI_Comm comm, KokkosCsrMatrix A, KokkosCsrMatrix B, PetscInt cstart, PetscInt cend, const PetscInt *garray1, PetscSF ownerSF, MatReuse reuse, PetscInt *map, MatMatStruct_AtB *mm)
676 auto &leafBuf = mm->leafBuf;
677 auto &rootBuf = mm->rootBuf;
678 auto &Fda = mm->Fd.values;
679 const auto &Fdjmap = mm->Fdjmap;
680 const auto &Fdjperm = mm->Fdjperm;
681 auto Fdnz = mm->Fd.nnz();
682 auto &Foa = mm->Fo.values;
683 const auto &Fojmap = mm->Fojmap;
684 const auto &Fojperm = mm->Fojperm;
685 auto Fonz = mm->Fo.nnz();
686 PetscSF reduceSF = mm->sf;
709 MatMPIAIJKokkosBcast - Bcast local rows of a MPIAIJKOKKOS matrix (E) to produce a local matrix (F, stored in mm) in split form
729 - mm - to stash matproduct intermediate data structures
733 - mm - contains various info, such as garray2[], Fd, Fo, etc.
739 static PetscErrorCode MatMPIAIJKokkosBcastBegin(Mat E, PetscSF ownerSF, MatReuse reuse, PetscInt *map, MatMatStruct_AB *mm)
913 mm->sf = bcastSF;
914 mm->E_NzLeft = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), E_NzLeft_h);
915 mm->F_NzLeft = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), F_NzLeft_h);
916 mm->irootloc = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), irootloc_h);
917 mm->rowoffset = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), rowoffset_h);
918 mm->rootBuf = MatScalarKokkosView(NoInit("rootBuf"), nroots);
919 mm->leafBuf = MatScalarKokkosView(NoInit("leafBuf"), nleaves);
920 mm->garray = garray2;
921 mm->n = n2;
930 PetscCallCXX(mm->Fd = KokkosCsrMatrix("Fd", Fm, cend - cstart, Fdnz, Fda_d, Fdi_d, Fdj_d));
931 PetscCallCXX(mm->Fo = KokkosCsrMatrix("Fo", Fm, n2, Fonz, Foa_d, Foi_d, Foj_d));
937 PetscCall(MatMergeGetLaunchParameters<DefaultExecutionSpace>(mm->irootloc.extent(0), mm->rootBuf.extent(0), -1, teamSize, vectorLength, rowsPerTeam));
938 mm->E_TeamSize = teamSize;
939 mm->E_VectorLength = vectorLength;
940 mm->E_RowsPerTeam = rowsPerTeam;
944 mm->F_TeamSize = teamSize;
945 mm->F_VectorLength = vectorLength;
946 mm->F_RowsPerTeam = rowsPerTeam;
960 PetscIntKokkosView &E_NzLeft = mm->E_NzLeft;
961 PetscSF &bcastSF = mm->sf;
962 MatScalarKokkosView &rootBuf = mm->rootBuf;
963 MatScalarKokkosView &leafBuf = mm->leafBuf;
964 PetscIntKokkosView &irootloc = mm->irootloc;
965 PetscIntKokkosView &rowoffset = mm->rowoffset;
967 PetscInt teamSize = mm->E_TeamSize;
968 PetscInt vectorLength = mm->E_VectorLength;
969 PetscInt rowsPerTeam = mm->E_RowsPerTeam;
1001 static PetscErrorCode MatMPIAIJKokkosBcastEnd(Mat E, PetscSF ownerSF, MatReuse reuse, PetscInt *map, MatMatStruct_AB *mm)
1004 const auto &Fd = mm->Fd;
1005 const auto &Fo = mm->Fo;
1012 PetscIntKokkosView &F_NzLeft = mm->F_NzLeft;
1013 PetscSF &bcastSF = mm->sf;
1014 MatScalarKokkosView &rootBuf = mm->rootBuf;
1015 MatScalarKokkosView &leafBuf = mm->leafBuf;
1016 PetscInt teamSize = mm->F_TeamSize;
1017 PetscInt vectorLength = mm->F_VectorLength;
1018 PetscInt rowsPerTeam = mm->F_RowsPerTeam;
1050 static PetscErrorCode MatProductSymbolic_MPIAIJKokkos_AtB(Mat_Product *product, Mat A, Mat B, MatMatStruct_AtB *mm)
1077 PetscCallCXX(mm->kh1.create_spgemm_handle(spgemm_alg));
1078 PetscCallCXX(mm->kh2.create_spgemm_handle(spgemm_alg));
1079 PetscCallCXX(mm->kh3.create_spgemm_handle(spgemm_alg));
1080 PetscCallCXX(mm->kh4.create_spgemm_handle(spgemm_alg));
1083 PetscCallCXX(KokkosSparse::spgemm_symbolic(mm->kh3, Aot, false, Bd, false, mm->C3));
1084 PetscCallCXX(KokkosSparse::spgemm_symbolic(mm->kh4, Aot, false, Bo, false, mm->C4));
1087 PetscCallCXX(KokkosSparse::spgemm_numeric(mm->kh3, Aot, false, Bd, false, mm->C3));
1088 PetscCallCXX(KokkosSparse::spgemm_numeric(mm->kh4, Aot, false, Bo, false, mm->C4));
1091 PetscCallCXX(sort_crs_matrix(mm->C3));
1092 PetscCallCXX(sort_crs_matrix(mm->C4));
1098 PetscCall(MatMPIAIJKokkosReduceBegin(comm, mm->C3, mm->C4, cstart, cend, bmpi->garray, ampi->Mvctx, MAT_INITIAL_MATRIX, map_h.data(), mm));
1101 PetscCallCXX(KokkosSparse::spgemm_symbolic(mm->kh1, Adt, false, Bd, false, mm->C1));
1102 PetscCallCXX(KokkosSparse::spgemm_symbolic(mm->kh2, Adt, false, Bo, false, mm->C2_mid));
1103 PetscCallCXX(KokkosSparse::spgemm_numeric(mm->kh1, Adt, false, Bd, false, mm->C1));
1104 PetscCallCXX(KokkosSparse::spgemm_numeric(mm->kh2, Adt, false, Bo, false, mm->C2_mid));
1106 PetscCallCXX(sort_crs_matrix(mm->C1));
1107 PetscCallCXX(sort_crs_matrix(mm->C2_mid));
1110 PetscCall(MatMPIAIJKokkosReduceEnd(comm, mm->C3, mm->C4, cstart, cend, bmpi->garray, ampi->Mvctx, MAT_INITIAL_MATRIX, map_h.data(), mm));
1113 MatColIdxKokkosView oldj = mm->C2_mid.graph.entries, newj(NoInit("j"), oldj.extent(0));
1116 PetscCallCXX(mm->C2 = KokkosCsrMatrix("C2", mm->C2_mid.numRows(), mm->n /*new column size*/, mm->C2_mid.nnz(), mm->C2_mid.values, mm->C2_mid.graph.row_map, newj));
1119 PetscCallCXX(mm->kh1.create_spadd_handle(true)); // C1, Fd are sorted
1120 PetscCallCXX(mm->kh2.create_spadd_handle(true)); // C2, Fo are sorted
1121 PetscCallCXX(KokkosSparse::spadd_symbolic(&mm->kh1, mm->C1, mm->Fd, mm->Cd));
1122 PetscCallCXX(KokkosSparse::spadd_symbolic(&mm->kh2, mm->C2, mm->Fo, mm->Co));
1123 PetscCallCXX(KokkosSparse::spadd_numeric(&mm->kh1, 1.0, mm->C1, 1.0, mm->Fd, mm->Cd));
1124 PetscCallCXX(KokkosSparse::spadd_numeric(&mm->kh2, 1.0, mm->C2, 1.0, mm->Fo, mm->Co));
1128 static PetscErrorCode MatProductNumeric_MPIAIJKokkos_AtB(Mat_Product *product, Mat A, Mat B, MatMatStruct_AtB *mm)
1143 PetscCallCXX(KokkosSparse::spgemm_numeric(mm->kh3, Aot, false, Bd, false, mm->C3));
1144 PetscCallCXX(KokkosSparse::spgemm_numeric(mm->kh4, Aot, false, Bo, false, mm->C4));
1147 PetscCall(MatMPIAIJKokkosReduceBegin(comm, mm->C3, mm->C4, 0, 0, NULL, NULL, MAT_REUSE_MATRIX, NULL, mm));
1150 PetscCallCXX(KokkosSparse::spgemm_numeric(mm->kh1, Adt, false, Bd, false, mm->C1));
1151 PetscCallCXX(KokkosSparse::spgemm_numeric(mm->kh2, Adt, false, Bo, false, mm->C2_mid));
1153 PetscCall(MatMPIAIJKokkosReduceEnd(comm, mm->C3, mm->C4, 0, 0, NULL, NULL, MAT_REUSE_MATRIX, NULL, mm));
1156 PetscCallCXX(KokkosSparse::spadd_numeric(&mm->kh1, 1.0, mm->C1, 1.0, mm->Fd, mm->Cd));
1157 PetscCallCXX(KokkosSparse::spadd_numeric(&mm->kh2, 1.0, mm->C2, 1.0, mm->Fo, mm->Co));
1167 - mm - a struct used to stash intermediate data when computing AB. Persist from symbolic to numeric operations.
1169 static PetscErrorCode MatProductSymbolic_MPIAIJKokkos_AB(Mat_Product *product, Mat A, Mat B, MatMatStruct_AB *mm)
1191 mm->kh1.create_spgemm_handle(spgemm_alg);
1192 mm->kh2.create_spgemm_handle(spgemm_alg);
1193 mm->kh3.create_spgemm_handle(spgemm_alg);
1194 mm->kh4.create_spgemm_handle(spgemm_alg);
1198 PetscCall(MatMPIAIJKokkosBcastBegin(B, ampi->Mvctx, MAT_INITIAL_MATRIX, map_h.data(), mm));
1201 PetscCallCXX(KokkosSparse::spgemm_symbolic(mm->kh1, Ad, false, Bd, false, mm->C1));
1202 PetscCallCXX(KokkosSparse::spgemm_symbolic(mm->kh2, Ad, false, Bo, false, mm->C2_mid)); // C2 aliases with C2_mid, except with new column indices
1205 PetscCallCXX(KokkosSparse::spgemm_numeric(mm->kh1, Ad, false, Bd, false, mm->C1));
1206 PetscCallCXX(KokkosSparse::spgemm_numeric(mm->kh2, Ad, false, Bo, false, mm->C2_mid));
1208 PetscCallCXX(sort_crs_matrix(mm->C1));
1209 PetscCallCXX(sort_crs_matrix(mm->C2_mid));
1212 PetscCall(MatMPIAIJKokkosBcastEnd(B, ampi->Mvctx, MAT_INITIAL_MATRIX, map_h.data(), mm));
1215 PetscCallCXX(KokkosSparse::spgemm_symbolic(mm->kh3, Ao, false, mm->Fd, false, mm->C3));
1216 PetscCallCXX(KokkosSparse::spgemm_numeric(mm->kh3, Ao, false, mm->Fd, false, mm->C3));
1217 PetscCallCXX(KokkosSparse::spgemm_symbolic(mm->kh4, Ao, false, mm->Fo, false, mm->C4));
1218 PetscCallCXX(KokkosSparse::spgemm_numeric(mm->kh4, Ao, false, mm->Fo, false, mm->C4));
1220 PetscCallCXX(sort_crs_matrix(mm->C3));
1221 PetscCallCXX(sort_crs_matrix(mm->C4));
1225 MatColIdxKokkosView oldj = mm->C2_mid.graph.entries, newj(NoInit("j"), oldj.extent(0));
1228 mm->C2 = KokkosCsrMatrix("C2", mm->C2_mid.numRows(), mm->n /*new column size*/, mm->C2_mid.nnz(), mm->C2_mid.values, mm->C2_mid.graph.row_map, newj);
1231 mm->kh1.create_spadd_handle(true); // C1, C3 are sorted
1232 mm->kh2.create_spadd_handle(true); // C2, C4 are sorted
1233 PetscCallCXX(KokkosSparse::spadd_symbolic(&mm->kh1, mm->C1, mm->C3, mm->Cd));
1234 PetscCallCXX(KokkosSparse::spadd_symbolic(&mm->kh2, mm->C2, mm->C4, mm->Co));
1235 PetscCallCXX(KokkosSparse::spadd_numeric(&mm->kh1, 1.0, mm->C1, 1.0, mm->C3, mm->Cd));
1236 PetscCallCXX(KokkosSparse::spadd_numeric(&mm->kh2, 1.0, mm->C2, 1.0, mm->C4, mm->Co));
1240 static PetscErrorCode MatProductNumeric_MPIAIJKokkos_AB(Mat_Product *product, Mat A, Mat B, MatMatStruct_AB *mm)
1253 PetscCall(MatMPIAIJKokkosBcastBegin(B, NULL, MAT_REUSE_MATRIX, NULL, mm));
1256 PetscCallCXX(KokkosSparse::spgemm_numeric(mm->kh1, Ad, false, Bd, false, mm->C1));
1257 PetscCallCXX(KokkosSparse::spgemm_numeric(mm->kh2, Ad, false, Bo, false, mm->C2_mid));
1259 PetscCall(MatMPIAIJKokkosBcastEnd(B, NULL, MAT_REUSE_MATRIX, NULL, mm));
1262 PetscCallCXX(KokkosSparse::spgemm_numeric(mm->kh3, Ao, false, mm->Fd, false, mm->C3));
1263 PetscCallCXX(KokkosSparse::spgemm_numeric(mm->kh4, Ao, false, mm->Fo, false, mm->C4));
1266 PetscCallCXX(KokkosSparse::spadd_numeric(&mm->kh1, 1.0, mm->C1, 1.0, mm->C3, mm->Cd));
1267 PetscCallCXX(KokkosSparse::spadd_numeric(&mm->kh2, 1.0, mm->C2, 1.0, mm->C4, mm->Co));
1314 MatMatStruct *mm = NULL;
1363 mm = pdata->mmAB = mmAB;
1367 mm = pdata->mmAtB = mmAtB;
1387 mm = pdata->mmAtB = mmAtB;
1390 PetscCall(MatCreateSeqAIJKokkosWithKokkosCsrMatrix(PETSC_COMM_SELF, mm->Cd, &Cd));
1391 PetscCall(MatCreateSeqAIJKokkosWithKokkosCsrMatrix(PETSC_COMM_SELF, mm->Co, &Co));
1396 mpiaij->garray = mm->garray;