Actual source code: mpidense.c

  1: /*
  2:    Basic functions for basic parallel dense matrices.
  3:    Portions of this code are under:
  4:    Copyright (c) 2022 Advanced Micro Devices, Inc. All rights reserved.
  5: */

  7: #include <../src/mat/impls/dense/mpi/mpidense.h>
  8: #include <../src/mat/impls/aij/mpi/mpiaij.h>
  9: #include <petscblaslapack.h>
 10: #include <petsc/private/vecimpl.h>
 11: #include <petsc/private/deviceimpl.h>
 12: #include <petsc/private/sfimpl.h>

 14: /*@
 15:   MatDenseGetLocalMatrix - For a `MATMPIDENSE` or `MATSEQDENSE` matrix returns the sequential
 16:   matrix that represents the operator. For sequential matrices it returns itself.

 18:   Input Parameter:
 19: . A - the sequential or MPI `MATDENSE` matrix

 21:   Output Parameter:
 22: . B - the inner matrix

 24:   Level: intermediate

 26: .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MATMPIDENSE`, `MATSEQDENSE`
 27: @*/
 28: PetscErrorCode MatDenseGetLocalMatrix(Mat A, Mat *B)
 29: {
 30:   Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
 31:   PetscBool     flg;

 33:   PetscFunctionBegin;
 35:   PetscAssertPointer(B, 2);
 36:   PetscCall(PetscObjectBaseTypeCompare((PetscObject)A, MATMPIDENSE, &flg));
 37:   if (flg) *B = mat->A;
 38:   else {
 39:     PetscCall(PetscObjectBaseTypeCompare((PetscObject)A, MATSEQDENSE, &flg));
 40:     PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "Not for matrix type %s", ((PetscObject)A)->type_name);
 41:     *B = A;
 42:   }
 43:   PetscFunctionReturn(PETSC_SUCCESS);
 44: }

 46: static PetscErrorCode MatCopy_MPIDense(Mat A, Mat B, MatStructure s)
 47: {
 48:   Mat_MPIDense *Amat = (Mat_MPIDense *)A->data;
 49:   Mat_MPIDense *Bmat = (Mat_MPIDense *)B->data;

 51:   PetscFunctionBegin;
 52:   PetscCall(MatCopy(Amat->A, Bmat->A, s));
 53:   PetscFunctionReturn(PETSC_SUCCESS);
 54: }

 56: PetscErrorCode MatShift_MPIDense(Mat A, PetscScalar alpha)
 57: {
 58:   Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
 59:   PetscInt      j, lda, rstart = A->rmap->rstart, rend = A->rmap->rend, rend2;
 60:   PetscScalar  *v;

 62:   PetscFunctionBegin;
 63:   PetscCall(MatDenseGetArray(mat->A, &v));
 64:   PetscCall(MatDenseGetLDA(mat->A, &lda));
 65:   rend2 = PetscMin(rend, A->cmap->N);
 66:   if (rend2 > rstart) {
 67:     for (j = rstart; j < rend2; j++) v[j - rstart + j * lda] += alpha;
 68:     PetscCall(PetscLogFlops(rend2 - rstart));
 69:   }
 70:   PetscCall(MatDenseRestoreArray(mat->A, &v));
 71:   PetscFunctionReturn(PETSC_SUCCESS);
 72: }

 74: static PetscErrorCode MatGetRow_MPIDense(Mat A, PetscInt row, PetscInt *nz, PetscInt **idx, PetscScalar **v)
 75: {
 76:   Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
 77:   PetscInt      lrow, rstart = A->rmap->rstart, rend = A->rmap->rend;

 79:   PetscFunctionBegin;
 80:   PetscCheck(row >= rstart && row < rend, PETSC_COMM_SELF, PETSC_ERR_SUP, "only local rows");
 81:   lrow = row - rstart;
 82:   PetscCall(MatGetRow(mat->A, lrow, nz, (const PetscInt **)idx, (const PetscScalar **)v));
 83:   PetscFunctionReturn(PETSC_SUCCESS);
 84: }

 86: static PetscErrorCode MatRestoreRow_MPIDense(Mat A, PetscInt row, PetscInt *nz, PetscInt **idx, PetscScalar **v)
 87: {
 88:   Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
 89:   PetscInt      lrow, rstart = A->rmap->rstart, rend = A->rmap->rend;

 91:   PetscFunctionBegin;
 92:   PetscCheck(row >= rstart && row < rend, PETSC_COMM_SELF, PETSC_ERR_SUP, "only local rows");
 93:   lrow = row - rstart;
 94:   PetscCall(MatRestoreRow(mat->A, lrow, nz, (const PetscInt **)idx, (const PetscScalar **)v));
 95:   PetscFunctionReturn(PETSC_SUCCESS);
 96: }

 98: static PetscErrorCode MatGetDiagonalBlock_MPIDense(Mat A, Mat *a)
 99: {
100:   Mat_MPIDense *mdn = (Mat_MPIDense *)A->data;
101:   PetscInt      m = A->rmap->n, rstart = A->rmap->rstart;
102:   PetscScalar  *array;
103:   MPI_Comm      comm;
104:   PetscBool     flg;
105:   Mat           B;

107:   PetscFunctionBegin;
108:   PetscCall(MatHasCongruentLayouts(A, &flg));
109:   PetscCheck(flg, PETSC_COMM_SELF, PETSC_ERR_SUP, "Only square matrices supported.");
110:   PetscCall(PetscObjectQuery((PetscObject)A, "DiagonalBlock", (PetscObject *)&B));
111:   if (!B) { /* This should use MatDenseGetSubMatrix (not create), but we would need a call like MatRestoreDiagonalBlock */
112: #if PetscDefined(HAVE_CUDA)
113:     PetscCall(PetscObjectTypeCompare((PetscObject)mdn->A, MATSEQDENSECUDA, &flg));
114:     PetscCheck(!flg, PETSC_COMM_SELF, PETSC_ERR_SUP, "Not coded for %s. Send an email to petsc-dev@mcs.anl.gov to request this feature", MATSEQDENSECUDA);
115: #elif PetscDefined(HAVE_HIP)
116:     PetscCall(PetscObjectTypeCompare((PetscObject)mdn->A, MATSEQDENSEHIP, &flg));
117:     PetscCheck(!flg, PETSC_COMM_SELF, PETSC_ERR_SUP, "Not coded for %s. Send an email to petsc-dev@mcs.anl.gov to request this feature", MATSEQDENSEHIP);
118: #endif
119:     PetscCall(PetscObjectGetComm((PetscObject)mdn->A, &comm));
120:     PetscCall(MatCreate(comm, &B));
121:     PetscCall(MatSetSizes(B, m, m, m, m));
122:     PetscCall(MatSetType(B, ((PetscObject)mdn->A)->type_name));
123:     PetscCall(MatDenseGetArrayRead(mdn->A, (const PetscScalar **)&array));
124:     PetscCall(MatSeqDenseSetPreallocation(B, array + m * rstart));
125:     PetscCall(MatDenseRestoreArrayRead(mdn->A, (const PetscScalar **)&array));
126:     PetscCall(PetscObjectCompose((PetscObject)A, "DiagonalBlock", (PetscObject)B));
127:     *a = B;
128:     PetscCall(MatDestroy(&B));
129:   } else *a = B;
130:   PetscFunctionReturn(PETSC_SUCCESS);
131: }

133: static PetscErrorCode MatSetValues_MPIDense(Mat mat, PetscInt m, const PetscInt idxm[], PetscInt n, const PetscInt idxn[], const PetscScalar v[], InsertMode addv)
134: {
135:   Mat_MPIDense *A = (Mat_MPIDense *)mat->data;
136:   PetscInt      i, j, rstart = mat->rmap->rstart, rend = mat->rmap->rend, row;
137:   PetscBool     roworiented = A->roworiented;

139:   PetscFunctionBegin;
140:   for (i = 0; i < m; i++) {
141:     if (idxm[i] < 0) continue;
142:     PetscCheck(idxm[i] < mat->rmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Row too large");
143:     if (idxm[i] >= rstart && idxm[i] < rend) {
144:       row = idxm[i] - rstart;
145:       if (roworiented) {
146:         PetscCall(MatSetValues(A->A, 1, &row, n, idxn, PetscSafePointerPlusOffset(v, i * n), addv));
147:       } else {
148:         for (j = 0; j < n; j++) {
149:           if (idxn[j] < 0) continue;
150:           PetscCheck(idxn[j] < mat->cmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Column too large");
151:           PetscCall(MatSetValues(A->A, 1, &row, 1, &idxn[j], PetscSafePointerPlusOffset(v, i + j * m), addv));
152:         }
153:       }
154:     } else if (!A->donotstash) {
155:       mat->assembled = PETSC_FALSE;
156:       if (roworiented) {
157:         PetscCall(MatStashValuesRow_Private(&mat->stash, idxm[i], n, idxn, PetscSafePointerPlusOffset(v, i * n), PETSC_FALSE));
158:       } else {
159:         PetscCall(MatStashValuesCol_Private(&mat->stash, idxm[i], n, idxn, PetscSafePointerPlusOffset(v, i), m, PETSC_FALSE));
160:       }
161:     }
162:   }
163:   PetscFunctionReturn(PETSC_SUCCESS);
164: }

166: static PetscErrorCode MatGetValues_MPIDense(Mat mat, PetscInt m, const PetscInt idxm[], PetscInt n, const PetscInt idxn[], PetscScalar v[])
167: {
168:   Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
169:   PetscInt      i, j, rstart = mat->rmap->rstart, rend = mat->rmap->rend, row;
170:   PetscBool     roworiented = mdn->roworiented;
171:   PetscScalar  *value;

173:   PetscFunctionBegin;
174:   for (i = 0; i < m; i++) {
175:     if (idxm[i] < 0) continue; /* negative row */
176:     PetscCheck(idxm[i] < mat->rmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Row too large");
177:     PetscCheck(idxm[i] >= rstart && idxm[i] < rend, PETSC_COMM_SELF, PETSC_ERR_SUP, "Only local values currently supported");
178:     row = idxm[i] - rstart;
179:     for (j = 0; j < n; j++) {
180:       if (idxn[j] < 0) continue; /* negative column */
181:       PetscCheck(idxn[j] < mat->cmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Column too large");
182:       value = roworiented ? &v[j + i * n] : &v[i + j * m];
183:       PetscCall(MatGetValues(mdn->A, 1, &row, 1, &idxn[j], value));
184:     }
185:   }
186:   PetscFunctionReturn(PETSC_SUCCESS);
187: }

189: static PetscErrorCode MatDenseGetLDA_MPIDense(Mat A, PetscInt *lda)
190: {
191:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;

193:   PetscFunctionBegin;
194:   PetscCall(MatDenseGetLDA(a->A, lda));
195:   PetscFunctionReturn(PETSC_SUCCESS);
196: }

198: static PetscErrorCode MatDenseSetLDA_MPIDense(Mat A, PetscInt lda)
199: {
200:   Mat_MPIDense *a     = (Mat_MPIDense *)A->data;
201:   MatType       mtype = MATSEQDENSE;

203:   PetscFunctionBegin;
204:   if (!a->A) {
205:     PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
206:     PetscCall(PetscLayoutSetUp(A->rmap));
207:     PetscCall(PetscLayoutSetUp(A->cmap));
208:     PetscCall(MatCreate(PETSC_COMM_SELF, &a->A));
209:     PetscCall(MatSetSizes(a->A, A->rmap->n, A->cmap->N, A->rmap->n, A->cmap->N));
210: #if PetscDefined(HAVE_CUDA)
211:     PetscBool iscuda;
212:     PetscCall(PetscObjectTypeCompare((PetscObject)A, MATMPIDENSECUDA, &iscuda));
213:     if (iscuda) mtype = MATSEQDENSECUDA;
214: #elif PetscDefined(HAVE_HIP)
215:     PetscBool iship;
216:     PetscCall(PetscObjectTypeCompare((PetscObject)A, MATMPIDENSEHIP, &iship));
217:     if (iship) mtype = MATSEQDENSEHIP;
218: #endif
219:     PetscCall(MatSetType(a->A, mtype));
220:   }
221:   PetscCall(MatDenseSetLDA(a->A, lda));
222:   PetscFunctionReturn(PETSC_SUCCESS);
223: }

225: static PetscErrorCode MatDenseGetArray_MPIDense(Mat A, PetscScalar **array)
226: {
227:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;

229:   PetscFunctionBegin;
230:   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
231:   PetscCall(MatDenseGetArray(a->A, array));
232:   PetscFunctionReturn(PETSC_SUCCESS);
233: }

235: static PetscErrorCode MatDenseGetArrayRead_MPIDense(Mat A, PetscScalar **array)
236: {
237:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;

239:   PetscFunctionBegin;
240:   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
241:   PetscCall(MatDenseGetArrayRead(a->A, (const PetscScalar **)array));
242:   PetscFunctionReturn(PETSC_SUCCESS);
243: }

245: static PetscErrorCode MatDenseGetArrayWrite_MPIDense(Mat A, PetscScalar **array)
246: {
247:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;

249:   PetscFunctionBegin;
250:   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
251:   PetscCall(MatDenseGetArrayWrite(a->A, array));
252:   PetscFunctionReturn(PETSC_SUCCESS);
253: }

255: static PetscErrorCode MatDensePlaceArray_MPIDense(Mat A, const PetscScalar *array)
256: {
257:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;

259:   PetscFunctionBegin;
260:   PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
261:   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
262:   PetscCall(MatDensePlaceArray(a->A, array));
263:   PetscFunctionReturn(PETSC_SUCCESS);
264: }

266: static PetscErrorCode MatDenseResetArray_MPIDense(Mat A)
267: {
268:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;

270:   PetscFunctionBegin;
271:   PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
272:   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
273:   PetscCall(MatDenseResetArray(a->A));
274:   PetscFunctionReturn(PETSC_SUCCESS);
275: }

277: static PetscErrorCode MatDenseReplaceArray_MPIDense(Mat A, const PetscScalar *array)
278: {
279:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;

281:   PetscFunctionBegin;
282:   PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
283:   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
284:   PetscCall(MatDenseReplaceArray(a->A, array));
285:   PetscFunctionReturn(PETSC_SUCCESS);
286: }

288: static PetscErrorCode MatCreateSubMatrix_MPIDense(Mat A, IS isrow, IS iscol, MatReuse scall, Mat *B)
289: {
290:   Mat_MPIDense      *mat = (Mat_MPIDense *)A->data, *newmatd;
291:   PetscInt           lda, i, j, rstart, rend, nrows, ncols, Ncols, nlrows, nlcols;
292:   const PetscInt    *irow, *icol;
293:   const PetscScalar *v;
294:   PetscScalar       *bv;
295:   Mat                newmat;
296:   IS                 iscol_local;
297:   MPI_Comm           comm_is, comm_mat;

299:   PetscFunctionBegin;
300:   PetscCall(PetscObjectGetComm((PetscObject)A, &comm_mat));
301:   PetscCall(PetscObjectGetComm((PetscObject)iscol, &comm_is));
302:   PetscCheck(comm_mat == comm_is, PETSC_COMM_SELF, PETSC_ERR_ARG_NOTSAMECOMM, "IS communicator must match matrix communicator");

304:   PetscCall(ISAllGather(iscol, &iscol_local));
305:   PetscCall(ISGetIndices(isrow, &irow));
306:   PetscCall(ISGetIndices(iscol_local, &icol));
307:   PetscCall(ISGetLocalSize(isrow, &nrows));
308:   PetscCall(ISGetLocalSize(iscol, &ncols));
309:   PetscCall(ISGetSize(iscol, &Ncols)); /* global number of columns, size of iscol_local */

311:   /* No parallel redistribution currently supported! Should really check each index set
312:      to confirm that it is OK.  ... Currently supports only submatrix same partitioning as
313:      original matrix! */

315:   PetscCall(MatGetLocalSize(A, &nlrows, &nlcols));
316:   PetscCall(MatGetOwnershipRange(A, &rstart, &rend));

318:   /* Check submatrix call */
319:   if (scall == MAT_REUSE_MATRIX) {
320:     /* SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Reused submatrix wrong size"); */
321:     /* Really need to test rows and column sizes! */
322:     newmat = *B;
323:   } else {
324:     /* Create and fill new matrix */
325:     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &newmat));
326:     PetscCall(MatSetSizes(newmat, nrows, ncols, PETSC_DECIDE, Ncols));
327:     PetscCall(MatSetType(newmat, ((PetscObject)A)->type_name));
328:     PetscCall(MatMPIDenseSetPreallocation(newmat, NULL));
329:   }

331:   /* Now extract the data pointers and do the copy, column at a time */
332:   newmatd = (Mat_MPIDense *)newmat->data;
333:   PetscCall(MatDenseGetArray(newmatd->A, &bv));
334:   PetscCall(MatDenseGetArrayRead(mat->A, &v));
335:   PetscCall(MatDenseGetLDA(mat->A, &lda));
336:   for (i = 0; i < Ncols; i++) {
337:     const PetscScalar *av = v + lda * icol[i];
338:     for (j = 0; j < nrows; j++) *bv++ = av[irow[j] - rstart];
339:   }
340:   PetscCall(MatDenseRestoreArrayRead(mat->A, &v));
341:   PetscCall(MatDenseRestoreArray(newmatd->A, &bv));

343:   /* Assemble the matrices so that the correct flags are set */
344:   PetscCall(MatAssemblyBegin(newmat, MAT_FINAL_ASSEMBLY));
345:   PetscCall(MatAssemblyEnd(newmat, MAT_FINAL_ASSEMBLY));

347:   /* Free work space */
348:   PetscCall(ISRestoreIndices(isrow, &irow));
349:   PetscCall(ISRestoreIndices(iscol_local, &icol));
350:   PetscCall(ISDestroy(&iscol_local));
351:   *B = newmat;
352:   PetscFunctionReturn(PETSC_SUCCESS);
353: }

355: static PetscErrorCode MatDenseRestoreArray_MPIDense(Mat A, PetscScalar **array)
356: {
357:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;

359:   PetscFunctionBegin;
360:   PetscCall(MatDenseRestoreArray(a->A, array));
361:   PetscFunctionReturn(PETSC_SUCCESS);
362: }

364: static PetscErrorCode MatDenseRestoreArrayRead_MPIDense(Mat A, PetscScalar **array)
365: {
366:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;

368:   PetscFunctionBegin;
369:   PetscCall(MatDenseRestoreArrayRead(a->A, (const PetscScalar **)array));
370:   PetscFunctionReturn(PETSC_SUCCESS);
371: }

373: static PetscErrorCode MatDenseRestoreArrayWrite_MPIDense(Mat A, PetscScalar **array)
374: {
375:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;

377:   PetscFunctionBegin;
378:   PetscCall(MatDenseRestoreArrayWrite(a->A, array));
379:   PetscFunctionReturn(PETSC_SUCCESS);
380: }

382: static PetscErrorCode MatAssemblyBegin_MPIDense(Mat mat, MatAssemblyType mode)
383: {
384:   Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
385:   PetscInt      nstash, reallocs;

387:   PetscFunctionBegin;
388:   if (mdn->donotstash || mat->nooffprocentries) PetscFunctionReturn(PETSC_SUCCESS);

390:   PetscCall(MatStashScatterBegin_Private(mat, &mat->stash, mat->rmap->range));
391:   PetscCall(MatStashGetInfo_Private(&mat->stash, &nstash, &reallocs));
392:   PetscCall(PetscInfo(mdn->A, "Stash has %" PetscInt_FMT " entries, uses %" PetscInt_FMT " mallocs.\n", nstash, reallocs));
393:   PetscFunctionReturn(PETSC_SUCCESS);
394: }

396: static PetscErrorCode MatAssemblyEnd_MPIDense(Mat mat, MatAssemblyType mode)
397: {
398:   Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
399:   PetscInt      i, *row, *col, flg, j, rstart, ncols;
400:   PetscMPIInt   n;
401:   PetscScalar  *val;

403:   PetscFunctionBegin;
404:   if (!mdn->donotstash && !mat->nooffprocentries) {
405:     /*  wait on receives */
406:     while (1) {
407:       PetscCall(MatStashScatterGetMesg_Private(&mat->stash, &n, &row, &col, &val, &flg));
408:       if (!flg) break;

410:       for (i = 0; i < n;) {
411:         /* Now identify the consecutive vals belonging to the same row */
412:         for (j = i, rstart = row[j]; j < n; j++) {
413:           if (row[j] != rstart) break;
414:         }
415:         if (j < n) ncols = j - i;
416:         else ncols = n - i;
417:         /* Now assemble all these values with a single function call */
418:         PetscCall(MatSetValues_MPIDense(mat, 1, row + i, ncols, col + i, val + i, mat->insertmode));
419:         i = j;
420:       }
421:     }
422:     PetscCall(MatStashScatterEnd_Private(&mat->stash));
423:   }

425:   PetscCall(MatAssemblyBegin(mdn->A, mode));
426:   PetscCall(MatAssemblyEnd(mdn->A, mode));
427:   PetscFunctionReturn(PETSC_SUCCESS);
428: }

430: static PetscErrorCode MatZeroEntries_MPIDense(Mat A)
431: {
432:   Mat_MPIDense *l = (Mat_MPIDense *)A->data;

434:   PetscFunctionBegin;
435:   PetscCall(MatZeroEntries(l->A));
436:   PetscFunctionReturn(PETSC_SUCCESS);
437: }

439: static PetscErrorCode MatZeroRows_MPIDense(Mat A, PetscInt n, const PetscInt rows[], PetscScalar diag, Vec x, Vec b)
440: {
441:   Mat_MPIDense *l = (Mat_MPIDense *)A->data;
442:   PetscInt      i, len, *lrows;

444:   PetscFunctionBegin;
445:   /* get locally owned rows */
446:   PetscCall(PetscLayoutMapLocal(A->rmap, n, rows, &len, &lrows, NULL));
447:   /* fix right-hand side if needed */
448:   if (x && b) {
449:     const PetscScalar *xx;
450:     PetscScalar       *bb;

452:     PetscCall(VecGetArrayRead(x, &xx));
453:     PetscCall(VecGetArrayWrite(b, &bb));
454:     for (i = 0; i < len; ++i) bb[lrows[i]] = diag * xx[lrows[i]];
455:     PetscCall(VecRestoreArrayRead(x, &xx));
456:     PetscCall(VecRestoreArrayWrite(b, &bb));
457:   }
458:   PetscCall(MatZeroRows(l->A, len, lrows, 0.0, NULL, NULL));
459:   if (diag != 0.0) {
460:     Vec d;

462:     PetscCall(MatCreateVecs(A, NULL, &d));
463:     PetscCall(VecSet(d, diag));
464:     PetscCall(MatDiagonalSet(A, d, INSERT_VALUES));
465:     PetscCall(VecDestroy(&d));
466:   }
467:   PetscCall(PetscFree(lrows));
468:   PetscFunctionReturn(PETSC_SUCCESS);
469: }

471: PETSC_INTERN PetscErrorCode MatMult_SeqDense(Mat, Vec, Vec);
472: PETSC_INTERN PetscErrorCode MatMultAdd_SeqDense(Mat, Vec, Vec, Vec);
473: PETSC_INTERN PetscErrorCode MatMultTranspose_SeqDense(Mat, Vec, Vec);
474: PETSC_INTERN PetscErrorCode MatMultTransposeAdd_SeqDense(Mat, Vec, Vec, Vec);

476: static PetscErrorCode MatMultColumnRange_MPIDense(Mat mat, Vec xx, Vec yy, PetscInt c_start, PetscInt c_end)
477: {
478:   Mat_MPIDense      *mdn = (Mat_MPIDense *)mat->data;
479:   const PetscScalar *ax;
480:   PetscScalar       *ay;
481:   PetscMemType       axmtype, aymtype;

483:   PetscFunctionBegin;
484:   if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat));
485:   PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype));
486:   PetscCall(VecGetArrayWriteAndMemType(mdn->lvec, &ay, &aymtype));
487:   PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE));
488:   PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE));
489:   PetscCall(VecRestoreArrayWriteAndMemType(mdn->lvec, &ay));
490:   PetscCall(VecRestoreArrayReadAndMemType(xx, &ax));
491:   PetscUseMethod(mdn->A, "MatMultColumnRange_C", (Mat, Vec, Vec, PetscInt, PetscInt), (mdn->A, mdn->lvec, yy, c_start, c_end));
492:   PetscFunctionReturn(PETSC_SUCCESS);
493: }

