Actual source code: mpibaij.c

  1: #include <../src/mat/impls/baij/mpi/mpibaij.h>

  3: #include <petsc/private/hashseti.h>
  4: #include <petscblaslapack.h>
  5: #include <petscsf.h>

  7: static PetscErrorCode MatDestroy_MPIBAIJ(Mat mat)
  8: {
  9:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)mat->data;

 11:   PetscFunctionBegin;
 12:   PetscCall(PetscLogObjectState((PetscObject)mat, "Rows=%" PetscInt_FMT ",Cols=%" PetscInt_FMT, mat->rmap->N, mat->cmap->N));
 13:   PetscCall(MatStashDestroy_Private(&mat->stash));
 14:   PetscCall(MatStashDestroy_Private(&mat->bstash));
 15:   PetscCall(MatDestroy(&baij->A));
 16:   PetscCall(MatDestroy(&baij->B));
 17: #if defined(PETSC_USE_CTABLE)
 18:   PetscCall(PetscHMapIDestroy(&baij->colmap));
 19: #else
 20:   PetscCall(PetscFree(baij->colmap));
 21: #endif
 22:   PetscCall(PetscFree(baij->garray));
 23:   PetscCall(VecDestroy(&baij->lvec));
 24:   PetscCall(VecScatterDestroy(&baij->Mvctx));
 25:   PetscCall(PetscFree2(baij->rowvalues, baij->rowindices));
 26:   PetscCall(PetscFree(baij->barray));
 27:   PetscCall(PetscFree2(baij->hd, baij->ht));
 28:   PetscCall(PetscFree(baij->rangebs));
 29:   PetscCall(PetscFree(mat->data));

 31:   PetscCall(PetscObjectChangeTypeName((PetscObject)mat, NULL));
 32:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatStoreValues_C", NULL));
 33:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatRetrieveValues_C", NULL));
 34:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatGetMultPetscSF_C", NULL));
 35:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMPIBAIJSetPreallocation_C", NULL));
 36:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMPIBAIJSetPreallocationCSR_C", NULL));
 37:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDiagonalScaleLocal_C", NULL));
 38:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatSetHashTableFactor_C", NULL));
 39:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpibaij_mpisbaij_C", NULL));
 40:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpibaij_mpiadj_C", NULL));
 41:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpibaij_mpiaij_C", NULL));
 42: #if defined(PETSC_HAVE_HYPRE)
 43:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpibaij_hypre_C", NULL));
 44: #endif
 45:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpibaij_is_C", NULL));
 46:   PetscFunctionReturn(PETSC_SUCCESS);
 47: }

 49: /* defines MatSetValues_MPI_Hash(), MatAssemblyBegin_MPI_Hash(), and  MatAssemblyEnd_MPI_Hash() */
 50: #define TYPE BAIJ
 51: #include "../src/mat/impls/aij/mpi/mpihashmat.h"
 52: #undef TYPE

 54: #if defined(PETSC_HAVE_HYPRE)
 55: PETSC_INTERN PetscErrorCode MatConvert_AIJ_HYPRE(Mat, MatType, MatReuse, Mat *);
 56: #endif

 58: static PetscErrorCode MatGetRowMaxAbs_MPIBAIJ(Mat A, Vec v, PetscInt idx[])
 59: {
 60:   Mat_MPIBAIJ       *a = (Mat_MPIBAIJ *)A->data;
 61:   PetscInt           i, *idxb = NULL, m = A->rmap->n, bs = A->cmap->bs;
 62:   PetscScalar       *vv;
 63:   Vec                vB, vA;
 64:   const PetscScalar *va, *vb;

 66:   PetscFunctionBegin;
 67:   PetscCall(MatCreateVecs(a->A, NULL, &vA));
 68:   PetscCall(MatGetRowMaxAbs(a->A, vA, idx));

 70:   PetscCall(VecGetArrayRead(vA, &va));
 71:   if (idx) {
 72:     for (i = 0; i < m; i++) {
 73:       if (PetscAbsScalar(va[i])) idx[i] += A->cmap->rstart;
 74:     }
 75:   }

 77:   PetscCall(MatCreateVecs(a->B, NULL, &vB));
 78:   PetscCall(PetscMalloc1(m, &idxb));
 79:   PetscCall(MatGetRowMaxAbs(a->B, vB, idxb));

 81:   PetscCall(VecGetArrayWrite(v, &vv));
 82:   PetscCall(VecGetArrayRead(vB, &vb));
 83:   for (i = 0; i < m; i++) {
 84:     if (PetscAbsScalar(va[i]) < PetscAbsScalar(vb[i])) {
 85:       vv[i] = vb[i];
 86:       if (idx) idx[i] = bs * a->garray[idxb[i] / bs] + (idxb[i] % bs);
 87:     } else {
 88:       vv[i] = va[i];
 89:       if (idx && PetscAbsScalar(va[i]) == PetscAbsScalar(vb[i]) && idxb[i] != -1 && idx[i] > bs * a->garray[idxb[i] / bs] + (idxb[i] % bs)) idx[i] = bs * a->garray[idxb[i] / bs] + (idxb[i] % bs);
 90:     }
 91:   }
 92:   PetscCall(VecRestoreArrayWrite(v, &vv));
 93:   PetscCall(VecRestoreArrayRead(vA, &va));
 94:   PetscCall(VecRestoreArrayRead(vB, &vb));
 95:   PetscCall(PetscFree(idxb));
 96:   PetscCall(VecDestroy(&vA));
 97:   PetscCall(VecDestroy(&vB));
 98:   PetscFunctionReturn(PETSC_SUCCESS);
 99: }

101: static PetscErrorCode MatGetRowSumAbs_MPIBAIJ(Mat A, Vec v)
102: {
103:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;
104:   Vec          vB, vA;

106:   PetscFunctionBegin;
107:   PetscCall(MatCreateVecs(a->A, NULL, &vA));
108:   PetscCall(MatGetRowSumAbs(a->A, vA));
109:   PetscCall(MatCreateVecs(a->B, NULL, &vB));
110:   PetscCall(MatGetRowSumAbs(a->B, vB));
111:   PetscCall(VecAXPY(vA, 1.0, vB));
112:   PetscCall(VecDestroy(&vB));
113:   PetscCall(VecCopy(vA, v));
114:   PetscCall(VecDestroy(&vA));
115:   PetscFunctionReturn(PETSC_SUCCESS);
116: }

118: static PetscErrorCode MatStoreValues_MPIBAIJ(Mat mat)
119: {
120:   Mat_MPIBAIJ *aij = (Mat_MPIBAIJ *)mat->data;

122:   PetscFunctionBegin;
123:   PetscCall(MatStoreValues(aij->A));
124:   PetscCall(MatStoreValues(aij->B));
125:   PetscFunctionReturn(PETSC_SUCCESS);
126: }

128: static PetscErrorCode MatRetrieveValues_MPIBAIJ(Mat mat)
129: {
130:   Mat_MPIBAIJ *aij = (Mat_MPIBAIJ *)mat->data;

132:   PetscFunctionBegin;
133:   PetscCall(MatRetrieveValues(aij->A));
134:   PetscCall(MatRetrieveValues(aij->B));
135:   PetscFunctionReturn(PETSC_SUCCESS);
136: }

138: /*
139:      Local utility routine that creates a mapping from the global column
140:    number to the local number in the off-diagonal part of the local
141:    storage of the matrix.  This is done in a non scalable way since the
142:    length of colmap equals the global matrix length.
143: */
144: PetscErrorCode MatCreateColmap_MPIBAIJ_Private(Mat mat)
145: {
146:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)mat->data;
147:   Mat_SeqBAIJ *B    = (Mat_SeqBAIJ *)baij->B->data;
148:   PetscInt     nbs = B->nbs, i, bs = mat->rmap->bs;

150:   PetscFunctionBegin;
151: #if defined(PETSC_USE_CTABLE)
152:   PetscCall(PetscHMapICreateWithSize(baij->nbs, &baij->colmap));
153:   for (i = 0; i < nbs; i++) PetscCall(PetscHMapISet(baij->colmap, baij->garray[i] + 1, i * bs + 1));
154: #else
155:   PetscCall(PetscCalloc1(baij->Nbs + 1, &baij->colmap));
156:   for (i = 0; i < nbs; i++) baij->colmap[baij->garray[i]] = i * bs + 1;
157: #endif
158:   PetscFunctionReturn(PETSC_SUCCESS);
159: }

