Actual source code: maij.c

  1: #include <../src/mat/impls/maij/maij.h>
  2: #include <../src/mat/utils/freespace.h>

  4: /*@
  5:   MatMAIJGetAIJ - Get the `MATAIJ` matrix describing the blockwise action of the `MATMAIJ` matrix

  7:   Not Collective, but if the `MATMAIJ` matrix is parallel, the `MATAIJ` matrix is also parallel

  9:   Input Parameter:
 10: . A - the `MATMAIJ` matrix

 12:   Output Parameter:
 13: . B - the `MATAIJ` matrix

 15:   Level: advanced

 17:   Note:
 18:   The reference count on the `MATAIJ` matrix is not increased so you should not destroy it.

 20: .seealso: [](ch_matrices), `Mat`, `MATMAIJ`, `MATAIJ`, `MatCreateMAIJ()`
 21: @*/
 22: PetscErrorCode MatMAIJGetAIJ(Mat A, Mat *B)
 23: {
 24:   PetscBool ismpimaij, isseqmaij;

 26:   PetscFunctionBegin;
 27:   PetscCall(PetscObjectTypeCompare((PetscObject)A, MATMPIMAIJ, &ismpimaij));
 28:   PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQMAIJ, &isseqmaij));
 29:   if (ismpimaij) {
 30:     Mat_MPIMAIJ *b = (Mat_MPIMAIJ *)A->data;

 32:     *B = b->A;
 33:   } else if (isseqmaij) {
 34:     Mat_SeqMAIJ *b = (Mat_SeqMAIJ *)A->data;

 36:     *B = b->AIJ;
 37:   } else {
 38:     *B = A;
 39:   }
 40:   PetscFunctionReturn(PETSC_SUCCESS);
 41: }

 43: /*@
 44:   MatMAIJRedimension - Get a new `MATMAIJ` matrix with the same action, but for a different block size

 46:   Logically Collective

 48:   Input Parameters:
 49: + A   - the `MATMAIJ` matrix
 50: - dof - the block size for the new matrix

 52:   Output Parameter:
 53: . B - the new `MATMAIJ` matrix

 55:   Level: advanced

 57: .seealso: [](ch_matrices), `Mat`, `MATMAIJ`, `MatCreateMAIJ()`
 58: @*/
 59: PetscErrorCode MatMAIJRedimension(Mat A, PetscInt dof, Mat *B)
 60: {
 61:   Mat Aij = NULL;

 63:   PetscFunctionBegin;
 65:   PetscCall(MatMAIJGetAIJ(A, &Aij));
 66:   PetscCall(MatCreateMAIJ(Aij, dof, B));
 67:   PetscFunctionReturn(PETSC_SUCCESS);
 68: }

 70: static PetscErrorCode MatDestroy_SeqMAIJ(Mat A)
 71: {
 72:   Mat_SeqMAIJ *b = (Mat_SeqMAIJ *)A->data;

 74:   PetscFunctionBegin;
 75:   PetscCall(MatDestroy(&b->AIJ));
 76:   PetscCall(PetscFree(A->data));
 77:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_seqmaij_seqaijcusparse_C", NULL));
 78:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_seqmaij_seqaijkokkos_C", NULL));
 79:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_seqmaij_seqaij_C", NULL));
 80:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaij_seqmaij_C", NULL));
 81:   PetscFunctionReturn(PETSC_SUCCESS);
 82: }

 84: static PetscErrorCode MatSetUp_MAIJ(Mat A)
 85: {
 86:   PetscFunctionBegin;
 87:   SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "Must use MatCreateMAIJ() to create MAIJ matrices");
 88: }

 90: static PetscErrorCode MatView_SeqMAIJ(Mat A, PetscViewer viewer)
 91: {
 92:   Mat B;

 94:   PetscFunctionBegin;
 95:   PetscCall(MatConvert(A, MATSEQAIJ, MAT_INITIAL_MATRIX, &B));
 96:   PetscCall(MatView(B, viewer));
 97:   PetscCall(MatDestroy(&B));
 98:   PetscFunctionReturn(PETSC_SUCCESS);
 99: }

101: static PetscErrorCode MatView_MPIMAIJ(Mat A, PetscViewer viewer)
102: {
103:   Mat B;

105:   PetscFunctionBegin;
106:   PetscCall(MatConvert(A, MATMPIAIJ, MAT_INITIAL_MATRIX, &B));
107:   PetscCall(MatView(B, viewer));
108:   PetscCall(MatDestroy(&B));
109:   PetscFunctionReturn(PETSC_SUCCESS);
110: }

112: static PetscErrorCode MatDestroy_MPIMAIJ(Mat A)
113: {
114:   Mat_MPIMAIJ *b = (Mat_MPIMAIJ *)A->data;

116:   PetscFunctionBegin;
117:   PetscCall(MatDestroy(&b->AIJ));
118:   PetscCall(MatDestroy(&b->OAIJ));
119:   PetscCall(MatDestroy(&b->A));
120:   PetscCall(VecScatterDestroy(&b->ctx));
121:   PetscCall(VecDestroy(&b->w));
122:   PetscCall(PetscFree(A->data));
123:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_mpimaij_mpiaijcusparse_C", NULL));
124:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_mpimaij_mpiaijkokkos_C", NULL));
125:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_mpimaij_mpiaij_C", NULL));
126:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_mpiaij_mpimaij_C", NULL));
127:   PetscCall(PetscObjectChangeTypeName((PetscObject)A, NULL));
128:   PetscFunctionReturn(PETSC_SUCCESS);
129: }

131: /*MC
132:   MATMAIJ - MATMAIJ = "maij" - A matrix type to be used for restriction and interpolation operations for
133:   multicomponent problems, interpolating or restricting each component the same way independently.
134:   The matrix type is based on `MATSEQAIJ` for sequential matrices, and `MATMPIAIJ` for distributed matrices.

136:   Operations provided:
137: .vb
138:     MatMult()
139:     MatMultTranspose()
140:     MatMultAdd()
141:     MatMultTransposeAdd()
142: .ve

144:   Level: advanced

146: .seealso: [](ch_matrices), `Mat`, `MATAIJ`, `MatMAIJGetAIJ()`, `MatMAIJRedimension()`, `MatCreateMAIJ()`
147: M*/

149: PETSC_EXTERN PetscErrorCode MatCreate_MAIJ(Mat A)
150: {
151:   Mat_MPIMAIJ *b;
152:   PetscMPIInt  size;

154:   PetscFunctionBegin;
155:   PetscCall(PetscNew(&b));
156:   A->data = (void *)b;

158:   PetscCall(PetscMemzero(A->ops, sizeof(struct _MatOps)));

160:   A->ops->setup = MatSetUp_MAIJ;

162:   b->AIJ  = NULL;
163:   b->dof  = 0;
164:   b->OAIJ = NULL;
165:   b->ctx  = NULL;
166:   b->w    = NULL;
167:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size));
168:   if (size == 1) PetscCall(PetscObjectChangeTypeName((PetscObject)A, MATSEQMAIJ));
169:   else PetscCall(PetscObjectChangeTypeName((PetscObject)A, MATMPIMAIJ));
170:   A->preallocated = PETSC_TRUE;
171:   A->assembled    = PETSC_TRUE;
172:   PetscFunctionReturn(PETSC_SUCCESS);
173: }

175: #if PetscHasAttribute(always_inline)
176:   #define PETSC_FORCE_INLINE __attribute__((always_inline))
177: #else
178:   #define PETSC_FORCE_INLINE
179: #endif

181: #if defined(__clang__)
182:   #define PETSC_PRAGMA_UNROLL _Pragma("unroll")
183: #else
184:   #define PETSC_PRAGMA_UNROLL
185: #endif

187: enum {
188:   MAT_SEQMAIJ_MAX_TEMPLATE_SIZE = 18
189: };