495: static PetscErrorCode MatMult_MPIDense(Mat mat, Vec xx, Vec yy)
496: {
497:   Mat_MPIDense      *mdn = (Mat_MPIDense *)mat->data;
498:   const PetscScalar *ax;
499:   PetscScalar       *ay;
500:   PetscMemType       axmtype, aymtype;

502:   PetscFunctionBegin;
503:   if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat));
504:   PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype));
505:   PetscCall(VecGetArrayWriteAndMemType(mdn->lvec, &ay, &aymtype));
506:   PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE));
507:   PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE));
508:   PetscCall(VecRestoreArrayWriteAndMemType(mdn->lvec, &ay));
509:   PetscCall(VecRestoreArrayReadAndMemType(xx, &ax));
510:   PetscUseTypeMethod(mdn->A, mult, mdn->lvec, yy);
511:   PetscFunctionReturn(PETSC_SUCCESS);
512: }

514: static PetscErrorCode MatGetMultPetscSF_MPIDense(Mat A, PetscSF *sf)
515: {
516:   Mat_MPIDense *mdn = (Mat_MPIDense *)A->data;

518:   PetscFunctionBegin;
519:   *sf = mdn->Mvctx;
520:   PetscFunctionReturn(PETSC_SUCCESS);
521: }

523: static PetscErrorCode MatMultAddColumnRange_MPIDense(Mat mat, Vec xx, Vec yy, Vec zz, PetscInt c_start, PetscInt c_end)
524: {
525:   Mat_MPIDense      *mdn = (Mat_MPIDense *)mat->data;
526:   const PetscScalar *ax;
527:   PetscScalar       *ay;
528:   PetscMemType       axmtype, aymtype;

530:   PetscFunctionBegin;
531:   if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat));
532:   PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype));
533:   PetscCall(VecGetArrayAndMemType(mdn->lvec, &ay, &aymtype));
534:   PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE));
535:   PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE));
536:   PetscCall(VecRestoreArrayAndMemType(mdn->lvec, &ay));
537:   PetscCall(VecRestoreArrayReadAndMemType(xx, &ax));
538:   PetscUseMethod(mdn->A, "MatMultAddColumnRange_C", (Mat, Vec, Vec, Vec, PetscInt, PetscInt), (mdn->A, mdn->lvec, yy, zz, c_start, c_end));
539:   PetscFunctionReturn(PETSC_SUCCESS);
540: }

542: static PetscErrorCode MatMultAdd_MPIDense(Mat mat, Vec xx, Vec yy, Vec zz)
543: {
544:   Mat_MPIDense      *mdn = (Mat_MPIDense *)mat->data;
545:   const PetscScalar *ax;
546:   PetscScalar       *ay;
547:   PetscMemType       axmtype, aymtype;

549:   PetscFunctionBegin;
550:   if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat));
551:   PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype));
552:   PetscCall(VecGetArrayAndMemType(mdn->lvec, &ay, &aymtype));
553:   PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE));
554:   PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE));
555:   PetscCall(VecRestoreArrayAndMemType(mdn->lvec, &ay));
556:   PetscCall(VecRestoreArrayReadAndMemType(xx, &ax));
557:   PetscUseTypeMethod(mdn->A, multadd, mdn->lvec, yy, zz);
558:   PetscFunctionReturn(PETSC_SUCCESS);
559: }

561: static PetscErrorCode MatMultHermitianTransposeColumnRange_MPIDense(Mat A, Vec xx, Vec yy, PetscInt c_start, PetscInt c_end)
562: {
563:   Mat_MPIDense      *a = (Mat_MPIDense *)A->data;
564:   const PetscScalar *ax;
565:   PetscScalar       *ay;
566:   PetscMemType       axmtype, aymtype;
567:   PetscInt           r_start, r_end;
568:   PetscInt           c_start_local, c_end_local;

570:   PetscFunctionBegin;
571:   if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
572:   PetscCall(VecZeroEntries(a->lvec));
573:   PetscCall(VecGetOwnershipRange(yy, &r_start, &r_end));
574:   c_start_local = PetscMax(c_start, r_start);
575:   c_end_local   = PetscMin(c_end, r_end);
576:   PetscCall(VecGetArrayAndMemType(yy, &ay, &aymtype));
577:   if (c_end_local > c_start_local) {
578:     if (PetscMemTypeHost(aymtype)) {
579:       PetscCall(PetscArrayzero(&ay[c_start_local], (size_t)(c_end_local - c_start_local)));
580:     } else {
581:       PetscCall(PetscDeviceRegisterMemory(ay, aymtype, sizeof(*ay) * ((size_t)(r_end - r_start))));
582:       PetscCall(PetscDeviceArrayZero(NULL, &ay[c_start_local], (size_t)(c_end_local - c_start_local)));
583:     }
584:   }
585:   PetscUseMethod(a->A, "MatMultHermitianTransposeColumnRange_C", (Mat, Vec, Vec, PetscInt, PetscInt), (a->A, xx, a->lvec, c_start, c_end));
586:   PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype));
587:   PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM));
588:   PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM));
589:   PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax));
590:   PetscCall(VecRestoreArrayAndMemType(yy, &ay));
591:   PetscFunctionReturn(PETSC_SUCCESS);
592: }

594: static PetscErrorCode MatMultTransposeKernel_MPIDense(Mat A, Vec xx, Vec yy, PetscBool herm)
595: {
596:   Mat_MPIDense      *a = (Mat_MPIDense *)A->data;
597:   const PetscScalar *ax;
598:   PetscScalar       *ay;
599:   PetscMemType       axmtype, aymtype;

601:   PetscFunctionBegin;
602:   if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
603:   PetscCall(VecSet(yy, 0.0));
604:   if (herm) PetscUseTypeMethod(a->A, multhermitiantranspose, xx, a->lvec);
605:   else PetscUseTypeMethod(a->A, multtranspose, xx, a->lvec);
606:   PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype));
607:   PetscCall(VecGetArrayAndMemType(yy, &ay, &aymtype));
608:   PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM));
609:   PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM));
610:   PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax));
611:   PetscCall(VecRestoreArrayAndMemType(yy, &ay));
612:   PetscFunctionReturn(PETSC_SUCCESS);
613: }

615: static PetscErrorCode MatMultHermitianTransposeAddColumnRange_MPIDense(Mat A, Vec xx, Vec yy, Vec zz, PetscInt c_start, PetscInt c_end)
616: {
617:   Mat_MPIDense      *a = (Mat_MPIDense *)A->data;
618:   const PetscScalar *ax;
619:   PetscScalar       *ay;
620:   PetscMemType       axmtype, aymtype;

622:   PetscFunctionBegin;
623:   if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
624:   PetscCall(VecCopy(yy, zz));
625:   PetscCall(VecZeroEntries(a->lvec));
626:   PetscUseMethod(a->A, "MatMultHermitianTransposeColumnRange_C", (Mat, Vec, Vec, PetscInt, PetscInt), (a->A, xx, a->lvec, c_start, c_end));
627:   PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype));
628:   PetscCall(VecGetArrayAndMemType(zz, &ay, &aymtype));
629:   PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM));
630:   PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM));
631:   PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax));
632:   PetscCall(VecRestoreArrayAndMemType(zz, &ay));
633:   PetscFunctionReturn(PETSC_SUCCESS);
634: }

636: static PetscErrorCode MatMultTransposeAddKernel_MPIDense(Mat A, Vec xx, Vec yy, Vec zz, PetscBool herm)
637: {
638:   Mat_MPIDense      *a = (Mat_MPIDense *)A->data;
639:   const PetscScalar *ax;
640:   PetscScalar       *ay;
641:   PetscMemType       axmtype, aymtype;

643:   PetscFunctionBegin;
644:   if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
645:   PetscCall(VecCopy(yy, zz));
646:   if (herm) PetscUseTypeMethod(a->A, multhermitiantranspose, xx, a->lvec);
647:   else PetscUseTypeMethod(a->A, multtranspose, xx, a->lvec);
648:   PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype));
649:   PetscCall(VecGetArrayAndMemType(zz, &ay, &aymtype));
650:   PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM));
651:   PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM));
652:   PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax));
653:   PetscCall(VecRestoreArrayAndMemType(zz, &ay));
654:   PetscFunctionReturn(PETSC_SUCCESS);
655: }

657: static PetscErrorCode MatMultTranspose_MPIDense(Mat A, Vec xx, Vec yy)
658: {
659:   PetscFunctionBegin;
660:   PetscCall(MatMultTransposeKernel_MPIDense(A, xx, yy, PETSC_FALSE));
661:   PetscFunctionReturn(PETSC_SUCCESS);
662: }

664: static PetscErrorCode MatMultTransposeAdd_MPIDense(Mat A, Vec xx, Vec yy, Vec zz)
665: {
666:   PetscFunctionBegin;
667:   PetscCall(MatMultTransposeAddKernel_MPIDense(A, xx, yy, zz, PETSC_FALSE));
668:   PetscFunctionReturn(PETSC_SUCCESS);
669: }

671: static PetscErrorCode MatMultHermitianTranspose_MPIDense(Mat A, Vec xx, Vec yy)
672: {
673:   PetscFunctionBegin;
674:   PetscCall(MatMultTransposeKernel_MPIDense(A, xx, yy, PETSC_TRUE));
675:   PetscFunctionReturn(PETSC_SUCCESS);
676: }

678: static PetscErrorCode MatMultHermitianTransposeAdd_MPIDense(Mat A, Vec xx, Vec yy, Vec zz)
679: {
680:   PetscFunctionBegin;
681:   PetscCall(MatMultTransposeAddKernel_MPIDense(A, xx, yy, zz, PETSC_TRUE));
682:   PetscFunctionReturn(PETSC_SUCCESS);
683: }