161: #define MatSetValues_SeqBAIJ_A_Private(row, col, value, addv, orow, ocol) \
162:   do { \
163:     brow = row / bs; \
164:     rp   = PetscSafePointerPlusOffset(aj, ai[brow]); \
165:     ap   = PetscSafePointerPlusOffset(aa, bs2 * ai[brow]); \
166:     rmax = aimax[brow]; \
167:     nrow = ailen[brow]; \
168:     bcol = col / bs; \
169:     ridx = row % bs; \
170:     cidx = col % bs; \
171:     low  = 0; \
172:     high = nrow; \
173:     while (high - low > 3) { \
174:       t = (low + high) / 2; \
175:       if (rp[t] > bcol) high = t; \
176:       else low = t; \
177:     } \
178:     for (_i = low; _i < high; _i++) { \
179:       if (rp[_i] > bcol) break; \
180:       if (rp[_i] == bcol) { \
181:         bap = ap + bs2 * _i + bs * cidx + ridx; \
182:         if (addv == ADD_VALUES) *bap += value; \
183:         else *bap = value; \
184:         goto a_noinsert; \
185:       } \
186:     } \
187:     if (a->nonew == 1) goto a_noinsert; \
188:     PetscCheck(a->nonew != -1, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Inserting a new nonzero at global row/column (%" PetscInt_FMT ", %" PetscInt_FMT ") into matrix", orow, ocol); \
189:     MatSeqXAIJReallocateAIJ(A, a->mbs, bs2, nrow, brow, bcol, rmax, aa, ai, aj, rp, ap, aimax, a->nonew, MatScalar); \
190:     N = nrow++ - 1; \
191:     /* shift up all the later entries in this row */ \
192:     PetscCall(PetscArraymove(rp + _i + 1, rp + _i, N - _i + 1)); \
193:     PetscCall(PetscArraymove(ap + bs2 * (_i + 1), ap + bs2 * _i, bs2 * (N - _i + 1))); \
194:     PetscCall(PetscArrayzero(ap + bs2 * _i, bs2)); \
195:     rp[_i]                          = bcol; \
196:     ap[bs2 * _i + bs * cidx + ridx] = value; \
197:   a_noinsert:; \
198:     ailen[brow] = nrow; \
199:   } while (0)

201: #define MatSetValues_SeqBAIJ_B_Private(row, col, value, addv, orow, ocol) \
202:   do { \
203:     brow = row / bs; \
204:     rp   = PetscSafePointerPlusOffset(bj, bi[brow]); \
205:     ap   = PetscSafePointerPlusOffset(ba, bs2 * bi[brow]); \
206:     rmax = bimax[brow]; \
207:     nrow = bilen[brow]; \
208:     bcol = col / bs; \
209:     ridx = row % bs; \
210:     cidx = col % bs; \
211:     low  = 0; \
212:     high = nrow; \
213:     while (high - low > 3) { \
214:       t = (low + high) / 2; \
215:       if (rp[t] > bcol) high = t; \
216:       else low = t; \
217:     } \
218:     for (_i = low; _i < high; _i++) { \
219:       if (rp[_i] > bcol) break; \
220:       if (rp[_i] == bcol) { \
221:         bap = ap + bs2 * _i + bs * cidx + ridx; \
222:         if (addv == ADD_VALUES) *bap += value; \
223:         else *bap = value; \
224:         goto b_noinsert; \
225:       } \
226:     } \
227:     if (b->nonew == 1) goto b_noinsert; \
228:     PetscCheck(b->nonew != -1, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Inserting a new nonzero at global row/column  (%" PetscInt_FMT ", %" PetscInt_FMT ") into matrix", orow, ocol); \
229:     MatSeqXAIJReallocateAIJ(B, b->mbs, bs2, nrow, brow, bcol, rmax, ba, bi, bj, rp, ap, bimax, b->nonew, MatScalar); \
230:     N = nrow++ - 1; \
231:     /* shift up all the later entries in this row */ \
232:     PetscCall(PetscArraymove(rp + _i + 1, rp + _i, N - _i + 1)); \
233:     PetscCall(PetscArraymove(ap + bs2 * (_i + 1), ap + bs2 * _i, bs2 * (N - _i + 1))); \
234:     PetscCall(PetscArrayzero(ap + bs2 * _i, bs2)); \
235:     rp[_i]                          = bcol; \
236:     ap[bs2 * _i + bs * cidx + ridx] = value; \
237:   b_noinsert:; \
238:     bilen[brow] = nrow; \
239:   } while (0)

241: PetscErrorCode MatSetValues_MPIBAIJ(Mat mat, PetscInt m, const PetscInt im[], PetscInt n, const PetscInt in[], const PetscScalar v[], InsertMode addv)
242: {
243:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)mat->data;
244:   MatScalar    value;
245:   PetscBool    roworiented = baij->roworiented;
246:   PetscInt     i, j, row, col;
247:   PetscInt     rstart_orig = mat->rmap->rstart;
248:   PetscInt     rend_orig = mat->rmap->rend, cstart_orig = mat->cmap->rstart;
249:   PetscInt     cend_orig = mat->cmap->rend, bs = mat->rmap->bs;

251:   /* Some Variables required in the macro */
252:   Mat          A     = baij->A;
253:   Mat_SeqBAIJ *a     = (Mat_SeqBAIJ *)A->data;
254:   PetscInt    *aimax = a->imax, *ai = a->i, *ailen = a->ilen, *aj = a->j;
255:   MatScalar   *aa = a->a;

257:   Mat          B     = baij->B;
258:   Mat_SeqBAIJ *b     = (Mat_SeqBAIJ *)B->data;
259:   PetscInt    *bimax = b->imax, *bi = b->i, *bilen = b->ilen, *bj = b->j;
260:   MatScalar   *ba = b->a;

262:   PetscInt  *rp, ii, nrow, _i, rmax, N, brow, bcol;
263:   PetscInt   low, high, t, ridx, cidx, bs2 = a->bs2;
264:   MatScalar *ap, *bap;

266:   PetscFunctionBegin;
267:   for (i = 0; i < m; i++) {
268:     if (im[i] < 0) continue;
269:     PetscCheck(im[i] < mat->rmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Row too large: row %" PetscInt_FMT " max %" PetscInt_FMT, im[i], mat->rmap->N - 1);
270:     if (im[i] >= rstart_orig && im[i] < rend_orig) {
271:       row = im[i] - rstart_orig;
272:       for (j = 0; j < n; j++) {
273:         if (in[j] >= cstart_orig && in[j] < cend_orig) {
274:           col = in[j] - cstart_orig;
275:           if (roworiented) value = v[i * n + j];
276:           else value = v[i + j * m];
277:           MatSetValues_SeqBAIJ_A_Private(row, col, value, addv, im[i], in[j]);
278:         } else if (in[j] < 0) {
279:           continue;
280:         } else {
281:           PetscCheck(in[j] < mat->cmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Column too large: col %" PetscInt_FMT " max %" PetscInt_FMT, in[j], mat->cmap->N - 1);
282:           if (mat->was_assembled) {
283:             if (!baij->colmap) PetscCall(MatCreateColmap_MPIBAIJ_Private(mat));
284: #if defined(PETSC_USE_CTABLE)
285:             PetscCall(PetscHMapIGetWithDefault(baij->colmap, in[j] / bs + 1, 0, &col));
286:             col = col - 1;
287: #else
288:             col = baij->colmap[in[j] / bs] - 1;
289: #endif
290:             if (col < 0 && !((Mat_SeqBAIJ *)baij->B->data)->nonew) {
291:               PetscCall(MatDisAssemble_MPIBAIJ(mat));
292:               col = in[j];
293:               /* Reinitialize the variables required by MatSetValues_SeqBAIJ_B_Private() */
294:               B     = baij->B;
295:               b     = (Mat_SeqBAIJ *)B->data;
296:               bimax = b->imax;
297:               bi    = b->i;
298:               bilen = b->ilen;
299:               bj    = b->j;
300:               ba    = b->a;
301:             } else {
302:               PetscCheck(col >= 0, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Inserting a new nonzero (%" PetscInt_FMT ", %" PetscInt_FMT ") into matrix", im[i], in[j]);
303:               col += in[j] % bs;
304:             }
305:           } else col = in[j];
306:           if (roworiented) value = v[i * n + j];
307:           else value = v[i + j * m];
308:           MatSetValues_SeqBAIJ_B_Private(row, col, value, addv, im[i], in[j]);
309:           /* PetscCall(MatSetValues_SeqBAIJ(baij->B,1,&row,1,&col,&value,addv)); */
310:         }
311:       }
312:     } else {
313:       PetscCheck(!mat->nooffprocentries, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Setting off process row %" PetscInt_FMT " even though MatSetOption(,MAT_NO_OFF_PROC_ENTRIES,PETSC_TRUE) was set", im[i]);
314:       if (!baij->donotstash) {
315:         mat->assembled = PETSC_FALSE;
316:         if (roworiented) {
317:           PetscCall(MatStashValuesRow_Private(&mat->stash, im[i], n, in, v + i * n, PETSC_FALSE));
318:         } else {
319:           PetscCall(MatStashValuesCol_Private(&mat->stash, im[i], n, in, v + i, m, PETSC_FALSE));
320:         }
321:       }
322:     }
323:   }
324:   PetscFunctionReturn(PETSC_SUCCESS);
325: }

327: static inline PetscErrorCode MatSetValuesBlocked_SeqBAIJ_Inlined(Mat A, PetscInt row, PetscInt col, const PetscScalar v[], InsertMode is, PetscInt orow, PetscInt ocol)
328: {
329:   Mat_SeqBAIJ       *a = (Mat_SeqBAIJ *)A->data;
330:   PetscInt          *rp, low, high, t, ii, jj, nrow, i, rmax, N;
331:   PetscInt          *imax = a->imax, *ai = a->i, *ailen = a->ilen;
332:   PetscInt          *aj = a->j, nonew = a->nonew, bs2 = a->bs2, bs = A->rmap->bs;
333:   PetscBool          roworiented = a->roworiented;
334:   const PetscScalar *value       = v;
335:   MatScalar         *ap, *aa = a->a, *bap;

337:   PetscFunctionBegin;
338:   rp    = aj + ai[row];
339:   ap    = aa + bs2 * ai[row];
340:   rmax  = imax[row];
341:   nrow  = ailen[row];
342:   value = v;
343:   low   = 0;
344:   high  = nrow;
345:   while (high - low > 7) {
346:     t = (low + high) / 2;
347:     if (rp[t] > col) high = t;
348:     else low = t;
349:   }
350:   for (i = low; i < high; i++) {
351:     if (rp[i] > col) break;
352:     if (rp[i] == col) {
353:       bap = ap + bs2 * i;
354:       if (roworiented) {
355:         if (is == ADD_VALUES) {
356:           for (ii = 0; ii < bs; ii++) {
357:             for (jj = ii; jj < bs2; jj += bs) bap[jj] += *value++;
358:           }
359:         } else {
360:           for (ii = 0; ii < bs; ii++) {
361:             for (jj = ii; jj < bs2; jj += bs) bap[jj] = *value++;
362:           }
363:         }
364:       } else {
365:         if (is == ADD_VALUES) {
366:           for (ii = 0; ii < bs; ii++, value += bs) {
367:             for (jj = 0; jj < bs; jj++) bap[jj] += value[jj];
368:             bap += bs;
369:           }
370:         } else {
371:           for (ii = 0; ii < bs; ii++, value += bs) {
372:             for (jj = 0; jj < bs; jj++) bap[jj] = value[jj];
373:             bap += bs;
374:           }
375:         }
376:       }
377:       goto noinsert2;
378:     }
379:   }
380:   if (nonew == 1) goto noinsert2;
381:   PetscCheck(nonew != -1, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Inserting a new global block indexed nonzero block (%" PetscInt_FMT ", %" PetscInt_FMT ") in the matrix", orow, ocol);
382:   MatSeqXAIJReallocateAIJ(A, a->mbs, bs2, nrow, row, col, rmax, aa, ai, aj, rp, ap, imax, nonew, MatScalar);
383:   N = nrow++ - 1;
384:   high++;
385:   /* shift up all the later entries in this row */
386:   PetscCall(PetscArraymove(rp + i + 1, rp + i, N - i + 1));
387:   PetscCall(PetscArraymove(ap + bs2 * (i + 1), ap + bs2 * i, bs2 * (N - i + 1)));
388:   rp[i] = col;
389:   bap   = ap + bs2 * i;
390:   if (roworiented) {
391:     for (ii = 0; ii < bs; ii++) {
392:       for (jj = ii; jj < bs2; jj += bs) bap[jj] = *value++;
393:     }
394:   } else {
395:     for (ii = 0; ii < bs; ii++) {
396:       for (jj = 0; jj < bs; jj++) *bap++ = *value++;
397:     }
398:   }
399: noinsert2:;
400:   ailen[row] = nrow;
401:   PetscFunctionReturn(PETSC_SUCCESS);
402: }

404: /*
405:     This routine should be optimized so that the block copy at ** Here a copy is required ** below is not needed
406:     by passing additional stride information into the MatSetValuesBlocked_SeqBAIJ_Inlined() routine
407: */
408: static PetscErrorCode MatSetValuesBlocked_MPIBAIJ(Mat mat, PetscInt m, const PetscInt im[], PetscInt n, const PetscInt in[], const PetscScalar v[], InsertMode addv)
409: {
410:   Mat_MPIBAIJ       *baij = (Mat_MPIBAIJ *)mat->data;
411:   const PetscScalar *value;
412:   MatScalar         *barray      = baij->barray;
413:   PetscBool          roworiented = baij->roworiented;
414:   PetscInt           i, j, ii, jj, row, col, rstart = baij->rstartbs;
415:   PetscInt           rend = baij->rendbs, cstart = baij->cstartbs, stepval;
416:   PetscInt           cend = baij->cendbs, bs = mat->rmap->bs, bs2 = baij->bs2;

418:   PetscFunctionBegin;
419:   if (!barray) {
420:     PetscCall(PetscMalloc1(bs2, &barray));
421:     baij->barray = barray;
422:   }

424:   if (roworiented) stepval = (n - 1) * bs;
425:   else stepval = (m - 1) * bs;

427:   for (i = 0; i < m; i++) {
428:     if (im[i] < 0) continue;
429:     PetscCheck(im[i] < baij->Mbs, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Block indexed row too large %" PetscInt_FMT " max %" PetscInt_FMT, im[i], baij->Mbs - 1);
430:     if (im[i] >= rstart && im[i] < rend) {
431:       row = im[i] - rstart;
432:       for (j = 0; j < n; j++) {
433:         /* If NumCol = 1 then a copy is not required */
434:         if ((roworiented) && (n == 1)) {
435:           barray = (MatScalar *)v + i * bs2;
436:         } else if ((!roworiented) && (m == 1)) {
437:           barray = (MatScalar *)v + j * bs2;
438:         } else { /* Here a copy is required */
439:           if (roworiented) {
440:             value = v + (i * (stepval + bs) + j) * bs;
441:           } else {
442:             value = v + (j * (stepval + bs) + i) * bs;
443:           }
444:           for (ii = 0; ii < bs; ii++, value += bs + stepval) {
445:             for (jj = 0; jj < bs; jj++) barray[jj] = value[jj];
446:             barray += bs;
447:           }
448:           barray -= bs2;
449:         }

451:         if (in[j] >= cstart && in[j] < cend) {
452:           col = in[j] - cstart;
453:           PetscCall(MatSetValuesBlocked_SeqBAIJ_Inlined(baij->A, row, col, barray, addv, im[i], in[j]));
454:         } else if (in[j] < 0) {
455:           continue;
456:         } else {
457:           PetscCheck(in[j] < baij->Nbs, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Block indexed column too large %" PetscInt_FMT " max %" PetscInt_FMT, in[j], baij->Nbs - 1);
458:           if (mat->was_assembled) {
459:             if (!baij->colmap) PetscCall(MatCreateColmap_MPIBAIJ_Private(mat));

461: #if defined(PETSC_USE_CTABLE)
462:             PetscCall(PetscHMapIGetWithDefault(baij->colmap, in[j] + 1, 0, &col));
463:             col = col < 1 ? -1 : (col - 1) / bs;
464: #else
465:             col = baij->colmap[in[j]] < 1 ? -1 : (baij->colmap[in[j]] - 1) / bs;
466: #endif
467:             if (col < 0 && !((Mat_SeqBAIJ *)baij->B->data)->nonew) {
468:               PetscCall(MatDisAssemble_MPIBAIJ(mat));
469:               col = in[j];
470:             } else PetscCheck(col >= 0, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Inserting a new blocked indexed nonzero block (%" PetscInt_FMT ", %" PetscInt_FMT ") into matrix", im[i], in[j]);
471:           } else col = in[j];
472:           PetscCall(MatSetValuesBlocked_SeqBAIJ_Inlined(baij->B, row, col, barray, addv, im[i], in[j]));
473:         }
474:       }
475:     } else {
476:       PetscCheck(!mat->nooffprocentries, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Setting off process block indexed row %" PetscInt_FMT " even though MatSetOption(,MAT_NO_OFF_PROC_ENTRIES,PETSC_TRUE) was set", im[i]);
477:       if (!baij->donotstash) {
478:         if (roworiented) {
479:           PetscCall(MatStashValuesRowBlocked_Private(&mat->bstash, im[i], n, in, v, m, n, i));
480:         } else {
481:           PetscCall(MatStashValuesColBlocked_Private(&mat->bstash, im[i], n, in, v, m, n, i));
482:         }
483:       }
484:     }
485:   }
486:   PetscFunctionReturn(PETSC_SUCCESS);
487: }

489: #define HASH_KEY             0.6180339887
490: #define HASH(size, key, tmp) (tmp = (key) * HASH_KEY, (PetscInt)((size) * (tmp - (PetscInt)tmp)))
491: /* #define HASH(size,key) ((PetscInt)((size)*fmod(((key)*HASH_KEY),1))) */
492: /* #define HASH(size,key,tmp) ((PetscInt)((size)*fmod(((key)*HASH_KEY),1))) */
493: static PetscErrorCode MatSetValues_MPIBAIJ_HT(Mat mat, PetscInt m, const PetscInt im[], PetscInt n, const PetscInt in[], const PetscScalar v[], InsertMode addv)
494: {
495:   Mat_MPIBAIJ *baij        = (Mat_MPIBAIJ *)mat->data;
496:   PetscBool    roworiented = baij->roworiented;
497:   PetscInt     i, j, row, col;
498:   PetscInt     rstart_orig = mat->rmap->rstart;
499:   PetscInt     rend_orig = mat->rmap->rend, Nbs = baij->Nbs;
500:   PetscInt     h1, key, size = baij->ht_size, bs = mat->rmap->bs, *HT = baij->ht, idx;
501:   PetscReal    tmp;
502:   MatScalar  **HD       = baij->hd, value;
503:   PetscInt     total_ct = baij->ht_total_ct, insert_ct = baij->ht_insert_ct;

505:   PetscFunctionBegin;
506:   for (i = 0; i < m; i++) {
507:     if (PetscDefined(USE_DEBUG)) {
508:       PetscCheck(im[i] >= 0, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Negative row");
509:       PetscCheck(im[i] < mat->rmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Row too large: row %" PetscInt_FMT " max %" PetscInt_FMT, im[i], mat->rmap->N - 1);
510:     }
511:     row = im[i];
512:     if (row >= rstart_orig && row < rend_orig) {
513:       for (j = 0; j < n; j++) {
514:         col = in[j];
515:         if (roworiented) value = v[i * n + j];
516:         else value = v[i + j * m];
517:         /* Look up PetscInto the Hash Table */
518:         key = (row / bs) * Nbs + (col / bs) + 1;
519:         h1  = HASH(size, key, tmp);

521:         idx = h1;
522:         if (PetscDefined(USE_DEBUG)) {
523:           insert_ct++;
524:           total_ct++;
525:           if (HT[idx] != key) {
526:             for (idx = h1; (idx < size) && (HT[idx] != key); idx++, total_ct++);
527:             if (idx == size) {
528:               for (idx = 0; (idx < h1) && (HT[idx] != key); idx++, total_ct++);
529:               PetscCheck(idx != h1, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "(%" PetscInt_FMT ",%" PetscInt_FMT ") has no entry in the hash table", row, col);
530:             }
531:           }
532:         } else if (HT[idx] != key) {
533:           for (idx = h1; (idx < size) && (HT[idx] != key); idx++);
534:           if (idx == size) {
535:             for (idx = 0; (idx < h1) && (HT[idx] != key); idx++);
536:             PetscCheck(idx != h1, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "(%" PetscInt_FMT ",%" PetscInt_FMT ") has no entry in the hash table", row, col);
537:           }
538:         }
539:         /* A HASH table entry is found, so insert the values at the correct address */
540:         if (addv == ADD_VALUES) *(HD[idx] + (col % bs) * bs + (row % bs)) += value;
541:         else *(HD[idx] + (col % bs) * bs + (row % bs)) = value;
542:       }
543:     } else if (!baij->donotstash) {
544:       if (roworiented) {
545:         PetscCall(MatStashValuesRow_Private(&mat->stash, im[i], n, in, v + i * n, PETSC_FALSE));
546:       } else {
547:         PetscCall(MatStashValuesCol_Private(&mat->stash, im[i], n, in, v + i, m, PETSC_FALSE));
548:       }
549:     }
550:   }
551:   if (PetscDefined(USE_DEBUG)) {
552:     baij->ht_total_ct += total_ct;
553:     baij->ht_insert_ct += insert_ct;
554:   }
555:   PetscFunctionReturn(PETSC_SUCCESS);
556: }

558: static PetscErrorCode MatSetValuesBlocked_MPIBAIJ_HT(Mat mat, PetscInt m, const PetscInt im[], PetscInt n, const PetscInt in[], const PetscScalar v[], InsertMode addv)
559: {
560:   Mat_MPIBAIJ       *baij        = (Mat_MPIBAIJ *)mat->data;
561:   PetscBool          roworiented = baij->roworiented;
562:   PetscInt           i, j, ii, jj, row, col;
563:   PetscInt           rstart = baij->rstartbs;
564:   PetscInt           rend = mat->rmap->rend, stepval, bs = mat->rmap->bs, bs2 = baij->bs2, nbs2 = n * bs2;
565:   PetscInt           h1, key, size = baij->ht_size, idx, *HT = baij->ht, Nbs = baij->Nbs;
566:   PetscReal          tmp;
567:   MatScalar        **HD = baij->hd, *baij_a;
568:   const PetscScalar *v_t, *value;
569:   PetscInt           total_ct = baij->ht_total_ct, insert_ct = baij->ht_insert_ct;

571:   PetscFunctionBegin;
572:   if (roworiented) stepval = (n - 1) * bs;
573:   else stepval = (m - 1) * bs;

575:   for (i = 0; i < m; i++) {
576:     if (PetscDefined(USE_DEBUG)) {
577:       PetscCheck(im[i] >= 0, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Negative row: %" PetscInt_FMT, im[i]);
578:       PetscCheck(im[i] < baij->Mbs, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Row too large: row %" PetscInt_FMT " max %" PetscInt_FMT, im[i], baij->Mbs - 1);
579:     }
580:     row = im[i];
581:     v_t = v + i * nbs2;
582:     if (row >= rstart && row < rend) {
583:       for (j = 0; j < n; j++) {
584:         col = in[j];

586:         /* Look up into the Hash Table */
587:         key = row * Nbs + col + 1;
588:         h1  = HASH(size, key, tmp);

590:         idx = h1;
591:         if (PetscDefined(USE_DEBUG)) {
592:           total_ct++;
593:           insert_ct++;
594:           if (HT[idx] != key) {
595:             for (idx = h1; (idx < size) && (HT[idx] != key); idx++, total_ct++);
596:             if (idx == size) {
597:               for (idx = 0; (idx < h1) && (HT[idx] != key); idx++, total_ct++);
598:               PetscCheck(idx != h1, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "(%" PetscInt_FMT ",%" PetscInt_FMT ") has no entry in the hash table", row, col);
599:             }
600:           }
601:         } else if (HT[idx] != key) {
602:           for (idx = h1; (idx < size) && (HT[idx] != key); idx++);
603:           if (idx == size) {
604:             for (idx = 0; (idx < h1) && (HT[idx] != key); idx++);
605:             PetscCheck(idx != h1, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "(%" PetscInt_FMT ",%" PetscInt_FMT ") has no entry in the hash table", row, col);
606:           }
607:         }
608:         baij_a = HD[idx];
609:         if (roworiented) {
610:           /*value = v + i*(stepval+bs)*bs + j*bs;*/
611:           /* value = v + (i*(stepval+bs)+j)*bs; */
612:           value = v_t;
613:           v_t += bs;
614:           if (addv == ADD_VALUES) {
615:             for (ii = 0; ii < bs; ii++, value += stepval) {
616:               for (jj = ii; jj < bs2; jj += bs) baij_a[jj] += *value++;
617:             }
618:           } else {
619:             for (ii = 0; ii < bs; ii++, value += stepval) {
620:               for (jj = ii; jj < bs2; jj += bs) baij_a[jj] = *value++;
621:             }
622:           }
623:         } else {
624:           value = v + j * (stepval + bs) * bs + i * bs;
625:           if (addv == ADD_VALUES) {
626:             for (ii = 0; ii < bs; ii++, value += stepval, baij_a += bs) {
627:               for (jj = 0; jj < bs; jj++) baij_a[jj] += *value++;
628:             }
629:           } else {
630:             for (ii = 0; ii < bs; ii++, value += stepval, baij_a += bs) {
631:               for (jj = 0; jj < bs; jj++) baij_a[jj] = *value++;
632:             }
633:           }
634:         }
635:       }
636:     } else {
637:       if (!baij->donotstash) {
638:         if (roworiented) {
639:           PetscCall(MatStashValuesRowBlocked_Private(&mat->bstash, im[i], n, in, v, m, n, i));
640:         } else {
641:           PetscCall(MatStashValuesColBlocked_Private(&mat->bstash, im[i], n, in, v, m, n, i));
642:         }
643:       }
644:     }
645:   }
646:   if (PetscDefined(USE_DEBUG)) {
647:     baij->ht_total_ct += total_ct;
648:     baij->ht_insert_ct += insert_ct;
649:   }
650:   PetscFunctionReturn(PETSC_SUCCESS);
651: }

653: static PetscErrorCode MatGetValues_MPIBAIJ(Mat mat, PetscInt m, const PetscInt idxm[], PetscInt n, const PetscInt idxn[], PetscScalar v[])
654: {
655:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)mat->data;
656:   PetscInt     bs = mat->rmap->bs, i, j, bsrstart = mat->rmap->rstart, bsrend = mat->rmap->rend;
657:   PetscInt     bscstart = mat->cmap->rstart, bscend = mat->cmap->rend, row, col, data;

659:   PetscFunctionBegin;
660:   for (i = 0; i < m; i++) {
661:     if (idxm[i] < 0) continue; /* negative row */
662:     PetscCheck(idxm[i] < mat->rmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Row too large: row %" PetscInt_FMT " max %" PetscInt_FMT, idxm[i], mat->rmap->N - 1);
663:     PetscCheck(idxm[i] >= bsrstart && idxm[i] < bsrend, PETSC_COMM_SELF, PETSC_ERR_SUP, "Only local values currently supported");
664:     row = idxm[i] - bsrstart;
665:     for (j = 0; j < n; j++) {
666:       if (idxn[j] < 0) continue; /* negative column */
667:       PetscCheck(idxn[j] < mat->cmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Column too large: col %" PetscInt_FMT " max %" PetscInt_FMT, idxn[j], mat->cmap->N - 1);
668:       if (idxn[j] >= bscstart && idxn[j] < bscend) {
669:         col = idxn[j] - bscstart;
670:         PetscCall(MatGetValues_SeqBAIJ(baij->A, 1, &row, 1, &col, v + i * n + j));
671:       } else {
672:         if (!baij->colmap) PetscCall(MatCreateColmap_MPIBAIJ_Private(mat));
673: #if defined(PETSC_USE_CTABLE)
674:         PetscCall(PetscHMapIGetWithDefault(baij->colmap, idxn[j] / bs + 1, 0, &data));
675:         data--;
676: #else
677:         data = baij->colmap[idxn[j] / bs] - 1;
678: #endif
679:         if (data < 0 || baij->garray[data / bs] != idxn[j] / bs) *(v + i * n + j) = 0.0;
680:         else {
681:           col = data + idxn[j] % bs;
682:           PetscCall(MatGetValues_SeqBAIJ(baij->B, 1, &row, 1, &col, v + i * n + j));
683:         }
684:       }
685:     }
686:   }
687:   PetscFunctionReturn(PETSC_SUCCESS);
688: }

690: static PetscErrorCode MatNorm_MPIBAIJ(Mat mat, NormType type, PetscReal *nrm)
691: {
692:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)mat->data;
693:   Mat_SeqBAIJ *amat = (Mat_SeqBAIJ *)baij->A->data, *bmat = (Mat_SeqBAIJ *)baij->B->data;
694:   PetscInt     i, j, bs2 = baij->bs2, bs = baij->A->rmap->bs, nz, row, col;
695:   PetscReal    sum = 0.0;
696:   MatScalar   *v;

698:   PetscFunctionBegin;
699:   if (baij->size == 1) {
700:     PetscCall(MatNorm(baij->A, type, nrm));
701:   } else {
702:     if (type == NORM_FROBENIUS) {
703:       v  = amat->a;
704:       nz = amat->nz * bs2;
705:       for (i = 0; i < nz; i++) {
706:         sum += PetscRealPart(PetscConj(*v) * (*v));
707:         v++;
708:       }
709:       v  = bmat->a;
710:       nz = bmat->nz * bs2;
711:       for (i = 0; i < nz; i++) {
712:         sum += PetscRealPart(PetscConj(*v) * (*v));
713:         v++;
714:       }
715:       PetscCallMPI(MPIU_Allreduce(&sum, nrm, 1, MPIU_REAL, MPIU_SUM, PetscObjectComm((PetscObject)mat)));
716:       *nrm = PetscSqrtReal(*nrm);
717:     } else if (type == NORM_1) { /* max column sum */
718:       Vec          col, bcol;
719:       PetscScalar *array;
720:       PetscInt    *jj, *garray = baij->garray;

722:       PetscCall(MatCreateVecs(mat, &col, NULL));
723:       PetscCall(VecGetArrayWrite(col, &array));
724:       v  = amat->a;
725:       jj = amat->j;
726:       for (i = 0; i < amat->nz; i++) {
727:         for (j = 0; j < bs; j++) {
728:           PetscInt col = bs * *jj + j; /* column index */

730:           for (row = 0; row < bs; row++) array[col] += PetscAbsScalar(*v++);
731:         }
732:         jj++;
733:       }
734:       PetscCall(VecRestoreArrayWrite(col, &array));
735:       PetscCall(MatCreateVecs(baij->B, &bcol, NULL));
736:       PetscCall(VecGetArrayWrite(bcol, &array));
737:       v  = bmat->a;
738:       jj = bmat->j;
739:       for (i = 0; i < bmat->nz; i++) {
740:         for (j = 0; j < bs; j++) {
741:           PetscInt col = bs * *jj + j; /* column index */

743:           for (row = 0; row < bs; row++) array[col] += PetscAbsScalar(*v++);
744:         }
745:         jj++;
746:       }
747:       PetscCall(VecSetValuesBlocked(col, bmat->nbs, garray, array, ADD_VALUES));
748:       PetscCall(VecRestoreArrayWrite(bcol, &array));
749:       PetscCall(VecDestroy(&bcol));
750:       PetscCall(VecAssemblyBegin(col));
751:       PetscCall(VecAssemblyEnd(col));
752:       PetscCall(VecNorm(col, NORM_INFINITY, nrm));
753:       PetscCall(VecDestroy(&col));
754:     } else if (type == NORM_INFINITY) { /* max row sum */
755:       PetscReal *sums;
756:       PetscCall(PetscMalloc1(bs, &sums));
757:       sum = 0.0;
758:       for (j = 0; j < amat->mbs; j++) {
759:         for (row = 0; row < bs; row++) sums[row] = 0.0;
760:         v  = amat->a + bs2 * amat->i[j];
761:         nz = amat->i[j + 1] - amat->i[j];
762:         for (i = 0; i < nz; i++) {
763:           for (col = 0; col < bs; col++) {
764:             for (row = 0; row < bs; row++) {
765:               sums[row] += PetscAbsScalar(*v);
766:               v++;
767:             }
768:           }
769:         }
770:         v  = bmat->a + bs2 * bmat->i[j];
771:         nz = bmat->i[j + 1] - bmat->i[j];
772:         for (i = 0; i < nz; i++) {
773:           for (col = 0; col < bs; col++) {
774:             for (row = 0; row < bs; row++) {
775:               sums[row] += PetscAbsScalar(*v);
776:               v++;
777:             }
778:           }
779:         }
780:         for (row = 0; row < bs; row++) {
781:           if (sums[row] > sum) sum = sums[row];
782:         }
783:       }
784:       PetscCallMPI(MPIU_Allreduce(&sum, nrm, 1, MPIU_REAL, MPIU_MAX, PetscObjectComm((PetscObject)mat)));
785:       PetscCall(PetscFree(sums));
786:     } else SETERRQ(PetscObjectComm((PetscObject)mat), PETSC_ERR_SUP, "No support for this norm yet");
787:   }
788:   PetscFunctionReturn(PETSC_SUCCESS);
789: }

791: /*
792:   Creates the hash table, and sets the table
793:   This table is created only once.
794:   If new entries need to be added to the matrix
795:   then the hash table has to be destroyed and
796:   recreated.
797: */
798: static PetscErrorCode MatCreateHashTable_MPIBAIJ_Private(Mat mat, PetscReal factor)
799: {
800:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)mat->data;
801:   Mat          A = baij->A, B = baij->B;
802:   Mat_SeqBAIJ *a = (Mat_SeqBAIJ *)A->data, *b = (Mat_SeqBAIJ *)B->data;
803:   PetscInt     i, j, k, nz = a->nz + b->nz, h1, *ai = a->i, *aj = a->j, *bi = b->i, *bj = b->j;
804:   PetscInt     ht_size, bs2 = baij->bs2, rstart = baij->rstartbs;
805:   PetscInt     cstart = baij->cstartbs, *garray = baij->garray, row, col, Nbs = baij->Nbs;
806:   PetscInt    *HT, key;
807:   MatScalar  **HD;
808:   PetscReal    tmp;
809: #if defined(PETSC_USE_INFO)
810:   PetscInt ct = 0, max = 0;
811: #endif

813:   PetscFunctionBegin;
814:   if (baij->ht) PetscFunctionReturn(PETSC_SUCCESS);

816:   baij->ht_size = (PetscInt)(factor * nz);
817:   ht_size       = baij->ht_size;

819:   /* Allocate Memory for Hash Table */
820:   PetscCall(PetscCalloc2(ht_size, &baij->hd, ht_size, &baij->ht));
821:   HD = baij->hd;
822:   HT = baij->ht;

824:   /* Loop Over A */
825:   for (i = 0; i < a->mbs; i++) {
826:     for (j = ai[i]; j < ai[i + 1]; j++) {
827:       row = i + rstart;
828:       col = aj[j] + cstart;

830:       key = row * Nbs + col + 1;
831:       h1  = HASH(ht_size, key, tmp);
832:       for (k = 0; k < ht_size; k++) {
833:         if (!HT[(h1 + k) % ht_size]) {
834:           HT[(h1 + k) % ht_size] = key;
835:           HD[(h1 + k) % ht_size] = a->a + j * bs2;
836:           break;
837: #if defined(PETSC_USE_INFO)
838:         } else {
839:           ct++;
840: #endif
841:         }
842:       }
843: #if defined(PETSC_USE_INFO)
844:       if (k > max) max = k;
845: #endif
846:     }
847:   }
848:   /* Loop Over B */
849:   for (i = 0; i < b->mbs; i++) {
850:     for (j = bi[i]; j < bi[i + 1]; j++) {
851:       row = i + rstart;
852:       col = garray[bj[j]];
853:       key = row * Nbs + col + 1;
854:       h1  = HASH(ht_size, key, tmp);
855:       for (k = 0; k < ht_size; k++) {
856:         if (!HT[(h1 + k) % ht_size]) {
857:           HT[(h1 + k) % ht_size] = key;
858:           HD[(h1 + k) % ht_size] = b->a + j * bs2;
859:           break;
860: #if defined(PETSC_USE_INFO)
861:         } else {
862:           ct++;
863: #endif
864:         }
865:       }
866: #if defined(PETSC_USE_INFO)
867:       if (k > max) max = k;
868: #endif
869:     }
870:   }

872:   /* Print Summary */
873: #if defined(PETSC_USE_INFO)
874:   for (i = 0, j = 0; i < ht_size; i++) {
875:     if (HT[i]) j++;
876:   }
877:   PetscCall(PetscInfo(mat, "Average Search = %5.2g,max search = %" PetscInt_FMT "\n", (!j) ? 0.0 : (double)(((PetscReal)(ct + j)) / j), max));
878: #endif
879:   PetscFunctionReturn(PETSC_SUCCESS);
880: }

882: static PetscErrorCode MatAssemblyBegin_MPIBAIJ(Mat mat, MatAssemblyType mode)
883: {
884:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)mat->data;
885:   PetscInt     nstash, reallocs;

887:   PetscFunctionBegin;
888:   if (baij->donotstash || mat->nooffprocentries) PetscFunctionReturn(PETSC_SUCCESS);

890:   PetscCall(MatStashScatterBegin_Private(mat, &mat->stash, mat->rmap->range));
891:   PetscCall(MatStashScatterBegin_Private(mat, &mat->bstash, baij->rangebs));
892:   PetscCall(MatStashGetInfo_Private(&mat->stash, &nstash, &reallocs));
893:   PetscCall(PetscInfo(mat, "Stash has %" PetscInt_FMT " entries,uses %" PetscInt_FMT " mallocs.\n", nstash, reallocs));
894:   PetscCall(MatStashGetInfo_Private(&mat->bstash, &nstash, &reallocs));
895:   PetscCall(PetscInfo(mat, "Block-Stash has %" PetscInt_FMT " entries, uses %" PetscInt_FMT " mallocs.\n", nstash, reallocs));
896:   PetscFunctionReturn(PETSC_SUCCESS);
897: }

899: static PetscErrorCode MatAssemblyEnd_MPIBAIJ(Mat mat, MatAssemblyType mode)
900: {
901:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)mat->data;
902:   Mat_SeqBAIJ *a    = (Mat_SeqBAIJ *)baij->A->data;
903:   PetscInt     i, j, rstart, ncols, flg, bs2 = baij->bs2;
904:   PetscInt    *row, *col;
905:   PetscBool    r1, r2, r3, all_assembled;
906:   MatScalar   *val;
907:   PetscMPIInt  n;

909:   PetscFunctionBegin;
910:   /* do not use 'b=(Mat_SeqBAIJ*)baij->B->data' as B can be reset in disassembly */
911:   if (!baij->donotstash && !mat->nooffprocentries) {
912:     while (1) {
913:       PetscCall(MatStashScatterGetMesg_Private(&mat->stash, &n, &row, &col, &val, &flg));
914:       if (!flg) break;

916:       for (i = 0; i < n;) {
917:         /* Now identify the consecutive vals belonging to the same row */
918:         for (j = i, rstart = row[j]; j < n; j++) {
919:           if (row[j] != rstart) break;
920:         }
921:         if (j < n) ncols = j - i;
922:         else ncols = n - i;
923:         /* Now assemble all these values with a single function call */
924:         PetscCall(MatSetValues_MPIBAIJ(mat, 1, row + i, ncols, col + i, val + i, mat->insertmode));
925:         i = j;
926:       }
927:     }
928:     PetscCall(MatStashScatterEnd_Private(&mat->stash));
929:     /* Now process the block-stash. Since the values are stashed column-oriented,
930:        set the row-oriented flag to column-oriented, and after MatSetValues()
931:        restore the original flags */
932:     r1 = baij->roworiented;
933:     r2 = a->roworiented;
934:     r3 = ((Mat_SeqBAIJ *)baij->B->data)->roworiented;

936:     baij->roworiented                           = PETSC_FALSE;
937:     a->roworiented                              = PETSC_FALSE;
938:     ((Mat_SeqBAIJ *)baij->B->data)->roworiented = PETSC_FALSE;
939:     while (1) {
940:       PetscCall(MatStashScatterGetMesg_Private(&mat->bstash, &n, &row, &col, &val, &flg));
941:       if (!flg) break;

943:       for (i = 0; i < n;) {
944:         /* Now identify the consecutive vals belonging to the same row */
945:         for (j = i, rstart = row[j]; j < n; j++) {
946:           if (row[j] != rstart) break;
947:         }
948:         if (j < n) ncols = j - i;
949:         else ncols = n - i;
950:         PetscCall(MatSetValuesBlocked_MPIBAIJ(mat, 1, row + i, ncols, col + i, val + i * bs2, mat->insertmode));
951:         i = j;
952:       }
953:     }
954:     PetscCall(MatStashScatterEnd_Private(&mat->bstash));

956:     baij->roworiented                           = r1;
957:     a->roworiented                              = r2;
958:     ((Mat_SeqBAIJ *)baij->B->data)->roworiented = r3;
959:   }

961:   PetscCall(MatAssemblyBegin(baij->A, mode));
962:   PetscCall(MatAssemblyEnd(baij->A, mode));

964:   /* determine if any process has disassembled, if so we must
965:      also disassemble ourselves, in order that we may reassemble. */
966:   /*
967:      if nonzero structure of submatrix B cannot change then we know that
968:      no process disassembled thus we can skip this stuff
969:   */
970:   if (!((Mat_SeqBAIJ *)baij->B->data)->nonew) {
971:     PetscCallMPI(MPIU_Allreduce(&mat->was_assembled, &all_assembled, 1, MPI_C_BOOL, MPI_LAND, PetscObjectComm((PetscObject)mat)));
972:     if (mat->was_assembled && !all_assembled) PetscCall(MatDisAssemble_MPIBAIJ(mat));
973:   }

975:   if (!mat->was_assembled && mode == MAT_FINAL_ASSEMBLY) PetscCall(MatSetUpMultiply_MPIBAIJ(mat));
976:   PetscCall(MatAssemblyBegin(baij->B, mode));
977:   PetscCall(MatAssemblyEnd(baij->B, mode));

979: #if defined(PETSC_USE_INFO)
980:   if (baij->ht && mode == MAT_FINAL_ASSEMBLY) {
981:     PetscCall(PetscInfo(mat, "Average Hash Table Search in MatSetValues = %5.2f\n", (double)((PetscReal)baij->ht_total_ct) / baij->ht_insert_ct));

983:     baij->ht_total_ct  = 0;
984:     baij->ht_insert_ct = 0;
985:   }
986: #endif
987:   if (baij->ht_flag && !baij->ht && mode == MAT_FINAL_ASSEMBLY) {
988:     PetscCall(MatCreateHashTable_MPIBAIJ_Private(mat, baij->ht_fact));

990:     mat->ops->setvalues        = MatSetValues_MPIBAIJ_HT;
991:     mat->ops->setvaluesblocked = MatSetValuesBlocked_MPIBAIJ_HT;
992:   }

994:   PetscCall(PetscFree2(baij->rowvalues, baij->rowindices));

996:   baij->rowvalues = NULL;

998:   /* if no new nonzero locations are allowed in matrix then only set the matrix state the first time through */
999:   if ((!mat->was_assembled && mode == MAT_FINAL_ASSEMBLY) || !((Mat_SeqBAIJ *)baij->A->data)->nonew) {
1000:     PetscObjectState state = baij->A->nonzerostate + baij->B->nonzerostate;
1001:     PetscCallMPI(MPIU_Allreduce(&state, &mat->nonzerostate, 1, MPIU_INT64, MPI_SUM, PetscObjectComm((PetscObject)mat)));
1002:   }
1003:   PetscFunctionReturn(PETSC_SUCCESS);
1004: }

1006: #include <petscdraw.h>
1007: static PetscErrorCode MatView_MPIBAIJ_ASCIIorDraworSocket(Mat mat, PetscViewer viewer)
1008: {
1009:   Mat_MPIBAIJ      *baij = (Mat_MPIBAIJ *)mat->data;
1010:   PetscMPIInt       rank = baij->rank;
1011:   PetscInt          bs   = mat->rmap->bs;
1012:   PetscBool         isascii, isdraw;
1013:   PetscViewer       sviewer;
1014:   PetscViewerFormat format;

1016:   PetscFunctionBegin;
1017:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
1018:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERDRAW, &isdraw));
1019:   if (isascii) {
1020:     PetscCall(PetscViewerGetFormat(viewer, &format));
1021:     if (format == PETSC_VIEWER_ASCII_INFO_DETAIL) {
1022:       MatInfo info;
1023:       PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)mat), &rank));
1024:       PetscCall(MatGetInfo(mat, MAT_LOCAL, &info));
1025:       PetscCall(PetscViewerASCIIPushSynchronized(viewer));
1026:       PetscCall(PetscViewerASCIISynchronizedPrintf(viewer, "[%d] Local rows %" PetscInt_FMT " nz %" PetscInt_FMT " nz alloced %" PetscInt_FMT " bs %" PetscInt_FMT " mem %g\n", rank, mat->rmap->n, (PetscInt)info.nz_used, (PetscInt)info.nz_allocated,
1027:                                                    mat->rmap->bs, info.memory));
1028:       PetscCall(MatGetInfo(baij->A, MAT_LOCAL, &info));
1029:       PetscCall(PetscViewerASCIISynchronizedPrintf(viewer, "[%d] on-diagonal part: nz %" PetscInt_FMT " \n", rank, (PetscInt)info.nz_used));
1030:       PetscCall(MatGetInfo(baij->B, MAT_LOCAL, &info));
1031:       PetscCall(PetscViewerASCIISynchronizedPrintf(viewer, "[%d] off-diagonal part: nz %" PetscInt_FMT " \n", rank, (PetscInt)info.nz_used));
1032:       PetscCall(PetscViewerFlush(viewer));
1033:       PetscCall(PetscViewerASCIIPopSynchronized(viewer));
1034:       PetscCall(PetscViewerASCIIPrintf(viewer, "Information on VecScatter used in matrix-vector product: \n"));
1035:       PetscCall(VecScatterView(baij->Mvctx, viewer));
1036:       PetscFunctionReturn(PETSC_SUCCESS);
1037:     } else if (format == PETSC_VIEWER_ASCII_INFO || format == PETSC_VIEWER_ASCII_FACTOR_INFO) PetscFunctionReturn(PETSC_SUCCESS);
1038:   }

1040:   if (isdraw) {
1041:     PetscDraw draw;
1042:     PetscBool isnull;
1043:     PetscCall(PetscViewerDrawGetDraw(viewer, 0, &draw));
1044:     PetscCall(PetscDrawIsNull(draw, &isnull));
1045:     if (isnull) PetscFunctionReturn(PETSC_SUCCESS);
1046:   }

1048:   {
1049:     /* assemble the entire matrix onto first processor. */
1050:     Mat          A;
1051:     Mat_SeqBAIJ *Aloc;
1052:     PetscInt     M = mat->rmap->N, N = mat->cmap->N, *ai, *aj, col, i, j, k, *rvals, mbs = baij->mbs;
1053:     MatScalar   *a;
1054:     const char  *matname;

1056:     /* Here we are creating a temporary matrix, so will assume MPIBAIJ is acceptable */
1057:     /* Perhaps this should be the type of mat? */
1058:     PetscCall(MatCreate(PetscObjectComm((PetscObject)mat), &A));
1059:     if (rank == 0) {
1060:       PetscCall(MatSetSizes(A, M, N, M, N));
1061:     } else {
1062:       PetscCall(MatSetSizes(A, 0, 0, M, N));
1063:     }
1064:     PetscCall(MatSetType(A, MATMPIBAIJ));
1065:     PetscCall(MatMPIBAIJSetPreallocation(A, mat->rmap->bs, 0, NULL, 0, NULL));
1066:     PetscCall(MatSetOption(A, MAT_NEW_NONZERO_LOCATION_ERR, PETSC_FALSE));

1068:     /* copy over the A part */
1069:     Aloc = (Mat_SeqBAIJ *)baij->A->data;
1070:     ai   = Aloc->i;
1071:     aj   = Aloc->j;
1072:     a    = Aloc->a;
1073:     PetscCall(PetscMalloc1(bs, &rvals));

1075:     for (i = 0; i < mbs; i++) {
1076:       rvals[0] = bs * (baij->rstartbs + i);
1077:       for (j = 1; j < bs; j++) rvals[j] = rvals[j - 1] + 1;
1078:       for (j = ai[i]; j < ai[i + 1]; j++) {
1079:         col = (baij->cstartbs + aj[j]) * bs;
1080:         for (k = 0; k < bs; k++) {
1081:           PetscCall(MatSetValues_MPIBAIJ(A, bs, rvals, 1, &col, a, INSERT_VALUES));
1082:           col++;
1083:           a += bs;
1084:         }
1085:       }
1086:     }
1087:     /* copy over the B part */
1088:     Aloc = (Mat_SeqBAIJ *)baij->B->data;
1089:     ai   = Aloc->i;
1090:     aj   = Aloc->j;
1091:     a    = Aloc->a;
1092:     for (i = 0; i < mbs; i++) {
1093:       rvals[0] = bs * (baij->rstartbs + i);
1094:       for (j = 1; j < bs; j++) rvals[j] = rvals[j - 1] + 1;
1095:       for (j = ai[i]; j < ai[i + 1]; j++) {
1096:         col = baij->garray[aj[j]] * bs;
1097:         for (k = 0; k < bs; k++) {
1098:           PetscCall(MatSetValues_MPIBAIJ(A, bs, rvals, 1, &col, a, INSERT_VALUES));
1099:           col++;
1100:           a += bs;
1101:         }
1102:       }
1103:     }
1104:     PetscCall(PetscFree(rvals));
1105:     PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
1106:     PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
1107:     /*
1108:        Everyone has to call to draw the matrix since the graphics waits are
1109:        synchronized across all processors that share the PetscDraw object
1110:     */
1111:     PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sviewer));
1112:     if (((PetscObject)mat)->name) PetscCall(PetscObjectGetName((PetscObject)mat, &matname));
1113:     if (rank == 0) {
1114:       if (((PetscObject)mat)->name) PetscCall(PetscObjectSetName((PetscObject)((Mat_MPIBAIJ *)A->data)->A, matname));
1115:       PetscCall(MatView_SeqBAIJ(((Mat_MPIBAIJ *)A->data)->A, sviewer));
1116:     }
1117:     PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sviewer));
1118:     PetscCall(MatDestroy(&A));
1119:   }
1120:   PetscFunctionReturn(PETSC_SUCCESS);
1121: }

1123: /* Used for both MPIBAIJ and MPISBAIJ matrices */
1124: PetscErrorCode MatView_MPIBAIJ_Binary(Mat mat, PetscViewer viewer)
1125: {
1126:   Mat_MPIBAIJ    *aij    = (Mat_MPIBAIJ *)mat->data;
1127:   Mat_SeqBAIJ    *A      = (Mat_SeqBAIJ *)aij->A->data;
1128:   Mat_SeqBAIJ    *B      = (Mat_SeqBAIJ *)aij->B->data;
1129:   const PetscInt *garray = aij->garray;
1130:   PetscInt        header[4], M, N, m, rs, cs, bs, cnt, i, j, ja, jb, k, l;
1131:   PetscCount      nz, hnz;
1132:   PetscInt       *rowlens, *colidxs;
1133:   PetscScalar    *matvals;
1134:   PetscMPIInt     rank;

1136:   PetscFunctionBegin;
1137:   PetscCall(PetscViewerSetUp(viewer));

1139:   M  = mat->rmap->N;
1140:   N  = mat->cmap->N;
1141:   m  = mat->rmap->n;
1142:   rs = mat->rmap->rstart;
1143:   cs = mat->cmap->rstart;
1144:   bs = mat->rmap->bs;
1145:   nz = bs * bs * (A->nz + B->nz);

1147:   /* write matrix header */
1148:   header[0] = MAT_FILE_CLASSID;
1149:   header[1] = M;
1150:   header[2] = N;
1151:   PetscCallMPI(MPI_Reduce(&nz, &hnz, 1, MPIU_COUNT, MPI_SUM, 0, PetscObjectComm((PetscObject)mat)));
1152:   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)mat), &rank));
1153:   if (rank == 0) PetscCall(PetscIntCast(hnz, &header[3]));
1154:   PetscCall(PetscViewerBinaryWrite(viewer, header, 4, PETSC_INT));