191: // try as hard as possible to get these "template"s inlined, GCC apparently does take 'inline'
192: // keyword into account for these...
193: PETSC_FORCE_INLINE static inline PetscErrorCode MatMult_MatMultAdd_SeqMAIJ_Template(Mat A, Vec xx, Vec yy, Vec zz, int N)
194: {
195:   const PetscBool    mult_add   = yy == NULL ? PETSC_FALSE : PETSC_TRUE;
196:   const Mat_SeqMAIJ *b          = (Mat_SeqMAIJ *)A->data;
197:   const Mat          baij       = b->AIJ;
198:   const Mat_SeqAIJ  *a          = (Mat_SeqAIJ *)baij->data;
199:   const PetscInt     m          = baij->rmap->n;
200:   const PetscInt     nz         = a->nz;
201:   const PetscInt    *idx        = a->j;
202:   const PetscInt    *ii         = a->i;
203:   const PetscScalar *v          = a->a;
204:   PetscInt           nonzerorow = 0;
205:   const PetscScalar *x;
206:   PetscScalar       *z;

208:   PetscFunctionBegin;
209:   PetscAssert(N <= MAT_SEQMAIJ_MAX_TEMPLATE_SIZE, PETSC_COMM_SELF, PETSC_ERR_PLIB, "%s() called with N = %d > max size %d", PETSC_FUNCTION_NAME, N, MAT_SEQMAIJ_MAX_TEMPLATE_SIZE);
210:   if (mult_add && yy != zz) PetscCall(VecCopy(yy, zz));
211:   PetscCall(VecGetArrayRead(xx, &x));
212:   if (mult_add) {
213:     PetscCall(VecGetArray(zz, &z));
214:   } else {
215:     PetscCall(VecGetArrayWrite(zz, &z));
216:   }

218:   for (PetscInt i = 0; i < m; ++i) {
219:     PetscInt       jrow = ii[i];
220:     const PetscInt n    = ii[i + 1] - jrow;
221:     // leave a line so clang-format does not align these decls
222:     PetscScalar sum[MAT_SEQMAIJ_MAX_TEMPLATE_SIZE] = {0};

224:     nonzerorow += n > 0;
225:     for (PetscInt j = 0; j < n; ++j, ++jrow) {
226:       const PetscScalar v_jrow     = v[jrow];
227:       const PetscInt    N_idx_jrow = N * idx[jrow];

229:       PETSC_PRAGMA_UNROLL
230:       for (int k = 0; k < N; ++k) sum[k] += v_jrow * x[N_idx_jrow + k];
231:     }

233:     PETSC_PRAGMA_UNROLL
234:     for (int k = 0; k < N; ++k) {
235:       const PetscInt z_idx = N * i + k;

237:       if (mult_add) {
238:         z[z_idx] += sum[k];
239:       } else {
240:         z[z_idx] = sum[k];
241:       }
242:     }
243:   }
244:   PetscCall(PetscLogFlops(2 * N * nz - (mult_add ? 0 : (N * nonzerorow))));
245:   PetscCall(VecRestoreArrayRead(xx, &x));
246:   if (mult_add) {
247:     PetscCall(VecRestoreArray(zz, &z));
248:   } else {
249:     PetscCall(VecRestoreArrayWrite(zz, &z));
250:   }
251:   PetscFunctionReturn(PETSC_SUCCESS);
252: }

254: PETSC_FORCE_INLINE static inline PetscErrorCode MatMultTranspose_MatMultTransposeAdd_SeqMAIJ_Template(Mat A, Vec xx, Vec yy, Vec zz, int N)
255: {
256:   const PetscBool    mult_add = yy == NULL ? PETSC_FALSE : PETSC_TRUE;
257:   const Mat_SeqMAIJ *b        = (Mat_SeqMAIJ *)A->data;
258:   const Mat          baij     = b->AIJ;
259:   const Mat_SeqAIJ  *a        = (Mat_SeqAIJ *)baij->data;
260:   const PetscInt     m        = baij->rmap->n;
261:   const PetscInt     nz       = a->nz;
262:   const PetscInt    *a_j      = a->j;
263:   const PetscInt    *a_i      = a->i;
264:   const PetscScalar *a_a      = a->a;
265:   const PetscScalar *x;
266:   PetscScalar       *z;

268:   PetscFunctionBegin;
269:   PetscAssert(N <= MAT_SEQMAIJ_MAX_TEMPLATE_SIZE, PETSC_COMM_SELF, PETSC_ERR_PLIB, "%s() called with N = %d > max size %d", PETSC_FUNCTION_NAME, N, MAT_SEQMAIJ_MAX_TEMPLATE_SIZE);
270:   if (mult_add) {
271:     if (yy != zz) PetscCall(VecCopy(yy, zz));
272:   } else {
273:     PetscCall(VecSet(zz, 0.0));
274:   }
275:   PetscCall(VecGetArrayRead(xx, &x));
276:   PetscCall(VecGetArray(zz, &z));

278:   for (PetscInt i = 0; i < m; i++) {
279:     const PetscInt     a_ii = a_i[i];
280:     const PetscInt    *idx  = PetscSafePointerPlusOffset(a_j, a_ii);
281:     const PetscScalar *v    = PetscSafePointerPlusOffset(a_a, a_ii);
282:     const PetscInt     n    = a_i[i + 1] - a_ii;
283:     PetscScalar        alpha[MAT_SEQMAIJ_MAX_TEMPLATE_SIZE];

285:     PETSC_PRAGMA_UNROLL
286:     for (int k = 0; k < N; ++k) alpha[k] = x[N * i + k];
287:     for (PetscInt j = 0; j < n; ++j) {
288:       const PetscInt    N_idx_j = N * idx[j];
289:       const PetscScalar v_j     = v[j];

291:       PETSC_PRAGMA_UNROLL
292:       for (int k = 0; k < N; ++k) z[N_idx_j + k] += alpha[k] * v_j;
293:     }
294:   }

296:   PetscCall(PetscLogFlops(2 * N * nz));
297:   PetscCall(VecRestoreArrayRead(xx, &x));
298:   PetscCall(VecRestoreArray(zz, &z));
299:   PetscFunctionReturn(PETSC_SUCCESS);
300: }

302: #define MAT_SEQ_MAIJ_INSTANTIATE_MATMULT_MATMULTADD_TEMPLATE(N) \
303:   static PetscErrorCode PetscConcat(MatMult_SeqMAIJ_, N)(Mat A, Vec xx, Vec yy) \
304:   { \
305:     PetscFunctionBegin; \
306:     PetscCall(MatMult_MatMultAdd_SeqMAIJ_Template(A, xx, NULL, yy, N)); \
307:     PetscFunctionReturn(PETSC_SUCCESS); \
308:   } \
309:   static PetscErrorCode PetscConcat(MatMultTranspose_SeqMAIJ_, N)(Mat A, Vec xx, Vec yy) \
310:   { \
311:     PetscFunctionBegin; \
312:     PetscCall(MatMultTranspose_MatMultTransposeAdd_SeqMAIJ_Template(A, xx, NULL, yy, N)); \
313:     PetscFunctionReturn(PETSC_SUCCESS); \
314:   } \
315:   static PetscErrorCode PetscConcat(MatMultAdd_SeqMAIJ_, N)(Mat A, Vec xx, Vec yy, Vec zz) \
316:   { \
317:     PetscFunctionBegin; \
318:     PetscCall(MatMult_MatMultAdd_SeqMAIJ_Template(A, xx, yy, zz, N)); \
319:     PetscFunctionReturn(PETSC_SUCCESS); \
320:   } \
321:   static PetscErrorCode PetscConcat(MatMultTransposeAdd_SeqMAIJ_, N)(Mat A, Vec xx, Vec yy, Vec zz) \
322:   { \
323:     PetscFunctionBegin; \
324:     PetscCall(MatMultTranspose_MatMultTransposeAdd_SeqMAIJ_Template(A, xx, yy, zz, N)); \
325:     PetscFunctionReturn(PETSC_SUCCESS); \
326:   }

328: MAT_SEQ_MAIJ_INSTANTIATE_MATMULT_MATMULTADD_TEMPLATE(2)
329: MAT_SEQ_MAIJ_INSTANTIATE_MATMULT_MATMULTADD_TEMPLATE(3)
330: MAT_SEQ_MAIJ_INSTANTIATE_MATMULT_MATMULTADD_TEMPLATE(4)
331: MAT_SEQ_MAIJ_INSTANTIATE_MATMULT_MATMULTADD_TEMPLATE(5)
332: MAT_SEQ_MAIJ_INSTANTIATE_MATMULT_MATMULTADD_TEMPLATE(6)
333: MAT_SEQ_MAIJ_INSTANTIATE_MATMULT_MATMULTADD_TEMPLATE(7)
334: MAT_SEQ_MAIJ_INSTANTIATE_MATMULT_MATMULTADD_TEMPLATE(8)
335: MAT_SEQ_MAIJ_INSTANTIATE_MATMULT_MATMULTADD_TEMPLATE(9)
336: MAT_SEQ_MAIJ_INSTANTIATE_MATMULT_MATMULTADD_TEMPLATE(10)
337: MAT_SEQ_MAIJ_INSTANTIATE_MATMULT_MATMULTADD_TEMPLATE(11)
338: MAT_SEQ_MAIJ_INSTANTIATE_MATMULT_MATMULTADD_TEMPLATE(16)
339: MAT_SEQ_MAIJ_INSTANTIATE_MATMULT_MATMULTADD_TEMPLATE(18)

