Actual source code: matnest.c

  1: #include <../src/mat/impls/nest/matnestimpl.h>
  2: #include <../src/mat/impls/aij/seq/aij.h>
  3: #include <../src/mat/impls/shell/shell.h>
  4: #include <petscsf.h>

  6: static PetscErrorCode MatSetUp_NestIS_Private(Mat, PetscInt, const IS[], PetscInt, const IS[]);
  7: static PetscErrorCode MatCreateVecs_Nest(Mat, Vec *, Vec *);
  8: static PetscErrorCode MatReset_Nest(Mat);

 10: PETSC_INTERN PetscErrorCode MatConvert_Nest_IS(Mat, MatType, MatReuse, Mat *);

 12: /* private functions */
 13: static PetscErrorCode MatNestGetSizes_Private(Mat A, PetscInt *m, PetscInt *n, PetscInt *M, PetscInt *N)
 14: {
 15:   Mat_Nest *bA = (Mat_Nest *)A->data;

 17:   PetscFunctionBegin;
 18:   *m = *n = *M = *N = 0;
 19:   for (PetscInt i = 0; i < bA->nr; i++) { /* rows */
 20:     PetscInt sm, sM;

 22:     PetscCall(ISGetLocalSize(bA->isglobal.row[i], &sm));
 23:     PetscCall(ISGetSize(bA->isglobal.row[i], &sM));
 24:     *m += sm;
 25:     *M += sM;
 26:   }
 27:   for (PetscInt j = 0; j < bA->nc; j++) { /* cols */
 28:     PetscInt sn, sN;

 30:     PetscCall(ISGetLocalSize(bA->isglobal.col[j], &sn));
 31:     PetscCall(ISGetSize(bA->isglobal.col[j], &sN));
 32:     *n += sn;
 33:     *N += sN;
 34:   }
 35:   PetscFunctionReturn(PETSC_SUCCESS);
 36: }

 38: /* operations */
 39: static PetscErrorCode MatMult_Nest(Mat A, Vec x, Vec y)
 40: {
 41:   Mat_Nest *bA = (Mat_Nest *)A->data;
 42:   Vec      *bx = bA->right, *by = bA->left;

 44:   PetscFunctionBegin;
 45:   for (PetscInt i = 0; i < bA->nr; i++) PetscCall(VecGetSubVector(y, bA->isglobal.row[i], &by[i]));
 46:   for (PetscInt i = 0; i < bA->nc; i++) PetscCall(VecGetSubVector(x, bA->isglobal.col[i], &bx[i]));
 47:   for (PetscInt i = 0; i < bA->nr; i++) {
 48:     PetscCall(VecZeroEntries(by[i]));
 49:     for (PetscInt j = 0; j < bA->nc; j++) {
 50:       if (!bA->m[i][j]) continue;
 51:       /* y[i] <- y[i] + A[i][j] * x[j] */
 52:       PetscCall(MatMultAdd(bA->m[i][j], bx[j], by[i], by[i]));
 53:     }
 54:   }
 55:   for (PetscInt i = 0; i < bA->nr; i++) PetscCall(VecRestoreSubVector(y, bA->isglobal.row[i], &by[i]));
 56:   for (PetscInt i = 0; i < bA->nc; i++) PetscCall(VecRestoreSubVector(x, bA->isglobal.col[i], &bx[i]));
 57:   PetscFunctionReturn(PETSC_SUCCESS);
 58: }

 60: static PetscErrorCode MatMultAdd_Nest(Mat A, Vec x, Vec y, Vec z)
 61: {
 62:   Mat_Nest *bA = (Mat_Nest *)A->data;
 63:   Vec      *bx = bA->right, *bz = bA->left;

 65:   PetscFunctionBegin;
 66:   for (PetscInt i = 0; i < bA->nr; i++) PetscCall(VecGetSubVector(z, bA->isglobal.row[i], &bz[i]));
 67:   for (PetscInt i = 0; i < bA->nc; i++) PetscCall(VecGetSubVector(x, bA->isglobal.col[i], &bx[i]));
 68:   for (PetscInt i = 0; i < bA->nr; i++) {
 69:     if (y != z) {
 70:       Vec by;
 71:       PetscCall(VecGetSubVector(y, bA->isglobal.row[i], &by));
 72:       PetscCall(VecCopy(by, bz[i]));
 73:       PetscCall(VecRestoreSubVector(y, bA->isglobal.row[i], &by));
 74:     }
 75:     for (PetscInt j = 0; j < bA->nc; j++) {
 76:       if (!bA->m[i][j]) continue;
 77:       /* y[i] <- y[i] + A[i][j] * x[j] */
 78:       PetscCall(MatMultAdd(bA->m[i][j], bx[j], bz[i], bz[i]));
 79:     }
 80:   }
 81:   for (PetscInt i = 0; i < bA->nr; i++) PetscCall(VecRestoreSubVector(z, bA->isglobal.row[i], &bz[i]));
 82:   for (PetscInt i = 0; i < bA->nc; i++) PetscCall(VecRestoreSubVector(x, bA->isglobal.col[i], &bx[i]));
 83:   PetscFunctionReturn(PETSC_SUCCESS);
 84: }

 86: typedef struct {
 87:   Mat         *workC;      /* array of Mat with specific containers depending on the underlying MatMatMult implementation */
 88:   PetscScalar *tarray;     /* buffer for storing all temporary products A[i][j] B[j] */
 89:   PetscInt    *dm, *dn, k; /* displacements and number of submatrices */
 90: } Nest_Dense;

 92: static PetscErrorCode MatProductNumeric_Nest_Dense(Mat C)
 93: {
 94:   Mat_Nest          *bA;
 95:   Nest_Dense        *contents;
 96:   Mat                viewB, viewC, productB, workC;
 97:   const PetscScalar *barray;
 98:   PetscScalar       *carray;
 99:   PetscInt           M, N, nr, nc, ldb, ldc;
100:   Mat                A, B;

102:   PetscFunctionBegin;
103:   MatCheckProduct(C, 1);
104:   A = C->product->A;
105:   B = C->product->B;
106:   PetscCall(MatGetSize(B, NULL, &N));
107:   if (!N) {
108:     PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
109:     PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
110:     PetscFunctionReturn(PETSC_SUCCESS);
111:   }
112:   contents = (Nest_Dense *)C->product->data;
113:   PetscCheck(contents, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
114:   bA = (Mat_Nest *)A->data;
115:   nr = bA->nr;
116:   nc = bA->nc;
117:   PetscCall(MatDenseGetLDA(B, &ldb));
118:   PetscCall(MatDenseGetLDA(C, &ldc));
119:   PetscCall(MatZeroEntries(C));
120:   PetscCall(MatDenseGetArrayRead(B, &barray));
121:   PetscCall(MatDenseGetArray(C, &carray));
122:   for (PetscInt i = 0; i < nr; i++) {
123:     PetscCall(ISGetSize(bA->isglobal.row[i], &M));
124:     PetscCall(MatCreateDense(PetscObjectComm((PetscObject)A), contents->dm[i + 1] - contents->dm[i], PETSC_DECIDE, M, N, PetscSafePointerPlusOffset(carray, contents->dm[i]), &viewC));
125:     PetscCall(MatDenseSetLDA(viewC, ldc));
126:     for (PetscInt j = 0; j < nc; j++) {
127:       if (!bA->m[i][j]) continue;
128:       PetscCall(ISGetSize(bA->isglobal.col[j], &M));
129:       PetscCall(MatCreateDense(PetscObjectComm((PetscObject)A), contents->dn[j + 1] - contents->dn[j], PETSC_DECIDE, M, N, PetscSafePointerPlusOffset((PetscScalar *)barray, contents->dn[j]), &viewB));
130:       PetscCall(MatDenseSetLDA(viewB, ldb));

132:       /* MatMatMultNumeric(bA->m[i][j],viewB,contents->workC[i*nc + j]); */
133:       workC             = contents->workC[i * nc + j];
134:       productB          = workC->product->B;
135:       workC->product->B = viewB; /* use newly created dense matrix viewB */
136:       PetscCall(MatProductNumeric(workC));
137:       PetscCall(MatDestroy(&viewB));
138:       workC->product->B = productB; /* resume original B */

140:       /* C[i] <- workC + C[i] */
141:       PetscCall(MatAXPY(viewC, 1.0, contents->workC[i * nc + j], SAME_NONZERO_PATTERN));
142:     }
143:     PetscCall(MatDestroy(&viewC));
144:   }
145:   PetscCall(MatDenseRestoreArray(C, &carray));
146:   PetscCall(MatDenseRestoreArrayRead(B, &barray));

148:   PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
149:   PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
150:   PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
151:   PetscFunctionReturn(PETSC_SUCCESS);
152: }

154: static PetscErrorCode MatNest_DenseDestroy(PetscCtxRt ctx)
155: {
156:   Nest_Dense *contents = *(Nest_Dense **)ctx;

158:   PetscFunctionBegin;
159:   PetscCall(PetscFree(contents->tarray));
160:   for (PetscInt i = 0; i < contents->k; i++) PetscCall(MatDestroy(contents->workC + i));
161:   PetscCall(PetscFree3(contents->dm, contents->dn, contents->workC));
162:   PetscCall(PetscFree(contents));
163:   PetscFunctionReturn(PETSC_SUCCESS);
164: }

166: static PetscErrorCode MatProductSymbolic_Nest_Dense(Mat C)
167: {
168:   Mat_Nest          *bA;
169:   Mat                viewB, workC;
170:   const PetscScalar *barray;
171:   PetscInt           M, N, m, n, nr, nc, maxm = 0, ldb;
172:   Nest_Dense        *contents = NULL;
173:   PetscBool          cisdense;
174:   Mat                A, B;
175:   PetscReal          fill;

177:   PetscFunctionBegin;
178:   MatCheckProduct(C, 1);
179:   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
180:   A    = C->product->A;
181:   B    = C->product->B;
182:   fill = C->product->fill;
183:   bA   = (Mat_Nest *)A->data;
184:   nr   = bA->nr;
185:   nc   = bA->nc;
186:   PetscCall(MatGetLocalSize(C, &m, &n));
187:   PetscCall(MatGetSize(C, &M, &N));
188:   if (m == PETSC_DECIDE || n == PETSC_DECIDE || M == PETSC_DECIDE || N == PETSC_DECIDE) {
189:     PetscCall(MatGetLocalSize(B, NULL, &n));
190:     PetscCall(MatGetSize(B, NULL, &N));
191:     PetscCall(MatGetLocalSize(A, &m, NULL));
192:     PetscCall(MatGetSize(A, &M, NULL));
193:     PetscCall(MatSetSizes(C, m, n, M, N));
194:   }
195:   PetscCall(PetscObjectTypeCompareAny((PetscObject)C, &cisdense, MATSEQDENSE, MATMPIDENSE, MATSEQDENSECUDA, MATMPIDENSECUDA, ""));
196:   if (!cisdense) PetscCall(MatSetType(C, ((PetscObject)B)->type_name));
197:   PetscCall(MatSetUp(C));
198:   if (!N) {
199:     C->ops->productnumeric = MatProductNumeric_Nest_Dense;
200:     PetscFunctionReturn(PETSC_SUCCESS);
201:   }

203:   PetscCall(PetscNew(&contents));
204:   C->product->data    = contents;
205:   C->product->destroy = MatNest_DenseDestroy;
206:   PetscCall(PetscCalloc3(nr + 1, &contents->dm, nc + 1, &contents->dn, nr * nc, &contents->workC));
207:   contents->k = nr * nc;
208:   for (PetscInt i = 0; i < nr; i++) {
209:     PetscCall(ISGetLocalSize(bA->isglobal.row[i], contents->dm + i + 1));
210:     maxm = PetscMax(maxm, contents->dm[i + 1]);
211:     contents->dm[i + 1] += contents->dm[i];
212:   }
213:   for (PetscInt i = 0; i < nc; i++) {
214:     PetscCall(ISGetLocalSize(bA->isglobal.col[i], contents->dn + i + 1));
215:     contents->dn[i + 1] += contents->dn[i];
216:   }
217:   PetscCall(PetscMalloc1(maxm * N, &contents->tarray));
218:   PetscCall(MatDenseGetLDA(B, &ldb));
219:   PetscCall(MatGetSize(B, NULL, &N));
220:   PetscCall(MatDenseGetArrayRead(B, &barray));
221:   /* loops are permuted compared to MatMatMultNumeric so that viewB is created only once per column of A */
222:   for (PetscInt j = 0; j < nc; j++) {
223:     PetscCall(ISGetSize(bA->isglobal.col[j], &M));
224:     PetscCall(MatCreateDense(PetscObjectComm((PetscObject)A), contents->dn[j + 1] - contents->dn[j], PETSC_DECIDE, M, N, PetscSafePointerPlusOffset((PetscScalar *)barray, contents->dn[j]), &viewB));
225:     PetscCall(MatDenseSetLDA(viewB, ldb));
226:     for (PetscInt i = 0; i < nr; i++) {
227:       if (!bA->m[i][j]) continue;
228:       /* MatMatMultSymbolic may attach a specific container (depending on MatType of bA->m[i][j]) to workC[i][j] */

230:       PetscCall(MatProductCreate(bA->m[i][j], viewB, NULL, &contents->workC[i * nc + j]));
231:       workC = contents->workC[i * nc + j];
232:       PetscCall(MatProductSetType(workC, MATPRODUCT_AB));
233:       PetscCall(MatProductSetAlgorithm(workC, "default"));
234:       PetscCall(MatProductSetFill(workC, fill));
235:       PetscCall(MatProductSetFromOptions(workC));
236:       PetscCall(MatProductSymbolic(workC));

238:       /* since tarray will be shared by all Mat */
239:       PetscCall(MatSeqDenseSetPreallocation(workC, contents->tarray));
240:       PetscCall(MatMPIDenseSetPreallocation(workC, contents->tarray));
241:     }
242:     PetscCall(MatDestroy(&viewB));
243:   }
244:   PetscCall(MatDenseRestoreArrayRead(B, &barray));

246:   C->ops->productnumeric = MatProductNumeric_Nest_Dense;
247:   PetscFunctionReturn(PETSC_SUCCESS);
248: }

250: static PetscErrorCode MatProductSetFromOptions_Nest_Dense(Mat C)
251: {
252:   Mat_Product *product = C->product;

254:   PetscFunctionBegin;
255:   if (product->type == MATPRODUCT_AB) C->ops->productsymbolic = MatProductSymbolic_Nest_Dense;
256:   PetscFunctionReturn(PETSC_SUCCESS);
257: }

259: static PetscErrorCode MatMultTransposeKernel_Nest(Mat A, Vec x, Vec y, PetscBool herm)
260: {
261:   Mat_Nest *bA = (Mat_Nest *)A->data;
262:   Vec      *bx = bA->left, *by = bA->right;

264:   PetscFunctionBegin;
265:   for (PetscInt i = 0; i < bA->nr; i++) PetscCall(VecGetSubVector(x, bA->isglobal.row[i], &bx[i]));
266:   for (PetscInt i = 0; i < bA->nc; i++) PetscCall(VecGetSubVector(y, bA->isglobal.col[i], &by[i]));
267:   for (PetscInt j = 0; j < bA->nc; j++) {
268:     PetscCall(VecZeroEntries(by[j]));
269:     for (PetscInt i = 0; i < bA->nr; i++) {
270:       if (!bA->m[i][j]) continue;
271:       if (herm) PetscCall(MatMultHermitianTransposeAdd(bA->m[i][j], bx[i], by[j], by[j])); /* y[j] <- y[j] + (A[i][j])^H * x[i] */
272:       else PetscCall(MatMultTransposeAdd(bA->m[i][j], bx[i], by[j], by[j]));               /* y[j] <- y[j] + (A[i][j])^T * x[i] */
273:     }
274:   }
275:   for (PetscInt i = 0; i < bA->nr; i++) PetscCall(VecRestoreSubVector(x, bA->isglobal.row[i], &bx[i]));
276:   for (PetscInt i = 0; i < bA->nc; i++) PetscCall(VecRestoreSubVector(y, bA->isglobal.col[i], &by[i]));
277:   PetscFunctionReturn(PETSC_SUCCESS);
278: }

280: static PetscErrorCode MatMultTranspose_Nest(Mat A, Vec x, Vec y)
281: {
282:   PetscFunctionBegin;
283:   PetscCall(MatMultTransposeKernel_Nest(A, x, y, PETSC_FALSE));
284:   PetscFunctionReturn(PETSC_SUCCESS);
285: }

287: static PetscErrorCode MatMultHermitianTranspose_Nest(Mat A, Vec x, Vec y)
288: {
289:   PetscFunctionBegin;
290:   PetscCall(MatMultTransposeKernel_Nest(A, x, y, PETSC_TRUE));
291:   PetscFunctionReturn(PETSC_SUCCESS);
292: }

294: static PetscErrorCode MatMultTransposeAddKernel_Nest(Mat A, Vec x, Vec y, Vec z, PetscBool herm)
295: {
296:   Mat_Nest *bA = (Mat_Nest *)A->data;
297:   Vec      *bx = bA->left, *bz = bA->right;

299:   PetscFunctionBegin;
300:   for (PetscInt i = 0; i < bA->nr; i++) PetscCall(VecGetSubVector(x, bA->isglobal.row[i], &bx[i]));
301:   for (PetscInt i = 0; i < bA->nc; i++) PetscCall(VecGetSubVector(z, bA->isglobal.col[i], &bz[i]));
302:   for (PetscInt j = 0; j < bA->nc; j++) {
303:     if (y != z) {
304:       Vec by;

306:       PetscCall(VecGetSubVector(y, bA->isglobal.col[j], &by));
307:       PetscCall(VecCopy(by, bz[j]));
308:       PetscCall(VecRestoreSubVector(y, bA->isglobal.col[j], &by));
309:     }
310:     for (PetscInt i = 0; i < bA->nr; i++) {
311:       if (!bA->m[i][j]) continue;
312:       if (herm) PetscCall(MatMultHermitianTransposeAdd(bA->m[i][j], bx[i], bz[j], bz[j])); /* z[j] <- y[j] + (A[i][j])^H * x[i] */
313:       else PetscCall(MatMultTransposeAdd(bA->m[i][j], bx[i], bz[j], bz[j]));               /* z[j] <- y[j] + (A[i][j])^T * x[i] */
314:     }
315:   }
316:   for (PetscInt i = 0; i < bA->nr; i++) PetscCall(VecRestoreSubVector(x, bA->isglobal.row[i], &bx[i]));
317:   for (PetscInt i = 0; i < bA->nc; i++) PetscCall(VecRestoreSubVector(z, bA->isglobal.col[i], &bz[i]));
318:   PetscFunctionReturn(PETSC_SUCCESS);
319: }

321: static PetscErrorCode MatMultTransposeAdd_Nest(Mat A, Vec x, Vec y, Vec z)
322: {
323:   PetscFunctionBegin;
324:   PetscCall(MatMultTransposeAddKernel_Nest(A, x, y, z, PETSC_FALSE));
325:   PetscFunctionReturn(PETSC_SUCCESS);
326: }

328: static PetscErrorCode MatMultHermitianTransposeAdd_Nest(Mat A, Vec x, Vec y, Vec z)
329: {
330:   PetscFunctionBegin;
331:   PetscCall(MatMultTransposeAddKernel_Nest(A, x, y, z, PETSC_TRUE));
332:   PetscFunctionReturn(PETSC_SUCCESS);
333: }

335: static PetscErrorCode MatTranspose_Nest(Mat A, MatReuse reuse, Mat *B)
336: {
337:   Mat_Nest *bA = (Mat_Nest *)A->data, *bC;
338:   Mat       C;
339:   PetscInt  i, j, nr = bA->nr, nc = bA->nc;

341:   PetscFunctionBegin;
342:   if (reuse == MAT_REUSE_MATRIX) PetscCall(MatTransposeCheckNonzeroState_Private(A, *B));
343:   PetscCheck(reuse != MAT_INPLACE_MATRIX || nr == nc, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_SIZ, "Square nested matrix only for in-place");

345:   if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_INPLACE_MATRIX) {
346:     Mat *subs;
347:     IS  *is_row, *is_col;

349:     PetscCall(PetscCalloc1(nr * nc, &subs));
350:     PetscCall(PetscMalloc2(nr, &is_row, nc, &is_col));
351:     PetscCall(MatNestGetISs(A, is_row, is_col));
352:     if (reuse == MAT_INPLACE_MATRIX) {
353:       for (i = 0; i < nr; i++) {
354:         for (j = 0; j < nc; j++) subs[i + nr * j] = bA->m[i][j];
355:       }
356:     }

358:     PetscCall(MatCreateNest(PetscObjectComm((PetscObject)A), nc, is_col, nr, is_row, subs, &C));
359:     PetscCall(PetscFree(subs));
360:     PetscCall(PetscFree2(is_row, is_col));
361:   } else {
362:     C = *B;
363:   }

365:   bC = (Mat_Nest *)C->data;
366:   for (i = 0; i < nr; i++) {
367:     for (j = 0; j < nc; j++) {
368:       if (bA->m[i][j]) {
369:         PetscCall(MatTranspose(bA->m[i][j], reuse, &bC->m[j][i]));
370:       } else {
371:         bC->m[j][i] = NULL;
372:       }
373:     }
374:   }

376:   if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_REUSE_MATRIX) {
377:     *B = C;
378:   } else {
379:     PetscCall(MatHeaderMerge(A, &C));
380:   }
381:   PetscFunctionReturn(PETSC_SUCCESS);
382: }

