Actual source code: sbaijov.c

  1: /*
  2:    Routines to compute overlapping regions of a parallel MPI matrix.
  3:    Used for finding submatrices that were shared across processors.
  4: */
  5: #include <../src/mat/impls/sbaij/mpi/mpisbaij.h>

  7: static PetscErrorCode MatIncreaseOverlap_MPISBAIJ_Once(Mat, PetscInt, IS *);
  8: static PetscErrorCode MatIncreaseOverlap_MPISBAIJ_Local(Mat, PetscInt *, PetscInt, PetscInt *, PetscBT *);

 10: PetscErrorCode MatIncreaseOverlap_MPISBAIJ(Mat C, PetscInt is_max, IS is[], PetscInt ov)
 11: {
 12:   PetscInt        i, N = C->cmap->N, bs = C->rmap->bs, nis;
 13:   IS             *is_new;
 14:   Mat             B;
 15:   const PetscInt *idx;
 16:   PetscBool       flg;

 18:   PetscFunctionBegin;
 19:   PetscCall(PetscMalloc1(is_max, &is_new));
 20:   /* Convert the indices into block format */
 21:   PetscCall(ISCompressIndicesGeneral(N, C->rmap->n, bs, is_max, is, is_new));
 22:   PetscCheck(ov >= 0, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Negative overlap specified");

 24:   /*  previous non-scalable implementation  */
 25:   flg = PETSC_FALSE;
 26:   PetscCall(PetscOptionsHasName(NULL, NULL, "-IncreaseOverlap_old", &flg));
 27:   if (flg) { /* previous non-scalable implementation */
 28:     printf("use previous non-scalable implementation...\n");
 29:     for (i = 0; i < ov; ++i) PetscCall(MatIncreaseOverlap_MPISBAIJ_Once(C, is_max, is_new));
 30:   } else if (ov && is_max) {
 31:     /* MPISBAIJ stores only the upper triangular part of the matrix. Build its full sparsity pattern once,
 32:        then use the MPIBAIJ overlap algorithm to find both row and column neighbors */
 33:     PetscCall(MatSBAIJCreateSymmetricStructure_Private(C, MATMPIBAIJ, PETSC_TRUE, &B));
 34:     for (i = 0; i < ov; ++i) PetscCall(MatIncreaseOverlap_MPIBAIJ_Once(B, is_max, is_new));
 35:     PetscCall(MatDestroy(&B));
 36:   }
 37:   for (i = 0; i < is_max; i++) {
 38:     PetscCall(ISDestroy(&is[i]));
 39:     PetscCall(ISGetLocalSize(is_new[i], &nis));
 40:     PetscCall(ISGetIndices(is_new[i], &idx));
 41:     PetscCall(ISCreateBlock(PetscObjectComm((PetscObject)is_new[i]), bs, nis, idx, PETSC_COPY_VALUES, &is[i]));
 42:     PetscCall(ISDestroy(&is_new[i]));
 43:   }
 44:   PetscCall(PetscFree(is_new));
 45:   PetscFunctionReturn(PETSC_SUCCESS);
 46: }

 48: typedef enum {
 49:   MINE,
 50:   OTHER
 51: } WhoseOwner;
 52: /*  data1, odata1 and odata2 are packed in the format (for communication):
 53:        data[0]          = is_max, no of is
 54:        data[1]          = size of is[0]
 55:         ...
 56:        data[is_max]     = size of is[is_max-1]
 57:        data[is_max + 1] = data(is[0])
 58:         ...
 59:        data[is_max+1+sum(size of is[k]), k=0,...,i-1] = data(is[i])
 60:         ...
 61:    data2 is packed in the format (for creating output is[]):
 62:        data[0]          = is_max, no of is
 63:        data[1]          = size of is[0]
 64:         ...
 65:        data[is_max]     = size of is[is_max-1]
 66:        data[is_max + 1] = data(is[0])
 67:         ...
 68:        data[is_max + 1 + Mbs*i) = data(is[i])
 69:         ...
 70: */
 71: static PetscErrorCode MatIncreaseOverlap_MPISBAIJ_Once(Mat C, PetscInt is_max, IS is[])
 72: {
 73:   Mat_MPISBAIJ   *c        = (Mat_MPISBAIJ *)C->data;
 74:   PetscMPIInt     proc_end = 0, size, rank, tag1, tag2, *len_s, nrqr, nrqs, *id_r1, *len_r1, flag, *iwork;
 75:   const PetscInt *idx_i;
 76:   PetscInt        idx, isz, col, *n, *data1, **data1_start, *data2, *data2_i, *data, *data_i, len;
 77:   PetscInt        Mbs, i, j, k, *odata1, *odata2;
 78:   PetscInt      **odata2_ptr, *ctable = NULL, *btable, len_max, len_est;
 79:   PetscInt        len_unused, nodata2;
 80:   PetscInt        ois_max; /* max no of is[] in each of processor */
 81:   PetscByte      *t_p;
 82:   MPI_Comm        comm;
 83:   MPI_Request    *s_waits1, *s_waits2, r_req;
 84:   MPI_Status     *s_status, r_status;
 85:   PetscBT        *table; /* mark indices of this processor's is[] */
 86:   PetscBT         table_i;
 87:   PetscBT         otable; /* mark indices of other processors' is[] */
 88:   PetscInt        bs = C->rmap->bs, Bn = c->B->cmap->n, Bnbs = Bn / bs, *Bowners;
 89:   IS              garray_local, garray_gl;

 91:   PetscFunctionBegin;
 92:   PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
 93:   size = c->size;
 94:   rank = c->rank;
 95:   Mbs  = c->Mbs;

 97:   PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag1));
 98:   PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag2));