341: #undef MAT_SEQ_MAIJ_INSTANTIATE_MATMULT_MATMULTADD_TEMPLATE

343: static PetscErrorCode MatMult_SeqMAIJ_N(Mat A, Vec xx, Vec yy)
344: {
345:   Mat_SeqMAIJ       *b = (Mat_SeqMAIJ *)A->data;
346:   Mat_SeqAIJ        *a = (Mat_SeqAIJ *)b->AIJ->data;
347:   const PetscScalar *x, *v;
348:   PetscScalar       *y, *sums;
349:   const PetscInt     m = b->AIJ->rmap->n, *idx, *ii;
350:   PetscInt           n, i, jrow, j, dof = b->dof, k;

352:   PetscFunctionBegin;
353:   PetscCall(VecGetArrayRead(xx, &x));
354:   PetscCall(VecSet(yy, 0.0));
355:   PetscCall(VecGetArray(yy, &y));
356:   idx = a->j;
357:   v   = a->a;
358:   ii  = a->i;

360:   for (i = 0; i < m; i++) {
361:     jrow = ii[i];
362:     n    = ii[i + 1] - jrow;
363:     sums = y + dof * i;
364:     for (j = 0; j < n; j++) {
365:       for (k = 0; k < dof; k++) sums[k] += v[jrow] * x[dof * idx[jrow] + k];
366:       jrow++;
367:     }
368:   }

370:   PetscCall(PetscLogFlops(2.0 * dof * a->nz));
371:   PetscCall(VecRestoreArrayRead(xx, &x));
372:   PetscCall(VecRestoreArray(yy, &y));
373:   PetscFunctionReturn(PETSC_SUCCESS);
374: }

376: static PetscErrorCode MatMultAdd_SeqMAIJ_N(Mat A, Vec xx, Vec yy, Vec zz)
377: {
378:   Mat_SeqMAIJ       *b = (Mat_SeqMAIJ *)A->data;
379:   Mat_SeqAIJ        *a = (Mat_SeqAIJ *)b->AIJ->data;
380:   const PetscScalar *x, *v;
381:   PetscScalar       *y, *sums;
382:   const PetscInt     m = b->AIJ->rmap->n, *idx, *ii;
383:   PetscInt           n, i, jrow, j, dof = b->dof, k;

385:   PetscFunctionBegin;
386:   if (yy != zz) PetscCall(VecCopy(yy, zz));
387:   PetscCall(VecGetArrayRead(xx, &x));
388:   PetscCall(VecGetArray(zz, &y));
389:   idx = a->j;
390:   v   = a->a;
391:   ii  = a->i;

393:   for (i = 0; i < m; i++) {
394:     jrow = ii[i];
395:     n    = ii[i + 1] - jrow;
396:     sums = y + dof * i;
397:     for (j = 0; j < n; j++) {
398:       for (k = 0; k < dof; k++) sums[k] += v[jrow] * x[dof * idx[jrow] + k];
399:       jrow++;
400:     }
401:   }

403:   PetscCall(PetscLogFlops(2.0 * dof * a->nz));
404:   PetscCall(VecRestoreArrayRead(xx, &x));
405:   PetscCall(VecRestoreArray(zz, &y));
406:   PetscFunctionReturn(PETSC_SUCCESS);
407: }

409: static PetscErrorCode MatMultTranspose_SeqMAIJ_N(Mat A, Vec xx, Vec yy)
410: {
411:   Mat_SeqMAIJ       *b = (Mat_SeqMAIJ *)A->data;
412:   Mat_SeqAIJ        *a = (Mat_SeqAIJ *)b->AIJ->data;
413:   const PetscScalar *x, *v, *alpha;
414:   PetscScalar       *y;
415:   const PetscInt     m = b->AIJ->rmap->n, *idx, dof = b->dof;
416:   PetscInt           n, i, k;

418:   PetscFunctionBegin;
419:   PetscCall(VecGetArrayRead(xx, &x));
420:   PetscCall(VecSet(yy, 0.0));
421:   PetscCall(VecGetArray(yy, &y));
422:   for (i = 0; i < m; i++) {
423:     idx   = PetscSafePointerPlusOffset(a->j, a->i[i]);
424:     v     = PetscSafePointerPlusOffset(a->a, a->i[i]);
425:     n     = a->i[i + 1] - a->i[i];
426:     alpha = x + dof * i;
427:     while (n-- > 0) {
428:       for (k = 0; k < dof; k++) y[dof * (*idx) + k] += alpha[k] * (*v);
429:       idx++;
430:       v++;
431:     }
432:   }
433:   PetscCall(PetscLogFlops(2.0 * dof * a->nz));
434:   PetscCall(VecRestoreArrayRead(xx, &x));
435:   PetscCall(VecRestoreArray(yy, &y));
436:   PetscFunctionReturn(PETSC_SUCCESS);
437: }

439: static PetscErrorCode MatMultTransposeAdd_SeqMAIJ_N(Mat A, Vec xx, Vec yy, Vec zz)
440: {
441:   Mat_SeqMAIJ       *b = (Mat_SeqMAIJ *)A->data;
442:   Mat_SeqAIJ        *a = (Mat_SeqAIJ *)b->AIJ->data;
443:   const PetscScalar *x, *v, *alpha;
444:   PetscScalar       *y;
445:   const PetscInt     m = b->AIJ->rmap->n, *idx, dof = b->dof;
446:   PetscInt           n, i, k;

448:   PetscFunctionBegin;
449:   if (yy != zz) PetscCall(VecCopy(yy, zz));
450:   PetscCall(VecGetArrayRead(xx, &x));
451:   PetscCall(VecGetArray(zz, &y));
452:   for (i = 0; i < m; i++) {
453:     idx   = a->j + a->i[i];
454:     v     = a->a + a->i[i];
455:     n     = a->i[i + 1] - a->i[i];
456:     alpha = x + dof * i;
457:     while (n-- > 0) {
458:       for (k = 0; k < dof; k++) y[dof * (*idx) + k] += alpha[k] * (*v);
459:       idx++;
460:       v++;
461:     }
462:   }
463:   PetscCall(PetscLogFlops(2.0 * dof * a->nz));
464:   PetscCall(VecRestoreArrayRead(xx, &x));
465:   PetscCall(VecRestoreArray(zz, &y));
466:   PetscFunctionReturn(PETSC_SUCCESS);
467: }

469: static PetscErrorCode MatMult_MPIMAIJ_dof(Mat A, Vec xx, Vec yy)
470: {
471:   Mat_MPIMAIJ *b = (Mat_MPIMAIJ *)A->data;

473:   PetscFunctionBegin;
474:   /* start the scatter */
475:   PetscCall(VecScatterBegin(b->ctx, xx, b->w, INSERT_VALUES, SCATTER_FORWARD));
476:   PetscUseTypeMethod(b->AIJ, mult, xx, yy);
477:   PetscCall(VecScatterEnd(b->ctx, xx, b->w, INSERT_VALUES, SCATTER_FORWARD));
478:   PetscUseTypeMethod(b->OAIJ, multadd, b->w, yy, yy);
479:   PetscFunctionReturn(PETSC_SUCCESS);
480: }

482: static PetscErrorCode MatMultTranspose_MPIMAIJ_dof(Mat A, Vec xx, Vec yy)
483: {
484:   Mat_MPIMAIJ *b = (Mat_MPIMAIJ *)A->data;

486:   PetscFunctionBegin;
487:   PetscUseTypeMethod(b->OAIJ, multtranspose, xx, b->w);
488:   PetscUseTypeMethod(b->AIJ, multtranspose, xx, yy);
489:   PetscCall(VecScatterBegin(b->ctx, b->w, yy, ADD_VALUES, SCATTER_REVERSE));
490:   PetscCall(VecScatterEnd(b->ctx, b->w, yy, ADD_VALUES, SCATTER_REVERSE));
491:   PetscFunctionReturn(PETSC_SUCCESS);
492: }