1156:   /* fill in and store row lengths */
1157:   PetscCall(PetscMalloc1(m, &rowlens));
1158:   for (cnt = 0, i = 0; i < A->mbs; i++)
1159:     for (j = 0; j < bs; j++) rowlens[cnt++] = bs * (A->i[i + 1] - A->i[i] + B->i[i + 1] - B->i[i]);
1160:   PetscCall(PetscViewerBinaryWriteAll(viewer, rowlens, m, rs, M, PETSC_INT));
1161:   PetscCall(PetscFree(rowlens));

1163:   /* fill in and store column indices */
1164:   PetscCall(PetscMalloc1(nz, &colidxs));
1165:   for (cnt = 0, i = 0; i < A->mbs; i++) {
1166:     for (k = 0; k < bs; k++) {
1167:       for (jb = B->i[i]; jb < B->i[i + 1]; jb++) {
1168:         if (garray[B->j[jb]] > cs / bs) break;
1169:         for (l = 0; l < bs; l++) colidxs[cnt++] = bs * garray[B->j[jb]] + l;
1170:       }
1171:       for (ja = A->i[i]; ja < A->i[i + 1]; ja++)
1172:         for (l = 0; l < bs; l++) colidxs[cnt++] = bs * A->j[ja] + l + cs;
1173:       for (; jb < B->i[i + 1]; jb++)
1174:         for (l = 0; l < bs; l++) colidxs[cnt++] = bs * garray[B->j[jb]] + l;
1175:     }
1176:   }
1177:   PetscCheck(cnt == nz, PETSC_COMM_SELF, PETSC_ERR_LIB, "Internal PETSc error: cnt = %" PetscInt_FMT " nz = %" PetscCount_FMT, cnt, nz);
1178:   PetscCall(PetscViewerBinaryWriteAll(viewer, colidxs, nz, PETSC_DECIDE, PETSC_DECIDE, PETSC_INT));
1179:   PetscCall(PetscFree(colidxs));

1181:   /* fill in and store nonzero values */
1182:   PetscCall(PetscMalloc1(nz, &matvals));
1183:   for (cnt = 0, i = 0; i < A->mbs; i++) {
1184:     for (k = 0; k < bs; k++) {
1185:       for (jb = B->i[i]; jb < B->i[i + 1]; jb++) {
1186:         if (garray[B->j[jb]] > cs / bs) break;
1187:         for (l = 0; l < bs; l++) matvals[cnt++] = B->a[bs * (bs * jb + l) + k];
1188:       }
1189:       for (ja = A->i[i]; ja < A->i[i + 1]; ja++)
1190:         for (l = 0; l < bs; l++) matvals[cnt++] = A->a[bs * (bs * ja + l) + k];
1191:       for (; jb < B->i[i + 1]; jb++)
1192:         for (l = 0; l < bs; l++) matvals[cnt++] = B->a[bs * (bs * jb + l) + k];
1193:     }
1194:   }
1195:   PetscCall(PetscViewerBinaryWriteAll(viewer, matvals, nz, PETSC_DECIDE, PETSC_DECIDE, PETSC_SCALAR));
1196:   PetscCall(PetscFree(matvals));

1198:   /* write block size option to the viewer's .info file */
1199:   PetscCall(MatView_Binary_BlockSizes(mat, viewer));
1200:   PetscFunctionReturn(PETSC_SUCCESS);
1201: }

1203: PetscErrorCode MatView_MPIBAIJ(Mat mat, PetscViewer viewer)
1204: {
1205:   PetscBool isascii, isdraw, issocket, isbinary;

1207:   PetscFunctionBegin;
1208:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
1209:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERDRAW, &isdraw));
1210:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERSOCKET, &issocket));
1211:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &isbinary));
1212:   if (isascii || isdraw || issocket) PetscCall(MatView_MPIBAIJ_ASCIIorDraworSocket(mat, viewer));
1213:   else if (isbinary) PetscCall(MatView_MPIBAIJ_Binary(mat, viewer));
1214:   PetscFunctionReturn(PETSC_SUCCESS);
1215: }

1217: static PetscErrorCode MatMult_MPIBAIJ(Mat A, Vec xx, Vec yy)
1218: {
1219:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;
1220:   PetscInt     nt;

1222:   PetscFunctionBegin;
1223:   PetscCall(VecGetLocalSize(xx, &nt));
1224:   PetscCheck(nt == A->cmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Incompatible partition of A and xx");
1225:   PetscCall(VecGetLocalSize(yy, &nt));
1226:   PetscCheck(nt == A->rmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Incompatible partition of A and yy");
1227:   PetscCall(VecScatterBegin(a->Mvctx, xx, a->lvec, INSERT_VALUES, SCATTER_FORWARD));
1228:   PetscUseTypeMethod(a->A, mult, xx, yy);
1229:   PetscCall(VecScatterEnd(a->Mvctx, xx, a->lvec, INSERT_VALUES, SCATTER_FORWARD));
1230:   PetscUseTypeMethod(a->B, multadd, a->lvec, yy, yy);
1231:   PetscFunctionReturn(PETSC_SUCCESS);
1232: }

1234: static PetscErrorCode MatMultAdd_MPIBAIJ(Mat A, Vec xx, Vec yy, Vec zz)
1235: {
1236:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

1238:   PetscFunctionBegin;
1239:   PetscCall(VecScatterBegin(a->Mvctx, xx, a->lvec, INSERT_VALUES, SCATTER_FORWARD));
1240:   PetscUseTypeMethod(a->A, multadd, xx, yy, zz);
1241:   PetscCall(VecScatterEnd(a->Mvctx, xx, a->lvec, INSERT_VALUES, SCATTER_FORWARD));
1242:   PetscUseTypeMethod(a->B, multadd, a->lvec, zz, zz);
1243:   PetscFunctionReturn(PETSC_SUCCESS);
1244: }

1246: static PetscErrorCode MatMultTranspose_MPIBAIJ(Mat A, Vec xx, Vec yy)
1247: {
1248:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

1250:   PetscFunctionBegin;
1251:   /* do nondiagonal part */
1252:   PetscUseTypeMethod(a->B, multtranspose, xx, a->lvec);
1253:   /* do local part */
1254:   PetscUseTypeMethod(a->A, multtranspose, xx, yy);
1255:   /* add partial results together */
1256:   PetscCall(VecScatterBegin(a->Mvctx, a->lvec, yy, ADD_VALUES, SCATTER_REVERSE));
1257:   PetscCall(VecScatterEnd(a->Mvctx, a->lvec, yy, ADD_VALUES, SCATTER_REVERSE));
1258:   PetscFunctionReturn(PETSC_SUCCESS);
1259: }

1261: static PetscErrorCode MatMultTransposeAdd_MPIBAIJ(Mat A, Vec xx, Vec yy, Vec zz)
1262: {
1263:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

1265:   PetscFunctionBegin;
1266:   /* do nondiagonal part */
1267:   PetscUseTypeMethod(a->B, multtranspose, xx, a->lvec);
1268:   /* do local part */
1269:   PetscUseTypeMethod(a->A, multtransposeadd, xx, yy, zz);
1270:   /* add partial results together */
1271:   PetscCall(VecScatterBegin(a->Mvctx, a->lvec, zz, ADD_VALUES, SCATTER_REVERSE));
1272:   PetscCall(VecScatterEnd(a->Mvctx, a->lvec, zz, ADD_VALUES, SCATTER_REVERSE));
1273:   PetscFunctionReturn(PETSC_SUCCESS);
1274: }

1276: /*
1277:   This only works correctly for square matrices where the subblock A->A is the
1278:    diagonal block
1279: */
1280: static PetscErrorCode MatGetDiagonal_MPIBAIJ(Mat A, Vec v)
1281: {
1282:   PetscFunctionBegin;
1283:   PetscCheck(A->rmap->N == A->cmap->N, PETSC_COMM_SELF, PETSC_ERR_SUP, "Supports only square matrix where A->A is diag block");
1284:   PetscCall(MatGetDiagonal(((Mat_MPIBAIJ *)A->data)->A, v));
1285:   PetscFunctionReturn(PETSC_SUCCESS);
1286: }

1288: static PetscErrorCode MatScale_MPIBAIJ(Mat A, PetscScalar aa)
1289: {
1290:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

1292:   PetscFunctionBegin;
1293:   PetscCall(MatScale(a->A, aa));
1294:   PetscCall(MatScale(a->B, aa));
1295:   PetscFunctionReturn(PETSC_SUCCESS);
1296: }

1298: static PetscErrorCode MatGetRow_MPIBAIJ(Mat matin, PetscInt row, PetscInt *nz, PetscInt **idx, PetscScalar **v)
1299: {
1300:   Mat_MPIBAIJ *mat = (Mat_MPIBAIJ *)matin->data;
1301:   PetscScalar *vworkA, *vworkB, **pvA, **pvB, *v_p;
1302:   PetscInt     bs = matin->rmap->bs, bs2 = mat->bs2, i, *cworkA, *cworkB, **pcA, **pcB;
1303:   PetscInt     nztot, nzA, nzB, lrow, brstart = matin->rmap->rstart, brend = matin->rmap->rend;
1304:   PetscInt    *cmap, *idx_p, cstart = mat->cstartbs;

1306:   PetscFunctionBegin;
1307:   PetscCheck(row >= brstart && row < brend, PETSC_COMM_SELF, PETSC_ERR_SUP, "Only local rows");
1308:   PetscCheck(!mat->getrowactive, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Already active");
1309:   mat->getrowactive = PETSC_TRUE;

1311:   if (!mat->rowvalues && (idx || v)) {
1312:     /*
1313:         allocate enough space to hold information from the longest row.
1314:     */
1315:     Mat_SeqBAIJ *Aa = (Mat_SeqBAIJ *)mat->A->data, *Ba = (Mat_SeqBAIJ *)mat->B->data;
1316:     PetscInt     max = 1, mbs = mat->mbs, tmp;
1317:     for (i = 0; i < mbs; i++) {
1318:       tmp = Aa->i[i + 1] - Aa->i[i] + Ba->i[i + 1] - Ba->i[i];
1319:       if (max < tmp) max = tmp;
1320:     }
1321:     PetscCall(PetscMalloc2(max * bs2, &mat->rowvalues, max * bs2, &mat->rowindices));
1322:   }
1323:   lrow = row - brstart;

1325:   pvA = &vworkA;
1326:   pcA = &cworkA;
1327:   pvB = &vworkB;
1328:   pcB = &cworkB;
1329:   if (!v) {
1330:     pvA = NULL;
1331:     pvB = NULL;
1332:   }
1333:   if (!idx) {
1334:     pcA = NULL;
1335:     if (!v) pcB = NULL;
1336:   }
1337:   PetscUseTypeMethod(mat->A, getrow, lrow, &nzA, pcA, pvA);
1338:   PetscUseTypeMethod(mat->B, getrow, lrow, &nzB, pcB, pvB);
1339:   nztot = nzA + nzB;

1341:   cmap = mat->garray;
1342:   if (v || idx) {
1343:     if (nztot) {
1344:       /* Sort by increasing column numbers, assuming A and B already sorted */
1345:       PetscInt imark = -1;
1346:       if (v) {
1347:         *v = v_p = mat->rowvalues;
1348:         for (i = 0; i < nzB; i++) {
1349:           if (cmap[cworkB[i] / bs] < cstart) v_p[i] = vworkB[i];
1350:           else break;
1351:         }
1352:         imark = i;
1353:         for (i = 0; i < nzA; i++) v_p[imark + i] = vworkA[i];
1354:         for (i = imark; i < nzB; i++) v_p[nzA + i] = vworkB[i];
1355:       }
1356:       if (idx) {
1357:         *idx = idx_p = mat->rowindices;
1358:         if (imark > -1) {
1359:           for (i = 0; i < imark; i++) idx_p[i] = cmap[cworkB[i] / bs] * bs + cworkB[i] % bs;
1360:         } else {
1361:           for (i = 0; i < nzB; i++) {
1362:             if (cmap[cworkB[i] / bs] < cstart) idx_p[i] = cmap[cworkB[i] / bs] * bs + cworkB[i] % bs;
1363:             else break;
1364:           }
1365:           imark = i;
1366:         }
1367:         for (i = 0; i < nzA; i++) idx_p[imark + i] = cstart * bs + cworkA[i];
1368:         for (i = imark; i < nzB; i++) idx_p[nzA + i] = cmap[cworkB[i] / bs] * bs + cworkB[i] % bs;
1369:       }
1370:     } else {
1371:       if (idx) *idx = NULL;
1372:       if (v) *v = NULL;
1373:     }
1374:   }
1375:   *nz = nztot;
1376:   PetscUseTypeMethod(mat->A, restorerow, lrow, &nzA, pcA, pvA);
1377:   PetscUseTypeMethod(mat->B, restorerow, lrow, &nzB, pcB, pvB);
1378:   PetscFunctionReturn(PETSC_SUCCESS);
1379: }

1381: static PetscErrorCode MatRestoreRow_MPIBAIJ(Mat mat, PetscInt row, PetscInt *nz, PetscInt **idx, PetscScalar **v)
1382: {
1383:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)mat->data;

1385:   PetscFunctionBegin;
1386:   PetscCheck(baij->getrowactive, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "MatGetRow not called");
1387:   baij->getrowactive = PETSC_FALSE;
1388:   PetscFunctionReturn(PETSC_SUCCESS);
1389: }

1391: static PetscErrorCode MatZeroEntries_MPIBAIJ(Mat A)
1392: {
1393:   Mat_MPIBAIJ *l = (Mat_MPIBAIJ *)A->data;

1395:   PetscFunctionBegin;
1396:   PetscCall(MatZeroEntries(l->A));
1397:   PetscCall(MatZeroEntries(l->B));
1398:   PetscFunctionReturn(PETSC_SUCCESS);
1399: }

1401: static PetscErrorCode MatGetInfo_MPIBAIJ(Mat matin, MatInfoType flag, MatInfo *info)
1402: {
1403:   Mat_MPIBAIJ   *a = (Mat_MPIBAIJ *)matin->data;
1404:   Mat            A = a->A, B = a->B;
1405:   PetscLogDouble isend[5], irecv[5];

1407:   PetscFunctionBegin;
1408:   info->block_size = (PetscReal)matin->rmap->bs;

1410:   PetscCall(MatGetInfo(A, MAT_LOCAL, info));

1412:   isend[0] = info->nz_used;
1413:   isend[1] = info->nz_allocated;
1414:   isend[2] = info->nz_unneeded;
1415:   isend[3] = info->memory;
1416:   isend[4] = info->mallocs;

1418:   PetscCall(MatGetInfo(B, MAT_LOCAL, info));

1420:   isend[0] += info->nz_used;
1421:   isend[1] += info->nz_allocated;
1422:   isend[2] += info->nz_unneeded;
1423:   isend[3] += info->memory;
1424:   isend[4] += info->mallocs;

1426:   if (flag == MAT_LOCAL) {
1427:     info->nz_used      = isend[0];
1428:     info->nz_allocated = isend[1];
1429:     info->nz_unneeded  = isend[2];
1430:     info->memory       = isend[3];
1431:     info->mallocs      = isend[4];
1432:   } else if (flag == MAT_GLOBAL_MAX) {
1433:     PetscCallMPI(MPIU_Allreduce(isend, irecv, 5, MPIU_PETSCLOGDOUBLE, MPI_MAX, PetscObjectComm((PetscObject)matin)));

1435:     info->nz_used      = irecv[0];
1436:     info->nz_allocated = irecv[1];
1437:     info->nz_unneeded  = irecv[2];
1438:     info->memory       = irecv[3];
1439:     info->mallocs      = irecv[4];
1440:   } else if (flag == MAT_GLOBAL_SUM) {
1441:     PetscCallMPI(MPIU_Allreduce(isend, irecv, 5, MPIU_PETSCLOGDOUBLE, MPI_SUM, PetscObjectComm((PetscObject)matin)));

1443:     info->nz_used      = irecv[0];
1444:     info->nz_allocated = irecv[1];
1445:     info->nz_unneeded  = irecv[2];
1446:     info->memory       = irecv[3];
1447:     info->mallocs      = irecv[4];
1448:   } else SETERRQ(PetscObjectComm((PetscObject)matin), PETSC_ERR_ARG_WRONG, "Unknown MatInfoType argument %d", (int)flag);
1449:   info->fill_ratio_given  = 0; /* no parallel LU/ILU/Cholesky */
1450:   info->fill_ratio_needed = 0;
1451:   info->factor_mallocs    = 0;
1452:   PetscFunctionReturn(PETSC_SUCCESS);
1453: }

1455: static PetscErrorCode MatSetOption_MPIBAIJ(Mat A, MatOption op, PetscBool flg)
1456: {
1457:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

1459:   PetscFunctionBegin;
1460:   switch (op) {
1461:   case MAT_NEW_NONZERO_LOCATIONS:
1462:   case MAT_NEW_NONZERO_ALLOCATION_ERR:
1463:   case MAT_UNUSED_NONZERO_LOCATION_ERR:
1464:   case MAT_KEEP_NONZERO_PATTERN:
1465:   case MAT_NEW_NONZERO_LOCATION_ERR:
1466:     MatCheckPreallocated(A, 1);
1467:     PetscCall(MatSetOption(a->A, op, flg));
1468:     PetscCall(MatSetOption(a->B, op, flg));
1469:     break;
1470:   case MAT_ROW_ORIENTED:
1471:     MatCheckPreallocated(A, 1);
1472:     a->roworiented = flg;

1474:     PetscCall(MatSetOption(a->A, op, flg));
1475:     PetscCall(MatSetOption(a->B, op, flg));
1476:     break;
1477:   case MAT_IGNORE_OFF_PROC_ENTRIES:
1478:     a->donotstash = flg;
1479:     break;
1480:   case MAT_USE_HASH_TABLE:
1481:     a->ht_flag = flg;
1482:     a->ht_fact = 1.39;
1483:     break;
1484:   case MAT_SPD:
1485:   case MAT_SYMMETRIC:
1486:   case MAT_STRUCTURALLY_SYMMETRIC:
1487:   case MAT_HERMITIAN:
1488:   case MAT_SYMMETRY_ETERNAL:
1489:   case MAT_STRUCTURAL_SYMMETRY_ETERNAL:
1490:   case MAT_SPD_ETERNAL:
1491:     /* if the diagonal matrix is square it inherits some of the properties above */
1492:     if (a->A && A->rmap->n == A->cmap->n) PetscCall(MatSetOption(a->A, op, flg));
1493:     break;
1494:   default:
1495:     break;
1496:   }
1497:   PetscFunctionReturn(PETSC_SUCCESS);
1498: }

1500: static PetscErrorCode MatTranspose_MPIBAIJ(Mat A, MatReuse reuse, Mat *matout)
1501: {
1502:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)A->data;
1503:   Mat_SeqBAIJ *Aloc;
1504:   Mat          B;
1505:   PetscInt     M = A->rmap->N, N = A->cmap->N, *ai, *aj, i, *rvals, j, k, col;
1506:   PetscInt     bs = A->rmap->bs, mbs = baij->mbs;
1507:   MatScalar   *a;

1509:   PetscFunctionBegin;
1510:   if (reuse == MAT_REUSE_MATRIX) PetscCall(MatTransposeCheckNonzeroState_Private(A, *matout));
1511:   if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_INPLACE_MATRIX) {
1512:     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
1513:     PetscCall(MatSetSizes(B, A->cmap->n, A->rmap->n, N, M));
1514:     PetscCall(MatSetType(B, ((PetscObject)A)->type_name));
1515:     /* Do not know preallocation information, but must set block size */
1516:     PetscCall(MatMPIBAIJSetPreallocation(B, A->rmap->bs, PETSC_DECIDE, NULL, PETSC_DECIDE, NULL));
1517:   } else {
1518:     B = *matout;
1519:   }

1521:   /* copy over the A part */
1522:   Aloc = (Mat_SeqBAIJ *)baij->A->data;
1523:   ai   = Aloc->i;
1524:   aj   = Aloc->j;
1525:   a    = Aloc->a;
1526:   PetscCall(PetscMalloc1(bs, &rvals));

1528:   for (i = 0; i < mbs; i++) {
1529:     rvals[0] = bs * (baij->rstartbs + i);
1530:     for (j = 1; j < bs; j++) rvals[j] = rvals[j - 1] + 1;
1531:     for (j = ai[i]; j < ai[i + 1]; j++) {
1532:       col = (baij->cstartbs + aj[j]) * bs;
1533:       for (k = 0; k < bs; k++) {
1534:         PetscCall(MatSetValues_MPIBAIJ(B, 1, &col, bs, rvals, a, INSERT_VALUES));

1536:         col++;
1537:         a += bs;
1538:       }
1539:     }
1540:   }
1541:   /* copy over the B part */
1542:   Aloc = (Mat_SeqBAIJ *)baij->B->data;
1543:   ai   = Aloc->i;
1544:   aj   = Aloc->j;
1545:   a    = Aloc->a;
1546:   for (i = 0; i < mbs; i++) {
1547:     rvals[0] = bs * (baij->rstartbs + i);
1548:     for (j = 1; j < bs; j++) rvals[j] = rvals[j - 1] + 1;
1549:     for (j = ai[i]; j < ai[i + 1]; j++) {
1550:       col = baij->garray[aj[j]] * bs;
1551:       for (k = 0; k < bs; k++) {
1552:         PetscCall(MatSetValues_MPIBAIJ(B, 1, &col, bs, rvals, a, INSERT_VALUES));
1553:         col++;
1554:         a += bs;
1555:       }
1556:     }
1557:   }
1558:   PetscCall(PetscFree(rvals));
1559:   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
1560:   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));

