xref: /petsc/src/mat/impls/nest/matnest.c (revision 4e8208cbcbc709572b8abe32f33c78b69c819375) !
1 #include <../src/mat/impls/nest/matnestimpl.h> /*I   "petscmat.h"   I*/
2 #include <../src/mat/impls/aij/seq/aij.h>
3 #include <../src/mat/impls/shell/shell.h>
4 #include <petscsf.h>
5 
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);
9 
10 PETSC_INTERN PetscErrorCode MatConvert_Nest_IS(Mat, MatType, MatReuse, Mat *);
11 
12 /* private functions */
MatNestGetSizes_Private(Mat A,PetscInt * m,PetscInt * n,PetscInt * M,PetscInt * N)13 static PetscErrorCode MatNestGetSizes_Private(Mat A, PetscInt *m, PetscInt *n, PetscInt *M, PetscInt *N)
14 {
15   Mat_Nest *bA = (Mat_Nest *)A->data;
16   PetscInt  i, j;
17 
18   PetscFunctionBegin;
19   *m = *n = *M = *N = 0;
20   for (i = 0; i < bA->nr; i++) { /* rows */
21     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 (j = 0; j < bA->nc; j++) { /* cols */
28     PetscInt sn, sN;
29     PetscCall(ISGetLocalSize(bA->isglobal.col[j], &sn));
30     PetscCall(ISGetSize(bA->isglobal.col[j], &sN));
31     *n += sn;
32     *N += sN;
33   }
34   PetscFunctionReturn(PETSC_SUCCESS);
35 }
36 
37 /* operations */
MatMult_Nest(Mat A,Vec x,Vec y)38 static PetscErrorCode MatMult_Nest(Mat A, Vec x, Vec y)
39 {
40   Mat_Nest *bA = (Mat_Nest *)A->data;
41   Vec      *bx = bA->right, *by = bA->left;
42   PetscInt  i, j, nr = bA->nr, nc = bA->nc;
43 
44   PetscFunctionBegin;
45   for (i = 0; i < nr; i++) PetscCall(VecGetSubVector(y, bA->isglobal.row[i], &by[i]));
46   for (i = 0; i < nc; i++) PetscCall(VecGetSubVector(x, bA->isglobal.col[i], &bx[i]));
47   for (i = 0; i < nr; i++) {
48     PetscCall(VecZeroEntries(by[i]));
49     for (j = 0; j < 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 (i = 0; i < nr; i++) PetscCall(VecRestoreSubVector(y, bA->isglobal.row[i], &by[i]));
56   for (i = 0; i < nc; i++) PetscCall(VecRestoreSubVector(x, bA->isglobal.col[i], &bx[i]));
57   PetscFunctionReturn(PETSC_SUCCESS);
58 }
59 
MatMultAdd_Nest(Mat A,Vec x,Vec y,Vec z)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;
64   PetscInt  i, j, nr = bA->nr, nc = bA->nc;
65 
66   PetscFunctionBegin;
67   for (i = 0; i < nr; i++) PetscCall(VecGetSubVector(z, bA->isglobal.row[i], &bz[i]));
68   for (i = 0; i < nc; i++) PetscCall(VecGetSubVector(x, bA->isglobal.col[i], &bx[i]));
69   for (i = 0; i < nr; i++) {
70     if (y != z) {
71       Vec by;
72       PetscCall(VecGetSubVector(y, bA->isglobal.row[i], &by));
73       PetscCall(VecCopy(by, bz[i]));
74       PetscCall(VecRestoreSubVector(y, bA->isglobal.row[i], &by));
75     }
76     for (j = 0; j < nc; j++) {
77       if (!bA->m[i][j]) continue;
78       /* y[i] <- y[i] + A[i][j] * x[j] */
79       PetscCall(MatMultAdd(bA->m[i][j], bx[j], bz[i], bz[i]));
80     }
81   }
82   for (i = 0; i < nr; i++) PetscCall(VecRestoreSubVector(z, bA->isglobal.row[i], &bz[i]));
83   for (i = 0; i < nc; i++) PetscCall(VecRestoreSubVector(x, bA->isglobal.col[i], &bx[i]));
84   PetscFunctionReturn(PETSC_SUCCESS);
85 }
86 
87 typedef struct {
88   Mat         *workC;      /* array of Mat with specific containers depending on the underlying MatMatMult implementation */
89   PetscScalar *tarray;     /* buffer for storing all temporary products A[i][j] B[j] */
90   PetscInt    *dm, *dn, k; /* displacements and number of submatrices */
91 } Nest_Dense;
92 
MatProductNumeric_Nest_Dense(Mat C)93 static PetscErrorCode MatProductNumeric_Nest_Dense(Mat C)
94 {
95   Mat_Nest          *bA;
96   Nest_Dense        *contents;
97   Mat                viewB, viewC, productB, workC;
98   const PetscScalar *barray;
99   PetscScalar       *carray;
100   PetscInt           i, j, M, N, nr, nc, ldb, ldc;
101   Mat                A, B;
102 
103   PetscFunctionBegin;
104   MatCheckProduct(C, 1);
105   A = C->product->A;
106   B = C->product->B;
107   PetscCall(MatGetSize(B, NULL, &N));
108   if (!N) {
109     PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
110     PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
111     PetscFunctionReturn(PETSC_SUCCESS);
112   }
113   contents = (Nest_Dense *)C->product->data;
114   PetscCheck(contents, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
115   bA = (Mat_Nest *)A->data;
116   nr = bA->nr;
117   nc = bA->nc;
118   PetscCall(MatDenseGetLDA(B, &ldb));
119   PetscCall(MatDenseGetLDA(C, &ldc));
120   PetscCall(MatZeroEntries(C));
121   PetscCall(MatDenseGetArrayRead(B, &barray));
122   PetscCall(MatDenseGetArray(C, &carray));
123   for (i = 0; i < nr; i++) {
124     PetscCall(ISGetSize(bA->isglobal.row[i], &M));
125     PetscCall(MatCreateDense(PetscObjectComm((PetscObject)A), contents->dm[i + 1] - contents->dm[i], PETSC_DECIDE, M, N, PetscSafePointerPlusOffset(carray, contents->dm[i]), &viewC));
126     PetscCall(MatDenseSetLDA(viewC, ldc));
127     for (j = 0; j < nc; j++) {
128       if (!bA->m[i][j]) continue;
129       PetscCall(ISGetSize(bA->isglobal.col[j], &M));
130       PetscCall(MatCreateDense(PetscObjectComm((PetscObject)A), contents->dn[j + 1] - contents->dn[j], PETSC_DECIDE, M, N, PetscSafePointerPlusOffset((PetscScalar *)barray, contents->dn[j]), &viewB));
131       PetscCall(MatDenseSetLDA(viewB, ldb));
132 
133       /* MatMatMultNumeric(bA->m[i][j],viewB,contents->workC[i*nc + j]); */
134       workC             = contents->workC[i * nc + j];
135       productB          = workC->product->B;
136       workC->product->B = viewB; /* use newly created dense matrix viewB */
137       PetscCall(MatProductNumeric(workC));
138       PetscCall(MatDestroy(&viewB));
139       workC->product->B = productB; /* resume original B */
140 
141       /* C[i] <- workC + C[i] */
142       PetscCall(MatAXPY(viewC, 1.0, contents->workC[i * nc + j], SAME_NONZERO_PATTERN));
143     }
144     PetscCall(MatDestroy(&viewC));
145   }
146   PetscCall(MatDenseRestoreArray(C, &carray));
147   PetscCall(MatDenseRestoreArrayRead(B, &barray));
148 
149   PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
150   PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
151   PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
152   PetscFunctionReturn(PETSC_SUCCESS);
153 }
154 
MatNest_DenseDestroy(PetscCtxRt ctx)155 static PetscErrorCode MatNest_DenseDestroy(PetscCtxRt ctx)
156 {
157   Nest_Dense *contents = *(Nest_Dense **)ctx;
158   PetscInt    i;
159 
160   PetscFunctionBegin;
161   PetscCall(PetscFree(contents->tarray));
162   for (i = 0; i < contents->k; i++) PetscCall(MatDestroy(contents->workC + i));
163   PetscCall(PetscFree3(contents->dm, contents->dn, contents->workC));
164   PetscCall(PetscFree(contents));
165   PetscFunctionReturn(PETSC_SUCCESS);
166 }
167 
MatProductSymbolic_Nest_Dense(Mat C)168 static PetscErrorCode MatProductSymbolic_Nest_Dense(Mat C)
169 {
170   Mat_Nest          *bA;
171   Mat                viewB, workC;
172   const PetscScalar *barray;
173   PetscInt           i, j, M, N, m, n, nr, nc, maxm = 0, ldb;
174   Nest_Dense        *contents = NULL;
175   PetscBool          cisdense;
176   Mat                A, B;
177   PetscReal          fill;
178 
179   PetscFunctionBegin;
180   MatCheckProduct(C, 1);
181   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
182   A    = C->product->A;
183   B    = C->product->B;
184   fill = C->product->fill;
185   bA   = (Mat_Nest *)A->data;
186   nr   = bA->nr;
187   nc   = bA->nc;
188   PetscCall(MatGetLocalSize(C, &m, &n));
189   PetscCall(MatGetSize(C, &M, &N));
190   if (m == PETSC_DECIDE || n == PETSC_DECIDE || M == PETSC_DECIDE || N == PETSC_DECIDE) {
191     PetscCall(MatGetLocalSize(B, NULL, &n));
192     PetscCall(MatGetSize(B, NULL, &N));
193     PetscCall(MatGetLocalSize(A, &m, NULL));
194     PetscCall(MatGetSize(A, &M, NULL));
195     PetscCall(MatSetSizes(C, m, n, M, N));
196   }
197   PetscCall(PetscObjectTypeCompareAny((PetscObject)C, &cisdense, MATSEQDENSE, MATMPIDENSE, MATSEQDENSECUDA, MATMPIDENSECUDA, ""));
198   if (!cisdense) PetscCall(MatSetType(C, ((PetscObject)B)->type_name));
199   PetscCall(MatSetUp(C));
200   if (!N) {
201     C->ops->productnumeric = MatProductNumeric_Nest_Dense;
202     PetscFunctionReturn(PETSC_SUCCESS);
203   }
204 
205   PetscCall(PetscNew(&contents));
206   C->product->data    = contents;
207   C->product->destroy = MatNest_DenseDestroy;
208   PetscCall(PetscCalloc3(nr + 1, &contents->dm, nc + 1, &contents->dn, nr * nc, &contents->workC));
209   contents->k = nr * nc;
210   for (i = 0; i < nr; i++) {
211     PetscCall(ISGetLocalSize(bA->isglobal.row[i], contents->dm + i + 1));
212     maxm = PetscMax(maxm, contents->dm[i + 1]);
213     contents->dm[i + 1] += contents->dm[i];
214   }
215   for (i = 0; i < nc; i++) {
216     PetscCall(ISGetLocalSize(bA->isglobal.col[i], contents->dn + i + 1));
217     contents->dn[i + 1] += contents->dn[i];
218   }
219   PetscCall(PetscMalloc1(maxm * N, &contents->tarray));
220   PetscCall(MatDenseGetLDA(B, &ldb));
221   PetscCall(MatGetSize(B, NULL, &N));
222   PetscCall(MatDenseGetArrayRead(B, &barray));
223   /* loops are permuted compared to MatMatMultNumeric so that viewB is created only once per column of A */
224   for (j = 0; j < nc; j++) {
225     PetscCall(ISGetSize(bA->isglobal.col[j], &M));
226     PetscCall(MatCreateDense(PetscObjectComm((PetscObject)A), contents->dn[j + 1] - contents->dn[j], PETSC_DECIDE, M, N, PetscSafePointerPlusOffset((PetscScalar *)barray, contents->dn[j]), &viewB));
227     PetscCall(MatDenseSetLDA(viewB, ldb));
228     for (i = 0; i < nr; i++) {
229       if (!bA->m[i][j]) continue;
230       /* MatMatMultSymbolic may attach a specific container (depending on MatType of bA->m[i][j]) to workC[i][j] */
231 
232       PetscCall(MatProductCreate(bA->m[i][j], viewB, NULL, &contents->workC[i * nc + j]));
233       workC = contents->workC[i * nc + j];
234       PetscCall(MatProductSetType(workC, MATPRODUCT_AB));
235       PetscCall(MatProductSetAlgorithm(workC, "default"));
236       PetscCall(MatProductSetFill(workC, fill));
237       PetscCall(MatProductSetFromOptions(workC));
238       PetscCall(MatProductSymbolic(workC));
239 
240       /* since tarray will be shared by all Mat */
241       PetscCall(MatSeqDenseSetPreallocation(workC, contents->tarray));
242       PetscCall(MatMPIDenseSetPreallocation(workC, contents->tarray));
243     }
244     PetscCall(MatDestroy(&viewB));
245   }
246   PetscCall(MatDenseRestoreArrayRead(B, &barray));
247 
248   C->ops->productnumeric = MatProductNumeric_Nest_Dense;
249   PetscFunctionReturn(PETSC_SUCCESS);
250 }
251 
MatProductSetFromOptions_Nest_Dense(Mat C)252 static PetscErrorCode MatProductSetFromOptions_Nest_Dense(Mat C)
253 {
254   Mat_Product *product = C->product;
255 
256   PetscFunctionBegin;
257   if (product->type == MATPRODUCT_AB) C->ops->productsymbolic = MatProductSymbolic_Nest_Dense;
258   PetscFunctionReturn(PETSC_SUCCESS);
259 }
260 
MatMultTransposeKernel_Nest(Mat A,Vec x,Vec y,PetscBool herm)261 static PetscErrorCode MatMultTransposeKernel_Nest(Mat A, Vec x, Vec y, PetscBool herm)
262 {
263   Mat_Nest *bA = (Mat_Nest *)A->data;
264   Vec      *bx = bA->left, *by = bA->right;
265   PetscInt  i, j, nr = bA->nr, nc = bA->nc;
266 
267   PetscFunctionBegin;
268   for (i = 0; i < nr; i++) PetscCall(VecGetSubVector(x, bA->isglobal.row[i], &bx[i]));
269   for (i = 0; i < nc; i++) PetscCall(VecGetSubVector(y, bA->isglobal.col[i], &by[i]));
270   for (j = 0; j < nc; j++) {
271     PetscCall(VecZeroEntries(by[j]));
272     for (i = 0; i < nr; i++) {
273       if (!bA->m[i][j]) continue;
274       if (herm) PetscCall(MatMultHermitianTransposeAdd(bA->m[i][j], bx[i], by[j], by[j])); /* y[j] <- y[j] + (A[i][j])^H * x[i] */
275       else PetscCall(MatMultTransposeAdd(bA->m[i][j], bx[i], by[j], by[j]));               /* y[j] <- y[j] + (A[i][j])^T * x[i] */
276     }
277   }
278   for (i = 0; i < nr; i++) PetscCall(VecRestoreSubVector(x, bA->isglobal.row[i], &bx[i]));
279   for (i = 0; i < nc; i++) PetscCall(VecRestoreSubVector(y, bA->isglobal.col[i], &by[i]));
280   PetscFunctionReturn(PETSC_SUCCESS);
281 }
282 
MatMultTranspose_Nest(Mat A,Vec x,Vec y)283 static PetscErrorCode MatMultTranspose_Nest(Mat A, Vec x, Vec y)
284 {
285   PetscFunctionBegin;
286   PetscCall(MatMultTransposeKernel_Nest(A, x, y, PETSC_FALSE));
287   PetscFunctionReturn(PETSC_SUCCESS);
288 }
289 
MatMultHermitianTranspose_Nest(Mat A,Vec x,Vec y)290 static PetscErrorCode MatMultHermitianTranspose_Nest(Mat A, Vec x, Vec y)
291 {
292   PetscFunctionBegin;
293   PetscCall(MatMultTransposeKernel_Nest(A, x, y, PETSC_TRUE));
294   PetscFunctionReturn(PETSC_SUCCESS);
295 }
296 
MatMultTransposeAddKernel_Nest(Mat A,Vec x,Vec y,Vec z,PetscBool herm)297 static PetscErrorCode MatMultTransposeAddKernel_Nest(Mat A, Vec x, Vec y, Vec z, PetscBool herm)
298 {
299   Mat_Nest *bA = (Mat_Nest *)A->data;
300   Vec      *bx = bA->left, *bz = bA->right;
301   PetscInt  i, j, nr = bA->nr, nc = bA->nc;
302 
303   PetscFunctionBegin;
304   for (i = 0; i < nr; i++) PetscCall(VecGetSubVector(x, bA->isglobal.row[i], &bx[i]));
305   for (i = 0; i < nc; i++) PetscCall(VecGetSubVector(z, bA->isglobal.col[i], &bz[i]));
306   for (j = 0; j < nc; j++) {
307     if (y != z) {
308       Vec by;
309       PetscCall(VecGetSubVector(y, bA->isglobal.col[j], &by));
310       PetscCall(VecCopy(by, bz[j]));
311       PetscCall(VecRestoreSubVector(y, bA->isglobal.col[j], &by));
312     }
313     for (i = 0; i < nr; i++) {
314       if (!bA->m[i][j]) continue;
315       if (herm) PetscCall(MatMultHermitianTransposeAdd(bA->m[i][j], bx[i], bz[j], bz[j])); /* z[j] <- y[j] + (A[i][j])^H * x[i] */
316       else PetscCall(MatMultTransposeAdd(bA->m[i][j], bx[i], bz[j], bz[j]));               /* z[j] <- y[j] + (A[i][j])^T * x[i] */
317     }
318   }
319   for (i = 0; i < nr; i++) PetscCall(VecRestoreSubVector(x, bA->isglobal.row[i], &bx[i]));
320   for (i = 0; i < nc; i++) PetscCall(VecRestoreSubVector(z, bA->isglobal.col[i], &bz[i]));
321   PetscFunctionReturn(PETSC_SUCCESS);
322 }
323 
MatMultTransposeAdd_Nest(Mat A,Vec x,Vec y,Vec z)324 static PetscErrorCode MatMultTransposeAdd_Nest(Mat A, Vec x, Vec y, Vec z)
325 {
326   PetscFunctionBegin;
327   PetscCall(MatMultTransposeAddKernel_Nest(A, x, y, z, PETSC_FALSE));
328   PetscFunctionReturn(PETSC_SUCCESS);
329 }
330 
MatMultHermitianTransposeAdd_Nest(Mat A,Vec x,Vec y,Vec z)331 static PetscErrorCode MatMultHermitianTransposeAdd_Nest(Mat A, Vec x, Vec y, Vec z)
332 {
333   PetscFunctionBegin;
334   PetscCall(MatMultTransposeAddKernel_Nest(A, x, y, z, PETSC_TRUE));
335   PetscFunctionReturn(PETSC_SUCCESS);
336 }
337 
MatTranspose_Nest(Mat A,MatReuse reuse,Mat * B)338 static PetscErrorCode MatTranspose_Nest(Mat A, MatReuse reuse, Mat *B)
339 {
340   Mat_Nest *bA = (Mat_Nest *)A->data, *bC;
341   Mat       C;
342   PetscInt  i, j, nr = bA->nr, nc = bA->nc;
343 
344   PetscFunctionBegin;
345   if (reuse == MAT_REUSE_MATRIX) PetscCall(MatTransposeCheckNonzeroState_Private(A, *B));
346   PetscCheck(reuse != MAT_INPLACE_MATRIX || nr == nc, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_SIZ, "Square nested matrix only for in-place");
347 
348   if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_INPLACE_MATRIX) {
349     Mat *subs;
350     IS  *is_row, *is_col;
351 
352     PetscCall(PetscCalloc1(nr * nc, &subs));
353     PetscCall(PetscMalloc2(nr, &is_row, nc, &is_col));
354     PetscCall(MatNestGetISs(A, is_row, is_col));
355     if (reuse == MAT_INPLACE_MATRIX) {
356       for (i = 0; i < nr; i++) {
357         for (j = 0; j < nc; j++) subs[i + nr * j] = bA->m[i][j];
358       }
359     }
360 
361     PetscCall(MatCreateNest(PetscObjectComm((PetscObject)A), nc, is_col, nr, is_row, subs, &C));
362     PetscCall(PetscFree(subs));
363     PetscCall(PetscFree2(is_row, is_col));
364   } else {
365     C = *B;
366   }
367 
368   bC = (Mat_Nest *)C->data;
369   for (i = 0; i < nr; i++) {
370     for (j = 0; j < nc; j++) {
371       if (bA->m[i][j]) {
372         PetscCall(MatTranspose(bA->m[i][j], reuse, &bC->m[j][i]));
373       } else {
374         bC->m[j][i] = NULL;
375       }
376     }
377   }
378 
379   if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_REUSE_MATRIX) {
380     *B = C;
381   } else {
382     PetscCall(MatHeaderMerge(A, &C));
383   }
384   PetscFunctionReturn(PETSC_SUCCESS);
385 }
386 
MatNestDestroyISList(PetscInt n,IS ** list)387 static PetscErrorCode MatNestDestroyISList(PetscInt n, IS **list)
388 {
389   IS      *lst = *list;
390   PetscInt i;
391 
392   PetscFunctionBegin;
393   if (!lst) PetscFunctionReturn(PETSC_SUCCESS);
394   for (i = 0; i < n; i++)
395     if (lst[i]) PetscCall(ISDestroy(&lst[i]));
396   PetscCall(PetscFree(lst));
397   *list = NULL;
398   PetscFunctionReturn(PETSC_SUCCESS);
399 }
400 
MatReset_Nest(Mat A)401 static PetscErrorCode MatReset_Nest(Mat A)
402 {
403   Mat_Nest *vs = (Mat_Nest *)A->data;
404   PetscInt  i, j;
405 
406   PetscFunctionBegin;
407   /* release the matrices and the place holders */
408   PetscCall(MatNestDestroyISList(vs->nr, &vs->isglobal.row));
409   PetscCall(MatNestDestroyISList(vs->nc, &vs->isglobal.col));
410   PetscCall(MatNestDestroyISList(vs->nr, &vs->islocal.row));
411   PetscCall(MatNestDestroyISList(vs->nc, &vs->islocal.col));
412 
413   PetscCall(PetscFree(vs->row_len));
414   PetscCall(PetscFree(vs->col_len));
415   PetscCall(PetscFree(vs->nnzstate));
416 
417   PetscCall(PetscFree2(vs->left, vs->right));
418 
419   /* release the matrices and the place holders */
420   if (vs->m) {
421     for (i = 0; i < vs->nr; i++) {
422       for (j = 0; j < vs->nc; j++) PetscCall(MatDestroy(&vs->m[i][j]));
423     }
424     PetscCall(PetscFree(vs->m[0]));
425     PetscCall(PetscFree(vs->m));
426   }
427 
428   /* restore defaults */
429   vs->nr            = 0;
430   vs->nc            = 0;
431   vs->splitassembly = PETSC_FALSE;
432   PetscFunctionReturn(PETSC_SUCCESS);
433 }
434 
MatDestroy_Nest(Mat A)435 static PetscErrorCode MatDestroy_Nest(Mat A)
436 {
437   PetscFunctionBegin;
438   PetscCall(MatReset_Nest(A));
439   PetscCall(PetscFree(A->data));
440   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetSubMat_C", NULL));
441   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestSetSubMat_C", NULL));
442   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetSubMats_C", NULL));
443   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetSize_C", NULL));
444   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetISs_C", NULL));
445   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetLocalISs_C", NULL));
446   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestSetVecType_C", NULL));
447   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestSetSubMats_C", NULL));
448   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_mpiaij_C", NULL));
449   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_seqaij_C", NULL));
450   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_aij_C", NULL));
451   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_is_C", NULL));
452   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_mpidense_C", NULL));
453   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_seqdense_C", NULL));
454   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_nest_seqdense_C", NULL));
455   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_nest_mpidense_C", NULL));
456   PetscFunctionReturn(PETSC_SUCCESS);
457 }
458 
MatAssemblyBegin_Nest(Mat A,MatAssemblyType type)459 static PetscErrorCode MatAssemblyBegin_Nest(Mat A, MatAssemblyType type)
460 {
461   Mat_Nest *vs = (Mat_Nest *)A->data;
462   PetscInt  i, j;
463   PetscBool nnzstate = PETSC_FALSE;
464 
465   PetscFunctionBegin;
466   for (i = 0; i < vs->nr; i++) {
467     for (j = 0; j < vs->nc; j++) {
468       PetscObjectState subnnzstate = 0;
469       if (vs->m[i][j]) {
470         PetscCall(MatAssemblyBegin(vs->m[i][j], type));
471         if (!vs->splitassembly) {
472           /* Note: split assembly will fail if the same block appears more than once (even indirectly through a nested
473            * sub-block). This could be fixed by adding a flag to Mat so that there was a way to check if a Mat was
474            * already performing an assembly, but the result would by more complicated and appears to offer less
475            * potential for diagnostics and correctness checking. Split assembly should be fixed once there is an
476            * interface for libraries to make asynchronous progress in "user-defined non-blocking collectives".
477            */
478           PetscCall(MatAssemblyEnd(vs->m[i][j], type));
479           PetscCall(MatGetNonzeroState(vs->m[i][j], &subnnzstate));
480         }
481       }
482       nnzstate                     = (PetscBool)(nnzstate || vs->nnzstate[i * vs->nc + j] != subnnzstate);
483       vs->nnzstate[i * vs->nc + j] = subnnzstate;
484     }
485   }
486   if (nnzstate) A->nonzerostate++;
487   PetscFunctionReturn(PETSC_SUCCESS);
488 }
489 
MatAssemblyEnd_Nest(Mat A,MatAssemblyType type)490 static PetscErrorCode MatAssemblyEnd_Nest(Mat A, MatAssemblyType type)
491 {
492   Mat_Nest *vs = (Mat_Nest *)A->data;
493   PetscInt  i, j;
494 
495   PetscFunctionBegin;
496   for (i = 0; i < vs->nr; i++) {
497     for (j = 0; j < vs->nc; j++) {
498       if (vs->m[i][j]) {
499         if (vs->splitassembly) PetscCall(MatAssemblyEnd(vs->m[i][j], type));
500       }
501     }
502   }
503   PetscFunctionReturn(PETSC_SUCCESS);
504 }
505 
MatNestFindNonzeroSubMatRow(Mat A,PetscInt row,Mat * B)506 static PetscErrorCode MatNestFindNonzeroSubMatRow(Mat A, PetscInt row, Mat *B)
507 {
508   Mat_Nest *vs = (Mat_Nest *)A->data;
509   PetscInt  j;
510   Mat       sub;
511 
512   PetscFunctionBegin;
513   sub = (row < vs->nc) ? vs->m[row][row] : (Mat)NULL; /* Prefer to find on the diagonal */
514   for (j = 0; !sub && j < vs->nc; j++) sub = vs->m[row][j];
515   if (sub) PetscCall(MatSetUp(sub)); /* Ensure that the sizes are available */
516   *B = sub;
517   PetscFunctionReturn(PETSC_SUCCESS);
518 }
519 
MatNestFindNonzeroSubMatCol(Mat A,PetscInt col,Mat * B)520 static PetscErrorCode MatNestFindNonzeroSubMatCol(Mat A, PetscInt col, Mat *B)
521 {
522   Mat_Nest *vs = (Mat_Nest *)A->data;
523   PetscInt  i;
524   Mat       sub;
525 
526   PetscFunctionBegin;
527   sub = (col < vs->nr) ? vs->m[col][col] : (Mat)NULL; /* Prefer to find on the diagonal */
528   for (i = 0; !sub && i < vs->nr; i++) sub = vs->m[i][col];
529   if (sub) PetscCall(MatSetUp(sub)); /* Ensure that the sizes are available */
530   *B = sub;
531   PetscFunctionReturn(PETSC_SUCCESS);
532 }
533 
MatNestFindISRange(Mat A,PetscInt n,const IS list[],IS is,PetscInt * begin,PetscInt * end)534 static PetscErrorCode MatNestFindISRange(Mat A, PetscInt n, const IS list[], IS is, PetscInt *begin, PetscInt *end)
535 {
536   PetscInt  i, j, size, m;
537   PetscBool flg;
538   IS        out, concatenate[2];
539 
540   PetscFunctionBegin;
541   PetscAssertPointer(list, 3);
542   PetscValidHeaderSpecific(is, IS_CLASSID, 4);
543   if (begin) {
544     PetscAssertPointer(begin, 5);
545     *begin = -1;
546   }
547   if (end) {
548     PetscAssertPointer(end, 6);
549     *end = -1;
550   }
551   for (i = 0; i < n; i++) {
552     if (!list[i]) continue;
553     PetscCall(ISEqualUnsorted(list[i], is, &flg));
554     if (flg) {
555       if (begin) *begin = i;
556       if (end) *end = i + 1;
557       PetscFunctionReturn(PETSC_SUCCESS);
558     }
559   }
560   PetscCall(ISGetSize(is, &size));
561   for (i = 0; i < n - 1; i++) {
562     if (!list[i]) continue;
563     m = 0;
564     PetscCall(ISConcatenate(PetscObjectComm((PetscObject)A), 2, list + i, &out));
565     PetscCall(ISGetSize(out, &m));
566     for (j = i + 2; j < n && m < size; j++) {
567       if (list[j]) {
568         concatenate[0] = out;
569         concatenate[1] = list[j];
570         PetscCall(ISConcatenate(PetscObjectComm((PetscObject)A), 2, concatenate, &out));
571         PetscCall(ISDestroy(concatenate));
572         PetscCall(ISGetSize(out, &m));
573       }
574     }
575     if (m == size) {
576       PetscCall(ISEqualUnsorted(out, is, &flg));
577       if (flg) {
578         if (begin) *begin = i;
579         if (end) *end = j;
580         PetscCall(ISDestroy(&out));
581         PetscFunctionReturn(PETSC_SUCCESS);
582       }
583     }
584     PetscCall(ISDestroy(&out));
585   }
586   PetscFunctionReturn(PETSC_SUCCESS);
587 }
588 
MatNestFillEmptyMat_Private(Mat A,PetscInt i,PetscInt j,Mat * B)589 static PetscErrorCode MatNestFillEmptyMat_Private(Mat A, PetscInt i, PetscInt j, Mat *B)
590 {
591   Mat_Nest *vs = (Mat_Nest *)A->data;
592   PetscInt  lr, lc;
593 
594   PetscFunctionBegin;
595   PetscCall(MatCreate(PetscObjectComm((PetscObject)A), B));
596   PetscCall(ISGetLocalSize(vs->isglobal.row[i], &lr));
597   PetscCall(ISGetLocalSize(vs->isglobal.col[j], &lc));
598   PetscCall(MatSetSizes(*B, lr, lc, PETSC_DECIDE, PETSC_DECIDE));
599   PetscCall(MatSetType(*B, MATAIJ));
600   PetscCall(MatSeqAIJSetPreallocation(*B, 0, NULL));
601   PetscCall(MatMPIAIJSetPreallocation(*B, 0, NULL, 0, NULL));
602   PetscCall(MatSetUp(*B));
603   PetscCall(MatSetOption(*B, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
604   PetscCall(MatAssemblyBegin(*B, MAT_FINAL_ASSEMBLY));
605   PetscCall(MatAssemblyEnd(*B, MAT_FINAL_ASSEMBLY));
606   PetscFunctionReturn(PETSC_SUCCESS);
607 }
608 
MatNestGetBlock_Private(Mat A,PetscInt rbegin,PetscInt rend,PetscInt cbegin,PetscInt cend,Mat * B)609 static PetscErrorCode MatNestGetBlock_Private(Mat A, PetscInt rbegin, PetscInt rend, PetscInt cbegin, PetscInt cend, Mat *B)
610 {
611   Mat_Nest  *vs = (Mat_Nest *)A->data;
612   Mat       *a;
613   PetscInt   i, j, k, l, nr = rend - rbegin, nc = cend - cbegin;
614   char       keyname[256];
615   PetscBool *b;
616   PetscBool  flg;
617 
618   PetscFunctionBegin;
619   *B = NULL;
620   PetscCall(PetscSNPrintf(keyname, sizeof(keyname), "NestBlock_%" PetscInt_FMT "-%" PetscInt_FMT "x%" PetscInt_FMT "-%" PetscInt_FMT, rbegin, rend, cbegin, cend));
621   PetscCall(PetscObjectQuery((PetscObject)A, keyname, (PetscObject *)B));
622   if (*B) PetscFunctionReturn(PETSC_SUCCESS);
623 
624   PetscCall(PetscMalloc2(nr * nc, &a, nr * nc, &b));
625   for (i = 0; i < nr; i++) {
626     for (j = 0; j < nc; j++) {
627       a[i * nc + j] = vs->m[rbegin + i][cbegin + j];
628       b[i * nc + j] = PETSC_FALSE;
629     }
630   }
631   if (nc != vs->nc && nr != vs->nr) {
632     for (i = 0; i < nr; i++) {
633       for (j = 0; j < nc; j++) {
634         flg = PETSC_FALSE;
635         for (k = 0; (k < nr && !flg); k++) {
636           if (a[j + k * nc]) flg = PETSC_TRUE;
637         }
638         if (flg) {
639           flg = PETSC_FALSE;
640           for (l = 0; (l < nc && !flg); l++) {
641             if (a[i * nc + l]) flg = PETSC_TRUE;
642           }
643         }
644         if (!flg) {
645           b[i * nc + j] = PETSC_TRUE;
646           PetscCall(MatNestFillEmptyMat_Private(A, rbegin + i, cbegin + j, a + i * nc + j));
647         }
648       }
649     }
650   }
651   PetscCall(MatCreateNest(PetscObjectComm((PetscObject)A), nr, nr != vs->nr ? NULL : vs->isglobal.row, nc, nc != vs->nc ? NULL : vs->isglobal.col, a, B));
652   for (i = 0; i < nr; i++) {
653     for (j = 0; j < nc; j++) {
654       if (b[i * nc + j]) PetscCall(MatDestroy(a + i * nc + j));
655     }
656   }
657   PetscCall(PetscFree2(a, b));
658   (*B)->assembled = A->assembled;
659   PetscCall(PetscObjectCompose((PetscObject)A, keyname, (PetscObject)*B));
660   PetscCall(PetscObjectDereference((PetscObject)*B)); /* Leave the only remaining reference in the composition */
661   PetscFunctionReturn(PETSC_SUCCESS);
662 }
663 
MatNestFindSubMat(Mat A,struct MatNestISPair * is,IS isrow,IS iscol,Mat * B)664 static PetscErrorCode MatNestFindSubMat(Mat A, struct MatNestISPair *is, IS isrow, IS iscol, Mat *B)
665 {
666   Mat_Nest *vs = (Mat_Nest *)A->data;
667   PetscInt  rbegin, rend, cbegin, cend;
668 
669   PetscFunctionBegin;
670   PetscCall(MatNestFindISRange(A, vs->nr, is->row, isrow, &rbegin, &rend));
671   PetscCall(MatNestFindISRange(A, vs->nc, is->col, iscol, &cbegin, &cend));
672   if (rend == rbegin + 1 && cend == cbegin + 1) {
673     if (!vs->m[rbegin][cbegin]) PetscCall(MatNestFillEmptyMat_Private(A, rbegin, cbegin, vs->m[rbegin] + cbegin));
674     *B = vs->m[rbegin][cbegin];
675   } else if (rbegin != -1 && cbegin != -1) {
676     PetscCall(MatNestGetBlock_Private(A, rbegin, rend, cbegin, cend, B));
677   } else SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_INCOMP, "Could not find index set");
678   PetscFunctionReturn(PETSC_SUCCESS);
679 }
680 
681 /*
682    TODO: This does not actually returns a submatrix we can modify
683 */
MatCreateSubMatrix_Nest(Mat A,IS isrow,IS iscol,MatReuse reuse,Mat * B)684 static PetscErrorCode MatCreateSubMatrix_Nest(Mat A, IS isrow, IS iscol, MatReuse reuse, Mat *B)
685 {
686   Mat_Nest *vs = (Mat_Nest *)A->data;
687   Mat       sub;
688 
689   PetscFunctionBegin;
690   PetscCall(MatNestFindSubMat(A, &vs->isglobal, isrow, iscol, &sub));
691   switch (reuse) {
692   case MAT_INITIAL_MATRIX:
693     PetscCall(PetscObjectReference((PetscObject)sub));
694     if (sub) PetscCall(PetscObjectStateIncrease((PetscObject)sub));
695     *B = sub;
696     break;
697   case MAT_REUSE_MATRIX:
698     PetscCheck(sub == *B, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Submatrix was not used before in this call");
699     if (sub) PetscCall(PetscObjectStateIncrease((PetscObject)sub));
700     break;
701   default:
702     break;
703   }
704   PetscFunctionReturn(PETSC_SUCCESS);
705 }
706 
MatGetLocalSubMatrix_Nest(Mat A,IS isrow,IS iscol,Mat * B)707 static PetscErrorCode MatGetLocalSubMatrix_Nest(Mat A, IS isrow, IS iscol, Mat *B)
708 {
709   Mat_Nest *vs = (Mat_Nest *)A->data;
710   Mat       sub;
711 
712   PetscFunctionBegin;
713   PetscCall(MatNestFindSubMat(A, &vs->islocal, isrow, iscol, &sub));
714   /* We allow the submatrix to be NULL, perhaps it would be better for the user to return an empty matrix instead */
715   if (sub) PetscCall(PetscObjectReference((PetscObject)sub));
716   *B = sub;
717   PetscFunctionReturn(PETSC_SUCCESS);
718 }
719 
MatRestoreLocalSubMatrix_Nest(Mat A,IS isrow,IS iscol,Mat * B)720 static PetscErrorCode MatRestoreLocalSubMatrix_Nest(Mat A, IS isrow, IS iscol, Mat *B)
721 {
722   Mat_Nest *vs = (Mat_Nest *)A->data;
723   Mat       sub;
724 
725   PetscFunctionBegin;
726   PetscCall(MatNestFindSubMat(A, &vs->islocal, isrow, iscol, &sub));
727   PetscCheck(*B == sub, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Local submatrix has not been gotten");
728   if (sub) {
729     PetscCheck(((PetscObject)sub)->refct > 1, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Local submatrix has had reference count decremented too many times");
730     PetscCall(MatDestroy(B));
731   }
732   PetscFunctionReturn(PETSC_SUCCESS);
733 }
734 
MatGetDiagonal_Nest(Mat A,Vec v)735 static PetscErrorCode MatGetDiagonal_Nest(Mat A, Vec v)
736 {
737   Mat_Nest *bA = (Mat_Nest *)A->data;
738   PetscInt  i;
739 
740   PetscFunctionBegin;
741   for (i = 0; i < bA->nr; i++) {
742     Vec bv;
743     PetscCall(VecGetSubVector(v, bA->isglobal.row[i], &bv));
744     if (bA->m[i][i]) {
745       PetscCall(MatGetDiagonal(bA->m[i][i], bv));
746     } else {
747       PetscCall(VecSet(bv, 0.0));
748     }
749     PetscCall(VecRestoreSubVector(v, bA->isglobal.row[i], &bv));
750   }
751   PetscFunctionReturn(PETSC_SUCCESS);
752 }
753 
MatDiagonalScale_Nest(Mat A,Vec l,Vec r)754 static PetscErrorCode MatDiagonalScale_Nest(Mat A, Vec l, Vec r)
755 {
756   Mat_Nest *bA = (Mat_Nest *)A->data;
757   Vec       bl, *br;
758   PetscInt  i, j;
759 
760   PetscFunctionBegin;
761   PetscCall(PetscCalloc1(bA->nc, &br));
762   if (r) {
763     for (j = 0; j < bA->nc; j++) PetscCall(VecGetSubVector(r, bA->isglobal.col[j], &br[j]));
764   }
765   bl = NULL;
766   for (i = 0; i < bA->nr; i++) {
767     if (l) PetscCall(VecGetSubVector(l, bA->isglobal.row[i], &bl));
768     for (j = 0; j < bA->nc; j++) {
769       if (bA->m[i][j]) PetscCall(MatDiagonalScale(bA->m[i][j], bl, br[j]));
770     }
771     if (l) PetscCall(VecRestoreSubVector(l, bA->isglobal.row[i], &bl));
772   }
773   if (r) {
774     for (j = 0; j < bA->nc; j++) PetscCall(VecRestoreSubVector(r, bA->isglobal.col[j], &br[j]));
775   }
776   PetscCall(PetscFree(br));
777   PetscFunctionReturn(PETSC_SUCCESS);
778 }
779 
MatScale_Nest(Mat A,PetscScalar a)780 static PetscErrorCode MatScale_Nest(Mat A, PetscScalar a)
781 {
782   Mat_Nest *bA = (Mat_Nest *)A->data;
783   PetscInt  i, j;
784 
785   PetscFunctionBegin;
786   for (i = 0; i < bA->nr; i++) {
787     for (j = 0; j < bA->nc; j++) {
788       if (bA->m[i][j]) PetscCall(MatScale(bA->m[i][j], a));
789     }
790   }
791   PetscFunctionReturn(PETSC_SUCCESS);
792 }
793 
MatShift_Nest(Mat A,PetscScalar a)794 static PetscErrorCode MatShift_Nest(Mat A, PetscScalar a)
795 {
796   Mat_Nest *bA = (Mat_Nest *)A->data;
797   PetscInt  i;
798   PetscBool nnzstate = PETSC_FALSE;
799 
800   PetscFunctionBegin;
801   for (i = 0; i < bA->nr; i++) {
802     PetscObjectState subnnzstate = 0;
803     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);
804     PetscCall(MatShift(bA->m[i][i], a));
805     PetscCall(MatGetNonzeroState(bA->m[i][i], &subnnzstate));
806     nnzstate                     = (PetscBool)(nnzstate || bA->nnzstate[i * bA->nc + i] != subnnzstate);
807     bA->nnzstate[i * bA->nc + i] = subnnzstate;
808   }
809   if (nnzstate) A->nonzerostate++;
810   PetscFunctionReturn(PETSC_SUCCESS);
811 }
812 
MatDiagonalSet_Nest(Mat A,Vec D,InsertMode is)813 static PetscErrorCode MatDiagonalSet_Nest(Mat A, Vec D, InsertMode is)
814 {
815   Mat_Nest *bA = (Mat_Nest *)A->data;
816   PetscInt  i;
817   PetscBool nnzstate = PETSC_FALSE;
818 
819   PetscFunctionBegin;
820   for (i = 0; i < bA->nr; i++) {
821     PetscObjectState subnnzstate = 0;
822     Vec              bv;
823     PetscCall(VecGetSubVector(D, bA->isglobal.row[i], &bv));
824     if (bA->m[i][i]) {
825       PetscCall(MatDiagonalSet(bA->m[i][i], bv, is));
826       PetscCall(MatGetNonzeroState(bA->m[i][i], &subnnzstate));
827     }
828     PetscCall(VecRestoreSubVector(D, bA->isglobal.row[i], &bv));
829     nnzstate                     = (PetscBool)(nnzstate || bA->nnzstate[i * bA->nc + i] != subnnzstate);
830     bA->nnzstate[i * bA->nc + i] = subnnzstate;
831   }
832   if (nnzstate) A->nonzerostate++;
833   PetscFunctionReturn(PETSC_SUCCESS);
834 }
835 
MatSetRandom_Nest(Mat A,PetscRandom rctx)836 static PetscErrorCode MatSetRandom_Nest(Mat A, PetscRandom rctx)
837 {
838   Mat_Nest *bA = (Mat_Nest *)A->data;
839   PetscInt  i, j;
840 
841   PetscFunctionBegin;
842   for (i = 0; i < bA->nr; i++) {
843     for (j = 0; j < bA->nc; j++) {
844       if (bA->m[i][j]) PetscCall(MatSetRandom(bA->m[i][j], rctx));
845     }
846   }
847   PetscFunctionReturn(PETSC_SUCCESS);
848 }
849 
MatCreateVecs_Nest(Mat A,Vec * right,Vec * left)850 static PetscErrorCode MatCreateVecs_Nest(Mat A, Vec *right, Vec *left)
851 {
852   Mat_Nest *bA = (Mat_Nest *)A->data;
853   Vec      *L, *R;
854   MPI_Comm  comm;
855   PetscInt  i, j;
856 
857   PetscFunctionBegin;
858   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
859   if (right) {
860     /* allocate R */
861     PetscCall(PetscMalloc1(bA->nc, &R));
862     /* Create the right vectors */
863     for (j = 0; j < bA->nc; j++) {
864       for (i = 0; i < bA->nr; i++) {
865         if (bA->m[i][j]) {
866           PetscCall(MatCreateVecs(bA->m[i][j], &R[j], NULL));
867           break;
868         }
869       }
870       PetscCheck(i != bA->nr, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Mat(Nest) contains a null column.");
871     }
872     PetscCall(VecCreateNest(comm, bA->nc, bA->isglobal.col, R, right));
873     /* hand back control to the nest vector */
874     for (j = 0; j < bA->nc; j++) PetscCall(VecDestroy(&R[j]));
875     PetscCall(PetscFree(R));
876   }
877 
878   if (left) {
879     /* allocate L */
880     PetscCall(PetscMalloc1(bA->nr, &L));
881     /* Create the left vectors */
882     for (i = 0; i < bA->nr; i++) {
883       for (j = 0; j < bA->nc; j++) {
884         if (bA->m[i][j]) {
885           PetscCall(MatCreateVecs(bA->m[i][j], NULL, &L[i]));
886           break;
887         }
888       }
889       PetscCheck(j != bA->nc, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Mat(Nest) contains a null row.");
890     }
891 
892     PetscCall(VecCreateNest(comm, bA->nr, bA->isglobal.row, L, left));
893     for (i = 0; i < bA->nr; i++) PetscCall(VecDestroy(&L[i]));
894 
895     PetscCall(PetscFree(L));
896   }
897   PetscFunctionReturn(PETSC_SUCCESS);
898 }
899 
MatView_Nest(Mat A,PetscViewer viewer)900 static PetscErrorCode MatView_Nest(Mat A, PetscViewer viewer)
901 {
902   Mat_Nest *bA = (Mat_Nest *)A->data;
903   PetscBool isascii, viewSub = PETSC_FALSE;
904   PetscInt  i, j;
905 
906   PetscFunctionBegin;
907   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
908   if (isascii) {
909     PetscViewerFormat format;
910 
911     PetscCall(PetscViewerGetFormat(viewer, &format));
912     if (format == PETSC_VIEWER_ASCII_MATLAB) {
913       Mat T;
914 
915       PetscCall(MatConvert(A, MATAIJ, MAT_INITIAL_MATRIX, &T));
916       PetscCall(MatView(T, viewer));
917       PetscCall(MatDestroy(&T));
918       PetscFunctionReturn(PETSC_SUCCESS);
919     }
920     PetscCall(PetscOptionsGetBool(((PetscObject)A)->options, ((PetscObject)A)->prefix, "-mat_view_nest_sub", &viewSub, NULL));
921     PetscCall(PetscViewerASCIIPushTab(viewer));
922     PetscCall(PetscViewerASCIIPrintf(viewer, "MatNest, rows=%" PetscInt_FMT ", cols=%" PetscInt_FMT ", structure:\n", bA->nr, bA->nc));
923     for (i = 0; i < bA->nr; i++) {
924       for (j = 0; j < bA->nc; j++) {
925         MatType   type;
926         char      name[256] = "", prefix[256] = "";
927         PetscInt  NR, NC;
928         PetscBool isNest = PETSC_FALSE;
929 
930         if (!bA->m[i][j]) {
931           PetscCall(PetscViewerASCIIPrintf(viewer, "(%" PetscInt_FMT ",%" PetscInt_FMT ") : NULL\n", i, j));
932           continue;
933         }
934         PetscCall(MatGetSize(bA->m[i][j], &NR, &NC));
935         PetscCall(MatGetType(bA->m[i][j], &type));
936         if (((PetscObject)bA->m[i][j])->name) PetscCall(PetscSNPrintf(name, sizeof(name), "name=\"%s\", ", ((PetscObject)bA->m[i][j])->name));
937         if (((PetscObject)bA->m[i][j])->prefix) PetscCall(PetscSNPrintf(prefix, sizeof(prefix), "prefix=\"%s\", ", ((PetscObject)bA->m[i][j])->prefix));
938         PetscCall(PetscObjectTypeCompare((PetscObject)bA->m[i][j], MATNEST, &isNest));
939 
940         PetscCall(PetscViewerASCIIPrintf(viewer, "(%" PetscInt_FMT ",%" PetscInt_FMT ") : %s%stype=%s, rows=%" PetscInt_FMT ", cols=%" PetscInt_FMT "\n", i, j, name, prefix, type, NR, NC));
941 
942         if (isNest || viewSub) {
943           PetscCall(PetscViewerASCIIPushTab(viewer)); /* push1 */
944           PetscCall(MatView(bA->m[i][j], viewer));
945           PetscCall(PetscViewerASCIIPopTab(viewer)); /* pop1 */
946         }
947       }
948     }
949     PetscCall(PetscViewerASCIIPopTab(viewer)); /* pop0 */
950   }
951   PetscFunctionReturn(PETSC_SUCCESS);
952 }
953 
MatZeroEntries_Nest(Mat A)954 static PetscErrorCode MatZeroEntries_Nest(Mat A)
955 {
956   Mat_Nest *bA = (Mat_Nest *)A->data;
957   PetscInt  i, j;
958 
959   PetscFunctionBegin;
960   for (i = 0; i < bA->nr; i++) {
961     for (j = 0; j < bA->nc; j++) {
962       if (!bA->m[i][j]) continue;
963       PetscCall(MatZeroEntries(bA->m[i][j]));
964     }
965   }
966   PetscFunctionReturn(PETSC_SUCCESS);
967 }
968 
MatCopy_Nest(Mat A,Mat B,MatStructure str)969 static PetscErrorCode MatCopy_Nest(Mat A, Mat B, MatStructure str)
970 {
971   Mat_Nest *bA = (Mat_Nest *)A->data, *bB = (Mat_Nest *)B->data;
972   PetscInt  i, j, nr = bA->nr, nc = bA->nc;
973   PetscBool nnzstate = PETSC_FALSE;
974 
975   PetscFunctionBegin;
976   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);
977   for (i = 0; i < nr; i++) {
978     for (j = 0; j < nc; j++) {
979       PetscObjectState subnnzstate = 0;
980       if (bA->m[i][j] && bB->m[i][j]) {
981         PetscCall(MatCopy(bA->m[i][j], bB->m[i][j], str));
982         PetscCall(MatGetNonzeroState(bB->m[i][j], &subnnzstate));
983         nnzstate                 = (PetscBool)(nnzstate || bB->nnzstate[i * nc + j] != subnnzstate);
984         bB->nnzstate[i * nc + j] = subnnzstate;
985       } else if (bA->m[i][j]) { // bB->m[i][j] is NULL
986         Mat M;
987 
988         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);
989         PetscCall(MatDuplicate(bA->m[i][j], MAT_COPY_VALUES, &M));
990         PetscCall(MatNestSetSubMat(B, i, j, M));
991         PetscCall(MatDestroy(&M));
992       } else if (bB->m[i][j]) { // bA->m[i][j] is NULL
993         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);
994         PetscCall(MatNestSetSubMat(B, i, j, NULL));
995       }
996     }
997   }
998   if (nnzstate) B->nonzerostate++;
999   PetscFunctionReturn(PETSC_SUCCESS);
1000 }
1001 
MatAXPY_Nest(Mat Y,PetscScalar a,Mat X,MatStructure str)1002 static PetscErrorCode MatAXPY_Nest(Mat Y, PetscScalar a, Mat X, MatStructure str)
1003 {
1004   Mat_Nest *bY = (Mat_Nest *)Y->data, *bX = (Mat_Nest *)X->data;
1005   PetscInt  i, j, nr = bY->nr, nc = bY->nc;
1006   PetscBool nnzstate = PETSC_FALSE;
1007 
1008   PetscFunctionBegin;
1009   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);
1010   for (i = 0; i < nr; i++) {
1011     for (j = 0; j < nc; j++) {
1012       PetscObjectState subnnzstate = 0;
1013       if (bY->m[i][j] && bX->m[i][j]) {
1014         PetscCall(MatAXPY(bY->m[i][j], a, bX->m[i][j], str));
1015       } else if (bX->m[i][j]) {
1016         Mat M;
1017 
1018         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);
1019         PetscCall(MatDuplicate(bX->m[i][j], MAT_COPY_VALUES, &M));
1020         PetscCall(MatNestSetSubMat(Y, i, j, M));
1021         PetscCall(MatDestroy(&M));
1022       }
1023       if (bY->m[i][j]) PetscCall(MatGetNonzeroState(bY->m[i][j], &subnnzstate));
1024       nnzstate                 = (PetscBool)(nnzstate || bY->nnzstate[i * nc + j] != subnnzstate);
1025       bY->nnzstate[i * nc + j] = subnnzstate;
1026     }
1027   }
1028   if (nnzstate) Y->nonzerostate++;
1029   PetscFunctionReturn(PETSC_SUCCESS);
1030 }
1031 
MatDuplicate_Nest(Mat A,MatDuplicateOption op,Mat * B)1032 static PetscErrorCode MatDuplicate_Nest(Mat A, MatDuplicateOption op, Mat *B)
1033 {
1034   Mat_Nest *bA = (Mat_Nest *)A->data;
1035   Mat      *b;
1036   PetscInt  i, j, nr = bA->nr, nc = bA->nc;
1037 
1038   PetscFunctionBegin;
1039   PetscCall(PetscMalloc1(nr * nc, &b));
1040   for (i = 0; i < nr; i++) {
1041     for (j = 0; j < nc; j++) {
1042       if (bA->m[i][j]) {
1043         PetscCall(MatDuplicate(bA->m[i][j], op, &b[i * nc + j]));
1044       } else {
1045         b[i * nc + j] = NULL;
1046       }
1047     }
1048   }
1049   PetscCall(MatCreateNest(PetscObjectComm((PetscObject)A), nr, bA->isglobal.row, nc, bA->isglobal.col, b, B));
1050   /* Give the new MatNest exclusive ownership */
1051   for (i = 0; i < nr * nc; i++) PetscCall(MatDestroy(&b[i]));
1052   PetscCall(PetscFree(b));
1053 
1054   PetscCall(MatAssemblyBegin(*B, MAT_FINAL_ASSEMBLY));
1055   PetscCall(MatAssemblyEnd(*B, MAT_FINAL_ASSEMBLY));
1056   PetscFunctionReturn(PETSC_SUCCESS);
1057 }
1058 
1059 /* nest api */
MatNestGetSubMat_Nest(Mat A,PetscInt idxm,PetscInt jdxm,Mat * mat)1060 static PetscErrorCode MatNestGetSubMat_Nest(Mat A, PetscInt idxm, PetscInt jdxm, Mat *mat)
1061 {
1062   Mat_Nest *bA = (Mat_Nest *)A->data;
1063 
1064   PetscFunctionBegin;
1065   PetscCheck(idxm < bA->nr, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Row too large: row %" PetscInt_FMT " max %" PetscInt_FMT, idxm, bA->nr - 1);
1066   PetscCheck(jdxm < bA->nc, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Col too large: row %" PetscInt_FMT " max %" PetscInt_FMT, jdxm, bA->nc - 1);
1067   *mat = bA->m[idxm][jdxm];
1068   PetscFunctionReturn(PETSC_SUCCESS);
1069 }
1070 
1071 /*@
1072   MatNestGetSubMat - Returns a single, sub-matrix from a `MATNEST`
1073 
1074   Not Collective
1075 
1076   Input Parameters:
1077 + A    - `MATNEST` matrix
1078 . idxm - index of the matrix within the nest matrix
1079 - jdxm - index of the matrix within the nest matrix
1080 
1081   Output Parameter:
1082 . sub - matrix at index `idxm`, `jdxm` within the nest matrix
1083 
1084   Level: developer
1085 
1086 .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestGetSize()`, `MatNestGetSubMats()`, `MatCreateNest()`, `MatNestSetSubMat()`,
1087           `MatNestGetLocalISs()`, `MatNestGetISs()`
1088 @*/
MatNestGetSubMat(Mat A,PetscInt idxm,PetscInt jdxm,Mat * sub)1089 PetscErrorCode MatNestGetSubMat(Mat A, PetscInt idxm, PetscInt jdxm, Mat *sub)
1090 {
1091   PetscFunctionBegin;
1092   PetscValidHeaderSpecific(A, MAT_CLASSID, 1);
1093   PetscValidLogicalCollectiveInt(A, idxm, 2);
1094   PetscValidLogicalCollectiveInt(A, jdxm, 3);
1095   PetscAssertPointer(sub, 4);
1096   PetscUseMethod(A, "MatNestGetSubMat_C", (Mat, PetscInt, PetscInt, Mat *), (A, idxm, jdxm, sub));
1097   PetscFunctionReturn(PETSC_SUCCESS);
1098 }
1099 
MatNestSetSubMat_Nest(Mat A,PetscInt idxm,PetscInt jdxm,Mat mat)1100 static PetscErrorCode MatNestSetSubMat_Nest(Mat A, PetscInt idxm, PetscInt jdxm, Mat mat)
1101 {
1102   Mat_Nest *bA = (Mat_Nest *)A->data;
1103   PetscInt  m, n, M, N, mi, ni, Mi, Ni;
1104 
1105   PetscFunctionBegin;
1106   PetscCheck(idxm < bA->nr, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Row too large: row %" PetscInt_FMT " max %" PetscInt_FMT, idxm, bA->nr - 1);
1107   PetscCheck(jdxm < bA->nc, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Col too large: row %" PetscInt_FMT " max %" PetscInt_FMT, jdxm, bA->nc - 1);
1108   if (mat) {
1109     PetscCall(MatGetLocalSize(mat, &m, &n));
1110     PetscCall(MatGetSize(mat, &M, &N));
1111     PetscCall(ISGetLocalSize(bA->isglobal.row[idxm], &mi));
1112     PetscCall(ISGetSize(bA->isglobal.row[idxm], &Mi));
1113     PetscCall(ISGetLocalSize(bA->isglobal.col[jdxm], &ni));
1114     PetscCall(ISGetSize(bA->isglobal.col[jdxm], &Ni));
1115     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);
1116     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);
1117   }
1118 
1119   /* do not increase object state */
1120   if (mat == bA->m[idxm][jdxm]) PetscFunctionReturn(PETSC_SUCCESS);
1121 
1122   PetscCall(PetscObjectReference((PetscObject)mat));
1123   PetscCall(MatDestroy(&bA->m[idxm][jdxm]));
1124   bA->m[idxm][jdxm] = mat;
1125   PetscCall(PetscObjectStateIncrease((PetscObject)A));
1126   if (mat) PetscCall(MatGetNonzeroState(mat, &bA->nnzstate[idxm * bA->nc + jdxm]));
1127   else bA->nnzstate[idxm * bA->nc + jdxm] = 0;
1128   A->nonzerostate++;
1129   PetscFunctionReturn(PETSC_SUCCESS);
1130 }
1131 
1132 /*@
1133   MatNestSetSubMat - Set a single submatrix in the `MATNEST`
1134 
1135   Logically Collective
1136 
1137   Input Parameters:
1138 + A    - `MATNEST` matrix
1139 . idxm - index of the matrix within the nest matrix
1140 . jdxm - index of the matrix within the nest matrix
1141 - sub  - matrix at index `idxm`, `jdxm` within the nest matrix
1142 
1143   Level: developer
1144 
1145   Notes:
1146   The new submatrix must have the same size and communicator as that block of the nest.
1147 
1148   This increments the reference count of the submatrix.
1149 
1150 .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestSetSubMats()`, `MatNestGetSubMats()`, `MatNestGetLocalISs()`, `MatCreateNest()`,
1151           `MatNestGetSubMat()`, `MatNestGetISs()`, `MatNestGetSize()`
1152 @*/
MatNestSetSubMat(Mat A,PetscInt idxm,PetscInt jdxm,Mat sub)1153 PetscErrorCode MatNestSetSubMat(Mat A, PetscInt idxm, PetscInt jdxm, Mat sub)
1154 {
1155   PetscFunctionBegin;
1156   PetscValidHeaderSpecific(A, MAT_CLASSID, 1);
1157   PetscValidLogicalCollectiveInt(A, idxm, 2);
1158   PetscValidLogicalCollectiveInt(A, jdxm, 3);
1159   if (sub) PetscValidHeaderSpecific(sub, MAT_CLASSID, 4);
1160   PetscTryMethod(A, "MatNestSetSubMat_C", (Mat, PetscInt, PetscInt, Mat), (A, idxm, jdxm, sub));
1161   PetscFunctionReturn(PETSC_SUCCESS);
1162 }
1163 
MatNestGetSubMats_Nest(Mat A,PetscInt * M,PetscInt * N,Mat *** mat)1164 static PetscErrorCode MatNestGetSubMats_Nest(Mat A, PetscInt *M, PetscInt *N, Mat ***mat)
1165 {
1166   Mat_Nest *bA = (Mat_Nest *)A->data;
1167 
1168   PetscFunctionBegin;
1169   if (M) *M = bA->nr;
1170   if (N) *N = bA->nc;
1171   if (mat) *mat = bA->m;
1172   PetscFunctionReturn(PETSC_SUCCESS);
1173 }
1174 
1175 /*@C
1176   MatNestGetSubMats - Returns the entire two dimensional array of matrices defining a `MATNEST` matrix.
1177 
1178   Not Collective
1179 
1180   Input Parameter:
1181 . A - nest matrix
1182 
1183   Output Parameters:
1184 + M   - number of submatrix rows in the nest matrix
1185 . N   - number of submatrix columns in the nest matrix
1186 - mat - array of matrices
1187 
1188   Level: developer
1189 
1190   Note:
1191   The user should not free the array `mat`.
1192 
1193   Fortran Notes:
1194   This routine has a calling sequence `call MatNestGetSubMats(A, M, N, mat, ierr)`
1195   where the space allocated for the optional argument `mat` is assumed large enough (if provided).
1196   Matrices in `mat` are returned in row-major order, see `MatCreateNest()` for an example.
1197 
1198 .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestGetSize()`, `MatNestGetSubMat()`, `MatNestGetLocalISs()`, `MatCreateNest()`,
1199           `MatNestSetSubMats()`, `MatNestGetISs()`, `MatNestSetSubMat()`
1200 @*/
MatNestGetSubMats(Mat A,PetscInt * M,PetscInt * N,Mat *** mat)1201 PetscErrorCode MatNestGetSubMats(Mat A, PetscInt *M, PetscInt *N, Mat ***mat)
1202 {
1203   PetscFunctionBegin;
1204   PetscValidHeaderSpecific(A, MAT_CLASSID, 1);
1205   PetscUseMethod(A, "MatNestGetSubMats_C", (Mat, PetscInt *, PetscInt *, Mat ***), (A, M, N, mat));
1206   PetscFunctionReturn(PETSC_SUCCESS);
1207 }
1208 
MatNestGetSize_Nest(Mat A,PetscInt * M,PetscInt * N)1209 static PetscErrorCode MatNestGetSize_Nest(Mat A, PetscInt *M, PetscInt *N)
1210 {
1211   Mat_Nest *bA = (Mat_Nest *)A->data;
1212 
1213   PetscFunctionBegin;
1214   if (M) *M = bA->nr;
1215   if (N) *N = bA->nc;
1216   PetscFunctionReturn(PETSC_SUCCESS);
1217 }
1218 
1219 /*@
1220   MatNestGetSize - Returns the size of the `MATNEST` matrix.
1221 
1222   Not Collective
1223 
1224   Input Parameter:
1225 . A - `MATNEST` matrix
1226 
1227   Output Parameters:
1228 + M - number of rows in the nested mat
1229 - N - number of cols in the nested mat
1230 
1231   Level: developer
1232 
1233   Note:
1234   `size` refers to the number of submatrices in the row and column directions of the nested matrix
1235 
1236 .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestGetSubMat()`, `MatNestGetSubMats()`, `MatCreateNest()`, `MatNestGetLocalISs()`,
1237           `MatNestGetISs()`
1238 @*/
MatNestGetSize(Mat A,PetscInt * M,PetscInt * N)1239 PetscErrorCode MatNestGetSize(Mat A, PetscInt *M, PetscInt *N)
1240 {
1241   PetscFunctionBegin;
1242   PetscValidHeaderSpecific(A, MAT_CLASSID, 1);
1243   PetscUseMethod(A, "MatNestGetSize_C", (Mat, PetscInt *, PetscInt *), (A, M, N));
1244   PetscFunctionReturn(PETSC_SUCCESS);
1245 }
1246 
MatNestGetISs_Nest(Mat A,IS rows[],IS cols[])1247 static PetscErrorCode MatNestGetISs_Nest(Mat A, IS rows[], IS cols[])
1248 {
1249   Mat_Nest *vs = (Mat_Nest *)A->data;
1250   PetscInt  i;
1251 
1252   PetscFunctionBegin;
1253   if (rows)
1254     for (i = 0; i < vs->nr; i++) rows[i] = vs->isglobal.row[i];
1255   if (cols)
1256     for (i = 0; i < vs->nc; i++) cols[i] = vs->isglobal.col[i];
1257   PetscFunctionReturn(PETSC_SUCCESS);
1258 }
1259 
1260 /*@
1261   MatNestGetISs - Returns the index sets partitioning the row and column spaces of a `MATNEST`
1262 
1263   Not Collective
1264 
1265   Input Parameter:
1266 . A - `MATNEST` matrix
1267 
1268   Output Parameters:
1269 + rows - array of row index sets (pass `NULL` to ignore)
1270 - cols - array of column index sets (pass `NULL` to ignore)
1271 
1272   Level: advanced
1273 
1274   Note:
1275   The user must have allocated arrays of the correct size. The reference count is not increased on the returned `IS`s.
1276 
1277 .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestGetSubMat()`, `MatNestGetSubMats()`, `MatNestGetSize()`, `MatNestGetLocalISs()`,
1278           `MatCreateNest()`, `MatNestSetSubMats()`
1279 @*/
MatNestGetISs(Mat A,IS rows[],IS cols[])1280 PetscErrorCode MatNestGetISs(Mat A, IS rows[], IS cols[])
1281 {
1282   PetscFunctionBegin;
1283   PetscValidHeaderSpecific(A, MAT_CLASSID, 1);
1284   PetscUseMethod(A, "MatNestGetISs_C", (Mat, IS[], IS[]), (A, rows, cols));
1285   PetscFunctionReturn(PETSC_SUCCESS);
1286 }
1287 
MatNestGetLocalISs_Nest(Mat A,IS rows[],IS cols[])1288 static PetscErrorCode MatNestGetLocalISs_Nest(Mat A, IS rows[], IS cols[])
1289 {
1290   Mat_Nest *vs = (Mat_Nest *)A->data;
1291   PetscInt  i;
1292 
1293   PetscFunctionBegin;
1294   if (rows)
1295     for (i = 0; i < vs->nr; i++) rows[i] = vs->islocal.row[i];
1296   if (cols)
1297     for (i = 0; i < vs->nc; i++) cols[i] = vs->islocal.col[i];
1298   PetscFunctionReturn(PETSC_SUCCESS);
1299 }
1300 
1301 /*@
1302   MatNestGetLocalISs - Returns the index sets partitioning the row and column spaces of a `MATNEST`
1303 
1304   Not Collective
1305 
1306   Input Parameter:
1307 . A - `MATNEST` matrix
1308 
1309   Output Parameters:
1310 + rows - array of row index sets (pass `NULL` to ignore)
1311 - cols - array of column index sets (pass `NULL` to ignore)
1312 
1313   Level: advanced
1314 
1315   Note:
1316   The user must have allocated arrays of the correct size. The reference count is not increased on the returned `IS`s.
1317 
1318 .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestGetSubMat()`, `MatNestGetSubMats()`, `MatNestGetSize()`, `MatNestGetISs()`, `MatCreateNest()`,
1319           `MatNestSetSubMats()`, `MatNestSetSubMat()`
1320 @*/
MatNestGetLocalISs(Mat A,IS rows[],IS cols[])1321 PetscErrorCode MatNestGetLocalISs(Mat A, IS rows[], IS cols[])
1322 {
1323   PetscFunctionBegin;
1324   PetscValidHeaderSpecific(A, MAT_CLASSID, 1);
1325   PetscUseMethod(A, "MatNestGetLocalISs_C", (Mat, IS[], IS[]), (A, rows, cols));
1326   PetscFunctionReturn(PETSC_SUCCESS);
1327 }
1328 
MatNestSetVecType_Nest(Mat A,VecType vtype)1329 static PetscErrorCode MatNestSetVecType_Nest(Mat A, VecType vtype)
1330 {
1331   PetscBool flg;
1332 
1333   PetscFunctionBegin;
1334   PetscCall(PetscStrcmp(vtype, VECNEST, &flg));
1335   /* In reality, this only distinguishes VECNEST and "other" */
1336   if (flg) A->ops->getvecs = MatCreateVecs_Nest;
1337   else A->ops->getvecs = NULL;
1338   PetscFunctionReturn(PETSC_SUCCESS);
1339 }
1340 
1341 /*@
1342   MatNestSetVecType - Sets the type of `Vec` returned by `MatCreateVecs()`
1343 
1344   Not Collective
1345 
1346   Input Parameters:
1347 + A     - `MATNEST` matrix
1348 - vtype - `VecType` to use for creating vectors
1349 
1350   Level: developer
1351 
1352 .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatCreateVecs()`, `MatCreateNest()`, `VecType`
1353 @*/
MatNestSetVecType(Mat A,VecType vtype)1354 PetscErrorCode MatNestSetVecType(Mat A, VecType vtype)
1355 {
1356   PetscFunctionBegin;
1357   PetscValidHeaderSpecific(A, MAT_CLASSID, 1);
1358   PetscTryMethod(A, "MatNestSetVecType_C", (Mat, VecType), (A, vtype));
1359   PetscFunctionReturn(PETSC_SUCCESS);
1360 }
1361 
MatNestSetSubMats_Nest(Mat A,PetscInt nr,const IS is_row[],PetscInt nc,const IS is_col[],const Mat a[])1362 static PetscErrorCode MatNestSetSubMats_Nest(Mat A, PetscInt nr, const IS is_row[], PetscInt nc, const IS is_col[], const Mat a[])
1363 {
1364   Mat_Nest *s = (Mat_Nest *)A->data;
1365   PetscInt  i, j, m, n, M, N;
1366   PetscBool cong, isstd, sametype = PETSC_FALSE;
1367   VecType   vtype, type;
1368 
1369   PetscFunctionBegin;
1370   PetscCall(MatReset_Nest(A));
1371 
1372   s->nr = nr;
1373   s->nc = nc;
1374 
1375   /* Create space for submatrices */
1376   PetscCall(PetscMalloc1(nr, &s->m));
1377   PetscCall(PetscMalloc1(nr * nc, &s->m[0]));
1378   for (i = 0; i < nr; i++) {
1379     s->m[i] = s->m[0] + i * nc;
1380     for (j = 0; j < nc; j++) {
1381       s->m[i][j] = a ? a[i * nc + j] : NULL;
1382       PetscCall(PetscObjectReference((PetscObject)s->m[i][j]));
1383     }
1384   }
1385   PetscCall(MatGetVecType(A, &vtype));
1386   PetscCall(PetscStrcmp(vtype, VECSTANDARD, &isstd));
1387   if (isstd) {
1388     /* check if all blocks have the same vectype */
1389     vtype = NULL;
1390     for (i = 0; i < nr; i++) {
1391       for (j = 0; j < nc; j++) {
1392         if (s->m[i][j]) {
1393           if (!vtype) { /* first visited block */
1394             PetscCall(MatGetVecType(s->m[i][j], &vtype));
1395             sametype = PETSC_TRUE;
1396           } else if (sametype) {
1397             PetscCall(MatGetVecType(s->m[i][j], &type));
1398             PetscCall(PetscStrcmp(vtype, type, &sametype));
1399           }
1400         }
1401       }
1402     }
1403     if (sametype) { /* propagate vectype */
1404       PetscCall(MatSetVecType(A, vtype));
1405     }
1406   }
1407 
1408   PetscCall(MatSetUp_NestIS_Private(A, nr, is_row, nc, is_col));
1409 
1410   PetscCall(PetscMalloc1(nr, &s->row_len));
1411   PetscCall(PetscMalloc1(nc, &s->col_len));
1412   for (i = 0; i < nr; i++) s->row_len[i] = -1;
1413   for (j = 0; j < nc; j++) s->col_len[j] = -1;
1414 
1415   PetscCall(PetscCalloc1(nr * nc, &s->nnzstate));
1416   for (i = 0; i < nr; i++) {
1417     for (j = 0; j < nc; j++) {
1418       if (s->m[i][j]) PetscCall(MatGetNonzeroState(s->m[i][j], &s->nnzstate[i * nc + j]));
1419     }
1420   }
1421 
1422   PetscCall(MatNestGetSizes_Private(A, &m, &n, &M, &N));
1423 
1424   PetscCall(PetscLayoutSetSize(A->rmap, M));
1425   PetscCall(PetscLayoutSetLocalSize(A->rmap, m));
1426   PetscCall(PetscLayoutSetSize(A->cmap, N));
1427   PetscCall(PetscLayoutSetLocalSize(A->cmap, n));
1428 
1429   PetscCall(PetscLayoutSetUp(A->rmap));
1430   PetscCall(PetscLayoutSetUp(A->cmap));
1431 
1432   /* disable operations that are not supported for non-square matrices,
1433      or matrices for which is_row != is_col  */
1434   PetscCall(MatHasCongruentLayouts(A, &cong));
1435   if (cong && nr != nc) cong = PETSC_FALSE;
1436   if (cong) {
1437     for (i = 0; cong && i < nr; i++) PetscCall(ISEqualUnsorted(s->isglobal.row[i], s->isglobal.col[i], &cong));
1438   }
1439   if (!cong) {
1440     A->ops->getdiagonal = NULL;
1441     A->ops->shift       = NULL;
1442     A->ops->diagonalset = NULL;
1443   }
1444 
1445   PetscCall(PetscCalloc2(nr, &s->left, nc, &s->right));
1446   PetscCall(PetscObjectStateIncrease((PetscObject)A));
1447   A->nonzerostate++;
1448   PetscFunctionReturn(PETSC_SUCCESS);
1449 }
1450 
1451 /*@
1452   MatNestSetSubMats - Sets the nested submatrices in a `MATNEST`
1453 
1454   Collective
1455 
1456   Input Parameters:
1457 + A      - `MATNEST` matrix
1458 . nr     - number of nested row blocks
1459 . is_row - index sets for each nested row block, or `NULL` to make contiguous
1460 . nc     - number of nested column blocks
1461 . is_col - index sets for each nested column block, or `NULL` to make contiguous
1462 - a      - array of $ nr \times nc$ submatrices, or `NULL`
1463 
1464   Level: advanced
1465 
1466   Notes:
1467   This always resets any block matrix information previously set.
1468 
1469   Pass `NULL` in the corresponding entry of `a` for an empty block.
1470 
1471   In both C and Fortran, `a` must be a one-dimensional array representing a two-dimensional row-major order array containing the matrices. See
1472   `MatCreateNest()` for an example.
1473 
1474   Fortran Note:
1475   Pass `PETSC_NULL_MAT` in the corresponding entry of `a` for an empty block
1476 
1477 .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatCreateNest()`, `MatNestSetSubMat()`, `MatNestGetSubMat()`, `MatNestGetSubMats()`
1478 @*/
MatNestSetSubMats(Mat A,PetscInt nr,const IS is_row[],PetscInt nc,const IS is_col[],const Mat a[])1479 PetscErrorCode MatNestSetSubMats(Mat A, PetscInt nr, const IS is_row[], PetscInt nc, const IS is_col[], const Mat a[]) PeNSS
1480 {
1481   PetscFunctionBegin;
1482   PetscValidHeaderSpecific(A, MAT_CLASSID, 1);
1483   PetscValidLogicalCollectiveInt(A, nr, 2);
1484   PetscCheck(nr >= 0, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Number of rows cannot be negative");
1485   if (nr && is_row) {
1486     PetscAssertPointer(is_row, 3);
1487     for (PetscInt i = 0; i < nr; i++) PetscValidHeaderSpecific(is_row[i], IS_CLASSID, 3);
1488   }
1489   PetscValidLogicalCollectiveInt(A, nc, 4);
1490   PetscCheck(nc >= 0, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Number of columns cannot be negative");
1491   if (nc && is_col) {
1492     PetscAssertPointer(is_col, 5);
1493     for (PetscInt i = 0; i < nc; i++) PetscValidHeaderSpecific(is_col[i], IS_CLASSID, 5);
1494   }
1495   PetscTryMethod(A, "MatNestSetSubMats_C", (Mat, PetscInt, const IS[], PetscInt, const IS[], const Mat[]), (A, nr, is_row, nc, is_col, a));
1496   PetscFunctionReturn(PETSC_SUCCESS);
1497 }
1498 
MatNestCreateAggregateL2G_Private(Mat A,PetscInt n,const IS islocal[],const IS isglobal[],PetscBool colflg,ISLocalToGlobalMapping * ltog)1499 static PetscErrorCode MatNestCreateAggregateL2G_Private(Mat A, PetscInt n, const IS islocal[], const IS isglobal[], PetscBool colflg, ISLocalToGlobalMapping *ltog)
1500 {
1501   PetscBool flg;
1502   PetscInt  i, j, m, mi, *ix;
1503 
1504   PetscFunctionBegin;
1505   *ltog = NULL;
1506   for (i = 0, m = 0, flg = PETSC_FALSE; i < n; i++) {
1507     if (islocal[i]) {
1508       PetscCall(ISGetLocalSize(islocal[i], &mi));
1509       flg = PETSC_TRUE; /* We found a non-trivial entry */
1510     } else {
1511       PetscCall(ISGetLocalSize(isglobal[i], &mi));
1512     }
1513     m += mi;
1514   }
1515   if (!flg) PetscFunctionReturn(PETSC_SUCCESS);
1516 
1517   PetscCall(PetscMalloc1(m, &ix));
1518   for (i = 0, m = 0; i < n; i++) {
1519     ISLocalToGlobalMapping smap = NULL;
1520     Mat                    sub  = NULL;
1521     PetscSF                sf;
1522     PetscLayout            map;
1523     const PetscInt        *ix2;
1524 
1525     if (!colflg) {
1526       PetscCall(MatNestFindNonzeroSubMatRow(A, i, &sub));
1527     } else {
1528       PetscCall(MatNestFindNonzeroSubMatCol(A, i, &sub));
1529     }
1530     if (sub) {
1531       if (!colflg) {
1532         PetscCall(MatGetLocalToGlobalMapping(sub, &smap, NULL));
1533       } else {
1534         PetscCall(MatGetLocalToGlobalMapping(sub, NULL, &smap));
1535       }
1536     }
1537     /*
1538        Now we need to extract the monolithic global indices that correspond to the given split global indices.
1539        In many/most cases, we only want MatGetLocalSubMatrix() to work, in which case we only need to know the size of the local spaces.
1540     */
1541     PetscCall(ISGetIndices(isglobal[i], &ix2));
1542     if (islocal[i]) {
1543       PetscInt *ilocal, *iremote;
1544       PetscInt  mil, nleaves;
1545 
1546       PetscCall(ISGetLocalSize(islocal[i], &mi));
1547       PetscCheck(smap, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing local to global map");
1548       for (j = 0; j < mi; j++) ix[m + j] = j;
1549       PetscCall(ISLocalToGlobalMappingApply(smap, mi, ix + m, ix + m));
1550 
1551       /* PetscSFSetGraphLayout does not like negative indices */
1552       PetscCall(PetscMalloc2(mi, &ilocal, mi, &iremote));
1553       for (j = 0, nleaves = 0; j < mi; j++) {
1554         if (ix[m + j] < 0) continue;
1555         ilocal[nleaves]  = j;
1556         iremote[nleaves] = ix[m + j];
1557         nleaves++;
1558       }
1559       PetscCall(ISGetLocalSize(isglobal[i], &mil));
1560       PetscCall(PetscSFCreate(PetscObjectComm((PetscObject)A), &sf));
1561       PetscCall(PetscLayoutCreate(PetscObjectComm((PetscObject)A), &map));
1562       PetscCall(PetscLayoutSetLocalSize(map, mil));
1563       PetscCall(PetscLayoutSetUp(map));
1564       PetscCall(PetscSFSetGraphLayout(sf, map, nleaves, ilocal, PETSC_USE_POINTER, iremote));
1565       PetscCall(PetscLayoutDestroy(&map));
1566       PetscCall(PetscSFBcastBegin(sf, MPIU_INT, ix2, ix + m, MPI_REPLACE));
1567       PetscCall(PetscSFBcastEnd(sf, MPIU_INT, ix2, ix + m, MPI_REPLACE));
1568       PetscCall(PetscSFDestroy(&sf));
1569       PetscCall(PetscFree2(ilocal, iremote));
1570     } else {
1571       PetscCall(ISGetLocalSize(isglobal[i], &mi));
1572       for (j = 0; j < mi; j++) ix[m + j] = ix2[i];
1573     }
1574     PetscCall(ISRestoreIndices(isglobal[i], &ix2));
1575     m += mi;
1576   }
1577   PetscCall(ISLocalToGlobalMappingCreate(PetscObjectComm((PetscObject)A), 1, m, ix, PETSC_OWN_POINTER, ltog));
1578   PetscFunctionReturn(PETSC_SUCCESS);
1579 }
1580 
1581 /* If an IS was provided, there is nothing Nest needs to do, otherwise Nest will build a strided IS */
1582 /*
1583   nprocessors = NP
1584   Nest x^T = ((g_0,g_1,...g_nprocs-1), (h_0,h_1,...h_NP-1))
1585        proc 0: => (g_0,h_0,)
1586        proc 1: => (g_1,h_1,)
1587        ...
1588        proc nprocs-1: => (g_NP-1,h_NP-1,)
1589 
1590             proc 0:                      proc 1:                    proc nprocs-1:
1591     is[0] = (0,1,2,...,nlocal(g_0)-1)  (0,1,...,nlocal(g_1)-1)  (0,1,...,nlocal(g_NP-1))
1592 
1593             proc 0:
1594     is[1] = (nlocal(g_0),nlocal(g_0)+1,...,nlocal(g_0)+nlocal(h_0)-1)
1595             proc 1:
1596     is[1] = (nlocal(g_1),nlocal(g_1)+1,...,nlocal(g_1)+nlocal(h_1)-1)
1597 
1598             proc NP-1:
1599     is[1] = (nlocal(g_NP-1),nlocal(g_NP-1)+1,...,nlocal(g_NP-1)+nlocal(h_NP-1)-1)
1600 */
MatSetUp_NestIS_Private(Mat A,PetscInt nr,const IS is_row[],PetscInt nc,const IS is_col[])1601 static PetscErrorCode MatSetUp_NestIS_Private(Mat A, PetscInt nr, const IS is_row[], PetscInt nc, const IS is_col[])
1602 {
1603   Mat_Nest *vs = (Mat_Nest *)A->data;
1604   PetscInt  i, j, offset, n, nsum, bs;
1605   Mat       sub = NULL;
1606 
1607   PetscFunctionBegin;
1608   PetscCall(PetscMalloc1(nr, &vs->isglobal.row));
1609   PetscCall(PetscMalloc1(nc, &vs->isglobal.col));
1610   if (is_row) { /* valid IS is passed in */
1611     /* refs on is[] are incremented */
1612     for (i = 0; i < vs->nr; i++) {
1613       PetscCall(PetscObjectReference((PetscObject)is_row[i]));
1614       vs->isglobal.row[i] = is_row[i];
1615     }
1616   } else { /* Create the ISs by inspecting sizes of a submatrix in each row */
1617     nsum = 0;
1618     for (i = 0; i < vs->nr; i++) { /* Add up the local sizes to compute the aggregate offset */
1619       PetscCall(MatNestFindNonzeroSubMatRow(A, i, &sub));
1620       PetscCheck(sub, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "No nonzero submatrix in row %" PetscInt_FMT, i);
1621       PetscCall(MatGetLocalSize(sub, &n, NULL));
1622       PetscCheck(n >= 0, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Sizes have not yet been set for submatrix");
1623       nsum += n;
1624     }
1625     PetscCallMPI(MPI_Scan(&nsum, &offset, 1, MPIU_INT, MPI_SUM, PetscObjectComm((PetscObject)A)));
1626     offset -= nsum;
1627     for (i = 0; i < vs->nr; i++) {
1628       PetscCall(MatNestFindNonzeroSubMatRow(A, i, &sub));
1629       PetscCall(MatGetLocalSize(sub, &n, NULL));
1630       PetscCall(MatGetBlockSizes(sub, &bs, NULL));
1631       PetscCall(ISCreateStride(PetscObjectComm((PetscObject)sub), n, offset, 1, &vs->isglobal.row[i]));
1632       PetscCall(ISSetBlockSize(vs->isglobal.row[i], bs));
1633       offset += n;
1634     }
1635   }
1636 
1637   if (is_col) { /* valid IS is passed in */
1638     /* refs on is[] are incremented */
1639     for (j = 0; j < vs->nc; j++) {
1640       PetscCall(PetscObjectReference((PetscObject)is_col[j]));
1641       vs->isglobal.col[j] = is_col[j];
1642     }
1643   } else { /* Create the ISs by inspecting sizes of a submatrix in each column */
1644     offset = A->cmap->rstart;
1645     nsum   = 0;
1646     for (j = 0; j < vs->nc; j++) {
1647       PetscCall(MatNestFindNonzeroSubMatCol(A, j, &sub));
1648       PetscCheck(sub, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "No nonzero submatrix in column %" PetscInt_FMT, i);
1649       PetscCall(MatGetLocalSize(sub, NULL, &n));
1650       PetscCheck(n >= 0, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Sizes have not yet been set for submatrix");
1651       nsum += n;
1652     }
1653     PetscCallMPI(MPI_Scan(&nsum, &offset, 1, MPIU_INT, MPI_SUM, PetscObjectComm((PetscObject)A)));
1654     offset -= nsum;
1655     for (j = 0; j < vs->nc; j++) {
1656       PetscCall(MatNestFindNonzeroSubMatCol(A, j, &sub));
1657       PetscCall(MatGetLocalSize(sub, NULL, &n));
1658       PetscCall(MatGetBlockSizes(sub, NULL, &bs));
1659       PetscCall(ISCreateStride(PetscObjectComm((PetscObject)sub), n, offset, 1, &vs->isglobal.col[j]));
1660       PetscCall(ISSetBlockSize(vs->isglobal.col[j], bs));
1661       offset += n;
1662     }
1663   }
1664 
1665   /* Set up the local ISs */
1666   PetscCall(PetscMalloc1(vs->nr, &vs->islocal.row));
1667   PetscCall(PetscMalloc1(vs->nc, &vs->islocal.col));
1668   for (i = 0, offset = 0; i < vs->nr; i++) {
1669     IS                     isloc;
1670     ISLocalToGlobalMapping rmap = NULL;
1671     PetscInt               nlocal, bs;
1672     PetscCall(MatNestFindNonzeroSubMatRow(A, i, &sub));
1673     if (sub) PetscCall(MatGetLocalToGlobalMapping(sub, &rmap, NULL));
1674     if (rmap) {
1675       PetscCall(MatGetBlockSizes(sub, &bs, NULL));
1676       PetscCall(ISLocalToGlobalMappingGetSize(rmap, &nlocal));
1677       PetscCall(ISCreateStride(PETSC_COMM_SELF, nlocal, offset, 1, &isloc));
1678       PetscCall(ISSetBlockSize(isloc, bs));
1679     } else {
1680       nlocal = 0;
1681       isloc  = NULL;
1682     }
1683     vs->islocal.row[i] = isloc;
1684     offset += nlocal;
1685   }
1686   for (i = 0, offset = 0; i < vs->nc; i++) {
1687     IS                     isloc;
1688     ISLocalToGlobalMapping cmap = NULL;
1689     PetscInt               nlocal, bs;
1690     PetscCall(MatNestFindNonzeroSubMatCol(A, i, &sub));
1691     if (sub) PetscCall(MatGetLocalToGlobalMapping(sub, NULL, &cmap));
1692     if (cmap) {
1693       PetscCall(MatGetBlockSizes(sub, NULL, &bs));
1694       PetscCall(ISLocalToGlobalMappingGetSize(cmap, &nlocal));
1695       PetscCall(ISCreateStride(PETSC_COMM_SELF, nlocal, offset, 1, &isloc));
1696       PetscCall(ISSetBlockSize(isloc, bs));
1697     } else {
1698       nlocal = 0;
1699       isloc  = NULL;
1700     }
1701     vs->islocal.col[i] = isloc;
1702     offset += nlocal;
1703   }
1704 
1705   /* Set up the aggregate ISLocalToGlobalMapping */
1706   {
1707     ISLocalToGlobalMapping rmap, cmap;
1708     PetscCall(MatNestCreateAggregateL2G_Private(A, vs->nr, vs->islocal.row, vs->isglobal.row, PETSC_FALSE, &rmap));
1709     PetscCall(MatNestCreateAggregateL2G_Private(A, vs->nc, vs->islocal.col, vs->isglobal.col, PETSC_TRUE, &cmap));
1710     if (rmap && cmap) PetscCall(MatSetLocalToGlobalMapping(A, rmap, cmap));
1711     PetscCall(ISLocalToGlobalMappingDestroy(&rmap));
1712     PetscCall(ISLocalToGlobalMappingDestroy(&cmap));
1713   }
1714 
1715   if (PetscDefined(USE_DEBUG)) {
1716     for (i = 0; i < vs->nr; i++) {
1717       for (j = 0; j < vs->nc; j++) {
1718         PetscInt m, n, M, N, mi, ni, Mi, Ni;
1719         Mat      B = vs->m[i][j];
1720         if (!B) continue;
1721         PetscCall(MatGetSize(B, &M, &N));
1722         PetscCall(MatGetLocalSize(B, &m, &n));
1723         PetscCall(ISGetSize(vs->isglobal.row[i], &Mi));
1724         PetscCall(ISGetSize(vs->isglobal.col[j], &Ni));
1725         PetscCall(ISGetLocalSize(vs->isglobal.row[i], &mi));
1726         PetscCall(ISGetLocalSize(vs->isglobal.col[j], &ni));
1727         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);
1728         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);
1729       }
1730     }
1731   }
1732 
1733   /* Set A->assembled if all non-null blocks are currently assembled */
1734   for (i = 0; i < vs->nr; i++) {
1735     for (j = 0; j < vs->nc; j++) {
1736       if (vs->m[i][j] && !vs->m[i][j]->assembled) PetscFunctionReturn(PETSC_SUCCESS);
1737     }
1738   }
1739   A->assembled = PETSC_TRUE;
1740   PetscFunctionReturn(PETSC_SUCCESS);
1741 }
1742 
1743 /*@C
1744   MatCreateNest - Creates a new `MATNEST` matrix containing several nested submatrices, each stored separately
1745 
1746   Collective
1747 
1748   Input Parameters:
1749 + comm   - Communicator for the new `MATNEST`
1750 . nr     - number of nested row blocks
1751 . is_row - index sets for each nested row block, or `NULL` to make contiguous
1752 . nc     - number of nested column blocks
1753 . is_col - index sets for each nested column block, or `NULL` to make contiguous
1754 - a      - array of $nr \times nc$ submatrices, empty submatrices can be passed using `NULL`
1755 
1756   Output Parameter:
1757 . B - new matrix
1758 
1759   Level: advanced
1760 
1761   Note:
1762   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.
1763   For instance, to represent the matrix
1764   $\begin{bmatrix} A_{11} & A_{12} \\ A_{21} & A_{22}\end{bmatrix}$
1765   one should use `Mat a[4]={A11,A12,A21,A22}`.
1766 
1767   Fortran Note:
1768   Pass `PETSC_NULL_MAT` in the corresponding entry of `a` for an empty block
1769 
1770 .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatCreate()`, `VecCreateNest()`, `DMCreateMatrix()`, `MatNestSetSubMat()`,
1771           `MatNestGetSubMat()`, `MatNestGetLocalISs()`, `MatNestGetSize()`,
1772           `MatNestGetISs()`, `MatNestSetSubMats()`, `MatNestGetSubMats()`
1773 @*/
MatCreateNest(MPI_Comm comm,PetscInt nr,const IS is_row[],PetscInt nc,const IS is_col[],const Mat a[],Mat * B)1774 PetscErrorCode MatCreateNest(MPI_Comm comm, PetscInt nr, const IS is_row[], PetscInt nc, const IS is_col[], const Mat a[], Mat *B) PeNSS
1775 {
1776   PetscFunctionBegin;
1777   PetscCall(MatCreate(comm, B));
1778   PetscCall(MatSetType(*B, MATNEST));
1779   (*B)->preallocated = PETSC_TRUE;
1780   PetscCall(MatNestSetSubMats(*B, nr, is_row, nc, is_col, a));
1781   PetscFunctionReturn(PETSC_SUCCESS);
1782 }
1783 
MatConvert_Nest_SeqAIJ_fast(Mat A,MatType newtype,MatReuse reuse,Mat * newmat)1784 static PetscErrorCode MatConvert_Nest_SeqAIJ_fast(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
1785 {
1786   Mat_Nest     *nest = (Mat_Nest *)A->data;
1787   Mat          *trans;
1788   PetscScalar **avv;
1789   PetscScalar  *vv;
1790   PetscInt    **aii, **ajj;
1791   PetscInt     *ii, *jj, *ci;
1792   PetscInt      nr, nc, nnz, i, j;
1793   PetscBool     done;
1794 
1795   PetscFunctionBegin;
1796   PetscCall(MatGetSize(A, &nr, &nc));
1797   if (reuse == MAT_REUSE_MATRIX) {
1798     PetscInt rnr;
1799 
1800     PetscCall(MatGetRowIJ(*newmat, 0, PETSC_FALSE, PETSC_FALSE, &rnr, (const PetscInt **)&ii, (const PetscInt **)&jj, &done));
1801     PetscCheck(done, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "MatGetRowIJ");
1802     PetscCheck(rnr == nr, PetscObjectComm((PetscObject)A), PETSC_ERR_USER, "Cannot reuse matrix, wrong number of rows");
1803     PetscCall(MatSeqAIJGetArray(*newmat, &vv));
1804   }
1805   /* extract CSR for nested SeqAIJ matrices */
1806   nnz = 0;
1807   PetscCall(PetscCalloc4(nest->nr * nest->nc, &aii, nest->nr * nest->nc, &ajj, nest->nr * nest->nc, &avv, nest->nr * nest->nc, &trans));
1808   for (i = 0; i < nest->nr; ++i) {
1809     for (j = 0; j < nest->nc; ++j) {
1810       Mat B = nest->m[i][j];
1811       if (B) {
1812         PetscScalar *naa;
1813         PetscInt    *nii, *njj, nnr;
1814         PetscBool    istrans;
1815 
1816         PetscCall(PetscObjectTypeCompare((PetscObject)B, MATTRANSPOSEVIRTUAL, &istrans));
1817         if (istrans) {
1818           Mat Bt;
1819 
1820           PetscCall(MatTransposeGetMat(B, &Bt));
1821           PetscCall(MatTranspose(Bt, MAT_INITIAL_MATRIX, &trans[i * nest->nc + j]));
1822           B = trans[i * nest->nc + j];
1823         } else {
1824           PetscCall(PetscObjectTypeCompare((PetscObject)B, MATHERMITIANTRANSPOSEVIRTUAL, &istrans));
1825           if (istrans) {
1826             Mat Bt;
1827 
1828             PetscCall(MatHermitianTransposeGetMat(B, &Bt));
1829             PetscCall(MatHermitianTranspose(Bt, MAT_INITIAL_MATRIX, &trans[i * nest->nc + j]));
1830             B = trans[i * nest->nc + j];
1831           }
1832         }
1833         PetscCall(MatGetRowIJ(B, 0, PETSC_FALSE, PETSC_FALSE, &nnr, (const PetscInt **)&nii, (const PetscInt **)&njj, &done));
1834         PetscCheck(done, PetscObjectComm((PetscObject)B), PETSC_ERR_PLIB, "MatGetRowIJ");
1835         PetscCall(MatSeqAIJGetArray(B, &naa));
1836         nnz += nii[nnr];
1837 
1838         aii[i * nest->nc + j] = nii;
1839         ajj[i * nest->nc + j] = njj;
1840         avv[i * nest->nc + j] = naa;
1841       }
1842     }
1843   }
1844   if (reuse != MAT_REUSE_MATRIX) {
1845     PetscCall(PetscMalloc1(nr + 1, &ii));
1846     PetscCall(PetscMalloc1(nnz, &jj));
1847     PetscCall(PetscMalloc1(nnz, &vv));
1848   } else {
1849     PetscCheck(nnz == ii[nr], PetscObjectComm((PetscObject)A), PETSC_ERR_USER, "Cannot reuse matrix, wrong number of nonzeros");
1850   }
1851 
1852   /* new row pointer */
1853   PetscCall(PetscArrayzero(ii, nr + 1));
1854   for (i = 0; i < nest->nr; ++i) {
1855     PetscInt ncr, rst;
1856 
1857     PetscCall(ISStrideGetInfo(nest->isglobal.row[i], &rst, NULL));
1858     PetscCall(ISGetLocalSize(nest->isglobal.row[i], &ncr));
1859     for (j = 0; j < nest->nc; ++j) {
1860       if (aii[i * nest->nc + j]) {
1861         PetscInt *nii = aii[i * nest->nc + j];
1862         PetscInt  ir;
1863 
1864         for (ir = rst; ir < ncr + rst; ++ir) {
1865           ii[ir + 1] += nii[1] - nii[0];
1866           nii++;
1867         }
1868       }
1869     }
1870   }
1871   for (i = 0; i < nr; i++) ii[i + 1] += ii[i];
1872 
1873   /* construct CSR for the new matrix */
1874   PetscCall(PetscCalloc1(nr, &ci));
1875   for (i = 0; i < nest->nr; ++i) {
1876     PetscInt ncr, rst;
1877 
1878     PetscCall(ISStrideGetInfo(nest->isglobal.row[i], &rst, NULL));
1879     PetscCall(ISGetLocalSize(nest->isglobal.row[i], &ncr));
1880     for (j = 0; j < nest->nc; ++j) {
1881       if (aii[i * nest->nc + j]) {
1882         PetscScalar *nvv = avv[i * nest->nc + j], vscale = 1.0, vshift = 0.0;
1883         PetscInt    *nii = aii[i * nest->nc + j];
1884         PetscInt    *njj = ajj[i * nest->nc + j];
1885         PetscInt     ir, cst;
1886 
1887         if (trans[i * nest->nc + j]) {
1888           vscale = ((Mat_Shell *)nest->m[i][j]->data)->vscale;
1889           vshift = ((Mat_Shell *)nest->m[i][j]->data)->vshift;
1890         }
1891         PetscCall(ISStrideGetInfo(nest->isglobal.col[j], &cst, NULL));
1892         for (ir = rst; ir < ncr + rst; ++ir) {
1893           PetscInt ij, rsize = nii[1] - nii[0], ist = ii[ir] + ci[ir];
1894 
1895           for (ij = 0; ij < rsize; ij++) {
1896             jj[ist + ij] = *njj + cst;
1897             vv[ist + ij] = vscale * *nvv;
1898             if (PetscUnlikely(vshift != 0.0 && *njj == ir - rst)) vv[ist + ij] += vshift;
1899             njj++;
1900             nvv++;
1901           }
1902           ci[ir] += rsize;
1903           nii++;
1904         }
1905       }
1906     }
1907   }
1908   PetscCall(PetscFree(ci));
1909 
1910   /* restore info */
1911   for (i = 0; i < nest->nr; ++i) {
1912     for (j = 0; j < nest->nc; ++j) {
1913       Mat B = nest->m[i][j];
1914       if (B) {
1915         PetscInt nnr = 0, k = i * nest->nc + j;
1916 
1917         B = (trans[k] ? trans[k] : B);
1918         PetscCall(MatRestoreRowIJ(B, 0, PETSC_FALSE, PETSC_FALSE, &nnr, (const PetscInt **)&aii[k], (const PetscInt **)&ajj[k], &done));
1919         PetscCheck(done, PetscObjectComm((PetscObject)B), PETSC_ERR_PLIB, "MatRestoreRowIJ");
1920         PetscCall(MatSeqAIJRestoreArray(B, &avv[k]));
1921         PetscCall(MatDestroy(&trans[k]));
1922       }
1923     }
1924   }
1925   PetscCall(PetscFree4(aii, ajj, avv, trans));
1926 
1927   /* finalize newmat */
1928   if (reuse == MAT_INITIAL_MATRIX) {
1929     PetscCall(MatCreateSeqAIJWithArrays(PetscObjectComm((PetscObject)A), nr, nc, ii, jj, vv, newmat));
1930   } else if (reuse == MAT_INPLACE_MATRIX) {
1931     Mat B;
1932 
1933     PetscCall(MatCreateSeqAIJWithArrays(PetscObjectComm((PetscObject)A), nr, nc, ii, jj, vv, &B));
1934     PetscCall(MatHeaderReplace(A, &B));
1935   }
1936   PetscCall(MatAssemblyBegin(*newmat, MAT_FINAL_ASSEMBLY));
1937   PetscCall(MatAssemblyEnd(*newmat, MAT_FINAL_ASSEMBLY));
1938   {
1939     Mat_SeqAIJ *a = (Mat_SeqAIJ *)((*newmat)->data);
1940     a->free_a     = PETSC_TRUE;
1941     a->free_ij    = PETSC_TRUE;
1942   }
1943   PetscFunctionReturn(PETSC_SUCCESS);
1944 }
1945 
MatAXPY_Dense_Nest(Mat Y,PetscScalar a,Mat X)1946 PETSC_INTERN PetscErrorCode MatAXPY_Dense_Nest(Mat Y, PetscScalar a, Mat X)
1947 {
1948   Mat_Nest *nest = (Mat_Nest *)X->data;
1949   PetscInt  i, j, k, rstart;
1950   PetscBool flg;
1951 
1952   PetscFunctionBegin;
1953   /* Fill by row */
1954   for (j = 0; j < nest->nc; ++j) {
1955     /* Using global column indices and ISAllGather() is not scalable. */
1956     IS              bNis;
1957     PetscInt        bN;
1958     const PetscInt *bNindices;
1959     PetscCall(ISAllGather(nest->isglobal.col[j], &bNis));
1960     PetscCall(ISGetSize(bNis, &bN));
1961     PetscCall(ISGetIndices(bNis, &bNindices));
1962     for (i = 0; i < nest->nr; ++i) {
1963       Mat             B = nest->m[i][j], D = NULL;
1964       PetscInt        bm, br;
1965       const PetscInt *bmindices;
1966       if (!B) continue;
1967       PetscCall(PetscObjectTypeCompareAny((PetscObject)B, &flg, MATTRANSPOSEVIRTUAL, MATHERMITIANTRANSPOSEVIRTUAL, ""));
1968       if (flg) {
1969         PetscTryMethod(B, "MatTransposeGetMat_C", (Mat, Mat *), (B, &D));
1970         PetscTryMethod(B, "MatHermitianTransposeGetMat_C", (Mat, Mat *), (B, &D));
1971         PetscCall(MatConvert(B, ((PetscObject)D)->type_name, MAT_INITIAL_MATRIX, &D));
1972         B = D;
1973       }
1974       PetscCall(PetscObjectTypeCompareAny((PetscObject)B, &flg, MATSEQSBAIJ, MATMPISBAIJ, ""));
1975       if (flg) {
1976         if (D) PetscCall(MatConvert(D, MATBAIJ, MAT_INPLACE_MATRIX, &D));
1977         else PetscCall(MatConvert(B, MATBAIJ, MAT_INITIAL_MATRIX, &D));
1978         B = D;
1979       }
1980       PetscCall(ISGetLocalSize(nest->isglobal.row[i], &bm));
1981       PetscCall(ISGetIndices(nest->isglobal.row[i], &bmindices));
1982       PetscCall(MatGetOwnershipRange(B, &rstart, NULL));
1983       for (br = 0; br < bm; ++br) {
1984         PetscInt           row = bmindices[br], brncols, *cols;
1985         const PetscInt    *brcols;
1986         const PetscScalar *brcoldata;
1987         PetscScalar       *vals = NULL;
1988         PetscCall(MatGetRow(B, br + rstart, &brncols, &brcols, &brcoldata));
1989         PetscCall(PetscMalloc1(brncols, &cols));
1990         for (k = 0; k < brncols; k++) cols[k] = bNindices[brcols[k]];
1991         /*
1992           Nest blocks are required to be nonoverlapping -- otherwise nest and monolithic index layouts wouldn't match.
1993           Thus, we could use INSERT_VALUES, but I prefer ADD_VALUES.
1994          */
1995         if (a != 1.0) {
1996           PetscCall(PetscMalloc1(brncols, &vals));
1997           for (k = 0; k < brncols; k++) vals[k] = a * brcoldata[k];
1998           PetscCall(MatSetValues(Y, 1, &row, brncols, cols, vals, ADD_VALUES));
1999           PetscCall(PetscFree(vals));
2000         } else {
2001           PetscCall(MatSetValues(Y, 1, &row, brncols, cols, brcoldata, ADD_VALUES));
2002         }
2003         PetscCall(MatRestoreRow(B, br + rstart, &brncols, &brcols, &brcoldata));
2004         PetscCall(PetscFree(cols));
2005       }
2006       if (D) PetscCall(MatDestroy(&D));
2007       PetscCall(ISRestoreIndices(nest->isglobal.row[i], &bmindices));
2008     }
2009     PetscCall(ISRestoreIndices(bNis, &bNindices));
2010     PetscCall(ISDestroy(&bNis));
2011   }
2012   PetscCall(MatAssemblyBegin(Y, MAT_FINAL_ASSEMBLY));
2013   PetscCall(MatAssemblyEnd(Y, MAT_FINAL_ASSEMBLY));
2014   PetscFunctionReturn(PETSC_SUCCESS);
2015 }
2016 
MatConvert_Nest_AIJ(Mat A,MatType newtype,MatReuse reuse,Mat * newmat)2017 static PetscErrorCode MatConvert_Nest_AIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
2018 {
2019   Mat_Nest   *nest = (Mat_Nest *)A->data;
2020   PetscInt    m, n, M, N, i, j, k, *dnnz, *onnz = NULL, rstart, cstart, cend;
2021   PetscMPIInt size;
2022   Mat         C;
2023 
2024   PetscFunctionBegin;
2025   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size));
2026   if (size == 1) { /* look for a special case with SeqAIJ matrices and strided-1, contiguous, blocks */
2027     PetscInt  nf;
2028     PetscBool fast;
2029 
2030     PetscCall(PetscStrcmp(newtype, MATAIJ, &fast));
2031     if (!fast) PetscCall(PetscStrcmp(newtype, MATSEQAIJ, &fast));
2032     for (i = 0; i < nest->nr && fast; ++i) {
2033       for (j = 0; j < nest->nc && fast; ++j) {
2034         Mat B = nest->m[i][j];
2035         if (B) {
2036           PetscCall(PetscObjectTypeCompare((PetscObject)B, MATSEQAIJ, &fast));
2037           if (!fast) {
2038             PetscBool istrans;
2039 
2040             PetscCall(PetscObjectTypeCompare((PetscObject)B, MATTRANSPOSEVIRTUAL, &istrans));
2041             if (istrans) {
2042               Mat Bt;
2043 
2044               PetscCall(MatTransposeGetMat(B, &Bt));
2045               PetscCall(PetscObjectTypeCompare((PetscObject)Bt, MATSEQAIJ, &fast));
2046             } else {
2047               PetscCall(PetscObjectTypeCompare((PetscObject)B, MATHERMITIANTRANSPOSEVIRTUAL, &istrans));
2048               if (istrans) {
2049                 Mat Bt;
2050 
2051                 PetscCall(MatHermitianTransposeGetMat(B, &Bt));
2052                 PetscCall(PetscObjectTypeCompare((PetscObject)Bt, MATSEQAIJ, &fast));
2053               }
2054             }
2055             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);
2056           }
2057         }
2058       }
2059     }
2060     for (i = 0, nf = 0; i < nest->nr && fast; ++i) {
2061       PetscCall(PetscObjectTypeCompare((PetscObject)nest->isglobal.row[i], ISSTRIDE, &fast));
2062       if (fast) {
2063         PetscInt f, s;
2064 
2065         PetscCall(ISStrideGetInfo(nest->isglobal.row[i], &f, &s));
2066         if (f != nf || s != 1) {
2067           fast = PETSC_FALSE;
2068         } else {
2069           PetscCall(ISGetSize(nest->isglobal.row[i], &f));
2070           nf += f;
2071         }
2072       }
2073     }
2074     for (i = 0, nf = 0; i < nest->nc && fast; ++i) {
2075       PetscCall(PetscObjectTypeCompare((PetscObject)nest->isglobal.col[i], ISSTRIDE, &fast));
2076       if (fast) {
2077         PetscInt f, s;
2078 
2079         PetscCall(ISStrideGetInfo(nest->isglobal.col[i], &f, &s));
2080         if (f != nf || s != 1) {
2081           fast = PETSC_FALSE;
2082         } else {
2083           PetscCall(ISGetSize(nest->isglobal.col[i], &f));
2084           nf += f;
2085         }
2086       }
2087     }
2088     if (fast) {
2089       PetscCall(MatConvert_Nest_SeqAIJ_fast(A, newtype, reuse, newmat));
2090       PetscFunctionReturn(PETSC_SUCCESS);
2091     }
2092   }
2093   PetscCall(MatGetSize(A, &M, &N));
2094   PetscCall(MatGetLocalSize(A, &m, &n));
2095   PetscCall(MatGetOwnershipRangeColumn(A, &cstart, &cend));
2096   if (reuse == MAT_REUSE_MATRIX) C = *newmat;
2097   else {
2098     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &C));
2099     PetscCall(MatSetType(C, newtype));
2100     PetscCall(MatSetSizes(C, m, n, M, N));
2101   }
2102   PetscCall(PetscMalloc1(2 * m, &dnnz));
2103   if (m) {
2104     onnz = dnnz + m;
2105     for (k = 0; k < m; k++) {
2106       dnnz[k] = 0;
2107       onnz[k] = 0;
2108     }
2109   }
2110   for (j = 0; j < nest->nc; ++j) {
2111     IS              bNis;
2112     PetscInt        bN;
2113     const PetscInt *bNindices;
2114     PetscBool       flg;
2115     /* Using global column indices and ISAllGather() is not scalable. */
2116     PetscCall(ISAllGather(nest->isglobal.col[j], &bNis));
2117     PetscCall(ISGetSize(bNis, &bN));
2118     PetscCall(ISGetIndices(bNis, &bNindices));
2119     for (i = 0; i < nest->nr; ++i) {
2120       PetscSF         bmsf;
2121       PetscSFNode    *iremote;
2122       Mat             B = nest->m[i][j], D = NULL;
2123       PetscInt        bm, *sub_dnnz, *sub_onnz, br;
2124       const PetscInt *bmindices;
2125       if (!B) continue;
2126       PetscCall(ISGetLocalSize(nest->isglobal.row[i], &bm));
2127       PetscCall(ISGetIndices(nest->isglobal.row[i], &bmindices));
2128       PetscCall(PetscSFCreate(PetscObjectComm((PetscObject)A), &bmsf));
2129       PetscCall(PetscMalloc1(bm, &iremote));
2130       PetscCall(PetscMalloc1(bm, &sub_dnnz));
2131       PetscCall(PetscMalloc1(bm, &sub_onnz));
2132       for (k = 0; k < bm; ++k) {
2133         sub_dnnz[k] = 0;
2134         sub_onnz[k] = 0;
2135       }
2136       PetscCall(PetscObjectTypeCompareAny((PetscObject)B, &flg, MATTRANSPOSEVIRTUAL, MATHERMITIANTRANSPOSEVIRTUAL, ""));
2137       if (flg) {
2138         PetscTryMethod(B, "MatTransposeGetMat_C", (Mat, Mat *), (B, &D));
2139         PetscTryMethod(B, "MatHermitianTransposeGetMat_C", (Mat, Mat *), (B, &D));
2140         PetscCall(MatConvert(B, ((PetscObject)D)->type_name, MAT_INITIAL_MATRIX, &D));
2141         B = D;
2142       }
2143       PetscCall(PetscObjectTypeCompareAny((PetscObject)B, &flg, MATSEQSBAIJ, MATMPISBAIJ, ""));
2144       if (flg) {
2145         if (D) PetscCall(MatConvert(D, MATBAIJ, MAT_INPLACE_MATRIX, &D));
2146         else PetscCall(MatConvert(B, MATBAIJ, MAT_INITIAL_MATRIX, &D));
2147         B = D;
2148       }
2149       /*
2150        Locate the owners for all of the locally-owned global row indices for this row block.
2151        These determine the roots of PetscSF used to communicate preallocation data to row owners.
2152        The roots correspond to the dnnz and onnz entries; thus, there are two roots per row.
2153        */
2154       PetscCall(MatGetOwnershipRange(B, &rstart, NULL));
2155       for (br = 0; br < bm; ++br) {
2156         PetscInt        row = bmindices[br], brncols, col;
2157         const PetscInt *brcols;
2158         PetscInt        rowrel   = 0; /* row's relative index on its owner rank */
2159         PetscMPIInt     rowowner = 0;
2160         PetscCall(PetscLayoutFindOwnerIndex(A->rmap, row, &rowowner, &rowrel));
2161         /* how many roots  */
2162         iremote[br].rank  = rowowner;
2163         iremote[br].index = rowrel; /* edge from bmdnnz to dnnz */
2164         /* get nonzero pattern */
2165         PetscCall(MatGetRow(B, br + rstart, &brncols, &brcols, NULL));
2166         for (k = 0; k < brncols; k++) {
2167           col = bNindices[brcols[k]];
2168           if (col >= A->cmap->range[rowowner] && col < A->cmap->range[rowowner + 1]) {
2169             sub_dnnz[br]++;
2170           } else {
2171             sub_onnz[br]++;
2172           }
2173         }
2174         PetscCall(MatRestoreRow(B, br + rstart, &brncols, &brcols, NULL));
2175       }
2176       if (D) PetscCall(MatDestroy(&D));
2177       PetscCall(ISRestoreIndices(nest->isglobal.row[i], &bmindices));
2178       /* bsf will have to take care of disposing of bedges. */
2179       PetscCall(PetscSFSetGraph(bmsf, m, bm, NULL, PETSC_OWN_POINTER, iremote, PETSC_OWN_POINTER));
2180       PetscCall(PetscSFReduceBegin(bmsf, MPIU_INT, sub_dnnz, dnnz, MPI_SUM));
2181       PetscCall(PetscSFReduceEnd(bmsf, MPIU_INT, sub_dnnz, dnnz, MPI_SUM));
2182       PetscCall(PetscSFReduceBegin(bmsf, MPIU_INT, sub_onnz, onnz, MPI_SUM));
2183       PetscCall(PetscSFReduceEnd(bmsf, MPIU_INT, sub_onnz, onnz, MPI_SUM));
2184       PetscCall(PetscFree(sub_dnnz));
2185       PetscCall(PetscFree(sub_onnz));
2186       PetscCall(PetscSFDestroy(&bmsf));
2187     }
2188     PetscCall(ISRestoreIndices(bNis, &bNindices));
2189     PetscCall(ISDestroy(&bNis));
2190   }
2191   /* Resize preallocation if overestimated */
2192   for (i = 0; i < m; i++) {
2193     dnnz[i] = PetscMin(dnnz[i], A->cmap->n);
2194     onnz[i] = PetscMin(onnz[i], A->cmap->N - A->cmap->n);
2195   }
2196   PetscCall(MatSeqAIJSetPreallocation(C, 0, dnnz));
2197   PetscCall(MatMPIAIJSetPreallocation(C, 0, dnnz, 0, onnz));
2198   PetscCall(PetscFree(dnnz));
2199   PetscCall(MatAXPY_Dense_Nest(C, 1.0, A));
2200   if (reuse == MAT_INPLACE_MATRIX) PetscCall(MatHeaderReplace(A, &C));
2201   else *newmat = C;
2202   PetscFunctionReturn(PETSC_SUCCESS);
2203 }
2204 
MatConvert_Nest_Dense(Mat A,MatType newtype,MatReuse reuse,Mat * newmat)2205 static PetscErrorCode MatConvert_Nest_Dense(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
2206 {
2207   Mat      B;
2208   PetscInt m, n, M, N;
2209 
2210   PetscFunctionBegin;
2211   PetscCall(MatGetSize(A, &M, &N));
2212   PetscCall(MatGetLocalSize(A, &m, &n));
2213   if (reuse == MAT_REUSE_MATRIX) {
2214     B = *newmat;
2215     PetscCall(MatZeroEntries(B));
2216   } else {
2217     PetscCall(MatCreateDense(PetscObjectComm((PetscObject)A), m, PETSC_DECIDE, M, N, NULL, &B));
2218   }
2219   PetscCall(MatAXPY_Dense_Nest(B, 1.0, A));
2220   if (reuse == MAT_INPLACE_MATRIX) PetscCall(MatHeaderReplace(A, &B));
2221   else if (reuse == MAT_INITIAL_MATRIX) *newmat = B;
2222   PetscFunctionReturn(PETSC_SUCCESS);
2223 }
2224 
MatHasOperation_Nest(Mat mat,MatOperation op,PetscBool * has)2225 static PetscErrorCode MatHasOperation_Nest(Mat mat, MatOperation op, PetscBool *has)
2226 {
2227   Mat_Nest    *bA = (Mat_Nest *)mat->data;
2228   MatOperation opAdd;
2229   PetscInt     i, j, nr = bA->nr, nc = bA->nc;
2230   PetscBool    flg;
2231 
2232   PetscFunctionBegin;
2233   *has = PETSC_FALSE;
2234   if (op == MATOP_MULT || op == MATOP_MULT_ADD || op == MATOP_MULT_TRANSPOSE || op == MATOP_MULT_TRANSPOSE_ADD) {
2235     opAdd = (op == MATOP_MULT || op == MATOP_MULT_ADD ? MATOP_MULT_ADD : MATOP_MULT_TRANSPOSE_ADD);
2236     for (j = 0; j < nc; j++) {
2237       for (i = 0; i < nr; i++) {
2238         if (!bA->m[i][j]) continue;
2239         PetscCall(MatHasOperation(bA->m[i][j], opAdd, &flg));
2240         if (!flg) PetscFunctionReturn(PETSC_SUCCESS);
2241       }
2242     }
2243   }
2244   if (((void **)mat->ops)[op]) *has = PETSC_TRUE;
2245   PetscFunctionReturn(PETSC_SUCCESS);
2246 }
2247 
2248 /*MC
2249   MATNEST -  "nest" - Matrix type consisting of nested submatrices, each stored separately.
2250 
2251   Level: intermediate
2252 
2253   Notes:
2254   This matrix type permits scalable use of `PCFIELDSPLIT` and avoids the large memory costs of extracting submatrices.
2255   It allows the use of symmetric and block formats for parts of multi-physics simulations.
2256   It is usually used with `DMCOMPOSITE` and `DMCreateMatrix()`
2257 
2258   Each of the submatrices lives on the same MPI communicator as the original nest matrix (though they can have zero
2259   rows/columns on some processes.) Thus this is not meant for cases where the submatrices live on far fewer processes
2260   than the nest matrix.
2261 
2262 .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatCreate()`, `MatType`, `MatCreateNest()`, `MatNestSetSubMat()`, `MatNestGetSubMat()`,
2263           `VecCreateNest()`, `DMCreateMatrix()`, `DMCOMPOSITE`, `MatNestSetVecType()`, `MatNestGetLocalISs()`,
2264           `MatNestGetISs()`, `MatNestSetSubMats()`, `MatNestGetSubMats()`
2265 M*/
MatCreate_Nest(Mat A)2266 PETSC_EXTERN PetscErrorCode MatCreate_Nest(Mat A)
2267 {
2268   Mat_Nest *s;
2269 
2270   PetscFunctionBegin;
2271   PetscCall(PetscNew(&s));
2272   A->data = (void *)s;
2273 
2274   s->nr            = -1;
2275   s->nc            = -1;
2276   s->m             = NULL;
2277   s->splitassembly = PETSC_FALSE;
2278 
2279   PetscCall(PetscMemzero(A->ops, sizeof(*A->ops)));
2280 
2281   A->ops->mult                      = MatMult_Nest;
2282   A->ops->multadd                   = MatMultAdd_Nest;
2283   A->ops->multtranspose             = MatMultTranspose_Nest;
2284   A->ops->multtransposeadd          = MatMultTransposeAdd_Nest;
2285   A->ops->transpose                 = MatTranspose_Nest;
2286   A->ops->multhermitiantranspose    = MatMultHermitianTranspose_Nest;
2287   A->ops->multhermitiantransposeadd = MatMultHermitianTransposeAdd_Nest;
2288   A->ops->assemblybegin             = MatAssemblyBegin_Nest;
2289   A->ops->assemblyend               = MatAssemblyEnd_Nest;
2290   A->ops->zeroentries               = MatZeroEntries_Nest;
2291   A->ops->copy                      = MatCopy_Nest;
2292   A->ops->axpy                      = MatAXPY_Nest;
2293   A->ops->duplicate                 = MatDuplicate_Nest;
2294   A->ops->createsubmatrix           = MatCreateSubMatrix_Nest;
2295   A->ops->destroy                   = MatDestroy_Nest;
2296   A->ops->view                      = MatView_Nest;
2297   A->ops->getvecs                   = NULL; /* Use VECNEST by calling MatNestSetVecType(A,VECNEST) */
2298   A->ops->getlocalsubmatrix         = MatGetLocalSubMatrix_Nest;
2299   A->ops->restorelocalsubmatrix     = MatRestoreLocalSubMatrix_Nest;
2300   A->ops->getdiagonal               = MatGetDiagonal_Nest;
2301   A->ops->diagonalscale             = MatDiagonalScale_Nest;
2302   A->ops->scale                     = MatScale_Nest;
2303   A->ops->shift                     = MatShift_Nest;
2304   A->ops->diagonalset               = MatDiagonalSet_Nest;
2305   A->ops->setrandom                 = MatSetRandom_Nest;
2306   A->ops->hasoperation              = MatHasOperation_Nest;
2307 
2308   A->spptr     = NULL;
2309   A->assembled = PETSC_FALSE;
2310 
2311   /* expose Nest api's */
2312   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetSubMat_C", MatNestGetSubMat_Nest));
2313   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestSetSubMat_C", MatNestSetSubMat_Nest));
2314   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetSubMats_C", MatNestGetSubMats_Nest));
2315   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetSize_C", MatNestGetSize_Nest));
2316   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetISs_C", MatNestGetISs_Nest));
2317   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetLocalISs_C", MatNestGetLocalISs_Nest));
2318   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestSetVecType_C", MatNestSetVecType_Nest));
2319   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestSetSubMats_C", MatNestSetSubMats_Nest));
2320   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_mpiaij_C", MatConvert_Nest_AIJ));
2321   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_seqaij_C", MatConvert_Nest_AIJ));
2322   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_aij_C", MatConvert_Nest_AIJ));
2323   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_is_C", MatConvert_Nest_IS));
2324   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_mpidense_C", MatConvert_Nest_Dense));
2325   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_seqdense_C", MatConvert_Nest_Dense));
2326   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_nest_seqdense_C", MatProductSetFromOptions_Nest_Dense));
2327   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_nest_mpidense_C", MatProductSetFromOptions_Nest_Dense));
2328 
2329   PetscCall(PetscObjectChangeTypeName((PetscObject)A, MATNEST));
2330   PetscFunctionReturn(PETSC_SUCCESS);
2331 }
2332