685: PetscErrorCode MatGetDiagonal_MPIDense(Mat A, Vec v)
686: {
687:   Mat_MPIDense      *a = (Mat_MPIDense *)A->data;
688:   PetscInt           lda, len, i, nl, ng, m = A->rmap->n, radd;
689:   PetscScalar       *x;
690:   const PetscScalar *av;

692:   PetscFunctionBegin;
693:   PetscCall(VecGetArray(v, &x));
694:   PetscCall(VecGetSize(v, &ng));
695:   PetscCheck(ng == A->rmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Nonconforming mat and vec");
696:   PetscCall(VecGetLocalSize(v, &nl));
697:   len  = PetscMin(a->A->rmap->n, a->A->cmap->n);
698:   radd = A->rmap->rstart * m;
699:   PetscCall(MatDenseGetArrayRead(a->A, &av));
700:   PetscCall(MatDenseGetLDA(a->A, &lda));
701:   for (i = 0; i < len; i++) x[i] = av[radd + i * lda + i];
702:   PetscCall(MatDenseRestoreArrayRead(a->A, &av));
703:   if (nl - i > 0) PetscCall(PetscArrayzero(x + i, nl - i));
704:   PetscCall(VecRestoreArray(v, &x));
705:   PetscFunctionReturn(PETSC_SUCCESS);
706: }

708: static PetscErrorCode MatDestroy_MPIDense(Mat mat)
709: {
710:   Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;

712:   PetscFunctionBegin;
713:   PetscCall(PetscLogObjectState((PetscObject)mat, "Rows=%" PetscInt_FMT ", Cols=%" PetscInt_FMT, mat->rmap->N, mat->cmap->N));
714:   PetscCall(MatStashDestroy_Private(&mat->stash));
715:   PetscCheck(!mdn->vecinuse, PetscObjectComm((PetscObject)mat), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
716:   PetscCheck(!mdn->matinuse, PetscObjectComm((PetscObject)mat), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
717:   PetscCall(MatDestroy(&mdn->A));
718:   PetscCall(VecDestroy(&mdn->lvec));
719:   PetscCall(PetscSFDestroy(&mdn->Mvctx));
720:   PetscCall(VecDestroy(&mdn->cvec));
721:   PetscCall(MatDestroy(&mdn->cmat));

723:   PetscCall(PetscFree(mat->data));
724:   PetscCall(PetscObjectChangeTypeName((PetscObject)mat, NULL));

726:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetLDA_C", NULL));
727:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseSetLDA_C", NULL));
728:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArray_C", NULL));
729:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArray_C", NULL));
730:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayRead_C", NULL));
731:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayRead_C", NULL));
732:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayWrite_C", NULL));
733:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayWrite_C", NULL));
734:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDensePlaceArray_C", NULL));
735:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseResetArray_C", NULL));
736:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseReplaceArray_C", NULL));
737:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpiaij_mpidense_C", NULL));
738:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpiaij_C", NULL));
739: #if PetscDefined(HAVE_ELEMENTAL)
740:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_elemental_C", NULL));
741: #endif
742: #if PetscDefined(HAVE_SCALAPACK) && (PetscDefined(USE_REAL_SINGLE) || PetscDefined(USE_REAL_DOUBLE))
743:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_scalapack_C", NULL));
744: #endif
745:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMPIDenseSetPreallocation_C", NULL));
746:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidense_C", NULL));
747:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaij_C", NULL));
748: #if PetscDefined(HAVE_CUDA)
749:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijcusparse_mpidense_C", NULL));
750:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijcusparse_C", NULL));
751:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensecuda_C", NULL));
752:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidensecuda_mpidense_C", NULL));
753:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidensecuda_C", NULL));
754:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijcusparse_mpidensecuda_C", NULL));
755:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensecuda_mpiaij_C", NULL));
756:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensecuda_mpiaijcusparse_C", NULL));
757:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAGetArray_C", NULL));
758:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAGetArrayRead_C", NULL));
759:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAGetArrayWrite_C", NULL));
760:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDARestoreArray_C", NULL));
761:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDARestoreArrayRead_C", NULL));
762:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDARestoreArrayWrite_C", NULL));
763:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAPlaceArray_C", NULL));
764:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAResetArray_C", NULL));
765:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAReplaceArray_C", NULL));
766:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDASetPreallocation_C", NULL));
767: #endif
768: #if PetscDefined(HAVE_HIP)
769:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijhipsparse_mpidense_C", NULL));
770:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijhipsparse_C", NULL));
771:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensehip_C", NULL));
772:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidensehip_mpidense_C", NULL));
773:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidensehip_C", NULL));
774:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijhipsparse_mpidensehip_C", NULL));
775:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensehip_mpiaij_C", NULL));
776:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensehip_mpiaijhipsparse_C", NULL));
777:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPGetArray_C", NULL));
778:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPGetArrayRead_C", NULL));
779:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPGetArrayWrite_C", NULL));
780:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPRestoreArray_C", NULL));
781:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPRestoreArrayRead_C", NULL));
782:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPRestoreArrayWrite_C", NULL));
783:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPPlaceArray_C", NULL));
784:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPResetArray_C", NULL));
785:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPReplaceArray_C", NULL));
786:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPSetPreallocation_C", NULL));
787: #endif
788:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumn_C", NULL));
789:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumn_C", NULL));
790:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVec_C", NULL));
791:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVec_C", NULL));
792:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecRead_C", NULL));
793:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecRead_C", NULL));
794:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecWrite_C", NULL));
795:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecWrite_C", NULL));
796:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetSubMatrix_C", NULL));
797:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreSubMatrix_C", NULL));
798:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultColumnRange_C", NULL));
799:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultAddColumnRange_C", NULL));
800:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeColumnRange_C", NULL));
801:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeAddColumnRange_C", NULL));
802:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatGetMultPetscSF_C", NULL));

804:   PetscCall(PetscObjectCompose((PetscObject)mat, "DiagonalBlock", NULL));
805:   PetscFunctionReturn(PETSC_SUCCESS);
806: }

808: #include <petscdraw.h>
809: static PetscErrorCode MatView_MPIDense_ASCIIorDraworSocket(Mat mat, PetscViewer viewer)
810: {
811:   Mat_MPIDense     *mdn = (Mat_MPIDense *)mat->data;
812:   PetscMPIInt       rank;
813:   PetscViewerType   vtype;
814:   PetscBool         isascii, isdraw;
815:   PetscViewer       sviewer;
816:   PetscViewerFormat format;

818:   PetscFunctionBegin;
819:   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)mat), &rank));
820:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
821:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERDRAW, &isdraw));
822:   if (isascii) {
823:     PetscCall(PetscViewerGetType(viewer, &vtype));
824:     PetscCall(PetscViewerGetFormat(viewer, &format));
825:     if (format == PETSC_VIEWER_ASCII_INFO_DETAIL) {
826:       MatInfo info;
827:       PetscCall(MatGetInfo(mat, MAT_LOCAL, &info));
828:       PetscCall(PetscViewerASCIIPushSynchronized(viewer));
829:       PetscCall(PetscViewerASCIISynchronizedPrintf(viewer, "  [%d] local rows %" PetscInt_FMT " nz %" PetscInt_FMT " nz alloced %" PetscInt_FMT " mem %" PetscInt_FMT " \n", rank, mat->rmap->n, (PetscInt)info.nz_used, (PetscInt)info.nz_allocated,
830:                                                    (PetscInt)info.memory));
831:       PetscCall(PetscViewerFlush(viewer));
832:       PetscCall(PetscViewerASCIIPopSynchronized(viewer));
833:       if (mdn->Mvctx) PetscCall(PetscSFView(mdn->Mvctx, viewer));
834:       PetscFunctionReturn(PETSC_SUCCESS);
835:     } else if (format == PETSC_VIEWER_ASCII_INFO) {
836:       PetscFunctionReturn(PETSC_SUCCESS);
837:     }
838:   } else if (isdraw) {
839:     PetscDraw draw;
840:     PetscBool isnull;

842:     PetscCall(PetscViewerDrawGetDraw(viewer, 0, &draw));
843:     PetscCall(PetscDrawIsNull(draw, &isnull));
844:     if (isnull) PetscFunctionReturn(PETSC_SUCCESS);
845:   }

847:   {
848:     /* assemble the entire matrix onto first processor. */
849:     Mat          A;
850:     PetscInt     M = mat->rmap->N, N = mat->cmap->N, m, row, i, nz;
851:     PetscInt    *cols;
852:     PetscScalar *vals;

854:     PetscCall(MatCreate(PetscObjectComm((PetscObject)mat), &A));
855:     if (rank == 0) {
856:       PetscCall(MatSetSizes(A, M, N, M, N));
857:     } else {
858:       PetscCall(MatSetSizes(A, 0, 0, M, N));
859:     }
860:     /* Since this is a temporary matrix, MATMPIDENSE instead of ((PetscObject)A)->type_name here is probably acceptable. */
861:     PetscCall(MatSetType(A, MATMPIDENSE));
862:     PetscCall(MatMPIDenseSetPreallocation(A, NULL));

864:     /* Copy the matrix ... This isn't the most efficient means,
865:        but it's quick for now */
866:     A->insertmode = INSERT_VALUES;

868:     row = mat->rmap->rstart;
869:     m   = mdn->A->rmap->n;
870:     for (i = 0; i < m; i++) {
871:       PetscCall(MatGetRow_MPIDense(mat, row, &nz, &cols, &vals));
872:       PetscCall(MatSetValues_MPIDense(A, 1, &row, nz, cols, vals, INSERT_VALUES));
873:       PetscCall(MatRestoreRow_MPIDense(mat, row, &nz, &cols, &vals));
874:       row++;
875:     }

877:     PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
878:     PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
879:     PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sviewer));
880:     if (rank == 0) {
881:       PetscCall(PetscObjectSetName((PetscObject)((Mat_MPIDense *)A->data)->A, ((PetscObject)mat)->name));
882:       PetscCall(MatView_SeqDense(((Mat_MPIDense *)A->data)->A, sviewer));
883:     }
884:     PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sviewer));
885:     PetscCall(MatDestroy(&A));
886:   }
887:   PetscFunctionReturn(PETSC_SUCCESS);
888: }

890: static PetscErrorCode MatView_MPIDense(Mat mat, PetscViewer viewer)
891: {
892:   PetscBool isascii, isbinary, isdraw, issocket;

894:   PetscFunctionBegin;
895:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
896:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &isbinary));
897:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERSOCKET, &issocket));
898:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERDRAW, &isdraw));

900:   if (isascii || issocket || isdraw) PetscCall(MatView_MPIDense_ASCIIorDraworSocket(mat, viewer));
901:   else if (isbinary) PetscCall(MatView_Dense_Binary(mat, viewer));
902:   PetscFunctionReturn(PETSC_SUCCESS);
903: }

905: static PetscErrorCode MatGetInfo_MPIDense(Mat A, MatInfoType flag, MatInfo *info)
906: {
907:   Mat_MPIDense  *mat = (Mat_MPIDense *)A->data;
908:   Mat            mdn = mat->A;
909:   PetscLogDouble irecv[5];

911:   PetscFunctionBegin;
912:   info->block_size = 1.0;

914:   PetscCall(MatGetInfo(mdn, MAT_LOCAL, info));

916:   irecv[0] = info->nz_used;
917:   irecv[1] = info->nz_allocated;
918:   irecv[2] = info->nz_unneeded;
919:   irecv[3] = info->memory;
920:   irecv[4] = info->mallocs;
921:   if (flag == MAT_LOCAL) {
922:     info->nz_used      = irecv[0];
923:     info->nz_allocated = irecv[1];
924:     info->nz_unneeded  = irecv[2];
925:     info->memory       = irecv[3];
926:     info->mallocs      = irecv[4];
927:   } else if (flag == MAT_GLOBAL_MAX) {
928:     PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, irecv, 5, MPIU_PETSCLOGDOUBLE, MPI_MAX, PetscObjectComm((PetscObject)A)));

930:     info->nz_used      = irecv[0];
931:     info->nz_allocated = irecv[1];
932:     info->nz_unneeded  = irecv[2];
933:     info->memory       = irecv[3];
934:     info->mallocs      = irecv[4];
935:   } else if (flag == MAT_GLOBAL_SUM) {
936:     PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, irecv, 5, MPIU_PETSCLOGDOUBLE, MPI_SUM, PetscObjectComm((PetscObject)A)));

938:     info->nz_used      = irecv[0];
939:     info->nz_allocated = irecv[1];
940:     info->nz_unneeded  = irecv[2];
941:     info->memory       = irecv[3];
942:     info->mallocs      = irecv[4];
943:   }
944:   info->fill_ratio_given  = 0; /* no parallel LU/ILU/Cholesky */
945:   info->fill_ratio_needed = 0;
946:   info->factor_mallocs    = 0;
947:   PetscFunctionReturn(PETSC_SUCCESS);
948: }

950: static PetscErrorCode MatSetOption_MPIDense(Mat A, MatOption op, PetscBool flg)
951: {
952:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;

954:   PetscFunctionBegin;
955:   switch (op) {
956:   case MAT_NEW_NONZERO_LOCATIONS:
957:   case MAT_NEW_NONZERO_LOCATION_ERR:
958:   case MAT_NEW_NONZERO_ALLOCATION_ERR:
959:     MatCheckPreallocated(A, 1);
960:     PetscCall(MatSetOption(a->A, op, flg));
961:     break;
962:   case MAT_ROW_ORIENTED:
963:     MatCheckPreallocated(A, 1);
964:     a->roworiented = flg;
965:     PetscCall(MatSetOption(a->A, op, flg));
966:     break;
967:   case MAT_IGNORE_OFF_PROC_ENTRIES:
968:     a->donotstash = flg;
969:     break;
970:   case MAT_SYMMETRIC:
971:   case MAT_STRUCTURALLY_SYMMETRIC:
972:   case MAT_HERMITIAN:
973:   case MAT_SYMMETRY_ETERNAL:
974:   case MAT_STRUCTURAL_SYMMETRY_ETERNAL:
975:   case MAT_SPD:
976:   case MAT_SPD_ETERNAL:
977:     /* if the diagonal matrix is square it inherits some of the properties above */
978:     if (a->A && A->rmap->n == A->cmap->n) PetscCall(MatSetOption(a->A, op, flg));
979:     break;
980:   default:
981:     break;
982:   }
983:   PetscFunctionReturn(PETSC_SUCCESS);
984: }

986: static PetscErrorCode MatDiagonalScale_MPIDense(Mat A, Vec ll, Vec rr)
987: {
988:   Mat_MPIDense      *mdn = (Mat_MPIDense *)A->data;
989:   const PetscScalar *l;
990:   PetscScalar        x, *v, *vv, *r;
991:   PetscInt           i, j, s2a, s3a, s2, s3, m = mdn->A->rmap->n, n = mdn->A->cmap->n, lda;

993:   PetscFunctionBegin;
994:   PetscCall(MatDenseGetArray(mdn->A, &vv));
995:   PetscCall(MatDenseGetLDA(mdn->A, &lda));
996:   PetscCall(MatGetLocalSize(A, &s2, &s3));
997:   if (ll) {
998:     PetscCall(VecGetLocalSize(ll, &s2a));
999:     PetscCheck(s2a == s2, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Left scaling vector non-conforming local size, %" PetscInt_FMT " != %" PetscInt_FMT, s2a, s2);
1000:     PetscCall(VecGetArrayRead(ll, &l));
1001:     for (i = 0; i < m; i++) {
1002:       x = l[i];
1003:       v = vv + i;
1004:       for (j = 0; j < n; j++) {
1005:         (*v) *= x;
1006:         v += lda;
1007:       }
1008:     }
1009:     PetscCall(VecRestoreArrayRead(ll, &l));
1010:     PetscCall(PetscLogFlops(1.0 * n * m));
1011:   }
1012:   if (rr) {
1013:     const PetscScalar *ar;

1015:     PetscCall(VecGetLocalSize(rr, &s3a));
1016:     PetscCheck(s3a == s3, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Right scaling vec non-conforming local size, %" PetscInt_FMT " != %" PetscInt_FMT ".", s3a, s3);
1017:     PetscCall(VecGetArrayRead(rr, &ar));
1018:     if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
1019:     PetscCall(VecGetArray(mdn->lvec, &r));
1020:     PetscCall(PetscSFBcastBegin(mdn->Mvctx, MPIU_SCALAR, ar, r, MPI_REPLACE));
1021:     PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ar, r, MPI_REPLACE));
1022:     PetscCall(VecRestoreArrayRead(rr, &ar));
1023:     for (i = 0; i < n; i++) {
1024:       x = r[i];
1025:       v = vv + i * lda;
1026:       for (j = 0; j < m; j++) (*v++) *= x;
1027:     }
1028:     PetscCall(VecRestoreArray(mdn->lvec, &r));
1029:     PetscCall(PetscLogFlops(1.0 * n * m));
1030:   }
1031:   PetscCall(MatDenseRestoreArray(mdn->A, &vv));
1032:   PetscFunctionReturn(PETSC_SUCCESS);
1033: }

1035: static PetscErrorCode MatNorm_MPIDense(Mat A, NormType type, PetscReal *nrm)
1036: {
1037:   Mat_MPIDense      *mdn = (Mat_MPIDense *)A->data;
1038:   PetscInt           i, j;
1039:   PetscMPIInt        size;
1040:   const PetscScalar *av, *v;

1042:   PetscFunctionBegin;
1043:   PetscCall(MatDenseGetArrayRead(mdn->A, &av));
1044:   v = av;
1045:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size));
1046:   if (size == 1) {
1047:     PetscCall(MatNorm(mdn->A, type, nrm));
1048:   } else {
1049:     if (type == NORM_FROBENIUS) {
1050:       *nrm = 0.0;
1051:       for (i = 0; i < mdn->A->cmap->n * mdn->A->rmap->n; i++) {
1052:         *nrm += PetscRealPart(PetscConj(*v) * (*v));
1053:         v++;
1054:       }
1055:       PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, nrm, 1, MPIU_REAL, MPIU_SUM, PetscObjectComm((PetscObject)A)));
1056:       *nrm = PetscSqrtReal(*nrm);
1057:       PetscCall(PetscLogFlops(2.0 * mdn->A->cmap->n * mdn->A->rmap->n));
1058:     } else if (type == NORM_1) {
1059:       PetscReal *tmp;

1061:       PetscCall(PetscCalloc1(A->cmap->N, &tmp));
1062:       *nrm = 0.0;
1063:       v    = av;
1064:       for (j = 0; j < mdn->A->cmap->n; j++) {
1065:         for (i = 0; i < mdn->A->rmap->n; i++) {
1066:           tmp[j] += PetscAbsScalar(*v);
1067:           v++;
1068:         }
1069:       }
1070:       PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, tmp, A->cmap->N, MPIU_REAL, MPIU_SUM, PetscObjectComm((PetscObject)A)));
1071:       for (j = 0; j < A->cmap->N; j++) {
1072:         if (tmp[j] > *nrm) *nrm = tmp[j];
1073:       }
1074:       PetscCall(PetscFree(tmp));
1075:       PetscCall(PetscLogFlops(A->cmap->n * A->rmap->n));
1076:     } else if (type == NORM_INFINITY) { /* max row norm */
1077:       PetscCall(MatNorm(mdn->A, type, nrm));
1078:       PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, nrm, 1, MPIU_REAL, MPIU_MAX, PetscObjectComm((PetscObject)A)));
1079:     } else SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "Unsupported norm type %s", NormTypes[type]);
1080:   }
1081:   PetscCall(MatDenseRestoreArrayRead(mdn->A, &av));
1082:   PetscFunctionReturn(PETSC_SUCCESS);
1083: }

