Actual source code: sfgatherv.c
2: #include <../src/vec/is/sf/impls/basic/gatherv/sfgatherv.h>
4: /* Reuse the type. The difference is some fields (displs, recvcounts) are only significant
5: on rank 0 in Gatherv. On other ranks they are harmless NULL.
6: */
7: typedef PetscSF_Allgatherv PetscSF_Gatherv;
9: PETSC_INTERN PetscErrorCode PetscSFBcastBegin_Gatherv(PetscSF sf, MPI_Datatype unit, PetscMemType rootmtype, const void *rootdata, PetscMemType leafmtype, void *leafdata, MPI_Op op)
10: {
11: PetscSFLink link;
12: PetscMPIInt sendcount;
13: MPI_Comm comm;
14: PetscSF_Gatherv *dat = (PetscSF_Gatherv *)sf->data;
15: void *rootbuf = NULL, *leafbuf = NULL; /* buffer seen by MPI */
16: MPI_Request *req;
18: PetscFunctionBegin;
19: PetscCall(PetscSFLinkCreate(sf, unit, rootmtype, rootdata, leafmtype, leafdata, op, PETSCSF_BCAST, &link));
20: PetscCall(PetscSFLinkPackRootData(sf, link, PETSCSF_REMOTE, rootdata));
21: PetscCall(PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host before sending */));
22: PetscCall(PetscObjectGetComm((PetscObject)sf, &comm));
23: PetscCall(PetscMPIIntCast(sf->nroots, &sendcount));
24: PetscCall(PetscSFLinkGetMPIBuffersAndRequests(sf, link, PETSCSF_../../../../../../..2LEAF, &rootbuf, &leafbuf, &req, NULL));
25: PetscCall(PetscSFLinkSyncStreamBeforeCallMPI(sf, link, PETSCSF_../../../../../../..2LEAF));
26: PetscCallMPI(MPIU_Igatherv(rootbuf, sendcount, unit, leafbuf, dat->recvcounts, dat->displs, unit, 0 /*rank 0*/, comm, req));
27: PetscFunctionReturn(PETSC_SUCCESS);
28: }
30: static PetscErrorCode PetscSFReduceBegin_Gatherv(PetscSF sf, MPI_Datatype unit, PetscMemType leafmtype, const void *leafdata, PetscMemType rootmtype, void *rootdata, MPI_Op op)
31: {
32: PetscSFLink link;
33: PetscMPIInt recvcount;
34: MPI_Comm comm;
35: PetscSF_Gatherv *dat = (PetscSF_Gatherv *)sf->data;
36: void *rootbuf = NULL, *leafbuf = NULL; /* buffer seen by MPI */
37: MPI_Request *req;
39: PetscFunctionBegin;
40: PetscCall(PetscSFLinkCreate(sf, unit, rootmtype, rootdata, leafmtype, leafdata, op, PETSCSF_REDUCE, &link));
41: PetscCall(PetscSFLinkPackLeafData(sf, link, PETSCSF_REMOTE, leafdata));
42: PetscCall(PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host before sending */));
43: PetscCall(PetscObjectGetComm((PetscObject)sf, &comm));
44: PetscCall(PetscMPIIntCast(sf->nroots, &recvcount));
45: PetscCall(PetscSFLinkGetMPIBuffersAndRequests(sf, link, PETSCSF_LEAF2../../../../../../.., &rootbuf, &leafbuf, &req, NULL));
46: PetscCall(PetscSFLinkSyncStreamBeforeCallMPI(sf, link, PETSCSF_LEAF2../../../../../../..));
47: PetscCallMPI(MPIU_Iscatterv(leafbuf, dat->recvcounts, dat->displs, unit, rootbuf, recvcount, unit, 0, comm, req));
48: PetscFunctionReturn(PETSC_SUCCESS);
49: }
51: PETSC_INTERN PetscErrorCode PetscSFFetchAndOpBegin_Gatherv(PetscSF sf, MPI_Datatype unit, PetscMemType rootmtype, void *rootdata, PetscMemType leafmtype, const void *leafdata, void *leafupdate, MPI_Op op)
52: {
53: PetscFunctionBegin;
54: /* In Gatherv, each root only has one leaf. So we just need to bcast rootdata to leafupdate and then reduce leafdata to rootdata */
55: PetscCall(PetscSFBcastBegin(sf, unit, rootdata, leafupdate, MPI_REPLACE));
56: PetscCall(PetscSFBcastEnd(sf, unit, rootdata, leafupdate, MPI_REPLACE));
57: PetscCall(PetscSFReduceBegin(sf, unit, leafdata, rootdata, op));
58: PetscFunctionReturn(PETSC_SUCCESS);
59: }
61: PETSC_INTERN PetscErrorCode PetscSFCreate_Gatherv(PetscSF sf)
62: {
63: PetscSF_Gatherv *dat = (PetscSF_Gatherv *)sf->data;
65: PetscFunctionBegin;
66: sf->ops->BcastEnd = PetscSFBcastEnd_Basic;
67: sf->ops->ReduceEnd = PetscSFReduceEnd_Basic;
69: /* Inherit from Allgatherv */
70: sf->ops->SetUp = PetscSFSetUp_Allgatherv;
71: sf->ops->Reset = PetscSFReset_Allgatherv;
72: sf->ops->Destroy = PetscSFDestroy_Allgatherv;
73: sf->ops->GetGraph = PetscSFGetGraph_Allgatherv;
74: sf->ops->GetLeafRanks = PetscSFGetLeafRanks_Allgatherv;
75: sf->ops->GetRootRanks = PetscSFGetRootRanks_Allgatherv;
76: sf->ops->FetchAndOpEnd = PetscSFFetchAndOpEnd_Allgatherv;
77: sf->ops->CreateLocalSF = PetscSFCreateLocalSF_Allgatherv;
79: /* Gatherv stuff */
80: sf->ops->BcastBegin = PetscSFBcastBegin_Gatherv;
81: sf->ops->ReduceBegin = PetscSFReduceBegin_Gatherv;
82: sf->ops->FetchAndOpBegin = PetscSFFetchAndOpBegin_Gatherv;
84: PetscCall(PetscNew(&dat));
85: sf->data = (void *)dat;
86: PetscFunctionReturn(PETSC_SUCCESS);
87: }