Actual source code: mpiov.c
1: /*
2: Routines to compute overlapping regions of a parallel MPI matrix
3: and to find submatrices that were shared across processors.
4: */
5: #include <../src/mat/impls/aij/seq/aij.h>
6: #include <../src/mat/impls/aij/mpi/mpiaij.h>
7: #include <petscsf.h>
9: static PetscErrorCode MatIncreaseOverlap_MPIAIJ_Once(Mat, PetscInt, IS *);
10: static PetscErrorCode MatIncreaseOverlap_MPIAIJ_Local(Mat, PetscInt, PetscBT *, PetscInt *, PetscInt **, PetscHMapI *);
11: static PetscErrorCode MatIncreaseOverlap_MPIAIJ_Receive(Mat, PetscInt, PetscInt **, PetscInt **, PetscInt *);
12: extern PetscErrorCode MatGetRow_MPIAIJ(Mat, PetscInt, PetscInt *, PetscInt **, PetscScalar **);
13: extern PetscErrorCode MatRestoreRow_MPIAIJ(Mat, PetscInt, PetscInt *, PetscInt **, PetscScalar **);
15: static PetscErrorCode MatIncreaseOverlap_MPIAIJ_Once_Scalable(Mat, PetscInt, IS *);
16: static PetscErrorCode MatIncreaseOverlap_MPIAIJ_Local_Scalable(Mat, PetscInt, IS *);
17: static PetscErrorCode MatIncreaseOverlap_MPIAIJ_Send_Scalable(Mat, PetscInt, PetscMPIInt, PetscMPIInt *, PetscInt *, PetscInt *, PetscInt **, PetscInt **);
18: static PetscErrorCode MatIncreaseOverlap_MPIAIJ_Receive_Scalable(Mat, PetscInt, IS *, PetscInt, PetscInt *);
20: /*
21: Takes a general IS and builds a block version of the IS that contains the given IS plus any needed values to fill out the blocks
23: The entire MatIncreaseOverlap_MPIAIJ() stack could be rewritten to respect the bs and it would offer higher performance but
24: that is a very major recoding job.
26: Possible scalability issues with this routine because it allocates space proportional to Nmax-Nmin
27: */
28: static PetscErrorCode ISAdjustForBlockSize(PetscInt bs, PetscInt imax, IS is[])
29: {
30: PetscFunctionBegin;
31: for (PetscInt i = 0; i < imax; i++) {
32: if (!is[i]) break;
33: PetscInt n = 0, N, Nmax, Nmin;
34: const PetscInt *idx;
35: PetscInt *nidx = NULL;
36: MPI_Comm comm;
37: PetscBT bt;
39: PetscCall(ISGetLocalSize(is[i], &N));
40: if (N > 0) { /* Nmax and Nmin are garbage for empty IS */
41: PetscCall(ISGetIndices(is[i], &idx));
42: PetscCall(ISGetMinMax(is[i], &Nmin, &Nmax));
43: Nmin = Nmin / bs;
44: Nmax = Nmax / bs;
45: PetscCall(PetscBTCreate(Nmax - Nmin, &bt));
46: for (PetscInt j = 0; j < N; j++) {
47: if (!PetscBTLookupSet(bt, idx[j] / bs - Nmin)) n++;
48: }
49: PetscCall(PetscMalloc1(n, &nidx));
50: n = 0;
51: PetscCall(PetscBTMemzero(Nmax - Nmin, bt));
52: for (PetscInt j = 0; j < N; j++) {
53: if (!PetscBTLookupSet(bt, idx[j] / bs - Nmin)) nidx[n++] = idx[j] / bs;
54: }
55: PetscCall(PetscBTDestroy(&bt));
56: PetscCall(ISRestoreIndices(is[i], &idx));
57: }
58: PetscCall(PetscObjectGetComm((PetscObject)is[i], &comm));
59: PetscCall(ISDestroy(is + i));
60: PetscCall(ISCreateBlock(comm, bs, n, nidx, PETSC_OWN_POINTER, is + i));
61: }
62: PetscFunctionReturn(PETSC_SUCCESS);
63: }
65: PetscErrorCode MatIncreaseOverlap_MPIAIJ(Mat C, PetscInt imax, IS is[], PetscInt ov)
66: {
67: PetscInt i;
69: PetscFunctionBegin;
70: PetscCheck(ov >= 0, PetscObjectComm((PetscObject)C), PETSC_ERR_ARG_OUTOFRANGE, "Negative overlap specified");
71: for (i = 0; i < ov; ++i) PetscCall(MatIncreaseOverlap_MPIAIJ_Once(C, imax, is));
72: if (C->rmap->bs > 1 && C->rmap->bs == C->cmap->bs) PetscCall(ISAdjustForBlockSize(C->rmap->bs, imax, is));
73: PetscFunctionReturn(PETSC_SUCCESS);
74: }
76: PetscErrorCode MatIncreaseOverlap_MPIAIJ_Scalable(Mat C, PetscInt imax, IS is[], PetscInt ov)
77: {
78: PetscInt i;
80: PetscFunctionBegin;
81: PetscCheck(ov >= 0, PetscObjectComm((PetscObject)C), PETSC_ERR_ARG_OUTOFRANGE, "Negative overlap specified");
82: for (i = 0; i < ov; ++i) PetscCall(MatIncreaseOverlap_MPIAIJ_Once_Scalable(C, imax, is));
83: if (C->rmap->bs > 1 && C->rmap->bs == C->cmap->bs) PetscCall(ISAdjustForBlockSize(C->rmap->bs, imax, is));
84: PetscFunctionReturn(PETSC_SUCCESS);
85: }
87: static PetscErrorCode MatIncreaseOverlap_MPIAIJ_Once_Scalable(Mat mat, PetscInt nidx, IS is[])
88: {
89: MPI_Comm comm;
90: PetscInt *length, length_i, tlength, *remoterows, nrrows, reducednrrows, *rrow_isids, j;
91: PetscInt *tosizes, *tosizes_temp, *toffsets, *fromsizes, *todata, *fromdata;
92: PetscInt nrecvrows, *sbsizes = NULL, *sbdata = NULL;
93: const PetscInt *indices_i, **indices;
94: PetscLayout rmap;
95: PetscMPIInt rank, size, *toranks, *fromranks, nto, nfrom, owner, *rrow_ranks;
96: PetscSF sf;
97: PetscSFNode *remote;
99: PetscFunctionBegin;
100: PetscCall(PetscObjectGetComm((PetscObject)mat, &comm));
101: PetscCallMPI(MPI_Comm_rank(comm, &rank));
102: PetscCallMPI(MPI_Comm_size(comm, &size));
103: /* get row map to determine where rows should be going */
104: PetscCall(MatGetLayouts(mat, &rmap, NULL));
105: /* retrieve IS data and put all together so that we
106: * can optimize communication
107: * */
108: PetscCall(PetscMalloc2(nidx, (PetscInt ***)&indices, nidx, &length));
109: tlength = 0;
110: for (PetscInt i = 0; i < nidx; i++) {
111: PetscCall(ISGetLocalSize(is[i], &length[i]));
112: tlength += length[i];
113: PetscCall(ISGetIndices(is[i], &indices[i]));
114: }
115: /* find these rows on remote processors */
116: PetscCall(PetscCalloc3(tlength, &remoterows, tlength, &rrow_ranks, tlength, &rrow_isids));
117: PetscCall(PetscCalloc3(size, &toranks, 2 * size, &tosizes, size, &tosizes_temp));
118: nrrows = 0;
119: for (PetscInt i = 0; i < nidx; i++) {
120: length_i = length[i];
121: indices_i = indices[i];
122: for (PetscInt j = 0; j < length_i; j++) {
123: owner = -1;
124: PetscCall(PetscLayoutFindOwner(rmap, indices_i[j], &owner));
125: /* remote processors */
126: if (owner != rank) {
127: tosizes_temp[owner]++; /* number of rows to owner */
128: rrow_ranks[nrrows] = owner; /* processor */
129: rrow_isids[nrrows] = i; /* is id */
130: remoterows[nrrows++] = indices_i[j]; /* row */
131: }
132: }
133: PetscCall(ISRestoreIndices(is[i], &indices[i]));
134: }
135: PetscCall(PetscFree2(*(PetscInt ***)&indices, length));
136: /* test if we need to exchange messages
137: * generally speaking, we do not need to exchange
138: * data when overlap is 1
139: * */
140: PetscCallMPI(MPIU_Allreduce(&nrrows, &reducednrrows, 1, MPIU_INT, MPI_MAX, comm));
141: /* we do not have any messages
142: * It usually corresponds to overlap 1
143: * */
144: if (!reducednrrows) {
145: PetscCall(PetscFree3(toranks, tosizes, tosizes_temp));
146: PetscCall(PetscFree3(remoterows, rrow_ranks, rrow_isids));
147: PetscCall(MatIncreaseOverlap_MPIAIJ_Local_Scalable(mat, nidx, is));
148: PetscFunctionReturn(PETSC_SUCCESS);
149: }
150: nto = 0;
151: /* send sizes and ranks for building a two-sided communication */
152: for (PetscMPIInt i = 0; i < size; i++) {
153: if (tosizes_temp[i]) {
154: tosizes[nto * 2] = tosizes_temp[i] * 2; /* size */
155: tosizes_temp[i] = nto; /* a map from processor to index */
156: toranks[nto++] = i; /* MPI process */
157: }
158: }
159: PetscCall(PetscMalloc1(nto + 1, &toffsets));
160: toffsets[0] = 0;
161: for (PetscInt i = 0; i < nto; i++) {
162: toffsets[i + 1] = toffsets[i] + tosizes[2 * i]; /* offsets */
163: tosizes[2 * i + 1] = toffsets[i]; /* offsets to send */
164: }
165: /* send information to other processors */
166: PetscCall(PetscCommBuildTwoSided(comm, 2, MPIU_INT, nto, toranks, tosizes, &nfrom, &fromranks, &fromsizes));
167: nrecvrows = 0;
168: for (PetscMPIInt i = 0; i < nfrom; i++) nrecvrows += fromsizes[2 * i];
169: PetscCall(PetscMalloc1(nrecvrows, &remote));
170: nrecvrows = 0;
171: for (PetscMPIInt i = 0; i < nfrom; i++) {
172: for (PetscInt j = 0; j < fromsizes[2 * i]; j++) {
173: remote[nrecvrows].rank = fromranks[i];
174: remote[nrecvrows++].index = fromsizes[2 * i + 1] + j;
175: }
176: }
177: PetscCall(PetscSFCreate(comm, &sf));
178: PetscCall(PetscSFSetGraph(sf, nrecvrows, nrecvrows, NULL, PETSC_OWN_POINTER, remote, PETSC_OWN_POINTER));
179: /* use two-sided communication by default since OPENMPI has some bugs for one-sided one */
180: PetscCall(PetscSFSetType(sf, PETSCSFBASIC));
181: PetscCall(PetscSFSetFromOptions(sf));
182: /* message pair <no of is, row> */
183: PetscCall(PetscCalloc2(2 * nrrows, &todata, nrecvrows, &fromdata));
184: for (PetscInt i = 0; i < nrrows; i++) {
185: owner = rrow_ranks[i]; /* process */
186: j = tosizes_temp[owner]; /* index */
187: todata[toffsets[j]++] = rrow_isids[i];
188: todata[toffsets[j]++] = remoterows[i];
189: }
190: PetscCall(PetscFree3(toranks, tosizes, tosizes_temp));
191: PetscCall(PetscFree3(remoterows, rrow_ranks, rrow_isids));
192: PetscCall(PetscFree(toffsets));
193: PetscCall(PetscSFBcastBegin(sf, MPIU_INT, todata, fromdata, MPI_REPLACE));
194: PetscCall(PetscSFBcastEnd(sf, MPIU_INT, todata, fromdata, MPI_REPLACE));
195: PetscCall(PetscSFDestroy(&sf));
196: /* send rows belonging to the remote so that then we could get the overlapping data back */
197: PetscCall(MatIncreaseOverlap_MPIAIJ_Send_Scalable(mat, nidx, nfrom, fromranks, fromsizes, fromdata, &sbsizes, &sbdata));
198: PetscCall(PetscFree2(todata, fromdata));
199: PetscCall(PetscFree(fromsizes));
200: PetscCall(PetscCommBuildTwoSided(comm, 2, MPIU_INT, nfrom, fromranks, sbsizes, &nto, &toranks, &tosizes));
201: PetscCall(PetscFree(fromranks));
202: nrecvrows = 0;
203: for (PetscInt i = 0; i < nto; i++) nrecvrows += tosizes[2 * i];
204: PetscCall(PetscCalloc1(nrecvrows, &todata));
205: PetscCall(PetscMalloc1(nrecvrows, &remote));
206: nrecvrows = 0;
207: for (PetscInt i = 0; i < nto; i++) {
208: for (PetscInt j = 0; j < tosizes[2 * i]; j++) {
209: remote[nrecvrows].rank = toranks[i];
210: remote[nrecvrows++].index = tosizes[2 * i + 1] + j;
211: }
212: }
213: PetscCall(PetscSFCreate(comm, &sf));
214: PetscCall(PetscSFSetGraph(sf, nrecvrows, nrecvrows, NULL, PETSC_OWN_POINTER, remote, PETSC_OWN_POINTER));
215: /* use two-sided communication by default since OPENMPI has some bugs for one-sided one */
216: PetscCall(PetscSFSetType(sf, PETSCSFBASIC));
217: PetscCall(PetscSFSetFromOptions(sf));
218: /* overlap communication and computation */
219: PetscCall(PetscSFBcastBegin(sf, MPIU_INT, sbdata, todata, MPI_REPLACE));
220: PetscCall(MatIncreaseOverlap_MPIAIJ_Local_Scalable(mat, nidx, is));
221: PetscCall(PetscSFBcastEnd(sf, MPIU_INT, sbdata, todata, MPI_REPLACE));
222: PetscCall(PetscSFDestroy(&sf));
223: PetscCall(PetscFree2(sbdata, sbsizes));
224: PetscCall(MatIncreaseOverlap_MPIAIJ_Receive_Scalable(mat, nidx, is, nrecvrows, todata));
225: PetscCall(PetscFree(toranks));
226: PetscCall(PetscFree(tosizes));
227: PetscCall(PetscFree(todata));
228: PetscFunctionReturn(PETSC_SUCCESS);
229: }
231: static PetscErrorCode MatIncreaseOverlap_MPIAIJ_Receive_Scalable(Mat mat, PetscInt nidx, IS is[], PetscInt nrecvs, PetscInt *recvdata)
232: {
233: PetscInt *isz, isz_i, i, j, is_id, data_size;
234: PetscInt col, lsize, max_lsize, *indices_temp, *indices_i;
235: const PetscInt *indices_i_temp;
236: MPI_Comm *iscomms;
238: PetscFunctionBegin;
239: max_lsize = 0;
240: PetscCall(PetscMalloc1(nidx, &isz));
241: for (i = 0; i < nidx; i++) {
242: PetscCall(ISGetLocalSize(is[i], &lsize));
243: max_lsize = lsize > max_lsize ? lsize : max_lsize;
244: isz[i] = lsize;
245: }
246: PetscCall(PetscMalloc2((max_lsize + nrecvs) * nidx, &indices_temp, nidx, &iscomms));
247: for (i = 0; i < nidx; i++) {
248: PetscCall(PetscCommDuplicate(PetscObjectComm((PetscObject)is[i]), &iscomms[i], NULL));
249: PetscCall(ISGetIndices(is[i], &indices_i_temp));
250: PetscCall(PetscArraycpy(PetscSafePointerPlusOffset(indices_temp, i * (max_lsize + nrecvs)), indices_i_temp, isz[i]));
251: PetscCall(ISRestoreIndices(is[i], &indices_i_temp));
252: PetscCall(ISDestroy(&is[i]));
253: }
254: /* retrieve information to get row id and its overlap */
255: for (i = 0; i < nrecvs;) {
256: is_id = recvdata[i++];
257: data_size = recvdata[i++];
258: indices_i = indices_temp + (max_lsize + nrecvs) * is_id;
259: isz_i = isz[is_id];
260: for (j = 0; j < data_size; j++) {
261: col = recvdata[i++];
262: indices_i[isz_i++] = col;
263: }
264: isz[is_id] = isz_i;
265: }
266: /* remove duplicate entities */
267: for (i = 0; i < nidx; i++) {
268: indices_i = PetscSafePointerPlusOffset(indices_temp, (max_lsize + nrecvs) * i);
269: isz_i = isz[i];
270: PetscCall(PetscSortRemoveDupsInt(&isz_i, indices_i));
271: PetscCall(ISCreateGeneral(iscomms[i], isz_i, indices_i, PETSC_COPY_VALUES, &is[i]));
272: PetscCall(PetscCommDestroy(&iscomms[i]));
273: }
274: PetscCall(PetscFree(isz));
275: PetscCall(PetscFree2(indices_temp, iscomms));
276: PetscFunctionReturn(PETSC_SUCCESS);
277: }
279: static PetscErrorCode MatIncreaseOverlap_MPIAIJ_Send_Scalable(Mat mat, PetscInt nidx, PetscMPIInt nfrom, PetscMPIInt *fromranks, PetscInt *fromsizes, PetscInt *fromrows, PetscInt **sbrowsizes, PetscInt **sbrows)
280: {
281: PetscLayout rmap, cmap;
282: PetscInt i, j, k, l, *rows_i, *rows_data_ptr, **rows_data, max_fszs, rows_pos, *rows_pos_i;
283: PetscInt is_id, tnz, an, bn, rstart, cstart, row, start, end, col, totalrows, *sbdata;
284: PetscInt *indv_counts, indvc_ij, *sbsizes, *indices_tmp, *offsets;
285: const PetscInt *gcols, *ai, *aj, *bi, *bj;
286: Mat amat, bmat;
287: PetscMPIInt rank;
288: PetscBool done;
289: MPI_Comm comm;
291: PetscFunctionBegin;
292: PetscCall(PetscObjectGetComm((PetscObject)mat, &comm));
293: PetscCallMPI(MPI_Comm_rank(comm, &rank));
294: PetscCall(MatMPIAIJGetSeqAIJ(mat, &amat, &bmat, &gcols));
295: /* Even if the mat is symmetric, we still assume it is not symmetric */
296: PetscCall(MatGetRowIJ(amat, 0, PETSC_FALSE, PETSC_FALSE, &an, &ai, &aj, &done));
297: PetscCheck(done, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "can not get row IJ ");
298: PetscCall(MatGetRowIJ(bmat, 0, PETSC_FALSE, PETSC_FALSE, &bn, &bi, &bj, &done));
299: PetscCheck(done, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "can not get row IJ ");
300: /* total number of nonzero values is used to estimate the memory usage in the next step */
301: tnz = ai[an] + bi[bn];
302: PetscCall(MatGetLayouts(mat, &rmap, &cmap));
303: PetscCall(PetscLayoutGetRange(rmap, &rstart, NULL));
304: PetscCall(PetscLayoutGetRange(cmap, &cstart, NULL));
305: /* to find the longest message */
306: max_fszs = 0;
307: for (i = 0; i < nfrom; i++) max_fszs = fromsizes[2 * i] > max_fszs ? fromsizes[2 * i] : max_fszs;
308: /* better way to estimate number of nonzero in the mat??? */
309: PetscCall(PetscCalloc5(max_fszs * nidx, &rows_data_ptr, nidx, &rows_data, nidx, &rows_pos_i, nfrom * nidx, &indv_counts, tnz, &indices_tmp));
310: for (i = 0; i < nidx; i++) rows_data[i] = PetscSafePointerPlusOffset(rows_data_ptr, max_fszs * i);
311: rows_pos = 0;
312: totalrows = 0;
313: for (i = 0; i < nfrom; i++) {
314: PetscCall(PetscArrayzero(rows_pos_i, nidx));
315: /* group data together */
316: for (j = 0; j < fromsizes[2 * i]; j += 2) {
317: is_id = fromrows[rows_pos++]; /* no of is */
318: rows_i = rows_data[is_id];
319: rows_i[rows_pos_i[is_id]++] = fromrows[rows_pos++]; /* row */
320: }
321: /* estimate a space to avoid multiple allocations */
322: for (j = 0; j < nidx; j++) {
323: indvc_ij = 0;
324: rows_i = rows_data[j];
325: for (l = 0; l < rows_pos_i[j]; l++) {
326: row = rows_i[l] - rstart;
327: start = ai[row];
328: end = ai[row + 1];
329: for (k = start; k < end; k++) { /* Amat */
330: col = aj[k] + cstart;
331: indices_tmp[indvc_ij++] = col; /* do not count the rows from the original rank */
332: }
333: start = bi[row];
334: end = bi[row + 1];
335: for (k = start; k < end; k++) { /* Bmat */
336: col = gcols[bj[k]];
337: indices_tmp[indvc_ij++] = col;
338: }
339: }
340: PetscCall(PetscSortRemoveDupsInt(&indvc_ij, indices_tmp));
341: indv_counts[i * nidx + j] = indvc_ij;
342: totalrows += indvc_ij;
343: }
344: }
345: /* message triple <no of is, number of rows, rows> */
346: PetscCall(PetscCalloc2(totalrows + nidx * nfrom * 2, &sbdata, 2 * nfrom, &sbsizes));
347: totalrows = 0;
348: rows_pos = 0;
349: /* use this code again */
350: for (i = 0; i < nfrom; i++) {
351: PetscCall(PetscArrayzero(rows_pos_i, nidx));
352: for (j = 0; j < fromsizes[2 * i]; j += 2) {
353: is_id = fromrows[rows_pos++];
354: rows_i = rows_data[is_id];
355: rows_i[rows_pos_i[is_id]++] = fromrows[rows_pos++];
356: }
357: /* add data */
358: for (j = 0; j < nidx; j++) {
359: if (!indv_counts[i * nidx + j]) continue;
360: indvc_ij = 0;
361: sbdata[totalrows++] = j;
362: sbdata[totalrows++] = indv_counts[i * nidx + j];
363: sbsizes[2 * i] += 2;
364: rows_i = rows_data[j];
365: for (l = 0; l < rows_pos_i[j]; l++) {
366: row = rows_i[l] - rstart;
367: start = ai[row];
368: end = ai[row + 1];
369: for (k = start; k < end; k++) { /* Amat */
370: col = aj[k] + cstart;
371: indices_tmp[indvc_ij++] = col;
372: }
373: start = bi[row];
374: end = bi[row + 1];
375: for (k = start; k < end; k++) { /* Bmat */
376: col = gcols[bj[k]];
377: indices_tmp[indvc_ij++] = col;
378: }
379: }
380: PetscCall(PetscSortRemoveDupsInt(&indvc_ij, indices_tmp));
381: sbsizes[2 * i] += indvc_ij;
382: PetscCall(PetscArraycpy(sbdata + totalrows, indices_tmp, indvc_ij));
383: totalrows += indvc_ij;
384: }
385: }
386: PetscCall(PetscMalloc1(nfrom + 1, &offsets));
387: offsets[0] = 0;
388: for (i = 0; i < nfrom; i++) {
389: offsets[i + 1] = offsets[i] + sbsizes[2 * i];
390: sbsizes[2 * i + 1] = offsets[i];
391: }
392: PetscCall(PetscFree(offsets));
393: if (sbrowsizes) *sbrowsizes = sbsizes;
394: if (sbrows) *sbrows = sbdata;
395: PetscCall(PetscFree5(rows_data_ptr, rows_data, rows_pos_i, indv_counts, indices_tmp));
396: PetscCall(MatRestoreRowIJ(amat, 0, PETSC_FALSE, PETSC_FALSE, &an, &ai, &aj, &done));
397: PetscCall(MatRestoreRowIJ(bmat, 0, PETSC_FALSE, PETSC_FALSE, &bn, &bi, &bj, &done));
398: PetscFunctionReturn(PETSC_SUCCESS);
399: }
401: static PetscErrorCode MatIncreaseOverlap_MPIAIJ_Local_Scalable(Mat mat, PetscInt nidx, IS is[])
402: {
403: const PetscInt *gcols, *ai, *aj, *bi, *bj, *indices;
404: PetscInt tnz, an, bn, i, j, row, start, end, rstart, cstart, col, k, *indices_temp;
405: PetscInt lsize, lsize_tmp;
406: PetscMPIInt rank, owner;
407: Mat amat, bmat;
408: PetscBool done;
409: PetscLayout cmap, rmap;
410: MPI_Comm comm;
412: PetscFunctionBegin;
413: PetscCall(PetscObjectGetComm((PetscObject)mat, &comm));
414: PetscCallMPI(MPI_Comm_rank(comm, &rank));
415: PetscCall(MatMPIAIJGetSeqAIJ(mat, &amat, &bmat, &gcols));
416: PetscCall(MatGetRowIJ(amat, 0, PETSC_FALSE, PETSC_FALSE, &an, &ai, &aj, &done));
417: PetscCheck(done, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "can not get row IJ ");
418: PetscCall(MatGetRowIJ(bmat, 0, PETSC_FALSE, PETSC_FALSE, &bn, &bi, &bj, &done));
419: PetscCheck(done, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "can not get row IJ ");
420: /* is it a safe way to compute number of nonzero values ? */
421: tnz = ai[an] + bi[bn];
422: PetscCall(MatGetLayouts(mat, &rmap, &cmap));
423: PetscCall(PetscLayoutGetRange(rmap, &rstart, NULL));
424: PetscCall(PetscLayoutGetRange(cmap, &cstart, NULL));
425: /* it is a better way to estimate memory than the old implementation
426: * where global size of matrix is used
427: * */
428: PetscCall(PetscMalloc1(tnz, &indices_temp));
429: for (i = 0; i < nidx; i++) {
430: MPI_Comm iscomm;
432: PetscCall(ISGetLocalSize(is[i], &lsize));
433: PetscCall(ISGetIndices(is[i], &indices));
434: lsize_tmp = 0;
435: for (j = 0; j < lsize; j++) {
436: owner = -1;
437: row = indices[j];
438: PetscCall(PetscLayoutFindOwner(rmap, row, &owner));
439: if (owner != rank) continue;
440: /* local number */
441: row -= rstart;
442: start = ai[row];
443: end = ai[row + 1];
444: for (k = start; k < end; k++) { /* Amat */
445: col = aj[k] + cstart;
446: indices_temp[lsize_tmp++] = col;
447: }
448: start = bi[row];
449: end = bi[row + 1];
450: for (k = start; k < end; k++) { /* Bmat */
451: col = gcols[bj[k]];
452: indices_temp[lsize_tmp++] = col;
453: }
454: }
455: PetscCall(ISRestoreIndices(is[i], &indices));
456: PetscCall(PetscCommDuplicate(PetscObjectComm((PetscObject)is[i]), &iscomm, NULL));
457: PetscCall(ISDestroy(&is[i]));
458: PetscCall(PetscSortRemoveDupsInt(&lsize_tmp, indices_temp));
459: PetscCall(ISCreateGeneral(iscomm, lsize_tmp, indices_temp, PETSC_COPY_VALUES, &is[i]));
460: PetscCall(PetscCommDestroy(&iscomm));
461: }
462: PetscCall(PetscFree(indices_temp));
463: PetscCall(MatRestoreRowIJ(amat, 0, PETSC_FALSE, PETSC_FALSE, &an, &ai, &aj, &done));
464: PetscCall(MatRestoreRowIJ(bmat, 0, PETSC_FALSE, PETSC_FALSE, &bn, &bi, &bj, &done));
465: PetscFunctionReturn(PETSC_SUCCESS);
466: }
468: /*
469: Sample message format:
470: If a processor A wants processor B to process some elements corresponding
471: to index sets is[1],is[5]
472: mesg [0] = 2 (no of index sets in the mesg)
473: -----------
474: mesg [1] = 1 => is[1]
475: mesg [2] = sizeof(is[1]);
476: -----------
477: mesg [3] = 5 => is[5]
478: mesg [4] = sizeof(is[5]);
479: -----------
480: mesg [5]
481: mesg [n] datas[1]
482: -----------
483: mesg[n+1]
484: mesg[m] data(is[5])
485: -----------
487: Notes:
488: nrqs - no of requests sent (or to be sent out)
489: nrqr - no of requests received (which have to be or which have been processed)
490: */
491: static PetscErrorCode MatIncreaseOverlap_MPIAIJ_Once(Mat C, PetscInt imax, IS is[])
492: {
493: Mat_MPIAIJ *c = (Mat_MPIAIJ *)C->data;
494: PetscMPIInt *w1, *w2, nrqr, *w3, *w4, *onodes1, *olengths1, *onodes2, *olengths2;
495: const PetscInt **idx, *idx_i;
496: PetscInt *n, **data, len;
497: #if PetscDefined(USE_CTABLE)
498: PetscHMapI *table_data, table_data_i;
499: PetscInt *tdata, tcount, tcount_max;
500: #else
501: PetscInt *data_i, *d_p;
502: #endif
503: PetscMPIInt size, rank, tag1, tag2, proc = 0, nrqs, *pa;
504: PetscInt M, k, **rbuf, row, msz, **outdat, **ptr, *isz1;
505: PetscInt *ctr, *tmp, *isz, **xdata, **rbuf2;
506: PetscBT *table;
507: MPI_Comm comm;
508: MPI_Request *s_waits1, *r_waits1, *s_waits2, *r_waits2;
509: MPI_Status *recv_status;
510: MPI_Comm *iscomms;
511: PetscByte *t_p;
513: PetscFunctionBegin;
514: PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
515: size = c->size;
516: rank = c->rank;
517: M = C->rmap->N;
519: PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag1));
520: PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag2));
522: PetscCall(PetscMalloc2(imax, (PetscInt ***)&idx, imax, &n));
524: for (PetscInt i = 0; i < imax; i++) {
525: PetscCall(ISGetIndices(is[i], &idx[i]));
526: PetscCall(ISGetLocalSize(is[i], &n[i]));
527: }
529: /* evaluate communication - mesg to who,length of mesg, and buffer space
530: required. Based on this, buffers are allocated, and data copied into them */
531: PetscCall(PetscCalloc4(size, &w1, size, &w2, size, &w3, size, &w4));
532: for (PetscInt i = 0; i < imax; i++) {
533: PetscCall(PetscArrayzero(w4, size)); /* initialise work vector*/
534: idx_i = idx[i];
535: len = n[i];
536: for (PetscInt j = 0; j < len; j++) {
537: row = idx_i[j];
538: PetscCheck(row >= 0, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Index set cannot have negative entries");
539: PetscCall(PetscLayoutFindOwner(C->rmap, row, &proc));
540: w4[proc]++;
541: }
542: for (PetscMPIInt j = 0; j < size; j++) {
543: if (w4[j]) {
544: w1[j] += w4[j];
545: w3[j]++;
546: }
547: }
548: }
550: nrqs = 0; /* no of outgoing messages */
551: msz = 0; /* total mesg length (for all proc */
552: w1[rank] = 0; /* no mesg sent to intself */
553: w3[rank] = 0;
554: for (PetscMPIInt i = 0; i < size; i++) {
555: if (w1[i]) {
556: w2[i] = 1;
557: nrqs++;
558: } /* there exists a message to proc i */
559: }
560: /* pa - is list of processors to communicate with */
561: PetscCall(PetscMalloc1(nrqs, &pa));
562: for (PetscMPIInt i = 0, j = 0; i < size; i++) {
563: if (w1[i]) {
564: pa[j] = i;
565: j++;
566: }
567: }
569: /* Each message would have a header = 1 + 2*(no of IS) + data */
570: for (PetscMPIInt i = 0; i < nrqs; i++) {
571: PetscMPIInt j = pa[i];
572: w1[j] += w2[j] + 2 * w3[j];
573: msz += w1[j];
574: }
576: /* Determine the number of messages to expect, their lengths, from from-ids */
577: PetscCall(PetscGatherNumberOfMessages(comm, w2, w1, &nrqr));
578: PetscCall(PetscGatherMessageLengths(comm, nrqs, nrqr, w1, &onodes1, &olengths1));
580: /* Now post the Irecvs corresponding to these messages */
581: PetscCall(PetscPostIrecvInt(comm, tag1, nrqr, onodes1, olengths1, &rbuf, &r_waits1));
583: /* Allocate Memory for outgoing messages */
584: PetscCall(PetscMalloc4(size, &outdat, size, &ptr, msz, &tmp, size, &ctr));
585: PetscCall(PetscArrayzero(outdat, size));
586: PetscCall(PetscArrayzero(ptr, size));
588: {
589: PetscInt *iptr = tmp, ict = 0;
590: for (PetscMPIInt i = 0; i < nrqs; i++) {
591: PetscMPIInt j = pa[i];
592: iptr += ict;
593: outdat[j] = iptr;
594: ict = w1[j];
595: }
596: }
598: /* Form the outgoing messages */
599: /* plug in the headers */
600: for (PetscMPIInt i = 0; i < nrqs; i++) {
601: PetscMPIInt j = pa[i];
602: outdat[j][0] = 0;
603: PetscCall(PetscArrayzero(outdat[j] + 1, 2 * w3[j]));
604: ptr[j] = outdat[j] + 2 * w3[j] + 1;
605: }
607: /* Memory for doing local proc's work */
608: {
609: PetscInt M_BPB_imax = 0;
610: #if PetscDefined(USE_CTABLE)
611: PetscCall(PetscIntMultError(M / PETSC_BITS_PER_BYTE + 1, imax, &M_BPB_imax));
612: PetscCall(PetscMalloc1(imax, &table_data));
613: for (PetscInt i = 0; i < imax; i++) PetscCall(PetscHMapICreateWithSize(n[i], table_data + i));
614: PetscCall(PetscCalloc4(imax, &table, imax, &data, imax, &isz, M_BPB_imax, &t_p));
615: for (PetscInt i = 0; i < imax; i++) table[i] = t_p + (M / PETSC_BITS_PER_BYTE + 1) * i;
616: #else
617: PetscInt Mimax = 0;
618: PetscCall(PetscIntMultError(M, imax, &Mimax));
619: PetscCall(PetscIntMultError(M / PETSC_BITS_PER_BYTE + 1, imax, &M_BPB_imax));
620: PetscCall(PetscCalloc5(imax, &table, imax, &data, imax, &isz, Mimax, &d_p, M_BPB_imax, &t_p));
621: for (PetscInt i = 0; i < imax; i++) {
622: table[i] = t_p + (M / PETSC_BITS_PER_BYTE + 1) * i;
623: data[i] = d_p + M * i;
624: }
625: #endif
626: }
628: /* Parse the IS and update local tables and the outgoing buf with the data */
629: {
630: PetscInt n_i, isz_i, *outdat_j, ctr_j;
631: PetscBT table_i;
633: for (PetscInt i = 0; i < imax; i++) {
634: PetscCall(PetscArrayzero(ctr, size));
635: n_i = n[i];
636: table_i = table[i];
637: idx_i = idx[i];
638: #if PetscDefined(USE_CTABLE)
639: table_data_i = table_data[i];
640: #else
641: data_i = data[i];
642: #endif
643: isz_i = isz[i];
644: for (PetscInt j = 0; j < n_i; j++) { /* parse the indices of each IS */
645: row = idx_i[j];
646: PetscCall(PetscLayoutFindOwner(C->rmap, row, &proc));
647: if (proc != rank) { /* copy to the outgoing buffer */
648: ctr[proc]++;
649: *ptr[proc] = row;
650: ptr[proc]++;
651: } else if (!PetscBTLookupSet(table_i, row)) {
652: #if PetscDefined(USE_CTABLE)
653: PetscCall(PetscHMapISet(table_data_i, row + 1, isz_i + 1));
654: #else
655: data_i[isz_i] = row; /* Update the local table */
656: #endif
657: isz_i++;
658: }
659: }
660: /* Update the headers for the current IS */
661: for (PetscMPIInt j = 0; j < size; j++) { /* Can Optimise this loop by using pa[] */
662: if ((ctr_j = ctr[j])) {
663: outdat_j = outdat[j];
664: k = ++outdat_j[0];
665: outdat_j[2 * k] = ctr_j;
666: outdat_j[2 * k - 1] = i;
667: }
668: }
669: isz[i] = isz_i;
670: }
671: }
673: /* Now post the sends */
674: PetscCall(PetscMalloc1(nrqs, &s_waits1));
675: for (PetscMPIInt i = 0; i < nrqs; ++i) {
676: PetscMPIInt j = pa[i];
677: PetscCallMPI(MPIU_Isend(outdat[j], w1[j], MPIU_INT, j, tag1, comm, s_waits1 + i));
678: }
680: /* No longer need the original indices */
681: PetscCall(PetscMalloc1(imax, &iscomms));
682: for (PetscInt i = 0; i < imax; ++i) {
683: PetscCall(ISRestoreIndices(is[i], idx + i));
684: PetscCall(PetscCommDuplicate(PetscObjectComm((PetscObject)is[i]), &iscomms[i], NULL));
685: }
686: PetscCall(PetscFree2(*(PetscInt ***)&idx, n));
688: for (PetscInt i = 0; i < imax; ++i) PetscCall(ISDestroy(&is[i]));
690: /* Do Local work */
691: #if PetscDefined(USE_CTABLE)
692: PetscCall(MatIncreaseOverlap_MPIAIJ_Local(C, imax, table, isz, NULL, table_data));
693: #else
694: PetscCall(MatIncreaseOverlap_MPIAIJ_Local(C, imax, table, isz, data, NULL));
695: #endif
697: /* Receive messages */
698: PetscCall(PetscMalloc1(nrqr, &recv_status));
699: PetscCallMPI(MPI_Waitall(nrqr, r_waits1, recv_status));
700: PetscCallMPI(MPI_Waitall(nrqs, s_waits1, MPI_STATUSES_IGNORE));
702: /* Phase 1 sends are complete - deallocate buffers */
703: PetscCall(PetscFree4(outdat, ptr, tmp, ctr));
704: PetscCall(PetscFree4(w1, w2, w3, w4));
706: PetscCall(PetscMalloc1(nrqr, &xdata));
707: PetscCall(PetscMalloc1(nrqr, &isz1));
708: PetscCall(MatIncreaseOverlap_MPIAIJ_Receive(C, nrqr, rbuf, xdata, isz1));
709: PetscCall(PetscFree(rbuf[0]));
710: PetscCall(PetscFree(rbuf));
712: /* Send the data back */
713: /* Do a global reduction to know the buffer space req for incoming messages */
714: {
715: PetscMPIInt *rw1;
717: PetscCall(PetscCalloc1(size, &rw1));
718: for (PetscMPIInt i = 0; i < nrqr; ++i) {
719: proc = recv_status[i].MPI_SOURCE;
720: PetscCheck(proc == onodes1[i], PETSC_COMM_SELF, PETSC_ERR_PLIB, "MPI_SOURCE mismatch");
721: PetscCall(PetscMPIIntCast(isz1[i], &rw1[proc]));
722: }
723: PetscCall(PetscFree(onodes1));
724: PetscCall(PetscFree(olengths1));
726: /* Determine the number of messages to expect, their lengths, from from-ids */
727: PetscCall(PetscGatherMessageLengths(comm, nrqr, nrqs, rw1, &onodes2, &olengths2));
728: PetscCall(PetscFree(rw1));
729: }
730: /* Now post the Irecvs corresponding to these messages */
731: PetscCall(PetscPostIrecvInt(comm, tag2, nrqs, onodes2, olengths2, &rbuf2, &r_waits2));
733: /* Now post the sends */
734: PetscCall(PetscMalloc1(nrqr, &s_waits2));
735: for (PetscMPIInt i = 0; i < nrqr; ++i) PetscCallMPI(MPIU_Isend(xdata[i], isz1[i], MPIU_INT, recv_status[i].MPI_SOURCE, tag2, comm, s_waits2 + i));
737: /* receive work done on other processors */
738: {
739: PetscInt is_no, ct1, max, *rbuf2_i, isz_i, jmax;
740: PetscMPIInt idex;
741: PetscBT table_i;
743: for (PetscMPIInt i = 0; i < nrqs; ++i) {
744: PetscCallMPI(MPI_Waitany(nrqs, r_waits2, &idex, MPI_STATUS_IGNORE));
745: /* Process the message */
746: rbuf2_i = rbuf2[idex];
747: ct1 = 2 * rbuf2_i[0] + 1;
748: jmax = rbuf2[idex][0];
749: for (PetscInt j = 1; j <= jmax; j++) {
750: max = rbuf2_i[2 * j];
751: is_no = rbuf2_i[2 * j - 1];
752: isz_i = isz[is_no];
753: table_i = table[is_no];
754: #if PetscDefined(USE_CTABLE)
755: table_data_i = table_data[is_no];
756: #else
757: data_i = data[is_no];
758: #endif
759: for (k = 0; k < max; k++, ct1++) {
760: row = rbuf2_i[ct1];
761: if (!PetscBTLookupSet(table_i, row)) {
762: #if PetscDefined(USE_CTABLE)
763: PetscCall(PetscHMapISet(table_data_i, row + 1, isz_i + 1));
764: #else
765: data_i[isz_i] = row;
766: #endif
767: isz_i++;
768: }
769: }
770: isz[is_no] = isz_i;
771: }
772: }
774: PetscCallMPI(MPI_Waitall(nrqr, s_waits2, MPI_STATUSES_IGNORE));
775: }
777: #if PetscDefined(USE_CTABLE)
778: tcount_max = 0;
779: for (PetscInt i = 0; i < imax; ++i) {
780: table_data_i = table_data[i];
781: PetscCall(PetscHMapIGetSize(table_data_i, &tcount));
782: if (tcount_max < tcount) tcount_max = tcount;
783: }
784: PetscCall(PetscMalloc1(tcount_max, &tdata));
785: #endif
787: for (PetscInt i = 0; i < imax; ++i) {
788: #if PetscDefined(USE_CTABLE)
789: PetscHashIter tpos;
790: PetscInt j;
792: table_data_i = table_data[i];
793: PetscHashIterBegin(table_data_i, tpos);
794: while (!PetscHashIterAtEnd(table_data_i, tpos)) {
795: PetscHashIterGetKey(table_data_i, tpos, k);
796: PetscHashIterGetVal(table_data_i, tpos, j);
797: PetscHashIterNext(table_data_i, tpos);
798: tdata[--j] = --k;
799: }
800: PetscCall(ISCreateGeneral(iscomms[i], isz[i], tdata, PETSC_COPY_VALUES, is + i));
801: #else
802: PetscCall(ISCreateGeneral(iscomms[i], isz[i], data[i], PETSC_COPY_VALUES, is + i));
803: #endif
804: PetscCall(PetscCommDestroy(&iscomms[i]));
805: }
807: PetscCall(PetscFree(iscomms));
808: PetscCall(PetscFree(onodes2));
809: PetscCall(PetscFree(olengths2));
811: PetscCall(PetscFree(pa));
812: PetscCall(PetscFree(rbuf2[0]));
813: PetscCall(PetscFree(rbuf2));
814: PetscCall(PetscFree(s_waits1));
815: PetscCall(PetscFree(r_waits1));
816: PetscCall(PetscFree(s_waits2));
817: PetscCall(PetscFree(r_waits2));
818: PetscCall(PetscFree(recv_status));
819: if (xdata) {
820: PetscCall(PetscFree(xdata[0]));
821: PetscCall(PetscFree(xdata));
822: }
823: PetscCall(PetscFree(isz1));
824: #if PetscDefined(USE_CTABLE)
825: for (PetscInt i = 0; i < imax; i++) PetscCall(PetscHMapIDestroy(table_data + i));
826: PetscCall(PetscFree(table_data));
827: PetscCall(PetscFree(tdata));
828: PetscCall(PetscFree4(table, data, isz, t_p));
829: #else
830: PetscCall(PetscFree5(table, data, isz, d_p, t_p));
831: #endif
832: PetscFunctionReturn(PETSC_SUCCESS);
833: }
835: /*
836: MatIncreaseOverlap_MPIAIJ_Local - Called by MatincreaseOverlap, to do
837: the work on the local processor.
839: Inputs:
840: C - MAT_MPIAIJ;
841: imax - total no of index sets processed at a time;
842: table - an array of char - size = m bits.
844: Output:
845: isz - array containing the count of the solution elements corresponding
846: to each index set;
847: data or table_data - pointer to the solutions
848: */
849: static PetscErrorCode MatIncreaseOverlap_MPIAIJ_Local(Mat C, PetscInt imax, PetscBT *table, PetscInt *isz, PetscInt **data, PetscHMapI *table_data)
850: {
851: Mat_MPIAIJ *c = (Mat_MPIAIJ *)C->data;
852: Mat A = c->A, B = c->B;
853: Mat_SeqAIJ *a = (Mat_SeqAIJ *)A->data, *b = (Mat_SeqAIJ *)B->data;
854: PetscInt start, end, val, max, rstart, cstart, *ai, *aj;
855: PetscInt *bi, *bj, *garray, i, j, k, row, isz_i;
856: PetscBT table_i;
857: #if PetscDefined(USE_CTABLE)
858: PetscHMapI table_data_i;
859: PetscHashIter tpos;
860: PetscInt tcount, *tdata;
861: #else
862: PetscInt *data_i;
863: #endif
865: PetscFunctionBegin;
866: rstart = C->rmap->rstart;
867: cstart = C->cmap->rstart;
868: ai = a->i;
869: aj = a->j;
870: bi = b->i;
871: bj = b->j;
872: garray = c->garray;
874: for (i = 0; i < imax; i++) {
875: #if PetscDefined(USE_CTABLE)
876: /* copy existing entries of table_data_i into tdata[] */
877: table_data_i = table_data[i];
878: PetscCall(PetscHMapIGetSize(table_data_i, &tcount));
879: PetscCheck(tcount == isz[i], PETSC_COMM_SELF, PETSC_ERR_PLIB, " tcount %" PetscInt_FMT " != isz[%" PetscInt_FMT "] %" PetscInt_FMT, tcount, i, isz[i]);
881: PetscCall(PetscMalloc1(tcount, &tdata));
882: PetscHashIterBegin(table_data_i, tpos);
883: while (!PetscHashIterAtEnd(table_data_i, tpos)) {
884: PetscHashIterGetKey(table_data_i, tpos, row);
885: PetscHashIterGetVal(table_data_i, tpos, j);
886: PetscHashIterNext(table_data_i, tpos);
887: tdata[--j] = --row;
888: PetscCheck(j <= tcount - 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, " j %" PetscInt_FMT " >= tcount %" PetscInt_FMT, j, tcount);
889: }
890: #else
891: data_i = data[i];
892: #endif
893: table_i = table[i];
894: isz_i = isz[i];
895: max = isz[i];
897: for (j = 0; j < max; j++) {
898: #if PetscDefined(USE_CTABLE)
899: row = tdata[j] - rstart;
900: #else
901: row = data_i[j] - rstart;
902: #endif
903: start = ai[row];
904: end = ai[row + 1];
905: for (k = start; k < end; k++) { /* Amat */
906: val = aj[k] + cstart;
907: if (!PetscBTLookupSet(table_i, val)) {
908: #if PetscDefined(USE_CTABLE)
909: PetscCall(PetscHMapISet(table_data_i, val + 1, isz_i + 1));
910: #else
911: data_i[isz_i] = val;
912: #endif
913: isz_i++;
914: }
915: }
916: start = bi[row];
917: end = bi[row + 1];
918: for (k = start; k < end; k++) { /* Bmat */
919: val = garray[bj[k]];
920: if (!PetscBTLookupSet(table_i, val)) {
921: #if PetscDefined(USE_CTABLE)
922: PetscCall(PetscHMapISet(table_data_i, val + 1, isz_i + 1));
923: #else
924: data_i[isz_i] = val;
925: #endif
926: isz_i++;
927: }
928: }
929: }
930: isz[i] = isz_i;
932: #if PetscDefined(USE_CTABLE)
933: PetscCall(PetscFree(tdata));
934: #endif
935: }
936: PetscFunctionReturn(PETSC_SUCCESS);
937: }
939: /*
940: MatIncreaseOverlap_MPIAIJ_Receive - Process the received messages,
941: and return the output
943: Input:
944: C - the matrix
945: nrqr - no of messages being processed.
946: rbuf - an array of pointers to the received requests
948: Output:
949: xdata - array of messages to be sent back
950: isz1 - size of each message
952: For better efficiency perhaps we should malloc separately each xdata[i],
953: then if a remalloc is required we need only copy the data for that one row
954: rather than all previous rows as it is now where a single large chunk of
955: memory is used.
957: */
958: static PetscErrorCode MatIncreaseOverlap_MPIAIJ_Receive(Mat C, PetscInt nrqr, PetscInt **rbuf, PetscInt **xdata, PetscInt *isz1)
959: {
960: Mat_MPIAIJ *c = (Mat_MPIAIJ *)C->data;
961: Mat A = c->A, B = c->B;
962: Mat_SeqAIJ *a = (Mat_SeqAIJ *)A->data, *b = (Mat_SeqAIJ *)B->data;
963: PetscInt rstart, cstart, *ai, *aj, *bi, *bj, *garray, i, j, k;
964: PetscInt row, total_sz, ct, ct1, ct2, ct3, mem_estimate, oct2, l, start, end;
965: PetscInt val, max1, max2, m, no_malloc = 0, *tmp, new_estimate, ctr;
966: PetscInt *rbuf_i, kmax, rbuf_0;
967: PetscBT xtable;
969: PetscFunctionBegin;
970: m = C->rmap->N;
971: rstart = C->rmap->rstart;
972: cstart = C->cmap->rstart;
973: ai = a->i;
974: aj = a->j;
975: bi = b->i;
976: bj = b->j;
977: garray = c->garray;
979: for (i = 0, ct = 0, total_sz = 0; i < nrqr; ++i) {
980: rbuf_i = rbuf[i];
981: rbuf_0 = rbuf_i[0];
982: ct += rbuf_0;
983: for (j = 1; j <= rbuf_0; j++) total_sz += rbuf_i[2 * j];
984: }
986: if (C->rmap->n) max1 = ct * (a->nz + b->nz) / C->rmap->n;
987: else max1 = 1;
988: mem_estimate = 3 * ((total_sz > max1 ? total_sz : max1) + 1);
989: if (nrqr) {
990: PetscCall(PetscMalloc1(mem_estimate, &xdata[0]));
991: ++no_malloc;
992: }
993: PetscCall(PetscBTCreate(m, &xtable));
994: PetscCall(PetscArrayzero(isz1, nrqr));
996: ct3 = 0;
997: for (i = 0; i < nrqr; i++) { /* for easch mesg from proc i */
998: rbuf_i = rbuf[i];
999: rbuf_0 = rbuf_i[0];
1000: ct1 = 2 * rbuf_0 + 1;
1001: ct2 = ct1;
1002: ct3 += ct1;
1003: for (j = 1; j <= rbuf_0; j++) { /* for each IS from proc i*/
1004: PetscCall(PetscBTMemzero(m, xtable));
1005: oct2 = ct2;
1006: kmax = rbuf_i[2 * j];
1007: for (k = 0; k < kmax; k++, ct1++) {
1008: row = rbuf_i[ct1];
1009: if (!PetscBTLookupSet(xtable, row)) {
1010: if (!(ct3 < mem_estimate)) {
1011: new_estimate = (PetscInt)(1.5 * mem_estimate) + 1;
1012: PetscCall(PetscMalloc1(new_estimate, &tmp));
1013: PetscCall(PetscArraycpy(tmp, xdata[0], mem_estimate));
1014: PetscCall(PetscFree(xdata[0]));
1015: xdata[0] = tmp;
1016: mem_estimate = new_estimate;
1017: ++no_malloc;
1018: for (ctr = 1; ctr <= i; ctr++) xdata[ctr] = xdata[ctr - 1] + isz1[ctr - 1];
1019: }
1020: xdata[i][ct2++] = row;
1021: ct3++;
1022: }
1023: }
1024: for (k = oct2, max2 = ct2; k < max2; k++) {
1025: row = xdata[i][k] - rstart;
1026: start = ai[row];
1027: end = ai[row + 1];
1028: for (l = start; l < end; l++) {
1029: val = aj[l] + cstart;
1030: if (!PetscBTLookupSet(xtable, val)) {
1031: if (!(ct3 < mem_estimate)) {
1032: new_estimate = (PetscInt)(1.5 * mem_estimate) + 1;
1033: PetscCall(PetscMalloc1(new_estimate, &tmp));
1034: PetscCall(PetscArraycpy(tmp, xdata[0], mem_estimate));
1035: PetscCall(PetscFree(xdata[0]));
1036: xdata[0] = tmp;
1037: mem_estimate = new_estimate;
1038: ++no_malloc;
1039: for (ctr = 1; ctr <= i; ctr++) xdata[ctr] = xdata[ctr - 1] + isz1[ctr - 1];
1040: }
1041: xdata[i][ct2++] = val;
1042: ct3++;
1043: }
1044: }
1045: start = bi[row];
1046: end = bi[row + 1];
1047: for (l = start; l < end; l++) {
1048: val = garray[bj[l]];
1049: if (!PetscBTLookupSet(xtable, val)) {
1050: if (!(ct3 < mem_estimate)) {
1051: new_estimate = (PetscInt)(1.5 * mem_estimate) + 1;
1052: PetscCall(PetscMalloc1(new_estimate, &tmp));
1053: PetscCall(PetscArraycpy(tmp, xdata[0], mem_estimate));
1054: PetscCall(PetscFree(xdata[0]));
1055: xdata[0] = tmp;
1056: mem_estimate = new_estimate;
1057: ++no_malloc;
1058: for (ctr = 1; ctr <= i; ctr++) xdata[ctr] = xdata[ctr - 1] + isz1[ctr - 1];
1059: }
1060: xdata[i][ct2++] = val;
1061: ct3++;
1062: }
1063: }
1064: }
1065: /* Update the header*/
1066: xdata[i][2 * j] = ct2 - oct2; /* Undo the vector isz1 and use only a var*/
1067: xdata[i][2 * j - 1] = rbuf_i[2 * j - 1];
1068: }
1069: xdata[i][0] = rbuf_0;
1070: if (i + 1 < nrqr) xdata[i + 1] = xdata[i] + ct2;
1071: isz1[i] = ct2; /* size of each message */
1072: }
1073: PetscCall(PetscBTDestroy(&xtable));
1074: PetscCall(PetscInfo(C, "Allocated %" PetscInt_FMT " bytes, required %" PetscInt_FMT " bytes, no of mallocs = %" PetscInt_FMT "\n", mem_estimate, ct3, no_malloc));
1075: PetscFunctionReturn(PETSC_SUCCESS);
1076: }
1078: extern PetscErrorCode MatCreateSubMatrices_MPIAIJ_Local(Mat, PetscInt, const IS[], const IS[], MatReuse, Mat *);
1079: /*
1080: Every processor gets the entire matrix
1081: */
1082: PetscErrorCode MatCreateSubMatrix_MPIAIJ_All(Mat A, MatCreateSubMatrixOption flag, MatReuse scall, Mat *Bin[])
1083: {
1084: Mat B;
1085: Mat_MPIAIJ *a = (Mat_MPIAIJ *)A->data;
1086: Mat_SeqAIJ *b, *ad = (Mat_SeqAIJ *)a->A->data, *bd = (Mat_SeqAIJ *)a->B->data;
1087: PetscMPIInt size, rank;
1088: PetscInt sendcount, *rstarts = A->rmap->range, n, cnt, j, nrecv = 0;
1089: PetscInt m, *b_sendj, *garray = a->garray, *lens, *jsendbuf, *a_jsendbuf, *b_jsendbuf;
1091: PetscFunctionBegin;
1092: PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size));
1093: PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)A), &rank));
1094: if (scall == MAT_INITIAL_MATRIX) {
1095: /* Tell every processor the number of nonzeros per row */
1096: PetscCall(PetscCalloc1(A->rmap->N, &lens));
1097: for (PetscInt i = A->rmap->rstart; i < A->rmap->rend; i++) lens[i] = ad->i[i - A->rmap->rstart + 1] - ad->i[i - A->rmap->rstart] + bd->i[i - A->rmap->rstart + 1] - bd->i[i - A->rmap->rstart];
1099: /* All MPI processes get the same matrix */
1100: PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, lens, A->rmap->N, MPIU_INT, MPI_SUM, PetscObjectComm((PetscObject)A)));
1102: /* Create the sequential matrix of the same type as the local block diagonal */
1103: PetscCall(MatCreate(PETSC_COMM_SELF, &B));
1104: PetscCall(MatSetSizes(B, A->rmap->N, A->cmap->N, PETSC_DETERMINE, PETSC_DETERMINE));
1105: PetscCall(MatSetBlockSizesFromMats(B, A, A));
1106: PetscCall(MatSetType(B, ((PetscObject)a->A)->type_name));
1107: PetscCall(MatSeqAIJSetPreallocation(B, 0, lens));
1108: PetscCall(PetscCalloc1(2, Bin));
1109: **Bin = B;
1110: b = (Mat_SeqAIJ *)B->data;
1112: /* zero column space */
1113: nrecv = 0;
1114: for (PetscMPIInt i = 0; i < size; i++) {
1115: for (j = A->rmap->range[i]; j < A->rmap->range[i + 1]; j++) nrecv += lens[j];
1116: }
1117: PetscCall(PetscArrayzero(b->j, nrecv));
1119: /* Copy my part of matrix column indices over */
1120: sendcount = ad->nz + bd->nz;
1121: jsendbuf = PetscSafePointerPlusOffset(b->j, b->i[rstarts[rank]]);
1122: a_jsendbuf = ad->j;
1123: b_jsendbuf = bd->j;
1124: n = A->rmap->rend - A->rmap->rstart;
1125: cnt = 0;
1126: for (PetscInt i = 0; i < n; i++) {
1127: /* put in lower diagonal portion */
1128: m = bd->i[i + 1] - bd->i[i];
1129: while (m > 0) {
1130: /* is it above diagonal (in bd (compressed) numbering) */
1131: if (garray[*b_jsendbuf] > A->rmap->rstart + i) break;
1132: jsendbuf[cnt++] = garray[*b_jsendbuf++];
1133: m--;
1134: }
1136: /* put in diagonal portion */
1137: for (PetscInt j = ad->i[i]; j < ad->i[i + 1]; j++) jsendbuf[cnt++] = A->rmap->rstart + *a_jsendbuf++;
1139: /* put in upper diagonal portion */
1140: while (m-- > 0) jsendbuf[cnt++] = garray[*b_jsendbuf++];
1141: }
1142: PetscCheck(cnt == sendcount, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Corrupted PETSc matrix: nz given %" PetscInt_FMT " actual nz %" PetscInt_FMT, sendcount, cnt);
1144: /* send column indices, b->j was zeroed */
1145: PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, b->j, nrecv, MPIU_INT, MPI_SUM, PetscObjectComm((PetscObject)A)));
1147: /* Assemble the matrix into useable form (numerical values not yet set) */
1148: /* set the b->ilen (length of each row) values */
1149: PetscCall(PetscArraycpy(b->ilen, lens, A->rmap->N));
1150: /* set the b->i indices */
1151: b->i[0] = 0;
1152: for (PetscInt i = 1; i <= A->rmap->N; i++) b->i[i] = b->i[i - 1] + lens[i - 1];
1153: PetscCall(PetscFree(lens));
1154: PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
1155: PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
1157: } else {
1158: B = **Bin;
1159: b = (Mat_SeqAIJ *)B->data;
1160: }
1162: /* Copy my part of matrix numerical values into the values location */
1163: if (flag == MAT_GET_VALUES) {
1164: const PetscScalar *ada, *bda, *a_sendbuf, *b_sendbuf;
1165: MatScalar *sendbuf;
1167: /* initialize b->a */
1168: PetscCall(PetscArrayzero(b->a, b->nz));
1170: PetscCall(MatSeqAIJGetArrayRead(a->A, &ada));
1171: PetscCall(MatSeqAIJGetArrayRead(a->B, &bda));
1172: sendcount = ad->nz + bd->nz;
1173: sendbuf = PetscSafePointerPlusOffset(b->a, b->i[rstarts[rank]]);
1174: a_sendbuf = ada;
1175: b_sendbuf = bda;
1176: b_sendj = bd->j;
1177: n = A->rmap->rend - A->rmap->rstart;
1178: cnt = 0;
1179: for (PetscInt i = 0; i < n; i++) {
1180: /* put in lower diagonal portion */
1181: m = bd->i[i + 1] - bd->i[i];
1182: while (m > 0) {
1183: /* is it above diagonal (in bd (compressed) numbering) */
1184: if (garray[*b_sendj] > A->rmap->rstart + i) break;
1185: sendbuf[cnt++] = *b_sendbuf++;
1186: m--;
1187: b_sendj++;
1188: }
1190: /* put in diagonal portion */
1191: for (PetscInt j = ad->i[i]; j < ad->i[i + 1]; j++) sendbuf[cnt++] = *a_sendbuf++;
1193: /* put in upper diagonal portion */
1194: while (m-- > 0) {
1195: sendbuf[cnt++] = *b_sendbuf++;
1196: b_sendj++;
1197: }
1198: }
1199: PetscCheck(cnt == sendcount, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Corrupted PETSc matrix: nz given %" PetscInt_FMT " actual nz %" PetscInt_FMT, sendcount, cnt);
1201: /* send values, b->a was zeroed */
1202: PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, b->a, b->nz, MPIU_SCALAR, MPIU_SUM, PetscObjectComm((PetscObject)A)));
1204: PetscCall(MatSeqAIJRestoreArrayRead(a->A, &ada));
1205: PetscCall(MatSeqAIJRestoreArrayRead(a->B, &bda));
1206: } /* endof (flag == MAT_GET_VALUES) */
1208: PetscCall(MatPropagateSymmetryOptions(A, B));
1209: PetscFunctionReturn(PETSC_SUCCESS);
1210: }
1212: PetscErrorCode MatCreateSubMatrices_MPIAIJ_SingleIS_Local(Mat C, PetscInt ismax, const IS isrow[], const IS iscol[], MatReuse scall, PetscBool allcolumns, Mat *submats)
1213: {
1214: Mat_MPIAIJ *c = (Mat_MPIAIJ *)C->data;
1215: Mat submat, A = c->A, B = c->B;
1216: Mat_SeqAIJ *a = (Mat_SeqAIJ *)A->data, *b = (Mat_SeqAIJ *)B->data, *subc;
1217: Mat_SubSppt *smatis1;
1218: const PetscInt *icol, *irow;
1219: PetscInt cstart = C->cmap->rstart, cend = C->cmap->rend, rstart = C->rmap->rstart;
1220: PetscInt nzA, nzB, nrow, ncol, start, k, ct1, ct2, ct3, row, msz, tcol, max1, nnz, rmax, ncols, Crow, ctr_j, kmax, jcnt, lwrite, ib, jb;
1221: PetscInt *ai = a->i, *aj = a->j, *bi = b->i, *bj = b->j, *bmap = c->garray;
1222: PetscInt *req_size, *ctr, *tmp, *iptr, *lens, *cols, *sbuf1_j, *sbuf_aj_i, *rbuf1_i, *sbuf1_i, *rbuf2_i, *rbuf3_i, *cworkB, *subcols;
1223: PetscInt **sbuf1, **sbuf2, **rbuf1, **ptr, **rbuf3, **sbuf_aj, **rbuf2;
1224: #if PetscDefined(USE_CTABLE)
1225: PetscInt *cmap_loc, *rmap_loc;
1226: PetscHMapI cmap, rmap;
1227: #else
1228: PetscInt *cmap, *rmap;
1229: #endif
1230: PetscMPIInt rank, size, tag1, tag2, tag3, tag4, nrqr, nrqs = 0, proc, idex, end;
1231: PetscMPIInt *req_source1, *req_source2, *w1, *w2, *pa, *onodes1, *olengths1, *row2proc;
1232: PetscScalar *vworkA, *vworkB, *a_a, *b_a, *subvals = NULL, *suba = NULL, *vals, *sbuf_aa_i, *rbuf4_i;
1233: PetscScalar **rbuf4, **sbuf_aa;
1234: PetscBool isrowsorted, iscolsorted, direct_csr_reuse = PETSC_FALSE;
1235: PetscObjectState submat_nonzerostate;
1236: MPI_Request *s_waits1, *r_waits1, *s_waits2, *r_waits2, *r_waits3, *r_waits4, *s_waits3 = NULL, *s_waits4;
1237: MPI_Status *r_status1, *r_status2, *s_status1, *s_status3 = NULL, *s_status2, *r_status3 = NULL, *r_status4, *s_status4;
1238: MPI_Comm comm;
1240: PetscFunctionBegin;
1243: PetscCheck(ismax == 1, PETSC_COMM_SELF, PETSC_ERR_SUP, "This routine only works when all processes have ismax=1");
1244: PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
1245: size = c->size;
1246: rank = c->rank;
1248: PetscCall(ISSorted(iscol[0], &iscolsorted));
1249: PetscCall(ISSorted(isrow[0], &isrowsorted));
1250: PetscCall(ISGetIndices(isrow[0], &irow));
1251: PetscCall(ISGetLocalSize(isrow[0], &nrow));
1252: if (allcolumns) {
1253: icol = NULL;
1254: ncol = C->cmap->N;
1255: } else {
1256: PetscCall(ISGetIndices(iscol[0], &icol));
1257: PetscCall(ISGetLocalSize(iscol[0], &ncol));
1258: }
1260: /* MAT_REUSE_MATRIX requires an unchanged pattern, so there are no values to update for a structure-only matrix */
1261: if (C->structure_only && scall == MAT_REUSE_MATRIX) {
1262: PetscCheck(submats[0]->rmap->n == nrow && submats[0]->cmap->n == ncol, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Cannot reuse matrix: wrong dimensions");
1263: PetscCall(ISRestoreIndices(isrow[0], &irow));
1264: if (!allcolumns) PetscCall(ISRestoreIndices(iscol[0], &icol));
1265: PetscFunctionReturn(PETSC_SUCCESS);
1266: }
1268: PetscCall(MatSeqAIJGetArrayRead(A, (const PetscScalar **)&a_a));
1269: PetscCall(MatSeqAIJGetArrayRead(B, (const PetscScalar **)&b_a));
1270: if (scall == MAT_INITIAL_MATRIX) {
1271: PetscInt *sbuf2_i, *cworkA, lwrite, ctmp;
1273: /* Get some new tags to keep the communication clean */
1274: tag1 = ((PetscObject)C)->tag;
1275: PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag2));
1276: PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag3));
1278: /* evaluate communication - mesg to who, length of mesg, and buffer space
1279: required. Based on this, buffers are allocated, and data copied into them */
1280: PetscCall(PetscCalloc2(size, &w1, size, &w2));
1281: PetscCall(PetscMalloc1(nrow, &row2proc));
1283: /* w1[proc] = num of rows owned by proc -- to be requested */
1284: proc = 0;
1285: nrqs = 0; /* num of outgoing messages */
1286: for (PetscInt j = 0; j < nrow; j++) {
1287: row = irow[j];
1288: if (!isrowsorted) proc = 0;
1289: while (row >= C->rmap->range[proc + 1]) proc++;
1290: w1[proc]++;
1291: row2proc[j] = proc; /* map row index to proc */
1293: if (proc != rank && !w2[proc]) {
1294: w2[proc] = 1;
1295: nrqs++;
1296: }
1297: }
1298: w1[rank] = 0; /* rows owned by self will not be requested */
1300: PetscCall(PetscMalloc1(nrqs, &pa)); /*(proc -array)*/
1301: for (PetscMPIInt proc = 0, j = 0; proc < size; proc++) {
1302: if (w1[proc]) pa[j++] = proc;
1303: }
1305: /* Each message would have a header = 1 + 2*(num of IS) + data (here,num of IS = 1) */
1306: msz = 0; /* total mesg length (for all procs) */
1307: for (PetscMPIInt i = 0; i < nrqs; i++) {
1308: w1[pa[i]] += 3;
1309: msz += w1[pa[i]];
1310: }
1311: PetscCall(PetscInfo(0, "Number of outgoing messages %d Total message length %" PetscInt_FMT "\n", nrqs, msz));
1313: /* Determine nrqr, the number of messages to expect, their lengths, from from-ids */
1314: /* if w2[proc]=1, a message of length w1[proc] will be sent to proc; */
1315: PetscCall(PetscGatherNumberOfMessages(comm, w2, w1, &nrqr));
1317: /* Input: nrqs: nsend; nrqr: nrecv; w1: msg length to be sent;
1318: Output: onodes1: recv node-ids; olengths1: corresponding recv message length */
1319: PetscCall(PetscGatherMessageLengths(comm, nrqs, nrqr, w1, &onodes1, &olengths1));
1321: /* Now post the Irecvs corresponding to these messages */
1322: PetscCall(PetscPostIrecvInt(comm, tag1, nrqr, onodes1, olengths1, &rbuf1, &r_waits1));
1324: PetscCall(PetscFree(onodes1));
1325: PetscCall(PetscFree(olengths1));
1327: /* Allocate Memory for outgoing messages */
1328: PetscCall(PetscMalloc4(size, &sbuf1, size, &ptr, 2 * msz, &tmp, size, &ctr));
1329: PetscCall(PetscArrayzero(sbuf1, size));
1330: PetscCall(PetscArrayzero(ptr, size));
1332: /* subf1[pa[0]] = tmp, subf1[pa[i]] = subf1[pa[i-1]] + w1[pa[i-1]] */
1333: iptr = tmp;
1334: for (PetscMPIInt i = 0; i < nrqs; i++) {
1335: sbuf1[pa[i]] = iptr;
1336: iptr += w1[pa[i]];
1337: }
1339: /* Form the outgoing messages */
1340: /* Initialize the header space */
1341: for (PetscMPIInt i = 0; i < nrqs; i++) {
1342: PetscCall(PetscArrayzero(sbuf1[pa[i]], 3));
1343: ptr[pa[i]] = sbuf1[pa[i]] + 3;
1344: }
1346: /* Parse the isrow and copy data into outbuf */
1347: PetscCall(PetscArrayzero(ctr, size));
1348: for (PetscInt j = 0; j < nrow; j++) { /* parse the indices of each IS */
1349: if (row2proc[j] != rank) { /* copy to the outgoing buf */
1350: *ptr[row2proc[j]] = irow[j];
1351: ctr[row2proc[j]]++;
1352: ptr[row2proc[j]]++;
1353: }
1354: }
1356: /* Update the headers for the current IS */
1357: for (PetscMPIInt j = 0; j < size; j++) { /* Can Optimise this loop too */
1358: if ((ctr_j = ctr[j])) {
1359: sbuf1_j = sbuf1[j];
1360: k = ++sbuf1_j[0];
1361: sbuf1_j[2 * k] = ctr_j;
1362: sbuf1_j[2 * k - 1] = 0;
1363: }
1364: }
1366: /* Now post the sends */
1367: PetscCall(PetscMalloc1(nrqs, &s_waits1));
1368: for (PetscMPIInt i = 0; i < nrqs; ++i) PetscCallMPI(MPIU_Isend(sbuf1[pa[i]], w1[pa[i]], MPIU_INT, pa[i], tag1, comm, s_waits1 + i));
1370: /* Post Receives to capture the buffer size */
1371: PetscCall(PetscMalloc4(nrqs, &r_status2, nrqr, &s_waits2, nrqs, &r_waits2, nrqr, &s_status2));
1372: PetscCall(PetscMalloc3(nrqs, &req_source2, nrqs, &rbuf2, nrqs, &rbuf3));
1374: if (nrqs) rbuf2[0] = tmp + msz;
1375: for (PetscMPIInt i = 1; i < nrqs; ++i) rbuf2[i] = rbuf2[i - 1] + w1[pa[i - 1]];
1377: for (PetscMPIInt i = 0; i < nrqs; ++i) PetscCallMPI(MPIU_Irecv(rbuf2[i], w1[pa[i]], MPIU_INT, pa[i], tag2, comm, r_waits2 + i));
1379: PetscCall(PetscFree2(w1, w2));
1381: /* Send to other procs the buf size they should allocate */
1382: /* Receive messages*/
1383: PetscCall(PetscMalloc1(nrqr, &r_status1));
1384: PetscCall(PetscMalloc3(nrqr, &sbuf2, nrqr, &req_size, nrqr, &req_source1));
1386: PetscCallMPI(MPI_Waitall(nrqr, r_waits1, r_status1));
1387: for (PetscMPIInt i = 0; i < nrqr; ++i) {
1388: req_size[i] = 0;
1389: rbuf1_i = rbuf1[i];
1390: start = 2 * rbuf1_i[0] + 1;
1391: PetscCallMPI(MPI_Get_count(r_status1 + i, MPIU_INT, &end));
1392: PetscCall(PetscMalloc1(end, &sbuf2[i]));
1393: sbuf2_i = sbuf2[i];
1394: for (PetscInt j = start; j < end; j++) {
1395: k = rbuf1_i[j] - rstart;
1396: ncols = ai[k + 1] - ai[k] + bi[k + 1] - bi[k];
1397: sbuf2_i[j] = ncols;
1398: req_size[i] += ncols;
1399: }
1400: req_source1[i] = r_status1[i].MPI_SOURCE;
1402: /* form the header */
1403: sbuf2_i[0] = req_size[i];
1404: for (PetscInt j = 1; j < start; j++) sbuf2_i[j] = rbuf1_i[j];
1406: PetscCallMPI(MPIU_Isend(sbuf2_i, end, MPIU_INT, req_source1[i], tag2, comm, s_waits2 + i));
1407: }
1409: PetscCall(PetscFree(r_status1));
1410: PetscCall(PetscFree(r_waits1));
1412: /* rbuf2 is received, Post recv column indices a->j */
1413: PetscCallMPI(MPI_Waitall(nrqs, r_waits2, r_status2));
1415: PetscCall(PetscMalloc4(nrqs, &r_waits3, nrqr, &s_waits3, nrqs, &r_status3, nrqr, &s_status3));
1416: for (PetscMPIInt i = 0; i < nrqs; ++i) {
1417: PetscCall(PetscMalloc1(rbuf2[i][0], &rbuf3[i]));
1418: req_source2[i] = r_status2[i].MPI_SOURCE;
1419: PetscCallMPI(MPIU_Irecv(rbuf3[i], rbuf2[i][0], MPIU_INT, req_source2[i], tag3, comm, r_waits3 + i));
1420: }
1422: /* Wait on sends1 and sends2 */
1423: PetscCall(PetscMalloc1(nrqs, &s_status1));
1424: PetscCallMPI(MPI_Waitall(nrqs, s_waits1, s_status1));
1425: PetscCall(PetscFree(s_waits1));
1426: PetscCall(PetscFree(s_status1));
1428: PetscCallMPI(MPI_Waitall(nrqr, s_waits2, s_status2));
1429: PetscCall(PetscFree4(r_status2, s_waits2, r_waits2, s_status2));
1431: /* Now allocate sending buffers for a->j, and send them off */
1432: PetscCall(PetscMalloc1(nrqr, &sbuf_aj));
1433: jcnt = 0;
1434: for (PetscMPIInt i = 0; i < nrqr; i++) jcnt += req_size[i];
1435: if (nrqr) PetscCall(PetscMalloc1(jcnt, &sbuf_aj[0]));
1436: for (PetscMPIInt i = 1; i < nrqr; i++) sbuf_aj[i] = sbuf_aj[i - 1] + req_size[i - 1];
1438: for (PetscMPIInt i = 0; i < nrqr; i++) { /* for each requested message */
1439: rbuf1_i = rbuf1[i];
1440: sbuf_aj_i = sbuf_aj[i];
1441: ct1 = 2 * rbuf1_i[0] + 1;
1442: ct2 = 0;
1443: /* max1=rbuf1_i[0]; PetscCheck(max1 == 1,PETSC_COMM_SELF,PETSC_ERR_PLIB,"max1 %d != 1",max1); */
1445: kmax = rbuf1[i][2];
1446: for (PetscInt k = 0; k < kmax; k++, ct1++) { /* for each row */
1447: row = rbuf1_i[ct1] - rstart;
1448: nzA = ai[row + 1] - ai[row];
1449: nzB = bi[row + 1] - bi[row];
1450: ncols = nzA + nzB;
1451: cworkA = PetscSafePointerPlusOffset(aj, ai[row]);
1452: cworkB = PetscSafePointerPlusOffset(bj, bi[row]);
1454: /* load the column indices for this row into cols*/
1455: cols = PetscSafePointerPlusOffset(sbuf_aj_i, ct2);
1457: lwrite = 0;
1458: for (PetscInt l = 0; l < nzB; l++) {
1459: if ((ctmp = bmap[cworkB[l]]) < cstart) cols[lwrite++] = ctmp;
1460: }
1461: for (PetscInt l = 0; l < nzA; l++) cols[lwrite++] = cstart + cworkA[l];
1462: for (PetscInt l = 0; l < nzB; l++) {
1463: if ((ctmp = bmap[cworkB[l]]) >= cend) cols[lwrite++] = ctmp;
1464: }
1466: ct2 += ncols;
1467: }
1468: PetscCallMPI(MPIU_Isend(sbuf_aj_i, req_size[i], MPIU_INT, req_source1[i], tag3, comm, s_waits3 + i));
1469: }
1471: /* create column map (cmap): global col of C -> local col of submat */
1472: #if PetscDefined(USE_CTABLE)
1473: if (!allcolumns) {
1474: PetscCall(PetscHMapICreateWithSize(ncol, &cmap));
1475: PetscCall(PetscCalloc1(C->cmap->n, &cmap_loc));
1476: for (PetscInt j = 0; j < ncol; j++) { /* use array cmap_loc[] for local col indices */
1477: if (icol[j] >= cstart && icol[j] < cend) {
1478: cmap_loc[icol[j] - cstart] = j + 1;
1479: } else { /* use PetscHMapI for non-local col indices */
1480: PetscCall(PetscHMapISet(cmap, icol[j] + 1, j + 1));
1481: }
1482: }
1483: } else {
1484: cmap = NULL;
1485: cmap_loc = NULL;
1486: }
1487: PetscCall(PetscCalloc1(C->rmap->n, &rmap_loc));
1488: #else
1489: if (!allcolumns) {
1490: PetscCall(PetscCalloc1(C->cmap->N, &cmap));
1491: for (PetscInt j = 0; j < ncol; j++) cmap[icol[j]] = j + 1;
1492: } else {
1493: cmap = NULL;
1494: }
1495: #endif
1497: /* Create lens for MatSeqAIJSetPreallocation() */
1498: PetscCall(PetscCalloc1(nrow, &lens));
1500: /* Compute lens from local part of C */
1501: for (PetscInt j = 0; j < nrow; j++) {
1502: row = irow[j];
1503: if (row2proc[j] == rank) {
1504: /* diagonal part A = c->A */
1505: ncols = ai[row - rstart + 1] - ai[row - rstart];
1506: cols = PetscSafePointerPlusOffset(aj, ai[row - rstart]);
1507: if (!allcolumns) {
1508: for (PetscInt k = 0; k < ncols; k++) {
1509: #if PetscDefined(USE_CTABLE)
1510: tcol = cmap_loc[cols[k]];
1511: #else
1512: tcol = cmap[cols[k] + cstart];
1513: #endif
1514: if (tcol) lens[j]++;
1515: }
1516: } else { /* allcolumns */
1517: lens[j] = ncols;
1518: }
1520: /* off-diagonal part B = c->B */
1521: ncols = bi[row - rstart + 1] - bi[row - rstart];
1522: cols = PetscSafePointerPlusOffset(bj, bi[row - rstart]);
1523: if (!allcolumns) {
1524: for (PetscInt k = 0; k < ncols; k++) {
1525: #if PetscDefined(USE_CTABLE)
1526: PetscCall(PetscHMapIGetWithDefault(cmap, bmap[cols[k]] + 1, 0, &tcol));
1527: #else
1528: tcol = cmap[bmap[cols[k]]];
1529: #endif
1530: if (tcol) lens[j]++;
1531: }
1532: } else { /* allcolumns */
1533: lens[j] += ncols;
1534: }
1535: }
1536: }
1538: /* Create row map (rmap): global row of C -> local row of submat */
1539: #if PetscDefined(USE_CTABLE)
1540: PetscCall(PetscHMapICreateWithSize(nrow, &rmap));
1541: for (PetscInt j = 0; j < nrow; j++) {
1542: row = irow[j];
1543: if (row2proc[j] == rank) { /* a local row */
1544: rmap_loc[row - rstart] = j;
1545: } else {
1546: PetscCall(PetscHMapISet(rmap, irow[j] + 1, j + 1));
1547: }
1548: }
1549: #else
1550: PetscCall(PetscCalloc1(C->rmap->N, &rmap));
1551: for (PetscInt j = 0; j < nrow; j++) rmap[irow[j]] = j;
1552: #endif
1554: /* Update lens from offproc data */
1555: /* recv a->j is done */
1556: PetscCallMPI(MPI_Waitall(nrqs, r_waits3, r_status3));
1557: for (PetscMPIInt i = 0; i < nrqs; i++) {
1558: sbuf1_i = sbuf1[pa[i]];
1559: /* jmax = sbuf1_i[0]; PetscCheck(jmax == 1,PETSC_COMM_SELF,PETSC_ERR_PLIB,"jmax !=1"); */
1560: ct1 = 2 + 1;
1561: ct2 = 0;
1562: rbuf2_i = rbuf2[i]; /* received length of C->j */
1563: rbuf3_i = rbuf3[i]; /* received C->j */
1565: /* is_no = sbuf1_i[2*j-1]; PetscCheck(is_no == 0,PETSC_COMM_SELF,PETSC_ERR_PLIB,"is_no !=0"); */
1566: max1 = sbuf1_i[2];
1567: for (PetscInt k = 0; k < max1; k++, ct1++) {
1568: #if PetscDefined(USE_CTABLE)
1569: PetscCall(PetscHMapIGetWithDefault(rmap, sbuf1_i[ct1] + 1, 0, &row));
1570: row--;
1571: PetscCheck(row >= 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "row not found in table");
1572: #else
1573: row = rmap[sbuf1_i[ct1]]; /* the row index in submat */
1574: #endif
1575: /* Now, store row index of submat in sbuf1_i[ct1] */
1576: sbuf1_i[ct1] = row;
1578: nnz = rbuf2_i[ct1];
1579: if (!allcolumns) {
1580: for (PetscMPIInt l = 0; l < nnz; l++, ct2++) {
1581: #if PetscDefined(USE_CTABLE)
1582: if (rbuf3_i[ct2] >= cstart && rbuf3_i[ct2] < cend) {
1583: tcol = cmap_loc[rbuf3_i[ct2] - cstart];
1584: } else {
1585: PetscCall(PetscHMapIGetWithDefault(cmap, rbuf3_i[ct2] + 1, 0, &tcol));
1586: }
1587: #else
1588: tcol = cmap[rbuf3_i[ct2]]; /* column index in submat */
1589: #endif
1590: if (tcol) lens[row]++;
1591: }
1592: } else { /* allcolumns */
1593: lens[row] += nnz;
1594: }
1595: }
1596: }
1597: PetscCallMPI(MPI_Waitall(nrqr, s_waits3, s_status3));
1598: PetscCall(PetscFree4(r_waits3, s_waits3, r_status3, s_status3));
1600: /* Create the submatrices */
1601: PetscCall(MatCreate(PETSC_COMM_SELF, &submat));
1602: PetscCall(MatSetSizes(submat, nrow, ncol, PETSC_DETERMINE, PETSC_DETERMINE));
1604: PetscCall(ISGetBlockSize(isrow[0], &ib));
1605: PetscCall(ISGetBlockSize(iscol[0], &jb));
1606: if (ib > 1 || jb > 1) PetscCall(MatSetBlockSizes(submat, ib, jb));
1607: PetscCall(MatSetType(submat, ((PetscObject)A)->type_name));
1608: PetscCall(MatSetOption(submat, MAT_STRUCTURE_ONLY, C->structure_only));
1609: PetscCall(MatSeqAIJSetPreallocation(submat, 0, lens));
1611: /* create struct Mat_SubSppt and attached it to submat */
1612: PetscCall(PetscNew(&smatis1));
1613: subc = (Mat_SeqAIJ *)submat->data;
1614: subc->submatis1 = smatis1;
1616: smatis1->id = 0;
1617: smatis1->nrqs = nrqs;
1618: smatis1->nrqr = nrqr;
1619: smatis1->rbuf1 = rbuf1;
1620: smatis1->rbuf2 = rbuf2;
1621: smatis1->rbuf3 = rbuf3;
1622: smatis1->sbuf2 = sbuf2;
1623: smatis1->req_source2 = req_source2;
1625: smatis1->sbuf1 = sbuf1;
1626: smatis1->ptr = ptr;
1627: smatis1->tmp = tmp;
1628: smatis1->ctr = ctr;
1630: smatis1->pa = pa;
1631: smatis1->req_size = req_size;
1632: smatis1->req_source1 = req_source1;
1634: smatis1->allcolumns = allcolumns;
1635: smatis1->singleis = PETSC_TRUE;
1636: smatis1->row2proc = row2proc;
1637: smatis1->rmap = rmap;
1638: smatis1->cmap = cmap;
1639: #if PetscDefined(USE_CTABLE)
1640: smatis1->rmap_loc = rmap_loc;
1641: smatis1->cmap_loc = cmap_loc;
1642: #endif
1644: smatis1->destroy = submat->ops->destroy;
1645: submat->ops->destroy = MatDestroySubMatrix_SeqAIJ;
1646: submat->factortype = C->factortype;
1648: /* compute rmax */
1649: rmax = 0;
1650: for (PetscMPIInt i = 0; i < nrow; i++) rmax = PetscMax(rmax, lens[i]);
1652: } else { /* scall == MAT_REUSE_MATRIX */
1653: submat = submats[0];
1654: PetscCheck(submat->rmap->n == nrow && submat->cmap->n == ncol, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Cannot reuse matrix. wrong size");
1656: subc = (Mat_SeqAIJ *)submat->data;
1657: rmax = subc->rmax;
1658: smatis1 = subc->submatis1;
1659: nrqs = smatis1->nrqs;
1660: nrqr = smatis1->nrqr;
1661: rbuf1 = smatis1->rbuf1;
1662: rbuf2 = smatis1->rbuf2;
1663: rbuf3 = smatis1->rbuf3;
1664: req_source2 = smatis1->req_source2;
1666: sbuf1 = smatis1->sbuf1;
1667: sbuf2 = smatis1->sbuf2;
1668: ptr = smatis1->ptr;
1669: tmp = smatis1->tmp;
1670: ctr = smatis1->ctr;
1672: pa = smatis1->pa;
1673: req_size = smatis1->req_size;
1674: req_source1 = smatis1->req_source1;
1676: allcolumns = smatis1->allcolumns;
1677: row2proc = smatis1->row2proc;
1678: rmap = smatis1->rmap;
1679: cmap = smatis1->cmap;
1680: #if PetscDefined(USE_CTABLE)
1681: rmap_loc = smatis1->rmap_loc;
1682: cmap_loc = smatis1->cmap_loc;
1683: #endif
1684: PetscCall(MatGetNonzeroState(submat, &submat_nonzerostate));
1685: // User changes to the submatrix graph invalidate the cached CSR positions.
1686: direct_csr_reuse = (PetscBool)(iscolsorted && !allcolumns && submat_nonzerostate == smatis1->nonzerostate);
1687: // Build owned-row maps only when reuse first needs them and the initial graph still matches.
1688: if (direct_csr_reuse && !smatis1->csrcached) {
1689: PetscInt nlocal_a = 0, nlocal_b = 0;
1691: for (PetscInt j = 0; j < nrow; j++) {
1692: row = irow[j];
1693: if (row2proc[j] != rank) continue;
1694: Crow = row - rstart;
1695: for (PetscInt k = ai[Crow]; k < ai[Crow + 1]; k++) {
1696: #if PetscDefined(USE_CTABLE)
1697: tcol = cmap_loc[aj[k]];
1698: #else
1699: tcol = cmap[aj[k] + cstart];
1700: #endif
1701: if (tcol) nlocal_a++;
1702: }
1703: for (PetscInt k = bi[Crow]; k < bi[Crow + 1]; k++) {
1704: #if PetscDefined(USE_CTABLE)
1705: PetscCall(PetscHMapIGetWithDefault(cmap, bmap[bj[k]] + 1, 0, &tcol));
1706: #else
1707: tcol = cmap[bmap[bj[k]]];
1708: #endif
1709: if (tcol) nlocal_b++;
1710: }
1711: }
1712: PetscCall(PetscMalloc2(nlocal_a, &smatis1->local_a_parent, nlocal_a, &smatis1->local_a_sub));
1713: PetscCall(PetscMalloc2(nlocal_b, &smatis1->local_b_parent, nlocal_b, &smatis1->local_b_sub));
1714: smatis1->nlocal_a = 0;
1715: smatis1->nlocal_b = 0;
1716: for (PetscInt j = 0; j < nrow; j++) {
1717: PetscInt subrow_start, subrow_nnz;
1719: row = irow[j];
1720: if (row2proc[j] != rank) continue;
1721: Crow = row - rstart;
1722: #if PetscDefined(USE_CTABLE)
1723: row = rmap_loc[Crow];
1724: #else
1725: row = rmap[row];
1726: #endif
1727: subrow_start = subc->i[row];
1728: subrow_nnz = subc->i[row + 1] - subrow_start;
1729: for (PetscInt k = ai[Crow]; k < ai[Crow + 1]; k++) {
1730: PetscInt loc;
1732: #if PetscDefined(USE_CTABLE)
1733: tcol = cmap_loc[aj[k]];
1734: #else
1735: tcol = cmap[aj[k] + cstart];
1736: #endif
1737: if (!tcol) continue;
1738: PetscCall(PetscFindInt(tcol - 1, subrow_nnz, subc->j + subrow_start, &loc));
1739: PetscCheck(loc >= 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Owned diagonal entry is missing from reused submatrix graph");
1740: smatis1->local_a_parent[smatis1->nlocal_a] = k;
1741: smatis1->local_a_sub[smatis1->nlocal_a++] = subrow_start + loc;
1742: }
1743: for (PetscInt k = bi[Crow]; k < bi[Crow + 1]; k++) {
1744: PetscInt loc;
1746: #if PetscDefined(USE_CTABLE)
1747: PetscCall(PetscHMapIGetWithDefault(cmap, bmap[bj[k]] + 1, 0, &tcol));
1748: #else
1749: tcol = cmap[bmap[bj[k]]];
1750: #endif
1751: if (!tcol) continue;
1752: PetscCall(PetscFindInt(tcol - 1, subrow_nnz, subc->j + subrow_start, &loc));
1753: PetscCheck(loc >= 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Owned off-diagonal entry is missing from reused submatrix graph");
1754: smatis1->local_b_parent[smatis1->nlocal_b] = k;
1755: smatis1->local_b_sub[smatis1->nlocal_b++] = subrow_start + loc;
1756: }
1757: }
1758: PetscCheck(smatis1->nlocal_a == nlocal_a && smatis1->nlocal_b == nlocal_b, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Cached submatrix CSR map size changed during construction");
1759: smatis1->csrcached = PETSC_TRUE;
1760: }
1761: }
1763: /* Post recv matrix values */
1764: PetscCall(PetscMalloc1(rmax, &subcols));
1765: if (!C->structure_only) {
1766: PetscCall(PetscMalloc2(nrqs, &rbuf4, rmax, &subvals));
1767: PetscCall(PetscMalloc4(nrqs, &r_waits4, nrqr, &s_waits4, nrqs, &r_status4, nrqr, &s_status4));
1768: PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag4));
1769: for (PetscMPIInt i = 0; i < nrqs; ++i) {
1770: PetscCall(PetscMalloc1(rbuf2[i][0], &rbuf4[i]));
1771: PetscCallMPI(MPIU_Irecv(rbuf4[i], rbuf2[i][0], MPIU_SCALAR, req_source2[i], tag4, comm, r_waits4 + i));
1772: }
1774: /* Allocate sending buffers for a->a, and send them off */
1775: PetscCall(PetscMalloc1(nrqr, &sbuf_aa));
1776: jcnt = 0;
1777: for (PetscMPIInt i = 0; i < nrqr; i++) jcnt += req_size[i];
1778: if (nrqr) PetscCall(PetscMalloc1(jcnt, &sbuf_aa[0]));
1779: for (PetscMPIInt i = 1; i < nrqr; i++) sbuf_aa[i] = sbuf_aa[i - 1] + req_size[i - 1];
1781: for (PetscMPIInt i = 0; i < nrqr; i++) {
1782: rbuf1_i = rbuf1[i];
1783: sbuf_aa_i = sbuf_aa[i];
1784: ct1 = 2 * rbuf1_i[0] + 1;
1785: ct2 = 0;
1786: /* max1=rbuf1_i[0]; PetscCheck(max1 == 1,PETSC_COMM_SELF,PETSC_ERR_PLIB,"max1 !=1"); */
1788: kmax = rbuf1_i[2];
1789: for (PetscInt k = 0; k < kmax; k++, ct1++) {
1790: row = rbuf1_i[ct1] - rstart;
1791: nzA = ai[row + 1] - ai[row];
1792: nzB = bi[row + 1] - bi[row];
1793: ncols = nzA + nzB;
1794: cworkB = PetscSafePointerPlusOffset(bj, bi[row]);
1795: vworkA = PetscSafePointerPlusOffset(a_a, ai[row]);
1796: vworkB = PetscSafePointerPlusOffset(b_a, bi[row]);
1798: /* load the column values for this row into vals*/
1799: vals = PetscSafePointerPlusOffset(sbuf_aa_i, ct2);
1801: lwrite = 0;
1802: for (PetscInt l = 0; l < nzB; l++) {
1803: if (bmap[cworkB[l]] < cstart) vals[lwrite++] = vworkB[l];
1804: }
1805: for (PetscInt l = 0; l < nzA; l++) vals[lwrite++] = vworkA[l];
1806: for (PetscInt l = 0; l < nzB; l++) {
1807: if (bmap[cworkB[l]] >= cend) vals[lwrite++] = vworkB[l];
1808: }
1810: ct2 += ncols;
1811: }
1812: PetscCallMPI(MPIU_Isend(sbuf_aa_i, req_size[i], MPIU_SCALAR, req_source1[i], tag4, comm, s_waits4 + i));
1813: }
1814: }
1816: /* Assemble submat */
1817: /* First assemble the local rows */
1818: if (direct_csr_reuse) {
1819: PetscCall(MatSeqAIJGetArray(submat, &suba));
1820: for (PetscInt j = 0; j < smatis1->nlocal_a; j++) suba[smatis1->local_a_sub[j]] = a_a[smatis1->local_a_parent[j]];
1821: for (PetscInt j = 0; j < smatis1->nlocal_b; j++) suba[smatis1->local_b_sub[j]] = b_a[smatis1->local_b_parent[j]];
1822: } else {
1823: for (PetscInt j = 0; j < nrow; j++) {
1824: row = irow[j];
1825: if (row2proc[j] != rank) continue;
1826: Crow = row - rstart; /* local row index of C */
1827: #if PetscDefined(USE_CTABLE)
1828: row = rmap_loc[Crow]; /* row index of submat */
1829: #else
1830: row = rmap[row];
1831: #endif
1833: if (allcolumns) {
1834: PetscInt ncol = 0;
1836: /* diagonal part A = c->A */
1837: ncols = ai[Crow + 1] - ai[Crow];
1838: cols = PetscSafePointerPlusOffset(aj, ai[Crow]);
1839: vals = PetscSafePointerPlusOffset(a_a, ai[Crow]);
1840: for (PetscInt k = 0; k < ncols; k++, ncol++) {
1841: subcols[ncol] = cols[k] + cstart;
1842: if (!C->structure_only) subvals[ncol] = vals[k];
1843: }
1845: /* off-diagonal part B = c->B */
1846: ncols = bi[Crow + 1] - bi[Crow];
1847: cols = PetscSafePointerPlusOffset(bj, bi[Crow]);
1848: vals = PetscSafePointerPlusOffset(b_a, bi[Crow]);
1849: for (PetscInt k = 0; k < ncols; k++, ncol++) {
1850: subcols[ncol] = bmap[cols[k]];
1851: if (!C->structure_only) subvals[ncol] = vals[k];
1852: }
1854: PetscCall(MatSetValues_SeqAIJ(submat, 1, &row, ncol, subcols, subvals, INSERT_VALUES));
1856: } else { /* !allcolumns */
1857: PetscInt ncol = 0;
1859: #if PetscDefined(USE_CTABLE)
1860: /* diagonal part A = c->A */
1861: ncols = ai[Crow + 1] - ai[Crow];
1862: cols = PetscSafePointerPlusOffset(aj, ai[Crow]);
1863: vals = PetscSafePointerPlusOffset(a_a, ai[Crow]);
1864: for (PetscInt k = 0; k < ncols; k++) {
1865: tcol = cmap_loc[cols[k]];
1866: if (tcol) {
1867: subcols[ncol] = --tcol;
1868: if (!C->structure_only) subvals[ncol] = vals[k];
1869: ncol++;
1870: }
1871: }
1873: /* off-diagonal part B = c->B */
1874: ncols = bi[Crow + 1] - bi[Crow];
1875: cols = PetscSafePointerPlusOffset(bj, bi[Crow]);
1876: vals = PetscSafePointerPlusOffset(b_a, bi[Crow]);
1877: for (PetscInt k = 0; k < ncols; k++) {
1878: PetscCall(PetscHMapIGetWithDefault(cmap, bmap[cols[k]] + 1, 0, &tcol));
1879: if (tcol) {
1880: subcols[ncol] = --tcol;
1881: if (!C->structure_only) subvals[ncol] = vals[k];
1882: ncol++;
1883: }
1884: }
1885: #else
1886: /* diagonal part A = c->A */
1887: ncols = ai[Crow + 1] - ai[Crow];
1888: cols = aj + ai[Crow];
1889: vals = PetscSafePointerPlusOffset(a_a, ai[Crow]);
1890: for (PetscInt k = 0; k < ncols; k++) {
1891: tcol = cmap[cols[k] + cstart];
1892: if (tcol) {
1893: subcols[ncol] = --tcol;
1894: if (!C->structure_only) subvals[ncol] = vals[k];
1895: ncol++;
1896: }
1897: }
1899: /* off-diagonal part B = c->B */
1900: ncols = bi[Crow + 1] - bi[Crow];
1901: cols = bj + bi[Crow];
1902: vals = PetscSafePointerPlusOffset(b_a, bi[Crow]);
1903: for (PetscInt k = 0; k < ncols; k++) {
1904: tcol = cmap[bmap[cols[k]]];
1905: if (tcol) {
1906: subcols[ncol] = --tcol;
1907: if (!C->structure_only) subvals[ncol] = vals[k];
1908: ncol++;
1909: }
1910: }
1911: #endif
1912: PetscCall(MatSetValues_SeqAIJ(submat, 1, &row, ncol, subcols, subvals, INSERT_VALUES));
1913: }
1914: }
1915: }
1917: /* Now assemble the off-proc rows */
1918: for (PetscMPIInt i = 0; i < nrqs; i++) { /* for each requested message */
1919: /* recv values from other processes */
1920: if (!C->structure_only) PetscCallMPI(MPI_Waitany(nrqs, r_waits4, &idex, r_status4 + i));
1921: else idex = i;
1922: sbuf1_i = sbuf1[pa[idex]];
1923: /* jmax = sbuf1_i[0]; PetscCheck(jmax == 1,PETSC_COMM_SELF,PETSC_ERR_PLIB,"jmax %d != 1",jmax); */
1924: ct1 = 2 + 1;
1925: ct2 = 0; /* count of received C->j */
1926: ct3 = 0; /* count of received C->j that will be inserted into submat */
1927: rbuf2_i = rbuf2[idex]; /* int** received length of C->j from other processes */
1928: rbuf3_i = rbuf3[idex]; /* int** received C->j from other processes */
1929: rbuf4_i = C->structure_only ? NULL : rbuf4[idex]; /* scalar** received C->a from other processes */
1931: /* is_no = sbuf1_i[2*j-1]; PetscCheck(is_no == 0,PETSC_COMM_SELF,PETSC_ERR_PLIB,"is_no !=0"); */
1932: max1 = sbuf1_i[2]; /* num of rows */
1933: for (PetscInt k = 0; k < max1; k++, ct1++) { /* for each recved row */
1934: row = sbuf1_i[ct1]; /* row index of submat */
1935: if (!allcolumns) {
1936: idex = 0;
1937: if (scall == MAT_INITIAL_MATRIX || !iscolsorted) {
1938: nnz = rbuf2_i[ct1]; /* num of C entries in this row */
1939: for (PetscInt l = 0; l < nnz; l++, ct2++) { /* for each recved column */
1940: #if PetscDefined(USE_CTABLE)
1941: if (rbuf3_i[ct2] >= cstart && rbuf3_i[ct2] < cend) {
1942: tcol = cmap_loc[rbuf3_i[ct2] - cstart];
1943: } else {
1944: PetscCall(PetscHMapIGetWithDefault(cmap, rbuf3_i[ct2] + 1, 0, &tcol));
1945: }
1946: #else
1947: tcol = cmap[rbuf3_i[ct2]];
1948: #endif
1949: if (tcol) {
1950: subcols[idex] = --tcol; /* may not be sorted */
1951: if (!C->structure_only) subvals[idex] = rbuf4_i[ct2];
1952: idex++;
1954: /* We receive an entire column of C, but a subset of it needs to be inserted into submat.
1955: For reuse, we replace received C->j with index that should be inserted to submat */
1956: if (iscolsorted) rbuf3_i[ct3++] = ct2;
1957: }
1958: }
1959: PetscCall(MatSetValues_SeqAIJ(submat, 1, &row, idex, subcols, subvals, INSERT_VALUES));
1960: } else { /* scall == MAT_REUSE_MATRIX */
1961: submat = submats[0];
1962: subc = (Mat_SeqAIJ *)submat->data;
1964: nnz = subc->i[row + 1] - subc->i[row]; /* num of submat entries in this row */
1965: for (PetscInt l = 0; l < nnz; l++) {
1966: ct2 = rbuf3_i[ct3++]; /* index of rbuf4_i[] which needs to be inserted into submat */
1967: if (direct_csr_reuse) suba[subc->i[row] + idex++] = rbuf4_i[ct2];
1968: else subvals[idex++] = rbuf4_i[ct2];
1969: }
1970: if (!direct_csr_reuse) {
1971: bj = subc->j + subc->i[row]; /* sorted column indices */
1972: PetscCall(MatSetValues_SeqAIJ(submat, 1, &row, nnz, bj, subvals, INSERT_VALUES));
1973: }
1974: }
1975: } else { /* allcolumns */
1976: nnz = rbuf2_i[ct1]; /* num of C entries in this row */
1977: PetscCall(MatSetValues_SeqAIJ(submat, 1, &row, nnz, PetscSafePointerPlusOffset(rbuf3_i, ct2), PetscSafePointerPlusOffset(rbuf4_i, ct2), INSERT_VALUES));
1978: ct2 += nnz;
1979: }
1980: }
1981: }
1983: /* sending a->a are done */
1984: if (!C->structure_only) {
1985: PetscCallMPI(MPI_Waitall(nrqr, s_waits4, s_status4));
1986: PetscCall(PetscFree4(r_waits4, s_waits4, r_status4, s_status4));
1987: }
1989: if (direct_csr_reuse) PetscCall(MatSeqAIJRestoreArray(submat, &suba));
1990: PetscCall(MatAssemblyBegin(submat, MAT_FINAL_ASSEMBLY));
1991: PetscCall(MatAssemblyEnd(submat, MAT_FINAL_ASSEMBLY));
1992: // Preserve the initial graph state even when no reuse maps have been built.
1993: if (scall == MAT_INITIAL_MATRIX && !C->structure_only && iscolsorted && !allcolumns) PetscCall(MatGetNonzeroState(submat, &smatis1->nonzerostate));
1994: submats[0] = submat;
1996: /* Restore the indices */
1997: PetscCall(ISRestoreIndices(isrow[0], &irow));
1998: if (!allcolumns) PetscCall(ISRestoreIndices(iscol[0], &icol));
2000: /* Destroy allocated memory */
2001: PetscCall(PetscFree(subcols));
2002: if (!C->structure_only) {
2003: for (PetscMPIInt i = 0; i < nrqs; ++i) PetscCall(PetscFree(rbuf4[i]));
2004: PetscCall(PetscFree2(rbuf4, subvals));
2005: if (sbuf_aa) {
2006: PetscCall(PetscFree(sbuf_aa[0]));
2007: PetscCall(PetscFree(sbuf_aa));
2008: }
2009: }
2011: if (scall == MAT_INITIAL_MATRIX) {
2012: PetscCall(PetscFree(lens));
2013: if (sbuf_aj) {
2014: PetscCall(PetscFree(sbuf_aj[0]));
2015: PetscCall(PetscFree(sbuf_aj));
2016: }
2017: }
2018: PetscCall(MatSeqAIJRestoreArrayRead(A, (const PetscScalar **)&a_a));
2019: PetscCall(MatSeqAIJRestoreArrayRead(B, (const PetscScalar **)&b_a));
2020: PetscFunctionReturn(PETSC_SUCCESS);
2021: }
2023: static PetscErrorCode MatCreateSubMatrices_MPIAIJ_SingleIS(Mat C, PetscInt ismax, const IS isrow[], const IS iscol[], MatReuse scall, Mat *submat[])
2024: {
2025: PetscInt ncol;
2026: PetscBool colflag, allcolumns = PETSC_FALSE;
2028: PetscFunctionBegin;
2029: /* Allocate memory to hold all the submatrices */
2030: if (scall == MAT_INITIAL_MATRIX) PetscCall(PetscCalloc1(2, submat));
2032: /* Check for special case: each processor gets entire matrix columns */
2033: PetscCall(ISIdentity(iscol[0], &colflag));
2034: PetscCall(ISGetLocalSize(iscol[0], &ncol));
2035: if (colflag && ncol == C->cmap->N) allcolumns = PETSC_TRUE;
2037: PetscCall(MatCreateSubMatrices_MPIAIJ_SingleIS_Local(C, ismax, isrow, iscol, scall, allcolumns, *submat));
2038: PetscFunctionReturn(PETSC_SUCCESS);
2039: }
2041: PetscErrorCode MatCreateSubMatrices_MPIAIJ(Mat C, PetscInt ismax, const IS isrow[], const IS iscol[], MatReuse scall, Mat *submat[])
2042: {
2043: PetscInt nmax, nstages = 0, max_no, nrow, ncol, out[2];
2044: PetscBool rowflag, colflag, wantallmatrix = PETSC_FALSE;
2045: Mat_SeqAIJ *subc;
2046: Mat_SubSppt *smat;
2048: PetscFunctionBegin;
2049: /* Check for special case: each processor has a single IS */
2050: if (C->submat_singleis) { /* flag is set in PCSetUp_ASM() to skip MPI_Allreduce() */
2051: PetscCall(MatCreateSubMatrices_MPIAIJ_SingleIS(C, ismax, isrow, iscol, scall, submat));
2052: C->submat_singleis = PETSC_FALSE; /* resume its default value in case C will be used for non-single IS */
2053: PetscFunctionReturn(PETSC_SUCCESS);
2054: }
2056: /* Collect global wantallmatrix and nstages */
2057: if (!C->cmap->N) nmax = 20 * 1000000 / sizeof(PetscInt);
2058: else nmax = 20 * 1000000 / (C->cmap->N * sizeof(PetscInt));
2059: if (!nmax) nmax = 1;
2061: if (scall == MAT_INITIAL_MATRIX) {
2062: /* Collect global wantallmatrix and nstages */
2063: if (ismax == 1 && C->rmap->N == C->cmap->N) {
2064: PetscCall(ISIdentity(*isrow, &rowflag));
2065: PetscCall(ISIdentity(*iscol, &colflag));
2066: PetscCall(ISGetLocalSize(*isrow, &nrow));
2067: PetscCall(ISGetLocalSize(*iscol, &ncol));
2068: if (rowflag && colflag && nrow == C->rmap->N && ncol == C->cmap->N) {
2069: wantallmatrix = PETSC_TRUE;
2071: PetscCall(PetscOptionsGetBool(((PetscObject)C)->options, ((PetscObject)C)->prefix, "-use_fast_submatrix", &wantallmatrix, NULL));
2072: }
2073: }
2075: /* Determine the number of stages through which submatrices are done
2076: Each stage will extract nmax submatrices.
2077: nmax is determined by the matrix column dimension.
2078: If the original matrix has 20M columns, only one submatrix per stage is allowed, etc.
2079: */
2080: nstages = ismax / nmax + ((ismax % nmax) ? 1 : 0); /* local nstages */
2082: out[0] = -1 * (PetscInt)wantallmatrix;
2083: out[1] = nstages;
2084: PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, out, 2, MPIU_INT, MPI_MAX, PetscObjectComm((PetscObject)C)));
2085: wantallmatrix = (PetscBool)(-out[0]);
2086: nstages = out[1]; /* Make sure every processor loops through the global nstages */
2088: } else { /* MAT_REUSE_MATRIX */
2089: if (ismax) {
2090: subc = (Mat_SeqAIJ *)(*submat)[0]->data;
2091: smat = subc->submatis1;
2092: } else { /* (*submat)[0] is a dummy matrix */
2093: smat = (Mat_SubSppt *)(*submat)[0]->data;
2094: }
2095: if (!smat) {
2096: /* smat is not generated by MatCreateSubMatrix_MPIAIJ_All(...,MAT_INITIAL_MATRIX,...) */
2097: wantallmatrix = PETSC_TRUE;
2098: } else if (smat->singleis) {
2099: PetscCall(MatCreateSubMatrices_MPIAIJ_SingleIS(C, ismax, isrow, iscol, scall, submat));
2100: PetscFunctionReturn(PETSC_SUCCESS);
2101: } else {
2102: nstages = smat->nstages;
2103: }
2104: }
2106: if (wantallmatrix) {
2107: PetscCall(MatCreateSubMatrix_MPIAIJ_All(C, MAT_GET_VALUES, scall, submat));
2108: PetscFunctionReturn(PETSC_SUCCESS);
2109: }
2111: /* Allocate memory to hold all the submatrices and dummy submatrices */
2112: if (scall == MAT_INITIAL_MATRIX) PetscCall(PetscCalloc1(ismax + nstages, submat));
2114: for (PetscInt i = 0, pos = 0; i < nstages; i++) {
2115: if (pos + nmax <= ismax) max_no = nmax;
2116: else if (pos >= ismax) max_no = 0;
2117: else max_no = ismax - pos;
2119: PetscCall(MatCreateSubMatrices_MPIAIJ_Local(C, max_no, PetscSafePointerPlusOffset(isrow, pos), PetscSafePointerPlusOffset(iscol, pos), scall, *submat + pos));
2120: if (!max_no) {
2121: if (scall == MAT_INITIAL_MATRIX) { /* submat[pos] is a dummy matrix */
2122: smat = (Mat_SubSppt *)(*submat)[pos]->data;
2123: smat->nstages = nstages;
2124: }
2125: pos++; /* advance to next dummy matrix if any */
2126: } else pos += max_no;
2127: }
2129: if (ismax && scall == MAT_INITIAL_MATRIX) {
2130: /* save nstages for reuse */
2131: subc = (Mat_SeqAIJ *)(*submat)[0]->data;
2132: smat = subc->submatis1;
2133: smat->nstages = nstages;
2134: }
2135: PetscFunctionReturn(PETSC_SUCCESS);
2136: }
2138: PetscErrorCode MatCreateSubMatrices_MPIAIJ_Local(Mat C, PetscInt ismax, const IS isrow[], const IS iscol[], MatReuse scall, Mat *submats)
2139: {
2140: Mat_MPIAIJ *c = (Mat_MPIAIJ *)C->data;
2141: Mat A = c->A;
2142: Mat_SeqAIJ *a = (Mat_SeqAIJ *)A->data, *b = (Mat_SeqAIJ *)c->B->data, *subc;
2143: const PetscInt **icol, **irow;
2144: PetscInt *nrow, *ncol, start;
2145: PetscMPIInt nrqs = 0, *pa, proc = -1;
2146: PetscMPIInt rank, size, tag0, tag2, tag3, tag4, *w1, *w2, *w3, *w4, nrqr, *req_source1 = NULL, *req_source2;
2147: PetscInt **sbuf1, **sbuf2, k, ct1, ct2, **rbuf1, row;
2148: PetscInt msz, **ptr = NULL, *req_size = NULL, *ctr = NULL, *tmp = NULL, tcol;
2149: PetscInt **rbuf3 = NULL, **sbuf_aj, **rbuf2 = NULL, max1, max2;
2150: PetscInt **lens, is_no, ncols, *cols, mat_i, *mat_j, tmp2, jmax;
2151: #if PetscDefined(USE_CTABLE)
2152: PetscHMapI *cmap, cmap_i = NULL, *rmap, rmap_i;
2153: #else
2154: PetscInt **cmap, *cmap_i = NULL, **rmap, *rmap_i;
2155: #endif
2156: const PetscInt *irow_i;
2157: PetscInt ctr_j, *sbuf1_j, *sbuf_aj_i, *rbuf1_i, kmax, *lens_i;
2158: MPI_Request *s_waits1, *r_waits1, *s_waits2, *r_waits2, *r_waits3;
2159: MPI_Request *r_waits4, *s_waits3, *s_waits4;
2160: MPI_Comm comm;
2161: PetscScalar **rbuf4 = NULL, *rbuf4_i, **sbuf_aa = NULL, *vals, *mat_a, *imat_a, *sbuf_aa_i;
2162: PetscMPIInt *onodes1, *olengths1, end, **row2proc, *row2proc_i;
2163: PetscInt ilen_row, *imat_ilen, *imat_j, *imat_i, old_row;
2164: Mat_SubSppt *smat_i;
2165: PetscBool *issorted, *allcolumns, colflag, iscsorted = PETSC_TRUE;
2166: PetscInt *sbuf1_i, *rbuf2_i, *rbuf3_i, ilen, jcnt;
2168: PetscFunctionBegin;
2169: PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
2170: size = c->size;
2171: rank = c->rank;
2173: PetscCall(PetscMalloc4(ismax, &row2proc, ismax, &cmap, ismax, &rmap, ismax + 1, &allcolumns));
2174: PetscCall(PetscMalloc5(ismax, (PetscInt ***)&irow, ismax, (PetscInt ***)&icol, ismax, &nrow, ismax, &ncol, ismax, &issorted));
2176: for (PetscInt i = 0; i < ismax; i++) {
2177: PetscCall(ISSorted(iscol[i], &issorted[i]));
2178: if (!issorted[i]) iscsorted = issorted[i];
2180: PetscCall(ISSorted(isrow[i], &issorted[i]));
2182: PetscCall(ISGetIndices(isrow[i], &irow[i]));
2183: PetscCall(ISGetLocalSize(isrow[i], &nrow[i]));
2185: /* Check for special case: allcolumn */
2186: PetscCall(ISIdentity(iscol[i], &colflag));
2187: PetscCall(ISGetLocalSize(iscol[i], &ncol[i]));
2188: if (colflag && ncol[i] == C->cmap->N) {
2189: allcolumns[i] = PETSC_TRUE;
2190: icol[i] = NULL;
2191: } else {
2192: allcolumns[i] = PETSC_FALSE;
2193: PetscCall(ISGetIndices(iscol[i], &icol[i]));
2194: }
2195: }
2197: if (scall == MAT_REUSE_MATRIX) {
2198: /* Assumes new rows are same length as the old rows */
2199: for (PetscInt i = 0; i < ismax; i++) {
2200: PetscCheck(submats[i], PETSC_COMM_SELF, PETSC_ERR_ARG_NULL, "submats[%" PetscInt_FMT "] is null, cannot reuse", i);
2201: subc = (Mat_SeqAIJ *)submats[i]->data;
2202: PetscCheck(!(submats[i]->rmap->n != nrow[i]) && !(submats[i]->cmap->n != ncol[i]), PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Cannot reuse matrix. wrong size");
2204: /* Initial matrix as if empty */
2205: PetscCall(PetscArrayzero(subc->ilen, submats[i]->rmap->n));
2207: smat_i = subc->submatis1;
2209: nrqs = smat_i->nrqs;
2210: nrqr = smat_i->nrqr;
2211: rbuf1 = smat_i->rbuf1;
2212: rbuf2 = smat_i->rbuf2;
2213: rbuf3 = smat_i->rbuf3;
2214: req_source2 = smat_i->req_source2;
2216: sbuf1 = smat_i->sbuf1;
2217: sbuf2 = smat_i->sbuf2;
2218: ptr = smat_i->ptr;
2219: tmp = smat_i->tmp;
2220: ctr = smat_i->ctr;
2222: pa = smat_i->pa;
2223: req_size = smat_i->req_size;
2224: req_source1 = smat_i->req_source1;
2226: allcolumns[i] = smat_i->allcolumns;
2227: row2proc[i] = smat_i->row2proc;
2228: rmap[i] = smat_i->rmap;
2229: cmap[i] = smat_i->cmap;
2230: }
2232: if (!ismax) { /* Get dummy submatrices and retrieve struct submatis1 */
2233: PetscCheck(submats[0], PETSC_COMM_SELF, PETSC_ERR_ARG_NULL, "submats are null, cannot reuse");
2234: smat_i = (Mat_SubSppt *)submats[0]->data;
2236: nrqs = smat_i->nrqs;
2237: nrqr = smat_i->nrqr;
2238: rbuf1 = smat_i->rbuf1;
2239: rbuf2 = smat_i->rbuf2;
2240: rbuf3 = smat_i->rbuf3;
2241: req_source2 = smat_i->req_source2;
2243: sbuf1 = smat_i->sbuf1;
2244: sbuf2 = smat_i->sbuf2;
2245: ptr = smat_i->ptr;
2246: tmp = smat_i->tmp;
2247: ctr = smat_i->ctr;
2249: pa = smat_i->pa;
2250: req_size = smat_i->req_size;
2251: req_source1 = smat_i->req_source1;
2253: allcolumns[0] = PETSC_FALSE;
2254: }
2255: } else { /* scall == MAT_INITIAL_MATRIX */
2256: /* Get some new tags to keep the communication clean */
2257: PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag2));
2258: PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag3));
2260: /* evaluate communication - mesg to who, length of mesg, and buffer space
2261: required. Based on this, buffers are allocated, and data copied into them*/
2262: PetscCall(PetscCalloc4(size, &w1, size, &w2, size, &w3, size, &w4)); /* mesg size, initialize work vectors */
2264: for (PetscInt i = 0; i < ismax; i++) {
2265: jmax = nrow[i];
2266: irow_i = irow[i];
2268: PetscCall(PetscMalloc1(jmax, &row2proc_i));
2269: row2proc[i] = row2proc_i;
2271: if (issorted[i]) proc = 0;
2272: for (PetscInt j = 0; j < jmax; j++) {
2273: if (!issorted[i]) proc = 0;
2274: row = irow_i[j];
2275: while (row >= C->rmap->range[proc + 1]) proc++;
2276: w4[proc]++;
2277: row2proc_i[j] = proc; /* map row index to proc */
2278: }
2279: for (PetscMPIInt j = 0; j < size; j++) {
2280: if (w4[j]) {
2281: w1[j] += w4[j];
2282: w3[j]++;
2283: w4[j] = 0;
2284: }
2285: }
2286: }
2288: nrqs = 0; /* no of outgoing messages */
2289: msz = 0; /* total mesg length (for all procs) */
2290: w1[rank] = 0; /* no mesg sent to self */
2291: w3[rank] = 0;
2292: for (PetscMPIInt i = 0; i < size; i++) {
2293: if (w1[i]) {
2294: w2[i] = 1;
2295: nrqs++;
2296: } /* there exists a message to proc i */
2297: }
2298: PetscCall(PetscMalloc1(nrqs, &pa)); /*(proc -array)*/
2299: for (PetscMPIInt i = 0, j = 0; i < size; i++) {
2300: if (w1[i]) {
2301: pa[j] = i;
2302: j++;
2303: }
2304: }
2306: /* Each message would have a header = 1 + 2*(no of IS) + data */
2307: for (PetscMPIInt i = 0; i < nrqs; i++) {
2308: PetscMPIInt j = pa[i];
2309: w1[j] += w2[j] + 2 * w3[j];
2310: msz += w1[j];
2311: }
2312: PetscCall(PetscInfo(0, "Number of outgoing messages %d Total message length %" PetscInt_FMT "\n", nrqs, msz));
2314: /* Determine the number of messages to expect, their lengths, from from-ids */
2315: PetscCall(PetscGatherNumberOfMessages(comm, w2, w1, &nrqr));
2316: PetscCall(PetscGatherMessageLengths(comm, nrqs, nrqr, w1, &onodes1, &olengths1));
2318: /* Now post the Irecvs corresponding to these messages */
2319: PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag0));
2320: PetscCall(PetscPostIrecvInt(comm, tag0, nrqr, onodes1, olengths1, &rbuf1, &r_waits1));
2322: /* Allocate Memory for outgoing messages */
2323: PetscCall(PetscMalloc4(size, &sbuf1, size, &ptr, 2 * msz, &tmp, size, &ctr));
2324: PetscCall(PetscArrayzero(sbuf1, size));
2325: PetscCall(PetscArrayzero(ptr, size));
2327: {
2328: PetscInt *iptr = tmp;
2329: k = 0;
2330: for (PetscMPIInt i = 0; i < nrqs; i++) {
2331: PetscMPIInt j = pa[i];
2332: iptr += k;
2333: sbuf1[j] = iptr;
2334: k = w1[j];
2335: }
2336: }
2338: /* Form the outgoing messages. Initialize the header space */
2339: for (PetscMPIInt i = 0; i < nrqs; i++) {
2340: PetscMPIInt j = pa[i];
2341: sbuf1[j][0] = 0;
2342: PetscCall(PetscArrayzero(sbuf1[j] + 1, 2 * w3[j]));
2343: ptr[j] = sbuf1[j] + 2 * w3[j] + 1;
2344: }
2346: /* Parse the isrow and copy data into outbuf */
2347: for (PetscInt i = 0; i < ismax; i++) {
2348: row2proc_i = row2proc[i];
2349: PetscCall(PetscArrayzero(ctr, size));
2350: irow_i = irow[i];
2351: jmax = nrow[i];
2352: for (PetscInt j = 0; j < jmax; j++) { /* parse the indices of each IS */
2353: proc = row2proc_i[j];
2354: if (proc != rank) { /* copy to the outgoing buf */
2355: ctr[proc]++;
2356: *ptr[proc] = irow_i[j];
2357: ptr[proc]++;
2358: }
2359: }
2360: /* Update the headers for the current IS */
2361: for (PetscMPIInt j = 0; j < size; j++) { /* Can Optimise this loop too */
2362: if ((ctr_j = ctr[j])) {
2363: sbuf1_j = sbuf1[j];
2364: k = ++sbuf1_j[0];
2365: sbuf1_j[2 * k] = ctr_j;
2366: sbuf1_j[2 * k - 1] = i;
2367: }
2368: }
2369: }
2371: /* Now post the sends */
2372: PetscCall(PetscMalloc1(nrqs, &s_waits1));
2373: for (PetscMPIInt i = 0; i < nrqs; ++i) {
2374: PetscMPIInt j = pa[i];
2375: PetscCallMPI(MPIU_Isend(sbuf1[j], w1[j], MPIU_INT, j, tag0, comm, s_waits1 + i));
2376: }
2378: /* Post Receives to capture the buffer size */
2379: PetscCall(PetscMalloc1(nrqs, &r_waits2));
2380: PetscCall(PetscMalloc3(nrqs, &req_source2, nrqs, &rbuf2, nrqs, &rbuf3));
2381: if (nrqs) rbuf2[0] = tmp + msz;
2382: for (PetscMPIInt i = 1; i < nrqs; ++i) rbuf2[i] = rbuf2[i - 1] + w1[pa[i - 1]];
2383: for (PetscMPIInt i = 0; i < nrqs; ++i) {
2384: PetscMPIInt j = pa[i];
2385: PetscCallMPI(MPIU_Irecv(rbuf2[i], w1[j], MPIU_INT, j, tag2, comm, r_waits2 + i));
2386: }
2388: /* Send to other procs the buf size they should allocate */
2389: /* Receive messages*/
2390: PetscCall(PetscMalloc1(nrqr, &s_waits2));
2391: PetscCall(PetscMalloc3(nrqr, &sbuf2, nrqr, &req_size, nrqr, &req_source1));
2392: {
2393: PetscInt *sAi = a->i, *sBi = b->i, id, rstart = C->rmap->rstart;
2394: PetscInt *sbuf2_i;
2396: PetscCallMPI(MPI_Waitall(nrqr, r_waits1, MPI_STATUSES_IGNORE));
2397: for (PetscMPIInt i = 0; i < nrqr; ++i) {
2398: req_size[i] = 0;
2399: rbuf1_i = rbuf1[i];
2400: start = 2 * rbuf1_i[0] + 1;
2401: end = olengths1[i];
2402: PetscCall(PetscMalloc1(end, &sbuf2[i]));
2403: sbuf2_i = sbuf2[i];
2404: for (PetscInt j = start; j < end; j++) {
2405: id = rbuf1_i[j] - rstart;
2406: ncols = sAi[id + 1] - sAi[id] + sBi[id + 1] - sBi[id];
2407: sbuf2_i[j] = ncols;
2408: req_size[i] += ncols;
2409: }
2410: req_source1[i] = onodes1[i];
2411: /* form the header */
2412: sbuf2_i[0] = req_size[i];
2413: for (PetscInt j = 1; j < start; j++) sbuf2_i[j] = rbuf1_i[j];
2415: PetscCallMPI(MPIU_Isend(sbuf2_i, end, MPIU_INT, req_source1[i], tag2, comm, s_waits2 + i));
2416: }
2417: }
2419: PetscCall(PetscFree(onodes1));
2420: PetscCall(PetscFree(olengths1));
2421: PetscCall(PetscFree(r_waits1));
2422: PetscCall(PetscFree4(w1, w2, w3, w4));
2424: /* Receive messages*/
2425: PetscCall(PetscMalloc1(nrqs, &r_waits3));
2426: PetscCallMPI(MPI_Waitall(nrqs, r_waits2, MPI_STATUSES_IGNORE));
2427: for (PetscMPIInt i = 0; i < nrqs; ++i) {
2428: PetscCall(PetscMalloc1(rbuf2[i][0], &rbuf3[i]));
2429: req_source2[i] = pa[i];
2430: PetscCallMPI(MPIU_Irecv(rbuf3[i], rbuf2[i][0], MPIU_INT, req_source2[i], tag3, comm, r_waits3 + i));
2431: }
2432: PetscCall(PetscFree(r_waits2));
2434: /* Wait on sends1 and sends2 */
2435: PetscCallMPI(MPI_Waitall(nrqs, s_waits1, MPI_STATUSES_IGNORE));
2436: PetscCallMPI(MPI_Waitall(nrqr, s_waits2, MPI_STATUSES_IGNORE));
2437: PetscCall(PetscFree(s_waits1));
2438: PetscCall(PetscFree(s_waits2));
2440: /* Now allocate sending buffers for a->j, and send them off */
2441: PetscCall(PetscMalloc1(nrqr, &sbuf_aj));
2442: jcnt = 0;
2443: for (PetscMPIInt i = 0; i < nrqr; i++) jcnt += req_size[i];
2444: if (nrqr) PetscCall(PetscMalloc1(jcnt, &sbuf_aj[0]));
2445: for (PetscMPIInt i = 1; i < nrqr; i++) sbuf_aj[i] = sbuf_aj[i - 1] + req_size[i - 1];
2447: PetscCall(PetscMalloc1(nrqr, &s_waits3));
2448: {
2449: PetscInt nzA, nzB, *a_i = a->i, *b_i = b->i, lwrite;
2450: PetscInt *cworkA, *cworkB, cstart = C->cmap->rstart, rstart = C->rmap->rstart, *bmap = c->garray;
2451: PetscInt cend = C->cmap->rend;
2452: PetscInt *a_j = a->j, *b_j = b->j, ctmp;
2454: for (PetscMPIInt i = 0; i < nrqr; i++) {
2455: rbuf1_i = rbuf1[i];
2456: sbuf_aj_i = sbuf_aj[i];
2457: ct1 = 2 * rbuf1_i[0] + 1;
2458: ct2 = 0;
2459: for (PetscInt j = 1, max1 = rbuf1_i[0]; j <= max1; j++) {
2460: kmax = rbuf1[i][2 * j];
2461: for (PetscInt k = 0; k < kmax; k++, ct1++) {
2462: row = rbuf1_i[ct1] - rstart;
2463: nzA = a_i[row + 1] - a_i[row];
2464: nzB = b_i[row + 1] - b_i[row];
2465: ncols = nzA + nzB;
2466: cworkA = PetscSafePointerPlusOffset(a_j, a_i[row]);
2467: cworkB = PetscSafePointerPlusOffset(b_j, b_i[row]);
2469: /* load the column indices for this row into cols */
2470: cols = sbuf_aj_i + ct2;
2472: lwrite = 0;
2473: for (PetscInt l = 0; l < nzB; l++) {
2474: if ((ctmp = bmap[cworkB[l]]) < cstart) cols[lwrite++] = ctmp;
2475: }
2476: for (PetscInt l = 0; l < nzA; l++) cols[lwrite++] = cstart + cworkA[l];
2477: for (PetscInt l = 0; l < nzB; l++) {
2478: if ((ctmp = bmap[cworkB[l]]) >= cend) cols[lwrite++] = ctmp;
2479: }
2481: ct2 += ncols;
2482: }
2483: }
2484: PetscCallMPI(MPIU_Isend(sbuf_aj_i, req_size[i], MPIU_INT, req_source1[i], tag3, comm, s_waits3 + i));
2485: }
2486: }
2488: /* create col map: global col of C -> local col of submatrices */
2489: {
2490: const PetscInt *icol_i;
2491: #if PetscDefined(USE_CTABLE)
2492: for (PetscInt i = 0; i < ismax; i++) {
2493: if (!allcolumns[i]) {
2494: PetscCall(PetscHMapICreateWithSize(ncol[i], cmap + i));
2496: jmax = ncol[i];
2497: icol_i = icol[i];
2498: cmap_i = cmap[i];
2499: for (PetscInt j = 0; j < jmax; j++) PetscCall(PetscHMapISet(cmap[i], icol_i[j] + 1, j + 1));
2500: } else cmap[i] = NULL;
2501: }
2502: #else
2503: for (PetscInt i = 0; i < ismax; i++) {
2504: if (!allcolumns[i]) {
2505: PetscCall(PetscCalloc1(C->cmap->N, &cmap[i]));
2506: jmax = ncol[i];
2507: icol_i = icol[i];
2508: cmap_i = cmap[i];
2509: for (PetscInt j = 0; j < jmax; j++) cmap_i[icol_i[j]] = j + 1;
2510: } else cmap[i] = NULL;
2511: }
2512: #endif
2513: }
2515: /* Create lens which is required for MatCreate... */
2516: jcnt = 0;
2517: for (PetscInt i = 0; i < ismax; i++) jcnt += nrow[i];
2518: PetscCall(PetscMalloc1(ismax, &lens));
2520: if (ismax) PetscCall(PetscCalloc1(jcnt, &lens[0]));
2521: for (PetscInt i = 1; i < ismax; i++) lens[i] = PetscSafePointerPlusOffset(lens[i - 1], nrow[i - 1]);
2523: /* Update lens from local data */
2524: for (PetscInt i = 0; i < ismax; i++) {
2525: row2proc_i = row2proc[i];
2526: jmax = nrow[i];
2527: if (!allcolumns[i]) cmap_i = cmap[i];
2528: irow_i = irow[i];
2529: lens_i = lens[i];
2530: for (PetscInt j = 0; j < jmax; j++) {
2531: row = irow_i[j];
2532: proc = row2proc_i[j];
2533: if (proc == rank) {
2534: PetscCall(MatGetRow_MPIAIJ(C, row, &ncols, &cols, NULL));
2535: if (!allcolumns[i]) {
2536: for (PetscInt k = 0; k < ncols; k++) {
2537: #if PetscDefined(USE_CTABLE)
2538: PetscCall(PetscHMapIGetWithDefault(cmap_i, cols[k] + 1, 0, &tcol));
2539: #else
2540: tcol = cmap_i[cols[k]];
2541: #endif
2542: if (tcol) lens_i[j]++;
2543: }
2544: } else { /* allcolumns */
2545: lens_i[j] = ncols;
2546: }
2547: PetscCall(MatRestoreRow_MPIAIJ(C, row, &ncols, &cols, NULL));
2548: }
2549: }
2550: }
2552: /* Create row map: global row of C -> local row of submatrices */
2553: #if PetscDefined(USE_CTABLE)
2554: for (PetscInt i = 0; i < ismax; i++) {
2555: PetscCall(PetscHMapICreateWithSize(nrow[i], rmap + i));
2556: irow_i = irow[i];
2557: jmax = nrow[i];
2558: for (PetscInt j = 0; j < jmax; j++) PetscCall(PetscHMapISet(rmap[i], irow_i[j] + 1, j + 1));
2559: }
2560: #else
2561: for (PetscInt i = 0; i < ismax; i++) {
2562: PetscCall(PetscCalloc1(C->rmap->N, &rmap[i]));
2563: rmap_i = rmap[i];
2564: irow_i = irow[i];
2565: jmax = nrow[i];
2566: for (PetscInt j = 0; j < jmax; j++) rmap_i[irow_i[j]] = j;
2567: }
2568: #endif
2570: /* Update lens from offproc data */
2571: {
2572: PetscInt *rbuf2_i, *rbuf3_i, *sbuf1_i;
2574: PetscCallMPI(MPI_Waitall(nrqs, r_waits3, MPI_STATUSES_IGNORE));
2575: for (tmp2 = 0; tmp2 < nrqs; tmp2++) {
2576: sbuf1_i = sbuf1[pa[tmp2]];
2577: jmax = sbuf1_i[0];
2578: ct1 = 2 * jmax + 1;
2579: ct2 = 0;
2580: rbuf2_i = rbuf2[tmp2];
2581: rbuf3_i = rbuf3[tmp2];
2582: for (PetscInt j = 1; j <= jmax; j++) {
2583: is_no = sbuf1_i[2 * j - 1];
2584: max1 = sbuf1_i[2 * j];
2585: lens_i = lens[is_no];
2586: if (!allcolumns[is_no]) cmap_i = cmap[is_no];
2587: rmap_i = rmap[is_no];
2588: for (PetscInt k = 0; k < max1; k++, ct1++) {
2589: #if PetscDefined(USE_CTABLE)
2590: PetscCall(PetscHMapIGetWithDefault(rmap_i, sbuf1_i[ct1] + 1, 0, &row));
2591: row--;
2592: PetscCheck(row >= 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "row not found in table");
2593: #else
2594: row = rmap_i[sbuf1_i[ct1]]; /* the val in the new matrix to be */
2595: #endif
2596: max2 = rbuf2_i[ct1];
2597: for (PetscInt l = 0; l < max2; l++, ct2++) {
2598: if (!allcolumns[is_no]) {
2599: #if PetscDefined(USE_CTABLE)
2600: PetscCall(PetscHMapIGetWithDefault(cmap_i, rbuf3_i[ct2] + 1, 0, &tcol));
2601: #else
2602: tcol = cmap_i[rbuf3_i[ct2]];
2603: #endif
2604: if (tcol) lens_i[row]++;
2605: } else { /* allcolumns */
2606: lens_i[row]++; /* lens_i[row] += max2 ? */
2607: }
2608: }
2609: }
2610: }
2611: }
2612: }
2613: PetscCall(PetscFree(r_waits3));
2614: PetscCallMPI(MPI_Waitall(nrqr, s_waits3, MPI_STATUSES_IGNORE));
2615: PetscCall(PetscFree(s_waits3));
2617: /* Create the submatrices */
2618: for (PetscInt i = 0; i < ismax; i++) {
2619: PetscInt rbs, cbs;
2621: PetscCall(ISGetBlockSize(isrow[i], &rbs));
2622: PetscCall(ISGetBlockSize(iscol[i], &cbs));
2624: PetscCall(MatCreate(PETSC_COMM_SELF, submats + i));
2625: PetscCall(MatSetSizes(submats[i], nrow[i], ncol[i], PETSC_DETERMINE, PETSC_DETERMINE));
2627: if (rbs > 1 || cbs > 1) PetscCall(MatSetBlockSizes(submats[i], rbs, cbs));
2628: PetscCall(MatSetType(submats[i], ((PetscObject)A)->type_name));
2629: PetscCall(MatSetOption(submats[i], MAT_STRUCTURE_ONLY, C->structure_only));
2630: PetscCall(MatSeqAIJSetPreallocation(submats[i], 0, lens[i]));
2632: /* create struct Mat_SubSppt and attached it to submat */
2633: PetscCall(PetscNew(&smat_i));
2634: subc = (Mat_SeqAIJ *)submats[i]->data;
2635: subc->submatis1 = smat_i;
2637: smat_i->destroy = submats[i]->ops->destroy;
2638: submats[i]->ops->destroy = MatDestroySubMatrix_SeqAIJ;
2639: submats[i]->factortype = C->factortype;
2641: smat_i->id = i;
2642: smat_i->nrqs = nrqs;
2643: smat_i->nrqr = nrqr;
2644: smat_i->rbuf1 = rbuf1;
2645: smat_i->rbuf2 = rbuf2;
2646: smat_i->rbuf3 = rbuf3;
2647: smat_i->sbuf2 = sbuf2;
2648: smat_i->req_source2 = req_source2;
2650: smat_i->sbuf1 = sbuf1;
2651: smat_i->ptr = ptr;
2652: smat_i->tmp = tmp;
2653: smat_i->ctr = ctr;
2655: smat_i->pa = pa;
2656: smat_i->req_size = req_size;
2657: smat_i->req_source1 = req_source1;
2659: smat_i->allcolumns = allcolumns[i];
2660: smat_i->singleis = PETSC_FALSE;
2661: smat_i->row2proc = row2proc[i];
2662: smat_i->rmap = rmap[i];
2663: smat_i->cmap = cmap[i];
2664: }
2666: if (!ismax) { /* Create dummy submats[0] for reuse struct subc */
2667: PetscCall(MatCreate(PETSC_COMM_SELF, &submats[0]));
2668: PetscCall(MatSetSizes(submats[0], 0, 0, PETSC_DETERMINE, PETSC_DETERMINE));
2669: PetscCall(MatSetType(submats[0], MATDUMMY));
2671: /* create struct Mat_SubSppt and attached it to submat */
2672: PetscCall(PetscNew(&smat_i));
2673: submats[0]->data = (void *)smat_i;
2675: smat_i->destroy = submats[0]->ops->destroy;
2676: submats[0]->ops->destroy = MatDestroySubMatrix_Dummy;
2677: submats[0]->factortype = C->factortype;
2679: smat_i->id = 0;
2680: smat_i->nrqs = nrqs;
2681: smat_i->nrqr = nrqr;
2682: smat_i->rbuf1 = rbuf1;
2683: smat_i->rbuf2 = rbuf2;
2684: smat_i->rbuf3 = rbuf3;
2685: smat_i->sbuf2 = sbuf2;
2686: smat_i->req_source2 = req_source2;
2688: smat_i->sbuf1 = sbuf1;
2689: smat_i->ptr = ptr;
2690: smat_i->tmp = tmp;
2691: smat_i->ctr = ctr;
2693: smat_i->pa = pa;
2694: smat_i->req_size = req_size;
2695: smat_i->req_source1 = req_source1;
2697: smat_i->allcolumns = PETSC_FALSE;
2698: smat_i->singleis = PETSC_FALSE;
2699: smat_i->row2proc = NULL;
2700: smat_i->rmap = NULL;
2701: smat_i->cmap = NULL;
2702: }
2704: if (ismax) PetscCall(PetscFree(lens[0]));
2705: PetscCall(PetscFree(lens));
2706: if (sbuf_aj) {
2707: PetscCall(PetscFree(sbuf_aj[0]));
2708: PetscCall(PetscFree(sbuf_aj));
2709: }
2711: } /* endof scall == MAT_INITIAL_MATRIX */
2713: if (!C->structure_only) {
2714: /* Post recv matrix values */
2715: PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag4));
2716: PetscCall(PetscMalloc1(nrqs, &rbuf4));
2717: PetscCall(PetscMalloc1(nrqs, &r_waits4));
2718: for (PetscMPIInt i = 0; i < nrqs; ++i) {
2719: PetscCall(PetscMalloc1(rbuf2[i][0], &rbuf4[i]));
2720: PetscCallMPI(MPIU_Irecv(rbuf4[i], rbuf2[i][0], MPIU_SCALAR, req_source2[i], tag4, comm, r_waits4 + i));
2721: }
2723: /* Allocate sending buffers for a->a, and send them off */
2724: PetscCall(PetscMalloc1(nrqr, &sbuf_aa));
2725: jcnt = 0;
2726: for (PetscMPIInt i = 0; i < nrqr; i++) jcnt += req_size[i];
2727: if (nrqr) PetscCall(PetscMalloc1(jcnt, &sbuf_aa[0]));
2728: for (PetscMPIInt i = 1; i < nrqr; i++) sbuf_aa[i] = sbuf_aa[i - 1] + req_size[i - 1];
2730: PetscCall(PetscMalloc1(nrqr, &s_waits4));
2731: {
2732: PetscInt nzA, nzB, *a_i = a->i, *b_i = b->i, *cworkB, lwrite;
2733: PetscInt cstart = C->cmap->rstart, rstart = C->rmap->rstart, *bmap = c->garray;
2734: PetscInt cend = C->cmap->rend;
2735: PetscInt *b_j = b->j;
2736: PetscScalar *vworkA, *vworkB, *a_a, *b_a;
2738: PetscCall(MatSeqAIJGetArrayRead(A, (const PetscScalar **)&a_a));
2739: PetscCall(MatSeqAIJGetArrayRead(c->B, (const PetscScalar **)&b_a));
2740: for (PetscMPIInt i = 0; i < nrqr; i++) {
2741: rbuf1_i = rbuf1[i];
2742: sbuf_aa_i = sbuf_aa[i];
2743: ct1 = 2 * rbuf1_i[0] + 1;
2744: ct2 = 0;
2745: for (PetscInt j = 1, max1 = rbuf1_i[0]; j <= max1; j++) {
2746: kmax = rbuf1_i[2 * j];
2747: for (PetscInt k = 0; k < kmax; k++, ct1++) {
2748: row = rbuf1_i[ct1] - rstart;
2749: nzA = a_i[row + 1] - a_i[row];
2750: nzB = b_i[row + 1] - b_i[row];
2751: ncols = nzA + nzB;
2752: cworkB = PetscSafePointerPlusOffset(b_j, b_i[row]);
2753: vworkA = PetscSafePointerPlusOffset(a_a, a_i[row]);
2754: vworkB = PetscSafePointerPlusOffset(b_a, b_i[row]);
2756: /* load the column values for this row into vals*/
2757: vals = sbuf_aa_i + ct2;
2759: lwrite = 0;
2760: for (PetscInt l = 0; l < nzB; l++) {
2761: if (bmap[cworkB[l]] < cstart) vals[lwrite++] = vworkB[l];
2762: }
2763: for (PetscInt l = 0; l < nzA; l++) vals[lwrite++] = vworkA[l];
2764: for (PetscInt l = 0; l < nzB; l++) {
2765: if (bmap[cworkB[l]] >= cend) vals[lwrite++] = vworkB[l];
2766: }
2768: ct2 += ncols;
2769: }
2770: }
2771: PetscCallMPI(MPIU_Isend(sbuf_aa_i, req_size[i], MPIU_SCALAR, req_source1[i], tag4, comm, s_waits4 + i));
2772: }
2773: PetscCall(MatSeqAIJRestoreArrayRead(A, (const PetscScalar **)&a_a));
2774: PetscCall(MatSeqAIJRestoreArrayRead(c->B, (const PetscScalar **)&b_a));
2775: }
2776: }
2778: /* Assemble the matrices */
2779: /* First assemble the local rows */
2780: for (PetscInt i = 0; i < ismax; i++) {
2781: row2proc_i = row2proc[i];
2782: subc = (Mat_SeqAIJ *)submats[i]->data;
2783: imat_ilen = subc->ilen;
2784: imat_j = subc->j;
2785: imat_i = subc->i;
2786: PetscCall(MatSeqAIJGetArrayWrite(submats[i], &imat_a));
2788: if (!allcolumns[i]) cmap_i = cmap[i];
2789: rmap_i = rmap[i];
2790: irow_i = irow[i];
2791: jmax = nrow[i];
2792: for (PetscInt j = 0; j < jmax; j++) {
2793: row = irow_i[j];
2794: proc = row2proc_i[j];
2795: if (proc == rank) {
2796: old_row = row;
2797: #if PetscDefined(USE_CTABLE)
2798: PetscCall(PetscHMapIGetWithDefault(rmap_i, row + 1, 0, &row));
2799: row--;
2800: #else
2801: row = rmap_i[row];
2802: #endif
2803: ilen_row = imat_ilen[row];
2804: PetscCall(MatGetRow_MPIAIJ(C, old_row, &ncols, &cols, C->structure_only ? NULL : &vals));
2805: mat_i = imat_i[row];
2806: mat_a = PetscSafePointerPlusOffset(imat_a, mat_i);
2807: mat_j = PetscSafePointerPlusOffset(imat_j, mat_i);
2808: if (!allcolumns[i]) {
2809: for (PetscInt k = 0; k < ncols; k++) {
2810: #if PetscDefined(USE_CTABLE)
2811: PetscCall(PetscHMapIGetWithDefault(cmap_i, cols[k] + 1, 0, &tcol));
2812: #else
2813: tcol = cmap_i[cols[k]];
2814: #endif
2815: if (tcol) {
2816: *mat_j++ = tcol - 1;
2817: if (!C->structure_only) *mat_a++ = vals[k];
2818: ilen_row++;
2819: }
2820: }
2821: } else { /* allcolumns */
2822: for (PetscInt k = 0; k < ncols; k++, ilen_row++) {
2823: *mat_j++ = cols[k]; /* global col index! */
2824: if (!C->structure_only) *mat_a++ = vals[k];
2825: }
2826: }
2827: PetscCall(MatRestoreRow_MPIAIJ(C, old_row, &ncols, &cols, C->structure_only ? NULL : &vals));
2829: imat_ilen[row] = ilen_row;
2830: }
2831: }
2832: PetscCall(MatSeqAIJRestoreArrayWrite(submats[i], &imat_a));
2833: }
2835: /* Now assemble the off proc rows */
2836: if (!C->structure_only) PetscCallMPI(MPI_Waitall(nrqs, r_waits4, MPI_STATUSES_IGNORE));
2837: for (tmp2 = 0; tmp2 < nrqs; tmp2++) {
2838: sbuf1_i = sbuf1[pa[tmp2]];
2839: jmax = sbuf1_i[0];
2840: ct1 = 2 * jmax + 1;
2841: ct2 = 0;
2842: rbuf2_i = rbuf2[tmp2];
2843: rbuf3_i = rbuf3[tmp2];
2844: rbuf4_i = C->structure_only ? NULL : rbuf4[tmp2];
2845: for (PetscInt j = 1; j <= jmax; j++) {
2846: is_no = sbuf1_i[2 * j - 1];
2847: rmap_i = rmap[is_no];
2848: if (!allcolumns[is_no]) cmap_i = cmap[is_no];
2849: subc = (Mat_SeqAIJ *)submats[is_no]->data;
2850: imat_ilen = subc->ilen;
2851: imat_j = subc->j;
2852: imat_i = subc->i;
2853: PetscCall(MatSeqAIJGetArrayWrite(submats[is_no], &imat_a));
2854: max1 = sbuf1_i[2 * j];
2855: for (PetscInt k = 0; k < max1; k++, ct1++) {
2856: row = sbuf1_i[ct1];
2857: #if PetscDefined(USE_CTABLE)
2858: PetscCall(PetscHMapIGetWithDefault(rmap_i, row + 1, 0, &row));
2859: row--;
2860: #else
2861: row = rmap_i[row];
2862: #endif
2863: ilen = imat_ilen[row];
2864: mat_i = imat_i[row];
2865: mat_a = PetscSafePointerPlusOffset(imat_a, mat_i);
2866: mat_j = PetscSafePointerPlusOffset(imat_j, mat_i);
2867: max2 = rbuf2_i[ct1];
2868: if (!allcolumns[is_no]) {
2869: for (PetscInt l = 0; l < max2; l++, ct2++) {
2870: #if PetscDefined(USE_CTABLE)
2871: PetscCall(PetscHMapIGetWithDefault(cmap_i, rbuf3_i[ct2] + 1, 0, &tcol));
2872: #else
2873: tcol = cmap_i[rbuf3_i[ct2]];
2874: #endif
2875: if (tcol) {
2876: *mat_j++ = tcol - 1;
2877: if (!C->structure_only) *mat_a++ = rbuf4_i[ct2];
2878: ilen++;
2879: }
2880: }
2881: } else { /* allcolumns */
2882: for (PetscInt l = 0; l < max2; l++, ct2++, ilen++) {
2883: *mat_j++ = rbuf3_i[ct2]; /* same global column index of C */
2884: if (!C->structure_only) *mat_a++ = rbuf4_i[ct2];
2885: }
2886: }
2887: imat_ilen[row] = ilen;
2888: }
2889: PetscCall(MatSeqAIJRestoreArrayWrite(submats[is_no], &imat_a));
2890: }
2891: }
2893: if (!iscsorted) { /* sort column indices of the rows */
2894: for (PetscInt i = 0; i < ismax; i++) {
2895: subc = (Mat_SeqAIJ *)submats[i]->data;
2896: imat_j = subc->j;
2897: imat_i = subc->i;
2898: imat_ilen = subc->ilen;
2900: if (allcolumns[i]) continue;
2901: PetscCall(MatSeqAIJGetArrayWrite(submats[i], &imat_a));
2902: jmax = nrow[i];
2903: for (PetscInt j = 0; j < jmax; j++) {
2904: mat_i = imat_i[j];
2905: mat_a = PetscSafePointerPlusOffset(imat_a, mat_i);
2906: mat_j = PetscSafePointerPlusOffset(imat_j, mat_i);
2907: if (C->structure_only) PetscCall(PetscSortInt(imat_ilen[j], mat_j));
2908: else PetscCall(PetscSortIntWithScalarArray(imat_ilen[j], mat_j, mat_a));
2909: }
2910: PetscCall(MatSeqAIJRestoreArrayWrite(submats[i], &imat_a));
2911: }
2912: }
2914: if (!C->structure_only) {
2915: PetscCall(PetscFree(r_waits4));
2916: PetscCallMPI(MPI_Waitall(nrqr, s_waits4, MPI_STATUSES_IGNORE));
2917: PetscCall(PetscFree(s_waits4));
2918: }
2920: /* Restore the indices */
2921: for (PetscInt i = 0; i < ismax; i++) {
2922: PetscCall(ISRestoreIndices(isrow[i], irow + i));
2923: if (!allcolumns[i]) PetscCall(ISRestoreIndices(iscol[i], icol + i));
2924: }
2926: for (PetscInt i = 0; i < ismax; i++) {
2927: PetscCall(MatAssemblyBegin(submats[i], MAT_FINAL_ASSEMBLY));
2928: PetscCall(MatAssemblyEnd(submats[i], MAT_FINAL_ASSEMBLY));
2929: }
2931: /* Destroy allocated memory */
2932: if (sbuf_aa) {
2933: PetscCall(PetscFree(sbuf_aa[0]));
2934: PetscCall(PetscFree(sbuf_aa));
2935: }
2936: PetscCall(PetscFree5(*(PetscInt ***)&irow, *(PetscInt ***)&icol, nrow, ncol, issorted));
2938: if (rbuf4) {
2939: for (PetscMPIInt i = 0; i < nrqs; ++i) PetscCall(PetscFree(rbuf4[i]));
2940: PetscCall(PetscFree(rbuf4));
2941: }
2943: PetscCall(PetscFree4(row2proc, cmap, rmap, allcolumns));
2944: PetscFunctionReturn(PETSC_SUCCESS);
2945: }
2947: /*
2948: Permute A & B into C's *local* index space using rowemb,dcolemb for A and rowemb,ocolemb for B.
2949: Embeddings are supposed to be injections and the above implies that the range of rowemb is a subset
2950: of [0,m), dcolemb is in [0,n) and ocolemb is in [N-n).
2951: If pattern == DIFFERENT_NONZERO_PATTERN, C is preallocated according to A&B.
2952: After that B's columns are mapped into C's global column space, so that C is in the "disassembled"
2953: state, and needs to be "assembled" later by compressing B's column space.
2955: This function may be called in lieu of preallocation, so C should not be expected to be preallocated.
2956: Following this call, C->A & C->B have been created, even if empty.
2957: */
2958: PetscErrorCode MatSetSeqMats_MPIAIJ(Mat C, IS rowemb, IS dcolemb, IS ocolemb, MatStructure pattern, Mat A, Mat B)
2959: {
2960: /* If making this function public, change the error returned in this function away from _PLIB. */
2961: Mat_MPIAIJ *aij;
2962: Mat_SeqAIJ *Baij;
2963: PetscBool seqaij, Bdisassembled;
2964: MatType type;
2965: PetscInt m, n, *nz, ngcol, col, cstart, cend, shift, count;
2966: PetscScalar v;
2967: const PetscInt *rowindices, *colindices;
2969: PetscFunctionBegin;
2970: /* Check to make sure the component matrices (and embeddings) are compatible with C. */
2971: if (A) {
2972: PetscCall(PetscObjectBaseTypeCompare((PetscObject)A, MATSEQAIJ, &seqaij));
2973: PetscCheck(seqaij, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Diagonal matrix is of wrong type");
2974: if (rowemb) {
2975: PetscCall(ISGetLocalSize(rowemb, &m));
2976: PetscCheck(m == A->rmap->n, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Row IS of size %" PetscInt_FMT " is incompatible with diag matrix row size %" PetscInt_FMT, m, A->rmap->n);
2977: } else PetscCheck(C->rmap->n == A->rmap->n, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Diag seq matrix is row-incompatible with the MPIAIJ matrix");
2978: if (dcolemb) {
2979: PetscCall(ISGetLocalSize(dcolemb, &n));
2980: PetscCheck(n == A->cmap->n, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Diag col IS of size %" PetscInt_FMT " is incompatible with diag matrix col size %" PetscInt_FMT, n, A->cmap->n);
2981: } else PetscCheck(C->cmap->n == A->cmap->n, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Diag seq matrix is col-incompatible with the MPIAIJ matrix");
2982: }
2983: if (B) {
2984: PetscCall(PetscObjectBaseTypeCompare((PetscObject)B, MATSEQAIJ, &seqaij));
2985: PetscCheck(seqaij, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Off-diagonal matrix is of wrong type");
2986: if (rowemb) {
2987: PetscCall(ISGetLocalSize(rowemb, &m));
2988: PetscCheck(m == B->rmap->n, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Row IS of size %" PetscInt_FMT " is incompatible with off-diag matrix row size %" PetscInt_FMT, m, B->rmap->n);
2989: } else PetscCheck(C->rmap->n == B->rmap->n, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Off-diag seq matrix is row-incompatible with the MPIAIJ matrix");
2990: if (ocolemb) {
2991: PetscCall(ISGetLocalSize(ocolemb, &n));
2992: PetscCheck(n == B->cmap->n, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Off-diag col IS of size %" PetscInt_FMT " is incompatible with off-diag matrix col size %" PetscInt_FMT, n, B->cmap->n);
2993: } else PetscCheck(C->cmap->N - C->cmap->n == B->cmap->n, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Off-diag seq matrix is col-incompatible with the MPIAIJ matrix");
2994: }
2996: PetscCall(MatGetRootType_Private(C, &type));
2997: aij = (Mat_MPIAIJ *)C->data;
2998: if (!aij->A) {
2999: /* Mimic parts of MatMPIAIJSetPreallocation() */
3000: PetscCall(MatCreate(PETSC_COMM_SELF, &aij->A));
3001: PetscCall(MatSetSizes(aij->A, C->rmap->n, C->cmap->n, C->rmap->n, C->cmap->n));
3002: PetscCall(MatSetBlockSizesFromMats(aij->A, C, C));
3003: PetscCall(MatSetType(aij->A, type));
3004: }
3005: if (A) PetscCall(MatSetSeqMat_SeqAIJ(aij->A, rowemb, dcolemb, pattern, A));
3006: else PetscCall(MatSetUp(aij->A));
3007: if (B) { /* Destroy the old matrix or the column map, depending on the sparsity pattern. */
3008: /*
3009: If pattern == DIFFERENT_NONZERO_PATTERN, we reallocate B and
3010: need to "disassemble" B -- convert it to using C's global indices.
3011: To insert the values we take the safer, albeit more expensive, route of MatSetValues().
3013: If pattern == SUBSET_NONZERO_PATTERN, we do not "disassemble" B and do not reallocate;
3014: we MatZeroValues(B) first, so there may be a bunch of zeros that, perhaps, could be compacted out.
3016: TODO: Put B's values into aij->B's aij structure in place using the embedding ISs?
3017: At least avoid calling MatSetValues() and the implied searches?
3018: */
3020: if (pattern == DIFFERENT_NONZERO_PATTERN) {
3021: #if PetscDefined(USE_CTABLE)
3022: PetscCall(PetscHMapIDestroy(&aij->colmap));
3023: #else
3024: PetscCall(PetscFree(aij->colmap));
3025: /* A bit of a HACK: ideally we should deal with case aij->B all in one code block below. */
3026: #endif
3027: ngcol = 0;
3028: if (aij->lvec) PetscCall(VecGetSize(aij->lvec, &ngcol));
3029: PetscCall(PetscFree(aij->garray));
3030: PetscCall(VecDestroy(&aij->lvec));
3031: PetscCall(VecScatterDestroy(&aij->Mvctx));
3032: }
3033: if (aij->B && pattern == DIFFERENT_NONZERO_PATTERN) PetscCall(MatDestroy(&aij->B));
3034: if (aij->B && pattern == SUBSET_NONZERO_PATTERN) PetscCall(MatZeroEntries(aij->B));
3035: }
3036: Bdisassembled = PETSC_FALSE;
3037: if (!aij->B) {
3038: PetscCall(MatCreate(PETSC_COMM_SELF, &aij->B));
3039: PetscCall(MatSetSizes(aij->B, C->rmap->n, C->cmap->N, C->rmap->n, C->cmap->N));
3040: PetscCall(MatSetBlockSizesFromMats(aij->B, B, B));
3041: PetscCall(MatSetType(aij->B, type));
3042: Bdisassembled = PETSC_TRUE;
3043: }
3044: if (B) {
3045: Baij = (Mat_SeqAIJ *)B->data;
3046: rowindices = NULL;
3047: if (rowemb) PetscCall(ISGetIndices(rowemb, &rowindices));
3048: if (pattern == DIFFERENT_NONZERO_PATTERN) {
3049: PetscCall(PetscMalloc1(C->rmap->n, &nz));
3050: if (rowemb) {
3051: PetscCall(PetscArrayzero(nz, C->rmap->n));
3052: for (PetscInt i = 0; i < B->rmap->n; i++) nz[rowindices[i]] = Baij->i[i + 1] - Baij->i[i];
3053: } else {
3054: for (PetscInt i = 0; i < B->rmap->n; i++) nz[i] = Baij->i[i + 1] - Baij->i[i];
3055: }
3056: PetscCall(MatSeqAIJSetPreallocation(aij->B, 0, nz));
3057: PetscCall(PetscFree(nz));
3058: }
3060: PetscCall(PetscLayoutGetRange(C->cmap, &cstart, &cend));
3061: shift = cend - cstart;
3062: count = 0;
3063: colindices = NULL;
3064: if (ocolemb) PetscCall(ISGetIndices(ocolemb, &colindices));
3065: for (PetscInt i = 0; i < B->rmap->n; i++) {
3066: PetscInt row;
3067: row = i;
3068: if (rowindices) row = rowindices[i];
3069: for (PetscInt j = Baij->i[i]; j < Baij->i[i + 1]; j++) {
3070: col = Baij->j[count];
3071: if (colindices) col = colindices[col];
3072: if (Bdisassembled && col >= cstart) col += shift;
3073: v = Baij->a[count];
3074: PetscCall(MatSetValues(aij->B, 1, &row, 1, &col, &v, INSERT_VALUES));
3075: ++count;
3076: }
3077: }
3078: if (ocolemb) PetscCall(ISRestoreIndices(ocolemb, &colindices));
3079: if (rowemb) PetscCall(ISRestoreIndices(rowemb, &rowindices));
3080: /* No assembly for aij->B is necessary. */
3081: /* FIXME: set aij->B's nonzerostate correctly. */
3082: } else PetscCall(MatSetUp(aij->B));
3083: C->preallocated = PETSC_TRUE;
3084: C->was_assembled = PETSC_FALSE;
3085: C->assembled = PETSC_FALSE;
3086: /*
3087: C will need to be assembled so that aij->B can be compressed into local form in MatSetUpMultiply_MPIAIJ().
3088: Furthermore, its nonzerostate will need to be based on that of aij->A's and aij->B's.
3089: */
3090: PetscFunctionReturn(PETSC_SUCCESS);
3091: }
3093: /*
3094: B uses local indices with column indices ranging between 0 and N-n; they must be interpreted using garray.
3095: */
3096: PetscErrorCode MatGetSeqMats_MPIAIJ(Mat C, Mat *A, Mat *B)
3097: {
3098: Mat_MPIAIJ *aij = (Mat_MPIAIJ *)C->data;
3100: PetscFunctionBegin;
3101: PetscAssertPointer(A, 2);
3102: PetscAssertPointer(B, 3);
3103: /* FIXME: make sure C is assembled */
3104: *A = aij->A;
3105: *B = aij->B;
3106: /* Note that we don't incref *A and *B, so be careful! */
3107: PetscFunctionReturn(PETSC_SUCCESS);
3108: }
3110: /*
3111: Extract MPI submatrices encoded by pairs of IS that may live on subcomms of C.
3112: NOT SCALABLE due to the use of ISGetNonlocalIS() (see below).
3113: */
3114: static PetscErrorCode MatCreateSubMatricesMPI_MPIXAIJ(Mat C, PetscInt ismax, const IS isrow[], const IS iscol[], MatReuse scall, Mat *submat[], PetscErrorCode (*getsubmats_seq)(Mat, PetscInt, const IS[], const IS[], MatReuse, Mat **), PetscErrorCode (*getlocalmats)(Mat, Mat *, Mat *), PetscErrorCode (*setseqmat)(Mat, IS, IS, MatStructure, Mat), PetscErrorCode (*setseqmats)(Mat, IS, IS, IS, MatStructure, Mat, Mat))
3115: {
3116: PetscMPIInt size, flag;
3117: PetscInt cismax, ispar;
3118: Mat *A, *B;
3119: IS *isrow_p, *iscol_p, *cisrow, *ciscol, *ciscol_p;
3121: PetscFunctionBegin;
3122: if (!ismax) PetscFunctionReturn(PETSC_SUCCESS);
3124: cismax = 0;
3125: for (PetscInt i = 0; i < ismax; ++i) {
3126: PetscCallMPI(MPI_Comm_compare(((PetscObject)isrow[i])->comm, ((PetscObject)iscol[i])->comm, &flag));
3127: PetscCheck(flag == MPI_IDENT, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Row and column index sets must have the same communicator");
3128: PetscCallMPI(MPI_Comm_size(((PetscObject)isrow[i])->comm, &size));
3129: if (size > 1) ++cismax;
3130: }
3132: /*
3133: If cismax is zero on all C's ranks, then and only then can we use purely sequential matrix extraction.
3134: ispar counts the number of parallel ISs across C's comm.
3135: */
3136: PetscCallMPI(MPIU_Allreduce(&cismax, &ispar, 1, MPIU_INT, MPI_MAX, PetscObjectComm((PetscObject)C)));
3137: if (!ispar) { /* Sequential ISs only across C's comm, so can call the sequential matrix extraction subroutine. */
3138: PetscCall((*getsubmats_seq)(C, ismax, isrow, iscol, scall, submat));
3139: PetscFunctionReturn(PETSC_SUCCESS);
3140: }
3142: /* if (ispar) */
3143: /*
3144: Construct the "complements" -- the off-processor indices -- of the iscol ISs for parallel ISs only.
3145: These are used to extract the off-diag portion of the resulting parallel matrix.
3146: The row IS for the off-diag portion is the same as for the diag portion,
3147: so we merely alias (without increfing) the row IS, while skipping those that are sequential.
3148: */
3149: PetscCall(PetscMalloc2(cismax, &cisrow, cismax, &ciscol));
3150: PetscCall(PetscMalloc1(cismax, &ciscol_p));
3151: for (PetscInt i = 0, ii = 0; i < ismax; ++i) {
3152: PetscCallMPI(MPI_Comm_size(((PetscObject)isrow[i])->comm, &size));
3153: if (size > 1) {
3154: /*
3155: TODO: This is the part that's ***NOT SCALABLE***.
3156: To fix this we need to extract just the indices of C's nonzero columns
3157: that lie on the intersection of isrow[i] and ciscol[ii] -- the nonlocal
3158: part of iscol[i] -- without actually computing ciscol[ii]. This also has
3159: to be done without serializing on the IS list, so, most likely, it is best
3160: done by rewriting MatCreateSubMatrices_MPIAIJ() directly.
3161: */
3162: PetscCall(ISGetNonlocalIS(iscol[i], &ciscol[ii]));
3163: /* Now we have to
3164: (a) make sure ciscol[ii] is sorted, since, even if the off-proc indices
3165: were sorted on each rank, concatenated they might no longer be sorted;
3166: (b) Use ISSortPermutation() to construct ciscol_p, the mapping from the
3167: indices in the nondecreasing order to the original index positions.
3168: If ciscol[ii] is strictly increasing, the permutation IS is NULL.
3169: */
3170: PetscCall(ISSortPermutation(ciscol[ii], PETSC_FALSE, ciscol_p + ii));
3171: if (ciscol_p[ii] != NULL) {
3172: IS sorted;
3174: /* preserve the cached ordering for subsequent extractions */
3175: PetscCall(ISDuplicate(ciscol[ii], &sorted));
3176: PetscCall(ISRestoreNonlocalIS(iscol[i], &ciscol[ii]));
3177: ciscol[ii] = sorted;
3178: PetscCall(ISSort(ciscol[ii]));
3179: }
3180: ++ii;
3181: }
3182: }
3183: PetscCall(PetscMalloc2(ismax, &isrow_p, ismax, &iscol_p));
3184: for (PetscInt i = 0, ii = 0; i < ismax; ++i) {
3185: PetscInt issize;
3186: const PetscInt *indices;
3188: /*
3189: Permute the indices into a nondecreasing order. Reject row and col indices with duplicates.
3190: */
3191: PetscCall(ISSortPermutation(isrow[i], PETSC_FALSE, isrow_p + i));
3192: PetscCall(ISSort(isrow[i]));
3193: PetscCall(ISGetLocalSize(isrow[i], &issize));
3194: PetscCall(ISGetIndices(isrow[i], &indices));
3195: for (PetscInt j = 1; j < issize; ++j) {
3196: PetscCheck(indices[j] != indices[j - 1], PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Repeated indices in row IS %" PetscInt_FMT ": indices at %" PetscInt_FMT " and %" PetscInt_FMT " are both %" PetscInt_FMT, i, j - 1, j, indices[j]);
3197: }
3198: PetscCall(ISRestoreIndices(isrow[i], &indices));
3199: if (isrow[i] == iscol[i]) {
3200: /* Row and column are the same IS; isrow[i] is already sorted, so reuse row permutation */
3201: iscol_p[i] = isrow_p[i];
3202: PetscCall(PetscObjectReference((PetscObject)iscol_p[i]));
3203: } else {
3204: PetscCall(ISSortPermutation(iscol[i], PETSC_FALSE, iscol_p + i));
3205: PetscCall(ISSort(iscol[i]));
3206: PetscCall(ISGetLocalSize(iscol[i], &issize));
3207: PetscCall(ISGetIndices(iscol[i], &indices));
3208: for (PetscInt j = 1; j < issize; ++j) {
3209: PetscCheck(indices[j - 1] != indices[j], PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Repeated indices in col IS %" PetscInt_FMT ": indices at %" PetscInt_FMT " and %" PetscInt_FMT " are both %" PetscInt_FMT, i, j - 1, j, indices[j]);
3210: }
3211: PetscCall(ISRestoreIndices(iscol[i], &indices));
3212: }
3213: PetscCallMPI(MPI_Comm_size(((PetscObject)isrow[i])->comm, &size));
3214: if (size > 1) {
3215: cisrow[ii] = isrow[i];
3216: ++ii;
3217: }
3218: }
3219: /*
3220: Allocate the necessary arrays to hold the resulting parallel matrices as well as the intermediate
3221: array of sequential matrices underlying the resulting parallel matrices.
3222: Which arrays to allocate is based on the value of MatReuse scall and whether ISs are sorted and/or
3223: contain duplicates.
3225: There are as many diag matrices as there are original index sets. There are only as many parallel
3226: and off-diag matrices, as there are parallel (comm size > 1) index sets.
3228: ARRAYS that can hold Seq matrices get allocated in any event -- either here or by getsubmats_seq():
3229: - If the array of MPI matrices already exists and is being reused, we need to allocate the array
3230: and extract the underlying seq matrices into it to serve as placeholders, into which getsubmats_seq
3231: will deposite the extracted diag and off-diag parts. Thus, we allocate the A&B arrays and fill them
3232: with A[i] and B[ii] extracted from the corresponding MPI submat.
3233: - However, if the rows, A's column indices or B's column indices are not sorted, the extracted A[i] & B[ii]
3234: will have a different order from what getsubmats_seq expects. To handle this case -- indicated
3235: by a nonzero isrow_p[i], iscol_p[i], or ciscol_p[ii] -- we duplicate A[i] --> AA[i], B[ii] --> BB[ii]
3236: (retrieve composed AA[i] or BB[ii]) and reuse them here. AA[i] and BB[ii] are then used to permute its
3237: values into A[i] and B[ii] sitting inside the corresponding submat.
3238: - If no reuse is taking place then getsubmats_seq will allocate the A&B arrays and create the corresponding
3239: A[i], B[ii], AA[i] or BB[ii] matrices.
3240: */
3241: /* Parallel matrix array is allocated here only if no reuse is taking place. If reused, it is passed in by the caller. */
3242: if (scall == MAT_INITIAL_MATRIX) PetscCall(PetscMalloc1(ismax, submat));
3244: /* Now obtain the sequential A and B submatrices separately. */
3245: /* scall=MAT_REUSE_MATRIX is not handled yet, because getsubmats_seq() requires reuse of A and B */
3246: PetscCall((*getsubmats_seq)(C, ismax, isrow, iscol, MAT_INITIAL_MATRIX, &A));
3247: PetscCall((*getsubmats_seq)(C, cismax, cisrow, ciscol, MAT_INITIAL_MATRIX, &B));
3249: /*
3250: If scall == MAT_REUSE_MATRIX AND the permutations are NULL, we are done, since the sequential
3251: matrices A & B have been extracted directly into the parallel matrices containing them, or
3252: simply into the sequential matrix identical with the corresponding A (if size == 1).
3253: Note that in that case colmap doesn't need to be rebuilt, since the matrices are expected
3254: to have the same sparsity pattern.
3255: Otherwise, A and/or B have to be properly embedded into C's index spaces and the correct colmap
3256: must be constructed for C. This is done by setseqmat(s).
3257: */
3258: for (PetscInt i = 0, ii = 0; i < ismax; ++i) {
3259: /*
3260: TODO: cache ciscol, permutation ISs and maybe cisrow? What about isrow & iscol?
3261: That way we can avoid sorting and computing permutations when reusing.
3262: To this end:
3263: - remove the old cache, if it exists, when extracting submatrices with MAT_INITIAL_MATRIX
3264: - if caching arrays to hold the ISs, make and compose a container for them so that it can
3265: be destroyed upon destruction of C (use PetscContainerUserDestroy() to clear out the contents).
3266: */
3267: MatStructure pattern = DIFFERENT_NONZERO_PATTERN;
3269: PetscCallMPI(MPI_Comm_size(((PetscObject)isrow[i])->comm, &size));
3270: /* Construct submat[i] from the Seq pieces A (and B, if necessary). */
3271: if (size > 1) {
3272: if (scall == MAT_INITIAL_MATRIX) {
3273: MatType type;
3275: PetscCall(MatGetRootType_Private(C, &type));
3276: PetscCall(MatCreate(((PetscObject)isrow[i])->comm, (*submat) + i));
3277: PetscCall(MatSetSizes((*submat)[i], A[i]->rmap->n, A[i]->cmap->n, PETSC_DETERMINE, PETSC_DETERMINE));
3278: PetscCall(MatSetBlockSizesFromMats((*submat)[i], A[i], A[i]));
3279: PetscCall(MatSetType((*submat)[i], type));
3280: PetscCall(PetscLayoutSetUp((*submat)[i]->rmap));
3281: PetscCall(PetscLayoutSetUp((*submat)[i]->cmap));
3282: }
3283: /*
3284: For each parallel isrow[i], insert the extracted sequential matrices into the parallel matrix.
3285: */
3286: {
3287: Mat AA = A[i], BB = B[ii];
3289: if (AA || BB) {
3290: PetscBool nooffprocentries = (*submat)[i]->nooffprocentries;
3292: PetscCall(setseqmats((*submat)[i], isrow_p[i], iscol_p[i], ciscol_p[ii], pattern, AA, BB));
3293: PetscCall(MatSetOption((*submat)[i], MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE)); /* only the local diagonal and off-diagonal blocks were filled */
3294: PetscCall(MatAssemblyBegin((*submat)[i], MAT_FINAL_ASSEMBLY));
3295: PetscCall(MatAssemblyEnd((*submat)[i], MAT_FINAL_ASSEMBLY));
3296: PetscCall(MatSetOption((*submat)[i], MAT_NO_OFF_PROC_ENTRIES, nooffprocentries));
3297: }
3298: PetscCall(MatDestroy(&AA));
3299: }
3300: if (ciscol_p[ii] != NULL) PetscCall(ISDestroy(ciscol + ii));
3301: else PetscCall(ISRestoreNonlocalIS(iscol[i], ciscol + ii));
3302: PetscCall(ISDestroy(ciscol_p + ii));
3303: ++ii;
3304: } else { /* if (size == 1) */
3305: if (scall == MAT_REUSE_MATRIX) PetscCall(MatDestroy(&(*submat)[i]));
3306: if (isrow_p[i] || iscol_p[i]) {
3307: PetscCall(MatDuplicate(A[i], MAT_DO_NOT_COPY_VALUES, (*submat) + i));
3308: PetscCall(setseqmat((*submat)[i], isrow_p[i], iscol_p[i], pattern, A[i]));
3309: /* Otherwise A is extracted straight into (*submats)[i]. */
3310: /* TODO: Compose A[i] on (*submat([i] for future use, if ((isrow_p[i] || iscol_p[i]) && MAT_INITIAL_MATRIX). */
3311: PetscCall(MatDestroy(A + i));
3312: } else (*submat)[i] = A[i];
3313: }
3314: PetscCall(ISDestroy(&isrow_p[i]));
3315: PetscCall(ISDestroy(&iscol_p[i]));
3316: }
3317: PetscCall(PetscFree2(cisrow, ciscol));
3318: PetscCall(PetscFree2(isrow_p, iscol_p));
3319: PetscCall(PetscFree(ciscol_p));
3320: PetscCall(PetscFree(A));
3321: PetscCall(MatDestroySubMatrices(cismax, &B));
3322: PetscFunctionReturn(PETSC_SUCCESS);
3323: }
3325: PetscErrorCode MatCreateSubMatricesMPI_MPIAIJ(Mat C, PetscInt ismax, const IS isrow[], const IS iscol[], MatReuse scall, Mat *submat[])
3326: {
3327: PetscFunctionBegin;
3328: PetscCall(MatCreateSubMatricesMPI_MPIXAIJ(C, ismax, isrow, iscol, scall, submat, MatCreateSubMatrices_MPIAIJ, MatGetSeqMats_MPIAIJ, MatSetSeqMat_SeqAIJ, MatSetSeqMats_MPIAIJ));
3329: PetscFunctionReturn(PETSC_SUCCESS);
3330: }