Actual source code: mpihashmat.h
1: /*
2: used by MPIAIJ, BAIJ and SBAIJ to reduce code duplication
4: define TYPE to AIJ BAIJ or SBAIJ
5: TYPE_SBAIJ for SBAIJ matrix
7: */
9: static PetscErrorCode MatSetValues_MPI_Hash(Mat A, PetscInt m, const PetscInt *rows, PetscInt n, const PetscInt *cols, const PetscScalar *values, InsertMode addv)
10: {
11: PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data;
12: const PetscInt rStart = A->rmap->rstart;
13: const PetscInt rEnd = A->rmap->rend;
14: const PetscInt cStart = A->cmap->rstart;
15: const PetscInt cEnd = A->cmap->rend;
16: #if defined(TYPE_SBAIJ)
17: const PetscInt bs = A->rmap->bs;
18: #endif
19: const PetscBool structure_only = a->A->structure_only;
20: const PetscBool ignorezeroentries = (PetscBool)(!structure_only && ((Mat_SeqAIJ *)a->A->data)->ignorezeroentries);
22: PetscFunctionBegin;
23: for (PetscInt r = 0; r < m; ++r) {
24: PetscScalar value;
25: if (rows[r] < 0) continue;
26: if (rows[r] < rStart || rows[r] >= rEnd) {
27: PetscCheck(!A->nooffprocentries, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Setting off process row %" PetscInt_FMT " even though MatSetOption(,MAT_NO_OFF_PROC_ENTRIES,PETSC_TRUE) was set", rows[r]);
28: if (!a->donotstash) {
29: A->assembled = PETSC_FALSE;
30: if (a->roworiented) {
31: PetscCall(MatStashValuesRow_Private(&A->stash, rows[r], n, cols, PetscSafePointerPlusOffset(values, r * n), (PetscBool)(ignorezeroentries && addv == ADD_VALUES)));
32: } else {
33: PetscCall(MatStashValuesCol_Private(&A->stash, rows[r], n, cols, PetscSafePointerPlusOffset(values, r), m, (PetscBool)(ignorezeroentries && addv == ADD_VALUES)));
34: }
35: }
36: } else {
37: for (PetscInt c = 0; c < n; ++c) {
38: #if defined(TYPE_SBAIJ)
39: if (cols[c] / bs < rows[r] / bs) continue;
40: #else
41: if (cols[c] < 0) continue;
42: #endif
43: value = values && !structure_only ? (a->roworiented ? values[r * n + c] : values[r + m * c]) : 0;
44: if (ignorezeroentries && value == 0.0 && addv == ADD_VALUES && rows[r] != cols[c]) continue;
45: if (cols[c] >= cStart && cols[c] < cEnd) {
46: PetscCall(MatSetValue(a->A, rows[r] - rStart, cols[c] - cStart, value, addv));
47: } else if (!ignorezeroentries || value != 0.0) {
48: /* MPIAIJ never inserts in B if it is 0 */
49: PetscCall(MatSetValue(a->B, rows[r] - rStart, cols[c], value, addv));
50: }
51: }
52: }
53: }
54: PetscFunctionReturn(PETSC_SUCCESS);
55: }
57: static PetscErrorCode MatAssemblyBegin_MPI_Hash(Mat A, PETSC_UNUSED MatAssemblyType type)
58: {
59: PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data;
60: PetscInt nstash, reallocs;
62: PetscFunctionBegin;
63: if (a->donotstash || A->nooffprocentries) PetscFunctionReturn(PETSC_SUCCESS);
64: PetscCall(MatStashScatterBegin_Private(A, &A->stash, A->rmap->range));
65: PetscCall(MatStashGetInfo_Private(&A->stash, &nstash, &reallocs));
66: PetscCall(PetscInfo(A, "Stash has %" PetscInt_FMT " entries, uses %" PetscInt_FMT " mallocs.\n", nstash, reallocs));
67: PetscFunctionReturn(PETSC_SUCCESS);
68: }
70: static PetscErrorCode MatFinishScatterAndSetValues_MPI_Hash(Mat A)
71: {
72: PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data;
73: PetscMPIInt n;
74: PetscScalar *val;
75: PetscInt *row, *col;
76: PetscInt j, ncols, flg, rstart;
78: PetscFunctionBegin;
79: if (!a->donotstash && !A->nooffprocentries) {
80: while (1) {
81: PetscCall(MatStashScatterGetMesg_Private(&A->stash, &n, &row, &col, &val, &flg));
82: if (!flg) break;
84: for (PetscInt i = 0; i < n;) {
85: /* Now identify the consecutive vals belonging to the same row */
86: for (j = i, rstart = row[j]; j < n; j++) {
87: if (row[j] != rstart) break;
88: }
89: if (j < n) ncols = j - i;
90: else ncols = n - i;
91: /* Now assemble all these values with a single function call */
92: PetscCall(MatSetValues_MPI_Hash(A, 1, row + i, ncols, col + i, val + i, A->insertmode));
93: i = j;
94: }
95: }
96: PetscCall(MatStashScatterEnd_Private(&A->stash));
97: }
98: PetscFunctionReturn(PETSC_SUCCESS);
99: }
101: static PetscErrorCode MatAssemblyEnd_MPI_Hash(Mat A, MatAssemblyType type)
102: {
103: PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data;
104: PetscBool nooffprocentries = A->nooffprocentries;
106: PetscFunctionBegin;
107: PetscCall(MatFinishScatterAndSetValues_MPI_Hash(A));
108: if (type != MAT_FINAL_ASSEMBLY) PetscFunctionReturn(PETSC_SUCCESS);
110: A->insertmode = NOT_SET_VALUES; /* this was set by the previous calls to MatSetValues() */
112: A->ops[0] = a->cops;
113: A->hash_active = PETSC_FALSE;
115: /* a->B must be in CSR before MatSetUpMultiply_MPIAIJ() runs inside MatAssemblyEnd_MPIAIJ().
116: a->A is assembled by MatAssemblyEnd_MPIAIJ() itself and does not need a separate call here. */
117: PetscCall(MatAssemblyBegin(a->B, MAT_FINAL_ASSEMBLY));
118: PetscCall(MatAssemblyEnd(a->B, MAT_FINAL_ASSEMBLY));
119: /* Off-process entries were already scattered in MatAssemblyBegin_MPI_Hash().
120: Suppress the stash scatter in MatAssemblyBegin_MPIAIJ() to avoid a second
121: PetscCommBuildTwoSidedFReq() call on an empty stash. */
122: A->nooffprocentries = PETSC_TRUE;
123: PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
124: PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
125: A->nooffprocentries = nooffprocentries;
126: PetscFunctionReturn(PETSC_SUCCESS);
127: }
129: static PetscErrorCode MatCopyHashToXAIJ_MPI_Hash(Mat A, Mat B)
130: {
131: PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data, *b = (PetscConcat(Mat_MPI, TYPE) *)B->data;
133: PetscFunctionBegin;
134: /* Let's figure there's no harm done in doing the scatters for A now even if A != B */
135: PetscCall(MatAssemblyBegin_MPI_Hash(A, /*unused*/ MAT_FINAL_ASSEMBLY));
136: PetscCall(MatFinishScatterAndSetValues_MPI_Hash(A));
138: PetscCall(MatCopyHashToXAIJ(a->A, b->A));
139: PetscCall(MatCopyHashToXAIJ(a->B, b->B));
140: PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
141: PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
142: PetscFunctionReturn(PETSC_SUCCESS);
143: }
145: static PetscErrorCode MatDestroy_MPI_Hash(Mat A)
146: {
147: PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data;
149: PetscFunctionBegin;
150: PetscCall(MatStashDestroy_Private(&A->stash));
151: PetscCall(MatDestroy(&a->A));
152: PetscCall(MatDestroy(&a->B));
153: PetscCall((*a->cops.destroy)(A));
154: PetscFunctionReturn(PETSC_SUCCESS);
155: }
157: static PetscErrorCode MatZeroEntries_MPI_Hash(PETSC_UNUSED Mat A)
158: {
159: PetscFunctionBegin;
160: PetscFunctionReturn(PETSC_SUCCESS);
161: }
163: static PetscErrorCode MatSetRandom_MPI_Hash(Mat A, PETSC_UNUSED PetscRandom r)
164: {
165: SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Must set preallocation first");
166: }
168: static PetscErrorCode MatSetUp_MPI_Hash(Mat A)
169: {
170: PetscConcat(Mat_MPI, TYPE) *a = (PetscConcat(Mat_MPI, TYPE) *)A->data;
171: PetscMPIInt size;
172: #if !defined(TYPE_AIJ)
173: PetscInt bs;
174: #endif
176: PetscFunctionBegin;
177: PetscCall(PetscInfo(A, "Using hash-based MatSetValues() for MATMPI" PetscStringize(TYPE) " because no preallocation provided\n"));
178: PetscCall(PetscLayoutSetUp(A->rmap));
179: PetscCall(PetscLayoutSetUp(A->cmap));
180: if (A->rmap->bs < 1) A->rmap->bs = 1;
181: if (A->cmap->bs < 1) A->cmap->bs = 1;
182: PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size));
184: #if !defined(TYPE_AIJ)
185: PetscCall(MatGetBlockSize(A, &bs));
186: /* these values are set in MatMPISBAIJSetPreallocation() */
187: a->bs2 = bs * bs;
188: a->mbs = A->rmap->n / bs;
189: a->nbs = A->cmap->n / bs;
190: a->Mbs = A->rmap->N / bs;
191: a->Nbs = A->cmap->N / bs;
193: for (PetscInt i = 0; i <= a->size; i++) a->rangebs[i] = A->rmap->range[i] / bs;
194: a->rstartbs = A->rmap->rstart / bs;
195: a->rendbs = A->rmap->rend / bs;
196: a->cstartbs = A->cmap->rstart / bs;
197: a->cendbs = A->cmap->rend / bs;
198: PetscCall(MatStashCreate_Private(PetscObjectComm((PetscObject)A), A->rmap->bs, &A->bstash));
199: #endif
201: PetscCall(MatCreate(PETSC_COMM_SELF, &a->A));
202: PetscCall(MatSetSizes(a->A, A->rmap->n, A->cmap->n, A->rmap->n, A->cmap->n));
203: PetscCall(MatSetBlockSizesFromMats(a->A, A, A));
204: #if defined(SUB_TYPE_CUSPARSE)
205: PetscCall(MatSetType(a->A, MATSEQAIJCUSPARSE));
206: #else
207: PetscCall(MatSetType(a->A, PetscConcat(MATSEQ, TYPE)));
208: #endif
209: PetscCall(MatSetOption(a->A, MAT_STRUCTURE_ONLY, A->structure_only));
210: PetscCall(MatSetUp(a->A));
212: PetscCall(MatCreate(PETSC_COMM_SELF, &a->B));
213: PetscCall(MatSetSizes(a->B, A->rmap->n, size > 1 ? A->cmap->N : 0, A->rmap->n, size > 1 ? A->cmap->N : 0));
214: PetscCall(MatSetBlockSizesFromMats(a->B, A, A));
215: #if defined(TYPE_SBAIJ)
216: PetscCall(MatSetType(a->B, MATSEQBAIJ));
217: #else
218: #if defined(SUB_TYPE_CUSPARSE)
219: PetscCall(MatSetType(a->B, MATSEQAIJCUSPARSE));
220: #else
221: PetscCall(MatSetType(a->B, PetscConcat(MATSEQ, TYPE)));
222: #endif
223: #endif
224: PetscCall(MatSetOption(a->B, MAT_STRUCTURE_ONLY, A->structure_only));
225: PetscCall(MatSetUp(a->B));
227: /* keep a record of the operations so they can be reset when the hash handling is complete */
228: a->cops = A->ops[0];
229: A->ops->assemblybegin = MatAssemblyBegin_MPI_Hash;
230: A->ops->assemblyend = MatAssemblyEnd_MPI_Hash;
231: A->ops->setvalues = MatSetValues_MPI_Hash;
232: A->ops->destroy = MatDestroy_MPI_Hash;
233: A->ops->zeroentries = MatZeroEntries_MPI_Hash;
234: A->ops->setrandom = MatSetRandom_MPI_Hash;
235: A->ops->copyhashtoxaij = MatCopyHashToXAIJ_MPI_Hash;
236: A->ops->setvaluesblocked = NULL;
238: A->preallocated = PETSC_TRUE;
239: A->hash_active = PETSC_TRUE;
240: PetscFunctionReturn(PETSC_SUCCESS);
241: }