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: #if PetscDefined(HAVE_LIBXSMM)
  8: PETSC_INTERN PetscErrorCode MatConvert_MPIBAIJ_MPIBAIJLIBXSMM(Mat, MatType, MatReuse, Mat *);
  9: #endif

 11: static PetscErrorCode MatDestroy_MPIBAIJ(Mat mat)
 12: {
 13:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)mat->data;

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

 35:   PetscCall(PetscObjectChangeTypeName((PetscObject)mat, NULL));
 36:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatStoreValues_C", NULL));
 37:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatRetrieveValues_C", NULL));
 38:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatGetMultPetscSF_C", NULL));
 39:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMPIBAIJSetPreallocation_C", NULL));
 40:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMPIBAIJSetPreallocationCSR_C", NULL));
 41:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDiagonalScaleLocal_C", NULL));
 42:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatSetHashTableFactor_C", NULL));
 43:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpibaij_mpidense_C", NULL));
 44:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpibaij_mpisbaij_C", NULL));
 45:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpibaij_mpiadj_C", NULL));
 46:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpibaij_mpiaij_C", NULL));
 47: #if PetscDefined(HAVE_HYPRE)
 48:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpibaij_hypre_C", NULL));
 49: #endif
 50:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpibaij_is_C", NULL));
 51: #if PetscDefined(HAVE_LIBXSMM)
 52:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpibaij_mpibaijlibxsmm_C", NULL));
 53:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpibaijlibxsmm_mpidense_C", NULL));
 54:   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpibaijlibxsmm_mpibaij_C", NULL));
 55: #endif
 56:   PetscFunctionReturn(PETSC_SUCCESS);
 57: }

 59: /* defines MatSetValues_MPI_Hash(), MatAssemblyBegin_MPI_Hash(), and  MatAssemblyEnd_MPI_Hash() */
 60: #define TYPE BAIJ
 61: #include "../src/mat/impls/aij/mpi/mpihashmat.h"
 62: #undef TYPE

 64: #if PetscDefined(HAVE_HYPRE)
 65: PETSC_INTERN PetscErrorCode MatConvert_AIJ_HYPRE(Mat, MatType, MatReuse, Mat *);
 66: #endif

 68: static PetscErrorCode MatGetRowMaxAbs_MPIBAIJ(Mat A, Vec v, PetscInt idx[])
 69: {
 70:   Mat_MPIBAIJ       *a = (Mat_MPIBAIJ *)A->data;
 71:   PetscInt           i, *idxb = NULL, m = A->rmap->n, bs = A->cmap->bs;
 72:   PetscScalar       *vv;
 73:   Vec                vB, vA;
 74:   const PetscScalar *va, *vb;

 76:   PetscFunctionBegin;
 77:   PetscCall(MatCreateVecs(a->A, NULL, &vA));
 78:   PetscCall(MatGetRowMaxAbs(a->A, vA, idx));

 80:   PetscCall(VecGetArrayRead(vA, &va));
 81:   if (idx) {
 82:     for (i = 0; i < m; i++) {
 83:       if (PetscAbsScalar(va[i])) idx[i] += A->cmap->rstart;
 84:     }
 85:   }

 87:   PetscCall(MatCreateVecs(a->B, NULL, &vB));
 88:   PetscCall(PetscMalloc1(m, &idxb));
 89:   PetscCall(MatGetRowMaxAbs(a->B, vB, idxb));

 91:   PetscCall(VecGetArrayWrite(v, &vv));
 92:   PetscCall(VecGetArrayRead(vB, &vb));
 93:   for (i = 0; i < m; i++) {
 94:     if (PetscAbsScalar(va[i]) < PetscAbsScalar(vb[i])) {
 95:       vv[i] = vb[i];
 96:       if (idx) idx[i] = bs * a->garray[idxb[i] / bs] + (idxb[i] % bs);
 97:     } else {
 98:       vv[i] = va[i];
 99:       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);
100:     }
101:   }
102:   PetscCall(VecRestoreArrayWrite(v, &vv));
103:   PetscCall(VecRestoreArrayRead(vA, &va));
104:   PetscCall(VecRestoreArrayRead(vB, &vb));
105:   PetscCall(PetscFree(idxb));
106:   PetscCall(VecDestroy(&vA));
107:   PetscCall(VecDestroy(&vB));
108:   PetscFunctionReturn(PETSC_SUCCESS);
109: }

111: static PetscErrorCode MatGetRowSumAbs_MPIBAIJ(Mat A, Vec v)
112: {
113:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;
114:   Vec          vB, vA;

116:   PetscFunctionBegin;
117:   PetscCall(MatCreateVecs(a->A, NULL, &vA));
118:   PetscCall(MatGetRowSumAbs(a->A, vA));
119:   PetscCall(MatCreateVecs(a->B, NULL, &vB));
120:   PetscCall(MatGetRowSumAbs(a->B, vB));
121:   PetscCall(VecAXPY(vA, 1.0, vB));
122:   PetscCall(VecDestroy(&vB));
123:   PetscCall(VecCopy(vA, v));
124:   PetscCall(VecDestroy(&vA));
125:   PetscFunctionReturn(PETSC_SUCCESS);
126: }

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

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

138: static PetscErrorCode MatRetrieveValues_MPIBAIJ(Mat mat)
139: {
140:   Mat_MPIBAIJ *aij = (Mat_MPIBAIJ *)mat->data;

142:   PetscFunctionBegin;
143:   PetscCall(MatRetrieveValues(aij->A));
144:   PetscCall(MatRetrieveValues(aij->B));
145:   PetscFunctionReturn(PETSC_SUCCESS);
146: }

148: /*
149:      Local utility routine that creates a mapping from the global column
150:    number to the local number in the off-diagonal part of the local
151:    storage of the matrix.  This is done in a non scalable way since the
152:    length of colmap equals the global matrix length.
153: */
154: PetscErrorCode MatCreateColmap_MPIBAIJ_Private(Mat mat)
155: {
156:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)mat->data;
157:   Mat_SeqBAIJ *B    = (Mat_SeqBAIJ *)baij->B->data;
158:   PetscInt     nbs = B->nbs, i, bs = mat->rmap->bs;

160:   PetscFunctionBegin;
161: #if PetscDefined(USE_CTABLE)
162:   PetscCall(PetscHMapICreateWithSize(baij->nbs, &baij->colmap));
163:   for (i = 0; i < nbs; i++) PetscCall(PetscHMapISet(baij->colmap, baij->garray[i] + 1, i * bs + 1));
164: #else
165:   PetscCall(PetscCalloc1(baij->Nbs + 1, &baij->colmap));
166:   for (i = 0; i < nbs; i++) baij->colmap[baij->garray[i]] = i * bs + 1;
167: #endif
168:   PetscFunctionReturn(PETSC_SUCCESS);
169: }

171: #define MatSetValues_SeqBAIJ_A_Private(row, col, value, addv, orow, ocol) \
172:   do { \
173:     brow = (row) / bs; \
174:     rp   = PetscSafePointerPlusOffset(aj, ai[brow]); \
175:     if (!A->structure_only) ap = PetscSafePointerPlusOffset(aa, bs2 * ai[brow]); \
176:     rmax = aimax[brow]; \
177:     nrow = ailen[brow]; \
178:     bcol = (col) / bs; \
179:     ridx = (row) % bs; \
180:     cidx = (col) % bs; \
181:     low  = 0; \
182:     high = nrow; \
183:     while (high - low > 3) { \
184:       t = (low + high) / 2; \
185:       if (rp[t] > bcol) high = t; \
186:       else low = t; \
187:     } \
188:     for (_i = low; _i < high; _i++) { \
189:       if (rp[_i] > bcol) break; \
190:       if (rp[_i] == bcol) { \
191:         if (A->structure_only) goto a_noinsert; \
192:         bap = ap + bs2 * _i + bs * cidx + ridx; \
193:         if (addv == ADD_VALUES) *bap += value; \
194:         else *bap = value; \
195:         goto a_noinsert; \
196:       } \
197:     } \
198:     if (a->nonew == 1) goto a_noinsert; \
199:     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); \
200:     if (A->structure_only) MatSeqXAIJReallocateAIJ_structure_only(A, a->mbs, bs2, nrow, brow, bcol, rmax, ai, aj, rp, aimax, a->nonew, MatScalar); \
201:     else MatSeqXAIJReallocateAIJ(A, a->mbs, bs2, nrow, brow, bcol, rmax, aa, ai, aj, rp, ap, aimax, a->nonew, MatScalar); \
202:     N = nrow++ - 1; \
203:     /* shift up all the later entries in this row */ \
204:     PetscCall(PetscArraymove(rp + _i + 1, rp + _i, N - _i + 1)); \
205:     rp[_i] = bcol; \
206:     if (!A->structure_only) { \
207:       PetscCall(PetscArraymove(ap + bs2 * (_i + 1), ap + bs2 * _i, bs2 * (N - _i + 1))); \
208:       PetscCall(PetscArrayzero(ap + bs2 * _i, bs2)); \
209:       ap[bs2 * _i + bs * cidx + ridx] = value; \
210:     } \
211:   a_noinsert:; \
212:     ailen[brow] = nrow; \
213:   } while (0)

215: #define MatSetValues_SeqBAIJ_B_Private(row, col, value, addv, orow, ocol) \
216:   do { \
217:     brow = (row) / bs; \
218:     rp   = PetscSafePointerPlusOffset(bj, bi[brow]); \
219:     if (!B->structure_only) ap = PetscSafePointerPlusOffset(ba, bs2 * bi[brow]); \
220:     rmax = bimax[brow]; \
221:     nrow = bilen[brow]; \
222:     bcol = (col) / bs; \
223:     ridx = (row) % bs; \
224:     cidx = (col) % bs; \
225:     low  = 0; \
226:     high = nrow; \
227:     while (high - low > 3) { \
228:       t = (low + high) / 2; \
229:       if (rp[t] > bcol) high = t; \
230:       else low = t; \
231:     } \
232:     for (_i = low; _i < high; _i++) { \
233:       if (rp[_i] > bcol) break; \
234:       if (rp[_i] == bcol) { \
235:         if (B->structure_only) goto b_noinsert; \
236:         bap = ap + bs2 * _i + bs * cidx + ridx; \
237:         if (addv == ADD_VALUES) *bap += value; \
238:         else *bap = value; \
239:         goto b_noinsert; \
240:       } \
241:     } \
242:     if (b->nonew == 1) goto b_noinsert; \
243:     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); \
244:     if (B->structure_only) MatSeqXAIJReallocateAIJ_structure_only(B, b->mbs, bs2, nrow, brow, bcol, rmax, bi, bj, rp, bimax, b->nonew, MatScalar); \
245:     else MatSeqXAIJReallocateAIJ(B, b->mbs, bs2, nrow, brow, bcol, rmax, ba, bi, bj, rp, ap, bimax, b->nonew, MatScalar); \
246:     N = nrow++ - 1; \
247:     /* shift up all the later entries in this row */ \
248:     PetscCall(PetscArraymove(rp + _i + 1, rp + _i, N - _i + 1)); \
249:     rp[_i] = bcol; \
250:     if (!B->structure_only) { \
251:       PetscCall(PetscArraymove(ap + bs2 * (_i + 1), ap + bs2 * _i, bs2 * (N - _i + 1))); \
252:       PetscCall(PetscArrayzero(ap + bs2 * _i, bs2)); \
253:       ap[bs2 * _i + bs * cidx + ridx] = value; \
254:     } \
255:   b_noinsert:; \
256:     bilen[brow] = nrow; \
257:   } while (0)

259: static PetscErrorCode MatSetValues_MPIBAIJ(Mat mat, PetscInt m, const PetscInt im[], PetscInt n, const PetscInt in[], const PetscScalar v[], InsertMode addv)
260: {
261:   Mat_MPIBAIJ *baij        = (Mat_MPIBAIJ *)mat->data;
262:   MatScalar    value       = 0.0;
263:   PetscBool    roworiented = baij->roworiented;
264:   PetscInt     i, j, row, col;
265:   PetscInt     rstart_orig = mat->rmap->rstart;
266:   PetscInt     rend_orig = mat->rmap->rend, cstart_orig = mat->cmap->rstart;
267:   PetscInt     cend_orig = mat->cmap->rend, bs = mat->rmap->bs;

269:   /* Some Variables required in the macro */
270:   Mat          A     = baij->A;
271:   Mat_SeqBAIJ *a     = (Mat_SeqBAIJ *)A->data;
272:   PetscInt    *aimax = a->imax, *ai = a->i, *ailen = a->ilen, *aj = a->j;
273:   MatScalar   *aa = a->a;

275:   Mat          B     = baij->B;
276:   Mat_SeqBAIJ *b     = (Mat_SeqBAIJ *)B->data;
277:   PetscInt    *bimax = b->imax, *bi = b->i, *bilen = b->ilen, *bj = b->j;
278:   MatScalar   *ba = b->a;

280:   PetscInt  *rp, ii, nrow, _i, rmax, N, brow, bcol;
281:   PetscInt   low, high, t, ridx, cidx, bs2 = a->bs2;
282:   MatScalar *ap = NULL, *bap;

284:   PetscFunctionBegin;
285:   for (i = 0; i < m; i++) {
286:     if (im[i] < 0) continue;
287:     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);
288:     if (im[i] >= rstart_orig && im[i] < rend_orig) {
289:       row = im[i] - rstart_orig;
290:       for (j = 0; j < n; j++) {
291:         if (in[j] >= cstart_orig && in[j] < cend_orig) {
292:           col = in[j] - cstart_orig;
293:           if (!mat->structure_only) {
294:             if (roworiented) value = v[i * n + j];
295:             else value = v[i + j * m];
296:           }
297:           MatSetValues_SeqBAIJ_A_Private(row, col, value, addv, im[i], in[j]);
298:         } else if (in[j] < 0) {
299:           continue;
300:         } else {
301:           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);
302:           if (mat->was_assembled) {
303:             if (!baij->colmap) PetscCall(MatCreateColmap_MPIBAIJ_Private(mat));
304: #if PetscDefined(USE_CTABLE)
305:             PetscCall(PetscHMapIGetWithDefault(baij->colmap, in[j] / bs + 1, 0, &col));
306:             col = col - 1;
307: #else
308:             col = baij->colmap[in[j] / bs] - 1;
309: #endif
310:             if (col < 0 && !((Mat_SeqBAIJ *)baij->B->data)->nonew) {
311:               PetscCall(MatDisAssemble_MPIBAIJ(mat));
312:               col = in[j];
313:               /* Reinitialize the variables required by MatSetValues_SeqBAIJ_B_Private() */
314:               B     = baij->B;
315:               b     = (Mat_SeqBAIJ *)B->data;
316:               bimax = b->imax;
317:               bi    = b->i;
318:               bilen = b->ilen;
319:               bj    = b->j;
320:               ba    = b->a;
321:             } else {
322:               PetscCheck(col >= 0, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Inserting a new nonzero (%" PetscInt_FMT ", %" PetscInt_FMT ") into matrix", im[i], in[j]);
323:               col += in[j] % bs;
324:             }
325:           } else col = in[j];
326:           if (!mat->structure_only) {
327:             if (roworiented) value = v[i * n + j];
328:             else value = v[i + j * m];
329:           }
330:           MatSetValues_SeqBAIJ_B_Private(row, col, value, addv, im[i], in[j]);
331:           /* PetscCall(MatSetValues_SeqBAIJ(baij->B,1,&row,1,&col,&value,addv)); */
332:         }
333:       }
334:     } else {
335:       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]);
336:       if (!baij->donotstash) {
337:         mat->assembled = PETSC_FALSE;
338:         if (roworiented) {
339:           PetscCall(MatStashValuesRow_Private(&mat->stash, im[i], n, in, PetscSafePointerPlusOffset(v, i * n), PETSC_FALSE));
340:         } else {
341:           PetscCall(MatStashValuesCol_Private(&mat->stash, im[i], n, in, PetscSafePointerPlusOffset(v, i), m, PETSC_FALSE));
342:         }
343:       }
344:     }
345:   }
346:   PetscFunctionReturn(PETSC_SUCCESS);
347: }

349: static inline PetscErrorCode MatSetValuesBlocked_SeqBAIJ_Inlined(Mat A, PetscInt row, PetscInt col, const PetscScalar v[], InsertMode is, PetscInt orow, PetscInt ocol)
350: {
351:   Mat_SeqBAIJ       *a = (Mat_SeqBAIJ *)A->data;
352:   PetscInt          *rp, low, high, t, ii, jj, nrow, i, rmax, N;
353:   PetscInt          *imax = a->imax, *ai = a->i, *ailen = a->ilen;
354:   PetscInt          *aj = a->j, nonew = a->nonew, bs2 = a->bs2, bs = A->rmap->bs;
355:   PetscBool          roworiented = a->roworiented;
356:   const PetscScalar *value       = v;
357:   MatScalar         *ap = NULL, *aa = a->a, *bap;

359:   PetscFunctionBegin;
360:   rp    = aj + ai[row];
361:   ap    = PetscSafePointerPlusOffset(aa, bs2 * ai[row]);
362:   rmax  = imax[row];
363:   nrow  = ailen[row];
364:   value = v;
365:   low   = 0;
366:   high  = nrow;
367:   while (high - low > 7) {
368:     t = (low + high) / 2;
369:     if (rp[t] > col) high = t;
370:     else low = t;
371:   }
372:   for (i = low; i < high; i++) {
373:     if (rp[i] > col) break;
374:     if (rp[i] == col) {
375:       if (A->structure_only) goto noinsert2;
376:       bap = ap + bs2 * i;
377:       if (roworiented) {
378:         if (is == ADD_VALUES) {
379:           for (ii = 0; ii < bs; ii++) {
380:             for (jj = ii; jj < bs2; jj += bs) bap[jj] += *value++;
381:           }
382:         } else {
383:           for (ii = 0; ii < bs; ii++) {
384:             for (jj = ii; jj < bs2; jj += bs) bap[jj] = *value++;
385:           }
386:         }
387:       } else {
388:         if (is == ADD_VALUES) {
389:           for (ii = 0; ii < bs; ii++, value += bs) {
390:             for (jj = 0; jj < bs; jj++) bap[jj] += value[jj];
391:             bap += bs;
392:           }
393:         } else {
394:           for (ii = 0; ii < bs; ii++, value += bs) {
395:             for (jj = 0; jj < bs; jj++) bap[jj] = value[jj];
396:             bap += bs;
397:           }
398:         }
399:       }
400:       goto noinsert2;
401:     }
402:   }
403:   if (nonew == 1) goto noinsert2;
404:   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);
405:   if (A->structure_only) MatSeqXAIJReallocateAIJ_structure_only(A, a->mbs, bs2, nrow, row, col, rmax, ai, aj, rp, imax, nonew, MatScalar);
406:   else MatSeqXAIJReallocateAIJ(A, a->mbs, bs2, nrow, row, col, rmax, aa, ai, aj, rp, ap, imax, nonew, MatScalar);
407:   N = nrow++ - 1;
408:   high++;
409:   /* shift up all the later entries in this row */
410:   PetscCall(PetscArraymove(rp + i + 1, rp + i, N - i + 1));
411:   rp[i] = col;
412:   if (!A->structure_only) {
413:     PetscCall(PetscArraymove(ap + bs2 * (i + 1), ap + bs2 * i, bs2 * (N - i + 1)));
414:     bap = ap + bs2 * i;
415:     if (roworiented) {
416:       for (ii = 0; ii < bs; ii++) {
417:         for (jj = ii; jj < bs2; jj += bs) bap[jj] = *value++;
418:       }
419:     } else {
420:       for (ii = 0; ii < bs; ii++) {
421:         for (jj = 0; jj < bs; jj++) *bap++ = *value++;
422:       }
423:     }
424:   }
425: noinsert2:;
426:   ailen[row] = nrow;
427:   PetscFunctionReturn(PETSC_SUCCESS);
428: }

430: /*
431:     This routine should be optimized so that the block copy at ** Here a copy is required ** below is not needed
432:     by passing additional stride information into the MatSetValuesBlocked_SeqBAIJ_Inlined() routine
433: */
434: static PetscErrorCode MatSetValuesBlocked_MPIBAIJ(Mat mat, PetscInt m, const PetscInt im[], PetscInt n, const PetscInt in[], const PetscScalar v[], InsertMode addv)
435: {
436:   Mat_MPIBAIJ       *baij = (Mat_MPIBAIJ *)mat->data;
437:   const PetscScalar *value;
438:   MatScalar         *barray      = baij->barray;
439:   PetscBool          roworiented = baij->roworiented;
440:   PetscInt           i, j, ii, jj, row, col, rstart = baij->rstartbs;
441:   PetscInt           rend = baij->rendbs, cstart = baij->cstartbs, stepval;
442:   PetscInt           cend = baij->cendbs, bs = mat->rmap->bs, bs2 = baij->bs2;

444:   PetscFunctionBegin;
445:   if (!mat->structure_only && !barray) {
446:     PetscCall(PetscMalloc1(bs2, &barray));
447:     baij->barray = barray;
448:   }

450:   if (roworiented) stepval = (n - 1) * bs;
451:   else stepval = (m - 1) * bs;

453:   for (i = 0; i < m; i++) {
454:     if (im[i] < 0) continue;
455:     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);
456:     if (im[i] >= rstart && im[i] < rend) {
457:       row = im[i] - rstart;
458:       for (j = 0; j < n; j++) {
459:         if (!mat->structure_only) {
460:           /* If NumCol = 1 then a copy is not required */
461:           if (roworiented && (n == 1)) {
462:             barray = (MatScalar *)v + i * bs2;
463:           } else if ((!roworiented) && (m == 1)) {
464:             barray = (MatScalar *)v + j * bs2;
465:           } else { /* Here a copy is required */
466:             if (roworiented) {
467:               value = v + (i * (stepval + bs) + j) * bs;
468:             } else {
469:               value = v + (j * (stepval + bs) + i) * bs;
470:             }
471:             for (ii = 0; ii < bs; ii++, value += bs + stepval) {
472:               for (jj = 0; jj < bs; jj++) barray[jj] = value[jj];
473:               barray += bs;
474:             }
475:             barray -= bs2;
476:           }
477:         }

479:         if (in[j] >= cstart && in[j] < cend) {
480:           col = in[j] - cstart;
481:           PetscCall(MatSetValuesBlocked_SeqBAIJ_Inlined(baij->A, row, col, barray, addv, im[i], in[j]));
482:         } else if (in[j] < 0) {
483:           continue;
484:         } else {
485:           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);
486:           if (mat->was_assembled) {
487:             if (!baij->colmap) PetscCall(MatCreateColmap_MPIBAIJ_Private(mat));

489: #if PetscDefined(USE_CTABLE)
490:             PetscCall(PetscHMapIGetWithDefault(baij->colmap, in[j] + 1, 0, &col));
491:             col = col < 1 ? -1 : (col - 1) / bs;
492: #else
493:             col = baij->colmap[in[j]] < 1 ? -1 : (baij->colmap[in[j]] - 1) / bs;
494: #endif
495:             if (col < 0 && !((Mat_SeqBAIJ *)baij->B->data)->nonew) {
496:               PetscCall(MatDisAssemble_MPIBAIJ(mat));
497:               col = in[j];
498:             } 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]);
499:           } else col = in[j];
500:           PetscCall(MatSetValuesBlocked_SeqBAIJ_Inlined(baij->B, row, col, barray, addv, im[i], in[j]));
501:         }
502:       }
503:     } else {
504:       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]);
505:       if (!baij->donotstash) {
506:         if (roworiented) {
507:           PetscCall(MatStashValuesRowBlocked_Private(&mat->bstash, im[i], n, in, v, m, n, i));
508:         } else {
509:           PetscCall(MatStashValuesColBlocked_Private(&mat->bstash, im[i], n, in, v, m, n, i));
510:         }
511:       }
512:     }
513:   }
514:   PetscFunctionReturn(PETSC_SUCCESS);
515: }