494: static PetscErrorCode MatMultAdd_MPIMAIJ_dof(Mat A, Vec xx, Vec yy, Vec zz)
495: {
496:   Mat_MPIMAIJ *b = (Mat_MPIMAIJ *)A->data;

498:   PetscFunctionBegin;
499:   /* start the scatter */
500:   PetscCall(VecScatterBegin(b->ctx, xx, b->w, INSERT_VALUES, SCATTER_FORWARD));
501:   PetscUseTypeMethod(b->AIJ, multadd, xx, yy, zz);
502:   PetscCall(VecScatterEnd(b->ctx, xx, b->w, INSERT_VALUES, SCATTER_FORWARD));
503:   PetscUseTypeMethod(b->OAIJ, multadd, b->w, zz, zz);
504:   PetscFunctionReturn(PETSC_SUCCESS);
505: }

507: static PetscErrorCode MatMultTransposeAdd_MPIMAIJ_dof(Mat A, Vec xx, Vec yy, Vec zz)
508: {
509:   Mat_MPIMAIJ *b = (Mat_MPIMAIJ *)A->data;

511:   PetscFunctionBegin;
512:   PetscUseTypeMethod(b->OAIJ, multtranspose, xx, b->w);
513:   PetscUseTypeMethod(b->AIJ, multtransposeadd, xx, yy, zz);
514:   PetscCall(VecScatterBegin(b->ctx, b->w, zz, ADD_VALUES, SCATTER_REVERSE));
515:   PetscCall(VecScatterEnd(b->ctx, b->w, zz, ADD_VALUES, SCATTER_REVERSE));
516:   PetscFunctionReturn(PETSC_SUCCESS);
517: }

519: static PetscErrorCode MatProductSetFromOptions_SeqAIJ_SeqMAIJ(Mat C)
520: {
521:   Mat_Product *product = C->product;

523:   PetscFunctionBegin;
524:   PetscCheck(product->type == MATPRODUCT_PtAP, PETSC_COMM_SELF, PETSC_ERR_SUP, "Mat Product type %s is not supported for SeqAIJ and SeqMAIJ matrices", MatProductTypes[product->type]);
525:   C->ops->productsymbolic = MatProductSymbolic_PtAP_SeqAIJ_SeqMAIJ;
526:   PetscFunctionReturn(PETSC_SUCCESS);
527: }

529: static PetscErrorCode MatProductSetFromOptions_MPIAIJ_MPIMAIJ(Mat C)
530: {
531:   Mat_Product *product = C->product;
532:   PetscBool    flg     = PETSC_FALSE;
533:   Mat          A = product->A, P = product->B;
534:   PetscInt     alg = 1; /* set default algorithm */
535: #if !PetscDefined(HAVE_HYPRE)
536:   const char *algTypes[4] = {"scalable", "nonscalable", "allatonce", "allatonce_merged"};
537:   PetscInt    nalg        = 4;
538: #else
539:   const char *algTypes[5] = {"scalable", "nonscalable", "allatonce", "allatonce_merged", "hypre"};
540:   PetscInt    nalg        = 5;
541: #endif

543:   PetscFunctionBegin;
544:   PetscCheck(product->type == MATPRODUCT_PtAP, PETSC_COMM_SELF, PETSC_ERR_SUP, "Mat Product type %s is not supported for MPIAIJ and MPIMAIJ matrices", MatProductTypes[product->type]);

546:   /* PtAP */
547:   /* Check matrix local sizes */
548:   PetscCheck(A->rmap->rstart == P->rmap->rstart && A->rmap->rend == P->rmap->rend, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Matrix local dimensions are incompatible, Arow (%" PetscInt_FMT ", %" PetscInt_FMT ") != Prow (%" PetscInt_FMT ",%" PetscInt_FMT ")",
549:              A->rmap->rstart, A->rmap->rend, P->rmap->rstart, P->rmap->rend);
550:   PetscCheck(A->cmap->rstart == P->rmap->rstart && A->cmap->rend == P->rmap->rend, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Matrix local dimensions are incompatible, Acol (%" PetscInt_FMT ", %" PetscInt_FMT ") != Prow (%" PetscInt_FMT ",%" PetscInt_FMT ")",
551:              A->cmap->rstart, A->cmap->rend, P->rmap->rstart, P->rmap->rend);

553:   /* Set the default algorithm */
554:   PetscCall(PetscStrcmp(C->product->alg, "default", &flg));
555:   if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));

557:   /* Get runtime option */
558:   PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatProduct_PtAP", "Mat");
559:   PetscCall(PetscOptionsEList("-mat_product_algorithm", "Algorithmic approach", "MatPtAP", algTypes, nalg, algTypes[alg], &alg, &flg));
560:   if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
561:   PetscOptionsEnd();

563:   PetscCall(PetscStrcmp(C->product->alg, "allatonce", &flg));
564:   if (flg) {
565:     C->ops->productsymbolic = MatProductSymbolic_PtAP_MPIAIJ_MPIMAIJ;
566:     PetscFunctionReturn(PETSC_SUCCESS);
567:   }

569:   PetscCall(PetscStrcmp(C->product->alg, "allatonce_merged", &flg));
570:   if (flg) {
571:     C->ops->productsymbolic = MatProductSymbolic_PtAP_MPIAIJ_MPIMAIJ;
572:     PetscFunctionReturn(PETSC_SUCCESS);
573:   }

575:   /* Convert P from MAIJ to AIJ matrix since implementation not available for MAIJ */
576:   PetscCall(PetscInfo(A, "Converting from MAIJ to AIJ matrix since implementation not available for MAIJ\n"));
577:   PetscCall(MatConvert(P, MATMPIAIJ, MAT_INPLACE_MATRIX, &P));
578:   PetscCall(MatProductSetFromOptions(C));
579:   PetscFunctionReturn(PETSC_SUCCESS);
580: }

582: static PetscErrorCode MatPtAPNumeric_SeqAIJ_SeqMAIJ(Mat A, Mat PP, Mat C)
583: {
584:   /* This routine requires testing -- first draft only */
585:   Mat_SeqMAIJ     *pp = (Mat_SeqMAIJ *)PP->data;
586:   Mat              P  = pp->AIJ;
587:   Mat_SeqAIJ      *a  = (Mat_SeqAIJ *)A->data;
588:   Mat_SeqAIJ      *p  = (Mat_SeqAIJ *)P->data;
589:   Mat_SeqAIJ      *c  = (Mat_SeqAIJ *)C->data;
590:   const PetscInt  *ai = a->i, *aj = a->j, *pi = p->i, *pj = p->j, *pJ, *pjj;
591:   const PetscInt  *ci = c->i, *cj = c->j, *cjj;
592:   const PetscInt   am = A->rmap->N, cn = C->cmap->N, cm = C->rmap->N, ppdof = pp->dof;
593:   PetscInt         i, j, k, pshift, poffset, anzi, pnzi, apnzj, nextap, pnzj, prow, crow, *apj, *apjdense;
594:   const MatScalar *aa = a->a, *pa = p->a, *pA, *paj;
595:   MatScalar       *ca = c->a, *caj, *apa;

597:   PetscFunctionBegin;
598:   /* Allocate temporary array for storage of one row of A*P */
599:   PetscCall(PetscCalloc3(cn, &apa, cn, &apj, cn, &apjdense));

601:   /* Clear old values in C */
602:   PetscCall(PetscArrayzero(ca, ci[cm]));

604:   for (i = 0; i < am; i++) {
605:     /* Form sparse row of A*P */
606:     anzi  = ai[i + 1] - ai[i];
607:     apnzj = 0;
608:     for (j = 0; j < anzi; j++) {
609:       /* Get offset within block of P */
610:       pshift = *aj % ppdof;
611:       /* Get block row of P */
612:       prow = *aj++ / ppdof; /* integer division */
613:       pnzj = pi[prow + 1] - pi[prow];
614:       pjj  = pj + pi[prow];
615:       paj  = pa + pi[prow];
616:       for (k = 0; k < pnzj; k++) {
617:         poffset = pjj[k] * ppdof + pshift;
618:         if (!apjdense[poffset]) {
619:           apjdense[poffset] = -1;
620:           apj[apnzj++]      = poffset;
621:         }
622:         apa[poffset] += (*aa) * paj[k];
623:       }
624:       PetscCall(PetscLogFlops(2.0 * pnzj));
625:       aa++;
626:     }

628:     /* Sort the j index array for quick sparse axpy. */
629:     /* Note: a array does not need sorting as it is in dense storage locations. */
630:     PetscCall(PetscSortInt(apnzj, apj));

632:     /* Compute P^T*A*P using outer product (P^T)[:,j]*(A*P)[j,:]. */
633:     prow    = i / ppdof; /* integer division */
634:     pshift  = i % ppdof;
635:     poffset = pi[prow];
636:     pnzi    = pi[prow + 1] - poffset;
637:     /* Reset pJ and pA so we can traverse the same row of P 'dof' times. */
638:     pJ = pj + poffset;
639:     pA = pa + poffset;
640:     for (j = 0; j < pnzi; j++) {
641:       crow = (*pJ) * ppdof + pshift;
642:       cjj  = cj + ci[crow];
643:       caj  = ca + ci[crow];
644:       pJ++;
645:       /* Perform sparse axpy operation.  Note cjj includes apj. */
646:       for (k = 0, nextap = 0; nextap < apnzj; k++) {
647:         if (cjj[k] == apj[nextap]) caj[k] += (*pA) * apa[apj[nextap++]];
648:       }
649:       PetscCall(PetscLogFlops(2.0 * apnzj));
650:       pA++;
651:     }

653:     /* Zero the current row info for A*P */
654:     for (j = 0; j < apnzj; j++) {
655:       apa[apj[j]]      = 0.;
656:       apjdense[apj[j]] = 0;
657:     }
658:   }

660:   /* Assemble the final matrix and clean up */
661:   PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
662:   PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
663:   PetscCall(PetscFree3(apa, apj, apjdense));
664:   PetscFunctionReturn(PETSC_SUCCESS);
665: }

