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), ¤t_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: }