100:   /* create tables used in
101:      step 1: table[i] - mark c->garray of proc [i]
102:      step 3: table[i] - mark indices of is[i] when whose=MINE
103:              table[0] - mark incideces of is[] when whose=OTHER */
104:   len = PetscMax(is_max, size);
105:   PetscCall(PetscMalloc2(len, &table, (Mbs / PETSC_BITS_PER_BYTE + 1) * len, &t_p));
106:   for (i = 0; i < len; i++) table[i] = t_p + (Mbs / PETSC_BITS_PER_BYTE + 1) * i;

108:   PetscCallMPI(MPIU_Allreduce(&is_max, &ois_max, 1, MPIU_INT, MPI_MAX, comm));

110:   /* 1. Send this processor's is[] to other processors */
111:   /* allocate spaces */
112:   PetscCall(PetscMalloc1(is_max, &n));
113:   len = 0;
114:   for (i = 0; i < is_max; i++) {
115:     PetscCall(ISGetLocalSize(is[i], &n[i]));
116:     len += n[i];
117:   }
118:   if (!len) {
119:     is_max = 0;
120:   } else {
121:     len += 1 + is_max; /* max length of data1 for one processor */
122:   }

124:   PetscCall(PetscMalloc1(size * len + 1, &data1));
125:   PetscCall(PetscMalloc1(size, &data1_start));
126:   for (i = 0; i < size; i++) data1_start[i] = data1 + i * len;

128:   PetscCall(PetscMalloc4(size, &len_s, size, &btable, size, &iwork, size + 1, &Bowners));

130:   /* gather c->garray from all processors */
131:   PetscCall(ISCreateGeneral(comm, Bnbs, c->garray, PETSC_COPY_VALUES, &garray_local));
132:   PetscCall(ISAllGather(garray_local, &garray_gl));
133:   PetscCall(ISDestroy(&garray_local));
134:   PetscCallMPI(MPI_Allgather(&Bnbs, 1, MPIU_INT, Bowners + 1, 1, MPIU_INT, comm));

136:   Bowners[0] = 0;
137:   for (i = 0; i < size; i++) Bowners[i + 1] += Bowners[i];