384: static PetscErrorCode MatNestDestroyISList(PetscInt n, IS **list)
385: {
386:   IS *lst = *list;

388:   PetscFunctionBegin;
389:   if (!lst) PetscFunctionReturn(PETSC_SUCCESS);
390:   for (PetscInt i = 0; i < n; i++) PetscCall(ISDestroy(&lst[i]));
391:   PetscCall(PetscFree(lst));
392:   *list = NULL;
393:   PetscFunctionReturn(PETSC_SUCCESS);
394: }

396: static PetscErrorCode MatReset_Nest(Mat A)
397: {
398:   Mat_Nest *vs = (Mat_Nest *)A->data;

400:   PetscFunctionBegin;
401:   /* release the matrices and the place holders */
402:   PetscCall(MatNestDestroyISList(vs->nr, &vs->isglobal.row));
403:   PetscCall(MatNestDestroyISList(vs->nc, &vs->isglobal.col));
404:   PetscCall(MatNestDestroyISList(vs->nr, &vs->islocal.row));
405:   PetscCall(MatNestDestroyISList(vs->nc, &vs->islocal.col));

407:   PetscCall(PetscFree(vs->row_len));
408:   PetscCall(PetscFree(vs->col_len));
409:   PetscCall(PetscFree(vs->nnzstate));

411:   PetscCall(PetscFree2(vs->left, vs->right));

413:   /* release the matrices and the place holders */
414:   if (vs->m) {
415:     for (PetscInt i = 0; i < vs->nr; i++) {
416:       for (PetscInt j = 0; j < vs->nc; j++) PetscCall(MatDestroy(&vs->m[i][j]));
417:     }
418:     PetscCall(PetscFree(vs->m[0]));
419:     PetscCall(PetscFree(vs->m));
420:   }

422:   /* restore defaults */
423:   vs->nr            = 0;
424:   vs->nc            = 0;
425:   vs->splitassembly = PETSC_FALSE;
426:   PetscFunctionReturn(PETSC_SUCCESS);
427: }

429: static PetscErrorCode MatDestroy_Nest(Mat A)
430: {
431:   PetscFunctionBegin;
432:   PetscCall(MatReset_Nest(A));
433:   PetscCall(PetscFree(A->data));
434:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetSubMat_C", NULL));
435:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestSetSubMat_C", NULL));
436:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetSubMats_C", NULL));
437:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetSize_C", NULL));
438:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetISs_C", NULL));
439:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetLocalISs_C", NULL));
440:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestSetVecType_C", NULL));
441:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestSetSubMats_C", NULL));
442:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_mpiaij_C", NULL));
443:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_seqaij_C", NULL));
444:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_aij_C", NULL));
445:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_is_C", NULL));
446:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_mpidense_C", NULL));
447:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_seqdense_C", NULL));
448:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_nest_seqdense_C", NULL));
449:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_nest_mpidense_C", NULL));
450:   PetscFunctionReturn(PETSC_SUCCESS);
451: }

453: static PetscErrorCode MatAssemblyBegin_Nest(Mat A, MatAssemblyType type)
454: {
455:   Mat_Nest *vs       = (Mat_Nest *)A->data;
456:   PetscBool nnzstate = PETSC_FALSE;

458:   PetscFunctionBegin;
459:   for (PetscInt i = 0; i < vs->nr; i++) {
460:     for (PetscInt j = 0; j < vs->nc; j++) {
461:       PetscObjectState subnnzstate = 0;
462:       if (vs->m[i][j]) {
463:         PetscCall(MatAssemblyBegin(vs->m[i][j], type));
464:         if (!vs->splitassembly) {
465:           /* Note: split assembly will fail if the same block appears more than once (even indirectly through a nested
466:            * sub-block). This could be fixed by adding a flag to Mat so that there was a way to check if a Mat was
467:            * already performing an assembly, but the result would by more complicated and appears to offer less
468:            * potential for diagnostics and correctness checking. Split assembly should be fixed once there is an
469:            * interface for libraries to make asynchronous progress in "user-defined non-blocking collectives".
470:            */
471:           PetscCall(MatAssemblyEnd(vs->m[i][j], type));
472:           PetscCall(MatGetNonzeroState(vs->m[i][j], &subnnzstate));
473:         }
474:       }
475:       nnzstate                     = (PetscBool)(nnzstate || vs->nnzstate[i * vs->nc + j] != subnnzstate);
476:       vs->nnzstate[i * vs->nc + j] = subnnzstate;
477:     }
478:   }
479:   if (nnzstate) A->nonzerostate++;
480:   PetscFunctionReturn(PETSC_SUCCESS);
481: }

483: static PetscErrorCode MatAssemblyEnd_Nest(Mat A, MatAssemblyType type)
484: {
485:   Mat_Nest *vs = (Mat_Nest *)A->data;

487:   PetscFunctionBegin;
488:   for (PetscInt i = 0; i < vs->nr; i++) {
489:     for (PetscInt j = 0; j < vs->nc; j++) {
490:       if (vs->m[i][j]) {
491:         if (vs->splitassembly) PetscCall(MatAssemblyEnd(vs->m[i][j], type));
492:       }
493:     }
494:   }
495:   PetscFunctionReturn(PETSC_SUCCESS);
496: }

498: static PetscErrorCode MatNestFindNonzeroSubMatRow(Mat A, PetscInt row, Mat *B)
499: {
500:   Mat_Nest *vs = (Mat_Nest *)A->data;
501:   Mat       sub;

503:   PetscFunctionBegin;
504:   sub = (row < vs->nc) ? vs->m[row][row] : (Mat)NULL; /* Prefer to find on the diagonal */
505:   for (PetscInt j = 0; !sub && j < vs->nc; j++) sub = vs->m[row][j];
506:   if (sub) PetscCall(MatSetUp(sub)); /* Ensure that the sizes are available */
507:   *B = sub;
508:   PetscFunctionReturn(PETSC_SUCCESS);
509: }

511: static PetscErrorCode MatNestFindNonzeroSubMatCol(Mat A, PetscInt col, Mat *B)
512: {
513:   Mat_Nest *vs = (Mat_Nest *)A->data;
514:   Mat       sub;

516:   PetscFunctionBegin;
517:   sub = (col < vs->nr) ? vs->m[col][col] : (Mat)NULL; /* Prefer to find on the diagonal */
518:   for (PetscInt i = 0; !sub && i < vs->nr; i++) sub = vs->m[i][col];
519:   if (sub) PetscCall(MatSetUp(sub)); /* Ensure that the sizes are available */
520:   *B = sub;
521:   PetscFunctionReturn(PETSC_SUCCESS);
522: }

524: static PetscErrorCode MatNestFindISRange(Mat A, PetscInt n, const IS list[], IS is, PetscInt *begin, PetscInt *end)
525: {
526:   PetscInt  i, j, size, m;
527:   PetscBool flg;
528:   IS        out, concatenate[2];

530:   PetscFunctionBegin;
531:   PetscAssertPointer(list, 3);
533:   if (begin) {
534:     PetscAssertPointer(begin, 5);
535:     *begin = -1;
536:   }
537:   if (end) {
538:     PetscAssertPointer(end, 6);
539:     *end = -1;
540:   }
541:   for (i = 0; i < n; i++) {
542:     if (!list[i]) continue;
543:     PetscCall(ISEqualUnsorted(list[i], is, &flg));
544:     if (flg) {
545:       if (begin) *begin = i;
546:       if (end) *end = i + 1;
547:       PetscFunctionReturn(PETSC_SUCCESS);
548:     }
549:   }
550:   PetscCall(ISGetSize(is, &size));
551:   for (i = 0; i < n - 1; i++) {
552:     if (!list[i]) continue;
553:     m = 0;
554:     PetscCall(ISConcatenate(PetscObjectComm((PetscObject)A), 2, list + i, &out));
555:     PetscCall(ISGetSize(out, &m));
556:     for (j = i + 2; j < n && m < size; j++) {
557:       if (list[j]) {
558:         concatenate[0] = out;
559:         concatenate[1] = list[j];
560:         PetscCall(ISConcatenate(PetscObjectComm((PetscObject)A), 2, concatenate, &out));
561:         PetscCall(ISDestroy(concatenate));
562:         PetscCall(ISGetSize(out, &m));
563:       }
564:     }
565:     if (m == size) {
566:       PetscCall(ISEqualUnsorted(out, is, &flg));
567:       if (flg) {
568:         if (begin) *begin = i;
569:         if (end) *end = j;
570:         PetscCall(ISDestroy(&out));
571:         PetscFunctionReturn(PETSC_SUCCESS);
572:       }
573:     }
574:     PetscCall(ISDestroy(&out));
575:   }
576:   PetscFunctionReturn(PETSC_SUCCESS);
577: }

