Actual source code: mpishm.c

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

  3: struct _n_PetscShmComm {
  4:   PetscMPIInt *globranks;         /* global ranks of each rank in the shared memory communicator */
  5:   PetscMPIInt  shmsize;           /* size of the shared memory communicator */
  6:   MPI_Comm     globcomm, shmcomm; /* global communicator and shared memory communicator (a sub-communicator of the former) */
  7: };

  9: /*
 10:    Private routine to delete internal shared memory communicator when a communicator is freed.

 12:    This is called by MPI, not by users. This is called by MPI_Comm_free() when the communicator that has this  data as an attribute is freed.

 14:    Note: this is declared extern "C" because it is passed to MPI_Comm_create_keyval()

 16: */
 17: PETSC_EXTERN PetscMPIInt MPIAPI Petsc_ShmComm_Attr_DeleteFn(MPI_Comm comm, PetscMPIInt keyval, void *val, void *extra_state)
 18: {
 19:   PetscShmComm p = (PetscShmComm)val;

 21:   PetscFunctionBegin;
 22:   PetscCallReturnMPI(PetscInfo(NULL, "Deleting shared memory subcommunicator in a MPI_Comm %ld\n", (long)comm));
 23:   PetscCallMPIReturnMPI(MPI_Comm_free(&p->shmcomm));
 24:   PetscCallReturnMPI(PetscFree(p->globranks));
 25:   PetscCallReturnMPI(PetscFree(val));
 26:   PetscFunctionReturn(MPI_SUCCESS);
 27: }

 29: #if PetscDefined(HAVE_MPI_PROCESS_SHARED_MEMORY)
 30:   /* Data structures to support freeing comms created in PetscShmCommGet().
 31:   Since we predict communicators passed to PetscShmCommGet() are very likely
 32:   either a PETSc inner communicator or an MPI communicator with a linked PETSc
 33:   inner communicator, we use a simple static array to store dupped communicators
 34:   on rare cases otherwise.
 35:  */
 36:   #define MAX_SHMCOMM_DUPPED_COMMS 16
 37: static PetscInt       num_dupped_comms = 0;
 38: static MPI_Comm       shmcomm_dupped_comms[MAX_SHMCOMM_DUPPED_COMMS];
 39: static PetscErrorCode PetscShmCommDestroyDuppedComms(void)
 40: {
 41:   PetscFunctionBegin;
 42:   for (PetscInt i = 0; i < num_dupped_comms; i++) PetscCall(PetscCommDestroy(&shmcomm_dupped_comms[i]));
 43:   num_dupped_comms = 0; /* reset so that PETSc can be reinitialized */
 44:   PetscFunctionReturn(PETSC_SUCCESS);
 45: }
 46: #endif

 48: /*@
 49:   PetscShmCommGet - Returns a sub-communicator of all ranks that share a common memory

 51:   Collective.

 53:   Input Parameter:
 54: . globcomm - `MPI_Comm`, which can be a user `MPI_Comm` or a PETSc inner `MPI_Comm`

 56:   Output Parameter:
 57: . pshmcomm - the PETSc shared memory communicator object

 59:   Level: developer

 61:   Note:
 62:   When used with MPICH, MPICH must be configured with `--download-mpich-device=ch3:nemesis`

 64: .seealso: `PetscShmCommGlobalToLocal()`, `PetscShmCommLocalToGlobal()`, `PetscShmCommGetMpiShmComm()`
 65: @*/
 66: PetscErrorCode PetscShmCommGet(MPI_Comm globcomm, PetscShmComm *pshmcomm)
 67: {
 68: #if PetscDefined(HAVE_MPI_PROCESS_SHARED_MEMORY)
 69:   MPI_Group         globgroup, shmgroup;
 70:   PetscMPIInt      *shmranks, i, flg;
 71:   PetscCommCounter *counter;

 73:   PetscFunctionBegin;
 74:   PetscAssertPointer(pshmcomm, 2);
 75:   /* Get a PETSc inner comm, since we always want to stash pshmcomm on PETSc inner comms */
 76:   PetscCallMPI(MPI_Comm_get_attr(globcomm, Petsc_Counter_keyval, &counter, &flg));
 77:   if (!flg) { /* globcomm is not a PETSc comm */
 78:     union
 79:     {
 80:       MPI_Comm comm;
 81:       void    *ptr;
 82:     } ucomm;
 83:     /* check if globcomm already has a linked PETSc inner comm */
 84:     PetscCallMPI(MPI_Comm_get_attr(globcomm, Petsc_InnerComm_keyval, &ucomm, &flg));
 85:     if (!flg) {
 86:       /* globcomm does not have a linked PETSc inner comm, so we create one and replace globcomm with it */
 87:       PetscCheck(num_dupped_comms < MAX_SHMCOMM_DUPPED_COMMS, globcomm, PETSC_ERR_PLIB, "PetscShmCommGet() is trying to dup more than %d MPI_Comms", MAX_SHMCOMM_DUPPED_COMMS);
 88:       PetscCall(PetscCommDuplicate(globcomm, &globcomm, NULL));
 89:       /* Register a function to free the dupped PETSc comms at PetscFinalize() at the first time */
 90:       if (num_dupped_comms == 0) PetscCall(PetscRegisterFinalize(PetscShmCommDestroyDuppedComms));
 91:       shmcomm_dupped_comms[num_dupped_comms] = globcomm;
 92:       num_dupped_comms++;
 93:     } else {
 94:       /* otherwise, we pull out the inner comm and use it as globcomm */
 95:       globcomm = ucomm.comm;
 96:     }
 97:   }

 99:   /* Check if globcomm already has an attached pshmcomm. If no, create one */
100:   PetscCallMPI(MPI_Comm_get_attr(globcomm, Petsc_ShmComm_keyval, pshmcomm, &flg));
101:   if (flg) PetscFunctionReturn(PETSC_SUCCESS);

103:   PetscCall(PetscNew(pshmcomm));
104:   (*pshmcomm)->globcomm = globcomm;

106:   PetscCallMPI(MPI_Comm_split_type(globcomm, MPI_COMM_TYPE_SHARED, 0, MPI_INFO_NULL, &(*pshmcomm)->shmcomm));

108:   PetscCallMPI(MPI_Comm_size((*pshmcomm)->shmcomm, &(*pshmcomm)->shmsize));
109:   PetscCallMPI(MPI_Comm_group(globcomm, &globgroup));
110:   PetscCallMPI(MPI_Comm_group((*pshmcomm)->shmcomm, &shmgroup));
111:   PetscCall(PetscMalloc1((*pshmcomm)->shmsize, &shmranks));
112:   PetscCall(PetscMalloc1((*pshmcomm)->shmsize, &(*pshmcomm)->globranks));
113:   for (i = 0; i < (*pshmcomm)->shmsize; i++) shmranks[i] = i;
114:   PetscCallMPI(MPI_Group_translate_ranks(shmgroup, (*pshmcomm)->shmsize, shmranks, globgroup, (*pshmcomm)->globranks));
115:   PetscCall(PetscFree(shmranks));
116:   PetscCallMPI(MPI_Group_free(&globgroup));
117:   PetscCallMPI(MPI_Group_free(&shmgroup));

119:   for (i = 0; i < (*pshmcomm)->shmsize; i++) PetscCall(PetscInfo(NULL, "Shared memory rank %d global rank %d\n", i, (*pshmcomm)->globranks[i]));
120:   PetscCallMPI(MPI_Comm_set_attr(globcomm, Petsc_ShmComm_keyval, *pshmcomm));
121:   PetscFunctionReturn(PETSC_SUCCESS);
122: #else
123:   SETERRQ(globcomm, PETSC_ERR_SUP, "Shared memory communicators need MPI-3 package support.\nPlease upgrade your MPI or reconfigure with --download-mpich.");
124: #endif
125: }

