Actual source code: sfalltoall.c
petsc-3.12.5 2020-03-29
1: #include <../src/vec/is/sf/impls/basic/allgatherv/sfallgatherv.h>
2: #include <../src/vec/is/sf/impls/basic/allgather/sfallgather.h>
3: #include <../src/vec/is/sf/impls/basic/gatherv/sfgatherv.h>
5: #define PetscSFPackGet_Alltoall PetscSFPackGet_Allgatherv
7: /* Reuse the type. The difference is some fields (i.e., displs, recvcounts) are not used, which is not a big deal */
8: typedef PetscSF_Allgatherv PetscSF_Alltoall;
10: /*===================================================================================*/
11: /* Implementations of SF public APIs */
12: /*===================================================================================*/
13: static PetscErrorCode PetscSFGetGraph_Alltoall(PetscSF sf,PetscInt *nroots,PetscInt *nleaves,const PetscInt **ilocal,const PetscSFNode **iremote)
14: {
16: PetscInt i;
19: if (nroots) *nroots = sf->nroots;
20: if (nleaves) *nleaves = sf->nleaves;
21: if (ilocal) *ilocal = NULL; /* Contiguous local indices */
22: if (iremote) {
23: if (!sf->remote) {
24: PetscMalloc1(sf->nleaves,&sf->remote);
25: sf->remote_alloc = sf->remote;
26: for (i=0; i<sf->nleaves; i++) {
27: sf->remote[i].rank = i;
28: sf->remote[i].index = i;
29: }
30: }
31: *iremote = sf->remote;
32: }
33: return(0);
34: }
36: static PetscErrorCode PetscSFBcastAndOpBegin_Alltoall(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,void *leafdata,MPI_Op op)
37: {
38: PetscErrorCode ierr;
39: PetscSFPack link;
40: MPI_Comm comm;
41: char *recvbuf;
44: PetscSFPackGet_Alltoall(sf,unit,rootmtype,rootdata,leafmtype,leafdata,&link);
45: PetscObjectGetComm((PetscObject)sf,&comm);
47: if (op != MPIU_REPLACE) {
48: if (!link->leafbuf[leafmtype]) {PetscMallocWithMemType(leafmtype,sf->nleaves*link->unitbytes,(void**)&link->leafbuf[leafmtype]);}
49: recvbuf = link->leafbuf[leafmtype];
50: } else {
51: recvbuf = (char*)leafdata;
52: }
53: MPIU_Ialltoall(rootdata,1,unit,recvbuf,1,unit,comm,link->rootreqs[PETSCSF_../../../../../..2LEAF_BCAST][rootmtype]);
54: return(0);
55: }
57: static PetscErrorCode PetscSFReduceBegin_Alltoall(PetscSF sf,MPI_Datatype unit,PetscMemType leafmtype,const void *leafdata,PetscMemType rootmtype,void *rootdata,MPI_Op op)
58: {
59: PetscErrorCode ierr;
60: PetscSFPack link;
61: MPI_Comm comm;
62: char *recvbuf;
65: PetscSFPackGet_Alltoall(sf,unit,rootmtype,rootdata,leafmtype,leafdata,&link);
66: PetscObjectGetComm((PetscObject)sf,&comm);
68: if (op != MPIU_REPLACE) {
69: if (!link->rootbuf[rootmtype]) {PetscMallocWithMemType(rootmtype,sf->nroots*link->unitbytes,(void**)&link->rootbuf[rootmtype]);}
70: recvbuf = link->rootbuf[rootmtype];
71: } else {
72: recvbuf = (char*)rootdata;
73: }
74: MPIU_Ialltoall(leafdata,1,unit,recvbuf,1,unit,comm,link->rootreqs[PETSCSF_LEAF2../../../../../.._REDUCE][rootmtype]);
75: return(0);
76: }
78: static PetscErrorCode PetscSFCreateLocalSF_Alltoall(PetscSF sf,PetscSF *out)
79: {
81: PetscInt nroots=1,nleaves=1,*ilocal;
82: PetscSFNode *iremote = NULL;
83: PetscSF lsf;
84: PetscMPIInt rank;
87: nroots = 1;
88: nleaves = 1;
89: MPI_Comm_rank(PetscObjectComm((PetscObject)sf),&rank);
90: PetscMalloc1(nleaves,&ilocal);
91: PetscMalloc1(nleaves,&iremote);
92: ilocal[0] = rank;
93: iremote[0].rank = 0; /* rank in PETSC_COMM_SELF */
94: iremote[0].index = rank; /* LocalSF is an embedded SF. Indices are not remapped */
96: PetscSFCreate(PETSC_COMM_SELF,&lsf);
97: PetscSFSetGraph(lsf,nroots,nleaves,NULL/*contiguous leaves*/,PETSC_OWN_POINTER,iremote,PETSC_OWN_POINTER);
98: PetscSFSetUp(lsf);
99: *out = lsf;
100: return(0);
101: }
103: static PetscErrorCode PetscSFCreateEmbeddedSF_Alltoall(PetscSF sf,PetscInt nselected,const PetscInt *selected,PetscSF *newsf)
104: {
106: PetscInt i,*tmproots,*ilocal,ndranks,ndiranks;
107: PetscSFNode *iremote;
108: PetscMPIInt nroots,*roots,nleaves,*leaves,rank;
109: MPI_Comm comm;
110: PetscSF_Basic *bas;
111: PetscSF esf;
114: PetscObjectGetComm((PetscObject)sf,&comm);
115: MPI_Comm_rank(comm,&rank);
117: /* Uniq selected[] and store the result in roots[] */
118: PetscMalloc1(nselected,&tmproots);
119: PetscArraycpy(tmproots,selected,nselected);
120: PetscSortRemoveDupsInt(&nselected,tmproots); /* nselected might be changed */
121: if (tmproots[0] < 0 || tmproots[nselected-1] >= sf->nroots) SETERRQ3(comm,PETSC_ERR_ARG_OUTOFRANGE,"Min/Max root indices %D/%D are not in [0,%D)",tmproots[0],tmproots[nselected-1],sf->nroots);
122: nroots = nselected; /* For Alltoall, we know root indices will not overflow MPI_INT */
123: PetscMalloc1(nselected,&roots);
124: for (i=0; i<nselected; i++) roots[i] = tmproots[i];
125: PetscFree(tmproots);
127: /* Find out which leaves are still connected to roots in the embedded sf. Expect PetscCommBuildTwoSided is more scalable than MPI_Alltoall */
128: PetscCommBuildTwoSided(comm,0/*empty msg*/,MPI_INT/*fake*/,nroots,roots,NULL/*todata*/,&nleaves,&leaves,NULL/*fromdata*/);
130: /* Move myself ahead if rank is in leaves[], since I am a distinguished rank */
131: ndranks = 0;
132: for (i=0; i<nleaves; i++) {
133: if (leaves[i] == rank) {leaves[i] = -rank; ndranks = 1; break;}
134: }
135: PetscSortMPIInt(nleaves,leaves);
136: if (nleaves && leaves[0] < 0) leaves[0] = rank;
138: /* Build esf and fill its fields manually (without calling PetscSFSetUp) */
139: PetscMalloc1(nleaves,&ilocal);
140: PetscMalloc1(nleaves,&iremote);
141: for (i=0; i<nleaves; i++) { /* 1:1 map from roots to leaves */
142: ilocal[i] = leaves[i];
143: iremote[i].rank = leaves[i];
144: iremote[i].index = leaves[i];
145: }
146: PetscSFCreate(comm,&esf);
147: PetscSFSetType(esf,PETSCSFBASIC); /* This optimized routine can only create a basic sf */
148: PetscSFSetGraph(esf,sf->nleaves,nleaves,ilocal,PETSC_OWN_POINTER,iremote,PETSC_OWN_POINTER);
150: /* As if we are calling PetscSFSetUpRanks(esf,self's group) */
151: PetscMalloc4(nleaves,&esf->ranks,nleaves+1,&esf->roffset,nleaves,&esf->rmine,nleaves,&esf->rremote);
152: esf->nranks = nleaves;
153: esf->ndranks = ndranks;
154: esf->roffset[0] = 0;
155: for (i=0; i<nleaves; i++) {
156: esf->ranks[i] = leaves[i];
157: esf->roffset[i+1] = i+1;
158: esf->rmine[i] = leaves[i];
159: esf->rremote[i] = leaves[i];
160: }
162: /* Set up esf->data, the incoming communication (i.e., recv info), which is usually done by PetscSFSetUp_Basic */
163: bas = (PetscSF_Basic*)esf->data;
164: PetscMalloc2(nroots,&bas->iranks,nroots+1,&bas->ioffset);
165: PetscMalloc1(nroots,&bas->irootloc);
166: /* Move myself ahead if rank is in roots[], since I am a distinguished irank */
167: ndiranks = 0;
168: for (i=0; i<nroots; i++) {
169: if (roots[i] == rank) {roots[i] = -rank; ndiranks = 1; break;}
170: }
171: PetscSortMPIInt(nroots,roots);
172: if (nroots && roots[0] < 0) roots[0] = rank;
174: bas->niranks = nroots;
175: bas->ndiranks = ndiranks;
176: bas->ioffset[0] = 0;
177: bas->itotal = nroots;
178: for (i=0; i<nroots; i++) {
179: bas->iranks[i] = roots[i];
180: bas->ioffset[i+1] = i+1;
181: bas->irootloc[i] = roots[i];
182: }
184: *newsf = esf;
185: return(0);
186: }
188: PETSC_INTERN PetscErrorCode PetscSFCreate_Alltoall(PetscSF sf)
189: {
190: PetscErrorCode ierr;
191: PetscSF_Alltoall *dat = (PetscSF_Alltoall*)sf->data;
194: /* Inherit from Allgatherv. It is astonishing Alltoall can inherit so much from Allgather(v) */
195: sf->ops->Destroy = PetscSFDestroy_Allgatherv;
196: sf->ops->Reset = PetscSFReset_Allgatherv;
197: sf->ops->BcastAndOpEnd = PetscSFBcastAndOpEnd_Allgatherv;
198: sf->ops->ReduceEnd = PetscSFReduceEnd_Allgatherv;
199: sf->ops->FetchAndOpEnd = PetscSFFetchAndOpEnd_Allgatherv;
200: sf->ops->GetRootRanks = PetscSFGetRootRanks_Allgatherv;
202: /* Inherit from Allgather. Every process gathers equal-sized data from others, which enables this inheritance. */
203: sf->ops->GetLeafRanks = PetscSFGetLeafRanks_Allgatherv;
205: /* Inherit from Gatherv. Each root has only one leaf connected, which enables this inheritance */
206: sf->ops->FetchAndOpBegin = PetscSFFetchAndOpBegin_Gatherv;
208: /* Alltoall stuff */
209: sf->ops->GetGraph = PetscSFGetGraph_Alltoall;
210: sf->ops->BcastAndOpBegin = PetscSFBcastAndOpBegin_Alltoall;
211: sf->ops->ReduceBegin = PetscSFReduceBegin_Alltoall;
212: sf->ops->CreateLocalSF = PetscSFCreateLocalSF_Alltoall;
213: sf->ops->CreateEmbeddedSF = PetscSFCreateEmbeddedSF_Alltoall;
215: PetscNewLog(sf,&dat);
216: sf->data = (void*)dat;
217: return(0);
218: }