579: static PetscErrorCode MatNestFillEmptyMat_Private(Mat A, PetscInt i, PetscInt j, Mat *B)
580: {
581:   Mat_Nest *vs = (Mat_Nest *)A->data;
582:   PetscInt  lr, lc;

584:   PetscFunctionBegin;
585:   PetscCall(MatCreate(PetscObjectComm((PetscObject)A), B));
586:   PetscCall(ISGetLocalSize(vs->isglobal.row[i], &lr));
587:   PetscCall(ISGetLocalSize(vs->isglobal.col[j], &lc));
588:   PetscCall(MatSetSizes(*B, lr, lc, PETSC_DECIDE, PETSC_DECIDE));
589:   PetscCall(MatSetType(*B, MATAIJ));
590:   PetscCall(MatSeqAIJSetPreallocation(*B, 0, NULL));
591:   PetscCall(MatMPIAIJSetPreallocation(*B, 0, NULL, 0, NULL));
592:   PetscCall(MatSetUp(*B));
593:   PetscCall(MatSetOption(*B, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
594:   PetscCall(MatAssemblyBegin(*B, MAT_FINAL_ASSEMBLY));
595:   PetscCall(MatAssemblyEnd(*B, MAT_FINAL_ASSEMBLY));
596:   PetscFunctionReturn(PETSC_SUCCESS);
597: }

599: static PetscErrorCode MatNestGetBlock_Private(Mat A, PetscInt rbegin, PetscInt rend, PetscInt cbegin, PetscInt cend, Mat *B)
600: {
601:   Mat_Nest  *vs = (Mat_Nest *)A->data;
602:   Mat       *a;
603:   PetscInt   i, j, k, l, nr = rend - rbegin, nc = cend - cbegin;
604:   char       keyname[256];
605:   PetscBool *b;
606:   PetscBool  flg;

608:   PetscFunctionBegin;
609:   *B = NULL;
610:   PetscCall(PetscSNPrintf(keyname, sizeof(keyname), "NestBlock_%" PetscInt_FMT "-%" PetscInt_FMT "x%" PetscInt_FMT "-%" PetscInt_FMT, rbegin, rend, cbegin, cend));
611:   PetscCall(PetscObjectQuery((PetscObject)A, keyname, (PetscObject *)B));
612:   if (*B) PetscFunctionReturn(PETSC_SUCCESS);

614:   PetscCall(PetscMalloc2(nr * nc, &a, nr * nc, &b));
615:   for (i = 0; i < nr; i++) {
616:     for (j = 0; j < nc; j++) {
617:       a[i * nc + j] = vs->m[rbegin + i][cbegin + j];
618:       b[i * nc + j] = PETSC_FALSE;
619:     }
620:   }
621:   if (nc != vs->nc && nr != vs->nr) {
622:     for (i = 0; i < nr; i++) {
623:       for (j = 0; j < nc; j++) {
624:         flg = PETSC_FALSE;
625:         for (k = 0; (k < nr && !flg); k++) {
626:           if (a[j + k * nc]) flg = PETSC_TRUE;
627:         }
628:         if (flg) {
629:           flg = PETSC_FALSE;
630:           for (l = 0; (l < nc && !flg); l++) {
631:             if (a[i * nc + l]) flg = PETSC_TRUE;
632:           }
633:         }
634:         if (!flg) {
635:           b[i * nc + j] = PETSC_TRUE;
636:           PetscCall(MatNestFillEmptyMat_Private(A, rbegin + i, cbegin + j, a + i * nc + j));
637:         }
638:       }
639:     }
640:   }
641:   PetscCall(MatCreateNest(PetscObjectComm((PetscObject)A), nr, nr != vs->nr ? NULL : vs->isglobal.row, nc, nc != vs->nc ? NULL : vs->isglobal.col, a, B));
642:   for (i = 0; i < nr; i++) {
643:     for (j = 0; j < nc; j++) {
644:       if (b[i * nc + j]) PetscCall(MatDestroy(a + i * nc + j));
645:     }
646:   }
647:   PetscCall(PetscFree2(a, b));
648:   (*B)->assembled = A->assembled;
649:   PetscCall(PetscObjectCompose((PetscObject)A, keyname, (PetscObject)*B));
650:   PetscCall(PetscObjectDereference((PetscObject)*B)); /* Leave the only remaining reference in the composition */
651:   PetscFunctionReturn(PETSC_SUCCESS);
652: }

654: static PetscErrorCode MatNestFindSubMat(Mat A, IS isrow, IS iscol, PetscBool global, PetscBool *found, Mat *B)
655: {
656:   Mat_Nest *vs = (Mat_Nest *)A->data;
657:   PetscInt  rbegin, rend, cbegin, cend;

659:   PetscFunctionBegin;
660:   *B = NULL;
661:   PetscCall(MatNestFindISRange(A, vs->nr, global ? vs->isglobal.row : vs->islocal.row, isrow, &rbegin, &rend));
662:   PetscCall(MatNestFindISRange(A, vs->nc, global ? vs->isglobal.col : vs->islocal.col, iscol, &cbegin, &cend));
663:   if (rend == rbegin + 1 && cend == cbegin + 1) {
664:     if (!vs->m[rbegin][cbegin]) PetscCall(MatNestFillEmptyMat_Private(A, rbegin, cbegin, vs->m[rbegin] + cbegin));
665:     *B = vs->m[rbegin][cbegin];
666:     if (found) *found = PETSC_TRUE;
667:   } else if (rbegin != -1 && cbegin != -1) {
668:     PetscCheck(global == PETSC_TRUE, PETSC_COMM_SELF, PETSC_ERR_SUP, "MATNEST local submatrix cannot select more than a single submatrix");
669:     PetscCall(MatNestGetBlock_Private(A, rbegin, rend, cbegin, cend, B));
670:     if (found) *found = PETSC_TRUE;
671:   } else if (found) *found = PETSC_FALSE;
672:   PetscFunctionReturn(PETSC_SUCCESS);
673: }

675: static PetscErrorCode MatNestFindFullBlocks_Private(Mat A, PetscInt n, const IS blockis[], IS is, const char axis[], PetscInt *nselected, PetscInt **selected, IS **isout)
676: {
677:   const PetscInt *idx;
678:   PetscInt       *blocks;
679:   IS             *out;
680:   PetscInt        N, bs, cursor = 0, i, nblock, nlocal, nout = 0, offset = 0, start;
681:   PetscBool       complete, match;

683:   PetscFunctionBegin;
684:   PetscCall(ISGetSize(is, &N));
685:   PetscCheck(N, PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "Empty %s index sets are not supported for MATNEST submatrices", axis);
686:   PetscCall(ISGetLocalSize(is, &nlocal));
687:   PetscCall(ISGetIndices(is, &idx));
688:   PetscCall(PetscMalloc1(n, &blocks));
689:   for (i = 0; i < n; i++) {
690:     const PetscInt *bidx;

692:     PetscCall(ISGetSize(blockis[i], &N));
693:     if (!N) continue;
694:     PetscCall(ISGetLocalSize(blockis[i], &nblock));
695:     match = (PetscBool)(cursor + nblock <= nlocal);
696:     if (match && nblock) {
697:       PetscCall(ISGetIndices(blockis[i], &bidx));
698:       PetscCall(PetscArraycmp(idx + cursor, bidx, nblock, &match));
699:       PetscCall(ISRestoreIndices(blockis[i], &bidx));
700:     }
701:     PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &match, 1, MPI_C_BOOL, MPI_LAND, PetscObjectComm((PetscObject)A)));
702:     if (match) {
703:       blocks[nout++] = i;
704:       cursor += nblock;
705:     }
706:   }
707:   complete = (PetscBool)(cursor == nlocal);
708:   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &complete, 1, MPI_C_BOOL, MPI_LAND, PetscObjectComm((PetscObject)A)));
709:   PetscCall(ISRestoreIndices(is, &idx));
710:   if (!complete || !nout) PetscCall(PetscFree(blocks));
711:   PetscCheck(complete && nout, PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "MATNEST submatrix %s index set must be an ordered union of complete MATNEST blocks", axis);

713:   PetscCallMPI(MPI_Scan(&nlocal, &start, 1, MPIU_INT, MPI_SUM, PetscObjectComm((PetscObject)A)));
714:   start -= nlocal;
715:   PetscCall(PetscMalloc1(nout, &out));
716:   for (i = 0; i < nout; i++) {
717:     PetscCall(ISGetLocalSize(blockis[blocks[i]], &nblock));
718:     PetscCall(ISGetBlockSize(blockis[blocks[i]], &bs));
719:     PetscCall(ISCreateStride(PetscObjectComm((PetscObject)A), nblock, start + offset, 1, out + i));
720:     PetscCall(ISSetBlockSize(out[i], bs));
721:     offset += nblock;
722:   }
723:   PetscCheck(offset == nlocal, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Inconsistent MATNEST submatrix %s layout", axis);
724:   *nselected = nout;
725:   *selected  = blocks;
726:   *isout     = out;
727:   PetscFunctionReturn(PETSC_SUCCESS);
728: }