1562:   if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_REUSE_MATRIX) *matout = B;
1563:   else PetscCall(MatHeaderMerge(A, &B));
1564:   PetscFunctionReturn(PETSC_SUCCESS);
1565: }

1567: static PetscErrorCode MatDiagonalScale_MPIBAIJ(Mat mat, Vec ll, Vec rr)
1568: {
1569:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)mat->data;
1570:   Mat          a = baij->A, b = baij->B;
1571:   PetscInt     s1, s2, s3;

1573:   PetscFunctionBegin;
1574:   PetscCall(MatGetLocalSize(mat, &s2, &s3));
1575:   if (rr) {
1576:     PetscCall(VecGetLocalSize(rr, &s1));
1577:     PetscCheck(s1 == s3, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "right vector non-conforming local size");
1578:     /* Overlap communication with computation. */
1579:     PetscCall(VecScatterBegin(baij->Mvctx, rr, baij->lvec, INSERT_VALUES, SCATTER_FORWARD));
1580:   }
1581:   if (ll) {
1582:     PetscCall(VecGetLocalSize(ll, &s1));
1583:     PetscCheck(s1 == s2, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "left vector non-conforming local size");
1584:     PetscUseTypeMethod(b, diagonalscale, ll, NULL);
1585:   }
1586:   /* scale  the diagonal block */
1587:   PetscUseTypeMethod(a, diagonalscale, ll, rr);

1589:   if (rr) {
1590:     /* Do a scatter end and then right scale the off-diagonal block */
1591:     PetscCall(VecScatterEnd(baij->Mvctx, rr, baij->lvec, INSERT_VALUES, SCATTER_FORWARD));
1592:     PetscUseTypeMethod(b, diagonalscale, NULL, baij->lvec);
1593:   }
1594:   PetscFunctionReturn(PETSC_SUCCESS);
1595: }

1597: static PetscErrorCode MatZeroRows_MPIBAIJ(Mat A, PetscInt N, const PetscInt rows[], PetscScalar diag, Vec x, Vec b)
1598: {
1599:   Mat_MPIBAIJ *l = (Mat_MPIBAIJ *)A->data;
1600:   PetscInt    *lrows;
1601:   PetscInt     r, len;
1602:   PetscBool    cong;

1604:   PetscFunctionBegin;
1605:   /* get locally owned rows */
1606:   PetscCall(MatZeroRowsMapLocal_Private(A, N, rows, &len, &lrows));
1607:   /* fix right-hand side if needed */
1608:   if (x && b) {
1609:     const PetscScalar *xx;
1610:     PetscScalar       *bb;

1612:     PetscCall(VecGetArrayRead(x, &xx));
1613:     PetscCall(VecGetArray(b, &bb));
1614:     for (r = 0; r < len; ++r) bb[lrows[r]] = diag * xx[lrows[r]];
1615:     PetscCall(VecRestoreArrayRead(x, &xx));
1616:     PetscCall(VecRestoreArray(b, &bb));
1617:   }

1619:   /* actually zap the local rows */
1620:   /*
1621:         Zero the required rows. If the "diagonal block" of the matrix
1622:      is square and the user wishes to set the diagonal we use separate
1623:      code so that MatSetValues() is not called for each diagonal allocating
1624:      new memory, thus calling lots of mallocs and slowing things down.

1626:   */
1627:   /* must zero l->B before l->A because the (diag) case below may put values into l->B*/
1628:   PetscCall(MatZeroRows_SeqBAIJ(l->B, len, lrows, 0.0, NULL, NULL));
1629:   PetscCall(MatHasCongruentLayouts(A, &cong));
1630:   if ((diag != 0.0) && cong) {
1631:     PetscCall(MatZeroRows_SeqBAIJ(l->A, len, lrows, diag, NULL, NULL));
1632:   } else if (diag != 0.0) {
1633:     PetscCall(MatZeroRows_SeqBAIJ(l->A, len, lrows, 0.0, NULL, NULL));
1634:     PetscCheck(!((Mat_SeqBAIJ *)l->A->data)->nonew, PETSC_COMM_SELF, PETSC_ERR_SUP, "MatZeroRows() on rectangular matrices cannot be used with the Mat options MAT_NEW_NONZERO_LOCATIONS, MAT_NEW_NONZERO_LOCATION_ERR, and MAT_NEW_NONZERO_ALLOCATION_ERR");
1635:     for (r = 0; r < len; ++r) {
1636:       const PetscInt row = lrows[r] + A->rmap->rstart;
1637:       PetscCall(MatSetValues(A, 1, &row, 1, &row, &diag, INSERT_VALUES));
1638:     }
1639:     PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
1640:     PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
1641:   } else {
1642:     PetscCall(MatZeroRows_SeqBAIJ(l->A, len, lrows, 0.0, NULL, NULL));
1643:   }
1644:   PetscCall(PetscFree(lrows));

1646:   /* only change matrix nonzero state if pattern was allowed to be changed */
1647:   if (!((Mat_SeqBAIJ *)l->A->data)->keepnonzeropattern || !((Mat_SeqBAIJ *)l->A->data)->nonew) {
1648:     PetscObjectState state = l->A->nonzerostate + l->B->nonzerostate;
1649:     PetscCallMPI(MPIU_Allreduce(&state, &A->nonzerostate, 1, MPIU_INT64, MPI_SUM, PetscObjectComm((PetscObject)A)));
1650:   }
1651:   PetscFunctionReturn(PETSC_SUCCESS);
1652: }

1654: static PetscErrorCode MatZeroRowsColumns_MPIBAIJ(Mat A, PetscInt N, const PetscInt rows[], PetscScalar diag, Vec x, Vec b)
1655: {
1656:   Mat_MPIBAIJ       *l = (Mat_MPIBAIJ *)A->data;
1657:   PetscMPIInt        n, p = 0;
1658:   PetscInt           i, j, k, r, len = 0, row, col, count;
1659:   PetscInt          *lrows, *owners = A->rmap->range;
1660:   PetscSFNode       *rrows;
1661:   PetscSF            sf;
1662:   const PetscScalar *xx;
1663:   PetscScalar       *bb, *mask;
1664:   Vec                xmask, lmask;
1665:   Mat_SeqBAIJ       *baij = (Mat_SeqBAIJ *)l->B->data;
1666:   PetscInt           bs = A->rmap->bs, bs2 = baij->bs2;
1667:   PetscScalar       *aa;

1669:   PetscFunctionBegin;
1670:   PetscCall(PetscMPIIntCast(A->rmap->n, &n));
1671:   /* create PetscSF where leaves are input rows and roots are owned rows */
1672:   PetscCall(PetscMalloc1(n, &lrows));
1673:   for (r = 0; r < n; ++r) lrows[r] = -1;
1674:   PetscCall(PetscMalloc1(N, &rrows));
1675:   for (r = 0; r < N; ++r) {
1676:     const PetscInt idx = rows[r];
1677:     PetscCheck(idx >= 0 && A->rmap->N > idx, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Row %" PetscInt_FMT " out of range [0,%" PetscInt_FMT ")", idx, A->rmap->N);
1678:     if (idx < owners[p] || owners[p + 1] <= idx) { /* short-circuit the search if the last p owns this row too */
1679:       PetscCall(PetscLayoutFindOwner(A->rmap, idx, &p));
1680:     }
1681:     rrows[r].rank  = p;
1682:     rrows[r].index = rows[r] - owners[p];
1683:   }
1684:   PetscCall(PetscSFCreate(PetscObjectComm((PetscObject)A), &sf));
1685:   PetscCall(PetscSFSetGraph(sf, n, N, NULL, PETSC_OWN_POINTER, rrows, PETSC_OWN_POINTER));
1686:   /* collect flags for rows to be zeroed */
1687:   PetscCall(PetscSFReduceBegin(sf, MPIU_INT, (PetscInt *)rows, lrows, MPI_LOR));
1688:   PetscCall(PetscSFReduceEnd(sf, MPIU_INT, (PetscInt *)rows, lrows, MPI_LOR));
1689:   PetscCall(PetscSFDestroy(&sf));
1690:   /* compress and put in row numbers */
1691:   for (r = 0; r < n; ++r)
1692:     if (lrows[r] >= 0) lrows[len++] = r;
1693:   /* zero diagonal part of matrix */
1694:   PetscCall(MatZeroRowsColumns(l->A, len, lrows, diag, x, b));
1695:   /* handle off-diagonal part of matrix */
1696:   PetscCall(MatCreateVecs(A, &xmask, NULL));
1697:   PetscCall(VecDuplicate(l->lvec, &lmask));
1698:   PetscCall(VecGetArray(xmask, &bb));
1699:   for (i = 0; i < len; i++) bb[lrows[i]] = 1;
1700:   PetscCall(VecRestoreArray(xmask, &bb));
1701:   PetscCall(VecScatterBegin(l->Mvctx, xmask, lmask, ADD_VALUES, SCATTER_FORWARD));
1702:   PetscCall(VecScatterEnd(l->Mvctx, xmask, lmask, ADD_VALUES, SCATTER_FORWARD));
1703:   PetscCall(VecDestroy(&xmask));
1704:   if (x) {
1705:     PetscCall(VecScatterBegin(l->Mvctx, x, l->lvec, INSERT_VALUES, SCATTER_FORWARD));
1706:     PetscCall(VecScatterEnd(l->Mvctx, x, l->lvec, INSERT_VALUES, SCATTER_FORWARD));
1707:     PetscCall(VecGetArrayRead(l->lvec, &xx));
1708:     PetscCall(VecGetArray(b, &bb));
1709:   }
1710:   PetscCall(VecGetArray(lmask, &mask));
1711:   /* remove zeroed rows of off-diagonal matrix */
1712:   for (i = 0; i < len; ++i) {
1713:     row   = lrows[i];
1714:     count = (baij->i[row / bs + 1] - baij->i[row / bs]) * bs;
1715:     aa    = baij->a + baij->i[row / bs] * bs2 + (row % bs);
1716:     for (k = 0; k < count; ++k) {
1717:       aa[0] = 0.0;
1718:       aa += bs;
1719:     }
1720:   }
1721:   /* loop over all elements of off process part of matrix zeroing removed columns */
1722:   for (i = 0; i < l->B->rmap->N; ++i) {
1723:     row = i / bs;
1724:     for (j = baij->i[row]; j < baij->i[row + 1]; ++j) {
1725:       for (k = 0; k < bs; ++k) {
1726:         col = bs * baij->j[j] + k;
1727:         if (PetscAbsScalar(mask[col])) {
1728:           aa = baij->a + j * bs2 + (i % bs) + bs * k;
1729:           if (x) bb[i] -= aa[0] * xx[col];
1730:           aa[0] = 0.0;
1731:         }
1732:       }
1733:     }
1734:   }
1735:   if (x) {
1736:     PetscCall(VecRestoreArray(b, &bb));
1737:     PetscCall(VecRestoreArrayRead(l->lvec, &xx));
1738:   }
1739:   PetscCall(VecRestoreArray(lmask, &mask));
1740:   PetscCall(VecDestroy(&lmask));
1741:   PetscCall(PetscFree(lrows));

1743:   /* only change matrix nonzero state if pattern was allowed to be changed */
1744:   if (!((Mat_SeqBAIJ *)l->A->data)->nonew) {
1745:     PetscObjectState state = l->A->nonzerostate + l->B->nonzerostate;
1746:     PetscCallMPI(MPIU_Allreduce(&state, &A->nonzerostate, 1, MPIU_INT64, MPI_SUM, PetscObjectComm((PetscObject)A)));
1747:   }
1748:   PetscFunctionReturn(PETSC_SUCCESS);
1749: }

1751: static PetscErrorCode MatSetUnfactored_MPIBAIJ(Mat A)
1752: {
1753:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

1755:   PetscFunctionBegin;
1756:   PetscCall(MatSetUnfactored(a->A));
1757:   PetscFunctionReturn(PETSC_SUCCESS);
1758: }

1760: static PetscErrorCode MatDuplicate_MPIBAIJ(Mat, MatDuplicateOption, Mat *);

1762: static PetscErrorCode MatEqual_MPIBAIJ(Mat A, Mat B, PetscBool *flag)
1763: {
1764:   Mat_MPIBAIJ *matB = (Mat_MPIBAIJ *)B->data, *matA = (Mat_MPIBAIJ *)A->data;
1765:   Mat          a, b, c, d;
1766:   PetscBool    flg;

1768:   PetscFunctionBegin;
1769:   a = matA->A;
1770:   b = matA->B;
1771:   c = matB->A;
1772:   d = matB->B;

1774:   PetscCall(MatEqual(a, c, &flg));
1775:   if (flg) PetscCall(MatEqual(b, d, &flg));
1776:   PetscCallMPI(MPIU_Allreduce(&flg, flag, 1, MPI_C_BOOL, MPI_LAND, PetscObjectComm((PetscObject)A)));
1777:   PetscFunctionReturn(PETSC_SUCCESS);
1778: }

1780: static PetscErrorCode MatCopy_MPIBAIJ(Mat A, Mat B, MatStructure str)
1781: {
1782:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;
1783:   Mat_MPIBAIJ *b = (Mat_MPIBAIJ *)B->data;

1785:   PetscFunctionBegin;
1786:   /* If the two matrices don't have the same copy implementation, they aren't compatible for fast copy. */
1787:   if ((str != SAME_NONZERO_PATTERN) || (A->ops->copy != B->ops->copy)) {
1788:     PetscCall(MatCopy_Basic(A, B, str));
1789:   } else {
1790:     PetscCall(MatCopy(a->A, b->A, str));
1791:     PetscCall(MatCopy(a->B, b->B, str));
1792:   }
1793:   PetscCall(PetscObjectStateIncrease((PetscObject)B));
1794:   PetscFunctionReturn(PETSC_SUCCESS);
1795: }

1797: PetscErrorCode MatAXPYGetPreallocation_MPIBAIJ(Mat Y, const PetscInt *yltog, Mat X, const PetscInt *xltog, PetscInt *nnz)
1798: {
1799:   PetscInt     bs = Y->rmap->bs, m = Y->rmap->N / bs;
1800:   Mat_SeqBAIJ *x = (Mat_SeqBAIJ *)X->data;
1801:   Mat_SeqBAIJ *y = (Mat_SeqBAIJ *)Y->data;

1803:   PetscFunctionBegin;
1804:   PetscCall(MatAXPYGetPreallocation_MPIX_private(m, x->i, x->j, xltog, y->i, y->j, yltog, nnz));
1805:   PetscFunctionReturn(PETSC_SUCCESS);
1806: }

1808: static PetscErrorCode MatAXPY_MPIBAIJ(Mat Y, PetscScalar a, Mat X, MatStructure str)
1809: {
1810:   Mat_MPIBAIJ *xx = (Mat_MPIBAIJ *)X->data, *yy = (Mat_MPIBAIJ *)Y->data;
1811:   PetscBLASInt bnz, one                         = 1;
1812:   Mat_SeqBAIJ *x, *y;
1813:   PetscInt     bs2 = Y->rmap->bs * Y->rmap->bs;

1815:   PetscFunctionBegin;
1816:   if (str == SAME_NONZERO_PATTERN) {
1817:     PetscScalar alpha = a;
1818:     x                 = (Mat_SeqBAIJ *)xx->A->data;
1819:     y                 = (Mat_SeqBAIJ *)yy->A->data;
1820:     PetscCall(PetscBLASIntCast(x->nz * bs2, &bnz));
1821:     PetscCallBLAS("BLASaxpy", BLASaxpy_(&bnz, &alpha, x->a, &one, y->a, &one));
1822:     x = (Mat_SeqBAIJ *)xx->B->data;
1823:     y = (Mat_SeqBAIJ *)yy->B->data;
1824:     PetscCall(PetscBLASIntCast(x->nz * bs2, &bnz));
1825:     PetscCallBLAS("BLASaxpy", BLASaxpy_(&bnz, &alpha, x->a, &one, y->a, &one));
1826:     PetscCall(PetscObjectStateIncrease((PetscObject)Y));
1827:   } else if (str == SUBSET_NONZERO_PATTERN) { /* nonzeros of X is a subset of Y's */
1828:     PetscCall(MatAXPY_Basic(Y, a, X, str));
1829:   } else {
1830:     Mat       B;
1831:     PetscInt *nnz_d, *nnz_o, bs = Y->rmap->bs;
1832:     PetscCall(PetscMalloc1(yy->A->rmap->N, &nnz_d));
1833:     PetscCall(PetscMalloc1(yy->B->rmap->N, &nnz_o));
1834:     PetscCall(MatCreate(PetscObjectComm((PetscObject)Y), &B));
1835:     PetscCall(PetscObjectSetName((PetscObject)B, ((PetscObject)Y)->name));
1836:     PetscCall(MatSetSizes(B, Y->rmap->n, Y->cmap->n, Y->rmap->N, Y->cmap->N));
1837:     PetscCall(MatSetBlockSizesFromMats(B, Y, Y));
1838:     PetscCall(MatSetType(B, MATMPIBAIJ));
1839:     PetscCall(MatAXPYGetPreallocation_SeqBAIJ(yy->A, xx->A, nnz_d));
1840:     PetscCall(MatAXPYGetPreallocation_MPIBAIJ(yy->B, yy->garray, xx->B, xx->garray, nnz_o));
1841:     PetscCall(MatMPIBAIJSetPreallocation(B, bs, 0, nnz_d, 0, nnz_o));
1842:     /* MatAXPY_BasicWithPreallocation() for BAIJ matrix is much slower than AIJ, even for bs=1 ! */
1843:     PetscCall(MatAXPY_BasicWithPreallocation(B, Y, a, X, str));
1844:     PetscCall(MatHeaderMerge(Y, &B));
1845:     PetscCall(PetscFree(nnz_d));
1846:     PetscCall(PetscFree(nnz_o));
1847:   }
1848:   PetscFunctionReturn(PETSC_SUCCESS);
1849: }

1851: static PetscErrorCode MatConjugate_MPIBAIJ(Mat mat)
1852: {
1853:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)mat->data;

1855:   PetscFunctionBegin;
1856:   PetscCall(MatConjugate_SeqBAIJ(a->A));
1857:   PetscCall(MatConjugate_SeqBAIJ(a->B));
1858:   PetscFunctionReturn(PETSC_SUCCESS);
1859: }

1861: static PetscErrorCode MatRealPart_MPIBAIJ(Mat A)
1862: {
1863:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

1865:   PetscFunctionBegin;
1866:   PetscCall(MatRealPart(a->A));
1867:   PetscCall(MatRealPart(a->B));
1868:   PetscFunctionReturn(PETSC_SUCCESS);
1869: }

