Date: 2025-03-29T05:58:46+00:00

Actual source code: psplit.c

      1: #include <petscsys.h>

      3: /*@
      4:   PetscSplitOwnershipBlock - Given a global (or local) length determines a local
      5:   (or global) length via a simple formula. Splits so each processors local size
      6:   is divisible by the block size.

      8:   Collective (if `N` is `PETSC_DECIDE`)

     10:   Input Parameters:
     11: + comm - MPI communicator that shares the object being divided
     12: . bs   - block size
     13: . n    - local length (or `PETSC_DECIDE` to have it set)
     14: - N    - global length (or `PETSC_DECIDE`)

     16:   Level: developer

     18:   Notes:
     19:   `n` and `N` cannot be both `PETSC_DECIDE`

     21:   If one processor calls this with `N` of `PETSC_DECIDE` then all processors
     22:   must, otherwise the program will hang.

     24: .seealso: `PetscSplitOwnership()`, `PetscSplitOwnershipEqual()`
     25: @*/
     26: PetscErrorCode PetscSplitOwnershipBlock(MPI_Comm comm, PetscInt bs, PetscInt *n, PetscInt *N)
     27: {
     28:   PetscMPIInt size, rank;

     30:   PetscFunctionBegin;
     31:   PetscCheck(*N != PETSC_DECIDE || *n != PETSC_DECIDE, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Both n and N cannot be PETSC_DECIDE");

     33:   if (*N == PETSC_DECIDE) {
     34:     PetscCheck(*n % bs == 0, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "local size %" PetscInt_FMT " not divisible by block size %" PetscInt_FMT, *n, bs);
     35:     PetscCallMPI(MPIU_Allreduce(n, N, 1, MPIU_INT, MPI_SUM, comm));
     36:   } else if (*n == PETSC_DECIDE) {
     37:     PetscInt Nbs = *N / bs;
     38:     PetscCallMPI(MPI_Comm_size(comm, &size));
     39:     PetscCallMPI(MPI_Comm_rank(comm, &rank));
     40:     *n = bs * (Nbs / size + ((Nbs % size) > rank));
     41:   }
     42:   PetscFunctionReturn(PETSC_SUCCESS);
     43: }

     45: /*@
     46:   PetscSplitOwnership - Given a global (or local) length determines a local
     47:   (or global) length via a simple formula

     49:   Collective (if `n` or `N` is `PETSC_DECIDE` or `PETSC_DETERMINE`)

     51:   Input Parameters:
     52: + comm - MPI communicator that shares the object being divided
     53: . n    - local length (or `PETSC_DECIDE` to have it set)
     54: - N    - global length (or `PETSC_DETERMINE` to have it set)

     56:   Level: developer

     58:   Notes:
     59:   `n` and `N` cannot be both `PETSC_DECIDE` and `PETSC_DETERMINE`

     61:   If one processor calls this with `n` of `PETSC_DECIDE` (or with `N` `PETSC_DETERMINE`) then all processors
     62:   must. Otherwise, an error is thrown in debug mode while the program will hang
     63:   in optimized (i.e. configured --with-debugging=0) mode.

     65: .seealso: `PetscSplitOwnershipBlock()`, `PetscSplitOwnershipEqual()`, `PETSC_DECIDE`, `PETSC_DETERMINE`
     66: @*/
     67: PetscErrorCode PetscSplitOwnership(MPI_Comm comm, PetscInt *n, PetscInt *N)
     68: {
     69:   PetscMPIInt size, rank;

     71:   PetscFunctionBegin;
     72:   PetscCheck(*N != PETSC_DECIDE || *n != PETSC_DECIDE, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Both n and N cannot be PETSC_DECIDE, likely a call to VecSetSizes() or MatSetSizes() is wrong. See https://petsc.org/release/faq/#split-ownership");
     73:   if (PetscDefined(USE_DEBUG)) {
     74:     PetscMPIInt l[2], g[2];
     75:     l[0] = (*n == PETSC_DECIDE) ? 1 : 0;
     76:     l[1] = (*N == PETSC_DECIDE) ? 1 : 0;
     77:     PetscCallMPI(MPI_Comm_size(comm, &size));
     78:     PetscCallMPI(MPIU_Allreduce(l, g, 2, MPI_INT, MPI_SUM, comm));
     79:     PetscCheck(!g[0] || g[0] == size, comm, PETSC_ERR_ARG_INCOMP, "All processes must supply PETSC_DECIDE for local size");
     80:     PetscCheck(!g[1] || g[1] == size, comm, PETSC_ERR_ARG_INCOMP, "All processes must supply PETSC_DECIDE for global size");
     81:   }

     83:   if (*N == PETSC_DECIDE) {
     84:     PetscCallMPI(MPIU_Allreduce(n, N, 1, MPIU_INT, MPI_SUM, comm));
     85:   } else if (*n == PETSC_DECIDE) {
     86:     PetscCallMPI(MPI_Comm_size(comm, &size));
     87:     PetscCallMPI(MPI_Comm_rank(comm, &rank));
     88:     *n = *N / size + ((*N % size) > rank);
     89:   } else if (PetscDefined(USE_DEBUG)) {
     90:     PetscInt tmp;
     91:     PetscCallMPI(MPIU_Allreduce(n, &tmp, 1, MPIU_INT, MPI_SUM, comm));
     92:     PetscCheck(tmp == *N, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Sum of local lengths %" PetscInt_FMT " does not equal global length %" PetscInt_FMT ", my local length %" PetscInt_FMT ", likely a call to VecSetSizes() or MatSetSizes() is wrong. See https://petsc.org/release/faq/#split-ownership", tmp, *N, *n);
     93:   }
     94:   PetscFunctionReturn(PETSC_SUCCESS);
     95: }

     97: /*@
     98:   PetscSplitOwnershipEqual - Given a global (or local) length determines a local
     99:   (or global) length via a simple formula, trying to have all local lengths equal

    101:   Collective (if `n` or `N` is `PETSC_DECIDE`)

    103:   Input Parameters:
    104: + comm - MPI communicator that shares the object being divided
    105: . n    - local length (or `PETSC_DECIDE` to have it set)
    106: - N    - global length (or `PETSC_DECIDE`)

    108:   Level: developer

    110:   Notes:
    111:   This is intended to be used with `MATSCALAPACK`, where the local size must
    112:   be equal in all processes (except possibly the last one). For instance,
    113:   the local sizes when splitting `N`=50 with 6 processes are 9,9,9,9,9,5

    115:   n and N cannot be both `PETSC_DECIDE`

    117:   If one processor calls this with `n` or `N` of `PETSC_DECIDE` then all processors
    118:   must. Otherwise, an error is thrown in debug mode while the program will hang
    119:   in optimized (i.e. configured --with-debugging=0) mode.

    121: .seealso: `PetscSplitOwnership()`, `PetscSplitOwnershipBlock()`
    122: @*/
    123: PetscErrorCode PetscSplitOwnershipEqual(MPI_Comm comm, PetscInt *n, PetscInt *N)
    124: {
    125:   PetscMPIInt size, rank;

    127:   PetscFunctionBegin;
    128:   PetscCheck(*N != PETSC_DECIDE || *n != PETSC_DECIDE, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Both n and N cannot be PETSC_DECIDE");
    129:   if (PetscDefined(USE_DEBUG)) {
    130:     PetscMPIInt l[2], g[2];
    131:     l[0] = (*n == PETSC_DECIDE) ? 1 : 0;
    132:     l[1] = (*N == PETSC_DECIDE) ? 1 : 0;
    133:     PetscCallMPI(MPI_Comm_size(comm, &size));
    134:     PetscCallMPI(MPIU_Allreduce(l, g, 2, MPI_INT, MPI_SUM, comm));
    135:     PetscCheck(!g[0] || g[0] == size, comm, PETSC_ERR_ARG_INCOMP, "All processes must supply PETSC_DECIDE for local size");
    136:     PetscCheck(!g[1] || g[1] == size, comm, PETSC_ERR_ARG_INCOMP, "All processes must supply PETSC_DECIDE for global size");
    137:   }

    139:   if (*N == PETSC_DECIDE) {
    140:     PetscCallMPI(MPIU_Allreduce(n, N, 1, MPIU_INT, MPI_SUM, comm));
    141:   } else if (*n == PETSC_DECIDE) {
    142:     PetscCallMPI(MPI_Comm_size(comm, &size));
    143:     PetscCallMPI(MPI_Comm_rank(comm, &rank));
    144:     *n = *N / size;
    145:     if (*N % size) {
    146:       if ((rank + 1) * (*n + 1) <= *N) *n = *n + 1;
    147:       else if (rank * (*n + 1) <= *N) *n = *N - rank * (*n + 1);
    148:       else *n = 0;
    149:     }
    150:   } else if (PetscDefined(USE_DEBUG)) {
    151:     PetscInt tmp;
    152:     PetscCallMPI(MPIU_Allreduce(n, &tmp, 1, MPIU_INT, MPI_SUM, comm));
    153:     PetscCheck(tmp == *N, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Sum of local lengths %" PetscInt_FMT " does not equal global length %" PetscInt_FMT ", my local length %" PetscInt_FMT, tmp, *N, *n);
    154:   }
    155:   PetscFunctionReturn(PETSC_SUCCESS);
    156: }