667: static PetscErrorCode MatPtAPSymbolic_SeqAIJ_SeqMAIJ(Mat A, Mat PP, PetscReal fill, Mat C)
668: {
669:   PetscFreeSpaceList free_space = NULL, current_space = NULL;
670:   Mat_SeqMAIJ       *pp = (Mat_SeqMAIJ *)PP->data;
671:   Mat                P  = pp->AIJ;
672:   Mat_SeqAIJ        *a = (Mat_SeqAIJ *)A->data, *p = (Mat_SeqAIJ *)P->data, *c;
673:   PetscInt          *pti, *ptj, *ptJ;
674:   PetscInt          *ci, *cj, *ptadenserow, *ptasparserow, *denserow, *sparserow, *ptaj;
675:   const PetscInt     an = A->cmap->N, am = A->rmap->N, pn = P->cmap->N, pm = P->rmap->N, ppdof = pp->dof;
676:   PetscInt           i, j, k, dof, pshift, ptnzi, arow, anzj, ptanzi, prow, pnzj, cnzi, cn;
677:   MatScalar         *ca;
678:   const PetscInt    *pi = p->i, *pj = p->j, *pjj, *ai = a->i, *aj = a->j, *ajj;

680:   PetscFunctionBegin;
681:   /* Get ij structure of P^T */
682:   PetscCall(MatGetSymbolicTranspose_SeqAIJ(P, &pti, &ptj));

684:   cn = pn * ppdof;
685:   /* Allocate ci array, arrays for fill computation and */
686:   /* free space for accumulating nonzero column info */
687:   PetscCall(PetscMalloc1(cn + 1, &ci));
688:   ci[0] = 0;

690:   /* Work arrays for rows of P^T*A */
691:   PetscCall(PetscMalloc4(an, &ptadenserow, an, &ptasparserow, cn, &denserow, cn, &sparserow));
692:   PetscCall(PetscArrayzero(ptadenserow, an));
693:   PetscCall(PetscArrayzero(denserow, cn));

695:   /* Set initial free space to be nnz(A) scaled by aspect ratio of P. */
696:   /* This should be reasonable if sparsity of PtAP is similar to that of A. */
697:   /* Note, aspect ratio of P is the same as the aspect ratio of SeqAIJ inside P */
698:   PetscCall(PetscFreeSpaceGet(PetscIntMultTruncate(ai[am] / pm, pn), &free_space));
699:   current_space = free_space;

701:   /* Determine symbolic info for each row of C: */
702:   for (i = 0; i < pn; i++) {
703:     ptnzi = pti[i + 1] - pti[i];
704:     ptJ   = ptj + pti[i];
705:     for (dof = 0; dof < ppdof; dof++) {
706:       ptanzi = 0;
707:       /* Determine symbolic row of PtA: */
708:       for (j = 0; j < ptnzi; j++) {
709:         /* Expand ptJ[j] by block size and shift by dof to get the right row of A */
710:         arow = ptJ[j] * ppdof + dof;
711:         /* Nonzeros of P^T*A will be in same locations as any element of A in that row */
712:         anzj = ai[arow + 1] - ai[arow];
713:         ajj  = aj + ai[arow];
714:         for (k = 0; k < anzj; k++) {
715:           if (!ptadenserow[ajj[k]]) {
716:             ptadenserow[ajj[k]]    = -1;
717:             ptasparserow[ptanzi++] = ajj[k];
718:           }
719:         }
720:       }
721:       /* Using symbolic info for row of PtA, determine symbolic info for row of C: */
722:       ptaj = ptasparserow;
723:       cnzi = 0;
724:       for (j = 0; j < ptanzi; j++) {
725:         /* Get offset within block of P */
726:         pshift = *ptaj % ppdof;
727:         /* Get block row of P */
728:         prow = (*ptaj++) / ppdof; /* integer division */
729:         /* P has same number of nonzeros per row as the compressed form */
730:         pnzj = pi[prow + 1] - pi[prow];
731:         pjj  = pj + pi[prow];
732:         for (k = 0; k < pnzj; k++) {
733:           /* Locations in C are shifted by the offset within the block */
734:           /* Note: we cannot use PetscLLAdd here because of the additional offset for the write location */
735:           if (!denserow[pjj[k] * ppdof + pshift]) {
736:             denserow[pjj[k] * ppdof + pshift] = -1;
737:             sparserow[cnzi++]                 = pjj[k] * ppdof + pshift;
738:           }
739:         }
740:       }

742:       /* sort sparserow */
743:       PetscCall(PetscSortInt(cnzi, sparserow));

745:       /* If free space is not available, make more free space */
746:       /* Double the amount of total space in the list */
747:       if (current_space->local_remaining < cnzi) PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(cnzi, current_space->total_array_size), &current_space));

749:       /* Copy data into free space, and zero out denserows */
750:       PetscCall(PetscArraycpy(current_space->array, sparserow, cnzi));

752:       current_space->array += cnzi;
753:       current_space->local_used += cnzi;
754:       current_space->local_remaining -= cnzi;

756:       for (j = 0; j < ptanzi; j++) ptadenserow[ptasparserow[j]] = 0;
757:       for (j = 0; j < cnzi; j++) denserow[sparserow[j]] = 0;

759:       /* Aside: Perhaps we should save the pta info for the numerical factorization. */
760:       /*        For now, we will recompute what is needed. */
761:       ci[i * ppdof + 1 + dof] = ci[i * ppdof + dof] + cnzi;
762:     }
763:   }
764:   /* nnz is now stored in ci[ptm], column indices are in the list of free space */
765:   /* Allocate space for cj, initialize cj, and */
766:   /* destroy list of free space and other temporary array(s) */
767:   PetscCall(PetscMalloc1(ci[cn], &cj));
768:   PetscCall(PetscFreeSpaceContiguous(&free_space, cj));
769:   PetscCall(PetscFree4(ptadenserow, ptasparserow, denserow, sparserow));

771:   /* Allocate space for ca */
772:   PetscCall(PetscCalloc1(ci[cn], &ca));

774:   /* put together the new matrix */
775:   PetscCall(MatSetSeqAIJWithArrays_private(PetscObjectComm((PetscObject)A), cn, cn, ci, cj, ca, NULL, C));
776:   PetscCall(MatSetBlockSize(C, pp->dof));

778:   /* MatCreateSeqAIJWithArrays flags matrix so PETSc doesn't free the user's arrays. */
779:   /* Since these are PETSc arrays, change flags to free them as necessary. */
780:   c          = (Mat_SeqAIJ *)C->data;
781:   c->free_a  = PETSC_TRUE;
782:   c->free_ij = PETSC_TRUE;
783:   c->nonew   = 0;

785:   C->ops->ptapnumeric    = MatPtAPNumeric_SeqAIJ_SeqMAIJ;
786:   C->ops->productnumeric = MatProductNumeric_PtAP;