1871: static PetscErrorCode MatImaginaryPart_MPIBAIJ(Mat A)
1872: {
1873:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

1875:   PetscFunctionBegin;
1876:   PetscCall(MatImaginaryPart(a->A));
1877:   PetscCall(MatImaginaryPart(a->B));
1878:   PetscFunctionReturn(PETSC_SUCCESS);
1879: }

1881: static PetscErrorCode MatCreateSubMatrix_MPIBAIJ(Mat mat, IS isrow, IS iscol, MatReuse call, Mat *newmat)
1882: {
1883:   IS       iscol_local;
1884:   PetscInt csize;

1886:   PetscFunctionBegin;
1887:   PetscCall(ISGetLocalSize(iscol, &csize));
1888:   if (call == MAT_REUSE_MATRIX) {
1889:     PetscCall(PetscObjectQuery((PetscObject)*newmat, "ISAllGather", (PetscObject *)&iscol_local));
1890:     PetscCheck(iscol_local, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Submatrix passed in was not used before, cannot reuse");
1891:   } else {
1892:     PetscCall(ISAllGather(iscol, &iscol_local));
1893:   }
1894:   PetscCall(MatCreateSubMatrix_MPIBAIJ_Private(mat, isrow, iscol_local, csize, call, newmat, PETSC_FALSE));
1895:   if (call == MAT_INITIAL_MATRIX) {
1896:     PetscCall(PetscObjectCompose((PetscObject)*newmat, "ISAllGather", (PetscObject)iscol_local));
1897:     PetscCall(ISDestroy(&iscol_local));
1898:   }
1899:   PetscFunctionReturn(PETSC_SUCCESS);
1900: }

1902: /*
1903:   Not great since it makes two copies of the submatrix, first an SeqBAIJ
1904:   in local and then by concatenating the local matrices the end result.
1905:   Writing it directly would be much like MatCreateSubMatrices_MPIBAIJ().
1906:   This routine is used for BAIJ and SBAIJ matrices (unfortunate dependency).
1907: */
1908: PetscErrorCode MatCreateSubMatrix_MPIBAIJ_Private(Mat mat, IS isrow, IS iscol, PetscInt csize, MatReuse call, Mat *newmat, PetscBool sym)
1909: {
1910:   PetscMPIInt  rank, size;
1911:   PetscInt     i, m, n, rstart, row, rend, nz, *cwork, j, bs;
1912:   PetscInt    *ii, *jj, nlocal, *dlens, *olens, dlen, olen, jend, mglobal;
1913:   Mat          M, Mreuse;
1914:   MatScalar   *vwork, *aa;
1915:   MPI_Comm     comm;
1916:   IS           isrow_new, iscol_new;
1917:   Mat_SeqBAIJ *aij;

1919:   PetscFunctionBegin;
1920:   PetscCall(PetscObjectGetComm((PetscObject)mat, &comm));
1921:   PetscCallMPI(MPI_Comm_rank(comm, &rank));
1922:   PetscCallMPI(MPI_Comm_size(comm, &size));
1923:   /* The compression and expansion should be avoided. Doesn't point
1924:      out errors, might change the indices, hence buggey */
1925:   PetscCall(ISCompressIndicesGeneral(mat->rmap->N, mat->rmap->n, mat->rmap->bs, 1, &isrow, &isrow_new));
1926:   if (isrow == iscol) {
1927:     iscol_new = isrow_new;
1928:     PetscCall(PetscObjectReference((PetscObject)iscol_new));
1929:   } else PetscCall(ISCompressIndicesGeneral(mat->cmap->N, mat->cmap->n, mat->cmap->bs, 1, &iscol, &iscol_new));

1931:   if (call == MAT_REUSE_MATRIX) {
1932:     PetscCall(PetscObjectQuery((PetscObject)*newmat, "SubMatrix", (PetscObject *)&Mreuse));
1933:     PetscCheck(Mreuse, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Submatrix passed in was not used before, cannot reuse");
1934:     PetscCall(MatCreateSubMatrices_MPIBAIJ_local(mat, 1, &isrow_new, &iscol_new, MAT_REUSE_MATRIX, &Mreuse, sym));
1935:   } else {
1936:     PetscCall(MatCreateSubMatrices_MPIBAIJ_local(mat, 1, &isrow_new, &iscol_new, MAT_INITIAL_MATRIX, &Mreuse, sym));
1937:   }
1938:   PetscCall(ISDestroy(&isrow_new));
1939:   PetscCall(ISDestroy(&iscol_new));
1940:   /*
1941:       m - number of local rows
1942:       n - number of columns (same on all processors)
1943:       rstart - first row in new global matrix generated
1944:   */
1945:   PetscCall(MatGetBlockSize(mat, &bs));
1946:   PetscCall(MatGetSize(Mreuse, &m, &n));
1947:   m = m / bs;
1948:   n = n / bs;

1950:   if (call == MAT_INITIAL_MATRIX) {
1951:     aij = (Mat_SeqBAIJ *)Mreuse->data;
1952:     ii  = aij->i;
1953:     jj  = aij->j;

1955:     /*
1956:         Determine the number of non-zeros in the diagonal and off-diagonal
1957:         portions of the matrix in order to do correct preallocation
1958:     */

1960:     /* first get start and end of "diagonal" columns */
1961:     if (csize == PETSC_DECIDE) {
1962:       PetscCall(ISGetSize(isrow, &mglobal));
1963:       if (mglobal == n * bs) { /* square matrix */
1964:         nlocal = m;
1965:       } else {
1966:         nlocal = n / size + ((n % size) > rank);
1967:       }
1968:     } else {
1969:       nlocal = csize / bs;
1970:     }
1971:     PetscCallMPI(MPI_Scan(&nlocal, &rend, 1, MPIU_INT, MPI_SUM, comm));
1972:     rstart = rend - nlocal;
1973:     PetscCheck(rank != size - 1 || rend == n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Local column sizes %" PetscInt_FMT " do not add up to total number of columns %" PetscInt_FMT, rend, n);

1975:     /* next, compute all the lengths */
1976:     PetscCall(PetscMalloc2(m + 1, &dlens, m + 1, &olens));
1977:     for (i = 0; i < m; i++) {
1978:       jend = ii[i + 1] - ii[i];
1979:       olen = 0;
1980:       dlen = 0;
1981:       for (j = 0; j < jend; j++) {
1982:         if (*jj < rstart || *jj >= rend) olen++;
1983:         else dlen++;
1984:         jj++;
1985:       }
1986:       olens[i] = olen;
1987:       dlens[i] = dlen;
1988:     }
1989:     PetscCall(MatCreate(comm, &M));
1990:     PetscCall(MatSetSizes(M, bs * m, bs * nlocal, PETSC_DECIDE, bs * n));
1991:     PetscCall(MatSetType(M, sym ? ((PetscObject)mat)->type_name : MATMPIBAIJ));
1992:     PetscCall(MatMPIBAIJSetPreallocation(M, bs, 0, dlens, 0, olens));
1993:     PetscCall(MatMPISBAIJSetPreallocation(M, bs, 0, dlens, 0, olens));
1994:     PetscCall(PetscFree2(dlens, olens));
1995:   } else {
1996:     PetscInt ml, nl;

1998:     M = *newmat;
1999:     PetscCall(MatGetLocalSize(M, &ml, &nl));
2000:     PetscCheck(ml == m, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Previous matrix must be same size/layout as request");
2001:     PetscCall(MatZeroEntries(M));
2002:     /*
2003:          The next two lines are needed so we may call MatSetValues_MPIAIJ() below directly,
2004:        rather than the slower MatSetValues().
2005:     */
2006:     M->was_assembled = PETSC_TRUE;
2007:     M->assembled     = PETSC_FALSE;
2008:   }
2009:   PetscCall(MatSetOption(M, MAT_ROW_ORIENTED, PETSC_FALSE));
2010:   PetscCall(MatGetOwnershipRange(M, &rstart, &rend));
2011:   aij = (Mat_SeqBAIJ *)Mreuse->data;
2012:   ii  = aij->i;
2013:   jj  = aij->j;
2014:   aa  = aij->a;
2015:   for (i = 0; i < m; i++) {
2016:     row   = rstart / bs + i;
2017:     nz    = ii[i + 1] - ii[i];
2018:     cwork = jj;
2019:     jj    = PetscSafePointerPlusOffset(jj, nz);
2020:     vwork = aa;
2021:     aa    = PetscSafePointerPlusOffset(aa, nz * bs * bs);
2022:     PetscCall(MatSetValuesBlocked_MPIBAIJ(M, 1, &row, nz, cwork, vwork, INSERT_VALUES));
2023:   }

2025:   PetscCall(MatAssemblyBegin(M, MAT_FINAL_ASSEMBLY));
2026:   PetscCall(MatAssemblyEnd(M, MAT_FINAL_ASSEMBLY));
2027:   *newmat = M;

2029:   /* save submatrix used in processor for next request */
2030:   if (call == MAT_INITIAL_MATRIX) {
2031:     PetscCall(PetscObjectCompose((PetscObject)M, "SubMatrix", (PetscObject)Mreuse));
2032:     PetscCall(PetscObjectDereference((PetscObject)Mreuse));
2033:   }
2034:   PetscFunctionReturn(PETSC_SUCCESS);
2035: }

2037: static PetscErrorCode MatPermute_MPIBAIJ(Mat A, IS rowp, IS colp, Mat *B)
2038: {
2039:   MPI_Comm        comm, pcomm;
2040:   PetscInt        clocal_size, nrows;
2041:   const PetscInt *rows;
2042:   PetscMPIInt     size;
2043:   IS              crowp, lcolp;

2045:   PetscFunctionBegin;
2046:   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
2047:   /* make a collective version of 'rowp' */
2048:   PetscCall(PetscObjectGetComm((PetscObject)rowp, &pcomm));
2049:   if (pcomm == comm) {
2050:     crowp = rowp;
2051:   } else {
2052:     PetscCall(ISGetSize(rowp, &nrows));
2053:     PetscCall(ISGetIndices(rowp, &rows));
2054:     PetscCall(ISCreateGeneral(comm, nrows, rows, PETSC_COPY_VALUES, &crowp));
2055:     PetscCall(ISRestoreIndices(rowp, &rows));
2056:   }
2057:   PetscCall(ISSetPermutation(crowp));
2058:   /* make a local version of 'colp' */
2059:   PetscCall(PetscObjectGetComm((PetscObject)colp, &pcomm));
2060:   PetscCallMPI(MPI_Comm_size(pcomm, &size));
2061:   if (size == 1) {
2062:     lcolp = colp;
2063:   } else {
2064:     PetscCall(ISAllGather(colp, &lcolp));
2065:   }
2066:   PetscCall(ISSetPermutation(lcolp));
2067:   /* now we just get the submatrix */
2068:   PetscCall(MatGetLocalSize(A, NULL, &clocal_size));
2069:   PetscCall(MatCreateSubMatrix_MPIBAIJ_Private(A, crowp, lcolp, clocal_size, MAT_INITIAL_MATRIX, B, PETSC_FALSE));
2070:   /* clean up */
2071:   if (pcomm != comm) PetscCall(ISDestroy(&crowp));
2072:   if (size > 1) PetscCall(ISDestroy(&lcolp));
2073:   PetscFunctionReturn(PETSC_SUCCESS);
2074: }

2076: static PetscErrorCode MatGetGhosts_MPIBAIJ(Mat mat, PetscInt *nghosts, const PetscInt *ghosts[])
2077: {
2078:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)mat->data;
2079:   Mat_SeqBAIJ *B    = (Mat_SeqBAIJ *)baij->B->data;

2081:   PetscFunctionBegin;
2082:   if (nghosts) *nghosts = B->nbs;
2083:   if (ghosts) *ghosts = baij->garray;
2084:   PetscFunctionReturn(PETSC_SUCCESS);
2085: }

2087: static PetscErrorCode MatGetSeqNonzeroStructure_MPIBAIJ(Mat A, Mat *newmat)
2088: {
2089:   Mat          B;
2090:   Mat_MPIBAIJ *a  = (Mat_MPIBAIJ *)A->data;
2091:   Mat_SeqBAIJ *ad = (Mat_SeqBAIJ *)a->A->data, *bd = (Mat_SeqBAIJ *)a->B->data;
2092:   Mat_SeqAIJ  *b;
2093:   PetscMPIInt  size, rank, *recvcounts = NULL, *displs = NULL;
2094:   PetscInt     sendcount, i, *rstarts = A->rmap->range, n, cnt, j, bs = A->rmap->bs;
2095:   PetscInt     m, *garray = a->garray, *lens, *jsendbuf, *a_jsendbuf, *b_jsendbuf;

2097:   PetscFunctionBegin;
2098:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size));
2099:   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)A), &rank));

2101:   /*   Tell every processor the number of nonzeros per row  */
2102:   PetscCall(PetscMalloc1(A->rmap->N / bs, &lens));
2103:   for (i = A->rmap->rstart / bs; i < A->rmap->rend / bs; i++) lens[i] = ad->i[i - A->rmap->rstart / bs + 1] - ad->i[i - A->rmap->rstart / bs] + bd->i[i - A->rmap->rstart / bs + 1] - bd->i[i - A->rmap->rstart / bs];
2104:   PetscCall(PetscMalloc1(2 * size, &recvcounts));
2105:   displs = recvcounts + size;
2106:   for (i = 0; i < size; i++) {
2107:     PetscCall(PetscMPIIntCast(A->rmap->range[i + 1] / bs - A->rmap->range[i] / bs, &recvcounts[i]));
2108:     PetscCall(PetscMPIIntCast(A->rmap->range[i] / bs, &displs[i]));
2109:   }
2110:   PetscCallMPI(MPI_Allgatherv(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, lens, recvcounts, displs, MPIU_INT, PetscObjectComm((PetscObject)A)));
2111:   /* Create the sequential matrix of the same type as the local block diagonal  */
2112:   PetscCall(MatCreate(PETSC_COMM_SELF, &B));
2113:   PetscCall(MatSetSizes(B, A->rmap->N / bs, A->cmap->N / bs, PETSC_DETERMINE, PETSC_DETERMINE));
2114:   PetscCall(MatSetType(B, MATSEQAIJ));
2115:   PetscCall(MatSeqAIJSetPreallocation(B, 0, lens));
2116:   b = (Mat_SeqAIJ *)B->data;

2118:   /*     Copy my part of matrix column indices over  */
2119:   sendcount  = ad->nz + bd->nz;
2120:   jsendbuf   = b->j + b->i[rstarts[rank] / bs];
2121:   a_jsendbuf = ad->j;
2122:   b_jsendbuf = bd->j;
2123:   n          = A->rmap->rend / bs - A->rmap->rstart / bs;
2124:   cnt        = 0;
2125:   for (i = 0; i < n; i++) {
2126:     /* put in lower diagonal portion */
2127:     m = bd->i[i + 1] - bd->i[i];
2128:     while (m > 0) {
2129:       /* is it above diagonal (in bd (compressed) numbering) */
2130:       if (garray[*b_jsendbuf] > A->rmap->rstart / bs + i) break;
2131:       jsendbuf[cnt++] = garray[*b_jsendbuf++];
2132:       m--;
2133:     }

2135:     /* put in diagonal portion */
2136:     for (j = ad->i[i]; j < ad->i[i + 1]; j++) jsendbuf[cnt++] = A->rmap->rstart / bs + *a_jsendbuf++;

2138:     /* put in upper diagonal portion */
2139:     while (m-- > 0) jsendbuf[cnt++] = garray[*b_jsendbuf++];
2140:   }
2141:   PetscCheck(cnt == sendcount, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Corrupted PETSc matrix: nz given %" PetscInt_FMT " actual nz %" PetscInt_FMT, sendcount, cnt);

2143:   /*  Gather all column indices to all processors  */
2144:   for (i = 0; i < size; i++) {
2145:     recvcounts[i] = 0;
2146:     for (j = A->rmap->range[i] / bs; j < A->rmap->range[i + 1] / bs; j++) recvcounts[i] += lens[j];
2147:   }
2148:   displs[0] = 0;
2149:   for (i = 1; i < size; i++) displs[i] = displs[i - 1] + recvcounts[i - 1];
2150:   PetscCallMPI(MPI_Allgatherv(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, b->j, recvcounts, displs, MPIU_INT, PetscObjectComm((PetscObject)A)));
2151:   /*  Assemble the matrix into usable form (note numerical values not yet set)  */
2152:   /* set the b->ilen (length of each row) values */
2153:   PetscCall(PetscArraycpy(b->ilen, lens, A->rmap->N / bs));
2154:   /* set the b->i indices */
2155:   b->i[0] = 0;
2156:   for (i = 1; i <= A->rmap->N / bs; i++) b->i[i] = b->i[i - 1] + lens[i - 1];
2157:   PetscCall(PetscFree(lens));
2158:   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
2159:   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
2160:   PetscCall(PetscFree(recvcounts));

2162:   PetscCall(MatPropagateSymmetryOptions(A, B));
2163:   *newmat = B;
2164:   PetscFunctionReturn(PETSC_SUCCESS);
2165: }

2167: static PetscErrorCode MatSOR_MPIBAIJ(Mat matin, Vec bb, PetscReal omega, MatSORType flag, PetscReal fshift, PetscInt its, PetscInt lits, Vec xx)
2168: {
2169:   Mat_MPIBAIJ *mat = (Mat_MPIBAIJ *)matin->data;
2170:   Vec          bb1 = NULL;

2172:   PetscFunctionBegin;
2173:   if (flag == SOR_APPLY_UPPER) {
2174:     PetscUseTypeMethod(mat->A, sor, bb, omega, flag, fshift, lits, 1, xx);
2175:     PetscFunctionReturn(PETSC_SUCCESS);
2176:   }

2178:   if (its > 1 || ~flag & SOR_ZERO_INITIAL_GUESS) PetscCall(VecDuplicate(bb, &bb1));

2180:   if ((flag & SOR_LOCAL_SYMMETRIC_SWEEP) == SOR_LOCAL_SYMMETRIC_SWEEP) {
2181:     if (flag & SOR_ZERO_INITIAL_GUESS) {
2182:       PetscUseTypeMethod(mat->A, sor, bb, omega, flag, fshift, lits, 1, xx);
2183:       its--;
2184:     }

2186:     while (its--) {
2187:       PetscCall(VecScatterBegin(mat->Mvctx, xx, mat->lvec, INSERT_VALUES, SCATTER_FORWARD));
2188:       PetscCall(VecScatterEnd(mat->Mvctx, xx, mat->lvec, INSERT_VALUES, SCATTER_FORWARD));

2190:       /* update rhs: bb1 = bb - B*x */
2191:       PetscCall(VecScale(mat->lvec, -1.0));
2192:       PetscUseTypeMethod(mat->B, multadd, mat->lvec, bb, bb1);

2194:       /* local sweep */
2195:       PetscUseTypeMethod(mat->A, sor, bb1, omega, SOR_SYMMETRIC_SWEEP, fshift, lits, 1, xx);
2196:     }
2197:   } else if (flag & SOR_LOCAL_FORWARD_SWEEP) {
2198:     if (flag & SOR_ZERO_INITIAL_GUESS) {
2199:       PetscUseTypeMethod(mat->A, sor, bb, omega, flag, fshift, lits, 1, xx);
2200:       its--;
2201:     }
2202:     while (its--) {
2203:       PetscCall(VecScatterBegin(mat->Mvctx, xx, mat->lvec, INSERT_VALUES, SCATTER_FORWARD));
2204:       PetscCall(VecScatterEnd(mat->Mvctx, xx, mat->lvec, INSERT_VALUES, SCATTER_FORWARD));

2206:       /* update rhs: bb1 = bb - B*x */
2207:       PetscCall(VecScale(mat->lvec, -1.0));
2208:       PetscUseTypeMethod(mat->B, multadd, mat->lvec, bb, bb1);

2210:       /* local sweep */
2211:       PetscUseTypeMethod(mat->A, sor, bb1, omega, SOR_FORWARD_SWEEP, fshift, lits, 1, xx);
2212:     }
2213:   } else if (flag & SOR_LOCAL_BACKWARD_SWEEP) {
2214:     if (flag & SOR_ZERO_INITIAL_GUESS) {
2215:       PetscUseTypeMethod(mat->A, sor, bb, omega, flag, fshift, lits, 1, xx);
2216:       its--;
2217:     }
2218:     while (its--) {
2219:       PetscCall(VecScatterBegin(mat->Mvctx, xx, mat->lvec, INSERT_VALUES, SCATTER_FORWARD));
2220:       PetscCall(VecScatterEnd(mat->Mvctx, xx, mat->lvec, INSERT_VALUES, SCATTER_FORWARD));

2222:       /* update rhs: bb1 = bb - B*x */
2223:       PetscCall(VecScale(mat->lvec, -1.0));
2224:       PetscUseTypeMethod(mat->B, multadd, mat->lvec, bb, bb1);

2226:       /* local sweep */
2227:       PetscUseTypeMethod(mat->A, sor, bb1, omega, SOR_BACKWARD_SWEEP, fshift, lits, 1, xx);
2228:     }
2229:   } else SETERRQ(PetscObjectComm((PetscObject)matin), PETSC_ERR_SUP, "Parallel version of SOR requested not supported");

2231:   PetscCall(VecDestroy(&bb1));
2232:   PetscFunctionReturn(PETSC_SUCCESS);
2233: }

2235: static PetscErrorCode MatGetColumnReductions_MPIBAIJ(Mat A, PetscInt type, PetscReal *reductions)
2236: {
2237:   Mat_MPIBAIJ *aij = (Mat_MPIBAIJ *)A->data;
2238:   PetscInt     m, N, i, *garray = aij->garray;
2239:   PetscInt     ib, jb, bs = A->rmap->bs;
2240:   Mat_SeqBAIJ *a_aij = (Mat_SeqBAIJ *)aij->A->data;
2241:   MatScalar   *a_val = a_aij->a;
2242:   Mat_SeqBAIJ *b_aij = (Mat_SeqBAIJ *)aij->B->data;
2243:   MatScalar   *b_val = b_aij->a;
2244:   PetscReal   *work;

2246:   PetscFunctionBegin;
2247:   PetscCall(MatGetSize(A, &m, &N));
2248:   PetscCall(PetscCalloc1(N, &work));
2249:   if (type == NORM_2) {
2250:     for (i = a_aij->i[0]; i < a_aij->i[aij->A->rmap->n / bs]; i++) {
2251:       for (jb = 0; jb < bs; jb++) {
2252:         for (ib = 0; ib < bs; ib++) {
2253:           work[A->cmap->rstart + a_aij->j[i] * bs + jb] += PetscAbsScalar(*a_val * *a_val);
2254:           a_val++;
2255:         }
2256:       }
2257:     }
2258:     for (i = b_aij->i[0]; i < b_aij->i[aij->B->rmap->n / bs]; i++) {
2259:       for (jb = 0; jb < bs; jb++) {
2260:         for (ib = 0; ib < bs; ib++) {
2261:           work[garray[b_aij->j[i]] * bs + jb] += PetscAbsScalar(*b_val * *b_val);
2262:           b_val++;
2263:         }
2264:       }
2265:     }
2266:   } else if (type == NORM_1) {
2267:     for (i = a_aij->i[0]; i < a_aij->i[aij->A->rmap->n / bs]; i++) {
2268:       for (jb = 0; jb < bs; jb++) {
2269:         for (ib = 0; ib < bs; ib++) {
2270:           work[A->cmap->rstart + a_aij->j[i] * bs + jb] += PetscAbsScalar(*a_val);
2271:           a_val++;
2272:         }
2273:       }
2274:     }
2275:     for (i = b_aij->i[0]; i < b_aij->i[aij->B->rmap->n / bs]; i++) {
2276:       for (jb = 0; jb < bs; jb++) {
2277:         for (ib = 0; ib < bs; ib++) {
2278:           work[garray[b_aij->j[i]] * bs + jb] += PetscAbsScalar(*b_val);
2279:           b_val++;
2280:         }
2281:       }
2282:     }
2283:   } else if (type == NORM_INFINITY) {
2284:     for (i = a_aij->i[0]; i < a_aij->i[aij->A->rmap->n / bs]; i++) {
2285:       for (jb = 0; jb < bs; jb++) {
2286:         for (ib = 0; ib < bs; ib++) {
2287:           PetscInt col = A->cmap->rstart + a_aij->j[i] * bs + jb;
2288:           work[col]    = PetscMax(PetscAbsScalar(*a_val), work[col]);
2289:           a_val++;
2290:         }
2291:       }
2292:     }
2293:     for (i = b_aij->i[0]; i < b_aij->i[aij->B->rmap->n / bs]; i++) {
2294:       for (jb = 0; jb < bs; jb++) {
2295:         for (ib = 0; ib < bs; ib++) {
2296:           PetscInt col = garray[b_aij->j[i]] * bs + jb;
2297:           work[col]    = PetscMax(PetscAbsScalar(*b_val), work[col]);
2298:           b_val++;
2299:         }
2300:       }
2301:     }
2302:   } else if (type == REDUCTION_SUM_REALPART || type == REDUCTION_MEAN_REALPART) {
2303:     for (i = a_aij->i[0]; i < a_aij->i[aij->A->rmap->n / bs]; i++) {
2304:       for (jb = 0; jb < bs; jb++) {
2305:         for (ib = 0; ib < bs; ib++) {
2306:           work[A->cmap->rstart + a_aij->j[i] * bs + jb] += PetscRealPart(*a_val);
2307:           a_val++;
2308:         }
2309:       }
2310:     }
2311:     for (i = b_aij->i[0]; i < b_aij->i[aij->B->rmap->n / bs]; i++) {
2312:       for (jb = 0; jb < bs; jb++) {
2313:         for (ib = 0; ib < bs; ib++) {
2314:           work[garray[b_aij->j[i]] * bs + jb] += PetscRealPart(*b_val);
2315:           b_val++;
2316:         }
2317:       }
2318:     }
2319:   } else if (type == REDUCTION_SUM_IMAGINARYPART || type == REDUCTION_MEAN_IMAGINARYPART) {
2320:     for (i = a_aij->i[0]; i < a_aij->i[aij->A->rmap->n / bs]; i++) {
2321:       for (jb = 0; jb < bs; jb++) {
2322:         for (ib = 0; ib < bs; ib++) {
2323:           work[A->cmap->rstart + a_aij->j[i] * bs + jb] += PetscImaginaryPart(*a_val);
2324:           a_val++;
2325:         }
2326:       }
2327:     }
2328:     for (i = b_aij->i[0]; i < b_aij->i[aij->B->rmap->n / bs]; i++) {
2329:       for (jb = 0; jb < bs; jb++) {
2330:         for (ib = 0; ib < bs; ib++) {
2331:           work[garray[b_aij->j[i]] * bs + jb] += PetscImaginaryPart(*b_val);
2332:           b_val++;
2333:         }
2334:       }
2335:     }
2336:   } else SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Unknown reduction type");
2337:   if (type == NORM_INFINITY) {
2338:     PetscCallMPI(MPIU_Allreduce(work, reductions, N, MPIU_REAL, MPIU_MAX, PetscObjectComm((PetscObject)A)));
2339:   } else {
2340:     PetscCallMPI(MPIU_Allreduce(work, reductions, N, MPIU_REAL, MPIU_SUM, PetscObjectComm((PetscObject)A)));
2341:   }
2342:   PetscCall(PetscFree(work));
2343:   if (type == NORM_2) {
2344:     for (i = 0; i < N; i++) reductions[i] = PetscSqrtReal(reductions[i]);
2345:   } else if (type == REDUCTION_MEAN_REALPART || type == REDUCTION_MEAN_IMAGINARYPART) {
2346:     for (i = 0; i < N; i++) reductions[i] /= m;
2347:   }
2348:   PetscFunctionReturn(PETSC_SUCCESS);
2349: }