139:   if (is_max) {
140:     /* hash table ctable which maps c->row to proc_id) */
141:     PetscCall(PetscMalloc1(Mbs, &ctable));
142:     j = 0;
143:     for (PetscMPIInt proc_id = 0; proc_id < size; proc_id++) {
144:       for (; j < C->rmap->range[proc_id + 1] / bs; j++) ctable[j] = proc_id;
145:     }

147:     /* hash tables marking c->garray */
148:     PetscCall(ISGetIndices(garray_gl, &idx_i));
149:     for (i = 0; i < size; i++) {
150:       table_i = table[i];
151:       PetscCall(PetscBTMemzero(Mbs, table_i));
152:       for (j = Bowners[i]; j < Bowners[i + 1]; j++) { /* go through B cols of proc[i]*/
153:         PetscCall(PetscBTSet(table_i, idx_i[j]));
154:       }
155:     }
156:     PetscCall(ISRestoreIndices(garray_gl, &idx_i));
157:   } /* if (is_max) */
158:   PetscCall(ISDestroy(&garray_gl));

160:   /* evaluate communication - mesg to who, length, and buffer space */
161:   for (i = 0; i < size; i++) len_s[i] = 0;

163:   /* header of data1 */
164:   for (PetscMPIInt proc_id = 0; proc_id < size; proc_id++) {
165:     iwork[proc_id]        = 0;
166:     *data1_start[proc_id] = is_max;
167:     data1_start[proc_id]++;
168:     for (j = 0; j < is_max; j++) {
169:       if (proc_id == rank) {
170:         *data1_start[proc_id] = n[j];
171:       } else {
172:         *data1_start[proc_id] = 0;
173:       }
174:       data1_start[proc_id]++;
175:     }
176:   }

178:   for (i = 0; i < is_max; i++) {
179:     PetscCall(ISGetIndices(is[i], &idx_i));
180:     for (j = 0; j < n[i]; j++) {
181:       idx                = idx_i[j];
182:       *data1_start[rank] = idx;
183:       data1_start[rank]++; /* for local processing */
184:       PetscCall(PetscMPIIntCast(ctable[idx], &proc_end));
185:       for (PetscMPIInt proc_id = 0; proc_id <= proc_end; proc_id++) {            /* for others to process */
186:         if (proc_id == rank) continue;                                           /* done before this loop */
187:         if (proc_id < proc_end && !PetscBTLookup(table[proc_id], idx)) continue; /* no need for sending idx to [proc_id] */
188:         *data1_start[proc_id] = idx;
189:         data1_start[proc_id]++;
190:         len_s[proc_id]++;
191:       }
192:     }
193:     /* update header data */
194:     for (PetscMPIInt proc_id = 0; proc_id < size; proc_id++) {
195:       if (proc_id == rank) continue;
196:       *(data1 + proc_id * len + 1 + i) = len_s[proc_id] - iwork[proc_id];
197:       iwork[proc_id]                   = len_s[proc_id];
198:     }
199:     PetscCall(ISRestoreIndices(is[i], &idx_i));
200:   }

202:   nrqs = 0;
203:   nrqr = 0;
204:   for (i = 0; i < size; i++) {
205:     data1_start[i] = data1 + i * len;
206:     if (len_s[i]) {
207:       nrqs++;
208:       len_s[i] += 1 + is_max; /* add no. of header msg */
209:     }
210:   }

212:   for (i = 0; i < is_max; i++) PetscCall(ISDestroy(&is[i]));
213:   PetscCall(PetscFree(n));
214:   PetscCall(PetscFree(ctable));

216:   /* Determine the number of messages to expect, their lengths, from from-ids */
217:   PetscCall(PetscGatherNumberOfMessages(comm, NULL, len_s, &nrqr));
218:   PetscCall(PetscGatherMessageLengths(comm, nrqs, nrqr, len_s, &id_r1, &len_r1));

220:   /*  Now  post the sends */
221:   PetscCall(PetscMalloc2(size, &s_waits1, size, &s_waits2));
222:   k = 0;
223:   for (PetscMPIInt proc_id = 0; proc_id < size; proc_id++) { /* send data1 to processor [proc_id] */
224:     if (len_s[proc_id]) {
225:       PetscCallMPI(MPIU_Isend(data1_start[proc_id], len_s[proc_id], MPIU_INT, proc_id, tag1, comm, s_waits1 + k));
226:       k++;
227:     }
228:   }

230:   /* 2. Receive other's is[] and process. Then send back */
231:   len = 0;
232:   for (i = 0; i < nrqr; i++) {
233:     if (len_r1[i] > len) len = len_r1[i];
234:   }
235:   PetscCall(PetscFree(len_r1));
236:   PetscCall(PetscFree(id_r1));

