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