2351: static PetscErrorCode MatInvertBlockDiagonal_MPIBAIJ(Mat A, const PetscScalar **values)
2352: {
2353:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

2355:   PetscFunctionBegin;
2356:   PetscCall(MatInvertBlockDiagonal(a->A, values));
2357:   A->factorerrortype             = a->A->factorerrortype;
2358:   A->factorerror_zeropivot_value = a->A->factorerror_zeropivot_value;
2359:   A->factorerror_zeropivot_row   = a->A->factorerror_zeropivot_row;
2360:   PetscFunctionReturn(PETSC_SUCCESS);
2361: }

2363: static PetscErrorCode MatShift_MPIBAIJ(Mat Y, PetscScalar a)
2364: {
2365:   Mat_MPIBAIJ *maij = (Mat_MPIBAIJ *)Y->data;
2366:   Mat_SeqBAIJ *aij  = (Mat_SeqBAIJ *)maij->A->data;

2368:   PetscFunctionBegin;
2369:   if (!Y->preallocated) {
2370:     PetscCall(MatMPIBAIJSetPreallocation(Y, Y->rmap->bs, 1, NULL, 0, NULL));
2371:   } else if (!aij->nz) {
2372:     PetscInt nonew = aij->nonew;
2373:     PetscCall(MatSeqBAIJSetPreallocation(maij->A, Y->rmap->bs, 1, NULL));
2374:     aij->nonew = nonew;
2375:   }
2376:   PetscCall(MatShift_Basic(Y, a));
2377:   PetscFunctionReturn(PETSC_SUCCESS);
2378: }

2380: static PetscErrorCode MatGetDiagonalBlock_MPIBAIJ(Mat A, Mat *a)
2381: {
2382:   PetscFunctionBegin;
2383:   *a = ((Mat_MPIBAIJ *)A->data)->A;
2384:   PetscFunctionReturn(PETSC_SUCCESS);
2385: }

2387: static PetscErrorCode MatEliminateZeros_MPIBAIJ(Mat A, PetscBool keep)
2388: {
2389:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

2391:   PetscFunctionBegin;
2392:   PetscCall(MatEliminateZeros_SeqBAIJ(a->A, keep));        // possibly keep zero diagonal coefficients
2393:   PetscCall(MatEliminateZeros_SeqBAIJ(a->B, PETSC_FALSE)); // never keep zero diagonal coefficients
2394:   PetscFunctionReturn(PETSC_SUCCESS);
2395: }

2397: static struct _MatOps MatOps_Values = {MatSetValues_MPIBAIJ,
2398:                                        MatGetRow_MPIBAIJ,
2399:                                        MatRestoreRow_MPIBAIJ,
2400:                                        MatMult_MPIBAIJ,
2401:                                        /* 4*/ MatMultAdd_MPIBAIJ,
2402:                                        MatMultTranspose_MPIBAIJ,
2403:                                        MatMultTransposeAdd_MPIBAIJ,
2404:                                        NULL,
2405:                                        NULL,
2406:                                        NULL,
2407:                                        /*10*/ NULL,
2408:                                        NULL,
2409:                                        NULL,
2410:                                        MatSOR_MPIBAIJ,
2411:                                        MatTranspose_MPIBAIJ,
2412:                                        /*15*/ MatGetInfo_MPIBAIJ,
2413:                                        MatEqual_MPIBAIJ,
2414:                                        MatGetDiagonal_MPIBAIJ,
2415:                                        MatDiagonalScale_MPIBAIJ,
2416:                                        MatNorm_MPIBAIJ,
2417:                                        /*20*/ MatAssemblyBegin_MPIBAIJ,
2418:                                        MatAssemblyEnd_MPIBAIJ,
2419:                                        MatSetOption_MPIBAIJ,
2420:                                        MatZeroEntries_MPIBAIJ,
2421:                                        /*24*/ MatZeroRows_MPIBAIJ,
2422:                                        NULL,
2423:                                        NULL,
2424:                                        NULL,
2425:                                        NULL,
2426:                                        /*29*/ MatSetUp_MPI_Hash,
2427:                                        NULL,
2428:                                        NULL,
2429:                                        MatGetDiagonalBlock_MPIBAIJ,
2430:                                        NULL,
2431:                                        /*34*/ MatDuplicate_MPIBAIJ,
2432:                                        NULL,
2433:                                        NULL,
2434:                                        NULL,
2435:                                        NULL,
2436:                                        /*39*/ MatAXPY_MPIBAIJ,
2437:                                        MatCreateSubMatrices_MPIBAIJ,
2438:                                        MatIncreaseOverlap_MPIBAIJ,
2439:                                        MatGetValues_MPIBAIJ,
2440:                                        MatCopy_MPIBAIJ,
2441:                                        /*44*/ NULL,
2442:                                        MatScale_MPIBAIJ,
2443:                                        MatShift_MPIBAIJ,
2444:                                        NULL,
2445:                                        MatZeroRowsColumns_MPIBAIJ,
2446:                                        /*49*/ NULL,
2447:                                        NULL,
2448:                                        NULL,
2449:                                        NULL,
2450:                                        NULL,
2451:                                        /*54*/ MatFDColoringCreate_MPIXAIJ,
2452:                                        NULL,
2453:                                        MatSetUnfactored_MPIBAIJ,
2454:                                        MatPermute_MPIBAIJ,
2455:                                        MatSetValuesBlocked_MPIBAIJ,
2456:                                        /*59*/ MatCreateSubMatrix_MPIBAIJ,
2457:                                        MatDestroy_MPIBAIJ,
2458:                                        MatView_MPIBAIJ,
2459:                                        NULL,
2460:                                        NULL,
2461:                                        /*64*/ NULL,
2462:                                        NULL,
2463:                                        NULL,
2464:                                        NULL,
2465:                                        MatGetRowMaxAbs_MPIBAIJ,
2466:                                        /*69*/ NULL,
2467:                                        NULL,
2468:                                        NULL,
2469:                                        MatFDColoringApply_BAIJ,
2470:                                        NULL,
2471:                                        /*74*/ NULL,
2472:                                        NULL,
2473:                                        NULL,
2474:                                        NULL,
2475:                                        MatLoad_MPIBAIJ,
2476:                                        /*79*/ NULL,
2477:                                        NULL,
2478:                                        NULL,
2479:                                        NULL,
2480:                                        NULL,
2481:                                        /*84*/ NULL,
2482:                                        NULL,
2483:                                        NULL,
2484:                                        NULL,
2485:                                        NULL,
2486:                                        /*89*/ NULL,
2487:                                        NULL,
2488:                                        NULL,
2489:                                        NULL,
2490:                                        MatConjugate_MPIBAIJ,
2491:                                        /*94*/ NULL,
2492:                                        NULL,
2493:                                        MatRealPart_MPIBAIJ,
2494:                                        MatImaginaryPart_MPIBAIJ,
2495:                                        NULL,
2496:                                        /*99*/ NULL,
2497:                                        NULL,
2498:                                        NULL,
2499:                                        NULL,
2500:                                        NULL,
2501:                                        /*104*/ MatGetSeqNonzeroStructure_MPIBAIJ,
2502:                                        NULL,
2503:                                        MatGetGhosts_MPIBAIJ,
2504:                                        NULL,
2505:                                        NULL,
2506:                                        /*109*/ NULL,
2507:                                        NULL,
2508:                                        NULL,
2509:                                        NULL,
2510:                                        MatGetMultiProcBlock_MPIBAIJ,
2511:                                        /*114*/ NULL,
2512:                                        MatGetColumnReductions_MPIBAIJ,
2513:                                        MatInvertBlockDiagonal_MPIBAIJ,
2514:                                        NULL,
2515:                                        NULL,
2516:                                        /*119*/ NULL,
2517:                                        NULL,
2518:                                        NULL,
2519:                                        NULL,
2520:                                        NULL,
2521:                                        /*124*/ NULL,
2522:                                        MatSetBlockSizes_Default,
2523:                                        NULL,
2524:                                        MatFDColoringSetUp_MPIXAIJ,
2525:                                        NULL,
2526:                                        /*129*/ MatCreateMPIMatConcatenateSeqMat_MPIBAIJ,
2527:                                        NULL,
2528:                                        NULL,
2529:                                        NULL,
2530:                                        NULL,
2531:                                        /*134*/ NULL,
2532:                                        MatEliminateZeros_MPIBAIJ,
2533:                                        MatGetRowSumAbs_MPIBAIJ,
2534:                                        NULL,
2535:                                        NULL,
2536:                                        /*139*/ NULL,
2537:                                        MatCopyHashToXAIJ_MPI_Hash,
2538:                                        NULL,
2539:                                        NULL,
2540:                                        MatADot_Default,
2541:                                        /*144*/ MatANorm_Default,
2542:                                        NULL,
2543:                                        NULL,
2544:                                        NULL};

2546: PETSC_INTERN PetscErrorCode MatConvert_MPIBAIJ_MPISBAIJ(Mat, MatType, MatReuse, Mat *);
2547: PETSC_INTERN PetscErrorCode MatConvert_XAIJ_IS(Mat, MatType, MatReuse, Mat *);

2549: static PetscErrorCode MatMPIBAIJSetPreallocationCSR_MPIBAIJ(Mat B, PetscInt bs, const PetscInt ii[], const PetscInt jj[], const PetscScalar V[])
2550: {
2551:   PetscInt        m, rstart, cstart, cend;
2552:   PetscInt        i, j, dlen, olen, nz, nz_max = 0, *d_nnz = NULL, *o_nnz = NULL;
2553:   const PetscInt *JJ          = NULL;
2554:   PetscScalar    *values      = NULL;
2555:   PetscBool       roworiented = ((Mat_MPIBAIJ *)B->data)->roworiented;
2556:   PetscBool       nooffprocentries;

2558:   PetscFunctionBegin;
2559:   PetscCall(PetscLayoutSetBlockSize(B->rmap, bs));
2560:   PetscCall(PetscLayoutSetBlockSize(B->cmap, bs));
2561:   PetscCall(PetscLayoutSetUp(B->rmap));
2562:   PetscCall(PetscLayoutSetUp(B->cmap));
2563:   PetscCall(PetscLayoutGetBlockSize(B->rmap, &bs));
2564:   m      = B->rmap->n / bs;
2565:   rstart = B->rmap->rstart / bs;
2566:   cstart = B->cmap->rstart / bs;
2567:   cend   = B->cmap->rend / bs;

2569:   PetscCheck(!ii[0], PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "ii[0] must be 0 but it is %" PetscInt_FMT, ii[0]);
2570:   PetscCall(PetscMalloc2(m, &d_nnz, m, &o_nnz));
2571:   for (i = 0; i < m; i++) {
2572:     nz = ii[i + 1] - ii[i];
2573:     PetscCheck(nz >= 0, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Local row %" PetscInt_FMT " has a negative number of columns %" PetscInt_FMT, i, nz);
2574:     nz_max = PetscMax(nz_max, nz);
2575:     dlen   = 0;
2576:     olen   = 0;
2577:     JJ     = jj + ii[i];
2578:     for (j = 0; j < nz; j++) {
2579:       if (*JJ < cstart || *JJ >= cend) olen++;
2580:       else dlen++;
2581:       JJ++;
2582:     }
2583:     d_nnz[i] = dlen;
2584:     o_nnz[i] = olen;
2585:   }
2586:   PetscCall(MatMPIBAIJSetPreallocation(B, bs, 0, d_nnz, 0, o_nnz));
2587:   PetscCall(PetscFree2(d_nnz, o_nnz));

2589:   values = (PetscScalar *)V;
2590:   if (!values) PetscCall(PetscCalloc1(bs * bs * nz_max, &values));
2591:   for (i = 0; i < m; i++) {
2592:     PetscInt        row   = i + rstart;
2593:     PetscInt        ncols = ii[i + 1] - ii[i];
2594:     const PetscInt *icols = jj + ii[i];
2595:     if (bs == 1 || !roworiented) { /* block ordering matches the non-nested layout of MatSetValues so we can insert entire rows */
2596:       const PetscScalar *svals = values + (V ? (bs * bs * ii[i]) : 0);
2597:       PetscCall(MatSetValuesBlocked_MPIBAIJ(B, 1, &row, ncols, icols, svals, INSERT_VALUES));
2598:     } else { /* block ordering does not match so we can only insert one block at a time. */
2599:       PetscInt j;
2600:       for (j = 0; j < ncols; j++) {
2601:         const PetscScalar *svals = values + (V ? (bs * bs * (ii[i] + j)) : 0);
2602:         PetscCall(MatSetValuesBlocked_MPIBAIJ(B, 1, &row, 1, &icols[j], svals, INSERT_VALUES));
2603:       }
2604:     }
2605:   }

2607:   if (!V) PetscCall(PetscFree(values));
2608:   nooffprocentries    = B->nooffprocentries;
2609:   B->nooffprocentries = PETSC_TRUE;
2610:   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
2611:   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
2612:   B->nooffprocentries = nooffprocentries;

2614:   PetscCall(MatSetOption(B, MAT_NEW_NONZERO_LOCATION_ERR, PETSC_TRUE));
2615:   PetscFunctionReturn(PETSC_SUCCESS);
2616: }

2618: /*@C
2619:   MatMPIBAIJSetPreallocationCSR - Creates a sparse parallel matrix in `MATBAIJ` format using the given nonzero structure and (optional) numerical values

2621:   Collective

2623:   Input Parameters:
2624: + B  - the matrix
2625: . bs - the block size
2626: . i  - the indices into `j` for the start of each local row (starts with zero)
2627: . j  - the column indices for each local row (starts with zero) these must be sorted for each row
2628: - v  - optional values in the matrix, use `NULL` if not provided

2630:   Level: advanced

2632:   Notes:
2633:   The `i`, `j`, and `v` arrays ARE copied by this routine into the internal format used by PETSc;
2634:   thus you CANNOT change the matrix entries by changing the values of `v` after you have
2635:   called this routine.

2637:   The order of the entries in values is specified by the `MatOption` `MAT_ROW_ORIENTED`.  For example, C programs
2638:   may want to use the default `MAT_ROW_ORIENTED` with value `PETSC_TRUE` and use an array v[nnz][bs][bs] where the second index is
2639:   over rows within a block and the last index is over columns within a block row.  Fortran programs will likely set
2640:   `MAT_ROW_ORIENTED` with value `PETSC_FALSE` and use a Fortran array v(bs,bs,nnz) in which the first index is over rows within a
2641:   block column and the second index is over columns within a block.

2643:   Though this routine has Preallocation() in the name it also sets the exact nonzero locations of the matrix entries and usually the numerical values as well

2645: .seealso: `Mat`, `MatCreate()`, `MatCreateSeqAIJ()`, `MatSetValues()`, `MatMPIBAIJSetPreallocation()`, `MatCreateAIJ()`, `MATMPIAIJ`, `MatCreateMPIBAIJWithArrays()`, `MATMPIBAIJ`
2646: @*/
2647: PetscErrorCode MatMPIBAIJSetPreallocationCSR(Mat B, PetscInt bs, const PetscInt i[], const PetscInt j[], const PetscScalar v[])
2648: {
2649:   PetscFunctionBegin;
2653:   PetscTryMethod(B, "MatMPIBAIJSetPreallocationCSR_C", (Mat, PetscInt, const PetscInt[], const PetscInt[], const PetscScalar[]), (B, bs, i, j, v));
2654:   PetscFunctionReturn(PETSC_SUCCESS);
2655: }

2657: PetscErrorCode MatMPIBAIJSetPreallocation_MPIBAIJ(Mat B, PetscInt bs, PetscInt d_nz, const PetscInt *d_nnz, PetscInt o_nz, const PetscInt *o_nnz)
2658: {
2659:   Mat_MPIBAIJ *b = (Mat_MPIBAIJ *)B->data;
2660:   PetscInt     i;
2661:   PetscMPIInt  size;

2663:   PetscFunctionBegin;
2664:   if (B->hash_active) {
2665:     B->ops[0]      = b->cops;
2666:     B->hash_active = PETSC_FALSE;
2667:   }
2668:   if (!B->preallocated) PetscCall(MatStashCreate_Private(PetscObjectComm((PetscObject)B), bs, &B->bstash));
2669:   PetscCall(MatSetBlockSize(B, bs));
2670:   PetscCall(PetscLayoutSetUp(B->rmap));
2671:   PetscCall(PetscLayoutSetUp(B->cmap));
2672:   PetscCall(PetscLayoutGetBlockSize(B->rmap, &bs));

2674:   if (d_nnz) {
2675:     for (i = 0; i < B->rmap->n / bs; i++) PetscCheck(d_nnz[i] >= 0, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "d_nnz cannot be less than -1: local row %" PetscInt_FMT " value %" PetscInt_FMT, i, d_nnz[i]);
2676:   }
2677:   if (o_nnz) {
2678:     for (i = 0; i < B->rmap->n / bs; i++) PetscCheck(o_nnz[i] >= 0, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "o_nnz cannot be less than -1: local row %" PetscInt_FMT " value %" PetscInt_FMT, i, o_nnz[i]);
2679:   }

2681:   b->bs2 = bs * bs;
2682:   b->mbs = B->rmap->n / bs;
2683:   b->nbs = B->cmap->n / bs;
2684:   b->Mbs = B->rmap->N / bs;
2685:   b->Nbs = B->cmap->N / bs;

2687:   for (i = 0; i <= b->size; i++) b->rangebs[i] = B->rmap->range[i] / bs;
2688:   b->rstartbs = B->rmap->rstart / bs;
2689:   b->rendbs   = B->rmap->rend / bs;
2690:   b->cstartbs = B->cmap->rstart / bs;
2691:   b->cendbs   = B->cmap->rend / bs;

2693: #if defined(PETSC_USE_CTABLE)
2694:   PetscCall(PetscHMapIDestroy(&b->colmap));
2695: #else
2696:   PetscCall(PetscFree(b->colmap));
2697: #endif
2698:   PetscCall(PetscFree(b->garray));
2699:   PetscCall(VecDestroy(&b->lvec));
2700:   PetscCall(VecScatterDestroy(&b->Mvctx));

2702:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)B), &size));

2704:   MatSeqXAIJGetOptions_Private(b->B);
2705:   PetscCall(MatDestroy(&b->B));
2706:   PetscCall(MatCreate(PETSC_COMM_SELF, &b->B));
2707:   PetscCall(MatSetSizes(b->B, B->rmap->n, size > 1 ? B->cmap->N : 0, B->rmap->n, size > 1 ? B->cmap->N : 0));
2708:   PetscCall(MatSetType(b->B, MATSEQBAIJ));
2709:   MatSeqXAIJRestoreOptions_Private(b->B);

2711:   MatSeqXAIJGetOptions_Private(b->A);
2712:   PetscCall(MatDestroy(&b->A));
2713:   PetscCall(MatCreate(PETSC_COMM_SELF, &b->A));
2714:   PetscCall(MatSetSizes(b->A, B->rmap->n, B->cmap->n, B->rmap->n, B->cmap->n));
2715:   PetscCall(MatSetType(b->A, MATSEQBAIJ));
2716:   MatSeqXAIJRestoreOptions_Private(b->A);

2718:   PetscCall(MatSeqBAIJSetPreallocation(b->A, bs, d_nz, d_nnz));
2719:   PetscCall(MatSeqBAIJSetPreallocation(b->B, bs, o_nz, o_nnz));
2720:   B->preallocated  = PETSC_TRUE;
2721:   B->was_assembled = PETSC_FALSE;
2722:   B->assembled     = PETSC_FALSE;
2723:   PetscFunctionReturn(PETSC_SUCCESS);
2724: }

2726: extern PetscErrorCode MatDiagonalScaleLocal_MPIBAIJ(Mat, Vec);
2727: extern PetscErrorCode MatSetHashTableFactor_MPIBAIJ(Mat, PetscReal);