788:   /* Clean up. */
789:   PetscCall(MatRestoreSymbolicTranspose_SeqAIJ(P, &pti, &ptj));
790:   PetscFunctionReturn(PETSC_SUCCESS);
791: }

793: PETSC_INTERN PetscErrorCode MatProductSymbolic_PtAP_SeqAIJ_SeqMAIJ(Mat C)
794: {
795:   Mat_Product *product = C->product;
796:   Mat          A = product->A, P = product->B;

798:   PetscFunctionBegin;
799:   PetscCall(MatPtAPSymbolic_SeqAIJ_SeqMAIJ(A, P, product->fill, C));
800:   PetscFunctionReturn(PETSC_SUCCESS);
801: }

803: PetscErrorCode MatPtAPNumeric_MPIAIJ_MPIXAIJ_allatonce(Mat, Mat, PetscInt, Mat);

805: PETSC_INTERN PetscErrorCode MatPtAPNumeric_MPIAIJ_MPIMAIJ_allatonce(Mat A, Mat P, Mat C)
806: {
807:   Mat_MPIMAIJ *maij = (Mat_MPIMAIJ *)P->data;

809:   PetscFunctionBegin;
810:   PetscCall(MatPtAPNumeric_MPIAIJ_MPIXAIJ_allatonce(A, maij->A, maij->dof, C));
811:   PetscFunctionReturn(PETSC_SUCCESS);
812: }

814: PetscErrorCode MatPtAPSymbolic_MPIAIJ_MPIXAIJ_allatonce(Mat, Mat, PetscInt, PetscReal, Mat);

816: PETSC_INTERN PetscErrorCode MatPtAPSymbolic_MPIAIJ_MPIMAIJ_allatonce(Mat A, Mat P, PetscReal fill, Mat C)
817: {
818:   Mat_MPIMAIJ *maij = (Mat_MPIMAIJ *)P->data;

820:   PetscFunctionBegin;
821:   PetscCall(MatPtAPSymbolic_MPIAIJ_MPIXAIJ_allatonce(A, maij->A, maij->dof, fill, C));
822:   C->ops->ptapnumeric = MatPtAPNumeric_MPIAIJ_MPIMAIJ_allatonce;
823:   PetscFunctionReturn(PETSC_SUCCESS);
824: }

826: PetscErrorCode MatPtAPNumeric_MPIAIJ_MPIXAIJ_allatonce_merged(Mat, Mat, PetscInt, Mat);

828: PETSC_INTERN PetscErrorCode MatPtAPNumeric_MPIAIJ_MPIMAIJ_allatonce_merged(Mat A, Mat P, Mat C)
829: {
830:   Mat_MPIMAIJ *maij = (Mat_MPIMAIJ *)P->data;

832:   PetscFunctionBegin;
833:   PetscCall(MatPtAPNumeric_MPIAIJ_MPIXAIJ_allatonce_merged(A, maij->A, maij->dof, C));
834:   PetscFunctionReturn(PETSC_SUCCESS);
835: }

837: PetscErrorCode MatPtAPSymbolic_MPIAIJ_MPIXAIJ_allatonce_merged(Mat, Mat, PetscInt, PetscReal, Mat);

839: PETSC_INTERN PetscErrorCode MatPtAPSymbolic_MPIAIJ_MPIMAIJ_allatonce_merged(Mat A, Mat P, PetscReal fill, Mat C)
840: {
841:   Mat_MPIMAIJ *maij = (Mat_MPIMAIJ *)P->data;

843:   PetscFunctionBegin;
844:   PetscCall(MatPtAPSymbolic_MPIAIJ_MPIXAIJ_allatonce_merged(A, maij->A, maij->dof, fill, C));
845:   C->ops->ptapnumeric = MatPtAPNumeric_MPIAIJ_MPIMAIJ_allatonce_merged;
846:   PetscFunctionReturn(PETSC_SUCCESS);
847: }

849: PETSC_INTERN PetscErrorCode MatProductSymbolic_PtAP_MPIAIJ_MPIMAIJ(Mat C)
850: {
851:   Mat_Product *product = C->product;
852:   Mat          A = product->A, P = product->B;
853:   PetscBool    flg;

855:   PetscFunctionBegin;
856:   PetscCall(PetscStrcmp(product->alg, "allatonce", &flg));
857:   if (flg) {
858:     PetscCall(MatPtAPSymbolic_MPIAIJ_MPIMAIJ_allatonce(A, P, product->fill, C));
859:     C->ops->productnumeric = MatProductNumeric_PtAP;
860:     PetscFunctionReturn(PETSC_SUCCESS);
861:   }

863:   PetscCall(PetscStrcmp(product->alg, "allatonce_merged", &flg));
864:   PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_SUP, "Mat Product Algorithm is not supported");
865:   PetscCall(MatPtAPSymbolic_MPIAIJ_MPIMAIJ_allatonce_merged(A, P, product->fill, C));
866:   C->ops->productnumeric = MatProductNumeric_PtAP;
867:   PetscFunctionReturn(PETSC_SUCCESS);
868: }

870: PETSC_INTERN PetscErrorCode MatConvert_SeqMAIJ_SeqAIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
871: {
872:   Mat_SeqMAIJ *b   = (Mat_SeqMAIJ *)A->data;
873:   Mat          a   = b->AIJ, B;
874:   Mat_SeqAIJ  *aij = (Mat_SeqAIJ *)a->data;
875:   PetscInt     m, n, i, ncols, *ilen, nmax = 0, *icols, j, k, ii, dof = b->dof;
876:   PetscInt    *cols;
877:   PetscScalar *vals;

879:   PetscFunctionBegin;
880:   PetscCall(MatGetSize(a, &m, &n));
881:   PetscCall(PetscMalloc1(dof * m, &ilen));
882:   for (i = 0; i < m; i++) {
883:     nmax = PetscMax(nmax, aij->ilen[i]);
884:     for (j = 0; j < dof; j++) ilen[dof * i + j] = aij->ilen[i];
885:   }
886:   PetscCall(MatCreate(PETSC_COMM_SELF, &B));
887:   PetscCall(MatSetSizes(B, dof * m, dof * n, dof * m, dof * n));
888:   PetscCall(MatSetType(B, newtype));
889:   PetscCall(MatSeqAIJSetPreallocation(B, 0, ilen));
890:   PetscCall(PetscFree(ilen));
891:   PetscCall(PetscMalloc1(nmax, &icols));
892:   ii = 0;
893:   for (i = 0; i < m; i++) {
894:     PetscCall(MatGetRow_SeqAIJ(a, i, &ncols, &cols, &vals));
895:     for (j = 0; j < dof; j++) {
896:       for (k = 0; k < ncols; k++) icols[k] = dof * cols[k] + j;
897:       PetscCall(MatSetValues_SeqAIJ(B, 1, &ii, ncols, icols, vals, INSERT_VALUES));
898:       ii++;
899:     }
900:     PetscCall(MatRestoreRow_SeqAIJ(a, i, &ncols, &cols, &vals));
901:   }
902:   PetscCall(PetscFree(icols));
903:   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
904:   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));

906:   if (reuse == MAT_INPLACE_MATRIX) {
907:     PetscCall(MatHeaderReplace(A, &B));
908:   } else {
909:     *newmat = B;
910:   }
911:   PetscFunctionReturn(PETSC_SUCCESS);
912: }

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

