Actual source code: mpimatmatmult.c
1: /*
2: Defines matrix-matrix product routines for pairs of MPIAIJ matrices
3: C = A * B
4: */
5: #include <../src/mat/impls/aij/seq/aij.h>
6: #include <../src/mat/utils/freespace.h>
7: #include <../src/mat/impls/aij/mpi/mpiaij.h>
8: #include <petscbt.h>
9: #include <../src/mat/impls/dense/mpi/mpidense.h>
10: #include <petsc/private/vecimpl.h>
11: #include <petsc/private/sfimpl.h>
13: #if PetscDefined(HAVE_HYPRE)
14: PETSC_INTERN PetscErrorCode MatMatMultSymbolic_AIJ_AIJ_wHYPRE(Mat, Mat, PetscReal, Mat);
15: #endif
17: PETSC_INTERN PetscErrorCode MatProductSymbolic_ABt_MPIAIJ_MPIAIJ(Mat C)
18: {
19: Mat_Product *product = C->product;
20: Mat B = product->B;
22: PetscFunctionBegin;
23: PetscCall(MatTranspose(B, MAT_INITIAL_MATRIX, &product->B));
24: PetscCall(MatDestroy(&B));
25: PetscCall(MatProductSymbolic_AB_MPIAIJ_MPIAIJ(C));
26: PetscFunctionReturn(PETSC_SUCCESS);
27: }
29: PETSC_INTERN PetscErrorCode MatProductSymbolic_AB_MPIAIJ_MPIAIJ(Mat C)
30: {
31: Mat_Product *product = C->product;
32: Mat A = product->A, B = product->B;
33: MatProductAlgorithm alg = product->alg;
34: PetscReal fill = product->fill;
35: PetscBool flg;
37: PetscFunctionBegin;
38: /* scalable */
39: PetscCall(PetscStrcmp(alg, "scalable", &flg));
40: if (flg) {
41: PetscCall(MatMatMultSymbolic_MPIAIJ_MPIAIJ(A, B, fill, C));
42: PetscFunctionReturn(PETSC_SUCCESS);
43: }
45: /* nonscalable */
46: PetscCall(PetscStrcmp(alg, "nonscalable", &flg));
47: if (flg) {
48: PetscCall(MatMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(A, B, fill, C));
49: PetscFunctionReturn(PETSC_SUCCESS);
50: }
52: /* seqmpi */
53: PetscCall(PetscStrcmp(alg, "seqmpi", &flg));
54: if (flg) {
55: PetscCall(MatMatMultSymbolic_MPIAIJ_MPIAIJ_seqMPI(A, B, fill, C));
56: PetscFunctionReturn(PETSC_SUCCESS);
57: }
59: /* backend general code */
60: PetscCall(PetscStrcmp(alg, "backend", &flg));
61: if (flg) {
62: PetscCall(MatProductSymbolic_MPIAIJBACKEND(C));
63: PetscFunctionReturn(PETSC_SUCCESS);
64: }
66: #if PetscDefined(HAVE_HYPRE)
67: PetscCall(PetscStrcmp(alg, "hypre", &flg));
68: if (flg) {
69: PetscCall(MatMatMultSymbolic_AIJ_AIJ_wHYPRE(A, B, fill, C));
70: PetscFunctionReturn(PETSC_SUCCESS);
71: }
72: #endif
73: SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_SUP, "Mat Product Algorithm is not supported");
74: }
76: PetscErrorCode MatProductCtxDestroy_MPIAIJ_MatMatMult(PetscCtxRt data)
77: {
78: MatProductCtx_APMPI *ptap = *(MatProductCtx_APMPI **)data;
80: PetscFunctionBegin;
81: PetscCall(PetscFree2(ptap->startsj_s, ptap->startsj_r));
82: PetscCall(PetscFree(ptap->bufa));
83: PetscCall(MatDestroy(&ptap->P_loc));
84: PetscCall(MatDestroy(&ptap->P_oth));
85: PetscCall(MatDestroy(&ptap->Pt));
86: PetscCall(PetscFree(ptap->api));
87: PetscCall(PetscFree(ptap->apj));
88: PetscCall(PetscFree(ptap->apa));
89: PetscCall(PetscFree(ptap));
90: PetscFunctionReturn(PETSC_SUCCESS);
91: }
93: PetscErrorCode MatMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable(Mat A, Mat P, Mat C)
94: {
95: Mat_MPIAIJ *a = (Mat_MPIAIJ *)A->data, *c = (Mat_MPIAIJ *)C->data;
96: Mat_SeqAIJ *ad = (Mat_SeqAIJ *)a->A->data, *ao = (Mat_SeqAIJ *)a->B->data;
97: Mat_SeqAIJ *cd = (Mat_SeqAIJ *)c->A->data, *co = (Mat_SeqAIJ *)c->B->data;
98: PetscScalar *cda, *coa;
99: Mat_SeqAIJ *p_loc, *p_oth;
100: PetscScalar *apa, *ca;
101: PetscInt cm = C->rmap->n;
102: MatProductCtx_APMPI *ptap;
103: PetscInt *api, *apj, *apJ, i, k;
104: PetscInt cstart = C->cmap->rstart;
105: PetscInt cdnz, conz, k0, k1;
106: const PetscScalar *dummy1, *dummy2, *dummy3, *dummy4;
107: MPI_Comm comm;
108: PetscMPIInt size;
110: PetscFunctionBegin;
111: MatCheckProduct(C, 3);
112: ptap = (MatProductCtx_APMPI *)C->product->data;
113: PetscCheck(ptap, PetscObjectComm((PetscObject)C), PETSC_ERR_ARG_WRONGSTATE, "PtAP cannot be computed. Missing data");
114: PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
115: PetscCallMPI(MPI_Comm_size(comm, &size));
116: PetscCheck(ptap->P_oth || size <= 1, PetscObjectComm((PetscObject)C), PETSC_ERR_ARG_WRONGSTATE, "AP cannot be reused. Do not call MatProductClear()");
118: /* flag CPU mask for C */
119: #if PetscDefined(HAVE_DEVICE)
120: if (C->offloadmask != PETSC_OFFLOAD_UNALLOCATED) C->offloadmask = PETSC_OFFLOAD_CPU;
121: if (c->A->offloadmask != PETSC_OFFLOAD_UNALLOCATED) c->A->offloadmask = PETSC_OFFLOAD_CPU;
122: if (c->B->offloadmask != PETSC_OFFLOAD_UNALLOCATED) c->B->offloadmask = PETSC_OFFLOAD_CPU;
123: #endif
125: /* 1) get P_oth = ptap->P_oth and P_loc = ptap->P_loc */
126: /* update numerical values of P_oth and P_loc */
127: PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A, P, MAT_REUSE_MATRIX, &ptap->startsj_s, &ptap->startsj_r, &ptap->bufa, &ptap->P_oth));
128: PetscCall(MatMPIAIJGetLocalMat(P, MAT_REUSE_MATRIX, &ptap->P_loc));
130: /* 2) compute numeric C_loc = A_loc*P = Ad*P_loc + Ao*P_oth */
131: /* get data from symbolic products */
132: p_loc = (Mat_SeqAIJ *)ptap->P_loc->data;
133: p_oth = NULL;
134: if (size > 1) p_oth = (Mat_SeqAIJ *)ptap->P_oth->data;
136: /* get apa for storing dense row A[i,:]*P */
137: apa = ptap->apa;
139: api = ptap->api;
140: apj = ptap->apj;
141: /* trigger copy to CPU */
142: PetscCall(MatSeqAIJGetArrayRead(a->A, &dummy1));
143: PetscCall(MatSeqAIJGetArrayRead(a->B, &dummy2));
144: PetscCall(MatSeqAIJGetArrayRead(ptap->P_loc, &dummy3));
145: if (ptap->P_oth) PetscCall(MatSeqAIJGetArrayRead(ptap->P_oth, &dummy4));
146: PetscCall(MatSeqAIJGetArrayWrite(c->A, &cda));
147: PetscCall(MatSeqAIJGetArrayWrite(c->B, &coa));
148: for (i = 0; i < cm; i++) {
149: /* compute apa = A[i,:]*P */
150: AProw_nonscalable(i, ad, ao, p_loc, p_oth, apa);
152: /* set values in C */
153: apJ = PetscSafePointerPlusOffset(apj, api[i]);
154: cdnz = cd->i[i + 1] - cd->i[i];
155: conz = co->i[i + 1] - co->i[i];
157: /* 1st off-diagonal part of C */
158: ca = PetscSafePointerPlusOffset(coa, co->i[i]);
159: k = 0;
160: for (k0 = 0; k0 < conz; k0++) {
161: if (apJ[k] >= cstart) break;
162: ca[k0] = apa[apJ[k]];
163: apa[apJ[k++]] = 0.0;
164: }
166: /* diagonal part of C */
167: ca = PetscSafePointerPlusOffset(cda, cd->i[i]);
168: for (k1 = 0; k1 < cdnz; k1++) {
169: ca[k1] = apa[apJ[k]];
170: apa[apJ[k++]] = 0.0;
171: }
173: /* 2nd off-diagonal part of C */
174: ca = PetscSafePointerPlusOffset(coa, co->i[i]);
175: for (; k0 < conz; k0++) {
176: ca[k0] = apa[apJ[k]];
177: apa[apJ[k++]] = 0.0;
178: }
179: }
180: PetscCall(MatSeqAIJRestoreArrayRead(a->A, &dummy1));
181: PetscCall(MatSeqAIJRestoreArrayRead(a->B, &dummy2));
182: PetscCall(MatSeqAIJRestoreArrayRead(ptap->P_loc, &dummy3));
183: if (ptap->P_oth) PetscCall(MatSeqAIJRestoreArrayRead(ptap->P_oth, &dummy4));
184: PetscCall(MatSeqAIJRestoreArrayWrite(c->A, &cda));
185: PetscCall(MatSeqAIJRestoreArrayWrite(c->B, &coa));
187: PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
188: PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
189: PetscFunctionReturn(PETSC_SUCCESS);
190: }
192: PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(Mat A, Mat P, PetscReal fill, Mat C)
193: {
194: MPI_Comm comm;
195: PetscMPIInt size;
196: MatProductCtx_APMPI *ptap;
197: PetscFreeSpaceList free_space = NULL, current_space = NULL;
198: Mat_MPIAIJ *a = (Mat_MPIAIJ *)A->data;
199: Mat_SeqAIJ *ad = (Mat_SeqAIJ *)a->A->data, *ao = (Mat_SeqAIJ *)a->B->data, *p_loc, *p_oth;
200: PetscInt *pi_loc, *pj_loc, *pi_oth, *pj_oth, *dnz, *onz;
201: PetscInt *adi = ad->i, *adj = ad->j, *aoi = ao->i, *aoj = ao->j, rstart = A->rmap->rstart;
202: PetscInt *lnk, i, pnz, row, *api, *apj, *Jptr, apnz, nspacedouble = 0, j, nzi;
203: PetscInt am = A->rmap->n, pN = P->cmap->N, pn = P->cmap->n, pm = P->rmap->n;
204: PetscBT lnkbt;
205: PetscReal afill;
206: MatType mtype;
208: PetscFunctionBegin;
209: MatCheckProduct(C, 4);
210: PetscCheck(!C->product->data, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Extra product struct not empty");
211: PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
212: PetscCallMPI(MPI_Comm_size(comm, &size));
214: /* create struct MatProductCtx_APMPI and attached it to C later */
215: PetscCall(PetscNew(&ptap));
217: /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */
218: PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A, P, MAT_INITIAL_MATRIX, &ptap->startsj_s, &ptap->startsj_r, &ptap->bufa, &ptap->P_oth));
220: /* get P_loc by taking all local rows of P */
221: PetscCall(MatMPIAIJGetLocalMat(P, MAT_INITIAL_MATRIX, &ptap->P_loc));
223: p_loc = (Mat_SeqAIJ *)ptap->P_loc->data;
224: pi_loc = p_loc->i;
225: pj_loc = p_loc->j;
226: if (size > 1) {
227: p_oth = (Mat_SeqAIJ *)ptap->P_oth->data;
228: pi_oth = p_oth->i;
229: pj_oth = p_oth->j;
230: } else {
231: p_oth = NULL;
232: pi_oth = NULL;
233: pj_oth = NULL;
234: }
236: /* first, compute symbolic AP = A_loc*P = A_diag*P_loc + A_off*P_oth */
237: PetscCall(PetscMalloc1(am + 1, &api));
238: ptap->api = api;
239: api[0] = 0;
241: /* create and initialize a linked list */
242: PetscCall(PetscLLCondensedCreate(pN, pN, &lnk, &lnkbt));
244: /* Initial FreeSpace size is fill*(nnz(A)+nnz(P)) */
245: PetscCall(PetscFreeSpaceGet(PetscRealIntMultTruncate(fill, PetscIntSumTruncate(adi[am], PetscIntSumTruncate(aoi[am], pi_loc[pm]))), &free_space));
246: current_space = free_space;
248: MatPreallocateBegin(comm, am, pn, dnz, onz);
249: for (i = 0; i < am; i++) {
250: /* diagonal portion of A */
251: nzi = adi[i + 1] - adi[i];
252: for (j = 0; j < nzi; j++) {
253: row = *adj++;
254: pnz = pi_loc[row + 1] - pi_loc[row];
255: Jptr = pj_loc + pi_loc[row];
256: /* add non-zero cols of P into the sorted linked list lnk */
257: PetscCall(PetscLLCondensedAddSorted(pnz, Jptr, lnk, lnkbt));
258: }
259: /* off-diagonal portion of A */
260: nzi = aoi[i + 1] - aoi[i];
261: for (j = 0; j < nzi; j++) {
262: row = *aoj++;
263: pnz = pi_oth[row + 1] - pi_oth[row];
264: Jptr = pj_oth + pi_oth[row];
265: PetscCall(PetscLLCondensedAddSorted(pnz, Jptr, lnk, lnkbt));
266: }
267: /* add possible missing diagonal entry */
268: if (C->force_diagonals) {
269: j = i + rstart; /* column index */
270: PetscCall(PetscLLCondensedAddSorted(1, &j, lnk, lnkbt));
271: }
273: apnz = lnk[0];
274: api[i + 1] = api[i] + apnz;
276: /* if free space is not available, double the total space in the list */
277: if (current_space->local_remaining < apnz) {
278: PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(apnz, current_space->total_array_size), ¤t_space));
279: nspacedouble++;
280: }
282: /* Copy data into free space, then initialize lnk */
283: PetscCall(PetscLLCondensedClean(pN, apnz, current_space->array, lnk, lnkbt));
284: PetscCall(MatPreallocateSet(i + rstart, apnz, current_space->array, dnz, onz));
286: current_space->array += apnz;
287: current_space->local_used += apnz;
288: current_space->local_remaining -= apnz;
289: }
291: /* Allocate space for apj, initialize apj, and */
292: /* destroy list of free space and other temporary array(s) */
293: PetscCall(PetscMalloc1(api[am], &ptap->apj));
294: apj = ptap->apj;
295: PetscCall(PetscFreeSpaceContiguous(&free_space, ptap->apj));
296: PetscCall(PetscLLDestroy(lnk, lnkbt));
298: /* malloc apa to store dense row A[i,:]*P */
299: PetscCall(PetscCalloc1(pN, &ptap->apa));
301: /* set and assemble symbolic parallel matrix C */
302: PetscCall(MatSetSizes(C, am, pn, PETSC_DETERMINE, PETSC_DETERMINE));
303: PetscCall(MatSetBlockSizesFromMats(C, A, P));
305: PetscCall(MatGetType(A, &mtype));
306: PetscCall(MatSetType(C, mtype));
307: PetscCall(MatMPIAIJSetPreallocation(C, 0, dnz, 0, onz));
308: MatPreallocateEnd(dnz, onz);
310: PetscCall(MatSetValues_MPIAIJ_CopyFromCSRFormat_Symbolic(C, apj, api));
311: PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
312: PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
313: PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
314: PetscCall(MatSetOption(C, MAT_NEW_NONZERO_LOCATION_ERR, PETSC_TRUE));
316: C->ops->matmultnumeric = MatMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable;
317: C->ops->productnumeric = MatProductNumeric_AB;
319: /* attach the supporting struct to C for reuse */
320: C->product->data = ptap;
321: C->product->destroy = MatProductCtxDestroy_MPIAIJ_MatMatMult;
323: /* set MatInfo */
324: afill = (PetscReal)api[am] / (adi[am] + aoi[am] + pi_loc[pm] + 1) + 1.e-5;
325: if (afill < 1.0) afill = 1.0;
326: C->info.mallocs = nspacedouble;
327: C->info.fill_ratio_given = fill;
328: C->info.fill_ratio_needed = afill;
330: if (PetscDefined(USE_INFO)) {
331: if (api[am]) {
332: PetscCall(PetscInfo(C, "Reallocs %" PetscInt_FMT "; Fill ratio: given %g needed %g.\n", nspacedouble, (double)fill, (double)afill));
333: PetscCall(PetscInfo(C, "Use MatMatMult(A,B,MatReuse,%g,&C) for best performance.;\n", (double)afill));
334: } else PetscCall(PetscInfo(C, "Empty matrix product\n"));
335: }
336: PetscFunctionReturn(PETSC_SUCCESS);
337: }
339: static PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIDense(Mat, Mat, PetscReal, Mat);
340: static PetscErrorCode MatMatMultNumeric_MPIAIJ_MPIDense(Mat, Mat, Mat);
342: static PetscErrorCode MatProductSetFromOptions_MPIAIJ_MPIDense_AB(Mat C)
343: {
344: Mat_Product *product = C->product;
345: Mat A = product->A, B = product->B;
347: PetscFunctionBegin;
348: if (A->cmap->rstart != B->rmap->rstart || A->cmap->rend != B->rmap->rend)
349: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Matrix local dimensions are incompatible, (%" PetscInt_FMT ", %" PetscInt_FMT ") != (%" PetscInt_FMT ",%" PetscInt_FMT ")", A->cmap->rstart, A->cmap->rend, B->rmap->rstart, B->rmap->rend);
351: C->ops->matmultsymbolic = MatMatMultSymbolic_MPIAIJ_MPIDense;
352: C->ops->productsymbolic = MatProductSymbolic_AB;
353: PetscFunctionReturn(PETSC_SUCCESS);
354: }
356: static PetscErrorCode MatProductSetFromOptions_MPIAIJ_MPIDense_AtB(Mat C)
357: {
358: Mat_Product *product = C->product;
359: Mat A = product->A, B = product->B;
361: PetscFunctionBegin;
362: if (A->rmap->rstart != B->rmap->rstart || A->rmap->rend != B->rmap->rend)
363: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Matrix local dimensions are incompatible, (%" PetscInt_FMT ", %" PetscInt_FMT ") != (%" PetscInt_FMT ",%" PetscInt_FMT ")", A->rmap->rstart, A->rmap->rend, B->rmap->rstart, B->rmap->rend);
365: C->ops->transposematmultsymbolic = MatTransposeMatMultSymbolic_MPIAIJ_MPIDense;
366: C->ops->productsymbolic = MatProductSymbolic_AtB;
367: PetscFunctionReturn(PETSC_SUCCESS);
368: }
370: PETSC_INTERN PetscErrorCode MatProductSetFromOptions_MPIAIJ_MPIDense(Mat C)
371: {
372: Mat_Product *product = C->product;
374: PetscFunctionBegin;
375: switch (product->type) {
376: case MATPRODUCT_AB:
377: PetscCall(MatProductSetFromOptions_MPIAIJ_MPIDense_AB(C));
378: break;
379: case MATPRODUCT_AtB:
380: PetscCall(MatProductSetFromOptions_MPIAIJ_MPIDense_AtB(C));
381: break;
382: default:
383: break;
384: }
385: PetscFunctionReturn(PETSC_SUCCESS);
386: }
388: PETSC_INTERN PetscErrorCode MatMPIDenseScatterDestroy_Private(MPIAIJ_MPIDense *contents)
389: {
390: PetscFunctionBegin;
391: PetscCall(MatDestroy(&contents->workB));
392: for (PetscInt i = 0; i < contents->nsends; i++) PetscCallMPI(MPI_Type_free(&contents->stype[i]));
393: for (PetscInt i = 0; i < contents->nrecvs; i++) PetscCallMPI(MPI_Type_free(&contents->rtype[i]));
394: PetscCall(PetscFree4(contents->stype, contents->rtype, contents->rwaits, contents->swaits));
395: PetscFunctionReturn(PETSC_SUCCESS);
396: }
398: static PetscErrorCode MatMPIAIJ_MPIDenseDestroy(PetscCtxRt ctx)
399: {
400: MPIAIJ_MPIDense *contents = *(MPIAIJ_MPIDense **)ctx;
402: PetscFunctionBegin;
403: PetscCall(MatMPIDenseScatterDestroy_Private(contents));
404: PetscCall(PetscFree(contents));
405: PetscFunctionReturn(PETSC_SUCCESS);
406: }
408: PETSC_INTERN PetscErrorCode MatMPIDenseScatterSetUp_Private(VecScatter ctx, PetscInt nrows, PetscInt bs, PetscInt Am, Mat B, Mat C, MPIAIJ_MPIDense *contents, PetscInt *batchSize, PetscInt *numBatches)
409: {
410: PetscInt Bm = B->rmap->n, BN = B->cmap->N, Bbn, Bbs, numBb;
411: MPI_Comm comm;
412: MPI_Datatype type1;
413: const PetscInt *sindices, *sstarts, *rstarts;
414: PetscMPIInt *disp;
415: PetscMPIInt nsends, nrecvs, nrows_to, nrows_from, bs_mpi;
417: PetscFunctionBegin;
418: PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
419: PetscCall(MatDenseGetLDA(B, &contents->blda));
420: PetscCall(VecScatterGetRemote_Private(ctx, PETSC_TRUE, &nsends, &sstarts, &sindices, NULL, NULL));
421: PetscCall(VecScatterGetRemoteOrdered_Private(ctx, PETSC_FALSE, &nrecvs, &rstarts, NULL, NULL, NULL));
423: /* Create column block of B and C for memory scalability when BN is too large */
424: /* Estimate Bbn, column size of Bb */
425: if (nrows) {
426: Bbn = 2 * Am * BN / nrows;
427: if (!Bbn) Bbn = 1;
428: } else Bbn = BN;
429: Bbs = B->cmap->bs;
430: Bbn = Bbn / Bbs * Bbs;
431: if (Bbn > BN) Bbn = BN;
432: PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &Bbn, 1, MPIU_INT, MPI_MAX, comm));
434: /* Enable runtime option for Bbn */
435: PetscOptionsBegin(comm, ((PetscObject)C)->prefix, "MatProduct", "Mat");
436: PetscCall(PetscOptionsDeprecated("-matmatmult_Bbn", "-matproduct_batch_size", "3.25", NULL));
437: PetscCall(PetscOptionsBoundedInt("-matproduct_batch_size", "Number of dense columns per batch", "MatProduct", Bbn, &Bbn, NULL, 0));
438: PetscOptionsEnd();
439: Bbn = PetscMin(Bbn, BN);
441: if (Bbn > 0 && Bbn < BN) numBb = BN / Bbn;
442: else numBb = 0;
443: if (numBb) PetscCall(PetscInfo(C, "Using column batches of size %" PetscInt_FMT " for %" PetscInt_FMT " dense columns\n", Bbn, BN));
444: /* Create work matrix used to store off processor rows of B needed for local product */
445: PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, nrows, Bbn ? Bbn : BN, NULL, &contents->workB));
447: /* Use MPI derived data type to reduce memory required by the send/recv buffers */
448: PetscCall(PetscMalloc4(nsends, &contents->stype, nrecvs, &contents->rtype, nrecvs, &contents->rwaits, nsends, &contents->swaits));
449: contents->nsends = nsends;
450: contents->nrecvs = nrecvs;
452: PetscCall(PetscMalloc1(PetscMax(Bm, 1), &disp));
453: PetscCall(PetscMPIIntCast(bs, &bs_mpi));
454: for (PetscMPIInt i = 0; i < nsends; i++) {
455: PetscCall(PetscMPIIntCast(sstarts[i + 1] - sstarts[i], &nrows_to));
456: for (PetscInt j = 0; j < nrows_to; j++) PetscCall(PetscMPIIntCast(sindices[sstarts[i] + j] * bs, &disp[j]));
457: PetscCallMPI(MPI_Type_create_indexed_block(nrows_to, bs_mpi, disp, MPIU_SCALAR, &type1));
458: PetscCallMPI(MPI_Type_create_resized(type1, 0, contents->blda * sizeof(PetscScalar), &contents->stype[i]));
459: PetscCallMPI(MPI_Type_commit(&contents->stype[i]));
460: PetscCallMPI(MPI_Type_free(&type1));
461: }
463: for (PetscMPIInt i = 0; i < nrecvs; i++) {
464: /* received values from a process form a (nrows_from x Bbn) row block in workB (column-wise) */
465: PetscCall(PetscMPIIntCast((rstarts[i + 1] - rstarts[i]) * bs, &nrows_from));
466: disp[0] = 0;
467: PetscCallMPI(MPI_Type_create_indexed_block(1, nrows_from, disp, MPIU_SCALAR, &type1));
468: PetscCallMPI(MPI_Type_create_resized(type1, 0, nrows * sizeof(PetscScalar), &contents->rtype[i]));
469: PetscCallMPI(MPI_Type_commit(&contents->rtype[i]));
470: PetscCallMPI(MPI_Type_free(&type1));
471: }
473: PetscCall(PetscFree(disp));
474: PetscCall(VecScatterRestoreRemote_Private(ctx, PETSC_TRUE /*send*/, &nsends, &sstarts, &sindices, NULL, NULL));
475: PetscCall(VecScatterRestoreRemoteOrdered_Private(ctx, PETSC_FALSE /*recv*/, &nrecvs, &rstarts, NULL, NULL, NULL));
476: if (batchSize) *batchSize = Bbn;
477: if (numBatches) *numBatches = numBb;
478: PetscFunctionReturn(PETSC_SUCCESS);
479: }
481: static PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIDense(Mat A, Mat B, PetscReal fill, Mat C)
482: {
483: Mat_MPIAIJ *aij = (Mat_MPIAIJ *)A->data;
484: MPIAIJ_MPIDense *contents;
485: PetscInt nz = aij->B->cmap->n, m, M, n, N;
486: VecScatter ctx = aij->Mvctx;
487: PetscInt Am = A->rmap->n, BN = B->cmap->N;
488: PetscBool cisdense;
490: PetscFunctionBegin;
491: MatCheckProduct(C, 4);
492: PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
493: PetscCall(PetscObjectBaseTypeCompare((PetscObject)C, MATMPIDENSE, &cisdense));
494: if (!cisdense) PetscCall(MatSetType(C, ((PetscObject)B)->type_name));
495: PetscCall(MatGetLocalSize(C, &m, &n));
496: PetscCall(MatGetSize(C, &M, &N));
497: if (m == PETSC_DECIDE || n == PETSC_DECIDE || M == PETSC_DECIDE || N == PETSC_DECIDE) PetscCall(MatSetSizes(C, Am, B->cmap->n, A->rmap->N, BN));
498: PetscCall(MatSetBlockSizesFromMats(C, A, B));
499: PetscCall(MatSetUp(C));
500: PetscCall(PetscNew(&contents));
501: PetscCall(MatMPIDenseScatterSetUp_Private(ctx, nz, 1, Am, B, C, contents, NULL, NULL));
502: PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
503: PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
504: PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
505: PetscCall(MatProductClear(aij->A));
506: PetscCall(MatProductClear(((Mat_MPIDense *)B->data)->A));
507: PetscCall(MatProductClear(((Mat_MPIDense *)C->data)->A));
508: PetscCall(MatProductCreateWithMat(aij->A, ((Mat_MPIDense *)B->data)->A, NULL, ((Mat_MPIDense *)C->data)->A));
509: PetscCall(MatProductSetType(((Mat_MPIDense *)C->data)->A, MATPRODUCT_AB));
510: PetscCall(MatProductSetFromOptions(((Mat_MPIDense *)C->data)->A));
511: PetscCall(MatProductSymbolic(((Mat_MPIDense *)C->data)->A));
512: C->product->data = contents;
513: C->product->destroy = MatMPIAIJ_MPIDenseDestroy;
514: C->ops->matmultnumeric = MatMatMultNumeric_MPIAIJ_MPIDense;
515: PetscFunctionReturn(PETSC_SUCCESS);
516: }
518: PETSC_INTERN PetscErrorCode MatMatMultNumericAdd_SeqAIJ_SeqDense(Mat, Mat, Mat, const PetscBool);
520: /*
521: Performs an efficient scatter on the rows of B needed by this process; this is
522: a modification of the VecScatterBegin_() routines.
523: */
525: PETSC_INTERN PetscErrorCode MatMPIDenseScatter_Private(VecScatter ctx, PetscInt nrows, PetscInt bs, Mat workB, MPIAIJ_MPIDense *contents, Mat B, Mat C)
526: {
527: const PetscScalar *b;
528: PetscScalar *rvalues;
529: const PetscInt *sindices, *sstarts, *rstarts;
530: const PetscMPIInt *sprocs, *rprocs;
531: PetscMPIInt nsends, nrecvs;
532: MPI_Comm comm;
533: PetscMPIInt tag = ((PetscObject)ctx)->tag, ncols, nsends_mpi, nrecvs_mpi;
534: PetscInt blda;
536: PetscFunctionBegin;
537: MatCheckProduct(C, 7);
538: PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
539: PetscCall(PetscMPIIntCast(B->cmap->N, &ncols));
540: PetscCall(VecScatterGetRemote_Private(ctx, PETSC_TRUE /*send*/, &nsends, &sstarts, &sindices, &sprocs, NULL /*bs*/));
541: PetscCall(VecScatterGetRemoteOrdered_Private(ctx, PETSC_FALSE /*recv*/, &nrecvs, &rstarts, NULL, &rprocs, NULL /*bs*/));
542: PetscCall(PetscMPIIntCast(nsends, &nsends_mpi));
543: PetscCall(PetscMPIIntCast(nrecvs, &nrecvs_mpi));
544: PetscCheck(nrows == workB->rmap->n, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Number of rows of workB %" PetscInt_FMT " not equal to columns of off-diagonal block %" PetscInt_FMT, workB->rmap->n, nrows);
546: PetscCall(MatDenseGetArrayRead(B, &b));
547: PetscCall(MatDenseGetLDA(B, &blda));
548: PetscCheck(blda == contents->blda, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Cannot reuse an input matrix with lda %" PetscInt_FMT " != %" PetscInt_FMT, blda, contents->blda);
549: PetscCall(MatDenseGetArray(workB, &rvalues));
551: /* Post recv, use MPI derived data type to save memory */
552: PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
553: for (PetscMPIInt i = 0; i < nrecvs; i++) PetscCallMPI(MPIU_Irecv(rvalues + ((rstarts[i] - rstarts[0]) * bs), ncols, contents->rtype[i], rprocs[i], tag, comm, contents->rwaits + i));
554: for (PetscMPIInt i = 0; i < nsends; i++) PetscCallMPI(MPIU_Isend(b, ncols, contents->stype[i], sprocs[i], tag, comm, contents->swaits + i));
556: if (nrecvs) PetscCallMPI(MPI_Waitall(nrecvs_mpi, contents->rwaits, MPI_STATUSES_IGNORE));
557: if (nsends) PetscCallMPI(MPI_Waitall(nsends_mpi, contents->swaits, MPI_STATUSES_IGNORE));
559: PetscCall(VecScatterRestoreRemote_Private(ctx, PETSC_TRUE /*send*/, &nsends, &sstarts, &sindices, &sprocs, NULL));
560: PetscCall(VecScatterRestoreRemoteOrdered_Private(ctx, PETSC_FALSE /*recv*/, &nrecvs, &rstarts, NULL, &rprocs, NULL));
561: PetscCall(MatDenseRestoreArrayRead(B, &b));
562: PetscCall(MatDenseRestoreArray(workB, &rvalues));
563: PetscFunctionReturn(PETSC_SUCCESS);
564: }
566: static PetscErrorCode MatMPIDenseScatter(Mat A, Mat B, Mat workB, Mat C)
567: {
568: Mat_MPIAIJ *aij = (Mat_MPIAIJ *)A->data;
569: MPIAIJ_MPIDense *contents;
571: PetscFunctionBegin;
572: contents = (MPIAIJ_MPIDense *)C->product->data;
573: PetscCall(MatMPIDenseScatter_Private(aij->Mvctx, aij->B->cmap->n, 1, workB, contents, B, C));
574: PetscFunctionReturn(PETSC_SUCCESS);
575: }
577: static PetscErrorCode MatMatMultNumeric_MPIAIJ_MPIDense(Mat A, Mat B, Mat C)
578: {
579: Mat_MPIAIJ *aij = (Mat_MPIAIJ *)A->data;
580: Mat_MPIDense *bdense = (Mat_MPIDense *)B->data;
581: Mat_MPIDense *cdense = (Mat_MPIDense *)C->data;
582: Mat workB;
583: MPIAIJ_MPIDense *contents;
585: PetscFunctionBegin;
586: MatCheckProduct(C, 3);
587: PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
588: contents = (MPIAIJ_MPIDense *)C->product->data;
589: /* diagonal block of A times all local rows of B, first make sure that everything is up-to-date */
590: if (!cdense->A->product) {
591: PetscCall(MatProductCreateWithMat(aij->A, bdense->A, NULL, cdense->A));
592: PetscCall(MatProductSetType(cdense->A, MATPRODUCT_AB));
593: PetscCall(MatProductSetFromOptions(cdense->A));
594: PetscCall(MatProductSymbolic(cdense->A));
595: } else PetscCall(MatProductReplaceMats(aij->A, bdense->A, NULL, cdense->A));
596: if (PetscDefined(HAVE_CUPM) && !cdense->A->product->clear) {
597: PetscBool flg;
599: PetscCall(PetscObjectTypeCompare((PetscObject)C, MATMPIDENSE, &flg));
600: if (flg) PetscCall(PetscObjectTypeCompare((PetscObject)A, MATMPIAIJ, &flg));
601: if (!flg) cdense->A->product->clear = PETSC_TRUE; /* if either A or C is a device Mat, make sure MatProductClear() is called */
602: }
603: PetscCall(MatProductNumeric(cdense->A));
604: if (contents->workB->cmap->n == B->cmap->N) {
605: /* get off processor parts of B needed to complete C=A*B */
606: workB = contents->workB;
607: PetscCall(MatMPIDenseScatter(A, B, workB, C));
609: /* off-diagonal block of A times nonlocal rows of B */
610: PetscCall(MatMatMultNumericAdd_SeqAIJ_SeqDense(aij->B, workB, cdense->A, PETSC_TRUE));
611: } else {
612: Mat Bb, Cb;
613: PetscInt BN = B->cmap->N, n = contents->workB->cmap->n, cols;
614: PetscBool ccpu;
616: PetscCheck(n > 0, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Column block size %" PetscInt_FMT " must be positive", n);
617: /* Prevent from unneeded copies back and forth from the GPU
618: when getting and restoring the submatrix
619: We need a proper GPU code for AIJ * dense in parallel */
620: PetscCall(MatBoundToCPU(C, &ccpu));
621: PetscCall(MatBindToCPU(C, PETSC_TRUE));
622: for (PetscInt i = 0; i < BN; i += n) {
623: cols = PetscMin(n, BN - i);
624: workB = contents->workB;
625: if (cols != n) PetscCall(MatDenseGetSubMatrix(contents->workB, PETSC_DECIDE, PETSC_DECIDE, 0, cols, &workB));
626: PetscCall(MatDenseGetSubMatrix(B, PETSC_DECIDE, PETSC_DECIDE, i, i + cols, &Bb));
627: PetscCall(MatDenseGetSubMatrix(C, PETSC_DECIDE, PETSC_DECIDE, i, i + cols, &Cb));
629: /* get off processor parts of B needed to complete C=A*B */
630: PetscCall(MatMPIDenseScatter(A, Bb, workB, C));
632: /* off-diagonal block of A times nonlocal rows of B */
633: cdense = (Mat_MPIDense *)Cb->data;
634: PetscCall(MatMatMultNumericAdd_SeqAIJ_SeqDense(aij->B, workB, cdense->A, PETSC_TRUE));
635: if (cols != n) PetscCall(MatDenseRestoreSubMatrix(contents->workB, &workB));
636: PetscCall(MatDenseRestoreSubMatrix(B, &Bb));
637: PetscCall(MatDenseRestoreSubMatrix(C, &Cb));
638: }
639: PetscCall(MatBindToCPU(C, ccpu));
640: }
641: PetscFunctionReturn(PETSC_SUCCESS);
642: }
644: PetscErrorCode MatMatMultNumeric_MPIAIJ_MPIAIJ(Mat A, Mat P, Mat C)
645: {
646: Mat_MPIAIJ *a = (Mat_MPIAIJ *)A->data, *c = (Mat_MPIAIJ *)C->data;
647: Mat_SeqAIJ *ad = (Mat_SeqAIJ *)a->A->data, *ao = (Mat_SeqAIJ *)a->B->data;
648: Mat_SeqAIJ *cd = (Mat_SeqAIJ *)c->A->data, *co = (Mat_SeqAIJ *)c->B->data;
649: PetscInt *adi = ad->i, *adj, *aoi = ao->i, *aoj;
650: PetscScalar *ada, *aoa, *cda = cd->a, *coa = co->a;
651: Mat_SeqAIJ *p_loc, *p_oth;
652: PetscInt *pi_loc, *pj_loc, *pi_oth, *pj_oth, *pj;
653: PetscScalar *pa_loc, *pa_oth, *pa, valtmp, *ca;
654: PetscInt cm = C->rmap->n, anz, pnz;
655: MatProductCtx_APMPI *ptap;
656: PetscScalar *apa_sparse;
657: const PetscScalar *dummy;
658: PetscInt *api, *apj, *apJ, i, j, k, row;
659: PetscInt cstart = C->cmap->rstart;
660: PetscInt cdnz, conz, k0, k1, nextp;
661: MPI_Comm comm;
662: PetscMPIInt size;
664: PetscFunctionBegin;
665: MatCheckProduct(C, 3);
666: ptap = (MatProductCtx_APMPI *)C->product->data;
667: PetscCheck(ptap, PetscObjectComm((PetscObject)C), PETSC_ERR_ARG_WRONGSTATE, "PtAP cannot be computed. Missing data");
668: PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
669: PetscCallMPI(MPI_Comm_size(comm, &size));
670: PetscCheck(ptap->P_oth || size <= 1, PetscObjectComm((PetscObject)C), PETSC_ERR_ARG_WRONGSTATE, "AP cannot be reused. Do not call MatProductClear()");
672: /* flag CPU mask for C */
673: #if PetscDefined(HAVE_DEVICE)
674: if (C->offloadmask != PETSC_OFFLOAD_UNALLOCATED) C->offloadmask = PETSC_OFFLOAD_CPU;
675: if (c->A->offloadmask != PETSC_OFFLOAD_UNALLOCATED) c->A->offloadmask = PETSC_OFFLOAD_CPU;
676: if (c->B->offloadmask != PETSC_OFFLOAD_UNALLOCATED) c->B->offloadmask = PETSC_OFFLOAD_CPU;
677: #endif
678: apa_sparse = ptap->apa;
680: /* 1) get P_oth = ptap->P_oth and P_loc = ptap->P_loc */
681: /* update numerical values of P_oth and P_loc */
682: PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A, P, MAT_REUSE_MATRIX, &ptap->startsj_s, &ptap->startsj_r, &ptap->bufa, &ptap->P_oth));
683: PetscCall(MatMPIAIJGetLocalMat(P, MAT_REUSE_MATRIX, &ptap->P_loc));
685: /* 2) compute numeric C_loc = A_loc*P = Ad*P_loc + Ao*P_oth */
686: /* get data from symbolic products */
687: p_loc = (Mat_SeqAIJ *)ptap->P_loc->data;
688: pi_loc = p_loc->i;
689: pj_loc = p_loc->j;
690: pa_loc = p_loc->a;
691: if (size > 1) {
692: p_oth = (Mat_SeqAIJ *)ptap->P_oth->data;
693: pi_oth = p_oth->i;
694: pj_oth = p_oth->j;
695: pa_oth = p_oth->a;
696: } else {
697: p_oth = NULL;
698: pi_oth = NULL;
699: pj_oth = NULL;
700: pa_oth = NULL;
701: }
703: /* trigger copy to CPU */
704: PetscCall(MatSeqAIJGetArrayRead(a->A, &dummy));
705: PetscCall(MatSeqAIJRestoreArrayRead(a->A, &dummy));
706: PetscCall(MatSeqAIJGetArrayRead(a->B, &dummy));
707: PetscCall(MatSeqAIJRestoreArrayRead(a->B, &dummy));
708: api = ptap->api;
709: apj = ptap->apj;
710: for (i = 0; i < cm; i++) {
711: apJ = apj + api[i];
713: /* diagonal portion of A */
714: anz = adi[i + 1] - adi[i];
715: adj = ad->j + adi[i];
716: ada = ad->a + adi[i];
717: for (j = 0; j < anz; j++) {
718: row = adj[j];
719: pnz = pi_loc[row + 1] - pi_loc[row];
720: pj = pj_loc + pi_loc[row];
721: pa = pa_loc + pi_loc[row];
722: /* perform sparse axpy */
723: valtmp = ada[j];
724: nextp = 0;
725: for (k = 0; nextp < pnz; k++) {
726: if (apJ[k] == pj[nextp]) { /* column of AP == column of P */
727: apa_sparse[k] += valtmp * pa[nextp++];
728: }
729: }
730: PetscCall(PetscLogFlops(2.0 * pnz));
731: }
733: /* off-diagonal portion of A */
734: anz = aoi[i + 1] - aoi[i];
735: aoj = PetscSafePointerPlusOffset(ao->j, aoi[i]);
736: aoa = PetscSafePointerPlusOffset(ao->a, aoi[i]);
737: for (j = 0; j < anz; j++) {
738: row = aoj[j];
739: pnz = pi_oth[row + 1] - pi_oth[row];
740: pj = pj_oth + pi_oth[row];
741: pa = pa_oth + pi_oth[row];
742: /* perform sparse axpy */
743: valtmp = aoa[j];
744: nextp = 0;
745: for (k = 0; nextp < pnz; k++) {
746: if (apJ[k] == pj[nextp]) { /* column of AP == column of P */
747: apa_sparse[k] += valtmp * pa[nextp++];
748: }
749: }
750: PetscCall(PetscLogFlops(2.0 * pnz));
751: }
753: /* set values in C */
754: cdnz = cd->i[i + 1] - cd->i[i];
755: conz = co->i[i + 1] - co->i[i];
757: /* 1st off-diagonal part of C */
758: ca = PetscSafePointerPlusOffset(coa, co->i[i]);
759: k = 0;
760: for (k0 = 0; k0 < conz; k0++) {
761: if (apJ[k] >= cstart) break;
762: ca[k0] = apa_sparse[k];
763: apa_sparse[k] = 0.0;
764: k++;
765: }
767: /* diagonal part of C */
768: ca = cda + cd->i[i];
769: for (k1 = 0; k1 < cdnz; k1++) {
770: ca[k1] = apa_sparse[k];
771: apa_sparse[k] = 0.0;
772: k++;
773: }
775: /* 2nd off-diagonal part of C */
776: ca = PetscSafePointerPlusOffset(coa, co->i[i]);
777: for (; k0 < conz; k0++) {
778: ca[k0] = apa_sparse[k];
779: apa_sparse[k] = 0.0;
780: k++;
781: }
782: }
783: PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
784: PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
785: PetscFunctionReturn(PETSC_SUCCESS);
786: }
788: /* same as MatMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(), except using LLCondensed to avoid O(BN) memory requirement */
789: PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIAIJ(Mat A, Mat P, PetscReal fill, Mat C)
790: {
791: MPI_Comm comm;
792: PetscMPIInt size;
793: MatProductCtx_APMPI *ptap;
794: PetscFreeSpaceList free_space = NULL, current_space = NULL;
795: Mat_MPIAIJ *a = (Mat_MPIAIJ *)A->data;
796: Mat_SeqAIJ *ad = (Mat_SeqAIJ *)a->A->data, *ao = (Mat_SeqAIJ *)a->B->data, *p_loc, *p_oth;
797: PetscInt *pi_loc, *pj_loc, *pi_oth, *pj_oth, *dnz, *onz;
798: PetscInt *adi = ad->i, *adj = ad->j, *aoi = ao->i, *aoj = ao->j, rstart = A->rmap->rstart;
799: PetscInt i, pnz, row, *api, *apj, *Jptr, apnz, nspacedouble = 0, j, nzi, *lnk, apnz_max = 1;
800: PetscInt am = A->rmap->n, pn = P->cmap->n, pm = P->rmap->n, lsize = pn + 20;
801: PetscReal afill;
802: MatType mtype;
804: PetscFunctionBegin;
805: MatCheckProduct(C, 4);
806: PetscCheck(!C->product->data, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Extra product struct not empty");
807: PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
808: PetscCallMPI(MPI_Comm_size(comm, &size));
810: /* create struct MatProductCtx_APMPI and attached it to C later */
811: PetscCall(PetscNew(&ptap));
813: /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */
814: PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A, P, MAT_INITIAL_MATRIX, &ptap->startsj_s, &ptap->startsj_r, &ptap->bufa, &ptap->P_oth));
816: /* get P_loc by taking all local rows of P */
817: PetscCall(MatMPIAIJGetLocalMat(P, MAT_INITIAL_MATRIX, &ptap->P_loc));
819: p_loc = (Mat_SeqAIJ *)ptap->P_loc->data;
820: pi_loc = p_loc->i;
821: pj_loc = p_loc->j;
822: if (size > 1) {
823: p_oth = (Mat_SeqAIJ *)ptap->P_oth->data;
824: pi_oth = p_oth->i;
825: pj_oth = p_oth->j;
826: } else {
827: p_oth = NULL;
828: pi_oth = NULL;
829: pj_oth = NULL;
830: }
832: /* first, compute symbolic AP = A_loc*P = A_diag*P_loc + A_off*P_oth */
833: PetscCall(PetscMalloc1(am + 1, &api));
834: ptap->api = api;
835: api[0] = 0;
837: PetscCall(PetscLLCondensedCreate_Scalable(lsize, &lnk));
839: /* Initial FreeSpace size is fill*(nnz(A)+nnz(P)) */
840: PetscCall(PetscFreeSpaceGet(PetscRealIntMultTruncate(fill, PetscIntSumTruncate(adi[am], PetscIntSumTruncate(aoi[am], pi_loc[pm]))), &free_space));
841: current_space = free_space;
842: MatPreallocateBegin(comm, am, pn, dnz, onz);
843: for (i = 0; i < am; i++) {
844: /* diagonal portion of A */
845: nzi = adi[i + 1] - adi[i];
846: for (j = 0; j < nzi; j++) {
847: row = *adj++;
848: pnz = pi_loc[row + 1] - pi_loc[row];
849: Jptr = pj_loc + pi_loc[row];
850: /* Expand list if it is not long enough */
851: if (pnz + apnz_max > lsize) {
852: lsize = pnz + apnz_max;
853: PetscCall(PetscLLCondensedExpand_Scalable(lsize, &lnk));
854: }
855: /* add non-zero cols of P into the sorted linked list lnk */
856: PetscCall(PetscLLCondensedAddSorted_Scalable(pnz, Jptr, lnk));
857: apnz = *lnk; /* The first element in the list is the number of items in the list */
858: api[i + 1] = api[i] + apnz;
859: if (apnz > apnz_max) apnz_max = apnz + 1; /* '1' for diagonal entry */
860: }
861: /* off-diagonal portion of A */
862: nzi = aoi[i + 1] - aoi[i];
863: for (j = 0; j < nzi; j++) {
864: row = *aoj++;
865: pnz = pi_oth[row + 1] - pi_oth[row];
866: Jptr = pj_oth + pi_oth[row];
867: /* Expand list if it is not long enough */
868: if (pnz + apnz_max > lsize) {
869: lsize = pnz + apnz_max;
870: PetscCall(PetscLLCondensedExpand_Scalable(lsize, &lnk));
871: }
872: /* add non-zero cols of P into the sorted linked list lnk */
873: PetscCall(PetscLLCondensedAddSorted_Scalable(pnz, Jptr, lnk));
874: apnz = *lnk; /* The first element in the list is the number of items in the list */
875: api[i + 1] = api[i] + apnz;
876: if (apnz > apnz_max) apnz_max = apnz + 1; /* '1' for diagonal entry */
877: }
879: /* add missing diagonal entry */
880: if (C->force_diagonals) {
881: j = i + rstart; /* column index */
882: PetscCall(PetscLLCondensedAddSorted_Scalable(1, &j, lnk));
883: }
885: apnz = *lnk;
886: api[i + 1] = api[i] + apnz;
887: if (apnz > apnz_max) apnz_max = apnz;
889: /* if free space is not available, double the total space in the list */
890: if (current_space->local_remaining < apnz) {
891: PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(apnz, current_space->total_array_size), ¤t_space));
892: nspacedouble++;
893: }
895: /* Copy data into free space, then initialize lnk */
896: PetscCall(PetscLLCondensedClean_Scalable(apnz, current_space->array, lnk));
897: PetscCall(MatPreallocateSet(i + rstart, apnz, current_space->array, dnz, onz));
899: current_space->array += apnz;
900: current_space->local_used += apnz;
901: current_space->local_remaining -= apnz;
902: }
904: /* Allocate space for apj, initialize apj, and */
905: /* destroy list of free space and other temporary array(s) */
906: PetscCall(PetscMalloc1(api[am], &ptap->apj));
907: apj = ptap->apj;
908: PetscCall(PetscFreeSpaceContiguous(&free_space, ptap->apj));
909: PetscCall(PetscLLCondensedDestroy_Scalable(lnk));
911: /* create and assemble symbolic parallel matrix C */
912: PetscCall(MatSetSizes(C, am, pn, PETSC_DETERMINE, PETSC_DETERMINE));
913: PetscCall(MatSetBlockSizesFromMats(C, A, P));
914: PetscCall(MatGetType(A, &mtype));
915: PetscCall(MatSetType(C, mtype));
916: PetscCall(MatMPIAIJSetPreallocation(C, 0, dnz, 0, onz));
917: MatPreallocateEnd(dnz, onz);
919: /* malloc apa for assembly C */
920: PetscCall(PetscCalloc1(apnz_max, &ptap->apa));
922: PetscCall(MatSetValues_MPIAIJ_CopyFromCSRFormat_Symbolic(C, apj, api));
923: PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
924: PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
925: PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
926: PetscCall(MatSetOption(C, MAT_NEW_NONZERO_LOCATION_ERR, PETSC_TRUE));
928: C->ops->matmultnumeric = MatMatMultNumeric_MPIAIJ_MPIAIJ;
929: C->ops->productnumeric = MatProductNumeric_AB;
931: /* attach the supporting struct to C for reuse */
932: C->product->data = ptap;
933: C->product->destroy = MatProductCtxDestroy_MPIAIJ_MatMatMult;
935: /* set MatInfo */
936: afill = (PetscReal)api[am] / (adi[am] + aoi[am] + pi_loc[pm] + 1) + 1.e-5;
937: if (afill < 1.0) afill = 1.0;
938: C->info.mallocs = nspacedouble;
939: C->info.fill_ratio_given = fill;
940: C->info.fill_ratio_needed = afill;
942: if (PetscDefined(USE_INFO)) {
943: if (api[am]) {
944: PetscCall(PetscInfo(C, "Reallocs %" PetscInt_FMT "; Fill ratio: given %g needed %g.\n", nspacedouble, (double)fill, (double)afill));
945: PetscCall(PetscInfo(C, "Use MatMatMult(A,B,MatReuse,%g,&C) for best performance.;\n", (double)afill));
946: } else PetscCall(PetscInfo(C, "Empty matrix product\n"));
947: }
948: PetscFunctionReturn(PETSC_SUCCESS);
949: }
951: /* This function is needed for the seqMPI matrix-matrix multiplication. */
952: /* Three input arrays are merged to one output array. The size of the */
953: /* output array is also output. Duplicate entries only show up once. */
954: static void Merge3SortedArrays(PetscInt size1, PetscInt *in1, PetscInt size2, PetscInt *in2, PetscInt size3, PetscInt *in3, PetscInt *size4, PetscInt *out)
955: {
956: int i = 0, j = 0, k = 0, l = 0;
958: /* Traverse all three arrays */
959: while (i < size1 && j < size2 && k < size3) {
960: if (in1[i] < in2[j] && in1[i] < in3[k]) {
961: out[l++] = in1[i++];
962: } else if (in2[j] < in1[i] && in2[j] < in3[k]) {
963: out[l++] = in2[j++];
964: } else if (in3[k] < in1[i] && in3[k] < in2[j]) {
965: out[l++] = in3[k++];
966: } else if (in1[i] == in2[j] && in1[i] < in3[k]) {
967: out[l++] = in1[i];
968: i++, j++;
969: } else if (in1[i] == in3[k] && in1[i] < in2[j]) {
970: out[l++] = in1[i];
971: i++, k++;
972: } else if (in3[k] == in2[j] && in2[j] < in1[i]) {
973: out[l++] = in2[j];
974: k++, j++;
975: } else if (in1[i] == in2[j] && in1[i] == in3[k]) {
976: out[l++] = in1[i];
977: i++, j++, k++;
978: }
979: }
981: /* Traverse two remaining arrays */
982: while (i < size1 && j < size2) {
983: if (in1[i] < in2[j]) {
984: out[l++] = in1[i++];
985: } else if (in1[i] > in2[j]) {
986: out[l++] = in2[j++];
987: } else {
988: out[l++] = in1[i];
989: i++, j++;
990: }
991: }
993: while (i < size1 && k < size3) {
994: if (in1[i] < in3[k]) {
995: out[l++] = in1[i++];
996: } else if (in1[i] > in3[k]) {
997: out[l++] = in3[k++];
998: } else {
999: out[l++] = in1[i];
1000: i++, k++;
1001: }
1002: }
1004: while (k < size3 && j < size2) {
1005: if (in3[k] < in2[j]) {
1006: out[l++] = in3[k++];
1007: } else if (in3[k] > in2[j]) {
1008: out[l++] = in2[j++];
1009: } else {
1010: out[l++] = in3[k];
1011: k++, j++;
1012: }
1013: }
1015: /* Traverse one remaining array */
1016: while (i < size1) out[l++] = in1[i++];
1017: while (j < size2) out[l++] = in2[j++];
1018: while (k < size3) out[l++] = in3[k++];
1020: *size4 = l;
1021: }
1023: /* This matrix-matrix multiplication algorithm divides the multiplication into three multiplications and */
1024: /* adds up the products. Two of these three multiplications are performed with existing (sequential) */
1025: /* matrix-matrix multiplications. */
1026: PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIAIJ_seqMPI(Mat A, Mat P, PetscReal fill, Mat C)
1027: {
1028: MPI_Comm comm;
1029: PetscMPIInt size;
1030: MatProductCtx_APMPI *ptap;
1031: PetscFreeSpaceList free_space_diag = NULL, current_space = NULL;
1032: Mat_MPIAIJ *a = (Mat_MPIAIJ *)A->data;
1033: Mat_SeqAIJ *ad = (Mat_SeqAIJ *)a->A->data, *ao = (Mat_SeqAIJ *)a->B->data, *p_loc;
1034: Mat_MPIAIJ *p = (Mat_MPIAIJ *)P->data;
1035: Mat_SeqAIJ *adpd_seq, *p_off, *aopoth_seq;
1036: PetscInt adponz, adpdnz;
1037: PetscInt *pi_loc, *dnz, *onz;
1038: PetscInt *adi = ad->i, *adj = ad->j, *aoi = ao->i, rstart = A->rmap->rstart;
1039: PetscInt *lnk, i, i1 = 0, pnz, row, *adpoi, *adpoj, *api, *adpoJ, *aopJ, *apJ, *Jptr, aopnz, nspacedouble = 0, j, nzi, *apj, apnz, *adpdi, *adpdj, *adpdJ, *poff_i, *poff_j, *j_temp, *aopothi, *aopothj;
1040: PetscInt am = A->rmap->n, pN = P->cmap->N, pn = P->cmap->n, pm = P->rmap->n, p_colstart, p_colend;
1041: PetscBT lnkbt;
1042: PetscReal afill;
1043: PetscMPIInt rank;
1044: Mat adpd, aopoth;
1045: MatType mtype;
1046: const char *prefix;
1048: PetscFunctionBegin;
1049: MatCheckProduct(C, 4);
1050: PetscCheck(!C->product->data, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Extra product struct not empty");
1051: PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
1052: PetscCallMPI(MPI_Comm_size(comm, &size));
1053: PetscCallMPI(MPI_Comm_rank(comm, &rank));
1054: PetscCall(MatGetOwnershipRangeColumn(P, &p_colstart, &p_colend));
1056: /* create struct MatProductCtx_APMPI and attached it to C later */
1057: PetscCall(PetscNew(&ptap));
1059: /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */
1060: PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A, P, MAT_INITIAL_MATRIX, &ptap->startsj_s, &ptap->startsj_r, &ptap->bufa, &ptap->P_oth));
1062: /* get P_loc by taking all local rows of P */
1063: PetscCall(MatMPIAIJGetLocalMat(P, MAT_INITIAL_MATRIX, &ptap->P_loc));
1065: p_loc = (Mat_SeqAIJ *)ptap->P_loc->data;
1066: pi_loc = p_loc->i;
1068: /* Allocate memory for the i arrays of the matrices A*P, A_diag*P_off and A_offd * P */
1069: PetscCall(PetscMalloc1(am + 1, &api));
1070: PetscCall(PetscMalloc1(am + 1, &adpoi));
1072: adpoi[0] = 0;
1073: ptap->api = api;
1074: api[0] = 0;
1076: /* create and initialize a linked list, will be used for both A_diag * P_loc_off and A_offd * P_oth */
1077: PetscCall(PetscLLCondensedCreate(pN, pN, &lnk, &lnkbt));
1078: MatPreallocateBegin(comm, am, pn, dnz, onz);
1080: /* Symbolic calc of A_loc_diag * P_loc_diag */
1081: PetscCall(MatGetOptionsPrefix(A, &prefix));
1082: PetscCall(MatProductCreate(a->A, p->A, NULL, &adpd));
1083: PetscCall(MatGetOptionsPrefix(A, &prefix));
1084: PetscCall(MatSetOptionsPrefix(adpd, prefix));
1085: PetscCall(MatAppendOptionsPrefix(adpd, "inner_diag_"));
1087: PetscCall(MatProductSetType(adpd, MATPRODUCT_AB));
1088: PetscCall(MatProductSetAlgorithm(adpd, "sorted"));
1089: PetscCall(MatProductSetFill(adpd, fill));
1090: PetscCall(MatProductSetFromOptions(adpd));
1092: adpd->force_diagonals = C->force_diagonals;
1093: PetscCall(MatProductSymbolic(adpd));
1095: adpd_seq = (Mat_SeqAIJ *)((adpd)->data);
1096: adpdi = adpd_seq->i;
1097: adpdj = adpd_seq->j;
1098: p_off = (Mat_SeqAIJ *)p->B->data;
1099: poff_i = p_off->i;
1100: poff_j = p_off->j;
1102: /* j_temp stores indices of a result row before they are added to the linked list */
1103: PetscCall(PetscMalloc1(pN, &j_temp));
1105: /* Symbolic calc of the A_diag * p_loc_off */
1106: /* Initial FreeSpace size is fill*(nnz(A)+nnz(P)) */
1107: PetscCall(PetscFreeSpaceGet(PetscRealIntMultTruncate(fill, PetscIntSumTruncate(adi[am], PetscIntSumTruncate(aoi[am], pi_loc[pm]))), &free_space_diag));
1108: current_space = free_space_diag;
1110: for (i = 0; i < am; i++) {
1111: /* A_diag * P_loc_off */
1112: nzi = adi[i + 1] - adi[i];
1113: for (j = 0; j < nzi; j++) {
1114: row = *adj++;
1115: pnz = poff_i[row + 1] - poff_i[row];
1116: Jptr = poff_j + poff_i[row];
1117: for (i1 = 0; i1 < pnz; i1++) j_temp[i1] = p->garray[Jptr[i1]];
1118: /* add non-zero cols of P into the sorted linked list lnk */
1119: PetscCall(PetscLLCondensedAddSorted(pnz, j_temp, lnk, lnkbt));
1120: }
1122: adponz = lnk[0];
1123: adpoi[i + 1] = adpoi[i] + adponz;
1125: /* if free space is not available, double the total space in the list */
1126: if (current_space->local_remaining < adponz) {
1127: PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(adponz, current_space->total_array_size), ¤t_space));
1128: nspacedouble++;
1129: }
1131: /* Copy data into free space, then initialize lnk */
1132: PetscCall(PetscLLCondensedClean(pN, adponz, current_space->array, lnk, lnkbt));
1134: current_space->array += adponz;
1135: current_space->local_used += adponz;
1136: current_space->local_remaining -= adponz;
1137: }
1139: /* Symbolic calc of A_off * P_oth */
1140: PetscCall(MatSetOptionsPrefix(a->B, prefix));
1141: PetscCall(MatAppendOptionsPrefix(a->B, "inner_offdiag_"));
1142: PetscCall(MatCreate(PETSC_COMM_SELF, &aopoth));
1143: PetscCall(MatMatMultSymbolic_SeqAIJ_SeqAIJ(a->B, ptap->P_oth, fill, aopoth));
1144: aopoth_seq = (Mat_SeqAIJ *)((aopoth)->data);
1145: aopothi = aopoth_seq->i;
1146: aopothj = aopoth_seq->j;
1148: /* Allocate space for apj, adpj, aopj, ... */
1149: /* destroy lists of free space and other temporary array(s) */
1151: PetscCall(PetscMalloc1(aopothi[am] + adpoi[am] + adpdi[am], &ptap->apj));
1152: PetscCall(PetscMalloc1(adpoi[am], &adpoj));
1154: /* Copy from linked list to j-array */
1155: PetscCall(PetscFreeSpaceContiguous(&free_space_diag, adpoj));
1156: PetscCall(PetscLLDestroy(lnk, lnkbt));
1158: adpoJ = adpoj;
1159: adpdJ = adpdj;
1160: aopJ = aopothj;
1161: apj = ptap->apj;
1162: apJ = apj; /* still empty */
1164: /* Merge j-arrays of A_off * P, A_diag * P_loc_off, and */
1165: /* A_diag * P_loc_diag to get A*P */
1166: for (i = 0; i < am; i++) {
1167: aopnz = aopothi[i + 1] - aopothi[i];
1168: adponz = adpoi[i + 1] - adpoi[i];
1169: adpdnz = adpdi[i + 1] - adpdi[i];
1171: /* Correct indices from A_diag*P_diag */
1172: for (i1 = 0; i1 < adpdnz; i1++) adpdJ[i1] += p_colstart;
1173: /* Merge j-arrays of A_diag * P_loc_off and A_diag * P_loc_diag and A_off * P_oth */
1174: Merge3SortedArrays(adponz, adpoJ, adpdnz, adpdJ, aopnz, aopJ, &apnz, apJ);
1175: PetscCall(MatPreallocateSet(i + rstart, apnz, apJ, dnz, onz));
1177: aopJ += aopnz;
1178: adpoJ += adponz;
1179: adpdJ += adpdnz;
1180: apJ += apnz;
1181: api[i + 1] = api[i] + apnz;
1182: }
1184: /* malloc apa to store dense row A[i,:]*P */
1185: PetscCall(PetscCalloc1(pN, &ptap->apa));
1187: /* create and assemble symbolic parallel matrix C */
1188: PetscCall(MatSetSizes(C, am, pn, PETSC_DETERMINE, PETSC_DETERMINE));
1189: PetscCall(MatSetBlockSizesFromMats(C, A, P));
1190: PetscCall(MatGetType(A, &mtype));
1191: PetscCall(MatSetType(C, mtype));
1192: PetscCall(MatMPIAIJSetPreallocation(C, 0, dnz, 0, onz));
1193: MatPreallocateEnd(dnz, onz);
1195: PetscCall(MatSetValues_MPIAIJ_CopyFromCSRFormat_Symbolic(C, apj, api));
1196: PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
1197: PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
1198: PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
1199: PetscCall(MatSetOption(C, MAT_NEW_NONZERO_LOCATION_ERR, PETSC_TRUE));
1201: C->ops->matmultnumeric = MatMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable;
1202: C->ops->productnumeric = MatProductNumeric_AB;
1204: /* attach the supporting struct to C for reuse */
1205: C->product->data = ptap;
1206: C->product->destroy = MatProductCtxDestroy_MPIAIJ_MatMatMult;
1208: /* set MatInfo */
1209: afill = (PetscReal)api[am] / (adi[am] + aoi[am] + pi_loc[pm] + 1) + 1.e-5;
1210: if (afill < 1.0) afill = 1.0;
1211: C->info.mallocs = nspacedouble;
1212: C->info.fill_ratio_given = fill;
1213: C->info.fill_ratio_needed = afill;
1215: if (PetscDefined(USE_INFO)) {
1216: if (api[am]) {
1217: PetscCall(PetscInfo(C, "Reallocs %" PetscInt_FMT "; Fill ratio: given %g needed %g.\n", nspacedouble, (double)fill, (double)afill));
1218: PetscCall(PetscInfo(C, "Use MatMatMult(A,B,MatReuse,%g,&C) for best performance.;\n", (double)afill));
1219: } else PetscCall(PetscInfo(C, "Empty matrix product\n"));
1220: }
1222: PetscCall(MatDestroy(&aopoth));
1223: PetscCall(MatDestroy(&adpd));
1224: PetscCall(PetscFree(j_temp));
1225: PetscCall(PetscFree(adpoj));
1226: PetscCall(PetscFree(adpoi));
1227: PetscFunctionReturn(PETSC_SUCCESS);
1228: }
1230: /* This routine only works when scall=MAT_REUSE_MATRIX! */
1231: PetscErrorCode MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_matmatmult(Mat P, Mat A, Mat C)
1232: {
1233: MatProductCtx_APMPI *ptap;
1234: Mat Pt;
1236: PetscFunctionBegin;
1237: MatCheckProduct(C, 3);
1238: ptap = (MatProductCtx_APMPI *)C->product->data;
1239: PetscCheck(ptap, PetscObjectComm((PetscObject)C), PETSC_ERR_ARG_WRONGSTATE, "PtAP cannot be computed. Missing data");
1240: PetscCheck(ptap->Pt, PetscObjectComm((PetscObject)C), PETSC_ERR_ARG_WRONGSTATE, "PtA cannot be reused. Do not call MatProductClear()");
1242: Pt = ptap->Pt;
1243: PetscCall(MatTransposeSetPrecursor(P, Pt));
1244: PetscCall(MatTranspose(P, MAT_REUSE_MATRIX, &Pt));
1245: PetscCall(MatMatMultNumeric_MPIAIJ_MPIAIJ(Pt, A, C));
1246: PetscFunctionReturn(PETSC_SUCCESS);
1247: }
1249: /* This routine is modified from MatPtAPSymbolic_MPIAIJ_MPIAIJ() */
1250: PetscErrorCode MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(Mat P, Mat A, PetscReal fill, Mat C)
1251: {
1252: MatProductCtx_APMPI *ptap;
1253: Mat_MPIAIJ *p = (Mat_MPIAIJ *)P->data;
1254: MPI_Comm comm;
1255: PetscMPIInt size, rank;
1256: PetscFreeSpaceList free_space = NULL, current_space = NULL;
1257: PetscInt pn = P->cmap->n, aN = A->cmap->N, an = A->cmap->n;
1258: PetscInt *lnk, i, k, rstart;
1259: PetscBT lnkbt;
1260: PetscMPIInt tagi, tagj, *len_si, *len_s, *len_ri, nrecv, proc, nsend;
1261: PETSC_UNUSED PetscMPIInt icompleted = 0;
1262: PetscInt **buf_rj, **buf_ri, **buf_ri_k, row, ncols, *cols;
1263: PetscInt len, *dnz, *onz, *owners, nzi;
1264: PetscInt nrows, *buf_s, *buf_si, *buf_si_i, **nextrow, **nextci;
1265: MPI_Request *swaits, *rwaits;
1266: MPI_Status *sstatus, rstatus;
1267: PetscLayout rowmap;
1268: PetscInt *owners_co, *coi, *coj; /* i and j array of (p->B)^T*A*P - used in the communication */
1269: PetscMPIInt *len_r, *id_r; /* array of length of comm->size, store send/recv matrix values */
1270: PetscInt *Jptr, *prmap = p->garray, con, j, Crmax;
1271: Mat_SeqAIJ *a_loc, *c_loc, *c_oth;
1272: PetscHMapI ta;
1273: MatType mtype;
1274: const char *prefix;
1276: PetscFunctionBegin;
1277: PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
1278: PetscCallMPI(MPI_Comm_size(comm, &size));
1279: PetscCallMPI(MPI_Comm_rank(comm, &rank));
1281: /* create symbolic parallel matrix C */
1282: PetscCall(MatGetType(A, &mtype));
1283: PetscCall(MatSetType(C, mtype));
1285: C->ops->transposematmultnumeric = MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable;
1287: /* create struct MatProductCtx_APMPI and attached it to C later */
1288: PetscCall(PetscNew(&ptap));
1290: /* (0) compute Rd = Pd^T, Ro = Po^T */
1291: PetscCall(MatTranspose(p->A, MAT_INITIAL_MATRIX, &ptap->Rd));
1292: PetscCall(MatTranspose(p->B, MAT_INITIAL_MATRIX, &ptap->Ro));
1294: /* (1) compute symbolic A_loc */
1295: PetscCall(MatMPIAIJGetLocalMat(A, MAT_INITIAL_MATRIX, &ptap->A_loc));
1297: /* (2-1) compute symbolic C_oth = Ro*A_loc */
1298: PetscCall(MatGetOptionsPrefix(A, &prefix));
1299: PetscCall(MatSetOptionsPrefix(ptap->Ro, prefix));
1300: PetscCall(MatAppendOptionsPrefix(ptap->Ro, "inner_offdiag_"));
1301: PetscCall(MatCreate(PETSC_COMM_SELF, &ptap->C_oth));
1302: PetscCall(MatMatMultSymbolic_SeqAIJ_SeqAIJ(ptap->Ro, ptap->A_loc, fill, ptap->C_oth));
1304: /* (3) send coj of C_oth to other processors */
1305: /* determine row ownership */
1306: PetscCall(PetscLayoutCreate(comm, &rowmap));
1307: rowmap->n = pn;
1308: rowmap->bs = 1;
1309: PetscCall(PetscLayoutSetUp(rowmap));
1310: owners = rowmap->range;
1312: /* determine the number of messages to send, their lengths */
1313: PetscCall(PetscMalloc4(size, &len_s, size, &len_si, size, &sstatus, size + 1, &owners_co));
1314: PetscCall(PetscArrayzero(len_s, size));
1315: PetscCall(PetscArrayzero(len_si, size));
1317: c_oth = (Mat_SeqAIJ *)ptap->C_oth->data;
1318: coi = c_oth->i;
1319: coj = c_oth->j;
1320: con = ptap->C_oth->rmap->n;
1321: proc = 0;
1322: for (i = 0; i < con; i++) {
1323: while (prmap[i] >= owners[proc + 1]) proc++;
1324: len_si[proc]++; /* num of rows in Co(=Pt*A) to be sent to [proc] */
1325: len_s[proc] += coi[i + 1] - coi[i]; /* num of nonzeros in Co to be sent to [proc] */
1326: }
1328: len = 0; /* max length of buf_si[], see (4) */
1329: owners_co[0] = 0;
1330: nsend = 0;
1331: for (proc = 0; proc < size; proc++) {
1332: owners_co[proc + 1] = owners_co[proc] + len_si[proc];
1333: if (len_s[proc]) {
1334: nsend++;
1335: len_si[proc] = 2 * (len_si[proc] + 1); /* length of buf_si to be sent to [proc] */
1336: len += len_si[proc];
1337: }
1338: }
1340: /* determine the number and length of messages to receive for coi and coj */
1341: PetscCall(PetscGatherNumberOfMessages(comm, NULL, len_s, &nrecv));
1342: PetscCall(PetscGatherMessageLengths2(comm, nsend, nrecv, len_s, len_si, &id_r, &len_r, &len_ri));
1344: /* post the Irecv and Isend of coj */
1345: PetscCall(PetscCommGetNewTag(comm, &tagj));
1346: PetscCall(PetscPostIrecvInt(comm, tagj, nrecv, id_r, len_r, &buf_rj, &rwaits));
1347: PetscCall(PetscMalloc1(nsend, &swaits));
1348: for (proc = 0, k = 0; proc < size; proc++) {
1349: if (!len_s[proc]) continue;
1350: i = owners_co[proc];
1351: PetscCallMPI(MPIU_Isend(coj + coi[i], len_s[proc], MPIU_INT, proc, tagj, comm, swaits + k));
1352: k++;
1353: }
1355: /* (2-2) compute symbolic C_loc = Rd*A_loc */
1356: PetscCall(MatSetOptionsPrefix(ptap->Rd, prefix));
1357: PetscCall(MatAppendOptionsPrefix(ptap->Rd, "inner_diag_"));
1358: PetscCall(MatCreate(PETSC_COMM_SELF, &ptap->C_loc));
1359: PetscCall(MatMatMultSymbolic_SeqAIJ_SeqAIJ(ptap->Rd, ptap->A_loc, fill, ptap->C_loc));
1360: c_loc = (Mat_SeqAIJ *)ptap->C_loc->data;
1362: /* receives coj are complete */
1363: for (i = 0; i < nrecv; i++) PetscCallMPI(MPI_Waitany(nrecv, rwaits, &icompleted, &rstatus));
1364: PetscCall(PetscFree(rwaits));
1365: if (nsend) PetscCallMPI(MPI_Waitall(nsend, swaits, sstatus));
1367: /* add received column indices into ta to update Crmax */
1368: a_loc = (Mat_SeqAIJ *)ptap->A_loc->data;
1370: /* create and initialize a linked list */
1371: PetscCall(PetscHMapICreateWithSize(an, &ta)); /* for compute Crmax */
1372: MatRowMergeMax_SeqAIJ(a_loc, ptap->A_loc->rmap->N, ta);
1374: for (k = 0; k < nrecv; k++) { /* k-th received message */
1375: Jptr = buf_rj[k];
1376: for (j = 0; j < len_r[k]; j++) PetscCall(PetscHMapISet(ta, *(Jptr + j) + 1, 1));
1377: }
1378: PetscCall(PetscHMapIGetSize(ta, &Crmax));
1379: PetscCall(PetscHMapIDestroy(&ta));
1381: /* (4) send and recv coi */
1382: PetscCall(PetscCommGetNewTag(comm, &tagi));
1383: PetscCall(PetscPostIrecvInt(comm, tagi, nrecv, id_r, len_ri, &buf_ri, &rwaits));
1384: PetscCall(PetscMalloc1(len, &buf_s));
1385: buf_si = buf_s; /* points to the beginning of k-th msg to be sent */
1386: for (proc = 0, k = 0; proc < size; proc++) {
1387: if (!len_s[proc]) continue;
1388: /* form outgoing message for i-structure:
1389: buf_si[0]: nrows to be sent
1390: [1:nrows]: row index (global)
1391: [nrows+1:2*nrows+1]: i-structure index
1392: */
1393: nrows = len_si[proc] / 2 - 1; /* num of rows in Co to be sent to [proc] */
1394: buf_si_i = buf_si + nrows + 1;
1395: buf_si[0] = nrows;
1396: buf_si_i[0] = 0;
1397: nrows = 0;
1398: for (i = owners_co[proc]; i < owners_co[proc + 1]; i++) {
1399: nzi = coi[i + 1] - coi[i];
1400: buf_si_i[nrows + 1] = buf_si_i[nrows] + nzi; /* i-structure */
1401: buf_si[nrows + 1] = prmap[i] - owners[proc]; /* local row index */
1402: nrows++;
1403: }
1404: PetscCallMPI(MPIU_Isend(buf_si, len_si[proc], MPIU_INT, proc, tagi, comm, swaits + k));
1405: k++;
1406: buf_si += len_si[proc];
1407: }
1408: for (i = 0; i < nrecv; i++) PetscCallMPI(MPI_Waitany(nrecv, rwaits, &icompleted, &rstatus));
1409: PetscCall(PetscFree(rwaits));
1410: if (nsend) PetscCallMPI(MPI_Waitall(nsend, swaits, sstatus));
1412: PetscCall(PetscFree4(len_s, len_si, sstatus, owners_co));
1413: PetscCall(PetscFree(len_ri));
1414: PetscCall(PetscFree(swaits));
1415: PetscCall(PetscFree(buf_s));
1417: /* (5) compute the local portion of C */
1418: /* set initial free space to be Crmax, sufficient for holding nonzeros in each row of C */
1419: PetscCall(PetscFreeSpaceGet(Crmax, &free_space));
1420: current_space = free_space;
1422: PetscCall(PetscMalloc3(nrecv, &buf_ri_k, nrecv, &nextrow, nrecv, &nextci));
1423: for (k = 0; k < nrecv; k++) {
1424: buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */
1425: nrows = *buf_ri_k[k];
1426: nextrow[k] = buf_ri_k[k] + 1; /* next row number of k-th recved i-structure */
1427: nextci[k] = buf_ri_k[k] + (nrows + 1); /* points to the next i-structure of k-th recved i-structure */
1428: }
1430: MatPreallocateBegin(comm, pn, an, dnz, onz);
1431: PetscCall(PetscLLCondensedCreate(Crmax, aN, &lnk, &lnkbt));
1432: for (i = 0; i < pn; i++) { /* for each local row of C */
1433: /* add C_loc into C */
1434: nzi = c_loc->i[i + 1] - c_loc->i[i];
1435: Jptr = c_loc->j + c_loc->i[i];
1436: PetscCall(PetscLLCondensedAddSorted(nzi, Jptr, lnk, lnkbt));
1438: /* add received col data into lnk */
1439: for (k = 0; k < nrecv; k++) { /* k-th received message */
1440: if (i == *nextrow[k]) { /* i-th row */
1441: nzi = *(nextci[k] + 1) - *nextci[k];
1442: Jptr = buf_rj[k] + *nextci[k];
1443: PetscCall(PetscLLCondensedAddSorted(nzi, Jptr, lnk, lnkbt));
1444: nextrow[k]++;
1445: nextci[k]++;
1446: }
1447: }
1449: /* add missing diagonal entry */
1450: if (C->force_diagonals) {
1451: k = i + owners[rank]; /* column index */
1452: PetscCall(PetscLLCondensedAddSorted(1, &k, lnk, lnkbt));
1453: }
1455: nzi = lnk[0];
1457: /* copy data into free space, then initialize lnk */
1458: PetscCall(PetscLLCondensedClean(aN, nzi, current_space->array, lnk, lnkbt));
1459: PetscCall(MatPreallocateSet(i + owners[rank], nzi, current_space->array, dnz, onz));
1460: }
1461: PetscCall(PetscFree3(buf_ri_k, nextrow, nextci));
1462: PetscCall(PetscLLDestroy(lnk, lnkbt));
1463: PetscCall(PetscFreeSpaceDestroy(free_space));
1465: /* local sizes and preallocation */
1466: PetscCall(MatSetSizes(C, pn, an, PETSC_DETERMINE, PETSC_DETERMINE));
1467: PetscCall(PetscLayoutSetBlockSize(C->rmap, P->cmap->bs));
1468: PetscCall(PetscLayoutSetBlockSize(C->cmap, A->cmap->bs));
1469: PetscCall(MatMPIAIJSetPreallocation(C, 0, dnz, 0, onz));
1470: MatPreallocateEnd(dnz, onz);
1472: /* add C_loc and C_oth to C */
1473: PetscCall(MatGetOwnershipRange(C, &rstart, NULL));
1474: for (i = 0; i < pn; i++) {
1475: ncols = c_loc->i[i + 1] - c_loc->i[i];
1476: cols = c_loc->j + c_loc->i[i];
1477: row = rstart + i;
1478: PetscCall(MatSetValues(C, 1, (const PetscInt *)&row, ncols, (const PetscInt *)cols, NULL, INSERT_VALUES));
1480: if (C->force_diagonals) PetscCall(MatSetValues(C, 1, (const PetscInt *)&row, 1, (const PetscInt *)&row, NULL, INSERT_VALUES));
1481: }
1482: for (i = 0; i < con; i++) {
1483: ncols = c_oth->i[i + 1] - c_oth->i[i];
1484: cols = c_oth->j + c_oth->i[i];
1485: row = prmap[i];
1486: PetscCall(MatSetValues(C, 1, (const PetscInt *)&row, ncols, (const PetscInt *)cols, NULL, INSERT_VALUES));
1487: }
1488: PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
1489: PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
1490: PetscCall(MatSetOption(C, MAT_NEW_NONZERO_LOCATION_ERR, PETSC_TRUE));
1492: /* members in merge */
1493: PetscCall(PetscFree(id_r));
1494: PetscCall(PetscFree(len_r));
1495: PetscCall(PetscFree(buf_ri[0]));
1496: PetscCall(PetscFree(buf_ri));
1497: PetscCall(PetscFree(buf_rj[0]));
1498: PetscCall(PetscFree(buf_rj));
1499: PetscCall(PetscLayoutDestroy(&rowmap));
1501: /* attach the supporting struct to C for reuse */
1502: C->product->data = ptap;
1503: C->product->destroy = MatProductCtxDestroy_MPIAIJ_PtAP;
1504: PetscFunctionReturn(PETSC_SUCCESS);
1505: }
1507: PetscErrorCode MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable(Mat P, Mat A, Mat C)
1508: {
1509: Mat_MPIAIJ *p = (Mat_MPIAIJ *)P->data;
1510: Mat_SeqAIJ *c_seq;
1511: MatProductCtx_APMPI *ptap;
1512: Mat A_loc, C_loc, C_oth;
1513: PetscInt i, rstart, rend, cm, ncols, row;
1514: const PetscInt *cols;
1515: const PetscScalar *vals;
1517: PetscFunctionBegin;
1518: MatCheckProduct(C, 3);
1519: ptap = (MatProductCtx_APMPI *)C->product->data;
1520: PetscCheck(ptap, PetscObjectComm((PetscObject)C), PETSC_ERR_ARG_WRONGSTATE, "PtAP cannot be computed. Missing data");
1521: PetscCheck(ptap->A_loc, PetscObjectComm((PetscObject)C), PETSC_ERR_ARG_WRONGSTATE, "PtA cannot be reused. Do not call MatProductClear()");
1522: PetscCall(MatZeroEntries(C));
1524: /* These matrices are obtained in MatTransposeMatMultSymbolic() */
1525: /* 1) get R = Pd^T, Ro = Po^T */
1526: PetscCall(MatTransposeSetPrecursor(p->A, ptap->Rd));
1527: PetscCall(MatTranspose(p->A, MAT_REUSE_MATRIX, &ptap->Rd));
1528: PetscCall(MatTransposeSetPrecursor(p->B, ptap->Ro));
1529: PetscCall(MatTranspose(p->B, MAT_REUSE_MATRIX, &ptap->Ro));
1531: /* 2) compute numeric A_loc */
1532: PetscCall(MatMPIAIJGetLocalMat(A, MAT_REUSE_MATRIX, &ptap->A_loc));
1534: /* 3) C_loc = Rd*A_loc, C_oth = Ro*A_loc */
1535: A_loc = ptap->A_loc;
1536: PetscCall(ptap->C_loc->ops->matmultnumeric(ptap->Rd, A_loc, ptap->C_loc));
1537: PetscCall(ptap->C_oth->ops->matmultnumeric(ptap->Ro, A_loc, ptap->C_oth));
1538: C_loc = ptap->C_loc;
1539: C_oth = ptap->C_oth;
1541: /* add C_loc and C_oth to C */
1542: PetscCall(MatGetOwnershipRange(C, &rstart, &rend));
1544: /* C_loc -> C */
1545: cm = C_loc->rmap->N;
1546: c_seq = (Mat_SeqAIJ *)C_loc->data;
1547: cols = c_seq->j;
1548: vals = c_seq->a;
1549: for (i = 0; i < cm; i++) {
1550: ncols = c_seq->i[i + 1] - c_seq->i[i];
1551: row = rstart + i;
1552: PetscCall(MatSetValues(C, 1, &row, ncols, cols, vals, ADD_VALUES));
1553: cols += ncols;
1554: vals += ncols;
1555: }
1557: /* Co -> C, off-processor part */
1558: cm = C_oth->rmap->N;
1559: c_seq = (Mat_SeqAIJ *)C_oth->data;
1560: cols = c_seq->j;
1561: vals = c_seq->a;
1562: for (i = 0; i < cm; i++) {
1563: ncols = c_seq->i[i + 1] - c_seq->i[i];
1564: row = p->garray[i];
1565: PetscCall(MatSetValues(C, 1, &row, ncols, cols, vals, ADD_VALUES));
1566: cols += ncols;
1567: vals += ncols;
1568: }
1569: PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
1570: PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
1571: PetscCall(MatSetOption(C, MAT_NEW_NONZERO_LOCATION_ERR, PETSC_TRUE));
1572: PetscFunctionReturn(PETSC_SUCCESS);
1573: }
1575: PetscErrorCode MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ(Mat P, Mat A, Mat C)
1576: {
1577: MatMergeSeqsToMPI *merge;
1578: Mat_MPIAIJ *p = (Mat_MPIAIJ *)P->data;
1579: Mat_SeqAIJ *pd = (Mat_SeqAIJ *)p->A->data, *po = (Mat_SeqAIJ *)p->B->data;
1580: MatProductCtx_APMPI *ap;
1581: PetscInt *adj;
1582: PetscInt i, j, k, anz, pnz, row, *cj, nexta;
1583: MatScalar *ada, *ca, valtmp;
1584: PetscInt am = A->rmap->n, cm = C->rmap->n, pon = (p->B)->cmap->n;
1585: MPI_Comm comm;
1586: PetscMPIInt size, rank, taga, *len_s, proc;
1587: PetscInt *owners, nrows, **buf_ri_k, **nextrow, **nextci;
1588: PetscInt **buf_ri, **buf_rj;
1589: PetscInt cnz = 0, *bj_i, *bi, *bj, bnz, nextcj; /* bi,bj,ba: local array of C(mpi mat) */
1590: MPI_Request *s_waits, *r_waits;
1591: MPI_Status *status;
1592: MatScalar **abuf_r, *ba_i, *pA, *coa, *ba;
1593: const PetscScalar *dummy;
1594: PetscInt *ai, *aj, *coi, *coj, *poJ, *pdJ;
1595: Mat A_loc;
1596: Mat_SeqAIJ *a_loc;
1598: PetscFunctionBegin;
1599: MatCheckProduct(C, 3);
1600: ap = (MatProductCtx_APMPI *)C->product->data;
1601: PetscCheck(ap, PetscObjectComm((PetscObject)C), PETSC_ERR_ARG_WRONGSTATE, "PtA cannot be computed. Missing data");
1602: PetscCheck(ap->A_loc, PetscObjectComm((PetscObject)C), PETSC_ERR_ARG_WRONGSTATE, "PtA cannot be reused. Do not call MatProductClear()");
1603: PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
1604: PetscCallMPI(MPI_Comm_size(comm, &size));
1605: PetscCallMPI(MPI_Comm_rank(comm, &rank));
1607: merge = ap->merge;
1609: /* 2) compute numeric C_seq = P_loc^T*A_loc */
1610: /* get data from symbolic products */
1611: coi = merge->coi;
1612: coj = merge->coj;
1613: PetscCall(PetscCalloc1(coi[pon], &coa));
1614: bi = merge->bi;
1615: bj = merge->bj;
1616: owners = merge->rowmap->range;
1617: PetscCall(PetscCalloc1(bi[cm], &ba));
1619: /* get A_loc by taking all local rows of A */
1620: A_loc = ap->A_loc;
1621: PetscCall(MatMPIAIJGetLocalMat(A, MAT_REUSE_MATRIX, &A_loc));
1622: a_loc = (Mat_SeqAIJ *)A_loc->data;
1623: ai = a_loc->i;
1624: aj = a_loc->j;
1626: /* trigger copy to CPU */
1627: PetscCall(MatSeqAIJGetArrayRead(p->A, &dummy));
1628: PetscCall(MatSeqAIJRestoreArrayRead(p->A, &dummy));
1629: PetscCall(MatSeqAIJGetArrayRead(p->B, &dummy));
1630: PetscCall(MatSeqAIJRestoreArrayRead(p->B, &dummy));
1631: for (i = 0; i < am; i++) {
1632: anz = ai[i + 1] - ai[i];
1633: adj = aj + ai[i];
1634: ada = a_loc->a + ai[i];
1636: /* 2-b) Compute Cseq = P_loc[i,:]^T*A[i,:] using outer product */
1637: /* put the value into Co=(p->B)^T*A (off-diagonal part, send to others) */
1638: pnz = po->i[i + 1] - po->i[i];
1639: poJ = po->j + po->i[i];
1640: pA = po->a + po->i[i];
1641: for (j = 0; j < pnz; j++) {
1642: row = poJ[j];
1643: cj = coj + coi[row];
1644: ca = coa + coi[row];
1645: /* perform sparse axpy */
1646: nexta = 0;
1647: valtmp = pA[j];
1648: for (k = 0; nexta < anz; k++) {
1649: if (cj[k] == adj[nexta]) {
1650: ca[k] += valtmp * ada[nexta];
1651: nexta++;
1652: }
1653: }
1654: PetscCall(PetscLogFlops(2.0 * anz));
1655: }
1657: /* put the value into Cd (diagonal part) */
1658: pnz = pd->i[i + 1] - pd->i[i];
1659: pdJ = pd->j + pd->i[i];
1660: pA = pd->a + pd->i[i];
1661: for (j = 0; j < pnz; j++) {
1662: row = pdJ[j];
1663: cj = bj + bi[row];
1664: ca = ba + bi[row];
1665: /* perform sparse axpy */
1666: nexta = 0;
1667: valtmp = pA[j];
1668: for (k = 0; nexta < anz; k++) {
1669: if (cj[k] == adj[nexta]) {
1670: ca[k] += valtmp * ada[nexta];
1671: nexta++;
1672: }
1673: }
1674: PetscCall(PetscLogFlops(2.0 * anz));
1675: }
1676: }
1678: /* 3) send and recv matrix values coa */
1679: buf_ri = merge->buf_ri;
1680: buf_rj = merge->buf_rj;
1681: len_s = merge->len_s;
1682: PetscCall(PetscCommGetNewTag(comm, &taga));
1683: PetscCall(PetscPostIrecvScalar(comm, taga, merge->nrecv, merge->id_r, merge->len_r, &abuf_r, &r_waits));
1685: PetscCall(PetscMalloc2(merge->nsend, &s_waits, size, &status));
1686: for (proc = 0, k = 0; proc < size; proc++) {
1687: if (!len_s[proc]) continue;
1688: i = merge->owners_co[proc];
1689: PetscCallMPI(MPIU_Isend(coa + coi[i], len_s[proc], MPIU_MATSCALAR, proc, taga, comm, s_waits + k));
1690: k++;
1691: }
1692: if (merge->nrecv) PetscCallMPI(MPI_Waitall(merge->nrecv, r_waits, status));
1693: if (merge->nsend) PetscCallMPI(MPI_Waitall(merge->nsend, s_waits, status));
1695: PetscCall(PetscFree2(s_waits, status));
1696: PetscCall(PetscFree(r_waits));
1697: PetscCall(PetscFree(coa));
1699: /* 4) insert local Cseq and received values into Cmpi */
1700: PetscCall(PetscMalloc3(merge->nrecv, &buf_ri_k, merge->nrecv, &nextrow, merge->nrecv, &nextci));
1701: for (k = 0; k < merge->nrecv; k++) {
1702: buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */
1703: nrows = *buf_ri_k[k];
1704: nextrow[k] = buf_ri_k[k] + 1; /* next row number of k-th recved i-structure */
1705: nextci[k] = buf_ri_k[k] + (nrows + 1); /* points to the next i-structure of k-th recved i-structure */
1706: }
1708: for (i = 0; i < cm; i++) {
1709: row = owners[rank] + i; /* global row index of C_seq */
1710: bj_i = bj + bi[i]; /* col indices of the i-th row of C */
1711: ba_i = ba + bi[i];
1712: bnz = bi[i + 1] - bi[i];
1713: /* add received vals into ba */
1714: for (k = 0; k < merge->nrecv; k++) { /* k-th received message */
1715: /* i-th row */
1716: if (i == *nextrow[k]) {
1717: cnz = *(nextci[k] + 1) - *nextci[k];
1718: cj = buf_rj[k] + *nextci[k];
1719: ca = abuf_r[k] + *nextci[k];
1720: nextcj = 0;
1721: for (j = 0; nextcj < cnz; j++) {
1722: if (bj_i[j] == cj[nextcj]) { /* bcol == ccol */
1723: ba_i[j] += ca[nextcj++];
1724: }
1725: }
1726: nextrow[k]++;
1727: nextci[k]++;
1728: PetscCall(PetscLogFlops(2.0 * cnz));
1729: }
1730: }
1731: PetscCall(MatSetValues(C, 1, &row, bnz, bj_i, ba_i, INSERT_VALUES));
1732: }
1733: PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
1734: PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
1736: PetscCall(PetscFree(ba));
1737: PetscCall(PetscFree(abuf_r[0]));
1738: PetscCall(PetscFree(abuf_r));
1739: PetscCall(PetscFree3(buf_ri_k, nextrow, nextci));
1740: PetscFunctionReturn(PETSC_SUCCESS);
1741: }
1743: PetscErrorCode MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ(Mat P, Mat A, PetscReal fill, Mat C)
1744: {
1745: Mat A_loc;
1746: MatProductCtx_APMPI *ap;
1747: PetscFreeSpaceList free_space = NULL, current_space = NULL;
1748: Mat_MPIAIJ *p = (Mat_MPIAIJ *)P->data, *a = (Mat_MPIAIJ *)A->data;
1749: PetscInt *pdti, *pdtj, *poti, *potj, *ptJ;
1750: PetscInt nnz;
1751: PetscInt *lnk, *owners_co, *coi, *coj, i, k, pnz, row;
1752: PetscInt am = A->rmap->n, pn = P->cmap->n;
1753: MPI_Comm comm;
1754: PetscMPIInt size, rank, tagi, tagj, *len_si, *len_s, *len_ri, proc;
1755: PetscInt **buf_rj, **buf_ri, **buf_ri_k;
1756: PetscInt len, *dnz, *onz, *owners;
1757: PetscInt nzi, *bi, *bj;
1758: PetscInt nrows, *buf_s, *buf_si, *buf_si_i, **nextrow, **nextci;
1759: MPI_Request *swaits, *rwaits;
1760: MPI_Status *sstatus, rstatus;
1761: MatMergeSeqsToMPI *merge;
1762: PetscInt *ai, *aj, *Jptr, anz, *prmap = p->garray, pon, nspacedouble = 0, j;
1763: PetscReal afill = 1.0, afill_tmp;
1764: PetscInt rstart = P->cmap->rstart, rmax, Armax;
1765: Mat_SeqAIJ *a_loc;
1766: PetscHMapI ta;
1767: MatType mtype;
1769: PetscFunctionBegin;
1770: PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
1771: /* check if matrix local sizes are compatible */
1772: PetscCheck(A->rmap->rstart == P->rmap->rstart && A->rmap->rend == P->rmap->rend, comm, PETSC_ERR_ARG_SIZ, "Matrix local dimensions are incompatible, A (%" PetscInt_FMT ", %" PetscInt_FMT ") != P (%" PetscInt_FMT ",%" PetscInt_FMT ")", A->rmap->rstart,
1773: A->rmap->rend, P->rmap->rstart, P->rmap->rend);
1775: PetscCallMPI(MPI_Comm_size(comm, &size));
1776: PetscCallMPI(MPI_Comm_rank(comm, &rank));
1778: /* create struct MatProductCtx_APMPI and attached it to C later */
1779: PetscCall(PetscNew(&ap));
1781: /* get A_loc by taking all local rows of A */
1782: PetscCall(MatMPIAIJGetLocalMat(A, MAT_INITIAL_MATRIX, &A_loc));
1784: ap->A_loc = A_loc;
1785: a_loc = (Mat_SeqAIJ *)A_loc->data;
1786: ai = a_loc->i;
1787: aj = a_loc->j;
1789: /* determine symbolic Co=(p->B)^T*A - send to others */
1790: PetscCall(MatGetSymbolicTranspose_SeqAIJ(p->A, &pdti, &pdtj));
1791: PetscCall(MatGetSymbolicTranspose_SeqAIJ(p->B, &poti, &potj));
1792: pon = (p->B)->cmap->n; /* total num of rows to be sent to other processors
1793: >= (num of nonzero rows of C_seq) - pn */
1794: PetscCall(PetscMalloc1(pon + 1, &coi));
1795: coi[0] = 0;
1797: /* set initial free space to be fill*(nnz(p->B) + nnz(A)) */
1798: nnz = PetscRealIntMultTruncate(fill, PetscIntSumTruncate(poti[pon], ai[am]));
1799: PetscCall(PetscFreeSpaceGet(nnz, &free_space));
1800: current_space = free_space;
1802: /* create and initialize a linked list */
1803: PetscCall(PetscHMapICreateWithSize(A->cmap->n + a->B->cmap->N, &ta));
1804: MatRowMergeMax_SeqAIJ(a_loc, am, ta);
1805: PetscCall(PetscHMapIGetSize(ta, &Armax));
1807: PetscCall(PetscLLCondensedCreate_Scalable(Armax, &lnk));
1809: for (i = 0; i < pon; i++) {
1810: pnz = poti[i + 1] - poti[i];
1811: ptJ = potj + poti[i];
1812: for (j = 0; j < pnz; j++) {
1813: row = ptJ[j]; /* row of A_loc == col of Pot */
1814: anz = ai[row + 1] - ai[row];
1815: Jptr = aj + ai[row];
1816: /* add non-zero cols of AP into the sorted linked list lnk */
1817: PetscCall(PetscLLCondensedAddSorted_Scalable(anz, Jptr, lnk));
1818: }
1819: nnz = lnk[0];
1821: /* If free space is not available, double the total space in the list */
1822: if (current_space->local_remaining < nnz) {
1823: PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(nnz, current_space->total_array_size), ¤t_space));
1824: nspacedouble++;
1825: }
1827: /* Copy data into free space, and zero out denserows */
1828: PetscCall(PetscLLCondensedClean_Scalable(nnz, current_space->array, lnk));
1830: current_space->array += nnz;
1831: current_space->local_used += nnz;
1832: current_space->local_remaining -= nnz;
1834: coi[i + 1] = coi[i] + nnz;
1835: }
1837: PetscCall(PetscMalloc1(coi[pon], &coj));
1838: PetscCall(PetscFreeSpaceContiguous(&free_space, coj));
1839: PetscCall(PetscLLCondensedDestroy_Scalable(lnk)); /* must destroy to get a new one for C */
1841: afill_tmp = (PetscReal)coi[pon] / (poti[pon] + ai[am] + 1);
1842: if (afill_tmp > afill) afill = afill_tmp;
1844: /* send j-array (coj) of Co to other processors */
1845: /* determine row ownership */
1846: PetscCall(PetscNew(&merge));
1847: PetscCall(PetscLayoutCreate(comm, &merge->rowmap));
1849: merge->rowmap->n = pn;
1850: merge->rowmap->bs = 1;
1852: PetscCall(PetscLayoutSetUp(merge->rowmap));
1853: owners = merge->rowmap->range;
1855: /* determine the number of messages to send, their lengths */
1856: PetscCall(PetscCalloc1(size, &len_si));
1857: PetscCall(PetscCalloc1(size, &merge->len_s));
1859: len_s = merge->len_s;
1860: merge->nsend = 0;
1862: PetscCall(PetscMalloc1(size + 1, &owners_co));
1864: proc = 0;
1865: for (i = 0; i < pon; i++) {
1866: while (prmap[i] >= owners[proc + 1]) proc++;
1867: len_si[proc]++; /* num of rows in Co to be sent to [proc] */
1868: len_s[proc] += coi[i + 1] - coi[i];
1869: }
1871: len = 0; /* max length of buf_si[] */
1872: owners_co[0] = 0;
1873: for (proc = 0; proc < size; proc++) {
1874: owners_co[proc + 1] = owners_co[proc] + len_si[proc];
1875: if (len_s[proc]) {
1876: merge->nsend++;
1877: len_si[proc] = 2 * (len_si[proc] + 1);
1878: len += len_si[proc];
1879: }
1880: }
1882: /* determine the number and length of messages to receive for coi and coj */
1883: PetscCall(PetscGatherNumberOfMessages(comm, NULL, len_s, &merge->nrecv));
1884: PetscCall(PetscGatherMessageLengths2(comm, merge->nsend, merge->nrecv, len_s, len_si, &merge->id_r, &merge->len_r, &len_ri));
1886: /* post the Irecv and Isend of coj */
1887: PetscCall(PetscCommGetNewTag(comm, &tagj));
1888: PetscCall(PetscPostIrecvInt(comm, tagj, merge->nrecv, merge->id_r, merge->len_r, &buf_rj, &rwaits));
1889: PetscCall(PetscMalloc1(merge->nsend, &swaits));
1890: for (proc = 0, k = 0; proc < size; proc++) {
1891: if (!len_s[proc]) continue;
1892: i = owners_co[proc];
1893: PetscCallMPI(MPIU_Isend(coj + coi[i], len_s[proc], MPIU_INT, proc, tagj, comm, swaits + k));
1894: k++;
1895: }
1897: /* receives and sends of coj are complete */
1898: PetscCall(PetscMalloc1(size, &sstatus));
1899: for (i = 0; i < merge->nrecv; i++) {
1900: PETSC_UNUSED PetscMPIInt icompleted;
1901: PetscCallMPI(MPI_Waitany(merge->nrecv, rwaits, &icompleted, &rstatus));
1902: }
1903: PetscCall(PetscFree(rwaits));
1904: if (merge->nsend) PetscCallMPI(MPI_Waitall(merge->nsend, swaits, sstatus));
1906: /* add received column indices into table to update Armax */
1907: /* Armax can be as large as aN if a P[row,:] is dense, see src/ksp/ksp/tutorials/ex56.c! */
1908: for (k = 0; k < merge->nrecv; k++) { /* k-th received message */
1909: Jptr = buf_rj[k];
1910: for (j = 0; j < merge->len_r[k]; j++) PetscCall(PetscHMapISet(ta, *(Jptr + j) + 1, 1));
1911: }
1912: PetscCall(PetscHMapIGetSize(ta, &Armax));
1914: /* send and recv coi */
1915: PetscCall(PetscCommGetNewTag(comm, &tagi));
1916: PetscCall(PetscPostIrecvInt(comm, tagi, merge->nrecv, merge->id_r, len_ri, &buf_ri, &rwaits));
1917: PetscCall(PetscMalloc1(len, &buf_s));
1918: buf_si = buf_s; /* points to the beginning of k-th msg to be sent */
1919: for (proc = 0, k = 0; proc < size; proc++) {
1920: if (!len_s[proc]) continue;
1921: /* form outgoing message for i-structure:
1922: buf_si[0]: nrows to be sent
1923: [1:nrows]: row index (global)
1924: [nrows+1:2*nrows+1]: i-structure index
1925: */
1926: nrows = len_si[proc] / 2 - 1;
1927: buf_si_i = buf_si + nrows + 1;
1928: buf_si[0] = nrows;
1929: buf_si_i[0] = 0;
1930: nrows = 0;
1931: for (i = owners_co[proc]; i < owners_co[proc + 1]; i++) {
1932: nzi = coi[i + 1] - coi[i];
1933: buf_si_i[nrows + 1] = buf_si_i[nrows] + nzi; /* i-structure */
1934: buf_si[nrows + 1] = prmap[i] - owners[proc]; /* local row index */
1935: nrows++;
1936: }
1937: PetscCallMPI(MPIU_Isend(buf_si, len_si[proc], MPIU_INT, proc, tagi, comm, swaits + k));
1938: k++;
1939: buf_si += len_si[proc];
1940: }
1941: i = merge->nrecv;
1942: while (i--) {
1943: PETSC_UNUSED PetscMPIInt icompleted;
1944: PetscCallMPI(MPI_Waitany(merge->nrecv, rwaits, &icompleted, &rstatus));
1945: }
1946: PetscCall(PetscFree(rwaits));
1947: if (merge->nsend) PetscCallMPI(MPI_Waitall(merge->nsend, swaits, sstatus));
1948: PetscCall(PetscFree(len_si));
1949: PetscCall(PetscFree(len_ri));
1950: PetscCall(PetscFree(swaits));
1951: PetscCall(PetscFree(sstatus));
1952: PetscCall(PetscFree(buf_s));
1954: /* compute the local portion of C (mpi mat) */
1955: /* allocate bi array and free space for accumulating nonzero column info */
1956: PetscCall(PetscMalloc1(pn + 1, &bi));
1957: bi[0] = 0;
1959: /* set initial free space to be fill*(nnz(P) + nnz(AP)) */
1960: nnz = PetscRealIntMultTruncate(fill, PetscIntSumTruncate(pdti[pn], PetscIntSumTruncate(poti[pon], ai[am])));
1961: PetscCall(PetscFreeSpaceGet(nnz, &free_space));
1962: current_space = free_space;
1964: PetscCall(PetscMalloc3(merge->nrecv, &buf_ri_k, merge->nrecv, &nextrow, merge->nrecv, &nextci));
1965: for (k = 0; k < merge->nrecv; k++) {
1966: buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */
1967: nrows = *buf_ri_k[k];
1968: nextrow[k] = buf_ri_k[k] + 1; /* next row number of k-th recved i-structure */
1969: nextci[k] = buf_ri_k[k] + (nrows + 1); /* points to the next i-structure of k-th received i-structure */
1970: }
1972: PetscCall(PetscLLCondensedCreate_Scalable(Armax, &lnk));
1973: MatPreallocateBegin(comm, pn, A->cmap->n, dnz, onz);
1974: rmax = 0;
1975: for (i = 0; i < pn; i++) {
1976: /* add pdt[i,:]*AP into lnk */
1977: pnz = pdti[i + 1] - pdti[i];
1978: ptJ = pdtj + pdti[i];
1979: for (j = 0; j < pnz; j++) {
1980: row = ptJ[j]; /* row of AP == col of Pt */
1981: anz = ai[row + 1] - ai[row];
1982: Jptr = aj + ai[row];
1983: /* add non-zero cols of AP into the sorted linked list lnk */
1984: PetscCall(PetscLLCondensedAddSorted_Scalable(anz, Jptr, lnk));
1985: }
1987: /* add received col data into lnk */
1988: for (k = 0; k < merge->nrecv; k++) { /* k-th received message */
1989: if (i == *nextrow[k]) { /* i-th row */
1990: nzi = *(nextci[k] + 1) - *nextci[k];
1991: Jptr = buf_rj[k] + *nextci[k];
1992: PetscCall(PetscLLCondensedAddSorted_Scalable(nzi, Jptr, lnk));
1993: nextrow[k]++;
1994: nextci[k]++;
1995: }
1996: }
1998: /* add missing diagonal entry */
1999: if (C->force_diagonals) {
2000: k = i + owners[rank]; /* column index */
2001: PetscCall(PetscLLCondensedAddSorted_Scalable(1, &k, lnk));
2002: }
2004: nnz = lnk[0];
2006: /* if free space is not available, make more free space */
2007: if (current_space->local_remaining < nnz) {
2008: PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(nnz, current_space->total_array_size), ¤t_space));
2009: nspacedouble++;
2010: }
2011: /* copy data into free space, then initialize lnk */
2012: PetscCall(PetscLLCondensedClean_Scalable(nnz, current_space->array, lnk));
2013: PetscCall(MatPreallocateSet(i + owners[rank], nnz, current_space->array, dnz, onz));
2015: current_space->array += nnz;
2016: current_space->local_used += nnz;
2017: current_space->local_remaining -= nnz;
2019: bi[i + 1] = bi[i] + nnz;
2020: if (nnz > rmax) rmax = nnz;
2021: }
2022: PetscCall(PetscFree3(buf_ri_k, nextrow, nextci));
2024: PetscCall(PetscMalloc1(bi[pn], &bj));
2025: PetscCall(PetscFreeSpaceContiguous(&free_space, bj));
2026: afill_tmp = (PetscReal)bi[pn] / (pdti[pn] + poti[pon] + ai[am] + 1);
2027: if (afill_tmp > afill) afill = afill_tmp;
2028: PetscCall(PetscLLCondensedDestroy_Scalable(lnk));
2029: PetscCall(PetscHMapIDestroy(&ta));
2030: PetscCall(MatRestoreSymbolicTranspose_SeqAIJ(p->A, &pdti, &pdtj));
2031: PetscCall(MatRestoreSymbolicTranspose_SeqAIJ(p->B, &poti, &potj));
2033: /* create symbolic parallel matrix C - why cannot be assembled in Numeric part */
2034: PetscCall(MatSetSizes(C, pn, A->cmap->n, PETSC_DETERMINE, PETSC_DETERMINE));
2035: PetscCall(MatSetBlockSizes(C, P->cmap->bs, A->cmap->bs));
2036: PetscCall(MatGetType(A, &mtype));
2037: PetscCall(MatSetType(C, mtype));
2038: PetscCall(MatMPIAIJSetPreallocation(C, 0, dnz, 0, onz));
2039: MatPreallocateEnd(dnz, onz);
2040: PetscCall(MatSetBlockSize(C, 1));
2041: PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
2042: for (i = 0; i < pn; i++) {
2043: row = i + rstart;
2044: nnz = bi[i + 1] - bi[i];
2045: Jptr = bj + bi[i];
2046: PetscCall(MatSetValues(C, 1, &row, nnz, Jptr, NULL, INSERT_VALUES));
2047: }
2048: PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
2049: PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
2050: PetscCall(MatSetOption(C, MAT_NEW_NONZERO_LOCATION_ERR, PETSC_TRUE));
2051: merge->bi = bi;
2052: merge->bj = bj;
2053: merge->coi = coi;
2054: merge->coj = coj;
2055: merge->buf_ri = buf_ri;
2056: merge->buf_rj = buf_rj;
2057: merge->owners_co = owners_co;
2059: /* attach the supporting struct to C for reuse */
2060: C->product->data = ap;
2061: C->product->destroy = MatProductCtxDestroy_MPIAIJ_PtAP;
2062: ap->merge = merge;
2064: C->ops->mattransposemultnumeric = MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ;
2066: if (PetscDefined(USE_INFO)) {
2067: if (bi[pn] != 0) {
2068: PetscCall(PetscInfo(C, "Reallocs %" PetscInt_FMT "; Fill ratio: given %g needed %g.\n", nspacedouble, (double)fill, (double)afill));
2069: PetscCall(PetscInfo(C, "Use MatTransposeMatMult(A,B,MatReuse,%g,&C) for best performance.\n", (double)afill));
2070: } else PetscCall(PetscInfo(C, "Empty matrix product\n"));
2071: }
2072: PetscFunctionReturn(PETSC_SUCCESS);
2073: }
2075: static PetscErrorCode MatProductSymbolic_AtB_MPIAIJ_MPIAIJ(Mat C)
2076: {
2077: Mat_Product *product = C->product;
2078: Mat A = product->A, B = product->B;
2079: PetscReal fill = product->fill;
2080: PetscBool flg;
2082: PetscFunctionBegin;
2083: /* scalable */
2084: PetscCall(PetscStrcmp(product->alg, "scalable", &flg));
2085: if (flg) {
2086: PetscCall(MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ(A, B, fill, C));
2087: goto next;
2088: }
2090: /* nonscalable */
2091: PetscCall(PetscStrcmp(product->alg, "nonscalable", &flg));
2092: if (flg) {
2093: PetscCall(MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(A, B, fill, C));
2094: goto next;
2095: }
2097: /* matmatmult */
2098: PetscCall(PetscStrcmp(product->alg, "at*b", &flg));
2099: if (flg) {
2100: Mat At;
2101: MatProductCtx_APMPI *ptap;
2103: PetscCall(MatTranspose(A, MAT_INITIAL_MATRIX, &At));
2104: PetscCall(MatMatMultSymbolic_MPIAIJ_MPIAIJ(At, B, fill, C));
2105: ptap = (MatProductCtx_APMPI *)C->product->data;
2106: if (ptap) {
2107: ptap->Pt = At;
2108: C->product->destroy = MatProductCtxDestroy_MPIAIJ_PtAP;
2109: }
2110: C->ops->transposematmultnumeric = MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_matmatmult;
2111: goto next;
2112: }
2114: /* backend general code */
2115: PetscCall(PetscStrcmp(product->alg, "backend", &flg));
2116: if (flg) {
2117: PetscCall(MatProductSymbolic_MPIAIJBACKEND(C));
2118: PetscFunctionReturn(PETSC_SUCCESS);
2119: }
2121: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "MatProduct type is not supported");
2123: next:
2124: C->ops->productnumeric = MatProductNumeric_AtB;
2125: PetscFunctionReturn(PETSC_SUCCESS);
2126: }
2128: /* Set options for MatMatMultxxx_MPIAIJ_MPIAIJ */
2129: static PetscErrorCode MatProductSetFromOptions_MPIAIJ_AB(Mat C)
2130: {
2131: Mat_Product *product = C->product;
2132: Mat A = product->A, B = product->B;
2133: #if PetscDefined(HAVE_HYPRE)
2134: const char *algTypes[5] = {"scalable", "nonscalable", "seqmpi", "backend", "hypre"};
2135: PetscInt nalg = 5;
2136: #else
2137: const char *algTypes[4] = {
2138: "scalable",
2139: "nonscalable",
2140: "seqmpi",
2141: "backend",
2142: };
2143: PetscInt nalg = 4;
2144: #endif
2145: PetscInt alg = 1; /* set nonscalable algorithm as default */
2146: PetscBool flg;
2147: MPI_Comm comm;
2149: PetscFunctionBegin;
2150: PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
2152: /* Set "nonscalable" as default algorithm */
2153: PetscCall(PetscStrcmp(C->product->alg, "default", &flg));
2154: if (flg) {
2155: PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
2157: /* Set "scalable" as default if BN and local nonzeros of A and B are large */
2158: if (B->cmap->N > 100000) { /* may switch to scalable algorithm as default */
2159: MatInfo Ainfo, Binfo;
2160: PetscInt nz_local;
2161: PetscBool alg_scalable = PETSC_FALSE;
2163: PetscCall(MatGetInfo(A, MAT_LOCAL, &Ainfo));
2164: PetscCall(MatGetInfo(B, MAT_LOCAL, &Binfo));
2165: nz_local = (PetscInt)(Ainfo.nz_allocated + Binfo.nz_allocated);
2167: if (B->cmap->N > product->fill * nz_local) alg_scalable = PETSC_TRUE;
2168: PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &alg_scalable, 1, MPI_C_BOOL, MPI_LOR, comm));
2170: if (alg_scalable) {
2171: alg = 0; /* scalable algorithm would 50% slower than nonscalable algorithm */
2172: PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
2173: PetscCall(PetscInfo(B, "Use scalable algorithm, BN %" PetscInt_FMT ", fill*nz_allocated %g\n", B->cmap->N, (double)(product->fill * nz_local)));
2174: }
2175: }
2176: }
2178: /* Get runtime option */
2179: if (product->api_user) {
2180: PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatMatMult", "Mat");
2181: PetscCall(PetscOptionsEList("-matmatmult_via", "Algorithmic approach", "MatMatMult", algTypes, nalg, algTypes[alg], &alg, &flg));
2182: PetscOptionsEnd();
2183: } else {
2184: PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatProduct_AB", "Mat");
2185: PetscCall(PetscOptionsEList("-mat_product_algorithm", "Algorithmic approach", "MatMatMult", algTypes, nalg, algTypes[alg], &alg, &flg));
2186: PetscOptionsEnd();
2187: }
2188: if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
2190: C->ops->productsymbolic = MatProductSymbolic_AB_MPIAIJ_MPIAIJ;
2191: PetscFunctionReturn(PETSC_SUCCESS);
2192: }
2194: static PetscErrorCode MatProductSetFromOptions_MPIAIJ_ABt(Mat C)
2195: {
2196: PetscFunctionBegin;
2197: PetscCall(MatProductSetFromOptions_MPIAIJ_AB(C));
2198: C->ops->productsymbolic = MatProductSymbolic_ABt_MPIAIJ_MPIAIJ;
2199: PetscFunctionReturn(PETSC_SUCCESS);
2200: }
2202: /* Set options for MatTransposeMatMultXXX_MPIAIJ_MPIAIJ */
2203: static PetscErrorCode MatProductSetFromOptions_MPIAIJ_AtB(Mat C)
2204: {
2205: Mat_Product *product = C->product;
2206: Mat A = product->A, B = product->B;
2207: const char *algTypes[4] = {"scalable", "nonscalable", "at*b", "backend"};
2208: PetscInt nalg = 4;
2209: PetscInt alg = 1; /* set default algorithm */
2210: PetscBool flg;
2211: MPI_Comm comm;
2213: PetscFunctionBegin;
2214: /* Check matrix local sizes */
2215: PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
2216: PetscCheck(A->rmap->rstart == B->rmap->rstart && A->rmap->rend == B->rmap->rend, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Matrix local dimensions are incompatible, A (%" PetscInt_FMT ", %" PetscInt_FMT ") != B (%" PetscInt_FMT ",%" PetscInt_FMT ")",
2217: A->rmap->rstart, A->rmap->rend, B->rmap->rstart, B->rmap->rend);
2219: /* Set default algorithm */
2220: PetscCall(PetscStrcmp(C->product->alg, "default", &flg));
2221: if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
2223: /* Set "scalable" as default if BN and local nonzeros of A and B are large */
2224: if (alg && B->cmap->N > 100000) { /* may switch to scalable algorithm as default */
2225: MatInfo Ainfo, Binfo;
2226: PetscInt nz_local;
2227: PetscBool alg_scalable = PETSC_FALSE;
2229: PetscCall(MatGetInfo(A, MAT_LOCAL, &Ainfo));
2230: PetscCall(MatGetInfo(B, MAT_LOCAL, &Binfo));
2231: nz_local = (PetscInt)(Ainfo.nz_allocated + Binfo.nz_allocated);
2233: if (B->cmap->N > product->fill * nz_local) alg_scalable = PETSC_TRUE;
2234: PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &alg_scalable, 1, MPI_C_BOOL, MPI_LOR, comm));
2236: if (alg_scalable) {
2237: alg = 0; /* scalable algorithm would 50% slower than nonscalable algorithm */
2238: PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
2239: PetscCall(PetscInfo(B, "Use scalable algorithm, BN %" PetscInt_FMT ", fill*nz_allocated %g\n", B->cmap->N, (double)(product->fill * nz_local)));
2240: }
2241: }
2243: /* Get runtime option */
2244: if (product->api_user) {
2245: PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatTransposeMatMult", "Mat");
2246: PetscCall(PetscOptionsEList("-mattransposematmult_via", "Algorithmic approach", "MatTransposeMatMult", algTypes, nalg, algTypes[alg], &alg, &flg));
2247: PetscOptionsEnd();
2248: } else {
2249: PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatProduct_AtB", "Mat");
2250: PetscCall(PetscOptionsEList("-mat_product_algorithm", "Algorithmic approach", "MatTransposeMatMult", algTypes, nalg, algTypes[alg], &alg, &flg));
2251: PetscOptionsEnd();
2252: }
2253: if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
2255: C->ops->productsymbolic = MatProductSymbolic_AtB_MPIAIJ_MPIAIJ;
2256: PetscFunctionReturn(PETSC_SUCCESS);
2257: }
2259: static PetscErrorCode MatProductSetFromOptions_MPIAIJ_PtAP(Mat C)
2260: {
2261: Mat_Product *product = C->product;
2262: Mat A = product->A, P = product->B;
2263: MPI_Comm comm;
2264: PetscBool flg;
2265: PetscInt alg = 1; /* set default algorithm */
2266: #if !PetscDefined(HAVE_HYPRE)
2267: const char *algTypes[5] = {"scalable", "nonscalable", "allatonce", "allatonce_merged", "backend"};
2268: PetscInt nalg = 5;
2269: #else
2270: const char *algTypes[6] = {"scalable", "nonscalable", "allatonce", "allatonce_merged", "backend", "hypre"};
2271: PetscInt nalg = 6;
2272: #endif
2273: PetscInt pN = P->cmap->N;
2275: PetscFunctionBegin;
2276: /* Check matrix local sizes */
2277: PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
2278: PetscCheck(A->rmap->rstart == P->rmap->rstart && A->rmap->rend == P->rmap->rend, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Matrix local dimensions are incompatible, Arow (%" PetscInt_FMT ", %" PetscInt_FMT ") != Prow (%" PetscInt_FMT ",%" PetscInt_FMT ")",
2279: A->rmap->rstart, A->rmap->rend, P->rmap->rstart, P->rmap->rend);
2280: PetscCheck(A->cmap->rstart == P->rmap->rstart && A->cmap->rend == P->rmap->rend, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Matrix local dimensions are incompatible, Acol (%" PetscInt_FMT ", %" PetscInt_FMT ") != Prow (%" PetscInt_FMT ",%" PetscInt_FMT ")",
2281: A->cmap->rstart, A->cmap->rend, P->rmap->rstart, P->rmap->rend);
2283: /* Set "nonscalable" as default algorithm */
2284: PetscCall(PetscStrcmp(C->product->alg, "default", &flg));
2285: if (flg) {
2286: PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
2288: /* Set "scalable" as default if BN and local nonzeros of A and B are large */
2289: if (pN > 100000) {
2290: MatInfo Ainfo, Pinfo;
2291: PetscInt nz_local;
2292: PetscBool alg_scalable = PETSC_FALSE;
2294: PetscCall(MatGetInfo(A, MAT_LOCAL, &Ainfo));
2295: PetscCall(MatGetInfo(P, MAT_LOCAL, &Pinfo));
2296: nz_local = (PetscInt)(Ainfo.nz_allocated + Pinfo.nz_allocated);
2298: if (pN > product->fill * nz_local) alg_scalable = PETSC_TRUE;
2299: PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &alg_scalable, 1, MPI_C_BOOL, MPI_LOR, comm));
2301: if (alg_scalable) {
2302: alg = 0; /* scalable algorithm would 50% slower than nonscalable algorithm */
2303: PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
2304: }
2305: }
2306: }
2308: /* Get runtime option */
2309: if (product->api_user) {
2310: PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatPtAP", "Mat");
2311: PetscCall(PetscOptionsEList("-matptap_via", "Algorithmic approach", "MatPtAP", algTypes, nalg, algTypes[alg], &alg, &flg));
2312: PetscOptionsEnd();
2313: } else {
2314: PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatProduct_PtAP", "Mat");
2315: PetscCall(PetscOptionsEList("-mat_product_algorithm", "Algorithmic approach", "MatPtAP", algTypes, nalg, algTypes[alg], &alg, &flg));
2316: PetscOptionsEnd();
2317: }
2318: if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
2320: C->ops->productsymbolic = MatProductSymbolic_PtAP_MPIAIJ_MPIAIJ;
2321: PetscFunctionReturn(PETSC_SUCCESS);
2322: }
2324: static PetscErrorCode MatProductSetFromOptions_MPIAIJ_RARt(Mat C)
2325: {
2326: Mat_Product *product = C->product;
2327: Mat A = product->A, R = product->B;
2329: PetscFunctionBegin;
2330: /* Check matrix local sizes */
2331: PetscCheck(A->cmap->n == R->cmap->n && A->rmap->n == R->cmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Matrix local dimensions are incompatible, A local (%" PetscInt_FMT ", %" PetscInt_FMT "), R local (%" PetscInt_FMT ",%" PetscInt_FMT ")", A->rmap->n,
2332: A->rmap->n, R->rmap->n, R->cmap->n);
2334: C->ops->productsymbolic = MatProductSymbolic_RARt_MPIAIJ_MPIAIJ;
2335: PetscFunctionReturn(PETSC_SUCCESS);
2336: }
2338: /*
2339: Set options for ABC = A*B*C = A*(B*C); ABC's algorithm must be chosen from AB's algorithm
2340: */
2341: static PetscErrorCode MatProductSetFromOptions_MPIAIJ_ABC(Mat C)
2342: {
2343: Mat_Product *product = C->product;
2344: PetscBool flg = PETSC_FALSE;
2345: PetscInt alg = 1; /* default algorithm */
2346: const char *algTypes[3] = {"scalable", "nonscalable", "seqmpi"};
2347: PetscInt nalg = 3;
2349: PetscFunctionBegin;
2350: /* Set default algorithm */
2351: PetscCall(PetscStrcmp(C->product->alg, "default", &flg));
2352: if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
2354: /* Get runtime option */
2355: if (product->api_user) {
2356: PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatMatMatMult", "Mat");
2357: PetscCall(PetscOptionsEList("-matmatmatmult_via", "Algorithmic approach", "MatMatMatMult", algTypes, nalg, algTypes[alg], &alg, &flg));
2358: PetscOptionsEnd();
2359: } else {
2360: PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatProduct_ABC", "Mat");
2361: PetscCall(PetscOptionsEList("-mat_product_algorithm", "Algorithmic approach", "MatProduct_ABC", algTypes, nalg, algTypes[alg], &alg, &flg));
2362: PetscOptionsEnd();
2363: }
2364: if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
2366: C->ops->matmatmultsymbolic = MatMatMatMultSymbolic_MPIAIJ_MPIAIJ_MPIAIJ;
2367: C->ops->productsymbolic = MatProductSymbolic_ABC;
2368: PetscFunctionReturn(PETSC_SUCCESS);
2369: }
2371: PETSC_INTERN PetscErrorCode MatProductSetFromOptions_MPIAIJ(Mat C)
2372: {
2373: Mat_Product *product = C->product;
2375: PetscFunctionBegin;
2376: switch (product->type) {
2377: case MATPRODUCT_AB:
2378: PetscCall(MatProductSetFromOptions_MPIAIJ_AB(C));
2379: break;
2380: case MATPRODUCT_ABt:
2381: PetscCall(MatProductSetFromOptions_MPIAIJ_ABt(C));
2382: break;
2383: case MATPRODUCT_AtB:
2384: PetscCall(MatProductSetFromOptions_MPIAIJ_AtB(C));
2385: break;
2386: case MATPRODUCT_PtAP:
2387: PetscCall(MatProductSetFromOptions_MPIAIJ_PtAP(C));
2388: break;
2389: case MATPRODUCT_RARt:
2390: PetscCall(MatProductSetFromOptions_MPIAIJ_RARt(C));
2391: break;
2392: case MATPRODUCT_ABC:
2393: PetscCall(MatProductSetFromOptions_MPIAIJ_ABC(C));
2394: break;
2395: default:
2396: break;
2397: }
2398: PetscFunctionReturn(PETSC_SUCCESS);
2399: }