1085: static PetscErrorCode MatTranspose_MPIDense(Mat A, MatReuse reuse, Mat *matout)
1086: {
1087:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1088:   Mat           B;
1089:   PetscInt      M = A->rmap->N, N = A->cmap->N, m, n, *rwork, rstart = A->rmap->rstart;
1090:   PetscInt      j, i, lda;
1091:   PetscScalar  *v;

1093:   PetscFunctionBegin;
1094:   if (reuse == MAT_REUSE_MATRIX) PetscCall(MatTransposeCheckNonzeroState_Private(A, *matout));
1095:   if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_INPLACE_MATRIX) {
1096:     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
1097:     PetscCall(MatSetSizes(B, A->cmap->n, A->rmap->n, N, M));
1098:     PetscCall(MatSetType(B, ((PetscObject)A)->type_name));
1099:     PetscCall(MatMPIDenseSetPreallocation(B, NULL));
1100:   } else B = *matout;

1102:   m = a->A->rmap->n;
1103:   n = a->A->cmap->n;
1104:   PetscCall(MatDenseGetArrayRead(a->A, (const PetscScalar **)&v));
1105:   PetscCall(MatDenseGetLDA(a->A, &lda));
1106:   PetscCall(PetscMalloc1(m, &rwork));
1107:   for (i = 0; i < m; i++) rwork[i] = rstart + i;
1108:   for (j = 0; j < n; j++) {
1109:     PetscCall(MatSetValues(B, 1, &j, m, rwork, v, INSERT_VALUES));
1110:     v = PetscSafePointerPlusOffset(v, lda);
1111:   }
1112:   PetscCall(MatDenseRestoreArrayRead(a->A, (const PetscScalar **)&v));
1113:   PetscCall(PetscFree(rwork));
1114:   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
1115:   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
1116:   if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_REUSE_MATRIX) {
1117:     *matout = B;
1118:   } else {
1119:     PetscCall(MatHeaderMerge(A, &B));
1120:   }
1121:   PetscFunctionReturn(PETSC_SUCCESS);
1122: }

1124: static PetscErrorCode       MatDuplicate_MPIDense(Mat, MatDuplicateOption, Mat *);
1125: PETSC_INTERN PetscErrorCode MatScale_MPIDense(Mat, PetscScalar);

1127: static PetscErrorCode MatSetUp_MPIDense(Mat A)
1128: {
1129:   PetscFunctionBegin;
1130:   PetscCall(PetscLayoutSetUp(A->rmap));
1131:   PetscCall(PetscLayoutSetUp(A->cmap));
1132:   if (!A->preallocated) PetscCall(MatMPIDenseSetPreallocation(A, NULL));
1133:   PetscFunctionReturn(PETSC_SUCCESS);
1134: }

1136: static PetscErrorCode MatAXPY_MPIDense(Mat Y, PetscScalar alpha, Mat X, MatStructure str)
1137: {
1138:   Mat_MPIDense *A = (Mat_MPIDense *)Y->data, *B = (Mat_MPIDense *)X->data;

1140:   PetscFunctionBegin;
1141:   PetscCall(MatAXPY(A->A, alpha, B->A, str));
1142:   PetscFunctionReturn(PETSC_SUCCESS);
1143: }

1145: static PetscErrorCode MatConjugate_MPIDense(Mat mat)
1146: {
1147:   Mat_MPIDense *a = (Mat_MPIDense *)mat->data;

1149:   PetscFunctionBegin;
1150:   PetscCall(MatConjugate(a->A));
1151:   PetscFunctionReturn(PETSC_SUCCESS);
1152: }

1154: static PetscErrorCode MatRealPart_MPIDense(Mat A)
1155: {
1156:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;

1158:   PetscFunctionBegin;
1159:   PetscCall(MatRealPart(a->A));
1160:   PetscFunctionReturn(PETSC_SUCCESS);
1161: }

1163: static PetscErrorCode MatImaginaryPart_MPIDense(Mat A)
1164: {
1165:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;

1167:   PetscFunctionBegin;
1168:   PetscCall(MatImaginaryPart(a->A));
1169:   PetscFunctionReturn(PETSC_SUCCESS);
1170: }

1172: static PetscErrorCode MatGetColumnVector_MPIDense(Mat A, Vec v, PetscInt col)
1173: {
1174:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;

1176:   PetscFunctionBegin;
1177:   PetscCheck(a->A, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Missing local matrix");
1178:   PetscCheck(a->A->ops->getcolumnvector, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Missing get column operation");
1179:   PetscUseTypeMethod(a->A, getcolumnvector, v, col);
1180:   PetscFunctionReturn(PETSC_SUCCESS);
1181: }

1183: PETSC_INTERN PetscErrorCode MatGetColumnReductions_SeqDense(Mat, PetscInt, PetscReal *);

1185: static PetscErrorCode MatGetColumnReductions_MPIDense(Mat A, PetscInt type, PetscReal *reductions)
1186: {
1187:   PetscInt      i, m, n;
1188:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;

1190:   PetscFunctionBegin;
1191:   PetscCall(MatGetSize(A, &m, &n));
1192:   if (type == REDUCTION_MEAN_REALPART) {
1193:     PetscCall(MatGetColumnReductions_SeqDense(a->A, (PetscInt)REDUCTION_SUM_REALPART, reductions));
1194:   } else if (type == REDUCTION_MEAN_IMAGINARYPART) {
1195:     PetscCall(MatGetColumnReductions_SeqDense(a->A, (PetscInt)REDUCTION_SUM_IMAGINARYPART, reductions));
1196:   } else {
1197:     PetscCall(MatGetColumnReductions_SeqDense(a->A, type, reductions));
1198:   }
1199:   if (type == NORM_2) {
1200:     for (i = 0; i < n; i++) reductions[i] *= reductions[i];
1201:   }
1202:   if (type == NORM_INFINITY) {
1203:     PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, reductions, n, MPIU_REAL, MPIU_MAX, A->hdr.comm));
1204:   } else {
1205:     PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, reductions, n, MPIU_REAL, MPIU_SUM, A->hdr.comm));
1206:   }
1207:   if (type == NORM_2) {
1208:     for (i = 0; i < n; i++) reductions[i] = PetscSqrtReal(reductions[i]);
1209:   } else if (type == REDUCTION_MEAN_REALPART || type == REDUCTION_MEAN_IMAGINARYPART) {
1210:     for (i = 0; i < n; i++) reductions[i] /= m;
1211:   }
1212:   PetscFunctionReturn(PETSC_SUCCESS);
1213: }

1215: static PetscErrorCode MatSetRandom_MPIDense(Mat x, PetscRandom rctx)
1216: {
1217:   Mat_MPIDense *d = (Mat_MPIDense *)x->data;

1219:   PetscFunctionBegin;
1220:   PetscCall(MatSetRandom(d->A, rctx));
1221: #if PetscDefined(HAVE_DEVICE)
1222:   x->offloadmask = d->A->offloadmask;
1223: #endif
1224:   PetscFunctionReturn(PETSC_SUCCESS);
1225: }

1227: static PetscErrorCode MatMatTransposeMultSymbolic_MPIDense_MPIDense(Mat, Mat, PetscReal, Mat);
1228: static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense(Mat, Mat, Mat);
1229: static PetscErrorCode MatTransposeMatMultSymbolic_MPIDense_MPIDense(Mat, Mat, PetscReal, Mat);
1230: static PetscErrorCode MatTransposeMatMultNumeric_MPIDense_MPIDense(Mat, Mat, Mat);
1231: static PetscErrorCode MatEqual_MPIDense(Mat, Mat, PetscBool *);
1232: static PetscErrorCode MatLoad_MPIDense(Mat, PetscViewer);
1233: static PetscErrorCode MatProductSetFromOptions_MPIDense(Mat);

1235: static struct _MatOps MatOps_Values = {MatSetValues_MPIDense,
1236:                                        MatGetRow_MPIDense,
1237:                                        MatRestoreRow_MPIDense,
1238:                                        MatMult_MPIDense,
1239:                                        /*  4*/ MatMultAdd_MPIDense,
1240:                                        MatMultTranspose_MPIDense,
1241:                                        MatMultTransposeAdd_MPIDense,
1242:                                        NULL,
1243:                                        NULL,
1244:                                        NULL,
1245:                                        /* 10*/ NULL,
1246:                                        NULL,
1247:                                        NULL,
1248:                                        NULL,
1249:                                        MatTranspose_MPIDense,
1250:                                        /* 15*/ MatGetInfo_MPIDense,
1251:                                        MatEqual_MPIDense,
1252:                                        MatGetDiagonal_MPIDense,
1253:                                        MatDiagonalScale_MPIDense,
1254:                                        MatNorm_MPIDense,
1255:                                        /* 20*/ MatAssemblyBegin_MPIDense,
1256:                                        MatAssemblyEnd_MPIDense,
1257:                                        MatSetOption_MPIDense,
1258:                                        MatZeroEntries_MPIDense,
1259:                                        /* 24*/ MatZeroRows_MPIDense,
1260:                                        NULL,
1261:                                        NULL,
1262:                                        NULL,
1263:                                        NULL,
1264:                                        /* 29*/ MatSetUp_MPIDense,
1265:                                        NULL,
1266:                                        NULL,
1267:                                        MatGetDiagonalBlock_MPIDense,
1268:                                        NULL,
1269:                                        /* 34*/ MatDuplicate_MPIDense,
1270:                                        NULL,
1271:                                        NULL,
1272:                                        NULL,
1273:                                        NULL,
1274:                                        /* 39*/ MatAXPY_MPIDense,
1275:                                        MatCreateSubMatrices_MPIDense,
1276:                                        NULL,
1277:                                        MatGetValues_MPIDense,
1278:                                        MatCopy_MPIDense,
1279:                                        /* 44*/ NULL,
1280:                                        MatScale_MPIDense,
1281:                                        MatShift_MPIDense,
1282:                                        NULL,
1283:                                        NULL,
1284:                                        /* 49*/ MatSetRandom_MPIDense,
1285:                                        NULL,
1286:                                        NULL,
1287:                                        NULL,
1288:                                        NULL,
1289:                                        /* 54*/ NULL,
1290:                                        NULL,
1291:                                        NULL,
1292:                                        NULL,
1293:                                        NULL,
1294:                                        /* 59*/ MatCreateSubMatrix_MPIDense,
1295:                                        MatDestroy_MPIDense,
1296:                                        MatView_MPIDense,
1297:                                        NULL,
1298:                                        NULL,
1299:                                        /* 64*/ NULL,
1300:                                        NULL,
1301:                                        NULL,
1302:                                        NULL,
1303:                                        NULL,
1304:                                        /* 69*/ NULL,
1305:                                        NULL,
1306:                                        NULL,
1307:                                        NULL,
1308:                                        NULL,
1309:                                        /* 74*/ NULL,
1310:                                        NULL,
1311:                                        NULL,
1312:                                        NULL,
1313:                                        MatLoad_MPIDense,
1314:                                        /* 79*/ NULL,
1315:                                        NULL,
1316:                                        NULL,
1317:                                        NULL,
1318:                                        /* 83*/ NULL,
1319:                                        NULL,
1320:                                        NULL,
1321:                                        NULL,
1322:                                        MatMatTransposeMultSymbolic_MPIDense_MPIDense,
1323:                                        MatMatTransposeMultNumeric_MPIDense_MPIDense,
1324:                                        /* 89*/ NULL,
1325:                                        MatProductSetFromOptions_MPIDense,
1326:                                        NULL,
1327:                                        NULL,
1328:                                        MatConjugate_MPIDense,
1329:                                        /* 94*/ NULL,
1330:                                        NULL,
1331:                                        MatRealPart_MPIDense,
1332:                                        MatImaginaryPart_MPIDense,
1333:                                        NULL,
1334:                                        /*99*/ NULL,
1335:                                        NULL,
1336:                                        NULL,
1337:                                        NULL,
1338:                                        MatGetColumnVector_MPIDense,
1339:                                        /*104*/ NULL,
1340:                                        NULL,
1341:                                        NULL,
1342:                                        NULL,
1343:                                        NULL,
1344:                                        /*109*/ NULL,
1345:                                        NULL,
1346:                                        MatMultHermitianTranspose_MPIDense,
1347:                                        MatMultHermitianTransposeAdd_MPIDense,
1348:                                        NULL,
1349:                                        /*114*/ NULL,
1350:                                        MatGetColumnReductions_MPIDense,
1351:                                        NULL,
1352:                                        NULL,
1353:                                        NULL,
1354:                                        /*120*/ MatTransposeMatMultSymbolic_MPIDense_MPIDense,
1355:                                        MatTransposeMatMultNumeric_MPIDense_MPIDense,
1356:                                        NULL,
1357:                                        NULL,
1358:                                        /*124*/ NULL,
1359:                                        NULL,
1360:                                        NULL,
1361:                                        NULL,
1362:                                        NULL,
1363:                                        /*129*/ NULL,
1364:                                        MatCreateMPIMatConcatenateSeqMat_MPIDense,
1365:                                        NULL,
1366:                                        NULL,
1367:                                        NULL,
1368:                                        /*134*/ NULL,
1369:                                        NULL,
1370:                                        NULL,
1371:                                        NULL,
1372:                                        NULL,
1373:                                        /*139*/ NULL,
1374:                                        NULL,
1375:                                        NULL,
1376:                                        NULL,
1377:                                        NULL,
1378:                                        MatADot_Default,
1379:                                        /*144*/ MatANorm_Default,
1380:                                        NULL,
1381:                                        NULL,
1382:                                        NULL};

1384: static PetscErrorCode MatMPIDenseSetPreallocation_MPIDense(Mat mat, PetscScalar *data)
1385: {
1386:   Mat_MPIDense *a     = (Mat_MPIDense *)mat->data;
1387:   MatType       mtype = MATSEQDENSE;

1389:   PetscFunctionBegin;
1390:   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)mat), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1391:   PetscCall(PetscLayoutSetUp(mat->rmap));
1392:   PetscCall(PetscLayoutSetUp(mat->cmap));
1393:   if (!a->A) {
1394:     PetscCall(MatCreate(PETSC_COMM_SELF, &a->A));
1395:     PetscCall(MatSetSizes(a->A, mat->rmap->n, mat->cmap->N, mat->rmap->n, mat->cmap->N));
1396:   }
1397: #if PetscDefined(HAVE_CUDA)
1398:   PetscBool iscuda;
1399:   PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATMPIDENSECUDA, &iscuda));
1400:   if (iscuda) mtype = MATSEQDENSECUDA;
1401: #endif
1402: #if PetscDefined(HAVE_HIP)
1403:   PetscBool iship;
1404:   PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATMPIDENSEHIP, &iship));
1405:   if (iship) mtype = MATSEQDENSEHIP;
1406: #endif
1407:   PetscCall(MatSetType(a->A, mtype));
1408:   PetscCall(MatSeqDenseSetPreallocation(a->A, data));
1409: #if PetscDefined(HAVE_CUDA) || PetscDefined(HAVE_HIP)
1410:   mat->offloadmask = a->A->offloadmask;
1411: #endif
1412:   mat->preallocated = PETSC_TRUE;
1413:   mat->assembled    = PETSC_TRUE;
1414:   PetscFunctionReturn(PETSC_SUCCESS);
1415: }

