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