Actual source code: mpishm.c
1: #include <petscsys.h>
2: #include <petsc/private/petscimpl.h>
4: struct _n_PetscShmComm {
5: PetscMPIInt *globranks; /* global ranks of each rank in the shared memory communicator */
6: PetscMPIInt shmsize; /* size of the shared memory communicator */
7: MPI_Comm globcomm, shmcomm; /* global communicator and shared memory communicator (a sub-communicator of the former) */
8: };
10: /*
11: Private routine to delete internal shared memory communicator when a communicator is freed.
13: 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.
15: Note: this is declared extern "C" because it is passed to MPI_Comm_create_keyval()
17: */
18: PETSC_EXTERN PetscMPIInt MPIAPI Petsc_ShmComm_Attr_DeleteFn(MPI_Comm comm, PetscMPIInt keyval, void *val, void *extra_state)
19: {
20: PetscShmComm p = (PetscShmComm)val;
22: PetscFunctionBegin;
23: PetscCallReturnMPI(PetscInfo(NULL, "Deleting shared memory subcommunicator in a MPI_Comm %ld\n", (long)comm));
24: PetscCallMPIReturnMPI(MPI_Comm_free(&p->shmcomm));
25: PetscCallReturnMPI(PetscFree(p->globranks));
26: PetscCallReturnMPI(PetscFree(val));
27: PetscFunctionReturn(MPI_SUCCESS);
28: }
30: #ifdef PETSC_HAVE_MPI_PROCESS_SHARED_MEMORY
31: /* Data structures to support freeing comms created in PetscShmCommGet().
32: Since we predict communicators passed to PetscShmCommGet() are very likely
33: either a PETSc inner communicator or an MPI communicator with a linked PETSc
34: inner communicator, we use a simple static array to store dupped communicators
35: on rare cases otherwise.
36: */
37: #define MAX_SHMCOMM_DUPPED_COMMS 16
38: static PetscInt num_dupped_comms = 0;
39: static MPI_Comm shmcomm_dupped_comms[MAX_SHMCOMM_DUPPED_COMMS];
40: static PetscErrorCode PetscShmCommDestroyDuppedComms(void)
41: {
42: PetscFunctionBegin;
43: for (PetscInt i = 0; i < num_dupped_comms; i++) PetscCall(PetscCommDestroy(&shmcomm_dupped_comms[i]));
44: num_dupped_comms = 0; /* reset so that PETSc can be reinitialized */
45: PetscFunctionReturn(PETSC_SUCCESS);
46: }
47: #endif
49: /*@C
50: PetscShmCommGet - Returns a sub-communicator of all ranks that share a common memory
52: Collective.
54: Input Parameter:
55: . globcomm - `MPI_Comm`, which can be a user `MPI_Comm` or a PETSc inner `MPI_Comm`
57: Output Parameter:
58: . pshmcomm - the PETSc shared memory communicator object
60: Level: developer
62: Note:
63: When used with MPICH, MPICH must be configured with `--download-mpich-device=ch3:nemesis`
65: .seealso: `PetscShmCommGlobalToLocal()`, `PetscShmCommLocalToGlobal()`, `PetscShmCommGetMpiShmComm()`
66: @*/
67: PetscErrorCode PetscShmCommGet(MPI_Comm globcomm, PetscShmComm *pshmcomm)
68: {
69: #ifdef PETSC_HAVE_MPI_PROCESS_SHARED_MEMORY
70: MPI_Group globgroup, shmgroup;
71: PetscMPIInt *shmranks, i, flg;
72: PetscCommCounter *counter;
74: PetscFunctionBegin;
75: PetscAssertPointer(pshmcomm, 2);
76: /* Get a PETSc inner comm, since we always want to stash pshmcomm on PETSc inner comms */
77: PetscCallMPI(MPI_Comm_get_attr(globcomm, Petsc_Counter_keyval, &counter, &flg));
78: if (!flg) { /* globcomm is not a PETSc comm */
79: union
80: {
81: MPI_Comm comm;
82: void *ptr;
83: } ucomm;
84: /* check if globcomm already has a linked PETSc inner comm */
85: PetscCallMPI(MPI_Comm_get_attr(globcomm, Petsc_InnerComm_keyval, &ucomm, &flg));
86: if (!flg) {
87: /* globcomm does not have a linked PETSc inner comm, so we create one and replace globcomm with it */
88: 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);
89: PetscCall(PetscCommDuplicate(globcomm, &globcomm, NULL));
90: /* Register a function to free the dupped PETSc comms at PetscFinalize() at the first time */
91: if (num_dupped_comms == 0) PetscCall(PetscRegisterFinalize(PetscShmCommDestroyDuppedComms));
92: shmcomm_dupped_comms[num_dupped_comms] = globcomm;
93: num_dupped_comms++;
94: } else {
95: /* otherwise, we pull out the inner comm and use it as globcomm */
96: globcomm = ucomm.comm;
97: }
98: }
100: /* Check if globcomm already has an attached pshmcomm. If no, create one */
101: PetscCallMPI(MPI_Comm_get_attr(globcomm, Petsc_ShmComm_keyval, pshmcomm, &flg));
102: if (flg) PetscFunctionReturn(PETSC_SUCCESS);
104: PetscCall(PetscNew(pshmcomm));
105: (*pshmcomm)->globcomm = globcomm;
107: PetscCallMPI(MPI_Comm_split_type(globcomm, MPI_COMM_TYPE_SHARED, 0, MPI_INFO_NULL, &(*pshmcomm)->shmcomm));
109: PetscCallMPI(MPI_Comm_size((*pshmcomm)->shmcomm, &(*pshmcomm)->shmsize));
110: PetscCallMPI(MPI_Comm_group(globcomm, &globgroup));
111: PetscCallMPI(MPI_Comm_group((*pshmcomm)->shmcomm, &shmgroup));
112: PetscCall(PetscMalloc1((*pshmcomm)->shmsize, &shmranks));
113: PetscCall(PetscMalloc1((*pshmcomm)->shmsize, &(*pshmcomm)->globranks));
114: for (i = 0; i < (*pshmcomm)->shmsize; i++) shmranks[i] = i;
115: PetscCallMPI(MPI_Group_translate_ranks(shmgroup, (*pshmcomm)->shmsize, shmranks, globgroup, (*pshmcomm)->globranks));
116: PetscCall(PetscFree(shmranks));
117: PetscCallMPI(MPI_Group_free(&globgroup));
118: PetscCallMPI(MPI_Group_free(&shmgroup));
120: for (i = 0; i < (*pshmcomm)->shmsize; i++) PetscCall(PetscInfo(NULL, "Shared memory rank %d global rank %d\n", i, (*pshmcomm)->globranks[i]));
121: PetscCallMPI(MPI_Comm_set_attr(globcomm, Petsc_ShmComm_keyval, *pshmcomm));
122: PetscFunctionReturn(PETSC_SUCCESS);
123: #else
124: SETERRQ(globcomm, PETSC_ERR_SUP, "Shared memory communicators need MPI-3 package support.\nPlease upgrade your MPI or reconfigure with --download-mpich.");
125: #endif
126: }
128: /*@C
129: PetscShmCommGlobalToLocal - Given a global rank returns the local rank in the shared memory communicator
131: Input Parameters:
132: + pshmcomm - the shared memory communicator object
133: - grank - the global rank
135: Output Parameter:
136: . lrank - the local rank, or `MPI_PROC_NULL` if it does not exist
138: Level: developer
140: Developer Notes:
141: Assumes the pshmcomm->globranks[] is sorted
143: It may be better to rewrite this to map multiple global ranks to local in the same function call
145: .seealso: `PetscShmCommGet()`, `PetscShmCommLocalToGlobal()`, `PetscShmCommGetMpiShmComm()`
146: @*/
147: PetscErrorCode PetscShmCommGlobalToLocal(PetscShmComm pshmcomm, PetscMPIInt grank, PetscMPIInt *lrank)
148: {
149: PetscMPIInt low, high, t, i;
150: PetscBool flg = PETSC_FALSE;
152: PetscFunctionBegin;
153: PetscAssertPointer(pshmcomm, 1);
154: PetscAssertPointer(lrank, 3);
155: *lrank = MPI_PROC_NULL;
156: if (grank < pshmcomm->globranks[0]) PetscFunctionReturn(PETSC_SUCCESS);
157: if (grank > pshmcomm->globranks[pshmcomm->shmsize - 1]) PetscFunctionReturn(PETSC_SUCCESS);
158: PetscCall(PetscOptionsGetBool(NULL, NULL, "-noshared", &flg, NULL));
159: if (flg) PetscFunctionReturn(PETSC_SUCCESS);
160: low = 0;
161: high = pshmcomm->shmsize;
162: while (high - low > 5) {
163: t = (low + high) / 2;
164: if (pshmcomm->globranks[t] > grank) high = t;
165: else low = t;
166: }
167: for (i = low; i < high; i++) {
168: if (pshmcomm->globranks[i] > grank) PetscFunctionReturn(PETSC_SUCCESS);
169: if (pshmcomm->globranks[i] == grank) {
170: *lrank = i;
171: PetscFunctionReturn(PETSC_SUCCESS);
172: }
173: }
174: PetscFunctionReturn(PETSC_SUCCESS);
175: }
177: /*@C
178: PetscShmCommLocalToGlobal - Given a local rank in the shared memory communicator returns the global rank
180: Input Parameters:
181: + pshmcomm - the shared memory communicator object
182: - lrank - the local rank in the shared memory communicator
184: Output Parameter:
185: . grank - the global rank in the global communicator where the shared memory communicator is built
187: Level: developer
189: .seealso: `PetscShmCommGlobalToLocal()`, `PetscShmCommGet()`, `PetscShmCommGetMpiShmComm()`
190: @*/
191: PetscErrorCode PetscShmCommLocalToGlobal(PetscShmComm pshmcomm, PetscMPIInt lrank, PetscMPIInt *grank)
192: {
193: PetscFunctionBegin;
194: PetscAssertPointer(pshmcomm, 1);
195: PetscAssertPointer(grank, 3);
196: PetscCheck(lrank >= 0 && lrank < pshmcomm->shmsize, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "No rank %d in the shared memory communicator", lrank);
197: *grank = pshmcomm->globranks[lrank];
198: PetscFunctionReturn(PETSC_SUCCESS);
199: }
201: /*@C
202: PetscShmCommGetMpiShmComm - Returns the MPI communicator that represents all processes with common shared memory
204: Input Parameter:
205: . pshmcomm - PetscShmComm object obtained with PetscShmCommGet()
207: Output Parameter:
208: . comm - the MPI communicator
210: Level: developer
212: .seealso: `PetscShmCommGlobalToLocal()`, `PetscShmCommGet()`, `PetscShmCommLocalToGlobal()`
213: @*/
214: PetscErrorCode PetscShmCommGetMpiShmComm(PetscShmComm pshmcomm, MPI_Comm *comm)
215: {
216: PetscFunctionBegin;
217: PetscAssertPointer(pshmcomm, 1);
218: PetscAssertPointer(comm, 2);
219: *comm = pshmcomm->shmcomm;
220: PetscFunctionReturn(PETSC_SUCCESS);
221: }