916: PETSC_INTERN PetscErrorCode MatConvert_MPIMAIJ_MPIAIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
917: {
918:   Mat_MPIMAIJ *maij    = (Mat_MPIMAIJ *)A->data;
919:   Mat          MatAIJ  = ((Mat_SeqMAIJ *)maij->AIJ->data)->AIJ, B;
920:   Mat          MatOAIJ = ((Mat_SeqMAIJ *)maij->OAIJ->data)->AIJ;
921:   Mat_SeqAIJ  *AIJ     = (Mat_SeqAIJ *)MatAIJ->data;
922:   Mat_SeqAIJ  *OAIJ    = (Mat_SeqAIJ *)MatOAIJ->data;
923:   Mat_MPIAIJ  *mpiaij  = (Mat_MPIAIJ *)maij->A->data;
924:   PetscInt     dof = maij->dof, i, j, *dnz = NULL, *onz = NULL, nmax = 0, onmax = 0;
925:   PetscInt    *oicols = NULL, *icols = NULL, ncols, *cols = NULL, oncols, *ocols = NULL;
926:   PetscInt     rstart, cstart, *garray, ii, k;
927:   PetscScalar *vals, *ovals;

929:   PetscFunctionBegin;
930:   PetscCall(PetscMalloc2(A->rmap->n, &dnz, A->rmap->n, &onz));
931:   for (i = 0; i < A->rmap->n / dof; i++) {
932:     nmax  = PetscMax(nmax, AIJ->ilen[i]);
933:     onmax = PetscMax(onmax, OAIJ->ilen[i]);
934:     for (j = 0; j < dof; j++) {
935:       dnz[dof * i + j] = AIJ->ilen[i];
936:       onz[dof * i + j] = OAIJ->ilen[i];
937:     }
938:   }
939:   PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
940:   PetscCall(MatSetSizes(B, A->rmap->n, A->cmap->n, A->rmap->N, A->cmap->N));
941:   PetscCall(MatSetType(B, newtype));
942:   PetscCall(MatMPIAIJSetPreallocation(B, 0, dnz, 0, onz));
943:   PetscCall(MatSetBlockSize(B, dof));
944:   PetscCall(PetscFree2(dnz, onz));

946:   PetscCall(PetscMalloc2(nmax, &icols, onmax, &oicols));
947:   rstart = dof * maij->A->rmap->rstart;
948:   cstart = dof * maij->A->cmap->rstart;
949:   garray = mpiaij->garray;

951:   ii = rstart;
952:   for (i = 0; i < A->rmap->n / dof; i++) {
953:     PetscCall(MatGetRow_SeqAIJ(MatAIJ, i, &ncols, &cols, &vals));
954:     PetscCall(MatGetRow_SeqAIJ(MatOAIJ, i, &oncols, &ocols, &ovals));
955:     for (j = 0; j < dof; j++) {
956:       for (k = 0; k < ncols; k++) icols[k] = cstart + dof * cols[k] + j;
957:       for (k = 0; k < oncols; k++) oicols[k] = dof * garray[ocols[k]] + j;
958:       PetscCall(MatSetValues_MPIAIJ(B, 1, &ii, ncols, icols, vals, INSERT_VALUES));
959:       PetscCall(MatSetValues_MPIAIJ(B, 1, &ii, oncols, oicols, ovals, INSERT_VALUES));
960:       ii++;
961:     }
962:     PetscCall(MatRestoreRow_SeqAIJ(MatAIJ, i, &ncols, &cols, &vals));
963:     PetscCall(MatRestoreRow_SeqAIJ(MatOAIJ, i, &oncols, &ocols, &ovals));
964:   }
965:   PetscCall(PetscFree2(icols, oicols));

967:   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
968:   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));

970:   if (reuse == MAT_INPLACE_MATRIX) {
971:     PetscInt refct          = ((PetscObject)A)->refct; /* save ((PetscObject)A)->refct */
972:     ((PetscObject)A)->refct = 1;

974:     PetscCall(MatHeaderReplace(A, &B));

976:     ((PetscObject)A)->refct = refct; /* restore ((PetscObject)A)->refct */
977:   } else {
978:     *newmat = B;
979:   }
980:   PetscFunctionReturn(PETSC_SUCCESS);
981: }

983: static PetscErrorCode MatCreateSubMatrix_MAIJ(Mat mat, IS isrow, IS iscol, MatReuse cll, Mat *newmat)
984: {
985:   Mat A;

987:   PetscFunctionBegin;
988:   PetscCall(MatConvert(mat, MATAIJ, MAT_INITIAL_MATRIX, &A));
989:   PetscCall(MatCreateSubMatrix(A, isrow, iscol, cll, newmat));
990:   PetscCall(MatDestroy(&A));
991:   PetscFunctionReturn(PETSC_SUCCESS);
992: }

994: static PetscErrorCode MatCreateSubMatrices_MAIJ(Mat mat, PetscInt n, const IS irow[], const IS icol[], MatReuse scall, Mat *submat[])
995: {
996:   Mat A;

998:   PetscFunctionBegin;
999:   PetscCall(MatConvert(mat, MATAIJ, MAT_INITIAL_MATRIX, &A));
1000:   PetscCall(MatCreateSubMatrices(A, n, irow, icol, scall, submat));
1001:   PetscCall(MatDestroy(&A));
1002:   PetscFunctionReturn(PETSC_SUCCESS);
1003: }