127: /*@
128:   PetscShmCommGlobalToLocal - Given a global rank returns the local rank in the shared memory communicator

130:   Input Parameters:
131: + pshmcomm - the shared memory communicator object
132: - grank    - the global rank

134:   Output Parameter:
135: . lrank - the local rank, or `MPI_PROC_NULL` if it does not exist

137:   Level: developer

139:   Developer Notes:
140:   Assumes the pshmcomm->globranks[] is sorted

142:   It may be better to rewrite this to map multiple global ranks to local in the same function call

144: .seealso: `PetscShmCommGet()`, `PetscShmCommLocalToGlobal()`, `PetscShmCommGetMpiShmComm()`
145: @*/
146: PetscErrorCode PetscShmCommGlobalToLocal(PetscShmComm pshmcomm, PetscMPIInt grank, PetscMPIInt *lrank)
147: {
148:   PetscMPIInt low, high, t, i;
149:   PetscBool   flg = PETSC_FALSE;

151:   PetscFunctionBegin;
152:   PetscAssertPointer(pshmcomm, 1);
153:   PetscAssertPointer(lrank, 3);
154:   *lrank = MPI_PROC_NULL;
155:   if (grank < pshmcomm->globranks[0]) PetscFunctionReturn(PETSC_SUCCESS);
156:   if (grank > pshmcomm->globranks[pshmcomm->shmsize - 1]) PetscFunctionReturn(PETSC_SUCCESS);
157:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-noshared", &flg, NULL));
158:   if (flg) PetscFunctionReturn(PETSC_SUCCESS);
159:   low  = 0;
160:   high = pshmcomm->shmsize;
161:   while (high - low > 5) {
162:     t = (low + high) / 2;
163:     if (pshmcomm->globranks[t] > grank) high = t;
164:     else low = t;
165:   }
166:   for (i = low; i < high; i++) {
167:     if (pshmcomm->globranks[i] > grank) PetscFunctionReturn(PETSC_SUCCESS);
168:     if (pshmcomm->globranks[i] == grank) {
169:       *lrank = i;
170:       PetscFunctionReturn(PETSC_SUCCESS);
171:     }
172:   }
173:   PetscFunctionReturn(PETSC_SUCCESS);
174: }