517: #define HASH_KEY             0.6180339887
518: #define HASH(size, key, tmp) (tmp = (key) * HASH_KEY, (PetscInt)((size) * ((tmp) - (PetscInt)(tmp))))
519: /* #define HASH(size,key) ((PetscInt)((size)*fmod(((key)*HASH_KEY),1))) */
520: /* #define HASH(size,key,tmp) ((PetscInt)((size)*fmod(((key)*HASH_KEY),1))) */
521: static PetscErrorCode MatSetValues_MPIBAIJ_HT(Mat mat, PetscInt m, const PetscInt im[], PetscInt n, const PetscInt in[], const PetscScalar v[], InsertMode addv)
522: {
523:   Mat_MPIBAIJ *baij        = (Mat_MPIBAIJ *)mat->data;
524:   PetscBool    roworiented = baij->roworiented;
525:   PetscInt     i, j, row, col;
526:   PetscInt     rstart_orig = mat->rmap->rstart;
527:   PetscInt     rend_orig = mat->rmap->rend, Nbs = baij->Nbs;
528:   PetscInt     h1, key, size = baij->ht_size, bs = mat->rmap->bs, *HT = baij->ht, idx;
529:   PetscReal    tmp;
530:   MatScalar  **HD       = baij->hd, value;
531:   PetscInt     total_ct = baij->ht_total_ct, insert_ct = baij->ht_insert_ct;

533:   PetscFunctionBegin;
534:   for (i = 0; i < m; i++) {
535:     if (PetscDefined(USE_DEBUG)) {
536:       PetscCheck(im[i] >= 0, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Negative row");
537:       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);
538:     }
539:     row = im[i];
540:     if (row >= rstart_orig && row < rend_orig) {
541:       for (j = 0; j < n; j++) {
542:         col = in[j];
543:         if (roworiented) value = v[i * n + j];
544:         else value = v[i + j * m];
545:         /* Look up PetscInto the Hash Table */
546:         key = (row / bs) * Nbs + (col / bs) + 1;
547:         h1  = HASH(size, key, tmp);

549:         idx = h1;
550:         if (PetscDefined(USE_DEBUG)) {
551:           insert_ct++;
552:           total_ct++;
553:           if (HT[idx] != key) {
554:             for (idx = h1; (idx < size) && (HT[idx] != key); idx++, total_ct++);
555:             if (idx == size) {
556:               for (idx = 0; (idx < h1) && (HT[idx] != key); idx++, total_ct++);
557:               PetscCheck(idx != h1, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "(%" PetscInt_FMT ",%" PetscInt_FMT ") has no entry in the hash table", row, col);
558:             }
559:           }
560:         } else if (HT[idx] != key) {
561:           for (idx = h1; (idx < size) && (HT[idx] != key); idx++);
562:           if (idx == size) {
563:             for (idx = 0; (idx < h1) && (HT[idx] != key); idx++);
564:             PetscCheck(idx != h1, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "(%" PetscInt_FMT ",%" PetscInt_FMT ") has no entry in the hash table", row, col);
565:           }
566:         }
567:         /* A HASH table entry is found, so insert the values at the correct address */
568:         if (addv == ADD_VALUES) *(HD[idx] + (col % bs) * bs + (row % bs)) += value;
569:         else *(HD[idx] + (col % bs) * bs + (row % bs)) = value;
570:       }
571:     } else if (!baij->donotstash) {
572:       if (roworiented) {
573:         PetscCall(MatStashValuesRow_Private(&mat->stash, im[i], n, in, v + i * n, PETSC_FALSE));
574:       } else {
575:         PetscCall(MatStashValuesCol_Private(&mat->stash, im[i], n, in, v + i, m, PETSC_FALSE));
576:       }
577:     }
578:   }
579:   if (PetscDefined(USE_DEBUG)) {
580:     baij->ht_total_ct += total_ct;
581:     baij->ht_insert_ct += insert_ct;
582:   }
583:   PetscFunctionReturn(PETSC_SUCCESS);
584: }

586: static PetscErrorCode MatSetValuesBlocked_MPIBAIJ_HT(Mat mat, PetscInt m, const PetscInt im[], PetscInt n, const PetscInt in[], const PetscScalar v[], InsertMode addv)
587: {
588:   Mat_MPIBAIJ       *baij        = (Mat_MPIBAIJ *)mat->data;
589:   PetscBool          roworiented = baij->roworiented;
590:   PetscInt           i, j, ii, jj, row, col;
591:   PetscInt           rstart = baij->rstartbs;
592:   PetscInt           rend = mat->rmap->rend, stepval, bs = mat->rmap->bs, bs2 = baij->bs2, nbs2 = n * bs2;
593:   PetscInt           h1, key, size = baij->ht_size, idx, *HT = baij->ht, Nbs = baij->Nbs;
594:   PetscReal          tmp;
595:   MatScalar        **HD = baij->hd, *baij_a;
596:   const PetscScalar *v_t, *value;
597:   PetscInt           total_ct = baij->ht_total_ct, insert_ct = baij->ht_insert_ct;

599:   PetscFunctionBegin;
600:   if (roworiented) stepval = (n - 1) * bs;
601:   else stepval = (m - 1) * bs;

603:   for (i = 0; i < m; i++) {
604:     if (PetscDefined(USE_DEBUG)) {
605:       PetscCheck(im[i] >= 0, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Negative row: %" PetscInt_FMT, im[i]);
606:       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);
607:     }
608:     row = im[i];
609:     v_t = v + i * nbs2;
610:     if (row >= rstart && row < rend) {
611:       for (j = 0; j < n; j++) {
612:         col = in[j];

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

618:         idx = h1;
619:         if (PetscDefined(USE_DEBUG)) {
620:           total_ct++;
621:           insert_ct++;
622:           if (HT[idx] != key) {
623:             for (idx = h1; (idx < size) && (HT[idx] != key); idx++, total_ct++);
624:             if (idx == size) {
625:               for (idx = 0; (idx < h1) && (HT[idx] != key); idx++, total_ct++);
626:               PetscCheck(idx != h1, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "(%" PetscInt_FMT ",%" PetscInt_FMT ") has no entry in the hash table", row, col);
627:             }
628:           }
629:         } else if (HT[idx] != key) {
630:           for (idx = h1; (idx < size) && (HT[idx] != key); idx++);
631:           if (idx == size) {
632:             for (idx = 0; (idx < h1) && (HT[idx] != key); idx++);
633:             PetscCheck(idx != h1, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "(%" PetscInt_FMT ",%" PetscInt_FMT ") has no entry in the hash table", row, col);
634:           }
635:         }
636:         baij_a = HD[idx];
637:         if (roworiented) {
638:           /*value = v + i*(stepval+bs)*bs + j*bs;*/
639:           /* value = v + (i*(stepval+bs)+j)*bs; */
640:           value = v_t;
641:           v_t += bs;
642:           if (addv == ADD_VALUES) {
643:             for (ii = 0; ii < bs; ii++, value += stepval) {
644:               for (jj = ii; jj < bs2; jj += bs) baij_a[jj] += *value++;
645:             }
646:           } else {
647:             for (ii = 0; ii < bs; ii++, value += stepval) {
648:               for (jj = ii; jj < bs2; jj += bs) baij_a[jj] = *value++;
649:             }
650:           }
651:         } else {
652:           value = v + j * (stepval + bs) * bs + i * bs;
653:           if (addv == ADD_VALUES) {
654:             for (ii = 0; ii < bs; ii++, value += stepval, baij_a += bs) {
655:               for (jj = 0; jj < bs; jj++) baij_a[jj] += *value++;
656:             }
657:           } else {
658:             for (ii = 0; ii < bs; ii++, value += stepval, baij_a += bs) {
659:               for (jj = 0; jj < bs; jj++) baij_a[jj] = *value++;
660:             }
661:           }
662:         }
663:       }
664:     } else {
665:       if (!baij->donotstash) {
666:         if (roworiented) {
667:           PetscCall(MatStashValuesRowBlocked_Private(&mat->bstash, im[i], n, in, v, m, n, i));
668:         } else {
669:           PetscCall(MatStashValuesColBlocked_Private(&mat->bstash, im[i], n, in, v, m, n, i));
670:         }
671:       }
672:     }
673:   }
674:   if (PetscDefined(USE_DEBUG)) {
675:     baij->ht_total_ct += total_ct;
676:     baij->ht_insert_ct += insert_ct;
677:   }
678:   PetscFunctionReturn(PETSC_SUCCESS);
679: }

681: static PetscErrorCode MatGetValues_MPIBAIJ(Mat mat, PetscInt m, const PetscInt idxm[], PetscInt n, const PetscInt idxn[], PetscScalar v[])
682: {
683:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)mat->data;
684:   PetscInt     bs = mat->rmap->bs, i, j, bsrstart = mat->rmap->rstart, bsrend = mat->rmap->rend;
685:   PetscInt     bscstart = mat->cmap->rstart, bscend = mat->cmap->rend, row, col, data;
686:   PetscBool    roworiented = baij->roworiented;
687:   PetscScalar *value;

689:   PetscFunctionBegin;
690:   for (i = 0; i < m; i++) {
691:     if (idxm[i] < 0) continue; /* negative row */
692:     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);
693:     PetscCheck(idxm[i] >= bsrstart && idxm[i] < bsrend, PETSC_COMM_SELF, PETSC_ERR_SUP, "Only local values currently supported");
694:     row = idxm[i] - bsrstart;
695:     for (j = 0; j < n; j++) {
696:       if (idxn[j] < 0) continue; /* negative column */
697:       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);
698:       value = roworiented ? &v[j + i * n] : &v[i + j * m];
699:       if (idxn[j] >= bscstart && idxn[j] < bscend) {
700:         col = idxn[j] - bscstart;
701:         PetscCall(MatGetValues_SeqBAIJ(baij->A, 1, &row, 1, &col, value));
702:       } else {
703:         if (!baij->colmap) PetscCall(MatCreateColmap_MPIBAIJ_Private(mat));
704: #if PetscDefined(USE_CTABLE)
705:         PetscCall(PetscHMapIGetWithDefault(baij->colmap, idxn[j] / bs + 1, 0, &data));
706:         data--;
707: #else
708:         data = baij->colmap[idxn[j] / bs] - 1;
709: #endif
710:         if (data < 0 || baij->garray[data / bs] != idxn[j] / bs) *value = 0.0;
711:         else {
712:           col = data + idxn[j] % bs;
713:           PetscCall(MatGetValues_SeqBAIJ(baij->B, 1, &row, 1, &col, value));
714:         }
715:       }
716:     }
717:   }
718:   PetscFunctionReturn(PETSC_SUCCESS);
719: }

721: static PetscErrorCode MatNorm_MPIBAIJ(Mat mat, NormType type, PetscReal *nrm)
722: {
723:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)mat->data;
724:   Mat_SeqBAIJ *amat = (Mat_SeqBAIJ *)baij->A->data, *bmat = (Mat_SeqBAIJ *)baij->B->data;
725:   PetscInt     i, j, bs2 = baij->bs2, bs = baij->A->rmap->bs, nz, row, col;
726:   PetscReal    sum = 0.0;
727:   MatScalar   *v;

729:   PetscFunctionBegin;
730:   if (baij->size == 1) {
731:     PetscCall(MatNorm(baij->A, type, nrm));
732:   } else {
733:     if (type == NORM_FROBENIUS) {
734:       v  = amat->a;
735:       nz = amat->nz * bs2;
736:       for (i = 0; i < nz; i++) {
737:         sum += PetscRealPart(PetscConj(*v) * (*v));
738:         v++;
739:       }
740:       v  = bmat->a;
741:       nz = bmat->nz * bs2;
742:       for (i = 0; i < nz; i++) {
743:         sum += PetscRealPart(PetscConj(*v) * (*v));
744:         v++;
745:       }
746:       PetscCallMPI(MPIU_Allreduce(&sum, nrm, 1, MPIU_REAL, MPIU_SUM, PetscObjectComm((PetscObject)mat)));
747:       *nrm = PetscSqrtReal(*nrm);
748:     } else if (type == NORM_1) { /* max column sum */
749:       Vec          col, bcol;
750:       PetscScalar *array;
751:       PetscInt    *jj, *garray = baij->garray;

753:       PetscCall(MatCreateVecs(mat, &col, NULL));
754:       PetscCall(VecGetArrayWrite(col, &array));
755:       v  = amat->a;
756:       jj = amat->j;
757:       for (i = 0; i < amat->nz; i++) {
758:         for (j = 0; j < bs; j++) {
759:           PetscInt col = bs * *jj + j; /* column index */

761:           for (row = 0; row < bs; row++) array[col] += PetscAbsScalar(*v++);
762:         }
763:         jj++;
764:       }
765:       PetscCall(VecRestoreArrayWrite(col, &array));
766:       PetscCall(MatCreateVecs(baij->B, &bcol, NULL));
767:       PetscCall(VecGetArrayWrite(bcol, &array));
768:       v  = bmat->a;
769:       jj = bmat->j;
770:       for (i = 0; i < bmat->nz; i++) {
771:         for (j = 0; j < bs; j++) {
772:           PetscInt col = bs * *jj + j; /* column index */

774:           for (row = 0; row < bs; row++) array[col] += PetscAbsScalar(*v++);
775:         }
776:         jj++;
777:       }
778:       PetscCall(VecSetValuesBlocked(col, bmat->nbs, garray, array, ADD_VALUES));
779:       PetscCall(VecRestoreArrayWrite(bcol, &array));
780:       PetscCall(VecDestroy(&bcol));
781:       PetscCall(VecAssemblyBegin(col));
782:       PetscCall(VecAssemblyEnd(col));
783:       PetscCall(VecNorm(col, NORM_INFINITY, nrm));
784:       PetscCall(VecDestroy(&col));
785:     } else if (type == NORM_INFINITY) { /* max row sum */
786:       PetscReal *sums;
787:       PetscCall(PetscMalloc1(bs, &sums));
788:       sum = 0.0;
789:       for (j = 0; j < amat->mbs; j++) {
790:         for (row = 0; row < bs; row++) sums[row] = 0.0;
791:         v  = amat->a + bs2 * amat->i[j];
792:         nz = amat->i[j + 1] - amat->i[j];
793:         for (i = 0; i < nz; i++) {
794:           for (col = 0; col < bs; col++) {
795:             for (row = 0; row < bs; row++) {
796:               sums[row] += PetscAbsScalar(*v);
797:               v++;
798:             }
799:           }
800:         }
801:         v  = bmat->a + bs2 * bmat->i[j];
802:         nz = bmat->i[j + 1] - bmat->i[j];
803:         for (i = 0; i < nz; i++) {
804:           for (col = 0; col < bs; col++) {
805:             for (row = 0; row < bs; row++) {
806:               sums[row] += PetscAbsScalar(*v);
807:               v++;
808:             }
809:           }
810:         }
811:         for (row = 0; row < bs; row++) {
812:           if (sums[row] > sum) sum = sums[row];
813:         }
814:       }
815:       PetscCallMPI(MPIU_Allreduce(&sum, nrm, 1, MPIU_REAL, MPIU_MAX, PetscObjectComm((PetscObject)mat)));
816:       PetscCall(PetscFree(sums));
817:     } else SETERRQ(PetscObjectComm((PetscObject)mat), PETSC_ERR_SUP, "No support for this norm yet");
818:   }
819:   PetscFunctionReturn(PETSC_SUCCESS);
820: }

822: /*
823:   Creates the hash table, and sets the table
824:   This table is created only once.
825:   If new entries need to be added to the matrix
826:   then the hash table has to be destroyed and
827:   recreated.
828: */
829: static PetscErrorCode MatCreateHashTable_MPIBAIJ_Private(Mat mat, PetscReal factor)
830: {
831:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)mat->data;
832:   Mat          A = baij->A, B = baij->B;
833:   Mat_SeqBAIJ *a = (Mat_SeqBAIJ *)A->data, *b = (Mat_SeqBAIJ *)B->data;
834:   PetscInt     i, j, k, nz = a->nz + b->nz, h1, *ai = a->i, *aj = a->j, *bi = b->i, *bj = b->j;
835:   PetscInt     ht_size, bs2 = baij->bs2, rstart = baij->rstartbs;
836:   PetscInt     cstart = baij->cstartbs, *garray = baij->garray, row, col, Nbs = baij->Nbs;
837:   PetscInt    *HT, key;
838:   MatScalar  **HD;
839:   PetscReal    tmp;
840:   PetscInt     ct = 0, max = 0;

842:   PetscFunctionBegin;
843:   if (baij->ht) PetscFunctionReturn(PETSC_SUCCESS);

845:   baij->ht_size = (PetscInt)(factor * nz);
846:   ht_size       = baij->ht_size;

848:   /* Allocate Memory for Hash Table */
849:   PetscCall(PetscCalloc2(ht_size, &baij->hd, ht_size, &baij->ht));
850:   HD = baij->hd;
851:   HT = baij->ht;

853:   /* Loop Over A */
854:   for (i = 0; i < a->mbs; i++) {
855:     for (j = ai[i]; j < ai[i + 1]; j++) {
856:       row = i + rstart;
857:       col = aj[j] + cstart;

859:       key = row * Nbs + col + 1;
860:       h1  = HASH(ht_size, key, tmp);
861:       for (k = 0; k < ht_size; k++) {
862:         if (!HT[(h1 + k) % ht_size]) {
863:           HT[(h1 + k) % ht_size] = key;
864:           HD[(h1 + k) % ht_size] = a->a + j * bs2;
865:           break;
866:         } else if (PetscDefined(USE_INFO)) ct++;
867:       }
868:       if (PetscDefined(USE_INFO) && k > max) max = k;
869:     }
870:   }
871:   /* Loop Over B */
872:   for (i = 0; i < b->mbs; i++) {
873:     for (j = bi[i]; j < bi[i + 1]; j++) {
874:       row = i + rstart;
875:       col = garray[bj[j]];
876:       key = row * Nbs + col + 1;
877:       h1  = HASH(ht_size, key, tmp);
878:       for (k = 0; k < ht_size; k++) {
879:         if (!HT[(h1 + k) % ht_size]) {
880:           HT[(h1 + k) % ht_size] = key;
881:           HD[(h1 + k) % ht_size] = b->a + j * bs2;
882:           break;
883:         } else if (PetscDefined(USE_INFO)) ct++;
884:       }
885:       if (PetscDefined(USE_INFO) && k > max) max = k;
886:     }
887:   }

889:   /* Print Summary */
890:   if (PetscDefined(USE_INFO)) {
891:     for (i = 0, j = 0; i < ht_size; i++) {
892:       if (HT[i]) j++;
893:     }
894:     PetscCall(PetscInfo(mat, "Average Search = %5.2g,max search = %" PetscInt_FMT "\n", (!j) ? 0.0 : (double)(((PetscReal)(ct + j)) / j), max));
895:   }
896:   PetscFunctionReturn(PETSC_SUCCESS);
897: }

899: static PetscErrorCode MatAssemblyBegin_MPIBAIJ(Mat mat, MatAssemblyType mode)
900: {
901:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)mat->data;
902:   PetscInt     nstash, reallocs;

904:   PetscFunctionBegin;
905:   if (baij->donotstash || mat->nooffprocentries) PetscFunctionReturn(PETSC_SUCCESS);

907:   PetscCall(MatStashScatterBegin_Private(mat, &mat->stash, mat->rmap->range));
908:   PetscCall(MatStashScatterBegin_Private(mat, &mat->bstash, baij->rangebs));
909:   PetscCall(MatStashGetInfo_Private(&mat->stash, &nstash, &reallocs));
910:   PetscCall(PetscInfo(mat, "Stash has %" PetscInt_FMT " entries, uses %" PetscInt_FMT " mallocs.\n", nstash, reallocs));
911:   PetscCall(MatStashGetInfo_Private(&mat->bstash, &nstash, &reallocs));
912:   PetscCall(PetscInfo(mat, "Block-Stash has %" PetscInt_FMT " entries, uses %" PetscInt_FMT " mallocs.\n", nstash, reallocs));
913:   PetscFunctionReturn(PETSC_SUCCESS);
914: }

916: static PetscErrorCode MatAssemblyEnd_MPIBAIJ(Mat mat, MatAssemblyType mode)
917: {
918:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)mat->data;
919:   Mat_SeqBAIJ *a    = (Mat_SeqBAIJ *)baij->A->data;
920:   PetscInt     i, j, rstart, ncols, flg, bs2 = baij->bs2;
921:   PetscInt    *row, *col;
922:   PetscBool    r1, r2, r3, all_assembled;
923:   MatScalar   *val;
924:   PetscMPIInt  n;

926:   PetscFunctionBegin;
927:   /* do not use 'b=(Mat_SeqBAIJ*)baij->B->data' as B can be reset in disassembly */
928:   if (!baij->donotstash && !mat->nooffprocentries) {
929:     while (1) {
930:       PetscCall(MatStashScatterGetMesg_Private(&mat->stash, &n, &row, &col, &val, &flg));
931:       if (!flg) break;

933:       for (i = 0; i < n;) {
934:         /* Now identify the consecutive vals belonging to the same row */
935:         for (j = i, rstart = row[j]; j < n; j++) {
936:           if (row[j] != rstart) break;
937:         }
938:         if (j < n) ncols = j - i;
939:         else ncols = n - i;
940:         /* Now assemble all these values with a single function call */
941:         PetscCall(MatSetValues_MPIBAIJ(mat, 1, row + i, ncols, col + i, val + i, mat->insertmode));
942:         i = j;
943:       }
944:     }
945:     PetscCall(MatStashScatterEnd_Private(&mat->stash));
946:     /* Now process the block-stash. Since the values are stashed column-oriented,
947:        set the row-oriented flag to column-oriented, and after MatSetValues()
948:        restore the original flags */
949:     r1 = baij->roworiented;
950:     r2 = a->roworiented;
951:     r3 = ((Mat_SeqBAIJ *)baij->B->data)->roworiented;

953:     baij->roworiented                           = PETSC_FALSE;
954:     a->roworiented                              = PETSC_FALSE;
955:     ((Mat_SeqBAIJ *)baij->B->data)->roworiented = PETSC_FALSE;
956:     while (1) {
957:       PetscCall(MatStashScatterGetMesg_Private(&mat->bstash, &n, &row, &col, &val, &flg));
958:       if (!flg) break;

960:       for (i = 0; i < n;) {
961:         /* Now identify the consecutive vals belonging to the same row */
962:         for (j = i, rstart = row[j]; j < n; j++) {
963:           if (row[j] != rstart) break;
964:         }
965:         if (j < n) ncols = j - i;
966:         else ncols = n - i;
967:         PetscCall(MatSetValuesBlocked_MPIBAIJ(mat, 1, row + i, ncols, col + i, val + i * bs2, mat->insertmode));
968:         i = j;
969:       }
970:     }
971:     PetscCall(MatStashScatterEnd_Private(&mat->bstash));

973:     baij->roworiented                           = r1;
974:     a->roworiented                              = r2;
975:     ((Mat_SeqBAIJ *)baij->B->data)->roworiented = r3;
976:   }

978:   PetscCall(MatAssemblyBegin(baij->A, mode));
979:   PetscCall(MatAssemblyEnd(baij->A, mode));

981:   /* determine if any process has disassembled, if so we must
982:      also disassemble ourselves, in order that we may reassemble. */
983:   /*
984:      if nonzero structure of submatrix B cannot change then we know that
985:      no process disassembled thus we can skip this stuff
986:   */
987:   if (!((Mat_SeqBAIJ *)baij->B->data)->nonew) {
988:     PetscCallMPI(MPIU_Allreduce(&mat->was_assembled, &all_assembled, 1, MPI_C_BOOL, MPI_LAND, PetscObjectComm((PetscObject)mat)));
989:     if (mat->was_assembled && !all_assembled) PetscCall(MatDisAssemble_MPIBAIJ(mat));
990:   }

992:   if (!mat->was_assembled && mode == MAT_FINAL_ASSEMBLY) PetscCall(MatSetUpMultiply_MPIBAIJ(mat));
993:   PetscCall(MatAssemblyBegin(baij->B, mode));
994:   PetscCall(MatAssemblyEnd(baij->B, mode));

996:   if (PetscDefined(USE_INFO) && baij->ht && mode == MAT_FINAL_ASSEMBLY) {
997:     PetscCall(PetscInfo(mat, "Average Hash Table Search in MatSetValues = %5.2f\n", (double)((PetscReal)baij->ht_total_ct) / baij->ht_insert_ct));

999:     baij->ht_total_ct  = 0;
1000:     baij->ht_insert_ct = 0;
1001:   }
1002:   if (baij->ht_flag && !baij->ht && mode == MAT_FINAL_ASSEMBLY) {
1003:     PetscCall(MatCreateHashTable_MPIBAIJ_Private(mat, baij->ht_fact));

1005:     mat->ops->setvalues        = MatSetValues_MPIBAIJ_HT;
1006:     mat->ops->setvaluesblocked = MatSetValuesBlocked_MPIBAIJ_HT;
1007:   }

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

1011:   baij->rowvalues = NULL;

