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 */ 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]; 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 @*/ 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 @*/ 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 @*/ 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 @*/ 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