xref: /petsc/src/vec/is/sf/impls/basic/neighbor/sfneighbor.c (revision f1cab4e1e924af9242bb4b99cbdea45bd01822e9)
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