1013:   /* if no new nonzero locations are allowed in matrix then only set the matrix state the first time through */
1014:   if ((!mat->was_assembled && mode == MAT_FINAL_ASSEMBLY) || !((Mat_SeqBAIJ *)baij->A->data)->nonew) {
1015:     mat->nonzerostate = baij->A->nonzerostate + baij->B->nonzerostate;
1016:     PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &mat->nonzerostate, 1, MPIU_INT64, MPI_SUM, PetscObjectComm((PetscObject)mat)));
1017:   }
1018:   PetscFunctionReturn(PETSC_SUCCESS);
1019: }

1021: #include <petscdraw.h>
1022: static PetscErrorCode MatView_MPIBAIJ_ASCIIorDraworSocket(Mat mat, PetscViewer viewer)
1023: {
1024:   Mat_MPIBAIJ      *baij = (Mat_MPIBAIJ *)mat->data;
1025:   PetscMPIInt       rank = baij->rank;
1026:   PetscBool         isascii, isdraw;
1027:   PetscViewer       sviewer;
1028:   PetscViewerFormat format;

1030:   PetscFunctionBegin;
1031:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
1032:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERDRAW, &isdraw));
1033:   if (isascii) {
1034:     PetscCall(PetscViewerGetFormat(viewer, &format));
1035:     if (format == PETSC_VIEWER_ASCII_INFO_DETAIL) {
1036:       MatInfo info;
1037:       PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)mat), &rank));
1038:       PetscCall(MatGetInfo(mat, MAT_LOCAL, &info));
1039:       PetscCall(PetscViewerASCIIPushSynchronized(viewer));
1040:       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,
1041:                                                    mat->rmap->bs, info.memory));
1042:       PetscCall(MatGetInfo(baij->A, MAT_LOCAL, &info));
1043:       PetscCall(PetscViewerASCIISynchronizedPrintf(viewer, "[%d] on-diagonal part: nz %" PetscInt_FMT " \n", rank, (PetscInt)info.nz_used));
1044:       PetscCall(MatGetInfo(baij->B, MAT_LOCAL, &info));
1045:       PetscCall(PetscViewerASCIISynchronizedPrintf(viewer, "[%d] off-diagonal part: nz %" PetscInt_FMT " \n", rank, (PetscInt)info.nz_used));
1046:       PetscCall(PetscViewerFlush(viewer));
1047:       PetscCall(PetscViewerASCIIPopSynchronized(viewer));
1048:       PetscCall(PetscViewerASCIIPrintf(viewer, "Information on VecScatter used in matrix-vector product: \n"));
1049:       PetscCall(VecScatterView(baij->Mvctx, viewer));
1050:       PetscFunctionReturn(PETSC_SUCCESS);
1051:     } else if (format == PETSC_VIEWER_ASCII_INFO || format == PETSC_VIEWER_ASCII_FACTOR_INFO) PetscFunctionReturn(PETSC_SUCCESS);
1052:   }

1054:   if (isdraw) {
1055:     PetscDraw draw;
1056:     PetscBool isnull;
1057:     PetscCall(PetscViewerDrawGetDraw(viewer, 0, &draw));
1058:     PetscCall(PetscDrawIsNull(draw, &isnull));
1059:     if (isnull) PetscFunctionReturn(PETSC_SUCCESS);
1060:   }

1062:   { /* assemble the entire matrix onto first process */
1063:     Mat A, Av;
1064:     IS  isrow, iscol;

1066:     PetscCall(ISCreateStride(PetscObjectComm((PetscObject)mat), rank == 0 ? mat->rmap->N : 0, 0, 1, &isrow));
1067:     PetscCall(ISCreateStride(PetscObjectComm((PetscObject)mat), rank == 0 ? mat->cmap->N : 0, 0, 1, &iscol));
1068:     PetscCall(MatCreateSubMatrix(mat, isrow, iscol, MAT_INITIAL_MATRIX, &A));
1069:     PetscCall(MatMPIBAIJGetSeqBAIJ(A, &Av, NULL, NULL));
1070:     PetscCall(ISDestroy(&isrow));
1071:     PetscCall(ISDestroy(&iscol));
1072:     /*
1073:        Everyone has to call to draw the matrix since the graphics waits are
1074:        synchronized across all processors that share the PetscDraw object
1075:     */
1076:     PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sviewer));
1077:     if (rank == 0) {
1078:       if (((PetscObject)mat)->name) PetscCall(PetscObjectSetName((PetscObject)Av, ((PetscObject)mat)->name));
1079:       PetscCall(MatView_SeqBAIJ(Av, sviewer));
1080:     }
1081:     PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sviewer));
1082:     PetscCall(MatDestroy(&A));
1083:   }
1084:   PetscFunctionReturn(PETSC_SUCCESS);
1085: }

1087: /* Used for both MPIBAIJ and MPISBAIJ matrices */
1088: PetscErrorCode MatView_MPIBAIJ_Binary(Mat mat, PetscViewer viewer)
1089: {
1090:   Mat_MPIBAIJ    *aij    = (Mat_MPIBAIJ *)mat->data;
1091:   Mat_SeqBAIJ    *A      = (Mat_SeqBAIJ *)aij->A->data;
1092:   Mat_SeqBAIJ    *B      = (Mat_SeqBAIJ *)aij->B->data;
1093:   const PetscInt *garray = aij->garray;
1094:   PetscInt        header[4], M, N, m, rs, cs, bs, cnt, i, j, ja, jb, k, l;
1095:   PetscCount      nz, hnz;
1096:   PetscInt       *rowlens, *colidxs;
1097:   PetscScalar    *matvals;
1098:   PetscMPIInt     rank;

1100:   PetscFunctionBegin;
1101:   PetscCall(PetscViewerSetUp(viewer));

1103:   M  = mat->rmap->N;
1104:   N  = mat->cmap->N;
1105:   m  = mat->rmap->n;
1106:   rs = mat->rmap->rstart;
1107:   cs = mat->cmap->rstart;
1108:   bs = mat->rmap->bs;
1109:   nz = bs * bs * (A->nz + B->nz);

1111:   /* write matrix header */
1112:   header[0] = MAT_FILE_CLASSID;
1113:   header[1] = M;
1114:   header[2] = N;
1115:   PetscCallMPI(MPI_Reduce(&nz, &hnz, 1, MPIU_COUNT, MPI_SUM, 0, PetscObjectComm((PetscObject)mat)));
1116:   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)mat), &rank));
1117:   if (rank == 0) PetscCall(PetscIntCast(hnz, &header[3]));
1118:   PetscCall(PetscViewerBinaryWrite(viewer, header, 4, PETSC_INT));

1120:   /* fill in and store row lengths */
1121:   PetscCall(PetscMalloc1(m, &rowlens));
1122:   for (cnt = 0, i = 0; i < A->mbs; i++)
1123:     for (j = 0; j < bs; j++) rowlens[cnt++] = bs * (A->i[i + 1] - A->i[i] + B->i[i + 1] - B->i[i]);
1124:   PetscCall(PetscViewerBinaryWriteAll(viewer, rowlens, m, rs, M, PETSC_INT));
1125:   PetscCall(PetscFree(rowlens));

1127:   /* fill in and store column indices */
1128:   PetscCall(PetscMalloc1(nz, &colidxs));
1129:   for (cnt = 0, i = 0; i < A->mbs; i++) {
1130:     for (k = 0; k < bs; k++) {
1131:       for (jb = B->i[i]; jb < B->i[i + 1]; jb++) {
1132:         if (garray[B->j[jb]] > cs / bs) break;
1133:         for (l = 0; l < bs; l++) colidxs[cnt++] = bs * garray[B->j[jb]] + l;
1134:       }
1135:       for (ja = A->i[i]; ja < A->i[i + 1]; ja++)
1136:         for (l = 0; l < bs; l++) colidxs[cnt++] = bs * A->j[ja] + l + cs;
1137:       for (; jb < B->i[i + 1]; jb++)
1138:         for (l = 0; l < bs; l++) colidxs[cnt++] = bs * garray[B->j[jb]] + l;
1139:     }
1140:   }
1141:   PetscCheck(cnt == nz, PETSC_COMM_SELF, PETSC_ERR_LIB, "Internal PETSc error: cnt = %" PetscInt_FMT " nz = %" PetscCount_FMT, cnt, nz);
1142:   PetscCall(PetscViewerBinaryWriteAll(viewer, colidxs, nz, PETSC_DECIDE, PETSC_DECIDE, PETSC_INT));
1143:   PetscCall(PetscFree(colidxs));

1145:   /* fill in and store nonzero values */
1146:   PetscCall(PetscMalloc1(nz, &matvals));
1147:   for (cnt = 0, i = 0; i < A->mbs; i++) {
1148:     for (k = 0; k < bs; k++) {
1149:       for (jb = B->i[i]; jb < B->i[i + 1]; jb++) {
1150:         if (garray[B->j[jb]] > cs / bs) break;
1151:         for (l = 0; l < bs; l++) matvals[cnt++] = B->a[bs * (bs * jb + l) + k];
1152:       }
1153:       for (ja = A->i[i]; ja < A->i[i + 1]; ja++)
1154:         for (l = 0; l < bs; l++) matvals[cnt++] = A->a[bs * (bs * ja + l) + k];
1155:       for (; jb < B->i[i + 1]; jb++)
1156:         for (l = 0; l < bs; l++) matvals[cnt++] = B->a[bs * (bs * jb + l) + k];
1157:     }
1158:   }
1159:   PetscCall(PetscViewerBinaryWriteAll(viewer, matvals, nz, PETSC_DECIDE, PETSC_DECIDE, PETSC_SCALAR));
1160:   PetscCall(PetscFree(matvals));

1162:   /* write block size option to the viewer's .info file */
1163:   PetscCall(MatView_Binary_BlockSizes(mat, viewer));
1164:   PetscFunctionReturn(PETSC_SUCCESS);
1165: }

1167: PetscErrorCode MatView_MPIBAIJ(Mat mat, PetscViewer viewer)
1168: {
1169:   PetscBool isascii, isdraw, issocket, isbinary;

1171:   PetscFunctionBegin;
1172:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
1173:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERDRAW, &isdraw));
1174:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERSOCKET, &issocket));
1175:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &isbinary));
1176:   if (isascii || isdraw || issocket) PetscCall(MatView_MPIBAIJ_ASCIIorDraworSocket(mat, viewer));
1177:   else if (isbinary) PetscCall(MatView_MPIBAIJ_Binary(mat, viewer));
1178:   PetscFunctionReturn(PETSC_SUCCESS);
1179: }

1181: static PetscErrorCode MatMult_MPIBAIJ(Mat A, Vec xx, Vec yy)
1182: {
1183:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;
1184:   PetscInt     nt;

1186:   PetscFunctionBegin;
1187:   PetscCall(VecGetLocalSize(xx, &nt));
1188:   PetscCheck(nt == A->cmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Incompatible partition of A and xx");
1189:   PetscCall(VecGetLocalSize(yy, &nt));
1190:   PetscCheck(nt == A->rmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Incompatible partition of A and yy");
1191:   PetscCall(VecScatterBegin(a->Mvctx, xx, a->lvec, INSERT_VALUES, SCATTER_FORWARD));
1192:   PetscUseTypeMethod(a->A, mult, xx, yy);
1193:   PetscCall(VecScatterEnd(a->Mvctx, xx, a->lvec, INSERT_VALUES, SCATTER_FORWARD));
1194:   PetscUseTypeMethod(a->B, multadd, a->lvec, yy, yy);
1195:   PetscFunctionReturn(PETSC_SUCCESS);
1196: }

1198: static PetscErrorCode MatMultAdd_MPIBAIJ(Mat A, Vec xx, Vec yy, Vec zz)
1199: {
1200:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

1202:   PetscFunctionBegin;
1203:   PetscCall(VecScatterBegin(a->Mvctx, xx, a->lvec, INSERT_VALUES, SCATTER_FORWARD));
1204:   PetscUseTypeMethod(a->A, multadd, xx, yy, zz);
1205:   PetscCall(VecScatterEnd(a->Mvctx, xx, a->lvec, INSERT_VALUES, SCATTER_FORWARD));
1206:   PetscUseTypeMethod(a->B, multadd, a->lvec, zz, zz);
1207:   PetscFunctionReturn(PETSC_SUCCESS);
1208: }

1210: static PetscErrorCode MatMultTranspose_MPIBAIJ(Mat A, Vec xx, Vec yy)
1211: {
1212:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

1214:   PetscFunctionBegin;
1215:   /* do nondiagonal part */
1216:   PetscUseTypeMethod(a->B, multtranspose, xx, a->lvec);
1217:   /* do local part */
1218:   PetscUseTypeMethod(a->A, multtranspose, xx, yy);
1219:   /* add partial results together */
1220:   PetscCall(VecScatterBegin(a->Mvctx, a->lvec, yy, ADD_VALUES, SCATTER_REVERSE));
1221:   PetscCall(VecScatterEnd(a->Mvctx, a->lvec, yy, ADD_VALUES, SCATTER_REVERSE));
1222:   PetscFunctionReturn(PETSC_SUCCESS);
1223: }

1225: static PetscErrorCode MatMultTransposeAdd_MPIBAIJ(Mat A, Vec xx, Vec yy, Vec zz)
1226: {
1227:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

1229:   PetscFunctionBegin;
1230:   /* do nondiagonal part */
1231:   PetscUseTypeMethod(a->B, multtranspose, xx, a->lvec);
1232:   /* do local part */
1233:   PetscUseTypeMethod(a->A, multtransposeadd, xx, yy, zz);
1234:   /* add partial results together */
1235:   PetscCall(VecScatterBegin(a->Mvctx, a->lvec, zz, ADD_VALUES, SCATTER_REVERSE));
1236:   PetscCall(VecScatterEnd(a->Mvctx, a->lvec, zz, ADD_VALUES, SCATTER_REVERSE));
1237:   PetscFunctionReturn(PETSC_SUCCESS);
1238: }

1240: /*
1241:   This only works correctly for square matrices where the subblock A->A is the
1242:    diagonal block
1243: */
1244: static PetscErrorCode MatGetDiagonal_MPIBAIJ(Mat A, Vec v)
1245: {
1246:   PetscFunctionBegin;
1247:   PetscCheck(A->rmap->N == A->cmap->N, PETSC_COMM_SELF, PETSC_ERR_SUP, "Supports only square matrix where A->A is diag block");
1248:   PetscCall(MatGetDiagonal(((Mat_MPIBAIJ *)A->data)->A, v));
1249:   PetscFunctionReturn(PETSC_SUCCESS);
1250: }

1252: static PetscErrorCode MatScale_MPIBAIJ(Mat A, PetscScalar aa)
1253: {
1254:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

1256:   PetscFunctionBegin;
1257:   PetscCall(MatScale(a->A, aa));
1258:   PetscCall(MatScale(a->B, aa));
1259:   PetscFunctionReturn(PETSC_SUCCESS);
1260: }

1262: static PetscErrorCode MatGetRow_MPIBAIJ(Mat matin, PetscInt row, PetscInt *nz, PetscInt **idx, PetscScalar **v)
1263: {
1264:   Mat_MPIBAIJ *mat = (Mat_MPIBAIJ *)matin->data;
1265:   PetscScalar *vworkA, *vworkB, **pvA, **pvB, *v_p;
1266:   PetscInt     bs = matin->rmap->bs, bs2 = mat->bs2, i, *cworkA, *cworkB, **pcA, **pcB;
1267:   PetscInt     nztot, nzA, nzB, lrow, brstart = matin->rmap->rstart, brend = matin->rmap->rend;
1268:   PetscInt    *cmap, *idx_p, cstart = mat->cstartbs;

1270:   PetscFunctionBegin;
1271:   PetscCheck(row >= brstart && row < brend, PETSC_COMM_SELF, PETSC_ERR_SUP, "Only local rows");
1272:   PetscCheck(!mat->getrowactive, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Already active");
1273:   mat->getrowactive = PETSC_TRUE;

1275:   if (!mat->rowvalues && (idx || v)) {
1276:     /*
1277:         allocate enough space to hold information from the longest row.
1278:     */
1279:     Mat_SeqBAIJ *Aa = (Mat_SeqBAIJ *)mat->A->data, *Ba = (Mat_SeqBAIJ *)mat->B->data;
1280:     PetscInt     max = 1, mbs = mat->mbs, tmp;
1281:     for (i = 0; i < mbs; i++) {
1282:       tmp = Aa->i[i + 1] - Aa->i[i] + Ba->i[i + 1] - Ba->i[i];
1283:       if (max < tmp) max = tmp;
1284:     }
1285:     PetscCall(PetscMalloc2(max * bs2, &mat->rowvalues, max * bs2, &mat->rowindices));
1286:   }
1287:   lrow = row - brstart;

1289:   pvA = &vworkA;
1290:   pcA = &cworkA;
1291:   pvB = &vworkB;
1292:   pcB = &cworkB;
1293:   if (!v) {
1294:     pvA = NULL;
1295:     pvB = NULL;
1296:   }
1297:   if (!idx) {
1298:     pcA = NULL;
1299:     if (!v) pcB = NULL;
1300:   }
1301:   PetscUseTypeMethod(mat->A, getrow, lrow, &nzA, pcA, pvA);
1302:   PetscUseTypeMethod(mat->B, getrow, lrow, &nzB, pcB, pvB);
1303:   nztot = nzA + nzB;

1305:   cmap = mat->garray;
1306:   if (v || idx) {
1307:     if (nztot) {
1308:       /* Sort by increasing column numbers, assuming A and B already sorted */
1309:       PetscInt imark = -1;
1310:       if (v) {
1311:         *v = v_p = mat->rowvalues;
1312:         for (i = 0; i < nzB; i++) {
1313:           if (cmap[cworkB[i] / bs] < cstart) v_p[i] = vworkB[i];
1314:           else break;
1315:         }
1316:         imark = i;
1317:         for (i = 0; i < nzA; i++) v_p[imark + i] = vworkA[i];
1318:         for (i = imark; i < nzB; i++) v_p[nzA + i] = vworkB[i];
1319:       }
1320:       if (idx) {
1321:         *idx = idx_p = mat->rowindices;
1322:         if (imark > -1) {
1323:           for (i = 0; i < imark; i++) idx_p[i] = cmap[cworkB[i] / bs] * bs + cworkB[i] % bs;
1324:         } else {
1325:           for (i = 0; i < nzB; i++) {
1326:             if (cmap[cworkB[i] / bs] < cstart) idx_p[i] = cmap[cworkB[i] / bs] * bs + cworkB[i] % bs;
1327:             else break;
1328:           }
1329:           imark = i;
1330:         }
1331:         for (i = 0; i < nzA; i++) idx_p[imark + i] = cstart * bs + cworkA[i];
1332:         for (i = imark; i < nzB; i++) idx_p[nzA + i] = cmap[cworkB[i] / bs] * bs + cworkB[i] % bs;
1333:       }
1334:     } else {
1335:       if (idx) *idx = NULL;
1336:       if (v) *v = NULL;
1337:     }
1338:   }
1339:   *nz = nztot;
1340:   PetscUseTypeMethod(mat->A, restorerow, lrow, &nzA, pcA, pvA);
1341:   PetscUseTypeMethod(mat->B, restorerow, lrow, &nzB, pcB, pvB);
1342:   PetscFunctionReturn(PETSC_SUCCESS);
1343: }

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

1349:   PetscFunctionBegin;
1350:   PetscCheck(baij->getrowactive, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "MatGetRow not called");
1351:   baij->getrowactive = PETSC_FALSE;
1352:   PetscFunctionReturn(PETSC_SUCCESS);
1353: }

1355: static PetscErrorCode MatZeroEntries_MPIBAIJ(Mat A)
1356: {
1357:   Mat_MPIBAIJ *l = (Mat_MPIBAIJ *)A->data;

1359:   PetscFunctionBegin;
1360:   PetscCall(MatZeroEntries(l->A));
1361:   PetscCall(MatZeroEntries(l->B));
1362:   PetscFunctionReturn(PETSC_SUCCESS);
1363: }

1365: static PetscErrorCode MatGetInfo_MPIBAIJ(Mat matin, MatInfoType flag, MatInfo *info)
1366: {
1367:   Mat_MPIBAIJ   *a = (Mat_MPIBAIJ *)matin->data;
1368:   Mat            A = a->A, B = a->B;
1369:   PetscLogDouble irecv[5];

1371:   PetscFunctionBegin;
1372:   info->block_size = (PetscReal)matin->rmap->bs;

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

1376:   irecv[0] = info->nz_used;
1377:   irecv[1] = info->nz_allocated;
1378:   irecv[2] = info->nz_unneeded;
1379:   irecv[3] = info->memory;
1380:   irecv[4] = info->mallocs;

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

1384:   irecv[0] += info->nz_used;
1385:   irecv[1] += info->nz_allocated;
1386:   irecv[2] += info->nz_unneeded;
1387:   irecv[3] += info->memory;
1388:   irecv[4] += info->mallocs;

1390:   if (flag == MAT_LOCAL) {
1391:     info->nz_used      = irecv[0];
1392:     info->nz_allocated = irecv[1];
1393:     info->nz_unneeded  = irecv[2];
1394:     info->memory       = irecv[3];
1395:     info->mallocs      = irecv[4];
1396:   } else if (flag == MAT_GLOBAL_MAX) {
1397:     PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, irecv, 5, MPIU_PETSCLOGDOUBLE, MPI_MAX, PetscObjectComm((PetscObject)matin)));

1399:     info->nz_used      = irecv[0];
1400:     info->nz_allocated = irecv[1];
1401:     info->nz_unneeded  = irecv[2];
1402:     info->memory       = irecv[3];
1403:     info->mallocs      = irecv[4];
1404:   } else if (flag == MAT_GLOBAL_SUM) {
1405:     PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, irecv, 5, MPIU_PETSCLOGDOUBLE, MPI_SUM, PetscObjectComm((PetscObject)matin)));

1407:     info->nz_used      = irecv[0];
1408:     info->nz_allocated = irecv[1];
1409:     info->nz_unneeded  = irecv[2];
1410:     info->memory       = irecv[3];
1411:     info->mallocs      = irecv[4];
1412:   } else SETERRQ(PetscObjectComm((PetscObject)matin), PETSC_ERR_ARG_WRONG, "Unknown MatInfoType argument %d", (int)flag);
1413:   info->fill_ratio_given  = 0; /* no parallel LU/ILU/Cholesky */
1414:   info->fill_ratio_needed = 0;
1415:   info->factor_mallocs    = 0;
1416:   PetscFunctionReturn(PETSC_SUCCESS);
1417: }

1419: static PetscErrorCode MatSetOption_MPIBAIJ(Mat A, MatOption op, PetscBool flg)
1420: {
1421:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

1423:   PetscFunctionBegin;
1424:   switch (op) {
1425:   case MAT_NEW_NONZERO_LOCATIONS:
1426:   case MAT_NEW_NONZERO_ALLOCATION_ERR:
1427:   case MAT_UNUSED_NONZERO_LOCATION_ERR:
1428:   case MAT_KEEP_NONZERO_PATTERN:
1429:   case MAT_NEW_NONZERO_LOCATION_ERR:
1430:   case MAT_ROW_ORIENTED:
1431:     MatCheckPreallocated(A, 1);
1432:     if (op == MAT_ROW_ORIENTED) a->roworiented = flg;
1433:     PetscCall(MatSetOption(a->A, op, flg));
1434:     PetscCall(MatSetOption(a->B, op, flg));
1435:     break;
1436:   case MAT_STRUCTURE_ONLY:
1437:     if (a->A) PetscCall(MatSetOption(a->A, op, flg));
1438:     if (a->B) PetscCall(MatSetOption(a->B, op, flg));
1439:     break;
1440:   case MAT_IGNORE_OFF_PROC_ENTRIES:
1441:     a->donotstash = flg;
1442:     break;
1443:   case MAT_USE_HASH_TABLE:
1444:     a->ht_flag = flg;
1445:     a->ht_fact = 1.39;
1446:     break;
1447:   case MAT_SPD:
1448:   case MAT_SYMMETRIC:
1449:   case MAT_STRUCTURALLY_SYMMETRIC:
1450:   case MAT_HERMITIAN:
1451:   case MAT_SYMMETRY_ETERNAL:
1452:   case MAT_STRUCTURAL_SYMMETRY_ETERNAL:
1453:   case MAT_SPD_ETERNAL:
1454:     /* if the diagonal matrix is square it inherits some of the properties above */
1455:     if (a->A && A->rmap->n == A->cmap->n) PetscCall(MatSetOption(a->A, op, flg));
1456:     break;
1457:   default:
1458:     break;
1459:   }
1460:   PetscFunctionReturn(PETSC_SUCCESS);
1461: }