1417: PETSC_INTERN PetscErrorCode MatConvert_MPIAIJ_MPIDense(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
1418: {
1419:   Mat B, C;

1421:   PetscFunctionBegin;
1422:   PetscCall(MatMPIAIJGetLocalMat(A, MAT_INITIAL_MATRIX, &C));
1423:   PetscCall(MatConvert_SeqAIJ_SeqDense(C, MATSEQDENSE, MAT_INITIAL_MATRIX, &B));
1424:   PetscCall(MatDestroy(&C));
1425:   if (reuse == MAT_REUSE_MATRIX) {
1426:     C = *newmat;
1427:   } else C = NULL;
1428:   PetscCall(MatCreateMPIMatConcatenateSeqMat(PetscObjectComm((PetscObject)A), B, A->cmap->n, !C ? MAT_INITIAL_MATRIX : MAT_REUSE_MATRIX, &C));
1429:   PetscCall(MatDestroy(&B));
1430:   if (reuse == MAT_INPLACE_MATRIX) PetscCall(MatHeaderReplace(A, &C));
1431:   else if (reuse == MAT_INITIAL_MATRIX) *newmat = C;
1432:   PetscFunctionReturn(PETSC_SUCCESS);
1433: }

1435: static PetscErrorCode MatConvert_MPIDense_MPIAIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
1436: {
1437:   Mat B, C;

1439:   PetscFunctionBegin;
1440:   PetscCall(MatDenseGetLocalMatrix(A, &C));
1441:   PetscCall(MatConvert_SeqDense_SeqAIJ(C, MATSEQAIJ, MAT_INITIAL_MATRIX, &B));
1442:   if (reuse == MAT_REUSE_MATRIX) {
1443:     C = *newmat;
1444:   } else C = NULL;
1445:   PetscCall(MatCreateMPIMatConcatenateSeqMat(PetscObjectComm((PetscObject)A), B, A->cmap->n, !C ? MAT_INITIAL_MATRIX : MAT_REUSE_MATRIX, &C));
1446:   PetscCall(MatDestroy(&B));
1447:   if (reuse == MAT_INPLACE_MATRIX) PetscCall(MatHeaderReplace(A, &C));
1448:   else if (reuse == MAT_INITIAL_MATRIX) *newmat = C;
1449:   PetscFunctionReturn(PETSC_SUCCESS);
1450: }

1452: #if PetscDefined(HAVE_ELEMENTAL)
1453: PETSC_INTERN PetscErrorCode MatConvert_MPIDense_Elemental(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
1454: {
1455:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1456:   Mat           mat_elemental;
1457:   PetscScalar  *v;
1458:   PetscInt      m = A->rmap->n, N = A->cmap->N, rstart = A->rmap->rstart, i, *rows, *cols, lda;

1460:   PetscFunctionBegin;
1461:   if (reuse == MAT_REUSE_MATRIX) {
1462:     mat_elemental = *newmat;
1463:     PetscCall(MatZeroEntries(*newmat));
1464:   } else {
1465:     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &mat_elemental));
1466:     PetscCall(MatSetSizes(mat_elemental, PETSC_DECIDE, PETSC_DECIDE, A->rmap->N, A->cmap->N));
1467:     PetscCall(MatSetType(mat_elemental, MATELEMENTAL));
1468:     PetscCall(MatSetUp(mat_elemental));
1469:     PetscCall(MatSetOption(mat_elemental, MAT_ROW_ORIENTED, PETSC_FALSE));
1470:   }

1472:   PetscCall(PetscMalloc2(m, &rows, N, &cols));
1473:   for (i = 0; i < N; i++) cols[i] = i;
1474:   for (i = 0; i < m; i++) rows[i] = rstart + i;

1476:   /* PETSc-Elemental interface uses axpy for setting off-processor entries, only ADD_VALUES is allowed */
1477:   PetscCall(MatDenseGetArray(A, &v));
1478:   PetscCall(MatDenseGetLDA(a->A, &lda));
1479:   if (lda == m) PetscCall(MatSetValues(mat_elemental, m, rows, N, cols, v, ADD_VALUES));
1480:   else {
1481:     for (i = 0; i < N; i++) PetscCall(MatSetValues(mat_elemental, m, rows, 1, &i, v + lda * i, ADD_VALUES));
1482:   }
1483:   PetscCall(MatAssemblyBegin(mat_elemental, MAT_FINAL_ASSEMBLY));
1484:   PetscCall(MatAssemblyEnd(mat_elemental, MAT_FINAL_ASSEMBLY));
1485:   PetscCall(MatDenseRestoreArray(A, &v));
1486:   PetscCall(PetscFree2(rows, cols));

1488:   if (reuse == MAT_INPLACE_MATRIX) {
1489:     PetscCall(MatHeaderReplace(A, &mat_elemental));
1490:   } else {
1491:     *newmat = mat_elemental;
1492:   }
1493:   PetscFunctionReturn(PETSC_SUCCESS);
1494: }
1495: #endif

1497: static PetscErrorCode MatDenseGetColumn_MPIDense(Mat A, PetscInt col, PetscScalar **vals)
1498: {
1499:   Mat_MPIDense *mat = (Mat_MPIDense *)A->data;

1501:   PetscFunctionBegin;
1502:   PetscCall(MatDenseGetColumn(mat->A, col, vals));
1503:   PetscFunctionReturn(PETSC_SUCCESS);
1504: }

1506: static PetscErrorCode MatDenseRestoreColumn_MPIDense(Mat A, PetscScalar **vals)
1507: {
1508:   Mat_MPIDense *mat = (Mat_MPIDense *)A->data;

1510:   PetscFunctionBegin;
1511:   PetscCall(MatDenseRestoreColumn(mat->A, vals));
1512:   PetscFunctionReturn(PETSC_SUCCESS);
1513: }

