Date: 2026-07-30T21:14:14+00:00

Actual source code: baijov.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/baij/mpi/mpibaij.h>
      6: #include <petscbt.h>

      8: static PetscErrorCode MatIncreaseOverlap_MPIBAIJ_Local(Mat, PetscInt, PetscBT *, PetscInt *, PetscInt **);
      9: static PetscErrorCode MatIncreaseOverlap_MPIBAIJ_Receive(Mat, PetscInt, PetscInt **, PetscInt **, PetscInt *);

     11: PetscErrorCode MatIncreaseOverlap_MPIBAIJ(Mat C, PetscInt imax, IS is[], PetscInt ov)
     12: {
     13:   PetscInt        i, N = C->cmap->N, bs = C->rmap->bs, n;
     14:   const PetscInt *idx;
     15:   IS             *is_new;

     17:   PetscFunctionBegin;
     18:   PetscCall(PetscMalloc1(imax, &is_new));
     19:   /* Convert the indices into block format */
     20:   PetscCall(ISCompressIndicesGeneral(N, C->rmap->n, bs, imax, is, is_new));
     21:   PetscCheck(ov >= 0, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Negative overlap specified");
     22:   for (i = 0; i < ov; ++i) PetscCall(MatIncreaseOverlap_MPIBAIJ_Once(C, imax, is_new));
     23:   for (i = 0; i < imax; i++) {
     24:     PetscCall(ISDestroy(&is[i]));
     25:     PetscCall(ISGetLocalSize(is_new[i], &n));
     26:     PetscCall(ISGetIndices(is_new[i], &idx));
     27:     PetscCall(ISCreateBlock(PetscObjectComm((PetscObject)is_new[i]), bs, n, idx, PETSC_COPY_VALUES, &is[i]));
     28:     PetscCall(ISDestroy(&is_new[i]));
     29:   }
     30:   PetscCall(PetscFree(is_new));
     31:   PetscFunctionReturn(PETSC_SUCCESS);
     32: }

     34: /*
     35:   Sample message format:
     36:   If a processor A wants processor B to process some elements corresponding
     37:   to index sets is[1], is[5]
     38:   mesg [0] = 2   (no of index sets in the mesg)
     39:   -----------
     40:   mesg [1] = 1 => is[1]
     41:   mesg [2] = sizeof(is[1]);
     42:   -----------
     43:   mesg [5] = 5  => is[5]
     44:   mesg [6] = sizeof(is[5]);
     45:   -----------
     46:   mesg [7]
     47:   mesg [n]  data(is[1])
     48:   -----------
     49:   mesg[n+1]
     50:   mesg[m]  data(is[5])
     51:   -----------

     53:   Notes:
     54:   nrqs - no of requests sent (or to be sent out)
     55:   nrqr - no of requests received (which have to be or which have been processed)
     56: */
     57: PetscErrorCode MatIncreaseOverlap_MPIBAIJ_Once(Mat C, PetscInt imax, IS is[])
     58: {
     59:   Mat_MPIBAIJ     *c = (Mat_MPIBAIJ *)C->data;
     60:   const PetscInt **idx, *idx_i;
     61:   PetscInt        *n, *w3, *w4, **data, len;
     62:   PetscMPIInt      size, rank, tag1, tag2, *w2, *w1, nrqs, nrqr, *pa;
     63:   PetscInt         Mbs, **rbuf, row, msz, **outdat, **ptr;
     64:   PetscInt        *ctr, *tmp, *isz, *isz1, **xdata, **rbuf2, *d_p;
     65:   PetscMPIInt     *onodes1, *olengths1, *onodes2, *olengths2, proc = -1;
     66:   PetscBT         *table;
     67:   MPI_Comm         comm, *iscomms;
     68:   MPI_Request     *s_waits1, *r_waits1, *s_waits2, *r_waits2;
     69:   PetscByte       *t_p;

     71:   PetscFunctionBegin;
     72:   PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
     73:   size = c->size;
     74:   rank = c->rank;
     75:   Mbs  = c->Mbs;

     77:   PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag1));
     78:   PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag2));

     80:   PetscCall(PetscMalloc2(imax, (PetscInt ***)&idx, imax, &n));

     82:   for (PetscInt i = 0; i < imax; i++) {
     83:     PetscCall(ISGetIndices(is[i], &idx[i]));
     84:     PetscCall(ISGetLocalSize(is[i], &n[i]));
     85:   }

     87:   /* evaluate communication - mesg to who,length of mesg, and buffer space
     88:      required. Based on this, buffers are allocated, and data copied into them*/
     89:   PetscCall(PetscCalloc4(size, &w1, size, &w2, size, &w3, size, &w4));
     90:   for (PetscInt i = 0; i < imax; i++) {
     91:     PetscCall(PetscArrayzero(w4, size)); /* initialise work vector*/
     92:     idx_i = idx[i];
     93:     len   = n[i];
     94:     for (PetscInt j = 0; j < len; j++) {
     95:       row = idx_i[j];
     96:       PetscCheck(row >= 0, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Index set cannot have negative entries");
     97:       PetscCall(PetscLayoutFindOwner(C->rmap, row * C->rmap->bs, &proc));
     98:       w4[proc]++;
     99:     }
    100:     for (PetscMPIInt j = 0; j < size; j++) {
    101:       if (w4[j]) {
    102:         w1[j] += w4[j];
    103:         w3[j]++;
    104:       }
    105:     }
    106:   }

    108:   nrqs     = 0; /* no of outgoing messages */
    109:   msz      = 0; /* total mesg length (for all proc */
    110:   w1[rank] = 0; /* no mesg sent to itself */
    111:   w3[rank] = 0;
    112:   for (PetscMPIInt i = 0; i < size; i++) {
    113:     if (w1[i]) {
    114:       w2[i] = 1;
    115:       nrqs++;
    116:     } /* there exists a message to proc i */
    117:   }
    118:   /* pa - is list of processors to communicate with */
    119:   PetscCall(PetscMalloc1(nrqs, &pa));
    120:   for (PetscMPIInt i = 0, j = 0; i < size; i++) {
    121:     if (w1[i]) {
    122:       pa[j] = i;
    123:       j++;
    124:     }
    125:   }

    127:   /* Each message would have a header = 1 + 2*(no of IS) + data */
    128:   for (PetscMPIInt i = 0; i < nrqs; i++) {
    129:     PetscMPIInt j = pa[i];
    130:     w1[j] += w2[j] + 2 * w3[j];
    131:     msz += w1[j];
    132:   }

    134:   /* Determine the number of messages to expect, their lengths, from from-ids */
    135:   PetscCall(PetscGatherNumberOfMessages(comm, w2, w1, &nrqr));
    136:   PetscCall(PetscGatherMessageLengths(comm, nrqs, nrqr, w1, &onodes1, &olengths1));

    138:   /* Now post the Irecvs corresponding to these messages */
    139:   PetscCall(PetscPostIrecvInt(comm, tag1, nrqr, onodes1, olengths1, &rbuf, &r_waits1));

    141:   /* Allocate Memory for outgoing messages */
    142:   PetscCall(PetscMalloc4(size, &outdat, size, &ptr, msz, &tmp, size, &ctr));
    143:   PetscCall(PetscArrayzero(outdat, size));
    144:   PetscCall(PetscArrayzero(ptr, size));
    145:   {
    146:     PetscInt *iptr = tmp, ict = 0;
    147:     for (PetscMPIInt i = 0; i < nrqs; i++) {
    148:       PetscMPIInt j = pa[i];
    149:       iptr += ict;
    150:       outdat[j] = iptr;
    151:       ict       = w1[j];
    152:     }
    153:   }

    155:   /* Form the outgoing messages */
    156:   /*plug in the headers*/
    157:   for (PetscMPIInt i = 0; i < nrqs; i++) {
    158:     PetscMPIInt j = pa[i];
    159:     outdat[j][0]  = 0;
    160:     PetscCall(PetscArrayzero(outdat[j] + 1, 2 * w3[j]));
    161:     ptr[j] = outdat[j] + 2 * w3[j] + 1;
    162:   }

    164:   /* Memory for doing local proc's work*/
    165:   {
    166:     PetscCall(PetscCalloc5(imax, &table, imax, &data, imax, &isz, Mbs * imax, &d_p, (Mbs / PETSC_BITS_PER_BYTE + 1) * imax, &t_p));

    168:     for (PetscInt i = 0; i < imax; i++) {
    169:       table[i] = t_p + (Mbs / PETSC_BITS_PER_BYTE + 1) * i;
    170:       data[i]  = d_p + (Mbs)*i;
    171:     }
    172:   }

    174:   /* Parse the IS and update local tables and the outgoing buf with the data*/
    175:   {
    176:     PetscInt n_i, *data_i, isz_i, *outdat_j, ctr_j, k;
    177:     PetscBT  table_i;

    179:     for (PetscInt i = 0; i < imax; i++) {
    180:       PetscCall(PetscArrayzero(ctr, size));
    181:       n_i     = n[i];
    182:       table_i = table[i];
    183:       idx_i   = idx[i];
    184:       data_i  = data[i];
    185:       isz_i   = isz[i];
    186:       for (PetscInt j = 0; j < n_i; j++) { /* parse the indices of each IS */
    187:         row = idx_i[j];
    188:         PetscCall(PetscLayoutFindOwner(C->rmap, row * C->rmap->bs, &proc));
    189:         if (proc != rank) { /* copy to the outgoing buffer */
    190:           ctr[proc]++;
    191:           *ptr[proc] = row;
    192:           ptr[proc]++;
    193:         } else { /* Update the local table */
    194:           if (!PetscBTLookupSet(table_i, row)) data_i[isz_i++] = row;
    195:         }
    196:       }
    197:       /* Update the headers for the current IS */
    198:       for (PetscMPIInt j = 0; j < size; j++) { /* Can Optimise this loop by using pa[] */
    199:         if ((ctr_j = ctr[j])) {
    200:           outdat_j            = outdat[j];
    201:           k                   = ++outdat_j[0];
    202:           outdat_j[2 * k]     = ctr_j;
    203:           outdat_j[2 * k - 1] = i;
    204:         }
    205:       }
    206:       isz[i] = isz_i;
    207:     }
    208:   }

    210:   /*  Now  post the sends */
    211:   PetscCall(PetscMalloc1(nrqs, &s_waits1));
    212:   for (PetscMPIInt i = 0; i < nrqs; ++i) {
    213:     PetscMPIInt j = pa[i];
    214:     PetscCallMPI(MPIU_Isend(outdat[j], w1[j], MPIU_INT, j, tag1, comm, s_waits1 + i));
    215:   }

    217:   /* No longer need the original indices*/
    218:   for (PetscInt i = 0; i < imax; ++i) PetscCall(ISRestoreIndices(is[i], idx + i));
    219:   PetscCall(PetscFree2(*(PetscInt ***)&idx, n));

    221:   PetscCall(PetscMalloc1(imax, &iscomms));
    222:   for (PetscInt i = 0; i < imax; ++i) {
    223:     PetscCall(PetscCommDuplicate(PetscObjectComm((PetscObject)is[i]), &iscomms[i], NULL));
    224:     PetscCall(ISDestroy(&is[i]));
    225:   }

    227:   /* Do Local work*/
    228:   PetscCall(MatIncreaseOverlap_MPIBAIJ_Local(C, imax, table, isz, data));

    230:   /* Receive messages*/
    231:   PetscCallMPI(MPI_Waitall(nrqr, r_waits1, MPI_STATUSES_IGNORE));
    232:   PetscCallMPI(MPI_Waitall(nrqs, s_waits1, MPI_STATUSES_IGNORE));

    234:   /* Phase 1 sends are complete - deallocate buffers */
    235:   PetscCall(PetscFree4(outdat, ptr, tmp, ctr));
    236:   PetscCall(PetscFree4(w1, w2, w3, w4));

    238:   PetscCall(PetscMalloc1(nrqr, &xdata));
    239:   PetscCall(PetscMalloc1(nrqr, &isz1));
    240:   PetscCall(MatIncreaseOverlap_MPIBAIJ_Receive(C, nrqr, rbuf, xdata, isz1));
    241:   if (rbuf) {
    242:     PetscCall(PetscFree(rbuf[0]));
    243:     PetscCall(PetscFree(rbuf));
    244:   }

    246:   /* Send the data back*/
    247:   /* Do a global reduction to know the buffer space req for incoming messages*/
    248:   {
    249:     PetscMPIInt *rw1;

    251:     PetscCall(PetscCalloc1(size, &rw1));
    252:     for (PetscMPIInt i = 0; i < nrqr; ++i) {
    253:       proc = onodes1[i];
    254:       PetscCall(PetscMPIIntCast(isz1[i], &rw1[proc]));
    255:     }

    257:     /* Determine the number of messages to expect, their lengths, from from-ids */
    258:     PetscCall(PetscGatherMessageLengths(comm, nrqr, nrqs, rw1, &onodes2, &olengths2));
    259:     PetscCall(PetscFree(rw1));
    260:   }
    261:   /* Now post the Irecvs corresponding to these messages */
    262:   PetscCall(PetscPostIrecvInt(comm, tag2, nrqs, onodes2, olengths2, &rbuf2, &r_waits2));

    264:   /*  Now  post the sends */
    265:   PetscCall(PetscMalloc1(nrqr, &s_waits2));
    266:   for (PetscMPIInt i = 0; i < nrqr; ++i) {
    267:     PetscMPIInt j = onodes1[i];
    268:     PetscCallMPI(MPIU_Isend(xdata[i], isz1[i], MPIU_INT, j, tag2, comm, s_waits2 + i));
    269:   }

    271:   PetscCall(PetscFree(onodes1));
    272:   PetscCall(PetscFree(olengths1));

    274:   /* receive work done on other processors*/
    275:   {
    276:     PetscMPIInt idex;
    277:     PetscInt    is_no, ct1, max, *rbuf2_i, isz_i, *data_i, jmax;
    278:     PetscBT     table_i;

    280:     for (PetscMPIInt i = 0; i < nrqs; ++i) {
    281:       PetscCallMPI(MPI_Waitany(nrqs, r_waits2, &idex, MPI_STATUS_IGNORE));
    282:       /* Process the message*/
    283:       rbuf2_i = rbuf2[idex];
    284:       ct1     = 2 * rbuf2_i[0] + 1;
    285:       jmax    = rbuf2[idex][0];
    286:       for (PetscInt j = 1; j <= jmax; j++) {
    287:         max     = rbuf2_i[2 * j];
    288:         is_no   = rbuf2_i[2 * j - 1];
    289:         isz_i   = isz[is_no];
    290:         data_i  = data[is_no];
    291:         table_i = table[is_no];
    292:         for (PetscInt k = 0; k < max; k++, ct1++) {
    293:           row = rbuf2_i[ct1];
    294:           if (!PetscBTLookupSet(table_i, row)) data_i[isz_i++] = row;
    295:         }
    296:         isz[is_no] = isz_i;
    297:       }
    298:     }
    299:     PetscCallMPI(MPI_Waitall(nrqr, s_waits2, MPI_STATUSES_IGNORE));
    300:   }

    302:   for (PetscInt i = 0; i < imax; ++i) {
    303:     PetscCall(ISCreateGeneral(iscomms[i], isz[i], data[i], PETSC_COPY_VALUES, is + i));
    304:     PetscCall(PetscCommDestroy(&iscomms[i]));
    305:   }

    307:   PetscCall(PetscFree(iscomms));
    308:   PetscCall(PetscFree(onodes2));
    309:   PetscCall(PetscFree(olengths2));

    311:   PetscCall(PetscFree(pa));
    312:   if (rbuf2) {
    313:     PetscCall(PetscFree(rbuf2[0]));
    314:     PetscCall(PetscFree(rbuf2));
    315:   }
    316:   PetscCall(PetscFree(s_waits1));
    317:   PetscCall(PetscFree(r_waits1));
    318:   PetscCall(PetscFree(s_waits2));
    319:   PetscCall(PetscFree(r_waits2));
    320:   PetscCall(PetscFree5(table, data, isz, d_p, t_p));
    321:   if (xdata) {
    322:     PetscCall(PetscFree(xdata[0]));
    323:     PetscCall(PetscFree(xdata));
    324:   }
    325:   PetscCall(PetscFree(isz1));
    326:   PetscFunctionReturn(PETSC_SUCCESS);
    327: }

    329: /*
    330:    MatIncreaseOverlap_MPIBAIJ_Local - Called by MatincreaseOverlap, to do
    331:        the work on the local processor.

    333:      Inputs:
    334:       C      - MAT_MPIBAIJ;
    335:       imax - total no of index sets processed at a time;
    336:       table  - an array of char - size = Mbs bits.

    338:      Output:
    339:       isz    - array containing the count of the solution elements corresponding
    340:                to each index set;
    341:       data   - pointer to the solutions
    342: */
    343: static PetscErrorCode MatIncreaseOverlap_MPIBAIJ_Local(Mat C, PetscInt imax, PetscBT *table, PetscInt *isz, PetscInt **data)
    344: {
    345:   Mat_MPIBAIJ *c = (Mat_MPIBAIJ *)C->data;
    346:   Mat          A = c->A, B = c->B;
    347:   Mat_SeqBAIJ *a = (Mat_SeqBAIJ *)A->data, *b = (Mat_SeqBAIJ *)B->data;
    348:   PetscInt     start, end, val, max, rstart, cstart, *ai, *aj;
    349:   PetscInt    *bi, *bj, *garray, i, j, k, row, *data_i, isz_i;
    350:   PetscBT      table_i;

    352:   PetscFunctionBegin;
    353:   rstart = c->rstartbs;
    354:   cstart = c->cstartbs;
    355:   ai     = a->i;
    356:   aj     = a->j;
    357:   bi     = b->i;
    358:   bj     = b->j;
    359:   garray = c->garray;

    361:   for (i = 0; i < imax; i++) {
    362:     data_i  = data[i];
    363:     table_i = table[i];
    364:     isz_i   = isz[i];
    365:     for (j = 0, max = isz[i]; j < max; j++) {
    366:       row   = data_i[j] - rstart;
    367:       start = ai[row];
    368:       end   = ai[row + 1];
    369:       for (k = start; k < end; k++) { /* Amat */
    370:         val = aj[k] + cstart;
    371:         if (!PetscBTLookupSet(table_i, val)) data_i[isz_i++] = val;
    372:       }
    373:       start = bi[row];
    374:       end   = bi[row + 1];
    375:       for (k = start; k < end; k++) { /* Bmat */
    376:         val = garray[bj[k]];
    377:         if (!PetscBTLookupSet(table_i, val)) data_i[isz_i++] = val;
    378:       }
    379:     }
    380:     isz[i] = isz_i;
    381:   }
    382:   PetscFunctionReturn(PETSC_SUCCESS);
    383: }

    385: /*
    386:       MatIncreaseOverlap_MPIBAIJ_Receive - Process the received messages,
    387:          and return the output

    389:          Input:
    390:            C    - the matrix
    391:            nrqr - no of messages being processed.
    392:            rbuf - an array of pointers to the received requests

    394:          Output:
    395:            xdata - array of messages to be sent back
    396:            isz1  - size of each message

    398:   For better efficiency perhaps we should malloc separately each xdata[i],
    399: then if a remalloc is required we need only copy the data for that one row
    400: rather than all previous rows as it is now where a single large chunk of
    401: memory is used.

    403: */
    404: static PetscErrorCode MatIncreaseOverlap_MPIBAIJ_Receive(Mat C, PetscInt nrqr, PetscInt **rbuf, PetscInt **xdata, PetscInt *isz1)
    405: {
    406:   Mat_MPIBAIJ *c = (Mat_MPIBAIJ *)C->data;
    407:   Mat          A = c->A, B = c->B;
    408:   Mat_SeqBAIJ *a = (Mat_SeqBAIJ *)A->data, *b = (Mat_SeqBAIJ *)B->data;
    409:   PetscInt     rstart, cstart, *ai, *aj, *bi, *bj, *garray, i, j, k;
    410:   PetscInt     row, total_sz, ct, ct1, ct2, ct3, mem_estimate, oct2, l, start, end;
    411:   PetscInt     val, max1, max2, Mbs, no_malloc = 0, *tmp, new_estimate, ctr;
    412:   PetscInt    *rbuf_i, kmax, rbuf_0;
    413:   PetscBT      xtable;

    415:   PetscFunctionBegin;
    416:   Mbs    = c->Mbs;
    417:   rstart = c->rstartbs;
    418:   cstart = c->cstartbs;
    419:   ai     = a->i;
    420:   aj     = a->j;
    421:   bi     = b->i;
    422:   bj     = b->j;
    423:   garray = c->garray;

    425:   for (i = 0, ct = 0, total_sz = 0; i < nrqr; ++i) {
    426:     rbuf_i = rbuf[i];
    427:     rbuf_0 = rbuf_i[0];
    428:     ct += rbuf_0;
    429:     for (j = 1; j <= rbuf_0; j++) total_sz += rbuf_i[2 * j];
    430:   }

    432:   if (c->Mbs) max1 = ct * (a->nz + b->nz) / c->Mbs;
    433:   else max1 = 1;
    434:   mem_estimate = 3 * ((total_sz > max1 ? total_sz : max1) + 1);
    435:   if (nrqr) {
    436:     PetscCall(PetscMalloc1(mem_estimate, &xdata[0]));
    437:     ++no_malloc;
    438:   }
    439:   PetscCall(PetscBTCreate(Mbs, &xtable));
    440:   PetscCall(PetscArrayzero(isz1, nrqr));

    442:   ct3 = 0;
    443:   for (i = 0; i < nrqr; i++) { /* for easch mesg from proc i */
    444:     rbuf_i = rbuf[i];
    445:     rbuf_0 = rbuf_i[0];
    446:     ct1    = 2 * rbuf_0 + 1;
    447:     ct2    = ct1;
    448:     ct3 += ct1;
    449:     for (j = 1; j <= rbuf_0; j++) { /* for each IS from proc i*/
    450:       PetscCall(PetscBTMemzero(Mbs, xtable));
    451:       oct2 = ct2;
    452:       kmax = rbuf_i[2 * j];
    453:       for (k = 0; k < kmax; k++, ct1++) {
    454:         row = rbuf_i[ct1];
    455:         if (!PetscBTLookupSet(xtable, row)) {
    456:           if (!(ct3 < mem_estimate)) {
    457:             new_estimate = (PetscInt)(1.5 * mem_estimate) + 1;
    458:             PetscCall(PetscMalloc1(new_estimate, &tmp));
    459:             PetscCall(PetscArraycpy(tmp, xdata[0], mem_estimate));
    460:             PetscCall(PetscFree(xdata[0]));
    461:             xdata[0]     = tmp;
    462:             mem_estimate = new_estimate;
    463:             ++no_malloc;
    464:             for (ctr = 1; ctr <= i; ctr++) xdata[ctr] = xdata[ctr - 1] + isz1[ctr - 1];
    465:           }
    466:           xdata[i][ct2++] = row;
    467:           ct3++;
    468:         }
    469:       }
    470:       for (k = oct2, max2 = ct2; k < max2; k++) {
    471:         row   = xdata[i][k] - rstart;
    472:         start = ai[row];
    473:         end   = ai[row + 1];
    474:         for (l = start; l < end; l++) {
    475:           val = aj[l] + cstart;
    476:           if (!PetscBTLookupSet(xtable, val)) {
    477:             if (!(ct3 < mem_estimate)) {
    478:               new_estimate = (PetscInt)(1.5 * mem_estimate) + 1;
    479:               PetscCall(PetscMalloc1(new_estimate, &tmp));
    480:               PetscCall(PetscArraycpy(tmp, xdata[0], mem_estimate));
    481:               PetscCall(PetscFree(xdata[0]));
    482:               xdata[0]     = tmp;
    483:               mem_estimate = new_estimate;
    484:               ++no_malloc;
    485:               for (ctr = 1; ctr <= i; ctr++) xdata[ctr] = xdata[ctr - 1] + isz1[ctr - 1];
    486:             }
    487:             xdata[i][ct2++] = val;
    488:             ct3++;
    489:           }
    490:         }
    491:         start = bi[row];
    492:         end   = bi[row + 1];
    493:         for (l = start; l < end; l++) {
    494:           val = garray[bj[l]];
    495:           if (!PetscBTLookupSet(xtable, val)) {
    496:             if (!(ct3 < mem_estimate)) {
    497:               new_estimate = (PetscInt)(1.5 * mem_estimate) + 1;
    498:               PetscCall(PetscMalloc1(new_estimate, &tmp));
    499:               PetscCall(PetscArraycpy(tmp, xdata[0], mem_estimate));
    500:               PetscCall(PetscFree(xdata[0]));
    501:               xdata[0]     = tmp;
    502:               mem_estimate = new_estimate;
    503:               ++no_malloc;
    504:               for (ctr = 1; ctr <= i; ctr++) xdata[ctr] = xdata[ctr - 1] + isz1[ctr - 1];
    505:             }
    506:             xdata[i][ct2++] = val;
    507:             ct3++;
    508:           }
    509:         }
    510:       }
    511:       /* Update the header*/
    512:       xdata[i][2 * j]     = ct2 - oct2; /* Undo the vector isz1 and use only a var*/
    513:       xdata[i][2 * j - 1] = rbuf_i[2 * j - 1];
    514:     }
    515:     xdata[i][0] = rbuf_0;
    516:     if (i + 1 < nrqr) xdata[i + 1] = xdata[i] + ct2;
    517:     isz1[i] = ct2; /* size of each message */
    518:   }
    519:   PetscCall(PetscBTDestroy(&xtable));
    520:   PetscCall(PetscInfo(C, "Allocated %" PetscInt_FMT " bytes, required %" PetscInt_FMT ", no of mallocs = %" PetscInt_FMT "\n", mem_estimate, ct3, no_malloc));
    521:   PetscFunctionReturn(PETSC_SUCCESS);
    522: }

    524: PetscErrorCode MatCreateSubMatrices_MPIBAIJ(Mat C, PetscInt ismax, const IS isrow[], const IS iscol[], MatReuse scall, Mat *submat[])
    525: {
    526:   IS          *isrow_block, *iscol_block;
    527:   Mat_MPIBAIJ *c = (Mat_MPIBAIJ *)C->data;
    528:   PetscInt     nmax, nstages, i, pos, max_no, N = C->cmap->N, bs = C->rmap->bs;
    529:   Mat_SeqBAIJ *subc;
    530:   Mat_SubSppt *smat;
    531:   PetscBool    sym = PETSC_FALSE, flg[2];

    533:   PetscFunctionBegin;
    534:   PetscCall(PetscObjectTypeCompare((PetscObject)C, MATMPISBAIJ, flg));
    535:   if (flg[0]) {
    536:     if (isrow == iscol) sym = PETSC_TRUE;
    537:     else {
    538:       flg[0] = flg[1] = PETSC_TRUE;
    539:       for (i = 0; i < ismax; i++) {
    540:         if (isrow[i] != iscol[i]) flg[0] = PETSC_FALSE;
    541:         PetscCall(ISGetLocalSize(iscol[0], &nmax));
    542:         if (nmax == C->cmap->N && flg[1]) PetscCall(ISIdentity(iscol[0], flg + 1));
    543:       }
    544:       sym = (PetscBool)(flg[0] || flg[1]);
    545:     }
    546:   }
    547:   /* The compression and expansion should be avoided. Doesn't point
    548:      out errors, might change the indices, hence buggey */
    549:   PetscCall(PetscMalloc2(ismax, &isrow_block, ismax, &iscol_block));
    550:   PetscCall(ISCompressIndicesGeneral(C->rmap->N, C->rmap->n, bs, ismax, isrow, isrow_block));
    551:   if (isrow == iscol) {
    552:     for (i = 0; i < ismax; i++) {
    553:       iscol_block[i] = isrow_block[i];
    554:       PetscCall(PetscObjectReference((PetscObject)iscol_block[i]));
    555:     }
    556:   } else PetscCall(ISCompressIndicesGeneral(N, C->cmap->n, bs, ismax, iscol, iscol_block));

    558:   /* Determine the number of stages through which submatrices are done */
    559:   if (!C->cmap->N) nmax = 20 * 1000000 / sizeof(PetscInt);
    560:   else nmax = 20 * 1000000 / (c->Nbs * sizeof(PetscInt));
    561:   if (!nmax) nmax = 1;

    563:   if (scall == MAT_INITIAL_MATRIX) {
    564:     nstages = ismax / nmax + ((ismax % nmax) ? 1 : 0); /* local nstages */

    566:     /* Make sure every processor loops through the nstages */
    567:     PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &nstages, 1, MPIU_INT, MPI_MAX, PetscObjectComm((PetscObject)C)));

    569:     /* Allocate memory to hold all the submatrices and dummy submatrices */
    570:     PetscCall(PetscCalloc1(ismax + nstages, submat));
    571:   } else { /* MAT_REUSE_MATRIX */
    572:     if (ismax) {
    573:       subc = (Mat_SeqBAIJ *)((*submat)[0]->data);
    574:       smat = subc->submatis1;
    575:     } else { /* (*submat)[0] is a dummy matrix */
    576:       smat = (Mat_SubSppt *)(*submat)[0]->data;
    577:     }
    578:     PetscCheck(smat, PETSC_COMM_SELF, PETSC_ERR_ARG_NULL, "MatCreateSubMatrices(...,MAT_REUSE_MATRIX,...) requires submat");
    579:     nstages = smat->nstages;
    580:   }

    582:   for (i = 0, pos = 0; i < nstages; i++) {
    583:     if (pos + nmax <= ismax) max_no = nmax;
    584:     else if (pos >= ismax) max_no = 0;
    585:     else max_no = ismax - pos;

    587:     PetscCall(MatCreateSubMatrices_MPIBAIJ_local(C, max_no, isrow_block + pos, iscol_block + pos, scall, *submat + pos, sym));
    588:     if (!max_no) {
    589:       if (scall == MAT_INITIAL_MATRIX) { /* submat[pos] is a dummy matrix */
    590:         smat          = (Mat_SubSppt *)(*submat)[pos]->data;
    591:         smat->nstages = nstages;
    592:       }
    593:       pos++; /* advance to next dummy matrix if any */
    594:     } else pos += max_no;
    595:   }

    597:   if (scall == MAT_INITIAL_MATRIX && ismax) {
    598:     /* save nstages for reuse */
    599:     subc          = (Mat_SeqBAIJ *)((*submat)[0]->data);
    600:     smat          = subc->submatis1;
    601:     smat->nstages = nstages;
    602:   }

    604:   for (i = 0; i < ismax; i++) {
    605:     PetscCall(ISDestroy(&isrow_block[i]));
    606:     PetscCall(ISDestroy(&iscol_block[i]));
    607:   }
    608:   PetscCall(PetscFree2(isrow_block, iscol_block));
    609:   PetscFunctionReturn(PETSC_SUCCESS);
    610: }

    612: /* This code is used for BAIJ and SBAIJ matrices (unfortunate dependency) */
    613: PetscErrorCode MatCreateSubMatrices_MPIBAIJ_local(Mat C, PetscInt ismax, const IS isrow[], const IS iscol[], MatReuse scall, Mat *submats, PetscBool sym)
    614: {
    615:   Mat_MPIBAIJ     *c = (Mat_MPIBAIJ *)C->data;
    616:   Mat              A = c->A;
    617:   Mat_SeqBAIJ     *a = (Mat_SeqBAIJ *)A->data, *b = (Mat_SeqBAIJ *)c->B->data, *subc;
    618:   const PetscInt **icol, **irow;
    619:   PetscInt        *nrow, *ncol, start;
    620:   PetscMPIInt     *pa, **row2proc, *row2proc_i, proc = -1;
    621:   PetscMPIInt      rank, size, tag0, tag2, tag3, tag4, *w1, *w2, *w3, *w4, nrqr, nrqs = 0, *req_source1 = NULL, *req_source2;
    622:   PetscInt       **sbuf1, **sbuf2, *sbuf2_i, ct1, ct2, **rbuf1, row;
    623:   PetscInt         msz, **ptr = NULL, *req_size = NULL, *ctr = NULL, *tmp = NULL, tcol;
    624:   PetscInt       **rbuf3 = NULL, **sbuf_aj, **rbuf2 = NULL, max1, max2;
    625:   PetscInt       **lens, is_no, ncols, *cols, mat_i, *mat_j, tmp2, jmax;
    626: #if defined(PETSC_USE_CTABLE)
    627:   PetscHMapI *cmap, cmap_i = NULL, *rmap, rmap_i;
    628: #else
    629:   PetscInt **cmap, *cmap_i = NULL, **rmap, *rmap_i;
    630: #endif
    631:   const PetscInt *irow_i, *icol_i;
    632:   PetscInt        ctr_j, *sbuf1_j, *sbuf_aj_i, *rbuf1_i, kmax, *lens_i, jcnt;
    633:   MPI_Request    *s_waits1, *r_waits1, *s_waits2, *r_waits2, *r_waits3;
    634:   MPI_Request    *r_waits4, *s_waits3, *s_waits4;
    635:   MPI_Comm        comm;
    636:   PetscScalar   **rbuf4, *rbuf4_i = NULL, **sbuf_aa, *vals, *mat_a = NULL, *imat_a = NULL, *sbuf_aa_i;
    637:   PetscMPIInt    *onodes1, *olengths1, end;
    638:   PetscInt       *imat_ilen, *imat_j, *imat_i;
    639:   Mat_SubSppt    *smat_i;
    640:   PetscBool      *issorted, colflag, iscsorted = PETSC_TRUE;
    641:   PetscInt       *sbuf1_i, *rbuf2_i, *rbuf3_i, ilen;
    642:   PetscInt        bs = C->rmap->bs, bs2 = c->bs2, rstart = c->rstartbs;
    643:   PetscBool       ijonly = c->ijonly; /* private flag indicates only matrix data structures are requested */
    644:   PetscInt        nzA, nzB, *a_i = a->i, *b_i = b->i, *a_j = a->j, *b_j = b->j, ctmp, imark, *cworkA, *cworkB;
    645:   PetscScalar    *vworkA = NULL, *vworkB = NULL, *a_a = a->a, *b_a = b->a;
    646:   PetscInt        cstart = c->cstartbs, *bmap = c->garray;
    647:   PetscBool      *allrows, *allcolumns;

    649:   PetscFunctionBegin;
    650:   PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
    651:   size = c->size;
    652:   rank = c->rank;

    654:   PetscCall(PetscMalloc5(ismax, &row2proc, ismax, &cmap, ismax, &rmap, ismax + 1, &allcolumns, ismax, &allrows));
    655:   PetscCall(PetscMalloc5(ismax, (PetscInt ***)&irow, ismax, (PetscInt ***)&icol, ismax, &nrow, ismax, &ncol, ismax, &issorted));

    657:   for (PetscInt i = 0; i < ismax; i++) {
    658:     PetscCall(ISSorted(iscol[i], &issorted[i]));
    659:     if (!issorted[i]) iscsorted = issorted[i]; /* columns are not sorted! */
    660:     PetscCall(ISSorted(isrow[i], &issorted[i]));

    662:     /* Check for special case: allcolumns */
    663:     PetscCall(ISIdentity(iscol[i], &colflag));
    664:     PetscCall(ISGetLocalSize(iscol[i], &ncol[i]));

    666:     if (colflag && ncol[i] == c->Nbs) {
    667:       allcolumns[i] = PETSC_TRUE;
    668:       icol[i]       = NULL;
    669:     } else {
    670:       allcolumns[i] = PETSC_FALSE;
    671:       PetscCall(ISGetIndices(iscol[i], &icol[i]));
    672:     }

    674:     /* Check for special case: allrows */
    675:     PetscCall(ISIdentity(isrow[i], &colflag));
    676:     PetscCall(ISGetLocalSize(isrow[i], &nrow[i]));
    677:     if (colflag && nrow[i] == c->Mbs) {
    678:       allrows[i] = PETSC_TRUE;
    679:       irow[i]    = NULL;
    680:     } else {
    681:       allrows[i] = PETSC_FALSE;
    682:       PetscCall(ISGetIndices(isrow[i], &irow[i]));
    683:     }
    684:   }

    686:   if (scall == MAT_REUSE_MATRIX) {
    687:     /* Assumes new rows are same length as the old rows */
    688:     for (PetscInt i = 0; i < ismax; i++) {
    689:       subc = (Mat_SeqBAIJ *)submats[i]->data;

    691:       PetscCheck(subc->mbs == nrow[i] && subc->nbs == ncol[i], PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Cannot reuse matrix. wrong size");

    693:       /* Initial matrix as if empty */
    694:       PetscCall(PetscArrayzero(subc->ilen, subc->mbs));

    696:       /* Initial matrix as if empty */
    697:       submats[i]->factortype = C->factortype;

    699:       smat_i = subc->submatis1;

    701:       nrqs        = smat_i->nrqs;
    702:       nrqr        = smat_i->nrqr;
    703:       rbuf1       = smat_i->rbuf1;
    704:       rbuf2       = smat_i->rbuf2;
    705:       rbuf3       = smat_i->rbuf3;
    706:       req_source2 = smat_i->req_source2;

    708:       sbuf1 = smat_i->sbuf1;
    709:       sbuf2 = smat_i->sbuf2;
    710:       ptr   = smat_i->ptr;
    711:       tmp   = smat_i->tmp;
    712:       ctr   = smat_i->ctr;

    714:       pa          = smat_i->pa;
    715:       req_size    = smat_i->req_size;
    716:       req_source1 = smat_i->req_source1;

    718:       allcolumns[i] = smat_i->allcolumns;
    719:       allrows[i]    = smat_i->allrows;
    720:       row2proc[i]   = smat_i->row2proc;
    721:       rmap[i]       = smat_i->rmap;
    722:       cmap[i]       = smat_i->cmap;
    723:     }

    725:     if (!ismax) { /* Get dummy submatrices and retrieve struct submatis1 */
    726:       PetscCheck(submats[0], PETSC_COMM_SELF, PETSC_ERR_ARG_NULL, "submats are null, cannot reuse");
    727:       smat_i = (Mat_SubSppt *)submats[0]->data;

    729:       nrqs        = smat_i->nrqs;
    730:       nrqr        = smat_i->nrqr;
    731:       rbuf1       = smat_i->rbuf1;
    732:       rbuf2       = smat_i->rbuf2;
    733:       rbuf3       = smat_i->rbuf3;
    734:       req_source2 = smat_i->req_source2;

    736:       sbuf1 = smat_i->sbuf1;
    737:       sbuf2 = smat_i->sbuf2;
    738:       ptr   = smat_i->ptr;
    739:       tmp   = smat_i->tmp;
    740:       ctr   = smat_i->ctr;

    742:       pa          = smat_i->pa;
    743:       req_size    = smat_i->req_size;
    744:       req_source1 = smat_i->req_source1;

    746:       allcolumns[0] = PETSC_FALSE;
    747:     }
    748:   } else { /* scall == MAT_INITIAL_MATRIX */
    749:     /* Get some new tags to keep the communication clean */
    750:     PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag2));
    751:     PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag3));

    753:     /* evaluate communication - mesg to who, length of mesg, and buffer space
    754:      required. Based on this, buffers are allocated, and data copied into them*/
    755:     PetscCall(PetscCalloc4(size, &w1, size, &w2, size, &w3, size, &w4)); /* mesg size, initialize work vectors */

    757:     for (PetscInt i = 0; i < ismax; i++) {
    758:       jmax   = nrow[i];
    759:       irow_i = irow[i];

    761:       PetscCall(PetscMalloc1(jmax, &row2proc_i));
    762:       row2proc[i] = row2proc_i;

    764:       if (issorted[i]) proc = 0;
    765:       for (PetscInt j = 0; j < jmax; j++) {
    766:         if (!issorted[i]) proc = 0;
    767:         if (allrows[i]) row = j;
    768:         else row = irow_i[j];

    770:         while (row >= c->rangebs[proc + 1]) proc++;
    771:         w4[proc]++;
    772:         row2proc_i[j] = proc; /* map row index to proc */
    773:       }
    774:       for (PetscMPIInt j = 0; j < size; j++) {
    775:         if (w4[j]) {
    776:           w1[j] += w4[j];
    777:           w3[j]++;
    778:           w4[j] = 0;
    779:         }
    780:       }
    781:     }

    783:     nrqs     = 0; /* no of outgoing messages */
    784:     msz      = 0; /* total mesg length (for all procs) */
    785:     w1[rank] = 0; /* no mesg sent to self */
    786:     w3[rank] = 0;
    787:     for (PetscMPIInt i = 0; i < size; i++) {
    788:       if (w1[i]) {
    789:         w2[i] = 1;
    790:         nrqs++;
    791:       } /* there exists a message to proc i */
    792:     }
    793:     PetscCall(PetscMalloc1(nrqs, &pa)); /*(proc -array)*/
    794:     for (PetscMPIInt i = 0, j = 0; i < size; i++) {
    795:       if (w1[i]) pa[j++] = i;
    796:     }

    798:     /* Each message would have a header = 1 + 2*(no of IS) + data */
    799:     for (PetscMPIInt i = 0; i < nrqs; i++) {
    800:       PetscMPIInt j = pa[i];
    801:       w1[j] += w2[j] + 2 * w3[j];
    802:       msz += w1[j];
    803:     }
    804:     PetscCall(PetscInfo(0, "Number of outgoing messages %d Total message length %" PetscInt_FMT "\n", nrqs, msz));

    806:     /* Determine the number of messages to expect, their lengths, from from-ids */
    807:     PetscCall(PetscGatherNumberOfMessages(comm, w2, w1, &nrqr));
    808:     PetscCall(PetscGatherMessageLengths(comm, nrqs, nrqr, w1, &onodes1, &olengths1));

    810:     /* Now post the Irecvs corresponding to these messages */
    811:     PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag0));
    812:     PetscCall(PetscPostIrecvInt(comm, tag0, nrqr, onodes1, olengths1, &rbuf1, &r_waits1));

    814:     /* Allocate Memory for outgoing messages */
    815:     PetscCall(PetscMalloc4(size, &sbuf1, size, &ptr, 2 * msz, &tmp, size, &ctr));
    816:     PetscCall(PetscArrayzero(sbuf1, size));
    817:     PetscCall(PetscArrayzero(ptr, size));

    819:     {
    820:       PetscInt   *iptr = tmp;
    821:       PetscMPIInt k    = 0;
    822:       for (PetscMPIInt i = 0; i < nrqs; i++) {
    823:         PetscMPIInt j = pa[i];
    824:         iptr += k;
    825:         sbuf1[j] = iptr;
    826:         k        = w1[j];
    827:       }
    828:     }

    830:     /* Form the outgoing messages. Initialize the header space */
    831:     for (PetscMPIInt i = 0; i < nrqs; i++) {
    832:       PetscMPIInt j = pa[i];

    834:       sbuf1[j][0] = 0;
    835:       PetscCall(PetscArrayzero(sbuf1[j] + 1, 2 * w3[j]));
    836:       ptr[j] = sbuf1[j] + 2 * w3[j] + 1;
    837:     }

    839:     /* Parse the isrow and copy data into outbuf */
    840:     for (PetscInt i = 0; i < ismax; i++) {
    841:       row2proc_i = row2proc[i];
    842:       PetscCall(PetscArrayzero(ctr, size));
    843:       irow_i = irow[i];
    844:       jmax   = nrow[i];
    845:       for (PetscInt j = 0; j < jmax; j++) { /* parse the indices of each IS */
    846:         proc = row2proc_i[j];
    847:         if (allrows[i]) row = j;
    848:         else row = irow_i[j];

    850:         if (proc != rank) { /* copy to the outgoing buf*/
    851:           ctr[proc]++;
    852:           *ptr[proc] = row;
    853:           ptr[proc]++;
    854:         }
    855:       }
    856:       /* Update the headers for the current IS */
    857:       for (PetscMPIInt j = 0; j < size; j++) { /* Can Optimise this loop too */
    858:         if ((ctr_j = ctr[j])) {
    859:           PetscInt k;

    861:           sbuf1_j            = sbuf1[j];
    862:           k                  = ++sbuf1_j[0];
    863:           sbuf1_j[2 * k]     = ctr_j;
    864:           sbuf1_j[2 * k - 1] = i;
    865:         }
    866:       }
    867:     }

    869:     /*  Now  post the sends */
    870:     PetscCall(PetscMalloc1(nrqs, &s_waits1));
    871:     for (PetscMPIInt i = 0; i < nrqs; ++i) {
    872:       PetscMPIInt j = pa[i];
    873:       PetscCallMPI(MPIU_Isend(sbuf1[j], w1[j], MPIU_INT, j, tag0, comm, s_waits1 + i));
    874:     }

    876:     /* Post Receives to capture the buffer size */
    877:     PetscCall(PetscMalloc1(nrqs, &r_waits2));
    878:     PetscCall(PetscMalloc3(nrqs, &req_source2, nrqs, &rbuf2, nrqs, &rbuf3));
    879:     if (nrqs) rbuf2[0] = tmp + msz;
    880:     for (PetscMPIInt i = 1; i < nrqs; ++i) rbuf2[i] = rbuf2[i - 1] + w1[pa[i - 1]];
    881:     for (PetscMPIInt i = 0; i < nrqs; ++i) {
    882:       PetscMPIInt j = pa[i];
    883:       PetscCallMPI(MPIU_Irecv(rbuf2[i], w1[j], MPIU_INT, j, tag2, comm, r_waits2 + i));
    884:     }

    886:     /* Send to other procs the buf size they should allocate */
    887:     /* Receive messages*/
    888:     PetscCall(PetscMalloc1(nrqr, &s_waits2));
    889:     PetscCall(PetscMalloc3(nrqr, &sbuf2, nrqr, &req_size, nrqr, &req_source1));

    891:     PetscCallMPI(MPI_Waitall(nrqr, r_waits1, MPI_STATUSES_IGNORE));
    892:     for (PetscMPIInt i = 0; i < nrqr; ++i) {
    893:       req_size[i] = 0;
    894:       rbuf1_i     = rbuf1[i];
    895:       start       = 2 * rbuf1_i[0] + 1;
    896:       end         = olengths1[i];
    897:       PetscCall(PetscMalloc1(end, &sbuf2[i]));
    898:       sbuf2_i = sbuf2[i];
    899:       for (PetscInt j = start; j < end; j++) {
    900:         row        = rbuf1_i[j] - rstart;
    901:         ncols      = a_i[row + 1] - a_i[row] + b_i[row + 1] - b_i[row];
    902:         sbuf2_i[j] = ncols;
    903:         req_size[i] += ncols;
    904:       }
    905:       req_source1[i] = onodes1[i];
    906:       /* form the header */
    907:       sbuf2_i[0] = req_size[i];
    908:       for (PetscInt j = 1; j < start; j++) sbuf2_i[j] = rbuf1_i[j];

    910:       PetscCallMPI(MPIU_Isend(sbuf2_i, end, MPIU_INT, req_source1[i], tag2, comm, s_waits2 + i));
    911:     }
    912:     PetscCall(PetscFree(onodes1));
    913:     PetscCall(PetscFree(olengths1));

    915:     PetscCall(PetscFree(r_waits1));
    916:     PetscCall(PetscFree4(w1, w2, w3, w4));

    918:     /* Receive messages*/
    919:     PetscCall(PetscMalloc1(nrqs, &r_waits3));

    921:     PetscCallMPI(MPI_Waitall(nrqs, r_waits2, MPI_STATUSES_IGNORE));
    922:     for (PetscMPIInt i = 0; i < nrqs; ++i) {
    923:       PetscCall(PetscMalloc1(rbuf2[i][0], &rbuf3[i]));
    924:       req_source2[i] = pa[i];
    925:       PetscCallMPI(MPIU_Irecv(rbuf3[i], rbuf2[i][0], MPIU_INT, req_source2[i], tag3, comm, r_waits3 + i));
    926:     }
    927:     PetscCall(PetscFree(r_waits2));

    929:     /* Wait on sends1 and sends2 */
    930:     PetscCallMPI(MPI_Waitall(nrqs, s_waits1, MPI_STATUSES_IGNORE));
    931:     PetscCallMPI(MPI_Waitall(nrqr, s_waits2, MPI_STATUSES_IGNORE));
    932:     PetscCall(PetscFree(s_waits1));
    933:     PetscCall(PetscFree(s_waits2));

    935:     /* Now allocate sending buffers for a->j, and send them off */
    936:     PetscCall(PetscMalloc1(nrqr, &sbuf_aj));
    937:     jcnt = 0;
    938:     for (PetscMPIInt i = 0; i < nrqr; i++) jcnt += req_size[i];
    939:     if (nrqr) PetscCall(PetscMalloc1(jcnt, &sbuf_aj[0]));
    940:     for (PetscMPIInt i = 1; i < nrqr; i++) sbuf_aj[i] = sbuf_aj[i - 1] + req_size[i - 1];

    942:     PetscCall(PetscMalloc1(nrqr, &s_waits3));
    943:     {
    944:       for (PetscMPIInt i = 0; i < nrqr; i++) {
    945:         rbuf1_i   = rbuf1[i];
    946:         sbuf_aj_i = sbuf_aj[i];
    947:         ct1       = 2 * rbuf1_i[0] + 1;
    948:         ct2       = 0;
    949:         for (PetscInt j = 1, max1 = rbuf1_i[0]; j <= max1; j++) {
    950:           kmax = rbuf1[i][2 * j];
    951:           for (PetscInt k = 0; k < kmax; k++, ct1++) {
    952:             PetscInt l;
    953:             row    = rbuf1_i[ct1] - rstart;
    954:             nzA    = a_i[row + 1] - a_i[row];
    955:             nzB    = b_i[row + 1] - b_i[row];
    956:             ncols  = nzA + nzB;
    957:             cworkA = PetscSafePointerPlusOffset(a_j, a_i[row]);
    958:             cworkB = PetscSafePointerPlusOffset(b_j, b_i[row]);

    960:             /* load the column indices for this row into cols */
    961:             cols = sbuf_aj_i + ct2;
    962:             for (l = 0; l < nzB; l++) {
    963:               if ((ctmp = bmap[cworkB[l]]) < cstart) cols[l] = ctmp;
    964:               else break;
    965:             }
    966:             imark = l;
    967:             for (l = 0; l < nzA; l++) cols[imark + l] = cstart + cworkA[l];
    968:             for (l = imark; l < nzB; l++) cols[nzA + l] = bmap[cworkB[l]];
    969:             ct2 += ncols;
    970:           }
    971:         }
    972:         PetscCallMPI(MPIU_Isend(sbuf_aj_i, req_size[i], MPIU_INT, req_source1[i], tag3, comm, s_waits3 + i));
    973:       }
    974:     }

    976:     /* create col map: global col of C -> local col of submatrices */
    977: #if defined(PETSC_USE_CTABLE)
    978:     for (PetscInt i = 0; i < ismax; i++) {
    979:       if (!allcolumns[i]) {
    980:         PetscCall(PetscHMapICreateWithSize(ncol[i], cmap + i));

    982:         jmax   = ncol[i];
    983:         icol_i = icol[i];
    984:         cmap_i = cmap[i];
    985:         for (PetscInt j = 0; j < jmax; j++) PetscCall(PetscHMapISet(cmap[i], icol_i[j] + 1, j + 1));
    986:       } else cmap[i] = NULL;
    987:     }
    988: #else
    989:     for (PetscInt i = 0; i < ismax; i++) {
    990:       if (!allcolumns[i]) {
    991:         PetscCall(PetscCalloc1(c->Nbs, &cmap[i]));
    992:         jmax   = ncol[i];
    993:         icol_i = icol[i];
    994:         cmap_i = cmap[i];
    995:         for (PetscInt j = 0; j < jmax; j++) cmap_i[icol_i[j]] = j + 1;
    996:       } else cmap[i] = NULL;
    997:     }
    998: #endif

    1000:     /* Create lens which is required for MatCreate... */
    1001:     jcnt = 0;
    1002:     for (PetscInt i = 0; i < ismax; i++) jcnt += nrow[i];
    1003:     PetscCall(PetscMalloc1(ismax, &lens));
    1004:     if (ismax) PetscCall(PetscCalloc1(jcnt, &lens[0]));
    1005:     for (PetscInt i = 1; i < ismax; i++) lens[i] = PetscSafePointerPlusOffset(lens[i - 1], nrow[i - 1]);

    1007:     /* Update lens from local data */
    1008:     for (PetscInt i = 0; i < ismax; i++) {
    1009:       row2proc_i = row2proc[i];
    1010:       jmax       = nrow[i];
    1011:       if (!allcolumns[i]) cmap_i = cmap[i];
    1012:       irow_i = irow[i];
    1013:       lens_i = lens[i];
    1014:       for (PetscInt j = 0; j < jmax; j++) {
    1015:         if (allrows[i]) row = j;
    1016:         else row = irow_i[j]; /* global blocked row of C */

    1018:         proc = row2proc_i[j];
    1019:         if (proc == rank) {
    1020:           /* Get indices from matA and then from matB */
    1021: #if defined(PETSC_USE_CTABLE)
    1022:           PetscInt tt;
    1023: #endif
    1024:           row    = row - rstart;
    1025:           nzA    = a_i[row + 1] - a_i[row];
    1026:           nzB    = b_i[row + 1] - b_i[row];
    1027:           cworkA = a_j + a_i[row];
    1028:           cworkB = PetscSafePointerPlusOffset(b_j, b_i[row]);

    1030:           if (!allcolumns[i]) {
    1031: #if defined(PETSC_USE_CTABLE)
    1032:             for (PetscInt k = 0; k < nzA; k++) {
    1033:               PetscCall(PetscHMapIGetWithDefault(cmap_i, cstart + cworkA[k] + 1, 0, &tt));
    1034:               if (tt) lens_i[j]++;
    1035:             }
    1036:             for (PetscInt k = 0; k < nzB; k++) {
    1037:               PetscCall(PetscHMapIGetWithDefault(cmap_i, bmap[cworkB[k]] + 1, 0, &tt));
    1038:               if (tt) lens_i[j]++;
    1039:             }

    1041: #else
    1042:             for (PetscInt k = 0; k < nzA; k++) {
    1043:               if (cmap_i[cstart + cworkA[k]]) lens_i[j]++;
    1044:             }
    1045:             for (PetscInt k = 0; k < nzB; k++) {
    1046:               if (cmap_i[bmap[cworkB[k]]]) lens_i[j]++;
    1047:             }
    1048: #endif
    1049:           } else { /* allcolumns */
    1050:             lens_i[j] = nzA + nzB;
    1051:           }
    1052:         }
    1053:       }
    1054:     }

    1056:     /* Create row map: global row of C -> local row of submatrices */
    1057:     for (PetscInt i = 0; i < ismax; i++) {
    1058:       if (!allrows[i]) {
    1059: #if defined(PETSC_USE_CTABLE)
    1060:         PetscCall(PetscHMapICreateWithSize(nrow[i], rmap + i));
    1061:         irow_i = irow[i];
    1062:         jmax   = nrow[i];
    1063:         for (PetscInt j = 0; j < jmax; j++) {
    1064:           if (allrows[i]) {
    1065:             PetscCall(PetscHMapISet(rmap[i], j + 1, j + 1));
    1066:           } else {
    1067:             PetscCall(PetscHMapISet(rmap[i], irow_i[j] + 1, j + 1));
    1068:           }
    1069:         }
    1070: #else
    1071:         PetscCall(PetscCalloc1(c->Mbs, &rmap[i]));
    1072:         rmap_i = rmap[i];
    1073:         irow_i = irow[i];
    1074:         jmax   = nrow[i];
    1075:         for (PetscInt j = 0; j < jmax; j++) {
    1076:           if (allrows[i]) rmap_i[j] = j;
    1077:           else rmap_i[irow_i[j]] = j;
    1078:         }
    1079: #endif
    1080:       } else rmap[i] = NULL;
    1081:     }

    1083:     /* Update lens from offproc data */
    1084:     {
    1085:       PetscInt *rbuf2_i, *rbuf3_i, *sbuf1_i;

    1087:       PetscCallMPI(MPI_Waitall(nrqs, r_waits3, MPI_STATUSES_IGNORE));
    1088:       for (tmp2 = 0; tmp2 < nrqs; tmp2++) {
    1089:         sbuf1_i = sbuf1[pa[tmp2]];
    1090:         jmax    = sbuf1_i[0];
    1091:         ct1     = 2 * jmax + 1;
    1092:         ct2     = 0;
    1093:         rbuf2_i = rbuf2[tmp2];
    1094:         rbuf3_i = rbuf3[tmp2];
    1095:         for (PetscInt j = 1; j <= jmax; j++) {
    1096:           is_no  = sbuf1_i[2 * j - 1];
    1097:           max1   = sbuf1_i[2 * j];
    1098:           lens_i = lens[is_no];
    1099:           if (!allcolumns[is_no]) cmap_i = cmap[is_no];
    1100:           rmap_i = rmap[is_no];
    1101:           for (PetscInt k = 0; k < max1; k++, ct1++) {
    1102:             if (allrows[is_no]) {
    1103:               row = sbuf1_i[ct1];
    1104:             } else {
    1105: #if defined(PETSC_USE_CTABLE)
    1106:               PetscCall(PetscHMapIGetWithDefault(rmap_i, sbuf1_i[ct1] + 1, 0, &row));
    1107:               row--;
    1108:               PetscCheck(row >= 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "row not found in table");
    1109: #else
    1110:               row = rmap_i[sbuf1_i[ct1]]; /* the val in the new matrix to be */
    1111: #endif
    1112:             }
    1113:             max2 = rbuf2_i[ct1];
    1114:             for (PetscInt l = 0; l < max2; l++, ct2++) {
    1115:               if (!allcolumns[is_no]) {
    1116: #if defined(PETSC_USE_CTABLE)
    1117:                 PetscCall(PetscHMapIGetWithDefault(cmap_i, rbuf3_i[ct2] + 1, 0, &tcol));
    1118: #else
    1119:                 tcol = cmap_i[rbuf3_i[ct2]];
    1120: #endif
    1121:                 if (tcol) lens_i[row]++;
    1122:               } else {         /* allcolumns */
    1123:                 lens_i[row]++; /* lens_i[row] += max2 ? */
    1124:               }
    1125:             }
    1126:           }
    1127:         }
    1128:       }
    1129:     }
    1130:     PetscCall(PetscFree(r_waits3));
    1131:     PetscCallMPI(MPI_Waitall(nrqr, s_waits3, MPI_STATUSES_IGNORE));
    1132:     PetscCall(PetscFree(s_waits3));

    1134:     /* Create the submatrices */
    1135:     for (PetscInt i = 0; i < ismax; i++) {
    1136:       PetscInt bs_tmp;
    1137:       if (ijonly) bs_tmp = 1;
    1138:       else bs_tmp = bs;

    1140:       PetscCall(MatCreate(PETSC_COMM_SELF, submats + i));
    1141:       PetscCall(MatSetSizes(submats[i], nrow[i] * bs_tmp, ncol[i] * bs_tmp, PETSC_DETERMINE, PETSC_DETERMINE));

    1143:       PetscCall(MatSetType(submats[i], sym ? ((PetscObject)A)->type_name : MATSEQBAIJ));
    1144:       PetscCall(MatSeqBAIJSetPreallocation(submats[i], bs_tmp, 0, lens[i]));
    1145:       PetscCall(MatSeqSBAIJSetPreallocation(submats[i], bs_tmp, 0, lens[i])); /* this subroutine is used by SBAIJ routines */

    1147:       /* create struct Mat_SubSppt and attached it to submat */
    1148:       PetscCall(PetscNew(&smat_i));
    1149:       subc            = (Mat_SeqBAIJ *)submats[i]->data;
    1150:       subc->submatis1 = smat_i;

    1152:       smat_i->destroy          = submats[i]->ops->destroy;
    1153:       submats[i]->ops->destroy = MatDestroySubMatrix_SeqBAIJ;
    1154:       submats[i]->factortype   = C->factortype;

    1156:       smat_i->id          = i;
    1157:       smat_i->nrqs        = nrqs;
    1158:       smat_i->nrqr        = nrqr;
    1159:       smat_i->rbuf1       = rbuf1;
    1160:       smat_i->rbuf2       = rbuf2;
    1161:       smat_i->rbuf3       = rbuf3;
    1162:       smat_i->sbuf2       = sbuf2;
    1163:       smat_i->req_source2 = req_source2;

    1165:       smat_i->sbuf1 = sbuf1;
    1166:       smat_i->ptr   = ptr;
    1167:       smat_i->tmp   = tmp;
    1168:       smat_i->ctr   = ctr;

    1170:       smat_i->pa          = pa;
    1171:       smat_i->req_size    = req_size;
    1172:       smat_i->req_source1 = req_source1;

    1174:       smat_i->allcolumns = allcolumns[i];
    1175:       smat_i->allrows    = allrows[i];
    1176:       smat_i->singleis   = PETSC_FALSE;
    1177:       smat_i->row2proc   = row2proc[i];
    1178:       smat_i->rmap       = rmap[i];
    1179:       smat_i->cmap       = cmap[i];
    1180:     }

    1182:     if (!ismax) { /* Create dummy submats[0] for reuse struct subc */
    1183:       PetscCall(MatCreate(PETSC_COMM_SELF, &submats[0]));
    1184:       PetscCall(MatSetSizes(submats[0], 0, 0, PETSC_DETERMINE, PETSC_DETERMINE));
    1185:       PetscCall(MatSetType(submats[0], MATDUMMY));

    1187:       /* create struct Mat_SubSppt and attached it to submat */
    1188:       PetscCall(PetscNew(&smat_i));
    1189:       submats[0]->data = (void *)smat_i;

    1191:       smat_i->destroy          = submats[0]->ops->destroy;
    1192:       submats[0]->ops->destroy = MatDestroySubMatrix_Dummy;
    1193:       submats[0]->factortype   = C->factortype;

    1195:       smat_i->id          = 0;
    1196:       smat_i->nrqs        = nrqs;
    1197:       smat_i->nrqr        = nrqr;
    1198:       smat_i->rbuf1       = rbuf1;
    1199:       smat_i->rbuf2       = rbuf2;
    1200:       smat_i->rbuf3       = rbuf3;
    1201:       smat_i->sbuf2       = sbuf2;
    1202:       smat_i->req_source2 = req_source2;

    1204:       smat_i->sbuf1 = sbuf1;
    1205:       smat_i->ptr   = ptr;
    1206:       smat_i->tmp   = tmp;
    1207:       smat_i->ctr   = ctr;

    1209:       smat_i->pa          = pa;
    1210:       smat_i->req_size    = req_size;
    1211:       smat_i->req_source1 = req_source1;

    1213:       smat_i->allcolumns = PETSC_FALSE;
    1214:       smat_i->singleis   = PETSC_FALSE;
    1215:       smat_i->row2proc   = NULL;
    1216:       smat_i->rmap       = NULL;
    1217:       smat_i->cmap       = NULL;
    1218:     }

    1220:     if (ismax) PetscCall(PetscFree(lens[0]));
    1221:     PetscCall(PetscFree(lens));
    1222:     if (sbuf_aj) {
    1223:       PetscCall(PetscFree(sbuf_aj[0]));
    1224:       PetscCall(PetscFree(sbuf_aj));
    1225:     }

    1227:   } /* endof scall == MAT_INITIAL_MATRIX */

    1229:   /* Post recv matrix values */
    1230:   if (!ijonly) {
    1231:     PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag4));
    1232:     PetscCall(PetscMalloc1(nrqs, &rbuf4));
    1233:     PetscCall(PetscMalloc1(nrqs, &r_waits4));
    1234:     for (PetscMPIInt i = 0; i < nrqs; ++i) {
    1235:       PetscCall(PetscMalloc1(rbuf2[i][0] * bs2, &rbuf4[i]));
    1236:       PetscCallMPI(MPIU_Irecv(rbuf4[i], rbuf2[i][0] * bs2, MPIU_SCALAR, req_source2[i], tag4, comm, r_waits4 + i));
    1237:     }

    1239:     /* Allocate sending buffers for a->a, and send them off */
    1240:     PetscCall(PetscMalloc1(nrqr, &sbuf_aa));
    1241:     jcnt = 0;
    1242:     for (PetscMPIInt i = 0; i < nrqr; i++) jcnt += req_size[i];
    1243:     if (nrqr) PetscCall(PetscMalloc1(jcnt * bs2, &sbuf_aa[0]));
    1244:     for (PetscMPIInt i = 1; i < nrqr; i++) sbuf_aa[i] = sbuf_aa[i - 1] + req_size[i - 1] * bs2;

    1246:     PetscCall(PetscMalloc1(nrqr, &s_waits4));

    1248:     for (PetscMPIInt i = 0; i < nrqr; i++) {
    1249:       rbuf1_i   = rbuf1[i];
    1250:       sbuf_aa_i = sbuf_aa[i];
    1251:       ct1       = 2 * rbuf1_i[0] + 1;
    1252:       ct2       = 0;
    1253:       for (PetscInt j = 1, max1 = rbuf1_i[0]; j <= max1; j++) {
    1254:         kmax = rbuf1_i[2 * j];
    1255:         for (PetscInt k = 0; k < kmax; k++, ct1++) {
    1256:           PetscInt l;

    1258:           row    = rbuf1_i[ct1] - rstart;
    1259:           nzA    = a_i[row + 1] - a_i[row];
    1260:           nzB    = b_i[row + 1] - b_i[row];
    1261:           ncols  = nzA + nzB;
    1262:           cworkB = PetscSafePointerPlusOffset(b_j, b_i[row]);
    1263:           vworkA = PetscSafePointerPlusOffset(a_a, a_i[row] * bs2);
    1264:           vworkB = PetscSafePointerPlusOffset(b_a, b_i[row] * bs2);

    1266:           /* load the column values for this row into vals*/
    1267:           vals = sbuf_aa_i + ct2 * bs2;
    1268:           for (l = 0; l < nzB; l++) {
    1269:             if ((bmap[cworkB[l]]) < cstart) PetscCall(PetscArraycpy(vals + l * bs2, vworkB + l * bs2, bs2));
    1270:             else break;
    1271:           }
    1272:           imark = l;
    1273:           for (l = 0; l < nzA; l++) PetscCall(PetscArraycpy(vals + (imark + l) * bs2, vworkA + l * bs2, bs2));
    1274:           for (l = imark; l < nzB; l++) PetscCall(PetscArraycpy(vals + (nzA + l) * bs2, vworkB + l * bs2, bs2));

    1276:           ct2 += ncols;
    1277:         }
    1278:       }
    1279:       PetscCallMPI(MPIU_Isend(sbuf_aa_i, req_size[i] * bs2, MPIU_SCALAR, req_source1[i], tag4, comm, s_waits4 + i));
    1280:     }
    1281:   }

    1283:   /* Assemble the matrices */
    1284:   /* First assemble the local rows */
    1285:   for (PetscInt i = 0; i < ismax; i++) {
    1286:     row2proc_i = row2proc[i];
    1287:     subc       = (Mat_SeqBAIJ *)submats[i]->data;
    1288:     imat_ilen  = subc->ilen;
    1289:     imat_j     = subc->j;
    1290:     imat_i     = subc->i;
    1291:     imat_a     = subc->a;

    1293:     if (!allcolumns[i]) cmap_i = cmap[i];
    1294:     rmap_i = rmap[i];
    1295:     irow_i = irow[i];
    1296:     jmax   = nrow[i];
    1297:     for (PetscInt j = 0; j < jmax; j++) {
    1298:       if (allrows[i]) row = j;
    1299:       else row = irow_i[j];
    1300:       proc = row2proc_i[j];

    1302:       if (proc == rank) {
    1303:         row    = row - rstart;
    1304:         nzA    = a_i[row + 1] - a_i[row];
    1305:         nzB    = b_i[row + 1] - b_i[row];
    1306:         cworkA = a_j + a_i[row];
    1307:         cworkB = PetscSafePointerPlusOffset(b_j, b_i[row]);
    1308:         if (!ijonly) {
    1309:           vworkA = a_a + a_i[row] * bs2;
    1310:           vworkB = PetscSafePointerPlusOffset(b_a, b_i[row] * bs2);
    1311:         }

    1313:         if (allrows[i]) {
    1314:           row = row + rstart;
    1315:         } else {
    1316: #if defined(PETSC_USE_CTABLE)
    1317:           PetscCall(PetscHMapIGetWithDefault(rmap_i, row + rstart + 1, 0, &row));
    1318:           row--;

    1320:           PetscCheck(row >= 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "row not found in table");
    1321: #else
    1322:           row = rmap_i[row + rstart];
    1323: #endif
    1324:         }
    1325:         mat_i = imat_i[row];
    1326:         if (!ijonly) mat_a = PetscSafePointerPlusOffset(imat_a, mat_i * bs2);
    1327:         mat_j = PetscSafePointerPlusOffset(imat_j, mat_i);
    1328:         ilen  = imat_ilen[row];

    1330:         /* load the column indices for this row into cols*/
    1331:         if (!allcolumns[i]) {
    1332:           PetscInt l;

    1334:           for (l = 0; l < nzB; l++) {
    1335:             if ((ctmp = bmap[cworkB[l]]) < cstart) {
    1336: #if defined(PETSC_USE_CTABLE)
    1337:               PetscCall(PetscHMapIGetWithDefault(cmap_i, ctmp + 1, 0, &tcol));
    1338:               if (tcol) {
    1339: #else
    1340:               if ((tcol = cmap_i[ctmp])) {
    1341: #endif
    1342:                 *mat_j++ = tcol - 1;
    1343:                 PetscCall(PetscArraycpy(mat_a, vworkB + l * bs2, bs2));
    1344:                 mat_a += bs2;
    1345:                 ilen++;
    1346:               }
    1347:             } else break;
    1348:           }
    1349:           imark = l;
    1350:           for (PetscInt l = 0; l < nzA; l++) {
    1351: #if defined(PETSC_USE_CTABLE)
    1352:             PetscCall(PetscHMapIGetWithDefault(cmap_i, cstart + cworkA[l] + 1, 0, &tcol));
    1353:             if (tcol) {
    1354: #else
    1355:             if ((tcol = cmap_i[cstart + cworkA[l]])) {
    1356: #endif
    1357:               *mat_j++ = tcol - 1;
    1358:               if (!ijonly) {
    1359:                 PetscCall(PetscArraycpy(mat_a, vworkA + l * bs2, bs2));
    1360:                 mat_a += bs2;
    1361:               }
    1362:               ilen++;
    1363:             }
    1364:           }
    1365:           for (l = imark; l < nzB; l++) {
    1366: #if defined(PETSC_USE_CTABLE)
    1367:             PetscCall(PetscHMapIGetWithDefault(cmap_i, bmap[cworkB[l]] + 1, 0, &tcol));
    1368:             if (tcol) {
    1369: #else
    1370:             if ((tcol = cmap_i[bmap[cworkB[l]]])) {
    1371: #endif
    1372:               *mat_j++ = tcol - 1;
    1373:               if (!ijonly) {
    1374:                 PetscCall(PetscArraycpy(mat_a, vworkB + l * bs2, bs2));
    1375:                 mat_a += bs2;
    1376:               }
    1377:               ilen++;
    1378:             }
    1379:           }
    1380:         } else { /* allcolumns */
    1381:           PetscInt l;
    1382:           for (l = 0; l < nzB; l++) {
    1383:             if ((ctmp = bmap[cworkB[l]]) < cstart) {
    1384:               *mat_j++ = ctmp;
    1385:               PetscCall(PetscArraycpy(mat_a, vworkB + l * bs2, bs2));
    1386:               mat_a += bs2;
    1387:               ilen++;
    1388:             } else break;
    1389:           }
    1390:           imark = l;
    1391:           for (l = 0; l < nzA; l++) {
    1392:             *mat_j++ = cstart + cworkA[l];
    1393:             if (!ijonly) {
    1394:               PetscCall(PetscArraycpy(mat_a, vworkA + l * bs2, bs2));
    1395:               mat_a += bs2;
    1396:             }
    1397:             ilen++;
    1398:           }
    1399:           for (l = imark; l < nzB; l++) {
    1400:             *mat_j++ = bmap[cworkB[l]];
    1401:             if (!ijonly) {
    1402:               PetscCall(PetscArraycpy(mat_a, vworkB + l * bs2, bs2));
    1403:               mat_a += bs2;
    1404:             }
    1405:             ilen++;
    1406:           }
    1407:         }
    1408:         imat_ilen[row] = ilen;
    1409:       }
    1410:     }
    1411:   }

    1413:   /* Now assemble the off proc rows */
    1414:   if (!ijonly) PetscCallMPI(MPI_Waitall(nrqs, r_waits4, MPI_STATUSES_IGNORE));
    1415:   for (tmp2 = 0; tmp2 < nrqs; tmp2++) {
    1416:     sbuf1_i = sbuf1[pa[tmp2]];
    1417:     jmax    = sbuf1_i[0];
    1418:     ct1     = 2 * jmax + 1;
    1419:     ct2     = 0;
    1420:     rbuf2_i = rbuf2[tmp2];
    1421:     rbuf3_i = rbuf3[tmp2];
    1422:     if (!ijonly) rbuf4_i = rbuf4[tmp2];
    1423:     for (PetscInt j = 1; j <= jmax; j++) {
    1424:       is_no  = sbuf1_i[2 * j - 1];
    1425:       rmap_i = rmap[is_no];
    1426:       if (!allcolumns[is_no]) cmap_i = cmap[is_no];
    1427:       subc      = (Mat_SeqBAIJ *)submats[is_no]->data;
    1428:       imat_ilen = subc->ilen;
    1429:       imat_j    = subc->j;
    1430:       imat_i    = subc->i;
    1431:       if (!ijonly) imat_a = subc->a;
    1432:       max1 = sbuf1_i[2 * j];
    1433:       for (PetscInt k = 0; k < max1; k++, ct1++) { /* for each recved block row */
    1434:         row = sbuf1_i[ct1];

    1436:         if (allrows[is_no]) {
    1437:           row = sbuf1_i[ct1];
    1438:         } else {
    1439: #if defined(PETSC_USE_CTABLE)
    1440:           PetscCall(PetscHMapIGetWithDefault(rmap_i, row + 1, 0, &row));
    1441:           row--;
    1442:           PetscCheck(row >= 0, PETSC_COMM_SELF, PETSC_ERR_PLIB, "row not found in table");
    1443: #else
    1444:           row = rmap_i[row];
    1445: #endif
    1446:         }
    1447:         ilen  = imat_ilen[row];
    1448:         mat_i = imat_i[row];
    1449:         if (!ijonly) mat_a = imat_a + mat_i * bs2;
    1450:         mat_j = imat_j + mat_i;
    1451:         max2  = rbuf2_i[ct1];
    1452:         if (!allcolumns[is_no]) {
    1453:           for (PetscInt l = 0; l < max2; l++, ct2++) {
    1454: #if defined(PETSC_USE_CTABLE)
    1455:             PetscCall(PetscHMapIGetWithDefault(cmap_i, rbuf3_i[ct2] + 1, 0, &tcol));
    1456: #else
    1457:             tcol = cmap_i[rbuf3_i[ct2]];
    1458: #endif
    1459:             if (tcol) {
    1460:               *mat_j++ = tcol - 1;
    1461:               if (!ijonly) {
    1462:                 PetscCall(PetscArraycpy(mat_a, rbuf4_i + ct2 * bs2, bs2));
    1463:                 mat_a += bs2;
    1464:               }
    1465:               ilen++;
    1466:             }
    1467:           }
    1468:         } else { /* allcolumns */
    1469:           for (PetscInt l = 0; l < max2; l++, ct2++) {
    1470:             *mat_j++ = rbuf3_i[ct2]; /* same global column index of C */
    1471:             if (!ijonly) {
    1472:               PetscCall(PetscArraycpy(mat_a, rbuf4_i + ct2 * bs2, bs2));
    1473:               mat_a += bs2;
    1474:             }
    1475:             ilen++;
    1476:           }
    1477:         }
    1478:         imat_ilen[row] = ilen;
    1479:       }
    1480:     }
    1481:   }

    1483:   if (!iscsorted) { /* sort column indices of the rows */
    1484:     MatScalar *work;

    1486:     PetscCall(PetscMalloc1(bs2, &work));
    1487:     for (PetscInt i = 0; i < ismax; i++) {
    1488:       subc      = (Mat_SeqBAIJ *)submats[i]->data;
    1489:       imat_ilen = subc->ilen;
    1490:       imat_j    = subc->j;
    1491:       imat_i    = subc->i;
    1492:       if (!ijonly) imat_a = subc->a;
    1493:       if (allcolumns[i]) continue;

    1495:       jmax = nrow[i];
    1496:       for (PetscInt j = 0; j < jmax; j++) {
    1497:         mat_i = imat_i[j];
    1498:         mat_j = imat_j + mat_i;
    1499:         ilen  = imat_ilen[j];
    1500:         if (ijonly) {
    1501:           PetscCall(PetscSortInt(ilen, mat_j));
    1502:         } else {
    1503:           mat_a = imat_a + mat_i * bs2;
    1504:           PetscCall(PetscSortIntWithDataArray(ilen, mat_j, mat_a, bs2 * sizeof(MatScalar), work));
    1505:         }
    1506:       }
    1507:     }
    1508:     PetscCall(PetscFree(work));
    1509:   }

    1511:   if (!ijonly) {
    1512:     PetscCall(PetscFree(r_waits4));
    1513:     PetscCallMPI(MPI_Waitall(nrqr, s_waits4, MPI_STATUSES_IGNORE));
    1514:     PetscCall(PetscFree(s_waits4));
    1515:   }

    1517:   /* Restore the indices */
    1518:   for (PetscInt i = 0; i < ismax; i++) {
    1519:     if (!allrows[i]) PetscCall(ISRestoreIndices(isrow[i], irow + i));
    1520:     if (!allcolumns[i]) PetscCall(ISRestoreIndices(iscol[i], icol + i));
    1521:   }

    1523:   for (PetscInt i = 0; i < ismax; i++) {
    1524:     PetscCall(MatAssemblyBegin(submats[i], MAT_FINAL_ASSEMBLY));
    1525:     PetscCall(MatAssemblyEnd(submats[i], MAT_FINAL_ASSEMBLY));
    1526:   }

    1528:   PetscCall(PetscFree5(*(PetscInt ***)&irow, *(PetscInt ***)&icol, nrow, ncol, issorted));
    1529:   PetscCall(PetscFree5(row2proc, cmap, rmap, allcolumns, allrows));

    1531:   if (!ijonly) {
    1532:     if (sbuf_aa) {
    1533:       PetscCall(PetscFree(sbuf_aa[0]));
    1534:       PetscCall(PetscFree(sbuf_aa));
    1535:     }

    1537:     for (PetscMPIInt i = 0; i < nrqs; ++i) PetscCall(PetscFree(rbuf4[i]));
    1538:     PetscCall(PetscFree(rbuf4));
    1539:   }
    1540:   c->ijonly = PETSC_FALSE; /* set back to the default */
    1541:   PetscFunctionReturn(PETSC_SUCCESS);
    1542: }