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 { 9 PetscSFLink link; 10 PetscMPIInt sendcount; 11 MPI_Comm comm; 12 void *rootbuf = NULL, *leafbuf = NULL; 13 MPI_Request *req; 14 15 PetscFunctionBegin; 16 PetscCall(PetscSFLinkCreate(sf, unit, rootmtype, rootdata, leafmtype, leafdata, op, PETSCSF_BCAST, &link)); 17 PetscCall(PetscSFLinkPackRootData(sf, link, PETSCSF_REMOTE, rootdata)); 18 PetscCall(PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host before sending */)); 19 PetscCall(PetscObjectGetComm((PetscObject)sf, &comm)); 20 PetscCall(PetscMPIIntCast(sf->nroots, &sendcount)); 21 PetscCall(PetscSFLinkGetMPIBuffersAndRequests(sf, link, PETSCSF_ROOT2LEAF, &rootbuf, &leafbuf, &req, NULL)); 22 PetscCall(PetscSFLinkSyncStreamBeforeCallMPI(sf, link, PETSCSF_ROOT2LEAF)); 23 PetscCallMPI(MPIU_Igather(rootbuf == leafbuf ? MPI_IN_PLACE : rootbuf, sendcount, unit, leafbuf, sendcount, unit, 0 /*rank 0*/, comm, req)); 24 PetscFunctionReturn(PETSC_SUCCESS); 25 } 26 27 static PetscErrorCode PetscSFReduceBegin_Gather(PetscSF sf, MPI_Datatype unit, PetscMemType leafmtype, const void *leafdata, PetscMemType rootmtype, void *rootdata, MPI_Op op) 28 { 29 PetscSFLink link; 30 PetscMPIInt recvcount; 31 MPI_Comm comm; 32 void *rootbuf = NULL, *leafbuf = NULL; 33 MPI_Request *req; 34 35 PetscFunctionBegin; 36 PetscCall(PetscSFLinkCreate(sf, unit, rootmtype, rootdata, leafmtype, leafdata, op, PETSCSF_REDUCE, &link)); 37 PetscCall(PetscSFLinkPackLeafData(sf, link, PETSCSF_REMOTE, leafdata)); 38 PetscCall(PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host before sending */)); 39 PetscCall(PetscObjectGetComm((PetscObject)sf, &comm)); 40 PetscCall(PetscMPIIntCast(sf->nroots, &recvcount)); 41 PetscCall(PetscSFLinkGetMPIBuffersAndRequests(sf, link, PETSCSF_LEAF2ROOT, &rootbuf, &leafbuf, &req, NULL)); 42 PetscCall(PetscSFLinkSyncStreamBeforeCallMPI(sf, link, PETSCSF_LEAF2ROOT)); 43 PetscCallMPI(MPIU_Iscatter(leafbuf, recvcount, unit, rootbuf == leafbuf ? MPI_IN_PLACE : rootbuf, recvcount, unit, 0 /*rank 0*/, comm, req)); 44 PetscFunctionReturn(PETSC_SUCCESS); 45 } 46 47 PETSC_INTERN PetscErrorCode PetscSFCreate_Gather(PetscSF sf) 48 { 49 PetscSF_Gather *dat = (PetscSF_Gather *)sf->data; 50 51 PetscFunctionBegin; 52 sf->ops->BcastEnd = PetscSFBcastEnd_Basic; 53 sf->ops->ReduceEnd = PetscSFReduceEnd_Basic; 54 55 /* Inherit from Allgatherv */ 56 sf->ops->Reset = PetscSFReset_Allgatherv; 57 sf->ops->Destroy = PetscSFDestroy_Allgatherv; 58 sf->ops->GetGraph = PetscSFGetGraph_Allgatherv; 59 sf->ops->GetRootRanks = PetscSFGetRootRanks_Allgatherv; 60 sf->ops->GetLeafRanks = PetscSFGetLeafRanks_Allgatherv; 61 sf->ops->FetchAndOpEnd = PetscSFFetchAndOpEnd_Allgatherv; 62 sf->ops->CreateLocalSF = PetscSFCreateLocalSF_Allgatherv; 63 64 /* Inherit from Allgather */ 65 sf->ops->SetUp = PetscSFSetUp_Allgather; 66 67 /* Inherit from Gatherv */ 68 sf->ops->FetchAndOpBegin = PetscSFFetchAndOpBegin_Gatherv; 69 70 /* Gather stuff */ 71 sf->ops->BcastBegin = PetscSFBcastBegin_Gather; 72 sf->ops->ReduceBegin = PetscSFReduceBegin_Gather; 73 74 PetscCall(PetscNew(&dat)); 75 sf->data = (void *)dat; 76 PetscFunctionReturn(PETSC_SUCCESS); 77 } 78