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 CHKERRQ(PetscSFLinkCreate(sf,unit,rootmtype,rootdata,leafmtype,leafdata,op,PETSCSF_BCAST,&link)); 17 CHKERRQ(PetscSFLinkPackRootData(sf,link,PETSCSF_REMOTE,rootdata)); 18 CHKERRQ(PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(sf,link,PETSC_TRUE/* device2host before sending */)); 19 CHKERRQ(PetscObjectGetComm((PetscObject)sf,&comm)); 20 CHKERRQ(PetscMPIIntCast(sf->nroots,&sendcount)); 21 CHKERRQ(PetscSFLinkGetMPIBuffersAndRequests(sf,link,PETSCSF_ROOT2LEAF,&rootbuf,&leafbuf,&req,NULL)); 22 CHKERRQ(PetscSFLinkSyncStreamBeforeCallMPI(sf,link,PETSCSF_ROOT2LEAF)); 23 CHKERRMPI(MPIU_Igather(rootbuf == leafbuf ? MPI_IN_PLACE : rootbuf,sendcount,unit,leafbuf,sendcount,unit,0/*rank 0*/,comm,req)); 24 PetscFunctionReturn(0); 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 CHKERRQ(PetscSFLinkCreate(sf,unit,rootmtype,rootdata,leafmtype,leafdata,op,PETSCSF_REDUCE,&link)); 37 CHKERRQ(PetscSFLinkPackLeafData(sf,link,PETSCSF_REMOTE,leafdata)); 38 CHKERRQ(PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(sf,link,PETSC_TRUE/* device2host before sending */)); 39 CHKERRQ(PetscObjectGetComm((PetscObject)sf,&comm)); 40 CHKERRQ(PetscMPIIntCast(sf->nroots,&recvcount)); 41 CHKERRQ(PetscSFLinkGetMPIBuffersAndRequests(sf,link,PETSCSF_LEAF2ROOT,&rootbuf,&leafbuf,&req,NULL)); 42 CHKERRQ(PetscSFLinkSyncStreamBeforeCallMPI(sf,link,PETSCSF_LEAF2ROOT)); 43 CHKERRMPI(MPIU_Iscatter(leafbuf,recvcount,unit,rootbuf == leafbuf ? MPI_IN_PLACE : rootbuf,recvcount,unit,0/*rank 0*/,comm,req)); 44 PetscFunctionReturn(0); 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 CHKERRQ(PetscNewLog(sf,&dat)); 75 sf->data = (void*)dat; 76 PetscFunctionReturn(0); 77 } 78