238:   for (PetscMPIInt proc_id = 0; proc_id < size; proc_id++) len_s[proc_id] = iwork[proc_id] = 0;

240:   PetscCall(PetscMalloc1(len + 1, &odata1));
241:   PetscCall(PetscMalloc1(size, &odata2_ptr));
242:   PetscCall(PetscBTCreate(Mbs, &otable));

244:   len_max = ois_max * (Mbs + 1); /* max space storing all is[] for each receive */
245:   len_est = 2 * len_max;         /* estimated space of storing is[] for all receiving messages */
246:   PetscCall(PetscMalloc1(len_est + 1, &odata2));
247:   nodata2 = 0; /* nodata2+1: num of PetscMalloc(,&odata2_ptr[]) called */

249:   odata2_ptr[nodata2] = odata2;

251:   len_unused = len_est; /* unused space in the array odata2_ptr[nodata2]-- needs to be >= len_max  */

253:   k = 0;
254:   while (k < nrqr) {
255:     PetscMPIInt ilen;

257:     /* Receive messages */
258:     PetscCallMPI(MPI_Iprobe(MPI_ANY_SOURCE, tag1, comm, &flag, &r_status));
259:     if (flag) {
260:       PetscMPIInt proc_id;

262:       PetscCallMPI(MPI_Get_count(&r_status, MPIU_INT, &ilen));
263:       proc_id = r_status.MPI_SOURCE;
264:       PetscCallMPI(MPIU_Irecv(odata1, ilen, MPIU_INT, proc_id, r_status.MPI_TAG, comm, &r_req));
265:       PetscCallMPI(MPI_Wait(&r_req, &r_status));

267:       /*  Process messages */
268:       /*  make sure there is enough unused space in odata2 array */
269:       if (len_unused < len_max) { /* allocate more space for odata2 */
270:         PetscCall(PetscMalloc1(len_est + 1, &odata2));
271:         odata2_ptr[++nodata2] = odata2;
272:         len_unused            = len_est;
273:       }

275:       PetscCall(MatIncreaseOverlap_MPISBAIJ_Local(C, odata1, OTHER, odata2, &otable));
276:       len = 1 + odata2[0];
277:       for (i = 0; i < odata2[0]; i++) len += odata2[1 + i];

279:       /* Send messages back */
280:       PetscCallMPI(MPIU_Isend(odata2, len, MPIU_INT, proc_id, tag2, comm, s_waits2 + k));
281:       k++;
282:       odata2 += len;
283:       len_unused -= len;
284:       PetscCall(PetscMPIIntCast(len, &iwork[proc_id])); /* length of message sending back to proc_id */
285:     }
286:   }
287:   PetscCall(PetscFree(odata1));
288:   PetscCall(PetscBTDestroy(&otable));

290:   /* 3. Do local work on this processor's is[] */
291:   /* make sure there is enough unused space in odata2(=data) array */
292:   len_max = is_max * (Mbs + 1); /* max space storing all is[] for this processor */
293:   if (len_unused < len_max) {   /* allocate more space for odata2 */
294:     PetscCall(PetscMalloc1(len_est + 1, &odata2));

296:     odata2_ptr[++nodata2] = odata2;
297:   }

299:   data = odata2;
300:   PetscCall(MatIncreaseOverlap_MPISBAIJ_Local(C, data1_start[rank], MINE, data, table));
301:   PetscCall(PetscFree(data1_start));

303:   /* 4. Receive work done on other processors, then merge */
304:   /* get max number of messages that this processor expects to recv */
305:   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, iwork, size, MPI_INT, MPI_MAX, comm));
306:   PetscCall(PetscMalloc1(iwork[rank] + 1, &data2));
307:   PetscCall(PetscFree4(len_s, btable, iwork, Bowners));