1005: /*@
1006:   MatCreateMAIJ - Creates a matrix type providing restriction and interpolation
1007:   operations for multicomponent problems.  It interpolates each component the same
1008:   way independently.  The matrix type is based on `MATSEQAIJ` for sequential matrices,
1009:   and `MATMPIAIJ` for distributed matrices.

1011:   Collective

1013:   Input Parameters:
1014: + A   - the `MATAIJ` matrix describing the action on blocks
1015: - dof - the block size (number of components per node)

1017:   Output Parameter:
1018: . maij - the new `MATMAIJ` matrix

1020:   Level: advanced

1022: .seealso: [](ch_matrices), `Mat`, `MATAIJ`, `MATMAIJ`, `MatMAIJGetAIJ()`, `MatMAIJRedimension()`
1023: @*/
1024: PetscErrorCode MatCreateMAIJ(Mat A, PetscInt dof, Mat *maij)
1025: {
1026:   PetscInt  n;
1027:   Mat       B;
1028:   PetscBool flg;
1029:   /* hack to prevent conversion to AIJ format for device types when used inside a parallel MAIJ */
1030:   PetscBool convert = dof < 0 ? PETSC_FALSE : PETSC_TRUE;

1032:   PetscFunctionBegin;
1033:   dof = PetscAbs(dof);
1034:   PetscCall(PetscObjectReference((PetscObject)A));

1036:   if (dof == 1) *maij = A;
1037:   else {
1038:     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
1039:     /* propagate vec type */
1040:     PetscCall(MatSetVecType(B, A->defaultvectype));
1041:     PetscCall(MatSetSizes(B, dof * A->rmap->n, dof * A->cmap->n, dof * A->rmap->N, dof * A->cmap->N));
1042:     PetscCall(PetscLayoutSetBlockSize(B->rmap, dof));
1043:     PetscCall(PetscLayoutSetBlockSize(B->cmap, dof));
1044:     PetscCall(PetscLayoutSetUp(B->rmap));
1045:     PetscCall(PetscLayoutSetUp(B->cmap));

1047:     B->assembled = PETSC_TRUE;

1049:     PetscCall(PetscObjectBaseTypeCompare((PetscObject)A, MATSEQAIJ, &flg));
1050:     if (flg) {
1051:       Mat_SeqMAIJ *b;

1053:       PetscCall(MatSetType(B, MATSEQMAIJ));

1055:       B->ops->setup   = NULL;
1056:       B->ops->destroy = MatDestroy_SeqMAIJ;
1057:       B->ops->view    = MatView_SeqMAIJ;

1059:       b      = (Mat_SeqMAIJ *)B->data;
1060:       b->dof = dof;
1061:       b->AIJ = A;

1063:       if (dof == 2) {
1064:         B->ops->mult             = MatMult_SeqMAIJ_2;
1065:         B->ops->multadd          = MatMultAdd_SeqMAIJ_2;
1066:         B->ops->multtranspose    = MatMultTranspose_SeqMAIJ_2;
1067:         B->ops->multtransposeadd = MatMultTransposeAdd_SeqMAIJ_2;
1068:       } else if (dof == 3) {
1069:         B->ops->mult             = MatMult_SeqMAIJ_3;
1070:         B->ops->multadd          = MatMultAdd_SeqMAIJ_3;
1071:         B->ops->multtranspose    = MatMultTranspose_SeqMAIJ_3;
1072:         B->ops->multtransposeadd = MatMultTransposeAdd_SeqMAIJ_3;
1073:       } else if (dof == 4) {
1074:         B->ops->mult             = MatMult_SeqMAIJ_4;
1075:         B->ops->multadd          = MatMultAdd_SeqMAIJ_4;
1076:         B->ops->multtranspose    = MatMultTranspose_SeqMAIJ_4;
1077:         B->ops->multtransposeadd = MatMultTransposeAdd_SeqMAIJ_4;
1078:       } else if (dof == 5) {
1079:         B->ops->mult             = MatMult_SeqMAIJ_5;
1080:         B->ops->multadd          = MatMultAdd_SeqMAIJ_5;
1081:         B->ops->multtranspose    = MatMultTranspose_SeqMAIJ_5;
1082:         B->ops->multtransposeadd = MatMultTransposeAdd_SeqMAIJ_5;
1083:       } else if (dof == 6) {
1084:         B->ops->mult             = MatMult_SeqMAIJ_6;
1085:         B->ops->multadd          = MatMultAdd_SeqMAIJ_6;
1086:         B->ops->multtranspose    = MatMultTranspose_SeqMAIJ_6;
1087:         B->ops->multtransposeadd = MatMultTransposeAdd_SeqMAIJ_6;
1088:       } else if (dof == 7) {
1089:         B->ops->mult             = MatMult_SeqMAIJ_7;
1090:         B->ops->multadd          = MatMultAdd_SeqMAIJ_7;
1091:         B->ops->multtranspose    = MatMultTranspose_SeqMAIJ_7;
1092:         B->ops->multtransposeadd = MatMultTransposeAdd_SeqMAIJ_7;
1093:       } else if (dof == 8) {
1094:         B->ops->mult             = MatMult_SeqMAIJ_8;
1095:         B->ops->multadd          = MatMultAdd_SeqMAIJ_8;
1096:         B->ops->multtranspose    = MatMultTranspose_SeqMAIJ_8;
1097:         B->ops->multtransposeadd = MatMultTransposeAdd_SeqMAIJ_8;
1098:       } else if (dof == 9) {
1099:         B->ops->mult             = MatMult_SeqMAIJ_9;
1100:         B->ops->multadd          = MatMultAdd_SeqMAIJ_9;
1101:         B->ops->multtranspose    = MatMultTranspose_SeqMAIJ_9;
1102:         B->ops->multtransposeadd = MatMultTransposeAdd_SeqMAIJ_9;
1103:       } else if (dof == 10) {
1104:         B->ops->mult             = MatMult_SeqMAIJ_10;
1105:         B->ops->multadd          = MatMultAdd_SeqMAIJ_10;
1106:         B->ops->multtranspose    = MatMultTranspose_SeqMAIJ_10;
1107:         B->ops->multtransposeadd = MatMultTransposeAdd_SeqMAIJ_10;
1108:       } else if (dof == 11) {
1109:         B->ops->mult             = MatMult_SeqMAIJ_11;
1110:         B->ops->multadd          = MatMultAdd_SeqMAIJ_11;
1111:         B->ops->multtranspose    = MatMultTranspose_SeqMAIJ_11;
1112:         B->ops->multtransposeadd = MatMultTransposeAdd_SeqMAIJ_11;
1113:       } else if (dof == 16) {
1114:         B->ops->mult             = MatMult_SeqMAIJ_16;
1115:         B->ops->multadd          = MatMultAdd_SeqMAIJ_16;
1116:         B->ops->multtranspose    = MatMultTranspose_SeqMAIJ_16;
1117:         B->ops->multtransposeadd = MatMultTransposeAdd_SeqMAIJ_16;
1118:       } else if (dof == 18) {
1119:         B->ops->mult             = MatMult_SeqMAIJ_18;
1120:         B->ops->multadd          = MatMultAdd_SeqMAIJ_18;
1121:         B->ops->multtranspose    = MatMultTranspose_SeqMAIJ_18;
1122:         B->ops->multtransposeadd = MatMultTransposeAdd_SeqMAIJ_18;
1123:       } else {
1124:         B->ops->mult             = MatMult_SeqMAIJ_N;
1125:         B->ops->multadd          = MatMultAdd_SeqMAIJ_N;
1126:         B->ops->multtranspose    = MatMultTranspose_SeqMAIJ_N;
1127:         B->ops->multtransposeadd = MatMultTransposeAdd_SeqMAIJ_N;
1128:       }
1129: #if PetscDefined(HAVE_CUDA)
1130:       PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_seqmaij_seqaijcusparse_C", MatConvert_SeqMAIJ_SeqAIJ));
1131: #endif
1132: #if PetscDefined(HAVE_KOKKOS_KERNELS)
1133:       PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_seqmaij_seqaijkokkos_C", MatConvert_SeqMAIJ_SeqAIJ));
1134: #endif
1135:       PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_seqmaij_seqaij_C", MatConvert_SeqMAIJ_SeqAIJ));
1136:       PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatProductSetFromOptions_seqaij_seqmaij_C", MatProductSetFromOptions_SeqAIJ_SeqMAIJ));
1137:     } else {
1138:       Mat_MPIAIJ  *mpiaij = (Mat_MPIAIJ *)A->data;
1139:       Mat_MPIMAIJ *b;
1140:       IS           from, to;
1141:       Vec          gvec;

1143:       PetscCall(MatSetType(B, MATMPIMAIJ));

1145:       B->ops->setup   = NULL;
1146:       B->ops->destroy = MatDestroy_MPIMAIJ;
1147:       B->ops->view    = MatView_MPIMAIJ;

1149:       b      = (Mat_MPIMAIJ *)B->data;
1150:       b->dof = dof;
1151:       b->A   = A;

1153:       PetscCall(MatCreateMAIJ(mpiaij->A, -dof, &b->AIJ));
1154:       PetscCall(MatCreateMAIJ(mpiaij->B, -dof, &b->OAIJ));

1156:       PetscCall(VecGetSize(mpiaij->lvec, &n));
1157:       PetscCall(VecCreate(PETSC_COMM_SELF, &b->w));
1158:       PetscCall(VecSetSizes(b->w, n * dof, n * dof));
1159:       PetscCall(VecSetBlockSize(b->w, dof));
1160:       PetscCall(VecSetType(b->w, VECSEQ));

1162:       /* create two temporary Index sets for build scatter gather */
1163:       PetscCall(ISCreateBlock(PetscObjectComm((PetscObject)A), dof, n, mpiaij->garray, PETSC_COPY_VALUES, &from));
1164:       PetscCall(ISCreateStride(PETSC_COMM_SELF, n * dof, 0, 1, &to));

1166:       /* create temporary global vector to generate scatter context */
1167:       PetscCall(VecCreateMPIWithArray(PetscObjectComm((PetscObject)A), dof, dof * A->cmap->n, dof * A->cmap->N, NULL, &gvec));

1169:       /* generate the scatter context */
1170:       PetscCall(VecScatterCreate(gvec, from, b->w, to, &b->ctx));

1172:       PetscCall(ISDestroy(&from));
1173:       PetscCall(ISDestroy(&to));
1174:       PetscCall(VecDestroy(&gvec));

1176:       B->ops->mult             = MatMult_MPIMAIJ_dof;
1177:       B->ops->multtranspose    = MatMultTranspose_MPIMAIJ_dof;
1178:       B->ops->multadd          = MatMultAdd_MPIMAIJ_dof;
1179:       B->ops->multtransposeadd = MatMultTransposeAdd_MPIMAIJ_dof;

1181: #if PetscDefined(HAVE_CUDA)
1182:       PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_mpimaij_mpiaijcusparse_C", MatConvert_MPIMAIJ_MPIAIJ));
1183: #endif
1184: #if PetscDefined(HAVE_KOKKOS_KERNELS)
1185:       PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_mpimaij_mpiaijkokkos_C", MatConvert_MPIMAIJ_MPIAIJ));
1186: #endif
1187:       PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_mpimaij_mpiaij_C", MatConvert_MPIMAIJ_MPIAIJ));
1188:       PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatProductSetFromOptions_mpiaij_mpimaij_C", MatProductSetFromOptions_MPIAIJ_MPIMAIJ));
1189:     }
1190:     B->ops->createsubmatrix   = MatCreateSubMatrix_MAIJ;
1191:     B->ops->createsubmatrices = MatCreateSubMatrices_MAIJ;
1192:     PetscCall(MatSetUp(B));
1193:     /* temporary until we have device implementations of MAIJ */
1194:     if (convert) {
1195:       PetscCall(PetscObjectTypeCompareAny((PetscObject)A, &flg, MATSEQAIJCUSPARSE, MATMPIAIJCUSPARSE, MATAIJCUSPARSE, MATSEQAIJKOKKOS, MATMPIAIJKOKKOS, MATAIJKOKKOS, ""));
1196:       if (flg) PetscCall(MatConvert(B, ((PetscObject)A)->type_name, MAT_INPLACE_MATRIX, &B));
1197:     }
1198:     *maij = B;
1199:   }
1200:   PetscFunctionReturn(PETSC_SUCCESS);
1201: }