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 /* SF Neighbor completely relies on the two sided info built by SF Basic. Therefore we build Neighbor as a subclass of Basic */ 7 8 typedef struct { 9 SPPACKBASICHEADER; 10 } *PetscSFPack_Neighbor; 11 12 typedef struct { 13 SFBASICHEADER; 14 MPI_Comm comms[2]; /* Communicators with distributed topology in both directions */ 15 PetscBool initialized[2]; /* Are the two communicators initialized? */ 16 PetscMPIInt *rootdispls,*rootcounts,*leafdispls,*leafcounts; /* displs/counts for non-distinguished ranks */ 17 } PetscSF_Neighbor; 18 19 /*===================================================================================*/ 20 /* Internal utility routines */ 21 /*===================================================================================*/ 22 23 /* Get the communicator with distributed graph topology, which is not cheap to build so we do it on demand (instead of at PetscSFSetUp time) */ 24 static PetscErrorCode PetscSFGetDistComm_Neighbor(PetscSF sf,PetscSFDirection direction,MPI_Comm *distcomm) 25 { 26 PetscErrorCode ierr; 27 PetscSF_Neighbor *dat = (PetscSF_Neighbor*)sf->data; 28 PetscInt nrootranks,ndrootranks,nleafranks,ndleafranks; 29 const PetscMPIInt *rootranks,*leafranks; 30 MPI_Comm comm; 31 32 PetscFunctionBegin; 33 ierr = PetscSFGetRootInfo_Basic(sf,&nrootranks,&ndrootranks,&rootranks,NULL,NULL);CHKERRQ(ierr); /* Which ranks will access my roots (I am a destination) */ 34 ierr = PetscSFGetLeafInfo_Basic(sf,&nleafranks,&ndleafranks,&leafranks,NULL,NULL,NULL);CHKERRQ(ierr); /* My leaves will access whose roots (I am a source) */ 35 36 if (!dat->initialized[direction]) { 37 const PetscMPIInt indegree = nrootranks-ndrootranks,*sources = rootranks+ndrootranks; 38 const PetscMPIInt outdegree = nleafranks-ndleafranks,*destinations = leafranks+ndleafranks; 39 MPI_Comm *mycomm = &dat->comms[direction]; 40 ierr = PetscObjectGetComm((PetscObject)sf,&comm);CHKERRQ(ierr); 41 if (direction == PETSCSF_LEAF2ROOT_REDUCE) { 42 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); 43 } else { /* PETSCSF_ROOT2LEAF_BCAST, reverse src & dest */ 44 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); 45 } 46 dat->initialized[direction] = PETSC_TRUE; 47 } 48 *distcomm = dat->comms[direction]; 49 PetscFunctionReturn(0); 50 } 51 52 static PetscErrorCode PetscSFPackGet_Neighbor(PetscSF sf,MPI_Datatype unit,const void *rootdata,const void *leafdata,PetscSFPack_Neighbor *mylink) 53 { 54 PetscErrorCode ierr; 55 56 PetscFunctionBegin; 57 /* Unlike SFBasic, SFNeighbor only needs one request. But we allocate two, to anticipate MPI will support persistent Neighbor requests in future. */ 58 ierr = PetscSFPackGet_Basic_Common(sf,unit,rootdata,leafdata,1/*half*/,(PetscSFPack_Basic*)mylink);CHKERRQ(ierr); 59 PetscFunctionReturn(0); 60 } 61 62 /*===================================================================================*/ 63 /* Implementations of SF public APIs */ 64 /*===================================================================================*/ 65 static PetscErrorCode PetscSFSetUp_Neighbor(PetscSF sf) 66 { 67 PetscErrorCode ierr; 68 PetscSF_Neighbor *dat = (PetscSF_Neighbor*)sf->data; 69 PetscInt i,j,nrootranks,ndrootranks,nleafranks,ndleafranks; 70 const PetscInt *rootoffset,*leafoffset; 71 PetscMPIInt m,n; 72 73 PetscFunctionBegin; 74 ierr = PetscSFSetUp_Basic(sf);CHKERRQ(ierr); 75 ierr = PetscSFGetRootInfo_Basic(sf,&nrootranks,&ndrootranks,NULL,&rootoffset,NULL);CHKERRQ(ierr); 76 ierr = PetscSFGetLeafInfo_Basic(sf,&nleafranks,&ndleafranks,NULL,&leafoffset,NULL,NULL);CHKERRQ(ierr); 77 78 /* Only setup MPI displs/counts for non-distinguished ranks. Distinguished ranks use shared memory */ 79 ierr = PetscMalloc4(nrootranks-ndrootranks,&dat->rootdispls,nrootranks-ndrootranks,&dat->rootcounts,nleafranks-ndleafranks,&dat->leafdispls,nleafranks-ndleafranks,&dat->leafcounts);CHKERRQ(ierr); 80 for (i=ndrootranks,j=0; i<nrootranks; i++,j++) { 81 ierr = PetscMPIIntCast(rootoffset[i]-rootoffset[ndrootranks],&m);CHKERRQ(ierr); dat->rootdispls[j] = m; 82 ierr = PetscMPIIntCast(rootoffset[i+1]-rootoffset[i], &n);CHKERRQ(ierr); dat->rootcounts[j] = n; 83 } 84 85 for (i=ndleafranks,j=0; i<nleafranks; i++,j++) { 86 ierr = PetscMPIIntCast(leafoffset[i]-leafoffset[ndleafranks],&m);CHKERRQ(ierr); dat->leafdispls[j] = m; 87 ierr = PetscMPIIntCast(leafoffset[i+1]-leafoffset[i], &n);CHKERRQ(ierr); dat->leafcounts[j] = n; 88 } 89 PetscFunctionReturn(0); 90 } 91 92 static PetscErrorCode PetscSFReset_Neighbor(PetscSF sf) 93 { 94 PetscErrorCode ierr; 95 PetscInt i; 96 PetscSF_Neighbor *dat = (PetscSF_Neighbor*)sf->data; 97 98 PetscFunctionBegin; 99 if (dat->inuse) SETERRQ(PetscObjectComm((PetscObject)sf),PETSC_ERR_ARG_WRONGSTATE,"Outstanding operation has not been completed"); 100 ierr = PetscFree4(dat->rootdispls,dat->rootcounts,dat->leafdispls,dat->leafcounts);CHKERRQ(ierr); 101 for (i=0; i<2; i++) { 102 if (dat->initialized[i]) { 103 ierr = MPI_Comm_free(&dat->comms[i]);CHKERRQ(ierr); 104 dat->initialized[i] = PETSC_FALSE; 105 } 106 } 107 ierr = PetscSFReset_Basic(sf);CHKERRQ(ierr); /* Common part */ 108 PetscFunctionReturn(0); 109 } 110 111 static PetscErrorCode PetscSFDestroy_Neighbor(PetscSF sf) 112 { 113 PetscErrorCode ierr; 114 115 PetscFunctionBegin; 116 ierr = PetscSFReset_Neighbor(sf);CHKERRQ(ierr); 117 ierr = PetscFree(sf->data);CHKERRQ(ierr); 118 PetscFunctionReturn(0); 119 } 120 121 static PetscErrorCode PetscSFBcastAndOpBegin_Neighbor(PetscSF sf,MPI_Datatype unit,const void *rootdata,void *leafdata,MPI_Op op) 122 { 123 PetscErrorCode ierr; 124 PetscSFPack_Neighbor link; 125 PetscInt nrootranks,ndrootranks,nleafranks,ndleafranks; 126 const PetscInt *rootoffset,*rootloc; 127 PetscMPIInt ind,outd; 128 PetscSF_Neighbor *dat = (PetscSF_Neighbor*)sf->data; 129 MPI_Comm distcomm; 130 131 PetscFunctionBegin; 132 ierr = PetscSFGetRootInfo_Basic(sf,&nrootranks,&ndrootranks,NULL,&rootoffset,&rootloc);CHKERRQ(ierr); 133 ierr = PetscSFGetLeafInfo_Basic(sf,&nleafranks,&ndleafranks,NULL,NULL,NULL,NULL);CHKERRQ(ierr); 134 ierr = PetscSFPackGet_Neighbor(sf,unit,rootdata,leafdata,&link);CHKERRQ(ierr); 135 ierr = PetscSFPackRootData(sf,(PetscSFPack)link,nrootranks,ndrootranks,rootoffset,rootloc,rootdata);CHKERRQ(ierr); 136 137 /* Do neighborhood alltoallv for non-distinguished ranks */ 138 ierr = PetscSFGetDistComm_Neighbor(sf,PETSCSF_ROOT2LEAF_BCAST,&distcomm);CHKERRQ(ierr); 139 outd = nrootranks - ndrootranks; 140 ind = nleafranks - ndleafranks; 141 ierr = MPI_Start_ineighbor_alltoallv(outd,ind,link->rootbuf,dat->rootcounts,dat->rootdispls,unit,link->leafbuf,dat->leafcounts,dat->leafdispls,unit,distcomm,&link->requests[PETSCSF_ROOT2LEAF_BCAST]);CHKERRQ(ierr); 142 PetscFunctionReturn(0); 143 } 144 145 static PetscErrorCode PetscSFReduceBegin_Neighbor(PetscSF sf,MPI_Datatype unit,const void *leafdata,void *rootdata,MPI_Op op) 146 { 147 PetscErrorCode ierr; 148 PetscInt nrootranks,ndrootranks,nleafranks,ndleafranks; 149 const PetscInt *leafoffset,*leafloc; 150 PetscMPIInt ind,outd; 151 PetscSFPack_Neighbor link; 152 PetscSF_Neighbor *dat = (PetscSF_Neighbor*)sf->data; 153 MPI_Comm distcomm; 154 155 PetscFunctionBegin; 156 ierr = PetscSFGetRootInfo_Basic(sf,&nrootranks,&ndrootranks,NULL,NULL,NULL);CHKERRQ(ierr); 157 ierr = PetscSFGetLeafInfo_Basic(sf,&nleafranks,&ndleafranks,NULL,&leafoffset,&leafloc,NULL);CHKERRQ(ierr); 158 ierr = PetscSFPackGet_Neighbor(sf,unit,rootdata,leafdata,&link);CHKERRQ(ierr); 159 ierr = PetscSFPackLeafData(sf,(PetscSFPack)link,nleafranks,ndleafranks,leafoffset,leafloc,leafdata);CHKERRQ(ierr); 160 161 /* Do neighborhood alltoallv for non-distinguished ranks */ 162 ierr = PetscSFGetDistComm_Neighbor(sf,PETSCSF_LEAF2ROOT_REDUCE,&distcomm);CHKERRQ(ierr); 163 ind = nrootranks - ndrootranks; 164 outd = nleafranks - ndleafranks; 165 ierr = MPI_Start_ineighbor_alltoallv(outd,ind,link->leafbuf,dat->leafcounts,dat->leafdispls,unit,link->rootbuf,dat->rootcounts,dat->rootdispls,unit,distcomm,&link->requests[PETSCSF_LEAF2ROOT_REDUCE]);CHKERRQ(ierr); 166 PetscFunctionReturn(0); 167 } 168 169 static PetscErrorCode PetscSFFetchAndOpEnd_Neighbor(PetscSF sf,MPI_Datatype unit,void *rootdata,const void *leafdata,void *leafupdate,MPI_Op op) 170 { 171 PetscErrorCode ierr; 172 PetscSFPack_Basic link; 173 PetscInt nrootranks,ndrootranks,nleafranks,ndleafranks; 174 const PetscInt *rootoffset,*leafoffset,*rootloc,*leafloc; 175 const PetscMPIInt *rootranks,*leafranks; 176 MPI_Comm comm; 177 PetscMPIInt ind,outd; 178 PetscSF_Neighbor *dat = (PetscSF_Neighbor*)sf->data; 179 180 PetscFunctionBegin; 181 ierr = PetscSFPackGetInUse(sf,unit,rootdata,leafdata,PETSC_OWN_POINTER,(PetscSFPack*)&link);CHKERRQ(ierr); 182 ierr = PetscSFPackWaitall_Basic(link,PETSCSF_LEAF2ROOT_REDUCE);CHKERRQ(ierr); 183 184 ierr = PetscSFGetRootInfo_Basic(sf,&nrootranks,&ndrootranks,&rootranks,&rootoffset,&rootloc);CHKERRQ(ierr); 185 ierr = PetscSFGetLeafInfo_Basic(sf,&nleafranks,&ndleafranks,&leafranks,&leafoffset,&leafloc,NULL);CHKERRQ(ierr); 186 187 /* Process local fetch-and-op */ 188 ierr = PetscSFFetchAndOpRootData(sf,(PetscSFPack)link,nrootranks,ndrootranks,rootoffset,rootloc,rootdata,op);CHKERRQ(ierr); 189 190 /* Bcast the updated root buffer back to leaves */ 191 ierr = PetscSFGetDistComm_Neighbor(sf,PETSCSF_ROOT2LEAF_BCAST,&comm);CHKERRQ(ierr); 192 outd = nrootranks - ndrootranks; 193 ind = nleafranks - ndleafranks; 194 ierr = MPI_Start_neighbor_alltoallv(outd,ind,link->rootbuf,dat->rootcounts,dat->rootdispls,unit,link->leafbuf,dat->leafcounts,dat->leafdispls,unit,comm);CHKERRQ(ierr); 195 ierr = PetscSFUnpackAndOpLeafData(sf,(PetscSFPack)link,nleafranks,ndleafranks,leafoffset,leafloc,leafupdate,MPIU_REPLACE);CHKERRQ(ierr); 196 ierr = PetscSFPackReclaim(sf,(PetscSFPack*)&link);CHKERRQ(ierr); 197 PetscFunctionReturn(0); 198 } 199 200 PETSC_INTERN PetscErrorCode PetscSFCreate_Neighbor(PetscSF sf) 201 { 202 PetscErrorCode ierr; 203 PetscSF_Neighbor *dat; 204 205 PetscFunctionBegin; 206 sf->ops->CreateEmbeddedSF = PetscSFCreateEmbeddedSF_Basic; 207 sf->ops->CreateEmbeddedLeafSF = PetscSFCreateEmbeddedLeafSF_Basic; 208 sf->ops->BcastAndOpEnd = PetscSFBcastAndOpEnd_Basic; 209 sf->ops->ReduceEnd = PetscSFReduceEnd_Basic; 210 sf->ops->FetchAndOpBegin = PetscSFFetchAndOpBegin_Basic; 211 sf->ops->GetLeafRanks = PetscSFGetLeafRanks_Basic; 212 sf->ops->View = PetscSFView_Basic; 213 214 sf->ops->SetUp = PetscSFSetUp_Neighbor; 215 sf->ops->Reset = PetscSFReset_Neighbor; 216 sf->ops->Destroy = PetscSFDestroy_Neighbor; 217 sf->ops->BcastAndOpBegin = PetscSFBcastAndOpBegin_Neighbor; 218 sf->ops->ReduceBegin = PetscSFReduceBegin_Neighbor; 219 sf->ops->FetchAndOpEnd = PetscSFFetchAndOpEnd_Neighbor; 220 221 ierr = PetscNewLog(sf,&dat);CHKERRQ(ierr); 222 sf->data = (void*)dat; 223 PetscFunctionReturn(0); 224 } 225 #endif 226