1515: PetscErrorCode MatCreateMPIMatConcatenateSeqMat_MPIDense(MPI_Comm comm, Mat inmat, PetscInt n, MatReuse scall, Mat *outmat)
1516: {
1517:   Mat_MPIDense *mat;
1518:   PetscInt      m, nloc, N;

1520:   PetscFunctionBegin;
1521:   PetscCall(MatGetSize(inmat, &m, &N));
1522:   PetscCall(MatGetLocalSize(inmat, NULL, &nloc));
1523:   if (scall == MAT_INITIAL_MATRIX) { /* symbolic phase */
1524:     PetscInt sum;

1526:     if (n == PETSC_DECIDE) PetscCall(PetscSplitOwnership(comm, &n, &N));
1527:     /* Check sum(n) = N */
1528:     PetscCallMPI(MPIU_Allreduce(&n, &sum, 1, MPIU_INT, MPI_SUM, comm));
1529:     PetscCheck(sum == N, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Sum of local columns %" PetscInt_FMT " != global columns %" PetscInt_FMT, sum, N);

1531:     PetscCall(MatCreateDense(comm, m, n, PETSC_DETERMINE, N, NULL, outmat));
1532:     PetscCall(MatSetOption(*outmat, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
1533:   }

1535:   /* numeric phase */
1536:   mat = (Mat_MPIDense *)(*outmat)->data;
1537:   PetscCall(MatCopy(inmat, mat->A, SAME_NONZERO_PATTERN));
1538:   PetscFunctionReturn(PETSC_SUCCESS);
1539: }

1541: PetscErrorCode MatDenseGetColumnVec_MPIDense(Mat A, PetscInt col, Vec *v)
1542: {
1543:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1544:   PetscInt      lda;

1546:   PetscFunctionBegin;
1547:   PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
1548:   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1549:   if (!a->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &a->cvec));
1550:   a->vecinuse = col + 1;
1551:   PetscCall(MatDenseGetLDA(a->A, &lda));
1552:   PetscCall(MatDenseGetArray(a->A, (PetscScalar **)&a->ptrinuse));
1553:   PetscCall(VecPlaceArray(a->cvec, PetscSafePointerPlusOffset(a->ptrinuse, (size_t)col * (size_t)lda)));
1554:   *v = a->cvec;
1555:   PetscFunctionReturn(PETSC_SUCCESS);
1556: }

1558: PetscErrorCode MatDenseRestoreColumnVec_MPIDense(Mat A, PetscInt col, Vec *v)
1559: {
1560:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;

1562:   PetscFunctionBegin;
1563:   PetscCheck(a->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
1564:   PetscCheck(a->cvec, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Missing internal column vector");
1565:   VecCheckAssembled(a->cvec);
1566:   a->vecinuse = 0;
1567:   PetscCall(MatDenseRestoreArray(a->A, (PetscScalar **)&a->ptrinuse));
1568:   PetscCall(VecResetArray(a->cvec));
1569:   if (v) *v = NULL;
1570:   PetscFunctionReturn(PETSC_SUCCESS);
1571: }

1573: PetscErrorCode MatDenseGetColumnVecRead_MPIDense(Mat A, PetscInt col, Vec *v)
1574: {
1575:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1576:   PetscInt      lda;

1578:   PetscFunctionBegin;
1579:   PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
1580:   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1581:   if (!a->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &a->cvec));
1582:   a->vecinuse = col + 1;
1583:   PetscCall(MatDenseGetLDA(a->A, &lda));
1584:   PetscCall(MatDenseGetArrayRead(a->A, &a->ptrinuse));
1585:   PetscCall(VecPlaceArray(a->cvec, PetscSafePointerPlusOffset(a->ptrinuse, (size_t)col * (size_t)lda)));
1586:   PetscCall(VecLockReadPush(a->cvec));
1587:   *v = a->cvec;
1588:   PetscFunctionReturn(PETSC_SUCCESS);
1589: }

1591: PetscErrorCode MatDenseRestoreColumnVecRead_MPIDense(Mat A, PetscInt col, Vec *v)
1592: {
1593:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;

1595:   PetscFunctionBegin;
1596:   PetscCheck(a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
1597:   PetscCheck(a->cvec, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing internal column vector");
1598:   VecCheckAssembled(a->cvec);
1599:   a->vecinuse = 0;
1600:   PetscCall(MatDenseRestoreArrayRead(a->A, &a->ptrinuse));
1601:   PetscCall(VecLockReadPop(a->cvec));
1602:   PetscCall(VecResetArray(a->cvec));
1603:   if (v) *v = NULL;
1604:   PetscFunctionReturn(PETSC_SUCCESS);
1605: }

1607: PetscErrorCode MatDenseGetColumnVecWrite_MPIDense(Mat A, PetscInt col, Vec *v)
1608: {
1609:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1610:   PetscInt      lda;

1612:   PetscFunctionBegin;
1613:   PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
1614:   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1615:   if (!a->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &a->cvec));
1616:   a->vecinuse = col + 1;
1617:   PetscCall(MatDenseGetLDA(a->A, &lda));
1618:   PetscCall(MatDenseGetArrayWrite(a->A, (PetscScalar **)&a->ptrinuse));
1619:   PetscCall(VecPlaceArray(a->cvec, PetscSafePointerPlusOffset(a->ptrinuse, (size_t)col * (size_t)lda)));
1620:   *v = a->cvec;
1621:   PetscFunctionReturn(PETSC_SUCCESS);
1622: }

1624: PetscErrorCode MatDenseRestoreColumnVecWrite_MPIDense(Mat A, PetscInt col, Vec *v)
1625: {
1626:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;

1628:   PetscFunctionBegin;
1629:   PetscCheck(a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
1630:   PetscCheck(a->cvec, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing internal column vector");
1631:   VecCheckAssembled(a->cvec);
1632:   a->vecinuse = 0;
1633:   PetscCall(MatDenseRestoreArrayWrite(a->A, (PetscScalar **)&a->ptrinuse));
1634:   PetscCall(VecResetArray(a->cvec));
1635:   if (v) *v = NULL;
1636:   PetscFunctionReturn(PETSC_SUCCESS);
1637: }

1639: static PetscErrorCode MatDenseGetSubMatrix_MPIDense(Mat A, PetscInt rbegin, PetscInt rend, PetscInt cbegin, PetscInt cend, Mat *v)
1640: {
1641:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1642:   Mat_MPIDense *c;
1643:   MPI_Comm      comm;
1644:   PetscInt      prbegin, prend, pcbegin, pcend;

1646:   PetscFunctionBegin;
1647:   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
1648:   PetscCheck(!a->vecinuse, comm, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
1649:   PetscCheck(!a->matinuse, comm, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1650:   prbegin = PetscMax(0, PetscMin(A->rmap->rend, rbegin) - A->rmap->rstart);
1651:   prend   = PetscMin(A->rmap->n, PetscMax(0, rend - A->rmap->rstart));
1652:   pcbegin = PetscMax(0, PetscMin(A->cmap->rend, cbegin) - A->cmap->rstart);
1653:   pcend   = PetscMin(A->cmap->n, PetscMax(0, cend - A->cmap->rstart));
1654:   if (!a->cmat) {
1655:     PetscCall(MatCreate(comm, &a->cmat));
1656:     PetscCall(MatSetType(a->cmat, ((PetscObject)A)->type_name));
1657:     if (rend - rbegin == A->rmap->N) PetscCall(PetscLayoutReference(A->rmap, &a->cmat->rmap));
1658:     else {
1659:       PetscCall(PetscLayoutSetLocalSize(a->cmat->rmap, prend - prbegin));
1660:       PetscCall(PetscLayoutSetSize(a->cmat->rmap, rend - rbegin));
1661:       PetscCall(PetscLayoutSetUp(a->cmat->rmap));
1662:     }
1663:     if (cend - cbegin == A->cmap->N) PetscCall(PetscLayoutReference(A->cmap, &a->cmat->cmap));
1664:     else {
1665:       PetscCall(PetscLayoutSetLocalSize(a->cmat->cmap, pcend - pcbegin));
1666:       PetscCall(PetscLayoutSetSize(a->cmat->cmap, cend - cbegin));
1667:       PetscCall(PetscLayoutSetUp(a->cmat->cmap));
1668:     }
1669:     c             = (Mat_MPIDense *)a->cmat->data;
1670:     c->sub_rbegin = rbegin;
1671:     c->sub_rend   = rend;
1672:     c->sub_cbegin = cbegin;
1673:     c->sub_cend   = cend;
1674:   }
1675:   c = (Mat_MPIDense *)a->cmat->data;
1676:   if (c->sub_rbegin != rbegin || c->sub_rend != rend) {
1677:     PetscCall(PetscLayoutDestroy(&a->cmat->rmap));
1678:     PetscCall(PetscLayoutCreate(comm, &a->cmat->rmap));
1679:     PetscCall(PetscLayoutSetLocalSize(a->cmat->rmap, prend - prbegin));
1680:     PetscCall(PetscLayoutSetSize(a->cmat->rmap, rend - rbegin));
1681:     PetscCall(PetscLayoutSetUp(a->cmat->rmap));
1682:     c->sub_rbegin = rbegin;
1683:     c->sub_rend   = rend;
1684:   }
1685:   if (c->sub_cbegin != cbegin || c->sub_cend != cend) {
1686:     // special optimization: check if all columns are owned by rank 0, in which case no communication is necessary
1687:     if ((cend - cbegin != a->cmat->cmap->N) || (A->cmap->range[1] != A->cmap->N)) {
1688:       PetscCall(PetscLayoutDestroy(&a->cmat->cmap));
1689:       PetscCall(PetscLayoutCreate(comm, &a->cmat->cmap));
1690:       PetscCall(PetscLayoutSetLocalSize(a->cmat->cmap, pcend - pcbegin));
1691:       PetscCall(PetscLayoutSetSize(a->cmat->cmap, cend - cbegin));
1692:       PetscCall(PetscLayoutSetUp(a->cmat->cmap));
1693:       PetscCall(VecDestroy(&c->lvec));
1694:       PetscCall(PetscSFDestroy(&c->Mvctx));
1695:     }
1696:     c->sub_cbegin = cbegin;
1697:     c->sub_cend   = cend;
1698:   }
1699:   PetscCheck(!c->A, comm, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1700:   PetscCall(MatDenseGetSubMatrix(a->A, prbegin, prend, cbegin, cend, &c->A));

1702:   a->cmat->preallocated = PETSC_TRUE;
1703:   a->cmat->assembled    = PETSC_TRUE;
1704: #if PetscDefined(HAVE_DEVICE)
1705:   a->cmat->offloadmask = c->A->offloadmask;
1706: #endif
1707:   a->matinuse = cbegin + 1;
1708:   *v          = a->cmat;
1709:   PetscFunctionReturn(PETSC_SUCCESS);
1710: }

1712: static PetscErrorCode MatDenseRestoreSubMatrix_MPIDense(Mat A, Mat *v)
1713: {
1714:   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1715:   Mat_MPIDense *c;

1717:   PetscFunctionBegin;
1718:   PetscCheck(a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseGetSubMatrix() first");
1719:   PetscCheck(a->cmat, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing internal matrix");
1720:   PetscCheck(*v == a->cmat, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Not the matrix obtained from MatDenseGetSubMatrix()");
1721:   a->matinuse = 0;
1722:   c           = (Mat_MPIDense *)a->cmat->data;
1723:   PetscCall(MatDenseRestoreSubMatrix(a->A, &c->A));
1724:   *v = NULL;
1725: #if PetscDefined(HAVE_DEVICE)
1726:   A->offloadmask = a->A->offloadmask;
1727: #endif
1728:   PetscFunctionReturn(PETSC_SUCCESS);
1729: }

1731: /*MC
1732:    MATMPIDENSE - MATMPIDENSE = "mpidense" - A matrix type to be used for distributed dense matrices.

1734:    Options Database Key:
1735: . -mat_type mpidense - sets the matrix type to `MATMPIDENSE` during a call to `MatSetFromOptions()`

1737:   Level: beginner

1739: .seealso: [](ch_matrices), `Mat`, `MatCreateDense()`, `MATSEQDENSE`, `MATDENSE`
1740: M*/
1741: PetscErrorCode MatCreate_MPIDense(Mat mat)
1742: {
1743:   Mat_MPIDense *a;

1745:   PetscFunctionBegin;
1746:   PetscCall(PetscNew(&a));
1747:   mat->data   = (void *)a;
1748:   mat->ops[0] = MatOps_Values;

1750:   mat->insertmode = NOT_SET_VALUES;

1752:   /* build cache for off array entries formed */
1753:   a->donotstash = PETSC_FALSE;

1755:   PetscCall(MatStashCreate_Private(PetscObjectComm((PetscObject)mat), 1, &mat->stash));

1757:   /* stuff used for matrix vector multiply */
1758:   a->lvec        = NULL;
1759:   a->Mvctx       = NULL;
1760:   a->roworiented = PETSC_TRUE;

1762:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetLDA_C", MatDenseGetLDA_MPIDense));
1763:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseSetLDA_C", MatDenseSetLDA_MPIDense));
1764:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArray_C", MatDenseGetArray_MPIDense));
1765:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArray_C", MatDenseRestoreArray_MPIDense));
1766:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayRead_C", MatDenseGetArrayRead_MPIDense));
1767:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayRead_C", MatDenseRestoreArrayRead_MPIDense));
1768:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayWrite_C", MatDenseGetArrayWrite_MPIDense));
1769:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayWrite_C", MatDenseRestoreArrayWrite_MPIDense));
1770:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDensePlaceArray_C", MatDensePlaceArray_MPIDense));
1771:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseResetArray_C", MatDenseResetArray_MPIDense));
1772:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseReplaceArray_C", MatDenseReplaceArray_MPIDense));
1773:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVec_C", MatDenseGetColumnVec_MPIDense));
1774:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVec_C", MatDenseRestoreColumnVec_MPIDense));
1775:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecRead_C", MatDenseGetColumnVecRead_MPIDense));
1776:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecRead_C", MatDenseRestoreColumnVecRead_MPIDense));
1777:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecWrite_C", MatDenseGetColumnVecWrite_MPIDense));
1778:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecWrite_C", MatDenseRestoreColumnVecWrite_MPIDense));
1779:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetSubMatrix_C", MatDenseGetSubMatrix_MPIDense));
1780:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreSubMatrix_C", MatDenseRestoreSubMatrix_MPIDense));
1781:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpiaij_mpidense_C", MatConvert_MPIAIJ_MPIDense));
1782:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpiaij_C", MatConvert_MPIDense_MPIAIJ));
1783: #if PetscDefined(HAVE_ELEMENTAL)
1784:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_elemental_C", MatConvert_MPIDense_Elemental));
1785: #endif
1786: #if PetscDefined(HAVE_SCALAPACK) && (PetscDefined(USE_REAL_SINGLE) || PetscDefined(USE_REAL_DOUBLE))
1787:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_scalapack_C", MatConvert_Dense_ScaLAPACK));
1788: #endif
1789: #if PetscDefined(HAVE_CUDA)
1790:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensecuda_C", MatConvert_MPIDense_MPIDenseCUDA));
1791: #endif
1792:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMPIDenseSetPreallocation_C", MatMPIDenseSetPreallocation_MPIDense));
1793:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidense_C", MatProductSetFromOptions_MPIAIJ_MPIDense));
1794:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaij_C", MatProductSetFromOptions_MPIDense_MPIAIJ));
1795: #if PetscDefined(HAVE_CUDA)
1796:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijcusparse_mpidense_C", MatProductSetFromOptions_MPIAIJ_MPIDense));
1797:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijcusparse_C", MatProductSetFromOptions_MPIDense_MPIAIJ));
1798: #endif
1799: #if PetscDefined(HAVE_HIP)
1800:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensehip_C", MatConvert_MPIDense_MPIDenseHIP));
1801:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijhipsparse_mpidense_C", MatProductSetFromOptions_MPIAIJ_MPIDense));
1802:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijhipsparse_C", MatProductSetFromOptions_MPIDense_MPIAIJ));
1803: #endif
1804:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumn_C", MatDenseGetColumn_MPIDense));
1805:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumn_C", MatDenseRestoreColumn_MPIDense));
1806:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultColumnRange_C", MatMultColumnRange_MPIDense));
1807:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultAddColumnRange_C", MatMultAddColumnRange_MPIDense));
1808:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeColumnRange_C", MatMultHermitianTransposeColumnRange_MPIDense));
1809:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeAddColumnRange_C", MatMultHermitianTransposeAddColumnRange_MPIDense));
1810:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatGetMultPetscSF_C", MatGetMultPetscSF_MPIDense));
1811:   PetscCall(PetscObjectChangeTypeName((PetscObject)mat, MATMPIDENSE));
1812:   PetscFunctionReturn(PETSC_SUCCESS);
1813: }

1815: /*MC
1816:    MATDENSE - MATDENSE = "dense" - A matrix type to be used for dense matrices.

1818:    This matrix type is identical to `MATSEQDENSE` when constructed with a single process communicator,
1819:    and `MATMPIDENSE` otherwise.

1821:    Options Database Key:
1822: . -mat_type dense - sets the matrix type to `MATDENSE` during a call to `MatSetFromOptions()`

1824:   Level: beginner

1826: .seealso: [](ch_matrices), `Mat`, `MATSEQDENSE`, `MATMPIDENSE`, `MATDENSECUDA`, `MATDENSEHIP`
1827: M*/

1829: /*@
1830:   MatMPIDenseSetPreallocation - Sets the array used to store the matrix entries

1832:   Collective

1834:   Input Parameters:
1835: + B    - the matrix
1836: - data - optional location of matrix data.  Set to `NULL` for PETSc
1837:          to control all matrix memory allocation.

1839:   Level: intermediate

1841:   Notes:
1842:   The dense format is fully compatible with standard Fortran
1843:   storage by columns.

1845:   The data input variable is intended primarily for Fortran programmers
1846:   who wish to allocate their own matrix memory space.  Most users should
1847:   set `data` to `NULL`.

1849: .seealso: [](ch_matrices), `Mat`, `MATMPIDENSE`, `MatCreate()`, `MatCreateSeqDense()`, `MatSetValues()`
1850: @*/
1851: PetscErrorCode MatMPIDenseSetPreallocation(Mat B, PetscScalar *data)
1852: {
1853:   PetscFunctionBegin;
1855:   PetscTryMethod(B, "MatMPIDenseSetPreallocation_C", (Mat, PetscScalar *), (B, data));
1856:   PetscFunctionReturn(PETSC_SUCCESS);
1857: }

1859: /*@
1860:   MatDensePlaceArray - Allows one to replace the array in a `MATDENSE` matrix with an
1861:   array provided by the user. This is useful to avoid copying an array
1862:   into a matrix

1864:   Not Collective

1866:   Input Parameters:
1867: + mat   - the matrix
1868: - array - the array in column major order

1870:   Level: developer

1872:   Note:
1873:   Adding `const` to `array` was an oversight, see notes in `VecPlaceArray()`.

1875:   You can return to the original array with a call to `MatDenseResetArray()`. The user is responsible for freeing this array; it will not be
1876:   freed when the matrix is destroyed.

1878: .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MatDenseGetArray()`, `MatDenseResetArray()`, `VecPlaceArray()`, `VecGetArray()`, `VecRestoreArray()`, `VecReplaceArray()`, `VecResetArray()`,
1879:           `MatDenseReplaceArray()`
1880: @*/
1881: PetscErrorCode MatDensePlaceArray(Mat mat, const PetscScalar *array)
1882: {
1883:   PetscFunctionBegin;
1885:   PetscUseMethod(mat, "MatDensePlaceArray_C", (Mat, const PetscScalar *), (mat, array));
1886:   PetscCall(PetscObjectStateIncrease((PetscObject)mat));
1887: #if PetscDefined(HAVE_CUDA) || PetscDefined(HAVE_HIP)
1888:   mat->offloadmask = PETSC_OFFLOAD_CPU;
1889: #endif
1890:   PetscFunctionReturn(PETSC_SUCCESS);
1891: }

1893: /*@
1894:   MatDenseResetArray - Resets the matrix array to that it previously had before the call to `MatDensePlaceArray()`

1896:   Not Collective

1898:   Input Parameter:
1899: . mat - the matrix

1901:   Level: developer

1903:   Note:
1904:   You can only call this after a call to `MatDensePlaceArray()`

1906: .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MatDenseGetArray()`, `MatDensePlaceArray()`, `VecPlaceArray()`, `VecGetArray()`, `VecRestoreArray()`, `VecReplaceArray()`, `VecResetArray()`
1907: @*/
1908: PetscErrorCode MatDenseResetArray(Mat mat)
1909: {
1910:   PetscFunctionBegin;
1912:   PetscUseMethod(mat, "MatDenseResetArray_C", (Mat), (mat));
1913:   PetscCall(PetscObjectStateIncrease((PetscObject)mat));
1914:   PetscFunctionReturn(PETSC_SUCCESS);
1915: }

1917: /*@
1918:   MatDenseReplaceArray - Allows one to replace the array in a dense matrix with an
1919:   array provided by the user. This is useful to avoid copying an array
1920:   into a matrix

1922:   Not Collective

1924:   Input Parameters:
1925: + mat   - the matrix
1926: - array - the array in column major order

1928:   Level: developer

1930:   Note:
1931:   Adding `const` to `array` was an oversight, see notes in `VecPlaceArray()`.

1933:   The memory passed in MUST be obtained with `PetscMalloc()` and CANNOT be
1934:   freed by the user. It will be freed when the matrix is destroyed.

1936: .seealso: [](ch_matrices), `Mat`, `MatDensePlaceArray()`, `MatDenseGetArray()`, `VecReplaceArray()`
1937: @*/
1938: PetscErrorCode MatDenseReplaceArray(Mat mat, const PetscScalar *array)
1939: {
1940:   PetscFunctionBegin;
1942:   PetscUseMethod(mat, "MatDenseReplaceArray_C", (Mat, const PetscScalar *), (mat, array));
1943:   PetscCall(PetscObjectStateIncrease((PetscObject)mat));
1944: #if PetscDefined(HAVE_CUDA) || PetscDefined(HAVE_HIP)
1945:   mat->offloadmask = PETSC_OFFLOAD_CPU;
1946: #endif
1947:   PetscFunctionReturn(PETSC_SUCCESS);
1948: }

1950: /*@
1951:   MatCreateDense - Creates a matrix in `MATDENSE` format.

1953:   Collective

1955:   Input Parameters:
1956: + comm - MPI communicator
1957: . m    - number of local rows (or `PETSC_DECIDE` to have calculated if `M` is given)
1958: . n    - number of local columns (or `PETSC_DECIDE` to have calculated if `N` is given)
1959: . M    - number of global rows (or `PETSC_DECIDE` to have calculated if `m` is given)
1960: . N    - number of global columns (or `PETSC_DECIDE` to have calculated if `n` is given)
1961: - data - optional location of matrix data.  Set data to `NULL` (`PETSC_NULL_SCALAR_ARRAY` for Fortran users) for PETSc
1962:    to control all matrix memory allocation.

1964:   Output Parameter:
1965: . A - the matrix

1967:   Level: intermediate

1969:   Notes:
1970:   The dense format is fully compatible with standard Fortran
1971:   storage by columns.

1973:   Although local portions of the matrix are stored in column-major
1974:   order, the matrix is partitioned across MPI ranks by row.

1976:   The data input variable is intended primarily for Fortran programmers
1977:   who wish to allocate their own matrix memory space.  Most users should
1978:   set `data` to `NULL` (`PETSC_NULL_SCALAR_ARRAY` for Fortran users).

1980:   The user MUST specify either the local or global matrix dimensions
1981:   (possibly both).

1983: .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MatCreate()`, `MatCreateSeqDense()`, `MatSetValues()`
1984: @*/
1985: PetscErrorCode MatCreateDense(MPI_Comm comm, PetscInt m, PetscInt n, PetscInt M, PetscInt N, PetscScalar data[], Mat *A)
1986: {
1987:   PetscFunctionBegin;
1988:   PetscCall(MatCreate(comm, A));
1989:   PetscCall(MatSetSizes(*A, m, n, M, N));
1990:   PetscCall(MatSetType(*A, MATDENSE));
1991:   PetscCall(MatSeqDenseSetPreallocation(*A, data));
1992:   PetscCall(MatMPIDenseSetPreallocation(*A, data));
1993:   PetscFunctionReturn(PETSC_SUCCESS);
1994: }

1996: static PetscErrorCode MatDuplicate_MPIDense(Mat A, MatDuplicateOption cpvalues, Mat *newmat)
1997: {
1998:   Mat           mat;
1999:   Mat_MPIDense *a, *oldmat = (Mat_MPIDense *)A->data;

2001:   PetscFunctionBegin;
2002:   *newmat = NULL;
2003:   PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &mat));
2004:   PetscCall(MatSetSizes(mat, A->rmap->n, A->cmap->n, A->rmap->N, A->cmap->N));
2005:   PetscCall(MatSetType(mat, ((PetscObject)A)->type_name));
2006:   a = (Mat_MPIDense *)mat->data;

2008:   mat->factortype   = A->factortype;
2009:   mat->assembled    = PETSC_TRUE;
2010:   mat->preallocated = PETSC_TRUE;

2012:   mat->insertmode = NOT_SET_VALUES;
2013:   a->donotstash   = oldmat->donotstash;

2015:   PetscCall(PetscLayoutReference(A->rmap, &mat->rmap));
2016:   PetscCall(PetscLayoutReference(A->cmap, &mat->cmap));

2018:   PetscCall(MatDuplicate(oldmat->A, cpvalues, &a->A));

2020:   *newmat = mat;
2021:   PetscFunctionReturn(PETSC_SUCCESS);
2022: }

2024: static PetscErrorCode MatLoad_MPIDense(Mat newMat, PetscViewer viewer)
2025: {
2026:   PetscBool isbinary;
2027: #if PetscDefined(HAVE_HDF5)
2028:   PetscBool ishdf5;
2029: #endif

2031:   PetscFunctionBegin;
2034:   /* force binary viewer to load .info file if it has not yet done so */
2035:   PetscCall(PetscViewerSetUp(viewer));
2036:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &isbinary));
2037: #if PetscDefined(HAVE_HDF5)
2038:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERHDF5, &ishdf5));
2039: #endif
2040:   if (isbinary) {
2041:     PetscCall(MatLoad_Dense_Binary(newMat, viewer));
2042: #if PetscDefined(HAVE_HDF5)
2043:   } else if (ishdf5) {
2044:     PetscCall(MatLoad_Dense_HDF5(newMat, viewer));
2045: #endif
2046:   } else SETERRQ(PetscObjectComm((PetscObject)newMat), PETSC_ERR_SUP, "Viewer type %s not yet supported for reading %s matrices", ((PetscObject)viewer)->type_name, ((PetscObject)newMat)->type_name);
2047:   PetscFunctionReturn(PETSC_SUCCESS);
2048: }

2050: static PetscErrorCode MatEqual_MPIDense(Mat A, Mat B, PetscBool *flag)
2051: {
2052:   Mat_MPIDense *matB = (Mat_MPIDense *)B->data, *matA = (Mat_MPIDense *)A->data;
2053:   Mat           a, b;

2055:   PetscFunctionBegin;
2056:   a = matA->A;
2057:   b = matB->A;
2058:   PetscCall(MatEqual(a, b, flag));
2059:   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, flag, 1, MPI_C_BOOL, MPI_LAND, PetscObjectComm((PetscObject)A)));
2060:   PetscFunctionReturn(PETSC_SUCCESS);
2061: }

2063: static PetscErrorCode MatProductCtxDestroy_MatTransMatMult_MPIDense_MPIDense(PetscCtxRt data)
2064: {
2065:   MatProductCtx_TransMatMultDense *atb = *(MatProductCtx_TransMatMultDense **)data;

2067:   PetscFunctionBegin;
2068:   PetscCall(PetscFree2(atb->sendbuf, atb->recvcounts));
2069:   PetscCall(MatDestroy(&atb->atb));
2070:   PetscCall(PetscFree(atb));
2071:   PetscFunctionReturn(PETSC_SUCCESS);
2072: }

2074: static PetscErrorCode MatProductCtxDestroy_MatMatTransMult_MPIDense_MPIDense(PetscCtxRt data)
2075: {
2076:   MatProductCtx_MatTransMultDense *abt = *(MatProductCtx_MatTransMultDense **)data;

2078:   PetscFunctionBegin;
2079:   PetscCall(PetscFree2(abt->buf[0], abt->buf[1]));
2080:   PetscCall(PetscFree2(abt->recvcounts, abt->recvdispls));
2081:   PetscCall(PetscFree(abt));
2082:   PetscFunctionReturn(PETSC_SUCCESS);
2083: }

2085: static PetscErrorCode MatTransposeMatMultNumeric_MPIDense_MPIDense(Mat A, Mat B, Mat C)
2086: {
2087:   Mat_MPIDense                    *a = (Mat_MPIDense *)A->data, *b = (Mat_MPIDense *)B->data, *c = (Mat_MPIDense *)C->data;
2088:   MatProductCtx_TransMatMultDense *atb;
2089:   MPI_Comm                         comm;
2090:   PetscMPIInt                      size, *recvcounts;
2091:   PetscScalar                     *carray, *sendbuf;
2092:   const PetscScalar               *atbarray;
2093:   PetscInt                         i, cN = C->cmap->N, proc, k, j, lda;
2094:   const PetscInt                  *ranges;

2096:   PetscFunctionBegin;
2097:   MatCheckProduct(C, 3);
2098:   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
2099:   atb        = (MatProductCtx_TransMatMultDense *)C->product->data;
2100:   recvcounts = atb->recvcounts;
2101:   sendbuf    = atb->sendbuf;

2103:   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
2104:   PetscCallMPI(MPI_Comm_size(comm, &size));

2106:   /* compute atbarray = aseq^T * bseq */
2107:   PetscCall(MatTransposeMatMult(a->A, b->A, atb->atb ? MAT_REUSE_MATRIX : MAT_INITIAL_MATRIX, PETSC_DETERMINE, &atb->atb));

2109:   PetscCall(MatGetOwnershipRanges(C, &ranges));

2111:   if (ranges[1] == C->rmap->N) {
2112:     /* all of the values are being reduced to rank 0: optimize this case to use MPI_Reduce and GPU aware MPI if available */
2113:     PetscInt           atb_lda, c_lda;
2114:     Mat                atb_local = atb->atb;
2115:     Mat                atb_alloc = NULL;
2116:     Mat                c_local   = c->A;
2117:     Mat                c_alloc   = NULL;
2118:     PetscMemType       atb_memtype, c_memtype;
2119:     const PetscScalar *atb_array = NULL;
2120:     MPI_Datatype       vector_type;
2121:     PetscScalar       *c_array = NULL;
2122:     PetscMPIInt        rank;

2124:     PetscCallMPI(MPI_Comm_rank(comm, &rank));

2126:     PetscCall(MatDenseGetLDA(atb_local, &atb_lda));
2127:     if (atb_lda != C->rmap->N) {
2128:       // copy atb to a matrix that will have lda == the number of rows
2129:       PetscCall(MatDuplicate(atb_local, MAT_DO_NOT_COPY_VALUES, &atb_alloc));
2130:       PetscCall(MatCopy(atb_local, atb_alloc, DIFFERENT_NONZERO_PATTERN));
2131:       atb_local = atb_alloc;
2132:     }

2134:     if (rank == 0) {
2135:       PetscCall(MatDenseGetLDA(c_local, &c_lda));
2136:       if (c_lda != C->rmap->N) {
2137:         // copy c to a matrix that will have lda == the number of rows
2138:         PetscCall(MatDuplicate(c_local, MAT_DO_NOT_COPY_VALUES, &c_alloc));
2139:         c_local = c_alloc;
2140:       }
2141:       PetscCall(MatZeroEntries(c_local));
2142:     }
2143:     /* atb_local and c_local have nrows = lda = A->cmap->N and ncols =
2144:      * B->cmap->N: use the a->Mvctx to use the best reduction method */
2145:     if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
2146:     vector_type = MPIU_SCALAR;
2147:     if (B->cmap->N > 1) {
2148:       PetscMPIInt mpi_N;

2150:       PetscCall(PetscMPIIntCast(B->cmap->N, &mpi_N));
2151:       PetscCallMPI(MPI_Type_contiguous(mpi_N, MPIU_SCALAR, &vector_type));
2152:       PetscCallMPI(MPI_Type_commit(&vector_type));
2153:     }
2154:     PetscCall(MatDenseGetArrayReadAndMemType(atb_local, &atb_array, &atb_memtype));
2155:     PetscCall(MatDenseGetArrayWriteAndMemType(c_local, &c_array, &c_memtype));
2156:     PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, vector_type, atb_memtype, atb_array, c_memtype, c_array, MPIU_SUM));
2157:     PetscCall(PetscSFReduceEnd(a->Mvctx, vector_type, atb_array, c_array, MPIU_SUM));
2158:     PetscCall(MatDenseRestoreArrayWriteAndMemType(c_local, &c_array));
2159:     PetscCall(MatDenseRestoreArrayReadAndMemType(atb_local, &atb_array));
2160:     if (rank == 0 && c_local != c->A) PetscCall(MatCopy(c_local, c->A, DIFFERENT_NONZERO_PATTERN));
2161:     if (B->cmap->N > 1) PetscCallMPI(MPI_Type_free(&vector_type));
2162:     PetscCall(MatDestroy(&atb_alloc));
2163:     PetscCall(MatDestroy(&c_alloc));
2164:     PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
2165:     PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
2166:     PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
2167:     PetscFunctionReturn(PETSC_SUCCESS);
2168:   }

2170:   /* arrange atbarray into sendbuf */
2171:   PetscCall(MatDenseGetArrayRead(atb->atb, &atbarray));
2172:   PetscCall(MatDenseGetLDA(atb->atb, &lda));
2173:   for (proc = 0, k = 0; proc < size; proc++) {
2174:     for (j = 0; j < cN; j++) {
2175:       for (i = ranges[proc]; i < ranges[proc + 1]; i++) sendbuf[k++] = atbarray[i + j * lda];
2176:     }
2177:   }
2178:   PetscCall(MatDenseRestoreArrayRead(atb->atb, &atbarray));

2180:   /* sum all atbarray to local values of C */
2181:   PetscCall(MatDenseGetArrayWrite(c->A, &carray));
2182:   PetscCallMPI(MPI_Reduce_scatter(sendbuf, carray, recvcounts, MPIU_SCALAR, MPIU_SUM, comm));
2183:   PetscCall(MatDenseRestoreArrayWrite(c->A, &carray));
2184:   PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
2185:   PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
2186:   PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
2187:   PetscFunctionReturn(PETSC_SUCCESS);
2188: }

2190: static PetscErrorCode MatTransposeMatMultSymbolic_MPIDense_MPIDense(Mat A, Mat B, PetscReal fill, Mat C)
2191: {
2192:   MPI_Comm                         comm;
2193:   PetscMPIInt                      size;
2194:   PetscInt                         cm = A->cmap->n, cM, cN = B->cmap->N;
2195:   MatProductCtx_TransMatMultDense *atb;
2196:   PetscBool                        cisdense = PETSC_FALSE;
2197:   const PetscInt                  *ranges;

2199:   PetscFunctionBegin;
2200:   MatCheckProduct(C, 4);
2201:   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
2202:   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
2203:   PetscCheck(A->rmap->rstart == B->rmap->rstart && A->rmap->rend == B->rmap->rend, comm, PETSC_ERR_ARG_SIZ, "Matrix local dimensions are incompatible, A (%" PetscInt_FMT ", %" PetscInt_FMT ") != B (%" PetscInt_FMT ",%" PetscInt_FMT ")", A->rmap->rstart,
2204:              A->rmap->rend, B->rmap->rstart, B->rmap->rend);

2206:   /* create matrix product C */
2207:   PetscCall(MatSetSizes(C, cm, B->cmap->n, A->cmap->N, B->cmap->N));
2208: #if PetscDefined(HAVE_CUDA)
2209:   PetscCall(PetscObjectTypeCompareAny((PetscObject)C, &cisdense, MATMPIDENSE, MATMPIDENSECUDA, ""));
2210: #endif
2211: #if PetscDefined(HAVE_HIP)
2212:   PetscCall(PetscObjectTypeCompareAny((PetscObject)C, &cisdense, MATMPIDENSE, MATMPIDENSEHIP, ""));
2213: #endif
2214:   if (!cisdense) PetscCall(MatSetType(C, ((PetscObject)A)->type_name));
2215:   PetscCall(MatSetUp(C));

2217:   /* create data structure for reuse C */
2218:   PetscCallMPI(MPI_Comm_size(comm, &size));
2219:   PetscCall(PetscNew(&atb));
2220:   cM = C->rmap->N;
2221:   PetscCall(PetscMalloc2(cM * cN, &atb->sendbuf, size, &atb->recvcounts));
2222:   PetscCall(MatGetOwnershipRanges(C, &ranges));
2223:   for (PetscMPIInt i = 0; i < size; i++) PetscCall(PetscMPIIntCast((ranges[i + 1] - ranges[i]) * cN, &atb->recvcounts[i]));
2224:   C->product->data    = atb;
2225:   C->product->destroy = MatProductCtxDestroy_MatTransMatMult_MPIDense_MPIDense;
2226:   PetscFunctionReturn(PETSC_SUCCESS);
2227: }

2229: static PetscErrorCode MatMatTransposeMultSymbolic_MPIDense_MPIDense(Mat A, Mat B, PetscReal fill, Mat C)
2230: {
2231:   MPI_Comm                         comm;
2232:   PetscMPIInt                      i, size;
2233:   PetscInt                         maxRows, bufsiz;
2234:   PetscMPIInt                      tag;
2235:   PetscInt                         alg;
2236:   MatProductCtx_MatTransMultDense *abt;
2237:   Mat_Product                     *product = C->product;
2238:   PetscBool                        flg;

2240:   PetscFunctionBegin;
2241:   MatCheckProduct(C, 4);
2242:   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
2243:   /* check local size of A and B */
2244:   PetscCheck(A->cmap->n == B->cmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Matrix local column dimensions are incompatible, A (%" PetscInt_FMT ") != B (%" PetscInt_FMT ")", A->cmap->n, B->cmap->n);

2246:   PetscCall(PetscStrcmp(product->alg, "allgatherv", &flg));
2247:   alg = flg ? 0 : 1;

2249:   /* setup matrix product C */
2250:   PetscCall(MatSetSizes(C, A->rmap->n, B->rmap->n, A->rmap->N, B->rmap->N));
2251:   PetscCall(MatSetType(C, MATMPIDENSE));
2252:   PetscCall(MatSetUp(C));
2253:   PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag));

2255:   /* create data structure for reuse C */
2256:   PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
2257:   PetscCallMPI(MPI_Comm_size(comm, &size));
2258:   PetscCall(PetscNew(&abt));
2259:   abt->tag = tag;
2260:   abt->alg = alg;
2261:   switch (alg) {
2262:   case 1: /* alg: "cyclic" */
2263:     for (maxRows = 0, i = 0; i < size; i++) maxRows = PetscMax(maxRows, B->rmap->range[i + 1] - B->rmap->range[i]);
2264:     bufsiz = A->cmap->N * maxRows;
2265:     PetscCall(PetscMalloc2(bufsiz, &abt->buf[0], bufsiz, &abt->buf[1]));
2266:     break;
2267:   default: /* alg: "allgatherv" */
2268:     PetscCall(PetscMalloc2(B->rmap->n * B->cmap->N, &abt->buf[0], B->rmap->N * B->cmap->N, &abt->buf[1]));
2269:     PetscCall(PetscMalloc2(size, &abt->recvcounts, size + 1, &abt->recvdispls));
2270:     for (i = 0; i <= size; i++) PetscCall(PetscMPIIntCast(B->rmap->range[i] * A->cmap->N, &abt->recvdispls[i]));
2271:     for (i = 0; i < size; i++) PetscCall(PetscMPIIntCast(abt->recvdispls[i + 1] - abt->recvdispls[i], &abt->recvcounts[i]));
2272:     break;
2273:   }

2275:   C->product->data                = abt;
2276:   C->product->destroy             = MatProductCtxDestroy_MatMatTransMult_MPIDense_MPIDense;
2277:   C->ops->mattransposemultnumeric = MatMatTransposeMultNumeric_MPIDense_MPIDense;
2278:   PetscFunctionReturn(PETSC_SUCCESS);
2279: }

2281: static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense_Cyclic(Mat A, Mat B, Mat C)
2282: {
2283:   Mat_MPIDense                    *a = (Mat_MPIDense *)A->data, *b = (Mat_MPIDense *)B->data, *c = (Mat_MPIDense *)C->data;
2284:   MatProductCtx_MatTransMultDense *abt;
2285:   MPI_Comm                         comm;
2286:   PetscMPIInt                      rank, size, sendto, recvfrom, recvisfrom;
2287:   PetscScalar                     *sendbuf, *recvbuf = NULL, *cv;
2288:   PetscInt                         i, cK             = A->cmap->N, sendsiz, recvsiz, k, j, bn;
2289:   PetscScalar                      _DOne = 1.0, _DZero = 0.0;
2290:   const PetscScalar               *av, *bv;
2291:   PetscBLASInt                     cm, cn, ck, alda, blda = 0, clda;
2292:   MPI_Request                      reqs[2];
2293:   const PetscInt                  *ranges;

2295:   PetscFunctionBegin;
2296:   MatCheckProduct(C, 3);
2297:   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
2298:   abt = (MatProductCtx_MatTransMultDense *)C->product->data;
2299:   PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
2300:   PetscCallMPI(MPI_Comm_rank(comm, &rank));
2301:   PetscCallMPI(MPI_Comm_size(comm, &size));
2302:   PetscCall(MatDenseGetArrayRead(a->A, &av));
2303:   PetscCall(MatDenseGetArrayRead(b->A, &bv));
2304:   PetscCall(MatDenseGetArrayWrite(c->A, &cv));
2305:   PetscCall(MatDenseGetLDA(a->A, &i));
2306:   PetscCall(PetscBLASIntCast(i, &alda));
2307:   PetscCall(MatDenseGetLDA(b->A, &i));
2308:   PetscCall(PetscBLASIntCast(i, &blda));
2309:   PetscCall(MatDenseGetLDA(c->A, &i));
2310:   PetscCall(PetscBLASIntCast(i, &clda));
2311:   PetscCall(MatGetOwnershipRanges(B, &ranges));
2312:   bn = B->rmap->n;
2313:   if (blda == bn) {
2314:     sendbuf = (PetscScalar *)bv;
2315:   } else {
2316:     sendbuf = abt->buf[0];
2317:     for (k = 0, i = 0; i < cK; i++) {
2318:       for (j = 0; j < bn; j++, k++) sendbuf[k] = bv[i * blda + j];
2319:     }
2320:   }
2321:   if (size > 1) {
2322:     sendto   = (rank + size - 1) % size;
2323:     recvfrom = (rank + size + 1) % size;
2324:   } else {
2325:     sendto = recvfrom = 0;
2326:   }
2327:   PetscCall(PetscBLASIntCast(cK, &ck));
2328:   PetscCall(PetscBLASIntCast(c->A->rmap->n, &cm));
2329:   recvisfrom = rank;
2330:   for (i = 0; i < size; i++) {
2331:     /* we have finished receiving in sending, bufs can be read/modified */
2332:     PetscMPIInt nextrecvisfrom = (recvisfrom + 1) % size; /* which process the next recvbuf will originate on */
2333:     PetscInt    nextbn         = ranges[nextrecvisfrom + 1] - ranges[nextrecvisfrom];

2335:     if (nextrecvisfrom != rank) {
2336:       /* start the cyclic sends from sendbuf, to recvbuf (which will switch to sendbuf) */
2337:       sendsiz = cK * bn;
2338:       recvsiz = cK * nextbn;
2339:       recvbuf = (i & 1) ? abt->buf[0] : abt->buf[1];
2340:       PetscCallMPI(MPIU_Isend(sendbuf, sendsiz, MPIU_SCALAR, sendto, abt->tag, comm, &reqs[0]));
2341:       PetscCallMPI(MPIU_Irecv(recvbuf, recvsiz, MPIU_SCALAR, recvfrom, abt->tag, comm, &reqs[1]));
2342:     }

2344:     /* local aseq * sendbuf^T */
2345:     PetscCall(PetscBLASIntCast(ranges[recvisfrom + 1] - ranges[recvisfrom], &cn));
2346:     if (cm && cn && ck) PetscCallBLAS("BLASgemm", BLASgemm_("N", "T", &cm, &cn, &ck, &_DOne, av, &alda, sendbuf, &cn, &_DZero, cv + clda * ranges[recvisfrom], &clda));

2348:     if (nextrecvisfrom != rank) {
2349:       /* wait for the sends and receives to complete, swap sendbuf and recvbuf */
2350:       PetscCallMPI(MPI_Waitall(2, reqs, MPI_STATUSES_IGNORE));
2351:     }
2352:     bn         = nextbn;
2353:     recvisfrom = nextrecvisfrom;
2354:     sendbuf    = recvbuf;
2355:   }
2356:   PetscCall(MatDenseRestoreArrayRead(a->A, &av));
2357:   PetscCall(MatDenseRestoreArrayRead(b->A, &bv));
2358:   PetscCall(MatDenseRestoreArrayWrite(c->A, &cv));
2359:   PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
2360:   PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
2361:   PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
2362:   PetscFunctionReturn(PETSC_SUCCESS);
2363: }

2365: static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense_Allgatherv(Mat A, Mat B, Mat C)
2366: {
2367:   Mat_MPIDense                    *a = (Mat_MPIDense *)A->data, *b = (Mat_MPIDense *)B->data, *c = (Mat_MPIDense *)C->data;
2368:   MatProductCtx_MatTransMultDense *abt;
2369:   MPI_Comm                         comm;
2370:   PetscMPIInt                      size, ibn;
2371:   PetscScalar                     *cv, *sendbuf, *recvbuf;
2372:   const PetscScalar               *av, *bv;
2373:   PetscInt                         blda, i, cK = A->cmap->N, k, j, bn;
2374:   PetscScalar                      _DOne = 1.0, _DZero = 0.0;
2375:   PetscBLASInt                     cm, cn, ck, alda, clda;

2377:   PetscFunctionBegin;
2378:   MatCheckProduct(C, 3);
2379:   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
2380:   abt = (MatProductCtx_MatTransMultDense *)C->product->data;
2381:   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
2382:   PetscCallMPI(MPI_Comm_size(comm, &size));
2383:   PetscCall(MatDenseGetArrayRead(a->A, &av));
2384:   PetscCall(MatDenseGetArrayRead(b->A, &bv));
2385:   PetscCall(MatDenseGetArrayWrite(c->A, &cv));
2386:   PetscCall(MatDenseGetLDA(a->A, &i));
2387:   PetscCall(PetscBLASIntCast(i, &alda));
2388:   PetscCall(MatDenseGetLDA(b->A, &blda));
2389:   PetscCall(MatDenseGetLDA(c->A, &i));
2390:   PetscCall(PetscBLASIntCast(i, &clda));
2391:   /* copy transpose of B into buf[0] */
2392:   bn      = B->rmap->n;
2393:   sendbuf = abt->buf[0];
2394:   recvbuf = abt->buf[1];
2395:   for (k = 0, j = 0; j < bn; j++) {
2396:     for (i = 0; i < cK; i++, k++) sendbuf[k] = bv[i * blda + j];
2397:   }
2398:   PetscCall(MatDenseRestoreArrayRead(b->A, &bv));
2399:   PetscCall(PetscMPIIntCast(bn * cK, &ibn));
2400:   PetscCallMPI(MPI_Allgatherv(sendbuf, ibn, MPIU_SCALAR, recvbuf, abt->recvcounts, abt->recvdispls, MPIU_SCALAR, comm));
2401:   PetscCall(PetscBLASIntCast(cK, &ck));
2402:   PetscCall(PetscBLASIntCast(c->A->rmap->n, &cm));
2403:   PetscCall(PetscBLASIntCast(c->A->cmap->n, &cn));
2404:   if (cm && cn && ck) PetscCallBLAS("BLASgemm", BLASgemm_("N", "N", &cm, &cn, &ck, &_DOne, av, &alda, recvbuf, &ck, &_DZero, cv, &clda));
2405:   PetscCall(MatDenseRestoreArrayRead(a->A, &av));
2406:   PetscCall(MatDenseRestoreArrayRead(b->A, &bv));
2407:   PetscCall(MatDenseRestoreArrayWrite(c->A, &cv));
2408:   PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
2409:   PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
2410:   PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
2411:   PetscFunctionReturn(PETSC_SUCCESS);
2412: }

2414: static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense(Mat A, Mat B, Mat C)
2415: {
2416:   MatProductCtx_MatTransMultDense *abt;

2418:   PetscFunctionBegin;
2419:   MatCheckProduct(C, 3);
2420:   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
2421:   abt = (MatProductCtx_MatTransMultDense *)C->product->data;
2422:   switch (abt->alg) {
2423:   case 1:
2424:     PetscCall(MatMatTransposeMultNumeric_MPIDense_MPIDense_Cyclic(A, B, C));
2425:     break;
2426:   default:
2427:     PetscCall(MatMatTransposeMultNumeric_MPIDense_MPIDense_Allgatherv(A, B, C));
2428:     break;
2429:   }
2430:   PetscFunctionReturn(PETSC_SUCCESS);
2431: }

2433: static PetscErrorCode MatProductCtxDestroy_MatMatMult_MPIDense_MPIDense(PetscCtxRt data)
2434: {
2435:   MatProductCtx_MatMultDense *ab = *(MatProductCtx_MatMultDense **)data;

2437:   PetscFunctionBegin;
2438:   PetscCall(MatDestroy(&ab->Ce));
2439:   PetscCall(MatDestroy(&ab->Ae));
2440:   PetscCall(MatDestroy(&ab->Be));
2441:   PetscCall(PetscFree(ab));
2442:   PetscFunctionReturn(PETSC_SUCCESS);
2443: }

2445: static PetscErrorCode MatMatMultNumeric_MPIDense_MPIDense(Mat A, Mat B, Mat C)
2446: {
2447:   MatProductCtx_MatMultDense *ab;
2448:   Mat_MPIDense               *mdn = (Mat_MPIDense *)A->data;
2449:   Mat_MPIDense               *b   = (Mat_MPIDense *)B->data;

2451:   PetscFunctionBegin;
2452:   MatCheckProduct(C, 3);
2453:   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Missing product data");
2454:   ab = (MatProductCtx_MatMultDense *)C->product->data;
2455:   if (ab->Ae && ab->Ce) {
2456: #if PetscDefined(HAVE_ELEMENTAL)
2457:     PetscCall(MatConvert_MPIDense_Elemental(A, MATELEMENTAL, MAT_REUSE_MATRIX, &ab->Ae));
2458:     PetscCall(MatConvert_MPIDense_Elemental(B, MATELEMENTAL, MAT_REUSE_MATRIX, &ab->Be));
2459:     PetscCall(MatMatMultNumeric_Elemental(ab->Ae, ab->Be, ab->Ce));
2460:     PetscCall(MatConvert(ab->Ce, MATMPIDENSE, MAT_REUSE_MATRIX, &C));
2461: #else
2462:     SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "PETSC_HAVE_ELEMENTAL not defined");
2463: #endif
2464:   } else {
2465:     MPI_Comm           comm;
2466:     const PetscScalar *read;
2467:     PetscScalar       *write;
2468:     PetscInt           lda;
2469:     const PetscInt    *ranges;
2470:     PetscMPIInt        size;

2472:     if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A)); /* cannot be done during the symbolic phase because of possible calls to MatProductReplaceMats() */
2473:     comm = PetscObjectComm((PetscObject)B);
2474:     PetscCallMPI(MPI_Comm_size(comm, &size));
2475:     PetscCall(PetscLayoutGetRanges(B->rmap, &ranges));
2476:     if (ranges[1] == ranges[size]) {
2477:       // optimize for the case where the B matrix is broadcast from rank 0
2478:       PetscInt           b_lda, be_lda;
2479:       Mat                b_local  = b->A;
2480:       Mat                b_alloc  = NULL;
2481:       Mat                be_local = ab->Be;
2482:       Mat                be_alloc = NULL;
2483:       PetscMemType       b_memtype, be_memtype;
2484:       const PetscScalar *b_array = NULL;
2485:       MPI_Datatype       vector_type;
2486:       PetscScalar       *be_array = NULL;
2487:       PetscMPIInt        rank;

2489:       PetscCallMPI(MPI_Comm_rank(comm, &rank));
2490:       PetscCall(MatDenseGetLDA(be_local, &be_lda));
2491:       if (be_lda != B->rmap->N) {
2492:         PetscCall(MatDuplicate(be_local, MAT_DO_NOT_COPY_VALUES, &be_alloc));
2493:         be_local = be_alloc;
2494:       }

2496:       if (rank == 0) {
2497:         PetscCall(MatDenseGetLDA(b_local, &b_lda));
2498:         if (b_lda != B->rmap->N) {
2499:           PetscCall(MatDuplicate(b_local, MAT_DO_NOT_COPY_VALUES, &b_alloc));
2500:           PetscCall(MatCopy(b_local, b_alloc, DIFFERENT_NONZERO_PATTERN));
2501:           b_local = b_alloc;
2502:         }
2503:       }
2504:       vector_type = MPIU_SCALAR;
2505:       if (B->cmap->N > 1) {
2506:         PetscMPIInt mpi_N;

2508:         PetscCall(PetscMPIIntCast(B->cmap->N, &mpi_N));
2509:         PetscCallMPI(MPI_Type_contiguous(mpi_N, MPIU_SCALAR, &vector_type));
2510:         PetscCallMPI(MPI_Type_commit(&vector_type));
2511:       }
2512:       PetscCall(MatDenseGetArrayReadAndMemType(b_local, &b_array, &b_memtype));
2513:       PetscCall(MatDenseGetArrayWriteAndMemType(be_local, &be_array, &be_memtype));
2514:       PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, vector_type, b_memtype, b_array, be_memtype, be_array, MPI_REPLACE));
2515:       PetscCall(PetscSFBcastEnd(mdn->Mvctx, vector_type, b_array, be_array, MPI_REPLACE));
2516:       PetscCall(MatDenseRestoreArrayWriteAndMemType(be_local, &be_array));
2517:       PetscCall(MatDenseRestoreArrayReadAndMemType(b_local, &b_array));
2518:       if (be_local != ab->Be) PetscCall(MatCopy(be_local, ab->Be, DIFFERENT_NONZERO_PATTERN));
2519:       if (B->cmap->N > 1) PetscCallMPI(MPI_Type_free(&vector_type));
2520:       PetscCall(MatDestroy(&be_alloc));
2521:       PetscCall(MatDestroy(&b_alloc));
2522:     } else {
2523:       PetscCall(MatDenseGetLDA(B, &lda));
2524:       PetscCall(MatDenseGetArrayRead(B, &read));
2525:       PetscCall(MatDenseGetArrayWrite(ab->Be, &write));
2526:       for (PetscInt i = 0; i < C->cmap->N; ++i) {
2527:         PetscCall(PetscSFBcastBegin(mdn->Mvctx, MPIU_SCALAR, read + i * lda, write + i * ab->Be->rmap->n, MPI_REPLACE));
2528:         PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, read + i * lda, write + i * ab->Be->rmap->n, MPI_REPLACE));
2529:       }
2530:       PetscCall(MatDenseRestoreArrayWrite(ab->Be, &write));
2531:       PetscCall(MatDenseRestoreArrayRead(B, &read));
2532:     }
2533:     PetscCall(MatMatMultNumeric_SeqDense_SeqDense(((Mat_MPIDense *)A->data)->A, ab->Be, ((Mat_MPIDense *)C->data)->A));
2534:   }
2535:   PetscFunctionReturn(PETSC_SUCCESS);
2536: }