176: /*@
177:   PetscShmCommLocalToGlobal - Given a local rank in the shared memory communicator returns the global rank

179:   Input Parameters:
180: + pshmcomm - the shared memory communicator object
181: - lrank    - the local rank in the shared memory communicator

183:   Output Parameter:
184: . grank - the global rank in the global communicator where the shared memory communicator is built

186:   Level: developer

188: .seealso: `PetscShmCommGlobalToLocal()`, `PetscShmCommGet()`, `PetscShmCommGetMpiShmComm()`
189: @*/
190: PetscErrorCode PetscShmCommLocalToGlobal(PetscShmComm pshmcomm, PetscMPIInt lrank, PetscMPIInt *grank)
191: {
192:   PetscFunctionBegin;
193:   PetscAssertPointer(pshmcomm, 1);
194:   PetscAssertPointer(grank, 3);
195:   PetscCheck(lrank >= 0 && lrank < pshmcomm->shmsize, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "No rank %d in the shared memory communicator", lrank);
196:   *grank = pshmcomm->globranks[lrank];
197:   PetscFunctionReturn(PETSC_SUCCESS);
198: }

200: /*@
201:   PetscShmCommGetMpiShmComm - Returns the MPI communicator that represents all processes with common shared memory

203:   Input Parameter:
204: . pshmcomm - PetscShmComm object obtained with PetscShmCommGet()

206:   Output Parameter:
207: . comm - the MPI communicator

209:   Level: developer

211: .seealso: `PetscShmCommGlobalToLocal()`, `PetscShmCommGet()`, `PetscShmCommLocalToGlobal()`
212: @*/
213: PetscErrorCode PetscShmCommGetMpiShmComm(PetscShmComm pshmcomm, MPI_Comm *comm)
214: {
215:   PetscFunctionBegin;
216:   PetscAssertPointer(pshmcomm, 1);
217:   PetscAssertPointer(comm, 2);
218:   *comm = pshmcomm->shmcomm;
219:   PetscFunctionReturn(PETSC_SUCCESS);
220: }