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 static PetscErrorCode PetscSFLinkStartCommunication_Gatherv(PetscSF sf, PetscSFLink link, PetscSFDirection direction) 9 { 10 MPI_Comm comm = MPI_COMM_NULL; 11 PetscMPIInt count; 12 PetscSF_Gatherv *dat = (PetscSF_Gatherv *)sf->data; 13 void *rootbuf = NULL, *leafbuf = NULL; /* buffer seen by MPI */ 14 MPI_Request *req = NULL; 15 MPI_Datatype unit = link->unit; 16 17 PetscFunctionBegin; 18 if (direction == PETSCSF_ROOT2LEAF) { 19 PetscCall(PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host before sending */)); 20 } else { 21 PetscCall(PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host */)); 22 } 23 PetscCall(PetscObjectGetComm((PetscObject)sf, &comm)); 24 PetscCall(PetscMPIIntCast(sf->nroots, &count)); 25 PetscCall(PetscSFLinkGetMPIBuffersAndRequests(sf, link, direction, &rootbuf, &leafbuf, &req, NULL)); 26 PetscCall(PetscSFLinkSyncStreamBeforeCallMPI(sf, link, direction)); 27 28 if (direction == PETSCSF_ROOT2LEAF) { 29 PetscCallMPI(MPIU_Igatherv(rootbuf, count, unit, leafbuf, dat->recvcounts, dat->displs, unit, 0 /*rank 0*/, comm, req)); 30 } else { 31 PetscCallMPI(MPIU_Iscatterv(leafbuf, dat->recvcounts, dat->displs, unit, rootbuf, count, unit, 0, comm, req)); 32 } 33 PetscFunctionReturn(PETSC_SUCCESS); 34 } 35 36 static PetscErrorCode PetscSFSetCommunicationOps_Gatherv(PetscSF sf, PetscSFLink link) 37 { 38 PetscFunctionBegin; 39 link->StartCommunication = PetscSFLinkStartCommunication_Gatherv; 40 PetscFunctionReturn(PETSC_SUCCESS); 41 } 42 43 PETSC_INTERN PetscErrorCode PetscSFFetchAndOpBegin_Gatherv(PetscSF sf, MPI_Datatype unit, PetscMemType rootmtype, void *rootdata, PetscMemType leafmtype, const void *leafdata, void *leafupdate, MPI_Op op) 44 { 45 PetscFunctionBegin; 46 /* In Gatherv, each root only has one leaf. So we just need to bcast rootdata to leafupdate and then reduce leafdata to rootdata */ 47 PetscCall(PetscSFBcastBegin(sf, unit, rootdata, leafupdate, MPI_REPLACE)); 48 PetscCall(PetscSFBcastEnd(sf, unit, rootdata, leafupdate, MPI_REPLACE)); 49 PetscCall(PetscSFReduceBegin(sf, unit, leafdata, rootdata, op)); 50 PetscFunctionReturn(PETSC_SUCCESS); 51 } 52 53 PETSC_INTERN PetscErrorCode PetscSFCreate_Gatherv(PetscSF sf) 54 { 55 PetscSF_Gatherv *dat = (PetscSF_Gatherv *)sf->data; 56 57 PetscFunctionBegin; 58 sf->ops->BcastBegin = PetscSFBcastBegin_Basic; 59 sf->ops->BcastEnd = PetscSFBcastEnd_Basic; 60 sf->ops->ReduceBegin = PetscSFReduceBegin_Basic; 61 sf->ops->ReduceEnd = PetscSFReduceEnd_Basic; 62 63 /* Inherit from Allgatherv */ 64 sf->ops->SetUp = PetscSFSetUp_Allgatherv; 65 sf->ops->Reset = PetscSFReset_Allgatherv; 66 sf->ops->Destroy = PetscSFDestroy_Allgatherv; 67 sf->ops->GetGraph = PetscSFGetGraph_Allgatherv; 68 sf->ops->GetLeafRanks = PetscSFGetLeafRanks_Allgatherv; 69 sf->ops->GetRootRanks = PetscSFGetRootRanks_Allgatherv; 70 sf->ops->FetchAndOpEnd = PetscSFFetchAndOpEnd_Allgatherv; 71 sf->ops->CreateLocalSF = PetscSFCreateLocalSF_Allgatherv; 72 73 /* Gatherv stuff */ 74 sf->ops->FetchAndOpBegin = PetscSFFetchAndOpBegin_Gatherv; 75 76 sf->ops->SetCommunicationOps = PetscSFSetCommunicationOps_Gatherv; 77 78 sf->collective = PETSC_TRUE; 79 80 PetscCall(PetscNew(&dat)); 81 sf->data = (void *)dat; 82 PetscFunctionReturn(PETSC_SUCCESS); 83 } 84