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