730: static PetscErrorCode MatCreateSubMatrix_Nest_Nontrivial(Mat A, IS isrow, IS iscol, MatReuse reuse, Mat *B)
731: {
732:   Mat_Nest *vs = (Mat_Nest *)A->data;
733:   Mat      *submats;
734:   IS       *rowis, *colis;
735:   PetscInt *rows, *cols;
736:   PetscInt  nr, nc;
737:   PetscBool flg;

739:   PetscFunctionBegin;
740:   PetscCall(MatNestFindFullBlocks_Private(A, vs->nr, vs->isglobal.row, isrow, "row", &nr, &rows, &rowis));
741:   PetscCall(MatNestFindFullBlocks_Private(A, vs->nc, vs->isglobal.col, iscol, "column", &nc, &cols, &colis));
742:   PetscCall(PetscMalloc1(nr * nc, &submats));
743:   for (PetscInt i = 0; i < nr; i++) {
744:     for (PetscInt j = 0; j < nc; j++) submats[i * nc + j] = vs->m[rows[i]][cols[j]];
745:   }
746:   if (reuse == MAT_INITIAL_MATRIX) {
747:     PetscCall(MatCreateNest(PetscObjectComm((PetscObject)A), nr, rowis, nc, colis, submats, B));
748:     (*B)->assembled = A->assembled;
749:   } else {
750:     PetscCheck(reuse == MAT_REUSE_MATRIX, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Invalid MatReuse %d", (int)reuse);
751:     PetscCall(PetscObjectTypeCompare((PetscObject)*B, MATNEST, &flg));
752:     PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Cannot reuse a non-MATNEST matrix for this MATNEST submatrix");
753:     vs = (Mat_Nest *)(*B)->data;
754:     PetscCheck(vs->nr == nr && vs->nc == nc, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Cannot reuse MATNEST submatrix with a different block layout");
755:     for (PetscInt i = 0; i < nr; i++) {
756:       PetscCall(ISEqualUnsorted(vs->isglobal.row[i], rowis[i], &flg));
757:       PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Cannot reuse MATNEST submatrix with a different row layout");
758:     }
759:     for (PetscInt j = 0; j < nc; j++) {
760:       PetscCall(ISEqualUnsorted(vs->isglobal.col[j], colis[j], &flg));
761:       PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Cannot reuse MATNEST submatrix with a different column layout");
762:     }
763:     PetscCall(MatNestSetSubMats(*B, nr, rowis, nc, colis, submats));
764:     (*B)->assembled = A->assembled;
765:   }
766:   PetscCall(PetscFree(submats));
767:   for (PetscInt i = 0; i < nr; i++) PetscCall(ISDestroy(rowis + i));
768:   for (PetscInt j = 0; j < nc; j++) PetscCall(ISDestroy(colis + j));
769:   PetscCall(PetscFree(rows));
770:   PetscCall(PetscFree(rowis));
771:   PetscCall(PetscFree(cols));
772:   PetscCall(PetscFree(colis));
773:   PetscFunctionReturn(PETSC_SUCCESS);
774: }

776: /*
777:    TODO: This does not actually returns a submatrix we can modify
778: */
779: static PetscErrorCode MatCreateSubMatrix_Nest(Mat A, IS isrow, IS iscol, MatReuse reuse, Mat *B)
780: {
781:   Mat       sub;
782:   PetscBool found;

784:   PetscFunctionBegin;
785:   PetscCall(MatNestFindSubMat(A, isrow, iscol, PETSC_TRUE, &found, &sub));
786:   if (!found) {
787:     PetscCall(MatCreateSubMatrix_Nest_Nontrivial(A, isrow, iscol, reuse, B));
788:     PetscFunctionReturn(PETSC_SUCCESS);
789:   }
790:   switch (reuse) {
791:   case MAT_INITIAL_MATRIX:
792:     PetscCall(PetscObjectReference((PetscObject)sub));
793:     if (sub) PetscCall(PetscObjectStateIncrease((PetscObject)sub));
794:     *B = sub;
795:     break;
796:   case MAT_REUSE_MATRIX:
797:     PetscCheck(sub == *B, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Submatrix was not used before in this call");
798:     if (sub) PetscCall(PetscObjectStateIncrease((PetscObject)sub));
799:     break;
800:   default:
801:     break;
802:   }
803:   PetscFunctionReturn(PETSC_SUCCESS);
804: }

806: static PetscErrorCode MatGetLocalSubMatrix_Nest(Mat A, IS isrow, IS iscol, Mat *B)
807: {
808:   Mat sub;

810:   PetscFunctionBegin;
811:   PetscCall(MatNestFindSubMat(A, isrow, iscol, PETSC_FALSE, NULL, &sub));
812:   /* We allow the submatrix to be NULL, perhaps it would be better for the user to return an empty matrix instead */
813:   PetscCall(PetscObjectReference((PetscObject)sub));
814:   *B = sub;
815:   PetscFunctionReturn(PETSC_SUCCESS);
816: }

818: static PetscErrorCode MatRestoreLocalSubMatrix_Nest(Mat A, IS isrow, IS iscol, Mat *B)
819: {
820:   Mat sub;

822:   PetscFunctionBegin;
823:   PetscCall(MatNestFindSubMat(A, isrow, iscol, PETSC_FALSE, NULL, &sub));
824:   PetscCheck(*B == sub, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Local submatrix has not been gotten");
825:   if (sub) {
826:     PetscCheck(((PetscObject)sub)->refct > 1, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Local submatrix has had reference count decremented too many times");
827:     PetscCall(MatDestroy(B));
828:   }
829:   PetscFunctionReturn(PETSC_SUCCESS);
830: }

832: static PetscErrorCode MatGetDiagonal_Nest(Mat A, Vec v)
833: {
834:   Mat_Nest *bA = (Mat_Nest *)A->data;

836:   PetscFunctionBegin;
837:   for (PetscInt i = 0; i < bA->nr; i++) {
838:     Vec bv;
839:     PetscCall(VecGetSubVector(v, bA->isglobal.row[i], &bv));
840:     if (bA->m[i][i]) PetscCall(MatGetDiagonal(bA->m[i][i], bv));
841:     else PetscCall(VecSet(bv, 0.0));
842:     PetscCall(VecRestoreSubVector(v, bA->isglobal.row[i], &bv));
843:   }
844:   PetscFunctionReturn(PETSC_SUCCESS);
845: }

847: static PetscErrorCode MatDiagonalScale_Nest(Mat A, Vec l, Vec r)
848: {
849:   Mat_Nest *bA = (Mat_Nest *)A->data;
850:   Vec       bl, *br;

852:   PetscFunctionBegin;
853:   PetscCall(PetscCalloc1(bA->nc, &br));
854:   if (r) {
855:     for (PetscInt j = 0; j < bA->nc; j++) PetscCall(VecGetSubVector(r, bA->isglobal.col[j], &br[j]));
856:   }
857:   bl = NULL;
858:   for (PetscInt i = 0; i < bA->nr; i++) {
859:     if (l) PetscCall(VecGetSubVector(l, bA->isglobal.row[i], &bl));
860:     for (PetscInt j = 0; j < bA->nc; j++) {
861:       if (bA->m[i][j]) PetscCall(MatDiagonalScale(bA->m[i][j], bl, br[j]));
862:     }
863:     if (l) PetscCall(VecRestoreSubVector(l, bA->isglobal.row[i], &bl));
864:   }
865:   if (r) {
866:     for (PetscInt j = 0; j < bA->nc; j++) PetscCall(VecRestoreSubVector(r, bA->isglobal.col[j], &br[j]));
867:   }
868:   PetscCall(PetscFree(br));
869:   PetscFunctionReturn(PETSC_SUCCESS);
870: }

872: static PetscErrorCode MatScale_Nest(Mat A, PetscScalar a)
873: {
874:   Mat_Nest *bA = (Mat_Nest *)A->data;

876:   PetscFunctionBegin;
877:   for (PetscInt i = 0; i < bA->nr; i++) {
878:     for (PetscInt j = 0; j < bA->nc; j++) {
879:       if (bA->m[i][j]) PetscCall(MatScale(bA->m[i][j], a));
880:     }
881:   }
882:   PetscFunctionReturn(PETSC_SUCCESS);
883: }

885: static PetscErrorCode MatShift_Nest(Mat A, PetscScalar a)
886: {
887:   Mat_Nest *bA       = (Mat_Nest *)A->data;
888:   PetscBool nnzstate = PETSC_FALSE;

890:   PetscFunctionBegin;
891:   for (PetscInt i = 0; i < bA->nr; i++) {
892:     PetscObjectState subnnzstate = 0;
893:     PetscCheck(bA->m[i][i], PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "No support for shifting an empty diagonal block, insert a matrix in block (%" PetscInt_FMT ",%" PetscInt_FMT ")", i, i);
894:     PetscCall(MatShift(bA->m[i][i], a));
895:     PetscCall(MatGetNonzeroState(bA->m[i][i], &subnnzstate));
896:     nnzstate                     = (PetscBool)(nnzstate || bA->nnzstate[i * bA->nc + i] != subnnzstate);
897:     bA->nnzstate[i * bA->nc + i] = subnnzstate;
898:   }
899:   if (nnzstate) A->nonzerostate++;
900:   PetscFunctionReturn(PETSC_SUCCESS);
901: }

903: static PetscErrorCode MatDiagonalSet_Nest(Mat A, Vec D, InsertMode is)
904: {
905:   Mat_Nest *bA       = (Mat_Nest *)A->data;
906:   PetscBool nnzstate = PETSC_FALSE;

908:   PetscFunctionBegin;
909:   for (PetscInt i = 0; i < bA->nr; i++) {
910:     PetscObjectState subnnzstate = 0;
911:     Vec              bv;
912:     PetscCall(VecGetSubVector(D, bA->isglobal.row[i], &bv));
913:     if (bA->m[i][i]) {
914:       PetscCall(MatDiagonalSet(bA->m[i][i], bv, is));
915:       PetscCall(MatGetNonzeroState(bA->m[i][i], &subnnzstate));
916:     }
917:     PetscCall(VecRestoreSubVector(D, bA->isglobal.row[i], &bv));
918:     nnzstate                     = (PetscBool)(nnzstate || bA->nnzstate[i * bA->nc + i] != subnnzstate);
919:     bA->nnzstate[i * bA->nc + i] = subnnzstate;
920:   }
921:   if (nnzstate) A->nonzerostate++;
922:   PetscFunctionReturn(PETSC_SUCCESS);
923: }

925: static PetscErrorCode MatSetRandom_Nest(Mat A, PetscRandom rctx)
926: {
927:   Mat_Nest *bA = (Mat_Nest *)A->data;

929:   PetscFunctionBegin;
930:   for (PetscInt i = 0; i < bA->nr; i++) {
931:     for (PetscInt j = 0; j < bA->nc; j++) {
932:       if (bA->m[i][j]) PetscCall(MatSetRandom(bA->m[i][j], rctx));
933:     }
934:   }
935:   PetscFunctionReturn(PETSC_SUCCESS);
936: }

938: static PetscErrorCode MatCreateVecs_Nest(Mat A, Vec *right, Vec *left)
939: {
940:   Mat_Nest *bA = (Mat_Nest *)A->data;
941:   Vec      *L, *R;
942:   MPI_Comm  comm;
943:   PetscInt  i, j;

945:   PetscFunctionBegin;
946:   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
947:   if (right) {
948:     /* allocate R */
949:     PetscCall(PetscMalloc1(bA->nc, &R));
950:     /* Create the right vectors */
951:     for (j = 0; j < bA->nc; j++) {
952:       for (i = 0; i < bA->nr; i++) {
953:         if (bA->m[i][j]) {
954:           PetscCall(MatCreateVecs(bA->m[i][j], &R[j], NULL));
955:           break;
956:         }
957:       }
958:       PetscCheck(i != bA->nr, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Mat(Nest) contains a null column.");
959:     }
960:     PetscCall(VecCreateNest(comm, bA->nc, bA->isglobal.col, R, right));
961:     /* hand back control to the nest vector */
962:     for (j = 0; j < bA->nc; j++) PetscCall(VecDestroy(&R[j]));
963:     PetscCall(PetscFree(R));
964:   }

966:   if (left) {
967:     /* allocate L */
968:     PetscCall(PetscMalloc1(bA->nr, &L));
969:     /* Create the left vectors */
970:     for (i = 0; i < bA->nr; i++) {
971:       for (j = 0; j < bA->nc; j++) {
972:         if (bA->m[i][j]) {
973:           PetscCall(MatCreateVecs(bA->m[i][j], NULL, &L[i]));
974:           break;
975:         }
976:       }
977:       PetscCheck(j != bA->nc, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Mat(Nest) contains a null row.");
978:     }

980:     PetscCall(VecCreateNest(comm, bA->nr, bA->isglobal.row, L, left));
981:     for (i = 0; i < bA->nr; i++) PetscCall(VecDestroy(&L[i]));

983:     PetscCall(PetscFree(L));
984:   }
985:   PetscFunctionReturn(PETSC_SUCCESS);
986: }

988: static PetscErrorCode MatView_Nest(Mat A, PetscViewer viewer)
989: {
990:   Mat_Nest *bA = (Mat_Nest *)A->data;
991:   PetscBool isascii, viewSub = PETSC_FALSE;

993:   PetscFunctionBegin;
994:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
995:   if (isascii) {
996:     PetscViewerFormat format;

998:     PetscCall(PetscViewerGetFormat(viewer, &format));
999:     if (format == PETSC_VIEWER_ASCII_MATLAB) {
1000:       Mat T;

1002:       PetscCall(MatConvert(A, MATAIJ, MAT_INITIAL_MATRIX, &T));
1003:       PetscCall(MatView(T, viewer));
1004:       PetscCall(MatDestroy(&T));
1005:       PetscFunctionReturn(PETSC_SUCCESS);
1006:     }
1007:     PetscCall(PetscOptionsGetBool(((PetscObject)A)->options, ((PetscObject)A)->prefix, "-mat_view_nest_sub", &viewSub, NULL));
1008:     PetscCall(PetscViewerASCIIPushTab(viewer));
1009:     PetscCall(PetscViewerASCIIPrintf(viewer, "MatNest, rows=%" PetscInt_FMT ", cols=%" PetscInt_FMT ", structure:\n", bA->nr, bA->nc));
1010:     for (PetscInt i = 0; i < bA->nr; i++) {
1011:       for (PetscInt j = 0; j < bA->nc; j++) {
1012:         MatType   type;
1013:         char      name[256] = "", prefix[256] = "";
1014:         PetscInt  NR, NC;
1015:         PetscBool isNest = PETSC_FALSE;

1017:         if (!bA->m[i][j]) {
1018:           PetscCall(PetscViewerASCIIPrintf(viewer, "(%" PetscInt_FMT ",%" PetscInt_FMT ") : NULL\n", i, j));
1019:           continue;
1020:         }
1021:         PetscCall(MatGetSize(bA->m[i][j], &NR, &NC));
1022:         PetscCall(MatGetType(bA->m[i][j], &type));
1023:         if (((PetscObject)bA->m[i][j])->name) PetscCall(PetscSNPrintf(name, sizeof(name), "name=\"%s\", ", ((PetscObject)bA->m[i][j])->name));
1024:         if (((PetscObject)bA->m[i][j])->prefix) PetscCall(PetscSNPrintf(prefix, sizeof(prefix), "prefix=\"%s\", ", ((PetscObject)bA->m[i][j])->prefix));
1025:         PetscCall(PetscObjectTypeCompare((PetscObject)bA->m[i][j], MATNEST, &isNest));

1027:         PetscCall(PetscViewerASCIIPrintf(viewer, "(%" PetscInt_FMT ",%" PetscInt_FMT ") : %s%stype=%s, rows=%" PetscInt_FMT ", cols=%" PetscInt_FMT "\n", i, j, name, prefix, type, NR, NC));

1029:         if (isNest || viewSub) {
1030:           PetscCall(PetscViewerASCIIPushTab(viewer)); /* push1 */
1031:           PetscCall(MatView(bA->m[i][j], viewer));
1032:           PetscCall(PetscViewerASCIIPopTab(viewer)); /* pop1 */
1033:         }
1034:       }
1035:     }
1036:     PetscCall(PetscViewerASCIIPopTab(viewer)); /* pop0 */
1037:   }
1038:   PetscFunctionReturn(PETSC_SUCCESS);
1039: }

1041: static PetscErrorCode MatZeroEntries_Nest(Mat A)
1042: {
1043:   Mat_Nest *bA = (Mat_Nest *)A->data;

1045:   PetscFunctionBegin;
1046:   for (PetscInt i = 0; i < bA->nr; i++) {
1047:     for (PetscInt j = 0; j < bA->nc; j++) {
1048:       if (!bA->m[i][j]) continue;
1049:       PetscCall(MatZeroEntries(bA->m[i][j]));
1050:     }
1051:   }
1052:   PetscFunctionReturn(PETSC_SUCCESS);
1053: }

1055: static PetscErrorCode MatCopy_Nest(Mat A, Mat B, MatStructure str)
1056: {
1057:   Mat_Nest *bA = (Mat_Nest *)A->data, *bB = (Mat_Nest *)B->data;
1058:   PetscInt  i, j, nr = bA->nr, nc = bA->nc;
1059:   PetscBool nnzstate = PETSC_FALSE;

1061:   PetscFunctionBegin;
1062:   PetscCheck(nr == bB->nr && nc == bB->nc, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_INCOMP, "Cannot copy a Mat_Nest of block size (%" PetscInt_FMT ",%" PetscInt_FMT ") to a Mat_Nest of block size (%" PetscInt_FMT ",%" PetscInt_FMT ")", bB->nr, bB->nc, nr, nc);
1063:   for (i = 0; i < nr; i++) {
1064:     for (j = 0; j < nc; j++) {
1065:       PetscObjectState subnnzstate = 0;
1066:       if (bA->m[i][j] && bB->m[i][j]) {
1067:         PetscCall(MatCopy(bA->m[i][j], bB->m[i][j], str));
1068:         PetscCall(MatGetNonzeroState(bB->m[i][j], &subnnzstate));
1069:         nnzstate                 = (PetscBool)(nnzstate || bB->nnzstate[i * nc + j] != subnnzstate);
1070:         bB->nnzstate[i * nc + j] = subnnzstate;
1071:       } else if (bA->m[i][j]) { // bB->m[i][j] is NULL
1072:         Mat M;

1074:         PetscCheck(str == DIFFERENT_NONZERO_PATTERN || str == UNKNOWN_NONZERO_PATTERN, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_INCOMP, "Matrix block does not exist at %" PetscInt_FMT ",%" PetscInt_FMT ". Use DIFFERENT_NONZERO_PATTERN or UNKNOWN_NONZERO_PATTERN", i, j);
1075:         PetscCall(MatDuplicate(bA->m[i][j], MAT_COPY_VALUES, &M));
1076:         PetscCall(MatNestSetSubMat(B, i, j, M));
1077:         PetscCall(MatDestroy(&M));
1078:       } else if (bB->m[i][j]) { // bA->m[i][j] is NULL
1079:         PetscCheck(str == DIFFERENT_NONZERO_PATTERN || str == SUBSET_NONZERO_PATTERN || str == UNKNOWN_NONZERO_PATTERN, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_INCOMP, "Matrix block does not exist at %" PetscInt_FMT ",%" PetscInt_FMT ". Use DIFFERENT_NONZERO_PATTERN, SUBSET_NONZERO_PATTERN or UNKNOWN_NONZERO_PATTERN", i, j);
1080:         PetscCall(MatNestSetSubMat(B, i, j, NULL));
1081:       }
1082:     }
1083:   }
1084:   if (nnzstate) B->nonzerostate++;
1085:   PetscFunctionReturn(PETSC_SUCCESS);
1086: }

1088: static PetscErrorCode MatAXPY_Nest(Mat Y, PetscScalar a, Mat X, MatStructure str)
1089: {
1090:   Mat_Nest *bY = (Mat_Nest *)Y->data, *bX = (Mat_Nest *)X->data;
1091:   PetscInt  i, j, nr = bY->nr, nc = bY->nc;
1092:   PetscBool nnzstate = PETSC_FALSE;

1094:   PetscFunctionBegin;
1095:   PetscCheck(nr == bX->nr && nc == bX->nc, PetscObjectComm((PetscObject)Y), PETSC_ERR_ARG_INCOMP, "Cannot AXPY a MatNest of block size (%" PetscInt_FMT ",%" PetscInt_FMT ") with a MatNest of block size (%" PetscInt_FMT ",%" PetscInt_FMT ")", bX->nr, bX->nc, nr, nc);
1096:   for (i = 0; i < nr; i++) {
1097:     for (j = 0; j < nc; j++) {
1098:       PetscObjectState subnnzstate = 0;
1099:       if (bY->m[i][j] && bX->m[i][j]) {
1100:         PetscCall(MatAXPY(bY->m[i][j], a, bX->m[i][j], str));
1101:       } else if (bX->m[i][j]) {
1102:         Mat M;

1104:         PetscCheck(str == DIFFERENT_NONZERO_PATTERN || str == UNKNOWN_NONZERO_PATTERN, PetscObjectComm((PetscObject)Y), PETSC_ERR_ARG_INCOMP, "Matrix block does not exist at %" PetscInt_FMT ",%" PetscInt_FMT ". Use DIFFERENT_NONZERO_PATTERN or UNKNOWN_NONZERO_PATTERN", i, j);
1105:         PetscCall(MatDuplicate(bX->m[i][j], MAT_COPY_VALUES, &M));
1106:         PetscCall(MatScale(M, a));
1107:         PetscCall(MatNestSetSubMat(Y, i, j, M));
1108:         PetscCall(MatDestroy(&M));
1109:       }
1110:       if (bY->m[i][j]) PetscCall(MatGetNonzeroState(bY->m[i][j], &subnnzstate));
1111:       nnzstate                 = (PetscBool)(nnzstate || bY->nnzstate[i * nc + j] != subnnzstate);
1112:       bY->nnzstate[i * nc + j] = subnnzstate;
1113:     }
1114:   }
1115:   if (nnzstate) Y->nonzerostate++;
1116:   PetscFunctionReturn(PETSC_SUCCESS);
1117: }

1119: static PetscErrorCode MatDuplicate_Nest(Mat A, MatDuplicateOption op, Mat *B)
1120: {
1121:   Mat_Nest *bA = (Mat_Nest *)A->data;
1122:   Mat      *b;
1123:   PetscInt  i, j, nr = bA->nr, nc = bA->nc;

1125:   PetscFunctionBegin;
1126:   PetscCall(PetscMalloc1(nr * nc, &b));
1127:   for (i = 0; i < nr; i++) {
1128:     for (j = 0; j < nc; j++) {
1129:       if (bA->m[i][j]) PetscCall(MatDuplicate(bA->m[i][j], op, &b[i * nc + j]));
1130:       else b[i * nc + j] = NULL;
1131:     }
1132:   }
1133:   PetscCall(MatCreateNest(PetscObjectComm((PetscObject)A), nr, bA->isglobal.row, nc, bA->isglobal.col, b, B));
1134:   /* Give the new MatNest exclusive ownership */
1135:   for (i = 0; i < nr * nc; i++) PetscCall(MatDestroy(&b[i]));
1136:   PetscCall(PetscFree(b));

1138:   PetscCall(MatAssemblyBegin(*B, MAT_FINAL_ASSEMBLY));
1139:   PetscCall(MatAssemblyEnd(*B, MAT_FINAL_ASSEMBLY));
1140:   PetscFunctionReturn(PETSC_SUCCESS);
1141: }

1143: /* nest api */
1144: static PetscErrorCode MatNestGetSubMat_Nest(Mat A, PetscInt idxm, PetscInt jdxm, Mat *mat)
1145: {
1146:   Mat_Nest *bA = (Mat_Nest *)A->data;

1148:   PetscFunctionBegin;
1149:   PetscCheck(idxm < bA->nr, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Row too large: row %" PetscInt_FMT " max %" PetscInt_FMT, idxm, bA->nr - 1);
1150:   PetscCheck(jdxm < bA->nc, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Col too large: row %" PetscInt_FMT " max %" PetscInt_FMT, jdxm, bA->nc - 1);
1151:   *mat = bA->m[idxm][jdxm];
1152:   PetscFunctionReturn(PETSC_SUCCESS);
1153: }

1155: /*@
1156:   MatNestGetSubMat - Returns a single, sub-matrix from a `MATNEST`

1158:   Not Collective

1160:   Input Parameters:
1161: + A    - `MATNEST` matrix
1162: . idxm - index of the matrix within the nest matrix
1163: - jdxm - index of the matrix within the nest matrix

1165:   Output Parameter:
1166: . sub - matrix at index `idxm`, `jdxm` within the nest matrix

1168:   Level: developer

1170: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestGetSize()`, `MatNestGetSubMats()`, `MatCreateNest()`, `MatNestSetSubMat()`,
1171:           `MatNestGetLocalISs()`, `MatNestGetISs()`
1172: @*/
1173: PetscErrorCode MatNestGetSubMat(Mat A, PetscInt idxm, PetscInt jdxm, Mat *sub)
1174: {
1175:   PetscFunctionBegin;
1179:   PetscAssertPointer(sub, 4);
1180:   PetscUseMethod(A, "MatNestGetSubMat_C", (Mat, PetscInt, PetscInt, Mat *), (A, idxm, jdxm, sub));
1181:   PetscFunctionReturn(PETSC_SUCCESS);
1182: }

1184: static PetscErrorCode MatNestSetSubMat_Nest(Mat A, PetscInt idxm, PetscInt jdxm, Mat mat)
1185: {
1186:   Mat_Nest *bA = (Mat_Nest *)A->data;
1187:   PetscInt  m, n, M, N, mi, ni, Mi, Ni;

1189:   PetscFunctionBegin;
1190:   PetscCheck(idxm < bA->nr, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Row too large: row %" PetscInt_FMT " max %" PetscInt_FMT, idxm, bA->nr - 1);
1191:   PetscCheck(jdxm < bA->nc, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Col too large: row %" PetscInt_FMT " max %" PetscInt_FMT, jdxm, bA->nc - 1);
1192:   if (mat) {
1193:     PetscCall(MatGetLocalSize(mat, &m, &n));
1194:     PetscCall(MatGetSize(mat, &M, &N));
1195:     PetscCall(ISGetLocalSize(bA->isglobal.row[idxm], &mi));
1196:     PetscCall(ISGetSize(bA->isglobal.row[idxm], &Mi));
1197:     PetscCall(ISGetLocalSize(bA->isglobal.col[jdxm], &ni));
1198:     PetscCall(ISGetSize(bA->isglobal.col[jdxm], &Ni));
1199:     PetscCheck(M == Mi && N == Ni, PetscObjectComm((PetscObject)mat), PETSC_ERR_ARG_INCOMP, "Submatrix dimension (%" PetscInt_FMT ",%" PetscInt_FMT ") incompatible with nest block (%" PetscInt_FMT ",%" PetscInt_FMT ")", M, N, Mi, Ni);
1200:     PetscCheck(m == mi && n == ni, PetscObjectComm((PetscObject)mat), PETSC_ERR_ARG_INCOMP, "Submatrix local dimension (%" PetscInt_FMT ",%" PetscInt_FMT ") incompatible with nest block (%" PetscInt_FMT ",%" PetscInt_FMT ")", m, n, mi, ni);
1201:   }

1203:   /* do not increase object state */
1204:   if (mat == bA->m[idxm][jdxm]) PetscFunctionReturn(PETSC_SUCCESS);

1206:   PetscCall(PetscObjectReference((PetscObject)mat));
1207:   PetscCall(MatDestroy(&bA->m[idxm][jdxm]));
1208:   bA->m[idxm][jdxm] = mat;
1209:   PetscCall(PetscObjectStateIncrease((PetscObject)A));
1210:   if (mat) PetscCall(MatGetNonzeroState(mat, &bA->nnzstate[idxm * bA->nc + jdxm]));
1211:   else bA->nnzstate[idxm * bA->nc + jdxm] = 0;
1212:   A->nonzerostate++;
1213:   PetscFunctionReturn(PETSC_SUCCESS);
1214: }

1216: /*@
1217:   MatNestSetSubMat - Set a single submatrix in the `MATNEST`

1219:   Logically Collective

1221:   Input Parameters:
1222: + A    - `MATNEST` matrix
1223: . idxm - index of the matrix within the nest matrix
1224: . jdxm - index of the matrix within the nest matrix
1225: - sub  - matrix at index `idxm`, `jdxm` within the nest matrix

1227:   Level: developer

1229:   Notes:
1230:   The new submatrix must have the same size and communicator as that block of the nest.

1232:   This increments the reference count of the submatrix.

1234: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestSetSubMats()`, `MatNestGetSubMats()`, `MatNestGetLocalISs()`, `MatCreateNest()`,
1235:           `MatNestGetSubMat()`, `MatNestGetISs()`, `MatNestGetSize()`
1236: @*/
1237: PetscErrorCode MatNestSetSubMat(Mat A, PetscInt idxm, PetscInt jdxm, Mat sub)
1238: {
1239:   PetscFunctionBegin;
1244:   PetscTryMethod(A, "MatNestSetSubMat_C", (Mat, PetscInt, PetscInt, Mat), (A, idxm, jdxm, sub));
1245:   PetscFunctionReturn(PETSC_SUCCESS);
1246: }

1248: static PetscErrorCode MatNestGetSubMats_Nest(Mat A, PetscInt *M, PetscInt *N, Mat ***mat)
1249: {
1250:   Mat_Nest *bA = (Mat_Nest *)A->data;

1252:   PetscFunctionBegin;
1253:   if (M) *M = bA->nr;
1254:   if (N) *N = bA->nc;
1255:   if (mat) *mat = bA->m;
1256:   PetscFunctionReturn(PETSC_SUCCESS);
1257: }

1259: /*@C
1260:   MatNestGetSubMats - Returns the entire two dimensional array of matrices defining a `MATNEST` matrix.

1262:   Not Collective

1264:   Input Parameter:
1265: . A - nest matrix

1267:   Output Parameters:
1268: + M   - number of submatrix rows in the nest matrix
1269: . N   - number of submatrix columns in the nest matrix
1270: - mat - array of matrices

1272:   Level: developer

1274:   Note:
1275:   The user should not free the array `mat`.

1277:   Fortran Notes:
1278:   This routine has a calling sequence `call MatNestGetSubMats(A, M, N, mat, ierr)`
1279:   where the space allocated for the optional argument `mat` is assumed large enough (if provided).
1280:   Matrices in `mat` are returned in row-major order, see `MatCreateNest()` for an example.

1282: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestGetSize()`, `MatNestGetSubMat()`, `MatNestGetLocalISs()`, `MatCreateNest()`,
1283:           `MatNestSetSubMats()`, `MatNestGetISs()`, `MatNestSetSubMat()`
1284: @*/
1285: PetscErrorCode MatNestGetSubMats(Mat A, PetscInt *M, PetscInt *N, Mat ***mat)
1286: {
1287:   PetscFunctionBegin;
1289:   PetscUseMethod(A, "MatNestGetSubMats_C", (Mat, PetscInt *, PetscInt *, Mat ***), (A, M, N, mat));
1290:   PetscFunctionReturn(PETSC_SUCCESS);
1291: }

1293: static PetscErrorCode MatNestGetSize_Nest(Mat A, PetscInt *M, PetscInt *N)
1294: {
1295:   Mat_Nest *bA = (Mat_Nest *)A->data;

1297:   PetscFunctionBegin;
1298:   if (M) *M = bA->nr;
1299:   if (N) *N = bA->nc;
1300:   PetscFunctionReturn(PETSC_SUCCESS);
1301: }

1303: /*@
1304:   MatNestGetSize - Returns the size of the `MATNEST` matrix.

1306:   Not Collective

1308:   Input Parameter:
1309: . A - `MATNEST` matrix

1311:   Output Parameters:
1312: + M - number of rows in the nested mat
1313: - N - number of cols in the nested mat

1315:   Level: developer

1317:   Note:
1318:   `size` refers to the number of submatrices in the row and column directions of the nested matrix

1320: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestGetSubMat()`, `MatNestGetSubMats()`, `MatCreateNest()`, `MatNestGetLocalISs()`,
1321:           `MatNestGetISs()`
1322: @*/
1323: PetscErrorCode MatNestGetSize(Mat A, PetscInt *M, PetscInt *N)
1324: {
1325:   PetscFunctionBegin;
1327:   PetscUseMethod(A, "MatNestGetSize_C", (Mat, PetscInt *, PetscInt *), (A, M, N));
1328:   PetscFunctionReturn(PETSC_SUCCESS);
1329: }

1331: static PetscErrorCode MatNestGetISs_Nest(Mat A, IS rows[], IS cols[])
1332: {
1333:   Mat_Nest *vs = (Mat_Nest *)A->data;

1335:   PetscFunctionBegin;
1336:   if (rows) {
1337:     for (PetscInt i = 0; i < vs->nr; i++) rows[i] = vs->isglobal.row[i];
1338:   }
1339:   if (cols) {
1340:     for (PetscInt i = 0; i < vs->nc; i++) cols[i] = vs->isglobal.col[i];
1341:   }
1342:   PetscFunctionReturn(PETSC_SUCCESS);
1343: }

1345: /*@
1346:   MatNestGetISs - Returns the index sets partitioning the row and column spaces of a `MATNEST`

1348:   Not Collective

1350:   Input Parameter:
1351: . A - `MATNEST` matrix

1353:   Output Parameters:
1354: + rows - array of row index sets (pass `NULL` to ignore)
1355: - cols - array of column index sets (pass `NULL` to ignore)

1357:   Level: advanced

1359:   Note:
1360:   The user must have allocated arrays of the correct size. The reference count is not increased on the returned `IS`s.

1362: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestGetSubMat()`, `MatNestGetSubMats()`, `MatNestGetSize()`, `MatNestGetLocalISs()`,
1363:           `MatCreateNest()`, `MatNestSetSubMats()`
1364: @*/
1365: PetscErrorCode MatNestGetISs(Mat A, IS rows[], IS cols[])
1366: {
1367:   PetscFunctionBegin;
1369:   PetscUseMethod(A, "MatNestGetISs_C", (Mat, IS[], IS[]), (A, rows, cols));
1370:   PetscFunctionReturn(PETSC_SUCCESS);
1371: }

1373: static PetscErrorCode MatNestGetLocalISs_Nest(Mat A, IS rows[], IS cols[])
1374: {
1375:   Mat_Nest *vs = (Mat_Nest *)A->data;

1377:   PetscFunctionBegin;
1378:   if (rows) {
1379:     for (PetscInt i = 0; i < vs->nr; i++) rows[i] = vs->islocal.row[i];
1380:   }
1381:   if (cols) {
1382:     for (PetscInt i = 0; i < vs->nc; i++) cols[i] = vs->islocal.col[i];
1383:   }
1384:   PetscFunctionReturn(PETSC_SUCCESS);
1385: }

1387: /*@
1388:   MatNestGetLocalISs - Returns the index sets partitioning the row and column spaces of a `MATNEST`

1390:   Not Collective

1392:   Input Parameter:
1393: . A - `MATNEST` matrix

1395:   Output Parameters:
1396: + rows - array of row index sets (pass `NULL` to ignore)
1397: - cols - array of column index sets (pass `NULL` to ignore)

1399:   Level: advanced

1401:   Note:
1402:   The user must have allocated arrays of the correct size. The reference count is not increased on the returned `IS`s.

1404: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestGetSubMat()`, `MatNestGetSubMats()`, `MatNestGetSize()`, `MatNestGetISs()`, `MatCreateNest()`,
1405:           `MatNestSetSubMats()`, `MatNestSetSubMat()`
1406: @*/
1407: PetscErrorCode MatNestGetLocalISs(Mat A, IS rows[], IS cols[])
1408: {
1409:   PetscFunctionBegin;
1411:   PetscUseMethod(A, "MatNestGetLocalISs_C", (Mat, IS[], IS[]), (A, rows, cols));
1412:   PetscFunctionReturn(PETSC_SUCCESS);
1413: }

1415: static PetscErrorCode MatNestSetVecType_Nest(Mat A, VecType vtype)
1416: {
1417:   PetscBool flg;

1419:   PetscFunctionBegin;
1420:   PetscCall(PetscStrcmp(vtype, VECNEST, &flg));
1421:   /* In reality, this only distinguishes VECNEST and "other" */
1422:   if (flg) A->ops->getvecs = MatCreateVecs_Nest;
1423:   else A->ops->getvecs = NULL;
1424:   PetscFunctionReturn(PETSC_SUCCESS);
1425: }

1427: /*@
1428:   MatNestSetVecType - Sets the type of `Vec` returned by `MatCreateVecs()`

1430:   Not Collective

1432:   Input Parameters:
1433: + A     - `MATNEST` matrix
1434: - vtype - `VecType` to use for creating vectors

1436:   Level: developer

1438: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatCreateVecs()`, `MatCreateNest()`, `VecType`
1439: @*/
1440: PetscErrorCode MatNestSetVecType(Mat A, VecType vtype)
1441: {
1442:   PetscFunctionBegin;
1444:   PetscTryMethod(A, "MatNestSetVecType_C", (Mat, VecType), (A, vtype));
1445:   PetscFunctionReturn(PETSC_SUCCESS);
1446: }

1448: static PetscErrorCode MatNestSetSubMats_Nest(Mat A, PetscInt nr, const IS is_row[], PetscInt nc, const IS is_col[], const Mat a[])
1449: {
1450:   Mat_Nest *s = (Mat_Nest *)A->data;
1451:   PetscInt  i, j, m, n, M, N;
1452:   PetscBool cong, isstd, sametype = PETSC_FALSE;
1453:   VecType   vtype, type;

1455:   PetscFunctionBegin;
1456:   PetscCall(MatReset_Nest(A));

1458:   s->nr = nr;
1459:   s->nc = nc;

1461:   /* Create space for submatrices */
1462:   PetscCall(PetscMalloc1(nr, &s->m));
1463:   PetscCall(PetscMalloc1(nr * nc, &s->m[0]));
1464:   for (i = 0; i < nr; i++) {
1465:     s->m[i] = s->m[0] + i * nc;
1466:     for (j = 0; j < nc; j++) {
1467:       s->m[i][j] = a ? a[i * nc + j] : NULL;
1468:       PetscCall(PetscObjectReference((PetscObject)s->m[i][j]));
1469:     }
1470:   }
1471:   PetscCall(MatGetVecType(A, &vtype));
1472:   PetscCall(PetscStrcmp(vtype, VECSTANDARD, &isstd));
1473:   if (isstd) {
1474:     /* check if all blocks have the same vectype */
1475:     vtype = NULL;
1476:     for (i = 0; i < nr; i++) {
1477:       for (j = 0; j < nc; j++) {
1478:         if (s->m[i][j]) {
1479:           if (!vtype) { /* first visited block */
1480:             PetscCall(MatGetVecType(s->m[i][j], &vtype));
1481:             sametype = PETSC_TRUE;
1482:           } else if (sametype) {
1483:             PetscCall(MatGetVecType(s->m[i][j], &type));
1484:             PetscCall(PetscStrcmp(vtype, type, &sametype));
1485:           }
1486:         }
1487:       }
1488:     }
1489:     if (sametype) { /* propagate vectype */
1490:       PetscCall(MatSetVecType(A, vtype));
1491:     }
1492:   }

1494:   PetscCall(MatSetUp_NestIS_Private(A, nr, is_row, nc, is_col));

1496:   PetscCall(PetscMalloc1(nr, &s->row_len));
1497:   PetscCall(PetscMalloc1(nc, &s->col_len));
1498:   for (i = 0; i < nr; i++) s->row_len[i] = -1;
1499:   for (j = 0; j < nc; j++) s->col_len[j] = -1;

1501:   PetscCall(PetscCalloc1(nr * nc, &s->nnzstate));
1502:   for (i = 0; i < nr; i++) {
1503:     for (j = 0; j < nc; j++) {
1504:       if (s->m[i][j]) PetscCall(MatGetNonzeroState(s->m[i][j], &s->nnzstate[i * nc + j]));
1505:     }
1506:   }

1508:   PetscCall(MatNestGetSizes_Private(A, &m, &n, &M, &N));

1510:   PetscCall(PetscLayoutSetSize(A->rmap, M));
1511:   PetscCall(PetscLayoutSetLocalSize(A->rmap, m));
1512:   PetscCall(PetscLayoutSetSize(A->cmap, N));
1513:   PetscCall(PetscLayoutSetLocalSize(A->cmap, n));

1515:   PetscCall(PetscLayoutSetUp(A->rmap));
1516:   PetscCall(PetscLayoutSetUp(A->cmap));

1518:   /* disable operations that are not supported for non-square matrices,
1519:      or matrices for which is_row != is_col  */
1520:   PetscCall(MatHasCongruentLayouts(A, &cong));
1521:   if (cong && nr != nc) cong = PETSC_FALSE;
1522:   if (cong) {
1523:     for (i = 0; cong && i < nr; i++) PetscCall(ISEqualUnsorted(s->isglobal.row[i], s->isglobal.col[i], &cong));
1524:   }
1525:   if (!cong) {
1526:     A->ops->getdiagonal = NULL;
1527:     A->ops->shift       = NULL;
1528:     A->ops->diagonalset = NULL;
1529:   }

1531:   PetscCall(PetscCalloc2(nr, &s->left, nc, &s->right));
1532:   PetscCall(PetscObjectStateIncrease((PetscObject)A));
1533:   A->nonzerostate++;
1534:   PetscFunctionReturn(PETSC_SUCCESS);
1535: }

1537: /*@
1538:   MatNestSetSubMats - Sets the nested submatrices in a `MATNEST`

1540:   Collective

1542:   Input Parameters:
1543: + A      - `MATNEST` matrix
1544: . nr     - number of nested row blocks
1545: . is_row - index sets for each nested row block, or `NULL` to make contiguous
1546: . nc     - number of nested column blocks
1547: . is_col - index sets for each nested column block, or `NULL` to make contiguous
1548: - a      - array of $ nr \times nc$ submatrices, or `NULL`

1550:   Level: advanced

1552:   Notes:
1553:   This always resets any block matrix information previously set.

1555:   Pass `NULL` in the corresponding entry of `a` for an empty block.

1557:   In both C and Fortran, `a` must be a one-dimensional array representing a two-dimensional row-major order array containing the matrices. See
1558:   `MatCreateNest()` for an example.

1560:   Fortran Note:
1561:   Pass `PETSC_NULL_MAT` in the corresponding entry of `a` for an empty block

1563: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatCreateNest()`, `MatNestSetSubMat()`, `MatNestGetSubMat()`, `MatNestGetSubMats()`
1564: @*/
1565: PetscErrorCode MatNestSetSubMats(Mat A, PetscInt nr, const IS is_row[], PetscInt nc, const IS is_col[], const Mat a[]) PeNSS
1566: {
1567:   PetscFunctionBegin;
1570:   PetscCheck(nr >= 0, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Number of rows cannot be negative");
1571:   if (nr && is_row) {
1572:     PetscAssertPointer(is_row, 3);
1574:   }
1576:   PetscCheck(nc >= 0, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Number of columns cannot be negative");
1577:   if (nc && is_col) {
1578:     PetscAssertPointer(is_col, 5);
1580:   }
1581:   PetscTryMethod(A, "MatNestSetSubMats_C", (Mat, PetscInt, const IS[], PetscInt, const IS[], const Mat[]), (A, nr, is_row, nc, is_col, a));
1582:   PetscFunctionReturn(PETSC_SUCCESS);
1583: }

1585: static PetscErrorCode MatNestCreateAggregateL2G_Private(Mat A, PetscInt n, const IS islocal[], const IS isglobal[], PetscBool colflg, ISLocalToGlobalMapping *ltog)
1586: {
1587:   PetscBool flg;
1588:   PetscInt  i, j, m, mi, *ix;

1590:   PetscFunctionBegin;
1591:   *ltog = NULL;
1592:   for (i = 0, m = 0, flg = PETSC_FALSE; i < n; i++) {
1593:     if (islocal[i]) {
1594:       PetscCall(ISGetLocalSize(islocal[i], &mi));
1595:       flg = PETSC_TRUE; /* We found a non-trivial entry */
1596:     } else {
1597:       PetscCall(ISGetLocalSize(isglobal[i], &mi));
1598:     }
1599:     m += mi;
1600:   }
1601:   if (!flg) PetscFunctionReturn(PETSC_SUCCESS);

1603:   PetscCall(PetscMalloc1(m, &ix));
1604:   for (i = 0, m = 0; i < n; i++) {
1605:     ISLocalToGlobalMapping smap = NULL;
1606:     Mat                    sub  = NULL;
1607:     PetscSF                sf;
1608:     PetscLayout            map;
1609:     const PetscInt        *ix2;

1611:     if (!colflg) {
1612:       PetscCall(MatNestFindNonzeroSubMatRow(A, i, &sub));
1613:     } else {
1614:       PetscCall(MatNestFindNonzeroSubMatCol(A, i, &sub));
1615:     }
1616:     if (sub) {
1617:       if (!colflg) PetscCall(MatGetLocalToGlobalMapping(sub, &smap, NULL));
1618:       else PetscCall(MatGetLocalToGlobalMapping(sub, NULL, &smap));
1619:     }
1620:     /*
1621:        Now we need to extract the monolithic global indices that correspond to the given split global indices.
1622:        In many/most cases, we only want MatGetLocalSubMatrix() to work, in which case we only need to know the size of the local spaces.
1623:     */
1624:     PetscCall(ISGetIndices(isglobal[i], &ix2));
1625:     if (islocal[i]) {
1626:       PetscInt *ilocal, *iremote;
1627:       PetscInt  mil, nleaves;

1629:       PetscCall(ISGetLocalSize(islocal[i], &mi));
1630:       PetscCheck(smap, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing local to global map");
1631:       for (j = 0; j < mi; j++) ix[m + j] = j;
1632:       PetscCall(ISLocalToGlobalMappingApply(smap, mi, ix + m, ix + m));

1634:       /* PetscSFSetGraphLayout does not like negative indices */
1635:       PetscCall(PetscMalloc2(mi, &ilocal, mi, &iremote));
1636:       for (j = 0, nleaves = 0; j < mi; j++) {
1637:         if (ix[m + j] < 0) continue;
1638:         ilocal[nleaves]  = j;
1639:         iremote[nleaves] = ix[m + j];
1640:         nleaves++;
1641:       }
1642:       PetscCall(ISGetLocalSize(isglobal[i], &mil));
1643:       PetscCall(PetscSFCreate(PetscObjectComm((PetscObject)A), &sf));
1644:       PetscCall(PetscLayoutCreate(PetscObjectComm((PetscObject)A), &map));
1645:       PetscCall(PetscLayoutSetLocalSize(map, mil));
1646:       PetscCall(PetscLayoutSetUp(map));
1647:       PetscCall(PetscSFSetGraphLayout(sf, map, nleaves, ilocal, PETSC_USE_POINTER, iremote));
1648:       PetscCall(PetscLayoutDestroy(&map));
1649:       PetscCall(PetscSFBcastBegin(sf, MPIU_INT, ix2, ix + m, MPI_REPLACE));
1650:       PetscCall(PetscSFBcastEnd(sf, MPIU_INT, ix2, ix + m, MPI_REPLACE));
1651:       PetscCall(PetscSFDestroy(&sf));
1652:       PetscCall(PetscFree2(ilocal, iremote));
1653:     } else {
1654:       PetscCall(ISGetLocalSize(isglobal[i], &mi));
1655:       for (j = 0; j < mi; j++) ix[m + j] = ix2[j];
1656:     }
1657:     PetscCall(ISRestoreIndices(isglobal[i], &ix2));
1658:     m += mi;
1659:   }
1660:   PetscCall(ISLocalToGlobalMappingCreate(PetscObjectComm((PetscObject)A), 1, m, ix, PETSC_OWN_POINTER, ltog));
1661:   PetscFunctionReturn(PETSC_SUCCESS);
1662: }

1664: /* If an IS was provided, there is nothing Nest needs to do, otherwise Nest will build a strided IS */
1665: /*
1666:   nprocessors = NP
1667:   Nest x^T = ((g_0,g_1,...g_nprocs-1), (h_0,h_1,...h_NP-1))
1668:        proc 0: => (g_0,h_0,)
1669:        proc 1: => (g_1,h_1,)
1670:        ...
1671:        proc nprocs-1: => (g_NP-1,h_NP-1,)

1673:             proc 0:                      proc 1:                    proc nprocs-1:
1674:     is[0] = (0,1,2,...,nlocal(g_0)-1)  (0,1,...,nlocal(g_1)-1)  (0,1,...,nlocal(g_NP-1))

1676:             proc 0:
1677:     is[1] = (nlocal(g_0),nlocal(g_0)+1,...,nlocal(g_0)+nlocal(h_0)-1)
1678:             proc 1:
1679:     is[1] = (nlocal(g_1),nlocal(g_1)+1,...,nlocal(g_1)+nlocal(h_1)-1)

1681:             proc NP-1:
1682:     is[1] = (nlocal(g_NP-1),nlocal(g_NP-1)+1,...,nlocal(g_NP-1)+nlocal(h_NP-1)-1)
1683: */
1684: static PetscErrorCode MatSetUp_NestIS_Private(Mat A, PetscInt nr, const IS is_row[], PetscInt nc, const IS is_col[])
1685: {
1686:   Mat_Nest *vs = (Mat_Nest *)A->data;
1687:   PetscInt  i, j, offset, n, nsum, bs;
1688:   Mat       sub = NULL;

1690:   PetscFunctionBegin;
1691:   PetscCall(PetscMalloc1(nr, &vs->isglobal.row));
1692:   PetscCall(PetscMalloc1(nc, &vs->isglobal.col));
1693:   if (is_row) { /* valid IS is passed in */
1694:     /* refs on is[] are incremented */
1695:     for (i = 0; i < vs->nr; i++) {
1696:       PetscCall(PetscObjectReference((PetscObject)is_row[i]));
1697:       vs->isglobal.row[i] = is_row[i];
1698:     }
1699:   } else { /* Create the ISs by inspecting sizes of a submatrix in each row */
1700:     nsum = 0;
1701:     for (i = 0; i < vs->nr; i++) { /* Add up the local sizes to compute the aggregate offset */
1702:       PetscCall(MatNestFindNonzeroSubMatRow(A, i, &sub));
1703:       PetscCheck(sub, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "No nonzero submatrix in row %" PetscInt_FMT, i);
1704:       PetscCall(MatGetLocalSize(sub, &n, NULL));
1705:       PetscCheck(n >= 0, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Sizes have not yet been set for submatrix");
1706:       nsum += n;
1707:     }
1708:     PetscCallMPI(MPI_Scan(&nsum, &offset, 1, MPIU_INT, MPI_SUM, PetscObjectComm((PetscObject)A)));
1709:     offset -= nsum;
1710:     for (i = 0; i < vs->nr; i++) {
1711:       PetscCall(MatNestFindNonzeroSubMatRow(A, i, &sub));
1712:       PetscCall(MatGetLocalSize(sub, &n, NULL));
1713:       PetscCall(MatGetBlockSizes(sub, &bs, NULL));
1714:       PetscCall(ISCreateStride(PetscObjectComm((PetscObject)sub), n, offset, 1, &vs->isglobal.row[i]));
1715:       PetscCall(ISSetBlockSize(vs->isglobal.row[i], bs));
1716:       offset += n;
1717:     }
1718:   }

1720:   if (is_col) { /* valid IS is passed in */
1721:     /* refs on is[] are incremented */
1722:     for (j = 0; j < vs->nc; j++) {
1723:       PetscCall(PetscObjectReference((PetscObject)is_col[j]));
1724:       vs->isglobal.col[j] = is_col[j];
1725:     }
1726:   } else { /* Create the ISs by inspecting sizes of a submatrix in each column */
1727:     offset = A->cmap->rstart;
1728:     nsum   = 0;
1729:     for (j = 0; j < vs->nc; j++) {
1730:       PetscCall(MatNestFindNonzeroSubMatCol(A, j, &sub));
1731:       PetscCheck(sub, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "No nonzero submatrix in column %" PetscInt_FMT, i);
1732:       PetscCall(MatGetLocalSize(sub, NULL, &n));
1733:       PetscCheck(n >= 0, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Sizes have not yet been set for submatrix");
1734:       nsum += n;
1735:     }
1736:     PetscCallMPI(MPI_Scan(&nsum, &offset, 1, MPIU_INT, MPI_SUM, PetscObjectComm((PetscObject)A)));
1737:     offset -= nsum;
1738:     for (j = 0; j < vs->nc; j++) {
1739:       PetscCall(MatNestFindNonzeroSubMatCol(A, j, &sub));
1740:       PetscCall(MatGetLocalSize(sub, NULL, &n));
1741:       PetscCall(MatGetBlockSizes(sub, NULL, &bs));
1742:       PetscCall(ISCreateStride(PetscObjectComm((PetscObject)sub), n, offset, 1, &vs->isglobal.col[j]));
1743:       PetscCall(ISSetBlockSize(vs->isglobal.col[j], bs));
1744:       offset += n;
1745:     }
1746:   }

1748:   /* Set up the local ISs */
1749:   PetscCall(PetscMalloc1(vs->nr, &vs->islocal.row));
1750:   PetscCall(PetscMalloc1(vs->nc, &vs->islocal.col));
1751:   for (i = 0, offset = 0; i < vs->nr; i++) {
1752:     IS                     isloc;
1753:     ISLocalToGlobalMapping rmap = NULL;
1754:     PetscInt               nlocal, bs;
1755:     PetscCall(MatNestFindNonzeroSubMatRow(A, i, &sub));
1756:     if (sub) PetscCall(MatGetLocalToGlobalMapping(sub, &rmap, NULL));
1757:     if (rmap) {
1758:       PetscCall(MatGetBlockSizes(sub, &bs, NULL));
1759:       PetscCall(ISLocalToGlobalMappingGetSize(rmap, &nlocal));
1760:       PetscCall(ISCreateStride(PETSC_COMM_SELF, nlocal, offset, 1, &isloc));
1761:       PetscCall(ISSetBlockSize(isloc, bs));
1762:     } else {
1763:       nlocal = 0;
1764:       isloc  = NULL;
1765:     }
1766:     vs->islocal.row[i] = isloc;
1767:     offset += nlocal;
1768:   }
1769:   for (i = 0, offset = 0; i < vs->nc; i++) {
1770:     IS                     isloc;
1771:     ISLocalToGlobalMapping cmap = NULL;
1772:     PetscInt               nlocal, bs;
1773:     PetscCall(MatNestFindNonzeroSubMatCol(A, i, &sub));
1774:     if (sub) PetscCall(MatGetLocalToGlobalMapping(sub, NULL, &cmap));
1775:     if (cmap) {
1776:       PetscCall(MatGetBlockSizes(sub, NULL, &bs));
1777:       PetscCall(ISLocalToGlobalMappingGetSize(cmap, &nlocal));
1778:       PetscCall(ISCreateStride(PETSC_COMM_SELF, nlocal, offset, 1, &isloc));
1779:       PetscCall(ISSetBlockSize(isloc, bs));
1780:     } else {
1781:       nlocal = 0;
1782:       isloc  = NULL;
1783:     }
1784:     vs->islocal.col[i] = isloc;
1785:     offset += nlocal;
1786:   }

1788:   /* Set up the aggregate ISLocalToGlobalMapping */
1789:   {
1790:     ISLocalToGlobalMapping rmap, cmap;
1791:     PetscCall(MatNestCreateAggregateL2G_Private(A, vs->nr, vs->islocal.row, vs->isglobal.row, PETSC_FALSE, &rmap));
1792:     PetscCall(MatNestCreateAggregateL2G_Private(A, vs->nc, vs->islocal.col, vs->isglobal.col, PETSC_TRUE, &cmap));
1793:     if (rmap && cmap) PetscCall(MatSetLocalToGlobalMapping(A, rmap, cmap));
1794:     PetscCall(ISLocalToGlobalMappingDestroy(&rmap));
1795:     PetscCall(ISLocalToGlobalMappingDestroy(&cmap));
1796:   }

1798:   if (PetscDefined(USE_DEBUG)) {
1799:     for (i = 0; i < vs->nr; i++) {
1800:       for (j = 0; j < vs->nc; j++) {
1801:         PetscInt m, n, M, N, mi, ni, Mi, Ni;
1802:         Mat      B = vs->m[i][j];
1803:         if (!B) continue;
1804:         PetscCall(MatGetSize(B, &M, &N));
1805:         PetscCall(MatGetLocalSize(B, &m, &n));
1806:         PetscCall(ISGetSize(vs->isglobal.row[i], &Mi));
1807:         PetscCall(ISGetSize(vs->isglobal.col[j], &Ni));
1808:         PetscCall(ISGetLocalSize(vs->isglobal.row[i], &mi));
1809:         PetscCall(ISGetLocalSize(vs->isglobal.col[j], &ni));
1810:         PetscCheck(M == Mi && N == Ni, PetscObjectComm((PetscObject)sub), PETSC_ERR_ARG_INCOMP, "Global sizes (%" PetscInt_FMT ",%" PetscInt_FMT ") of nested submatrix (%" PetscInt_FMT ",%" PetscInt_FMT ") do not agree with space defined by index sets (%" PetscInt_FMT ",%" PetscInt_FMT ")", M, N, i, j, Mi, Ni);
1811:         PetscCheck(m == mi && n == ni, PetscObjectComm((PetscObject)sub), PETSC_ERR_ARG_INCOMP, "Local sizes (%" PetscInt_FMT ",%" PetscInt_FMT ") of nested submatrix (%" PetscInt_FMT ",%" PetscInt_FMT ") do not agree with space defined by index sets (%" PetscInt_FMT ",%" PetscInt_FMT ")", m, n, i, j, mi, ni);
1812:       }
1813:     }
1814:   }

1816:   /* Set A->assembled if all non-null blocks are currently assembled */
1817:   for (i = 0; i < vs->nr; i++) {
1818:     for (j = 0; j < vs->nc; j++) {
1819:       if (vs->m[i][j] && !vs->m[i][j]->assembled) PetscFunctionReturn(PETSC_SUCCESS);
1820:     }
1821:   }
1822:   A->assembled = PETSC_TRUE;
1823:   PetscFunctionReturn(PETSC_SUCCESS);
1824: }

1826: /*@C
1827:   MatCreateNest - Creates a new `MATNEST` matrix containing several nested submatrices, each stored separately

1829:   Collective

1831:   Input Parameters:
1832: + comm   - Communicator for the new `MATNEST`
1833: . nr     - number of nested row blocks
1834: . is_row - index sets for each nested row block, or `NULL` to make contiguous
1835: . nc     - number of nested column blocks
1836: . is_col - index sets for each nested column block, or `NULL` to make contiguous
1837: - a      - array of $nr \times nc$ submatrices, empty submatrices can be passed using `NULL`

1839:   Output Parameter:
1840: . B - new matrix

1842:   Level: advanced

1844:   Note:
1845:   In both C and Fortran, `a` must be a one-dimensional array representing a two-dimensional row-major order array holding references to the matrices.
1846:   For instance, to represent the matrix
1847:   $\begin{bmatrix} A_{11} & A_{12} \\ A_{21} & A_{22}\end{bmatrix}$
1848:   one should use `Mat a[4]={A11,A12,A21,A22}`.

1850:   Fortran Note:
1851:   Pass `PETSC_NULL_MAT` in the corresponding entry of `a` for an empty block

1853: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatCreate()`, `VecCreateNest()`, `DMCreateMatrix()`, `MatNestSetSubMat()`,
1854:           `MatNestGetSubMat()`, `MatNestGetLocalISs()`, `MatNestGetSize()`,
1855:           `MatNestGetISs()`, `MatNestSetSubMats()`, `MatNestGetSubMats()`
1856: @*/
1857: PetscErrorCode MatCreateNest(MPI_Comm comm, PetscInt nr, const IS is_row[], PetscInt nc, const IS is_col[], const Mat a[], Mat *B) PeNSS
1858: {
1859:   PetscFunctionBegin;
1860:   PetscCall(MatCreate(comm, B));
1861:   PetscCall(MatSetType(*B, MATNEST));
1862:   (*B)->preallocated = PETSC_TRUE;
1863:   PetscCall(MatNestSetSubMats(*B, nr, is_row, nc, is_col, a));
1864:   PetscFunctionReturn(PETSC_SUCCESS);
1865: }

1867: static PetscErrorCode MatConvert_Nest_SeqAIJ_fast(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
1868: {
1869:   Mat_Nest     *nest = (Mat_Nest *)A->data;
1870:   Mat          *trans;
1871:   PetscScalar **avv;
1872:   PetscScalar  *vv;
1873:   PetscInt    **aii, **ajj;
1874:   PetscInt     *ii, *jj, *ci;
1875:   PetscInt      nr, nc, nnz, i, j;
1876:   PetscBool     done;

1878:   PetscFunctionBegin;
1879:   PetscCall(MatGetSize(A, &nr, &nc));
1880:   if (reuse == MAT_REUSE_MATRIX) {
1881:     PetscInt rnr;

1883:     PetscCall(MatGetRowIJ(*newmat, 0, PETSC_FALSE, PETSC_FALSE, &rnr, (const PetscInt **)&ii, (const PetscInt **)&jj, &done));
1884:     PetscCheck(done, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "MatGetRowIJ");
1885:     PetscCheck(rnr == nr, PetscObjectComm((PetscObject)A), PETSC_ERR_USER, "Cannot reuse matrix, wrong number of rows");
1886:     PetscCall(MatSeqAIJGetArray(*newmat, &vv));
1887:   }
1888:   /* extract CSR for nested SeqAIJ matrices */
1889:   nnz = 0;
1890:   PetscCall(PetscCalloc4(nest->nr * nest->nc, &aii, nest->nr * nest->nc, &ajj, nest->nr * nest->nc, &avv, nest->nr * nest->nc, &trans));
1891:   for (i = 0; i < nest->nr; ++i) {
1892:     for (j = 0; j < nest->nc; ++j) {
1893:       Mat B = nest->m[i][j];
1894:       if (B) {
1895:         PetscScalar *naa;
1896:         PetscInt    *nii, *njj, nnr;
1897:         PetscBool    istrans;

1899:         PetscCall(PetscObjectTypeCompare((PetscObject)B, MATTRANSPOSEVIRTUAL, &istrans));
1900:         if (istrans) {
1901:           Mat Bt;

1903:           PetscCall(MatTransposeGetMat(B, &Bt));
1904:           PetscCall(MatTranspose(Bt, MAT_INITIAL_MATRIX, &trans[i * nest->nc + j]));
1905:           B = trans[i * nest->nc + j];
1906:         } else {
1907:           PetscCall(PetscObjectTypeCompare((PetscObject)B, MATHERMITIANTRANSPOSEVIRTUAL, &istrans));
1908:           if (istrans) {
1909:             Mat Bt;

1911:             PetscCall(MatHermitianTransposeGetMat(B, &Bt));
1912:             PetscCall(MatHermitianTranspose(Bt, MAT_INITIAL_MATRIX, &trans[i * nest->nc + j]));
1913:             B = trans[i * nest->nc + j];
1914:           }
1915:         }
1916:         PetscCall(MatGetRowIJ(B, 0, PETSC_FALSE, PETSC_FALSE, &nnr, (const PetscInt **)&nii, (const PetscInt **)&njj, &done));
1917:         PetscCheck(done, PetscObjectComm((PetscObject)B), PETSC_ERR_PLIB, "MatGetRowIJ");
1918:         PetscCall(MatSeqAIJGetArray(B, &naa));
1919:         nnz += nii[nnr];

1921:         aii[i * nest->nc + j] = nii;
1922:         ajj[i * nest->nc + j] = njj;
1923:         avv[i * nest->nc + j] = naa;
1924:       }
1925:     }
1926:   }
1927:   if (reuse != MAT_REUSE_MATRIX) {
1928:     PetscCall(PetscMalloc1(nr + 1, &ii));
1929:     PetscCall(PetscMalloc1(nnz, &jj));
1930:     PetscCall(PetscMalloc1(nnz, &vv));
1931:   } else {
1932:     PetscCheck(nnz == ii[nr], PetscObjectComm((PetscObject)A), PETSC_ERR_USER, "Cannot reuse matrix, wrong number of nonzeros");
1933:   }

1935:   /* new row pointer */
1936:   PetscCall(PetscArrayzero(ii, nr + 1));
1937:   for (i = 0; i < nest->nr; ++i) {
1938:     PetscInt ncr, rst;

1940:     PetscCall(ISStrideGetInfo(nest->isglobal.row[i], &rst, NULL));
1941:     PetscCall(ISGetLocalSize(nest->isglobal.row[i], &ncr));
1942:     for (j = 0; j < nest->nc; ++j) {
1943:       if (aii[i * nest->nc + j]) {
1944:         PetscInt *nii = aii[i * nest->nc + j];

1946:         for (PetscInt ir = rst; ir < ncr + rst; ++ir) {
1947:           ii[ir + 1] += nii[1] - nii[0];
1948:           nii++;
1949:         }
1950:       }
1951:     }
1952:   }
1953:   for (i = 0; i < nr; i++) ii[i + 1] += ii[i];

1955:   /* construct CSR for the new matrix */
1956:   PetscCall(PetscCalloc1(nr, &ci));
1957:   for (i = 0; i < nest->nr; ++i) {
1958:     PetscInt ncr, rst;

1960:     PetscCall(ISStrideGetInfo(nest->isglobal.row[i], &rst, NULL));
1961:     PetscCall(ISGetLocalSize(nest->isglobal.row[i], &ncr));
1962:     for (j = 0; j < nest->nc; ++j) {
1963:       if (aii[i * nest->nc + j]) {
1964:         PetscScalar *nvv = avv[i * nest->nc + j], vscale = 1.0, vshift = 0.0;
1965:         PetscInt    *nii = aii[i * nest->nc + j];
1966:         PetscInt    *njj = ajj[i * nest->nc + j];
1967:         PetscInt     cst;

1969:         if (trans[i * nest->nc + j]) {
1970:           vscale = ((Mat_Shell *)nest->m[i][j]->data)->vscale;
1971:           vshift = ((Mat_Shell *)nest->m[i][j]->data)->vshift;
1972:         }
1973:         PetscCall(ISStrideGetInfo(nest->isglobal.col[j], &cst, NULL));
1974:         for (PetscInt ir = rst; ir < ncr + rst; ++ir) {
1975:           PetscInt ij, rsize = nii[1] - nii[0], ist = ii[ir] + ci[ir];

1977:           for (ij = 0; ij < rsize; ij++) {
1978:             jj[ist + ij] = *njj + cst;
1979:             vv[ist + ij] = vscale * *nvv;
1980:             if (PetscUnlikely(vshift != 0.0 && *njj == ir - rst)) vv[ist + ij] += vshift;
1981:             njj++;
1982:             nvv++;
1983:           }
1984:           ci[ir] += rsize;
1985:           nii++;
1986:         }
1987:       }
1988:     }
1989:   }
1990:   PetscCall(PetscFree(ci));

1992:   /* restore info */
1993:   for (i = 0; i < nest->nr; ++i) {
1994:     for (j = 0; j < nest->nc; ++j) {
1995:       Mat B = nest->m[i][j];
1996:       if (B) {
1997:         PetscInt nnr = 0, k = i * nest->nc + j;

1999:         B = (trans[k] ? trans[k] : B);
2000:         PetscCall(MatRestoreRowIJ(B, 0, PETSC_FALSE, PETSC_FALSE, &nnr, (const PetscInt **)&aii[k], (const PetscInt **)&ajj[k], &done));
2001:         PetscCheck(done, PetscObjectComm((PetscObject)B), PETSC_ERR_PLIB, "MatRestoreRowIJ");
2002:         PetscCall(MatSeqAIJRestoreArray(B, &avv[k]));
2003:         PetscCall(MatDestroy(&trans[k]));
2004:       }
2005:     }
2006:   }
2007:   PetscCall(PetscFree4(aii, ajj, avv, trans));

2009:   /* finalize newmat */
2010:   if (reuse == MAT_INITIAL_MATRIX) {
2011:     PetscCall(MatCreateSeqAIJWithArrays(PetscObjectComm((PetscObject)A), nr, nc, ii, jj, vv, newmat));
2012:   } else if (reuse == MAT_INPLACE_MATRIX) {
2013:     Mat B;

2015:     PetscCall(MatCreateSeqAIJWithArrays(PetscObjectComm((PetscObject)A), nr, nc, ii, jj, vv, &B));
2016:     PetscCall(MatHeaderReplace(A, &B));
2017:   }
2018:   PetscCall(MatAssemblyBegin(*newmat, MAT_FINAL_ASSEMBLY));
2019:   PetscCall(MatAssemblyEnd(*newmat, MAT_FINAL_ASSEMBLY));
2020:   {
2021:     Mat_SeqAIJ *a = (Mat_SeqAIJ *)((*newmat)->data);
2022:     a->free_a     = PETSC_TRUE;
2023:     a->free_ij    = PETSC_TRUE;
2024:   }
2025:   PetscFunctionReturn(PETSC_SUCCESS);
2026: }

2028: PETSC_INTERN PetscErrorCode MatAXPY_Dense_Nest(Mat Y, PetscScalar a, Mat X)
2029: {
2030:   Mat_Nest *nest = (Mat_Nest *)X->data;
2031:   PetscInt  i, j, k, rstart;
2032:   PetscBool flg;

2034:   PetscFunctionBegin;
2035:   /* Fill by row */
2036:   for (j = 0; j < nest->nc; ++j) {
2037:     /* Using global column indices and ISAllGather() is not scalable. */
2038:     IS              bNis;
2039:     PetscInt        bN;
2040:     const PetscInt *bNindices;
2041:     PetscCall(ISAllGather(nest->isglobal.col[j], &bNis));
2042:     PetscCall(ISGetSize(bNis, &bN));
2043:     PetscCall(ISGetIndices(bNis, &bNindices));
2044:     for (i = 0; i < nest->nr; ++i) {
2045:       Mat             B = nest->m[i][j], D = NULL;
2046:       PetscInt        bm, br;
2047:       const PetscInt *bmindices;
2048:       if (!B) continue;
2049:       PetscCall(PetscObjectTypeCompareAny((PetscObject)B, &flg, MATTRANSPOSEVIRTUAL, MATHERMITIANTRANSPOSEVIRTUAL, ""));
2050:       if (flg) {
2051:         PetscTryMethod(B, "MatTransposeGetMat_C", (Mat, Mat *), (B, &D));
2052:         PetscTryMethod(B, "MatHermitianTransposeGetMat_C", (Mat, Mat *), (B, &D));
2053:         PetscCall(MatConvert(B, ((PetscObject)D)->type_name, MAT_INITIAL_MATRIX, &D));
2054:         B = D;
2055:       }
2056:       PetscCall(PetscObjectTypeCompareAny((PetscObject)B, &flg, MATSEQSBAIJ, MATMPISBAIJ, ""));
2057:       if (flg) {
2058:         if (D) PetscCall(MatConvert(D, MATBAIJ, MAT_INPLACE_MATRIX, &D));
2059:         else PetscCall(MatConvert(B, MATBAIJ, MAT_INITIAL_MATRIX, &D));
2060:         B = D;
2061:       }
2062:       PetscCall(ISGetLocalSize(nest->isglobal.row[i], &bm));
2063:       PetscCall(ISGetIndices(nest->isglobal.row[i], &bmindices));
2064:       PetscCall(MatGetOwnershipRange(B, &rstart, NULL));
2065:       for (br = 0; br < bm; ++br) {
2066:         PetscInt           row = bmindices[br], brncols, *cols;
2067:         const PetscInt    *brcols;
2068:         const PetscScalar *brcoldata;
2069:         PetscScalar       *vals = NULL;
2070:         PetscCall(MatGetRow(B, br + rstart, &brncols, &brcols, &brcoldata));
2071:         PetscCall(PetscMalloc1(brncols, &cols));
2072:         for (k = 0; k < brncols; k++) cols[k] = bNindices[brcols[k]];
2073:         /*
2074:           Nest blocks are required to be nonoverlapping -- otherwise nest and monolithic index layouts wouldn't match.
2075:           Thus, we could use INSERT_VALUES, but I prefer ADD_VALUES.
2076:          */
2077:         if (a != 1.0) {
2078:           PetscCall(PetscMalloc1(brncols, &vals));
2079:           for (k = 0; k < brncols; k++) vals[k] = a * brcoldata[k];
2080:           PetscCall(MatSetValues(Y, 1, &row, brncols, cols, vals, ADD_VALUES));
2081:           PetscCall(PetscFree(vals));
2082:         } else {
2083:           PetscCall(MatSetValues(Y, 1, &row, brncols, cols, brcoldata, ADD_VALUES));
2084:         }
2085:         PetscCall(MatRestoreRow(B, br + rstart, &brncols, &brcols, &brcoldata));
2086:         PetscCall(PetscFree(cols));
2087:       }
2088:       PetscCall(MatDestroy(&D));
2089:       PetscCall(ISRestoreIndices(nest->isglobal.row[i], &bmindices));
2090:     }
2091:     PetscCall(ISRestoreIndices(bNis, &bNindices));
2092:     PetscCall(ISDestroy(&bNis));
2093:   }
2094:   PetscCall(MatAssemblyBegin(Y, MAT_FINAL_ASSEMBLY));
2095:   PetscCall(MatAssemblyEnd(Y, MAT_FINAL_ASSEMBLY));
2096:   PetscFunctionReturn(PETSC_SUCCESS);
2097: }

2099: static PetscErrorCode MatConvert_Nest_AIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
2100: {
2101:   Mat_Nest   *nest = (Mat_Nest *)A->data;
2102:   PetscInt    m, n, M, N, i, j, k, *dnnz, *onnz = NULL, rstart, cstart, cend;
2103:   PetscMPIInt size;
2104:   Mat         C;

2106:   PetscFunctionBegin;
2107:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size));
2108:   if (size == 1) { /* look for a special case with SeqAIJ matrices and strided-1, contiguous, blocks */
2109:     PetscInt  nf;
2110:     PetscBool fast;

2112:     PetscCall(PetscStrcmp(newtype, MATAIJ, &fast));
2113:     if (!fast) PetscCall(PetscStrcmp(newtype, MATSEQAIJ, &fast));
2114:     for (i = 0; i < nest->nr && fast; ++i) {
2115:       for (j = 0; j < nest->nc && fast; ++j) {
2116:         Mat B = nest->m[i][j];
2117:         if (B) {
2118:           PetscCall(PetscObjectTypeCompare((PetscObject)B, MATSEQAIJ, &fast));
2119:           if (!fast) {
2120:             PetscBool istrans;

2122:             PetscCall(PetscObjectTypeCompare((PetscObject)B, MATTRANSPOSEVIRTUAL, &istrans));
2123:             if (istrans) {
2124:               Mat Bt;

2126:               PetscCall(MatTransposeGetMat(B, &Bt));
2127:               PetscCall(PetscObjectTypeCompare((PetscObject)Bt, MATSEQAIJ, &fast));
2128:             } else {
2129:               PetscCall(PetscObjectTypeCompare((PetscObject)B, MATHERMITIANTRANSPOSEVIRTUAL, &istrans));
2130:               if (istrans) {
2131:                 Mat Bt;

2133:                 PetscCall(MatHermitianTransposeGetMat(B, &Bt));
2134:                 PetscCall(PetscObjectTypeCompare((PetscObject)Bt, MATSEQAIJ, &fast));
2135:               }
2136:             }
2137:             if (fast) fast = (PetscBool)(!((Mat_Shell *)B->data)->zrows && !((Mat_Shell *)B->data)->zcols && !((Mat_Shell *)B->data)->axpy && !((Mat_Shell *)B->data)->left && !((Mat_Shell *)B->data)->right && !((Mat_Shell *)B->data)->dshift);
2138:           }
2139:         }
2140:       }
2141:     }
2142:     for (i = 0, nf = 0; i < nest->nr && fast; ++i) {
2143:       PetscCall(PetscObjectTypeCompare((PetscObject)nest->isglobal.row[i], ISSTRIDE, &fast));
2144:       if (fast) {
2145:         PetscInt f, s;

2147:         PetscCall(ISStrideGetInfo(nest->isglobal.row[i], &f, &s));
2148:         if (f != nf || s != 1) {
2149:           fast = PETSC_FALSE;
2150:         } else {
2151:           PetscCall(ISGetSize(nest->isglobal.row[i], &f));
2152:           nf += f;
2153:         }
2154:       }
2155:     }
2156:     for (i = 0, nf = 0; i < nest->nc && fast; ++i) {
2157:       PetscCall(PetscObjectTypeCompare((PetscObject)nest->isglobal.col[i], ISSTRIDE, &fast));
2158:       if (fast) {
2159:         PetscInt f, s;

2161:         PetscCall(ISStrideGetInfo(nest->isglobal.col[i], &f, &s));
2162:         if (f != nf || s != 1) {
2163:           fast = PETSC_FALSE;
2164:         } else {
2165:           PetscCall(ISGetSize(nest->isglobal.col[i], &f));
2166:           nf += f;
2167:         }
2168:       }
2169:     }
2170:     if (fast) {
2171:       PetscCall(MatConvert_Nest_SeqAIJ_fast(A, newtype, reuse, newmat));
2172:       PetscFunctionReturn(PETSC_SUCCESS);
2173:     }
2174:   }
2175:   PetscCall(MatGetSize(A, &M, &N));
2176:   PetscCall(MatGetLocalSize(A, &m, &n));
2177:   PetscCall(MatGetOwnershipRangeColumn(A, &cstart, &cend));
2178:   if (reuse == MAT_REUSE_MATRIX) C = *newmat;
2179:   else {
2180:     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &C));
2181:     PetscCall(MatSetType(C, newtype));
2182:     PetscCall(MatSetSizes(C, m, n, M, N));
2183:   }
2184:   PetscCall(PetscMalloc1(2 * m, &dnnz));
2185:   if (m) {
2186:     onnz = dnnz + m;
2187:     for (k = 0; k < m; k++) {
2188:       dnnz[k] = 0;
2189:       onnz[k] = 0;
2190:     }
2191:   }
2192:   for (j = 0; j < nest->nc; ++j) {
2193:     IS              bNis;
2194:     PetscInt        bN;
2195:     const PetscInt *bNindices;
2196:     PetscBool       flg;
2197:     /* Using global column indices and ISAllGather() is not scalable. */
2198:     PetscCall(ISAllGather(nest->isglobal.col[j], &bNis));
2199:     PetscCall(ISGetSize(bNis, &bN));
2200:     PetscCall(ISGetIndices(bNis, &bNindices));
2201:     for (i = 0; i < nest->nr; ++i) {
2202:       PetscSF         bmsf;
2203:       PetscSFNode    *iremote;
2204:       Mat             B = nest->m[i][j], D = NULL;
2205:       PetscInt        bm, *sub_dnnz, *sub_onnz, br;
2206:       const PetscInt *bmindices;
2207:       if (!B) continue;
2208:       PetscCall(ISGetLocalSize(nest->isglobal.row[i], &bm));
2209:       PetscCall(ISGetIndices(nest->isglobal.row[i], &bmindices));
2210:       PetscCall(PetscSFCreate(PetscObjectComm((PetscObject)A), &bmsf));
2211:       PetscCall(PetscMalloc1(bm, &iremote));
2212:       PetscCall(PetscMalloc1(bm, &sub_dnnz));
2213:       PetscCall(PetscMalloc1(bm, &sub_onnz));
2214:       for (k = 0; k < bm; ++k) {
2215:         sub_dnnz[k] = 0;
2216:         sub_onnz[k] = 0;
2217:       }
2218:       PetscCall(PetscObjectTypeCompareAny((PetscObject)B, &flg, MATTRANSPOSEVIRTUAL, MATHERMITIANTRANSPOSEVIRTUAL, ""));
2219:       if (flg) {
2220:         PetscTryMethod(B, "MatTransposeGetMat_C", (Mat, Mat *), (B, &D));
2221:         PetscTryMethod(B, "MatHermitianTransposeGetMat_C", (Mat, Mat *), (B, &D));
2222:         PetscCall(MatConvert(B, ((PetscObject)D)->type_name, MAT_INITIAL_MATRIX, &D));
2223:         B = D;
2224:       }
2225:       PetscCall(PetscObjectTypeCompareAny((PetscObject)B, &flg, MATSEQSBAIJ, MATMPISBAIJ, ""));
2226:       if (flg) {
2227:         if (D) PetscCall(MatConvert(D, MATBAIJ, MAT_INPLACE_MATRIX, &D));
2228:         else PetscCall(MatConvert(B, MATBAIJ, MAT_INITIAL_MATRIX, &D));
2229:         B = D;
2230:       }
2231:       /*
2232:        Locate the owners for all of the locally-owned global row indices for this row block.
2233:        These determine the roots of PetscSF used to communicate preallocation data to row owners.
2234:        The roots correspond to the dnnz and onnz entries; thus, there are two roots per row.
2235:        */
2236:       PetscCall(MatGetOwnershipRange(B, &rstart, NULL));
2237:       for (br = 0; br < bm; ++br) {
2238:         PetscInt        row = bmindices[br], brncols, col;
2239:         const PetscInt *brcols;
2240:         PetscInt        rowrel   = 0; /* row's relative index on its owner rank */
2241:         PetscMPIInt     rowowner = 0;
2242:         PetscCall(PetscLayoutFindOwnerIndex(A->rmap, row, &rowowner, &rowrel));
2243:         /* how many roots  */
2244:         iremote[br].rank  = rowowner;
2245:         iremote[br].index = rowrel; /* edge from bmdnnz to dnnz */
2246:         /* get nonzero pattern */
2247:         PetscCall(MatGetRow(B, br + rstart, &brncols, &brcols, NULL));
2248:         for (k = 0; k < brncols; k++) {
2249:           col = bNindices[brcols[k]];
2250:           if (col >= A->cmap->range[rowowner] && col < A->cmap->range[rowowner + 1]) {
2251:             sub_dnnz[br]++;
2252:           } else {
2253:             sub_onnz[br]++;
2254:           }
2255:         }
2256:         PetscCall(MatRestoreRow(B, br + rstart, &brncols, &brcols, NULL));
2257:       }
2258:       PetscCall(MatDestroy(&D));
2259:       PetscCall(ISRestoreIndices(nest->isglobal.row[i], &bmindices));
2260:       /* bsf will have to take care of disposing of bedges. */
2261:       PetscCall(PetscSFSetGraph(bmsf, m, bm, NULL, PETSC_OWN_POINTER, iremote, PETSC_OWN_POINTER));
2262:       PetscCall(PetscSFReduceBegin(bmsf, MPIU_INT, sub_dnnz, dnnz, MPI_SUM));
2263:       PetscCall(PetscSFReduceEnd(bmsf, MPIU_INT, sub_dnnz, dnnz, MPI_SUM));
2264:       PetscCall(PetscSFReduceBegin(bmsf, MPIU_INT, sub_onnz, onnz, MPI_SUM));
2265:       PetscCall(PetscSFReduceEnd(bmsf, MPIU_INT, sub_onnz, onnz, MPI_SUM));
2266:       PetscCall(PetscFree(sub_dnnz));
2267:       PetscCall(PetscFree(sub_onnz));
2268:       PetscCall(PetscSFDestroy(&bmsf));
2269:     }
2270:     PetscCall(ISRestoreIndices(bNis, &bNindices));
2271:     PetscCall(ISDestroy(&bNis));
2272:   }
2273:   /* Resize preallocation if overestimated */
2274:   for (i = 0; i < m; i++) {
2275:     dnnz[i] = PetscMin(dnnz[i], A->cmap->n);
2276:     onnz[i] = PetscMin(onnz[i], A->cmap->N - A->cmap->n);
2277:   }
2278:   PetscCall(MatSeqAIJSetPreallocation(C, 0, dnnz));
2279:   PetscCall(MatMPIAIJSetPreallocation(C, 0, dnnz, 0, onnz));
2280:   PetscCall(PetscFree(dnnz));
2281:   PetscCall(MatAXPY_Dense_Nest(C, 1.0, A));
2282:   if (reuse == MAT_INPLACE_MATRIX) PetscCall(MatHeaderReplace(A, &C));
2283:   else *newmat = C;
2284:   PetscFunctionReturn(PETSC_SUCCESS);
2285: }

2287: static PetscErrorCode MatConvert_Nest_Dense(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
2288: {
2289:   Mat      B;
2290:   PetscInt m, n, M, N;

2292:   PetscFunctionBegin;
2293:   PetscCall(MatGetSize(A, &M, &N));
2294:   PetscCall(MatGetLocalSize(A, &m, &n));
2295:   if (reuse == MAT_REUSE_MATRIX) {
2296:     B = *newmat;
2297:     PetscCall(MatZeroEntries(B));
2298:   } else {
2299:     PetscCall(MatCreateDense(PetscObjectComm((PetscObject)A), m, PETSC_DECIDE, M, N, NULL, &B));
2300:   }
2301:   PetscCall(MatAXPY_Dense_Nest(B, 1.0, A));
2302:   if (reuse == MAT_INPLACE_MATRIX) PetscCall(MatHeaderReplace(A, &B));
2303:   else if (reuse == MAT_INITIAL_MATRIX) *newmat = B;
2304:   PetscFunctionReturn(PETSC_SUCCESS);
2305: }

2307: static PetscErrorCode MatHasOperation_Nest(Mat mat, MatOperation op, PetscBool *has)
2308: {
2309:   Mat_Nest *bA  = (Mat_Nest *)mat->data;
2310:   PetscBool flg = PETSC_TRUE;

2312:   PetscFunctionBegin;
2313:   *has = PETSC_FALSE;
2314:   if (op == MATOP_MULT || op == MATOP_MULT_ADD || op == MATOP_MULT_TRANSPOSE || op == MATOP_MULT_TRANSPOSE_ADD || op == MATOP_MULT_HERMITIAN_TRANSPOSE || op == MATOP_MULT_HERMITIAN_TRANS_ADD) {
2315:     MatOperation opAdd;

2317:     if (op == MATOP_MULT || op == MATOP_MULT_ADD) opAdd = MATOP_MULT_ADD;
2318:     else if (op == MATOP_MULT_TRANSPOSE || op == MATOP_MULT_TRANSPOSE_ADD) opAdd = MATOP_MULT_TRANSPOSE_ADD;
2319:     else opAdd = MATOP_MULT_HERMITIAN_TRANS_ADD;
2320:     for (PetscInt j = 0; j < bA->nc && flg; j++) {
2321:       for (PetscInt i = 0; i < bA->nr; i++) {
2322:         if (!bA->m[i][j]) continue;
2323:         PetscCall(MatHasOperation(bA->m[i][j], opAdd, &flg));
2324:         if (!flg) break;
2325:       }
2326:     }
2327:   }
2328:   if (flg && ((void **)mat->ops)[op]) *has = PETSC_TRUE;
2329:   PetscFunctionReturn(PETSC_SUCCESS);
2330: }

2332: /*MC
2333:   MATNEST -  "nest" - Matrix type consisting of nested submatrices, each stored separately.

2335:   Level: intermediate

2337:   Notes:
2338:   This matrix type permits scalable use of `PCFIELDSPLIT` and avoids the large memory costs of extracting submatrices.
2339:   It allows the use of symmetric and block formats for parts of multi-physics simulations.
2340:   It is usually used with `DMCOMPOSITE` and `DMCreateMatrix()`

2342:   Each of the submatrices lives on the same MPI communicator as the original nest matrix (though they can have zero
2343:   rows/columns on some processes.) Thus this is not meant for cases where the submatrices live on far fewer processes
2344:   than the nest matrix.

2346: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatCreate()`, `MatType`, `MatCreateNest()`, `MatNestSetSubMat()`, `MatNestGetSubMat()`,
2347:           `VecCreateNest()`, `DMCreateMatrix()`, `DMCOMPOSITE`, `MatNestSetVecType()`, `MatNestGetLocalISs()`,
2348:           `MatNestGetISs()`, `MatNestSetSubMats()`, `MatNestGetSubMats()`
2349: M*/
2350: PETSC_EXTERN PetscErrorCode MatCreate_Nest(Mat A)
2351: {
2352:   Mat_Nest *s;

2354:   PetscFunctionBegin;
2355:   PetscCall(PetscNew(&s));
2356:   A->data = (void *)s;

2358:   s->nr            = -1;
2359:   s->nc            = -1;
2360:   s->m             = NULL;
2361:   s->splitassembly = PETSC_FALSE;

2363:   PetscCall(PetscMemzero(A->ops, sizeof(*A->ops)));

2365:   A->ops->mult                      = MatMult_Nest;
2366:   A->ops->multadd                   = MatMultAdd_Nest;
2367:   A->ops->multtranspose             = MatMultTranspose_Nest;
2368:   A->ops->multtransposeadd          = MatMultTransposeAdd_Nest;
2369:   A->ops->transpose                 = MatTranspose_Nest;
2370:   A->ops->multhermitiantranspose    = MatMultHermitianTranspose_Nest;
2371:   A->ops->multhermitiantransposeadd = MatMultHermitianTransposeAdd_Nest;
2372:   A->ops->assemblybegin             = MatAssemblyBegin_Nest;
2373:   A->ops->assemblyend               = MatAssemblyEnd_Nest;
2374:   A->ops->zeroentries               = MatZeroEntries_Nest;
2375:   A->ops->copy                      = MatCopy_Nest;
2376:   A->ops->axpy                      = MatAXPY_Nest;
2377:   A->ops->duplicate                 = MatDuplicate_Nest;
2378:   A->ops->createsubmatrix           = MatCreateSubMatrix_Nest;
2379:   A->ops->destroy                   = MatDestroy_Nest;
2380:   A->ops->view                      = MatView_Nest;
2381:   A->ops->getvecs                   = NULL; /* Use VECNEST by calling MatNestSetVecType(A,VECNEST) */
2382:   A->ops->getlocalsubmatrix         = MatGetLocalSubMatrix_Nest;
2383:   A->ops->restorelocalsubmatrix     = MatRestoreLocalSubMatrix_Nest;
2384:   A->ops->getdiagonal               = MatGetDiagonal_Nest;
2385:   A->ops->diagonalscale             = MatDiagonalScale_Nest;
2386:   A->ops->scale                     = MatScale_Nest;
2387:   A->ops->shift                     = MatShift_Nest;
2388:   A->ops->diagonalset               = MatDiagonalSet_Nest;
2389:   A->ops->setrandom                 = MatSetRandom_Nest;
2390:   A->ops->hasoperation              = MatHasOperation_Nest;

2392:   A->spptr     = NULL;
2393:   A->assembled = PETSC_FALSE;

2395:   /* expose Nest api's */
2396:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetSubMat_C", MatNestGetSubMat_Nest));
2397:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestSetSubMat_C", MatNestSetSubMat_Nest));
2398:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetSubMats_C", MatNestGetSubMats_Nest));
2399:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetSize_C", MatNestGetSize_Nest));
2400:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetISs_C", MatNestGetISs_Nest));
2401:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetLocalISs_C", MatNestGetLocalISs_Nest));
2402:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestSetVecType_C", MatNestSetVecType_Nest));
2403:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestSetSubMats_C", MatNestSetSubMats_Nest));
2404:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_mpiaij_C", MatConvert_Nest_AIJ));
2405:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_seqaij_C", MatConvert_Nest_AIJ));
2406:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_aij_C", MatConvert_Nest_AIJ));
2407:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_is_C", MatConvert_Nest_IS));
2408:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_mpidense_C", MatConvert_Nest_Dense));
2409:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_seqdense_C", MatConvert_Nest_Dense));
2410:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_nest_seqdense_C", MatProductSetFromOptions_Nest_Dense));
2411:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_nest_mpidense_C", MatProductSetFromOptions_Nest_Dense));

2413:   PetscCall(PetscObjectChangeTypeName((PetscObject)A, MATNEST));
2414:   PetscFunctionReturn(PETSC_SUCCESS);
2415: }