1 #include <../src/vec/is/sf/impls/basic/sfpack.h> 2 #include <../src/vec/is/sf/impls/basic/sfbasic.h> 3 4 #if defined(PETSC_HAVE_MPI_NEIGHBORHOOD_COLLECTIVES) 5 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; /* Number of non-distinguished root/leaf ranks, equal to outdegree or indegree in neigborhood collectives, depending on PetscSFDirection */ 12 } PetscSF_Neighbor; 13 14 /*===================================================================================*/ 15 /* Internal utility routines */ 16 /*===================================================================================*/ 17 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; 26 27 PetscFunctionBegin; 28 ierr = PetscSFGetRootInfo_Basic(sf,&nrootranks,&ndrootranks,&rootranks,NULL,NULL);CHKERRQ(ierr); /* Which ranks will access my roots (I am a destination) */ 29 ierr = PetscSFGetLeafInfo_Basic(sf,&nleafranks,&ndleafranks,&leafranks,NULL,NULL,NULL);CHKERRQ(ierr); /* My leaves will access whose roots (I am a source) */ 30 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 ierr = PetscObjectGetComm((PetscObject)sf,&comm);CHKERRQ(ierr); 36 if (direction == PETSCSF_LEAF2ROOT_REDUCE) { 37 ierr = MPI_Dist_graph_create_adjacent(comm,indegree,sources,dat->rootcounts/*src weights*/,outdegree,destinations,dat->leafcounts/*dest weights*/,MPI_INFO_NULL,1/*reorder*/,mycomm);CHKERRQ(ierr); 38 } else { /* PETSCSF_ROOT2LEAF_BCAST, reverse src & dest */ 39 ierr = MPI_Dist_graph_create_adjacent(comm,outdegree,destinations,dat->leafcounts/*src weights*/,indegree,sources,dat->rootcounts/*dest weights*/,MPI_INFO_NULL,1/*reorder*/,mycomm);CHKERRQ(ierr); 40 } 41 dat->initialized[direction] = PETSC_TRUE; 42 } 43 *distcomm = dat->comms[direction]; 44 PetscFunctionReturn(0); 45 } 46 47 static PetscErrorCode PetscSFPackGet_Neighbor(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,const void *leafdata,PetscSFPack *mylink) 48 { 49 PetscErrorCode ierr; 50 51 PetscFunctionBegin; 52 ierr = PetscSFPackGet_Basic_Common(sf,unit,rootmtype,rootdata,leafmtype,leafdata,1/*nrootreqs*/,1/*nleafreqs*/,mylink);CHKERRQ(ierr); 53 PetscFunctionReturn(0); 54 } 55 56 /*===================================================================================*/ 57 /* Implementations of SF public APIs */ 58 /*===================================================================================*/ 59 static PetscErrorCode PetscSFSetUp_Neighbor(PetscSF sf) 60 { 61 PetscErrorCode ierr; 62 PetscSF_Neighbor *dat = (PetscSF_Neighbor*)sf->data; 63 PetscInt i,j,nrootranks,ndrootranks,nleafranks,ndleafranks; 64 const PetscInt *rootoffset,*leafoffset; 65 PetscMPIInt m,n; 66 67 PetscFunctionBegin; 68 ierr = PetscSFSetUp_Basic(sf);CHKERRQ(ierr); 69 ierr = PetscSFGetRootInfo_Basic(sf,&nrootranks,&ndrootranks,NULL,&rootoffset,NULL);CHKERRQ(ierr); 70 ierr = PetscSFGetLeafInfo_Basic(sf,&nleafranks,&ndleafranks,NULL,&leafoffset,NULL,NULL);CHKERRQ(ierr); 71 72 dat->rootdegree = nrootranks-ndrootranks; 73 dat->leafdegree = nleafranks-ndleafranks; 74 75 /* Only setup MPI displs/counts for non-distinguished ranks. Distinguished ranks use shared memory */ 76 ierr = PetscMalloc4(dat->rootdegree,&dat->rootdispls,dat->rootdegree,&dat->rootcounts,dat->leafdegree,&dat->leafdispls,dat->leafdegree,&dat->leafcounts);CHKERRQ(ierr); 77 for (i=ndrootranks,j=0; i<nrootranks; i++,j++) { 78 ierr = PetscMPIIntCast(rootoffset[i]-rootoffset[ndrootranks],&m);CHKERRQ(ierr); dat->rootdispls[j] = m; 79 ierr = PetscMPIIntCast(rootoffset[i+1]-rootoffset[i], &n);CHKERRQ(ierr); dat->rootcounts[j] = n; 80 } 81 82 for (i=ndleafranks,j=0; i<nleafranks; i++,j++) { 83 ierr = PetscMPIIntCast(leafoffset[i]-leafoffset[ndleafranks],&m);CHKERRQ(ierr); dat->leafdispls[j] = m; 84 ierr = PetscMPIIntCast(leafoffset[i+1]-leafoffset[i], &n);CHKERRQ(ierr); dat->leafcounts[j] = n; 85 } 86 PetscFunctionReturn(0); 87 } 88 89 static PetscErrorCode PetscSFReset_Neighbor(PetscSF sf) 90 { 91 PetscErrorCode ierr; 92 PetscInt i; 93 PetscSF_Neighbor *dat = (PetscSF_Neighbor*)sf->data; 94 95 PetscFunctionBegin; 96 if (dat->inuse) SETERRQ(PetscObjectComm((PetscObject)sf),PETSC_ERR_ARG_WRONGSTATE,"Outstanding operation has not been completed"); 97 ierr = PetscFree4(dat->rootdispls,dat->rootcounts,dat->leafdispls,dat->leafcounts);CHKERRQ(ierr); 98 for (i=0; i<2; i++) { 99 if (dat->initialized[i]) { 100 ierr = MPI_Comm_free(&dat->comms[i]);CHKERRQ(ierr); 101 dat->initialized[i] = PETSC_FALSE; 102 } 103 } 104 ierr = PetscSFReset_Basic(sf);CHKERRQ(ierr); /* Common part */ 105 PetscFunctionReturn(0); 106 } 107 108 static PetscErrorCode PetscSFDestroy_Neighbor(PetscSF sf) 109 { 110 PetscErrorCode ierr; 111 112 PetscFunctionBegin; 113 ierr = PetscSFReset_Neighbor(sf);CHKERRQ(ierr); 114 ierr = PetscFree(sf->data);CHKERRQ(ierr); 115 PetscFunctionReturn(0); 116 } 117 118 static PetscErrorCode PetscSFBcastAndOpBegin_Neighbor(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,void *leafdata,MPI_Op op) 119 { 120 PetscErrorCode ierr; 121 PetscSFPack link; 122 const PetscInt *rootloc = NULL; 123 PetscSF_Neighbor *dat = (PetscSF_Neighbor*)sf->data; 124 MPI_Comm distcomm; 125 PetscMemType rootmtype_mpi,leafmtype_mpi; /* memtypes seen by MPI */ 126 127 PetscFunctionBegin; 128 ierr = PetscSFPackGet_Neighbor(sf,unit,rootmtype,rootdata,leafmtype,leafdata,&link);CHKERRQ(ierr); 129 ierr = PetscSFGetRootIndicesWithMemType_Basic(sf,rootmtype,&rootloc);CHKERRQ(ierr); 130 ierr = PetscSFPackRootData(sf,link,rootloc,rootdata,PETSC_TRUE);CHKERRQ(ierr); 131 132 /* Do neighborhood alltoallv for non-distinguished ranks */ 133 ierr = PetscSFGetDistComm_Neighbor(sf,PETSCSF_ROOT2LEAF_BCAST,&distcomm);CHKERRQ(ierr); 134 if (use_gpu_aware_mpi) { 135 rootmtype_mpi = rootmtype; 136 leafmtype_mpi = leafmtype; 137 } else { 138 rootmtype_mpi = PETSC_MEMTYPE_HOST; 139 leafmtype_mpi = PETSC_MEMTYPE_HOST; 140 } 141 ierr = MPI_Start_ineighbor_alltoallv(dat->rootdegree,dat->leafdegree,link->rootbuf[rootmtype_mpi],dat->rootcounts,dat->rootdispls,unit,link->leafbuf[leafmtype_mpi],dat->leafcounts,dat->leafdispls,unit,distcomm,link->rootreqs[PETSCSF_ROOT2LEAF_BCAST][rootmtype_mpi]);CHKERRQ(ierr); 142 if (rootmtype != leafmtype) {ierr = PetscMemcpyWithMemType(leafmtype,rootmtype,link->selfbuf[leafmtype],link->selfbuf[rootmtype],link->selfbuflen*link->unitbytes);CHKERRQ(ierr);} 143 PetscFunctionReturn(0); 144 } 145 146 static PetscErrorCode PetscSFReduceBegin_Neighbor(PetscSF sf,MPI_Datatype unit,PetscMemType leafmtype,const void *leafdata,PetscMemType rootmtype,void *rootdata,MPI_Op op) 147 { 148 PetscErrorCode ierr; 149 const PetscInt *leafloc = NULL; 150 PetscSFPack link; 151 PetscSF_Neighbor *dat = (PetscSF_Neighbor*)sf->data; 152 MPI_Comm distcomm = MPI_COMM_NULL; 153 PetscMemType rootmtype_mpi,leafmtype_mpi; /* memtypes seen by MPI */ 154 155 PetscFunctionBegin; 156 ierr = PetscSFGetLeafIndicesWithMemType_Basic(sf,leafmtype,&leafloc);CHKERRQ(ierr); 157 ierr = PetscSFPackGet_Neighbor(sf,unit,rootmtype,rootdata,leafmtype,leafdata,&link);CHKERRQ(ierr); 158 ierr = PetscSFPackLeafData(sf,link,leafloc,leafdata,PETSC_TRUE);CHKERRQ(ierr); 159 160 /* Do neighborhood alltoallv for non-distinguished ranks */ 161 ierr = PetscSFGetDistComm_Neighbor(sf,PETSCSF_LEAF2ROOT_REDUCE,&distcomm);CHKERRQ(ierr); 162 if (use_gpu_aware_mpi) { 163 rootmtype_mpi = rootmtype; 164 leafmtype_mpi = leafmtype; 165 } else { 166 rootmtype_mpi = PETSC_MEMTYPE_HOST; 167 leafmtype_mpi = PETSC_MEMTYPE_HOST; 168 } 169 ierr = MPI_Start_ineighbor_alltoallv(dat->leafdegree,dat->rootdegree,link->leafbuf[leafmtype_mpi],dat->leafcounts,dat->leafdispls,unit,link->rootbuf[rootmtype_mpi],dat->rootcounts,dat->rootdispls,unit,distcomm,link->rootreqs[PETSCSF_LEAF2ROOT_REDUCE][rootmtype_mpi]);CHKERRQ(ierr); 170 if (rootmtype != leafmtype) {ierr = PetscMemcpyWithMemType(rootmtype,leafmtype,link->selfbuf[rootmtype],link->selfbuf[leafmtype],link->selfbuflen*link->unitbytes);CHKERRQ(ierr);} 171 PetscFunctionReturn(0); 172 } 173 174 static PetscErrorCode PetscSFFetchAndOpEnd_Neighbor(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,void *rootdata,PetscMemType leafmtype,const void *leafdata,void *leafupdate,MPI_Op op) 175 { 176 PetscErrorCode ierr; 177 PetscSFPack link; 178 const PetscInt *rootloc = NULL,*leafloc = NULL; 179 PetscSF_Neighbor *dat = (PetscSF_Neighbor*)sf->data; 180 PetscMemType rootmtype_mpi,leafmtype_mpi; /* memtypes seen by MPI */ 181 MPI_Comm distcomm = MPI_COMM_NULL; 182 183 PetscFunctionBegin; 184 ierr = PetscSFPackGetInUse(sf,unit,rootdata,leafdata,PETSC_OWN_POINTER,&link);CHKERRQ(ierr); 185 ierr = PetscSFPackWaitall(link,PETSCSF_LEAF2ROOT_REDUCE);CHKERRQ(ierr); 186 ierr = PetscSFGetRootIndicesWithMemType_Basic(sf,rootmtype,&rootloc);CHKERRQ(ierr); 187 ierr = PetscSFGetLeafIndicesWithMemType_Basic(sf,leafmtype,&leafloc);CHKERRQ(ierr); 188 /* Process local fetch-and-op */ 189 ierr = PetscSFFetchAndOpRootData(sf,link,rootloc,rootdata,op,PETSC_TRUE);CHKERRQ(ierr); 190 191 /* Bcast the updated rootbuf back to leaves */ 192 ierr = PetscSFGetDistComm_Neighbor(sf,PETSCSF_ROOT2LEAF_BCAST,&distcomm);CHKERRQ(ierr); 193 if (use_gpu_aware_mpi) { 194 rootmtype_mpi = rootmtype; 195 leafmtype_mpi = leafmtype; 196 } else { 197 rootmtype_mpi = PETSC_MEMTYPE_HOST; 198 leafmtype_mpi = PETSC_MEMTYPE_HOST; 199 } 200 ierr = MPI_Start_neighbor_alltoallv(dat->rootdegree,dat->leafdegree,link->rootbuf[rootmtype_mpi],dat->rootcounts,dat->rootdispls,unit,link->leafbuf[leafmtype_mpi],dat->leafcounts,dat->leafdispls,unit,distcomm);CHKERRQ(ierr); 201 if (rootmtype != leafmtype) {ierr = PetscMemcpyWithMemType(leafmtype,rootmtype,link->selfbuf[leafmtype],link->selfbuf[rootmtype],link->selfbuflen*link->unitbytes);CHKERRQ(ierr);} 202 ierr = PetscSFUnpackAndOpLeafData(sf,link,leafloc,leafupdate,MPIU_REPLACE,PETSC_TRUE);CHKERRQ(ierr); 203 ierr = PetscSFPackReclaim(sf,&link);CHKERRQ(ierr); 204 PetscFunctionReturn(0); 205 } 206 207 PETSC_INTERN PetscErrorCode PetscSFCreate_Neighbor(PetscSF sf) 208 { 209 PetscErrorCode ierr; 210 PetscSF_Neighbor *dat; 211 212 PetscFunctionBegin; 213 sf->ops->CreateEmbeddedSF = PetscSFCreateEmbeddedSF_Basic; 214 sf->ops->CreateEmbeddedLeafSF = PetscSFCreateEmbeddedLeafSF_Basic; 215 sf->ops->BcastAndOpEnd = PetscSFBcastAndOpEnd_Basic; 216 sf->ops->ReduceEnd = PetscSFReduceEnd_Basic; 217 sf->ops->FetchAndOpBegin = PetscSFFetchAndOpBegin_Basic; 218 sf->ops->GetLeafRanks = PetscSFGetLeafRanks_Basic; 219 sf->ops->View = PetscSFView_Basic; 220 sf->ops->SetFromOptions = PetscSFSetFromOptions_Basic; 221 222 sf->ops->SetUp = PetscSFSetUp_Neighbor; 223 sf->ops->Reset = PetscSFReset_Neighbor; 224 sf->ops->Destroy = PetscSFDestroy_Neighbor; 225 sf->ops->BcastAndOpBegin = PetscSFBcastAndOpBegin_Neighbor; 226 sf->ops->ReduceBegin = PetscSFReduceBegin_Neighbor; 227 sf->ops->FetchAndOpEnd = PetscSFFetchAndOpEnd_Neighbor; 228 229 ierr = PetscNewLog(sf,&dat);CHKERRQ(ierr); 230 sf->data = (void*)dat; 231 PetscFunctionReturn(0); 232 } 233 #endif 234