Date: 2025-03-29T05:58:46+00:00

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: }