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:     PetscCall(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`)

 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_DECIDE`)

 56:   Level: developer

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

 61:   If one processor calls this with `n` or `N` of `PETSC_DECIDE` 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()`
 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:     PetscCall(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:     PetscInt64 m = *n, M;
 85:     PetscCall(MPIU_Allreduce(&m, &M, 1, MPIU_INT64, MPI_SUM, comm));
 86:     PetscCheck(M <= PETSC_MAX_INT, comm, PETSC_ERR_INT_OVERFLOW, "Global size overflow %" PetscInt64_FMT ". You may consider ./configure PETSc with --with-64-bit-indices for the case you are running", M);
 87:     *N = (PetscInt)M;
 88:   } else if (*n == PETSC_DECIDE) {
 89:     PetscCallMPI(MPI_Comm_size(comm, &size));
 90:     PetscCallMPI(MPI_Comm_rank(comm, &rank));
 91:     *n = *N / size + ((*N % size) > rank);
 92:   } else if (PetscDefined(USE_DEBUG)) {
 93:     PetscInt tmp;
 94:     PetscCall(MPIU_Allreduce(n, &tmp, 1, MPIU_INT, MPI_SUM, comm));
 95:     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);
 96:   }
 97:   PetscFunctionReturn(PETSC_SUCCESS);
 98: }

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

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

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

111:   Level: developer

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

118:   n and N cannot be both `PETSC_DECIDE`

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

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

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

142:   if (*N == PETSC_DECIDE) {
143:     PetscInt64 m = *n, M;
144:     PetscCall(MPIU_Allreduce(&m, &M, 1, MPIU_INT64, MPI_SUM, comm));
145:     PetscCheck(M <= PETSC_MAX_INT, comm, PETSC_ERR_INT_OVERFLOW, "Global size overflow %" PetscInt64_FMT ". You may consider ./configure PETSc with --with-64-bit-indices for the case you are running", M);
146:     *N = (PetscInt)M;
147:   } else if (*n == PETSC_DECIDE) {
148:     PetscCallMPI(MPI_Comm_size(comm, &size));
149:     PetscCallMPI(MPI_Comm_rank(comm, &rank));
150:     *n = *N / size;
151:     if (*N % size) {
152:       if ((rank + 1) * (*n + 1) <= *N) *n = *n + 1;
153:       else if (rank * (*n + 1) <= *N) *n = *N - rank * (*n + 1);
154:       else *n = 0;
155:     }
156:   } else if (PetscDefined(USE_DEBUG)) {
157:     PetscInt tmp;
158:     PetscCall(MPIU_Allreduce(n, &tmp, 1, MPIU_INT, MPI_SUM, comm));
159:     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);
160:   }
161:   PetscFunctionReturn(PETSC_SUCCESS);
162: }