309:   k = 0;
310:   while (k < nrqs) {
311:     /* Receive messages */
312:     PetscCallMPI(MPI_Iprobe(MPI_ANY_SOURCE, tag2, comm, &flag, &r_status));
313:     if (flag) {
314:       PetscMPIInt proc_id, ilen;
315:       PetscCallMPI(MPI_Get_count(&r_status, MPIU_INT, &ilen));
316:       proc_id = r_status.MPI_SOURCE;
317:       PetscCallMPI(MPIU_Irecv(data2, ilen, MPIU_INT, proc_id, r_status.MPI_TAG, comm, &r_req));
318:       PetscCallMPI(MPI_Wait(&r_req, &r_status));
319:       if (ilen > 1 + is_max) { /* Add data2 into data */
320:         data2_i = data2 + 1 + is_max;
321:         for (i = 0; i < is_max; i++) {
322:           table_i = table[i];
323:           data_i  = data + 1 + is_max + Mbs * i;
324:           isz     = data[1 + i];
325:           for (j = 0; j < data2[1 + i]; j++) {
326:             col = data2_i[j];
327:             if (!PetscBTLookupSet(table_i, col)) data_i[isz++] = col;
328:           }
329:           data[1 + i] = isz;
330:           if (i < is_max - 1) data2_i += data2[1 + i];
331:         }
332:       }
333:       k++;
334:     }
335:   }
336:   PetscCall(PetscFree(data2));
337:   PetscCall(PetscFree2(table, t_p));

339:   /* phase 1 sends are complete */
340:   PetscCall(PetscMalloc1(size, &s_status));
341:   if (nrqs) PetscCallMPI(MPI_Waitall(nrqs, s_waits1, s_status));
342:   PetscCall(PetscFree(data1));

344:   /* phase 2 sends are complete */
345:   if (nrqr) PetscCallMPI(MPI_Waitall(nrqr, s_waits2, s_status));
346:   PetscCall(PetscFree2(s_waits1, s_waits2));
347:   PetscCall(PetscFree(s_status));

349:   /* 5. Create new is[] */
350:   for (i = 0; i < is_max; i++) {
351:     data_i = data + 1 + is_max + Mbs * i;
352:     PetscCall(ISCreateGeneral(PETSC_COMM_SELF, data[1 + i], data_i, PETSC_COPY_VALUES, is + i));
353:   }
354:   for (k = 0; k <= nodata2; k++) PetscCall(PetscFree(odata2_ptr[k]));
355:   PetscCall(PetscFree(odata2_ptr));
356:   PetscFunctionReturn(PETSC_SUCCESS);
357: }