2729: PETSC_INTERN PetscErrorCode MatConvert_MPIBAIJ_MPIAdj(Mat B, MatType newtype, MatReuse reuse, Mat *adj)
2730: {
2731:   Mat_MPIBAIJ    *b = (Mat_MPIBAIJ *)B->data;
2732:   Mat_SeqBAIJ    *d = (Mat_SeqBAIJ *)b->A->data, *o = (Mat_SeqBAIJ *)b->B->data;
2733:   PetscInt        M = B->rmap->n / B->rmap->bs, i, *ii, *jj, cnt, j, k, rstart = B->rmap->rstart / B->rmap->bs;
2734:   const PetscInt *id = d->i, *jd = d->j, *io = o->i, *jo = o->j, *garray = b->garray;

2736:   PetscFunctionBegin;
2737:   PetscCall(PetscMalloc1(M + 1, &ii));
2738:   ii[0] = 0;
2739:   for (i = 0; i < M; i++) {
2740:     PetscCheck((id[i + 1] - id[i]) >= 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Indices wrong %" PetscInt_FMT " %" PetscInt_FMT " %" PetscInt_FMT, i, id[i], id[i + 1]);
2741:     PetscCheck((io[i + 1] - io[i]) >= 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Indices wrong %" PetscInt_FMT " %" PetscInt_FMT " %" PetscInt_FMT, i, io[i], io[i + 1]);
2742:     ii[i + 1] = ii[i] + id[i + 1] - id[i] + io[i + 1] - io[i];
2743:     /* remove one from count of matrix has diagonal */
2744:     for (j = id[i]; j < id[i + 1]; j++) {
2745:       if (jd[j] == i) {
2746:         ii[i + 1]--;
2747:         break;
2748:       }
2749:     }
2750:   }
2751:   PetscCall(PetscMalloc1(ii[M], &jj));
2752:   cnt = 0;
2753:   for (i = 0; i < M; i++) {
2754:     for (j = io[i]; j < io[i + 1]; j++) {
2755:       if (garray[jo[j]] > rstart) break;
2756:       jj[cnt++] = garray[jo[j]];
2757:     }
2758:     for (k = id[i]; k < id[i + 1]; k++) {
2759:       if (jd[k] != i) jj[cnt++] = rstart + jd[k];
2760:     }
2761:     for (; j < io[i + 1]; j++) jj[cnt++] = garray[jo[j]];
2762:   }
2763:   PetscCall(MatCreateMPIAdj(PetscObjectComm((PetscObject)B), M, B->cmap->N / B->rmap->bs, ii, jj, NULL, adj));
2764:   PetscFunctionReturn(PETSC_SUCCESS);
2765: }

2767: #include <../src/mat/impls/aij/mpi/mpiaij.h>

2769: PETSC_INTERN PetscErrorCode MatConvert_SeqBAIJ_SeqAIJ(Mat, MatType, MatReuse, Mat *);

2771: PETSC_INTERN PetscErrorCode MatConvert_MPIBAIJ_MPIAIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
2772: {
2773:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;
2774:   Mat_MPIAIJ  *b;
2775:   Mat          B;

2777:   PetscFunctionBegin;
2778:   PetscCheck(A->assembled, PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "Matrix must be assembled");

2780:   if (reuse == MAT_REUSE_MATRIX) {
2781:     B = *newmat;
2782:   } else {
2783:     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
2784:     PetscCall(MatSetType(B, MATMPIAIJ));
2785:     PetscCall(MatSetSizes(B, A->rmap->n, A->cmap->n, A->rmap->N, A->cmap->N));
2786:     PetscCall(MatSetBlockSizes(B, A->rmap->bs, A->cmap->bs));
2787:     PetscCall(MatSeqAIJSetPreallocation(B, 0, NULL));
2788:     PetscCall(MatMPIAIJSetPreallocation(B, 0, NULL, 0, NULL));
2789:   }
2790:   b = (Mat_MPIAIJ *)B->data;

2792:   if (reuse == MAT_REUSE_MATRIX) {
2793:     PetscCall(MatConvert_SeqBAIJ_SeqAIJ(a->A, MATSEQAIJ, MAT_REUSE_MATRIX, &b->A));
2794:     PetscCall(MatConvert_SeqBAIJ_SeqAIJ(a->B, MATSEQAIJ, MAT_REUSE_MATRIX, &b->B));
2795:   } else {
2796:     PetscInt   *garray = a->garray;
2797:     Mat_SeqAIJ *bB;
2798:     PetscInt    bs, nnz;
2799:     PetscCall(MatDestroy(&b->A));
2800:     PetscCall(MatDestroy(&b->B));
2801:     /* just clear out the data structure */
2802:     PetscCall(MatDisAssemble_MPIAIJ(B, PETSC_FALSE));
2803:     PetscCall(MatConvert_SeqBAIJ_SeqAIJ(a->A, MATSEQAIJ, MAT_INITIAL_MATRIX, &b->A));
2804:     PetscCall(MatConvert_SeqBAIJ_SeqAIJ(a->B, MATSEQAIJ, MAT_INITIAL_MATRIX, &b->B));

2806:     /* Global numbering for b->B columns */
2807:     bB  = (Mat_SeqAIJ *)b->B->data;
2808:     bs  = A->rmap->bs;
2809:     nnz = bB->i[A->rmap->n];
2810:     for (PetscInt k = 0; k < nnz; k++) {
2811:       PetscInt bj = bB->j[k] / bs;
2812:       PetscInt br = bB->j[k] % bs;
2813:       bB->j[k]    = garray[bj] * bs + br;
2814:     }
2815:   }
2816:   PetscCall(MatSetOption(B, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
2817:   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
2818:   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
2819:   PetscCall(MatSetOption(B, MAT_NO_OFF_PROC_ENTRIES, PETSC_FALSE));

2821:   if (reuse == MAT_INPLACE_MATRIX) {
2822:     PetscCall(MatHeaderReplace(A, &B));
2823:   } else {
2824:     *newmat = B;
2825:   }
2826:   PetscFunctionReturn(PETSC_SUCCESS);
2827: }

2829: /*MC
2830:    MATMPIBAIJ - MATMPIBAIJ = "mpibaij" - A matrix type to be used for distributed block sparse matrices.

2832:    Options Database Keys:
2833: + -mat_type mpibaij              - sets the matrix type to `MATMPIBAIJ` during a call to `MatSetFromOptions()`
2834: . -mat_block_size bs             - set the blocksize used to store the matrix
2835: . -mat_baij_mult_version version - indicate the version of the matrix-vector product to use  (0 often indicates using BLAS)
2836: - -mat_use_hash_table fact       - set hash table factor

2838:    Level: beginner

2840:    Note:
2841:     `MatSetOption(A, MAT_STRUCTURE_ONLY, PETSC_TRUE)` may be called for this matrix type. In this no
2842:     space is allocated for the nonzero entries and any entries passed with `MatSetValues()` are ignored

2844: .seealso: `Mat`, `MATBAIJ`, `MATSEQBAIJ`, `MatCreateBAIJ`
2845: M*/

2847: PETSC_INTERN PetscErrorCode MatConvert_MPIBAIJ_MPIBSTRM(Mat, MatType, MatReuse, Mat *);

2849: static PetscErrorCode MatGetMultPetscSF_MPIBAIJ(Mat A, PetscSF *sf)
2850: {
2851:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

2853:   PetscFunctionBegin;
2854:   *sf = a->Mvctx;
2855:   PetscFunctionReturn(PETSC_SUCCESS);
2856: }

2858: PETSC_EXTERN PetscErrorCode MatCreate_MPIBAIJ(Mat B)
2859: {
2860:   Mat_MPIBAIJ *b;
2861:   PetscBool    flg = PETSC_FALSE;

2863:   PetscFunctionBegin;
2864:   PetscCall(PetscNew(&b));
2865:   B->data      = (void *)b;
2866:   B->ops[0]    = MatOps_Values;
2867:   B->assembled = PETSC_FALSE;

2869:   B->insertmode = NOT_SET_VALUES;
2870:   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)B), &b->rank));
2871:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)B), &b->size));

2873:   /* build local table of row and column ownerships */
2874:   PetscCall(PetscMalloc1(b->size + 1, &b->rangebs));

2876:   /* build cache for off array entries formed */
2877:   PetscCall(MatStashCreate_Private(PetscObjectComm((PetscObject)B), 1, &B->stash));

2879:   b->donotstash  = PETSC_FALSE;
2880:   b->colmap      = NULL;
2881:   b->garray      = NULL;
2882:   b->roworiented = PETSC_TRUE;

2884:   /* stuff used in block assembly */
2885:   b->barray = NULL;

2887:   /* stuff used for matrix vector multiply */
2888:   b->lvec  = NULL;
2889:   b->Mvctx = NULL;

2891:   /* stuff for MatGetRow() */
2892:   b->rowindices   = NULL;
2893:   b->rowvalues    = NULL;
2894:   b->getrowactive = PETSC_FALSE;

2896:   /* hash table stuff */
2897:   b->ht           = NULL;
2898:   b->hd           = NULL;
2899:   b->ht_size      = 0;
2900:   b->ht_flag      = PETSC_FALSE;
2901:   b->ht_fact      = 0;
2902:   b->ht_total_ct  = 0;
2903:   b->ht_insert_ct = 0;

2905:   /* stuff for MatCreateSubMatrices_MPIBAIJ_local() */
2906:   b->ijonly = PETSC_FALSE;

2908:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_mpibaij_mpiadj_C", MatConvert_MPIBAIJ_MPIAdj));
2909:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_mpibaij_mpiaij_C", MatConvert_MPIBAIJ_MPIAIJ));
2910:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_mpibaij_mpisbaij_C", MatConvert_MPIBAIJ_MPISBAIJ));
2911: #if defined(PETSC_HAVE_HYPRE)
2912:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_mpibaij_hypre_C", MatConvert_AIJ_HYPRE));
2913: #endif
2914:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatStoreValues_C", MatStoreValues_MPIBAIJ));
2915:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatRetrieveValues_C", MatRetrieveValues_MPIBAIJ));
2916:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatMPIBAIJSetPreallocation_C", MatMPIBAIJSetPreallocation_MPIBAIJ));
2917:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatMPIBAIJSetPreallocationCSR_C", MatMPIBAIJSetPreallocationCSR_MPIBAIJ));
2918:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatDiagonalScaleLocal_C", MatDiagonalScaleLocal_MPIBAIJ));
2919:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatSetHashTableFactor_C", MatSetHashTableFactor_MPIBAIJ));
2920:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_mpibaij_is_C", MatConvert_XAIJ_IS));
2921:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatGetMultPetscSF_C", MatGetMultPetscSF_MPIBAIJ));
2922:   PetscCall(PetscObjectChangeTypeName((PetscObject)B, MATMPIBAIJ));

2924:   PetscOptionsBegin(PetscObjectComm((PetscObject)B), NULL, "Options for loading MPIBAIJ matrix 1", "Mat");
2925:   PetscCall(PetscOptionsName("-mat_use_hash_table", "Use hash table to save time in constructing matrix", "MatSetOption", &flg));
2926:   if (flg) {
2927:     PetscReal fact = 1.39;
2928:     PetscCall(MatSetOption(B, MAT_USE_HASH_TABLE, PETSC_TRUE));
2929:     PetscCall(PetscOptionsReal("-mat_use_hash_table", "Use hash table factor", "MatMPIBAIJSetHashTableFactor", fact, &fact, NULL));
2930:     if (fact <= 1.0) fact = 1.39;
2931:     PetscCall(MatMPIBAIJSetHashTableFactor(B, fact));
2932:     PetscCall(PetscInfo(B, "Hash table Factor used %5.2g\n", (double)fact));
2933:   }
2934:   PetscOptionsEnd();
2935:   PetscFunctionReturn(PETSC_SUCCESS);
2936: }

2938: // PetscClangLinter pragma disable: -fdoc-section-header-unknown
2939: /*MC
2940:    MATBAIJ - MATBAIJ = "baij" - A matrix type to be used for block sparse matrices.

2942:    This matrix type is identical to `MATSEQBAIJ` when constructed with a single process communicator,
2943:    and `MATMPIBAIJ` otherwise.

2945:    Options Database Keys:
2946: . -mat_type baij - sets the matrix type to `MATBAIJ` during a call to `MatSetFromOptions()`

2948:   Level: beginner

2950: .seealso: `Mat`, `MatCreateBAIJ()`, `MATSEQBAIJ`, `MATMPIBAIJ`, `MatMPIBAIJSetPreallocation()`, `MatMPIBAIJSetPreallocationCSR()`
2951: M*/

2953: /*@
2954:   MatMPIBAIJSetPreallocation - Allocates memory for a sparse parallel matrix in `MATMPIBAIJ` format
2955:   (block compressed row).

2957:   Collective

2959:   Input Parameters:
2960: + B     - the matrix
2961: . bs    - size of block, the blocks are ALWAYS square. One can use `MatSetBlockSizes()` to set a different row and column blocksize but the row
2962:           blocksize always defines the size of the blocks. The column blocksize sets the blocksize of the vectors obtained with `MatCreateVecs()`
2963: . d_nz  - number of block nonzeros per block row in diagonal portion of local
2964:            submatrix  (same for all local rows)
2965: . d_nnz - array containing the number of block nonzeros in the various block rows
2966:            of the in diagonal portion of the local (possibly different for each block
2967:            row) or `NULL`.  If you plan to factor the matrix you must leave room for the diagonal entry and
2968:            set it even if it is zero.
2969: . o_nz  - number of block nonzeros per block row in the off-diagonal portion of local
2970:            submatrix (same for all local rows).
2971: - o_nnz - array containing the number of nonzeros in the various block rows of the
2972:            off-diagonal portion of the local submatrix (possibly different for
2973:            each block row) or `NULL`.

2975:    If the *_nnz parameter is given then the *_nz parameter is ignored

2977:   Options Database Keys:
2978: + -mat_block_size          - size of the blocks to use
2979: - -mat_use_hash_table fact - set hash table factor

2981:   Level: intermediate

2983:   Notes:
2984:   For good matrix assembly performance
2985:   the user should preallocate the matrix storage by setting the parameters
2986:   `d_nz` (or `d_nnz`) and `o_nz` (or `o_nnz`).  By setting these parameters accurately,
2987:   performance can be increased by more than a factor of 50.

2989:   If `PETSC_DECIDE` or  `PETSC_DETERMINE` is used for a particular argument on one processor
2990:   than it must be used on all processors that share the object for that argument.

2992:   Storage Information:
2993:   For a square global matrix we define each processor's diagonal portion
2994:   to be its local rows and the corresponding columns (a square submatrix);
2995:   each processor's off-diagonal portion encompasses the remainder of the
2996:   local matrix (a rectangular submatrix).

2998:   The user can specify preallocated storage for the diagonal part of
2999:   the local submatrix with either `d_nz` or `d_nnz` (not both).  Set
3000:   `d_nz` = `PETSC_DEFAULT` and `d_nnz` = `NULL` for PETSc to control dynamic
3001:   memory allocation.  Likewise, specify preallocated storage for the
3002:   off-diagonal part of the local submatrix with `o_nz` or `o_nnz` (not both).

3004:   Consider a processor that owns rows 3, 4 and 5 of a parallel matrix. In
3005:   the figure below we depict these three local rows and all columns (0-11).

3007: .vb
3008:            0 1 2 3 4 5 6 7 8 9 10 11
3009:           --------------------------
3010:    row 3  |o o o d d d o o o o  o  o
3011:    row 4  |o o o d d d o o o o  o  o
3012:    row 5  |o o o d d d o o o o  o  o
3013:           --------------------------
3014: .ve

3016:   Thus, any entries in the d locations are stored in the d (diagonal)
3017:   submatrix, and any entries in the o locations are stored in the
3018:   o (off-diagonal) submatrix.  Note that the d and the o submatrices are
3019:   stored simply in the `MATSEQBAIJ` format for compressed row storage.

3021:   Now `d_nz` should indicate the number of block nonzeros per row in the d matrix,
3022:   and `o_nz` should indicate the number of block nonzeros per row in the o matrix.
3023:   In general, for PDE problems in which most nonzeros are near the diagonal,
3024:   one expects `d_nz` >> `o_nz`.

3026:   You can call `MatGetInfo()` to get information on how effective the preallocation was;
3027:   for example the fields mallocs,nz_allocated,nz_used,nz_unneeded;
3028:   You can also run with the option `-info` and look for messages with the string
3029:   malloc in them to see if additional memory allocation was needed.

3031: .seealso: `Mat`, `MATMPIBAIJ`, `MatCreate()`, `MatCreateSeqBAIJ()`, `MatSetValues()`, `MatCreateBAIJ()`, `MatMPIBAIJSetPreallocationCSR()`, `PetscSplitOwnership()`
3032: @*/
3033: PetscErrorCode MatMPIBAIJSetPreallocation(Mat B, PetscInt bs, PetscInt d_nz, const PetscInt d_nnz[], PetscInt o_nz, const PetscInt o_nnz[])
3034: {
3035:   PetscFunctionBegin;
3039:   PetscTryMethod(B, "MatMPIBAIJSetPreallocation_C", (Mat, PetscInt, PetscInt, const PetscInt[], PetscInt, const PetscInt[]), (B, bs, d_nz, d_nnz, o_nz, o_nnz));
3040:   PetscFunctionReturn(PETSC_SUCCESS);
3041: }

3043: // PetscClangLinter pragma disable: -fdoc-section-header-unknown
3044: /*@
3045:   MatCreateBAIJ - Creates a sparse parallel matrix in `MATBAIJ` format
3046:   (block compressed row).

3048:   Collective

3050:   Input Parameters:
3051: + comm  - MPI communicator
3052: . bs    - size of block, the blocks are ALWAYS square. One can use `MatSetBlockSizes()` to set a different row and column blocksize but the row
3053:           blocksize always defines the size of the blocks. The column blocksize sets the blocksize of the vectors obtained with `MatCreateVecs()`
3054: . m     - number of local rows (or `PETSC_DECIDE` to have calculated if M is given)
3055:           This value should be the same as the local size used in creating the
3056:           y vector for the matrix-vector product y = Ax.
3057: . n     - number of local columns (or `PETSC_DECIDE` to have calculated if N is given)
3058:           This value should be the same as the local size used in creating the
3059:           x vector for the matrix-vector product y = Ax.
3060: . M     - number of global rows (or `PETSC_DETERMINE` to have calculated if m is given)
3061: . N     - number of global columns (or `PETSC_DETERMINE` to have calculated if n is given)
3062: . d_nz  - number of nonzero blocks per block row in diagonal portion of local
3063:           submatrix  (same for all local rows)
3064: . d_nnz - array containing the number of nonzero blocks in the various block rows
3065:           of the in diagonal portion of the local (possibly different for each block
3066:           row) or NULL.  If you plan to factor the matrix you must leave room for the diagonal entry
3067:           and set it even if it is zero.
3068: . o_nz  - number of nonzero blocks per block row in the off-diagonal portion of local
3069:           submatrix (same for all local rows).
3070: - o_nnz - array containing the number of nonzero blocks in the various block rows of the
3071:           off-diagonal portion of the local submatrix (possibly different for
3072:           each block row) or NULL.

3074:   Output Parameter:
3075: . A - the matrix

3077:   Options Database Keys:
3078: + -mat_block_size          - size of the blocks to use
3079: - -mat_use_hash_table fact - set hash table factor

3081:   Level: intermediate

3083:   Notes:
3084:   It is recommended that one use `MatCreateFromOptions()` or the `MatCreate()`, `MatSetType()` and/or `MatSetFromOptions()`,
3085:   MatXXXXSetPreallocation() paradigm instead of this routine directly.
3086:   [MatXXXXSetPreallocation() is, for example, `MatSeqBAIJSetPreallocation()`]

3088:   For good matrix assembly performance
3089:   the user should preallocate the matrix storage by setting the parameters
3090:   `d_nz` (or `d_nnz`) and `o_nz` (or `o_nnz`).  By setting these parameters accurately,
3091:   performance can be increased by more than a factor of 50.

3093:   If the *_nnz parameter is given then the *_nz parameter is ignored

3095:   A nonzero block is any block that as 1 or more nonzeros in it

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

3100:   If `PETSC_DECIDE` or  `PETSC_DETERMINE` is used for a particular argument on one processor
3101:   than it must be used on all processors that share the object for that argument.

3103:   If `m` and `n` are not `PETSC_DECIDE`, then the values determine the `PetscLayout` of the matrix and the ranges returned by
3104:   `MatGetOwnershipRange()`,  `MatGetOwnershipRanges()`, `MatGetOwnershipRangeColumn()`, and `MatGetOwnershipRangesColumn()`.

3106:   Storage Information:
3107:   For a square global matrix we define each processor's diagonal portion
3108:   to be its local rows and the corresponding columns (a square submatrix);
3109:   each processor's off-diagonal portion encompasses the remainder of the
3110:   local matrix (a rectangular submatrix).

3112:   The user can specify preallocated storage for the diagonal part of
3113:   the local submatrix with either d_nz or d_nnz (not both).  Set
3114:   `d_nz` = `PETSC_DEFAULT` and `d_nnz` = `NULL` for PETSc to control dynamic
3115:   memory allocation.  Likewise, specify preallocated storage for the
3116:   off-diagonal part of the local submatrix with `o_nz` or `o_nnz` (not both).

3118:   Consider a processor that owns rows 3, 4 and 5 of a parallel matrix. In
3119:   the figure below we depict these three local rows and all columns (0-11).

3121: .vb
3122:            0 1 2 3 4 5 6 7 8 9 10 11
3123:           --------------------------
3124:    row 3  |o o o d d d o o o o  o  o
3125:    row 4  |o o o d d d o o o o  o  o
3126:    row 5  |o o o d d d o o o o  o  o
3127:           --------------------------
3128: .ve

3130:   Thus, any entries in the d locations are stored in the d (diagonal)
3131:   submatrix, and any entries in the o locations are stored in the
3132:   o (off-diagonal) submatrix.  Note that the d and the o submatrices are
3133:   stored simply in the `MATSEQBAIJ` format for compressed row storage.

3135:   Now `d_nz` should indicate the number of block nonzeros per row in the d matrix,
3136:   and `o_nz` should indicate the number of block nonzeros per row in the o matrix.
3137:   In general, for PDE problems in which most nonzeros are near the diagonal,
3138:   one expects `d_nz` >> `o_nz`.

3140: .seealso: `Mat`, `MatCreate()`, `MatCreateSeqBAIJ()`, `MatSetValues()`, `MatMPIBAIJSetPreallocation()`, `MatMPIBAIJSetPreallocationCSR()`,
3141:           `MatGetOwnershipRange()`, `MatGetOwnershipRanges()`, `MatGetOwnershipRangeColumn()`, `MatGetOwnershipRangesColumn()`, `PetscLayout`
3142: @*/
3143: PetscErrorCode MatCreateBAIJ(MPI_Comm comm, PetscInt bs, PetscInt m, PetscInt n, PetscInt M, PetscInt N, PetscInt d_nz, const PetscInt d_nnz[], PetscInt o_nz, const PetscInt o_nnz[], Mat *A)
3144: {
3145:   PetscMPIInt size;

3147:   PetscFunctionBegin;
3148:   PetscCall(MatCreate(comm, A));
3149:   PetscCall(MatSetSizes(*A, m, n, M, N));
3150:   PetscCallMPI(MPI_Comm_size(comm, &size));
3151:   if (size > 1) {
3152:     PetscCall(MatSetType(*A, MATMPIBAIJ));
3153:     PetscCall(MatMPIBAIJSetPreallocation(*A, bs, d_nz, d_nnz, o_nz, o_nnz));
3154:   } else {
3155:     PetscCall(MatSetType(*A, MATSEQBAIJ));
3156:     PetscCall(MatSeqBAIJSetPreallocation(*A, bs, d_nz, d_nnz));
3157:   }
3158:   PetscFunctionReturn(PETSC_SUCCESS);
3159: }

3161: static PetscErrorCode MatDuplicate_MPIBAIJ(Mat matin, MatDuplicateOption cpvalues, Mat *newmat)
3162: {
3163:   Mat          mat;
3164:   Mat_MPIBAIJ *a, *oldmat = (Mat_MPIBAIJ *)matin->data;
3165:   PetscInt     len = 0;

3167:   PetscFunctionBegin;
3168:   *newmat = NULL;
3169:   PetscCall(MatCreate(PetscObjectComm((PetscObject)matin), &mat));
3170:   PetscCall(MatSetSizes(mat, matin->rmap->n, matin->cmap->n, matin->rmap->N, matin->cmap->N));
3171:   PetscCall(MatSetType(mat, ((PetscObject)matin)->type_name));

3173:   PetscCall(PetscLayoutReference(matin->rmap, &mat->rmap));
3174:   PetscCall(PetscLayoutReference(matin->cmap, &mat->cmap));
3175:   if (matin->hash_active) PetscCall(MatSetUp(mat));
3176:   else {
3177:     mat->factortype   = matin->factortype;
3178:     mat->preallocated = PETSC_TRUE;
3179:     mat->assembled    = PETSC_TRUE;
3180:     mat->insertmode   = NOT_SET_VALUES;

3182:     a             = (Mat_MPIBAIJ *)mat->data;
3183:     mat->rmap->bs = matin->rmap->bs;
3184:     a->bs2        = oldmat->bs2;
3185:     a->mbs        = oldmat->mbs;
3186:     a->nbs        = oldmat->nbs;
3187:     a->Mbs        = oldmat->Mbs;
3188:     a->Nbs        = oldmat->Nbs;

3190:     a->size         = oldmat->size;
3191:     a->rank         = oldmat->rank;
3192:     a->donotstash   = oldmat->donotstash;
3193:     a->roworiented  = oldmat->roworiented;
3194:     a->rowindices   = NULL;
3195:     a->rowvalues    = NULL;
3196:     a->getrowactive = PETSC_FALSE;
3197:     a->barray       = NULL;
3198:     a->rstartbs     = oldmat->rstartbs;
3199:     a->rendbs       = oldmat->rendbs;
3200:     a->cstartbs     = oldmat->cstartbs;
3201:     a->cendbs       = oldmat->cendbs;

3203:     /* hash table stuff */
3204:     a->ht           = NULL;
3205:     a->hd           = NULL;
3206:     a->ht_size      = 0;
3207:     a->ht_flag      = oldmat->ht_flag;
3208:     a->ht_fact      = oldmat->ht_fact;
3209:     a->ht_total_ct  = 0;
3210:     a->ht_insert_ct = 0;

3212:     PetscCall(PetscArraycpy(a->rangebs, oldmat->rangebs, a->size + 1));
3213:     if (oldmat->colmap) {
3214: #if defined(PETSC_USE_CTABLE)
3215:       PetscCall(PetscHMapIDuplicate(oldmat->colmap, &a->colmap));
3216: #else
3217:       PetscCall(PetscMalloc1(a->Nbs, &a->colmap));
3218:       PetscCall(PetscArraycpy(a->colmap, oldmat->colmap, a->Nbs));
3219: #endif
3220:     } else a->colmap = NULL;

3222:     if (oldmat->garray && (len = ((Mat_SeqBAIJ *)oldmat->B->data)->nbs)) {
3223:       PetscCall(PetscMalloc1(len, &a->garray));
3224:       PetscCall(PetscArraycpy(a->garray, oldmat->garray, len));
3225:     } else a->garray = NULL;

3227:     PetscCall(MatStashCreate_Private(PetscObjectComm((PetscObject)matin), matin->rmap->bs, &mat->bstash));
3228:     PetscCall(VecDuplicate(oldmat->lvec, &a->lvec));
3229:     PetscCall(VecScatterCopy(oldmat->Mvctx, &a->Mvctx));

3231:     PetscCall(MatDuplicate(oldmat->A, cpvalues, &a->A));
3232:     PetscCall(MatDuplicate(oldmat->B, cpvalues, &a->B));
3233:   }
3234:   PetscCall(PetscFunctionListDuplicate(((PetscObject)matin)->qlist, &((PetscObject)mat)->qlist));
3235:   *newmat = mat;
3236:   PetscFunctionReturn(PETSC_SUCCESS);
3237: }

