Date: 2026-06-01T23:13:52+00:00

Actual source code: psort.c

      1: #include <petsc/private/petscimpl.h>
      2: #include <petscis.h>

      4: /* This is the bitonic merge that works on non-power-of-2 sizes found at http://www.iti.fh-flensburg.de/lang/algorithmen/sortieren/bitonic/oddn.htm */
      5: static PetscErrorCode PetscParallelSortInt_Bitonic_Merge(MPI_Comm comm, PetscMPIInt tag, PetscMPIInt rankStart, PetscMPIInt rankEnd, PetscMPIInt rank, PetscMPIInt n, PetscInt keys[], PetscInt buffer[], PetscBool forward)
      6: {
      7:   PetscInt    diff;
      8:   PetscMPIInt split, mid, partner;

     10:   PetscFunctionBegin;
     11:   diff = rankEnd - rankStart;
     12:   if (diff <= 0) PetscFunctionReturn(PETSC_SUCCESS);
     13:   if (diff == 1) {
     14:     if (forward) {
     15:       PetscCall(PetscSortInt(n, keys));
     16:     } else {
     17:       PetscCall(PetscSortReverseInt(n, keys));
     18:     }
     19:     PetscFunctionReturn(PETSC_SUCCESS);
     20:   }
     21:   split = 1;
     22:   while (2 * split < diff) split *= 2;
     23:   mid = rankStart + split;
     24:   if (rank < mid) {
     25:     partner = rank + split;
     26:   } else {
     27:     partner = rank - split;
     28:   }
     29:   if (partner < rankEnd) {
     30:     PetscMPIInt i;

     32:     PetscCallMPI(MPI_Sendrecv(keys, n, MPIU_INT, partner, tag, buffer, n, MPIU_INT, partner, tag, comm, MPI_STATUS_IGNORE));
     33:     if ((rank < partner) == (forward == PETSC_TRUE)) {
     34:       for (i = 0; i < n; i++) keys[i] = (keys[i] <= buffer[i]) ? keys[i] : buffer[i];
     35:     } else {
     36:       for (i = 0; i < n; i++) keys[i] = (keys[i] > buffer[i]) ? keys[i] : buffer[i];
     37:     }
     38:   }
     39:   /* divide and conquer */
     40:   if (rank < mid) {
     41:     PetscCall(PetscParallelSortInt_Bitonic_Merge(comm, tag, rankStart, mid, rank, n, keys, buffer, forward));
     42:   } else {
     43:     PetscCall(PetscParallelSortInt_Bitonic_Merge(comm, tag, mid, rankEnd, rank, n, keys, buffer, forward));
     44:   }
     45:   PetscFunctionReturn(PETSC_SUCCESS);
     46: }

     48: /* This is the bitonic sort that works on non-power-of-2 sizes found at http://www.iti.fh-flensburg.de/lang/algorithmen/sortieren/bitonic/oddn.htm */
     49: static PetscErrorCode PetscParallelSortInt_Bitonic_Recursive(MPI_Comm comm, PetscMPIInt tag, PetscMPIInt rankStart, PetscMPIInt rankEnd, PetscMPIInt rank, PetscMPIInt n, PetscInt keys[], PetscInt buffer[], PetscBool forward)
     50: {
     51:   PetscMPIInt diff, mid;

     53:   PetscFunctionBegin;
     54:   diff = rankEnd - rankStart;
     55:   if (diff <= 0) PetscFunctionReturn(PETSC_SUCCESS);
     56:   if (diff == 1) {
     57:     if (forward) {
     58:       PetscCall(PetscSortInt(n, keys));
     59:     } else {
     60:       PetscCall(PetscSortReverseInt(n, keys));
     61:     }
     62:     PetscFunctionReturn(PETSC_SUCCESS);
     63:   }
     64:   mid = rankStart + diff / 2;
     65:   /* divide and conquer */
     66:   if (rank < mid) {
     67:     PetscCall(PetscParallelSortInt_Bitonic_Recursive(comm, tag, rankStart, mid, rank, n, keys, buffer, (PetscBool)!forward));
     68:   } else {
     69:     PetscCall(PetscParallelSortInt_Bitonic_Recursive(comm, tag, mid, rankEnd, rank, n, keys, buffer, forward));
     70:   }
     71:   /* bitonic merge */
     72:   PetscCall(PetscParallelSortInt_Bitonic_Merge(comm, tag, rankStart, rankEnd, rank, n, keys, buffer, forward));
     73:   PetscFunctionReturn(PETSC_SUCCESS);
     74: }

     76: static PetscErrorCode PetscParallelSortInt_Bitonic(MPI_Comm comm, PetscInt n, PetscInt keys[])
     77: {
     78:   PetscMPIInt size, rank, tag, mpin;
     79:   PetscInt   *buffer;

     81:   PetscFunctionBegin;
     82:   PetscAssertPointer(keys, 3);
     83:   PetscCall(PetscCommGetNewTag(comm, &tag));
     84:   PetscCallMPI(MPI_Comm_size(comm, &size));
     85:   PetscCallMPI(MPI_Comm_rank(comm, &rank));
     86:   PetscCall(PetscMPIIntCast(n, &mpin));
     87:   PetscCall(PetscMalloc1(n, &buffer));
     88:   PetscCall(PetscParallelSortInt_Bitonic_Recursive(comm, tag, 0, size, rank, mpin, keys, buffer, PETSC_TRUE));
     89:   PetscCall(PetscFree(buffer));
     90:   PetscFunctionReturn(PETSC_SUCCESS);
     91: }

     93: static PetscErrorCode PetscParallelSampleSelect(PetscLayout mapin, PetscLayout mapout, PetscInt keysin[], PetscInt *outpivots[])
     94: {
     95:   PetscMPIInt  size, rank;
     96:   PetscInt    *pivots, *finalpivots, i;
     97:   PetscInt     non_empty, my_first, count;
     98:   PetscMPIInt *keys_per, max_keys_per;

    100:   PetscFunctionBegin;
    101:   PetscCallMPI(MPI_Comm_size(mapin->comm, &size));
    102:   PetscCallMPI(MPI_Comm_rank(mapin->comm, &rank));

    104:   /* Choose P - 1 pivots that would be ideal for the distribution on this process */
    105:   PetscCall(PetscMalloc1(size - 1, &pivots));
    106:   for (i = 0; i < size - 1; i++) {
    107:     PetscInt index;

    109:     if (!mapin->n) {
    110:       /* if this rank is empty, put "infinity" in the list.  each process knows how many empty
    111:        * processes are in the layout, so those values will be ignored from the end of the sorted
    112:        * pivots */
    113:       pivots[i] = PETSC_INT_MAX;
    114:     } else {
    115:       /* match the distribution to the desired output described by mapout by getting the index
    116:        * that is approximately the appropriate fraction through the list */
    117:       index     = ((PetscReal)mapout->range[i + 1]) * ((PetscReal)mapin->n) / ((PetscReal)mapout->N);
    118:       index     = PetscMin(index, mapin->n - 1);
    119:       index     = PetscMax(index, 0);
    120:       pivots[i] = keysin[index];
    121:     }
    122:   }
    123:   /* sort the pivots in parallel */
    124:   PetscCall(PetscParallelSortInt_Bitonic(mapin->comm, size - 1, pivots));
    125:   if (PetscDefined(USE_DEBUG)) {
    126:     PetscBool sorted;

    128:     PetscCall(PetscParallelSortedInt(mapin->comm, size - 1, pivots, &sorted));
    129:     PetscCheck(sorted, mapin->comm, PETSC_ERR_PLIB, "bitonic sort failed");
    130:   }

    132:   /* if there are Z nonempty processes, we have (P - 1) * Z real pivots, and we want to select
    133:    * at indices Z - 1, 2*Z - 1, ... (P - 1) * Z - 1 */
    134:   non_empty = size;
    135:   for (i = 0; i < size; i++)
    136:     if (mapout->range[i] == mapout->range[i + 1]) non_empty--;
    137:   PetscCall(PetscCalloc1(size + 1, &keys_per));
    138:   my_first = -1;
    139:   if (non_empty) {
    140:     for (i = 0; i < size - 1; i++) {
    141:       size_t sample      = (i + 1) * non_empty - 1;
    142:       size_t sample_rank = sample / (size - 1);

    144:       keys_per[sample_rank]++;
    145:       if (my_first < 0 && (PetscMPIInt)sample_rank == rank) my_first = (PetscInt)(sample - sample_rank * (size - 1));
    146:     }
    147:   }
    148:   for (i = 0, max_keys_per = 0; i < size; i++) max_keys_per = PetscMax(keys_per[i], max_keys_per);
    149:   PetscCall(PetscMalloc1(size * max_keys_per, &finalpivots));
    150:   /* now that we know how many pivots each process will provide, gather the selected pivots at the start of the array
    151:    * and allgather them */
    152:   for (i = 0; i < keys_per[rank]; i++) pivots[i] = pivots[my_first + i * non_empty];
    153:   for (i = keys_per[rank]; i < max_keys_per; i++) pivots[i] = PETSC_INT_MAX;
    154:   PetscCallMPI(MPI_Allgather(pivots, max_keys_per, MPIU_INT, finalpivots, max_keys_per, MPIU_INT, mapin->comm));
    155:   for (i = 0, count = 0; i < size; i++) {
    156:     PetscInt j;

    158:     for (j = 0; j < max_keys_per; j++) {
    159:       if (j < keys_per[i]) finalpivots[count++] = finalpivots[i * max_keys_per + j];
    160:     }
    161:   }
    162:   *outpivots = finalpivots;
    163:   PetscCall(PetscFree(keys_per));
    164:   PetscCall(PetscFree(pivots));
    165:   PetscFunctionReturn(PETSC_SUCCESS);
    166: }

    168: static PetscErrorCode PetscParallelRedistribute(PetscLayout map, PetscInt n, PetscInt arrayin[], PetscInt arrayout[])
    169: {
    170:   PetscMPIInt  size, rank;
    171:   PetscInt     myOffset, nextOffset;
    172:   PetscCount   total;
    173:   PetscMPIInt  nfirst, nsecond;
    174:   PetscMPIInt  firsttag, secondtag;
    175:   MPI_Request  firstreqrcv;
    176:   MPI_Request *firstreqs;
    177:   MPI_Request *secondreqs;

    179:   PetscFunctionBegin;
    180:   PetscCallMPI(MPI_Comm_size(map->comm, &size));
    181:   PetscCallMPI(MPI_Comm_rank(map->comm, &rank));
    182:   PetscCall(PetscCommGetNewTag(map->comm, &firsttag));
    183:   PetscCall(PetscCommGetNewTag(map->comm, &secondtag));
    184:   PetscCall(PetscMalloc2(size, &firstreqs, size, &secondreqs));
    185:   PetscCallMPI(MPI_Scan(&n, &nextOffset, 1, MPIU_INT, MPI_SUM, map->comm));
    186:   myOffset = nextOffset - n;
    187:   total    = map->range[rank + 1] - map->range[rank];
    188:   if (total > 0) PetscCallMPI(MPIU_Irecv(arrayout, total, MPIU_INT, MPI_ANY_SOURCE, firsttag, map->comm, &firstreqrcv));
    189:   nsecond = 0;
    190:   nfirst  = 0;
    191:   for (PetscMPIInt i = 0; i < size; i++) {
    192:     PetscInt itotal;
    193:     PetscInt overlap, oStart, oEnd;

    195:     itotal = map->range[i + 1] - map->range[i];
    196:     if (itotal <= 0) continue;
    197:     oStart  = PetscMax(myOffset, map->range[i]);
    198:     oEnd    = PetscMin(nextOffset, map->range[i + 1]);
    199:     overlap = oEnd - oStart;
    200:     if (map->range[i] >= myOffset && map->range[i] < nextOffset) {
    201:       /* send first message */
    202:       PetscCallMPI(MPIU_Isend(&arrayin[map->range[i] - myOffset], overlap, MPIU_INT, i, firsttag, map->comm, &firstreqs[nfirst++]));
    203:     } else if (overlap > 0) {
    204:       /* send second message */
    205:       PetscCallMPI(MPIU_Isend(&arrayin[oStart - myOffset], overlap, MPIU_INT, i, secondtag, map->comm, &secondreqs[nsecond++]));
    206:     } else if (overlap == 0 && myOffset > map->range[i] && myOffset < map->range[i + 1]) {
    207:       /* send empty second message */
    208:       PetscCallMPI(MPIU_Isend(&arrayin[oStart - myOffset], 0, MPIU_INT, i, secondtag, map->comm, &secondreqs[nsecond++]));
    209:     }
    210:   }
    211:   if (total > 0) {
    212:     MPI_Status  status;
    213:     PetscMPIInt sender = -1;
    214:     PetscCount  filled = 0;

    216:     PetscCallMPI(MPI_Wait(&firstreqrcv, &status));
    217:     sender = status.MPI_SOURCE;
    218:     PetscCallMPI(MPIU_Get_count(&status, MPIU_INT, &filled));
    219:     while (filled < total) {
    220:       PetscCount mfilled;

    222:       sender++;
    223:       PetscCallMPI(MPIU_Recv(&arrayout[filled], total - filled, MPIU_INT, sender, secondtag, map->comm, &status));
    224:       PetscCallMPI(MPIU_Get_count(&status, MPIU_INT, &mfilled));
    225:       filled += mfilled;
    226:     }
    227:   }
    228:   PetscCallMPI(MPI_Waitall(nfirst, firstreqs, MPI_STATUSES_IGNORE));
    229:   PetscCallMPI(MPI_Waitall(nsecond, secondreqs, MPI_STATUSES_IGNORE));
    230:   PetscCall(PetscFree2(firstreqs, secondreqs));
    231:   PetscFunctionReturn(PETSC_SUCCESS);
    232: }

    234: static PetscErrorCode PetscParallelSortInt_Samplesort(PetscLayout mapin, PetscLayout mapout, PetscInt keysin[], PetscInt keysout[])
    235: {
    236:   PetscMPIInt  size, rank;
    237:   PetscInt    *pivots = NULL, *buffer;
    238:   PetscInt     j;
    239:   PetscMPIInt *keys_per_snd, *keys_per_rcv, *offsets_snd, *offsets_rcv, nrecv;

    241:   PetscFunctionBegin;
    242:   PetscCallMPI(MPI_Comm_size(mapin->comm, &size));
    243:   PetscCallMPI(MPI_Comm_rank(mapin->comm, &rank));
    244:   PetscCall(PetscMalloc4(size, &keys_per_snd, size, &keys_per_rcv, size + 1, &offsets_snd, size + 1, &offsets_rcv));
    245:   /* sort locally */
    246:   PetscCall(PetscSortInt(mapin->n, keysin));
    247:   /* get P - 1 pivots */
    248:   PetscCall(PetscParallelSampleSelect(mapin, mapout, keysin, &pivots));
    249:   /* determine which entries in the sorted array go to which other processes based on the pivots */
    250:   j = 0;
    251:   for (PetscMPIInt i = 0; i < size - 1; i++) {
    252:     PetscInt prev = j;

    254:     while ((j < mapin->n) && (keysin[j] < pivots[i])) j++;
    255:     PetscCall(PetscMPIIntCast(prev, &offsets_snd[i]));
    256:     PetscCall(PetscMPIIntCast(j - prev, &keys_per_snd[i]));
    257:   }
    258:   PetscCall(PetscMPIIntCast(j, &offsets_snd[size - 1]));
    259:   PetscCall(PetscMPIIntCast(mapin->n - j, &keys_per_snd[size - 1]));
    260:   PetscCall(PetscMPIIntCast(mapin->n, &offsets_snd[size]));
    261:   /* get the incoming sizes */
    262:   PetscCallMPI(MPI_Alltoall(keys_per_snd, 1, MPI_INT, keys_per_rcv, 1, MPI_INT, mapin->comm));
    263:   offsets_rcv[0] = 0;
    264:   for (PetscMPIInt i = 0; i < size; i++) offsets_rcv[i + 1] = offsets_rcv[i] + keys_per_rcv[i];
    265:   nrecv = offsets_rcv[size];
    266:   /* all to all exchange */
    267:   PetscCall(PetscMalloc1(nrecv, &buffer));
    268:   PetscCallMPI(MPI_Alltoallv(keysin, keys_per_snd, offsets_snd, MPIU_INT, buffer, keys_per_rcv, offsets_rcv, MPIU_INT, mapin->comm));
    269:   PetscCall(PetscFree(pivots));
    270:   PetscCall(PetscFree4(keys_per_snd, keys_per_rcv, offsets_snd, offsets_rcv));

    272:   /* local sort */
    273:   PetscCall(PetscSortInt(nrecv, buffer));
    274: #if defined(PETSC_USE_DEBUG)
    275:   {
    276:     PetscBool sorted;

    278:     PetscCall(PetscParallelSortedInt(mapin->comm, nrecv, buffer, &sorted));
    279:     PetscCheck(sorted, mapin->comm, PETSC_ERR_PLIB, "samplesort (pre-redistribute) sort failed");
    280:   }
    281: #endif

    283:   /* redistribute to the desired order */
    284:   PetscCall(PetscParallelRedistribute(mapout, nrecv, buffer, keysout));
    285:   PetscCall(PetscFree(buffer));
    286:   PetscFunctionReturn(PETSC_SUCCESS);
    287: }

    289: /*@
    290:   PetscParallelSortInt - Globally sort a distributed array of integers

    292:   Collective

    294:   Input Parameters:
    295: + mapin  - `PetscLayout` describing the distribution of the input keys
    296: . mapout - `PetscLayout` describing the desired distribution of the output keys
    297: - keysin - the pre-sorted array of integers

    299:   Output Parameter:
    300: . keysout - the array in which the sorted integers will be stored.  If `mapin` == `mapout`, then `keysin` may be equal to `keysout`.

    302:   Level: developer

    304:   Notes:

    306:   This implements a distributed samplesort, which, with local array sizes n_in and n_out,
    307:   global size N, and global number of MPI processes P, does\:
    308: .vb
    309:   - sorts locally
    310:   - chooses pivots by sorting (in parallel) (P-1) pivot suggestions from each process using bitonic sort and allgathering a subset of (P-1) of those
    311:   - using to the pivots to repartition the keys by all-to-all exchange
    312:   - sorting the repartitioned keys locally (the array is now globally sorted, but does not match the mapout layout)
    313:   - redistributing to match the mapout layout
    314: .ve

    316:   If `keysin` != `keysout`, then `keysin` will not be changed during `PetscParallelSortInt()`.

    318: .seealso: `PetscSortInt()`, `PetscParallelSortedInt()`
    319: @*/
    320: PetscErrorCode PetscParallelSortInt(PetscLayout mapin, PetscLayout mapout, PetscInt keysin[], PetscInt keysout[])
    321: {
    322:   PetscMPIInt size;
    323:   PetscMPIInt result;
    324:   PetscInt   *keysincopy = NULL;

    326:   PetscFunctionBegin;
    327:   PetscAssertPointer(mapin, 1);
    328:   PetscAssertPointer(mapout, 2);
    329:   PetscCallMPI(MPI_Comm_compare(mapin->comm, mapout->comm, &result));
    330:   PetscCheck(result == MPI_IDENT || result == MPI_CONGRUENT, mapin->comm, PETSC_ERR_ARG_NOTSAMECOMM, "layouts are not on the same communicator");
    331:   PetscCall(PetscLayoutSetUp(mapin));
    332:   PetscCall(PetscLayoutSetUp(mapout));
    333:   if (mapin->n) PetscAssertPointer(keysin, 3);
    334:   if (mapout->n) PetscAssertPointer(keysout, 4);
    335:   PetscCheck(mapin->N == mapout->N, mapin->comm, PETSC_ERR_ARG_SIZ, "Input and output layouts have different global sizes (%" PetscInt_FMT " != %" PetscInt_FMT ")", mapin->N, mapout->N);
    336:   PetscCallMPI(MPI_Comm_size(mapin->comm, &size));
    337:   if (size == 1) {
    338:     if (keysout != keysin) PetscCall(PetscArraycpy(keysout, keysin, mapin->n));
    339:     PetscCall(PetscSortInt(mapout->n, keysout));
    340:     if (size == 1) PetscFunctionReturn(PETSC_SUCCESS);
    341:   }
    342:   if (keysout != keysin) {
    343:     PetscCall(PetscMalloc1(mapin->n, &keysincopy));
    344:     PetscCall(PetscArraycpy(keysincopy, keysin, mapin->n));
    345:     keysin = keysincopy;
    346:   }
    347:   PetscCall(PetscParallelSortInt_Samplesort(mapin, mapout, keysin, keysout));
    348: #if defined(PETSC_USE_DEBUG)
    349:   {
    350:     PetscBool sorted;

    352:     PetscCall(PetscParallelSortedInt(mapout->comm, mapout->n, keysout, &sorted));
    353:     PetscCheck(sorted, mapout->comm, PETSC_ERR_PLIB, "samplesort sort failed");
    354:   }
    355: #endif
    356:   PetscCall(PetscFree(keysincopy));
    357:   PetscFunctionReturn(PETSC_SUCCESS);
    358: }