359: /*
360:    MatIncreaseOverlap_MPISBAIJ_Local - Called by MatIncreaseOverlap, to do
361:        the work on the local processor.

363:      Inputs:
364:       C      - MAT_MPISBAIJ;
365:       data   - holds is[]. See MatIncreaseOverlap_MPISBAIJ_Once() for the format.
366:       whose  - whose is[] to be processed,
367:                MINE:  this processor's is[]
368:                OTHER: other processor's is[]
369:      Output:
370:        nidx  - whose = MINE:
371:                      holds input and newly found indices in the same format as data
372:                whose = OTHER:
373:                      only holds the newly found indices
374:        table - table[i]: mark the indices of is[i], i=0,...,is_max. Used only in the case 'whose=MINE'.
375: */
376: /* Would computation be reduced by swapping the loop 'for each is' and 'for each row'? */
377: static PetscErrorCode MatIncreaseOverlap_MPISBAIJ_Local(Mat C, PetscInt *data, PetscInt whose, PetscInt *nidx, PetscBT *table)
378: {
379:   Mat_MPISBAIJ *c = (Mat_MPISBAIJ *)C->data;
380:   Mat_SeqSBAIJ *a = (Mat_SeqSBAIJ *)c->A->data;
381:   Mat_SeqBAIJ  *b = (Mat_SeqBAIJ *)c->B->data;
382:   PetscInt      row, mbs, Mbs, *nidx_i, col, col_max, isz, isz0, *ai, *aj, *bi, *bj, *garray, rstart, l;
383:   PetscInt      a_start, a_end, b_start, b_end, i, j, k, is_max, *idx_i, n;
384:   PetscBT       table0;  /* mark the indices of input is[] for look up */
385:   PetscBT       table_i; /* points to i-th table. When whose=OTHER, a single table is used for all is[] */

387:   PetscFunctionBegin;
388:   Mbs    = c->Mbs;
389:   mbs    = a->mbs;
390:   ai     = a->i;
391:   aj     = a->j;
392:   bi     = b->i;
393:   bj     = b->j;
394:   garray = c->garray;
395:   rstart = c->rstartbs;
396:   is_max = data[0];

398:   PetscCall(PetscBTCreate(Mbs, &table0));

400:   nidx[0] = is_max;
401:   idx_i   = data + is_max + 1;   /* ptr to input is[0] array */
402:   nidx_i  = nidx + is_max + 1;   /* ptr to output is[0] array */
403:   for (i = 0; i < is_max; i++) { /* for each is */
404:     isz = 0;
405:     n   = data[1 + i]; /* size of input is[i] */

407:     /* initialize and set table_i(mark idx and nidx) and table0(only mark idx) */
408:     if (whose == MINE) { /* process this processor's is[] */
409:       table_i = table[i];
410:       nidx_i  = nidx + 1 + is_max + Mbs * i;
411:     } else { /* process other processor's is[] - only use one temp table */
412:       table_i = table[0];
413:     }
414:     PetscCall(PetscBTMemzero(Mbs, table_i));
415:     PetscCall(PetscBTMemzero(Mbs, table0));
416:     if (n == 0) {
417:       nidx[1 + i] = 0; /* size of new is[i] */
418:       continue;
419:     }

421:     isz0    = 0;
422:     col_max = 0;
423:     for (j = 0; j < n; j++) {
424:       col = idx_i[j];
425:       PetscCheck(col < Mbs, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "index col %" PetscInt_FMT " >= Mbs %" PetscInt_FMT, col, Mbs);
426:       if (!PetscBTLookupSet(table_i, col)) {
427:         PetscCall(PetscBTSet(table0, col));
428:         if (whose == MINE) nidx_i[isz0] = col;
429:         if (col_max < col) col_max = col;
430:         isz0++;
431:       }
432:     }

434:     if (whose == MINE) isz = isz0;
435:     k = 0; /* no. of indices from input is[i] that have been examined */
436:     for (row = 0; row < mbs; row++) {
437:       a_start = ai[row];
438:       a_end   = ai[row + 1];
439:       b_start = bi[row];
440:       b_end   = bi[row + 1];
441:       if (PetscBTLookup(table0, row + rstart)) { /* row is on input is[i]:
442:                                                 do row search: collect all col in this row */
443:         for (l = a_start; l < a_end; l++) {      /* Amat */
444:           col = aj[l] + rstart;
445:           if (!PetscBTLookupSet(table_i, col)) nidx_i[isz++] = col;
446:         }
447:         for (l = b_start; l < b_end; l++) { /* Bmat */
448:           col = garray[bj[l]];
449:           if (!PetscBTLookupSet(table_i, col)) nidx_i[isz++] = col;
450:         }
451:         k++;
452:         if (k >= isz0) break;               /* for (row=0; row<mbs; row++) */
453:       } else {                              /* row is not on input is[i]:
454:                   do col search: add row onto nidx_i if there is a col in nidx_i */
455:         for (l = a_start; l < a_end; l++) { /* Amat */
456:           col = aj[l] + rstart;
457:           if (col > col_max) break;
458:           if (PetscBTLookup(table0, col)) {
459:             if (!PetscBTLookupSet(table_i, row + rstart)) nidx_i[isz++] = row + rstart;
460:             break; /* for l = start; l<end ; l++) */
461:           }
462:         }
463:         for (l = b_start; l < b_end; l++) { /* Bmat */
464:           col = garray[bj[l]];
465:           if (col > col_max) break;
466:           if (PetscBTLookup(table0, col)) {
467:             if (!PetscBTLookupSet(table_i, row + rstart)) nidx_i[isz++] = row + rstart;
468:             break; /* for l = start; l<end ; l++) */
469:           }
470:         }
471:       }
472:     }

474:     if (i < is_max - 1) {
475:       idx_i += n;    /* ptr to input is[i+1] array */
476:       nidx_i += isz; /* ptr to output is[i+1] array */
477:     }
478:     nidx[1 + i] = isz; /* size of new is[i] */
479:   } /* for each is */
480:   PetscCall(PetscBTDestroy(&table0));
481:   PetscFunctionReturn(PETSC_SUCCESS);
482: }