3239: /* Used for both MPIBAIJ and MPISBAIJ matrices */
3240: PetscErrorCode MatLoad_MPIBAIJ_Binary(Mat mat, PetscViewer viewer)
3241: {
3242:   PetscInt     header[4], M, N, nz, bs, m, n, mbs, nbs, rows, cols, sum, i, j, k;
3243:   PetscInt    *rowidxs, *colidxs, rs, cs, ce;
3244:   PetscScalar *matvals;

3246:   PetscFunctionBegin;
3247:   PetscCall(PetscViewerSetUp(viewer));

3249:   /* read in matrix header */
3250:   PetscCall(PetscViewerBinaryRead(viewer, header, 4, NULL, PETSC_INT));
3251:   PetscCheck(header[0] == MAT_FILE_CLASSID, PetscObjectComm((PetscObject)viewer), PETSC_ERR_FILE_UNEXPECTED, "Not a matrix object in file");
3252:   M  = header[1];
3253:   N  = header[2];
3254:   nz = header[3];
3255:   PetscCheck(M >= 0, PetscObjectComm((PetscObject)viewer), PETSC_ERR_FILE_UNEXPECTED, "Matrix row size (%" PetscInt_FMT ") in file is negative", M);
3256:   PetscCheck(N >= 0, PetscObjectComm((PetscObject)viewer), PETSC_ERR_FILE_UNEXPECTED, "Matrix column size (%" PetscInt_FMT ") in file is negative", N);
3257:   PetscCheck(nz >= 0, PETSC_COMM_SELF, PETSC_ERR_FILE_UNEXPECTED, "Matrix stored in special format on disk, cannot load as MPIBAIJ");

3259:   /* set block sizes from the viewer's .info file */
3260:   PetscCall(MatLoad_Binary_BlockSizes(mat, viewer));
3261:   /* set local sizes if not set already */
3262:   if (mat->rmap->n < 0 && M == N) mat->rmap->n = mat->cmap->n;
3263:   if (mat->cmap->n < 0 && M == N) mat->cmap->n = mat->rmap->n;
3264:   /* set global sizes if not set already */
3265:   if (mat->rmap->N < 0) mat->rmap->N = M;
3266:   if (mat->cmap->N < 0) mat->cmap->N = N;
3267:   PetscCall(PetscLayoutSetUp(mat->rmap));
3268:   PetscCall(PetscLayoutSetUp(mat->cmap));

3270:   /* check if the matrix sizes are correct */
3271:   PetscCall(MatGetSize(mat, &rows, &cols));
3272:   PetscCheck(M == rows && N == cols, PETSC_COMM_SELF, PETSC_ERR_FILE_UNEXPECTED, "Matrix in file of different sizes (%" PetscInt_FMT ", %" PetscInt_FMT ") than the input matrix (%" PetscInt_FMT ", %" PetscInt_FMT ")", M, N, rows, cols);
3273:   PetscCall(MatGetBlockSize(mat, &bs));
3274:   PetscCall(MatGetLocalSize(mat, &m, &n));
3275:   PetscCall(PetscLayoutGetRange(mat->rmap, &rs, NULL));
3276:   PetscCall(PetscLayoutGetRange(mat->cmap, &cs, &ce));
3277:   mbs = m / bs;
3278:   nbs = n / bs;

3280:   /* read in row lengths and build row indices */
3281:   PetscCall(PetscMalloc1(m + 1, &rowidxs));
3282:   PetscCall(PetscViewerBinaryReadAll(viewer, rowidxs + 1, m, PETSC_DECIDE, M, PETSC_INT));
3283:   rowidxs[0] = 0;
3284:   for (i = 0; i < m; i++) rowidxs[i + 1] += rowidxs[i];
3285:   PetscCallMPI(MPIU_Allreduce(&rowidxs[m], &sum, 1, MPIU_INT, MPI_SUM, PetscObjectComm((PetscObject)viewer)));
3286:   PetscCheck(sum == nz, PetscObjectComm((PetscObject)viewer), PETSC_ERR_FILE_UNEXPECTED, "Inconsistent matrix data in file: nonzeros = %" PetscInt_FMT ", sum-row-lengths = %" PetscInt_FMT, nz, sum);

3288:   /* read in column indices and matrix values */
3289:   PetscCall(PetscMalloc2(rowidxs[m], &colidxs, rowidxs[m], &matvals));
3290:   PetscCall(PetscViewerBinaryReadAll(viewer, colidxs, rowidxs[m], PETSC_DETERMINE, PETSC_DETERMINE, PETSC_INT));
3291:   PetscCall(PetscViewerBinaryReadAll(viewer, matvals, rowidxs[m], PETSC_DETERMINE, PETSC_DETERMINE, PETSC_SCALAR));

3293:   {                /* preallocate matrix storage */
3294:     PetscBT    bt; /* helper bit set to count diagonal nonzeros */
3295:     PetscHSetI ht; /* helper hash set to count off-diagonal nonzeros */
3296:     PetscBool  sbaij, done;
3297:     PetscInt  *d_nnz, *o_nnz;

3299:     PetscCall(PetscBTCreate(nbs, &bt));
3300:     PetscCall(PetscHSetICreate(&ht));
3301:     PetscCall(PetscCalloc2(mbs, &d_nnz, mbs, &o_nnz));
3302:     PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATMPISBAIJ, &sbaij));
3303:     for (i = 0; i < mbs; i++) {
3304:       PetscCall(PetscBTMemzero(nbs, bt));
3305:       PetscCall(PetscHSetIClear(ht));
3306:       for (k = 0; k < bs; k++) {
3307:         PetscInt row = bs * i + k;
3308:         for (j = rowidxs[row]; j < rowidxs[row + 1]; j++) {
3309:           PetscInt col = colidxs[j];
3310:           if (!sbaij || col >= row) {
3311:             if (col >= cs && col < ce) {
3312:               if (!PetscBTLookupSet(bt, (col - cs) / bs)) d_nnz[i]++;
3313:             } else {
3314:               PetscCall(PetscHSetIQueryAdd(ht, col / bs, &done));
3315:               if (done) o_nnz[i]++;
3316:             }
3317:           }
3318:         }
3319:       }
3320:     }
3321:     PetscCall(PetscBTDestroy(&bt));
3322:     PetscCall(PetscHSetIDestroy(&ht));
3323:     PetscCall(MatMPIBAIJSetPreallocation(mat, bs, 0, d_nnz, 0, o_nnz));
3324:     PetscCall(MatMPISBAIJSetPreallocation(mat, bs, 0, d_nnz, 0, o_nnz));
3325:     PetscCall(PetscFree2(d_nnz, o_nnz));
3326:   }

3328:   /* store matrix values */
3329:   for (i = 0; i < m; i++) {
3330:     PetscInt row = rs + i, s = rowidxs[i], e = rowidxs[i + 1];
3331:     PetscUseTypeMethod(mat, setvalues, 1, &row, e - s, colidxs + s, matvals + s, INSERT_VALUES);
3332:   }

3334:   PetscCall(PetscFree(rowidxs));
3335:   PetscCall(PetscFree2(colidxs, matvals));
3336:   PetscCall(MatAssemblyBegin(mat, MAT_FINAL_ASSEMBLY));
3337:   PetscCall(MatAssemblyEnd(mat, MAT_FINAL_ASSEMBLY));
3338:   PetscFunctionReturn(PETSC_SUCCESS);
3339: }

3341: PetscErrorCode MatLoad_MPIBAIJ(Mat mat, PetscViewer viewer)
3342: {
3343:   PetscBool isbinary;

3345:   PetscFunctionBegin;
3346:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &isbinary));
3347:   PetscCheck(isbinary, PetscObjectComm((PetscObject)viewer), PETSC_ERR_SUP, "Viewer type %s not yet supported for reading %s matrices", ((PetscObject)viewer)->type_name, ((PetscObject)mat)->type_name);
3348:   PetscCall(MatLoad_MPIBAIJ_Binary(mat, viewer));
3349:   PetscFunctionReturn(PETSC_SUCCESS);
3350: }

3352: /*@
3353:   MatMPIBAIJSetHashTableFactor - Sets the factor required to compute the size of the matrices hash table

3355:   Input Parameters:
3356: + mat  - the matrix
3357: - fact - factor

3359:   Options Database Key:
3360: . -mat_use_hash_table fact - provide the factor

3362:   Level: advanced

3364: .seealso: `Mat`, `MATMPIBAIJ`, `MatSetOption()`
3365: @*/
3366: PetscErrorCode MatMPIBAIJSetHashTableFactor(Mat mat, PetscReal fact)
3367: {
3368:   PetscFunctionBegin;
3369:   PetscTryMethod(mat, "MatSetHashTableFactor_C", (Mat, PetscReal), (mat, fact));
3370:   PetscFunctionReturn(PETSC_SUCCESS);
3371: }

3373: PetscErrorCode MatSetHashTableFactor_MPIBAIJ(Mat mat, PetscReal fact)
3374: {
3375:   Mat_MPIBAIJ *baij;

3377:   PetscFunctionBegin;
3378:   baij          = (Mat_MPIBAIJ *)mat->data;
3379:   baij->ht_fact = fact;
3380:   PetscFunctionReturn(PETSC_SUCCESS);
3381: }

3383: /*@C
3384:   MatMPIBAIJGetSeqBAIJ - Get the on-process (diagonal block) and off-process (off-diagonal block) `MATSEQBAIJ`
3385:   matrices that make up an `MATMPIBAIJ` matrix, together with the local-to-global column map for the off-diagonal block.

3387:   Not Collective

3389:   Input Parameter:
3390: . A - the `MATMPIBAIJ` matrix

3392:   Output Parameters:
3393: + Ad     - the diagonal block `MATSEQBAIJ`, or `NULL` if not needed
3394: . Ao     - the off-diagonal block `MATSEQBAIJ`, or `NULL` if not needed
3395: - colmap - the local-to-global column index map for `Ao`, or `NULL` if not needed

3397:   Level: advanced

3399: .seealso: `Mat`, `MATMPIBAIJ`, `MATSEQBAIJ`, `MatMPIAIJGetSeqAIJ()`
3400: @*/
3401: PetscErrorCode MatMPIBAIJGetSeqBAIJ(Mat A, Mat *Ad, Mat *Ao, const PetscInt *colmap[])
3402: {
3403:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;
3404:   PetscBool    flg;

3406:   PetscFunctionBegin;
3407:   PetscCall(PetscObjectTypeCompare((PetscObject)A, MATMPIBAIJ, &flg));
3408:   PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "This function requires a MATMPIBAIJ matrix as input");
3409:   if (Ad) *Ad = a->A;
3410:   if (Ao) *Ao = a->B;
3411:   if (colmap) *colmap = a->garray;
3412:   PetscFunctionReturn(PETSC_SUCCESS);
3413: }

3415: /*
3416:     Special version for direct calls from Fortran (to eliminate two function call overheads
3417: */
3418: #if defined(PETSC_HAVE_FORTRAN_CAPS)
3419:   #define matmpibaijsetvaluesblocked_ MATMPIBAIJSETVALUESBLOCKED
3420: #elif !defined(PETSC_HAVE_FORTRAN_UNDERSCORE)
3421:   #define matmpibaijsetvaluesblocked_ matmpibaijsetvaluesblocked
3422: #endif

3424: // PetscClangLinter pragma disable: -fdoc-synopsis-matching-symbol-name
3425: /*@C
3426:   MatMPIBAIJSetValuesBlocked - Direct Fortran call to replace call to `MatSetValuesBlocked()`

3428:   Collective

3430:   Input Parameters:
3431: + matin  - the matrix
3432: . min    - number of input rows
3433: . im     - input rows
3434: . nin    - number of input columns
3435: . in     - input columns
3436: . v      - numerical values input
3437: - addvin - `INSERT_VALUES` or `ADD_VALUES`

3439:   Level: advanced

3441:   Developer Notes:
3442:   This has a complete copy of `MatSetValuesBlocked_MPIBAIJ()` which is terrible code un-reuse.

3444: .seealso: `Mat`, `MatSetValuesBlocked()`
3445: @*/
3446: PETSC_EXTERN PetscErrorCode matmpibaijsetvaluesblocked_(Mat *matin, PetscInt *min, const PetscInt im[], PetscInt *nin, const PetscInt in[], const MatScalar v[], InsertMode *addvin)
3447: {
3448:   /* convert input arguments to C version */
3449:   Mat        mat = *matin;
3450:   PetscInt   m = *min, n = *nin;
3451:   InsertMode addv = *addvin;

3453:   Mat_MPIBAIJ     *baij = (Mat_MPIBAIJ *)mat->data;
3454:   const MatScalar *value;
3455:   MatScalar       *barray      = baij->barray;
3456:   PetscBool        roworiented = baij->roworiented;
3457:   PetscInt         i, j, ii, jj, row, col, rstart = baij->rstartbs;
3458:   PetscInt         rend = baij->rendbs, cstart = baij->cstartbs, stepval;
3459:   PetscInt         cend = baij->cendbs, bs = mat->rmap->bs, bs2 = baij->bs2;

3461:   PetscFunctionBegin;
3462:   /* tasks normally handled by MatSetValuesBlocked() */
3463:   if (mat->insertmode == NOT_SET_VALUES) mat->insertmode = addv;
3464:   else PetscCheck(mat->insertmode == addv, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Cannot mix add values and insert values");
3465:   PetscCheck(!mat->factortype, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Not for factored matrix");
3466:   if (mat->assembled) {
3467:     mat->was_assembled = PETSC_TRUE;
3468:     mat->assembled     = PETSC_FALSE;
3469:   }
3470:   PetscCall(PetscLogEventBegin(MAT_SetValues, mat, 0, 0, 0));

3472:   if (!barray) {
3473:     PetscCall(PetscMalloc1(bs2, &barray));
3474:     baij->barray = barray;
3475:   }

3477:   if (roworiented) stepval = (n - 1) * bs;
3478:   else stepval = (m - 1) * bs;

3480:   for (i = 0; i < m; i++) {
3481:     if (im[i] < 0) continue;
3482:     PetscCheck(im[i] < baij->Mbs, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Row too large, row %" PetscInt_FMT " max %" PetscInt_FMT, im[i], baij->Mbs - 1);
3483:     if (im[i] >= rstart && im[i] < rend) {
3484:       row = im[i] - rstart;
3485:       for (j = 0; j < n; j++) {
3486:         /* If NumCol = 1 then a copy is not required */
3487:         if ((roworiented) && (n == 1)) {
3488:           barray = (MatScalar *)v + i * bs2;
3489:         } else if ((!roworiented) && (m == 1)) {
3490:           barray = (MatScalar *)v + j * bs2;
3491:         } else { /* Here a copy is required */
3492:           if (roworiented) {
3493:             value = v + i * (stepval + bs) * bs + j * bs;
3494:           } else {
3495:             value = v + j * (stepval + bs) * bs + i * bs;
3496:           }
3497:           for (ii = 0; ii < bs; ii++, value += stepval) {
3498:             for (jj = 0; jj < bs; jj++) *barray++ = *value++;
3499:           }
3500:           barray -= bs2;
3501:         }

3503:         if (in[j] >= cstart && in[j] < cend) {
3504:           col = in[j] - cstart;
3505:           PetscCall(MatSetValuesBlocked_SeqBAIJ_Inlined(baij->A, row, col, barray, addv, im[i], in[j]));
3506:         } else if (in[j] < 0) {
3507:           continue;
3508:         } else {
3509:           PetscCheck(in[j] < baij->Nbs, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Column too large, col %" PetscInt_FMT " max %" PetscInt_FMT, in[j], baij->Nbs - 1);
3510:           if (mat->was_assembled) {
3511:             if (!baij->colmap) PetscCall(MatCreateColmap_MPIBAIJ_Private(mat));

3513: #if defined(PETSC_USE_DEBUG)
3514:   #if defined(PETSC_USE_CTABLE)
3515:             {
3516:               PetscInt data;
3517:               PetscCall(PetscHMapIGetWithDefault(baij->colmap, in[j] + 1, 0, &data));
3518:               PetscCheck((data - 1) % bs == 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Incorrect colmap");
3519:             }
3520:   #else
3521:             PetscCheck((baij->colmap[in[j]] - 1) % bs == 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Incorrect colmap");
3522:   #endif
3523: #endif
3524: #if defined(PETSC_USE_CTABLE)
3525:             PetscCall(PetscHMapIGetWithDefault(baij->colmap, in[j] + 1, 0, &col));
3526:             col = (col - 1) / bs;
3527: #else
3528:             col = (baij->colmap[in[j]] - 1) / bs;
3529: #endif
3530:             if (col < 0 && !((Mat_SeqBAIJ *)baij->A->data)->nonew) {
3531:               PetscCall(MatDisAssemble_MPIBAIJ(mat));
3532:               col = in[j];
3533:             }
3534:           } else col = in[j];
3535:           PetscCall(MatSetValuesBlocked_SeqBAIJ_Inlined(baij->B, row, col, barray, addv, im[i], in[j]));
3536:         }
3537:       }
3538:     } else {
3539:       if (!baij->donotstash) {
3540:         if (roworiented) {
3541:           PetscCall(MatStashValuesRowBlocked_Private(&mat->bstash, im[i], n, in, v, m, n, i));
3542:         } else {
3543:           PetscCall(MatStashValuesColBlocked_Private(&mat->bstash, im[i], n, in, v, m, n, i));
3544:         }
3545:       }
3546:     }
3547:   }

3549:   /* task normally handled by MatSetValuesBlocked() */
3550:   PetscCall(PetscLogEventEnd(MAT_SetValues, mat, 0, 0, 0));
3551:   PetscFunctionReturn(PETSC_SUCCESS);
3552: }

3554: /*@
3555:   MatCreateMPIBAIJWithArrays - creates a `MATMPIBAIJ` matrix using arrays that contain in standard block CSR format for the local rows.

3557:   Collective

3559:   Input Parameters:
3560: + comm - MPI communicator
3561: . bs   - the block size, only a block size of 1 is supported
3562: . m    - number of local rows (Cannot be `PETSC_DECIDE`)
3563: . n    - This value should be the same as the local size used in creating the
3564:          x vector for the matrix-vector product $ y = Ax $. (or `PETSC_DECIDE` to have
3565:          calculated if `N` is given) For square matrices `n` is almost always `m`.
3566: . M    - number of global rows (or `PETSC_DETERMINE` to have calculated if `m` is given)
3567: . N    - number of global columns (or `PETSC_DETERMINE` to have calculated if `n` is given)
3568: . i    - row indices; that is i[0] = 0, i[row] = i[row-1] + number of block elements in that rowth block row of the matrix
3569: . j    - column indices
3570: - a    - matrix values

3572:   Output Parameter:
3573: . mat - the matrix

3575:   Level: intermediate

3577:   Notes:
3578:   The `i`, `j`, and `a` arrays ARE copied by this routine into the internal format used by PETSc;
3579:   thus you CANNOT change the matrix entries by changing the values of a[] after you have
3580:   called this routine. Use `MatCreateMPIAIJWithSplitArrays()` to avoid needing to copy the arrays.

3582:   The order of the entries in values is the same as the block compressed sparse row storage format; that is, it is
3583:   the same as a three dimensional array in Fortran values(bs,bs,nnz) that contains the first column of the first
3584:   block, followed by the second column of the first block etc etc.  That is, the blocks are contiguous in memory
3585:   with column-major ordering within blocks.

3587:   The `i` and `j` indices are 0 based, and `i` indices are indices corresponding to the local `j` array.

3589: .seealso: `Mat`, `MatCreate()`, `MatCreateSeqAIJ()`, `MatSetValues()`, `MatMPIAIJSetPreallocation()`, `MatMPIAIJSetPreallocationCSR()`,
3590:           `MATMPIAIJ`, `MatCreateAIJ()`, `MatCreateMPIAIJWithSplitArrays()`
3591: @*/
3592: PetscErrorCode MatCreateMPIBAIJWithArrays(MPI_Comm comm, PetscInt bs, PetscInt m, PetscInt n, PetscInt M, PetscInt N, const PetscInt i[], const PetscInt j[], const PetscScalar a[], Mat *mat)
3593: {
3594:   PetscFunctionBegin;
3595:   PetscCheck(!i[0], PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "i (row indices) must start with 0");
3596:   PetscCheck(m >= 0, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "local number of rows (m) cannot be PETSC_DECIDE, or negative");
3597:   PetscCall(MatCreate(comm, mat));
3598:   PetscCall(MatSetSizes(*mat, m, n, M, N));
3599:   PetscCall(MatSetType(*mat, MATMPIBAIJ));
3600:   PetscCall(MatSetBlockSize(*mat, bs));
3601:   PetscCall(MatSetUp(*mat));
3602:   PetscCall(MatSetOption(*mat, MAT_ROW_ORIENTED, PETSC_FALSE));
3603:   PetscCall(MatMPIBAIJSetPreallocationCSR(*mat, bs, i, j, a));
3604:   PetscCall(MatSetOption(*mat, MAT_ROW_ORIENTED, PETSC_TRUE));
3605:   PetscFunctionReturn(PETSC_SUCCESS);
3606: }

3608: PetscErrorCode MatCreateMPIMatConcatenateSeqMat_MPIBAIJ(MPI_Comm comm, Mat inmat, PetscInt n, MatReuse scall, Mat *outmat)
3609: {
3610:   PetscInt     m, N, i, rstart, nnz, Ii, bs, cbs;
3611:   PetscInt    *indx;
3612:   PetscScalar *values;

3614:   PetscFunctionBegin;
3615:   PetscCall(MatGetSize(inmat, &m, &N));
3616:   if (scall == MAT_INITIAL_MATRIX) { /* symbolic phase */
3617:     Mat_SeqBAIJ *a = (Mat_SeqBAIJ *)inmat->data;
3618:     PetscInt    *dnz, *onz, mbs, Nbs, nbs;
3619:     PetscInt    *bindx, rmax = a->rmax, j;
3620:     PetscMPIInt  rank, size;

3622:     PetscCall(MatGetBlockSizes(inmat, &bs, &cbs));
3623:     mbs = m / bs;
3624:     Nbs = N / cbs;
3625:     if (n == PETSC_DECIDE) PetscCall(PetscSplitOwnershipBlock(comm, cbs, &n, &N));
3626:     nbs = n / cbs;

3628:     PetscCall(PetscMalloc1(rmax, &bindx));
3629:     MatPreallocateBegin(comm, mbs, nbs, dnz, onz); /* inline function, output __end and __rstart are used below */

3631:     PetscCallMPI(MPI_Comm_rank(comm, &rank));
3632:     PetscCallMPI(MPI_Comm_size(comm, &size));
3633:     if (rank == size - 1) {
3634:       /* Check sum(nbs) = Nbs */
3635:       PetscCheck(__end == Nbs, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Sum of local block columns %" PetscInt_FMT " != global block columns %" PetscInt_FMT, __end, Nbs);
3636:     }

3638:     rstart = __rstart; /* block rstart of *outmat; see inline function MatPreallocateBegin */
3639:     for (i = 0; i < mbs; i++) {
3640:       PetscCall(MatGetRow_SeqBAIJ(inmat, i * bs, &nnz, &indx, NULL)); /* non-blocked nnz and indx */
3641:       nnz = nnz / bs;
3642:       for (j = 0; j < nnz; j++) bindx[j] = indx[j * bs] / bs;
3643:       PetscCall(MatPreallocateSet(i + rstart, nnz, bindx, dnz, onz));
3644:       PetscCall(MatRestoreRow_SeqBAIJ(inmat, i * bs, &nnz, &indx, NULL));
3645:     }
3646:     PetscCall(PetscFree(bindx));

3648:     PetscCall(MatCreate(comm, outmat));
3649:     PetscCall(MatSetSizes(*outmat, m, n, PETSC_DETERMINE, PETSC_DETERMINE));
3650:     PetscCall(MatSetBlockSizes(*outmat, bs, cbs));
3651:     PetscCall(MatSetType(*outmat, MATBAIJ));
3652:     PetscCall(MatSeqBAIJSetPreallocation(*outmat, bs, 0, dnz));
3653:     PetscCall(MatMPIBAIJSetPreallocation(*outmat, bs, 0, dnz, 0, onz));
3654:     MatPreallocateEnd(dnz, onz);
3655:     PetscCall(MatSetOption(*outmat, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
3656:   }

3658:   /* numeric phase */
3659:   PetscCall(MatGetBlockSizes(inmat, &bs, &cbs));
3660:   PetscCall(MatGetOwnershipRange(*outmat, &rstart, NULL));

3662:   for (i = 0; i < m; i++) {
3663:     PetscCall(MatGetRow_SeqBAIJ(inmat, i, &nnz, &indx, &values));
3664:     Ii = i + rstart;
3665:     PetscCall(MatSetValues(*outmat, 1, &Ii, nnz, indx, values, INSERT_VALUES));
3666:     PetscCall(MatRestoreRow_SeqBAIJ(inmat, i, &nnz, &indx, &values));
3667:   }
3668:   PetscCall(MatAssemblyBegin(*outmat, MAT_FINAL_ASSEMBLY));
3669:   PetscCall(MatAssemblyEnd(*outmat, MAT_FINAL_ASSEMBLY));
3670:   PetscFunctionReturn(PETSC_SUCCESS);
3671: }