2538: static PetscErrorCode MatMatMultSymbolic_MPIDense_MPIDense(Mat A, Mat B, PetscReal fill, Mat C)
2539: {
2540:   Mat_Product                *product = C->product;
2541:   PetscInt                    alg;
2542:   MatProductCtx_MatMultDense *ab;
2543:   PetscBool                   flg;

2545:   PetscFunctionBegin;
2546:   MatCheckProduct(C, 4);
2547:   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
2548:   /* check local size of A and B */
2549:   PetscCheck(A->cmap->rstart == B->rmap->rstart && A->cmap->rend == B->rmap->rend, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_SIZ, "Matrix local dimensions are incompatible, A (%" PetscInt_FMT ", %" PetscInt_FMT ") != B (%" PetscInt_FMT ", %" PetscInt_FMT ")",
2550:              A->rmap->rstart, A->rmap->rend, B->rmap->rstart, B->rmap->rend);

2552:   PetscCall(PetscStrcmp(product->alg, "petsc", &flg));
2553:   alg = flg ? 0 : 1;

2555:   /* setup C */
2556:   PetscCall(MatSetSizes(C, A->rmap->n, B->cmap->n, A->rmap->N, B->cmap->N));
2557:   PetscCall(MatSetType(C, MATMPIDENSE));
2558:   PetscCall(MatSetUp(C));

2560:   /* create data structure for reuse Cdense */
2561:   PetscCall(PetscNew(&ab));

2563:   switch (alg) {
2564:   case 1: /* alg: "elemental" */
2565: #if PetscDefined(HAVE_ELEMENTAL)
2566:     /* create elemental matrices Ae and Be */
2567:     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &ab->Ae));
2568:     PetscCall(MatSetSizes(ab->Ae, PETSC_DECIDE, PETSC_DECIDE, A->rmap->N, A->cmap->N));
2569:     PetscCall(MatSetType(ab->Ae, MATELEMENTAL));
2570:     PetscCall(MatSetUp(ab->Ae));
2571:     PetscCall(MatSetOption(ab->Ae, MAT_ROW_ORIENTED, PETSC_FALSE));