1463: static PetscErrorCode MatTranspose_MPIBAIJ(Mat A, MatReuse reuse, Mat *matout)
1464: {
1465:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)A->data;
1466:   Mat_SeqBAIJ *Aloc;
1467:   Mat          B;
1468:   PetscInt     M = A->rmap->N, N = A->cmap->N, *ai, *aj, i, *rvals, j, k, col;
1469:   PetscInt     bs = A->rmap->bs, mbs = baij->mbs;
1470:   MatScalar   *a;

1472:   PetscFunctionBegin;
1473:   if (reuse == MAT_REUSE_MATRIX) PetscCall(MatTransposeCheckNonzeroState_Private(A, *matout));
1474:   if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_INPLACE_MATRIX) {
1475:     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
1476:     PetscCall(MatSetSizes(B, A->cmap->n, A->rmap->n, N, M));
1477:     PetscCall(MatSetType(B, ((PetscObject)A)->type_name));
1478:     /* Do not know preallocation information, but must set block size */
1479:     PetscCall(MatMPIBAIJSetPreallocation(B, A->rmap->bs, PETSC_DECIDE, NULL, PETSC_DECIDE, NULL));
1480:   } else {
1481:     B = *matout;
1482:   }

1484:   /* copy over the A part */
1485:   Aloc = (Mat_SeqBAIJ *)baij->A->data;
1486:   ai   = Aloc->i;
1487:   aj   = Aloc->j;
1488:   a    = Aloc->a;
1489:   PetscCall(PetscMalloc1(bs, &rvals));

1491:   for (i = 0; i < mbs; i++) {
1492:     rvals[0] = bs * (baij->rstartbs + i);
1493:     for (j = 1; j < bs; j++) rvals[j] = rvals[j - 1] + 1;
1494:     for (j = ai[i]; j < ai[i + 1]; j++) {
1495:       col = (baij->cstartbs + aj[j]) * bs;
1496:       for (k = 0; k < bs; k++) {
1497:         PetscCall(MatSetValues_MPIBAIJ(B, 1, &col, bs, rvals, a, INSERT_VALUES));

1499:         col++;
1500:         a += bs;
1501:       }
1502:     }
1503:   }
1504:   /* copy over the B part */
1505:   Aloc = (Mat_SeqBAIJ *)baij->B->data;
1506:   ai   = Aloc->i;
1507:   aj   = Aloc->j;
1508:   a    = Aloc->a;
1509:   for (i = 0; i < mbs; i++) {
1510:     rvals[0] = bs * (baij->rstartbs + i);
1511:     for (j = 1; j < bs; j++) rvals[j] = rvals[j - 1] + 1;
1512:     for (j = ai[i]; j < ai[i + 1]; j++) {
1513:       col = baij->garray[aj[j]] * bs;
1514:       for (k = 0; k < bs; k++) {
1515:         PetscCall(MatSetValues_MPIBAIJ(B, 1, &col, bs, rvals, a, INSERT_VALUES));
1516:         col++;
1517:         a += bs;
1518:       }
1519:     }
1520:   }
1521:   PetscCall(PetscFree(rvals));
1522:   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
1523:   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));

1525:   if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_REUSE_MATRIX) *matout = B;
1526:   else PetscCall(MatHeaderMerge(A, &B));
1527:   PetscFunctionReturn(PETSC_SUCCESS);
1528: }

1530: static PetscErrorCode MatDiagonalScale_MPIBAIJ(Mat mat, Vec ll, Vec rr)
1531: {
1532:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)mat->data;
1533:   Mat          a = baij->A, b = baij->B;
1534:   PetscInt     s1, s2, s3;

1536:   PetscFunctionBegin;
1537:   PetscCall(MatGetLocalSize(mat, &s2, &s3));
1538:   if (rr) {
1539:     PetscCall(VecGetLocalSize(rr, &s1));
1540:     PetscCheck(s1 == s3, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "right vector non-conforming local size");
1541:     /* Overlap communication with computation. */
1542:     PetscCall(VecScatterBegin(baij->Mvctx, rr, baij->lvec, INSERT_VALUES, SCATTER_FORWARD));
1543:   }
1544:   if (ll) {
1545:     PetscCall(VecGetLocalSize(ll, &s1));
1546:     PetscCheck(s1 == s2, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "left vector non-conforming local size");
1547:     PetscUseTypeMethod(b, diagonalscale, ll, NULL);
1548:   }
1549:   /* scale  the diagonal block */
1550:   PetscUseTypeMethod(a, diagonalscale, ll, rr);

1552:   if (rr) {
1553:     /* Do a scatter end and then right scale the off-diagonal block */
1554:     PetscCall(VecScatterEnd(baij->Mvctx, rr, baij->lvec, INSERT_VALUES, SCATTER_FORWARD));
1555:     PetscUseTypeMethod(b, diagonalscale, NULL, baij->lvec);
1556:   }
1557:   /* MatDiagonalScale() cannot be used on the blocks: they are on PETSC_COMM_SELF while ll and rr
1558:      are parallel, so the interface's communicator check rejects them. Advance the block states
1559:      here instead, as the interface would. MatDiagonalScale_MPIAIJ() does not need this because
1560:      MatSeqAIJRestoreArray() advances the state for it. */
1561:   PetscCall(PetscObjectStateIncrease((PetscObject)a));
1562:   PetscCall(PetscObjectStateIncrease((PetscObject)b));
1563:   PetscFunctionReturn(PETSC_SUCCESS);
1564: }

1566: static PetscErrorCode MatZeroRows_MPIBAIJ(Mat A, PetscInt N, const PetscInt rows[], PetscScalar diag, Vec x, Vec b)
1567: {
1568:   Mat_MPIBAIJ *l = (Mat_MPIBAIJ *)A->data;
1569:   PetscInt    *lrows;
1570:   PetscInt     r, len;
1571:   PetscBool    cong;

1573:   PetscFunctionBegin;
1574:   /* get locally owned rows */
1575:   PetscCall(MatZeroRowsMapLocal_Private(A, N, rows, &len, &lrows));
1576:   /* fix right-hand side if needed */
1577:   if (x && b) {
1578:     const PetscScalar *xx;
1579:     PetscScalar       *bb;

1581:     PetscCall(VecGetArrayRead(x, &xx));
1582:     PetscCall(VecGetArray(b, &bb));
1583:     for (r = 0; r < len; ++r) bb[lrows[r]] = diag * xx[lrows[r]];
1584:     PetscCall(VecRestoreArrayRead(x, &xx));
1585:     PetscCall(VecRestoreArray(b, &bb));
1586:   }

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

1595:   */
1596:   /* must zero l->B before l->A because the (diag) case below may put values into l->B*/
1597:   PetscCall(MatZeroRows_SeqBAIJ(l->B, len, lrows, 0.0, NULL, NULL));
1598:   PetscCall(MatHasCongruentLayouts(A, &cong));
1599:   if ((diag != 0.0) && cong) {
1600:     PetscCall(MatZeroRows_SeqBAIJ(l->A, len, lrows, diag, NULL, NULL));
1601:   } else if (diag != 0.0) {
1602:     PetscCall(MatZeroRows_SeqBAIJ(l->A, len, lrows, 0.0, NULL, NULL));
1603:     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");
1604:     for (r = 0; r < len; ++r) {
1605:       const PetscInt row = lrows[r] + A->rmap->rstart;
1606:       PetscCall(MatSetValues(A, 1, &row, 1, &row, &diag, INSERT_VALUES));
1607:     }
1608:     PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
1609:     PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
1610:   } else {
1611:     PetscCall(MatZeroRows_SeqBAIJ(l->A, len, lrows, 0.0, NULL, NULL));
1612:   }
1613:   /* MatZeroRows() cannot be used on the blocks: it honors -mat_view, which would print each
1614:      sequential block as well (see mat_tests-ex12_5). Advance the diagonal block's state here
1615:      instead, as the interface would; MatInvertBlockDiagonal_MPIBAIJ() caches on that state. */
1616:   PetscCall(PetscObjectStateIncrease((PetscObject)l->A));
1617:   PetscCall(PetscFree(lrows));

1619:   /* only change matrix nonzero state if pattern was allowed to be changed */
1620:   if (!((Mat_SeqBAIJ *)l->A->data)->keepnonzeropattern || !((Mat_SeqBAIJ *)l->A->data)->nonew) {
1621:     A->nonzerostate = l->A->nonzerostate + l->B->nonzerostate;
1622:     PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &A->nonzerostate, 1, MPIU_INT64, MPI_SUM, PetscObjectComm((PetscObject)A)));
1623:   }
1624:   PetscFunctionReturn(PETSC_SUCCESS);
1625: }

1627: static PetscErrorCode MatZeroRowsColumns_MPIBAIJ(Mat A, PetscInt N, const PetscInt rows[], PetscScalar diag, Vec x, Vec b)
1628: {
1629:   Mat_MPIBAIJ       *l = (Mat_MPIBAIJ *)A->data;
1630:   PetscMPIInt        n, p = 0;
1631:   PetscInt           i, j, k, r, len = 0, row, col, count;
1632:   PetscInt          *lrows, *owners = A->rmap->range;
1633:   PetscSFNode       *rrows;
1634:   PetscSF            sf;
1635:   const PetscScalar *xx;
1636:   PetscScalar       *bb, *mask;
1637:   Vec                xmask, lmask;
1638:   Mat_SeqBAIJ       *baij = (Mat_SeqBAIJ *)l->B->data;
1639:   PetscInt           bs = A->rmap->bs, bs2 = baij->bs2;
1640:   PetscScalar       *aa;

1642:   PetscFunctionBegin;
1643:   PetscCall(PetscMPIIntCast(A->rmap->n, &n));
1644:   /* create PetscSF where leaves are input rows and roots are owned rows */
1645:   PetscCall(PetscMalloc1(n, &lrows));
1646:   for (r = 0; r < n; ++r) lrows[r] = -1;
1647:   PetscCall(PetscMalloc1(N, &rrows));
1648:   for (r = 0; r < N; ++r) {
1649:     const PetscInt idx = rows[r];
1650:     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);
1651:     if (idx < owners[p] || owners[p + 1] <= idx) { /* short-circuit the search if the last p owns this row too */
1652:       PetscCall(PetscLayoutFindOwner(A->rmap, idx, &p));
1653:     }
1654:     rrows[r].rank  = p;
1655:     rrows[r].index = rows[r] - owners[p];
1656:   }
1657:   PetscCall(PetscSFCreate(PetscObjectComm((PetscObject)A), &sf));
1658:   PetscCall(PetscSFSetGraph(sf, n, N, NULL, PETSC_OWN_POINTER, rrows, PETSC_OWN_POINTER));
1659:   /* collect flags for rows to be zeroed */
1660:   PetscCall(PetscSFReduceBegin(sf, MPIU_INT, (PetscInt *)rows, lrows, MPI_LOR));
1661:   PetscCall(PetscSFReduceEnd(sf, MPIU_INT, (PetscInt *)rows, lrows, MPI_LOR));
1662:   PetscCall(PetscSFDestroy(&sf));
1663:   /* compress and put in row numbers */
1664:   for (r = 0; r < n; ++r)
1665:     if (lrows[r] >= 0) lrows[len++] = r;
1666:   /* zero diagonal part of matrix */
1667:   PetscCall(MatZeroRowsColumns(l->A, len, lrows, diag, x, b));
1668:   /* handle off-diagonal part of matrix */
1669:   PetscCall(MatCreateVecs(A, &xmask, NULL));
1670:   PetscCall(VecDuplicate(l->lvec, &lmask));
1671:   PetscCall(VecGetArray(xmask, &bb));
1672:   for (i = 0; i < len; i++) bb[lrows[i]] = 1;
1673:   PetscCall(VecRestoreArray(xmask, &bb));
1674:   PetscCall(VecScatterBegin(l->Mvctx, xmask, lmask, ADD_VALUES, SCATTER_FORWARD));
1675:   PetscCall(VecScatterEnd(l->Mvctx, xmask, lmask, ADD_VALUES, SCATTER_FORWARD));
1676:   PetscCall(VecDestroy(&xmask));
1677:   if (x) {
1678:     PetscCall(VecScatterBegin(l->Mvctx, x, l->lvec, INSERT_VALUES, SCATTER_FORWARD));
1679:     PetscCall(VecScatterEnd(l->Mvctx, x, l->lvec, INSERT_VALUES, SCATTER_FORWARD));
1680:     PetscCall(VecGetArrayRead(l->lvec, &xx));
1681:     PetscCall(VecGetArray(b, &bb));
1682:   }
1683:   PetscCall(VecGetArray(lmask, &mask));
1684:   /* remove zeroed rows of off-diagonal matrix */
1685:   for (i = 0; i < len; ++i) {
1686:     row   = lrows[i];
1687:     count = (baij->i[row / bs + 1] - baij->i[row / bs]) * bs;
1688:     aa    = PetscSafePointerPlusOffset(baij->a, baij->i[row / bs] * bs2 + (row % bs));
1689:     for (k = 0; k < count; ++k) {
1690:       aa[0] = 0.0;
1691:       aa += bs;
1692:     }
1693:   }
1694:   /* loop over all elements of off process part of matrix zeroing removed columns */
1695:   for (i = 0; i < l->B->rmap->N; ++i) {
1696:     row = i / bs;
1697:     for (j = baij->i[row]; j < baij->i[row + 1]; ++j) {
1698:       for (k = 0; k < bs; ++k) {
1699:         col = bs * baij->j[j] + k;
1700:         if (PetscAbsScalar(mask[col])) {
1701:           aa = baij->a + j * bs2 + (i % bs) + bs * k;
1702:           if (x) bb[i] -= aa[0] * xx[col];
1703:           aa[0] = 0.0;
1704:         }
1705:       }
1706:     }
1707:   }
1708:   if (x) {
1709:     PetscCall(VecRestoreArray(b, &bb));
1710:     PetscCall(VecRestoreArrayRead(l->lvec, &xx));
1711:   }
1712:   PetscCall(VecRestoreArray(lmask, &mask));
1713:   PetscCall(VecDestroy(&lmask));
1714:   PetscCall(PetscFree(lrows));

1716:   /* only change matrix nonzero state if pattern was allowed to be changed */
1717:   if (!((Mat_SeqBAIJ *)l->A->data)->nonew) {
1718:     A->nonzerostate = l->A->nonzerostate + l->B->nonzerostate;
1719:     PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &A->nonzerostate, 1, MPIU_INT64, MPI_SUM, PetscObjectComm((PetscObject)A)));
1720:   }
1721:   PetscFunctionReturn(PETSC_SUCCESS);
1722: }

1724: static PetscErrorCode MatSetUnfactored_MPIBAIJ(Mat A)
1725: {
1726:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

1728:   PetscFunctionBegin;
1729:   PetscCall(MatSetUnfactored(a->A));
1730:   PetscFunctionReturn(PETSC_SUCCESS);
1731: }

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

1735: static PetscErrorCode MatEqual_MPIBAIJ(Mat A, Mat B, PetscBool *flag)
1736: {
1737:   Mat_MPIBAIJ *matB = (Mat_MPIBAIJ *)B->data, *matA = (Mat_MPIBAIJ *)A->data;
1738:   Mat          a, b, c, d;

1740:   PetscFunctionBegin;
1741:   a = matA->A;
1742:   b = matA->B;
1743:   c = matB->A;
1744:   d = matB->B;

1746:   PetscCall(MatEqual(a, c, flag));
1747:   if (*flag) PetscCall(MatEqual(b, d, flag));
1748:   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, flag, 1, MPI_C_BOOL, MPI_LAND, PetscObjectComm((PetscObject)A)));
1749:   PetscFunctionReturn(PETSC_SUCCESS);
1750: }

1752: static PetscErrorCode MatCopy_MPIBAIJ(Mat A, Mat B, MatStructure str)
1753: {
1754:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;
1755:   Mat_MPIBAIJ *b = (Mat_MPIBAIJ *)B->data;

1757:   PetscFunctionBegin;
1758:   /* If the two matrices don't have the same copy implementation, they aren't compatible for fast copy. */
1759:   if (str != SAME_NONZERO_PATTERN || A->ops->copy != B->ops->copy) {
1760:     PetscCall(MatCopy_Basic(A, B, str));
1761:   } else {
1762:     PetscCall(MatCopy(a->A, b->A, str));
1763:     PetscCall(MatCopy(a->B, b->B, str));
1764:   }
1765:   PetscCall(PetscObjectStateIncrease((PetscObject)B));
1766:   PetscFunctionReturn(PETSC_SUCCESS);
1767: }

1769: PetscErrorCode MatAXPYGetPreallocation_MPIBAIJ(Mat Y, const PetscInt *yltog, Mat X, const PetscInt *xltog, PetscInt *nnz)
1770: {
1771:   PetscInt     bs = Y->rmap->bs, m = Y->rmap->N / bs;
1772:   Mat_SeqBAIJ *x = (Mat_SeqBAIJ *)X->data;
1773:   Mat_SeqBAIJ *y = (Mat_SeqBAIJ *)Y->data;

1775:   PetscFunctionBegin;
1776:   PetscCall(MatAXPYGetPreallocation_MPIX_private(m, x->i, x->j, xltog, y->i, y->j, yltog, nnz));
1777:   PetscFunctionReturn(PETSC_SUCCESS);
1778: }

1780: static PetscErrorCode MatAXPY_MPIBAIJ(Mat Y, PetscScalar a, Mat X, MatStructure str)
1781: {
1782:   Mat_MPIBAIJ *xx = (Mat_MPIBAIJ *)X->data, *yy = (Mat_MPIBAIJ *)Y->data;
1783:   PetscBLASInt bnz, one                         = 1;
1784:   Mat_SeqBAIJ *x, *y;
1785:   PetscInt     bs2 = Y->rmap->bs * Y->rmap->bs;

1787:   PetscFunctionBegin;
1788:   if (str == SAME_NONZERO_PATTERN) {
1789:     PetscScalar alpha = a;
1790:     x                 = (Mat_SeqBAIJ *)xx->A->data;
1791:     y                 = (Mat_SeqBAIJ *)yy->A->data;
1792:     PetscCall(PetscBLASIntCast(x->nz * bs2, &bnz));
1793:     PetscCallBLAS("BLASaxpy", BLASaxpy_(&bnz, &alpha, x->a, &one, y->a, &one));
1794:     x = (Mat_SeqBAIJ *)xx->B->data;
1795:     y = (Mat_SeqBAIJ *)yy->B->data;
1796:     PetscCall(PetscBLASIntCast(x->nz * bs2, &bnz));
1797:     PetscCallBLAS("BLASaxpy", BLASaxpy_(&bnz, &alpha, x->a, &one, y->a, &one));
1798:     /* the blocks' values were changed directly, so advance their states as MatAXPY() on each
1799:        block would; MatInvertBlockDiagonal_SeqBAIJ() caches on the diagonal block's state */
1800:     PetscCall(PetscObjectStateIncrease((PetscObject)yy->A));
1801:     PetscCall(PetscObjectStateIncrease((PetscObject)yy->B));
1802:     PetscCall(PetscObjectStateIncrease((PetscObject)Y));
1803:   } else if (str == SUBSET_NONZERO_PATTERN) { /* nonzeros of X is a subset of Y's */
1804:     PetscCall(MatAXPY_Basic(Y, a, X, str));
1805:   } else {
1806:     Mat       B;
1807:     PetscInt *nnz_d, *nnz_o, bs = Y->rmap->bs;
1808:     PetscCall(PetscMalloc1(yy->A->rmap->N, &nnz_d));
1809:     PetscCall(PetscMalloc1(yy->B->rmap->N, &nnz_o));
1810:     PetscCall(MatCreate(PetscObjectComm((PetscObject)Y), &B));
1811:     PetscCall(PetscObjectSetName((PetscObject)B, ((PetscObject)Y)->name));
1812:     PetscCall(MatSetSizes(B, Y->rmap->n, Y->cmap->n, Y->rmap->N, Y->cmap->N));
1813:     PetscCall(MatSetBlockSizesFromMats(B, Y, Y));
1814:     PetscCall(MatSetType(B, MATMPIBAIJ));
1815:     PetscCall(MatAXPYGetPreallocation_SeqBAIJ(yy->A, xx->A, nnz_d));
1816:     PetscCall(MatAXPYGetPreallocation_MPIBAIJ(yy->B, yy->garray, xx->B, xx->garray, nnz_o));
1817:     PetscCall(MatMPIBAIJSetPreallocation(B, bs, 0, nnz_d, 0, nnz_o));
1818:     /* MatAXPY_BasicWithPreallocation() for BAIJ matrix is much slower than AIJ, even for bs=1 ! */
1819:     PetscCall(MatAXPY_BasicWithPreallocation(B, Y, a, X, str));
1820:     PetscCall(MatHeaderMerge(Y, &B));
1821:     PetscCall(PetscFree(nnz_d));
1822:     PetscCall(PetscFree(nnz_o));
1823:   }
1824:   PetscFunctionReturn(PETSC_SUCCESS);
1825: }

1827: static PetscErrorCode MatConjugate_MPIBAIJ(Mat mat)
1828: {
1829:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)mat->data;

1831:   PetscFunctionBegin;
1832:   PetscCall(MatConjugate(a->A));
1833:   PetscCall(MatConjugate(a->B));
1834:   PetscFunctionReturn(PETSC_SUCCESS);
1835: }

1837: static PetscErrorCode MatRealPart_MPIBAIJ(Mat A)
1838: {
1839:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

1841:   PetscFunctionBegin;
1842:   PetscCall(MatRealPart(a->A));
1843:   PetscCall(MatRealPart(a->B));
1844:   PetscFunctionReturn(PETSC_SUCCESS);
1845: }

1847: static PetscErrorCode MatImaginaryPart_MPIBAIJ(Mat A)
1848: {
1849:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

1851:   PetscFunctionBegin;
1852:   PetscCall(MatImaginaryPart(a->A));
1853:   PetscCall(MatImaginaryPart(a->B));
1854:   PetscFunctionReturn(PETSC_SUCCESS);
1855: }

