1 #include <../src/vec/is/sf/impls/basic/gatherv/sfgatherv.h> 2 #include <../src/vec/is/sf/impls/basic/allgather/sfallgather.h> 3 4 /* Reuse the type. The difference is some fields (i.e., displs, recvcounts) are not used in Gather, which is not a big deal */ 5 typedef PetscSF_Allgatherv PetscSF_Gather; 6 7 PETSC_INTERN PetscErrorCode PetscSFBcastBegin_Gather(PetscSF sf, MPI_Datatype unit, PetscMemType rootmtype, const void *rootdata, PetscMemType leafmtype, void *leafdata, MPI_Op op) { 8 PetscSFLink link; 9 PetscMPIInt sendcount; 10 MPI_Comm comm; 11 void *rootbuf = NULL, *leafbuf = NULL; 12 MPI_Request *req; 13 14 PetscFunctionBegin; 15 PetscCall(PetscSFLinkCreate(sf, unit, rootmtype, rootdata, leafmtype, leafdata, op, PETSCSF_BCAST, &link)); 16 PetscCall(PetscSFLinkPackRootData(sf, link, PETSCSF_REMOTE, rootdata)); 17 PetscCall(PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host before sending */)); 18 PetscCall(PetscObjectGetComm((PetscObject)sf, &comm)); 19 PetscCall(PetscMPIIntCast(sf->nroots, &sendcount)); 20 PetscCall(PetscSFLinkGetMPIBuffersAndRequests(sf, link, PETSCSF_ROOT2LEAF, &rootbuf, &leafbuf, &req, NULL)); 21 PetscCall(PetscSFLinkSyncStreamBeforeCallMPI(sf, link, PETSCSF_ROOT2LEAF)); 22 PetscCallMPI(MPIU_Igather(rootbuf == leafbuf ? MPI_IN_PLACE : rootbuf, sendcount, unit, leafbuf, sendcount, unit, 0 /*rank 0*/, comm, req)); 23 PetscFunctionReturn(0); 24 } 25 26 static PetscErrorCode PetscSFReduceBegin_Gather(PetscSF sf, MPI_Datatype unit, PetscMemType leafmtype, const void *leafdata, PetscMemType rootmtype, void *rootdata, MPI_Op op) { 27 PetscSFLink link; 28 PetscMPIInt recvcount; 29 MPI_Comm comm; 30 void *rootbuf = NULL, *leafbuf = NULL; 31 MPI_Request *req; 32 33 PetscFunctionBegin; 34 PetscCall(PetscSFLinkCreate(sf, unit, rootmtype, rootdata, leafmtype, leafdata, op, PETSCSF_REDUCE, &link)); 35 PetscCall(PetscSFLinkPackLeafData(sf, link, PETSCSF_REMOTE, leafdata)); 36 PetscCall(PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host before sending */)); 37 PetscCall(PetscObjectGetComm((PetscObject)sf, &comm)); 38 PetscCall(PetscMPIIntCast(sf->nroots, &recvcount)); 39 PetscCall(PetscSFLinkGetMPIBuffersAndRequests(sf, link, PETSCSF_LEAF2ROOT, &rootbuf, &leafbuf, &req, NULL)); 40 PetscCall(PetscSFLinkSyncStreamBeforeCallMPI(sf, link, PETSCSF_LEAF2ROOT)); 41 PetscCallMPI(MPIU_Iscatter(leafbuf, recvcount, unit, rootbuf == leafbuf ? MPI_IN_PLACE : rootbuf, recvcount, unit, 0 /*rank 0*/, comm, req)); 42 PetscFunctionReturn(0); 43 } 44 45 PETSC_INTERN PetscErrorCode PetscSFCreate_Gather(PetscSF sf) { 46 PetscSF_Gather *dat = (PetscSF_Gather *)sf->data; 47 48 PetscFunctionBegin; 49 sf->ops->BcastEnd = PetscSFBcastEnd_Basic; 50 sf->ops->ReduceEnd = PetscSFReduceEnd_Basic; 51 52 /* Inherit from Allgatherv */ 53 sf->ops->Reset = PetscSFReset_Allgatherv; 54 sf->ops->Destroy = PetscSFDestroy_Allgatherv; 55 sf->ops->GetGraph = PetscSFGetGraph_Allgatherv; 56 sf->ops->GetRootRanks = PetscSFGetRootRanks_Allgatherv; 57 sf->ops->GetLeafRanks = PetscSFGetLeafRanks_Allgatherv; 58 sf->ops->FetchAndOpEnd = PetscSFFetchAndOpEnd_Allgatherv; 59 sf->ops->CreateLocalSF = PetscSFCreateLocalSF_Allgatherv; 60 61 /* Inherit from Allgather */ 62 sf->ops->SetUp = PetscSFSetUp_Allgather; 63 64 /* Inherit from Gatherv */ 65 sf->ops->FetchAndOpBegin = PetscSFFetchAndOpBegin_Gatherv; 66 67 /* Gather stuff */ 68 sf->ops->BcastBegin = PetscSFBcastBegin_Gather; 69 sf->ops->ReduceBegin = PetscSFReduceBegin_Gather; 70 71 PetscCall(PetscNewLog(sf, &dat)); 72 sf->data = (void *)dat; 73 PetscFunctionReturn(0); 74 } 75