2573:     PetscCall(MatCreate(PetscObjectComm((PetscObject)B), &ab->Be));
2574:     PetscCall(MatSetSizes(ab->Be, PETSC_DECIDE, PETSC_DECIDE, B->rmap->N, B->cmap->N));
2575:     PetscCall(MatSetType(ab->Be, MATELEMENTAL));
2576:     PetscCall(MatSetUp(ab->Be));
2577:     PetscCall(MatSetOption(ab->Be, MAT_ROW_ORIENTED, PETSC_FALSE));

2579:     /* compute symbolic Ce = Ae*Be */
2580:     PetscCall(MatCreate(PetscObjectComm((PetscObject)C), &ab->Ce));
2581:     PetscCall(MatMatMultSymbolic_Elemental(ab->Ae, ab->Be, fill, ab->Ce));
2582: #else
2583:     SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "PETSC_HAVE_ELEMENTAL not defined");
2584: #endif
2585:     break;
2586:   default: /* alg: "petsc" */
2587:     ab->Ae = NULL;
2588:     PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, A->cmap->N, B->cmap->N, NULL, &ab->Be));
2589:     ab->Ce = NULL;
2590:     break;
2591:   }

2593:   C->product->data       = ab;
2594:   C->product->destroy    = MatProductCtxDestroy_MatMatMult_MPIDense_MPIDense;
2595:   C->ops->matmultnumeric = MatMatMultNumeric_MPIDense_MPIDense;
2596:   PetscFunctionReturn(PETSC_SUCCESS);
2597: }

2599: static PetscErrorCode MatProductSetFromOptions_MPIDense_AB(Mat C)
2600: {
2601:   Mat_Product *product     = C->product;
2602:   const char  *algTypes[2] = {"petsc", "elemental"};
2603:   PetscInt     alg, nalg = PetscDefined(HAVE_ELEMENTAL) ? 2 : 1;
2604:   PetscBool    flg = PETSC_FALSE;

2606:   PetscFunctionBegin;
2607:   /* Set default algorithm */
2608:   alg = 0; /* default is PETSc */
2609:   PetscCall(PetscStrcmp(product->alg, "default", &flg));
2610:   if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));

2612:   /* Get runtime option */
2613:   PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatProduct_AB", "Mat");
2614:   PetscCall(PetscOptionsEList("-mat_product_algorithm", "Algorithmic approach", "MatProduct_AB", algTypes, nalg, algTypes[alg], &alg, &flg));
2615:   PetscOptionsEnd();
2616:   if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));

2618:   C->ops->matmultsymbolic = MatMatMultSymbolic_MPIDense_MPIDense;
2619:   C->ops->productsymbolic = MatProductSymbolic_AB;
2620:   PetscFunctionReturn(PETSC_SUCCESS);
2621: }

2623: static PetscErrorCode MatProductSetFromOptions_MPIDense_AtB(Mat C)
2624: {
2625:   Mat_Product *product = C->product;
2626:   Mat          A = product->A, B = product->B;

2628:   PetscFunctionBegin;
2629:   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, (%" PetscInt_FMT ", %" PetscInt_FMT ") != (%" PetscInt_FMT ",%" PetscInt_FMT ")",
2630:              A->rmap->rstart, A->rmap->rend, B->rmap->rstart, B->rmap->rend);
2631:   C->ops->transposematmultsymbolic = MatTransposeMatMultSymbolic_MPIDense_MPIDense;
2632:   C->ops->productsymbolic          = MatProductSymbolic_AtB;
2633:   PetscFunctionReturn(PETSC_SUCCESS);
2634: }

2636: static PetscErrorCode MatProductSetFromOptions_MPIDense_ABt(Mat C)
2637: {
2638:   Mat_Product *product     = C->product;
2639:   const char  *algTypes[2] = {"allgatherv", "cyclic"};
2640:   PetscInt     alg, nalg = 2;
2641:   PetscBool    flg = PETSC_FALSE;

2643:   PetscFunctionBegin;
2644:   /* Set default algorithm */
2645:   alg = 0; /* default is allgatherv */
2646:   PetscCall(PetscStrcmp(product->alg, "default", &flg));
2647:   if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));

2649:   /* Get runtime option */
2650:   if (product->api_user) {
2651:     PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatMatTransposeMult", "Mat");
2652:     PetscCall(PetscOptionsEList("-matmattransmult_mpidense_mpidense_via", "Algorithmic approach", "MatMatTransposeMult", algTypes, nalg, algTypes[alg], &alg, &flg));
2653:     PetscOptionsEnd();
2654:   } else {
2655:     PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatProduct_ABt", "Mat");
2656:     PetscCall(PetscOptionsEList("-mat_product_algorithm", "Algorithmic approach", "MatProduct_ABt", algTypes, nalg, algTypes[alg], &alg, &flg));
2657:     PetscOptionsEnd();
2658:   }
2659:   if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));

2661:   C->ops->mattransposemultsymbolic = MatMatTransposeMultSymbolic_MPIDense_MPIDense;
2662:   C->ops->productsymbolic          = MatProductSymbolic_ABt;
2663:   PetscFunctionReturn(PETSC_SUCCESS);
2664: }

2666: static PetscErrorCode MatProductSetFromOptions_MPIDense(Mat C)
2667: {
2668:   Mat_Product *product = C->product;

2670:   PetscFunctionBegin;
2671:   switch (product->type) {
2672:   case MATPRODUCT_AB:
2673:     PetscCall(MatProductSetFromOptions_MPIDense_AB(C));
2674:     break;
2675:   case MATPRODUCT_AtB:
2676:     PetscCall(MatProductSetFromOptions_MPIDense_AtB(C));
2677:     break;
2678:   case MATPRODUCT_ABt:
2679:     PetscCall(MatProductSetFromOptions_MPIDense_ABt(C));
2680:     break;
2681:   default:
2682:     break;
2683:   }
2684:   PetscFunctionReturn(PETSC_SUCCESS);
2685: }

2687: PetscErrorCode MatDenseScatter_Private(PetscSF sf, Mat X, Mat Y, InsertMode mode, ScatterMode smode)
2688: {
2689:   const PetscScalar *in;
2690:   PetscScalar       *out;
2691:   PetscSF            vsf;
2692:   PetscInt           N, ny, rld, lld;
2693:   PetscMemType       mtype[2];
2694:   MPI_Op             op = MPI_OP_NULL;

2696:   PetscFunctionBegin;
2700:   if (mode == INSERT_VALUES) op = MPI_REPLACE;
2701:   else if (mode == ADD_VALUES) op = MPIU_SUM;
2702:   else if (mode == MAX_VALUES) op = MPIU_MAX;
2703:   else if (mode == MIN_VALUES) op = MPIU_MIN;
2704:   PetscCheck(op != MPI_OP_NULL, PetscObjectComm((PetscObject)sf), PETSC_ERR_SUP, "Unsupported InsertMode %d in MatDenseScatter_Private()", mode);
2705:   PetscCheck(smode == SCATTER_FORWARD || smode == SCATTER_REVERSE, PetscObjectComm((PetscObject)sf), PETSC_ERR_SUP, "Unsupported ScatterMode %d in MatDenseScatter_Private()", smode);
2706:   PetscCall(MatGetSize(X, NULL, &N));
2707:   PetscCall(MatGetSize(Y, NULL, &ny));
2708:   PetscCheck(N == ny, PetscObjectComm((PetscObject)sf), PETSC_ERR_ARG_SIZ, "Matrix column sizes must match: %" PetscInt_FMT " != %" PetscInt_FMT, N, ny);
2709:   PetscCall(MatDenseGetLDA(X, &rld));
2710:   PetscCall(MatDenseGetLDA(Y, &lld));
2711:   /* get cached or create new strided PetscSF when the number of columns is greater than one */
2712:   if (N > 1) {
2713:     PetscCall(PetscObjectQuery((PetscObject)sf, "_MatDenseScatter_StridedSF", (PetscObject *)&vsf));
2714:     if (vsf) {
2715:       PetscInt nr[2], nl[2];

2717:       PetscCall(PetscSFGetGraph(sf, nr, nl, NULL, NULL));
2718:       PetscCall(PetscSFGetGraph(vsf, nr + 1, nl + 1, NULL, NULL));
2719:       if (N * nr[0] != nr[1] || N * nl[0] != nl[1]) vsf = NULL;
2720:     }
2721:     if (!vsf) {
2722:       PetscCall(PetscSFCreateStridedSF(sf, N, rld, lld, &vsf));
2723:       PetscCall(PetscObjectCompose((PetscObject)sf, "_MatDenseScatter_StridedSF", (PetscObject)vsf));
2724:       PetscCall(PetscObjectDereference((PetscObject)vsf));
2725:     }
2726:   } else vsf = sf;
2727:   /* the output array is accessed in read and write mode,
2728:     but write-only in the INSERT_VALUES case could be worth exploring */
2729:   PetscCall(MatDenseGetArrayReadAndMemType(X, &in, &mtype[0]));
2730:   PetscCall(MatDenseGetArrayAndMemType(Y, &out, &mtype[1]));
2731:   if (smode == SCATTER_FORWARD) {
2732:     PetscCall(PetscSFBcastWithMemTypeBegin(vsf, vsf->vscat.unit, mtype[0], in, mtype[1], out, op));
2733:     PetscCall(PetscSFBcastEnd(vsf, vsf->vscat.unit, in, out, op));
2734:   } else {
2735:     PetscCall(PetscSFReduceWithMemTypeBegin(vsf, vsf->vscat.unit, mtype[0], in, mtype[1], out, op));
2736:     PetscCall(PetscSFReduceEnd(vsf, vsf->vscat.unit, in, out, op));
2737:   }
2738:   PetscCall(MatDenseRestoreArrayAndMemType(Y, &out));
2739:   PetscCall(MatDenseRestoreArrayReadAndMemType(X, &in));
2740:   PetscFunctionReturn(PETSC_SUCCESS);
2741: }