1857: static PetscErrorCode MatCreateSubMatrix_MPIBAIJ(Mat mat, IS isrow, IS iscol, MatReuse call, Mat *newmat)
1858: {
1859:   IS       iscol_local;
1860:   PetscInt csize;

1862:   PetscFunctionBegin;
1863:   PetscCall(ISGetLocalSize(iscol, &csize));
1864:   if (call == MAT_REUSE_MATRIX) {
1865:     PetscCall(PetscObjectQuery((PetscObject)*newmat, "ISAllGather", (PetscObject *)&iscol_local));
1866:     PetscCheck(iscol_local, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Submatrix passed in was not used before, cannot reuse");
1867:   } else {
1868:     PetscCall(ISAllGather(iscol, &iscol_local));
1869:   }
1870:   PetscCall(MatCreateSubMatrix_MPIBAIJ_Private(mat, isrow, iscol_local, csize, call, newmat, PETSC_FALSE));
1871:   if (call == MAT_INITIAL_MATRIX) {
1872:     PetscCall(PetscObjectCompose((PetscObject)*newmat, "ISAllGather", (PetscObject)iscol_local));
1873:     PetscCall(ISDestroy(&iscol_local));
1874:   }
1875:   PetscFunctionReturn(PETSC_SUCCESS);
1876: }

1878: /*
1879:   Not great since it makes two copies of the submatrix, first an SeqBAIJ
1880:   in local and then by concatenating the local matrices the end result.
1881:   Writing it directly would be much like MatCreateSubMatrices_MPIBAIJ().
1882:   This routine is used for BAIJ and SBAIJ matrices (unfortunate dependency).
1883: */
1884: PetscErrorCode MatCreateSubMatrix_MPIBAIJ_Private(Mat mat, IS isrow, IS iscol, PetscInt csize, MatReuse call, Mat *newmat, PetscBool sym)
1885: {
1886:   PetscMPIInt  rank, size;
1887:   PetscInt     i, m, n, rstart, row, rend, nz, *cwork, j, bs;
1888:   PetscInt    *ii, *jj, nlocal, *dlens, *olens, dlen, olen, jend, mglobal;
1889:   Mat          M, Mreuse;
1890:   MatScalar   *vwork, *aa;
1891:   MPI_Comm     comm;
1892:   IS           isrow_new, iscol_new;
1893:   Mat_SeqBAIJ *aij;

1895:   PetscFunctionBegin;
1896:   PetscCall(PetscObjectGetComm((PetscObject)mat, &comm));
1897:   PetscCallMPI(MPI_Comm_rank(comm, &rank));
1898:   PetscCallMPI(MPI_Comm_size(comm, &size));
1899:   /* The compression and expansion should be avoided. Doesn't point
1900:      out errors, might change the indices, hence buggey */
1901:   PetscCall(ISCompressIndicesGeneral(mat->rmap->N, mat->rmap->n, mat->rmap->bs, 1, &isrow, &isrow_new));
1902:   if (isrow == iscol) {
1903:     iscol_new = isrow_new;
1904:     PetscCall(PetscObjectReference((PetscObject)iscol_new));
1905:   } else PetscCall(ISCompressIndicesGeneral(mat->cmap->N, mat->cmap->n, mat->cmap->bs, 1, &iscol, &iscol_new));

1907:   if (call == MAT_REUSE_MATRIX) {
1908:     PetscCall(PetscObjectQuery((PetscObject)*newmat, "SubMatrix", (PetscObject *)&Mreuse));
1909:     PetscCheck(Mreuse, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Submatrix passed in was not used before, cannot reuse");
1910:     PetscCall(MatCreateSubMatrices_MPIBAIJ_local(mat, 1, &isrow_new, &iscol_new, MAT_REUSE_MATRIX, &Mreuse, sym));
1911:   } else {
1912:     PetscCall(MatCreateSubMatrices_MPIBAIJ_local(mat, 1, &isrow_new, &iscol_new, MAT_INITIAL_MATRIX, &Mreuse, sym));
1913:   }
1914:   PetscCall(ISDestroy(&isrow_new));
1915:   PetscCall(ISDestroy(&iscol_new));
1916:   /*
1917:       m - number of local rows
1918:       n - number of columns (same on all processors)
1919:       rstart - first row in new global matrix generated
1920:   */
1921:   PetscCall(MatGetBlockSize(mat, &bs));
1922:   PetscCall(MatGetSize(Mreuse, &m, &n));
1923:   m = m / bs;
1924:   n = n / bs;

1926:   if (call == MAT_INITIAL_MATRIX) {
1927:     aij = (Mat_SeqBAIJ *)Mreuse->data;
1928:     ii  = aij->i;
1929:     jj  = aij->j;

1931:     /*
1932:         Determine the number of non-zeros in the diagonal and off-diagonal
1933:         portions of the matrix in order to do correct preallocation
1934:     */

1936:     /* first get start and end of "diagonal" columns */
1937:     if (csize == PETSC_DECIDE) {
1938:       PetscCall(ISGetSize(isrow, &mglobal));
1939:       if (mglobal == n * bs) { /* square matrix */
1940:         nlocal = m;
1941:       } else {
1942:         nlocal = n / size + ((n % size) > rank);
1943:       }
1944:     } else {
1945:       nlocal = csize / bs;
1946:     }
1947:     PetscCallMPI(MPI_Scan(&nlocal, &rend, 1, MPIU_INT, MPI_SUM, comm));
1948:     rstart = rend - nlocal;
1949:     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);

1951:     /* next, compute all the lengths */
1952:     PetscCall(PetscMalloc2(m + 1, &dlens, m + 1, &olens));
1953:     for (i = 0; i < m; i++) {
1954:       jend = ii[i + 1] - ii[i];
1955:       olen = 0;
1956:       dlen = 0;
1957:       for (j = 0; j < jend; j++) {
1958:         if (*jj < rstart || *jj >= rend) olen++;
1959:         else dlen++;
1960:         jj++;
1961:       }
1962:       olens[i] = olen;
1963:       dlens[i] = dlen;
1964:     }
1965:     PetscCall(MatCreate(comm, &M));
1966:     PetscCall(MatSetSizes(M, bs * m, bs * nlocal, PETSC_DECIDE, bs * n));
1967:     PetscCall(MatSetType(M, sym ? ((PetscObject)mat)->type_name : MATMPIBAIJ));
1968:     PetscCall(MatSetOption(M, MAT_STRUCTURE_ONLY, mat->structure_only));
1969:     PetscCall(MatMPIBAIJSetPreallocation(M, bs, 0, dlens, 0, olens));
1970:     PetscCall(MatMPISBAIJSetPreallocation(M, bs, 0, dlens, 0, olens));
1971:     PetscCall(PetscFree2(dlens, olens));
1972:   } else {
1973:     PetscInt ml, nl;

1975:     M = *newmat;
1976:     PetscCall(MatGetLocalSize(M, &ml, &nl));
1977:     PetscCheck(ml == m * bs, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Previous matrix must be same size/layout as request");
1978:     PetscCall(MatZeroEntries(M));
1979:     /*
1980:          The next two lines are needed so we may call MatSetValues_MPIAIJ() below directly,
1981:        rather than the slower MatSetValues().
1982:     */
1983:     M->was_assembled = PETSC_TRUE;
1984:     M->assembled     = PETSC_FALSE;
1985:   }
1986:   PetscCall(MatSetOption(M, MAT_ROW_ORIENTED, PETSC_FALSE));
1987:   PetscCall(MatGetOwnershipRange(M, &rstart, &rend));
1988:   aij = (Mat_SeqBAIJ *)Mreuse->data;
1989:   ii  = aij->i;
1990:   jj  = aij->j;
1991:   aa  = aij->a;
1992:   for (i = 0; i < m; i++) {
1993:     row   = rstart / bs + i;
1994:     nz    = ii[i + 1] - ii[i];
1995:     cwork = jj;
1996:     jj    = PetscSafePointerPlusOffset(jj, nz);
1997:     vwork = aa;
1998:     aa    = PetscSafePointerPlusOffset(aa, nz * bs * bs);
1999:     PetscUseTypeMethod(M, setvaluesblocked, 1, &row, nz, cwork, vwork, INSERT_VALUES);
2000:   }

2002:   PetscCall(MatAssemblyBegin(M, MAT_FINAL_ASSEMBLY));
2003:   PetscCall(MatAssemblyEnd(M, MAT_FINAL_ASSEMBLY));
2004:   *newmat = M;

2006:   /* save submatrix used in processor for next request */
2007:   if (call == MAT_INITIAL_MATRIX) {
2008:     PetscCall(PetscObjectCompose((PetscObject)M, "SubMatrix", (PetscObject)Mreuse));
2009:     PetscCall(PetscObjectDereference((PetscObject)Mreuse));
2010:   }
2011:   PetscFunctionReturn(PETSC_SUCCESS);
2012: }

2014: static PetscErrorCode MatPermute_MPIBAIJ(Mat A, IS rowp, IS colp, Mat *B)
2015: {
2016:   MPI_Comm        comm, pcomm;
2017:   PetscInt        clocal_size, nrows;
2018:   const PetscInt *rows;
2019:   PetscMPIInt     size;
2020:   IS              crowp, lcolp;

2022:   PetscFunctionBegin;
2023:   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
2024:   /* make a collective version of 'rowp' */
2025:   PetscCall(PetscObjectGetComm((PetscObject)rowp, &pcomm));
2026:   if (pcomm == comm) {
2027:     crowp = rowp;
2028:   } else {
2029:     PetscCall(ISGetSize(rowp, &nrows));
2030:     PetscCall(ISGetIndices(rowp, &rows));
2031:     PetscCall(ISCreateGeneral(comm, nrows, rows, PETSC_COPY_VALUES, &crowp));
2032:     PetscCall(ISRestoreIndices(rowp, &rows));
2033:   }
2034:   PetscCall(ISSetPermutation(crowp));
2035:   /* make a local version of 'colp' */
2036:   PetscCall(PetscObjectGetComm((PetscObject)colp, &pcomm));
2037:   PetscCallMPI(MPI_Comm_size(pcomm, &size));
2038:   if (size == 1) {
2039:     lcolp = colp;
2040:   } else {
2041:     PetscCall(ISAllGather(colp, &lcolp));
2042:   }
2043:   PetscCall(ISSetPermutation(lcolp));
2044:   /* now we just get the submatrix */
2045:   PetscCall(MatGetLocalSize(A, NULL, &clocal_size));
2046:   PetscCall(MatCreateSubMatrix_MPIBAIJ_Private(A, crowp, lcolp, clocal_size, MAT_INITIAL_MATRIX, B, PETSC_FALSE));
2047:   /* clean up */
2048:   if (pcomm != comm) PetscCall(ISDestroy(&crowp));
2049:   if (size > 1) PetscCall(ISDestroy(&lcolp));
2050:   PetscFunctionReturn(PETSC_SUCCESS);
2051: }

2053: static PetscErrorCode MatGetGhosts_MPIBAIJ(Mat mat, PetscInt *nghosts, const PetscInt *ghosts[])
2054: {
2055:   Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)mat->data;
2056:   Mat_SeqBAIJ *B    = (Mat_SeqBAIJ *)baij->B->data;

2058:   PetscFunctionBegin;
2059:   if (nghosts) *nghosts = B->nbs;
2060:   if (ghosts) *ghosts = baij->garray;
2061:   PetscFunctionReturn(PETSC_SUCCESS);
2062: }

2064: static PetscErrorCode MatGetSeqNonzeroStructure_MPIBAIJ(Mat A, Mat *newmat)
2065: {
2066:   Mat          B;
2067:   Mat_MPIBAIJ *a  = (Mat_MPIBAIJ *)A->data;
2068:   Mat_SeqBAIJ *ad = (Mat_SeqBAIJ *)a->A->data, *bd = (Mat_SeqBAIJ *)a->B->data;
2069:   Mat_SeqAIJ  *b;
2070:   PetscMPIInt  size, rank, *recvcounts = NULL, *displs = NULL;
2071:   PetscInt     sendcount, i, *rstarts = A->rmap->range, n, cnt, j, bs = A->rmap->bs;
2072:   PetscInt     m, *garray = a->garray, *lens, *jsendbuf, *a_jsendbuf, *b_jsendbuf;

2074:   PetscFunctionBegin;
2075:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size));
2076:   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)A), &rank));

2078:   /*   Tell every processor the number of nonzeros per row  */
2079:   PetscCall(PetscMalloc1(A->rmap->N / bs, &lens));
2080:   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];
2081:   PetscCall(PetscMalloc1(2 * size, &recvcounts));
2082:   displs = recvcounts + size;
2083:   for (i = 0; i < size; i++) {
2084:     PetscCall(PetscMPIIntCast(A->rmap->range[i + 1] / bs - A->rmap->range[i] / bs, &recvcounts[i]));
2085:     PetscCall(PetscMPIIntCast(A->rmap->range[i] / bs, &displs[i]));
2086:   }
2087:   PetscCallMPI(MPI_Allgatherv(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, lens, recvcounts, displs, MPIU_INT, PetscObjectComm((PetscObject)A)));
2088:   /* Create the sequential matrix of the same type as the local block diagonal  */
2089:   PetscCall(MatCreate(PETSC_COMM_SELF, &B));
2090:   PetscCall(MatSetSizes(B, A->rmap->N / bs, A->cmap->N / bs, PETSC_DETERMINE, PETSC_DETERMINE));
2091:   PetscCall(MatSetType(B, MATSEQAIJ));
2092:   PetscCall(MatSetOption(B, MAT_STRUCTURE_ONLY, PETSC_TRUE));
2093:   PetscCall(MatSeqAIJSetPreallocation(B, 0, lens));
2094:   b = (Mat_SeqAIJ *)B->data;

2096:   /*     Copy my part of matrix column indices over  */
2097:   sendcount  = ad->nz + bd->nz;
2098:   jsendbuf   = b->j + b->i[rstarts[rank] / bs];
2099:   a_jsendbuf = ad->j;
2100:   b_jsendbuf = bd->j;
2101:   n          = A->rmap->rend / bs - A->rmap->rstart / bs;
2102:   cnt        = 0;
2103:   for (i = 0; i < n; i++) {
2104:     /* put in lower diagonal portion */
2105:     m = bd->i[i + 1] - bd->i[i];
2106:     while (m > 0) {
2107:       /* is it above diagonal (in bd (compressed) numbering) */
2108:       if (garray[*b_jsendbuf] > A->rmap->rstart / bs + i) break;
2109:       jsendbuf[cnt++] = garray[*b_jsendbuf++];
2110:       m--;
2111:     }

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

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

2121:   /*  Gather all column indices to all processors  */
2122:   for (i = 0; i < size; i++) {
2123:     recvcounts[i] = 0;
2124:     for (j = A->rmap->range[i] / bs; j < A->rmap->range[i + 1] / bs; j++) recvcounts[i] += lens[j];
2125:   }
2126:   displs[0] = 0;
2127:   for (i = 1; i < size; i++) displs[i] = displs[i - 1] + recvcounts[i - 1];
2128:   PetscCallMPI(MPI_Allgatherv(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, b->j, recvcounts, displs, MPIU_INT, PetscObjectComm((PetscObject)A)));
2129:   /*  Assemble the matrix into usable form (note numerical values not yet set)  */
2130:   /* set the b->ilen (length of each row) values */
2131:   PetscCall(PetscArraycpy(b->ilen, lens, A->rmap->N / bs));
2132:   /* set the b->i indices */
2133:   b->i[0] = 0;
2134:   for (i = 1; i <= A->rmap->N / bs; i++) b->i[i] = b->i[i - 1] + lens[i - 1];
2135:   PetscCall(PetscFree(lens));
2136:   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
2137:   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
2138:   PetscCall(PetscFree(recvcounts));

2140:   PetscCall(MatPropagateSymmetryOptions(A, B));
2141:   *newmat = B;
2142:   PetscFunctionReturn(PETSC_SUCCESS);
2143: }

