1dd5b3ca6SJunchao Zhang #include <../src/vec/is/sf/impls/basic/gatherv/sfgatherv.h> 2cd620004SJunchao Zhang #include <../src/vec/is/sf/impls/basic/allgather/sfallgather.h> 3dd5b3ca6SJunchao Zhang 4dd5b3ca6SJunchao Zhang /* Reuse the type. The difference is some fields (i.e., displs, recvcounts) are not used in Gather, which is not a big deal */ 5dd5b3ca6SJunchao Zhang typedef PetscSF_Allgatherv PetscSF_Gather; 6dd5b3ca6SJunchao Zhang 7f5d27ee7SJunchao Zhang static PetscErrorCode PetscSFLinkStartCommunication_Gather(PetscSF sf, PetscSFLink link, PetscSFDirection direction) 8d71ae5a4SJacob Faibussowitsch { 9f5d27ee7SJunchao Zhang MPI_Comm comm = MPI_COMM_NULL; 10cd620004SJunchao Zhang void *rootbuf = NULL, *leafbuf = NULL; 11f5d27ee7SJunchao Zhang MPI_Request *req = NULL; 12f5d27ee7SJunchao Zhang PetscMPIInt count; 13f5d27ee7SJunchao Zhang MPI_Datatype unit = link->unit; 14dd5b3ca6SJunchao Zhang 15dd5b3ca6SJunchao Zhang PetscFunctionBegin; 16f5d27ee7SJunchao Zhang if (direction == PETSCSF_ROOT2LEAF) { 179566063dSJacob Faibussowitsch PetscCall(PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host before sending */)); 18f5d27ee7SJunchao Zhang } else { 19f5d27ee7SJunchao Zhang PetscCall(PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host */)); 20f5d27ee7SJunchao Zhang } 219566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)sf, &comm)); 22f5d27ee7SJunchao Zhang PetscCall(PetscMPIIntCast(sf->nroots, &count)); 23f5d27ee7SJunchao Zhang PetscCall(PetscSFLinkGetMPIBuffersAndRequests(sf, link, direction, &rootbuf, &leafbuf, &req, NULL)); 24*646b835dSJunchao Zhang PetscCall(PetscSFLinkSyncStreamBeforeCallMPI(sf, link)); 25f5d27ee7SJunchao Zhang 26f5d27ee7SJunchao Zhang if (direction == PETSCSF_ROOT2LEAF) { 27f5d27ee7SJunchao Zhang PetscCallMPI(MPIU_Igather(rootbuf == leafbuf ? MPI_IN_PLACE : rootbuf, count, unit, leafbuf, count, unit, 0 /*rank 0*/, comm, req)); 28f5d27ee7SJunchao Zhang } else { 29f5d27ee7SJunchao Zhang PetscCallMPI(MPIU_Iscatter(leafbuf, count, unit, rootbuf == leafbuf ? MPI_IN_PLACE : rootbuf, count, unit, 0 /*rank 0*/, comm, req)); 30f5d27ee7SJunchao Zhang } 313ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 32dd5b3ca6SJunchao Zhang } 33dd5b3ca6SJunchao Zhang 34f5d27ee7SJunchao Zhang static PetscErrorCode PetscSFSetCommunicationOps_Gather(PetscSF sf, PetscSFLink link) 35d71ae5a4SJacob Faibussowitsch { 36dd5b3ca6SJunchao Zhang PetscFunctionBegin; 37f5d27ee7SJunchao Zhang link->StartCommunication = PetscSFLinkStartCommunication_Gather; 383ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 39dd5b3ca6SJunchao Zhang } 40dd5b3ca6SJunchao Zhang 41d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode PetscSFCreate_Gather(PetscSF sf) 42d71ae5a4SJacob Faibussowitsch { 43dd5b3ca6SJunchao Zhang PetscSF_Gather *dat = (PetscSF_Gather *)sf->data; 44dd5b3ca6SJunchao Zhang 45dd5b3ca6SJunchao Zhang PetscFunctionBegin; 46f5d27ee7SJunchao Zhang sf->ops->BcastBegin = PetscSFBcastBegin_Basic; 47ad227feaSJunchao Zhang sf->ops->BcastEnd = PetscSFBcastEnd_Basic; 48f5d27ee7SJunchao Zhang sf->ops->ReduceBegin = PetscSFReduceBegin_Basic; 49cd620004SJunchao Zhang sf->ops->ReduceEnd = PetscSFReduceEnd_Basic; 50cd620004SJunchao Zhang 51dd5b3ca6SJunchao Zhang /* Inherit from Allgatherv */ 52dd5b3ca6SJunchao Zhang sf->ops->Reset = PetscSFReset_Allgatherv; 53dd5b3ca6SJunchao Zhang sf->ops->Destroy = PetscSFDestroy_Allgatherv; 54dd5b3ca6SJunchao Zhang sf->ops->GetGraph = PetscSFGetGraph_Allgatherv; 55dd5b3ca6SJunchao Zhang sf->ops->GetRootRanks = PetscSFGetRootRanks_Allgatherv; 56dd5b3ca6SJunchao Zhang sf->ops->GetLeafRanks = PetscSFGetLeafRanks_Allgatherv; 57dd5b3ca6SJunchao Zhang sf->ops->FetchAndOpEnd = PetscSFFetchAndOpEnd_Allgatherv; 58dd5b3ca6SJunchao Zhang sf->ops->CreateLocalSF = PetscSFCreateLocalSF_Allgatherv; 59dd5b3ca6SJunchao Zhang 60cd620004SJunchao Zhang /* Inherit from Allgather */ 61cd620004SJunchao Zhang sf->ops->SetUp = PetscSFSetUp_Allgather; 62cd620004SJunchao Zhang 63dd5b3ca6SJunchao Zhang /* Inherit from Gatherv */ 64dd5b3ca6SJunchao Zhang sf->ops->FetchAndOpBegin = PetscSFFetchAndOpBegin_Gatherv; 65dd5b3ca6SJunchao Zhang 66f5d27ee7SJunchao Zhang sf->ops->SetCommunicationOps = PetscSFSetCommunicationOps_Gather; 67dd5b3ca6SJunchao Zhang 686677b1c1SJunchao Zhang sf->collective = PETSC_TRUE; 696677b1c1SJunchao Zhang 704dfa11a4SJacob Faibussowitsch PetscCall(PetscNew(&dat)); 71dd5b3ca6SJunchao Zhang sf->data = (void *)dat; 723ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 73dd5b3ca6SJunchao Zhang } 74