1 #include <petscsys.h> /*I "petscsys.h" I*/
2 #include <petsc/private/petscimpl.h>
3
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 };
9
10 /*
11 Private routine to delete internal shared memory communicator when a communicator is freed.
12
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.
14
15 Note: this is declared extern "C" because it is passed to MPI_Comm_create_keyval()
16
17 */
Petsc_ShmComm_Attr_DeleteFn(MPI_Comm comm,PetscMPIInt keyval,void * val,void * extra_state)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;
21
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 }
29
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];
PetscShmCommDestroyDuppedComms(void)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
48
49 /*@C
50 PetscShmCommGet - Returns a sub-communicator of all ranks that share a common memory
51
52 Collective.
53
54 Input Parameter:
55 . globcomm - `MPI_Comm`, which can be a user `MPI_Comm` or a PETSc inner `MPI_Comm`
56
57 Output Parameter:
58 . pshmcomm - the PETSc shared memory communicator object
59
60 Level: developer
61
62 Note:
63 When used with MPICH, MPICH must be configured with `--download-mpich-device=ch3:nemesis`
64
65 .seealso: `PetscShmCommGlobalToLocal()`, `PetscShmCommLocalToGlobal()`, `PetscShmCommGetMpiShmComm()`
66 @*/
PetscShmCommGet(MPI_Comm globcomm,PetscShmComm * pshmcomm)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;
73
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 }
99
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);
103
104 PetscCall(PetscNew(pshmcomm));
105 (*pshmcomm)->globcomm = globcomm;
106
107 PetscCallMPI(MPI_Comm_split_type(globcomm, MPI_COMM_TYPE_SHARED, 0, MPI_INFO_NULL, &(*pshmcomm)->shmcomm));
108
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));
119
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 }
127
128 /*@C
129 PetscShmCommGlobalToLocal - Given a global rank returns the local rank in the shared memory communicator
130
131 Input Parameters:
132 + pshmcomm - the shared memory communicator object
133 - grank - the global rank
134
135 Output Parameter:
136 . lrank - the local rank, or `MPI_PROC_NULL` if it does not exist
137
138 Level: developer
139
140 Developer Notes:
141 Assumes the pshmcomm->globranks[] is sorted
142
143 It may be better to rewrite this to map multiple global ranks to local in the same function call
144
145 .seealso: `PetscShmCommGet()`, `PetscShmCommLocalToGlobal()`, `PetscShmCommGetMpiShmComm()`
146 @*/
PetscShmCommGlobalToLocal(PetscShmComm pshmcomm,PetscMPIInt grank,PetscMPIInt * lrank)147 PetscErrorCode PetscShmCommGlobalToLocal(PetscShmComm pshmcomm, PetscMPIInt grank, PetscMPIInt *lrank)
148 {
149 PetscMPIInt low, high, t, i;
150 PetscBool flg = PETSC_FALSE;
151
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 }
176
177 /*@C
178 PetscShmCommLocalToGlobal - Given a local rank in the shared memory communicator returns the global rank
179
180 Input Parameters:
181 + pshmcomm - the shared memory communicator object
182 - lrank - the local rank in the shared memory communicator
183
184 Output Parameter:
185 . grank - the global rank in the global communicator where the shared memory communicator is built
186
187 Level: developer
188
189 .seealso: `PetscShmCommGlobalToLocal()`, `PetscShmCommGet()`, `PetscShmCommGetMpiShmComm()`
190 @*/
PetscShmCommLocalToGlobal(PetscShmComm pshmcomm,PetscMPIInt lrank,PetscMPIInt * grank)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 }
200
201 /*@C
202 PetscShmCommGetMpiShmComm - Returns the MPI communicator that represents all processes with common shared memory
203
204 Input Parameter:
205 . pshmcomm - PetscShmComm object obtained with PetscShmCommGet()
206
207 Output Parameter:
208 . comm - the MPI communicator
209
210 Level: developer
211
212 .seealso: `PetscShmCommGlobalToLocal()`, `PetscShmCommGet()`, `PetscShmCommLocalToGlobal()`
213 @*/
PetscShmCommGetMpiShmComm(PetscShmComm pshmcomm,MPI_Comm * comm)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 }
222