2145: static PetscErrorCode MatSOR_MPIBAIJ(Mat matin, Vec bb, PetscReal omega, MatSORType flag, PetscReal fshift, PetscInt its, PetscInt lits, Vec xx)
2146: {
2147:   Mat_MPIBAIJ *mat = (Mat_MPIBAIJ *)matin->data;
2148:   Vec          bb1 = NULL;

2150:   PetscFunctionBegin;
2151:   if (flag == SOR_APPLY_UPPER) {
2152:     PetscUseTypeMethod(mat->A, sor, bb, omega, flag, fshift, lits, 1, xx);
2153:     PetscFunctionReturn(PETSC_SUCCESS);
2154:   }

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

2158:   if ((flag & SOR_LOCAL_SYMMETRIC_SWEEP) == SOR_LOCAL_SYMMETRIC_SWEEP) {
2159:     if (flag & SOR_ZERO_INITIAL_GUESS) {
2160:       PetscUseTypeMethod(mat->A, sor, bb, omega, flag, fshift, lits, 1, xx);
2161:       its--;
2162:     }

2164:     while (its--) {
2165:       PetscCall(VecScatterBegin(mat->Mvctx, xx, mat->lvec, INSERT_VALUES, SCATTER_FORWARD));
2166:       PetscCall(VecScatterEnd(mat->Mvctx, xx, mat->lvec, INSERT_VALUES, SCATTER_FORWARD));

2168:       /* update rhs: bb1 = bb - B*x */
2169:       PetscCall(VecScale(mat->lvec, -1.0));
2170:       PetscUseTypeMethod(mat->B, multadd, mat->lvec, bb, bb1);

2172:       /* local sweep */
2173:       PetscUseTypeMethod(mat->A, sor, bb1, omega, SOR_SYMMETRIC_SWEEP, fshift, lits, 1, xx);
2174:     }
2175:   } else if (flag & SOR_LOCAL_FORWARD_SWEEP) {
2176:     if (flag & SOR_ZERO_INITIAL_GUESS) {
2177:       PetscUseTypeMethod(mat->A, sor, bb, omega, flag, fshift, lits, 1, xx);
2178:       its--;
2179:     }
2180:     while (its--) {
2181:       PetscCall(VecScatterBegin(mat->Mvctx, xx, mat->lvec, INSERT_VALUES, SCATTER_FORWARD));
2182:       PetscCall(VecScatterEnd(mat->Mvctx, xx, mat->lvec, INSERT_VALUES, SCATTER_FORWARD));

2184:       /* update rhs: bb1 = bb - B*x */
2185:       PetscCall(VecScale(mat->lvec, -1.0));
2186:       PetscUseTypeMethod(mat->B, multadd, mat->lvec, bb, bb1);

2188:       /* local sweep */
2189:       PetscUseTypeMethod(mat->A, sor, bb1, omega, SOR_FORWARD_SWEEP, fshift, lits, 1, xx);
2190:     }
2191:   } else if (flag & SOR_LOCAL_BACKWARD_SWEEP) {
2192:     if (flag & SOR_ZERO_INITIAL_GUESS) {
2193:       PetscUseTypeMethod(mat->A, sor, bb, omega, flag, fshift, lits, 1, xx);
2194:       its--;
2195:     }
2196:     while (its--) {
2197:       PetscCall(VecScatterBegin(mat->Mvctx, xx, mat->lvec, INSERT_VALUES, SCATTER_FORWARD));
2198:       PetscCall(VecScatterEnd(mat->Mvctx, xx, mat->lvec, INSERT_VALUES, SCATTER_FORWARD));

2200:       /* update rhs: bb1 = bb - B*x */
2201:       PetscCall(VecScale(mat->lvec, -1.0));
2202:       PetscUseTypeMethod(mat->B, multadd, mat->lvec, bb, bb1);

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

2209:   PetscCall(VecDestroy(&bb1));
2210:   PetscFunctionReturn(PETSC_SUCCESS);
2211: }

2213: static PetscErrorCode MatGetColumnReductions_MPIBAIJ(Mat A, PetscInt type, PetscReal *reductions)
2214: {
2215:   Mat_MPIBAIJ *aij = (Mat_MPIBAIJ *)A->data;
2216:   PetscInt     m, N, i, *garray = aij->garray;
2217:   PetscInt     ib, jb, bs = A->rmap->bs;
2218:   Mat_SeqBAIJ *a_aij = (Mat_SeqBAIJ *)aij->A->data;
2219:   MatScalar   *a_val = a_aij->a;
2220:   Mat_SeqBAIJ *b_aij = (Mat_SeqBAIJ *)aij->B->data;
2221:   MatScalar   *b_val = b_aij->a;

2223:   PetscFunctionBegin;
2224:   PetscCall(MatGetSize(A, &m, &N));
2225:   PetscCall(PetscArrayzero(reductions, N));
2226:   if (type == NORM_2) {
2227:     for (i = a_aij->i[0]; i < a_aij->i[aij->A->rmap->n / bs]; i++) {
2228:       for (jb = 0; jb < bs; jb++) {
2229:         for (ib = 0; ib < bs; ib++) {
2230:           reductions[A->cmap->rstart + a_aij->j[i] * bs + jb] += PetscAbsScalar(*a_val * *a_val);
2231:           a_val++;
2232:         }
2233:       }
2234:     }
2235:     for (i = b_aij->i[0]; i < b_aij->i[aij->B->rmap->n / bs]; i++) {
2236:       for (jb = 0; jb < bs; jb++) {
2237:         for (ib = 0; ib < bs; ib++) {
2238:           reductions[garray[b_aij->j[i]] * bs + jb] += PetscAbsScalar(*b_val * *b_val);
2239:           b_val++;
2240:         }
2241:       }
2242:     }
2243:   } else if (type == NORM_1) {
2244:     for (i = a_aij->i[0]; i < a_aij->i[aij->A->rmap->n / bs]; i++) {
2245:       for (jb = 0; jb < bs; jb++) {
2246:         for (ib = 0; ib < bs; ib++) {
2247:           reductions[A->cmap->rstart + a_aij->j[i] * bs + jb] += PetscAbsScalar(*a_val);
2248:           a_val++;
2249:         }
2250:       }
2251:     }
2252:     for (i = b_aij->i[0]; i < b_aij->i[aij->B->rmap->n / bs]; i++) {
2253:       for (jb = 0; jb < bs; jb++) {
2254:         for (ib = 0; ib < bs; ib++) {
2255:           reductions[garray[b_aij->j[i]] * bs + jb] += PetscAbsScalar(*b_val);
2256:           b_val++;
2257:         }
2258:       }
2259:     }
2260:   } else if (type == NORM_INFINITY) {
2261:     for (i = a_aij->i[0]; i < a_aij->i[aij->A->rmap->n / bs]; i++) {
2262:       for (jb = 0; jb < bs; jb++) {
2263:         for (ib = 0; ib < bs; ib++) {
2264:           PetscInt col    = A->cmap->rstart + a_aij->j[i] * bs + jb;
2265:           reductions[col] = PetscMax(PetscAbsScalar(*a_val), reductions[col]);
2266:           a_val++;
2267:         }
2268:       }
2269:     }
2270:     for (i = b_aij->i[0]; i < b_aij->i[aij->B->rmap->n / bs]; i++) {
2271:       for (jb = 0; jb < bs; jb++) {
2272:         for (ib = 0; ib < bs; ib++) {
2273:           PetscInt col    = garray[b_aij->j[i]] * bs + jb;
2274:           reductions[col] = PetscMax(PetscAbsScalar(*b_val), reductions[col]);
2275:           b_val++;
2276:         }
2277:       }
2278:     }
2279:   } else if (type == REDUCTION_SUM_REALPART || type == REDUCTION_MEAN_REALPART) {
2280:     for (i = a_aij->i[0]; i < a_aij->i[aij->A->rmap->n / bs]; i++) {
2281:       for (jb = 0; jb < bs; jb++) {
2282:         for (ib = 0; ib < bs; ib++) {
2283:           reductions[A->cmap->rstart + a_aij->j[i] * bs + jb] += PetscRealPart(*a_val);
2284:           a_val++;
2285:         }
2286:       }
2287:     }
2288:     for (i = b_aij->i[0]; i < b_aij->i[aij->B->rmap->n / bs]; i++) {
2289:       for (jb = 0; jb < bs; jb++) {
2290:         for (ib = 0; ib < bs; ib++) {
2291:           reductions[garray[b_aij->j[i]] * bs + jb] += PetscRealPart(*b_val);
2292:           b_val++;
2293:         }
2294:       }
2295:     }
2296:   } else {
2297:     PetscCheck(type == REDUCTION_SUM_IMAGINARYPART || type == REDUCTION_MEAN_IMAGINARYPART, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Unknown reduction type");
2298:     for (i = a_aij->i[0]; i < a_aij->i[aij->A->rmap->n / bs]; i++) {
2299:       for (jb = 0; jb < bs; jb++) {
2300:         for (ib = 0; ib < bs; ib++) {
2301:           reductions[A->cmap->rstart + a_aij->j[i] * bs + jb] += PetscImaginaryPart(*a_val);
2302:           a_val++;
2303:         }
2304:       }
2305:     }
2306:     for (i = b_aij->i[0]; i < b_aij->i[aij->B->rmap->n / bs]; i++) {
2307:       for (jb = 0; jb < bs; jb++) {
2308:         for (ib = 0; ib < bs; ib++) {
2309:           reductions[garray[b_aij->j[i]] * bs + jb] += PetscImaginaryPart(*b_val);
2310:           b_val++;
2311:         }
2312:       }
2313:     }
2314:   }
2315:   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, reductions, N, MPIU_REAL, type == NORM_INFINITY ? MPIU_MAX : MPIU_SUM, PetscObjectComm((PetscObject)A)));
2316:   if (type == NORM_2) {
2317:     for (i = 0; i < N; i++) reductions[i] = PetscSqrtReal(reductions[i]);
2318:   } else if (type == REDUCTION_MEAN_REALPART || type == REDUCTION_MEAN_IMAGINARYPART) {
2319:     for (i = 0; i < N; i++) reductions[i] /= m;
2320:   }
2321:   PetscFunctionReturn(PETSC_SUCCESS);
2322: }

2324: static PetscErrorCode MatInvertBlockDiagonal_MPIBAIJ(Mat A, const PetscScalar **values)
2325: {
2326:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

2328:   PetscFunctionBegin;
2329:   PetscCall(MatInvertBlockDiagonal(a->A, values));
2330:   A->factorerrortype             = a->A->factorerrortype;
2331:   A->factorerror_zeropivot_value = a->A->factorerror_zeropivot_value;
2332:   A->factorerror_zeropivot_row   = a->A->factorerror_zeropivot_row;
2333:   PetscFunctionReturn(PETSC_SUCCESS);
2334: }

2336: static PetscErrorCode MatShift_MPIBAIJ(Mat Y, PetscScalar a)
2337: {
2338:   Mat_MPIBAIJ *maij = (Mat_MPIBAIJ *)Y->data;
2339:   Mat_SeqBAIJ *aij  = (Mat_SeqBAIJ *)maij->A->data;

2341:   PetscFunctionBegin;
2342:   if (!Y->preallocated) {
2343:     PetscCall(MatMPIBAIJSetPreallocation(Y, Y->rmap->bs, 1, NULL, 0, NULL));
2344:   } else if (!aij->nz) {
2345:     PetscInt nonew = aij->nonew;
2346:     PetscCall(MatSeqBAIJSetPreallocation(maij->A, Y->rmap->bs, 1, NULL));
2347:     aij->nonew = nonew;
2348:   }
2349:   PetscCall(MatShift_Basic(Y, a));
2350:   PetscFunctionReturn(PETSC_SUCCESS);
2351: }

2353: static PetscErrorCode MatGetDiagonalBlock_MPIBAIJ(Mat A, Mat *a)
2354: {
2355:   PetscFunctionBegin;
2356:   *a = ((Mat_MPIBAIJ *)A->data)->A;
2357:   PetscFunctionReturn(PETSC_SUCCESS);
2358: }

2360: static PetscErrorCode MatEliminateZeros_MPIBAIJ(Mat A, PetscBool keep)
2361: {
2362:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

2364:   PetscFunctionBegin;
2365:   PetscCall(MatEliminateZeros_SeqBAIJ(a->A, keep));        // possibly keep zero diagonal coefficients
2366:   PetscCall(MatEliminateZeros_SeqBAIJ(a->B, PETSC_FALSE)); // never keep zero diagonal coefficients
2367:   PetscFunctionReturn(PETSC_SUCCESS);
2368: }

2370: static struct _MatOps MatOps_Values = {MatSetValues_MPIBAIJ,
2371:                                        MatGetRow_MPIBAIJ,
2372:                                        MatRestoreRow_MPIBAIJ,
2373:                                        MatMult_MPIBAIJ,
2374:                                        /* 4*/ MatMultAdd_MPIBAIJ,
2375:                                        MatMultTranspose_MPIBAIJ,
2376:                                        MatMultTransposeAdd_MPIBAIJ,
2377:                                        NULL,
2378:                                        NULL,
2379:                                        NULL,
2380:                                        /*10*/ NULL,
2381:                                        NULL,
2382:                                        NULL,
2383:                                        MatSOR_MPIBAIJ,
2384:                                        MatTranspose_MPIBAIJ,
2385:                                        /*15*/ MatGetInfo_MPIBAIJ,
2386:                                        MatEqual_MPIBAIJ,
2387:                                        MatGetDiagonal_MPIBAIJ,
2388:                                        MatDiagonalScale_MPIBAIJ,
2389:                                        MatNorm_MPIBAIJ,
2390:                                        /*20*/ MatAssemblyBegin_MPIBAIJ,
2391:                                        MatAssemblyEnd_MPIBAIJ,
2392:                                        MatSetOption_MPIBAIJ,
2393:                                        MatZeroEntries_MPIBAIJ,
2394:                                        /*24*/ MatZeroRows_MPIBAIJ,
2395:                                        NULL,
2396:                                        NULL,
2397:                                        NULL,
2398:                                        NULL,
2399:                                        /*29*/ MatSetUp_MPI_Hash,
2400:                                        NULL,
2401:                                        NULL,
2402:                                        MatGetDiagonalBlock_MPIBAIJ,
2403:                                        NULL,
2404:                                        /*34*/ MatDuplicate_MPIBAIJ,
2405:                                        NULL,
2406:                                        NULL,
2407:                                        NULL,
2408:                                        NULL,
2409:                                        /*39*/ MatAXPY_MPIBAIJ,
2410:                                        MatCreateSubMatrices_MPIBAIJ,
2411:                                        MatIncreaseOverlap_MPIBAIJ,
2412:                                        MatGetValues_MPIBAIJ,
2413:                                        MatCopy_MPIBAIJ,
2414:                                        /*44*/ NULL,
2415:                                        MatScale_MPIBAIJ,
2416:                                        MatShift_MPIBAIJ,
2417:                                        NULL,
2418:                                        MatZeroRowsColumns_MPIBAIJ,
2419:                                        /*49*/ NULL,
2420:                                        NULL,
2421:                                        NULL,
2422:                                        NULL,
2423:                                        NULL,
2424:                                        /*54*/ MatFDColoringCreate_MPIXAIJ,
2425:                                        NULL,
2426:                                        MatSetUnfactored_MPIBAIJ,
2427:                                        MatPermute_MPIBAIJ,
2428:                                        MatSetValuesBlocked_MPIBAIJ,
2429:                                        /*59*/ MatCreateSubMatrix_MPIBAIJ,
2430:                                        MatDestroy_MPIBAIJ,
2431:                                        MatView_MPIBAIJ,
2432:                                        NULL,
2433:                                        NULL,
2434:                                        /*64*/ NULL,
2435:                                        NULL,
2436:                                        NULL,
2437:                                        NULL,
2438:                                        MatGetRowMaxAbs_MPIBAIJ,
2439:                                        /*69*/ NULL,
2440:                                        NULL,
2441:                                        NULL,
2442:                                        MatFDColoringApply_BAIJ,
2443:                                        NULL,
2444:                                        /*74*/ NULL,
2445:                                        NULL,
2446:                                        NULL,
2447:                                        NULL,
2448:                                        MatLoad_MPIBAIJ,
2449:                                        /*79*/ NULL,
2450:                                        NULL,
2451:                                        NULL,
2452:                                        NULL,
2453:                                        NULL,
2454:                                        /*84*/ NULL,
2455:                                        NULL,
2456:                                        NULL,
2457:                                        NULL,
2458:                                        NULL,
2459:                                        /*89*/ NULL,
2460:                                        NULL,
2461:                                        NULL,
2462:                                        NULL,
2463:                                        MatConjugate_MPIBAIJ,
2464:                                        /*94*/ NULL,
2465:                                        NULL,
2466:                                        MatRealPart_MPIBAIJ,
2467:                                        MatImaginaryPart_MPIBAIJ,
2468:                                        NULL,
2469:                                        /*99*/ NULL,
2470:                                        NULL,
2471:                                        NULL,
2472:                                        NULL,
2473:                                        NULL,
2474:                                        /*104*/ MatGetSeqNonzeroStructure_MPIBAIJ,
2475:                                        NULL,
2476:                                        MatGetGhosts_MPIBAIJ,
2477:                                        NULL,
2478:                                        NULL,
2479:                                        /*109*/ NULL,
2480:                                        NULL,
2481:                                        NULL,
2482:                                        NULL,
2483:                                        MatGetMultiProcBlock_MPIBAIJ,
2484:                                        /*114*/ NULL,
2485:                                        MatGetColumnReductions_MPIBAIJ,
2486:                                        MatInvertBlockDiagonal_MPIBAIJ,
2487:                                        NULL,
2488:                                        NULL,
2489:                                        /*119*/ NULL,
2490:                                        NULL,
2491:                                        NULL,
2492:                                        NULL,
2493:                                        NULL,
2494:                                        /*124*/ NULL,
2495:                                        MatSetBlockSizes_Default,
2496:                                        NULL,
2497:                                        MatFDColoringSetUp_MPIXAIJ,
2498:                                        NULL,
2499:                                        /*129*/ MatCreateMPIMatConcatenateSeqMat_MPIBAIJ,
2500:                                        NULL,
2501:                                        NULL,
2502:                                        NULL,
2503:                                        NULL,
2504:                                        /*134*/ NULL,
2505:                                        MatEliminateZeros_MPIBAIJ,
2506:                                        MatGetRowSumAbs_MPIBAIJ,
2507:                                        NULL,
2508:                                        NULL,
2509:                                        /*139*/ NULL,
2510:                                        MatCopyHashToXAIJ_MPI_Hash,
2511:                                        NULL,
2512:                                        NULL,
2513:                                        NULL,
2514:                                        /*144*/ NULL,
2515:                                        NULL,
2516:                                        NULL,
2517:                                        NULL};

2519: PETSC_INTERN PetscErrorCode MatConvert_MPIBAIJ_MPISBAIJ(Mat, MatType, MatReuse, Mat *);
2520: PETSC_INTERN PetscErrorCode MatConvert_XAIJ_IS(Mat, MatType, MatReuse, Mat *);

2522: static PetscErrorCode MatMPIBAIJSetPreallocationCSR_MPIBAIJ(Mat B, PetscInt bs, const PetscInt ii[], const PetscInt jj[], const PetscScalar V[])
2523: {
2524:   PetscInt        m, rstart, cstart, cend;
2525:   PetscInt        i, j, dlen, olen, nz, nz_max = 0, *d_nnz = NULL, *o_nnz = NULL;
2526:   const PetscInt *JJ          = NULL;
2527:   PetscScalar    *values      = NULL;
2528:   PetscBool       roworiented = ((Mat_MPIBAIJ *)B->data)->roworiented;
2529:   PetscBool       nooffprocentries;

2531:   PetscFunctionBegin;
2532:   PetscCall(PetscLayoutSetBlockSize(B->rmap, bs));
2533:   PetscCall(PetscLayoutSetBlockSize(B->cmap, bs));
2534:   PetscCall(PetscLayoutSetUp(B->rmap));
2535:   PetscCall(PetscLayoutSetUp(B->cmap));
2536:   PetscCall(PetscLayoutGetBlockSize(B->rmap, &bs));
2537:   m      = B->rmap->n / bs;
2538:   rstart = B->rmap->rstart / bs;
2539:   cstart = B->cmap->rstart / bs;
2540:   cend   = B->cmap->rend / bs;

2542:   PetscCheck(!ii[0], PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "ii[0] must be 0 but it is %" PetscInt_FMT, ii[0]);
2543:   PetscCall(PetscMalloc2(m, &d_nnz, m, &o_nnz));
2544:   for (i = 0; i < m; i++) {
2545:     nz = ii[i + 1] - ii[i];
2546:     PetscCheck(nz >= 0, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Local row %" PetscInt_FMT " has a negative number of columns %" PetscInt_FMT, i, nz);
2547:     nz_max = PetscMax(nz_max, nz);
2548:     dlen   = 0;
2549:     olen   = 0;
2550:     JJ     = jj + ii[i];
2551:     for (j = 0; j < nz; j++) {
2552:       if (*JJ < cstart || *JJ >= cend) olen++;
2553:       else dlen++;
2554:       JJ++;
2555:     }
2556:     d_nnz[i] = dlen;
2557:     o_nnz[i] = olen;
2558:   }
2559:   PetscCall(MatMPIBAIJSetPreallocation(B, bs, 0, d_nnz, 0, o_nnz));
2560:   PetscCall(PetscFree2(d_nnz, o_nnz));

2562:   values = (PetscScalar *)V;
2563:   if (!values) PetscCall(PetscCalloc1(bs * bs * nz_max, &values));
2564:   for (i = 0; i < m; i++) {
2565:     PetscInt        row   = i + rstart;
2566:     PetscInt        ncols = ii[i + 1] - ii[i];
2567:     const PetscInt *icols = jj + ii[i];
2568:     if (bs == 1 || !roworiented) { /* block ordering matches the non-nested layout of MatSetValues so we can insert entire rows */
2569:       const PetscScalar *svals = values + (V ? (bs * bs * ii[i]) : 0);
2570:       PetscCall(MatSetValuesBlocked_MPIBAIJ(B, 1, &row, ncols, icols, svals, INSERT_VALUES));
2571:     } else { /* block ordering does not match so we can only insert one block at a time. */
2572:       PetscInt j;
2573:       for (j = 0; j < ncols; j++) {
2574:         const PetscScalar *svals = values + (V ? (bs * bs * (ii[i] + j)) : 0);
2575:         PetscCall(MatSetValuesBlocked_MPIBAIJ(B, 1, &row, 1, &icols[j], svals, INSERT_VALUES));
2576:       }
2577:     }
2578:   }

2580:   if (!V) PetscCall(PetscFree(values));
2581:   nooffprocentries    = B->nooffprocentries;
2582:   B->nooffprocentries = PETSC_TRUE;
2583:   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
2584:   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
2585:   B->nooffprocentries = nooffprocentries;

2587:   PetscCall(MatSetOption(B, MAT_NEW_NONZERO_LOCATION_ERR, PETSC_TRUE));
2588:   PetscFunctionReturn(PETSC_SUCCESS);
2589: }

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

2594:   Collective

2596:   Input Parameters:
2597: + B  - the matrix
2598: . bs - the block size
2599: . i  - the indices into `j` for the start of each local row (starts with zero)
2600: . j  - the column indices for each local row (starts with zero) these must be sorted for each row
2601: - v  - optional values in the matrix, use `NULL` if not provided

2603:   Level: advanced

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

2610:   The order of the entries in values is specified by the `MatOption` `MAT_ROW_ORIENTED`.  For example, C programs
2611:   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
2612:   over rows within a block and the last index is over columns within a block row.  Fortran programs will likely set
2613:   `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
2614:   block column and the second index is over columns within a block.

2616:   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

2618: .seealso: `Mat`, `MatCreate()`, `MatCreateSeqAIJ()`, `MatSetValues()`, `MatMPIBAIJSetPreallocation()`, `MatCreateAIJ()`, `MATMPIAIJ`, `MatCreateMPIBAIJWithArrays()`, `MATMPIBAIJ`
2619: @*/
2620: PetscErrorCode MatMPIBAIJSetPreallocationCSR(Mat B, PetscInt bs, const PetscInt i[], const PetscInt j[], const PetscScalar v[])
2621: {
2622:   PetscFunctionBegin;
2626:   PetscTryMethod(B, "MatMPIBAIJSetPreallocationCSR_C", (Mat, PetscInt, const PetscInt[], const PetscInt[], const PetscScalar[]), (B, bs, i, j, v));
2627:   PetscFunctionReturn(PETSC_SUCCESS);
2628: }

2630: PetscErrorCode MatMPIBAIJSetPreallocation_MPIBAIJ(Mat B, PetscInt bs, PetscInt d_nz, const PetscInt *d_nnz, PetscInt o_nz, const PetscInt *o_nnz)
2631: {
2632:   Mat_MPIBAIJ *b = (Mat_MPIBAIJ *)B->data;
2633:   PetscInt     i;
2634:   PetscMPIInt  size;

2636:   PetscFunctionBegin;
2637:   if (B->hash_active) {
2638:     B->ops[0]      = b->cops;
2639:     B->hash_active = PETSC_FALSE;
2640:   }
2641:   if (!B->preallocated) PetscCall(MatStashCreate_Private(PetscObjectComm((PetscObject)B), bs, &B->bstash));
2642:   PetscCall(MatSetBlockSize(B, bs));
2643:   PetscCall(PetscLayoutSetUp(B->rmap));
2644:   PetscCall(PetscLayoutSetUp(B->cmap));
2645:   PetscCall(PetscLayoutGetBlockSize(B->rmap, &bs));

2647:   if (d_nnz) {
2648:     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]);
2649:   }
2650:   if (o_nnz) {
2651:     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]);
2652:   }

2654:   b->bs2 = bs * bs;
2655:   b->mbs = B->rmap->n / bs;
2656:   b->nbs = B->cmap->n / bs;
2657:   b->Mbs = B->rmap->N / bs;
2658:   b->Nbs = B->cmap->N / bs;

2660:   for (i = 0; i <= b->size; i++) b->rangebs[i] = B->rmap->range[i] / bs;
2661:   b->rstartbs = B->rmap->rstart / bs;
2662:   b->rendbs   = B->rmap->rend / bs;
2663:   b->cstartbs = B->cmap->rstart / bs;
2664:   b->cendbs   = B->cmap->rend / bs;

2666: #if PetscDefined(USE_CTABLE)
2667:   PetscCall(PetscHMapIDestroy(&b->colmap));
2668: #else
2669:   PetscCall(PetscFree(b->colmap));
2670: #endif
2671:   PetscCall(PetscFree(b->garray));
2672:   PetscCall(VecDestroy(&b->lvec));
2673:   PetscCall(VecScatterDestroy(&b->Mvctx));

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

2677:   MatSeqXAIJGetOptions_Private(b->B);
2678:   PetscCall(MatDestroy(&b->B));
2679:   PetscCall(MatCreate(PETSC_COMM_SELF, &b->B));
2680:   PetscCall(MatSetSizes(b->B, B->rmap->n, size > 1 ? B->cmap->N : 0, B->rmap->n, size > 1 ? B->cmap->N : 0));
2681:   PetscCall(MatSetType(b->B, MATSEQBAIJ));
2682:   MatSeqXAIJRestoreOptions_Private(b->B);
2683:   PetscCall(MatSetOption(b->B, MAT_STRUCTURE_ONLY, B->structure_only));

2685:   MatSeqXAIJGetOptions_Private(b->A);
2686:   PetscCall(MatDestroy(&b->A));
2687:   PetscCall(MatCreate(PETSC_COMM_SELF, &b->A));
2688:   PetscCall(MatSetSizes(b->A, B->rmap->n, B->cmap->n, B->rmap->n, B->cmap->n));
2689:   PetscCall(MatSetType(b->A, MATSEQBAIJ));
2690:   MatSeqXAIJRestoreOptions_Private(b->A);
2691:   PetscCall(MatSetOption(b->A, MAT_STRUCTURE_ONLY, B->structure_only));

2693:   PetscCall(MatSeqBAIJSetPreallocation(b->A, bs, d_nz, d_nnz));
2694:   PetscCall(MatSeqBAIJSetPreallocation(b->B, bs, o_nz, o_nnz));
2695:   B->preallocated  = PETSC_TRUE;
2696:   B->was_assembled = PETSC_FALSE;
2697:   B->assembled     = PETSC_FALSE;
2698:   PetscFunctionReturn(PETSC_SUCCESS);
2699: }

2701: extern PetscErrorCode MatDiagonalScaleLocal_MPIBAIJ(Mat, Vec);
2702: extern PetscErrorCode MatSetHashTableFactor_MPIBAIJ(Mat, PetscReal);

2704: PETSC_INTERN PetscErrorCode MatConvert_MPIBAIJ_MPIAdj(Mat B, MatType newtype, MatReuse reuse, Mat *adj)
2705: {
2706:   Mat_MPIBAIJ    *b = (Mat_MPIBAIJ *)B->data;
2707:   Mat_SeqBAIJ    *d = (Mat_SeqBAIJ *)b->A->data, *o = (Mat_SeqBAIJ *)b->B->data;
2708:   PetscInt        M = B->rmap->n / B->rmap->bs, i, *ii, *jj, cnt, j, k, rstart = B->rmap->rstart / B->rmap->bs;
2709:   const PetscInt *id = d->i, *jd = d->j, *io = o->i, *jo = o->j, *garray = b->garray;

2711:   PetscFunctionBegin;
2712:   PetscCall(PetscMalloc1(M + 1, &ii));
2713:   ii[0] = 0;
2714:   for (i = 0; i < M; i++) {
2715:     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]);
2716:     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]);
2717:     ii[i + 1] = ii[i] + id[i + 1] - id[i] + io[i + 1] - io[i];
2718:     /* remove one from count of matrix has diagonal */
2719:     for (j = id[i]; j < id[i + 1]; j++) {
2720:       if (jd[j] == i) {
2721:         ii[i + 1]--;
2722:         break;
2723:       }
2724:     }
2725:   }
2726:   PetscCall(PetscMalloc1(ii[M], &jj));
2727:   cnt = 0;
2728:   for (i = 0; i < M; i++) {
2729:     for (j = io[i]; j < io[i + 1]; j++) {
2730:       if (garray[jo[j]] > rstart) break;
2731:       jj[cnt++] = garray[jo[j]];
2732:     }
2733:     for (k = id[i]; k < id[i + 1]; k++) {
2734:       if (jd[k] != i) jj[cnt++] = rstart + jd[k];
2735:     }
2736:     for (; j < io[i + 1]; j++) jj[cnt++] = garray[jo[j]];
2737:   }
2738:   PetscCall(MatCreateMPIAdj(PetscObjectComm((PetscObject)B), M, B->cmap->N / B->rmap->bs, ii, jj, NULL, adj));
2739:   PetscFunctionReturn(PETSC_SUCCESS);
2740: }

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

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

2746: PETSC_INTERN PetscErrorCode MatConvert_MPIBAIJ_MPIAIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
2747: {
2748:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;
2749:   Mat_MPIAIJ  *b;
2750:   Mat          B;

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

2755:   if (reuse == MAT_REUSE_MATRIX) {
2756:     B = *newmat;
2757:   } else {
2758:     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
2759:     PetscCall(MatSetType(B, MATMPIAIJ));
2760:     PetscCall(MatSetSizes(B, A->rmap->n, A->cmap->n, A->rmap->N, A->cmap->N));
2761:     PetscCall(MatSetBlockSizes(B, A->rmap->bs, A->cmap->bs));
2762:     PetscCall(MatSeqAIJSetPreallocation(B, 0, NULL));
2763:     PetscCall(MatMPIAIJSetPreallocation(B, 0, NULL, 0, NULL));
2764:   }
2765:   b = (Mat_MPIAIJ *)B->data;

2767:   if (reuse == MAT_REUSE_MATRIX) {
2768:     PetscCall(MatConvert_SeqBAIJ_SeqAIJ(a->A, MATSEQAIJ, MAT_REUSE_MATRIX, &b->A));
2769:     PetscCall(MatConvert_SeqBAIJ_SeqAIJ(a->B, MATSEQAIJ, MAT_REUSE_MATRIX, &b->B));
2770:   } else {
2771:     PetscInt   *garray = a->garray;
2772:     Mat_SeqAIJ *bB;
2773:     PetscInt    bs, nnz;
2774:     PetscCall(MatDestroy(&b->A));
2775:     PetscCall(MatDestroy(&b->B));
2776:     /* just clear out the data structure */
2777:     PetscCall(MatDisAssemble_MPIAIJ(B, PETSC_FALSE));
2778:     PetscCall(MatConvert_SeqBAIJ_SeqAIJ(a->A, MATSEQAIJ, MAT_INITIAL_MATRIX, &b->A));
2779:     PetscCall(MatConvert_SeqBAIJ_SeqAIJ(a->B, MATSEQAIJ, MAT_INITIAL_MATRIX, &b->B));

2781:     /* Global numbering for b->B columns */
2782:     bB  = (Mat_SeqAIJ *)b->B->data;
2783:     bs  = A->rmap->bs;
2784:     nnz = bB->i[A->rmap->n];
2785:     for (PetscInt k = 0; k < nnz; k++) {
2786:       PetscInt bj = bB->j[k] / bs;
2787:       PetscInt br = bB->j[k] % bs;
2788:       bB->j[k]    = garray[bj] * bs + br;
2789:     }
2790:   }
2791:   PetscCall(MatSetOption(B, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
2792:   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
2793:   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
2794:   PetscCall(MatSetOption(B, MAT_NO_OFF_PROC_ENTRIES, PETSC_FALSE));

2796:   if (reuse == MAT_INPLACE_MATRIX) {
2797:     PetscCall(MatHeaderReplace(A, &B));
2798:   } else {
2799:     *newmat = B;
2800:   }
2801:   PetscFunctionReturn(PETSC_SUCCESS);
2802: }

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

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

2813:    Level: beginner

2815:    Notes:
2816:     Call `MatSetOption(A, MAT_STRUCTURE_ONLY, PETSC_TRUE)` before preallocation or `MatSetUp()` to store only the nonzero pattern.
2817:     The assembled matrix has no numerical value array. Row and column indices supplied during insertion are retained, while numerical values are ignored.
2818:     Such matrices can be used for structural operations, but not for numerical operations.

2820: .seealso: `Mat`, `MATBAIJ`, `MATSEQBAIJ`, `MatCreateBAIJ`
2821: M*/

2823: typedef struct {
2824:   MPIAIJ_MPIDense scatter;
2825:   Mat             workC;
2826: } MPIBAIJ_MPIDense;

2828: static PetscErrorCode MatMPIBAIJ_MPIDenseDestroy(PetscCtxRt ctx)
2829: {
2830:   MPIBAIJ_MPIDense *data = *(MPIBAIJ_MPIDense **)ctx;

2832:   PetscFunctionBegin;
2833:   PetscCall(MatDestroy(&data->workC));
2834:   PetscCall(MatMPIDenseScatterDestroy_Private(&data->scatter));
2835:   PetscCall(PetscFree(data));
2836:   PetscFunctionReturn(PETSC_SUCCESS);
2837: }

2839: static PetscErrorCode MatMPIDenseScatter_MPIBAIJ(Mat A, Mat B, Mat workB, Mat C)
2840: {
2841:   Mat_MPIBAIJ      *baij = (Mat_MPIBAIJ *)A->data;
2842:   MPIBAIJ_MPIDense *data = (MPIBAIJ_MPIDense *)C->product->data;
2843:   PetscInt          bs;

2845:   PetscFunctionBegin;
2846:   PetscCall(MatGetBlockSize(A, &bs));
2847:   PetscCall(MatMPIDenseScatter_Private(baij->Mvctx, baij->B->cmap->n, bs, workB, &data->scatter, B, C));
2848:   PetscFunctionReturn(PETSC_SUCCESS);
2849: }

2851: static PetscErrorCode MatMatMultNumeric_MPIBAIJ_MPIDense(Mat A, Mat B, Mat C)
2852: {
2853:   Mat_MPIBAIJ      *baij   = (Mat_MPIBAIJ *)A->data;
2854:   Mat_MPIDense     *bdense = (Mat_MPIDense *)B->data;
2855:   Mat_MPIDense     *cdense = (Mat_MPIDense *)C->data;
2856:   Mat               workB;
2857:   MPIBAIJ_MPIDense *data;

2859:   PetscFunctionBegin;
2860:   MatCheckProduct(C, 3);
2861:   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
2862:   data = (MPIBAIJ_MPIDense *)C->product->data;
2863:   if (!cdense->A->product) {
2864:     PetscCall(MatProductCreateWithMat(baij->A, bdense->A, NULL, cdense->A));
2865:     PetscCall(MatProductSetType(cdense->A, MATPRODUCT_AB));
2866:     PetscCall(MatProductSetFromOptions(cdense->A));
2867:     PetscCall(MatProductSymbolic(cdense->A));
2868:   } else PetscCall(MatProductReplaceMats(baij->A, bdense->A, NULL, cdense->A));
2869:   PetscCall(MatProductNumeric(cdense->A));

2871:   if (data->scatter.workB->cmap->n == B->cmap->N) {
2872:     workB = data->scatter.workB;
2873:     PetscCall(MatMPIDenseScatter_MPIBAIJ(A, B, workB, C));
2874:     if (data->workC) {
2875:       PetscCall(MatProductReplaceMats(baij->B, workB, NULL, data->workC));
2876:       PetscCall(MatProductNumeric(data->workC));
2877:       PetscCall(MatAXPY(cdense->A, 1.0, data->workC, SAME_NONZERO_PATTERN));
2878:     }
2879:   } else {
2880:     Mat           Bb, Cb, workC;
2881:     Mat_MPIDense *cbdense;
2882:     PetscInt      BN = B->cmap->N, n = data->scatter.workB->cmap->n, cols;

2884:     PetscCheck(n > 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Column batch size must be positive");
2885:     for (PetscInt i = 0; i < BN; i += n) {
2886:       cols  = PetscMin(n, BN - i);
2887:       workB = data->scatter.workB;
2888:       workC = data->workC;
2889:       if (cols != n) {
2890:         PetscCall(MatDenseGetSubMatrix(data->scatter.workB, PETSC_DECIDE, PETSC_DECIDE, 0, cols, &workB));
2891:         if (workC) PetscCall(MatDenseGetSubMatrix(data->workC, PETSC_DECIDE, PETSC_DECIDE, 0, cols, &workC));
2892:       }
2893:       PetscCall(MatDenseGetSubMatrix(B, PETSC_DECIDE, PETSC_DECIDE, i, i + cols, &Bb));
2894:       PetscCall(MatDenseGetSubMatrix(C, PETSC_DECIDE, PETSC_DECIDE, i, i + cols, &Cb));
2895:       PetscCall(MatMPIDenseScatter_MPIBAIJ(A, Bb, workB, C));
2896:       if (workC) {
2897:         cbdense = (Mat_MPIDense *)Cb->data;
2898:         PetscCall(MatProductReplaceMats(baij->B, workB, NULL, workC));
2899:         PetscCall(MatProductNumeric(workC));
2900:         PetscCall(MatAXPY(cbdense->A, 1.0, workC, SAME_NONZERO_PATTERN));
2901:       }
2902:       if (cols != n) {
2903:         if (workC) PetscCall(MatDenseRestoreSubMatrix(data->workC, &workC));
2904:         PetscCall(MatDenseRestoreSubMatrix(data->scatter.workB, &workB));
2905:       }
2906:       PetscCall(MatDenseRestoreSubMatrix(B, &Bb));
2907:       PetscCall(MatDenseRestoreSubMatrix(C, &Cb));
2908:     }
2909:   }
2910:   PetscFunctionReturn(PETSC_SUCCESS);
2911: }

2913: static PetscErrorCode MatMatMultSymbolic_MPIBAIJ_MPIDense(Mat A, Mat B, PetscReal fill, Mat C)
2914: {
2915:   Mat_MPIBAIJ      *baij = (Mat_MPIBAIJ *)A->data;
2916:   MPIBAIJ_MPIDense *data;
2917:   VecScatter        ctx = baij->Mvctx;
2918:   PetscInt          nz  = baij->B->cmap->n, bs;
2919:   PetscInt          Am = A->rmap->n, BN = B->cmap->N, Bbn, numBb;
2920:   Mat               workB1, workC1;
2921:   PetscBool         cisdense;

2923:   PetscFunctionBegin;
2924:   MatCheckProduct(C, 4);
2925:   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
2926:   PetscCall(PetscObjectBaseTypeCompare((PetscObject)C, MATMPIDENSE, &cisdense));
2927:   if (!cisdense) {
2928:     PetscCall(MatSetType(C, ((PetscObject)B)->type_name));
2929:     PetscCall(MatSetVecType(C, B->defaultvectype));
2930:   }
2931:   PetscCall(MatSetSizes(C, Am, B->cmap->n, A->rmap->N, BN));
2932:   PetscCall(MatSetBlockSizesFromMats(C, A, B));
2933:   PetscCall(MatSetUp(C));
2934:   PetscCall(MatGetBlockSize(A, &bs));
2935:   PetscCall(PetscNew(&data));
2936:   PetscCall(MatMPIDenseScatterSetUp_Private(ctx, nz, bs, Am, B, C, &data->scatter, &Bbn, &numBb));

2938:   PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
2939:   PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
2940:   PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
2941:   PetscCall(MatProductClear(baij->A));
2942:   PetscCall(MatProductClear(((Mat_MPIDense *)B->data)->A));
2943:   PetscCall(MatProductClear(((Mat_MPIDense *)C->data)->A));
2944:   PetscCall(MatProductCreateWithMat(baij->A, ((Mat_MPIDense *)B->data)->A, NULL, ((Mat_MPIDense *)C->data)->A));
2945:   PetscCall(MatProductSetType(((Mat_MPIDense *)C->data)->A, MATPRODUCT_AB));
2946:   PetscCall(MatProductSetFromOptions(((Mat_MPIDense *)C->data)->A));
2947:   PetscCall(MatProductSymbolic(((Mat_MPIDense *)C->data)->A));

2949:   if (nz) {
2950:     PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, Am, Bbn ? Bbn : BN, NULL, &data->workC));
2951:     PetscCall(MatProductCreateWithMat(baij->B, data->scatter.workB, NULL, data->workC));
2952:     PetscCall(MatProductSetType(data->workC, MATPRODUCT_AB));
2953:     PetscCall(MatProductSetFromOptions(data->workC));
2954:     PetscCall(MatProductSymbolic(data->workC));
2955:     if (numBb && BN % Bbn) {
2956:       PetscCall(MatDenseGetSubMatrix(data->scatter.workB, PETSC_DECIDE, PETSC_DECIDE, 0, BN % Bbn, &workB1));
2957:       PetscCall(MatDenseGetSubMatrix(data->workC, PETSC_DECIDE, PETSC_DECIDE, 0, BN % Bbn, &workC1));
2958:       PetscCall(MatProductCreateWithMat(baij->B, workB1, NULL, workC1));
2959:       PetscCall(MatProductSetType(workC1, MATPRODUCT_AB));
2960:       PetscCall(MatProductSetFromOptions(workC1));
2961:       PetscCall(MatProductSymbolic(workC1));
2962:       PetscCall(MatDenseRestoreSubMatrix(data->workC, &workC1));
2963:       PetscCall(MatDenseRestoreSubMatrix(data->scatter.workB, &workB1));
2964:     }
2965:   }

2967:   C->product->data       = data;
2968:   C->product->destroy    = MatMPIBAIJ_MPIDenseDestroy;
2969:   C->ops->matmultnumeric = MatMatMultNumeric_MPIBAIJ_MPIDense;
2970:   PetscFunctionReturn(PETSC_SUCCESS);
2971: }

2973: PETSC_INTERN PetscErrorCode MatProductSetFromOptions_MPIBAIJ_MPIDense(Mat C)
2974: {
2975:   Mat_Product *product = C->product;

2977:   PetscFunctionBegin;
2978:   MatCheckProduct(C, 1);
2979:   if (product->type == MATPRODUCT_AB) {
2980:     C->ops->matmultsymbolic = MatMatMultSymbolic_MPIBAIJ_MPIDense;
2981:     C->ops->productsymbolic = MatProductSymbolic_AB;
2982:   }
2983:   PetscFunctionReturn(PETSC_SUCCESS);
2984: }

2986: static PetscErrorCode MatGetMultPetscSF_MPIBAIJ(Mat A, PetscSF *sf)
2987: {
2988:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;

2990:   PetscFunctionBegin;
2991:   *sf = a->Mvctx;
2992:   PetscFunctionReturn(PETSC_SUCCESS);
2993: }

2995: PETSC_EXTERN PetscErrorCode MatCreate_MPIBAIJ(Mat B)
2996: {
2997:   Mat_MPIBAIJ *b;
2998:   PetscBool    flg = PETSC_FALSE;

3000:   PetscFunctionBegin;
3001:   PetscCall(PetscNew(&b));
3002:   B->data      = (void *)b;
3003:   B->ops[0]    = MatOps_Values;
3004:   B->assembled = PETSC_FALSE;

3006:   B->insertmode = NOT_SET_VALUES;
3007:   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)B), &b->rank));
3008:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)B), &b->size));

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

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

3016:   b->donotstash  = PETSC_FALSE;
3017:   b->colmap      = NULL;
3018:   b->garray      = NULL;
3019:   b->roworiented = PETSC_TRUE;

3021:   /* stuff used in block assembly */
3022:   b->barray = NULL;

3024:   /* stuff used for matrix vector multiply */
3025:   b->lvec  = NULL;
3026:   b->Mvctx = NULL;

3028:   /* stuff for MatGetRow() */
3029:   b->rowindices   = NULL;
3030:   b->rowvalues    = NULL;
3031:   b->getrowactive = PETSC_FALSE;

3033:   /* hash table stuff */
3034:   b->ht           = NULL;
3035:   b->hd           = NULL;
3036:   b->ht_size      = 0;
3037:   b->ht_flag      = PETSC_FALSE;
3038:   b->ht_fact      = 0;
3039:   b->ht_total_ct  = 0;
3040:   b->ht_insert_ct = 0;

3042:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_mpibaij_mpiadj_C", MatConvert_MPIBAIJ_MPIAdj));
3043:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_mpibaij_mpiaij_C", MatConvert_MPIBAIJ_MPIAIJ));
3044:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_mpibaij_mpisbaij_C", MatConvert_MPIBAIJ_MPISBAIJ));
3045: #if PetscDefined(HAVE_HYPRE)
3046:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_mpibaij_hypre_C", MatConvert_AIJ_HYPRE));
3047: #endif
3048:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatStoreValues_C", MatStoreValues_MPIBAIJ));
3049:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatRetrieveValues_C", MatRetrieveValues_MPIBAIJ));
3050:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatMPIBAIJSetPreallocation_C", MatMPIBAIJSetPreallocation_MPIBAIJ));
3051:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatMPIBAIJSetPreallocationCSR_C", MatMPIBAIJSetPreallocationCSR_MPIBAIJ));
3052:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatDiagonalScaleLocal_C", MatDiagonalScaleLocal_MPIBAIJ));
3053:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatSetHashTableFactor_C", MatSetHashTableFactor_MPIBAIJ));
3054:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_mpibaij_is_C", MatConvert_XAIJ_IS));
3055:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatGetMultPetscSF_C", MatGetMultPetscSF_MPIBAIJ));
3056:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatProductSetFromOptions_mpibaij_mpidense_C", MatProductSetFromOptions_MPIBAIJ_MPIDense));
3057: #if PetscDefined(HAVE_LIBXSMM)
3058:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_mpibaij_mpibaijlibxsmm_C", MatConvert_MPIBAIJ_MPIBAIJLIBXSMM));
3059: #endif
3060:   PetscCall(PetscObjectChangeTypeName((PetscObject)B, MATMPIBAIJ));

3062:   PetscOptionsBegin(PetscObjectComm((PetscObject)B), NULL, "Options for loading MPIBAIJ matrix 1", "Mat");
3063:   PetscCall(PetscOptionsName("-mat_use_hash_table", "Use hash table to save time in constructing matrix", "MatSetOption", &flg));
3064:   if (flg) {
3065:     PetscReal fact = 1.39;
3066:     PetscCall(MatSetOption(B, MAT_USE_HASH_TABLE, PETSC_TRUE));
3067:     PetscCall(PetscOptionsReal("-mat_use_hash_table", "Use hash table factor", "MatMPIBAIJSetHashTableFactor", fact, &fact, NULL));
3068:     if (fact <= 1.0) fact = 1.39;
3069:     PetscCall(MatMPIBAIJSetHashTableFactor(B, fact));
3070:     PetscCall(PetscInfo(B, "Hash table Factor used %5.2g\n", (double)fact));
3071:   }
3072:   PetscOptionsEnd();
3073:   PetscFunctionReturn(PETSC_SUCCESS);
3074: }

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

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

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

3086:   Level: beginner

3088:   Notes:
3089:   Call `MatSetOption(A, MAT_STRUCTURE_ONLY, PETSC_TRUE)` before preallocation or `MatSetUp()` to store only the nonzero pattern.
3090:   The assembled matrix has no numerical value array. Row and column indices supplied during insertion are retained, while numerical values are ignored.
3091:   Such matrices can be used for structural operations, but not for numerical operations.

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

3096: /*@
3097:   MatMPIBAIJSetPreallocation - Allocates memory for a sparse parallel matrix in `MATMPIBAIJ` format
3098:   (block compressed row).

3100:   Collective

3102:   Input Parameters:
3103: + B     - the matrix
3104: . bs    - size of block, the blocks are ALWAYS square. One can use `MatSetBlockSizes()` to set a different row and column blocksize but the row
3105:           blocksize always defines the size of the blocks. The column blocksize sets the blocksize of the vectors obtained with `MatCreateVecs()`
3106: . d_nz  - number of block nonzeros per block row in diagonal portion of local
3107:            submatrix  (same for all local rows)
3108: . d_nnz - array containing the number of block nonzeros in the various block rows
3109:            of the in diagonal portion of the local (possibly different for each block
3110:            row) or `NULL`.  If you plan to factor the matrix you must leave room for the diagonal entry and
3111:            set it even if it is zero.
3112: . o_nz  - number of block nonzeros per block row in the off-diagonal portion of local
3113:            submatrix (same for all local rows).
3114: - o_nnz - array containing the number of nonzeros in the various block rows of the
3115:            off-diagonal portion of the local submatrix (possibly different for
3116:            each block row) or `NULL`.

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

3120:   Options Database Keys:
3121: + -mat_block_size          - size of the blocks to use
3122: - -mat_use_hash_table fact - set hash table factor

3124:   Level: intermediate

3126:   Notes:
3127:   For good matrix assembly performance
3128:   the user should preallocate the matrix storage by setting the parameters
3129:   `d_nz` (or `d_nnz`) and `o_nz` (or `o_nnz`).  By setting these parameters accurately,
3130:   performance can be increased by more than a factor of 50.

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

3135:   Storage Information:
3136:   For a square global matrix we define each processor's diagonal portion
3137:   to be its local rows and the corresponding columns (a square submatrix);
3138:   each processor's off-diagonal portion encompasses the remainder of the
3139:   local matrix (a rectangular submatrix).

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

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

3150: .vb
3151:            0 1 2 3 4 5 6 7 8 9 10 11
3152:           --------------------------
3153:    row 3  |o o o d d d o o o o  o  o
3154:    row 4  |o o o d d d o o o o  o  o
3155:    row 5  |o o o d d d o o o o  o  o
3156:           --------------------------
3157: .ve

3159:   Thus, any entries in the d locations are stored in the d (diagonal)
3160:   submatrix, and any entries in the o locations are stored in the
3161:   o (off-diagonal) submatrix.  Note that the d and the o submatrices are
3162:   stored simply in the `MATSEQBAIJ` format for compressed row storage.

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

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

3174: .seealso: `Mat`, `MATMPIBAIJ`, `MatCreate()`, `MatCreateSeqBAIJ()`, `MatSetValues()`, `MatCreateBAIJ()`, `MatMPIBAIJSetPreallocationCSR()`, `PetscSplitOwnership()`
3175: @*/
3176: PetscErrorCode MatMPIBAIJSetPreallocation(Mat B, PetscInt bs, PetscInt d_nz, const PetscInt d_nnz[], PetscInt o_nz, const PetscInt o_nnz[])
3177: {
3178:   PetscFunctionBegin;
3182:   PetscTryMethod(B, "MatMPIBAIJSetPreallocation_C", (Mat, PetscInt, PetscInt, const PetscInt[], PetscInt, const PetscInt[]), (B, bs, d_nz, d_nnz, o_nz, o_nnz));
3183:   PetscFunctionReturn(PETSC_SUCCESS);
3184: }

3186: // PetscClangLinter pragma disable: -fdoc-section-header-unknown
3187: /*@
3188:   MatCreateBAIJ - Creates a sparse parallel matrix in `MATBAIJ` format
3189:   (block compressed row).

3191:   Collective

3193:   Input Parameters:
3194: + comm  - MPI communicator
3195: . bs    - size of block, the blocks are ALWAYS square. One can use `MatSetBlockSizes()` to set a different row and column blocksize but the row
3196:           blocksize always defines the size of the blocks. The column blocksize sets the blocksize of the vectors obtained with `MatCreateVecs()`
3197: . m     - number of local rows (or `PETSC_DECIDE` to have calculated if M is given)
3198:           This value should be the same as the local size used in creating the
3199:           y vector for the matrix-vector product y = Ax.
3200: . n     - number of local columns (or `PETSC_DECIDE` to have calculated if N is given)
3201:           This value should be the same as the local size used in creating the
3202:           x vector for the matrix-vector product y = Ax.
3203: . M     - number of global rows (or `PETSC_DETERMINE` to have calculated if m is given)
3204: . N     - number of global columns (or `PETSC_DETERMINE` to have calculated if n is given)
3205: . d_nz  - number of nonzero blocks per block row in diagonal portion of local
3206:           submatrix  (same for all local rows)
3207: . d_nnz - array containing the number of nonzero blocks in the various block rows
3208:           of the in diagonal portion of the local (possibly different for each block
3209:           row) or NULL.  If you plan to factor the matrix you must leave room for the diagonal entry
3210:           and set it even if it is zero.
3211: . o_nz  - number of nonzero blocks per block row in the off-diagonal portion of local
3212:           submatrix (same for all local rows).
3213: - o_nnz - array containing the number of nonzero blocks in the various block rows of the
3214:           off-diagonal portion of the local submatrix (possibly different for
3215:           each block row) or NULL.

3217:   Output Parameter:
3218: . A - the matrix

3220:   Options Database Keys:
3221: + -mat_block_size          - size of the blocks to use
3222: - -mat_use_hash_table fact - set hash table factor

3224:   Level: intermediate

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

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

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

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

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

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

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

3249:   Storage Information:
3250:   For a square global matrix we define each processor's diagonal portion
3251:   to be its local rows and the corresponding columns (a square submatrix);
3252:   each processor's off-diagonal portion encompasses the remainder of the
3253:   local matrix (a rectangular submatrix).

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

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

3264: .vb
3265:            0 1 2 3 4 5 6 7 8 9 10 11
3266:           --------------------------
3267:    row 3  |o o o d d d o o o o  o  o
3268:    row 4  |o o o d d d o o o o  o  o
3269:    row 5  |o o o d d d o o o o  o  o
3270:           --------------------------
3271: .ve

3273:   Thus, any entries in the d locations are stored in the d (diagonal)
3274:   submatrix, and any entries in the o locations are stored in the
3275:   o (off-diagonal) submatrix.  Note that the d and the o submatrices are
3276:   stored simply in the `MATSEQBAIJ` format for compressed row storage.

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

3283: .seealso: `Mat`, `MatCreate()`, `MatCreateSeqBAIJ()`, `MatSetValues()`, `MatMPIBAIJSetPreallocation()`, `MatMPIBAIJSetPreallocationCSR()`,
3284:           `MatGetOwnershipRange()`, `MatGetOwnershipRanges()`, `MatGetOwnershipRangeColumn()`, `MatGetOwnershipRangesColumn()`, `PetscLayout`
3285: @*/
3286: 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)
3287: {
3288:   PetscMPIInt size;

3290:   PetscFunctionBegin;
3291:   PetscCall(MatCreate(comm, A));
3292:   PetscCall(MatSetSizes(*A, m, n, M, N));
3293:   PetscCallMPI(MPI_Comm_size(comm, &size));
3294:   if (size > 1) {
3295:     PetscCall(MatSetType(*A, MATMPIBAIJ));
3296:     PetscCall(MatMPIBAIJSetPreallocation(*A, bs, d_nz, d_nnz, o_nz, o_nnz));
3297:   } else {
3298:     PetscCall(MatSetType(*A, MATSEQBAIJ));
3299:     PetscCall(MatSeqBAIJSetPreallocation(*A, bs, d_nz, d_nnz));
3300:   }
3301:   PetscFunctionReturn(PETSC_SUCCESS);
3302: }

3304: static PetscErrorCode MatDuplicate_MPIBAIJ(Mat matin, MatDuplicateOption cpvalues, Mat *newmat)
3305: {
3306:   Mat          mat;
3307:   Mat_MPIBAIJ *a, *oldmat = (Mat_MPIBAIJ *)matin->data;
3308:   PetscInt     len = 0;

3310:   PetscFunctionBegin;
3311:   *newmat = NULL;
3312:   PetscCall(MatCreate(PetscObjectComm((PetscObject)matin), &mat));
3313:   PetscCall(MatSetSizes(mat, matin->rmap->n, matin->cmap->n, matin->rmap->N, matin->cmap->N));
3314:   PetscCall(MatSetType(mat, ((PetscObject)matin)->type_name));
3315:   PetscCall(MatSetOption(mat, MAT_STRUCTURE_ONLY, matin->structure_only));

3317:   PetscCall(PetscLayoutReference(matin->rmap, &mat->rmap));
3318:   PetscCall(PetscLayoutReference(matin->cmap, &mat->cmap));
3319:   if (matin->hash_active) PetscCall(MatSetUp(mat));
3320:   else {
3321:     mat->factortype   = matin->factortype;
3322:     mat->preallocated = PETSC_TRUE;
3323:     mat->assembled    = PETSC_TRUE;
3324:     mat->insertmode   = NOT_SET_VALUES;

3326:     a             = (Mat_MPIBAIJ *)mat->data;
3327:     mat->rmap->bs = matin->rmap->bs;
3328:     a->bs2        = oldmat->bs2;
3329:     a->mbs        = oldmat->mbs;
3330:     a->nbs        = oldmat->nbs;
3331:     a->Mbs        = oldmat->Mbs;
3332:     a->Nbs        = oldmat->Nbs;

3334:     a->size         = oldmat->size;
3335:     a->rank         = oldmat->rank;
3336:     a->donotstash   = oldmat->donotstash;
3337:     a->roworiented  = oldmat->roworiented;
3338:     a->rowindices   = NULL;
3339:     a->rowvalues    = NULL;
3340:     a->getrowactive = PETSC_FALSE;
3341:     a->barray       = NULL;
3342:     a->rstartbs     = oldmat->rstartbs;
3343:     a->rendbs       = oldmat->rendbs;
3344:     a->cstartbs     = oldmat->cstartbs;
3345:     a->cendbs       = oldmat->cendbs;

3347:     /* hash table stuff */
3348:     a->ht           = NULL;
3349:     a->hd           = NULL;
3350:     a->ht_size      = 0;
3351:     a->ht_flag      = oldmat->ht_flag;
3352:     a->ht_fact      = oldmat->ht_fact;
3353:     a->ht_total_ct  = 0;
3354:     a->ht_insert_ct = 0;

3356:     PetscCall(PetscArraycpy(a->rangebs, oldmat->rangebs, a->size + 1));
3357:     if (oldmat->colmap) {
3358: #if PetscDefined(USE_CTABLE)
3359:       PetscCall(PetscHMapIDuplicate(oldmat->colmap, &a->colmap));
3360: #else
3361:       PetscCall(PetscMalloc1(a->Nbs, &a->colmap));
3362:       PetscCall(PetscArraycpy(a->colmap, oldmat->colmap, a->Nbs));
3363: #endif
3364:     } else a->colmap = NULL;

3366:     if (oldmat->garray && (len = ((Mat_SeqBAIJ *)oldmat->B->data)->nbs)) {
3367:       PetscCall(PetscMalloc1(len, &a->garray));
3368:       PetscCall(PetscArraycpy(a->garray, oldmat->garray, len));
3369:     } else a->garray = NULL;

3371:     PetscCall(MatStashCreate_Private(PetscObjectComm((PetscObject)matin), matin->rmap->bs, &mat->bstash));
3372:     PetscCall(VecDuplicate(oldmat->lvec, &a->lvec));
3373:     PetscCall(VecScatterCopy(oldmat->Mvctx, &a->Mvctx));

3375:     PetscCall(MatDuplicate(oldmat->A, cpvalues, &a->A));
3376:     PetscCall(MatDuplicate(oldmat->B, cpvalues, &a->B));
3377:   }
3378:   PetscCall(PetscFunctionListDuplicate(((PetscObject)matin)->qlist, &((PetscObject)mat)->qlist));
3379:   *newmat = mat;
3380:   PetscFunctionReturn(PETSC_SUCCESS);
3381: }

3383: /* Used for both MPIBAIJ and MPISBAIJ matrices */
3384: PetscErrorCode MatLoad_MPIBAIJ_Binary(Mat mat, PetscViewer viewer)
3385: {
3386:   PetscInt     header[4], M, N, nz, bs, m, n, mbs, nbs, rows, cols, sum, i, j, k;
3387:   PetscInt    *rowidxs, *colidxs, rs, cs, ce;
3388:   PetscScalar *matvals;
3389:   PetscBool    nooffprocentries = mat->nooffprocentries;

3391:   PetscFunctionBegin;
3392:   PetscCall(PetscViewerSetUp(viewer));

3394:   /* read in matrix header */
3395:   PetscCall(PetscViewerBinaryRead(viewer, header, 4, NULL, PETSC_INT));
3396:   PetscCheck(header[0] == MAT_FILE_CLASSID, PetscObjectComm((PetscObject)viewer), PETSC_ERR_FILE_UNEXPECTED, "Not a matrix object in file");
3397:   M  = header[1];
3398:   N  = header[2];
3399:   nz = header[3];
3400:   PetscCheck(M >= 0, PetscObjectComm((PetscObject)viewer), PETSC_ERR_FILE_UNEXPECTED, "Matrix row size (%" PetscInt_FMT ") in file is negative", M);
3401:   PetscCheck(N >= 0, PetscObjectComm((PetscObject)viewer), PETSC_ERR_FILE_UNEXPECTED, "Matrix column size (%" PetscInt_FMT ") in file is negative", N);
3402:   PetscCheck(nz >= 0, PETSC_COMM_SELF, PETSC_ERR_FILE_UNEXPECTED, "Matrix stored in special format on disk, cannot load as MPIBAIJ");

3404:   /* set block sizes from the viewer's .info file */
3405:   PetscCall(MatLoad_Binary_BlockSizes(mat, viewer));
3406:   /* set local sizes if not set already */
3407:   if (mat->rmap->n < 0 && M == N) mat->rmap->n = mat->cmap->n;
3408:   if (mat->cmap->n < 0 && M == N) mat->cmap->n = mat->rmap->n;
3409:   /* set global sizes if not set already */
3410:   if (mat->rmap->N < 0) mat->rmap->N = M;
3411:   if (mat->cmap->N < 0) mat->cmap->N = N;
3412:   PetscCall(PetscLayoutSetUp(mat->rmap));
3413:   PetscCall(PetscLayoutSetUp(mat->cmap));

3415:   /* check if the matrix sizes are correct */
3416:   PetscCall(MatGetSize(mat, &rows, &cols));
3417:   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);
3418:   PetscCall(MatGetBlockSize(mat, &bs));
3419:   PetscCall(MatGetLocalSize(mat, &m, &n));
3420:   PetscCall(PetscLayoutGetRange(mat->rmap, &rs, NULL));
3421:   PetscCall(PetscLayoutGetRange(mat->cmap, &cs, &ce));
3422:   mbs = m / bs;
3423:   nbs = n / bs;

3425:   /* read in row lengths and build row indices */
3426:   PetscCall(PetscMalloc1(m + 1, &rowidxs));
3427:   PetscCall(PetscViewerBinaryReadAll(viewer, rowidxs + 1, m, PETSC_DECIDE, M, PETSC_INT));
3428:   rowidxs[0] = 0;
3429:   for (i = 0; i < m; i++) rowidxs[i + 1] += rowidxs[i];
3430:   PetscCallMPI(MPIU_Allreduce(&rowidxs[m], &sum, 1, MPIU_INT, MPI_SUM, PetscObjectComm((PetscObject)viewer)));
3431:   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);

3433:   /* read in column indices and matrix values */
3434:   PetscCall(PetscMalloc2(rowidxs[m], &colidxs, rowidxs[m], &matvals));
3435:   PetscCall(PetscViewerBinaryReadAll(viewer, colidxs, rowidxs[m], PETSC_DETERMINE, PETSC_DETERMINE, PETSC_INT));
3436:   PetscCall(PetscViewerBinaryReadAll(viewer, matvals, rowidxs[m], PETSC_DETERMINE, PETSC_DETERMINE, PETSC_SCALAR));

3438:   {                /* preallocate matrix storage */
3439:     PetscBT    bt; /* helper bit set to count diagonal nonzeros */
3440:     PetscHSetI ht; /* helper hash set to count off-diagonal nonzeros */
3441:     PetscBool  sbaij, done;
3442:     PetscInt  *d_nnz, *o_nnz;

3444:     PetscCall(PetscBTCreate(nbs, &bt));
3445:     PetscCall(PetscHSetICreate(&ht));
3446:     PetscCall(PetscCalloc2(mbs, &d_nnz, mbs, &o_nnz));
3447:     PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATMPISBAIJ, &sbaij));
3448:     for (i = 0; i < mbs; i++) {
3449:       PetscCall(PetscBTMemzero(nbs, bt));
3450:       PetscCall(PetscHSetIClear(ht));
3451:       for (k = 0; k < bs; k++) {
3452:         const PetscInt row = bs * i + k;

3454:         for (j = rowidxs[row]; j < rowidxs[row + 1]; j++) {
3455:           const PetscInt col = colidxs[j];

3457:           if (!sbaij || col / bs >= rs / bs + i) {
3458:             if (col >= cs && col < ce) {
3459:               if (!PetscBTLookupSet(bt, (col - cs) / bs)) d_nnz[i]++;
3460:             } else {
3461:               PetscCall(PetscHSetIQueryAdd(ht, col / bs, &done));
3462:               if (done) o_nnz[i]++;
3463:             }
3464:           }
3465:         }
3466:       }
3467:     }
3468:     PetscCall(PetscBTDestroy(&bt));
3469:     PetscCall(PetscHSetIDestroy(&ht));
3470:     PetscCall(MatMPIBAIJSetPreallocation(mat, bs, 0, d_nnz, 0, o_nnz));
3471:     PetscCall(MatMPISBAIJSetPreallocation(mat, bs, 0, d_nnz, 0, o_nnz));
3472:     PetscCall(PetscFree2(d_nnz, o_nnz));
3473:   }

3475:   /* store matrix values */
3476:   for (i = 0; i < m; i++) {
3477:     PetscInt row = rs + i, s = rowidxs[i], e = rowidxs[i + 1];
3478:     PetscUseTypeMethod(mat, setvalues, 1, &row, e - s, colidxs + s, matvals + s, INSERT_VALUES);
3479:   }

3481:   PetscCall(PetscFree(rowidxs));
3482:   PetscCall(PetscFree2(colidxs, matvals));
3483:   mat->nooffprocentries = PETSC_TRUE;
3484:   PetscCall(MatAssemblyBegin(mat, MAT_FINAL_ASSEMBLY));
3485:   PetscCall(MatAssemblyEnd(mat, MAT_FINAL_ASSEMBLY));
3486:   mat->nooffprocentries = nooffprocentries;
3487:   PetscFunctionReturn(PETSC_SUCCESS);
3488: }

3490: PetscErrorCode MatLoad_MPIBAIJ(Mat mat, PetscViewer viewer)
3491: {
3492:   PetscBool isbinary;

3494:   PetscFunctionBegin;
3495:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &isbinary));
3496:   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);
3497:   PetscCall(MatLoad_MPIBAIJ_Binary(mat, viewer));
3498:   PetscFunctionReturn(PETSC_SUCCESS);
3499: }

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

3504:   Input Parameters:
3505: + mat  - the matrix
3506: - fact - factor

3508:   Options Database Key:
3509: . -mat_use_hash_table fact - provide the factor

3511:   Level: advanced

3513: .seealso: `Mat`, `MATMPIBAIJ`, `MatSetOption()`
3514: @*/
3515: PetscErrorCode MatMPIBAIJSetHashTableFactor(Mat mat, PetscReal fact)
3516: {
3517:   PetscFunctionBegin;
3518:   PetscTryMethod(mat, "MatSetHashTableFactor_C", (Mat, PetscReal), (mat, fact));
3519:   PetscFunctionReturn(PETSC_SUCCESS);
3520: }

3522: PetscErrorCode MatSetHashTableFactor_MPIBAIJ(Mat mat, PetscReal fact)
3523: {
3524:   Mat_MPIBAIJ *baij;

3526:   PetscFunctionBegin;
3527:   baij          = (Mat_MPIBAIJ *)mat->data;
3528:   baij->ht_fact = fact;
3529:   PetscFunctionReturn(PETSC_SUCCESS);
3530: }

3532: /*@
3533:   MatMPIBAIJGetSeqBAIJ - Get the on-process (diagonal block) and off-process (off-diagonal block) sequential matrices
3534:   that make up a `MATMPIBAIJ` or `MATMPISBAIJ` matrix, together with the local-to-global column map for the off-diagonal block.

3536:   Not Collective

3538:   Input Parameter:
3539: . A - the `MATMPIBAIJ` or `MATMPISBAIJ` matrix

3541:   Output Parameters:
3542: + Ad     - the diagonal block (`MATSEQBAIJ` or `MATSEQSBAIJ`), or `NULL` if not needed
3543: . Ao     - the off-diagonal block `MATSEQBAIJ`, or `NULL` if not needed
3544: - colmap - the local-to-global column index map for `Ao`, or `NULL` if not needed

3546:   Level: advanced

3548: .seealso: `Mat`, `MATMPIBAIJ`, `MATMPISBAIJ`, `MATSEQBAIJ`, `MATSEQSBAIJ`, `MatMPIAIJGetSeqAIJ()`
3549: @*/
3550: PetscErrorCode MatMPIBAIJGetSeqBAIJ(Mat A, Mat *Ad, Mat *Ao, const PetscInt *colmap[])
3551: {
3552:   Mat_MPIBAIJ *a = (Mat_MPIBAIJ *)A->data;
3553:   PetscBool    flg;

3555:   PetscFunctionBegin;
3556:   PetscCall(PetscObjectTypeCompareAny((PetscObject)A, &flg, MATMPIBAIJ, MATMPISBAIJ, ""));
3557:   PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "This function requires a MATMPIBAIJ or MATMPISBAIJ matrix as input");
3558:   if (Ad) *Ad = a->A;
3559:   if (Ao) *Ao = a->B;
3560:   if (colmap) *colmap = a->garray;
3561:   PetscFunctionReturn(PETSC_SUCCESS);
3562: }

3564: /*
3565:     Special version for direct calls from Fortran (to eliminate two function call overheads
3566: */
3567: #if PetscDefined(HAVE_FORTRAN_CAPS)
3568:   #define matmpibaijsetvaluesblocked_ MATMPIBAIJSETVALUESBLOCKED
3569: #elif !PetscDefined(HAVE_FORTRAN_UNDERSCORE)
3570:   #define matmpibaijsetvaluesblocked_ matmpibaijsetvaluesblocked
3571: #endif

3573: // PetscClangLinter pragma disable: -fdoc-synopsis-matching-symbol-name
3574: /*@
3575:   MatMPIBAIJSetValuesBlocked - Direct Fortran call to replace call to `MatSetValuesBlocked()`

3577:   Collective

3579:   Input Parameters:
3580: + matin  - the matrix
3581: . min    - number of input rows
3582: . im     - input rows
3583: . nin    - number of input columns
3584: . in     - input columns
3585: . v      - numerical values input
3586: - addvin - `INSERT_VALUES` or `ADD_VALUES`

3588:   Level: advanced

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

3593: .seealso: `Mat`, `MatSetValuesBlocked()`
3594: @*/
3595: PETSC_EXTERN PetscErrorCode matmpibaijsetvaluesblocked_(Mat *matin, PetscInt *min, const PetscInt im[], PetscInt *nin, const PetscInt in[], const MatScalar v[], InsertMode *addvin)
3596: {
3597:   /* convert input arguments to C version */
3598:   Mat        mat = *matin;
3599:   PetscInt   m = *min, n = *nin;
3600:   InsertMode addv = *addvin;

3602:   Mat_MPIBAIJ     *baij = (Mat_MPIBAIJ *)mat->data;
3603:   const MatScalar *value;
3604:   MatScalar       *barray      = baij->barray;
3605:   PetscBool        roworiented = baij->roworiented;
3606:   PetscInt         i, j, ii, jj, row, col, rstart = baij->rstartbs;
3607:   PetscInt         rend = baij->rendbs, cstart = baij->cstartbs, stepval;
3608:   PetscInt         cend = baij->cendbs, bs = mat->rmap->bs, bs2 = baij->bs2;

3610:   PetscFunctionBegin;
3611:   /* tasks normally handled by MatSetValuesBlocked() */
3612:   if (mat->insertmode == NOT_SET_VALUES) mat->insertmode = addv;
3613:   else PetscCheck(mat->insertmode == addv, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Cannot mix add values and insert values");
3614:   PetscCheck(!mat->factortype, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Not for factored matrix");
3615:   if (mat->assembled) {
3616:     mat->was_assembled = PETSC_TRUE;
3617:     mat->assembled     = PETSC_FALSE;
3618:   }
3619:   PetscCall(PetscLogEventBegin(MAT_SetValues, mat, 0, 0, 0));

3621:   if (!barray) {
3622:     PetscCall(PetscMalloc1(bs2, &barray));
3623:     baij->barray = barray;
3624:   }

3626:   if (roworiented) stepval = (n - 1) * bs;
3627:   else stepval = (m - 1) * bs;

3629:   for (i = 0; i < m; i++) {
3630:     if (im[i] < 0) continue;
3631:     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);
3632:     if (im[i] >= rstart && im[i] < rend) {
3633:       row = im[i] - rstart;
3634:       for (j = 0; j < n; j++) {
3635:         /* If NumCol = 1 then a copy is not required */
3636:         if (roworiented && (n == 1)) {
3637:           barray = (MatScalar *)v + i * bs2;
3638:         } else if ((!roworiented) && (m == 1)) {
3639:           barray = (MatScalar *)v + j * bs2;
3640:         } else { /* Here a copy is required */
3641:           if (roworiented) {
3642:             value = v + i * (stepval + bs) * bs + j * bs;
3643:           } else {
3644:             value = v + j * (stepval + bs) * bs + i * bs;
3645:           }
3646:           for (ii = 0; ii < bs; ii++, value += stepval) {
3647:             for (jj = 0; jj < bs; jj++) *barray++ = *value++;
3648:           }
3649:           barray -= bs2;
3650:         }

3652:         if (in[j] >= cstart && in[j] < cend) {
3653:           col = in[j] - cstart;
3654:           PetscCall(MatSetValuesBlocked_SeqBAIJ_Inlined(baij->A, row, col, barray, addv, im[i], in[j]));
3655:         } else if (in[j] < 0) {
3656:           continue;
3657:         } else {
3658:           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);
3659:           if (mat->was_assembled) {
3660:             if (!baij->colmap) PetscCall(MatCreateColmap_MPIBAIJ_Private(mat));

3662: #if PetscDefined(USE_CTABLE)
3663:             if (PetscDefined(USE_DEBUG)) {
3664:               PetscInt data;
3665:               PetscCall(PetscHMapIGetWithDefault(baij->colmap, in[j] + 1, 0, &data));
3666:               PetscCheck((data - 1) % bs == 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Incorrect colmap");
3667:             }
3668: #else
3669:             if (PetscDefined(USE_DEBUG)) PetscCheck((baij->colmap[in[j]] - 1) % bs == 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Incorrect colmap");
3670: #endif
3671: #if PetscDefined(USE_CTABLE)
3672:             PetscCall(PetscHMapIGetWithDefault(baij->colmap, in[j] + 1, 0, &col));
3673:             col = (col - 1) / bs;
3674: #else
3675:             col = (baij->colmap[in[j]] - 1) / bs;
3676: #endif
3677:             if (col < 0 && !((Mat_SeqBAIJ *)baij->A->data)->nonew) {
3678:               PetscCall(MatDisAssemble_MPIBAIJ(mat));
3679:               col = in[j];
3680:             }
3681:           } else col = in[j];
3682:           PetscCall(MatSetValuesBlocked_SeqBAIJ_Inlined(baij->B, row, col, barray, addv, im[i], in[j]));
3683:         }
3684:       }
3685:     } else {
3686:       if (!baij->donotstash) {
3687:         if (roworiented) {
3688:           PetscCall(MatStashValuesRowBlocked_Private(&mat->bstash, im[i], n, in, v, m, n, i));
3689:         } else {
3690:           PetscCall(MatStashValuesColBlocked_Private(&mat->bstash, im[i], n, in, v, m, n, i));
3691:         }
3692:       }
3693:     }
3694:   }

3696:   /* task normally handled by MatSetValuesBlocked() */
3697:   PetscCall(PetscLogEventEnd(MAT_SetValues, mat, 0, 0, 0));
3698:   PetscFunctionReturn(PETSC_SUCCESS);
3699: }

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

3704:   Collective

3706:   Input Parameters:
3707: + comm - MPI communicator
3708: . bs   - the block size, only a block size of 1 is supported
3709: . m    - number of local rows (Cannot be `PETSC_DECIDE`)
3710: . n    - This value should be the same as the local size used in creating the
3711:          x vector for the matrix-vector product $ y = Ax $. (or `PETSC_DECIDE` to have
3712:          calculated if `N` is given) For square matrices `n` is almost always `m`.
3713: . M    - number of global rows (or `PETSC_DETERMINE` to have calculated if `m` is given)
3714: . N    - number of global columns (or `PETSC_DETERMINE` to have calculated if `n` is given)
3715: . 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
3716: . j    - column indices
3717: - a    - matrix values

3719:   Output Parameter:
3720: . mat - the matrix

3722:   Level: intermediate

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

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

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

3736: .seealso: `Mat`, `MatCreate()`, `MatCreateSeqAIJ()`, `MatSetValues()`, `MatMPIAIJSetPreallocation()`, `MatMPIAIJSetPreallocationCSR()`,
3737:           `MATMPIAIJ`, `MatCreateAIJ()`, `MatCreateMPIAIJWithSplitArrays()`
3738: @*/
3739: 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)
3740: {
3741:   PetscFunctionBegin;
3742:   PetscCheck(!i[0], PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "i (row indices) must start with 0");
3743:   PetscCheck(m >= 0, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "local number of rows (m) cannot be PETSC_DECIDE, or negative");
3744:   PetscCall(MatCreate(comm, mat));
3745:   PetscCall(MatSetSizes(*mat, m, n, M, N));
3746:   PetscCall(MatSetType(*mat, MATMPIBAIJ));
3747:   PetscCall(MatSetBlockSize(*mat, bs));
3748:   PetscCall(MatSetUp(*mat));
3749:   PetscCall(MatSetOption(*mat, MAT_ROW_ORIENTED, PETSC_FALSE));
3750:   PetscCall(MatMPIBAIJSetPreallocationCSR(*mat, bs, i, j, a));
3751:   PetscCall(MatSetOption(*mat, MAT_ROW_ORIENTED, PETSC_TRUE));
3752:   PetscFunctionReturn(PETSC_SUCCESS);
3753: }

3755: PetscErrorCode MatCreateMPIMatConcatenateSeqMat_MPIBAIJ(MPI_Comm comm, Mat inmat, PetscInt n, MatReuse scall, Mat *outmat)
3756: {
3757:   PetscInt     m, N, i, rstart, nnz, Ii, bs, cbs;
3758:   PetscInt    *indx;
3759:   PetscScalar *values;

3761:   PetscFunctionBegin;
3762:   PetscCall(MatGetSize(inmat, &m, &N));
3763:   if (scall == MAT_INITIAL_MATRIX) { /* symbolic phase */
3764:     Mat_SeqBAIJ *a = (Mat_SeqBAIJ *)inmat->data;
3765:     PetscInt    *dnz, *onz, mbs, Nbs, nbs;
3766:     PetscInt    *bindx, rmax = a->rmax, j;
3767:     PetscMPIInt  rank, size;

3769:     PetscCall(MatGetBlockSizes(inmat, &bs, &cbs));
3770:     mbs = m / bs;
3771:     Nbs = N / cbs;
3772:     if (n == PETSC_DECIDE) PetscCall(PetscSplitOwnershipBlock(comm, cbs, &n, &N));
3773:     nbs = n / cbs;

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

3778:     PetscCallMPI(MPI_Comm_rank(comm, &rank));
3779:     PetscCallMPI(MPI_Comm_size(comm, &size));
3780:     if (rank == size - 1) {
3781:       /* Check sum(nbs) = Nbs */
3782:       PetscCheck(__end == Nbs, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Sum of local block columns %" PetscInt_FMT " != global block columns %" PetscInt_FMT, __end, Nbs);
3783:     }

3785:     rstart = __rstart; /* block rstart of *outmat; see inline function MatPreallocateBegin */
3786:     for (i = 0; i < mbs; i++) {
3787:       PetscCall(MatGetRow_SeqBAIJ(inmat, i * bs, &nnz, &indx, NULL)); /* non-blocked nnz and indx */
3788:       nnz = nnz / bs;
3789:       for (j = 0; j < nnz; j++) bindx[j] = indx[j * bs] / bs;
3790:       PetscCall(MatPreallocateSet(i + rstart, nnz, bindx, dnz, onz));
3791:       PetscCall(MatRestoreRow_SeqBAIJ(inmat, i * bs, &nnz, &indx, NULL));
3792:     }
3793:     PetscCall(PetscFree(bindx));

3795:     PetscCall(MatCreate(comm, outmat));
3796:     PetscCall(MatSetSizes(*outmat, m, n, PETSC_DETERMINE, PETSC_DETERMINE));
3797:     PetscCall(MatSetBlockSizes(*outmat, bs, cbs));
3798:     PetscCall(MatSetType(*outmat, MATBAIJ));
3799:     PetscCall(MatSeqBAIJSetPreallocation(*outmat, bs, 0, dnz));
3800:     PetscCall(MatMPIBAIJSetPreallocation(*outmat, bs, 0, dnz, 0, onz));
3801:     MatPreallocateEnd(dnz, onz);
3802:     PetscCall(MatSetOption(*outmat, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
3803:   }

3805:   /* numeric phase */
3806:   PetscCall(MatGetBlockSizes(inmat, &bs, &cbs));
3807:   PetscCall(MatGetOwnershipRange(*outmat, &rstart, NULL));

3809:   for (i = 0; i < m; i++) {
3810:     PetscCall(MatGetRow_SeqBAIJ(inmat, i, &nnz, &indx, &values));
3811:     Ii = i + rstart;
3812:     PetscCall(MatSetValues(*outmat, 1, &Ii, nnz, indx, values, INSERT_VALUES));
3813:     PetscCall(MatRestoreRow_SeqBAIJ(inmat, i, &nnz, &indx, &values));
3814:   }
3815:   PetscCall(MatAssemblyBegin(*outmat, MAT_FINAL_ASSEMBLY));
3816:   PetscCall(MatAssemblyEnd(*outmat, MAT_FINAL_ASSEMBLY));
3817:   PetscFunctionReturn(PETSC_SUCCESS);
3818: }