Date: 2022-03-04T17:37:35+00:00

Actual source code: sfneighbor.c

      1: #include <../src/vec/is/sf/impls/basic/sfpack.h>
      2: #include <../src/vec/is/sf/impls/basic/sfbasic.h>

      4: #if defined(PETSC_HAVE_MPI_NEIGHBORHOOD_COLLECTIVES)

      6: typedef struct {
      7:   SFBASICHEADER;
      8:   MPI_Comm      comms[2];       /* Communicators with distributed topology in both directions */
      9:   PetscBool     initialized[2]; /* Are the two communicators initialized? */
     10:   PetscMPIInt   *rootdispls,*rootcounts,*leafdispls,*leafcounts; /* displs/counts for non-distinguished ranks */
     11:   PetscInt      rootdegree,leafdegree;
     12: } PetscSF_Neighbor;

     14: /*===================================================================================*/
     15: /*              Internal utility routines                                            */
     16: /*===================================================================================*/

     18: /* Get the communicator with distributed graph topology, which is not cheap to build so we do it on demand (instead of at PetscSFSetUp time) */
     19: static PetscErrorCode PetscSFGetDistComm_Neighbor(PetscSF sf,PetscSFDirection direction,MPI_Comm *distcomm)
     20: {
     21:   PetscErrorCode    ierr;
     22:   PetscSF_Neighbor  *dat = (PetscSF_Neighbor*)sf->data;
     23:   PetscInt          nrootranks,ndrootranks,nleafranks,ndleafranks;
     24:   const PetscMPIInt *rootranks,*leafranks;
     25:   MPI_Comm          comm;

     28:   PetscSFGetRootInfo_Basic(sf,&nrootranks,&ndrootranks,&rootranks,NULL,NULL);      /* Which ranks will access my roots (I am a destination) */
     29:   PetscSFGetLeafInfo_Basic(sf,&nleafranks,&ndleafranks,&leafranks,NULL,NULL,NULL); /* My leaves will access whose roots (I am a source) */

     31:   if (!dat->initialized[direction]) {
     32:     const PetscMPIInt indegree  = nrootranks-ndrootranks,*sources      = rootranks+ndrootranks;
     33:     const PetscMPIInt outdegree = nleafranks-ndleafranks,*destinations = leafranks+ndleafranks;
     34:     MPI_Comm          *mycomm   = &dat->comms[direction];
     35:     PetscObjectGetComm((PetscObject)sf,&comm);
     36:     if (direction == PETSCSF_LEAF2../../../../../..) {
     37:       MPI_Dist_graph_create_adjacent(comm,indegree,sources,dat->rootcounts/*src weights*/,outdegree,destinations,dat->leafcounts/*dest weights*/,MPI_INFO_NULL,1/*reorder*/,mycomm);
     38:     } else { /* PETSCSF_../../../../../..2LEAF, reverse src & dest */
     39:       MPI_Dist_graph_create_adjacent(comm,outdegree,destinations,dat->leafcounts/*src weights*/,indegree,sources,dat->rootcounts/*dest weights*/,MPI_INFO_NULL,1/*reorder*/,mycomm);
     40:     }
     41:     dat->initialized[direction] = PETSC_TRUE;
     42:   }
     43:   *distcomm = dat->comms[direction];
     44:   return(0);
     45: }

     47: /*===================================================================================*/
     48: /*              Implementations of SF public APIs                                    */
     49: /*===================================================================================*/
     50: static PetscErrorCode PetscSFSetUp_Neighbor(PetscSF sf)
     51: {
     52:   PetscErrorCode   ierr;
     53:   PetscSF_Neighbor *dat = (PetscSF_Neighbor*)sf->data;
     54:   PetscInt         i,j,nrootranks,ndrootranks,nleafranks,ndleafranks;
     55:   const PetscInt   *rootoffset,*leafoffset;
     56:   PetscMPIInt      m,n;

     59:   /* SFNeighbor inherits from Basic */
     60:   PetscSFSetUp_Basic(sf);
     61:   /* SFNeighbor specific */
     62:   sf->persistent  = PETSC_FALSE;
     63:   PetscSFGetRootInfo_Basic(sf,&nrootranks,&ndrootranks,NULL,&rootoffset,NULL);
     64:   PetscSFGetLeafInfo_Basic(sf,&nleafranks,&ndleafranks,NULL,&leafoffset,NULL,NULL);
     65:   dat->rootdegree = nrootranks-ndrootranks;
     66:   dat->leafdegree = nleafranks-ndleafranks;
     67:   sf->nleafreqs   = 0;
     68:   dat->nrootreqs  = 1;

     70:   /* Only setup MPI displs/counts for non-distinguished ranks. Distinguished ranks use shared memory */
     71:   PetscMalloc4(dat->rootdegree,&dat->rootdispls,dat->rootdegree,&dat->rootcounts,dat->leafdegree,&dat->leafdispls,dat->leafdegree,&dat->leafcounts);
     72:   for (i=ndrootranks,j=0; i<nrootranks; i++,j++) {
     73:     PetscMPIIntCast(rootoffset[i]-rootoffset[ndrootranks],&m); dat->rootdispls[j] = m;
     74:     PetscMPIIntCast(rootoffset[i+1]-rootoffset[i],        &n); dat->rootcounts[j] = n;
     75:   }

     77:   for (i=ndleafranks,j=0; i<nleafranks; i++,j++) {
     78:     PetscMPIIntCast(leafoffset[i]-leafoffset[ndleafranks],&m); dat->leafdispls[j] = m;
     79:     PetscMPIIntCast(leafoffset[i+1]-leafoffset[i],        &n); dat->leafcounts[j] = n;
     80:   }
     81:   return(0);
     82: }

     84: static PetscErrorCode PetscSFReset_Neighbor(PetscSF sf)
     85: {
     86:   PetscErrorCode       ierr;
     87:   PetscInt             i;
     88:   PetscSF_Neighbor     *dat = (PetscSF_Neighbor*)sf->data;

     91:   if (dat->inuse) SETERRQ(PetscObjectComm((PetscObject)sf),PETSC_ERR_ARG_WRONGSTATE,"Outstanding operation has not been completed");
     92:   PetscFree4(dat->rootdispls,dat->rootcounts,dat->leafdispls,dat->leafcounts);
     93:   for (i=0; i<2; i++) {
     94:     if (dat->initialized[i]) {
     95:       MPI_Comm_free(&dat->comms[i]);
     96:       dat->initialized[i] = PETSC_FALSE;
     97:     }
     98:   }
     99:   PetscSFReset_Basic(sf); /* Common part */
    100:   return(0);
    101: }

    103: static PetscErrorCode PetscSFDestroy_Neighbor(PetscSF sf)
    104: {

    108:   PetscSFReset_Neighbor(sf);
    109:   PetscFree(sf->data);
    110:   return(0);
    111: }

    113: static PetscErrorCode PetscSFBcastBegin_Neighbor(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,void *leafdata,MPI_Op op)
    114: {
    115:   PetscErrorCode       ierr;
    116:   PetscSFLink          link;
    117:   PetscSF_Neighbor     *dat = (PetscSF_Neighbor*)sf->data;
    118:   MPI_Comm             distcomm = MPI_COMM_NULL;
    119:   void                 *rootbuf = NULL,*leafbuf = NULL;
    120:   MPI_Request          *req;

    123:   PetscSFLinkCreate(sf,unit,rootmtype,rootdata,leafmtype,leafdata,op,PETSCSF_BCAST,&link);
    124:   PetscSFLinkPackRootData(sf,link,PETSCSF_REMOTE,rootdata);
    125:   /* Do neighborhood alltoallv for remote ranks */
    126:   PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(sf,link,PETSC_TRUE/* device2host before sending */);
    127:   PetscSFGetDistComm_Neighbor(sf,PETSCSF_../../../../../..2LEAF,&distcomm);
    128:   PetscSFLinkGetMPIBuffersAndRequests(sf,link,PETSCSF_../../../../../..2LEAF,&rootbuf,&leafbuf,&req,NULL);
    129:   PetscSFLinkSyncStreamBeforeCallMPI(sf,link,PETSCSF_../../../../../..2LEAF);
    130:   MPI_Start_ineighbor_alltoallv(dat->rootdegree,dat->leafdegree,rootbuf,dat->rootcounts,dat->rootdispls,unit,leafbuf,dat->leafcounts,dat->leafdispls,unit,distcomm,req);
    131:   PetscSFLinkScatterLocal(sf,link,PETSCSF_../../../../../..2LEAF,(void*)rootdata,leafdata,op);
    132:   return(0);
    133: }

    135: PETSC_STATIC_INLINE PetscErrorCode PetscSFLeafToRootBegin_Neighbor(PetscSF sf,MPI_Datatype unit,PetscMemType leafmtype,const void *leafdata,PetscMemType rootmtype,void *rootdata,MPI_Op op,PetscSFOperation sfop,PetscSFLink *out)
    136: {
    137:   PetscErrorCode       ierr;
    138:   PetscSFLink          link;
    139:   PetscSF_Neighbor     *dat = (PetscSF_Neighbor*)sf->data;
    140:   MPI_Comm             distcomm = MPI_COMM_NULL;
    141:   void                 *rootbuf = NULL,*leafbuf = NULL;
    142:   MPI_Request          *req = NULL;

    145:   PetscSFLinkCreate(sf,unit,rootmtype,rootdata,leafmtype,leafdata,op,sfop,&link);
    146:   PetscSFLinkPackLeafData(sf,link,PETSCSF_REMOTE,leafdata);
    147:   /* Do neighborhood alltoallv for remote ranks */
    148:   PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(sf,link,PETSC_TRUE/* device2host before sending */);
    149:   PetscSFGetDistComm_Neighbor(sf,PETSCSF_LEAF2../../../../../..,&distcomm);
    150:   PetscSFLinkGetMPIBuffersAndRequests(sf,link,PETSCSF_LEAF2../../../../../..,&rootbuf,&leafbuf,&req,NULL);
    151:   PetscSFLinkSyncStreamBeforeCallMPI(sf,link,PETSCSF_LEAF2../../../../../..);
    152:   MPI_Start_ineighbor_alltoallv(dat->leafdegree,dat->rootdegree,leafbuf,dat->leafcounts,dat->leafdispls,unit,rootbuf,dat->rootcounts,dat->rootdispls,unit,distcomm,req);
    153:   *out = link;
    154:   return(0);
    155: }

    157: static PetscErrorCode PetscSFReduceBegin_Neighbor(PetscSF sf,MPI_Datatype unit,PetscMemType leafmtype,const void *leafdata,PetscMemType rootmtype,void *rootdata,MPI_Op op)
    158: {
    159:   PetscErrorCode       ierr;
    160:   PetscSFLink          link = NULL;

    163:   PetscSFLeafToRootBegin_Neighbor(sf,unit,leafmtype,leafdata,rootmtype,rootdata,op,PETSCSF_REDUCE,&link);
    164:   PetscSFLinkScatterLocal(sf,link,PETSCSF_LEAF2../../../../../..,rootdata,(void*)leafdata,op);
    165:   return(0);
    166: }

    168: static PetscErrorCode PetscSFFetchAndOpBegin_Neighbor(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,void *rootdata,PetscMemType leafmtype,const void *leafdata,void *leafupdate,MPI_Op op)
    169: {
    170:   PetscErrorCode       ierr;
    171:   PetscSFLink          link = NULL;

    174:   PetscSFLeafToRootBegin_Neighbor(sf,unit,leafmtype,leafdata,rootmtype,rootdata,op,PETSCSF_FETCH,&link);
    175:   PetscSFLinkFetchAndOpLocal(sf,link,rootdata,leafdata,leafupdate,op);
    176:   return(0);
    177: }

    179: static PetscErrorCode PetscSFFetchAndOpEnd_Neighbor(PetscSF sf,MPI_Datatype unit,void *rootdata,const void *leafdata,void *leafupdate,MPI_Op op)
    180: {
    181:   PetscErrorCode    ierr;
    182:   PetscSFLink       link = NULL;
    183:   MPI_Comm          comm = MPI_COMM_NULL;
    184:   PetscSF_Neighbor  *dat = (PetscSF_Neighbor*)sf->data;
    185:   void              *rootbuf = NULL,*leafbuf = NULL;

    188:   PetscSFLinkGetInUse(sf,unit,rootdata,leafdata,PETSC_OWN_POINTER,&link);
    189:   PetscSFLinkFinishCommunication(sf,link,PETSCSF_LEAF2../../../../../..);
    190:   /* Process remote fetch-and-op */
    191:   PetscSFLinkFetchAndOpRemote(sf,link,rootdata,op);
    192:   /* Bcast the updated rootbuf back to leaves */
    193:   PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(sf,link,PETSC_TRUE/* device2host before sending */);
    194:   PetscSFGetDistComm_Neighbor(sf,PETSCSF_../../../../../..2LEAF,&comm);
    195:   PetscSFLinkGetMPIBuffersAndRequests(sf,link,PETSCSF_../../../../../..2LEAF,&rootbuf,&leafbuf,NULL,NULL);
    196:   PetscSFLinkSyncStreamBeforeCallMPI(sf,link,PETSCSF_../../../../../..2LEAF);
    197:   MPI_Start_neighbor_alltoallv(dat->rootdegree,dat->leafdegree,rootbuf,dat->rootcounts,dat->rootdispls,unit,leafbuf,dat->leafcounts,dat->leafdispls,unit,comm);
    198:   PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(sf,link,PETSC_FALSE/* host2device after recving */);
    199:   PetscSFLinkUnpackLeafData(sf,link,PETSCSF_REMOTE,leafupdate,MPI_REPLACE);
    200:   PetscSFLinkReclaim(sf,&link);
    201:   return(0);
    202: }

    204: PETSC_INTERN PetscErrorCode PetscSFCreate_Neighbor(PetscSF sf)
    205: {
    206:   PetscErrorCode   ierr;
    207:   PetscSF_Neighbor *dat;

    210:   sf->ops->CreateEmbeddedRootSF = PetscSFCreateEmbeddedRootSF_Basic;
    211:   sf->ops->BcastEnd             = PetscSFBcastEnd_Basic;
    212:   sf->ops->ReduceEnd            = PetscSFReduceEnd_Basic;
    213:   sf->ops->GetLeafRanks         = PetscSFGetLeafRanks_Basic;
    214:   sf->ops->View                 = PetscSFView_Basic;

    216:   sf->ops->SetUp                = PetscSFSetUp_Neighbor;
    217:   sf->ops->Reset                = PetscSFReset_Neighbor;
    218:   sf->ops->Destroy              = PetscSFDestroy_Neighbor;
    219:   sf->ops->BcastBegin           = PetscSFBcastBegin_Neighbor;
    220:   sf->ops->ReduceBegin          = PetscSFReduceBegin_Neighbor;
    221:   sf->ops->FetchAndOpBegin      = PetscSFFetchAndOpBegin_Neighbor;
    222:   sf->ops->FetchAndOpEnd        = PetscSFFetchAndOpEnd_Neighbor;

    224:   PetscNewLog(sf,&dat);
    225:   sf->data = (void*)dat;
    226:   return(0);
    227: }
    228: #endif