1 #include <../src/vec/is/sf/impls/basic/gatherv/sfgatherv.h> 2 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; 7 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; 16 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 } 28 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; 37 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 } 49 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 } 59 60 PETSC_INTERN PetscErrorCode PetscSFCreate_Gatherv(PetscSF sf) 61 { 62 PetscSF_Gatherv *dat = (PetscSF_Gatherv *)sf->data; 63 64 PetscFunctionBegin; 65 sf->ops->BcastEnd = PetscSFBcastEnd_Basic; 66 sf->ops->ReduceEnd = PetscSFReduceEnd_Basic; 67 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; 77 78 /* Gatherv stuff */ 79 sf->ops->BcastBegin = PetscSFBcastBegin_Gatherv; 80 sf->ops->ReduceBegin = PetscSFReduceBegin_Gatherv; 81 sf->ops->FetchAndOpBegin = PetscSFFetchAndOpBegin_Gatherv; 82 83 PetscCall(PetscNew(&dat)); 84 sf->data = (void *)dat; 85 PetscFunctionReturn(PETSC_SUCCESS); 86 } 87