xref: /petsc/src/vec/is/sf/impls/basic/allgatherv/sfallgatherv.c (revision 661006243e0904bc232e219bc5c02c8d300d1d6d)
1dd5b3ca6SJunchao Zhang #include <../src/vec/is/sf/impls/basic/allgatherv/sfallgatherv.h>
2dd5b3ca6SJunchao Zhang 
3ad227feaSJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFBcastBegin_Gatherv(PetscSF, MPI_Datatype, PetscMemType, const void *, PetscMemType, void *, MPI_Op);
4dd5b3ca6SJunchao Zhang 
5dd5b3ca6SJunchao Zhang /* PetscSFGetGraph is non-collective. An implementation should not have collective calls */
6d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode PetscSFGetGraph_Allgatherv(PetscSF sf, PetscInt *nroots, PetscInt *nleaves, const PetscInt **ilocal, const PetscSFNode **iremote)
7d71ae5a4SJacob Faibussowitsch {
8dd5b3ca6SJunchao Zhang   PetscInt        i, j, k;
9dd5b3ca6SJunchao Zhang   const PetscInt *range;
10dd5b3ca6SJunchao Zhang   PetscMPIInt     size;
11dd5b3ca6SJunchao Zhang 
12dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
139566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)sf), &size));
14dd5b3ca6SJunchao Zhang   if (nroots) *nroots = sf->nroots;
15dd5b3ca6SJunchao Zhang   if (nleaves) *nleaves = sf->nleaves;
16dd5b3ca6SJunchao Zhang   if (ilocal) *ilocal = NULL; /* Contiguous leaves */
17dd5b3ca6SJunchao Zhang   if (iremote) {
18dd5b3ca6SJunchao Zhang     if (!sf->remote && sf->nleaves) { /* The && sf->nleaves makes sfgatherv able to inherit this routine */
199566063dSJacob Faibussowitsch       PetscCall(PetscLayoutGetRanges(sf->map, &range));
209566063dSJacob Faibussowitsch       PetscCall(PetscMalloc1(sf->nleaves, &sf->remote));
21dd5b3ca6SJunchao Zhang       sf->remote_alloc = sf->remote;
22dd5b3ca6SJunchao Zhang       for (i = 0; i < size; i++) {
23dd5b3ca6SJunchao Zhang         for (j = range[i], k = 0; j < range[i + 1]; j++, k++) {
24dd5b3ca6SJunchao Zhang           sf->remote[j].rank  = i;
25dd5b3ca6SJunchao Zhang           sf->remote[j].index = k;
26dd5b3ca6SJunchao Zhang         }
27dd5b3ca6SJunchao Zhang       }
28dd5b3ca6SJunchao Zhang     }
29dd5b3ca6SJunchao Zhang     *iremote = sf->remote;
30dd5b3ca6SJunchao Zhang   }
31dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
32dd5b3ca6SJunchao Zhang }
33dd5b3ca6SJunchao Zhang 
34d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode PetscSFSetUp_Allgatherv(PetscSF sf)
35d71ae5a4SJacob Faibussowitsch {
36dd5b3ca6SJunchao Zhang   PetscSF_Allgatherv *dat = (PetscSF_Allgatherv *)sf->data;
37dd5b3ca6SJunchao Zhang   PetscMPIInt         size;
38dd5b3ca6SJunchao Zhang   PetscInt            i;
39dd5b3ca6SJunchao Zhang   const PetscInt     *range;
40*66100624SStefano Zampini   MPI_Comm            comm;
41dd5b3ca6SJunchao Zhang 
42dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
439566063dSJacob Faibussowitsch   PetscCall(PetscSFSetUp_Allgather(sf));
44*66100624SStefano Zampini   PetscCall(PetscObjectGetComm((PetscObject)sf, &comm));
45*66100624SStefano Zampini   PetscCallMPI(MPI_Comm_size(comm, &size));
46dd5b3ca6SJunchao Zhang   if (sf->nleaves) { /* This if (sf->nleaves) test makes sfgatherv able to inherit this routine */
47*66100624SStefano Zampini     PetscBool isallgatherv = PETSC_FALSE;
48*66100624SStefano Zampini 
499566063dSJacob Faibussowitsch     PetscCall(PetscMalloc1(size, &dat->recvcounts));
509566063dSJacob Faibussowitsch     PetscCall(PetscMalloc1(size, &dat->displs));
519566063dSJacob Faibussowitsch     PetscCall(PetscLayoutGetRanges(sf->map, &range));
52dd5b3ca6SJunchao Zhang 
53dd5b3ca6SJunchao Zhang     for (i = 0; i < size; i++) {
549566063dSJacob Faibussowitsch       PetscCall(PetscMPIIntCast(range[i], &dat->displs[i]));
559566063dSJacob Faibussowitsch       PetscCall(PetscMPIIntCast(range[i + 1] - range[i], &dat->recvcounts[i]));
56dd5b3ca6SJunchao Zhang     }
57*66100624SStefano Zampini 
58*66100624SStefano Zampini     /* check if we actually have a one-to-all pattern */
59*66100624SStefano Zampini     PetscCall(PetscObjectTypeCompare((PetscObject)sf, PETSCSFALLGATHERV, &isallgatherv));
60*66100624SStefano Zampini     if (isallgatherv) {
61*66100624SStefano Zampini       PetscMPIInt rank, nRanksWithZeroRoots;
62*66100624SStefano Zampini 
63*66100624SStefano Zampini       nRanksWithZeroRoots = (sf->nroots == 0) ? 1 : 0; /* I have no roots */
64*66100624SStefano Zampini       PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &nRanksWithZeroRoots, 1, MPI_INT, MPI_SUM, comm));
65*66100624SStefano Zampini       if (nRanksWithZeroRoots == size - 1) { /* Only one rank has roots, which indicates a bcast pattern */
66*66100624SStefano Zampini         dat->bcast_pattern = PETSC_TRUE;
67*66100624SStefano Zampini         PetscCallMPI(MPI_Comm_rank(comm, &rank));
68*66100624SStefano Zampini         dat->bcast_root = sf->nroots > 0 ? rank : -1;
69*66100624SStefano Zampini         PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &dat->bcast_root, 1, MPI_INT, MPI_MAX, comm));
70*66100624SStefano Zampini       }
71*66100624SStefano Zampini     }
72dd5b3ca6SJunchao Zhang   }
73dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
74dd5b3ca6SJunchao Zhang }
75dd5b3ca6SJunchao Zhang 
76d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode PetscSFReset_Allgatherv(PetscSF sf)
77d71ae5a4SJacob Faibussowitsch {
78eb02082bSJunchao Zhang   PetscSF_Allgatherv *dat  = (PetscSF_Allgatherv *)sf->data;
7971438e86SJunchao Zhang   PetscSFLink         link = dat->avail, next;
80dd5b3ca6SJunchao Zhang 
81dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
829566063dSJacob Faibussowitsch   PetscCall(PetscFree(dat->iranks));
839566063dSJacob Faibussowitsch   PetscCall(PetscFree(dat->ioffset));
849566063dSJacob Faibussowitsch   PetscCall(PetscFree(dat->irootloc));
859566063dSJacob Faibussowitsch   PetscCall(PetscFree(dat->recvcounts));
869566063dSJacob Faibussowitsch   PetscCall(PetscFree(dat->displs));
8728b400f6SJacob Faibussowitsch   PetscCheck(!dat->inuse, PetscObjectComm((PetscObject)sf), PETSC_ERR_ARG_WRONGSTATE, "Outstanding operation has not been completed");
889371c9d4SSatish Balay   for (; link; link = next) {
899371c9d4SSatish Balay     next = link->next;
909371c9d4SSatish Balay     PetscCall(PetscSFLinkDestroy(sf, link));
919371c9d4SSatish Balay   }
9271438e86SJunchao Zhang   dat->avail = NULL;
93dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
94dd5b3ca6SJunchao Zhang }
95dd5b3ca6SJunchao Zhang 
96d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode PetscSFDestroy_Allgatherv(PetscSF sf)
97d71ae5a4SJacob Faibussowitsch {
98dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
999566063dSJacob Faibussowitsch   PetscCall(PetscSFReset_Allgatherv(sf));
1009566063dSJacob Faibussowitsch   PetscCall(PetscFree(sf->data));
101dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
102dd5b3ca6SJunchao Zhang }
103dd5b3ca6SJunchao Zhang 
104d71ae5a4SJacob Faibussowitsch static PetscErrorCode PetscSFBcastBegin_Allgatherv(PetscSF sf, MPI_Datatype unit, PetscMemType rootmtype, const void *rootdata, PetscMemType leafmtype, void *leafdata, MPI_Op op)
105d71ae5a4SJacob Faibussowitsch {
106cd620004SJunchao Zhang   PetscSFLink         link;
107*66100624SStefano Zampini   PetscMPIInt         sendcount, rank;
108dd5b3ca6SJunchao Zhang   MPI_Comm            comm;
109cd620004SJunchao Zhang   void               *rootbuf = NULL, *leafbuf = NULL;
110cd620004SJunchao Zhang   MPI_Request        *req;
111dd5b3ca6SJunchao Zhang   PetscSF_Allgatherv *dat = (PetscSF_Allgatherv *)sf->data;
112dd5b3ca6SJunchao Zhang 
113dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
1149566063dSJacob Faibussowitsch   PetscCall(PetscSFLinkCreate(sf, unit, rootmtype, rootdata, leafmtype, leafdata, op, PETSCSF_BCAST, &link));
1159566063dSJacob Faibussowitsch   PetscCall(PetscSFLinkPackRootData(sf, link, PETSCSF_REMOTE, rootdata));
1169566063dSJacob Faibussowitsch   PetscCall(PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host before sending */));
1179566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)sf, &comm));
118*66100624SStefano Zampini   PetscCallMPI(MPI_Comm_rank(comm, &rank));
1199566063dSJacob Faibussowitsch   PetscCall(PetscMPIIntCast(sf->nroots, &sendcount));
1209566063dSJacob Faibussowitsch   PetscCall(PetscSFLinkGetMPIBuffersAndRequests(sf, link, PETSCSF_ROOT2LEAF, &rootbuf, &leafbuf, &req, NULL));
121*66100624SStefano Zampini 
122*66100624SStefano Zampini   if (dat->bcast_pattern && rank == dat->bcast_root) PetscCall((*link->Memcpy)(link, link->leafmtype_mpi, leafbuf, link->rootmtype_mpi, rootbuf, (size_t)sendcount * link->unitbytes));
123*66100624SStefano Zampini   /* Ready the buffers for MPI */
1249566063dSJacob Faibussowitsch   PetscCall(PetscSFLinkSyncStreamBeforeCallMPI(sf, link, PETSCSF_ROOT2LEAF));
125*66100624SStefano Zampini   if (dat->bcast_pattern) PetscCallMPI(MPIU_Ibcast(leafbuf, sf->nleaves, unit, dat->bcast_root, comm, req));
126*66100624SStefano Zampini   else PetscCallMPI(MPIU_Iallgatherv(rootbuf, sendcount, unit, leafbuf, dat->recvcounts, dat->displs, unit, comm, req));
127855db38dSJunchao Zhang   PetscFunctionReturn(0);
128855db38dSJunchao Zhang }
129855db38dSJunchao Zhang 
130d71ae5a4SJacob Faibussowitsch static PetscErrorCode PetscSFReduceBegin_Allgatherv(PetscSF sf, MPI_Datatype unit, PetscMemType leafmtype, const void *leafdata, PetscMemType rootmtype, void *rootdata, MPI_Op op)
131d71ae5a4SJacob Faibussowitsch {
132cd620004SJunchao Zhang   PetscSFLink         link;
133dd5b3ca6SJunchao Zhang   PetscSF_Allgatherv *dat = (PetscSF_Allgatherv *)sf->data;
134dd5b3ca6SJunchao Zhang   PetscInt            rstart;
135cd620004SJunchao Zhang   PetscMPIInt         rank, count, recvcount;
136dd5b3ca6SJunchao Zhang   MPI_Comm            comm;
137cd620004SJunchao Zhang   void               *rootbuf = NULL, *leafbuf = NULL;
138cd620004SJunchao Zhang   MPI_Request        *req;
139dd5b3ca6SJunchao Zhang 
140dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
1419566063dSJacob Faibussowitsch   PetscCall(PetscSFLinkCreate(sf, unit, rootmtype, rootdata, leafmtype, leafdata, op, PETSCSF_REDUCE, &link));
14283df288dSJunchao Zhang   if (op == MPI_REPLACE) {
143cd620004SJunchao Zhang     /* REPLACE is only meaningful when all processes have the same leafdata to reduce. Therefore copying from local leafdata is fine */
1449566063dSJacob Faibussowitsch     PetscCall(PetscLayoutGetRange(sf->map, &rstart, NULL));
1459566063dSJacob Faibussowitsch     PetscCall((*link->Memcpy)(link, rootmtype, rootdata, leafmtype, (const char *)leafdata + (size_t)rstart * link->unitbytes, (size_t)sf->nroots * link->unitbytes));
1469566063dSJacob Faibussowitsch     if (PetscMemTypeDevice(leafmtype) && PetscMemTypeHost(rootmtype)) PetscCall((*link->SyncStream)(link));
147dd5b3ca6SJunchao Zhang   } else {
1489566063dSJacob Faibussowitsch     PetscCall(PetscObjectGetComm((PetscObject)sf, &comm));
1499566063dSJacob Faibussowitsch     PetscCall(PetscSFLinkPackLeafData(sf, link, PETSCSF_REMOTE, leafdata));
1509566063dSJacob Faibussowitsch     PetscCall(PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host before sending */));
1519566063dSJacob Faibussowitsch     PetscCall(PetscSFLinkGetMPIBuffersAndRequests(sf, link, PETSCSF_LEAF2ROOT, &rootbuf, &leafbuf, &req, NULL));
152*66100624SStefano Zampini     PetscCall(PetscSFLinkSyncStreamBeforeCallMPI(sf, link, PETSCSF_LEAF2ROOT));
153*66100624SStefano Zampini     if (dat->bcast_pattern) {
154*66100624SStefano Zampini #if defined(PETSC_HAVE_OMPI_MAJOR_VERSION) /* Workaround: cuda-aware OpenMPI-4.1.3 does not support MPI_Ireduce() with device buffers */
155*66100624SStefano Zampini       *req = MPI_REQUEST_NULL;             /* Set NULL so that we can safely MPI_Wait(req) */
156*66100624SStefano Zampini       PetscCallMPI(MPI_Reduce(leafbuf, rootbuf, sf->nleaves, unit, op, dat->bcast_root, comm));
157*66100624SStefano Zampini #else
158*66100624SStefano Zampini       PetscCallMPI(MPIU_Ireduce(leafbuf, rootbuf, sf->nleaves, unit, op, dat->bcast_root, comm, req));
159*66100624SStefano Zampini #endif
160*66100624SStefano Zampini     } else { /* Reduce leafdata, then scatter to rootdata */
161*66100624SStefano Zampini       PetscCallMPI(MPI_Comm_rank(comm, &rank));
1629566063dSJacob Faibussowitsch       PetscCall(PetscMPIIntCast(dat->rootbuflen[PETSCSF_REMOTE], &recvcount));
163cd620004SJunchao Zhang       /* Allocate a separate leaf buffer on rank 0 */
164dd400576SPatrick Sanan       if (rank == 0 && !link->leafbuf_alloc[PETSCSF_REMOTE][link->leafmtype_mpi]) {
1659566063dSJacob Faibussowitsch         PetscCall(PetscSFMalloc(sf, link->leafmtype_mpi, sf->leafbuflen[PETSCSF_REMOTE] * link->unitbytes, (void **)&link->leafbuf_alloc[PETSCSF_REMOTE][link->leafmtype_mpi]));
166dd5b3ca6SJunchao Zhang       }
167cd620004SJunchao Zhang       /* In case we already copied leafdata from device to host (i.e., no use_gpu_aware_mpi), we need to adjust leafbuf on rank 0 */
168dd400576SPatrick Sanan       if (rank == 0 && link->leafbuf_alloc[PETSCSF_REMOTE][link->leafmtype_mpi] == leafbuf) leafbuf = MPI_IN_PLACE;
1699566063dSJacob Faibussowitsch       PetscCall(PetscMPIIntCast(sf->nleaves * link->bs, &count));
170*66100624SStefano Zampini       PetscCallMPI(MPI_Reduce(leafbuf, link->leafbuf_alloc[PETSCSF_REMOTE][link->leafmtype_mpi], count, link->basicunit, op, 0, comm)); /* Must do reduce with MPI builtin datatype basicunit */
1719566063dSJacob Faibussowitsch       PetscCallMPI(MPIU_Iscatterv(link->leafbuf_alloc[PETSCSF_REMOTE][link->leafmtype_mpi], dat->recvcounts, dat->displs, unit, rootbuf, recvcount, unit, 0, comm, req));
172dd5b3ca6SJunchao Zhang     }
173*66100624SStefano Zampini   }
174eb02082bSJunchao Zhang   PetscFunctionReturn(0);
175eb02082bSJunchao Zhang }
176eb02082bSJunchao Zhang 
177d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode PetscSFReduceEnd_Allgatherv(PetscSF sf, MPI_Datatype unit, const void *leafdata, void *rootdata, MPI_Op op)
178d71ae5a4SJacob Faibussowitsch {
1799319200aSJunchao Zhang   PetscSFLink link;
1809319200aSJunchao Zhang 
1819319200aSJunchao Zhang   PetscFunctionBegin;
1829319200aSJunchao Zhang   if (op == MPI_REPLACE) {
1839319200aSJunchao Zhang     /* A rare case happens when op is MPI_REPLACE, using GPUs but no GPU aware MPI. In PetscSFReduceBegin_Allgather(v),
1849319200aSJunchao Zhang       we did a device to device copy and in effect finished the communication. But in PetscSFLinkFinishCommunication()
1859319200aSJunchao Zhang       of PetscSFReduceEnd_Basic(), it thinks since there is rootbuf, it calls PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI().
1866aad120cSJose E. Roman       It does a host to device memory copy on rootbuf, wrongly overwriting the results. So we don't overload
1879319200aSJunchao Zhang       PetscSFReduceEnd_Basic() in this case, and just reclaim the link.
1889319200aSJunchao Zhang      */
1899566063dSJacob Faibussowitsch     PetscCall(PetscSFLinkGetInUse(sf, unit, rootdata, leafdata, PETSC_OWN_POINTER, &link));
1909566063dSJacob Faibussowitsch     PetscCall(PetscSFLinkReclaim(sf, &link));
1919319200aSJunchao Zhang   } else {
1929566063dSJacob Faibussowitsch     PetscCall(PetscSFReduceEnd_Basic(sf, unit, leafdata, rootdata, op));
1939319200aSJunchao Zhang   }
1949319200aSJunchao Zhang   PetscFunctionReturn(0);
1959319200aSJunchao Zhang }
1969319200aSJunchao Zhang 
197d71ae5a4SJacob Faibussowitsch static PetscErrorCode PetscSFBcastToZero_Allgatherv(PetscSF sf, MPI_Datatype unit, PetscMemType rootmtype, const void *rootdata, PetscMemType leafmtype, void *leafdata)
198d71ae5a4SJacob Faibussowitsch {
199cd620004SJunchao Zhang   PetscSFLink link;
200855db38dSJunchao Zhang   PetscMPIInt rank;
201eb02082bSJunchao Zhang 
202eb02082bSJunchao Zhang   PetscFunctionBegin;
2039566063dSJacob Faibussowitsch   PetscCall(PetscSFBcastBegin_Gatherv(sf, unit, rootmtype, rootdata, leafmtype, leafdata, MPI_REPLACE));
2049566063dSJacob Faibussowitsch   PetscCall(PetscSFLinkGetInUse(sf, unit, rootdata, leafdata, PETSC_OWN_POINTER, &link));
2059566063dSJacob Faibussowitsch   PetscCall(PetscSFLinkFinishCommunication(sf, link, PETSCSF_ROOT2LEAF));
2069566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)sf), &rank));
20748a46eb9SPierre Jolivet   if (rank == 0 && PetscMemTypeDevice(leafmtype) && !sf->use_gpu_aware_mpi) PetscCall((*link->Memcpy)(link, PETSC_MEMTYPE_DEVICE, leafdata, PETSC_MEMTYPE_HOST, link->leafbuf[PETSC_MEMTYPE_HOST], sf->leafbuflen[PETSCSF_REMOTE] * link->unitbytes));
2089566063dSJacob Faibussowitsch   PetscCall(PetscSFLinkReclaim(sf, &link));
209dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
210dd5b3ca6SJunchao Zhang }
211dd5b3ca6SJunchao Zhang 
212dd5b3ca6SJunchao Zhang /* This routine is very tricky (I believe it is rarely used with this kind of graph so just provide a simple but not-optimal implementation).
213dd5b3ca6SJunchao Zhang 
214dd5b3ca6SJunchao Zhang    Suppose we have three ranks. Rank 0 has a root with value 1. Rank 0,1,2 has a leaf with value 2,3,4 respectively. The leaves are connected
215dd5b3ca6SJunchao Zhang    to the root on rank 0. Suppose op=MPI_SUM and rank 0,1,2 gets root state in their rank order. By definition of this routine, rank 0 sees 1
216dd5b3ca6SJunchao Zhang    in root, fetches it into its leafupate, then updates root to 1 + 2 = 3; rank 1 sees 3 in root, fetches it into its leafupate, then updates
217dd5b3ca6SJunchao Zhang    root to 3 + 3 = 6; rank 2 sees 6 in root, fetches it into its leafupdate, then updates root to 6 + 4 = 10.  At the end, leafupdate on rank
218dd5b3ca6SJunchao Zhang    0,1,2 is 1,3,6 respectively. root is 10.
219dd5b3ca6SJunchao Zhang 
220dd5b3ca6SJunchao Zhang    We use a simpler implementation. From the same initial state, we copy leafdata to leafupdate
221dd5b3ca6SJunchao Zhang              rank-0   rank-1    rank-2
222dd5b3ca6SJunchao Zhang         Root     1
223dd5b3ca6SJunchao Zhang         Leaf     2       3         4
224dd5b3ca6SJunchao Zhang      Leafupdate  2       3         4
225dd5b3ca6SJunchao Zhang 
226dd5b3ca6SJunchao Zhang    Do MPI_Exscan on leafupdate,
227dd5b3ca6SJunchao Zhang              rank-0   rank-1    rank-2
228dd5b3ca6SJunchao Zhang         Root     1
229dd5b3ca6SJunchao Zhang         Leaf     2       3         4
230dd5b3ca6SJunchao Zhang      Leafupdate  2       2         5
231dd5b3ca6SJunchao Zhang 
232dd5b3ca6SJunchao Zhang    BcastAndOp from root to leafupdate,
233dd5b3ca6SJunchao Zhang              rank-0   rank-1    rank-2
234dd5b3ca6SJunchao Zhang         Root     1
235dd5b3ca6SJunchao Zhang         Leaf     2       3         4
236dd5b3ca6SJunchao Zhang      Leafupdate  3       3         6
237dd5b3ca6SJunchao Zhang 
238dd5b3ca6SJunchao Zhang    Copy root to leafupdate on rank-0
239dd5b3ca6SJunchao Zhang              rank-0   rank-1    rank-2
240dd5b3ca6SJunchao Zhang         Root     1
241dd5b3ca6SJunchao Zhang         Leaf     2       3         4
242dd5b3ca6SJunchao Zhang      Leafupdate  1       3         6
243dd5b3ca6SJunchao Zhang 
244dd5b3ca6SJunchao Zhang    Reduce from leaf to root,
245dd5b3ca6SJunchao Zhang              rank-0   rank-1    rank-2
246dd5b3ca6SJunchao Zhang         Root     10
247dd5b3ca6SJunchao Zhang         Leaf     2       3         4
248dd5b3ca6SJunchao Zhang      Leafupdate  1       3         6
249dd5b3ca6SJunchao Zhang */
250d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode PetscSFFetchAndOpBegin_Allgatherv(PetscSF sf, MPI_Datatype unit, PetscMemType rootmtype, void *rootdata, PetscMemType leafmtype, const void *leafdata, void *leafupdate, MPI_Op op)
251d71ae5a4SJacob Faibussowitsch {
252cd620004SJunchao Zhang   PetscSFLink link;
253dd5b3ca6SJunchao Zhang   MPI_Comm    comm;
254dd5b3ca6SJunchao Zhang   PetscMPIInt count;
255dd5b3ca6SJunchao Zhang 
256dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
2579566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)sf, &comm));
25808401ef6SPierre Jolivet   PetscCheck(!PetscMemTypeDevice(rootmtype) && !PetscMemTypeDevice(leafmtype), comm, PETSC_ERR_SUP, "Do FetchAndOp on device");
259dd5b3ca6SJunchao Zhang   /* Copy leafdata to leafupdate */
2609566063dSJacob Faibussowitsch   PetscCall(PetscSFLinkCreate(sf, unit, rootmtype, rootdata, leafmtype, leafdata, op, PETSCSF_FETCH, &link));
2619566063dSJacob Faibussowitsch   PetscCall(PetscSFLinkPackLeafData(sf, link, PETSCSF_REMOTE, leafdata)); /* Sync the device */
2629566063dSJacob Faibussowitsch   PetscCall((*link->Memcpy)(link, leafmtype, leafupdate, leafmtype, leafdata, sf->nleaves * link->unitbytes));
2639566063dSJacob Faibussowitsch   PetscCall(PetscSFLinkGetInUse(sf, unit, rootdata, leafdata, PETSC_OWN_POINTER, &link));
264dd5b3ca6SJunchao Zhang 
265dd5b3ca6SJunchao Zhang   /* Exscan on leafupdate and then BcastAndOp rootdata to leafupdate */
26683df288dSJunchao Zhang   if (op == MPI_REPLACE) {
267dd5b3ca6SJunchao Zhang     PetscMPIInt size, rank, prev, next;
2689566063dSJacob Faibussowitsch     PetscCallMPI(MPI_Comm_rank(comm, &rank));
2699566063dSJacob Faibussowitsch     PetscCallMPI(MPI_Comm_size(comm, &size));
270dd5b3ca6SJunchao Zhang     prev = rank ? rank - 1 : MPI_PROC_NULL;
271dd5b3ca6SJunchao Zhang     next = (rank < size - 1) ? rank + 1 : MPI_PROC_NULL;
2729566063dSJacob Faibussowitsch     PetscCall(PetscMPIIntCast(sf->nleaves, &count));
2739566063dSJacob Faibussowitsch     PetscCallMPI(MPI_Sendrecv_replace(leafupdate, count, unit, next, link->tag, prev, link->tag, comm, MPI_STATUSES_IGNORE));
274cd620004SJunchao Zhang   } else {
2759566063dSJacob Faibussowitsch     PetscCall(PetscMPIIntCast(sf->nleaves * link->bs, &count));
2769566063dSJacob Faibussowitsch     PetscCallMPI(MPI_Exscan(MPI_IN_PLACE, leafupdate, count, link->basicunit, op, comm));
277cd620004SJunchao Zhang   }
2789566063dSJacob Faibussowitsch   PetscCall(PetscSFLinkReclaim(sf, &link));
2799566063dSJacob Faibussowitsch   PetscCall(PetscSFBcastBegin(sf, unit, rootdata, leafupdate, op));
2809566063dSJacob Faibussowitsch   PetscCall(PetscSFBcastEnd(sf, unit, rootdata, leafupdate, op));
281dd5b3ca6SJunchao Zhang 
282dd5b3ca6SJunchao Zhang   /* Bcast roots to rank 0's leafupdate */
2839566063dSJacob Faibussowitsch   PetscCall(PetscSFBcastToZero_Private(sf, unit, rootdata, leafupdate)); /* Using this line makes Allgather SFs able to inherit this routine */
284dd5b3ca6SJunchao Zhang 
285dd5b3ca6SJunchao Zhang   /* Reduce leafdata to rootdata */
2869566063dSJacob Faibussowitsch   PetscCall(PetscSFReduceBegin(sf, unit, leafdata, rootdata, op));
287dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
288dd5b3ca6SJunchao Zhang }
289dd5b3ca6SJunchao Zhang 
290d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode PetscSFFetchAndOpEnd_Allgatherv(PetscSF sf, MPI_Datatype unit, void *rootdata, const void *leafdata, void *leafupdate, MPI_Op op)
291d71ae5a4SJacob Faibussowitsch {
292dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
2939566063dSJacob Faibussowitsch   PetscCall(PetscSFReduceEnd(sf, unit, leafdata, rootdata, op));
294dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
295dd5b3ca6SJunchao Zhang }
296dd5b3ca6SJunchao Zhang 
297dd5b3ca6SJunchao Zhang /* Get root ranks accessing my leaves */
298d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode PetscSFGetRootRanks_Allgatherv(PetscSF sf, PetscInt *nranks, const PetscMPIInt **ranks, const PetscInt **roffset, const PetscInt **rmine, const PetscInt **rremote)
299d71ae5a4SJacob Faibussowitsch {
300dd5b3ca6SJunchao Zhang   PetscInt        i, j, k, size;
301dd5b3ca6SJunchao Zhang   const PetscInt *range;
302dd5b3ca6SJunchao Zhang 
303dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
304dd5b3ca6SJunchao Zhang   /* Lazily construct these large arrays if users really need them for this type of SF. Very likely, they do not */
305dd5b3ca6SJunchao Zhang   if (sf->nranks && !sf->ranks) { /* On rank!=0, sf->nranks=0. The sf->nranks test makes this routine also works for sfgatherv */
306dd5b3ca6SJunchao Zhang     size = sf->nranks;
3079566063dSJacob Faibussowitsch     PetscCall(PetscLayoutGetRanges(sf->map, &range));
3089566063dSJacob Faibussowitsch     PetscCall(PetscMalloc4(size, &sf->ranks, size + 1, &sf->roffset, sf->nleaves, &sf->rmine, sf->nleaves, &sf->rremote));
309dd5b3ca6SJunchao Zhang     for (i = 0; i < size; i++) sf->ranks[i] = i;
3109566063dSJacob Faibussowitsch     PetscCall(PetscArraycpy(sf->roffset, range, size + 1));
311dd5b3ca6SJunchao Zhang     for (i = 0; i < sf->nleaves; i++) sf->rmine[i] = i; /*rmine are never NULL even for contiguous leaves */
312dd5b3ca6SJunchao Zhang     for (i = 0; i < size; i++) {
313dd5b3ca6SJunchao Zhang       for (j = range[i], k = 0; j < range[i + 1]; j++, k++) sf->rremote[j] = k;
314dd5b3ca6SJunchao Zhang     }
315dd5b3ca6SJunchao Zhang   }
316dd5b3ca6SJunchao Zhang 
317dd5b3ca6SJunchao Zhang   if (nranks) *nranks = sf->nranks;
318dd5b3ca6SJunchao Zhang   if (ranks) *ranks = sf->ranks;
319dd5b3ca6SJunchao Zhang   if (roffset) *roffset = sf->roffset;
320dd5b3ca6SJunchao Zhang   if (rmine) *rmine = sf->rmine;
321dd5b3ca6SJunchao Zhang   if (rremote) *rremote = sf->rremote;
322dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
323dd5b3ca6SJunchao Zhang }
324dd5b3ca6SJunchao Zhang 
325dd5b3ca6SJunchao Zhang /* Get leaf ranks accessing my roots */
326d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode PetscSFGetLeafRanks_Allgatherv(PetscSF sf, PetscInt *niranks, const PetscMPIInt **iranks, const PetscInt **ioffset, const PetscInt **irootloc)
327d71ae5a4SJacob Faibussowitsch {
328dd5b3ca6SJunchao Zhang   PetscSF_Allgatherv *dat = (PetscSF_Allgatherv *)sf->data;
329dd5b3ca6SJunchao Zhang   MPI_Comm            comm;
330dd5b3ca6SJunchao Zhang   PetscMPIInt         size, rank;
331dd5b3ca6SJunchao Zhang   PetscInt            i, j;
332dd5b3ca6SJunchao Zhang 
333dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
334dd5b3ca6SJunchao Zhang   /* Lazily construct these large arrays if users really need them for this type of SF. Very likely, they do not */
3359566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)sf, &comm));
3369566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(comm, &size));
3379566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_rank(comm, &rank));
338dd5b3ca6SJunchao Zhang   if (niranks) *niranks = size;
339dd5b3ca6SJunchao Zhang 
340dd5b3ca6SJunchao Zhang   /* PetscSF_Basic has distinguished incoming ranks. Here we do not need that. But we must put self as the first and
341dd5b3ca6SJunchao Zhang      sort other ranks. See comments in PetscSFSetUp_Basic about MatGetBrowsOfAoCols_MPIAIJ on why.
342dd5b3ca6SJunchao Zhang    */
343dd5b3ca6SJunchao Zhang   if (iranks) {
344dd5b3ca6SJunchao Zhang     if (!dat->iranks) {
3459566063dSJacob Faibussowitsch       PetscCall(PetscMalloc1(size, &dat->iranks));
346dd5b3ca6SJunchao Zhang       dat->iranks[0] = rank;
3479371c9d4SSatish Balay       for (i = 0, j = 1; i < size; i++) {
3489371c9d4SSatish Balay         if (i == rank) continue;
3499371c9d4SSatish Balay         dat->iranks[j++] = i;
3509371c9d4SSatish Balay       }
351dd5b3ca6SJunchao Zhang     }
352*66100624SStefano Zampini     *iranks = dat->iranks; /* dat->iranks was init'ed to NULL by PetscNew */
353dd5b3ca6SJunchao Zhang   }
354dd5b3ca6SJunchao Zhang 
355dd5b3ca6SJunchao Zhang   if (ioffset) {
356dd5b3ca6SJunchao Zhang     if (!dat->ioffset) {
3579566063dSJacob Faibussowitsch       PetscCall(PetscMalloc1(size + 1, &dat->ioffset));
358dd5b3ca6SJunchao Zhang       for (i = 0; i <= size; i++) dat->ioffset[i] = i * sf->nroots;
359dd5b3ca6SJunchao Zhang     }
360dd5b3ca6SJunchao Zhang     *ioffset = dat->ioffset;
361dd5b3ca6SJunchao Zhang   }
362dd5b3ca6SJunchao Zhang 
363dd5b3ca6SJunchao Zhang   if (irootloc) {
364dd5b3ca6SJunchao Zhang     if (!dat->irootloc) {
3659566063dSJacob Faibussowitsch       PetscCall(PetscMalloc1(sf->nleaves, &dat->irootloc));
366dd5b3ca6SJunchao Zhang       for (i = 0; i < size; i++) {
367dd5b3ca6SJunchao Zhang         for (j = 0; j < sf->nroots; j++) dat->irootloc[i * sf->nroots + j] = j;
368dd5b3ca6SJunchao Zhang       }
369dd5b3ca6SJunchao Zhang     }
370dd5b3ca6SJunchao Zhang     *irootloc = dat->irootloc;
371dd5b3ca6SJunchao Zhang   }
372dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
373dd5b3ca6SJunchao Zhang }
374dd5b3ca6SJunchao Zhang 
375d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode PetscSFCreateLocalSF_Allgatherv(PetscSF sf, PetscSF *out)
376d71ae5a4SJacob Faibussowitsch {
377dd5b3ca6SJunchao Zhang   PetscInt     i, nroots, nleaves, rstart, *ilocal;
378dd5b3ca6SJunchao Zhang   PetscSFNode *iremote;
379dd5b3ca6SJunchao Zhang   PetscSF      lsf;
380dd5b3ca6SJunchao Zhang 
381dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
382eb02082bSJunchao Zhang   nleaves = sf->nleaves ? sf->nroots : 0; /* sf->nleaves can be zero with SFGather(v) */
383eb02082bSJunchao Zhang   nroots  = nleaves;
3849566063dSJacob Faibussowitsch   PetscCall(PetscMalloc1(nleaves, &ilocal));
3859566063dSJacob Faibussowitsch   PetscCall(PetscMalloc1(nleaves, &iremote));
3869566063dSJacob Faibussowitsch   PetscCall(PetscLayoutGetRange(sf->map, &rstart, NULL));
387dd5b3ca6SJunchao Zhang 
388dd5b3ca6SJunchao Zhang   for (i = 0; i < nleaves; i++) {
389dd5b3ca6SJunchao Zhang     ilocal[i]        = rstart + i; /* lsf does not change leave indices */
390dd5b3ca6SJunchao Zhang     iremote[i].rank  = 0;          /* rank in PETSC_COMM_SELF */
391dd5b3ca6SJunchao Zhang     iremote[i].index = i;          /* root index */
392dd5b3ca6SJunchao Zhang   }
393dd5b3ca6SJunchao Zhang 
3949566063dSJacob Faibussowitsch   PetscCall(PetscSFCreate(PETSC_COMM_SELF, &lsf));
3959566063dSJacob Faibussowitsch   PetscCall(PetscSFSetGraph(lsf, nroots, nleaves, ilocal, PETSC_OWN_POINTER, iremote, PETSC_OWN_POINTER));
3969566063dSJacob Faibussowitsch   PetscCall(PetscSFSetUp(lsf));
397dd5b3ca6SJunchao Zhang   *out = lsf;
398dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
399dd5b3ca6SJunchao Zhang }
400dd5b3ca6SJunchao Zhang 
401d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode PetscSFCreate_Allgatherv(PetscSF sf)
402d71ae5a4SJacob Faibussowitsch {
403dd5b3ca6SJunchao Zhang   PetscSF_Allgatherv *dat = (PetscSF_Allgatherv *)sf->data;
404dd5b3ca6SJunchao Zhang 
405dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
406ad227feaSJunchao Zhang   sf->ops->BcastEnd  = PetscSFBcastEnd_Basic;
4079319200aSJunchao Zhang   sf->ops->ReduceEnd = PetscSFReduceEnd_Allgatherv;
408cd620004SJunchao Zhang 
409dd5b3ca6SJunchao Zhang   sf->ops->SetUp           = PetscSFSetUp_Allgatherv;
410dd5b3ca6SJunchao Zhang   sf->ops->Reset           = PetscSFReset_Allgatherv;
411dd5b3ca6SJunchao Zhang   sf->ops->Destroy         = PetscSFDestroy_Allgatherv;
412dd5b3ca6SJunchao Zhang   sf->ops->GetRootRanks    = PetscSFGetRootRanks_Allgatherv;
413dd5b3ca6SJunchao Zhang   sf->ops->GetLeafRanks    = PetscSFGetLeafRanks_Allgatherv;
414dd5b3ca6SJunchao Zhang   sf->ops->GetGraph        = PetscSFGetGraph_Allgatherv;
415ad227feaSJunchao Zhang   sf->ops->BcastBegin      = PetscSFBcastBegin_Allgatherv;
416dd5b3ca6SJunchao Zhang   sf->ops->ReduceBegin     = PetscSFReduceBegin_Allgatherv;
417dd5b3ca6SJunchao Zhang   sf->ops->FetchAndOpBegin = PetscSFFetchAndOpBegin_Allgatherv;
418dd5b3ca6SJunchao Zhang   sf->ops->FetchAndOpEnd   = PetscSFFetchAndOpEnd_Allgatherv;
419dd5b3ca6SJunchao Zhang   sf->ops->CreateLocalSF   = PetscSFCreateLocalSF_Allgatherv;
420dd5b3ca6SJunchao Zhang   sf->ops->BcastToZero     = PetscSFBcastToZero_Allgatherv;
421dd5b3ca6SJunchao Zhang 
4224dfa11a4SJacob Faibussowitsch   PetscCall(PetscNew(&dat));
423*66100624SStefano Zampini   dat->bcast_root = -1;
424dd5b3ca6SJunchao Zhang   sf->data        = (void *)dat;
425dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
426dd5b3ca6SJunchao Zhang }
427